fix(#712): forward MLPS_/MLPSTORAGE_ env vars via mpirun -x#713
Merged
Conversation
…MPI ranks 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.
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.
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #712 — object-store checkpoint writes drop the
s3://scheme for every model-parallel shard that lands on a remote host.Root cause
CheckpointingBenchmark.add_checkpoint_params(mlpstorage_py/benchmarks/dlio.py:963-988) strips the URI scheme fromcheckpoint.checkpoint_folderto sidestep DLIO'sObjStoreLibStorage._preflightdouble-prepend (fix from #583) and stashes the scheme inos.environ['MLPSTORAGE_CHECKPOINT_URI_SCHEME'].StorageWriterFactoryandStorageReaderFactoryreconstruct the URI at dispatch time 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-xallowlist coveredAWS_*/S3DLIO_*/STORAGE_LIBRARY/BUCKET(#592) but notMLPSTORAGE_*. Under multi-host runs with--map-by node(the default ingenerate_mpi_prefix_cmdfor multi-host) and llama3-70b at 64 procs (TP=8, PP=1):MLPSTORAGE_CHECKPOINT_URI_SCHEME→_normalize_checkpoint_urino-ops →S3DLIOStorageWriterraisesUnsupported URI schemeUnder
--map-by node,mp_rank = my_rank % 8lines up with node index, so mp_rank_0 = my_rank ∈ {0, 8, 16, …} all live on the head node and succeed, while mp_rank_1..7 all live on remote nodes and fail. Matches the reporter's observation exactly. Single-node subset runs (8 procs, one host) are unaffected because env is inherited without-x.Change
Extend the
-xallowlist to also matchMLPS_andMLPSTORAGE_prefixes. Covers the concrete #712 trigger (MLPSTORAGE_CHECKPOINT_URI_SCHEME) plus the latent siblingMLPS_CHECKPOINT_MP_START_METHOD(streaming_checkpoint.py:88) that would fail the same way if a user overrode it on the driver shell.The HPE Cray PALS
mpiexecbranch (utils.py:580-601) propagates env by default and doesn't touch this allowlist, so no PALS-side change is needed.Symmetric coverage — read + write
Reads and writes both run inside the same DLIO invocation under one
mpirun, so one-xaddition covers both._normalize_checkpoint_uriis symmetrically wired intoStorageReaderFactory(fix from #641 / #649). This PR closes the multi-host read-path latent bug at the same time as the write-path reported one.Test plan
RED-first, following the convention from #641:
test(#712): RED) — 2 forwarding assertions + 2 absent-case invariants intests/unit/test_mpi_s3_env_forwarding.py. Verified failing againstorigin/main.fix(#712)) — 4-line allowlist extension inmlpstorage_py/benchmarks/dlio.py. Verified GREEN.fix/712:tests— 2740 passed, 1 skipped, 13 deselectedmlpstorage_py/tests— 828 passedvdb_benchmark/tests— 174 passedkv_cache_benchmark/tests— 238 passed, 2 deselected