[RetinaNet / DLIO] High-concurrency read path is restricted to object storage; POSIX gets 1×, creating up to 64× I/O-concurrency asymmetry under CLOSED

[ByoungjunSeo]

Summary
For 1-sample-per-file formats such as JPEG (RetinaNet), the high-concurrency prefetch path (TorchIterableDatasetSimple, effective depth 64 × num_workers) is enabled only for object-storage backends, while POSIX / local filesystems fall back to the map-style path (1 × num_workers). Under the same CLOSED rules this gives object storage up to 64× higher I/O concurrency than POSIX, even when the underlying storage is equally capable — so the dataloader path, not the storage, can dominate the result.

Where in the code
dlio_benchmark/data_loader/torch_data_loader.py (verified on HEAD f16fe43, 2026-06-29):

• TorchIterableDatasetSimple (~L432-461): S3 → s3dlio.get_many(64 in-flight per worker); POSIX → ThreadPoolExecutor(64 workers); comment states "Effective pipeline depth: 64 × num_workers (vs 1 × num_workers before)."
• Selection gate (~L629-632):

python  use_simple_iterable_dataset = (
      self.format_type in _simple_iterable_formats
      and not use_rg_iterable_dataset
      and self._args.storage_type in _s3_types   # S3 / AISTORE / DIRECT_FS only
  )

where _s3_types = (StorageType.S3, StorageType.AISTORE, StorageType.DIRECT_FS). Plain POSIX is excluded and falls through to the map-style TorchDataset (1 in-flight per worker).

• A nearby comment notes the restriction exists for fork-safety of the prefetch thread — suggesting this is a technical limitation rather than a deliberate benchmarking choice.

Impact
RetinaNet reads many small (~315 KB) JPEGs, so it is highly sensitive to read concurrency. A POSIX/distributed-FS (e.g., CephFS) submission is capped at 1 in-flight per worker, while an object-storage submission gets up to 64, making CLOSED comparisons unfair.

Question / request
Is this asymmetry intended? If it is a fork-safety limitation, could the POSIX path also be given a high-concurrency prefetch option (e.g., the ThreadPoolExecutor(64) path) so POSIX and object-storage submissions compete on equal footing under CLOSED?

[FileSystemGuy]

Thanks for the detailed writeup — the asymmetry reproduces exactly as described and the root cause is confirmed. Tagging @russfellows for review of the plan below.

Confirmed against the pinned DLIO source

The gate at torch_data_loader.py:629-633 is exactly as reported:

_s3_types = (StorageType.S3, StorageType.AISTORE, StorageType.DIRECT_FS)
use_simple_iterable_dataset = (
    self.format_type in _simple_iterable_formats
    and not use_rg_iterable_dataset
    and self._args.storage_type in _s3_types  # LOCAL_FS excluded
)

LOCAL_FS falls through to TorchDataset(map-style, N workers) — 1 in-flight per worker. On the ~315 KB JPEG RetinaNet workload this is a real ~64× I/O-concurrency delta versus S3/AISTORE/DIRECT_FS. Not intentional — the inline comment at L622-628 explicitly frames this as a fork-safety workaround.

The code lives in DLIO_local_changes, not mlpstorage

storage's pin (pyproject.toml:148):

dlio-benchmark = { git = "https://github.com/mlcommons/DLIO_local_changes.git", rev = "f16fe43..." }

So the fix belongs in DLIO_local_changes; storage would only bump the pin once merged.

Root cause is narrower than the comment suggests

The _prefetch_buffered path (_local_fs_iterable_mixin.py:164) is already fork-safe — it creates a fresh ThreadPoolExecutor(max_workers=n_workers) inside a with block, so the pool is instantiated post-fork inside the DataLoader worker and torn down cleanly. That's not the blocker.

The actual blocker is the module-level pipelined-prefetch pool at _local_fs_iterable_mixin.py:79:

_PREFETCH_POOL = ThreadPoolExecutor(max_workers=1, thread_name_prefix="localfs_prefetch")

used at lines 293/299 to overlap chunk N+1 while yielding chunk N. Being module-level, its background worker thread is created in the parent process at import time and does not survive os.fork(). DataLoader workers inherit a half-dead executor and .submit() deadlocks. That's the classic ThreadPoolExecutor-across-fork foot-gun that motivated excluding LOCAL_FS from the gate.

Proposed fix (for DLIO_local_changes)

Any of the standard fork-safe patterns unlocks LOCAL_FS without giving up the 64× concurrency:

1. Lazy per-worker pool — replace the module singleton with a _get_prefetch_pool() helper keyed on os.getpid(), so each worker re-creates its own pool on first use after fork.
2. multiprocessing.util.register_after_fork — register a callback that nulls the child's inherited _PREFETCH_POOL handle and re-inits on first use.
3. Per-instance pool created in worker_init() — move the executor onto the reader instance, since TorchIterableDatasetSimple.worker_init() already runs post-fork on the worker process.

Option 3 is the cleanest structurally (no module-level shared state at all) and mirrors how S3 already scopes s3dlio calls per instance. Once applied, the gate collapses to just format-type checking and LOCAL_FS joins S3/AISTORE/DIRECT_FS in the high-concurrency path — RetinaNet POSIX vs. object-storage submissions then compete on equal footing under CLOSED.

Note: the S3 mixin has the same module-level pattern at _s3_iterable_mixin.py:80. It works today only because s3dlio's Rust runtime sidesteps the Python-thread fork issue, but it's the same latent shape — worth fixing in the same PR for consistency.

Plan

We're planning to implement this in DLIO_local_changes (option 3, unless review surfaces a reason to prefer another) as soon as @russfellows has had a chance to review the plan. Once merged there, storage will bump the DLIO pin and the concurrency asymmetry should be closed.

[russfellows]

I agree, in that I think we SHOULD definitely permit more parallelism in the filesystem data loader, and make object and file as equal as possible. I think the only way EITHER data loader can succeed for such small files / objects is to have VERY high parallelism. SO no limitations should be artificially imposed.

Moreover, the more equivalent we can make file and object, the better in my mind. If that means "Do Option 3" then so be it.