[auraCodesKM]

What this PR does / why we need it:

The runtime mlPolicy.numNodes is stored in the count of the PodSet labeled as the trainer ancestor. In the built-in deepspeed-distributed and mlx-distributed runtimes that is the launcher, and the MPI plugin only updates the node count when the TrainJob sets numNodes. With runtime numNodes greater than 1, the launcher Job got numNodes pods and the launcher endpoints were added to the hostfile. Details and a table of results are in #4210.

EnforceMLPolicy now:

• uses the TrainJob numNodes when set, and otherwise the trainer PodSet count (the runtime numNodes),
• sets the node count from that value, max(numNodes - 1, 1) with runLauncherAsNode and numNodes otherwise, as described in docs/operator-guides/migration.md,
• keeps the launcher at one pod.

For the shipped runtimes with the default numNodes: 1 the generated objects do not change, with or without the TrainJob numNodes.

One behavior change: for a runtime that labels node as the trainer ancestor, with runLauncherAsNode: true and runtime numNodes greater than 1, a TrainJob without numNodes now gets numNodes - 1 node pods instead of numNodes, so the launcher counts as a node like it does for the TrainJob value.

Relationship to #4045 (#3954): that PR changes the lower bound of max(numNodes - 1, 1) to 0 for an explicit numNodes: 1. This PR does not change the bound, and the new tests do not depend on it. The two touch the same lines in EnforceMLPolicy, so whichever merges second needs a small rebase. I checked locally that both changes together pass the tests in this PR. With both, a TrainJob without numNodes on the default runtime also runs launcher only.

Testing (rebased on master at c53a08e, which includes #4029):

• TestMPIRuntimePodCounts in pkg/runtime/core builds the objects from manifests/base/runtimes/deepspeed_distributed.yaml and mlx_distributed.yaml and checks the Job parallelism and hostfile entries. On master 6 of the 10 cases fail, for example "launcher": 3 instead of 1. All pass with this change.
• TestEnforceMLPolicyPodSetCounts in the MPI plugin covers runtimes that label node as the trainer ancestor. On master the runLauncherAsNode case fails with node=3, want 2.
• The OMPI_MCA_* and PRTE_MCA_* env vars generated for the shipped runtimes are the same as on master.
• Unit tests for the make test package list pass (24 packages).
• Integration tests (ginkgo ./test/integration/..., envtest 1.37.0): controller 51/51, statusserver 5/5, webhooks 83/83.
• go vet and gofmt report no issues, golangci-lint v2.14.0 reports 0 issues.
• Not tested on a cluster.

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

Checklist:

• Docs included if any changes are user facing

Assisted-by: Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[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 astefanutti for approval. For more information see the Kubernetes Code Review Process.

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

Details

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

[auraCodesKM]

Rebased on master (081bd75) to resolve an import conflict in pkg/runtime/core/trainingruntime_test.go after #4059. The change is the same as before: mpi.go and mpi_test.go have the same patch, and the test file only keeps both imports.

Local checks on the rebased commit f413d0d:

• go test ./pkg/... (24 packages), go vet ./pkg/..., and golangci-lint v2.14.0 on ./pkg/runtime/...
• envtest: controller 54/54, webhooks 88/88, statusserver 5/5

In the first controller run, the new spec "Should delete the PodGroup when the TrainJob reaches the Failed condition" failed once because the TrainJob webhook did not find the TrainingRuntime yet. It passed in two more full runs and ten focused runs, on this branch and on master. Not tested on a cluster.

Could a maintainer approve the workflow runs so CI can run? The PR fixes the pod counts of the built-in deepspeed-distributed and mlx-distributed runtimes when mlPolicy.numNodes is greater than 1 (#4210).