kubeflow / trainer / 41473986622

80%

push

github

fix(framework): set MPI resourcesPerNode on the nodes (#4364)

* fix(framework): set MPI resourcesPerNode on the nodes

With runLauncherAsNode false, resourcesPerNode was merged into the launcher, which only runs mpirun, while the node Pods that run the MPI processes kept the runtime resources. For MPI, set the per node resources and claims on the node replicatedJob always, and on the launcher only when it runs as a node.

Fixes #4363

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>

* Update pkg/runtime/core/trainingruntime_test.go

Co-authored-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>

* chore(test): compare the full runtime.Info in TestNewRuntimeInfo

Each case now carries the TrainJob, the JobSet template and the policies, built with the testingutil wrappers, and the expected runtime.Info, with the DRA gate enabled. Add cases for runLauncherAsNode unset, MPI resourceClaimsPerNode, and resourcesPerNode merged with the runtime resources. Remove the builder TODO that this fix resolves.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>

* chore(test): set the TestNewRuntimeInfo inputs in each case

Inline the TrainJob, JobSet and MLPolicy wrappers in every case so each case reads on its own, and update the MPI guide: resourcesPerNode now always reaches the nodes and reaches the launcher only with runLauncherAsNode.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>

* chore(test): pass scheduling to newRuntimeInfo in TestNewRuntimeInfo

Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>

---------

Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anth... (continued)

5 of 5 new or added lines in 1 file covered. (100.0%)

3946 of 4942 relevant lines covered (79.85%)

0.8 hits per line