Updating annotations of pods belonging Ray cluster in order to adopting Yunikorn Gang scheduling #5594

0yukali0 · 2024-07-26T13:38:12Z

Tracking issue

Related to #5575

Why are the changes needed?

Ability to classify pods to different gang scheduling group based on their role and worker group number.
Currently, Pods under ray cluster share same gang scheduling task group via podTemplate.
With same pod spec, pods share same gang scheduling group name when adopting yunikorn.

What changes were proposed in this pull request?

For enable gang scheduling with yunikorn, updating annotations pods belonging to ray cluster based on role and number of worker group.

How was this patch tested?

run make test_unit and make lint in flyteplugins.

Setup process

adopting yaml including scheduler name and parameters

tasks:
  task-plugins:
    enabled-plugins:
      - container
      - sidecar
      - K8S-ARRAY
      - ray
    default-for-task-types:
      - container: container
      - container_array: K8S-ARRAY
      - ray: ray
    
plugins:
  ray:
    batchScheduler:
      scheduler: yunikorn
      parameters: "placeholderTimeoutInSeconds=10 gangSchedulingStyle=Soft"

Set flyte dev env
Yunikorn installation

embedAdmissionController: false
enableSchedulerPlugin: false
enableWebService: false

helm install yunikorn yunikorn/yunikorn -f values.yaml -n yunikorn --create-namespace
4. Set flyte plugin ray
5.Submit ray example

import typing
from flytekit import ImageSpec, Resources, task, workflow

custom_image = ImageSpec(
        registry="localhost:30000",
    packages=["flytekitplugins-ray"],
    apt_packages=["wget"],
)

import ray
from flytekitplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

@ray.remote
def f(x):
    return x * x

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=1), WorkerNodeConfig(group_name="ray-group2", replicas=1)],
    runtime_env={"pip": ["numpy", "pandas"]},  # or runtime_env="./requirements.txt"
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=3600,
)

@task(
    task_config=ray_config,
    requests=Resources(mem="2Gi", cpu="2"),
    container_image=custom_image,
)
def ray_task(n: int) -> typing.List[int]:
    futures = [f.remote(i) for i in range(n)]
    return ray.get(futures)

@workflow
def ray_workflow(n: int) -> typing.List[int]:
    return ray_task(n=n)

if __name__ == "__main__":
    print(ray_workflow(n=10))

Screenshots

chments/assets/d3d7d765-598f-4fcc-ae40-86bb8d684681)

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Related PRs

Docs link

Yunikorn gang scheduling doc

welcome · 2024-07-26T13:38:15Z

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
Sign off your commits (Reference: DCO Guide).

codecov · 2024-07-26T22:35:51Z

Codecov Report

Attention: Patch coverage is 10.00000% with 144 lines in your changes missing coverage. Please review.

Project coverage is 34.44%. Comparing base (bed761c) to head (bd74f60).
Report is 68 commits behind head on master.

Files with missing lines	Patch %	Lines
...asks/plugins/k8s/batchscheduler/yunikorn/helper.go	0.00%	88 Missing ⚠️
flyteplugins/go/tasks/plugins/k8s/ray/ray.go	2.85%	34 Missing ⚠️
...o/tasks/plugins/k8s/batchscheduler/utils/helper.go	0.00%	12 Missing ⚠️
...er/pkg/controller/nodes/task/k8s/plugin_manager.go	0.00%	4 Missing and 2 partials ⚠️
...tasks/pluginmachinery/workqueue/mocks/processor.go	0.00%	2 Missing ⚠️
...o/tasks/plugins/k8s/batchscheduler/kueue/helper.go	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5594      +/-   ##
==========================================
- Coverage   36.17%   34.44%   -1.73%     
==========================================
  Files        1302     1143     -159     
  Lines      109492   102898    -6594     
==========================================
- Hits        39606    35446    -4160     
+ Misses      65748    63776    -1972     
+ Partials     4138     3676     -462

