Refactor RayJob submitter template handling

Signed-off-by: You-Cheng Lin (Owen) <[email protected]>

Author: owenowenisme

ray-operator/controllers/ray/batchscheduler/yunikorn/yunikorn_task_groups.go

@@ -71,13 +71,7 @@ func newTaskGroupsFromRayClusterSpec(rayClusterSpec *v1.RayClusterSpec) *TaskGro
 func newTaskGroupsFromRayJobSpec(rayJobSpec *v1.RayJobSpec) *TaskGroups {
 	taskGroups := newTaskGroupsFromRayClusterSpec(rayJobSpec.RayClusterSpec)
 
-	// submitter group
-	var submitterGroupSpec corev1.PodSpec
-	if rayJobSpec.SubmitterPodTemplate != nil {
-		submitterGroupSpec = rayJobSpec.SubmitterPodTemplate.Spec
-	} else {
-		submitterGroupSpec = common.GetDefaultSubmitterTemplate(rayJobSpec.RayClusterSpec).Spec
-	}
+	submitterGroupSpec := common.GetSubmitterTemplate(rayJobSpec, rayJobSpec.RayClusterSpec).Spec
 
 	submitterPodMinResource := utils.CalculatePodResource(submitterGroupSpec)
 	taskGroups.addTaskGroup(

ray-operator/controllers/ray/common/job.go

@@ -198,8 +198,11 @@ func BuildJobSubmitCommand(rayJobInstance *rayv1.RayJob, submissionMode rayv1.Jo
 	return cmd, nil
 }
 
-// GetDefaultSubmitterTemplate creates a default submitter template for the Ray job.
-func GetDefaultSubmitterTemplate(rayClusterSpec *rayv1.RayClusterSpec) corev1.PodTemplateSpec {
+// GetSubmitterTemplate creates a default submitter template for the Ray job.
+func GetSubmitterTemplate(rayJobSpec *rayv1.RayJobSpec, rayClusterSpec *rayv1.RayClusterSpec) corev1.PodTemplateSpec {
+	if rayJobSpec.SubmitterPodTemplate != nil {
+		return *rayJobSpec.SubmitterPodTemplate.DeepCopy()
+	}
 	return corev1.PodTemplateSpec{
 		Spec: corev1.PodSpec{
 			Containers: []corev1.Container{

ray-operator/controllers/ray/common/job_test.go

@@ -218,7 +218,10 @@ func TestMetadataRaisesErrorBeforeRay26(t *testing.T) {
 	require.Error(t, err)
 }
 
-func TestGetDefaultSubmitterTemplate(t *testing.T) {
+func TestGetSubmitterTemplate(t *testing.T) {
+	rayJob := &rayv1.RayJob{
+		Spec: rayv1.RayJobSpec{},
+	}
 	rayCluster := &rayv1.RayCluster{
 		Spec: rayv1.RayClusterSpec{
 			HeadGroupSpec: rayv1.HeadGroupSpec{
@@ -234,6 +237,6 @@ func TestGetDefaultSubmitterTemplate(t *testing.T) {
 			},
 		},
 	}
-	template := GetDefaultSubmitterTemplate(&rayCluster.Spec)
+	template := GetSubmitterTemplate(&rayJob.Spec, &rayCluster.Spec)
 	assert.Equal(t, template.Spec.Containers[0].Image, rayCluster.Spec.HeadGroupSpec.Template.Spec.Containers[utils.RayContainerIndex].Image)
 }

ray-operator/controllers/ray/rayjob_controller.go

@@ -576,7 +576,7 @@ func (r *RayJobReconciler) createK8sJobIfNeed(ctx context.Context, rayJobInstanc
 	namespacedName := common.RayJobK8sJobNamespacedName(rayJobInstance)
 	if err := r.Client.Get(ctx, namespacedName, job); err != nil {
 		if errors.IsNotFound(err) {
-			submitterTemplate, err := getSubmitterTemplate(ctx, rayJobInstance, rayClusterInstance)
+			submitterTemplate, err := getSubmitterTemplate(rayJobInstance, rayClusterInstance)
 			if err != nil {
 				return err
 			}
@@ -597,18 +597,9 @@ func (r *RayJobReconciler) createK8sJobIfNeed(ctx context.Context, rayJobInstanc
 }
 
 // getSubmitterTemplate builds the submitter pod template for the Ray job.
-func getSubmitterTemplate(ctx context.Context, rayJobInstance *rayv1.RayJob, rayClusterInstance *rayv1.RayCluster) (corev1.PodTemplateSpec, error) {
-	logger := ctrl.LoggerFrom(ctx)
-	var submitterTemplate corev1.PodTemplateSpec
-
+func getSubmitterTemplate(rayJobInstance *rayv1.RayJob, rayClusterInstance *rayv1.RayCluster) (corev1.PodTemplateSpec, error) {
 	// Set the default value for the optional field SubmitterPodTemplate if not provided.
-	if rayJobInstance.Spec.SubmitterPodTemplate == nil {
-		submitterTemplate = common.GetDefaultSubmitterTemplate(&rayClusterInstance.Spec)
-		logger.Info("default submitter template is used")
-	} else {
-		submitterTemplate = *rayJobInstance.Spec.SubmitterPodTemplate.DeepCopy()
-		logger.Info("user-provided submitter template is used; the first container is assumed to be the submitter")
-	}
+	submitterTemplate := common.GetSubmitterTemplate(&rayJobInstance.Spec, &rayClusterInstance.Spec)
 
 	if err := configureSubmitterContainer(&submitterTemplate.Spec.Containers[utils.RayContainerIndex], rayJobInstance, rayv1.K8sJobMode); err != nil {
 		return corev1.PodTemplateSpec{}, err

ray-operator/controllers/ray/rayjob_controller_test.go

@@ -357,7 +357,7 @@ var _ = Context("RayJob with different submission modes", func() {
 			namespace := "default"
 			rayJob := rayJobTemplate("rayjob-invalid-test", namespace)
 			rayCluster := &rayv1.RayCluster{Spec: *rayJob.Spec.RayClusterSpec}
-			template := common.GetDefaultSubmitterTemplate(&rayCluster.Spec)
+			template := common.GetSubmitterTemplate(&rayJob.Spec, &rayCluster.Spec)
 			template.Spec.RestartPolicy = "" // Make it invalid to create a submitter. Ref: https://github.com/ray-project/kuberay/pull/2389#issuecomment-2359564334
 			rayJob.Spec.SubmitterPodTemplate = &template
 

ray-operator/controllers/ray/rayjob_controller_unit_test.go

@@ -160,37 +160,35 @@ func TestGetSubmitterTemplate(t *testing.T) {
 		},
 	}
 
-	ctx := context.Background()
-
 	// Test 1: User provided template with command
-	submitterTemplate, err := getSubmitterTemplate(ctx, rayJobInstanceWithTemplate, nil)
+	submitterTemplate, err := getSubmitterTemplate(rayJobInstanceWithTemplate, rayClusterInstance)
 	require.NoError(t, err)
 	assert.Equal(t, "user-command", submitterTemplate.Spec.Containers[utils.RayContainerIndex].Command[0])
 
 	// Test 2: User provided template without command
 	rayJobInstanceWithTemplate.Spec.SubmitterPodTemplate.Spec.Containers[utils.RayContainerIndex].Command = []string{}
-	submitterTemplate, err = getSubmitterTemplate(ctx, rayJobInstanceWithTemplate, nil)
+	submitterTemplate, err = getSubmitterTemplate(rayJobInstanceWithTemplate, rayClusterInstance)
 	require.NoError(t, err)
 	assert.Equal(t, []string{"/bin/bash", "-ce", "--"}, submitterTemplate.Spec.Containers[utils.RayContainerIndex].Command)
 	assert.Equal(t, []string{"if ! ray job status --address http://test-url test-job-id >/dev/null 2>&1 ; then ray job submit --address http://test-url --no-wait --submission-id test-job-id -- echo no quote 'single quote' \"double quote\" ; fi ; ray job logs --address http://test-url --follow test-job-id"}, submitterTemplate.Spec.Containers[utils.RayContainerIndex].Args)
 
 	// Test 3: User did not provide template, should use the image of the Ray Head
-	submitterTemplate, err = getSubmitterTemplate(ctx, rayJobInstanceWithoutTemplate, rayClusterInstance)
+	submitterTemplate, err = getSubmitterTemplate(rayJobInstanceWithoutTemplate, rayClusterInstance)
 	require.NoError(t, err)
 	assert.Equal(t, []string{"/bin/bash", "-ce", "--"}, submitterTemplate.Spec.Containers[utils.RayContainerIndex].Command)
 	assert.Equal(t, []string{"if ! ray job status --address http://test-url test-job-id >/dev/null 2>&1 ; then ray job submit --address http://test-url --no-wait --submission-id test-job-id -- echo no quote 'single quote' \"double quote\" ; fi ; ray job logs --address http://test-url --follow test-job-id"}, submitterTemplate.Spec.Containers[utils.RayContainerIndex].Args)
 	assert.Equal(t, "rayproject/ray:custom-version", submitterTemplate.Spec.Containers[utils.RayContainerIndex].Image)
 
 	// Test 4: Check default PYTHONUNBUFFERED setting
-	submitterTemplate, err = getSubmitterTemplate(ctx, rayJobInstanceWithoutTemplate, rayClusterInstance)
+	submitterTemplate, err = getSubmitterTemplate(rayJobInstanceWithoutTemplate, rayClusterInstance)
 	require.NoError(t, err)
 
 	envVar, found := utils.EnvVarByName(PythonUnbufferedEnvVarName, submitterTemplate.Spec.Containers[utils.RayContainerIndex].Env)
 	assert.True(t, found)
 	assert.Equal(t, "1", envVar.Value)
 
 	// Test 5: Check default RAY_DASHBOARD_ADDRESS env var
-	submitterTemplate, err = getSubmitterTemplate(ctx, rayJobInstanceWithTemplate, nil)
+	submitterTemplate, err = getSubmitterTemplate(rayJobInstanceWithTemplate, rayClusterInstance)
 	require.NoError(t, err)
 
 	envVar, found = utils.EnvVarByName(utils.RAY_DASHBOARD_ADDRESS, submitterTemplate.Spec.Containers[utils.RayContainerIndex].Env)