Bug fixes and performance enhancements: object storage, checkpointing, Parquet loading

[russfellows]

PR Summary: branch-3-0-2/bug-fixes-perf-enhancements

Branch: branch-3-0-2/bug-fixes-perf-enhancements
Base: main (mlcommons/storage)
Date: May 13, 2026
Tests: 127 passed, 0 failed (was 112 passed, 13 failed on clean main)

---

Issues Addressed

Of the 7 most recent open issues on mlcommons/storage, 6 are fixed by this branch.
Issue #369 was determined to be an environment/OpenMPI configuration problem with
no code fix applicable.

Issue	Title	Status	Fix location
#362	Training stuck at epoch 1, no NVMe reads	✅ Fixed	dlio_benchmark — reader_factory.py
#363	collect_cluster_info() missing required results_dir	✅ Fixed	benchmarks/base.py
#364	Flux AU limited by Parquet deserialization throughput	✅ Fixed	dlio_benchmark — reader_factory.py + s3dlio
#365	Checkpointing split-phase reports wrong operation counts	✅ Fixed	benchmarks/base.py
#367	reportgen crashes with AttributeError on Namespace.file	✅ Fixed	cli_parser.py
#369	orte_init failed — No permission (-17)	⚪ Not a code bug	OpenMPI environment/permissions issue
#371	--params storage.storage_type=direct_fs silently uses pagecache	✅ Fixed	dlio_benchmark — pytorch_checkpointing.py
#372	32 GB hard cap blocks large-memory runs	✅ Fixed (pending commit)	dlio_benchmark — utils/config.py

---

Commit History (above main)

Commit 1 — 022820b

Author: Devasena Inupakutika <devasena.i@samsung.com>
Message: cli_parser: guard --file/--object consolidation for non-benchmark subcommands
Fixes: #367
Cherry-picked from: PR #368

Problem: The reportgen, history, and lockfile subcommands do not call
add_storage_type_arguments(), so their Namespace objects have no .file or
.object attribute. The unconditional read and del in parse_arguments()
crashed with AttributeError.

Changes — mlpstorage_py/cli_parser.py:

• Guard the --file/--object consolidation block with
  if hasattr(parsed_args, "file") or hasattr(parsed_args, "object"):
• Use getattr(parsed_args, "file", False) instead of direct attribute access
• Replace bare del parsed_args.file / del parsed_args.object with a
  for _attr in ("file", "object"): if hasattr(...): delattr(...) loop
  so neither attribute is required to be present

Also includes new unit tests in tests/unit/test_cli.py covering the
parser behaviour for all subcommand types.

---

Commit 2 — 03765a2

Author: Devasena Inupakutika <devasena.i@samsung.com>
Message: Remove unwanted file
Cherry-picked from: PR #368

Removes a requirements.txt that was accidentally included in the
previous commit.

---

Commit 3 — 7e4245b

Author: Devasena Inupakutika <devasena.i@samsung.com>
Message: Fix #363: pass results_dir to collect_cluster_info
Fixes: #363
Cherry-picked from: PR #366

Problem: Benchmark._collect_cluster_information() called
collect_cluster_info() without the required positional argument
results_dir. This caused a TypeError at runtime:

WARNING: MPI cluster info collection failed: collect_cluster_info()
missing 1 required positional argument: 'results_dir'

The missing cluster info then propagated as None into reportgen,
causing a downstream crash:

[INVALID] None: Check check_num_files_train failed with error:
'NoneType' object has no attribute 'total_memory_bytes'

Changes — mlpstorage_py/benchmarks/base.py:

• Extract ssh_username and shared_staging_dir from self.args via
  getattr(..., None) before the call
• Pass results_dir=self.run_result_output (the benchmark's computed
  output directory) to collect_cluster_info()
• Pass shared_staging_dir=shared_staging_dir and
  ssh_username=ssh_username so SSH-based collection uses the correct
  credentials and staging path

Changes — mlpstorage_py/tests/test_benchmarks.py:

• Set benchmark.run_result_output = '/tmp/results/run-001' in the
  test fixture (previously missing; the call site needs this attribute)
• Update assert_called_once_with to expect results_dir,
  shared_staging_dir, and ssh_username
• Add TestCollectClusterInfoSignatureBinding regression test class (2
  new tests) that binds the actual kwargs against inspect.signature()
  of the real collect_cluster_info function, so future signature drift
  is caught at unit-test time rather than at runtime

---

Commit 4 — 2431011

Author: Russell Fellows
Co-authored-by: Devasena Inupakutika <devasena.i@samsung.com>
Message: Fix #365, #372: metadata override propagation, test suite fixes, env lock
Fixes: #365