Flag	Coverage Δ
unittests-datacatalog	`51.37% <ø> (ø)`
unittests-flyteadmin	`55.60% <ø> (+0.27%)`	⬆️
unittests-flytecopilot	`12.17% <ø> (ø)`
unittests-flytectl	`?`
unittests-flyteidl	`7.17% <ø> (+0.08%)`	⬆️
unittests-flyteplugins	`52.90% <10.38%> (-0.42%)`	⬇️
unittests-flytepropeller	`42.01% <0.00%> (+0.27%)`	⬆️
unittests-flytestdlib	`55.29% <ø> (-0.07%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

flyteplugins/go/tasks/plugins/k8s/ray/batchscheduler/config/config.go

flyteplugins/go/tasks/plugins/k8s/ray/config.go

flyteplugins/go/tasks/plugins/k8s/ray/batchscheduler/config/config_test.go

flyteplugins/go/tasks/plugins/k8s/ray/batchscheduler/scheduler/yunikorn/yunikorn.go

flyteplugins/go/tasks/plugins/k8s/ray/batchscheduler/plugins.go

0yukali0 · 2024-08-15T12:58:21Z

Hi @pingsutw, thanks for review.
I am about to improve this approach in these days.

Signed-off-by: yuteng <[email protected]>

davidmirror-ops · 2024-08-27T18:52:31Z

@0yukali0 thank you, this is great. Let us know if you need help taking this to the finish line. This is important for several Flyte users

0yukali0 · 2024-09-01T10:12:01Z

Thanks @davidmirror-ops , i will.
Please feel free to mark any areas where i can improve this commit.
BTW, there's some exciting news regarding Spark.
The Spark operator has supported gang scheduling with Yunikorn since last week. Therefore, the future plugin for spark in Flyte can focus on generating Yunikorn application IDs and updating gang scheduling parameters when a new application is submitted.

The commit in spark operator

Signed-off-by: yuteng <[email protected]>

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/kubernetes/default.go

Signed-off-by: Kevin Su <[email protected]>

Signed-off-by: yuteng <[email protected]>

kumare3 · 2024-09-22T16:27:01Z

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn/rayhandler.go

+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+func MutateRayJob(parameters string, app *rayv1.RayJob) error {


Lets add some comments explaining what this does?

kumare3 · 2024-09-22T16:27:21Z

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/noop_scheduler.go

@@ -0,0 +1,16 @@
+package scheduler


Lets add some module level comments that will help in generated docs

kumare3 · 2024-09-22T16:28:13Z

flyteplugins/go/tasks/plugins/k8s/batchscheduler/plugins.go

+)
+
+type SchedulerManager interface {
+	// Mutate is responsible for mutating the object to be scheduled.


Lets explain that the Mutate method will mutate the object in place.

kumare3 · 2024-09-22T16:32:42Z

flyteplugins/go/tasks/plugins/k8s/batchscheduler/plugins.go

+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+type SchedulerManager interface {


How about docs that say,something to this effect)

Suggested change

type SchedulerManager interface {

// SchedulerManager is an interface that allows plugging in custom schedulers

// such as Yunikorn, Kueue, or any other scheduler that manages pod scheduling

// with advanced features like gang scheduling, preemption, and priority.

type SchedulerManager interface {

kumare3 · 2024-09-22T16:34:31Z

flyteplugins/go/tasks/plugins/k8s/batchscheduler/plugins.go

@@ -0,0 +1,23 @@
+package batchscheduler


Suggested change

package batchscheduler

// Flyte often deals with complex, distributed workloads that may require advanced scheduling features. Native Kubernetes scheduling may not suffice for scenarios that involve job-level scheduling policies like gang scheduling or job preemption. The SchedulerManager interface provides a flexible mechanism for integrating such schedulers into Flyte’s plugin system by allowing mutation of the scheduling-related properties of Kubernetes objects.

package batchscheduler

also how about we call this scheduler?

kumare3 · 2024-09-22T16:36:24Z

flyteplugins/go/tasks/plugins/k8s/batchscheduler/plugins.go

+type SchedulerManager interface {
+	// Mutate is responsible for mutating the object to be scheduled.
+	// It will add the necessary annotations, labels, etc. to the object.
+	Mutate(ctx context.Context, object client.Object) error


Suggested change

Mutate(ctx context.Context, object client.Object) error

// Mutate allows the custom scheduler to modify the given Kubernetes object (e.g., a Pod or

// a custom resource representing a job) before it is submitted to the Kubernetes API.

// This may include adding annotations, labels, or modifying fields related to scheduling.

// The object is typically a Kubernetes Pod or a custom resource.

// It returns an error if the mutation fails.

Mutate(ctx context.Context, object client.Object) error

kumare3 · 2024-09-22T16:38:26Z

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn/yunikorn.go

+
+func (y *YunikornSchedulerManager) Mutate(ctx context.Context, object client.Object) error {
+	switch object.(type) {
+	case *rayv1.RayJob:


this is weird isnt it?
WE have a common interface, but have a switch for each?
In this case, we should allow this overriding to exist next to ray job plugin and not with yunikorn. this will break when things update pretty quickly.

Why have a common interface?

kumare3 · 2024-09-22T16:39:48Z

flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager.go

@@ -203,6 +205,12 @@ func (e *PluginManager) launchResource(ctx context.Context, tCtx pluginsCore.Tas

 	key := backoff.ComposeResourceKey(o)

+	err = e.schedulerMgr.Mutate(ctx, o)


I see this is why... So here is my suggestion, how about we handle the scheduler related changes directly to the plugin itself.

For example
MutateForGangScheduling

Also how about preemption etc if we add that in the future?

kumare3 · 2024-09-22T16:41:29Z

I love this PR, @0yukali0 but there is probably a better way to introduce this interface. I do not love the idea of adding cases in a separate place from the actual ray job plugin. My preference is to keep all that code under one module, also avoids imports

Also can you share an example of how preemption and priority works with yunikorn?

fg91 · 2024-10-04T17:09:34Z

Adding support for a gang scheduler is a great effort, thank you 🙏

I personally am not a huge fan of the folder/module structure proposed in this PR because for each k8s plugin, we'd need to create another module under scheduler/yunikorn.

Is my understanding correct that we "only" have to add yunikorn annotations to the respective pod template specs?

If yes, would it be possible to have one central helper function defined that can be imported in the plugins and, if configured in the respective plugin's config, is called in the BuildResource function to append the annotations to the pod template spec?

Signed-off-by: yuteng <[email protected]>

0yukali0 force-pushed the gangScheduling branch from 39cf088 to 2144c34 Compare July 26, 2024 14:06

0yukali0 marked this pull request as draft July 27, 2024 05:43

0yukali0 marked this pull request as ready for review July 29, 2024 21:48

0yukali0 force-pushed the gangScheduling branch from a202a11 to a805af9 Compare July 29, 2024 21:55

0yukali0 marked this pull request as draft July 30, 2024 14:04

0yukali0 marked this pull request as ready for review July 30, 2024 15:28

pingsutw reviewed Aug 13, 2024

View reviewed changes

0yukali0 marked this pull request as draft August 15, 2024 12:57

0yukali0 added 16 commits August 15, 2024 21:29

gangscheduling annotations

516978e

Signed-off-by: yuteng <[email protected]>

test

0e00f61

Signed-off-by: yuteng <[email protected]>

test

9309006

Signed-off-by: yuteng <[email protected]>

return error when parse fail

dc0e84e

Signed-off-by: yuteng <[email protected]>

update when no assign

20f3646

Signed-off-by: yuteng <[email protected]>

fix bug that worker groups share same task group name

f86948e

Signed-off-by: yuteng <[email protected]>

make test pass

75ccabc

Signed-off-by: yuteng <[email protected]>

refactor

e3f3631

Signed-off-by: yuteng <[email protected]>

unit tests

ca4a4eb

Signed-off-by: yuteng <[email protected]>

lint:scheduler config

d8ba6bc

Signed-off-by: yuteng <[email protected]>

gci format

b4e6467

Signed-off-by: yuteng <[email protected]>

Update go.sum to origin one

3fbe18d

Signed-off-by: yuteng <[email protected]>

refactor

20c91fe

Signed-off-by: yuteng <[email protected]>

update test in ray

c14d265

Signed-off-by: yuteng <[email protected]>

pflag

b047ecd

Signed-off-by: yuteng <[email protected]>

codespell

43b88b7

Signed-off-by: yuteng <[email protected]>

0yukali0 force-pushed the gangScheduling branch from 1bbd46b to 43b88b7 Compare August 15, 2024 13:29

0yukali0 added 2 commits August 17, 2024 12:34

scheduler config description

9e73735

Signed-off-by: yuteng <[email protected]>

use empty scheduler config

58b9a0b

Signed-off-by: yuteng <[email protected]>

0yukali0 added 5 commits August 19, 2024 21:57

move to k8s repo

9c00cbb

Signed-off-by: yuteng <[email protected]>

move to k8s repo

f5474ef

Signed-off-by: yuteng <[email protected]>

refactoring

ecefa77

Signed-off-by: yuteng <[email protected]>

refactor

189f618

Signed-off-by: yuteng <[email protected]>

fix wrong format

3f83a1a

Signed-off-by: yuteng <[email protected]>

0yukali0 added 2 commits September 1, 2024 22:07

fix ray sidecar problem and ready to refactor based on this commit

9cb0a5e

Signed-off-by: yuteng <[email protected]>

fix resource computation

ad68f48

Signed-off-by: yuteng <[email protected]>

pingsutw reviewed Sep 6, 2024

View reviewed changes

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/kubernetes/default.go Outdated Show resolved Hide resolved

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/kubernetes/default.go Outdated Show resolved Hide resolved

pingsutw and others added 4 commits September 8, 2024 02:09

kevin refactor

64b8655

Signed-off-by: Kevin Su <[email protected]>

kevin refactor

a4cca9d

Signed-off-by: Kevin Su <[email protected]>

nit

4387f64

Signed-off-by: Kevin Su <[email protected]>

build completed tg

7515dd8

Signed-off-by: yuteng <[email protected]>

kumare3 reviewed Sep 22, 2024

View reviewed changes

refactoring proposal

bd74f60

Signed-off-by: yuteng <[email protected]>

0yukali0 mentioned this pull request Nov 13, 2024

[Newbie] Supporting Yunikorn and Kueue #5915

Open

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Updating annotations of pods belonging Ray cluster in order to adopting Yunikorn Gang scheduling #5594

Updating annotations of pods belonging Ray cluster in order to adopting Yunikorn Gang scheduling #5594

0yukali0 commented Jul 26, 2024 •

edited

Loading

welcome bot commented Jul 26, 2024

codecov bot commented Jul 26, 2024 •

edited

Loading

0yukali0 commented Aug 15, 2024

davidmirror-ops commented Aug 27, 2024

0yukali0 commented Sep 1, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 Sep 22, 2024

kumare3 commented Sep 22, 2024

fg91 commented Oct 4, 2024

-type SchedulerManager interface {
+// SchedulerManager is an interface that allows plugging in custom schedulers
+// such as Yunikorn, Kueue, or any other scheduler that manages pod scheduling
+// with advanced features like gang scheduling, preemption, and priority.
+type SchedulerManager interface {

	package batchscheduler
	// Flyte often deals with complex, distributed workloads that may require advanced scheduling features. Native Kubernetes scheduling may not suffice for scenarios that involve job-level scheduling policies like gang scheduling or job preemption. The SchedulerManager interface provides a flexible mechanism for integrating such schedulers into Flyte’s plugin system by allowing mutation of the scheduling-related properties of Kubernetes objects.
	package batchscheduler

-	Mutate(ctx context.Context, object client.Object) error
+ // Mutate allows the custom scheduler to modify the given Kubernetes object (e.g., a Pod or
+    // a custom resource representing a job) before it is submitted to the Kubernetes API.
+    // This may include adding annotations, labels, or modifying fields related to scheduling.
+    // The object is typically a Kubernetes Pod or a custom resource.
+    // It returns an error if the mutation fails.
+	Mutate(ctx context.Context, object client.Object) error

		@@ -203,6 +205,12 @@ func (e *PluginManager) launchResource(ctx context.Context, tCtx pluginsCore.Tas

		key := backoff.ComposeResourceKey(o)

		err = e.schedulerMgr.Mutate(ctx, o)

Updating annotations of pods belonging Ray cluster in order to adopting Yunikorn Gang scheduling #5594

Are you sure you want to change the base?

Updating annotations of pods belonging Ray cluster in order to adopting Yunikorn Gang scheduling #5594

Conversation

0yukali0 commented Jul 26, 2024 • edited Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

welcome bot commented Jul 26, 2024

codecov bot commented Jul 26, 2024 • edited Loading

Codecov Report

0yukali0 commented Aug 15, 2024

davidmirror-ops commented Aug 27, 2024

0yukali0 commented Sep 1, 2024

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

kumare3 commented Sep 22, 2024

fg91 commented Oct 4, 2024

0yukali0 commented Jul 26, 2024 •

edited

Loading

codecov bot commented Jul 26, 2024 •

edited

Loading