refactor(umbp): make distributed mode backend- and transport-agnostic

.claude/skills/umbp-add-transfer-engine/SKILL.md

@@ -0,0 +1,327 @@
+---
+name: umbp-add-transfer-engine
+description: >-
+  Add a new byte-moving path to UMBP's distributed mode as a TransferEngine
+  (GPU copy, a second RDMA transport, a file/GDS engine, a fake for tests).
+  Covers CanHandle pair-dispatch and why engines must be disjoint, the
+  Plan/Submit/Wait split, bounce buffers, PeerDirectory for transports that need
+  a handshake, and how to extend TransferRef for a non-memory endpoint. Use when
+  the user wants to add or debug a transfer engine in src/umbp/distributed/
+  transfer, mentions TransferEngine, CanHandle, TransferRef, MoriIoEngine,
+  LocalCopyEngine or CompositeTransferEngine, or asks why a transfer is
+  "unplannable" / no engine claimed a pair.
+---
+
+# Adding a transfer engine to UMBP distributed mode
+
+**There is exactly one byte-moving path in the system and it is the transfer
+layer.** A backend owns bytes and publishes endpoints; an engine decides how to
+move bytes between two endpoints. Admitting a second byte-moving path at the
+backend level is the coupling the whole backend-agnostic refactor exists to
+remove.
+
+The companion skill `umbp-add-backend` covers the other side of that boundary.
+
+## The one idea: selection is a function of the PAIR
+
+Never of either endpoint alone. This is the rule that makes everything else
+work, and the most common way to get an engine wrong is to write `CanHandle` as
+a property of `src` or `dst` in isolation.
+
+The engines in tree partition the `(src, dst)` space with **no overlap**:
+
+| pair | engine |
+|---|---|
+| both local, both `loc == CPU` | `LocalCopyEngine` (NT-AVX2 memcpy) |
+| both local, either `loc == GPU` | `HbmCopyEngine` (`hipMemcpy`) |
+| exactly one endpoint remote | `MoriIoEngine` (RDMA via mori-io) |
+
+Because they are disjoint, `CompositeTransferEngine`'s first-match registration
+order is *documentation*, not a tie-break. **Keep it that way.** If two engines
+can both claim a pair, registration order silently becomes a performance
+decision that nothing tests. Assert disjointness:
+
+```cpp
+EXPECT_TRUE(local.CanHandle(h1, h2));
+EXPECT_FALSE(hbm.CanHandle(h1, h2));   // both-CPU stays with the NT-AVX2 path
+```
+
+### Finding the gap you are filling
+
+Before writing an engine, prove no existing one claims your pair. Read the
+`CanHandle` bodies — they are short and each has a hard precondition:
+
+- `local_copy_engine.cpp`: `src.HasHostPtr() && dst.HasHostPtr() && both loc == CPU`
+- `mori_io_engine.cpp`: `if (src_remote == dst_remote) return false;` — refuses
+  both-local *and* both-remote
+
+That second line is why `HbmCopyEngine` had to exist: a local pair with a GPU
+endpoint was claimed by nobody, so an HBM backend's local Put/Get could not
+complete at all. Write that as your first test — it is the engine's reason for
+existing, and it should be asserted against the real composite rather than by
+inspection.
+
+## `TransferRef`: what an endpoint is
+
+```cpp
+struct TransferRef {
+  void* host_ptr;                    // process-local view; null for a peer's memory
+  uint64_t size;
+  mori::io::MemoryLocationType loc;  // CPU | GPU
+  int device;
+  mori::io::MemoryDesc mem;          // mori-io registration; .size == 0 if none
+};
+```
+
+**`host_ptr` means "addressable in this process", not "host memory".** For an
+`hipMalloc`'d buffer the *device* pointer is that view, since it is what
+`hipMemcpy` accepts. So `HbmCopyEngine` needed no new field — `loc = GPU` plus
+the existing pointer was enough. Check whether that is true for you before
+adding anything.
+
+It is deliberately **not** a `std::variant`. Registration fans *out*: the same
+buffer is a raw pointer **and** an RDMA MR at the same time, and which one is
+used is a property of the pair. A variant makes the both-local-but-also-
+registered case inexpressible.
+
+Empty handles are how an endpoint says "not reachable that way":
+`host_ptr == nullptr` → not addressable here; `mem.size == 0` → not registered
+with mori-io.
+
+### If you genuinely need a non-memory endpoint
+
+A file or object endpoint needs a `kind` tag plus a second handle set in
+`transfer_engine.h` — a one-file change, and `CanHandle` already exists to route
+it. But note what comes with it: `CompositeTransferEngine` does **not** implement
+chaining, where no single engine spans the pair and the layer composes two hops
+through a bounce buffer (a file on node A, a reader on node B). You would be
+adding that too. `SsdBackend` sidesteps the whole question by staging through
+registered host memory — see `umbp-add-backend`.
+
+## The interface
+
+```cpp
+class TransferEngine : public MemoryRegistrar {
+  virtual const char* Name() const = 0;
+  virtual TransferRef RegisterMemory(void*, size_t, MemoryLocationType, int device) = 0;
+  virtual void Deregister(const TransferRef&) = 0;
+  virtual bool CanHandle(const TransferRef& src, const TransferRef& dst) const = 0;
+  virtual TransferPlanSet Plan(const std::vector<TransferItem>&) const = 0;
+  virtual std::unique_ptr<TransferHandle> Submit(std::vector<TransferPlan>) = 0;
+  bool Transfer(const std::vector<TransferItem>&, std::vector<size_t>* failed_tags);
+};
+```
+
+### `RegisterMemory`
+
+If your engine needs no pinning, return the endpoint unchanged:
+
+```cpp
+TransferRef RegisterMemory(void* base, size_t size, MemoryLocationType loc, int dev) override {
+  return TransferRef::HostBytes(base, size, loc, dev);
+}
+void Deregister(const TransferRef&) override {}
+```
+
+`CompositeTransferEngine` fans registration out to **every** sub-engine and
+merges the handles, so a buffer can be a raw pointer and an MR at once. Your
+engine seeing a buffer it cannot pin is normal, not an error.
+
+### `Plan` — pure, and virtual for a reason
+
+`Plan` must issue no IO and change no engine state, so a caller may plan on one
+thread and submit on another. It is virtual rather than a base helper because
+grouping is engine-specific: pair-grouping is an RDMA optimization (it cuts CQE
+and post count), a memcpy engine gains little, and a file engine would rather
+group by file and merge adjacent offset ranges.
+
+Two things every `Plan` must do:
+
+1. **Reject what you cannot carry**, into `rejected_tags`. This is not an error
+   path the caller can ignore — those items were dropped and their keys must be
+   failed. `PoolClient` logs them as "transfer unplannable (no engine claimed the
+   endpoints)".
+
+2. **Bounds-check once, here**, not per copy:
+
+   ```cpp
+   if (item.size > item.src.size || item.src_offset > item.src.size - item.size ||
+       item.size > item.dst.size || item.dst_offset > item.dst.size - item.size) {
+     out.rejected_tags.push_back(item.tag);
+     continue;
+   }
+   ```
+
+   Written that way to avoid overflow. A page index that does not fit its buffer
+   is a backend bookkeeping bug, and the key must be failed rather than copied
+   past the end — on a GPU endpoint an overrun is not a segfault but silent
+   corruption of whatever else is in the pool.
+
+Coalescing adjacent segments is worth it if your per-call cost is non-trivial:
+`HbmCopyEngine` coalesces because each `hipMemcpy` carries launch overhead a
+`memcpy` does not.
+
+`tag` is the caller's per-entry index. It must survive grouping so a failed plan
+maps back to the keys that rode in it; de-duplicate as you append.
+
+### `Submit` — and the one blocking exception
+
+`Submit` does **not** block on the wire, with one documented exception: a plan
+that needs the engine's **bounce buffer** completes *inside* `Submit` and comes
+back already settled. That keeps the staging lock from ever being held across a
+return, which is what makes a submit-all-then-wait loop over several peers
+deadlock-free without the caller knowing staging exists.
+
+If your engine completes inline (any synchronous copy), return an
+already-settled handle that just replays the outcome:
+
+```cpp
+class SettledHandle final : public TransferHandle {
+  void Wait(std::vector<TransferFailure>* failures) override {
+    if (reported_) return;              // idempotent: a second call appends nothing
+    reported_ = true;
+    for (auto& f : failures_) failures->push_back(std::move(f));
+    failures_.clear();
+  }
+};
+```
+
+Return `nullptr` if nothing was posted — the caller must then treat every tag as
+failed.
+
+**Failures are per PLAN, not per segment.** The plan is the unit a tag set maps
+back to, and a partially-copied key is failed wholesale anyway. Fill in
+`TransferFailure{tags, code, message, endpoint}`; `endpoint` is for diagnosis.
+
+### `Wait` lives on the handle, not the engine
+
+Because mori-io's RDMA backend keeps raw `TransferStatus*` into the handle's
+status vector: the handle must outlive the post and must never move its
+statuses. Its destructor drains as a **safety net** for an early or exceptional
+destroy, so a completion callback never writes freed memory. `Wait` never breaks
+early — every plan is waited, or a status is left live.
+
+## Threading and ambient state
+
+`Submit` is called concurrently from `PoolClient`'s
+`UMBP_DRAM_{READ,WRITE}_THREADS` executor threads. Two consequences:
+
+- Shared mutable engine state needs its own synchronization. `HbmCopyEngine`
+  avoids a shared HIP stream for exactly this reason — the parallelism that
+  matters is *across keys* and already exists above the engine, so a stream
+  would add synchronization with no matching win at KV-block sizes.
+- **Restore any ambient state you touch.** Those threads are not yours to leave
+  re-pointed at another GPU:
+
+  ```cpp
+  int entry_device = -1;
+  if (hipGetDevice(&entry_device) != hipSuccess) entry_device = -1;
+  // ... hipSetDevice(want) per plan ...
+  if (touched && entry_device >= 0) (void)hipSetDevice(entry_device);
+  ```
+
+## Transports that need a handshake: `PeerDirectory`
+
+Some transports cannot address a peer until told about it — mori-io needs the
+peer's `EngineDesc` registered and its published `MemoryDesc`s unpacked before a
+byte moves. That handshake is not a transfer, so it is **not** on
+`TransferEngine`; it is its own interface, implemented by the engines that need
+it. `PoolClient` drives peer setup through a `PeerDirectory*` obtained at
+composition time, so a second remote transport implements the interface and
+`pool_client.cpp` does not change. A local copy engine has no peers and
+implements nothing.
+
+Note "directory", not "cache": these calls are the authority on what this
+process knows about a peer, and `ForgetRemote` is how a failed handshake is
+undone.
+
+## Wiring it in
+
+`PoolClient::Init` is the composition root and the only place any concrete
+engine is named:
+
+```cpp
+auto composite = std::make_unique<CompositeTransferEngine>();
+composite->AddEngine(std::make_unique<LocalCopyEngine>());
+composite->AddEngine(std::make_unique<HbmCopyEngine>());
+if (!config_.io_engine.host.empty()) {
+  auto rdma = std::make_unique<MoriIoEngine>(...);
+  if (!rdma->Init()) { /* fail */ }
+  peer_directory_ = rdma.get();
+  composite->AddEngine(std::move(rdma));
+}
+transfer_engine_ = std::move(composite);
+```
+
+Local engines go **before** the wire engine on principle: a pair that reached
+mori-io wrongly would be refused rather than served slowly. Add your source to
+`src/umbp/CMakeLists.txt` under `umbp_common`.
+
+New external dependencies are usually already transitively available —
+`umbp_common` links `mori_io`, which links `hip::host`, which is why
+`HbmCopyEngine` could call `hipMemcpy` with no CMake change beyond the source
+listing.
+
+That `AddEngine` call is also all the observability wiring there is.
+`CompositeTransferEngine` is the dispatch point, so it is the measurement point:
+it charges bytes, plan counts, failures and in-flight time to whichever engine
+carried each plan and publishes them under `engine=<your Name()>,
+direction=push|pull|local`. Your engine appears in the transfer panels of
+`examples/monitoring/grafana/dashboards/umbp_backends.json` without a line of
+metrics code.
+
+Override `SampleMetrics()` (from `MetricSource`) only for transport-internal
+state the dispatcher cannot see — bounce-pool pressure, queue depth. Name the
+metric generically and put the specifics in a label, the way a storage backend
+does; `engine=` is stamped for you. See
+`umbp/distributed/metrics/component_metrics.h`.
+
+## Testing
+
+Split the suite the way the layer splits:
+
+**Selection and planning need no hardware and should always run.** Cover:
+the gap your engine fills (assert the *other* engines refuse the pair);
+disjointness; composite routing via `SelectEngine`, which is exposed for exactly
+this; coalescing; out-of-bounds rejection; and rejection of pairs you cannot
+carry.
+
+```cpp
+TransferEngine* chosen = composite.SelectEngine(host_ref, gpu_ref);
+ASSERT_NE(chosen, nullptr);
+EXPECT_STREQ(chosen->Name(), "HbmCopyEngine");
+```
+
+**Byte movement needs the hardware**, so `GTEST_SKIP()` without it. Use
+`Transfer()` (plan+submit+wait) for round-trip tests — it is there for callers
+with nothing to overlap, and tests are such a caller.
+
+A trap worth knowing: when building a "scattered pages" test, a stride that is
+not coprime with the page count is not a permutation and silently leaves pages
+uncopied — `(i * 2) % 4` gives `{0,2,0,2}`. That failure looks like an engine
+bug and is not.
+
+`MoriIoEngine` is not unit-testable (needs a real fabric); it is covered
+end-to-end by `test_cross_node_smoke`, which is labeled `integration`.
+
+## Build and run
+
+The umbp build needs `protoc`/`grpc_cpp_plugin` from the mori Docker image, not
+the bare host:
+
+```bash
+docker exec <mori-container> bash -c "cd <repo> && ninja -C build && \
+  cd build && ctest -E '^cco_' -LE integration"
+```
+
+## Reference
+
+- `transfer/transfer_engine.h` — `TransferRef`, `MemoryRegistrar`,
+  `PeerDirectory`, `TransferItem/Plan/Handle`, `TransferEngine`; also records
+  what was deliberately left out and why
+- `transfer/local_copy_engine.{h,cpp}` — the simplest complete engine
+- `transfer/hbm_copy_engine.{h,cpp}` — a second engine closing a real gap
+- `transfer/mori_io_engine.{h,cpp}` — RDMA, bounce buffers, `PeerDirectory`
+- `transfer/composite_transfer_engine.{h,cpp}` — fan-out registration,
+  per-pair dispatch; the "no chaining" note is here
+- `doc/design-backend-agnostic-refactor.md` §4

