diff --git a/.review-findings.json b/.review-findings.json new file mode 100644 index 0000000..262598c --- /dev/null +++ b/.review-findings.json @@ -0,0 +1,157 @@ +{ + "level": "max", + "summary": "The 53 verified findings collapse to ~15 distinct defects, heavily duplicated across the staging pipeline. The most serious are two silent-data-corruption bugs in datastage (the drainer thread never calls torch.cuda.set_device, so its copy-completion events record on GPU 0 and writers race the D2H DMA on 7 of 8 ranks; and dist.new_group sorts its rank list, so all-gather slot order stops matching node_index under --distribution=arbitrary), plus three silent-data-loss/hang paths in planning (os.walk with no onerror drops unreadable subtrees, a missing destination-is-a-directory check makes plan_copy_operations discard sources[1:] or plan a directory as a 4 KiB file, and an unguarded os.stat in build_plan kills rank 0 while every peer blocks in the broadcast). Below those sit real but lower-severity issues: a chmod loop that spins forever on dest_root \"/\" and widens pre-existing directories to 0777, leaked .datastage.tmp files, an fastmd5 walk regression on directory symlinks, unclamped all-gather/D2H traffic, a memory-budget error whose suggested --buffer-size is itself over budget, no dedicated CUDA stream so the advertised overlap never happens, and test doubles that make exactly these corruption modes impossible to catch.", + "findings": [ + { + "file": "mlperf_common/fileio/datastage.py", + "line": 476, + "summary": "The drainer thread never calls torch.cuda.set_device, so the CUDA events it creates bind to device 0's stream instead of the rank's GPU, and `_write_one`'s `ready.synchronize()` (line 338) and `last_copy.synchronize()` (line 494) wait on nothing. [same root cause also at: mlperf_common/fileio/datastage.py:477, mlperf_common/fileio/datastage.py:477]", + "failure_scenario": "`main()` calls `torch.cuda.set_device(LOCAL_RANK)` (line 669) on the main thread only, but CUDA's current device is per-host-thread, so `drain_thread` (started at line 516) has current device 0. `chunk_view.copy_(recv_dev[...], non_blocking=True)` is issued on the rank's real device (the tensor's device guard picks it), but `torch.cuda.Event(blocking=True)` + `copied.record()` resolves `torch.cuda.current_stream()` against the *thread's* current device, so on any rank with LOCAL_RANK != 0 the event is lazily created and recorded on device 0's idle default stream. `_write_one` then calls `ready.synchronize()`, which returns immediately, and `direct_io.pwrite` reads the pinned chunk before the D2H copy has landed; `last_copy.synchronize()` likewise releases `recv_dev[dev_slot]` back to `recv_free_q` while copies are still reading it, so the next all-gather overwrites it. On an 8-GPU node, 7 of every 8 slices of every staged file get zeros or stale bytes \u2014 a file of exactly the right length with wrong contents, no exception, and training runs on the corrupt dataset. `tests/stubs.py` makes `FakeEvent.synchronize()` a no-op, so `test_pipeline.py` cannot go red on this.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 199, + "summary": "`dist.new_group(ranks)` sorts its rank list, so a member's group rank is its position in sorted order, not its position in `group_ranks[l]` \u2014 but the drainer maps all-gather output position i to `layout.segment(i, ...)` assuming position == node_index. [same root cause also at: mlperf_common/fileio/datastage.py:199]", + "failure_scenario": "Confirmed against torch/distributed/distributed_c10d.py (`ranks = sorted(ranks)` then `group_rank = ranks.index(global_rank)`). `Topology` indexes `group_ranks[local_rank]` by `node_index` (order of first hostname appearance in the `all_gather_object` result), but that list is only sorted when RANK happens to be block/cyclic distributed. Under `--distribution=arbitrary` with hostfile `A B B A`: identities are rank0=(A,0), rank1=(B,0), rank2=(B,1), rank3=(A,1); `group_ranks[1] = [3, 2]`, which `new_group` sorts to `[2, 3]`, so group rank 0 is rank 2 on host B (node_index 1) and group rank 1 is rank 3 on host A (node_index 0). `all_gather_into_tensor` therefore lays B's sub-shard at window offset `0 * piece` and A's at `1 * piece`, while the drainer writes offset `first+index` to `layout.segment(first+index, round_index)` (line 482). Every file's LOCAL_RANK-1 slice is written with the two nodes' sub-shards swapped: right-sized file, wrong bytes, no error \u2014 precisely the mis-grouping corruption `Topology` was written to prevent.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/copyplan.py", + "line": 82, + "summary": "os.walk is called with the default onerror=None, so a directory whose contents cannot be listed is silently skipped instead of raising UnreadableEntries \u2014 the exact silent-partial-copy failure this module was written to prevent. [same root cause also at: mlperf_common/fileio/copyplan.py:82, mlperf_common/fileio/copyplan.py:82, mlperf_common/fileio/copyplan.py:82]", + "failure_scenario": "Verified locally: a tree src/ containing good/a.bin and secret/b.bin where secret/ is mode 000 makes list_relative_files('src') return ['good/a.bin'] and plan_copy_operations return a single job, with no exception. If root itself is unlistable the result is [] \u2014 also silent. In production, one dataset subdirectory with a bad mode / EIO / stale Lustre handle makes `datastage -r $SLOW_DATADIR/$DATASET $DATADIR` stage a partial dataset and exit 0. Because fastmd5 now walks through the same function, checksumming the staged tree omits the same files, so verification passes too, and training runs on a silently incomplete dataset.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 649, + "summary": "`parse_args` validates sources but never checks that the destination is an existing directory, while the shared `plan_copy_operations` silently falls back to a single-file copy of `sources[0]` when it is not (copyplan.py:112) \u2014 a precondition fastcp enforces at its own call site and datastage does not. [same root cause also at: mlperf_common/fileio/copyplan.py:112, mlperf_common/fileio/datastage.py:649, mlperf_common/fileio/datastage.py:649, mlperf_common/fileio/copyplan.py:113, mlperf_common/fileio/copyplan.py:112]", + "failure_scenario": "Reproduced: `datastage -r /lustre/dataset /raid/scratch/dataset --dry-run` with the destination not yet created prints one job `\"/lustre/dataset -> /raid/scratch/dataset (4096 bytes)\"` \u2014 the whole tree collapses to a single bogus job for the source *directory*. In a real run `open_maybe_direct(src_dir, O_RDONLY)` gets EINVAL for O_DIRECT on a directory, which line 130 swallows, so it falls back to a buffered directory open; the reader thread then dies with `preadv failed with error: [Errno 21] Is a directory` while every peer sits in the collective until the NCCL watchdog fires, leaving a stray `.datastage.tmp.`. With more than one source (`datastage a.bin b.bin MISSING_DIR`) `sources[1:]` are dropped with no message and the job exits 0 having staged one of two files. `client/fastcp` blocks both cases in `parse_and_validate_args` (lines 172-179); datastage's `parse_args` omits the equivalent check.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 572, + "summary": "The os.stat() call that builds the broadcast payload sits inside a try that only catches UnreadableEntries, so any other rank-0 planning error escapes before the broadcast and parks every other rank until the watchdog fires \u2014 the precise hang the comment above it says must not happen. [same root cause also at: mlperf_common/fileio/datastage.py:572, mlperf_common/fileio/datastage.py:572]", + "failure_scenario": "plan_copy_operations already stat'd every source, and this re-stats each one to pick up st_mtime_ns. If a file is unlinked or a Lustre OST goes away between the two stats, os.stat raises FileNotFoundError/OSError, which is not UnreadableEntries; build_plan propagates it, rank 0 runs `finally: dist.destroy_process_group()` and exits, while ranks 1..W-1 are already blocked in dist.broadcast_object_list. They hang until the NCCL watchdog timeout (10+ minutes) and the job dies with an opaque collective-timeout message instead of \"cannot stat \", burning the allocation.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "client/fastmd5", + "line": 111, + "summary": "Replacing `Path.rglob('*')` with `copyplan.list_relative_files` (`os.walk(followlinks=True)`) removed fastmd5's immunity to directory-symlink cycles; it now aborts the whole run with zero checksum output. [same root cause also at: client/fastmd5:116]", + "failure_scenario": "Verified on a tree `data/{f1.bin, sub/{f2.bin, up -> ..}}`: the old `Path.rglob('*')` walk (which does not descend into symlinked directories) enumerated 3 entries and fastmd5 printed checksums for the real files. The new walk descends `up` ~40 levels until the kernel returns ELOOP, and `list_relative_files` raises `UnreadableEntries: 1 unreadable entry: data/sub/up/sub/up/.../up: Too many levels of symbolic links`, which `client/fastmd5:113` turns into `sys.exit(f\"fastmd5: {exc}\")`. So a dataset containing any self- or ancestor-referential directory symlink (`latest -> .`, `current -> ..`) can no longer be checksummed at all \u2014 fastmd5 exits 1 having printed not one line, leaving no way to verify a staged copy of that tree.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 406, + "summary": "_chmod_parents loops forever when dest_root is \"/\", because os.path.dirname(\"/\") returns \"/\" and the loop condition stays true. [same root cause also at: mlperf_common/fileio/datastage.py:406]", + "failure_scenario": "`datastage a.bin /` makes dest_root == \"/\" and dst == \"/a.bin\". _chmod_parents starts at path = \"/\", which satisfies `path == self.dest_root`, chmods \"/\" to 0777, then sets path = os.path.dirname(\"/\") == \"/\" and repeats \u2014 an unbounded loop that keeps chmod'ing the container root while every peer rank blocks in the dist.barrier on the next line, so the job hangs until the wall clock limit rather than failing.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 408, + "summary": "_chmod_parents chmods every directory from the file's parent up to and including dest_root, not just the ones it created, so a pre-existing destination tree is silently widened to --chmod (0777 by default).", + "failure_scenario": "With a pre-existing destination root created as `mkdir -m 0750 /raid/scratch/ds && mkdir -m 0750 /raid/scratch/ds/a`, running `datastage -r SRC /raid/scratch/ds` (default --chmod 0777) walks up from dirname(dst) and chmods every level whose path is dest_root or below it. Verified with a stubbed Stager: dest=0o750 and dest/a=0o750 before, dest=0o777 and dest/a=0o777 after. The docstring says the function exists to fix 'the directories we created' because os.makedirs applies the umask, but the loop's terminating condition (`path.startswith(dest_root + os.sep) or path == dest_root`) makes no distinction between directories datastage created and directories that were already there. Under --container-remap-root the process is root, so the chmod always succeeds and there is no error to notice.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 343, + "summary": "The .datastage.tmp. file is preallocated to the full source size and is only removed by the success-path os.rename; no failure path unlinks it, and the job-ID suffix means each retry orphans a fresh full-size copy.", + "failure_scenario": "stage_file creates `{dst}.datastage.tmp.{SLURM_JOB_ID}` and immediately `os.ftruncate(fd, size)`, then writes real data into it. The only thing that ever removes it is `os.rename(tmp, dst)` at the end of the success path; `grep -n 'unlink|os.remove'` over datastage.py returns nothing. If staging dies partway (NCCL error, ENOSPC, a rank raising out of _run_pipeline), main() prints a diagnostic and re-raises without cleanup, leaving a partially written temp file of up to `size` bytes on every node's local scratch. Because the name is scoped to SLURM_JOB_ID, a resubmitted job gets a new suffix rather than reusing/truncating the old file, so a resubmit loop against a flaky fabric accumulates one orphaned near-full copy of the dataset per attempt until node-local NVMe fills -- and the next attempt then fails for a different reason (ENOSPC) than the original. The tests show the authors care about this (stage_once asserts on leftover `.datastage.tmp.` files) but only ever exercise the success path.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 537, + "summary": "The all-gather always transfers a full --buffer-size per rank per round regardless of how many bytes the round actually holds, so every file costs buffer_size x node_count of fabric traffic even when it is a few KiB. [same root cause also at: mlperf_common/fileio/datastage.py:537, mlperf_common/fileio/datastage.py:474]", + "failure_scenario": "FileLayout gives shard = align_up(ceil(slice_len/N), 2MiB) >= 2MiB whenever a slice is non-empty, so rounds == 1 for any file below buffer_size, and send_dev is always `piece` bytes wide. Staging a 1 KiB file at N=2048 nodes with the default -b 8M issues an all_gather that delivers 8 MiB x 2048 = 16 GiB into every node's recv_dev, and the drainer then D2H-copies the whole window (nbytes = count * layout.piece is unconditional at line 474; only the pwrite skips zero-length segments) before discarding all but 1 KiB. For a dataset of 100k small files this moves petabytes of NCCL and PCIe traffic to stage a few GB, plus two world barriers per file \u2014 datastage ends up far slower than the rsync path it replaces on small-file datasets.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 277, + "summary": "The `--buffer-size` value suggested by the device-memory budget error is itself still over budget, so following the error's own advice reproduces the identical error; the formula it replaced (`align_down(budget // node_count)` against a `window > budget` check) was exact. [same root cause also at: mlperf_common/fileio/datastage.py:277, mlperf_common/fileio/datastage.py:277, mlperf_common/fileio/datastage.py:277]", + "failure_scenario": "Reproduced against the repo's own stubs with `total_memory` = 79.65 GiB (H100 80GB) and 64 nodes: `-b 512M` is rejected with \"needs 64.5 GiB of device memory, over the 47.8 GiB budget. Lower --buffer-size to at most 382M.\" Re-running with `-b 382M` prints \"needs 48.1 GiB ... Lower --buffer-size to at most 382M.\" \u2014 byte-identical advice. `-b 380M` is also rejected. The check is `piece + RECV_DEV_SLOTS*node_count*piece > budget` (i.e. `piece*(2N+1)`), but `per_window` divides only by `RECV_DEV_SLOTS*node_count` (`2N`), dropping the `+1`; the correct answer here is 378M. Every node count up to ~64 always mis-suggests (I checked N=16/32/64 on 80/94/141 GB devices \u2014 all STILL REJECTED), so an operator burns repeated multi-node job allocations, each dying in `Stager.__init__` before any staging, with no workable value ever offered.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 475, + "summary": "The drainer's device-to-host copies are issued on the same default stream as the H2D copy and the collective, so the overlap the design is built around never happens.", + "failure_scenario": "The comment at line 444 claims \"the collective for round k+1 can overlap the copy-back and writes for round k\", but chunk_view.copy_(recv_dev[...], non_blocking=True) is enqueued on the default stream from the drain thread while the main thread enqueues H2D(k+1) and all_gather(k+1) on that same stream; single-stream work serialises in enqueue order (and ProcessGroupNCCL makes its internal stream wait on the default stream), so the copy-back and the next collective run back to back. All the machinery of RECV_DEV_SLOTS=2, the drain thread and the CUDA events buys only host-side write overlap; staging runs at collective+copyback serialised instead of max(). A dedicated torch.cuda.Stream for the drainer, with the existing events used as cross-stream waits, is what delivers the intended overlap.", + "category": "cleanup", + "verdict": "CONFIRMED" + }, + { + "file": "client/direct_io.py", + "line": 35, + "summary": "The shim's sys.path fallback cannot recover when a stale `mlperf_common` is already importable, because the retry re-resolves through the package object cached in sys.modules whose __path__ still points at the old install.", + "failure_scenario": "Verified: with a stub `mlperf_common` package (no `fileio` subpackage) on sys.path -- i.e. a container that pip-installed an older mlperf-common from a benchmark's requirements.txt -- a caller that only puts client/ on sys.path and does `import direct_io` fails with `ImportError: No module named 'mlperf_common.fileio'. direct_io lives in the mlperf_common package. Either install mlperf-common, or keep this script in a tree with mlperf_common/ alongside client/`, even though the script is already in exactly such a tree. The first `from mlperf_common.fileio.direct_io import *` leaves `mlperf_common` in sys.modules bound to the site-packages location, so inserting the repo root at sys.path[0] changes nothing on the retry. The user is told to do the one thing they have already done, and the compatibility path the shim exists to serve is dead.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "tests/test_pipeline.py", + "line": 90, + "summary": "The fake all-gather ignores its `inp` argument and re-reads each segment from the source fd, so all 8 multi-node pipeline cases pass even when the reader produces completely wrong bytes.", + "failure_scenario": "Verified empirically: wrapping `Stager._read_piece` so it overwrites everything it read with 0xEE and running `python3 tests/run_tests.py` still reports 0 failures for all 8 MULTINODE configurations (only the 7 single-node cases go red), because `fake_all_gather` fills the window from `os.pread(STATE['fd'], ...)` rather than from the buffer the reader filled. So the read path -- `_read_piece`'s offset math, its multi-threaded split, send-slot recycling and the H2D copy -- is unverified in every multi-node case; instrumentation shows the threaded branch of `_read_piece` is entered exactly once in the whole suite, and that one call is in a multi-node case whose result is discarded. A regression that reads the wrong offsets or recycles a send slot before the H2D copy lands ships with a green suite and corrupts the staged dataset on the cluster.", + "category": "correctness", + "verdict": "CONFIRMED" + }, + { + "file": "tests/stubs.py", + "line": 106, + "summary": "The fake torch.distributed omits `all_gather_object` and `new_group` (and torch omits `tensor`), so `Topology` and `build_plan` -- the pieces CLAUDE.md flags as silently corrupting -- can never be constructed by any test.", + "failure_scenario": "`stubs.install()` defines only barrier/get_rank/broadcast_object_list/all_reduce/all_gather_into_tensor, so instantiating `ds.Topology` raises AttributeError ('module torch.distributed has no attribute all_gather_object') and `ds.build_plan` raises AttributeError on `torch.tensor`. Both test files consequently hand-roll their own `Topology` class (test_layout.py:47, test_pipeline.py:95) and never touch the real one. Reintroduce the exact bug CLAUDE.md warns about -- grouping nodes by `rank // local_world_size` instead of by reported hostname -- and `python3 tests/run_tests.py` still prints 'all tests passed', while a real `--distribution=arbitrary` run silently stages a corrupt dataset. The same holds for build_plan's cross-rank size/mtime consistency check.", + "category": "correctness", + "verdict": "CONFIRMED" + } + ], + "refuted": [ + { + "file": "client/fastmd5", + "line": 111, + "summary": "Replacing fastmd5's rglob walk with copyplan.list_relative_files (os.walk with followlinks=True) makes it descend symlinked directories, so a self-referential symlink now loops forever instead of finishing." + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 660, + "summary": "The copy plan is printed by two separate code paths that must be kept in step by hand." + }, + { + "file": "mlperf_common/fileio/__init__.py", + "line": 1, + "summary": "The commit that introduced this package uses a subject that does not follow the repo's commit-message convention." + }, + { + "file": "mlperf_common/fileio/datastage.py", + "line": 245, + "summary": "FileLayout's documented alignment invariant silently depends on an unenforced precondition (piece % align == 0); nothing in FileLayout or Stager asserts it and test_layout.py only ever passes multiples of BUFFER_ALIGN." + } + ], + "stats": { + "level": "max", + "finders": 6, + "candidates": 57, + "verifierAgents": 40, + "verified": 57, + "refuted": 4, + "reported": 15 + } +} \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..828400e --- /dev/null +++ b/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. diff --git a/IB-BINDING-NOTES.md b/IB-BINDING-NOTES.md new file mode 100644 index 0000000..78400c8 --- /dev/null +++ b/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`. diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md new file mode 100644 index 0000000..bb165d2 --- /dev/null +++ b/REVIEW-FINDINGS.md @@ -0,0 +1,427 @@ +# Code review findings — branch `mfrank/py-data-staging` + +Working file, not part of the deliverable. Delete it when the list is worked +through. Raw verifier output is in `.review-findings.json` next to this file. + +Reviewed diff: `git diff main...HEAD` (merge-base `c02c93c`, 13 files, ++1766/-212). Review ran at `max` effort: 49 agents, 53 raw findings deduplicated +to the 15 below, each independently verified. Verdicts were `CONFIRMED` across +the board, but see the premise correction immediately below — it demotes one of +them, and the verifiers did not know it. + +## Premise correction (2026-07-31) + +**datastage does not need to support `srun --distribution=arbitrary`.** The +review assumed it did, because `Topology`'s comments say so. That requirement +was a mistake by the original author of the C version and was carried into the +Python port unexamined; it is not a real constraint. + +Consequences: + +* **F2 is demoted** from live corruption to latent fragility — its repro needs a + non-ascending `group_ranks[l]`, which block/cyclic distribution never + produces. Worth an assert, not a panic. +* `Topology`'s `all_gather_object` hostname derivation existed only to serve + this non-requirement. **Removed 2026-07-31** in favour of block-distribution + arithmetic with a local consistency check — see F2. +* `CLAUDE.md` stated the arbitrary-distribution rationale as a design invariant. + Corrected in the same change that created this file. + +## Status legend + +`[ ]` open · `[x]` fixed · `[~]` intentionally skipped · `[-]` no change needed + +--- + +## Tier 1 — silent data corruption + +- [x] **F1 · `mlperf_common/fileio/datastage.py:476` · set_device missing on drainer thread** — fixed, see "F1 fix" below + + The drain thread never calls `torch.cuda.set_device`. CUDA's current device is + per-host-thread, so on any rank with `LOCAL_RANK != 0` the thread sits on + device 0. The `copy_` itself lands on the right device via the tensor's device + guard, but `torch.cuda.Event()` + `record()` resolve `current_stream()` against + the *thread's* device, so the event records on device 0's idle default stream. + `_write_one`'s `ready.synchronize()` (:338) and `last_copy.synchronize()` + (:494) then return immediately. + + Effect: `pwrite` reads the pinned chunk before the D2H copy lands, and + `recv_dev` slots are recycled while still being read. On an 8-GPU node, 7 of 8 + slices of every staged file get zeros or stale bytes. Right length, wrong + contents, exit 0. + + Fires on every multi-GPU node — independent of the distribution correction + above. Fix first. + + **F1 fix (2026-07-31).** `torch.cuda.set_device(self.device)` at the top of + `drainer()`. Guarded by `tests/test_device.py`, which models a per-thread + current device in the stubs and asserts every recorded event belongs to this + rank's device. Confirmed red before the fix (3 of 9 events on device 0 — the + drainer's, one per round) and green after. + + Still unverified on hardware: no GPU was available. The test checks that the + threads agree on which device they are on, not CUDA ordering semantics. Worth + a multi-GPU confirmation run when a node is free. + +- [x] **F2 · `mlperf_common/fileio/datastage.py:199` · new_group sorts its rank list** — resolved by simplifying `Topology` + + `dist.new_group` does `ranks = sorted(ranks)` then + `group_rank = ranks.index(global_rank)`, so a member's group rank is its + position in *sorted* order. The drainer maps all-gather output position `i` to + `layout.segment(i, ...)`, assuming position == `node_index`, where + `node_index` is hostname first-appearance order. + + These agree whenever RANK is block- or cyclic-distributed, which is now the + only supported case. Under a hypothetical `A B B A` layout `group_ranks[1] = + [3, 2]` sorts to `[2, 3]` and two nodes' sub-shards swap. + + **Resolved (2026-07-31) by deleting the machinery.** `Topology` now computes + the layout from slurm's default block distribution — node `i` holds ranks + `i*L .. i*L+L-1` — so `group_ranks[l]` is `[node*L + l for node in ...]`, + ascending by construction, and `new_group`'s sort is provably a no-op. The + `all_gather_object` hostname discovery and the `socket` import are gone. + + The assumption is verified rather than assumed, locally and for free: slurm + reports RANK (`SLURM_PROCID`) and LOCAL_RANK (`SLURM_LOCALID`) independently, + and they agree only under a block distribution, so `rank % L != local_rank` + rejects cyclic and arbitrary launches on the affected ranks with no + collective. `LOCAL_WORLD_SIZE` is now required rather than defaulting to 1, + which used to turn one 8-GPU node into eight single-rank "nodes". + + `tests/test_topology.py` covers six layouts, confirmed able to go red by + reintroducing a mis-grouping. + +## Tier 2 — silent data loss / job hangs + +- [x] **F3 · `mlperf_common/fileio/copyplan.py:82` · os.walk swallows unreadable subtrees** — fixed, see "F3 fix" below + + Default `onerror=None` means a directory that can't be listed is skipped + silently, instead of raising `UnreadableEntries` — the exact failure this + module was written to prevent. + + Verified: `src/{good/a.bin, secret/b.bin}` with `secret/` at mode 000 returns + `['good/a.bin']`, no exception. Unlistable root returns `[]`, also silent. + + Compounding: fastmd5 now walks through the same function, so verification of a + partially staged tree omits the same files and passes. + + **Not a regression.** `os.walk(followlinks=True)` with no `onerror` was + already in `client/fastcp` at `c02c93c`, before the fileio refactor; copyplan + inherited it verbatim. `4f1c063` made *stat* failures noisy and did so + correctly — the miss is that an unlistable directory produces no entry to + stat, so it fails by omission and there is nothing for a stat-oriented check + to collect. What the branch changed is blast radius: one shared walk means + fastmd5 now inherits fastcp's blind spot. + + **F3 fix (2026-07-31).** An `onerror` callback appending to the same + `problems` list, so an unlistable directory joins the existing report. Covers + the root of the walk too. `tests/test_copyplan.py` grew two checks, confirmed + red first: the subtree case returned `['readable/good.bin']` with the hidden + file silently absent, and the unlistable-root case returned `[]`. The check + verifies its own premise and skips if run as a user who can list a 0o000 + directory, rather than passing vacuously. + +- [x] **F4 · `mlperf_common/fileio/datastage.py:649` · no destination-is-a-directory check** — fixed, see "F4 fix" below + + `parse_args` validates sources but never checks the destination. + `plan_copy_operations` (copyplan.py:112) silently falls back to a single-file + copy of `sources[0]` when the destination isn't an existing directory. + + Reproduced: `datastage -r /lustre/dataset /raid/scratch/dataset --dry-run` + with the destination absent emits one bogus job for the source *directory* + (`4096 bytes`). In a real run the O_DIRECT open gets EINVAL, falls back to + buffered, and the reader dies with `[Errno 21] Is a directory` while peers + block in the collective until the watchdog fires. With multiple sources, + `sources[1:]` are dropped with no message and the job exits 0. + + **The review was wrong about fastcp.** It claimed fastcp blocks both cases. + It blocks the multi-source one, and correctly *permits* `-r src newdir` — + which GNU cp permits too. The planner then mis-plans that permitted case, so + `fastcp -r src newdir` was broken in the same way. Verified against real `cp`: + + | invocation | cp | + | --- | --- | + | `cp -r src newdir` (absent) | legal — contents land directly in `newdir` | + | `cp f1 f2 nodir` (absent) | `target 'nodir': No such file or directory` | + | `cp -r src f1` (a file) | `cannot overwrite non-directory ...` | + + The real defect was architectural: the refactor moved *planning* into + copyplan and left *validation* in fastcp's CLI, and the planner's docstring + stated its `not isdir(dst)` fallback as a fact when it was a precondition the + caller had to establish. datastage was the first caller to arrive without + that knowledge. + + **F4 fix (2026-07-31).** `copyplan.validate_copy_args` (+ `CopyArgumentError`) + now owns the cp argument rules, and `plan_copy_operations` applies them itself + so no caller can skip them. The planner handles `cp -r src newdir` properly: + a directory source with a non-directory destination copies its *contents* + under the destination, with no basename level. + + fastcp's destination block is deleted in favour of the shared call, and + datastage's `parse_args` gained the same call — deliberately there rather than + in `build_plan`, since parse_args runs on every rank before the process group + exists, so a bad invocation kills the job uniformly instead of leaving rank 0 + exiting while peers block (which is F5's failure mode). + + fastmd5 is unaffected: read-only, no destination. + + All six cp cases are pinned in `tests/test_copyplan.py`, and both CLIs were + exercised end to end — fastcp copying real bytes, datastage's parse_args via + the stub harness — with error text matching cp's wording. + +- [x] **F5 · `mlperf_common/fileio/datastage.py:572` · unguarded os.stat before the broadcast** — fixed, see "F5 fix" below + + The `try` around the payload build catches only `UnreadableEntries`. A file + unlinked between `plan_copy_operations`'s stat and this one raises + `FileNotFoundError`; rank 0 exits while every peer is already blocked in + `broadcast_object_list`. Job dies on a 10+ minute watchdog timeout with an + opaque collective message instead of `cannot stat `. The comment above + the call says this must not happen. + + **F5 fix (2026-07-31).** The `except UnreadableEntries` became `except + Exception`, carrying the type name into the broadcast payload. The bar is + reaching the broadcast, not anticipating the cause — narrowing the catch to + the failure you went looking for is what created the hang. + + Two live paths, not one: the `os.stat` race, and `CopyArgumentError` from + `plan_copy_operations`, which became reachable here when F4 moved the cp + rules into the planner. The new `tests/test_buildplan.py` caught the second + on its own — I had only written the test for the first. + + Result, raised by every rank together instead of a watchdog timeout: + + cannot stage the source tree: FileNotFoundError: [Errno 2] No such + file or directory: '/.../src/vanished.bin' + + `torch.tensor` was added to the stubs to reach `build_plan`'s success path, + which shaves a little off F15. + + **Same class, not fixed here:** `Stager.__init__`'s memory-budget check + raises per-rank. It is uniform when every node has the same GPU, so it fails + cleanly in practice, but it is the same shape of hazard. `stage_file` already + documents that a mid-file failure leaves peers to the watchdog. + +## Tier 3 — real, lower severity + +- [ ] **F6 · `client/fastmd5:111` · directory-symlink cycles now abort the run** *(priority questioned — see note)* + + Swapping `Path.rglob('*')` (does not descend symlinked dirs) for + `copyplan.list_relative_files` (`os.walk(followlinks=True)`) means a + `latest -> .` or `current -> ..` symlink recurses to ELOOP, raises + `UnreadableEntries`, and `fastmd5:113` turns that into `sys.exit`. Zero + checksums printed — the tree can't be verified at all. Regression against the + pre-branch behaviour. + + **Priority questioned (2026-07-31, Matt).** Not convinced this is worth + fixing: it needs a dataset containing a self- or ancestor-referential + directory symlink, and unlike the rest of the list it fails loudly — a + nonzero exit with no output — rather than silently producing wrong results. + Left open, not scheduled. + + **Measured, since both this file and the code's own FIXME described it + wrongly.** It does not recurse forever and cannot exhaust the stack + (`os.walk` is iterative). The kernel's 40-symlink limit bounds the descent: + 82 directories, depth 81, under 10 ms on a `up -> ..` tree. The mechanism is + not the walk — CPython wraps `entry.is_dir()` in `try/except OSError` and + treats failure as "not a directory", so the un-openable link is reclassified + as a *file* and the walk reports nothing, not even via the `onerror` hook F3 + added. The ELOOP surfaces from `_stat_or_problem`, as one unreadable entry. + + **Trap for whoever does fix it:** the walk enumerates the files under the + cycle once per level on the way down, so "fixing" this by tolerating or + skipping the unreadable entry converts a loud refusal into ~40 redundant + copies of everything beneath it — moving it from the loud category into the + silent one. A real fix detects the cycle (depth-first walk tracking visited + `(st_dev, st_ino)`), which also removes the duplicates. Documented in the + `list_relative_files` docstring. + +- [x] **F7 · `mlperf_common/fileio/datastage.py:406` · `_chmod_parents` infinite loop on `/`** — fixed + + `datastage a.bin /` makes `dest_root == "/"`; `os.path.dirname("/") == "/"`, so + the loop never terminates, chmod'ing `/` forever while peers block in the + following barrier. Hangs to wall-clock limit. + + **Fixed (2026-07-31).** Break when `dirname(path) == path`, which is the + general "walked off the top of the filesystem" condition rather than a + special case for `/`. `tests/test_stager.py` drives `_chmod_parents` against + a stub with a call-count ceiling, so a runaway reports as a failure instead + of hanging the suite; confirmed red by removing the guard. + +- [ ] **F8 · `mlperf_common/fileio/datastage.py:408` · `_chmod_parents` widens pre-existing dirs** + + Loop condition (`path.startswith(dest_root + os.sep) or path == dest_root`) + doesn't distinguish directories datastage created from ones already there. + Verified: a 0750 destination tree becomes 0777. Docstring says the function + exists to fix the umask on *directories we created*. Under + `--container-remap-root` the process is root, so it always succeeds silently. + +- [x] **F9 · `mlperf_common/fileio/datastage.py:343` · `.datastage.tmp` files leak** — fixed + + Preallocated to full source size via `ftruncate`, removed only by the + success-path `os.rename`. No `unlink` anywhere in the file. Every failure + leaves a near-full copy, and the `SLURM_JOB_ID` suffix means each resubmit + orphans a fresh one until node-local NVMe fills — at which point attempts fail + with ENOSPC rather than the original error. Tests assert on leftover temp + files but only exercise the success path. + + **Fixed (2026-07-31).** `stage_file` now wraps the work in `try/except + BaseException`, unlinks the temp best-effort, and re-raises; the body moved + to `_stage_to_temp`. Any rank that raises does the unlink — peers still hold + the file open, but POSIX keeps the inode alive until they close, so the + space returns when they die, and the rename that would have published it is + not going to happen. Ranks parked in a collective never reach the handler, + which is what the NCCL watchdog is for; that residual case is commented. + + `tests/test_pipeline.py` injects a failure into `_run_pipeline` and asserts + no temp survives — and records, from inside the failure, that the temp + existed at that moment, so the test cannot pass by never creating one. + Confirmed red by reverting the unlink. + +- [ ] **F10 · `mlperf_common/fileio/datastage.py:537` · all-gather ignores actual round length** + + `FileLayout` gives `shard >= 2 MiB` for any non-empty slice, so `rounds == 1` + for any file below `buffer_size`, and `send_dev` is always `piece` wide. A + 1 KiB file at N=2048 with default `-b 8M` moves a 16 GiB window into every + node and D2H-copies all of it (`nbytes = count * layout.piece` is + unconditional at :474) to keep 1 KiB. For 100k small files this is petabytes + of fabric traffic plus two world barriers per file — slower than the rsync + path it replaces. + +- [x] **F11 · `mlperf_common/fileio/datastage.py:277` · memory-budget error suggests an over-budget value** — fixed + + Check is `piece * (2N + 1) > budget`; the suggestion divides by `2N`, dropping + the `+1`. Reproduced at 79.65 GiB / 64 nodes: `-b 512M` rejected with "lower to + at most 382M", and `-b 382M` reprints byte-identical advice. Correct answer is + 378M. Mis-suggests at every node count checked (16/32/64 on 80/94/141 GB), so + an operator burns repeated allocations with no workable value ever offered. + +- [x] **F12 · `mlperf_common/fileio/datastage.py:475` · no dedicated stream, so no overlap** — fixed, see "F12 fix" below + + The comment at :444 claims round k+1's collective overlaps round k's copy-back. + It doesn't: the drainer's `copy_` is enqueued on the default stream, same as + H2D(k+1) and all_gather(k+1), and ProcessGroupNCCL makes its internal stream + wait on the default stream. `RECV_DEV_SLOTS=2`, the drain thread and the events + buy only host-side write overlap. Fix is a dedicated `torch.cuda.Stream` for + the drainer with the existing events as cross-stream waits — which also + interacts with F1, so do them together. + + **F12 fix (2026-07-31).** `Stager.drain_stream`, with the copies and their + events issued inside `torch.cuda.stream(self.drain_stream)` and the host-side + `assembled.synchronize()` replaced by `drain_stream.wait_event(assembled)` — + the drain thread now queues a round's copies while that round's collective is + still running, instead of blocking until it finishes. + + `last_copy.synchronize()` before `recv_free_q.put` is now load-bearing for a + second reason: the next all-gather is issued on the default stream, which has + no ordering against `drain_stream`, so that host-side wait is the only thing + keeping the collective off a window still being copied. Commented in place. + + `test_device.py` gained two checks: no drainer operation may be issued on the + default stream, and a side stream must `wait_event` before its first copy. + Both confirmed red before the change; the second was re-confirmed red by + deleting the `wait_event` line and re-running. + + **Unmeasured.** The expected win is per-round `T_collective + T_d2h` becoming + roughly `max(...)`, but whether the GPU side is the bottleneck at all depends + on Lustre read and NVMe write rates. The stubs make every copy instantaneous, + so the suite cannot see throughput. Profile on a real node before claiming a + speedup. + +- [ ] **F13 · `client/direct_io.py:35` · shim fallback can't displace a stale install** + + When an older `mlperf_common` is already importable (a container that + pip-installed one from a benchmark's `requirements.txt`), the first import + binds `mlperf_common` in `sys.modules` to site-packages. Inserting the repo + root at `sys.path[0]` doesn't change the cached package's `__path__`, so the + retry re-resolves to the same place and the error tells the user to do the one + thing they have already done. The compatibility path the shim exists for is + dead in exactly the case it was written for. + +## Found during the work, not by the review + +- [x] **F16 · `mlperf_common/fileio/datastage.py` · the cross-rank source check was a metadata storm** — removed + + `build_plan` had every rank stat every file and all-reduce a mismatch count, + to catch ranks seeing a different view of shared storage. That is one stat + per rank per file — `W × F`, where W is the world size (one rank per GPU, so + nodes × 8) and F is the number of files, matching the W/L/N notation in + `datastage.py`'s module docstring — issued as a synchronised burst: 5.1M at + 64 nodes / 10k files, 41M at + 512 nodes, 1.6 **billion** at 2048 nodes / 100k files. At a generous 50k + ops/s for one MDT that is ~14 minutes of pure metadata at 512 nodes before a + byte moves, and at 2048 it takes the MDT down for every other job on the + filesystem too. + + The cost scaled with the file count; the detection power did not. The + condition it looked for — a stale handle, a failed mount, the wrong dataset + at the same path — is per-mount, and affects every file on that node + identically. It was also imperfect: size and mtime miss same-size, + same-mtime content differences. + + **Removed entirely (2026-07-31, Matt's call).** Reasons: mount verification + already belongs to `mountcheck.py`, once per job, where a sparse SHA256 + fingerprint is cheaper *and* stronger; subtle mount errors of this kind have + not occurred in six years on these clusters; and the planned direction is for + small files to be handled by a single rank each, which makes any per-rank + per-file metadata work the wrong shape regardless. + + Rationale recorded in `build_plan`'s docstring so it does not get reinvented. + `torch.tensor`, `FakeCounter` and `dist.all_reduce` left the stubs with it. + + Worth noting the review did not find this — 49 agents, none reasoning about + filesystem load. Staging still opens every file on every rank — `W × F` + again — which is inherent to having W disjoint readers; this check was pure + addition on top of that. + +## Tier 4 — the tests can't catch the above + +- [ ] **F14 · `tests/test_pipeline.py:90` · fake all-gather re-reads the source** + + `fake_all_gather` ignores its `inp` argument and fills the window from + `os.pread(STATE['fd'], ...)`. Verified: making `_read_piece` overwrite + everything it reads with `0xEE` still passes all 8 multi-node configurations + (only the 7 single-node cases go red). So `_read_piece`'s offset math, its + threaded split, send-slot recycling and the H2D copy are unverified in every + multi-node case. Instrumentation shows the threaded branch is entered exactly + once in the whole suite, in a case whose result is discarded. + +- [x] **F15 · `tests/stubs.py:106` · stubs can't construct `Topology` or `build_plan`** — resolved + + `install()` defines only barrier/get_rank/broadcast_object_list/all_reduce/ + all_gather_into_tensor. `ds.Topology` raises AttributeError on + `all_gather_object`; `ds.build_plan` raises on `torch.tensor`. Both test files + therefore hand-roll their own `Topology` (test_layout.py:47, + test_pipeline.py:95) and never touch the real one. Reintroduce the node + grouping bug and the suite still prints "all tests passed". + + Note `tests/README.md` tells you to validate pipeline changes by reintroducing + a bug and confirming red. Per F14 that instruction does not currently hold for + the multi-node path. + + **Resolved (2026-07-31), largely by deleting the untestable parts.** + `all_gather_object` is gone from `Topology` (F2) and `torch.tensor` is gone + from `build_plan` (F16), so the two blockers stopped existing. `new_group` is + now stubbed, returning its rank list so a test can inspect it. Both functions + have direct tests: `test_topology.py` and `test_buildplan.py`. The hand-rolled + `Topology` classes in `test_layout.py` and `test_pipeline.py` remain, and are + fine — they parameterise `node_count` directly, which is what those tests + need; the real one is exercised on its own. + +--- + +## Suggested order + +1. **F1** — fires unconditionally, corrupts data, blocks trusting anything else. +2. **F14 + F15** — without these, an F1 fix can't be shown to work. +3. **F3, F4, F5** — silent loss and hangs; all small, independent fixes. +4. **F12** with F1 (same code, same stream reasoning). +5. Remainder in any order. **F2** likely reduces to an assert. + +## Provenance + +* `.review-findings.json` — verifier output, this directory (durable copy). +* Workflow run `wf_46ab6788-524`, session `5451bab2-8557-4276-b6fb-39cb3735c04e`. + Per-agent transcripts under + `~/.claude-nvidia-account/projects/-home-matt-work-jun-2026-mlperf-common//subagents/workflows//`. + Session-scoped; will not survive into a new session. +* The original `/tmp/claude-1000/.../tasks/w0ajzd4ot.output` is volatile. diff --git a/client/direct_io.py b/client/direct_io.py index 99424ed..e3ff479 100644 --- a/client/direct_io.py +++ b/client/direct_io.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,137 +14,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import ctypes - -def print_memoryview(mv): - """Print the address and size of a memoryview.""" - address = ctypes.addressof(ctypes.c_char.from_buffer(mv)) - size = mv.nbytes - print(f"Address: {hex(address)}, Size: {hex(size)}") - -def allocate_aligned_buffers(buffer_size, alignment, num_bufs): - """Allocate a sequence of aligned buffers. - - Args: - buffer_size (int): Size of each buffer in bytes. - alignment (int): Alignment boundary in bytes. - num_bufs (int): Number of buffers to allocate. - - Returns: - list[memoryview]: A list of memoryview objects, each representing an aligned buffer - raw_buf: the ctypes raw string_buffer. This *must* be kept to avoid the - gc prematurely collecting the buffers. - """ - assert buffer_size % alignment == 0, "buffer_size must be a multiple of alignment" - - # Allocate a single large buffer that can accommodate all buffers with proper alignment - total_size = (num_bufs * buffer_size) + alignment - raw_buf = ctypes.create_string_buffer(total_size) - base_address = ctypes.addressof(raw_buf) - - # Calculate the offset to align the first buffer - first_offset = (alignment - (base_address % alignment)) % alignment - aligned_base_address = base_address + first_offset - - # Create memoryviews for each buffer - buffers = [] - for i in range(num_bufs): - buf_address = aligned_base_address + i * buffer_size - aligned_buf = (ctypes.c_char * buffer_size).from_address(buf_address) - mv = memoryview(aligned_buf) - # print_memoryview(mv) - buffers.append(mv) - return buffers, raw_buf - -def round_up(value, multiple): - """Round up `value` to the next multiple of `multiple`.""" - if multiple == 0: - raise ValueError("Multiple must be greater than 0.") - return ((value + multiple - 1) // multiple) * multiple - -def pread(fd, aligned_memview, count, offset, fs_block_size, thread_id): - """Perform a direct pread. - - Returns: - number of bytes read. Unlike Posix pread this routine handles - interrupts, so return value should always be equal to count unless - there is an unrecoverable error, in which case it throws rather than returns - - Required: `count` <= aligned_memview size - """ - padded_count = round_up(count, fs_block_size) - assert padded_count <= aligned_memview.nbytes, "memview too small for requested read" - - while True: - try: - bytes_read = os.preadv(fd, [aligned_memview[:padded_count]], offset) +# Compatibility shim. direct_io now lives in the mlperf_common package so that +# it can be imported by package modules (mlperf_common.fileio.datastage) and not only +# by the scripts installed alongside it in bin/. This shim keeps `import +# direct_io` working for anything that still expects a module next to fastcp. - if bytes_read == count: - # Expected case: all requested bytes were read. - return bytes_read +# The package is found either because mlperf-common is installed (setup.py +# installs the package and these scripts together, so this holds for any pip +# install) or because we are running from a tree with mlperf_common/ next to +# client/ -- which covers the source checkout and the "deploy the repo to a +# shared filesystem and run client/ scripts in place" pattern. Copying +# individual scripts out of client/ on their own does not work. - if bytes_read > 0: - # Retry case: partial read, retry the entire read to maintain alignment restrictions - assert bytes_read < count, "bytes_read cannot exceed count" - continue - - if bytes_read == 0: - # Unexpected EOF - raise RuntimeError(f"Unexpected EOF encountered at offset {offset}") - - if bytes_read < 0: - raise OSError("preadv returned a negative value") - - except InterruptedError: - # Retry on EINTR - continue - except OSError as e: - # Non-retriable errors - if e.errno == 14: # Errno 14 corresponds to "Bad address" - raise RuntimeError(f"preadv failed with error: {e}, this usually means out of memory") - raise RuntimeError(f"preadv failed with error: {e}") - -def pwrite(fd, aligned_memview, count, offset, fs_block_size): - """Perform a direct pwrite - - Returns: - number of bytes written. Unlike Posix pwrite this routine handles - interrupts, so return value should always be equal to count unless - there is an unrecoverable error, in which case it throws rather than - returns - - Required: `count <= aligned_memview size - """ - padded_count = round_up(count, fs_block_size) - assert padded_count <= aligned_memview.nbytes, "memview too small for requested write" - while True: - try: - bytes_written = os.pwritev(fd, [aligned_memview[:padded_count]], offset) - - if bytes_written == padded_count: - # Expected case: all requested bytes were written. - # in the case of the last block in the file we may have padded up, so return - # the bytes _requested_ rather than the padded count - return count - - if bytes_written > 0: - # Retry case: partial write, retry the entire write to maintain alignment restrictions - assert bytes_written < count, "bytes_written cannot exceed count" - continue - - if bytes_written == 0: - # Unexpected failure to write - raise RuntimeError(f"Unexpected failure to write at offset {offset}") - - if bytes_written < 0: - raise OSError("pwritev returned a negative value") - - except InterruptedError: - # Retry on EINTR - continue - except OSError as e: - # Non-retriable errors - if e.errno == 14: # Errno 14 corresponds to "Bad address" - raise RuntimeError(f"pwritev failed with error: {e}, this usually means out of memory") - raise RuntimeError(f"pwritev failed with error: {e}") +import os +import sys + +try: + from mlperf_common.fileio.direct_io import * # noqa: F401,F403 +except ImportError: + sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir)) + try: + from mlperf_common.fileio.direct_io import * # noqa: F401,F403 + except ImportError as exc: + raise ImportError( + f"{exc}. direct_io lives in the mlperf_common package. Either install " + "mlperf-common, or keep this script in a tree with mlperf_common/ " + "alongside client/." + ) from exc diff --git a/client/fastcp b/client/fastcp index b7787ab..e48b027 100755 --- a/client/fastcp +++ b/client/fastcp @@ -21,10 +21,26 @@ import tempfile import queue import threading -# find direct_io in the same directory as this program -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +# Find direct_io next to this program, and mlperf_common one level up. That +# covers an install (setup.py installs the package and these scripts together, +# so the package is importable) and a tree with mlperf_common/ next to client/, +# which is both the source checkout and the "deploy the repo to a shared +# filesystem and run client/ scripts in place" pattern. +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) +sys.path.insert(1, os.path.join(_HERE, os.pardir)) import direct_io +try: + from mlperf_common.fileio.copyplan import ( + CopyArgumentError, UnreadableEntries, plan_copy_operations, validate_copy_args) +except ImportError as exc: + sys.exit( + f"fastcp: cannot import mlperf_common ({exc}).\n" + "fastcp needs the mlperf_common package: either install mlperf-common, " + "or run this script from a tree with mlperf_common/ alongside client/." + ) + def copy_worker(fd_src, fd_dst, buffer_size, workpile, thread_buffer, fs_block_size, thread_id): """Worker function to copy chunks from the workpile.""" while True: @@ -147,68 +163,17 @@ def parse_and_validate_args(): # usually meant that to be the source file name, and forgot to give a target parser.error(f"{prog_name}: missing destination file operand after '{args.destination}'") - # Validate sources - for src in args.sources: - if not os.path.exists(src): - sys.exit(f"{prog_name}: cannot stat '{src}': No such file or directory") - if os.path.isdir(src) and not args.recursive: - sys.exit(f"{prog_name}: -r not specified; omitting directory '{src}'") - - # Validate destination - if args.target_directory or len(args.sources) > 1: - if not os.path.exists(args.destination) or not os.path.isdir(args.destination): - sys.exit(f"{prog_name}: target '{args.destination}' is not a directory") - else: - if os.path.exists(args.destination) and os.path.isfile(args.destination): - if os.path.isdir(args.sources[0]): # Only valid if both are files - sys.exit(f"{prog_name}: cannot overwrite '{args.destination}' with directory '{args.sources[0]}'") + # The cp argument rules live in copyplan, next to the mapping they make + # well defined, so that fastcp and datastage cannot drift apart on them. + try: + validate_copy_args(args.sources, args.destination, + recursive=args.recursive, + into_directory=bool(args.target_directory)) + except CopyArgumentError as exc: + sys.exit(f"{prog_name}: {exc}") return args -def list_relative_files(root): - """Return all file paths under 'root' as relative paths, sorted alphabetically. - os.walk(followlinks=True) allows following symlinked directories, - but it does not guard against cycles. This may cause infinite loops - if symlinks form a directory cycle. - """ - file_list = [] - # FIXME: os.walk(followlinks) doesn't protect against cycles - # to fix this we'd need to write our own version that did a depth-first - # spanning tree. - for dirpath, _, filenames in os.walk(root, followlinks=True): - for fname in filenames: - full_path = os.path.join(dirpath, fname) - rel_path = os.path.relpath(full_path, root) - file_list.append(rel_path) - return sorted(file_list) - -def plan_copy_operations(args): - """Return list of (src_abs, dst_abs, size_bytes) file tuples to copy.""" - file_jobs = [] - dst_root = os.path.abspath(args.destination) - - if not os.path.isdir(dst_root): # case 1: single file copy - src_abs = os.path.abspath(args.sources[0]) - dst_abs = os.path.abspath(args.destination) - size = os.path.getsize(src_abs) - file_jobs.append((src_abs, dst_abs, size)) - else: - for src in args.sources: - src_abs = os.path.abspath(src) - base = os.path.basename(src.rstrip("/")) - if os.path.isdir(src): - for relpath in list_relative_files(src): - full_src = os.path.join(src_abs, relpath) - full_dst = os.path.join(dst_root, base, relpath) - size = os.path.getsize(full_src) - file_jobs.append((full_src, full_dst, size)) - else: - dst_path = os.path.join(dst_root, base) - size = os.path.getsize(src_abs) - file_jobs.append((src_abs, dst_path, size)) - - return file_jobs - if __name__ == "__main__": args = parse_and_validate_args() # round buffer size up to next multiple of 2 MiB @@ -219,7 +184,10 @@ if __name__ == "__main__": print(f"Sources: {args.sources}") print(f"Destination: {args.destination}") - file_jobs = plan_copy_operations(args) + try: + file_jobs = plan_copy_operations(args.sources, args.destination) + except UnreadableEntries as exc: + sys.exit(f"fastcp: {exc}") if not args.force: for _, dst, _ in file_jobs: diff --git a/client/fastmd5 b/client/fastmd5 index adf7644..35e17df 100755 --- a/client/fastmd5 +++ b/client/fastmd5 @@ -6,12 +6,26 @@ import sys import queue import hashlib import threading -from pathlib import Path -# find direct_io in the same directory as this program -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +# Find direct_io next to this program, and mlperf_common one level up. That +# covers an install (setup.py installs the package and these scripts together, +# so the package is importable) and a tree with mlperf_common/ next to client/, +# which is both the source checkout and the "deploy the repo to a shared +# filesystem and run client/ scripts in place" pattern. +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) +sys.path.insert(1, os.path.join(_HERE, os.pardir)) import direct_io # for pread(), allocate_aligned_buffers(), round_up() +try: + from mlperf_common.fileio.copyplan import UnreadableEntries, list_relative_files +except ImportError as exc: + sys.exit( + f"fastmd5: cannot import mlperf_common ({exc}).\n" + "fastmd5 needs the mlperf_common package: either install mlperf-common, " + "or run this script from a tree with mlperf_common/ alongside client/." + ) + def parse_and_validate_args(): """Parse command-line arguments. Always assume recursive for directories.""" parser = argparse.ArgumentParser( @@ -47,30 +61,41 @@ def enqueue_file_chunks(display_name, actual_path, chunk_size, workpile): this_chunk = min(chunk_size, size - offset) workpile.put((display_name, offset, this_chunk, block_size)) -def checksum_worker(thread_id, buffer, workpile, buffer_size, print_lock): - """Worker function to compute md5 checksums of file chunks.""" +def checksum_worker(thread_id, buffer, workpile, buffer_size, print_lock, errors): + """Worker function to compute md5 checksums of file chunks. + + A failure here must not be silent: an exception escaping a thread does not + affect the process exit status, so a caller comparing two trees would see + a short but successful-looking result. Record it instead and let main() + exit nonzero. + """ while True: try: filepath, offset, size, fs_block_size = workpile.get_nowait() except queue.Empty: return - md5 = hashlib.md5() - - with open(filepath, 'rb') as f: - fd = f.fileno() - remaining = size - local_offset = offset - while remaining > 0: - this_read = min(remaining, buffer_size) - bytes_read = direct_io.pread(fd, buffer, this_read, local_offset, fs_block_size, thread_id) - assert bytes_read == this_read, f"pread returned {bytes_read}, expected {this_read}" - md5.update(buffer[:this_read]) - local_offset += this_read - remaining -= this_read - - with print_lock: - print(f"{filepath}\t{offset}\t{size}\t{md5.hexdigest()}") + try: + md5 = hashlib.md5() + + with open(filepath, 'rb') as f: + fd = f.fileno() + remaining = size + local_offset = offset + while remaining > 0: + this_read = min(remaining, buffer_size) + bytes_read = direct_io.pread(fd, buffer, this_read, local_offset, fs_block_size, thread_id) + assert bytes_read == this_read, f"pread returned {bytes_read}, expected {this_read}" + md5.update(buffer[:this_read]) + local_offset += this_read + remaining -= this_read + + with print_lock: + print(f"{filepath}\t{offset}\t{size}\t{md5.hexdigest()}") + except Exception as exc: # noqa: BLE001 - reported and re-surfaced by main() + errors.append((filepath, offset, exc)) + with print_lock: + print(f"fastmd5: {filepath} at offset {offset}: {exc}", file=sys.stderr) workpile.task_done() def main(): @@ -78,16 +103,19 @@ def main(): workpile = queue.Queue() CHUNK_SIZE = 1 << 30 # 1 GiB - for root in args.paths: - original_arg = root - path = Path(original_arg) - if path.is_dir(): - for file in path.rglob('*'): - if file.is_file() or file.is_symlink(): - display_name = os.path.join(original_arg, os.path.relpath(file, path)) - enqueue_file_chunks(display_name, file, CHUNK_SIZE, workpile) - elif path.is_file() or path.is_symlink(): - enqueue_file_chunks(original_arg, path, CHUNK_SIZE, workpile) + # Share the tree walk with fastcp and datastage, so that checksumming a + # staged copy compares the same set of files the copy was planned from. + for original_arg in args.paths: + if os.path.isdir(original_arg): + try: + relpaths = list_relative_files(original_arg) + except UnreadableEntries as exc: + sys.exit(f"fastmd5: {exc}") + for relpath in relpaths: + full_path = os.path.join(original_arg, relpath) + enqueue_file_chunks(full_path, full_path, CHUNK_SIZE, workpile) + else: + enqueue_file_chunks(original_arg, original_arg, CHUNK_SIZE, workpile) # Prepare aligned buffers for each thread alignment = 2 * 1024 * 1024 # 2 MiB @@ -100,14 +128,20 @@ def main(): # Launch worker threads print_lock = threading.Lock() + errors = [] # list.append is atomic, so no lock needed threads = [] for i in range(args.num_threads): - t = threading.Thread(target=checksum_worker, args=(i, thread_buffers[i], workpile, args.buffer_size, print_lock)) + t = threading.Thread(target=checksum_worker, args=(i, thread_buffers[i], workpile, args.buffer_size, print_lock, errors)) t.start() threads.append(t) for t in threads: t.join() + if errors: + print(f"fastmd5: {len(errors)} chunk(s) failed; output is incomplete", + file=sys.stderr) + sys.exit(1) + if __name__ == '__main__': main() diff --git a/mlperf_common/dist_env.py b/mlperf_common/dist_env.py new file mode 100644 index 0000000..39575e6 --- /dev/null +++ b/mlperf_common/dist_env.py @@ -0,0 +1,371 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Derive torch.distributed's env:// variables from slurm or mpirun. + +torch.distributed.init_process_group(init_method="env://") wants RANK, +WORLD_SIZE, MASTER_ADDR and MASTER_PORT in the environment, and by convention +LOCAL_RANK and LOCAL_WORLD_SIZE alongside them. Neither slurm nor mpirun sets +those names, so something has to translate. `client/slurm2pytorch` does it in +bash and execs the real command; this is the same translation in Python, so a +program can do it for itself instead of being wrapped. + +The two must stay behaviourally compatible: a program launched under +slurm2pytorch finds the variables already set, and takes them as given. + +Two deliberate divergences from the script, both in the direction of failing +loudly rather than proceeding on a guess: + + * slurm2pytorch defaults LOCAL_WORLD_SIZE to 1. But SLURM_NTASKS_PER_NODE is + only set when --ntasks-per-node was actually passed, so `srun -N2 -n16` + has no source for it and that default turns two 8-GPU nodes into sixteen + single-rank "nodes" -- which reads as a valid job right up until the copy + comes out wrong. We consult SLURM_TASKS_PER_NODE, which srun always sets, + and refuse to guess on a multi-rank job. + * slurm2pytorch falls back to MASTER_ADDR=127.0.0.1, which its own comment + notes "will fail for multinode" -- as a rendezvous that hangs until the + job's wall clock runs out. We parse the address out of slurm's nodelist, + and if that fails on a multi-node job we say so immediately. + +There is deliberately no PYTORCH_VERSION gate. slurm2pytorch has one because +it wraps arbitrary commands in arbitrary containers and must be a no-op outside +a pytorch one. A caller importing this module is already torch code. +""" + +import collections +import os + +__all__ = ["DistEnv", "DistEnvError", "MASTER_PORT_DEFAULT", "configure", + "describe", "detect", "first_hostname"] + +# Static rendezvous port. torch.distributed's TCPStore does not handle races +# for dynamic port assignment gracefully, so a fixed port is safer than an +# ephemeral one. 29500 is what torch.distributed.run uses, is unassigned by +# IANA, and sits outside both the IANA ephemeral range (49152-65535) and +# Ubuntu's default (32768-60999). If your site puts a service on 29500, or +# includes it in the local ephemeral range (`cat +# /proc/sys/net/ipv4/ip_local_port_range`), set MASTER_PORT explicitly; +# 61000-65535 is usually a good alternative on Linux. +MASTER_PORT_DEFAULT = 29500 + +_LOCALHOST = "127.0.0.1" + +DistEnv = collections.namedtuple( + "DistEnv", + "rank world_size local_rank local_world_size master_addr master_port source") + + +class DistEnvError(RuntimeError): + """The environment does not describe a job we can join. + + Carries a message without a program name; callers prefix their own. + """ + + +def _value(environ, *names): + """First of `names` with a non-empty value, as bash ${X:-...} would treat it. + + environ.get() alone is not enough: an exported-but-empty variable is a + real thing to find in a job script, and bash's :- skips it. + """ + for name in names: + text = environ.get(name) + if text is not None and text.strip(): + return name, text.strip() + return None, None + + +def _int(environ, *names): + """_value(), converted, naming the variable that was actually wrong.""" + name, text = _value(environ, *names) + if name is None: + return None, None + try: + return name, int(text) + except ValueError: + raise DistEnvError(f"{name}={text!r} is not an integer") from None + + +def first_hostname(nodelist): + """First hostname of a slurm hostlist, or None if it cannot be read. + + Slurm hands the allocation over compressed -- "dgx[001-004,007]", + "a,b", "rack[1-2]node[3-4]" -- and `scontrol show hostnames` is not + available inside a container. Only the first name is ever needed (rank 0 + runs there under a block distribution), and that is much less work than + expanding the list. + + Two things a naive nodelist.split(",")[0] gets wrong: + + * commas inside brackets. "dgx[001-004,007]" is one entry, not two, and + splitting yields "dgx[001". + * zero padding. The low bound is appended as the literal token it was + written as, never through int(), because "dgx[001-004]" is dgx001 and + dgx1 does not resolve. + + Returns None rather than raising: an unparseable nodelist is only fatal if + nothing else supplies an address, and that is the caller's judgement. + """ + if not nodelist or not nodelist.strip(): + return None + + text = nodelist.strip() + out = [] + i = 0 + while i < len(text): + char = text[i] + if char == ",": # top-level comma: end of the first entry + break + if char == "[": + close = text.find("]", i) + if close < 0: + return None # unbalanced + low = text[i + 1:close].split(",")[0].split("-")[0].strip() + if not low.isdigit(): + return None + out.append(low) + i = close + 1 + continue + if char == "]": + return None # stray close bracket + out.append(char) + i += 1 + + name = "".join(out).strip() + if not name or not all(c.isalnum() or c in "-_." for c in name): + return None + return name + + +def _tasks_per_node(text): + """Expand SLURM_TASKS_PER_NODE ("8", "8(x2)", "4,8", "2(x3),1") to a list. + + Unlike SLURM_NTASKS_PER_NODE this is set whenever srun runs, which is what + makes it a usable fallback. Returns None if it cannot be parsed. + """ + counts = [] + for part in text.split(","): + part = part.strip() + if not part: + return None + if "(" in part: + count, _, repeat = part.partition("(") + repeat = repeat.strip().rstrip(")").lstrip("xX") + if not count.strip().isdigit() or not repeat.isdigit(): + return None + counts.extend([int(count)] * int(repeat)) + else: + if not part.isdigit(): + return None + counts.append(int(part)) + return counts or None + + +def _local_world_size(environ, world_size): + """Ranks per node, refusing to guess it on a multi-rank job.""" + name, value = _int(environ, "LOCAL_WORLD_SIZE", "SLURM_NTASKS_PER_NODE", + "OMPI_COMM_WORLD_LOCAL_SIZE") + if name is not None: + return value + + text = environ.get("SLURM_TASKS_PER_NODE", "").strip() + if text: + counts = _tasks_per_node(text) + if counts and len(set(counts)) == 1: + return counts[0] + if counts: + raise DistEnvError( + f"SLURM_TASKS_PER_NODE={text!r} describes a ragged allocation " + f"({sorted(set(counts))} tasks per node); datastage needs the " + "same number of ranks on every node" + ) + + if world_size == 1: + return 1 + raise DistEnvError( + f"cannot determine LOCAL_WORLD_SIZE for a WORLD_SIZE={world_size} job; " + "launch with a uniform --ntasks-per-node (one task per GPU), or set " + "LOCAL_WORLD_SIZE explicitly" + ) + + +def _master_addr(environ, world_size, local_world_size): + """Where rank 0 is, which is the first node of the allocation.""" + _, addr = _value(environ, "MASTER_ADDR", "MLPERF_SLURM_FIRSTNODE") + if addr: + return addr + + # The step's nodelist first: a step may run on a subset of the job's nodes, + # and rank 0 of *this step* is on the first of those. SLURM_NODELIST is + # the older alias, kept last. + for name in ("SLURM_STEP_NODELIST", "SLURM_JOB_NODELIST", "SLURM_NODELIST"): + host = first_hostname(environ.get(name)) + if host: + return host + + if world_size <= local_world_size: + return _LOCALHOST + + seen = {name: environ.get(name) for name in + ("SLURM_STEP_NODELIST", "SLURM_JOB_NODELIST", "SLURM_NODELIST") + if environ.get(name)} + raise DistEnvError( + f"cannot determine MASTER_ADDR for a multi-node job " + f"(WORLD_SIZE={world_size}, LOCAL_WORLD_SIZE={local_world_size}); " + f"nodelist variables are {seen or 'unset'}. Set MASTER_ADDR, or set " + 'MLPERF_SLURM_FIRSTNODE from the host with $(scontrol show hostnames ' + '"$SLURM_JOB_NODELIST" | head -n1)' + ) + + +def _check_sources_agree(environ, rank, world_size): + """Catch an environment describing two different jobs at once. + + slurm and mpirun each set their own variables, and a command line that + mixes launchers leaves both -- with only one of them true. Preferring + slurm silently would pick a rank ordering the other half of the job does + not share. + """ + _, ompi_rank = _int(environ, "OMPI_COMM_WORLD_RANK") + _, ompi_world = _int(environ, "OMPI_COMM_WORLD_SIZE") + _, slurm_rank = _int(environ, "SLURM_PROCID") + _, slurm_world = _int(environ, "SLURM_NTASKS") + + if slurm_rank is None or ompi_rank is None: + return + if slurm_rank != ompi_rank or ( + slurm_world is not None and ompi_world is not None + and slurm_world != ompi_world): + raise DistEnvError( + f"slurm and mpirun disagree about this process: " + f"SLURM_PROCID={slurm_rank} of {slurm_world}, " + f"OMPI_COMM_WORLD_RANK={ompi_rank} of {ompi_world}. Set RANK and " + "WORLD_SIZE explicitly to say which is right" + ) + + +def _check_slurm_arithmetic(environ): + """SLURM_NNODES * SLURM_NTASKS_PER_NODE has to be SLURM_NTASKS.""" + _, nodes = _int(environ, "SLURM_NNODES", "SLURM_JOB_NUM_NODES") + _, per_node = _int(environ, "SLURM_NTASKS_PER_NODE") + _, tasks = _int(environ, "SLURM_NTASKS") + if None in (nodes, per_node, tasks): + return + if nodes * per_node != tasks: + raise DistEnvError( + f"slurm's own numbers do not agree: SLURM_NNODES={nodes} x " + f"SLURM_NTASKS_PER_NODE={per_node} is {nodes * per_node}, but " + f"SLURM_NTASKS={tasks}" + ) + + +def detect(environ=None): + """Work out the job layout without touching the environment. + + Validation here is about the *sources*: that they are present, agree with + each other, and describe a coherent job. Whether the resulting layout is + the block distribution datastage needs is Topology's business, and it + still checks -- someone can set these variables by hand without coming + through here. + """ + environ = os.environ if environ is None else environ + + _check_slurm_arithmetic(environ) + + world_name, world_size = _int(environ, "WORLD_SIZE", "SLURM_NTASKS", + "OMPI_COMM_WORLD_SIZE") + if world_name is None: + world_size = 1 + rank_name, rank = _int(environ, "RANK", "SLURM_PROCID", "OMPI_COMM_WORLD_RANK") + if rank_name is None: + rank = 0 + local_name, local_rank = _int(environ, "LOCAL_RANK", "SLURM_LOCALID", + "OMPI_COMM_WORLD_LOCAL_RANK") + if local_name is None: + local_rank = 0 + + if rank_name != "RANK": + _check_sources_agree(environ, rank, world_size) + + source = {"RANK": "preset", "SLURM_PROCID": "slurm", + "OMPI_COMM_WORLD_RANK": "ompi"}.get(rank_name, "single") + + local_world_size = _local_world_size(environ, world_size) + + if world_size < 1: + raise DistEnvError(f"WORLD_SIZE={world_size} must be at least 1") + if not 0 <= rank < world_size: + raise DistEnvError(f"RANK={rank} is outside WORLD_SIZE={world_size}") + if local_world_size < 1: + raise DistEnvError( + f"LOCAL_WORLD_SIZE={local_world_size} must be at least 1") + if not 0 <= local_rank < local_world_size: + raise DistEnvError( + f"LOCAL_RANK={local_rank} is outside " + f"LOCAL_WORLD_SIZE={local_world_size}") + if world_size % local_world_size: + raise DistEnvError( + f"WORLD_SIZE={world_size} is not a multiple of " + f"LOCAL_WORLD_SIZE={local_world_size}; launch with a uniform " + "--ntasks-per-node" + ) + + _, master_port = _int(environ, "MASTER_PORT") + if master_port is None: + master_port = MASTER_PORT_DEFAULT + if not 1 <= master_port <= 65535: + raise DistEnvError(f"MASTER_PORT={master_port} is not a valid port") + + return DistEnv( + rank=rank, + world_size=world_size, + local_rank=local_rank, + local_world_size=local_world_size, + master_addr=_master_addr(environ, world_size, local_world_size), + master_port=master_port, + source=source, + ) + + +def configure(environ=None): + """detect(), then write the result where torch and Topology will read it. + + Setting the environment is the point, not a side effect: + init_process_group(init_method="env://") reads os.environ, and so does + Topology. + """ + environ = os.environ if environ is None else environ + env = detect(environ) + + environ["RANK"] = str(env.rank) + environ["WORLD_SIZE"] = str(env.world_size) + environ["LOCAL_RANK"] = str(env.local_rank) + environ["LOCAL_WORLD_SIZE"] = str(env.local_world_size) + environ["MASTER_ADDR"] = env.master_addr + environ["MASTER_PORT"] = str(env.master_port) + + # As torch.distributed.run does: without this every rank on a node spawns + # a full set of OpenMP threads and they fight over the same cores. + if env.local_world_size > 1 and not environ.get("OMP_NUM_THREADS", "").strip(): + environ["OMP_NUM_THREADS"] = "1" + + return env + + +def describe(env): + """One line, worth printing before the rendezvous rather than after.""" + nodes = env.world_size // env.local_world_size + return (f"{env.source} launch: {env.world_size} ranks = {nodes} nodes x " + f"{env.local_world_size} ranks/node, rendezvous at " + f"{env.master_addr}:{env.master_port}") diff --git a/mlperf_common/fileio/__init__.py b/mlperf_common/fileio/__init__.py new file mode 100644 index 0000000..82eacfb --- /dev/null +++ b/mlperf_common/fileio/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fast file I/O: O_DIRECT primitives, copy planning, and collective staging. + + direct_io O_DIRECT pread/pwrite with aligned buffers (fastcp, fastmd5, datastage) + copyplan source-tree walk and src->dst mapping (fastcp, datastage) + datastage NCCL-collective dataset staging onto node-local storage + +Only datastage requires torch; the other two are dependency-free so that the +single-node client scripts do not pull in a training stack. +""" diff --git a/mlperf_common/fileio/copyplan.py b/mlperf_common/fileio/copyplan.py new file mode 100644 index 0000000..263ed17 --- /dev/null +++ b/mlperf_common/fileio/copyplan.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Source-tree walk and src->dst mapping shared by fastcp and datastage. + +Both tools take the same cp/rsync-shaped arguments and need the same answer: +given SOURCE(s) and a DEST, which files get copied where. Keeping that in one +place means the single-node and the collective stager can never disagree about +what "copy this directory" means. + +The argument rules live here too, in validate_copy_args, rather than in each +tool's CLI. They are not decoration: the mapping below is only well defined +once they hold, and a caller that skipped them used to get a silently wrong +plan instead of an error. plan_copy_operations therefore applies them itself, +and the CLIs call them early only to fail before doing any other setup. + +The semantics are GNU cp's, which is what fastcp set out to match: + + cp -r src newdir newdir absent -> create it, src's *contents* inside + cp -r src existingdir -> existingdir/src/... + cp a b existingdir -> existingdir/a, existingdir/b + cp a b newdir newdir absent -> error, target is not a directory + cp -r src file -> error, cannot overwrite non-directory +""" + +import os + +__all__ = ["CopyArgumentError", "UnreadableEntries", "list_relative_files", + "plan_copy_operations", "validate_copy_args"] + +# How many bad paths to name before summarising the rest. +_MAX_REPORTED = 20 + + +class UnreadableEntries(Exception): + """Some entries under a source tree could not be stat'd. + + Raised instead of letting the first bad entry throw, so a caller learns the + full extent in one pass rather than fixing them one job at a time. The + usual cause is a dangling symlink: os.walk lists it among the filenames + because it is not a directory, and stat then follows it to nothing. + """ + + def __init__(self, entries): + self.entries = list(entries) + super().__init__(self._describe()) + + def _describe(self): + count = len(self.entries) + noun = "entry" if count == 1 else "entries" + lines = [f"{count} unreadable {noun}:"] + for path, reason in self.entries[:_MAX_REPORTED]: + lines.append(f" {path}: {reason}") + if count > _MAX_REPORTED: + lines.append(f" ... and {count - _MAX_REPORTED} more") + return "\n".join(lines) + + +class CopyArgumentError(Exception): + """The SOURCE/DEST combination is not one cp would accept. + + Carries a message in cp's wording, without a program name; callers prefix + their own, as cp does. + """ + + +def validate_copy_args(sources, destination, recursive=True, into_directory=False): + """Check a SOURCE(s)/DEST combination, raising CopyArgumentError if bad. + + `recursive` is the caller's -r flag and `into_directory` its -t. Both only + tighten the check, so plan_copy_operations can re-apply this with the + defaults on arguments a CLI has already accepted and never disagree. + """ + if not sources: + raise CopyArgumentError( + f"missing destination file operand after '{destination}'") + + for src in sources: + if not os.path.exists(src): + raise CopyArgumentError(f"cannot stat '{src}': No such file or directory") + if os.path.isdir(src) and not recursive: + raise CopyArgumentError(f"-r not specified; omitting directory '{src}'") + + if os.path.isdir(destination): + return + + # Not a directory, so there is nowhere to put a second source, and no + # basename to place anything under. -t says the destination is meant to be + # a directory to copy into, so it has to already be one. + if into_directory or len(sources) > 1: + if os.path.exists(destination): + raise CopyArgumentError(f"target '{destination}' is not a directory") + raise CopyArgumentError(f"target '{destination}': No such file or directory") + + if os.path.exists(destination) and os.path.isdir(sources[0]): + raise CopyArgumentError( + f"cannot overwrite non-directory '{destination}' " + f"with directory '{sources[0]}'") + + +def _stat_or_problem(path, problems): + """stat(), following symlinks. On failure record it and return None.""" + try: + return os.stat(path) + except OSError as exc: + problems.append((path, exc.strerror or str(exc))) + return None + + +def list_relative_files(root): + """Return all file paths under 'root' as relative paths, sorted alphabetically. + + os.walk(followlinks=True) allows following symlinked directories, and + symlinks to files appear in filenames, so both are dereferenced and copied + as content rather than recreated as links. + + Raises UnreadableEntries if anything under root cannot be stat'd, or if a + directory under it cannot be listed, reporting every such entry rather than + dying on the first one. + + Nothing here detects symlink cycles (`latest -> .`, `up -> ..`), but they + do not hang: the kernel allows 40 symlink traversals per path resolution, + so the descent stops around 40 levels down, in milliseconds. What that + looks like is worth knowing, because os.walk hides it -- it wraps + entry.is_dir() in try/except OSError and treats a failure as "not a + directory", so the un-openable link is reclassified as a *file* and the + walk reports no error at all, not even through onerror. The ELOOP then + surfaces here, from the stat, as an unreadable entry, and the copy is + refused. + + That refusal is the safe outcome, so do NOT "fix" this by skipping the + unreadable entry: the walk enumerates the files under the cycle once per + level on the way down, so tolerating the ELOOP would turn a loud refusal + into ~40 redundant copies of everything beneath it. A real fix means + detecting the cycle -- a depth-first walk tracking visited (st_dev, st_ino) + -- which also drops the duplicates. Not worth it until a dataset actually + contains such a link. + """ + file_list = [] + problems = [] + + def unlistable(exc): + # The one unreadable-entry case with no entry to report: files under a + # directory we cannot list never reach the walk at all, so there is + # nothing to stat and nothing to collect. os.walk's default is to + # swallow this and carry on, which makes an unlistable subtree + # indistinguishable from an empty one -- a partial copy that exits 0, + # and, since fastmd5 enumerates through here too, a partial checksum + # that agrees with it. + problems.append((exc.filename or root, exc.strerror or str(exc))) + + for dirpath, _, filenames in os.walk(root, followlinks=True, onerror=unlistable): + for fname in filenames: + full_path = os.path.join(dirpath, fname) + if _stat_or_problem(full_path, problems) is not None: + file_list.append(os.path.relpath(full_path, root)) + if problems: + raise UnreadableEntries(problems) + return sorted(file_list) + + +def plan_copy_operations(sources, destination): + """Return a list of (src_abs, dst_abs, size_bytes) tuples to copy. + + If `destination` is an existing directory each source is placed inside it + under its own basename (recursing into directories). Otherwise there is + exactly one source and `destination` names it directly: a file is copied to + that name, and a directory has its *contents* copied in under it, which is + what `cp -r src newdir` does when newdir does not yet exist. The result is + sorted by destination path so that every rank of a collective copy walks + the files in the same order. + + Raises CopyArgumentError if the arguments are not a combination cp would + accept, and UnreadableEntries if any source cannot be stat'd or any + directory under it cannot be listed, reporting all of them together. + """ + validate_copy_args(sources, destination) + + file_jobs = [] + problems = [] + dst_root = os.path.abspath(destination) + + def add(src_abs, dst_abs): + st = _stat_or_problem(src_abs, problems) + if st is not None: + file_jobs.append((src_abs, dst_abs, st.st_size)) + + if not os.path.isdir(dst_root): + # Validation has established there is exactly one source and, if it is + # a directory, that nothing is in the way. `destination` *is* the + # copy, so a directory's contents go directly under it -- no basename + # level, which is where cp -r and cp differ. + src_abs = os.path.abspath(sources[0]) + if os.path.isdir(src_abs): + try: + for relpath in list_relative_files(src_abs): + add(os.path.join(src_abs, relpath), os.path.join(dst_root, relpath)) + except UnreadableEntries as exc: + problems.extend(exc.entries) + else: + add(src_abs, dst_root) + else: + for src in sources: + src_abs = os.path.abspath(src) + base = os.path.basename(src.rstrip("/")) + if os.path.isdir(src): + try: + relpaths = list_relative_files(src) + except UnreadableEntries as exc: + problems.extend(exc.entries) + continue + for relpath in relpaths: + add(os.path.join(src_abs, relpath), + os.path.join(dst_root, base, relpath)) + else: + add(src_abs, os.path.join(dst_root, base)) + + if problems: + raise UnreadableEntries(problems) + return sorted(file_jobs, key=lambda job: job[1]) diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py new file mode 100644 index 0000000..c83d421 --- /dev/null +++ b/mlperf_common/fileio/datastage.py @@ -0,0 +1,828 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Stage a dataset from shared storage onto node-local storage, collectively. + +Every node needs the same bytes. Having each node read the whole dataset from +Lustre is N-times redundant and runs at the speed of the slowest reader, so +instead each rank reads a disjoint shard and the bytes cross the fabric once, +via NCCL. + +Decomposition (W = world size, L = ranks per node, N = nodes): + + * The world is split into L process groups, group `l` holding the rank with + LOCAL_RANK == l on every node. Group `l` owns slice `l` of the file -- a + contiguous 1/L of it. + * Inside group `l`, each of the N ranks reads a disjoint 1/N sub-shard of + that slice from shared storage, and an all-gather assembles the whole slice + on every node. Each rank then writes slice `l` to its node-local copy. + * The node's L ranks together write the whole file, so nothing is exchanged + or written twice within a node. + +That gives W disjoint readers on the shared filesystem, one fabric crossing per +byte per node, and L concurrent all-gathers which between them drive every NIC +without any hand-tuned per-cluster transport configuration. + +Streaming is windowed so memory stays bounded, and reads, the collective, and +writes are pipelined across rounds. + +Usage mirrors cp/rsync and `fastcp`: + + python3 -m mlperf_common.fileio.datastage -r SRC... DST + +Launch with one task per GPU. No wrapper is needed: mlperf_common.dist_env +derives the rendezvous variables (RANK, WORLD_SIZE, LOCAL_RANK, +LOCAL_WORLD_SIZE, MASTER_ADDR, MASTER_PORT) from slurm or mpirun, including +parsing MASTER_ADDR out of slurm's nodelist, which is the one value slurm does +not hand over directly. + + srun --ntasks-per-node=${DGXNGPU} ... \\ + python3 -m mlperf_common.fileio.datastage -r "${SLOW_DATADIR}/${DATASET}" "${DATADIR}" + +--ntasks-per-node is effectively required for a multi-node copy: it is what +tells us how many ranks share a node, and rather than guess we refuse. +Launching under `slurm2pytorch` still works -- the variables are then already +set, and are taken as given. With no launcher at all, --dry-run prints the +copy plan without touching CUDA. +""" + +import argparse +import os +import queue +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import torch +import torch.distributed as dist + +from mlperf_common import dist_env +from mlperf_common.fileio import direct_io +from mlperf_common.fileio.copyplan import ( + CopyArgumentError, UnreadableEntries, plan_copy_operations, validate_copy_args) + +# Buffer alignment. 2 MiB is the Linux huge page size, comfortably above any +# filesystem block size we will meet, so O_DIRECT is always happy. Same +# constant fastcp uses. +BUFFER_ALIGN = 2 * 1024 * 1024 + +# Don't bother splitting a read across threads below this; the syscall is +# already large enough to saturate a reader. +MIN_READ_PIECE = 4 * 1024 * 1024 + +# Pipeline depth. Send slots let the reader run ahead of the collective. +SEND_SLOTS = 3 + +# Assembled windows held on the device. Two, so the next round's all-gather +# can run while the previous window is still being copied back. Device memory +# is the cheap resource here: this runs as its own job step, before training, +# and everything is released when the process exits. +RECV_DEV_SLOTS = 2 + +# Host staging for the write side. An assembled window is a concatenation of +# one segment per node, each bound for a different file offset, so it can be +# copied back and written a few segments at a time instead of whole. Keeping +# this pool small and fixed is what stops pinned host memory -- which is +# page-locked and reserved up front -- from scaling with the node count. +CHUNK_TARGET = 64 << 20 +CHUNK_SLOTS = 3 + + +def ceil_div(a, b): + return -(-a // b) + + +def align_up(value, multiple): + return ceil_div(value, multiple) * multiple + + +def align_down(value, multiple): + return (value // multiple) * multiple + + +def parse_size(text): + """Parse a size with an optional K/M/G/T suffix.""" + text = str(text).strip() + multipliers = {"k": 1024, "m": 1024 ** 2, "g": 1024 ** 3, "t": 1024 ** 4} + multiplier = 1 + if text and text[-1].lower() in multipliers: + multiplier = multipliers[text[-1].lower()] + text = text[:-1] + value = int(float(text) * multiplier) + if value <= 0: + raise argparse.ArgumentTypeError(f"size must be positive: {text}") + return value + + +def open_maybe_direct(path, flags): + """Open with O_DIRECT, falling back to buffered if the filesystem says no. + + Reads and writes stay block-padded either way, which is harmless for a + buffered descriptor, so nothing downstream has to care which we got. + """ + try: + return os.open(path, flags | os.O_DIRECT) + except OSError as exc: + if exc.errno not in (22, 95): # EINVAL, EOPNOTSUPP + raise + return os.open(path, flags) + + +def pinned_aligned(nbytes, alignment): + """Allocate a pinned host buffer aligned for O_DIRECT. + + cudaHostAlloc returns page-aligned memory in practice, but O_DIRECT is + unforgiving and the failure mode (EINVAL deep in a worker thread) is + miserable to debug, so over-allocate and slice to a known-aligned offset + rather than trusting it. + + Returns (owner, view, memoryview); `owner` must be kept alive. + """ + owner = torch.empty(nbytes + alignment, dtype=torch.uint8, pin_memory=True) + offset = (-owner.data_ptr()) % alignment + view = owner[offset:offset + nbytes] + assert view.data_ptr() % alignment == 0, "failed to align pinned buffer" + return owner, view, memoryview(view.numpy()) + + +class Topology: + """Rank/node layout and the per-LOCAL_RANK process groups. + + Ranks are laid out in slurm's default block distribution: node `i` holds + ranks `i*L` through `i*L + L - 1`, so a rank's node is `RANK // L` and its + position on that node is `RANK % L`. That is arithmetic, not a discovery + problem, and it needs no collective to work out. + + The one thing worth checking is that the assumption holds, and slurm hands + us everything needed to check it locally: RANK is SLURM_PROCID and + LOCAL_RANK is SLURM_LOCALID, two independently reported numbers that agree + only under a block distribution. Comparing them costs a modulo on each + rank and no communication, and turns a launch this code cannot handle -- + `--distribution=cyclic` or `=arbitrary`, or a ragged --ntasks-per-node -- + into an immediate error on the ranks affected, instead of groups whose + all-gather assembles the right bytes in the wrong order. + """ + + def __init__(self): + self.rank = int(os.environ["RANK"]) + self.world_size = int(os.environ["WORLD_SIZE"]) + self.local_rank = int(os.environ["LOCAL_RANK"]) + self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) + + L = self.local_world_size + if L < 1: + raise RuntimeError(f"LOCAL_WORLD_SIZE={L} must be at least 1") + if self.world_size % L: + raise RuntimeError( + f"WORLD_SIZE={self.world_size} is not a multiple of " + f"LOCAL_WORLD_SIZE={L}; launch with a uniform --ntasks-per-node" + ) + if self.rank % L != self.local_rank: + raise RuntimeError( + f"rank {self.rank} reports LOCAL_RANK={self.local_rank}, but a " + f"block distribution of {L} ranks per node puts it at " + f"{self.rank % L}; datastage needs slurm's default block " + "distribution (no --distribution=cyclic or =arbitrary)" + ) + + self.node_count = self.world_size // L + self.node_index = self.rank // L + + # group_ranks[l][node] = global rank of LOCAL_RANK l on that node. + # Ascending by construction, which matters: dist.new_group sorts the + # list it is given and derives each member's position in the group from + # that order, and the drainer maps all-gather output position to + # node_index. Those agree only while this stays sorted. + group_ranks = [[node * L + l for node in range(self.node_count)] + for l in range(L)] + + # new_group is collective: every rank creates every group, in the same + # order, but only ever uses its own. + self.groups = [dist.new_group(ranks) for ranks in group_ranks] + self.group = self.groups[self.local_rank] + + def describe(self): + return ( + f"{self.world_size} ranks = {self.node_count} nodes x " + f"{self.local_world_size} ranks/node" + ) + + +class FileLayout: + """Deterministic, identical-on-every-rank mapping of a file onto ranks. + + Slice `l` of the file belongs to LOCAL_RANK l's group; within that group, + node `i` reads sub-shard `i` and the all-gather hands every node the whole + slice, one `window` at a time. + + Every offset and length here is a multiple of `align` except the very last + byte range of the file. That matters: O_DIRECT writes are padded up to the + block size, so a padded write must never reach into a range some other rank + owns. Because only the final range of the final slice is unaligned, the + only padding that ever happens runs off the end of the file, and the + closing ftruncate trims it. + """ + + def __init__(self, size, topology, piece, align): + self.size = size + self.align = align + self.piece = piece + L = topology.local_world_size + N = topology.node_count + + self.slice_size = align_up(ceil_div(size, L), align) if size else 0 + self.slice_start = min(topology.local_rank * self.slice_size, size) + slice_end = min(self.slice_start + self.slice_size, size) + self.slice_len = max(slice_end - self.slice_start, 0) + + # Per-node sub-shard of this slice. + self.shard = align_up(ceil_div(self.slice_len, N), align) if self.slice_len else 0 + self.rounds = ceil_div(self.shard, self.piece) if self.shard else 0 + + def segment(self, node, round_index): + """(file_offset, length) that `node` contributes in `round_index`.""" + shard_start = min(node * self.shard, self.slice_len) + shard_len = min(self.shard, self.slice_len - shard_start) + base = round_index * self.piece + if base >= shard_len: + return self.slice_start + shard_start + base, 0 + length = min(self.piece, shard_len - base) + return self.slice_start + shard_start + base, length + + +class Stager: + def __init__(self, args, topology): + self.args = args + self.topo = topology + self.device = torch.device("cuda", torch.cuda.current_device()) + self.pool = ThreadPoolExecutor(max_workers=max(args.num_threads, 4)) + + self.align = BUFFER_ALIGN + self.dest_root = os.path.abspath(args.destination) + # buffer_size is what each rank reads per round, exactly as fastcp's + # --buffer-size is what each thread reads. The assembled window that + # every rank must hold is node_count of those, so memory scales with the + # job: an all-gather delivers the whole window to every participant. + self.piece = args.buffer_size + window = self.piece * topology.node_count + # How many segments come back from the device per host staging chunk. + self.group_nodes = max(1, min(topology.node_count, CHUNK_TARGET // self.piece)) + chunk_bytes = self.group_nodes * self.piece + + # Fail before allocating rather than dying inside a CUDA OOM. Only the + # device side scales with the node count now, so this is the check that + # matters; 60% leaves room for the NCCL buffers and the CUDA context. + device_bytes = self.piece + RECV_DEV_SLOTS * window + budget = int(torch.cuda.get_device_properties(self.device).total_memory * 0.60) + if device_bytes > budget: + # device_bytes is piece * (RECV_DEV_SLOTS * node_count + 1) -- the + # send buffer is the +1. Dividing by RECV_DEV_SLOTS * node_count + # drops it and suggests a size that is itself over budget, so the + # advice reproduces the error it came with. + per_rank = align_down( + budget // (RECV_DEV_SLOTS * topology.node_count + 1), self.align) + if per_rank < self.align: + raise RuntimeError( + f"{topology.node_count} nodes need " + f"{self.align * (RECV_DEV_SLOTS * topology.node_count + 1) / 1024 ** 3:.1f} " + f"GiB of device memory even at the {self.align >> 20} MiB minimum " + f"--buffer-size, over the {budget / 1024 ** 3:.1f} GiB budget. " + "This node count does not fit on this GPU." + ) + per_window = per_rank + raise RuntimeError( + f"--buffer-size {self.piece >> 20}M across {topology.node_count} nodes " + f"needs {device_bytes / 1024 ** 3:.1f} GiB of device memory, over the " + f"{budget / 1024 ** 3:.1f} GiB budget. Lower --buffer-size to at most " + f"{per_window >> 20}M." + ) + + self.send_host = [pinned_aligned(self.piece, self.align) for _ in range(SEND_SLOTS)] + self.chunk_host = [pinned_aligned(chunk_bytes, self.align) for _ in range(CHUNK_SLOTS)] + self.send_dev = torch.empty(self.piece, dtype=torch.uint8, device=self.device) + self.recv_dev = [torch.empty(window, dtype=torch.uint8, device=self.device) + for _ in range(RECV_DEV_SLOTS)] + + # The copy-back gets a stream of its own. non_blocking=True only means + # the host does not wait; the copy still queues on whatever stream it + # was issued to, so on the default stream it runs after the next round's + # collective rather than beside it on the copy engine -- which leaves + # the second window RECV_DEV_SLOTS pays for doing nothing. + self.drain_stream = torch.cuda.Stream(device=self.device) + + if topology.rank == 0: + pinned = SEND_SLOTS * self.piece + CHUNK_SLOTS * chunk_bytes + print(f"datastage: {self.piece >> 20} MiB per rank per round, " + f"{window / 1024 ** 3:.2f} GiB window; per rank " + f"{device_bytes / 1024 ** 3:.2f} GiB device, " + f"{pinned / 1024 ** 3:.2f} GiB pinned host " + f"({self.group_nodes} segments per {chunk_bytes >> 20} MiB chunk)", + flush=True) + + def _read_piece(self, fd, mview, offset, length, block_size): + """Fill mview[:length] from fd at offset, split across reader threads.""" + if length == 0: + return + nthreads = max(1, min(self.args.num_threads, length // MIN_READ_PIECE)) + if nthreads == 1: + direct_io.pread(fd, mview, length, offset, block_size, 0) + return + # Aligning the per-thread piece to `align` (not to some larger read + # block) is what keeps this from silently collapsing to one thread when + # the round is small. + per = max(align_down(length // nthreads, self.align), self.align) + futures = [] + start = 0 + tid = 0 + while start < length: + count = min(per, length - start) + # Hand each read the whole remaining buffer, not just `count`: + # direct_io pads the request up to the block size and asserts the + # view is big enough to hold the padding. + futures.append( + self.pool.submit( + direct_io.pread, fd, mview[start:], count, + offset + start, block_size, tid, + ) + ) + start += count + tid += 1 + for future in futures: + future.result() + + def _write_one(self, ready, fd, mview, base, stride, length, offset, block_size): + """Write one segment, once the copy that filled its chunk has landed. + + Waiting here rather than on the draining thread means the next chunk's + device-to-host copy can be issued while this one is still being written. + """ + ready.synchronize() + return direct_io.pwrite(fd, mview[base:base + stride], length, offset, block_size) + + def stage_file(self, src, dst, size, mtime_ns): + tmp = f"{dst}.datastage.tmp.{os.environ.get('SLURM_JOB_ID', 'nojob')}" + try: + self._stage_to_temp(src, tmp, size, mtime_ns, dst) + except BaseException: + # The temp file is preallocated to the full source size, and only + # the success path's rename ever removed it, so a failed copy left + # a near-full file on node-local scratch -- one per attempt, since + # the name is scoped to the job id, until a resubmit loop filled + # the NVMe and started failing for a different reason. + # + # Any rank that raises unlinks it. Peers still hold it open, but + # POSIX keeps the inode alive until they close, so the space comes + # back when they die; and the rename that would have published it + # is not going to happen now anyway. Ranks parked in a collective + # never reach this, which is what the NCCL watchdog is for. + try: + os.unlink(tmp) + except OSError: + pass + raise + + def _stage_to_temp(self, src, tmp, size, mtime_ns, dst): + topo = self.topo + + # One rank per node creates the file, so the others can open it without + # racing on O_CREAT and without truncating each other's writes. + if topo.local_rank == 0: + os.makedirs(os.path.dirname(tmp), exist_ok=True) + self._chmod_parents(dst) + fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, self.args.chmod) + try: + if size: + os.ftruncate(fd, size) + os.fchmod(fd, self.args.chmod) + finally: + os.close(fd) + dist.barrier() + + fd_src = open_maybe_direct(src, os.O_RDONLY) + fd_dst = open_maybe_direct(tmp, os.O_WRONLY) + try: + src_block = os.fstatvfs(fd_src).f_bsize + dst_block = os.fstatvfs(fd_dst).f_bsize + assert self.align >= max(src_block, dst_block), "alignment below fs block size" + + layout = FileLayout(size, topo, self.piece, self.align) + if layout.rounds: + self._run_pipeline(fd_src, fd_dst, layout, src_block, dst_block) + os.fsync(fd_dst) + finally: + os.close(fd_src) + os.close(fd_dst) + + # Publish only once every rank on every node has written and synced. + dist.barrier() + if topo.local_rank == 0: + # O_DIRECT writes are padded up to the block size, so the last write + # of the file runs past its end. Trim it back -- after the barrier, + # because until then another rank may still be padding. + fd = os.open(tmp, os.O_WRONLY) + try: + os.ftruncate(fd, size) + os.fsync(fd) + finally: + os.close(fd) + os.utime(tmp, ns=(mtime_ns, mtime_ns)) + os.rename(tmp, dst) + parent = os.open(os.path.dirname(dst), os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(parent) + finally: + os.close(parent) + + def _chmod_parents(self, dst): + """Make the directories we created world-writable. + + os.makedirs applies the umask, which would leave subdirectories 0755 on + node-local scratch that is shared between users -- the same trap the + rsync path avoided with `umask 0000 --chmod=ugo+rwx`. + + Bounded to the destination root: these steps run with + --container-remap-root, so walking further up would happily chmod real + directories outside the staging tree. + """ + path = os.path.dirname(dst) + while path.startswith(self.dest_root + os.sep) or path == self.dest_root: + try: + os.chmod(path, self.args.chmod) + except OSError: + break + parent = os.path.dirname(path) + if parent == path: + # dirname("/") is "/", so a destination root of "/" would + # otherwise satisfy the loop condition forever -- chmod'ing the + # container root while every peer waits in the barrier below, + # until the job hits its wall clock. + break + path = parent + + def _run_pipeline(self, fd_src, fd_dst, layout, src_block, dst_block): + topo = self.topo + free_q = queue.Queue() + filled_q = queue.Queue() + recv_free_q = queue.Queue() + drain_q = queue.Queue() + for slot in range(SEND_SLOTS): + free_q.put(slot) + for slot in range(RECV_DEV_SLOTS): + recv_free_q.put(slot) + failure = [] + stop = threading.Event() + + def reader(): + try: + for round_index in range(layout.rounds): + slot = free_q.get() + if stop.is_set(): + break + offset, length = layout.segment(topo.node_index, round_index) + _, _, mview = self.send_host[slot] + self._read_piece(fd_src, mview, offset, length, src_block) + filled_q.put((round_index, slot, length)) + except BaseException as exc: # noqa: BLE001 - re-raised on main thread + failure.append(exc) + finally: + filled_q.put(None) + + def drainer(): + """Copy assembled windows back a chunk at a time and write them. + + Runs off the main thread so that the collective for round k+1 can + overlap the copy-back and writes for round k. A device window is + only released once all of its copies have landed; the writes carry + on behind that against the host chunk pool. + """ + # CUDA's current device is per host thread, and main() set it on + # the main thread only, so this one starts out on device 0. The + # copies below would still run on the right device -- the tensors + # carry it -- but torch.cuda.Event binds to the *calling thread's* + # current device when recorded, so the events would land on device + # 0's idle stream and every synchronize() on them would return + # immediately, letting the writers read chunks the copies have not + # filled yet. + torch.cuda.set_device(self.device) + + chunk_pending = [[] for _ in range(CHUNK_SLOTS)] + try: # noqa: PLR1702 + while True: + item = drain_q.get() + if item is None: + break + dev_slot, round_index, assembled = item + # Order the copies behind the collective on the device + # rather than on the host: the copies below can be queued + # while the all-gather that fills this window is still + # running, and the GPU holds them until it finishes. + self.drain_stream.wait_event(assembled) + last_copy = None + for group, first in enumerate( + range(0, topo.node_count, self.group_nodes)): + # Deliberately not `stop`: that only means the main loop + # has stopped feeding us, which is the normal end of a + # file. Windows already queued still have to be written. + if failure: + break + count = min(self.group_nodes, topo.node_count - first) + nbytes = count * layout.piece + chunk = group % CHUNK_SLOTS + for future in chunk_pending[chunk]: + future.result() + chunk_pending[chunk] = [] + + _, chunk_view, chunk_mview = self.chunk_host[chunk] + base = first * layout.piece + with torch.cuda.stream(self.drain_stream): + chunk_view[:nbytes].copy_( + self.recv_dev[dev_slot][base:base + nbytes], + non_blocking=True) + copied = torch.cuda.Event(blocking=True) + copied.record() + last_copy = copied + + submitted = [] + for index in range(count): + offset, length = layout.segment(first + index, round_index) + if length == 0: + continue + submitted.append(self.pool.submit( + self._write_one, copied, fd_dst, chunk_mview, + index * layout.piece, layout.piece, length, offset, + dst_block)) + chunk_pending[chunk] = submitted + + # The device window is reusable once its copies have landed, + # which is well before the writes behind them finish. + # + # This is also what keeps the window safe now that the + # copies read it from their own stream: the next all-gather + # is issued on the default stream, which has no ordering + # against drain_stream, so the only thing stopping it from + # overwriting a window still being copied is that the slot + # does not go back on recv_free_q until the host has seen + # the last copy complete. + if last_copy is not None: + last_copy.synchronize() + recv_free_q.put(dev_slot) + for futures in chunk_pending: + for future in futures: + future.result() + except BaseException as exc: # noqa: BLE001 - re-raised on main thread + failure.append(exc) + stop.set() + # Never leave the main thread parked waiting for a window. + for _ in range(RECV_DEV_SLOTS): + recv_free_q.put(0) + finally: + # No write may still be in flight when we return: stage_file + # closes the destination fd as soon as _run_pipeline does. + for futures in chunk_pending: + for future in futures: + try: + future.result() + except BaseException: # noqa: BLE001 - already recorded + pass + + reader_thread = threading.Thread(target=reader, name="datastage-reader") + drain_thread = threading.Thread(target=drainer, name="datastage-drainer") + reader_thread.start() + drain_thread.start() + + try: + for round_index in range(layout.rounds): + item = filled_q.get() + if item is None or failure: + break + _, slot, length = item + + _, send_view, _ = self.send_host[slot] + self.send_dev.copy_(send_view, non_blocking=True) + h2d_done = torch.cuda.Event(blocking=True) + h2d_done.record() + + # Safe to launch before the copy has landed: the collective is + # queued behind it on the same stream. Waiting on this one + # event, rather than the whole device, is what lets the reader + # refill the slot while the collective is still running. + dev_slot = recv_free_q.get() + dist.all_gather_into_tensor(self.recv_dev[dev_slot], self.send_dev, + group=topo.group) + assembled = torch.cuda.Event(blocking=True) + assembled.record() + h2d_done.synchronize() + free_q.put(slot) + + # Hand the window off; the main thread does no host copies and + # no writes, so it goes straight back to the next collective. + drain_q.put((dev_slot, round_index, assembled)) + finally: + # Let the drainer finish the windows already queued before anything + # is torn down -- it lags the main loop by design. + drain_q.put(None) + drain_thread.join() + # If we left the loop early the reader may be parked on free_q; + # release it and let it observe the stop flag. filled_q and drain_q + # are unbounded, so neither thread can block on the other side. + stop.set() + for slot in range(SEND_SLOTS): + free_q.put(slot) + reader_thread.join() + + if failure: + raise failure[0] + + +def build_plan(args): + """Rank 0 walks the source tree; everyone else takes its answer verbatim. + + Verbatim, and unverified: there is deliberately no cross-rank check that + the other ranks see the same sizes and mtimes rank 0 saw. Such a check + costs one stat per rank per file -- at 2048 nodes and 100k files, 1.6 + billion metadata operations issued as a synchronised burst, which is the + worst pattern there is for a Lustre MDT and would take it down for every + other job on the filesystem, not just this one. + + What it would detect is a node with a stale handle or the wrong dataset + mounted, and that condition is per-mount: it affects every file on that + node identically, so the cost scales with the file count while the + detection power does not. Verifying mounts belongs to mountcheck.py, once + per job, where a sparse SHA256 fingerprint is both cheaper and stronger + than size and mtime. + """ + if dist.get_rank() == 0: + # A planning failure has to be broadcast rather than raised here: every + # other rank is already waiting in the broadcast below, so anything + # raised instead of broadcast does not surface as an error at all -- + # rank 0 exits, its peers wait for a message that never comes, and the + # job dies on an NCCL watchdog timeout naming neither the file nor the + # reason, ten minutes of allocation later. + # + # So catch everything, not just the failure we went looking for. The + # stat below re-stats files plan_copy_operations already stat'd, and + # loses that race against anything modifying shared storage; and + # plan_copy_operations rejects bad arguments outright. Neither is an + # UnreadableEntries. The bar is reaching the broadcast, not + # anticipating the cause. + try: + jobs = plan_copy_operations(args.sources, args.destination) + payload = [[(src, dst, size, os.stat(src).st_mtime_ns) + for src, dst, size in jobs]] + except Exception as exc: # noqa: BLE001 - re-raised on every rank below + payload = [{"error": f"{type(exc).__name__}: {exc}"}] + else: + payload = [None] + dist.broadcast_object_list(payload, src=0) + if isinstance(payload[0], dict): + raise RuntimeError(f"cannot stage the source tree: {payload[0]['error']}") + return payload[0] + + +def parse_args(argv=None): + prog = "python3 -m mlperf_common.fileio.datastage" + parser = argparse.ArgumentParser( + prog=prog, + description="Collectively stage SOURCE(s) onto node-local storage on every node.", + usage=f"""{prog} [OPTION]... SOURCE DEST + {prog} [OPTION]... SOURCE... DIRECTORY + {prog} [OPTION]... -t DIRECTORY SOURCE...""", + ) + parser.add_argument("-t", "--target-directory", metavar="DIRECTORY", + help="copy all SOURCE arguments into DIRECTORY") + parser.add_argument("-r", "--recursive", action="store_true", + help="copy directories recursively") + parser.add_argument("files", nargs="+", + help="source file(s) and destination (or just sources if -t is used)") + parser.add_argument("-n", "--num-threads", type=int, default=16, + help="reader/writer threads per rank (default: 16). Only helps " + "while the per-round read is large; at high node counts the " + "window is divided thinly enough that reads are single-threaded.") + parser.add_argument("-b", "--buffer-size", type=parse_size, default=parse_size("8M"), + help="bytes each rank reads per round (default: 8M, must be a " + "multiple of 2MiB). The all-gather assembles one of these " + "per node into a window that every rank holds, so device " + "and pinned host memory scale with the node count: at 2048 " + "nodes 8M gives a 16 GiB window.") + parser.add_argument("--chmod", type=lambda v: int(v, 8), default=0o777, + help="octal mode for staged files and directories (default: 0777)") + parser.add_argument("--dry-run", action="store_true", + help="print the copy plan and exit") + + args = parser.parse_args(argv) + + if args.target_directory: + args.sources = args.files + args.destination = args.target_directory + else: + args.sources = args.files[:-1] + args.destination = args.files[-1] + if not args.sources: + parser.error(f"missing destination file operand after '{args.destination}'") + # O_DIRECT wants alignment, and the all-gather wants a uniform count, so + # round up rather than rejecting -- same as fastcp. + if args.buffer_size % BUFFER_ALIGN != 0: + args.buffer_size = align_up(args.buffer_size, BUFFER_ALIGN) + print(f"{prog}: rounding buffer size up to {args.buffer_size >> 20} MiB", + file=sys.stderr) + # Shared with fastcp, and applied here rather than in build_plan because + # parse_args runs on every rank before the process group exists: a bad + # invocation has to kill the whole job uniformly, not leave rank 0 exiting + # while its peers block in a collective. + try: + validate_copy_args(args.sources, args.destination, + recursive=args.recursive, + into_directory=bool(args.target_directory)) + except CopyArgumentError as exc: + sys.exit(f"{prog}: {exc}") + return args + + +def main(argv=None): + args = parse_args(argv) + + # Same reasoning as the argument validation in parse_args: this runs on + # every rank before the process group exists and derives from environment + # alone, so a bad launch fails the whole job identically rather than + # leaving some ranks to block in a collective. + try: + env = dist_env.configure() + except dist_env.DistEnvError as exc: + sys.exit(f"datastage: {exc}") + + if args.dry_run and env.source == "single": + try: + jobs = plan_copy_operations(args.sources, args.destination) + except UnreadableEntries as exc: + sys.exit(f"datastage: {exc}") + for src, dst, size in jobs: + print(f"{src} -> {dst} ({size} bytes)") + return 0 + + # Before init_process_group, not after: a wrong MASTER_ADDR hangs inside + # the rendezvous, so anything printed afterwards never appears -- and the + # rendezvous is exactly what this line is for diagnosing. + if env.rank == 0 or os.environ.get("NV_MLPERF_DEBUG"): + print(f"datastage: {dist_env.describe(env)}", flush=True) + + torch.cuda.set_device(env.local_rank) + dist.init_process_group(backend="nccl", init_method="env://") + try: + topology = Topology() + is_root = topology.rank == 0 + if is_root: + print(f"datastage: {topology.describe()}", flush=True) + + jobs = build_plan(args) + total_bytes = sum(size for _, _, size, _ in jobs) + if is_root: + print(f"datastage: {len(jobs)} files, {total_bytes / 1e9:.2f} GB", flush=True) + if args.dry_run: + if is_root: + for src, dst, size, _ in jobs: + print(f"{src} -> {dst} ({size} bytes)") + return 0 + + stager = Stager(args, topology) + started = time.monotonic() + for src, dst, size, mtime_ns in jobs: + file_started = time.monotonic() + try: + stager.stage_file(src, dst, size, mtime_ns) + except BaseException as exc: # noqa: BLE001 - annotate then re-raise + # Peers are blocked in a collective and will be torn down by the + # NCCL watchdog, so say which rank and which file went wrong. + print(f"datastage: rank {topology.rank} failed on {src}: {exc}", + file=sys.stderr, flush=True) + raise + if is_root: + elapsed = time.monotonic() - file_started + rate = size / 1e9 / elapsed if elapsed > 0 else 0.0 + print(f"STAGE {src} bytes={size} secs={elapsed:.3f} GB/s={rate:.2f}", + flush=True) + dist.barrier() + if is_root: + elapsed = time.monotonic() - started + rate = total_bytes / 1e9 / elapsed if elapsed > 0 else 0.0 + print(f"DONE files={len(jobs)} secs={elapsed:.3f} GB/s={rate:.2f}", flush=True) + finally: + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/mlperf_common/fileio/direct_io.py b/mlperf_common/fileio/direct_io.py new file mode 100644 index 0000000..99424ed --- /dev/null +++ b/mlperf_common/fileio/direct_io.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import ctypes + +def print_memoryview(mv): + """Print the address and size of a memoryview.""" + address = ctypes.addressof(ctypes.c_char.from_buffer(mv)) + size = mv.nbytes + print(f"Address: {hex(address)}, Size: {hex(size)}") + +def allocate_aligned_buffers(buffer_size, alignment, num_bufs): + """Allocate a sequence of aligned buffers. + + Args: + buffer_size (int): Size of each buffer in bytes. + alignment (int): Alignment boundary in bytes. + num_bufs (int): Number of buffers to allocate. + + Returns: + list[memoryview]: A list of memoryview objects, each representing an aligned buffer + raw_buf: the ctypes raw string_buffer. This *must* be kept to avoid the + gc prematurely collecting the buffers. + """ + assert buffer_size % alignment == 0, "buffer_size must be a multiple of alignment" + + # Allocate a single large buffer that can accommodate all buffers with proper alignment + total_size = (num_bufs * buffer_size) + alignment + raw_buf = ctypes.create_string_buffer(total_size) + base_address = ctypes.addressof(raw_buf) + + # Calculate the offset to align the first buffer + first_offset = (alignment - (base_address % alignment)) % alignment + aligned_base_address = base_address + first_offset + + # Create memoryviews for each buffer + buffers = [] + for i in range(num_bufs): + buf_address = aligned_base_address + i * buffer_size + aligned_buf = (ctypes.c_char * buffer_size).from_address(buf_address) + mv = memoryview(aligned_buf) + # print_memoryview(mv) + buffers.append(mv) + return buffers, raw_buf + +def round_up(value, multiple): + """Round up `value` to the next multiple of `multiple`.""" + if multiple == 0: + raise ValueError("Multiple must be greater than 0.") + return ((value + multiple - 1) // multiple) * multiple + +def pread(fd, aligned_memview, count, offset, fs_block_size, thread_id): + """Perform a direct pread. + + Returns: + number of bytes read. Unlike Posix pread this routine handles + interrupts, so return value should always be equal to count unless + there is an unrecoverable error, in which case it throws rather than returns + + Required: `count` <= aligned_memview size + """ + padded_count = round_up(count, fs_block_size) + assert padded_count <= aligned_memview.nbytes, "memview too small for requested read" + + while True: + try: + bytes_read = os.preadv(fd, [aligned_memview[:padded_count]], offset) + + if bytes_read == count: + # Expected case: all requested bytes were read. + return bytes_read + + if bytes_read > 0: + # Retry case: partial read, retry the entire read to maintain alignment restrictions + assert bytes_read < count, "bytes_read cannot exceed count" + continue + + if bytes_read == 0: + # Unexpected EOF + raise RuntimeError(f"Unexpected EOF encountered at offset {offset}") + + if bytes_read < 0: + raise OSError("preadv returned a negative value") + + except InterruptedError: + # Retry on EINTR + continue + except OSError as e: + # Non-retriable errors + if e.errno == 14: # Errno 14 corresponds to "Bad address" + raise RuntimeError(f"preadv failed with error: {e}, this usually means out of memory") + raise RuntimeError(f"preadv failed with error: {e}") + +def pwrite(fd, aligned_memview, count, offset, fs_block_size): + """Perform a direct pwrite + + Returns: + number of bytes written. Unlike Posix pwrite this routine handles + interrupts, so return value should always be equal to count unless + there is an unrecoverable error, in which case it throws rather than + returns + + Required: `count <= aligned_memview size + """ + padded_count = round_up(count, fs_block_size) + assert padded_count <= aligned_memview.nbytes, "memview too small for requested write" + while True: + try: + bytes_written = os.pwritev(fd, [aligned_memview[:padded_count]], offset) + + if bytes_written == padded_count: + # Expected case: all requested bytes were written. + # in the case of the last block in the file we may have padded up, so return + # the bytes _requested_ rather than the padded count + return count + + if bytes_written > 0: + # Retry case: partial write, retry the entire write to maintain alignment restrictions + assert bytes_written < count, "bytes_written cannot exceed count" + continue + + if bytes_written == 0: + # Unexpected failure to write + raise RuntimeError(f"Unexpected failure to write at offset {offset}") + + if bytes_written < 0: + raise OSError("pwritev returned a negative value") + + except InterruptedError: + # Retry on EINTR + continue + except OSError as e: + # Non-retriable errors + if e.errno == 14: # Errno 14 corresponds to "Bad address" + raise RuntimeError(f"pwritev failed with error: {e}, this usually means out of memory") + raise RuntimeError(f"pwritev failed with error: {e}") diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..330965d --- /dev/null +++ b/tests/README.md @@ -0,0 +1,93 @@ +# tests + +Stdlib only, no pytest, no numpy. Run them with: + + python3 tests/run_tests.py + +or run any one directly: + + python3 tests/test_pipeline.py + +Each test runs in its own interpreter, because each installs its own fake +`torch` into `sys.modules` and patches module-level names in `datastage`. + +`test_dist_env.py` is the exception: `mlperf_common.dist_env` is stdlib-only, +so it imports the real module and passes plain dicts where the environment +would go. No stubs involved. + +## What is here + +| file | covers | +| --- | --- | +| `test_layout.py` | `FileLayout` tiles each file exactly once, with aligned boundaries | +| `test_copyplan.py` | tree walk, src→dst mapping, refusal to plan an unreadable tree, cp argument semantics | +| `test_buildplan.py` | rank 0 broadcasts planning failures instead of raising them past its blocked peers | +| `test_topology.py` | block-distribution grouping: one rank per node per group, ascending, and non-block launches refused | +| `test_dist_env.py` | SLURM/OMPI → torch `env://` translation, slurm hostlist parsing, and the launches we refuse to guess at | +| `test_pipeline.py` | `stage_file` end to end: bytes in == bytes out, and no temp file left after a failure | +| `test_stager.py` | `_chmod_parents` terminates (including on `/`), and the memory-budget error suggests a size that works | +| `test_device.py` | events land on this rank's device; the drainer copies on its own stream, ordered behind the collective | +| `stubs.py` | fake `torch` / `torch.distributed` so the above run on a CPU | + +## Why the bytes are compared, not just the exit status + +`datastage` moves buffers between a reader thread, a collective, a draining +thread and a writer pool through four queues. Getting that handoff wrong does +not raise — it produces a file of exactly the right length containing the wrong +bytes. Two such bugs were caught this way: + +* the drain queue was emptied on shutdown before it had been processed, so + windows already queued were silently dropped. How much of a file survived + depended on timing. +* nothing trimmed the O_DIRECT write padding off the end of a file, so any file + whose size was not a multiple of the filesystem block size ended up longer + than the source, with garbage on the end. + +Both pass a "did it throw?" check and both fail a "do the bytes match?" check. +If you change the pipeline, confirm a deliberately reintroduced bug still makes +`test_pipeline.py` fail — a green test that cannot go red is worth nothing. + +## What these do NOT cover + +`stubs.py` replaces the GPU and the job. So none of this exercises: + +* real NCCL, or more than one process — the multi-node cases fake the + all-gather by filling each node's segment with the bytes that node would have + read, which is what the real collective delivers, but no data crosses a wire +* real CUDA events, so no ordering bug between a copy and its consumer can show + up here; on the CPU every "copy" has already landed. `test_device.py` covers + the one part of this that *is* checkable without a GPU — which device each + event was recorded against, since `torch.cuda`'s current device is per host + thread and a thread that forgets `set_device` gets device 0. That caught a + real corruption bug in the drainer. It says nothing about whether the events + order the work correctly once they are on the right device +* pinned memory, and whether it is aligned enough for a real filesystem block + size +* O_DIRECT itself — the tests reopen files buffered, since the fake buffers + cannot promise the alignment it demands +* Lustre, node-local NVMe, and anything about throughput + +A green run means the arithmetic and the choreography are right. It says +nothing about whether staging works on a cluster. + +## The other half: `cluster-selftest.sh` + +Not part of `run_tests.py` — it needs an allocation. It builds a dataset with +the shapes that have historically broken things (sizes either side of the +2 MiB boundary, a zero-length file, files large enough to need several rounds, +64 small ones, a symlink), stages it with a real multi-node `srun`, and then +checks on **every node** that the node-local copy matches the source: + +```bash +salloc -N2 --ntasks-per-node=4 ... +tests/cluster-selftest.sh /lustre/scratch/me/selftest /raid/scratch/me/selftest +``` + +It compares two things, not one. `fastmd5` emits one line per GB-chunk, so a +zero-length file produces no lines at all and a destination missing it would +compare equal on checksums alone — the size-and-path listing is what catches +that. It also asserts no `.datastage.tmp.*` files survive and that everything +is mode 0777. + +This is the only thing that exercises real NCCL, real CUDA events, pinned +memory against a real block size, and O_DIRECT. Nothing above does. diff --git a/tests/cluster-selftest.sh b/tests/cluster-selftest.sh new file mode 100755 index 0000000..22afb39 --- /dev/null +++ b/tests/cluster-selftest.sh @@ -0,0 +1,248 @@ +#!/bin/bash + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +############################################################################### +# End-to-end check of datastage on real hardware. +# +# The suite in this directory is stdlib-only and single-process: it fakes torch +# before importing datastage, so it never loads CUDA or NCCL and never talks to +# another rank. Running it on a GPU node changes nothing -- launching it under +# srun just gets you N independent copies of the same CPU test. Nothing in it +# exercises a real collective, real CUDA events, pinned memory against a real +# block size, or O_DIRECT. +# +# This script is the other half: one real staging job, verified by comparing +# checksums of what came out against what went in, on every node. +# +# Usage, from inside an allocation: +# +# salloc -N2 --ntasks-per-node=4 ... +# tests/cluster-selftest.sh /lustre/scratch/me/selftest /raid/scratch/me/selftest +# +# arg 1 is a directory on the shared filesystem (the source, and where results +# are collected) +# arg 2 is a directory on node-local storage (the destination) +# +# Everything under both is deleted and rebuilt unless --keep-dataset is passed. +############################################################################### + +set -euo pipefail + +usage() { + sed -n '18,40p' "$0" | sed 's/^# \?//' | grep -v '^#*$' + exit "${1:-1}" +} + +keep_dataset=0 +args=() +for arg in "$@"; do + case "${arg}" in + --keep-dataset) keep_dataset=1 ;; + -h|--help) usage 0 ;; + -*) echo "unknown option ${arg}" >&2; usage ;; + *) args+=("${arg}") ;; + esac +done +[[ "${#args[@]}" -eq 2 ]] || usage + +readonly SHARED="${args[0]}" +readonly LOCAL="${args[1]}" +readonly REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +readonly DATASET="${SHARED}/dataset" +readonly RESULTS="${SHARED}/results" + +if [[ -z "${SLURM_JOB_ID:-}" ]]; then + echo "ERROR: no SLURM_JOB_ID; run this inside an salloc or sbatch" >&2 + exit 1 +fi + +# One task per GPU is how datastage is meant to run. DGXNGPU if the site sets +# it, otherwise count what this node has. +readonly NGPUS="${DGXNGPU:-$(nvidia-smi -L | wc -l)}" +readonly NNODES="${SLURM_JOB_NUM_NODES:-1}" +export PYTHONPATH="${REPO}${PYTHONPATH:+:${PYTHONPATH}}" + +echo "==============================================================" +echo "datastage cluster self-test" +echo " repo ${REPO}" +echo " nodes ${NNODES} x ${NGPUS} ranks" +echo " source ${DATASET} (shared)" +echo " destination ${LOCAL} (node-local)" +echo "==============================================================" + +############################################################################### +# 1. Build a dataset with the shapes that have historically broken things. +############################################################################### +if [[ "${keep_dataset}" -eq 1 && -d "${DATASET}" ]]; then + echo "--- 1. reusing existing dataset" +else + echo "--- 1. building dataset" + rm -rf "${DATASET}" + mkdir -p "${DATASET}" + python3 - "${DATASET}" <<'PY' +import os, sys +root = sys.argv[1] +MiB = 1024 ** 2 + +# Sizes clustered around the 2 MiB alignment boundary, where the O_DIRECT write +# padding and the closing ftruncate interact, plus two large enough to need +# several rounds through the pipeline at any plausible node count. +sizes = { + "empty.bin": 0, + "one.bin": 1, + "small.bin": 1000, + "align-1.bin": 2 * MiB - 1, + "align.bin": 2 * MiB, + "align+1.bin": 2 * MiB + 1, + "ragged.bin": 5 * MiB + 12345, + "sub/medium.bin": 64 * MiB + 7, + "sub/deep/big.bin": 200 * MiB, +} +for name, size in sizes.items(): + path = os.path.join(root, name) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "wb") as handle: + remaining = size + while remaining: + block = min(remaining, 8 * MiB) + handle.write(os.urandom(block)) + remaining -= block + +# A pile of small files: one metadata operation each, and at high node counts +# each costs a whole window of collective traffic regardless of its size. +os.makedirs(os.path.join(root, "many"), exist_ok=True) +for i in range(64): + with open(os.path.join(root, "many", f"f{i:03d}.bin"), "wb") as handle: + handle.write(os.urandom(4096 + i)) + +# Symlinks are dereferenced and copied as content, not recreated as links. +os.symlink("align.bin", os.path.join(root, "link-to-file")) + +total = sum(sizes.values()) + sum(4096 + i for i in range(64)) +print(f" {len(sizes) + 64 + 1} files, {total / 1e6:.1f} MB") +PY +fi + +############################################################################### +# 2. Checksum the source once, from the shared filesystem. +############################################################################### +echo "--- 2. checksumming source" +rm -rf "${RESULTS}" +mkdir -p "${RESULTS}" +"${REPO}/client/fastmd5" "${DATASET}" \ + | sed "s|^${DATASET}/||" | sort > "${RESULTS}/source.md5" + +# Checksums alone are not enough. fastmd5 emits one line per GB-chunk, so a +# zero-length file produces no lines at all -- and a destination missing +# empty.bin entirely would compare equal. The size+path listing closes that, +# and catches extra files and truncations too. +# +# mtime is deliberately not compared: datastage does set it from the source, +# but timestamp granularity varies by filesystem and a false failure here on a +# first run would cost more than the check is worth. +( cd "${DATASET}" && find . -type f -printf '%s\t%p\n' | sort ) \ + > "${RESULTS}/source.list" +echo " $(wc -l < "${RESULTS}/source.md5") chunk checksums, \ +$(wc -l < "${RESULTS}/source.list") files" + +############################################################################### +# 3. Clear the destination on every node. +############################################################################### +echo "--- 3. clearing destinations" +srun --ntasks-per-node=1 bash -c "rm -rf '${LOCAL}' && mkdir -p '${LOCAL}'" + +############################################################################### +# 4. Stage. No wrapper: datastage derives the rendezvous itself. +############################################################################### +echo "--- 4. staging" +staged_ok=1 +srun --ntasks-per-node="${NGPUS}" \ + python3 -m mlperf_common.fileio.datastage -r "${DATASET}" "${LOCAL}" \ + || staged_ok=0 +if [[ "${staged_ok}" -eq 0 ]]; then + echo "FAIL: datastage exited nonzero" >&2 + exit 1 +fi + +############################################################################### +# 5. Verify on every node, against the node-local copy. +############################################################################### +echo "--- 5. verifying every node" +srun --ntasks-per-node=1 bash -c " + set -e + host=\$(hostname -s) + '${REPO}/client/fastmd5' '${LOCAL}/dataset' \ + | sed 's|^${LOCAL}/dataset/||' | sort > '${RESULTS}/'\${host}.md5 + ( cd '${LOCAL}/dataset' && find . -type f -printf '%s\t%p\n' | sort ) \ + > '${RESULTS}/'\${host}.list + # F9: nothing should be left of the temp files staging writes through. + find '${LOCAL}' -name '*.datastage.tmp.*' > '${RESULTS}/'\${host}.temps + # --chmod defaults to 0777; anything tighter breaks a shared scratch dir. + find '${LOCAL}/dataset' \\! -perm -0777 > '${RESULTS}/'\${host}.modes +" + +############################################################################### +# 6. Compare. +############################################################################### +echo "--- 6. results" +failures=0 +for md5 in "${RESULTS}"/*.md5; do + host="$(basename "${md5}" .md5)" + [[ "${host}" == "source" ]] && continue + + if diff -q "${RESULTS}/source.list" "${RESULTS}/${host}.list" > /dev/null; then + : # same files, same sizes + else + echo " ${host}: FILE LIST DIFFERS (< missing, > unexpected)" + diff "${RESULTS}/source.list" "${RESULTS}/${host}.list" \ + | head -20 | sed 's/^/ /' + failures=$((failures + 1)) + fi + + if diff -q "${RESULTS}/source.md5" "${md5}" > /dev/null; then + echo " ${host}: bytes match" + else + echo " ${host}: CONTENT MISMATCH" + diff "${RESULTS}/source.md5" "${md5}" | head -20 | sed 's/^/ /' + failures=$((failures + 1)) + fi + + if [[ -s "${RESULTS}/${host}.temps" ]]; then + echo " ${host}: left temp files:" + sed 's/^/ /' "${RESULTS}/${host}.temps" + failures=$((failures + 1)) + fi + if [[ -s "${RESULTS}/${host}.modes" ]]; then + echo " ${host}: entries not mode 0777:" + head -10 "${RESULTS}/${host}.modes" | sed 's/^/ /' + failures=$((failures + 1)) + fi +done + +nodes_checked=$(ls -1 "${RESULTS}"/*.md5 | grep -cv 'source\.md5$') +if [[ "${nodes_checked}" -ne "${NNODES}" ]]; then + echo " only ${nodes_checked} of ${NNODES} nodes reported" >&2 + failures=$((failures + 1)) +fi + +echo "==============================================================" +if [[ "${failures}" -eq 0 ]]; then + echo "PASS: ${nodes_checked} node(s) hold byte-identical copies" +else + echo "FAIL: ${failures} problem(s); details under ${RESULTS}" +fi +echo "==============================================================" +exit "$((failures == 0 ? 0 : 1))" diff --git a/tests/run_tests.py b/tests/run_tests.py new file mode 100644 index 0000000..15f530c --- /dev/null +++ b/tests/run_tests.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run every test_*.py here in its own interpreter; exit nonzero if any failed. + +Separate processes because each test installs its own fake torch into +sys.modules and patches module-level names in datastage. +""" + +import glob +import os +import subprocess +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def main(): + failed = [] + for path in sorted(glob.glob(os.path.join(HERE, "test_*.py"))): + name = os.path.basename(path) + completed = subprocess.run([sys.executable, path], cwd=HERE, check=False) + if completed.returncode != 0: + failed.append(name) + print() + if failed: + print(f"FAILED: {', '.join(failed)}") + return 1 + print("all tests passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/stubs.py b/tests/stubs.py new file mode 100644 index 0000000..f069476 --- /dev/null +++ b/tests/stubs.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Enough of torch and torch.distributed to run mlperf_common.fileio on a CPU. + +datastage's real dependencies are a GPU and a NCCL job, neither of which is +available where these tests run. The parts worth testing without them are the +sharding arithmetic, the thread choreography, and whether the bytes that come +out match the bytes that went in -- all of which are independent of the device. + +So torch tensors become memoryviews over ctypes buffers, CUDA events become +no-ops, and all_gather_into_tensor becomes whatever the test wants it to be. +What this deliberately does NOT cover: real CUDA events, real NCCL, pinned +memory alignment against a real filesystem block size, and O_DIRECT itself. +""" + +import contextlib +import ctypes +import importlib.util +import os +import sys +import threading +import types + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +class FakeTensor: + """A flat byte buffer with the slice/copy_/numpy surface datastage uses. + + Slicing returns a view sharing the parent's memory, as torch does, which is + what makes the pinned_aligned() offset dance behave the same way here. + """ + + def __init__(self, nbytes=None, buf=None, offset=0, size=None): + if buf is None: + buf = ctypes.create_string_buffer(nbytes) + size = nbytes + self._buf = buf + self._offset = offset + self._size = size + + def _view(self): + return memoryview(self._buf).cast("B")[self._offset:self._offset + self._size] + + def __len__(self): + return self._size + + def __getitem__(self, item): + start, stop, step = item.indices(self._size) + assert step == 1, "datastage only takes contiguous slices" + return FakeTensor(buf=self._buf, offset=self._offset + start, size=stop - start) + + def copy_(self, other, non_blocking=False): + record_op("copy", self) + source = other._view() + self._view()[:len(source)] = source + return self + + def numpy(self): + return self._view() + + def data_ptr(self): + return ctypes.addressof(self._buf) + self._offset + + +class FakeDevice: + """Stand-in for torch.device('cuda', i). datastage only ever passes it on.""" + + def __init__(self, kind="cuda", index=0): + self.type = kind + self.index = index + + def __repr__(self): + return f"device(type={self.type!r}, index={self.index})" + + +# CUDA's current device and current stream are both per *host thread*, and a +# thread that never set either gets device 0 and that device's default stream, +# no matter what any other thread did. Modelling that faithfully is the whole +# point: it is the state a thread doing CUDA work can silently get wrong, and +# no amount of comparing staged bytes on a CPU will show it up. +_CURRENT = threading.local() + +# Every event that has been recorded, in order. test_device.py checks which +# device each one landed on; other tests ignore it. +EVENTS = [] + +# Every stream-ordered operation, in issue order, tagged with the thread that +# issued it and the stream it went to. What makes overlap possible is which +# stream work lands on, and that is visible here even though concurrency is +# not. +OPS = [] + + +class Op: + __slots__ = ("kind", "thread", "stream", "obj") + + def __init__(self, kind, stream, obj): + self.kind = kind + self.thread = threading.current_thread().name + self.stream = stream + self.obj = obj + + def __repr__(self): + return f"Op({self.kind}, thread={self.thread!r}, stream={self.stream})" + + +def record_op(kind, obj=None): + OPS.append(Op(kind, current_stream(), obj)) + + +class FakeStream: + """Stand-in for torch.cuda.Stream. + + Only the ordering surface datastage needs: work issued inside a stream + context belongs to that stream, and wait_event makes this stream wait for + an event recorded on another one without blocking the host. + """ + + def __init__(self, device=None, default=False): + self.device = device.index if isinstance(device, FakeDevice) else int(device or 0) + self.default = default + self.waited = [] + + def wait_event(self, event): + self.waited.append(event) + OPS.append(Op("wait", self, event)) + + def __repr__(self): + kind = "default" if self.default else "side" + return f"{kind}-stream(device={self.device})" + + +_DEFAULT_STREAMS = {} + + +def current_device(): + return getattr(_CURRENT, "index", 0) + + +def set_device(device): + _CURRENT.index = device.index if isinstance(device, FakeDevice) else int(device) + + +def current_stream(device=None): + explicit = getattr(_CURRENT, "stream", None) + if explicit is not None: + return explicit + index = current_device() if device is None else device + return _DEFAULT_STREAMS.setdefault(index, FakeStream(index, default=True)) + + +@contextlib.contextmanager +def stream(target): + """torch.cuda.stream(): make `target` the calling thread's current stream.""" + previous = getattr(_CURRENT, "stream", None) + _CURRENT.stream = target + try: + yield + finally: + _CURRENT.stream = previous + + +class FakeEvent: + """CUDA event stand-in: everything is synchronous on the CPU already. + + It does track one thing that is not synchronous, though. A real + torch.cuda.Event is created lazily and binds to the *calling thread's* + current device when it is recorded -- not to the device the work it is + meant to track ran on. So an event recorded by a thread that forgot + set_device belongs to an idle stream on device 0, and synchronize() on it + returns immediately while the copy it stands for is still in flight. + Remembering the device here is what lets a test see that. + """ + + def __init__(self, blocking=False): + self.device = None + self.stream = None + + def record(self): + self.device = current_device() + self.stream = current_stream() + EVENTS.append(self) + record_op("record", self) + + def synchronize(self): + pass + + +def install(total_memory=288 * 1024 ** 3): + """Put fake torch / torch.distributed modules into sys.modules.""" + torch = types.ModuleType("torch") + dist = types.ModuleType("torch.distributed") + torch.uint8 = "uint8" + torch.int64 = "int64" + torch.distributed = dist + torch.empty = lambda n, dtype=None, device=None, pin_memory=False: FakeTensor(n) + torch.device = lambda kind="cuda", index=0: FakeDevice(kind, index) + torch.cuda = types.SimpleNamespace( + Event=FakeEvent, + Stream=lambda device=None: FakeStream(device), + stream=stream, + current_stream=current_stream, + current_device=current_device, + set_device=set_device, + get_device_properties=lambda device: types.SimpleNamespace( + total_memory=total_memory), + ) + # Returns the rank list itself as the group handle, so a test can see which + # ranks a group was built from. Real new_group is collective and sorts its + # argument; Topology depends on already handing it a sorted list, and + # test_topology.py checks that. + dist.new_group = lambda ranks: list(ranks) + dist.barrier = lambda *args, **kwargs: None + dist.get_rank = lambda: 0 + dist.broadcast_object_list = lambda payload, src=0: None + dist.all_gather_into_tensor = lambda out, inp, group=None: out.copy_(inp) + sys.modules["torch"] = torch + sys.modules["torch.distributed"] = dist + return torch, dist + + +def load_fileio(): + """Import mlperf_common.fileio.* from this checkout, without installing it.""" + package_root = os.path.join(REPO_ROOT, "mlperf_common") + package = types.ModuleType("mlperf_common") + package.__path__ = [package_root] + sys.modules["mlperf_common"] = package + fileio = types.ModuleType("mlperf_common.fileio") + fileio.__path__ = [os.path.join(package_root, "fileio")] + sys.modules["mlperf_common.fileio"] = fileio + + loaded = {} + for name in ("direct_io", "copyplan", "datastage"): + spec = importlib.util.spec_from_file_location( + f"mlperf_common.fileio.{name}", + os.path.join(package_root, "fileio", f"{name}.py")) + module = importlib.util.module_from_spec(spec) + sys.modules[f"mlperf_common.fileio.{name}"] = module + spec.loader.exec_module(module) + setattr(fileio, name, module) + loaded[name] = module + return loaded diff --git a/tests/test_buildplan.py b/tests/test_buildplan.py new file mode 100644 index 0000000..edcddd3 --- /dev/null +++ b/tests/test_buildplan.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Rank 0 must never leave the planning step without telling the other ranks. + +build_plan has one rank walk the tree and broadcast the answer. Every other +rank is already blocked in that broadcast, so anything rank 0 raises instead of +broadcasting does not surface as an error: rank 0 exits, its peers wait for a +message that will never come, and the job dies ten minutes later on an NCCL +watchdog timeout naming neither the file nor the reason. + +So the requirement is not "handle the expected failure" but "reach the +broadcast no matter what". These tests drive rank 0's side and check that +failures arrive as a broadcastable payload rather than an exception. +""" + +import os +import shutil +import sys +import tempfile + +import stubs + +torch, dist = stubs.install() +modules = stubs.load_fileio() +ds = modules["datastage"] +copyplan = modules["copyplan"] + +results = [] + + +def check(name, condition, detail=""): + results.append((name, condition, detail)) + + +class Args: + def __init__(self, sources, destination): + self.sources = sources + self.destination = destination + + +def build_plan_outcome(args): + """Run rank 0's build_plan; report how it came back. + + Returns ("broadcast", message) if the failure was turned into something + every rank raises together, ("escaped", exc) if it got out of build_plan + before the broadcast, or ("ok", jobs). + """ + try: + return "ok", ds.build_plan(args) + except RuntimeError as exc: + # What build_plan raises after a broadcast payload carrying an error; + # every rank reaches this together. + return "broadcast", str(exc) + except BaseException as exc: # noqa: BLE001 - the failure under test + return "escaped", exc + + +def main(): + root = tempfile.mkdtemp(prefix="buildplan-") + original_plan = ds.plan_copy_operations + try: + source = os.path.join(root, "src") + os.makedirs(source) + for name in ("a.bin", "b.bin"): + with open(os.path.join(source, name), "wb") as handle: + handle.write(b"x" * 10) + destination = os.path.join(root, "dst") + os.makedirs(destination) + + kind, value = build_plan_outcome(Args([source], destination)) + check("a good tree plans", kind == "ok", f"got {kind}: {value}") + if kind == "ok": + check("the plan carries size and mtime", + len(value) == 2 and all(len(job) == 4 for job in value), + f"got {value}") + + # A file that disappears between plan_copy_operations' stat and the + # mtime stat below it. Rare, but it is exactly what a shared + # filesystem does under a concurrent job, and os.stat raises + # FileNotFoundError, which is not UnreadableEntries. + vanished = os.path.join(source, "gone.bin") + ds.plan_copy_operations = lambda sources, dest: [ + (vanished, os.path.join(destination, "gone.bin"), 10)] + kind, value = build_plan_outcome(Args([source], destination)) + check("a file vanishing mid-plan is broadcast, not raised", + kind == "broadcast", f"got {kind}: {value!r}") + if kind == "broadcast": + check("the broadcast error names the file", + "gone.bin" in value, f"got {value!r}") + ds.plan_copy_operations = original_plan + + # The failure build_plan already expected still works. + os.symlink("/nowhere", os.path.join(source, "dead")) + kind, value = build_plan_outcome(Args([source], destination)) + check("an unreadable entry is broadcast", kind == "broadcast", + f"got {kind}: {value!r}") + if kind == "broadcast": + check("the broadcast error names the unreadable entry", + "dead" in value, f"got {value!r}") + os.remove(os.path.join(source, "dead")) + + # Bad arguments reach here too, now that plan_copy_operations applies + # the cp rules itself -- CopyArgumentError is not UnreadableEntries + # either. + kind, value = build_plan_outcome( + Args([source, source], os.path.join(root, "absent"))) + check("a rejected argument combination is broadcast", + kind == "broadcast", f"got {kind}: {value!r}") + finally: + ds.plan_copy_operations = original_plan + shutil.rmtree(root, ignore_errors=True) + + failures = [r for r in results if not r[1]] + for name, _, detail in failures: + print(f" FAIL {name}{': ' + detail if detail else ''}") + print(f"test_buildplan: {len(results)} checks, {len(failures)} failures") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_copyplan.py b/tests/test_copyplan.py new file mode 100644 index 0000000..d07ef69 --- /dev/null +++ b/tests/test_copyplan.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""copyplan's walk, mapping, and its refusal to plan an unreadable tree. + +fastcp, fastmd5 and datastage all enumerate through here, so a disagreement +about which files a copy covers would mean verifying a staged tree against a +different set of files than was staged. +""" + +import os +import shutil +import sys +import tempfile + +import stubs + +stubs.install() +copyplan = stubs.load_fileio()["copyplan"] + +results = [] + + +def check(name, condition, detail=""): + results.append((name, condition, detail)) + + +def build_tree(root): + """A tree with each kind of entry the walk has to have an answer for.""" + os.makedirs(os.path.join(root, "sub")) + os.makedirs(os.path.join(root, "realdir")) + for path in ("a.bin", "sub/b.bin", "realdir/c.bin"): + with open(os.path.join(root, path), "wb") as handle: + handle.write(b"x" * 10) + os.symlink("a.bin", os.path.join(root, "link_to_file")) + os.symlink("realdir", os.path.join(root, "link_to_dir")) + + +def check_cp_semantics(root): + """The SOURCE/DEST rules, against what GNU cp actually does. + + fastcp was written to match cp, and datastage inherits the mapping through + this module, so these are the cases where "what does copy mean here" has to + have exactly one answer. The rows marked ERROR are the ones cp refuses. + """ + base = os.path.join(root, "sem") + src = os.path.join(base, "src") + os.makedirs(os.path.join(src, "sub")) + for path in ("a.bin", "sub/b.bin"): + with open(os.path.join(src, path), "wb") as handle: + handle.write(b"x" * 10) + for name in ("f1", "f2"): + with open(os.path.join(base, name), "wb") as handle: + handle.write(b"y" * 10) + existing = os.path.join(base, "existingdir") + os.makedirs(existing) + + f1, f2 = os.path.join(base, "f1"), os.path.join(base, "f2") + absent = os.path.join(base, "newdir") + + def plan(sources, destination, **kwargs): + jobs = copyplan.plan_copy_operations(sources, destination, **kwargs) + return sorted(os.path.relpath(dst, base) for _, dst, _ in jobs) + + # cp -r src newdir, newdir absent: the contents land directly in newdir, + # with no basename level. This is the case that used to plan a single + # 4 KiB "copy the directory as a file" job. + check("cp -r src newdir puts contents directly under newdir", + plan([src], absent) == ["newdir/a.bin", "newdir/sub/b.bin"], + f"got {plan([src], absent)}") + + # cp -r src existingdir: basename level is kept. + check("cp -r src existingdir keeps the basename level", + plan([src], existing) == ["existingdir/src/a.bin", + "existingdir/src/sub/b.bin"], + f"got {plan([src], existing)}") + + check("cp a b existingdir places both under it", + plan([f1, f2], existing) == ["existingdir/f1", "existingdir/f2"]) + + check("cp a newname renames a single file", + plan([f1], os.path.join(base, "renamed.bin")) == ["renamed.bin"]) + + for name, sources, destination, kwargs in ( + ("cp a b newdir is refused when newdir is absent", [f1, f2], absent, {}), + ("cp -r src file is refused", [src], f1, {}), + ("cp -t newdir src is refused when newdir is absent", + [src], absent, {"into_directory": True}), + ("a directory source without -r is refused", + [src], existing, {"recursive": False}), + ): + try: + copyplan.validate_copy_args(sources, destination, **kwargs) + check(name, False, "no exception raised") + except copyplan.CopyArgumentError: + check(name, True) + + # The planner must not depend on its caller having validated first. + try: + copyplan.plan_copy_operations([f1, f2], absent) + check("planning applies the rules itself", False, "no exception raised") + except copyplan.CopyArgumentError: + check("planning applies the rules itself", True) + + +def check_unlistable_directory(root): + """A subtree we cannot list must be reported, not quietly left out. + + This is the one unreadable-entry case with no entry to report: the files + under an unlistable directory never reach the walk, so there is nothing to + stat and nothing to collect. os.walk's default is to swallow the error and + carry on, which makes such a subtree indistinguishable from an empty one -- + a partial copy that exits 0, and a partial checksum that agrees with it. + """ + source = os.path.join(root, "blocked") + os.makedirs(os.path.join(source, "readable")) + os.makedirs(os.path.join(source, "secret")) + for path in ("readable/good.bin", "secret/hidden.bin"): + with open(os.path.join(source, path), "wb") as handle: + handle.write(b"x" * 10) + os.chmod(os.path.join(source, "secret"), 0o000) + + try: + # Root ignores the mode bits, so the premise would not hold and the + # check below would pass without testing anything. + try: + os.listdir(os.path.join(source, "secret")) + except PermissionError: + pass + else: + check("unlistable directory is reported", True, + "SKIPPED: this user can list a 0o000 directory") + return + + try: + found = copyplan.list_relative_files(source) + check("unlistable directory is reported", False, + f"no exception; returned {found}") + except copyplan.UnreadableEntries as exc: + check("unlistable directory is reported", True) + check("the message names the directory it could not list", + "secret" in str(exc), f"got {exc}") + + # Same hole one level up: the root of the walk itself. + os.chmod(source, 0o000) + try: + found = copyplan.list_relative_files(source) + check("unlistable root is reported", False, + f"no exception; returned {found}") + except copyplan.UnreadableEntries: + check("unlistable root is reported", True) + finally: + # Otherwise rmtree cannot clean up after us. + os.chmod(source, 0o755) + os.chmod(os.path.join(source, "secret"), 0o755) + + +def main(): + root = tempfile.mkdtemp(prefix="copyplan-") + try: + source = os.path.join(root, "src") + build_tree(source) + + found = copyplan.list_relative_files(source) + check("symlink to a file is followed and listed", "link_to_file" in found) + check("symlink to a directory is descended into", + "link_to_dir/c.bin" in found, f"got {found}") + check("the symlinked directory is not itself listed as a file", + "link_to_dir" not in found) + check("result is sorted", found == sorted(found)) + + destination = os.path.join(root, "dst") + os.makedirs(destination) + jobs = copyplan.plan_copy_operations([source], destination) + check("every planned file lands under dest//", + all(dst.startswith(os.path.join(destination, "src") + os.sep) + for _, dst, _ in jobs)) + check("plan is sorted by destination", + [j[1] for j in jobs] == sorted(j[1] for j in jobs)) + check("sizes come back with the plan", all(size == 10 for _, _, size in jobs)) + + single = copyplan.plan_copy_operations( + [os.path.join(source, "a.bin")], os.path.join(destination, "renamed.bin")) + check("file-to-file copy keeps the given destination name", + len(single) == 1 and single[0][1].endswith("renamed.bin")) + + # Dangling symlinks: os.walk lists them among the filenames because they + # are not directories, and stat then follows them to nothing. + os.symlink("/nowhere", os.path.join(source, "dead1")) + os.symlink("/also/nowhere", os.path.join(source, "sub", "dead2")) + try: + copyplan.list_relative_files(source) + check("dangling symlinks are refused", False, "no exception raised") + except copyplan.UnreadableEntries as exc: + check("dangling symlinks are refused", True) + check("every bad entry is reported, not just the first", + len(exc.entries) == 2, f"reported {len(exc.entries)}") + check("the message names the bad paths", + "dead1" in str(exc) and "dead2" in str(exc)) + try: + copyplan.plan_copy_operations([source], destination) + check("planning refuses an unreadable tree", False, "no exception raised") + except copyplan.UnreadableEntries: + check("planning refuses an unreadable tree", True) + + check_cp_semantics(root) + check_unlistable_directory(root) + finally: + shutil.rmtree(root, ignore_errors=True) + + failures = [r for r in results if not r[1]] + for name, _, detail in failures: + print(f" FAIL {name}{': ' + detail if detail else ''}") + print(f"test_copyplan: {len(results)} checks, {len(failures)} failures") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_device.py b/tests/test_device.py new file mode 100644 index 0000000..6e6af59 --- /dev/null +++ b/tests/test_device.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CUDA bookkeeping the pipeline gets to be wrong about on a CPU. + +Two properties, both invisible in the staged bytes and both able to corrupt a +real run: which *device* each event was recorded against, and which *stream* +the drainer's copies were issued on. + +Every CUDA event must be recorded against this rank's device, not device 0. + +torch.cuda's current device is per host thread. main() sets it once, on the +main thread, so every thread datastage starts afterwards inherits device 0 -- +and torch.cuda.Event binds to the *calling thread's* current device when it is +recorded. An event recorded on the wrong device belongs to an idle stream, so +Event.synchronize() on it returns immediately: the writer pool reads a pinned +chunk the device-to-host copy has not filled, and a device window is recycled +while copies are still reading it. Right-sized file, wrong bytes, exit 0, on +every rank except LOCAL_RANK 0. + +None of that is visible in the staged bytes here, because on a CPU every +"copy" has already landed by the time anything looks -- which is exactly why +test_pipeline.py cannot catch it. What is visible is the device each event was +recorded against, so that is what this checks. + +The drainer's copy-back must go to a stream of its own. On the default stream +it queues behind the next round's collective instead of running beside it on +the copy engine, which makes the second device window RECV_DEV_SLOTS buys pure +waste. That much is only a throughput bug -- but a side stream that does not +first wait on the collective's event is a correctness one, because the copy +would read a window the all-gather has not filled. Both are checkable here: +concurrency is not observable on a CPU, but which stream work was issued on, +and in what order, is. + +This is not a test of CUDA semantics. It is a test that the threads which +touch CUDA agree about which GPU they are on and which queue they are feeding. +""" + +import os +import shutil +import sys +import tempfile + +import stubs + +MiB = 1024 ** 2 + +# Anything but 0. Device 0 is what a thread that never called set_device gets, +# so pinning this rank to device 0 would pass no matter what. +# +# Stager takes its device from torch.cuda.current_device(), not from +# topo.local_rank, so the two are free to differ here: local_rank stays 0 so +# that this single process is the one that creates the destination file, while +# the device it is actually on is 3. +DEVICE = 3 + +# Set by _run_pipeline when it starts the drain thread; this is how the ops +# below are attributed to the drainer rather than to the main loop. +DRAIN_THREAD = "datastage-drainer" + +torch, dist = stubs.install() +modules = stubs.load_fileio() +ds = modules["datastage"] + +# As in test_pipeline: O_DIRECT would demand alignment the fake buffers cannot +# promise, and nothing here depends on it. +ds.open_maybe_direct = lambda path, flags: os.open(path, flags) + + +class Topology: + """One node, one rank, sitting on GPU DEVICE.""" + + node_count = 1 + local_world_size = 1 + local_rank = 0 # this process creates the destination file + rank = 1 # not 0: keeps the stager's banner out of the way + node_index = 0 + group = None + + +class Args: + num_threads = 4 + chmod = 0o777 + buffer_size = 2 * MiB + + def __init__(self, destination): + self.destination = destination + + +def main(): + root = tempfile.mkdtemp(prefix="datastage-device-") + try: + src = os.path.join(root, "src") + # Several rounds, with a ragged last one. + size = 5 * MiB + 12345 + with open(src, "wb") as handle: + handle.write(os.urandom(size)) + + # What main() does, on the thread main() runs on -- and only there. + torch.cuda.set_device(torch.device("cuda", DEVICE)) + stager = ds.Stager(Args(root), Topology()) + + del stubs.EVENTS[:] + del stubs.OPS[:] + stager.stage_file(src, os.path.join(root, "dst"), size, + os.stat(src).st_mtime_ns) + events = list(stubs.EVENTS) + ops = list(stubs.OPS) + finally: + shutil.rmtree(root, ignore_errors=True) + + if not events: + print(" FAIL no events recorded at all; this test is exercising nothing") + return 1 + + problems = [] + + stray = [event for event in events if event.device != DEVICE] + if stray: + wrong = sorted({event.device for event in stray}) + problems.append( + f"{len(stray)}/{len(events)} events recorded against device {wrong} " + f"instead of {DEVICE}: a thread that records events never called " + f"torch.cuda.set_device, so synchronizing on them waits for nothing") + + # The drain thread is named where it is started, in _run_pipeline. + drained = [op for op in ops if op.thread == DRAIN_THREAD] + if not drained: + problems.append( + f"no operations issued from a thread named {DRAIN_THREAD!r}; either " + f"the drainer was renamed or the copy-back moved somewhere else, and " + f"this test is no longer looking at it") + else: + on_default = [op for op in drained if op.kind != "wait" and op.stream.default] + if on_default: + problems.append( + f"{len(on_default)}/{len(drained)} drainer operations were issued " + f"on the default stream, where they queue behind the next round's " + f"collective instead of overlapping it") + + # Only meaningful once the copies are on a stream of their own: work + # queued on the default stream is already ordered behind the collective + # by the stream itself, and a host-side wait orders it too. A side + # stream has neither, so it has to be told explicitly. + side = [op for op in drained if not op.stream.default] + if side: + copies = [i for i, op in enumerate(side) if op.kind == "copy"] + waits = [i for i, op in enumerate(side) if op.kind == "wait"] + if copies and (not waits or waits[0] > copies[0]): + problems.append( + "the drainer copied from a device window on a side stream " + "without first waiting on the event that says the all-gather " + "filled it; nothing orders the copy behind the collective") + + print(f"test_device: {len(events)} events, {len(drained)} drainer operations " + f"on {len({op.stream for op in drained})} stream(s), {len(problems)} problems") + for problem in problems: + print(f" FAIL {problem}") + return 1 if problems else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_dist_env.py b/tests/test_dist_env.py new file mode 100644 index 0000000..fe72b7f --- /dev/null +++ b/tests/test_dist_env.py @@ -0,0 +1,341 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The slurm/mpirun -> torch env:// translation datastage does for itself. + +Two things are worth pinning down here. The first is that this agrees with +`client/slurm2pytorch`, which other benchmarks still use: a program launched +under the wrapper must take the wrapper's values as given, not recompute them +into something different. + +The second is the failure modes. Everything this module derives feeds a +rendezvous, and a rendezvous that is wrong does not return an error -- it +hangs until the job's wall clock runs out, which is the worst way to find out +that --ntasks-per-node was missing. So the cases that must raise get as much +attention as the cases that must work. + +This is the one test file that does not use stubs.py. dist_env is stdlib +only, so it can be imported directly, and detect() takes the environment as an +argument -- meaning almost everything here is a plain dict rather than a dance +around os.environ. +""" + +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from mlperf_common import dist_env # noqa: E402 + +results = [] + + +def check(name, condition, detail=""): + results.append((name, condition, detail)) + + +def check_raises(name, environ): + try: + env = dist_env.detect(environ) + check(name, False, f"no exception; got {env}") + return None + except dist_env.DistEnvError as exc: + check(name, True) + return str(exc) + + +def check_hostnames(): + """Slurm's compressed hostlist, of which we only ever want the first name.""" + cases = [ + ("dgx001", "dgx001"), + ("dgx[001-004,007]", "dgx001"), # comma inside brackets + ("dgx[007,001-004]", "dgx007"), # literal first, not sorted + ("dgx[001]", "dgx001"), + ("a,b", "a"), + ("node[01-04],other[7-9]", "node01"), + ("nid[1-2]-ib", "nid1-ib"), # suffix after the bracket + ("rack[1-2]node[3-4]", "rack1node3"), # two groups, one name + (" dgx001 ", "dgx001"), + ("dgx001,", "dgx001"), + ("", None), + (" ", None), + (None, None), + ("dgx[001-004", None), # unbalanced + ("dgx[]", None), + ("dgx[-4]", None), + ("dgx[a-b]", None), + ("dgx001]", None), # stray close + ] + for nodelist, expected in cases: + got = dist_env.first_hostname(nodelist) + check(f"first_hostname({nodelist!r}) == {expected!r}", got == expected, + f"got {got!r}") + + # Zero padding has to survive: dgx1 does not resolve. + check("zero padding is preserved", + dist_env.first_hostname("dgx[0007-0009]") == "dgx0007") + + +def check_srun(): + """A plain srun launch, no wrapper -- the case this module exists for.""" + environ = { + "SLURM_PROCID": "9", + "SLURM_NTASKS": "16", + "SLURM_LOCALID": "1", + "SLURM_NTASKS_PER_NODE": "8", + "SLURM_NNODES": "2", + "SLURM_JOB_NODELIST": "dgx[001-002]", + } + env = dist_env.detect(environ) + check("srun: rank", env.rank == 9, f"got {env.rank}") + check("srun: world size", env.world_size == 16) + check("srun: local rank", env.local_rank == 1) + check("srun: local world size", env.local_world_size == 8) + check("srun: master addr from the nodelist", env.master_addr == "dgx001", + f"got {env.master_addr}") + check("srun: default port", env.master_port == 29500) + check("srun: source", env.source == "slurm", f"got {env.source}") + check("srun: describe names the rendezvous", + "dgx001:29500" in dist_env.describe(env), + dist_env.describe(env)) + + +def check_precedence(): + base = { + "SLURM_PROCID": "0", "SLURM_NTASKS": "16", "SLURM_LOCALID": "0", + "SLURM_NTASKS_PER_NODE": "8", "SLURM_JOB_NODELIST": "job[001-002]", + } + check("job nodelist is used when nothing better exists", + dist_env.detect(base).master_addr == "job001") + + with_step = dict(base, SLURM_STEP_NODELIST="step[005-006]") + check("step nodelist beats job nodelist", + dist_env.detect(with_step).master_addr == "step005") + + with_mlperf = dict(with_step, MLPERF_SLURM_FIRSTNODE="fromhost") + check("MLPERF_SLURM_FIRSTNODE beats both nodelists", + dist_env.detect(with_mlperf).master_addr == "fromhost") + + with_master = dict(with_mlperf, MASTER_ADDR="explicit") + check("MASTER_ADDR beats everything", + dist_env.detect(with_master).master_addr == "explicit") + + +def check_mpirun(): + environ = { + "OMPI_COMM_WORLD_RANK": "3", + "OMPI_COMM_WORLD_SIZE": "8", + "OMPI_COMM_WORLD_LOCAL_RANK": "3", + "OMPI_COMM_WORLD_LOCAL_SIZE": "8", + } + env = dist_env.detect(environ) + check("mpirun: source", env.source == "ompi", f"got {env.source}") + check("mpirun: rank", env.rank == 3) + check("mpirun: single node gets localhost", env.master_addr == "127.0.0.1") + + +def check_under_slurm2pytorch(): + """The wrapper still works: its values are taken as given, not recomputed. + + The slurm variables here deliberately disagree with the preset ones. If + anything recomputed rather than deferred, this is where it would show. + """ + environ = { + "RANK": "5", "WORLD_SIZE": "16", "LOCAL_RANK": "5", + "LOCAL_WORLD_SIZE": "8", "MASTER_ADDR": "wrapper-host", + "MASTER_PORT": "29500", + "SLURM_PROCID": "11", "SLURM_NTASKS": "32", "SLURM_LOCALID": "3", + "SLURM_NTASKS_PER_NODE": "4", "SLURM_JOB_NODELIST": "other[001-008]", + } + env = dist_env.detect(environ) + check("wrapper: source is preset", env.source == "preset", f"got {env.source}") + check("wrapper: rank taken as given", env.rank == 5, f"got {env.rank}") + check("wrapper: world size taken as given", env.world_size == 16) + check("wrapper: local rank taken as given", env.local_rank == 5) + check("wrapper: local world size taken as given", env.local_world_size == 8) + check("wrapper: master addr taken as given", + env.master_addr == "wrapper-host") + + +def check_empty_and_single(): + check("an exported-but-empty variable counts as unset", + dist_env.detect({"RANK": "", "SLURM_PROCID": "3", + "SLURM_NTASKS": "8", "SLURM_LOCALID": "3", + "SLURM_NTASKS_PER_NODE": "8"}).rank == 3) + + env = dist_env.detect({}) + check("no launcher at all: source is single", env.source == "single", + f"got {env.source}") + check("no launcher at all: rank 0 of 1", (env.rank, env.world_size) == (0, 1)) + check("no launcher at all: localhost", env.master_addr == "127.0.0.1") + + env = dist_env.detect({"SLURM_PROCID": "2", "SLURM_NTASKS": "8", + "SLURM_LOCALID": "2", "SLURM_NTASKS_PER_NODE": "8"}) + check("single node with no nodelist falls back to localhost", + env.master_addr == "127.0.0.1") + + +def check_refusals(): + """The launches we cannot serve, which must not turn into a hung rendezvous.""" + message = check_raises( + "a multi-node job with no derivable master is refused", + {"SLURM_PROCID": "0", "SLURM_NTASKS": "16", "SLURM_LOCALID": "0", + "SLURM_NTASKS_PER_NODE": "8", "SLURM_JOB_NODELIST": "dgx[001"}) + if message: + check("that message says how to fix it", + "MLPERF_SLURM_FIRSTNODE" in message, message) + + message = check_raises( + "srun -N2 -n16 with no --ntasks-per-node is refused", + {"SLURM_PROCID": "0", "SLURM_NTASKS": "16", "SLURM_LOCALID": "0", + "SLURM_JOB_NODELIST": "dgx[001-002]"}) + if message: + check("that message names --ntasks-per-node", + "--ntasks-per-node" in message, message) + + # Guarded, because the interesting way for this to break is for detect() + # to start raising -- and an exception thrown inside a check() argument + # takes the whole file down with a traceback instead of naming itself. + try: + supplied = dist_env.detect( + {"SLURM_PROCID": "0", "SLURM_NTASKS": "16", "SLURM_LOCALID": "0", + "SLURM_TASKS_PER_NODE": "8(x2)", + "SLURM_JOB_NODELIST": "dgx[001-002]"}).local_world_size + check("SLURM_TASKS_PER_NODE supplies the missing ranks-per-node", + supplied == 8, f"got {supplied}") + except dist_env.DistEnvError as exc: + check("SLURM_TASKS_PER_NODE supplies the missing ranks-per-node", + False, f"raised {exc}") + + check_raises( + "a ragged SLURM_TASKS_PER_NODE is refused", + {"SLURM_PROCID": "0", "SLURM_NTASKS": "12", "SLURM_LOCALID": "0", + "SLURM_TASKS_PER_NODE": "4,8", "SLURM_JOB_NODELIST": "dgx[001-002]"}) + + check_raises( + "slurm's own numbers not multiplying out is refused", + {"SLURM_PROCID": "0", "SLURM_NTASKS": "15", "SLURM_LOCALID": "0", + "SLURM_NNODES": "2", "SLURM_NTASKS_PER_NODE": "8", + "SLURM_JOB_NODELIST": "dgx[001-002]"}) + + check_raises( + "slurm and mpirun disagreeing is refused", + {"SLURM_PROCID": "0", "SLURM_NTASKS": "16", "SLURM_LOCALID": "0", + "SLURM_NTASKS_PER_NODE": "8", "SLURM_JOB_NODELIST": "dgx[001-002]", + "OMPI_COMM_WORLD_RANK": "4", "OMPI_COMM_WORLD_SIZE": "16"}) + + check_raises("a non-integer count is refused", + {"SLURM_PROCID": "0", "SLURM_NTASKS": "x"}) + + check_raises("a rank outside the world is refused", + {"RANK": "9", "WORLD_SIZE": "8", "LOCAL_RANK": "1", + "LOCAL_WORLD_SIZE": "8"}) + + check_raises("a local rank outside its node is refused", + {"RANK": "1", "WORLD_SIZE": "8", "LOCAL_RANK": "9", + "LOCAL_WORLD_SIZE": "8"}) + + check_raises("a world that is not a multiple of the node size is refused", + {"RANK": "1", "WORLD_SIZE": "15", "LOCAL_RANK": "1", + "LOCAL_WORLD_SIZE": "8", "MASTER_ADDR": "h"}) + + check_raises("an impossible port is refused", + {"RANK": "0", "WORLD_SIZE": "1", "LOCAL_RANK": "0", + "LOCAL_WORLD_SIZE": "1", "MASTER_PORT": "99999"}) + + +def check_configure(): + environ = { + "SLURM_PROCID": "9", "SLURM_NTASKS": "16", "SLURM_LOCALID": "1", + "SLURM_NTASKS_PER_NODE": "8", "SLURM_JOB_NODELIST": "dgx[001-002]", + } + dist_env.configure(environ) + expected = {"RANK": "9", "WORLD_SIZE": "16", "LOCAL_RANK": "1", + "LOCAL_WORLD_SIZE": "8", "MASTER_ADDR": "dgx001", + "MASTER_PORT": "29500"} + for key, value in expected.items(): + check(f"configure sets {key}", environ.get(key) == value, + f"got {environ.get(key)!r}") + check("configure sets OMP_NUM_THREADS on a multi-rank node", + environ.get("OMP_NUM_THREADS") == "1") + + kept = dict(environ, OMP_NUM_THREADS="4") + dist_env.configure(kept) + check("configure leaves an existing OMP_NUM_THREADS alone", + kept["OMP_NUM_THREADS"] == "4") + + single = {"RANK": "0", "WORLD_SIZE": "1", "LOCAL_RANK": "0", + "LOCAL_WORLD_SIZE": "1"} + dist_env.configure(single) + check("configure leaves OMP_NUM_THREADS unset for one rank per node", + "OMP_NUM_THREADS" not in single) + + +def check_topology_contract(): + """What configure() writes is what Topology reads back. + + Topology does int(os.environ[name]) for these four and validates the block + layout itself; this only checks the handover, not the layout. + """ + keys = ("RANK", "WORLD_SIZE", "LOCAL_RANK", "LOCAL_WORLD_SIZE", + "MASTER_ADDR", "MASTER_PORT", "OMP_NUM_THREADS", + "SLURM_PROCID", "SLURM_NTASKS", "SLURM_LOCALID", + "SLURM_NTASKS_PER_NODE", "SLURM_JOB_NODELIST") + saved = {key: os.environ.get(key) for key in keys} + try: + for key in keys: + os.environ.pop(key, None) + os.environ.update({ + "SLURM_PROCID": "9", "SLURM_NTASKS": "16", "SLURM_LOCALID": "1", + "SLURM_NTASKS_PER_NODE": "8", "SLURM_JOB_NODELIST": "dgx[001-002]", + }) + env = dist_env.configure() + readback = (int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]), + int(os.environ["LOCAL_RANK"]), + int(os.environ["LOCAL_WORLD_SIZE"])) + check("os.environ round-trip agrees with the returned values", + readback == (env.rank, env.world_size, env.local_rank, + env.local_world_size), + f"got {readback}") + check("the round-trip satisfies Topology's block-layout check", + readback[0] % readback[3] == readback[2]) + finally: + for key, value in saved.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def main(): + check_hostnames() + check_srun() + check_precedence() + check_mpirun() + check_under_slurm2pytorch() + check_empty_and_single() + check_refusals() + check_configure() + check_topology_contract() + + failures = [r for r in results if not r[1]] + for name, _, detail in failures: + print(f" FAIL {name}{': ' + detail if detail else ''}") + print(f"test_dist_env: {len(results)} checks, {len(failures)} failures") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_layout.py b/tests/test_layout.py new file mode 100644 index 0000000..6dc7378 --- /dev/null +++ b/tests/test_layout.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FileLayout must tile every file exactly once, with aligned boundaries. + +Two invariants, both of which corrupt data silently if broken: + + * every byte of the file is claimed by exactly one (local_rank, node, round), + with no gap and no overlap -- a gap leaves stale bytes, an overlap means two + ranks writing the same range; + * every segment boundary is alignment-aligned except the file's own end. + direct_io.pwrite pads writes up to the block size, so an unaligned interior + boundary would let one segment's padding overwrite the next segment's bytes. +""" + +import sys + +import stubs + +stubs.install() +ds = stubs.load_fileio()["datastage"] + +ALIGN = ds.BUFFER_ALIGN +MiB = 1024 ** 2 +GiB = 1024 ** 3 + +SIZES = [0, 1, 4096, ALIGN - 1, ALIGN, ALIGN + 1, 3 * ALIGN, + 100 * MiB, GiB, 7 * GiB + 12345, 137 * GiB + 999] +NODE_COUNTS = [1, 2, 3, 8, 64, 2048] +RANKS_PER_NODE = [1, 4, 8] +BUFFER_SIZES = [2 * MiB, 8 * MiB, 64 * MiB] + + +class Topology: + def __init__(self, node_count, local_world_size, local_rank): + self.node_count = node_count + self.local_world_size = local_world_size + self.local_rank = local_rank + + +def check(size, nodes, ranks_per_node, buffer_size): + """Return an error string, or None if this configuration tiles correctly.""" + covered = [] + for local_rank in range(ranks_per_node): + topology = Topology(nodes, ranks_per_node, local_rank) + layout = ds.FileLayout(size, topology, buffer_size, ALIGN) + for round_index in range(layout.rounds): + for node in range(nodes): + offset, length = layout.segment(node, round_index) + if length: + covered.append((offset, length)) + + covered.sort() + cursor = 0 + for offset, length in covered: + if offset != cursor: + kind = "gap" if offset > cursor else "overlap" + return f"{kind} at {cursor} (next segment starts at {offset})" + cursor += length + if cursor != size: + return f"covered {cursor} of {size} bytes" + + for offset, length in covered: + end = offset + length + if offset % ALIGN: + return f"segment starts unaligned at {offset}" + if end % ALIGN and end != size: + return f"interior segment ends unaligned at {end}" + return None + + +def main(): + checked = 0 + failures = [] + for size in SIZES: + for nodes in NODE_COUNTS: + for ranks_per_node in RANKS_PER_NODE: + for buffer_size in BUFFER_SIZES: + checked += 1 + problem = check(size, nodes, ranks_per_node, buffer_size) + if problem: + failures.append( + f"size={size} nodes={nodes} ranks/node={ranks_per_node} " + f"-b={buffer_size >> 20}M: {problem}") + for failure in failures: + print(f" FAIL {failure}") + print(f"test_layout: {checked} configurations, {len(failures)} failures") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..6519f6e --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run stage_file end to end on the CPU and compare the bytes. + +The reader thread, the collective, the draining thread and the writer pool hand +buffers between each other through four queues. Getting that wrong does not +raise -- it produces a file of exactly the right length holding the wrong bytes, +which is why this compares content rather than just checking for exceptions. +Two such bugs were found this way: a drain queue emptied on shutdown before it +had been processed, and a missing final ftruncate that left O_DIRECT write +padding on the end of every unaligned file. + +The multi-node case fakes the all-gather by filling each node's segment with +the bytes that node would have read, which is what the real collective delivers. +""" + +import os +import random +import shutil +import sys +import tempfile + +import stubs + +MiB = 1024 ** 2 +torch, dist = stubs.install() +modules = stubs.load_fileio() +ds = modules["datastage"] + +# The tests run on a normal filesystem, where O_DIRECT would demand alignment +# the fake buffers cannot promise; the padding behaviour under test is +# direct_io's, and that is identical on a buffered descriptor. +ds.open_maybe_direct = lambda path, flags: os.open(path, flags) + +# Sizes chosen around the alignment boundary, since that is where the write +# padding and the final truncate interact. +SIZES = [0, 1, 1000, 2 * MiB - 1, 2 * MiB, 2 * MiB + 1, 5 * MiB + 12345, 20 * MiB] +MULTINODE = [ + # (nodes, buffer_size, file size) + (1, 2 * MiB, 5_000_000), + (3, 2 * MiB, 5_000_000), + (3, 2 * MiB, 100_000_000), + (8, 2 * MiB, 100_000_000), + (8, 8 * MiB, 100_000_000), + (64, 2 * MiB, 100_000_000), # several host chunks per window + (64, 8 * MiB, 100_000_000), + (64, 2 * MiB, 999), +] + +STATE = {} +_FileLayout = ds.FileLayout + + +class RecordingLayout(_FileLayout): + """Expose the layout the stager built, so the fake collective can use it.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + STATE["layout"] = self + STATE["round"] = 0 + + +ds.FileLayout = RecordingLayout + + +def fake_all_gather(out, inp, group=None): + """Deliver what a real all-gather would: every node's segment, filled in.""" + layout, round_index = STATE["layout"], STATE["round"] + STATE["round"] += 1 + view = out.numpy() + view[:] = b"\0" * len(view) + for node in range(STATE["nodes"]): + offset, length = layout.segment(node, round_index) + if not length: + continue + chunk = os.pread(STATE["fd"], length, offset) + base = node * layout.piece + view[base:base + len(chunk)] = chunk + + +class Topology: + def __init__(self, nodes): + self.node_count = nodes + self.local_world_size = 1 + self.local_rank = 0 + self.rank = 1 # not 0: keeps the stager's banner out of the way + self.node_index = 0 + self.group = None + + +class Args: + num_threads = 4 + chmod = 0o777 + + def __init__(self, destination, buffer_size): + self.destination = destination + self.buffer_size = buffer_size + + +def stage_once(root, nodes, buffer_size, size, tag): + """Stage one generated file and report whether it arrived intact.""" + src = os.path.join(root, f"src_{tag}") + dst = os.path.join(root, f"dst_{tag}") + payload = random.Random(size + nodes).randbytes(size) + with open(src, "wb") as handle: + handle.write(payload) + + STATE["nodes"] = nodes + STATE["fd"] = os.open(src, os.O_RDONLY) + try: + stager = ds.Stager(Args(root, buffer_size), Topology(nodes)) + stager.stage_file(src, dst, size, os.stat(src).st_mtime_ns) + finally: + os.close(STATE["fd"]) + + with open(dst, "rb") as handle: + got = handle.read() + leftovers = [n for n in os.listdir(root) if ".datastage.tmp." in n] + problems = [] + if len(got) != size: + problems.append(f"length {len(got)} != {size}") + elif got != payload: + first = next(i for i in range(size) if got[i] != payload[i]) + problems.append(f"bytes differ from offset {first}") + if os.stat(dst).st_mode & 0o777 != 0o777: + problems.append(f"mode {oct(os.stat(dst).st_mode & 0o777)}") + if leftovers: + problems.append(f"left {len(leftovers)} temp files") + return problems + + +def stage_failure_leaves_nothing(root): + """A failed copy must not leave its preallocated temp file behind. + + The temp file is ftruncate'd to the full source size before any data is + written, and only the success path's rename ever removed it -- so a failure + used to leave a near-full copy on node-local scratch, one per attempt, + until a resubmit loop against a flaky fabric filled the NVMe. + """ + src = os.path.join(root, "src_fail") + dst = os.path.join(root, "dst_fail") + size = 5 * MiB + with open(src, "wb") as handle: + handle.write(b"z" * size) + + def temps(): + return [n for n in os.listdir(root) if ".datastage.tmp." in n] + + # Recorded from inside the failure, so a test that stopped reaching the + # temp file at all would show up as a failure rather than a pass. + existed = [] + + def boom(*args, **kwargs): + existed.extend(temps()) + raise RuntimeError("injected failure") + + STATE["nodes"] = 1 + STATE["fd"] = os.open(src, os.O_RDONLY) + real_pipeline = ds.Stager._run_pipeline + ds.Stager._run_pipeline = boom + problems = [] + try: + stager = ds.Stager(Args(root, 2 * MiB), Topology(1)) + try: + stager.stage_file(src, dst, size, os.stat(src).st_mtime_ns) + problems.append("stage_file did not raise") + except RuntimeError: + pass + finally: + ds.Stager._run_pipeline = real_pipeline + os.close(STATE["fd"]) + + if not existed: + problems.append("the temp file was never created; test proves nothing") + if temps(): + problems.append(f"left {temps()}") + if os.path.exists(dst): + problems.append("published a destination despite failing") + return problems + + +def main(): + root = tempfile.mkdtemp(prefix="datastage-") + failures = 0 + checked = 0 + try: + dist.all_gather_into_tensor = lambda out, inp, group=None: out.copy_(inp) + STATE["nodes"] = 1 + for size in SIZES: + checked += 1 + STATE["fd"] = -1 + problems = stage_once(root, 1, 2 * MiB, size, f"s{size}") + if problems: + failures += 1 + print(f" FAIL single-node size={size}: {'; '.join(problems)}") + + dist.all_gather_into_tensor = fake_all_gather + for nodes, buffer_size, size in MULTINODE: + checked += 1 + problems = stage_once(root, nodes, buffer_size, + size, f"m{nodes}_{buffer_size}_{size}") + if problems: + failures += 1 + print(f" FAIL nodes={nodes} -b={buffer_size >> 20}M size={size}: " + f"{'; '.join(problems)}") + + dist.all_gather_into_tensor = lambda out, inp, group=None: out.copy_(inp) + checked += 1 + problems = stage_failure_leaves_nothing(root) + if problems: + failures += 1 + print(f" FAIL cleanup after a failed copy: {'; '.join(problems)}") + finally: + shutil.rmtree(root, ignore_errors=True) + + print(f"test_pipeline: {checked} staged files, {failures} failures") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_stager.py b/tests/test_stager.py new file mode 100644 index 0000000..da3aa76 --- /dev/null +++ b/tests/test_stager.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Two Stager behaviours that are not about bytes. + +`_chmod_parents` walks up from a staged file towards the destination root, and +its terminating condition has to cope with a destination root of "/", where +os.path.dirname stops making progress. A runaway there does not crash: it +chmods the container root forever while every peer rank waits in the barrier +on the next line, so the job burns its whole wall-clock allocation. + +The device-memory budget check refuses a --buffer-size that will not fit and +suggests a smaller one. If the suggestion is itself over budget the operator +follows the advice, gets the identical error, and burns another multi-node +allocation -- so what matters is not that a number is printed but that the +number works. +""" + +import os +import re +import sys +import types + +import stubs + +GiB = 1024 ** 3 +MiB = 1024 ** 2 + +torch, dist = stubs.install() +ds = stubs.load_fileio()["datastage"] + +results = [] + + +def check(name, condition, detail=""): + results.append((name, condition, detail)) + + +class Topology: + def __init__(self, nodes): + self.node_count = nodes + self.local_world_size = 1 + self.local_rank = 0 + self.rank = 1 # not 0: keeps the stager's banner quiet + self.node_index = 0 + self.group = None + + +class Args: + num_threads = 4 + chmod = 0o777 + + def __init__(self, destination, buffer_size): + self.destination = destination + self.buffer_size = buffer_size + + +def chmod_parents_calls(dest_root, dst, limit=64): + """Run _chmod_parents against a stub, reporting the paths it chmod'd. + + Raises RuntimeError if it exceeds `limit` calls, which is how a + non-terminating walk shows up as a failure instead of hanging the suite. + """ + calls = [] + stager = types.SimpleNamespace(dest_root=dest_root, + args=types.SimpleNamespace(chmod=0o777)) + + def counting_chmod(path, mode): + calls.append(path) + if len(calls) > limit: + raise RuntimeError(f"_chmod_parents did not terminate: {limit}+ calls") + + real_chmod = os.chmod + os.chmod = counting_chmod + try: + ds.Stager._chmod_parents(stager, dst) + finally: + os.chmod = real_chmod + return calls + + +def check_chmod_parents(): + try: + calls = chmod_parents_calls("/", "/a.bin") + check("a destination root of '/' terminates", True) + check("'/' is chmod'd exactly once", calls == ["/"], f"got {calls}") + except RuntimeError as exc: + check("a destination root of '/' terminates", False, str(exc)) + check("'/' is chmod'd exactly once", False, "did not terminate") + + calls = chmod_parents_calls("/raid/scratch/ds", "/raid/scratch/ds/a/b/c.bin") + check("it walks up to the destination root and stops", + calls == ["/raid/scratch/ds/a/b", "/raid/scratch/ds/a", "/raid/scratch/ds"], + f"got {calls}") + + calls = chmod_parents_calls("/raid/scratch/ds", "/raid/scratch/ds/c.bin") + check("a file directly in the root chmods only the root", + calls == ["/raid/scratch/ds"], f"got {calls}") + + # dest_root is a prefix of the path but not a parent directory of it; the + # startswith test must not be fooled by the shared prefix. + calls = chmod_parents_calls("/raid/scratch/ds", "/raid/scratch/ds-other/c.bin") + check("a sibling directory sharing the prefix is not touched", + calls == [], f"got {calls}") + + +def budget_error(total_memory, nodes, buffer_size): + """Build a Stager and return the RuntimeError text, or None if it fit. + + The device size is patched into the stub in place rather than by calling + stubs.install() again: datastage bound `torch` at import, so a fresh module + from install() would not be the one it consults. + + Allocation is made free for the same duration. The budget check runs + before any buffers are reserved, and has to keep doing so -- but an + *accepted* 64-node case would otherwise really allocate tens of GiB, since + the fake tensors are backed by actual host memory. + """ + saved = (ds.torch.cuda.get_device_properties, ds.torch.empty, ds.pinned_aligned) + ds.torch.cuda.get_device_properties = lambda device: types.SimpleNamespace( + total_memory=int(total_memory)) + ds.torch.empty = lambda n, **kwargs: stubs.FakeTensor(0) + ds.pinned_aligned = lambda nbytes, alignment: ( + None, stubs.FakeTensor(0), memoryview(bytearray(1))) + try: + ds.Stager(Args("/tmp", buffer_size), Topology(nodes)) + return None + except RuntimeError as exc: + return str(exc) + finally: + (ds.torch.cuda.get_device_properties, ds.torch.empty, + ds.pinned_aligned) = saved + + +def check_budget_advice(): + """The suggested --buffer-size has to be one that actually works.""" + # Device sizes against node counts that overflow them. The requirement is + # piece * (RECV_DEV_SLOTS * nodes + 1) against 60% of the device, so e.g. + # 512M at 64 nodes wants 64.5 GiB of an 80 GB card's 47.8 GiB budget. + cases = [ + (79.65 * GiB, 64, 512 * MiB), + (79.65 * GiB, 32, 1024 * MiB), + (79.65 * GiB, 16, 2048 * MiB), + (94.0 * GiB, 64, 512 * MiB), + (141.0 * GiB, 64, 1024 * MiB), + (141.0 * GiB, 128, 1024 * MiB), + (79.65 * GiB, 2048, 32 * MiB), + (8.0 * GiB, 2048, 8 * MiB), # cannot fit even at the 2 MiB minimum + ] + for total_memory, nodes, buffer_size in cases: + label = f"{total_memory / GiB:.0f}GiB/{nodes}n/{buffer_size >> 20}M" + message = budget_error(total_memory, nodes, buffer_size) + if message is None: + check(f"{label}: expected a refusal", False, "it was accepted") + continue + + match = re.search(r"at most (\d+)M", message) + if not match: + # No suggestion means it claimed the node count cannot fit at all. + check(f"{label}: says the node count does not fit", + "does not fit on this GPU" in message, message) + minimum = budget_error(total_memory, nodes, 2 * MiB) + check(f"{label}: and the minimum really is refused", + minimum is not None, "the 2M minimum was accepted") + continue + + suggested = int(match.group(1)) * MiB + again = budget_error(total_memory, nodes, suggested) + check(f"{label}: the suggested {suggested >> 20}M is accepted", + again is None, f"still refused: {again}") + + # And it should be close to the largest that fits, not needlessly small. + bigger = budget_error(total_memory, nodes, suggested + 2 * MiB) + check(f"{label}: {(suggested >> 20) + 2}M would not have fit", + bigger is not None, "the suggestion was more than 2 MiB too low") + + +def main(): + check_chmod_parents() + check_budget_advice() + + failures = [r for r in results if not r[1]] + for name, _, detail in failures: + print(f" FAIL {name}{': ' + detail if detail else ''}") + print(f"test_stager: {len(results)} checks, {len(failures)} failures") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_topology.py b/tests/test_topology.py new file mode 100644 index 0000000..e0b1416 --- /dev/null +++ b/tests/test_topology.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Topology's node grouping, which nothing could construct before. + +datastage assumes slurm's default block distribution: node i holds ranks +i*L .. i*L+L-1. Getting the grouping wrong does not raise -- it builds groups +whose all-gather assembles the right bytes in the wrong order, so every file +comes out the right length with its slices permuted. + +Two things are checked. That the grouping is right: every rank agrees with +every other about which global ranks make up each group, each group holds one +rank per node, and every rank appears exactly once across the groups. And +that the group lists are ascending -- dist.new_group sorts what it is given +and derives each member's position from that order, while the drainer maps +all-gather output position to node_index, so those two agree only while the +lists are sorted to begin with. +""" + +import os +import sys + +import stubs + +torch, dist = stubs.install() +ds = stubs.load_fileio()["datastage"] + +results = [] + + +def check(name, condition, detail=""): + results.append((name, condition, detail)) + + +def topology_for(rank, nodes, ranks_per_node, **overrides): + """Build Topology as `rank` would see it in a nodes x ranks_per_node job.""" + env = { + "RANK": str(rank), + "WORLD_SIZE": str(nodes * ranks_per_node), + "LOCAL_RANK": str(rank % ranks_per_node), + "LOCAL_WORLD_SIZE": str(ranks_per_node), + } + env.update({k: str(v) for k, v in overrides.items()}) + saved = {k: os.environ.get(k) for k in env} + os.environ.update(env) + try: + return ds.Topology() + finally: + for key, value in saved.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def check_layout(nodes, ranks_per_node): + world = nodes * ranks_per_node + label = f"{nodes}x{ranks_per_node}" + views = [topology_for(rank, nodes, ranks_per_node) for rank in range(world)] + + check(f"{label}: every rank sees the same groups", + len({tuple(tuple(g) for g in view.groups) for view in views}) == 1) + + groups = views[0].groups + check(f"{label}: one group per local rank", len(groups) == ranks_per_node) + check(f"{label}: each group has one rank per node", + all(len(g) == nodes for g in groups)) + check(f"{label}: every rank appears exactly once", + sorted(r for g in groups for r in g) == list(range(world))) + check(f"{label}: group lists are ascending", + all(list(g) == sorted(g) for g in groups), + f"got {groups}") + + for rank, view in enumerate(views): + if view.node_index != rank // ranks_per_node: + check(f"{label}: rank {rank} knows its node", False, + f"node_index {view.node_index}") + return + if view.rank not in view.group: + check(f"{label}: rank {rank} is a member of its own group", False) + return + # The position the all-gather will place this rank's segment at has to + # be its node_index, or the drainer writes segments to other nodes' + # offsets. + if list(view.group).index(rank) != view.node_index: + check(f"{label}: rank {rank} sits at its node_index in its group", + False, f"position {list(view.group).index(rank)}") + return + check(f"{label}: every rank sits at its node_index in its own group", True) + + +def check_rejected(name, **kwargs): + try: + topology_for(**kwargs) + check(name, False, "no exception raised") + except RuntimeError: + check(name, True) + + +def main(): + for nodes, ranks_per_node in ((1, 1), (1, 8), (2, 8), (3, 4), (64, 8), (17, 3)): + check_layout(nodes, ranks_per_node) + + # A cyclic distribution reports a LOCAL_RANK that disagrees with RANK % L. + # Rank 1 of a 2x2 job is local rank 0 on node 1 under cyclic, but block + # arithmetic puts it at local rank 1. + check_rejected("a cyclic distribution is refused", + rank=1, nodes=2, ranks_per_node=2, LOCAL_RANK=0) + check_rejected("a ragged --ntasks-per-node is refused", + rank=0, nodes=2, ranks_per_node=8, WORLD_SIZE=15) + check_rejected("LOCAL_WORLD_SIZE of 0 is refused", + rank=0, nodes=1, ranks_per_node=1, LOCAL_WORLD_SIZE=0) + + # A missing LOCAL_WORLD_SIZE used to default to 1, silently turning an + # 8-GPU node into 8 single-rank nodes. + try: + saved = os.environ.pop("LOCAL_WORLD_SIZE", None) + os.environ.update({"RANK": "3", "WORLD_SIZE": "8", "LOCAL_RANK": "3"}) + ds.Topology() + check("a missing LOCAL_WORLD_SIZE is refused", False, "no exception") + except KeyError: + check("a missing LOCAL_WORLD_SIZE is refused", True) + except RuntimeError: + check("a missing LOCAL_WORLD_SIZE is refused", True) + finally: + if saved is not None: + os.environ["LOCAL_WORLD_SIZE"] = saved + + failures = [r for r in results if not r[1]] + for name, _, detail in failures: + print(f" FAIL {name}{': ' + detail if detail else ''}") + print(f"test_topology: {len(results)} checks, {len(failures)} failures") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main())