[RayJob] Yunikorn Integration

[owenowenisme]

Why are these changes needed?

This PR integrates RayJob with Apache Yunikorn as a batch scheduler, enabling RayJob to leverage Yunikorn’s features such as gang-scheduling. This allows both the submitter pod and RayCluster to be included in gang-scheduling.

Special thanks to the original contribution from @troychiu! 🙏 (#3379)

In this PR:

• Added batchSchedulerManager to RayJobReconcileOption in main.go
• Modified the RayJob controller to add essential labels and annotations to both RayCluster and the submitter pod
• Added ray-job.yunikorn-scheduler.yaml as a demonstration script for gang-scheduling with RayJob
• Added unit tests for new functions in yunikorn_scheduler.go

Validation

1. Set the Yunikorn queue to 6Gi memory and 4 cores, same as for RayCluster.
   

2. Deploy the KubeRay operator with the batch-scheduler set to Yunikorn:

   ./bin/manager -leader-election-namespace default -use-kubernetes-proxy -batch-scheduler=yunikorn
   

3. Deploy a sample RayJob:

   k apply -f config/samples/ray-job.yunikorn-scheduler.yaml
   

   Each RayJob consists of:

   • 1 Head pod (CPU: 1, Memory: 2Gi)
   • 1 Worker pod (CPU: 1, Memory: 2Gi)
   • 1 Submitter pod (CPU: 500m, Memory: 200Mi)

   Total: 2.5 CPU and 4.2Gi memory. Therefore, our queue cannot fit two RayJobs at the same time.

4. Deploy another RayJob using the same script but change the RayJob name and app ID from job-0 to job-1.
   We can see that the resources for RayJob1 (including the submitter pod from job-1) are being held by Yunikorn.

Note that even though the submitter pod of ‎job-1 has sufficient resources for scheduling, gang-scheduling requires that all pods be scheduled together.

5. Delete RayJob-0.
   After deleting RayJob-0, the resources should be sufficient for RayJob-1, so RayJob-1 is now up and running.
   

Testing compatibility with RayJob HTTPMode

To make sure RayJob Yunikorn works well with HTTPMode submission, I did the following steps:

1. Use the same sample yaml but set the submissionMode to HTTPMode

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-yunikorn-0
  labels:
    ray.io/gang-scheduling-enabled: "true"
    yunikorn.apache.org/app-id: rayjob-yunikorn-0
    yunikorn.apache.org/queue: root.test
spec:
  # submissionMode specifies how RayJob submits the Ray job to the RayCluster.
  # The default value is "K8sJobMode", meaning RayJob will submit the Ray job via a submitter Kubernetes Job.
  # The alternative value is "HTTPMode", indicating that KubeRay will submit the Ray job by sending an HTTP request to the RayCluster.
  submissionMode: "HTTPMode"
  entrypoint: python /home/ray/samples/sample_code.py

2. Check the task group annotation, the submitter group shouldn't exist, since we use HTTP request instead of submitter pod to submit job.

yunikorn.apache.org/task-groups: '[{"minResource":{"cpu":"1","memory":"2Gi"},"name":"headgroup","minMember":1},{"minResource":{"cpu":"1","memory":"2Gi"},"name":"small-group","minMember":1}]'

3. Check the job result and log
   I enter the shell of head pod and query the job using ray job list and ray job logs JobID, it should be successful with the correct log.

Related issue number

Closes #3284

Checks

• I've made sure the tests are passing.
• Testing Strategy
  • Unit tests
  • Manual tests
  • This PR is not tested :(

[win5923]

Need to set yunikorn.apache.org/app-id and RayJob.Name to the same name, so that when renaming, the app-id is updated together. This makes it easier for users to understand when a newly created RayJob is stuck in the Accepted state but not running yet.

Ref: https://docs.ray.io/en/latest/cluster/kubernetes/k8s-ecosystem/yunikorn.html#step-4-use-apache-yunikorn-for-gang-scheduling

[owenowenisme]

Thanks, just updated.

[win5923]

Nit: GetScheduler always returns nil. For consistency, we can change it to return an error like in other places.

[owenowenisme]

Nice advice, we could open a follow-up PR after this one, since RayCluster controller need to be modified as well.

[win5923]

Overall LGTM, Thanks for your hard work! I’ll take another look at the test part.

cc @troychiu for review

[Future-Outlier]

it works on my laptop

[Future-Outlier]

it looks good to me, thank you

[win5923]

We should also check whether RayJobSpec.RayClusterSpec is nil before calling newTaskGroupsFromRayJobSpec, otherwise it can cause a nil pointer panic when RayJob uses an existing cluster with gang-scheduling enabled.

Alternatively, we could add this check in rayjob_controller.go, just prevent user use both batch scheduler and clusterSelector.

example like:

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-use-existing-raycluster
  labels:
    ray.io/gang-scheduling-enabled: "true"
    yunikorn.apache.org/app-id: test-yunikorn-0
    yunikorn.apache.org/queue: root.test
spec:
  entrypoint: python -c "import ray; ray.init(); print(ray.cluster_resources())"
  clusterSelector:
    ray.io/cluster: raycluster-kuberay

[owenowenisme]

Thanks!
But why would the user enable gang-scheduling with clusterSelector?
The RayJob will only have one submitter pod when using clusterSelector which will never need gang-scheduling IIUC?

[win5923]

Normally, this wouldn’t happen, but we want to prevent users from actually setting it this way, because with the current logic it would cause a panic.

[Future-Outlier]

I think we should avoid using gang scheduling + cluster selector, maybe we have to add this to validation rule in rayjob.

[Future-Outlier]

this can be a follow-up pr

[Future-Outlier]

cc @owenowenisme for solving the resolve conflict

[Future-Outlier]

thank you for solving the merge conflict

[win5923]

Do we need to keep these comments?

[owenowenisme]

Good catch, added them back. Thanks!

// propagateTaskGroupsAnnotation is a helper function that propagates the task groups annotation to the child
// if the parent has the task groups annotation, it will be copied to the child
// if the parent doesn't have the task groups annotation, a new one will be created
// TODO: remove the legacy labels, i.e "applicationId" and "queue", directly populate labels
// RayApplicationIDLabelName and RayApplicationQueueLabelName to pod labels.
// Currently we use this function to translate labels "yunikorn.apache.org/app-id" and "yunikorn.apache.org/queue"
// to legacy labels "applicationId" and "queue", this is for the better compatibilities to support older yunikorn
// versions.

[troychiu]

This logic looks like from

kuberay/ray-operator/controllers/ray/rayjob_controller.go

Line 606 in 7cc7c2d

submitterTemplate = common.GetDefaultSubmitterTemplate(&rayClusterInstance.Spec)

Is it possible to use the submitter template passed in so that we don't have duplicate logic?

[Future-Outlier]

How about create a new public function?

// old one
func GetDefaultSubmitterTemplate(rayClusterSpec *rayv1.RayClusterSpec) corev1.PodTemplateSpec {
	return corev1.PodTemplateSpec{
		Spec: corev1.PodSpec{
			Containers: []corev1.Container{
				GetDefaultSubmitterContainer(rayClusterSpec),
			},
			RestartPolicy: corev1.RestartPolicyNever,
		},
	}
}

// new one
func GetDefaultSubmitterTemplate(rayClusterSpec *rayv1.RayClusterSpec) corev1.PodTemplateSpec {
	return corev1.PodTemplateSpec{
		Spec: GetSubmitterPodSpec(rayClusterSpec)
	}
}

func GetSubmitterPodSpec(rayClusterSpec *rayv1.RayClusterSpec) corev1.PodSpec {
	return corev1.PodSpec{
				Containers: []corev1.Container{
					GetDefaultSubmitterContainer(rayClusterSpec),
				},
				RestartPolicy: corev1.RestartPolicyNever,
			}
}

In this case, we can call GetSubmitterPodSpec directly.

[Future-Outlier]

cc @troychiu for the final review!

[troychiu]

LGTM Thank you! cc @rueian

[Future-Outlier]

I tested this again, this can be merged IMO.