KEP-3015: Support Workload API for TrainJob Scheduling
- Dominant language
- Go
- Stars
- 2.2k
- Forks
- 1.1k
- Avg merge
- 3d 22h
- Merged PRs (30d)
- 39
Description
### What you would like to be added?
Kubernetes recently introduced support for Workload API: https://github.com/kubernetes/enhancements/issues/4671
In v1.35 this API allows to schedule group of Pods together (e.g. Gang-Scheduling).
We should create dedicated `PodGroupPolicy` plugin to support creation of Workload object.
### API Changes
We can design an initial API as follows:
```yaml
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
name: mpi-runtime
spec:
podGroupPolicy:
workload: {}
```
If this plugin is enabled, it should create Workload API:
```yaml
apiVersion: scheduling.k8s.io/v1alpha2
kind: Workload
metadata:
name:
ownerReferences:
- apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
name:
spec:
controllerRef:
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
name:
podGroupTemplate:
- name: trainer
schedulingPolicy:
gang:
minCount: 8 # Equal to trainJob.spec.trainer.numNodes
---
apiVersion: scheduling.k8s.io/v1alpha1
kind: PodGroup
metadata:
name: --
ownerReferences:
- apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
name:
- apiVersion: scheduling.k8s.io/v1alpha1
kind: Workload
name:
spec:
podGroupTemplateRef:
workloadName:
podGroupTemplateName: trainer
schedulingPolicy:
gang:
minCount: 8 # Equal to trainJob.spec.trainer.numNodes
```
And update the Pod's spec schedulingGroup:
```yaml
spec:
schedulingGroup:
podGroupName: --
```
### JobSet Integrations
@imreddy13 has an open KEP to support Workload creation in the JobSet controller as well: https://github.com/kubernetes-sigs/jobset/issues/969
We should discuss whether we want to re-use this API once it is available, or TrainJob controller should always be responsible to create Workload object.
Since we have plans to support non-JobSet based jobs, like Flux MiniCluster: https://github.com/kubeflow/trainer/pull/2909 (cc @vsoch)
cc @kannon92 @kubeflow/kubeflow-trainer-team @macsko @wojtek-t @erictune
/area api
/area controller
### Why is this needed?
Enable gang-scheduling in TrainJob using the Workload API.
### Love this feature?
Give it a 👍 We prioritize the features with most 👍
Contributor guide
Assessment
This issue has not been assessed yet.