[sakethalladaaa]

What this PR does / why we need it:

Fixes MPI runtime node scaling when mlPolicy.numNodes is set on the runtime and TrainJob.spec.trainer.numNodes is omitted.

Previously, the runtime numNodes value was applied to the trainer ancestor PodSet. For built-in MPI runtimes such as DeepSpeed and MLX, the trainer ancestor is the launcher, which caused the launcher Job to scale while the worker node Job kept its template count.

This change:

• derives the effective runtime numNodes from the trainer ancestor PodSet,
• preserves TrainJob.spec.trainer.numNodes as the higher-priority override,
• keeps the MPI launcher count at 1,
• applies the effective node count to the worker PodSet,
• adds regression coverage for the runtime-level numNodes fallback.

Validation:

• make test
• make generate
• targeted MPI and runtime/core tests

Assisted-by: ChatGPT

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

Checklist:

• Docs included if any changes are user facing

[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 gaocegege 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! 🙏

[harshrajdebug]

I ran this against master and against #4211, which fixes the same issue.

On the shipped runtimes the two give the same result. #4211's TestMPIRuntimePodCounts (deepspeed_distributed and mlx_distributed, 10 cases) and its MPI plugin tests pass on this branch, and TestTrainingRuntimeNewObjectsMPIRuntimeNumNodes from this PR passes on #4211.

They differ when no replicated Job has the trainer.kubeflow.org/trainjob-ancestor-step: trainer label, which the TrainingRuntime webhook does not require. The PodSet counts then come from each Job's parallelism, and since numNodes starts at 1 here, a TrainJob without numNodes gets a single node pod. I took deepspeed_distributed.yaml, removed the ancestor label from both Jobs, set the node Job's parallelism to 4 and created a TrainJob without numNodes:

	node parallelism
master	4
#4211	4
this PR	1

The result is the same with runLauncherAsNode true and false. Leaving the node count unchanged when there is neither a trainer PodSet nor a TrainJob numNodes, as #4211 does by keeping numNodes nil, preserves master's behavior for these runtimes.

cc @robert-bell since both PRs fix #4210.