Fix #380:issue 380 dlio mpi rank by idevasena · Pull Request #396 · mlcommons/storage

idevasena · 2026-05-27T00:41:09Z

Fix #380: portable MPI rank in test_dlio_mpi.py + correct DLIO launch commands

Fixes #380

Summary

The integration test tests/integration/test_dlio_mpi.py, run as documented, crashed before doing anything useful, and the follow-on dlio_benchmark command it printed was itself invalid. This PR fixes both the test and the documented
commands, and adds a regression test.

Root causes

1. KeyError on non-OpenMPI launchers. The test selected its endpoint with int(os.environ['OMPI_COMM_WORLD_RANK']). That env var is OpenMPI-specific and is unset under plain python, MPICH/mpiexec, or Slurm srun, so the test died with a bare KeyError: 'OMPI_COMM_WORLD_RANK'. The sibling test_mpi_basic.py already does this correctly via comm.Get_rank() and a safe os.environ.get(...).

2. Docs launched an MPI program without MPI. tests/README.md told users to run python tests/integration/test_dlio_mpi.py, which yields a single rank (not a meaningful multi-endpoint test) and triggers the KeyError above. It also gave no guidance for the "not enough slots available" error that OpenMPI raises when
-np exceeds the host core count.

3. Invalid dlio_benchmark invocation. Both the test's printed "Next steps" and the trailing comment in configs/dlio/workload/multi_endpoint_mpi.yaml instructed dlio_benchmark --config multi_endpoint_mpi.yaml. dlio_benchmark is a Hydra app and rejects --config as an ambiguous abbreviation of --config-path / --config-name / --config-dir. The correct form selects the workload via a Hydra override.

Changes

tests/integration/test_dlio_mpi.py — select the endpoint from rank = comm.Get_rank() instead of the OpenMPI env var (read now only for display, via os.environ.get(..., 'not set')); add a size == 1 guard that prints the correct mpirun invocation; correct the printed "Next steps" command to dlio_benchmark workload=multi_endpoint_mpi --config-dir=configs/dlio.
tests/README.md — launch the test with mpirun -np 8 ... and document --oversubscribe for hosts with fewer cores than -np.
configs/dlio/workload/multi_endpoint_mpi.yaml — fix the how-to-run comment to the valid Hydra workload=... --config-dir=... form, with a note on why --config <file> fails.
tests/integration/test_issue_380_dlio_mpi_rank.py (new) — regression tests that reproduce the original KeyError, pin rank-based round-robin endpoint selection, and guard (via source checks) that neither the hard env-var subscript
nor the ambiguous dlio_benchmark --config pattern can return. No live MPI runtime or mpi4py required.

The corrected command form matches the repo's own harness in mlpstorage_py/benchmarks/dlio.py and the sibling configs/dlio/workload/llama3_8b_checkpoint.yaml.

Validation

pytest tests/integration/test_issue_380_dlio_mpi_rank.py -v — 15 passed.
python tests/integration/test_dlio_mpi.py (no launcher) — exits 0 with a launcher hint instead of crashing.
mpirun -np 8 python tests/integration/test_dlio_mpi.py — ranks 0–7 map round-robin to endpoints 0–3; rank N selects endpoint[N % 4].

Tests

smrc@dskbd029:~/Storage_Repo_Tests/storage_vdb375$ mpirun -np 8 uv run python tests/integration/test_dlio_mpi.py



============================================================

DLIO Multi-Endpoint Test with MPI

============================================================

Total MPI processes: 8

Endpoint assignment will be: rank % 4

============================================================



Rank  4: OMPI_COMM_WORLD_RANK=4 → endpoint[0] = http://endpoint1:9000

Rank  6: OMPI_COMM_WORLD_RANK=6 → endpoint[2] = http://endpoint3:9000

Rank  3: OMPI_COMM_WORLD_RANK=3 → endpoint[3] = http://endpoint4:9000

Rank  0: OMPI_COMM_WORLD_RANK=0 → endpoint[0] = http://endpoint1:9000

Rank  2: OMPI_COMM_WORLD_RANK=2 → endpoint[2] = http://endpoint3:9000

Rank  1: OMPI_COMM_WORLD_RANK=1 → endpoint[1] = http://endpoint2:9000

Rank  7: OMPI_COMM_WORLD_RANK=7 → endpoint[3] = http://endpoint4:9000

Rank  5: OMPI_COMM_WORLD_RANK=5 → endpoint[1] = http://endpoint2:9000



============================================================

✅ DLIO multi-endpoint MPI test completed!

============================================================



Next steps:

  1. Use configs/dlio/workload/multi_endpoint_mpi.yaml

  2. Run: mpirun -np 8 dlio_benchmark --config multi_endpoint_mpi.yaml

============================================================

smrc@dskbd029:~/Storage_Repo_Tests/storage_vdb375$ uv run pytest tests/integration/test_issue_380_dlio_mpi_rank.py -v

================================================================ test session starts =================================================================

platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0 -- /home/smrc/Storage_Repo_Tests/storage_vdb375/.venv/bin/python

cachedir: .pytest_cache

rootdir: /home/smrc/Storage_Repo_Tests/storage_vdb375

configfile: pyproject.toml

plugins: hydra-core-1.3.2

collected 11 items                                                                                                                                   



tests/integration/test_issue_380_dlio_mpi_rank.py::test_hard_env_subscript_was_the_crash 

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_hard_env_subscript_was_the_crashPASSED

