Support Workload API for TrainJob Scheduling

[andreyvelich]

What you would like to be added?

Kubernetes recently introduced support for Workload API: kubernetes/enhancements#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:

apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: mpi-runtime
spec:
  podGroupPolicy:
    workload: {}

If this plugin is enabled, it should create Workload API:

apiVersion: scheduling.k8s.io/v1alpha2
kind: Workload
metadata:
  name: <job-name>
  ownerReferences:
  - apiVersion: trainer.kubeflow.org/v1alpha1
    kind: TrainJob
    name: <job-name>
spec:
  controllerRef:
    apiVersion: trainer.kubeflow.org/v1alpha1
    kind: TrainJob
    name: <job-name>
  podGroupTemplate:
  - name: trainer
    schedulingPolicy:
      gang:
        minCount: 8  # Equal to trainJob.spec.trainer.numNodes
---
apiVersion: scheduling.k8s.io/v1alpha1
kind: PodGroup
metadata:
  name: <job-name>-<podGroup-template-name>-<hash>
  ownerReferences:
  - apiVersion: trainer.kubeflow.org/v1alpha1
    kind: TrainJob
    name: <job-name>
  - apiVersion: scheduling.k8s.io/v1alpha1
    kind: Workload
    name: <job-name>
spec:
  podGroupTemplateRef:
    workloadName: <job-name>
    podGroupTemplateName: trainer
  schedulingPolicy:
    gang:
      minCount: 8  # Equal to trainJob.spec.trainer.numNodes

And update the Pod's spec schedulingGroup:

spec:
  schedulingGroup:
    podGroupName: <workload-name>-<podGroup-template-name>-<hash>

JobSet Integrations

@imreddy13 has an open KEP to support Workload creation in the JobSet controller as well: kubernetes-sigs/jobset#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: #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 👍

[vsoch]

Since we have plans to support non-JobSet based jobs, like Flux MiniCluster: #2909 (cc @vsoch)

To clarify, the implementation that I've done for Kubeflow (the PR isn't open yet because I'm waiting on the HPC Policy KEP) does use JobSet. And we don't deploy an official MiniCluster, but we create the underlying abstractions (and sub the Indexed Job for the native JobSet).

[andreyvelich]

And we don't deploy an official MiniCluster, but we create the underlying abstractions (and sub the Indexed Job for the native JobSet).

That makes sense, as we discussed initially we will support Flux with JobSet implementation.

I can imagine in the future we can support other runtimes, which are not JobSet based. Let's say we want to give users stateful HPC cluster environment to which they can connect to and execute mpirun interactively.

[TusharMohapatra07]

Hey maintainers! can i work on this ??

[kannon92]

I don't think if this is a great feature to work on right now imo.

There is a lot of work in flight around WAS in upstream Kubernetes.

[TusharMohapatra07]

Makes sense @kannon92! Thanks for clarifying :)

[andreyvelich]

I think, we can leverage JobSet's API for Workload object creation for now: kubernetes-sigs/jobset#1111
If the approach of the workloadTemplate API works well, we can integrate similar functionality into TrainJob API later.