kubeflow / trainer / 40387943167

78%

push

github

fix(framework): BREAKING CHANGE: Set the TrainJob image on the MPI launcher and nodes (#4301)

* chore(framework): set the trainer image in the MLPolicy plugins

The JobSet builder set spec.trainer.image on the node container of the trainer ancestor replicatedJob for every framework. Move it to EnforceMLPolicy, where each MLPolicy plugin sets the image on its trainer PodSet container, so a plugin can decide which containers run the TrainJob image. PlainML covers XGBoost, Flux and runtimes without an MLPolicy, as it does for env. No behaviour change.

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

* fix(framework): set the TrainJob image on the MPI nodes

The MPI plugin now also sets spec.trainer.image on the node container of the node replicatedJob, so the MPI processes run the same image as mpirun on the launcher. command and args still reach only the launcher, so the nodes keep running sshd. Other frameworks are unchanged.

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

* chore(framework): set the trainer image in the Flux and XGBoost plugins

Rename ApplyTrainerImage to SetContainerImageFromTrainer, and set the TrainJob image on the trainer container in the Flux and XGBoost plugins too, instead of relying on the PlainML fallback.

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

* address review on the trainer image tests and docs

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

* Update pkg/runtime/framework/plugins/torch/torch.go

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

* Update pkg/runtime/framework/plugins/jobset/builder.go

Co-authored-by: An... (continued)

14 of 14 new or added lines in 7 files covered. (100.0%)

3309 of 4279 relevant lines covered (77.33%)

0.78 hits per line