Support gang scheduling with Yunikorn

[jacobsalway]

Purpose

Add a scheduler implementation for YuniKorn to support task group annotations on the driver pod for gang scheduling and queue labels on both driver and executor pods.

Completes the first dot point of #2098. Will add docs for the Spark operator in a separate PR, but will add docs to YuniKorn as well.

Changes

• Add a new batch scheduler with the name yunikorn
• Calculate the resources required to schedule both driver and executor pods i.e. requests and executor counts, and use this to create task groups annotations on the driver pod to support gang scheduling
• Use the existing batchSchedulerOptions.queue to allow queue configuration. This is not a required field as Yunikorn supports placement rules that allow an app to be placed without explicitly specifying a queue

Testing

# clone my fork and checkout the branch
git clone git@github.com:jacobsalway/spark-operator.git
cd spark-operator && git checkout add-yunikorn-scheduler

# build the docker image
docker build . -t docker.io/kubeflow/spark-operator:add-yunikorn-scheduler

# create a kind cluster
kind create cluster --wait=1m

# load the operator image
kind load docker-image docker.io/kubeflow/spark-operator:add-yunikorn-scheduler

# install the operator with the forked build and the batch scheduler enabled
helm install spark-operator ./charts/spark-operator-chart \
  --create-namespace \
  --namespace spark-operator \
  --set image.tag=add-yunikorn-scheduler \
  --set controller.batchScheduler.enable=true

# install yunikorn
helm install yunikorn yunikorn \
  --repo https://apache.github.io/yunikorn-release \
  --create-namespace \
  --namespace yunikorn

# create the example application
kubectl apply -f examples/spark-pi-yunikorn.yaml

# observe that yunikorn will gang schedule the application
# task group pods prefixed with `tg-spark-...-spark-{driver/executor}-...` with requests
# that match the real pods will be created

recording.mp4

Change Category

Indicate the type of change by marking the applicable boxes:

• Bugfix (non-breaking change which fixes an issue)
• Feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that could affect existing functionality)
• [ ] Documentation update will do in a separate PR to avoid this one getting any larger

Checklist

Before submitting your PR, please review the following:

• I have conducted a self-review of my own code.
• [ ] I have updated documentation accordingly. will do in a separate PR to avoid this one getting any larger
• I have added tests that prove my changes are effective or that my feature works.
• Existing unit tests pass locally with my changes.

[jacobsalway]

Still working on this, just sorting out some issues with converting the java byte string to Kubernetes resource values. Have re-read through the upstream code and have a better approach in mind

[jacobsalway]

Stacked PR also available to allow the default batch scheduler to be set if not specified by the user jacobsalway#1

[ChenYi015]

Thanks for your contributions, I have left some comments.

[ChenYi015]

Since the mergeMaps func is an util function, maybe we can move it to pkg/utill/util.go.

[jacobsalway]

I've moved the function into scheduler.go and renamed it to mergeNodeSelector because I think that better describes what it does. Since it returns nil if the map is empty to make the marshalled result a bit nicer, I'm not sure if that would be expected behaviour for a global utils function

[ChenYi015]

The logic of calculating initial executors is not the same as spark core. Ref: https://github.com/apache/spark/blob/899fad4710bef174684deee64314ac483c16c494/core/src/main/scala/org/apache/spark/scheduler/cluster/SchedulerBackendUtils.scala#L23-L47.

[ChenYi015]

Maybe we can move this funcNumInitialExecutors to pkg/util/sparkapplication.go as this is an util function for SparkApplication.

[jacobsalway]

Ah nice, good catch. I think I'd only been looking at https://github.com/apache/spark/blob/899fad4710bef174684deee64314ac483c16c494/core/src/main/scala/org/apache/spark/util/Utils.scala#L2534-L2557 when I wrote this implementation.

Will fix and move to utils tomorrow

[jacobsalway]

Fixed to match the core implementation and moved to pkg/utils/sparkapplication.go

[ChenYi015]

We can add license comment in every new file, like

spark-operator/internal/controller/sparkapplication/controller.go

Lines 1 to 15 in 5972482

	/*
	Copyright 2024 The Kubeflow authors.
	Licensed under the Apache License, Version 2.0 (the "License");
	you may not use this file except in compliance with the License.
	You may obtain a copy of the License at
	https://www.apache.org/licenses/LICENSE-2.0
	Unless required by applicable law or agreed to in writing, software
	distributed under the License is distributed on an "AS IS" BASIS,
	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	See the License for the specific language governing permissions and
	limitations under the License.
	*/

.

[jacobsalway]

I'll add the license comment to the top of all my new files, but what do you think if in a separate PR we add a license comment check as part of CI, and maybe a Makefile target for checking and adding locally?

[ChenYi015]

I think it is a good point to add license check as part of CI. cc @andreyvelich @vara-bonthu @yuchaoran2011 .

[jacobsalway]

Created an issue for it #2139

[ChenYi015]

I see that in Volcano scheduler, the PodGroup names are different between SparkApplications. So I am wondering whether we need to set distinct task group names for different SparkApplications?

[jacobsalway]

For Yunikorn, a task group is a property of an application, so the task group names only need to be unique within the application they're a part of.

Application IDs must be unique in the same way that a Volcano PodGroup name has to be, however there's a nice integration in Yunikorn to use the spark-app-selector label for the application ID.

[yuchaoran2011]

/lgtm
Great to see that all the edge cases including min overhead memory are covered.

[google-oss-prow]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: ChenYi015, yuchaoran2011

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

• OWNERS [ChenYi015,yuchaoran2011]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment