Skip to content

Commit 44de5bb

Browse files
dmitshdmitsh
authored
add support for kai-scheduler (#120)
Signed-off-by: Dmitry Shmulevich <[email protected]>
1 parent 19a0808 commit 44de5bb

File tree

7 files changed

+208
-3
lines changed

7 files changed

+208
-3
lines changed

‎docs/examples/kai/kai.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
## Example of running `KAI` with `knavigator`
2+
3+
### Running workflows with `MPI job`
4+
5+
Install [KAI scheduler](https://github.com/NVIDIA/KAI-Scheduler/blob/main/README.md) in your cluster.
6+
7+
Run an MPI job:
8+
```shell
9+
./bin/knavigator -workflow resources/workflows/kai/test-mpijob.yaml
10+
```

‎docs/getting_started.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,3 +27,4 @@ We have tested several of these and offer templates and workflows to support the
2727
* [Kueue](./examples/kueue/kueue.md)
2828
* [YuniKorn](./examples/yunikorn/yunikorn.md)
2929
* [Run:ai](./examples/runai/runai.md)
30+
* [Kai](./examples/kai/kai.md)

‎resources/templates/kai/mpijob.yaml‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: kubeflow.org/v2beta1
16+
kind: MPIJob
17+
metadata:
18+
name: "{{._NAME_}}"
19+
namespace: "{{.namespace}}"
20+
labels:
21+
runai/queue: "{{.queue}}"
22+
spec:
23+
slotsPerWorker: 1
24+
runPolicy:
25+
cleanPodPolicy: Running
26+
mpiReplicaSpecs:
27+
Launcher:
28+
replicas: 1
29+
template:
30+
metadata:
31+
annotations:
32+
pod-complete.stage.kwok.x-k8s.io/delay: {{.ttl}}
33+
pod-complete.stage.kwok.x-k8s.io/jitter-delay: {{.ttl}}
34+
spec:
35+
schedulerName: kai-scheduler
36+
containers:
37+
- image: {{.image}}
38+
name: mpi-launcher
39+
resources:
40+
limits:
41+
cpu: "{{.cpu}}"
42+
memory: {{.memory}}
43+
nvidia.com/gpu: "{{.gpu}}"
44+
Worker:
45+
replicas: {{.workers}}
46+
template:
47+
metadata:
48+
annotations:
49+
pod-complete.stage.kwok.x-k8s.io/delay: {{.ttl}}
50+
pod-complete.stage.kwok.x-k8s.io/jitter-delay: {{.ttl}}
51+
labels:
52+
app: {{._NAME_}}
53+
spec:
54+
schedulerName: kai-scheduler
55+
containers:
56+
- image: {{.image}}
57+
name: mpi-worker
58+
resources:
59+
limits:
60+
cpu: "{{.cpu}}"
61+
memory: {{.memory}}
62+
nvidia.com/gpu: "{{.gpu}}"

‎resources/templates/kai/queue.yaml‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: scheduling.run.ai/v2
16+
kind: Queue
17+
metadata:
18+
name: "{{.name}}"
19+
spec:
20+
{{- if .parentQueue }}
21+
parentQueue: "{{.parentQueue}}"
22+
{{- end }}
23+
resources:
24+
cpu:
25+
quota: -1
26+
limit: -1
27+
overQuotaWeight: 1
28+
gpu:
29+
quota: -1
30+
limit: -1
31+
overQuotaWeight: 1
32+
memory:
33+
quota: -1
34+
limit: -1
35+
overQuotaWeight: 1

‎resources/workflows/kai/test-mpijob.yaml‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
name: test-kai-mpijob
16+
description: register, deploy and configure run:ai custom resources
17+
tasks:
18+
- id: register-queue
19+
type: RegisterObj
20+
params:
21+
template: "resources/templates/kai/queue.yaml"
22+
- id: register-mpijob
23+
type: RegisterObj
24+
params:
25+
template: "resources/templates/kai/mpijob.yaml"
26+
nameFormat: "mpijob{{._ENUM_}}"
27+
podNameFormat: "{{._NAME_}}-(launcher-[a-z0-9]+|worker-[0-9]+)"
28+
podCount: "{{.workers}} + 1"
29+
- id: configure
30+
type: Configure
31+
params:
32+
nodes:
33+
- type: dgxa100.80g
34+
count: 3
35+
labels:
36+
nvidia.com/gpu.count: "8"
37+
timeout: 1m
38+
- id: default-queue
39+
type: SubmitObj
40+
params:
41+
refTaskId: register-queue
42+
params:
43+
name: default
44+
- id: test-queue
45+
type: SubmitObj
46+
params:
47+
refTaskId: register-queue
48+
params:
49+
name: test
50+
parentQueue: default
51+
- id: mpijob
52+
type: SubmitObj
53+
params:
54+
refTaskId: register-mpijob
55+
count: 1
56+
params:
57+
namespace: default
58+
queue: test
59+
workers: 2
60+
image: ubuntu
61+
cpu: 100m
62+
memory: 250M
63+
gpu: 8
64+
ttl: "20s"
65+
- id: status
66+
type: CheckPod
67+
params:
68+
refTaskId: mpijob
69+
status: Running
70+
timeout: 10s

‎scripts/create-test-cluster.sh‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,12 @@ if kind get clusters > /dev/null 2>&1; then
3333
read -p "> " choice
3434
if [[ "$choice" == "y" ]]; then
3535
kind delete cluster
36-
kind create cluster --image=kindest/node:v1.29.7
36+
kind create cluster
37+
# --image=kindest/node:v1.29.7
3738
fi
3839
else
39-
kind create cluster --image=kindest/node:v1.29.7
40+
kind create cluster
41+
# --image=kindest/node:v1.29.7
4042
fi
4143

4244
deploy_prometheus
@@ -52,7 +54,8 @@ cat << EOF
5254
3: volcano (https://github.com/volcano-sh/volcano)
5355
4: yunikorn (https://github.com/apache/yunikorn-core)
5456
5: run:ai (https://www.run.ai)
55-
6: combined: coscheduler plugin + jobset + kueue
57+
6: kai (https://github.com/NVIDIA/KAI-Scheduler)
58+
7: combined: coscheduler plugin + jobset + kueue
5659
EOF
5760
read -p "> " choice
5861

@@ -73,6 +76,9 @@ case "$choice" in
7376
deploy_runai
7477
;;
7578
6)
79+
deploy_kai
80+
;;
81+
7)
7682
deploy_scheduler_plugins
7783
deploy_jobset
7884
deploy_kueue

‎scripts/env.sh‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,27 @@ Run:ai deployment requires environment variables:
231231
--set-json 'affinity={"nodeAffinity":{"requiredDuringSchedulingIgnoredDuringExecution":{"nodeSelectorTerms":[{"matchExpressions":[{"key":"type","operator":"NotIn","values":["kwok"]}]}]}}}'
232232
}
233233

234+
# https://github.com/NVIDIA/KAI-Scheduler/
235+
TRAINING_OPERATOR_VERSION=v1.8.0
236+
MPI_OPERATOR_VERSION=v0.4.0
237+
function deploy_kai() {
238+
printGreen Deploying kai
239+
240+
kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=$TRAINING_OPERATOR_VERSION"
241+
242+
kubectl patch deployment training-operator -n kubeflow --type='json' \
243+
-p='[{"op": "add", "path": "/spec/template/spec/containers/0/args", "value": ["--enable-scheme=tfjob", "--enable-scheme=pytorchjob", "--enable-scheme=xgboostjob"]}]'
244+
245+
kubectl delete crd mpijobs.kubeflow.org
246+
247+
kubectl apply -f https://raw.githubusercontent.com/kubeflow/mpi-operator/$MPI_OPERATOR_VERSION/deploy/v2beta1/mpi-operator.yaml
248+
249+
helm repo add --force-update nvidia-k8s https://helm.ngc.nvidia.com/nvidia/k8s
250+
helm repo update
251+
helm upgrade --install kai-scheduler nvidia-k8s/kai-scheduler -n kai-scheduler \
252+
--create-namespace --wait --set "global.registry=nvcr.io/nvidia/k8s"
253+
}
254+
234255
SCHEDULER_PLUGINS_VERSION=v0.29.7
235256
function deploy_scheduler_plugins() {
236257
printGreen Deploying scheduler-plugins

0 commit comments

Comments
 (0)