Mfrank/py data staging

CLAUDE.md

@@ -0,0 +1,199 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## What this is
+
+`mlperf-common` is a grab-bag of utilities shared across NVIDIA's MLPerf benchmark
+submissions. Benchmarks consume it by adding
+`git+https://github.com/NVIDIA/mlperf-common.git` to their `requirements.txt`.
+
+Two mostly independent halves live here:
+
+* **Logging / profiling** (`mlperf_common/logging.py`, `scaleoutbridge.py`,
+  `frameworks/`, `callbacks/`) — wraps the official `mlperf_logging` package.
+* **Fast file I/O** (`mlperf_common/fileio/`, `client/`, `src/`) — O_DIRECT copy
+  and checksum tools, plus collective dataset staging onto node-local storage.
+
+## Commands
+
+```bash
+python3 tests/run_tests.py          # whole suite (stdlib only, no pytest, no GPU)
+python3 tests/test_pipeline.py      # one test file, directly
+
+# real hardware, from inside an allocation; not part of run_tests.py
+tests/cluster-selftest.sh /lustre/scratch/me/selftest /raid/scratch/me/selftest
+
+make -C src                         # build the C++ tools
+make -C src install prefix=/usr/local
+
+pip install .                       # installs the package + client/ scripts into bin/
+```
+
+There is no linter or formatter configured, and no CI in the repo.
+
+Tests deliberately run each file in a **separate interpreter**: each installs its
+own fake `torch` (`tests/stubs.py`) into `sys.modules` and patches module-level
+names in `datastage`. Don't collapse them into one process. See `tests/README.md`
+for what the stubs do and do not cover — notably not NCCL, not real CUDA events,
+not O_DIRECT itself.
+
+## Architecture
+
+### Logging stack
+
+`MLLoggerWrapper` (`mlperf_common/logging.py`) is the entry point. It takes a
+`CommunicationHandler` so the same wrapper works under either `torch.distributed`
+or MPI; `frameworks/base.py` defines the `CommunicationHandler` /
+`ProfilerHandler` interfaces and `frameworks/{pyt,mxnet,hugectr,base_mpi}.py`
+implement them per framework. Adding framework support means implementing those
+two interfaces, not touching the wrapper.
+
+`scaleoutbridge.py` layers profiling on top: `init_bridge(prof_handler,
+comm_handler, mllogger)` picks a bridge implementation from env vars —
+`TIME_TAGS` / `NVTX_FLAG` select `ScaleoutBridgeIterwise`, `EPOCH_PROF` selects
+`ScaleoutBridgeEpochwise`, and with none set you get the no-op
+`ScaleoutBridgeBase`.
+
+`callbacks/logging.py` is the Lightning/NeMo layer (`LoggingCallback`,
+`MLPerfLogger`, `StatsLogCallback`). It imports `lightning.pytorch` with a
+fallback to `pytorch_lightning`, and is knob-driven by env vars
+(`FORCE_SUCCESS_STATUS`, `REDUCE_TP`, `LOG_EVERY_N_BATCHES`, `RUN_N_ITERS`,
+`SEED`). This is the only part that pulls in Lightning and numpy.
+
+### fileio
+
+Three layers, with a deliberate dependency rule stated in
+`mlperf_common/fileio/__init__.py`: **only `datastage` may import torch.**
+`direct_io` and `copyplan` stay dependency-free so the single-node `client/`
+scripts don't drag in a training stack.
+
+* `direct_io.py` — aligned-buffer `pread`/`pwrite` with retry loops.
+* `copyplan.py` — source-tree walk and src→dst mapping. `plan_copy_operations`
+  raises `UnreadableEntries` listing *every* bad entry rather than dying on the
+  first.
+* `datastage.py` — collective staging (below).
+
+`BUFFER_ALIGN = 2 MiB` (the huge-page size) is the shared alignment constant
+across `fastcp`, `fastmd5`, and `datastage`.
+
+### datastage — the one thing worth reading before editing
+
+Stages a dataset from shared storage (Lustre) to node-local storage across a
+whole job. With W ranks = N nodes × L ranks/node:
+
+* The world splits into L process groups; group `l` holds every node's
+  `LOCAL_RANK == l` rank and owns slice `l` of each file.
+* Within group `l`, each of the N ranks reads a disjoint 1/N of that slice, and
+  an all-gather assembles the full slice on every node. Each rank writes slice
+  `l` locally.
+* Result: W disjoint readers on the shared FS, one fabric crossing per byte per
+  node, L concurrent all-gathers.
+
+`Topology` assumes **slurm's default block distribution**: node `i` holds ranks
+`i*L .. i*L+L-1`, so a rank's node is `RANK // L` and its slot is `RANK % L`.
+That is arithmetic, requiring no collective. It verifies the assumption locally
+— slurm reports RANK (`SLURM_PROCID`) and LOCAL_RANK (`SLURM_LOCALID`)
+independently, and they agree only under a block distribution, so a modulo
+comparison on each rank rejects `--distribution=cyclic`/`=arbitrary` and ragged
+`--ntasks-per-node` immediately.
+
+An earlier version derived the layout from gathered hostnames to support
+`--distribution=arbitrary`. That was never a real requirement — a mistake
+carried over from the C version — and it could not have worked anyway, since
+`MASTER_ADDR` is the first node in the nodelist and rank 0 must be there. Don't
+reintroduce discovery here.
+
+`group_ranks[l]` must stay **ascending**: `dist.new_group` sorts the list it is
+given and derives each member's group position from that order, while the
+drainer maps all-gather output position to `node_index`. Those agree only while
+the lists are sorted. `tests/test_topology.py` checks it.
+
+`FileLayout` keeps every offset and length aligned
+*except* the final range of the final slice, so O_DIRECT write padding can only
+ever run off the end of the file, where a closing `ftruncate` trims it.
+
+The pipeline (`Stager._run_pipeline`) moves buffers between a reader thread, the
+collective, a drain thread, and a writer pool through four queues, with CUDA
+events for ordering. Getting a handoff wrong does not raise — it writes a
+correctly-sized file containing wrong bytes. That is why the tests compare bytes,
+not exit status. If you change the pipeline, deliberately reintroduce a bug and
+confirm `test_pipeline.py` goes red.
+
+Device and pinned-host memory scale with node count (the window is
+`--buffer-size × N`), so `Stager.__init__` budget-checks against 60% of device
+memory and fails with a suggested `--buffer-size` rather than OOMing inside CUDA.
+
+Run it as one task per GPU. No wrapper needed:
+
+```bash
+srun --ntasks-per-node=${DGXNGPU} ... \
+    python3 -m mlperf_common.fileio.datastage -r "${SLOW_DATADIR}/${DATASET}" "${DATADIR}"
+```
+
+`--dry-run` outside any launcher (`dist_env` reports `source == "single"`)
+prints the copy plan without touching CUDA.
+
+### dist_env
+
+`mlperf_common/dist_env.py` (stdlib-only, top level so `affinity/` can use it
+too) derives torch's `env://` variables from SLURM or OMPI — the same
+translation `client/slurm2pytorch` does in bash, which is why datastage no
+longer needs the wrapper. Launching under the wrapper still works: the
+variables are already set and taken as given (`source == "preset"`).
+
+The two must stay behaviourally compatible. `tests/test_dist_env.py` pins that
+down, and the claim was checked by running the real script and `configure()`
+against the same pre-wrapper environment and diffing all seven variables.
+
+Two deliberate divergences, both refusing to guess where the script defaults:
+
+* **`LOCAL_WORLD_SIZE`** — `SLURM_NTASKS_PER_NODE` is only set when
+  `--ntasks-per-node` was passed, so `srun -N2 -n16` has no source for it and
+  the script's `:-1` turns two 8-GPU nodes into sixteen single-rank "nodes".
+  We fall back to `SLURM_TASKS_PER_NODE` (always set under srun) and refuse to
+  guess on a multi-rank job.
+* **`MASTER_ADDR`** — the script falls back to `127.0.0.1`, which its own
+  comment says "will fail for multinode", as a rendezvous that hangs to the
+  wall clock. We parse the first hostname out of slurm's compressed nodelist
+  (`dgx[001-004,007]` → `dgx001`, zero padding preserved — `dgx1` doesn't
+  resolve), and error immediately if that fails on a multi-node job.
+
+There is deliberately **no `PYTORCH_VERSION` gate**; the script has one because
+it wraps arbitrary commands in arbitrary containers.
+
+Validation splits: `dist_env` checks that the *sources* are present, agree, and
+are arithmetically coherent; `Topology` checks that the resulting *layout* is
+block-distributed. Don't duplicate one into the other.
+
+### client/ and src/
+
+`client/` holds scripts installed onto `PATH` by `setup.py`: `bindpcie` (NUMA/IB
+affinity binding), `mgpurun`, `slurm2pytorch` (derives PyTorch rendezvous env
+from SLURM), `fastcp` / `fastmd5` (threaded O_DIRECT copy and per-GB checksum),
+`dropcache`, plus log/telemetry shell helpers.
+
+**Don't delete `slurm2pytorch`.** datastage no longer needs it, but benchmarks
+outside this repo do, and it stays installed. `mlperf_common/dist_env.py` is the
+Python equivalent; the two must stay behaviourally compatible.
+
+These scripts import `direct_io` and `mlperf_common.fileio.copyplan` via a
+`sys.path` dance that works both for a pip install and for a source tree with
+`mlperf_common/` alongside `client/` (the "deploy the repo to a shared filesystem
+and run in place" pattern). `client/direct_io.py` is a compatibility shim
+re-exporting `mlperf_common.fileio.direct_io`. Copying individual scripts out of
+`client/` on their own does not work — preserve that when touching imports, and
+if you add a script remember to add it to `scripts=` in `setup.py`.
+
+`src/` is a separate, older C++ implementation (`fastcp`, `cp-into`,
+`alloc-empty-file-buffer`) sharing `cmdline.h`. It is not built or installed by
+`setup.py`; use the `Makefile` directly.
+
+## Notes
+
+* Commit subjects follow `area: lowercase imperative summary`.
+* New files carry the Apache 2.0 header with an NVIDIA copyright line.
+* The README's "Mount check" section documents `get-mount-info.sh` /
+  `verify-mounts.sh`, which no longer exist — `mlperf_common/mountcheck.py`
+  (`--initialize` / verify against `expected-mounts.csv`, sparse SHA256
+  fingerprints) replaced them. Trust the module over the README there.

IB-BINDING-NOTES.md

@@ -0,0 +1,154 @@
+# Notes: per-rank NIC binding, and whether MPI could replace NCCL in datastage
+
+Working notes for a follow-up session that has a real GPU node. Everything in
+"What the code does" is read off `client/bindpcie` and is solid; everything in
+"Predictions" is inference from reading and **has not been run on hardware**.
+The point of the on-node session is to settle the predictions.
+
+Delete this file once the questions are answered.
+
+## Why this matters
+
+`mlperf_common/fileio/datastage.py` uses NCCL, via torch, purely as a network
+transport. Trace the bytes: Lustre → host → **GPU → fabric → GPU** → host →
+NVMe. The data is host-resident at both ends and nothing computes on it, so the
+GPU round trip costs two PCIe crossings a host-side transfer would not pay.
+
+The justification in datastage's module docstring is that NCCL drives every NIC
+"without any hand-tuned per-cluster transport configuration". That is an
+argument about **NIC aggregation**, not about GPUs. If an MPI program can get
+one NIC per rank — eight ranks, eight NICs, host memory — it should aggregate
+the same fabric over a shorter path, and datastage could drop CUDA entirely.
+
+What that would buy, beyond the shorter path:
+
+* No torch, so no 20 GB container for a program that copies files. (Matt: the
+  container copy cost is real but livable, so this is a secondary benefit, not
+  the driver.)
+* The whole class of bug fixed in `9035aa4` and `4ab9a13` — drainer thread
+  device affinity, stream ordering, CUDA event semantics — stops existing,
+  because there are no CUDA events.
+
+What it risks: NCCL's out-of-the-box multi-NIC behaviour is genuinely good and
+uniform across systems; MPI quality varies by site. That is a measurement, not
+an argument, which is what this file is for.
+
+## What `bindpcie --ib=single` actually does
+
+All line numbers are `client/bindpcie`.
+
+* **Device list** (`:107-111`): `ibv_devinfo --list | tail -n+2 | cut -f2`,
+  in whatever order `ibv_devinfo` prints. `num_ibdevs` is that count.
+* **GPU count** (`:79`): `nvidia-smi -i 0 --query-gpu=count`, so it reflects
+  what is *visible* to the container.
+* **Guard** (`:184-187`): if `num_ibdevs > num_gpus` or
+  `num_gpus % num_ibdevs != 0`, print an error naming
+  `MELLANOX_VISIBLE_DEVICES` and **`exit 1`**. A hard failure, not a warning.
+* **Mapping** (`:189`):
+  ```bash
+  ibdev="${ibdevs[$(( local_rank * num_ibdevs / num_gpus ))]}"
+  ```
+* **Exports** (`:190-191`):
+  ```bash
+  export OMPI_MCA_btl_openib_if_include="${OMPI_MCA_btl_openib_if_include-$ibdev}"
+  export UCX_NET_DEVICES="${UCX_NET_DEVICES-$ibdev:1}"
+  ```
+
+Note `MELLANOX_VISIBLE_DEVICES` is only ever *mentioned*, in the error message.
+The script never reads it. It is an enroot/pyxis hook that filters which IB
+devices the container sees, so it acts on `ibv_devinfo`'s output upstream of
+this code.
+
+## Predictions to verify on hardware
+
+**1. The "near its GPU" claim is not implemented.** `--help` says `--ib=single`
+binds "each rank to a single IB device near its GPU", but the mapping is pure
+index arithmetic over `ibv_devinfo` order. There is no topology query anywhere
+in the IB path — contrast the CPU path, which does interrogate `nvidia-smi` and
+`lscpu`. Locality holds only if `ibv_devinfo` enumerates in GPU order.
+
+This is the finding that would make the flag *harmful* rather than merely
+useless: a mis-ordered list pins each rank to a NIC that may be across the root
+complex, which is worse than letting UCX choose by locality.
+
+With `num_ibdevs == num_gpus` the mapping reduces to `ibdevs[local_rank]`, so
+the whole question collapses to: **does `ibv_devinfo --list` order match GPU
+order on this platform?** That is directly checkable against
+`nvidia-smi topo -m`.
+
+**2. The guard probably passes under enroot, and probably would not bare.** A
+stock H100 node reports compute *and* storage NICs to `ibv_devinfo` — typically
+8 + 2 against 8 GPUs, so `10 > 8` and the script exits 1. Matt notes enroot does
+set `MELLANOX_VISIBLE_DEVICES`, which should filter to the compute NICs and make
+`num_ibdevs == 8`. Worth recording what `num_ibdevs` actually is in both
+contexts rather than assuming.
+
+**3. One of the two exported variables is dead.**
+`OMPI_MCA_btl_openib_if_include` targets the openib BTL, deprecated in OpenMPI
+4.0 and removed in 5.0. On anything modern it is ignored, and `UCX_NET_DEVICES`
+is doing all the work.
+
+**4. The UCX port is hardcoded** to `:1`. Correct for single-port cards, wrong
+for a dual-port card whose second port carries the traffic.
+
+**5. Diagnostics go to stdout.** `:183`, `:185`, `:186`, `:197` all use
+`echo "..." 2>&1`, which is a no-op for `echo` — the intent was `>&2`. Compare
+`:80`, which gets it right. So these errors and warnings land on **stdout** and
+will interleave with the wrapped program's own output. Minor, but it is one
+reason a broken `--ib=single` could go unnoticed for years.
+
+## What to check on the node
+
+```bash
+# 1. What does the container actually see?
+echo "MELLANOX_VISIBLE_DEVICES=${MELLANOX_VISIBLE_DEVICES:-unset}"
+ibv_devinfo --list
+nvidia-smi -i 0 --query-gpu=count --format=csv,noheader,nounits
+
+# 2. Does ibv_devinfo order match GPU order?  This is the crux of prediction 1.
+nvidia-smi topo -m          # look for PIX/PXB between GPU i and each NIC
+#    then compare against ibdevs[i] for i in 0..num_gpus-1
+
+# 3. Does the binding take at all?
+srun ... bindpcie --ib=single -- bash -c 'echo "$SLURM_LOCALID $UCX_NET_DEVICES"'
+```
+
+## Observing NIC usage — do not infer it from bandwidth
+
+Ground truth, per device, independent of what any library claims:
+
+```bash
+cat /sys/class/infiniband/mlx5_*/ports/1/counters/port_xmit_data
+```
+
+Sample before and after a transfer and diff. That answers "did all eight NICs
+move bytes" directly. (Units are 4-byte lanes, which does not matter for a
+did-it-move check.)
+
+For the other half — "did *this rank* select the NIC we told it to" —
+`UCX_LOG_LEVEL=info`, or `UCX_PROTO_INFO=y` on newer UCX, prints each rank's
+selected transports and devices. For OpenMPI, `--mca pml_ucx_verbose 10`.
+
+**Caveat for the comparison:** UCX defaults to `UCX_MAX_RNDV_RAILS=2`, so an
+*unbound* rank may already use two NICs for large transfers. Pinning each rank
+to exactly one device can therefore lower per-rank bandwidth while raising
+aggregate spread. Measure **aggregate across all eight ranks**, bound versus
+unbound — a single rank's number will mislead.
+
+## The question this feeds
+
+If eight ranks each on their own NIC, in host memory, get within ~10% of what
+datastage's current all-gather achieves (its per-file `GB/s` line, or the final
+`DONE ... GB/s`), then the GPU is pure cost in this program and an MPI transport
+is strictly better: shorter path, smaller container, and a large category of
+correctness surface deleted.
+
+If they do not, the NCCL path is earning its keep and the right move is instead
+to drop *torch* while keeping NCCL — `tests/stubs.py` is already an interface
+specification for exactly the surface datastage uses (`torch.empty`, `device`,
+`cuda.{Event,Stream,stream,current_stream,current_device,set_device,
+get_device_properties}`, `dist.{barrier,get_rank,new_group,
+broadcast_object_list,all_gather_into_tensor}`), and a ctypes NCCL backend is a
+second implementation of it. The rendezvous falls out too: datastage always has
+a shared filesystem, so rank 0 can drop the 128-byte `ncclUniqueId` in a file
+instead of standing up a `TCPStore`.