fix(apiserver): env MEMORY_* use memory ResourceField (#2438)

Signed-off-by: Abirdcfly <[email protected]>
ray-project · Oct 21, 2024 · d73c12e · d73c12e
1 parent 4e5a916
commit d73c12e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/apiserver/pkg/util/cluster.go b/apiserver/pkg/util/cluster.go
@@ -484,7 +484,7 @@ func buildWorkerPodTemplate(imageVersion string, envs *api.EnvironmentVariables,
 							ValueFrom: &corev1.EnvVarSource{
 								ResourceFieldRef: &corev1.ResourceFieldSelector{
 									ContainerName: "ray-worker",
-									Resource:      "requests.cpu",
+									Resource:      "requests.memory",
 								},
 							},
 						},
@@ -493,7 +493,7 @@ func buildWorkerPodTemplate(imageVersion string, envs *api.EnvironmentVariables,
 							ValueFrom: &corev1.EnvVarSource{
 								ResourceFieldRef: &corev1.ResourceFieldSelector{
 									ContainerName: "ray-worker",
-									Resource:      "limits.cpu",
+									Resource:      "limits.memory",
 								},
 							},
 						},

diff --git a/apiserver/pkg/util/cluster_test.go b/apiserver/pkg/util/cluster_test.go
@@ -618,6 +618,12 @@ func TestBuilWorkerPodTemplate(t *testing.T) {
 	assert.Equal(t, expectedToleration, podSpec.Spec.Tolerations[0], "failed to propagate tolerations")
 	assert.Equal(t, "bar", podSpec.Annotations["foo"], "failed to convert annotations")
 	assert.Equal(t, expectedLabels, podSpec.Labels, "failed to convert labels")
+	assert.True(t, containsEnvValueFrom(podSpec.Spec.Containers[0].Env, "CPU_REQUEST", &corev1.EnvVarSource{ResourceFieldRef: &corev1.ResourceFieldSelector{ContainerName: "ray-worker", Resource: "requests.cpu"}}), "failed to propagate environment variable: CPU_REQUEST")
+	assert.True(t, containsEnvValueFrom(podSpec.Spec.Containers[0].Env, "CPU_LIMITS", &corev1.EnvVarSource{ResourceFieldRef: &corev1.ResourceFieldSelector{ContainerName: "ray-worker", Resource: "limits.cpu"}}), "failed to propagate environment variable: CPU_LIMITS")
+	assert.True(t, containsEnvValueFrom(podSpec.Spec.Containers[0].Env, "MEMORY_REQUESTS", &corev1.EnvVarSource{ResourceFieldRef: &corev1.ResourceFieldSelector{ContainerName: "ray-worker", Resource: "requests.memory"}}), "failed to propagate environment variable: MEMORY_REQUESTS")
+	assert.True(t, containsEnvValueFrom(podSpec.Spec.Containers[0].Env, "MEMORY_LIMITS", &corev1.EnvVarSource{ResourceFieldRef: &corev1.ResourceFieldSelector{ContainerName: "ray-worker", Resource: "limits.memory"}}), "failed to propagate environment variable: MEMORY_LIMITS")
+	assert.True(t, containsEnvValueFrom(podSpec.Spec.Containers[0].Env, "MY_POD_NAME", &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.name"}}), "failed to propagate environment variable: MY_POD_NAME")
+	assert.True(t, containsEnvValueFrom(podSpec.Spec.Containers[0].Env, "MY_POD_IP", &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "status.podIP"}}), "failed to propagate environment variable: MY_POD_IP")
 
 	// Check Resources
 	container := podSpec.Spec.Containers[0]
@@ -645,6 +651,15 @@ func containsEnv(envs []corev1.EnvVar, key string, val string) bool {
 	return false
 }
 
+func containsEnvValueFrom(envs []corev1.EnvVar, key string, valFrom *corev1.EnvVarSource) bool {
+	for _, env := range envs {
+		if env.Name == key && reflect.DeepEqual(env.ValueFrom, valFrom) {
+			return true
+		}
+	}
+	return false
+}
+
 func tolerationToString(toleration *corev1.Toleration) string {
 	return "Key: " + toleration.Key + " Operator: " + string(toleration.Operator) + " Effect: " + string(toleration.Effect)
 }