Add EnableUsageStats start parameter

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>
flyteorg · Sep 21, 2023 · 71c70f1 · 71c70f1
1 parent 370d13e
commit 71c70f1
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 24 deletions.
diff --git a/go/tasks/plugins/k8s/ray/config.go b/go/tasks/plugins/k8s/ray/config.go
@@ -18,18 +18,19 @@ var (
 		ServiceType:              "NodePort",
 		IncludeDashboard:         true,
 		DashboardHost:            "0.0.0.0",
+		EnableUsageStats:         false,
 		Defaults: DefaultConfig{
 			HeadNode: NodeConfig{
 				StartParameters: map[string]string{
 					// Disable usage reporting by default: https://docs.ray.io/en/latest/cluster/usage-stats.html
-					"disable-usage-stats": "true",
+					DisableUsageStatsStartParameter: "true",
 				},
 				IPAddress: "$MY_POD_IP",
 			},
 			WorkerNode: NodeConfig{
 				StartParameters: map[string]string{
 					// Disable usage reporting by default: https://docs.ray.io/en/latest/cluster/usage-stats.html
-					"disable-usage-stats": "true",
+					DisableUsageStatsStartParameter: "true",
 				},
 				IPAddress: "$MY_POD_IP",
 			},
@@ -80,6 +81,8 @@ type Config struct {
 	Logs logs.LogConfig `json:"logs" pflag:"-,Log configuration for ray jobs"`
 
 	Defaults DefaultConfig `json:"defaults" pflag:"-,Default configuration for ray jobs"`
+
+	EnableUsageStats bool `json:"enableUsageStats" pflag:",Enable usage stats for ray jobs. These stats are submitted to usage-stats.ray.io per https://docs.ray.io/en/latest/cluster/usage-stats.html"`
 }
 
 type DefaultConfig struct {

diff --git a/go/tasks/plugins/k8s/ray/config_flags.go b/go/tasks/plugins/k8s/ray/config_flags.go
diff --git a/go/tasks/plugins/k8s/ray/config_flags_test.go b/go/tasks/plugins/k8s/ray/config_flags_test.go
diff --git a/go/tasks/plugins/k8s/ray/ray.go b/go/tasks/plugins/k8s/ray/ray.go
@@ -28,11 +28,12 @@ import (
 )
 
 const (
-	rayTaskType      = "ray"
-	KindRayJob       = "RayJob"
-	IncludeDashboard = "include-dashboard"
-	NodeIPAddress    = "node-ip-address"
-	DashboardHost    = "dashboard-host"
+	rayTaskType                     = "ray"
+	KindRayJob                      = "RayJob"
+	IncludeDashboard                = "include-dashboard"
+	NodeIPAddress                   = "node-ip-address"
+	DashboardHost                   = "dashboard-host"
+	DisableUsageStatsStartParameter = "disable-usage-stats"
 )
 
 type rayJobResourceHandler struct {
@@ -98,6 +99,10 @@ func (rayJobResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsC
 		headNodeRayStartParams[DashboardHost] = cfg.DashboardHost
 	}
 
+	if _, exists := headNodeRayStartParams[DisableUsageStatsStartParameter]; !exists && !cfg.EnableUsageStats {
+		headNodeRayStartParams[DisableUsageStatsStartParameter] = "true"
+	}
+
 	enableIngress := true
 	rayClusterSpec := rayv1alpha1.RayClusterSpec{
 		HeadGroupSpec: rayv1alpha1.HeadGroupSpec{
@@ -133,6 +138,10 @@ func (rayJobResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsC
 			workerNodeRayStartParams[NodeIPAddress] = cfg.Defaults.WorkerNode.IPAddress
 		}
 
+		if _, exists := workerNodeRayStartParams[DisableUsageStatsStartParameter]; !exists && !cfg.EnableUsageStats {
+			workerNodeRayStartParams[DisableUsageStatsStartParameter] = "true"
+		}
+
 		workerNodeSpec := rayv1alpha1.WorkerGroupSpec{
 			GroupName:      spec.GroupName,
 			MinReplicas:    &minReplicas,

diff --git a/go/tasks/plugins/k8s/ray/ray_test.go b/go/tasks/plugins/k8s/ray/ray_test.go
@@ -169,7 +169,9 @@ func TestBuildResourceRay(t *testing.T) {
 	assert.Equal(t, ray.Spec.RayClusterSpec.HeadGroupSpec.Replicas, &headReplica)
 	assert.Equal(t, ray.Spec.RayClusterSpec.HeadGroupSpec.Template.Spec.ServiceAccountName, serviceAccount)
 	assert.Equal(t, ray.Spec.RayClusterSpec.HeadGroupSpec.RayStartParams,
-		map[string]string{"dashboard-host": "0.0.0.0", "include-dashboard": "true", "node-ip-address": "$MY_POD_IP", "num-cpus": "1"})
+		map[string]string{
+			"dashboard-host": "0.0.0.0", "disable-usage-stats": "true", "include-dashboard": "true",
+			"node-ip-address": "$MY_POD_IP", "num-cpus": "1"})
 	assert.Equal(t, ray.Spec.RayClusterSpec.HeadGroupSpec.Template.Annotations, map[string]string{"annotation-1": "val1"})
 	assert.Equal(t, ray.Spec.RayClusterSpec.HeadGroupSpec.Template.Labels, map[string]string{"label-1": "val1"})
 	assert.Equal(t, ray.Spec.RayClusterSpec.HeadGroupSpec.Template.Spec.Tolerations, toleration)
@@ -180,7 +182,7 @@ func TestBuildResourceRay(t *testing.T) {
 	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].MaxReplicas, &workerReplica)
 	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].GroupName, workerGroupName)
 	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Spec.ServiceAccountName, serviceAccount)
-	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].RayStartParams, map[string]string{"node-ip-address": "$MY_POD_IP"})
+	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].RayStartParams, map[string]string{"disable-usage-stats": "true", "node-ip-address": "$MY_POD_IP"})
 	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Annotations, map[string]string{"annotation-1": "val1"})
 	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Labels, map[string]string{"label-1": "val1"})
 	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Tolerations, toleration)