File-based handshake

README.md

@@ -11,7 +11,7 @@ Each MPI rank holds a shard of the full dataset in memory. DDStore exposes a glo
 | Dependency | Notes |
 |---|---|
 | MPI (OpenMPI / MPICH) | `mpicc` and `mpicxx` must be on `PATH` |
-| libfabric | Required for RDMA backend (`method=1`) |
+| libfabric | Required for the RDMA backends (`method=1` and `method=2`) |
 | Python ≥ 3.6 | |
 | NumPy, mpi4py, Cython | Python build dependencies |
 
@@ -71,13 +71,26 @@ mpirun -n 4 python my_script.py
 
 ## API Reference
 
-### `PyDDStore(comm, method=0, ddstore_width=None)`
+### `PyDDStore(comm_or_none=None, method=0, handshake_dir="", n_core=0, nic_map=None)`
 
 | Parameter | Type | Description |
 |---|---|---|
-| `comm` | `mpi4py.MPI.Comm` | MPI communicator covering all ranks |
-| `method` | `int` | `0` = MPI RMA (default), `1` = libfabric RDMA |
-| `ddstore_width` | `int` or `None` | Ranks per DDStore group. `None` uses all ranks in `comm` as a single group |
+| `comm_or_none` | `mpi4py.MPI.Comm` or `None` | MPI communicator covering all ranks. `None` only for a `method=2` extra member |
+| `method` | `int` | `0` = MPI RMA (default), `1` = libfabric RDMA, `2` = file-based handshake (see [below](#file-based-handshake-method2)) |
+| `handshake_dir` | `str` | Required for `method=2`: shared-filesystem directory used to exchange RDMA addresses |
+| `n_core` | `int` | Required for a `method=2` extra member: number of core ranks that published data |
+| `nic_map` | `str` or `None` | Optional, `method=1`/`2` only: a precomputed CPU→NIC map string (see [`DDSTORE_NIC_MAP`](#libfabric-rdma-method1) below) to use instead of the environment variable. Ignored if `FABRIC_IFACE` is already set |
+
+Four call shapes:
+
+```python
+PyDDStore(comm)                                     # method 0, MPI RMA
+PyDDStore(comm, method=1)                            # method 1, libfabric RDMA
+PyDDStore(comm, method=2, handshake_dir="/path")      # method 2, core member (n_core == comm size)
+PyDDStore(None, method=2, handshake_dir="/path", n_core=N)  # method 2, extra member (no comm)
+```
+
+Note: grouping ranks into independent stores (the "sub-communicator" pattern below) is done by splitting `comm` yourself before constructing `PyDDStore` — there is no `ddstore_width` constructor parameter. `DistDataset` in [examples/vae/distdataset.py](examples/vae/distdataset.py) shows the pattern (`comm.Split()` then `PyDDStore(sub_comm)`).
 
 ---
 
@@ -129,6 +142,22 @@ Read `arr.shape[0]` consecutive rows starting at global index `start` into `arr`
 
 ---
 
+### `join(name)`
+
+`method=2` extra member only. Discovers a variable published by the core group by polling the handshake directory until the combined record file (`{name}.bin`) written by core rank 0 reaches its expected size (up to `DDSTORE_HANDSHAKE_TIMEOUT_S` seconds), then registers it for `get()`.
+
+| Parameter | Type | Description |
+|---|---|---|
+| `name` | `str` | Variable identifier, matching the `name` used in the core group's `add()` |
+
+---
+
+### `info(name)`
+
+Returns `(total_rows, disp, itemsize)` for a variable that has been `add()`-ed or `join()`-ed. Useful on the extra side to size output buffers without hardcoding shapes.
+
+---
+
 ### `epoch_begin()` / `epoch_end()`
 
 Open and close an MPI RMA access epoch (calls `MPI_Win_fence`). **Collective**. Required around `get()` calls when using `method=0`. No-op for `method=1`.
@@ -147,32 +176,90 @@ Uses `MPI_Win_create` and `MPI_Get` for one-sided remote reads. Works on any MPI
 
 ### libfabric RDMA (`method=1`)
 
-Uses `fi_read` for true RDMA transfers over high-speed interconnects (Infiniband/verbs, Cray GNI, Intel PSM2). Lower latency than MPI RMA on supported hardware. `epoch_begin`/`epoch_end` are no-ops with this backend.
+Uses `fi_read` for true RDMA transfers over high-speed interconnects (Infiniband/verbs, Cray GNI, Intel PSM2, Cray Slingshot). Lower latency than MPI RMA on supported hardware. `epoch_begin`/`epoch_end` are no-ops with this backend.
+
+**`DDSTORE_FABRIC`** selects which libfabric provider to open, for `method=1`/`2`:
+
+- `hsn` (default, unset) — Frontier: opens the `tcp;ofi_rxm` domain over Cray Slingshot.
+- `cxi` — Perlmutter: opens the native `cxi` domain over Cray Slingshot.
+
+The two are independent code paths (not runtime auto-detection), so set this explicitly per system rather than relying on a guess:
 
-Set `FABRIC_IFACE` to select a specific network interface when the automatic selection picks the wrong one:
 ```bash
-export FABRIC_IFACE=hsn0   # e.g. Cray Slingshot
+export DDSTORE_FABRIC=hsn   # Frontier (default; usually not needed)
+export DDSTORE_FABRIC=cxi   # Perlmutter
 ```
 
+`PyDDStore` picks the network interface (`FABRIC_IFACE`) automatically for `method=1`/`2`, based on each rank's real CPU affinity (`os.sched_getaffinity`) — no changes needed in your code:
+
+- **`DDSTORE_NIC_MAP`** — a precomputed CPU→NIC map, used directly if set (no NIC discovery at construction time). Generate it once from a context with reliable NIC visibility, e.g. an `sbatch` batch step's own shell (not a nested `srun` task — NIC/PCI discovery has been observed to fail there), and export it before launching ranks so every one inherits it:
+  ```bash
+  export DDSTORE_NIC_MAP=$(python3 -m cpu_nic_map --env)
+  srun ... python train.py
+  ```
+- If `DDSTORE_NIC_MAP` isn't set, each rank falls back to a live `hwloc-calc`/`lstopo` query against its own CPU affinity (`cpu_nic_map.allocated_nics()`, also runnable standalone as `python3 cpu_nic_map.py --allocated`) to find the nearest NIC.
+- Set `FABRIC_IFACE` explicitly to override both and force a specific interface, e.g. when the automatic selection picks the wrong one:
+  ```bash
+  export FABRIC_IFACE=hsn0   # e.g. Cray Slingshot
+  ```
+- Or skip the environment entirely and pass a map straight to the constructor: `PyDDStore(comm, method=1, nic_map="hsn0=0-15,64-79;hsn1=...")`.
+
+### File-based handshake (`method=2`)
+
+Splits the dataset-holding job from the training job entirely: a **core** group loads and publishes data, and a separate **extra** group reads it over RDMA (`fi_read`, same transport as `method=1`) — the two are independent MPI jobs (e.g. two separate `srun`/`mpirun` launches, possibly on different node allocations) that never share a communicator. They rendezvous only through record files written to a shared-filesystem directory (must be visible to all nodes, e.g. Lustre):
+
+- **Core member** — has an MPI communicator, publishes with `add()`/`init()`. Core ranks exchange records via `MPI_Allgather`, and rank 0 writes the combined set to a single `{name}.bin` file (fabric address, MR key, base pointer, row count, dtype per rank) into `handshake_dir`.
+- **Extra member** — no MPI communicator; constructed with `comm_or_none=None` and an explicit `n_core`. Calls `join(name)` to poll for and read all `n_core` core-rank records, then `get()` works exactly as on the core side, reading directly from core-rank memory over RDMA.
+
+```python
+# core side — one MPI job
+store = dds.PyDDStore(comm, method=2, handshake_dir="/lustre/.../ddstore_hs")
+store.add("x", data)
+...                       # wait for the extra side to finish (e.g. a sentinel file)
+store.free()
+
+# extra side — a separate MPI job, no comm needed
+store = dds.PyDDStore(None, method=2, handshake_dir="/lustre/.../ddstore_hs", n_core=4)
+store.join("x")
+out = np.zeros((1, ncols), dtype=np.float32)
+store.get("x", out, start=global_idx)
+store.free()
+```
+
+Environment variables:
+
+| Variable | Default | Description |
+|---|---|---|
+| `DDSTORE_HANDSHAKE_DIR` | `./ddstore_hs` | Shared directory for handshake record files |
+| `DDSTORE_HANDSHAKE_TIMEOUT_S` | `300` | Seconds to poll for core records / a join before raising a timeout |
+| `DDSTORE_NIC_MAP` | unset | CPU→NIC map for `FABRIC_IFACE` auto-selection — see [libfabric RDMA](#libfabric-rdma-method1) above |
+| `DDSTORE_FABRIC` | `hsn` | `hsn` (Frontier) or `cxi` (Perlmutter) — see [libfabric RDMA](#libfabric-rdma-method1) above |
+
+See [test/test_method2_core.py](test/test_method2_core.py) / [test/test_method2_extra.py](test/test_method2_extra.py) for a minimal runnable pair, and [examples/vae/vae_core_server.py](examples/vae/vae_core_server.py) / [examples/vae/vae_extra_train.py](examples/vae/vae_extra_train.py) for a full DDP training example using this split.
+
+`ddstore_width` grouping (below) is not currently supported with `method=2` — every core rank in `comm` is treated as one group.
+
 ## Partitioned / Sub-communicator Usage
 
-`ddstore_width` controls how many MPI ranks form a single DDStore group. The global communicator is split so that each group of `ddstore_width` consecutive ranks shares one independent store, with each group holding a full replica of the dataset partitioned across its members.
+`PyDDStore` itself always spans the full communicator you pass it — there is no built-in "ranks per group" option. To run several independent stores side by side (e.g. one per node), split `comm` yourself before constructing `PyDDStore`, giving each group its own sub-communicator. Each group then holds a full replica of the dataset, partitioned across its own members.
 
-**Example — 16 ranks, `ddstore_width=4`:**
+**Example — 16 ranks split into groups of 4:**
 ```
 ranks  0– 3  →  DDStore group 0
 ranks  4– 7  →  DDStore group 1
 ranks  8–11  →  DDStore group 2
 ranks 12–15  →  DDStore group 3
 ```
 
-This is useful when you want one store per node (e.g. 4 GPUs per node → `ddstore_width=4`), limiting cross-node RDMA traffic to the dataset replication step at startup rather than every sample fetch.
+This is useful when you want one store per node (e.g. 4 GPUs per node), limiting cross-node RDMA traffic to the dataset replication step at startup rather than every sample fetch.
 
 ```python
-store = dds.PyDDStore(comm, ddstore_width=4)   # e.g. 4 GPUs per store
+width = 4                                       # ranks per group, e.g. GPUs per node
+sub_comm = comm.Split(rank // width, rank)
+store = dds.PyDDStore(sub_comm)                  # one independent store per group
 ```
 
-If `ddstore_width` is omitted, all ranks in `comm` form a single store.
+`DistDataset` in [examples/vae/distdataset.py](examples/vae/distdataset.py) wraps exactly this pattern behind a `ddstore_width` constructor argument — pass `ddstore_width=None` (default) for a single store across all ranks in `comm`, or an integer to split into groups of that size.
 
 ## PyTorch Dataset Integration
 
@@ -227,6 +314,18 @@ Optional arguments for `examples/scripts/demo.py` and `examples/scripts/test.py`
 | `--dim` | `64` | Elements per row |
 | `--nbatch` | `32` | Number of random reads |
 
+### Method 2 (file-based handshake)
+
+Two separate launches sharing a handshake directory on a shared filesystem — not a single `mpirun`, since core and extra are independent jobs:
+
+```bash
+# Terminal 1 — core (data-holding) side
+mpirun -n 4 python test/test_method2_core.py /path/to/shared/ddstore_hs
+
+# Terminal 2 — extra (reader) side, after or while the core side is running
+python test/test_method2_extra.py /path/to/shared/ddstore_hs 4
+```
+
 ## Citation
 
 If you use DDStore in your research, please cite:

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

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()
 

examples/vae/ddp_utils.py

@@ -0,0 +1,169 @@
+import os
+import re
+import socket
+
+import psutil
+import torch
+import torch.distributed as dist
+
+"""
+Functions for DDP on HPC
+"""
+
+
+def init_comm_size_and_rank():
+    world_size = None
+    world_rank = 0
+
+    if os.getenv("OMPI_COMM_WORLD_SIZE") and os.getenv("OMPI_COMM_WORLD_RANK"):
+        ## Summit
+        world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
+        world_rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
+    elif os.getenv("SLURM_NPROCS") and os.getenv("SLURM_PROCID"):
+        ## CADES
+        world_size = int(os.environ["SLURM_NPROCS"])
+        world_rank = int(os.environ["SLURM_PROCID"])
+    else:
+        from mpi4py import MPI
+
+        world_size = MPI.COMM_WORLD.Get_size()
+        world_rank = MPI.COMM_WORLD.Get_rank()
+
+    ## Fall back to default
+    if world_size is None:
+        world_size = 1
+
+    return int(world_size), int(world_rank)
+
+
+def get_local_rank(rank):
+    """
+    Determine which GPU on the local node this rank should use.
+    Falls back to rank % device_count when no launcher-provided local rank
+    is available (e.g. plain mpirun without per-rank GPU visibility).
+    """
+    if os.getenv("OMPI_COMM_WORLD_LOCAL_RANK") is not None:
+        return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
+    elif os.getenv("SLURM_LOCALID") is not None:
+        return int(os.environ["SLURM_LOCALID"])
+    return 0
+
+
+def find_ifname(myaddr):
+    """
+    Find socket ifname for a given ip adress. This is for "GLOO" ddp setup.
+    Usage example:
+        find_ifname("127.0.0.1") will return a network interface name, such as "lo". "lo0", etc.
+    """
+    ipaddr = socket.gethostbyname(myaddr)
+    ifname = None
+    for nic, addrs in psutil.net_if_addrs().items():
+        for addr in addrs:
+            if addr.address == ipaddr:
+                ifname = nic
+                break
+        if ifname is not None:
+            break
+
+    return ifname
+
+
+def parse_slurm_nodelist(nodelist):
+    """
+    Parse SLURM_NODELIST env string to get list of nodes.
+    Usage example:
+        parse_slurm_nodelist(os.environ["SLURM_NODELIST"])
+    Input examples:
+        "or-condo-g04"
+        "or-condo-g[05,07-08,13]"
+        "or-condo-g[05,07-08,13],or-condo-h[01,12]"
+    """
+    nlist = list()
+    for block, _ in re.findall(r"([\w-]+(\[[\d\-,]+\])*)", nodelist):
+        m = re.match(r"^(?P<prefix>[\w\-]+)\[(?P<group>.*)\]", block)
+        if m is None:
+            ## single node
+            nlist.append(block)
+        else:
+            ## multiple nodes
+            g = m.groups()
+            prefix = g[0]
+            for sub in g[1].split(","):
+                if "-" in sub:
+                    start, end = re.match(r"(\d+)-(\d+)", sub).groups()
+                    fmt = "%%0%dd" % (len(start))
+                    for i in range(int(start), int(end) + 1):
+                        node = prefix + fmt % i
+                        nlist.append(node)
+                else:
+                    node = prefix + sub
+                    nlist.append(node)
+
+    return nlist
+
+
+def setup_ddp():
+    """ "Initialize DDP"""
+
+    if os.getenv("DDSTORE_BACKEND") is not None:
+        backend = os.environ["DDSTORE_BACKEND"]
+    elif dist.is_nccl_available() and torch.cuda.is_available():
+        backend = "nccl"
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        backend = "xccl"
+    elif torch.distributed.is_gloo_available():
+        backend = "gloo"
+    else:
+        raise RuntimeError("No parallel backends available")
+
+    world_size, world_rank = init_comm_size_and_rank()
+    print(f"DDP: Hi from rank {world_rank} of {world_size}.")
+
+    ## Default setting
+    master_addr = "127.0.0.1"
+    master_port = os.getenv("MASTER_PORT", "2345")
+
+    if os.getenv("LSB_HOSTS") is not None:
+        master_addr = os.environ["LSB_HOSTS"].split()[1]
+    elif os.getenv("LSB_MCPU_HOSTS") is not None:
+        master_addr = os.environ["LSB_MCPU_HOSTS"].split()[2]
+    elif os.getenv("SLURM_STEP_NODELIST") is not None:
+        master_addr = parse_slurm_nodelist(os.environ["SLURM_STEP_NODELIST"])[0]
+    elif os.getenv("SLURM_NODELIST") is not None:
+        master_addr = parse_slurm_nodelist(os.environ["SLURM_NODELIST"])[0]
+    elif os.getenv("PBS_O_HOST") is not None:
+        if os.environ["PBS_O_HOST"][-19:] == "aurora.alcf.anl.gov":
+            from mpi4py import MPI
+
+            RANK = MPI.COMM_WORLD.Get_rank()
+            MASTER_ADDR = socket.gethostname() if RANK == 0 else None
+            MASTER_ADDR = MPI.COMM_WORLD.bcast(MASTER_ADDR, root=0)
+            master_addr = f"{MASTER_ADDR}.hsn.cm.aurora.alcf.anl.gov"
+        else:
+            ## The following is CADES specific
+            master_addr = parse_slurm_nodelist(os.environ["PBS_O_HOST"])[0]
+
+    try:
+        if backend in ["nccl", "gloo", "xccl"]:
+            os.environ["MASTER_ADDR"] = master_addr
+            os.environ["MASTER_PORT"] = str(master_port)
+            os.environ["WORLD_SIZE"] = str(world_size)
+            os.environ["RANK"] = str(world_rank)
+
+        if (backend == "gloo") and ("GLOO_SOCKET_IFNAME" not in os.environ):
+            ifname = find_ifname(master_addr)
+            if ifname is not None:
+                os.environ["GLOO_SOCKET_IFNAME"] = ifname
+
+        print(
+            "Distributed data parallel: %s master at %s:%s"
+            % (backend, master_addr, master_port),
+        )
+
+        if not dist.is_initialized():
+            dist.init_process_group(backend=backend, init_method="env://")
+
+    except KeyError:
+        print("DDP has to be initialized within a job - Running in sequential mode")
+
+    return world_size, world_rank