File-based handshake

[jychoi-hpc]

Adding file-based handshake

[Copilot]

Pull request overview

Adds a new “method=2” file-based handshake mode to DDStore to enable RDMA reads between two independent jobs (core publishers + extra readers) that do not share an MPI communicator, plus supporting infrastructure for libfabric interface selection and updated documentation/examples.

Changes:

• Implement method=2 in C++ (core publishes combined record files; extra member joins via polling + RDMA reads) and expose it in the Python bindings (PyDDStore(..., method=2, handshake_dir=..., n_core=...)) with new join() and info().
• Add CPU→NIC mapping utility (cpu_nic_map.py) and integrate automatic FABRIC_IFACE selection for libfabric backends.
• Update README and VAE/demo scripts/tests to demonstrate the new workflow and improve formatting.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
test/test_single.py	Formatting-only updates (blank lines / parametrize layout).
test/test_multirank.py	Formatting-only updates; one expected-value expression reformatted.
test/test_method2_extra.py	New standalone extra-member script to join/get via method=2 and write sentinel.
test/test_method2_core.py	New standalone core-member script to publish via method=2 and wait for sentinel.
test/conftest.py	Minor formatting (blank line).
src/pyddstore.pyx	Python/Cython API extended for method=2 + NIC selection + new join/info.
src/ddstore.cxx	Method=2 constructors, join() implementation, and method=2 cleanup integration.
src/cpu_nic_map.py	New utility for CPU affinity → NIC selection and env serialization.
src/common.cxx	Adds CXI provider path + file-based handshake helpers (write/join/dir resolve).
setup.py	Packages cpu_nic_map alongside the extension; formatting cleanup.
README.md	Documents method=2, join/info, fabric selection, and updated grouping guidance.
include/ddstore.hpp	Method=2 fields/constructors, join/size APIs, and method=2 add/init/get paths.
include/common.h	Adds CoreRecord + method=2 handshake APIs and CXI-related helpers.
examples/vae/vae-ddp.py	Refactor: extract model/DDP helpers, fix import ordering, update behavior.
examples/vae/vae_model.py	New extracted VAE model + loss function module.
examples/vae/vae_extra_train.py	New method=2 “extra” training script using DistDatasetReader + sentinel.
examples/vae/vae_core_server.py	New method=2 “core” data server script publishing MNIST via DDStore.
examples/vae/distdataset.py	Method=2 support + reader class; fixes data shaping for disp inference.
examples/vae/ddp_utils.py	New shared DDP initialization utilities.
examples/scripts/test.py	Formatting-only CLI argument layout changes.
examples/scripts/demo.py	Formatting-only CLI argument layout changes.

Suppressed comments (2)

src/common.cxx:795

• handshake_write() is declared inside an extern "C" block in include/common.h, but the definition here is missing extern "C". This can cause a linkage mismatch (undefined reference / conflicting declaration) when called from other translation units.

 * -------------------------------------------------------------------------- */
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)

src/common.cxx:922

• handshake_join() is declared with C linkage in include/common.h, but the definition here lacks extern "C". This can break linking (name mangling mismatch) or cause a conflicting linkage error because src/common.cxx includes common.h.

 * -------------------------------------------------------------------------- */
int handshake_join(struct fabric_state *fs,
                   const char *dir, const char *varname,
                   int n_core,
                   long *lenlist, int *out_disp, int *out_itemsize)

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 21 out of 21 changed files in this pull request and generated 2 comments.

Suppressed comments (4)

Previously missed (3) — in code that hasn't changed since the last review.

src/ddstore.cxx:155

• DDStore::join() allocates and partially initializes a fabric_state, but throws on init_fabric()/handshake_join() failure without releasing any allocated libfabric resources (and the fabric_state itself). This leaks endpoints, domains, and memory, and will accumulate across retries or multiple join attempts.

    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<long> 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);

src/ddstore.cxx:193

• DDStore::free() releases MPI windows / libfabric state but never frees the data buffers allocated via MPI_Alloc_mem() in add()/init(), so long-running processes that create/free stores will leak potentially large allocations. Free the per-variable base pointer once the window/MR is released.

This issue also appears on line 200 of the same file.

        for (auto &x : this->varlist)
        {
            if (x.second.active)
            {
                MPI_Win_free(&x.second.win);
            }
            x.second.active = false;

src/pyddstore.pyx:87

• PyDDStore.cinit calls cpu_nic_map.select_fabric_iface() before validating method=2 parameters. If construction fails (e.g. missing handshake_dir or n_core), this still mutates FABRIC_IFACE in the process environment, creating surprising side effects for callers. Validate inputs first, then perform NIC selection.

        cdef MPI.Comm mpi_comm
        if method != 0:
            cpu_nic_map.select_fabric_iface(nic_map=nic_map)
        if method == 2:

src/ddstore.cxx:216

• For method=1/2, DDStore::free() closes fabric resources but also never frees the MPI_Alloc_mem() buffer stored in VarInfo_t::base. After closing the MR (fs->mr) and before discarding the VarInfo, the backing buffer should be released to avoid large leaks in repeated publish/free cycles.

            if (x.second.active && x.second.fabric_state)
            {
                struct fabric_state *fs = x.second.fabric_state;
                if (fs->recv_mr)   fi_close(&fs->recv_mr->fid);
                if (fs->mr)        fi_close(&fs->mr->fid);
                if (fs->signal)    fi_close(&fs->signal->fid);
                if (fs->cq_signal) fi_close(&fs->cq_signal->fid);
                if (fs->av)        fi_close(&fs->av->fid);
                if (fs->domain)    fi_close(&fs->domain->fid);
                if (fs->fabric)    fi_close(&fs->fabric->fid);
                if (fs->info)      fi_freeinfo(fs->info);
                if (fs->ctx)       ::free(fs->ctx);
                ::free(fs->comm_partner);
                ::free(fs->remote_key);
                ::free(fs->remote_address);
                ::free(fs);
                x.second.fabric_state = NULL;