Added init containers(optional) configuration in worker-goroup spec

apiserver/pkg/util/cluster.go

@@ -418,6 +418,7 @@ func constructRayImage(containerImage string, version string) string {
 func buildWorkerPodTemplate(imageVersion string, envs *api.EnvironmentVariables, spec *api.WorkerGroupSpec, computeRuntime *api.ComputeTemplate) (*corev1.PodTemplateSpec, error) {
 	// If user doesn't provide the image, let's use the default image instead.
 	// TODO: verify the versions in the range
+	fmt.Println(spec)
 	image := constructRayImage(RayClusterDefaultImageRepository, imageVersion)
 	if len(spec.Image) != 0 {
 		image = spec.Image
@@ -441,6 +442,30 @@ func buildWorkerPodTemplate(imageVersion string, envs *api.EnvironmentVariables,
 		return nil, err
 	}
 
+	// Build init containers
+	var initContainers []corev1.Container
+	for _, initContainerSpec := range spec.InitContainers {
+		// Map WorkerGroupSpec.InitContainer to corev1.Container
+		initContainer := corev1.Container{
+			Name:            initContainerSpec.Name,
+			Image:           initContainerSpec.Image,
+			ImagePullPolicy: corev1.PullIfNotPresent, // Default to IfNotPresent
+
+			// Add volumes
+			VolumeMounts: buildVolumeMounts(initContainerSpec.Volumes),
+
+			// Add environment variables, if present
+			Env: convertEnvironmentVariables(initContainerSpec.Environment),
+		}
+
+		// Set image pull policy if explicitly given
+		if len(initContainerSpec.ImagePullPolicy) > 0 && strings.ToLower(initContainerSpec.ImagePullPolicy) == "always" {
+			initContainer.ImagePullPolicy = corev1.PullAlways
+		}
+
+		initContainers = append(initContainers, initContainer)
+	}
+
 	podTemplateSpec := corev1.PodTemplateSpec{
 		ObjectMeta: metav1.ObjectMeta{
 			Annotations: buildNodeGroupAnnotations(computeRuntime, spec.Image),
@@ -543,7 +568,8 @@ func buildWorkerPodTemplate(imageVersion string, envs *api.EnvironmentVariables,
 					SecurityContext: buildSecurityContext(spec.SecurityContext),
 				},
 			},
-			Volumes: vols,
+			InitContainers: initContainers,
+			Volumes:        vols,
 		},
 	}
 

apiserver/test/cluster/cluster/cluster

@@ -42,6 +42,32 @@ curl -X POST 'localhost:8888/apis/v1/namespaces/default/clusters' \
             "source": "ray-job-code-sample",
             "items": {"sample_code.py" : "sample_code.py"}
           }
+        ],
+        "initContainers": [
+          {
+            "name": "init-container-demo",
+            "image": "busybox",
+            "volumes": [
+              {
+                "mountPath": "/data",
+                "volumeType": "PERSISTENT_VOLUME_CLAIM",
+                "name": "init-data-volume",
+                "source": "data-pvc",
+                "readOnly": false,
+                "accessMode": "RWO",
+                "hostPathType": "DIRECTORY",
+                "storageClassName": "standard",
+                "storage": "1Gi"
+              }
+            ],
+            "environment": {
+              "values": {
+                "INIT_JOB_TYPE": "worker",
+                "INIT_COMMAND":   "echo Initializing worker node; sleep 5",
+              }
+            },
+            "imagePullPolicy": "IfNotPresent"
+          }
         ]
       }
     ]

apiserver/test/e2e/cluster_server_e2e_test.go

@@ -426,6 +426,62 @@ func TestCreateClusterEndpoint(t *testing.T) {
 				HTTPStatusCode: http.StatusBadRequest,
 			},
 		},
+		{
+			Name: "Create cluster with init container in worker group",
+			Input: &api.CreateClusterRequest{
+				Cluster: &api.Cluster{
+					Name:        tCtx.GetNextName(),
+					Namespace:   tCtx.GetNamespaceName(),
+					User:        "boris",
+					Version:     tCtx.GetRayVersion(),
+					Environment: api.Cluster_DEV,
+					ClusterSpec: &api.ClusterSpec{
+						HeadGroupSpec: &api.HeadGroupSpec{
+							ComputeTemplate: tCtx.GetComputeTemplateName(),
+							Image:           tCtx.GetRayImage(),
+							ServiceType:     "NodePort",
+							RayStartParams: map[string]string{
+								"dashboard-host":      "0.0.0.0",
+								"metrics-export-port": "8080",
+							},
+						},
+						WorkerGroupSpec: []*api.WorkerGroupSpec{
+							{
+								GroupName:       "small-wg",
+								ComputeTemplate: tCtx.GetComputeTemplateName(),
+								Image:           tCtx.GetRayImage(),
+								Replicas:        1,
+								MinReplicas:     1,
+								MaxReplicas:     5,
+								InitContainers: []*api.InitContainer{
+									{
+										Name:  "init-container-demo",
+										Image: "busybox",
+										Volumes: []*api.Volume{
+											{
+												MountPath:  "/data",
+												VolumeType: api.Volume_PERSISTENT_VOLUME_CLAIM,
+												Name:       "init-data-volume",
+												Source:     "data-pvc",
+												ReadOnly:   false,
+											},
+										},
+										Environment: &api.EnvironmentVariables{
+											Values: map[string]string{
+												"INIT_TASK_TYPE": "worker",
+												"INIT_COMMAND":   "echo Initializing worker node; sleep 5",
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+				Namespace: tCtx.GetNamespaceName(),
+			},
+			ExpectedError: nil,
+		},
 	}
 	// Execute tests sequentially
 	for _, tc := range tests {

proto/cluster.proto

@@ -352,6 +352,21 @@ message WorkerGroupSpec {
   string imagePullPolicy = 14;
   // Optional. Configure the security context for the worker container for debugging etc.
   SecurityContext security_context = 15;
+  // Optional. init container specs
+  repeated InitContainer init_containers = 16;
+}
+
+message InitContainer {
+  // Required. Group name of the current worker group
+  string name = 1 [(google.api.field_behavior) = REQUIRED];
+  // Required. This field will be used to retrieve right init container
+  string image = 2 [(google.api.field_behavior) = REQUIRED];
+  // Optional. The volumes mount to init container
+  repeated Volume volumes = 3;
+  // Optional. Environment variables for init container
+  EnvironmentVariables environment = 5;
+  // Optional image pull policy We only support Always and ifNotPresent
+  string imagePullPolicy = 6;
 }
 
 message ClusterEvent {