Fix #365 — CLI override_parameters not reflected in metadata.json

Problem: The submission checker reads num_checkpoints_write /
num_checkpoints_read from metadata['parameters'] (the YAML
defaults). For split-phase submissions (write-only or read-only runs),
the correct counts are passed as CLI overrides such as:

override_parameters.num_checkpoints_write=10

These overrides landed in metadata['override_parameters'] only, which
the checker ignores. As a result, a 10-write + 10-read split-phase run
would aggregate to 20 writes + 20 reads and be marked INVALID.

Changes — mlpstorage_py/benchmarks/base.py:

• Add _apply_dotted_overrides(params, overrides) static method that
  deep-copies params and merges dotted-key overrides into the nested
  dict structure
• In the metadata property, call _apply_dotted_overrides() so
  metadata['parameters'] reflects the effective runtime configuration
• metadata['override_parameters'] is still emitted unchanged for a
  full audit trail

Note: PR #370 (crossmeta/zettalane) addresses the same root cause.
That PR is blocked pending CLA signature from @zettalane. This
implementation is carried independently; the two fixes are
functionally equivalent.

Fix — DLIOResultParser system info fallback

Problem: When a DLIO summary.json does not contain a system_info
block (e.g. runs from older DLIO versions), DLIOResultParser.parse()
returned None for ClusterInformation, breaking BenchmarkRun
validation.

Changes — mlpstorage_py/rules/models.py:

• DLIOResultParser.parse() now accepts an optional metadata kwarg
• When ClusterInformation.from_dlio_summary_json() returns None,
  fall back to metadata['cluster_information'] if present and
  reconstruct via ClusterInformation.from_dict()
• BenchmarkRun.__init__ passes the run's metadata object to
  parser.parse() to enable the fallback

Fix — 13 pre-existing test failures

mlpstorage_py/tests/test_cluster_collector.py (10 tests):

• All MPIClusterCollector(...) constructor calls and
  collect_cluster_info(...) call sites in failing tests were missing
  the now-required results_dir argument — added results_dir='/tmp'
  to all 10 affected call sites
• test_collector_returns_valid_data_without_error_marker: rewrote to
  use the current shared_staging_dir=tmpdir pattern instead of the
  obsolete UUID-based staging directory approach

mlpstorage_py/tests/test_rules.py (3 tests):

• TestBenchmarkRunSystemInfoFallback tests were failing with
  ValueError: No summary.json found in /tmp/test_run because they
  attempted real filesystem I/O
• Patched DLIOResultParser._load_summary and
  DLIOResultParser._load_hydra_configs to return in-memory mock data,
  removing the filesystem dependency

pyproject.toml / uv.lock

• Add [tool.uv] environments = ["sys_platform == 'linux'"] to
  pyproject.toml so uv lock does not attempt to resolve non-Linux
  platform markers (s3dlio only publishes Linux wheels)
• Regenerate uv.lock accordingly

---

dlio_benchmark Fixes (russfellows/dlio_benchmark — feat/parquet-dgen-streaming)

The following fixes are in the dlio_benchmark fork that is pinned by this
branch's pyproject.toml. They are already committed in the fork; issue #372
has an additional local change that is pending commit/push.

---

Fix #362 / #364 — Training stuck at epoch 1; Flux AU limited by CPU Parquet deserialization

Files: dlio_benchmark/reader/reader_factory.py,
dlio_benchmark/reader/parquet_reader_file_iterable.py (new),
dlio_benchmark/reader/parquet_reader_s3dlio.py
Commit: 1635b79 (feat: s3dlio-gen streaming, iterable dataloader, file iterable reader)

Issue #362 — Stuck at epoch 1, no NVMe reads:
reader_factory.py routed LOCAL_FS + Parquet to the legacy ParquetReader,
which calls pf.read_row_group() — full PyArrow deserialization on every read.
This is entirely CPU-bound and saturates the Python GIL, starving DLIO's
DataLoader workers of CPU time. Observed symptom: benchmark reaches
"Starting epoch 1" and then makes no measurable NVMe I/O while CPU pegs at
88-95%.

Issue #364 — Flux AU limited by per-process Parquet deserialization:
Same root cause. Even on a 192-vCPU Zen 4 machine, PyArrow's
read_row_group(use_threads=True) spawns additional decode threads per call.
Under DLIO's model (e.g. 4 MPI × 8 read_threads = 32 workers), hundreds of
threads contend on the GIL. AU on Skylake with data in tmpfs (zero I/O latency):
21% — storage is provably not the bottleneck; CPU decode is.

