Repository: opendatahub-io/distributed-workloads
Pull request #399: Adjust existing kfto e2e tests to use kueue
Author: login: abhijeet-dhumal
Labels: name: lgtm
Created: 2025-05-20T10:44:00Z | Merged: 2025-05-21T06:31:44Z

<!--- Provide a general summary of your changes in the Title above -->

## Description
<!--- Describe your changes in detail -->

## How Has This Been Tested?
<!--- Please describe in detail how you tested your changes. -->
<!--- Include details of your testing environment, and the tests you ran to -->
<!--- see how your change affects other areas of the code, etc. -->
```
export NOTEBOOK_IMAGE=quay.io/modh/odh-generic-data-science-notebook@sha256:1c585c849cc466fce6443f5c186fb9b642d21b5f3581515ada08bba2fbc3591f \
TEST_TRAINING_IMAGE=quay.io/modh/training@sha256:1d0caea3e5d56ff7d672954b1ad511e661df9bdb364d56879961169a4ca8dae0 # CUDA

# quay.io/modh/training@sha256:f98d70edc49203537d4fec005aaff68a4830b9c9ce2390bb4d7910b89c312b97 ROCm

go test  -timeout 45m -v ./tests/kfto -run TestMnistSDK 
```


## Merge criteria:
<!--- This PR will be merged by any repository approver when it meets all the points in the checklist -->
<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->

- [ ] The commits are squashed in a cohesive manner and have meaningful messages.
- [ ] Testing instructions have been added in the PR body (for PRs involving changes that are not immediately obvious).
- [ ] The developer has manually tested the changes and verified that the changes work



