Wait for workers before submitted jobs in sidecar mode

marosset · marosset · commit 26006312a6f3 · 2026-01-21T15:35:33.000-08:00
Signed-off-by: Mark Rossett &lt;marosset@microsoft.com&gt;
diff --git a/ray-operator/controllers/ray/common/job.go b/ray-operator/controllers/ray/common/job.go
@@ -81,6 +81,26 @@ func GetMetadataJson(metadata map[string]string, rayVersion string) (string, err
 	return pkgutils.ConvertByteSliceToString(metadataBytes), nil
 }
 
+// GetMinReplicasFromSpec calculates the minimum expected worker replicas from the RayClusterSpec.
+// This is used in SidecarMode to determine how many workers should be registered before submitting the job.
+func GetMinReplicasFromSpec(rayClusterSpec *rayv1.RayClusterSpec) int32 {
+	if rayClusterSpec == nil {
+		return 0
+	}
+	count := int32(0)
+	for _, nodeGroup := range rayClusterSpec.WorkerGroupSpecs {
+		if nodeGroup.Suspend != nil && *nodeGroup.Suspend {
+			continue
+		}
+		minReplicas := int32(0)
+		if nodeGroup.MinReplicas != nil {
+			minReplicas = *nodeGroup.MinReplicas
+		}
+		count += minReplicas * nodeGroup.NumOfHosts
+	}
+	return count
+}
+
 // BuildJobSubmitCommand builds the `ray job submit` command based on submission mode.
 func BuildJobSubmitCommand(rayJobInstance *rayv1.RayJob, submissionMode rayv1.JobSubmissionMode) ([]string, error) {
 	var address string
@@ -139,6 +159,29 @@ func BuildJobSubmitCommand(rayJobInstance *rayv1.RayJob, submissionMode rayv1.Jo
 			"do", "echo", strconv.Quote("Waiting for Ray Dashboard GCS to become healthy at " + address + " ..."), ";", "sleep", "2", ";", "done", ";",
 		}
 		cmd = append(cmd, waitLoop...)
+
+		// Wait for the expected number of worker nodes to register for the Ray cluster.
+		// RAY_EXPECTED_MIN_WORKERS is set by the controller based on the MinReplicas in the RayClusterSpec.
+		// The loop queries the Ray Dashboard API to get the number of alive nodes and
+		// continues until the number of alive nodes is equal to (expected_workers + 1) for head node.
+		// This ensures that worker pods are connected before the job is submitted otherwise
+		// the jobs may run on the Head node.
+		//
+		// Note: This loop will never timeout and will wait indefinately if workers never register.
+		// This can be mitigated by setting the RayJob's `activeDeadlineSeconds` field
+		// to enforce a maximum job execution time.
+		waitForNodesLoop := []string{
+			"if", "[", "-n", "\"$" + utils.RAY_EXPECTED_MIN_WORKERS + "\"", "]", "&&", "[", "\"$" + utils.RAY_EXPECTED_MIN_WORKERS + "\"", "-gt", "\"0\"", "]", ";", "then",
+			"EXPECTED_NODES=$(($" + utils.RAY_EXPECTED_MIN_WORKERS + " + 1))", ";",
+			"echo", strconv.Quote("Waiting for $EXPECTED_NODES nodes (1 head + $" + utils.RAY_EXPECTED_MIN_WORKERS + " workers) to register..."), ";",
+			"until", "[",
+			"\"$(wget -q -O- " + address + "/nodes?view=summary 2>/dev/null | python3 -c \"import sys,json; d=json.load(sys.stdin); print(len([n for n in d.get('data',{}).get('summary',[]) if n.get('raylet',{}).get('state','')=='ALIVE']))\" 2>/dev/null || echo 0)\"",
+			"-ge", "\"$EXPECTED_NODES\"", "]", ";",
+			"do", "echo", strconv.Quote("Waiting for Ray nodes to register. Expected: $EXPECTED_NODES ..."), ";", "sleep", "2", ";", "done", ";",
+			"echo", strconv.Quote("All expected nodes are registered."), ";",
+			"fi", ";",
+		}
+		cmd = append(cmd, waitForNodesLoop...)
 	}
 
 	// In Sidecar mode, we only support RayJob level retry, which means that the submitter retry won't happen,
