Support gang scheduling with Apache YuniKorn #2396

yangwwei · 2024-09-21T19:59:34Z

Add Gang Scheduling support with Apache YuniKorn.

This follows the doc from: https://yunikorn.apache.org/docs/next/user_guide/gang_scheduling/.

What's implemented in this PR

When yunikron is enabled, and RayCluster CR has label "ray.io/gang-scheduling-enabled" defined, this means the Gang scheduling is enabled for this RayCluster. The operator will add the following labels to pods, these are the task groups info attached to each pod, yunikorn.apache.org/task-groups:

{
 "name": "headgroup",
 "minMember": 1,
 "minResource": {
   "cpu": "100m",
   "memory": "50M"
  },
  "nodeSelector": {},
  "tolerations": [],
  "affinity": {},
  "topologySpreadConstraints": []
}, 
{
 "name": "worker-group",
 "minMember": 2,
 "minResource": {
   "cpu": "100m",
   "memory": "50M"
  },
  "nodeSelector": {},
  "tolerations": [],
  "affinity": {},
  "topologySpreadConstraints": []
},

and then the head pod in the headgroup will be assigned with one label:

yunikorn.apache.org/task-group-name=headgroup

and the worker pod in the worker-group will be assigned with one more label:

yunikorn.apache.org/task-group-name=worker-group

underneath, yunikorn will be able to differenciate the worker group pods from head group pods, and then count their minimal resources and num of members before scheduling them. This is to ensure yunikorn can reserve the slots for all these minimal number of pods before actually scheduling them.

To use this feature

Prerequisite:

Start Ray operator with "--batch-scheduler=yunikorn"
Install yunikorn https://yunikorn.apache.org/docs/#install

Steps

Set the following labels in the Ray cluster object

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  labels:
    ray.io/gang-scheduling-enabled: "true"
    yunikorn.apache.org/app-id: my-ray-cluster-0001
    yunikorn.apache.org/queue: root.default
    ...

Create the Ray cluster on Kubernetes
The Ray cluster will be scheduled in root.default queue with gang scheduling enforced.

I will put a doc PR to add a tutorial how to run this with Apache YuniKorn. This PR includes the implementation and UT. The logic is very simple, when the scheduler finds the gang scheduling is enabled, it will popluate the taskGroups info to the pod annotations. The task group format is:

[ 
  {
    "name":"headgroup",
    "minMember":1,
    "minResource": {"cpu":"1"},
    ...
  }, 
 { 
  "name":"workergroup",
  "minMember":1,
  "minResource":{"cpu":"1"},
  ...
 },
]

and head pod will belong to group "headgroup", the each worker group will have its own group.

yangwwei · 2024-09-23T04:26:51Z

@kevin85421 pls help to review, thanks!

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

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

MortalHappiness · 2024-09-26T10:31:56Z

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

 	y.populatePodLabels(app, pod, RayClusterApplicationIDLabelName, YuniKornPodApplicationIDLabelName)
 	y.populatePodLabels(app, pod, RayClusterQueueLabelName, YuniKornPodQueueLabelName)


I don't know why the label names need to be translated from RayClusterApplicationIDLabelName to YuniKornPodApplicationIDLabelName and from RayClusterQueueLabelName to YuniKornPodQueueLabelName.

Does Yunikorn read the labels on resources other than Pods? If not, why not set YuniKornPodApplicationIDLabelName to yunikorn.apache.org/app-id and YuniKornPodQueueLabelName to yunikorn.apache.org/queue and set them directly on RayCluster, and then apply the same labels to all the Pods that controlled by RayCluster?

the caveat is, yunikorn depends on labels to identify apps, but they are currently using very simple format:

YuniKornPodApplicationIDLabelName: applicationId

YuniKornPodQueueLabelName: queue

but in Ray, I think it makes sense to add the domain name to the label, to indicate these lables are set for the scheduler.

yunikorn.apache.org/application-id

yunikorn.apache.org/queue-name

the code logic is very simple, just translate these labels and set them on pods, without introducing any inconsistency in Ray cluster spec. Hope that makes sense.

But I viewed the doc here

https://yunikorn.apache.org/docs/next/user_guide/labels_and_annotations_in_yunikorn

there are two labels yunikorn.apache.org/app-id and yunikorn.apache.org/queue. And applicationId and queue labels are legacy now.

Oh, good catch. Seems like the community is moving to a domain scoped label approach recently. Let me take a look, and do some tests, will update the PR once confims these labels works. Thanks!

hi @MortalHappiness I just found these labels are only supported in v1.6 which was released just a few weeks ago. For better compatibility with some older version of yunikorn deployment, I think it's better to continue to support the legacy label names. However, I have modified RayClusterApplicationIDLabelName and RayClusterQueueLabelName in the package to be the same as yunikorn upstream. We can remove this translation layer and directly populate the labels to pods after a few releases. Does this sound good?

SGTM. Thanks!

