feat(api): allow spec.trainer.numNodes to be mutable when TrainJob is suspended

[kapil27]

What you would like to be added?

Allow spec.trainer.numNodes to be modified on a suspended TrainJob. Currently,
spec.trainer is fully immutable after creation (added in #3157), which prevents
users from scaling the number of training nodes on a paused job.

Why is this needed?

Before #3157 (Trainer v2.1), spec.trainer was mutable. Users and dashboards
could pause a TrainJob, change numNodes, and resume. The controller would
reconcile the JobSet with the updated replica count on unsuspend.
PR #3157 added blanket self == oldSelf CEL on spec.trainer to prevent mutations
that would conflict with the running JobSet (e.g., changing image or command mid-run).
This was the right call for most fields, but it also blocked the legitimate use case
of scaling numNodes on a suspended job where no pods are running and the JobSet
can safely be rebuilt.

Use case

1. User creates a TrainJob with numNodes: 1
2. Training runs, user pauses the job (checkpoints saved)
3. User wants to scale to numNodes: 2 for faster training
4. User edits spec.trainer.numNodes on the suspended TrainJob
5. User resumes → controller rebuilds JobSet with 2 replicas → training resumes from checkpoint

What should stay immutable

All other spec.trainer fields should remain immutable:

• image, command, args — changing these on a running distributed job would
  cause pod restarts with inconsistent binaries
• resourcesPerNode — changing resource requests requires pod recreation
• env — changing environment mid-training can corrupt state
• numProcPerNode — tied to GPU topology, unsafe to change
  Only numNodes should be exempted, and only when suspended.

References

• feat(api): Fix immutability of the TrainJob APIs #3157 — feat(api): Fix immutability of the TrainJob APIs (added the blanket immutability)
• fix(api): allow atomic RuntimePatches update on unsuspend #3469 — fix(api): allow atomic RuntimePatches update on unsuspend (precedent for suspend-aware exemption)
• bug(api): CRD CEL immutability rules on runtimePatches sub-fields block Kueue unsuspend #3888 — CEL immutability blocks Kueue unsuspend (related, different root cause)

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[robert-bell]

/kind enhancement
/area crd

@andreyvelich is this something that is no longer supported by JobSet? This worked previously in KFT 2.1. Do you know what would be required to get this working again? This is a use case our downstream previously supported.

[robert-bell]

/assign

[andreyvelich]

JobSet was blocking modification of parallelism and completions before @aniket2405's proposal on elastic JobSet: https://github.com/kubernetes-sigs/jobset/tree/main/keps/463-ElasticJobsets

So we might want to graduate this feature to beta as @kannon92 suggested here: kubernetes-sigs/jobset#1296, so it is enabled by default.

/remove-label needs-triage

[robert-bell]

@andreyvelich do you know- was this something that was removed in the version of jobset trainer uses since 2.2? We had this working previously in Trainer 2.1 -- the TrainJob had to be suspended.

[andreyvelich]

We had this working previously in Trainer 2.1 -- the TrainJob had to be suspended.

Did it also update the values in JobSet and Job spec after TrainJob is suspended?

[robert-bell]

We had this working previously in Trainer 2.1 -- the TrainJob had to be suspended.

Did it also update the values in JobSet and Job spec after TrainJob is suspended?

No I don't believe we did. Not to worry. I'm happy to investigate.