fix(framework): BREAKING CHANGE: Set the TrainJob image on the MPI launcher and nodes by matbun · Pull Request #4301 · kubeflow/trainer

matbun · 2026-10-23T20:47:22Z

What this PR does / why we need it:

Follow up of the discussion in #4050. The JobSet builder set spec.trainer.image on the node container of the replicatedJob with the trainer ancestor label, the same container for every framework. This PR moves it to EnforceMLPolicy: each MLPolicy plugin sets the image on its trainer containers with SetContainerImageFromTrainer, and the JobSet plugin syncs it into the JobSet like the other PodSet fields.

Torch, JAX, XGBoost, Flux and runtimes without an MLPolicy keep their current behaviour. The MPI plugin now sets the image on both the launcher and the nodes, so all MPI processes run the same application and MPI versions, while command and args still reach only the launcher and the nodes keep running the SSH server. MPI runtimes whose nodes used a different image than the TrainJob now run the TrainJob image on the nodes too, so the migration guide is updated. This supersedes #4050.

Tested with the unit tests and with the controller and webhook integration tests.

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

Checklist:

Docs included if any changes are user facing

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>

google-oss-prow · 2026-10-23T20:47:30Z

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

andreyvelich

/ok-to-test

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>

matbun · 2026-10-23T21:34:45Z

+	// The TrainJob image runs mpirun on the launcher, the trainer step, and the MPI processes
+	// on the nodes, which must ship the same application and MPI versions as the launcher.
 	runtime.ApplyTrainerImage(info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node), trainJob.Spec.Trainer)
+	runtime.ApplyTrainerImage(info.FindContainerByPodSetName(constants.Node, constants.Node), trainJob.Spec.Trainer)


Relevant for #4030

andreyvelich · 2026-10-23T21:46:22Z

+// ApplyTrainerImage sets the image of the given trainer container from the TrainJob trainer.
+// It is called by the MLPolicy plugins in EnforceMLPolicy, so each plugin decides which
+// containers run the TrainJob image.
+func ApplyTrainerImage(container *Container, jobTrainer *trainer.Trainer) {


Suggested change

// ApplyTrainerImage sets the image of the given trainer container from the TrainJob trainer.

// It is called by the MLPolicy plugins in EnforceMLPolicy, so each plugin decides which

// containers run the TrainJob image.

func ApplyTrainerImage(container *Container, jobTrainer *trainer.Trainer) {

// SetContainerImageFromTrainer sets the container image from the TrainJob trainer image.

// It is a no-op if the container is nil or the TrainJob trainer image is not set.

func SetContainerImageFromTrainer(container *Container, jobTrainer *trainer.Trainer) {

andreyvelich · 2026-10-23T21:47:24Z

 		return nil
 	}

+	// The TrainJob image runs mpirun on the launcher, the trainer step, and the MPI processes


Could you also update Flux plugin?
https://github.com/matbun/trainer/blob/9f331529afbeac8c36659c1207a37e655ae75e0a/pkg/runtime/framework/plugins/flux/flux.go

and XGBoost

done for both

andreyvelich · 2026-10-23T21:50:51Z

+				Scheduler: &runtime.Scheduler{PodLabels: make(map[string]string)},
+			},
+		},
+		"trainJob trainer image is set on the trainer container of XGBoost runtimes": {


This should go to XGBoost plugin, right?

yes, removed it from the PlainML tests. The XGBoost tests already cover it: XGBoost training with GPU resources, resources set in TrainJob only and resources set in both Runtime and TrainJob now expect the TrainJob image on the trainer container

andreyvelich · 2026-10-23T21:51:57Z

 	}
 }
+
+func TestMPITrainerImage(t *testing.T) {


Rather than adding another test func, let's just add image in this test: https://github.com/matbun/trainer/blob/9f331529afbeac8c36659c1207a37e655ae75e0a/pkg/runtime/framework/plugins/mpi/mpi_test.go#L89

Done, removed TestMPITrainerImage and moved the image into TestMPI.

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>

andreyvelich · 2026-10-23T22:04:45Z

+	// The TrainJob image runs mpirun on the launcher, the trainer step, and the MPI processes
+	// on the nodes, which must ship the same application and MPI versions as the launcher.
+	runtime.SetContainerImageFromTrainer(info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node), trainJob.Spec.Trainer)
+	runtime.SetContainerImageFromTrainer(info.FindContainerByPodSetName(constants.Node, constants.Node), trainJob.Spec.Trainer)


Look up both PodSets by name, consistent with the rest of the MPI plugin (e.g. the volumes loop below):

Suggested change

// The TrainJob image runs mpirun on the launcher, the trainer step, and the MPI processes

// on the nodes, which must ship the same application and MPI versions as the launcher.

runtime.SetContainerImageFromTrainer(info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node), trainJob.Spec.Trainer)

runtime.SetContainerImageFromTrainer(info.FindContainerByPodSetName(constants.Node, constants.Node), trainJob.Spec.Trainer)

// Set the TrainJob image on the launcher and the nodes, since mpirun on the launcher starts

// the MPI processes on the nodes, which must ship the same application and MPI versions.

