fix: correctly set GPUs in rayStartParams (#497)

This is a follow up fix to #328. Code in that PR breaks out of a loop that checks for GPU resources on the first iteration.
ray-project · Aug 19, 2022 · 92d1ffb · 92d1ffb
1 parent 40ea21c
commit 92d1ffb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/ray-operator/controllers/ray/common/pod.go b/ray-operator/controllers/ray/common/pod.go
@@ -670,9 +670,9 @@ func concatenateContainerCommand(nodeType rayiov1alpha1.RayNodeType, rayStartPar
 		for resourceKey, resource := range resource.Limits {
 			if strings.HasSuffix(string(resourceKey), "gpu") && !resource.IsZero() {
 				rayStartParams["num-gpus"] = strconv.FormatInt(resource.Value(), 10)
+				// For now, only support one GPU type. Break on first match.
+				break
 			}
-			// For now, only support one GPU type. Break on first match.
-			break
 		}
 	}
 

diff --git a/ray-operator/controllers/ray/common/pod_test.go b/ray-operator/controllers/ray/common/pod_test.go
@@ -110,7 +110,9 @@ var instance = rayiov1alpha1.RayCluster{
 								Image: "repo/image:custom",
 								Resources: v1.ResourceRequirements{
 									Limits: v1.ResourceList{
-										"nvidia.com/gpu": resource.MustParse("3"),
+										v1.ResourceCPU:    resource.MustParse("1"),
+										v1.ResourceMemory: testMemoryLimit,
+										"nvidia.com/gpu":  resource.MustParse("3"),
 									},
 								},
 								Env: []v1.EnvVar{
@@ -358,9 +360,10 @@ 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 --num-gpus=3 --address=raycluster-sample-head-svc:6379 --port=6379 --redis-password=LetMeInRay --metrics-export-port=8080")
-	if !reflect.DeepEqual(expectedCommandArg, splitAndSort(pod.Spec.Containers[0].Args[0])) {
-		t.Fatalf("Expected `%v` but got `%v`", expectedCommandArg, pod.Spec.Containers[0].Args[0])
+	expectedCommandArg := splitAndSort("ulimit -n 65536; ray start --block --memory=1073741824 --num-cpus=1 --num-gpus=3 --address=raycluster-sample-head-svc:6379 --port=6379 --redis-password=LetMeInRay --metrics-export-port=8080")
+	actualCommandArg := splitAndSort(pod.Spec.Containers[0].Args[0])
+	if !reflect.DeepEqual(expectedCommandArg, actualCommandArg) {
+		t.Fatalf("Expected `%v` but got `%v`", expectedCommandArg, actualCommandArg)
 	}
 
 	// Check Envs