fix(#712): forward MLPS_/MLPSTORAGE_ env vars via mpirun -x

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}"

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