kubeflow / kubeflow/trainer

KEP-3015: Support Workload API for TrainJob Scheduling

Open
#3,015 12 comments 8 reactions 1 assignee Claimed by @andreyvelich View on GitHub
area/api area/controller kind/feature kind/plan-kep
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.