diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..12f1491 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,42 @@ +name: docs + +# Build the Sphinx docs on every push and pull request; publish them to +# GitHub Pages from main. Needs Settings > Pages > Source: "GitHub Actions". +on: + push: + pull_request: + workflow_dispatch: + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: pages-${{ github.ref }} + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - run: pip install -r docs/requirements.txt + - run: sphinx-build -W --keep-going -b html docs docs/_build/html + - uses: actions/upload-pages-artifact@v3 + with: + path: docs/_build/html + + deploy: + if: github.ref == 'refs/heads/main' && github.event_name != 'pull_request' + needs: build + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - id: deployment + uses: actions/deploy-pages@v4 diff --git a/.gitignore b/.gitignore index 22e6d26..b85a2c8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,13 @@ # Cython build output build/ src/pyddstore.cpp +src/pyddstore/_core.cpp +src/pyddstore/_core.numpy-version # built module PyDDStore.egg-info/ pyddstore.cpython-*.so +_core.cpython-*.so +__pycache__/ +*.pyc +# Sphinx output +docs/_build/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 00f5caf..72e5a29 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,14 +28,14 @@ target_include_directories(demo PUBLIC ddstore) # option(BUILD_PYTHON_BINDINGS "Build Python bindings using cython" ON) # find_package(Python3 REQUIRED COMPONENTS Interpreter Development NumPy) # add_custom_command( -# OUTPUT ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore.cpp -# COMMAND ${Python3_EXECUTABLE} -m cython -3 --cplus ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore.pyx +# OUTPUT ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore/_core.cpp +# COMMAND ${Python3_EXECUTABLE} -m cython -3 --cplus ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore/_core.pyx # WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src # DEPENDS -# ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore.pyx +# ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore/_core.pyx # ) # add_library(pyddstore SHARED -# ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore.cpp +# ${CMAKE_CURRENT_SOURCE_DIR}/src/pyddstore/_core.cpp # ) # set_target_properties(pyddstore PROPERTIES PREFIX "") # set_target_properties(pyddstore PROPERTIES POSITION_INDEPENDENT_CODE ON) diff --git a/README.md b/README.md index 64fe9d4..e1a9922 100644 --- a/README.md +++ b/README.md @@ -1,329 +1,90 @@ # DDStore +DDStore logo + Efficient distributed data loading for distributed data-parallel (DDP) training. Each MPI rank holds a shard of the full dataset in memory. DDStore exposes a global index space so any rank can read any sample via one-sided remote memory access — either MPI RMA (default) or libfabric RDMA — without coordinator synchronization. +- **Batched reads**: `get_batch()` fetches a whole training batch in one call (one-sided RDMA reads in flight together, or an MPI collective for `method=0`). +- **GPUDirect RDMA**: data can live in, and be read straight into, GPU memory. +- **PyTorch integration**: `pyddstore.torch` turns any map-style dataset into a distributed one (`DistDataset`), reads samples made of several stored rows (`WindowedDataset`), and provides a thread-based `ThreadDataLoader` that is safe with MPI and GPU buffers. +- **Thread-safe** reads, a profiler for where read time goes, and a split mode (`method=2`) where a separate job reads data published by another. + DDStore architecture +**Documentation: ** (source in [docs/](docs/index.md)). + ## Prerequisites | Dependency | Notes | |---|---| | MPI (OpenMPI / MPICH) | `mpicc` and `mpicxx` must be on `PATH` | | libfabric | Required for the RDMA backends (`method=1` and `method=2`) | -| Python ≥ 3.6 | | +| Python ≥ 3.9 | | | NumPy, mpi4py, Cython | Python build dependencies | +| PyTorch (optional) | For `pyddstore.torch` and GPU buffers (CUDA or ROCm build) | ## Installation ```bash -# Install Python build dependencies pip install numpy mpi4py Cython - -# Build in-place (use with PYTHONPATH=$PWD:$PYTHONPATH) -CC=mpicc CXX=mpicxx python setup.py build_ext --inplace - -# Or install into the active virtual environment -CC=mpicc CXX=mpicxx pip install . - -# Or install in editable/development mode -CC=mpicc CXX=mpicxx pip install -e . - -# Or install directly from GitHub -CC=mpicc CXX=mpicxx pip install git+https://github.com/ORNL/DDStore.git +CC=mpicc CXX=mpicxx pip install . # or ".[torch]" to also pull PyTorch +CC=mpicc CXX=mpicxx pip install -e . # editable, for development ``` -## Quick Start +On Cray systems, building against the environment's own `mpi4py` (`--no-build-isolation`), and other build details: [Installation](docs/installation.md). -```python -import mpi4py -mpi4py.rc.thread_level = "serialized" -mpi4py.rc.threads = False +## Quick start +```python import numpy as np from mpi4py import MPI import pyddstore as dds comm = MPI.COMM_WORLD -rank = comm.Get_rank() - -# Each rank contributes its own shard -store = dds.PyDDStore(comm) # MPI RMA backend (default) -# store = dds.PyDDStore(comm, method=1) # libfabric RDMA backend +store = dds.PyDDStore(comm) # MPI RMA; method=1 for libfabric RDMA data = np.random.rand(1024, 64).astype(np.float32) -store.add("features", data) # collective — all ranks must call +store.add("features", data) # collective: each rank adds its shard -# Read any global sample index out = np.zeros((1, 64), dtype=np.float32) store.epoch_begin() -store.get("features", out, start=2048) # global index across all shards +store.get("features", out, start=2048) # any global row, from any rank store.epoch_end() - -store.free() -``` - -Run with: -```bash -mpirun -n 4 python my_script.py -``` - -## API Reference - -### `PyDDStore(comm_or_none=None, method=0, handshake_dir="", n_core=0, nic_map=None)` - -| Parameter | Type | Description | -|---|---|---| -| `comm_or_none` | `mpi4py.MPI.Comm` or `None` | MPI communicator covering all ranks. `None` only for a `method=2` extra member | -| `method` | `int` | `0` = MPI RMA (default), `1` = libfabric RDMA, `2` = file-based handshake (see [below](#file-based-handshake-method2)) | -| `handshake_dir` | `str` | Required for `method=2`: shared-filesystem directory used to exchange RDMA addresses | -| `n_core` | `int` | Required for a `method=2` extra member: number of core ranks that published data | -| `nic_map` | `str` or `None` | Optional, `method=1`/`2` only: a precomputed CPU→NIC map string (see [`DDSTORE_NIC_MAP`](#libfabric-rdma-method1) below) to use instead of the environment variable. Ignored if `FABRIC_IFACE` is already set | - -Four call shapes: - -```python -PyDDStore(comm) # method 0, MPI RMA -PyDDStore(comm, method=1) # method 1, libfabric RDMA -PyDDStore(comm, method=2, handshake_dir="/path") # method 2, core member (n_core == comm size) -PyDDStore(None, method=2, handshake_dir="/path", n_core=N) # method 2, extra member (no comm) -``` - -Note: grouping ranks into independent stores (the "sub-communicator" pattern below) is done by splitting `comm` yourself before constructing `PyDDStore` — there is no `ddstore_width` constructor parameter. `DistDataset` in [examples/vae/distdataset.py](examples/vae/distdataset.py) shows the pattern (`comm.Split()` then `PyDDStore(sub_comm)`). - ---- - -### `init(name, nrows, disp, itemsize=1)` - -Pre-allocate a named variable without providing data yet. Use `update()` to fill it in afterwards. **Collective**. - -| Parameter | Type | Description | -|---|---|---| -| `name` | `str` | Variable identifier | -| `nrows` | `int` | Number of rows in this rank's shard | -| `disp` | `int` | Number of elements per row | -| `itemsize` | `int` | Bytes per element (default `1`) | - ---- - -### `add(name, arr)` - -Register a NumPy array as a named variable. Each rank contributes its local shard; the global index space is the concatenation of all shards in rank order. **Collective** — all ranks in `comm` must call with the same `name`. - -| Parameter | Type | Description | -|---|---|---| -| `name` | `str` | Variable identifier | -| `arr` | `np.ndarray` | C-contiguous 2-D (or 1-D) array. Supported dtypes: `int32`, `int64`, `uint8`, `float32`, `float64`, `bool_` | - ---- - -### `update(name, arr, offset=0)` - -Overwrite a region of the local shard for a variable registered with `init()`. Local operation — does not require epoch or barrier. - -| Parameter | Type | Description | -|---|---|---| -| `name` | `str` | Variable identifier | -| `arr` | `np.ndarray` | Data to write | -| `offset` | `int` | Row offset within the local shard | - ---- - -### `get(name, arr, start=0)` - -Read `arr.shape[0]` consecutive rows starting at global index `start` into `arr`. The range must fall within a single rank's shard. Must be called inside an `epoch_begin` / `epoch_end` pair when using the MPI backend. - -| Parameter | Type | Description | -|---|---|---| -| `name` | `str` | Variable identifier | -| `arr` | `np.ndarray` | Pre-allocated, C-contiguous output buffer | -| `start` | `int` | Global row index | - ---- - -### `join(name)` - -`method=2` extra member only. Discovers a variable published by the core group by polling the handshake directory until the combined record file (`{name}.bin`) written by core rank 0 reaches its expected size (up to `DDSTORE_HANDSHAKE_TIMEOUT_S` seconds), then registers it for `get()`. - -| Parameter | Type | Description | -|---|---|---| -| `name` | `str` | Variable identifier, matching the `name` used in the core group's `add()` | - ---- - -### `info(name)` - -Returns `(total_rows, disp, itemsize)` for a variable that has been `add()`-ed or `join()`-ed. Useful on the extra side to size output buffers without hardcoding shapes. - ---- - -### `epoch_begin()` / `epoch_end()` - -Open and close an MPI RMA access epoch (calls `MPI_Win_fence`). **Collective**. Required around `get()` calls when using `method=0`. No-op for `method=1`. - ---- - -### `free()` - -Release all MPI windows and allocated memory. Safe to call after `MPI_Finalize`. - -## Backends - -### MPI RMA (`method=0`, default) - -Uses `MPI_Win_create` and `MPI_Get` for one-sided remote reads. Works on any MPI-capable cluster without additional hardware. `epoch_begin`/`epoch_end` are required to delimit access epochs. - -### libfabric RDMA (`method=1`) - -Uses `fi_read` for true RDMA transfers over high-speed interconnects (Infiniband/verbs, Cray GNI, Intel PSM2, Cray Slingshot). Lower latency than MPI RMA on supported hardware. `epoch_begin`/`epoch_end` are no-ops with this backend. - -**`DDSTORE_FABRIC`** selects which libfabric provider to open, for `method=1`/`2`: - -- `hsn` (default, unset) — Frontier: opens the `tcp;ofi_rxm` domain over Cray Slingshot. -- `cxi` — Perlmutter: opens the native `cxi` domain over Cray Slingshot. - -The two are independent code paths (not runtime auto-detection), so set this explicitly per system rather than relying on a guess: - -```bash -export DDSTORE_FABRIC=hsn # Frontier (default; usually not needed) -export DDSTORE_FABRIC=cxi # Perlmutter -``` - -`PyDDStore` picks the network interface (`FABRIC_IFACE`) automatically for `method=1`/`2`, based on each rank's real CPU affinity (`os.sched_getaffinity`) — no changes needed in your code: - -- **`DDSTORE_NIC_MAP`** — a precomputed CPU→NIC map, used directly if set (no NIC discovery at construction time). Generate it once from a context with reliable NIC visibility, e.g. an `sbatch` batch step's own shell (not a nested `srun` task — NIC/PCI discovery has been observed to fail there), and export it before launching ranks so every one inherits it: - ```bash - export DDSTORE_NIC_MAP=$(python3 -m cpu_nic_map --env) - srun ... python train.py - ``` -- If `DDSTORE_NIC_MAP` isn't set, each rank falls back to a live `hwloc-calc`/`lstopo` query against its own CPU affinity (`cpu_nic_map.allocated_nics()`, also runnable standalone as `python3 cpu_nic_map.py --allocated`) to find the nearest NIC. -- Set `FABRIC_IFACE` explicitly to override both and force a specific interface, e.g. when the automatic selection picks the wrong one: - ```bash - export FABRIC_IFACE=hsn0 # e.g. Cray Slingshot - ``` -- Or skip the environment entirely and pass a map straight to the constructor: `PyDDStore(comm, method=1, nic_map="hsn0=0-15,64-79;hsn1=...")`. - -### File-based handshake (`method=2`) - -Splits the dataset-holding job from the training job entirely: a **core** group loads and publishes data, and a separate **extra** group reads it over RDMA (`fi_read`, same transport as `method=1`) — the two are independent MPI jobs (e.g. two separate `srun`/`mpirun` launches, possibly on different node allocations) that never share a communicator. They rendezvous only through record files written to a shared-filesystem directory (must be visible to all nodes, e.g. Lustre): - -- **Core member** — has an MPI communicator, publishes with `add()`/`init()`. Core ranks exchange records via `MPI_Allgather`, and rank 0 writes the combined set to a single `{name}.bin` file (fabric address, MR key, base pointer, row count, dtype per rank) into `handshake_dir`. -- **Extra member** — no MPI communicator; constructed with `comm_or_none=None` and an explicit `n_core`. Calls `join(name)` to poll for and read all `n_core` core-rank records, then `get()` works exactly as on the core side, reading directly from core-rank memory over RDMA. - -```python -# core side — one MPI job -store = dds.PyDDStore(comm, method=2, handshake_dir="/lustre/.../ddstore_hs") -store.add("x", data) -... # wait for the extra side to finish (e.g. a sentinel file) store.free() - -# extra side — a separate MPI job, no comm needed -store = dds.PyDDStore(None, method=2, handshake_dir="/lustre/.../ddstore_hs", n_core=4) -store.join("x") -out = np.zeros((1, ncols), dtype=np.float32) -store.get("x", out, start=global_idx) -store.free() -``` - -Environment variables: - -| Variable | Default | Description | -|---|---|---| -| `DDSTORE_HANDSHAKE_DIR` | `./ddstore_hs` | Shared directory for handshake record files | -| `DDSTORE_HANDSHAKE_TIMEOUT_S` | `300` | Seconds to poll for core records / a join before raising a timeout | -| `DDSTORE_NIC_MAP` | unset | CPU→NIC map for `FABRIC_IFACE` auto-selection — see [libfabric RDMA](#libfabric-rdma-method1) above | -| `DDSTORE_FABRIC` | `hsn` | `hsn` (Frontier) or `cxi` (Perlmutter) — see [libfabric RDMA](#libfabric-rdma-method1) above | - -See [test/test_method2_core.py](test/test_method2_core.py) / [test/test_method2_extra.py](test/test_method2_extra.py) for a minimal runnable pair, and [examples/vae/vae_core_server.py](examples/vae/vae_core_server.py) / [examples/vae/vae_extra_train.py](examples/vae/vae_extra_train.py) for a full DDP training example using this split. - -`ddstore_width` grouping (below) is not currently supported with `method=2` — every core rank in `comm` is treated as one group. - -## Partitioned / Sub-communicator Usage - -`PyDDStore` itself always spans the full communicator you pass it — there is no built-in "ranks per group" option. To run several independent stores side by side (e.g. one per node), split `comm` yourself before constructing `PyDDStore`, giving each group its own sub-communicator. Each group then holds a full replica of the dataset, partitioned across its own members. - -**Example — 16 ranks split into groups of 4:** -``` -ranks 0– 3 → DDStore group 0 -ranks 4– 7 → DDStore group 1 -ranks 8–11 → DDStore group 2 -ranks 12–15 → DDStore group 3 ``` -This is useful when you want one store per node (e.g. 4 GPUs per node), limiting cross-node RDMA traffic to the dataset replication step at startup rather than every sample fetch. +With PyTorch: ```python -width = 4 # ranks per group, e.g. GPUs per node -sub_comm = comm.Split(rank // width, rank) -store = dds.PyDDStore(sub_comm) # one independent store per group -``` - -`DistDataset` in [examples/vae/distdataset.py](examples/vae/distdataset.py) wraps exactly this pattern behind a `ddstore_width` constructor argument — pass `ddstore_width=None` (default) for a single store across all ranks in `comm`, or an integer to split into groups of that size. - -## PyTorch Dataset Integration - -See [examples/vae/distdataset.py](examples/vae/distdataset.py) for a `torch.utils.data.Dataset` wrapper and [examples/vae/vae-ddp.py](examples/vae/vae-ddp.py) for a full DDP training example. - -```bash -mpirun -n 4 python examples/vae/vae-ddp.py -``` - -## Testing - -### Unit tests (pytest) - -Install test dependencies: - -```bash -pip install pytest pytest-mpi -``` - -**Single-rank** — no cluster required, covers all dtypes, `add`/`get`/`init`/`update`, and error cases: - -```bash -mpirun -n 1 python -m pytest test/test_single.py -v -``` - -**Multi-rank** — verifies remote reads across all rank pairs and sub-communicator grouping: - -```bash -mpirun -n 4 python -m pytest test/test_multirank.py -v -``` - -| Test file | Min ranks | What is tested | -|---|---|---| -| `test/test_single.py` | 1 | All dtypes, `add`/`get`, `init`/`update`/`get`, error handling, double `free()` | -| `test/test_multirank.py` | 2 (4 recommended) | Remote reads, shard boundaries, multiple variables, `ddstore_width` grouping | - -### Integration scripts - -```bash -# Basic functional test (MPI RMA) -mpirun -n 4 python examples/scripts/demo.py +import torch # import torch before MPI starts +from mpi4py import MPI +from pyddstore.torch import DistDataset, ThreadDataLoader -# Integration test with PyTorch DDP -mpirun -n 4 python examples/scripts/test.py +trainset = DistDataset(my_dataset, "train", MPI.COMM_WORLD) # each rank loads only its share +sampler = torch.utils.data.distributed.DistributedSampler(trainset) +loader = ThreadDataLoader(trainset, batch_size=128, sampler=sampler, num_workers=1) +for x, y in loader: + ... ``` -Optional arguments for `examples/scripts/demo.py` and `examples/scripts/test.py`: +Run with `mpirun -n 4 python my_script.py` (or `srun`). -| Flag | Default | Description | -|---|---|---| -| `--num` | `1048576` | Rows per rank | -| `--dim` | `64` | Elements per row | -| `--nbatch` | `32` | Number of random reads | +## Documentation -### Method 2 (file-based handshake) +| | | +|---|---| +| Getting started | [Installation](docs/installation.md), [Quick start](docs/quickstart.md) | +| User guide | [Backends](docs/backends.md) (MPI RMA, libfabric, file-based handshake, partitioned stores), [GPUDirect RDMA](docs/gpudirect.md), [PyTorch integration](docs/pytorch.md), [HPC systems](docs/hpc.md) (Slurm, Slingshot, Frontier, Perlmutter), [Performance](docs/performance.md), [Concurrency](docs/concurrency.md) | +| Reference | [`PyDDStore`](docs/api-pyddstore.md), [`pyddstore.torch`](docs/api-torch.md), [Environment variables](docs/environment.md) | +| More | [Testing](docs/testing.md), [Measurements](docs/results.md) | -Two separate launches sharing a handshake directory on a shared filesystem — not a single `mpirun`, since core and extra are independent jobs: +To build the documentation locally: ```bash -# Terminal 1 — core (data-holding) side -mpirun -n 4 python test/test_method2_core.py /path/to/shared/ddstore_hs - -# Terminal 2 — extra (reader) side, after or while the core side is running -python test/test_method2_extra.py /path/to/shared/ddstore_hs 4 +pip install -r docs/requirements.txt +sphinx-build -b html docs docs/_build/html ``` ## Citation @@ -340,6 +101,17 @@ If you use DDStore in your research, please cite: } ``` +```bibtex +@inproceedings{bae2024mdloader, + title={MDLoader: A Hybrid Model-Driven Data Loader for Distributed Graph Neural Network Training}, + author={Bae, Jonghyun and Choi, Jong Youl and Lupo Pasini, Massimiliano and Mehta, Kshitij and Zhang, Pei and Ibrahim, Khaled}, + booktitle={SC24-W: Workshops of the International Conference for High Performance Computing, Networking, Storage and Analysis}, + year={2024}, + month={nov}, + doi={10.1109/SCW63240.2024.00145} +} +``` + ## License See [LICENSE](LICENSE). diff --git a/docs/api-pyddstore.md b/docs/api-pyddstore.md new file mode 100644 index 0000000..52d2e5b --- /dev/null +++ b/docs/api-pyddstore.md @@ -0,0 +1,135 @@ +# `PyDDStore` reference + +## `PyDDStore(comm_or_none=None, method=0, handshake_dir="", n_core=0, nic_map=None)` + +| Parameter | Type | Description | +|---|---|---| +| `comm_or_none` | `mpi4py.MPI.Comm` or `None` | MPI communicator covering all ranks. `None` only for a `method=2` extra member | +| `method` | `int` | `0` = MPI RMA (default), `1` = libfabric RDMA, `2` = file-based handshake (see [below](backends.md#file-based-handshake-method2)) | +| `handshake_dir` | `str` | Required for `method=2`: shared-filesystem directory used to exchange RDMA addresses | +| `n_core` | `int` | Required for a `method=2` extra member: number of core ranks that published data | +| `nic_map` | `str` or `None` | Optional, `method=1`/`2` only: a precomputed CPU→NIC map string (see [`DDSTORE_NIC_MAP`](backends.md#libfabric-rdma-method1) below) to use instead of the environment variable. Ignored if `FABRIC_IFACE` is already set | + +Four call shapes: + +```python +PyDDStore(comm) # method 0, MPI RMA +PyDDStore(comm, method=1) # method 1, libfabric RDMA +PyDDStore(comm, method=2, handshake_dir="/path") # method 2, core member (n_core == comm size) +PyDDStore(None, method=2, handshake_dir="/path", n_core=N) # method 2, extra member (no comm) +``` + +Note: grouping ranks into independent stores (the "sub-communicator" pattern below) is done by splitting `comm` yourself before constructing `PyDDStore` — there is no `ddstore_width` constructor parameter. `pyddstore.torch.DistDataset` does this for you (`ddstore_width`) and shows the pattern (`comm.Split()` then `PyDDStore(sub_comm)`). + +--- + +## `init(name, nrows, disp, itemsize=1)` + +Pre-allocate a named variable without providing data yet. Use `update()` to fill it in afterwards. **Collective**. + +| Parameter | Type | Description | +|---|---|---| +| `name` | `str` | Variable identifier | +| `nrows` | `int` | Number of rows in this rank's shard | +| `disp` | `int` | Number of elements per row | +| `itemsize` | `int` | Bytes per element (default `1`) | + +--- + +## `add(name, arr)` + +Register a NumPy array as a named variable. Each rank contributes its local shard; the global index space is the concatenation of all shards in rank order. **Collective** — all ranks in `comm` must call with the same `name`. + +| Parameter | Type | Description | +|---|---|---| +| `name` | `str` | Variable identifier | +| `arr` | `np.ndarray` or `torch.Tensor` | C-contiguous 2-D (or 1-D) array/tensor. Supported dtypes: `int32`, `int64`, `uint8`, `float32`, `float64`, `bool_`/`bool`. A CUDA/HIP tensor registers GPU memory directly — see [GPUDirect RDMA](gpudirect.md) below | + +--- + +## `update(name, arr, offset)` + +Overwrite a region of the local shard for a variable registered with `init()`. Local operation — does not require epoch or barrier. + +| Parameter | Type | Description | +|---|---|---| +| `name` | `str` | Variable identifier | +| `arr` | `np.ndarray` | Data to write | +| `offset` | `int` | Row offset within the local shard | + +--- + +## `get(name, arr, start=0)` + +Read `arr.shape[0]` consecutive rows starting at global index `start` into `arr`. The range must fall within a single rank's shard. Must be called inside an `epoch_begin` / `epoch_end` pair when using the MPI backend. + +| Parameter | Type | Description | +|---|---|---| +| `name` | `str` | Variable identifier | +| `arr` | `np.ndarray` or `torch.Tensor` | Pre-allocated, C-contiguous output buffer. A CUDA/HIP tensor writes the RDMA transfer directly into GPU memory — see [GPUDirect RDMA](gpudirect.md) below | +| `start` | `int` | Global row index | + +--- + +## `get_batch(name, arr, indices)` + +Read rows `indices` (global row ids; any order, any ranks, repeats allowed) into `arr`: row `i` of `arr` receives row `indices[i]`, so `arr.shape[0]` must equal `len(indices)` (else `ValueError`). Same buffer rules as `get()` (NumPy array or CUDA/HIP tensor). Every index is checked before anything is read; an out-of-range one raises `IndexError` and leaves the store usable. + +For `method=1`/`2` the whole batch is one call: one lock acquisition, at most one memory registration (none into a [registered](api-pyddstore.md#register_recvname-arr--unregister_recvname-arr) buffer), and (GPU destination) one device sync, with all of the batch's `fi_read`s posted before any is waited for, so the reads overlap on the network. It is one `fi_read` per row, or several for a row longer than `DDSTORE_MAX_READ_BYTES` (default 1 GiB). + +For `method=0`, `get_batch()` is **collective**, after the collective module of [MDLoader](https://ieeexplore.ieee.org/abstract/document/10820758) (see [Citation](citation.md)): every rank all-gathers all ranks' indices (`MPI_Allgatherv`), packs the rows it owns for each requester, and one `MPI_Alltoallv` delivers them, on a private duplicate of the store's communicator, in rounds of at most `DDSTORE_ALLTOALL_MAX_BYTES` (default 2 MiB) received per rank so large rows don't turn into one huge exchange. So every rank must call it for the variable the same number of times, in the same order, from one thread at a time; the number of indices may differ per rank (including 0). Indices are checked on the gathered list, so a bad index raises on every rank together. `DistributedSampler` gives every rank the same number of batches, and `vae-ddp.py` allows no worker threads with `method=0`, so the data loaders meet this automatically. + +```python +idx = np.array([2048, 7, 4096, 7]) +out = np.zeros((len(idx), 64), dtype=np.float32) +store.get_batch("features", out, idx) +``` + +`DistDataset`/`DistDatasetReader` use it by default through `__getitems__`, which PyTorch's `DataLoader` (and `ThreadDataLoader`) calls with a whole batch's indices, so the VAE examples and job scripts read in batches with no extra flag. Set `DDSTORE_BATCH_GET=0` to fall back to one `get()` per sample. + +--- + +## `register_recv(name, arr)` / `unregister_recv(name, arr)` + +Register `arr` (a C-contiguous NumPy array or CUDA/HIP tensor) once as a destination for `get()`/`get_batch()` of `name`. Reads into `arr` or any slice of it then skip memory registration, which otherwise happens whenever the destination isn't the buffer registered by the previous read. For large rows, registration can cost more than the transfer. Use it for buffers you reuse, such as a pool per loader thread: several can be registered per variable and none is evicted. The store holds a reference to `arr` until `unregister_recv()` or `free()`. No-op for `method=0`. + +```python +pool = np.empty((batch_size, ncols), dtype=np.float32) +store.register_recv("features", pool) +for idx in batches: + store.get_batch("features", pool[: len(idx)], idx) # no registration +``` + +--- + +## `get_profile(name)` + +With `DDSTORE_PROFILE=1` set before the process starts: timing of `get()`/`get_batch()` for `name` (`method=1`/`2`), in seconds unless noted. C++ counters for this variable: `calls` (get + get_batch), `rows`, `lock_wait`, `mr` (receive-buffer registration, including cache checks), `mr_miss` (registrations, a count), `read` (posting `fi_read`), `cq` (waiting for completions). Python counters for the whole store: `py_gets`, `py_get` (whole calls), `py_sync` (`torch.cuda.synchronize()` on the GPU path). All zero for `method=0` or without profiling. See [Performance](performance.md). + +--- + +## `join(name)` + +`method=2` extra member only. Discovers a variable published by the core group by polling the handshake directory until the combined record file (`{name}.bin`) written by core rank 0 reaches its expected size (up to `DDSTORE_HANDSHAKE_TIMEOUT_S` seconds), then registers it for `get()`. + +| Parameter | Type | Description | +|---|---|---| +| `name` | `str` | Variable identifier, matching the `name` used in the core group's `add()` | + +--- + +## `info(name)` + +Returns `(total_rows, disp, itemsize)` for a variable that has been `add()`-ed or `join()`-ed. Useful on the extra side to size output buffers without hardcoding shapes. + +--- + +## `epoch_begin()` / `epoch_end()` + +Open and close an MPI RMA access epoch (calls `MPI_Win_fence`). **Collective**. Required around `get()` calls when using `method=0`. No-op for `method=1`/`2`. + +--- + +## `free()` + +Release every variable's MPI window (`method=0`) or libfabric endpoints and memory registrations, including [`register_recv()`](api-pyddstore.md#register_recvname-arr--unregister_recvname-arr) buffers (`method=1`/`2`), then the host buffer DDStore allocated for it in `add()`/`init()` (a GPU tensor passed to `add()` is the caller's and is not freed). Safe to call more than once. After `MPI_Finalize` the MPI window and buffer can no longer be released and are skipped. diff --git a/docs/api-torch.md b/docs/api-torch.md new file mode 100644 index 0000000..dbb9a76 --- /dev/null +++ b/docs/api-torch.md @@ -0,0 +1,22 @@ +# `pyddstore.torch` reference + +Generated from the docstrings. For how the pieces fit together, see +[PyTorch integration](pytorch.md). + +```{eval-rst} +.. automodule:: pyddstore.torch + :no-members: + +.. autoclass:: pyddstore.torch.DistDataset(source, name, comm=None, ddstore_width=None, device=None, add_device=None, method=None, handshake_dir=None, chunk_size=None, encode=None, decode=None, fields=None) + :members: read_rows, alloc, release, shapes, dtypes, __getitems__ + +.. autoclass:: pyddstore.torch.DistDatasetReader(name, handshake_dir=None, n_core=None, device=None, decode=None) + :members: read_rows, alloc, release, shapes, dtypes, __getitems__ + +.. autoclass:: pyddstore.torch.WindowedDataset(ds, window, stride=1, dilation=1, starts=None, fields=None) + +.. autofunction:: pyddstore.torch.row_of + +.. autoclass:: pyddstore.torch.ThreadDataLoader(dataset, reuse_buffers=False, collate_copies=False, **DataLoader_kwargs) + :members: close +``` diff --git a/docs/backends.md b/docs/backends.md new file mode 100644 index 0000000..97d8d27 --- /dev/null +++ b/docs/backends.md @@ -0,0 +1,76 @@ +# Backends + +## MPI RMA (`method=0`, default) + +Uses `MPI_Win_create` and `MPI_Get` for one-sided remote reads. Works on any MPI-capable cluster without additional hardware. `epoch_begin`/`epoch_end` are required to delimit access epochs. + +## libfabric RDMA (`method=1`) + +Uses `fi_read` for true RDMA transfers over high-speed interconnects (Infiniband/verbs, Cray GNI, Intel PSM2, Cray Slingshot). Lower latency than MPI RMA on supported hardware. `epoch_begin`/`epoch_end` are no-ops with this backend. + +**`DDSTORE_FABRIC`** selects which libfabric provider to open, for `method=1`/`2`: + +- `hsn` (default, unset) — opens the `tcp;ofi_rxm` domain over Cray Slingshot (Frontier). +- `cxi` — opens the native `cxi` domain over Cray Slingshot (Frontier and Perlmutter; Perlmutter is CXI-only). Required for [GPUDirect RDMA](gpudirect.md), and the default in the Frontier job scripts. + +The two are independent code paths (not runtime auto-detection), so set this explicitly per system rather than relying on a guess: + +```bash +export DDSTORE_FABRIC=hsn # tcp;ofi_rxm (default) +export DDSTORE_FABRIC=cxi # native CXI: Perlmutter, or Frontier with GPUDirect +``` + +`PyDDStore` picks the network interface (`FABRIC_IFACE`) automatically for `method=1`/`2`, based on each rank's real CPU affinity (`os.sched_getaffinity`) — no changes needed in your code: + +- **`DDSTORE_NIC_MAP`** — a precomputed CPU→NIC map, used directly if set (no NIC discovery at construction time). Generate it once from a context with reliable NIC visibility, e.g. an `sbatch` batch step's own shell (not a nested `srun` task — NIC/PCI discovery has been observed to fail there), and export it before launching ranks so every one inherits it: + ```bash + export DDSTORE_NIC_MAP=$(python3 -m cpu_nic_map --env) + srun ... python train.py + ``` +- If `DDSTORE_NIC_MAP` isn't set, each rank falls back to a live `hwloc-calc`/`lstopo` query against its own CPU affinity (`cpu_nic_map.allocated_nics()`, also runnable standalone as `python3 cpu_nic_map.py --allocated`) to find the nearest NIC. +- Set `FABRIC_IFACE` explicitly to override both and force a specific interface, e.g. when the automatic selection picks the wrong one: + ```bash + export FABRIC_IFACE=hsn0 # e.g. Cray Slingshot + ``` +- Or skip the environment entirely and pass a map straight to the constructor: `PyDDStore(comm, method=1, nic_map="hsn0=0-15,64-79;hsn1=...")`. + +## File-based handshake (`method=2`) + +Splits the dataset-holding job from the training job entirely: a **core** group loads and publishes data, and a separate **extra** group reads it over RDMA (`fi_read`, same transport as `method=1`) — the two are independent MPI jobs (e.g. two separate `srun`/`mpirun` launches, possibly on different node allocations) that never share a communicator. They rendezvous only through record files written to a shared-filesystem directory (must be visible to all nodes, e.g. Lustre): + +- **Core member** — has an MPI communicator, publishes with `add()`/`init()`. Core ranks exchange records via `MPI_Allgather`, and rank 0 writes the combined set to a single `{name}.bin` file (fabric address, MR key, base pointer, row count, dtype per rank) into `handshake_dir`. +- **Extra member** — no MPI communicator; constructed with `comm_or_none=None` and an explicit `n_core`. Calls `join(name)` to poll for and read all `n_core` core-rank records, then `get()` works exactly as on the core side, reading directly from core-rank memory over RDMA. + +```python +# core side — one MPI job +store = dds.PyDDStore(comm, method=2, handshake_dir="/lustre/.../ddstore_hs") +store.add("x", data) +... # wait for the extra side to finish (e.g. a sentinel file) +store.free() + +# extra side — a separate MPI job, no comm needed +store = dds.PyDDStore(None, method=2, handshake_dir="/lustre/.../ddstore_hs", n_core=4) +store.join("x") +out = np.zeros((1, ncols), dtype=np.float32) +store.get("x", out, start=global_idx) +store.free() +``` + +Environment variables: `DDSTORE_HANDSHAKE_DIR`, `DDSTORE_HANDSHAKE_TIMEOUT_S`, `DDSTORE_FABRIC` and `DDSTORE_NIC_MAP` — see [Environment variables](environment.md). + +See [test/test_method2_core.py](https://github.com/ORNL/DDStore/blob/main/test/test_method2_core.py) / [test/test_method2_extra.py](https://github.com/ORNL/DDStore/blob/main/test/test_method2_extra.py) for a minimal runnable pair, and [examples/vae/vae_core_server.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae_core_server.py) / [examples/vae/vae_extra_train.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae_extra_train.py) for a full DDP training example using this split. + +`ddstore_width` grouping (below) is not currently supported with `method=2` — every core rank in `comm` is treated as one group. + + +## Partitioned / Sub-communicator Usage + +`PyDDStore` always spans the whole communicator you pass it. To run several independent stores side by side (e.g. one per node), split `comm` first; each group then holds a full replica of the dataset, partitioned across its own members: + +```python +width = 4 # ranks per group, e.g. GPUs per node +sub_comm = comm.Split(rank // width, rank) +store = dds.PyDDStore(sub_comm) # one independent store per group +``` + +This keeps sample fetches inside a node at the cost of replicating the data per group. `DistDataset` exposes it as `ddstore_width` (`None` = one store across all ranks). Not supported with `method=2`. diff --git a/docs/citation.md b/docs/citation.md new file mode 100644 index 0000000..7fcf368 --- /dev/null +++ b/docs/citation.md @@ -0,0 +1,28 @@ +# Citation and license + +If you use DDStore in your research, please cite: + +```bibtex +@inproceedings{choi2023ddstore, + title={DDStore: Distributed data store for scalable training of graph neural networks on large atomistic modeling datasets}, + author={Choi, Jong Youl and Lupo Pasini, Massimiliano and Zhang, Pei and Mehta, Kshitij and Liu, Frank and Bae, Jonghyun and Ibrahim, Khaled}, + booktitle={Proceedings of the SC'23 Workshops of the International Conference on High Performance Computing, Network, Storage, and Analysis}, + pages={941--950}, + year={2023} +} +``` + +```bibtex +@inproceedings{bae2024mdloader, + title={MDLoader: A Hybrid Model-Driven Data Loader for Distributed Graph Neural Network Training}, + author={Bae, Jonghyun and Choi, Jong Youl and Lupo Pasini, Massimiliano and Mehta, Kshitij and Zhang, Pei and Ibrahim, Khaled}, + booktitle={SC24-W: Workshops of the International Conference for High Performance Computing, Networking, Storage and Analysis}, + year={2024}, + month={nov}, + doi={10.1109/SCW63240.2024.00145} +} +``` + +## License + +See [LICENSE](https://github.com/ORNL/DDStore/blob/main/LICENSE). diff --git a/docs/concurrency.md b/docs/concurrency.md new file mode 100644 index 0000000..5d3847b --- /dev/null +++ b/docs/concurrency.md @@ -0,0 +1,5 @@ +# Concurrency + +- `get()` and `get_batch()` are thread-safe. For `method=1`/`2` a per-variable lock in `DDStore::get()` serializes calls on one variable (it guards shared receive state; without it concurrent calls crashed). Both release the GIL during the transfer. +- `method=0` `get_batch()` is **collective**: every rank calls it the same number of times, in the same order, from one thread. `vae-ddp.py` therefore allows no worker threads with `method=0`. +- Only the main thread calls MPI (setup, `epoch_begin`/`epoch_end`, `method=0` reads); mpi4py's default `MPI_THREAD_MULTIPLE` is fine, `FUNNELED` is the minimum. If you call MPI from your own worker threads, keep `MULTIPLE`. diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..2a9f62f --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,39 @@ +"""Sphinx configuration for the DDStore documentation (MyST Markdown).""" + +import os +import sys + +# Document pyddstore from the source tree; its compiled core, MPI and torch +# are mocked, so the docs build needs neither a compiler, MPI nor libfabric. +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src"))) + +project = "DDStore" +author = "Oak Ridge National Laboratory" +copyright = "UT-Battelle, LLC" + +extensions = [ + "myst_parser", + "sphinx.ext.autodoc", + "sphinx.ext.napoleon", +] +source_suffix = {".md": "markdown"} +exclude_patterns = ["_build"] + +# GitHub-style anchors on headings, so links like page.md#get_batchname-arr-indices work +myst_heading_anchors = 3 +myst_enable_extensions = ["colon_fence"] + +autodoc_mock_imports = ["mpi4py", "torch", "pyddstore._core"] +autodoc_member_order = "bysource" +autodoc_typehints = "none" +napoleon_google_docstring = True +napoleon_numpy_docstring = False + +html_theme = "furo" +html_title = "DDStore" +html_logo = "../images/DDStore-logo.png" +html_theme_options = { + "source_repository": "https://github.com/ORNL/DDStore/", + "source_branch": "main", + "source_directory": "docs/", +} diff --git a/docs/environment.md b/docs/environment.md new file mode 100644 index 0000000..5ac8e4a --- /dev/null +++ b/docs/environment.md @@ -0,0 +1,39 @@ +# Environment variables + +**Read by DDStore itself** (the C++ library, `pyddstore`, `cpu_nic_map`): + +| Variable | Default | Effect | +|---|---|---| +| `DDSTORE_FABRIC` | `hsn` | libfabric provider for `method=1`/`2`: `hsn` (`tcp;ofi_rxm`) or `cxi` (native Slingshot; required for [GPUDirect RDMA](gpudirect.md)). See [libfabric RDMA](backends.md#libfabric-rdma-method1). | +| `FABRIC_IFACE` | auto | Network interface (libfabric domain, e.g. `cxi0`, `hsn0`) for `method=1`/`2`. Set it to force one; otherwise picked from the rank's CPU affinity. | +| `DDSTORE_NIC_MAP` | unset | Precomputed CPU→NIC map used for that automatic pick instead of a live hwloc query (`python3 -m cpu_nic_map --env`). The constructor's `nic_map=` argument takes priority. | +| `DDSTORE_HANDSHAKE_DIR` | `./ddstore_hs` | `method=2` handshake directory when none is given (C++ API; `PyDDStore` requires `handshake_dir`, and the examples fill it from this variable). Must be on a shared filesystem. | +| `DDSTORE_HANDSHAKE_TIMEOUT_S` | `300` | Seconds a `method=2` extra member's `join()` polls for the core group's record file. | +| `DDSTORE_PROFILE` | off | `1` turns on `get()`/`get_batch()` timing counters, read with `get_profile(name)`. See [Performance](performance.md). | +| `DDSTORE_MAX_READ_BYTES` | `1073741824` (1 GiB) | `method=1`/`2`: largest single `fi_read`; longer rows are read in pieces (on Perlmutter's `cxi` one 5 GB read fails with `EMSGSIZE`, 2.5 GB works, and the provider doesn't report the limit). Lowered to the endpoint's `max_msg_size` when the provider reports one. | +| `DDSTORE_ALLTOALL_MAX_BYTES` | `2097152` (2 MiB) | `method=0` `get_batch()`: bytes each rank receives per exchange round. Must be equal on all ranks. | + +**Read by `pyddstore.torch`** (defaults for arguments not given): + +| Variable | Default | Effect | +|---|---|---| +| `DDSTORE_METHOD` | `0` | `DistDataset`'s backend when `method=` isn't passed: `0` MPI RMA, `1` libfabric, `2` file-based handshake. (`PyDDStore` itself takes `method=` only.) | +| `DDSTORE_BATCH_GET` | `1` | `DistDataset.__getitems__` reads a whole batch with one `get_batch()` per field; `0` reads one sample at a time. | +| `DDSTORE_HANDSHAKE_DIR`, `DDSTORE_HANDSHAKE_TIMEOUT_S` | `./ddstore_hs`, `300` | `method=2` directory, and how long `DistDatasetReader` waits for the core group to publish. | +| `DDSTORE_N_CORE` | unset | `DistDatasetReader`'s number of core ranks when `n_core=` isn't passed (the examples default it to 4). | +| `DDSTORE_AFFINITY_WIDTH` / `DDSTORE_AFFINITY_OFFSET` | `0` / `0` | `ThreadDataLoader`: pin worker thread *i* to CPUs `[offset + i·width, offset + (i+1)·width)` of the process's affinity; width `0` = no pinning. | + +**Read by the examples** (`examples/vae/`, `examples/scripts/`, job scripts): + +| Variable | Default | Effect | +|---|---|---| +| `DDSTORE_BACKEND` | auto | `torch.distributed` backend for the examples' DDP setup (`nccl`, `gloo`, `xccl`). | +| `VAE_PROFILE` | off | `1`: `vae-ddp.py` prints per-epoch fetch vs compute time. | +| `MASTER_PORT` | `2345` | DDP rendezvous port; the core/extra job script gives each step its own. | + +**System settings that matter on Frontier:** + +| Variable | Effect | +|---|---| +| `SLINGSHOT_VNIS` | Set by Slurm per step. With `--network=job_vni`, keep only the last (job-wide) entry before starting Python so separate `srun` steps can reach each other — see [Multiple `srun` steps](hpc.md#multiple-srun-steps-in-one-job-method2-cxi). | +| `GPU_MAX_HW_QUEUES` | ROCm hardware queues per GPU per process (default 4); raise it if data-loading threads use their own streams — see [HIP streams](results.md#hip-streams-and-hardware-queues-frontier-rocm-72). | diff --git a/docs/gpudirect.md b/docs/gpudirect.md new file mode 100644 index 0000000..bd4c650 --- /dev/null +++ b/docs/gpudirect.md @@ -0,0 +1,19 @@ +# GPUDirect RDMA + +`add()`, `get()` and `get_batch()` accept a CUDA/HIP `torch.Tensor` in place of a NumPy array, so RDMA reads from or writes directly into GPU memory, with no `.cpu()`/`.to(device)` copy. Requires `method=1` or `2`, **`DDSTORE_FABRIC=cxi`**, and a CUDA- or ROCm-enabled PyTorch. A GPU tensor with `DDSTORE_FABRIC=hsn` (the default) or `method=0` raises a clear error instead of silently copying through the host. + +```python +import torch +data = torch.rand(1024, 64, dtype=torch.float32, device="cuda") +store.add("features", data) # GPU source, no host copy + +out = torch.empty((1, 64), dtype=torch.float32, device="cuda") +store.get("features", out, start=2048) # GPU destination, no host copy +``` + +- **`add()` with a GPU tensor registers your tensor's own memory; no copy is made.** Keep it alive and unmodified until `free()`. `PyDDStore` holds a reference as a safety net, and adding the same name again with a GPU tensor is rejected. (With NumPy, `add()` copies and the array can be reused right away.) +- **The device is synchronized before each GPU transfer** (`torch.cuda.synchronize()`, once per `get()` / `get_batch()` call): the NIC writes outside PyTorch's stream ordering, and without the sync training hit GPU memory faults. Prefer `get_batch()` on the GPU path so this costs one sync per batch, not per sample. +- `init()`/`update()` stay host-only. +- Whether GPU destinations are faster than host ones depends on the machine: on Frontier they win from ~12.5 KB rows up, on Perlmutter host destinations win at every size ([results](results.md#bench_getpy-µs-per-row)). + +Examples: [test/test_gpu_rdma.py](https://github.com/ORNL/DDStore/blob/main/test/test_gpu_rdma.py), and `--gpu-dest`/`--gpu-source` on [vae-ddp.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae-ddp.py), [vae_extra_train.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae_extra_train.py) and [vae_core_server.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae_core_server.py). diff --git a/docs/hpc.md b/docs/hpc.md new file mode 100644 index 0000000..8826d83 --- /dev/null +++ b/docs/hpc.md @@ -0,0 +1,60 @@ +# HPC systems (Slurm, Slingshot) + +Running DDStore under Slurm on Cray Slingshot systems (Frontier, Perlmutter): +the example job scripts, which `--network` options a layout needs, and what +to do when RDMA can't connect. + +## Do you need `--network` flags? + +Every `srun` step gets its own Slingshot VNI (network isolation ID), and two +endpoints can only reach each other on the same VNI. Whether you need +`#SBATCH --network=single_node_vni,job_vni` depends on how steps talk: + +| Layout | Perlmutter, no flags | Perlmutter, `single_node_vni,job_vni` | Frontier | +|---|---|---|---| +| One `srun` step (any number of nodes): normal training, `method=0`/`1` | works | works | works; `single_node_vni` needed if the step has one node | +| Core and extra as separate one-node steps (`--layout=split-node`) | works: one-node steps get no VNI and share the default one | needs `job_vni` + the wrapper below | needs both flags + the wrapper | +| Core and extra on the same nodes, or steps spanning several nodes | extra can't reach core | needs `job_vni` + the wrapper + `srun --overlap` | second step fails to launch | + +The Perlmutter columns were measured on 2 nodes (job scripts and logs in the +[results](results.md#perlmutter-validation-2-nodes--4-a100-cuda-13-cxi)). +## Slurm job scripts + +[job-vae-single.sh](https://github.com/ORNL/DDStore/blob/main/examples/vae/script/job-vae-single.sh) runs `vae-ddp.py` as one `srun` step; [job-vae-core-extra.sh](https://github.com/ORNL/DDStore/blob/main/examples/vae/script/job-vae-core-extra.sh) runs the core/extra split as two steps. Run either with `--help` for all options. Their `#SBATCH` lines target Frontier (`-A FUS184`, 8 ranks × 7 cores per node); see below for Perlmutter. + +```bash +sbatch examples/vae/script/job-vae-single.sh --method=1 --num-workers=1 +sbatch examples/vae/script/job-vae-single.sh --method=1 --gpudirect --image-scale=2 +sbatch examples/vae/script/job-vae-core-extra.sh # split-node: 1 core node, the rest extra +sbatch examples/vae/script/job-vae-core-extra.sh --gpudirect --core-nnodes=2 +``` + +`job-vae-core-extra.sh` sets up Slingshot networking for its two steps (see [Multiple `srun` steps](hpc.md#multiple-srun-steps-in-one-job-method2-cxi)). `--layout=colocate` (both steps on the same nodes) works on Perlmutter only. + +Both scripts also run on Perlmutter: they detect the machine (`NERSC_HOST`) and use 4 ranks per node, with the training ranks seeing all 4 GPUs of their node (NCCL needs that there). Override the Frontier `#SBATCH` lines when submitting: + +```bash +sbatch -A -C gpu --gpus-per-node=4 examples/vae/script/job-vae-single.sh --method=1 +``` + + +## Multiple `srun` steps in one job (`method=2`, `cxi`) + +On Slingshot every `srun` step gets its own VNI (network isolation ID), and two endpoints can only talk on the same VNI. For core and extra running as separate steps, [job-vae-core-extra.sh](https://github.com/ORNL/DDStore/blob/main/examples/vae/script/job-vae-core-extra.sh) does both of these: + +1. `#SBATCH --network=single_node_vni,job_vni`: `job_vni` adds a job-wide VNI to every step (`SLINGSHOT_VNIS=,`); on Frontier, `single_node_vni` is also what gives single-node steps a CXI service at all (without it `fi_domain()` fails with `-38`). +2. In each task, before Python starts: `export SLINGSHOT_VNIS=${SLINGSHOT_VNIS##*,}`. libfabric's cxi provider uses only the first VNI listed (the step's own), so without this reads fail with `VNI_NOT_FOUND`. + +Colocating both steps on the same nodes: + +| | `job_vni` + wrapper | `job_vni` + wrapper + `srun --overlap` | no `--network` | +|---|---|---|---| +| Frontier | second step fails to launch: `Error configuring interconnect` | same failure | each step has only its own VNI: extra cannot reach core | +| Perlmutter | second step does not start | **works** | extra cannot reach core | + +So the script defaults to `--layout=split-node`; `--layout=colocate` (with `--overlap`) is for Perlmutter. Perlmutter's single-node steps also work without the `--network` flags. + + +## Troubleshooting: RDMA fails to connect (`cxi`) + +If `fi_domain()` fails with `-38 (Function not implemented)`, the step has no CXI service: add `#SBATCH --network=single_node_vni` (needed on Frontier for any single-node step, i.e. a `-N 1` job or a one-node step inside a larger job). If ranks in different `srun` steps can't reach each other (`VNI_NOT_FOUND`), see [Multiple `srun` steps](hpc.md#multiple-srun-steps-in-one-job-method2-cxi) above. diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..4265377 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,52 @@ +# DDStore + +DDStore logo + +Efficient distributed data loading for distributed data-parallel (DDP) training. + +Each MPI rank holds a shard of the full dataset in memory. DDStore exposes a global index space so any rank can read any sample via one-sided remote memory access — either MPI RMA (default) or libfabric RDMA — without coordinator synchronization. + +- **Batched reads**: [`get_batch()`](api-pyddstore.md#get_batchname-arr-indices) fetches a whole training batch in one call (one-sided RDMA reads in flight together, or an MPI collective for `method=0`). +- **GPUDirect RDMA**: data can live in, and be read straight into, GPU memory ([details](gpudirect.md)). +- **PyTorch integration**: [`pyddstore.torch`](pytorch.md) turns any map-style dataset into a distributed one (`DistDataset`) and provides a thread-based `ThreadDataLoader` that is safe with MPI and GPU buffers. +- **Thread-safe** reads, a [profiler](performance.md) for where read time goes, and a split mode (`method=2`) where a separate job reads data published by another. + +DDStore architecture + +```{toctree} +:caption: Getting started +:maxdepth: 2 + +installation +quickstart +``` + +```{toctree} +:caption: User guide +:maxdepth: 2 + +backends +gpudirect +pytorch +hpc +performance +concurrency +``` + +```{toctree} +:caption: Reference +:maxdepth: 2 + +api-pyddstore +api-torch +environment +``` + +```{toctree} +:caption: More +:maxdepth: 1 + +testing +results +citation +``` diff --git a/docs/installation.md b/docs/installation.md new file mode 100644 index 0000000..bb088df --- /dev/null +++ b/docs/installation.md @@ -0,0 +1,47 @@ +# Installation + +## Prerequisites + +| Dependency | Notes | +|---|---| +| MPI (OpenMPI / MPICH) | `mpicc` and `mpicxx` must be on `PATH` | +| libfabric | Required for the RDMA backends (`method=1` and `method=2`) | +| Python ≥ 3.9 | | +| NumPy, mpi4py, Cython | Python build dependencies | +| PyTorch (optional) | For `pyddstore.torch` and GPU buffers (CUDA or ROCm build) | + +## Building and installing + +```bash +# Install Python build dependencies +pip install numpy mpi4py Cython + +# Build in-place (use with PYTHONPATH=$PWD/src:$PYTHONPATH) +CC=mpicc CXX=mpicxx python setup.py build_ext --inplace + +# Or install into the active virtual environment +CC=mpicc CXX=mpicxx pip install . +CC=mpicc CXX=mpicxx pip install ".[torch]" # also pulls PyTorch, for pyddstore.torch + +# Or install in editable/development mode +CC=mpicc CXX=mpicxx pip install -e . + +# Or install directly from GitHub +CC=mpicc CXX=mpicxx pip install git+https://github.com/ORNL/DDStore.git +``` + +To build against the packages already in the current environment (e.g. an `mpi4py` built against Cray MPICH) instead of letting pip fetch fresh build dependencies into an isolated build environment, disable build isolation: + +```bash +CC=cc CXX=CC pip install --no-build-isolation --no-deps -e . +``` + +If that fails with `ModuleNotFoundError: No module named 'distutils.msvccompiler'` (newer setuptools combined with an older system NumPy, e.g. `cray-python/3.11.7` on Frontier), point setuptools at the standard-library `distutils` for the build: + +```bash +SETUPTOOLS_USE_DISTUTILS=stdlib CC=cc CXX=CC pip install --no-build-isolation --no-deps -e . +``` + +The package is `pyddstore` (compiled core `pyddstore._core`, plus `pyddstore.torch`). After updating from a 1.x checkout, rebuild; an old `src/pyddstore.cpython-*.so` or `src/pyddstore.cpp` left behind is unused and can be deleted (the build warns about them). + +Editable and in-place builds keep the generated `src/pyddstore/_core.cpp` in the checkout, shared by every environment that builds from it. `setup.py` regenerates it whenever the NumPy major version differs from the previous build's, because a file generated against NumPy 2 doesn't compile against NumPy 1.x headers. diff --git a/docs/performance.md b/docs/performance.md new file mode 100644 index 0000000..1c98047 --- /dev/null +++ b/docs/performance.md @@ -0,0 +1,8 @@ +# Performance + +- Use **batched reads** (the default with `DistDataset`, or `get_batch()` directly). They cut per-sample cost by 10–27× for small rows and make the GPU path insensitive to worker threads; in the VAE every configuration got 1.2–3.9× faster per epoch. +- **Reuse destination buffers** and [`register_recv()`](api-pyddstore.md#register_recvname-arr--unregister_recvname-arr) them. Reading into a fresh buffer every time re-registers memory on every read; `get_profile(name)["mr_miss"]` counts those registrations. +- `method=1` (one-sided `fi_read`) is the fastest backend; `method=0` with batching (collective) comes close for small rows. +- `DDSTORE_PROFILE=1` + `get_profile(name)` shows where `get()`/`get_batch()` time goes: lock wait, memory registration, posting and completing `fi_read`, GPU sync. `vae-ddp.py` prints an all-rank summary when it is set. [examples/scripts/bench_get.py](https://github.com/ORNL/DDStore/blob/main/examples/scripts/bench_get.py) measures per-row latency and throughput vs row size, destination, batch size and threads. + +Measurements, profiles and the experiments behind these choices: [docs/results.md](results.md). diff --git a/docs/pytorch.md b/docs/pytorch.md new file mode 100644 index 0000000..b02dfb3 --- /dev/null +++ b/docs/pytorch.md @@ -0,0 +1,62 @@ +# PyTorch integration + +`pyddstore.torch` (needs PyTorch: `pip install .[torch]` or an existing PyTorch) turns any map-style dataset into a distributed one: + +```python +import torch # import torch before MPI starts +from mpi4py import MPI +from pyddstore.torch import DistDataset, ThreadDataLoader + +trainset = DistDataset(my_dataset, "train", MPI.COMM_WORLD) # each rank loads only its share +sampler = torch.utils.data.distributed.DistributedSampler(trainset) +loader = ThreadDataLoader(trainset, batch_size=128, sampler=sampler, num_workers=1) +for x, y in loader: + ... +``` + +- **`DistDataset(source, name, comm=None, ddstore_width=None, device=None, add_device=None, method=None, handshake_dir=None, chunk_size=None, encode=None, decode=None, fields=None)`**: each rank loads its contiguous share of `source` (anything with `len()` and `[i]`) into DDStore; every rank can then read every sample. Samples keep the source's structure (a tensor, numpy array or number; a tuple or list of them; or a dict of them), with each field's shape and dtype. Fields must have the same shape and dtype in every sample; supported dtypes are bool, uint8, int32, int64, float32 and float64. A field (or the whole sample) may also be a numpy structured record (`np.void`, a structured `ndarray`, or `np.recarray`) of any field types: it is stored as raw bytes and comes back as the same kind of object with the same layout (`default_collate` can't batch records, so pass a `collate_fn`). `ds.shapes` / `ds.dtypes` describe the fields, `ds.ddstore` is the underlying `PyDDStore`. + - `method` (default `DDSTORE_METHOD` or 0) picks the backend; `ddstore_width` splits `comm` into independent stores (see [Partitioned usage](backends.md#partitioned--sub-communicator-usage)). + - `device` puts tensor fields of read samples on a GPU and `add_device` keeps each rank's share there ([GPUDirect](gpudirect.md)). + - `chunk_size` loads each rank's share that many samples at a time, writing each chunk into the store before reading the next: peak memory is about the share plus one chunk, instead of about three times the share (400 MiB share: 461 vs 1202 MiB). Host storage only (not with `add_device`). + - `encode` / `decode` / `fields` handle samples that can't be stored as they are (strings, labels, metadata objects, data shared by a group of samples). `encode(sample)` runs on every source sample before it is stored and returns what to store; `decode(stored, index)` runs on every sample read (`ds[i]`, `__getitems__`, so also in loaders) and rebuilds the full sample, e.g. adding constants or looking up tables by index or by a stored id. `decode` runs on the reading rank, so anything it looks up must exist on every rank. `fields=[...]` keeps only those keys (dict samples) or positions (tuple/list samples); it can't be combined with `encode`. `read_rows()` and `WindowedDataset` return stored rows without `decode`. `DistDatasetReader` takes `decode` too. + +```python +LABELS = ["cat", "dog", "owl"] +ds = DistDataset(src, "pets", comm, + encode=lambda s: {"x": s["x"], "label": LABELS.index(s["label"]), "group": s["group"]}, + decode=lambda d, i: {**d, "label": LABELS[d["label"]], "meta": GROUP_INFO[d["group"]]}) +``` + + - **Batched by default**: `__getitems__` reads a whole batch with one [`get_batch()`](api-pyddstore.md#get_batchname-arr-indices) per field, which `DataLoader` and `ThreadDataLoader` call automatically; `DDSTORE_BATCH_GET=0` reads one sample at a time. With `method=0` batched reads are collective, so every rank must iterate the same number of batches from one thread (`DistributedSampler` does). +- **Reading rows and reusing buffers** (`DistDataset` and `DistDatasetReader`): + - `ds.read_rows(rows, fields=None, out=None)` reads stored rows `rows` (any order, repeats allowed) of the selected fields (default all), one `get_batch()` per field, and returns a dict of values shaped `(len(rows), *field_shape)`. Keys are the sample's: dict keys, tuple/list positions, or `0` for a single value. With `method=0` it is collective, like `get_batch()`. + - `ds.alloc(n, fields=None)` returns buffers for `n` rows, one per field and keyed the same way, each [registered](api-pyddstore.md#register_recvname-arr--unregister_recvname-arr) once. Pass them as `out=` to `read_rows()` or `ds.__getitems__(idx, out=)`: reads then skip memory registration, and the results are views into the buffers, so you decide when a buffer can be reused. `ds.release(bufs)` unregisters them; `free()` does too. +- **`WindowedDataset(ds, window, stride=1, dilation=1, starts=None, fields=None)`**: samples made of several stored rows of `ds` (time windows, clips, sequences), each stored row held once. Sample `i` is rows `s, s + dilation, …, s + (window - 1)·dilation` with `s = i·stride`, or `s = starts[i]` when `starts` is given; use `starts` to keep only windows that don't cross a trajectory or file boundary. Fields come back stacked, `(window, *field_shape)`, in the structure of `ds`'s samples (a dict when `fields` is given). A batch of windows is one `read_rows()`. +- **`row_of(concat, source, index)`**: with several sources in one store (a `DistDataset` over a `torch.utils.data.ConcatDataset`), the row of sample `index` of source `source`, e.g. to map (file, trajectory, step) to `starts`: + +```python +from torch.utils.data import ConcatDataset +from pyddstore.torch import DistDataset, WindowedDataset, row_of + +files = ConcatDataset([StepsOf(f) for f in paths]) # one sample per time step +frames = DistDataset(files, "frames", comm, method=1) +starts = [row_of(files, k, t) for k, f in enumerate(paths) # windows inside each file + for t in range(0, len(files.datasets[k]) - 2 * dt)] +pairs = WindowedDataset(frames, window=3, dilation=dt, starts=starts) # (t, t+dt, t+2dt) +``` + +- **`DistDatasetReader(name, handshake_dir=None, n_core=None, device=None, decode=None)`**: the same dataset read by a separate `method=2` extra job; it learns the fields from a `{name}.meta.json` file the core group writes next to the handshake records. +- **`ThreadDataLoader(dataset, reuse_buffers=False, collate_copies=False, **DataLoader args)`**: a `DataLoader` whose workers are threads, not forked processes, so it is safe with MPI and GPU buffers. Each batch is fetched, collated and optionally pinned in a worker thread; random draws match `DataLoader`'s. As with `DataLoader`, `iter(loader)` returns a separate iterator for one epoch (`list(it)` or `islice(it, …)` after `next(it)` continue the epoch; iterators over one loader are independent), and while the training step holds a batch, `num_workers * prefetch_factor` more are being fetched, so that many plus one are in memory. One or two workers are enough: reads on one variable are serialized by its lock, and one batched read already keeps the network busy. + - `reuse_buffers=True` reads every batch into one of a fixed pool of `num_workers` buffer sets from `dataset.alloc(batch_size)`, registered once, instead of fresh buffers registered on every read (for large rows, registration can cost more than the transfer). A worker reads and collates its batch, then returns the set, so the collate must copy: it needs `batch_size` and the default `collate_fn`, or `collate_copies=True` to declare that your `collate_fn` copies; other setups raise. `loader.close()` (or deleting the loader) waits for running fetches and unregisters the pool. `DDSTORE_AFFINITY_WIDTH` / `DDSTORE_AFFINITY_OFFSET` pin worker threads to CPUs. + +[examples/vae/vae-ddp.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae-ddp.py) trains a VAE with DDP on top of it: + +```bash +DDSTORE_METHOD=1 DDSTORE_FABRIC=cxi mpirun -n 4 python examples/vae/vae-ddp.py --num-workers=1 +DDSTORE_METHOD=1 DDSTORE_FABRIC=cxi mpirun -n 4 python examples/vae/vae-ddp.py --num-workers=1 --gpu-dest --gpu-source +``` + +- **`--num-workers`** (default 0): `0` uses PyTorch's standard `DataLoader`; `> 0` uses `ThreadDataLoader` (needs `method=1`/`2`). +- **`--gpu-dest` / `--gpu-source`**: `DistDataset`'s `device` / `add_device`. +- **`--replicate R`** repeats the training set R times (longer epochs); **`--image-scale S`** upscales images to (28·S)² so each row is S² larger. Both default to 1, the original example. +- The [method=2 split](backends.md#file-based-handshake-method2) variant is [vae_core_server.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae_core_server.py) (a `DistDataset` core group) + [vae_extra_train.py](https://github.com/ORNL/DDStore/blob/main/examples/vae/vae_extra_train.py) (a `DistDatasetReader`), with the same options. diff --git a/docs/quickstart.md b/docs/quickstart.md new file mode 100644 index 0000000..a75f64b --- /dev/null +++ b/docs/quickstart.md @@ -0,0 +1,32 @@ +# Quick start + +```python +import numpy as np +from mpi4py import MPI +import pyddstore as dds + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() + +# Each rank contributes its own shard +store = dds.PyDDStore(comm) # MPI RMA backend (default) +# store = dds.PyDDStore(comm, method=1) # libfabric RDMA backend + +data = np.random.rand(1024, 64).astype(np.float32) +store.add("features", data) # collective — all ranks must call + +# Read any global sample index +out = np.zeros((1, 64), dtype=np.float32) +store.epoch_begin() +store.get("features", out, start=2048) # global index across all shards +store.epoch_end() + +store.free() +``` + +Run with: +```bash +mpirun -n 4 python my_script.py +``` + +With PyTorch, [`pyddstore.torch.DistDataset`](pytorch.md) does the sharding, `add()` and batched reads for you. diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000..d1cf123 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,4 @@ +sphinx>=7 +myst-parser>=2 +furo +numpy diff --git a/docs/results.md b/docs/results.md new file mode 100644 index 0000000..1d8176c --- /dev/null +++ b/docs/results.md @@ -0,0 +1,160 @@ +# DDStore measurements and findings + +Measurements behind the recommendations in the [documentation](index.md), +collected on the `check-thread` branch in October 2026. Unless noted: +`method=1`, `DDSTORE_FABRIC=cxi`, `vae-ddp.py` with `VAE_PROFILE=1`, epoch +times averaged over all epochs but the first. Epochs are short (0.1–0.6 s), +so treat differences under ~10% as noise; numbers from different jobs vary by +up to ~2× (different nodes), comparisons within one table are from one job. + +Machines: **Frontier** (AMD MI250X, ROCm 7.2, 8 ranks × 7 cores per node) and +**Perlmutter** (NVIDIA A100, CUDA 13, 4 ranks × 32 cores per node). + +## Batched reads (`get_batch`) vs one `get()` per sample + +VAE, seconds per epoch, per-sample (`DDSTORE_BATCH_GET=0`) → batched (`1`). +Losses are identical within every row. + +**Frontier, 2 nodes × 8 ranks, 8 epochs** (loss 8.8960 at S=1, 30.0178 at S=2): + +| `--image-scale` | host, 0 workers | host, 1 worker | GPU, 0 workers | GPU, 1 worker | GPU, 2 workers | +|---|---|---|---|---|---| +| 1 | 0.247 → 0.132 | 0.156 → 0.126 | 0.263 → 0.135 | 0.259 → 0.127 | 0.403 → 0.134 | +| 2 | 0.395 → 0.286 | 0.293 → 0.265 | 0.428 → 0.277 | 0.416 → 0.264 | 0.518 → 0.297 | + +"GPU" = `--gpu-dest --gpu-source`. + +**Frontier, 4 nodes × 8 ranks, `job-vae-single.sh` (3 epochs)** (loss 6.7155 / 24.9987): + +| | S=1 | S=2 | +|---|---|---| +| method 0, 0 workers | 0.241 → 0.088 | 0.356 → 0.203 | +| method 1, 0 workers | 0.128 → 0.088 | 0.254 → 0.185 | +| method 1, 2 workers | 0.096 → 0.072 | 0.218 → 0.175 | +| method 1, 4 workers | 0.113 → 0.067 | 0.193 → 0.174 | + +**Perlmutter, 2 nodes × 4 ranks, 8 epochs** (loss 15.3781 at S=1): + +| | per-sample → batched | +|---|---| +| method 1, host, 0 workers | 0.398 → 0.208 | +| method 1, host, 2 workers | 0.380 → 0.181 | +| method 1, GPU, 0 workers | 0.508 → 0.204 | +| method 1, GPU, 2 workers | 0.678 → 0.183 | +| method 0, host, 0 workers | 0.849 → 0.219 | + +## Where `get()` time goes (`DDSTORE_PROFILE=1`) + +Frontier, 2 nodes × 8 ranks, VAE, µs per call, one `get()` per sample: + +| run | total | GPU sync | lock wait | MR | read post | CQ wait | other | +|---|---|---|---|---|---|---|---| +| host, 0/1 workers (S=1) | 7.9 / 9.3 | – | 0.0 | 0.1 | 0.6 | 3.6 | 3.6 / 5.0 | +| GPU, 0 workers (S=1 / S=2) | 13.6 / 13.6 | 3.8 / 3.7 | 0.1 | 0.3 | 0.6 | 3.7 / 4.0 | ~5 | +| GPU, 1 worker (S=1 / S=2) | 142.7 / 198.7 | **131.6 / 187.3** | 0.1 | 0.4 | 0.6 | 3.7 / 4.0 | ~6 | +| GPU, 2 workers (S=1 / S=2) | 310.7 / 452.8 | **287.2 / 429.3** | 0.2 | 0.7 | 0.7 | 3.7 / 4.1 | ~18 | + +- The per-call whole-device `torch.cuda.synchronize()` dominates the GPU path + once a worker thread runs next to training: the worker's sync waits for the + training kernels. With no workers the GPU is idle and the sync costs ~4 µs. +- Lock contention (≤0.2 µs) and MR registration (<1 µs, even at 100% cache + misses; libfabric caches registrations) are negligible. The RDMA round trip + is ~4–5 µs. +- With `get_batch()` the per-row sync drops to ~0.1 µs and compute time + recovers (training no longer waits behind the workers' syncs). + +## `bench_get.py`: µs per row + +Single-row `get()` vs `get_batch()` of 128 rows, 1 thread, host vs fresh GPU +destination. + +| row | Frontier host, 1 / 128 | Frontier GPU, 1 / 128 | Perlmutter host, 1 / 128 | Perlmutter GPU, 1 / 128 | +|---|---|---|---|---| +| 3 KB | 9.0 / 0.68 | 20.9 / 0.78 | 8.6 / 0.82 | 21.6 / 0.99 | +| 12.5 KB | 10.1 / 1.82 | 21.9 / 1.38 | 9.3 / 1.19 | 22.4 / 1.80 | +| 200 KB | 32.5 / 27.5 | 37.9 / 19.6 | – | – | +| 1 MB | 160 / 189 | 134 / 99 | 88 / 75 | 142 / 142 | + +- Batching makes small rows 10–27× cheaper per row on both machines. +- Host vs GPU destination is platform-dependent: on Frontier GPUDirect wins + from ~12.5 KB rows (up to ~10.6 GB/s per rank); on Perlmutter host + destinations win at every size (1 MB: ~14 vs ~7.4 GB/s per rank). +- On Frontier, host-destination batches of 1 MB rows are slower than single + reads (not investigated). +- A second thread adds no per-rank throughput: the per-variable lock + serializes transfers on one variable (lock wait ≈ transfer time at large + rows). One batch already keeps up to 128 reads in flight. + +## `method=0`: per-row `MPI_Get` vs collective `get_batch` + +Frontier, 16 ranks, host, µs per row (per-row `get()` → batch 128, 2 MiB rounds): +3 KB 28.8 → 3.0; 12.5 KB 34.9 → 7.5; 200 KB 170 → 111; 1 MB 739 → 685. + +Round size (`DDSTORE_ALLTOALL_MAX_BYTES`) at batch 128: 200 KB rows — no cap +272, 2 MiB 111, 8 MiB 125, 32 MiB 267; 1 MB rows — no cap 1765, 2 MiB 685, +8 MiB 696, 32 MiB 1262. One unbounded exchange of large rows is slower than +per-row reads; 2–8 MiB rounds fix it. One-sided `method=1` batching is still +faster at every size (3 KB: 1.6 µs/row in the same job). + +## Worker threads (`ThreadDataLoader`), before batching + +Frontier, 2 nodes × 8 ranks, one `get()` per sample, s/epoch (fetch / total): + +| `--num-workers` | host | GPU (`--gpu-dest --gpu-source`) | +|---|---|---| +| 0 (`DataLoader`) | 0.11 / 0.22 | 0.15 / 0.25 | +| 1 | 0.02 / 0.14–0.16 | 0.09 / 0.26 | +| 2 | 0.03 / 0.16 | 0.11 / 0.31 | +| 4 | 0.04 / 0.17 | 0.12 / 0.30 | +| 8 | 0.08 / 0.21 | 0.15 / 0.35–0.38 | + +One worker hid the fetch on the host path; more workers only contended for +the lock and the GIL. On the GPU path threads did not help (per-call device +sync). Moving collation into the worker thread later brought host, 1 worker +to fetch ≈ 0.006 s / total ≈ 0.15 s. Pinned memory with +`.to(device, non_blocking=True)` gave no gain (and was ~3× slower with 0 +workers, where pinning runs on the training thread). + +## Correctness experiments + +- **GPU sync in `get()`**: removing it made every `--gpu-dest` VAE run abort + in the first epoch with `HSA_STATUS_ERROR_EXCEPTION ... code: 0x1016` (GPU + memory fault) on every rank (Frontier); host and `--gpu-source`-only runs + were unaffected. The sync stays. +- **Per-variable lock**: without it, concurrent `get()` calls crashed with + `double free or corruption`; it protects the shared recv fields and MR cache. +- **GPU destination-buffer pool** (removed): round-robin slices of one + pre-registered buffer corrupted data under `ThreadDataLoader` with more than + one worker (slot order followed lock acquisition, not batch order). Each + `get()` now uses its own fresh tensor; MR registration is cheap enough. +- **MPI thread level**: only the main thread calls MPI. `vae-ddp.py` and the + test suites gave identical results and timing with `MPI_THREAD_SINGLE`, + `FUNNELED` and `MULTIPLE` (Frontier, Cray MPICH). + +## HIP streams and hardware queues (Frontier, ROCm 7.2) + +HIP maps streams onto `GPU_MAX_HW_QUEUES` hardware queues per GPU per process +(default 4); extra streams share a queue round-robin, and work in a shared +queue runs in order. With the default stream kept busy, 12 of 16 new streams +were independent of it by default (every 4th collided, including the first +created), 14 of 16 with `GPU_MAX_HW_QUEUES=8`, 15 of 16 with `16`. Relevant if +data-loading threads get their own streams. + +## Perlmutter validation (2 nodes × 4 A100, CUDA 13, cxi) + +One 2-node debug job covering what Frontier could not: + +- **Tests**: `test_single` 14/14, `test_multirank` 5/5, `test_get_batch` + 20/20 on 8 ranks, `test_gpu_rdma` 14/14 (CUDA GPUDirect, `FI_HMEM_CUDA`). +- **VAE**: identical losses for every variant (method 0/1, host/GPU, 0/2 + workers, per-sample/batched): 15.3781 at S=1, 54.7914 at S=2. +- **Slingshot**: each step's `SLINGSHOT_VNIS` is `,`, job + VNI last, as on Frontier, so the same wrapper works. Ranks spread over + `cxi0`–`cxi3`. Single-node steps work even without `--network` flags (no + `SLINGSHOT_*` variables; cxi falls back to a default CXI service). +- **core/extra**: split-node passes (host and `--gpu-dest`); colocate works + only with `job_vni` + the wrapper + `srun --overlap`. +- **DDP setup**: training ranks must see all 4 GPUs of their node + (`--gpus-per-node=4`, each picks `cuda:$SLURM_LOCALID`); with + `--gpus-per-task=1`, NCCL 2.29 fails in DDP setup with "Cuda failure 101 + 'invalid device ordinal'". The job scripts handle this. diff --git a/docs/testing.md b/docs/testing.md new file mode 100644 index 0000000..a1762fa --- /dev/null +++ b/docs/testing.md @@ -0,0 +1,72 @@ +# Testing + +## Unit tests (pytest) + +Install test dependencies: + +```bash +pip install pytest pytest-mpi +``` + +**Single-rank** — no cluster required, covers all dtypes, `add`/`get`/`init`/`update`, and error cases: + +```bash +mpirun -n 1 python -m pytest test/test_single.py -v +``` + +**Multi-rank** — verifies remote reads across all rank pairs and sub-communicator grouping: + +```bash +mpirun -n 4 python -m pytest test/test_multirank.py -v +``` + +**Batched reads and the PyTorch layer** — method 0 everywhere; method 1, method 2 and GPU cases run inside a Slurm step with a CXI device (provider from `DDSTORE_FABRIC`, default `cxi`): + +```bash +mpirun -n 4 python -m pytest test/test_get_batch.py test/test_torch.py -v +``` + +**GPUDirect RDMA** — requires a live `cxi` fabric and a CUDA/HIP GPU per rank (skipped automatically otherwise); see [GPUDirect RDMA](gpudirect.md): + +```bash +DDSTORE_FABRIC=cxi mpirun -n 2 python -m pytest test/test_gpu_rdma.py -v +``` + +| Test file | Min ranks | What is tested | +|---|---|---| +| `test/test_single.py` | 1 | All dtypes, `add`/`get`, `init`/`update`/`get`, error handling, double `free()` | +| `test/test_multirank.py` | 2 (4 recommended) | Remote reads, shard boundaries, multiple variables, `ddstore_width` grouping | +| `test/test_gpu_rdma.py` | 2 | GPU-resident `add()`/`get()` in both directions, both libfabric methods, negative/error cases | +| `test/test_get_batch.py` | 2 (4 recommended) | `get_batch()`: shuffled indices across ranks with repeats, single row, dtypes, error recovery, GPU destination, concurrent threads, registered destination buffers, wide rows (with `DDSTORE_MAX_READ_BYTES=4096` they are read in pieces); method 0, plus method 1 over `cxi` inside a Slurm step | +| `test/test_torch.py` | 2 (4 recommended) | `pyddstore.torch`: tuple/dict/single samples of every field kind, numpy records, loaders vs the plain source, `ThreadDataLoader` iterators and prefetch depth, `read_rows`/`alloc` buffers, `reuse_buffers` (also with GPU buffers), `WindowedDataset`, `row_of`, `encode`/`decode`/`fields`, chunked loading, error handling, `ddstore_width`, GPU placement, `DistDatasetReader` | + +## Integration scripts + +```bash +# Basic functional test (libfabric, method=1) +mpirun -n 4 python examples/scripts/demo.py + +# Integration test with PyTorch DDP (libfabric, method=1) +mpirun -n 4 python examples/scripts/test.py +``` + +Optional arguments for `examples/scripts/demo.py` and `examples/scripts/test.py`: + +| Flag | Default | Description | +|---|---|---| +| `--num` | `1048576` | Rows per rank | +| `--dim` | `64` | Elements per row | +| `--nbatch` | `32` | Number of random reads | +| `--gloo` / `--nccl` | `--gloo` | `test.py` only: `torch.distributed` backend | + +## Method 2 (file-based handshake) + +Two separate launches sharing a handshake directory on a shared filesystem — not a single `mpirun`, since core and extra are independent jobs: + +```bash +# Terminal 1 — core (data-holding) side +mpirun -n 4 python test/test_method2_core.py /path/to/shared/ddstore_hs + +# Terminal 2 — extra (reader) side, after or while the core side is running +python test/test_method2_extra.py /path/to/shared/ddstore_hs 4 +``` diff --git a/examples/scripts/bench_get.py b/examples/scripts/bench_get.py new file mode 100644 index 0000000..d581901 --- /dev/null +++ b/examples/scripts/bench_get.py @@ -0,0 +1,179 @@ +"""Microbenchmark for DDStore get(): per-call latency and throughput vs row size. + +Each rank adds a shard of float32 rows, then reads --nget uniformly random +global rows (most of them remote), either one get() per row (--batch 1) or +get_batch() of --batch rows per call, and reports time per row, per-rank +throughput, and the DDSTORE_PROFILE breakdown per row (lock wait / MR / +fi_read post / CQ wait / GPU sync). + +Run (method 1, cxi), e.g. on 2 nodes: + DDSTORE_PROFILE=1 DDSTORE_FABRIC=cxi srun -N2 -n16 -c7 --gpus-per-task=1 \\ + python examples/scripts/bench_get.py --row-floats 784,3136 --dest host,gpu + +Destinations: "host" reads into a reused numpy row; "gpu" allocates a fresh +torch.empty() per call, like DistDataset.get() with --gpu-dest; "gpu-reuse" +reuses one GPU row (no allocator churn). +""" + +import argparse +import os +import threading +import time + +## torch must load before mpi4py triggers MPI_Init (see vae-ddp.py). +import torch +import numpy as np +from mpi4py import MPI + +import pyddstore as dds + +parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) +parser.add_argument( + "--row-floats", + default="784,3136", + help="comma-separated row widths in float32 (784 = MNIST " + "28x28 / --image-scale 1, 3136 = 56x56 / --image-scale 2)", +) +parser.add_argument("--rows-per-rank", type=int, default=4096) +parser.add_argument( + "--nget", type=int, default=4000, help="get() calls per rank per configuration" +) +parser.add_argument( + "--dest", default="host,gpu", help="comma-separated: host, gpu, gpu-reuse" +) +parser.add_argument( + "--gpu-source", + action="store_true", + help="add() the shard as a GPU tensor instead of numpy", +) +parser.add_argument( + "--batch", + default="1", + help="comma-separated rows per call: 1 = get(), >1 = get_batch()", +) +parser.add_argument( + "--threads", + default="1", + help="comma-separated thread counts issuing get()s concurrently", +) +parser.add_argument( + "--method", type=int, default=int(os.environ.get("DDSTORE_METHOD", "1")) +) +args = parser.parse_args() + +comm = MPI.COMM_WORLD +rank, size = comm.Get_rank(), comm.Get_size() +ngpu = torch.cuda.device_count() +device = ( + torch.device(f"cuda:{int(os.environ.get('SLURM_LOCALID', 0)) % ngpu}") + if ngpu + else None +) +if device is not None: + torch.cuda.set_device(device) + +row_floats = [int(x) for x in args.row_floats.split(",")] +dests = args.dest.split(",") +thread_counts = [int(x) for x in args.threads.split(",")] +batches = [int(x) for x in args.batch.split(",")] +total_rows = args.rows_per_rank * size + +if rank == 0: + print( + f"ranks={size} rows/rank={args.rows_per_rank} nget/rank={args.nget} " + f"method={args.method} fabric={os.environ.get('DDSTORE_FABRIC', 'hsn')} " + f"gpu_source={args.gpu_source} profile={os.environ.get('DDSTORE_PROFILE', '0')}", + flush=True, + ) + print( + f"{'row_B':>8} {'dest':>9} {'batch':>5} {'thr':>3} {'us/row':>8} {'MB/s/rank':>9} | " + f"{'sync':>6} {'lock':>6} {'mr':>6} {'miss%':>6} {'read':>6} {'cq':>6} {'other':>6} (us/row)", + flush=True, + ) + +for nf in row_floats: + for dest in dests: + if dest != "host" and device is None: + continue + for B, nthr in [(b, t) for b in batches for t in thread_counts]: + store = dds.PyDDStore(comm, method=args.method) + rng = np.random.default_rng(rank) + shard = np.full((args.rows_per_rank, nf), float(rank), dtype=np.float32) + if args.gpu_source: + shard = torch.from_numpy(shard).to(device) + store.add("x", shard) + comm.Barrier() + store.epoch_begin() + + idx = rng.integers(0, total_rows, size=args.nget) + chunks = np.array_split(idx, nthr) + + def worker(ids): + reuse_host = np.empty((B, nf), dtype=np.float32) + reuse_gpu = ( + torch.empty((B, nf), dtype=torch.float32, device=device) + if device is not None + else None + ) + for s0 in range(0, len(ids), B): + chunk = ids[s0 : s0 + B] + k = len(chunk) + if dest == "host": + out = reuse_host[:k] + elif dest == "gpu": + out = torch.empty((k, nf), dtype=torch.float32, device=device) + else: + out = reuse_gpu[:k] + if B == 1: + store.get("x", out, int(chunk[0])) + else: + store.get_batch("x", out, chunk) + + # One warm-up call so first-call registration isn't timed as typical. + worker(idx[:B]) + comm.Barrier() + p0 = store.get_profile("x") + comm.Barrier() + t0 = time.perf_counter() + ths = [threading.Thread(target=worker, args=(c,)) for c in chunks] + for t in ths: + t.start() + for t in ths: + t.join() + if device is not None: + torch.cuda.synchronize() + dt = time.perf_counter() - t0 + p1 = store.get_profile("x") + store.epoch_end() + + d = {k: p1[k] - p0[k] for k in p1} + vals = comm.gather((dt, d), root=0) + if rank == 0: + ngets = args.nget * size + # rows comes from the C++ counters (methods 1/2 only); method 0 + # has none, so fall back to the rows this run requested. + n = sum(v[1]["rows"] for v in vals) or ngets + us_get = 1e6 * sum(v[0] for v in vals) / ngets + mbps = nf * 4 * args.nget / (sum(v[0] for v in vals) / size) / 1e6 + s = lambda k: 1e6 * sum(v[1][k] for v in vals) / n + other = ( + s("py_get") + - s("py_sync") + - s("lock_wait") + - s("mr") + - s("read") + - s("cq") + ) + miss = 100.0 * sum(v[1]["mr_miss"] for v in vals) / n + print( + f"{nf * 4:>8} {dest:>9} {B:>5} {nthr:>3} {us_get:>8.2f} {mbps:>9.1f} | " + f"{s('py_sync'):>6.2f} {s('lock_wait'):>6.2f} {s('mr'):>6.2f} " + f"{miss:>6.1f} {s('read'):>6.2f} {s('cq'):>6.2f} {other:>6.2f}", + flush=True, + ) + # Every rank must finish reading before any rank tears down its + # endpoint (gather() doesn't synchronize non-root ranks). + comm.Barrier() + store.free() + del store + comm.Barrier() diff --git a/examples/scripts/demo.py b/examples/scripts/demo.py index bdc53fc..311d67e 100644 --- a/examples/scripts/demo.py +++ b/examples/scripts/demo.py @@ -1,8 +1,3 @@ -import mpi4py - -mpi4py.rc.thread_level = "serialized" -mpi4py.rc.threads = False - import numpy as np from mpi4py import MPI import argparse diff --git a/examples/scripts/test.py b/examples/scripts/test.py index 1019623..29c6d66 100644 --- a/examples/scripts/test.py +++ b/examples/scripts/test.py @@ -1,9 +1,3 @@ -import mpi4py - -## (2024/12): got an assert error on osx without the following line -mpi4py.rc.thread_level = "serialized" -mpi4py.rc.threads = False - import numpy as np from mpi4py import MPI import argparse diff --git a/examples/vae/README.md b/examples/vae/README.md index b903830..2c71631 100644 --- a/examples/vae/README.md +++ b/examples/vae/README.md @@ -2,5 +2,5 @@ Make sure the `pyddstore` module is installed correctly -To run this script and distrubute data across 4 processes: +To run this script and distribute data across 4 processes: `$ mpirun -n 4 python -u vae-ddp.py` diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index 12b4e69..946a2d0 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -38,9 +38,9 @@ def init_comm_size_and_rank(): def get_local_rank(rank): """ - Determine which GPU on the local node this rank should use. - Falls back to rank % device_count when no launcher-provided local rank - is available (e.g. plain mpirun without per-rank GPU visibility). + Determine which GPU on the local node this rank should use, from the + launcher's local rank (OMPI_COMM_WORLD_LOCAL_RANK or SLURM_LOCALID). + Returns 0 when neither is set; `rank` is currently unused. """ if os.getenv("OMPI_COMM_WORLD_LOCAL_RANK") is not None: return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) @@ -103,7 +103,7 @@ def parse_slurm_nodelist(nodelist): def setup_ddp(): - """ "Initialize DDP""" + """Initialize DDP""" if os.getenv("DDSTORE_BACKEND") is not None: backend = os.environ["DDSTORE_BACKEND"] diff --git a/examples/vae/distdataset.py b/examples/vae/distdataset.py deleted file mode 100644 index 6f07f58..0000000 --- a/examples/vae/distdataset.py +++ /dev/null @@ -1,163 +0,0 @@ -from mpi4py import MPI -import numpy as np -import os - -import torch -from torch.utils.data import Dataset - -import pyddstore as dds - - -def nsplit(a, n): - k, m = divmod(len(a), n) - return (a[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(n)) - - -class DistDataset(Dataset): - """Distributed dataset class""" - - def __init__(self, data, label, comm=MPI.COMM_WORLD, ddstore_width=None): - super().__init__() - - self.dataset = list() - self.label = label - self.comm = comm - self.rank = self.comm.Get_rank() - self.comm_size = self.comm.Get_size() - print("init", self.rank, self.comm_size) - self.ddstore_width = ( - ddstore_width if ddstore_width is not None else self.comm_size - ) - self.ddstore_comm = self.comm.Split(self.rank // self.ddstore_width, self.rank) - self.ddstore_comm_rank = self.ddstore_comm.Get_rank() - self.ddstore_comm_size = self.ddstore_comm.Get_size() - - ddstore_method = int(os.getenv("DDSTORE_METHOD", "0")) - print("DDStore method:", ddstore_method) - handshake_dir = os.getenv("DDSTORE_HANDSHAKE_DIR", "./ddstore_hs") - - if ddstore_method == 2 and self.ddstore_width != self.comm_size: - # File-based handshake: each Split group would publish into the - # same shared {varname}.bin file, so more than one group sharing - # a handshake_dir would silently collide. - raise NotImplementedError( - "method=2 does not yet support ddstore_width < comm_size " - "(multiple core groups would collide on the same " - "handshake_dir)" - ) - - self.ddstore = dds.PyDDStore( - self.ddstore_comm, method=ddstore_method, handshake_dir=handshake_dir - ) - print("FABRIC_IFACE:", os.environ.get("FABRIC_IFACE", "n/a (method=0)")) - - ## set total before set subset - self.total_ns = len(data) - print("init", self.total_ns) - - # WHEN READY FOR WHOLE DATA SET CHANGE THE RANGE TO range(len(data)) - rx = list(nsplit(range(len(data)), self.ddstore_comm_size))[ - self.ddstore_comm_rank - ] - - for i in rx: - self.dataset.append(data[i]) - - print(self.rank, len(self.dataset)) - self.data = list() - self.labels = list() - - nbytes = 0 - for data, label in self.dataset: - val = data.cpu().numpy() - val = val.flatten() - self.data.append(val) - self.labels.append(label) - - # np.stack (not concatenate) keeps one row per image (nrows, 784) so - # ddstore.add() infers disp=784 instead of flattening into a single - # (nrows*784,) vector, which it would read back as disp=1. - self.data = np.stack(self.data) - self.data = np.ascontiguousarray(self.data) - - self.labels = np.array(self.labels, dtype=np.int32) - self.labels = np.ascontiguousarray(self.labels) - - self.ddstore.add(f"{self.label}data", self.data) - self.ddstore.add(f"{self.label}labels", self.labels) - - def len(self): - return self.total_ns - - def __len__(self): - return self.len() - - def get(self, idx): - ## first dim must be the row count (1), not the flattened feature - ## width, since ddstore.get() infers count from arr.shape[0] - val = np.zeros((1, 28 * 28), dtype=np.float32) - label = np.zeros(1, dtype=np.int32) - val = np.ascontiguousarray(val) - assert val.data.contiguous - self.ddstore.get(f"{self.label}data", val, idx) - self.ddstore.get(f"{self.label}labels", label, idx) - # print("rank", self.rank, "fetching idx", idx) - val = torch.tensor(val) - val = torch.reshape(val, (1, 28, 28)) - return (val, label[0]) - - def __getitem__(self, idx): - return self.get(idx) - - -class DistDatasetReader(Dataset): - """Distributed dataset class — extra (read-only) member. - - Joins a variable published by a core group (see DistDataset) via - DDStore method=2's file-based handshake. Owns no MPI communicator and no - local copy of the data — every __getitem__ is an RDMA read against a - core rank's memory. - """ - - def __init__(self, label, handshake_dir, n_core): - super().__init__() - self.label = label - - self.ddstore = dds.PyDDStore( - None, method=2, handshake_dir=handshake_dir, n_core=n_core - ) - print("FABRIC_IFACE:", os.environ.get("FABRIC_IFACE", "n/a")) - self.ddstore.join(f"{label}data") - self.ddstore.join(f"{label}labels") - - self.total_ns, self.data_disp, self.data_itemsize = self.ddstore.info( - f"{label}data" - ) - self.side = int(round(self.data_disp**0.5)) - if self.side * self.side != self.data_disp: - raise ValueError( - f"joined '{label}data' has disp={self.data_disp}, " - "which is not a perfect square (expected a flattened square image)" - ) - - def len(self): - return self.total_ns - - def __len__(self): - return self.len() - - def get(self, idx): - ## first dim must be the row count (1), not the flattened feature - ## width, since ddstore.get() infers count from arr.shape[0] - val = np.zeros((1, self.data_disp), dtype=np.float32) - label = np.zeros(1, dtype=np.int32) - val = np.ascontiguousarray(val) - assert val.data.contiguous - self.ddstore.get(f"{self.label}data", val, idx) - self.ddstore.get(f"{self.label}labels", label, idx) - val = torch.tensor(val) - val = torch.reshape(val, (1, self.side, self.side)) - return (val, label[0]) - - def __getitem__(self, idx): - return self.get(idx) diff --git a/examples/vae/script/job-vae-core-extra.sh b/examples/vae/script/job-vae-core-extra.sh new file mode 100755 index 0000000..e324458 --- /dev/null +++ b/examples/vae/script/job-vae-core-extra.sh @@ -0,0 +1,157 @@ +#!/bin/bash +#SBATCH -A FUS184 +#SBATCH -J GX-core-extra +#SBATCH -o job-%j.out +#SBATCH -e job-%j.out +#SBATCH -N 4 +#SBATCH -t 30:00 +#SBATCH -q debug +#SBATCH --network=single_node_vni,job_vni +# +# Core/extra split VAE DDP run (two independent srun steps). +# +# Slingshot networking (Frontier): every srun step gets its own VNI, and two +# endpoints can only talk on the same VNI. --network=job_vni adds a job-wide +# VNI to every step (SLINGSHOT_VNIS=,), and single_node_vni +# makes single-node steps (split-node layout) get a CXI service at all -- +# without it fi_domain() fails with -38 (ENOSYS). libfabric's cxi provider +# uses only the FIRST VNI listed, so each task below restricts +# SLINGSHOT_VNIS to the job VNI (last entry) so core and extra share it. +# +# Runs on Frontier and Perlmutter; the rank layout is picked from the machine +# (see PLATFORM below). The #SBATCH lines above are for Frontier. On +# Perlmutter, override them on the command line: +# sbatch -A -C gpu --gpus-per-node=4 examples/vae/script/job-vae-core-extra.sh + +usage() { + cat < 0 + switches to ThreadDataLoader with that many worker threads. + Default: 0. + --replicate=R Repeat the MNIST training set R times on the core side + (longer epochs, same per-sample cost). Default: 1. + --image-scale=S Upscale images to (28*S)x(28*S) on the core side (the + extra side follows the published size). Default: 1. + -h, --help Show this help message and exit. +EOF +} + +for arg in "$@"; do + case "$arg" in + -h|--help) usage; exit 0 ;; + esac +done + +METHOD= +FABRIC= +GPUDIRECT=0 +LAYOUT= +CORE_NNODES_OPT= +NUM_WORKERS= +REPLICATE= +IMAGE_SCALE= +for arg in "$@"; do + case "$arg" in + --method=*) METHOD="${arg#--method=}" ;; + --fabric=*) FABRIC="${arg#--fabric=}" ;; + --gpudirect) GPUDIRECT=1 ;; + --layout=*) LAYOUT="${arg#--layout=}" ;; + --core-nnodes=*) CORE_NNODES_OPT="${arg#--core-nnodes=}" ;; + --num-workers=*) NUM_WORKERS="${arg#--num-workers=}" ;; + --replicate=*) REPLICATE="${arg#--replicate=}" ;; + --image-scale=*) IMAGE_SCALE="${arg#--image-scale=}" ;; + esac +done +LAYOUT="${LAYOUT:-split-node}" +NUM_WORKERS="${NUM_WORKERS:-0}" +REPLICATE="${REPLICATE:-1}" +IMAGE_SCALE="${IMAGE_SCALE:-1}" + +export DDSTORE_FABRIC="${FABRIC:-cxi}" +export DDSTORE_METHOD="${METHOD:-2}" + +if [ "$GPUDIRECT" == "1" ]; then + CORE_GPUS_PER_TASK=1 + CORE_EXTRA_ARGS="--gpu-source" + EXTRA_EXTRA_ARGS="--gpu-dest" +else + CORE_GPUS_PER_TASK=0 + CORE_EXTRA_ARGS="" + EXTRA_EXTRA_ARGS="" +fi + +# Per-task wrapper: use only the job VNI (see the header comment). +JOB_VNI_WRAP='export SLINGSHOT_VNIS=${SLINGSHOT_VNIS##*,}; exec "$@"' + +rm -rf ddstore_hs_vae +mkdir -p results +sleep 2 + +if [ "$LAYOUT" == "colocate" ]; then + CORE_NNODES=$SLURM_NNODES + EXTRA_NNODES=$SLURM_NNODES + # Both steps share the nodes; Perlmutter needs --overlap for that. + OVERLAP=--overlap +else + CORE_NNODES="${CORE_NNODES_OPT:-1}" + EXTRA_NNODES=$((SLURM_NNODES - CORE_NNODES)) +fi +# Perlmutter: the extra (training) ranks see all 4 GPUs of their node and +# vae_extra_train.py picks cuda:$SLURM_LOCALID. With --gpus-per-task=1 NCCL +# (2.29, pytorch/2.13.0) fails in DDP init with "Cuda failure 101 'invalid +# device ordinal'". The core step does no NCCL, so --gpus-per-task is fine +# there; 4 core ranks per node fit the 4 GPUs with --gpudirect. +PLATFORM="${NERSC_HOST:-${LMOD_SYSTEM_NAME:-frontier}}" +case "$PLATFORM" in + perlmutter) + CORE_NR=4 + EXTRA_NR=4 + EXTRA_CPUS_PER_TASK=16 + EXTRA_GPU_ARGS="--gpus-per-node=4" ;; + *) + CORE_NR=8 + EXTRA_NR=8 + EXTRA_CPUS_PER_TASK=6 + EXTRA_GPU_ARGS="--gpus-per-task=1" ;; +esac +CORE_NTASKS=$((CORE_NNODES * CORE_NR)) +EXTRA_NTASKS=$((EXTRA_NNODES * EXTRA_NR)) + +echo "PLATFORM=$PLATFORM DDSTORE_METHOD=$DDSTORE_METHOD DDSTORE_FABRIC=$DDSTORE_FABRIC LAYOUT=$LAYOUT GPUDIRECT=$GPUDIRECT" +echo "CORE_NNODES=$CORE_NNODES CORE_NTASKS=$CORE_NTASKS CORE_GPUS_PER_TASK=$CORE_GPUS_PER_TASK CORE_EXTRA_ARGS=\"$CORE_EXTRA_ARGS\" REPLICATE=$REPLICATE IMAGE_SCALE=$IMAGE_SCALE" +echo "EXTRA_NNODES=$EXTRA_NNODES EXTRA_NTASKS=$EXTRA_NTASKS EXTRA_EXTRA_ARGS=\"$EXTRA_EXTRA_ARGS\" NUM_WORKERS=$NUM_WORKERS" + +MASTER_PORT=8889 srun ${OVERLAP:-} -N$CORE_NNODES -n$CORE_NTASKS -c1 --gpus-per-task=$CORE_GPUS_PER_TASK --cpu-bind=verbose,core -l \ + bash -c "$JOB_VNI_WRAP" _ python -u examples/vae/vae_core_server.py ddstore_hs_vae --replicate=$REPLICATE --image-scale=$IMAGE_SCALE $CORE_EXTRA_ARGS \ + > >(sed 's/^/[core] /') 2> >(sed 's/^/[core] /') & +sleep 5 + +MASTER_PORT=8891 DDSTORE_HANDSHAKE_TIMEOUT_S=60 srun ${OVERLAP:-} -N$EXTRA_NNODES -n$EXTRA_NTASKS -c$EXTRA_CPUS_PER_TASK $EXTRA_GPU_ARGS --cpu-bind=verbose,core -l \ + bash -c "$JOB_VNI_WRAP" _ python -u examples/vae/vae_extra_train.py --handshake-dir ddstore_hs_vae --n-core $CORE_NTASKS --epochs 3 --num-workers=$NUM_WORKERS $EXTRA_EXTRA_ARGS \ + > >(sed 's/^/[extr] /') 2> >(sed 's/^/[extr] /') +sleep 5 + +wait diff --git a/examples/vae/script/job-vae-single.sh b/examples/vae/script/job-vae-single.sh new file mode 100755 index 0000000..18848c4 --- /dev/null +++ b/examples/vae/script/job-vae-single.sh @@ -0,0 +1,113 @@ +#!/bin/bash +#SBATCH -A FUS184 +#SBATCH -J GX-single +#SBATCH -o job-%j.out +#SBATCH -e job-%j.out +#SBATCH -N 4 +#SBATCH -t 30:00 +#SBATCH -q debug +# +# Baseline VAE DDP run. Runs on Frontier and Perlmutter; the rank layout is +# picked from the machine (see PLATFORM below). The #SBATCH lines above are +# for Frontier. On Perlmutter, override them on the command line: +# sbatch -A -C gpu --gpus-per-node=4 examples/vae/script/job-vae-single.sh + +usage() { + cat < 0 switches to ThreadDataLoader with that many worker + threads (requires --method=1 or 2). Default: 0. + --replicate=R Repeat the MNIST training set R times (longer epochs, + same per-sample cost). Default: 1. + --image-scale=S Upscale images to (28*S)x(28*S): S^2 larger samples. + Default: 1. + --ranks-per-node=N Default: 8 on Frontier, 4 on Perlmutter. + --cpus-per-task=N Default: 7 on Frontier, 32 on Perlmutter. + -h, --help Show this help message and exit. + +Examples: + $(basename "$0") # method=0, cxi + $(basename "$0") --method=1 --gpudirect # GPUDirect over libfabric + $(basename "$0") --method=2 --gpudirect # GPUDirect over file-based handshake + $(basename "$0") --fabric=hsn # baseline over hsn instead + $(basename "$0") --method=1 --num-workers=4 # ThreadDataLoader, no GPU buffers + $(basename "$0") --method=1 --gpudirect --num-workers=4 # GPUDirect + ThreadDataLoader +EOF +} + +for arg in "$@"; do + case "$arg" in + -h|--help) usage; exit 0 ;; + esac +done + +rm -rf ddstore_hs* +mkdir -p results +sleep 2 + +export VAE_PROFILE=1 + +METHOD= +FABRIC= +GPUDIRECT_ARGS="" +NUM_WORKERS= +REPLICATE= +IMAGE_SCALE= +RANKS_PER_NODE= +CPUS_PER_TASK= +for arg in "$@"; do + case "$arg" in + --method=*) METHOD="${arg#--method=}" ;; + --fabric=*) FABRIC="${arg#--fabric=}" ;; + --gpudirect) GPUDIRECT_ARGS="--gpu-dest --gpu-source" ;; + --num-workers=*) NUM_WORKERS="${arg#--num-workers=}" ;; + --replicate=*) REPLICATE="${arg#--replicate=}" ;; + --image-scale=*) IMAGE_SCALE="${arg#--image-scale=}" ;; + --ranks-per-node=*) RANKS_PER_NODE="${arg#--ranks-per-node=}" ;; + --cpus-per-task=*) CPUS_PER_TASK="${arg#--cpus-per-task=}" ;; + esac +done + +export DDSTORE_FABRIC="${FABRIC:-cxi}" +METHOD="${METHOD:-0}" +NUM_WORKERS="${NUM_WORKERS:-0}" +REPLICATE="${REPLICATE:-1}" +IMAGE_SCALE="${IMAGE_SCALE:-1}" + +EXTRA_ARGS="$GPUDIRECT_ARGS --num-workers=$NUM_WORKERS --replicate=$REPLICATE --image-scale=$IMAGE_SCALE" + +echo "DDSTORE_METHOD=$METHOD DDSTORE_FABRIC=$DDSTORE_FABRIC EXTRA_ARGS=\"$EXTRA_ARGS\"" + +# Perlmutter: every rank sees all 4 GPUs of its node and vae-ddp.py picks +# cuda:$SLURM_LOCALID. With --gpus-per-task=1 each rank sees only its own GPU +# and NCCL (2.29, pytorch/2.13.0) fails in DDP init with "Cuda failure 101 +# 'invalid device ordinal'" (transport/p2p.cc). +PLATFORM="${NERSC_HOST:-${LMOD_SYSTEM_NAME:-frontier}}" +case "$PLATFORM" in + perlmutter) + RANKS_PER_NODE="${RANKS_PER_NODE:-4}" + CPUS_PER_TASK="${CPUS_PER_TASK:-32}" + GPU_ARGS="--gpus-per-node=4" ;; + *) + RANKS_PER_NODE="${RANKS_PER_NODE:-8}" + CPUS_PER_TASK="${CPUS_PER_TASK:-7}" + GPU_ARGS="--gpus-per-task=1" ;; +esac +NNODES="${SLURM_NNODES:-1}" + +echo "PLATFORM=$PLATFORM NNODES=$NNODES RANKS_PER_NODE=$RANKS_PER_NODE CPUS_PER_TASK=$CPUS_PER_TASK GPU_ARGS=$GPU_ARGS" + +DDSTORE_METHOD=$METHOD srun -N$NNODES -n$((NNODES*RANKS_PER_NODE)) -c$CPUS_PER_TASK $GPU_ARGS -l \ + python -u examples/vae/vae-ddp.py --epochs 3 $EXTRA_ARGS \ + > >(sed 's/^/[core] /') 2> >(sed 's/^/[core] /') diff --git a/examples/vae/vae-ddp.py b/examples/vae/vae-ddp.py index f7bed6e..5fc2891 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -4,6 +4,7 @@ ## Do not reorder these imports. import argparse import os +import time import torch import torch.utils.data from torch import optim @@ -11,17 +12,12 @@ from torchvision.utils import save_image import torch.distributed as dist -import mpi4py - -mpi4py.rc.thread_level = "serialized" -mpi4py.rc.threads = False from mpi4py import MPI -import distdataset -from distdataset import DistDataset +from pyddstore.torch import DistDataset, ThreadDataLoader from ddp_utils import setup_ddp, get_local_rank -from vae_model import VAE, loss_function +from vae_model import VAE, loss_function, mnist_transform parser = argparse.ArgumentParser(description="VAE MNIST Example") parser.add_argument( @@ -54,6 +50,52 @@ metavar="N", help="how many batches to wait before logging training status", ) +parser.add_argument( + "--gpu-dest", + action="store_true", + default=False, + help="Allocate the DDStore get() destination buffer directly on the " + "training device (GPUDirect RDMA, Phase 1), skipping the " + "host->device copy. Requires DDSTORE_METHOD in (1, 2) and " + "DDSTORE_FABRIC=cxi (e.g. job-vae-single.sh --method=1 --gpudirect).", +) +parser.add_argument( + "--gpu-source", + action="store_true", + default=False, + help="Stack this rank's local shard directly on the training device " + "and add() it in place (GPUDirect RDMA source, Phase 2), skipping " + "the host round-trip. Same DDSTORE_METHOD/DDSTORE_FABRIC " + "requirements as --gpu-dest; independent of it -- use either or " + "both.", +) +parser.add_argument( + "--num-workers", + type=int, + default=0, + metavar="N", + help="Number of DataLoader workers. 0 uses PyTorch's standard " + "DataLoader in the main process (no worker processes, no fork). " + "> 0 switches to ThreadDataLoader " + "(pyddstore.torch), with that many worker threads " + "-- forked processes can't safely own GPU state or MPI's live state, " + "so any --num-workers > 0 goes through threads, never a fork. " + "Requires DDSTORE_METHOD 1 or 2 when > 0. Default: 0.", +) +parser.add_argument( + "--replicate", + type=int, + default=1, + metavar="R", + help="Repeat the MNIST training set this many times (torch ConcatDataset), for longer epochs with the same per-sample cost. Default: 1.", +) +parser.add_argument( + "--image-scale", + type=int, + default=1, + metavar="S", + help="Upscale MNIST images to (28*S)x(28*S) (bilinear), so each sample is S^2 times larger; the VAE hidden layer grows to 400*S. Default: 1 (plain 28x28).", +) args = parser.parse_args() args.cuda = not args.no_cuda and torch.cuda.is_available() use_mps = not args.no_mps and torch.backends.mps.is_available() @@ -82,46 +124,92 @@ else: device = torch.device("cpu") -print("DDP setup:", comm_size, rank, device) +print( + "DDP setup:", + comm_size, + rank, + device, + "gpu_dest:", + args.gpu_dest, + "gpu_source:", + args.gpu_source, +) if rank == 0: os.makedirs("results", exist_ok=True) comm.Barrier() -model = VAE().to(device) +side = 28 * args.image_scale +model = VAE(input_dim=side * side, hidden=400 * args.image_scale).to(device) model = torch.nn.parallel.DistributedDataParallel(model) optimizer = optim.Adam(model.parameters(), lr=1e-3) -# kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} -# kwargs = {'pin_memory': True} if args.cuda else {} -kwargs = {} - +mnist_train = datasets.MNIST( + "data", train=True, download=True, transform=mnist_transform(args.image_scale) +) trainset = DistDataset( - datasets.MNIST("data", train=True, download=True, transform=transforms.ToTensor()), + torch.utils.data.ConcatDataset([mnist_train] * args.replicate), "train", comm, + device=device if args.gpu_dest else None, + add_device=device if args.gpu_source else None, ) # trainset = datasets.MNIST('data', train=True, download=True,transform=transforms.ToTensor()) comm.Barrier() sampler = torch.utils.data.distributed.DistributedSampler(trainset) -train_loader = torch.utils.data.DataLoader( - trainset, batch_size=args.batch_size, shuffle=False, **kwargs, sampler=sampler +if args.num_workers > 0: + if int(os.environ.get("DDSTORE_METHOD", "0")) == 0: + raise RuntimeError("--num-workers > 0 requires DDSTORE_METHOD=1 or 2") + train_loader = ThreadDataLoader( + trainset, + batch_size=args.batch_size, + shuffle=False, + sampler=sampler, + num_workers=args.num_workers, + ) +else: + train_loader = torch.utils.data.DataLoader( + trainset, batch_size=args.batch_size, shuffle=False, sampler=sampler + ) + +print( + f"train_loader: {type(train_loader).__name__}, num_workers={train_loader.num_workers}" ) testset = datasets.MNIST( - "data", train=False, download=True, transform=transforms.ToTensor() + "data", train=False, download=True, transform=mnist_transform(args.image_scale) ) test_loader = torch.utils.data.DataLoader( - testset, batch_size=args.batch_size, shuffle=False, **kwargs + testset, batch_size=args.batch_size, shuffle=False ) +# VAE_PROFILE=1 splits each epoch's wall time into "fetch" (time spent +# inside the DataLoader producing a batch -- __getitem__/get()/collate) vs +# "compute" (forward/backward/optimizer.step()), to see whether a data- +# loading change (e.g. --gpu-dest/--gpu-source) is actually moving the +# needle relative to the rest of the step, rather than guessing from total +# wall time alone. Off by default -- adds one time.perf_counter() pair per +# batch, negligible but not zero. +PROFILE = os.environ.get("VAE_PROFILE") == "1" + + def train(epoch): model.train() train_loss = 0 + fetch_time = 0.0 + compute_time = 0.0 train_loader.dataset.ddstore.epoch_begin() - for batch_idx, (data, _) in enumerate(train_loader): + data_iter = iter(train_loader) + batch_idx = 0 + while True: + t0 = time.perf_counter() + try: + data, _ = next(data_iter) + except StopIteration: + break + t1 = time.perf_counter() train_loader.dataset.ddstore.epoch_end() # print(rank, device) data = data.to(device) @@ -137,6 +225,14 @@ def train(epoch): # print(rank, "backward") optimizer.step() # print(rank, "step") + # Skip epoch 1: CUDA/HIP kernel compilation, MIOpen/cuDNN algo + # selection, and allocator warmup make it dominated by one-time + # costs unrelated to steady-state fetch/compute timing. + if PROFILE and epoch > 1: + torch.cuda.synchronize(device=device) + t2 = time.perf_counter() + fetch_time += t1 - t0 + compute_time += t2 - t1 if batch_idx % args.log_interval == 0: print( "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( @@ -149,6 +245,7 @@ def train(epoch): ) train_loader.dataset.ddstore.epoch_begin() + batch_idx += 1 train_loader.dataset.ddstore.epoch_end() if rank == 0: @@ -157,6 +254,13 @@ def train(epoch): epoch, train_loss / len(train_loader.dataset) ) ) + if PROFILE and epoch > 1: + print( + "[profile] epoch {}: fetch={:.3f}s compute={:.3f}s".format( + epoch, fetch_time, compute_time + ), + flush=True, + ) def test(epoch): @@ -170,7 +274,7 @@ def test(epoch): if i == 0: n = min(data.size(0), 8) comparison = torch.cat( - [data[:n], recon_batch.view(args.batch_size, 1, 28, 28)[:n]] + [data[:n], recon_batch.view(-1, 1, side, side)[:n]] ) save_image( comparison.cpu(), @@ -192,7 +296,64 @@ def test(epoch): sample = torch.randn(64, 20).to(device) sample = model.module.decode(sample).cpu() save_image( - sample.view(64, 1, 28, 28), "results/sample_" + str(epoch) + ".png" + sample.view(64, 1, side, side), + "results/sample_" + str(epoch) + ".png", ) + # DDSTORE_PROFILE=1: where get() time goes, summed over all ranks and + # both variables (data + labels), all epochs. + if os.environ.get("DDSTORE_PROFILE", "0") not in ("", "0"): + ds = trainset.ddstore + tot = {} + for var in ("traindata", "trainlabels"): + for k, v in ds.get_profile(var).items(): + if not k.startswith("py_"): + tot[k] = tot.get(k, 0) + v + prof = ds.get_profile("traindata") + for k in ("py_gets", "py_get", "py_sync"): + tot[k] = prof[k] + tot = {k: comm.allreduce(v) for k, v in tot.items()} + if rank == 0: + n = max(tot["calls"], 1) + us = lambda x: 1e6 * x / n + us_row = lambda x: 1e6 * x / max(tot["rows"], 1) + other = ( + tot["py_get"] + - tot["py_sync"] + - tot["lock_wait"] + - tot["mr"] + - tot["read"] + - tot["cq"] + ) + print( + "[ddstore-profile] all ranks: calls={} rows={} py_calls={} mr_miss={} ({:.1%})".format( + tot["calls"], + tot["rows"], + tot["py_gets"], + tot["mr_miss"], + tot["mr_miss"] / n, + ) + ) + print( + "[ddstore-profile] per call (us): total={:.1f} sync={:.1f} lock_wait={:.1f} " + "mr={:.1f} read={:.1f} cq={:.1f} other={:.1f}".format( + us(tot["py_get"]), + us(tot["py_sync"]), + us(tot["lock_wait"]), + us(tot["mr"]), + us(tot["read"]), + us(tot["cq"]), + us(other), + ), + flush=True, + ) + print( + "[ddstore-profile] per row (us): total={:.2f} sync={:.2f} read+cq={:.2f}".format( + us_row(tot["py_get"]), + us_row(tot["py_sync"]), + us_row(tot["read"] + tot["cq"]), + ), + flush=True, + ) + dist.destroy_process_group() diff --git a/examples/vae/vae_core_server.py b/examples/vae/vae_core_server.py index 5956f8e..2728eaa 100644 --- a/examples/vae/vae_core_server.py +++ b/examples/vae/vae_core_server.py @@ -8,7 +8,19 @@ is done. Usage: - srun -n python examples/vae/vae_core_server.py [handshake_dir] + srun -n python examples/vae/vae_core_server.py [handshake_dir] [--gpu-source] [--replicate=R] [--image-scale=S] + + --gpu-source Stack this rank's shard directly on the GPU and add() it in + place (GPUDirect RDMA source, Phase 2), skipping the host + round-trip. Requires DDSTORE_METHOD=2 (already set below) + and DDSTORE_FABRIC=cxi, and one visible GPU per rank + (--gpus-per-task=1, unlike the --gpus-per-task=0 this + script normally runs with). + --replicate=R Repeat the MNIST training set R times (default 1), for + longer epochs with the same per-sample cost. The extra side + picks up the row count from the published variable. + --image-scale=S Upscale images to (28*S)x(28*S) (default 1). The extra + side derives S from the published row width. Environment: DDSTORE_HANDSHAKE_DIR overrides handshake_dir positional arg @@ -23,19 +35,17 @@ import sys import time -## torch (pulled in below via torchvision/distdataset) must finish loading +## torch (pulled in below via torchvision/pyddstore.torch) must finish loading ## before mpi4py triggers MPI_Init, or - if GPU/NCCL use is ever added here - ## their static destructors run in the wrong order at interpreter exit and ## corrupt the heap. Do not reorder these imports. +import torch from torchvision import datasets, transforms -import mpi4py - -mpi4py.rc.thread_level = "serialized" -mpi4py.rc.threads = False from mpi4py import MPI -from distdataset import DistDataset +from pyddstore.torch import DistDataset +from vae_model import mnist_transform def _resolve_dir(arg): @@ -44,27 +54,58 @@ def _resolve_dir(arg): return os.environ.get("DDSTORE_HANDSHAKE_DIR", "./ddstore_hs") -hs_dir = _resolve_dir(sys.argv[1] if len(sys.argv) > 1 else "") +gpu_source = "--gpu-source" in sys.argv +replicate = 1 +image_scale = 1 +for a in sys.argv[1:]: + if a.startswith("--replicate="): + replicate = int(a.split("=", 1)[1]) + elif a.startswith("--image-scale="): + image_scale = int(a.split("=", 1)[1]) +positional_args = [a for a in sys.argv[1:] if not a.startswith("--")] +hs_dir = _resolve_dir(positional_args[0] if positional_args else "") os.environ["DDSTORE_METHOD"] = "2" os.environ["DDSTORE_HANDSHAKE_DIR"] = hs_dir comm = MPI.COMM_WORLD rank = comm.Get_rank() +add_device = None +if gpu_source: + if not torch.cuda.is_available(): + raise RuntimeError("--gpu-source requires a visible CUDA/HIP GPU") + add_device = torch.device("cuda") + if rank == 0: os.makedirs(hs_dir, exist_ok=True) for fname in os.listdir(hs_dir): - if fname.endswith(".bin") or fname == "done_extra": + if fname.endswith((".bin", ".meta.json")) or fname == "done_extra": os.remove(os.path.join(hs_dir, fname)) print(f"[core] handshake_dir={hs_dir}", flush=True) comm.Barrier() trainset = datasets.MNIST( - "data", train=True, download=True, transform=transforms.ToTensor() + "data", train=True, download=True, transform=mnist_transform(image_scale) +) +dds_trainset = DistDataset( + torch.utils.data.ConcatDataset([trainset] * replicate), + "train", + comm, + add_device=add_device, ) -dds_trainset = DistDataset(trainset, "train", comm) comm.Barrier() +if rank == 0: + print( + "gpu_source:", + gpu_source, + "replicate:", + replicate, + "image_scale:", + image_scale, + flush=True, + ) + if rank == 0: print( f"[core] published {len(dds_trainset)} training rows, waiting for extras...", @@ -87,7 +128,7 @@ def _resolve_dir(arg): if rank == 0: for fname in os.listdir(hs_dir): - if fname.endswith(".bin"): + if fname.endswith((".bin", ".meta.json")): try: os.remove(os.path.join(hs_dir, fname)) except OSError: diff --git a/examples/vae/vae_extra_train.py b/examples/vae/vae_extra_train.py index 9cda3a2..967764a 100644 --- a/examples/vae/vae_extra_train.py +++ b/examples/vae/vae_extra_train.py @@ -36,15 +36,11 @@ from torchvision.utils import save_image import torch.distributed as dist -import mpi4py - -mpi4py.rc.thread_level = "serialized" -mpi4py.rc.threads = False from mpi4py import MPI from ddp_utils import setup_ddp, get_local_rank -from distdataset import DistDatasetReader -from vae_model import VAE, loss_function +from pyddstore.torch import DistDatasetReader, ThreadDataLoader +from vae_model import VAE, loss_function, mnist_transform parser = argparse.ArgumentParser(description="VAE MNIST Example - extra (reader) group") parser.add_argument( @@ -89,6 +85,27 @@ default=int(os.environ.get("DDSTORE_N_CORE", "4")), help="number of core ranks that published the data", ) +parser.add_argument( + "--gpu-dest", + action="store_true", + default=False, + help="Allocate the DDStore get() destination buffer directly on the " + "training device (GPUDirect RDMA, Phase 1), skipping the " + "host->device copy. Requires DDSTORE_FABRIC=cxi and a " + "libfabric-backed method (already the case for this script).", +) +parser.add_argument( + "--num-workers", + type=int, + default=0, + metavar="N", + help="Number of DataLoader workers. 0 uses PyTorch's standard " + "DataLoader in the main process (no worker processes, no fork). " + "> 0 switches to ThreadDataLoader " + "(pyddstore.torch), with that many worker threads " + "-- forked processes can't safely own GPU state, so any " + "--num-workers > 0 goes through threads, never a fork. Default: 0.", +) args = parser.parse_args() args.cuda = not args.no_cuda and torch.cuda.is_available() use_mps = not args.no_mps and torch.backends.mps.is_available() @@ -115,26 +132,48 @@ else: device = torch.device("cpu") -print("DDP setup:", comm_size, rank, device) +print("DDP setup:", comm_size, rank, device, "gpu_dest:", args.gpu_dest) -model = VAE().to(device) +trainset = DistDatasetReader( + "train", + args.handshake_dir, + args.n_core, + device=device if args.gpu_dest else None, +) + +# Image size comes from the core side's published data (vae_core_server.py +# --image-scale); the model and test set must match it. +side = trainset.shapes[0][-1] # samples are (image (1, side, side), label) +image_scale = side // 28 +model = VAE(input_dim=side * side, hidden=400 * image_scale).to(device) model = torch.nn.parallel.DistributedDataParallel(model) optimizer = optim.Adam(model.parameters(), lr=1e-3) - -kwargs = {} - -trainset = DistDatasetReader("train", args.handshake_dir, args.n_core) sampler = torch.utils.data.distributed.DistributedSampler(trainset) -train_loader = torch.utils.data.DataLoader( - trainset, batch_size=args.batch_size, shuffle=False, **kwargs, sampler=sampler +if args.num_workers > 0: + # DistDatasetReader always joins via method=2 (file-based handshake), + # so no DDSTORE_METHOD=0 guard is needed here (unlike vae-ddp.py). + train_loader = ThreadDataLoader( + trainset, + batch_size=args.batch_size, + shuffle=False, + sampler=sampler, + num_workers=args.num_workers, + ) +else: + train_loader = torch.utils.data.DataLoader( + trainset, batch_size=args.batch_size, shuffle=False, sampler=sampler + ) + +print( + f"train_loader: {type(train_loader).__name__}, num_workers={train_loader.num_workers}" ) testset = datasets.MNIST( - "data", train=False, download=True, transform=transforms.ToTensor() + "data", train=False, download=True, transform=mnist_transform(image_scale) ) test_loader = torch.utils.data.DataLoader( - testset, batch_size=args.batch_size, shuffle=False, **kwargs + testset, batch_size=args.batch_size, shuffle=False ) @@ -183,7 +222,7 @@ def test(epoch): if i == 0: n = min(data.size(0), 8) comparison = torch.cat( - [data[:n], recon_batch.view(args.batch_size, 1, 28, 28)[:n]] + [data[:n], recon_batch.view(-1, 1, side, side)[:n]] ) save_image( comparison.cpu(), @@ -203,7 +242,7 @@ def test(epoch): sample = torch.randn(64, 20).to(device) sample = model.module.decode(sample).cpu() save_image( - sample.view(64, 1, 28, 28), + sample.view(64, 1, side, side), "results/extra_sample_" + str(epoch) + ".png", ) diff --git a/examples/vae/vae_model.py b/examples/vae/vae_model.py index b8f435e..1759c0b 100644 --- a/examples/vae/vae_model.py +++ b/examples/vae/vae_model.py @@ -1,17 +1,19 @@ import torch from torch import nn from torch.nn import functional as F +from torchvision import transforms class VAE(nn.Module): - def __init__(self): + def __init__(self, input_dim=784, hidden=400): super(VAE, self).__init__() - self.fc1 = nn.Linear(784, 400) - self.fc21 = nn.Linear(400, 20) - self.fc22 = nn.Linear(400, 20) - self.fc3 = nn.Linear(20, 400) - self.fc4 = nn.Linear(400, 784) + self.input_dim = input_dim + self.fc1 = nn.Linear(input_dim, hidden) + self.fc21 = nn.Linear(hidden, 20) + self.fc22 = nn.Linear(hidden, 20) + self.fc3 = nn.Linear(20, hidden) + self.fc4 = nn.Linear(hidden, input_dim) def encode(self, x): h1 = F.relu(self.fc1(x)) @@ -27,14 +29,14 @@ def decode(self, z): return torch.sigmoid(self.fc4(h3)) def forward(self, x): - mu, logvar = self.encode(x.view(-1, 784)) + mu, logvar = self.encode(x.view(-1, self.input_dim)) z = self.reparameterize(mu, logvar) return self.decode(z), mu, logvar def loss_function(recon_x, x, mu, logvar): # Reconstruction + KL divergence losses summed over all elements and batch - BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction="sum") + BCE = F.binary_cross_entropy(recon_x, x.view(-1, recon_x.shape[1]), reduction="sum") # see Appendix B from VAE paper: # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014 @@ -43,3 +45,10 @@ def loss_function(recon_x, x, mu, logvar): KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) return BCE + KLD + + +def mnist_transform(scale): + """ToTensor(), preceded by a bilinear upscale to (28*scale)^2 if scale > 1.""" + if scale == 1: + return transforms.ToTensor() + return transforms.Compose([transforms.Resize(28 * scale), transforms.ToTensor()]) diff --git a/images/DDStore-logo.png b/images/DDStore-logo.png new file mode 100644 index 0000000..e2d4780 Binary files /dev/null and b/images/DDStore-logo.png differ diff --git a/include/common.h b/include/common.h index 18318ef..742f207 100644 --- a/include/common.h +++ b/include/common.h @@ -4,7 +4,10 @@ #include #include #include +#include +#include #include +#include #include #define DP_AV_DEF_SIZE 512 @@ -33,6 +36,15 @@ extern "C" { #endif + /* One caller-registered recv buffer (see fabric_state::pinned). */ + struct recv_region + { + struct fid_mr *mr; + char *base; + size_t len; + int hmem_iface; + }; + struct fabric_state { struct fi_context *ctx; @@ -46,18 +58,104 @@ extern "C" fi_addr_t *comm_partner; char *send_data; size_t send_data_len; + /* FI_HMEM_SYSTEM (0) if send_data is host memory; otherwise the + * fi_hmem_iface value identifying what kind of GPU memory it is. + * Set by the caller (see ddstore.hpp's add()) and forwarded into + * fi_mr_regattr()'s attr.iface in handshake() (method=1) / add()'s + * inline registration (method=2). A separate field from + * recv_hmem_iface: one fabric_state can be simultaneously the send + * side (registered once at add() time) and the recv side + * (re-registered per get() call, including self-reads) -- these + * are independent MRs with independent lifetimes. */ + int send_hmem_iface; char *recv_data; size_t recv_data_len; + /* FI_HMEM_SYSTEM (0) if recv_data is host memory; otherwise the + * fi_hmem_iface value (FI_HMEM_CUDA, FI_HMEM_ROCR, ...) identifying + * what kind of GPU memory it is. Set by the caller (see + * ddstore.hpp's get()); the value is opaque here, just forwarded + * into fi_mr_regattr()'s attr.iface in read_from_remote(). */ + int recv_hmem_iface; struct fid_mr *mr; struct fid_mr *recv_mr; + /* Cached recv-side MR region: the registered range is + * [recv_mr_base, recv_mr_base + recv_mr_reg_len). Any recv_data + * pointer that falls within this range with recv_data_len bytes + * fitting inside it can reuse recv_mr without re-registration. + * + * Hits when the same buffer (or a sub-range of it) is passed again -- + * e.g. PyTorch's caching allocator returning the same block for a + * same-shape torch.empty() -- so get() doesn't re-register every call. + * + * Initialised to NULL/0 so the first call always registers. */ + char *recv_mr_base; + size_t recv_mr_reg_len; + /* Caller-registered recv regions (register_recv_region()), checked + * before the one-slot cache above and never evicted: the caller owns + * these buffers and keeps them alive until unregister / free(). */ + struct recv_region *pinned; + int n_pinned; + /* Largest single fi_read() the endpoint accepts (FI_OPT_MAX_MSG_SIZE + * or ep_attr->max_msg_size); 0 if unknown. Longer rows are split. */ + size_t max_msg_size; uint64_t key; uint64_t *remote_key; uint64_t *remote_address; int world_size; int rank; + + /* Serializes concurrent get() calls on THIS variable's + * fabric_state -- see fabric_state_lock_guard below. libfabric + * itself doesn't guarantee thread safety unless the domain is + * opened with FI_THREAD_SAFE (it isn't here; see + * init_fabric_hsn()'s FI_THREAD_DOMAIN hint and init_fabric_cxi()'s + * unconstrained NULL-hints query), and even then that would only + * cover libfabric's own objects, not the plain fields above + * (recv_data/recv_mr/recv_mr_base/recv_mr_reg_len) that read_from_ + * remote() reads and writes as a cache. + * Confirmed necessary by direct experiment: concurrent get() calls + * without this crashed with "double free or corruption". + * Zero-initialized by calloc() at every allocation site below, but + * explicitly pthread_mutex_init()'d right after each one anyway -- + * relying on zero-initialized pthread_mutex_t being equivalent to + * PTHREAD_MUTEX_INITIALIZER is a common but implementation-defined + * assumption; init explicitly instead. */ + pthread_mutex_t recv_lock; + + /* DDSTORE_PROFILE=1: cumulative get() timing for this variable, all + * updated while recv_lock is held (see ddstore_profile_enabled()). + * prof_lock_wait_ns is the time spent waiting to acquire recv_lock; + * mr = recv-MR cache check / (re)registration; read = posting + * fi_read(); cq = polling the CQ until the read completes. */ + uint64_t prof_calls; /* get() / get_batch() calls */ + uint64_t prof_rows; /* rows read by those calls */ + uint64_t prof_lock_wait_ns; + uint64_t prof_mr_ns; + uint64_t prof_mr_miss; + uint64_t prof_read_ns; + uint64_t prof_cq_ns; }; + static inline uint64_t ddstore_now_ns(void) + { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (uint64_t)ts.tv_sec * 1000000000ull + (uint64_t)ts.tv_nsec; + } + + /* True if DDSTORE_PROFILE is set to a non-"0" value (read once). */ + static inline bool ddstore_profile_enabled(void) + { + static int enabled = -1; + if (enabled < 0) + { + const char *e = getenv("DDSTORE_PROFILE"); + enabled = (e && e[0] && strcmp(e, "0") != 0) ? 1 : 0; + } + return enabled == 1; + } + static bool is_local_mr_req(struct fabric_state *f) { return (f->info->mode & FI_LOCAL_MR) != 0; @@ -66,11 +164,10 @@ extern "C" /* CXI (and some other providers) use FI_MR_ENDPOINT: after fi_mr_reg the * MR must be bound to the endpoint and enabled before it can be used, and * the key is only valid after fi_mr_enable(). - * On Perlmutter, fi_getinfo with NULL hints returns mr_mode=0 even for - * CXI, so we detect by provider name instead of mr_mode flags. False - * (no-op) for every provider dev-file2 already supports (hsn/verbs/ - * gni/psm2), since none of those set mr_mode & FI_MR_ENDPOINT and none - * are named "cxi". */ + * With NULL hints fi_getinfo returns mr_mode=0 even for CXI (seen on + * Perlmutter), so we detect by provider name instead of mr_mode flags. + * False (no-op) for hsn/verbs/gni/psm2, since none of those set + * mr_mode & FI_MR_ENDPOINT and none are named "cxi". */ static bool is_mr_endpoint(struct fabric_state *f) { return (f->info->domain_attr->mr_mode & FI_MR_ENDPOINT) != 0 || @@ -81,25 +178,51 @@ extern "C" /* With FI_MR_VIRT_ADDR the fi_read remote addr is the virtual address. * CXI does NOT use virtual addresses — offset is 0-based from MR base. * - * NOTE: this is deliberately NOT a mr_mode bit check. dev-file2's - * init_fabric_hsn() sets mr_mode to the legacy FI_MR_BASIC sentinel, + * NOTE: this is deliberately NOT a mr_mode bit check. init_fabric_hsn() + * sets mr_mode to the legacy FI_MR_BASIC sentinel, * which on this system's libfabric (2.3.1) is bit 0 (value 1) — a * completely different bit than FI_MR_VIRT_ADDR (bit 4). A `mr_mode & * FI_MR_VIRT_ADDR` check would therefore silently resolve to false for * hsn, breaking address exchange for the already-proven path. Before - * this helper existed, dev-file2 unconditionally used the real pointer - * for every provider it supported (hsn/verbs/gni/psm2) — no virt-addr/ - * prov-key distinction existed at all — so preserve that unconditional - * behavior for anything that isn't cxi. */ + * cxi support, the real pointer was used unconditionally for every + * provider (hsn/verbs/gni/psm2), so preserve that for anything that + * isn't cxi. */ static bool is_virt_addr(struct fabric_state *f) { return !(f->info->fabric_attr->prov_name && strcmp(f->info->fabric_attr->prov_name, "cxi") == 0); } + /* True only for the cxi provider (the real Slingshot/HW path). GPU + * (ROCr HMEM) buffer registration is only attempted when this is true — + * verified empirically that cxi's NULL-hints fi_getinfo() already + * reports FI_HMEM in caps by default on Frontier; hsn (tcp;ofi_rxm) + * has no such support and would otherwise fail with a confusing + * low-level libfabric error instead of a clear one. */ + static bool is_hmem_capable(struct fabric_state *f) + { + return f->info && f->info->fabric_attr->prov_name && + strcmp(f->info->fabric_attr->prov_name, "cxi") == 0; + } + void init_fabric(struct fabric_state *fabric); int handshake(struct fabric_state *fabric_state, MPI_Comm comm); int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset); + /* n rows of row_len bytes into recv_data (recv_data_len == n * row_len); + * row i from rank src[i] at byte offset offset[i]. All reads are posted + * before any is waited for. 0 on success. See common.cxx. */ + int read_batch_from_remote(struct fabric_state *fabric_state, long n, + const int *src, const uint64_t *offset, size_t row_len); + /* Register [base, base + len) once as a recv buffer (hmem_iface as for + * recv_hmem_iface); reads into it then skip registration. Registering a + * region already registered is a no-op. 0 on success. Caller holds + * recv_lock and keeps the buffer alive until unregister / free. */ + int register_recv_region(struct fabric_state *fs, char *base, size_t len, int hmem_iface); + /* Undo register_recv_region() for the region starting at base. 0 on + * success, 1 if no such region. Caller holds recv_lock. */ + int unregister_recv_region(struct fabric_state *fs, char *base); + /* Close every registered recv region (free()). */ + void close_recv_regions(struct fabric_state *fs); /* --- Method 2: file-based handshake ---------------------------------- */ @@ -132,4 +255,24 @@ extern "C" #ifdef __cplusplus } + +/* RAII guard for struct fabric_state::recv_lock -- locks on construction, + * unlocks on destruction (including when leaving via an exception), so + * every exit path of the critical section it wraps is covered without + * having to hand-place lock/unlock calls on each one. See the comment on + * recv_lock above for what this protects and why. */ +struct fabric_state_lock_guard +{ + struct fabric_state *fs; + explicit fabric_state_lock_guard(struct fabric_state *fs) : fs(fs) + { + pthread_mutex_lock(&fs->recv_lock); + } + ~fabric_state_lock_guard() + { + pthread_mutex_unlock(&fs->recv_lock); + } + fabric_state_lock_guard(const fabric_state_lock_guard &) = delete; + fabric_state_lock_guard &operator=(const fabric_state_lock_guard &) = delete; +}; #endif diff --git a/include/ddstore.hpp b/include/ddstore.hpp index 2c582ee..0747152 100644 --- a/include/ddstore.hpp +++ b/include/ddstore.hpp @@ -1,5 +1,8 @@ #include #include +#include +#include +#include #include #include #include @@ -21,6 +24,10 @@ struct VarInfo bool active; bool fence_active; void *base; + /* true if base came from MPI_Alloc_mem() in add()/init() and free() + * must release it; false for a caller-owned GPU buffer (add() with + * hmem_iface != 0) or a joined variable (base == NULL). */ + bool owns_base; struct fabric_state *fabric_state; }; typedef struct VarInfo VarInfo_t; @@ -63,17 +70,64 @@ class DDStore /* Method 2 extra member: discover variable published by core members. */ void join(std::string name); + /* DDSTORE_PROFILE=1 counters for `name` (methods 1/2), as + * {calls, lock_wait_ns, mr_ns, mr_miss, read_ns, cq_ns, rows}; all zero + * for method 0 or when profiling is off. Takes the variable's lock. */ + void profile(std::string name, unsigned long long out[7]) + { + for (int i = 0; i < 7; i++) + out[i] = 0; + const VarInfo_t &varinfo = this->varlist.at(name); + struct fabric_state *fs = varinfo.fabric_state; + if (!fs) + return; + fabric_state_lock_guard lock(fs); + out[0] = fs->prof_calls; + out[1] = fs->prof_lock_wait_ns; + out[2] = fs->prof_mr_ns; + out[3] = fs->prof_mr_miss; + out[4] = fs->prof_read_ns; + out[5] = fs->prof_cq_ns; + out[6] = fs->prof_rows; + } + + /* hmem_iface: 0 (FI_HMEM_SYSTEM) for a host buffer, or an fi_hmem_iface + * value (FI_HMEM_CUDA, FI_HMEM_ROCR, ...) identifying what kind of GPU + * memory `buffer` is. Mirrors get()'s hmem_iface parameter. + * + * LIFETIME CONTRACT: for hmem_iface == 0 (host), DDStore makes its own + * private copy of `buffer` (as it always has) -- the caller's buffer + * may be freed/reused immediately after add() returns. For + * hmem_iface != 0 (GPU), DDStore does NOT copy -- it registers the + * caller's own device pointer directly. The caller must keep that GPU + * allocation alive (not garbage-collected, not reused) for as long as + * this variable stays registered, i.e. until free() or this DDStore's + * destruction. pyddstore/_core.pyx enforces this for Python callers via a + * keepalive dict; direct C++ callers must manage it themselves. */ template - void add(std::string name, T *buffer, long nrows, int disp) + void add(std::string name, T *buffer, long nrows, int disp, int hmem_iface = 0) { + if (this->method == 0 && hmem_iface != 0) + throw std::runtime_error("GPU source buffer is not supported with method=0 (MPI_Win)"); + void *base = NULL; - // (2025/03) jyc: necessary to avoid memory error - int err = MPI_Alloc_mem((MPI_Aint)(nrows * disp * sizeof(T)), MPI_INFO_NULL, &base); - if (err) + if (hmem_iface == 0) { - exit(1); + // (2025/03) jyc: necessary to avoid memory error + int err = MPI_Alloc_mem((MPI_Aint)(nrows * disp * sizeof(T)), MPI_INFO_NULL, &base); + if (err) + { + exit(1); + } + memcpy(base, buffer, nrows * disp * sizeof(T)); + } + else + { + /* GPU source buffer: register the caller's own pointer + * directly. No MPI_Alloc_mem, no copy -- see lifetime + * contract above. */ + base = (void *)buffer; } - memcpy(base, buffer, nrows * disp * sizeof(T)); MPI_Win win = MPI_WIN_NULL; struct fabric_state *fabric_state = NULL; @@ -90,40 +144,74 @@ class DDStore else if (this->method == 1) { fabric_state = (struct fabric_state *)calloc(1, sizeof(struct fabric_state)); - fabric_state->send_data = (char *)base; - fabric_state->send_data_len = nrows * disp * sizeof(T); - fabric_state->world_size = this->comm_size; - fabric_state->rank = this->rank; + pthread_mutex_init(&fabric_state->recv_lock, NULL); + fabric_state->send_data = (char *)base; + fabric_state->send_data_len = nrows * disp * sizeof(T); + fabric_state->send_hmem_iface = hmem_iface; + fabric_state->world_size = this->comm_size; + fabric_state->rank = this->rank; init_fabric(fabric_state); if (!fabric_state->info) throw std::runtime_error("init_fabric failed: no suitable fabric found"); + if (hmem_iface != 0 && !is_hmem_capable(fabric_state)) + throw std::runtime_error( + "GPU source buffer requires DDSTORE_FABRIC=cxi " + "(current fabric does not support FI_HMEM)"); if (handshake(fabric_state, this->comm) != 0) throw std::runtime_error("handshake failed (method=1)"); } else if (this->method == 2) { fabric_state = (struct fabric_state *)calloc(1, sizeof(struct fabric_state)); - fabric_state->send_data = (char *)base; - fabric_state->send_data_len = nrows * disp * sizeof(T); - fabric_state->world_size = this->n_core; - fabric_state->rank = this->rank; + pthread_mutex_init(&fabric_state->recv_lock, NULL); + fabric_state->send_data = (char *)base; + fabric_state->send_data_len = nrows * disp * sizeof(T); + fabric_state->send_hmem_iface = hmem_iface; + fabric_state->world_size = this->n_core; + fabric_state->rank = this->rank; init_fabric(fabric_state); if (!fabric_state->info) throw std::runtime_error("init_fabric failed: no suitable fabric found"); - - /* Register the send buffer as an MR before writing the record. */ - int mr_rc = fi_mr_reg( - fabric_state->domain, - fabric_state->send_data, - fabric_state->send_data_len, - FI_WRITE | FI_REMOTE_READ, - 0, 0, 0, - &fabric_state->mr, - NULL); + if (hmem_iface != 0 && !is_hmem_capable(fabric_state)) + throw std::runtime_error( + "GPU source buffer requires DDSTORE_FABRIC=cxi " + "(current fabric does not support FI_HMEM)"); + + /* Register the send buffer as an MR before writing the record -- + * same fi_mr_reg-vs-fi_mr_regattr branch as handshake(), + * duplicated here the same way the host-only version already + * is (method=2 registers inline instead of via handshake()). */ + bool send_is_hmem = hmem_iface != 0; + int mr_rc; + if (send_is_hmem) + { + struct iovec iov = {fabric_state->send_data, fabric_state->send_data_len}; + struct fi_mr_attr attr; + memset(&attr, 0, sizeof(attr)); + attr.mr_iov = &iov; + attr.iov_count = 1; + attr.access = FI_WRITE | FI_REMOTE_READ; + attr.iface = (enum fi_hmem_iface)hmem_iface; + attr.device.reserved = 0; + mr_rc = fi_mr_regattr(fabric_state->domain, &attr, 0, &fabric_state->mr); + } + else + { + mr_rc = fi_mr_reg( + fabric_state->domain, + fabric_state->send_data, + fabric_state->send_data_len, + FI_WRITE | FI_REMOTE_READ, + 0, 0, 0, + &fabric_state->mr, + NULL); + } if (mr_rc != FI_SUCCESS) - throw std::runtime_error(std::string("fi_mr_reg failed: ") + fi_strerror(mr_rc)); + throw std::runtime_error( + std::string(send_is_hmem ? "fi_mr_regattr failed: " : "fi_mr_reg failed: ") + + fi_strerror(mr_rc)); /* CXI (FI_MR_ENDPOINT): bind MR to endpoint and enable it before * use. The provider-assigned key is only valid after @@ -168,6 +256,7 @@ class DDStore var.active = true; var.fence_active = false; var.base = base; + var.owns_base = (hmem_iface == 0); var.fabric_state = fabric_state; this->varlist.insert(std::pair(name, var)); return; /* lenlist already stored; skip the MPI_Allgather block below */ @@ -198,6 +287,7 @@ class DDStore var.active = true; var.fence_active = false; var.base = base; + var.owns_base = (hmem_iface == 0); var.fabric_state = fabric_state; this->varlist.insert(std::pair(name, var)); @@ -230,6 +320,7 @@ class DDStore else if (this->method == 1) { fabric_state = (struct fabric_state *)calloc(1, sizeof(struct fabric_state)); + pthread_mutex_init(&fabric_state->recv_lock, NULL); fabric_state->send_data = (char *)base; fabric_state->send_data_len = nrows * disp * itemsize; fabric_state->world_size = this->comm_size; @@ -244,6 +335,7 @@ class DDStore else if (this->method == 2) { fabric_state = (struct fabric_state *)calloc(1, sizeof(struct fabric_state)); + pthread_mutex_init(&fabric_state->recv_lock, NULL); fabric_state->send_data = (char *)base; fabric_state->send_data_len = nrows * disp * itemsize; fabric_state->world_size = this->n_core; @@ -303,6 +395,7 @@ class DDStore var.active = true; var.fence_active = false; var.base = base; + var.owns_base = true; var.fabric_state = fabric_state; this->varlist.insert(std::pair(name, var)); return; @@ -333,6 +426,7 @@ class DDStore var.active = true; var.fence_active = false; var.base = base; + var.owns_base = true; var.fabric_state = fabric_state; this->varlist.insert(std::pair(name, var)); @@ -352,8 +446,13 @@ class DDStore memcpy((char*)base + offset * disp * itemsize, buffer, nrows * disp * itemsize); } + /* hmem_iface: 0 (FI_HMEM_SYSTEM) for a host buffer, or an fi_hmem_iface + * value (FI_HMEM_CUDA, FI_HMEM_ROCR, ...) identifying what kind of GPU + * memory `buffer` is. Left as a plain int (not the enum) so the Cython + * binding (pyddstore/_core.pyx) can pass it without cimporting the enum; + * read_from_remote() in common.cxx casts it back before use. */ template - void get(std::string name, long start, long count, T *buffer) + void get(std::string name, long start, long count, T *buffer, int hmem_iface = 0) { const VarInfo_t& varinfo = this->varlist.at(name); @@ -374,7 +473,11 @@ class DDStore // std::cout << "target,offset,start,count: " << target << "," << offset << "," << start << "," << count << // std::endl; - if (this->method == 0) + if (this->method == 0 && hmem_iface != 0) + { + throw std::runtime_error("GPU destination buffer is not supported with method=0 (MPI_Win)"); + } + else if (this->method == 0) { MPI_Win win = varinfo.win; MPI_Win_lock(MPI_LOCK_SHARED, target, 0, win); @@ -396,9 +499,27 @@ class DDStore } else if (this->method == 1 || this->method == 2) { - /* Methods 1 and 2 both use libfabric fi_read — same path. */ + /* Methods 1 and 2 both use libfabric fi_read — same path. + * Locked for the whole branch: the recv_data/recv_data_len/ + * recv_hmem_iface writes below are themselves racy across + * concurrent get() calls on this variable, not just the + * read_from_remote() call that follows them -- see recv_lock's + * comment in common.h. */ + const bool prof = ddstore_profile_enabled(); + uint64_t t_wait = prof ? ddstore_now_ns() : 0; + fabric_state_lock_guard lock(varinfo.fabric_state); + if (prof) + { + varinfo.fabric_state->prof_calls++; + varinfo.fabric_state->prof_lock_wait_ns += ddstore_now_ns() - t_wait; + } + if (hmem_iface != 0 && !is_hmem_capable(varinfo.fabric_state)) + throw std::runtime_error( + "GPU destination buffer requires DDSTORE_FABRIC=cxi " + "(current fabric does not support FI_HMEM)"); varinfo.fabric_state->recv_data = (char *)buffer; - varinfo.fabric_state->recv_data_len = varinfo.disp * varinfo.itemsize * count; + varinfo.fabric_state->recv_data_len = (size_t)varinfo.disp * varinfo.itemsize * count; + varinfo.fabric_state->recv_hmem_iface = hmem_iface; int rc = read_from_remote(varinfo.fabric_state, target, (start - offset) * varinfo.disp * varinfo.itemsize); if (rc != 0) throw std::runtime_error( @@ -407,8 +528,229 @@ class DDStore } } + /* Register `len` bytes at `buffer` once as a destination for get() / + * get_batch() of `name` (hmem_iface as in get()), so reads into it or + * any part of it skip memory registration. Several buffers can be + * registered per variable (e.g. one per loader thread); none is ever + * evicted. The caller keeps the buffer alive until unregister_recv() or + * free(). No-op for method 0. */ + void register_recv(std::string name, void *buffer, size_t len, int hmem_iface = 0) + { + const VarInfo_t &varinfo = this->varlist.at(name); + if (this->method == 0) + return; + if (hmem_iface != 0 && !is_hmem_capable(varinfo.fabric_state)) + throw std::runtime_error( + "GPU destination buffer requires DDSTORE_FABRIC=cxi " + "(current fabric does not support FI_HMEM)"); + fabric_state_lock_guard lock(varinfo.fabric_state); + if (register_recv_region(varinfo.fabric_state, (char *)buffer, len, hmem_iface) != 0) + throw std::runtime_error("register_recv failed for " + name); + } + + /* Undo register_recv() for the buffer starting at `buffer`. */ + void unregister_recv(std::string name, void *buffer) + { + const VarInfo_t &varinfo = this->varlist.at(name); + if (this->method == 0) + return; + fabric_state_lock_guard lock(varinfo.fabric_state); + if (unregister_recv_region(varinfo.fabric_state, (char *)buffer) != 0) + throw std::invalid_argument("buffer is not registered for " + name); + } + + /* Batched get: row i of `buffer` (n contiguous rows) receives global + * row idx[i]; rows may come from any ranks, in any order, with repeats. + * Every index is validated before anything is read. hmem_iface as in get(). + * + * Methods 1/2 (one-sided): take the variable's lock once and post all n + * fi_read()s before waiting for any (read_batch_from_remote()). + * + * Method 0 (collective, MDLoader-style): COLLECTIVE over this store's + * communicator — every rank must call get_batch() for the same variable + * the same number of times, in the same order (n may differ per rank, + * including 0). See get_batch_alltoall(). */ + template + void get_batch(std::string name, const long *idx, long n, T *buffer, int hmem_iface = 0) + { + const VarInfo_t& varinfo = this->varlist.at(name); + + if (varinfo.itemsize != sizeof(T)) + throw std::invalid_argument("Invalid data type"); + + if (this->method == 0) + { + if (hmem_iface != 0) + throw std::runtime_error("GPU destination buffer is not supported with method=0 (MPI_Win)"); + this->get_batch_alltoall(varinfo, idx, n, (char *)buffer); + return; + } + + if (n <= 0) + return; + + size_t row_bytes = (size_t)varinfo.disp * varinfo.itemsize; + std::vector target(n); + std::vector offset(n); + for (long i = 0; i < n; i++) + { + int t = sortedsearch(varinfo.lenlist, idx[i]); /* throws if out of range */ + long first = t > 0 ? varinfo.lenlist[t - 1] : 0; + target[i] = t; + offset[i] = (uint64_t)(idx[i] - first) * row_bytes; + } + + /* Methods 1 and 2: one lock acquisition for the whole batch. */ + const bool prof = ddstore_profile_enabled(); + uint64_t t_wait = prof ? ddstore_now_ns() : 0; + fabric_state_lock_guard lock(varinfo.fabric_state); + if (prof) + { + varinfo.fabric_state->prof_calls++; + varinfo.fabric_state->prof_lock_wait_ns += ddstore_now_ns() - t_wait; + } + if (hmem_iface != 0 && !is_hmem_capable(varinfo.fabric_state)) + throw std::runtime_error( + "GPU destination buffer requires DDSTORE_FABRIC=cxi " + "(current fabric does not support FI_HMEM)"); + varinfo.fabric_state->recv_data = (char *)buffer; + varinfo.fabric_state->recv_data_len = n * row_bytes; + varinfo.fabric_state->recv_hmem_iface = hmem_iface; + int rc = read_batch_from_remote(varinfo.fabric_state, n, target.data(), + offset.data(), row_bytes); + if (rc != 0) + throw std::runtime_error( + "read_batch_from_remote failed with code " + std::to_string(rc) + + " (" + std::to_string(n) + " rows)"); + } + private: + /* Method 0 batched get, after MDLoader's collective module (Bae et al., + * IPDPSW 2024): every rank all-gathers the batch indices of all ranks, + * packs the rows it owns for each requester, and one MPI_Alltoallv + * delivers them; each rank then puts its rows in request order. Uses a + * private duplicate of the store's communicator (coll_comm), so it never + * matches the caller's own collectives or the windows' fences. The + * indices are validated after the gather, on the global list, so every + * rank throws the same error together instead of one rank leaving the + * others blocked in the exchange. */ + void get_batch_alltoall(const VarInfo_t &varinfo, const long *idx, long n, char *out) + { + std::lock_guard guard(this->coll_mutex); + if (this->coll_comm == MPI_COMM_NULL) + MPI_Comm_dup(this->comm, &this->coll_comm); + + const int P = this->comm_size; + const int me = this->rank; + const size_t row = (size_t)varinfo.disp * varinfo.itemsize; + + /* 1. Every rank learns every rank's requests. */ + int nloc = (int)n; + std::vector nreq(P), rbase(P + 1, 0); + MPI_Allgather(&nloc, 1, MPI_INT, nreq.data(), 1, MPI_INT, this->coll_comm); + for (int p = 0; p < P; p++) + rbase[p + 1] = rbase[p] + nreq[p]; + std::vector all(rbase[P] > 0 ? rbase[P] : 1); + MPI_Allgatherv(idx, nloc, MPI_LONG, all.data(), nreq.data(), rbase.data(), + MPI_LONG, this->coll_comm); + + const long total_rows = varinfo.lenlist.empty() ? 0 : varinfo.lenlist.back(); + for (int j = 0; j < rbase[P]; j++) + if (all[j] < 0 || all[j] >= total_rows) + throw std::out_of_range( + "Global index " + std::to_string(all[j]) + + " is out of range [0, " + std::to_string(total_rows) + ")"); + + /* 2. Split the exchange into rounds of at most `cap` bytes received + * per rank (DDSTORE_ALLTOALL_MAX_BYTES, default 2 MiB): one huge + * Alltoallv of large rows is slower than per-row reads. Every rank + * derives the same round count from the gathered request counts; + * round k moves each rank's k-th slice of `per` requests. */ + const long cap = alltoall_max_bytes(); + const long per = row >= (size_t)cap ? 1 : (long)(cap / (long)row); + long max_req = 0; + for (int p = 0; p < P; p++) + max_req = nreq[p] > max_req ? nreq[p] : max_req; + const long rounds = (max_req + per - 1) / per; + + const long my_first = me > 0 ? varinfo.lenlist[me - 1] : 0; + const long my_end = varinfo.lenlist[me]; + std::vector owner(n > 0 ? n : 1); + for (long i = 0; i < n; i++) + owner[i] = sortedsearch(varinfo.lenlist, idx[i]); + + MPI_Datatype rowtype; + MPI_Type_contiguous((int)row, MPI_BYTE, &rowtype); + MPI_Type_commit(&rowtype); + std::vector scount(P), sdispl(P), rcount(P), rdispl(P), next(P); + std::vector sendbuf, recvbuf; + for (long k = 0; k < rounds; k++) + { + /* Rows this rank owns from every requester's slice, packed per + * requester in request order. */ + for (int p = 0; p < P; p++) + { + scount[p] = 0; + long lo = rbase[p] + k * per, hi = rbase[p] + std::min((long)nreq[p], (k + 1) * per); + for (long j = lo; j < hi; j++) + if (all[j] >= my_first && all[j] < my_end) + scount[p]++; + } + sdispl[0] = 0; + for (int p = 1; p < P; p++) + sdispl[p] = sdispl[p - 1] + scount[p - 1]; + const long nsend = sdispl[P - 1] + scount[P - 1]; + sendbuf.resize((size_t)(nsend > 0 ? nsend : 1) * row); + size_t ks = 0; + for (int p = 0; p < P; p++) + { + long lo = rbase[p] + k * per, hi = rbase[p] + std::min((long)nreq[p], (k + 1) * per); + for (long j = lo; j < hi; j++) + if (all[j] >= my_first && all[j] < my_end) + memcpy(sendbuf.data() + (ks++) * row, + (char *)varinfo.base + (size_t)(all[j] - my_first) * row, row); + } + + /* Where this rank's own slice comes from. */ + const long ilo = std::min(n, k * per), ihi = std::min(n, (k + 1) * per); + std::fill(rcount.begin(), rcount.end(), 0); + for (long i = ilo; i < ihi; i++) + rcount[owner[i]]++; + rdispl[0] = 0; + for (int p = 1; p < P; p++) + rdispl[p] = rdispl[p - 1] + rcount[p - 1]; + recvbuf.resize((size_t)(ihi > ilo ? ihi - ilo : 1) * row); + + MPI_Alltoallv(sendbuf.data(), scount.data(), sdispl.data(), rowtype, + recvbuf.data(), rcount.data(), rdispl.data(), rowtype, this->coll_comm); + + /* Received rows are grouped by owner, each group in request + * order: put them back in this rank's request order. */ + next = rdispl; + for (long i = ilo; i < ihi; i++) + memcpy(out + (size_t)i * row, recvbuf.data() + (size_t)(next[owner[i]]++) * row, row); + } + MPI_Type_free(&rowtype); + } + + /* DDSTORE_ALLTOALL_MAX_BYTES: method 0 get_batch() bytes received per + * rank per exchange round (default 2 MiB; read once). Must be the same + * on every rank: the round count is derived from it. */ + static long alltoall_max_bytes() + { + static long cap = -1; + if (cap < 0) + { + const char *e = getenv("DDSTORE_ALLTOALL_MAX_BYTES"); + long v = e ? atol(e) : 0; + cap = v > 0 ? v : 2L * 1024 * 1024; + } + return cap; + } + int method; // 0: MPI, 1: libfabric, 2: file-based handshake (libfabric transport) + MPI_Comm coll_comm = MPI_COMM_NULL; /* method 0 get_batch; see above */ + std::mutex coll_mutex; /* one get_batch_alltoall at a time */ MPI_Comm comm; int comm_size; diff --git a/pyproject.toml b/pyproject.toml index 9e40331..cdb8d2d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,10 @@ build-backend = "setuptools.build_meta" [project] name = "PyDDStore" -version = "0.1" +version = "2.0" description = "Distributed Data Store" -requires-python = ">=3.6" -dependencies = ["numpy", "mpi4py"] \ No newline at end of file +requires-python = ">=3.9" +dependencies = ["numpy", "mpi4py"] + +[project.optional-dependencies] +torch = ["torch"] diff --git a/setup.py b/setup.py index f91dcac..f4f1351 100644 --- a/setup.py +++ b/setup.py @@ -32,8 +32,8 @@ include_dirs.append("include") extending = Extension( - "pyddstore", - sources=["src/pyddstore.pyx", "src/ddstore.cxx", "src/common.cxx"], + "pyddstore._core", + sources=["src/pyddstore/_core.pyx", "src/ddstore.cxx", "src/common.cxx"], include_dirs=include_dirs, extra_compile_args=["-std=c++11"], define_macros=defs, @@ -45,11 +45,35 @@ extending, ] +# The generated _core.cpp depends on the NumPy headers it was generated +# against: one generated with NumPy 2 does not compile against NumPy 1.x +# headers (PyDataType_ELSIZE). An editable or in-place build keeps it in +# src/, shared by every environment that builds from this checkout, so +# regenerate it whenever the NumPy major version differs from last time. +numpy_major = np.__version__.split(".")[0] +stamp = join("src", "pyddstore", "_core.numpy-version") +try: + with open(stamp) as f: + regenerate = f.read().strip() != numpy_major +except OSError: + regenerate = True +with open(stamp, "w") as f: + f.write(numpy_major + "\n") + +# Left over from the layout before pyddstore became a package. +for old in ("src/pyddstore.cpp",) + tuple( + join("src", f) for f in os.listdir("src") + if f.startswith("pyddstore.") and f.endswith(".so") +): + if os.path.exists(old): + print(f"warning: stale build output {old} from the old layout; remove it") + setup( name="PyDDStore", - version="0.1", + version="2.0", description="Distributed Data Store", package_dir={"": "src"}, + packages=["pyddstore"], py_modules=["cpu_nic_map"], - ext_modules=cythonize(extensions), + ext_modules=cythonize(extensions, force=regenerate), ) diff --git a/src/common.cxx b/src/common.cxx index 3885c3e..ed48f15 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -15,8 +15,8 @@ #include #include -/* hsn (tcp;ofi_rxm over Slingshot) path — Frontier. Unchanged from the - * already-proven dev-file2 implementation; only renamed (was init_fabric). */ +/* hsn (tcp;ofi_rxm over Slingshot) path — DDSTORE_FABRIC=hsn (default). + * The original init_fabric(), unchanged apart from the rename. */ static void init_fabric_hsn(struct fabric_state *fabric) { struct fi_info *hints, *info, *originfo, *useinfo; @@ -186,6 +186,8 @@ static void init_fabric_hsn(struct fabric_state *fabric) fi_strerror(result)); return; } + if (info->ep_attr && info->ep_attr->max_msg_size > 0) + fabric->max_msg_size = info->ep_attr->max_msg_size; av_attr.type = FI_AV_MAP; av_attr.count = DP_AV_DEF_SIZE; @@ -257,9 +259,10 @@ static void init_fabric_hsn(struct fabric_state *fabric) fi_freeinfo(originfo); } -/* cxi path — Perlmutter. Ported from dev-cxi@7cb110b (confirmed working on - * Perlmutter). Kept structurally separate from init_fabric_hsn() above - * rather than unified, so cxi support cannot change hsn's behavior. */ +/* cxi (native Slingshot) path — DDSTORE_FABRIC=cxi. First written for + * Perlmutter; also used on Frontier, where it is required for GPUDirect + * RDMA. Kept structurally separate from init_fabric_hsn() above rather + * than unified, so cxi support cannot change hsn's behavior. */ static void init_fabric_cxi(struct fabric_state *fabric) { struct fi_info *info, *originfo, *useinfo; @@ -439,6 +442,7 @@ static void init_fabric_cxi(struct fabric_state *fabric) && max_msg_size > 0) { info->ep_attr->max_msg_size = max_msg_size; + fabric->max_msg_size = max_msg_size; fprintf(stderr, "endpoint max_msg_size=%zu\n", max_msg_size); } } @@ -542,19 +546,48 @@ int handshake(struct fabric_state *fabric_state, MPI_Comm comm) int world_size = fabric_state->world_size; int rank = fabric_state->rank; - int mr_rc = fi_mr_reg( - fabric_state->domain, - fabric_state->send_data, - fabric_state->send_data_len, - FI_WRITE | FI_REMOTE_READ, - 0, - 0, - 0, - &fabric_state->mr, - NULL); + bool send_is_hmem = fabric_state->send_hmem_iface != FI_HMEM_SYSTEM; + if (send_is_hmem && !is_hmem_capable(fabric_state)) + { + fprintf(stderr, "GPU (HMEM) send buffer requested but fabric is not cxi\n"); + return 1; + } + + int mr_rc; + if (send_is_hmem) + { + /* GPU source buffer. This MR is only ever the passive TARGET of + * other ranks' fi_read() (via remote_key/remote_address, exchanged + * below) -- never the local operand of a local fi_read/fi_write on + * this rank -- so unlike recv_mr in read_from_remote(), no + * fi_mr_desc()/local descriptor is needed here at all. */ + struct iovec iov = {fabric_state->send_data, fabric_state->send_data_len}; + struct fi_mr_attr attr; + memset(&attr, 0, sizeof(attr)); + attr.mr_iov = &iov; + attr.iov_count = 1; + attr.access = FI_WRITE | FI_REMOTE_READ; + attr.iface = (enum fi_hmem_iface)fabric_state->send_hmem_iface; + attr.device.reserved = 0; + mr_rc = fi_mr_regattr(fabric_state->domain, &attr, 0, &fabric_state->mr); + } + else + { + mr_rc = fi_mr_reg( + fabric_state->domain, + fabric_state->send_data, + fabric_state->send_data_len, + FI_WRITE | FI_REMOTE_READ, + 0, + 0, + 0, + &fabric_state->mr, + NULL); + } if (mr_rc != FI_SUCCESS) { - fprintf(stderr, "fi_mr_reg failed: %s\n", fi_strerror(mr_rc)); + fprintf(stderr, "%s (send) failed: %s\n", + send_is_hmem ? "fi_mr_regattr" : "fi_mr_reg", fi_strerror(mr_rc)); return 1; } @@ -629,102 +662,340 @@ int handshake(struct fabric_state *fabric_state, MPI_Comm comm) return 0; } -int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset) +/* Register [base, base + len) for receiving (FI_READ) into *out; hmem_iface + * as for recv_hmem_iface. Returns 0, or 1 after printing the libfabric error + * (*out is then NULL). */ +static int register_recv_mr(struct fabric_state *fabric_state, char *base, size_t len, + int hmem_iface, struct fid_mr **out) { - // register dest buffer; close previous recv MR first to avoid leaking it - if (fabric_state->recv_mr) - fi_close(&fabric_state->recv_mr->fid); - fi_mr_reg( - fabric_state->domain, - fabric_state->recv_data, - fabric_state->recv_data_len, - FI_READ, - 0, - 0, - 0, - &fabric_state->recv_mr, - NULL); + int mr_rc; + *out = NULL; + if (hmem_iface != FI_HMEM_SYSTEM) + { + /* GPU destination buffer (ROCr on AMD, CUDA on NVIDIA -- whichever + * iface the caller set). No host-staged fallback exists: either + * fi_mr_regattr succeeds and fi_read() DMAs straight into device + * memory, or it fails loudly here (checked below). */ + struct iovec iov = {base, len}; + struct fi_mr_attr attr; + memset(&attr, 0, sizeof(attr)); + attr.mr_iov = &iov; + attr.iov_count = 1; + attr.access = FI_READ; + attr.iface = (enum fi_hmem_iface)hmem_iface; + attr.device.reserved = 0; /* ROCr/CUDA both resolve the device from the pointer */ + mr_rc = fi_mr_regattr(fabric_state->domain, &attr, 0, out); + } + else + { + mr_rc = fi_mr_reg(fabric_state->domain, base, len, FI_READ, 0, 0, 0, out, NULL); + } + if (mr_rc != FI_SUCCESS) + { + fprintf(stderr, "%s failed: %s\n", + hmem_iface != FI_HMEM_SYSTEM ? "fi_mr_regattr" : "fi_mr_reg", + fi_strerror(mr_rc)); + *out = NULL; + return 1; + } /* CXI (FI_MR_ENDPOINT): bind and enable recv MR before use. No-op for - * hsn/verbs/gni/psm2 (is_mr_endpoint() is false for those). */ + * hsn/verbs/gni/psm2 (is_mr_endpoint() is false for those). */ if (is_mr_endpoint(fabric_state)) { - int rc_mr = fi_mr_bind(fabric_state->recv_mr, &fabric_state->signal->fid, 0); + int rc_mr = fi_mr_bind(*out, &fabric_state->signal->fid, 0); + if (rc_mr == FI_SUCCESS) + rc_mr = fi_mr_enable(*out); if (rc_mr != FI_SUCCESS) { - fprintf(stderr, "fi_mr_bind (recv) failed: %s\n", fi_strerror(rc_mr)); + fprintf(stderr, "fi_mr_bind/fi_mr_enable (recv) failed: %s\n", fi_strerror(rc_mr)); + fi_close(&(*out)->fid); + *out = NULL; return 1; } - rc_mr = fi_mr_enable(fabric_state->recv_mr); - if (rc_mr != FI_SUCCESS) + } + return 0; +} + +/* Find the recv MR covering recv_data..recv_data+recv_data_len, into *use. + * Returns 0, or 1 after printing the libfabric error. + * + * Caller-registered regions (register_recv_region()) are checked first; + * they are never evicted. Otherwise the one-slot cache: register on a miss + * and reuse the MR while later buffers lie inside [recv_mr_base, + * recv_mr_base + recv_mr_reg_len). Hits when the caller passes the same + * buffer again -- e.g. PyTorch's caching allocator returning the same block + * for a same-shape torch.empty() -- so each get() needn't re-register. */ +static int ensure_recv_mr(struct fabric_state *fabric_state, struct fid_mr **use) +{ + char *cur_base = fabric_state->recv_data; + size_t cur_len = fabric_state->recv_data_len; + + for (int i = 0; i < fabric_state->n_pinned; i++) + { + const struct recv_region *r = &fabric_state->pinned[i]; + if (r->hmem_iface == fabric_state->recv_hmem_iface && + cur_base >= r->base && cur_base + cur_len <= r->base + r->len) { - fprintf(stderr, "fi_mr_enable (recv) failed: %s\n", fi_strerror(rc_mr)); - return 1; + *use = r->mr; + return 0; } } - void *memory_descriptor = NULL; - if (is_local_mr_req(fabric_state)) + bool in_cached_region = + (fabric_state->recv_mr != NULL) && + (cur_base >= fabric_state->recv_mr_base) && + (cur_base + cur_len <= fabric_state->recv_mr_base + fabric_state->recv_mr_reg_len); + if (in_cached_region) { - memory_descriptor = fi_mr_desc(fabric_state->recv_mr); + *use = fabric_state->recv_mr; + return 0; } - size_t rc; - // fprintf(stderr, "fabric_state->remote_address: %llu\n", fabric_state->remote_address[src]); - do + if (ddstore_profile_enabled()) + fabric_state->prof_mr_miss++; + /* Close the stale registration before creating a new one. */ + if (fabric_state->recv_mr) { - rc = fi_read( - fabric_state->signal, - fabric_state->recv_data, - fabric_state->recv_data_len, - memory_descriptor, - fabric_state->comm_partner[src], - fabric_state->remote_address[src] + offset, - fabric_state->remote_key[src], - NULL); - } while (rc == -EAGAIN); - if (rc != 0) + fi_close(&fabric_state->recv_mr->fid); + fabric_state->recv_mr = NULL; + } + if (register_recv_mr(fabric_state, cur_base, cur_len, + fabric_state->recv_hmem_iface, &fabric_state->recv_mr) != 0) + return 1; + + /* Record the registered region for future range checks. */ + fabric_state->recv_mr_base = cur_base; + fabric_state->recv_mr_reg_len = cur_len; + *use = fabric_state->recv_mr; + return 0; +} + +int register_recv_region(struct fabric_state *fs, char *base, size_t len, int hmem_iface) +{ + for (int i = 0; i < fs->n_pinned; i++) + if (fs->pinned[i].base == base && fs->pinned[i].len == len && + fs->pinned[i].hmem_iface == hmem_iface) + return 0; + struct recv_region *grown = (struct recv_region *)realloc( + fs->pinned, (fs->n_pinned + 1) * sizeof(struct recv_region)); + if (!grown) + return 1; + fs->pinned = grown; + struct fid_mr *mr; + if (register_recv_mr(fs, base, len, hmem_iface, &mr) != 0) + return 1; + fs->pinned[fs->n_pinned++] = (struct recv_region){mr, base, len, hmem_iface}; + return 0; +} + +int unregister_recv_region(struct fabric_state *fs, char *base) +{ + for (int i = 0; i < fs->n_pinned; i++) { - fprintf(stderr, "fi_read failed with code %zu.\n", rc); - return (rc); + if (fs->pinned[i].base != base) + continue; + fi_close(&fs->pinned[i].mr->fid); + fs->pinned[i] = fs->pinned[--fs->n_pinned]; + return 0; } + return 1; +} - // (2025/09) segfault when using providers other than sockets - // struct fi_cq_data_entry CQEntry = {0}; - // rc = fi_cq_sread(fabric_state->cq_signal, &CQEntry, 1, NULL, -1); - // if (rc < 1) - // { - // fprintf(stderr, "Received no completion event for remote read\n"); - // return 1; - // } +void close_recv_regions(struct fabric_state *fs) +{ + for (int i = 0; i < fs->n_pinned; i++) + fi_close(&fs->pinned[i].mr->fid); + free(fs->pinned); + fs->pinned = NULL; + fs->n_pinned = 0; +} - for (;;) +/* Largest piece one fi_read() may carry: DDSTORE_MAX_READ_BYTES (default + * 1 GiB), lowered to the endpoint's max_msg_size when that is known. On + * Perlmutter's cxi one 5 GB read fails with EMSGSIZE (2.5 GB works) and + * FI_OPT_MAX_MSG_SIZE reports nothing, hence a default well below that. */ +static size_t max_read_bytes(const struct fabric_state *fabric_state) +{ + static size_t env_limit = 0; + if (env_limit == 0) { - struct fi_cq_data_entry CQEntry = {0}; - rc = fi_cq_read(fabric_state->cq_signal, &CQEntry, 1); - if (rc == 1) - break; - if (rc == -FI_EAVAIL) + const char *v = getenv("DDSTORE_MAX_READ_BYTES"); + long long n = v ? atoll(v) : 0; + env_limit = n > 0 ? (size_t)n : ((size_t)1 << 30); + } + size_t limit = env_limit; + if (fabric_state->max_msg_size > 0 && fabric_state->max_msg_size < limit) + limit = fabric_state->max_msg_size; + return limit; +} + +/* Reap at most one completion from the CQ. Returns 1 for a successful + * completion, -1 for an error completion (printed and consumed, so it + * still counts as one finished read), 0 if none is ready yet, and -2 if + * the CQ itself failed (nothing consumed; the caller cannot keep waiting). + * + * NOTE: CQEntry.len is NOT a reliable success signal on this provider/CQ + * format — it reads 0 even for host-to-host transfers independently + * verified to deliver correct data, so it can't be used to distinguish a + * real silent-no-op (observed once, for an HMEM/ROCr destination) from a + * normal completion. A hard check on it was tried and reverted: it + * false-positived on the working host path. Left unchecked deliberately. */ +static int reap_one(struct fabric_state *fabric_state) +{ + struct fi_cq_data_entry CQEntry = {0}; + ssize_t rc = fi_cq_read(fabric_state->cq_signal, &CQEntry, 1); + if (rc == 1) + return 1; + if (rc == -FI_EAGAIN) + return 0; + if (rc == -FI_EAVAIL) + { + struct fi_cq_err_entry ee = {0}; + fi_cq_readerr(fabric_state->cq_signal, &ee, 0); + /* prov_errno is provider-specific; fi_strerror() is only valid + * for generic fi_errno values. Use fi_cq_strerror() to get the + * correct provider-aware error string (provider-agnostic fix, + * applies to every provider, not just cxi). */ + char errbuf[256]; + const char *errstr = fi_cq_strerror(fabric_state->cq_signal, + ee.prov_errno, ee.err_data, + errbuf, sizeof(errbuf)); + fprintf(stderr, + "fi_cq_read failed: err=%d (%s) prov_errno=%d (%s)\n", + ee.err, fi_strerror(ee.err), ee.prov_errno, + errstr ? errstr : "(unknown)"); + return -1; + } + fprintf(stderr, "fi_cq_read failed: %zd (%s)\n", rc, fi_strerror((int)-rc)); + return -2; +} + +/* Read n rows of row_len bytes into recv_data (recv_data_len must be + * n * row_len): row i comes from rank src[i] at byte offset offset[i] of + * its registered buffer, into recv_data + i * row_len. + * + * All n fi_read()s are posted back to back and only then waited for, so + * the round trips overlap on the network; one recv MR covers the whole + * buffer. A row longer than max_read_bytes() is read in several pieces. + * Returns 0, or non-zero if any read failed. + * + * Every read that was posted is waited for before returning, even after + * an error: a completion left in the CQ would be taken as the next call's. + * This also makes the call blocking, which is load-bearing for a GPU + * (recv_hmem_iface != FI_HMEM_SYSTEM) destination: it keeps the caller's + * device buffer alive (still referenced on the Python stack, so PyTorch's + * caching allocator cannot reuse its storage) for the whole in-flight RDMA + * window. If this is ever made asynchronous, GPU buffer safety must be + * re-examined. */ +int read_batch_from_remote(struct fabric_state *fabric_state, long n, + const int *src, const uint64_t *offset, size_t row_len) +{ + if (fabric_state->recv_hmem_iface != FI_HMEM_SYSTEM && !is_hmem_capable(fabric_state)) + { + fprintf(stderr, "GPU (HMEM) recv buffer requested but fabric is not cxi\n"); + return 1; + } + + const bool prof = ddstore_profile_enabled(); + uint64_t t_mr = prof ? ddstore_now_ns() : 0; + struct fid_mr *recv_mr = NULL; + if (ensure_recv_mr(fabric_state, &recv_mr) != 0) + return 1; + + void *memory_descriptor = NULL; + /* HMEM (device) buffers need their local descriptor passed to fi_read() + * regardless of FI_LOCAL_MR: CXI uses it to route the transfer into + * device memory. FI_LOCAL_MR is deprecated and unset on this libfabric + * build (is_local_mr_req() is always false here), so without this the + * descriptor stayed NULL for HMEM too and fi_read() silently no-op'd + * instead of DMAing into the GPU buffer. */ + if (is_local_mr_req(fabric_state) || fabric_state->recv_hmem_iface != FI_HMEM_SYSTEM) + memory_descriptor = fi_mr_desc(recv_mr); + + uint64_t t_read = 0; + if (prof) + { + t_read = ddstore_now_ns(); + fabric_state->prof_mr_ns += t_read - t_mr; + } + + /* Rows longer than max_read_bytes() go as several fi_read()s. */ + const size_t chunk = max_read_bytes(fabric_state); + long posted = 0, done = 0; + int failed = 0; + for (long i = 0; i < n && !failed; i++) + { + for (size_t off = 0; off < row_len && !failed; off += chunk) { - struct fi_cq_err_entry ee = {0}; - fi_cq_readerr(fabric_state->cq_signal, &ee, 0); - /* prov_errno is provider-specific; fi_strerror() is only valid - * for generic fi_errno values. Use fi_cq_strerror() to get the - * correct provider-aware error string (provider-agnostic fix, - * applies to every provider, not just cxi). */ - char errbuf[256]; - const char *errstr = fi_cq_strerror(fabric_state->cq_signal, - ee.prov_errno, ee.err_data, - errbuf, sizeof(errbuf)); - fprintf(stderr, - "fi_cq_read failed: err=%d (%s) prov_errno=%d (%s)\n", - ee.err, fi_strerror(ee.err), ee.prov_errno, - errstr ? errstr : "(unknown)"); - return 1; + size_t len = row_len - off < chunk ? row_len - off : chunk; + for (;;) + { + ssize_t rc = fi_read( + fabric_state->signal, + fabric_state->recv_data + (size_t)i * row_len + off, + len, + memory_descriptor, + fabric_state->comm_partner[src[i]], + fabric_state->remote_address[src[i]] + offset[i] + off, + fabric_state->remote_key[src[i]], + NULL); + if (rc == 0) + { + posted++; + break; + } + if (rc != -FI_EAGAIN) + { + fprintf(stderr, "fi_read failed: %zd (%s)\n", rc, fi_strerror((int)-rc)); + failed = 1; + break; + } + /* Transmit queue full: make progress by reaping a completion. */ + int r = reap_one(fabric_state); + if (r == -2) + return 1; /* CQ broken: cannot account for in-flight reads */ + if (r != 0) + done++; + if (r < 0) + failed = 1; + } } } - return 0; + uint64_t t_cq = 0; + if (prof) + { + t_cq = ddstore_now_ns(); + fabric_state->prof_read_ns += t_cq - t_read; + } + + /* Wait for every posted read, successful or not. */ + while (done < posted) + { + int r = reap_one(fabric_state); + if (r == -2) + return 1; + if (r != 0) + done++; + if (r < 0) + failed = 1; + } + + if (prof) + { + fabric_state->prof_cq_ns += ddstore_now_ns() - t_cq; + fabric_state->prof_rows += n; + } + return failed; +} + +/* Single-row read into recv_data/recv_data_len from rank src at byte + * offset `offset`: a one-row batch. */ +int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset) +{ + return read_batch_from_remote(fabric_state, 1, &src, &offset, + fabric_state->recv_data_len); } /* ========================================================================= diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index 4682b2c..7830022 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -9,20 +9,20 @@ - allocated_nics(): affinity-aware wrapper — which NIC(s) is *this process*, given its actual pinning (os.sched_getaffinity), closest to. - select_fabric_iface(): called automatically by PyDDStore.__cinit__ - (src/pyddstore.pyx) for method=1/2 to set FABRIC_IFACE if not already set. + (src/pyddstore/_core.pyx) for method=1/2 to set FABRIC_IFACE if not already set. Kernel NIC names are always hsnN under /sys/class/net, on Frontier and -Perlmutter alike -- there is no per-system glob pattern to choose. Perlmutter -just exposes each hsnN NIC's libfabric domain under a different name (cxiN); -pass --fabric cxi (or set DDSTORE_FABRIC=cxi) to see that -translated name instead of the raw kernel one. +Perlmutter alike -- there is no per-system glob pattern to choose. The cxi +libfabric provider names each hsnN NIC's domain cxiN instead; pass --fabric +cxi (or set DDSTORE_FABRIC=cxi) to see that translated name instead of the +raw kernel one. CLI: cpu_nic_map.py print the full CPU -> nearest HSN NIC table cpu_nic_map.py 42 print only the nearest HSN NIC for cpu 42 cpu_nic_map.py --env print the compact DDSTORE_NIC_MAP env-var value cpu_nic_map.py --allocated print this process's allocated CPUs and nearest NIC(s) - cpu_nic_map.py --env --fabric cxi show the Perlmutter-translated (cxiN) names + cpu_nic_map.py --env --fabric cxi show the cxi-provider (cxiN) names export DDSTORE_NIC_MAP=$(python3 cpu_nic_map.py --env) srun --threads-per-core=2 -n8 -c14 python cpu_nic_map.py --allocated @@ -134,8 +134,8 @@ def build_map(pattern): for nic, part in zip(sorted(group), partitions): nic_closest[nic] = part - # multiple NICs can share a NUMA node; pick the numerically/PCI-closest - # NIC as the "same-NUMA fallback owner" for cores not exactly local to any NIC + # multiple NICs can share a NUMA node; the first one by name is the + # "same-NUMA fallback owner" for cores not exactly local to any NIC numa_to_nics = {} for n, numa in nic_numa.items(): numa_to_nics.setdefault(numa, []).append(n) @@ -172,8 +172,8 @@ def compress_ranges(values): def translate_iface(name, provider="hsn"): """Translate a kernel NIC name (hsnN) to the libfabric domain name for - `provider`. 'cxi' -> cxiN (Perlmutter exposes hsnN's libfabric domain - under this name); 'hsn' (default) or anything else -> unchanged.""" + `provider`. 'cxi' -> cxiN (the cxi provider's domain name for hsnN); + 'hsn' (default) or anything else -> unchanged.""" if provider == "cxi": m = re.match(r"hsn(\d+)$", name) if m: @@ -253,7 +253,7 @@ def select_fabric_iface(nic_map=None): else DDSTORE_NIC_MAP when set, or a live hwloc-calc/lstopo query (build_map) against this process's real CPU affinity when neither is. - Called automatically by PyDDStore.__cinit__ (src/pyddstore.pyx) for + Called automatically by PyDDStore.__cinit__ (src/pyddstore/_core.pyx) for method=1/2. nic_map: an explicit precomputed map string (serialize_env()/--env @@ -261,11 +261,10 @@ def select_fabric_iface(nic_map=None): than the process environment. DDSTORE_FABRIC selects hsn (default) or cxi: - - hsn: Frontier's unchanged, already-proven behavior -- the kernel NIC - name (hsnN) is used as-is. - - cxi: Perlmutter's behavior, ported from dev-cxi@7cb110b. The kernel - NIC names are hsn0-hsn3 there too, but libfabric only exposes them - as cxi0-cxi3, so the result is translated hsnN -> cxiN. Also adds a + - hsn: the kernel NIC name (hsnN) is used as-is (tcp;ofi_rxm). + - cxi: the kernel NIC names are hsn0-hsn3, but the cxi provider + exposes them as cxi0-cxi3, so the result is translated hsnN -> cxiN + (Frontier and Perlmutter alike). Also adds a SLURM_LOCALID round-robin fallback for when hwloc can't map this rank's CPU affinity to a NIC (common inside srun tasks with limited PCI visibility). @@ -330,7 +329,7 @@ def main(): " cpu_nic_map.py 42 print only the nearest HSN NIC for cpu 42\n" " export DDSTORE_NIC_MAP=$(cpu_nic_map.py --env) compute once, share via env\n" " srun ... python cpu_nic_map.py --allocated show this task's allocated CPUs + nearest NIC(s)\n" - " cpu_nic_map.py --env --fabric cxi show the Perlmutter-translated (cxiN) names\n" + " cpu_nic_map.py --env --fabric cxi show the cxi-provider (cxiN) names\n" ), ) parser.add_argument( @@ -345,7 +344,7 @@ def main(): choices=["hsn", "cxi"], help="translate printed NIC names to this fabric's libfabric " "domain name (default: $DDSTORE_FABRIC, or hsn if unset) " - "-- hsn: unchanged (e.g. hsn0); cxi: hsnN -> cxiN (Perlmutter)", + "-- hsn: unchanged (e.g. hsn0); cxi: hsnN -> cxiN", ) parser.add_argument( "--env", diff --git a/src/ddstore.cxx b/src/ddstore.cxx index 5e91ae2..3ac33c6 100644 --- a/src/ddstore.cxx +++ b/src/ddstore.cxx @@ -118,8 +118,9 @@ void DDStore::epoch_end() /* -------------------------------------------------------------------------- * join() — extra member: discover a variable published by core members. * - * Calls handshake_join() which polls for all CoreRecord files, then - * populates a fabric_state and builds the lenlist for get() calls. + * Calls handshake_join(), which polls for the combined {name}.bin record + * file written by core rank 0, then populates a fabric_state and builds + * the lenlist for get() calls. * -------------------------------------------------------------------------- */ void DDStore::join(std::string name) { @@ -130,6 +131,7 @@ void DDStore::join(std::string name) struct fabric_state *fs = (struct fabric_state *)calloc(1, sizeof(struct fabric_state)); + pthread_mutex_init(&fs->recv_lock, NULL); fs->world_size = this->n_core; fs->rank = -1; /* extra members have no core rank */ @@ -137,9 +139,9 @@ void DDStore::join(std::string name) if (!fs->info) throw std::runtime_error("init_fabric failed for extra member"); - /* Extra member has no send buffer to register as MR — set a dummy - * zero-length registration so handshake_join doesn't need special-casing. - * We only need fi_read capability, not FI_REMOTE_READ on our side. */ + /* Extra member has no send buffer, so nothing is registered for remote + * access (mr stays NULL, key 0). It only issues fi_read()s; get() + * registers each destination buffer as usual in read_from_remote(). */ fs->send_data = NULL; fs->send_data_len = 0; fs->mr = NULL; @@ -171,51 +173,61 @@ void DDStore::join(std::string name) var.active = true; var.fence_active = false; var.base = NULL; /* extra member owns no data */ + var.owns_base = false; var.fabric_state = fs; this->varlist.insert(std::pair(name, var)); } /* -------------------------------------------------------------------------- * free() — release all resources. + * + * Per variable: the MPI window (method 0) or the libfabric objects (methods + * 1/2) first, since they reference the buffer, then the buffer itself if + * DDStore allocated it (owns_base). MPI_Win_free/MPI_Free_mem are skipped + * after MPI_Finalize (no longer callable). Idempotent via `active`. * -------------------------------------------------------------------------- */ void DDStore::free() { - int flag; - MPI_Finalized(&flag); - if (!this->method && !flag) + int finalized; + MPI_Finalized(&finalized); + for (auto &x : this->varlist) { - for (auto &x : this->varlist) + VarInfo_t &var = x.second; + if (!var.active) + continue; + + if (this->method == 0) { - if (x.second.active) - { - MPI_Win_free(&x.second.win); - } - x.second.active = false; + if (!finalized) + MPI_Win_free(&var.win); } - } - else if (this->method == 1 || this->method == 2) - { - for (auto &x : this->varlist) + else if (var.fabric_state) { - if (x.second.active && x.second.fabric_state) - { - struct fabric_state *fs = x.second.fabric_state; - if (fs->recv_mr) fi_close(&fs->recv_mr->fid); - if (fs->mr) fi_close(&fs->mr->fid); - if (fs->signal) fi_close(&fs->signal->fid); - if (fs->cq_signal) fi_close(&fs->cq_signal->fid); - if (fs->av) fi_close(&fs->av->fid); - if (fs->domain) fi_close(&fs->domain->fid); - if (fs->fabric) fi_close(&fs->fabric->fid); - if (fs->info) fi_freeinfo(fs->info); - if (fs->ctx) ::free(fs->ctx); - ::free(fs->comm_partner); - ::free(fs->remote_key); - ::free(fs->remote_address); - ::free(fs); - x.second.fabric_state = NULL; - } - x.second.active = false; + struct fabric_state *fs = var.fabric_state; + close_recv_regions(fs); + if (fs->recv_mr) fi_close(&fs->recv_mr->fid); + if (fs->mr) fi_close(&fs->mr->fid); + if (fs->signal) fi_close(&fs->signal->fid); + if (fs->cq_signal) fi_close(&fs->cq_signal->fid); + if (fs->av) fi_close(&fs->av->fid); + if (fs->domain) fi_close(&fs->domain->fid); + if (fs->fabric) fi_close(&fs->fabric->fid); + if (fs->info) fi_freeinfo(fs->info); + if (fs->ctx) ::free(fs->ctx); + ::free(fs->comm_partner); + ::free(fs->remote_key); + ::free(fs->remote_address); + pthread_mutex_destroy(&fs->recv_lock); + ::free(fs); + var.fabric_state = NULL; } + + if (var.owns_base && var.base && !finalized) + MPI_Free_mem(var.base); + var.base = NULL; + var.owns_base = false; + var.active = false; } + if (this->coll_comm != MPI_COMM_NULL && !finalized) + MPI_Comm_free(&this->coll_comm); } diff --git a/src/pyddstore.pyx b/src/pyddstore.pyx deleted file mode 100644 index f85b976..0000000 --- a/src/pyddstore.pyx +++ /dev/null @@ -1,197 +0,0 @@ -# distutils: language=c++ -# cython: language_level=3 -# cython: language=c++ - -import mpi4py.MPI as MPI -cimport mpi4py.MPI as MPI -cimport mpi4py.libmpi as libmpi - -import cpu_nic_map - -import numpy as np -cimport numpy as np - -from libcpp.string cimport string -from libcpp.typeinfo cimport type_info - -from cpython.version cimport PY_MAJOR_VERSION - -cpdef str b2s(bytes x): - if PY_MAJOR_VERSION < 3: - return str(x) - else: - return x.decode() - -cpdef bytes s2b(str x): - if PY_MAJOR_VERSION < 3: - return x - else: - return x.encode() - -cdef extern from "ddstore.hpp": - ctypedef struct VarInfo: - string name - int disp - int itemsize - - - cdef cppclass DDStore: - DDStore() - DDStore(libmpi.MPI_Comm comm) - DDStore(int method, libmpi.MPI_Comm comm) - # Method 2: core member (with MPI communicator; n_core == comm size) - DDStore(int method, libmpi.MPI_Comm comm, - string handshake_dir) - # Method 2: extra member (no MPI communicator) - DDStore(int method, string handshake_dir, int n_core) - void add[T](string name, T* buffer, long nrows, int disp) except + - void get[T](string name, long start, long count, T* buffer) except + - void epoch_begin() - void epoch_end() - void free() - void init(string name, long nrows, int disp, int itemsize) except + - void update[T](string name, T* buffer, long nrows, long offset) except + - void join(string name) except + - void query(string name, VarInfo &varinfo) except + - long size(string name) except + - -cdef class PyDDstoreVarinfo: - cdef VarInfo c_varinfo - - def __cinit__(self): - pass - -cdef class PyDDStore: - cdef DDStore *c_ddstore - - def __cinit__(self, comm_or_none=None, int method=0, - str handshake_dir="", int n_core=0, nic_map=None): - """ - Constructors: - PyDDStore(comm) — method 0, MPI - PyDDStore(comm, method=1) — method 1, libfabric+MPI - PyDDStore(comm, method=2, — method 2, core member - handshake_dir="/path") (n_core == comm size) - PyDDStore(None, method=2, — method 2, extra member - handshake_dir="/path", n_core=N) - - nic_map: optional precomputed CPU->NIC map string (see - cpu_nic_map.py --env) used to select FABRIC_IFACE for this rank's - CPU affinity, for method=1/2. Takes priority over the - DDSTORE_NIC_MAP env var. Only used if FABRIC_IFACE isn't already - set in the environment. - """ - cdef MPI.Comm mpi_comm - if method != 0: - cpu_nic_map.select_fabric_iface(nic_map=nic_map) - if method == 2: - if not handshake_dir: - raise ValueError( - "method=2 requires handshake_dir (got handshake_dir=%r)" - % handshake_dir) - if comm_or_none is None: - # Extra member: no MPI communicator, n_core must be given - if n_core <= 0: - raise ValueError( - "method=2 extra member requires n_core > 0 " - "(got n_core=%d)" % n_core) - self.c_ddstore = new DDStore(method, - s2b(handshake_dir), n_core) - else: - # Core member with file-based handshake; n_core is derived - # from the communicator size. - mpi_comm = comm_or_none - self.c_ddstore = new DDStore(method, mpi_comm.ob_mpi, - s2b(handshake_dir)) - else: - # Methods 0 and 1: standard MPI constructor - if comm_or_none is None: - raise ValueError( - "method=%d requires a valid MPI communicator " - "(got comm_or_none=None)" % method) - mpi_comm = comm_or_none - self.c_ddstore = new DDStore(method, mpi_comm.ob_mpi) - - def __dealloc__(self): - if self.c_ddstore != NULL: - del self.c_ddstore - self.c_ddstore = NULL - - def add(self, str name, np.ndarray arr): - assert arr.flags.c_contiguous - cdef long nrows = arr.shape[0] - cdef int disp = arr.size // arr.shape[0] - if arr.dtype == np.int32: - self.c_ddstore.add(s2b(name), arr.data, nrows, disp) - elif arr.dtype == np.int64: - self.c_ddstore.add(s2b(name), arr.data, nrows, disp) - elif arr.dtype == np.uint8: - self.c_ddstore.add(s2b(name), arr.data, nrows, disp) - elif arr.dtype == np.float32: - self.c_ddstore.add(s2b(name), arr.data, nrows, disp) - elif arr.dtype == np.float64: - self.c_ddstore.add(s2b(name), arr.data, nrows, disp) - elif arr.dtype == np.bool_: - self.c_ddstore.add(s2b(name), arr.data, nrows, disp) - else: - raise NotImplementedError - - def get(self, str name, np.ndarray arr, long start=0): - assert arr.flags.c_contiguous - cdef long count = arr.shape[0] - assert arr.shape[0] >= count - if arr.dtype == np.int32: - self.c_ddstore.get(s2b(name), start, count, arr.data) - elif arr.dtype == np.int64: - self.c_ddstore.get(s2b(name), start, count, arr.data) - elif arr.dtype == np.uint8: - self.c_ddstore.get(s2b(name), start, count, arr.data) - elif arr.dtype == np.float32: - self.c_ddstore.get(s2b(name), start, count, arr.data) - elif arr.dtype == np.float64: - self.c_ddstore.get(s2b(name), start, count, arr.data) - elif arr.dtype == np.bool_: - self.c_ddstore.get(s2b(name), start, count, arr.data) - else: - raise NotImplementedError - - def epoch_begin(self): - self.c_ddstore.epoch_begin() - - def epoch_end(self): - self.c_ddstore.epoch_end() - - def free(self): - self.c_ddstore.free() - - def init(self, str name, long nrows, int disp, int itemsize=1): - self.c_ddstore.init(s2b(name), nrows, disp, itemsize) - - def update(self, str name, np.ndarray arr, long offset): - assert arr.flags.c_contiguous - cdef long nrows = arr.shape[0] - if arr.dtype == np.int32: - self.c_ddstore.update(s2b(name), arr.data, nrows, offset) - elif arr.dtype == np.int64: - self.c_ddstore.update(s2b(name), arr.data, nrows, offset) - elif arr.dtype == np.uint8: - self.c_ddstore.update(s2b(name), arr.data, nrows, offset) - elif arr.dtype == np.float32: - self.c_ddstore.update(s2b(name), arr.data, nrows, offset) - elif arr.dtype == np.float64: - self.c_ddstore.update(s2b(name), arr.data, nrows, offset) - elif arr.dtype == np.bool_: - self.c_ddstore.update(s2b(name), arr.data, nrows, offset) - else: - raise NotImplementedError - - def join(self, str name): - """Method 2 extra member: discover variable published by core members.""" - self.c_ddstore.join(s2b(name)) - - def info(self, str name): - """Return (total_rows, disp, itemsize) for an added or joined variable.""" - cdef VarInfo vi - self.c_ddstore.query(s2b(name), vi) - total_rows = self.c_ddstore.size(s2b(name)) - return (total_rows, vi.disp, vi.itemsize) diff --git a/src/pyddstore/__init__.py b/src/pyddstore/__init__.py new file mode 100644 index 0000000..870b071 --- /dev/null +++ b/src/pyddstore/__init__.py @@ -0,0 +1,29 @@ +"""DDStore: distributed in-memory data store for data-parallel training. + +The store itself is ``PyDDStore`` (compiled extension ``pyddstore._core``). +PyTorch helpers (``DistDataset``, ``DistDatasetReader``, ``ThreadDataLoader``) +live in ``pyddstore.torch`` and need PyTorch; importing ``pyddstore`` alone +does not. + +The extension (which initializes MPI through mpi4py) is loaded on first use +of ``PyDDStore``, not on ``import pyddstore``: ``from pyddstore.torch import +...`` then imports torch before MPI starts, the order torch + RCCL/NCCL need +to shut down cleanly. +""" + +try: + from importlib.metadata import version as _version + + __version__ = _version("PyDDStore") +except Exception: # pragma: no cover - not installed as a distribution + __version__ = "unknown" + +__all__ = ["PyDDStore", "PyDDstoreVarinfo", "__version__"] + + +def __getattr__(name): + if name in ("PyDDStore", "PyDDstoreVarinfo"): + from . import _core + + return getattr(_core, name) + raise AttributeError(f"module 'pyddstore' has no attribute {name!r}") diff --git a/src/pyddstore/_core.pyx b/src/pyddstore/_core.pyx new file mode 100644 index 0000000..1ce2d56 --- /dev/null +++ b/src/pyddstore/_core.pyx @@ -0,0 +1,468 @@ +# distutils: language=c++ +# cython: language_level=3 +# cython: language=c++ + +import os +import time + +import mpi4py.MPI as MPI +cimport mpi4py.MPI as MPI +cimport mpi4py.libmpi as libmpi + +import cpu_nic_map + +import numpy as np +cimport numpy as np + +from libcpp.string cimport string +from libcpp.typeinfo cimport type_info + +from cpython.version cimport PY_MAJOR_VERSION + +cpdef str b2s(bytes x): + if PY_MAJOR_VERSION < 3: + return str(x) + else: + return x.decode() + +cpdef bytes s2b(str x): + if PY_MAJOR_VERSION < 3: + return x + else: + return x.encode() + +def _is_cuda_tensor(obj): + """True if obj is a torch.Tensor on a CUDA/HIP device. + + Torch is optional — imported lazily so DDStore2 has no hard dependency + on it. If unimportable, no object is ever considered a CUDA tensor. + """ + try: + import torch + except ImportError: + return False + return isinstance(obj, torch.Tensor) and obj.is_cuda + +# Mirrors libfabric's enum fi_hmem_iface (rdma/fi_domain.h): FI_HMEM_SYSTEM=0, +# FI_HMEM_CUDA=1, FI_HMEM_ROCR=2. Kept as plain ints here (rather than +# cimporting the C enum) since only these two values are ever produced by +# _hmem_iface_for() below -- torch itself is either a CUDA or a ROCm build, +# never both. +_FI_HMEM_CUDA = 1 +_FI_HMEM_ROCR = 2 + +def _hmem_iface_for(tensor): + """fi_hmem_iface value for a CUDA tensor: ROCr on a ROCm/HIP build of + torch (AMD), CUDA otherwise (NVIDIA). Only call when _is_cuda_tensor() + is already True. + """ + import torch + return _FI_HMEM_ROCR if torch.version.hip is not None else _FI_HMEM_CUDA + + +def _check_gpu_fabric_preconditions(int method, str what): + """Shared method=1/2 + DDSTORE_FABRIC=cxi precondition check for a GPU + (CUDA/HIP) buffer passed to add() or get(). `what` customizes the error + wording ("GPU source buffer" / "GPU destination buffer"). + """ + if method not in (1, 2): + raise RuntimeError( + "%s requires method=1 or 2 (libfabric), got method=%d" % (what, method)) + provider = os.environ.get("DDSTORE_FABRIC", "hsn") + if provider != "cxi": + raise RuntimeError( + "%s requires DDSTORE_FABRIC=cxi (current DDSTORE_FABRIC=%r); " + "the hsn (tcp;ofi_rxm) path does not support FI_HMEM. Set " + "DDSTORE_FABRIC=cxi or pass a host (CPU) numpy array instead." + % (what, provider)) + +def _check_dtype(arr, bint is_gpu): + """Raise NotImplementedError unless arr's dtype is one DDStore supports. + add()/get() dispatch on item size alone (1/4/8 bytes), so this is what + keeps e.g. float16 or complex64 from slipping through on a size match. + """ + if is_gpu: + import torch + ok = arr.dtype in (torch.int32, torch.int64, torch.uint8, + torch.float32, torch.float64, torch.bool) + else: + ok = arr.dtype in (np.int32, np.int64, np.uint8, + np.float32, np.float64, np.bool_) + if not ok: + raise NotImplementedError("unsupported dtype: %s" % arr.dtype) + +cdef extern from "ddstore.hpp": + ctypedef struct VarInfo: + string name + int disp + int itemsize + + + cdef cppclass DDStore: + DDStore() + DDStore(libmpi.MPI_Comm comm) + DDStore(int method, libmpi.MPI_Comm comm) + # Method 2: core member (with MPI communicator; n_core == comm size) + DDStore(int method, libmpi.MPI_Comm comm, + string handshake_dir) + # Method 2: extra member (no MPI communicator) + DDStore(int method, string handshake_dir, int n_core) + void add[T](string name, T* buffer, long nrows, int disp, int hmem_iface) except + + void get[T](string name, long start, long count, T* buffer, int hmem_iface) except + nogil + void get_batch[T](string name, const long *idx, long n, T* buffer, int hmem_iface) except + nogil + void register_recv(string name, void *buffer, size_t len, int hmem_iface) except + + void unregister_recv(string name, void *buffer) except + + void epoch_begin() + void epoch_end() + void free() + void init(string name, long nrows, int disp, int itemsize) except + + void update[T](string name, T* buffer, long nrows, long offset) except + + void join(string name) except + + void query(string name, VarInfo &varinfo) except + + long size(string name) except + + void profile(string name, unsigned long long *out) except + + +cdef class PyDDstoreVarinfo: + cdef VarInfo c_varinfo + + def __cinit__(self): + pass + +cdef class PyDDStore: + cdef DDStore *c_ddstore + cdef int method + # Keepalive for GPU tensors passed to add(): C++ holds a raw pointer + # into them with no copy and no refcounting (see ddstore.hpp add()'s + # lifetime-contract comment) -- this dict keeps the Python reference + # alive for as long as the variable stays registered. + cdef dict _gpu_owned_buffers + cdef dict _recv_buffers + # DDSTORE_PROFILE=1: Python-side get() timing (see get_profile()). + cdef bint _prof + cdef double _prof_get_s + cdef double _prof_sync_s + cdef long _prof_gets + + def __cinit__(self, comm_or_none=None, int method=0, + str handshake_dir="", int n_core=0, nic_map=None): + """ + Constructors: + PyDDStore(comm) — method 0, MPI + PyDDStore(comm, method=1) — method 1, libfabric+MPI + PyDDStore(comm, method=2, — method 2, core member + handshake_dir="/path") (n_core == comm size) + PyDDStore(None, method=2, — method 2, extra member + handshake_dir="/path", n_core=N) + + nic_map: optional precomputed CPU->NIC map string (see + cpu_nic_map.py --env) used to select FABRIC_IFACE for this rank's + CPU affinity, for method=1/2. Takes priority over the + DDSTORE_NIC_MAP env var. Only used if FABRIC_IFACE isn't already + set in the environment. + """ + cdef MPI.Comm mpi_comm + self.method = method + self._gpu_owned_buffers = {} + self._recv_buffers = {} + self._prof = os.environ.get("DDSTORE_PROFILE", "0") not in ("", "0") + self._prof_get_s = 0.0 + self._prof_sync_s = 0.0 + self._prof_gets = 0 + if method != 0: + cpu_nic_map.select_fabric_iface(nic_map=nic_map) + if method == 2: + if not handshake_dir: + raise ValueError( + "method=2 requires handshake_dir (got handshake_dir=%r)" + % handshake_dir) + if comm_or_none is None: + # Extra member: no MPI communicator, n_core must be given + if n_core <= 0: + raise ValueError( + "method=2 extra member requires n_core > 0 " + "(got n_core=%d)" % n_core) + self.c_ddstore = new DDStore(method, + s2b(handshake_dir), n_core) + else: + # Core member with file-based handshake; n_core is derived + # from the communicator size. + mpi_comm = comm_or_none + self.c_ddstore = new DDStore(method, mpi_comm.ob_mpi, + s2b(handshake_dir)) + else: + # Methods 0 and 1: standard MPI constructor + if comm_or_none is None: + raise ValueError( + "method=%d requires a valid MPI communicator " + "(got comm_or_none=None)" % method) + mpi_comm = comm_or_none + self.c_ddstore = new DDStore(method, mpi_comm.ob_mpi) + + def __dealloc__(self): + if self.c_ddstore != NULL: + del self.c_ddstore + self.c_ddstore = NULL + self._gpu_owned_buffers.clear() + self._recv_buffers.clear() + + def add(self, str name, arr): + cdef size_t ptr + cdef int itemsize + cdef int iface + cdef long nrows = arr.shape[0] + cdef int disp + cdef bint is_gpu = _is_cuda_tensor(arr) + _check_dtype(arr, is_gpu) + if is_gpu: + _check_gpu_fabric_preconditions(self.method, "GPU source buffer") + assert arr.is_contiguous() + if name in self._gpu_owned_buffers: + raise RuntimeError( + "add() called again for variable '%s' with a GPU source " + "buffer; re-adding an existing variable name is not " + "supported (the original registration would remain " + "active in C++ while its Python keepalive reference is " + "replaced here, risking a dangling pointer)" % name) + import torch + # Flush any pending/async GPU compute-kernel writes to `arr` + # before handing it to RDMA -- otherwise the transfer can be + # silently masked by stale GPU cache content from a preceding, + # not-yet-retired write to the same memory. + torch.cuda.synchronize(device=arr.device) + ptr = arr.data_ptr() + itemsize = arr.element_size() + disp = arr.numel() // nrows + iface = _hmem_iface_for(arr) + else: + assert arr.flags.c_contiguous + ptr = arr.ctypes.data + itemsize = arr.itemsize + disp = arr.size // nrows + iface = 0 + + # DDStore::add() only uses T through sizeof(T), so dispatching on + # item size is enough. + cdef string cname = s2b(name) + if itemsize == 1: + self.c_ddstore.add(cname, ptr, nrows, disp, iface) + elif itemsize == 4: + self.c_ddstore.add(cname, ptr, nrows, disp, iface) + else: + self.c_ddstore.add(cname, ptr, nrows, disp, iface) + + if is_gpu: + # Keepalive: DDStore now holds a raw pointer into arr's storage + # with no copy and no C++-level refcounting -- see ddstore.hpp + # add()'s lifetime-contract doc comment. Must outlive this + # variable's registration; cleared in free()/__dealloc__. + self._gpu_owned_buffers[name] = arr + + def get(self, str name, arr, long start=0): + cdef double t_get = time.perf_counter() if self._prof else 0.0 + cdef double t_sync + cdef long count = arr.shape[0] + cdef size_t ptr + cdef int itemsize + cdef int iface + cdef bint is_gpu = _is_cuda_tensor(arr) + _check_dtype(arr, is_gpu) + if is_gpu: + _check_gpu_fabric_preconditions(self.method, "GPU destination buffer") + assert arr.is_contiguous() + import torch + # See the matching comment in add() for what this guards against. + if self._prof: + t_sync = time.perf_counter() + torch.cuda.synchronize(device=arr.device) + self._prof_sync_s += time.perf_counter() - t_sync + else: + torch.cuda.synchronize(device=arr.device) + ptr = arr.data_ptr() + itemsize = arr.element_size() + iface = _hmem_iface_for(arr) + else: + assert arr.flags.c_contiguous + ptr = arr.ctypes.data + itemsize = arr.itemsize + iface = 0 + + # DDStore::get() only uses T for its sizeof(T) == itemsize check + # and the pointer cast -- the transfer itself is a byte copy -- so + # dispatching on item size is enough. The read runs without the GIL, + # so other Python threads (e.g. a training loop while a background + # thread prefetches) keep running. Method 1/2 reads make no MPI calls. + cdef string cname = s2b(name) + with nogil: + if itemsize == 1: + self.c_ddstore.get(cname, start, count, ptr, iface) + elif itemsize == 4: + self.c_ddstore.get(cname, start, count, ptr, iface) + else: + self.c_ddstore.get(cname, start, count, ptr, iface) + if self._prof: + self._prof_get_s += time.perf_counter() - t_get + self._prof_gets += 1 + + def get_batch(self, str name, arr, indices): + """Read rows `indices` (global row ids, any order, repeats allowed) + into `arr`, whose first dimension must equal len(indices): row i of + `arr` receives row indices[i]. Same buffer rules as get(); for + method 1/2 all reads of the batch are in flight together, under one + lock acquisition and (GPU destination) one device sync. + + Method 0 is COLLECTIVE (MDLoader-style Allgatherv of indices + + Alltoallv of rows): every rank of the store must call get_batch() + for the variable the same number of times, in the same order, from + one thread at a time (len(indices) may differ, including 0).""" + cdef double t_get = time.perf_counter() if self._prof else 0.0 + cdef double t_sync + cdef np.ndarray idx = np.ascontiguousarray(indices, dtype=np.int64) + if idx.ndim != 1: + raise ValueError("indices must be one-dimensional") + cdef long n = idx.shape[0] + if arr.shape[0] != n: + raise ValueError( + "arr has %d rows but %d indices were given" % (arr.shape[0], n)) + cdef const long *cidx = idx.data + cdef size_t ptr + cdef int itemsize + cdef int iface + cdef bint is_gpu = _is_cuda_tensor(arr) + _check_dtype(arr, is_gpu) + if is_gpu: + _check_gpu_fabric_preconditions(self.method, "GPU destination buffer") + assert arr.is_contiguous() + import torch + # Same reason as get(), once per batch. + if self._prof: + t_sync = time.perf_counter() + torch.cuda.synchronize(device=arr.device) + self._prof_sync_s += time.perf_counter() - t_sync + else: + torch.cuda.synchronize(device=arr.device) + ptr = arr.data_ptr() + itemsize = arr.element_size() + iface = _hmem_iface_for(arr) + else: + assert arr.flags.c_contiguous + ptr = arr.ctypes.data + itemsize = arr.itemsize + iface = 0 + + cdef string cname = s2b(name) + with nogil: + if itemsize == 1: + self.c_ddstore.get_batch(cname, cidx, n, ptr, iface) + elif itemsize == 4: + self.c_ddstore.get_batch(cname, cidx, n, ptr, iface) + else: + self.c_ddstore.get_batch(cname, cidx, n, ptr, iface) + if self._prof: + self._prof_get_s += time.perf_counter() - t_get + self._prof_gets += 1 + + def register_recv(self, str name, arr): + """Register `arr` (contiguous host numpy array or GPU tensor) once as + a destination for get()/get_batch() of `name`: reads into it, or into + any slice of it, then skip memory registration. Use for buffers that + are reused across reads (e.g. a per-thread pool); several can be + registered per variable and none is evicted. The store keeps `arr` + alive until unregister_recv() or free(). No-op for method 0.""" + cdef size_t ptr + cdef size_t nbytes + cdef int iface + cdef bint is_gpu = _is_cuda_tensor(arr) + if is_gpu: + _check_gpu_fabric_preconditions(self.method, "GPU destination buffer") + assert arr.is_contiguous() + ptr = arr.data_ptr() + nbytes = arr.numel() * arr.element_size() + iface = _hmem_iface_for(arr) + else: + assert arr.flags.c_contiguous + ptr = arr.ctypes.data + nbytes = arr.nbytes + iface = 0 + if nbytes == 0: + return + self.c_ddstore.register_recv(s2b(name), ptr, nbytes, iface) + self._recv_buffers[(name, ptr)] = arr + + def unregister_recv(self, str name, arr): + """Undo register_recv(name, arr).""" + cdef size_t ptr = arr.data_ptr() if _is_cuda_tensor(arr) else arr.ctypes.data + if (name, ptr) not in self._recv_buffers: + if self.method == 0: + return + raise ValueError("buffer is not registered for %r" % name) + self.c_ddstore.unregister_recv(s2b(name), ptr) + del self._recv_buffers[(name, ptr)] + + def get_profile(self, str name): + """DDSTORE_PROFILE=1 timing for `name` (methods 1/2), in seconds. + + C++ counters for this variable: calls (get + get_batch), rows, + lock_wait, mr (recv-MR cache check/registration), mr_miss + (re-registrations), read (posting fi_read), cq (waiting for + completion). Python counters for this + store, across all variables: py_gets, py_get (whole get() calls), + py_sync (torch.cuda.synchronize on the GPU-destination path). + """ + cdef unsigned long long c[7] + self.c_ddstore.profile(s2b(name), c) + return { + "calls": c[0], "lock_wait": c[1] * 1e-9, "mr": c[2] * 1e-9, + "mr_miss": c[3], "read": c[4] * 1e-9, "cq": c[5] * 1e-9, + "rows": c[6], + "py_gets": self._prof_gets, "py_get": self._prof_get_s, + "py_sync": self._prof_sync_s, + } + + def epoch_begin(self): + self.c_ddstore.epoch_begin() + + def epoch_end(self): + self.c_ddstore.epoch_end() + + def free(self): + self.c_ddstore.free() + self._gpu_owned_buffers.clear() + self._recv_buffers.clear() + + def init(self, str name, long nrows, int disp, int itemsize=1): + self.c_ddstore.init(s2b(name), nrows, disp, itemsize) + + def update(self, str name, arr, long offset): + if _is_cuda_tensor(arr): + raise NotImplementedError( + "update() only supports host (numpy) buffers -- the " + "init()/update() path is host-only; pass arr.cpu().numpy() " + "instead, or add() the GPU tensor directly") + cdef np.ndarray np_arr = arr + assert np_arr.flags.c_contiguous + cdef long nrows = np_arr.shape[0] + if np_arr.dtype == np.int32: + self.c_ddstore.update(s2b(name), np_arr.data, nrows, offset) + elif np_arr.dtype == np.int64: + self.c_ddstore.update(s2b(name), np_arr.data, nrows, offset) + elif np_arr.dtype == np.uint8: + self.c_ddstore.update(s2b(name), np_arr.data, nrows, offset) + elif np_arr.dtype == np.float32: + self.c_ddstore.update(s2b(name), np_arr.data, nrows, offset) + elif np_arr.dtype == np.float64: + self.c_ddstore.update(s2b(name), np_arr.data, nrows, offset) + elif np_arr.dtype == np.bool_: + self.c_ddstore.update(s2b(name), np_arr.data, nrows, offset) + else: + raise NotImplementedError + + def join(self, str name): + """Method 2 extra member: discover variable published by core members.""" + self.c_ddstore.join(s2b(name)) + + def info(self, str name): + """Return (total_rows, disp, itemsize) for an added or joined variable.""" + cdef VarInfo vi + self.c_ddstore.query(s2b(name), vi) + total_rows = self.c_ddstore.size(s2b(name)) + return (total_rows, vi.disp, vi.itemsize) diff --git a/src/pyddstore/torch.py b/src/pyddstore/torch.py new file mode 100644 index 0000000..33f9e85 --- /dev/null +++ b/src/pyddstore/torch.py @@ -0,0 +1,1013 @@ +"""PyTorch integration for DDStore. + +- ``DistDataset``: a map-style ``torch.utils.data.Dataset`` backed by DDStore. + Each rank loads its share of any map-style source dataset into the store; + every rank can then read any sample. Samples keep the source's structure + (a tensor/array/number, a tuple or list of them, or a dict of them) and + each field's shape and dtype. ``__getitems__`` reads a whole batch with one + ``PyDDStore.get_batch()`` per field, which PyTorch's ``DataLoader`` uses + automatically. +- ``DistDatasetReader``: the same, as a ``method=2`` extra member that joins a + dataset published by a ``DistDataset`` core group through a shared + handshake directory. +- ``ThreadDataLoader``: a ``DataLoader`` whose workers are threads instead of + forked processes (safe with MPI and GPU-resident buffers). +- ``WindowedDataset``: samples made of several stored rows (time windows, + clips), read with ``read_rows()``; ``row_of()`` maps a sample of one source + in a ``ConcatDataset`` to its row. + +Fields must have the same shape and dtype in every sample (fixed-shape). +Supported dtypes: bool, uint8, int32, int64, float32, float64. + +Import torch before mpi4py/MPI starts; ``from pyddstore.torch import ...`` +does that by itself. +""" + +import json +import logging +import multiprocessing as mp +import os +import queue +import socket +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset + +logger = logging.getLogger(__name__) + +__all__ = [ + "DistDataset", + "DistDatasetReader", + "ThreadDataLoader", + "WindowedDataset", + "row_of", +] + +_SUPPORTED = { + np.dtype(np.bool_), + np.dtype(np.uint8), + np.dtype(np.int32), + np.dtype(np.int64), + np.dtype(np.float32), + np.dtype(np.float64), +} +_TORCH_DTYPE = { + "bool": torch.bool, + "uint8": torch.uint8, + "int32": torch.int32, + "int64": torch.int64, + "float32": torch.float32, + "float64": torch.float64, +} + + +def _nsplit(n, parts): + """Contiguous index ranges [lo, hi) splitting range(n) into `parts`.""" + k, m = divmod(n, parts) + return [(i * k + min(i, m), (i + 1) * k + min(i + 1, m)) for i in range(parts)] + + +def _handshake_dir(handshake_dir): + return handshake_dir or os.environ.get("DDSTORE_HANDSHAKE_DIR") or "./ddstore_hs" + + +def _meta_path(handshake_dir, name): + # same sanitization as the C library's record files + safe = name.replace("/", "_").replace(".", "_") + return os.path.join(handshake_dir, f"{safe}.meta.json") + + +# --------------------------------------------------------------------------- +# sample structure <-> flat fields +# --------------------------------------------------------------------------- + + +def _record_descr(dtype): + """JSON-safe description of a structured dtype (see _record_dtype).""" + return json.loads(json.dumps(np.lib.format.dtype_to_descr(dtype))) + + +def _record_dtype(descr): + """Structured dtype from _record_descr's output (JSON turns tuples into + lists; numpy needs them back as tuples).""" + + def fix(d): + if isinstance(d, str): + return d + out = [] + for item in d: + name = tuple(item[0]) if isinstance(item[0], list) else item[0] + entry = (name, fix(item[1])) + out.append(entry + (tuple(item[2]),) if len(item) > 2 else entry) + return out + + return np.lib.format.descr_to_dtype(fix(descr)) + + +def _field_spec(value, where): + """(kind, numpy dtype, shape) of one leaf value.""" + # numpy structured records: stored as raw bytes (uint8), rebuilt on read + if isinstance(value, (np.ndarray, np.void)) and value.dtype.names is not None: + if isinstance(value, np.void): + kind, shape = "record", () + else: + kind = "recarray" if isinstance(value, np.recarray) else "structarray" + shape = value.shape + return { + "kind": kind, + "dtype": "uint8", + "shape": list(shape), + "record": _record_descr(value.dtype), + } + if isinstance(value, torch.Tensor): + kind, dtype, shape = ( + "torch", + np.dtype(str(value.dtype).replace("torch.", "")), + tuple(value.shape), + ) + elif isinstance(value, np.ndarray): + kind, dtype, shape = "numpy", value.dtype, value.shape + elif isinstance(value, np.generic): + kind, dtype, shape = "npscalar", value.dtype, () + elif isinstance(value, bool): + kind, dtype, shape = "py", np.dtype(np.bool_), () + elif isinstance(value, int): + kind, dtype, shape = "py", np.dtype(np.int64), () + elif isinstance(value, float): + kind, dtype, shape = "py", np.dtype(np.float64), () + else: + raise TypeError( + f"{where}: unsupported value of type {type(value).__name__}; " + "use tensors, numpy arrays or Python/numpy numbers" + ) + if dtype not in _SUPPORTED: + raise TypeError( + f"{where}: dtype {dtype} is not supported " + f"(supported: {', '.join(sorted(str(d) for d in _SUPPORTED))})" + ) + return {"kind": kind, "dtype": dtype.name, "shape": list(shape)} + + +def _flatten(sample): + """(structure, keys, leaf values) of one sample.""" + if isinstance(sample, dict): + keys = [str(k) for k in sample.keys()] + return "dict", keys, list(sample.values()) + if isinstance(sample, (tuple, list)): + structure = "tuple" if isinstance(sample, tuple) else "list" + return structure, [str(i) for i in range(len(sample))], list(sample) + return "single", ["0"], [sample] + + +def _schema_of(sample, where): + structure, keys, values = _flatten(sample) + for v in values: + if isinstance(v, (dict, list, tuple)): + raise TypeError(f"{where}: nested containers are not supported") + fields = [_field_spec(v, f"{where} field {k!r}") for k, v in zip(keys, values)] + return {"structure": structure, "keys": keys, "fields": fields} + + +def _to_numpy_row(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu().numpy().reshape(-1) + if isinstance(value, (np.ndarray, np.void)) and value.dtype.names is not None: + return np.frombuffer(np.array(value, dtype=value.dtype).tobytes(), np.uint8) + return np.asarray(value).reshape(-1) + + +# --------------------------------------------------------------------------- +# datasets +# --------------------------------------------------------------------------- + + +def _selector(fields): + """encode() keeping `fields` of a dict sample, or positions of a + tuple/list sample, in that order.""" + + def select(sample): + try: + if isinstance(sample, dict): + return {k: sample[k] for k in fields} + if isinstance(sample, (tuple, list)): + return type(sample)(sample[k] for k in fields) + except (KeyError, IndexError, TypeError) as exc: + raise ValueError(f"fields={fields}: not in the sample ({exc!r})") from None + raise ValueError("fields= needs dict, tuple or list samples") + + return select + + +class _Encoded(Dataset): + """`source` with `encode` applied to every sample.""" + + def __init__(self, source, encode): + self.source, self.encode = source, encode + + def __len__(self): + return len(self.source) + + def __getitem__(self, i): + return self.encode(self.source[i]) + + +class _StoreDataset(Dataset): + """Reads samples described by `self.schema` from `self.ddstore`.""" + + def _setup_fields(self, schema, name, device): + self.schema = schema + self.name = name + self.device = device + self._var = [f"{name}/{k}" for k in schema["keys"]] + self._record = [ + _record_dtype(f["record"]) if "record" in f else None + for f in schema["fields"] + ] + self._size = [ + int(np.prod(f["shape"], dtype=np.int64)) + * (rec.itemsize if rec is not None else 1) + for f, rec in zip(schema["fields"], self._record) + ] + self._batch_get = os.environ.get("DDSTORE_BATCH_GET", "1") != "0" + if not hasattr(self, "_decode"): + self._decode = None + + # -- shapes as the user sees them (same structure as a sample) -------- + @property + def shapes(self): + return self._rebuild([tuple(f["shape"]) for f in self.schema["fields"]]) + + @property + def dtypes(self): + # plain fields: dtype name; record fields: the structured numpy dtype + return self._rebuild( + [ + rec if rec is not None else f["dtype"] + for f, rec in zip(self.schema["fields"], self._record) + ] + ) + + def _rebuild(self, values): + s = self.schema["structure"] + if s == "single": + return values[0] + if s == "dict": + return dict(zip(self.schema["keys"], values)) + return tuple(values) if s == "tuple" else list(values) + + def _alloc(self, n, j): + f = self.schema["fields"][j] + if self.device is not None and f["kind"] == "torch": + return torch.empty( + (n, self._size[j]), dtype=_TORCH_DTYPE[f["dtype"]], device=self.device + ) + return np.empty((n, self._size[j]), dtype=f["dtype"]) + + def _value(self, row, j): + """One sample's field from its flat row (a view, no copy).""" + f = self.schema["fields"][j] + shape = tuple(f["shape"]) + if f["kind"] == "torch": + t = row if isinstance(row, torch.Tensor) else torch.from_numpy(row) + return t.reshape(shape) + if f["kind"] == "numpy": + return row.reshape(shape) + if f["kind"] == "record": + return row.view(self._record[j])[0] + if f["kind"] in ("structarray", "recarray"): + arr = row.view(self._record[j]).reshape(shape) + return arr.view(np.recarray) if f["kind"] == "recarray" else arr + if f["kind"] == "npscalar": + return row[0] + return row[0].item() + + def _batch_value(self, buf, j): + """All rows of field j in buf as one value of shape (n, *shape) (a + view, no copy). Scalar fields give a 1-D array.""" + f = self.schema["fields"][j] + n = buf.shape[0] + shape = (n,) + tuple(f["shape"]) + if f["kind"] == "torch": + t = buf if isinstance(buf, torch.Tensor) else torch.from_numpy(buf) + return t.reshape(shape) + if f["kind"] == "record": + return buf.view(self._record[j]).reshape(n) + if f["kind"] in ("structarray", "recarray"): + arr = buf.view(self._record[j]).reshape(shape) + return arr.view(np.recarray) if f["kind"] == "recarray" else arr + return buf.reshape(shape) + + def _key(self, j): + """Field j's key as in a sample: dict key, tuple/list position, or 0.""" + k = self.schema["keys"][j] + return k if self.schema["structure"] == "dict" else int(k) + + def _field_ids(self, fields): + if fields is None: + return list(range(len(self._var))) + index = {self._key(j): j for j in range(len(self._var))} + try: + return [index[k] for k in fields] + except KeyError as exc: + raise KeyError( + f"{self.name}: no field {exc.args[0]!r} (fields: {list(index)})" + ) from None + + def alloc(self, n, fields=None): + """Buffers for reads of `n` rows: a dict keyed like the sample (dict + key, or tuple/list position, or 0 for a single value) of one buffer + per field, each registered once with ``register_recv()`` so reads + into it skip memory registration. Pass to ``read_rows(out=)`` or + ``__getitems__(out=)``; reads may use the first rows only. Samples + read into these buffers are views into them: the caller decides when + a buffer can be reused. ``release()`` unregisters them.""" + bufs = {} + for j in self._field_ids(fields): + buf = self._alloc(n, j) + self.ddstore.register_recv(self._var[j], buf) + bufs[self._key(j)] = buf + return bufs + + def release(self, bufs): + """Unregister buffers from ``alloc()`` (``free()`` does it too).""" + for j in self._field_ids(list(bufs)): + self.ddstore.unregister_recv(self._var[j], bufs[self._key(j)]) + + def _read_field(self, j, idx, out): + n = len(idx) + if out is None: + buf = self._alloc(n, j) + else: + buf = out[self._key(j)] + if buf.shape[0] < n or tuple(buf.shape[1:]) != (self._size[j],): + raise ValueError( + f"{self.name}: out[{self._key(j)!r}] has shape {tuple(buf.shape)}, " + f"need at least ({n}, {self._size[j]}) (use alloc())" + ) + buf = buf[:n] + self.ddstore.get_batch(self._var[j], buf, idx) + return buf + + def read_rows(self, rows, fields=None, out=None): + """Read stored rows `rows` (global sample indices; any order, repeats + allowed) of the selected fields (default: all) with one + ``get_batch()`` per field. Returns a dict keyed like ``alloc()`` of + values shaped ``(len(rows), *field_shape)``; scalar fields give 1-D + arrays. `out`: buffers from ``alloc()`` (at least ``len(rows)`` + rows); the values are then views into them. + + Rows come back as stored: ``decode`` is not applied. + + With ``method=0`` this is collective, like ``get_batch()``: every + rank calls it the same number of times, in the same order.""" + idx = np.asarray(rows, dtype=np.int64).reshape(-1) + return { + self._key(j): self._batch_value(self._read_field(j, idx, out), j) + for j in self._field_ids(fields) + } + + def __len__(self): + return self.total_ns + + def len(self): + return self.total_ns + + def get(self, idx): + values = [] + for j, var in enumerate(self._var): + buf = self._alloc(1, j) + self.ddstore.get(var, buf, int(idx)) + values.append(self._value(buf[0], j)) + sample = self._rebuild(values) + return sample if self._decode is None else self._decode(sample, int(idx)) + + def __getitem__(self, idx): + return self.get(idx) + + def __getitems__(self, indices, out=None): + """A whole batch: one get_batch() per field (DDSTORE_BATCH_GET=0: + one get() per sample). Called by DataLoader and ThreadDataLoader. + `out`: buffers from ``alloc()``; the samples are then views into + them.""" + if not self._batch_get and out is None: + return [self.get(i) for i in indices] + idx = np.asarray(indices, dtype=np.int64) + columns = [] + for j in range(len(self._var)): + buf = self._read_field(j, idx, out) + columns.append([self._value(buf[i], j) for i in range(len(idx))]) + samples = [self._rebuild([col[i] for col in columns]) for i in range(len(idx))] + if self._decode is not None: + samples = [self._decode(sm, int(i)) for sm, i in zip(samples, idx)] + return samples + + +def _rows(values, dtype): + """Stack per-sample values into contiguous (n, size) rows of `dtype`.""" + return np.ascontiguousarray( + np.stack([_to_numpy_row(v) for v in values]).astype(dtype, copy=False) + ) + + +class DistDataset(_StoreDataset): + """A map-style dataset stored in DDStore across the ranks of `comm`. + + Args: + source: any map-style dataset (``len(source)``, ``source[i]``). Each + rank loads only its contiguous share. Every sample must have the + same structure, and each field the same shape and dtype. + name: dataset name; field ``k`` is stored as variable ``name/k``. + comm: MPI communicator (default ``MPI.COMM_WORLD``). All its ranks + must construct the dataset together. + ddstore_width: ranks per independent store (default: all of + ``comm``); each group holds a full copy of the dataset. + device: put tensor fields of read samples on this device + (GPUDirect RDMA; needs ``method`` 1/2 and ``DDSTORE_FABRIC=cxi``). + add_device: keep this rank's share of tensor fields on this device. + method: DDStore backend (default ``DDSTORE_METHOD`` or 0). + handshake_dir: ``method=2`` directory (default + ``DDSTORE_HANDSHAKE_DIR`` or ``./ddstore_hs``). + chunk_size: load this rank's share ``chunk_size`` samples at a time, + writing each chunk into the store before reading the next, so + only one chunk is held in memory besides the store (default: + read the whole share, then add it). Host storage only (no + ``add_device`` for tensor fields). + encode: ``encode(sample) -> sample`` applied to every source sample + before it is stored: pick and convert what to store (e.g. drop + metadata objects, turn a label string into an id). Its result + must meet the rules above. + decode: ``decode(stored, index) -> sample`` applied to every sample + read (``ds[i]``, ``__getitems__``), with its index: add back + what wasn't stored (constants, tables looked up by index or by + a stored id). Runs on the reading rank, in the loader thread; + anything it looks up must be on every rank. Not applied by + ``read_rows()`` or ``WindowedDataset`` (row-level reads). + fields: store only these keys (dict samples) or positions + (tuple/list samples), in this order: shorthand for an + ``encode`` that selects them. Not together with ``encode``. + + ``ds[i]`` returns a sample with the source's structure: tensors stay + tensors (on ``device`` if given), numpy arrays stay arrays, numbers stay + numbers. ``ds.ddstore`` is the underlying ``PyDDStore``. + + With ``method=0``, batched reads are collective: every rank must read the + same number of batches (DistributedSampler does that) from one thread. + """ + + def __init__( + self, + source, + name, + comm=None, + ddstore_width=None, + device=None, + add_device=None, + method=None, + handshake_dir=None, + chunk_size=None, + encode=None, + decode=None, + fields=None, + ): + super().__init__() + from mpi4py import MPI + + if fields is not None: + if encode is not None: + raise ValueError("pass either fields or encode, not both") + encode = _selector(list(fields)) + if encode is not None: + source = _Encoded(source, encode) + self._decode = decode + + from ._core import PyDDStore + + self.comm = comm if comm is not None else MPI.COMM_WORLD + self.rank = self.comm.Get_rank() + self.comm_size = self.comm.Get_size() + self.add_device = add_device + self.ddstore_width = ( + ddstore_width if ddstore_width is not None else self.comm_size + ) + self.method = ( + int(os.environ.get("DDSTORE_METHOD", "0")) + if method is None + else int(method) + ) + if self.method == 2 and self.ddstore_width != self.comm_size: + raise NotImplementedError( + "method=2 does not support ddstore_width < comm size (groups would " + "collide on the same handshake directory)" + ) + self.ddstore_comm = self.comm.Split(self.rank // self.ddstore_width, self.rank) + group_rank = self.ddstore_comm.Get_rank() + group_size = self.ddstore_comm.Get_size() + + # This rank's share of the source. Errors found locally are raised + # only after comparing with every rank (_raise_on_all), so a bad + # sample on some ranks raises on all of them instead of leaving the + # others blocked in a collective. + self.total_ns = len(source) + lo, hi = _nsplit(self.total_ns, group_size)[group_rank] + first, schema, error = None, None, None + try: + if hi > lo: + first = source[lo] + schema = _schema_of(first, f"{name}[{lo}]") + except (TypeError, ValueError) as exc: + error = exc + schemas = self._raise_on_all(error, schema) + if any(s is None for s in schemas): + raise ValueError( + f"{name}: every rank needs at least one sample (dataset has {self.total_ns})" + ) + if any(s != schemas[0] for s in schemas): + raise ValueError( + f"{name}: samples differ in structure, shape or dtype across ranks" + ) + self._setup_fields(schemas[0], name, device) + if chunk_size is not None: + if chunk_size < 1: + raise ValueError(f"chunk_size must be >= 1 (got {chunk_size})") + if add_device is not None and any( + f["kind"] == "torch" for f in self.schema["fields"] + ): + raise ValueError( + "chunk_size needs host storage: it can't be combined with add_device" + ) + + def check(i, sample): + other = _schema_of(sample, f"{name}[{i}]") + if other != self.schema: + raise ValueError( + f"{name}[{i}]: structure, shape or dtype differs from {name}[{lo}] " + f"({other} vs {self.schema}); DistDataset needs fixed-shape samples" + ) + + hs = _handshake_dir(handshake_dir) + if self.method == 2: + self.ddstore = PyDDStore(self.ddstore_comm, method=2, handshake_dir=hs) + if group_rank == 0: + # published before the variables, so a reader that sees a + # variable's record file always finds the schema too + path = _meta_path(hs, name) + tmp = f"{path}.tmp.{os.getpid()}" + with open(tmp, "w") as fh: + json.dump({"total_ns": self.total_ns, **self.schema}, fh) + os.replace(tmp, path) + else: + self.ddstore = PyDDStore(self.ddstore_comm, method=self.method) + + fields = self.schema["fields"] + if chunk_size is None: + # Whole share at once: read, check, then add() each field. + samples, error = [first], None + try: + for i in range(lo + 1, hi): + samples.append(source[i]) + check(i, samples[-1]) + except (TypeError, ValueError) as exc: + error = exc + self._raise_on_all(error) + for j, var in enumerate(self._var): + values = [_flatten(s)[2][j] for s in samples] + if add_device is not None and fields[j]["kind"] == "torch": + rows = ( + torch.stack([v.reshape(-1) for v in values]) + .to(add_device) + .contiguous() + ) + else: + rows = _rows(values, fields[j]["dtype"]) + self.ddstore.add(var, rows) + else: + # Chunked: allocate every field (init, collective), then copy the + # share in chunks of chunk_size samples (update, local), so at most + # one chunk is held in memory besides the store itself. + for j, var in enumerate(self._var): + itemsize = np.dtype(fields[j]["dtype"]).itemsize + self.ddstore.init(var, hi - lo, self._size[j], itemsize) + error = None + try: + for start in range(lo, hi, chunk_size): + stop = min(start + chunk_size, hi) + chunk = [ + first if i == lo else source[i] for i in range(start, stop) + ] + for i, sample in zip(range(start, stop), chunk): + if i != lo: + check(i, sample) + for j, var in enumerate(self._var): + values = [_flatten(sm)[2][j] for sm in chunk] + self.ddstore.update( + var, _rows(values, fields[j]["dtype"]), start - lo + ) + if start == lo: + first = None # held only for the first chunk + except (TypeError, ValueError) as exc: + error = exc + # also makes sure every rank has filled its share before any reads + self._raise_on_all(error) + logger.debug( + "DistDataset %s: rank %d holds [%d, %d) of %d", + name, + self.rank, + lo, + hi, + self.total_ns, + ) + + def _raise_on_all(self, error, value=None): + """Allgather (value, error) over comm; if any rank had an error, + raise it on every rank. Returns the gathered values.""" + local = None if error is None else (type(error).__name__, str(error)) + gathered = self.comm.allgather((value, local)) + errors = [e for _, e in gathered if e is not None] + if errors: + cls = TypeError if errors[0][0] == "TypeError" else ValueError + raise cls(errors[0][1]) + return [v for v, _ in gathered] + + +class DistDatasetReader(_StoreDataset): + """A ``DistDataset`` published by a ``method=2`` core group, joined from a + separate job (no MPI communicator needed). + + Args: + name: the core group's dataset name. + handshake_dir: shared directory (default ``DDSTORE_HANDSHAKE_DIR`` or + ``./ddstore_hs``). + n_core: number of core ranks (default ``DDSTORE_N_CORE``). + device: put tensor fields of read samples on this device. + decode: as for ``DistDataset`` (the core group's ``encode`` already + ran before storing). + + Waits up to ``DDSTORE_HANDSHAKE_TIMEOUT_S`` (default 300 s) for the core + group to publish. + """ + + def __init__(self, name, handshake_dir=None, n_core=None, device=None, decode=None): + super().__init__() + self._decode = decode + from ._core import PyDDStore + + hs = _handshake_dir(handshake_dir) + if n_core is None: + if "DDSTORE_N_CORE" not in os.environ: + raise ValueError( + "DistDatasetReader needs the number of core ranks: pass n_core= " + "or set DDSTORE_N_CORE" + ) + n_core = int(os.environ["DDSTORE_N_CORE"]) + timeout = float(os.environ.get("DDSTORE_HANDSHAKE_TIMEOUT_S", "300")) + path = _meta_path(hs, name) + t0 = time.monotonic() + while not os.path.exists(path): + if time.monotonic() - t0 > timeout: + raise TimeoutError( + f"no dataset {name!r} published in {hs} after {timeout:.0f} s" + ) + time.sleep(0.05) + with open(path) as fh: + meta = json.load(fh) + self.total_ns = meta.pop("total_ns") + self._setup_fields(meta, name, device) + self.ddstore = PyDDStore(None, method=2, handshake_dir=hs, n_core=n_core) + for var in self._var: + self.ddstore.join(var) + + +class WindowedDataset(Dataset): + """Samples made of several stored rows of `ds` (a ``DistDataset`` or + ``DistDatasetReader``): time windows, clips, sequences. Each stored row + is held once, however many windows use it. + + Sample ``i`` is rows ``s, s + dilation, ..., s + (window - 1) * dilation`` + with ``s = starts[i]`` if `starts` is given, else ``s = i * stride``. + Use `starts` to keep only windows that don't cross a boundary between + trajectories or files. Each field comes back stacked, shaped + ``(window, *field_shape)``, in the structure of `ds`'s samples, or as a + dict of the selected `fields`. + + ``__getitems__`` reads a whole batch of windows with one + ``read_rows()`` (one ``get_batch()`` per field), so with ``method=0`` the + same collective rule applies as for ``ds``. + """ + + def __init__(self, ds, window, stride=1, dilation=1, starts=None, fields=None): + if window < 1 or stride < 1 or dilation < 1: + raise ValueError("window, stride and dilation must be >= 1") + self.ds = ds + self.window = int(window) + self.dilation = int(dilation) + self.fields = None if fields is None else list(fields) + self._offsets = np.arange(self.window, dtype=np.int64) * self.dilation + span = int(self._offsets[-1]) + 1 + if starts is not None: + self.starts = np.asarray(starts, dtype=np.int64).reshape(-1) + bad = (self.starts < 0) | (self.starts + span > len(ds)) + if bad.any(): + raise IndexError( + f"start {int(self.starts[bad][0])} + window span {span} is out " + f"of range for {len(ds)} rows" + ) + else: + n = (len(ds) - span) // int(stride) + 1 if len(ds) >= span else 0 + self.starts = np.arange(n, dtype=np.int64) * int(stride) + + def __len__(self): + return len(self.starts) + + def _rows(self, indices): + return (self.starts[np.asarray(indices, dtype=np.int64)][:, None] + + self._offsets).reshape(-1) + + def _sample(self, values): + if self.fields is not None: + return values + ds = self.ds + return ds._rebuild([values[ds._key(j)] for j in range(len(ds._var))]) + + def __getitem__(self, i): + if not -len(self) <= i < len(self): + raise IndexError(f"window {i} out of range ({len(self)} windows)") + return self._sample(self.ds.read_rows(self._rows([i]), self.fields)) + + def __getitems__(self, indices): + cols = self.ds.read_rows(self._rows(indices), self.fields) + w = self.window + return [ + self._sample({k: v[b * w : (b + 1) * w] for k, v in cols.items()}) + for b in range(len(indices)) + ] + + +def row_of(concat, source, index): + """The row of sample `index` of source `source` in + ``torch.utils.data.ConcatDataset`` `concat`, i.e. its index in a + ``DistDataset`` built over `concat`. Use it to map (file, trajectory, + step) to a row when several sources share one store.""" + sizes = concat.cumulative_sizes + if not 0 <= source < len(sizes): + raise IndexError(f"source {source} out of range ({len(sizes)} sources)") + first = sizes[source - 1] if source > 0 else 0 + if not 0 <= index < sizes[source] - first: + raise IndexError( + f"index {index} out of range for source {source} " + f"({sizes[source] - first} samples)" + ) + return first + index + + +# --------------------------------------------------------------------------- +# loader +# --------------------------------------------------------------------------- + + +class ThreadDataLoader(DataLoader): + """A ``DataLoader`` that fetches batches in a thread pool instead of forked + worker processes. Threads share the process's MPI state, CUDA context and + Python objects, so it is safe with DDStore and GPU-resident buffers, where + forked workers are not. Takes the same arguments as ``DataLoader``; + ``num_workers`` is the number of threads (0 means 1). + + Each batch is fetched (via ``dataset.__getitems__`` when present), + collated and optionally pinned in a worker thread; at most + ``num_workers * prefetch_factor`` batches are in flight. Random draws + match ``DataLoader``'s. ``DDSTORE_AFFINITY_WIDTH`` / + ``DDSTORE_AFFINITY_OFFSET`` pin worker thread *i* to CPUs + ``[offset + i*width, offset + (i+1)*width)`` of the process's affinity. + + ``reuse_buffers=True`` (dataset with ``alloc()``, e.g. ``DistDataset``): + read every batch into one of a fixed pool of ``num_workers`` buffer sets + from ``dataset.alloc(batch_size)``, registered once, instead of fresh + buffers that are registered on every read. A worker takes a set, reads + and collates the batch, and returns the set, so the collate must copy: + needs ``batch_size`` and the default ``collate_fn``, or + ``collate_copies=True`` to declare that a custom ``collate_fn`` copies. + ``close()`` (or deleting the loader) waits for running fetches and + unregisters the pool. + """ + + def __init__(self, dataset, reuse_buffers=False, collate_copies=False, **kwargs): + super().__init__(dataset, **kwargs) + + # Fixed pool of registered read buffers, one set per worker thread, + # owned by this loader (see reuse_buffers in the class docstring). + self._pool, self._pool_sets = None, [] + if reuse_buffers: + if not hasattr(dataset, "alloc"): + raise TypeError( + "reuse_buffers needs a dataset with alloc() (DistDataset, " + "DistDatasetReader)" + ) + if self.batch_size is None: + raise ValueError( + "reuse_buffers needs batch_size: without auto-collation the " + "batch is not copied out of the reused buffer" + ) + if ( + self.collate_fn is not torch.utils.data.default_collate + and not collate_copies + ): + raise ValueError( + "reuse_buffers with a custom collate_fn: pass collate_copies=True " + "if it copies the samples (the buffer is reused right after it)" + ) + self._pool = queue.Queue() + for _ in range(self.num_workers or 1): + bufs = dataset.alloc(self.batch_size) + self._pool_sets.append(bufs) + self._pool.put(bufs) + + # Persistent across epochs -- recreating the pool in every __iter__() + # would leak OS threads since the old pool is never shut down. + self._counter = mp.Value("i", 0) + self.executor = ThreadPoolExecutor( + max_workers=self.num_workers or 1, + initializer=self.worker_init, + initargs=(self._counter,), + ) + + logger.debug("num_workers: %s", self.num_workers) + logger.debug("len: %s", len(self._index_sampler)) + + @staticmethod + def worker_init(counter): + core_width = int(os.environ.get("DDSTORE_AFFINITY_WIDTH", "0")) + core_offset = int(os.environ.get("DDSTORE_AFFINITY_OFFSET", "0")) + if core_width <= 0 or not hasattr(os, "sched_getaffinity"): + return 0 + + with counter.get_lock(): + wid = counter.value + counter.value += 1 + + affinity = list(os.sched_getaffinity(0)) + affinity_mask = set( + affinity[ + core_width * wid + core_offset : core_width * (wid + 1) + core_offset + ] + ) + if affinity_mask: + os.sched_setaffinity(0, affinity_mask) + logger.debug( + "Worker: pid=%s hostname=%s ID=%s affinity=%s", + os.getpid(), + socket.gethostname(), + wid, + os.sched_getaffinity(0), + ) + return 0 + + @staticmethod + def fetch( + dataset, + ibatch, + index, + collate_fn=None, + pin_memory=False, + auto_collation=True, + pool=None, + ): + if pool is not None: + # Read into one of the loader's registered buffer sets; collate + # copies the batch out, then the set goes back to the pool. For a + # GPU buffer, get_batch() synchronizes the device before reading, + # so the collate's copy has finished before the set is refilled. + bufs = pool.get() + try: + batch = collate_fn(dataset.__getitems__(index, out=bufs)) + finally: + pool.put(bufs) + if pin_memory: + batch = torch.utils.data._utils.pin_memory.pin_memory(batch) + return (ibatch, batch) + # Collate here, in the worker, before pinning: pinning per-sample + # tensors and collating afterwards would just torch.stack them into + # a new, unpinned tensor. Use the dataset's whole-batch fetch when it + # has one, like torch's own map-style fetcher. With batch_size=None + # (no auto-collation) the sampler's index goes to dataset[index] as + # is, also as torch does (e.g. samplers that yield whole batches). + if not auto_collation: + batch = dataset[index] + elif getattr(dataset, "__getitems__", None): + batch = dataset.__getitems__(index) + else: + batch = [dataset[i] for i in index] + if collate_fn is not None: + batch = collate_fn(batch) + if pin_memory: + batch = torch.utils.data._utils.pin_memory.pin_memory(batch) + return (ibatch, batch) + + def __iter__(self): + """A new iterator over one epoch. Each has its own sampler position + and queue (the thread pool is shared), so several iterators over one + loader don't interfere, as with ``DataLoader``.""" + return _ThreadLoaderIter(self) + + def close(self): + """Stop the worker threads; with ``reuse_buffers``, wait for running + fetches first, then unregister the buffer pool. Called on deletion.""" + executor = getattr(self, "executor", None) + sets = getattr(self, "_pool_sets", []) + if executor is not None: + executor.shutdown(wait=bool(sets), cancel_futures=True) + for bufs in sets: + try: + self.dataset.release(bufs) + except (ValueError, KeyError): + pass # the store was freed first: nothing left to unregister + self._pool_sets = [] + + def __del__(self): + self.close() + + +class _ThreadLoaderIter: + """One epoch of a ThreadDataLoader; ``__iter__`` returns itself.""" + + def __init__(self, loader): + self.loader = loader + self._sampler_iter = iter(loader._index_sampler) + # torch's DataLoader iterator draws a base seed from the global RNG + # here, every epoch; draw it too, so the training loop's later random + # draws are the same as with DataLoader + torch.empty((), dtype=torch.int64).random_(generator=loader.generator) + self.fs = queue.Queue() + self.fs_iter = iter(self.fs.get, None) + self._num_yielded = 0 + self._next_batch_i = 0 + self._inflight = 0 + self._sampler_exhausted = False + # Bound how many batches can be in flight (submitted but not yet + # consumed via __next__) at once, instead of submitting the whole + # epoch up front -- keeps memory use (GPU tensors included) bounded + # regardless of dataset size. Mirrors torch's own prefetch_factor + # (default 2 per worker). + self._max_inflight = max( + 1, (loader.num_workers or 1) * (loader.prefetch_factor or 2) + ) + self._refill() + + def __iter__(self): + return self + + def __len__(self): + return len(self.loader) + + def _refill(self): + loader = self.loader + while self._inflight < self._max_inflight: + try: + index = next(self._sampler_iter) + except StopIteration: + if not self._sampler_exhausted: + self._sampler_exhausted = True + self.fs.put(None) + return + future = loader.executor.submit( + loader.fetch, + loader.dataset, + self._next_batch_i, + index, + collate_fn=loader.collate_fn, + pin_memory=loader.pin_memory, + auto_collation=loader._auto_collation, + pool=loader._pool, + ) + self.fs.put(future) + self._next_batch_i += 1 + self._inflight += 1 + + def __next__(self): + # Submit the replacement as soon as this batch is taken, as torch's + # DataLoader does, so num_workers * prefetch_factor batches are being + # fetched while the caller's training step runs (plus the one it + # holds). Refilling at the start of the *next* call instead would leave + # one slot idle during every step. + future = next(self.fs_iter) + ibatch, data = future.result() + self._inflight -= 1 + self._num_yielded += 1 + self._refill() + return data + + def close(self): + """Cancel the batches still queued (an epoch stopped early).""" + # Without blocking: the end marker (None) is queued only once the + # sampler is exhausted, so an epoch that stopped early has none, and + # waiting for it (iter(self.fs.get, None)) would block forever. Only + # this thread puts into fs, so qsize() is exact here. + while self.fs.qsize() > 0: + future = self.fs.get_nowait() + if future is not None: + future.cancel() + + def __del__(self): + if hasattr(self, "fs"): + self.close() diff --git a/test/conftest.py b/test/conftest.py index 09ff14a..0112905 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,12 +1,16 @@ -import mpi4py - -mpi4py.rc.thread_level = "serialized" -mpi4py.rc.threads = False - import pytest from mpi4py import MPI @pytest.fixture(scope="function") def comm(): - return MPI.COMM_WORLD + """Provide MPI.COMM_WORLD and barrier after each test. + + The barrier ensures all ranks finish the current test (including + store.free() and any fabric endpoint teardown) before any rank + begins the next test's handshake/MPI_Allgather. Without this, + CXI endpoint cleanup on a fast rank can desynchronize the ranks + enough to deadlock the next test's collective in add(). + """ + yield MPI.COMM_WORLD + MPI.COMM_WORLD.Barrier() diff --git a/test/test_get_batch.py b/test/test_get_batch.py new file mode 100644 index 0000000..8c59bc5 --- /dev/null +++ b/test/test_get_batch.py @@ -0,0 +1,329 @@ +""" +Batched get (PyDDStore.get_batch) tests — run with 2+ ranks, e.g.: + mpirun -n 4 pytest test/test_get_batch.py -v + +Each case runs with method 0 (MPI RMA) and, where a CXI device is present, +method 1 over libfabric: cxi by default, or the provider named by +DDSTORE_FABRIC if it is set (e.g. DDSTORE_FABRIC=hsn). Every rank fills its shard with values that encode the +global row id, so any misplaced or missing row is caught exactly. +""" + +import glob +import os +import threading + +import numpy as np +import pytest +from mpi4py import MPI + +import pyddstore as dds + +# A CXI device alone isn't enough (login nodes have one but can't open the +# fabric); also require running inside a Slurm job step. +HAVE_CXI = bool(glob.glob("/dev/cxi*")) and "SLURM_STEP_ID" in os.environ +FABRIC = os.environ.get("DDSTORE_FABRIC", "cxi") +METHODS = [ + 0, + pytest.param(1, marks=pytest.mark.skipif(not HAVE_CXI, reason="no CXI device")), +] + +try: + import torch + + HAVE_GPU = torch.cuda.is_available() +except ImportError: # pragma: no cover + torch = None + HAVE_GPU = False + +NROWS, NCOLS = 16, 5 + + +def all_passed(comm, local_ok): + return comm.allreduce(int(local_ok), op=MPI.LAND) + + +def make_store(comm, method, monkeypatch, dtype=np.float32): + if method != 0: + monkeypatch.setenv("DDSTORE_FABRIC", FABRIC) + rank = comm.Get_rank() + store = dds.PyDDStore(comm, method=method) + # row r (global) holds r*100 + column, so every element is identifiable + first = rank * NROWS + data = (np.arange(first, first + NROWS)[:, None] * 100 + np.arange(NCOLS)).astype( + dtype + ) + store.add("x", data) + store.epoch_begin() + return store + + +def expected_rows(idx, dtype=np.float32): + idx = np.asarray(idx) + return (idx[:, None] * 100 + np.arange(NCOLS)).astype(dtype) + + +def finish(store): + store.epoch_end() + store.free() + + +@pytest.mark.parametrize("method", METHODS) +def test_batch_all_ranks_shuffled_with_repeats(comm, monkeypatch, method): + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks") + store = make_store(comm, method, monkeypatch) + rng = np.random.default_rng(comm.Get_rank()) + idx = rng.integers(0, NROWS * size, size=200) # spans every rank, with repeats + out = np.zeros((len(idx), NCOLS), dtype=np.float32) + store.get_batch("x", out, idx) + ok = np.array_equal(out, expected_rows(idx)) + comm.Barrier() + finish(store) + assert all_passed(comm, ok) + + +@pytest.mark.parametrize("method", METHODS) +def test_batch_matches_per_row_get(comm, monkeypatch, method): + size = comm.Get_size() + store = make_store(comm, method, monkeypatch) + idx = list(range(NROWS * size))[::-1] + out = np.zeros((len(idx), NCOLS), dtype=np.float32) + store.get_batch("x", out, idx) + row = np.zeros((1, NCOLS), dtype=np.float32) + ok = True + for i, g in enumerate(idx): + store.get("x", row, g) + ok &= np.array_equal(row[0], out[i]) + comm.Barrier() + finish(store) + assert all_passed(comm, ok) + + +@pytest.mark.parametrize("method", METHODS) +def test_batch_single_row(comm, monkeypatch, method): + size = comm.Get_size() + store = make_store(comm, method, monkeypatch) + g = (comm.Get_rank() + 1) % size * NROWS + 3 + out = np.zeros((1, NCOLS), dtype=np.float32) + store.get_batch("x", out, [g]) + ok = np.array_equal(out, expected_rows([g])) + comm.Barrier() + finish(store) + assert all_passed(comm, ok) + + +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize( + "dtype", [np.uint8, np.int32, np.float32, np.int64, np.float64] +) +def test_batch_dtypes(comm, monkeypatch, method, dtype): + size = comm.Get_size() + store = make_store(comm, method, monkeypatch, dtype=dtype) + # row r holds r*100 + col: only rows 0..2 fit in uint8 + idx = [2, 0, 1] if dtype == np.uint8 else [NROWS * size - 1, 0, NROWS * size // 2] + out = np.zeros((len(idx), NCOLS), dtype=dtype) + store.get_batch("x", out, idx) + ok = np.array_equal(out, expected_rows(idx, dtype)) + comm.Barrier() + finish(store) + assert all_passed(comm, ok) + + +@pytest.mark.parametrize("method", METHODS) +def test_batch_errors_leave_store_usable(comm, monkeypatch, method): + size = comm.Get_size() + store = make_store(comm, method, monkeypatch) + out = np.zeros((2, NCOLS), dtype=np.float32) + with pytest.raises(IndexError): + store.get_batch("x", out, [0, NROWS * size]) # second index out of range + with pytest.raises(ValueError): + store.get_batch("x", out, [0, 1, 2]) # 3 indices for 2 rows + with pytest.raises(Exception): + store.get_batch("x", out.astype(np.float64), [0, 1]) # wrong item size + idx = [NROWS * size - 1, 0] + store.get_batch("x", out, idx) + ok = np.array_equal(out, expected_rows(idx)) + comm.Barrier() + finish(store) + assert all_passed(comm, ok) + + +@pytest.mark.skipif( + not (HAVE_CXI and HAVE_GPU and FABRIC == "cxi"), + reason="requires the cxi provider and a GPU", +) +def test_batch_into_gpu_tensor(comm, monkeypatch): + size = comm.Get_size() + store = make_store(comm, 1, monkeypatch) + rng = np.random.default_rng(100 + comm.Get_rank()) + ok = True + for _ in range(20): + idx = rng.integers(0, NROWS * size, size=64) + out = torch.empty((len(idx), NCOLS), dtype=torch.float32, device="cuda") + store.get_batch("x", out, idx) + # compute-kernel read, like a training step would do + diff = (out - torch.from_numpy(expected_rows(idx)).cuda()).abs().sum().item() + ok &= diff == 0.0 + comm.Barrier() + finish(store) + assert all_passed(comm, ok) + + +@pytest.mark.parametrize( + "method", + [pytest.param(1, marks=pytest.mark.skipif(not HAVE_CXI, reason="no CXI device"))], +) +def test_batch_concurrent_threads(comm, monkeypatch, method): + """get_batch from several threads at once on one variable: the + per-variable lock must keep each batch's rows and completions apart.""" + size = comm.Get_size() + store = make_store(comm, method, monkeypatch) + errors, bad = [], [] + + def worker(seed): + rng = np.random.default_rng(seed) + try: + for _ in range(50): + idx = rng.integers(0, NROWS * size, size=32) + out = np.zeros((len(idx), NCOLS), dtype=np.float32) + store.get_batch("x", out, idx) + if not np.array_equal(out, expected_rows(idx)): + bad.append(seed) + except Exception as exc: # noqa: BLE001 - surface any thread exception + errors.append(exc) + + threads = [ + threading.Thread(target=worker, args=(1000 * comm.Get_rank() + t,)) + for t in range(4) + ] + for t in threads: + t.start() + for t in threads: + t.join() + comm.Barrier() + finish(store) + assert not errors, f"worker thread(s) raised: {errors}" + assert all_passed(comm, not bad) + + +def _mr_miss(store): + """recv registrations so far (None unless DDSTORE_PROFILE was set at start).""" + if os.environ.get("DDSTORE_PROFILE", "0") in ("", "0"): + return None + return store.get_profile("x")["mr_miss"] + + +@pytest.mark.parametrize("method", METHODS) +def test_register_recv_pool(comm, monkeypatch, method): + """Reads into registered buffers, from several threads, each with its own + buffer: every row correct and (method 1) no registration per read.""" + size = comm.Get_size() + store = make_store(comm, method, monkeypatch) + nthreads, batch = 2, 16 + pools = [np.zeros((4 * batch, NCOLS), dtype=np.float32) for _ in range(nthreads)] + for pool in pools: + store.register_recv("x", pool) + store.register_recv("x", pools[0]) # registering twice is a no-op + miss0 = _mr_miss(store) + errors, bad = [], [] + + def worker(t): + rng = np.random.default_rng(1000 * comm.Get_rank() + t) + pool = pools[t] + try: + for it in range(40): + k = it % 4 + out = pool[k * batch : (k + 1) * batch] # a slice of the pool + idx = rng.integers(0, NROWS * size, size=batch) + store.get_batch("x", out, idx) + if not np.array_equal(out, expected_rows(idx)): + bad.append((t, "batch")) + one = pool[k * batch : k * batch + 1] + store.get("x", one, int(idx[0])) + if not np.array_equal(one, expected_rows(idx[:1])): + bad.append((t, "get")) + except Exception as exc: # noqa: BLE001 - surface any thread exception + errors.append(exc) + + if method == 0: # get_batch is collective there: one thread at a time + for t in range(nthreads): + worker(t) + else: + threads = [threading.Thread(target=worker, args=(t,)) for t in range(nthreads)] + for t in threads: + t.start() + for t in threads: + t.join() + ok = not errors and not bad + if method != 0 and miss0 is not None: + ok &= _mr_miss(store) == miss0 + for pool in pools: + store.unregister_recv("x", pool) + if method != 0: + with pytest.raises(ValueError): + store.unregister_recv("x", pools[0]) + # unregistered buffers still work (through the one-slot cache) + idx = np.arange(batch) % (NROWS * size) + store.get_batch("x", pools[1][:batch], idx) + ok &= np.array_equal(pools[1][:batch], expected_rows(idx)) + comm.Barrier() + finish(store) + assert not errors, f"worker thread(s) raised: {errors}" + assert all_passed(comm, ok), bad + + +@pytest.mark.skipif( + not (HAVE_CXI and HAVE_GPU and FABRIC == "cxi"), + reason="requires the cxi provider and a GPU", +) +def test_register_recv_gpu(comm, monkeypatch): + size = comm.Get_size() + store = make_store(comm, 1, monkeypatch) + pool = torch.empty((64, NCOLS), dtype=torch.float32, device="cuda") + store.register_recv("x", pool) + miss0 = _mr_miss(store) + rng = np.random.default_rng(200 + comm.Get_rank()) + ok = True + for it in range(20): + k = it % 4 + out = pool[k * 16 : (k + 1) * 16] + idx = rng.integers(0, NROWS * size, size=16) + store.get_batch("x", out, idx) + diff = (out - torch.from_numpy(expected_rows(idx)).cuda()).abs().sum().item() + ok &= diff == 0.0 + if miss0 is not None: + ok &= _mr_miss(store) == miss0 + store.unregister_recv("x", pool) + comm.Barrier() + finish(store) + assert all_passed(comm, ok) + + +WIDE = 3001 # float32 columns: 12004-byte rows + + +@pytest.mark.parametrize("method", METHODS) +def test_wide_rows(comm, monkeypatch, method): + """Rows of 12004 bytes. Run with DDSTORE_MAX_READ_BYTES=4096 (method 1) + to check that rows longer than one read are split, remainder included.""" + size = comm.Get_size() + if method != 0: + monkeypatch.setenv("DDSTORE_FABRIC", FABRIC) + rank = comm.Get_rank() + store = dds.PyDDStore(comm, method=method) + first = rank * 4 + rows = np.arange(first, first + 4)[:, None] * 10000.0 + np.arange(WIDE) + store.add("w", rows.astype(np.float32)) + store.epoch_begin() + idx = np.array([(rank + 1) % size * 4 + 3, 0, size * 4 - 1]) + out = np.zeros((len(idx), WIDE), dtype=np.float32) + store.get_batch("w", out, idx) + want = (idx[:, None] * 10000.0 + np.arange(WIDE)).astype(np.float32) + ok = np.array_equal(out, want) + one = np.zeros((1, WIDE), dtype=np.float32) + store.get("w", one, int(idx[0])) + ok &= np.array_equal(one, want[:1]) + comm.Barrier() + finish(store) + assert all_passed(comm, ok) diff --git a/test/test_gpu_rdma.py b/test/test_gpu_rdma.py new file mode 100644 index 0000000..f836f01 --- /dev/null +++ b/test/test_gpu_rdma.py @@ -0,0 +1,561 @@ +""" +GPUDirect RDMA tests: GPU destination (Phase 1), GPU source (Phase 2), both, +negative paths, and concurrent get() from multiple threads. + +Positive path — run with: DDSTORE_FABRIC=cxi mpirun -n 2 pytest test/test_gpu_rdma.py -v +requires a live cxi/Slingshot fabric and at least one visible GPU per rank. +Negative-path tests need neither and always run. +""" + +import threading + +import numpy as np +import pytest +from mpi4py import MPI + +import pyddstore as dds + +torch = pytest.importorskip("torch") + +gpu_required = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires a ROCm/HIP GPU" +) + + +def all_passed(comm, local_ok): + return comm.allreduce(int(local_ok), op=MPI.LAND) + + +# --------------------------------------------------------------------------- +# positive path: host source -> poisoned GPU destination, over cxi +# --------------------------------------------------------------------------- + + +@gpu_required +def test_get_into_gpu_tensor_cxi(comm, monkeypatch): + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + + store = dds.PyDDStore(comm, method=1) + data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + store.add("x", data) # host source (Phase 1 scope — unchanged) + + store.epoch_begin() + local_ok = True + for target_rank in range(size): + # poison, not zeros: a silent no-op/host-staged-fallback bug would + # leave this value in place instead of the real remote data. + out = torch.full((1, ncols), -999.0, dtype=torch.float32, device="cuda") + store.get("x", out, start=target_rank * nrows) + expected = float(target_rank + 1) + ok = bool(torch.all(out.cpu() == expected)) + print( + f"[rank {rank}] target_rank={target_rank} expected={expected} " + f"got={out.cpu().tolist()} ok={ok}", + flush=True, + ) + if not ok: + local_ok = False + store.epoch_end() + + assert all_passed(comm, local_ok) + store.free() + + +@gpu_required +def test_get_into_gpu_tensor_cxi_compute_kernel_read(comm, monkeypatch): + """Diagnostic: does reading the RDMA destination via a GPU COMPUTE + KERNEL (not a .cpu() DMA-engine copy) crash/fault, unlike every other + test in this file which always reads back via .cpu()? A real training + loop (examples/vae/vae-ddp.py --gpu-dest) feeds the destination tensor + directly into model(data) -- a compute-kernel read -- and hit + HSA_STATUS_ERROR_EXCEPTION hardware faults at real batch-loop scale, + something none of the .cpu()-based tests in this file have ever + reproduced. Hypothesis: the NIC's P2P write into GPU memory isn't + visible/coherent to compute cores (cache/TLB gap) the way it is to the + DMA engine .cpu() uses -- untested by every other test here. Also loops + many iterations back-to-back (no pauses) to mirror DataLoader's rapid + per-sample get() calls, in case repeated register/deregister at a + reused address (PyTorch's allocator likely returns the same block each + time for same-shape torch.empty() in a tight loop) is a contributing + factor rather than the compute-kernel read alone. + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + n_iters = 200 + + store = dds.PyDDStore(comm, method=1) + data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + store.add("x", data) + + store.epoch_begin() + target_rank = (rank + 1) % size + expected = float(target_rank + 1) + for i in range(n_iters): + # torch.empty (not full/poisoned): mirrors DistDataset.get()'s real + # allocation exactly, and a fresh, uninitialized block is what a + # compute kernel would actually read if the transfer no-op'd -- + # closer to the real crash scenario than a poisoned buffer. + out = torch.empty((1, ncols), dtype=torch.float32, device="cuda") + store.get("x", out, start=target_rank * nrows) + # GPU compute-kernel read (not .cpu()): elementwise op launches a + # real kernel touching `out`'s memory from the compute cores. + diff = (out - expected).abs().sum() + # Forces the host to wait for the kernel and surfaces any async + # HIP error at this point (torch raises a RuntimeError mentioning + # the HIP error, or the process aborts, same as the real crash). + torch.cuda.synchronize() + if i % 50 == 0: + print(f"[rank {rank}] iter={i} diff={diff.item()}", flush=True) + store.epoch_end() + print( + f"[rank {rank}] completed {n_iters} iterations without a HIP error", flush=True + ) + # Wait for every rank's reads before tearing down: otherwise a fast rank + # can free its endpoint while a peer is still reading (PTLTE_NOT_FOUND). + comm.Barrier() + store.free() + + +@gpu_required +def test_get_into_gpu_tensor_cxi_matrix(comm, monkeypatch): + """Isolates which factor actually determines pass/fail: buffer + allocation method (torch.empty, uninitialized vs torch.full, poisoned + via a GPU compute-kernel write) crossed with readback method (.cpu() + DMA copy vs GPU compute-kernel read + torch.cuda.synchronize()). + Written when test_get_into_gpu_tensor_cxi (poison + .cpu()) reliably + failed on Frontier while test_get_into_gpu_tensor_cxi_compute_kernel_read + (empty + compute-kernel read) passed; runs all 4 combinations to show + which axis (allocation vs readback) matters. Both now pass, since + PyDDStore.get() synchronizes the device before every GPU-destination + transfer (see test_get_into_gpu_tensor_cxi_sync_before_get). + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + POISON = -999.0 + + store = dds.PyDDStore(comm, method=1) + data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + store.add("x", data) + store.epoch_begin() + + target_rank = (rank + 1) % size + expected = float(target_rank + 1) + results = {} + for alloc in ("empty", "poison"): + for readback in ("cpu", "kernel"): + if alloc == "empty": + out = torch.empty((1, ncols), dtype=torch.float32, device="cuda") + else: + out = torch.full((1, ncols), POISON, dtype=torch.float32, device="cuda") + store.get("x", out, start=target_rank * nrows) + if readback == "cpu": + snapshot = out.cpu() + ok = bool(torch.all(snapshot == expected)) + detail = snapshot.tolist() + else: + diff = (out - expected).abs().sum() + torch.cuda.synchronize() + ok = bool(diff.item() == 0.0) + detail = f"diff={diff.item()}" + key = f"alloc={alloc},readback={readback}" + results[key] = ok + print(f"[rank {rank}] {key} ok={ok} detail={detail}", flush=True) + + store.epoch_end() + gathered = comm.gather(results, root=0) + if rank == 0: + print(f"[rank 0] ALL RESULTS: {gathered}", flush=True) + # gather() doesn't hold back non-root ranks; wait for every rank's reads + # before tearing down (see test_get_into_gpu_tensor_cxi_compute_kernel_read). + comm.Barrier() + store.free() + # Fail loudly with the full matrix visible in the log even if only one + # combination is wrong -- this test is diagnostic, not a pass/fail gate. + assert all(results.values()), f"[rank {rank}] matrix results: {results}" + + +@gpu_required +def test_get_into_gpu_tensor_cxi_sync_before_get(comm, monkeypatch): + """Follow-up to test_get_into_gpu_tensor_cxi_matrix's finding: RDMA into + a torch.full()-poisoned (GPU-kernel-written) buffer fails, but into a + torch.empty() (untouched) buffer works -- readback method is + irrelevant. Hypothesis: a cache-coherency gap where the NIC's RDMA + write doesn't invalidate whatever the GPU cache still holds from the + prior compute-kernel write. Tests whether an explicit + torch.cuda.synchronize() between the poisoning kernel and the RDMA + get() call (forcing the kernel write to fully retire/flush first) is + enough to fix it. + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + POISON = -999.0 + + store = dds.PyDDStore(comm, method=1) + data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + store.add("x", data) + store.epoch_begin() + + target_rank = (rank + 1) % size + expected = float(target_rank + 1) + + out = torch.full((1, ncols), POISON, dtype=torch.float32, device="cuda") + torch.cuda.synchronize() # <-- the fix under test: flush the poison write first + store.get("x", out, start=target_rank * nrows) + snapshot = out.cpu() + ok = bool(torch.all(snapshot == expected)) + print(f"[rank {rank}] sync-before-get: ok={ok} got={snapshot.tolist()}", flush=True) + + store.epoch_end() + # Wait for every rank's reads before tearing down (PTLTE_NOT_FOUND otherwise). + comm.Barrier() + store.free() + assert all_passed(comm, ok) + + +@gpu_required +def test_get_into_gpu_tensor_cxi_large(comm, monkeypatch): + """Diagnostic: same as test_get_into_gpu_tensor_cxi but with a transfer + well over FI_CXI_SAFE_DEVMEM_COPY_THRESHOLD (default 4096 bytes), to + check whether CXI's small-transfer 'safe load/store' HMEM path is what's + silently no-op'ing, vs. the registration approach itself being broken. + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4096 # 4096 floats/row = 16384 bytes >> 4096-byte threshold + + store = dds.PyDDStore(comm, method=1) + data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + store.add("x", data) + + store.epoch_begin() + local_ok = True + for target_rank in range(size): + out = torch.full((1, ncols), -999.0, dtype=torch.float32, device="cuda") + store.get("x", out, start=target_rank * nrows) + expected = float(target_rank + 1) + ok = bool(torch.all(out.cpu() == expected)) + nonpoison = int((out.cpu() != -999.0).sum()) + print( + f"[rank {rank}] target_rank={target_rank} expected={expected} " + f"ok={ok} nonpoison_count={nonpoison}/{ncols} " + f"sample={out.cpu().flatten()[:8].tolist()}", + flush=True, + ) + if not ok: + local_ok = False + store.epoch_end() + + assert all_passed(comm, local_ok) + store.free() + + +@gpu_required +def test_get_host_to_host_cxi(comm, monkeypatch): + """Diagnostic: same transfer as above, but into a host numpy buffer + (bypasses HMEM entirely) -- isolates whether plain host-to-host RDMA + over cxi works correctly in this environment, independent of GPU support. + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + + store = dds.PyDDStore(comm, method=1) + data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + store.add("x", data) + + store.epoch_begin() + local_ok = True + for target_rank in range(size): + out = np.full((1, ncols), -999.0, dtype=np.float32) + store.get("x", out, start=target_rank * nrows) + expected = float(target_rank + 1) + ok = bool(np.all(out == expected)) + print( + f"[rank {rank}] target_rank={target_rank} expected={expected} " + f"got={out.tolist()} ok={ok}", + flush=True, + ) + if not ok: + local_ok = False + store.epoch_end() + + assert all_passed(comm, local_ok) + store.free() + + +# --------------------------------------------------------------------------- +# negative paths: clear, early errors -- no live cxi fabric required +# --------------------------------------------------------------------------- + + +@gpu_required +def test_gpu_buffer_rejected_on_hsn(comm, monkeypatch): + # Set up with cxi so that add() and init_fabric succeed (hsn is not + # available on all machines, e.g. Perlmutter which is CXI-only). + # Then switch DDSTORE_FABRIC to hsn before get() — the Python-level check + # in pyddstore/_core.pyx reads the env var at get() time and rejects GPU buffers + # with a clear error before touching the fabric. + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + store = dds.PyDDStore(comm, method=1) + data = np.ones((4, 4), dtype=np.float32) + store.add("x", data) + + monkeypatch.setenv("DDSTORE_FABRIC", "hsn") + out = torch.zeros((1, 4), dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match="cxi"): + store.get("x", out, start=0) + store.free() + + +@gpu_required +def test_gpu_buffer_rejected_on_method0(comm): + store = dds.PyDDStore(comm, method=0) + data = np.ones((4, 4), dtype=np.float32) + store.add("x", data) + + out = torch.zeros((1, 4), dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match="method"): + store.get("x", out, start=0) + store.free() + + +@gpu_required +def test_gpu_source_rejected_on_method0(comm): + store = dds.PyDDStore(comm, method=0) + data = torch.ones((4, 4), dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match="method"): + store.add("x", data) + + +@gpu_required +def test_gpu_source_rejected_on_hsn(comm, monkeypatch): + monkeypatch.setenv("DDSTORE_FABRIC", "hsn") + store = dds.PyDDStore(comm, method=1) + data = torch.ones((4, 4), dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match="cxi"): + store.add("x", data) + + +# --------------------------------------------------------------------------- +# Phase 2: GPU-resident producer (add()) -- host/GPU destination, over cxi +# --------------------------------------------------------------------------- +# +# These pass on Frontier (ROCm + cxi). + + +@gpu_required +def test_add_from_gpu_tensor_host_dest_cxi(comm, monkeypatch): + """GPU source -> host (poisoned) destination. Isolates that the SEND + side specifically works, independent of Phase 1's already-proven + receive side. + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + + store = dds.PyDDStore(comm, method=1) + data = torch.full( + (nrows, ncols), float(rank + 1), dtype=torch.float32, device="cuda" + ) + store.add("x", data) # GPU source -- Phase 2 + + store.epoch_begin() + local_ok = True + for target_rank in range(size): + out = np.full((1, ncols), -999.0, dtype=np.float32) + store.get("x", out, start=target_rank * nrows) + ok = bool(np.all(out == float(target_rank + 1))) + if not ok: + local_ok = False + store.epoch_end() + + assert all_passed(comm, local_ok) + store.free() + + +@gpu_required +def test_add_from_gpu_tensor_gpu_dest_cxi(comm, monkeypatch): + """Full Phase 2 scenario: both ends device memory, method=1.""" + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + + store = dds.PyDDStore(comm, method=1) + data = torch.full( + (nrows, ncols), float(rank + 1), dtype=torch.float32, device="cuda" + ) + store.add("x", data) + + store.epoch_begin() + local_ok = True + for target_rank in range(size): + out = torch.full((1, ncols), -999.0, dtype=torch.float32, device="cuda") + store.get("x", out, start=target_rank * nrows) + expected = float(target_rank + 1) + ok = bool(torch.all(out.cpu() == expected)) + if not ok: + local_ok = False + store.epoch_end() + + assert all_passed(comm, local_ok) + store.free() + + +@gpu_required +def test_add_from_gpu_tensor_gpu_dest_cxi_method2(comm, monkeypatch, tmp_path): + """Same as test_add_from_gpu_tensor_gpu_dest_cxi but method=2 + (file-based handshake, core+extra split) -- the transport + pyddstore.torch's DistDatasetReader actually uses. Includes + a self-read check (core rank both add()s and get()s its own data, + mirroring test_method2_core.py's self-check pattern) to verify the + independent send_hmem_iface/mr vs recv_hmem_iface/recv_mr fields don't + interfere with each other. + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + + hs_dir = comm.bcast( + str(tmp_path / "ddstore_hs_add_method2") if rank == 0 else None, root=0 + ) + + core_store = dds.PyDDStore(comm, method=2, handshake_dir=hs_dir) + data = torch.full( + (nrows, ncols), float(rank + 1), dtype=torch.float32, device="cuda" + ) + core_store.add("x", data) # GPU source -- Phase 2 + comm.Barrier() + + # Self-read: every core rank reads its own just-added shard back. + local_ok = True + out_self = torch.full((1, ncols), -999.0, dtype=torch.float32, device="cuda") + core_store.get("x", out_self, start=rank * nrows) + if not bool(torch.all(out_self.cpu() == float(rank + 1))): + local_ok = False + + # Extra member: a separate instance joins and reads every rank's shard. + if rank == 0: + extra_store = dds.PyDDStore(None, method=2, handshake_dir=hs_dir, n_core=size) + extra_store.join("x") + for target_rank in range(size): + out = torch.full((1, ncols), -999.0, dtype=torch.float32, device="cuda") + extra_store.get("x", out, start=target_rank * nrows) + expected = float(target_rank + 1) + if not bool(torch.all(out.cpu() == expected)): + local_ok = False + extra_store.free() + + comm.Barrier() + assert all_passed(comm, local_ok) + core_store.free() + + +# --------------------------------------------------------------------------- +# thread-safety: concurrent get() calls from multiple Python threads +# --------------------------------------------------------------------------- + + +def test_concurrent_get_thread_safety(comm, monkeypatch): + """DDStore::get() releases the GIL for its blocking transfer (see the + `with nogil:` block in pyddstore/_core.pyx), so multiple Python threads can + genuinely be inside DDStore::get() at the same time. Without + synchronization, concurrent calls on the same variable would race on + the CQ poll loop and the recv-MR region cache in common.cxx (confirmed + by direct experiment: disabling the protection crashed with "double + free or corruption"). + + That protection now lives inside DDStore itself -- a per-variable + `pthread_mutex_t` on `struct fabric_state` (include/common.h), taken + via the `fabric_state_lock_guard` RAII helper around get()'s critical + section in include/ddstore.hpp. This test calls PyDDStore.get() + directly from multiple threads with **no lock at the Python level at + all** -- it would catch a regression if + that C++-level protection were ever removed or narrowed. + """ + monkeypatch.setenv("DDSTORE_FABRIC", "cxi") + rank = comm.Get_rank() + size = comm.Get_size() + if size < 2: + pytest.skip("requires at least 2 ranks for a genuine remote read") + nrows, ncols = 8, 4 + + store = dds.PyDDStore(comm, method=1) + data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + store.add("x", data) + comm.Barrier() + + store.epoch_begin() + results = {} + errors = [] + + def unlocked_get(target_rank): + out = np.full((1, ncols), -999.0, dtype=np.float32) + store.get("x", out, start=target_rank * nrows) + results[target_rank] = out.copy() + + def worker(target_ranks): + try: + for target_rank in target_ranks: + unlocked_get(target_rank) + except Exception as exc: # noqa: BLE001 - surface any thread exception + errors.append(exc) + + n_threads = 4 + threads = [ + threading.Thread(target=worker, args=(list(range(t, size, n_threads)),)) + for t in range(min(n_threads, size)) + ] + for t in threads: + t.start() + for t in threads: + t.join() + store.epoch_end() + + assert not errors, f"worker thread(s) raised: {errors}" + local_ok = True + for target_rank in range(size): + expected = float(target_rank + 1) + got = results[target_rank] + ok = bool(np.all(got == expected)) + if not ok: + local_ok = False + print( + f"[rank {rank}] target_rank={target_rank} expected={expected} " + f"got={got.tolist()} ok={ok}", + flush=True, + ) + + assert all_passed(comm, local_ok) + store.free() diff --git a/test/test_multirank.py b/test/test_multirank.py index ad89109..e62efcc 100644 --- a/test/test_multirank.py +++ b/test/test_multirank.py @@ -2,7 +2,8 @@ Multi-rank tests — run with: mpirun -n 4 pytest test/test_multirank.py -v Each rank stores a distinct value; tests verify cross-rank remote reads. -Requires at least 2 ranks; some tests require exactly 4. +Requires at least 2 ranks; the ddstore_width test needs at least 4 +(skipped otherwise). """ import numpy as np diff --git a/test/test_torch.py b/test/test_torch.py new file mode 100644 index 0000000..97e844c --- /dev/null +++ b/test/test_torch.py @@ -0,0 +1,692 @@ +""" +pyddstore.torch (DistDataset, DistDatasetReader, ThreadDataLoader) tests — +run with 2+ ranks, e.g.: + mpirun -n 4 pytest test/test_torch.py -v + +Method 0 always; method 1 and the method-2 reader where a CXI device is +present inside a Slurm step (provider: DDSTORE_FABRIC, default cxi). +""" + +import glob +import os + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +from torch.utils.data import DataLoader, Dataset # noqa: E402 + +from pyddstore.torch import ( + DistDataset, + DistDatasetReader, + ThreadDataLoader, + WindowedDataset, + row_of, +) # noqa: E402 + +HAVE_CXI = bool(glob.glob("/dev/cxi*")) and "SLURM_STEP_ID" in os.environ +FABRIC = os.environ.get("DDSTORE_FABRIC", "cxi") +HAVE_GPU = torch.cuda.is_available() +METHODS = [ + 0, + pytest.param(1, marks=pytest.mark.skipif(not HAVE_CXI, reason="no CXI device")), +] +N = 37 # not a multiple of the rank count + + +class TupleSource(Dataset): + """Every field kind, values derived from the index.""" + + def __len__(self): + return N + + def __getitem__(self, i): + return ( + torch.arange(12, dtype=torch.float32).reshape(3, 4) + + 100 * i, # torch tensor + i, # Python int + np.full(2, i / 3, dtype=np.float64), # numpy array + np.int32(-i), # numpy scalar + i % 2 == 0, # Python bool + float(i) * 0.5, # Python float + torch.tensor([i, i + 1], dtype=torch.uint8), # small dtype + ) + + +class DictSource(Dataset): + def __len__(self): + return N + + def __getitem__(self, i): + return { + "x": torch.full((2, 3), float(i)), + "y": torch.tensor(i, dtype=torch.int64), + } + + +class SingleSource(Dataset): + def __len__(self): + return N + + def __getitem__(self, i): + return np.full((5,), i, dtype=np.int32) + + +REC = np.dtype( + [ + ("x_modules", np.float32, (4, 3)), + ("mask", np.bool_, (4,)), + ("params", np.int64, (2,)), + ] +) +# padded (align=True) and nested layout +REC_NESTED = np.dtype( + [ + ("a", np.uint8), + ("b", np.float64), + ("sub", [("c", np.int32, (2,)), ("d", np.bool_)]), + ], + align=True, +) + + +def _record(i, dtype=REC): + a = np.zeros((), dtype=dtype) + if dtype is REC: + a["x_modules"], a["mask"], a["params"] = i, i % 2 == 0, (i, -i) + else: + a["a"], a["b"], a["sub"]["c"], a["sub"]["d"] = ( + i % 256, + i / 7, + (i, 2 * i), + i % 3 == 0, + ) + return a[()] # np.void + + +class RecordSource(Dataset): + """Items are numpy structured records, in the forms projects use.""" + + def __init__(self, form): + self.form = form + + def __len__(self): + return N + + def __getitem__(self, i): + if self.form == "void": + return _record(i) + if self.form == "array": # 1-element structured ndarray + return np.array([_record(i)], dtype=REC) + if self.form == "recarray": # np.recarray of shape (2,) + return np.array([_record(i), _record(i + 1)], dtype=REC).view(np.recarray) + if self.form == "nested": + return _record(i, REC_NESTED) + return {"rec": _record(i), "t": torch.full((3,), float(i))} # mixed dict + + +def same(a, b): + """Equal structure, types and values.""" + if type(a) is not type(b): + return False + if isinstance(a, (tuple, list)): + return len(a) == len(b) and all(same(x, y) for x, y in zip(a, b)) + if isinstance(a, dict): + return a.keys() == b.keys() and all(same(a[k], b[k]) for k in a) + if isinstance(a, torch.Tensor): + return ( + a.dtype == b.dtype + and a.shape == b.shape + and bool(torch.equal(a.cpu(), b.cpu())) + ) + if isinstance(a, np.ndarray): + return a.dtype == b.dtype and a.shape == b.shape and np.array_equal(a, b) + if isinstance(a, np.generic): + return a.dtype == b.dtype and a == b + return a == b + + +def all_ok(comm, ok): + return comm.allreduce(int(bool(ok)), op=__import__("mpi4py").MPI.LAND) + + +def make(comm, monkeypatch, source, method, **kw): + if method != 0: + monkeypatch.setenv("DDSTORE_FABRIC", FABRIC) + return DistDataset(source, f"t{method}", comm, method=method, **kw) + + +def finish(comm, ds): + comm.Barrier() # every rank done reading before any rank tears down + ds.ddstore.free() + + +@pytest.mark.parametrize("chunk_size", [None, 1, 4]) +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize("source_cls", [TupleSource, DictSource, SingleSource]) +def test_items_match_source(comm, monkeypatch, method, source_cls, chunk_size): + src = source_cls() + ds = make(comm, monkeypatch, src, method, chunk_size=chunk_size) + rng = np.random.default_rng(comm.Get_rank()) + idx = rng.integers(0, N, size=20) + ok = len(ds) == N + ok &= all(same(ds[int(i)], src[int(i)]) for i in idx[:5]) # per-sample get() + batch = ds.__getitems__(idx) # one get_batch() per field (collective for method 0) + ok &= all(same(b, src[int(i)]) for b, i in zip(batch, idx)) + finish(comm, ds) + assert all_ok(comm, ok) + + +def test_shapes_and_dtypes(comm, monkeypatch): + ds = make(comm, monkeypatch, DictSource(), 0) + ok = ds.shapes == {"x": (2, 3), "y": ()} and ds.dtypes == { + "x": "float32", + "y": "int64", + } + finish(comm, ds) + assert all_ok(comm, ok) + + +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize("loader", ["DataLoader", "ThreadDataLoader"]) +def test_loaders_match_plain_source(comm, monkeypatch, method, loader): + """A whole epoch through the loader equals the same loader over the plain + source (same order, same collation). Collective for method 0: every rank + iterates the same number of batches.""" + src = TupleSource() + ds = make(comm, monkeypatch, src, method) + cls = DataLoader if loader == "DataLoader" else ThreadDataLoader + kw = {} if cls is DataLoader else {"num_workers": 1 if method == 0 else 2} + got = list(cls(ds, batch_size=8, **kw)) + ref = list(DataLoader(src, batch_size=8)) + ok = len(got) == len(ref) and all(same(g, r) for g, r in zip(got, ref)) + finish(comm, ds) + assert all_ok(comm, ok) + + +def test_thread_loader_iterators(comm): + """iter(loader) is a separate iterator, as with DataLoader: iterating it + again continues the epoch (list(it), islice), two iterators over one + loader are independent, and an epoch stopped early doesn't leak into the + next one.""" + import itertools + + src = TupleSource() + ref = list(DataLoader(src, batch_size=4)) + loader = ThreadDataLoader(src, batch_size=4, num_workers=2) + + def eq(got, want): + return len(got) == len(want) and all(same(g, r) for g, r in zip(got, want)) + + it = iter(loader) + ok = iter(it) is it and len(it) == len(ref) + first = next(it) + two = list(itertools.islice(it, 2)) + rest = list(it) # continues, does not restart + ok &= eq([first] + two + rest, ref) + a, b = iter(loader), iter(loader) + got_a, got_b = [], [] + for x, y in zip(a, b): # interleaved + got_a.append(x) + got_b.append(y) + ok &= eq(got_a, ref) and eq(got_b, ref) + for i, _ in enumerate(loader): # stop early + if i == 1: + break + ok &= eq(list(loader), ref) + assert all_ok(comm, ok) + + +def test_thread_loader_keeps_prefetch_full(comm): + """While the caller holds a batch, num_workers * prefetch_factor more are + being fetched (as with DataLoader), not one fewer.""" + import threading + import time + + class Counting(Dataset): + def __init__(self): + self.lock, self.batches = threading.Lock(), 0 + + def __len__(self): + return 64 + + def __getitems__(self, idx): + with self.lock: + self.batches += 1 + return [torch.tensor(i) for i in idx] + + src = Counting() + loader = ThreadDataLoader(src, batch_size=4, num_workers=2, prefetch_factor=1) + it = iter(loader) + next(it) # held by the "training step" + want = 1 + 2 * 1 + deadline = time.time() + 10 + while src.batches < want and time.time() < deadline: + time.sleep(0.01) + time.sleep(0.2) # nothing beyond the bound should start + ok = src.batches == want + it.close() + assert all_ok(comm, ok), f"{src.batches} batches fetched, want {want}" + + +def stacked(src, rows, key): + """Field `key` of src[r] for r in rows, stacked as read_rows returns it.""" + vals = [src[int(r)] if key == 0 and not isinstance(src[0], (tuple, list, dict)) + else src[int(r)][key] for r in rows] + if isinstance(vals[0], torch.Tensor): + return torch.stack(vals) + return np.stack([np.asarray(v) for v in vals]) + + +def same_values(a, b): + a = a.cpu().numpy() if isinstance(a, torch.Tensor) else np.asarray(a) + b = b.cpu().numpy() if isinstance(b, torch.Tensor) else np.asarray(b) + return a.dtype == b.dtype and a.shape == b.shape and np.array_equal(a, b) + + +def shares(value, buf): + if isinstance(buf, torch.Tensor): + return value.data_ptr() == buf.data_ptr() + return np.shares_memory(np.asarray(value), buf) + + +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize("source_cls", [TupleSource, DictSource, SingleSource]) +def test_read_rows_and_out(comm, monkeypatch, method, source_cls): + """read_rows (all fields / a subset), alloc() buffers reused across + reads, __getitems__(out=). Same number of calls on every rank (method 0 + is collective).""" + src = source_cls() + ds = make(comm, monkeypatch, src, method) + keys = [ds._key(j) for j in range(len(ds._var))] + rng = np.random.default_rng(comm.Get_rank()) + rows = rng.integers(0, N, size=10) # any order, repeats + got = ds.read_rows(rows) + ok = list(got) == keys + ok &= all(same_values(got[k], stacked(src, rows, k)) for k in keys) + sub = ds.read_rows(rows[:3], fields=keys[-1:]) + ok &= list(sub) == keys[-1:] and same_values(sub[keys[-1]], stacked(src, rows[:3], keys[-1])) + with pytest.raises(KeyError): + ds.read_rows(rows, fields=["nope"]) + + bufs = ds.alloc(16) + miss0 = None + if os.environ.get("DDSTORE_PROFILE", "0") not in ("", "0") and method != 0: + miss0 = [ds.ddstore.get_profile(v)["mr_miss"] for v in ds._var] + for it in range(3): # the same buffers, different rows each time + r = rng.integers(0, N, size=12) + got = ds.read_rows(r, out=bufs) + ok &= all(same_values(got[k], stacked(src, r, k)) for k in keys) + ok &= all(shares(got[k], bufs[k]) for k in keys) + batch = ds.__getitems__(r[:5], out=bufs) + ok &= all(same(b, src[int(i)]) for b, i in zip(batch, r[:5])) + if miss0 is not None: + ok &= [ds.ddstore.get_profile(v)["mr_miss"] for v in ds._var] == miss0 + with pytest.raises(ValueError): + ds.read_rows(np.arange(17) % N, out=bufs) # more rows than the buffers + ds.release(bufs) + finish(comm, ds) + assert all_ok(comm, ok) + + +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize("loader", ["DataLoader", "ThreadDataLoader"]) +def test_windowed_dataset(comm, monkeypatch, method, loader): + """Windows with stride and dilation, explicit starts, a field subset, and + whole batches through a loader (one read_rows per batch).""" + src = TupleSource() + ds = make(comm, monkeypatch, src, method) + nf = len(ds._var) + + def window(rows): + return tuple(stacked(src, rows, j) for j in range(nf)) + + wd = WindowedDataset(ds, window=3, stride=2, dilation=2) # rows s, s+2, s+4 + ok = len(wd) == (N - 5) // 2 + 1 + ok &= same(wd[4], window([8, 10, 12])) and same(wd[-1], window([32, 34, 36])) + cls = DataLoader if loader == "DataLoader" else ThreadDataLoader + kw = {} if cls is DataLoader else {"num_workers": 1 if method == 0 else 2} + got = list(cls(wd, batch_size=4, **kw)) + ref = [ + torch.utils.data.default_collate([window([s, s + 2, s + 4]) for s in range(b, min(b + 8, len(wd) * 2), 2)]) + for b in range(0, len(wd) * 2, 8) + ] + ok &= len(got) == len(ref) and all(same(g, r) for g, r in zip(got, ref)) + + starts = [0, 10, 30] # e.g. one window per trajectory + ws = WindowedDataset(ds, window=2, starts=starts, fields=[0, 2]) + ok &= len(ws) == 3 + w = ws[1] + ok &= list(w) == [0, 2] and same_values(w[2], stacked(src, [10, 11], 2)) + with pytest.raises(IndexError): + WindowedDataset(ds, window=2, starts=[N - 1]) + finish(comm, ds) + assert all_ok(comm, ok) + + +class Offset(Dataset): + def __init__(self, n, base): + self.n, self.base = n, base + + def __len__(self): + return self.n + + def __getitem__(self, i): + return np.full((3,), self.base + i, dtype=np.int64) + + +def test_row_of_concat(comm, monkeypatch): + """Several sources in one store: row_of maps (source, index) to the row.""" + parts = [Offset(11, 0), Offset(7, 1000), Offset(19, 2000)] + concat = torch.utils.data.ConcatDataset(parts) + ds = make(comm, monkeypatch, concat, 0) + ok = row_of(concat, 0, 3) == 3 and row_of(concat, 2, 0) == 18 + rows = [row_of(concat, s, i) for s, i in [(1, 6), (2, 18), (0, 0)]] + got = ds.read_rows(rows)[0] + ok &= same_values(got, np.stack([parts[1][6], parts[2][18], parts[0][0]])) + for bad in [(3, 0), (1, 7), (0, -1)]: + with pytest.raises(IndexError): + row_of(concat, *bad) + finish(comm, ds) + assert all_ok(comm, ok) + + +@pytest.mark.parametrize("method", METHODS) +def test_thread_loader_reuse_buffers(comm, monkeypatch, method): + """reuse_buffers: two epochs equal the plain source, no registration per + read (method 1, DDSTORE_PROFILE=1), setups that can't copy are refused, + and close() unregisters the pool.""" + src = TupleSource() + ds = make(comm, monkeypatch, src, method) + nw = 1 if method == 0 else 2 + ref = list(DataLoader(src, batch_size=8)) # last batch is short (37 = 4*8 + 5) + loader = ThreadDataLoader(ds, batch_size=8, num_workers=nw, reuse_buffers=True) + ok = len(loader._pool_sets) == nw + prof = os.environ.get("DDSTORE_PROFILE", "0") not in ("", "0") and method != 0 + miss0 = [ds.ddstore.get_profile(v)["mr_miss"] for v in ds._var] if prof else None + for _ in range(2): + got = list(loader) + ok &= len(got) == len(ref) and all(same(g, r) for g, r in zip(got, ref)) + if prof: + ok &= [ds.ddstore.get_profile(v)["mr_miss"] for v in ds._var] == miss0 + sets = loader._pool_sets + loader.close() + if method != 0: # unregistered: releasing again raises + for bufs in sets: + with pytest.raises(ValueError): + ds.release(bufs) + + with pytest.raises(ValueError): # no auto-collation: nothing copies + ThreadDataLoader(ds, batch_size=None, reuse_buffers=True) + with pytest.raises(ValueError): # custom collate not declared as copying + ThreadDataLoader(ds, batch_size=8, collate_fn=lambda b: b, reuse_buffers=True) + with pytest.raises(TypeError): # dataset without alloc() + ThreadDataLoader(src, batch_size=8, reuse_buffers=True) + copying = ThreadDataLoader( + ds, batch_size=8, num_workers=nw, reuse_buffers=True, collate_copies=True, + collate_fn=lambda b: [tuple(x.clone() if isinstance(x, torch.Tensor) else + np.array(x, copy=True) for x in s) for s in b], + ) + got = [s for b in copying for s in b] + ok &= len(got) == N and all( + same_values(g[0], src[i][0]) and same_values(g[2], src[i][2]) + for i, g in enumerate(got) + ) + copying.close() + finish(comm, ds) + assert all_ok(comm, ok) + + +LABELS = ["cat", "dog", "owl"] +GROUP_INFO = {g: {"name": f"group-{g}", "scale": 1.5 * g} for g in range(4)} + + +class Labeled(Dataset): + """Samples with a string label and an object that can't be stored.""" + + def __len__(self): + return N + + def __getitem__(self, i): + return { + "x": torch.full((3,), float(i)), + "label": LABELS[i % 3], + "group": i // 10, + "meta": GROUP_INFO[i // 10], # per-group object + } + + +def labeled_encode(s): + return {"x": s["x"], "label": LABELS.index(s["label"]), "group": s["group"]} + + +def labeled_decode(d, i): + return { + "x": d["x"], + "label": LABELS[d["label"]], + "group": d["group"], + "meta": GROUP_INFO[d["group"]], + "index": i, + } + + +def labeled_ok(sample, i): + want = Labeled()[i] + return ( + torch.equal(sample["x"], want["x"]) + and sample["label"] == want["label"] + and sample["group"] == want["group"] + and sample["meta"] is GROUP_INFO[want["group"]] + and sample["index"] == i + ) + + +@pytest.mark.parametrize("method", METHODS) +def test_encode_decode(comm, monkeypatch, method): + """encode stores ids instead of strings/objects, decode rebuilds the + sample with its index; ds[i], __getitems__ and a loader all decode; + read_rows stays row-level.""" + with pytest.raises(TypeError): # strings can't be stored as they are + make(comm, monkeypatch, Labeled(), method) + ds = make(comm, monkeypatch, Labeled(), method, encode=labeled_encode, + decode=labeled_decode) + ok = labeled_ok(ds[7], 7) and labeled_ok(ds[N - 1], N - 1) + idx = list(range(N))[::-2] + ok &= all(labeled_ok(sm, i) for sm, i in zip(ds.__getitems__(idx), idx)) + kw = {"num_workers": 1 if method == 0 else 2} + seen = 0 + for batch in ThreadDataLoader(ds, batch_size=8, collate_fn=lambda b: b, **kw): + ok &= all(labeled_ok(sm, sm["index"]) for sm in batch) + seen += len(batch) + ok &= seen == N + raw = ds.read_rows([4, 5]) # stored form: label ids, no meta/index + ok &= sorted(raw) == ["group", "label", "x"] and raw["label"].tolist() == [1, 2] + finish(comm, ds) + assert all_ok(comm, ok) + + +def test_fields_selection(comm, monkeypatch): + """fields= keeps the given keys / positions, in that order.""" + ds = make(comm, monkeypatch, DictSource(), 0, fields=["y"]) + ok = same(ds[5], {"y": DictSource()[5]["y"]}) + finish(comm, ds) + src = TupleSource() + dt = make(comm, monkeypatch, src, 0, fields=[2, 0]) + ok &= same(dt[9], (src[9][2], src[9][0])) + finish(comm, dt) + with pytest.raises(ValueError): + make(comm, monkeypatch, src, 0, fields=[0], encode=lambda s: s) + with pytest.raises(ValueError): # missing key, raised on every rank + make(comm, monkeypatch, DictSource(), 0, fields=["nope"]) + assert all_ok(comm, ok) + + +def test_per_sample_fallback(comm, monkeypatch): + monkeypatch.setenv("DDSTORE_BATCH_GET", "0") + src = TupleSource() + ds = make(comm, monkeypatch, src, 0) + idx = list(range(N))[::-3] + ok = all(same(b, src[i]) for b, i in zip(ds.__getitems__(idx), idx)) + finish(comm, ds) + assert all_ok(comm, ok) + + +def test_ddstore_width_groups(comm, monkeypatch): + if comm.Get_size() < 4: + pytest.skip("requires at least 4 ranks") + src = SingleSource() + ds = make(comm, monkeypatch, src, 0, ddstore_width=2) + idx = list(range(N)) + ok = all(same(b, src[i]) for b, i in zip(ds.__getitems__(idx), idx)) + finish(comm, ds) + assert all_ok(comm, ok) + + +class _Bad(Dataset): + def __init__(self, kind): + self.kind = kind + + def __len__(self): + return N + + def __getitem__(self, i): + if self.kind == "shape": + return np.zeros(3 if i % 5 else 4, dtype=np.float32) + if self.kind == "dtype": + return torch.zeros(2, dtype=torch.float16) + if self.kind == "nested": + return (np.zeros(2), (1, 2)) + return object() + + +@pytest.mark.parametrize( + "kind,exc", + [ + ("shape", ValueError), + ("dtype", TypeError), + ("nested", TypeError), + ("object", TypeError), + ], +) +@pytest.mark.parametrize("chunk_size", [None, 3]) +def test_unsupported_samples_raise(comm, kind, exc, chunk_size): + with pytest.raises(exc): + DistDataset(_Bad(kind), "bad", comm, method=0, chunk_size=chunk_size) + comm.Barrier() + + +def test_chunk_size_needs_host_storage(comm): + with pytest.raises(ValueError): + DistDataset(TupleSource(), "c", comm, method=0, chunk_size=4, add_device="cpu") + with pytest.raises(ValueError): + DistDataset(TupleSource(), "c", comm, method=0, chunk_size=0) + comm.Barrier() + + +@pytest.mark.skipif( + not (HAVE_CXI and HAVE_GPU and FABRIC == "cxi"), + reason="requires the cxi provider and a GPU", +) +def test_gpu_device_and_add_device(comm, monkeypatch): + src = TupleSource() + ds = make(comm, monkeypatch, src, 1, device="cuda", add_device="cuda") + idx = list(range(N))[::-1] + batch = ds.__getitems__(idx) + ok = batch[0][0].is_cuda and isinstance( + batch[0][2], np.ndarray + ) # tensors on GPU, numpy stays host + ok &= all(same(b, src[i]) for b, i in zip(batch, idx)) + finish(comm, ds) + assert all_ok(comm, ok) + + +@pytest.mark.skipif( + not (HAVE_CXI and HAVE_GPU and FABRIC == "cxi"), + reason="requires the cxi provider and a GPU", +) +def test_gpu_reuse_buffers(comm, monkeypatch): + """reuse_buffers with GPU read buffers: the collate's GPU copy must finish + before a buffer is refilled (several epochs, 2 threads, small pool).""" + src = TupleSource() + ds = make(comm, monkeypatch, src, 1, device="cuda") + ref = list(DataLoader(src, batch_size=4)) + loader = ThreadDataLoader(ds, batch_size=4, num_workers=2, reuse_buffers=True) + ok = True + for _ in range(3): + got = list(loader) + ok &= got[0][0].is_cuda and len(got) == len(ref) + ok &= all(same(g, r) for g, r in zip(got, ref)) + loader.close() + finish(comm, ds) + assert all_ok(comm, ok) + + +@pytest.mark.skipif(not HAVE_CXI, reason="no CXI device") +def test_method2_reader(comm, monkeypatch, tmp_path): + monkeypatch.setenv("DDSTORE_FABRIC", FABRIC) + hs = comm.bcast(str(tmp_path / "hs") if comm.Get_rank() == 0 else None, root=0) + src = RecordSource("dict") # records + tensors: layout round-trips via meta.json + core = DistDataset(src, "rd", comm, method=2, handshake_dir=hs) + comm.Barrier() + ok = True + if comm.Get_rank() == 0: + reader = DistDatasetReader("rd", handshake_dir=hs, n_core=comm.Get_size()) + idx = list(range(N)) + ok = len(reader) == N and reader.shapes == core.shapes + ok &= all(same(b, src[i]) for b, i in zip(reader.__getitems__(idx), idx)) + reader.ddstore.free() + finish(comm, core) + assert all_ok(comm, ok) + + +@pytest.mark.skipif(not HAVE_CXI, reason="no CXI device") +def test_method2_reader_decode(comm, monkeypatch, tmp_path): + """The core group encodes, a DistDatasetReader decodes.""" + monkeypatch.setenv("DDSTORE_FABRIC", FABRIC) + hs = comm.bcast(str(tmp_path / "hs") if comm.Get_rank() == 0 else None, root=0) + core = DistDataset(Labeled(), "lab", comm, method=2, handshake_dir=hs, + encode=labeled_encode) + comm.Barrier() + ok = True + if comm.Get_rank() == 0: + reader = DistDatasetReader("lab", handshake_dir=hs, n_core=comm.Get_size(), + decode=labeled_decode) + idx = list(range(N)) + ok = all(labeled_ok(sm, i) for sm, i in zip(reader.__getitems__(idx), idx)) + ok &= labeled_ok(reader[3], 3) + reader.ddstore.free() + finish(comm, core) + assert all_ok(comm, ok) + + +@pytest.mark.parametrize("chunk_size", [None, 4]) +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize("form", ["void", "array", "recarray", "nested", "dict"]) +def test_record_items(comm, monkeypatch, method, form, chunk_size): + """numpy structured records as items (or a field) come back as the same + kind of object with the same layout and values.""" + src = RecordSource(form) + ds = make(comm, monkeypatch, src, method, chunk_size=chunk_size) + idx = list(range(N))[::-2] + ok = all(same(ds[i], src[i]) for i in idx[:4]) + ok &= all(same(b, src[i]) for b, i in zip(ds.__getitems__(idx), idx)) + # records don't collate with default_collate: pass collate_fn through + batches = list( + ThreadDataLoader(ds, batch_size=5, num_workers=1, collate_fn=lambda b: b) + ) + ok &= all( + same(b, src[i]) for i, b in zip(range(N), [x for bt in batches for x in bt]) + ) + finish(comm, ds) + assert all_ok(comm, ok) + + +def test_record_dtypes_property(comm, monkeypatch): + ds = make(comm, monkeypatch, RecordSource("dict"), 0) + ok = ( + ds.dtypes["rec"] == REC + and ds.dtypes["t"] == "float32" + and ds.shapes["rec"] == () + ) + finish(comm, ds) + assert all_ok(comm, ok)