diff --git a/ray-operator/controllers/ray/common/job_test.go b/ray-operator/controllers/ray/common/job_test.go
@@ -112,7 +112,9 @@ func TestBuildJobSubmitCommandWithSidecarMode(t *testing.T) {
 		},
 	}
 
+	address := "http://127.0.0.1:8265"
 	expected := []string{
+		// Wait for Dashboard GCS health
 		"until",
 		fmt.Sprintf(
 			utils.BaseWgetHealthCommand,
@@ -121,8 +123,19 @@ func TestBuildJobSubmitCommandWithSidecarMode(t *testing.T) {
 			utils.RayDashboardGCSHealthPath,
 		),
 		">/dev/null", "2>&1", ";",
-		"do", "echo", strconv.Quote("Waiting for Ray Dashboard GCS to become healthy at http://127.0.0.1:8265 ..."), ";", "sleep", "2", ";", "done", ";",
-		"ray", "job", "submit", "--address", "http://127.0.0.1:8265",
+		"do", "echo", strconv.Quote("Waiting for Ray Dashboard GCS to become healthy at " + address + " ..."), ";", "sleep", "2", ";", "done", ";",
+		// Wait for expected nodes to register
+		"if", "[", "-n", "\"$" + utils.RAY_EXPECTED_MIN_WORKERS + "\"", "]", "&&", "[", "\"$" + utils.RAY_EXPECTED_MIN_WORKERS + "\"", "-gt", "\"0\"", "]", ";", "then",
+		"EXPECTED_NODES=$(($" + utils.RAY_EXPECTED_MIN_WORKERS + " + 1))", ";",
+		"echo", strconv.Quote("Waiting for $EXPECTED_NODES nodes (1 head + $" + utils.RAY_EXPECTED_MIN_WORKERS + " workers) to register..."), ";",
+		"until", "[",
+		"\"$(wget -q -O- " + address + "/nodes?view=summary 2>/dev/null | python3 -c \"import sys,json; d=json.load(sys.stdin); print(len([n for n in d.get('data',{}).get('summary',[]) if n.get('raylet',{}).get('state','')=='ALIVE']))\" 2>/dev/null || echo 0)\"",
+		"-ge", "\"$EXPECTED_NODES\"", "]", ";",
+		"do", "echo", strconv.Quote("Waiting for Ray nodes to register. Expected: $EXPECTED_NODES ..."), ";", "sleep", "2", ";", "done", ";",
+		"echo", strconv.Quote("All expected nodes are registered."), ";",
+		"fi", ";",
+		// Job submit command
+		"ray", "job", "submit", "--address", address,
 		"--runtime-env-json", strconv.Quote(`{"test":"test"}`),
 		"--metadata-json", strconv.Quote(`{"testKey":"testValue"}`),
 		"--submission-id", "testJobId",
@@ -240,3 +253,120 @@ func TestGetSubmitterTemplate(t *testing.T) {
 	template := GetSubmitterTemplate(&rayJob.Spec, &rayCluster.Spec)
 	assert.Equal(t, template.Spec.Containers[0].Image, rayCluster.Spec.HeadGroupSpec.Template.Spec.Containers[utils.RayContainerIndex].Image)
 }
+
+func TestGetMinReplicasFromSpec(t *testing.T) {
+	tests := []struct {
+		spec     *rayv1.RayClusterSpec
+		name     string
+		expected int32
+	}{
+		{
+			name:     "nil spec returns 0",
+			spec:     nil,
+			expected: 0,
+		},
+		{
+			name: "no worker groups returns 0",
+			spec: &rayv1.RayClusterSpec{
+				WorkerGroupSpecs: []rayv1.WorkerGroupSpec{},
+			},
+			expected: 0,
+		},
+		{
+			name: "single worker group with minReplicas",
+			spec: &rayv1.RayClusterSpec{
+				WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+					{
+						MinReplicas: ptrInt32(2),
+						NumOfHosts:  1,
+					},
+				},
+			},
+			expected: 2,
+		},
+		{
+			name: "multiple worker groups",
+			spec: &rayv1.RayClusterSpec{
+				WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+					{
+						MinReplicas: ptrInt32(2),
+						NumOfHosts:  1,
+					},
+					{
+						MinReplicas: ptrInt32(3),
+						NumOfHosts:  1,
+					},
+				},
+			},
+			expected: 5,
+		},
+		{
+			name: "worker group with NumOfHosts > 1",
+			spec: &rayv1.RayClusterSpec{
+				WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+					{
+						MinReplicas: ptrInt32(2),
+						NumOfHosts:  2,
+					},
+				},
+			},
+			expected: 4,
+		},
+		{
+			name: "suspended worker group is skipped",
+			spec: &rayv1.RayClusterSpec{
+				WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+					{
+						MinReplicas: ptrInt32(2),
+						NumOfHosts:  1,
+						Suspend:     ptrBool(true),
+					},
+					{
+						MinReplicas: ptrInt32(3),
+						NumOfHosts:  1,
+					},
+				},
+			},
+			expected: 3,
+		},
+		{
+			name: "nil minReplicas defaults to 0",
+			spec: &rayv1.RayClusterSpec{
+				WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+					{
+						MinReplicas: nil,
+						NumOfHosts:  1,
+					},
+				},
+			},
+			expected: 0,
+		},
+		{
+			name: "NumOfHosts 0 results in 0 workers for that group",
+			spec: &rayv1.RayClusterSpec{
+				WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+					{
+						MinReplicas: ptrInt32(2),
+						NumOfHosts:  0,
+					},
+				},
+			},
+			expected: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := GetMinReplicasFromSpec(tt.spec)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func ptrInt32(v int32) *int32 {
+	return &v
+}
+
+func ptrBool(v bool) *bool {
+	return &v
+}
diff --git a/ray-operator/controllers/ray/rayjob_controller.go b/ray-operator/controllers/ray/rayjob_controller.go
@@ -613,6 +613,13 @@ func configureSubmitterContainer(container *corev1.Container, rayJobInstance *ra
 	// ray job submit --address=http://$RAY_DASHBOARD_ADDRESS --submission-id=$RAY_JOB_SUBMISSION_ID ...
 	container.Env = append(container.Env, corev1.EnvVar{Name: utils.RAY_DASHBOARD_ADDRESS, Value: rayJobInstance.Status.DashboardURL})
 	container.Env = append(container.Env, corev1.EnvVar{Name: utils.RAY_JOB_SUBMISSION_ID, Value: rayJobInstance.Status.JobId})
+
+	// In SidecarMode, pass the expected minimum worker count so the submitter can wait for workers to register
+	if submissionMode == rayv1.SidecarMode && rayJobInstance.Spec.RayClusterSpec != nil {
+		minWorkers := common.GetMinReplicasFromSpec(rayJobInstance.Spec.RayClusterSpec)
+		container.Env = append(container.Env, corev1.EnvVar{Name: utils.RAY_EXPECTED_MIN_WORKERS, Value: strconv.Itoa(int(minWorkers))})
+	}
+
 	if rayClusterInstance != nil && utils.IsAuthEnabled(&rayClusterInstance.Spec) {
 		common.SetContainerTokenAuthEnvVars(rayClusterInstance.Name, container)
 	}
diff --git a/ray-operator/controllers/ray/utils/constant.go b/ray-operator/controllers/ray/utils/constant.go
@@ -151,8 +151,9 @@ const (
 
 	// Environment variables for RayJob submitter Kubernetes Job.
 	// Example: ray job submit --address=http://$RAY_DASHBOARD_ADDRESS --submission-id=$RAY_JOB_SUBMISSION_ID ...
-	RAY_DASHBOARD_ADDRESS = "RAY_DASHBOARD_ADDRESS"
-	RAY_JOB_SUBMISSION_ID = "RAY_JOB_SUBMISSION_ID"
+	RAY_DASHBOARD_ADDRESS    = "RAY_DASHBOARD_ADDRESS"
+	RAY_EXPECTED_MIN_WORKERS = "RAY_EXPECTED_MIN_WORKERS"
+	RAY_JOB_SUBMISSION_ID    = "RAY_JOB_SUBMISSION_ID"
 
 	// Environment variables for Ray Autoscaler V2.
 	// The value of RAY_CLOUD_INSTANCE_ID is the Pod name for Autoscaler V2 alpha. This may change in the future.
diff --git a/ray-operator/test/e2erayjob/rayjob_sidecar_mode_test.go b/ray-operator/test/e2erayjob/rayjob_sidecar_mode_test.go
@@ -290,4 +290,61 @@ env_vars:
 
 		LogWithTimestamp(test.T(), "RayJob %s/%s completed successfully with auth token", rayJob.Namespace, rayJob.Name)
 	})
