[Telemetry] Inject env identifying KubeRay. ray-project#562

Signed-off-by: Dmitri Gekhtman <[email protected]> This change has the KubeRay operator inject an env var identifying a Ray container as having been deployed using KubeRay. The purpose is to refine Ray usage stats collected by Ray's telemetry service (https://docs.ray.io/en/latest/cluster/usage-stats.html). See ray-project/ray#28490
lowang-bh · Sep 14, 2022 · b6ad236 · b6ad236
1 parent 8e37457
commit b6ad236
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 0 deletions.
diff --git a/ray-operator/controllers/ray/common/constant.go b/ray-operator/controllers/ray/common/constant.go
@@ -74,6 +74,7 @@ const (
 	RAY_SERVE_KV_TIMEOUT_S                  = "RAY_SERVE_KV_TIMEOUT_S"
 	SERVE_DEPLOYMENT_HANDLE_IS_SYNC         = "SERVE_DEPLOYMENT_HANDLE_IS_SYNC"
 	SERVE_CONTROLLER_PIN_ON_NODE            = "RAY_INTERNAL_SERVE_CONTROLLER_PIN_ON_NODE"
+	RAY_USAGE_STATS_KUBERAY_IN_USE          = "RAY_USAGE_STATS_KUBERAY_IN_USE"
 
 	// Ray core default configurations
 	DefaultRedisPassword = "5241590000000000"

diff --git a/ray-operator/controllers/ray/common/pod.go b/ray-operator/controllers/ray/common/pod.go
@@ -588,6 +588,10 @@ func setContainerEnvVars(pod *v1.Pod, rayContainerIndex int, rayNodeType rayiov1
 		addressEnv := v1.EnvVar{Name: RAY_ADDRESS, Value: rayAddress}
 		container.Env = append(container.Env, addressEnv)
 	}
+	if !envVarExists(RAY_USAGE_STATS_KUBERAY_IN_USE, container.Env) {
+		usageEnv := v1.EnvVar{Name: RAY_USAGE_STATS_KUBERAY_IN_USE, Value: "1"}
+		container.Env = append(container.Env, usageEnv)
+	}
 	if !envVarExists(REDIS_PASSWORD, container.Env) {
 		// setting the REDIS_PASSWORD env var from the params
 		port := v1.EnvVar{Name: REDIS_PASSWORD}

diff --git a/ray-operator/controllers/ray/common/pod_test.go b/ray-operator/controllers/ray/common/pod_test.go
@@ -311,6 +311,8 @@ func TestBuildPod(t *testing.T) {
 
 	// Check RAY_ADDRESS env.
 	checkPodEnv(t, pod, RAY_ADDRESS, "127.0.0.1:6379")
+	// Check usage stats env.
+	checkPodEnv(t, pod, RAY_USAGE_STATS_KUBERAY_IN_USE, "1")
 
 	// Check labels.
 	actualResult := pod.Labels[RayClusterLabelKey]