feat(runtimes): make MPI launcher depend on worker readiness

[Amir380-A]

What this PR does / why we need it:

This PR updates the MPI-based ClusterTrainingRuntimes (DeepSpeed and MLX) to ensure the launcher replicated job is created only after the node replicated job becomes Ready by using JobSet's dependsOn API.

Which issue(s) this PR fixes
Fixes #2751

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

• OWNERS

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

[github-actions]

🎉 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:

• Slack: Join our #kubeflow-trainer Slack channel.
• Meetings: Attend the Kubeflow AutoML and Training Working Group bi-weekly meetings.

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

[Copilot]

Pull request overview

This PR updates the MPI-based ClusterTrainingRuntime definitions for DeepSpeed and MLX so the launcher replicated job is only created after the worker (“node”) replicated job reaches Ready, using JobSet’s dependsOn API. This directly addresses issue #2751 by preventing the launcher from starting (and failing) before MPI workers are ready to accept connections.

Changes:

• Reordered replicatedJobs so node comes before launcher (required because the first replicated job cannot use dependsOn).
• Added dependsOn: [{ name: node, status: Ready }] to the launcher replicated job for both DeepSpeed and MLX.
• Updated Helm unittest assertions to reflect the new replicated job ordering and dependency.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
manifests/base/runtimes/mlx_distributed.yaml	Makes MLX launcher creation depend on node readiness via dependsOn.
manifests/base/runtimes/deepspeed_distributed.yaml	Makes DeepSpeed launcher creation depend on node readiness via dependsOn.
charts/kubeflow-trainer/templates/runtimes/mlx-distributed.yaml	Helm template update to emit the dependsOn relationship and new ordering.
charts/kubeflow-trainer/templates/runtimes/deepspeed-distributed.yaml	Helm template update to emit the dependsOn relationship and new ordering.
charts/kubeflow-trainer/tests/runtimes/mlx_distributed_test.yaml	Updates Helm unittest expectations for replicated job order and dependsOn.
charts/kubeflow-trainer/tests/runtimes/deepspeed_distributed_test.yaml	Updates Helm unittest expectations for replicated job order and dependsOn.

[Amir380-A]

/ok-to-test

[google-oss-prow]

@Amir380-A: Cannot trigger testing until a trusted user reviews the PR and leaves an /ok-to-test message.

Details

In response to this:

/ok-to-test

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

[andreyvelich]

/ok-to-test
/retest

[andreyvelich]

Thank you for this @Amir380-A!
Please can you also update the Integration Test:
https://github.com/Amir380-A/trainer/blob/411c212a8a95c566678bda243a7c1927742f3b22/test/integration/controller/trainjob_controller_test.go#L1040
To ensure we have similar DependsOn configuration there.

/ok-to-test
/retest

[andreyvelich]

@Amir380-A Could you check why E2Es are failing please?

[andreyvelich]

/milestone v2.3

[Amir380-A]

/retest

[andreyvelich]

Why do you require to introduce Security Context as part of this PR?
I was able to run DeepSpeed runtime without it and with DependsOn on my local Kind cluster.

[Amir380-A]

I couldn't run DeepSpeed locally because of the massive size of the images required.
In GitHub Actions however, the DeepSpeed E2E was consistently failing with:

Unable to load host key: /home/mpiuser/.ssh/id_rsa
sshd: no hostkeys available -- exiting.

After adding the pod securityContext with fsGroup, the SSH key became readable by the container, sshd started successfully and the Helm E2E passed.

[andreyvelich]

I think, runAsGroup: 1000 might fail the E2Es.
Can you try to remove that alongside the securityContext? It should be working fine.

[Amir380-A]

ok

[Amir380-A]

/test E2E Tests / GPU E2E Test (1.33.1) (pull_request)

[google-oss-prow]

@Amir380-A: No presubmit jobs available for kubeflow/trainer@master

Details

In response to this:

/test E2E Tests / GPU E2E Test (1.33.1) (pull_request)

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

[google-oss-prow]

@Amir380-A: No presubmit jobs available for kubeflow/trainer@master

Details

In response to this:

/test E2E Tests / GPU E2E Test

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.