From 5ef94c45e100ff9f53da880caa5a45b2693803e6 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Mon, 7 Sep 2026 01:26:16 +0800 Subject: [PATCH 1/3] Run as a single process when there is no launcher and no MPI init_distributed fills in the distributed environment when a launcher did not, and its only route for that was mpi_discovery, which imports mpi4py. So `python train.py` on a machine with one accelerator and no launcher ends at ModuleNotFoundError: No module named 'mpi4py' before deepspeed.initialize returns. Reproduced on both an Apple M5 (MPS, ZeRO 0-3) and an H20 (CUDA, ZeRO-1); neither had mpi4py installed, which is the default for a machine that was never going to run an MPI job. comm/utils.py already reads a missing launcher as rank 0 of a world of 1 - get_local_rank_from_launcher and its siblings say so in a comment - so put that same reading in the environment the backend is initialized from. Nothing changes for anyone who has mpi4py: mpi_discovery still runs first and only an ImportError from it is caught. If that fires while an MPI launcher's rank variable is set, the run stops with a message naming mpi4py rather than silently becoming one rank of a many-rank job. Signed-off-by: alanhuangyoo --- deepspeed/comm/comm.py | 48 +++++++++++- .../comm/test_single_process_discovery.py | 74 +++++++++++++++++++ 2 files changed, 121 insertions(+), 1 deletion(-) create mode 100644 tests/unit/comm/test_single_process_discovery.py diff --git a/deepspeed/comm/comm.py b/deepspeed/comm/comm.py index 635f64fef8b8..d2b63259d7ed 100755 --- a/deepspeed/comm/comm.py +++ b/deepspeed/comm/comm.py @@ -843,7 +843,14 @@ def init_distributed(dist_backend: Optional[str] = None, elif in_aws_sm(): patch_aws_sm_env_for_torch_nccl_backend(verbose=verbose) else: - mpi_discovery(distributed_port=distributed_port, verbose=verbose) + try: + mpi_discovery(distributed_port=distributed_port, verbose=verbose) + except ImportError as err: + if in_mpi_job(): + raise ImportError("An MPI job is running but mpi4py is not installed, so the rank and " + "world size cannot be discovered from it. Install mpi4py, or set RANK, " + "WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and MASTER_PORT yourself.") from err + single_process_discovery(distributed_port=distributed_port, verbose=verbose) if cdb is not None and cdb.is_initialized(): if int(os.getenv('RANK', '0')) == 0: @@ -858,6 +865,45 @@ def init_distributed(dist_backend: Optional[str] = None, cdb = TorchBackend(dist_backend, timeout, init_method, rank, world_size) +# Rank variables the MPI launchers export: OpenMPI, MPICH and Intel MPI, PMIx, MVAPICH, and +# Slurm's srun. Used only to tell an MPI job that is missing mpi4py from a machine that has no +# launcher at all, so an unrecognized launcher still reaches mpi_discovery as before. +MPI_RANK_ENV_VARS = ("OMPI_COMM_WORLD_RANK", "PMI_RANK", "PMIX_RANK", "MV2_COMM_WORLD_RANK", "SLURM_PROCID") + + +def in_mpi_job(): + """Whether an MPI launcher started this process.""" + return any(var in os.environ for var in MPI_RANK_ENV_VARS) + + +def single_process_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True): + """Fill in the distributed environment for one process on one device. + + Reached when no launcher set the variables and no MPI job is running - `python train.py` on + a single-accelerator machine, which is how DeepSpeed is used on a laptop. `mpi_discovery` + cannot serve that case: it imports mpi4py, so without that package the run ends at + `ModuleNotFoundError: No module named 'mpi4py'` before `deepspeed.initialize` returns. + + `comm/utils.py` already reads a missing launcher as rank 0 of a world of 1; this puts the + same reading in the environment the backend is initialized from. + """ + defaults = { + "RANK": "0", + "WORLD_SIZE": "1", + "LOCAL_RANK": "0", + "MASTER_ADDR": "127.0.0.1", + "MASTER_PORT": str(distributed_port), + } + for name, value in defaults.items(): + os.environ.setdefault(name, value) + + if verbose: + utils.logger.info("No launcher and no MPI job detected; running as a single process with world_rank={}, " + "local_rank={}, world_size={}, master_addr={}, master_port={}".format( + os.environ["RANK"], os.environ["LOCAL_RANK"], os.environ["WORLD_SIZE"], + os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"])) + + def mpi_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True): ''' Discovery MPI environment via mpi4py and map to relevant dist state diff --git a/tests/unit/comm/test_single_process_discovery.py b/tests/unit/comm/test_single_process_discovery.py new file mode 100644 index 000000000000..e37bf0ec7b2e --- /dev/null +++ b/tests/unit/comm/test_single_process_discovery.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""`deepspeed.initialize` on a machine with no launcher must not require mpi4py. + +`init_distributed` fills in the distributed environment when a launcher did not, and its only +route for that was `mpi_discovery`, which imports mpi4py. Running `python train.py` on a single +accelerator - no launcher, no MPI, no mpi4py - therefore ended at `ModuleNotFoundError: No +module named 'mpi4py'` before `deepspeed.initialize` returned. +""" + +import os + +import pytest + +from deepspeed.comm.comm import MPI_RANK_ENV_VARS, in_mpi_job, single_process_discovery + +LAUNCHER_ENV = ("RANK", "WORLD_SIZE", "LOCAL_RANK", "MASTER_ADDR", "MASTER_PORT") + + +@pytest.fixture +def clean_env(monkeypatch): + for name in LAUNCHER_ENV + MPI_RANK_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +def test_no_mpi_variables_is_not_an_mpi_job(clean_env): + assert in_mpi_job() is False + + +@pytest.mark.parametrize("var", MPI_RANK_ENV_VARS) +def test_each_launcher_rank_variable_marks_an_mpi_job(clean_env, monkeypatch, var): + """OpenMPI, MPICH/Intel MPI, PMIx, MVAPICH and srun each export a different one.""" + monkeypatch.setenv(var, "0") + + assert in_mpi_job() is True + + +def test_single_process_discovery_fills_the_environment(clean_env): + single_process_discovery(distributed_port=29501, verbose=False) + + assert os.environ["RANK"] == "0" + assert os.environ["LOCAL_RANK"] == "0" + assert os.environ["WORLD_SIZE"] == "1" + assert os.environ["MASTER_ADDR"] == "127.0.0.1" + assert os.environ["MASTER_PORT"] == "29501" + + +def test_single_process_discovery_leaves_what_the_caller_set(clean_env, monkeypatch): + """A partially set environment is the caller's, not something to overwrite.""" + monkeypatch.setenv("MASTER_PORT", "12345") + monkeypatch.setenv("MASTER_ADDR", "10.0.0.7") + + single_process_discovery(distributed_port=29501, verbose=False) + + assert os.environ["MASTER_PORT"] == "12345" + assert os.environ["MASTER_ADDR"] == "10.0.0.7" + assert os.environ["WORLD_SIZE"] == "1" + + +def test_an_mpi_job_without_mpi4py_says_so(clean_env, monkeypatch): + """Falling back to a single process there would silently run one rank of a many-rank job.""" + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "3") + + with pytest.raises(ImportError, match="mpi4py"): + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert "WORLD_SIZE" not in os.environ, "the environment must not be filled in for an MPI job" From 114ba8517e90cc6465cc72421943b23e11c451fa Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Mon, 7 Sep 2026 13:42:06 +0800 Subject: [PATCH 2/3] Key the mpi4py error on the world size, not on a rank variable A rank variable says a launcher is present, not that the world is bigger than one. srun -n1 sets SLURM_PROCID for a single-task step, which is one process on one device and is exactly the case this PR exists to serve, so keying on the rank sent it to the error instead of the fallback. Read the size variables the same launchers export alongside their rank ones - OMPI_COMM_WORLD_SIZE, PMI_SIZE, PMIX_SIZE, MV2_COMM_WORLD_SIZE, SLURM_NTASKS - and only refuse when one of them is above 1. srun -n1 trains srun -n4 ImportError naming mpi4py mpirun -n2 ImportError naming mpi4py no launcher trains Caught by ebarkhordar in review. Signed-off-by: alanhuangyoo --- deepspeed/comm/comm.py | 31 +++++--- .../comm/test_single_process_discovery.py | 70 ++++++++++++++++--- 2 files changed, 79 insertions(+), 22 deletions(-) diff --git a/deepspeed/comm/comm.py b/deepspeed/comm/comm.py index d2b63259d7ed..029fb0b76618 100755 --- a/deepspeed/comm/comm.py +++ b/deepspeed/comm/comm.py @@ -846,10 +846,11 @@ def init_distributed(dist_backend: Optional[str] = None, try: mpi_discovery(distributed_port=distributed_port, verbose=verbose) except ImportError as err: - if in_mpi_job(): - raise ImportError("An MPI job is running but mpi4py is not installed, so the rank and " - "world size cannot be discovered from it. Install mpi4py, or set RANK, " - "WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and MASTER_PORT yourself.") from err + if in_multi_rank_mpi_job(): + raise ImportError("A multi-rank MPI job is running but mpi4py is not installed, so the " + "rank and world size cannot be discovered from it. Install mpi4py, or " + "set RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and MASTER_PORT " + "yourself.") from err single_process_discovery(distributed_port=distributed_port, verbose=verbose) if cdb is not None and cdb.is_initialized(): @@ -865,15 +866,23 @@ def init_distributed(dist_backend: Optional[str] = None, cdb = TorchBackend(dist_backend, timeout, init_method, rank, world_size) -# Rank variables the MPI launchers export: OpenMPI, MPICH and Intel MPI, PMIx, MVAPICH, and -# Slurm's srun. Used only to tell an MPI job that is missing mpi4py from a machine that has no -# launcher at all, so an unrecognized launcher still reaches mpi_discovery as before. -MPI_RANK_ENV_VARS = ("OMPI_COMM_WORLD_RANK", "PMI_RANK", "PMIX_RANK", "MV2_COMM_WORLD_RANK", "SLURM_PROCID") +# World-size variables the MPI launchers export: OpenMPI, MPICH and Intel MPI, PMIx, MVAPICH, +# and Slurm's srun. The size rather than the rank, because a rank variable only says a launcher +# is present - `srun -n1` sets SLURM_PROCID for a single-task step, which is one process on one +# device and wants the fallback below, not an error about mpi4py. Every launcher listed sets its +# size variable alongside its rank one. +MPI_WORLD_SIZE_ENV_VARS = ("OMPI_COMM_WORLD_SIZE", "PMI_SIZE", "PMIX_SIZE", "MV2_COMM_WORLD_SIZE", "SLURM_NTASKS") -def in_mpi_job(): - """Whether an MPI launcher started this process.""" - return any(var in os.environ for var in MPI_RANK_ENV_VARS) +def in_multi_rank_mpi_job(): + """Whether a launcher started this process as one of several ranks.""" + for var in MPI_WORLD_SIZE_ENV_VARS: + try: + if int(os.environ[var]) > 1: + return True + except (KeyError, ValueError): + continue + return False def single_process_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True): diff --git a/tests/unit/comm/test_single_process_discovery.py b/tests/unit/comm/test_single_process_discovery.py index e37bf0ec7b2e..ba24007e3a9d 100644 --- a/tests/unit/comm/test_single_process_discovery.py +++ b/tests/unit/comm/test_single_process_discovery.py @@ -13,27 +13,49 @@ import pytest -from deepspeed.comm.comm import MPI_RANK_ENV_VARS, in_mpi_job, single_process_discovery +from deepspeed.comm.comm import MPI_WORLD_SIZE_ENV_VARS, in_multi_rank_mpi_job, single_process_discovery LAUNCHER_ENV = ("RANK", "WORLD_SIZE", "LOCAL_RANK", "MASTER_ADDR", "MASTER_PORT") @pytest.fixture def clean_env(monkeypatch): - for name in LAUNCHER_ENV + MPI_RANK_ENV_VARS: + for name in LAUNCHER_ENV + MPI_WORLD_SIZE_ENV_VARS + ("SLURM_PROCID", "OMPI_COMM_WORLD_RANK"): monkeypatch.delenv(name, raising=False) -def test_no_mpi_variables_is_not_an_mpi_job(clean_env): - assert in_mpi_job() is False +def test_no_mpi_variables_is_not_a_multi_rank_job(clean_env): + assert in_multi_rank_mpi_job() is False -@pytest.mark.parametrize("var", MPI_RANK_ENV_VARS) -def test_each_launcher_rank_variable_marks_an_mpi_job(clean_env, monkeypatch, var): +@pytest.mark.parametrize("var", MPI_WORLD_SIZE_ENV_VARS) +def test_each_launcher_size_variable_above_one_marks_a_multi_rank_job(clean_env, monkeypatch, var): """OpenMPI, MPICH/Intel MPI, PMIx, MVAPICH and srun each export a different one.""" - monkeypatch.setenv(var, "0") + monkeypatch.setenv(var, "4") - assert in_mpi_job() is True + assert in_multi_rank_mpi_job() is True + + +@pytest.mark.parametrize("var", MPI_WORLD_SIZE_ENV_VARS) +def test_a_launcher_reporting_one_task_is_not_a_multi_rank_job(clean_env, monkeypatch, var): + """`srun -n1` is a launcher and a single process at once; it wants the fallback, not an error.""" + monkeypatch.setenv(var, "1") + + assert in_multi_rank_mpi_job() is False + + +def test_a_rank_variable_alone_does_not_make_it_multi_rank(clean_env, monkeypatch): + """A rank says a launcher is present, not that the world is bigger than one.""" + monkeypatch.setenv("SLURM_PROCID", "0") + monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "0") + + assert in_multi_rank_mpi_job() is False + + +def test_an_unparseable_size_is_not_taken_as_multi_rank(clean_env, monkeypatch): + monkeypatch.setenv("SLURM_NTASKS", "") + + assert in_multi_rank_mpi_job() is False def test_single_process_discovery_fills_the_environment(clean_env): @@ -58,7 +80,7 @@ def test_single_process_discovery_leaves_what_the_caller_set(clean_env, monkeypa assert os.environ["WORLD_SIZE"] == "1" -def test_an_mpi_job_without_mpi4py_says_so(clean_env, monkeypatch): +def test_a_multi_rank_job_without_mpi4py_says_so(clean_env, monkeypatch): """Falling back to a single process there would silently run one rank of a many-rank job.""" import deepspeed.comm.comm as comm @@ -66,9 +88,35 @@ def no_mpi4py(*args, **kwargs): raise ImportError("No module named 'mpi4py'") monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) - monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "3") + monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", "4") with pytest.raises(ImportError, match="mpi4py"): comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) - assert "WORLD_SIZE" not in os.environ, "the environment must not be filled in for an MPI job" + assert os.environ.get("WORLD_SIZE") != "1", "the environment must not be filled in for a multi-rank job" + + +def test_a_single_task_slurm_step_without_mpi4py_falls_back(clean_env, monkeypatch): + """`srun -n1 python train.py` with no mpi4py: the case ebarkhordar raised on the PR.""" + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + monkeypatch.setenv("SLURM_PROCID", "0") + monkeypatch.setenv("SLURM_NTASKS", "1") + + reached = {} + real = comm.single_process_discovery + + def wrapped(*args, **kwargs): + reached["yes"] = True + return real(*args, **kwargs) + + monkeypatch.setattr(comm, "single_process_discovery", wrapped) + + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert reached.get("yes"), "a one-task step took the mpi4py error instead of the fallback" + assert os.environ["WORLD_SIZE"] == "1" From 90820bf93bf0d1a58f058966faa5c8d1817f59b4 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Mon, 7 Sep 2026 18:02:06 +0800 Subject: [PATCH 3/3] Separate no launcher from a launcher that reports no size PMIX_SIZE is not an environment variable - it is a size_t type name registered with PMIX_REGISTER_TYPE - so the entry was dead, and dropping it lost the coverage PMIX_RANK had been giving. A PMIx-direct launch, prterun or prun, sets PMIX_RANK and PMIX_NAMESPACE and no size at all, so reading only sizes turned a multi-rank job with no mpi4py into every rank coming up as world size 1. prterun -n1 and prterun -n4 are indistinguishable from the environment, so this refuses both rather than guessing: falling back silently splits a four-rank job into four world-size-1 runs, while refusing costs the one-rank case an error that names mpi4py and is recoverable. no launcher trains srun -n1 SLURM_NTASKS=1 trains srun -n4 SLURM_NTASKS=4 ImportError naming the size mpirun -n1 OMPI_..._SIZE=1 trains mpirun -n2 OMPI_..._SIZE=2 ImportError naming the size prterun PMIX_RANK, no size ImportError naming the missing size Also adds an end-to-end test for the bare environment. The helper tests all passed against a version of this that shadowed init_distributed's own world_size parameter and handed None to the backend, because none of them reached the backend; the new test does, and fails on that version with 'world_size must be an integer. None'. PMIX_SIZE identified as a dead entry by ebarkhordar, who enumerated the setenv literals in openpmix and PRRTE. Signed-off-by: alanhuangyoo --- deepspeed/comm/comm.py | 48 ++++++++---- .../comm/test_single_process_discovery.py | 77 +++++++++++++++---- 2 files changed, 93 insertions(+), 32 deletions(-) diff --git a/deepspeed/comm/comm.py b/deepspeed/comm/comm.py index 029fb0b76618..592a239d754d 100755 --- a/deepspeed/comm/comm.py +++ b/deepspeed/comm/comm.py @@ -846,11 +846,20 @@ def init_distributed(dist_backend: Optional[str] = None, try: mpi_discovery(distributed_port=distributed_port, verbose=verbose) except ImportError as err: - if in_multi_rank_mpi_job(): - raise ImportError("A multi-rank MPI job is running but mpi4py is not installed, so the " - "rank and world size cannot be discovered from it. Install mpi4py, or " - "set RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and MASTER_PORT " - "yourself.") from err + launcher_world_size = mpi_world_size_from_env() + if launcher_world_size is not None and launcher_world_size > 1: + raise ImportError( + f"A launcher reports a world size of {launcher_world_size} but mpi4py is not " + "installed, so " + "the rank cannot be discovered from it. Install mpi4py, or set RANK, WORLD_SIZE, " + "LOCAL_RANK, MASTER_ADDR and MASTER_PORT yourself.") from err + if launcher_world_size is None and launched_by_mpi(): + raise ImportError( + "A launcher started this process but does not report a world size in the " + "environment - PMIx launched directly, prterun or prun, sets a rank and no size - " + "so whether this is one rank of several cannot be determined without mpi4py. " + "Install mpi4py, or set RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and MASTER_PORT " + "yourself.") from err single_process_discovery(distributed_port=distributed_port, verbose=verbose) if cdb is not None and cdb.is_initialized(): @@ -866,23 +875,30 @@ def init_distributed(dist_backend: Optional[str] = None, cdb = TorchBackend(dist_backend, timeout, init_method, rank, world_size) -# World-size variables the MPI launchers export: OpenMPI, MPICH and Intel MPI, PMIx, MVAPICH, -# and Slurm's srun. The size rather than the rank, because a rank variable only says a launcher -# is present - `srun -n1` sets SLURM_PROCID for a single-task step, which is one process on one -# device and wants the fallback below, not an error about mpi4py. Every launcher listed sets its -# size variable alongside its rank one. -MPI_WORLD_SIZE_ENV_VARS = ("OMPI_COMM_WORLD_SIZE", "PMI_SIZE", "PMIX_SIZE", "MV2_COMM_WORLD_SIZE", "SLURM_NTASKS") +# World sizes the launchers export. The size rather than the rank, because a rank only says a +# launcher is present: `srun -n1` sets SLURM_PROCID for a single-task step, which is one process +# on one device and wants the fallback below rather than an error about mpi4py. +MPI_WORLD_SIZE_ENV_VARS = ("OMPI_COMM_WORLD_SIZE", "PMI_SIZE", "MV2_COMM_WORLD_SIZE", "SLURM_NTASKS") +# Ranks the launchers export, used only to tell "no launcher" from "a launcher that reports no +# size". PMIx launched directly - prterun, prun - is the case that needs it: it sets PMIX_RANK +# and PMIX_NAMESPACE and no size at all, so the world cannot be read from the environment. +MPI_RANK_ENV_VARS = ("OMPI_COMM_WORLD_RANK", "PMI_RANK", "PMIX_RANK", "MV2_COMM_WORLD_RANK", "SLURM_PROCID") -def in_multi_rank_mpi_job(): - """Whether a launcher started this process as one of several ranks.""" + +def mpi_world_size_from_env(): + """The launcher's world size, or None when no launcher reports one.""" for var in MPI_WORLD_SIZE_ENV_VARS: try: - if int(os.environ[var]) > 1: - return True + return int(os.environ[var]) except (KeyError, ValueError): continue - return False + return None + + +def launched_by_mpi(): + """Whether a launcher started this process, whatever world size it reports.""" + return any(var in os.environ for var in MPI_RANK_ENV_VARS) def single_process_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True): diff --git a/tests/unit/comm/test_single_process_discovery.py b/tests/unit/comm/test_single_process_discovery.py index ba24007e3a9d..5064dfb68c2d 100644 --- a/tests/unit/comm/test_single_process_discovery.py +++ b/tests/unit/comm/test_single_process_discovery.py @@ -13,49 +13,52 @@ import pytest -from deepspeed.comm.comm import MPI_WORLD_SIZE_ENV_VARS, in_multi_rank_mpi_job, single_process_discovery +from deepspeed.comm.comm import (MPI_RANK_ENV_VARS, MPI_WORLD_SIZE_ENV_VARS, launched_by_mpi, mpi_world_size_from_env, + single_process_discovery) LAUNCHER_ENV = ("RANK", "WORLD_SIZE", "LOCAL_RANK", "MASTER_ADDR", "MASTER_PORT") @pytest.fixture def clean_env(monkeypatch): - for name in LAUNCHER_ENV + MPI_WORLD_SIZE_ENV_VARS + ("SLURM_PROCID", "OMPI_COMM_WORLD_RANK"): + for name in LAUNCHER_ENV + MPI_WORLD_SIZE_ENV_VARS + MPI_RANK_ENV_VARS + ("PMIX_NAMESPACE", ): monkeypatch.delenv(name, raising=False) -def test_no_mpi_variables_is_not_a_multi_rank_job(clean_env): - assert in_multi_rank_mpi_job() is False +def test_a_bare_environment_reports_no_launcher_and_no_size(clean_env): + assert mpi_world_size_from_env() is None + assert launched_by_mpi() is False @pytest.mark.parametrize("var", MPI_WORLD_SIZE_ENV_VARS) -def test_each_launcher_size_variable_above_one_marks_a_multi_rank_job(clean_env, monkeypatch, var): - """OpenMPI, MPICH/Intel MPI, PMIx, MVAPICH and srun each export a different one.""" +def test_each_launcher_size_variable_is_read(clean_env, monkeypatch, var): + """OpenMPI, MPICH/Intel MPI, MVAPICH and srun each export a different one.""" monkeypatch.setenv(var, "4") - assert in_multi_rank_mpi_job() is True + assert mpi_world_size_from_env() == 4 @pytest.mark.parametrize("var", MPI_WORLD_SIZE_ENV_VARS) -def test_a_launcher_reporting_one_task_is_not_a_multi_rank_job(clean_env, monkeypatch, var): +def test_a_launcher_reporting_one_task_reports_one(clean_env, monkeypatch, var): """`srun -n1` is a launcher and a single process at once; it wants the fallback, not an error.""" monkeypatch.setenv(var, "1") - assert in_multi_rank_mpi_job() is False + assert mpi_world_size_from_env() == 1 -def test_a_rank_variable_alone_does_not_make_it_multi_rank(clean_env, monkeypatch): - """A rank says a launcher is present, not that the world is bigger than one.""" - monkeypatch.setenv("SLURM_PROCID", "0") - monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "0") +@pytest.mark.parametrize("var", MPI_RANK_ENV_VARS) +def test_a_rank_variable_marks_a_launcher_but_gives_no_size(clean_env, monkeypatch, var): + """A rank says a launcher is present, not how big the world is.""" + monkeypatch.setenv(var, "0") - assert in_multi_rank_mpi_job() is False + assert launched_by_mpi() is True + assert mpi_world_size_from_env() is None -def test_an_unparseable_size_is_not_taken_as_multi_rank(clean_env, monkeypatch): +def test_an_unparseable_size_falls_through_to_the_next_variable(clean_env, monkeypatch): monkeypatch.setenv("SLURM_NTASKS", "") - assert in_multi_rank_mpi_job() is False + assert mpi_world_size_from_env() is None def test_single_process_discovery_fills_the_environment(clean_env): @@ -120,3 +123,45 @@ def wrapped(*args, **kwargs): assert reached.get("yes"), "a one-task step took the mpi4py error instead of the fallback" assert os.environ["WORLD_SIZE"] == "1" + + +def test_a_launcher_that_reports_no_size_is_refused(clean_env, monkeypatch): + """PMIx launched directly sets PMIX_RANK and no size at all. + + `prterun -n4` and `prterun -n1` are indistinguishable from the environment, so falling back + would turn the four-rank case into four separate world-size-1 runs. Refusing costs the + one-rank case an error naming mpi4py, which is the recoverable half of that trade. + """ + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + monkeypatch.setenv("PMIX_RANK", "0") + monkeypatch.setenv("PMIX_NAMESPACE", "prterun-host-1234@1") + + with pytest.raises(ImportError, match="does not report a world size"): + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert "WORLD_SIZE" not in os.environ + + +def test_a_bare_environment_initializes_end_to_end(clean_env, monkeypatch): + """`python train.py` with nothing set, all the way through init_distributed. + + The helper tests above pass on a version of this that shadows `init_distributed`'s own + `world_size` parameter and hands `None` to the backend, because they never reach the + backend. This one does. + """ + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert os.environ["WORLD_SIZE"] == "1" + assert os.environ["RANK"] == "0"