for _, psName := range []string{constants.Launcher, constants.Node} {

runtime.SetContainerImageFromTrainer(info.FindContainerByPodSetName(psName, constants.Node), trainJob.Spec.Trainer)

}

applied, thanks!

andreyvelich · 2026-10-23T22:04:45Z

+	// The TrainJob image runs mpirun on the launcher, the trainer step, and the MPI processes
+	// on the nodes, which must ship the same application and MPI versions as the launcher.
+	runtime.SetContainerImageFromTrainer(info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node), trainJob.Spec.Trainer)
+	runtime.SetContainerImageFromTrainer(info.FindContainerByPodSetName(constants.Node, constants.Node), trainJob.Spec.Trainer)


Test follow-ups for the node image:

When moving the cases into TestMPI, please also cover the runLauncherAsNode: true layout.

In the "MPI launcher runs as a node" integration test (test/integration/controller/trainjob_controller_test.go ~L1334), add g.Expect(podSpec.Containers[cIdx].Image).Should(gomega.Equal("test:trainjob")) to verify the node image end to end.

Done: TestMPI now checks the image on the launcher and the nodes in both the default layout (trainJob numNodes is respected rather than mlPolicy one) and runLauncherAsNode is true, and the launcher as a node integration test asserts the node image.

andreyvelich · 2026-10-23T22:04:46Z

 						Containers: []runtime.Container{{
-							Name: constants.Node,
+							Name:  constants.Node,
+							Image: "test:image",


These runtime containers have no image, so the tests only cover filling an empty image, not replacing one. Please set .WithImage("test:runtime") on the runtime container in at least one case (e.g. L420 here and jax_test.go L422). Similarly, the JAX/XGBoost/OpenMPI/Flux integration fixtures (trainjob_controller_test.go L1998, L2321, L1387, L2664) use test:trainjob for the runtime image too, so they can't detect a missing override.

the runtime containers now use test:runtime in that Torch case, the JAX case, both MPI cases, and the JAX, XGBoost, OpenMPI and Flux integration fixtures, so the tests check that the TrainJob image replaces the runtime one

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

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

Set the TrainJob image on the MPI launcher and nodes in one loop, as suggested in the review. Cover the image in both the default and the runLauncherAsNode layouts in TestMPI, give the runtime containers their own image in the Torch, JAX and MPI tests and in the JAX, XGBoost, OpenMPI and Flux integration fixtures, so the tests check that the TrainJob image replaces it, and assert the node image in the MPI launcher as a node integration test. Co-authored-by: Andrey Velichkevich <andrey.velichkevich@gmail.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>

…gins Give the runtime trainer container its own image in the XGBoost and Flux tests, so they check that the TrainJob image replaces it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Matteo Bunino <48362942+matbun@users.noreply.github.com>

andreyvelich

Thanks for this @matbun!
/lgtm
/assign @tenzen-y @robert-bell @akshaychitneni @abhijeet-dhumal

andreyvelich · 2026-10-23T22:43:32Z

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

github-project-automation Bot added this to Kubeflow Trainer Oct 23, 2026

google-oss-prow Bot added the do-not-merge/work-in-progress label Oct 23, 2026

github-project-automation Bot moved this to Needs Triage in Kubeflow Trainer Oct 23, 2026

google-oss-prow Bot requested a review from akshaychitneni October 23, 2026 20:47

google-oss-prow Bot requested a review from robert-bell October 23, 2026 20:47

google-oss-prow Bot added the size/L label Oct 23, 2026

andreyvelich reviewed Oct 23, 2026

View reviewed changes

google-oss-prow Bot added the ok-to-test label Oct 23, 2026

matbun marked this pull request as ready for review October 23, 2026 21:33

google-oss-prow Bot removed the do-not-merge/work-in-progress label Oct 23, 2026

google-oss-prow Bot requested a review from jinchihe October 23, 2026 21:33

matbun commented Oct 23, 2026

View reviewed changes

andreyvelich reviewed Oct 23, 2026

View reviewed changes

matbun and others added 5 commits October 24, 2026 00:06

address review on the trainer image tests and docs

de38b9c

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

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

2eb4de5

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

af840cd

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

Update pkg/runtime/runtime_test.go

b208f57

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

matbun changed the title ~~chore(framework): set the trainer image in the MLPolicy plugins~~ fix(framework): set the TrainJob image in the MLPolicy plugins and on the MPI nodes Oct 23, 2026

andreyvelich reviewed Oct 23, 2026

View reviewed changes

google-oss-prow Bot assigned abhijeet-dhumal, akshaychitneni and robert-bell Oct 23, 2026

google-oss-prow Bot assigned tenzen-y and andreyvelich Oct 23, 2026

google-oss-prow Bot added the lgtm label Oct 23, 2026

google-oss-prow Bot changed the title ~~fix(framework): set the TrainJob image in the MLPolicy plugins and on the MPI nodes~~ fix(framework): BREAKING CHANGE: Set the TrainJob image on the MPI launcher and nodes Oct 23, 2026

Conversation

matbun commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

google-oss-prow Bot commented Oct 23, 2026

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andreyvelich Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

andreyvelich commented Oct 23, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

matbun commented Oct 23, 2026 •

edited

Loading

andreyvelich Oct 23, 2026 •

edited

Loading