diff --git a/mlpstorage_py/benchmarks/dlio.py b/mlpstorage_py/benchmarks/dlio.py index 81d80d11..1421bd72 100755 --- a/mlpstorage_py/benchmarks/dlio.py +++ b/mlpstorage_py/benchmarks/dlio.py @@ -599,11 +599,20 @@ def generate_dlio_command(self): mpi_btl=getattr(self.args, 'mpi_btl', 'auto')) # Forward env vars to ranks — OpenMPI does not propagate arbitrary # env vars to remote hosts by default; -x VAR opts each one in. + # (HPE Cray PALS mpiexec propagates env by default and takes the + # PALS-native branch in generate_mpi_prefix_cmd, so it doesn't + # need this list.) if 'DLIO_DROP_CACHES_TIMEOUT' in os.environ: mpi_prefix += " -x DLIO_DROP_CACHES_TIMEOUT" # S3/object-storage vars required for multi-host runs (storage #592). + # MLPS_/MLPSTORAGE_ vars required for multi-host checkpointing — + # notably MLPSTORAGE_CHECKPOINT_URI_SCHEME (storage #712 / #583), + # without which remote-rank object-store writes fail with + # "Unsupported URI scheme" once the model-parallel shards land off + # the head node. for _v in sorted(os.environ): if (_v.startswith('AWS_') or _v.startswith('S3DLIO_') + or _v.startswith('MLPS_') or _v.startswith('MLPSTORAGE_') or _v in ('STORAGE_LIBRARY', 'BUCKET')): mpi_prefix += f" -x {_v}" cmd = f"{mpi_prefix} {cmd}" diff --git a/tests/unit/test_mpi_s3_env_forwarding.py b/tests/unit/test_mpi_s3_env_forwarding.py index 04f6859c..6c6fb2a1 100644 --- a/tests/unit/test_mpi_s3_env_forwarding.py +++ b/tests/unit/test_mpi_s3_env_forwarding.py @@ -124,6 +124,46 @@ def test_bucket_forwarded(self, _mock, monkeypatch): ) +# --------------------------------------------------------------------------- +# BUG #712: MLPS_/MLPSTORAGE_ vars not forwarded to remote ranks +# --------------------------------------------------------------------------- + +class TestMlpsEnvForwardingMissing: + """MLPS_/MLPSTORAGE_-prefixed env vars must reach remote MPI ranks. + + Concrete trigger (issue #712): `MLPSTORAGE_CHECKPOINT_URI_SCHEME` is set + by `CheckpointingBenchmark.add_checkpoint_params` (see #583) so the + writer/reader factories can put the object-store scheme back on the + URI. Without `-x` forwarding, remote ranks never see it and the + `S3DLIOStorageWriter` fails with "Unsupported URI scheme" on every + mp_rank that landed on a non-head host. + + `MLPS_CHECKPOINT_MP_START_METHOD` is the same class of latent bug for + the streaming-checkpoint subprocess start method. + """ + + @patch(_MPI_PREFIX_PATCH, return_value=_FIXED_MPI_PREFIX) + def test_mlpstorage_checkpoint_uri_scheme_forwarded(self, _mock, monkeypatch): + """#712: MLPSTORAGE_CHECKPOINT_URI_SCHEME must reach remote ranks.""" + monkeypatch.setenv("MLPSTORAGE_CHECKPOINT_URI_SCHEME", "s3") + monkeypatch.delenv("DLIO_DROP_CACHES_TIMEOUT", raising=False) + cmd = _make_dlio_for_cmd().generate_dlio_command() + assert "-x MLPSTORAGE_CHECKPOINT_URI_SCHEME" in cmd, ( + "BUG #712: MLPSTORAGE_CHECKPOINT_URI_SCHEME not forwarded; " + "remote-rank object-store writes fail with 'Unsupported URI scheme'" + ) + + @patch(_MPI_PREFIX_PATCH, return_value=_FIXED_MPI_PREFIX) + def test_mlps_checkpoint_mp_start_method_forwarded(self, _mock, monkeypatch): + """#712 sibling: MLPS_CHECKPOINT_MP_START_METHOD must reach remote ranks.""" + monkeypatch.setenv("MLPS_CHECKPOINT_MP_START_METHOD", "forkserver") + monkeypatch.delenv("DLIO_DROP_CACHES_TIMEOUT", raising=False) + cmd = _make_dlio_for_cmd().generate_dlio_command() + assert "-x MLPS_CHECKPOINT_MP_START_METHOD" in cmd, ( + "BUG #712: MLPS_CHECKPOINT_MP_START_METHOD not forwarded" + ) + + # --------------------------------------------------------------------------- # Invariant: absent vars must NOT produce -x flags (no spurious forwarding) # --------------------------------------------------------------------------- @@ -145,6 +185,20 @@ def test_absent_s3dlio_var_not_forwarded(self, _mock, monkeypatch): cmd = _make_dlio_for_cmd().generate_dlio_command() assert "-x S3DLIO_CONNECT_TIMEOUT_SECS" not in cmd + @patch(_MPI_PREFIX_PATCH, return_value=_FIXED_MPI_PREFIX) + def test_absent_mlpstorage_var_not_forwarded(self, _mock, monkeypatch): + monkeypatch.delenv("MLPSTORAGE_CHECKPOINT_URI_SCHEME", raising=False) + monkeypatch.delenv("DLIO_DROP_CACHES_TIMEOUT", raising=False) + cmd = _make_dlio_for_cmd().generate_dlio_command() + assert "-x MLPSTORAGE_CHECKPOINT_URI_SCHEME" not in cmd + + @patch(_MPI_PREFIX_PATCH, return_value=_FIXED_MPI_PREFIX) + def test_absent_mlps_var_not_forwarded(self, _mock, monkeypatch): + monkeypatch.delenv("MLPS_CHECKPOINT_MP_START_METHOD", raising=False) + monkeypatch.delenv("DLIO_DROP_CACHES_TIMEOUT", raising=False) + cmd = _make_dlio_for_cmd().generate_dlio_command() + assert "-x MLPS_CHECKPOINT_MP_START_METHOD" not in cmd + # --------------------------------------------------------------------------- # Existing DLIO_DROP_CACHES_TIMEOUT forwarding must be preserved