Fix: reader_factory.py now routes LOCAL_FS + Parquet to
ParquetReaderFileIterable — a new reader that performs raw byte-range reads
via a 64-thread ThreadPoolExecutor without any PyArrow decode. Data is
returned as raw bytes to the training loop. For S3/object storage, the s3dlio
Rust-based reader (ParquetReaderS3dlio) is used, which similarly bypasses
Python-side decode.

# Before (reader_factory.py):
# LOCAL_FS + Parquet → ParquetReader → pf.read_row_group() — full PyArrow decode

# After:
elif _args.storage_type in (StorageType.LOCAL_FS,):
    from dlio_benchmark.reader.parquet_reader_file_iterable import ParquetReaderFileIterable
    return ParquetReaderFileIterable(dataset_type, thread_index, epoch_number)

Result (from issue #364 testing, c6in.16xlarge, data on tmpfs):

Accelerators	use_threads	AU	Throughput	Result
4	True (before)	54.38%	~77 MB/s	❌ FAIL
4	False (workaround)	99.79%	141.80 MB/s	✅ PASS
8	False (workaround)	99.68%	283.07 MB/s	✅ PASS

The ParquetReaderFileIterable path goes further — no decode at all — giving
even better scaling on older CPU generations (Skylake, Cascade Lake) that lack
AVX-512 Parquet acceleration.

---

Fix #371 — --params storage.storage_type=direct_fs silently uses page cache

File: dlio_benchmark/checkpointing/pytorch_checkpointing.py
Commit: present in fork on branch feat/parquet-dgen-streaming

Problem: After PR #359 renamed the Python package from mlpstorage →
mlpstorage_py, one import path in dlio_benchmark was missed:

# Before (bug — old package name):
try:
    from mlpstorage.checkpointing import StreamingCheckpointing as _SC  # always fails
except ImportError:
    from dlio_benchmark.checkpointing.simple_streaming_checkpointing import (
        SimpleStreamingCheckpointing as _SC,   # silently falls back here
    )

SimpleStreamingCheckpointing ignores the backend='direct_fs' argument
entirely and uses plain open(path, "wb"). The result: when a user passes
--params storage.storage_type=direct_fs, page cache is never bypassed.
This was confirmed with free -h showing page cache growing during the write
phase and Lustre client cache filling up on a Lustre-backed mount.

Fix (one line):

# After:
from mlpstorage_py.checkpointing import StreamingCheckpointing as _SC

This ensures direct_fs checkpointing correctly uses O_DIRECT via s3dlio's
direct:// URI scheme, bypassing the page cache as intended.

---

Fix #372 — 32 GB hard cap blocks large-memory runs

File: dlio_benchmark/utils/config.py
Status: Modified locally in russfellows/dlio_benchmark — pending commit/push

Problem: BUDGET_MB was hard-coded to 32 * 1024 (32 GB). On hosts with
more than 32 GB of RAM this cap artificially constrains the number of DataLoader
workers. The error manifests as:

Exception: Memory budget exceeded: reader.read_threads=2 x comm_size=64 = 128
worker processes, estimated ~64 GB (hard cap: 32 GB). Reduce reader.read_threads
to at most 1 for this run.

On a 377 GB host trying to run 64 accelerators × 2 read_threads, the cap
prevents any run above 32 B200 ranks × 2 threads = 32 GB, limiting throughput
to ~2.3 GB/s regardless of storage capability (well below a Gen5 NVMe's 14 GB/s).

Fix:

# Before:
BUDGET_MB = 32 * 1024  # 32 GB hard cap

# After:
BUDGET_MB = psutil.virtual_memory().total // (1024 * 1024)  # actual host RAM

The budget now scales with actual installed RAM, which is the correct
upper bound for in-memory dataset caching.

---

Issue #369 — orte_init failed: No permission (-17) (No code fix)

Problem: OpenMPI orte_init fails with getting local rank failed → Returned value No permission (-17). This occurs when MPI processes are
launched as root without passing --allow-run-as-root to mpirun, or
when running inside a container with restricted Linux namespaces that
prevent OpenMPI's process management layer from initializing.

Assessment: This is an environment and OpenMPI configuration issue,
not a bug in mlpstorage or dlio_benchmark. The fix is to add
--allow-run-as-root to the mpirun invocation, or to configure the
container/namespace permissions to allow OpenMPI's process manager. No
code change is warranted.

---

Test Results

Before (clean main):  112 passed, 13 failed
After  (this branch): 127 passed,  0 failed

The net gain of 15 passing tests breaks down as:

• +13 pre-existing failures fixed (test_cluster_collector: 10, test_rules: 3)
• +2 new regression tests added by PR Fix #363: pass results_dir to collect_cluster_info #366 (TestCollectClusterInfoSignatureBinding)

[github-actions]

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

[russfellows]

Closing — opened in error, not ready for upstream review.