Implement PodSets inference for RayClusters

Author: tardieu

internal/webhook/appwrapper_fixtures_test.go

@@ -396,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

internal/webhook/appwrapper_webhook_test.go

@@ -191,13 +191,23 @@ var _ = Describe("AppWrapper Webhook Tests", func() {
 			Expect(k8sClient.Delete(ctx, aw)).To(Succeed())
 		})
 
-		It("PodSets are inferred for known GVKs", func() {
-			aw := toAppWrapper(pod(100), deploymentForInference(1, 100), podForInference(100),
-				jobForInference(2, 4, 100), jobForInference(8, 4, 100), pytorchJobForInference(100, 4, 100))
+		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())
+			})
 
-			Expect(k8sClient.Create(ctx, aw)).To(Succeed(), "PodSets for deployments and pods should be inferred")
-			Expect(aw.Spec.Suspend).Should(BeTrue())
-			Expect(k8sClient.Delete(ctx, aw)).To(Succeed())
+			It("PodSets are inferred for PyTorchJobs and RayClusters", func() {
+				aw := toAppWrapper(pytorchJobForInference(100, 4, 100), rayClusterForInference(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())
+			})
 		})
 	})
 

pkg/utils/utils.go

@@ -237,13 +237,34 @@ func InferPodSets(obj *unstructured.Unstructured) ([]workloadv1beta2.AppWrapperP
 		for _, replicaType := range []string{"Master", "Worker"} {
 			prefix := "template.spec.pytorchReplicaSpecs." + replicaType + "."
 			// validate path to replica template
-			if _, err := getValueAtPath(obj.UnstructuredContent(), prefix+"template"); err == nil {
+			if _, err := getValueAtPath(obj.UnstructuredContent(), prefix+templateString); err == nil {
 				// infer replica count
 				replicas, err := InferReplicas(obj.UnstructuredContent(), prefix+"replicas")
 				if err != nil {
 					return nil, err
 				}
-				podSets = append(podSets, workloadv1beta2.AppWrapperPodSet{Replicas: ptr.To(replicas), Path: prefix + "template"})
+				podSets = append(podSets, workloadv1beta2.AppWrapperPodSet{Replicas: ptr.To(replicas), Path: prefix + templateString})
+			}
+		}
+
+	case schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayCluster"}:
+		if _, err := getValueAtPath(obj.UnstructuredContent(), "template.spec.headGroupSpec.template"); err == nil {
+			podSets = append(podSets, workloadv1beta2.AppWrapperPodSet{Replicas: ptr.To(int32(1)), Path: "template.spec.headGroupSpec.template"})
+		}
+		if workers, err := getValueAtPath(obj.UnstructuredContent(), "template.spec.workerGroupSpecs"); err == nil {
+			if workers, ok := workers.([]interface{}); ok {
+				for i := range workers {
+					prefix := fmt.Sprintf("template.spec.workerGroupSpecs[%v].", i)
+					// validate path to replica template
+					if _, err := getValueAtPath(obj.UnstructuredContent(), prefix+templateString); err == nil {
+						// infer replica count
+						replicas, err := InferReplicas(obj.UnstructuredContent(), prefix+"replicas")
+						if err != nil {
+							return nil, err
+						}
+						podSets = append(podSets, workloadv1beta2.AppWrapperPodSet{Replicas: ptr.To(replicas), Path: prefix + templateString})
+					}
+				}
 			}
 		}