From 8f48c7d2c34429788717a11a0a69cf09f54a63f7 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Sat, 22 Aug 2026 11:06:18 -0400 Subject: [PATCH 01/23] add file-based method --- README.md | 95 ++++++++-- examples/vae/distdataset.py | 80 ++++++++- examples/vae/vae-ddp.py | 300 ++++++++++---------------------- examples/vae/vae_core_server.py | 96 ++++++++++ examples/vae/vae_extra_train.py | 200 +++++++++++++++++++++ examples/vae/vae_model.py | 45 +++++ include/common.h | 45 +++++ include/ddstore.hpp | 178 ++++++++++++++++--- src/common.cxx | 271 +++++++++++++++++++++++++++++ src/ddstore.cxx | 129 ++++++++++++-- src/pyddstore.pyx | 66 ++++++- test/test_method2_core.py | 137 +++++++++++++++ test/test_method2_extra.py | 92 ++++++++++ 13 files changed, 1469 insertions(+), 265 deletions(-) create mode 100644 examples/vae/vae_core_server.py create mode 100644 examples/vae/vae_extra_train.py create mode 100644 examples/vae/vae_model.py create mode 100644 test/test_method2_core.py create mode 100644 test/test_method2_extra.py diff --git a/README.md b/README.md index 536a832..2ba698b 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,25 @@ 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)` | 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 | + +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 +141,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 every core rank's record file appears (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`. @@ -154,11 +182,44 @@ Set `FABRIC_IFACE` to select a specific network interface when the automatic sel export FABRIC_IFACE=hsn0 # e.g. Cray Slingshot ``` +### 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()`. Each rank writes a `{name}_rank{N}.bin` record (fabric address, MR key, base pointer, row count, dtype) 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 | + +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 +227,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 +290,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/vae/distdataset.py b/examples/vae/distdataset.py index 4295169..84c3900 100644 --- a/examples/vae/distdataset.py +++ b/examples/vae/distdataset.py @@ -33,12 +33,25 @@ 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")) + gpu_id = int(os.getenv("SLURM_LOCALID", "0")) 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 under the + # same {varname}_rank{N}.bin filenames, 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 + ) ## set total before set subset self.total_ns = len(data) @@ -63,7 +76,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 +95,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 +110,55 @@ 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 + ) + 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..bc2b93d 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -1,163 +1,58 @@ -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 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 +import mpi4py - 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 +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 +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() @@ -171,37 +66,9 @@ def setup_ddp(): 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 - - 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) +print("DDP setup:", comm_size, rank, device) model = VAE().to(device) model = torch.nn.parallel.DistributedDataParallel(model) @@ -211,27 +78,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 +120,24 @@ 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))) + print( + "====> Epoch: {} Average loss: {:.4f}".format( + epoch, train_loss / len(train_loader.dataset) + ) + ) def test(epoch): @@ -277,13 +150,18 @@ 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) @@ -293,6 +171,8 @@ def 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') + 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..512a081 --- /dev/null +++ b/examples/vae/vae_core_server.py @@ -0,0 +1,96 @@ +""" +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) +""" + +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..ef6c24c --- /dev/null +++ b/examples/vae/vae_extra_train.py @@ -0,0 +1,200 @@ +""" +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) +""" + +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 +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) + +if args.cuda: + device = torch.device("cuda") +elif use_mps: + device = torch.device("mps") +else: + device = torch.device("cpu") + +comm_size, rank = setup_ddp() +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..37f8028 100644 --- a/include/common.h +++ b/include/common.h @@ -9,6 +9,22 @@ #define DP_AV_DEF_SIZE 512 #define COMM_FILE_WRITER_TO_READER "./writer_address.bin" +/* ----------------------------------------------------------------------- + * Method 2: file-based handshake record written by each core rank. + * One file per variable per core rank: + * {handshake_dir}/{varname}_rank{rank}.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" { @@ -48,6 +64,35 @@ extern "C" 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: write own CoreRecord to {dir}/{varname}_rank{rank}.bin. + * Uses write-to-tmp + rename for atomicity. */ + int handshake_write(struct fabric_state *fs, + const char *dir, const char *varname, int rank, + long nrows, int disp, int itemsize); + + /* Core rank: poll until all n_core record files are present, then read + * them all and populate fs->comm_partner[], remote_key[], + * remote_address[], and fill lenlist[0..n_core-1] (raw row counts, + * NOT yet prefix-summed). Also sets *out_disp and *out_itemsize. */ + int handshake_read(struct fabric_state *fs, + const char *dir, const char *varname, + int n_core, int my_rank, + long *lenlist, int *out_disp, int *out_itemsize); + + /* Extra member: same as handshake_read but does NOT write anything. + * Blocks until all n_core files appear (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..ae3a831 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) { @@ -73,6 +100,68 @@ class DDStore throw std::runtime_error("init_fabric failed: no suitable fabric found"); handshake(fabric_state, this->comm); } + 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)); + fabric_state->key = fi_mr_key(fabric_state->mr); + + /* Write this rank's record to the handshake directory. */ + if (handshake_write(fabric_state, + this->handshake_dir.c_str(), name.c_str(), this->rank, + nrows, disp, (int)sizeof(T)) != 0) + throw std::runtime_error("handshake_write failed"); + + /* Wait for all core ranks, then read all records. */ + std::vector raw_lens(this->n_core); + int file_disp = 0, file_itemsize = 0; + if (handshake_read(fabric_state, + this->handshake_dir.c_str(), name.c_str(), + this->n_core, this->rank, + raw_lens.data(), &file_disp, &file_itemsize) != 0) + throw std::runtime_error("handshake_read 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); MPI_Allgather(&nrows, 1, MPI_LONG, lenlist.data(), 1, MPI_LONG, this->comm); @@ -89,11 +178,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; @@ -146,6 +230,64 @@ class DDStore throw std::runtime_error("init_fabric failed: no suitable fabric found"); handshake(fabric_state, this->comm); } + 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)); + fabric_state->key = fi_mr_key(fabric_state->mr); + + if (handshake_write(fabric_state, + this->handshake_dir.c_str(), name.c_str(), this->rank, + nrows, disp, itemsize) != 0) + throw std::runtime_error("handshake_write failed"); + + std::vector raw_lens(this->n_core); + int file_disp = 0, file_itemsize = 0; + if (handshake_read(fabric_state, + this->handshake_dir.c_str(), name.c_str(), + this->n_core, this->rank, + raw_lens.data(), &file_disp, &file_itemsize) != 0) + throw std::runtime_error("handshake_read 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); MPI_Allgather(&nrows, 1, MPI_LONG, lenlist.data(), 1, MPI_LONG, this->comm); @@ -162,11 +304,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 +330,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,11 +375,9 @@ 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); @@ -252,11 +385,16 @@ class DDStore } 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/src/common.cxx b/src/common.cxx index 7935496..2ab4997 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -9,6 +9,10 @@ #include #include #include +#include +#include +#include +#include void init_fabric(struct fabric_state *fabric) { @@ -386,4 +390,271 @@ int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset } return 0; +} + +/* ========================================================================= + * Method 2: file-based handshake helpers + * ========================================================================= + * + * File naming convention: + * {dir}/{varname}_rank{rank}.bin — one CoreRecord per core rank + * + * 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. */ +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 rank i's record file into `buf`. */ +static void record_path(char *buf, size_t bufsz, + const char *dir, const char *varname, int rank) +{ + snprintf(buf, bufsz, "%s/%s_rank%d.bin", dir, varname, rank); +} + +/* 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(). + * Writes a CoreRecord to {resolved_dir}/{varname}_rank{rank}.bin. + * The directory is resolved via resolve_handshake_dir(). + * -------------------------------------------------------------------------- */ +int handshake_write(struct fabric_state *fs, + const char *dir, const char *varname, int rank, + long nrows, int disp, int itemsize) +{ + const char *rdir = resolve_handshake_dir(dir); + + /* Build the CoreRecord for this rank. */ + 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; + rec.base_address = (uint64_t)(uintptr_t)fs->send_data; + rec.nrows = nrows; + rec.disp = disp; + rec.itemsize = itemsize; + + /* 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. */ + char bin_path[4096]; + char tmp_path[4096 + 32]; + record_path(bin_path, sizeof(bin_path), rdir, varname, rank); + 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(&rec, sizeof(rec), 1, f) != 1) + { + 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] rank %d wrote %s\n", rank, bin_path); + return 0; +} + +/* -------------------------------------------------------------------------- + * handshake_read() + * + * Called by each core rank after handshake_write(). + * Polls until all n_core record files are present in the resolved directory, + * then reads them 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 rank-0 record; assumed uniform) + * + * Returns 0 on success, non-zero on error or timeout. + * -------------------------------------------------------------------------- */ +int handshake_read(struct fabric_state *fs, + const char *dir, const char *varname, + int n_core, int my_rank, + 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); + + /* Allocate arrays sized for n_core peers. */ + 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_read] malloc failed\n"); + return 1; + } + + /* Poll until all n_core files are fully written (size == sizeof CoreRecord). */ + for (;;) + { + int ready = 0; + for (int i = 0; i < n_core; i++) + { + char path[4096]; + record_path(path, sizeof(path), rdir, varname, i); + struct stat st; + if (stat(path, &st) == 0 && + st.st_size == (off_t)sizeof(struct CoreRecord)) + ready++; + } + if (ready == n_core) + 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_read] timeout after %.0f s waiting for " + "%d/%d core records (var=%s, dir=%s)\n", + elapsed, ready, n_core, varname, rdir); + return 1; + } + usleep(50000); /* 50 ms */ + } + + /* Read all records. */ + for (int i = 0; i < n_core; i++) + { + char path[4096]; + record_path(path, sizeof(path), rdir, varname, i); + + FILE *f = fopen(path, "rb"); + if (!f) + { + fprintf(stderr, "[handshake_read] cannot open %s: ", path); + perror(""); + return 1; + } + + struct CoreRecord rec; + if (fread(&rec, sizeof(rec), 1, f) != 1) + { + fprintf(stderr, "[handshake_read] fread failed for %s\n", path); + fclose(f); + return 1; + } + fclose(f); + + /* Insert fabric address into the address vector. */ + int rc = fi_av_insert(fs->av, + rec.fabric_address, 1, + &fs->comm_partner[i], 0, NULL); + if (rc != 1) + { + fprintf(stderr, + "[handshake_read] fi_av_insert failed for rank %d (rc=%d)\n", + i, rc); + return 1; + } + + fs->remote_key[i] = rec.key; + fs->remote_address[i] = rec.base_address; + lenlist[i] = rec.nrows; + + if (i == 0) + { + if (out_disp) *out_disp = rec.disp; + if (out_itemsize) *out_itemsize = rec.itemsize; + } + } + + fs->world_size = n_core; + (void)my_rank; + return 0; +} + +/* -------------------------------------------------------------------------- + * handshake_join() + * + * Called by extra members (no MPI, no data to publish). + * Identical to handshake_read() except it never writes anything. + * Blocks until all n_core record files are present (or 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) +{ + return handshake_read(fs, dir, varname, n_core, -1, + lenlist, out_disp, out_itemsize); } \ No newline at end of file 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..b755efa 100644 --- a/src/pyddstore.pyx +++ b/src/pyddstore.pyx @@ -37,6 +37,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 +49,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 +60,49 @@ 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): + """ + 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) + """ + cdef MPI.Comm mpi_comm + 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 + 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 +170,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/test_method2_core.py b/test/test_method2_core.py new file mode 100644 index 0000000..fdb1eeb --- /dev/null +++ b/test/test_method2_core.py @@ -0,0 +1,137 @@ +""" +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..530d476 --- /dev/null +++ b/test/test_method2_extra.py @@ -0,0 +1,92 @@ +""" +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) From b21a0dbe7deb56f1e2dfe7ae5dda3609d1622407 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Sat, 22 Aug 2026 11:08:18 -0400 Subject: [PATCH 02/23] add --- examples/vae/ddp_utils.py | 138 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 138 insertions(+) create mode 100644 examples/vae/ddp_utils.py diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py new file mode 100644 index 0000000..d88723b --- /dev/null +++ b/examples/vae/ddp_utils.py @@ -0,0 +1,138 @@ +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"]) + + ## 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 = os.getenv("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), + ) + + 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 From 307cfb7c35c39b7dac712230213615f3bcf9a5ca Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Mon, 24 Aug 2026 12:02:08 -0400 Subject: [PATCH 03/23] Add FABRIC_IFACE selection from CPU affinity via hwloc --- README.md | 21 ++- examples/vae/distdataset.py | 5 +- examples/vae/vae_core_server.py | 4 + examples/vae/vae_extra_train.py | 4 + setup.py | 2 + src/cpu_nic_map.py | 233 ++++++++++++++++++++++++++++++++ src/pyddstore.pyx | 12 +- 7 files changed, 272 insertions(+), 9 deletions(-) create mode 100644 src/cpu_nic_map.py diff --git a/README.md b/README.md index 2ba698b..0ff8e59 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,7 @@ mpirun -n 4 python my_script.py ## API Reference -### `PyDDStore(comm_or_none=None, method=0, handshake_dir="", n_core=0)` +### `PyDDStore(comm_or_none=None, method=0, handshake_dir="", n_core=0, nic_map=None)` | Parameter | Type | Description | |---|---|---| @@ -79,6 +79,7 @@ mpirun -n 4 python my_script.py | `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: @@ -177,10 +178,19 @@ Uses `MPI_Win_create` and `MPI_Get` for one-sided remote reads. Works on any MPI 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. -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 -``` +`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`) @@ -210,6 +220,7 @@ Environment variables: |---|---|---| | `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 | 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. diff --git a/examples/vae/distdataset.py b/examples/vae/distdataset.py index 84c3900..cef7230 100644 --- a/examples/vae/distdataset.py +++ b/examples/vae/distdataset.py @@ -33,10 +33,7 @@ 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", "0")) - os.environ["FABRIC_IFACE"] = f"hsn{gpu_id//2}" print("DDStore method:", ddstore_method) - print("FABRIC_IFACE:", os.environ["FABRIC_IFACE"]) handshake_dir = os.getenv("DDSTORE_HANDSHAKE_DIR", "./ddstore_hs") if ddstore_method == 2 and self.ddstore_width != self.comm_size: @@ -52,6 +49,7 @@ def __init__(self, data, label, comm=MPI.COMM_WORLD, ddstore_width=None): 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) @@ -128,6 +126,7 @@ def __init__(self, label, handshake_dir, n_core): 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") diff --git a/examples/vae/vae_core_server.py b/examples/vae/vae_core_server.py index 512a081..5956f8e 100644 --- a/examples/vae/vae_core_server.py +++ b/examples/vae/vae_core_server.py @@ -13,6 +13,10 @@ 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 diff --git a/examples/vae/vae_extra_train.py b/examples/vae/vae_extra_train.py index ef6c24c..72f5d31 100644 --- a/examples/vae/vae_extra_train.py +++ b/examples/vae/vae_extra_train.py @@ -14,6 +14,10 @@ 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 diff --git a/setup.py b/setup.py index 75ed248..08aeeae 100644 --- a/setup.py +++ b/setup.py @@ -46,5 +46,7 @@ name="PyDDStore", version="0.1", description="Distributed Data Store", + package_dir={"": "src"}, + py_modules=["cpu_nic_map"], ext_modules=cythonize(extensions) ) diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py new file mode 100644 index 0000000..f8721bb --- /dev/null +++ b/src/cpu_nic_map.py @@ -0,0 +1,233 @@ +#!/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. + +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 -p 'ens*' 0 match a different NIC name pattern + 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) + + 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 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 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: set(hcalc(f"os={n}", "PU")) for n in nics} + nic_numa = {n: (hcalc(f"os={n}", "NUMA") or [None])[0] for n in nics} + if not any(nic_closest.values()): + sys.exit( + "hwloc-calc found no PUs local to any NIC (os= lookups came back " + "empty) -- this process's hwloc topology view has no visibility into the " + "NICs, so any nearest-NIC answer would be a silent guess, not real data. " + "This has been observed when running directly inside a plain `srun` task; " + "try running from the sbatch batch step's own shell instead." + ) + + # 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 serialize_env(pattern="hsn*"): + all_pus, nearest = build_map(pattern) + by_nic = {} + for pu in all_pus: + by_nic.setdefault(nearest(pu)[0], []).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} + 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. + """ + if "FABRIC_IFACE" in os.environ: + return os.environ["FABRIC_IFACE"] + + allocated, nics = allocated_nics(nic_map=nic_map) + if not nics: + 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 len(nics) > 1: + print(f"FABRIC_IFACE: affinity spans {sorted(nics)}, 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" + " cpu_nic_map.py -p 'ens*' 0 match a different NIC name pattern\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" + ), + ) + parser.add_argument("cpu", nargs="?", type=int, + help="CPU (PU) id to look up; omit to print the full table") + parser.add_argument("-p", "--pattern", default="hsn*", + help="glob pattern for NIC names to consider (default: hsn*)") + 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(args.pattern)) + return + + if args.allocated: + allocated, nics = allocated_nics(args.pattern) + print(f"allocated CPUs: {allocated}") + print(f"nearest NIC(s): {sorted(nics)}") + return + + all_pus, nearest = build_map(args.pattern) + + 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(owner) + return + + print(f"{'CPU':>4} {'NUMA':>4} {'nearest NIC':>11} exact") + for pu in all_pus: + owner, exact, numa = nearest(pu) + print(f"{pu:>4} {numa!s:>4} {owner:>11} {exact}") + + +if __name__ == "__main__": + main() diff --git a/src/pyddstore.pyx b/src/pyddstore.pyx index b755efa..31566ad 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 @@ -63,7 +65,7 @@ cdef class PyDDStore: cdef DDStore *c_ddstore def __cinit__(self, comm_or_none=None, int method=0, - str handshake_dir="", int n_core=0): + str handshake_dir="", int n_core=0, nic_map=None): """ Constructors: PyDDStore(comm) — method 0, MPI @@ -72,8 +74,16 @@ cdef class PyDDStore: 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( From f40f182fdbf254d3aa925490ea745e645aa2617b Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Mon, 24 Aug 2026 12:02:12 -0400 Subject: [PATCH 04/23] Add SLURM_STEP_NODELIST for DDP master_addr resolution --- examples/vae/ddp_utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index d88723b..f077501 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -103,13 +103,12 @@ def setup_ddp(): master_port = os.getenv("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_STEP_NODELIST") is not None: + master_addr = parse_slurm_nodelist(os.environ["SLURM_STEP_NODELIST"])[0] elif os.getenv("SLURM_NODELIST") is not None: - ## The following is CADES specific master_addr = parse_slurm_nodelist(os.environ["SLURM_NODELIST"])[0] try: From 5403814fbdc77856db9cf6c3ce91496ad37f93be Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Tue, 25 Aug 2026 21:23:15 -0400 Subject: [PATCH 05/23] Apply black formatting to Python files Co-Authored-By: Claude Sonnet 5 --- examples/scripts/demo.py | 9 ++++-- examples/scripts/test.py | 16 +++++++--- setup.py | 25 ++++++++------- src/cpu_nic_map.py | 65 +++++++++++++++++++++++++++----------- test/conftest.py | 1 + test/test_method2_core.py | 16 +++++++--- test/test_method2_extra.py | 8 ++++- test/test_multirank.py | 12 +++++-- test/test_single.py | 22 ++++++++++--- 9 files changed, 125 insertions(+), 49 deletions(-) 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/setup.py b/setup.py index 08aeeae..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,16 +31,19 @@ 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", @@ -48,5 +51,5 @@ description="Distributed Data Store", package_dir={"": "src"}, py_modules=["cpu_nic_map"], - ext_modules=cythonize(extensions) + ext_modules=cythonize(extensions), ) diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index f8721bb..1533279 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -21,6 +21,7 @@ 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 @@ -38,17 +39,25 @@ def hcalc(loc, itype): # 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() + 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 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)) + 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") @@ -104,7 +113,9 @@ def serialize_env(pattern="hsn*"): by_nic = {} for pu in all_pus: by_nic.setdefault(nearest(pu)[0], []).append(pu) - return ";".join(f"{nic}={compress_ranges(pus)}" for nic, pus in sorted(by_nic.items())) + return ";".join( + f"{nic}={compress_ranges(pus)}" for nic, pus in sorted(by_nic.items()) + ) def parse_env(s): @@ -182,7 +193,7 @@ def select_fabric_iface(nic_map=None): def main(): parser = argparse.ArgumentParser( description="Find the nearest HSN (Slingshot) NIC for a given CPU (PU) id, " - "based on hwloc PCI/NUMA locality.", + "based on hwloc PCI/NUMA locality.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=( "examples:\n" @@ -193,15 +204,29 @@ def main(): " srun ... python cpu_nic_map.py --allocated show this task's allocated CPUs + nearest NIC(s)\n" ), ) - parser.add_argument("cpu", nargs="?", type=int, - help="CPU (PU) id to look up; omit to print the full table") - parser.add_argument("-p", "--pattern", default="hsn*", - help="glob pattern for NIC names to consider (default: hsn*)") - 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") + parser.add_argument( + "cpu", + nargs="?", + type=int, + help="CPU (PU) id to look up; omit to print the full table", + ) + parser.add_argument( + "-p", + "--pattern", + default="hsn*", + help="glob pattern for NIC names to consider (default: hsn*)", + ) + 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: @@ -218,7 +243,9 @@ def main(): 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)})") + sys.exit( + f"cpu id {args.cpu} not found (valid range: {min(all_pus)}-{max(all_pus)})" + ) owner, exact, numa = nearest(args.cpu) print(owner) return 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 index fdb1eeb..94324f3 100644 --- a/test/test_method2_core.py +++ b/test/test_method2_core.py @@ -35,6 +35,7 @@ # configuration — resolve directory with same priority as the C library # --------------------------------------------------------------------------- + def _resolve_dir(arg): if arg: return arg @@ -43,12 +44,17 @@ def _resolve_dir(arg): 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")) +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() +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +size = comm.Get_size() assert size == n_core, f"Expected {n_core} core ranks, got {size}" @@ -70,7 +76,7 @@ def _resolve_dir(arg): nrows = 8 ncols = 4 -data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) +data = np.full((nrows, ncols), float(rank + 1), dtype=np.float32) # --------------------------------------------------------------------------- # create DDStore (method=2, core member) diff --git a/test/test_method2_extra.py b/test/test_method2_extra.py index 530d476..c6c0937 100644 --- a/test/test_method2_extra.py +++ b/test/test_method2_extra.py @@ -27,6 +27,7 @@ # configuration — resolve directory with same priority as the C library # --------------------------------------------------------------------------- + def _resolve_dir(arg): if arg: return arg @@ -35,8 +36,13 @@ def _resolve_dir(arg): 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")) +n_core = ( + int(sys.argv[2]) + if len(sys.argv) > 2 + else int(os.environ.get("DDSTORE_N_CORE", "4")) +) nrows = 8 ncols = 4 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) From 923934f8ddfc9523f29f05c8a0e2d09626c8e23f Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Tue, 25 Aug 2026 21:24:12 -0400 Subject: [PATCH 06/23] Fix multi-GPU DDP setup in VAE example Assign each rank its own local GPU (via OMPI/SLURM local-rank env vars, falling back to rank % device_count) instead of always using cuda:0, and restrict per-epoch logging, testing, and sample checkpointing to rank 0 to avoid duplicated output and file-write races across ranks. Co-Authored-By: Claude Sonnet 5 --- examples/vae/ddp_utils.py | 15 +++++++++++++ examples/vae/vae-ddp.py | 47 ++++++++++++++++++++++++++------------- 2 files changed, 46 insertions(+), 16 deletions(-) diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index f077501..472c5cc 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -31,6 +31,21 @@ def init_comm_size_and_rank(): 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"]) + elif torch.cuda.is_available() and torch.cuda.device_count() > 0: + return rank % torch.cuda.device_count() + return 0 + + def find_ifname(myaddr): """ Find socket ifname for a given ip adress. This is for "GLOO" ddp setup. diff --git a/examples/vae/vae-ddp.py b/examples/vae/vae-ddp.py index bc2b93d..d6d44fb 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -3,6 +3,7 @@ ## 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 optim @@ -19,7 +20,7 @@ import distdataset from distdataset import DistDataset -from ddp_utils import setup_ddp +from ddp_utils import setup_ddp, get_local_rank from vae_model import VAE, loss_function parser = argparse.ArgumentParser(description="VAE MNIST Example") @@ -59,19 +60,31 @@ torch.manual_seed(args.seed) +comm = MPI.COMM_WORLD +comm_size, rank = setup_ddp() + if args.cuda: - device = torch.device("cuda") + local_rank = get_local_rank(rank) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") elif use_mps: device = torch.device("mps") else: device = torch.device("cpu") -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) +if args.cuda: + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[local_rank], output_device=local_rank + ) +else: + model = torch.nn.parallel.DistributedDataParallel(model) optimizer = optim.Adam(model.parameters(), lr=1e-3) # kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} @@ -133,11 +146,12 @@ def train(epoch): 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): @@ -167,12 +181,13 @@ def test(epoch): # 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() From 8f25ac4cec66527440bbfdfead62f0ba1d70f031 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Tue, 25 Aug 2026 21:24:22 -0400 Subject: [PATCH 07/23] Combine method=2 handshake into a single file per variable --- include/common.h | 40 +++---- include/ddstore.hpp | 37 +++---- src/common.cxx | 260 +++++++++++++++++++++++--------------------- 3 files changed, 171 insertions(+), 166 deletions(-) diff --git a/include/common.h b/include/common.h index 37f8028..04b0c28 100644 --- a/include/common.h +++ b/include/common.h @@ -10,9 +10,11 @@ #define COMM_FILE_WRITER_TO_READER "./writer_address.bin" /* ----------------------------------------------------------------------- - * Method 2: file-based handshake record written by each core rank. - * One file per variable per core rank: - * {handshake_dir}/{varname}_rank{rank}.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 { @@ -71,23 +73,23 @@ extern "C" * Returns pointer to a static buffer — copy if needed across calls. */ const char *resolve_handshake_dir(const char *user_dir); - /* Core rank: write own CoreRecord to {dir}/{varname}_rank{rank}.bin. - * Uses write-to-tmp + rename for atomicity. */ - int handshake_write(struct fabric_state *fs, - const char *dir, const char *varname, int rank, - long nrows, int disp, int itemsize); + /* 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); - /* Core rank: poll until all n_core record files are present, then read - * them all and populate fs->comm_partner[], remote_key[], - * remote_address[], and fill lenlist[0..n_core-1] (raw row counts, - * NOT yet prefix-summed). Also sets *out_disp and *out_itemsize. */ - int handshake_read(struct fabric_state *fs, - const char *dir, const char *varname, - int n_core, int my_rank, - long *lenlist, int *out_disp, int *out_itemsize); - - /* Extra member: same as handshake_read but does NOT write anything. - * Blocks until all n_core files appear (with timeout). */ + /* 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, diff --git a/include/ddstore.hpp b/include/ddstore.hpp index ae3a831..48a0bd6 100644 --- a/include/ddstore.hpp +++ b/include/ddstore.hpp @@ -125,20 +125,15 @@ class DDStore throw std::runtime_error(std::string("fi_mr_reg failed: ") + fi_strerror(mr_rc)); fabric_state->key = fi_mr_key(fabric_state->mr); - /* Write this rank's record to the handshake directory. */ - if (handshake_write(fabric_state, - this->handshake_dir.c_str(), name.c_str(), this->rank, - nrows, disp, (int)sizeof(T)) != 0) - throw std::runtime_error("handshake_write failed"); - - /* Wait for all core ranks, then read all records. */ + /* 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); - int file_disp = 0, file_itemsize = 0; - if (handshake_read(fabric_state, - this->handshake_dir.c_str(), name.c_str(), - this->n_core, this->rank, - raw_lens.data(), &file_disp, &file_itemsize) != 0) - throw std::runtime_error("handshake_read failed"); + 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; @@ -254,18 +249,12 @@ class DDStore throw std::runtime_error(std::string("fi_mr_reg failed: ") + fi_strerror(mr_rc)); fabric_state->key = fi_mr_key(fabric_state->mr); - if (handshake_write(fabric_state, - this->handshake_dir.c_str(), name.c_str(), this->rank, - nrows, disp, itemsize) != 0) - throw std::runtime_error("handshake_write failed"); - std::vector raw_lens(this->n_core); - int file_disp = 0, file_itemsize = 0; - if (handshake_read(fabric_state, - this->handshake_dir.c_str(), name.c_str(), - this->n_core, this->rank, - raw_lens.data(), &file_disp, &file_itemsize) != 0) - throw std::runtime_error("handshake_read failed"); + 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); diff --git a/src/common.cxx b/src/common.cxx index 2ab4997..9f68d0e 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -13,6 +13,7 @@ #include #include #include +#include void init_fabric(struct fabric_state *fabric) { @@ -397,7 +398,9 @@ int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset * ========================================================================= * * File naming convention: - * {dir}/{varname}_rank{rank}.bin — one CoreRecord per core rank + * {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) @@ -436,11 +439,11 @@ const char *resolve_handshake_dir(const char *user_dir) return resolved; } -/* Build the canonical path for rank i's record file into `buf`. */ +/* Build the canonical path for a variable's combined record file into `buf`. */ static void record_path(char *buf, size_t bufsz, - const char *dir, const char *varname, int rank) + const char *dir, const char *varname) { - snprintf(buf, bufsz, "%s/%s_rank%d.bin", dir, varname, rank); + snprintf(buf, bufsz, "%s/%s.bin", dir, varname); } /* Return the configured timeout in seconds (default 300). */ @@ -455,16 +458,22 @@ static int handshake_timeout_s(void) * handshake_write() * * Called by each core rank after init_fabric() and fi_mr_reg(). - * Writes a CoreRecord to {resolved_dir}/{varname}_rank{rank}.bin. - * The directory is resolved via resolve_handshake_dir(). + * 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. * -------------------------------------------------------------------------- */ -int handshake_write(struct fabric_state *fs, - const char *dir, const char *varname, int rank, - long nrows, int disp, int itemsize) +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) { - const char *rdir = resolve_handshake_dir(dir); + int rank = 0; + MPI_Comm_rank(comm, &rank); - /* Build the CoreRecord for this rank. */ + /* Build this rank's CoreRecord. */ struct CoreRecord rec; memset(&rec, 0, sizeof(rec)); @@ -484,67 +493,106 @@ int handshake_write(struct fabric_state *fs, rec.disp = disp; rec.itemsize = itemsize; - /* 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. */ - char bin_path[4096]; - char tmp_path[4096 + 32]; - record_path(bin_path, sizeof(bin_path), rdir, varname, rank); - snprintf(tmp_path, sizeof(tmp_path), "%s.tmp.%d", bin_path, (int)getpid()); + /* 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); - FILE *f = fopen(tmp_path, "wb"); - if (!f) + /* 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] cannot open %s: ", tmp_path); - perror(""); + fprintf(stderr, "[handshake_write] malloc failed\n"); return 1; } - if (fwrite(&rec, sizeof(rec), 1, f) != 1) + for (int i = 0; i < n_core; i++) { - fprintf(stderr, "[handshake_write] fwrite failed for %s\n", tmp_path); - fclose(f); - unlink(tmp_path); - return 1; + 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; } - if (fflush(f) != 0 || fsync(fileno(f)) != 0) + 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) { - fprintf(stderr, "[handshake_write] fsync failed for %s: ", tmp_path); - perror(""); + 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); - 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; + 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); } - fprintf(stderr, "[handshake_write] rank %d wrote %s\n", rank, bin_path); return 0; } /* -------------------------------------------------------------------------- - * handshake_read() + * handshake_join() * - * Called by each core rank after handshake_write(). - * Polls until all n_core record files are present in the resolved directory, - * then reads them and populates: + * 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 rank-0 record; assumed uniform) + * *out_disp, *out_itemsize (from record 0; assumed uniform) * * Returns 0 on success, non-zero on error or timeout. * -------------------------------------------------------------------------- */ -int handshake_read(struct fabric_state *fs, +int handshake_join(struct fabric_state *fs, const char *dir, const char *varname, - int n_core, int my_rank, + int n_core, long *lenlist, int *out_disp, int *out_itemsize) { const char *rdir = resolve_handshake_dir(dir); @@ -552,30 +600,15 @@ int handshake_read(struct fabric_state *fs, struct timespec ts_start, ts_now; clock_gettime(CLOCK_MONOTONIC, &ts_start); - /* Allocate arrays sized for n_core peers. */ - 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_read] malloc failed\n"); - return 1; - } + char path[4096]; + record_path(path, sizeof(path), rdir, varname); + size_t expected_size = (size_t)n_core * sizeof(struct CoreRecord); - /* Poll until all n_core files are fully written (size == sizeof CoreRecord). */ + /* Poll until the combined record file is fully written. */ for (;;) { - int ready = 0; - for (int i = 0; i < n_core; i++) - { - char path[4096]; - record_path(path, sizeof(path), rdir, varname, i); - struct stat st; - if (stat(path, &st) == 0 && - st.st_size == (off_t)sizeof(struct CoreRecord)) - ready++; - } - if (ready == n_core) + struct stat st; + if (stat(path, &st) == 0 && st.st_size == (off_t)expected_size) break; clock_gettime(CLOCK_MONOTONIC, &ts_now); @@ -584,77 +617,58 @@ int handshake_read(struct fabric_state *fs, if (elapsed > timeout_s) { fprintf(stderr, - "[handshake_read] timeout after %.0f s waiting for " - "%d/%d core records (var=%s, dir=%s)\n", - elapsed, ready, n_core, varname, rdir); + "[handshake_join] timeout after %.0f s waiting for " + "%s (var=%s, dir=%s)\n", + elapsed, path, varname, rdir); return 1; } usleep(50000); /* 50 ms */ } - /* Read all records. */ - for (int i = 0; i < n_core; i++) + std::vector all_recs(n_core); + FILE *f = fopen(path, "rb"); + if (!f) { - char path[4096]; - record_path(path, sizeof(path), rdir, varname, i); - - FILE *f = fopen(path, "rb"); - if (!f) - { - fprintf(stderr, "[handshake_read] cannot open %s: ", path); - perror(""); - return 1; - } - - struct CoreRecord rec; - if (fread(&rec, sizeof(rec), 1, f) != 1) - { - fprintf(stderr, "[handshake_read] fread failed for %s\n", path); - fclose(f); - return 1; - } + 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); - /* Insert fabric address into the address vector. */ - int rc = fi_av_insert(fs->av, - rec.fabric_address, 1, + 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_read] fi_av_insert failed for rank %d (rc=%d)\n", + "[handshake_join] fi_av_insert failed for rank %d (rc=%d)\n", i, rc); return 1; } - - fs->remote_key[i] = rec.key; - fs->remote_address[i] = rec.base_address; - lenlist[i] = rec.nrows; - - if (i == 0) - { - if (out_disp) *out_disp = rec.disp; - if (out_itemsize) *out_itemsize = rec.itemsize; - } + 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; - (void)my_rank; return 0; -} - -/* -------------------------------------------------------------------------- - * handshake_join() - * - * Called by extra members (no MPI, no data to publish). - * Identical to handshake_read() except it never writes anything. - * Blocks until all n_core record files are present (or 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) -{ - return handshake_read(fs, dir, varname, n_core, -1, - lenlist, out_disp, out_itemsize); } \ No newline at end of file From 6b5c92d14f59f09e6bd8797fcce134ea17640198 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Tue, 25 Aug 2026 21:24:12 -0400 Subject: [PATCH 08/23] Fix multi-GPU DDP setup in VAE example --- examples/vae/ddp_utils.py | 15 +++++++++++++ examples/vae/vae-ddp.py | 47 ++++++++++++++++++++++++++------------- 2 files changed, 46 insertions(+), 16 deletions(-) diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index f077501..472c5cc 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -31,6 +31,21 @@ def init_comm_size_and_rank(): 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"]) + elif torch.cuda.is_available() and torch.cuda.device_count() > 0: + return rank % torch.cuda.device_count() + return 0 + + def find_ifname(myaddr): """ Find socket ifname for a given ip adress. This is for "GLOO" ddp setup. diff --git a/examples/vae/vae-ddp.py b/examples/vae/vae-ddp.py index bc2b93d..d6d44fb 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -3,6 +3,7 @@ ## 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 optim @@ -19,7 +20,7 @@ import distdataset from distdataset import DistDataset -from ddp_utils import setup_ddp +from ddp_utils import setup_ddp, get_local_rank from vae_model import VAE, loss_function parser = argparse.ArgumentParser(description="VAE MNIST Example") @@ -59,19 +60,31 @@ torch.manual_seed(args.seed) +comm = MPI.COMM_WORLD +comm_size, rank = setup_ddp() + if args.cuda: - device = torch.device("cuda") + local_rank = get_local_rank(rank) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") elif use_mps: device = torch.device("mps") else: device = torch.device("cpu") -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) +if args.cuda: + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[local_rank], output_device=local_rank + ) +else: + model = torch.nn.parallel.DistributedDataParallel(model) optimizer = optim.Adam(model.parameters(), lr=1e-3) # kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} @@ -133,11 +146,12 @@ def train(epoch): 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): @@ -167,12 +181,13 @@ def test(epoch): # 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() From 6117437a09b3806a5402f84ee90a202493413d1e Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Tue, 25 Aug 2026 21:24:22 -0400 Subject: [PATCH 09/23] Combine method=2 handshake into a single file per variable --- include/common.h | 40 +++---- include/ddstore.hpp | 37 +++---- src/common.cxx | 260 +++++++++++++++++++++++--------------------- 3 files changed, 171 insertions(+), 166 deletions(-) diff --git a/include/common.h b/include/common.h index 37f8028..04b0c28 100644 --- a/include/common.h +++ b/include/common.h @@ -10,9 +10,11 @@ #define COMM_FILE_WRITER_TO_READER "./writer_address.bin" /* ----------------------------------------------------------------------- - * Method 2: file-based handshake record written by each core rank. - * One file per variable per core rank: - * {handshake_dir}/{varname}_rank{rank}.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 { @@ -71,23 +73,23 @@ extern "C" * Returns pointer to a static buffer — copy if needed across calls. */ const char *resolve_handshake_dir(const char *user_dir); - /* Core rank: write own CoreRecord to {dir}/{varname}_rank{rank}.bin. - * Uses write-to-tmp + rename for atomicity. */ - int handshake_write(struct fabric_state *fs, - const char *dir, const char *varname, int rank, - long nrows, int disp, int itemsize); + /* 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); - /* Core rank: poll until all n_core record files are present, then read - * them all and populate fs->comm_partner[], remote_key[], - * remote_address[], and fill lenlist[0..n_core-1] (raw row counts, - * NOT yet prefix-summed). Also sets *out_disp and *out_itemsize. */ - int handshake_read(struct fabric_state *fs, - const char *dir, const char *varname, - int n_core, int my_rank, - long *lenlist, int *out_disp, int *out_itemsize); - - /* Extra member: same as handshake_read but does NOT write anything. - * Blocks until all n_core files appear (with timeout). */ + /* 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, diff --git a/include/ddstore.hpp b/include/ddstore.hpp index ae3a831..48a0bd6 100644 --- a/include/ddstore.hpp +++ b/include/ddstore.hpp @@ -125,20 +125,15 @@ class DDStore throw std::runtime_error(std::string("fi_mr_reg failed: ") + fi_strerror(mr_rc)); fabric_state->key = fi_mr_key(fabric_state->mr); - /* Write this rank's record to the handshake directory. */ - if (handshake_write(fabric_state, - this->handshake_dir.c_str(), name.c_str(), this->rank, - nrows, disp, (int)sizeof(T)) != 0) - throw std::runtime_error("handshake_write failed"); - - /* Wait for all core ranks, then read all records. */ + /* 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); - int file_disp = 0, file_itemsize = 0; - if (handshake_read(fabric_state, - this->handshake_dir.c_str(), name.c_str(), - this->n_core, this->rank, - raw_lens.data(), &file_disp, &file_itemsize) != 0) - throw std::runtime_error("handshake_read failed"); + 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; @@ -254,18 +249,12 @@ class DDStore throw std::runtime_error(std::string("fi_mr_reg failed: ") + fi_strerror(mr_rc)); fabric_state->key = fi_mr_key(fabric_state->mr); - if (handshake_write(fabric_state, - this->handshake_dir.c_str(), name.c_str(), this->rank, - nrows, disp, itemsize) != 0) - throw std::runtime_error("handshake_write failed"); - std::vector raw_lens(this->n_core); - int file_disp = 0, file_itemsize = 0; - if (handshake_read(fabric_state, - this->handshake_dir.c_str(), name.c_str(), - this->n_core, this->rank, - raw_lens.data(), &file_disp, &file_itemsize) != 0) - throw std::runtime_error("handshake_read failed"); + 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); diff --git a/src/common.cxx b/src/common.cxx index 2ab4997..9f68d0e 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -13,6 +13,7 @@ #include #include #include +#include void init_fabric(struct fabric_state *fabric) { @@ -397,7 +398,9 @@ int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset * ========================================================================= * * File naming convention: - * {dir}/{varname}_rank{rank}.bin — one CoreRecord per core rank + * {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) @@ -436,11 +439,11 @@ const char *resolve_handshake_dir(const char *user_dir) return resolved; } -/* Build the canonical path for rank i's record file into `buf`. */ +/* Build the canonical path for a variable's combined record file into `buf`. */ static void record_path(char *buf, size_t bufsz, - const char *dir, const char *varname, int rank) + const char *dir, const char *varname) { - snprintf(buf, bufsz, "%s/%s_rank%d.bin", dir, varname, rank); + snprintf(buf, bufsz, "%s/%s.bin", dir, varname); } /* Return the configured timeout in seconds (default 300). */ @@ -455,16 +458,22 @@ static int handshake_timeout_s(void) * handshake_write() * * Called by each core rank after init_fabric() and fi_mr_reg(). - * Writes a CoreRecord to {resolved_dir}/{varname}_rank{rank}.bin. - * The directory is resolved via resolve_handshake_dir(). + * 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. * -------------------------------------------------------------------------- */ -int handshake_write(struct fabric_state *fs, - const char *dir, const char *varname, int rank, - long nrows, int disp, int itemsize) +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) { - const char *rdir = resolve_handshake_dir(dir); + int rank = 0; + MPI_Comm_rank(comm, &rank); - /* Build the CoreRecord for this rank. */ + /* Build this rank's CoreRecord. */ struct CoreRecord rec; memset(&rec, 0, sizeof(rec)); @@ -484,67 +493,106 @@ int handshake_write(struct fabric_state *fs, rec.disp = disp; rec.itemsize = itemsize; - /* 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. */ - char bin_path[4096]; - char tmp_path[4096 + 32]; - record_path(bin_path, sizeof(bin_path), rdir, varname, rank); - snprintf(tmp_path, sizeof(tmp_path), "%s.tmp.%d", bin_path, (int)getpid()); + /* 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); - FILE *f = fopen(tmp_path, "wb"); - if (!f) + /* 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] cannot open %s: ", tmp_path); - perror(""); + fprintf(stderr, "[handshake_write] malloc failed\n"); return 1; } - if (fwrite(&rec, sizeof(rec), 1, f) != 1) + for (int i = 0; i < n_core; i++) { - fprintf(stderr, "[handshake_write] fwrite failed for %s\n", tmp_path); - fclose(f); - unlink(tmp_path); - return 1; + 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; } - if (fflush(f) != 0 || fsync(fileno(f)) != 0) + 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) { - fprintf(stderr, "[handshake_write] fsync failed for %s: ", tmp_path); - perror(""); + 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); - 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; + 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); } - fprintf(stderr, "[handshake_write] rank %d wrote %s\n", rank, bin_path); return 0; } /* -------------------------------------------------------------------------- - * handshake_read() + * handshake_join() * - * Called by each core rank after handshake_write(). - * Polls until all n_core record files are present in the resolved directory, - * then reads them and populates: + * 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 rank-0 record; assumed uniform) + * *out_disp, *out_itemsize (from record 0; assumed uniform) * * Returns 0 on success, non-zero on error or timeout. * -------------------------------------------------------------------------- */ -int handshake_read(struct fabric_state *fs, +int handshake_join(struct fabric_state *fs, const char *dir, const char *varname, - int n_core, int my_rank, + int n_core, long *lenlist, int *out_disp, int *out_itemsize) { const char *rdir = resolve_handshake_dir(dir); @@ -552,30 +600,15 @@ int handshake_read(struct fabric_state *fs, struct timespec ts_start, ts_now; clock_gettime(CLOCK_MONOTONIC, &ts_start); - /* Allocate arrays sized for n_core peers. */ - 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_read] malloc failed\n"); - return 1; - } + char path[4096]; + record_path(path, sizeof(path), rdir, varname); + size_t expected_size = (size_t)n_core * sizeof(struct CoreRecord); - /* Poll until all n_core files are fully written (size == sizeof CoreRecord). */ + /* Poll until the combined record file is fully written. */ for (;;) { - int ready = 0; - for (int i = 0; i < n_core; i++) - { - char path[4096]; - record_path(path, sizeof(path), rdir, varname, i); - struct stat st; - if (stat(path, &st) == 0 && - st.st_size == (off_t)sizeof(struct CoreRecord)) - ready++; - } - if (ready == n_core) + struct stat st; + if (stat(path, &st) == 0 && st.st_size == (off_t)expected_size) break; clock_gettime(CLOCK_MONOTONIC, &ts_now); @@ -584,77 +617,58 @@ int handshake_read(struct fabric_state *fs, if (elapsed > timeout_s) { fprintf(stderr, - "[handshake_read] timeout after %.0f s waiting for " - "%d/%d core records (var=%s, dir=%s)\n", - elapsed, ready, n_core, varname, rdir); + "[handshake_join] timeout after %.0f s waiting for " + "%s (var=%s, dir=%s)\n", + elapsed, path, varname, rdir); return 1; } usleep(50000); /* 50 ms */ } - /* Read all records. */ - for (int i = 0; i < n_core; i++) + std::vector all_recs(n_core); + FILE *f = fopen(path, "rb"); + if (!f) { - char path[4096]; - record_path(path, sizeof(path), rdir, varname, i); - - FILE *f = fopen(path, "rb"); - if (!f) - { - fprintf(stderr, "[handshake_read] cannot open %s: ", path); - perror(""); - return 1; - } - - struct CoreRecord rec; - if (fread(&rec, sizeof(rec), 1, f) != 1) - { - fprintf(stderr, "[handshake_read] fread failed for %s\n", path); - fclose(f); - return 1; - } + 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); - /* Insert fabric address into the address vector. */ - int rc = fi_av_insert(fs->av, - rec.fabric_address, 1, + 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_read] fi_av_insert failed for rank %d (rc=%d)\n", + "[handshake_join] fi_av_insert failed for rank %d (rc=%d)\n", i, rc); return 1; } - - fs->remote_key[i] = rec.key; - fs->remote_address[i] = rec.base_address; - lenlist[i] = rec.nrows; - - if (i == 0) - { - if (out_disp) *out_disp = rec.disp; - if (out_itemsize) *out_itemsize = rec.itemsize; - } + 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; - (void)my_rank; return 0; -} - -/* -------------------------------------------------------------------------- - * handshake_join() - * - * Called by extra members (no MPI, no data to publish). - * Identical to handshake_read() except it never writes anything. - * Blocks until all n_core record files are present (or 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) -{ - return handshake_read(fs, dir, varname, n_core, -1, - lenlist, out_disp, out_itemsize); } \ No newline at end of file From e50ecd10c64dc647b7ad81f57da4c2ef41992051 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 15:26:51 -0400 Subject: [PATCH 10/23] fix setdevice --- examples/vae/vae-ddp.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/examples/vae/vae-ddp.py b/examples/vae/vae-ddp.py index d6d44fb..51adcc8 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -65,8 +65,11 @@ if args.cuda: local_rank = get_local_rank(rank) - torch.cuda.set_device(local_rank) - device = torch.device(f"cuda:{local_rank}") + 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 use_mps: device = torch.device("mps") else: @@ -80,8 +83,9 @@ model = VAE().to(device) if args.cuda: + cur_device = torch.cuda.current_device() model = torch.nn.parallel.DistributedDataParallel( - model, device_ids=[local_rank], output_device=local_rank + model, device_ids=[cur_device], output_device=cur_device ) else: model = torch.nn.parallel.DistributedDataParallel(model) From 5991499107d0da53d6c2c443c0e5a28f3c64710e Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 17:21:55 -0400 Subject: [PATCH 11/23] Add cxi (Perlmutter) fabric support alongside hsn (Frontier), selected via DDSTORE_FABRIC_PROVIDER --- README.md | 17 +- build-evn.sh | 30 ++++ examples/vae/distdataset.py | 6 +- include/common.h | 35 ++++ include/ddstore.hpp | 40 ++++- module-to-load-frontier.sh | 15 ++ module-to-load-perlmutter.sh | 10 ++ src/common.cxx | 336 ++++++++++++++++++++++++++++++++++- src/cpu_nic_map.py | 46 ++++- 9 files changed, 521 insertions(+), 14 deletions(-) create mode 100644 build-evn.sh create mode 100755 module-to-load-frontier.sh create mode 100644 module-to-load-perlmutter.sh diff --git a/README.md b/README.md index 0ff8e59..bd37795 100644 --- a/README.md +++ b/README.md @@ -176,7 +176,19 @@ 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_PROVIDER`** selects which libfabric provider to open, for `method=1`/`2`: + +- `hsn` (default, unset) — Frontier: opens the `tcp;ofi_rxm` domain over Cray Slingshot. +- `cxi` — Perlmutter: opens the native `cxi` domain over Cray Slingshot. + +The two are independent code paths (not runtime auto-detection), so set this explicitly per system rather than relying on a guess: + +```bash +export DDSTORE_FABRIC_PROVIDER=hsn # Frontier (default; usually not needed) +export DDSTORE_FABRIC_PROVIDER=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: @@ -196,7 +208,7 @@ Uses `fi_read` for true RDMA transfers over high-speed interconnects (Infiniband 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()`. Each rank writes a `{name}_rank{N}.bin` record (fabric address, MR key, base pointer, row count, dtype) into `handshake_dir`. +- **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 @@ -221,6 +233,7 @@ Environment variables: | `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_PROVIDER` | `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. diff --git a/build-evn.sh b/build-evn.sh new file mode 100644 index 0000000..42a7b53 --- /dev/null +++ b/build-evn.sh @@ -0,0 +1,30 @@ +#!/bin/bash + +# Load system modules +module load pytorch/2.13.0 +module unload darshan + +export PYTHONNOUSERSITE=1 +export MPICH_GPU_SUPPORT_ENABLED=0 +export DDSTORE_FABRIC_PROVIDER=cxi + +VENV_DIR=.venv +python -m venv --system-site-packages "${VENV_DIR}" + +source .venv/bin/activate +pip install --upgrade pip + +# Install build and runtime dependencies +pip install wheel Cython +MPICC="cc -shared" pip install --no-cache-dir --no-binary=mpi4py mpi4py==4.1.1 +pip install pytest pytest-mpi +pip install psutil + +# Build and install PyDDStore using Cray compiler wrappers +CC=cc CXX=CC pip install -e . + +# Install PyTorch with ROCm support +pip3 install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.2 + +echo "Environment is ready at '.venv'." +echo "To activate: source .venv/bin/activate" diff --git a/examples/vae/distdataset.py b/examples/vae/distdataset.py index cef7230..6f07f58 100644 --- a/examples/vae/distdataset.py +++ b/examples/vae/distdataset.py @@ -37,9 +37,9 @@ def __init__(self, data, label, comm=MPI.COMM_WORLD, ddstore_width=None): 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 under the - # same {varname}_rank{N}.bin filenames, so more than one group - # sharing a handshake_dir would silently collide. + # 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 " diff --git a/include/common.h b/include/common.h index 04b0c28..18318ef 100644 --- a/include/common.h +++ b/include/common.h @@ -4,6 +4,7 @@ #include #include #include +#include #include #define DP_AV_DEF_SIZE 512 @@ -62,6 +63,40 @@ 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); diff --git a/include/ddstore.hpp b/include/ddstore.hpp index 48a0bd6..2c582ee 100644 --- a/include/ddstore.hpp +++ b/include/ddstore.hpp @@ -98,7 +98,8 @@ 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) { @@ -123,6 +124,20 @@ class DDStore 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 @@ -223,7 +238,8 @@ 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) { @@ -247,6 +263,20 @@ class DDStore 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); @@ -369,7 +399,11 @@ class DDStore /* 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) + ")"); } } diff --git a/module-to-load-frontier.sh b/module-to-load-frontier.sh new file mode 100755 index 0000000..72a2aac --- /dev/null +++ b/module-to-load-frontier.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +module load PrgEnv-gnu +module load cpe/24.11 +module load libfabric +module unload darshan-runtime +module load cray-python/3.11.7 +module load rocm/7.2.0 + +export PYTHONNOUSERSITE=1 +export DDSTORE_FABRIC_PROVIDER=hsn + +## python env +source .venv/bin/activate + diff --git a/module-to-load-perlmutter.sh b/module-to-load-perlmutter.sh new file mode 100644 index 0000000..fe0b710 --- /dev/null +++ b/module-to-load-perlmutter.sh @@ -0,0 +1,10 @@ +#!/bin/bash + +module load pytorch/2.13.0 + +export PYTHONNOUSERSITE=1 +export MPICH_GPU_SUPPORT_ENABLED=0 +export DDSTORE_FABRIC_PROVIDER=cxi + +## python env +source .venv/bin/activate diff --git a/src/common.cxx b/src/common.cxx index 9f68d0e..f0f4214 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -15,7 +15,9 @@ #include #include -void init_fabric(struct fabric_state *fabric) +/* 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}; @@ -255,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_PROVIDER 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_PROVIDER"); + 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]; @@ -277,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); @@ -310,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++) @@ -339,6 +644,15 @@ 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)) + { + fi_mr_bind(fabric_state->recv_mr, &fabric_state->signal->fid, 0); + fi_mr_enable(fabric_state->recv_mr); + } + void *memory_descriptor = NULL; if (is_local_mr_req(fabric_state)) { @@ -384,8 +698,18 @@ 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; } } @@ -488,7 +812,9 @@ int handshake_write(struct fabric_state *fs, MPI_Comm comm, } rec.key = fs->key; - rec.base_address = (uint64_t)(uintptr_t)fs->send_data; + /* 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; diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index 1533279..db4cbf0 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -26,6 +26,7 @@ import fnmatch import glob import os +import re import subprocess import sys @@ -172,12 +173,49 @@ def select_fabric_iface(nic_map=None): 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_PROVIDER 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_PROVIDER") == "cxi" + if "FABRIC_IFACE" in os.environ: - return os.environ["FABRIC_IFACE"] + iface = os.environ["FABRIC_IFACE"] + if use_cxi: + m = re.match(r"hsn(\d+)$", iface) + if m: + iface = f"cxi{m.group(1)}" + 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)] + m = re.match(r"hsn(\d+)$", hsn) + iface = f"cxi{m.group(1)}" if m else hsn + 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 " @@ -186,6 +224,12 @@ def select_fabric_iface(nic_map=None): iface = sorted(nics)[0] if len(nics) > 1: print(f"FABRIC_IFACE: affinity spans {sorted(nics)}, picking {iface}") + + if use_cxi: + m = re.match(r"hsn(\d+)$", iface) + if m: + iface = f"cxi{m.group(1)}" + os.environ["FABRIC_IFACE"] = iface return iface From e6805207959df20aa57a7e558ec03ad2f804201d Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 17:41:25 -0400 Subject: [PATCH 12/23] Add --provider flag to cpu_nic_map.py CLI to preview cxi-translated NIC names --- src/cpu_nic_map.py | 48 +++++++++++++++++++++++++++++++++------------- 1 file changed, 35 insertions(+), 13 deletions(-) diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index db4cbf0..4e413eb 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -109,11 +109,23 @@ def compress_ranges(values): return ",".join(out) -def serialize_env(pattern="hsn*"): +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: - by_nic.setdefault(nearest(pu)[0], []).append(pu) + 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()) ) @@ -189,9 +201,9 @@ def select_fabric_iface(nic_map=None): if "FABRIC_IFACE" in os.environ: iface = os.environ["FABRIC_IFACE"] if use_cxi: - m = re.match(r"hsn(\d+)$", iface) - if m: - iface = f"cxi{m.group(1)}" + translated = translate_iface(iface, "cxi") + if translated != iface: + iface = translated os.environ["FABRIC_IFACE"] = iface return iface @@ -211,8 +223,7 @@ def select_fabric_iface(nic_map=None): ) local_rank = int(os.environ.get("SLURM_LOCALID", 0)) hsn = cxi_domains[local_rank % len(cxi_domains)] - m = re.match(r"hsn(\d+)$", hsn) - iface = f"cxi{m.group(1)}" if m else hsn + iface = translate_iface(hsn, "cxi") print(f"FABRIC_IFACE: fallback (SLURM_LOCALID={local_rank}) -> {iface}") os.environ["FABRIC_IFACE"] = iface return iface @@ -226,9 +237,7 @@ def select_fabric_iface(nic_map=None): print(f"FABRIC_IFACE: affinity spans {sorted(nics)}, picking {iface}") if use_cxi: - m = re.match(r"hsn(\d+)$", iface) - if m: - iface = f"cxi{m.group(1)}" + iface = translate_iface(iface, "cxi") os.environ["FABRIC_IFACE"] = iface return iface @@ -246,6 +255,7 @@ def main(): " cpu_nic_map.py -p 'ens*' 0 match a different NIC name pattern\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 --provider cxi show the Perlmutter-translated (cxiN) names\n" ), ) parser.add_argument( @@ -258,7 +268,17 @@ def main(): "-p", "--pattern", default="hsn*", - help="glob pattern for NIC names to consider (default: hsn*)", + help="glob pattern for kernel NIC names to consider (default: hsn*) " + "-- always hsn*, on Frontier and Perlmutter alike; see --provider " + "to see the libfabric domain name a given system will actually use", + ) + parser.add_argument( + "--provider", + default=os.environ.get("DDSTORE_FABRIC_PROVIDER", "hsn"), + choices=["hsn", "cxi"], + help="translate printed NIC names to this provider's libfabric " + "domain name (default: $DDSTORE_FABRIC_PROVIDER, or hsn if unset) " + "-- hsn: unchanged (e.g. hsn0); cxi: hsnN -> cxiN (Perlmutter)", ) parser.add_argument( "--env", @@ -274,11 +294,12 @@ def main(): args = parser.parse_args() if args.env: - print(serialize_env(args.pattern)) + print(serialize_env(args.pattern, provider=args.provider)) return if args.allocated: allocated, nics = allocated_nics(args.pattern) + nics = {translate_iface(n, args.provider) for n in nics} print(f"allocated CPUs: {allocated}") print(f"nearest NIC(s): {sorted(nics)}") return @@ -291,12 +312,13 @@ def main(): f"cpu id {args.cpu} not found (valid range: {min(all_pus)}-{max(all_pus)})" ) owner, exact, numa = nearest(args.cpu) - print(owner) + print(translate_iface(owner, args.provider)) 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.provider) print(f"{pu:>4} {numa!s:>4} {owner:>11} {exact}") From bf84628c48b0441d66b6b4ae7d4f9282e3e9a86d Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 17:42:29 -0400 Subject: [PATCH 13/23] Remove -p/--pattern CLI flag from cpu_nic_map.py; kernel NIC names are always hsn* on both systems --- src/cpu_nic_map.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index 4e413eb..bb8a438 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -11,12 +11,18 @@ - 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 --provider cxi (or set DDSTORE_FABRIC_PROVIDER=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 -p 'ens*' 0 match a different NIC name pattern 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 --provider 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 @@ -252,7 +258,6 @@ def main(): "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" - " cpu_nic_map.py -p 'ens*' 0 match a different NIC name pattern\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 --provider cxi show the Perlmutter-translated (cxiN) names\n" @@ -264,14 +269,6 @@ def main(): type=int, help="CPU (PU) id to look up; omit to print the full table", ) - parser.add_argument( - "-p", - "--pattern", - default="hsn*", - help="glob pattern for kernel NIC names to consider (default: hsn*) " - "-- always hsn*, on Frontier and Perlmutter alike; see --provider " - "to see the libfabric domain name a given system will actually use", - ) parser.add_argument( "--provider", default=os.environ.get("DDSTORE_FABRIC_PROVIDER", "hsn"), @@ -294,17 +291,17 @@ def main(): args = parser.parse_args() if args.env: - print(serialize_env(args.pattern, provider=args.provider)) + print(serialize_env(provider=args.provider)) return if args.allocated: - allocated, nics = allocated_nics(args.pattern) + allocated, nics = allocated_nics() nics = {translate_iface(n, args.provider) for n in nics} print(f"allocated CPUs: {allocated}") print(f"nearest NIC(s): {sorted(nics)}") return - all_pus, nearest = build_map(args.pattern) + all_pus, nearest = build_map("hsn*") if args.cpu is not None: if args.cpu not in all_pus: From 2f750e37de8de419cb4f64f1c8a3f53123f2fff8 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 17:57:08 -0400 Subject: [PATCH 14/23] Simplify vae-ddp.py device selection to match vae_extra_train.py (unconditional cuda device, no per-rank set_device) --- examples/vae/ddp_utils.py | 15 --------------- examples/vae/vae-ddp.py | 17 +++-------------- 2 files changed, 3 insertions(+), 29 deletions(-) diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index 472c5cc..f077501 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -31,21 +31,6 @@ def init_comm_size_and_rank(): 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"]) - elif torch.cuda.is_available() and torch.cuda.device_count() > 0: - return rank % torch.cuda.device_count() - return 0 - - def find_ifname(myaddr): """ Find socket ifname for a given ip adress. This is for "GLOO" ddp setup. diff --git a/examples/vae/vae-ddp.py b/examples/vae/vae-ddp.py index 51adcc8..5d63f79 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -20,7 +20,7 @@ import distdataset from distdataset import DistDataset -from ddp_utils import setup_ddp, get_local_rank +from ddp_utils import setup_ddp from vae_model import VAE, loss_function parser = argparse.ArgumentParser(description="VAE MNIST Example") @@ -64,12 +64,7 @@ comm_size, rank = setup_ddp() if args.cuda: - local_rank = get_local_rank(rank) - if torch.cuda.device_count() > 1: - torch.cuda.set_device(local_rank) - device = torch.device(f"cuda:{local_rank}") - else: - device = torch.device("cuda") + device = torch.device("cuda") elif use_mps: device = torch.device("mps") else: @@ -82,13 +77,7 @@ comm.Barrier() model = VAE().to(device) -if args.cuda: - cur_device = torch.cuda.current_device() - model = torch.nn.parallel.DistributedDataParallel( - model, device_ids=[cur_device], output_device=cur_device - ) -else: - model = torch.nn.parallel.DistributedDataParallel(model) +model = torch.nn.parallel.DistributedDataParallel(model) optimizer = optim.Adam(model.parameters(), lr=1e-3) # kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} From ce86fa85cd2c6c8fc5e034766d3d031ed5270cc0 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 17:59:35 -0400 Subject: [PATCH 15/23] Set CUDA device by local-rank id only when multiple GPUs are visible, identically in vae-ddp.py and vae_extra_train.py --- examples/vae/ddp_utils.py | 15 +++++++++++++++ examples/vae/vae-ddp.py | 9 +++++++-- examples/vae/vae_extra_train.py | 12 +++++++++--- 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index f077501..472c5cc 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -31,6 +31,21 @@ def init_comm_size_and_rank(): 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"]) + elif torch.cuda.is_available() and torch.cuda.device_count() > 0: + return rank % torch.cuda.device_count() + return 0 + + def find_ifname(myaddr): """ Find socket ifname for a given ip adress. This is for "GLOO" ddp setup. diff --git a/examples/vae/vae-ddp.py b/examples/vae/vae-ddp.py index 5d63f79..94c8d6d 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -20,7 +20,7 @@ import distdataset from distdataset import DistDataset -from ddp_utils import setup_ddp +from ddp_utils import setup_ddp, get_local_rank from vae_model import VAE, loss_function parser = argparse.ArgumentParser(description="VAE MNIST Example") @@ -64,7 +64,12 @@ comm_size, rank = setup_ddp() 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 use_mps: device = torch.device("mps") else: diff --git a/examples/vae/vae_extra_train.py b/examples/vae/vae_extra_train.py index 72f5d31..a71ab56 100644 --- a/examples/vae/vae_extra_train.py +++ b/examples/vae/vae_extra_train.py @@ -42,7 +42,7 @@ mpi4py.rc.threads = False from mpi4py import MPI -from ddp_utils import setup_ddp +from ddp_utils import setup_ddp, get_local_rank from distdataset import DistDatasetReader from vae_model import VAE, loss_function @@ -95,14 +95,20 @@ torch.manual_seed(args.seed) +comm_size, rank = setup_ddp() + 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 use_mps: device = torch.device("mps") else: device = torch.device("cpu") -comm_size, rank = setup_ddp() print("DDP setup:", comm_size, rank, device) model = VAE().to(device) From b26fdd64b67ba86f8dc9257207aa0348215816de Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 18:26:32 -0400 Subject: [PATCH 16/23] Rename DDSTORE_FABRIC_PROVIDER to DDSTORE_FABRIC (hsn isn't a real libfabric provider name) --- README.md | 8 ++++---- build-evn.sh | 2 +- module-to-load-frontier.sh | 2 +- module-to-load-perlmutter.sh | 2 +- src/common.cxx | 4 ++-- src/cpu_nic_map.py | 10 +++++----- 6 files changed, 14 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index bd37795..1539cdf 100644 --- a/README.md +++ b/README.md @@ -178,7 +178,7 @@ Uses `MPI_Win_create` and `MPI_Get` for one-sided remote reads. Works on any MPI 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_PROVIDER`** selects which libfabric provider to open, for `method=1`/`2`: +**`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. @@ -186,8 +186,8 @@ Uses `fi_read` for true RDMA transfers over high-speed interconnects (Infiniband The two are independent code paths (not runtime auto-detection), so set this explicitly per system rather than relying on a guess: ```bash -export DDSTORE_FABRIC_PROVIDER=hsn # Frontier (default; usually not needed) -export DDSTORE_FABRIC_PROVIDER=cxi # Perlmutter +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: @@ -233,7 +233,7 @@ Environment variables: | `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_PROVIDER` | `hsn` | `hsn` (Frontier) or `cxi` (Perlmutter) — 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. diff --git a/build-evn.sh b/build-evn.sh index 42a7b53..b47aee5 100644 --- a/build-evn.sh +++ b/build-evn.sh @@ -6,7 +6,7 @@ module unload darshan export PYTHONNOUSERSITE=1 export MPICH_GPU_SUPPORT_ENABLED=0 -export DDSTORE_FABRIC_PROVIDER=cxi +export DDSTORE_FABRIC=cxi VENV_DIR=.venv python -m venv --system-site-packages "${VENV_DIR}" diff --git a/module-to-load-frontier.sh b/module-to-load-frontier.sh index 72a2aac..223b7d4 100755 --- a/module-to-load-frontier.sh +++ b/module-to-load-frontier.sh @@ -8,7 +8,7 @@ module load cray-python/3.11.7 module load rocm/7.2.0 export PYTHONNOUSERSITE=1 -export DDSTORE_FABRIC_PROVIDER=hsn +export DDSTORE_FABRIC=hsn ## python env source .venv/bin/activate diff --git a/module-to-load-perlmutter.sh b/module-to-load-perlmutter.sh index fe0b710..556bbad 100644 --- a/module-to-load-perlmutter.sh +++ b/module-to-load-perlmutter.sh @@ -4,7 +4,7 @@ module load pytorch/2.13.0 export PYTHONNOUSERSITE=1 export MPICH_GPU_SUPPORT_ENABLED=0 -export DDSTORE_FABRIC_PROVIDER=cxi +export DDSTORE_FABRIC=cxi ## python env source .venv/bin/activate diff --git a/src/common.cxx b/src/common.cxx index f0f4214..c42c75c 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -520,7 +520,7 @@ static void init_fabric_cxi(struct fabric_state *fabric) if (originfo) fi_freeinfo(originfo); } -/* Dispatch: DDSTORE_FABRIC_PROVIDER selects the fabric-open implementation. +/* 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 @@ -528,7 +528,7 @@ static void init_fabric_cxi(struct fabric_state *fabric) * as two independent, individually-proven implementations instead. */ void init_fabric(struct fabric_state *fabric) { - const char *provider = getenv("DDSTORE_FABRIC_PROVIDER"); + const char *provider = getenv("DDSTORE_FABRIC"); if (provider && strcmp(provider, "cxi") == 0) init_fabric_cxi(fabric); else diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index bb8a438..dd6c4e8 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -14,7 +14,7 @@ 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 --provider cxi (or set DDSTORE_FABRIC_PROVIDER=cxi) to see that +pass --provider cxi (or set DDSTORE_FABRIC=cxi) to see that translated name instead of the raw kernel one. CLI: @@ -192,7 +192,7 @@ def select_fabric_iface(nic_map=None): format), for callers that already have the map from somewhere other than the process environment. - DDSTORE_FABRIC_PROVIDER selects hsn (default) or cxi: + 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 @@ -202,7 +202,7 @@ def select_fabric_iface(nic_map=None): rank's CPU affinity to a NIC (common inside srun tasks with limited PCI visibility). """ - use_cxi = os.environ.get("DDSTORE_FABRIC_PROVIDER") == "cxi" + use_cxi = os.environ.get("DDSTORE_FABRIC") == "cxi" if "FABRIC_IFACE" in os.environ: iface = os.environ["FABRIC_IFACE"] @@ -271,10 +271,10 @@ def main(): ) parser.add_argument( "--provider", - default=os.environ.get("DDSTORE_FABRIC_PROVIDER", "hsn"), + default=os.environ.get("DDSTORE_FABRIC", "hsn"), choices=["hsn", "cxi"], help="translate printed NIC names to this provider's libfabric " - "domain name (default: $DDSTORE_FABRIC_PROVIDER, or hsn if unset) " + "domain name (default: $DDSTORE_FABRIC, or hsn if unset) " "-- hsn: unchanged (e.g. hsn0); cxi: hsnN -> cxiN (Perlmutter)", ) parser.add_argument( From 40bc4ba330aa7966748e7d3ad8d6c15c7c03e903 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Wed, 26 Aug 2026 18:27:43 -0400 Subject: [PATCH 17/23] Rename cpu_nic_map.py CLI flag --provider to --fabric, matching DDSTORE_FABRIC --- src/cpu_nic_map.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index dd6c4e8..a7adbf9 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -14,7 +14,7 @@ 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 --provider cxi (or set DDSTORE_FABRIC=cxi) to see that +pass --fabric cxi (or set DDSTORE_FABRIC=cxi) to see that translated name instead of the raw kernel one. CLI: @@ -22,7 +22,7 @@ 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 --provider cxi show the Perlmutter-translated (cxiN) names + 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 @@ -260,7 +260,7 @@ def main(): " cpu_nic_map.py 42 print only the nearest HSN NIC for cpu 42\n" " export DDSTORE_NIC_MAP=$(cpu_nic_map.py --env) compute once, share via env\n" " srun ... python cpu_nic_map.py --allocated show this task's allocated CPUs + nearest NIC(s)\n" - " cpu_nic_map.py --env --provider cxi show the Perlmutter-translated (cxiN) names\n" + " cpu_nic_map.py --env --fabric cxi show the Perlmutter-translated (cxiN) names\n" ), ) parser.add_argument( @@ -270,10 +270,10 @@ def main(): help="CPU (PU) id to look up; omit to print the full table", ) parser.add_argument( - "--provider", + "--fabric", default=os.environ.get("DDSTORE_FABRIC", "hsn"), choices=["hsn", "cxi"], - help="translate printed NIC names to this provider's libfabric " + 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)", ) @@ -291,12 +291,12 @@ def main(): args = parser.parse_args() if args.env: - print(serialize_env(provider=args.provider)) + print(serialize_env(provider=args.fabric)) return if args.allocated: allocated, nics = allocated_nics() - nics = {translate_iface(n, args.provider) for n in nics} + nics = {translate_iface(n, args.fabric) for n in nics} print(f"allocated CPUs: {allocated}") print(f"nearest NIC(s): {sorted(nics)}") return @@ -309,13 +309,13 @@ def main(): 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.provider)) + 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.provider) + owner = translate_iface(owner, args.fabric) print(f"{pu:>4} {numa!s:>4} {owner:>11} {exact}") From befa9eaf341c2ec6d4c4c9dfe30e075fa4f65bff Mon Sep 17 00:00:00 2001 From: Jong Youl Choi Date: Thu, 27 Aug 2026 03:55:56 +0000 Subject: [PATCH 18/23] fix for aurora --- build-evn.sh | 30 ----------------------------- examples/vae/ddp_utils.py | 34 ++++++++++++++++++++++++++------- examples/vae/vae-ddp.py | 7 +++++++ examples/vae/vae_extra_train.py | 8 +++++++- module-to-load-perlmutter.sh | 0 5 files changed, 41 insertions(+), 38 deletions(-) delete mode 100644 build-evn.sh mode change 100644 => 100755 module-to-load-perlmutter.sh diff --git a/build-evn.sh b/build-evn.sh deleted file mode 100644 index b47aee5..0000000 --- a/build-evn.sh +++ /dev/null @@ -1,30 +0,0 @@ -#!/bin/bash - -# Load system modules -module load pytorch/2.13.0 -module unload darshan - -export PYTHONNOUSERSITE=1 -export MPICH_GPU_SUPPORT_ENABLED=0 -export DDSTORE_FABRIC=cxi - -VENV_DIR=.venv -python -m venv --system-site-packages "${VENV_DIR}" - -source .venv/bin/activate -pip install --upgrade pip - -# Install build and runtime dependencies -pip install wheel Cython -MPICC="cc -shared" pip install --no-cache-dir --no-binary=mpi4py mpi4py==4.1.1 -pip install pytest pytest-mpi -pip install psutil - -# Build and install PyDDStore using Cray compiler wrappers -CC=cc CXX=CC pip install -e . - -# Install PyTorch with ROCm support -pip3 install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.2 - -echo "Environment is ready at '.venv'." -echo "To activate: source .venv/bin/activate" diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index 472c5cc..919f47e 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -1,3 +1,6 @@ +from mpi4py import MPI +import os, socket + import os import re import socket @@ -23,6 +26,11 @@ def init_comm_size_and_rank(): ## 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: @@ -41,8 +49,6 @@ def get_local_rank(rank): return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) elif os.getenv("SLURM_LOCALID") is not None: return int(os.environ["SLURM_LOCALID"]) - elif torch.cuda.is_available() and torch.cuda.device_count() > 0: - return rank % torch.cuda.device_count() return 0 @@ -102,20 +108,23 @@ def parse_slurm_nodelist(nodelist): def setup_ddp(): """ "Initialize DDP""" - if os.getenv("HYDRAGNN_BACKEND") is not None: - backend = os.environ["HYDRAGNN_BACKEND"] + 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", "8889") + master_port = os.getenv("MASTER_PORT", "2345") if os.getenv("LSB_HOSTS") is not None: master_addr = os.environ["LSB_HOSTS"].split()[1] @@ -125,11 +134,22 @@ def setup_ddp(): 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"]: + if backend in ["nccl", "gloo", "xccl"]: os.environ["MASTER_ADDR"] = master_addr - os.environ["MASTER_PORT"] = master_port + os.environ["MASTER_PORT"] = str(master_port) os.environ["WORLD_SIZE"] = str(world_size) os.environ["RANK"] = str(world_rank) diff --git a/examples/vae/vae-ddp.py b/examples/vae/vae-ddp.py index 94c8d6d..f7bed6e 100644 --- a/examples/vae/vae-ddp.py +++ b/examples/vae/vae-ddp.py @@ -62,6 +62,7 @@ comm = MPI.COMM_WORLD comm_size, rank = setup_ddp() +local_rank = get_local_rank(rank) if args.cuda: if torch.cuda.device_count() > 1: @@ -70,6 +71,12 @@ 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: diff --git a/examples/vae/vae_extra_train.py b/examples/vae/vae_extra_train.py index a71ab56..7cca4aa 100644 --- a/examples/vae/vae_extra_train.py +++ b/examples/vae/vae_extra_train.py @@ -96,14 +96,20 @@ 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: - 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(localrank) + device = torch.device(f"xpu:{local_rank}") + else: + device = torch.device("xpu") elif use_mps: device = torch.device("mps") else: diff --git a/module-to-load-perlmutter.sh b/module-to-load-perlmutter.sh old mode 100644 new mode 100755 From 1e91ae4ce2247aea6cd376cf751c5df6ac31f88e Mon Sep 17 00:00:00 2001 From: Jong Youl Choi Date: Thu, 27 Aug 2026 14:00:43 +0000 Subject: [PATCH 19/23] fix for aurora --- src/cpu_nic_map.py | 102 ++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 92 insertions(+), 10 deletions(-) diff --git a/src/cpu_nic_map.py b/src/cpu_nic_map.py index a7adbf9..4682b2c 100644 --- a/src/cpu_nic_map.py +++ b/src/cpu_nic_map.py @@ -58,6 +58,34 @@ def hcalc(loc, itype): 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) @@ -68,17 +96,44 @@ def build_map(pattern): if not nics: sys.exit(f"no NICs matching '{pattern}' found under /sys/class/net") - nic_closest = {n: set(hcalc(f"os={n}", "PU")) for n in nics} - nic_numa = {n: (hcalc(f"os={n}", "NUMA") or [None])[0] for n in nics} + 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( - "hwloc-calc found no PUs local to any NIC (os= lookups came back " - "empty) -- this process's hwloc topology view has no visibility into the " - "NICs, so any nearest-NIC answer would be a silent guess, not real data. " - "This has been observed when running directly inside a plain `srun` task; " - "try running from the sbatch batch step's own shell instead." + "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 = {} @@ -176,6 +231,19 @@ def allocated_nics(pattern="hsn*", nic_map=None): 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 @@ -239,11 +307,13 @@ def select_fabric_iface(nic_map=None): f"around this" ) iface = sorted(nics)[0] - if len(nics) > 1: - print(f"FABRIC_IFACE: affinity spans {sorted(nics)}, picking {iface}") - 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 @@ -317,6 +387,18 @@ def main(): 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__": From 3ea9ff1baf59ecb1d2454be61640ad474b7e025b Mon Sep 17 00:00:00 2001 From: Jong Youl Choi Date: Thu, 27 Aug 2026 13:02:55 -0700 Subject: [PATCH 20/23] remove --- module-to-load-frontier.sh | 15 --------------- module-to-load-perlmutter.sh | 10 ---------- 2 files changed, 25 deletions(-) delete mode 100755 module-to-load-frontier.sh delete mode 100755 module-to-load-perlmutter.sh diff --git a/module-to-load-frontier.sh b/module-to-load-frontier.sh deleted file mode 100755 index 223b7d4..0000000 --- a/module-to-load-frontier.sh +++ /dev/null @@ -1,15 +0,0 @@ -#!/bin/bash - -module load PrgEnv-gnu -module load cpe/24.11 -module load libfabric -module unload darshan-runtime -module load cray-python/3.11.7 -module load rocm/7.2.0 - -export PYTHONNOUSERSITE=1 -export DDSTORE_FABRIC=hsn - -## python env -source .venv/bin/activate - diff --git a/module-to-load-perlmutter.sh b/module-to-load-perlmutter.sh deleted file mode 100755 index 556bbad..0000000 --- a/module-to-load-perlmutter.sh +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/bash - -module load pytorch/2.13.0 - -export PYTHONNOUSERSITE=1 -export MPICH_GPU_SUPPORT_ENABLED=0 -export DDSTORE_FABRIC=cxi - -## python env -source .venv/bin/activate From 5a2089b1d84bcc30684479a903e482eb67ac9865 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:46:05 +0000 Subject: [PATCH 21/23] Fix review comments: extern C linkage, localrank typo, README docs, path traversal, comm validation Co-authored-by: jychoi-hpc <3661063+jychoi-hpc@users.noreply.github.com> --- README.md | 2 +- examples/vae/vae_extra_train.py | 2 +- src/common.cxx | 22 +++++++++++++++++----- src/pyddstore.pyx | 4 ++++ 4 files changed, 23 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 1539cdf..64fe9d4 100644 --- a/README.md +++ b/README.md @@ -144,7 +144,7 @@ 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 every core rank's record file appears (up to `DDSTORE_HANDSHAKE_TIMEOUT_S` seconds), then registers it for `get()`. +`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 | |---|---|---| diff --git a/examples/vae/vae_extra_train.py b/examples/vae/vae_extra_train.py index 7cca4aa..9cda3a2 100644 --- a/examples/vae/vae_extra_train.py +++ b/examples/vae/vae_extra_train.py @@ -106,7 +106,7 @@ device = torch.device("cuda") elif hasattr(torch, "xpu") and torch.xpu.is_available(): if torch.xpu.device_count() > 1: - torch.xpu.set_device(localrank) + torch.xpu.set_device(local_rank) device = torch.device(f"xpu:{local_rank}") else: device = torch.device("xpu") diff --git a/src/common.cxx b/src/common.cxx index c42c75c..1e16dc6 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -740,7 +740,7 @@ int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset /* Resolve the handshake directory. * Returns a pointer to a static buffer — copy before next call. */ -const char *resolve_handshake_dir(const char *user_dir) +extern "C" const char *resolve_handshake_dir(const char *user_dir) { static char resolved[4096]; @@ -763,11 +763,23 @@ const char *resolve_handshake_dir(const char *user_dir) return resolved; } -/* Build the canonical path for a variable's combined record file into `buf`. */ +/* 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) { - snprintf(buf, bufsz, "%s/%s.bin", dir, 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). */ @@ -789,7 +801,7 @@ static int handshake_timeout_s(void) * the combined array to {resolved_dir}/{varname}.bin (tmp + fsync + rename) * so extra members can join later. * -------------------------------------------------------------------------- */ -int handshake_write(struct fabric_state *fs, MPI_Comm comm, +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) @@ -916,7 +928,7 @@ int handshake_write(struct fabric_state *fs, MPI_Comm comm, * * Returns 0 on success, non-zero on error or timeout. * -------------------------------------------------------------------------- */ -int handshake_join(struct fabric_state *fs, +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) diff --git a/src/pyddstore.pyx b/src/pyddstore.pyx index 31566ad..f85b976 100644 --- a/src/pyddstore.pyx +++ b/src/pyddstore.pyx @@ -105,6 +105,10 @@ cdef class PyDDStore: 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) From e074d41f994dd6e23b4d21ab4451159e997fcfb6 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Thu, 27 Aug 2026 21:44:42 -0400 Subject: [PATCH 22/23] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- examples/vae/ddp_utils.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/vae/ddp_utils.py b/examples/vae/ddp_utils.py index 919f47e..12b4e69 100644 --- a/examples/vae/ddp_utils.py +++ b/examples/vae/ddp_utils.py @@ -1,6 +1,3 @@ -from mpi4py import MPI -import os, socket - import os import re import socket From d6263fa437eceb91696f51e012851a6b51bb4db5 Mon Sep 17 00:00:00 2001 From: Jong Choi Date: Thu, 27 Aug 2026 21:45:09 -0400 Subject: [PATCH 23/23] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/common.cxx | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/common.cxx b/src/common.cxx index 1e16dc6..3885c3e 100644 --- a/src/common.cxx +++ b/src/common.cxx @@ -649,8 +649,18 @@ int read_from_remote(struct fabric_state *fabric_state, int src, uint64_t offset * hsn/verbs/gni/psm2 (is_mr_endpoint() is false for those). */ if (is_mr_endpoint(fabric_state)) { - fi_mr_bind(fabric_state->recv_mr, &fabric_state->signal->fid, 0); - fi_mr_enable(fabric_state->recv_mr); + 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;