add resource command (ray-project#170)

* auto fill resource information into command Co-authored-by: zhuang <[email protected]>
lowang-bh · Mar 10, 2022 · a6c6395 · a6c6395
1 parent af5964a
commit a6c6395
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 4 deletions.
diff --git a/ray-operator/config/samples/ray-cluster.mini.yaml b/ray-operator/config/samples/ray-cluster.mini.yaml
@@ -53,6 +53,13 @@ spec:
             valueFrom:
               fieldRef:
                 fieldPath: status.podIP
+          resources:
+            limits:
+              cpu: 1
+              memory: 2Gi
+            requests:
+              cpu: 1
+              memory: 2Gi
           ports:
           - containerPort: 6379
             name: redis

diff --git a/ray-operator/controllers/common/pod.go b/ray-operator/controllers/common/pod.go
@@ -109,7 +109,7 @@ func BuildPod(podTemplateSpec v1.PodTemplateSpec, rayNodeType rayiov1alpha1.RayN
 		cmd += convertCmdToString(pod.Spec.Containers[index].Args)
 	}
 	if !strings.Contains(cmd, "ray start") {
-		cont := concatenateContainerCommand(rayNodeType, rayStartParams)
+		cont := concatenateContainerCommand(rayNodeType, rayStartParams, pod.Spec.Containers[index].Resources)
 		// replacing the old command
 		pod.Spec.Containers[index].Command = []string{"/bin/bash", "-c", "--"}
 		if cmd != "" {
@@ -334,7 +334,31 @@ func setMissingRayStartParams(rayStartParams map[string]string, nodeType rayiov1
 }
 
 // concatenateContainerCommand with ray start
-func concatenateContainerCommand(nodeType rayiov1alpha1.RayNodeType, rayStartParams map[string]string) (fullCmd string) {
+func concatenateContainerCommand(nodeType rayiov1alpha1.RayNodeType, rayStartParams map[string]string, resource v1.ResourceRequirements) (fullCmd string) {
+
+	if _, ok := rayStartParams["num-cpus"]; !ok {
+		cpu := resource.Limits[v1.ResourceCPU]
+		if !cpu.IsZero() {
+			rayStartParams["num-cpus"] = strconv.FormatInt(cpu.Value(), 10)
+		}
+	}
+
+	if _, ok := rayStartParams["memory"]; !ok {
+		memory := resource.Limits[v1.ResourceMemory]
+		if !memory.IsZero() {
+			rayStartParams["memory"] = strconv.FormatInt(memory.Value(), 10)
+		}
+	}
+
+	if _, ok := rayStartParams["num-gpus"]; !ok {
+		gpu := resource.Limits["gpu"]
+		if !gpu.IsZero() {
+			rayStartParams["num-gpus"] = strconv.FormatInt(gpu.Value(), 10)
+		}
+	}
+
+	log.V(10).Info("concatenate container command", "ray start params", rayStartParams)
+
 	switch nodeType {
 	case rayiov1alpha1.HeadNode:
 		return fmt.Sprintf("ulimit -n 65536; ray start --head %s", convertParamMap(rayStartParams))

diff --git a/ray-operator/controllers/common/pod_test.go b/ray-operator/controllers/common/pod_test.go
@@ -191,9 +191,9 @@ func TestBuildPod(t *testing.T) {
 		t.Fatalf("Expected `%v` but got `%v`", expectedResult, actualResult)
 	}
 
-	expectedCommandArg := splitAndSort("ulimit -n 65536; ray start --block --num-cpus=1  --address=raycluster-sample-head-svc:6379  --port=6379  --redis-password=LetMeInRay")
+	expectedCommandArg := splitAndSort("ulimit -n 65536; ray start --block --num-cpus=1 --address=raycluster-sample-head-svc:6379  --port=6379  --redis-password=LetMeInRay")
 	if !reflect.DeepEqual(expectedCommandArg, splitAndSort(pod.Spec.Containers[0].Args[0])) {
-		t.Fatalf("Expected `%v` but got `%v`", expectedCommandArg, pod.Spec.Containers[0].Args)
+		t.Fatalf("Expected `%v` but got `%v`", expectedCommandArg, pod.Spec.Containers[0].Args[0])
 	}
 }