From 4c8c2f7156d111eb316aa2198b6f8785544e5302 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Thu, 30 Jul 2026 22:08:32 -0500 Subject: [PATCH 01/21] Add mlperf_common.fileio: O_DIRECT primitives, copy planning, NCCL staging Adds datastage, a collective dataset stager. Every node in a training job needs the same dataset, so copying it onto each node reads the whole thing N times from shared storage and runs at the speed of the slowest reader. Instead each rank reads a disjoint shard and NCCL fans the bytes out, so the data crosses the fabric once per node. The world is split into one process group per LOCAL_RANK, each holding one rank per node and owning a contiguous 1/L slice of the file. Within a group each rank reads a disjoint sub-shard and an all-gather assembles the slice on every node. A node's ranks together write the whole file, so nothing is exchanged or written twice within a node, and the concurrent all-gathers drive every NIC without per-cluster transport tuning. Groups the three related modules under mlperf_common/fileio rather than adding them to the package root: direct_io moved verbatim from client/, so package modules can import it (setup.py installs it as a script, which is not importable) copyplan source-tree walk and src->dst mapping, lifted out of fastcp so fastcp and datastage cannot disagree about what a copy covers datastage the stager and its cp/rsync-shaped CLI Only datastage needs torch, so the single-node client scripts do not pull in a training stack. client/direct_io.py becomes a shim re-exporting the package module, keeping `import direct_io` working for fastcp and fastmd5. Both resolve the package either from an install or from a tree with mlperf_common/ next to client/, and fail with an actionable message otherwise. Co-Authored-By: Claude Opus 5 --- client/direct_io.py | 159 ++------- client/fastcp | 65 +--- mlperf_common/fileio/__init__.py | 23 ++ mlperf_common/fileio/copyplan.py | 78 ++++ mlperf_common/fileio/datastage.py | 571 ++++++++++++++++++++++++++++++ mlperf_common/fileio/direct_io.py | 150 ++++++++ 6 files changed, 866 insertions(+), 180 deletions(-) create mode 100644 mlperf_common/fileio/__init__.py create mode 100644 mlperf_common/fileio/copyplan.py create mode 100644 mlperf_common/fileio/datastage.py create mode 100644 mlperf_common/fileio/direct_io.py 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..b5376da 100755 --- a/client/fastcp +++ b/client/fastcp @@ -21,10 +21,25 @@ 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 plan_copy_operations +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: @@ -165,50 +180,6 @@ def parse_and_validate_args(): 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 +190,7 @@ if __name__ == "__main__": print(f"Sources: {args.sources}") print(f"Destination: {args.destination}") - file_jobs = plan_copy_operations(args) + file_jobs = plan_copy_operations(args.sources, args.destination) if not args.force: for _, dst, _ in file_jobs: 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..4ded020 --- /dev/null +++ b/mlperf_common/fileio/copyplan.py @@ -0,0 +1,78 @@ +#!/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. +""" + +import os + +__all__ = ["list_relative_files", "plan_copy_operations"] + + +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. + + FIXME: os.walk(followlinks=True) doesn't protect against cycles. To fix + this we'd need to write our own version that did a depth-first spanning + tree. + """ + file_list = [] + 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(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 this is a + single file-to-file copy. The result is sorted by destination path so that + every rank of a collective copy walks the files in the same order. + """ + file_jobs = [] + dst_root = os.path.abspath(destination) + + if not os.path.isdir(dst_root): # case 1: single file copy + src_abs = os.path.abspath(sources[0]) + size = os.path.getsize(src_abs) + file_jobs.append((src_abs, dst_root, size)) + else: + for src in 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) + file_jobs.append((full_src, full_dst, os.path.getsize(full_src))) + else: + dst_path = os.path.join(dst_root, base) + file_jobs.append((src_abs, dst_path, os.path.getsize(src_abs))) + + 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..9ada2cc --- /dev/null +++ b/mlperf_common/fileio/datastage.py @@ -0,0 +1,571 @@ +#!/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, under slurm2pytorch so that RANK / WORLD_SIZE / +LOCAL_RANK / MASTER_ADDR are set: + + srun --ntasks-per-node=${DGXNGPU} ... slurm2pytorch \\ + python3 -m mlperf_common.fileio.datastage -r "${SLOW_DATADIR}/${DATASET}" "${DATADIR}" +""" + +import argparse +import os +import queue +import socket +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import torch +import torch.distributed as dist + +from mlperf_common.fileio import direct_io +from mlperf_common.fileio.copyplan import plan_copy_operations + +# 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; +# receive slots let the writers drain behind it. +SEND_SLOTS = 3 +RECV_SLOTS = 2 + + +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.""" + + 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.get("LOCAL_WORLD_SIZE", 1)) + + # Derive the node layout from what the ranks actually report rather + # than assuming a block distribution of RANK: run.sub launches some + # steps with --distribution=arbitrary, and silently mis-grouping would + # produce a corrupt copy rather than an error. + identities = [None] * self.world_size + dist.all_gather_object(identities, (socket.gethostname(), self.local_rank)) + + hosts = [] + for host, _ in identities: + if host not in hosts: + hosts.append(host) + self.node_count = len(hosts) + node_index = {host: i for i, host in enumerate(hosts)} + + # group_ranks[l][node] = global rank of LOCAL_RANK l on that node + group_ranks = [[None] * self.node_count for _ in range(self.local_world_size)] + for global_rank, (host, local_rank) in enumerate(identities): + if local_rank >= self.local_world_size: + raise RuntimeError( + f"rank {global_rank} reports LOCAL_RANK={local_rank} with " + f"LOCAL_WORLD_SIZE={self.local_world_size}" + ) + slot = group_ranks[local_rank][node_index[host]] + if slot is not None: + raise RuntimeError( + f"ranks {slot} and {global_rank} both claim LOCAL_RANK=" + f"{local_rank} on {host}" + ) + group_ranks[local_rank][node_index[host]] = global_rank + for local_rank, ranks in enumerate(group_ranks): + if any(r is None for r in ranks): + raise RuntimeError( + f"not every node has a rank with LOCAL_RANK={local_rank}; " + "launch with a uniform --ntasks-per-node" + ) + + # 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] + self.node_index = node_index[socket.gethostname()] + + 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) + # One piece per rank per round; the assembled window is node_count of + # them. Computed once here and handed to every FileLayout so the two + # can never drift apart. + self.piece = max(align_down(args.window // topology.node_count, self.align), self.align) + window = self.piece * topology.node_count + + self.send_host = [pinned_aligned(self.piece, self.align) for _ in range(SEND_SLOTS)] + self.recv_host = [pinned_aligned(window, self.align) for _ in range(RECV_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) + + 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_segments(self, fd, mview, layout, round_index, block_size): + """Scatter one assembled window to its true file offsets.""" + futures = [] + for node in range(self.topo.node_count): + offset, length = layout.segment(node, round_index) + if length == 0: + continue + base = node * layout.piece + futures.append( + self.pool.submit( + direct_io.pwrite, fd, + mview[base:base + layout.piece], length, offset, block_size, + ) + ) + return futures + + def stage_file(self, src, dst, size, mtime_ns): + topo = self.topo + tmp = f"{dst}.datastage.tmp.{os.environ.get('SLURM_JOB_ID', 'nojob')}" + + # 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: + 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 + path = os.path.dirname(path) + + def _run_pipeline(self, fd_src, fd_dst, layout, src_block, dst_block): + topo = self.topo + free_q = queue.Queue() + filled_q = queue.Queue() + for slot in range(SEND_SLOTS): + 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) + + reader_thread = threading.Thread(target=reader, name="datastage-reader") + reader_thread.start() + + pending = [[] for _ in range(RECV_SLOTS)] + try: + for round_index in range(layout.rounds): + item = filled_q.get() + if item is None: + break + _, slot, length = item + + _, send_view, _ = self.send_host[slot] + self.send_dev.copy_(send_view, non_blocking=True) + torch.cuda.synchronize() + free_q.put(slot) + + dist.all_gather_into_tensor(self.recv_dev, self.send_dev, group=topo.group) + + recv_slot = round_index % RECV_SLOTS + for future in pending[recv_slot]: + future.result() + pending[recv_slot] = [] + + _, recv_view, recv_mview = self.recv_host[recv_slot] + recv_view.copy_(self.recv_dev, non_blocking=True) + torch.cuda.synchronize() + + pending[recv_slot] = self._write_segments( + fd_dst, recv_mview, layout, round_index, dst_block + ) + for futures in pending: + for future in futures: + future.result() + finally: + # 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 is + # unbounded, so the reader can never 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.""" + if dist.get_rank() == 0: + jobs = plan_copy_operations(args.sources, args.destination) + payload = [[(src, dst, size, os.stat(src).st_mtime_ns) for src, dst, size in jobs]] + else: + payload = [None] + dist.broadcast_object_list(payload, src=0) + jobs = payload[0] + + # Guard against ranks seeing a different view of shared storage. + mismatches = 0 + for src, _, size, mtime_ns in jobs: + try: + st = os.stat(src) + except OSError: + mismatches += 1 + continue + if st.st_size != size or st.st_mtime_ns != mtime_ns: + mismatches += 1 + counter = torch.tensor([mismatches], dtype=torch.int64, device="cuda") + dist.all_reduce(counter) + if counter.item(): + raise RuntimeError( + f"{counter.item()} rank/file pairs disagree with rank 0 about the " + "source tree; shared storage is inconsistent or changing" + ) + return jobs + + +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("-w", "--window", type=parse_size, default=parse_size("2G"), + help="assembled all-gather window (default: 2G). Divided by the " + "node count to give the per-rank read size, so raise it at " + "high node counts. Costs this much device memory and twice " + "as much pinned host memory.") + 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}'") + for src in args.sources: + if not os.path.exists(src): + sys.exit(f"{prog}: cannot stat '{src}': No such file or directory") + if os.path.isdir(src) and not args.recursive: + sys.exit(f"{prog}: -r not specified; omitting directory '{src}'") + return args + + +def main(argv=None): + args = parse_args(argv) + + if args.dry_run and "RANK" not in os.environ: + for src, dst, size in plan_copy_operations(args.sources, args.destination): + print(f"{src} -> {dst} ({size} bytes)") + return 0 + + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) + 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}") From 5bcbc19755bc69693475865ed63b970d1ef5abf8 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Thu, 30 Jul 2026 22:29:34 -0500 Subject: [PATCH 02/21] fastmd5: share the tree walk, and fail loudly on worker errors fastmd5 walked the tree itself with Path.rglob and an `is_file() or is_symlink()` predicate. That predicate admitted entries that are not readable files: is_file() already follows symlinks, so the extra clause only added broken symlinks and symlinks to directories. A symlink to a directory then raised IsADirectoryError inside a worker thread, and an exception escaping a thread does not affect the process exit status -- so fastmd5 printed a traceback on stderr, skipped the file, and exited 0. Anything comparing two trees by parsing stdout saw a short but successful-looking result. Use copyplan.list_relative_files instead, so checksumming a staged copy enumerates the same files the copy was planned from. datastage dereferences symlinks, so a staged tree has different link structure from its source; the two walks have to agree for the comparison to mean anything. Also collect exceptions raised in worker threads and exit nonzero if there were any. fastmd5 still checksums everything it can and reports the partial result, it just no longer reports success. Co-Authored-By: Claude Opus 5 --- client/fastmd5 | 94 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 62 insertions(+), 32 deletions(-) diff --git a/client/fastmd5 b/client/fastmd5 index adf7644..d717011 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 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,15 @@ 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): + for relpath in list_relative_files(original_arg): + 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 +124,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() From 4f1c0632e86d4e53b245356e9aabc26b1b86596b Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Thu, 30 Jul 2026 22:33:47 -0500 Subject: [PATCH 03/21] copyplan: report every unreadable source entry instead of dying on the first A dangling symlink is listed by os.walk among the filenames, because it is not a directory, and stat then follows it to nothing. Every caller reached that FileNotFoundError during enumeration, on the main thread, before any work started -- so one dead link aborted the whole run with a raw traceback and zero output, even when the rest of the tree was fine. For datastage that meant rank 0 dying during planning, after the allocation was granted, and only naming whichever bad entry os.walk happened to reach first. Collect them instead and raise UnreadableEntries naming all of them (capped at 20 with a count of the remainder). Behaviour is still to fail rather than to skip: a staging tool that quietly omits files is the worse failure. plan_copy_operations now stats each source once and reuses the result for the size, rather than calling getsize separately. fastcp and fastmd5 turn it into a clean message and a nonzero exit. datastage has to broadcast the failure from rank 0 instead of raising, since every other rank is already parked in the broadcast and would otherwise hang until the NCCL watchdog fired. Co-Authored-By: Claude Opus 5 --- client/fastcp | 7 ++- client/fastmd5 | 8 ++- mlperf_common/fileio/copyplan.py | 84 +++++++++++++++++++++++++------ mlperf_common/fileio/datastage.py | 21 ++++++-- 4 files changed, 97 insertions(+), 23 deletions(-) diff --git a/client/fastcp b/client/fastcp index b5376da..6ada473 100755 --- a/client/fastcp +++ b/client/fastcp @@ -32,7 +32,7 @@ sys.path.insert(1, os.path.join(_HERE, os.pardir)) import direct_io try: - from mlperf_common.fileio.copyplan import plan_copy_operations + from mlperf_common.fileio.copyplan import UnreadableEntries, plan_copy_operations except ImportError as exc: sys.exit( f"fastcp: cannot import mlperf_common ({exc}).\n" @@ -190,7 +190,10 @@ if __name__ == "__main__": print(f"Sources: {args.sources}") print(f"Destination: {args.destination}") - file_jobs = plan_copy_operations(args.sources, args.destination) + 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 d717011..35e17df 100755 --- a/client/fastmd5 +++ b/client/fastmd5 @@ -18,7 +18,7 @@ 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 list_relative_files + from mlperf_common.fileio.copyplan import UnreadableEntries, list_relative_files except ImportError as exc: sys.exit( f"fastmd5: cannot import mlperf_common ({exc}).\n" @@ -107,7 +107,11 @@ def main(): # 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): - for relpath in list_relative_files(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: diff --git a/mlperf_common/fileio/copyplan.py b/mlperf_common/fileio/copyplan.py index 4ded020..7659177 100644 --- a/mlperf_common/fileio/copyplan.py +++ b/mlperf_common/fileio/copyplan.py @@ -24,7 +24,43 @@ import os -__all__ = ["list_relative_files", "plan_copy_operations"] +__all__ = ["UnreadableEntries", "list_relative_files", "plan_copy_operations"] + +# 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) + + +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): @@ -34,16 +70,22 @@ def list_relative_files(root): symlinks to files appear in filenames, so both are dereferenced and copied as content rather than recreated as links. - FIXME: os.walk(followlinks=True) doesn't protect against cycles. To fix - this we'd need to write our own version that did a depth-first spanning - tree. + Raises UnreadableEntries if anything under root cannot be stat'd, reporting + every such entry rather than dying on the first one. + + 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. """ file_list = [] + problems = [] 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) + 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) @@ -54,25 +96,37 @@ def plan_copy_operations(sources, destination): under its own basename (recursing into directories); otherwise this is a single file-to-file copy. The result is sorted by destination path so that every rank of a collective copy walks the files in the same order. + + Raises UnreadableEntries if any source cannot be stat'd, reporting all of + them together. """ 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): # case 1: single file copy - src_abs = os.path.abspath(sources[0]) - size = os.path.getsize(src_abs) - file_jobs.append((src_abs, dst_root, size)) + add(os.path.abspath(sources[0]), dst_root) else: for src in 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) - file_jobs.append((full_src, full_dst, os.path.getsize(full_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: - dst_path = os.path.join(dst_root, base) - file_jobs.append((src_abs, dst_path, os.path.getsize(src_abs))) + 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 index 9ada2cc..ea6b6b4 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -63,7 +63,7 @@ import torch.distributed as dist from mlperf_common.fileio import direct_io -from mlperf_common.fileio.copyplan import plan_copy_operations +from mlperf_common.fileio.copyplan import UnreadableEntries, plan_copy_operations # 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 @@ -440,11 +440,20 @@ def reader(): def build_plan(args): """Rank 0 walks the source tree; everyone else takes its answer verbatim.""" if dist.get_rank() == 0: - jobs = plan_copy_operations(args.sources, args.destination) - payload = [[(src, dst, size, os.stat(src).st_mtime_ns) for src, dst, size in jobs]] + # A planning failure has to be broadcast rather than raised here: every + # other rank is already waiting in the broadcast below and would hang + # until the NCCL watchdog fired. + 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 UnreadableEntries as exc: + payload = [{"error": str(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']}") jobs = payload[0] # Guard against ranks seeing a different view of shared storage. @@ -518,7 +527,11 @@ def main(argv=None): args = parse_args(argv) if args.dry_run and "RANK" not in os.environ: - for src, dst, size in plan_copy_operations(args.sources, args.destination): + 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 From d4edc959246a70ef4011e0f4bee3ab6bec8f06df Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 09:51:54 -0500 Subject: [PATCH 04/21] datastage: overlap the copies with CUDA events, and take --buffer-size Two fixes to the staging pipeline. torch.cuda.synchronize() after each host/device copy is a device-wide barrier, so the host-to-device copy, the all-gather, the device-to-host copy and the NVMe writes ran one after another -- most of what the slot pipeline exists for. Record an event per copy instead. The collective is launched before waiting on the host-to-device event, since it is queued behind the copy on the same stream anyway, and the writers now wait on the device-to-host event themselves so the main thread returns to the next round while that copy is still in flight. --window named the assembled all-gather buffer, and divided it by the node count to get the per-rank read size. That silently inflated: the per-rank piece cannot go below the 2MiB O_DIRECT alignment, so at 2048 nodes a requested 2G window became 4G, and 8G at 4096. Take --buffer-size instead, the bytes each rank reads per round, exactly as fastcp's --buffer-size is what each thread reads. The window is then derived rather than requested, so nothing is silently exceeded, and it is rounded up to a 2MiB multiple the way fastcp does. Memory now visibly scales with the job, because an all-gather delivers the whole window to every participant: at 2048 nodes 8M gives a 16 GiB window. Rank 0 reports the footprint, and an oversized request is rejected up front with the largest workable value rather than failing in a CUDA OOM. Co-Authored-By: Claude Opus 5 --- mlperf_common/fileio/datastage.py | 88 +++++++++++++++++++++++-------- 1 file changed, 65 insertions(+), 23 deletions(-) diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index ea6b6b4..f6c5eb8 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -246,17 +246,35 @@ def __init__(self, args, topology): self.align = BUFFER_ALIGN self.dest_root = os.path.abspath(args.destination) - # One piece per rank per round; the assembled window is node_count of - # them. Computed once here and handed to every FileLayout so the two - # can never drift apart. - self.piece = max(align_down(args.window // topology.node_count, self.align), self.align) + # 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 + # Fail before allocating rather than dying inside a CUDA OOM. + budget = int(torch.cuda.get_device_properties(self.device).total_memory * 0.25) + if window > budget: + raise RuntimeError( + f"--buffer-size {self.piece >> 20}M across {topology.node_count} nodes " + f"needs a {window / 1024 ** 3:.1f} GiB window per rank, over the " + f"{budget / 1024 ** 3:.1f} GiB budget. Lower --buffer-size to at most " + f"{max(align_down(budget // topology.node_count, self.align), self.align) >> 20}M." + ) + self.send_host = [pinned_aligned(self.piece, self.align) for _ in range(SEND_SLOTS)] self.recv_host = [pinned_aligned(window, self.align) for _ in range(RECV_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) + if topology.rank == 0: + pinned = SEND_SLOTS * self.piece + RECV_SLOTS * window + print(f"datastage: {self.piece >> 20} MiB per rank per round, " + f"{window / 1024 ** 3:.2f} GiB window; per rank " + f"{(self.piece + window) / 1024 ** 3:.2f} GiB device, " + f"{pinned / 1024 ** 3:.2f} GiB pinned host", 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: @@ -288,20 +306,25 @@ def _read_piece(self, fd, mview, offset, length, block_size): for future in futures: future.result() - def _write_segments(self, fd, mview, layout, round_index, block_size): - """Scatter one assembled window to its true file offsets.""" + def _write_segments(self, fd, mview, layout, round_index, block_size, ready): + """Scatter one assembled window to its true file offsets. + + `ready` is the event marking the end of the device-to-host copy that + filled `mview`. Each writer waits on it rather than the main thread, so + the main thread can go on and launch the next round's collective while + the copy is still in flight. + """ + def write_one(base, length, offset): + ready.synchronize() + return direct_io.pwrite(fd, mview[base:base + layout.piece], + length, offset, block_size) + futures = [] for node in range(self.topo.node_count): offset, length = layout.segment(node, round_index) if length == 0: continue - base = node * layout.piece - futures.append( - self.pool.submit( - direct_io.pwrite, fd, - mview[base:base + layout.piece], length, offset, block_size, - ) - ) + futures.append(self.pool.submit(write_one, node * layout.piece, length, offset)) return futures def stage_file(self, src, dst, size, mtime_ns): @@ -404,10 +427,16 @@ def reader(): _, send_view, _ = self.send_host[slot] self.send_dev.copy_(send_view, non_blocking=True) - torch.cuda.synchronize() - free_q.put(slot) + 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. dist.all_gather_into_tensor(self.recv_dev, self.send_dev, group=topo.group) + h2d_done.synchronize() + free_q.put(slot) recv_slot = round_index % RECV_SLOTS for future in pending[recv_slot]: @@ -416,10 +445,16 @@ def reader(): _, recv_view, recv_mview = self.recv_host[recv_slot] recv_view.copy_(self.recv_dev, non_blocking=True) - torch.cuda.synchronize() - + d2h_done = torch.cuda.Event(blocking=True) + d2h_done.record() + + # The writers wait on d2h_done themselves, so the main thread + # returns to the top of the loop immediately. recv_dev is safe + # to reuse because the next round's collective is queued behind + # this copy, and recv_host[recv_slot] is held by the pending + # futures checked above. pending[recv_slot] = self._write_segments( - fd_dst, recv_mview, layout, round_index, dst_block + fd_dst, recv_mview, layout, round_index, dst_block, d2h_done ) for futures in pending: for future in futures: @@ -495,11 +530,12 @@ def parse_args(argv=None): 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("-w", "--window", type=parse_size, default=parse_size("2G"), - help="assembled all-gather window (default: 2G). Divided by the " - "node count to give the per-rank read size, so raise it at " - "high node counts. Costs this much device memory and twice " - "as much pinned host memory.") + 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", @@ -515,6 +551,12 @@ def parse_args(argv=None): 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) for src in args.sources: if not os.path.exists(src): sys.exit(f"{prog}: cannot stat '{src}': No such file or directory") From 888c323b1635efd2200251ffabbe45c65f5fabc5 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 10:15:38 -0500 Subject: [PATCH 05/21] datastage: stream the window back to host, and fix two corruption bugs Copying the whole assembled window into pinned host memory before writing any of it made pinned host memory scale with the node count: two windows per rank, times DGXNGPU ranks per node, is 512 GiB per node at 2048 nodes with -b 32M. That was the binding limit, not GPU memory, which sat at 22%. The window is a concatenation of one segment per node, each bound for a different file offset, so it does not have to land in host memory whole. Copy it back a few segments at a time through a small fixed pool instead, and give the drain its own thread so the next round's collective overlaps the copy-back and the writes. Pinned host is now constant; the device holds two windows so the collective and the copy-back can overlap, which is the cheap resource here -- this is its own job step and exits before training. The device budget goes to 60% accordingly. Two bugs found by an end-to-end test of the pipeline with CUDA and NCCL stubbed out, both of which silently corrupted the staged file: _run_pipeline set the stop flag as soon as the main loop finished feeding the drain queue, and the drainer treated that as "abort" -- so every window still queued was dropped. The drain lags the main loop by design, so stop now means abort only, and the drainer is joined before teardown. Timing decided how much of a file survived. O_DIRECT writes are padded up to the block size, so the last write of a file runs past its end and nothing trimmed it back; any file whose size was not a multiple of the block size ended up padded with garbage. ftruncate after the barrier, where fastcp does the same after its copy. It has to be after the barrier because until then another rank may still be padding. Co-Authored-By: Claude Opus 5 --- mlperf_common/fileio/datastage.py | 205 +++++++++++++++++++++--------- 1 file changed, 147 insertions(+), 58 deletions(-) diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index f6c5eb8..9580899 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -74,10 +74,22 @@ # 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; -# receive slots let the writers drain behind it. +# Pipeline depth. Send slots let the reader run ahead of the collective. SEND_SLOTS = 3 -RECV_SLOTS = 2 + +# 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): @@ -252,28 +264,39 @@ def __init__(self, args, topology): # job: an all-gather delivers the whole window to every participant. self.piece = args.buffer_size window = self.piece * topology.node_count - - # Fail before allocating rather than dying inside a CUDA OOM. - budget = int(torch.cuda.get_device_properties(self.device).total_memory * 0.25) - if window > budget: + # 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: + per_window = max(align_down(budget // (RECV_DEV_SLOTS * topology.node_count), + self.align), self.align) raise RuntimeError( f"--buffer-size {self.piece >> 20}M across {topology.node_count} nodes " - f"needs a {window / 1024 ** 3:.1f} GiB window per rank, over the " + 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"{max(align_down(budget // topology.node_count, self.align), self.align) >> 20}M." + f"{per_window >> 20}M." ) self.send_host = [pinned_aligned(self.piece, self.align) for _ in range(SEND_SLOTS)] - self.recv_host = [pinned_aligned(window, self.align) for _ in range(RECV_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) + self.recv_dev = [torch.empty(window, dtype=torch.uint8, device=self.device) + for _ in range(RECV_DEV_SLOTS)] if topology.rank == 0: - pinned = SEND_SLOTS * self.piece + RECV_SLOTS * window + 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"{(self.piece + window) / 1024 ** 3:.2f} GiB device, " - f"{pinned / 1024 ** 3:.2f} GiB pinned host", flush=True) + 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.""" @@ -306,26 +329,14 @@ def _read_piece(self, fd, mview, offset, length, block_size): for future in futures: future.result() - def _write_segments(self, fd, mview, layout, round_index, block_size, ready): - """Scatter one assembled window to its true file offsets. + 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. - `ready` is the event marking the end of the device-to-host copy that - filled `mview`. Each writer waits on it rather than the main thread, so - the main thread can go on and launch the next round's collective while - the copy is still in flight. + 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. """ - def write_one(base, length, offset): - ready.synchronize() - return direct_io.pwrite(fd, mview[base:base + layout.piece], - length, offset, block_size) - - futures = [] - for node in range(self.topo.node_count): - offset, length = layout.segment(node, round_index) - if length == 0: - continue - futures.append(self.pool.submit(write_one, node * layout.piece, length, offset)) - return futures + ready.synchronize() + return direct_io.pwrite(fd, mview[base:base + stride], length, offset, block_size) def stage_file(self, src, dst, size, mtime_ns): topo = self.topo @@ -363,6 +374,15 @@ def stage_file(self, src, dst, size, mtime_ns): # 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) @@ -394,8 +414,12 @@ 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() @@ -414,14 +438,89 @@ def reader(): 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. + """ + 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 + assembled.synchronize() + 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 + 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. + 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() - pending = [[] for _ in range(RECV_SLOTS)] try: for round_index in range(layout.rounds): item = filled_q.get() - if item is None: + if item is None or failure: break _, slot, length = item @@ -434,35 +533,25 @@ def reader(): # 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. - dist.all_gather_into_tensor(self.recv_dev, self.send_dev, group=topo.group) + 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) - recv_slot = round_index % RECV_SLOTS - for future in pending[recv_slot]: - future.result() - pending[recv_slot] = [] - - _, recv_view, recv_mview = self.recv_host[recv_slot] - recv_view.copy_(self.recv_dev, non_blocking=True) - d2h_done = torch.cuda.Event(blocking=True) - d2h_done.record() - - # The writers wait on d2h_done themselves, so the main thread - # returns to the top of the loop immediately. recv_dev is safe - # to reuse because the next round's collective is queued behind - # this copy, and recv_host[recv_slot] is held by the pending - # futures checked above. - pending[recv_slot] = self._write_segments( - fd_dst, recv_mview, layout, round_index, dst_block, d2h_done - ) - for futures in pending: - for future in futures: - future.result() + # 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 is - # unbounded, so the reader can never block on the other side. + # 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) From 07b9870777ef73388cdb6b2a6178faf47da778f6 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 10:26:13 -0500 Subject: [PATCH 06/21] tests: cover the fileio layout, copy planning, and the staging pipeline Stdlib only, no pytest and no numpy, so `python3 tests/run_tests.py` works anywhere. torch and torch.distributed are stubbed, which is what lets datastage run without a GPU or a job: tensors become memoryviews over ctypes buffers, CUDA events become no-ops, and the multi-node cases fake the all-gather by filling each node's segment with the bytes that node would have read. test_pipeline compares content rather than checking for exceptions, because the failure mode of a mishandled buffer handoff is a file of exactly the right length holding the wrong bytes. The two bugs fixed in the previous commit both pass a "did it throw?" check and both fail this one; reintroducing either makes it go red, which is the property worth keeping. test_layout checks the two invariants that corrupt data silently if broken: every byte claimed exactly once, and every interior segment boundary aligned, since direct_io pads writes up to the block size. tests/README.md records what this deliberately does not cover -- real NCCL, real CUDA events, pinned memory alignment, O_DIRECT, and anything about throughput. A green run says the arithmetic and the choreography are right, not that staging works on a cluster. tests/ has no __init__.py, so find_packages() does not pick it up. Co-Authored-By: Claude Opus 5 --- tests/README.md | 57 +++++++++++++ tests/run_tests.py | 47 +++++++++++ tests/stubs.py | 136 +++++++++++++++++++++++++++++++ tests/test_copyplan.py | 110 +++++++++++++++++++++++++ tests/test_layout.py | 105 ++++++++++++++++++++++++ tests/test_pipeline.py | 178 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 633 insertions(+) create mode 100644 tests/README.md create mode 100644 tests/run_tests.py create mode 100644 tests/stubs.py create mode 100644 tests/test_copyplan.py create mode 100644 tests/test_layout.py create mode 100644 tests/test_pipeline.py diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..5d5fd56 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,57 @@ +# 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`. + +## 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 | +| `test_pipeline.py` | `stage_file` end to end: bytes in == bytes out | +| `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 +* 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. 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..454e735 --- /dev/null +++ b/tests/stubs.py @@ -0,0 +1,136 @@ +#!/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 ctypes +import importlib.util +import os +import sys +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): + 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 FakeEvent: + """CUDA event stand-in: everything is synchronous on the CPU already.""" + + def __init__(self, blocking=False): + pass + + def record(self): + pass + + 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 *args, **kwargs: "cpu" + torch.cuda = types.SimpleNamespace( + Event=FakeEvent, + current_device=lambda: 0, + set_device=lambda device: None, + get_device_properties=lambda device: types.SimpleNamespace( + total_memory=total_memory), + ) + dist.barrier = lambda *args, **kwargs: None + dist.get_rank = lambda: 0 + dist.broadcast_object_list = lambda payload, src=0: None + dist.all_reduce = lambda tensor: 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_copyplan.py b/tests/test_copyplan.py new file mode 100644 index 0000000..6985fc4 --- /dev/null +++ b/tests/test_copyplan.py @@ -0,0 +1,110 @@ +#!/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 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) + 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_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..40023c9 --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,178 @@ +#!/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 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)}") + 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()) From b12ff41f83a46f012baab8cc5a5e15492c3c2e76 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 12:16:37 -0500 Subject: [PATCH 07/21] docs: add CLAUDE.md and record the branch review findings CLAUDE.md documents the two halves of the repo, the fileio dependency rule (only datastage may import torch), and the datastage decomposition, so a reader does not have to reconstruct the collective staging design from comments spread across 715 lines. It also corrects the record on srun --distribution=arbitrary. Topology's comments present support for it as a design invariant; it is not one. The requirement was a mistake carried over from the C version, and ranks are block- or cyclic-distributed in every supported launch. REVIEW-FINDINGS.md and .review-findings.json carry the results of a review of this branch: 15 verified findings, tiered, with repro steps and status boxes. They are a working record, not a deliverable -- delete them once the list is worked through. The headline item is a missing torch.cuda.set_device on the drainer thread, which records the copy-completion events on GPU 0 and lets writers race the D2H DMA on every rank but LOCAL_RANK 0. Co-Authored-By: Claude Opus 5 --- .review-findings.json | 157 +++++++++++++++++++++++++++++ CLAUDE.md | 148 +++++++++++++++++++++++++++ REVIEW-FINDINGS.md | 227 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 532 insertions(+) create mode 100644 .review-findings.json create mode 100644 CLAUDE.md create mode 100644 REVIEW-FINDINGS.md 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..b80f07e --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,148 @@ +# 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 + +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` derives node grouping from what ranks actually report via +`all_gather_object` rather than assuming RANK is block-distributed. Its comments +justify this by the need to support `srun --distribution=arbitrary` — **that is +not a real requirement.** It was a mistake by the original author of the C +version, carried into the Python port unexamined. Ranks are block- or +cyclic-distributed in every supported launch. Do not add complexity to serve the +arbitrary case, and treat the existing hostname-derivation machinery as open to +simplification rather than as an invariant to preserve. `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, under `slurm2pytorch` so RANK/WORLD_SIZE/LOCAL_RANK/ +MASTER_ADDR are set: + +```bash +srun --ntasks-per-node=${DGXNGPU} ... slurm2pytorch \ + python3 -m mlperf_common.fileio.datastage -r "${SLOW_DATADIR}/${DATASET}" "${DATADIR}" +``` + +`--dry-run` without RANK set prints the copy plan on a single node. + +### 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. + +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/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md new file mode 100644 index 0000000..d12f06f --- /dev/null +++ b/REVIEW-FINDINGS.md @@ -0,0 +1,227 @@ +# 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 exists to serve this + non-requirement. Simplifying it is open for discussion; it is not in scope of + any finding below. +* `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 + +- [ ] **F1 · `mlperf_common/fileio/datastage.py:476` · set_device missing on drainer thread** + + 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. + + Not catchable by the suite today: `stubs.py` makes `FakeEvent.synchronize()` a + no-op. See F14/F15. + +- [ ] **F2 · `mlperf_common/fileio/datastage.py:199` · new_group sorts its rank list** *(demoted — see premise correction)* + + `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. + + Recommended: assert `group_ranks[l] == sorted(group_ranks[l])` and fail loudly, + rather than carrying the machinery that pretends to handle the general case. + +## Tier 2 — silent data loss / job hangs + +- [ ] **F3 · `mlperf_common/fileio/copyplan.py:82` · os.walk swallows unreadable subtrees** + + 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. + +- [ ] **F4 · `mlperf_common/fileio/datastage.py:649` · no destination-is-a-directory check** + + `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. + + `client/fastcp` blocks both cases at `parse_and_validate_args:172-179`. + Port that check. + +- [ ] **F5 · `mlperf_common/fileio/datastage.py:572` · unguarded os.stat before the broadcast** + + 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. + +## Tier 3 — real, lower severity + +- [ ] **F6 · `client/fastmd5:111` · directory-symlink cycles now abort the run** + + 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. + +- [ ] **F7 · `mlperf_common/fileio/datastage.py:406` · `_chmod_parents` infinite loop on `/`** + + `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. + +- [ ] **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. + +- [ ] **F9 · `mlperf_common/fileio/datastage.py:343` · `.datastage.tmp` files leak** + + 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. + +- [ ] **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. + +- [ ] **F11 · `mlperf_common/fileio/datastage.py:277` · memory-budget error suggests an over-budget value** + + 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. + +- [ ] **F12 · `mlperf_common/fileio/datastage.py:475` · no dedicated stream, so no overlap** + + 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. + +- [ ] **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. + +## 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. + +- [ ] **F15 · `tests/stubs.py:106` · stubs can't construct `Topology` or `build_plan`** + + `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. + +--- + +## 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. From 9035aa405938e272496a9ae215dc34a5573e2f99 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 12:26:35 -0500 Subject: [PATCH 08/21] datastage: set the drainer thread's CUDA device CUDA's current device is per host thread, and main() sets it on the main thread only, so the drain thread started by _run_pipeline was on device 0. The device-to-host copies still ran on the right device -- the tensors carry it -- but torch.cuda.Event binds to the calling thread's current device when recorded, so the copy-completion events landed on device 0's idle stream. Both synchronize() calls that depend on them then returned immediately: _write_one's, so the writer pool read a pinned chunk before the copy filled it, and the one guarding recv_free_q, so a device window was recycled while copies were still reading it and the next all-gather overwrote it. The result was a file of exactly the right length holding wrong bytes, no exception, on every rank except LOCAL_RANK 0 -- 7 of 8 slices per file on a DGX node. test_device.py guards it. The stubs now model a per-thread current device and have events remember which one they were recorded against, which is the part of this that is checkable without a GPU: it goes red on the unfixed code with 3 of 9 events on device 0, and green after. It does not check CUDA ordering semantics, so a multi-GPU confirmation is still worth doing. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 13 +++- mlperf_common/fileio/datastage.py | 10 +++ tests/README.md | 8 +- tests/stubs.py | 54 +++++++++++-- tests/test_device.py | 121 ++++++++++++++++++++++++++++++ 5 files changed, 196 insertions(+), 10 deletions(-) create mode 100644 tests/test_device.py diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index d12f06f..59c094c 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -35,7 +35,7 @@ Consequences: ## Tier 1 — silent data corruption -- [ ] **F1 · `mlperf_common/fileio/datastage.py:476` · set_device missing on drainer thread** +- [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 @@ -53,8 +53,15 @@ Consequences: Fires on every multi-GPU node — independent of the distribution correction above. Fix first. - Not catchable by the suite today: `stubs.py` makes `FakeEvent.synchronize()` a - no-op. See F14/F15. + **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. - [ ] **F2 · `mlperf_common/fileio/datastage.py:199` · new_group sorts its rank list** *(demoted — see premise correction)* diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 9580899..9597551 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -446,6 +446,16 @@ def drainer(): 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: diff --git a/tests/README.md b/tests/README.md index 5d5fd56..6604cff 100644 --- a/tests/README.md +++ b/tests/README.md @@ -18,6 +18,7 @@ Each test runs in its own interpreter, because each installs its own fake | `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 | | `test_pipeline.py` | `stage_file` end to end: bytes in == bytes out | +| `test_device.py` | every CUDA event is recorded against this rank's device | | `stubs.py` | fake `torch` / `torch.distributed` so the above run on a CPU | ## Why the bytes are compared, not just the exit status @@ -46,7 +47,12 @@ If you change the pipeline, confirm a deliberately reintroduced bug still makes 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 + 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 diff --git a/tests/stubs.py b/tests/stubs.py index 454e735..9f76850 100644 --- a/tests/stubs.py +++ b/tests/stubs.py @@ -31,6 +31,7 @@ import importlib.util import os import sys +import threading import types REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -74,14 +75,55 @@ 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 is per *host thread*, and a thread that never called +# set_device gets device 0 no matter what any other thread did. Modelling that +# faithfully is the whole point: it is the property 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 = [] + + +def current_device(): + return getattr(_CURRENT, "index", 0) + + +def set_device(device): + _CURRENT.index = device.index if isinstance(device, FakeDevice) else int(device) + + class FakeEvent: - """CUDA event stand-in: everything is synchronous on the CPU already.""" + """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): - pass + self.device = None def record(self): - pass + self.device = current_device() + EVENTS.append(self) def synchronize(self): pass @@ -95,11 +137,11 @@ def install(total_memory=288 * 1024 ** 3): torch.int64 = "int64" torch.distributed = dist torch.empty = lambda n, dtype=None, device=None, pin_memory=False: FakeTensor(n) - torch.device = lambda *args, **kwargs: "cpu" + torch.device = lambda kind="cuda", index=0: FakeDevice(kind, index) torch.cuda = types.SimpleNamespace( Event=FakeEvent, - current_device=lambda: 0, - set_device=lambda device: None, + current_device=current_device, + set_device=set_device, get_device_properties=lambda device: types.SimpleNamespace( total_memory=total_memory), ) diff --git a/tests/test_device.py b/tests/test_device.py new file mode 100644 index 0000000..cfe0094 --- /dev/null +++ b/tests/test_device.py @@ -0,0 +1,121 @@ +#!/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. + +"""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. + +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. +""" + +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 + +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[:] + stager.stage_file(src, os.path.join(root, "dst"), size, + os.stat(src).st_mtime_ns) + events = list(stubs.EVENTS) + finally: + shutil.rmtree(root, ignore_errors=True) + + if not events: + print(" FAIL no events recorded at all; this test is exercising nothing") + return 1 + + stray = [event for event in events if event.device != DEVICE] + print(f"test_device: {len(events)} events recorded, {len(stray)} on the wrong device") + if stray: + wrong = sorted({event.device for event in stray}) + print(f" FAIL {len(stray)}/{len(events)} events recorded against device " + f"{wrong} instead of {DEVICE}: a thread that records events " + f"never called torch.cuda.set_device, so synchronizing on them " + f"waits for nothing") + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 4ab9a131ae5c23a7ce761b91a96c246f02a22484 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 13:28:38 -0500 Subject: [PATCH 09/21] datastage: give the drainer its own CUDA stream The copy-back was issued on the default stream, the same one carrying the H2D copy and the all-gather, so it could not run beside them. non_blocking =True only promises the host will not wait; the copy still queues on the stream it was issued to, and ProcessGroupNCCL orders the collective against that stream in both directions. Every round therefore ran H2D, then the collective, then the copy-back, strictly in turn -- which made the second device window RECV_DEV_SLOTS allocates, and fights the memory budget check over, do nothing at all. The copies and their events now go to a dedicated stream, and the host-side assembled.synchronize() becomes drain_stream.wait_event(assembled), so the drain thread queues a round's copies while that round's collective is still in flight rather than blocking until it lands. That leaves the window-reuse handoff as the one piece of cross-stream safety with nothing but a host wait behind it: the next all-gather goes to the default stream, which has no ordering against drain_stream, so last_copy.synchronize() before recv_free_q.put is what stops it overwriting a window still being read. Commented in place. test_device.py now also requires that no drainer operation is issued on the default stream, and that a side stream waits on an event before its first copy -- the omission that would turn this from a throughput change into a race. Both confirmed red before the change. The speedup itself is unmeasured: the stubs make every copy instantaneous, and whether the GPU side is the bottleneck depends on Lustre and NVMe rates. Worth a profile on a real node. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 24 ++++++++- mlperf_common/fileio/datastage.py | 31 +++++++++-- tests/README.md | 2 +- tests/stubs.py | 85 +++++++++++++++++++++++++++++-- tests/test_device.py | 73 ++++++++++++++++++++++---- 5 files changed, 194 insertions(+), 21 deletions(-) diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index 59c094c..43dda1f 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -169,7 +169,7 @@ Consequences: 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. -- [ ] **F12 · `mlperf_common/fileio/datastage.py:475` · no dedicated stream, so no overlap** +- [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 @@ -179,6 +179,28 @@ Consequences: 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 diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 9597551..598f388 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -289,6 +289,13 @@ def __init__(self, args, topology): 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, " @@ -463,7 +470,11 @@ def drainer(): if item is None: break dev_slot, round_index, assembled = item - assembled.synchronize() + # 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)): @@ -481,10 +492,12 @@ def drainer(): _, chunk_view, chunk_mview = self.chunk_host[chunk] base = first * layout.piece - chunk_view[:nbytes].copy_( - self.recv_dev[dev_slot][base:base + nbytes], non_blocking=True) - copied = torch.cuda.Event(blocking=True) - copied.record() + 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 = [] @@ -500,6 +513,14 @@ def drainer(): # 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) diff --git a/tests/README.md b/tests/README.md index 6604cff..d9ec70a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -18,7 +18,7 @@ Each test runs in its own interpreter, because each installs its own fake | `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 | | `test_pipeline.py` | `stage_file` end to end: bytes in == bytes out | -| `test_device.py` | every CUDA event is recorded against this rank's device | +| `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 diff --git a/tests/stubs.py b/tests/stubs.py index 9f76850..ceda347 100644 --- a/tests/stubs.py +++ b/tests/stubs.py @@ -27,6 +27,7 @@ memory alignment against a real filesystem block size, and O_DIRECT itself. """ +import contextlib import ctypes import importlib.util import os @@ -64,6 +65,7 @@ def __getitem__(self, item): 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 @@ -86,17 +88,65 @@ def __repr__(self): return f"device(type={self.type!r}, index={self.index})" -# CUDA's current device is per *host thread*, and a thread that never called -# set_device gets device 0 no matter what any other thread did. Modelling that -# faithfully is the whole point: it is the property a thread doing CUDA work -# can silently get wrong, and no amount of comparing staged bytes on a CPU will -# show it up. +# 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) @@ -106,6 +156,25 @@ 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. @@ -120,10 +189,13 @@ class FakeEvent: 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 @@ -140,6 +212,9 @@ def install(total_memory=288 * 1024 ** 3): 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( diff --git a/tests/test_device.py b/tests/test_device.py index cfe0094..6e6af59 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -14,7 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Every CUDA event must be recorded against this rank's device, not device 0. +"""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 -- @@ -30,8 +36,17 @@ 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. +touch CUDA agree about which GPU they are on and which queue they are feeding. """ import os @@ -52,6 +67,10 @@ # 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"] @@ -95,9 +114,11 @@ def main(): 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) @@ -105,16 +126,50 @@ def main(): 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] - print(f"test_device: {len(events)} events recorded, {len(stray)} on the wrong device") if stray: wrong = sorted({event.device for event in stray}) - print(f" FAIL {len(stray)}/{len(events)} events recorded against device " - f"{wrong} instead of {DEVICE}: a thread that records events " - f"never called torch.cuda.set_device, so synchronizing on them " - f"waits for nothing") - return 1 - return 0 + 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__": From 8f66d5fd3b311dff68adbda369920534b33c5af7 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 13:38:43 -0500 Subject: [PATCH 10/21] copyplan: report directories the walk cannot list list_relative_files collected every entry it could not stat, but a directory it could not list produced nothing to collect: the files underneath never reach the walk, so there is no entry to stat and no problem to record. os.walk's default onerror swallows the error and carries on, which makes an unlistable subtree indistinguishable from an empty one. One dataset directory with a bad mode, an EIO, or a stale Lustre handle was therefore enough to stage a partial dataset and exit 0 -- and because fastmd5 enumerates through the same function, checksumming the staged tree omitted exactly the same files, so verification agreed. An unlistable root returned an empty list, equally silently. An onerror callback now appends to the same problems list, so these join the existing report and planning refuses the tree. This hole predates the fileio refactor: the same walk was in client/fastcp before copyplan existed. What changed is its reach, since fastcp, fastmd5 and datastage now share one enumeration and so share one blind spot. test_copyplan.py covers both the subtree and the root case, and verifies its own premise -- a user who can list a 0o000 directory gets a skip rather than a vacuous pass. Confirmed red first: the subtree case returned only the readable file, the root case returned nothing at all. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 29 +++++++++++++++-- mlperf_common/fileio/copyplan.py | 18 +++++++++-- tests/test_copyplan.py | 54 ++++++++++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 5 deletions(-) diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index 43dda1f..d5078db 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -80,7 +80,7 @@ Consequences: ## Tier 2 — silent data loss / job hangs -- [ ] **F3 · `mlperf_common/fileio/copyplan.py:82` · os.walk swallows unreadable subtrees** +- [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 @@ -92,6 +92,22 @@ Consequences: 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. + - [ ] **F4 · `mlperf_common/fileio/datastage.py:649` · no destination-is-a-directory check** `parse_args` validates sources but never checks the destination. @@ -119,7 +135,7 @@ Consequences: ## Tier 3 — real, lower severity -- [ ] **F6 · `client/fastmd5:111` · directory-symlink cycles now abort the run** +- [ ] **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 @@ -128,6 +144,15 @@ Consequences: 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. The proper fix is the depth-first walk with + visited-inode tracking that the standing `FIXME` in `list_relative_files` + describes, which would also close the cycle hole `os.walk(followlinks=True)` + leaves open. + - [ ] **F7 · `mlperf_common/fileio/datastage.py:406` · `_chmod_parents` infinite loop on `/`** `datastage a.bin /` makes `dest_root == "/"`; `os.path.dirname("/") == "/"`, so diff --git a/mlperf_common/fileio/copyplan.py b/mlperf_common/fileio/copyplan.py index 7659177..12b3d26 100644 --- a/mlperf_common/fileio/copyplan.py +++ b/mlperf_common/fileio/copyplan.py @@ -70,8 +70,9 @@ def list_relative_files(root): 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, reporting - every such entry rather than dying on the first one. + 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. 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 @@ -79,7 +80,18 @@ def list_relative_files(root): """ file_list = [] problems = [] - for dirpath, _, filenames in os.walk(root, followlinks=True): + + 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: diff --git a/tests/test_copyplan.py b/tests/test_copyplan.py index 6985fc4..bafb0e2 100644 --- a/tests/test_copyplan.py +++ b/tests/test_copyplan.py @@ -49,6 +49,58 @@ def build_tree(root): os.symlink("realdir", os.path.join(root, "link_to_dir")) +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: @@ -96,6 +148,8 @@ def main(): check("planning refuses an unreadable tree", False, "no exception raised") except copyplan.UnreadableEntries: check("planning refuses an unreadable tree", True) + + check_unlistable_directory(root) finally: shutil.rmtree(root, ignore_errors=True) From 023586ae3ee9b7ca88e7f1790e9360165d858bea Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 13:51:15 -0500 Subject: [PATCH 11/21] copyplan: own the cp argument rules, and fix cp -r src newdir Planning moved into this module when fileio was extracted; the argument rules that make the planning well defined stayed behind in fastcp's CLI. plan_copy_operations documented its fallback as a fact -- "otherwise this is a single file-to-file copy" -- when it was really a precondition the caller had to establish. datastage was the first caller to arrive without it, and so planned one job and silently dropped sources[1:] when the destination was not an existing directory. The rules now live in validate_copy_args next to the mapping they govern, and plan_copy_operations applies them itself, so a caller cannot skip them. fastcp's destination block is deleted in favour of the shared call and datastage's parse_args gained it -- there rather than in build_plan, because parse_args runs on every rank before the process group exists, so a bad invocation has to kill the job uniformly instead of leaving rank 0 exiting while its peers block in a collective. That also fixes a case fastcp got wrong on its own. cp -r src newdir, with newdir absent, is legal and copies src's *contents* into newdir. fastcp correctly permitted it and the planner then treated the source directory as a 4 KiB file, so the invocation was broken in both tools -- and it is the natural way to stage onto empty node-local scratch. A directory source with a non-directory destination now maps contents directly under it, with no basename level. fastmd5 is unaffected: read-only, no destination. test_copyplan pins all six cases against what GNU cp actually does, and both CLIs were exercised end to end, error wording included. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 38 ++++++++++++- client/fastcp | 26 ++++----- mlperf_common/fileio/copyplan.py | 93 ++++++++++++++++++++++++++++--- mlperf_common/fileio/datastage.py | 18 ++++-- tests/test_copyplan.py | 68 ++++++++++++++++++++++ 5 files changed, 209 insertions(+), 34 deletions(-) diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index d5078db..a2bab63 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -108,7 +108,7 @@ Consequences: verifies its own premise and skips if run as a user who can list a 0o000 directory, rather than passing vacuously. -- [ ] **F4 · `mlperf_common/fileio/datastage.py:649` · no destination-is-a-directory check** +- [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 @@ -121,8 +121,40 @@ Consequences: block in the collective until the watchdog fires. With multiple sources, `sources[1:]` are dropped with no message and the job exits 0. - `client/fastcp` blocks both cases at `parse_and_validate_args:172-179`. - Port that check. + **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. - [ ] **F5 · `mlperf_common/fileio/datastage.py:572` · unguarded os.stat before the broadcast** diff --git a/client/fastcp b/client/fastcp index 6ada473..e48b027 100755 --- a/client/fastcp +++ b/client/fastcp @@ -32,7 +32,8 @@ sys.path.insert(1, os.path.join(_HERE, os.pardir)) import direct_io try: - from mlperf_common.fileio.copyplan import UnreadableEntries, plan_copy_operations + 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" @@ -162,21 +163,14 @@ 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 diff --git a/mlperf_common/fileio/copyplan.py b/mlperf_common/fileio/copyplan.py index 12b3d26..c61b6b5 100644 --- a/mlperf_common/fileio/copyplan.py +++ b/mlperf_common/fileio/copyplan.py @@ -20,11 +20,26 @@ 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__ = ["UnreadableEntries", "list_relative_files", "plan_copy_operations"] +__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 @@ -54,6 +69,48 @@ def _describe(self): 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: @@ -105,13 +162,19 @@ 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 this is a - single file-to-file copy. The result is sorted by destination path so that - every rank of a collective copy walks the files in the same order. - - Raises UnreadableEntries if any source cannot be stat'd, reporting all of - them together. + 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) @@ -121,8 +184,20 @@ def add(src_abs, dst_abs): if st is not None: file_jobs.append((src_abs, dst_abs, st.st_size)) - if not os.path.isdir(dst_root): # case 1: single file copy - add(os.path.abspath(sources[0]), dst_root) + 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) diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 598f388..46bc24d 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -63,7 +63,8 @@ import torch.distributed as dist from mlperf_common.fileio import direct_io -from mlperf_common.fileio.copyplan import UnreadableEntries, plan_copy_operations +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 @@ -677,11 +678,16 @@ def parse_args(argv=None): 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) - for src in args.sources: - if not os.path.exists(src): - sys.exit(f"{prog}: cannot stat '{src}': No such file or directory") - if os.path.isdir(src) and not args.recursive: - sys.exit(f"{prog}: -r not specified; omitting directory '{src}'") + # 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 diff --git a/tests/test_copyplan.py b/tests/test_copyplan.py index bafb0e2..d07ef69 100644 --- a/tests/test_copyplan.py +++ b/tests/test_copyplan.py @@ -49,6 +49,73 @@ def build_tree(root): 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. @@ -149,6 +216,7 @@ def main(): 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) From ee7abb24e95011e2b6b3a4674b015e59c0ac9205 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 14:02:21 -0500 Subject: [PATCH 12/21] copyplan: say what a symlink cycle actually does The FIXME warned that os.walk(followlinks) does not protect against cycles and implied a custom depth-first traversal was needed to avoid running away. Measured, that is not what happens: the kernel allows 40 symlink traversals per path resolution, so a `up -> ..` tree stops at 82 directories and depth 81 in under 10 ms, and os.walk is iterative so there is no stack to blow. The path to the error is also not the obvious one, which is why it was worth writing down. os.walk 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 nothing at all -- not even through the onerror hook added in 8f66d5f. The ELOOP surfaces from the stat instead, as a single unreadable entry, and the copy is refused. The note now also records the trap: the walk lists the files under the cycle once per level on the way down, so tolerating the ELOOP rather than detecting the cycle would replace a loud refusal with ~40 redundant copies of everything beneath it -- turning the one loud failure in this area into a silent one. Comment only; no behaviour change. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 22 ++++++++++++++++++---- mlperf_common/fileio/copyplan.py | 20 +++++++++++++++++--- 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index a2bab63..c4ad429 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -180,10 +180,24 @@ Consequences: 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. The proper fix is the depth-first walk with - visited-inode tracking that the standing `FIXME` in `list_relative_files` - describes, which would also close the cycle hole `os.walk(followlinks=True)` - leaves open. + 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. - [ ] **F7 · `mlperf_common/fileio/datastage.py:406` · `_chmod_parents` infinite loop on `/`** diff --git a/mlperf_common/fileio/copyplan.py b/mlperf_common/fileio/copyplan.py index c61b6b5..263ed17 100644 --- a/mlperf_common/fileio/copyplan.py +++ b/mlperf_common/fileio/copyplan.py @@ -131,9 +131,23 @@ def list_relative_files(root): directory under it cannot be listed, reporting every such entry rather than dying on the first one. - 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. + 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 = [] From 429d952463389c89a27dd124b565cc0099c9c270 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 14:05:58 -0500 Subject: [PATCH 13/21] datastage: broadcast any planning failure, not just the expected one build_plan has rank 0 walk the tree and broadcast the answer, and the comment above it already said a failure has to be broadcast rather than raised, because every other rank is blocked in that broadcast. The catch was then narrowed to UnreadableEntries, which is the one failure the author went looking for. Anything else escaped: rank 0 unwound through destroy_process_group and exited while its peers waited for a message that would never come, so the job died on an NCCL watchdog timeout naming neither the file nor the reason, ten minutes of allocation later. The os.stat that collects st_mtime_ns re-stats files plan_copy_operations already stat'd and loses that race against anything modifying shared storage, which is the realistic trigger. Catching Exception fixes it. The bar is reaching the broadcast, not anticipating the cause -- narrowing the catch is what created the hang in the first place. There were two live paths, not one: the stat race, and CopyArgumentError, which became reachable here when 023586a moved the cp argument rules into plan_copy_operations. The new test caught the second on its own; it had only been written for the first. An operator now gets, on every rank together: cannot stage the source tree: FileNotFoundError: [Errno 2] No such file or directory: '/.../src/vanished.bin' test_buildplan.py drives rank 0's side and distinguishes a failure that was broadcast from one that escaped. torch.tensor joins the stubs so the success path is reachable. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 25 +++++- mlperf_common/fileio/datastage.py | 18 +++- tests/README.md | 3 +- tests/stubs.py | 11 +++ tests/test_buildplan.py | 135 ++++++++++++++++++++++++++++++ 5 files changed, 186 insertions(+), 6 deletions(-) create mode 100644 tests/test_buildplan.py diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index c4ad429..00eb5c9 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -156,7 +156,7 @@ Consequences: exercised end to end — fastcp copying real bytes, datastage's parse_args via the stub harness — with error text matching cp's wording. -- [ ] **F5 · `mlperf_common/fileio/datastage.py:572` · unguarded os.stat before the broadcast** +- [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 @@ -165,6 +165,29 @@ Consequences: 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)* diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 46bc24d..836873d 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -597,14 +597,24 @@ def build_plan(args): """Rank 0 walks the source tree; everyone else takes its answer verbatim.""" 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 and would hang - # until the NCCL watchdog fired. + # 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 UnreadableEntries as exc: - payload = [{"error": str(exc)}] + 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) diff --git a/tests/README.md b/tests/README.md index d9ec70a..b55bf67 100644 --- a/tests/README.md +++ b/tests/README.md @@ -16,7 +16,8 @@ Each test runs in its own interpreter, because each installs its own fake | 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 | +| `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_pipeline.py` | `stage_file` end to end: bytes in == bytes out | | `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 | diff --git a/tests/stubs.py b/tests/stubs.py index ceda347..97c7708 100644 --- a/tests/stubs.py +++ b/tests/stubs.py @@ -77,6 +77,16 @@ def data_ptr(self): return ctypes.addressof(self._buf) + self._offset +class FakeCounter: + """What torch.tensor([n]) is used for here: a scalar to all_reduce.""" + + def __init__(self, values): + self._values = list(values) + + def item(self): + return self._values[0] + + class FakeDevice: """Stand-in for torch.device('cuda', i). datastage only ever passes it on.""" @@ -209,6 +219,7 @@ def install(total_memory=288 * 1024 ** 3): torch.int64 = "int64" torch.distributed = dist torch.empty = lambda n, dtype=None, device=None, pin_memory=False: FakeTensor(n) + torch.tensor = lambda values, dtype=None, device=None: FakeCounter(values) torch.device = lambda kind="cuda", index=0: FakeDevice(kind, index) torch.cuda = types.SimpleNamespace( Event=FakeEvent, 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()) From 6799f47b6526836e25c669f791d45e5f995216ac Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 14:25:38 -0500 Subject: [PATCH 14/21] datastage: drop the cross-rank source check, it is a metadata storm 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 W x F stats, issued as a synchronised burst right after the broadcast: 5.1M at 64 nodes and 10k files, 41M at 512 nodes, 1.6 billion at 2048 nodes and 100k files. At a generous 50k ops/s for a single MDT that is roughly fourteen minutes of pure metadata at 512 nodes before a byte of data moves, and at 2048 it is an outage for every other job on the filesystem, not just this one. The cost scaled with the file count and the detection power did not. What it looked for -- a stale handle, a failed mount, the wrong dataset at the same path -- is a property of a mount, and affects every file on that node identically, so the ten-thousandth stat says nothing the first did not. It was also imperfect, since size and mtime agree on same-size same-mtime content differences. Removed rather than reduced. Mount verification already belongs to mountcheck.py, once per job, where a sparse SHA256 fingerprint is both cheaper and stronger; errors of this kind have not been seen in six years on these clusters; and small file copies are to be parallelised so that each small file is handled, and its metadata touched, by a single rank -- which makes per-rank per-file metadata work the wrong shape whatever its cost. The reasoning is in build_plan's docstring so it does not get reinvented. torch.tensor, FakeCounter and dist.all_reduce leave the stubs with it, having no remaining user. Staging still opens every file on every rank, which is inherent to having W disjoint readers. This was pure addition on top of that. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 32 ++++++++++++++++++++++++++ mlperf_common/fileio/datastage.py | 38 ++++++++++++++----------------- tests/stubs.py | 12 ---------- 3 files changed, 49 insertions(+), 33 deletions(-) diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index 00eb5c9..6e531bd 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -305,6 +305,38 @@ Consequences: 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 `W × F` + stats 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 does `W × F` opens, which is inherent to + having W disjoint readers; this was pure addition on top. + ## Tier 4 — the tests can't catch the above - [ ] **F14 · `tests/test_pipeline.py:90` · fake all-gather re-reads the source** diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 836873d..77fdc4c 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -594,7 +594,22 @@ def drainer(): def build_plan(args): - """Rank 0 walks the source tree; everyone else takes its answer verbatim.""" + """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 @@ -620,26 +635,7 @@ def build_plan(args): dist.broadcast_object_list(payload, src=0) if isinstance(payload[0], dict): raise RuntimeError(f"cannot stage the source tree: {payload[0]['error']}") - jobs = payload[0] - - # Guard against ranks seeing a different view of shared storage. - mismatches = 0 - for src, _, size, mtime_ns in jobs: - try: - st = os.stat(src) - except OSError: - mismatches += 1 - continue - if st.st_size != size or st.st_mtime_ns != mtime_ns: - mismatches += 1 - counter = torch.tensor([mismatches], dtype=torch.int64, device="cuda") - dist.all_reduce(counter) - if counter.item(): - raise RuntimeError( - f"{counter.item()} rank/file pairs disagree with rank 0 about the " - "source tree; shared storage is inconsistent or changing" - ) - return jobs + return payload[0] def parse_args(argv=None): diff --git a/tests/stubs.py b/tests/stubs.py index 97c7708..361c2dc 100644 --- a/tests/stubs.py +++ b/tests/stubs.py @@ -77,16 +77,6 @@ def data_ptr(self): return ctypes.addressof(self._buf) + self._offset -class FakeCounter: - """What torch.tensor([n]) is used for here: a scalar to all_reduce.""" - - def __init__(self, values): - self._values = list(values) - - def item(self): - return self._values[0] - - class FakeDevice: """Stand-in for torch.device('cuda', i). datastage only ever passes it on.""" @@ -219,7 +209,6 @@ def install(total_memory=288 * 1024 ** 3): torch.int64 = "int64" torch.distributed = dist torch.empty = lambda n, dtype=None, device=None, pin_memory=False: FakeTensor(n) - torch.tensor = lambda values, dtype=None, device=None: FakeCounter(values) torch.device = lambda kind="cuda", index=0: FakeDevice(kind, index) torch.cuda = types.SimpleNamespace( Event=FakeEvent, @@ -234,7 +223,6 @@ def install(total_memory=288 * 1024 ** 3): dist.barrier = lambda *args, **kwargs: None dist.get_rank = lambda: 0 dist.broadcast_object_list = lambda payload, src=0: None - dist.all_reduce = lambda tensor: None dist.all_gather_into_tensor = lambda out, inp, group=None: out.copy_(inp) sys.modules["torch"] = torch sys.modules["torch.distributed"] = dist From 56e4c49090ea2c1e2b66d875a4d9427c4f027c58 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 14:37:01 -0500 Subject: [PATCH 15/21] docs: define W and F where the shorthand is used REVIEW-FINDINGS used W x F cold. W is datastage's own notation, from the module docstring, but F was invented in the F16 write-up and defined nowhere, so the numbers that justify removing the check were unreadable without guessing at them. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index 6e531bd..a01d795 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -310,8 +310,11 @@ Consequences: - [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 `W × F` - stats issued as a synchronised burst: 5.1M at 64 nodes / 10k files, 41M at + 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 @@ -334,8 +337,9 @@ Consequences: `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 does `W × F` opens, which is inherent to - having W disjoint readers; this was pure addition on top. + 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 From b0746624589b12e6aad104dd39040eef95d56d63 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 14:50:07 -0500 Subject: [PATCH 16/21] datastage: derive the node layout arithmetically, and check it Topology gathered every rank's hostname and LOCAL_RANK to work out which ranks shared a node, so that --distribution=arbitrary would group correctly. That was never a real requirement: it came from the C version by mistake, and it could not have worked anyway, since MASTER_ADDR is the first node in the nodelist and rank 0 has to be there for the rendezvous to come up at all. Under slurm's default block distribution the layout is arithmetic. Node i holds ranks i*L through i*L+L-1, so a rank's node is RANK // L, its slot is RANK % L, and group_ranks[l] is [node*L + l for node in range(N)]. No collective, no hostnames, no socket import. The assumption is verified rather than assumed, and the check is free: slurm reports RANK as SLURM_PROCID and LOCAL_RANK as SLURM_LOCALID, two independently derived numbers that agree only under a block distribution. A modulo comparison on each rank rejects --distribution=cyclic and =arbitrary, and a ragged --ntasks-per-node, with no communication -- turning a launch this code cannot handle into an immediate error instead of groups whose all-gather assembles the right bytes in the wrong order. LOCAL_WORLD_SIZE is now required rather than defaulting to 1, which had turned a single 8-GPU node into eight single-rank "nodes". This also settles the new_group question. group_ranks[l] is ascending by construction, so new_group's internal sort is provably a no-op and each member's group position is its node_index, which is what the drainer assumes when it maps all-gather output position to a file offset. Noted where it matters, and checked. Topology needed all_gather_object, which the stubs could not sensibly fake, so nothing had ever constructed it. Now it needs only new_group. test_topology.py covers six layouts and the rejected launches, and was confirmed able to go red by reintroducing a mis-grouping. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 28 ++++-- REVIEW-FINDINGS.md | 36 +++++-- mlperf_common/fileio/datastage.py | 84 +++++++++-------- tests/README.md | 1 + tests/stubs.py | 5 + tests/test_topology.py | 150 ++++++++++++++++++++++++++++++ 6 files changed, 249 insertions(+), 55 deletions(-) create mode 100644 tests/test_topology.py diff --git a/CLAUDE.md b/CLAUDE.md index b80f07e..cd48f13 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -87,14 +87,26 @@ whole job. With W ranks = N nodes × L ranks/node: * Result: W disjoint readers on the shared FS, one fabric crossing per byte per node, L concurrent all-gathers. -`Topology` derives node grouping from what ranks actually report via -`all_gather_object` rather than assuming RANK is block-distributed. Its comments -justify this by the need to support `srun --distribution=arbitrary` — **that is -not a real requirement.** It was a mistake by the original author of the C -version, carried into the Python port unexamined. Ranks are block- or -cyclic-distributed in every supported launch. Do not add complexity to serve the -arbitrary case, and treat the existing hostname-derivation machinery as open to -simplification rather than as an invariant to preserve. `FileLayout` keeps every offset and length aligned +`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. diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index a01d795..f61b5f3 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -21,9 +21,9 @@ 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 exists to serve this - non-requirement. Simplifying it is open for discussion; it is not in scope of - any finding below. +* `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. @@ -63,7 +63,7 @@ Consequences: threads agree on which device they are on, not CUDA ordering semantics. Worth a multi-GPU confirmation run when a node is free. -- [ ] **F2 · `mlperf_common/fileio/datastage.py:199` · new_group sorts its rank list** *(demoted — see premise correction)* +- [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 @@ -75,8 +75,21 @@ Consequences: 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. - Recommended: assert `group_ranks[l] == sorted(group_ranks[l])` and fail loudly, - rather than carrying the machinery that pretends to handle the general case. + **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 @@ -353,7 +366,7 @@ Consequences: multi-node case. Instrumentation shows the threaded branch is entered exactly once in the whole suite, in a case whose result is discarded. -- [ ] **F15 · `tests/stubs.py:106` · stubs can't construct `Topology` or `build_plan`** +- [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 @@ -366,6 +379,15 @@ Consequences: 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 diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 77fdc4c..7af7efa 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -53,7 +53,6 @@ import argparse import os import queue -import socket import sys import threading import time @@ -151,55 +150,60 @@ def pinned_aligned(nbytes, alignment): class Topology: - """Rank/node layout and the per-LOCAL_RANK process groups.""" + """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.get("LOCAL_WORLD_SIZE", 1)) - - # Derive the node layout from what the ranks actually report rather - # than assuming a block distribution of RANK: run.sub launches some - # steps with --distribution=arbitrary, and silently mis-grouping would - # produce a corrupt copy rather than an error. - identities = [None] * self.world_size - dist.all_gather_object(identities, (socket.gethostname(), self.local_rank)) - - hosts = [] - for host, _ in identities: - if host not in hosts: - hosts.append(host) - self.node_count = len(hosts) - node_index = {host: i for i, host in enumerate(hosts)} - - # group_ranks[l][node] = global rank of LOCAL_RANK l on that node - group_ranks = [[None] * self.node_count for _ in range(self.local_world_size)] - for global_rank, (host, local_rank) in enumerate(identities): - if local_rank >= self.local_world_size: - raise RuntimeError( - f"rank {global_rank} reports LOCAL_RANK={local_rank} with " - f"LOCAL_WORLD_SIZE={self.local_world_size}" - ) - slot = group_ranks[local_rank][node_index[host]] - if slot is not None: - raise RuntimeError( - f"ranks {slot} and {global_rank} both claim LOCAL_RANK=" - f"{local_rank} on {host}" - ) - group_ranks[local_rank][node_index[host]] = global_rank - for local_rank, ranks in enumerate(group_ranks): - if any(r is None for r in ranks): - raise RuntimeError( - f"not every node has a rank with LOCAL_RANK={local_rank}; " - "launch with a uniform --ntasks-per-node" - ) + 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] - self.node_index = node_index[socket.gethostname()] def describe(self): return ( diff --git a/tests/README.md b/tests/README.md index b55bf67..5c02cdb 100644 --- a/tests/README.md +++ b/tests/README.md @@ -18,6 +18,7 @@ Each test runs in its own interpreter, because each installs its own fake | `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_pipeline.py` | `stage_file` end to end: bytes in == bytes out | | `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 | diff --git a/tests/stubs.py b/tests/stubs.py index 361c2dc..f069476 100644 --- a/tests/stubs.py +++ b/tests/stubs.py @@ -220,6 +220,11 @@ def install(total_memory=288 * 1024 ** 3): 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 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()) From ebf234463fa25b3361c6380adac5002879283678 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 15:56:32 -0500 Subject: [PATCH 17/21] dist_env: derive the torch rendezvous environment from slurm torch.distributed's env:// rendezvous wants RANK, WORLD_SIZE, MASTER_ADDR and MASTER_PORT, and by convention LOCAL_RANK and LOCAL_WORLD_SIZE beside 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 rather than being wrapped. The two are behaviourally compatible, and that is checked rather than asserted: fed the same pre-wrapper environment, the script and configure() produce identical values for all seven variables. A program launched under the wrapper finds them already set and takes them as given (source == "preset"), which is the case the test deliberately feeds contradictory slurm variables to catch. Two divergences, both refusing to proceed where the script would guess: slurm2pytorch defaults LOCAL_WORLD_SIZE to 1, but SLURM_NTASKS_PER_NODE is only set when --ntasks-per-node was actually passed. `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 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 says "will fail for multinode" -- as a rendezvous that hangs to the wall clock. We parse the address out of slurm's compressed nodelist instead, which is available in the container where `scontrol show hostnames` is not. Only the first name is needed, which is much less work than expanding the list, and it has to keep the zero padding: dgx[001-004] is dgx001, and dgx1 does not resolve. If that fails on a multi-node job we say so immediately, and name MLPERF_SLURM_FIRSTNODE as the fix. Beyond that it cross-checks what it is given, since everything here feeds a rendezvous and a wrong rendezvous does not error, it hangs: slurm's own NNODES x NTASKS_PER_NODE must equal NTASKS, slurm and mpirun must not disagree about which process this is, ranks must fall inside their worlds, and the world must divide evenly by the node size. Whether the resulting layout is the block distribution datastage needs stays Topology's business. Co-Authored-By: Claude Opus 5 --- mlperf_common/dist_env.py | 371 ++++++++++++++++++++++++++++++++++++++ tests/test_dist_env.py | 341 +++++++++++++++++++++++++++++++++++ 2 files changed, 712 insertions(+) create mode 100644 mlperf_common/dist_env.py create mode 100644 tests/test_dist_env.py 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/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()) From 7a0d7c1c5fc7fcb61901e37d7f8db3bf687ed1f2 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 15:58:31 -0500 Subject: [PATCH 18/21] datastage: run under srun without the slurm2pytorch wrapper main() now calls dist_env.configure() before anything else, so the rendezvous variables are derived in-process and the bash wrapper is no longer part of the launch. Running under it still works: the variables are then already set and taken as given. srun --ntasks-per-node=${DGXNGPU} ... python3 -m ...fileio.datastage ... configure() runs on every rank before the process group exists and derives from environment alone, so a bad launch fails the whole job identically instead of leaving some ranks blocked in a collective -- the same reasoning already written above the argument validation in parse_args. Three details worth their comments: The dry-run shortcut keyed on `"RANK" not in os.environ`, which this change would have silently broken by always populating RANK. It now asks dist_env which launcher it found and takes the shortcut only for source == "single". Same behaviour, stated rather than inferred from a missing variable. The rendezvous banner prints before init_process_group, not after. A wrong MASTER_ADDR hangs inside the rendezvous, so a line printed afterwards never appears -- and diagnosing that hang is the only reason to print it. Rank 0 always, every rank under NV_MLPERF_DEBUG, as slurm2pytorch's debug echo did. set_device takes env.local_rank rather than os.environ.get("LOCAL_RANK", 0). That silent default is the same class of guess this whole change removes. Verified through the stub harness: with no launcher, --dry-run prints the plan and returns 0; as `srun -N2 -n16` with no --ntasks-per-node, it exits naming the flag instead of proceeding with sixteen imaginary single-rank nodes. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 44 ++++++++++++++++++++++++++++--- mlperf_common/fileio/datastage.py | 35 ++++++++++++++++++++---- tests/README.md | 5 ++++ 3 files changed, 75 insertions(+), 9 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index cd48f13..47c6a88 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -121,15 +121,47 @@ 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, under `slurm2pytorch` so RANK/WORLD_SIZE/LOCAL_RANK/ -MASTER_ADDR are set: +Run it as one task per GPU. No wrapper needed: ```bash -srun --ntasks-per-node=${DGXNGPU} ... slurm2pytorch \ +srun --ntasks-per-node=${DGXNGPU} ... \ python3 -m mlperf_common.fileio.datastage -r "${SLOW_DATADIR}/${DATASET}" "${DATADIR}" ``` -`--dry-run` without RANK set prints the copy plan on a single node. +`--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/ @@ -138,6 +170,10 @@ 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 diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 7af7efa..42753cd 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -43,11 +43,20 @@ python3 -m mlperf_common.fileio.datastage -r SRC... DST -Launch with one task per GPU, under slurm2pytorch so that RANK / WORLD_SIZE / -LOCAL_RANK / MASTER_ADDR are set: +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} ... slurm2pytorch \\ + 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 @@ -61,6 +70,7 @@ 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) @@ -704,7 +714,16 @@ def parse_args(argv=None): def main(argv=None): args = parse_args(argv) - if args.dry_run and "RANK" not in os.environ: + # 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: @@ -713,7 +732,13 @@ def main(argv=None): print(f"{src} -> {dst} ({size} bytes)") return 0 - torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 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() diff --git a/tests/README.md b/tests/README.md index 5c02cdb..10aa483 100644 --- a/tests/README.md +++ b/tests/README.md @@ -11,6 +11,10 @@ or run any one directly: 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 | @@ -19,6 +23,7 @@ Each test runs in its own interpreter, because each installs its own fake | `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 | | `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 | From 5cc0cfe8ac35ad8ecc57d55854ae70aec54318e1 Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 16:21:15 -0500 Subject: [PATCH 19/21] docs: notes on per-rank NIC binding and a possible MPI transport datastage uses NCCL, through torch, purely as a network transport: the data is host-resident at both ends and nothing computes on it, so the trip through the GPU costs two PCIe crossings a host-side transfer would not. The stated reason for NCCL is that it drives every NIC without per-cluster tuning, which is an argument about NIC aggregation rather than about GPUs -- so if an MPI program can get one NIC per rank, the GPU may be droppable entirely. bindpcie --ib=single is supposed to provide exactly that, and has probably not been exercised in years. These are notes for a session on a real node. Read off the script, and solid: the ibdev list comes from ibv_devinfo order, the guard hard-exits when the device and GPU counts do not divide, and the mapping is ibdevs[local_rank * num_ibdevs / num_gpus]. Inferred from reading, and needing hardware to confirm: that mapping is index arithmetic with no topology query anywhere, despite --help promising a device "near its GPU", so locality holds only if ibv_devinfo happens to enumerate in GPU order -- which is the way this flag would be actively harmful rather than merely useless. OMPI_MCA_btl_openib_if_include targets a BTL removed in OpenMPI 5, so UCX_NET_DEVICES is doing all the work, with the port hardcoded to :1. The guard likely passes under enroot, which sets MELLANOX_VISIBLE_DEVICES, and likely would not bare, where storage NICs push the count past the GPU count. And every diagnostic in that block uses `2>&1` where `>&2` was meant, so the errors land on stdout -- one reason a broken --ib=single could go unnoticed. Also records how to observe this properly: per-device port_xmit_data counters rather than inferred bandwidth, and a warning that UCX_MAX_RNDV_RAILS defaults to 2, so an unbound rank may already use two NICs and only the aggregate across ranks is a fair comparison. Co-Authored-By: Claude Opus 5 --- IB-BINDING-NOTES.md | 154 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 IB-BINDING-NOTES.md 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`. From 0a1533ea71a7fc9555c0d1a8e5118c141d1bcb6f Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 16:37:10 -0500 Subject: [PATCH 20/21] datastage: fix three ways staging fails badly rather than loudly Three independent bugs, none of which reports itself as one. _chmod_parents walked up towards the destination root with `path = os.path.dirname(path)`, which stops making progress at "/". A destination root of "/" therefore satisfied the loop condition forever, chmod'ing the container root while every peer rank waited in the barrier on the next line, until the job hit its wall clock. Break when dirname(path) == path -- the general "walked off the top" condition, not a special case for "/". 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 failed copy left a near-full file on node-local scratch -- one per attempt, since the name carries the job id, until a resubmit loop against a flaky fabric filled the NVMe and attempts started failing with ENOSPC instead of the original error. stage_file now unlinks it best-effort on the way out. Any rank that raises does the unlink; peers still hold it open, but the inode survives until they close and 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 watchdog is for. The device-memory budget check computes device_bytes as piece * (2*nodes + 1) -- the +1 being the send buffer -- but suggested a replacement --buffer-size computed by dividing by 2*nodes, dropping it. The suggestion was therefore itself over budget, so an operator following the error's own advice got the identical error back, byte for byte, and burned another multi-node allocation. It also had no way to say "no size works here": the max() floor handed back 2 MiB even when 2 MiB could not fit either. Tests for all three, each confirmed able to go red. test_stager.py drives _chmod_parents against a stub with a call ceiling, so a runaway reports as a failure rather than hanging the suite, and checks the budget advice by the property that matters -- parse the suggested size out of the message, rebuild with it, require acceptance, and require that 2 MiB more would have been refused so the advice is not needlessly small. test_pipeline.py injects a failure into _run_pipeline and asserts no temp file survives, recording from inside the failure that one existed, so it cannot pass by never creating one. Co-Authored-By: Claude Opus 5 --- REVIEW-FINDINGS.md | 25 +++- mlperf_common/fileio/datastage.py | 49 +++++++- tests/README.md | 3 +- tests/test_pipeline.py | 57 +++++++++ tests/test_stager.py | 203 ++++++++++++++++++++++++++++++ 5 files changed, 329 insertions(+), 8 deletions(-) create mode 100644 tests/test_stager.py diff --git a/REVIEW-FINDINGS.md b/REVIEW-FINDINGS.md index f61b5f3..bb165d2 100644 --- a/REVIEW-FINDINGS.md +++ b/REVIEW-FINDINGS.md @@ -235,12 +235,18 @@ Consequences: `(st_dev, st_ino)`), which also removes the duplicates. Documented in the `list_relative_files` docstring. -- [ ] **F7 · `mlperf_common/fileio/datastage.py:406` · `_chmod_parents` infinite loop on `/`** +- [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`) @@ -249,7 +255,7 @@ Consequences: exists to fix the umask on *directories we created*. Under `--container-remap-root` the process is root, so it always succeeds silently. -- [ ] **F9 · `mlperf_common/fileio/datastage.py:343` · `.datastage.tmp` files leak** +- [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 @@ -258,6 +264,19 @@ Consequences: 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` @@ -268,7 +287,7 @@ Consequences: of fabric traffic plus two world barriers per file — slower than the rsync path it replaces. -- [ ] **F11 · `mlperf_common/fileio/datastage.py:277` · memory-budget error suggests an over-budget value** +- [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 diff --git a/mlperf_common/fileio/datastage.py b/mlperf_common/fileio/datastage.py index 42753cd..c83d421 100644 --- a/mlperf_common/fileio/datastage.py +++ b/mlperf_common/fileio/datastage.py @@ -289,8 +289,21 @@ def __init__(self, args, topology): 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: - per_window = max(align_down(budget // (RECV_DEV_SLOTS * topology.node_count), - self.align), self.align) + # 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 " @@ -361,8 +374,29 @@ def _write_one(self, ready, fd, mview, base, stride, length, offset, block_size) return direct_io.pwrite(fd, mview[base:base + stride], length, offset, block_size) def stage_file(self, src, dst, size, mtime_ns): - topo = self.topo 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. @@ -430,7 +464,14 @@ def _chmod_parents(self, dst): os.chmod(path, self.args.chmod) except OSError: break - path = os.path.dirname(path) + 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 diff --git a/tests/README.md b/tests/README.md index 10aa483..5c62f3b 100644 --- a/tests/README.md +++ b/tests/README.md @@ -24,7 +24,8 @@ would go. No stubs involved. | `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 | +| `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 | diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 40023c9..6519f6e 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -143,6 +143,56 @@ def stage_once(root, nodes, buffer_size, size, tag): 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 @@ -167,6 +217,13 @@ def main(): 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) 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()) From 3e303edfdd8cb4beaea462cd1deed2a1c86ec75b Mon Sep 17 00:00:00 2001 From: Matt Frank Date: Fri, 31 Jul 2026 17:00:23 -0500 Subject: [PATCH 21/21] tests: add an end-to-end cluster self-test The suite in tests/ 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 N independent copies of the same CPU test. Real NCCL, real CUDA events, pinned memory against a real block size, and O_DIRECT have never been exercised by anything in this repo. cluster-selftest.sh is that missing half. From inside an allocation it builds a dataset with the shapes that have historically broken things -- sizes either side of the 2 MiB alignment boundary where the write padding and the closing ftruncate interact, a zero-length file, two large enough to need several rounds at any plausible node count, 64 small ones, and a symlink -- stages it with a real multi-node srun, and checks on every node that what landed on node-local storage matches what was read from shared storage. It compares two things rather than one, which rehearsing it locally turned out to matter for: fastmd5 emits one line per GB-chunk, so a zero-length file produces no lines at all, and a destination missing empty.bin entirely compares equal on checksums alone. A size-and-path listing catches that, plus truncations and unexpected extra files. Verified by injecting both failure modes into a stand-in copy: the missing empty file shows up in the listing and a single flipped byte in the checksums, independently. It also asserts the things recently fixed stay fixed on real hardware: no .datastage.tmp.* files survive a successful run, and everything is mode 0777. mtime is deliberately not compared. datastage does set it from the source, but timestamp granularity varies by filesystem and a false failure on a first hardware run would cost more than the check is worth. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 3 + tests/README.md | 22 ++++ tests/cluster-selftest.sh | 248 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 273 insertions(+) create mode 100755 tests/cluster-selftest.sh diff --git a/CLAUDE.md b/CLAUDE.md index 47c6a88..828400e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -21,6 +21,9 @@ Two mostly independent halves live here: 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 diff --git a/tests/README.md b/tests/README.md index 5c62f3b..330965d 100644 --- a/tests/README.md +++ b/tests/README.md @@ -69,3 +69,25 @@ If you change the pipeline, confirm a deliberately reintroduced bug still makes 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))"