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

[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

[andreyvelich]

@astefanutti Should we bump k8s libs to 1.32 in this PR to remove toolchain ?

[astefanutti]

@andreyvelich yes, do you prefer the upgrade to be done in this PR or separately?

[andreyvelich]

Yeah, let's maybe first create PR to migrate to 1.32 if that is easier for you.

[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

[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.

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

Using client.Object may work as an approximation to compute SSA patches for simpler objects, but it needs some tweaks like the following one that feel like work-arounds and that may not "scale" well for more complex objects like JobSets:
https://github.com/kubernetes-sigs/kueue/blob/90ef60760b849e475f7cf32d07669bb91bbb479f/pkg/workload/workload.go#L509

I think the target should ultimately be to rely on the runtime.ApplyConfiguration interface that's going to be introduced: kubernetes/kubernetes#129313.

For now, I've moved the conversion to unstructured just before the call to client.Patch, but that'll go away once controller-runtime integrates runtime.ApplyConfiguration and exposes a first-class client.Apply method for SSA.

[astefanutti]

/unhold

All the tests pass now.

[andreyvelich]

We need to rebase PR before the merge.

[tenzen-y]

@astefanutti Could you check commits signs?

[astefanutti]

@astefanutti Could you check commits signs?

@tenzen-y should be good now, I need more coffee :)

[tenzen-y]

Thank you!
/lgtm
/approve
/hold cancel

[google-oss-prow]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: tenzen-y

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

The pull request process is described here

Needs approval from an approver in each of these files:

• OWNERS [tenzen-y]

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