mlcommons / mlcommons/storage

DLIO summary.json: host_memory_GB / host_cpu_count arrays are per-nnodes-slot, not per-host — positional access is broken

Open
#676 0 comments 0 reactions 1 assignee View on GitHub

@FileSystemGuy is already working on this.

Since Jul 5, 2026.

DLIO or mlpstorage Future
Dominant language
Python
Stars
203
Forks
67
Avg merge
20m
Merged PRs (30d)
8

Description

Follow-up to #669. The storage-side aggregation fix in #675 defenses against the current shape by aggregating via `sum()` / `sum() / num_hosts`, but this issue tracks the upstream DLIO fix: the arrays should be shaped so positional access is meaningful in the first place.

Where in DLIO

`dlio_benchmark/utils/statscounter.py` `:97-103` (host_memory_GB) and `:106-112` (host_cpu_count). Both use the same MPI `SUM`-Reduce pattern keyed on `self.MPI.node()`:

```python

host_memory_GB (line 97-103)

self.output['host_memory_GB'] = psutil.virtual_memory().total/1024./1024./1024
host_memory = np.zeros(self.MPI.nnodes())
host_memory_agg = np.zeros(self.MPI.nnodes())
if self.MPI.local_rank() == 0:
host_memory[self.MPI.node()] = self.output['host_memory_GB']
self.MPI.comm().Reduce(host_memory, host_memory_agg, op=MPI.SUM, root=0)
self.summary['host_memory_GB'] = list(host_memory_agg)
```

Actual shape observed on real hardware (from #669)

Reporter's cluster: 15 hosts, 30 accelerators, 2 ranks/host, 188.09 GiB/host verified via `/proc/meminfo`. Produced:

```
host_memory_GB = [376.18, 376.18, 376.18, 376.18, 376.18, 376.18, 376.18,
187.90, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

sum(array) = 2821 GiB ← correct cluster total (15 × 188.09)
array[0] = 376.18 GiB ← doubled (two hosts' local-rank-0 hit the same slot)
array[8..] = 0.0 ← empty slots (no host's local-rank-0 hit those)
```

The intended shape was `[188.09] × 15`. The actual shape is per-nnodes-slot where `self.MPI.node()` non-monotonically maps ranks to slots on multi-rank-per-host clusters.

Root cause candidate

On many MPI stacks (OpenMPI in particular), `self.MPI.node()` returns a per-rank node identifier whose numeric value does not line up cleanly with `0..nnodes-1` when there are multiple ranks per host. Two hosts can map to the same integer slot; other slots go untouched. The exact behavior depends on how DLIO's `MPI.node()` is derived from the launcher's rank/node bookkeeping. Anyone who has already dug into DLIO's `MPI` singleton is likely to know the right axis.

Proposed fix

Populate the per-host arrays by having each local-rank-0 append its own value, then gather the contributions to rank 0. Same result shape but no dependency on `self.MPI.node()` being a proper 0..nnodes-1 ordinal:

```python

On each local-rank-0, contribute this host's value; other ranks contribute None.

local_host_memory = self.output['host_memory_GB'] if self.MPI.local_rank() == 0 else None
per_host = self.MPI.comm().gather(local_host_memory, root=0)
if self.MPI.rank() == 0:
self.summary['host_memory_GB'] = [v for v in per_host if v is not None]
```

Resulting list: length `num_hosts` (unique physical hosts), one value per host, no zeros, no doubling. Positional access becomes meaningful.

Alternative (cleaner, more invasive): use `MPI_Comm_split_type(MPI_COMM_TYPE_SHARED)` to build a proper per-node communicator explicitly.

Applies equally to host_cpu_count

Same pattern at `utils/statscounter.py:106-112`. Fix both in the same DLIO PR to avoid a second follow-up.

Test suggestion

A regression test with a simulated `MPI` singleton where `self.MPI.node()` returns a non-monotonic mapping — e.g. ranks 0/1/2/3 all report `node=0` and ranks 4/5 report `node=5`, mimicking the reporter's observed shape — would pin the correct output.

Storage-side status

The storage-side fix in #675 aggregates via `sum()` / `sum() / num_hosts` and is defensively correct today AND after the DLIO fix lands. When this DLIO change ships, storage doesn't need to change further — the helper in `mlpstorage_py/submission_checker/dlio_summary_helpers.py` keeps working on the well-formed array too.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.