[Feature][RayJob] Use Use RayContainerIndex instead of 0 (ray-project…

…#1427) Use Use RayContainerIndex instead of 0
opendatahub-io · Sep 16, 2023 · 6d5020f · 6d5020f
1 parent 4492fe2
commit 6d5020f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/ray-operator/controllers/ray/rayjob_controller.go b/ray-operator/controllers/ray/rayjob_controller.go
@@ -360,7 +360,7 @@ func (r *RayJobReconciler) getSubmitterTemplate(rayJobInstance *rayv1alpha1.RayJ
 	}
 
 	// If the command in the submitter pod template isn't set, use the default command.
-	if len(submitterTemplate.Spec.Containers[0].Command) == 0 {
+	if len(submitterTemplate.Spec.Containers[common.RayContainerIndex].Command) == 0 {
 		// Check for deprecated 'runtimeEnv' field usage and log a warning.
 		if len(rayJobInstance.Spec.RuntimeEnv) > 0 {
 			r.Log.Info("Warning: The 'runtimeEnv' field is deprecated. Please use 'runtimeEnvYAML' instead.")
@@ -370,14 +370,14 @@ func (r *RayJobReconciler) getSubmitterTemplate(rayJobInstance *rayv1alpha1.RayJ
 		if err != nil {
 			return v1.PodTemplateSpec{}, err
 		}
-		submitterTemplate.Spec.Containers[0].Command = k8sJobCommand
+		submitterTemplate.Spec.Containers[common.RayContainerIndex].Command = k8sJobCommand
 		r.Log.Info("No command is specified in the user-provided template. Default command is used", "command", k8sJobCommand)
 	} else {
-		r.Log.Info("User-provided command is used", "command", submitterTemplate.Spec.Containers[0].Command)
+		r.Log.Info("User-provided command is used", "command", submitterTemplate.Spec.Containers[common.RayContainerIndex].Command)
 	}
 
 	// Set PYTHONUNBUFFERED=1 for real-time logging
-	submitterTemplate.Spec.Containers[0].Env = append(submitterTemplate.Spec.Containers[0].Env, v1.EnvVar{
+	submitterTemplate.Spec.Containers[common.RayContainerIndex].Env = append(submitterTemplate.Spec.Containers[common.RayContainerIndex].Env, v1.EnvVar{
 		Name:  PythonUnbufferedEnvVarName,
 		Value: "1",
 	})

diff --git a/ray-operator/controllers/ray/rayjob_controller_unit_test.go b/ray-operator/controllers/ray/rayjob_controller_unit_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	rayv1alpha1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1alpha1"
+	"github.com/ray-project/kuberay/ray-operator/controllers/ray/common"
 	"github.com/stretchr/testify/assert"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
@@ -121,25 +122,25 @@ func TestGetSubmitterTemplate(t *testing.T) {
 	// Test 1: User provided template with command
 	submitterTemplate, err := r.getSubmitterTemplate(rayJobInstanceWithTemplate)
 	assert.NoError(t, err)
-	assert.Equal(t, "user-command", submitterTemplate.Spec.Containers[0].Command[0])
+	assert.Equal(t, "user-command", submitterTemplate.Spec.Containers[common.RayContainerIndex].Command[0])
 
 	// Test 2: User provided template without command
-	rayJobInstanceWithTemplate.Spec.SubmitterPodTemplate.Spec.Containers[0].Command = []string{}
+	rayJobInstanceWithTemplate.Spec.SubmitterPodTemplate.Spec.Containers[common.RayContainerIndex].Command = []string{}
 	submitterTemplate, err = r.getSubmitterTemplate(rayJobInstanceWithTemplate)
 	assert.NoError(t, err)
-	assert.Equal(t, ([]string{"ray", "job", "submit", "--address", "http://test-url", "--", "echo", "hello", "world"}), submitterTemplate.Spec.Containers[0].Command)
+	assert.Equal(t, ([]string{"ray", "job", "submit", "--address", "http://test-url", "--", "echo", "hello", "world"}), submitterTemplate.Spec.Containers[common.RayContainerIndex].Command)
 
 	// Test 3: User did not provide template, should use the image of the Ray Head
 	submitterTemplate, err = r.getSubmitterTemplate(rayJobInstanceWithoutTemplate)
 	assert.NoError(t, err)
-	assert.Equal(t, ([]string{"ray", "job", "submit", "--address", "http://test-url", "--", "echo", "hello", "world"}), submitterTemplate.Spec.Containers[0].Command)
-	assert.Equal(t, "rayproject/ray:custom-version", submitterTemplate.Spec.Containers[0].Image)
+	assert.Equal(t, ([]string{"ray", "job", "submit", "--address", "http://test-url", "--", "echo", "hello", "world"}), submitterTemplate.Spec.Containers[common.RayContainerIndex].Command)
+	assert.Equal(t, "rayproject/ray:custom-version", submitterTemplate.Spec.Containers[common.RayContainerIndex].Image)
 
 	// Test 4: Check default PYTHONUNBUFFERED setting
 	submitterTemplate, err = r.getSubmitterTemplate(rayJobInstanceWithoutTemplate)
 	assert.NoError(t, err)
 	found := false
-	for _, envVar := range submitterTemplate.Spec.Containers[0].Env {
+	for _, envVar := range submitterTemplate.Spec.Containers[common.RayContainerIndex].Env {
 		if envVar.Name == PythonUnbufferedEnvVarName {
 			assert.Equal(t, "1", envVar.Value)
 			found = true