+
+	test.T().Run("SidecarMode waits for workers before job submission", func(_ *testing.T) {
+		// This test verifies that SidecarMode waits for workers to register before submitting
+		// the job. We create a RayJob with a worker that has a slow init container (30s delay).
+		// The inline Python script verifies that workers are available when it runs. If the
+		// sidecar submitted the job before workers were ready, the job would fail because no
+		// workers would be found.
+
+		// Create a worker template with a slow init container to simulate slow worker startup
+		workerTemplate := WorkerPodTemplateApplyConfiguration()
+		workerTemplate.Spec.WithInitContainers(corev1ac.Container().
+			WithName("init-delay").
+			WithImage(GetRayImage()).
+			WithCommand("bash", "-c", "echo 'Simulating slow worker startup...'; sleep 30"))
+
+		// Simple inline Python script that verifies workers are available
+		// Uses semicolons to avoid multi-line issues with shell escaping
+		verifyWorkersScript := `import ray; ray.init(); nodes = ray.nodes(); worker_count = sum(1 for n in nodes if n["Alive"] and not n.get("Resources", {}).get("node:__internal_head__")); print("Worker nodes:", worker_count); assert worker_count >= 1`
+
+		rayJobAC := rayv1ac.RayJob("wait-for-workers", namespace.Name).
+			WithSpec(rayv1ac.RayJobSpec().
+				WithSubmissionMode(rayv1.SidecarMode).
+				WithEntrypoint("python -c '" + verifyWorkersScript + "'").
+				WithShutdownAfterJobFinishes(true).
+				WithRayClusterSpec(rayv1ac.RayClusterSpec().
+					WithRayVersion(GetRayVersion()).
+					WithHeadGroupSpec(rayv1ac.HeadGroupSpec().
+						WithRayStartParams(map[string]string{"dashboard-host": "0.0.0.0"}).
+						WithTemplate(HeadPodTemplateApplyConfiguration())).
+					WithWorkerGroupSpecs(rayv1ac.WorkerGroupSpec().
+						WithReplicas(1).
+						WithMinReplicas(1).
+						WithMaxReplicas(1).
+						WithGroupName("slow-worker").
+						WithRayStartParams(map[string]string{"num-cpus": "1"}).
+						WithTemplate(workerTemplate))))
+
+		rayJob, err := test.Client().Ray().RayV1().RayJobs(namespace.Name).Apply(test.Ctx(), rayJobAC, TestApplyOptions)
+		g.Expect(err).NotTo(HaveOccurred())
+		LogWithTimestamp(test.T(), "Created RayJob %s/%s successfully", rayJob.Namespace, rayJob.Name)
+
+		// Wait for the RayJob to complete. The job will fail if workers aren't available.
+		LogWithTimestamp(test.T(), "Waiting for RayJob %s/%s to complete (this may take ~40s due to worker init delay)", rayJob.Namespace, rayJob.Name)
+		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutLong).
+			Should(WithTransform(RayJobStatus, Satisfy(rayv1.IsJobTerminal)))
+
+		// Assert the RayJob has completed successfully
+		// If the sidecar didn't wait for workers, the job would fail with assertion error
+		g.Expect(GetRayJob(test, rayJob.Namespace, rayJob.Name)).
+			To(WithTransform(RayJobStatus, Equal(rayv1.JobStatusSucceeded)))
+
+		// And the RayJob deployment status is updated accordingly
+		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name)).
+			Should(WithTransform(RayJobDeploymentStatus, Equal(rayv1.JobDeploymentStatusComplete)))
+
+		LogWithTimestamp(test.T(), "RayJob %s/%s completed successfully - workers were available before job submission", rayJob.Namespace, rayJob.Name)
+	})
 }
