Inject cluster name as an environment variable into head and worker p…

…ods (ray-project#934) Inject cluster name as an environment variable into head and worker pods. With this PR, users can use $RAY_CLUSTER_NAME to get RayCluster's name in Pods.
lowang-bh · Mar 14, 2023 · e46c46a · e46c46a
1 parent 6ee405d
commit e46c46a
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 7 deletions.
diff --git a/ray-operator/controllers/ray/common/constant.go b/ray-operator/controllers/ray/common/constant.go
@@ -72,7 +72,7 @@ const (
 
 	// Use as container env variable
 	NAMESPACE                               = "NAMESPACE"
-	CLUSTER_NAME                            = "CLUSTER_NAME"
+	RAY_CLUSTER_NAME                        = "RAY_CLUSTER_NAME"
 	RAY_IP                                  = "RAY_IP"
 	FQ_RAY_IP                               = "FQ_RAY_IP"
 	RAY_PORT                                = "RAY_PORT"

diff --git a/ray-operator/controllers/ray/common/pod.go b/ray-operator/controllers/ray/common/pod.go
@@ -380,10 +380,10 @@ func BuildAutoscalerContainer(autoscalerImage string) v1.Container {
 		ImagePullPolicy: v1.PullAlways,
 		Env: []v1.EnvVar{
 			{
-				Name: "RAY_CLUSTER_NAME",
+				Name: RAY_CLUSTER_NAME,
 				ValueFrom: &v1.EnvVarSource{
 					FieldRef: &v1.ObjectFieldSelector{
-						FieldPath: "metadata.labels['ray.io/cluster']",
+						FieldPath: fmt.Sprintf("metadata.labels['%s']", RayClusterLabelKey),
 					},
 				},
 			},
@@ -549,6 +549,7 @@ func setInitContainerEnvVars(container *v1.Container, fqdnRayIP string) {
 }
 
 func setContainerEnvVars(pod *v1.Pod, rayContainerIndex int, rayNodeType rayiov1alpha1.RayNodeType, rayStartParams map[string]string, fqdnRayIP string, headPort string, creator string) {
+	// TODO: Audit all environment variables to identify which should not be modified by users.
 	// set the port RAY_PORT
 	// set the password?
 	container := &pod.Spec.Containers[rayContainerIndex]
@@ -568,10 +569,22 @@ func setContainerEnvVars(pod *v1.Pod, rayContainerIndex int, rayNodeType rayiov1
 		)
 	}
 
+	// The RAY_CLUSTER_NAME environment variable is managed by KubeRay and should not be set by the user.
+	clusterNameEnv := v1.EnvVar{
+		Name: RAY_CLUSTER_NAME,
+		ValueFrom: &v1.EnvVarSource{
+			FieldRef: &v1.ObjectFieldSelector{
+				FieldPath: fmt.Sprintf("metadata.labels['%s']", RayClusterLabelKey),
+			},
+		},
+	}
+	container.Env = append(container.Env, clusterNameEnv)
+
 	if !envVarExists(RAY_PORT, container.Env) {
 		portEnv := v1.EnvVar{Name: RAY_PORT, Value: headPort}
 		container.Env = append(container.Env, portEnv)
 	}
+
 	if strings.ToLower(creator) == RayServiceCreatorLabelValue {
 		// Only add this env for Ray Service cluster to improve service SLA.
 		if !envVarExists(RAY_TIMEOUT_MS_TASK_WAIT_FOR_DEATH_INFO, container.Env) {

diff --git a/ray-operator/controllers/ray/common/pod_test.go b/ray-operator/controllers/ray/common/pod_test.go
@@ -191,10 +191,10 @@ var autoscalerContainer = v1.Container{
 	ImagePullPolicy: v1.PullAlways,
 	Env: []v1.EnvVar{
 		{
-			Name: "RAY_CLUSTER_NAME",
+			Name: RAY_CLUSTER_NAME,
 			ValueFrom: &v1.EnvVarSource{
 				FieldRef: &v1.ObjectFieldSelector{
-					FieldPath: "metadata.labels['ray.io/cluster']",
+					FieldPath: fmt.Sprintf("metadata.labels['%s']", RayClusterLabelKey),
 				},
 			},
 		},
@@ -326,9 +326,18 @@ func checkPodEnv(t *testing.T, pod v1.Pod, envName string, expectedValue string)
 	foundEnv := false
 	for _, env := range pod.Spec.Containers[0].Env {
 		if env.Name == envName {
-			if !(env.Value == expectedValue) {
-				t.Fatalf("Expected `%v` but got `%v`", expectedValue, env.Value)
+			if env.Value != "" {
+				if env.Value != expectedValue {
+					t.Fatalf("Expected `%v` but got `%v`", expectedValue, env.Value)
+				}
+			} else {
+				// env.ValueFrom is the source for the environment variable's value. Cannot be used if value is not empty.
+				// See https://pkg.go.dev/k8s.io/api/core/v1#EnvVar for more details.
+				if env.ValueFrom.FieldRef.FieldPath != expectedValue {
+					t.Fatalf("Expected `%v` but got `%v`", expectedValue, env.ValueFrom.FieldRef.FieldPath)
+				}
 			}
+
 			foundEnv = true
 			break
 		}
@@ -393,6 +402,7 @@ func TestBuildPod(t *testing.T) {
 	checkPodEnv(t, pod, RAY_ADDRESS, "raycluster-sample-head-svc.default.svc.cluster.local:6379")
 	checkPodEnv(t, pod, FQ_RAY_IP, "raycluster-sample-head-svc.default.svc.cluster.local")
 	checkPodEnv(t, pod, RAY_IP, "raycluster-sample-head-svc")
+	checkPodEnv(t, pod, RAY_CLUSTER_NAME, fmt.Sprintf("metadata.labels['%s']", RayClusterLabelKey))
 
 	// Check RayStartParams
 	expectedResult = fmt.Sprintf("%s:6379", fqdnRayIP)