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:

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:

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

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

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 std::unique_ptr<TransferPlanBuilder> CreatePlanBuilder(size_t expected_items = 0) const = 0;
  virtual std::unique_ptr<TransferHandle> Submit(std::vector<TransferPlan>) = 0;
  TransferPlanSet Plan(const std::vector<TransferItem>&) const;  // your builder over a vector
  bool Transfer(const std::vector<TransferItem>&, std::vector<size_t>* failed_tags);
};

RegisterMemory

If your engine needs no pinning, return the endpoint unchanged:

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.

Planning — CreatePlanBuilder, the one hook

An engine plans through a TransferPlanBuilder fed one item at a time, so a ranged get never materializes its ~15k items (720 B each) as a vector. Plan() is not virtual: it is your builder run over a vector, so batch and incremental planning cannot differ. Grouping is still yours: 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.

The rules every builder keeps (documented on TransferPlanBuilder; callers rely on them):

1. Add() issues no IO and changes no engine state, so a caller may plan on one thread and submit on another.
2. An empty item (size 0) is dropped: no plan, no rejection.
3. Reject what you cannot carry into rejected_tags: a pair your CanHandle() refuses, or a range past the end of a buffer you can check. This is not an error path the caller can ignore — those keys are failed, and PoolClient logs them as "transfer unplannable". Bounds-check once, at planning, not per copy: on a GPU endpoint an overrun is not a segfault but silent corruption of whatever else is in the pool.
4. Every other item lands in one plan whose tags name it, de-duplicated per run of its segments, so a failed plan maps back to the keys that rode in it.
5. Plans come out in a deterministic order (first item first).

How to write one — transfer/plan_builder.h has the shared pieces:

• One plan per (src, dst) buffer pair (a memcpy-like engine): don't write a builder at all. MakePairPlanBuilder(*this, append) admits by your CanHandle() plus ItemInBounds(), groups by pair (with a fast path for consecutive items of one pair), and calls append(plan, item) for the segments. LocalCopyEngine passes AppendSegment; HbmCopyEngine passes one that splits at host registration pieces.
• Anything else: write a small TransferPlanBuilder (Add / Finish) using ItemInBounds / RangeFits, PlanBetween and AppendSegment (merge adjacent segments, optional cap — MORI-IO caps at ibv_sge.length — and tag de-duplication). GdsEngine (one plan per item) and MoriIoEngine (groups by memory ids, bounce pool) are the examples.
• Plan as items arrive. Buffer them only if you genuinely must see the whole batch first (e.g. sorting by file offset), and say why in a comment.

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.

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:

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:

  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 MemoryDescs 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:

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. Run ExpectPlanContract(items, planned) (in test_transfer_engine.cpp) on your planning tests: every non-empty item rejected or named by a plan, and the accepted bytes planned exactly once.

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's transfers need a real fabric and are covered end-to-end by test_cross_node_smoke (labeled integration); its planner needs nothing and is unit-tested on an engine that is never Init()ed.

Build and run

The umbp build needs protoc/grpc_cpp_plugin from the mori Docker image, not the bare host:

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/plan_builder.h — PairPlanBuilder and the planning helpers
• 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