KEP-2170: Use SSA to reconcile TrainJob components

[astefanutti]

What this PR does / why we need it:

Replace Get/Create/Update logic with SSA to reconcile TrainJob components.

I've successfully tested it with https://github.com/kubeflow/trainer/blob/master/examples/pytorch/image-classification/mnist.ipynb.

Depends on kubernetes-sigs/jobset#782.

Which issue(s) this PR fixes (optional, in Fixes #<issue number>, #<issue number>, ... format, will close the issue(s) when PR gets merged):

Fixes #2297

Checklist:

• Docs included if any changes are user facing

[google-oss-prow]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign terrytangyuan for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[astefanutti]

/hold

Requires kubernetes-sigs/jobset#782

[astefanutti]

@andreyvelich @tenzen-y @varshaprasad96 please take a look when you can.

[tenzen-y]

Could you specify main branch, and then let us check if it should work fine as expected.

[astefanutti]

Sure, I assumed you preferred to wait until the upcoming JobSet release :).

[tenzen-y]

The current trainer development is under alpha stage. So, I don't care the library stability :)

[andreyvelich]

@tenzen-y @astefanutti Should we merge this PR once JobSet release v0.8.0 ?
I think, we are planning to release it this week: kubernetes-sigs/jobset#774

cc @kannon92 @ahg-g

[astefanutti]

I've replaced the JobSet API version to v0.8.0-devel for now.

I'm working on adapting the unit tests but the main controller should work fine.

[tenzen-y]

Thanks for doing this.
My first question is could we avoid unstructured and then keep using client.Object typed?
I was wondering if we can take a similar approach as https://github.com/kubernetes-sigs/kueue/blob/90ef60760b849e475f7cf32d07669bb91bbb479f/pkg/workload/workload.go#L503 and https://github.com/kubernetes-sigs/kueue/blob/90ef60760b849e475f7cf32d07669bb91bbb479f/pkg/workload/workload.go#L430

[tenzen-y]

Could we use these toolchains? Who is bringing this? JobSet?

[astefanutti]

With the replace added to override the transitive dependencies coming from the newer JobSet version, go mod tidy keeps adding this. I'd need to dig into it more to understand why.

[tenzen-y]

Uhm, interesting. Actually, Go < 1.23.4, they enforce to specify toolchain. But after the Go 1.23.4, the restrictions are relaxed. So, IIUC, this technically could be removed.

Or, we may just bump to Kube 1.32 as library. @andreyvelich Do you want to avoid bumping kube lib to 1.32 here?

[andreyvelich]

I think, we should try to avoid toolchain here if that is possible.
Let's bump k8s to 1.32

[astefanutti]

Thanks for doing this. My first question is could we avoid unstructured and then keep using client.Object typed? I was wondering if we can take a similar approach as https://github.com/kubernetes-sigs/kueue/blob/90ef60760b849e475f7cf32d07669bb91bbb479f/pkg/workload/workload.go#L503 and https://github.com/kubernetes-sigs/kueue/blob/90ef60760b849e475f7cf32d07669bb91bbb479f/pkg/workload/workload.go#L430

Ultimately the client.Patch method converts the object into an unstructured object for SSA, because that's the only way to guarantee no default fields get introduced.

I can look at it again, but from my experience and understanding, relying on apply configuration -> unstructured is the most reliable and canonical way do to SSA.

[astefanutti]

Note I still need to update the newly introduced MPI plugin to rely on apply configurations.

[tenzen-y]

Thanks for doing this. My first question is could we avoid unstructured and then keep using client.Object typed? I was wondering if we can take a similar approach as https://github.com/kubernetes-sigs/kueue/blob/90ef60760b849e475f7cf32d07669bb91bbb479f/pkg/workload/workload.go#L503 and https://github.com/kubernetes-sigs/kueue/blob/90ef60760b849e475f7cf32d07669bb91bbb479f/pkg/workload/workload.go#L430

Ultimately the client.Patch method converts the object into an unstructured object for SSA, because that's the only way to guarantee no default fields get introduced.

I can look at it again, but from my experience and understanding, relying on apply configuration -> unstructured is the most reliable and canonical way do to SSA.

Thank you for checking that. If you think introducing the non unstructured objects with SSA has any potential "risks", please let me know. My understanding might be incorrect for SSA

[tenzen-y]

Note I still need to update the newly introduced MPI plugin to rely on apply configurations.

Yeah, absolutely