tests/integration/test_issue_380_dlio_mpi_rank.py::test_safe_env_read_does_not_crash 

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_safe_env_read_does_not_crashPASSED

tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[0-0] 

        SETUP    F rank[0]

        SETUP    F expected_index[0]

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[0-0] (fixtures used: expected_index, rank)PASSED

        TEARDOWN F expected_index[0]

        TEARDOWN F rank[0]

tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[1-1] 

        SETUP    F rank[1]

        SETUP    F expected_index[1]

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[1-1] (fixtures used: expected_index, rank)PASSED

        TEARDOWN F expected_index[1]

        TEARDOWN F rank[1]

tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[2-2] 

        SETUP    F rank[2]

        SETUP    F expected_index[2]

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[2-2] (fixtures used: expected_index, rank)PASSED

        TEARDOWN F expected_index[2]

        TEARDOWN F rank[2]

tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[3-3] 

        SETUP    F rank[3]

        SETUP    F expected_index[3]

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[3-3] (fixtures used: expected_index, rank)PASSED

        TEARDOWN F expected_index[3]

        TEARDOWN F rank[3]

tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[4-0] 

        SETUP    F rank[4]

        SETUP    F expected_index[0]

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[4-0] (fixtures used: expected_index, rank)PASSED

        TEARDOWN F expected_index[0]

        TEARDOWN F rank[4]

tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[7-3] 

        SETUP    F rank[7]

        SETUP    F expected_index[3]

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[7-3] (fixtures used: expected_index, rank)PASSED

        TEARDOWN F expected_index[3]

        TEARDOWN F rank[7]

tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[8-0] 

        SETUP    F rank[8]

        SETUP    F expected_index[0]

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_rank_selects_endpoint_round_robin[8-0] (fixtures used: expected_index, rank)PASSED

        TEARDOWN F expected_index[0]

        TEARDOWN F rank[8]

tests/integration/test_issue_380_dlio_mpi_rank.py::test_script_has_no_hard_env_subscript 

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_script_has_no_hard_env_subscriptPASSED

tests/integration/test_issue_380_dlio_mpi_rank.py::test_script_selects_endpoint_from_rank 

        tests/integration/test_issue_380_dlio_mpi_rank.py::test_script_selects_endpoint_from_rankPASSED



================================================================= 11 passed in 0.02s =================================================================

smrc@dskbd029:~/Storage_Repo_Tests/storage_vdb375$ uv run python tests/integration/test_dlio_mpi.py



============================================================

DLIO Multi-Endpoint Test with MPI

============================================================

Total MPI processes: 1

Endpoint assignment will be: rank % 4

============================================================



NOTE: only 1 MPI process detected. This test is meant to run under an MPI launcher.

      Example: mpirun -np 8 python tests/integration/test_dlio_mpi.py

      If the host has fewer cores than -np, add --oversubscribe (or --use-hwthread-cpus).



Rank  0: OMPI_COMM_WORLD_RANK=not set → endpoint[0] = http://endpoint1:9000



============================================================

✅ DLIO multi-endpoint MPI test completed!

============================================================



Next steps:

  1. Use configs/dlio/workload/multi_endpoint_mpi.yaml

  2. Run: mpirun -np 8 dlio_benchmark --config multi_endpoint_mpi.yaml

============================================================

smrc@dskbd029:~/Storage_Repo_Tests/storage_vdb375$ mpirun -np 8 uv rub pytuv run^Cn tests/integration/test_dlio_mpi.py

smrc@dskbd029:~/Storage_Repo_Tests/storage_vdb375$ ^C

smrc@dskbd029:~/Storage_Repo_Tests/storage_vdb375$ uv run python tests/integration/test_mpi_basic.py

Rank 0/1: OMPI_COMM_WORLD_RANK=not set, OMPI_COMM_WORLD_SIZE=not set



============================================================

Testing Multi-Endpoint Distribution

============================================================

Rank  0 → endpoint[0] = http://endpoint1:9000

============================================================

✅ MPI test completed successfully!

============================================================

…y keys

…nt mpirun launch test_dlio_mpi.py selected its endpoint via int(os.environ['OMPI_COMM_WORLD_RANK']), which raises KeyError whenever the script is not launched by OpenMPI mpirun (plain python, MPICH, srun). Switch to the portable comm.Get_rank() (matching test_mpi_basic.py) and read the OMPI var only for diagnostics. Update tests/README.md to launch with mpirun and document --oversubscribe for under-provisioned hosts. Add regression test.

github-actions · 2026-05-27T00:41:19Z

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

idevasena added 2 commits May 26, 2026 03:22

Fix mlcommons#375: chunked vector loading uses globally-unique primar…

06a198d

…y keys

idevasena requested a review from a team May 27, 2026 00:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix #380:issue 380 dlio mpi rank#396

Fix #380:issue 380 dlio mpi rank#396
idevasena wants to merge 2 commits into
mlcommons:mainfrom
idevasena:fix/issue-380-dlio-mpi-rank

idevasena commented May 27, 2026

Uh oh!

github-actions Bot commented May 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

idevasena commented May 27, 2026

Fix #380: portable MPI rank in test_dlio_mpi.py + correct DLIO launch commands

Summary

Root causes

Changes

Validation

Tests

Uh oh!

github-actions Bot commented May 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant