kubeflow / trainer / 38738721619

77%

push

github

fix(runtimes): emit PRTE env vars so MPI plugin works with Open MPI 5 (#4029)

* fix(runtime): emit PRTE env vars so MPI plugin works with Open MPI 5

The MPI plugin configures OpenMPI runtimes exclusively through ORTE-era
MCA environment variables (OMPI_MCA_orte_default_hostfile,
OMPI_MCA_orte_keep_fqdn_hostnames, OMPI_MCA_orte_set_default_slots and
OMPI_MCA_plm_rsh_args).

Open MPI 5.0 removed ORTE and replaced it with PRTE, so those keys are
no longer recognised. Open MPI ignores unknown OMPI_MCA_* variables
silently, so on an Open MPI >= 5 image the generated hostfile is never
read: mpirun falls back to localhost with one slot per core and runs
every rank on the launcher. The TrainJob still reports Complete, which
makes a multi-node job silently degrade to single-node instead of
failing.

Reproduced on Kubeflow Trainer v2.1.0 with mpioperator/mpi-pi:openmpi
(Open MPI 5.0.7) and numNodes: 3 -- all 20 ranks landed on the launcher
while both worker pods idled. Passing --hostfile explicitly restored the
expected 3 ranks across 3 pods.

mpiImplementation has no version axis and the image's MPI version is not
knowable from the manifest, so rather than detect a version this emits
both families. Open MPI 4 ignores the PRTE_MCA_* keys and Open MPI 5
ignores the OMPI_MCA_orte_* ones, so a single TrainingRuntime is correct
on both. The PRTE names were verified against `prte_info --all` from
Open MPI 5.0.7.

PRTE's keep_fqdn_hostnames is deliberately left disabled, unlike its
ORTE counterpart. The generated hostfile always uses JobSet's short-form
Pod DNS name (<pod>.<service>), which never string-matches what PRTE
resolves for itself; with that option on and runLauncherAsNode enabled,
PRTE decided the launcher's own hostfile entry was a remote host and
tried (and failed) to SSH to itself -- caught by re-running the fix on
a live TrainJob before this branch was pushed anywhere.

The new keys are added to MPIReservedEnvNames for consistency with the
existing v... (continued)

24 of 24 new or added lines in 1 file covered. (100.0%)

3159 of 4314 relevant lines covered (73.23%)

0.74 hits per line