Skip to content

fix(runtimes): Apply trainer overrides to all node containers - #4050

Open
Harsh-2005d wants to merge 1 commit into
kubeflow:masterfrom
Harsh-2005d:mpi-fix
Open

Harsh-2005d wants to merge 1 commit into
kubeflow:masterfrom
Harsh-2005d:mpi-fix

Conversation

@Harsh-2005d

Copy link
Copy Markdown

What this PR does / why we need it:

  • The TrainJob Image is now applied to any container named node where the ancestor label is either trainer or empty (which applies to the MPI worker).

  • The TrainJob Command and Args are strictly applied only if the ancestor label is trainer (the Launcher).

Which issue(s) this PR fixes
Fixes #4030

Checklist:

  • Docs included if any changes are user facing

@google-oss-prow

Copy link
Copy Markdown
Contributor

[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 electronic-waste 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:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

@github-actions

github-actions Bot commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

🎉 Welcome to the Kubeflow Trainer! 🎉

Thanks for opening your first PR! We're happy to have you as part of our community 🚀

Here's what happens next:

  • If you haven't already, please check out our Contributing Guide for repo-specific guidelines and the Kubeflow Contributor Guide for general community standards.
  • Our team will review your PR soon! cc @kubeflow/kubeflow-trainer-team

Join the community:

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

Comment on lines +127 to +135
for j, container := range rJob.Template.Spec.Template.Spec.Containers {
if *container.Name != constants.Node {
continue
}
// Update values from the TrainJob trainer.
if ancestor == constants.AncestorTrainer || ancestor == "" {
if jobTrainer := trainJob.Spec.Trainer; jobTrainer != nil {
if image := jobTrainer.Image; image != nil {
b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Image = image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Harsh-2005d @robert-bell @matbun @tenzen-y Any thoughts on merging this PR before 2.4 release?
Technically, it is a breaking change, but without it users are unable to update Runtime's image using TrainJob API: #4030 (comment)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This changes image propagation for non-MPI and custom runtimes, which is very wide reaching. I'd rather not make this change as I think we'll need to undo it at some point.

What about making a more targeted fix to just at MPI workloads and have trainer.image set the image for both the launcher and worker? @matbun would that solve #4030 for you? trainer.command etc. would still only be applied to the launcher.

I think we might also be able to do this implementation in the mpi plugin EnforceMLPolicy to avoid leaking mpi implementation into the jobset plugin.

For the broader behaviour, I think we need to explore a broader design for 2.5 like Andrey proposed in #4030 (comment).

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, doing this just for MPI plugin sounds good.
We currently only set image for sidecar container in Flux plugin: https://github.com/kubeflow/trainer/blob/master/pkg/runtime/framework/plugins/flux/flux.go#L151

@robert-bell Any thoughts how we can configure image in other plugins during EnforceMLPolicy?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have any other images to set other than the node and launcher?
Flux lets you configure the image with an envar.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

doing this just for MPI plugin sounds good to me as well! Thanks @robert-bell !

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have any other images to set other than the node and launcher?

Not at the moment, but we might need to refactor all plugins to set PodSet image via EnforceMLPolicy interface, rather than setting it via trainJob.trainer.image: https://github.com/Harsh-2005d/trainer/blob/mpi-fix/pkg/runtime/framework/plugins/jobset/builder.go#L134-L136

@Harsh-2005d @matbun Do you want to pick up this task as part of this PR?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Happy to take the EnforceMLPolicy refactor for all plugins in a follow-up PR if @Harsh-2005d prefers to keep this one to the MPI fix

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's do that in the one PR, so feel free to take this over @matbun!

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd be happy to take on the broader refactor to move image configuration into EnforceMLPolicy as part of this PR, if that's okay with you. @matbun @andreyvelich

@andreyvelich

Copy link
Copy Markdown
Member

@Harsh-2005d Could you also rebase this PR please?

@andreyvelich

Copy link
Copy Markdown
Member

/ok-to-test

Fixes kubeflow#4030

Decouple command injection from image injection for MPI workers.

Signed-off-by: Harsh Dahiya <dahiyaharsh2005@gmail.com>
@Harsh-2005d Harsh-2005d changed the title fix(runtime): Apply trainer overrides to all node containers fix(runtimes): Apply trainer overrides to all node containers Oct 21, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: Needs Triage

Development

Successfully merging this pull request may close these issues.

MPI: spec.trainer.image and spec.trainer.command cannot both be used once the worker job carries the trainer ancestor label

4 participants