Diff:
diff --git a/tests/kfto/kfto_mnist_sdk_test.go b/tests/kfto/kfto_mnist_sdk_test.go
index 5458c957a..98f59b127 100644
--- a/tests/kfto/kfto_mnist_sdk_test.go
+++ b/tests/kfto/kfto_mnist_sdk_test.go
@@ -17,13 +17,18 @@ limitations under the License.
 package kfto
 
 import (
+	"fmt"
 	"strings"
 	"testing"
 	"time"
 
 	. "github.com/onsi/gomega"
 
+	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/kueue/apis/kueue/v1beta1"
 
 	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 	. "github.com/opendatahub-io/distributed-workloads/tests/common/support"
@@ -42,6 +47,41 @@ func TestMnistSDK(t *testing.T) {
 	// Create role binding with Namespace specific admin cluster role
 	CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin")
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceCPU, corev1.ResourceMemory, corev1.ResourceName(NVIDIA.ResourceLabel)},
+				Flavors: []v1beta1.FlavorQuotas{
+					{
+						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []v1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("1"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("4Gi"),
+							},
+							{
+								Name:         corev1.ResourceName(NVIDIA.ResourceLabel),
+								NominalQuota: resource.MustParse(fmt.Sprint(0)),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
+	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
+	CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name, AsDefaultQueue)
+
 	requiredChangesInNotebook := map[string]string{
 		"${api_url}":        GetOpenShiftApiUrl(test),
 		"${password}":       userToken,
@@ -80,7 +120,7 @@ func TestMnistSDK(t *testing.T) {
 	}()
 
 	// Make sure pytorch job is created
-	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp")).
+	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp"), TestTimeoutDouble).
 		Should(WithTransform(PyTorchJobConditionRunning, Equal(v1.ConditionTrue)))
 
 	// Make sure that the job eventually succeeds
diff --git a/tests/kfto/kfto_mnist_training_test.go b/tests/kfto/kfto_mnist_training_test.go
index f98111f08..1422634cb 100644
--- a/tests/kfto/kfto_mnist_training_test.go
+++ b/tests/kfto/kfto_mnist_training_test.go
@@ -27,6 +27,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/kueue/apis/kueue/v1beta1"
 
 	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 	. "github.com/opendatahub-io/distributed-workloads/tests/common/support"
@@ -82,10 +83,65 @@ func runKFTOPyTorchMnistJob(t *testing.T, accelerator Accelerator, image string,
 		"requirements.txt":           requirementsFileName,
 	})
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory")},
+				Flavors: []v1beta1.FlavorQuotas{
+					{
+						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []v1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("8"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("18Gi"),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	if accelerator.IsGpu() {
+		numGpus := (workerReplicas + 1) * numProcPerNode
+		cqSpec.ResourceGroups[0].CoveredResources = append(
+			cqSpec.ResourceGroups[0].CoveredResources,
+			corev1.ResourceName(accelerator.ResourceLabel),
+		)
+		cqSpec.ResourceGroups[0].Flavors[0].Resources = append(
+			cqSpec.ResourceGroups[0].Flavors[0].Resources,
+			v1beta1.ResourceQuota{
+				Name:         corev1.ResourceName(accelerator.ResourceLabel),
+				NominalQuota: resource.MustParse(fmt.Sprint(numGpus)),
+			},
+		)
+	}
+
+	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
+	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
+	localQueue := CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name, AsDefaultQueue)
+
 	// Create training PyTorch job
-	tuningJob := createKFTOPyTorchMnistJob(test, namespace.Name, *config, accelerator, workerReplicas, numProcPerNode, image)
+	tuningJob := createKFTOPyTorchMnistJob(test, namespace.Name, *config, accelerator, workerReplicas, numProcPerNode, image, localQueue)
 	defer test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace.Name).Delete(test.Ctx(), tuningJob.Name, *metav1.NewDeleteOptions(0))
 
+	// Make sure the Workload is created and running
+	test.Eventually(GetKueueWorkloads(test, namespace.Name), TestTimeoutMedium).
+		Should(
+			And(
+				HaveLen(1),
+				ContainElement(WithTransform(KueueWorkloadAdmitted, BeTrueBecause("Workload failed to be admitted"))),
+			),
+		)
+
 	// Make sure the PyTorch job is running
 	test.Eventually(PyTorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutDouble).
 		Should(WithTransform(PyTorchJobConditionRunning, Equal(corev1.ConditionTrue)))
@@ -96,7 +152,7 @@ func runKFTOPyTorchMnistJob(t *testing.T, accelerator Accelerator, image string,
 
 }
 
-func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.ConfigMap, accelerator Accelerator, workerReplicas int, numProcPerNode int, baseImage string) *kftov1.PyTorchJob {
+func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.ConfigMap, accelerator Accelerator, workerReplicas int, numProcPerNode int, baseImage string, localQueue *v1beta1.LocalQueue) *kftov1.PyTorchJob {
 	var backend string
 	if accelerator.IsGpu() {
 		backend = "nccl"
@@ -117,6 +173,9 @@ func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.Config
 		},
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: "kfto-mnist-",
+			Labels: map[string]string{
+				"kueue.x-k8s.io/queue-name": localQueue.Name,
+			},
 		},
 		Spec: kftov1.PyTorchJobSpec{
 			PyTorchReplicaSpecs: map[kftov1.ReplicaType]*kftov1.ReplicaSpec{
@@ -177,11 +236,11 @@ func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.Config
 									Resources: corev1.ResourceRequirements{
 										Requests: corev1.ResourceList{
 											corev1.ResourceCPU:    resource.MustParse(fmt.Sprintf("%d", numProcPerNode)),
-											corev1.ResourceMemory: resource.MustParse("6Gi"),
+											corev1.ResourceMemory: resource.MustParse("4Gi"),
 										},
 										Limits: corev1.ResourceList{
 											corev1.ResourceCPU:    resource.MustParse(fmt.Sprintf("%d", numProcPerNode)),
-											corev1.ResourceMemory: resource.MustParse("6Gi"),
+											corev1.ResourceMemory: resource.MustParse("4Gi"),
 										},
 									},
 								},
@@ -273,11 +332,11 @@ func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.Config
 									Resources: corev1.ResourceRequirements{
 										Requests: corev1.ResourceList{
 											corev1.ResourceCPU:    resource.MustParse(fmt.Sprintf("%d", numProcPerNode)),
-											corev1.ResourceMemory: resource.MustParse("6Gi"),
+											corev1.ResourceMemory: resource.MustParse("4Gi"),
 										},
 										Limits: corev1.ResourceList{
 											corev1.ResourceCPU:    resource.MustParse(fmt.Sprintf("%d", numProcPerNode)),
-											corev1.ResourceMemory: resource.MustParse("6Gi"),
+											corev1.ResourceMemory: resource.MustParse("4Gi"),
 										},
 									},
 								},
diff --git a/tests/kfto/kfto_pytorchjob_failed_test.go b/tests/kfto/kfto_pytorchjob_failed_test.go
index a8ae1d842..be8c3f63d 100644
--- a/tests/kfto/kfto_pytorchjob_failed_test.go
+++ b/tests/kfto/kfto_pytorchjob_failed_test.go
@@ -9,6 +9,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/kueue/apis/kueue/v1beta1"
 
 	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 	. "github.com/opendatahub-io/distributed-workloads/tests/common/support"
@@ -30,15 +31,46 @@ func runFailedPyTorchJobTest(t *testing.T, image string) {
 	// Create a namespace
 	namespace := test.NewTestNamespace()
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory")},
+				Flavors: []v1beta1.FlavorQuotas{
+					{
+						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []v1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("8"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("18Gi"),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
+	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
+	localQueue := CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name, AsDefaultQueue)
+
 	// Create training PyTorch job
-	tuningJob := createFailedPyTorchJob(test, namespace.Name, image)
+	tuningJob := createFailedPyTorchJob(test, namespace.Name, image, localQueue)
 
 	// Make sure the PyTorch job is failed
 	test.Eventually(PyTorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutDouble).
 		Should(WithTransform(PyTorchJobConditionFailed, Equal(corev1.ConditionTrue)))
 }
 
-func createFailedPyTorchJob(test Test, namespace string, baseImage string) *kftov1.PyTorchJob {
+func createFailedPyTorchJob(test Test, namespace string, baseImage string, localQueue *v1beta1.LocalQueue) *kftov1.PyTorchJob {
 	tuningJob := &kftov1.PyTorchJob{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: corev1.SchemeGroupVersion.String(),
@@ -46,6 +78,9 @@ func createFailedPyTorchJob(test Test, namespace string, baseImage string) *kfto
 		},
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: "kfto-sft-",
+			Labels: map[string]string{
+				"kueue.x-k8s.io/queue-name": localQueue.Name,
+			},
 		},
 		Spec: kftov1.PyTorchJobSpec{
 			PyTorchReplicaSpecs: map[kftov1.ReplicaType]*kftov1.ReplicaSpec{
diff --git a/tests/kfto/kfto_training_test.go b/tests/kfto/kfto_training_test.go
index 2484d62c8..e226b6c63 100644
--- a/tests/kfto/kfto_training_test.go
+++ b/tests/kfto/kfto_training_test.go
@@ -27,6 +27,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/kueue/apis/kueue/v1beta1"
 
 	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 	. "github.com/opendatahub-io/distributed-workloads/tests/common/support"
@@ -78,6 +79,53 @@ func runKFTOPyTorchJob(t *testing.T, image string, gpu Accelerator, numGpus, num
 	// Create a namespace
 	namespace := test.CreateOrGetTestNamespace().Name
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	fmt.Sprintln(gpu.ResourceLabel)
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory")},
+				Flavors: []v1beta1.FlavorQuotas{
+					{
+						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []v1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("8"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("32Gi"),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	if gpu.IsGpu() {
+		numberOfGpus := (numberOfWorkerNodes + 1) * numGpus
+		cqSpec.ResourceGroups[0].CoveredResources = append(
+			cqSpec.ResourceGroups[0].CoveredResources,
+			corev1.ResourceName(gpu.ResourceLabel),
+		)
+		cqSpec.ResourceGroups[0].Flavors[0].Resources = append(
+			cqSpec.ResourceGroups[0].Flavors[0].Resources,
+			v1beta1.ResourceQuota{
+				Name:         corev1.ResourceName(gpu.ResourceLabel),
+				NominalQuota: resource.MustParse(fmt.Sprint(numberOfGpus)),
+			},
+		)
+	}
+
+	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
+	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
+	localQueue := CreateKueueLocalQueue(test, namespace, clusterQueue.Name, AsDefaultQueue)
+
 	// Create a ConfigMap with training script
 	configData := map[string][]byte{
 		"hf_llm_training.py": readFile(test, "resources/hf_llm_training.py"),
@@ -89,7 +137,7 @@ func runKFTOPyTorchJob(t *testing.T, image string, gpu Accelerator, numGpus, num
 	defer test.Client().Core().CoreV1().PersistentVolumeClaims(namespace).Delete(test.Ctx(), outputPvc.Name, metav1.DeleteOptions{})
 
 	// Create training PyTorch job
-	tuningJob := createKFTOPyTorchJob(test, namespace, *config, gpu, numGpus, numberOfWorkerNodes, outputPvc.Name, image)
+	tuningJob := createKFTOPyTorchJob(test, namespace, *config, gpu, numGpus, numberOfWorkerNodes, outputPvc.Name, image, localQueue)
 	defer test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Delete(test.Ctx(), tuningJob.Name, *metav1.NewDeleteOptions(0))
 
 	// Make sure the PyTorch job is running
@@ -122,7 +170,7 @@ func runKFTOPyTorchJob(t *testing.T, image string, gpu Accelerator, numGpus, num
 	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
 }
 
-func createKFTOPyTorchJob(test Test, namespace string, config corev1.ConfigMap, gpu Accelerator, numGpus, numberOfWorkerNodes int, outputPvcName string, baseImage string) *kftov1.PyTorchJob {
+func createKFTOPyTorchJob(test Test, namespace string, config corev1.ConfigMap, gpu Accelerator, numGpus, numberOfWorkerNodes int, outputPvcName string, baseImage string, localQueue *v1beta1.LocalQueue) *kftov1.PyTorchJob {
 	tuningJob := &kftov1.PyTorchJob{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: corev1.SchemeGroupVersion.String(),
@@ -130,6 +178,9 @@ func createKFTOPyTorchJob(test Test, namespace string, config corev1.ConfigMap,
 		},
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: "kfto-llm-",
+			Labels: map[string]string{
+				"kueue.x-k8s.io/default-queue": localQueue.Name,
+			},
 		},
 		Spec: kftov1.PyTorchJobSpec{
 			PyTorchReplicaSpecs: map[kftov1.ReplicaType]*kftov1.ReplicaSpec{
diff --git a/tests/kfto/kfto_upgrade_sleep_test.go b/tests/kfto/kfto_upgrade_sleep_test.go
index bcf49800e..6f9b73142 100644
--- a/tests/kfto/kfto_upgrade_sleep_test.go
+++ b/tests/kfto/kfto_upgrade_sleep_test.go
@@ -24,7 +24,9 @@ import (
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/kueue/apis/kueue/v1beta1"
 
 	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 	. "github.com/opendatahub-io/distributed-workloads/tests/common/support"
@@ -42,8 +44,39 @@ func TestSetupSleepPytorchjob(t *testing.T) {
 	// Create a namespace
 	createOrGetUpgradeTestNamespace(test, sleepNamespaceName)
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory")},
+				Flavors: []v1beta1.FlavorQuotas{
+					{
+						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []v1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("8"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("18Gi"),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
+	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
+	localQueue := CreateKueueLocalQueue(test, sleepNamespaceName, clusterQueue.Name, AsDefaultQueue)
+
 	// Create training PyTorch job
-	createSleepPyTorchJob(test, sleepNamespaceName)
+	createSleepPyTorchJob(test, sleepNamespaceName, localQueue)
 
 	// Make sure the PyTorch job is running, waiting for Training operator upgrade
 	test.Eventually(PyTorchJob(test, sleepNamespaceName, sleepPyTorchJobName), TestTimeoutShort).
@@ -72,7 +105,7 @@ func TestVerifySleepPytorchjob(t *testing.T) {
 		)
 }
 
-func createSleepPyTorchJob(test Test, namespace string) *kftov1.PyTorchJob {
+func createSleepPyTorchJob(test Test, namespace string, localQueue *v1beta1.LocalQueue) *kftov1.PyTorchJob {
 	// Does PyTorchJob already exist?
 	_, err := test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Get(test.Ctx(), sleepPyTorchJobName, metav1.GetOptions{})
 	if err == nil {
@@ -87,6 +120,9 @@ func createSleepPyTorchJob(test Test, namespace string) *kftov1.PyTorchJob {
 	tuningJob := &kftov1.PyTorchJob{
 		ObjectMeta: metav1.ObjectMeta{
 			Name: sleepPyTorchJobName,
+			Labels: map[string]string{
+				"kueue.x-k8s.io/queue-name": localQueue.Name,
+			},
 		},
 		Spec: kftov1.PyTorchJobSpec{
 			PyTorchReplicaSpecs: map[kftov1.ReplicaType]*kftov1.ReplicaSpec{
diff --git a/tests/kfto/resources/mnist_kfto.ipynb b/tests/kfto/resources/mnist_kfto.ipynb
index 903bbff32..aea0bc632 100644
--- a/tests/kfto/resources/mnist_kfto.ipynb
+++ b/tests/kfto/resources/mnist_kfto.ipynb
@@ -2,7 +2,17 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
+   "id": "61061396",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip show kubeflow-training  # to check installed kubeflow-traning SDK version during test-run"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
    "id": "b55bc3ea-4ce3-49bf-bb1f-e209de8ca47a",
    "metadata": {
     "tags": []
@@ -19,7 +29,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 41,
    "id": "72dd1751",
    "metadata": {},
    "outputs": [],
@@ -34,7 +44,53 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
+   "id": "aadaa4a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def GetDefaultLocalQueue(namespace: str):\n",
+    "    \"\"\"\n",
+    "    Fetches the LocalQueue in the given namespace whose annotation\n",
+    "    \"kueue.x-k8s.io/default-queue\" == \"true\". Returns the dict\n",
+    "    for the queue, or None.\n",
+    "    \"\"\"\n",
+    "    group   = \"kueue.x-k8s.io\"\n",
+    "    version = \"v1beta1\"\n",
+    "    plural  = \"localqueues\"\n",
+    "\n",
+    "    conf = c.Configuration()\n",
+    "    conf.host = openshift_api_url\n",
+    "    conf.verify_ssl = False\n",
+    "    conf.api_key = {\"authorization\": f\"Bearer {token}\"}\n",
+    "\n",
+    "    api_client = c.ApiClient(configuration=conf)\n",
+    "    api = c.CustomObjectsApi(api_client)\n",
+    "\n",
+    "    resp = api.list_namespaced_custom_object(\n",
+    "        group=group, version=version, namespace=namespace, plural=plural\n",
+    "    )\n",
+    "\n",
+    "    default_q = None\n",
+    "    for item in resp.get(\"items\", []):\n",
+    "        ann = item.get(\"metadata\", {}).get(\"annotations\") or {}\n",
+    "        if ann.get(\"kueue.x-k8s.io/default-queue\") == \"true\":\n",
+    "            if default_q is not None:\n",
+    "                raise RuntimeError(\n",
+    "                    f\"multiple LocalQueues annotated as default in {namespace}: \"\n",
+    "                    f\"{default_q['metadata']['name']} and {item['metadata']['name']}\"\n",
+    "                )\n",
+    "            default_q = item['metadata']['name']\n",
+    "\n",
+    "    if default_q is None:\n",
+    "        raise RuntimeError(f\"no LocalQueue annotated as default in namespace {namespace}\")\n",
+    "\n",
+    "    return default_q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "4ca70b20",
    "metadata": {},
    "outputs": [],
@@ -42,7 +98,10 @@
     "api_key = {\"authorization\": f\"Bearer {token}\"}\n",
     "# config = c.Configuration(host=openshift_api_url, api_key=token)\n",
     "# config.verify_ssl = False\n",
-    "tc = TrainingClient()"
+    "tc = TrainingClient()\n",
+    "\n",
+    "# get default local queue\n",
+    "default_local_queue=GetDefaultLocalQueue(namespace)"
    ]
   },
   {
@@ -67,6 +126,9 @@
     "      \"TORCH_DISTRIBUTED_DEBUG\": \"DETAIL\", \n",
     "      \"DEFAULT_PIP_INDEX_URL\": os.environ.get(\"PIP_INDEX_URL\"),\n",
     "      \"PIP_TRUSTED_HOST\": os.environ.get(\"PIP_TRUSTED_HOST\")\n",
+    "   },\n",
+    "   labels={\n",
+    "       \"kueue.x-k8s.io/queue-name\": default_local_queue,\n",
     "   }\n",
     ")"
    ]
@@ -96,7 +158,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "kfto2",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -110,7 +172,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.10"
+   "version": "3.11.0"
   }
  },
  "nbformat": 4,


Reviews:
user: login: sutaakar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T11:46:51Z
user: login: sutaakar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T11:47:27Z
user: login: sutaakar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T11:59:17Z
user: login: sutaakar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T12:09:14Z
user: login: sutaakar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T12:20:54Z
user: login: sutaakar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T12:28:40Z
user: login: abhijeet-dhumal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-20T12:32:11Z
user: login: abhijeet-dhumal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-20T12:36:49Z
user: login: abhijeet-dhumal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-20T12:38:04Z
user: login: sutaakar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T13:23:06Z
user: login: ChughShilpa
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T05:42:51Z

Review comments:
diff_hunk: @@ -60,13 +119,16 @@
     "   num_workers=1,\n",
     "   resources_per_worker={\"gpu\": num_gpus},\n",
     "   base_image=training_image,\n",
-    "   packages_to_install=[\"torchvision==0.19.0\",\"minio==7.2.13\"],\n",
+    "   packages_to_install=[\"--no-cache-dir\", \"torchvision==0.19.0\",\"minio==7.2.13\"],\n",
path: tests/kfto/resources/mnist_kfto.ipynb
user: login: sutaakar
body: Why is `--no-cache-dir` needed?
created_at: 2025-05-20T11:46:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,15 +44,64 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
+   "id": "aadaa4a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def GetDefaultLocalQueue(namespace: str):\n",
+    "    \"\"\"\n",
+    "    Fetches the LocalQueue in the given namespace whose annotation\n",
+    "    \"kueue.x-k8s.io/default-queue\" == \"true\". Returns the dict\n",
+    "    for the queue, or None.\n",
+    "    \"\"\"\n",
+    "    group   = \"kueue.x-k8s.io\"\n",
+    "    version = \"v1beta1\"\n",
+    "    plural  = \"localqueues\"\n",
+    "\n",
+    "    conf = c.Configuration()\n",
+    "    conf.host = openshift_api_url\n",
+    "    conf.verify_ssl = False\n",
+    "    conf.api_key = {\"authorization\": f\"Bearer {token}\"}\n",
+    "\n",
+    "    api_client = c.ApiClient(configuration=conf)\n",
+    "    api = c.CustomObjectsApi(api_client)\n",
+    "\n",
+    "    resp = api.list_namespaced_custom_object(\n",
+    "        group=group, version=version, namespace=namespace, plural=plural\n",
+    "    )\n",
+    "\n",
+    "    default_q = None\n",
+    "    for item in resp.get(\"items\", []):\n",
+    "        ann = item.get(\"metadata\", {}).get(\"annotations\") or {}\n",
+    "        if ann.get(\"kueue.x-k8s.io/default-queue\") == \"true\":\n",
+    "            if default_q is not None:\n",
+    "                raise RuntimeError(\n",
+    "                    f\"multiple LocalQueues annotated as default in {namespace}: \"\n",
+    "                    f\"{default_q['metadata']['name']} and {item['metadata']['name']}\"\n",
+    "                )\n",
+    "            default_q = item['metadata']['name']\n",
+    "\n",
+    "    if default_q is None:\n",
+    "        raise RuntimeError(f\"no LocalQueue annotated as default in namespace {namespace}\")\n",
+    "\n",
+    "    return default_q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "4ca70b20",
    "metadata": {},
    "outputs": [],
    "source": [
     "api_key = {\"authorization\": f\"Bearer {token}\"}\n",
     "# config = c.Configuration(host=openshift_api_url, api_key=token)\n",
     "# config.verify_ssl = False\n",
-    "tc = TrainingClient()"
+    "tc = TrainingClient()\n",
+    "\n",
+    "# get default local queue\n",
+    "default_local_queue=GetDefaultLocalQueue(namespace)"
path: tests/kfto/resources/mnist_kfto.ipynb
user: login: sutaakar
body: Wouldn't it be easier to pass local queue as parameter for the notebook?
created_at: 2025-05-20T11:47:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -80,7 +120,7 @@ func TestMnistSDK(t *testing.T) {
 	}()
 
 	// Make sure pytorch job is created
-	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp")).
+	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp"), TestTimeoutDouble).
path: tests/kfto/kfto_mnist_sdk_test.go
user: login: sutaakar
body: Is `TestTimeoutDouble` needed here?
created_at: 2025-05-20T11:59:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -42,6 +47,41 @@ func TestMnistSDK(t *testing.T) {
 	// Create role binding with Namespace specific admin cluster role
 	CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin")
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory"), corev1.ResourceName(NVIDIA.ResourceLabel)},
+				Flavors: []v1beta1.FlavorQuotas{
+					{
+						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []v1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("8"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("18Gi"),
+							},
path: tests/kfto/kfto_mnist_sdk_test.go
user: login: sutaakar
body: That looks like quite a lot of resources, is there a specific reason for such values?
created_at: 2025-05-20T12:09:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -42,6 +47,41 @@ func TestMnistSDK(t *testing.T) {
 	// Create role binding with Namespace specific admin cluster role
 	CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin")
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory"), corev1.ResourceName(NVIDIA.ResourceLabel)},
path: tests/kfto/kfto_mnist_sdk_test.go
user: login: sutaakar
body: Nitpick
```suggestion
				CoveredResources: []corev1.ResourceName{corev1.ResourceName(corev1.ResourceCPU), corev1.ResourceName(corev1.ResourceMemory), corev1.ResourceName(NVIDIA.ResourceLabel)},
```
created_at: 2025-05-20T12:20:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -122,7 +179,7 @@ func runKFTOPyTorchJob(t *testing.T, image string, gpu Accelerator, numGpus, num
 	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
 }
 
-func createKFTOPyTorchJob(test Test, namespace string, config corev1.ConfigMap, gpu Accelerator, numGpus, numberOfWorkerNodes int, outputPvcName string, baseImage string) *kftov1.PyTorchJob {
+func createKFTOPyTorchJob(test Test, namespace string, config corev1.ConfigMap, gpu Accelerator, numGpus, numberOfWorkerNodes int, outputPvcName string, baseImage string, localQueue *v1beta1.LocalQueue) *kftov1.PyTorchJob {
path: tests/kfto/kfto_training_test.go
user: login: sutaakar
body: Missing Kueue label in PyTorchjob.
created_at: 2025-05-20T12:28:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -80,7 +120,7 @@ func TestMnistSDK(t *testing.T) {
 	}()
 
 	// Make sure pytorch job is created
-	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp")).
+	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp"), TestTimeoutDouble).
path: tests/kfto/kfto_mnist_sdk_test.go
user: login: abhijeet-dhumal
body: Yes, i noticed the training image pull was taking more time than expected.. the same timeout is being used in other tests
created_at: 2025-05-20T12:32:11Z
author_association: MEMBER
diff_hunk: @@ -34,15 +44,64 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
+   "id": "aadaa4a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def GetDefaultLocalQueue(namespace: str):\n",
+    "    \"\"\"\n",
+    "    Fetches the LocalQueue in the given namespace whose annotation\n",
+    "    \"kueue.x-k8s.io/default-queue\" == \"true\". Returns the dict\n",
+    "    for the queue, or None.\n",
+    "    \"\"\"\n",
+    "    group   = \"kueue.x-k8s.io\"\n",
+    "    version = \"v1beta1\"\n",
+    "    plural  = \"localqueues\"\n",
+    "\n",
+    "    conf = c.Configuration()\n",
+    "    conf.host = openshift_api_url\n",
+    "    conf.verify_ssl = False\n",
+    "    conf.api_key = {\"authorization\": f\"Bearer {token}\"}\n",
+    "\n",
+    "    api_client = c.ApiClient(configuration=conf)\n",
+    "    api = c.CustomObjectsApi(api_client)\n",
+    "\n",
+    "    resp = api.list_namespaced_custom_object(\n",
+    "        group=group, version=version, namespace=namespace, plural=plural\n",
+    "    )\n",
+    "\n",
+    "    default_q = None\n",
+    "    for item in resp.get(\"items\", []):\n",
+    "        ann = item.get(\"metadata\", {}).get(\"annotations\") or {}\n",
+    "        if ann.get(\"kueue.x-k8s.io/default-queue\") == \"true\":\n",
+    "            if default_q is not None:\n",
+    "                raise RuntimeError(\n",
+    "                    f\"multiple LocalQueues annotated as default in {namespace}: \"\n",
+    "                    f\"{default_q['metadata']['name']} and {item['metadata']['name']}\"\n",
+    "                )\n",
+    "            default_q = item['metadata']['name']\n",
+    "\n",
+    "    if default_q is None:\n",
+    "        raise RuntimeError(f\"no LocalQueue annotated as default in namespace {namespace}\")\n",
+    "\n",
+    "    return default_q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "4ca70b20",
    "metadata": {},
    "outputs": [],
    "source": [
     "api_key = {\"authorization\": f\"Bearer {token}\"}\n",
     "# config = c.Configuration(host=openshift_api_url, api_key=token)\n",
     "# config.verify_ssl = False\n",
-    "tc = TrainingClient()"
+    "tc = TrainingClient()\n",
+    "\n",
+    "# get default local queue\n",
+    "default_local_queue=GetDefaultLocalQueue(namespace)"
path: tests/kfto/resources/mnist_kfto.ipynb
user: login: abhijeet-dhumal
body: The create notebook method is being used in ODH tests too, to pass localQueue as a parameter to notebook that would need to update createNotebook method too..
I think we can use it as a temporary solution, as I'm planning to add provision in upstream to detect default local queue from namespace similar to which is done in codeflare-sdk : https://github.com/project-codeflare/codeflare-sdk/blob/028d0d5511246a5ddd586b5d1cc2950288c3ca20/src/codeflare_sdk/ray/cluster/build_ray_cluster.py#L487
created_at: 2025-05-20T12:36:49Z
author_association: MEMBER
diff_hunk: @@ -42,6 +47,41 @@ func TestMnistSDK(t *testing.T) {
 	// Create role binding with Namespace specific admin cluster role
 	CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin")
 
+	// Create Kueue resources
+	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := v1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []v1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory"), corev1.ResourceName(NVIDIA.ResourceLabel)},
+				Flavors: []v1beta1.FlavorQuotas{
+					{
+						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []v1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("8"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("18Gi"),
+							},
path: tests/kfto/kfto_mnist_sdk_test.go
user: login: abhijeet-dhumal
body: Ohh right, for this test it doesn't need these much resources 👀 , Apologies, updating it now!
created_at: 2025-05-20T12:38:04Z
author_association: MEMBER
diff_hunk: @@ -80,7 +120,7 @@ func TestMnistSDK(t *testing.T) {
 	}()
 
 	// Make sure pytorch job is created
-	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp")).
+	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp"), TestTimeoutDouble).
path: tests/kfto/kfto_mnist_sdk_test.go
user: login: sutaakar
body: ok
created_at: 2025-05-20T13:23:06Z
author_association: CONTRIBUTOR

PR comments:
user: login: sutaakar
created_at: 2025-05-20T13:59:22Z
body: /lgtm
author_association: CONTRIBUTOR
user: login: openshift-ci[bot]
created_at: 2025-05-21T05:42:56Z
body: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/opendatahub-io/distributed-workloads/pull/399#pullrequestreview-2856360187" title="Approved">ChughShilpa</a>*
**Once this PR has been reviewed and has the lgtm label**, please ask for approval from [sutaakar](https://github.com/sutaakar). For more information see [the Code Review Process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process).

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=opendatahub-io%2Fdistributed-workloads).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/opendatahub-io/distributed-workloads/blob/main/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["sutaakar"]} -->
author_association: NONE