Support Kubeflow TrainJob (v2)

[tenzen-y]

What would you like to be added:
I would like to support the Kubeflow TrainJob, which is Kubeflow TrainingOperator v2.

The new TrainJob AP relies on the JobSet API, but we should not enqueue the entire TrainJob or JobSet since the TrainJob uses the JobSet DependOn feature, which has some Kueue collaboration problems, as we mentioned there

Hence, in the short term, the TrainJob queueing should rely on the Kueue batch/v1 Job integration to mitigate and avoid the problems once the TrainJob is submitted. This concept is similar to Deployment integration, which depends on the Pod integration.

Additionally, we might be able to extend the JobSet integration so that the JobSet with dependsOn will be enqueued by batch/v1 Job integration to avoid the long-term unused computing resources lock.

Why is this needed:
Currently, we are supporting the Kubeflow v1 APIs, but the v1 APIs will be stopped support within the next year.
We will probably completely stop the v1 support in the next 2 or 3 Training Operator releases.

Completion requirements:

This enhancement requires the following artifacts:

• Design doc
• API change
• Docs update

The artifacts should be linked in subsequent comments.

[tenzen-y]

cc: @andreyvelich @mimowo @kannon92

[mimowo]

@mwielgus @mwysokin

[kannon92]

Hence, in the short term, the TrainJob queueing should rely on the Kueue batch/v1 Job integration to mitigate and avoid the problems once the TrainJob is submitted. This concept is similar to Deployment integration, which depends on the Pod integration.

TrainJob -> JobSet -> (Job A + JobB)

I think you'd want to explore integration with JobSet before you started with jobs.

What exactly is the end state of integration? Do we want DependsOn or are we fine with support via JobSet?

If we want to support DependsOn I think Kubeflow community should focus on implementing that in JobSet. Whatever we add we will have to support. So "Short Term" features would end up being long term features as we can't really drop support once customers use it.

[tenzen-y]

If we want to support DependsOn I think Kubeflow community should focus on implementing that in JobSet. Whatever we add we will have to support. So "Short Term" features would end up being long term features as we can't really drop support once customers use it.

I did not indicate that we aim to implement DependsOn in the Kubeflow repository, we will implement the DependsOn in the JobSet side, first.
The motivation for using batch/v1 Job Kueue integration is https://github.com/andreyvelich/jobset/blob/ea622a6d6ba8cfd876a7d662da2e141c977ecddc/keps/672-serial-job-execution/README.md#risks-and-mitigations.

Even if Kubeflow uses the JobSet DependsOn, we should implement the TrainJob integration as the Kueue batch/v1 Job depending integration in the 1st step since there is no way to enqueue the JobSet step by step.
In the long term, we need to consider how we can separately enqueue each ReplicatedJob as we mentioned in the JobSet DependsOn KEP.

After Kueue supports the separate admission Job, the TrainJob Kueue integration will be implemented as JobSet depending integration.

[k8s-triage-robot]

The Kubernetes project currently lacks enough contributors to adequately respond to all issues.

This bot triages un-triaged issues according to the following rules:

• After 90d of inactivity, lifecycle/stale is applied
• After 30d of inactivity since lifecycle/stale was applied, lifecycle/rotten is applied
• After 30d of inactivity since lifecycle/rotten was applied, the issue is closed

You can:

• Mark this issue as fresh with /remove-lifecycle stale
• Close this issue with /close
• Offer to help out with Issue Triage

Please send feedback to sig-contributor-experience at kubernetes/community.

/lifecycle stale

[andreyvelich]

/remove-lifecycle stale

[kannon92]

@andreyvelich @tenzen-y is there any idea of timeline when this can be implemented?

I know of some people that are interested in using the finetuning functionality with Kueue so it would be worth prioritizing this effort if possible.

[tenzen-y]

@andreyvelich @tenzen-y is there any idea of timeline when this can be implemented?

I know of some people that are interested in using the finetuning functionality with Kueue so it would be worth prioritizing this effort if possible.

Before we start natively supporting TrainJob in Kueue, we want to bump JobSet API version to Beta or GA to avoid multi API versions supporting.
Additionally, users can now enqueue TrainJob via Kueue JobSet integration, as Kueue supports limited integration.

@kannon92 Do you see any lack of Kueue features for TrainJob via JobSet Kueue integration?

[andreyvelich]

Additionally, users can now enqueue TrainJob via Kueue JobSet integration, as Kueue supports limited integration.

@tenzen-y Given that we have suspend and managedBy API on the TrainJob side, we still require integration with Kueue Job framework, isn't?
Also, TrainJob handles some custom orchestration of Kubernetes resources (e.g. MPI Plugin).

Additionally, TrainJob might introduce different state machine compare to JobSet, which needs to be integrated with Kueue and MultiKueue.

[tenzen-y]

Additionally, users can now enqueue TrainJob via Kueue JobSet integration, as Kueue supports limited integration.

@tenzen-y Given that we have suspend and managedBy API on the TrainJob side, we still require integration with Kueue Job framework, isn't? Also, TrainJob handles some custom orchestration of Kubernetes resources (e.g. MPI Plugin).

Additionally, TrainJob might introduce different state machine compare to JobSet, which needs to be integrated with Kueue and MultiKueue.

Yes, we are still eager to natively support TrainJob in Kueue since the JobSet integration is not able to MPI TrainJob, as you said. However, before doing that, KF and JobSet community should mature the APIs.

[andreyvelich]

Shall we search for contributors to support TrainJob within Kueue after we make the 2.0 release ?
Given that many customers will migrate from Training Operator V1 to V2, I feel like the integration with Kueue is a high priority.