Infer pod sets for known GVKs (#108)

tardieu · web-flow · commit b87fa228e2e8 · 2024-04-25T08:06:50.000-04:00
diff --git a/internal/webhook/appwrapper_fixtures_test.go b/internal/webhook/appwrapper_fixtures_test.go
@@ -86,6 +86,18 @@ func pod(milliCPU int64) workloadv1beta2.AppWrapperComponent {
 	}
 }
 
+func podForInference(milliCPU int64) workloadv1beta2.AppWrapperComponent {
+	yamlString := fmt.Sprintf(podYAML,
+		randName("pod"),
+		resource.NewMilliQuantity(milliCPU, resource.DecimalSI))
+
+	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
+	Expect(err).NotTo(HaveOccurred())
+	return workloadv1beta2.AppWrapperComponent{
+		Template: runtime.RawExtension{Raw: jsonBytes},
+	}
+}
+
 const namespacedPodYAML = `
 apiVersion: v1
 kind: Pod
@@ -179,6 +191,19 @@ func deployment(replicaCount int, milliCPU int64) workloadv1beta2.AppWrapperComp
 	}
 }
 
+func deploymentForInference(replicaCount int, milliCPU int64) workloadv1beta2.AppWrapperComponent {
+	yamlString := fmt.Sprintf(deploymentYAML,
+		randName("deployment"),
+		replicaCount,
+		resource.NewMilliQuantity(milliCPU, resource.DecimalSI))
+
+	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
+	Expect(err).NotTo(HaveOccurred())
+	return workloadv1beta2.AppWrapperComponent{
+		Template: runtime.RawExtension{Raw: jsonBytes},
+	}
+}
+
 const rayClusterYAML = `
 apiVersion: ray.io/v1
 kind: RayCluster
@@ -371,6 +396,20 @@ func rayCluster(workerCount int, milliCPU int64) workloadv1beta2.AppWrapperCompo
 	}
 }
 