.github/workflows/ci.yml

@@ -452,7 +452,12 @@ jobs:
         include:
           - platform: MI355X_AINIC
             runner: [self-hosted, MI355X-AINIC-TW]
-    timeout-minutes: 30
+            rdma_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
+            rdma_sl: 3
+            rdma_tc: 104
+    # Raised from 30 for the primitive concurrency gate, which sweeps five
+    # primitives over four rank counts against two server deployments.
+    timeout-minutes: 60
     steps:
       - name: Fix workspace permissions
         run: |
@@ -494,12 +499,42 @@ jobs:
 
       - name: UMBP C++ unit tests (local + distributed)
         run: |
-          $CT exec $CONTAINER bash -c "
+          $CT exec -e MORI_UMBP_LOG_LEVEL=warn $CONTAINER bash -c "
             cd $GITHUB_WORKSPACE/build
             timeout 300 ctest --output-on-failure \
               -R '^(umbp_|test_dummy_ssd_tier|test_prefix_aware_eviction|test_ssd_tier)'
           "
 
+      # Five client primitives -- put / get / exists / put_ranges / get_ranges
+      # -- swept over 1, 2, 4 and 8 concurrent client processes, against a
+      # standalone server run both without a master and with one.  Gates on
+      # correctness at every rank count, and on batch_exists throughput still
+      # growing with concurrency: a control plane that serializes reads looks
+      # healthy at one rank and collapses at eight, which no single-concurrency
+      # latency number can see.  See tests/python/umbp/umbp_prim_ci.py.
+      - name: UMBP primitive concurrency gate (standalone, with and without master)
+        run: |
+          $CT exec \
+            -e MORI_RDMA_DEVICES=${{ matrix.rdma_devices }} \
+            -e MORI_RDMA_SL=${{ matrix.rdma_sl }} \
+            -e MORI_RDMA_TC=${{ matrix.rdma_tc }} \
+            $CONTAINER bash -c "
+              cd $GITHUB_WORKSPACE
+              timeout 2400 python3 tests/python/umbp/umbp_prim_ci.py \
+                --arms nomaster master \
+                --ranks 1 2 4 8 \
+                --out-dir $GITHUB_WORKSPACE/umbp_prim_ci_results
+            "
+
+      - name: Upload primitive gate results
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: umbp-prim-ci-${{ github.run_id }}
+          path: |
+            ${{ github.workspace }}/umbp_prim_ci_results
+          if-no-files-found: warn
+
       - name: Cleanup
         if: always()
         run: |