How about add a TODO comment here to let us know we need to remove this in the future release?

Sure, updated.

@yangwwei Would you mind briefly adding comments for the context you mentioned in https://github.com/ray-project/kuberay/pull/2396/files#r1777710831?

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

kevin85421 · 2024-09-27T01:23:05Z

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

 	y.populatePodLabels(app, pod, RayClusterApplicationIDLabelName, YuniKornPodApplicationIDLabelName)
 	y.populatePodLabels(app, pod, RayClusterQueueLabelName, YuniKornPodQueueLabelName)


@yangwwei Would you mind briefly adding comments for the context you mentioned in https://github.com/ray-project/kuberay/pull/2396/files#r1777710831?

kevin85421 · 2024-09-27T01:42:05Z

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

+	}
+	pod.Annotations[YuniKornTaskGroupsAnnotationName] = taskGroupsAnnotationValue
+
+	y.log.Info("Gang Scheduling enabled for RayCluster",


I am not familiar with Yunikorn. Can we infer that gang scheduling is enabled if the annotation yunikorn.apache.org/task-groups is set? If not, should we move this log to:

if y.isGangSchedulingEnabled(app) { ... }

instead?

I may not fully understand this. Basically "yunikorn.apache.org/task-groups" only needs to be set to Ray pods by the operator, it is internal metadata. We only use one flag "ray.io/gang-scheduling-enabled" in Ray to tell if we want to enable Gang scheduling or not for this Ray cluster. Only when enabled, the operator populates the task groups info to the pods, then yunikorn will honor and apply gang scheduling.

I found it odd that we print the log saying "Gang Scheduling enabled for RayCluster" in populateTaskGroupsAnnotationToPod, a function that may also be used by other scheduling algorithms, which seems a bit strange.

That's why I’m asking if we should move the log to:

Copy code if y.isGangSchedulingEnabled(app) { ... }

instead?

I see what you mean. Basically populateTaskGroupsAnnotationToPod will ONLY be used for gang scheduling case, so it is safe to print the log inside of function. Also this function doesn't return any error, it only logs where or not gang scheduling is enabled (by checking if required task groups added to the pod).

Got it. Thanks!

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

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

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

kevin85421

Summarize this PR:

Every head Pod or worker Pod has the same applicationId and queue label values if they belong to the same RayCluster CR.
If gang scheduling is enabled, every head Pod or worker Pod has the same values of task-groups annotation if they belong to the same RayCluster CR. Pods belong to different Pod groups (head group + worker groups) have different values of task-group annotation.

Is this correct? In addition, what's the behavior if gang scheduling is not enabled?

MortalHappiness · 2024-09-27T03:34:33Z

I manually run it and it seems working.

Config:

  queues.yaml: |
    partitions:
      - name: default
        queues:
          - name: root
            queues:
              - name: test
                submitacl: "*"
                parent: false
                resources:
                  guaranteed:
                    memory: 6G
                    vcore: 4
                  max:
                    memory: 6G
                    vcore: 4

RayCluster:

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: test-yunikorn-0
  labels:
    ray.io/gang-scheduling-enabled: "true"
    yunikorn.apache.org/app-id: test-yunikorn-0
    yunikorn.apache.org/queue: root.test
spec:
  rayVersion: "2.34.0"
  headGroupSpec:
    rayStartParams: {}
    template:
      spec:
        containers:
          - name: ray-head
            image: rayproject/ray:2.34.0
            resources:
              limits:
                cpu: "1"
                memory: "2Gi"
              requests:
                cpu: "1"
                memory: "2Gi"
  workerGroupSpecs:
    - groupName: worker
      rayStartParams: {}
      replicas: 2
      minReplicas: 2
      maxReplicas: 2
      template:
        spec:
          containers:
            - name: ray-head
              image: rayproject/ray:2.34.0
              resources:
                limits:
                  cpu: "1"
                  memory: "1Gi"
                requests:
                  cpu: "1"
                  memory: "1Gi"

Apply the RayCluster yaml above, and then replace test-yunikorn-0 with test-yunikorn-1 and then apply again.

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

yangwwei · 2024-09-27T04:17:29Z

Summarize this PR:

Every head Pod or worker Pod has the same applicationId and queue label values if they belong to the same RayCluster CR.

If gang scheduling is enabled, every head Pod or worker Pod has the same values of task-groups annotation if they belong to the same RayCluster CR. Pods belong to different Pod groups (head group + worker groups) have different values of task-group annotation.

Is this correct? In addition, what's the behavior if gang scheduling is not enabled?

This is correct. Let me update the description of this PR too.
If gang scheduling is not enabled (i.e ray.io/gang-scheduling-enabled not set), then all task groups info will not be added to Ray pods, and the Ray cluster will still be scheduled, just like normal application, pods will get scheduled as long as there is resources in the queue and the cluster.

MortalHappiness

LGTM. I think this PR can be merged if @kevin85421 's comments are resolved.

cc @andrewsykim

kevin85421

the last comment: #2396 (comment). Others LGTM

yangwwei added 5 commits September 21, 2024 12:49

Support gang scheduling with Apache YuniKorn

b561e0a

Support gang scheduling with Apache YuniKorn

d72d73a

Support gang scheduling with Apache YuniKorn

708b86f

Support gang scheduling with Apache YuniKorn

1ceada8

Support gang scheduling with Apache YuniKorn

2e962f9

andrewsykim reviewed Sep 23, 2024

View reviewed changes

yangwwei added 3 commits September 23, 2024 13:25

address review comments

4ee2405

address review comments

f9b68db

address review comments

08cc7b3

MortalHappiness requested changes Sep 24, 2024

View reviewed changes

ray-operator/controllers/ray/batchscheduler/yunikorn/yunikorn_scheduler.go Outdated Show resolved Hide resolved

ray-operator/controllers/ray/batchscheduler/yunikorn/yunikorn_scheduler_test.go Outdated Show resolved Hide resolved

address review comments

6a06417

MortalHappiness reviewed Sep 26, 2024

View reviewed changes

yangwwei added 4 commits September 26, 2024 13:55

address review comments

6a3d044

Merge remote-tracking branch 'upstream/master' into gs

aae020e

address review comments

10119ed

address review comments

3ea3b64

kevin85421 assigned kevin85421, andrewsykim and MortalHappiness Sep 27, 2024

kevin85421 reviewed Sep 27, 2024

View reviewed changes

kevin85421 unassigned andrewsykim Sep 27, 2024

kevin85421 reviewed Sep 27, 2024

View reviewed changes

MortalHappiness reviewed Sep 27, 2024

View reviewed changes

ray-operator/controllers/ray/batchscheduler/yunikorn/yunikorn_scheduler_test.go Outdated Show resolved Hide resolved

yangwwei added 2 commits September 26, 2024 21:28

address review comments

31ade1c

address review comments

d932703

MortalHappiness approved these changes Sep 27, 2024

View reviewed changes

kevin85421 approved these changes Sep 27, 2024

View reviewed changes

kevin85421 merged commit df5577f into ray-project:master Sep 27, 2024
27 checks passed

MortalHappiness mentioned this pull request Sep 28, 2024

[KubeRay] Add YuniKorn Gang scheduling doc ray-project/ray#47848

Closed

kevin85421 mentioned this pull request Sep 28, 2024

[Docs][KubeRay] Add YuniKorn Gang scheduling doc ray-project/ray#47850

Merged

8 tasks

ryankert01 mentioned this pull request Oct 16, 2024

[YUNIKORN-2894] Update KubeRay operator documentation for YuniKorn integration apache/yunikorn-site#489

Closed

10 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support gang scheduling with Apache YuniKorn #2396

Support gang scheduling with Apache YuniKorn #2396

yangwwei commented Sep 21, 2024 •

edited

Loading

yangwwei commented Sep 23, 2024

MortalHappiness Sep 26, 2024

yangwwei Sep 26, 2024

MortalHappiness Sep 26, 2024 •

edited

Loading

yangwwei Sep 26, 2024

yangwwei Sep 27, 2024

MortalHappiness Sep 27, 2024

MortalHappiness Sep 27, 2024

yangwwei Sep 27, 2024

kevin85421 Sep 27, 2024

yangwwei Sep 27, 2024

kevin85421 Sep 27, 2024

kevin85421 Sep 27, 2024

yangwwei Sep 27, 2024

kevin85421 Sep 27, 2024

yangwwei Sep 27, 2024

kevin85421 Sep 27, 2024

kevin85421 left a comment

MortalHappiness commented Sep 27, 2024 •

edited

Loading

yangwwei commented Sep 27, 2024

MortalHappiness left a comment •

edited

Loading

kevin85421 left a comment

		y.populatePodLabels(app, pod, RayClusterApplicationIDLabelName, YuniKornPodApplicationIDLabelName)
		y.populatePodLabels(app, pod, RayClusterQueueLabelName, YuniKornPodQueueLabelName)

Support gang scheduling with Apache YuniKorn #2396

Support gang scheduling with Apache YuniKorn #2396

Conversation

yangwwei commented Sep 21, 2024 • edited Loading

What's implemented in this PR

Prerequisite:

Steps

yangwwei commented Sep 23, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

MortalHappiness Sep 26, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kevin85421 left a comment

Choose a reason for hiding this comment

MortalHappiness commented Sep 27, 2024 • edited Loading

yangwwei commented Sep 27, 2024

MortalHappiness left a comment • edited Loading

Choose a reason for hiding this comment

kevin85421 left a comment

Choose a reason for hiding this comment

yangwwei commented Sep 21, 2024 •

edited

Loading

MortalHappiness Sep 26, 2024 •

edited

Loading

MortalHappiness commented Sep 27, 2024 •

edited

Loading

MortalHappiness left a comment •

edited

Loading