+func rayClusterForInference(workerCount int, milliCPU int64) workloadv1beta2.AppWrapperComponent {
+	workerCPU := resource.NewMilliQuantity(milliCPU, resource.DecimalSI)
+	yamlString := fmt.Sprintf(rayClusterYAML,
+		randName("raycluster"),
+		workerCount, workerCount, workerCount,
+		workerCPU)
+
+	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
+	Expect(err).NotTo(HaveOccurred())
+	return workloadv1beta2.AppWrapperComponent{
+		Template: runtime.RawExtension{Raw: jsonBytes},
+	}
+}
+
 const jobSetYAML = `
 apiVersion: jobset.x-k8s.io/v1alpha2
 kind: JobSet
@@ -426,3 +465,128 @@ func jobSet(replicasWorker int, milliCPUWorker int64) workloadv1beta2.AppWrapper
 		Template: runtime.RawExtension{Raw: jsonBytes},
 	}
 }
+
+const jobYAML = `
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: %v
+spec:
+  parallelism: %v
+  completions: %v
+  template:
+    spec:
+      restartPolicy: Never
+      containers:
+      - name: busybox
+        image: quay.io/project-codeflare/busybox:1.36
+        command: ["sh", "-c", "sleep 30"]
+        resources:
+          requests:
+            cpu: %v`
+
+func jobForInference(parallelism int, completions int, milliCPU int64) workloadv1beta2.AppWrapperComponent {
+	yamlString := fmt.Sprintf(jobYAML,
+		randName("job"),
+		parallelism,
+		completions,
+		resource.NewMilliQuantity(milliCPU, resource.DecimalSI))
+
+	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
+	Expect(err).NotTo(HaveOccurred())
+	return workloadv1beta2.AppWrapperComponent{
+		Template: runtime.RawExtension{Raw: jsonBytes},
+	}
+}
+
+const pytorchJobYAML = `
+apiVersion: "kubeflow.org/v1"
+kind: PyTorchJob
+metadata:
+  name: %v
+spec:
+  pytorchReplicaSpecs:
+    Master:
+      restartPolicy: OnFailure
+      template:
+        spec:
+          containers:
+          - name: pytorch
+            image: docker.io/kubeflowkatib/pytorch-mnist-cpu:v1beta1-fc858d1
+            command:
+            - "python3"
+            - "/opt/pytorch-mnist/mnist.py"
+            - "--epochs=1"
+            resources:
+              requests:
+                cpu: %v
+    Worker:
+      replicas: %v
+      restartPolicy: OnFailure
+      template:
+        spec:
+          containers:
+          - name: pytorch
+            image: docker.io/kubeflowkatib/pytorch-mnist-cpu:v1beta1-fc858d1
+            command:
+            - "python3"
+            - "/opt/pytorch-mnist/mnist.py"
+            - "--epochs=1"
+            resources:
+              requests:
+                cpu: %v`
+
+func pytorchJobForInference(masterMilliCPU int64, workerReplicas int, workerMilliCPU int64) workloadv1beta2.AppWrapperComponent {
+	yamlString := fmt.Sprintf(pytorchJobYAML,
+		randName("pytorch-job"),
+		resource.NewMilliQuantity(masterMilliCPU, resource.DecimalSI),
+		workerReplicas,
+		resource.NewMilliQuantity(workerMilliCPU, resource.DecimalSI))
+
+	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
+	Expect(err).NotTo(HaveOccurred())
+	return workloadv1beta2.AppWrapperComponent{
+		Template: runtime.RawExtension{Raw: jsonBytes},
+	}
+}
+
+const rayJobYAML = `
+apiVersion: ray.io/v1
+kind: RayJob
+metadata:
+  name: %v
+spec:
+  rayClusterSpec:
+    headGroupSpec:
+      template:
+        spec:
+          containers:
+            - name: ray-head
+              image: rayproject/ray:2.9.0
+              resources:
+                requests:
+                  cpu: 1
+    workerGroupSpecs:
+      - replicas: %v
+        template:
+          spec:
+            containers:
+              - name: ray-worker
+                image: rayproject/ray:2.9.0
+                resources:
+                  requests:
+                    cpu: %v
+`
+
+func rayJobForInference(workerCount int, milliCPU int64) workloadv1beta2.AppWrapperComponent {
+	yamlString := fmt.Sprintf(rayJobYAML,
+		randName("rayjob"),
+		workerCount,
+		resource.NewMilliQuantity(milliCPU, resource.DecimalSI))
+
+	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
+	Expect(err).NotTo(HaveOccurred())
+	return workloadv1beta2.AppWrapperComponent{
+		Template: runtime.RawExtension{Raw: jsonBytes},
+	}
+}
diff --git a/internal/webhook/appwrapper_webhook.go b/internal/webhook/appwrapper_webhook.go
@@ -62,6 +62,10 @@ func (w *AppWrapperWebhook) Default(ctx context.Context, obj runtime.Object) err
 	if w.Config.EnableKueueIntegrations {
 		jobframework.ApplyDefaultForSuspend((*wlc.AppWrapper)(aw), w.Config.ManageJobsWithoutQueueName)
 	}
+	if err := inferPodSets(ctx, aw); err != nil {
+		log.FromContext(ctx).Info("Error raised during podSet inference", "job", aw)
+		return err
+	}
 	return nil
 }
 
@@ -98,6 +102,30 @@ func (w *AppWrapperWebhook) ValidateDelete(context.Context, runtime.Object) (adm
 	return nil, nil
 }
 
+// inferPodSets infers the AppWrapper's PodSets
+func inferPodSets(_ context.Context, aw *workloadv1beta2.AppWrapper) error {
+	components := aw.Spec.Components
+	componentsPath := field.NewPath("spec").Child("components")
+	for idx, component := range components {
+		compPath := componentsPath.Index(idx)
+
+		// Automatically create elided PodSets for known GVKs
+		if len(component.PodSets) == 0 {
+			unstruct := &unstructured.Unstructured{}
+			_, _, err := unstructured.UnstructuredJSONScheme.Decode(component.Template.Raw, nil, unstruct)
+			if err != nil {
+				return field.Invalid(compPath.Child("template"), component.Template, "failed to decode as JSON")
+			}
+			podSets, err := utils.InferPodSets(unstruct)
+			if err != nil {
+				return err
+			}
+			components[idx].PodSets = podSets
+		}
+	}
+	return nil
+}
+
 // rbacs required to enable SubjectAccessReview
 //+kubebuilder:rbac:groups=authorization.k8s.io,resources=subjectaccessreviews,verbs=create
 //+kubebuilder:rbac:groups=apiextensions.k8s.io,resources=customresourcedefinitions,verbs=list
@@ -182,9 +210,15 @@ func (w *AppWrapperWebhook) validateAppWrapperCreate(ctx context.Context, aw *wo
 			}
 			podSpecCount += 1
 		}
+
+		// 5. Validate PodSets for known GVKs
+		if err := utils.ValidatePodSets(unstruct, component.PodSets); err != nil {
+			allErrors = append(allErrors, field.Invalid(podSetsPath, component.PodSets, err.Error()))
+		}
+
 	}
 
-	// 5. Enforce Kueue limitation that 0 < podSpecCount <= 8
+	// 6. Enforce Kueue limitation that 0 < podSpecCount <= 8
 	if podSpecCount == 0 {
 		allErrors = append(allErrors, field.Invalid(componentsPath, components, "components contains no podspecs"))
 	}
diff --git a/internal/webhook/appwrapper_webhook_test.go b/internal/webhook/appwrapper_webhook_test.go
@@ -190,6 +190,25 @@ var _ = Describe("AppWrapper Webhook Tests", func() {
 			Expect(aw.Spec.Suspend).Should(BeTrue())
 			Expect(k8sClient.Delete(ctx, aw)).To(Succeed())
 		})
+
+		Context("PodSets are inferred for known GVKs", func() {
+			It("PodSets are inferred for common kinds", func() {
+				aw := toAppWrapper(pod(100), deploymentForInference(1, 100), podForInference(100),
+					jobForInference(2, 4, 100), jobForInference(8, 4, 100))
+
+				Expect(k8sClient.Create(ctx, aw)).To(Succeed(), "PodSets should be inferred")
+				Expect(aw.Spec.Suspend).Should(BeTrue())
+				Expect(k8sClient.Delete(ctx, aw)).To(Succeed())
+			})
+
+			It("PodSets are inferred for PyTorchJobs, RayClusters, and RayJobs", func() {
+				aw := toAppWrapper(pytorchJobForInference(100, 4, 100), rayClusterForInference(7, 100), rayJobForInference(7, 100))
+
+				Expect(k8sClient.Create(ctx, aw)).To(Succeed(), "PodSets should be inferred")
+				Expect(aw.Spec.Suspend).Should(BeTrue())
+				Expect(k8sClient.Delete(ctx, aw)).To(Succeed())
+			})
+		})
 	})
 
 })
diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
