From f169e4d37b90153a5c0fe35121773fa79820eb59 Mon Sep 17 00:00:00 2001 From: Curtis Anderson Date: Tue, 7 Jul 2026 17:47:57 -0700 Subject: [PATCH 1/2] =?UTF-8?q?test(#712):=20RED=20=E2=80=94=20MLPS=5F/MLP?= =?UTF-8?q?STORAGE=5F=20env=20vars=20not=20forwarded=20to=20remote=20MPI?= =?UTF-8?q?=20ranks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two RED tests + two invariants for the class of bug where mlpstorage-side env vars set on the driver process don't reach remote MPI ranks because generate_dlio_command's -x allowlist only covers AWS_/S3DLIO_/STORAGE_LIBRARY/BUCKET. Concrete trigger (#712): MLPSTORAGE_CHECKPOINT_URI_SCHEME is set by CheckpointingBenchmark.add_checkpoint_params (see #583) so the writer and reader factories can put the object-store scheme back on the URI that mlpstorage stripped for DLIO's ObjStoreLibStorage._preflight. Without -x forwarding, remote ranks never see the env var and S3DLIOStorageWriter raises 'Unsupported URI scheme' on every model-parallel shard that lands off the head node. MLPS_CHECKPOINT_MP_START_METHOD (streaming_checkpoint.py) is the same latent bug — user overrides on the driver don't reach remote ranks. The two invariants (absent MLPS_/MLPSTORAGE_ vars must not appear as -x flags) pass without the fix and guard against overzealous forwarding regressions. --- tests/unit/test_mpi_s3_env_forwarding.py | 54 ++++++++++++++++++++++++ 1 file changed, 54 insertions(+) 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 From e4d36c1852262339fdbaa2fc11141035a2860359 Mon Sep 17 00:00:00 2001 From: Curtis Anderson Date: Tue, 7 Jul 2026 17:48:18 -0700 Subject: [PATCH 2/2] fix(#712): forward MLPS_/MLPSTORAGE_ env vars via mpirun -x CheckpointingBenchmark.add_checkpoint_params (dlio.py:963-988) strips the URI scheme from checkpoint.checkpoint_folder to work around DLIO's ObjStoreLibStorage._preflight double-prepend (see #583) and stashes the scheme in os.environ[MLPSTORAGE_CHECKPOINT_URI_SCHEME]. The StorageWriterFactory and StorageReaderFactory reconstruct the URI via _normalize_checkpoint_uri, which reads that env var. OpenMPI does not propagate arbitrary env vars to remote ranks by default; generate_dlio_command's existing -x allowlist covered AWS_/S3DLIO_/STORAGE_LIBRARY/BUCKET (#592) but not MLPSTORAGE_. Consequence on multi-host runs (llama3-70b at 64 procs with TP=8 and --map-by node): the head node's ranks inherit the env directly and their scheme reconstruction works; every remote rank no-ops _normalize_checkpoint_uri and S3DLIOStorageWriter raises 'Unsupported URI scheme' on every mp_rank>0 shard. 56 of 64 writes fail, surviving ranks hang in the next collective, whole job wedges. Extend the allowlist to also match MLPS_ and MLPSTORAGE_ prefixes. Covers the concrete #712 trigger (MLPSTORAGE_CHECKPOINT_URI_SCHEME) plus the latent sibling MLPS_CHECKPOINT_MP_START_METHOD (writer start method knob in streaming_checkpoint.py) that would fail the same way if a user overrode it on the driver. The HPE Cray PALS mpiexec branch in generate_mpi_prefix_cmd propagates env by default and doesn't touch this allowlist, so no PALS-side change is needed. --- mlpstorage_py/benchmarks/dlio.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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}"