diff --git a/README.md b/README.md index 536a832..64fe9d4 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ Each MPI rank holds a shard of the full dataset in memory. DDStore exposes a glo | Dependency | Notes | |---|---| | MPI (OpenMPI / MPICH) | `mpicc` and `mpicxx` must be on `PATH` | -| libfabric | Required for RDMA backend (`method=1`) | +| libfabric | Required for the RDMA backends (`method=1` and `method=2`) | | Python ≥ 3.6 | | | NumPy, mpi4py, Cython | Python build dependencies | @@ -71,13 +71,26 @@ mpirun -n 4 python my_script.py ## API Reference -### `PyDDStore(comm, method=0, ddstore_width=None)` +### `PyDDStore(comm_or_none=None, method=0, handshake_dir="", n_core=0, nic_map=None)` | Parameter | Type | Description | |---|---|---| -| `comm` | `mpi4py.MPI.Comm` | MPI communicator covering all ranks | -| `method` | `int` | `0` = MPI RMA (default), `1` = libfabric RDMA | -| `ddstore_width` | `int` or `None` | Ranks per DDStore group. `None` uses all ranks in `comm` as a single group | +| `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)`). --- @@ -129,6 +142,22 @@ Read `arr.shape[0]` consecutive rows starting at global index `start` into `arr` --- +### `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`. @@ -147,18 +176,74 @@ Uses `MPI_Win_create` and `MPI_Get` for one-sided remote reads. Works on any MPI ### libfabric RDMA (`method=1`) -Uses `fi_read` for true RDMA transfers over high-speed interconnects (Infiniband/verbs, Cray GNI, Intel PSM2). Lower latency than MPI RMA on supported hardware. `epoch_begin`/`epoch_end` are no-ops with this backend. +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: -Set `FABRIC_IFACE` to select a specific network interface when the automatic selection picks the wrong one: ```bash -export FABRIC_IFACE=hsn0 # e.g. Cray Slingshot +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 -`ddstore_width` controls how many MPI ranks form a single DDStore group. The global communicator is split so that each group of `ddstore_width` consecutive ranks shares one independent store, with each group holding a full replica of the dataset partitioned across its members. +`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, `ddstore_width=4`:** +**Example — 16 ranks split into groups of 4:** ``` ranks 0– 3 → DDStore group 0 ranks 4– 7 → DDStore group 1 @@ -166,13 +251,15 @@ 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 → `ddstore_width=4`), limiting cross-node RDMA traffic to the dataset replication step at startup rather than every sample fetch. +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. ```python -store = dds.PyDDStore(comm, ddstore_width=4) # e.g. 4 GPUs per store +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 ``` -If `ddstore_width` is omitted, all ranks in `comm` form a single store. +`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 @@ -227,6 +314,18 @@ Optional arguments for `examples/scripts/demo.py` and `examples/scripts/test.py` | `--dim` | `64` | Elements per row | | `--nbatch` | `32` | Number of random reads | +### 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 +``` + ## Citation If you use DDStore in your research, please cite: diff --git a/examples/scripts/demo.py b/examples/scripts/demo.py index 642129a..bdc53fc 100644 --- a/examples/scripts/demo.py +++ b/examples/scripts/demo.py @@ -9,7 +9,6 @@ import pyddstore as dds import sys - if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( @@ -18,8 +17,12 @@ help="num. of data (default: %(default)s)", default=1024 * 1024, ) - parser.add_argument("--dim", type=int, help="dim (default: %(default)s)", default=64) - parser.add_argument("--nbatch", type=int, help="nbatch (default: %(default)s)", default=32) + parser.add_argument( + "--dim", type=int, help="dim (default: %(default)s)", default=64 + ) + parser.add_argument( + "--nbatch", type=int, help="nbatch (default: %(default)s)", default=32 + ) args = parser.parse_args() comm = MPI.COMM_WORLD diff --git a/examples/scripts/test.py b/examples/scripts/test.py index 7e62dbe..1019623 100644 --- a/examples/scripts/test.py +++ b/examples/scripts/test.py @@ -79,11 +79,19 @@ def parse_slurm_nodelist(nodelist): help="num. of data (default: %(default)s)", default=1024 * 1024, ) - parser.add_argument("--dim", type=int, help="dim (default: %(default)s)", default=64) - parser.add_argument("--nbatch", type=int, help="nbatch (default: %(default)s)", default=32) + parser.add_argument( + "--dim", type=int, help="dim (default: %(default)s)", default=64 + ) + parser.add_argument( + "--nbatch", type=int, help="nbatch (default: %(default)s)", default=32 + ) group = parser.add_mutually_exclusive_group() - group.add_argument("--gloo", help="gloo", action="store_const", dest="backend", const="gloo") - group.add_argument("--nccl", help="nccl", action="store_const", dest="backend", const="nccl") + group.add_argument( + "--gloo", help="gloo", action="store_const", dest="backend", const="gloo" + ) + group.add_argument( + "--nccl", help="nccl", action="store_const", dest="backend", const="nccl" + ) parser.set_defaults(backend="gloo") args = parser.parse_args() diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py new file mode 100644 index 0000000..12b4e69 --- /dev/null +++ b/examples/vae/ddp_utils.py @@ -0,0 +1,169 @@ +import os +import re +import socket + +import psutil +import torch +import torch.distributed as dist + +""" +Functions for DDP on HPC +""" + + +def init_comm_size_and_rank(): + world_size = None + world_rank = 0 + + if os.getenv("OMPI_COMM_WORLD_SIZE") and os.getenv("OMPI_COMM_WORLD_RANK"): + ## Summit + world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"]) + world_rank = int(os.environ["OMPI_COMM_WORLD_RANK"]) + elif os.getenv("SLURM_NPROCS") and os.getenv("SLURM_PROCID"): + ## CADES + world_size = int(os.environ["SLURM_NPROCS"]) + world_rank = int(os.environ["SLURM_PROCID"]) + else: + from mpi4py import MPI + + world_size = MPI.COMM_WORLD.Get_size() + world_rank = MPI.COMM_WORLD.Get_rank() + + ## Fall back to default + if world_size is None: + world_size = 1 + + return int(world_size), int(world_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). + """ + if os.getenv("OMPI_COMM_WORLD_LOCAL_RANK") is not None: + return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) + elif os.getenv("SLURM_LOCALID") is not None: + return int(os.environ["SLURM_LOCALID"]) + return 0 + + +def find_ifname(myaddr): + """ + Find socket ifname for a given ip adress. This is for "GLOO" ddp setup. + Usage example: + find_ifname("127.0.0.1") will return a network interface name, such as "lo". "lo0", etc. + """ + ipaddr = socket.gethostbyname(myaddr) + ifname = None + for nic, addrs in psutil.net_if_addrs().items(): + for addr in addrs: + if addr.address == ipaddr: + ifname = nic + break + if ifname is not None: + break + + return ifname + + +def parse_slurm_nodelist(nodelist): + """ + Parse SLURM_NODELIST env string to get list of nodes. + Usage example: + parse_slurm_nodelist(os.environ["SLURM_NODELIST"]) + Input examples: + "or-condo-g04" + "or-condo-g[05,07-08,13]" + "or-condo-g[05,07-08,13],or-condo-h[01,12]" + """ + nlist = list() + for block, _ in re.findall(r"([\w-]+(\[[\d\-,]+\])*)", nodelist): + m = re.match(r"^(?P[\w\-]+)\[(?P.*)\]", block) + if m is None: + ## single node + nlist.append(block) + else: + ## multiple nodes + g = m.groups() + prefix = g[0] + for sub in g[1].split(","): + if "-" in sub: + start, end = re.match(r"(\d+)-(\d+)", sub).groups() + fmt = "%%0%dd" % (len(start)) + for i in range(int(start), int(end) + 1): + node = prefix + fmt % i + nlist.append(node) + else: + node = prefix + sub + nlist.append(node) + + return nlist + + +def setup_ddp(): + """ "Initialize DDP""" + + if os.getenv("DDSTORE_BACKEND") is not None: + backend = os.environ["DDSTORE_BACKEND"] + elif dist.is_nccl_available() and torch.cuda.is_available(): + backend = "nccl" + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + backend = "xccl" + elif torch.distributed.is_gloo_available(): + backend = "gloo" + else: + raise RuntimeError("No parallel backends available") + + world_size, world_rank = init_comm_size_and_rank() + print(f"DDP: Hi from rank {world_rank} of {world_size}.") + + ## Default setting + master_addr = "127.0.0.1" + master_port = os.getenv("MASTER_PORT", "2345") + + if os.getenv("LSB_HOSTS") is not None: + master_addr = os.environ["LSB_HOSTS"].split()[1] + elif os.getenv("LSB_MCPU_HOSTS") is not None: + master_addr = os.environ["LSB_MCPU_HOSTS"].split()[2] + elif os.getenv("SLURM_STEP_NODELIST") is not None: + master_addr = parse_slurm_nodelist(os.environ["SLURM_STEP_NODELIST"])[0] + elif os.getenv("SLURM_NODELIST") is not None: + master_addr = parse_slurm_nodelist(os.environ["SLURM_NODELIST"])[0] + elif os.getenv("PBS_O_HOST") is not None: + if os.environ["PBS_O_HOST"][-19:] == "aurora.alcf.anl.gov": + from mpi4py import MPI + + RANK = MPI.COMM_WORLD.Get_rank() + MASTER_ADDR = socket.gethostname() if RANK == 0 else None + MASTER_ADDR = MPI.COMM_WORLD.bcast(MASTER_ADDR, root=0) + master_addr = f"{MASTER_ADDR}.hsn.cm.aurora.alcf.anl.gov" + else: + ## The following is CADES specific + master_addr = parse_slurm_nodelist(os.environ["PBS_O_HOST"])[0] + + try: + if backend in ["nccl", "gloo", "xccl"]: + os.environ["MASTER_ADDR"] = master_addr + os.environ["MASTER_PORT"] = str(master_port) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["RANK"] = str(world_rank) + + if (backend == "gloo") and ("GLOO_SOCKET_IFNAME" not in os.environ): + ifname = find_ifname(master_addr) + if ifname is not None: + os.environ["GLOO_SOCKET_IFNAME"] = ifname + + print( + "Distributed data parallel: %s master at %s:%s" + % (backend, master_addr, master_port), + ) + + if not dist.is_initialized(): + dist.init_process_group(backend=backend, init_method="env://") + + except KeyError: + print("DDP has to be initialized within a job - Running in sequential mode") + + return world_size, world_rank diff --git a/examples/vae/distdataset.py b/examples/vae/distdataset.py index 4295169..6f07f58 100644 --- a/examples/vae/distdataset.py +++ b/examples/vae/distdataset.py @@ -33,12 +33,23 @@ def __init__(self, data, label, comm=MPI.COMM_WORLD, ddstore_width=None): self.ddstore_comm_size = self.ddstore_comm.Get_size() ddstore_method = int(os.getenv("DDSTORE_METHOD", "0")) - gpu_id = int(os.getenv("SLURM_LOCALID")) - os.environ["FABRIC_IFACE"] = f"hsn{gpu_id//2}" print("DDStore method:", ddstore_method) - print("FABRIC_IFACE:", os.environ["FABRIC_IFACE"]) - - self.ddstore = dds.PyDDStore(self.ddstore_comm, 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) @@ -63,7 +74,10 @@ def __init__(self, data, label, comm=MPI.COMM_WORLD, ddstore_width=None): self.data.append(val) self.labels.append(label) - self.data = np.concatenate(self.data) + # 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) @@ -79,7 +93,9 @@ def __len__(self): return self.len() def get(self, idx): - val = np.zeros(28 * 28, dtype=np.float32) + ## 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 @@ -92,3 +108,56 @@ def get(self, idx): 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/vae-ddp.py b/examples/vae/vae-ddp.py index 458fce3..f7bed6e 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -1,207 +1,92 @@ -from __future__ import print_function -from mpi4py import MPI - +## torch (and the RCCL/HIP shared libraries it pulls in) must finish loading +## before mpi4py triggers MPI_Init, or their static destructors run in the +## wrong order at interpreter exit and corrupt the heap. +## Do not reorder these imports. import argparse +import os import torch import torch.utils.data -from torch import nn, optim -from torch.nn import functional as F +from torch import optim from torchvision import datasets, transforms from torchvision.utils import save_image - -import distdataset -from distdataset import DistDataset - import torch.distributed as dist -import os -import socket -import psutil -import re - -""" -Functions for DDP on HPC -""" -def init_comm_size_and_rank(): - world_size = None - world_rank = 0 - - if os.getenv("OMPI_COMM_WORLD_SIZE") and os.getenv("OMPI_COMM_WORLD_RANK"): - ## Summit - world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"]) - world_rank = int(os.environ["OMPI_COMM_WORLD_RANK"]) - elif os.getenv("SLURM_NPROCS") and os.getenv("SLURM_PROCID"): - ## CADES - world_size = int(os.environ["SLURM_NPROCS"]) - world_rank = int(os.environ["SLURM_PROCID"]) - - ## Fall back to default - if world_size is None: - world_size = 1 - - return int(world_size), int(world_rank) - -def find_ifname(myaddr): - """ - Find socket ifname for a given ip adress. This is for "GLOO" ddp setup. - Usage example: - find_ifname("127.0.0.1") will return a network interface name, such as "lo". "lo0", etc. - """ - ipaddr = socket.gethostbyname(myaddr) - ifname = None - for nic, addrs in psutil.net_if_addrs().items(): - for addr in addrs: - if addr.address == ipaddr: - ifname = nic - break - if ifname is not None: - break - return ifname - -def parse_slurm_nodelist(nodelist): - """ - Parse SLURM_NODELIST env string to get list of nodes. - Usage example: - parse_slurm_nodelist(os.environ["SLURM_NODELIST"]) - Input examples: - "or-condo-g04" - "or-condo-g[05,07-08,13]" - "or-condo-g[05,07-08,13],or-condo-h[01,12]" - """ - nlist = list() - for block, _ in re.findall(r"([\w-]+(\[[\d\-,]+\])*)", nodelist): - m = re.match(r"^(?P[\w\-]+)\[(?P.*)\]", block) - if m is None: - ## single node - nlist.append(block) - else: - ## multiple nodes - g = m.groups() - prefix = g[0] - for sub in g[1].split(","): - if "-" in sub: - start, end = re.match(r"(\d+)-(\d+)", sub).groups() - fmt = "%%0%dd" % (len(start)) - for i in range(int(start), int(end) + 1): - node = prefix + fmt % i - nlist.append(node) - else: - node = prefix + sub - nlist.append(node) - - return nlist - -def setup_ddp(): - """ "Initialize DDP""" - - if os.getenv("HYDRAGNN_BACKEND") is not None: - backend = os.environ["HYDRAGNN_BACKEND"] - elif dist.is_nccl_available() and torch.cuda.is_available(): - backend = "nccl" - elif torch.distributed.is_gloo_available(): - backend = "gloo" - else: - raise RuntimeError("No parallel backends available") - - world_size, world_rank = init_comm_size_and_rank() - - ## Default setting - master_addr = "127.0.0.1" - master_port = "8889" - - if os.getenv("LSB_HOSTS") is not None: - ## source: https://www.olcf.ornl.gov/wp-content/uploads/2019/12/Scaling-DL-on-Summit.pdf - ## The following is Summit specific - master_addr = os.environ["LSB_HOSTS"].split()[1] - elif os.getenv("LSB_MCPU_HOSTS") is not None: - master_addr = os.environ["LSB_MCPU_HOSTS"].split()[2] - elif os.getenv("SLURM_NODELIST") is not None: - ## The following is CADES specific - master_addr = parse_slurm_nodelist(os.environ["SLURM_NODELIST"])[0] - - try: - if backend in ["nccl", "gloo"]: - os.environ["MASTER_ADDR"] = master_addr - os.environ["MASTER_PORT"] = master_port - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["RANK"] = str(world_rank) - - if (backend == "gloo") and ("GLOO_SOCKET_IFNAME" not in os.environ): - ifname = find_ifname(master_addr) - if ifname is not None: - os.environ["GLOO_SOCKET_IFNAME"] = ifname - - print( - "Distributed data parallel: %s master at %s:%s" - % (backend, master_addr, master_port), - ) +import mpi4py - if not dist.is_initialized(): - dist.init_process_group(backend=backend, init_method="env://") - - except KeyError: - print("DDP has to be initialized within a job - Running in sequential mode") - - return world_size, world_rank +mpi4py.rc.thread_level = "serialized" +mpi4py.rc.threads = False +from mpi4py import MPI +import distdataset +from distdataset import DistDataset -parser = argparse.ArgumentParser(description='VAE MNIST Example') -parser.add_argument('--batch-size', type=int, default=128, metavar='N', - help='input batch size for training (default: 128)') -parser.add_argument('--epochs', type=int, default=10, metavar='N', - help='number of epochs to train (default: 10)') -parser.add_argument('--no-cuda', action='store_true', default=False, - help='disables CUDA training') -parser.add_argument('--no-mps', action='store_true', default=False, - help='disables macOS GPU training') -parser.add_argument('--seed', type=int, default=1, metavar='S', - help='random seed (default: 1)') -parser.add_argument('--log-interval', type=int, default=10, metavar='N', - help='how many batches to wait before logging training status') +from ddp_utils import setup_ddp, get_local_rank +from vae_model import VAE, loss_function + +parser = argparse.ArgumentParser(description="VAE MNIST Example") +parser.add_argument( + "--batch-size", + type=int, + default=128, + metavar="N", + help="input batch size for training (default: 128)", +) +parser.add_argument( + "--epochs", + type=int, + default=10, + metavar="N", + help="number of epochs to train (default: 10)", +) +parser.add_argument( + "--no-cuda", action="store_true", default=False, help="disables CUDA training" +) +parser.add_argument( + "--no-mps", action="store_true", default=False, help="disables macOS GPU training" +) +parser.add_argument( + "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)" +) +parser.add_argument( + "--log-interval", + type=int, + default=10, + metavar="N", + help="how many batches to wait before logging training status", +) 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() torch.manual_seed(args.seed) +comm = MPI.COMM_WORLD +comm_size, rank = setup_ddp() +local_rank = get_local_rank(rank) + if args.cuda: - device = torch.device("cuda") + if torch.cuda.device_count() > 1: + local_rank = get_local_rank(rank) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + device = torch.device("cuda") +elif hasattr(torch, "xpu") and torch.xpu.is_available(): + if torch.xpu.device_count() > 1: + torch.xpu.set_device(local_rank) + device = torch.device(f"xpu:{local_rank}") + else: + device = torch.device("xpu") elif use_mps: device = torch.device("mps") else: device = torch.device("cpu") -class VAE(nn.Module): - def __init__(self): - 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) - - def encode(self, x): - h1 = F.relu(self.fc1(x)) - return self.fc21(h1), self.fc22(h1) - - def reparameterize(self, mu, logvar): - std = torch.exp(0.5*logvar) - eps = torch.randn_like(std) - return mu + eps*std +print("DDP setup:", comm_size, rank, device) - def decode(self, z): - h3 = F.relu(self.fc3(z)) - return torch.sigmoid(self.fc4(h3)) - - def forward(self, x): - mu, logvar = self.encode(x.view(-1, 784)) - z = self.reparameterize(mu, logvar) - return self.decode(z), mu, logvar - -comm = MPI.COMM_WORLD -comm_size, rank = setup_ddp() -print ("DDP setup:", comm_size, rank, device) +if rank == 0: + os.makedirs("results", exist_ok=True) +comm.Barrier() model = VAE().to(device) model = torch.nn.parallel.DistributedDataParallel(model) @@ -211,27 +96,25 @@ def forward(self, x): # kwargs = {'pin_memory': True} if args.cuda else {} kwargs = {} -trainset = DistDataset(datasets.MNIST('data', train=True, download=True,transform=transforms.ToTensor()), "train", comm) +trainset = DistDataset( + datasets.MNIST("data", train=True, download=True, transform=transforms.ToTensor()), + "train", + comm, +) # 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) - -testset = datasets.MNIST('data', train=False, download=True,transform=transforms.ToTensor()) -test_loader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=False, **kwargs) - -# Reconstruction + KL divergence losses summed over all elements and batch -def loss_function(recon_x, x, mu, logvar): - BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum') +train_loader = torch.utils.data.DataLoader( + trainset, batch_size=args.batch_size, shuffle=False, **kwargs, sampler=sampler +) - # see Appendix B from VAE paper: - # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014 - # https://arxiv.org/abs/1312.6114 - # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2) - KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) - - return BCE + KLD +testset = datasets.MNIST( + "data", train=False, download=True, transform=transforms.ToTensor() +) +test_loader = torch.utils.data.DataLoader( + testset, batch_size=args.batch_size, shuffle=False, **kwargs +) def train(epoch): @@ -255,16 +138,25 @@ def train(epoch): optimizer.step() # print(rank, "step") if batch_idx % args.log_interval == 0: - print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( - epoch, batch_idx * len(data), len(train_loader.dataset), - 100. * batch_idx / len(train_loader), - loss.item() / len(data))) + print( + "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( + epoch, + batch_idx * len(data), + len(train_loader.dataset), + 100.0 * batch_idx / len(train_loader), + loss.item() / len(data), + ) + ) train_loader.dataset.ddstore.epoch_begin() - + train_loader.dataset.ddstore.epoch_end() - print('====> Epoch: {} Average loss: {:.4f}'.format( - epoch, train_loss / len(train_loader.dataset))) + if rank == 0: + print( + "====> Epoch: {} Average loss: {:.4f}".format( + epoch, train_loss / len(train_loader.dataset) + ) + ) def test(epoch): @@ -277,22 +169,30 @@ def test(epoch): test_loss += loss_function(recon_batch, data, mu, logvar).item() if i == 0: n = min(data.size(0), 8) - comparison = torch.cat([data[:n], - recon_batch.view(args.batch_size, 1, 28, 28)[:n]]) - save_image(comparison.cpu(), - 'results/reconstruction_' + str(epoch) + '.png', nrow=n) + comparison = torch.cat( + [data[:n], recon_batch.view(args.batch_size, 1, 28, 28)[:n]] + ) + save_image( + comparison.cpu(), + "results/reconstruction_" + str(epoch) + ".png", + nrow=n, + ) test_loss /= len(test_loader.dataset) - print('====> Test set loss: {:.4f}'.format(test_loss)) + print("====> Test set loss: {:.4f}".format(test_loss)) + if __name__ == "__main__": # print("main", rank) for epoch in range(1, args.epochs + 1): train(epoch) - test(epoch) - with torch.no_grad(): - 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') - + if rank == 0: + test(epoch) + with torch.no_grad(): + 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" + ) + + dist.destroy_process_group() diff --git a/examples/vae/vae_core_server.py b/examples/vae/vae_core_server.py new file mode 100644 index 0000000..5956f8e --- /dev/null +++ b/examples/vae/vae_core_server.py @@ -0,0 +1,100 @@ +""" +VAE example — core (data-holding) group. + +Loads the MNIST training set, shards it across the ranks of MPI.COMM_WORLD, +and publishes it via DDStore method=2 (file-based handshake) so a separate +extra group can train against it over RDMA. This script does no training +itself — it just holds the data in memory until the extra group signals it +is done. + +Usage: + srun -n python examples/vae/vae_core_server.py [handshake_dir] + +Environment: + DDSTORE_HANDSHAKE_DIR overrides handshake_dir positional arg + DDSTORE_HANDSHAKE_TIMEOUT_S poll timeout in seconds (default: 300) + DDSTORE_NIC_MAP optional precomputed CPU->NIC map (see + cpu_nic_map.py --env) for FABRIC_IFACE + auto-selection; falls back to a live + hwloc-calc query if unset +""" + +import os +import sys +import time + +## torch (pulled in below via torchvision/distdataset) 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. +from torchvision import datasets, transforms + +import mpi4py + +mpi4py.rc.thread_level = "serialized" +mpi4py.rc.threads = False +from mpi4py import MPI + +from distdataset import DistDataset + + +def _resolve_dir(arg): + if arg: + return arg + return os.environ.get("DDSTORE_HANDSHAKE_DIR", "./ddstore_hs") + + +hs_dir = _resolve_dir(sys.argv[1] if len(sys.argv) > 1 else "") +os.environ["DDSTORE_METHOD"] = "2" +os.environ["DDSTORE_HANDSHAKE_DIR"] = hs_dir + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() + +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": + 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() +) +dds_trainset = DistDataset(trainset, "train", comm) +comm.Barrier() + +if rank == 0: + print( + f"[core] published {len(dds_trainset)} training rows, waiting for extras...", + flush=True, + ) + +sentinel = os.path.join(hs_dir, "done_extra") +timeout_s = int(os.environ.get("DDSTORE_HANDSHAKE_TIMEOUT_S", "300")) +t0 = time.monotonic() + +if rank == 0: + while not os.path.exists(sentinel): + if time.monotonic() - t0 > timeout_s: + raise TimeoutError("timed out waiting for done_extra sentinel") + time.sleep(0.5) + print("[core] sentinel received, shutting down", flush=True) +comm.Barrier() + +dds_trainset.ddstore.free() + +if rank == 0: + for fname in os.listdir(hs_dir): + if fname.endswith(".bin"): + try: + os.remove(os.path.join(hs_dir, fname)) + except OSError: + pass + try: + os.remove(sentinel) + except OSError: + pass + +print(f"[core rank {rank}] done", flush=True) diff --git a/examples/vae/vae_extra_train.py b/examples/vae/vae_extra_train.py new file mode 100644 index 0000000..9cda3a2 --- /dev/null +++ b/examples/vae/vae_extra_train.py @@ -0,0 +1,216 @@ +""" +VAE example — extra (read-only) group. + +Trains the VAE model using data it never loads itself: every batch is read +over RDMA from a core group (see vae_core_server.py) via DDStore method=2's +file-based handshake. The extra ranks form their own DDP process group +among themselves, entirely separate from the core group's communicator. + +Usage: + srun -n python examples/vae/vae_extra_train.py \\ + --handshake-dir ddstore_hs --n-core 4 --epochs 10 + +Environment (used as defaults if the matching --flag is not given): + DDSTORE_HANDSHAKE_DIR shared handshake directory + DDSTORE_N_CORE number of core ranks that published the data + DDSTORE_HANDSHAKE_TIMEOUT_S poll timeout in seconds (default: 300) + DDSTORE_NIC_MAP optional precomputed CPU->NIC map (see + cpu_nic_map.py --env) for FABRIC_IFACE + auto-selection; falls back to a live + hwloc-calc query if unset +""" + +from __future__ import print_function + +import argparse +import os + +## torch (and the RCCL/HIP shared libraries it pulls in) must finish loading +## before mpi4py triggers MPI_Init, or their static destructors run in the +## wrong order at interpreter exit and corrupt the heap. +## Do not reorder these imports. +import torch +import torch.utils.data +from torch import optim +from torchvision import datasets, transforms +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 + +parser = argparse.ArgumentParser(description="VAE MNIST Example - extra (reader) group") +parser.add_argument( + "--batch-size", + type=int, + default=128, + metavar="N", + help="input batch size for training (default: 128)", +) +parser.add_argument( + "--epochs", + type=int, + default=10, + metavar="N", + help="number of epochs to train (default: 10)", +) +parser.add_argument( + "--no-cuda", action="store_true", default=False, help="disables CUDA training" +) +parser.add_argument( + "--no-mps", action="store_true", default=False, help="disables macOS GPU training" +) +parser.add_argument( + "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)" +) +parser.add_argument( + "--log-interval", + type=int, + default=10, + metavar="N", + help="how many batches to wait before logging training status", +) +parser.add_argument( + "--handshake-dir", + type=str, + default=os.environ.get("DDSTORE_HANDSHAKE_DIR", "./ddstore_hs"), + help="shared directory published by vae_core_server.py", +) +parser.add_argument( + "--n-core", + type=int, + default=int(os.environ.get("DDSTORE_N_CORE", "4")), + help="number of core ranks that published the data", +) +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() + +torch.manual_seed(args.seed) + +comm_size, rank = setup_ddp() +local_rank = get_local_rank(rank) + +if args.cuda: + if torch.cuda.device_count() > 1: + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + device = torch.device("cuda") +elif hasattr(torch, "xpu") and torch.xpu.is_available(): + if torch.xpu.device_count() > 1: + torch.xpu.set_device(local_rank) + device = torch.device(f"xpu:{local_rank}") + else: + device = torch.device("xpu") +elif use_mps: + device = torch.device("mps") +else: + device = torch.device("cpu") + +print("DDP setup:", comm_size, rank, device) + +model = VAE().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 +) + +testset = datasets.MNIST( + "data", train=False, download=True, transform=transforms.ToTensor() +) +test_loader = torch.utils.data.DataLoader( + testset, batch_size=args.batch_size, shuffle=False, **kwargs +) + + +def train(epoch): + model.train() + train_loss = 0 + train_loader.dataset.ddstore.epoch_begin() + for batch_idx, (data, _) in enumerate(train_loader): + train_loader.dataset.ddstore.epoch_end() + data = data.to(device) + optimizer.zero_grad() + recon_batch, mu, logvar = model(data) + loss = loss_function(recon_batch, data, mu, logvar) + loss.backward() + train_loss += loss.item() + optimizer.step() + if batch_idx % args.log_interval == 0: + print( + "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( + epoch, + batch_idx * len(data), + len(train_loader.dataset), + 100.0 * batch_idx / len(train_loader), + loss.item() / len(data), + ) + ) + + train_loader.dataset.ddstore.epoch_begin() + + train_loader.dataset.ddstore.epoch_end() + print( + "====> Epoch: {} Average loss: {:.4f}".format( + epoch, train_loss / len(train_loader.dataset) + ) + ) + + +def test(epoch): + model.eval() + test_loss = 0 + with torch.no_grad(): + for i, (data, _) in enumerate(test_loader): + data = data.to(device) + recon_batch, mu, logvar = model(data) + test_loss += loss_function(recon_batch, data, mu, logvar).item() + if i == 0: + n = min(data.size(0), 8) + comparison = torch.cat( + [data[:n], recon_batch.view(args.batch_size, 1, 28, 28)[:n]] + ) + save_image( + comparison.cpu(), + "results/extra_reconstruction_" + str(epoch) + ".png", + nrow=n, + ) + + test_loss /= len(test_loader.dataset) + print("====> Test set loss: {:.4f}".format(test_loss)) + + +if __name__ == "__main__": + for epoch in range(1, args.epochs + 1): + train(epoch) + test(epoch) + with torch.no_grad(): + sample = torch.randn(64, 20).to(device) + sample = model.module.decode(sample).cpu() + save_image( + sample.view(64, 1, 28, 28), + "results/extra_sample_" + str(epoch) + ".png", + ) + + if rank == 0: + sentinel = os.path.join(args.handshake_dir, "done_extra") + with open(sentinel, "w") as f: + f.write("done\n") + print("[extra] sentinel written, exiting", flush=True) + + dist.destroy_process_group() diff --git a/examples/vae/vae_model.py b/examples/vae/vae_model.py new file mode 100644 index 0000000..b8f435e --- /dev/null +++ b/examples/vae/vae_model.py @@ -0,0 +1,45 @@ +import torch +from torch import nn +from torch.nn import functional as F + + +class VAE(nn.Module): + def __init__(self): + 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) + + def encode(self, x): + h1 = F.relu(self.fc1(x)) + return self.fc21(h1), self.fc22(h1) + + def reparameterize(self, mu, logvar): + std = torch.exp(0.5 * logvar) + eps = torch.randn_like(std) + return mu + eps * std + + def decode(self, z): + h3 = F.relu(self.fc3(z)) + return torch.sigmoid(self.fc4(h3)) + + def forward(self, x): + mu, logvar = self.encode(x.view(-1, 784)) + 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") + + # see Appendix B from VAE paper: + # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014 + # https://arxiv.org/abs/1312.6114 + # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2) + KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) + + return BCE + KLD diff --git a/include/common.h b/include/common.h index 0ce4aa7..18318ef 100644 --- a/include/common.h +++ b/include/common.h @@ -4,11 +4,30 @@ #include #include #include +#include #include #define DP_AV_DEF_SIZE 512 #define COMM_FILE_WRITER_TO_READER "./writer_address.bin" +/* ----------------------------------------------------------------------- + * Method 2: file-based handshake record, one per core rank. + * All n_core records for a variable are gathered in memory (via MPI among + * core ranks) and published as a single combined file, written once by + * core rank 0: + * {handshake_dir}/{varname}.bin + * ----------------------------------------------------------------------- */ +struct CoreRecord +{ + char fabric_address[DP_AV_DEF_SIZE]; /* raw fi_getname output */ + size_t fabric_address_len; /* actual bytes used */ + uint64_t key; /* MR key from fi_mr_key() */ + uint64_t base_address; /* virtual address of send_data */ + long nrows; /* rows owned by this core rank */ + int disp; /* elements per row */ + int itemsize; /* bytes per element */ +}; + #ifdef __cplusplus extern "C" { @@ -44,10 +63,73 @@ extern "C" return (f->info->mode & FI_LOCAL_MR) != 0; } + /* 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". */ + static bool is_mr_endpoint(struct fabric_state *f) + { + return (f->info->domain_attr->mr_mode & FI_MR_ENDPOINT) != 0 || + (f->info->fabric_attr->prov_name && + strcmp(f->info->fabric_attr->prov_name, "cxi") == 0); + } + + /* 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, + * 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. */ + 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); + } + 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); + /* --- Method 2: file-based handshake ---------------------------------- */ + + /* Resolve the handshake directory (priority: user_dir > env var > cwd). + * Creates the directory if it does not exist. + * Returns pointer to a static buffer — copy if needed across calls. */ + const char *resolve_handshake_dir(const char *user_dir); + + /* Core rank: exchange CoreRecords with all other core ranks via + * MPI_Allgather over `comm` (no filesystem round-trip needed for + * core-to-core discovery), populate this rank's fs->comm_partner[], + * remote_key[], remote_address[], and fill lenlist[0..n_core-1] (raw + * row counts, NOT yet prefix-summed). Rank 0 additionally publishes + * the combined record set to {dir}/{varname}.bin (tmp + fsync + rename) + * so extra members can join later. */ + int handshake_write(struct fabric_state *fs, MPI_Comm comm, + const char *dir, const char *varname, + int n_core, long nrows, int disp, int itemsize, + long *lenlist); + + /* Extra member: poll for {dir}/{varname}.bin (single file holding all + * n_core CoreRecords), read it, and populate this process's + * fs->comm_partner[], remote_key[], remote_address[], and + * lenlist[0..n_core-1] (raw row counts, NOT yet prefix-summed). Does + * NOT write anything. Blocks until the file appears (with timeout). */ + int handshake_join(struct fabric_state *fs, + const char *dir, const char *varname, + int n_core, + long *lenlist, int *out_disp, int *out_itemsize); + #ifdef __cplusplus } #endif diff --git a/include/ddstore.hpp b/include/ddstore.hpp index 6543cb4..2c582ee 100644 --- a/include/ddstore.hpp +++ b/include/ddstore.hpp @@ -5,6 +5,10 @@ #include #include #include +#include +#include +#include +#include #include "common.h" struct VarInfo @@ -29,13 +33,36 @@ class DDStore DDStore(); DDStore(MPI_Comm comm); DDStore(int method, MPI_Comm comm); + + /* Method 2: core member constructor. + * handshake_dir: shared directory visible to all processes. + * n_core is derived from the communicator size (all ranks in `comm` + * are assumed to be core members). */ + DDStore(int method, MPI_Comm comm, + const std::string &handshake_dir); + + /* Method 2: extra member constructor (no MPI communicator required). + * The extra member calls join() per variable to discover the data. */ + DDStore(int method, const std::string &handshake_dir, int n_core); + ~DDStore(); void query(std::string name, VarInfo_t &varinfo); + + /* Total row count across all core ranks for a variable (added or joined). */ + long size(std::string name) + { + const VarInfo_t& varinfo = this->varlist.at(name); + return varinfo.lenlist.empty() ? 0 : varinfo.lenlist.back(); + } + void epoch_begin(); void epoch_end(); void free(); + /* Method 2 extra member: discover variable published by core members. */ + void join(std::string name); + template void add(std::string name, T *buffer, long nrows, int disp) { @@ -71,7 +98,79 @@ class DDStore init_fabric(fabric_state); if (!fabric_state->info) throw std::runtime_error("init_fabric failed: no suitable fabric found"); - handshake(fabric_state, this->comm); + 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; + + 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 (mr_rc != FI_SUCCESS) + throw std::runtime_error(std::string("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 + * fi_mr_enable() — same requirement as method=1's handshake(). + * No-op for hsn (is_mr_endpoint() is false). */ + if (is_mr_endpoint(fabric_state)) + { + int rc = fi_mr_bind(fabric_state->mr, &fabric_state->signal->fid, 0); + if (rc != FI_SUCCESS) + throw std::runtime_error(std::string("fi_mr_bind failed: ") + fi_strerror(rc)); + rc = fi_mr_enable(fabric_state->mr); + if (rc != FI_SUCCESS) + throw std::runtime_error(std::string("fi_mr_enable failed: ") + fi_strerror(rc)); + } + fabric_state->key = fi_mr_key(fabric_state->mr); + + /* Exchange records with all core ranks via MPI_Allgather, and + * (rank 0 only) publish the combined record set for extra + * members to join later. */ + std::vector raw_lens(this->n_core); + if (handshake_write(fabric_state, this->comm, + this->handshake_dir.c_str(), name.c_str(), + this->n_core, nrows, disp, (int)sizeof(T), + raw_lens.data()) != 0) + throw std::runtime_error("handshake_write failed"); + + /* Build prefix-sum lenlist from the raw per-rank row counts. */ + long sum = 0; + std::vector lenlist(this->n_core); + for (int i = 0; i < this->n_core; i++) + { + sum += raw_lens[i]; + lenlist[i] = sum; + } + + VarInfo_t var; + var.name = name; + var.itemsize = (int)sizeof(T); + var.disp = disp; + var.win = MPI_WIN_NULL; + var.lenlist = lenlist; + var.active = true; + var.fence_active = false; + var.base = base; + var.fabric_state = fabric_state; + this->varlist.insert(std::pair(name, var)); + return; /* lenlist already stored; skip the MPI_Allgather block below */ } std::vector lenlist(this->comm_size); @@ -89,11 +188,6 @@ class DDStore sum += lenlist[i]; lenlist[i] = sum; } - // for (long unsigned int i = 0; i < lenlist.size(); i++) - // { - // std::cout << "lenlist[" << i << "]: " << lenlist[i] << std::endl; - // } - // std::cout << "sum: " << sum << std::endl; VarInfo_t var; var.name = name; @@ -144,7 +238,74 @@ class DDStore init_fabric(fabric_state); if (!fabric_state->info) throw std::runtime_error("init_fabric failed: no suitable fabric found"); - handshake(fabric_state, this->comm); + 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 * itemsize; + 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"); + + 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 (mr_rc != FI_SUCCESS) + throw std::runtime_error(std::string("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 + * fi_mr_enable() — same requirement as method=1's handshake(). + * No-op for hsn (is_mr_endpoint() is false). */ + if (is_mr_endpoint(fabric_state)) + { + int rc = fi_mr_bind(fabric_state->mr, &fabric_state->signal->fid, 0); + if (rc != FI_SUCCESS) + throw std::runtime_error(std::string("fi_mr_bind failed: ") + fi_strerror(rc)); + rc = fi_mr_enable(fabric_state->mr); + if (rc != FI_SUCCESS) + throw std::runtime_error(std::string("fi_mr_enable failed: ") + fi_strerror(rc)); + } + fabric_state->key = fi_mr_key(fabric_state->mr); + + std::vector raw_lens(this->n_core); + if (handshake_write(fabric_state, this->comm, + this->handshake_dir.c_str(), name.c_str(), + this->n_core, nrows, disp, itemsize, + raw_lens.data()) != 0) + throw std::runtime_error("handshake_write failed"); + + long sum = 0; + std::vector lenlist(this->n_core); + for (int i = 0; i < this->n_core; i++) + { + sum += raw_lens[i]; + lenlist[i] = sum; + } + + VarInfo_t var; + var.name = name; + var.itemsize = itemsize; + var.disp = disp; + var.win = MPI_WIN_NULL; + var.lenlist = lenlist; + var.active = true; + var.fence_active = false; + var.base = base; + var.fabric_state = fabric_state; + this->varlist.insert(std::pair(name, var)); + return; } std::vector lenlist(this->comm_size); @@ -162,11 +323,6 @@ class DDStore sum += lenlist[i]; lenlist[i] = sum; } - // for (long unsigned int i = 0; i < lenlist.size(); i++) - // { - // std::cout << "lenlist[" << i << "]: " << lenlist[i] << std::endl; - // } - // std::cout << "sum: " << sum << std::endl; VarInfo_t var; var.name = name; @@ -193,8 +349,6 @@ class DDStore if (itemsize != sizeof(T)) throw std::invalid_argument("Invalid data type"); - // std::cout << "Update: " << name << ", nrows: " << nrows << ", offset: " << offset << std::endl; - // std::cout << "memcpy: " << (nrows * disp * itemsize)/1024/1024/1024 << " GB" << std::endl; memcpy((char*)base + offset * disp * itemsize, buffer, nrows * disp * itemsize); } @@ -240,23 +394,30 @@ class DDStore win /* window object */); MPI_Win_unlock(target, win); } - else if (this->method == 1) + else if (this->method == 1 || this->method == 2) { - // printf("varinfo.disp, T, count: %d %d %d\n", varinfo.disp, sizeof(T), count); - // printf("target, offset: %d %d\n", target, offset); - + /* Methods 1 and 2 both use libfabric fi_read — same path. */ varinfo.fabric_state->recv_data = (char *)buffer; varinfo.fabric_state->recv_data_len = varinfo.disp * varinfo.itemsize * count; - read_from_remote(varinfo.fabric_state, target, (start - offset) * varinfo.disp * varinfo.itemsize); + int rc = read_from_remote(varinfo.fabric_state, target, (start - offset) * varinfo.disp * varinfo.itemsize); + if (rc != 0) + throw std::runtime_error( + "read_from_remote failed with code " + std::to_string(rc) + + " (target=" + std::to_string(target) + ")"); } } private: - int method; // 0: MPI, 1: libfabric + int method; // 0: MPI, 1: libfabric, 2: file-based handshake (libfabric transport) + + MPI_Comm comm; + int comm_size; + int rank; - MPI_Comm comm; - int comm_size; - int rank; + /* Method 2 fields */ + std::string handshake_dir; /* shared directory for CoreRecord files */ + int n_core; /* number of core ranks */ + bool is_extra; /* true if this is an extra (read-only) node */ std::unordered_map varlist; }; diff --git a/setup.py b/setup.py index 75ed248..f91dcac 100644 --- a/setup.py +++ b/setup.py @@ -9,7 +9,7 @@ import os import subprocess -defs = [('NPY_NO_DEPRECATED_API', 0)] +defs = [("NPY_NO_DEPRECATED_API", 0)] include_dirs = list() library_dirs = list() libraries = list() @@ -31,20 +31,25 @@ include_dirs.append(np.get_include()) include_dirs.append("include") -extending = Extension("pyddstore", - sources=["src/pyddstore.pyx", "src/ddstore.cxx", "src/common.cxx"], - include_dirs=include_dirs, - extra_compile_args=["-std=c++11"], - define_macros=defs, - library_dirs=library_dirs, - libraries=libraries, - ) +extending = Extension( + "pyddstore", + sources=["src/pyddstore.pyx", "src/ddstore.cxx", "src/common.cxx"], + include_dirs=include_dirs, + extra_compile_args=["-std=c++11"], + define_macros=defs, + library_dirs=library_dirs, + libraries=libraries, +) -extensions = [extending,] +extensions = [ + extending, +] setup( name="PyDDStore", version="0.1", description="Distributed Data Store", - ext_modules=cythonize(extensions) + package_dir={"": "src"}, + py_modules=["cpu_nic_map"], + ext_modules=cythonize(extensions), ) diff --git a/src/common.cxx b/src/common.cxx index 7935496..3885c3e 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -9,8 +9,15 @@ #include #include #include - -void init_fabric(struct fabric_state *fabric) +#include +#include +#include +#include +#include + +/* hsn (tcp;ofi_rxm over Slingshot) path — Frontier. Unchanged from the + * already-proven dev-file2 implementation; only renamed (was init_fabric). */ +static void init_fabric_hsn(struct fabric_state *fabric) { struct fi_info *hints, *info, *originfo, *useinfo; struct fi_av_attr av_attr = {FI_AV_UNSPEC}; @@ -250,6 +257,284 @@ void init_fabric(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. */ +static void init_fabric_cxi(struct fabric_state *fabric) +{ + struct fi_info *info, *originfo, *useinfo; + struct fi_av_attr av_attr = {FI_AV_UNSPEC}; + struct fi_cq_attr cq_attr = {0}; + char *ifname; + int result; + + ifname = getenv("FABRIC_IFACE"); + fabric->info = NULL; + + int version = fi_version(); + + /* IMPORTANT: fi_getinfo() must be called with a literal NULL hints + * pointer. Passing ANY hints struct (even one with only ep_attr->type + * set) causes fi_getinfo to return NULL inside this process (torch + + * NCCL + mpi4py all loaded) on Perlmutter compute nodes, even though a + * plain MPI-only C program does not exhibit this. With a literal NULL + * hints pointer, fi_getinfo returns fully-populated fi_info entries + * (mr_mode, max_msg_size, tx_attr, rx_attr all correctly set) — verified + * on Perlmutter via a standalone test binary run under the same srun + * job. Do NOT "improve" this by adding hints back without re-verifying + * against a real running job (not just fi_info on the login node). */ + fi_getinfo(version, NULL, NULL, 0, NULL, &info); + if (!info) + { + fprintf(stderr, "no fabrics detected.\n"); + return; + } + + originfo = info; + useinfo = NULL; + while (info) + { + char *prov_name = info->fabric_attr->prov_name; + char *domain_name = info->domain_attr->name; + + /* If FABRIC_IFACE is set, match by domain name against cxi or tcp;ofi_rxm. */ + if (ifname && domain_name && (strcmp(ifname, domain_name) == 0) && + (strcmp(prov_name, "cxi") == 0 || strcmp(prov_name, "tcp;ofi_rxm") == 0)) + { + fprintf(stderr, "using interface set by FABRIC_IFACE: %s (%s).\n", + domain_name, prov_name); + useinfo = info; + break; + } + if ((strcmp(prov_name, "cxi") == 0 || + (strcmp(prov_name, "verbs") == 0 && info->src_addr) || + strcmp(prov_name, "gni") == 0 || + strcmp(prov_name, "psm2") == 0) && + (!useinfo || (ifname && domain_name && + strcmp(useinfo->domain_attr->name, ifname) != 0))) + { + useinfo = info; + } + info = info->next; + } + + info = useinfo; + + if (!info) + { + fprintf( + stderr, + "none of the usable system fabrics are supported high speed " + "interfaces (cxi, verbs, gni, psm2.) To use a compatible fabric " + "that is being ignored (probably sockets), set the environment " + "variable FABRIC_IFACE to the interface name. Check the output " + "of fi_info to troubleshoot this message.\n"); + fi_freeinfo(originfo); + return; + } + + if (info->mode & FI_CONTEXT2) + { + fabric->ctx = (fi_context*) calloc(2, sizeof(*fabric->ctx)); + } + else if (info->mode & FI_CONTEXT) + { + fabric->ctx = (fi_context*) calloc(1, sizeof(*fabric->ctx)); + } + else + { + fabric->ctx = NULL; + } + + /* For non-CXI providers, clear FI_MR_BASIC (legacy flag). For CXI, the + * real fi_info from NULL-hints fi_getinfo already has the correct + * mr_mode/max_msg_size/tx_attr/rx_attr — do NOT override them. */ + if (!info->fabric_attr->prov_name || + strcmp(info->fabric_attr->prov_name, "cxi") != 0) + { + info->domain_attr->mr_mode = 0; + } +#ifdef SST_HAVE_CRAY_DRC + if (strstr(info->fabric_attr->prov_name, "gni") && fabric->auth_key) + { + info->domain_attr->auth_key = (uint8_t *)fabric->auth_key; + info->domain_attr->auth_key_size = sizeof(struct fi_gni_raw_auth_key); + } +#endif /* SST_HAVE_CRAY_DRC */ + + fabric->info = fi_dupinfo(info); + if (!fabric->info) + { + fprintf(stderr, "copying the fabric info failed.\n"); + fi_freeinfo(originfo); + return; + } + info = fabric->info; + + result = fi_fabric(info->fabric_attr, &fabric->fabric, fabric->ctx); + if (result != FI_SUCCESS) + { + fprintf( + stderr, + "opening fabric access failed with %d (%s). This is fatal.\n", + result, + fi_strerror(result)); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + result = fi_domain(fabric->fabric, info, &fabric->domain, fabric->ctx); + if (result != FI_SUCCESS) + { + fprintf( + stderr, + "accessing domain failed with %d (%s). This is fatal.\n", + result, + fi_strerror(result)); + fprintf( + stderr, + "fi_domain() has failed, which may mean that libfabric is " + "defaulting to the wrong interface, or that no CXI service is " + "available to this user on this node. Check your FABRIC_IFACE " + "environment variable (or specify one).\n"); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + + /* CRITICAL: the fi_info returned by fi_getinfo has tx_attr->op_flags / + * rx_attr->op_flags with FI_INJECT set by default for CXI. libfabric's + * simple (non-msg) calls like fi_read()/fi_write() implicitly use the + * endpoint's default op_flags, so leaving FI_INJECT set causes CXI to + * treat every fi_read/fi_write as an inject operation, capping transfer + * size at tx_attr->inject_size (192 bytes on this system) and failing + * larger transfers with EMSGSIZE ("Message too long"). Clear op_flags + * before fi_endpoint() so normal-sized RMA reads/writes work. */ + info->tx_attr->op_flags = 0; + info->rx_attr->op_flags = 0; + /* Do NOT override ep_attr->type here — the fi_info already has the + * correct type from fi_getinfo and changing it after fi_domain causes + * fi_endpoint to fail with EINVAL (CXI). */ + + result = fi_endpoint(fabric->domain, info, &fabric->signal, fabric->ctx); + if (result != FI_SUCCESS || !fabric->signal) + { + fprintf( + stderr, + "opening endpoint failed with %d (%s). This is fatal.\n", + result, + fi_strerror(result)); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + + /* Query real max_msg_size from the created endpoint — necessary for CXI + * where fi_getinfo with NULL hints returns max_msg_size=0. */ + { + size_t max_msg_size = 0; + size_t optlen = sizeof(max_msg_size); + if (fi_getopt(&fabric->signal->fid, FI_OPT_ENDPOINT, + FI_OPT_MAX_MSG_SIZE, &max_msg_size, &optlen) == FI_SUCCESS + && max_msg_size > 0) + { + info->ep_attr->max_msg_size = max_msg_size; + fprintf(stderr, "endpoint max_msg_size=%zu\n", max_msg_size); + } + } + + av_attr.type = FI_AV_MAP; + av_attr.count = DP_AV_DEF_SIZE; + av_attr.ep_per_node = 0; + result = fi_av_open(fabric->domain, &av_attr, &fabric->av, fabric->ctx); + if (result != FI_SUCCESS) + { + fprintf( + stderr, + "could not initialize address vector, failed with %d " + "(%s). This is fatal.\n", + result, + fi_strerror(result)); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + result = fi_ep_bind(fabric->signal, &fabric->av->fid, 0); + if (result != FI_SUCCESS) + { + fprintf( + stderr, + "could not bind endpoint to address vector, failed with " + "%d (%s). This is fatal.\n", + result, + fi_strerror(result)); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + + cq_attr.size = 0; + cq_attr.format = FI_CQ_FORMAT_DATA; + result = + fi_cq_open(fabric->domain, &cq_attr, &fabric->cq_signal, fabric->ctx); + if (result != FI_SUCCESS) + { + fprintf( + stderr, + "opening completion queue failed with %d (%s). This is fatal.\n", + result, + fi_strerror(result)); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + + result = fi_ep_bind( + fabric->signal, &fabric->cq_signal->fid, FI_TRANSMIT | FI_RECV); + if (result != FI_SUCCESS) + { + fprintf( + stderr, + "could not bind endpoint to completion queue, failed " + "with %d (%s). This is fatal.\n", + result, + fi_strerror(result)); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + + result = fi_enable(fabric->signal); + if (result != FI_SUCCESS) + { + fprintf( + stderr, + "enable endpoint, failed with %d (%s). This is fatal.\n", + result, + fi_strerror(result)); + fi_freeinfo(originfo); + fabric->info = NULL; + return; + } + + if (originfo) fi_freeinfo(originfo); +} + +/* Dispatch: DDSTORE_FABRIC selects the fabric-open implementation. + * Default (unset) is "hsn", so Frontier's behavior is unaffected unless + * something explicitly opts into "cxi". Not runtime fi_getinfo detection — + * an earlier attempt to unify both paths via runtime detection compiled + * cleanly but crashed hsn on Frontier at runtime, so hsn and cxi are kept + * as two independent, individually-proven implementations instead. */ +void init_fabric(struct fabric_state *fabric) +{ + const char *provider = getenv("DDSTORE_FABRIC"); + if (provider && strcmp(provider, "cxi") == 0) + init_fabric_cxi(fabric); + else + init_fabric_hsn(fabric); +} + int handshake(struct fabric_state *fabric_state, MPI_Comm comm) { char address[DP_AV_DEF_SIZE]; @@ -272,6 +557,25 @@ int handshake(struct fabric_state *fabric_state, MPI_Comm comm) fprintf(stderr, "fi_mr_reg failed: %s\n", fi_strerror(mr_rc)); return 1; } + + /* CXI (FI_MR_ENDPOINT): bind MR to endpoint and enable it before use. + * The provider-assigned key is only valid after fi_mr_enable(). No-op + * for hsn/verbs/gni/psm2 (is_mr_endpoint() is false for those). */ + if (is_mr_endpoint(fabric_state)) + { + int rc = fi_mr_bind(fabric_state->mr, &fabric_state->signal->fid, 0); + if (rc != FI_SUCCESS) + { + fprintf(stderr, "fi_mr_bind (send) failed: %s\n", fi_strerror(rc)); + return 1; + } + rc = fi_mr_enable(fabric_state->mr); + if (rc != FI_SUCCESS) + { + fprintf(stderr, "fi_mr_enable (send) failed: %s\n", fi_strerror(rc)); + return 1; + } + } fabric_state->key = fi_mr_key(fabric_state->mr); int status = fi_getname((fid_t)fabric_state->signal, address, &address_len); @@ -305,7 +609,13 @@ int handshake(struct fabric_state *fabric_state, MPI_Comm comm) } size_t *pointer_addr_data = (size_t *)malloc(world_size * sizeof(size_t)); - pointer_addr_data[rank] = (size_t)fabric_state->send_data; + /* With FI_MR_VIRT_ADDR the remote offset in fi_read is the virtual + * address; without it (e.g. CXI with FI_MR_PROV_KEY) it is 0-based + * from the MR registration base, so exchange 0. Always true (the real + * pointer) for hsn/verbs/gni/psm2. */ + pointer_addr_data[rank] = is_virt_addr(fabric_state) + ? (size_t)fabric_state->send_data + : 0; MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, pointer_addr_data, 1, MPI_UNSIGNED_LONG, comm); for (int i = 0; i < world_size; i++) @@ -334,6 +644,25 @@ int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset 0, &fabric_state->recv_mr, NULL); + + /* 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). */ + if (is_mr_endpoint(fabric_state)) + { + int rc_mr = fi_mr_bind(fabric_state->recv_mr, &fabric_state->signal->fid, 0); + if (rc_mr != FI_SUCCESS) + { + fprintf(stderr, "fi_mr_bind (recv) failed: %s\n", fi_strerror(rc_mr)); + return 1; + } + rc_mr = fi_mr_enable(fabric_state->recv_mr); + if (rc_mr != FI_SUCCESS) + { + fprintf(stderr, "fi_mr_enable (recv) failed: %s\n", fi_strerror(rc_mr)); + return 1; + } + } + void *memory_descriptor = NULL; if (is_local_mr_req(fabric_state)) { @@ -379,11 +708,315 @@ int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset { struct fi_cq_err_entry ee = {0}; fi_cq_readerr(fabric_state->cq_signal, &ee, 0); - fprintf(stderr, "fi_cq_read failed with error: prov_errno=%d (%s)\n", - ee.prov_errno, fi_strerror(ee.prov_errno)); + /* 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; + } + } + + return 0; +} + +/* ========================================================================= + * Method 2: file-based handshake helpers + * ========================================================================= + * + * File naming convention: + * {dir}/{varname}.bin — one combined file per variable, holding all + * n_core CoreRecords, written once by core rank 0 + * after an MPI_Allgather among core ranks. + * + * Directory resolution priority: + * 1. user_dir argument (non-empty string) + * 2. DDSTORE_HANDSHAKE_DIR environment variable + * 3. "./ddstore_hs" (current working directory fallback) + * + * The resolved directory must be on a shared filesystem (e.g. Lustre) + * visible to all nodes. It is created automatically if it does not exist. + * + * Default poll timeout: DDSTORE_HANDSHAKE_TIMEOUT_S env var, default 300 s. + * Poll interval: 50 ms. + * ========================================================================= */ + +/* Resolve the handshake directory. + * Returns a pointer to a static buffer — copy before next call. */ +extern "C" const char *resolve_handshake_dir(const char *user_dir) +{ + static char resolved[4096]; + + if (user_dir && user_dir[0] != '\0') + { + snprintf(resolved, sizeof(resolved), "%s", user_dir); + } + else + { + const char *env = getenv("DDSTORE_HANDSHAKE_DIR"); + if (env && env[0] != '\0') + snprintf(resolved, sizeof(resolved), "%s", env); + else + snprintf(resolved, sizeof(resolved), "./ddstore_hs"); + } + + /* Create the directory if it does not exist (best-effort; ignore EEXIST). */ + mkdir(resolved, 0755); + + return resolved; +} + +/* Build the canonical path for a variable's combined record file into `buf`. + * varname is sanitised: any '/' or '.' characters are replaced with '_' so + * that a caller-supplied name cannot escape the handshake directory. */ +static void record_path(char *buf, size_t bufsz, + const char *dir, const char *varname) +{ + /* Copy and sanitise varname into a local buffer. */ + char safe[256]; + size_t i; + for (i = 0; i < sizeof(safe) - 1 && varname[i] != '\0'; i++) + { + char c = varname[i]; + safe[i] = (c == '/' || c == '.') ? '_' : c; + } + safe[i] = '\0'; + + snprintf(buf, bufsz, "%s/%s.bin", dir, safe); +} + +/* Return the configured timeout in seconds (default 300). */ +static int handshake_timeout_s(void) +{ + const char *env = getenv("DDSTORE_HANDSHAKE_TIMEOUT_S"); + if (env) return atoi(env); + return 300; +} + +/* -------------------------------------------------------------------------- + * handshake_write() + * + * Called by each core rank after init_fabric() and fi_mr_reg(). + * Exchanges CoreRecords with all other core ranks via MPI_Allgather over + * `comm` — no filesystem round-trip needed, since core ranks already share + * a communicator. Builds this rank's address vector / remote key / remote + * address arrays from the gathered records. Rank 0 additionally publishes + * the combined array to {resolved_dir}/{varname}.bin (tmp + fsync + rename) + * so extra members can join later. + * -------------------------------------------------------------------------- */ +extern "C" int handshake_write(struct fabric_state *fs, MPI_Comm comm, + const char *dir, const char *varname, + int n_core, long nrows, int disp, int itemsize, + long *lenlist) +{ + int rank = 0; + MPI_Comm_rank(comm, &rank); + + /* Build this rank's CoreRecord. */ + struct CoreRecord rec; + memset(&rec, 0, sizeof(rec)); + + rec.fabric_address_len = DP_AV_DEF_SIZE; + int status = fi_getname((fid_t)fs->signal, + rec.fabric_address, &rec.fabric_address_len); + if (status != FI_SUCCESS) + { + fprintf(stderr, "[handshake_write] fi_getname failed: %s\n", + fi_strerror(status)); + return 1; + } + + rec.key = fs->key; + /* Same rule as handshake(): 0 for CXI's provider-key mode, the real + * pointer otherwise (always true for hsn/verbs/gni/psm2). */ + rec.base_address = is_virt_addr(fs) ? (uint64_t)(uintptr_t)fs->send_data : 0; + rec.nrows = nrows; + rec.disp = disp; + rec.itemsize = itemsize; + + /* Exchange records among core ranks directly over MPI. */ + std::vector all_recs(n_core); + MPI_Allgather(&rec, sizeof(rec), MPI_BYTE, + all_recs.data(), sizeof(rec), MPI_BYTE, comm); + + /* Build this rank's address vector / remote key / remote address + * arrays from the gathered records. */ + fs->comm_partner = (fi_addr_t *)malloc(n_core * sizeof(fi_addr_t)); + fs->remote_key = (uint64_t *)malloc(n_core * sizeof(uint64_t)); + fs->remote_address = (uint64_t *)malloc(n_core * sizeof(uint64_t)); + if (!fs->comm_partner || !fs->remote_key || !fs->remote_address) + { + fprintf(stderr, "[handshake_write] malloc failed\n"); + return 1; + } + for (int i = 0; i < n_core; i++) + { + int rc = fi_av_insert(fs->av, all_recs[i].fabric_address, 1, + &fs->comm_partner[i], 0, NULL); + if (rc != 1) + { + fprintf(stderr, + "[handshake_write] fi_av_insert failed for rank %d (rc=%d)\n", + i, rc); + return 1; + } + fs->remote_key[i] = all_recs[i].key; + fs->remote_address[i] = all_recs[i].base_address; + lenlist[i] = all_recs[i].nrows; + } + fs->world_size = n_core; + + /* Rank 0 publishes the combined record array for extra members. Write + * to a temp file first, fsync, then atomically rename into place so + * readers polling the final path never observe a partially-written or + * truncated-then-being-rewritten record. */ + if (rank == 0) + { + const char *rdir = resolve_handshake_dir(dir); + char bin_path[4096]; + char tmp_path[4096 + 32]; + record_path(bin_path, sizeof(bin_path), rdir, varname); + snprintf(tmp_path, sizeof(tmp_path), "%s.tmp.%d", bin_path, (int)getpid()); + + FILE *f = fopen(tmp_path, "wb"); + if (!f) + { + fprintf(stderr, "[handshake_write] cannot open %s: ", tmp_path); + perror(""); + return 1; + } + if (fwrite(all_recs.data(), sizeof(struct CoreRecord), n_core, f) != (size_t)n_core) + { + fprintf(stderr, "[handshake_write] fwrite failed for %s\n", tmp_path); + fclose(f); + unlink(tmp_path); + return 1; + } + if (fflush(f) != 0 || fsync(fileno(f)) != 0) + { + fprintf(stderr, "[handshake_write] fsync failed for %s: ", tmp_path); + perror(""); + fclose(f); + unlink(tmp_path); return 1; } + fclose(f); + + if (rename(tmp_path, bin_path) != 0) + { + fprintf(stderr, "[handshake_write] rename %s -> %s failed: ", tmp_path, bin_path); + perror(""); + unlink(tmp_path); + return 1; + } + + fprintf(stderr, "[handshake_write] wrote %s (%d records)\n", bin_path, n_core); + } + + return 0; +} + +/* -------------------------------------------------------------------------- + * handshake_join() + * + * Called by extra members (no MPI communicator, no data to publish). + * Polls for {resolved_dir}/{varname}.bin (the combined record file written + * by core rank 0 in handshake_write()) to reach its expected size, reads + * it, and populates: + * fs->comm_partner[0..n_core-1] (fi_addr_t, via fi_av_insert) + * fs->remote_key[0..n_core-1] + * fs->remote_address[0..n_core-1] + * lenlist[0..n_core-1] (raw nrows, NOT prefix-summed) + * *out_disp, *out_itemsize (from record 0; assumed uniform) + * + * Returns 0 on success, non-zero on error or timeout. + * -------------------------------------------------------------------------- */ +extern "C" int handshake_join(struct fabric_state *fs, + const char *dir, const char *varname, + int n_core, + long *lenlist, int *out_disp, int *out_itemsize) +{ + const char *rdir = resolve_handshake_dir(dir); + int timeout_s = handshake_timeout_s(); + struct timespec ts_start, ts_now; + clock_gettime(CLOCK_MONOTONIC, &ts_start); + + char path[4096]; + record_path(path, sizeof(path), rdir, varname); + size_t expected_size = (size_t)n_core * sizeof(struct CoreRecord); + + /* Poll until the combined record file is fully written. */ + for (;;) + { + struct stat st; + if (stat(path, &st) == 0 && st.st_size == (off_t)expected_size) + break; + + clock_gettime(CLOCK_MONOTONIC, &ts_now); + double elapsed = (ts_now.tv_sec - ts_start.tv_sec) + + (ts_now.tv_nsec - ts_start.tv_nsec) * 1e-9; + if (elapsed > timeout_s) + { + fprintf(stderr, + "[handshake_join] timeout after %.0f s waiting for " + "%s (var=%s, dir=%s)\n", + elapsed, path, varname, rdir); + return 1; + } + usleep(50000); /* 50 ms */ + } + + std::vector all_recs(n_core); + FILE *f = fopen(path, "rb"); + if (!f) + { + fprintf(stderr, "[handshake_join] cannot open %s: ", path); + perror(""); + return 1; } + if (fread(all_recs.data(), sizeof(struct CoreRecord), n_core, f) != (size_t)n_core) + { + fprintf(stderr, "[handshake_join] fread failed for %s\n", path); + fclose(f); + return 1; + } + fclose(f); + + fs->comm_partner = (fi_addr_t *)malloc(n_core * sizeof(fi_addr_t)); + fs->remote_key = (uint64_t *)malloc(n_core * sizeof(uint64_t)); + fs->remote_address = (uint64_t *)malloc(n_core * sizeof(uint64_t)); + if (!fs->comm_partner || !fs->remote_key || !fs->remote_address) + { + fprintf(stderr, "[handshake_join] malloc failed\n"); + return 1; + } + + for (int i = 0; i < n_core; i++) + { + int rc = fi_av_insert(fs->av, all_recs[i].fabric_address, 1, + &fs->comm_partner[i], 0, NULL); + if (rc != 1) + { + fprintf(stderr, + "[handshake_join] fi_av_insert failed for rank %d (rc=%d)\n", + i, rc); + return 1; + } + fs->remote_key[i] = all_recs[i].key; + fs->remote_address[i] = all_recs[i].base_address; + lenlist[i] = all_recs[i].nrows; + } + + if (out_disp) *out_disp = all_recs[0].disp; + if (out_itemsize) *out_itemsize = all_recs[0].itemsize; + fs->world_size = n_core; return 0; } \ No newline at end of file diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py new file mode 100644 index 0000000..4682b2c --- /dev/null +++ b/src/cpu_nic_map.py @@ -0,0 +1,405 @@ +#!/usr/bin/env python3 +"""CPU <-> nearest HSN (Slingshot) NIC topology, and FABRIC_IFACE auto-selection. + +Requires: module load hwloc (or lstopo/hwloc-calc on PATH) + +Three layers, one file: + - build_map()/serialize_env()/parse_env(): pure hwloc-calc topology query, + independent of any particular process's CPU affinity. + - 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. + +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. + +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 + + export DDSTORE_NIC_MAP=$(python3 cpu_nic_map.py --env) + srun --threads-per-core=2 -n8 -c14 python cpu_nic_map.py --allocated +""" + +import argparse +import fnmatch +import glob +import os +import re +import subprocess +import sys + + +def hcalc(loc, itype): + # --disallowed (must come first): include cores excluded by SLURM core + # specialization (e.g. cpu 0, 8, 16, ...) so the map covers all 128 + # CPUs, not just the ~112 currently allocatable to jobs -- a rank's + # affinity mask isn't guaranteed to avoid them (e.g. -S0 jobs). + # -p: report physical (OS/kernel) indices, matching os.sched_getaffinity() + # ids. Without it, hwloc-calc reports logical indices, which silently + # diverge from real CPU ids whenever core specialization or other + # exclusions shift the logical numbering (see lstopo's "P#" vs "L#"). + out = ( + subprocess.run( + ["hwloc-calc", "--disallowed", "-p", "-I", itype, loc], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + .stdout.decode() + .strip() + ) + return [int(x) for x in out.split(",")] if out else [] + + +def _sysfs_cpulist(name): + """Parse /sys/class/net//device/local_cpulist into a set of PU ids.""" + path = f"/sys/class/net/{name}/device/local_cpulist" + try: + text = open(path).read().strip() + except OSError: + return set() + pus = set() + for part in text.split(","): + part = part.strip() + if "-" in part: + lo, hi = part.split("-", 1) + pus.update(range(int(lo), int(hi) + 1)) + elif part: + pus.add(int(part)) + return pus + + +def _sysfs_numa(name): + """Read /sys/class/net//device/numa_node; returns int or None.""" + path = f"/sys/class/net/{name}/device/numa_node" + try: + val = int(open(path).read().strip()) + return val if val >= 0 else None + except (OSError, ValueError): + return None + + +def build_map(pattern): + nics = sorted( + os.path.basename(p) + for p in glob.glob("/sys/class/net/*") + if os.path.exists(os.path.join(p, "device")) + and fnmatch.fnmatch(os.path.basename(p), pattern) + ) + if not nics: + sys.exit(f"no NICs matching '{pattern}' found under /sys/class/net") + + nic_closest = {n: _sysfs_cpulist(n) for n in nics} + nic_numa = {n: _sysfs_numa(n) for n in nics} + if not any(nic_closest.values()): + sys.exit( + "sysfs local_cpulist came back empty for all NICs -- " + "check /sys/class/net/hsn*/device/local_cpulist" + ) + + # When multiple NICs share the same local_cpulist (e.g. 4 NICs per NUMA + # node on Aurora all report identical affinity), split the CPU set evenly + # so each NIC owns a unique partition and nearest() returns distinct NICs. + # Split each contiguous range separately so every NIC gets a share of both + # physical cores and hyperthreads (not just one or the other). + cpuset_to_nics = {} + for n, cpus in nic_closest.items(): + key = frozenset(cpus) + cpuset_to_nics.setdefault(key, []).append(n) + for key, group in cpuset_to_nics.items(): + if len(group) > 1: + sorted_cpus = sorted(key) + k = len(group) + # Find contiguous ranges within the shared CPU set + ranges, start, prev = [], sorted_cpus[0], sorted_cpus[0] + for c in sorted_cpus[1:]: + if c != prev + 1: + ranges.append(list(range(start, prev + 1))) + start = c + prev = c + ranges.append(list(range(start, prev + 1))) + # Assign each NIC a chunk from every range + partitions = [set() for _ in range(k)] + for seg in ranges: + chunk = (len(seg) + k - 1) // k + for i in range(k): + partitions[i].update(seg[i * chunk : (i + 1) * chunk]) + 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 + numa_to_nics = {} + for n, numa in nic_numa.items(): + numa_to_nics.setdefault(numa, []).append(n) + + all_pus = hcalc("all", "PU") + pu_numa = {} + for numa in numa_to_nics: + for pu in hcalc(f"NUMA:{numa}", "PU"): + pu_numa[pu] = numa + + def nearest(pu): + numa = pu_numa.get(pu) + exact_owner = next((n for n in nics if pu in nic_closest[n]), None) + if exact_owner: + return exact_owner, "yes", numa + candidates = numa_to_nics.get(numa, nics) + return candidates[0], "same-NUMA", numa + + return all_pus, nearest + + +def compress_ranges(values): + values = sorted(values) + out = [] + i = 0 + while i < len(values): + j = i + while j + 1 < len(values) and values[j + 1] == values[j] + 1: + j += 1 + out.append(str(values[i]) if i == j else f"{values[i]}-{values[j]}") + i = j + 1 + return ",".join(out) + + +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.""" + if provider == "cxi": + m = re.match(r"hsn(\d+)$", name) + if m: + return f"cxi{m.group(1)}" + return name + + +def serialize_env(pattern="hsn*", provider="hsn"): + all_pus, nearest = build_map(pattern) + by_nic = {} + for pu in all_pus: + nic = translate_iface(nearest(pu)[0], provider) + by_nic.setdefault(nic, []).append(pu) + return ";".join( + f"{nic}={compress_ranges(pus)}" for nic, pus in sorted(by_nic.items()) + ) + + +def parse_env(s): + cpu_to_nic = {} + for segment in s.split(";"): + nic, ranges = segment.split("=", 1) + for r in ranges.split(","): + if "-" in r: + lo, hi = r.split("-") + cpu_to_nic.update((cpu, nic) for cpu in range(int(lo), int(hi) + 1)) + else: + cpu_to_nic[int(r)] = nic + return cpu_to_nic + + +def allocated_nics(pattern="hsn*", nic_map=None): + """Which NIC(s) this process's actual CPU affinity is nearest to. + + Uses os.sched_getaffinity(0) to get the real CPU set the process is + bound to (respects SLURM --cpus-per-task/--cpu-bind and cgroups). + + nic_map: an explicit precomputed map string (see serialize_env()/--env + format), taking priority over the DDSTORE_NIC_MAP env var. Pass this + when the caller already has the map from somewhere other than the + process environment. Falls back to a live hwloc-calc query when neither + is available. + + Prefer a precomputed map: hwloc-calc's NIC/PCI visibility has been + observed to fail silently inside some srun tasks, so computing the map + once (e.g. in the sbatch batch step's own shell, where NIC visibility is + reliable) and sharing it -- via DDSTORE_NIC_MAP or the nic_map argument + -- is more robust than querying hwloc fresh in every rank. + """ + allocated = sorted(os.sched_getaffinity(0)) + map_str = nic_map if nic_map is not None else os.environ.get("DDSTORE_NIC_MAP") + if map_str: + cpu_to_nic = parse_env(map_str) + nics = {cpu_to_nic[c] for c in allocated if c in cpu_to_nic} + else: + all_pus, nearest = build_map(pattern) + nics = {nearest(c)[0] for c in allocated if c in all_pus} + fabric = os.environ.get("DDSTORE_FABRIC", "hsn") + by_nic = {} + for pu in all_pus: + nic = translate_iface(nearest(pu)[0], fabric) + by_nic.setdefault(nic, []).append(pu) + map_str = ";".join( + f"{nic}={compress_ranges(pus)}" for nic, pus in sorted(by_nic.items()) + ) + print( + "Tip: cache this map for future runs to avoid recomputing it per rank:\n" + f" export DDSTORE_NIC_MAP={map_str}", + file=sys.stderr, + ) + return allocated, nics + + +def select_fabric_iface(nic_map=None): + """Pick the libfabric NIC (FABRIC_IFACE) for this process, if not + already set. Defers to allocated_nics(), which uses nic_map when given, + 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 + method=1/2. + + nic_map: an explicit precomputed map string (serialize_env()/--env + format), for callers that already have the map from somewhere other + 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 + 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). + """ + use_cxi = os.environ.get("DDSTORE_FABRIC") == "cxi" + + if "FABRIC_IFACE" in os.environ: + iface = os.environ["FABRIC_IFACE"] + if use_cxi: + translated = translate_iface(iface, "cxi") + if translated != iface: + iface = translated + os.environ["FABRIC_IFACE"] = iface + return iface + + allocated, nics = allocated_nics(nic_map=nic_map) + if not nics: + if use_cxi: + cxi_domains = sorted( + os.path.basename(p) + for p in glob.glob("/sys/class/net/hsn*") + if re.match(r"hsn\d+$", os.path.basename(p)) + ) + if not cxi_domains: + raise RuntimeError( + f"could not determine a nearest HSN NIC for this rank's " + f"CPU affinity {allocated} and no hsn* devices found; " + f"set FABRIC_IFACE explicitly to work around this" + ) + local_rank = int(os.environ.get("SLURM_LOCALID", 0)) + hsn = cxi_domains[local_rank % len(cxi_domains)] + iface = translate_iface(hsn, "cxi") + print(f"FABRIC_IFACE: fallback (SLURM_LOCALID={local_rank}) -> {iface}") + os.environ["FABRIC_IFACE"] = iface + return iface + raise RuntimeError( + f"could not determine a nearest HSN NIC for this rank's CPU " + f"affinity {allocated}; set FABRIC_IFACE explicitly to work " + f"around this" + ) + iface = sorted(nics)[0] + if use_cxi: + iface = translate_iface(iface, "cxi") + if len(nics) > 1: + translated = sorted( + translate_iface(n, "cxi" if use_cxi else "hsn") for n in nics + ) + print(f"FABRIC_IFACE: affinity spans {translated}, picking {iface}") + + os.environ["FABRIC_IFACE"] = iface + return iface + + +def main(): + parser = argparse.ArgumentParser( + description="Find the nearest HSN (Slingshot) NIC for a given CPU (PU) id, " + "based on hwloc PCI/NUMA locality.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "examples:\n" + " cpu_nic_map.py print the full CPU -> nearest HSN NIC table\n" + " 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" + ), + ) + parser.add_argument( + "cpu", + nargs="?", + type=int, + help="CPU (PU) id to look up; omit to print the full table", + ) + parser.add_argument( + "--fabric", + default=os.environ.get("DDSTORE_FABRIC", "hsn"), + 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)", + ) + parser.add_argument( + "--env", + action="store_true", + help="print only the compact DDSTORE_NIC_MAP env-var value", + ) + parser.add_argument( + "--allocated", + action="store_true", + help="print this process's allocated CPUs (os.sched_getaffinity) " + "and their nearest NIC(s), instead of the full table", + ) + args = parser.parse_args() + + if args.env: + print(serialize_env(provider=args.fabric)) + return + + if args.allocated: + allocated, nics = allocated_nics() + nics = {translate_iface(n, args.fabric) for n in nics} + print(f"allocated CPUs: {allocated}") + print(f"nearest NIC(s): {sorted(nics)}") + return + + all_pus, nearest = build_map("hsn*") + + if args.cpu is not None: + if args.cpu not in all_pus: + sys.exit( + f"cpu id {args.cpu} not found (valid range: {min(all_pus)}-{max(all_pus)})" + ) + owner, exact, numa = nearest(args.cpu) + print(translate_iface(owner, args.fabric)) + return + + print(f"{'CPU':>4} {'NUMA':>4} {'nearest NIC':>11} exact") + for pu in all_pus: + owner, exact, numa = nearest(pu) + owner = translate_iface(owner, args.fabric) + print(f"{pu:>4} {numa!s:>4} {owner:>11} {exact}") + by_nic = {} + for pu in all_pus: + nic = translate_iface(nearest(pu)[0], args.fabric) + by_nic.setdefault(nic, []).append(pu) + map_str = ";".join( + f"{nic}={compress_ranges(pus)}" for nic, pus in sorted(by_nic.items()) + ) + print( + "\nTip: cache this map for future runs:\n" + f" export DDSTORE_NIC_MAP={map_str}", + file=sys.stderr, + ) + + +if __name__ == "__main__": + main() diff --git a/src/ddstore.cxx b/src/ddstore.cxx index db4bd86..5e91ae2 100644 --- a/src/ddstore.cxx +++ b/src/ddstore.cxx @@ -6,6 +6,15 @@ #include #include +/* Forward-declare the C helper from common.cxx so we can call it here. */ +extern "C" const char *resolve_handshake_dir(const char *user_dir); + +/* Convenience wrapper: resolve and return as std::string. */ +static std::string resolve_dir(const std::string &user_dir) +{ + return std::string(resolve_handshake_dir(user_dir.c_str())); +} + int sortedsearch(const std::vector &vec, long num) { if (vec.empty() || num < 0 || num >= vec.back()) @@ -17,21 +26,21 @@ int sortedsearch(const std::vector &vec, long num) std::upper_bound(vec.begin(), vec.end(), num)); } -DDStore::DDStore() : method(0) +DDStore::DDStore() : method(0), comm_size(1), rank(0), n_core(0), is_extra(false) { this->comm = MPI_COMM_SELF; MPI_Comm_size(this->comm, &this->comm_size); MPI_Comm_rank(this->comm, &this->rank); } -DDStore::DDStore(MPI_Comm comm) : method(0) +DDStore::DDStore(MPI_Comm comm) : method(0), n_core(0), is_extra(false) { this->comm = comm; MPI_Comm_size(this->comm, &this->comm_size); MPI_Comm_rank(this->comm, &this->rank); } -DDStore::DDStore(int method, MPI_Comm comm) +DDStore::DDStore(int method, MPI_Comm comm) : n_core(0), is_extra(false) { this->method = method; this->comm = comm; @@ -39,6 +48,33 @@ DDStore::DDStore(int method, MPI_Comm comm) MPI_Comm_rank(this->comm, &this->rank); } +/* Method 2: core member constructor. n_core is derived from comm_size — + * every rank in `comm` is assumed to be a core member. */ +DDStore::DDStore(int method, MPI_Comm comm, + const std::string &handshake_dir) + : is_extra(false) +{ + this->method = method; + this->comm = comm; + this->handshake_dir = resolve_dir(handshake_dir); + MPI_Comm_size(this->comm, &this->comm_size); + MPI_Comm_rank(this->comm, &this->rank); + this->n_core = this->comm_size; + fprintf(stderr, "[DDStore] method=2 core: handshake_dir=%s, n_core=%d\n", + this->handshake_dir.c_str(), this->n_core); +} + +/* Method 2: extra member constructor (no MPI communicator needed). */ +DDStore::DDStore(int method, const std::string &handshake_dir, int n_core) + : comm(MPI_COMM_SELF), comm_size(1), rank(0), is_extra(true) +{ + this->method = method; + this->handshake_dir = resolve_dir(handshake_dir); + this->n_core = n_core; + fprintf(stderr, "[DDStore] method=2 extra: handshake_dir=%s\n", + this->handshake_dir.c_str()); +} + DDStore::~DDStore() { this->free(); @@ -61,6 +97,7 @@ void DDStore::epoch_begin() x.second.fence_active = true; } } + /* Methods 1 and 2 use libfabric — no fence needed. */ } void DDStore::epoch_end() @@ -75,8 +112,72 @@ void DDStore::epoch_end() x.second.fence_active = false; } } + /* Methods 1 and 2 use libfabric — no fence needed. */ +} + +/* -------------------------------------------------------------------------- + * 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. + * -------------------------------------------------------------------------- */ +void DDStore::join(std::string name) +{ + if (!this->is_extra) + throw std::logic_error("join() is only valid for extra members"); + if (this->method != 2) + throw std::logic_error("join() requires method=2"); + + struct fabric_state *fs = + (struct fabric_state *)calloc(1, sizeof(struct fabric_state)); + fs->world_size = this->n_core; + fs->rank = -1; /* extra members have no core rank */ + + init_fabric(fs); + 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. */ + fs->send_data = NULL; + fs->send_data_len = 0; + fs->mr = NULL; + fs->key = 0; + + std::vector raw_lens(this->n_core); + int out_disp = 0, out_itemsize = 0; + if (handshake_join(fs, + this->handshake_dir.c_str(), name.c_str(), + this->n_core, + raw_lens.data(), &out_disp, &out_itemsize) != 0) + throw std::runtime_error("handshake_join failed for variable: " + name); + + /* Build prefix-sum lenlist. */ + long sum = 0; + std::vector lenlist(this->n_core); + for (int i = 0; i < this->n_core; i++) + { + sum += raw_lens[i]; + lenlist[i] = sum; + } + + VarInfo_t var; + var.name = name; + var.itemsize = out_itemsize; + var.disp = out_disp; + var.win = MPI_WIN_NULL; + var.lenlist = lenlist; + var.active = true; + var.fence_active = false; + var.base = NULL; /* extra member owns no data */ + var.fabric_state = fs; + this->varlist.insert(std::pair(name, var)); } +/* -------------------------------------------------------------------------- + * free() — release all resources. + * -------------------------------------------------------------------------- */ void DDStore::free() { int flag; @@ -88,28 +189,26 @@ void DDStore::free() if (x.second.active) { MPI_Win_free(&x.second.win); - // (2024/12) no need as using the user pointer - // MPI_Free_mem(x.second.base); } x.second.active = false; } } - else if (this->method == 1) + else if (this->method == 1 || this->method == 2) { for (auto &x : this->varlist) { 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); + 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); diff --git a/src/pyddstore.pyx b/src/pyddstore.pyx index c09046f..f85b976 100644 --- a/src/pyddstore.pyx +++ b/src/pyddstore.pyx @@ -6,6 +6,8 @@ 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 @@ -37,6 +39,11 @@ cdef extern from "ddstore.hpp": 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() @@ -44,6 +51,9 @@ cdef extern from "ddstore.hpp": 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 @@ -52,12 +62,61 @@ cdef class PyDDstoreVarinfo: pass cdef class PyDDStore: - cdef DDStore c_ddstore + 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 __cinit__(self, MPI.Comm comm, int method = 0): - # print("PyDDStore init method:", method) - self.c_ddstore = DDStore(method, comm.ob_mpi) - def add(self, str name, np.ndarray arr): assert arr.flags.c_contiguous cdef long nrows = arr.shape[0] @@ -125,3 +184,14 @@ cdef class PyDDStore: 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/test/conftest.py b/test/conftest.py index 6a592d4..09ff14a 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,4 +1,5 @@ import mpi4py + mpi4py.rc.thread_level = "serialized" mpi4py.rc.threads = False diff --git a/test/test_method2_core.py b/test/test_method2_core.py new file mode 100644 index 0000000..94324f3 --- /dev/null +++ b/test/test_method2_core.py @@ -0,0 +1,143 @@ +""" +Method 2 — core member side. + +Run together with test_method2_extra.py in two separate mpirun launches +that share a handshake directory on a shared filesystem (e.g. Lustre). + +Usage (two terminals / two job steps): + + Terminal 1 (core): + mpirun -n 4 python test/test_method2_core.py [handshake_dir] [n_core] + + Terminal 2 (extra): + python test/test_method2_extra.py [handshake_dir] [n_core] + +Directory resolution priority (same as the C library): + 1. Command-line argument + 2. DDSTORE_HANDSHAKE_DIR environment variable + 3. ./ddstore_hs (current working directory — must be on shared filesystem) + +The handshake directory MUST be on a shared filesystem visible to all nodes. + +Environment: + DDSTORE_HANDSHAKE_DIR — shared Lustre path for handshake files + DDSTORE_HANDSHAKE_TIMEOUT_S — poll timeout in seconds (default: 300) +""" + +import os +import sys +import time +import numpy as np +from mpi4py import MPI +import pyddstore as dds + +# --------------------------------------------------------------------------- +# configuration — resolve directory with same priority as the C library +# --------------------------------------------------------------------------- + + +def _resolve_dir(arg): + if arg: + return arg + env = os.environ.get("DDSTORE_HANDSHAKE_DIR", "") + if env: + return env + return "./ddstore_hs" + + +hs_dir = _resolve_dir(sys.argv[1] if len(sys.argv) > 1 else "") +n_core = ( + int(sys.argv[2]) + if len(sys.argv) > 2 + else int(os.environ.get("DDSTORE_N_CORE", "4")) +) + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +size = comm.Get_size() + +assert size == n_core, f"Expected {n_core} core ranks, got {size}" + +# --------------------------------------------------------------------------- +# prepare handshake directory (rank 0 creates it and removes stale files) +# --------------------------------------------------------------------------- + +if rank == 0: + os.makedirs(hs_dir, exist_ok=True) + for fname in os.listdir(hs_dir): + if fname.endswith(".bin"): + os.remove(os.path.join(hs_dir, fname)) + print(f"[core] handshake_dir={hs_dir}", flush=True) +comm.Barrier() + +# --------------------------------------------------------------------------- +# build per-rank data: each rank holds `nrows` rows of `ncols` float32s +# --------------------------------------------------------------------------- + +nrows = 8 +ncols = 4 +data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) + +# --------------------------------------------------------------------------- +# create DDStore (method=2, core member) +# --------------------------------------------------------------------------- + +store = dds.PyDDStore(comm, method=2, handshake_dir=hs_dir) +store.add("x", data) + +if rank == 0: + print(f"[core] all {n_core} records written to {hs_dir}", flush=True) + +# --------------------------------------------------------------------------- +# self-check: core members can also read from each other via method=2 +# --------------------------------------------------------------------------- + +out = np.zeros((1, ncols), dtype=np.float32) +for target_rank in range(n_core): + global_idx = target_rank * nrows + store.get("x", out, start=global_idx) + expected = float(target_rank + 1) + assert np.all(out == expected), ( + f"[core rank {rank}] get from target {target_rank}: " + f"expected {expected}, got {out}" + ) + +if rank == 0: + print("[core] self-check passed", flush=True) + +# --------------------------------------------------------------------------- +# wait for extra member to signal it is done +# --------------------------------------------------------------------------- + +sentinel = os.path.join(hs_dir, "done_extra") +timeout_s = int(os.environ.get("DDSTORE_HANDSHAKE_TIMEOUT_S", "300")) +t0 = time.monotonic() + +if rank == 0: + print("[core] waiting for extra member sentinel ...", flush=True) + while not os.path.exists(sentinel): + if time.monotonic() - t0 > timeout_s: + raise TimeoutError("timed out waiting for done_extra sentinel") + time.sleep(0.2) + print("[core] sentinel received, shutting down", flush=True) +comm.Barrier() + +# --------------------------------------------------------------------------- +# cleanup +# --------------------------------------------------------------------------- + +store.free() + +if rank == 0: + for fname in os.listdir(hs_dir): + if fname.endswith(".bin"): + try: + os.remove(os.path.join(hs_dir, fname)) + except OSError: + pass + try: + os.remove(sentinel) + except OSError: + pass + +print(f"[core rank {rank}] done", flush=True) diff --git a/test/test_method2_extra.py b/test/test_method2_extra.py new file mode 100644 index 0000000..c6c0937 --- /dev/null +++ b/test/test_method2_extra.py @@ -0,0 +1,98 @@ +""" +Method 2 — extra member side (read-only, no MPI communicator needed). + +Run after (or concurrently with) test_method2_core.py: + + python test/test_method2_extra.py [handshake_dir] [n_core] + +Directory resolution priority (same as the C library): + 1. Command-line argument + 2. DDSTORE_HANDSHAKE_DIR environment variable + 3. ./ddstore_hs (current working directory — must be on shared filesystem) + +The handshake directory MUST be on a shared filesystem visible to all nodes. + +Environment: + DDSTORE_HANDSHAKE_DIR — shared Lustre path for handshake files + DDSTORE_N_CORE — number of core ranks (default: 4) + DDSTORE_HANDSHAKE_TIMEOUT_S — poll timeout in seconds (default: 300) +""" + +import os +import sys +import numpy as np +import pyddstore as dds + +# --------------------------------------------------------------------------- +# configuration — resolve directory with same priority as the C library +# --------------------------------------------------------------------------- + + +def _resolve_dir(arg): + if arg: + return arg + env = os.environ.get("DDSTORE_HANDSHAKE_DIR", "") + if env: + return env + return "./ddstore_hs" + + +hs_dir = _resolve_dir(sys.argv[1] if len(sys.argv) > 1 else "") +n_core = ( + int(sys.argv[2]) + if len(sys.argv) > 2 + else int(os.environ.get("DDSTORE_N_CORE", "4")) +) + +nrows = 8 +ncols = 4 + +# --------------------------------------------------------------------------- +# create DDStore as extra member (no MPI comm) +# --------------------------------------------------------------------------- + +print(f"[extra] joining store at {hs_dir}, n_core={n_core}", flush=True) +store = dds.PyDDStore(None, method=2, handshake_dir=hs_dir, n_core=n_core) + +# discover variable "x" published by core members +store.join("x") +print("[extra] join('x') succeeded", flush=True) + +# --------------------------------------------------------------------------- +# read every row from every core rank and verify +# --------------------------------------------------------------------------- + +out = np.zeros((1, ncols), dtype=np.float32) +errors = [] + +for target_rank in range(n_core): + for row in range(nrows): + global_idx = target_rank * nrows + row + store.get("x", out, start=global_idx) + expected = float(target_rank + 1) + if not np.all(out == expected): + errors.append( + f" global_idx={global_idx} (rank={target_rank} row={row}): " + f"expected {expected}, got {out[0]}" + ) + +if errors: + print("[extra] FAILED — mismatches:", flush=True) + for e in errors: + print(e, flush=True) + raise AssertionError(f"{len(errors)} mismatches") + +print(f"[extra] all {n_core * nrows} reads verified correctly", flush=True) + +# --------------------------------------------------------------------------- +# cleanup +# --------------------------------------------------------------------------- + +store.free() + +# signal core side that we are done +sentinel = os.path.join(hs_dir, "done_extra") +with open(sentinel, "w") as f: + f.write("done\n") + +print("[extra] sentinel written, exiting", flush=True) diff --git a/test/test_multirank.py b/test/test_multirank.py index 25a2924..ad89109 100644 --- a/test/test_multirank.py +++ b/test/test_multirank.py @@ -4,6 +4,7 @@ Each rank stores a distinct value; tests verify cross-rank remote reads. Requires at least 2 ranks; some tests require exactly 4. """ + import numpy as np import pytest from mpi4py import MPI @@ -19,6 +20,7 @@ def all_passed(comm, local_ok): # remote get: each rank reads from every other rank # --------------------------------------------------------------------------- + def test_remote_get_all_ranks(comm): rank = comm.Get_rank() size = comm.Get_size() @@ -47,6 +49,7 @@ def test_remote_get_all_ranks(comm): # boundary: last row of each rank's shard # --------------------------------------------------------------------------- + def test_remote_get_last_row(comm): rank = comm.Get_rank() size = comm.Get_size() @@ -65,8 +68,10 @@ def test_remote_get_last_row(comm): store.get("x", out, start=last_global_idx) expected = data[nrows - 1] + (target_rank - rank) * 1000 # recompute expected from target rank's perspective - expected = np.arange((nrows - 1) * ncols, nrows * ncols, - dtype=np.float32) + target_rank * 1000 + expected = ( + np.arange((nrows - 1) * ncols, nrows * ncols, dtype=np.float32) + + target_rank * 1000 + ) if not np.allclose(out[0], expected): local_ok = False store.epoch_end() @@ -79,6 +84,7 @@ def test_remote_get_last_row(comm): # init + update on each rank, then remote get # --------------------------------------------------------------------------- + def test_init_update_remote_get(comm): rank = comm.Get_rank() size = comm.Get_size() @@ -108,6 +114,7 @@ def test_init_update_remote_get(comm): # multiple variables — independent correctness # --------------------------------------------------------------------------- + def test_multiple_variables_remote(comm): rank = comm.Get_rank() size = comm.Get_size() @@ -140,6 +147,7 @@ def test_multiple_variables_remote(comm): # ddstore_width: sub-communicator grouping (requires size >= 4) # --------------------------------------------------------------------------- + def test_ddstore_width(comm): size = comm.Get_size() if size < 4: diff --git a/test/test_single.py b/test/test_single.py index f40563c..a29440f 100644 --- a/test/test_single.py +++ b/test/test_single.py @@ -3,17 +3,18 @@ No remote memory access; all data is local. """ + import numpy as np import pytest from mpi4py import MPI import pyddstore as dds - # --------------------------------------------------------------------------- # helpers # --------------------------------------------------------------------------- + def make_store(comm, method=0): return dds.PyDDStore(comm, method=method) @@ -44,9 +45,18 @@ def roundtrip(store, name, data): # dtype coverage # --------------------------------------------------------------------------- -@pytest.mark.parametrize("dtype", [ - np.int32, np.int64, np.uint8, np.float32, np.float64, np.bool_, -]) + +@pytest.mark.parametrize( + "dtype", + [ + np.int32, + np.int64, + np.uint8, + np.float32, + np.float64, + np.bool_, + ], +) def test_add_get_dtypes(comm, dtype): store = make_store(comm) data = make_data(dtype) @@ -59,6 +69,7 @@ def test_add_get_dtypes(comm, dtype): # init / update / get # --------------------------------------------------------------------------- + def test_init_update_get(comm): store = make_store(comm) nrows, ncols = 4, 8 @@ -100,6 +111,7 @@ def test_update_partial(comm): # multiple variables # --------------------------------------------------------------------------- + def test_multiple_variables(comm): store = make_store(comm) a = np.ones((4, 4), dtype=np.float32) @@ -121,6 +133,7 @@ def test_multiple_variables(comm): # error cases # --------------------------------------------------------------------------- + def test_get_out_of_range(comm): store = make_store(comm) data = np.ones((4, 4), dtype=np.float32) @@ -177,6 +190,7 @@ def test_get_wrong_dtype(comm): # double free safety # --------------------------------------------------------------------------- + def test_double_free(comm): store = make_store(comm) data = np.ones((4, 4), dtype=np.float32)