Add test coverage for sequential admission of PytorchJob

Author: sutaakar

tests/kfto/kfto_kueue_sft_test.go

@@ -46,6 +46,10 @@ func PytorchJobConditionSucceeded(job *kftov1.PyTorchJob) corev1.ConditionStatus
 	return PytorchJobCondition(job, kftov1.JobSucceeded)
 }
 
+func PytorchJobConditionSuspended(job *kftov1.PyTorchJob) corev1.ConditionStatus {
+	return PytorchJobCondition(job, kftov1.JobSuspended)
+}
+
 func PytorchJobCondition(job *kftov1.PyTorchJob, conditionType kftov1.JobConditionType) corev1.ConditionStatus {
 	for _, condition := range job.Status.Conditions {
 		if condition.Type == conditionType {
@@ -55,9 +59,12 @@ func PytorchJobCondition(job *kftov1.PyTorchJob, conditionType kftov1.JobConditi
 	return corev1.ConditionUnknown
 }
 
+func OwnerReferenceName(meta metav1.Object) string {
+	return meta.GetOwnerReferences()[0].Name
+}
+
 func TestPytorchjobWithSFTtrainer(t *testing.T) {
 	test := With(t)
-	test.T().Parallel()
 
 	// Create a namespace
 	namespace := test.NewTestNamespace()
@@ -99,17 +106,107 @@ func TestPytorchjobWithSFTtrainer(t *testing.T) {
 	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
 	localQueue := CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name)
 
-	// Run training PyTorch job
+	// Create training PyTorch job
+	tuningJob := createPyTorchJob(test, namespace.Name, localQueue.Name, *config)
+
+	// Make sure the Kueue Workload is admitted
+	test.Eventually(KueueWorkloads(test, namespace.Name), TestTimeoutLong).
+		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), TestTimeoutShort).
+		Should(WithTransform(PytorchJobConditionRunning, Equal(corev1.ConditionTrue)))
+
+	// Make sure the PyTorch job succeed
+	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutShort).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
+	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
+}
+
+func TestPytorchjobUsingKueueQuota(t *testing.T) {
+	test := With(t)
+
+	// Create a namespace
+	namespace := test.NewTestNamespace()
+
+	// Create a ConfigMap with training dataset and configuration
+	configData := map[string][]byte{
+		"config.json":                   ReadFile(test, "config.json"),
+		"twitter_complaints_small.json": ReadFile(test, "twitter_complaints_small.json"),
+	}
+	config := CreateConfigMap(test, namespace.Name, configData)
+
+	// Create limited Kueue resources to run just one Pytorchjob at a time
+	resourceFlavor := CreateKueueResourceFlavor(test, kueuev1beta1.ResourceFlavorSpec{})
+	defer test.Client().Kueue().KueueV1beta1().ResourceFlavors().Delete(test.Ctx(), resourceFlavor.Name, metav1.DeleteOptions{})
+	cqSpec := kueuev1beta1.ClusterQueueSpec{
+		NamespaceSelector: &metav1.LabelSelector{},
+		ResourceGroups: []kueuev1beta1.ResourceGroup{
+			{
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory")},
+				Flavors: []kueuev1beta1.FlavorQuotas{
+					{
+						Name: kueuev1beta1.ResourceFlavorReference(resourceFlavor.Name),
+						Resources: []kueuev1beta1.ResourceQuota{
+							{
+								Name:         corev1.ResourceCPU,
+								NominalQuota: resource.MustParse("3"),
+							},
+							{
+								Name:         corev1.ResourceMemory,
+								NominalQuota: resource.MustParse("6Gi"),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
+	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
+	localQueue := CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name)
+
+	// Create first training PyTorch job
+	tuningJob := createPyTorchJob(test, namespace.Name, localQueue.Name, *config)
+
+	// Make sure the PyTorch job is running
+	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutShort).
+		Should(WithTransform(PytorchJobConditionRunning, Equal(corev1.ConditionTrue)))
+
+	// Create second training PyTorch job
+	secondTuningJob := createPyTorchJob(test, namespace.Name, localQueue.Name, *config)
+
+	// Make sure the second PyTorch job is suspended, waiting for first job to finish
+	test.Eventually(PytorchJob(test, namespace.Name, secondTuningJob.Name), TestTimeoutShort).
+		Should(WithTransform(PytorchJobConditionSuspended, Equal(corev1.ConditionTrue)))
+
+	// Make sure the first PyTorch job succeed
+	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutLong).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
+	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
+
+	// Second PyTorch job should be started now
+	test.Eventually(PytorchJob(test, namespace.Name, secondTuningJob.Name), TestTimeoutShort).
+		Should(WithTransform(PytorchJobConditionRunning, Equal(corev1.ConditionTrue)))
+
+	// Make sure the second PyTorch job succeed
+	test.Eventually(PytorchJob(test, namespace.Name, secondTuningJob.Name), TestTimeoutLong).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
+	test.T().Logf("PytorchJob %s/%s ran successfully", secondTuningJob.Namespace, secondTuningJob.Name)
+}
+
+func createPyTorchJob(test Test, namespace, localQueueName string, config corev1.ConfigMap) *kftov1.PyTorchJob {
 	tuningJob := &kftov1.PyTorchJob{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: corev1.SchemeGroupVersion.String(),
 			Kind:       "PyTorchJob",
 		},
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      "kfto-sft",
-			Namespace: namespace.Name,
+			GenerateName: "kfto-sft-",
 			Labels: map[string]string{
-				"kueue.x-k8s.io/queue-name": localQueue.Name,
+				"kueue.x-k8s.io/queue-name": localQueueName,
 			},
 		},
 		Spec: kftov1.PyTorchJobSpec{
@@ -174,24 +271,9 @@ func TestPytorchjobWithSFTtrainer(t *testing.T) {
 		},
 	}
 
-	tuningJob, err := test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace.Name).Create(test.Ctx(), tuningJob, metav1.CreateOptions{})
+	tuningJob, err := test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Create(test.Ctx(), tuningJob, metav1.CreateOptions{})
 	test.Expect(err).NotTo(HaveOccurred())
 	test.T().Logf("Created PytorchJob %s/%s successfully", tuningJob.Namespace, tuningJob.Name)
 
-	// Make sure the Kueue Workload is admitted
-	test.Eventually(KueueWorkloads(test, namespace.Name), TestTimeoutLong).
-		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), TestTimeoutShort).
-		Should(WithTransform(PytorchJobConditionRunning, Equal(corev1.ConditionTrue)))
-
-	// Make sure the PyTorch job succeed
-	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutLong).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
-	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
+	return tuningJob
 }