diff --git a/.claude/skills/umbp-add-backend/SKILL.md b/.claude/skills/umbp-add-backend/SKILL.md new file mode 100644 index 000000000..b55bed28b --- /dev/null +++ b/.claude/skills/umbp-add-backend/SKILL.md @@ -0,0 +1,341 @@ +--- +name: umbp-add-backend +description: >- + Add a new storage medium to UMBP's distributed mode as a MediumBackend (HBM, + SSD, CXL, a remote object store, a fake for tests). Covers the two shapes a + backend can take — a PageMemorySource for anything page-addressable, or a full + MediumBackend for a medium whose bytes are not directly addressable — plus + registration, the heartbeat/event contract, and the failure modes that only + show up under the peer service. Use when the user wants to add or debug a + storage tier/medium/backend in src/umbp/distributed, mentions MediumBackend, + PageBackend, PageMemorySource or BackendRegistry, or asks why a new tier's + puts or gets are not working. +--- + +# Adding a storage medium to UMBP distributed mode + +A **backend owns bytes and publishes descriptors for them. It does not move +them.** Everything in this skill follows from that one rule, which is stated at +the top of `src/umbp/include/umbp/distributed/peer/backend/medium_backend.h` +and enforced by the type system since Phase 6 (`Init` receives a +`MemoryRegistrar`, which has no `Submit` on it). + +If you find yourself wanting to copy bytes inside a backend, stop — that is the +transfer layer's job, and the sibling skill `umbp-add-transfer-engine` covers +it. + +## First decide which of the two shapes you need + +This is the whole design decision, and getting it wrong costs a rewrite. + +**Is your medium page-addressable — can you hand out a raw pointer that the +transfer layer can read and write directly?** + +| | Yes (DRAM, HBM, CXL, a pmem mapping) | No (SSD, S3, a network filesystem) | +|---|---|---| +| Implement | `PageMemorySource` — **5 methods** | `MediumBackend` — 23 methods | +| Reuse | All of `PageBackend` | Nothing; but see "staging" below | +| Effort | ~100 lines | ~450 lines | +| Example | `hbm_backend.{h,cpp}` | `ssd_backend.{h,cpp}` | + +Most media are the first case. `PageBackend`'s slot lifecycle, bitmap +allocator, event outbox, reaper, read leases and copy pins **never consult the +tier**, so a second paged medium needs no second copy of any of it. + +--- + +## Shape 1: a paged medium (implement `PageMemorySource`) + +Five methods, in `page_backend.h`: + +```cpp +class PageMemorySource { + virtual bool Allocate(const std::vector& sizes, std::vector* out) = 0; + virtual void Release() = 0; + virtual mori::io::MemoryLocationType LocationType() const = 0; + virtual int Device() const = 0; + virtual const char* Name() const = 0; +}; +``` + +`LocationType()` and `Device()` are **the facts a descriptor cannot recover**. +They are the reason this interface exists: `PageBackend` mirrors them into every +`TransferRef` it publishes, and that is what makes a local transfer against your +medium select the right engine. Get them wrong and nothing fails at build or +init time — a transfer just silently picks the wrong engine, or no engine at +all. + +### The recipe + +1. **Write the source.** Put it in its own file if it needs a new dependency + (`hbm_backend.cpp` exists so HIP stays out of `page_backend.cpp`). + + ```cpp + bool MyPageMemorySource::Allocate(const std::vector& sizes, + std::vector* out) { + std::vector taken; + std::vector staged; + for (uint64_t size : sizes) { + if (size == 0) continue; // skip, do not fail + MyHandle h = MyAlloc(size); + if (!h.valid()) { + for (auto& t : taken) MyFree(t); // unwind THIS call only + return false; // leave `out` untouched + } + staged.push_back(Buffer{h.ptr, h.usable_size}); + taken.push_back(h); + } + handles_.insert(handles_.end(), taken.begin(), taken.end()); + out->insert(out->end(), staged.begin(), staged.end()); + return true; + } + ``` + + Report the **usable** size in `Buffer::size`, not the requested one — host + hugepage rounding makes the extra genuinely allocatable, and `PageBackend` + publishes what you report. + +2. **Add a factory** returning `std::unique_ptr`, next to + `MakePageBackend` / `MakeHbmBackend`. It must return the *interface*: Phase 5 + Rule A says only `PoolClient::Init` may name a concrete backend. + + ```cpp + std::unique_ptr MakeMyBackend(uint64_t page_size, ...) { + return std::make_unique(TierType::MY_TIER, page_size, + std::make_unique(...), + std::move(buffer_sizes), pending_ttl, + read_lease_ttl); + } + ``` + +3. **Add a config struct** in `distributed/config.h` holding your medium's own + knobs — no `enabled` flag: `PoolClientConfig::medium` is the single selector. + Do **not** extend `DramOwnershipConfig` — hugepages/NUMA/prefault are + meaningless for HBM, and a device ordinal is meaningless for host memory. + Each medium brings its own knobs; that asymmetry is why the seam is a class + and not an options struct. + +4. **Add a case to the medium switch in `PoolClient::Init`.** A node registers + **exactly one** backend, chosen by `config_.medium`, so you add a case rather + than another `if`. This is the only file that changes outside your own: + + ```cpp + case TierType::MY_TIER: { + backend = MakeMyBackend(page_size, config_.my, ...); + break; + } + ``` + + The shared `Init` / `Register` / error handling below the switch is written + once for every medium. + + **Why one medium, not a tier stack.** The routing plane does not tier: master + treats every advertised tier as an equally valid put target (Phase 4 deleted + the hardcoded tier orders), so a node registering two backends *mirrors* + across them rather than promoting/demoting between them. Heterogeneity comes + from different **nodes** picking different media. If you find yourself + wanting two live backends on one node, you are asking for a local tiering + policy that does not exist yet — that is a routing-plane change, not a + backend one. + +5. **Add the tier to `TierType`** (`types.h`) if it is genuinely new. `HBM=1, + DRAM=2, SSD=3` already exist. `BackendRegistry` is a `map`, so + one instance == one medium and the enum value is the identity. Also add a + matching `UMBPMedium` value in `common/config.h` and map it in `ToTierType`, + or no user-facing config can select your medium. + +That is the whole change. Routing, the peer service, the heartbeat and the batch +executors were all written against `BackendRegistry` and need no edit. + +### Metrics: do not write any + +`PoolClient::Init` wraps whatever the medium switch produced in +`InstrumentedBackend` before registering it, and that decorator derives the +whole generic series — operations by outcome, bytes committed / resolved / +freed, batch depth, time spent inside the medium — from the `MediumBackend` +calls themselves. Your backend is charted the moment it is composed in, under +`tier=`, in the panels of +`examples/monitoring/grafana/dashboards/umbp_backends.json` that already exist. +There is no metric to register, no dashboard to edit, and **no counter to add to +your backend**. + +Override `SampleMetrics()` (from `MetricSource`) only for state the interface +cannot show from outside — what the device itself did, how full an internal +arena is. When you do, obey the one rule that keeps the single dashboard +working: **publish under the generic name and put your specifics in a label.** + +```cpp +std::vector MyBackend::SampleMetrics() const { + return {MetricSample{MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL, + MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL_HELP, + {{"event", "device_read_error"}}, // YOU name the event + device_read_errors_.load(std::memory_order_relaxed)}, + MetricSample{MORI_UMBP_METRIC_BACKEND_MEDIUM_STATE, + MORI_UMBP_METRIC_BACKEND_MEDIUM_STATE_HELP, + {{"state", "queue_depth"}}, + depth, MetricKind::kGauge}}; +} +``` + +A metric named after your medium (`mori_umbp_myssd_reads_total`) compiles and +scrapes fine, and is still wrong: it needs its own panel, which is how UMBP +ended up with a dashboard per medium and an SSD dashboard wired to counters +that lost their publisher. `tier=` and `backend=` are stamped by the publisher — +do not set them. Counter values must be monotonic; gauges may move either way. +See `umbp/distributed/metrics/component_metrics.h`. + +--- + +## Shape 2: a medium whose bytes are not addressable + +If you cannot hand out a pointer, you have two options, and the cheap one is +almost always right. + +### Stage through registered host memory (what `SsdBackend` does) + +Publish an ordinary registered host DRAM arena as your buffer, and move bytes +between it and your device **inside** the backend: + +- `BatchAllocate` reserves a staging page and publishes it. The writer RDMAs + into it knowing nothing about your medium. +- `BatchCommit` spills that page to your device and **returns the page to the + arena**. The page is borrowed for the write, not the key's home. +- `BatchResolve` fills a staging page from your device and publishes it under a + **read lease**; a reaper reclaims the page when the lease expires. + +The cost is one host copy per side. The benefit is that your medium reaches the +data plane with **zero new transfer-layer concepts** — no new `TransferRef` +kind, no new engine, no chaining — and a remote peer reading from your node sees +a perfectly ordinary registered buffer and needs no code at all. + +### Add a real endpoint kind (bigger; not yet done in tree) + +`transfer_engine.h` reserves this: a `FileRef`/`ObjectRef` plus a `kind` tag, a +new engine, **and chaining in `CompositeTransferEngine`** for the remote reader +(device → bounce → wire), which that class explicitly does not implement. Take +this path only when you need zero-copy/GDS; staging does not block it. + +### Things that bite in shape 2 + +- **Exhaustion has no honest encoding.** A `Resolve` that cannot get a staging + page must report `found=false`, which makes the client exclude your node and + retry elsewhere — wrong, because your node *does* hold the key. + `medium_backend.h` records that the "not ready, retry here" state was proposed + and rejected. Size the arena for read concurrency, and pin the behavior in a + test so a future control-plane fix changes it deliberately. +- **Contiguity.** `PeerSsdManager::PrepareRead` takes one `(ptr, capacity)`, not + a scatter list, so a staged key must fit one contiguous page. `SsdBackend` + therefore allocates **one buffer** of `staging_pages * page_size` and refuses + keys larger than a page. In distributed mode master's `page_size` *is* the KV + block size, so 1 key == 1 page is the normal case. +- **Never hold the arena lock across IO.** Take the slot out under the lock, + release it, do the spill/fill, then re-lock to return the page. Holding it + stalls every concurrent `Allocate` and `Resolve`. + +--- + +## The contracts that are easy to get wrong + +These are the ones with no compile-time protection. + +**Events go in ONE bundle under ONE seq.** The heartbeat concatenates every +backend's events (`DrainAllBackends`). Never emit one bundle or one seq per +medium — that breaks the ack / seq-gap full-sync recovery. + +**`SnapshotOwnedKeysForFullSync` must clear the outbox in the SAME critical +section** as the snapshot. Two separate locks drop events committed in between. +`SnapshotOwnedKeys` is `const` and must not mutate; the full-sync variant is +not. + +**`kFailedNoSpace` vs `kFailed`.** `kFailedNoSpace` means "medium exhausted, +retry on another peer". Use it only when another peer could plausibly succeed. A +key too large for your page size is `kFailed` — no peer would do better, and the +writer must not keep hunting. + +**`Evict` returns one result per key, in request order.** The peer service sums +freed bytes for a key mirrored across media and relies on the positional match. +Return `bytes_freed = 0` (not an error) for unknown, already-freed, or protected +keys; master retries protected ones next round. + +**`BatchAbort` is idempotent** — a slot already reaped or never seen still +reports `true`. + +**`Shutdown` must tolerate a failed `Init`** and must not run concurrently with +anything else. Deregister **before** releasing memory: the registrar may still +hold an MR over those pages. + +**`Init` is idempotent**, and a backend is fully live when it returns — start +your own reaper thread there, so no caller has to know you have one. + +**Allocation is gated between `ClearLocal()` and `ClearFullSyncAcked()`.** No +new owned key may appear in that window. + +**Threading**: every method may be called from the peer service's gRPC handler +threads and the heartbeat thread concurrently. + +**`SetAutoFlushHook`'s callback runs under your lock** — it must be cheap and +must not re-enter the backend. Signal the heartbeat thread and return. + +--- + +## Testing + +Two levels, both worth having: + +1. **Without the transfer layer** — drive `BatchAllocate` / `BatchCommit` / + `BatchResolve` / `Evict` directly and assert the bookkeeping. A + `LocalOnlyRegistrar` (returns `TransferRef::HostBytes`, counts + register/deregister calls) is all the `MemoryRegistrar` you need; it is also + exactly what `CompositeTransferEngine` degrades to on a node with no RDMA. + +2. **Through a real `CompositeTransferEngine`** — a Put and a Get moving actual + bytes. This is what catches a wrong `LocationType()`/`Device()`, because a + mislabeled endpoint selects the wrong engine and the copy either fails or + silently does nothing. + +Assert the registration contract explicitly: + +```cpp +EXPECT_EQ(registrar.last_loc, mori::io::MemoryLocationType::GPU); +EXPECT_EQ(backend->BufferRef(0).device, 0); +``` + +For a staging backend, the highest-value test is that **Commit returns the +page**: configure 2 staging pages, do 10 sequential puts, assert all succeed. A +leaked page wedges the backend after `staging_pages` puts and nothing else +catches it. + +Skip rather than fail when hardware is absent (`if (!HaveGpu()) GTEST_SKIP()`), +so the suite stays runnable on a CPU-only box. + +Register the test in `tests/cpp/umbp/distributed/CMakeLists.txt` with +`add_test(NAME ... COMMAND ...)`, **not** `gtest_discover_tests` — the file +explains why (discovery bakes in build-time cmake paths and goes red when the +image's ctest runs it). + +## Building and running + +The umbp C++ build needs `protoc` and `grpc_cpp_plugin`, which live in the mori +Docker image, not on the bare host: + +```bash +docker exec bash -c "cd && \ + cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_UMBP=ON -DBUILD_TESTS=ON -DGPU_TARGETS=gfx950 \ + -DCMAKE_CXX_COMPILER=/opt/rocm/lib/llvm/bin/clang++ \ + -DCMAKE_HIP_COMPILER=/opt/rocm/lib/llvm/bin/clang++ && \ + ninja -C build && cd build && ctest -E '^cco_' -LE integration" +``` + +Exclude `^cco_` — those are unrelated GPU collective tests and some hang without +a full fabric. `-LE integration` skips the tests needing real RDMA. + +## Reference + +- `medium_backend.h` — the interface, and a list of things deliberately **not** + on it (with reasons, so they are not re-proposed) +- `page_backend.h` — `PageMemorySource`, `HostPageMemorySource`, `PageBackend` +- `hbm_backend.{h,cpp}` — the minimal paged medium (shape 1) +- `ssd_backend.{h,cpp}` — the staged, non-addressable medium (shape 2) +- `doc/design-backend-agnostic-refactor.md` — §2 the descriptor/pointer rule, + §3 the three-component split, §8 what none of it fixes diff --git a/.claude/skills/umbp-add-transfer-engine/SKILL.md b/.claude/skills/umbp-add-transfer-engine/SKILL.md new file mode 100644 index 000000000..0ba2e0f61 --- /dev/null +++ b/.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&) const = 0; + virtual std::unique_ptr Submit(std::vector) = 0; + bool Transfer(const std::vector&, std::vector* 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* 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(); +composite->AddEngine(std::make_unique()); +composite->AddEngine(std::make_unique()); +if (!config_.io_engine.host.empty()) { + auto rdma = std::make_unique(...); + 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=, +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 bash -c "cd && 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 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6706148c5..649f9d3e4 100644 --- a/.github/workflows/ci.yml +++ b/.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: | diff --git a/docs/UMBP-TIER-MANAGEMENT-BENCHMARK.md b/docs/UMBP-TIER-MANAGEMENT-BENCHMARK.md new file mode 100644 index 000000000..6a009658e --- /dev/null +++ b/docs/UMBP-TIER-MANAGEMENT-BENCHMARK.md @@ -0,0 +1,197 @@ +# UMBP Distributed Tier-Management Benchmark + +`umbp_tier_bench` runs deterministic synthetic or recorded workloads through an +in-process Master and one or more PoolClients. It measures the distributed UMBP +PUT/GET path, including routing, named backend placement, peer-local offload, +payload validation, throughput, latency, and backend capacity. + +## Build + +```bash +cmake -S . -B build -DBUILD_UMBP=ON -DBUILD_TESTS=ON +cmake --build build --target \ + umbp_tier_bench test_umbp_workload_trace test_umbp_workload_runner +``` + +The benchmark is a manual target and is not run by plain `ctest`. + +## Run, generate, and replay + +Run a workload against the default topology, one DRAM backend per peer: + +```bash +./build/src/umbp/umbp_tier_bench run \ + --profile mixed --clients 4 --operations 10000 --keys 2000 \ + --read-ratio 0.7 --min-value-bytes 4096 --max-value-bytes 4096 \ + --qps 20000 --backend-capacity 2GiB +``` + +Generate a reusable trace and replay it under a policy: + +```bash +./build/src/umbp/umbp_tier_bench generate \ + --trace workload.umbptrace --profile hotset --seed 42 \ + --operations 100000 --keys 10000 --clients 8 --qps 50000 + +./build/src/umbp/umbp_tier_bench replay \ + --trace workload.umbptrace --config policy.json +``` + +Use `umbp_tier_bench --help` for the complete option list. + +## JSON backend and tier policy + +Anything beyond one DRAM backend per peer — several backends, other media, +weights, logical tiers, watermarks — is described by `--config`, which is the +same policy format the product reads: + +```bash +./build/src/umbp/umbp_tier_bench run \ + --config examples/umbp/multi_backend_policy.json \ + --profile capacity-pressure --operations 1000 --page-size 2MiB +``` + +That file is the complete example: HBM in front of DRAM and two SSDs, three +tiers, weighted members. The `umbp_backend_policy_config` test loads it, so the +example and the parser cannot drift apart. A single tier from it, which is where +the placement and migration knobs live: + +```json +{ + "name": "warm", + "backends": { "dram": 70, "ssd_a": 30 }, + "offload_to": ["cold"], + "offload_trigger": "watermark", + "high_watermark": 0.9, + "low_watermark": 0.7 +} +``` + +New PUTs use `entry_tier` (the first tier when omitted) and its deterministic +backend weights. `NO_SPACE` recursively spills through `offload_to`; every +destination tier applies its own weights. +`on_evict` copies the key to the first available target before deleting its +source copy. `watermark` queues asynchronous LRU offload after aggregate tier +use reaches `high_watermark`. Migration continues toward `low_watermark` and +retries transient failures with bounded backoff. + +`promote_trigger` says when a tier sends a key back up, mirroring +`offload_trigger`: `never` (the default), `on_read` on the first hit, or +`on_hits` once `promote_hits` reads have been served for that key from this +tier. `promote_hits` is required by `on_hits`, rejected by the others, and must +be at least 2 — a threshold of 1 is `on_read`. The entry tier may not declare a +trigger, having no upstream tier to promote into. `promote_mode` decides whether +the promoted key is copied or moved; under `move` the source tier is reclaimed +when the master next names that key as an eviction victim, not at promotion time, +because the delete defers behind the read lease the triggering read holds. + +Backend and logical-tier names must be unique, and every backend must belong to +exactly one logical tier. Offload targets may name a backend or logical tier +and must resolve to a later tier, which makes the +topology acyclic. A multi-device HBM definition expands to one backend per +device while preserving the logical backend's aggregate weight and capacity. +The current per-peer limit is eight concrete backends. + +Production distributed clients may set +`UMBPDistributedConfig::backend_policy_path`; Python exposes the same +`backend_policy_path` property. Distributed-backed standalone servers also +accept `UMBP_BACKEND_POLICY`. + +Each benchmark run prints `backend_placement` and `tier_transitions` CSV +sections. Production clients can record successful PUT/GET traffic with +`UMBP_WORKLOAD_TRACE_PATH`; replay it with payload validation disabled because +production payload bytes are not embedded in the trace. + +**Page size vs value size.** Each value occupies one or more whole pages. The +defaults (`page_size=2MiB`, `backend_capacity=256MiB`) only hold 128 one-page +keys per backend. A 4KiB value still consumes 2MiB. For small-value policy +sweeps use `--page-size 64KiB --backend-capacity 2GiB` (or size capacity as +`page_size * key_count`). + +**GET affinity.** Synthetic PUT/GET pairs for one key stay on one client +stream. `--affinity none` plus `--get-strategy local` can still miss: Master +`most_available` may place the PUT on another node, and the producing client's +local GET then waits out the publication retry. For GET-heavy same-client +profiles use `--affinity local`. To exercise the remote RDMA path, PUT from +one client and GET from another (see `test_umbp_tier_benchmark`). + +**RDMA QP depth.** `mori_io` defaults `max_send_wr=8192`. On some mlx5 devices +that depth times SGE/inline exceeds the per-QP work-queue budget, so +`ibv_create_qp` returns EINVAL even though `max_qp_wr` is larger. Set a smaller +depth before running multi-client benchmarks or integration tests: + +```bash +export MORI_IO_QP_MAX_SEND_WR=1024 +export MORI_IO_QP_MAX_CQE=4096 +``` + +Also set `LD_LIBRARY_PATH` to the build tree (`build/src/application:build/src/io:build/src/metrics`) +when running binaries from `build/`. + +**SSD.** A value may span several contiguous staging pages. It must fit the +arena (`page_size * staging_pages`); committed SSD capacity is charged by the +value's actual bytes. `max_allocatable_bytes` is +`min(available_bytes, staging_arena_bytes)`. + +The workload controls are profile, seed, operation and key counts, minimum and +maximum value size, value-size distribution, read ratio, client count, batch +size, and QPS. Cluster controls select DRAM, HBM, or SSD; backend count, +capacity, and page size; single or weighted placement; and Master put strategy, +affinity, and get strategy. + +`run` defaults to max-throughput scheduling. `replay` defaults to open-loop +scheduling using trace timestamps. `--time-scale` scales those timestamps, and +`--max-throughput` ignores them while preserving order within each client. +Different clients execute concurrently. + +`--settle-ms` controls both the registration and post-run heartbeat settle. +Master eviction uses production defaults. + +## Synthetic profiles + +- `sequential`: write-only traversal of the key space. +- `uniform`: uniformly selected keys with the configured read ratio. +- `hotset`: Zipf-distributed access within the default hot set. +- `read-after-write`: alternating PUT and GET for the same key. +- `mixed`: deterministic key traversal with the configured read ratio. +- `capacity-pressure`: non-repeating writes intended to exceed capacity. + +Random choices and values are derived from the seed. UMBP keys are immutable: +repeated logical writes create versioned keys, and reads target the latest +version. Each key's dependency chain stays on one client stream. + +## Trace contract + +Trace files start with the versioned `UMBPTRCE` envelope and contain +length-delimited protobuf records defined in +`src/umbp/distributed/proto/umbp_workload.proto`: one +`WorkloadTraceHeader`, followed by zero or more `WorkloadEvent` records. + +The header records the schema version, nanosecond time unit, payload seed, and +all synthetic generation settings. Events record relative time, client id, +operation id, PUT/GET kind, key, value size, and batch id. Readers reject +unsupported versions and malformed records. + +Payload bytes are omitted. PUT data is generated from +`(key, operation_id, seed)`, and GET data can be checked against the same +sequence. Disable checking with `--no-payload-validation` when replaying a +future externally recorded trace without payload identity. + +Successful PUTs flush heartbeat publication. A GET for a key produced by the +same run retries for a fixed five seconds while that publication is pending, so +asynchronous Master visibility is not reported as a storage miss. + +## Results + +The `summary` CSV section reports aggregate operation and byte counts, +failures, misses, validation failures, wall time, throughput, latency +percentiles, and scheduling lag. + +After the measured interval and heartbeat settle, `backend_placement` reports +each node/backend's tier, owned key count, total and available capacity, and +maximum allocatable block. Placement inspection is outside the timed region and +uses side-effect-free backend counters. + +For policy comparisons, retain the trace and change only topology or routing +arguments. Keep the settle interval, heartbeat environment, and hardware +configuration identical. diff --git a/examples/monitoring/README.md b/examples/monitoring/README.md new file mode 100644 index 000000000..6600a0766 --- /dev/null +++ b/examples/monitoring/README.md @@ -0,0 +1,61 @@ +# UMBP monitoring + +A Prometheus scrape fragment and the Grafana dashboards for a UMBP deployment. +Point Grafana at `grafana/dashboards/` and Prometheus at +`umbp_master.scrape.yml`; `example/docker-compose.yml` is a working stack that +does both. + +## The dashboards + +| Dashboard | Covers | +| --- | --- | +| `umbp_backends.json` | **Every storage medium and transfer engine**, in shared panels | +| `umbp_data_rate_bandwidth.json` | Per-client inbound/outbound data rate and BatchPut/BatchGet bandwidth | +| `umbp_rpc_call_rates.json` | Client capacity, routing and batch API call rates | +| `umbp_master_client_rpc_latency.json` | Master RPC latency, QPS and error rate by RPC | +| `umbp_external_kv.json` | The external-KV index: report/revoke/match and live block counts | + +The last four are backend-agnostic — they measure the client, the master and +the KV index, which look the same whatever medium is underneath — and are +unchanged. + +## One dashboard for every medium + +`umbp_backends.json` is the one that got merged. It replaced the per-medium +dashboards (`umbp_ssd_tier.json`, and whatever the next medium would have +needed): every storage panel groups by `tier`, and every transfer panel by +`engine`, so DRAM, HBM and SSD appear as series in the same panel and can be +read against each other. There is deliberately **no dashboard per medium**. + +That is a property of how the metrics are emitted, not a convention this file +asks you to follow. A medium does not name its own metrics: `InstrumentedBackend` +wraps every `MediumBackend` and derives the generic series from the interface +calls, and `CompositeTransferEngine` does the same for every engine it +dispatches to. **A backend or engine added to the code shows up in these panels +with no dashboard change and no metrics code of its own.** + +The one thing a component publishes for itself is state its interface cannot +show from outside — what a device actually did, how full an internal arena is. +Those ride generic names with the specifics in a label +(`mori_umbp_backend_medium_events_total{tier="SSD", event="single_flight_dup"}`), +so they land in the existing "medium-internal" panels rather than needing new +ones. See `src/umbp/include/umbp/distributed/metrics/component_metrics.h` for +the vocabulary and the rule. + +## Editing the dashboards + +All five are plain Grafana JSON, edited the way Grafana dashboards normally +are: change the panel in the Grafana UI, export the dashboard, commit the JSON. + +When you add a panel to `umbp_backends.json`, write the query against +**labels** — `sum by (tier) (...)`, `sum by (engine, direction) (...)` — never +against a metric name or a pinned matcher that spells one medium or one +transport (`mori_umbp_ssd_read_total`, `tier="SSD"`). A panel written the other +way stops covering the system the day someone adds a backend, which is exactly +the failure this dashboard replaced. + +## Template variables + +`umbp_backends.json` uses `node` and `tier`, populated from +`mori_umbp_client_capacity_total_bytes` — a metric every live deployment +publishes, so the pickers fill in as soon as one client registers. diff --git a/examples/monitoring/grafana/dashboards/umbp_backends.json b/examples/monitoring/grafana/dashboards/umbp_backends.json new file mode 100644 index 000000000..57a6b2cc7 --- /dev/null +++ b/examples/monitoring/grafana/dashboards/umbp_backends.json @@ -0,0 +1,1671 @@ +{ + "uid": "umbp-backends", + "title": "UMBP Master - Backend Metrics", + "description": "One dashboard for every UMBP storage medium and transfer engine, replacing the per-medium dashboards (the SSD one and whatever the next medium would have needed). Panels are written against labels (tier, backend, engine) rather than metric names, so a backend or engine added to the code appears as another series here with no dashboard change. Everything that is NOT per-backend \u2014 client data rates, master RPC, the external-KV index \u2014 keeps its own dashboard in this directory.", + "tags": [ + "umbp", + "mori" + ], + "timezone": "browser", + "schemaVersion": 36, + "version": 1, + "editable": true, + "refresh": "10s", + "time": { + "from": "now-1h", + "to": "now" + }, + "templating": { + "list": [ + { + "type": "datasource", + "name": "datasource", + "label": "Data source", + "pluginId": "prometheus", + "query": "prometheus", + "hide": 0, + "current": { + "text": "Prometheus", + "value": "Prometheus" + } + }, + { + "type": "query", + "name": "node", + "label": "Node", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "query": "label_values(mori_umbp_client_capacity_total_bytes, node)", + "refresh": 2, + "includeAll": true, + "multi": true, + "allValue": ".+", + "current": { + "text": "All", + "value": "$__all" + } + }, + { + "type": "query", + "name": "tier", + "label": "Medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "query": "label_values(mori_umbp_client_capacity_total_bytes, tier)", + "refresh": 2, + "includeAll": true, + "multi": true, + "allValue": ".+", + "current": { + "text": "All", + "value": "$__all" + } + } + ] + }, + "panels": [ + { + "type": "row", + "title": "Fleet", + "collapsed": false, + "gridPos": { + "h": 1, + "w": 24, + "x": 0, + "y": 0 + }, + "panels": [], + "id": 1 + }, + { + "type": "stat", + "title": "Alive clients", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 4, + "w": 3, + "x": 0, + "y": 1 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "text", + "value": null + } + ] + } + }, + "overrides": [] + }, + "options": { + "colorMode": "value", + "graphMode": "area", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + }, + "targets": [ + { + "expr": "mori_umbp_client_count", + "legendFormat": "clients", + "instant": true, + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Clients currently registered with the master.", + "id": 2 + }, + { + "type": "stat", + "title": "Media in use", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 4, + "w": 3, + "x": 3, + "y": 1 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "text", + "value": null + } + ] + } + }, + "overrides": [] + }, + "options": { + "colorMode": "value", + "graphMode": "area", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + }, + "targets": [ + { + "expr": "count(count by (tier) (mori_umbp_client_capacity_total_bytes{node=~\"$node\", tier=~\"$tier\"}))", + "legendFormat": "tiers", + "instant": true, + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Distinct storage media advertised across the fleet. A new backend raises this with no dashboard change.", + "id": 3 + }, + { + "type": "stat", + "title": "Pool capacity", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 4, + "w": 3, + "x": 6, + "y": 1 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "text", + "value": null + } + ] + }, + "unit": "bytes" + }, + "overrides": [] + }, + "options": { + "colorMode": "value", + "graphMode": "area", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + }, + "targets": [ + { + "expr": "sum(mori_umbp_client_capacity_total_bytes{node=~\"$node\", tier=~\"$tier\"})", + "legendFormat": "total", + "instant": true, + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Total advertised capacity across every selected node and medium.", + "id": 4 + }, + { + "type": "stat", + "title": "Pool used", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 4, + "w": 3, + "x": 9, + "y": 1 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "text", + "value": null + } + ] + }, + "unit": "bytes" + }, + "overrides": [] + }, + "options": { + "colorMode": "value", + "graphMode": "area", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + }, + "targets": [ + { + "expr": "sum(mori_umbp_client_capacity_used_bytes{node=~\"$node\", tier=~\"$tier\"})", + "legendFormat": "used", + "instant": true, + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "id": 5 + }, + { + "type": "stat", + "title": "Live KV keys", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 4, + "w": 3, + "x": 12, + "y": 1 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "text", + "value": null + } + ] + } + }, + "overrides": [] + }, + "options": { + "colorMode": "value", + "graphMode": "area", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + }, + "targets": [ + { + "expr": "sum(mori_umbp_client_kv_live_count{node=~\"$node\", tier=~\"$tier\"})", + "legendFormat": "keys", + "instant": true, + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "id": 6 + }, + { + "type": "stat", + "title": "Write bandwidth", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 4, + "w": 4, + "x": 15, + "y": 1 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "text", + "value": null + } + ] + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "colorMode": "value", + "graphMode": "area", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + }, + "targets": [ + { + "expr": "sum(rate(mori_umbp_backend_bytes_total{node=~\"$node\", tier=~\"$tier\", op=\"commit\"}[$__rate_interval]))", + "legendFormat": "commit", + "instant": true, + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Bytes landing in local media across every backend.", + "id": 7 + }, + { + "type": "stat", + "title": "Read bandwidth", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 4, + "w": 5, + "x": 19, + "y": 1 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "text", + "value": null + } + ] + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "colorMode": "value", + "graphMode": "area", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + }, + "targets": [ + { + "expr": "sum(rate(mori_umbp_backend_bytes_total{node=~\"$node\", tier=~\"$tier\", op=\"resolve\"}[$__rate_interval]))", + "legendFormat": "resolve", + "instant": true, + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "id": 8 + }, + { + "type": "timeseries", + "title": "Capacity used / total by medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 0, + "y": 5 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "bytes" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier) (mori_umbp_client_capacity_used_bytes{node=~\"$node\", tier=~\"$tier\"})", + "legendFormat": "{{tier}} used", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + }, + { + "expr": "sum by (tier) (mori_umbp_client_capacity_total_bytes{node=~\"$node\", tier=~\"$tier\"})", + "legendFormat": "{{tier}} total", + "refId": "B", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Grouped by tier, so every medium on every selected node shares one panel.", + "id": 9 + }, + { + "type": "timeseries", + "title": "Capacity utilization by medium and node", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 8, + "y": 5 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "percentunit", + "min": 0, + "max": 1 + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "mori_umbp_client_capacity_utilization_ratio{node=~\"$node\", tier=~\"$tier\"}", + "legendFormat": "{{node}} {{tier}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Per node so a single full peer is visible behind a healthy fleet average.", + "id": 10 + }, + { + "type": "timeseries", + "title": "Live KV keys by medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 16, + "y": 5 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + } + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier) (mori_umbp_client_kv_live_count{node=~\"$node\", tier=~\"$tier\"})", + "legendFormat": "{{tier}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Keys each medium currently owns, as reported by its own backend through the heartbeat.", + "id": 11 + }, + { + "type": "row", + "title": "Storage backends \u2014 every medium, same panels", + "collapsed": false, + "gridPos": { + "h": 1, + "w": 24, + "x": 0, + "y": 13 + }, + "panels": [], + "id": 12 + }, + { + "type": "timeseries", + "title": "Operation rate by medium and op", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 0, + "y": 14 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "ops" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier, op) (rate(mori_umbp_backend_ops_total{node=~\"$node\", tier=~\"$tier\"}[$__rate_interval]))", + "legendFormat": "{{tier}} {{op}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "mori_umbp_backend_ops_total is emitted by the instrumentation decorator that wraps every MediumBackend, so a medium added later appears here without a dashboard edit.", + "id": 13 + }, + { + "type": "timeseries", + "title": "Write bandwidth by medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 8, + "y": 14 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier) (rate(mori_umbp_backend_bytes_total{node=~\"$node\", tier=~\"$tier\", op=\"commit\"}[$__rate_interval]))", + "legendFormat": "{{tier}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Bytes committed into each medium. Logical bytes \u2014 see the medium-internal panel below for what actually reached a device.", + "id": 14 + }, + { + "type": "timeseries", + "title": "Read bandwidth by medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 16, + "y": 14 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier) (rate(mori_umbp_backend_bytes_total{node=~\"$node\", tier=~\"$tier\", op=\"resolve\"}[$__rate_interval]))", + "legendFormat": "{{tier}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Bytes handed to readers out of each medium.", + "id": 15 + }, + { + "type": "timeseries", + "title": "Resolve hit rate by medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 0, + "y": 22 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "percentunit", + "min": 0, + "max": 1 + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier) (rate(mori_umbp_backend_ops_total{node=~\"$node\", tier=~\"$tier\", op=\"resolve\", status=\"ok\"}[$__rate_interval])) / clamp_min(sum by (tier) (rate(mori_umbp_backend_ops_total{node=~\"$node\", tier=~\"$tier\", op=\"resolve\"}[$__rate_interval])), 1e-9)", + "legendFormat": "{{tier}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Hits over hits+misses at the backend boundary. A medium whose hit rate collapses is either being asked for keys it never held or is losing them to eviction.", + "id": 16 + }, + { + "type": "timeseries", + "title": "Mean call latency by medium and op", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 8, + "y": 22 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "s" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier, op) (rate(mori_umbp_backend_op_seconds_total{node=~\"$node\", tier=~\"$tier\"}[$__rate_interval])) / clamp_min(sum by (tier, op) (rate(mori_umbp_backend_batches_total{node=~\"$node\", tier=~\"$tier\"}[$__rate_interval])), 1e-9)", + "legendFormat": "{{tier}} {{op}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Seconds spent inside the medium per batch call. Measured by the decorator, so it excludes the network and isolates the medium itself.", + "id": 17 + }, + { + "type": "timeseries", + "title": "Mean batch depth by medium and op", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 16, + "y": 22 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + } + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier, op) (rate(mori_umbp_backend_ops_total{node=~\"$node\", tier=~\"$tier\"}[$__rate_interval])) / clamp_min(sum by (tier, op) (rate(mori_umbp_backend_batches_total{node=~\"$node\", tier=~\"$tier\"}[$__rate_interval])), 1e-9)", + "legendFormat": "{{tier}} {{op}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Keys per call. A depth that drops toward 1 means the caller stopped batching, which shows up as latency long before it shows up as an error.", + "id": 18 + }, + { + "type": "timeseries", + "title": "Failed and refused operations by medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 0, + "y": 30 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "ops" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier, op, status) (rate(mori_umbp_backend_ops_total{node=~\"$node\", tier=~\"$tier\", status=~\"failed|no_space\"}[$__rate_interval]))", + "legendFormat": "{{tier}} {{op}} {{status}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "no_space is the medium being full \u2014 the writer retries elsewhere \u2014 while failed is the medium refusing. Empty is the healthy state.", + "id": 19 + }, + { + "type": "timeseries", + "title": "Eviction rate by medium", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 8, + "y": 30 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "ops" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier) (rate(mori_umbp_backend_ops_total{node=~\"$node\", tier=~\"$tier\", op=\"evict\", status=\"ok\"}[$__rate_interval]))", + "legendFormat": "{{tier}} evicted", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + }, + { + "expr": "sum by (tier) (rate(mori_umbp_backend_ops_total{node=~\"$node\", tier=~\"$tier\", op=\"evict\", status=\"miss\"}[$__rate_interval]))", + "legendFormat": "{{tier}} not freed", + "refId": "B", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "A key that freed no bytes was unknown, already gone, or protected by a live read lease; master retries those on a later round.", + "id": 20 + }, + { + "type": "timeseries", + "title": "Bytes reclaimed by eviction", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 16, + "y": 30 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier) (rate(mori_umbp_backend_bytes_total{node=~\"$node\", tier=~\"$tier\", op=\"evict\"}[$__rate_interval]))", + "legendFormat": "{{tier}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "id": 21 + }, + { + "type": "timeseries", + "title": "Medium-internal events", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 38 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "ops" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier, event) (rate(mori_umbp_backend_medium_events_total{node=~\"$node\", tier=~\"$tier\"}[$__rate_interval]))", + "legendFormat": "{{tier}} {{event}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "What a medium reports about itself \u2014 drive read outcomes, single-flight coalescing, eviction rounds, staging pressure. The medium names the EVENT, never the metric, so each new backend's internals land in this panel automatically.", + "id": 22 + }, + { + "type": "timeseries", + "title": "Medium-internal byte counters", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 38 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (tier, event) (rate(mori_umbp_backend_medium_bytes_total{node=~\"$node\", tier=~\"$tier\"}[$__rate_interval]))", + "legendFormat": "{{tier}} {{event}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Bytes that actually crossed a device boundary, as opposed to the logical bytes above. The gap between them is what caching and single-flight merging saved.", + "id": 23 + }, + { + "type": "timeseries", + "title": "Medium-internal state", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 6, + "w": 24, + "x": 0, + "y": 46 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + } + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "mori_umbp_backend_medium_state{node=~\"$node\", tier=~\"$tier\"}", + "legendFormat": "{{node}} {{tier}} {{state}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Live gauges a medium publishes (staging arena occupancy, queue depth). Backends with nothing to report contribute no series.", + "id": 24 + }, + { + "type": "row", + "title": "Transfer engines", + "collapsed": false, + "gridPos": { + "h": 1, + "w": 24, + "x": 0, + "y": 52 + }, + "panels": [], + "id": 25 + }, + { + "type": "timeseries", + "title": "Bytes moved by engine and direction", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 0, + "y": 53 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (engine, direction) (rate(mori_umbp_transfer_bytes_total{node=~\"$node\"}[$__rate_interval]))", + "legendFormat": "{{engine}} {{direction}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Measured by the composite engine that dispatches to them, so an engine added with AddEngine() is charted with no metrics code of its own.", + "id": 26 + }, + { + "type": "timeseries", + "title": "Achieved bandwidth by engine and direction", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 8, + "y": 53 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "Bps" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (engine, direction) (rate(mori_umbp_transfer_bytes_total{node=~\"$node\"}[$__rate_interval])) / clamp_min(sum by (engine, direction) (rate(mori_umbp_transfer_seconds_total{node=~\"$node\"}[$__rate_interval])), 1e-9)", + "legendFormat": "{{engine}} {{direction}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Bytes divided by time actually in flight \u2014 the rate while transferring, not averaged over idle gaps.", + "id": 27 + }, + { + "type": "timeseries", + "title": "Plan rate by engine and direction", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 16, + "y": 53 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "ops" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (engine, direction) (rate(mori_umbp_transfer_ops_total{node=~\"$node\", status=\"ok\"}[$__rate_interval]))", + "legendFormat": "{{engine}} {{direction}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "A plan is one grouped unit an engine posts, not one key.", + "id": 28 + }, + { + "type": "timeseries", + "title": "Transfer failures and unroutable items", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 61 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "ops" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (engine, direction) (rate(mori_umbp_transfer_ops_total{node=~\"$node\", status=\"failed\"}[$__rate_interval]))", + "legendFormat": "{{engine}} {{direction}} failed", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + }, + { + "expr": "sum by (node) (rate(mori_umbp_transfer_ops_total{node=~\"$node\", status=\"rejected\"}[$__rate_interval]))", + "legendFormat": "{{node}} unroutable", + "refId": "B", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "Unroutable items are counted against engine=\"none\": no engine would take the (src, dst) pair, which is a routing or registration bug rather than a transport failure.", + "id": 29 + }, + { + "type": "timeseries", + "title": "Ranged remote installs that could not be cached", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 61 + }, + "fieldConfig": { + "defaults": { + "custom": { + "fillOpacity": 0, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": true + }, + "color": { + "mode": "palette-classic" + }, + "unit": "ops" + }, + "overrides": [] + }, + "options": { + "legend": { + "displayMode": "table", + "placement": "bottom", + "calcs": [ + "mean", + "max", + "lastNotNull" + ] + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "expr": "sum by (node) (rate(mori_umbp_ranged_remote_install_failures_total{node=~\"$node\"}[$__rate_interval]))", + "legendFormat": "{{node}}", + "refId": "A", + "datasource": { + "type": "prometheus", + "uid": "${datasource}" + } + } + ], + "description": "A remote ranged object arrived but could not be installed in the local medium, so the next read pays the wire again.", + "id": 30 + } + ] +} diff --git a/examples/monitoring/grafana/dashboards/umbp_data_rate_bandwidth.json b/examples/monitoring/grafana/dashboards/umbp_data_rate_bandwidth.json index 47a2f7a15..6e51484b1 100644 --- a/examples/monitoring/grafana/dashboards/umbp_data_rate_bandwidth.json +++ b/examples/monitoring/grafana/dashboards/umbp_data_rate_bandwidth.json @@ -498,6 +498,158 @@ "refId": "E" } ] + }, + { + "id": 104, + "title": "Ranged Bandwidth", + "type": "row", + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 68 }, + "panels": [] + }, + { + "id": 33, + "title": "BatchPutRanges Per-Call Bandwidth - Local (GiB/s)", + "description": "Mean BatchPutRanges bandwidth (GiB/s, e2e call duration) for local traffic, per client. Bytes are the RANGE bytes committed, not whole objects.", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 69 }, + "fieldConfig": { "defaults": { "unit": "none", "min": 0 } }, + "targets": [ + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_sum{traffic=\"local\"}[30s]) / rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_count{traffic=\"local\"}[30s])", + "legendFormat": "{{client}}", + "refId": "A" + } + ] + }, + { + "id": 34, + "title": "BatchPutRanges Per-Call Bandwidth - Remote (GiB/s)", + "description": "Mean BatchPutRanges bandwidth (GiB/s, e2e call duration) for remote traffic, per client. Bytes are the RANGE bytes committed, not whole objects -- a remote ranged put assembles the object in the scratch arena before sending, so ranges and object size agree here.", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 69 }, + "fieldConfig": { "defaults": { "unit": "none", "min": 0 } }, + "targets": [ + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_sum{traffic=\"remote\"}[30s]) / rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_count{traffic=\"remote\"}[30s])", + "legendFormat": "{{client}}", + "refId": "A" + } + ] + }, + { + "id": 35, + "title": "BatchGetRanges Per-Call Bandwidth - Local (GiB/s)", + "description": "Mean BatchGetRanges bandwidth (GiB/s, e2e call duration) for local traffic, per client. Bytes are the RANGE bytes delivered, not whole objects.", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 77 }, + "fieldConfig": { "defaults": { "unit": "none", "min": 0 } }, + "targets": [ + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_sum{traffic=\"local\"}[30s]) / rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_count{traffic=\"local\"}[30s])", + "legendFormat": "{{client}}", + "refId": "A" + } + ] + }, + { + "id": 36, + "title": "BatchGetRanges Per-Call Bandwidth - Remote (GiB/s)", + "description": "Mean BatchGetRanges bandwidth (GiB/s, e2e call duration) for remote traffic, per client. Bytes are the RANGE bytes delivered, not whole objects -- a remote ranged get still pulls the FULL object over the wire into the scratch arena, so this understates wire traffic; see the inbound/outbound get byte counters above for that.", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 77 }, + "fieldConfig": { "defaults": { "unit": "none", "min": 0 } }, + "targets": [ + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_sum{traffic=\"remote\"}[30s]) / rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_count{traffic=\"remote\"}[30s])", + "legendFormat": "{{client}}", + "refId": "A" + } + ] + }, + { + "id": 37, + "title": "Global BatchPutRanges Per-Call Bandwidth (GiB/s, all clients)", + "description": "Global ranged-put bandwidth aggregated across all clients and traffic types (local+remote). Mean is exact (sum/count); percentiles use histogram_quantile over aggregated buckets.", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 85 }, + "fieldConfig": { "defaults": { "unit": "none", "min": 0 } }, + "targets": [ + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "sum(rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_sum[30s])) / sum(rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_count[30s]))", + "legendFormat": "mean", + "refId": "A" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.50, sum by (le)(rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p50", + "refId": "B" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.90, sum by (le)(rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p90", + "refId": "C" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.95, sum by (le)(rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p95", + "refId": "D" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.99, sum by (le)(rate(mori_umbp_client_batch_put_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p99", + "refId": "E" + } + ] + }, + { + "id": 38, + "title": "Global BatchGetRanges Per-Call Bandwidth (GiB/s, all clients)", + "description": "Global ranged-get bandwidth aggregated across all clients and traffic types (local+remote). Mean is exact (sum/count); percentiles use histogram_quantile over aggregated buckets. Compare against sglang:tree_connector_load_bandwidth_gb_s, which counts the same range bytes on the caller side.", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 85 }, + "fieldConfig": { "defaults": { "unit": "none", "min": 0 } }, + "targets": [ + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "sum(rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_sum[30s])) / sum(rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_count[30s]))", + "legendFormat": "mean", + "refId": "A" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.50, sum by (le)(rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p50", + "refId": "B" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.90, sum by (le)(rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p90", + "refId": "C" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.95, sum by (le)(rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p95", + "refId": "D" + } +, { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "expr": "histogram_quantile(0.99, sum by (le)(rate(mori_umbp_client_batch_get_ranges_bandwidth_gibps_bucket[30s])))", + "legendFormat": "p99", + "refId": "E" + } + ] } ] } diff --git a/examples/monitoring/grafana/dashboards/umbp_ssd_tier.json b/examples/monitoring/grafana/dashboards/umbp_ssd_tier.json deleted file mode 100644 index 21af055d7..000000000 --- a/examples/monitoring/grafana/dashboards/umbp_ssd_tier.json +++ /dev/null @@ -1,142 +0,0 @@ -{ - "title": "UMBP Master - SSD Metrics", - "uid": "9b1d4e2a-ssd-tier-umbp-0001", - "tags": ["umbp", "ssd", "tier", "cold-tier"], - "schemaVersion": 36, - "refresh": "5s", - "time": { "from": "now-1h", "to": "now" }, - "templating": { - "list": [ - { - "type": "datasource", - "name": "datasource", - "label": "Data Source", - "pluginId": "prometheus", - "current": { "text": "Prometheus", "value": "PBFA97CFB590B2093" }, - "hide": 0, - "query": "prometheus" - }, - { - "type": "query", - "name": "node", - "label": "Client (node label)", - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "query": "label_values(mori_umbp_ssd_copy_enqueued_total, node)", - "refresh": 2, - "includeAll": true, - "multi": true, - "allValue": ".+", - "current": { "text": "All", "value": "$__all" } - } - ] - }, - "panels": [ - { - "type": "stat", - "title": "SSD reads served (ok, total)", - "gridPos": { "h": 4, "w": 6, "x": 0, "y": 0 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "color": { "mode": "thresholds" }, "thresholds": { "steps": [ { "color": "green", "value": null } ] } } }, - "targets": [ { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(mori_umbp_ssd_read_total{node=~\"$node\", status=\"ok\"})", "legendFormat": "read ok" } ] - }, - { - "type": "stat", - "title": "SSD copies succeeded (total)", - "gridPos": { "h": 4, "w": 6, "x": 6, "y": 0 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "targets": [ { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(mori_umbp_ssd_copy_succeeded_total{node=~\"$node\"})", "legendFormat": "copy ok" } ] - }, - { - "type": "stat", - "title": "SSD capacity used", - "gridPos": { "h": 4, "w": 6, "x": 12, "y": 0 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "unit": "bytes" } }, - "targets": [ { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(mori_umbp_client_capacity_used_bytes{node=~\"$node\", tier=\"SSD\"})", "legendFormat": "used" } ] - }, - { - "type": "stat", - "title": "Copy failed + dropped (total)", - "gridPos": { "h": 4, "w": 6, "x": 18, "y": 0 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "color": { "mode": "thresholds" }, "thresholds": { "steps": [ { "color": "green", "value": null }, { "color": "red", "value": 1 } ] } } }, - "targets": [ { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(mori_umbp_ssd_copy_failed_total{node=~\"$node\"}) + sum(mori_umbp_ssd_copy_dropped_total{node=~\"$node\"})", "legendFormat": "failed+dropped" } ] - }, - { - "type": "timeseries", - "title": "SSD copy-on-commit throughput (ops/s)", - "description": "enqueued vs succeeded vs failed copy tasks (rate). enqueued≈succeeded with ~0 failed is healthy.", - "gridPos": { "h": 8, "w": 12, "x": 0, "y": 4 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "unit": "ops", "custom": { "drawStyle": "line", "fillOpacity": 10 } } }, - "targets": [ - { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_copy_enqueued_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "enqueued" }, - { "refId": "B", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_copy_succeeded_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "succeeded" }, - { "refId": "C", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_copy_failed_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "failed" }, - { "refId": "D", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum by (reason) (rate(mori_umbp_ssd_copy_dropped_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "dropped {{reason}}" } - ] - }, - { - "type": "timeseries", - "title": "SSD read outcomes (ops/s, by status)", - "description": "ok / not_found / no_slot / size_too_large / error. no_slot is transient (not a miss).", - "gridPos": { "h": 8, "w": 12, "x": 12, "y": 4 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "unit": "ops", "custom": { "drawStyle": "line", "fillOpacity": 20, "stacking": { "mode": "normal" } } } }, - "targets": [ - { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum by (status) (rate(mori_umbp_ssd_read_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "{{status}}" } - ] - }, - { - "type": "timeseries", - "title": "SSD capacity used / total (bytes)", - "gridPos": { "h": 8, "w": 12, "x": 0, "y": 12 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "unit": "bytes", "custom": { "drawStyle": "line", "fillOpacity": 10 } } }, - "targets": [ - { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "mori_umbp_client_capacity_used_bytes{node=~\"$node\", tier=\"SSD\"}", "legendFormat": "used {{node}}" }, - { "refId": "B", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "mori_umbp_client_capacity_total_bytes{node=~\"$node\", tier=\"SSD\"}", "legendFormat": "total {{node}}" } - ] - }, - { - "type": "timeseries", - "title": "SSD local eviction (ops/s & bytes/s)", - "description": "Rounds / victims / backend failures (left), bytes freed (right). Stays flat until SSD hits the high watermark.", - "gridPos": { "h": 8, "w": 12, "x": 12, "y": 12 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "unit": "ops", "custom": { "drawStyle": "line", "fillOpacity": 10 } }, "overrides": [ { "matcher": { "id": "byName", "options": "bytes_freed/s" }, "properties": [ { "id": "unit", "value": "Bps" }, { "id": "custom.axisPlacement", "value": "right" } ] } ] }, - "targets": [ - { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_eviction_rounds_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "rounds/s" }, - { "refId": "B", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_eviction_victims_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "victims/s" }, - { "refId": "C", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_eviction_backend_failed_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "backend_failed/s" }, - { "refId": "D", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_eviction_bytes_freed_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "bytes_freed/s" } - ] - }, - { - "type": "timeseries", - "title": "SSD IO throughput (offered, bytes/s)", - "description": "Achieved/offered SSD throughput = SSD IO bytes / wall-clock over the window (rate()), summed across ranks. This is a LOAD signal (how much the workload drove the SSD tier), NOT the device's per-IO bandwidth: it averages in idle gaps so it under-reads bursty copy. For true device bandwidth (bytes / actual IO time) use fio or add ssd_*_io_seconds_total counters.", - "gridPos": { "h": 8, "w": 24, "x": 0, "y": 20 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "unit": "Bps", "custom": { "drawStyle": "line", "fillOpacity": 15 } } }, - "targets": [ - { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_copy_bytes_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "SSD write (copy)" }, - { "refId": "B", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_read_bytes_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "SSD read" } - ] - }, - { - "type": "timeseries", - "title": "SSD read staging health", - "description": "slots in use (gauge), plus slot_full_rejects / expired_reclaims / reader-local transient (ops/s). Non-zero rejects/transient = staging pressure.", - "gridPos": { "h": 8, "w": 24, "x": 0, "y": 28 }, - "datasource": { "type": "prometheus", "uid": "${datasource}" }, - "fieldConfig": { "defaults": { "custom": { "drawStyle": "line", "fillOpacity": 10 } } }, - "targets": [ - { "refId": "A", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(mori_umbp_ssd_staging_slots_in_use{node=~\"$node\"})", "legendFormat": "slots in use" }, - { "refId": "B", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_staging_slot_full_rejects_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "slot_full_rejects/s" }, - { "refId": "C", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_staging_expired_reclaims_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "expired_reclaims/s" }, - { "refId": "D", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "expr": "sum(rate(mori_umbp_ssd_read_client_transient_total{node=~\"$node\"}[$__rate_interval]))", "legendFormat": "client_transient/s" } - ] - } - ] -} diff --git a/examples/umbp/multi_backend_policy.json b/examples/umbp/multi_backend_policy.json new file mode 100644 index 000000000..6992dda27 --- /dev/null +++ b/examples/umbp/multi_backend_policy.json @@ -0,0 +1,55 @@ +{ + "schema_version": 1, + "entry_tier": "hot", + "backends": { + "hbm": { + "type": "hbm", + "capacity": "80GiB", + "devices": [0, 1] + }, + "dram": { + "type": "dram", + "capacity": "512GiB", + "numa_node": 0 + }, + "ssd_a": { + "type": "ssd", + "capacity": "1TiB", + "path": "/mnt/kvcache/hot" + }, + "ssd_b": { + "type": "ssd", + "capacity": "3TiB", + "path": "/mnt/kvcache/cold" + } + }, + "tiers": [ + { + "name": "hot", + "backends": { + "hbm": 100 + }, + "offload_to": ["warm"], + "offload_trigger": "on_evict" + }, + { + "name": "warm", + "backends": { + "dram": 70, + "ssd_a": 30 + }, + "offload_to": ["cold"], + "offload_trigger": "watermark", + "high_watermark": 0.9, + "low_watermark": 0.7 + }, + { + "name": "cold", + "backends": { + "ssd_b": 100 + }, + "promote_trigger": "on_read", + "promote_mode": "copy" + } + ] +} diff --git a/include/mori/metrics/prometheus_metrics_server.hpp b/include/mori/metrics/prometheus_metrics_server.hpp index ca71633a8..b7e7594be 100644 --- a/include/mori/metrics/prometheus_metrics_server.hpp +++ b/include/mori/metrics/prometheus_metrics_server.hpp @@ -102,9 +102,16 @@ class MetricsServer { void setGauge(std::string_view name, std::string_view help, const Labels& labels, double value); // Add `delta` to a monotonically-increasing counter (default delta = 1). - void addCounter(std::string_view name, std::string_view help, uint64_t delta = 1); + // + // `delta` is a double because a Prometheus counter is a double, and because + // not every counter counts whole things: a *_seconds_total accumulates + // fractions of a second, and taking a uint64 here silently floored every one + // of them to zero. Whole-number counters are unaffected — see FormatValue in + // the .cpp, which serialises an integral double exactly rather than through + // the six-significant-digit default that would round a byte total. + void addCounter(std::string_view name, std::string_view help, double delta = 1); void addCounter(std::string_view name, std::string_view help, const Labels& labels, - uint64_t delta = 1); + double delta = 1); // Record one histogram observation. // `bounds` is an ascending list of finite upper-bound values; the implicit @@ -149,12 +156,12 @@ class MetricsServer { struct CounterEntry { std::string help; - uint64_t value{0}; + double value{0.0}; }; struct LabeledCounterFamily { std::string help; - std::map series; + std::map series; }; struct LabeledGaugeFamily { diff --git a/python/mori/umbp/__init__.py b/python/mori/umbp/__init__.py index 5332f4e3e..a2482462e 100644 --- a/python/mori/umbp/__init__.py +++ b/python/mori/umbp/__init__.py @@ -100,11 +100,13 @@ def _configure_packaged_umbp_standalone_server() -> None: UMBPDistributedConfig, UMBPDramConfig, UMBPDurabilityMode, + UMBPHbmConfig, UMBPHostBufferBacking, UMBPHostBufferHandle, UMBPHostMemAllocator, UMBPIoBackend, UMBPIoConfig, + UMBPMedium, UMBPRole, UMBPStandaloneProcessConfig, UMBPSsdConfig, @@ -125,11 +127,13 @@ def _configure_packaged_umbp_standalone_server() -> None: "UMBPDramConfig", "UMBPDurabilityConfig", "UMBPDurabilityMode", + "UMBPHbmConfig", "UMBPHostBufferBacking", "UMBPHostBufferHandle", "UMBPHostMemAllocator", "UMBPIoBackend", "UMBPIoConfig", + "UMBPMedium", "UMBPRole", "UMBPStandaloneProcessConfig", "UMBPSsdConfig", diff --git a/src/metrics/prometheus_metrics_server.cpp b/src/metrics/prometheus_metrics_server.cpp index 5fefb5f5f..208843b7d 100644 --- a/src/metrics/prometheus_metrics_server.cpp +++ b/src/metrics/prometheus_metrics_server.cpp @@ -49,9 +49,13 @@ #include #include +#include #include +#include +#include #include #include +#include #include "mori/utils/mori_log.hpp" @@ -137,7 +141,7 @@ void MetricsServer::setGauge(std::string_view name, std::string_view help, doubl entry.value = value; } -void MetricsServer::addCounter(std::string_view name, std::string_view help, uint64_t delta) { +void MetricsServer::addCounter(std::string_view name, std::string_view help, double delta) { std::lock_guard lk(mutex_); auto& entry = counters_[std::string(name)]; entry.help = help; @@ -256,8 +260,28 @@ static std::string FormatLabels(const mori::metrics::MetricsServer::Labels& labe return s; } +namespace { + +// Render a metric value for the exposition format. +// +// `ostream <<` on a double defaults to six significant digits, which turns a +// byte total into 8.58993e+09 and throws away everything past the seventh +// digit. Counters and gauges here carry exact byte counts and key counts, so +// an integral value is printed as an integer and only a genuinely fractional +// one (a *_seconds_total) falls back to full round-trip precision. +std::string FormatValue(double v) { + if (std::isfinite(v) && v == std::floor(v) && std::fabs(v) < 1e18) { + return std::to_string(static_cast(v)); + } + std::ostringstream ss; + ss << std::setprecision(std::numeric_limits::max_digits10) << v; + return ss.str(); +} + +} // namespace + void MetricsServer::addCounter(std::string_view name, std::string_view help, const Labels& labels, - uint64_t delta) { + double delta) { auto label_str = FormatLabels(labels); std::lock_guard lk(mutex_); auto& family = labeled_counters_[std::string(name)]; @@ -291,14 +315,14 @@ std::string MetricsServer::SerializeMaps( for (const auto& [name, g] : gauges) { out << "# HELP " << name << " " << g.help << "\n"; out << "# TYPE " << name << " gauge\n"; - out << name << " " << g.value << "\n\n"; + out << name << " " << FormatValue(g.value) << "\n\n"; } // Counters for (const auto& [name, c] : counters) { out << "# HELP " << name << " " << c.help << "\n"; out << "# TYPE " << name << " counter\n"; - out << name << " " << c.value << "\n\n"; + out << name << " " << FormatValue(c.value) << "\n\n"; } // Labeled gauges @@ -306,7 +330,7 @@ std::string MetricsServer::SerializeMaps( out << "# HELP " << name << " " << family.help << "\n"; out << "# TYPE " << name << " gauge\n"; for (const auto& [label_str, value] : family.series) { - out << name << label_str << " " << value << "\n"; + out << name << label_str << " " << FormatValue(value) << "\n"; } out << "\n"; } @@ -316,7 +340,7 @@ std::string MetricsServer::SerializeMaps( out << "# HELP " << name << " " << family.help << "\n"; out << "# TYPE " << name << " counter\n"; for (const auto& [label_str, value] : family.series) { - out << name << label_str << " " << value << "\n"; + out << name << label_str << " " << FormatValue(value) << "\n"; } out << "\n"; } @@ -332,7 +356,7 @@ std::string MetricsServer::SerializeMaps( } // +Inf bucket == total observation count. out << name << "_bucket{le=\"+Inf\"} " << h.count << "\n"; - out << name << "_sum " << h.sum << "\n"; + out << name << "_sum " << FormatValue(h.sum) << "\n"; out << name << "_count " << h.count << "\n\n"; } @@ -349,7 +373,7 @@ std::string MetricsServer::SerializeMaps( << h.bucket_counts[i] << "\n"; } out << name << "_bucket" << bucket_prefix << "le=\"+Inf\"} " << h.count << "\n"; - out << name << "_sum" << label_str << " " << h.sum << "\n"; + out << name << "_sum" << label_str << " " << FormatValue(h.sum) << "\n"; out << name << "_count" << label_str << " " << h.count << "\n"; } out << "\n"; diff --git a/src/pybind/pybind_umbp.cpp b/src/pybind/pybind_umbp.cpp index 5e7065564..be4b49eca 100644 --- a/src/pybind/pybind_umbp.cpp +++ b/src/pybind/pybind_umbp.cpp @@ -166,6 +166,16 @@ void RegisterMoriUmbp(py::module_& m) { .def_readwrite("capacity_bytes", &UMBPSsdConfig::capacity_bytes) .def_readwrite("layout_mode", &UMBPSsdConfig::layout_mode) .def_readwrite("segment_size_bytes", &UMBPSsdConfig::segment_size_bytes) + // Multi-drive: storage_dir accepts a comma-separated list (one per drive); + // shard_io_threads=0 means one worker per drive. + .def_readwrite("shard_io_threads", &UMBPSsdConfig::shard_io_threads) + .def_readwrite("tier_io_threads", &UMBPSsdConfig::tier_io_threads) + .def_readwrite("direct_io", &UMBPSsdConfig::direct_io) + .def_readwrite("verify_crc", &UMBPSsdConfig::verify_crc) + // sglang's UMBPStore builds UMBPConfig directly rather than via + // FromEnvironment(), so UMBP_SSD_SINGLE_FLIGHT never reaches it — this is + // the only way to turn coalescing off on that path. + .def_readwrite("single_flight_reads", &UMBPSsdConfig::single_flight_reads) .def_readwrite("high_watermark", &UMBPSsdConfig::high_watermark) .def_readwrite("low_watermark", &UMBPSsdConfig::low_watermark) .def_readwrite("io", &UMBPSsdConfig::io) @@ -223,18 +233,39 @@ void RegisterMoriUmbp(py::module_& m) { .value("NEVER", CacheRemoteAdmission::NEVER) .value("ALWAYS", CacheRemoteAdmission::ALWAYS); + py::enum_(m, "UMBPMedium") + .value("DRAM", UMBPMedium::DRAM) + .value("HBM", UMBPMedium::HBM) + .value("SSD", UMBPMedium::SSD); + + py::class_(m, "UMBPHbmConfig") + .def(py::init<>()) + .def_readwrite("device", &UMBPHbmConfig::device) + .def_readwrite("capacity_bytes", &UMBPHbmConfig::capacity_bytes); + py::class_(m, "UMBPDistributedConfig") .def(py::init<>()) .def_readwrite("master_config", &UMBPDistributedConfig::master_config) .def_readwrite("io_engine", &UMBPDistributedConfig::io_engine) .def_readwrite("staging_buffer_size", &UMBPDistributedConfig::staging_buffer_size) + .def_readwrite("ranged_scratch_size", &UMBPDistributedConfig::ranged_scratch_size) .def_readwrite("ssd_staging_buffer_size", &UMBPDistributedConfig::ssd_staging_buffer_size) .def_readwrite("ssd_staging_buffer_slots", &UMBPDistributedConfig::ssd_staging_buffer_slots) + .def_readwrite("ssd_staging_use_hugepages", &UMBPDistributedConfig::ssd_staging_use_hugepages) + .def_readwrite("ssd_staging_hugepage_size", &UMBPDistributedConfig::ssd_staging_hugepage_size) .def_readwrite("peer_service_port", &UMBPDistributedConfig::peer_service_port) .def_readwrite("cache_remote_fetches", &UMBPDistributedConfig::cache_remote_fetches) + .def_readwrite("ranged_locality_prefetch", &UMBPDistributedConfig::ranged_locality_prefetch) + .def_readwrite("local_first", &UMBPDistributedConfig::local_first) .def_readwrite("cache_remote_admission", &UMBPDistributedConfig::cache_remote_admission) .def_readwrite("admission_max_block_bytes", &UMBPDistributedConfig::admission_max_block_bytes) - .def_readwrite("dram_page_size", &UMBPDistributedConfig::dram_page_size); + .def_readwrite("dram_page_size", &UMBPDistributedConfig::dram_page_size) + .def_readwrite("backend_policy_path", &UMBPDistributedConfig::backend_policy_path) + .def_readwrite("workload_trace_path", &UMBPDistributedConfig::workload_trace_path) + .def_readwrite("workload_trace_client_id", &UMBPDistributedConfig::workload_trace_client_id) + .def_readwrite("workload_trace_seed", &UMBPDistributedConfig::workload_trace_seed) + .def_readwrite("medium", &UMBPDistributedConfig::medium) + .def_readwrite("hbm", &UMBPDistributedConfig::hbm); py::class_(m, "UMBPStandaloneProcessConfig") .def(py::init<>()) @@ -275,6 +306,12 @@ void RegisterMoriUmbp(py::module_& m) { py::call_guard()) .def("batch_get_into_ptr", &IUMBPClient::BatchGet, py::arg("keys"), py::arg("ptrs"), py::arg("sizes"), py::call_guard()) + .def("batch_get_ranges_into_ptr", &IUMBPClient::BatchGetRanges, py::arg("keys"), + py::arg("ptrs"), py::arg("sizes"), py::arg("src_offsets"), + py::call_guard()) + .def("batch_put_ranges_from_ptr", &IUMBPClient::BatchPutRanges, py::arg("keys"), + py::arg("object_sizes"), py::arg("ptrs"), py::arg("sizes"), py::arg("dst_offsets"), + py::call_guard()) .def("batch_exists", &IUMBPClient::BatchExists, py::arg("keys"), py::call_guard()) .def("batch_exists_consecutive", &IUMBPClient::BatchExistsConsecutive, py::arg("keys"), @@ -283,8 +320,20 @@ void RegisterMoriUmbp(py::module_& m) { .def("flush", &IUMBPClient::Flush, py::call_guard()) .def("is_distributed", &IUMBPClient::IsDistributed) // pure getter, no I/O .def("get_deployment_mode", &IUMBPClient::GetDeploymentMode) // pure getter, no I/O - .def("register_memory", &IUMBPClient::RegisterMemory, py::arg("ptr"), py::arg("size"), - py::call_guard()) + .def("get_backend_mode", &IUMBPClient::GetBackendMode) // pure getter, no I/O + .def("supports_ranged_io", &IUMBPClient::SupportsRangedIO) // pure getter, no I/O + // Bound to the pinned form on purpose. MemoryRegistration selects + // whether a region is pinned for RDMA or merely recorded for local + // copies, and that is decided by the deployment -- the standalone server + // picks it from its inner backend mode. A caller reaching this binding + // is registering its own buffers and always wants them pinned, so the + // knob is kept off the Python surface rather than offered and ignored. + .def( + "register_memory", + [](IUMBPClient& self, uintptr_t ptr, size_t size, mori::io::MemoryLocationType loc, + int device) { return self.RegisterMemory(ptr, size, loc, device); }, + py::arg("ptr"), py::arg("size"), py::arg("loc") = mori::io::MemoryLocationType::CPU, + py::arg("device") = -1, py::call_guard()) .def("deregister_memory", &IUMBPClient::DeregisterMemory, py::arg("ptr"), py::call_guard()) .def("report_external_kv_blocks", &IUMBPClient::ReportExternalKvBlocks, py::arg("hashes"), diff --git a/src/umbp/CMakeLists.txt b/src/umbp/CMakeLists.txt index 8a0d79110..d710741d1 100644 --- a/src/umbp/CMakeLists.txt +++ b/src/umbp/CMakeLists.txt @@ -91,22 +91,22 @@ endif() # Core static library set(UMBP_CORE_SRCS - local/block_index/local_block_index.cpp + common/device_copy.cpp storage/io/storage_io_driver.cpp storage/io/posix_storage_io_driver.cpp storage/io/io_uring_storage_io_driver.cpp - local/tiers/copy_pipeline.cpp local/tiers/segment/segment_format.cpp local/tiers/segment/segment_index.cpp local/tiers/segment/segment_scanner.cpp local/tiers/segment/segment_writer.cpp local/tiers/tier_backend.cpp - local/tiers/dram_tier.cpp + local/tiers/host_registration.cpp + local/tiers/device_gather.hip local/host_mem_allocator.cpp local/tiers/ssd_tier.cpp + local/tiers/sharded_ssd_tier.cpp + local/tiers/ssd_backend_factory.cpp local/tiers/dummy_ssd_tier.cpp - local/tiers/local_storage_manager.cpp - local/standalone_client.cpp storage/spdk/offset_allocator.cpp) # SPDK Proxy client (shared memory IPC, no SPDK dependency — Linux only) @@ -123,6 +123,10 @@ endif() add_library(umbp_core STATIC ${UMBP_CORE_SRCS}) target_include_directories(umbp_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) target_link_libraries(umbp_core PRIVATE pthread) +# The tiers register their regions for GPU access and gather scattered copies +# in a kernel, so the core library needs the HIP runtime. PUBLIC because no +# tier header exposes HIP types but device_gather.hip is linked in. +target_link_libraries(umbp_core PUBLIC hip::host) if(UNIX) target_link_libraries(umbp_core PRIVATE rt) endif() @@ -255,6 +259,18 @@ elseif(NOT _PROTOBUF_MM) message(FATAL_ERROR "[UMBP] cannot determine libprotobuf header version.") endif() +# umbp_peer.proto declares a proto3 `optional` field. protoc only accepts that +# behind an opt-in flag before 3.15, and Dockerfile.dev still pins Ubuntu +# jammy's 3.12. +set(_PROTOC_PROTO3_OPTIONAL_ARGS) +if(_PROTOC_MM VERSION_LESS "3.15") + list(APPEND _PROTOC_PROTO3_OPTIONAL_ARGS --experimental_allow_proto3_optional) + message( + STATUS + "[UMBP] protoc ${_PROTOC_MM} < 3.15 — enabling --experimental_allow_proto3_optional" + ) +endif() + message( STATUS "[UMBP] gRPC + Protobuf found — distributed control plane ENABLED") @@ -264,6 +280,7 @@ set(UMBP_PROTO_DIR "${CMAKE_CURRENT_SOURCE_DIR}/distributed/proto") set(UMBP_PROTO_FILE "${UMBP_PROTO_DIR}/umbp.proto") set(UMBP_PEER_PROTO_FILE "${UMBP_PROTO_DIR}/umbp_peer.proto") set(UMBP_STANDALONE_PROTO_FILE "${UMBP_PROTO_DIR}/umbp_standalone.proto") +set(UMBP_WORKLOAD_PROTO_FILE "${UMBP_PROTO_DIR}/umbp_workload.proto") set(UMBP_PROTO_GEN_DIR "${CMAKE_CURRENT_BINARY_DIR}/proto_gen") file(MAKE_DIRECTORY ${UMBP_PROTO_GEN_DIR}) @@ -282,13 +299,17 @@ set(UMBP_STANDALONE_PROTO_HDRS "${UMBP_PROTO_GEN_DIR}/umbp_standalone.pb.h") set(UMBP_STANDALONE_GRPC_SRCS "${UMBP_PROTO_GEN_DIR}/umbp_standalone.grpc.pb.cc") set(UMBP_STANDALONE_GRPC_HDRS "${UMBP_PROTO_GEN_DIR}/umbp_standalone.grpc.pb.h") +set(UMBP_WORKLOAD_PROTO_SRCS "${UMBP_PROTO_GEN_DIR}/umbp_workload.pb.cc") +set(UMBP_WORKLOAD_PROTO_HDRS "${UMBP_PROTO_GEN_DIR}/umbp_workload.pb.h") + add_custom_command( OUTPUT ${UMBP_PROTO_SRCS} ${UMBP_PROTO_HDRS} ${UMBP_GRPC_SRCS} ${UMBP_GRPC_HDRS} COMMAND ${_PROTOBUF_PROTOC} --proto_path=${UMBP_PROTO_DIR} --cpp_out=${UMBP_PROTO_GEN_DIR} --grpc_out=${UMBP_PROTO_GEN_DIR} - --plugin=protoc-gen-grpc=${_GRPC_CPP_PLUGIN} ${UMBP_PROTO_FILE} + --plugin=protoc-gen-grpc=${_GRPC_CPP_PLUGIN} + ${_PROTOC_PROTO3_OPTIONAL_ARGS} ${UMBP_PROTO_FILE} DEPENDS ${UMBP_PROTO_FILE} COMMENT "Generating UMBP protobuf/gRPC C++ sources" VERBATIM) @@ -299,7 +320,8 @@ add_custom_command( COMMAND ${_PROTOBUF_PROTOC} --proto_path=${UMBP_PROTO_DIR} --cpp_out=${UMBP_PROTO_GEN_DIR} --grpc_out=${UMBP_PROTO_GEN_DIR} - --plugin=protoc-gen-grpc=${_GRPC_CPP_PLUGIN} ${UMBP_PEER_PROTO_FILE} + --plugin=protoc-gen-grpc=${_GRPC_CPP_PLUGIN} + ${_PROTOC_PROTO3_OPTIONAL_ARGS} ${UMBP_PEER_PROTO_FILE} # umbp_peer.proto imports umbp.proto for shared types (TierType, PageLocation, # BufferMemoryDesc) — protoc reads it at generation time and the generated # peer headers include "umbp.pb.h", so the peer generation depends on the @@ -314,11 +336,25 @@ add_custom_command( COMMAND ${_PROTOBUF_PROTOC} --proto_path=${UMBP_PROTO_DIR} --cpp_out=${UMBP_PROTO_GEN_DIR} --grpc_out=${UMBP_PROTO_GEN_DIR} - --plugin=protoc-gen-grpc=${_GRPC_CPP_PLUGIN} ${UMBP_STANDALONE_PROTO_FILE} + --plugin=protoc-gen-grpc=${_GRPC_CPP_PLUGIN} + ${_PROTOC_PROTO3_OPTIONAL_ARGS} ${UMBP_STANDALONE_PROTO_FILE} DEPENDS ${UMBP_STANDALONE_PROTO_FILE} ${UMBP_PROTO_FILE} ${UMBP_PROTO_HDRS} COMMENT "Generating UMBP Standalone protobuf/gRPC C++ sources" VERBATIM) +add_custom_command( + OUTPUT ${UMBP_WORKLOAD_PROTO_SRCS} ${UMBP_WORKLOAD_PROTO_HDRS} + COMMAND + ${_PROTOBUF_PROTOC} --proto_path=${UMBP_PROTO_DIR} + --cpp_out=${UMBP_PROTO_GEN_DIR} ${_PROTOC_PROTO3_OPTIONAL_ARGS} + ${UMBP_WORKLOAD_PROTO_FILE} + DEPENDS ${UMBP_WORKLOAD_PROTO_FILE} + COMMENT "Generating UMBP workload trace protobuf C++ sources" + VERBATIM) +add_custom_target(umbp_workload_proto + DEPENDS ${UMBP_WORKLOAD_PROTO_SRCS} + ${UMBP_WORKLOAD_PROTO_HDRS}) + # ---------- Static library: umbp_common --------------------------------------- # Shared between PoolClient integration and standalone control-plane tools. @@ -330,6 +366,13 @@ add_library( ${UMBP_PEER_GRPC_SRCS} ${UMBP_STANDALONE_PROTO_SRCS} ${UMBP_STANDALONE_GRPC_SRCS} + ${UMBP_WORKLOAD_PROTO_SRCS} + # Trace capture is a client feature: PoolClient records live traffic when + # UMBP_WORKLOAD_TRACE_PATH is set, so the writer ships with the library. The + # rest of the workload harness drives benchmarks and lives in + # umbp_workload_bench instead. + distributed/benchmark/workload_trace.cpp + distributed/benchmark/workload_trace_recorder.cpp distributed/master/in_memory_master_metadata_store.cpp distributed/master/master_metadata_store_factory.cpp distributed/master/master_server.cpp @@ -340,9 +383,22 @@ add_library( distributed/routing/router.cpp distributed/routing/route_get_strategy.cpp distributed/routing/route_put_strategy.cpp - distributed/peer/peer_dram_allocator.cpp - distributed/peer/peer_ssd_manager.cpp - distributed/peer/ssd_copy_pipeline.cpp + distributed/metrics/component_metrics.cpp + distributed/peer/backend/page_backend.cpp + distributed/peer/backend/hbm_backend.cpp + distributed/peer/backend/ssd_backend.cpp + distributed/peer/backend/instrumented_backend.cpp + distributed/transfer/transfer_engine.cpp + distributed/transfer/mori_io_engine.cpp + distributed/transfer/local_copy_engine.cpp + distributed/transfer/hbm_copy_engine.cpp + distributed/transfer/composite_transfer_engine.cpp + distributed/peer/ssd/peer_ssd_manager.cpp + distributed/peer/ssd/ssd_copy_pipeline.cpp + distributed/pool/policy_config.cpp + distributed/pool/tier_graph.cpp + distributed/pool/tier_transition.cpp + distributed/pool/peer_pool.cpp distributed/peer/peer_service.cpp distributed/pool_client.cpp distributed/distributed_client.cpp @@ -350,6 +406,7 @@ add_library( standalone/ipc.cpp standalone/standalone_process_client.cpp standalone/standalone_server.cpp) +add_dependencies(umbp_common umbp_workload_proto) target_include_directories( umbp_common PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${UMBP_PROTO_GEN_DIR} @@ -374,6 +431,20 @@ target_link_libraries( target_compile_features(umbp_common PUBLIC cxx_std_17) +# ---------- Optional hipfile GDS transfer engine ----------------------------- +# GdsEngine reads an SSD segment's byte range straight into GPU memory via +# hipFileRead (no host bounce). Auto-enabled when hipfile is installed; a build +# without it keeps only the memory engines and the host-staged SSD path. +find_package(hipfile QUIET CONFIG GLOBAL) +if(TARGET hip::hipfile) + message(STATUS "[UMBP] hipfile found — GDS transfer engine ENABLED") + target_sources(umbp_common PRIVATE distributed/transfer/gds_engine.cpp) + target_link_libraries(umbp_common PUBLIC hip::hipfile) + target_compile_definitions(umbp_common PUBLIC UMBP_ENABLE_GDS) +else() + message(STATUS "[UMBP] hipfile not found — GDS transfer engine disabled") +endif() + # ---------- Optional Redis / RESP master metadata backend -------------------- # Makes the master stateless by storing all metadata in an external RESP store # (Redis / Dragonfly / Valkey). OFF by default so existing builds are @@ -425,14 +496,16 @@ if(MORI_UMBP_TESTING) target_compile_definitions(umbp_common PUBLIC MORI_UMBP_TESTING) endif() -# NOTE: umbp_core no longer depends on umbp_common. StandaloneClient has no -# gRPC/Protobuf dependency. Targets that need distributed features should link -# umbp_common directly. +# NOTE: umbp_core no longer depends on umbp_common. It holds the storage tiers +# (SSD/SPDK/segment) and the host allocator, none of which need gRPC/Protobuf; +# umbp_common's PeerSsdManager builds on them. Targets that need the clients +# should link umbp_common directly. target_link_libraries(umbp_core PUBLIC mori_logging) -# Factory: references both StandaloneClient and DistributedClient, so it belongs -# in umbp_common (which has the distributed symbols). Runtime switching is via -# config.distributed.has_value(). +# Factory: every client it can return (DistributedClient, StandaloneProcessClient) +# lives in umbp_common, so it belongs there too. Which one a config selects -- +# including the embedded deployment it synthesizes when a config names none -- +# is decided at runtime in WithEmbeddedDefaults/CreateUMBPClient. target_sources(umbp_common PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/umbp_client_factory.cpp) @@ -460,6 +533,35 @@ set_target_properties(umbp_client PROPERTIES BUILD_RPATH "$ORIGIN;${ROCM_PATH}/l INSTALL_RPATH "$ORIGIN;${ROCM_PATH}/lib" BUILD_WITH_INSTALL_RPATH TRUE) +# ---------- Workload harness for benchmarks ---------------------------------- +# Generates synthetic workloads, replays traces, and drives them against a +# client. Nothing in the data path needs it, so it stays out of umbp_common and +# is linked only by the benchmark tool and the tests that exercise it. +add_library( + umbp_workload_bench STATIC + distributed/benchmark/payload.cpp + distributed/benchmark/workload_source.cpp + distributed/benchmark/workload_runner.cpp) +add_dependencies(umbp_workload_bench umbp_workload_proto) +target_link_libraries(umbp_workload_bench PUBLIC umbp_common) +target_compile_features(umbp_workload_bench PUBLIC cxx_std_17) + +# ---------- Distributed tier-management benchmark --------------------------- +# Manual synthetic/trace-replay tool. It starts an in-process master and real +# PoolClients, so it exercises routing, peer placement, transfer, heartbeat, +# capacity, and eviction without depending on an inference framework. +add_executable(umbp_tier_bench distributed/bin/tier_benchmark_main.cpp) +add_dependencies(umbp_tier_bench umbp_workload_proto) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(umbp_tier_bench PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries(umbp_tier_bench PRIVATE umbp_workload_bench + ${_PROTOBUF_LIBS} + ${_GRPCPP_LIB}) +set_target_properties(umbp_tier_bench PROPERTIES BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE) + # ---------- Standalone-process executable ------------------------------------ add_executable(umbp_standalone_server standalone/bin/standalone_server_main.cpp) diff --git a/src/umbp/common/device_copy.cpp b/src/umbp/common/device_copy.cpp new file mode 100644 index 000000000..2697e2101 --- /dev/null +++ b/src/umbp/common/device_copy.cpp @@ -0,0 +1,80 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#include "umbp/common/device_copy.h" + +#include "mori/utils/mori_log.hpp" + +namespace mori::umbp { + +PointerLocation DetectPointerLocation(const void* ptr) { + if (ptr == nullptr) return {}; + + hipPointerAttribute_t attributes{}; + const hipError_t status = hipPointerGetAttributes(&attributes, ptr); + if (status == hipSuccess) { + if (attributes.type == hipMemoryTypeDevice && attributes.device >= 0) { + return {true, attributes.device}; + } + return {}; + } + + // An attribute miss is normal for malloc/mmap memory. Clear HIP's + // thread-local error so it cannot affect a later runtime call. + (void)hipGetLastError(); + return {}; +} + +ScopedHipDevice::ScopedHipDevice(int device_id) { + if (device_id < 0) return; + if (hipGetDevice(&previous_device_) != hipSuccess) { + (void)hipGetLastError(); + return; + } + if (previous_device_ != device_id) { + if (hipSetDevice(device_id) != hipSuccess) { + (void)hipGetLastError(); + return; + } + changed_ = true; + } + valid_ = true; +} + +ScopedHipDevice::~ScopedHipDevice() { + if (!changed_) return; + if (hipSetDevice(previous_device_) != hipSuccess) (void)hipGetLastError(); +} + +bool DeviceCopy(void* dst, const void* src, size_t size, hipMemcpyKind kind, int device_id) { + const hipError_t status = hipMemcpy(dst, src, size, kind); + if (status != hipSuccess) { + MORI_UMBP_ERROR("[UMBP] hipMemcpy failed on device {}: {}", device_id, + hipGetErrorString(status)); + return false; + } + return true; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/benchmark/payload.cpp b/src/umbp/distributed/benchmark/payload.cpp new file mode 100644 index 000000000..38364f309 --- /dev/null +++ b/src/umbp/distributed/benchmark/payload.cpp @@ -0,0 +1,98 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include "umbp/distributed/benchmark/payload.h" + +#include +#include + +namespace mori::umbp::benchmark { +namespace { + +uint64_t HashKey(std::string_view key) { + uint64_t hash = 1469598103934665603ULL; + for (unsigned char byte : key) { + hash ^= byte; + hash *= 1099511628211ULL; + } + return hash; +} + +uint64_t SplitMix64(uint64_t value) { + value += 0x9e3779b97f4a7c15ULL; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); +} + +uint64_t InitialState(std::string_view key, uint64_t operation_id, uint64_t seed) { + return SplitMix64(HashKey(key) ^ SplitMix64(operation_id) ^ SplitMix64(seed)); +} + +} // namespace + +void FillDeterministicPayload(std::string_view key, uint64_t operation_id, uint64_t seed, + void* destination, size_t size) { + if (size != 0 && destination == nullptr) { + throw std::invalid_argument("payload destination must not be null"); + } + auto* output = static_cast(destination); + uint64_t state = InitialState(key, operation_id, seed); + size_t offset = 0; + while (offset < size) { + state = SplitMix64(state); + const size_t count = std::min(sizeof(state), size - offset); + for (size_t byte = 0; byte < count; ++byte) { + output[offset + byte] = static_cast(state >> (byte * 8)); + } + offset += count; + } +} + +std::vector GenerateDeterministicPayload(std::string_view key, uint64_t operation_id, + uint64_t seed, size_t size) { + std::vector payload(size); + FillDeterministicPayload(key, operation_id, seed, payload.data(), payload.size()); + return payload; +} + +bool ValidateDeterministicPayload(std::string_view key, uint64_t operation_id, uint64_t seed, + const void* data, size_t size) { + if (size != 0 && data == nullptr) return false; + const auto* input = static_cast(data); + uint64_t state = InitialState(key, operation_id, seed); + size_t offset = 0; + while (offset < size) { + state = SplitMix64(state); + const size_t count = std::min(sizeof(state), size - offset); + for (size_t byte = 0; byte < count; ++byte) { + if (input[offset + byte] != static_cast(state >> (byte * 8))) { + return false; + } + } + offset += count; + } + return true; +} + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/distributed/benchmark/workload_runner.cpp b/src/umbp/distributed/benchmark/workload_runner.cpp new file mode 100644 index 000000000..65757544c --- /dev/null +++ b/src/umbp/distributed/benchmark/workload_runner.cpp @@ -0,0 +1,301 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include "umbp/distributed/benchmark/workload_runner.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/payload.h" + +namespace mori::umbp::benchmark { +namespace { + +using Event = ::umbp::benchmark::WorkloadEvent; + +OperationMetrics& MetricsFor(Event::Operation operation, WorkloadMetrics* metrics) { + return operation == Event::PUT ? metrics->puts : metrics->gets; +} + +void AddAttempt(const Event& event, WorkloadMetrics* metrics) { + ++MetricsFor(event.operation(), metrics).attempted; + ++metrics->total.attempted; +} + +void AddResult(const Event& event, bool success, WorkloadMetrics* metrics) { + OperationMetrics& operation = MetricsFor(event.operation(), metrics); + if (success) { + ++operation.succeeded; + operation.succeeded_bytes += event.value_size(); + ++metrics->total.succeeded; + metrics->total.succeeded_bytes += event.value_size(); + } else { + ++operation.failed; + ++metrics->total.failed; + } +} + +void MergeOperation(const OperationMetrics& source, OperationMetrics* destination) { + destination->attempted += source.attempted; + destination->succeeded += source.succeeded; + destination->failed += source.failed; + destination->succeeded_bytes += source.succeeded_bytes; +} + +void MergeMetrics(const WorkloadMetrics& source, WorkloadMetrics* destination) { + MergeOperation(source.total, &destination->total); + MergeOperation(source.puts, &destination->puts); + MergeOperation(source.gets, &destination->gets); + destination->get_misses += source.get_misses; + destination->get_validation_failures += source.get_validation_failures; + destination->latency.samples_ns.insert(destination->latency.samples_ns.end(), + source.latency.samples_ns.begin(), + source.latency.samples_ns.end()); + destination->schedule_lag.samples_ns.insert(destination->schedule_lag.samples_ns.end(), + source.schedule_lag.samples_ns.begin(), + source.schedule_lag.samples_ns.end()); +} + +void FinalizeDistribution(DistributionMetrics* distribution) { + if (distribution->samples_ns.empty()) return; + std::vector sorted = distribution->samples_ns; + std::sort(sorted.begin(), sorted.end()); + auto percentile = [&sorted](double fraction) { + const size_t rank = + std::max(1, static_cast(std::ceil(fraction * sorted.size()))); + return sorted[rank - 1]; + }; + distribution->p50_ns = percentile(0.50); + distribution->p95_ns = percentile(0.95); + distribution->p99_ns = percentile(0.99); + distribution->max_ns = sorted.back(); +} + +uint64_t Nanoseconds(std::chrono::steady_clock::duration duration) { + if (duration <= std::chrono::steady_clock::duration::zero()) return 0; + return static_cast( + std::chrono::duration_cast(duration).count()); +} + +} // namespace + +std::vector WorkloadClient::BatchPut( + uint32_t client_id, const std::vector& keys, + const std::vector>& values) { + if (keys.size() != values.size()) { + throw std::invalid_argument("BatchPut keys and values must have equal sizes"); + } + std::vector results; + results.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + results.push_back(Put(client_id, keys[i], values[i].data(), values[i].size())); + } + return results; +} + +std::vector WorkloadClient::BatchGet(uint32_t client_id, + const std::vector& keys, + const std::vector& sizes, + std::vector>* values) { + if (keys.size() != sizes.size() || values == nullptr) { + throw std::invalid_argument("BatchGet arguments are invalid"); + } + values->clear(); + values->resize(keys.size()); + std::vector results; + results.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + (*values)[i].resize(sizes[i]); + results.push_back(Get(client_id, keys[i], (*values)[i].data(), sizes[i])); + } + return results; +} + +WorkloadRunner::WorkloadRunner(WorkloadClient* client, WorkloadRunnerOptions options) + : client_(client), options_(options) { + if (client_ == nullptr) throw std::invalid_argument("client must not be null"); + if (options_.time_scale <= 0.0 || !std::isfinite(options_.time_scale)) { + throw std::invalid_argument("time_scale must be finite and positive"); + } +} + +WorkloadMetrics WorkloadRunner::Run(WorkloadSource* source) { + if (source == nullptr) throw std::invalid_argument("source must not be null"); + + WorkloadMetrics metrics; + std::map> client_events; + std::unordered_map> put_identity; + Event event; + while (source->Next(&event)) { + if (event.outcome() == Event::OUTCOME_MISS) ++metrics.recorded_misses; + if (event.outcome() == Event::OUTCOME_FAILURE) ++metrics.recorded_failures; + if (event.operation() != Event::PUT && event.operation() != Event::GET) { + throw std::invalid_argument("workload source produced an invalid operation"); + } + if (event.key().empty()) { + throw std::invalid_argument("workload source produced an empty key"); + } + if (event.value_size() > std::numeric_limits::max()) { + throw std::length_error("workload value does not fit in size_t"); + } + // Payload validation derives the expected bytes from (key, operation_id), + // which only works if a key is written once. A recorded production trace + // legitimately overwrites keys, so the invariant applies to generated + // traces being validated, not to replay in general. + if (event.operation() == Event::PUT && options_.validate_get_payloads) { + const auto [found, inserted] = put_identity.emplace( + event.key(), std::make_pair(event.operation_id(), event.value_size())); + if (!inserted) { + const bool conflicting = found->second.first != event.operation_id() || + found->second.second != event.value_size(); + throw std::invalid_argument( + conflicting + ? "workload contains conflicting PUT payload identities for an immutable key" + : "workload contains a duplicate PUT for an immutable key"); + } + } + auto& events = client_events[event.client_id()]; + if (!events.empty() && event.relative_timestamp_ns() < events.back().relative_timestamp_ns()) { + throw std::invalid_argument("per-client workload timestamps must be nondecreasing"); + } + events.push_back(event); + } + + const auto start = Clock::now(); + const uint64_t payload_seed = source->seed(); + std::mutex metrics_mutex; + std::exception_ptr worker_error; + std::vector workers; + workers.reserve(client_events.size()); + + for (auto& entry : client_events) { + const uint32_t client_id = entry.first; + workers.emplace_back([&, client_id, events = std::move(entry.second)] { + try { + WorkloadMetrics local; + size_t begin = 0; + while (begin < events.size()) { + size_t end = begin + 1; + if (events[begin].batch_id() != 0) { + while (end < events.size() && events[end].batch_id() == events[begin].batch_id() && + events[end].operation() == events[begin].operation()) { + ++end; + } + } + + if (!options_.max_throughput) { + const long double scaled = + static_cast(events[end - 1].relative_timestamp_ns()) * + options_.time_scale; + if (scaled > static_cast(std::numeric_limits::max())) { + throw std::overflow_error("scaled workload timestamp is too large"); + } + const auto target = start + std::chrono::nanoseconds(static_cast(scaled)); + std::this_thread::sleep_until(target); + } + + const auto dispatch_time = Clock::now(); + for (size_t i = begin; i < end; ++i) { + AddAttempt(events[i], &local); + if (!options_.max_throughput) { + const auto target = + start + std::chrono::nanoseconds(static_cast( + static_cast(events[i].relative_timestamp_ns()) * + options_.time_scale)); + local.schedule_lag.samples_ns.push_back(Nanoseconds(dispatch_time - target)); + } + } + + std::vector keys; + std::vector sizes; + keys.reserve(end - begin); + sizes.reserve(end - begin); + for (size_t i = begin; i < end; ++i) { + keys.push_back(events[i].key()); + sizes.push_back(static_cast(events[i].value_size())); + } + + std::vector results; + std::vector> values; + const auto operation_start = Clock::now(); + if (events[begin].operation() == Event::PUT) { + values.reserve(end - begin); + for (size_t i = begin; i < end; ++i) { + values.push_back(GenerateDeterministicPayload( + events[i].key(), events[i].operation_id(), payload_seed, sizes[i - begin])); + } + results = client_->BatchPut(client_id, keys, values); + } else { + results = client_->BatchGet(client_id, keys, sizes, &values); + } + const uint64_t latency = Nanoseconds(Clock::now() - operation_start); + + for (size_t i = begin; i < end; ++i) { + const size_t result_index = i - begin; + ClientResult result = + result_index < results.size() ? results[result_index] : ClientResult::kFailed; + bool success = result == ClientResult::kSuccess; + if (events[i].operation() == Event::GET) { + if (result == ClientResult::kNotFound) { + ++local.get_misses; + } else if (success && options_.validate_get_payloads) { + if (result_index >= values.size() || + values[result_index].size() != sizes[result_index] || + !ValidateDeterministicPayload(events[i].key(), events[i].operation_id(), + payload_seed, values[result_index].data(), + values[result_index].size())) { + ++local.get_validation_failures; + success = false; + } + } + } + AddResult(events[i], success, &local); + local.latency.samples_ns.push_back(latency); + } + begin = end; + } + std::lock_guard lock(metrics_mutex); + MergeMetrics(local, &metrics); + } catch (...) { + std::lock_guard lock(metrics_mutex); + if (!worker_error) worker_error = std::current_exception(); + } + }); + } + + for (std::thread& worker : workers) worker.join(); + if (worker_error) std::rethrow_exception(worker_error); + FinalizeDistribution(&metrics.latency); + FinalizeDistribution(&metrics.schedule_lag); + return metrics; +} + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/distributed/benchmark/workload_source.cpp b/src/umbp/distributed/benchmark/workload_source.cpp new file mode 100644 index 000000000..4920c7c13 --- /dev/null +++ b/src/umbp/distributed/benchmark/workload_source.cpp @@ -0,0 +1,175 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include "umbp/distributed/benchmark/workload_source.h" + +#include +#include +#include +#include +#include + +namespace mori::umbp::benchmark { + +TraceWorkloadSource::TraceWorkloadSource(const std::string& path) : reader_(path) {} + +bool TraceWorkloadSource::Next(::umbp::benchmark::WorkloadEvent* event) { + return reader_.ReadNext(event); +} + +SyntheticWorkloadSource::SyntheticWorkloadSource(SyntheticWorkloadConfig config) + : config_(std::move(config)), random_(config_.seed) { + if (config_.key_count == 0) throw std::invalid_argument("key_count must be positive"); + if (config_.client_count == 0) throw std::invalid_argument("client_count must be positive"); + if (config_.batch_size == 0) throw std::invalid_argument("batch_size must be positive"); + if (config_.min_value_size > config_.max_value_size) { + throw std::invalid_argument("min_value_size must not exceed max_value_size"); + } + if (config_.read_ratio < 0.0 || config_.read_ratio > 1.0) { + throw std::invalid_argument("read_ratio must be in [0, 1]"); + } + if (config_.qps < 0.0 || !std::isfinite(config_.qps)) { + throw std::invalid_argument("qps must be finite and non-negative"); + } + if (config_.hotset_fraction <= 0.0 || config_.hotset_fraction > 1.0) { + throw std::invalid_argument("hotset_fraction must be in (0, 1]"); + } + if (config_.zipf_exponent <= 0.0 || !std::isfinite(config_.zipf_exponent)) { + throw std::invalid_argument("zipf_exponent must be finite and positive"); + } + + key_versions_.assign(config_.key_count, 0); + last_put_operation_.assign(config_.key_count, 0); + last_put_size_.assign(config_.key_count, 0); + if (config_.profile == SyntheticProfile::kHotsetZipf) { + const uint64_t hot_keys = std::max( + 1, static_cast( + std::ceil(static_cast(config_.key_count) * config_.hotset_fraction))); + zipf_cdf_.reserve(hot_keys); + double sum = 0.0; + for (uint64_t rank = 1; rank <= hot_keys; ++rank) { + sum += 1.0 / std::pow(static_cast(rank), config_.zipf_exponent); + zipf_cdf_.push_back(sum); + } + for (double& value : zipf_cdf_) value /= sum; + } +} + +uint64_t SyntheticWorkloadSource::SelectKeyIndex(uint64_t sequence) { + switch (config_.profile) { + case SyntheticProfile::kSequential: + case SyntheticProfile::kReadAfterWrite: + case SyntheticProfile::kMixed: + return (config_.profile == SyntheticProfile::kReadAfterWrite ? sequence / 2 : sequence) % + config_.key_count; + case SyntheticProfile::kCapacityPressure: + // Pressure uses a non-repeating key stream so successful PUTs consume + // fresh capacity instead of becoming backend dedup hits. + return sequence; + case SyntheticProfile::kHotsetZipf: { + const double sample = std::generate_canonical(random_); + return static_cast(std::lower_bound(zipf_cdf_.begin(), zipf_cdf_.end(), sample) - + zipf_cdf_.begin()); + } + case SyntheticProfile::kUniform: + return std::uniform_int_distribution(0, config_.key_count - 1)(random_); + } + throw std::logic_error("unknown synthetic workload profile"); +} + +uint64_t SyntheticWorkloadSource::SelectValueSize() { + if (config_.value_size_distribution == ValueSizeDistribution::kFixed || + config_.min_value_size == config_.max_value_size) { + return config_.min_value_size; + } + if (config_.value_size_distribution == ValueSizeDistribution::kUniform) { + return std::uniform_int_distribution(config_.min_value_size, + config_.max_value_size)(random_); + } + const double low = std::log(static_cast(std::max(1, config_.min_value_size))); + const double high = std::log(static_cast(std::max(1, config_.max_value_size))); + const double sample = std::uniform_real_distribution(low, high)(random_); + return std::clamp(static_cast(std::llround(std::exp(sample))), + config_.min_value_size, config_.max_value_size); +} + +bool SyntheticWorkloadSource::SelectRead(uint64_t sequence) { + if (config_.profile == SyntheticProfile::kSequential || + config_.profile == SyntheticProfile::kCapacityPressure) { + return false; + } + if (config_.profile == SyntheticProfile::kReadAfterWrite) return sequence % 2 == 1; + return std::bernoulli_distribution(config_.read_ratio)(random_); +} + +bool SyntheticWorkloadSource::Next(::umbp::benchmark::WorkloadEvent* event) { + if (event == nullptr) throw std::invalid_argument("event must not be null"); + if (generated_ >= config_.operation_count) return false; + + const uint64_t sequence = generated_; + const uint64_t key_index = SelectKeyIndex(sequence); + bool is_read = SelectRead(sequence); + const bool tracks_working_set = config_.profile != SyntheticProfile::kCapacityPressure; + if (tracks_working_set && is_read && last_put_operation_[key_index] == 0) is_read = false; + + event->Clear(); + if (config_.qps > 0.0) { + const uint64_t events_per_batch = + static_cast(config_.batch_size) * config_.client_count; + const uint64_t batch_sequence = (sequence / events_per_batch) * events_per_batch; + const long double timestamp = + static_cast(batch_sequence) * 1000000000.0L / config_.qps; + if (timestamp > static_cast(std::numeric_limits::max())) { + throw std::overflow_error("synthetic timestamp exceeds uint64_t"); + } + event->set_relative_timestamp_ns(static_cast(timestamp)); + } + // Keep every key's dependency chain on one logical client. The runner may + // execute different clients concurrently, so assigning by event sequence + // could otherwise let a GET overtake the PUT that established its payload. + event->set_client_id(static_cast(key_index % config_.client_count)); + if (!is_read && tracks_working_set) ++key_versions_[key_index]; + event->set_key(config_.key_prefix + std::to_string(key_index) + "-v" + + std::to_string(tracks_working_set ? key_versions_[key_index] : 1)); + event->set_batch_id( + sequence / (static_cast(config_.batch_size) * config_.client_count) + 1); + + const uint64_t new_operation_id = sequence + 1; + if (is_read) { + event->set_operation(::umbp::benchmark::WorkloadEvent::GET); + event->set_operation_id(last_put_operation_[key_index]); + event->set_value_size(last_put_size_[key_index]); + } else { + event->set_operation(::umbp::benchmark::WorkloadEvent::PUT); + event->set_operation_id(new_operation_id); + event->set_value_size(SelectValueSize()); + if (tracks_working_set) { + last_put_operation_[key_index] = new_operation_id; + last_put_size_[key_index] = event->value_size(); + } + } + ++generated_; + return true; +} + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/distributed/benchmark/workload_trace.cpp b/src/umbp/distributed/benchmark/workload_trace.cpp new file mode 100644 index 000000000..7b5a3759c --- /dev/null +++ b/src/umbp/distributed/benchmark/workload_trace.cpp @@ -0,0 +1,187 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include "umbp/distributed/benchmark/workload_trace.h" + +#include +#include +#include +#include + +namespace mori::umbp::benchmark { +namespace { + +constexpr std::array kMagic = {'U', 'M', 'B', 'P', 'T', 'R', 'C', 'E'}; +constexpr uint32_t kEnvelopeVersion = 1; + +// Ceilings that stop a truncated or hostile trace from being read into +// unbounded memory. Fixed on purpose: a per-instance limit was threaded through +// every reader and writer and no caller ever set anything but the defaults. +constexpr size_t kMaxHeaderBytes = 1 << 20; +constexpr size_t kMaxEventBytes = 2 << 20; +constexpr size_t kMaxKeyBytes = 1 << 20; +constexpr uint64_t kMaxValueBytes = uint64_t{1} << 40; + +void WriteU32(std::ostream* stream, uint32_t value) { + std::array bytes{}; + for (size_t i = 0; i < bytes.size(); ++i) { + bytes[i] = static_cast((value >> (i * 8)) & 0xff); + } + stream->write(bytes.data(), bytes.size()); +} + +uint32_t DecodeU32(const std::array& bytes) { + uint32_t value = 0; + for (size_t i = 0; i < bytes.size(); ++i) { + value |= static_cast(static_cast(bytes[i])) << (i * 8); + } + return value; +} + +void ValidateHeader(const ::umbp::benchmark::WorkloadTraceHeader& header) { + if (header.schema_version() != kWorkloadTraceSchemaVersion) { + throw TraceError("unsupported workload trace schema version: " + + std::to_string(header.schema_version())); + } + if (header.time_unit() != ::umbp::benchmark::WorkloadTraceHeader::TIME_UNIT_NANOSECONDS) { + throw TraceError("unsupported workload trace time unit"); + } +} + +void ValidateEvent(const ::umbp::benchmark::WorkloadEvent& event) { + if (event.operation() != ::umbp::benchmark::WorkloadEvent::PUT && + event.operation() != ::umbp::benchmark::WorkloadEvent::GET) { + throw TraceError("workload event has an invalid operation"); + } + if (event.key().empty()) throw TraceError("workload event has an empty key"); + if (event.key().size() > kMaxKeyBytes) { + throw TraceError("workload event key exceeds configured limit"); + } + if (event.value_size() > kMaxValueBytes) { + throw TraceError("workload event value exceeds configured limit"); + } +} + +} // namespace + +TraceWriter::TraceWriter(const std::string& path, + const ::umbp::benchmark::WorkloadTraceHeader& header) + : stream_(path, std::ios::binary | std::ios::trunc) { + if (!stream_) throw TraceError("failed to open workload trace for writing: " + path); + ValidateHeader(header); + stream_.write(kMagic.data(), kMagic.size()); + WriteU32(&stream_, kEnvelopeVersion); + if (!stream_) throw TraceError("failed to write workload trace envelope"); + WriteRecord(header, kMaxHeaderBytes); +} + +TraceWriter::~TraceWriter() { + if (!closed_) { + stream_.flush(); + stream_.close(); + } +} + +void TraceWriter::WriteRecord(const google::protobuf::MessageLite& record, size_t max_bytes) { + const size_t size = record.ByteSizeLong(); + if (size > max_bytes || size > std::numeric_limits::max()) { + throw TraceError("workload trace record exceeds configured limit"); + } + std::string bytes; + if (!record.SerializeToString(&bytes)) { + throw TraceError("failed to serialize workload trace record"); + } + WriteU32(&stream_, static_cast(bytes.size())); + stream_.write(bytes.data(), static_cast(bytes.size())); + if (!stream_) throw TraceError("failed to write workload trace record"); +} + +void TraceWriter::Write(const ::umbp::benchmark::WorkloadEvent& event) { + if (closed_) throw TraceError("cannot write to a closed workload trace"); + ValidateEvent(event); + WriteRecord(event, kMaxEventBytes); +} + +void TraceWriter::Close() { + if (closed_) return; + stream_.flush(); + if (!stream_) throw TraceError("failed to flush workload trace"); + stream_.close(); + if (stream_.fail()) throw TraceError("failed to close workload trace"); + closed_ = true; +} + +TraceReader::TraceReader(const std::string& path) : stream_(path, std::ios::binary) { + if (!stream_) throw TraceError("failed to open workload trace for reading: " + path); + + std::array magic{}; + stream_.read(magic.data(), magic.size()); + if (stream_.gcount() != static_cast(magic.size()) || magic != kMagic) { + throw TraceError("invalid or truncated workload trace magic"); + } + std::array version{}; + stream_.read(version.data(), version.size()); + if (stream_.gcount() != static_cast(version.size())) { + throw TraceError("truncated workload trace envelope"); + } + const uint32_t envelope_version = DecodeU32(version); + if (envelope_version != kEnvelopeVersion) { + throw TraceError("unsupported workload trace envelope version: " + + std::to_string(envelope_version)); + } + ReadRecord(&header_, kMaxHeaderBytes, false); + ValidateHeader(header_); +} + +bool TraceReader::ReadRecord(google::protobuf::MessageLite* record, size_t max_bytes, + bool allow_eof) { + std::array length{}; + stream_.read(length.data(), length.size()); + const std::streamsize length_read = stream_.gcount(); + if (length_read == 0 && allow_eof && stream_.eof()) return false; + if (length_read != static_cast(length.size())) { + throw TraceError("truncated workload trace record length"); + } + const uint32_t size = DecodeU32(length); + if (size > max_bytes) throw TraceError("workload trace record exceeds configured limit"); + + std::string bytes(size, '\0'); + stream_.read(bytes.data(), static_cast(size)); + if (stream_.gcount() != static_cast(size)) { + throw TraceError("truncated workload trace record"); + } + if (!record->ParseFromArray(bytes.data(), static_cast(bytes.size()))) { + throw TraceError("malformed workload trace protobuf record"); + } + return true; +} + +bool TraceReader::ReadNext(::umbp::benchmark::WorkloadEvent* event) { + if (event == nullptr) throw std::invalid_argument("event must not be null"); + event->Clear(); + if (!ReadRecord(event, kMaxEventBytes, true)) return false; + ValidateEvent(*event); + return true; +} + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/distributed/benchmark/workload_trace_recorder.cpp b/src/umbp/distributed/benchmark/workload_trace_recorder.cpp new file mode 100644 index 000000000..e6e559fee --- /dev/null +++ b/src/umbp/distributed/benchmark/workload_trace_recorder.cpp @@ -0,0 +1,143 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include "umbp/distributed/benchmark/workload_trace_recorder.h" + +#include +#include + +namespace mori::umbp::benchmark { +namespace { + +::umbp::benchmark::WorkloadTraceHeader MakeHeader(const WorkloadTraceRecorderOptions& options) { + ::umbp::benchmark::WorkloadTraceHeader header; + header.set_schema_version(kWorkloadTraceSchemaVersion); + header.set_time_unit(::umbp::benchmark::WorkloadTraceHeader::TIME_UNIT_NANOSECONDS); + header.set_seed(options.seed); + auto& metadata = *header.mutable_metadata(); + metadata["source"] = "production"; + metadata["payload_mode"] = "external"; + metadata["node_id"] = options.node_id; + metadata["backend_policy"] = options.backend_policy; + return header; +} + +void ValidateBatch(const std::vector& keys, const std::vector& sizes, + const std::vector& outcomes) { + if (keys.size() != sizes.size() || keys.size() != outcomes.size()) { + throw std::invalid_argument( + "workload trace batch keys, sizes, and outcomes must have equal lengths"); + } +} + +::umbp::benchmark::WorkloadEvent::Outcome ToProtoOutcome(WorkloadTraceOutcome outcome) { + switch (outcome) { + case WorkloadTraceOutcome::kSuccess: + return ::umbp::benchmark::WorkloadEvent::OUTCOME_SUCCESS; + case WorkloadTraceOutcome::kMiss: + return ::umbp::benchmark::WorkloadEvent::OUTCOME_MISS; + case WorkloadTraceOutcome::kFailure: + return ::umbp::benchmark::WorkloadEvent::OUTCOME_FAILURE; + } + return ::umbp::benchmark::WorkloadEvent::OUTCOME_UNSPECIFIED; +} + +} // namespace + +WorkloadTraceRecorder::WorkloadTraceRecorder(Options options) + : client_id_(options.client_id), + writer_(options.path, MakeHeader(options)), + start_(Clock::now()) {} + +WorkloadTraceRecorder::~WorkloadTraceRecorder() { + try { + Close(); + } catch (...) { + } +} + +uint64_t WorkloadTraceRecorder::RelativeTimestampNs() const { + return static_cast( + std::chrono::duration_cast(Clock::now() - start_).count()); +} + +bool WorkloadTraceRecorder::WriteLocked(const ::umbp::benchmark::WorkloadEvent& event) { + try { + writer_.Write(event); + } catch (const std::exception& exception) { + failed_ = true; + error_ = exception.what(); + return false; + } catch (...) { + failed_ = true; + error_ = "unknown workload trace write error"; + return false; + } + ++event_count_; + return true; +} + +void WorkloadTraceRecorder::RecordBatch(::umbp::benchmark::WorkloadEvent::Operation operation, + const std::vector& keys, + const std::vector& sizes, + const std::vector& outcomes) { + ValidateBatch(keys, sizes, outcomes); + std::lock_guard lock(mutex_); + if (failed_) return; + const uint64_t batch_id = ++next_batch_id_; + const uint64_t timestamp_ns = RelativeTimestampNs(); + + for (size_t i = 0; i < keys.size(); ++i) { + ::umbp::benchmark::WorkloadEvent event; + event.set_relative_timestamp_ns(timestamp_ns); + event.set_client_id(client_id_); + event.set_operation_id(++next_operation_id_); + event.set_operation(operation); + event.set_key(keys[i]); + event.set_value_size(static_cast(sizes[i])); + event.set_batch_id(batch_id); + event.set_outcome(ToProtoOutcome(outcomes[i])); + if (!WriteLocked(event)) return; + } +} + +void WorkloadTraceRecorder::Close() { + std::lock_guard lock(mutex_); + try { + writer_.Close(); + } catch (const std::exception& exception) { + failed_ = true; + if (error_.empty()) error_ = exception.what(); + } catch (...) { + failed_ = true; + if (error_.empty()) error_ = "unknown workload trace close error"; + } + if (failed_) throw TraceError(error_); +} + +uint64_t WorkloadTraceRecorder::event_count() const { + std::lock_guard lock(mutex_); + return event_count_; +} + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/distributed/bin/tier_benchmark_main.cpp b/src/umbp/distributed/bin/tier_benchmark_main.cpp new file mode 100644 index 000000000..9d3471406 --- /dev/null +++ b/src/umbp/distributed/bin/tier_benchmark_main.cpp @@ -0,0 +1,706 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/pool_workload_client.h" +#include "umbp/distributed/benchmark/workload_runner.h" +#include "umbp/distributed/benchmark/workload_source.h" +#include "umbp/distributed/benchmark/workload_trace.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/pool/policy_config.h" +#include "umbp/distributed/pool_client.h" +#include "umbp/distributed/routing/route_get_strategy.h" +#include "umbp/distributed/routing/route_put_strategy.h" +namespace { +namespace bench = mori::umbp::benchmark; +using mori::umbp::BackendInstanceConfig; +using mori::umbp::ConfigurableRoutePutStrategy; +using mori::umbp::LocalPreferringRouteGetStrategy; +using mori::umbp::MasterServer; +using mori::umbp::MasterServerConfig; +using mori::umbp::MediumBackend; +using mori::umbp::PoolClient; +using mori::umbp::PoolClientConfig; +using mori::umbp::PoolPlacementPolicy; +using mori::umbp::RandomRouteGetStrategy; +using mori::umbp::TierType; +constexpr uint64_t kDefaultCapacity = 256ULL << 20; +constexpr uint64_t kDefaultPageSize = 2ULL << 20; +constexpr auto kServerStartTimeout = std::chrono::seconds(10); +struct Options { + std::string command; + std::string trace_path; + std::string policy_path; + // Parsed once by LoadPolicy; the clients lower it per node rather than + // reading and reparsing the file. + std::optional policy; + // Largest object that can fit one SSD backend's contiguous staging arena. + uint64_t max_ssd_staging_object_bytes = std::numeric_limits::max(); + bench::SyntheticWorkloadConfig workload; + bench::WorkloadRunnerOptions runner; + std::string profile = "mixed"; + std::string size_distribution = "fixed"; + std::string put_strategy = "most_available"; + std::string affinity = "none"; + std::string get_strategy = "local"; + uint64_t backend_capacity = kDefaultCapacity; + uint64_t page_size = kDefaultPageSize; + uint64_t settle_ms = 500; + bool clients_explicit = false; + bool scheduling_explicit = false; + // Set for a trace recorded from production: the workload is whatever the + // trace contains, so the synthetic generator's knobs describe nothing. + bool trace_defines_workload = false; +}; +[[noreturn]] void UsageError(const std::string& message) { + throw std::invalid_argument(message + " (use --help for usage)"); +} +void PrintHelp() { + std::cout << R"(UMBP distributed tier-management benchmark +Usage: + umbp_tier_bench generate --trace PATH [OPTIONS] + umbp_tier_bench run [OPTIONS] + umbp_tier_bench replay --trace PATH [OPTIONS] +Workload: + --profile NAME --seed N --operations N --keys N + --min-value-bytes SIZE --max-value-bytes SIZE + --size-distribution fixed|uniform|log-uniform + --read-ratio R --clients N --batch N --qps Q +Cluster: + --config PATH topology: named backends, logical tiers, weights, watermarks. + Without it every peer gets one DRAM backend sized by + --backend-capacity. + --backend-capacity SIZE --page-size SIZE + --put-strategy most_available|random --affinity none|same|local + --get-strategy local|random +Execution: + --max-throughput | --open-loop --time-scale R + --payload-validation BOOL --no-payload-validation --settle-ms N +SIZE may use B, KiB, MiB, or GiB. Options accept --name value or --name=value. +)"; +} +template +UInt ParseUnsigned(std::string_view text, const char* option) { + static_assert(std::is_unsigned_v); + uint64_t value = 0; + const auto result = std::from_chars(text.data(), text.data() + text.size(), value); + if (text.empty() || text.front() == '-' || result.ec != std::errc{} || + result.ptr != text.data() + text.size() || + value > static_cast(std::numeric_limits::max())) { + UsageError(std::string(option) + " expects an unsigned integer"); + } + return static_cast(value); +} +double ParseDouble(std::string_view text, const char* option) { + const std::string owned(text); + char* end = nullptr; + errno = 0; + const double value = std::strtod(owned.c_str(), &end); + if (errno == ERANGE || end == owned.c_str() || *end != '\0' || !std::isfinite(value)) { + UsageError(std::string(option) + " expects a finite number"); + } + return value; +} +bool ParseBool(std::string_view text, const char* option) { + if (text == "true" || text == "1") return true; + if (text == "false" || text == "0") return false; + UsageError(std::string(option) + " expects true or false"); +} +uint64_t ParseSize(std::string_view text, const char* option) { + uint64_t scale = 1; + for (const auto& suffix : {std::pair{"GiB", 1ULL << 30}, + {"MiB", 1ULL << 20}, + {"KiB", 1ULL << 10}, + {"B", 1}}) { + if (text.size() >= suffix.first.size() && + text.substr(text.size() - suffix.first.size()) == suffix.first) { + text.remove_suffix(suffix.first.size()); + scale = suffix.second; + break; + } + } + const uint64_t value = ParseUnsigned(text, option); + if (value > std::numeric_limits::max() / scale) { + UsageError(std::string(option) + " is too large"); + } + return value * scale; +} +Options ParseOptions(int argc, char** argv) { + if (argc < 2) UsageError("missing subcommand"); + Options options; + options.command = argv[1]; + options.runner.max_throughput = true; + if (options.command == "help" || options.command == "--help" || options.command == "-h") { + PrintHelp(); + std::exit(0); + } + if (options.command != "run" && options.command != "generate" && options.command != "replay") { + UsageError("unknown subcommand: " + options.command); + } + for (int i = 2; i < argc; ++i) { + std::string name = argv[i]; + std::string inline_value; + bool has_inline_value = false; + if (const size_t equals = name.find('='); equals != std::string::npos) { + inline_value = name.substr(equals + 1); + name.resize(equals); + has_inline_value = true; + } + if (name == "--help" || name == "-h") { + PrintHelp(); + std::exit(0); + } + auto value = [&]() { + if (has_inline_value) { + if (inline_value.empty()) UsageError(name + " requires a value"); + return inline_value; + } + if (i + 1 >= argc) UsageError(name + " requires a value"); + return std::string(argv[++i]); + }; + if (name == "--max-throughput") { + if (has_inline_value) UsageError(name + " does not take a value"); + options.runner.max_throughput = true; + options.scheduling_explicit = true; + } else if (name == "--open-loop") { + if (has_inline_value) UsageError(name + " does not take a value"); + options.runner.max_throughput = false; + options.scheduling_explicit = true; + } else if (name == "--no-payload-validation") { + if (has_inline_value) UsageError(name + " does not take a value"); + options.runner.validate_get_payloads = false; + } else if (name == "--trace") { + options.trace_path = value(); + } else if (name == "--config") { + options.policy_path = value(); + } else if (name == "--profile") { + options.profile = value(); + } else if (name == "--seed") { + options.workload.seed = ParseUnsigned(value(), "--seed"); + } else if (name == "--operations") { + options.workload.operation_count = ParseUnsigned(value(), "--operations"); + } else if (name == "--keys") { + options.workload.key_count = ParseUnsigned(value(), "--keys"); + } else if (name == "--min-value-bytes") { + options.workload.min_value_size = ParseSize(value(), "--min-value-bytes"); + } else if (name == "--max-value-bytes") { + options.workload.max_value_size = ParseSize(value(), "--max-value-bytes"); + } else if (name == "--size-distribution") { + options.size_distribution = value(); + } else if (name == "--read-ratio") { + options.workload.read_ratio = ParseDouble(value(), "--read-ratio"); + } else if (name == "--clients") { + options.workload.client_count = ParseUnsigned(value(), "--clients"); + options.clients_explicit = true; + } else if (name == "--batch") { + options.workload.batch_size = ParseUnsigned(value(), "--batch"); + } else if (name == "--qps") { + options.workload.qps = ParseDouble(value(), "--qps"); + } else if (name == "--backend-capacity") { + options.backend_capacity = ParseSize(value(), "--backend-capacity"); + } else if (name == "--page-size") { + options.page_size = ParseSize(value(), "--page-size"); + } else if (name == "--put-strategy") { + options.put_strategy = value(); + } else if (name == "--affinity") { + options.affinity = value(); + } else if (name == "--get-strategy") { + options.get_strategy = value(); + } else if (name == "--time-scale") { + options.runner.time_scale = ParseDouble(value(), "--time-scale"); + } else if (name == "--payload-validation") { + options.runner.validate_get_payloads = ParseBool(value(), "--payload-validation"); + } else if (name == "--settle-ms") { + options.settle_ms = ParseUnsigned(value(), "--settle-ms"); + } else { + UsageError("unknown option: " + name); + } + } + if (options.command == "replay" && !options.scheduling_explicit) { + options.runner.max_throughput = false; + } + return options; +} +bench::SyntheticProfile ParseProfile(const std::string& name) { + if (name == "sequential") return bench::SyntheticProfile::kSequential; + if (name == "uniform") return bench::SyntheticProfile::kUniform; + if (name == "hotset") return bench::SyntheticProfile::kHotsetZipf; + if (name == "read-after-write") return bench::SyntheticProfile::kReadAfterWrite; + if (name == "mixed") return bench::SyntheticProfile::kMixed; + if (name == "capacity-pressure") return bench::SyntheticProfile::kCapacityPressure; + UsageError("invalid --profile: " + name); +} +void ValidateOptions(Options* options) { + if (options->page_size == 0) UsageError("--page-size must be positive"); + options->workload.profile = ParseProfile(options->profile); + if (options->size_distribution == "fixed") { + options->workload.value_size_distribution = bench::ValueSizeDistribution::kFixed; + } else if (options->size_distribution == "uniform") { + options->workload.value_size_distribution = bench::ValueSizeDistribution::kUniform; + } else if (options->size_distribution == "log-uniform") { + options->workload.value_size_distribution = bench::ValueSizeDistribution::kLogUniform; + } else { + UsageError("--size-distribution must be fixed, uniform, or log-uniform"); + } + if ((options->command == "generate" || options->command == "replay") && + options->trace_path.empty()) { + UsageError(options->command + " requires --trace"); + } + if (options->workload.client_count == 0 || options->workload.key_count == 0 || + options->workload.batch_size == 0 || options->workload.min_value_size == 0) { + UsageError("clients, keys, batch, and value sizes must be positive"); + } + if (options->workload.min_value_size > options->workload.max_value_size) { + UsageError("--min-value-bytes must not exceed --max-value-bytes"); + } + if (options->workload.read_ratio < 0 || options->workload.read_ratio > 1 || + options->workload.qps < 0 || options->runner.time_scale <= 0) { + UsageError("read ratio, QPS, or time scale is out of range"); + } + if (options->put_strategy != "most_available" && options->put_strategy != "random") { + UsageError("--put-strategy must be most_available or random"); + } + if (options->affinity != "none" && options->affinity != "same" && options->affinity != "local") { + UsageError("--affinity must be none, same, or local"); + } + if (options->get_strategy != "local" && options->get_strategy != "random") { + UsageError("--get-strategy must be local or random"); + } + if (options->policy_path.empty() && options->backend_capacity < options->page_size) { + UsageError("--backend-capacity must be at least --page-size"); + } + // For a replay the authoritative sizes are in the events, and the trace scan + // checks them there. + if (options->command != "replay" && + options->workload.max_value_size > options->max_ssd_staging_object_bytes) { + UsageError("SSD values must not exceed the contiguous staging arena"); + } +} +// Reads --config once. The clients lower it again per node for their own +// instance names; this pass exists to reject a bad topology before anything +// starts, and to tell the rest of the run what media it will be writing to. +void LoadPolicy(Options* options) { + if (options->policy_path.empty()) return; + if (options->page_size == 0) UsageError("--page-size must be positive"); + auto loaded = mori::umbp::LoadBackendPolicyFile(options->policy_path); + if (!loaded.ok()) UsageError("invalid --config: " + loaded.error); + options->policy = std::move(*loaded.config); + + PoolClientConfig lowered; + lowered.dram_page_size = options->page_size; + std::string error; + if (!mori::umbp::ApplyBackendPolicy(*options->policy, &lowered, &error)) { + UsageError("invalid --config: " + error); + } + if (lowered.backends.size() > mori::umbp::kMaxBackendsPerPeer) { + UsageError("--config expands beyond the per-peer backend limit"); + } + for (const auto& backend : lowered.backends) { + if (backend.tier != TierType::SSD) continue; + const uint64_t pages = backend.ssd_staging_buffer_slots > 0 + ? static_cast(backend.ssd_staging_buffer_slots) + : 16; + if (pages > std::numeric_limits::max() / options->page_size) { + UsageError("SSD staging arena size overflows uint64"); + } + const uint64_t ssd_capacity = backend.ssd.ssd.capacity_bytes; + options->max_ssd_staging_object_bytes = + std::min({options->max_ssd_staging_object_bytes, pages * options->page_size, ssd_capacity}); + } +} +void AddTraceMetadata(const Options& options, ::umbp::benchmark::WorkloadTraceHeader* header) { + auto& metadata = *header->mutable_metadata(); + metadata["profile"] = options.profile; + metadata["seed"] = std::to_string(options.workload.seed); + metadata["operations"] = std::to_string(options.workload.operation_count); + metadata["keys"] = std::to_string(options.workload.key_count); + metadata["min_value_bytes"] = std::to_string(options.workload.min_value_size); + metadata["max_value_bytes"] = std::to_string(options.workload.max_value_size); + metadata["size_distribution"] = options.size_distribution; + metadata["read_ratio"] = std::to_string(options.workload.read_ratio); + metadata["clients"] = std::to_string(options.workload.client_count); + metadata["batch"] = std::to_string(options.workload.batch_size); + metadata["qps"] = std::to_string(options.workload.qps); + metadata["hotset_fraction"] = std::to_string(options.workload.hotset_fraction); + metadata["zipf_exponent"] = std::to_string(options.workload.zipf_exponent); + metadata["key_prefix"] = options.workload.key_prefix; +} +int GenerateTrace(const Options& options) { + ::umbp::benchmark::WorkloadTraceHeader header; + header.set_schema_version(bench::kWorkloadTraceSchemaVersion); + header.set_time_unit(::umbp::benchmark::WorkloadTraceHeader::TIME_UNIT_NANOSECONDS); + header.set_seed(options.workload.seed); + AddTraceMetadata(options, &header); + bench::SyntheticWorkloadSource source(options.workload); + bench::TraceWriter writer(options.trace_path, header); + ::umbp::benchmark::WorkloadEvent event; + uint64_t events = 0; + while (source.Next(&event)) { + writer.Write(event); + ++events; + } + writer.Close(); + std::cout << "trace_path,events,seed,schema_version\n\"" << options.trace_path << "\"," << events + << ',' << options.workload.seed << ',' << bench::kWorkloadTraceSchemaVersion << '\n'; + return 0; +} +class MasterHarness { + public: + explicit MasterHarness(const Options& options) { + MasterServerConfig config; + config.listen_address = "127.0.0.1:0"; + config.registry_config.default_dram_page_size = options.page_size; + const auto algorithm = options.put_strategy == "random" + ? ConfigurableRoutePutStrategy::SelectAlgo::kRandom + : ConfigurableRoutePutStrategy::SelectAlgo::kMostAvailable; + auto affinity = ConfigurableRoutePutStrategy::NodeAffinity::kNone; + if (options.affinity == "same") { + affinity = ConfigurableRoutePutStrategy::NodeAffinity::kSame; + } else if (options.affinity == "local") { + affinity = ConfigurableRoutePutStrategy::NodeAffinity::kLocal; + } + config.put_strategy = + std::make_unique(algorithm, affinity, options.workload.seed); + if (options.get_strategy == "random") { + config.get_strategy = std::make_unique(); + } else { + config.get_strategy = std::make_unique(); + } + server_ = std::make_unique(std::move(config)); + thread_ = std::thread([this] { + try { + server_->Run(); + } catch (...) { + error_ = std::current_exception(); + failed_.store(true, std::memory_order_release); + } + }); + const auto deadline = std::chrono::steady_clock::now() + kServerStartTimeout; + while (server_->GetBoundPort() == 0 && std::chrono::steady_clock::now() < deadline && + !failed_.load(std::memory_order_acquire)) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + if (server_->GetBoundPort() == 0) { + Stop(); + if (error_) std::rethrow_exception(error_); + throw std::runtime_error("timed out starting in-process master"); + } + } + + ~MasterHarness() { Stop(); } + std::string Address() const { return "127.0.0.1:" + std::to_string(server_->GetBoundPort()); } + + private: + void Stop() { + if (server_) server_->Shutdown(); + if (thread_.joinable()) thread_.join(); + } + std::unique_ptr server_; + std::thread thread_; + std::exception_ptr error_; + std::atomic failed_{false}; +}; +// Topology for a run without --config: the smallest thing that can serve the +// workload. Anything richer — several backends, other media, logical tiers, +// weights, watermarks — is what the policy file is for. +BackendInstanceConfig DefaultBackend(uint64_t capacity) { + BackendInstanceConfig backend; + backend.name = "dram-0"; + backend.tier = TierType::DRAM; + backend.dram.buffer_sizes = {capacity}; + return backend; +} +std::vector> StartClients(const Options& options, + const std::string& master_address) { + std::vector> clients; + const auto& policy = options.policy; + for (uint32_t node = 0; node < options.workload.client_count; ++node) { + PoolClientConfig config; + config.master_config.master_address = master_address; + config.master_config.node_id = "tier-benchmark-node-" + std::to_string(node); + config.master_config.node_address = "127.0.0.1"; + config.master_config.auto_heartbeat = true; + config.auto_peer_service_port = true; + if (options.workload.client_count > 1) { + config.io_engine.host = "127.0.0.1"; + config.io_engine.port = 0; + } + config.dram_page_size = options.page_size; + config.cache_remote_fetches = false; + if (policy.has_value()) { + std::string error; + if (!mori::umbp::ApplyBackendPolicy(*policy, &config, &error, + "-tier-benchmark-node-" + std::to_string(node))) { + throw std::runtime_error("failed to apply backend policy: " + error); + } + } else { + config.placement_policy = PoolPlacementPolicy::SINGLE_BACKEND; + config.backends.push_back(DefaultBackend(options.backend_capacity)); + } + auto client = std::make_unique(std::move(config)); + if (!client->Init()) { + throw std::runtime_error("failed to initialize PoolClient " + std::to_string(node)); + } + clients.push_back(std::move(client)); + } + return clients; +} +void Settle(const std::vector>& clients, uint64_t ms) { + for (const auto& client : clients) client->Master().FlushHeartbeat(); + if (ms != 0) std::this_thread::sleep_for(std::chrono::milliseconds(ms)); +} +std::string Csv(std::string_view value) { + std::string result = "\""; + for (char c : value) { + result += c; + if (c == '"') result += c; + } + return result + '"'; +} +void PrintSummary(const Options& options, const bench::WorkloadMetrics& metrics, uint64_t seed, + uint64_t wall_ns) { + const double seconds = static_cast(wall_ns) / 1e9; + const double ops_per_second = seconds == 0 ? 0 : metrics.total.succeeded / seconds; + const double gib_per_second = + seconds == 0 ? 0 : metrics.total.succeeded_bytes / static_cast(1ULL << 30) / seconds; + // When a recorded trace supplies the workload, the generator's knobs kept + // their defaults and describe nothing; reporting them invents a shape the run + // never had. The topology is always reported from the pool itself below. + const auto generated = [&](std::string_view value) { + return options.trace_defines_workload ? Csv("") : Csv(value); + }; + const auto generated_number = [&](auto value) { + if (options.trace_defines_workload) return std::string{}; + std::ostringstream text; + text << value; + return text.str(); + }; + std::cout << "config\n" + << "command,trace,policy,seed,profile,operations,keys,min_value_bytes,max_value_bytes," + "size_distribution,read_ratio,clients,batch,qps,backend_capacity,page_size," + "put_strategy,affinity,get_strategy,scheduling,time_scale,payload_validation," + "settle_ms\n" + << Csv(options.command) << ',' << Csv(options.trace_path) << ',' + << Csv(options.policy_path) << ',' << seed << ',' << generated(options.profile) << ',' + << options.workload.operation_count << ',' + << generated_number(options.workload.key_count) << ',' + << generated_number(options.workload.min_value_size) << ',' + << generated_number(options.workload.max_value_size) << ',' + << generated(options.size_distribution) << ',' + << generated_number(options.workload.read_ratio) << ',' << options.workload.client_count + << ',' << generated_number(options.workload.batch_size) << ',' + << generated_number(options.workload.qps) << ',' + << (options.policy_path.empty() ? std::to_string(options.backend_capacity) + : std::string{}) + << ',' << options.page_size << ',' << Csv(options.put_strategy) << ',' + << Csv(options.affinity) << ',' << Csv(options.get_strategy) << ',' + << Csv(options.runner.max_throughput ? "max-throughput" : "open-loop") << ',' + << options.runner.time_scale << ',' + << (options.runner.validate_get_payloads ? "true" : "false") << ',' << options.settle_ms + << '\n' + << "summary\n" + << "seed,attempted,succeeded,failed,succeeded_bytes,puts,puts_failed,gets,gets_failed," + "get_misses,validation_failures,recorded_misses,recorded_failures," + "wall_time_ns,ops_per_s,GiB_per_s,latency_p50_ns," + "latency_p95_ns,latency_p99_ns,latency_max_ns,schedule_lag_p99_ns\n" + << seed << ',' << metrics.total.attempted << ',' << metrics.total.succeeded << ',' + << metrics.total.failed << ',' << metrics.total.succeeded_bytes << ',' + << metrics.puts.succeeded << ',' << metrics.puts.failed << ',' << metrics.gets.succeeded + << ',' << metrics.gets.failed << ',' << metrics.get_misses << ',' + << metrics.get_validation_failures << ',' << metrics.recorded_misses << ',' + << metrics.recorded_failures << ',' << wall_ns << ',' << std::setprecision(10) + << ops_per_second << ',' << gib_per_second << ',' << metrics.latency.p50_ns << ',' + << metrics.latency.p95_ns << ',' << metrics.latency.p99_ns << ',' + << metrics.latency.max_ns << ',' << metrics.schedule_lag.p99_ns << '\n'; +} +void PrintBackendPlacement(const std::vector>& clients) { + // logical_tier is the topology as the pool actually built it, which is the + // only description that stays true when a policy file expands one named + // backend into several instances. + std::cout << "backend_placement\n" + << "node,backend,logical_tier,tier,owned_keys,total_bytes,available_bytes," + "max_allocatable_bytes\n"; + for (const auto& client : clients) { + auto& registry = client->Backends(); + for (MediumBackend* backend : registry.All()) { + const auto* entry = registry.GetEntry(backend); + const auto capacity = backend->Capacity(); + std::cout << Csv(client->NodeId()) << ',' + << Csv(entry == nullptr ? backend->Name() : entry->name) << ',' + << Csv(client->LogicalTierForBackend(registry.BackendId(backend))) << ',' + << Csv(mori::umbp::TierTypeName(backend->Tier())) << ',' << backend->OwnedKeyCount() + << ',' << capacity.total_bytes << ',' << capacity.available_bytes << ',' + << capacity.max_allocatable_bytes << '\n'; + } + } + std::cout << "tier_transitions\n" + << "node,attempted,succeeded,failed,offloaded_bytes,promoted_bytes\n"; + for (const auto& client : clients) { + const auto metrics = client->TransitionMetrics(); + std::cout << Csv(client->NodeId()) << ',' << metrics.attempted << ',' << metrics.succeeded + << ',' << metrics.failed << ',' << metrics.offloaded_bytes << ',' + << metrics.promoted_bytes << '\n'; + } + // Which tier answered the reads: a hierarchy whose hot tier serves few of them + // is misconfigured even when the throughput number looks acceptable. + std::cout << "tier_reads\n" << "node,logical_tier,reads_served\n"; + for (const auto& client : clients) { + for (const auto& [tier, reads] : client->TierReadHits()) { + std::cout << Csv(client->NodeId()) << ',' << Csv(tier) << ',' << reads << '\n'; + } + } +} +void InspectReplayTrace(Options* options) { + bench::TraceReader reader(options->trace_path); + options->workload.seed = reader.header().seed(); + const auto& metadata = reader.header().metadata(); + auto metadata_value = [&](const char* key) -> const std::string* { + const auto found = metadata.find(key); + return found == metadata.end() ? nullptr : &found->second; + }; + if (const auto* value = metadata_value("profile")) options->profile = *value; + // A recorded trace carries keys and sizes but not payloads, so the expected + // bytes cannot be rederived and GET validation would fail on every hit. + if (const auto* value = metadata_value("payload_mode"); + value != nullptr && *value == "external") { + options->runner.validate_get_payloads = false; + } + if (const auto* value = metadata_value("source"); value != nullptr && *value == "production") { + options->trace_defines_workload = true; + } + if (const auto* value = metadata_value("keys")) { + options->workload.key_count = ParseUnsigned(*value, "trace metadata keys"); + } + if (const auto* value = metadata_value("min_value_bytes")) { + options->workload.min_value_size = + ParseUnsigned(*value, "trace metadata min_value_bytes"); + } + if (const auto* value = metadata_value("max_value_bytes")) { + options->workload.max_value_size = + ParseUnsigned(*value, "trace metadata max_value_bytes"); + } + if (const auto* value = metadata_value("size_distribution")) { + options->size_distribution = *value; + } + if (const auto* value = metadata_value("read_ratio")) { + options->workload.read_ratio = ParseDouble(*value, "trace metadata read_ratio"); + } + if (const auto* value = metadata_value("batch")) { + options->workload.batch_size = ParseUnsigned(*value, "trace metadata batch"); + } + if (const auto* value = metadata_value("qps")) { + options->workload.qps = ParseDouble(*value, "trace metadata qps"); + } + if (!options->clients_explicit) { + const auto found = metadata.find("clients"); + if (found != metadata.end()) { + options->workload.client_count = + ParseUnsigned(found->second, "trace metadata clients"); + } + } + bool has_events = false; + uint32_t max_client = 0; + uint64_t event_count = 0; + ::umbp::benchmark::WorkloadEvent event; + while (reader.ReadNext(&event)) { + has_events = true; + ++event_count; + max_client = std::max(max_client, event.client_id()); + if (event.value_size() > options->max_ssd_staging_object_bytes) { + UsageError("trace contains an SSD value larger than the contiguous staging arena"); + } + } + options->workload.operation_count = event_count; + if (!options->clients_explicit && metadata.find("clients") == metadata.end() && has_events) { + if (max_client == std::numeric_limits::max()) { + UsageError("trace client id is too large"); + } + options->workload.client_count = max_client + 1; + } + if (has_events && max_client >= options->workload.client_count) { + UsageError("trace client id exceeds --clients"); + } +} +int RunBenchmark(Options options) { + if (options.command == "replay") { + InspectReplayTrace(&options); + ValidateOptions(&options); + } + std::unique_ptr source = + options.command == "replay" + ? std::unique_ptr( + std::make_unique(options.trace_path)) + : std::unique_ptr( + std::make_unique(options.workload)); + MasterHarness master(options); + auto clients = StartClients(options, master.Address()); + bench::PoolWorkloadClient adapter(&clients); + Settle(clients, options.settle_ms); + bench::WorkloadRunner runner(&adapter, options.runner); + const auto start = std::chrono::steady_clock::now(); + const auto metrics = runner.Run(source.get()); + const uint64_t wall_ns = + std::chrono::duration_cast(std::chrono::steady_clock::now() - start) + .count(); + Settle(clients, options.settle_ms); + PrintSummary(options, metrics, source->seed(), wall_ns); + PrintBackendPlacement(clients); + return metrics.total.failed == 0 ? 0 : 2; +} +} // namespace +int main(int argc, char** argv) { + try { + Options options = ParseOptions(argc, argv); + LoadPolicy(&options); + ValidateOptions(&options); + return options.command == "generate" ? GenerateTrace(options) + : RunBenchmark(std::move(options)); + } catch (const std::exception& error) { + std::cerr << "umbp_tier_bench: " << error.what() << '\n'; + return 1; + } +} diff --git a/src/umbp/distributed/distributed_client.cpp b/src/umbp/distributed/distributed_client.cpp index ff615b1d4..858469690 100644 --- a/src/umbp/distributed/distributed_client.cpp +++ b/src/umbp/distributed/distributed_client.cpp @@ -22,7 +22,9 @@ #include "umbp/distributed/distributed_client.h" #include +#include #include +#include #include "mori/io/engine.hpp" #include "mori/utils/mori_log.hpp" @@ -37,58 +39,108 @@ DistributedClient::DistributedClient(const UMBPConfig& config) : config_(config) } const auto& dc = config.distributed.value(); - + local_only_ = dc.master_config.master_address.empty(); + + // All three ownership blocks are lowered unconditionally; dc.medium selects + // the one PoolClient::Init actually builds a backend from (exactly one medium + // per node — see UMBPMedium in common/config.h). Lowering all three keeps + // this function free of medium branching and lets one deployment template + // carry dram/hbm/ssd sizing while choosing between them with a single field. + // + // Ownership (ptr, hugepages, NUMA, prefault) moved into PoolClient's DRAM + // PageBackend — it self-allocates at Init() instead of DistributedClient + // calling HostMemAllocator and handing over a buffer pointer + // (backend-agnostic refactor Phase 2b, design doc §1 item 4). Only the + // sizing/policy knobs cross this boundary now. + DramOwnershipConfig dram_ownership; + dram_ownership.buffer_sizes = {config.dram.capacity_bytes}; + dram_ownership.use_hugepages = config.dram.use_hugepages; + dram_ownership.hugepage_size = config.dram.hugepage_size; + dram_ownership.numa_node = config.dram.numa_node; + dram_ownership.prefault = config.dram.prefault; + + // SSD came back to the distributed data plane as a MediumBackend: Phase 0 + // unwired the old PeerSsdManager special case, and PoolClient::Init builds an + // SsdBackend from PoolClientConfig::ssd the same way it builds DRAM and HBM. + // + // PeerSsdConfig::enabled is set by ToPoolClientConfig from dc.medium, NOT + // from config.ssd.enabled — that flag defaults to true and describes the + // LOCAL-mode tier, so keying the distributed medium off it would make every + // existing distributed deployment start advertising SSD capacity it had never + // opted into. + PeerSsdConfig ssd_ownership; + ssd_ownership.ssd = config.ssd; + + // HBM is different again: dc.hbm carries everything PoolClient's HBM + // PageBackend needs (device/capacity), and ToPoolClientConfig lowers it + // directly from dc — no ownership struct crosses this boundary the way dram's + // hugepages/NUMA/prefault knobs do, because hipMalloc has none of those + // dimensions. + // The ranged scratch arena. Plain anonymous host pages: it is registered with + // the transfer engine below, and unlike a medium pool it has no hugepage/NUMA + // policy to honour — it is touched once per remote ranged operation, not held + // as a cache. + // + // Zero-sized is the default and means "this deployment does not do ranged + // I/O": no arena, no registration, and SupportsRangedIO() reports false, so + // a client that never issues ranged operations stops paying for the arena. HostMemAllocator allocator; - HostBufferOptions opts; - opts.backing = config.dram.use_hugepages ? HostBufferBacking::kAnonymousHugetlb - : HostBufferBacking::kAnonymous; - opts.hugepage_size = config.dram.hugepage_size; - opts.numa_node = config.dram.numa_node; - opts.prefault = config.dram.prefault; - - dram_pool_handle_ = allocator.Alloc(config.dram.capacity_bytes, opts); - if (!dram_pool_handle_.valid()) { - throw std::runtime_error("DistributedClient: memory allocation failed for DRAM pool"); - } - dram_pool_ = dram_pool_handle_.ptr; - // Use mapped_size (>= capacity_bytes, rounded up to page/hugepage boundary) - // so that RDMA registration, PeerDramAllocator capacity, and master-reported - // tier_capacities all agree on a single value. This means the effective - // pool size may exceed config.dram.capacity_bytes by up to one hugepage. - // NOTE: if hugepage_size is not a multiple of dram_page_size, the tail - // bytes that don't form a complete dram_page are reported in - // tier_capacities but never allocated by PeerDramAllocator; heartbeat's - // TierCapacitiesSnapshot() will correct master's view. Both default to - // 2 MiB, so this only matters with non-default page size combinations. - dram_pool_size_ = dram_pool_handle_.mapped_size; - - // Lower SSD config to the peer. When ssd.enabled, the peer builds a - // PeerSsdManager (SSDTier backend) from the SSD config (UMBPSsdConfig) and - // reports SSD capacity via TierType::SSD; when disabled, behavior is exactly - // DRAM-only (no PeerSsdManager, no SSD capacity, no SSD event source). - std::map tier_capacities = { - {TierType::DRAM, {dram_pool_size_, dram_pool_size_}}}; - PeerSsdConfig ssd_cfg; - if (config_.ssd.enabled) { - ssd_cfg.enabled = true; - ssd_cfg.ssd = config_.ssd; - const uint64_t ssd_cap = config_.ssd.capacity_bytes; - tier_capacities[TierType::SSD] = {ssd_cap, ssd_cap}; + HostBufferOptions scratch_opts; + // Two separate arenas — GET and PUT — each of ranged_scratch_size, so the two + // directions never share a buffer or a lock and can run concurrently. + if (dc.ranged_scratch_size > 0) { + ranged_get_scratch_handle_ = allocator.Alloc(dc.ranged_scratch_size, scratch_opts); + ranged_put_scratch_handle_ = allocator.Alloc(dc.ranged_scratch_size, scratch_opts); + if (!ranged_get_scratch_handle_.valid() || !ranged_put_scratch_handle_.valid()) { + // HostBufferHandle is not RAII and the destructor/Close() won't run when a + // constructor throws, so free whatever did allocate before bailing out + // (Free() is a no-op on an invalid handle). + allocator.Free(ranged_get_scratch_handle_); + allocator.Free(ranged_put_scratch_handle_); + throw std::runtime_error("DistributedClient: memory allocation failed for ranged scratch"); + } + ranged_get_scratch_ = ranged_get_scratch_handle_.ptr; + ranged_get_scratch_size_ = ranged_get_scratch_handle_.mapped_size; + ranged_put_scratch_ = ranged_put_scratch_handle_.ptr; + ranged_put_scratch_size_ = ranged_put_scratch_handle_.mapped_size; } - auto pc_config = ToPoolClientConfig(dc, - /*dram_buffers=*/{{dram_pool_, dram_pool_size_}}, - std::move(tier_capacities), std::move(ssd_cfg)); + + auto pc_config = ToPoolClientConfig(dc, std::move(dram_ownership), std::move(ssd_ownership), + config.dram.high_watermark, config.dram.low_watermark); + pc_config.ranged_get_scratch_buffer = ranged_get_scratch_; + pc_config.ranged_get_scratch_size = ranged_get_scratch_size_; + pc_config.ranged_put_scratch_buffer = ranged_put_scratch_; + pc_config.ranged_put_scratch_size = ranged_put_scratch_size_; pc_config.copy_pipeline = config_.copy_pipeline; + auto release_scratch = [this] { + HostMemAllocator cleanup; + cleanup.Free(ranged_get_scratch_handle_); + cleanup.Free(ranged_put_scratch_handle_); + ranged_get_scratch_ = nullptr; + ranged_get_scratch_size_ = 0; + ranged_put_scratch_ = nullptr; + ranged_put_scratch_size_ = 0; + }; + pool_client_ = std::make_unique(std::move(pc_config)); if (!pool_client_->Init()) { pool_client_.reset(); - HostMemAllocator cleanup_allocator; - cleanup_allocator.Free(dram_pool_handle_); - dram_pool_ = nullptr; - dram_pool_size_ = 0; + release_scratch(); throw std::runtime_error("DistributedClient: PoolClient::Init() failed"); } + // Registered explicitly rather than through a backend: each arena is a remote + // endpoint for RDMA reads and writes, so it needs a memory region even though + // no medium owns it. Skipped entirely when the deployment did not opt in. + if ((ranged_get_scratch_ && + !pool_client_->RegisterMemory(ranged_get_scratch_, ranged_get_scratch_size_)) || + (ranged_put_scratch_ && + !pool_client_->RegisterMemory(ranged_put_scratch_, ranged_put_scratch_size_))) { + pool_client_->Shutdown(); + pool_client_.reset(); + release_scratch(); + throw std::runtime_error("DistributedClient: ranged scratch registration failed"); + } std::string tags_str; for (const auto& t : dc.master_config.tags) { @@ -96,17 +148,49 @@ DistributedClient::DistributedClient(const UMBPConfig& config) : config_(config) tags_str += t; } + // Only the live medium's sizing is logged: the other two blocks were lowered + // but never allocated from, and printing them invites reading a config that + // is not in effect (the exact confusion the single selector removes). + const auto mb = [](uint64_t bytes) { return std::to_string(bytes / (1024 * 1024)); }; + std::string medium_desc; + switch (dc.medium) { + case UMBPMedium::DRAM: + medium_desc = "DRAM pool=" + mb(config_.dram.capacity_bytes) + + "MB hugepages=" + (config_.dram.use_hugepages ? "true" : "false") + + " hugepage_size=" + mb(config_.dram.hugepage_size) + + "MB numa_node=" + std::to_string(config_.dram.numa_node); + break; + case UMBPMedium::HBM: + medium_desc = + "HBM pool=" + mb(dc.hbm.capacity_bytes) + "MB device=" + std::to_string(dc.hbm.device); + break; + case UMBPMedium::SSD: + // Slot COUNT only, not ssd_staging_buffer_size. The arena SsdBackend + // actually allocates is staging_pages * page_size (see its own + // "[SsdBackend] Init ... arena_bytes=" line, which is authoritative); + // ssd_staging_buffer_size does not size it, so printing that number here + // stated a capacity the node did not have — 6144MB against a real 4096MB + // arena in a 2 MiB-page run. + // + // The slot count is the number worth showing, because it is the SSD + // medium's read-concurrency limit. Transient shortfalls now surface as + // BUSY and are retried; a batch whose own working set exceeds the arena + // is a permanent resolve failure. + medium_desc = "SSD pool=" + mb(config_.ssd.capacity_bytes) + + "MB backend=" + config_.ssd.ssd_backend + " dir=" + config_.ssd.storage_dir + + " staging_slots=" + std::to_string(dc.ssd_staging_buffer_slots); + break; + } + MORI_UMBP_INFO( "[DistributedClient] initialized — " - "node_id={} node_address={} master={} " - "dram_pool={}MB hugepages={} hugepage_size={}MB numa_node={} " - "dram_page_size={}KB staging_buffer={}MB peer_port={} cache_remote={} " + "node_id={} node_address={} master={} medium=[{}] " + "page_size={}KB staging_buffer={}MB peer_port={} cache_remote={} " "io_engine={}:{} tags=[{}]", dc.master_config.node_id, dc.master_config.node_address, dc.master_config.master_address, - dram_pool_size_ / (1024 * 1024), config_.dram.use_hugepages, - config_.dram.hugepage_size / (1024 * 1024), config_.dram.numa_node, dc.dram_page_size / 1024, - dc.staging_buffer_size / (1024 * 1024), dc.peer_service_port, dc.cache_remote_fetches, - dc.io_engine.host, dc.io_engine.port, tags_str); + medium_desc, dc.dram_page_size / 1024, dc.staging_buffer_size / (1024 * 1024), + dc.peer_service_port, dc.cache_remote_fetches, dc.io_engine.host, dc.io_engine.port, + tags_str); } DistributedClient::~DistributedClient() { Close(); } @@ -185,6 +269,36 @@ std::vector DistributedClient::BatchGet(const std::vector& ke return pool_client_->BatchGet(keys, dst_ptrs, sizes); } +std::vector DistributedClient::BatchGetRanges( + const std::vector& keys, const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& src_offsets) { + if (closing_) return std::vector(keys.size(), false); + std::shared_lock lk(op_mutex_); + if (closed_ || !pool_client_) return std::vector(keys.size(), false); + std::vector> dst_ptrs(dsts.size()); + for (size_t i = 0; i < dsts.size(); ++i) { + dst_ptrs[i].reserve(dsts[i].size()); + for (uintptr_t ptr : dsts[i]) dst_ptrs[i].push_back(reinterpret_cast(ptr)); + } + return pool_client_->BatchGetRanges(keys, dst_ptrs, sizes, src_offsets); +} + +std::vector DistributedClient::BatchPutRanges( + const std::vector& keys, const std::vector& object_sizes, + const std::vector>& srcs, const std::vector>& sizes, + const std::vector>& dst_offsets) { + if (closing_) return std::vector(keys.size(), false); + std::shared_lock lk(op_mutex_); + if (closed_ || !pool_client_) return std::vector(keys.size(), false); + std::vector> src_ptrs(srcs.size()); + for (size_t i = 0; i < srcs.size(); ++i) { + src_ptrs[i].reserve(srcs[i].size()); + for (uintptr_t ptr : srcs[i]) src_ptrs[i].push_back(reinterpret_cast(ptr)); + } + return pool_client_->BatchPutRanges(keys, object_sizes, src_ptrs, sizes, dst_offsets); +} + std::vector DistributedClient::BatchExists(const std::vector& keys) const { if (closing_) return std::vector(keys.size(), false); std::shared_lock lk(op_mutex_); @@ -216,11 +330,12 @@ size_t DistributedClient::BatchExistsConsecutive(const std::vector& // RegisterMemory / DeregisterMemory // --------------------------------------------------------------------------- -bool DistributedClient::RegisterMemory(uintptr_t ptr, size_t size) { +bool DistributedClient::RegisterMemory(uintptr_t ptr, size_t size, mori::io::MemoryLocationType loc, + int device, MemoryRegistration mode) { if (closing_) return false; std::shared_lock lk(op_mutex_); if (closed_) return false; - return pool_client_->RegisterMemory(reinterpret_cast(ptr), size); + return pool_client_->RegisterMemory(reinterpret_cast(ptr), size, loc, device, mode); } void DistributedClient::DeregisterMemory(uintptr_t ptr) { @@ -259,7 +374,9 @@ bool DistributedClient::Flush() { if (closing_) return true; std::shared_lock lk(op_mutex_); if (closed_ || !pool_client_) return true; - pool_client_->Master().FlushHeartbeat(); + // Nothing to flush to on a node with no master; the local media are already + // authoritative for everything they hold. + if (pool_client_->HasMaster()) pool_client_->Master().FlushHeartbeat(); return true; } @@ -274,11 +391,16 @@ void DistributedClient::Close() { pool_client_.reset(); } - if (dram_pool_) { + // Only after Shutdown has deregistered the regions — the arenas are + // caller-owned and PoolClient holds a memory region over each until then. + if (ranged_get_scratch_ || ranged_put_scratch_) { HostMemAllocator allocator; - allocator.Free(dram_pool_handle_); - dram_pool_ = nullptr; - dram_pool_size_ = 0; + allocator.Free(ranged_get_scratch_handle_); + allocator.Free(ranged_put_scratch_handle_); + ranged_get_scratch_ = nullptr; + ranged_get_scratch_size_ = 0; + ranged_put_scratch_ = nullptr; + ranged_put_scratch_size_ = 0; } MORI_UMBP_INFO("[DistributedClient] closed"); @@ -286,6 +408,40 @@ void DistributedClient::Close() { bool DistributedClient::IsDistributed() const { return true; } +bool DistributedClient::SupportsRangedIO() const { + // The scratch arenas are the whole opt-in. They are purely a remote-path + // resource -- remote gets are staged in them, remote puts assembled in them -- + // and they default to zero, so a deployment that never issues ranged I/O + // allocates and registers nothing. + // + // Deliberately NOT gated on the medium. SSD used to be excluded here on the + // reasoning that its bytes are not addressable by the transfer layer, but + // that is exactly what SsdBackend already solves for every other operation: + // it publishes a registered host staging arena and spills behind it, so a + // resolved SSD key reaches BuildLocalRangeTransfers as ordinary pages and the + // object-range -> page-range arithmetic never learns which medium it is + // walking (ssd_backend.h). All four ranged paths therefore work on an SSD + // node: get and put, each local and remote. + // + // What the medium changes is the saving, not the correctness. A resolve + // stages the whole object before anyone says which bytes they want, so a + // ranged get off SSD still reads the full object from the device and saves + // only the final copy into the caller's buffers. Ranged put is unaffected: + // it tiles its object, so the single whole-object write was always going to + // happen. Making the device read follow the requested extent needs an extent + // on the resolve path and is a separate change (doc/design-ssd-ranged-io.md, + // D1) -- this flag was never what stood in its way. + // + // The one exception is a node with no master. Nothing can route a key off + // this node, so every ranged operation it will ever serve is the local path + // -- which reads and writes the medium's own pages directly and never + // touches an arena. Requiring one there would make an embedded deployment + // either report ranged I/O as unsupported or allocate two arenas it cannot + // use. + if (local_only_) return true; + return ranged_get_scratch_size_ != 0 && ranged_put_scratch_size_ != 0; +} + bool DistributedClient::ReportExternalKvBlocks(const std::vector& hashes, TierType tier) { if (!pool_client_) return false; diff --git a/src/umbp/distributed/master/eviction_manager.cpp b/src/umbp/distributed/master/eviction_manager.cpp index 323f1dc7d..7e732528d 100644 --- a/src/umbp/distributed/master/eviction_manager.cpp +++ b/src/umbp/distributed/master/eviction_manager.cpp @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -31,6 +32,7 @@ #include #include "mori/utils/mori_log.hpp" +#include "umbp/common/grpc_limits.h" #include "umbp/distributed/master/evict_strategy.h" #include "umbp/distributed/master/master_metadata_store.h" #include "umbp/distributed/types.h" @@ -96,11 +98,12 @@ void EvictionManager::RunOnce() { for (const auto& client : clients) { for (const auto& [tier, cap] : client.tier_capacities) { - // SSD eviction is purely peer-local. Master must NOT turn an SSD - // overload into an EvictKey: EvictKey only acts on the peer's - // PeerDramAllocator, so it would wrongly evict the DRAM copy of the same - // key while leaving SSD untouched. - if (tier == TierType::SSD) continue; + // No tier is special-cased here any more (backend-agnostic refactor + // Phase 4). The old `if (tier == SSD) continue` existed because EvictKey + // only ever reached the peer's DRAM allocator, so an SSD overload would + // have evicted the DRAM copy of the key instead; since Phase 3 EvictKey + // fans out to every backend by key, so an overloaded medium evicts from + // the medium that is actually overloaded. if (cap.total_bytes == 0) continue; uint64_t used = cap.total_bytes - cap.available_bytes; double usage = static_cast(used) / static_cast(cap.total_bytes); @@ -120,15 +123,15 @@ void EvictionManager::RunOnce() { // Ask the store for eviction-eligible candidates in the overloaded buckets. // The store is policy-neutral: it returns rows ordered by the hint (LRU here, // so an indexed backend can push the ordering down) but makes no eviction - // decision and never sees the byte budget. max_per_bucket=0 → no cap, so the - // strategy gets full candidate visibility. + // decision and never sees the byte budget. Limit each bucket to the configured + // round size; a still-overloaded bucket is revisited on the next interval. std::vector buckets; buckets.reserve(bytes_to_free.size()); for (const auto& [ntk, _bytes] : bytes_to_free) buckets.push_back(ntk); - auto candidates_by_bucket = - store_.EnumerateEvictionCandidates(buckets, EvictionOrder::kLeastRecentlyAccessed, - /*max_per_bucket=*/0, std::chrono::system_clock::now()); + auto candidates_by_bucket = store_.EnumerateEvictionCandidates( + buckets, EvictionOrder::kLeastRecentlyAccessed, config_.evict_batch_size, + std::chrono::system_clock::now()); if (candidates_by_bucket.empty()) { MORI_UMBP_DEBUG("[EvictionManager] No eviction candidates found"); return; @@ -174,10 +177,32 @@ void EvictionManager::RunOnce() { node_id, keys.size()); continue; } - // Master state is unchanged here — REMOVE events on the peer's - // next heartbeat are what shrink the index. Re-eviction next - // round is safe because peer Evict is idempotent. - dispatcher_->DispatchEvictKey(node_id, it->second, std::move(keys)); + // Chunked to the shared gRPC message limit. An eviction round frees a + // FRACTION of the tier, so the victim list grows with the deployment while + // the limit does not: at 512 GiB a round selected 155,157 keys, which is + // 18.7 MB of `repeated string` and was refused whole on all 81 rounds -- + // protobuf has no partial delivery, so being over the limit freed nothing + // rather than freeing most of it, and the tier could not un-fill itself. + // + // Splitting is safe in a way that most batch APIs are not: peer Evict is + // idempotent, master state is unchanged here (REMOVE events on the peer's + // next heartbeat are what shrink the index), and a chunk that fails is + // retried next round. So N chunks reach the same state as one message + // would have, and a partial round still makes progress. + size_t sent = 0; + size_t chunks = 0; + while (sent < keys.size()) { + const size_t take = GrpcMaxItemsPerBatch(keys, sent); + std::vector chunk(std::make_move_iterator(keys.begin() + sent), + std::make_move_iterator(keys.begin() + sent + take)); + sent += take; + ++chunks; + dispatcher_->DispatchEvictKey(node_id, it->second, std::move(chunk)); + } + if (chunks > 1) { + MORI_UMBP_DEBUG("[EvictionManager] node={} {} keys dispatched in {} chunks", node_id, + keys.size(), chunks); + } } } diff --git a/src/umbp/distributed/master/in_memory_master_metadata_store.cpp b/src/umbp/distributed/master/in_memory_master_metadata_store.cpp index 08c5c24c1..d51bd85a2 100644 --- a/src/umbp/distributed/master/in_memory_master_metadata_store.cpp +++ b/src/umbp/distributed/master/in_memory_master_metadata_store.cpp @@ -67,11 +67,14 @@ size_t IndexShardCount() { // Caller MUST hold the shard's unique lock. Returns a pointer into `locations` // that's stable until the next mutation. std::pair FindOrInsertLocation(std::vector& locations, - const std::string& node_id, TierType tier) { + const std::string& node_id, TierType tier, + const std::string& logical_tier) { for (auto& loc : locations) { - if (loc.node_id == node_id && loc.tier == tier) return {&loc, false}; + if (loc.node_id == node_id && loc.tier == tier && loc.logical_tier == logical_tier) { + return {&loc, false}; + } } - locations.push_back(Location{node_id, /*size=*/0, tier}); + locations.push_back(Location{node_id, /*size=*/0, tier, logical_tier}); return {&locations.back(), true}; } @@ -103,7 +106,7 @@ size_t InMemoryMasterMetadataStore::ApplyAddOrRemoveLocked( entry.last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); entry.atomic_access_count.store(0, std::memory_order_relaxed); } - auto [loc, inserted] = FindOrInsertLocation(entry.locations, node_id, ev.tier); + auto [loc, inserted] = FindOrInsertLocation(entry.locations, node_id, ev.tier, ev.logical_tier); // Idempotent; must run on duplicate ADDs too. shard.node_to_keys[node_id].insert(ev.key); if (!inserted) { @@ -121,10 +124,12 @@ size_t InMemoryMasterMetadataStore::ApplyAddOrRemoveLocked( if (it == shard.entries.end()) return 0; auto& locs = it->second.locations; const size_t before = locs.size(); - locs.erase( - std::remove_if(locs.begin(), locs.end(), - [&](const Location& l) { return l.node_id == node_id && l.tier == ev.tier; }), - locs.end()); + locs.erase(std::remove_if(locs.begin(), locs.end(), + [&](const Location& l) { + return l.node_id == node_id && l.tier == ev.tier && + l.logical_tier == ev.logical_tier; + }), + locs.end()); if (locs.size() == before) return 0; // find(), not operator[]: don't grow an empty bucket for strangers. if (!HasLocationForNode(it->second.locations, node_id)) { @@ -217,7 +222,8 @@ void InMemoryMasterMetadataStore::ReplaceNodeLocationsInShards( entry.last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); entry.atomic_access_count.store(0, std::memory_order_relaxed); } - auto [loc, inserted] = FindOrInsertLocation(entry.locations, node_id, ev->tier); + auto [loc, inserted] = + FindOrInsertLocation(entry.locations, node_id, ev->tier, ev->logical_tier); (void)inserted; loc->size = ev->size; sh.node_to_keys[node_id].insert(ev->key); @@ -284,6 +290,7 @@ bool InMemoryMasterMetadataStore::RegisterClient(const ClientRegistration& regis record.last_heartbeat = now; record.registered_at = now; record.tier_capacities = registration.tier_capacities; + record.logical_tier_capacities = registration.logical_tier_capacities; record.peer_address = registration.peer_address; record.engine_desc_bytes = registration.engine_desc_bytes; record.last_applied_seq = 0; @@ -323,7 +330,7 @@ void InMemoryMasterMetadataStore::UnregisterClient(const std::string& node_id) { HeartbeatResult InMemoryMasterMetadataStore::ApplyHeartbeat( const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, const std::map& caps, const std::vector& events, - bool is_full_sync) { + bool is_full_sync, const std::map& logical_caps) { // Cross-domain write (approach A), hot path: seq-CAS + client-record update // under meta_mutex_, then the block apply AFTER releasing it so heartbeats to // different shards apply in parallel (PR #440's win) instead of serializing on @@ -356,6 +363,7 @@ HeartbeatResult InMemoryMasterMetadataStore::ApplyHeartbeat( record.last_heartbeat = now; record.status = ClientStatus::ALIVE; record.tier_capacities = caps; + record.logical_tier_capacities = logical_caps; record.last_applied_seq = seq; apply = true; } diff --git a/src/umbp/distributed/master/master_client.cpp b/src/umbp/distributed/master/master_client.cpp index 9a6528df8..f15b5112d 100644 --- a/src/umbp/distributed/master/master_client.cpp +++ b/src/umbp/distributed/master/master_client.cpp @@ -23,20 +23,22 @@ #include +#include #include #include #include #include +#include #include #include #include "mori/utils/mori_log.hpp" #include "umbp.grpc.pb.h" #include "umbp/common/env_time.h" +#include "umbp/common/grpc_limits.h" #include "umbp/distributed/master/master_metrics.h" #include "umbp/distributed/master/rpc_latency_timer.h" -#include "umbp/distributed/peer/peer_dram_allocator.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" +#include "umbp/distributed/pool/peer_pool.h" namespace mori::umbp { @@ -53,6 +55,36 @@ int RpcShutdownTimeoutMs() { return v; } +int HeartbeatRpcTimeoutMs() { + static const int timeout = [] { + const char* raw = std::getenv("UMBP_HEARTBEAT_RPC_TIMEOUT_MS"); + if (!raw || raw[0] == '\0') return 3000; + const int parsed = std::atoi(raw); + return parsed > 0 ? parsed : 3000; + }(); + return timeout; +} + +size_t HeartbeatEventsPerBundle() { + static const size_t value = [] { + const char* raw = std::getenv("UMBP_HEARTBEAT_EVENTS_PER_BUNDLE"); + if (!raw || raw[0] == '\0') return size_t{1024}; + const unsigned long long parsed = std::strtoull(raw, nullptr, 10); + return parsed > 0 ? static_cast(parsed) : size_t{1024}; + }(); + return value; +} + +size_t HeartbeatMaxBundlesPerRpc() { + static const size_t value = [] { + const char* raw = std::getenv("UMBP_HEARTBEAT_MAX_BUNDLES_PER_RPC"); + if (!raw || raw[0] == '\0') return size_t{16}; + const unsigned long long parsed = std::strtoull(raw, nullptr, 10); + return parsed > 0 ? static_cast(parsed) : size_t{16}; + }(); + return value; +} + uint64_t MetricsReportIntervalMs() { static const uint64_t v = static_cast(GetEnvMilliseconds("UMBP_METRICS_REPORT_INTERVAL_MS", @@ -66,6 +98,11 @@ ::umbp::UMBPMaster::Stub* GetStub(void* ptr) { return static_cast<::umbp::UMBPMa ::umbp::TierType ToProtoTier(TierType t) { return static_cast<::umbp::TierType>(t); } TierType FromProtoTier(::umbp::TierType t) { return static_cast(t); } +uint64_t SaturatingAdd(uint64_t lhs, uint64_t rhs) { + return rhs > std::numeric_limits::max() - lhs ? std::numeric_limits::max() + : lhs + rhs; +} + ::umbp::KvEvent::Kind ToProtoEventKind(KvEvent::Kind kind) { switch (kind) { case KvEvent::Kind::ADD: @@ -83,6 +120,22 @@ void FillTierCapacities(::google::protobuf::RepeatedPtrField<::umbp::TierCapacit tc->set_tier(ToProtoTier(tier)); tc->set_total_capacity_bytes(cap.total_bytes); tc->set_available_capacity_bytes(cap.available_bytes); + tc->set_max_allocatable_bytes(cap.max_allocatable_bytes); + } +} + +void FillLogicalTierCapacities( + ::google::protobuf::RepeatedPtrField<::umbp::LogicalTierCapacity>* dst, + const std::map& src) { + for (const auto& [name, logical] : src) { + auto* cap = dst->Add(); + cap->set_name(name); + cap->set_representative_tier(ToProtoTier(logical.representative_tier)); + cap->set_total_capacity_bytes(logical.capacity.total_bytes); + cap->set_available_capacity_bytes(logical.capacity.available_bytes); + cap->set_max_allocatable_bytes(logical.capacity.max_allocatable_bytes); + cap->set_put_eligible(logical.put_eligible); + cap->set_peak_member_utilization(logical.peak_member_utilization); } } @@ -100,7 +153,7 @@ void FillTierKvCounts(::google::protobuf::RepeatedPtrField<::umbp::TierKvCount>* } } -void FillBundle(::umbp::EventBundle* dst, const EventBundle& src) { +void FillBundle(::umbp::EventBundle* dst, const EventBundle& src, bool include_logical_tiers) { dst->set_seq(src.seq); for (const auto& ev : src.events) { auto* pe = dst->add_events(); @@ -108,6 +161,7 @@ void FillBundle(::umbp::EventBundle* dst, const EventBundle& src) { pe->set_key(ev.key); pe->set_tier(ToProtoTier(ev.tier)); pe->set_size(ev.size); + if (include_logical_tiers) pe->set_logical_tier(ev.logical_tier); } } @@ -117,8 +171,7 @@ MasterClient::MasterClient(const UMBPMasterClientConfig& config) : config_(config), stub_(nullptr, [](void* p) { delete static_cast<::umbp::UMBPMaster::Stub*>(p); }) { grpc::ChannelArguments args; - args.SetMaxReceiveMessageSize(64 * 1024 * 1024); - args.SetMaxSendMessageSize(64 * 1024 * 1024); + ApplyGrpcLimits(&args); channel_ = grpc::CreateCustomChannel(config.master_address, grpc::InsecureChannelCredentials(), args); stub_.reset(::umbp::UMBPMaster::NewStub(channel_).release()); @@ -134,9 +187,10 @@ MasterClient::~MasterClient() { } } -grpc::Status MasterClient::RegisterSelf(const std::map& tier_capacities, - const std::string& peer_address, - const std::vector& engine_desc_bytes) { +grpc::Status MasterClient::RegisterSelf( + const std::map& tier_capacities, const std::string& peer_address, + const std::vector& engine_desc_bytes, + const std::map& logical_tier_capacities) { ScopedRpcTimer _rpc_timer(this, "RegisterClient"); if (registered_) { return grpc::Status(grpc::StatusCode::ALREADY_EXISTS, "node is already registered"); @@ -148,6 +202,7 @@ grpc::Status MasterClient::RegisterSelf(const std::map& req.set_peer_address(peer_address); req.set_engine_desc(engine_desc_bytes.data(), engine_desc_bytes.size()); FillTierCapacities(req.mutable_tier_capacities(), tier_capacities); + FillLogicalTierCapacities(req.mutable_logical_tier_capacities(), logical_tier_capacities); for (const auto& tag : config_.tags) req.add_tags(tag); ::umbp::RegisterClientResponse resp; @@ -160,17 +215,24 @@ grpc::Status MasterClient::RegisterSelf(const std::map& } if (resp.heartbeat_interval_ms() > 0) heartbeat_interval_ms_ = resp.heartbeat_interval_ms(); + supports_max_allocatable_capacity_ = resp.supports_max_allocatable_bytes(); + supports_logical_tiers_ = resp.supports_logical_tiers(); + aggregate_backend_capacities_ = + aggregate_backend_capacities_requested_ && supports_max_allocatable_capacity_; { std::lock_guard lock(hb_state_mutex_); hb_last_acked_seq_ = 0; next_bundle_seq_ = 1; outbox_.clear(); - full_sync_pending_ = false; + full_sync_pending_ = supports_logical_tiers_ && peer_pool_ != nullptr; } { std::lock_guard lock(caps_mutex_); current_capacities_ = tier_capacities; } + registration_capacities_ = tier_capacities; + registered_peer_address_ = peer_address; + registered_engine_desc_bytes_ = engine_desc_bytes; registered_ = true; MORI_UMBP_INFO("[Client] Registered with master (heartbeat={}ms)", heartbeat_interval_ms_); StartMetricsReporting(); @@ -220,6 +282,7 @@ grpc::Status MasterClient::RoutePut(const std::string& key, uint64_t block_size, .node_id = resp.node_id(), .peer_address = resp.peer_address(), .tier = FromProtoTier(resp.tier()), + .logical_tier = resp.logical_tier(), }; break; case ::umbp::ROUTE_PUT_OUTCOME_ALREADY_EXISTS: @@ -259,6 +322,7 @@ grpc::Status MasterClient::RouteGet(const std::string& key, r.tier = FromProtoTier(resp.tier()); r.size = resp.size(); r.peer_address = resp.peer_address(); + r.logical_tier = resp.logical_tier(); *out_result = std::move(r); return grpc::Status::OK; } @@ -297,6 +361,7 @@ grpc::Status MasterClient::BatchRoutePut(const std::vector& keys, .node_id = e.node_id(), .peer_address = e.peer_address(), .tier = FromProtoTier(e.tier()), + .logical_tier = e.logical_tier(), }; break; case ::umbp::ROUTE_PUT_OUTCOME_ALREADY_EXISTS: @@ -333,7 +398,8 @@ grpc::Status MasterClient::BatchRouteGet(const std::vector& keys, // Columnar response: node_ref[i] is a 1-based index into resp.nodes() // (0 = not found); tier[i] / size[i] are parallel per-key arrays. const int n = resp.node_ref_size(); - if (resp.tier_size() != n || resp.size_size() != n) { + if (resp.tier_size() != n || resp.size_size() != n || + (resp.logical_tier_size() != 0 && resp.logical_tier_size() != n)) { return grpc::Status(grpc::StatusCode::INTERNAL, "BatchRouteGet: malformed columnar response (array length mismatch)"); } @@ -350,6 +416,7 @@ grpc::Status MasterClient::BatchRouteGet(const std::vector& keys, r.tier = FromProtoTier(resp.tier(i)); r.size = resp.size(i); r.peer_address = node.peer_address(); + if (resp.logical_tier_size() == n) r.logical_tier = resp.logical_tier(i); (*out)[i] = std::move(r); } return grpc::Status::OK; @@ -399,22 +466,25 @@ size_t AutoFlushEventThreshold() { } } // namespace -void MasterClient::SetPeerDramAllocator(PeerDramAllocator* dram_alloc) { - peer_alloc_ = dram_alloc; - AddOwnedLocationSource(dram_alloc); - if (dram_alloc != nullptr) { - dram_alloc->SetAutoFlushHook(AutoFlushEventThreshold(), [this] { FlushHeartbeat(); }); +void MasterClient::SetBackendRegistry(BackendRegistry* registry) { + registry_ = registry; + // Every medium's outbox feeds the same single-bundle heartbeat, so each one + // gets the same auto-flush hook — a backend added to the registry starts + // flushing without any change here. + const size_t threshold = AutoFlushEventThreshold(); + for (auto* backend : Backends()) { + backend->SetAutoFlushHook(threshold, [this] { FlushHeartbeat(); }); } } -void MasterClient::SetPeerSsdManager(PeerSsdManager* ssd_manager) { - ssd_manager_ = ssd_manager; - AddOwnedLocationSource(ssd_manager); +std::vector MasterClient::Backends() const { + return registry_ == nullptr ? std::vector{} : registry_->All(); } -void MasterClient::AddOwnedLocationSource(OwnedLocationSource* source) { - if (source == nullptr) return; - owned_sources_.push_back(source); +std::map MasterClient::OwnedKeyCountsByTier() const { + std::map counts; + for (auto* backend : Backends()) counts[backend->Tier()] += backend->OwnedKeyCount(); + return counts; } bool MasterClient::ClearFullSync() { @@ -422,13 +492,16 @@ bool MasterClient::ClearFullSync() { if (!registered_) return false; auto caps = SnapshotAndCacheTierCapacities(); - std::map kv_counts; - if (peer_alloc_ != nullptr) kv_counts = peer_alloc_->OwnedKeyCountByTier(); + auto kv_counts = OwnedKeyCountsByTier(); ::umbp::HeartbeatRequest req; req.set_node_id(config_.node_id); req.set_is_full_sync(true); FillTierCapacities(req.mutable_tier_capacities(), caps); + if (supports_logical_tiers_ && peer_pool_ != nullptr) { + FillLogicalTierCapacities(req.mutable_logical_tier_capacities(), + peer_pool_->LogicalTierCapacities()); + } FillTierKvCounts(req.mutable_tier_kv_counts(), kv_counts); { std::lock_guard state_lock(hb_state_mutex_); @@ -449,9 +522,11 @@ bool MasterClient::ClearFullSync() { return false; } - if (peer_alloc_ != nullptr) { - peer_alloc_->ClearFullSyncAcked(); - MORI_UMBP_INFO("[Client] Clear full-sync acked by master; allocator writes re-enabled"); + auto backends = Backends(); + for (auto* backend : backends) backend->ClearFullSyncAcked(); + if (!backends.empty()) { + MORI_UMBP_INFO("[Client] Clear full-sync acked by master; writes re-enabled on {} backend(s)", + backends.size()); } return true; } @@ -509,17 +584,23 @@ void MasterClient::HeartbeatLoop() { std::map MasterClient::SnapshotAndCacheTierCapacities() { std::map caps; bool have_live = false; - if (peer_alloc_ != nullptr) { - caps = peer_alloc_->TierCapacitiesSnapshot(); // DRAM/HBM, bitmap-derived - have_live = true; - } - if (ssd_manager_ != nullptr) { - auto [used, total] = ssd_manager_->Capacity(); - if (total > 0) { - const uint64_t avail = used < total ? total - used : 0; - caps[TierType::SSD] = TierCapacity{total, avail}; - have_live = true; + for (auto* backend : Backends()) { + auto backend_cap = backend->Capacity(); + if (aggregate_backend_capacities_) { + backend_cap.max_allocatable_bytes = + backend_cap.max_allocatable_bytes == 0 + ? backend_cap.available_bytes + : std::min(backend_cap.max_allocatable_bytes, backend_cap.available_bytes); + } + auto [it, inserted] = caps.try_emplace(backend->Tier(), backend_cap); + if (!inserted && aggregate_backend_capacities_) { + it->second.total_bytes = SaturatingAdd(it->second.total_bytes, backend_cap.total_bytes); + it->second.available_bytes = + SaturatingAdd(it->second.available_bytes, backend_cap.available_bytes); + it->second.max_allocatable_bytes = + std::max(it->second.max_allocatable_bytes, backend_cap.max_allocatable_bytes); } + have_live = true; } std::lock_guard lock(caps_mutex_); if (!have_live) return current_capacities_; @@ -536,8 +617,7 @@ bool MasterClient::SendHeartbeatOnce() { if (!registered_) return false; auto caps = SnapshotAndCacheTierCapacities(); - std::map kv_counts; - if (peer_alloc_ != nullptr) kv_counts = peer_alloc_->OwnedKeyCountByTier(); + auto kv_counts = OwnedKeyCountsByTier(); bool do_full_sync; { @@ -560,6 +640,10 @@ bool MasterClient::SendFullSyncHeartbeatLocked(const std::mapLogicalTierCapacities()); + } FillTierKvCounts(req.mutable_tier_kv_counts(), kv_counts); EventBundle snapshot; @@ -568,8 +652,10 @@ bool MasterClient::SendFullSyncHeartbeatLocked(const std::mapSnapshotOwnedKeysForFullSync() + : SnapshotAllBackendsForFullSync(Backends()); + FillBundle(req.add_bundles(), snapshot, supports_logical_tiers_); ::umbp::HeartbeatResponse resp; grpc::Status status = SendHeartbeatRpcLocked(req, &resp); @@ -585,21 +671,40 @@ bool MasterClient::SendDeltaHeartbeatLocked(const std::mapDrainPendingEvents() + : DrainAllBackends(Backends()); if (!new_events.empty()) { std::lock_guard state_lock(hb_state_mutex_); - outbox_.push_back(EventBundle{next_bundle_seq_++, std::move(new_events)}); + const size_t per_bundle = HeartbeatEventsPerBundle(); + for (size_t begin = 0; begin < new_events.size(); begin += per_bundle) { + const size_t end = std::min(new_events.size(), begin + per_bundle); + EventBundle bundle; + bundle.seq = next_bundle_seq_++; + bundle.events.reserve(end - begin); + std::move(new_events.begin() + begin, new_events.begin() + end, + std::back_inserter(bundle.events)); + outbox_.push_back(std::move(bundle)); + } } ::umbp::HeartbeatRequest req; req.set_node_id(config_.node_id); req.set_is_full_sync(false); FillTierCapacities(req.mutable_tier_capacities(), caps); + if (supports_logical_tiers_ && peer_pool_ != nullptr) { + FillLogicalTierCapacities(req.mutable_logical_tier_capacities(), + peer_pool_->LogicalTierCapacities()); + } FillTierKvCounts(req.mutable_tier_kv_counts(), kv_counts); { std::lock_guard state_lock(hb_state_mutex_); + size_t added = 0; for (const auto& bundle : outbox_) { - if (bundle.seq > hb_last_acked_seq_) FillBundle(req.add_bundles(), bundle); + if (bundle.seq <= hb_last_acked_seq_) continue; + if (added >= HeartbeatMaxBundlesPerRpc()) break; + FillBundle(req.add_bundles(), bundle, supports_logical_tiers_); + ++added; } } @@ -617,7 +722,25 @@ bool MasterClient::SendDeltaHeartbeatLocked(const std::map conservative_caps; + if (live_backends.empty()) { + conservative_caps = registration_capacities_; + } else { + for (auto* backend : live_backends) { + conservative_caps.try_emplace(backend->Tier(), backend->Capacity()); + } + } + FillTierCapacities(re_req.mutable_tier_capacities(), conservative_caps); + if (supports_logical_tiers_ && peer_pool_ != nullptr) { + FillLogicalTierCapacities(re_req.mutable_logical_tier_capacities(), + peer_pool_->LogicalTierCapacities()); + } for (const auto& tag : config_.tags) re_req.add_tags(tag); ::umbp::RegisterClientResponse re_resp; grpc::ClientContext re_ctx; @@ -628,6 +751,11 @@ bool MasterClient::SendDeltaHeartbeatLocked(const std::mapacked_seq(); - while (!outbox_.empty() && outbox_.front().seq <= hb_last_acked_seq_) { - outbox_.pop_front(); + bool more_pending = false; + bool made_progress = false; + { + std::lock_guard state_lock(hb_state_mutex_); + const uint64_t previous_ack = hb_last_acked_seq_; + hb_last_acked_seq_ = std::max(hb_last_acked_seq_, resp->acked_seq()); + made_progress = hb_last_acked_seq_ > previous_ack; + while (!outbox_.empty() && outbox_.front().seq <= hb_last_acked_seq_) { + outbox_.pop_front(); + } + more_pending = !outbox_.empty(); + } + // Drain a bounded backlog immediately in consecutive RPCs instead of paying + // one heartbeat interval per chunk. Each RPC remains independently bounded. + if (more_pending && made_progress && heartbeat_running_) { + flush_requested_ = true; + hb_cv_.notify_one(); } return status; } @@ -821,8 +962,8 @@ void MasterClient::MetricsLoop() { // Final flush: ship the last sub-interval of provider deltas before the // thread exits. PoolClient::Shutdown calls StopMetricsReporting() BEFORE // UnregisterSelf, so the master is still reachable here; without this final - // flush the last ( #include #include +#include #include #include #include @@ -34,6 +35,7 @@ #include "mori/utils/mori_log.hpp" #include "umbp.grpc.pb.h" #include "umbp/common/env_time.h" +#include "umbp/common/grpc_limits.h" #include "umbp/distributed/master/evict_strategy.h" #include "umbp/distributed/master/in_memory_master_metadata_store.h" #include "umbp/distributed/master/master_metadata_store.h" @@ -90,9 +92,27 @@ KvEvent FromProtoEvent(const ::umbp::KvEvent& pe) { ev.key = pe.key(); ev.tier = static_cast(pe.tier()); ev.size = pe.size(); + ev.logical_tier = pe.logical_tier(); return ev; } +std::map LogicalCapsFromProto( + const ::google::protobuf::RepeatedPtrField<::umbp::LogicalTierCapacity>& source) { + std::map out; + for (const auto& value : source) { + if (value.name().empty()) continue; + LogicalTierCapacity cap; + cap.representative_tier = static_cast(value.representative_tier()); + cap.capacity.total_bytes = value.total_capacity_bytes(); + cap.capacity.available_bytes = value.available_capacity_bytes(); + cap.capacity.max_allocatable_bytes = value.max_allocatable_bytes(); + cap.put_eligible = value.put_eligible(); + cap.peak_member_utilization = value.peak_member_utilization(); + out[value.name()] = cap; + } + return out; +} + int EvictKeyDeadlineMs() { static const int v = static_cast(GetEnvMilliseconds("UMBP_EVICTKEY_DEADLINE_MS", @@ -152,7 +172,8 @@ class MasterPeerStubPool : public EvictKeyDispatcher { auto it = entries_.find(node_id); if (it != entries_.end() && it->second.peer_address == peer_address) return it->second.stub; - auto channel = grpc::CreateChannel(peer_address, grpc::InsecureChannelCredentials()); + auto channel = grpc::CreateCustomChannel(peer_address, grpc::InsecureChannelCredentials(), + GrpcChannelArgs()); if (channel == nullptr) return nullptr; std::shared_ptr<::umbp::UMBPPeer::Stub> stub(::umbp::UMBPPeer::NewStub(channel).release()); entries_[node_id] = Entry{peer_address, stub}; @@ -234,6 +255,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser TierCapacity c; c.total_bytes = tc.total_capacity_bytes(); c.available_bytes = tc.available_capacity_bytes(); + c.max_allocatable_bytes = tc.max_allocatable_bytes(); caps[static_cast(tc.tier())] = c; } @@ -243,6 +265,10 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser registration.node_id = request->node_id(); registration.node_address = request->node_address(); registration.tier_capacities = caps; + if (config_.advertise_logical_tiers) { + registration.logical_tier_capacities = + LogicalCapsFromProto(request->logical_tier_capacities()); + } registration.peer_address = request->peer_address(); registration.engine_desc_bytes.assign(engine_desc_str.begin(), engine_desc_str.end()); registration.tags.assign(request->tags().begin(), request->tags().end()); @@ -260,11 +286,14 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser UpdateClientCountMetric(); UpdateClientCapacityMetrics(request->node_id(), caps); + UpdateClientLogicalTierMetrics(request->node_id(), registration.logical_tier_capacities); auto interval_ms = static_cast(config_.heartbeat_ttl.count() * 1000) / HeartbeatIntervalDivisor(); response->set_heartbeat_interval_ms(interval_ms); response->set_ack_seq(0); + response->set_supports_max_allocatable_bytes(config_.advertise_max_allocatable_bytes); + response->set_supports_logical_tiers(config_.advertise_logical_tiers); return grpc::Status::OK; }); } @@ -289,8 +318,12 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser TierCapacity c; c.total_bytes = tc.total_capacity_bytes(); c.available_bytes = tc.available_capacity_bytes(); + c.max_allocatable_bytes = tc.max_allocatable_bytes(); caps[static_cast(tc.tier())] = c; } + auto logical_caps = config_.advertise_logical_tiers + ? LogicalCapsFromProto(request->logical_tier_capacities()) + : std::map{}; std::vector bundles; bundles.reserve(static_cast(request->bundles_size())); @@ -298,7 +331,11 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser EventBundle bundle; bundle.seq = pb.seq(); bundle.events.reserve(static_cast(pb.events_size())); - for (const auto& pe : pb.events()) bundle.events.push_back(FromProtoEvent(pe)); + for (const auto& pe : pb.events()) { + auto event = FromProtoEvent(pe); + if (!config_.advertise_logical_tiers) event.logical_tier.clear(); + bundle.events.push_back(std::move(event)); + } bundles.push_back(std::move(bundle)); } @@ -320,7 +357,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser for (const auto& ev : bundle.events) events.push_back(ev); } auto result = store_.ApplyHeartbeat(request->node_id(), request->delta_seq_baseline(), now, - caps, events, /*is_full_sync=*/true); + caps, events, /*is_full_sync=*/true, logical_caps); if (result.status == HeartbeatResult::UNKNOWN) { client_status = ClientStatus::UNKNOWN; } else { @@ -333,7 +370,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser // without advancing last_applied_seq; the gap is ignored (no full-sync // request) because there is nothing to recover. auto result = store_.ApplyHeartbeat(request->node_id(), /*seq=*/0, now, caps, - /*events=*/{}, /*is_full_sync=*/false); + /*events=*/{}, /*is_full_sync=*/false, logical_caps); if (result.status == HeartbeatResult::UNKNOWN) { client_status = ClientStatus::UNKNOWN; } else { @@ -346,7 +383,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser // across the batch). for (const auto& bundle : bundles) { auto result = store_.ApplyHeartbeat(request->node_id(), bundle.seq, now, caps, - bundle.events, /*is_full_sync=*/false); + bundle.events, /*is_full_sync=*/false, logical_caps); if (result.status == HeartbeatResult::UNKNOWN) { client_status = ClientStatus::UNKNOWN; break; @@ -365,6 +402,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->set_request_full_sync(request_full_sync); UpdateClientCapacityMetrics(request->node_id(), caps); + UpdateClientLogicalTierMetrics(request->node_id(), logical_caps); if (metrics_ != nullptr && request->tier_kv_counts_size() > 0) { mori::metrics::MetricsServer::Labels base = {{"node", request->node_id()}}; @@ -428,6 +466,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->set_node_id(result->node_id); response->set_tier(static_cast<::umbp::TierType>(result->tier)); response->set_peer_address(result->peer_address); + response->set_logical_tier(result->logical_tier); if (metrics_) { metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_PUT, @@ -455,6 +494,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->set_tier(static_cast<::umbp::TierType>(result->location.tier)); response->set_size(result->location.size); response->set_peer_address(result->peer_address); + response->set_logical_tier(result->location.logical_tier); if (metrics_) { metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_GET, @@ -491,6 +531,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser entry->set_node_id(opt->node_id); entry->set_tier(static_cast<::umbp::TierType>(opt->tier)); entry->set_peer_address(opt->peer_address); + entry->set_logical_tier(opt->logical_tier); if (metrics_) { metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_PUT, MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_PUT_HELP, @@ -516,6 +557,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->mutable_node_ref()->Reserve(static_cast(results.size())); response->mutable_tier()->Reserve(static_cast(results.size())); response->mutable_size()->Reserve(static_cast(results.size())); + response->mutable_logical_tier()->Reserve(static_cast(results.size())); // Maps "node_id\0peer_address" -> 1-based index into response->nodes(). std::unordered_map node_index; for (auto& opt : results) { @@ -523,6 +565,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->add_node_ref(0); response->add_tier(::umbp::TIER_UNKNOWN); response->add_size(0); + response->add_logical_tier(""); continue; } std::string node_key = opt->location.node_id; @@ -538,6 +581,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->add_node_ref(it->second); response->add_tier(static_cast<::umbp::TierType>(opt->location.tier)); response->add_size(opt->location.size); + response->add_logical_tier(opt->location.logical_tier); if (metrics_) { metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET, MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET_HELP, @@ -755,8 +799,10 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser for (const auto& l : s.labels()) labels.push_back({l.name(), l.value()}); switch (s.value_case()) { case ::umbp::MetricSample::kCounterDelta: - metrics_->addCounter(s.name(), s.help(), labels, - static_cast(s.counter_delta())); + // No cast: the wire already carries a double, and flooring it here + // is what silently zeroed every fractional counter a component + // reported (a *_seconds_total never survives a uint64 truncation). + metrics_->addCounter(s.name(), s.help(), labels, s.counter_delta()); break; case ::umbp::MetricSample::kGaugeValue: metrics_->setGauge(s.name(), s.help(), labels, s.gauge_value()); @@ -811,6 +857,41 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser } } + void UpdateClientLogicalTierMetrics( + const std::string& node_id, const std::map& logical_caps) { + if (!metrics_) return; + for (const auto& [name, logical] : logical_caps) { + const auto& cap = logical.capacity; + mori::metrics::MetricsServer::Labels labels = { + {"node", node_id}, + {"logical_tier", name}, + {"tier", TierTypeName(logical.representative_tier)}}; + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_TOTAL, + MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_TOTAL_HELP, labels, + static_cast(cap.total_bytes)); + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_AVAIL, + MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_AVAIL_HELP, labels, + static_cast(cap.available_bytes)); + const uint64_t used_bytes = + cap.total_bytes >= cap.available_bytes ? cap.total_bytes - cap.available_bytes : 0; + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_USED, + MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_USED_HELP, labels, + static_cast(used_bytes)); + const double utilization = cap.total_bytes > 0 ? static_cast(used_bytes) / + static_cast(cap.total_bytes) + : 0.0; + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_UTILIZATION, + MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_UTILIZATION_HELP, labels, + utilization); + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_LOGICAL_PUT_ELIGIBLE, + MORI_UMBP_METRIC_CLIENT_LOGICAL_PUT_ELIGIBLE_HELP, labels, + logical.put_eligible ? 1.0 : 0.0); + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_LOGICAL_PEAK_UTILIZATION, + MORI_UMBP_METRIC_CLIENT_LOGICAL_PEAK_UTILIZATION_HELP, labels, + logical.peak_member_utilization); + } + } + IMasterMetadataStore& store_; Router& router_; ClientRegistryConfig config_; @@ -853,8 +934,7 @@ void MasterServer::Run() { StartHitIndexGc(); grpc::ServerBuilder builder; - builder.SetMaxReceiveMessageSize(64 * 1024 * 1024); - builder.SetMaxSendMessageSize(64 * 1024 * 1024); + ApplyGrpcLimits(&builder); // Size the sync-server poller/handler thread pool to the client fan-out. // The default sync server runs only a couple of poller threads, which also // execute the RPC handlers; under many concurrent clients a burst of large diff --git a/src/umbp/distributed/master/redis_master_metadata_store.cpp b/src/umbp/distributed/master/redis_master_metadata_store.cpp index b41780217..6ed7ebd7b 100644 --- a/src/umbp/distributed/master/redis_master_metadata_store.cpp +++ b/src/umbp/distributed/master/redis_master_metadata_store.cpp @@ -22,6 +22,7 @@ #include "umbp/distributed/master/redis_master_metadata_store.h" #include +#include #include #include #include @@ -82,7 +83,8 @@ std::chrono::system_clock::time_point FromEpochMs(int64_t ms) { return std::chrono::system_clock::time_point(std::chrono::milliseconds(ms)); } -// caps: "tier:total:avail;tier:total:avail;..." +// caps: "tier:total:avail:max_alloc;tier:total:avail:max_alloc;..." +// Decode also accepts the legacy three-field form. std::string EncodeCaps(const std::map& caps) { std::string out; for (const auto& [tier, cap] : caps) { @@ -91,6 +93,8 @@ std::string EncodeCaps(const std::map& caps) { out += std::to_string(cap.total_bytes); out += ':'; out += std::to_string(cap.available_bytes); + out += ':'; + out += std::to_string(cap.max_allocatable_bytes); out += ';'; } return out; @@ -107,12 +111,15 @@ std::map DecodeCaps(const std::string& blob) { if (tok.empty()) continue; const size_t c1 = tok.find(':'); const size_t c2 = (c1 == std::string::npos) ? std::string::npos : tok.find(':', c1 + 1); + const size_t c3 = (c2 == std::string::npos) ? std::string::npos : tok.find(':', c2 + 1); if (c1 == std::string::npos || c2 == std::string::npos) continue; try { const int tier = std::stoi(tok.substr(0, c1)); const uint64_t total = std::stoull(tok.substr(c1 + 1, c2 - c1 - 1)); - const uint64_t avail = std::stoull(tok.substr(c2 + 1)); - caps[static_cast(tier)] = TierCapacity{total, avail}; + const uint64_t avail = + std::stoull(tok.substr(c2 + 1, c3 == std::string::npos ? c3 : c3 - c2 - 1)); + const uint64_t max_alloc = c3 == std::string::npos ? 0 : std::stoull(tok.substr(c3 + 1)); + caps[static_cast(tier)] = TierCapacity{total, avail, max_alloc}; } catch (const std::exception&) { // best-effort decode; skip malformed token } @@ -120,6 +127,66 @@ std::map DecodeCaps(const std::string& blob) { return caps; } +std::string EncodeLogicalCaps(const std::map& caps) { + std::string out; + for (const auto& [name, logical] : caps) { + out += std::to_string(name.size()) + ':' + name + ':'; + out += std::to_string(static_cast(logical.representative_tier)) + ':'; + out += std::to_string(logical.capacity.total_bytes) + ':'; + out += std::to_string(logical.capacity.available_bytes) + ':'; + out += std::to_string(logical.capacity.max_allocatable_bytes) + ':'; + out += logical.put_eligible ? "1:" : "0:"; + out += std::to_string(logical.peak_member_utilization) + ';'; + } + return out; +} + +std::map DecodeLogicalCaps(const std::string& blob) { + std::map caps; + size_t pos = 0; + while (pos < blob.size()) { + try { + const size_t length_end = blob.find(':', pos); + if (length_end == std::string::npos) break; + const size_t name_size = std::stoull(blob.substr(pos, length_end - pos)); + const size_t name_begin = length_end + 1; + if (name_size > blob.size() - name_begin || name_begin + name_size >= blob.size() || + blob[name_begin + name_size] != ':') { + break; + } + const std::string name = blob.substr(name_begin, name_size); + pos = name_begin + name_size + 1; + const size_t record_end = blob.find(';', pos); + if (record_end == std::string::npos) { + throw std::invalid_argument("logical caps"); + } + std::vector fields; + while (pos < record_end) { + const size_t field_end = std::min(blob.find(':', pos), record_end); + fields.push_back(blob.substr(pos, field_end - pos)); + pos = field_end + 1; + } + if (fields.size() != 5 && fields.size() != 6) { + throw std::invalid_argument("logical caps"); + } + pos = record_end + 1; + LogicalTierCapacity logical; + logical.representative_tier = static_cast(std::stoi(fields[0])); + logical.capacity.total_bytes = std::stoull(fields[1]); + logical.capacity.available_bytes = std::stoull(fields[2]); + logical.capacity.max_allocatable_bytes = std::stoull(fields[3]); + logical.put_eligible = fields[4] == "1"; + if (fields.size() == 6) { + logical.peak_member_utilization = std::stod(fields[5]); + } + caps[name] = logical; + } catch (const std::exception&) { + break; + } + } + return caps; +} + std::string JoinTags(const std::vector& tags) { std::string out; for (size_t i = 0; i < tags.size(); ++i) { @@ -204,6 +271,9 @@ ClientRecord DecodeRecord(const std::string& node_id, } } if (auto* v = get("caps")) rec.tier_capacities = DecodeCaps(*v); + if (auto* v = get("lcaps")) { + rec.logical_tier_capacities = DecodeLogicalCaps(*v); + } if (auto* v = get("engine")) rec.engine_desc_bytes.assign(v->begin(), v->end()); if (auto* v = get("tags")) rec.tags = SplitTags(*v); return rec; @@ -312,7 +382,7 @@ void RedisMasterMetadataStore::ApplyBlockEventsMulti(const std::string& node_id, if (!is_full_sync && shard_events.empty()) continue; std::vector args; - args.reserve(5 + shard_events.size() * 4); + args.reserve(5 + shard_events.size() * 5); args.push_back(keys_.NodeBlocks(node_id, shard)); args.push_back(node_prefix); args.push_back(is_full_sync ? "1" : "0"); @@ -323,6 +393,7 @@ void RedisMasterMetadataStore::ApplyBlockEventsMulti(const std::string& node_id, args.push_back(keys_.Block(ev->key)); args.push_back(std::to_string(static_cast(ev->tier))); args.push_back(std::to_string(ev->size)); + args.push_back(ev->logical_tier); } // KEYS[1] = this shard's reverse-index key (a shard-tag key) so the cluster // client can route to the shard's slot; the script reads all its keys from @@ -392,7 +463,8 @@ bool RedisMasterMetadataStore::RegisterClient(const ClientRegistration& registra redis::kRegisterClientLua, {keys_.Node(registration.node_id)}, {keys_.Tag(), registration.node_id, std::to_string(ToEpochMs(now)), std::to_string(ToMs(stale_after)), registration.node_address, registration.peer_address, - EncodeCaps(registration.tier_capacities), engine, JoinTags(registration.tags)}); + EncodeCaps(registration.tier_capacities), engine, JoinTags(registration.tags), + EncodeLogicalCaps(registration.logical_tier_capacities)}); if (r.is_error()) throw std::runtime_error("[RedisStore] RegisterClient: " + r.str); return r.integer == 1; } @@ -424,19 +496,20 @@ void RedisMasterMetadataStore::UnregisterClient(const std::string& node_id) { HeartbeatResult RedisMasterMetadataStore::ApplyHeartbeat( const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, const std::map& caps, const std::vector& events, - bool is_full_sync) { + bool is_full_sync, const std::map& logical_caps) { ScopedStoreOp _op(metrics_, "ApplyHeartbeat"); RespValue r; if (!split_writes()) { // Single instance: one atomic script does seq-CAS + record + blocks. std::vector args; - args.reserve(7 + events.size() * 4); + args.reserve(8 + events.size() * 5); args.push_back(keys_.Tag()); args.push_back(node_id); args.push_back(std::to_string(seq)); args.push_back(std::to_string(ToEpochMs(now))); args.push_back(is_full_sync ? "1" : "0"); args.push_back(EncodeCaps(caps)); + args.push_back(EncodeLogicalCaps(logical_caps)); args.push_back(std::to_string(events.size())); for (const auto& ev : events) { args.push_back(ev.kind == KvEvent::Kind::ADD ? "0" : "1"); @@ -445,14 +518,16 @@ HeartbeatResult RedisMasterMetadataStore::ApplyHeartbeat( args.push_back(keys_.Block(ev.key)); args.push_back(std::to_string(static_cast(ev.tier))); args.push_back(std::to_string(ev.size)); + args.push_back(ev.logical_tier); } r = control().Eval(redis::kApplyHeartbeatLua, {keys_.Node(node_id)}, args); } else { // Multi-endpoint: control step (seq-CAS + record + alive/peers) only; block // events are applied per shard afterwards if this heartbeat is APPLIED. - r = control().Eval(redis::kApplyHeartbeatControlLua, {keys_.Node(node_id)}, - {keys_.Tag(), node_id, std::to_string(seq), std::to_string(ToEpochMs(now)), - is_full_sync ? "1" : "0", EncodeCaps(caps)}); + r = control().Eval( + redis::kApplyHeartbeatControlLua, {keys_.Node(node_id)}, + {keys_.Tag(), node_id, std::to_string(seq), std::to_string(ToEpochMs(now)), + is_full_sync ? "1" : "0", EncodeCaps(caps), EncodeLogicalCaps(logical_caps)}); } if (r.is_error()) throw std::runtime_error("[RedisStore] ApplyHeartbeat: " + r.str); @@ -660,8 +735,14 @@ std::vector RedisMasterMetadataStore::LookupBlock(const std::string& k if (sep == std::string::npos) continue; Location loc; loc.node_id = rest.substr(0, sep); + const std::string tier_and_logical = rest.substr(sep + 1); + const size_t logical_sep = tier_and_logical.find('|'); + const std::string tier_text = tier_and_logical.substr(0, logical_sep); + if (logical_sep != std::string::npos) { + loc.logical_tier = tier_and_logical.substr(logical_sep + 1); + } try { - loc.tier = static_cast(std::stoi(rest.substr(sep + 1))); + loc.tier = static_cast(std::stoi(tier_text)); loc.size = std::stoull(r.elements[i + 1].str); } catch (...) { continue; @@ -702,12 +783,13 @@ std::vector> RedisMasterMetadataStore::BatchLookupBlockFor [&](size_t group) { return &client_for_shard(batch.shard_of_group[group]); }, [&](size_t orig_index, const RespValue& locs) { if (!locs.is_array()) return; - for (size_t j = 0; j + 3 <= locs.elements.size(); j += 3) { + for (size_t j = 0; j + 4 <= locs.elements.size(); j += 4) { Location loc; loc.node_id = locs.elements[j].str; try { loc.size = std::stoull(locs.elements[j + 1].str); loc.tier = static_cast(std::stoi(locs.elements[j + 2].str)); + loc.logical_tier = locs.elements[j + 3].str; } catch (...) { continue; } diff --git a/src/umbp/distributed/metrics/component_metrics.cpp b/src/umbp/distributed/metrics/component_metrics.cpp new file mode 100644 index 000000000..6404d3e28 --- /dev/null +++ b/src/umbp/distributed/metrics/component_metrics.cpp @@ -0,0 +1,88 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/metrics/component_metrics.h" + +namespace mori::umbp { +namespace { + +// Baseline key = source + metric + label set. '\0' separates the parts so no +// combination of label text can forge another key's identity. +std::string BaselineKey(const std::string& source_id, const char* name, + const MetricLabels& source_labels, const MetricLabels& sample_labels) { + std::string id = source_id; + id += '\0'; + id += name; + auto append = [&id](const MetricLabels& labels) { + for (const auto& [k, v] : labels) { + id += '\0'; + id += k; + id += '='; + id += v; + } + }; + append(source_labels); + append(sample_labels); + return id; +} + +MetricLabels MergeLabels(const MetricLabels& source_labels, const MetricLabels& sample_labels) { + MetricLabels merged; + merged.reserve(source_labels.size() + sample_labels.size()); + merged.insert(merged.end(), source_labels.begin(), source_labels.end()); + merged.insert(merged.end(), sample_labels.begin(), sample_labels.end()); + return merged; +} + +} // namespace + +void MetricPublisher::Publish(const std::string& source_id, const MetricLabels& source_labels, + const MetricSource& source, const Sink& sink) { + Publish(source_id, source_labels, source.SampleMetrics(), sink); +} + +void MetricPublisher::Publish(const std::string& source_id, const MetricLabels& source_labels, + const std::vector& samples, const Sink& sink) { + for (const MetricSample& s : samples) { + if (s.name == nullptr) continue; + + if (s.kind == MetricKind::kGauge) { + // No baseline: a gauge is whatever it reads now. Shipped every tick so a + // value that stops changing still keeps its series alive. + if (sink.gauge) { + sink.gauge(s.name, s.help, MergeLabels(source_labels, s.labels), + static_cast(s.value) * s.scale); + } + continue; + } + + uint64_t& last = last_[BaselineKey(source_id, s.name, source_labels, s.labels)]; + if (s.value > last && sink.counter) { + sink.counter(s.name, s.help, MergeLabels(source_labels, s.labels), + static_cast(s.value - last) * s.scale); + } + // Updated even on a zero or negative delta so the next tick stays correct + // against a counter that was rebuilt underneath us. + last = s.value; + } +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/backend/hbm_backend.cpp b/src/umbp/distributed/peer/backend/hbm_backend.cpp new file mode 100644 index 000000000..f3b65e8f7 --- /dev/null +++ b/src/umbp/distributed/peer/backend/hbm_backend.cpp @@ -0,0 +1,116 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/peer/backend/hbm_backend.h" + +#include + +#include + +#include "mori/utils/mori_log.hpp" + +namespace mori::umbp { + +namespace { + +// Set the current device for the duration of a scope and put it back. Peer +// service threads do not own a current device (see the header), so every HIP +// call that depends on one has to establish it explicitly and leave no trace. +class ScopedDevice { + public: + explicit ScopedDevice(int device) { + if (hipGetDevice(&previous_) != hipSuccess) { + previous_ = -1; + return; + } + if (device == previous_) return; // nothing to restore + if (hipSetDevice(device) != hipSuccess) { + previous_ = -1; + return; + } + restore_ = true; + } + ~ScopedDevice() { + if (restore_ && previous_ >= 0) (void)hipSetDevice(previous_); + } + + ScopedDevice(const ScopedDevice&) = delete; + ScopedDevice& operator=(const ScopedDevice&) = delete; + + private: + int previous_ = -1; + bool restore_ = false; +}; + +} // namespace + +bool HbmPageMemorySource::Allocate(const std::vector& sizes, std::vector* out) { + ScopedDevice scope(device_); + + std::vector taken; + std::vector staged; + for (uint64_t size : sizes) { + if (size == 0) continue; + void* ptr = nullptr; + const hipError_t err = hipMalloc(&ptr, static_cast(size)); + if (err != hipSuccess || ptr == nullptr) { + MORI_UMBP_ERROR("[HbmPageMemorySource] hipMalloc({}) on device {} failed: {}", size, device_, + hipGetErrorString(err)); + // All-or-nothing, matching the PageMemorySource contract: unwind only + // what this call took, leaving `out` and any earlier pool untouched. + for (void* p : taken) (void)hipFree(p); + return false; + } + // hipMalloc returns exactly the requested size — no rounding to report, in + // contrast to the host source's hugepage mapped_size. + staged.push_back(Buffer{ptr, size}); + taken.push_back(ptr); + } + + ptrs_.insert(ptrs_.end(), taken.begin(), taken.end()); + out->insert(out->end(), staged.begin(), staged.end()); + MORI_UMBP_INFO("[HbmPageMemorySource] allocated {} buffers on device {}", staged.size(), device_); + return true; +} + +void HbmPageMemorySource::Release() { + if (ptrs_.empty()) return; + ScopedDevice scope(device_); + for (void* p : ptrs_) { + const hipError_t err = hipFree(p); + if (err != hipSuccess) { + MORI_UMBP_ERROR("[HbmPageMemorySource] hipFree on device {} failed: {}", device_, + hipGetErrorString(err)); + } + } + ptrs_.clear(); +} + +std::unique_ptr MakeHbmBackend(uint64_t page_size, int device, + std::vector buffer_sizes, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl) { + return std::make_unique(TierType::HBM, page_size, + std::make_unique(device), + std::move(buffer_sizes), pending_ttl, read_lease_ttl); +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/backend/instrumented_backend.cpp b/src/umbp/distributed/peer/backend/instrumented_backend.cpp new file mode 100644 index 000000000..8f0f6db47 --- /dev/null +++ b/src/umbp/distributed/peer/backend/instrumented_backend.cpp @@ -0,0 +1,263 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/peer/backend/instrumented_backend.h" + +#include +#include + +namespace mori::umbp { +namespace { + +constexpr double kNanosPerSecond = 1e9; + +// Label text for the enum indices. These strings ARE the dashboard's contract: +// a panel groups by op= and status=, so changing one renames a series. +const char* OpName(int op) { + switch (op) { + case 0: + return "allocate"; + case 1: + return "commit"; + case 2: + return "abort"; + case 3: + return "resolve"; + case 4: + return "evict"; + default: + return "unknown"; + } +} + +const char* StatusName(int status) { + switch (status) { + case 0: + return "ok"; + case 1: + return "exists"; + case 2: + return "no_space"; + case 3: + return "miss"; + case 4: + return "failed"; + default: + return "unknown"; + } +} + +// Scoped timer that charges elapsed nanoseconds to an op on destruction, so an +// early return inside a forwarded call still gets measured. +class ScopedNanos { + public: + explicit ScopedNanos(std::atomic& sink) + : sink_(sink), start_(std::chrono::steady_clock::now()) {} + ~ScopedNanos() { + const auto elapsed = std::chrono::steady_clock::now() - start_; + sink_.fetch_add(static_cast( + std::chrono::duration_cast(elapsed).count()), + std::memory_order_relaxed); + } + + private: + std::atomic& sink_; + std::chrono::steady_clock::time_point start_; +}; + +} // namespace + +InstrumentedBackend::InstrumentedBackend(std::unique_ptr inner) + : inner_(std::move(inner)) { + for (auto& op : entries_) { + for (auto& status : op) status.store(0, std::memory_order_relaxed); + } +} + +bool InstrumentedBackend::Init(MemoryRegistrar* registrar) { return inner_->Init(registrar); } + +std::vector InstrumentedBackend::BatchAllocate( + const std::vector& reqs) { + ops_[kAllocate].batches.fetch_add(1, std::memory_order_relaxed); + std::vector out; + { + ScopedNanos timer(ops_[kAllocate].nanos); + out = inner_->BatchAllocate(reqs); + } + // Outcome mapping is the medium-agnostic half of AllocateOutcome: kFailedNoSpace + // stays distinct because it is the one failure a writer retries elsewhere, and + // "is this pool full" is the first question a capacity panel has to answer. + for (const AllocateResult& r : out) { + switch (r.outcome) { + case AllocateOutcome::kSuccessAllocated: + Add(kAllocate, kOk, 1); + break; + case AllocateOutcome::kSuccessAlreadyExists: + Add(kAllocate, kExists, 1); + break; + case AllocateOutcome::kFailedNoSpace: + Add(kAllocate, kNoSpace, 1); + break; + case AllocateOutcome::kFailed: + Add(kAllocate, kFailed, 1); + break; + } + } + return out; +} + +std::vector InstrumentedBackend::BatchCommit(const std::vector& reqs) { + ops_[kCommit].batches.fetch_add(1, std::memory_order_relaxed); + std::vector out; + { + ScopedNanos timer(ops_[kCommit].nanos); + out = inner_->BatchCommit(reqs); + } + uint64_t bytes = 0; + for (const CommitResult& r : out) { + if (r.success) { + Add(kCommit, kOk, 1); + bytes += r.bytes_committed; + } else { + // A slot that was reaped, aborted, or never existed. Not "medium broken" + // but the writer's bytes are gone either way, so it counts as failed. + Add(kCommit, kFailed, 1); + } + } + if (bytes > 0) ops_[kCommit].bytes.fetch_add(bytes, std::memory_order_relaxed); + return out; +} + +std::vector InstrumentedBackend::BatchAbort(const std::vector& slot_ids) { + ops_[kAbort].batches.fetch_add(1, std::memory_order_relaxed); + std::vector out; + { + ScopedNanos timer(ops_[kAbort].nanos); + out = inner_->BatchAbort(slot_ids); + } + for (bool ok : out) Add(kAbort, ok ? kOk : kFailed, 1); + return out; +} + +std::vector InstrumentedBackend::BatchResolve(const std::vector& keys, + bool include_descs, + bool allow_file_refs) { + ops_[kResolve].batches.fetch_add(1, std::memory_order_relaxed); + std::vector out; + { + ScopedNanos timer(ops_[kResolve].nanos); + out = inner_->BatchResolve(keys, include_descs, allow_file_refs); + } + uint64_t bytes = 0; + for (const ResolvedEntry& r : out) { + if (r.found) { + Add(kResolve, kOk, 1); + bytes += r.size; + } else { + Add(kResolve, kMiss, 1); + } + } + if (bytes > 0) ops_[kResolve].bytes.fetch_add(bytes, std::memory_order_relaxed); + return out; +} + +std::vector InstrumentedBackend::Evict(const std::vector& keys) { + ops_[kEvict].batches.fetch_add(1, std::memory_order_relaxed); + std::vector out; + { + ScopedNanos timer(ops_[kEvict].nanos); + out = inner_->Evict(keys); + } + uint64_t bytes = 0; + for (const EvictResult& r : out) { + // bytes_freed == 0 means the key was unknown, already gone, or protected by + // a live lease. That is a miss, not a failure: master retries it later. + if (r.bytes_freed > 0) { + Add(kEvict, kOk, 1); + bytes += r.bytes_freed; + } else { + Add(kEvict, kMiss, 1); + } + } + if (bytes > 0) ops_[kEvict].bytes.fetch_add(bytes, std::memory_order_relaxed); + return out; +} + +std::vector InstrumentedBackend::SampleMetrics() const { + std::vector out; + out.reserve(kOpCount * 4 + 8); + + for (int op = 0; op < kOpCount; ++op) { + const char* op_name = OpName(op); + + for (int status = 0; status < kStatusCount; ++status) { + const uint64_t v = entries_[op][status].load(std::memory_order_relaxed); + // Skip combinations this op never produces (an abort has no "no_space") + // so the series set stays the shape of what actually happens. + if (v == 0) continue; + out.push_back(MetricSample{MORI_UMBP_METRIC_BACKEND_OPS_TOTAL, + MORI_UMBP_METRIC_BACKEND_OPS_TOTAL_HELP, + {{"op", op_name}, {"status", StatusName(status)}}, + v}); + } + + const uint64_t batches = ops_[op].batches.load(std::memory_order_relaxed); + if (batches == 0) continue; + out.push_back(MetricSample{MORI_UMBP_METRIC_BACKEND_BATCHES_TOTAL, + MORI_UMBP_METRIC_BACKEND_BATCHES_TOTAL_HELP, + {{"op", op_name}}, + batches}); + + // Accumulated in nanoseconds because that is what an atomic add can carry; + // MetricSample::scale converts once, in the publisher, so the metric keeps + // the seconds its name promises. + const uint64_t nanos = ops_[op].nanos.load(std::memory_order_relaxed); + if (nanos > 0) { + out.push_back(MetricSample{MORI_UMBP_METRIC_BACKEND_OP_SECONDS_TOTAL, + MORI_UMBP_METRIC_BACKEND_OP_SECONDS_TOTAL_HELP, + {{"op", op_name}}, + nanos, + MetricKind::kCounter, + 1.0 / kNanosPerSecond}); + } + + const uint64_t bytes = ops_[op].bytes.load(std::memory_order_relaxed); + if (bytes > 0) { + out.push_back(MetricSample{MORI_UMBP_METRIC_BACKEND_BYTES_TOTAL, + MORI_UMBP_METRIC_BACKEND_BYTES_TOTAL_HELP, + {{"op", op_name}}, + bytes}); + } + } + + // The medium's own view, appended unchanged. The publisher stamps tier= and + // backend= on both halves, so a medium counter and a decorator counter are + // indistinguishable to a dashboard — which is the point. + for (MetricSample& s : inner_->SampleMetrics()) out.push_back(std::move(s)); + return out; +} + +std::unique_ptr MakeInstrumentedBackend(std::unique_ptr inner) { + if (inner == nullptr) return nullptr; + return std::make_unique(std::move(inner)); +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/backend/page_backend.cpp b/src/umbp/distributed/peer/backend/page_backend.cpp new file mode 100644 index 000000000..7ade4ab15 --- /dev/null +++ b/src/umbp/distributed/peer/backend/page_backend.cpp @@ -0,0 +1,1060 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/peer/backend/page_backend.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/distributed/transfer/transfer_engine.h" + +namespace mori::umbp { + +// --------------------------------------------------------------------------- +// File-local helpers +// --------------------------------------------------------------------------- + +namespace { + +// Round size up to whole pages of `page_size`. Returns 0 only when +// size == 0; the allocator treats num_pages == 0 as ENOSPC anyway. +uint32_t SizeToPages(uint64_t size, uint64_t page_size) { + if (page_size == 0 || size == 0) return 0; + uint64_t pages = (size + page_size - 1) / page_size; + // Cap at uint32_t — PageBitmapAllocator uses uint32_t for page counts. + if (pages > std::numeric_limits::max()) { + return 0; + } + return static_cast(pages); +} + +} // namespace + +// --------------------------------------------------------------------------- +// Construction +// --------------------------------------------------------------------------- + +PageBackend::PageBackend(TierType tier, uint64_t page_size, TierConfig cfg, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl, + std::chrono::milliseconds reaper_interval) + : tier_(tier), + page_size_(page_size), + pending_ttl_(pending_ttl), + read_lease_ttl_(read_lease_ttl), + reaper_interval_(reaper_interval) { + if (page_size == 0) { + throw std::invalid_argument("PageBackend: page_size must be > 0"); + } + InstallTierConfig(std::move(cfg)); +} + +// Sugar: build the host source these knobs describe and delegate, so there is +// one self-allocating constructor body and one Init path. +PageBackend::PageBackend(TierType tier, uint64_t page_size, OwnershipConfig ownership, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl, + std::chrono::milliseconds reaper_interval) + : PageBackend(tier, page_size, + std::make_unique(HostPageMemorySource::Options{ + ownership.use_hugepages, ownership.hugepage_size, ownership.numa_node, + ownership.prefault}), + ownership.buffer_sizes, pending_ttl, read_lease_ttl, reaper_interval) {} + +PageBackend::PageBackend(TierType tier, uint64_t page_size, + std::unique_ptr source, + std::vector buffer_sizes, std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl, + std::chrono::milliseconds reaper_interval) + : tier_(tier), + page_size_(page_size), + pending_ttl_(pending_ttl), + read_lease_ttl_(read_lease_ttl), + reaper_interval_(reaper_interval), + source_(std::move(source)), + source_buffer_sizes_(std::move(buffer_sizes)) { + if (page_size == 0) { + throw std::invalid_argument("PageBackend: page_size must be > 0"); + } + if (source_ == nullptr) { + throw std::invalid_argument("PageBackend: PageMemorySource must not be null"); + } + // Captured now rather than read through source_ at BuildBufferRefs time: that + // runs for the TierConfig constructor too, where there is no source. + buffer_loc_ = source_->LocationType(); + buffer_device_ = source_->Device(); + // Not yet configured — allocator_ stays null until Init() self-allocates. +} + +PageBackend::~PageBackend() { + StopReaper(); + Shutdown(); +} + +void PageBackend::InstallTierConfig(TierConfig cfg) { + if (cfg.buffer_sizes.empty() && cfg.buffer_descs.empty()) { + allocator_ = std::make_unique(page_size_, std::vector{}); + return; + } + if (cfg.buffer_sizes.size() != cfg.buffer_descs.size()) { + throw std::invalid_argument("PageBackend: buffer_sizes / buffer_descs length mismatch"); + } + // buffer_bases is optional (deployments / tests that never pin leave it + // empty); when present it must line up with the buffers one-to-one. + if (!cfg.buffer_bases.empty() && cfg.buffer_bases.size() != cfg.buffer_sizes.size()) { + throw std::invalid_argument("PageBackend: buffer_bases / buffer_sizes length mismatch"); + } + allocator_ = std::make_unique(page_size_, cfg.buffer_sizes); + buffer_descs_ = std::move(cfg.buffer_descs); + buffer_bases_ = std::move(cfg.buffer_bases); + BuildBufferRefs(); +} + +// Materialize the per-buffer transfer endpoints once, so BufferRef() on the +// local hot path is an index rather than an msgpack unpack. The published size +// is the ALLOCATABLE size (total_pages * page_size), not the mapped size: it is +// the bound a same-process copy is checked against, and it is the tighter of +// the two. +void PageBackend::BuildBufferRefs() { + buffer_refs_.clear(); + if (!allocator_) return; + const size_t buffer_count = allocator_->NumBuffers(); + buffer_refs_.resize(buffer_count); + for (size_t i = 0; i < buffer_count; ++i) { + TransferRef ref; + ref.size = static_cast(allocator_->BufferPageCount(i)) * page_size_; + // From the source that allocated these bytes, not assumed. This is what + // makes a local transfer against this backend select the right engine: + // HbmCopyEngine claims the pair on loc == GPU, LocalCopyEngine on both CPU. + ref.loc = buffer_loc_; + ref.device = buffer_device_; + if (i < buffer_bases_.size()) ref.host_ptr = buffer_bases_[i]; + if (i < buffer_descs_.size() && !buffer_descs_[i].empty()) { + try { + auto handle = msgpack::unpack(reinterpret_cast(buffer_descs_[i].data()), + buffer_descs_[i].size()); + ref.mem = handle.get().as(); + } catch (const std::exception& e) { + MORI_UMBP_ERROR("[PageBackend] BuildBufferRefs: bad desc for buffer {}: {}", i, e.what()); + } + } + buffer_refs_[i] = std::move(ref); + } +} + +// --------------------------------------------------------------------------- +// MediumBackend — identity + ownership +// --------------------------------------------------------------------------- + +bool PageBackend::Init(MemoryRegistrar* registrar) { + if (source_ == nullptr) { + // Constructed with a pre-built TierConfig (test / legacy-direct path) — + // already configured, nothing to do. + return true; + } + if (owns_memory_) return true; // idempotent + + registrar_ = registrar; + + // The medium-specific half: what kind of memory, and where. Everything after + // this call is identical for DRAM, HBM, or any future paged medium. + std::vector buffers; + if (!source_->Allocate(source_buffer_sizes_, &buffers)) { + MORI_UMBP_ERROR("[PageBackend] Init: {} allocation failed for tier={}", source_->Name(), + static_cast(tier_)); + registrar_ = nullptr; + return false; + } + + const mori::io::MemoryLocationType loc = source_->LocationType(); + const int device = source_->Device(); + + TierConfig cfg; + cfg.buffer_sizes.reserve(buffers.size()); + cfg.buffer_descs.reserve(buffers.size()); + cfg.buffer_bases.reserve(buffers.size()); + + for (const auto& buffer : buffers) { + std::vector desc_bytes; + if (registrar != nullptr) { + // The backend supplies the facts a descriptor cannot recover — location + // type and device — because the backend is the allocator (design doc §6). + // They come from the source, which is the thing that actually knows. + TransferRef ref = registrar->RegisterMemory(buffer.base, buffer.size, loc, device); + owned_refs_.push_back(ref); + if (ref.HasMemoryDesc()) { + msgpack::sbuffer sbuf; + msgpack::pack(sbuf, ref.mem); + desc_bytes.assign(sbuf.data(), sbuf.data() + sbuf.size()); + } + } + + cfg.buffer_sizes.push_back(buffer.size); + cfg.buffer_descs.push_back(std::move(desc_bytes)); + cfg.buffer_bases.push_back(buffer.base); + } + + const size_t buffer_count = cfg.buffer_sizes.size(); + InstallTierConfig(std::move(cfg)); + owns_memory_ = true; + StartReaper(); + MORI_UMBP_INFO("[PageBackend] Init tier={} source={} loc={} device={} buffers={} total_bytes={}", + static_cast(tier_), source_->Name(), static_cast(loc), device, + buffer_count, Capacity().total_bytes); + return true; +} + +void PageBackend::Shutdown() { + if (!owns_memory_) return; + if (registrar_ != nullptr) { + for (const auto& ref : owned_refs_) registrar_->Deregister(ref); + } + owned_refs_.clear(); + + // Deregistration must precede release: the registrar may still hold an MR + // over these pages, and freeing under it is what the ordering here prevents. + if (source_ != nullptr) source_->Release(); + + owns_memory_ = false; + registrar_ = nullptr; +} + +// --------------------------------------------------------------------------- +// HostPageMemorySource — host DRAM, verbatim what Init() used to do inline +// --------------------------------------------------------------------------- + +bool HostPageMemorySource::Allocate(const std::vector& sizes, std::vector* out) { + HostMemAllocator allocator; + HostBufferOptions opts; + opts.backing = + opts_.use_hugepages ? HostBufferBacking::kAnonymousHugetlb : HostBufferBacking::kAnonymous; + opts.hugepage_size = opts_.hugepage_size; + opts.numa_node = opts_.numa_node; + opts.prefault = opts_.prefault; + + std::vector taken; + std::vector staged; + for (uint64_t size : sizes) { + if (size == 0) continue; + HostBufferHandle handle = allocator.Alloc(size, opts); + if (!handle.valid()) { + MORI_UMBP_ERROR("[HostPageMemorySource] host allocation failed for size={}", size); + // All-or-nothing: unwind only what THIS call took, leaving any earlier + // successful Allocate (and `out`) untouched. + for (auto& h : taken) allocator.Free(h); + return false; + } + // mapped_size, not the request: hugepage rounding makes the extra usable. + staged.push_back(Buffer{handle.ptr, handle.mapped_size}); + taken.push_back(handle); + } + + handles_.insert(handles_.end(), taken.begin(), taken.end()); + out->insert(out->end(), staged.begin(), staged.end()); + return true; +} + +void HostPageMemorySource::Release() { + HostMemAllocator allocator; + for (auto& handle : handles_) allocator.Free(handle); + handles_.clear(); +} + +// --------------------------------------------------------------------------- +// Allocation (reserve pages into a pending slot) +// --------------------------------------------------------------------------- + +AllocateResult PageBackend::Allocate(const std::string& key, uint64_t size) { + AllocateResult out; + { + std::unique_lock lock(mutex_); + out = AllocateLocked(key, size); + } + // A full pool is the one state the commit-time round cannot have prevented + // (one object larger than the headroom, or a burst that never reached + // commit). Evicting here does not rescue THIS allocation -- it is not + // retried -- but it stops the pool staying wedged for every put after it. + if (out.outcome == AllocateOutcome::kFailedNoSpace) MaybeEvictToLowWatermark(); + return out; +} + +AllocateResult PageBackend::AllocateLocked(const std::string& key, uint64_t size) { + auto fail = [&](AllocateOutcome outcome, const char* reason) { + MORI_UMBP_WARN("[PageBackend] Allocate reason={} key='{}' size={} tier={}", reason, key, size, + static_cast(tier_)); + AllocateResult r; + r.outcome = outcome; + return r; + }; + + const uint32_t num_pages = SizeToPages(size, page_size_); + if (num_pages == 0) return fail(AllocateOutcome::kFailed, "ZERO_SIZE"); + + // Block new allocations between ClearLocal() and full-sync ack — any + // owned key created here would miss the empty snapshot to master. + if (clear_full_sync_pending_.load(std::memory_order_acquire)) { + return fail(AllocateOutcome::kFailed, "CLEAR_PENDING"); + } + + // owned_ dedup (master-index-lag fallback). pending_ deliberately + // not checked — same-key pending race is absorbed by Commit(). + if (owned_.find(key) != owned_.end()) { + AllocateResult out; + out.outcome = AllocateOutcome::kSuccessAlreadyExists; + return out; + } + + if (!allocator_) return fail(AllocateOutcome::kFailed, "NOT_CONFIGURED"); + auto pages = allocator_->Allocate(num_pages); + if (!pages) return fail(AllocateOutcome::kFailedNoSpace, "NO_SPACE"); + + PendingSlot slot; + slot.slot_id = next_slot_id_.fetch_add(1, std::memory_order_relaxed); + slot.pages = std::move(*pages); + slot.size = size; + slot.deadline = std::chrono::steady_clock::now() + pending_ttl_; + slot.generation = allocator_generation_; + const uint64_t slot_id = slot.slot_id; + const auto slot_pages = slot.pages; + pending_[slot_id] = std::move(slot); + + AllocateResult out; + out.outcome = AllocateOutcome::kSuccessAllocated; + out.slot_id = slot_id; + out.pages = slot_pages; + out.size = size; + out.page_size = page_size_; + out.pending_ttl_ms = PendingTtlMs(); + return out; +} + +std::vector PageBackend::BatchAllocate( + const std::vector& entries) { + std::vector out(entries.size()); + if (entries.empty()) return out; + bool no_space = false; + { + std::unique_lock lock(mutex_); + for (size_t i = 0; i < entries.size(); ++i) { + out[i] = AllocateLocked(entries[i].key, entries[i].size); + if (out[i].outcome == AllocateOutcome::kSuccessAllocated) { + out[i].descs = BuildBufferDescsLocked(out[i].pages); + } + no_space = no_space || out[i].outcome == AllocateOutcome::kFailedNoSpace; + } + } + if (no_space) MaybeEvictToLowWatermark(); + return out; +} + +// --------------------------------------------------------------------------- +// Commit / abort (promote a pending slot to owned, or release it) +// --------------------------------------------------------------------------- + +bool PageBackend::Commit(uint64_t slot_id, const std::string& key, uint64_t& bytes_committed) { + bool ok = false; + { + std::unique_lock lock(mutex_); + ok = CommitLocked(slot_id, key, bytes_committed); + } + // Check-after-write, on the committing thread: there is no eviction thread, + // and a commit is the only event that grows usage. Outside the lock because + // a round takes mutex_ itself. No-op unless local eviction is enabled. + MaybeEvictToLowWatermark(); + return ok; +} + +bool PageBackend::CommitLocked(uint64_t slot_id, const std::string& key, + uint64_t& bytes_committed) { + bytes_committed = 0; + auto it = pending_.find(slot_id); + if (it == pending_.end()) { + MORI_UMBP_WARN("[PageBackend] Commit reason=SLOT_GONE key='{}' slot_id={}", key, slot_id); + return false; + } + + // Pre-clear pending slot: free pages, report Put failure, no ADD. + if (it->second.generation != allocator_generation_) { + MORI_UMBP_WARN("[PageBackend] Commit reason=PRE_CLEAR key='{}' slot_id={}", key, slot_id); + if (allocator_) allocator_->Deallocate(it->second.pages); + pending_.erase(it); + return false; + } + + // Race-window safety net: two writers passed Allocate() before either + // committed. Keep first, drop new pages, idempotent success. + auto existing = owned_.find(key); + if (existing != owned_.end()) { + MORI_UMBP_WARN( + "[PageBackend] duplicate Commit for key='{}' (existing size={}, new size={}) — keeping " + "prior slot", + key, existing->second.size, it->second.size); + if (allocator_) allocator_->Deallocate(it->second.pages); + bytes_committed = existing->second.size; + pending_.erase(it); + return true; + } + + OwnedSlot owned; + owned.pages = std::move(it->second.pages); + owned.size = it->second.size; + QueueEventLocked(KvEvent{KvEvent::Kind::ADD, key, tier_, owned.size}); + owned_[key] = std::move(owned); + TouchLru(key, owned_[key]); + pending_.erase(it); + bytes_committed = owned_[key].size; + return true; +} + +std::vector PageBackend::BatchCommit(const std::vector& entries) { + std::vector out(entries.size()); + if (entries.empty()) return out; + { + std::unique_lock lock(mutex_); + for (size_t i = 0; i < entries.size(); ++i) { + out[i].success = CommitLocked(entries[i].slot_id, entries[i].key, out[i].bytes_committed); + } + } + // One round for the whole batch, not one per key: a batch is a single burst + // of growth and a round already frees down to the low watermark. + MaybeEvictToLowWatermark(); + return out; +} + +bool PageBackend::Abort(uint64_t slot_id) { + std::unique_lock lock(mutex_); + return AbortLocked(slot_id); +} + +bool PageBackend::AbortLocked(uint64_t slot_id) { + auto it = pending_.find(slot_id); + if (it == pending_.end()) return true; // already reaped / aborted — idempotent + if (allocator_) allocator_->Deallocate(it->second.pages); + pending_.erase(it); + return true; +} + +std::vector PageBackend::BatchAbort(const std::vector& slot_ids) { + std::vector out(slot_ids.size(), false); + if (slot_ids.empty()) return out; + std::unique_lock lock(mutex_); + for (size_t i = 0; i < slot_ids.size(); ++i) out[i] = AbortLocked(slot_ids[i]); + return out; +} + +// --------------------------------------------------------------------------- +// Resolve (read path; grants a short read lease to fence eviction) +// --------------------------------------------------------------------------- + +ResolvedEntry PageBackend::Resolve(const std::string& key) { + // Shared: this reads owned_ only. The lease is an atomic and the LRU has + // its own lock, so neither needs exclusion here. + std::shared_lock lock(mutex_); + ResolvedEntry r; + auto it = owned_.find(key); + if (it == owned_.end()) return r; + r.outcome = ResolveOutcome::kFound; + r.found = true; + r.pages = it->second.pages; + r.size = it->second.size; + r.page_size = page_size_; + // Extend the read lease so concurrent Evict reports bytes_freed=0 for + // this key. CAS-max rather than a plain store: with a shared lock two + // readers can renew the same key at once, and neither may shorten it. + RenewLeaseAtomic(it->second, SteadyNs(std::chrono::steady_clock::now() + read_lease_ttl_)); + TouchLru(key, it->second); + return r; +} + +std::vector PageBackend::BatchResolve(const std::vector& keys, + bool include_descs, bool allow_file_refs) { + (void)allow_file_refs; // pages are this medium's storage; there is no file to name + std::vector out(keys.size()); + if (keys.empty()) return out; + // Shared: the whole batch -- hash lookups, page-vector copies and descriptor + // building -- is read-only against owned_. This is the concurrency that the + // gated exists sweep measures, so it must not serialize. + std::shared_lock lock(mutex_); + // One now() for the batch: it runs in well under a millisecond against a + // TTL measured in seconds, so the last key loses nothing worth a syscall. + const int64_t lease_deadline = SteadyNs(std::chrono::steady_clock::now() + read_lease_ttl_); + // Collect the hits, then touch the LRU once for the whole batch instead of + // taking lru_mutex_ per key. Nothing here can invalidate the pointers: a + // shared lock excludes every writer of owned_. + std::vector> touched; + if (local_evict_enabled_) touched.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + auto it = owned_.find(keys[i]); + if (it == owned_.end()) continue; + auto& entry = out[i]; + entry.outcome = ResolveOutcome::kFound; + entry.found = true; + entry.pages = it->second.pages; + entry.size = it->second.size; + entry.page_size = page_size_; + if (include_descs) entry.descs = BuildBufferDescsLocked(it->second.pages); + RenewLeaseAtomic(it->second, lease_deadline); + if (local_evict_enabled_) touched.emplace_back(&keys[i], &it->second); + } + if (!touched.empty()) { + std::lock_guard lru_lock(lru_mutex_); + for (auto& [key, slot] : touched) TouchLruLocked(*key, *slot); + } + return out; +} + +bool PageBackend::AcquireMigrationRead(const std::string& key, ResolvedEntry* resolved) { + if (resolved == nullptr) return false; + std::unique_lock lock(mutex_); + auto owned = owned_.find(key); + if (owned == owned_.end() || pins_.find(key) != pins_.end()) return false; + resolved->outcome = ResolveOutcome::kFound; + resolved->found = true; + resolved->pages = owned->second.pages; + resolved->size = owned->second.size; + resolved->page_size = page_size_; + pins_[key] = PinState{next_pin_token_++, std::chrono::steady_clock::now()}; + return true; +} + +void PageBackend::ReleaseMigrationRead(const std::string& key) { + std::unique_lock lock(mutex_); + pins_.erase(key); +} + +bool PageBackend::Contains(const std::string& key) const { + std::shared_lock lock(mutex_); + return owned_.find(key) != owned_.end(); +} + +std::vector PageBackend::BatchContains(const std::vector& keys) const { + std::vector out(keys.size(), false); + if (keys.empty()) return out; + // One lock and one hash lookup per key, and nothing allocated per key: an + // existence answer needs no page vector, no lease renewal and no LRU touch. + std::shared_lock lock(mutex_); + for (size_t i = 0; i < keys.size(); ++i) out[i] = owned_.find(keys[i]) != owned_.end(); + return out; +} + +// --------------------------------------------------------------------------- +// Eviction (skips leased / copy-pinned keys; emits REMOVE) +// --------------------------------------------------------------------------- + +std::vector PageBackend::Evict(const std::vector& keys) { + std::vector out; + out.reserve(keys.size()); + std::unique_lock lock(mutex_); + for (const auto& key : keys) { + EvictResult r; + r.key = key; + auto it = owned_.find(key); + if (it == owned_.end()) { + out.push_back(std::move(r)); + continue; + } + if (HasActiveReadLeaseLocked(it->second, std::chrono::steady_clock::now())) { + // Master will retry next round once the lease expires. Emit no event. + out.push_back(std::move(r)); + continue; + } + if (HasActivePinLocked(key)) { + // An SSD copy worker is reading these pages. Do NOT free them, do + // NOT emit REMOVE, keep the key owned. bytes_freed=0 tells master to + // retry; the pin is released when the copy finishes. + out.push_back(std::move(r)); + continue; + } + if (allocator_) allocator_->Deallocate(it->second.pages); + r.bytes_freed = it->second.size; + QueueEventLocked(KvEvent{KvEvent::Kind::REMOVE, key, tier_, 0}); + RemoveFromLru(it->second); + owned_.erase(it); + out.push_back(std::move(r)); + } + return out; +} + +// --------------------------------------------------------------------------- +// DRAM copy pins (protect owned pages while the SSD copy pipeline reads them) +// --------------------------------------------------------------------------- + +std::optional PageBackend::AcquireDramCopyPin(const std::string& key) { + std::unique_lock lock(mutex_); + auto it = owned_.find(key); + if (it == owned_.end()) return std::nullopt; // already evicted -> drop task + if (pins_.find(key) != pins_.end()) return std::nullopt; // duplicate task + + DramCopyPin pin; + pin.total_size = it->second.size; + pin.segments = BuildCopySegmentsLocked(it->second.pages, it->second.size); + pin.pin_token = next_pin_token_++; + pins_[key] = PinState{pin.pin_token, std::chrono::steady_clock::now()}; + return pin; +} + +void PageBackend::ReleaseDramCopyPin(const std::string& key, uint64_t pin_token) { + std::unique_lock lock(mutex_); + auto it = pins_.find(key); + if (it == pins_.end() || it->second.token != pin_token) return; // tolerate late/dup release + pins_.erase(it); +} + +bool PageBackend::HasActivePinLocked(const std::string& key) const { + return pins_.find(key) != pins_.end(); +} + +std::vector> PageBackend::BuildCopySegmentsLocked( + const std::vector& pages, uint64_t total_size) const { + std::vector> segments; + if (buffer_bases_.empty()) return segments; + segments.reserve(pages.size()); + uint64_t remaining = total_size; + for (const auto& p : pages) { + if (p.buffer_index >= buffer_bases_.size() || buffer_bases_[p.buffer_index] == nullptr) { + MORI_UMBP_ERROR("[PageBackend] copy segment: bad buffer_index {} (bases={})", p.buffer_index, + buffer_bases_.size()); + return {}; + } + // Last page may be partial; earlier pages are full page_size. + const uint64_t bytes = std::min(page_size_, remaining); + const char* base = static_cast(buffer_bases_[p.buffer_index]); + segments.emplace_back(base + static_cast(p.page_index) * page_size_, bytes); + remaining -= bytes; + } + return segments; +} + +// --------------------------------------------------------------------------- +// Clear (distributed clear; gates writes until the full-sync ack) +// --------------------------------------------------------------------------- + +void PageBackend::ClearLocal() { + std::unique_lock lock(mutex_); + const auto now = std::chrono::steady_clock::now(); + + // Pending slots become pre-clear via generation mismatch; their pages + // stay reserved (RDMA write may still be in flight) and are freed by the + // writer's Commit or by the reaper's TTL path. + ++allocator_generation_; + + // pins_ should be empty because PoolClient::Clear quiesces the copy + // pipeline. If not, this is a caller bug; log loudly. We do not + // support clearing with active copy pins and make no attempt to salvage + // them here. A debug assert turns this into a hard failure under + // test/CI; release builds keep running (the freed pages cannot be + // reused until ClearFullSyncAcked re-enables Allocate, so this stays + // UAF-safe in practice). + if (!pins_.empty()) { + MORI_UMBP_ERROR( + "[PageBackend] ClearLocal with {} active copy pin(s) — caller did not quiesce the copy " + "pipeline (bug)", + pins_.size()); + assert(pins_.empty() && "ClearLocal called with active copy pins; quiesce the pipeline first"); + } + + // Owned: defer pages with an active read lease (an RDMA read may still be + // in flight) until their lease deadline; free the rest immediately. No + // REMOVE events — the upcoming full-sync empty snapshot collapses + // master's index. + const int64_t now_ns = SteadyNs(now); + for (auto& [key, slot] : owned_) { + const int64_t lease_ns = slot.read_lease_ns.load(std::memory_order_relaxed); + if (lease_ns > now_ns) { + DeferredFree df; + df.key = key; + df.pages = std::move(slot.pages); + df.release_at = std::chrono::steady_clock::time_point(std::chrono::nanoseconds(lease_ns)); + deferred_frees_.push_back(std::move(df)); + continue; + } + if (allocator_) allocator_->Deallocate(slot.pages); + } + owned_.clear(); + { + std::lock_guard lru_lock(lru_mutex_); + lru_.clear(); + } + + // Deadlines that mattered are already in deferred_frees_; the rest went + // with owned_. + pins_.clear(); + + // Drop any queued ADD/REMOVE that the heartbeat hasn't shipped yet: the + // snapshot we're about to send is the authoritative state. + pending_events_.clear(); + + clear_full_sync_pending_.store(true, std::memory_order_release); + MORI_UMBP_INFO("[PageBackend] ClearLocal — pending writes will be rejected until ack"); +} + +void PageBackend::ClearFullSyncAcked() { + clear_full_sync_pending_.store(false, std::memory_order_release); +} + +// --------------------------------------------------------------------------- +// Heartbeat events (drain / snapshot) — MediumBackend + OwnedLocationSource +// --------------------------------------------------------------------------- + +void PageBackend::QueueEventLocked(KvEvent event) { + // No master, no consumer: DrainPendingEvents is called only by the + // heartbeat, so queueing here would grow by one entry per put forever. + if (!event_publishing_) return; + pending_events_.push_back(std::move(event)); + // Wake the heartbeat the moment the outbox first reaches the threshold. + // Events are appended one at a time, so `==` fires at most once per + // batch. Called under mutex_: auto_flush_cb_ must be cheap and MUST NOT + // re-enter this object (it only signals the heartbeat thread). + // Exactness isn't required — the heartbeat interval is the backstop. + // threshold_ == 0 disables size-based auto-flush entirely (ADDs then + // ship only on the interval / explicit Flush). + if (auto_flush_threshold_ > 0 && pending_events_.size() == auto_flush_threshold_ && + auto_flush_cb_) { + auto_flush_cb_(); + } +} + +std::vector PageBackend::DrainPendingEvents() { + std::unique_lock lock(mutex_); + std::vector drained; + drained.swap(pending_events_); + return drained; +} + +void PageBackend::SetEventPublishing(bool enabled) { + std::unique_lock lock(mutex_); + event_publishing_ = enabled; + // Whatever is already queued has no consumer either; drop it rather than + // hold it for a master that will never register. + if (!enabled) pending_events_.clear(); +} + +// --------------------------------------------------------------------------- +// Local eviction (no-master deployments; see EnableLocalEviction) +// --------------------------------------------------------------------------- + +void PageBackend::EnableLocalEviction(double high_watermark, double low_watermark) { + if (!(high_watermark > 0.0 && high_watermark <= 1.0 && low_watermark > 0.0 && + low_watermark < high_watermark)) { + MORI_UMBP_WARN( + "[PageBackend] ignoring local eviction watermarks high={} low={} " + "(require 0 < low < high <= 1) — local eviction stays OFF", + high_watermark, low_watermark); + return; + } + std::unique_lock lock(mutex_); + local_evict_enabled_ = true; + local_evict_high_wm_ = high_watermark; + local_evict_low_wm_ = low_watermark; + // Keys committed before this call are not in lru_ and would therefore be + // invisible to a round. Enabling before Init is the contract, but seed + // anyway so a late enable degrades to "wrong order", not "never evicted". + { + std::lock_guard lru_lock(lru_mutex_); + for (auto& [key, slot] : owned_) TouchLruLocked(key, slot); + } + MORI_UMBP_INFO("[PageBackend] local eviction ON tier={} high={} low={}", static_cast(tier_), + high_watermark, low_watermark); +} + +void PageBackend::TouchLruLocked(const std::string& key, OwnedSlot& slot) { + if (!local_evict_enabled_) return; + if (slot.in_lru) { + lru_.splice(lru_.begin(), lru_, slot.lru_it); + return; + } + lru_.push_front(key); + slot.lru_it = lru_.begin(); + slot.in_lru = true; +} + +void PageBackend::RemoveFromLruLocked(OwnedSlot& slot) { + if (!slot.in_lru) return; + lru_.erase(slot.lru_it); + slot.in_lru = false; +} + +void PageBackend::TouchLru(const std::string& key, OwnedSlot& slot) { + if (!local_evict_enabled_) return; + std::lock_guard lru_lock(lru_mutex_); + TouchLruLocked(key, slot); +} + +void PageBackend::RemoveFromLru(OwnedSlot& slot) { + std::lock_guard lru_lock(lru_mutex_); + RemoveFromLruLocked(slot); +} + +size_t PageBackend::MaybeEvictToLowWatermark() { + // Only one round at a time. A second committing thread backs off rather + // than over-evicting, and try_lock keeps the put path non-blocking. + std::unique_lock round(local_evict_round_mutex_, std::try_to_lock); + if (!round.owns_lock()) return 0; + + std::unique_lock lock(mutex_); + if (!local_evict_enabled_ || !allocator_) return 0; + + const uint64_t total = allocator_->TotalBytes(); + if (total == 0) return 0; + const uint64_t used = allocator_->UsedBytes(); + if (static_cast(used) < local_evict_high_wm_ * static_cast(total)) return 0; + + const auto low_bytes = static_cast(local_evict_low_wm_ * static_cast(total)); + if (used <= low_bytes) return 0; + const uint64_t to_free = used - low_bytes; + + // Oldest first. Two passes so the list is not mutated under its own + // reverse iterator; the selection pass is where the skips happen. + const auto now = std::chrono::steady_clock::now(); + uint64_t selected_bytes = 0; + std::vector victims; + { + std::lock_guard lru_lock(lru_mutex_); + for (auto it = lru_.rbegin(); it != lru_.rend() && selected_bytes < to_free; ++it) { + auto owned_it = owned_.find(*it); + if (owned_it == owned_.end()) continue; // defensive; lru_/owned_ stay in sync + // Same two protections master-driven Evict honours: an in-flight read + // owns the pages until its lease expires, and a copy pin until the + // worker is done. A round that finds only protected keys frees what it + // can and stops — it must not spin waiting for a reader. + if (HasActiveReadLeaseLocked(owned_it->second, now)) continue; + if (HasActivePinLocked(*it)) continue; + victims.push_back(*it); + selected_bytes += static_cast(owned_it->second.pages.size()) * page_size_; + } + } + + size_t freed = 0; + for (const auto& key : victims) { + auto owned_it = owned_.find(key); + if (owned_it == owned_.end()) continue; + if (allocator_) allocator_->Deallocate(owned_it->second.pages); + // A REMOVE is queued for symmetry with master-driven Evict; with no master + // event publishing is off and this is a no-op, so nothing accumulates. + QueueEventLocked(KvEvent{KvEvent::Kind::REMOVE, key, tier_, 0}); + RemoveFromLru(owned_it->second); + owned_.erase(owned_it); + ++freed; + } + if (freed > 0) { + MORI_UMBP_DEBUG("[PageBackend] local eviction freed {} key(s) ({} bytes) tier={}", freed, + selected_bytes, static_cast(tier_)); + local_evict_warned_ = false; + } else if (!local_evict_warned_) { + // Latched: a wedged pool would otherwise log this on every commit. + local_evict_warned_ = true; + MORI_UMBP_WARN( + "[PageBackend] local eviction found no reclaimable key at {}/{} bytes used " + "(all leased or pinned) tier={}", + used, total, static_cast(tier_)); + } + return freed; +} + +void PageBackend::SetAutoFlushHook(size_t threshold, std::function cb) { + std::unique_lock lock(mutex_); + auto_flush_threshold_ = threshold; + auto_flush_cb_ = std::move(cb); +} + +std::vector PageBackend::SnapshotOwnedKeysLocked() const { + std::vector out; + out.reserve(owned_.size()); + for (const auto& kv : owned_) { + out.push_back(KvEvent{KvEvent::Kind::ADD, kv.first, tier_, kv.second.size}); + } + return out; +} + +// Test-only: a pure read-only snapshot with no outbox side effects. +// Production full-sync uses SnapshotOwnedKeysForFullSync() below; tests use +// this to assert owned state without disturbing the event outbox. +std::vector PageBackend::SnapshotOwnedKeys() const { + std::unique_lock lock(mutex_); + return SnapshotOwnedKeysLocked(); +} + +std::vector PageBackend::SnapshotOwnedKeysForFullSync() { + std::unique_lock lock(mutex_); + auto out = SnapshotOwnedKeysLocked(); + // The snapshot is now authoritative: drop the queued delta (already + // reflected in it) so the next delta carries only new events. + pending_events_.clear(); + return out; +} + +// --------------------------------------------------------------------------- +// Capacity & buffer-descriptor queries +// --------------------------------------------------------------------------- + +uint64_t PageBackend::OwnedKeyCount() const { + std::unique_lock lock(mutex_); + return owned_.size(); +} + +TierCapacity PageBackend::Capacity() const { + std::unique_lock lock(mutex_); + TierCapacity cap; + if (allocator_) { + cap.total_bytes = allocator_->TotalBytes(); + cap.available_bytes = allocator_->AvailableBytes(); + cap.max_allocatable_bytes = cap.available_bytes; + } + return cap; +} + +std::vector PageBackend::AllBufferDescs() const { + std::unique_lock lock(mutex_); + std::vector out; + out.reserve(buffer_descs_.size()); + for (size_t i = 0; i < buffer_descs_.size(); ++i) { + out.push_back({static_cast(i), buffer_descs_[i]}); + } + return out; +} + +std::vector PageBackend::BufferDescsForPages( + const std::vector& pages) const { + std::unique_lock lock(mutex_); + return BuildBufferDescsLocked(pages); +} + +// No lock: buffer_refs_ is built once by InstallTierConfig and immutable +// afterwards, and the local copy path must not contend for mutex_ (the same +// lock every Allocate/Commit/Resolve and the heartbeat snapshot take) per page. +TransferRef PageBackend::BufferRef(uint32_t buffer_index) const { + if (buffer_index >= buffer_refs_.size()) return TransferRef{}; + return buffer_refs_[buffer_index]; +} + +std::vector PageBackend::BuildBufferDescsLocked( + const std::vector& pages) const { + std::vector out; + std::vector seen; + seen.reserve(pages.size()); + for (const auto& p : pages) { + if (std::find(seen.begin(), seen.end(), p.buffer_index) != seen.end()) continue; + if (p.buffer_index >= buffer_descs_.size()) continue; // defensive: skip dangling page refs + seen.push_back(p.buffer_index); + } + std::sort(seen.begin(), seen.end()); + out.reserve(seen.size()); + for (uint32_t idx : seen) out.push_back({idx, buffer_descs_[idx]}); + return out; +} + +// --------------------------------------------------------------------------- +// Read-lease helper +// --------------------------------------------------------------------------- + +bool PageBackend::HasActiveReadLeaseLocked(const OwnedSlot& slot, + std::chrono::steady_clock::time_point now) { + return slot.read_lease_ns.load(std::memory_order_relaxed) > SteadyNs(now); +} + +// --------------------------------------------------------------------------- +// Reaper (background sweep: pending TTL, expired leases, deferred frees) +// --------------------------------------------------------------------------- + +void PageBackend::StartReaper() { + bool expected = false; + if (!reaper_running_.compare_exchange_strong(expected, true)) return; + reaper_thread_ = std::thread([this] { ReaperLoop(); }); +} + +void PageBackend::StopReaper() { + if (!reaper_running_.exchange(false)) return; + reaper_cv_.notify_all(); + if (reaper_thread_.joinable()) reaper_thread_.join(); +} + +void PageBackend::ReaperLoop() { + while (reaper_running_.load()) { + { + std::unique_lock lk(reaper_cv_mutex_); + reaper_cv_.wait_for(lk, reaper_interval_, [this] { return !reaper_running_.load(); }); + } + if (!reaper_running_.load()) break; + ReaperSweep(); + } +} + +void PageBackend::ReaperSweep() { + std::unique_lock lock(mutex_); + const auto now = std::chrono::steady_clock::now(); + + // Expire pending slots whose deadline has passed. No event is emitted: + // the slot was never owned, master never indexed it. + for (auto it = pending_.begin(); it != pending_.end();) { + if (it->second.deadline <= now) { + if (allocator_) allocator_->Deallocate(it->second.pages); + MORI_UMBP_DEBUG("[PageBackend] reaped pending slot {} ({} bytes)", it->first, + it->second.size); + it = pending_.erase(it); + } else { + ++it; + } + } + + // Warn about copy pins held far longer than any healthy copy should + // take. We never force-free a pin (a worker may still be reading its + // segments — freeing would be a use-after-free); this is purely an + // observability signal that a copy worker is stuck. + constexpr std::chrono::seconds kLongRunningPinWarn{30}; + for (const auto& [key, pin] : pins_) { + if (now - pin.acquired_at > kLongRunningPinWarn) { + MORI_UMBP_WARN("[PageBackend] copy pin for key='{}' held >{}s — copy worker stuck?", key, + kLongRunningPinWarn.count()); + } + } + + // Release ClearLocal()-deferred pages whose lease deadline has passed. + for (auto it = deferred_frees_.begin(); it != deferred_frees_.end();) { + if (it->release_at <= now) { + if (allocator_) allocator_->Deallocate(it->pages); + MORI_UMBP_DEBUG("[PageBackend] released deferred key='{}' pages={}", it->key, + it->pages.size()); + it = deferred_frees_.erase(it); + } else { + ++it; + } + } +} + +// --------------------------------------------------------------------------- +// Factory +// --------------------------------------------------------------------------- + +std::unique_ptr MakePageBackend(TierType tier, uint64_t page_size, + PageBackend::OwnershipConfig ownership, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl) { + return std::make_unique(tier, page_size, std::move(ownership), pending_ttl, + read_lease_ttl); +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/backend/ssd_backend.cpp b/src/umbp/distributed/peer/backend/ssd_backend.cpp new file mode 100644 index 000000000..6e14857e8 --- /dev/null +++ b/src/umbp/distributed/peer/backend/ssd_backend.cpp @@ -0,0 +1,888 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/peer/backend/ssd_backend.h" + +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/distributed/master/master_metrics.h" + +namespace mori::umbp { + +namespace { +constexpr uint32_t kStagingBufferIndex = 0; + +uint32_t SizeToPages(uint64_t size, uint64_t page_size) { + if (size == 0 || page_size == 0) return 0; + const uint64_t pages = 1 + (size - 1) / page_size; + return pages > UINT32_MAX ? 0 : static_cast(pages); +} +} // namespace + +SsdBackend::SsdBackend(Config cfg) : cfg_(std::move(cfg)) { + if (cfg_.page_size == 0) { + throw std::invalid_argument("SsdBackend: page_size must be > 0"); + } + if (cfg_.staging_pages == 0) { + throw std::invalid_argument("SsdBackend: staging_pages must be > 0"); + } + if (cfg_.page_size > std::numeric_limits::max() / cfg_.staging_pages || + cfg_.page_size > std::numeric_limits::max() / cfg_.staging_pages) { + throw std::invalid_argument("SsdBackend: staging arena size overflows"); + } + ssd_ = std::make_unique(cfg_.ssd); +} + +SsdBackend::~SsdBackend() { + StopReaper(); + Shutdown(); +} + +// --------------------------------------------------------------------------- +// Ownership +// --------------------------------------------------------------------------- + +bool SsdBackend::Init(MemoryRegistrar* registrar) { + if (initialized_) return true; + if (!cfg_.ssd.enabled) { + MORI_UMBP_ERROR("[SsdBackend] Init: SSD tier is disabled in config"); + return false; + } + + registrar_ = registrar; + + // Opt-in switch for the file->GPU (GDS) read path. Off by default even when + // the build has hipfile; set UMBP_ENABLE_GDS=1 to route O_DIRECT SSD reads + // through the GdsEngine instead of the staging arena (no rebuild needed). + if (const char* env = std::getenv("UMBP_ENABLE_GDS")) { + const std::string v(env); + gds_enabled_ = (v == "1" || v == "true" || v == "TRUE" || v == "on" || v == "ON"); + } + + // ONE buffer covering every staging page, so each page is contiguous and a + // single registration serves the arena. Host memory deliberately: this is + // what the backend publishes, and publishing ordinary registered DRAM is the + // entire reason SSD needs no transfer-layer change (see the header). + staging_size_ = static_cast(cfg_.staging_pages) * cfg_.page_size; + // Every byte this backend moves crosses the arena twice (device <-> arena, + // arena <-> wire), so hugepage backing matters here more than for an ordinary + // buffer. HostMemAllocator falls back to 4 KiB pages on its own when no + // hugetlb pages are free, which is what we want: a node must still come up. + HostPageMemorySource::Options staging_opts; + staging_opts.use_hugepages = cfg_.staging_use_hugepages; + staging_opts.hugepage_size = cfg_.staging_hugepage_size; + staging_source_ = std::make_unique(staging_opts); + + std::vector buffers; + if (!staging_source_->Allocate({staging_size_}, &buffers) || buffers.empty()) { + MORI_UMBP_ERROR("[SsdBackend] Init: staging arena allocation failed ({} bytes)", staging_size_); + registrar_ = nullptr; + return false; + } + staging_base_ = buffers[0].base; + staging_size_ = buffers[0].size; + + if (registrar != nullptr) { + buffer_ref_ = registrar->RegisterMemory( + staging_base_, staging_size_, staging_source_->LocationType(), staging_source_->Device()); + if (buffer_ref_.HasMemoryDesc()) { + msgpack::sbuffer sbuf; + msgpack::pack(sbuf, buffer_ref_.mem); + buffer_desc_.assign(sbuf.data(), sbuf.data() + sbuf.size()); + } + } else { + buffer_ref_ = TransferRef::HostBytes( + staging_base_, staging_size_, staging_source_->LocationType(), staging_source_->Device()); + } + + // A bitmap lets a multi-page object reserve one contiguous run while keeping + // the transfer-visible page size unchanged. + const uint32_t usable_pages = + static_cast(std::min(cfg_.staging_pages, staging_size_ / cfg_.page_size)); + staging_page_used_.assign(usable_pages, false); + + // Crash-restart leftover: metadata is gone but files may remain, so used + // capacity would diverge from owned_. Must run before any IO. + ssd_->DiscardLeftoverOnStartup(); + + initialized_ = true; + StartReaper(); + MORI_UMBP_INFO( + "[SsdBackend] Init staging_pages={} page_size={} arena_bytes={} hugepages={} " + "single_flight={}", + usable_pages, cfg_.page_size, staging_size_, cfg_.staging_use_hugepages, + cfg_.ssd.ssd.single_flight_reads); + return true; +} + +void SsdBackend::Shutdown() { + if (!initialized_) return; + StopReaper(); + + if (registrar_ != nullptr && buffer_ref_.Valid()) registrar_->Deregister(buffer_ref_); + buffer_ref_ = TransferRef{}; + buffer_desc_.clear(); + + // Release the GDS file handles obtained lazily during resolves. + { + std::lock_guard lock(gds_mutex_); + if (registrar_ != nullptr) { + for (auto& [fd, handle] : gds_handles_) { + if (handle != nullptr) registrar_->Deregister(TransferRef::File(fd, 0, 0, handle)); + } + } + gds_handles_.clear(); + } + + // Deregister before release, same ordering rule as PageBackend::Shutdown. + if (staging_source_ != nullptr) staging_source_->Release(); + staging_source_.reset(); + staging_base_ = nullptr; + + { + std::lock_guard lock(mutex_); + pending_.clear(); + migration_reads_.clear(); + read_leases_.clear(); + staging_page_used_.clear(); + } + + initialized_ = false; + registrar_ = nullptr; +} + +// --------------------------------------------------------------------------- +// Staging arena +// --------------------------------------------------------------------------- + +SsdBackend::StagingSpan SsdBackend::AcquireStagingSpanLocked(uint32_t page_count) { + if (page_count == 0 || page_count > staging_page_used_.size()) return {}; + + uint32_t run_start = 0; + uint32_t run_length = 0; + for (uint32_t i = 0; i < staging_page_used_.size(); ++i) { + if (staging_page_used_[i]) { + run_length = 0; + continue; + } + if (run_length == 0) run_start = i; + if (++run_length != page_count) continue; + for (uint32_t page = run_start; page < run_start + page_count; ++page) { + staging_page_used_[page] = true; + } + // Mirrored for the metrics tick, which must not take mutex_. Correctness + // still lives entirely in staging_page_used_. + staging_pages_in_use_.fetch_add(page_count, std::memory_order_relaxed); + return StagingSpan{run_start, page_count}; + } + return {}; +} + +void SsdBackend::ReleaseStagingSpanLocked(StagingSpan span) { + if (span.page_count == 0 || span.first_page >= staging_page_used_.size()) return; + const uint32_t end = std::min(staging_page_used_.size(), + static_cast(span.first_page) + span.page_count); + for (uint32_t page = span.first_page; page < end; ++page) { + staging_page_used_[page] = false; + } + staging_pages_in_use_.fetch_sub(end - span.first_page, std::memory_order_relaxed); +} + +void* SsdBackend::StagingPagePtr(uint32_t page_index) const { + if (staging_base_ == nullptr) return nullptr; + return static_cast(staging_base_) + static_cast(page_index) * cfg_.page_size; +} + +// --------------------------------------------------------------------------- +// Control plane +// --------------------------------------------------------------------------- + +TierCapacity SsdBackend::Capacity() const { + const auto [used, total] = ssd_->Capacity(); + TierCapacity cap; + cap.total_bytes = static_cast(total); + cap.available_bytes = total > used ? static_cast(total - used) : 0; + const uint64_t staging_capacity = + static_cast(staging_page_used_.size()) * cfg_.page_size; + cap.max_allocatable_bytes = std::min({cap.available_bytes, cap.total_bytes, staging_capacity}); + return cap; +} + +uint64_t SsdBackend::OwnedKeyCount() const { + return static_cast(ssd_->SnapshotOwnedKeys().size()); +} + +std::vector SsdBackend::DrainPendingEvents() { + { + std::lock_guard lock(mutex_); + unshipped_events_ = 0; + } + return ssd_->DrainPendingEvents(); +} + +std::vector SsdBackend::SnapshotOwnedKeys() const { return ssd_->SnapshotOwnedKeys(); } + +std::vector SsdBackend::SnapshotOwnedKeysForFullSync() { + { + std::lock_guard lock(mutex_); + unshipped_events_ = 0; + } + return ssd_->SnapshotOwnedKeysForFullSync(); +} + +void SsdBackend::ClearLocal() { + std::vector migration_keys; + { + std::lock_guard lock(mutex_); + // Invalidate outstanding slots so a late Commit fails rather than + // resurrecting a key the cluster believes is gone. + for (auto& [slot_id, slot] : pending_) ReleaseStagingSpanLocked(slot.span); + pending_.clear(); + for (auto& [key, lease] : read_leases_) ReleaseStagingSpanLocked(lease.span); + read_leases_.clear(); + for (auto& [key, lease] : migration_reads_) { + ReleaseStagingSpanLocked(lease.span); + migration_keys.push_back(key); + } + migration_reads_.clear(); + unshipped_events_ = 0; + } + for (const auto& key : migration_keys) ssd_->UnpinForMigration(key); + // The SSD wipe below can close and reopen segment fds, so drop the cached GDS + // handles (they are re-registered lazily on the next resolve) rather than risk + // a stale handle if an fd number is later reused. + { + std::lock_guard lock(gds_mutex_); + if (registrar_ != nullptr) { + for (auto& [fd, handle] : gds_handles_) { + if (handle != nullptr) registrar_->Deregister(TransferRef::File(fd, 0, 0, handle)); + } + } + gds_handles_.clear(); + } + clear_full_sync_pending_.store(true, std::memory_order_release); + ssd_->ClearLocal(); +} + +void SsdBackend::ClearFullSyncAcked() { + clear_full_sync_pending_.store(false, std::memory_order_release); +} + +void SsdBackend::SetEventPublishing(bool enabled) { + if (ssd_) ssd_->SetEventPublishing(enabled); +} + +void SsdBackend::SetAutoFlushHook(size_t threshold, std::function cb) { + std::lock_guard lock(mutex_); + auto_flush_threshold_ = threshold == 0 ? 1 : threshold; + auto_flush_cb_ = std::move(cb); +} + +// Counted here rather than read out of PeerSsdManager: its outbox is private +// and the hook only needs to know "enough has happened to be worth a flush". +void SsdBackend::NoteEventQueuedLocked() { + ++unshipped_events_; + if (auto_flush_cb_ && unshipped_events_ >= auto_flush_threshold_) { + unshipped_events_ = 0; + auto_flush_cb_(); + } +} + +// --------------------------------------------------------------------------- +// Slot lifecycle +// --------------------------------------------------------------------------- + +std::vector SsdBackend::BatchAllocate(const std::vector& entries) { + std::vector results(entries.size()); + if (clear_full_sync_pending_.load(std::memory_order_acquire)) { + // Gated until master has acked the empty full sync, same rule as + // PageBackend: no new owned key may appear in that window. + return results; + } + + const auto now = std::chrono::steady_clock::now(); + std::lock_guard lock(mutex_); + const uint64_t ssd_capacity = static_cast(ssd_->Capacity().second); + const uint64_t staging_capacity = + static_cast(staging_page_used_.size()) * cfg_.page_size; + + for (size_t i = 0; i < entries.size(); ++i) { + const auto& req = entries[i]; + AllocateResult& out = results[i]; + + const uint32_t page_count = SizeToPages(req.size, cfg_.page_size); + if (page_count == 0 || req.size > ssd_capacity || page_count > staging_page_used_.size()) { + // The object can never fit this backend's staging arena. This is a shape + // failure, not transient arena pressure, so another identical peer would + // not do better. + MORI_UMBP_WARN("[SsdBackend] key '{}' size {} exceeds permanent limit {}", req.key, req.size, + std::min(ssd_capacity, staging_capacity)); + out.outcome = AllocateOutcome::kFailed; + continue; + } + if (ssd_->Exists(req.key)) { + out.outcome = AllocateOutcome::kSuccessAlreadyExists; + continue; + } + + const StagingSpan span = AcquireStagingSpanLocked(page_count); + if (span.page_count == 0) { + // The arena is a real capacity limit, so this IS retry-elsewhere-worthy. + out.outcome = AllocateOutcome::kFailedNoSpace; + continue; + } + + const uint64_t slot_id = next_slot_id_.fetch_add(1, std::memory_order_relaxed); + pending_.emplace(slot_id, + PendingSlot{slot_id, req.key, span, req.size, now + cfg_.pending_ttl}); + + out.outcome = AllocateOutcome::kSuccessAllocated; + out.slot_id = slot_id; + out.pages.reserve(span.page_count); + for (uint32_t page = span.first_page; page < span.first_page + span.page_count; ++page) { + out.pages.push_back(PageLocation{kStagingBufferIndex, page}); + } + out.size = req.size; + out.page_size = cfg_.page_size; + out.pending_ttl_ms = static_cast(cfg_.pending_ttl.count()); + if (!buffer_desc_.empty()) { + out.descs = {BufferMemoryDescBytes{kStagingBufferIndex, buffer_desc_}}; + } + } + return results; +} + +std::vector SsdBackend::BatchCommit(const std::vector& entries) { + std::vector results(entries.size()); + if (entries.empty()) return results; + + // Phase 1 (ONE lock): take every slot out of pending_. An entry whose slot is + // gone — reaped, aborted, cleared, or never there — fails here and takes no + // further part. + std::vector slots(entries.size()); + std::vector todo; + todo.reserve(entries.size()); + { + std::lock_guard lock(mutex_); + for (size_t i = 0; i < entries.size(); ++i) { + auto it = pending_.find(entries[i].slot_id); + if (it == pending_.end()) continue; // results[i].success stays false + slots[i] = it->second; + pending_.erase(it); + todo.push_back(i); + } + } + if (todo.empty()) return results; + + // Phase 2 (no lock): ONE batched SSD write. A per-key Write loop would idle + // every drive but one of a multi-drive tier, so keeping the batch intact all + // the way down to the backend is the whole point of this path. An entry whose + // staging page cannot be resolved is dropped from the batch and fails. + std::vector keys; + std::vector srcs; + std::vector sizes; + std::vector batched; // indices into `entries`, parallel to the above + keys.reserve(todo.size()); + srcs.reserve(todo.size()); + sizes.reserve(todo.size()); + batched.reserve(todo.size()); + for (size_t i : todo) { + const void* src = StagingPagePtr(slots[i].span.first_page); + if (src == nullptr) continue; + keys.push_back(entries[i].key); + srcs.push_back(src); + sizes.push_back(slots[i].size); + batched.push_back(i); + } + + std::vector ok; + if (!batched.empty()) { + ok = ssd_->WriteBatch(keys, srcs, sizes); + ok.resize(batched.size(), false); + } + + // Phase 3 (ONE lock): return every borrowed staging page — on success AND on + // failure, since the page was only ever borrowed for the write — and queue one + // event per key that landed. + { + std::lock_guard lock(mutex_); + for (size_t i : todo) ReleaseStagingSpanLocked(slots[i].span); + for (size_t j = 0; j < batched.size(); ++j) { + if (ok[j]) NoteEventQueuedLocked(); + } + } + + for (size_t j = 0; j < batched.size(); ++j) { + const size_t i = batched[j]; + if (!ok[j]) { + MORI_UMBP_ERROR("[SsdBackend] Commit: SSD write failed for key '{}' ({} bytes)", + entries[i].key, slots[i].size); + continue; + } + results[i].success = true; + results[i].bytes_committed = slots[i].size; + } + return results; +} + +std::vector SsdBackend::BatchAbort(const std::vector& slot_ids) { + std::vector results(slot_ids.size(), true); + std::lock_guard lock(mutex_); + for (uint64_t slot_id : slot_ids) { + auto it = pending_.find(slot_id); + if (it == pending_.end()) continue; // idempotent: already reaped or unknown + ReleaseStagingSpanLocked(it->second.span); + pending_.erase(it); + } + return results; +} + +std::vector SsdBackend::BatchResolve(const std::vector& keys, + bool include_descs, bool allow_file_refs) { + std::vector results(keys.size()); + if (keys.empty()) return results; + const auto now = std::chrono::steady_clock::now(); + + // Fills `out` from a lease the caller has already located. Caller holds + // mutex_. Factored out because the lease-hit and read-completed paths publish + // byte-identical results and drifting between them would be a silent bug. + auto publish = [&](ResolvedEntry& out, const ReadLease& lease) { + out.outcome = ResolveOutcome::kFound; + out.found = true; + out.pages.reserve(lease.span.page_count); + for (uint32_t page = lease.span.first_page; + page < lease.span.first_page + lease.span.page_count; ++page) { + out.pages.push_back(PageLocation{kStagingBufferIndex, page}); + } + out.size = lease.size; + out.page_size = cfg_.page_size; + if (include_descs && !buffer_desc_.empty()) { + out.descs = {BufferMemoryDescBytes{kStagingBufferIndex, buffer_desc_}}; + } + }; + + // Phase 1 (ONE lock): serve every key that is already staged, size every + // unique remaining key, then reserve ALL required spans atomically. A + // partial resolve is unusable to layer-wise callers and, worse, pins the + // successful subset while retries compete for the remainder. On transient + // pressure this call therefore publishes kBusy and keeps none of its new + // reservations. + std::vector todo; // keys needing a device read + std::vector page_counts; // parallel to todo + std::vector spans; // parallel to todo + std::vector dsts; // parallel to todo + std::vector caps; // parallel to todo + std::vector read_keys; // parallel to todo + std::vector> result_groups; // duplicate input indices per device read + std::unordered_map scheduled; + todo.reserve(keys.size()); + { + std::lock_guard lock(mutex_); + for (size_t i = 0; i < keys.size(); ++i) { + auto it = read_leases_.find(keys[i]); + if (it != read_leases_.end()) { + it->second.expires_at = now + cfg_.read_lease_ttl; + publish(results[i], it->second); + continue; + } + + auto duplicate = scheduled.find(keys[i]); + if (duplicate != scheduled.end()) { + result_groups[duplicate->second].push_back(i); + continue; + } + + // Zero-copy GDS (UMBP_ENABLE_GDS): when the caller can read a file ref, + // the switch is on, a file engine is present, and the record is on an + // O_DIRECT fd, publish a FileRef and skip the staging arena — the reader + // (GdsEngine) DMAs the range into device memory. + // + // allow_file_refs is tested FIRST on purpose. A caller that cannot use a + // file ref — a peer serving the wire, or an existence check reading only + // `found` — must fall through to the staging path and get real pages, and + // testing it first also spares it LocateRecord and the hipFile + // registration GdsHandleForFd performs as a side effect. + std::optional loc; + if (allow_file_refs && gds_enabled_ && (loc = ssd_->LocateRecord(keys[i])) && + loc->direct_io && loc->fd >= 0) { + if (void* handle = GdsHandleForFd(loc->fd)) { + results[i].outcome = ResolveOutcome::kFound; + results[i].found = true; + results[i].size = loc->value_size; + results[i].page_size = cfg_.page_size; + results[i].file_ref = + TransferRef::File(loc->fd, loc->value_offset, loc->readable_size, handle); + continue; + } + } + + const uint64_t size = ssd_->SizeOf(keys[i]); + const uint32_t page_count = SizeToPages(size, cfg_.page_size); + if (size == 0) continue; + if (page_count == 0 || page_count > staging_page_used_.size()) { + results[i].outcome = ResolveOutcome::kFailed; + MORI_UMBP_ERROR("[SsdBackend] Resolve: key '{}' size {} cannot fit staging arena", keys[i], + size); + continue; + } + todo.push_back(i); + page_counts.push_back(page_count); + read_keys.push_back(keys[i]); + scheduled.emplace(keys[i], todo.size() - 1); + result_groups.push_back({i}); + } + + uint64_t requested_pages = 0; + for (uint32_t page_count : page_counts) requested_pages += page_count; + if (requested_pages > staging_page_used_.size()) { + for (const auto& group : result_groups) { + for (size_t index : group) results[index].outcome = ResolveOutcome::kFailed; + } + MORI_UMBP_ERROR( + "[SsdBackend] Resolve: batch requires {} staging pages but arena has {}; not retryable", + requested_pages, staging_page_used_.size()); + return results; + } + + spans.reserve(todo.size()); + for (size_t j = 0; j < todo.size(); ++j) { + const StagingSpan span = AcquireStagingSpanLocked(page_counts[j]); + if (span.page_count == 0) { + for (const StagingSpan reserved : spans) ReleaseStagingSpanLocked(reserved); + for (const auto& group : result_groups) { + for (size_t index : group) results[index].outcome = ResolveOutcome::kBusy; + } + slot_full_rejects_.fetch_add(1, std::memory_order_relaxed); + MORI_UMBP_WARN( + "[SsdBackend] Resolve: staging arena busy; rolled back {} reservations for {} keys", + spans.size(), todo.size()); + return results; + } + spans.push_back(span); + } + for (const StagingSpan span : spans) { + dsts.push_back(StagingPagePtr(span.first_page)); + caps.push_back(static_cast(span.page_count) * cfg_.page_size); + } + } + if (todo.empty()) return results; + + // Phase 2 (no lock): ONE batched SSD read. ShardedSsdTier turns this into + // concurrent IO on every drive holding part of the batch. Duplicate keys in + // this call were collapsed above; PeerSsdManager additionally coalesces reads + // racing across calls. It marks each key in-flight for the window, so + // eviction already skips them. + const std::vector outcomes = ssd_->PrepareReadBatch(read_keys, dsts, caps); + + // Phase 3 (ONE lock): install a lease per key that landed, hand back the pages + // of the ones that did not. + { + std::lock_guard lock(mutex_); + for (size_t j = 0; j < todo.size(); ++j) { + const size_t primary = todo[j]; + if (outcomes[j].status != SsdReadStatus::kOk) { + ReleaseStagingSpanLocked(spans[j]); + const ResolveOutcome outcome = outcomes[j].status == SsdReadStatus::kNotFound + ? ResolveOutcome::kMissing + : ResolveOutcome::kFailed; + for (size_t result_index : result_groups[j]) { + results[result_index].outcome = outcome; + } + continue; + } + // A concurrent Resolve of the same key may have staged it while this read + // was in flight; keep the winner and give this page back. + const auto lease_deadline = std::chrono::steady_clock::now() + cfg_.read_lease_ttl; + auto [it, inserted] = read_leases_.emplace( + keys[primary], ReadLease{spans[j], outcomes[j].size, lease_deadline}); + if (!inserted) { + ReleaseStagingSpanLocked(spans[j]); + it->second.expires_at = lease_deadline; + } + for (size_t result_index : result_groups[j]) { + publish(results[result_index], it->second); + } + } + } + return results; +} + +bool SsdBackend::AcquireMigrationRead(const std::string& key, ResolvedEntry* resolved) { + if (resolved == nullptr) return false; + if (!ssd_->PinForMigration(key)) return false; + const uint64_t size = ssd_->SizeOf(key); + const uint32_t page_count = SizeToPages(size, cfg_.page_size); + if (page_count == 0 || page_count > staging_page_used_.size()) { + ssd_->UnpinForMigration(key); + return false; + } + + StagingSpan span; + { + std::lock_guard lock(mutex_); + if (migration_reads_.count(key) != 0) { + ssd_->UnpinForMigration(key); + return false; + } + span = AcquireStagingSpanLocked(page_count); + if (span.page_count == 0) { + ssd_->UnpinForMigration(key); + return false; + } + migration_reads_[key] = ReadLease{span, size, std::chrono::steady_clock::time_point::max()}; + } + const SsdReadOutcome outcome = ssd_->PrepareRead( + key, StagingPagePtr(span.first_page), static_cast(span.page_count) * cfg_.page_size); + if (outcome.status != SsdReadStatus::kOk) { + ReleaseMigrationRead(key); + return false; + } + resolved->outcome = ResolveOutcome::kFound; + resolved->found = true; + resolved->pages.reserve(span.page_count); + for (uint32_t page = span.first_page; page < span.first_page + span.page_count; ++page) { + resolved->pages.push_back(PageLocation{kStagingBufferIndex, page}); + } + resolved->size = outcome.size; + resolved->page_size = cfg_.page_size; + return true; +} + +void SsdBackend::ReleaseMigrationRead(const std::string& key) { + bool released = false; + { + std::lock_guard lock(mutex_); + auto it = migration_reads_.find(key); + if (it == migration_reads_.end()) return; + ReleaseStagingSpanLocked(it->second.span); + migration_reads_.erase(it); + released = true; + } + if (released) ssd_->UnpinForMigration(key); +} + +bool SsdBackend::Contains(const std::string& key) const { + return ssd_ != nullptr && ssd_->Exists(key); +} + +std::vector SsdBackend::Evict(const std::vector& keys) { + std::vector results; + results.reserve(keys.size()); + + for (const std::string& key : keys) { + EvictResult r; + r.key = key; + + // A staged read holds the bytes live for its lease; PeerSsdManager::Evict + // would refuse anyway (inflight_reads_), but reporting 0 here keeps the + // master's retry behavior explicit rather than incidental. + { + std::lock_guard lock(mutex_); + if (read_leases_.count(key) != 0 || migration_reads_.count(key) != 0) { + results.push_back(r); // bytes_freed = 0 -> master retries next round + continue; + } + } + + const uint64_t size = ssd_->SizeOf(key); + if (size != 0 && ssd_->Evict(key)) { + r.bytes_freed = size; + std::lock_guard lock(mutex_); + NoteEventQueuedLocked(); // PeerSsdManager queued the REMOVE + } + results.push_back(r); + } + return results; +} + +// --------------------------------------------------------------------------- +// Bootstrap +// --------------------------------------------------------------------------- + +std::vector SsdBackend::AllBufferDescs() const { + if (buffer_desc_.empty()) return {}; + return {BufferMemoryDescBytes{kStagingBufferIndex, buffer_desc_}}; +} + +TransferRef SsdBackend::BufferRef(uint32_t buffer_index) const { + if (buffer_index != kStagingBufferIndex) return TransferRef{}; + return buffer_ref_; +} + +void* SsdBackend::GdsHandleForFd(int fd) { + std::lock_guard lock(gds_mutex_); + auto it = gds_handles_.find(fd); + if (it != gds_handles_.end()) return it->second; // cached, possibly nullptr + // A zero-length RegisterFile just obtains (and ref-counts) the fd's handle; + // the per-key ranges are built by BatchResolve. An invalid ref means no file + // engine is configured — cache nullptr so the fd is not probed again. + void* handle = nullptr; + if (registrar_ != nullptr) { + TransferRef ref = registrar_->RegisterFile(fd, 0, 0); + if (ref.IsFile()) handle = ref.gds_handle; + } + gds_handles_[fd] = handle; + return handle; +} + +// --------------------------------------------------------------------------- +// Reaper +// --------------------------------------------------------------------------- + +void SsdBackend::StartReaper() { + bool expected = false; + if (!reaper_running_.compare_exchange_strong(expected, true)) return; + reaper_thread_ = std::thread([this] { ReaperLoop(); }); +} + +void SsdBackend::StopReaper() { + if (!reaper_running_.exchange(false)) return; + reaper_cv_.notify_all(); + if (reaper_thread_.joinable()) reaper_thread_.join(); +} + +void SsdBackend::ReaperLoop() { + while (reaper_running_.load(std::memory_order_acquire)) { + { + std::unique_lock lock(reaper_cv_mutex_); + reaper_cv_.wait_for(lock, cfg_.reaper_interval, + [this] { return !reaper_running_.load(std::memory_order_acquire); }); + } + if (!reaper_running_.load(std::memory_order_acquire)) break; + ReaperSweep(); + } +} + +// Returns staging pages whose borrower is gone: a writer that died between +// Allocate and Commit, and a reader whose lease has elapsed. Both are pure +// arena bookkeeping — no SSD state is touched, because neither case ever put +// anything on the SSD. +void SsdBackend::ReaperSweep() { + const auto now = std::chrono::steady_clock::now(); + std::lock_guard lock(mutex_); + + for (auto it = pending_.begin(); it != pending_.end();) { + if (it->second.deadline <= now) { + MORI_UMBP_WARN("[SsdBackend] reaping expired slot {} (key '{}')", it->second.slot_id, + it->second.key); + ReleaseStagingSpanLocked(it->second.span); + staging_expired_reclaims_.fetch_add(1, std::memory_order_relaxed); + it = pending_.erase(it); + } else { + ++it; + } + } + + for (auto it = read_leases_.begin(); it != read_leases_.end();) { + if (it->second.expires_at <= now) { + ReleaseStagingSpanLocked(it->second.span); + staging_expired_reclaims_.fetch_add(1, std::memory_order_relaxed); + it = read_leases_.erase(it); + } else { + ++it; + } + } +} + +std::unique_ptr MakeSsdBackend(SsdBackend::Config cfg) { + return std::make_unique(std::move(cfg)); +} + +std::vector SsdBackend::SampleMetrics() const { + // Everything a caller can see from OUTSIDE this class — resolve hits and + // misses, bytes handed out, eviction volume, per-call latency — is already + // measured by InstrumentedBackend against the MediumBackend interface, and is + // NOT repeated here. What is left is exactly what the interface hides: what + // the drive itself did, and how full the staging arena is. + // + // These go out under the generic MORI_UMBP_METRIC_BACKEND_MEDIUM_* names with + // the specifics in an event= / state= label, which is what lets one dashboard + // panel render this medium beside every other one. tier= and backend= are + // added by the publisher. + std::vector out; + if (ssd_ == nullptr) return out; + + auto event = [&out](const char* name, uint64_t v) { + if (v == 0) return; + out.push_back(MetricSample{MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL, + MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL_HELP, + {{"event", name}}, + v}); + }; + auto bytes = [&out](const char* name, uint64_t v) { + if (v == 0) return; + out.push_back(MetricSample{MORI_UMBP_METRIC_BACKEND_MEDIUM_BYTES_TOTAL, + MORI_UMBP_METRIC_BACKEND_MEDIUM_BYTES_TOTAL_HELP, + {{"event", name}}, + v}); + }; + auto state = [&out](const char* name, uint64_t v) { + out.push_back(MetricSample{MORI_UMBP_METRIC_BACKEND_MEDIUM_STATE, + MORI_UMBP_METRIC_BACKEND_MEDIUM_STATE_HELP, + {{"state", name}}, + v, + MetricKind::kGauge}); + }; + + // Drive-level read outcomes. These differ from the resolve counts the + // decorator produces: a resolve that a single-flight merge served never + // reached the drive, and a size_too_large never became a miss upstream. + event("device_read_ok", ssd_->ReadOk()); + event("device_read_not_found", ssd_->ReadNotFound()); + event("device_read_size_too_large", ssd_->ReadSizeTooLarge()); + event("device_read_error", ssd_->ReadError()); + + // Single-flight: dup is the opportunity, merged is what was actually taken. + // dup - merged flatlining at dup means coalescing is off or the merge window + // is being missed, which is invisible from the read outcomes alone. + event("single_flight_lead", ssd_->ReadLead()); + event("single_flight_dup", ssd_->ReadDup()); + event("single_flight_merged", ssd_->ReadMerged()); + + // Local high-watermark eviction. The victims and the freed bytes are the + // decorator's (Evict is an interface call); the ROUNDS are not — master never + // asked for them — and neither is a backend delete that refused. + event("eviction_round", ssd_->EvictionRounds()); + event("eviction_backend_failed", ssd_->EvictionBackendFailures()); + + // Staging pressure. slot_full_reject is the one to watch: it counts resolves + // that reported a miss for a key this node HOLDS, purely because the arena + // was full — indistinguishable from a real miss anywhere else. + event("staging_slot_full_reject", slot_full_rejects_.load(std::memory_order_relaxed)); + event("staging_expired_reclaim", staging_expired_reclaims_.load(std::memory_order_relaxed)); + + // Bytes that actually crossed the device boundary, as opposed to the logical + // bytes in mori_umbp_backend_bytes_total. Reads exclude single-flight merges, + // so this stays a true measure of device traffic. + bytes("device_write", ssd_->CopyBytes()); + bytes("device_read", ssd_->ReadBytes()); + + state("staging_pages_in_use", staging_pages_in_use_.load(std::memory_order_relaxed)); + state("staging_pages_total", cfg_.staging_pages); + return out; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/peer_dram_allocator.cpp b/src/umbp/distributed/peer/peer_dram_allocator.cpp deleted file mode 100644 index 034cd4278..000000000 --- a/src/umbp/distributed/peer/peer_dram_allocator.cpp +++ /dev/null @@ -1,683 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/distributed/peer/peer_dram_allocator.h" - -#include -#include -#include -#include -#include - -#include "mori/utils/mori_log.hpp" - -namespace mori::umbp { - -// --------------------------------------------------------------------------- -// File-local helpers -// --------------------------------------------------------------------------- - -namespace { - -// Round size up to whole pages of `page_size`. Returns 0 only when -// size == 0; the allocator treats num_pages == 0 as ENOSPC anyway. -uint32_t SizeToPages(uint64_t size, uint64_t page_size) { - if (page_size == 0 || size == 0) return 0; - uint64_t pages = (size + page_size - 1) / page_size; - // Cap at uint32_t — PageBitmapAllocator uses uint32_t for page counts. - if (pages > std::numeric_limits::max()) { - return 0; - } - return static_cast(pages); -} - -} // namespace - -// --------------------------------------------------------------------------- -// Construction -// --------------------------------------------------------------------------- - -PeerDramAllocator::PeerDramAllocator(uint64_t page_size, TierConfig dram, TierConfig hbm, - std::chrono::milliseconds pending_ttl, - std::chrono::milliseconds read_lease_ttl, - std::chrono::milliseconds reaper_interval) - : page_size_(page_size), - pending_ttl_(pending_ttl), - read_lease_ttl_(read_lease_ttl), - reaper_interval_(reaper_interval) { - if (page_size == 0) { - throw std::invalid_argument("PeerDramAllocator: page_size must be > 0"); - } - auto install_tier = [&](TierType tier, TierConfig& cfg) { - if (cfg.buffer_sizes.empty() && cfg.buffer_descs.empty()) return; - if (cfg.buffer_sizes.size() != cfg.buffer_descs.size()) { - throw std::invalid_argument("PeerDramAllocator: buffer_sizes / buffer_descs length mismatch"); - } - // buffer_bases is optional (deployments / tests that never pin leave it - // empty); when present it must line up with the buffers one-to-one. - if (!cfg.buffer_bases.empty() && cfg.buffer_bases.size() != cfg.buffer_sizes.size()) { - throw std::invalid_argument("PeerDramAllocator: buffer_bases / buffer_sizes length mismatch"); - } - allocators_.emplace(tier, std::make_unique(page_size, cfg.buffer_sizes)); - tier_descs_.emplace(tier, std::move(cfg.buffer_descs)); - tier_bases_.emplace(tier, std::move(cfg.buffer_bases)); - }; - install_tier(TierType::DRAM, dram); - install_tier(TierType::HBM, hbm); -} - -PeerDramAllocator::~PeerDramAllocator() { StopReaper(); } - -// --------------------------------------------------------------------------- -// Allocation (reserve pages into a pending slot) -// --------------------------------------------------------------------------- - -PageBitmapAllocator* PeerDramAllocator::AllocatorForLocked(TierType tier) { - auto it = allocators_.find(tier); - return it == allocators_.end() ? nullptr : it->second.get(); -} -const PageBitmapAllocator* PeerDramAllocator::AllocatorForLocked(TierType tier) const { - auto it = allocators_.find(tier); - return it == allocators_.end() ? nullptr : it->second.get(); -} - -PeerDramAllocator::AllocateResult PeerDramAllocator::Allocate(const std::string& key, uint64_t size, - TierType tier) { - std::lock_guard lock(mutex_); - return AllocateLocked(key, size, tier); -} - -PeerDramAllocator::AllocateResult PeerDramAllocator::AllocateLocked(const std::string& key, - uint64_t size, TierType tier) { - auto fail = [&](Outcome outcome, const char* reason) { - MORI_UMBP_WARN("[PeerDramAllocator] Allocate reason={} key='{}' size={} tier={}", reason, key, - size, static_cast(tier)); - AllocateResult r; - r.outcome = outcome; - return r; - }; - - const uint32_t num_pages = SizeToPages(size, page_size_); - if (num_pages == 0) return fail(Outcome::kFailed, "ZERO_SIZE"); - - // Block new allocations between ClearLocal() and full-sync ack — any - // owned key created here would miss the empty snapshot to master. - if (clear_full_sync_pending_.load(std::memory_order_acquire)) { - return fail(Outcome::kFailed, "CLEAR_PENDING"); - } - - // owned_ dedup (master-index-lag fallback). pending_ deliberately - // not checked — same-key pending race is absorbed by Commit(). - if (owned_.find(key) != owned_.end()) { - AllocateResult out; - out.outcome = Outcome::kSuccessAlreadyExists; - return out; - } - - PageBitmapAllocator* alloc = AllocatorForLocked(tier); - if (alloc == nullptr) return fail(Outcome::kFailed, "BAD_TIER"); - - auto pages = alloc->Allocate(num_pages); - if (!pages) return fail(Outcome::kFailedNoSpace, "NO_SPACE"); - - AllocateResult out; - - PendingSlot slot; - slot.slot_id = next_slot_id_.fetch_add(1, std::memory_order_relaxed); - slot.tier = tier; - slot.pages = std::move(*pages); - slot.size = size; - slot.deadline = std::chrono::steady_clock::now() + pending_ttl_; - slot.generation = allocator_generation_; - pending_[slot.slot_id] = slot; - out.outcome = Outcome::kSuccessAllocated; - out.slot = std::move(slot); - return out; -} - -std::vector PeerDramAllocator::BatchAllocate( - const std::vector& entries) { - std::vector out(entries.size()); - if (entries.empty()) return out; - std::lock_guard lock(mutex_); - for (size_t i = 0; i < entries.size(); ++i) { - const auto& entry = entries[i]; - auto result = AllocateLocked(entry.key, entry.size, entry.tier); - out[i].outcome = result.outcome; - out[i].slot = std::move(result.slot); - if (out[i].outcome == Outcome::kSuccessAllocated && out[i].slot.has_value()) { - out[i].descs = BuildBufferDescsLocked(out[i].slot->tier, out[i].slot->pages); - } - } - return out; -} - -// --------------------------------------------------------------------------- -// Commit / abort (promote a pending slot to owned, or release it) -// --------------------------------------------------------------------------- - -bool PeerDramAllocator::Commit(uint64_t slot_id, const std::string& key, - uint64_t& bytes_committed) { - std::lock_guard lock(mutex_); - return CommitLocked(slot_id, key, bytes_committed); -} - -bool PeerDramAllocator::CommitLocked(uint64_t slot_id, const std::string& key, - uint64_t& bytes_committed) { - bytes_committed = 0; - auto it = pending_.find(slot_id); - if (it == pending_.end()) { - MORI_UMBP_WARN("[PeerDramAllocator] Commit reason=SLOT_GONE key='{}' slot_id={}", key, slot_id); - return false; - } - - // Pre-clear pending slot: free pages, report Put failure, no ADD. - if (it->second.generation != allocator_generation_) { - MORI_UMBP_WARN("[PeerDramAllocator] Commit reason=PRE_CLEAR key='{}' slot_id={}", key, slot_id); - if (auto* alloc = AllocatorForLocked(it->second.tier)) { - alloc->Deallocate(it->second.pages); - } - pending_.erase(it); - return false; - } - - // Race-window safety net: two writers passed Allocate() before either - // committed. Keep first, drop new pages, idempotent success. - auto existing = owned_.find(key); - if (existing != owned_.end()) { - MORI_UMBP_WARN( - "[PeerDramAllocator] duplicate Commit for key='{}' " - "(existing tier={} size={}, new size={}) — keeping prior slot", - key, static_cast(existing->second.tier), existing->second.size, it->second.size); - if (auto* alloc = AllocatorForLocked(it->second.tier)) { - alloc->Deallocate(it->second.pages); - } - bytes_committed = existing->second.size; - pending_.erase(it); - return true; - } - - OwnedSlot owned; - owned.tier = it->second.tier; - owned.pages = std::move(it->second.pages); - owned.size = it->second.size; - QueueEventLocked(KvEvent{KvEvent::Kind::ADD, key, owned.tier, owned.size}); - owned_[key] = std::move(owned); - pending_.erase(it); - bytes_committed = owned_[key].size; - return true; -} - -std::vector PeerDramAllocator::BatchCommit( - const std::vector& entries) { - std::vector out(entries.size()); - if (entries.empty()) return out; - std::lock_guard lock(mutex_); - for (size_t i = 0; i < entries.size(); ++i) { - const auto& entry = entries[i]; - out[i].success = CommitLocked(entry.slot_id, entry.key, out[i].bytes_committed); - } - return out; -} - -bool PeerDramAllocator::Abort(uint64_t slot_id) { - std::lock_guard lock(mutex_); - return AbortLocked(slot_id); -} - -bool PeerDramAllocator::AbortLocked(uint64_t slot_id) { - auto it = pending_.find(slot_id); - if (it == pending_.end()) return true; // already reaped / aborted — idempotent - if (auto* alloc = AllocatorForLocked(it->second.tier)) { - alloc->Deallocate(it->second.pages); - } - pending_.erase(it); - return true; -} - -std::vector PeerDramAllocator::BatchAbort(const std::vector& slot_ids) { - std::vector out(slot_ids.size(), false); - if (slot_ids.empty()) return out; - std::lock_guard lock(mutex_); - for (size_t i = 0; i < slot_ids.size(); ++i) { - out[i] = AbortLocked(slot_ids[i]); - } - return out; -} - -// --------------------------------------------------------------------------- -// Resolve (read path; grants a short read lease to fence eviction) -// --------------------------------------------------------------------------- - -PeerDramAllocator::ResolveResult PeerDramAllocator::Resolve(const std::string& key) { - std::lock_guard lock(mutex_); - ResolveResult r; - auto it = owned_.find(key); - if (it == owned_.end()) return r; - r.found = true; - r.tier = it->second.tier; - r.pages = it->second.pages; - r.size = it->second.size; - // Extend the read lease so concurrent Evict reports bytes_freed=0 for - // this key. steady_clock is monotonic and read_lease_ttl_ is fixed, - // so this assignment is always >= any previous deadline for the key. - read_lease_until_[key] = std::chrono::steady_clock::now() + read_lease_ttl_; - return r; -} - -std::vector PeerDramAllocator::BatchResolve( - const std::vector& keys, bool include_descs) { - std::vector out(keys.size()); - if (keys.empty()) return out; - std::lock_guard lock(mutex_); - for (size_t i = 0; i < keys.size(); ++i) { - const auto& key = keys[i]; - auto it = owned_.find(key); - if (it == owned_.end()) continue; - auto& entry = out[i]; - entry.found = true; - entry.tier = it->second.tier; - entry.pages = it->second.pages; - entry.size = it->second.size; - if (include_descs) { - entry.descs = BuildBufferDescsLocked(it->second.tier, it->second.pages); - } - // Per-key now(): matches single-key Resolve() so the last key in - // a large batch isn't shortchanged by earlier keys' work. - read_lease_until_[key] = std::chrono::steady_clock::now() + read_lease_ttl_; - } - return out; -} - -// --------------------------------------------------------------------------- -// Eviction (skips leased / copy-pinned keys; emits REMOVE) -// --------------------------------------------------------------------------- - -std::vector PeerDramAllocator::Evict( - const std::vector& keys) { - std::vector out; - out.reserve(keys.size()); - std::lock_guard lock(mutex_); - for (const auto& key : keys) { - EvictResult r; - r.key = key; - auto it = owned_.find(key); - if (it == owned_.end()) { - out.push_back(std::move(r)); - continue; - } - if (HasActiveReadLeaseLocked(key)) { - // Master will retry next round once the lease expires. Emit no event. - out.push_back(std::move(r)); - continue; - } - if (HasActivePinLocked(key)) { - // An SSD copy worker is reading these pages. Do NOT free them, do - // NOT emit REMOVE DRAM, keep the key owned. bytes_freed=0 tells - // master to retry; the pin is released when the copy finishes. - out.push_back(std::move(r)); - continue; - } - if (auto* alloc = AllocatorForLocked(it->second.tier)) { - alloc->Deallocate(it->second.pages); - } - r.bytes_freed = it->second.size; - QueueEventLocked(KvEvent{KvEvent::Kind::REMOVE, key, it->second.tier, 0}); - owned_.erase(it); - out.push_back(std::move(r)); - } - return out; -} - -// --------------------------------------------------------------------------- -// DRAM copy pins (protect owned pages while the SSD copy pipeline reads them) -// --------------------------------------------------------------------------- - -std::optional PeerDramAllocator::AcquireDramCopyPin( - const std::string& key) { - std::lock_guard lock(mutex_); - auto it = owned_.find(key); - if (it == owned_.end()) return std::nullopt; // already evicted -> drop task - if (pins_.find(key) != pins_.end()) return std::nullopt; // duplicate task - - DramCopyPin pin; - pin.total_size = it->second.size; - pin.segments = BuildCopySegmentsLocked(it->second.tier, it->second.pages, it->second.size); - pin.pin_token = next_pin_token_++; - pins_[key] = PinState{pin.pin_token, std::chrono::steady_clock::now()}; - return pin; -} - -void PeerDramAllocator::ReleaseDramCopyPin(const std::string& key, uint64_t pin_token) { - std::lock_guard lock(mutex_); - auto it = pins_.find(key); - if (it == pins_.end() || it->second.token != pin_token) return; // tolerate late/dup release - pins_.erase(it); -} - -bool PeerDramAllocator::HasActivePinLocked(const std::string& key) const { - return pins_.find(key) != pins_.end(); -} - -std::vector> PeerDramAllocator::BuildCopySegmentsLocked( - TierType tier, const std::vector& pages, uint64_t total_size) const { - std::vector> segments; - auto it = tier_bases_.find(tier); - if (it == tier_bases_.end() || it->second.empty()) return segments; - const auto& bases = it->second; - segments.reserve(pages.size()); - uint64_t remaining = total_size; - for (const auto& p : pages) { - if (p.buffer_index >= bases.size() || bases[p.buffer_index] == nullptr) { - MORI_UMBP_ERROR("[PeerDramAllocator] copy segment: bad buffer_index {} (bases={})", - p.buffer_index, bases.size()); - return {}; - } - // Last page may be partial; earlier pages are full page_size. - const uint64_t bytes = std::min(page_size_, remaining); - const char* base = static_cast(bases[p.buffer_index]); - segments.emplace_back(base + static_cast(p.page_index) * page_size_, bytes); - remaining -= bytes; - } - return segments; -} - -// --------------------------------------------------------------------------- -// Clear (distributed clear; gates writes until the full-sync ack) -// --------------------------------------------------------------------------- - -void PeerDramAllocator::ClearLocal() { - std::lock_guard lock(mutex_); - const auto now = std::chrono::steady_clock::now(); - - // Pending slots become pre-clear via generation mismatch; their - // pages stay reserved (RDMA write may still be in flight) and are - // freed by the writer's Commit or by the reaper's TTL path. - ++allocator_generation_; - - // pins_ should be empty because PoolClient::Clear quiesces the copy pipeline. - // If not, this is a caller bug; log loudly. We do not support clearing with - // active copy pins and make no attempt to salvage them here. A debug assert - // turns this into a hard failure under test/CI; release builds keep running - // (the freed pages cannot be reused until ClearFullSyncAcked re-enables - // Allocate, so this stays UAF-safe in practice). - if (!pins_.empty()) { - MORI_UMBP_ERROR( - "[PeerDramAllocator] ClearLocal with {} active copy pin(s) — caller did not quiesce " - "the copy pipeline (bug)", - pins_.size()); - assert(pins_.empty() && "ClearLocal called with active copy pins; quiesce the pipeline first"); - } - - // Owned: defer pages with an active read lease (an RDMA read may still be in - // flight) until their lease deadline; free the rest immediately. No REMOVE - // events — the upcoming full-sync empty snapshot collapses master's index. - for (auto& [key, slot] : owned_) { - auto lease_it = read_lease_until_.find(key); - if (lease_it != read_lease_until_.end() && lease_it->second > now) { - DeferredFree df; - df.key = key; - df.tier = slot.tier; - df.pages = std::move(slot.pages); - df.release_at = lease_it->second; - deferred_frees_.push_back(std::move(df)); - continue; - } - if (auto* alloc = AllocatorForLocked(slot.tier)) { - alloc->Deallocate(slot.pages); - } - } - owned_.clear(); - - // Active read-lease deadlines that mattered are already in deferred_frees_. - read_lease_until_.clear(); - pins_.clear(); - - // Drop any queued ADD/REMOVE that the heartbeat hasn't shipped yet: - // the snapshot we're about to send is the authoritative state. - pending_events_.clear(); - - clear_full_sync_pending_.store(true, std::memory_order_release); - MORI_UMBP_INFO("[PeerDramAllocator] ClearLocal — pending writes will be rejected until ack"); -} - -void PeerDramAllocator::ClearFullSyncAcked() { - clear_full_sync_pending_.store(false, std::memory_order_release); -} - -// --------------------------------------------------------------------------- -// OwnedLocationSource (heartbeat events: queue / drain / snapshot) -// --------------------------------------------------------------------------- - -void PeerDramAllocator::QueueEventLocked(KvEvent event) { - pending_events_.push_back(std::move(event)); - // Wake the heartbeat the moment the outbox first reaches the threshold. Events - // are appended one at a time, so `==` fires at most once per batch. Called - // under mutex_: auto_flush_cb_ must be cheap and MUST NOT re-enter the - // allocator (it only signals the heartbeat thread). Exactness isn't required — - // the heartbeat interval is the backstop. threshold_ == 0 disables size-based - // auto-flush entirely (ADDs then ship only on the interval / explicit Flush). - if (auto_flush_threshold_ > 0 && pending_events_.size() == auto_flush_threshold_ && - auto_flush_cb_) { - auto_flush_cb_(); - } -} - -std::vector PeerDramAllocator::DrainPendingEvents() { - std::lock_guard lock(mutex_); - std::vector drained; - drained.swap(pending_events_); - return drained; -} - -void PeerDramAllocator::SetAutoFlushHook(size_t threshold, std::function cb) { - std::lock_guard lock(mutex_); - auto_flush_threshold_ = threshold; - auto_flush_cb_ = std::move(cb); -} - -std::vector PeerDramAllocator::SnapshotOwnedKeysLocked() const { - std::vector out; - out.reserve(owned_.size()); - for (const auto& kv : owned_) { - out.push_back(KvEvent{KvEvent::Kind::ADD, kv.first, kv.second.tier, kv.second.size}); - } - return out; -} - -// Test-only: a pure read-only snapshot with no outbox side effects. Production -// full-sync uses SnapshotOwnedKeysForFullSync() below; tests use this to assert -// owned state without disturbing the event outbox. Kept to satisfy the -// OwnedLocationSource interface and the unit tests; not on any production path. -std::vector PeerDramAllocator::SnapshotOwnedKeys() const { - std::lock_guard lock(mutex_); - return SnapshotOwnedKeysLocked(); -} - -std::vector PeerDramAllocator::SnapshotOwnedKeysForFullSync() { - std::lock_guard lock(mutex_); - auto out = SnapshotOwnedKeysLocked(); - // The snapshot is now authoritative: drop the queued delta (already reflected - // in it) so the next delta carries only new events. - pending_events_.clear(); - return out; -} - -// --------------------------------------------------------------------------- -// Capacity & buffer-descriptor queries -// --------------------------------------------------------------------------- - -std::map PeerDramAllocator::OwnedKeyCountByTier() const { - std::lock_guard lock(mutex_); - std::map result; - for (TierType t : {TierType::HBM, TierType::DRAM, TierType::SSD}) result[t] = 0; - for (const auto& [key, slot] : owned_) result[slot.tier]++; - return result; -} - -std::map PeerDramAllocator::TierCapacitiesSnapshot() const { - std::lock_guard lock(mutex_); - std::map out; - for (const auto& kv : allocators_) { - TierCapacity cap; - cap.total_bytes = kv.second->TotalBytes(); - cap.available_bytes = kv.second->AvailableBytes(); - out[kv.first] = cap; - } - return out; -} - -std::vector PeerDramAllocator::AllBufferDescs(TierType tier) const { - std::lock_guard lock(mutex_); - std::vector out; - auto it = tier_descs_.find(tier); - if (it == tier_descs_.end()) return out; - out.reserve(it->second.size()); - for (size_t i = 0; i < it->second.size(); ++i) { - out.push_back({static_cast(i), it->second[i]}); - } - return out; -} - -std::vector PeerDramAllocator::BufferDescsForPages( - TierType tier, const std::vector& pages) const { - std::lock_guard lock(mutex_); - return BuildBufferDescsLocked(tier, pages); -} - -std::vector PeerDramAllocator::BuildBufferDescsLocked( - TierType tier, const std::vector& pages) const { - std::vector out; - auto it = tier_descs_.find(tier); - if (it == tier_descs_.end()) return out; - std::vector seen; - seen.reserve(pages.size()); - for (const auto& p : pages) { - if (std::find(seen.begin(), seen.end(), p.buffer_index) != seen.end()) continue; - if (p.buffer_index >= it->second.size()) continue; // defensive: skip dangling page refs - seen.push_back(p.buffer_index); - } - std::sort(seen.begin(), seen.end()); - out.reserve(seen.size()); - for (uint32_t idx : seen) { - out.push_back({idx, it->second[idx]}); - } - return out; -} - -// --------------------------------------------------------------------------- -// Read-lease helper -// --------------------------------------------------------------------------- - -bool PeerDramAllocator::HasActiveReadLeaseLocked(const std::string& key) { - auto it = read_lease_until_.find(key); - if (it == read_lease_until_.end()) return false; - if (it->second <= std::chrono::steady_clock::now()) { - read_lease_until_.erase(it); - return false; - } - return true; -} - -// --------------------------------------------------------------------------- -// Reaper (background sweep: pending TTL, expired leases, deferred frees) -// --------------------------------------------------------------------------- - -void PeerDramAllocator::StartReaper() { - bool expected = false; - if (!reaper_running_.compare_exchange_strong(expected, true)) return; - reaper_thread_ = std::thread([this] { ReaperLoop(); }); -} - -void PeerDramAllocator::StopReaper() { - if (!reaper_running_.exchange(false)) return; - reaper_cv_.notify_all(); - if (reaper_thread_.joinable()) reaper_thread_.join(); -} - -void PeerDramAllocator::ReaperLoop() { - while (reaper_running_.load()) { - { - std::unique_lock lk(reaper_cv_mutex_); - reaper_cv_.wait_for(lk, reaper_interval_, [this] { return !reaper_running_.load(); }); - } - if (!reaper_running_.load()) break; - ReaperSweep(); - } -} - -void PeerDramAllocator::ReaperSweep() { - std::lock_guard lock(mutex_); - const auto now = std::chrono::steady_clock::now(); - - // Expire pending slots whose deadline has passed. No event is - // emitted: the slot was never owned, master never indexed it. - for (auto it = pending_.begin(); it != pending_.end();) { - if (it->second.deadline <= now) { - if (auto* alloc = AllocatorForLocked(it->second.tier)) { - alloc->Deallocate(it->second.pages); - } - MORI_UMBP_DEBUG("[PeerDramAllocator] reaped pending slot {} ({} bytes)", it->first, - it->second.size); - it = pending_.erase(it); - } else { - ++it; - } - } - - // Drop expired read leases so they stop blocking eviction and the - // map size stays bounded. - for (auto it = read_lease_until_.begin(); it != read_lease_until_.end();) { - if (it->second <= now) { - it = read_lease_until_.erase(it); - } else { - ++it; - } - } - - // Warn about copy pins held far longer than any healthy copy should - // take. We never force-free a pin (a worker may still be reading its - // segments — freeing would be a use-after-free); this is purely an - // observability signal that a copy worker is stuck. - constexpr std::chrono::seconds kLongRunningPinWarn{30}; - for (const auto& [key, pin] : pins_) { - if (now - pin.acquired_at > kLongRunningPinWarn) { - MORI_UMBP_WARN("[PeerDramAllocator] copy pin for key='{}' held >{}s — copy worker stuck?", - key, kLongRunningPinWarn.count()); - } - } - - // Release ClearLocal()-deferred pages whose lease deadline has passed. - for (auto it = deferred_frees_.begin(); it != deferred_frees_.end();) { - if (it->release_at <= now) { - if (auto* alloc = AllocatorForLocked(it->tier)) { - alloc->Deallocate(it->pages); - } - MORI_UMBP_DEBUG("[PeerDramAllocator] released deferred key='{}' pages={}", it->key, - it->pages.size()); - it = deferred_frees_.erase(it); - } else { - ++it; - } - } -} - -} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/peer_service.cpp b/src/umbp/distributed/peer/peer_service.cpp index 4cafb0c25..7e68b9813 100644 --- a/src/umbp/distributed/peer/peer_service.cpp +++ b/src/umbp/distributed/peer/peer_service.cpp @@ -23,19 +23,20 @@ #include -#include #include -#include +#include #include +#include +#include #include "mori/utils/mori_log.hpp" #include "umbp/common/env_time.h" +#include "umbp/common/grpc_limits.h" #include "umbp/distributed/master/master_client.h" #include "umbp/distributed/master/master_metrics.h" +#include "umbp/distributed/peer/backend/medium_backend.h" #include "umbp/distributed/peer/batch_resolve_codec.h" -#include "umbp/distributed/peer/peer_dram_allocator.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" -#include "umbp/distributed/peer/ssd_copy_pipeline.h" +#include "umbp/distributed/pool/peer_pool.h" #include "umbp/distributed/types.h" #include "umbp_peer.grpc.pb.h" @@ -49,58 +50,6 @@ std::chrono::seconds GrpcShutdownDeadline() { return v; } -// Free -> Preparing -> Leased. Preparing slots (IO in flight) are never -// reclaimed, so a slow IO can't have its slot reassigned mid-write. The TTL is -// anchored at request receipt (leased_at = received_at, set on promotion) to -// align the peer's reclaim point (received_at + ttl) with the reader's deadline -// (t_send + ttl, t_send < received_at). A reader trusting bytes from a slot -// reclaimed mid-use is prevented by Preparing (IO safety) + the reader's own -// lease gating (see ssd_read_lease.h). -enum class SlotState { kFree, kPreparing, kLeased }; - -struct StagingSlot { - SlotState state = SlotState::kFree; - uint64_t lease_id = 0; - size_t allocated_size = 0; - std::chrono::steady_clock::time_point leased_at; // valid only while kLeased -}; - -// Reclaim TTL-expired leased slots, then claim a free one as Preparing (TTL not -// yet started) and return its index, or -1 if none free. -int ClaimStagingSlot(std::vector& slots, std::atomic& next_lease_id, - std::chrono::milliseconds lease_timeout, StagingMetrics& metrics) { - auto now = std::chrono::steady_clock::now(); - for (auto& slot : slots) { - if (slot.state == SlotState::kLeased && now - slot.leased_at > lease_timeout) { - metrics.expired_reclaims.fetch_add(1, std::memory_order_relaxed); - MORI_UMBP_WARN("[PeerService] Reclaiming expired slot (lease_id={})", slot.lease_id); - slot.state = SlotState::kFree; - } - } - for (auto& slot : slots) { - if (slot.state == SlotState::kFree) { - slot.state = SlotState::kPreparing; - slot.lease_id = next_lease_id.fetch_add(1, std::memory_order_relaxed); - slot.allocated_size = 0; - return static_cast(&slot - &slots[0]); - } - } - return -1; -} - -bool ReleaseSlotByLeaseId(std::vector& slots, uint64_t lease_id) { - for (auto& slot : slots) { - if (slot.state == SlotState::kLeased && slot.lease_id == lease_id) { - slot.state = SlotState::kFree; - return true; - } - } - return false; -} -} // namespace - -namespace { - // Translate proto TierType <-> umbp::TierType. Defined inline because // only the peer service handlers need them. TierType FromProtoTier(::umbp::TierType t) { @@ -115,52 +64,56 @@ TierType FromProtoTier(::umbp::TierType t) { return TierType::UNKNOWN; } } -::umbp::TierType ToProtoTier(TierType t) { - switch (t) { - case TierType::HBM: - return ::umbp::TIER_HBM; - case TierType::DRAM: - return ::umbp::TIER_DRAM; - case TierType::SSD: - return ::umbp::TIER_SSD; - default: - return ::umbp::TIER_UNKNOWN; - } -} -// Map a PeerSsdManager read outcome (data-level: ok/not_found/size/error) to the -// wire status. NO_SLOT is a staging-layer concern owned by the peer service, -// not PeerSsdManager, and is set directly by the handler. Lease expiry is not -// a wire status: the reader decides it locally and the peer reclaims by TTL. -::umbp::SsdReadStatus ToProtoReadStatus(SsdReadStatus s) { - switch (s) { - case SsdReadStatus::kOk: - return ::umbp::SSD_READ_OK; - case SsdReadStatus::kNotFound: - return ::umbp::SSD_READ_NOT_FOUND; - case SsdReadStatus::kSizeTooLarge: - return ::umbp::SSD_READ_SIZE_TOO_LARGE; - case SsdReadStatus::kError: - return ::umbp::SSD_READ_ERROR; - } - return ::umbp::SSD_READ_ERROR; +// --------------------------------------------------------------------------- +// Slot-id backend tagging +// +// Commit / Abort carry only a slot_id — the proto has no tier field on them, +// and umbp_peer.proto documents the id as "opaque; echoed back by Commit / +// Abort". Every backend numbers its own slots from 1 independently, so a bare +// id is ambiguous the moment a second medium is live. +// +// So the peer service tags the backend instance id into the high byte on the +// way out and strips it on the way back in. TierType is not sufficient: a +// topology may contain ssd_a and ssd_b. The tag is peer-local — PoolClient's +// local fast path talks to a backend directly and never sees a tagged id. +// +// next_slot_id_ is a per-backend counter starting at 1, so the low 56 bits +// cannot reach the tag in any realistic process lifetime. +// --------------------------------------------------------------------------- +constexpr int kSlotBackendShift = 56; +constexpr uint64_t kSlotLocalMask = (1ULL << kSlotBackendShift) - 1; + +uint64_t TagSlotId(uint32_t backend_id, uint64_t local_id) { + return (static_cast(backend_id) << kSlotBackendShift) | (local_id & kSlotLocalMask); } +uint32_t BackendIdFromSlotId(uint64_t tagged) { + return static_cast(tagged >> kSlotBackendShift); +} +uint64_t LocalSlotId(uint64_t tagged) { return tagged & kSlotLocalMask; } // Drop a (pages, page_size, descs) tuple into a slot-shaped response // that exposes those fields directly. Templated so the same body // covers AllocateSlotResponse and ResolveKeyResponse. +// +// `backend_id` is stamped here rather than inside the backend: a backend +// numbers its buffers from 0 and knows nothing of its siblings, so the owning +// backend is a fact only this boundary has. See BufferMemoryDesc in +// umbp.proto. template void FillPagesAndDescs(Response* resp, const std::vector& pages, uint64_t page_size, - const std::vector& descs) { + const std::vector& descs, uint32_t backend_id) { for (const auto& p : pages) { auto* pl = resp->add_pages(); pl->set_buffer_index(p.buffer_index); pl->set_page_index(p.page_index); } resp->set_page_size(page_size); + resp->set_backend_id(backend_id); for (const auto& d : descs) { auto* desc = resp->add_descs(); desc->set_buffer_index(d.buffer_index); + desc->set_backend_id(backend_id); desc->set_desc(std::string(d.desc_bytes.begin(), d.desc_bytes.end())); } } @@ -169,142 +122,53 @@ void FillPagesAndDescs(Response* resp, const std::vector& pages, u class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Service { public: - UMBPPeerServiceImpl(void* ssd_staging_base, size_t ssd_staging_size, - const std::vector& ssd_staging_mem_desc_bytes, - PeerSsdManager* peer_ssd, PeerDramAllocator* dram_alloc, - MasterClient* master_client, StagingMetrics& metrics, int num_read_slots, - std::chrono::milliseconds lease_timeout, - const std::vector& engine_desc_bytes, SsdCopyPipeline* copy_pipeline) - : ssd_staging_base_(ssd_staging_base), - ssd_staging_size_(ssd_staging_size), - ssd_staging_mem_desc_bytes_(ssd_staging_mem_desc_bytes), - peer_ssd_(peer_ssd), - dram_alloc_(dram_alloc), - copy_pipeline_(copy_pipeline), + UMBPPeerServiceImpl(PeerPool* pool, MasterClient* master_client, + const std::vector& engine_desc_bytes) + : pool_(pool), + registry_(pool == nullptr ? nullptr : pool->Backends()), + // All() allocates, and Resolve is the hot read path — so the medium + // list is snapshotted once here rather than rebuilt per RPC. This is + // why the registry MUST be fully populated before the peer service is + // constructed (PoolClient::Init registers every backend first, then + // starts this server). + media_(registry_ == nullptr ? std::vector{} : registry_->All()), engine_desc_bytes_(engine_desc_bytes), - metrics_(metrics), - lease_timeout_(std::max(lease_timeout, std::chrono::milliseconds{1})), - num_read_slots_(std::max(num_read_slots, 1)), - // The whole staging buffer is the read region now: the lower half used - // to be direct-put write staging, removed in the SSD-tier redesign and - // reclaimed here so reads get more / larger slots. - read_region_base_(0), - read_slot_size_(ssd_staging_size / static_cast(std::max(num_read_slots, 1))), - read_slots_(std::max(num_read_slots, 1)), - master_client_(master_client) { - if (num_read_slots <= 0) { - MORI_UMBP_ERROR("[PeerService] num_read_slots={} invalid, clamped to 1", num_read_slots); - } - } + master_client_(master_client) {} grpc::Status GetPeerInfo(grpc::ServerContext* /*context*/, const ::umbp::GetPeerInfoRequest* /*request*/, ::umbp::GetPeerInfoResponse* response) override { - response->set_ssd_staging_mem_desc( - std::string(ssd_staging_mem_desc_bytes_.begin(), ssd_staging_mem_desc_bytes_.end())); - response->set_ssd_staging_size(ssd_staging_size_); if (!engine_desc_bytes_.empty()) { response->set_engine_desc(std::string(engine_desc_bytes_.begin(), engine_desc_bytes_.end())); } - if (dram_alloc_ != nullptr) { - // Ship every configured DRAM/HBM buffer's desc so first-contact - // writers can hydrate without a follow-up Allocate / Resolve. - // DRAM and HBM share a single page_size in this design, so the - // single field on the response is sufficient. - auto dram_descs = dram_alloc_->AllBufferDescs(TierType::DRAM); - auto hbm_descs = dram_alloc_->AllBufferDescs(TierType::HBM); - for (const auto& d : dram_descs) { - auto* out = response->add_dram_memory_descs(); + // Ship EVERY backend's buffer descs so first-contact writers can hydrate + // without a follow-up Allocate / Resolve. + // + // This used to publish one medium and log an error about the rest, because + // the wire carried a flat buffer_index space with no way to say which + // backend an index belonged to. That was not merely incomplete: the reader + // caches descriptors by index and asks the peer to omit them once it + // believes it has them (BatchResolveKeysRequest.omit_descs), so an + // unadvertised medium's pages were resolved against the advertised + // medium's memory — a hit, with bytes from the wrong pool. Every backend + // publishes exactly one buffer today, so index 0 collided on every + // mixed-media node and the corruption was deterministic. + // + // BufferMemoryDesc.backend_id closes it. buffer_index stays backend-local + // and no backend changed; the id is stamped here, at the boundary that + // knows it. + for (auto* backend : Media()) { + const uint32_t backend_id = registry_->BackendId(backend); + for (const auto& d : backend->AllBufferDescs()) { + auto* out = response->add_buffer_descs(); + out->set_backend_id(backend_id); out->set_buffer_index(d.buffer_index); out->set_desc(std::string(d.desc_bytes.begin(), d.desc_bytes.end())); } - for (const auto& d : hbm_descs) { - auto* out = response->add_dram_memory_descs(); - out->set_buffer_index(d.buffer_index); - out->set_desc(std::string(d.desc_bytes.begin(), d.desc_bytes.end())); - } - response->set_dram_page_size(dram_alloc_->PageSize()); - } - return grpc::Status::OK; - } - - // ---- SSD read staging: PrepareSsdRead + ReleaseSsdLease ---- - // Key-based: claim a slot -> PeerSsdManager::PrepareRead fills it -> reader - // RDMAs the bytes out of the published staging buffer -> best-effort release. - grpc::Status PrepareSsdRead(grpc::ServerContext* /*context*/, - const ::umbp::PrepareSsdReadRequest* request, - ::umbp::PrepareSsdReadResponse* response) override { - if (!SsdRpcAvailable()) { - response->set_status(::umbp::SSD_READ_ERROR); - return grpc::Status::OK; - } - - // Anchor the lease TTL at request receipt (see SlotState comment): the slot - // is promoted to Leased with this timestamp once the IO completes, so the - // peer's reclaim point stays aligned with the reader's t_send-based - // deadline rather than starting only after the (variable) SSD IO. - const auto received_at = std::chrono::steady_clock::now(); - - // Claim a slot first (Preparing — not yet TTL-tracked). NO_SLOT is a - // transient/retryable condition, not a miss. - int slot_idx; - uint64_t offset, lease_id; - { - std::lock_guard lock(read_slots_mutex_); - slot_idx = ClaimStagingSlot(read_slots_, next_lease_id_, lease_timeout_, metrics_); - if (slot_idx < 0) { - metrics_.slot_full_rejects.fetch_add(1, std::memory_order_relaxed); - MORI_UMBP_WARN("[PeerService] PrepareSsdRead: no free staging slots"); - response->set_status(::umbp::SSD_READ_NO_SLOT); - return grpc::Status::OK; - } - offset = read_region_base_ + static_cast(slot_idx) * read_slot_size_; - lease_id = read_slots_[slot_idx].lease_id; } - - // Cap the read at min(reader capacity, slot size); PrepareRead rejects an - // over-cap key BEFORE doing any SSD IO, so an oversized key costs no read. - void* dst = static_cast(ssd_staging_base_) + offset; - size_t cap = std::min(request->max_size(), read_slot_size_); - SsdReadOutcome outcome = peer_ssd_->PrepareRead(request->key(), dst, cap); - - if (outcome.status != SsdReadStatus::kOk) { - std::lock_guard lock(read_slots_mutex_); - read_slots_[slot_idx].state = SlotState::kFree; // give the slot straight back - response->set_status(ToProtoReadStatus(outcome.status)); - response->set_size(outcome.size); - return grpc::Status::OK; - } - - // Data is in the slot: promote Preparing -> Leased with the request-receipt - // TTL anchor (leased_at = received_at). - { - std::lock_guard lock(read_slots_mutex_); - auto& slot = read_slots_[slot_idx]; - slot.state = SlotState::kLeased; - slot.leased_at = received_at; - slot.allocated_size = outcome.size; - } - response->set_status(::umbp::SSD_READ_OK); - response->set_staging_offset(offset); - response->set_size(outcome.size); - response->set_lease_id(lease_id); - response->set_lease_ttl_ms(static_cast( - std::chrono::duration_cast(lease_timeout_).count())); - MORI_UMBP_DEBUG( - "[PeerService] PrepareSsdRead: key={}, slot={}, offset={}, size={}, lease_id={}", - request->key(), slot_idx, offset, outcome.size, lease_id); - return grpc::Status::OK; - } - - grpc::Status ReleaseSsdLease(grpc::ServerContext* /*context*/, - const ::umbp::ReleaseSsdLeaseRequest* request, - ::umbp::ReleaseSsdLeaseResponse* response) override { - // Best-effort fast release; correctness does not depend on it (the slot is - // also reclaimed by the lease TTL). Returns false only when the lease is - // already gone (released earlier or TTL-reclaimed). - std::lock_guard lock(read_slots_mutex_); - response->set_success(ReleaseSlotByLeaseId(read_slots_, request->lease_id())); + // Uniform across backends; BackendRegistry::Register refuses a backend that + // disagrees, so there is one page size to report. + if (registry_ != nullptr) response->set_page_size(registry_->PageSize()); return grpc::Status::OK; } @@ -312,91 +176,108 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se // DRAM/HBM allocator + key map (master-as-advisor design) // ============================================================ + // The single-key RPCs below are served by one-element batches (design doc §3) + // — there is no separate single-key path on MediumBackend to keep in sync. + grpc::Status AllocateSlot(grpc::ServerContext* /*ctx*/, const ::umbp::AllocateSlotRequest* request, ::umbp::AllocateSlotResponse* response) override { - if (dram_alloc_ == nullptr) { + if (pool_ == nullptr) { response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); return grpc::Status::OK; } - auto result = - dram_alloc_->Allocate(request->key(), request->size(), FromProtoTier(request->tier())); + PoolPlacementRequest entry; + entry.key = request->key(); + entry.size = request->size(); + entry.tier = FromProtoTier(request->tier()); + entry.backend_name = request->backend_name(); + entry.logical_tier = request->logical_tier(); + auto pool_result = pool_->BatchAllocate({entry}).front(); + auto* backend = Backend(pool_result.backend_id); + const auto& result = pool_result.allocation; switch (result.outcome) { - case PeerDramAllocator::Outcome::kSuccessAlreadyExists: + case AllocateOutcome::kSuccessAlreadyExists: response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALREADY_EXISTS); return grpc::Status::OK; - case PeerDramAllocator::Outcome::kFailed: + case AllocateOutcome::kFailed: response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); return grpc::Status::OK; - case PeerDramAllocator::Outcome::kFailedNoSpace: + case AllocateOutcome::kFailedNoSpace: response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED_NO_SPACE); return grpc::Status::OK; - case PeerDramAllocator::Outcome::kSuccessAllocated: + case AllocateOutcome::kSuccessAllocated: break; } - const auto& pending = *result.slot; - auto descs = dram_alloc_->BufferDescsForPages(pending.tier, pending.pages); + if (backend == nullptr) { + response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); + return grpc::Status::OK; + } response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); - response->set_slot_id(pending.slot_id); - FillPagesAndDescs(response, pending.pages, dram_alloc_->PageSize(), descs); - response->set_pending_ttl_ms(dram_alloc_->PendingTtlMs()); + response->set_slot_id(TagSlotId(pool_result.backend_id, result.slot_id)); + FillPagesAndDescs(response, result.pages, result.page_size, result.descs, + pool_result.backend_id); + response->set_pending_ttl_ms(result.pending_ttl_ms); return grpc::Status::OK; } grpc::Status CommitSlot(grpc::ServerContext* /*ctx*/, const ::umbp::CommitSlotRequest* request, ::umbp::CommitSlotResponse* response) override { - if (dram_alloc_ == nullptr) { + if (pool_ == nullptr) { response->set_success(false); return grpc::Status::OK; } - uint64_t committed_bytes = 0; - const bool ok = dram_alloc_->Commit(request->slot_id(), request->key(), committed_bytes); - response->set_success(ok); - if (ok) { - RecordInboundPut(committed_bytes, "remote"); - EnqueueSsdCopy(request->key(), committed_bytes); + PoolCommitRequest entry; + entry.slot.backend_id = BackendIdFromSlotId(request->slot_id()); + entry.slot.local_slot_id = LocalSlotId(request->slot_id()); + entry.key = request->key(); + const auto result = pool_->BatchCommit({entry}).front().commit; + response->set_success(result.success); + if (result.success) { + RecordInboundPut(result.bytes_committed, "remote"); } return grpc::Status::OK; } grpc::Status AbortSlot(grpc::ServerContext* /*ctx*/, const ::umbp::AbortSlotRequest* request, ::umbp::AbortSlotResponse* response) override { - if (dram_alloc_ == nullptr) { + if (pool_ == nullptr) { response->set_success(true); // idempotent: nothing to drop return grpc::Status::OK; } - response->set_success(dram_alloc_->Abort(request->slot_id())); + PoolSlotRef slot{BackendIdFromSlotId(request->slot_id()), LocalSlotId(request->slot_id())}; + auto results = pool_->BatchAbort({slot}); + response->set_success(results.front()); return grpc::Status::OK; } grpc::Status ResolveKey(grpc::ServerContext* /*ctx*/, const ::umbp::ResolveKeyRequest* request, ::umbp::ResolveKeyResponse* response) override { - if (dram_alloc_ == nullptr) { - response->set_found(false); - return grpc::Status::OK; - } - auto r = dram_alloc_->Resolve(request->key()); - response->set_found(r.found); - if (!r.found) return grpc::Status::OK; - auto descs = dram_alloc_->BufferDescsForPages(r.tier, r.pages); - FillPagesAndDescs(response, r.pages, dram_alloc_->PageSize(), descs); - response->set_size(r.size); - RecordInboundGet(r.size, "remote"); + if (pool_ == nullptr) return grpc::Status::OK; + auto result = pool_->BatchResolve({request->key()}, /*include_descs=*/true).front(); + if (!result.resolved.found) return grpc::Status::OK; + response->set_found(true); + FillPagesAndDescs(response, result.resolved.pages, result.resolved.page_size, + result.resolved.descs, result.backend_id); + response->set_size(result.resolved.size); + RecordInboundGet(result.resolved.size, "remote"); return grpc::Status::OK; } grpc::Status EvictKey(grpc::ServerContext* /*ctx*/, const ::umbp::EvictKeyRequest* request, ::umbp::EvictKeyResponse* response) override { - if (dram_alloc_ == nullptr) { - // No DRAM/HBM tier on this peer — nothing to evict, treat as success. - return grpc::Status::OK; - } + // Eviction carries no tier either. Master drives this to reclaim capacity + // it measured per (node, tier), so the pool is free to demote where a tier + // configures on_evict offload; otherwise a key mirrored across media is + // dropped from ALL of them. Either way the freed bytes are summed per key + // — master sizes its next eviction round off this total. std::vector keys(request->keys().begin(), request->keys().end()); - auto results = dram_alloc_->Evict(keys); - for (const auto& r : results) { + if (keys.empty()) return grpc::Status::OK; + auto evicted = + pool_ == nullptr ? std::vector{} : pool_->Evict(keys, PoolEvictMode::kReclaim); + for (size_t i = 0; i < keys.size(); ++i) { auto* entry = response->add_evicted(); - entry->set_key(r.key); - entry->set_bytes_freed(r.bytes_freed); + entry->set_key(keys[i]); + entry->set_bytes_freed(i < evicted.size() ? evicted[i].bytes_freed : 0); } return grpc::Status::OK; } @@ -406,45 +287,44 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se grpc::Status BatchAllocateSlots(grpc::ServerContext* /*ctx*/, const ::umbp::BatchAllocateSlotsRequest* request, ::umbp::BatchAllocateSlotsResponse* response) override { - if (dram_alloc_ == nullptr) { - for (int i = 0; i < request->entries_size(); ++i) { - auto* out = response->add_entries(); - out->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); - } - return grpc::Status::OK; + const int n = request->entries_size(); + for (int i = 0; i < n; ++i) { + response->add_entries()->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); } - - std::vector entries; - entries.reserve(request->entries_size()); - for (const auto& entry : request->entries()) { - PeerDramAllocator::AllocateRequest alloc_entry; - alloc_entry.key = entry.key(); - alloc_entry.size = entry.size(); - alloc_entry.tier = FromProtoTier(entry.tier()); - entries.push_back(std::move(alloc_entry)); + if (pool_ == nullptr) return grpc::Status::OK; + + std::vector entries; + entries.reserve(n); + for (int i = 0; i < n; ++i) { + PoolPlacementRequest entry; + entry.key = request->entries(i).key(); + entry.size = request->entries(i).size(); + entry.tier = FromProtoTier(request->entries(i).tier()); + entry.backend_name = request->entries(i).backend_name(); + entry.logical_tier = request->entries(i).logical_tier(); + entries.push_back(std::move(entry)); } - - auto results = dram_alloc_->BatchAllocate(entries); - for (const auto& result : results) { - auto* out = response->add_entries(); + auto results = pool_->BatchAllocate(entries); + for (size_t i = 0; i < results.size() && i < static_cast(n); ++i) { + const auto& result = results[i].allocation; + auto* out = response->mutable_entries(static_cast(i)); switch (result.outcome) { - case PeerDramAllocator::Outcome::kSuccessAlreadyExists: + case AllocateOutcome::kSuccessAlreadyExists: out->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALREADY_EXISTS); continue; - case PeerDramAllocator::Outcome::kFailed: + case AllocateOutcome::kFailed: out->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); continue; - case PeerDramAllocator::Outcome::kFailedNoSpace: + case AllocateOutcome::kFailedNoSpace: out->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED_NO_SPACE); continue; - case PeerDramAllocator::Outcome::kSuccessAllocated: + case AllocateOutcome::kSuccessAllocated: break; } - const auto& pending = *result.slot; out->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); - out->set_slot_id(pending.slot_id); - FillPagesAndDescs(out, pending.pages, dram_alloc_->PageSize(), result.descs); - out->set_pending_ttl_ms(dram_alloc_->PendingTtlMs()); + out->set_slot_id(TagSlotId(results[i].backend_id, result.slot_id)); + FillPagesAndDescs(out, result.pages, result.page_size, result.descs, results[i].backend_id); + out->set_pending_ttl_ms(result.pending_ttl_ms); } return grpc::Status::OK; } @@ -452,29 +332,25 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se grpc::Status BatchCommitSlots(grpc::ServerContext* /*ctx*/, const ::umbp::BatchCommitSlotsRequest* request, ::umbp::BatchCommitSlotsResponse* response) override { - if (dram_alloc_ == nullptr) { - for (int i = 0; i < request->entries_size(); ++i) response->add_success(false); - return grpc::Status::OK; - } - - std::vector entries; - entries.reserve(request->entries_size()); - for (const auto& entry : request->entries()) { - PeerDramAllocator::CommitRequest commit_entry; - commit_entry.slot_id = entry.slot_id(); - commit_entry.key = entry.key(); - entries.push_back(std::move(commit_entry)); + const int n = request->entries_size(); + for (int i = 0; i < n; ++i) response->add_success(false); + if (pool_ == nullptr) return grpc::Status::OK; + + std::vector entries; + entries.reserve(n); + for (int i = 0; i < n; ++i) { + PoolCommitRequest entry; + entry.slot.backend_id = BackendIdFromSlotId(request->entries(i).slot_id()); + entry.slot.local_slot_id = LocalSlotId(request->entries(i).slot_id()); + entry.key = request->entries(i).key(); + entries.push_back(std::move(entry)); } - auto results = dram_alloc_->BatchCommit(entries); uint64_t total_committed = 0; - for (size_t i = 0; i < results.size(); ++i) { - const auto& result = results[i]; - response->add_success(result.success); - if (result.success) { - total_committed += result.bytes_committed; - EnqueueSsdCopy(entries[i].key, result.bytes_committed); - } + auto results = pool_->BatchCommit(entries); + for (size_t i = 0; i < results.size() && i < static_cast(n); ++i) { + response->set_success(static_cast(i), results[i].commit.success); + if (results[i].commit.success) total_committed += results[i].commit.bytes_committed; } if (total_committed > 0) RecordInboundPut(total_committed, "remote"); return grpc::Status::OK; @@ -483,15 +359,22 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se grpc::Status BatchAbortSlots(grpc::ServerContext* /*ctx*/, const ::umbp::BatchAbortSlotsRequest* request, ::umbp::BatchAbortSlotsResponse* response) override { - if (dram_alloc_ == nullptr) { - for (int i = 0; i < request->slot_ids_size(); ++i) response->add_success(true); - return grpc::Status::OK; + const int n = request->slot_ids_size(); + // Abort is idempotent: an unknown slot (including one whose backend is no + // longer live) reports true — there is nothing left to drop. + for (int i = 0; i < n; ++i) response->add_success(true); + if (pool_ == nullptr) return grpc::Status::OK; + + std::vector slots; + slots.reserve(n); + for (int i = 0; i < n; ++i) { + slots.push_back(PoolSlotRef{BackendIdFromSlotId(request->slot_ids(i)), + LocalSlotId(request->slot_ids(i))}); } - std::vector slot_ids(request->slot_ids().begin(), request->slot_ids().end()); - auto results = dram_alloc_->BatchAbort(slot_ids); - for (bool ok : results) { - response->add_success(ok); + auto results = pool_->BatchAbort(slots); + for (size_t i = 0; i < results.size() && i < static_cast(n); ++i) { + response->set_success(static_cast(i), results[i]); } return grpc::Status::OK; } @@ -499,61 +382,62 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se grpc::Status BatchResolveKeys(grpc::ServerContext* /*ctx*/, const ::umbp::BatchResolveKeysRequest* request, ::umbp::BatchResolveKeysResponse* response) override { - if (dram_alloc_ == nullptr) { - std::vector misses(request->keys_size()); - EncodeBatchResolveResponse(misses, /*page_size=*/0, /*descs=*/{}, response); - return grpc::Status::OK; - } // The client can suppress descs entirely via omit_descs when it already - // hydrated them from GetPeerInfo — skip BuildBufferDescsLocked for the - // whole batch in that case rather than computing and discarding it. + // hydrated them from GetPeerInfo — skip building them for the whole batch + // in that case rather than computing and discarding them. const bool omit_descs = request->omit_descs(); std::vector keys(request->keys().begin(), request->keys().end()); - auto resolved = dram_alloc_->BatchResolve(keys, /*include_descs=*/!omit_descs); - // Project to the encoder's host shape and deduplicate the buffer - // descriptors once for the whole batch (their bytes are identical across - // keys for a given buffer_index). - std::vector entries; - entries.reserve(resolved.size()); + std::vector entries(keys.size()); std::vector batch_descs; - std::vector desc_seen; // indexed by buffer_index + std::vector> desc_seen(BackendRegistry::kMaxBackends); uint64_t total_bytes = 0; - for (auto& r : resolved) { - ResolvedKeyEntry e; - e.found = r.found; - if (r.found) { - e.tier = r.tier; - e.size = r.size; - e.pages = std::move(r.pages); - total_bytes += r.size; - if (!omit_descs) { - for (const auto& d : r.descs) { - if (d.buffer_index >= desc_seen.size()) desc_seen.resize(d.buffer_index + 1, false); - if (desc_seen[d.buffer_index]) continue; - desc_seen[d.buffer_index] = true; - batch_descs.push_back(d); - } - } + auto resolved = pool_ == nullptr ? std::vector{} + : pool_->BatchResolve(keys, /*include_descs=*/!omit_descs); + for (size_t i = 0; i < resolved.size() && i < entries.size(); ++i) { + auto& result = resolved[i]; + auto& r = result.resolved; + entries[i].outcome = EffectiveResolveOutcome(r); + if (result.backend_id < BackendRegistry::kMaxBackends) { + entries[i].tier = result.tier; + entries[i].backend_id = result.backend_id; + } + if (!r.found || result.backend_id >= BackendRegistry::kMaxBackends) continue; + entries[i].found = true; + entries[i].outcome = ResolveOutcome::kFound; + entries[i].tier = result.tier; + entries[i].backend_id = result.backend_id; + entries[i].size = r.size; + entries[i].pages = std::move(r.pages); + total_bytes += r.size; + + if (omit_descs) continue; + auto& seen = desc_seen[result.backend_id]; + for (auto& d : r.descs) { + if (d.buffer_index >= seen.size()) seen.resize(d.buffer_index + 1, false); + if (seen[d.buffer_index]) continue; + seen[d.buffer_index] = true; + d.backend_id = result.backend_id; + batch_descs.push_back(std::move(d)); } - entries.push_back(std::move(e)); } - EncodeBatchResolveResponse(entries, dram_alloc_->PageSize(), batch_descs, response); + + // Page size is uniform across backends (BackendRegistry::Register), so a + // batch spanning media still has one to report. + EncodeBatchResolveResponse(entries, registry_ == nullptr ? 0 : registry_->PageSize(), + batch_descs, response); RecordInboundGet(total_bytes, "remote"); return grpc::Status::OK; } private: - bool SsdRpcAvailable() const { - return peer_ssd_ != nullptr && ssd_staging_base_ != nullptr && ssd_staging_size_ > 0; + MediumBackend* Backend(uint32_t backend_id) const { + return registry_ == nullptr ? nullptr : registry_->Get(backend_id); } - // Best-effort enqueue of an async SSD copy after a successful DRAM commit. - // No-op when SSD is disabled; never blocks (a full/stopped queue drops). - void EnqueueSsdCopy(const std::string& key, uint64_t bytes) { - if (copy_pipeline_ == nullptr) return; - copy_pipeline_->Enqueue(SsdCopyTask{key, TierType::DRAM, static_cast(bytes)}); - } + // Every backend on this peer, snapshotted at construction for bootstrap + // descriptor publication. Placement/read order lives in PoolPolicy. + const std::vector& Media() const { return media_; } void RecordInboundPut(uint64_t bytes, const char* traffic) { if (master_client_ == nullptr || bytes == 0) return; @@ -571,69 +455,72 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se static_cast(bytes)); } - void* ssd_staging_base_; - size_t ssd_staging_size_; - const std::vector& ssd_staging_mem_desc_bytes_; - PeerSsdManager* peer_ssd_; - PeerDramAllocator* dram_alloc_; - SsdCopyPipeline* copy_pipeline_; - MasterClient* master_client_; + PeerPool* pool_; + BackendRegistry* registry_; + const std::vector media_; const std::vector& engine_desc_bytes_; - StagingMetrics& metrics_; - - const std::chrono::milliseconds lease_timeout_; - const int num_read_slots_; - const uint64_t read_region_base_; - const size_t read_slot_size_; - - std::mutex read_slots_mutex_; - std::vector read_slots_; - std::atomic next_lease_id_{1}; - - public: - // SSD read staging slots currently busy (Preparing or Leased). Sampled once - // per metrics flush by PoolClient's provider; a brief lock + small scan - // (num_read_slots, default 16) — not a busy loop. - size_t ReadSlotsInUse() { - std::lock_guard lock(read_slots_mutex_); - size_t in_use = 0; - for (const auto& slot : read_slots_) { - if (slot.state != SlotState::kFree) ++in_use; - } - return in_use; - } + MasterClient* master_client_; }; -PeerServiceServer::PeerServiceServer(PeerDramAllocator* dram_alloc, PeerSsdManager* peer_ssd, - void* ssd_staging_base, size_t ssd_staging_size, - std::vector ssd_staging_mem_desc_bytes, - int num_read_slots, std::chrono::milliseconds lease_timeout, +PeerServiceServer::PeerServiceServer(PeerPool& pool, std::vector engine_desc_bytes, + MasterClient* master_client) + : pool_(&pool), + master_client_(master_client), + engine_desc_bytes_(std::move(engine_desc_bytes)) { + service_ = std::make_unique(pool_, master_client_, engine_desc_bytes_); +} + +PeerServiceServer::PeerServiceServer(BackendRegistry* registry, std::vector engine_desc_bytes, - MasterClient* master_client, SsdCopyPipeline* copy_pipeline) - : ssd_staging_base_(ssd_staging_base), - ssd_staging_size_(ssd_staging_size), - peer_ssd_(peer_ssd), - dram_alloc_(dram_alloc), + MasterClient* master_client) + : owned_pool_(registry == nullptr + ? nullptr + : std::make_unique(registry, MakeSingleBackendPolicy())), + pool_(owned_pool_.get()), master_client_(master_client), - copy_pipeline_(copy_pipeline), - ssd_staging_mem_desc_bytes_(std::move(ssd_staging_mem_desc_bytes)), engine_desc_bytes_(std::move(engine_desc_bytes)) { - service_ = std::make_unique( - ssd_staging_base_, ssd_staging_size_, ssd_staging_mem_desc_bytes_, peer_ssd_, dram_alloc_, - master_client_, metrics_, num_read_slots, lease_timeout, engine_desc_bytes_, copy_pipeline_); + service_ = std::make_unique(pool_, master_client_, engine_desc_bytes_); } PeerServiceServer::~PeerServiceServer() { Stop(); } -size_t PeerServiceServer::SnapshotReadSlotsInUse() const { - return service_ ? service_->ReadSlotsInUse() : 0; -} - bool PeerServiceServer::Start(uint16_t port) { std::string address = "0.0.0.0:" + std::to_string(port); grpc::ServerBuilder builder; - builder.AddListeningPort(address, grpc::InsecureServerCredentials()); + ApplyGrpcLimits(&builder); + // gRPC turns SO_REUSEPORT ON by default for TCP servers on Linux, which means + // a SECOND server binding this same port SUCCEEDS and the kernel then splits + // incoming connections between the two. For a peer service that is never + // right: a client dialing this address would be answered, some fraction of + // the time, by a different process's peer service with different backends + // registered — a wrong answer rather than a connection error. It also made + // the "port may be in use" check below dead code, since BuildAndStart could + // not fail that way. Exactly one peer service owns a port. + builder.AddChannelArgument(GRPC_ARG_ALLOW_REUSEPORT, 0); + + // gRPC's default sync poller pool is small, and every peer RPC that stages + // SSD bytes occupies one of those threads for the whole device read. On a + // loopback read (reader and owner are the same node) the default pool is the + // binding constraint rather than the drive — the same head-of-line blocking + // the master server widens its pool for. Tunable for large deployments. + { + auto env_pollers = [](const char* name, int def) -> int { + const char* v = std::getenv(name); + if (v == nullptr) return def; + char* end = nullptr; + long n = std::strtol(v, &end, 10); + return (end == v || n <= 0) ? def : static_cast(n); + }; + const int min_pollers = env_pollers("UMBP_PEER_MIN_POLLERS", 8); + int max_pollers = env_pollers("UMBP_PEER_MAX_POLLERS", 64); + if (max_pollers < min_pollers) max_pollers = min_pollers; + builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::MIN_POLLERS, min_pollers); + builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::MAX_POLLERS, max_pollers); + } + + int selected_port = 0; + builder.AddListeningPort(address, grpc::InsecureServerCredentials(), &selected_port); builder.RegisterService(service_.get()); server_ = builder.BuildAndStart(); @@ -641,7 +528,15 @@ bool PeerServiceServer::Start(uint16_t port) { MORI_UMBP_ERROR("[PeerService] Failed to start on {} (port may be in use)", address); return false; } - MORI_UMBP_INFO("[PeerService] Listening on {}", address); + if (selected_port <= 0 || selected_port > std::numeric_limits::max()) { + MORI_UMBP_ERROR("[PeerService] invalid selected port {} for {}", selected_port, address); + server_->Shutdown(); + server_->Wait(); + server_.reset(); + return false; + } + bound_port_ = static_cast(selected_port); + MORI_UMBP_INFO("[PeerService] Listening on 0.0.0.0:{}", bound_port_); return true; } @@ -652,10 +547,11 @@ void PeerServiceServer::Stop() { server_->Shutdown(deadline); // Block until every in-flight handler has returned (Shutdown's deadline // force-cancels any that overrun). This guarantees no RPC handler is still - // touching borrowed state (dram_alloc_ / copy_pipeline_) after Stop() - // returns, so PoolClient can safely tear those down next. + // touching borrowed state (pool_ and its backends) after Stop() + // returns, so PoolClient can safely tear it down next. server_->Wait(); server_.reset(); + bound_port_ = 0; } } diff --git a/src/umbp/distributed/peer/peer_ssd_manager.cpp b/src/umbp/distributed/peer/peer_ssd_manager.cpp deleted file mode 100644 index 8a8229b73..000000000 --- a/src/umbp/distributed/peer/peer_ssd_manager.cpp +++ /dev/null @@ -1,439 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/distributed/peer/peer_ssd_manager.h" - -#include -#include - -#include "mori/utils/mori_log.hpp" -#include "umbp/local/tiers/spdk_proxy_tier.h" -#include "umbp/local/tiers/ssd_tier.h" -#include "umbp/local/tiers/tier_backend.h" - -namespace mori::umbp { - -namespace { -// Watermarks must satisfy 0 < low < high <= 1. We fail fast on a bad value -// rather than silently clamping, so a misconfigured env surfaces immediately. -bool WatermarksValid(double high, double low) { - return high > 0.0 && high <= 1.0 && low > 0.0 && low < high; -} -} // namespace - -// --------------------------------------------------------------------------- -// Construction -// --------------------------------------------------------------------------- - -PeerSsdManager::PeerSsdManager(const PeerSsdConfig& cfg) { - if (!cfg.enabled) { - MORI_UMBP_INFO("[PeerSsdManager] constructed disabled (no SSD backend)"); - return; - } - const auto& ssd = cfg.ssd; - - if (!WatermarksValid(ssd.high_watermark, ssd.low_watermark)) { - throw std::runtime_error( - "[PeerSsdManager] invalid SSD watermarks (require 0 < low_watermark < " - "high_watermark <= 1)"); - } - high_watermark_ = ssd.high_watermark; - low_watermark_ = ssd.low_watermark; - - // Explicit backend selection — an unknown value is a configuration error, not - // a reason to silently fall back to the file backend. - if (ssd.ssd_backend == "file") { - backend_ = std::make_unique(ssd.storage_dir, ssd.capacity_bytes, ssd); - MORI_UMBP_INFO("[PeerSsdManager] SSDTier ready dir={} capacity={}B", ssd.storage_dir, - ssd.capacity_bytes); - return; - } - - // The distributed peer shares one physical device across processes, so SPDK - // is reached through the spdk_proxy daemon, never the single-process direct - // SpdkSsdTier. Both "spdk" and "spdk_proxy" therefore map to SpdkProxyTier: - // "spdk" here means "use SPDK via the proxy". - if (ssd.ssd_backend == "spdk" || ssd.ssd_backend == "spdk_proxy") { - MORI_UMBP_INFO("[PeerSsdManager] distributed ssd_backend={} uses SpdkProxyTier", - ssd.ssd_backend); - auto proxy = std::make_unique(ssd); - // PeerSsdManager only connects to an already-ready proxy; it does not spawn - // the daemon. A connect failure with an explicit SPDK config is fatal — we - // must not silently disable SSD or fall back to POSIX. - if (!proxy->IsValid()) { - throw std::runtime_error("[PeerSsdManager] ssd_backend=" + ssd.ssd_backend + - ": SPDK proxy connect failed (shm=" + ssd.spdk_proxy_shm_name + - "). Ensure the spdk_proxy daemon is running and READY."); - } - backend_ = std::move(proxy); - MORI_UMBP_INFO("[PeerSsdManager] SpdkProxyTier ready (shm={})", ssd.spdk_proxy_shm_name); - return; - } - - throw std::runtime_error("[PeerSsdManager] unknown ssd_backend='" + ssd.ssd_backend + - "' (expected one of: file, spdk, spdk_proxy)"); -} - -PeerSsdManager::PeerSsdManager(std::unique_ptr backend, double high_watermark, - double low_watermark) - : backend_(std::move(backend)), high_watermark_(high_watermark), low_watermark_(low_watermark) { - if (!WatermarksValid(high_watermark_, low_watermark_)) { - throw std::runtime_error( - "[PeerSsdManager] invalid SSD watermarks (require 0 < low_watermark < " - "high_watermark <= 1)"); - } -} - -PeerSsdManager::~PeerSsdManager() = default; - -// --------------------------------------------------------------------------- -// Capacity & queries -// --------------------------------------------------------------------------- - -std::pair PeerSsdManager::Capacity() const { - if (!backend_) return {0, 0}; - return backend_->Capacity(); -} - -bool PeerSsdManager::Exists(const std::string& key) const { - std::lock_guard lock(mutex_); - return owned_.find(key) != owned_.end(); -} - -void PeerSsdManager::TouchLocked(const std::string& key) { - auto it = owned_.find(key); - if (it == owned_.end()) return; - lru_.splice(lru_.begin(), lru_, it->second.lru_it); // move-to-front; iterator stays valid - it->second.lru_it = lru_.begin(); -} - -// --------------------------------------------------------------------------- -// Write (copy-on-commit landing) -// --------------------------------------------------------------------------- - -bool PeerSsdManager::Write(const std::string& key, - const std::vector>& segments, - size_t total_size) { - if (!backend_) return false; - - // Optimization, not just defense: the DRAM pin only dedups *concurrent* - // copies, so a sequential re-put of an already-resident key would otherwise - // repeat the whole SSD write. Skip it and refresh LRU instead. - { - std::lock_guard lock(mutex_); - if (owned_.find(key) != owned_.end()) { - TouchLocked(key); - return true; - } - } - - // Assemble (possibly non-contiguous) DRAM source segments into one contiguous - // buffer for the backend. A single right-sized segment is written directly; - // otherwise we memcpy into scratch (a writev-style path could avoid the copy). - const void* data = nullptr; - std::vector scratch; - if (segments.size() == 1 && segments[0].second == total_size) { - data = segments[0].first; - } else { - scratch.resize(total_size); - size_t off = 0; - for (const auto& [ptr, len] : segments) { - if (off + len > total_size) { - MORI_UMBP_ERROR("[PeerSsdManager] Write key={} segments exceed total_size={}", key, - total_size); - return false; - } - if (len > 0) std::memcpy(scratch.data() + off, ptr, len); - off += len; - } - if (off != total_size) { - MORI_UMBP_ERROR("[PeerSsdManager] Write key={} assembled {} != total_size={}", key, off, - total_size); - return false; - } - data = scratch.data(); - } - - // Backend IO outside our mutex_ — SSDTier is internally synchronized. A - // failure may be ENOSPC: run one eviction round to reclaim space and retry - // once before giving up (best-effort, no event/record on final failure). - if (!backend_->Write(key, data, total_size)) { - EvictToLowWatermark(); - if (!backend_->Write(key, data, total_size)) { - MORI_UMBP_WARN("[PeerSsdManager] backend Write failed key={} size={}", key, total_size); - return false; - } - } - - // Bytes physically written to the SSD device (write IO bandwidth source). - // Counted on every successful backend Write, including the rare dup-content - // re-write by a second copy worker (real device IO happened either way). - metrics_.copy_bytes.fetch_add(total_size, std::memory_order_relaxed); - - // Record the location. Re-check owned_: with > 1 copy worker, two of them - // could have both seen "absent" above and written the same (identical, - // content-addressed) bytes — only the first records it + emits ADD SSD. - { - std::lock_guard lock(mutex_); - if (owned_.find(key) != owned_.end()) { - TouchLocked(key); - } else { - lru_.push_front(key); - owned_.emplace(key, OwnedEntry{total_size, lru_.begin()}); - pending_events_.push_back(KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, total_size}); - } - } - - // Check-after-write trigger, on this copy worker (no dedicated thread). - auto [used, total] = Capacity(); - if (total > 0 && static_cast(used) >= high_watermark_ * static_cast(total)) { - EvictToLowWatermark(); - } - return true; -} - -// --------------------------------------------------------------------------- -// Eviction (local, read-priority) -// --------------------------------------------------------------------------- - -bool PeerSsdManager::Evict(const std::string& key) { - if (!backend_) return false; - - // Reserve under the lock: skip if a read is in flight (read priority), if the - // key is gone, or if another worker is already evicting it. owned_ is kept - // until the backend confirms the delete. - { - std::lock_guard lock(mutex_); - if (owned_.find(key) == owned_.end()) return false; - if (inflight_reads_.count(key) != 0) return false; - if (!evicting_.insert(key).second) return false; - } - - bool ok = backend_->Evict(key); // backend IO outside the lock - - // Commit only if the backend freed the bytes; on failure keep owned_ for a - // later retry, so REMOVE SSD is never emitted while the bytes still exist. - std::lock_guard lock(mutex_); - evicting_.erase(key); - if (!ok) { - metrics_.evict_backend_failures.fetch_add(1, std::memory_order_relaxed); - MORI_UMBP_WARN("[PeerSsdManager] backend Evict failed key={} — keeping for retry", key); - return false; - } - auto it = owned_.find(key); - if (it != owned_.end()) { // still present (a racing ClearLocal could have dropped it) - metrics_.evict_victims.fetch_add(1, std::memory_order_relaxed); - metrics_.evict_bytes_freed.fetch_add(it->second.size, std::memory_order_relaxed); - lru_.erase(it->second.lru_it); - owned_.erase(it); - pending_events_.push_back(KvEvent{KvEvent::Kind::REMOVE, key, TierType::SSD, 0}); - } - return true; -} - -std::vector PeerSsdManager::SelectVictims(size_t bytes_to_free) { - std::lock_guard lock(mutex_); - std::vector victims; - if (bytes_to_free == 0) return victims; - - // Oldest first (LRU tail -> MRU front), skipping keys being read or evicted. - size_t freed = 0; - for (auto it = lru_.rbegin(); it != lru_.rend(); ++it) { - const std::string& key = *it; - if (inflight_reads_.count(key) != 0) continue; - if (evicting_.count(key) != 0) continue; - auto owned_it = owned_.find(key); - if (owned_it == owned_.end()) continue; // defensive; lru_/owned_ stay in sync - victims.push_back(key); - freed += owned_it->second.size; - if (freed >= bytes_to_free) break; - } - return victims; -} - -void PeerSsdManager::EvictToLowWatermark() { - if (!backend_) return; - - // Only one eviction round at a time: a second concurrent worker (or a worker - // racing ClearLocal) backs off instead of over-evicting. try_lock keeps the - // copy path non-blocking. - std::unique_lock round(eviction_mu_, std::try_to_lock); - if (!round.owns_lock()) return; - - auto [used, total] = Capacity(); - if (total == 0) return; - double low_bytes = low_watermark_ * static_cast(total); - if (static_cast(used) <= low_bytes) return; - size_t bytes_to_free = used - static_cast(low_bytes); - - // A real round is about to run (we own the round lock and are over the low - // watermark). Count it before selecting victims. - metrics_.evict_rounds.fetch_add(1, std::memory_order_relaxed); - - // Single pass, no retry loop: if everything reclaimable is in use we free - // what we can and stop, so a fully-pinned tier cannot starve the worker. - for (const auto& key : SelectVictims(bytes_to_free)) { - Evict(key); - } -} - -// --------------------------------------------------------------------------- -// Clear (drop metadata + wipe physical bytes) -// --------------------------------------------------------------------------- - -void PeerSsdManager::ClearLocal() { - // CALLER INVARIANT: quiesce the copy pipeline before calling this. We drain - // in-flight reads (below) but not in-flight Writes, so a racing Write could - // re-add owned_/ADD SSD after backend->Clear() wipes its bytes. - // PoolClient::Clear() Quiesce()s first; new callers must too. - // - // Exclude eviction rounds for the whole operation (lock order: eviction_mu_ - // before mutex_, matching EvictToLowWatermark). - std::lock_guard round(eviction_mu_); - { - std::unique_lock lock(mutex_); - owned_.clear(); - lru_.clear(); - pending_events_.clear(); - evicting_.clear(); - // Read priority: new reads already miss (owned_ is empty), but a read that - // already started holds inflight_reads_ and is doing backend IO. Wait for - // it to finish before wiping the bytes — SSD reads cannot be safely aborted. - reads_drained_cv_.wait(lock, [this] { return inflight_reads_.empty(); }); - } - if (backend_) backend_->Clear(); // delete physical SSD bytes (user Clear = discard cache) -} - -void PeerSsdManager::DiscardLeftoverOnStartup() { - // owned_ is empty on a fresh process, so there is nothing to reconcile: just - // wipe any orphan bytes the backend loaded from disk so used capacity starts - // at 0 (cache is re-fetchable; safe to drop). See header for the policy. - if (!backend_) return; - auto [used, total] = backend_->Capacity(); - if (used == 0) { - MORI_UMBP_INFO("[PeerSsdManager] startup discard: no SSD leftover (used=0)"); - return; - } - MORI_UMBP_INFO("[PeerSsdManager] startup discard: wiping {}B SSD leftover (total={}B)", used, - total); - backend_->Clear(); -} - -// --------------------------------------------------------------------------- -// Read -// --------------------------------------------------------------------------- - -SsdReadOutcome PeerSsdManager::PrepareRead(const std::string& key, void* staging_ptr, - size_t staging_cap) { - if (!backend_) { - metrics_.read_not_found.fetch_add(1, std::memory_order_relaxed); - return SsdReadOutcome{SsdReadStatus::kNotFound, 0}; - } - - // Resolve size and mark the read in flight under the lock, then run the - // blocking SSD IO outside it (so a concurrent copy Write is not serialized). - size_t size = 0; - { - std::lock_guard lock(mutex_); - auto it = owned_.find(key); - if (it == owned_.end()) { - metrics_.read_not_found.fetch_add(1, std::memory_order_relaxed); - return SsdReadOutcome{SsdReadStatus::kNotFound, 0}; - } - // Being evicted: bytes about to vanish — treat a new read as a stale-route - // miss rather than racing the backend delete. - if (evicting_.count(key) != 0) { - metrics_.read_not_found.fetch_add(1, std::memory_order_relaxed); - return SsdReadOutcome{SsdReadStatus::kNotFound, 0}; - } - size = it->second.size; - // Reject over-capacity before touching the device (no in-flight mark, no IO). - if (size > staging_cap) { - metrics_.read_size_too_large.fetch_add(1, std::memory_order_relaxed); - // staging_cap = per-slot cap (ssd_staging_buffer_size / ssd_staging_buffer_slots). - MORI_UMBP_WARN( - "[PeerSsdManager] remote SSD read key={} size={}B exceeds per-slot staging cap {}B; " - "raise ssd_staging_buffer_size or lower ssd_staging_buffer_slots", - key, size, staging_cap); - return SsdReadOutcome{SsdReadStatus::kSizeTooLarge, size}; - } - ++inflight_reads_[key]; - } - - bool read_ok = backend_->ReadIntoPtr(key, reinterpret_cast(staging_ptr), size); - - // Always release the in-flight mark (even on error), refresh LRU on success, - // and wake a waiting ClearLocal once the last read drains. - { - std::lock_guard lock(mutex_); - auto rit = inflight_reads_.find(key); - if (rit != inflight_reads_.end() && --rit->second <= 0) inflight_reads_.erase(rit); - if (read_ok && owned_.find(key) != owned_.end()) TouchLocked(key); - if (inflight_reads_.empty()) reads_drained_cv_.notify_all(); - } - - if (!read_ok) { - // owned_ had the key but the backend couldn't serve it (e.g. a local evict - // raced us, or a corrupt record): kError, not a definitive miss. - metrics_.read_error.fetch_add(1, std::memory_order_relaxed); - MORI_UMBP_WARN("[PeerSsdManager] PrepareRead backend read failed key={} size={}", key, size); - return SsdReadOutcome{SsdReadStatus::kError, size}; - } - metrics_.read_ok.fetch_add(1, std::memory_order_relaxed); - metrics_.read_bytes.fetch_add(size, std::memory_order_relaxed); // SSD read IO bandwidth source - return SsdReadOutcome{SsdReadStatus::kOk, size}; -} - -// --------------------------------------------------------------------------- -// OwnedLocationSource (heartbeat event drain / snapshot) -// --------------------------------------------------------------------------- - -std::vector PeerSsdManager::DrainPendingEvents() { - std::lock_guard lock(mutex_); - std::vector drained; - drained.swap(pending_events_); - return drained; -} - -std::vector PeerSsdManager::SnapshotOwnedKeysLocked() const { - std::vector out; - out.reserve(owned_.size()); - for (const auto& [key, entry] : owned_) { - out.push_back(KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, entry.size}); - } - return out; -} - -std::vector PeerSsdManager::SnapshotOwnedKeys() const { - std::lock_guard lock(mutex_); - return SnapshotOwnedKeysLocked(); -} - -std::vector PeerSsdManager::SnapshotOwnedKeysForFullSync() { - std::lock_guard lock(mutex_); - auto out = SnapshotOwnedKeysLocked(); - // Snapshot is authoritative now: drop the queued delta (already reflected in - // it) so the next delta carries only events produced after this snapshot. - pending_events_.clear(); - return out; -} - -} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/ssd/peer_ssd_manager.cpp b/src/umbp/distributed/peer/ssd/peer_ssd_manager.cpp new file mode 100644 index 000000000..b6e3f823a --- /dev/null +++ b/src/umbp/distributed/peer/ssd/peer_ssd_manager.cpp @@ -0,0 +1,947 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/peer/ssd/peer_ssd_manager.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/common/ssd_perf.h" +#include "umbp/local/tiers/spdk_proxy_tier.h" +#include "umbp/local/tiers/ssd_backend_factory.h" +#include "umbp/local/tiers/ssd_tier.h" +#include "umbp/local/tiers/tier_backend.h" + +namespace mori::umbp { + +namespace { +// Watermarks must satisfy 0 < low < high <= 1. We fail fast on a bad value +// rather than silently clamping, so a misconfigured env surfaces immediately. +bool WatermarksValid(double high, double low) { + return high > 0.0 && high <= 1.0 && low > 0.0 && low < high; +} +} // namespace + +// --------------------------------------------------------------------------- +// Construction +// --------------------------------------------------------------------------- + +PeerSsdManager::PeerSsdManager(const PeerSsdConfig& cfg) { + if (!cfg.enabled) { + MORI_UMBP_INFO("[PeerSsdManager] constructed disabled (no SSD backend)"); + return; + } + const auto& ssd = cfg.ssd; + + if (!WatermarksValid(ssd.high_watermark, ssd.low_watermark)) { + throw std::runtime_error( + "[PeerSsdManager] invalid SSD watermarks (require 0 < low_watermark < " + "high_watermark <= 1)"); + } + high_watermark_ = ssd.high_watermark; + low_watermark_ = ssd.low_watermark; + single_flight_ = ssd.single_flight_reads; + + // Explicit backend selection — an unknown value is a configuration error, not + // a reason to silently fall back to the file backend. + if (ssd.ssd_backend == "file") { + backend_ = MakeFileSsdBackend(ssd); + MORI_UMBP_INFO("[PeerSsdManager] SSDTier ready dir={} capacity={}B", ssd.storage_dir, + ssd.capacity_bytes); + return; + } + + // The distributed peer shares one physical device across processes, so SPDK + // is reached through the spdk_proxy daemon, never the single-process direct + // SpdkSsdTier. Both "spdk" and "spdk_proxy" therefore map to SpdkProxyTier: + // "spdk" here means "use SPDK via the proxy". + if (ssd.ssd_backend == "spdk" || ssd.ssd_backend == "spdk_proxy") { + MORI_UMBP_INFO("[PeerSsdManager] distributed ssd_backend={} uses SpdkProxyTier", + ssd.ssd_backend); + auto proxy = std::make_unique(ssd); + // PeerSsdManager only connects to an already-ready proxy; it does not spawn + // the daemon. A connect failure with an explicit SPDK config is fatal — we + // must not silently disable SSD or fall back to POSIX. + if (!proxy->IsValid()) { + throw std::runtime_error("[PeerSsdManager] ssd_backend=" + ssd.ssd_backend + + ": SPDK proxy connect failed (shm=" + ssd.spdk_proxy_shm_name + + "). Ensure the spdk_proxy daemon is running and READY."); + } + backend_ = std::move(proxy); + MORI_UMBP_INFO("[PeerSsdManager] SpdkProxyTier ready (shm={})", ssd.spdk_proxy_shm_name); + return; + } + + throw std::runtime_error("[PeerSsdManager] unknown ssd_backend='" + ssd.ssd_backend + + "' (expected one of: file, spdk, spdk_proxy)"); +} + +PeerSsdManager::PeerSsdManager(std::unique_ptr backend, double high_watermark, + double low_watermark, bool single_flight) + : backend_(std::move(backend)), + high_watermark_(high_watermark), + low_watermark_(low_watermark), + single_flight_(single_flight) { + if (!WatermarksValid(high_watermark_, low_watermark_)) { + throw std::runtime_error( + "[PeerSsdManager] invalid SSD watermarks (require 0 < low_watermark < " + "high_watermark <= 1)"); + } +} + +PeerSsdManager::~PeerSsdManager() = default; + +// --------------------------------------------------------------------------- +// Capacity & queries +// --------------------------------------------------------------------------- + +std::pair PeerSsdManager::Capacity() const { + if (!backend_) return {0, 0}; + return backend_->Capacity(); +} + +bool PeerSsdManager::Exists(const std::string& key) const { + std::lock_guard lock(mutex_); + return owned_.find(key) != owned_.end(); +} + +uint64_t PeerSsdManager::SizeOf(const std::string& key) const { + std::lock_guard lock(mutex_); + auto it = owned_.find(key); + return it == owned_.end() ? 0 : it->second.size; +} + +bool PeerSsdManager::PinForMigration(const std::string& key) { + std::lock_guard lock(mutex_); + if (owned_.find(key) == owned_.end() || evicting_.count(key) != 0) return false; + auto& slot = inflight_reads_[key]; + if (!slot) slot = std::make_unique(); + ++slot->refs; + return true; +} + +void PeerSsdManager::UnpinForMigration(const std::string& key) { + std::lock_guard lock(mutex_); + auto it = inflight_reads_.find(key); + if (it != inflight_reads_.end() && --it->second->refs <= 0) inflight_reads_.erase(it); + if (inflight_reads_.empty()) reads_drained_cv_.notify_all(); +} + +std::optional PeerSsdManager::LocateRecord(const std::string& key) const { + std::lock_guard lock(mutex_); + if (backend_ == nullptr) return std::nullopt; + // A key mid-eviction is a miss, exactly as PrepareRead treats it. + if (evicting_.count(key) != 0) return std::nullopt; + if (owned_.find(key) == owned_.end()) return std::nullopt; + return backend_->LocateRecord(key); +} + +void PeerSsdManager::TouchLocked(const std::string& key) { + auto it = owned_.find(key); + if (it == owned_.end()) return; + lru_.splice(lru_.begin(), lru_, it->second.lru_it); // move-to-front; iterator stays valid + it->second.lru_it = lru_.begin(); +} + +// --------------------------------------------------------------------------- +// Write (copy-on-commit landing) +// --------------------------------------------------------------------------- + +bool PeerSsdManager::Write(const std::string& key, + const std::vector>& segments, + size_t total_size) { + if (!backend_) return false; + + // Optimization, not just defense: the DRAM pin only dedups *concurrent* + // copies, so a sequential re-put of an already-resident key would otherwise + // repeat the whole SSD write. Skip it and refresh LRU instead. + { + std::lock_guard lock(mutex_); + if (owned_.find(key) != owned_.end()) { + TouchLocked(key); + return true; + } + } + + // Assemble (possibly non-contiguous) DRAM source segments into one contiguous + // buffer for the backend. A single right-sized segment is written directly; + // otherwise we memcpy into scratch (a writev-style path could avoid the copy). + const void* data = nullptr; + std::vector scratch; + if (segments.size() == 1 && segments[0].second == total_size) { + data = segments[0].first; + } else { + scratch.resize(total_size); + size_t off = 0; + for (const auto& [ptr, len] : segments) { + if (off + len > total_size) { + MORI_UMBP_ERROR("[PeerSsdManager] Write key={} segments exceed total_size={}", key, + total_size); + return false; + } + if (len > 0) std::memcpy(scratch.data() + off, ptr, len); + off += len; + } + if (off != total_size) { + MORI_UMBP_ERROR("[PeerSsdManager] Write key={} assembled {} != total_size={}", key, off, + total_size); + return false; + } + data = scratch.data(); + } + + // Backend IO outside our mutex_ — SSDTier is internally synchronized. A + // failure may be ENOSPC: run one eviction round to reclaim space and retry + // once before giving up (best-effort, no event/record on final failure). + if (!backend_->Write(key, data, total_size)) { + EvictToLowWatermark(); + if (!backend_->Write(key, data, total_size)) { + MORI_UMBP_WARN("[PeerSsdManager] backend Write failed key={} size={}", key, total_size); + return false; + } + } + + // Bytes physically written to the SSD device (write IO bandwidth source). + // Counted on every successful backend Write, including the rare dup-content + // re-write by a second copy worker (real device IO happened either way). + metrics_.copy_bytes.fetch_add(total_size, std::memory_order_relaxed); + + // Record the location. Re-check owned_: with > 1 copy worker, two of them + // could have both seen "absent" above and written the same (identical, + // content-addressed) bytes — only the first records it + emits ADD SSD. + { + std::lock_guard lock(mutex_); + if (owned_.find(key) != owned_.end()) { + TouchLocked(key); + } else { + lru_.push_front(key); + owned_.emplace(key, OwnedEntry{total_size, lru_.begin()}); + if (event_publishing_) { + pending_events_.push_back(KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, total_size}); + } + } + // Keep the just-written key out of this write's own watermark victim set. + // A successful Write must imply the key still exists when it returns. + auto& slot = inflight_reads_[key]; + if (!slot) slot = std::make_unique(); + ++slot->refs; + } + + // Check-after-write trigger, on this copy worker (no dedicated thread). + auto [used, total] = Capacity(); + if (total > 0 && static_cast(used) >= high_watermark_ * static_cast(total)) { + EvictToLowWatermark(); + } + UnpinForMigration(key); + return true; +} + +std::vector PeerSsdManager::WriteBatch(const std::vector& keys, + const std::vector& srcs, + const std::vector& sizes) { + std::vector results(keys.size(), false); + if (!backend_ || keys.empty()) return results; + if (srcs.size() != keys.size() || sizes.size() != keys.size()) { + MORI_UMBP_ERROR("[PeerSsdManager] WriteBatch: mismatched input lengths ({} / {} / {})", + keys.size(), srcs.size(), sizes.size()); + return results; + } + + // Stage timers for the [SsdPerf/peer] PUT breakdown (no-ops unless + // UMBP_SSD_TIMING is set). + const auto t_begin = ssdperf::Now(); + double dedup_ms = 0.0, backend_ms = 0.0, retry_ms = 0.0, record_ms = 0.0, evict_ms = 0.0; + + // Pass 1 — dedup under one lock. Already-owned keys succeed with no IO (same + // rationale as Write: a sequential re-put must not repeat the device write). + std::vector todo; + todo.reserve(keys.size()); + { + std::lock_guard lock(mutex_); + for (size_t i = 0; i < keys.size(); ++i) { + if (sizes[i] == 0) continue; // meaningless put; leave results[i] == false + if (owned_.find(keys[i]) != owned_.end()) { + TouchLocked(keys[i]); + results[i] = true; + continue; + } + todo.push_back(i); + } + } + dedup_ms = ssdperf::MsSince(t_begin); + const size_t dedup_hits = keys.size() - todo.size(); + if (todo.empty()) return results; + + auto gather = [&](const std::vector& idxs) { + std::vector k; + std::vector d; + std::vector z; + k.reserve(idxs.size()); + d.reserve(idxs.size()); + z.reserve(idxs.size()); + for (size_t i : idxs) { + k.push_back(keys[i]); + d.push_back(srcs[i]); + z.push_back(sizes[i]); + } + return std::make_tuple(std::move(k), std::move(d), std::move(z)); + }; + + // Pass 2 — one batched backend write outside the lock. Failures are most + // likely ENOSPC, so retry exactly those once after a single eviction round + // (mirrors Write's evict-and-retry, amortized over the whole batch). + const auto t_backend = ssdperf::Now(); + auto [bk, bd, bz] = gather(todo); + auto ok = backend_->BatchWrite(bk, bd, bz); + ok.resize(todo.size(), false); + backend_ms = ssdperf::MsSince(t_backend); + + std::vector retry; + for (size_t j = 0; j < todo.size(); ++j) { + if (!ok[j]) retry.push_back(todo[j]); + } + if (!retry.empty()) { + const auto t_retry = ssdperf::Now(); + EvictToLowWatermark(); + auto [rk, rd, rz] = gather(retry); + auto retry_ok = backend_->BatchWrite(rk, rd, rz); + retry_ok.resize(retry.size(), false); + for (size_t j = 0, r = 0; j < todo.size(); ++j) { + if (!ok[j]) ok[j] = retry_ok[r++]; + } + retry_ms = ssdperf::MsSince(t_retry); + } + + // Pass 3 — record locations + ADD events under one lock. The owned_ re-check + // matches Write: a concurrent writer may have landed the same content-addressed + // bytes, and only the first recorder emits ADD SSD. + const auto t_record = ssdperf::Now(); + uint64_t written_bytes = 0; + { + std::lock_guard lock(mutex_); + for (size_t j = 0; j < todo.size(); ++j) { + const size_t i = todo[j]; + if (!ok[j]) continue; + results[i] = true; + written_bytes += sizes[i]; + if (owned_.find(keys[i]) != owned_.end()) { + TouchLocked(keys[i]); + } else { + lru_.push_front(keys[i]); + owned_.emplace(keys[i], OwnedEntry{sizes[i], lru_.begin()}); + if (event_publishing_) { + pending_events_.push_back(KvEvent{KvEvent::Kind::ADD, keys[i], TierType::SSD, sizes[i]}); + } + } + } + } + metrics_.copy_bytes.fetch_add(written_bytes, std::memory_order_relaxed); + record_ms = ssdperf::MsSince(t_record); + + for (size_t j = 0; j < todo.size(); ++j) { + if (!ok[j]) { + MORI_UMBP_WARN("[PeerSsdManager] WriteBatch: backend write failed key={} size={}", + keys[todo[j]], sizes[todo[j]]); + } + } + + // One check-after-write for the whole batch instead of one per key. + const auto t_evict = ssdperf::Now(); + auto [used, total] = Capacity(); + if (total > 0 && static_cast(used) >= high_watermark_ * static_cast(total)) { + EvictToLowWatermark(); + } + evict_ms = ssdperf::MsSince(t_evict); + + if (ssdperf::Enabled()) { + const double total_ms = ssdperf::MsSince(t_begin); + MORI_UMBP_INFO( + "[SsdPerf/peer] PUT keys={} dedup_hits={} written={} bytes={} total_ms={:.3f} GB_s={:.2f} " + "| dedup={:.3f}ms backend={:.3f}ms ({:.0f}%) evict_retry={:.3f}ms record={:.3f}ms " + "evict_check={:.3f}ms | used={}/{}B", + keys.size(), dedup_hits, todo.size(), written_bytes, total_ms, + ssdperf::GbPerSec(written_bytes, total_ms), dedup_ms, backend_ms, + ssdperf::Pct(backend_ms, total_ms), retry_ms, record_ms, evict_ms, used, total); + } + return results; +} + +// --------------------------------------------------------------------------- +// Eviction (local, read-priority) +// --------------------------------------------------------------------------- + +bool PeerSsdManager::Evict(const std::string& key) { + if (!backend_) return false; + + // Reserve under the lock: skip if a read is in flight (read priority), if the + // key is gone, or if another worker is already evicting it. owned_ is kept + // until the backend confirms the delete. + { + std::lock_guard lock(mutex_); + if (owned_.find(key) == owned_.end()) return false; + if (inflight_reads_.count(key) != 0) return false; + if (!evicting_.insert(key).second) return false; + } + + bool ok = backend_->Evict(key); // backend IO outside the lock + + // Commit only if the backend freed the bytes; on failure keep owned_ for a + // later retry, so REMOVE SSD is never emitted while the bytes still exist. + std::lock_guard lock(mutex_); + evicting_.erase(key); + if (!ok) { + metrics_.evict_backend_failures.fetch_add(1, std::memory_order_relaxed); + MORI_UMBP_WARN("[PeerSsdManager] backend Evict failed key={} — keeping for retry", key); + return false; + } + auto it = owned_.find(key); + if (it != owned_.end()) { // still present (a racing ClearLocal could have dropped it) + metrics_.evict_victims.fetch_add(1, std::memory_order_relaxed); + metrics_.evict_bytes_freed.fetch_add(it->second.size, std::memory_order_relaxed); + lru_.erase(it->second.lru_it); + owned_.erase(it); + if (event_publishing_) { + pending_events_.push_back(KvEvent{KvEvent::Kind::REMOVE, key, TierType::SSD, 0}); + } + } + return true; +} + +std::vector PeerSsdManager::SelectVictims(size_t bytes_to_free) { + std::lock_guard lock(mutex_); + std::vector victims; + if (bytes_to_free == 0) return victims; + + // Oldest first (LRU tail -> MRU front), skipping keys being read or evicted. + size_t freed = 0; + for (auto it = lru_.rbegin(); it != lru_.rend(); ++it) { + const std::string& key = *it; + if (inflight_reads_.count(key) != 0) continue; + if (evicting_.count(key) != 0) continue; + auto owned_it = owned_.find(key); + if (owned_it == owned_.end()) continue; // defensive; lru_/owned_ stay in sync + victims.push_back(key); + freed += owned_it->second.size; + if (freed >= bytes_to_free) break; + } + return victims; +} + +void PeerSsdManager::EvictToLowWatermark() { + if (!backend_) return; + + // Only one eviction round at a time: a second concurrent worker (or a worker + // racing ClearLocal) backs off instead of over-evicting. try_lock keeps the + // copy path non-blocking. + std::unique_lock round(eviction_mu_, std::try_to_lock); + if (!round.owns_lock()) return; + + auto [used, total] = Capacity(); + if (total == 0) return; + double low_bytes = low_watermark_ * static_cast(total); + if (static_cast(used) <= low_bytes) return; + size_t bytes_to_free = used - static_cast(low_bytes); + + // A real round is about to run (we own the round lock and are over the low + // watermark). Count it before selecting victims. + metrics_.evict_rounds.fetch_add(1, std::memory_order_relaxed); + + // Single pass, no retry loop: if everything reclaimable is in use we free + // what we can and stop, so a fully-pinned tier cannot starve the worker. + for (const auto& key : SelectVictims(bytes_to_free)) { + Evict(key); + } +} + +// --------------------------------------------------------------------------- +// Clear (drop metadata + wipe physical bytes) +// --------------------------------------------------------------------------- + +void PeerSsdManager::ClearLocal() { + // CALLER INVARIANT: quiesce the copy pipeline before calling this. We drain + // in-flight reads (below) but not in-flight Writes, so a racing Write could + // re-add owned_/ADD SSD after backend->Clear() wipes its bytes. + // PoolClient::Clear() Quiesce()s first; new callers must too. + // + // Exclude eviction rounds for the whole operation (lock order: eviction_mu_ + // before mutex_, matching EvictToLowWatermark). + std::lock_guard round(eviction_mu_); + { + std::unique_lock lock(mutex_); + owned_.clear(); + lru_.clear(); + pending_events_.clear(); + evicting_.clear(); + // Read priority: new reads already miss (owned_ is empty), but a read that + // already started holds inflight_reads_ and is doing backend IO. Wait for + // it to finish before wiping the bytes — SSD reads cannot be safely aborted. + reads_drained_cv_.wait(lock, [this] { return inflight_reads_.empty(); }); + } + if (backend_) backend_->Clear(); // delete physical SSD bytes (user Clear = discard cache) +} + +void PeerSsdManager::DiscardLeftoverOnStartup() { + // owned_ is empty on a fresh process, so there is nothing to reconcile: just + // wipe any orphan bytes the backend loaded from disk so used capacity starts + // at 0 (cache is re-fetchable; safe to drop). See header for the policy. + if (!backend_) return; + auto [used, total] = backend_->Capacity(); + if (used == 0) { + MORI_UMBP_INFO("[PeerSsdManager] startup discard: no SSD leftover (used=0)"); + return; + } + MORI_UMBP_INFO("[PeerSsdManager] startup discard: wiping {}B SSD leftover (total={}B)", used, + total); + backend_->Clear(); +} + +// --------------------------------------------------------------------------- +// Read +// --------------------------------------------------------------------------- + +SsdReadOutcome PeerSsdManager::PrepareRead(const std::string& key, void* staging_ptr, + size_t staging_cap) { + if (!backend_) { + metrics_.read_not_found.fetch_add(1, std::memory_order_relaxed); + return SsdReadOutcome{SsdReadStatus::kNotFound, 0}; + } + + // Resolve size and mark the read in flight under the lock, then run the + // blocking SSD IO outside it (so a concurrent copy Write is not serialized). + size_t size = 0; + bool leading = false; // we own the device read and the fan-out + bool following = false; // a leader will fill our buffer; we only wait + std::shared_ptr episode; + // Our own slot in the leader's fan-out list, kept because the leader swaps + // that list out of the episode when it seals. + std::shared_ptr my_target; + { + std::lock_guard lock(mutex_); + auto it = owned_.find(key); + if (it == owned_.end()) { + metrics_.read_not_found.fetch_add(1, std::memory_order_relaxed); + return SsdReadOutcome{SsdReadStatus::kNotFound, 0}; + } + // Being evicted: bytes about to vanish — treat a new read as a stale-route + // miss rather than racing the backend delete. + if (evicting_.count(key) != 0) { + metrics_.read_not_found.fetch_add(1, std::memory_order_relaxed); + return SsdReadOutcome{SsdReadStatus::kNotFound, 0}; + } + size = it->second.size; + // Reject over-capacity before touching the device (no in-flight mark, no IO). + if (size > staging_cap) { + metrics_.read_size_too_large.fetch_add(1, std::memory_order_relaxed); + // staging_cap = per-slot cap (ssd_staging_buffer_size / ssd_staging_buffer_slots). + MORI_UMBP_WARN( + "[PeerSsdManager] remote SSD read key={} size={}B exceeds per-slot staging cap {}B; " + "raise ssd_staging_buffer_size or lower ssd_staging_buffer_slots", + key, size, staging_cap); + return SsdReadOutcome{SsdReadStatus::kSizeTooLarge, size}; + } + // Same leader/follower protocol as PrepareReadBatch, over the same map, so + // a single-key read and a batched one coalesce with each other. No log + // here: this path is one key per call, so the counters are the only sane + // reporting channel. + auto& slot = inflight_reads_[key]; + if (!slot) slot = std::make_unique(); + const int prior = slot->refs++; + metrics_.read_dup.fetch_add(prior > 0 ? 1 : 0, std::memory_order_relaxed); + metrics_.read_lead.fetch_add(prior > 0 ? 0 : 1, std::memory_order_relaxed); + if (!slot->episode) { + slot->episode = std::make_shared(); + episode = slot->episode; + leading = true; + } else if (single_flight_ && !slot->episode->sealed) { + my_target = std::make_shared(); + my_target->dst = staging_ptr; + my_target->cap = staging_cap; + slot->episode->followers.push_back(my_target); + episode = slot->episode; + following = true; + metrics_.read_merged.fetch_add(1, std::memory_order_relaxed); + } + // else: missed the merge window — fall through and read it ourselves. + } + + bool read_ok = false; + if (following) { + std::unique_lock lock(mutex_); + const bool signalled = episode->cv.wait_for(lock, std::chrono::seconds(kFollowerWaitSeconds), + [&episode] { return episode->done; }); + if (!signalled) { + MORI_UMBP_WARN( + "[PeerSsdManager] single-flight follower timed out after {}s waiting on leader key={}", + kFollowerWaitSeconds, key); + } + read_ok = signalled && my_target->ok; + } else { + read_ok = backend_->ReadIntoPtr(key, reinterpret_cast(staging_ptr), size); + if (leading) { + // Publish to whoever attached while we were on the device. Failures are + // published too, so a follower is never left to time out. + std::vector> fanout; + { + std::lock_guard lock(mutex_); + episode->sealed = true; + fanout.swap(episode->followers); + auto it = inflight_reads_.find(key); + if (it != inflight_reads_.end() && it->second->episode == episode) { + it->second->episode.reset(); + } + } + for (const auto& f : fanout) { + f->ok = read_ok; + if (read_ok) std::memcpy(f->dst, staging_ptr, std::min(size, f->cap)); + } + { + std::lock_guard lock(mutex_); + episode->done = true; + episode->cv.notify_all(); + } + } + } + + // Always release the in-flight mark (even on error), refresh LRU on success, + // and wake a waiting ClearLocal once the last read drains. + { + std::lock_guard lock(mutex_); + auto rit = inflight_reads_.find(key); + if (rit != inflight_reads_.end() && --rit->second->refs <= 0) inflight_reads_.erase(rit); + if (read_ok && owned_.find(key) != owned_.end()) TouchLocked(key); + if (inflight_reads_.empty()) reads_drained_cv_.notify_all(); + } + + if (!read_ok) { + // owned_ had the key but the backend couldn't serve it (e.g. a local evict + // raced us, or a corrupt record): kError, not a definitive miss. + metrics_.read_error.fetch_add(1, std::memory_order_relaxed); + MORI_UMBP_WARN("[PeerSsdManager] PrepareRead backend read failed key={} size={}", key, size); + return SsdReadOutcome{SsdReadStatus::kError, size}; + } + metrics_.read_ok.fetch_add(1, std::memory_order_relaxed); + metrics_.read_bytes.fetch_add(size, std::memory_order_relaxed); // SSD read IO bandwidth source + return SsdReadOutcome{SsdReadStatus::kOk, size}; +} + +std::vector PeerSsdManager::PrepareReadBatch(const std::vector& keys, + const std::vector& dsts, + const std::vector& caps) { + std::vector out(keys.size(), SsdReadOutcome{SsdReadStatus::kNotFound, 0}); + if (keys.empty()) return out; + if (!backend_) { + metrics_.read_not_found.fetch_add(keys.size(), std::memory_order_relaxed); + return out; + } + if (dsts.size() != keys.size() || caps.size() != keys.size()) { + MORI_UMBP_ERROR("[PeerSsdManager] PrepareReadBatch: mismatched input lengths ({} / {} / {})", + keys.size(), dsts.size(), caps.size()); + for (auto& o : out) o = SsdReadOutcome{SsdReadStatus::kError, 0}; + return out; + } + + // Stage timers for the [SsdPerf/peer] GET breakdown (no-ops unless + // UMBP_SSD_TIMING is set). + const auto t_begin = ssdperf::Now(); + double resolve_ms = 0.0, backend_ms = 0.0, release_ms = 0.0; + // fanout_ms: Pass 2b, single-threaded memcpy + per-key lock/unlock/notify_all + // serving OTHER concurrent callers' follower reads off of our own device read. + // wait_ms: Pass 2c, time spent blocked on someone else's episode as a follower. + // Both are 0 for calls that neither lead anyone nor follow anyone. + double fanout_ms = 0.0, wait_ms = 0.0; + uint64_t fanout_bytes = 0; + size_t fanout_count = 0; + + // Pass 1 — resolve sizes and mark every servable key in flight under ONE lock + // acquisition (the per-key path pays two per key). + std::vector todo; // I lead (or independently read) these + std::vector waits; // I follow these; a leader will fill my buffer + std::vector is_leader(keys.size(), false); + // Per-request handle on the episode we lead or follow. Held as a shared_ptr + // so it stays valid even after the slot moves on to a new leader. + std::vector> episodes(keys.size()); + // Our own entry in a leader's fan-out list, for the keys we follow. Kept + // per-request because the leader swaps its list out of the episode on seal. + std::vector> my_targets(keys.size()); + todo.reserve(keys.size()); + // dup = requests that found an identical read already in flight; + // merged = those actually served without touching the drive; + // max_concurrent = largest same-key fan-out seen this batch. + size_t dup = 0, merged = 0; + int max_concurrent = 0; + { + std::lock_guard lock(mutex_); + for (size_t i = 0; i < keys.size(); ++i) { + auto it = owned_.find(keys[i]); + if (it == owned_.end() || evicting_.count(keys[i]) != 0) { + metrics_.read_not_found.fetch_add(1, std::memory_order_relaxed); + continue; // stays kNotFound + } + const size_t size = it->second.size; + if (size > caps[i]) { + metrics_.read_size_too_large.fetch_add(1, std::memory_order_relaxed); + MORI_UMBP_WARN( + "[PeerSsdManager] SSD read key={} size={}B exceeds per-slot staging cap {}B; " + "raise ssd_staging_buffer_size or lower ssd_staging_buffer_slots", + keys[i], size, caps[i]); + out[i] = SsdReadOutcome{SsdReadStatus::kSizeTooLarge, size}; + continue; + } + out[i].size = size; + auto& slot = inflight_reads_[keys[i]]; + if (!slot) slot = std::make_unique(); + const int prior = slot->refs++; + if (prior > 0) { + ++dup; + if (prior > max_concurrent) max_concurrent = prior; + } + if (!slot->episode) { + // First in: this request owns the device read and the fan-out. + slot->episode = std::make_shared(); + episodes[i] = slot->episode; + is_leader[i] = true; + todo.push_back(i); + } else if (single_flight_ && !slot->episode->sealed) { + // Merge window open: hand the leader our buffer and wait on its episode. + my_targets[i] = std::make_shared(); + my_targets[i]->dst = dsts[i]; + my_targets[i]->cap = caps[i]; + slot->episode->followers.push_back(my_targets[i]); + episodes[i] = slot->episode; + waits.push_back(i); + ++merged; + } else { + // Missed the window (or single flight off): read it ourselves. This is + // the pre-single-flight path and is always correct, just not merged. + todo.push_back(i); + } + } + } + resolve_ms = ssdperf::MsSince(t_begin); + metrics_.read_dup.fetch_add(dup, std::memory_order_relaxed); + metrics_.read_merged.fetch_add(merged, std::memory_order_relaxed); + metrics_.read_lead.fetch_add(todo.size() + waits.size() - dup, std::memory_order_relaxed); + const size_t not_owned = keys.size() - todo.size() - waits.size(); + if (todo.empty() && waits.empty()) { + if (ssdperf::Enabled()) { + MORI_UMBP_INFO( + "[SsdPerf/peer] GET keys={} served=0 not_owned={} total_ms={:.3f} | resolve={:.3f}ms " + "(no backend IO: nothing owned locally)", + keys.size(), not_owned, ssdperf::MsSince(t_begin), resolve_ms); + } + return out; + } + + // Pass 2 — one batched backend read outside the lock; ShardedSsdTier turns + // this into concurrent IO on every drive holding part of the batch. + std::vector k; + std::vector d; + std::vector z; + k.reserve(todo.size()); + d.reserve(todo.size()); + z.reserve(todo.size()); + for (size_t i : todo) { + k.push_back(keys[i]); + d.push_back(reinterpret_cast(dsts[i])); + z.push_back(out[i].size); + } + std::vector ok; + if (!todo.empty()) { + const auto t_backend = ssdperf::Now(); + ok = backend_->BatchReadIntoPtr(k, d, z); + ok.resize(todo.size(), false); + backend_ms = ssdperf::MsSince(t_backend); + } + + // Pass 2b — fan out to followers. MUST run before Pass 2c: if this batch + // leads key A while another thread leads B and follows A, waiting first would + // close a cycle. Issuing all our own reads and publishing their results + // before we block on anyone else's makes that cycle impossible. + // + // A leader whose read FAILED still publishes (ok=false) — followers must be + // released with an error, never left to time out. + const auto t_fanout = ssdperf::Now(); + for (size_t j = 0; j < todo.size(); ++j) { + const size_t i = todo[j]; + if (!is_leader[i]) continue; // independent read: nobody is waiting on it + const auto& ep = episodes[i]; + std::vector> fanout; + { + std::lock_guard lock(mutex_); + ep->sealed = true; // late arrivals now go independent instead of waiting + fanout.swap(ep->followers); + // Retire the episode so the next arrival starts a fresh one rather than + // attaching to a read that is already past its snapshot. + auto it = inflight_reads_.find(keys[i]); + if (it != inflight_reads_.end() && it->second->episode == ep) it->second->episode.reset(); + } + // Deliver outside the lock: holding mutex_ across a page-sized memcpy would + // stall every other read/evict here. + for (const auto& f : fanout) { + f->ok = ok[j]; + if (!ok[j]) continue; + const size_t n = std::min(out[i].size, f->cap); + std::memcpy(f->dst, dsts[i], n); + fanout_bytes += n; + ++fanout_count; + } + { + std::lock_guard lock(mutex_); + ep->done = true; + ep->cv.notify_all(); + } + } + fanout_ms = ssdperf::MsSince(t_fanout); + + // Pass 2c — collect the keys we followed. Our staging buffer is still alive + // (we never returned), so a leader's memcpy landed in a valid destination. + const auto t_wait = ssdperf::Now(); + std::vector wait_ok(waits.size(), false); + for (size_t w = 0; w < waits.size(); ++w) { + const size_t i = waits[w]; + const auto& ep = episodes[i]; + std::unique_lock lock(mutex_); + const bool signalled = ep->cv.wait_for(lock, std::chrono::seconds(kFollowerWaitSeconds), + [&ep] { return ep->done; }); + if (!signalled) { + MORI_UMBP_WARN( + "[PeerSsdManager] single-flight follower timed out after {}s waiting on leader key={}", + kFollowerWaitSeconds, keys[i]); + } + wait_ok[w] = signalled && my_targets[i]->ok; + } + wait_ms = ssdperf::MsSince(t_wait); + + // Pass 3 — release the in-flight marks, refresh LRU for the reads that landed, + // and wake a waiting ClearLocal once the last read drains. + const auto t_release = ssdperf::Now(); + uint64_t read_bytes = 0; + { + std::lock_guard lock(mutex_); + auto release = [&](const std::string& key, bool landed) { + auto rit = inflight_reads_.find(key); + if (rit != inflight_reads_.end() && --rit->second->refs <= 0) inflight_reads_.erase(rit); + if (landed && owned_.find(key) != owned_.end()) TouchLocked(key); + }; + for (size_t j = 0; j < todo.size(); ++j) release(keys[todo[j]], ok[j]); + for (size_t w = 0; w < waits.size(); ++w) release(keys[waits[w]], wait_ok[w]); + if (inflight_reads_.empty()) reads_drained_cv_.notify_all(); + } + + auto record = [&](size_t i, bool landed, bool from_device) { + if (landed) { + out[i].status = SsdReadStatus::kOk; + // read_bytes is device traffic only, so merged reads are excluded — the + // counter stays a true measure of what the drive moved. + if (from_device) read_bytes += out[i].size; + metrics_.read_ok.fetch_add(1, std::memory_order_relaxed); + } else { + // owned_ had the key but it couldn't be served (a local evict raced us, a + // corrupt record, or our leader failed): kError, not a definitive miss. + out[i].status = SsdReadStatus::kError; + metrics_.read_error.fetch_add(1, std::memory_order_relaxed); + MORI_UMBP_WARN("[PeerSsdManager] PrepareReadBatch read failed key={} size={} (merged={})", + keys[i], out[i].size, !from_device); + } + }; + for (size_t j = 0; j < todo.size(); ++j) record(todo[j], ok[j], /*from_device=*/true); + for (size_t w = 0; w < waits.size(); ++w) record(waits[w], wait_ok[w], /*from_device=*/false); + metrics_.read_bytes.fetch_add(read_bytes, std::memory_order_relaxed); + release_ms = ssdperf::MsSince(t_release); + + if (ssdperf::Enabled()) { + const double total_ms = ssdperf::MsSince(t_begin); + size_t served = 0; + for (size_t j = 0; j < todo.size(); ++j) { + if (ok[j]) ++served; + } + for (size_t w = 0; w < waits.size(); ++w) { + if (wait_ok[w]) ++served; + } + const uint64_t dup_total = metrics_.read_dup.load(std::memory_order_relaxed); + const uint64_t lead_total = metrics_.read_lead.load(std::memory_order_relaxed); + const uint64_t merged_total = metrics_.read_merged.load(std::memory_order_relaxed); + const uint64_t probe_total = dup_total + lead_total; + MORI_UMBP_INFO( + "[SsdPerf/peer] GET keys={} device_reads={} served={} not_owned={} bytes={} " + "total_ms={:.3f} GB_s={:.2f} | resolve={:.3f}ms backend={:.3f}ms ({:.0f}%) " + "fanout={:.3f}ms ({:.0f}%) wait={:.3f}ms ({:.0f}%) release_lru={:.3f}ms | " + "fanout_bytes={} fanout_count={} fanout_GB_s={:.2f} | " + "singleflight={} dup={} merged={} max_concurrent={} " + "dup_total={}/{} ({:.1f}%) merged_total={}", + keys.size(), todo.size(), served, not_owned, read_bytes, total_ms, + ssdperf::GbPerSec(read_bytes, total_ms), resolve_ms, backend_ms, + ssdperf::Pct(backend_ms, total_ms), fanout_ms, ssdperf::Pct(fanout_ms, total_ms), wait_ms, + ssdperf::Pct(wait_ms, total_ms), release_ms, fanout_bytes, fanout_count, + ssdperf::GbPerSec(fanout_bytes, fanout_ms), single_flight_ ? "on" : "off", dup, merged, + max_concurrent, dup_total, probe_total, + ssdperf::Pct(static_cast(dup_total), probe_total), merged_total); + } + return out; +} + +// --------------------------------------------------------------------------- +// OwnedLocationSource (heartbeat event drain / snapshot) +// --------------------------------------------------------------------------- + +void PeerSsdManager::SetEventPublishing(bool enabled) { + std::lock_guard lock(mutex_); + event_publishing_ = enabled; + if (!enabled) pending_events_.clear(); +} + +std::vector PeerSsdManager::DrainPendingEvents() { + std::lock_guard lock(mutex_); + std::vector drained; + drained.swap(pending_events_); + return drained; +} + +std::vector PeerSsdManager::SnapshotOwnedKeysLocked() const { + std::vector out; + out.reserve(owned_.size()); + for (const auto& [key, entry] : owned_) { + out.push_back(KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, entry.size}); + } + return out; +} + +std::vector PeerSsdManager::SnapshotOwnedKeys() const { + std::lock_guard lock(mutex_); + return SnapshotOwnedKeysLocked(); +} + +std::vector PeerSsdManager::SnapshotOwnedKeysForFullSync() { + std::lock_guard lock(mutex_); + auto out = SnapshotOwnedKeysLocked(); + // Snapshot is authoritative now: drop the queued delta (already reflected in + // it) so the next delta carries only events produced after this snapshot. + pending_events_.clear(); + return out; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/peer/ssd_copy_pipeline.cpp b/src/umbp/distributed/peer/ssd/ssd_copy_pipeline.cpp similarity index 94% rename from src/umbp/distributed/peer/ssd_copy_pipeline.cpp rename to src/umbp/distributed/peer/ssd/ssd_copy_pipeline.cpp index ff223db6b..3bf703357 100644 --- a/src/umbp/distributed/peer/ssd_copy_pipeline.cpp +++ b/src/umbp/distributed/peer/ssd/ssd_copy_pipeline.cpp @@ -19,14 +19,14 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. -#include "umbp/distributed/peer/ssd_copy_pipeline.h" +#include "umbp/distributed/peer/ssd/ssd_copy_pipeline.h" #include #include #include "mori/utils/mori_log.hpp" -#include "umbp/distributed/peer/peer_dram_allocator.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" +#include "umbp/distributed/peer/backend/page_backend.h" +#include "umbp/distributed/peer/ssd/peer_ssd_manager.h" namespace mori::umbp { @@ -37,7 +37,7 @@ namespace { // pages against eviction. class DramCopyPinGuard { public: - DramCopyPinGuard(PeerDramAllocator* dram, std::string key, uint64_t token) + DramCopyPinGuard(PageBackend* dram, std::string key, uint64_t token) : dram_(dram), key_(std::move(key)), token_(token) {} ~DramCopyPinGuard() { dram_->ReleaseDramCopyPin(key_, token_); } @@ -45,14 +45,14 @@ class DramCopyPinGuard { DramCopyPinGuard& operator=(const DramCopyPinGuard&) = delete; private: - PeerDramAllocator* dram_; + PageBackend* dram_; std::string key_; uint64_t token_; }; } // namespace -SsdCopyPipeline::SsdCopyPipeline(PeerDramAllocator* dram, PeerSsdManager* ssd, size_t queue_depth, +SsdCopyPipeline::SsdCopyPipeline(PageBackend* dram, PeerSsdManager* ssd, size_t queue_depth, size_t worker_threads) : dram_(dram), ssd_(ssd), diff --git a/src/umbp/distributed/pool/peer_pool.cpp b/src/umbp/distributed/pool/peer_pool.cpp new file mode 100644 index 000000000..898000485 --- /dev/null +++ b/src/umbp/distributed/pool/peer_pool.cpp @@ -0,0 +1,1044 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#include "umbp/distributed/pool/peer_pool.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mori::umbp { +namespace { + +// One watermark scan queues at most this many keys; the worker tops the queue +// up after each completed offload, so a tier still drains all the way down to +// its low watermark without ever queueing its whole key set at once. +constexpr size_t kWatermarkCandidateBatch = 64; +// Floor between scans. Without it every commit batch arriving above the high +// watermark would re-walk the placements. +constexpr auto kWatermarkScanInterval = std::chrono::milliseconds(5); + +} // namespace + +PeerPool::PeerPool(BackendRegistry* backends, std::unique_ptr policy, + TransferEngine* transfer_engine) + : backends_(backends), + policy_(std::move(policy)), + transfer_engine_(transfer_engine), + tier_graph_(policy_ == nullptr ? nullptr : policy_->TierGraph()), + transition_executor_(backends, transfer_engine) { + if (tier_graph_ != nullptr) { + // A tier graph without an engine would accept offload and promotion rules + // and then silently fail every one of them, leaving the topology looking + // configured while nothing ever moves. + if (transfer_engine == nullptr) { + throw std::invalid_argument( + "PeerPool: a tiered policy requires a TransferEngine to move bytes between " + "backends"); + } + tier_scan_backoff_.assign(tier_graph_->TierCount(), std::chrono::steady_clock::time_point{}); + queued_by_tier_.assign(tier_graph_->TierCount(), 0); + for (size_t tier = 0; tier < tier_graph_->TierCount(); ++tier) { + const auto& node = tier_graph_->NodeAt(tier); + if (node.trigger == PoolOffloadTrigger::kWatermark && !node.offload_to.empty()) { + track_access_order_ = true; + break; + } + } + transition_worker_ = std::thread(&PeerPool::TransitionWorkerLoop, this); + } +} + +PeerPool::~PeerPool() { StopTransitionWorker(); } + +void PeerPool::EnqueueTransition(TransitionJob job) { + std::lock_guard lock(transition_mutex_); + if (stop_transition_worker_ || !queued_transition_keys_.insert(job.key).second) return; + if (job.source_tier < queued_by_tier_.size()) ++queued_by_tier_[job.source_tier]; + transition_queue_.push_back(std::move(job)); + transition_cv_.notify_one(); +} + +void PeerPool::StopTransitionWorker() { + { + std::lock_guard lock(transition_mutex_); + stop_transition_worker_ = true; + transition_queue_.clear(); + queued_transition_keys_.clear(); + std::fill(queued_by_tier_.begin(), queued_by_tier_.end(), 0); + } + transition_cv_.notify_all(); + if (transition_worker_.joinable()) transition_worker_.join(); +} + +void PeerPool::TouchLocked(const std::string& key) { + if (!track_access_order_) return; + auto access = last_access_.find(key); + if (access != last_access_.end()) { + access_order_.splice(access_order_.begin(), access_order_, access->second); + return; + } + access_order_.push_front(key); + last_access_.emplace(key, access_order_.begin()); +} + +void PeerPool::ForgetAccessLocked(const std::string& key) { + promote_hit_counts_.erase(key); + auto access = last_access_.find(key); + if (access == last_access_.end()) return; + access_order_.erase(access->second); + last_access_.erase(access); +} + +bool PeerPool::PromoteOnReadLocked(const std::string& key, + LogicalTierGraph::TierIndex source_tier) { + const auto& tier = tier_graph_->NodeAt(source_tier); + switch (tier.promote_trigger) { + case PoolPromoteTrigger::kNever: + return false; + case PoolPromoteTrigger::kOnRead: + return true; + case PoolPromoteTrigger::kOnHits: { + const uint32_t hits = ++promote_hit_counts_[key]; + if (hits < tier.promote_hits) return false; + // The count has done its job; leaving it behind would promote again on + // the next read if this promotion does not move the key. + promote_hit_counts_.erase(key); + return true; + } + } + return false; +} + +uint32_t PeerPool::FindOwnerLocked(const std::string& key) const { + if (policy_ == nullptr || backends_ == nullptr) return BackendRegistry::kMaxBackends; + for (uint32_t backend_id : policy_->ReadOrder(*backends_)) { + auto* backend = backends_->Get(backend_id); + if (backend != nullptr && backend->Contains(key)) return backend_id; + } + return BackendRegistry::kMaxBackends; +} + +size_t PeerPool::EnqueueWatermarkCandidatesLocked(LogicalTierGraph::TierIndex tier_index, + size_t max_count) { + if (tier_graph_ == nullptr || max_count == 0) return 0; + const auto& tier = tier_graph_->NodeAt(tier_index); + if (tier.trigger != PoolOffloadTrigger::kWatermark || tier.offload_to.empty()) return 0; + + std::vector candidates; + const auto consider = [&](const std::string& key, uint32_t backend_id) { + if (tier_graph_->TierIndexForBackendId(backend_id) != tier_index) return; + // Draining a member that is not itself under pressure frees nothing that + // was scarce, and in a mixed-medium tier it would move keys off the roomy + // member while the full one stays full. + if (!tier_graph_->MemberAtOrAbove(backend_id, tier.low_watermark)) return; + if (migrating_keys_.count(key) != 0) return; + candidates.push_back({TierTransitionKind::kOffload, key, backend_id, tier_index}); + }; + + // Coldest first: access_order_ keeps the hottest key at the front, so walking + // it backwards reaches the keys least likely to be read again first. + for (auto it = access_order_.rbegin(); it != access_order_.rend(); ++it) { + const std::string& key = *it; + auto placement = placements_.find(key); + if (placement == placements_.end()) continue; + consider(key, placement->second); + if (candidates.size() >= max_count) break; + } + for (auto& candidate : candidates) EnqueueTransition(std::move(candidate)); + return candidates.size(); +} + +void PeerPool::MaybeEnqueueWatermarkOffloadLocked() { + if (tier_graph_ == nullptr) return; + const auto now = std::chrono::steady_clock::now(); + for (size_t tier_index = 0; tier_index < tier_graph_->TierCount(); ++tier_index) { + const auto& tier = tier_graph_->NodeAt(tier_index); + if (tier.trigger != PoolOffloadTrigger::kWatermark || tier.offload_to.empty() || + now < tier_scan_backoff_[tier_index] || + !tier_graph_->AtOrAbove(tier_index, tier.high_watermark)) { + continue; + } + { + // Work is already queued for this tier; the worker tops it up as it + // drains, so there is nothing to decide here. + std::lock_guard lock(transition_mutex_); + if (queued_by_tier_[tier_index] > 0) continue; + } + if (EnqueueWatermarkCandidatesLocked(tier_index, kWatermarkCandidateBatch) == 0) { + tier_scan_backoff_[tier_index] = now + kWatermarkScanInterval; + } + } +} + +bool PeerPool::WatermarkDrivenLocked(const TransitionJob& job) const { + return job.kind == TierTransitionKind::kOffload && tier_graph_ != nullptr && + tier_graph_->NodeAt(job.source_tier).trigger == PoolOffloadTrigger::kWatermark; +} + +PeerPool::TransitionPlan PeerPool::PlanTransitionLocked(const TransitionJob& job) { + TransitionPlan plan; + if (tier_graph_ == nullptr) return plan; + auto placement = placements_.find(job.key); + if (placement == placements_.end() || placement->second != job.source_backend_id) { + return plan; + } + if (!migrating_keys_.insert(job.key).second) return plan; + + const bool promotion = job.kind == TierTransitionKind::kPromotion; + const auto& source_tier = tier_graph_->NodeAt(job.source_tier); + plan.remove_source = !promotion || source_tier.promote_mode == PoolTransitionMode::kMove; + plan.targets = promotion ? tier_graph_->PromoteTargetOrder(job.source_tier, job.key) + : tier_graph_->TransitionTargetOrder(job.source_tier, job.key); + plan.valid = true; + ++transition_metrics_.attempted; + return plan; +} + +void PeerPool::FinishTransitionLocked(const TransitionJob& job, + const TierTransitionResult& result) { + migrating_keys_.erase(job.key); + if (migrating_keys_.empty()) transition_idle_cv_.notify_all(); + if (!result.success) { + ++transition_metrics_.failed; + return; + } + ++transition_metrics_.succeeded; + if (job.kind == TierTransitionKind::kPromotion) { + transition_metrics_.promoted_bytes += result.bytes_moved; + } else { + transition_metrics_.offloaded_bytes += result.bytes_moved; + } + // The copy ran with the pool lock released, so another operation may have + // re-placed the key meanwhile. Only the placement this transition started + // from may be advanced to the target. + auto placement = placements_.find(job.key); + if (placement != placements_.end() && placement->second == job.source_backend_id) { + placement->second = result.target_backend_id; + } + if (result.source_draining) { + draining_sources_[job.key].insert(job.source_backend_id); + } else { + auto draining = draining_sources_.find(job.key); + if (draining != draining_sources_.end()) { + draining->second.erase(job.source_backend_id); + if (draining->second.empty()) draining_sources_.erase(draining); + } + } +} + +TierTransitionResult PeerPool::RunTransition(std::unique_lock& lock, + const TransitionJob& job) { + TierTransitionResult result; + const auto plan = PlanTransitionLocked(job); + if (!plan.valid) return result; + + lock.unlock(); + result = transition_executor_.Execute(job.key, job.source_backend_id, plan.targets, + plan.remove_source); + lock.lock(); + FinishTransitionLocked(job, result); + return result; +} + +void PeerPool::TransitionWorkerLoop() { + for (;;) { + TransitionJob job; + { + std::unique_lock lock(transition_mutex_); + transition_cv_.wait(lock, + [&] { return stop_transition_worker_ || !transition_queue_.empty(); }); + if (stop_transition_worker_ && transition_queue_.empty()) return; + job = std::move(transition_queue_.front()); + transition_queue_.pop_front(); + } + bool retry = false; + { + std::unique_lock lock(operation_mutex_); + if (tier_graph_ != nullptr) { + const auto& tier = tier_graph_->NodeAt(job.source_tier); + const bool watermark_driven = WatermarkDrivenLocked(job); + if (!watermark_driven || tier_graph_->AtOrAbove(job.source_tier, tier.low_watermark)) { + const auto result = RunTransition(lock, job); + const bool still_pressured = + watermark_driven && tier_graph_->AtOrAbove(job.source_tier, tier.low_watermark); + retry = still_pressured && !result.success && job.attempts < 3; + // Keep the drain going without waiting for the next commit batch to + // notice the tier is still above its low watermark. The scan skips + // whatever is still queued, so this tops the batch back up rather + // than duplicating it. + if (still_pressured && result.success) { + tier_scan_backoff_[job.source_tier] = std::chrono::steady_clock::time_point{}; + EnqueueWatermarkCandidatesLocked(job.source_tier, kWatermarkCandidateBatch); + } + } + } + } + if (retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(uint64_t{1} << job.attempts)); + } + { + std::lock_guard lock(transition_mutex_); + queued_transition_keys_.erase(job.key); + const size_t source_tier = job.source_tier; + if (retry && !stop_transition_worker_) { + ++job.attempts; + queued_transition_keys_.insert(job.key); + transition_queue_.push_back(std::move(job)); + transition_cv_.notify_one(); + } else if (source_tier < queued_by_tier_.size() && queued_by_tier_[source_tier] > 0) { + --queued_by_tier_[source_tier]; + } + } + } +} + +std::vector PeerPool::BatchAllocate( + const std::vector& requests) { + std::lock_guard operation_lock(operation_mutex_); + std::vector out(requests.size()); + if (backends_ == nullptr || policy_ == nullptr) return out; + + std::vector> put_orders(requests.size()); + std::vector next_candidate(requests.size(), 0); + std::vector active; + active.reserve(requests.size()); + for (size_t i = 0; i < requests.size(); ++i) { + auto existing = placements_.find(requests[i].key); + if (existing != placements_.end()) { + auto* existing_backend = backends_->Get(existing->second); + if (existing_backend != nullptr && existing_backend->Contains(requests[i].key)) { + out[i].backend_id = existing->second; + out[i].allocation.outcome = AllocateOutcome::kSuccessAlreadyExists; + continue; + } + placements_.erase(existing); + } + auto pending = pending_keys_.find(requests[i].key); + if (pending != pending_keys_.end()) { + if (pending->second.expires_at <= std::chrono::steady_clock::now()) { + auto* pending_backend = backends_->Get(pending->second.backend_id); + if (pending_backend != nullptr && pending->second.local_slot_id != 0) { + pending_backend->BatchAbort({pending->second.local_slot_id}); + } + pending_slots_.erase({pending->second.backend_id, pending->second.local_slot_id}); + pending_keys_.erase(pending); + } else { + // A second writer must retry until the reservation commits, aborts, or + // reaches the same TTL as the backend slot. ALREADY_EXISTS would expose + // uncommitted bytes. + out[i].allocation.outcome = AllocateOutcome::kFailed; + continue; + } + } + + // Rebuild the index before making a new placement decision: a persistent + // backend may still own the key even though this process has no record. + const uint32_t owner = FindOwnerLocked(requests[i].key); + if (owner != BackendRegistry::kMaxBackends) { + placements_[requests[i].key] = owner; + out[i].backend_id = owner; + out[i].allocation.outcome = AllocateOutcome::kSuccessAlreadyExists; + continue; + } + + put_orders[i] = policy_->PutOrder(*backends_, requests[i]); + if (!put_orders[i].empty()) { + const uint32_t selected = put_orders[i].front(); + pending_keys_[requests[i].key] = + PendingPlacement{selected, 0, std::chrono::steady_clock::time_point::max()}; + active.push_back(i); + } + } + + // Each round batches the requests targeting the same candidate. Only + // NO_SPACE advances to the next deterministic policy candidate; generic + // failures preserve their meaning and do not hide backend errors. + while (!active.empty()) { + std::map> by_backend; + for (size_t index : active) { + by_backend[put_orders[index][next_candidate[index]]].push_back(index); + } + std::vector retry; + // Moves a request onto its next policy candidate, or drops the reservation + // once the candidates run out. + auto advance = [&](size_t index) { + ++next_candidate[index]; + if (next_candidate[index] >= put_orders[index].size()) { + pending_keys_.erase(requests[index].key); + return; + } + pending_keys_[requests[index].key] = + PendingPlacement{put_orders[index][next_candidate[index]], 0, + std::chrono::steady_clock::time_point::max()}; + retry.push_back(index); + }; + + for (const auto& [backend_id, indices] : by_backend) { + auto* backend = backends_->Get(backend_id); + if (backend == nullptr) { + for (size_t index : indices) advance(index); + continue; + } + + std::vector backend_requests; + backend_requests.reserve(indices.size()); + for (size_t index : indices) { + backend_requests.push_back(AllocateRequest{requests[index].key, requests[index].size}); + } + + auto results = backend->BatchAllocate(backend_requests); + for (size_t i = 0; i < indices.size(); ++i) { + const size_t index = indices[i]; + if (i >= results.size()) { + pending_keys_.erase(requests[index].key); + continue; + } + out[index].backend_id = backend_id; + out[index].allocation = std::move(results[i]); + switch (out[index].allocation.outcome) { + case AllocateOutcome::kSuccessAllocated: { + const uint64_t slot_id = out[index].allocation.slot_id; + const auto ttl_ms = out[index].allocation.pending_ttl_ms; + const auto expires_at = + ttl_ms == 0 ? std::chrono::steady_clock::time_point::max() + : std::chrono::steady_clock::now() + std::chrono::milliseconds(ttl_ms); + pending_keys_[requests[index].key] = PendingPlacement{backend_id, slot_id, expires_at}; + pending_slots_[{backend_id, slot_id}] = requests[index].key; + break; + } + case AllocateOutcome::kSuccessAlreadyExists: + placements_[requests[index].key] = backend_id; + pending_keys_.erase(requests[index].key); + break; + case AllocateOutcome::kFailed: + pending_keys_.erase(requests[index].key); + break; + case AllocateOutcome::kFailedNoSpace: + advance(index); + break; + } + } + } + active = std::move(retry); + } + return out; +} + +std::vector PeerPool::BatchCommit( + const std::vector& requests) { + std::lock_guard operation_lock(operation_mutex_); + std::vector out(requests.size()); + if (backends_ == nullptr) return out; + + const auto release_pending = [&](const PoolSlotRef& slot) { + auto pending = pending_slots_.find({slot.backend_id, slot.local_slot_id}); + if (pending == pending_slots_.end()) return; + pending_keys_.erase(pending->second); + pending_slots_.erase(pending); + }; + + std::map> by_backend; + std::set> accepted_slots; + for (size_t i = 0; i < requests.size(); ++i) { + const auto slot_key = + std::make_pair(requests[i].slot.backend_id, requests[i].slot.local_slot_id); + if (requests[i].slot.backend_id >= BackendRegistry::kMaxBackends || + accepted_slots.find(slot_key) != accepted_slots.end()) { + continue; + } + auto slot = pending_slots_.find(slot_key); + if (slot == pending_slots_.end() || slot->second != requests[i].key) continue; + auto pending = pending_keys_.find(slot->second); + if (pending == pending_keys_.end() || + pending->second.backend_id != requests[i].slot.backend_id || + pending->second.local_slot_id != requests[i].slot.local_slot_id) { + continue; + } + if (pending->second.expires_at <= std::chrono::steady_clock::now()) { + auto* backend = backends_->Get(requests[i].slot.backend_id); + if (backend != nullptr) backend->BatchAbort({requests[i].slot.local_slot_id}); + release_pending(requests[i].slot); + continue; + } + accepted_slots.insert(slot_key); + by_backend[requests[i].slot.backend_id].push_back(i); + } + + for (const auto& [backend_id, indices] : by_backend) { + auto* backend = backends_->Get(backend_id); + if (backend == nullptr) { + for (size_t index : indices) release_pending(requests[index].slot); + continue; + } + + std::vector backend_requests; + backend_requests.reserve(indices.size()); + for (size_t index : indices) { + backend_requests.push_back( + CommitRequest{requests[index].slot.local_slot_id, requests[index].key}); + } + + auto results = backend->BatchCommit(backend_requests); + for (size_t i = 0; i < indices.size(); ++i) { + const size_t index = indices[i]; + if (i >= results.size()) { + release_pending(requests[index].slot); + continue; + } + out[index].backend_id = backend_id; + out[index].commit = results[i]; + if (results[i].success) { + placements_[requests[index].key] = backend_id; + TouchLocked(requests[index].key); + } + release_pending(requests[index].slot); + } + } + MaybeEnqueueWatermarkOffloadLocked(); + return out; +} + +std::vector PeerPool::BatchAbort(const std::vector& slots) { + std::lock_guard operation_lock(operation_mutex_); + std::vector out(slots.size(), true); + if (backends_ == nullptr) return out; + + const auto release_pending = [&](const PoolSlotRef& slot) { + auto pending = pending_slots_.find({slot.backend_id, slot.local_slot_id}); + if (pending == pending_slots_.end()) return; + pending_keys_.erase(pending->second); + pending_slots_.erase(pending); + }; + + std::map> by_backend; + for (size_t i = 0; i < slots.size(); ++i) { + if (slots[i].backend_id < BackendRegistry::kMaxBackends) { + by_backend[slots[i].backend_id].push_back(i); + } + } + + for (const auto& [backend_id, indices] : by_backend) { + auto* backend = backends_->Get(backend_id); + if (backend == nullptr) { + for (size_t index : indices) release_pending(slots[index]); + continue; + } + + std::vector local_slot_ids; + local_slot_ids.reserve(indices.size()); + for (size_t index : indices) local_slot_ids.push_back(slots[index].local_slot_id); + + auto results = backend->BatchAbort(local_slot_ids); + for (size_t i = 0; i < indices.size(); ++i) { + if (i < results.size()) out[indices[i]] = results[i]; + release_pending(slots[indices[i]]); + } + } + return out; +} + +std::vector PeerPool::BatchResolve(const std::vector& keys, + bool include_descs, bool allow_file_refs) { + std::shared_lock lifecycle_lock(lifecycle_mutex_); + std::vector out(keys.size()); + if (keys.empty()) return out; + + // A no-tier, single-backend pool needs placement metadata for allocate, + // commit, and eviction, but a hit whose placement already names that sole + // backend cannot repair or transition anything. Resolve it directly and + // avoid building the preferred/attempted/group vectors plus taking the + // metadata lock again after backend IO. This is the common no-master, + // 100%-hit restore shape. + uint32_t direct_backend_id = BackendRegistry::kMaxBackends; + MediumBackend* direct_backend = nullptr; + { + std::lock_guard operation_lock(operation_mutex_); + if (backends_ != nullptr && policy_ != nullptr && tier_graph_ == nullptr && + !track_access_order_ && backends_->Size() == 1) { + const std::vector order = policy_->ReadOrder(*backends_); + if (order.size() == 1 && order.front() < BackendRegistry::kMaxBackends) { + const uint32_t backend_id = order.front(); + bool all_placed_here = true; + for (const auto& key : keys) { + const auto placement = placements_.find(key); + if (placement == placements_.end() || placement->second != backend_id) { + all_placed_here = false; + break; + } + } + if (all_placed_here) { + direct_backend_id = backend_id; + direct_backend = backends_->Get(backend_id); + } + } + } + } + if (direct_backend != nullptr) { + auto resolved = direct_backend->BatchResolve(keys, include_descs, allow_file_refs); + bool batch_busy = false; + bool all_found = resolved.size() == keys.size(); + std::vector failed_but_owned(keys.size(), false); + for (size_t i = 0; i < resolved.size() && i < out.size(); ++i) { + const ResolveOutcome outcome = EffectiveResolveOutcome(resolved[i]); + batch_busy |= outcome == ResolveOutcome::kBusy; + all_found &= outcome == ResolveOutcome::kFound; + if (outcome == ResolveOutcome::kFailed) { + failed_but_owned[i] = direct_backend->Contains(keys[i]); + } + out[i].backend_id = direct_backend_id; + out[i].tier = direct_backend->Tier(); + out[i].resolved = std::move(resolved[i]); + out[i].resolved.outcome = outcome; + } + if (batch_busy || all_found) return out; + // A plain miss is authoritative, while kFailed needs one ownership probe: + // a backend can still own a key whose descriptor shape it cannot return. + // Finalize stale placements here instead of running the whole routing + // pipeline and resolving the same backend a second time. + std::lock_guard operation_lock(operation_mutex_); + for (size_t i = 0; i < out.size(); ++i) { + if (out[i].resolved.found || failed_but_owned[i]) continue; + const auto placement = placements_.find(keys[i]); + if (placement == placements_.end() || placement->second != direct_backend_id) continue; + placements_.erase(placement); + ForgetAccessLocked(keys[i]); + } + return out; + } + + uint64_t planned_clear_epoch = 0; + std::vector> preferred(keys.size()); + std::vector read_order; + { + std::lock_guard operation_lock(operation_mutex_); + if (backends_ == nullptr || policy_ == nullptr) return out; + planned_clear_epoch = clear_epoch_; + for (size_t i = 0; i < keys.size(); ++i) { + auto it = placements_.find(keys[i]); + if (it != placements_.end()) preferred[i] = it->second; + } + read_order = policy_->ReadOrder(*backends_); + } + + // A flat mask avoids one heap allocation per key in the common 1023-key + // ranged-get batch. + static_assert(BackendRegistry::kMaxBackends <= 16, "attempted mask is 16 bits wide"); + std::vector attempted(keys.size(), 0); + + const auto resolve = [&](uint32_t backend_id, const std::vector& indices) { + auto* backend = backends_->Get(backend_id); + if (backend == nullptr || indices.empty()) return; + bool identity = indices.size() == keys.size(); + for (size_t i = 0; identity && i < indices.size(); ++i) identity = indices[i] == i; + std::vector backend_keys; + if (!identity) backend_keys.reserve(indices.size()); + for (size_t index : indices) { + attempted[index] |= static_cast(1u << backend_id); + if (!identity) backend_keys.push_back(keys[index]); + } + auto results = + backend->BatchResolve(identity ? keys : backend_keys, include_descs, allow_file_refs); + for (size_t i = 0; i < indices.size() && i < results.size(); ++i) { + const size_t index = indices[i]; + const ResolveOutcome outcome = EffectiveResolveOutcome(results[i]); + if (outcome == ResolveOutcome::kFound) { + if (out[index].resolved.found) continue; + out[index].backend_id = backend_id; + out[index].tier = backend->Tier(); + out[index].resolved = std::move(results[i]); + out[index].resolved.outcome = ResolveOutcome::kFound; + continue; + } + if (out[index].resolved.found) continue; + // Preserve a retryable owner over misses from fallback media. A + // permanent backend failure is useful only when no backend reported + // BUSY; BUSY has priority because a later attempt can still succeed. + const ResolveOutcome current = EffectiveResolveOutcome(out[index].resolved); + if (outcome == ResolveOutcome::kBusy || current == ResolveOutcome::kMissing) { + out[index].backend_id = backend_id; + out[index].tier = backend->Tier(); + out[index].resolved.outcome = outcome; + } + } + }; + + std::map> preferred_groups; + for (size_t i = 0; i < preferred.size(); ++i) { + if (preferred[i].has_value() && *preferred[i] < BackendRegistry::kMaxBackends) { + preferred_groups[*preferred[i]].push_back(i); + } + } + for (const auto& [backend_id, indices] : preferred_groups) resolve(backend_id, indices); + + for (uint32_t backend_id : read_order) { + if (backend_id >= BackendRegistry::kMaxBackends) continue; + std::vector unresolved; + for (size_t i = 0; i < keys.size(); ++i) { + if (!out[i].resolved.found && (attempted[i] & (1u << backend_id)) == 0) { + unresolved.push_back(i); + } + } + resolve(backend_id, unresolved); + } + + // Classify once. A 100%-hit restore must not allocate a key-sized repair + // vector merely to discover that there is nothing to repair. + bool batch_busy = false; + bool all_found = true; + for (const auto& entry : out) { + batch_busy |= EffectiveResolveOutcome(entry.resolved) == ResolveOutcome::kBusy; + all_found &= entry.resolved.found; + } + // A caller retries the whole batch on BUSY and discards this response. Do + // not count hits, repair placements, or enqueue promotions for a response + // whose page locations will never be consumed. + if (batch_busy) return out; + + // Contains is backend work too; keep it off the pool metadata lock. A + // placement snapshot check below prevents this repair from overwriting a + // commit or migration that publishes while these probes are running. + std::vector> repaired; + if (!all_found) { + repaired.resize(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + if (out[i].resolved.found) continue; + // kMissing is an authoritative negative result from every attempted + // backend. Only kFailed is ambiguous: the selected backend may still own + // the key but be unable to materialize its descriptor. + if (EffectiveResolveOutcome(out[i].resolved) != ResolveOutcome::kFailed) continue; + auto* backend = backends_->Get(out[i].backend_id); + if (backend != nullptr && backend->Contains(keys[i])) { + repaired[i] = out[i].backend_id; + } + } + } + + std::lock_guard operation_lock(operation_mutex_); + // ClearLocal takes lifecycle_mutex_ exclusively, so this is defensive + // against future lifecycle paths that may invalidate an in-flight plan. + if (clear_epoch_ != planned_clear_epoch) return out; + + const auto placement_unchanged = [&](size_t index) { + auto current = placements_.find(keys[index]); + if (!preferred[index].has_value()) return current == placements_.end(); + return current != placements_.end() && current->second == *preferred[index]; + }; + + for (size_t i = 0; i < keys.size(); ++i) { + if (out[i].resolved.found) { + // A migration or commit may have published a newer owner while backend + // IO ran. Repair only the placement snapshot this resolve actually used; + // never move the logical owner backwards to a source that is draining. + const bool metadata_current = placement_unchanged(i); + if (metadata_current) placements_[keys[i]] = out[i].backend_id; + // A concurrent discard may have removed the key entirely. Do not + // resurrect only its LRU entry from a resolve that started earlier. + const bool still_placed = metadata_current || placements_.find(keys[i]) != placements_.end(); + if (still_placed) TouchLocked(keys[i]); + if (tier_graph_ != nullptr) { + ++tier_read_hits_[tier_graph_->NameForBackend(out[i].backend_id)]; + const auto source_tier = tier_graph_->TierIndexForBackendId(out[i].backend_id); + if (metadata_current && source_tier.has_value() && + *source_tier != tier_graph_->EntryTierIndex() && migrating_keys_.count(keys[i]) == 0 && + !tier_graph_->AtOrAbove( + tier_graph_->EntryTierIndex(), + tier_graph_->NodeAt(tier_graph_->EntryTierIndex()).high_watermark) && + PromoteOnReadLocked(keys[i], *source_tier)) { + EnqueueTransition( + {TierTransitionKind::kPromotion, keys[i], out[i].backend_id, *source_tier}); + } + } + } else { + // A concurrent commit or transition changed the placement after the + // snapshot; its publication is newer than this miss and must win. + if (!placement_unchanged(i)) continue; + if (repaired[i].has_value()) { + placements_[keys[i]] = *repaired[i]; + continue; + } + // A key the backends no longer hold must leave the access index too, or + // it accumulates there for the life of the process and skews LRU order. + placements_.erase(keys[i]); + ForgetAccessLocked(keys[i]); + } + } + return out; +} + +std::vector PeerPool::BatchContains(const std::vector& keys) { + std::shared_lock lifecycle_lock(lifecycle_mutex_); + std::vector out(keys.size(), false); + if (keys.empty()) return out; + + std::vector read_order; + { + std::lock_guard operation_lock(operation_mutex_); + if (backends_ == nullptr || policy_ == nullptr) return out; + read_order = policy_->ReadOrder(*backends_); + } + + // Carry only the still-unanswered indices from one backend to the next, so a + // second medium is asked about exactly what the first one missed. Same shape + // as the resolve paths, and for the same reason: one call per backend, not + // one per key. The backend calls stay outside operation_mutex_. + std::vector pending(keys.size()); + std::iota(pending.begin(), pending.end(), 0); + std::vector batch; + std::vector still_missing; + for (uint32_t backend_id : read_order) { + if (pending.empty()) break; + auto* backend = backends_->Get(backend_id); + if (backend == nullptr) continue; + batch.clear(); + batch.reserve(pending.size()); + for (size_t i : pending) batch.push_back(keys[i]); + + const auto hits = backend->BatchContains(batch); + still_missing.clear(); + for (size_t j = 0; j < pending.size(); ++j) { + if (j < hits.size() && hits[j]) { + out[pending[j]] = true; + } else { + still_missing.push_back(pending[j]); + } + } + pending.swap(still_missing); + } + return out; +} + +std::vector PeerPool::Evict(const std::vector& keys, PoolEvictMode mode) { + std::unique_lock operation_lock(operation_mutex_); + std::vector out; + out.reserve(keys.size()); + for (const auto& key : keys) out.push_back(EvictResult{key, 0}); + if (backends_ == nullptr || keys.empty()) return out; + + std::vector handled(keys.size(), false); + for (size_t i = 0; i < keys.size(); ++i) { + // A copy already owns this key. Freeing nothing makes the caller retry, + // which is the same contract as a read lease delaying a source delete, and + // it keeps a discard from racing a target commit that would resurrect it. + if (migrating_keys_.count(keys[i]) != 0) { + handled[i] = true; + continue; + } + // A discard has to leave no copy behind, so it skips demotion and the + // drain shortcut and goes straight to the delete pass below. + if (mode == PoolEvictMode::kDiscard) continue; + + auto draining = draining_sources_.find(keys[i]); + if (draining != draining_sources_.end()) { + std::vector drained; + for (uint32_t source_id : draining->second) { + auto* source = backends_->Get(source_id); + if (source != nullptr) { + auto retried = source->Evict({keys[i]}); + if (!retried.empty()) out[i].bytes_freed += retried.front().bytes_freed; + } + if (source == nullptr || !source->Contains(keys[i])) { + drained.push_back(source_id); + } + } + for (uint32_t source_id : drained) draining->second.erase(source_id); + if (draining->second.empty()) draining_sources_.erase(draining); + handled[i] = true; + continue; + } + + auto placement = placements_.find(keys[i]); + const uint32_t source_id = + placement != placements_.end() ? placement->second : FindOwnerLocked(keys[i]); + if (tier_graph_ == nullptr) continue; + const auto tier_index = tier_graph_->TierIndexForBackendId(source_id); + if (!tier_index.has_value()) continue; + const auto& tier = tier_graph_->NodeAt(*tier_index); + if (tier.trigger != PoolOffloadTrigger::kOnEvict || tier.offload_to.empty()) { + continue; + } + auto migrated = RunTransition(operation_lock, + {TierTransitionKind::kOffload, keys[i], source_id, *tier_index}); + handled[i] = migrated.success; + out[i].bytes_freed = migrated.bytes_freed; + } + + std::vector delete_keys; + std::vector delete_indices; + for (size_t i = 0; i < keys.size(); ++i) { + if (!handled[i]) { + delete_keys.push_back(keys[i]); + delete_indices.push_back(i); + } + } + for (auto* backend : backends_->All()) { + auto results = backend->Evict(delete_keys); + for (size_t i = 0; i < results.size() && i < delete_indices.size(); ++i) { + out[delete_indices[i]].bytes_freed += results[i].bytes_freed; + } + } + + for (size_t i = 0; i < keys.size(); ++i) { + if (handled[i]) continue; + const uint32_t owner = FindOwnerLocked(keys[i]); + if (owner != BackendRegistry::kMaxBackends) { + placements_[keys[i]] = owner; + } else { + placements_.erase(keys[i]); + ForgetAccessLocked(keys[i]); + draining_sources_.erase(keys[i]); + } + } + return out; +} + +void PeerPool::ClearLocal() { + std::unique_lock lifecycle_lock(lifecycle_mutex_); + std::unique_lock operation_lock(operation_mutex_); + { + std::lock_guard lock(transition_mutex_); + transition_queue_.clear(); + queued_transition_keys_.clear(); + } + // Wiping state underneath an in-flight copy would let its target commit + // resurrect a key the user asked to discard, so drain the copies first. The + // queue is already empty, so no new key can enter the set. + transition_idle_cv_.wait(operation_lock, [&] { return migrating_keys_.empty(); }); + + if (backends_ != nullptr) { + for (auto* backend : backends_->All()) backend->ClearLocal(); + } + placements_.clear(); + draining_sources_.clear(); + pending_keys_.clear(); + pending_slots_.clear(); + last_access_.clear(); + access_order_.clear(); + promote_hit_counts_.clear(); + ++clear_epoch_; +} + +std::vector PeerPool::DrainPendingEvents() { + std::lock_guard operation_lock(operation_mutex_); + if (backends_ == nullptr) return {}; + struct EventState { + std::optional add; + bool saw_remove = false; + }; + std::map, EventState> states; + const auto all_backends = backends_->All(); + for (auto* backend : all_backends) { + auto drained = backend->DrainPendingEvents(); + const std::string logical = tier_graph_ == nullptr + ? std::string{} + : tier_graph_->NameForBackend(backends_->BackendId(backend)); + for (auto& event : drained) { + event.logical_tier = logical; + auto& state = states[{logical, event.tier, event.key}]; + if (event.kind == KvEvent::Kind::ADD) { + state.add = std::move(event); + } else { + state.saw_remove = true; + } + } + } + std::vector events; + events.reserve(states.size()); + for (auto& [identity, state] : states) { + const auto& [logical, tier, key] = identity; + const bool still_owned = + std::any_of(all_backends.begin(), all_backends.end(), [&](MediumBackend* backend) { + if (backend == nullptr || backend->Tier() != tier || !backend->Contains(key)) { + return false; + } + return tier_graph_ == nullptr || + tier_graph_->NameForBackend(backends_->BackendId(backend)) == logical; + }); + if (state.add.has_value() && (still_owned || !state.saw_remove)) { + events.push_back(std::move(*state.add)); + } else if (state.saw_remove && !still_owned) { + events.push_back(KvEvent{KvEvent::Kind::REMOVE, key, tier, 0, logical}); + } + } + return events; +} + +std::vector PeerPool::SnapshotOwnedKeysForFullSync() { + std::lock_guard operation_lock(operation_mutex_); + if (backends_ == nullptr) return {}; + std::map, KvEvent> unique; + for (auto* backend : backends_->All()) { + auto snapshot = backend->SnapshotOwnedKeysForFullSync(); + const std::string logical = tier_graph_ == nullptr + ? std::string{} + : tier_graph_->NameForBackend(backends_->BackendId(backend)); + for (auto& event : snapshot) { + event.logical_tier = logical; + unique[{logical, event.tier, event.key}] = std::move(event); + } + } + std::vector events; + events.reserve(unique.size()); + for (auto& item : unique) events.push_back(std::move(item.second)); + return events; +} + +std::map PeerPool::LogicalTierCapacities() const { + std::lock_guard lock(operation_mutex_); + return tier_graph_ == nullptr ? std::map{} + : tier_graph_->CapacitySnapshot(); +} + +std::string PeerPool::LogicalTierForBackend(uint32_t backend_id) const { + std::lock_guard lock(operation_mutex_); + return tier_graph_ == nullptr ? std::string{} : tier_graph_->NameForBackend(backend_id); +} + +std::optional PeerPool::PlacementBackend(const std::string& key) const { + std::lock_guard lock(operation_mutex_); + auto it = placements_.find(key); + return it == placements_.end() ? std::nullopt : std::optional{it->second}; +} + +size_t PeerPool::PlacementCount() const { + std::lock_guard lock(operation_mutex_); + return placements_.size(); +} + +TierTransitionMetrics PeerPool::TransitionMetrics() const { + std::lock_guard lock(operation_mutex_); + return transition_metrics_; +} + +std::map PeerPool::TierReadHits() const { + std::lock_guard lock(operation_mutex_); + return tier_read_hits_; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/pool/policy_config.cpp b/src/umbp/distributed/pool/policy_config.cpp new file mode 100644 index 000000000..8479e2c4c --- /dev/null +++ b/src/umbp/distributed/pool/policy_config.cpp @@ -0,0 +1,723 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License + +#include "umbp/distributed/pool/policy_config.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/config.h" + +namespace mori::umbp { +namespace { + +using google::protobuf::ListValue; +using google::protobuf::Struct; +using google::protobuf::Value; + +[[noreturn]] void Invalid(const std::string& message) { throw std::runtime_error(message); } + +const char* ValueTypeName(const Value& value) { + switch (value.kind_case()) { + case Value::kNullValue: + return "null"; + case Value::kNumberValue: + return "number"; + case Value::kStringValue: + return "string"; + case Value::kBoolValue: + return "boolean"; + case Value::kStructValue: + return "object"; + case Value::kListValue: + return "array"; + case Value::KIND_NOT_SET: + return "unset"; + } + return "unknown"; +} + +const Value& RequiredField(const Struct& object, const std::string& field, + const std::string& context) { + const auto it = object.fields().find(field); + if (it == object.fields().end()) Invalid(context + ": missing required field '" + field + "'"); + return it->second; +} + +const Value* OptionalField(const Struct& object, const std::string& field) { + const auto it = object.fields().find(field); + return it == object.fields().end() ? nullptr : &it->second; +} + +void RejectUnknownFields(const Struct& object, std::initializer_list allowed, + const std::string& context) { + for (const auto& [name, value] : object.fields()) { + (void)value; + const bool known = std::any_of(allowed.begin(), allowed.end(), + [&](std::string_view candidate) { return candidate == name; }); + if (!known) Invalid(context + ": unknown field '" + name + "'"); + } +} + +const Struct& AsObject(const Value& value, const std::string& context) { + if (value.kind_case() != Value::kStructValue) { + Invalid(context + ": expected object, got " + ValueTypeName(value)); + } + return value.struct_value(); +} + +const ListValue& AsArray(const Value& value, const std::string& context) { + if (value.kind_case() != Value::kListValue) { + Invalid(context + ": expected array, got " + ValueTypeName(value)); + } + return value.list_value(); +} + +std::string AsString(const Value& value, const std::string& context) { + if (value.kind_case() != Value::kStringValue) { + Invalid(context + ": expected string, got " + ValueTypeName(value)); + } + return value.string_value(); +} + +double AsNumber(const Value& value, const std::string& context) { + if (value.kind_case() != Value::kNumberValue || !std::isfinite(value.number_value())) { + Invalid(context + ": expected finite number"); + } + return value.number_value(); +} + +template +Enum AsEnum(const Value& value, const std::string& context, + std::initializer_list> allowed) { + const std::string text = AsString(value, context); + std::string expected; + for (const auto& [name, mapped] : allowed) { + if (text == name) return mapped; + if (!expected.empty()) expected += " or "; + expected += "'" + std::string(name) + "'"; + } + Invalid(context + ": expected " + expected); +} + +uint32_t AsPositiveUint32(const Value& value, const std::string& context) { + const double number = AsNumber(value, context); + if (number <= 0.0 || std::trunc(number) != number || + number > static_cast(std::numeric_limits::max())) { + Invalid(context + ": expected positive integer no greater than uint32 max"); + } + return static_cast(number); +} + +int AsInt(const Value& value, int minimum, const std::string& context) { + const double number = AsNumber(value, context); + if (std::trunc(number) != number || number < static_cast(minimum) || + number > static_cast(std::numeric_limits::max())) { + Invalid(context + ": expected integer >= " + std::to_string(minimum)); + } + return static_cast(number); +} + +uint64_t ParseCapacity(std::string_view text, const std::string& context) { + if (text.empty()) Invalid(context + ": capacity must not be empty"); + + size_t split = 0; + uint64_t amount = 0; + while (split < text.size() && text[split] >= '0' && text[split] <= '9') { + const uint64_t digit = static_cast(text[split] - '0'); + if (amount > (std::numeric_limits::max() - digit) / 10) { + Invalid(context + ": capacity overflows uint64"); + } + amount = amount * 10 + digit; + ++split; + } + if (split == 0 || amount == 0) { + Invalid(context + ": capacity must be a positive integer followed by B, KiB, MiB, GiB, or TiB"); + } + + static constexpr std::pair kUnits[] = { + {"B", 1}, {"KiB", 1ULL << 10}, {"MiB", 1ULL << 20}, {"GiB", 1ULL << 30}, {"TiB", 1ULL << 40}}; + const std::string_view unit = text.substr(split); + uint64_t multiplier = 0; + for (const auto& [name, scale] : kUnits) { + if (unit == name) { + multiplier = scale; + break; + } + } + if (multiplier == 0) { + Invalid(context + ": capacity unit must be B, KiB, MiB, GiB, or TiB"); + } + if (amount > std::numeric_limits::max() / multiplier) { + Invalid(context + ": capacity overflows uint64"); + } + return amount * multiplier; +} + +PolicyBackendSpec ParseBackend(const std::string& name, const Value& value) { + const std::string context = "backend '" + name + "'"; + if (name.empty()) Invalid("backend names must not be empty"); + const Struct& object = AsObject(value, context); + + const std::string type = AsString(RequiredField(object, "type", context), context + ".type"); + PolicyBackendSpec backend; + backend.name = name; + backend.capacity_bytes = + ParseCapacity(AsString(RequiredField(object, "capacity", context), context + ".capacity"), + context + ".capacity"); + + if (type == "hbm") { + RejectUnknownFields(object, {"type", "capacity", "devices"}, context); + backend.tier = TierType::HBM; + const ListValue& devices = + AsArray(RequiredField(object, "devices", context), context + ".devices"); + if (devices.values().empty()) Invalid(context + ".devices: must not be empty"); + std::unordered_set seen; + for (int i = 0; i < devices.values_size(); ++i) { + const int device = + AsInt(devices.values(i), 0, context + ".devices[" + std::to_string(i) + "]"); + if (!seen.insert(device).second) { + Invalid(context + ".devices: duplicate device " + std::to_string(device)); + } + backend.devices.push_back(device); + } + } else if (type == "dram") { + RejectUnknownFields(object, {"type", "capacity", "numa_node"}, context); + backend.tier = TierType::DRAM; + if (const Value* numa = OptionalField(object, "numa_node")) { + backend.numa_node = AsInt(*numa, -1, context + ".numa_node"); + } + } else if (type == "ssd") { + RejectUnknownFields(object, {"type", "capacity", "path", "staging_slots"}, context); + backend.tier = TierType::SSD; + backend.path = AsString(RequiredField(object, "path", context), context + ".path"); + if (backend.path.empty()) Invalid(context + ".path: must not be empty"); + // Minimum 1, not 0: 0 is how a config assembled in code says "leave the + // default", which is not a thing a written-out field can mean. + if (const Value* slots = OptionalField(object, "staging_slots")) { + backend.staging_slots = AsInt(*slots, 1, context + ".staging_slots"); + } + } else { + Invalid(context + ".type: expected 'dram', 'hbm', or 'ssd'"); + } + return backend; +} + +LogicalTierConfig ParseLogicalTier(const Value& value, size_t index) { + const std::string context = "tiers[" + std::to_string(index) + "]"; + const Struct& object = AsObject(value, context); + RejectUnknownFields( + object, + {"name", "backends", "placement_policy", "offload_to", "offload_trigger", "high_watermark", + "low_watermark", "promote_trigger", "promote_hits", "promote_mode"}, + context); + + LogicalTierConfig tier; + const Value* name = OptionalField(object, "name"); + tier.name = + name == nullptr ? "tier_" + std::to_string(index) : AsString(*name, context + ".name"); + if (tier.name.empty()) Invalid(context + ".name: must not be empty"); + + // Placement within a tier is weighted by construction; the field exists so a + // policy can say so, not so it can ask for something else. + if (const Value* placement = OptionalField(object, "placement_policy")) { + if (AsString(*placement, context + ".placement_policy") != "weighted") { + Invalid(context + ".placement_policy: expected 'weighted'"); + } + } + + const Struct& members = + AsObject(RequiredField(object, "backends", context), context + ".backends"); + if (members.fields().empty()) Invalid(context + ".backends: must not be empty"); + for (const auto& [backend_name, weight] : members.fields()) { + if (backend_name.empty()) Invalid(context + ".backends: backend name must not be empty"); + tier.members.push_back( + {backend_name, AsPositiveUint32(weight, context + ".backends." + backend_name)}); + } + std::sort(tier.members.begin(), tier.members.end(), [](const auto& left, const auto& right) { + return left.backend_name < right.backend_name; + }); + + if (const Value* targets_field = OptionalField(object, "offload_to")) { + const ListValue& targets = AsArray(*targets_field, context + ".offload_to"); + for (int i = 0; i < targets.values_size(); ++i) { + tier.offload_to.push_back( + AsString(targets.values(i), context + ".offload_to[" + std::to_string(i) + "]")); + } + } + + if (const Value* trigger = OptionalField(object, "offload_trigger")) { + tier.trigger = AsEnum(*trigger, context + ".offload_trigger", + {{"on_evict", PoolOffloadTrigger::kOnEvict}, + {"watermark", PoolOffloadTrigger::kWatermark}}); + } else if (!tier.offload_to.empty()) { + Invalid(context + ": offload_trigger is required when offload_to is not empty"); + } + + if (const Value* high = OptionalField(object, "high_watermark")) { + tier.high_watermark = AsNumber(*high, context + ".high_watermark"); + } + if (const Value* low = OptionalField(object, "low_watermark")) { + tier.low_watermark = AsNumber(*low, context + ".low_watermark"); + } + if (const Value* promote = OptionalField(object, "promote_trigger")) { + tier.promote_trigger = AsEnum(*promote, context + ".promote_trigger", + {{"never", PoolPromoteTrigger::kNever}, + {"on_read", PoolPromoteTrigger::kOnRead}, + {"on_hits", PoolPromoteTrigger::kOnHits}}); + } + // A hit threshold has no meaningful default, so accepting it where nothing + // reads it would hide the misconfiguration instead of reporting it. The + // watermark fields above are lenient because they do have defaults. + if (const Value* hits = OptionalField(object, "promote_hits")) { + if (tier.promote_trigger != PoolPromoteTrigger::kOnHits) { + Invalid(context + ".promote_hits: only valid when promote_trigger is 'on_hits'"); + } + tier.promote_hits = static_cast(AsInt(*hits, 2, context + ".promote_hits")); + } else if (tier.promote_trigger == PoolPromoteTrigger::kOnHits) { + Invalid(context + ": promote_hits is required when promote_trigger is 'on_hits'"); + } + if (const Value* mode = OptionalField(object, "promote_mode")) { + tier.promote_mode = AsEnum( + *mode, context + ".promote_mode", + {{"copy", PoolTransitionMode::kCopy}, {"move", PoolTransitionMode::kMove}}); + } + return tier; +} + +// Promotion moves a key upstream and the entry tier has no upstream, so a +// trigger declared there can never fire: the read path tests the same condition +// and skips it. Checked wherever the entry tier is finally known, because +// naming it in `entry_tier` can move it after the per-tier pass. +std::string RejectEntryTierPromotion(const std::vector& tiers, size_t entry_tier, + const std::string& name) { + if (entry_tier >= tiers.size() || + tiers[entry_tier].promote_trigger == PoolPromoteTrigger::kNever) { + return {}; + } + return "logical tier '" + name + "': the entry tier has no upstream tier to promote into"; +} + +// The single definition of a well formed policy. The JSON front end and +// ApplyBackendPolicy both run it: a config assembled in code reaches the same +// bar as one that came from a file, and a rule only has to be stated once. +LogicalTierIndex ValidateBackendPolicy(const BackendPolicyConfig& policy) { + if (policy.schema_version != 1) { + Invalid("policy: only backend policy schema version 1 is supported"); + } + if (policy.backends.empty()) Invalid("policy: must define at least one backend"); + if (policy.logical_tiers.empty()) { + Invalid("policy: must define at least one logical tier"); + } + + std::unordered_set defined; + for (const auto& backend : policy.backends) { + if (backend.name.empty()) Invalid("policy.backends: names must not be empty"); + const std::string context = "backend '" + backend.name + "'"; + if (!defined.insert(backend.name).second) Invalid("duplicate " + context); + if (backend.capacity_bytes == 0) Invalid(context + ": capacity must be positive"); + if (backend.tier == TierType::HBM) { + if (backend.devices.empty()) Invalid(context + ": devices must not be empty"); + if (backend.capacity_bytes < backend.devices.size()) { + Invalid(context + ": per-device capacity would be zero"); + } + std::unordered_set devices; + for (int device : backend.devices) { + if (device < 0 || !devices.insert(device).second) { + Invalid(context + ": devices must be unique and nonnegative"); + } + } + } else if (backend.tier == TierType::DRAM) { + if (backend.numa_node < -1) Invalid(context + ": numa_node must be >= -1"); + } else if (backend.tier == TierType::SSD) { + if (backend.path.empty()) Invalid(context + ": path must not be empty"); + if (backend.staging_slots < 0) Invalid(context + ": staging_slots must be >= 0"); + } else { + Invalid(context + ": unsupported tier"); + } + } + + LogicalTierIndex index; + if (const std::string error = IndexLogicalTiers(policy.logical_tiers, &index); !error.empty()) { + Invalid(error); + } + for (size_t i = 0; i < policy.logical_tiers.size(); ++i) { + for (const auto& member : policy.logical_tiers[i].members) { + if (defined.count(member.backend_name) == 0) { + Invalid("logical tier '" + index.names[i] + "': undefined backend '" + member.backend_name + + "'"); + } + } + } + for (const auto& backend : policy.backends) { + if (index.by_backend.count(backend.name) == 0) { + Invalid("backend '" + backend.name + "' is not a member of any logical tier"); + } + } + + // A policy may name its entry tier instead of flagging it; the two must agree + // when both are present. + if (!policy.entry_tier.empty()) { + const auto named = index.by_name.find(policy.entry_tier); + if (named == index.by_name.end()) { + Invalid("policy.entry_tier: undefined logical tier '" + policy.entry_tier + "'"); + } + if (index.entry_flagged && index.entry_tier != named->second) { + Invalid("policy.entry_tier names '" + policy.entry_tier + + "' but a different logical tier is flagged as the entry tier"); + } + index.entry_tier = named->second; + if (const std::string error = RejectEntryTierPromotion(policy.logical_tiers, index.entry_tier, + index.names[index.entry_tier]); + !error.empty()) { + Invalid(error); + } + } + return index; +} + +BackendPolicyConfig ParsePolicy(const Struct& root) { + RejectUnknownFields(root, {"schema_version", "entry_tier", "backends", "tiers"}, "policy"); + + BackendPolicyConfig policy; + if (const Value* schema = OptionalField(root, "schema_version")) { + policy.schema_version = static_cast(AsInt(*schema, 1, "policy.schema_version")); + } + if (const Value* entry = OptionalField(root, "entry_tier")) { + policy.entry_tier = AsString(*entry, "policy.entry_tier"); + if (policy.entry_tier.empty()) Invalid("policy.entry_tier: must not be empty"); + } + const Struct& backends = AsObject(RequiredField(root, "backends", "policy"), "policy.backends"); + if (backends.fields().empty()) Invalid("policy.backends: must not be empty"); + policy.backends.reserve(backends.fields_size()); + for (const auto& [name, value] : backends.fields()) { + policy.backends.push_back(ParseBackend(name, value)); + } + + const ListValue& tiers = AsArray(RequiredField(root, "tiers", "policy"), "policy.tiers"); + if (tiers.values().empty()) Invalid("policy.tiers: must not be empty"); + policy.logical_tiers.reserve(tiers.values_size()); + for (int i = 0; i < tiers.values_size(); ++i) { + policy.logical_tiers.push_back(ParseLogicalTier(tiers.values(i), static_cast(i))); + } + + policy.logical_tiers[ValidateBackendPolicy(policy).entry_tier].entry = true; + return policy; +} + +BackendPolicyLoadResult ErrorResult(std::string error) { + BackendPolicyLoadResult result; + result.error = std::move(error); + return result; +} + +bool SetApplyError(std::string* error, std::string message) { + if (error != nullptr) *error = std::move(message); + return false; +} + +} // namespace + +std::string IndexLogicalTiers(const std::vector& tiers, + LogicalTierIndex* index) { + if (tiers.empty()) return "logical tier graph must not be empty"; + + *index = {}; + index->names.reserve(tiers.size()); + index->offload_to.resize(tiers.size()); + + for (size_t i = 0; i < tiers.size(); ++i) { + const auto& tier = tiers[i]; + std::string name = tier.name.empty() ? "tier_" + std::to_string(i) : tier.name; + const std::string context = "logical tier '" + name + "'"; + // Slot tokens join a tier name with a backend name, so '|' cannot appear in + // either half. + if (name.find('|') != std::string::npos) return context + ": name must not contain '|'"; + if (!index->by_name.emplace(name, i).second) { + return "duplicate logical tier name '" + name + "'"; + } + index->names.push_back(std::move(name)); + + if (tier.entry) { + if (index->entry_flagged) return "exactly one logical tier may be the entry tier"; + index->entry_flagged = true; + index->entry_tier = i; + } + if (!std::isfinite(tier.low_watermark) || !std::isfinite(tier.high_watermark) || + !(tier.low_watermark > 0.0 && tier.low_watermark < tier.high_watermark && + tier.high_watermark <= 1.0)) { + return context + ": watermarks must satisfy 0 < low_watermark < high_watermark <= 1"; + } + if (tier.trigger != PoolOffloadTrigger::kOnEvict && + tier.trigger != PoolOffloadTrigger::kWatermark) { + return context + ": unsupported offload trigger"; + } + switch (tier.promote_trigger) { + case PoolPromoteTrigger::kNever: + case PoolPromoteTrigger::kOnRead: + if (tier.promote_hits != 0) { + return context + ": promote_hits is only meaningful for the 'on_hits' trigger"; + } + break; + case PoolPromoteTrigger::kOnHits: + // 1 would be the on_read trigger under another name, and one policy + // meaning two things is how the two drift apart. + if (tier.promote_hits < 2) { + return context + ": promote_hits must be at least 2"; + } + break; + default: + return context + ": unsupported promote trigger"; + } + if (tier.members.empty()) return context + ": must have at least one backend"; + for (const auto& member : tier.members) { + if (member.backend_name.empty()) return context + ": backend names must not be empty"; + if (member.weight == 0) { + return context + ": weight for '" + member.backend_name + "' must be positive"; + } + if (!index->by_backend.emplace(member.backend_name, i).second) { + return "backend '" + member.backend_name + "' is a member of more than one tier"; + } + } + } + + for (size_t i = 0; i < tiers.size(); ++i) { + const std::string context = "logical tier '" + index->names[i] + "'"; + std::unordered_set seen; + for (const auto& target : tiers[i].offload_to) { + if (target.empty()) return context + ": offload target must not be empty"; + const auto by_backend = index->by_backend.find(target); + const auto by_name = index->by_name.find(target); + if (by_backend == index->by_backend.end() && by_name == index->by_name.end()) { + return context + ": undefined offload target '" + target + "'"; + } + const size_t resolved = + by_backend != index->by_backend.end() ? by_backend->second : by_name->second; + // Offload only ever runs downstream; an edge back up would let a key + // cycle between two tiers forever. + if (resolved <= i) { + return context + ": offload target '" + target + "' must belong to a strictly later tier"; + } + if (seen.insert(resolved).second) index->offload_to[i].push_back(resolved); + } + } + return RejectEntryTierPromotion(tiers, index->entry_tier, index->names[index->entry_tier]); +} + +BackendPolicyLoadResult LoadBackendPolicyJson(std::string_view json) { + try { + Struct root; + google::protobuf::util::JsonParseOptions options; + options.ignore_unknown_fields = false; + const auto status = + google::protobuf::util::JsonStringToMessage(std::string(json), &root, options); + if (!status.ok()) return ErrorResult("invalid backend policy JSON: " + status.ToString()); + + BackendPolicyLoadResult result; + result.config = ParsePolicy(root); + return result; + } catch (const std::exception& exception) { + return ErrorResult(exception.what()); + } catch (...) { + return ErrorResult("unknown error while loading backend policy JSON"); + } +} + +BackendPolicyLoadResult LoadBackendPolicyFile(std::string_view path) { + try { + const std::string file_path(path); + std::ifstream input(file_path, std::ios::in | std::ios::binary); + if (!input) return ErrorResult("failed to open backend policy file '" + file_path + "'"); + std::ostringstream contents; + contents << input.rdbuf(); + if (input.bad()) return ErrorResult("failed to read backend policy file '" + file_path + "'"); + return LoadBackendPolicyJson(contents.str()); + } catch (const std::exception& exception) { + return ErrorResult(exception.what()); + } catch (...) { + return ErrorResult("unknown error while reading backend policy file"); + } +} + +bool ApplyBackendPolicy(const BackendPolicyConfig& policy, PoolClientConfig* config, + std::string* error, std::string storage_path_suffix) { + if (config == nullptr) return SetApplyError(error, "PoolClientConfig output must not be null"); + + try { + const LogicalTierIndex index = ValidateBackendPolicy(policy); + + std::vector concrete_backends; + std::unordered_map> expansions; + + for (const auto& backend : policy.backends) { + std::vector concrete_names; + if (backend.tier == TierType::HBM) { + const uint64_t base = backend.capacity_bytes / backend.devices.size(); + const uint64_t remainder = backend.capacity_bytes % backend.devices.size(); + for (size_t i = 0; i < backend.devices.size(); ++i) { + const int device = backend.devices[i]; + BackendInstanceConfig instance; + instance.name = backend.devices.size() == 1 ? backend.name + : backend.name + "@" + std::to_string(device); + instance.tier = TierType::HBM; + instance.hbm.device = device; + instance.hbm.buffer_sizes = {base + (i < remainder ? 1 : 0)}; + concrete_names.push_back(instance.name); + concrete_backends.push_back(std::move(instance)); + } + } else { + BackendInstanceConfig instance; + instance.name = backend.name; + instance.tier = backend.tier; + if (backend.tier == TierType::DRAM) { + instance.dram.buffer_sizes = {backend.capacity_bytes}; + instance.dram.numa_node = backend.numa_node; + } else { + if (backend.capacity_bytes > std::numeric_limits::max()) { + Invalid("backend '" + backend.name + "': capacity does not fit size_t"); + } + instance.ssd.enabled = true; + instance.ssd.ssd.enabled = true; + instance.ssd.ssd.capacity_bytes = static_cast(backend.capacity_bytes); + instance.ssd.ssd.storage_dir = backend.path + storage_path_suffix; + instance.ssd.ssd.ssd_backend = "file"; + if (backend.staging_slots > 0) { + instance.ssd_staging_buffer_slots = backend.staging_slots; + } + } + concrete_names.push_back(instance.name); + concrete_backends.push_back(std::move(instance)); + } + expansions.emplace(backend.name, std::move(concrete_names)); + } + + std::unordered_set concrete_names; + for (const auto& backend : concrete_backends) { + if (!concrete_names.insert(backend.name).second) { + Invalid("backend expansion produces duplicate concrete name '" + backend.name + "'"); + } + } + + std::vector concrete_tiers; + concrete_tiers.reserve(policy.logical_tiers.size()); + for (size_t tier_index = 0; tier_index < policy.logical_tiers.size(); ++tier_index) { + const auto& source = policy.logical_tiers[tier_index]; + + // Members are weighted per named backend, but placement happens on the + // concrete instances a name expands to. Scaling every weight by the LCM + // of the expansion sizes keeps the configured ratio exact in integers. + uint64_t expansion_lcm = 1; + for (const auto& member : source.members) { + const uint64_t count = expansions.at(member.backend_name).size(); + const uint64_t divisor = std::gcd(expansion_lcm, count); + if (expansion_lcm > std::numeric_limits::max() / (count / divisor)) { + Invalid("logical tier " + std::to_string(tier_index) + ": expansion LCM overflows"); + } + expansion_lcm *= count / divisor; + } + + LogicalTierConfig lowered; + lowered.trigger = source.trigger; + lowered.high_watermark = source.high_watermark; + lowered.low_watermark = source.low_watermark; + lowered.name = index.names[tier_index]; + lowered.entry = tier_index == index.entry_tier; + lowered.promote_trigger = source.promote_trigger; + lowered.promote_hits = source.promote_hits; + lowered.promote_mode = source.promote_mode; + for (const auto& member : source.members) { + const auto& names = expansions.at(member.backend_name); + const uint64_t factor = expansion_lcm / names.size(); + if (factor > std::numeric_limits::max() / member.weight) { + Invalid("logical tier " + std::to_string(tier_index) + + ": expanded backend weight overflows uint32"); + } + const uint32_t concrete_weight = static_cast(factor * member.weight); + for (const auto& name : names) { + lowered.members.push_back({name, concrete_weight}); + } + } + for (const auto& target : source.offload_to) { + const auto expansion = expansions.find(target); + if (expansion != expansions.end()) { + lowered.offload_to.insert(lowered.offload_to.end(), expansion->second.begin(), + expansion->second.end()); + continue; + } + const size_t target_tier = index.by_name.at(target); + for (const auto& member : policy.logical_tiers[target_tier].members) { + const auto& concrete = expansions.at(member.backend_name); + lowered.offload_to.insert(lowered.offload_to.end(), concrete.begin(), concrete.end()); + } + } + concrete_tiers.push_back(std::move(lowered)); + } + + if (concrete_backends.size() > kMaxBackendsPerPeer) { + Invalid("backend policy expands beyond the per-peer backend limit of " + + std::to_string(kMaxBackendsPerPeer)); + } + const uint64_t page_size = + config->dram_page_size == 0 ? 2ULL * 1024 * 1024 : config->dram_page_size; + for (const auto& backend : concrete_backends) { + uint64_t capacity = 0; + if (backend.tier == TierType::DRAM && !backend.dram.buffer_sizes.empty()) { + capacity = backend.dram.buffer_sizes.front(); + } else if (backend.tier == TierType::HBM && !backend.hbm.buffer_sizes.empty()) { + capacity = backend.hbm.buffer_sizes.front(); + } else if (backend.tier == TierType::SSD) { + capacity = backend.ssd.ssd.capacity_bytes; + } + if (capacity < page_size) { + Invalid("backend '" + backend.name + "' capacity is smaller than page size"); + } + } + + config->backends.swap(concrete_backends); + config->logical_tiers.swap(concrete_tiers); + config->placement_policy = PoolPlacementPolicy::TIERED; + if (error != nullptr) error->clear(); + return true; + } catch (const std::exception& exception) { + return SetApplyError(error, exception.what()); + } catch (...) { + return SetApplyError(error, "unknown error while applying backend policy"); + } +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/pool/tier_graph.cpp b/src/umbp/distributed/pool/tier_graph.cpp new file mode 100644 index 000000000..031ef96e3 --- /dev/null +++ b/src/umbp/distributed/pool/tier_graph.cpp @@ -0,0 +1,295 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License + +#include "umbp/distributed/pool/tier_graph.h" + +#include +#include +#include +#include + +namespace mori::umbp { +namespace { + +LogicalTierGraph::CompileResult CompileError(std::string error) { + return {nullptr, std::move(error)}; +} + +uint64_t SaturatingAdd(uint64_t lhs, uint64_t rhs) { + return rhs > std::numeric_limits::max() - lhs ? std::numeric_limits::max() + : lhs + rhs; +} + +} // namespace + +uint64_t StablePlacementHash(std::string_view key) { + uint64_t hash = 14695981039346656037ULL; + for (unsigned char byte : key) { + hash ^= byte; + hash *= 1099511628211ULL; + } + return hash; +} + +std::vector WeightedBackendOrder(const std::vector& members, + const BackendRegistry& backends, std::string_view key) { + struct Candidate { + uint32_t backend_id; + uint32_t weight; + }; + + std::vector candidates; + candidates.reserve(members.size()); + uint64_t total_weight = 0; + for (const auto& member : members) { + const auto* entry = backends.GetEntry(member.backend_name); + if (entry == nullptr || member.weight == 0) continue; + candidates.push_back({entry->backend_id, member.weight}); + total_weight += member.weight; + } + if (candidates.empty()) return {}; + + uint64_t bucket = StablePlacementHash(key) % total_weight; + size_t primary = 0; + for (size_t i = 0; i < candidates.size(); ++i) { + if (bucket < candidates[i].weight) { + primary = i; + break; + } + bucket -= candidates[i].weight; + } + + std::vector order; + order.reserve(candidates.size()); + for (size_t offset = 0; offset < candidates.size(); ++offset) { + order.push_back(candidates[(primary + offset) % candidates.size()].backend_id); + } + return order; +} + +LogicalTierGraph::LogicalTierGraph(const BackendRegistry& backends, std::vector nodes, + std::vector> upstream, + std::unordered_map tier_by_backend_id, + TierIndex entry_tier) + : backends_(&backends), + nodes_(std::move(nodes)), + upstream_(std::move(upstream)), + tier_by_backend_id_(std::move(tier_by_backend_id)), + entry_tier_(entry_tier) {} + +LogicalTierGraph::CompileResult LogicalTierGraph::Compile( + const std::vector& tiers, const BackendRegistry& backends) { + LogicalTierIndex index; + if (std::string error = IndexLogicalTiers(tiers, &index); !error.empty()) { + return CompileError(std::move(error)); + } + + // Everything the tier list can say about itself has been checked; what is + // left is binding member names to the backends this peer actually registered. + std::vector nodes; + nodes.reserve(tiers.size()); + std::unordered_map tier_by_backend_id; + std::vector> upstream(tiers.size()); + + for (TierIndex i = 0; i < tiers.size(); ++i) { + Node node; + node.name = index.names[i]; + node.members = tiers[i].members; + node.offload_to = index.offload_to[i]; + node.trigger = tiers[i].trigger; + node.high_watermark = tiers[i].high_watermark; + node.low_watermark = tiers[i].low_watermark; + node.promote_trigger = tiers[i].promote_trigger; + node.promote_hits = tiers[i].promote_hits; + node.promote_mode = tiers[i].promote_mode; + for (const auto& member : node.members) { + const auto* entry = backends.GetEntry(member.backend_name); + if (entry == nullptr) { + return CompileError("logical tier '" + node.name + "' references unknown backend '" + + member.backend_name + "'"); + } + tier_by_backend_id.emplace(entry->backend_id, i); + } + for (TierIndex target : node.offload_to) upstream[target].push_back(i); + nodes.push_back(std::move(node)); + } + + auto* graph = new LogicalTierGraph(backends, std::move(nodes), std::move(upstream), + std::move(tier_by_backend_id), index.entry_tier); + return {std::shared_ptr(graph), {}}; +} + +std::optional LogicalTierGraph::TierIndexForBackendId( + uint32_t backend_id) const { + const auto it = tier_by_backend_id_.find(backend_id); + return it == tier_by_backend_id_.end() ? std::nullopt : std::optional{it->second}; +} + +std::string LogicalTierGraph::NameForBackend(uint32_t backend_id) const { + auto tier = TierIndexForBackendId(backend_id); + return tier.has_value() ? NodeAt(*tier).name : std::string{}; +} + +std::map LogicalTierGraph::CapacitySnapshot() const { + std::map snapshot; + for (TierIndex tier = 0; tier < nodes_.size(); ++tier) { + LogicalTierCapacity logical; + logical.put_eligible = tier == entry_tier_; + for (const auto& member : NodeAt(tier).members) { + const auto* backend = backends_->Get(member.backend_name); + if (backend == nullptr) continue; + const auto cap = backend->Capacity(); + if (logical.representative_tier == TierType::UNKNOWN) { + logical.representative_tier = backend->Tier(); + } + logical.capacity.total_bytes = SaturatingAdd(logical.capacity.total_bytes, cap.total_bytes); + logical.capacity.available_bytes = + SaturatingAdd(logical.capacity.available_bytes, cap.available_bytes); + const uint64_t allocatable = cap.max_allocatable_bytes == 0 + ? cap.available_bytes + : std::min(cap.available_bytes, cap.max_allocatable_bytes); + logical.capacity.max_allocatable_bytes = + std::max(logical.capacity.max_allocatable_bytes, allocatable); + } + logical.peak_member_utilization = PeakMemberUtilization(tier); + snapshot.emplace(NodeAt(tier).name, logical); + } + TierCapacity entry_pool; + for (TierIndex tier : ReachableTierOrder(entry_tier_, false, true)) { + const auto& capacity = snapshot.at(NodeAt(tier).name).capacity; + entry_pool.total_bytes = SaturatingAdd(entry_pool.total_bytes, capacity.total_bytes); + entry_pool.available_bytes = + SaturatingAdd(entry_pool.available_bytes, capacity.available_bytes); + entry_pool.max_allocatable_bytes = + std::max(entry_pool.max_allocatable_bytes, capacity.max_allocatable_bytes); + } + snapshot.at(NodeAt(entry_tier_).name).capacity = entry_pool; + return snapshot; +} + +namespace { + +double UtilizationOf(const TierCapacity& capacity) { + if (capacity.total_bytes == 0) return 0.0; + const uint64_t used = + capacity.total_bytes - std::min(capacity.total_bytes, capacity.available_bytes); + return static_cast(used) / static_cast(capacity.total_bytes); +} + +} // namespace + +double LogicalTierGraph::MemberUtilization(uint32_t backend_id) const { + const auto* backend = backends_->Get(backend_id); + return backend == nullptr ? 0.0 : UtilizationOf(backend->Capacity()); +} + +double LogicalTierGraph::PeakMemberUtilization(TierIndex tier) const { + double peak = 0.0; + for (const auto& member : NodeAt(tier).members) { + const auto* backend = backends_->Get(member.backend_name); + if (backend == nullptr) continue; + peak = std::max(peak, UtilizationOf(backend->Capacity())); + } + return peak; +} + +bool LogicalTierGraph::AtOrAbove(TierIndex tier, double watermark) const { + return PeakMemberUtilization(tier) >= watermark; +} + +bool LogicalTierGraph::MemberAtOrAbove(uint32_t backend_id, double watermark) const { + return MemberUtilization(backend_id) >= watermark; +} + +std::vector LogicalTierGraph::WeightedMemberOrder(TierIndex tier, + std::string_view key) const { + return WeightedBackendOrder(NodeAt(tier).members, *backends_, key); +} + +std::vector LogicalTierGraph::ReachableTierOrder( + TierIndex source, bool upstream, bool include_source) const { + if (source >= nodes_.size()) return {}; + std::vector order; + std::vector visited(nodes_.size(), false); + std::function visit = [&](TierIndex tier) { + if (visited[tier]) return; + visited[tier] = true; + order.push_back(tier); + const auto& edges = upstream ? upstream_[tier] : nodes_[tier].offload_to; + for (TierIndex target : edges) visit(target); + }; + visit(source); + if (!include_source) order.erase(order.begin()); + if (upstream) std::sort(order.begin(), order.end()); + return order; +} + +std::vector LogicalTierGraph::BackendOrder(const std::vector& tiers, + std::string_view key) const { + std::vector order; + for (TierIndex tier : tiers) { + auto members = WeightedMemberOrder(tier, key); + order.insert(order.end(), members.begin(), members.end()); + } + return order; +} + +std::vector LogicalTierGraph::PutOrder(std::string_view key) const { + return BackendOrder(ReachableTierOrder(entry_tier_, false, true), key); +} + +std::vector LogicalTierGraph::PutOrderFromTier(std::string_view name, + std::string_view key) const { + for (TierIndex tier = 0; tier < nodes_.size(); ++tier) { + if (NodeAt(tier).name == name) { + return BackendOrder(ReachableTierOrder(tier, false, true), key); + } + } + return {}; +} + +std::vector LogicalTierGraph::TransitionTargetOrder(TierIndex source, + std::string_view key) const { + return BackendOrder(ReachableTierOrder(source, false, false), key); +} + +std::vector LogicalTierGraph::PromoteTargetOrder(TierIndex source, + std::string_view key) const { + return BackendOrder(ReachableTierOrder(source, true, false), key); +} + +std::vector LogicalTierGraph::ReadOrder() const { + std::vector order; + order.reserve(tier_by_backend_id_.size()); + for (const auto& node : nodes_) { + for (const auto& member : node.members) { + const auto* entry = backends_->GetEntry(member.backend_name); + if (entry != nullptr) order.push_back(entry->backend_id); + } + } + return order; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/pool/tier_transition.cpp b/src/umbp/distributed/pool/tier_transition.cpp new file mode 100644 index 000000000..9350de8cf --- /dev/null +++ b/src/umbp/distributed/pool/tier_transition.cpp @@ -0,0 +1,138 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#include "umbp/distributed/pool/tier_transition.h" + +#include +#include + +namespace mori::umbp { + +bool TierTransitionExecutor::BuildTransfers(MediumBackend* source, const ResolvedEntry& resolved, + MediumBackend* target, const AllocateResult& allocation, + std::vector* items) { + if (source == nullptr || target == nullptr || items == nullptr || resolved.size == 0 || + resolved.page_size == 0 || allocation.page_size == 0 || resolved.pages.empty() || + allocation.pages.empty() || allocation.size != resolved.size) { + return false; + } + + items->clear(); + for (uint64_t copied = 0; copied < resolved.size;) { + const uint64_t source_page_number = copied / resolved.page_size; + const uint64_t target_page_number = copied / allocation.page_size; + if (source_page_number >= resolved.pages.size() || + target_page_number >= allocation.pages.size()) { + return false; + } + + const auto& source_page = resolved.pages[static_cast(source_page_number)]; + const auto& target_page = allocation.pages[static_cast(target_page_number)]; + TransferRef source_ref = source->BufferRef(source_page.buffer_index); + TransferRef target_ref = target->BufferRef(target_page.buffer_index); + if (!source_ref.Valid() || !target_ref.Valid()) return false; + + const uint64_t source_offset = copied % resolved.page_size; + const uint64_t target_offset = copied % allocation.page_size; + const uint64_t size = std::min({resolved.size - copied, resolved.page_size - source_offset, + allocation.page_size - target_offset}); + TransferItem item; + item.tag = items->size(); + item.src = std::move(source_ref); + item.src_offset = + static_cast(source_page.page_index) * resolved.page_size + source_offset; + item.dst = std::move(target_ref); + item.dst_offset = + static_cast(target_page.page_index) * allocation.page_size + target_offset; + item.size = size; + items->push_back(std::move(item)); + copied += size; + } + return true; +} + +TierTransitionResult TierTransitionExecutor::Execute( + const std::string& key, uint32_t source_backend_id, + const std::vector& target_backend_ids, bool remove_source) const { + TierTransitionResult result; + if (backends_ == nullptr || transfer_engine_ == nullptr || target_backend_ids.empty()) { + return result; + } + + MediumBackend* source = backends_->Get(source_backend_id); + ResolvedEntry resolved; + if (source == nullptr || !source->AcquireMigrationRead(key, &resolved) || !resolved.found) { + return result; + } + + for (uint32_t target_backend_id : target_backend_ids) { + MediumBackend* target = backends_->Get(target_backend_id); + if (target == nullptr || target == source) continue; + + auto allocations = target->BatchAllocate({AllocateRequest{key, resolved.size}}); + if (allocations.empty()) continue; + AllocateResult& allocation = allocations.front(); + + if (allocation.outcome == AllocateOutcome::kSuccessAllocated) { + std::vector transfers; + if (!BuildTransfers(source, resolved, target, allocation, &transfers) || + !transfer_engine_->Transfer(transfers, nullptr)) { + target->BatchAbort({allocation.slot_id}); + continue; + } + auto committed = target->BatchCommit({CommitRequest{allocation.slot_id, key}}); + if (committed.empty() || !committed.front().success) { + target->BatchAbort({allocation.slot_id}); + continue; + } + } else if (allocation.outcome != AllocateOutcome::kSuccessAlreadyExists) { + continue; + } + + ResolvedEntry target_read; + if (!target->AcquireMigrationRead(key, &target_read) || target_read.size != resolved.size) { + if (target_read.found) target->ReleaseMigrationRead(key); + if (allocation.outcome == AllocateOutcome::kSuccessAllocated) target->Evict({key}); + continue; + } + + result.success = true; + result.target_backend_id = target_backend_id; + result.bytes_moved = resolved.size; + source->ReleaseMigrationRead(key); + if (remove_source) { + auto evicted = source->Evict({key}); + const uint64_t freed = evicted.empty() ? 0 : evicted.front().bytes_freed; + result.source_draining = freed == 0 && source->Contains(key); + result.bytes_freed = freed; + } + target->ReleaseMigrationRead(key); + return result; + } + + source->ReleaseMigrationRead(key); + return result; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/pool_client.cpp b/src/umbp/distributed/pool_client.cpp index f960063f2..c0ce59e65 100644 --- a/src/umbp/distributed/pool_client.cpp +++ b/src/umbp/distributed/pool_client.cpp @@ -24,65 +24,54 @@ #include #include +#include +#include #include #include #include #include #include +#include #include +#include #include #include #include #include #include #include -#if defined(__x86_64__) || defined(__i386__) -#include -#endif #include +#include -#include "mori/io/backend.hpp" #include "mori/utils/mori_log.hpp" +#include "umbp/common/device_copy.h" #include "umbp/common/env_time.h" +#include "umbp/common/grpc_limits.h" +#include "umbp/common/parallel_for.h" +#include "umbp/common/range_utils.h" +#include "umbp/distributed/benchmark/workload_trace_recorder.h" #include "umbp/distributed/master/master_metrics.h" +#include "umbp/distributed/peer/backend/hbm_backend.h" +#include "umbp/distributed/peer/backend/instrumented_backend.h" +#include "umbp/distributed/peer/backend/page_backend.h" +#include "umbp/distributed/peer/backend/ssd_backend.h" #include "umbp/distributed/peer/batch_resolve_codec.h" -#include "umbp/distributed/peer/peer_dram_allocator.h" #include "umbp/distributed/peer/peer_service.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" -#include "umbp/distributed/peer/ssd_copy_pipeline.h" -#include "umbp/distributed/ssd_read_lease.h" +#include "umbp/distributed/pool/peer_pool.h" +#include "umbp/distributed/range_map.h" +#include "umbp/distributed/transfer/composite_transfer_engine.h" +#include "umbp/distributed/transfer/hbm_copy_engine.h" +#include "umbp/distributed/transfer/local_copy_engine.h" +#include "umbp/distributed/transfer/mori_io_engine.h" +#ifdef UMBP_ENABLE_GDS +#include "umbp/distributed/transfer/gds_engine.h" +#endif #include "umbp_peer.grpc.pb.h" namespace mori::umbp { namespace { -bool IsValidMemoryDesc(const mori::io::MemoryDesc& desc) { return desc.size > 0; } - -// Key for grouping page transfers by their (local MR, remote MR) pair. -struct PairKey { - mori::io::MemoryUniqueId local; - mori::io::MemoryUniqueId remote; - bool operator==(const PairKey& o) const noexcept { - return local == o.local && remote == o.remote; - } -}; -struct PairKeyHash { - size_t operator()(const PairKey& k) const noexcept { - // hash_combine (boost-style): independent of size_t width. - size_t h = std::hash{}(k.local); - h ^= std::hash{}(k.remote) + 0x9e3779b97f4a7c15ULL + (h << 6) + - (h >> 2); - return h; - } -}; - -// True iff [next, ...) is exactly adjacent after [base, base+len), with no -// size_t overflow in base+len. Used to coalesce contiguous SG segments. -inline bool AdjacentNoOverflow(size_t base, size_t len, size_t next) { - return len <= std::numeric_limits::max() - base && base + len == next; -} - // --------------------------------------------------------------------------- // Bandwidth metrics // --------------------------------------------------------------------------- @@ -127,71 +116,386 @@ BatchBandwidthSplit ComputeBatchBandwidthBytes(const std::vector& result return acc; } -void ObserveBatchBandwidth(MasterClient& master_client, double bytes, double seconds, +// `master_client` is null on a node running without a master; the histogram +// simply has nowhere to go, and the caller should not have to know that. +void ObserveBatchBandwidth(MasterClient* master_client, double bytes, double seconds, const char* metric_name, const char* metric_help, std::string_view traffic) { - if (bytes <= 0.0 || seconds <= 0.0) return; + if (master_client == nullptr || bytes <= 0.0 || seconds <= 0.0) return; const double gibps = (bytes / seconds) / kGiB; if (gibps <= 0.0) return; MasterClient::Labels labels = {{"traffic", std::string(traffic)}}; - master_client.Observe(metric_name, metric_help, std::move(labels), BatchBandwidthBucketsGiBps(), - gibps); + master_client->Observe(metric_name, metric_help, std::move(labels), BatchBandwidthBucketsGiBps(), + gibps); } // --------------------------------------------------------------------------- -// Page / size math +// Ranged-call sub-timers (UMBP_RANGED_CALL_DEBUG=1) +// +// ScopedBatchBandwidth already times the WHOLE ranged call, and HbmCopyEngine's +// debug mode times the copies inside it. Neither answers the question those +// two together raise: where does the rest of the call go? Subtracting one +// aggregate from the other only bounds it, because a batch fans keys across +// executor threads, so summed copy time is not the call's copy wall-clock. +// +// These timers close that gap by splitting a ranged call into the phases that +// can actually stall it: +// +// resolve per-key BatchResolve, scanning every backend in the registry +// route BatchRouteGet / BatchRoutePut RPC to the master +// xfer transfer_engine_->Transfer (the part HbmCopyEngine reports) +// lock waiting on the get/put scratch mutex, which serializes arena users +// remote the arena loop: fetch/assemble + install +// +// Note the paths are NOT symmetric, and the timers are what proves it: a get +// resolves locally first and only routes on a miss, while a put issues +// BatchRoutePut on EVERY call including fully-local ones. +// +// Off by default; knobs mirror UMBP_HBM_COPY_DEBUG: +// UMBP_RANGED_CALL_DEBUG=1 enable +// UMBP_RANGED_CALL_DEBUG_SAMPLE=N per-call line every Nth call (1) +// UMBP_RANGED_CALL_DEBUG_SUMMARY_SEC=N cumulative rollup cadence (10; 0=off) // --------------------------------------------------------------------------- -// Bytes belonging to the i-th logical page of a Put/Get spread across -// `num_pages` pages of `page_size` bytes. Last page may be partial. -// --- Parallel + AVX2 NT block copy for self-target (local) pages. ---------- -// In distributed mode 1 key == 1 page (master page_size == KV block size), so -// the distributed self-target path copies one ~4 MiB block per call. The local -// tier's (DRAMTier) multi-thread/AVX2 optimization never applied here because -// it parallelizes *within* a call's pages (always 1). We instead parallelize -// across the many keys of one BatchPut/BatchGet (different threads -> different -// keys); per-key NT-AVX2 copy gives the cache-bypass win. Threads via -// UMBP_DRAM_{READ,WRITE}_THREADS (same envs as DRAMTier). -#if defined(__x86_64__) || defined(__i386__) -__attribute__((target("avx2"))) inline void LocalNtCopyAvx2(char* d, const char* s, size_t n) { - size_t head = (32 - (reinterpret_cast(d) & 31)) & 31; - if (head > n) head = n; - std::memcpy(d, s, head); - size_t i = head; - for (; i + 128 <= n; i += 128) { - __m256i a = _mm256_loadu_si256(reinterpret_cast(s + i)); - __m256i b = _mm256_loadu_si256(reinterpret_cast(s + i + 32)); - __m256i c = _mm256_loadu_si256(reinterpret_cast(s + i + 64)); - __m256i e = _mm256_loadu_si256(reinterpret_cast(s + i + 96)); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i), a); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 32), b); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 64), c); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 96), e); - } - for (; i + 32 <= n; i += 32) { - __m256i a = _mm256_loadu_si256(reinterpret_cast(s + i)); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i), a); - } - if (i < n) std::memcpy(d + i, s + i, n - i); - _mm_sfence(); -} -inline bool LocalAvx2Supported() { return __builtin_cpu_supports("avx2"); } -#else -inline void LocalNtCopyAvx2(char* d, const char* s, size_t n) { std::memcpy(d, s, n); } -inline bool LocalAvx2Supported() { return false; } -#endif +bool RangedDebugEnvFlag(const char* name, bool fallback) { + const char* value = std::getenv(name); + if (value == nullptr) return fallback; + const std::string text(value); + return text != "0" && text != "off" && text != "false"; +} -inline void LocalCopyBlock(void* dst, const void* src, size_t size) { - static const bool kNt = LocalAvx2Supported() && !(std::getenv("UMBP_DRAM_NT_COPY") && - std::getenv("UMBP_DRAM_NT_COPY")[0] == '0'); - static const size_t kNtMinBytes = 256ull << 10; - if (kNt && size >= kNtMinBytes) { - LocalNtCopyAvx2(static_cast(dst), static_cast(src), size); - } else { - std::memcpy(dst, src, size); +size_t RangedDebugEnvSize(const char* name, size_t fallback) { + const char* value = std::getenv(name); + if (value == nullptr) return fallback; + char* end = nullptr; + const unsigned long long parsed = std::strtoull(value, &end, 10); + if (end == value) return fallback; + return static_cast(parsed); +} + +bool RangedDebugEnabled() { + static const bool enabled = RangedDebugEnvFlag("UMBP_RANGED_CALL_DEBUG", false); + return enabled; +} + +size_t RangedDebugSampleEvery() { + static const size_t n = + std::max(RangedDebugEnvSize("UMBP_RANGED_CALL_DEBUG_SAMPLE", 1), 1); + return n; +} + +size_t RangedDebugSummarySeconds() { + static const size_t n = RangedDebugEnvSize("UMBP_RANGED_CALL_DEBUG_SUMMARY_SEC", 10); + return n; +} + +bool RangedDebugSampleThisCall() { + static std::atomic counter{0}; + return (counter.fetch_add(1, std::memory_order_relaxed) % RangedDebugSampleEvery()) == 0; +} + +inline double RangedSeconds(std::chrono::steady_clock::time_point start, + std::chrono::steady_clock::time_point end) { + return std::chrono::duration_cast>(end - start).count(); +} + +struct RangedPhases { + double resolve = 0.0; + // Split out of `resolve`, which used to cover the whole per-key loop and so + // could not say whether the time went to the medium's index or to describing + // the caller's buffers. `classify` is the latter: one pointer lookup per + // RANGE, and a ranged call carries thousands of them. + double classify = 0.0; + double build = 0.0; + double validate = 0.0; // put only: checking the ranges tile their object + double commit = 0.0; // put only: slot commit + double route = 0.0; + double xfer = 0.0; + // xfer split three ways. Planning scales with items, the move with bytes, so + // a batch of many small ranges needs them apart to say which it is paying. + double xfer_plan = 0.0; + double xfer_submit = 0.0; + double xfer_wait = 0.0; + double lock = 0.0; + double remote = 0.0; + size_t keys = 0; + // TransferItems handed to the engine. Grows with ranges, not keys, and is + // what makes the item-per-range cost visible next to the phase times. + size_t items = 0; + size_t local_keys = 0; + size_t remote_keys = 0; + double bytes = 0.0; + // Component identity. A ranged call is always single-pool -- the tree + // connector groups by PoolName and chunks inside a per-pool loop, so one + // sample key names the component for the whole batch. Verified empirically + // too: 60/60 gather buckets carried a single object size. + // + // `object_size` is the grouping key rather than the string: the six DSv4-Pro + // components have six distinct object sizes, so it is an exact label that + // costs no parsing on the hot path. `key0` is carried alongside so the + // size->component mapping is provable from the log rather than inferred. + const std::string* key0 = nullptr; + size_t object_size = 0; +}; + +// Adds its lifetime to a sink. A null sink makes it inert, which is how the +// hot path pays nothing when debug is off. +class PhaseTimer { + public: + explicit PhaseTimer(double* sink) + : sink_(sink), + start_(sink ? std::chrono::steady_clock::now() : std::chrono::steady_clock::time_point{}) {} + ~PhaseTimer() { Stop(); } + PhaseTimer(const PhaseTimer&) = delete; + PhaseTimer& operator=(const PhaseTimer&) = delete; + void Stop() { + if (sink_ == nullptr) return; + *sink_ += RangedSeconds(start_, std::chrono::steady_clock::now()); + sink_ = nullptr; + } + + private: + double* sink_; + std::chrono::steady_clock::time_point start_; +}; + +class RangedStats { + public: + void Record(const char* op, const RangedPhases& p, double total) { + const bool is_get = (op[0] == 'g'); + std::lock_guard lock(mutex_); + Totals& t = is_get ? get_ : put_; + t.calls += 1; + t.total += total; + t.resolve += p.resolve; + t.classify += p.classify; + t.build += p.build; + t.validate += p.validate; + t.commit += p.commit; + t.route += p.route; + t.xfer += p.xfer; + t.xfer_plan += p.xfer_plan; + t.xfer_submit += p.xfer_submit; + t.xfer_wait += p.xfer_wait; + t.lock += p.lock; + t.remote += p.remote; + t.bytes += p.bytes; + t.items += p.items; + + // Per-component totals, keyed by object size (one call = one pool). The + // sample key is stored once per (op, size) so the log proves the + // size->component mapping instead of leaving it inferred. + if (p.object_size != 0) { + Component& c = components_[{is_get, p.object_size}]; + c.calls += 1; + c.keys += p.keys; + c.bytes += p.bytes; + c.total += total; + c.xfer += p.xfer; + if (c.sample_key.empty() && p.key0 != nullptr) c.sample_key = *p.key0; + } + + const size_t cadence = RangedDebugSummarySeconds(); + if (cadence == 0) return; + const auto now = std::chrono::steady_clock::now(); + if (RangedSeconds(last_report_, now) < static_cast(cadence)) return; + last_report_ = now; + Emit("GET", get_); + Emit("PUT", put_); + EmitComponentsLocked(); + } + + // The cadence emit only fires when a call lands more than one window after + // the previous one, so a measured phase shorter than the cadence reports no + // line at all -- indistinguishable from the instrumentation being off -- and + // the final partial window of a long run is lost the same way. + void FlushFinal() { + std::lock_guard lock(mutex_); + if (get_.calls == 0 && put_.calls == 0) return; + Emit("GET", get_); + Emit("PUT", put_); + EmitComponentsLocked(); + } + + private: + struct Totals { + uint64_t calls = 0; + uint64_t items = 0; + double total = 0, resolve = 0, classify = 0, build = 0, validate = 0, commit = 0, route = 0, + xfer = 0, xfer_plan = 0, xfer_submit = 0, xfer_wait = 0, lock = 0, remote = 0, bytes = 0; + }; + struct Component { + uint64_t calls = 0; + uint64_t keys = 0; + double bytes = 0, total = 0, xfer = 0; + std::string sample_key; + }; + static void Emit(const char* name, const Totals& t) { + if (t.calls == 0) return; + const double other = t.total - (t.resolve + t.classify + t.build + t.validate + t.commit + + t.route + t.xfer + t.lock + t.remote); + auto share = [&](double v) { return t.total > 0 ? 100.0 * v / t.total : 0.0; }; + MORI_UMBP_INFO( + "[RangedCall][dbg] SUMMARY {} calls={} total={:.3f}s mean_call={:.1f}us bytes={:.2f}GiB " + "items_per_call={:.0f} | resolve={:.1f}% classify={:.1f}% build={:.1f}% validate={:.1f}% " + "commit={:.1f}% route={:.1f}% xfer={:.1f}%(plan={:.1f}% submit={:.1f}% wait={:.1f}%) " + "lock={:.1f}% remote={:.1f}% " + "other={:.1f}% | xfer_only={:.2f}GiB/s end2end={:.2f}GiB/s", + name, t.calls, t.total, 1e6 * t.total / t.calls, t.bytes / (1024.0 * 1024 * 1024), + static_cast(t.items) / t.calls, share(t.resolve), share(t.classify), share(t.build), + share(t.validate), share(t.commit), share(t.route), share(t.xfer), share(t.xfer_plan), + share(t.xfer_submit), share(t.xfer_wait), share(t.lock), share(t.remote), share(other), + t.xfer > 0 ? (t.bytes / t.xfer) / (1024.0 * 1024 * 1024) : 0.0, + t.total > 0 ? (t.bytes / t.total) / (1024.0 * 1024 * 1024) : 0.0); + } + void EmitComponentsLocked() const { + for (const auto& [id, c] : components_) { + if (c.calls == 0) continue; + MORI_UMBP_INFO( + "[RangedCall][dbg] COMPONENT {} obj={} key0='{}' calls={} keys={} bytes={:.3f}GiB " + "total={:.3f}s mean_call={:.1f}us keys_per_call={:.1f} xfer_only={:.2f}GiB/s " + "end2end={:.2f}GiB/s", + id.first ? "GET" : "PUT", id.second, c.sample_key, c.calls, c.keys, + c.bytes / (1024.0 * 1024 * 1024), c.total, 1e6 * c.total / c.calls, + static_cast(c.keys) / c.calls, + c.xfer > 0 ? (c.bytes / c.xfer) / (1024.0 * 1024 * 1024) : 0.0, + c.total > 0 ? (c.bytes / c.total) / (1024.0 * 1024 * 1024) : 0.0); + } } + + std::mutex mutex_; + Totals get_; + Totals put_; + // key = (is_get, object_size); ordered so the rollup prints deterministically. + std::map, Component> components_; + std::chrono::steady_clock::time_point last_report_ = std::chrono::steady_clock::now(); +}; + +RangedStats& RangedStatsInstance() { + // Deliberately leaked: a destructor here would run at static-destruction + // time, after singletons it logs through may be gone. atexit gives back the + // final flush the leak would otherwise cost, and runs before those teardowns. + static RangedStats* stats = [] { + auto* created = new RangedStats(); + std::atexit([] { RangedStatsInstance().FlushFinal(); }); + return created; + }(); + return *stats; } +// Emits at scope exit so every early return in a ranged call is covered, the +// same reason ScopedBatchBandwidth is a destructor. +class ScopedRangedReport { + public: + ScopedRangedReport(const char* op, RangedPhases& phases, bool enabled, const double& local_bytes, + const double& remote_bytes) + : op_(op), + phases_(phases), + enabled_(enabled), + local_bytes_(local_bytes), + remote_bytes_(remote_bytes), + start_(std::chrono::steady_clock::now()) {} + ~ScopedRangedReport() { + if (!enabled_) return; + // Read at exit: the remote phase keeps adding after any mid-call snapshot. + phases_.bytes = local_bytes_ + remote_bytes_; + const double total = RangedSeconds(start_, std::chrono::steady_clock::now()); + RangedStatsInstance().Record(op_, phases_, total); + if (!RangedDebugSampleThisCall()) return; + const double other = + total - (phases_.resolve + phases_.classify + phases_.build + phases_.validate + + phases_.commit + phases_.route + phases_.xfer + phases_.lock + phases_.remote); + MORI_UMBP_INFO( + "[RangedCall][dbg] op={} obj={} key0='{}' keys={} items={} local={} remote={} bytes={} " + "total={:.1f}us resolve={:.1f}us classify={:.1f}us build={:.1f}us validate={:.1f}us " + "commit={:.1f}us route={:.1f}us xfer={:.1f}us(plan={:.1f} submit={:.1f} wait={:.1f}) " + "lock={:.1f}us remote_phase={:.1f}us other={:.1f}us", + op_, phases_.object_size, phases_.key0 != nullptr ? *phases_.key0 : std::string("?"), + phases_.keys, phases_.items, phases_.local_keys, phases_.remote_keys, phases_.bytes, + total * 1e6, phases_.resolve * 1e6, phases_.classify * 1e6, phases_.build * 1e6, + phases_.validate * 1e6, phases_.commit * 1e6, phases_.route * 1e6, phases_.xfer * 1e6, + phases_.xfer_plan * 1e6, phases_.xfer_submit * 1e6, phases_.xfer_wait * 1e6, + phases_.lock * 1e6, phases_.remote * 1e6, other * 1e6); + } + + private: + const char* op_; + RangedPhases& phases_; + bool enabled_; + const double& local_bytes_; + const double& remote_bytes_; + std::chrono::steady_clock::time_point start_; +}; + +// Runs a callable at scope exit. Same reason the two Scoped* classes below are +// destructor-based: a ranged call has several early returns that have still +// moved bytes, and end-of-call work must not depend on each `return` having +// remembered to do it. +template +class ScopeExit { + public: + explicit ScopeExit(Fn fn) : fn_(std::move(fn)) {} + ~ScopeExit() { fn_(); } + ScopeExit(const ScopeExit&) = delete; + ScopeExit& operator=(const ScopeExit&) = delete; + + private: + Fn fn_; +}; +template +ScopeExit(Fn) -> ScopeExit; + +// Ranged calls have several early returns that can still have moved bytes (a +// batch fully served from the local medium returns before the remote phase is +// even considered). Emitting from a destructor keeps every exit covered, +// including ones added later, instead of relying on each `return` remembering +// to observe first. The accumulators are held by reference and read at scope +// exit, so callers just add to them as keys succeed. +class ScopedBatchBandwidth { + public: + ScopedBatchBandwidth(MasterClient* master_client, const char* metric_name, + const char* metric_help, const double& local_bytes, + const double& remote_bytes) + : master_client_(master_client), + metric_name_(metric_name), + metric_help_(metric_help), + local_bytes_(local_bytes), + remote_bytes_(remote_bytes), + start_(std::chrono::steady_clock::now()) {} + + ScopedBatchBandwidth(const ScopedBatchBandwidth&) = delete; + ScopedBatchBandwidth& operator=(const ScopedBatchBandwidth&) = delete; + + ~ScopedBatchBandwidth() { + const double seconds = std::chrono::duration_cast>( + std::chrono::steady_clock::now() - start_) + .count(); + ObserveBatchBandwidth(master_client_, local_bytes_, seconds, metric_name_, metric_help_, + "local"); + ObserveBatchBandwidth(master_client_, remote_bytes_, seconds, metric_name_, metric_help_, + "remote"); + } + + private: + MasterClient* master_client_; + const char* metric_name_; + const char* metric_help_; + const double& local_bytes_; + const double& remote_bytes_; + std::chrono::steady_clock::time_point start_; +}; + +// --------------------------------------------------------------------------- +// Page / size math +// --------------------------------------------------------------------------- + +// --- Cross-key parallelism for the self-target (local) paths. -------------- +// In distributed mode 1 key == 1 page (master page_size == KV block size), so +// a self-target Put/Get copies one ~MiB-scale block per key. The parallelism +// that pays is therefore ACROSS the many keys of one BatchPut/BatchGet +// (different threads -> different keys), not within one key's pages. The +// per-key copy itself lives in LocalCopyEngine since Phase 6. Threads via +// UMBP_DRAM_{READ,WRITE}_THREADS (same envs as local mode's DRAMTier). inline int LocalCopyThreads(const char* env_name) { int t = 4; if (const char* e = std::getenv(env_name)) { @@ -204,125 +508,109 @@ inline int LocalCopyThreads(const char* env_name) { return t; } -template -inline void LocalParallelFor(size_t n, int num_threads, Fn&& fn) { - if (n == 0) return; - if (num_threads > static_cast(n)) num_threads = static_cast(n); - if (num_threads <= 1) { - for (size_t i = 0; i < n; ++i) fn(i); - return; +inline std::chrono::milliseconds ResolveBusyRetryTimeout() { + uint64_t ms = 30000; + if (const char* value = std::getenv("UMBP_RESOLVE_BUSY_TIMEOUT_MS")) { + char* end = nullptr; + const unsigned long long parsed = std::strtoull(value, &end, 10); + if (end != value && *end == '\0' && parsed > 0) ms = parsed; } - std::atomic next{0}; - auto worker = [&]() { - size_t i; - while ((i = next.fetch_add(1)) < n) fn(i); - }; - std::vector pool; - pool.reserve(num_threads); - for (int t = 0; t < num_threads; ++t) pool.emplace_back(worker); - for (auto& th : pool) th.join(); + return std::chrono::milliseconds(std::min(ms, 300000)); } -inline uint64_t LogicalPageBytes(size_t i, size_t num_pages, uint64_t page_size, - size_t total_size) { - return (i + 1 == num_pages) ? (total_size - i * page_size) : page_size; +std::vector ResolveLocalBatchWithBusyRetry(PeerPool* pool, + const std::vector& keys, + bool include_descs, + bool allow_file_refs) { + if (pool == nullptr) return std::vector(keys.size()); + const auto deadline = std::chrono::steady_clock::now() + ResolveBusyRetryTimeout(); + std::chrono::milliseconds backoff{1}; + size_t attempts = 0; + while (true) { + auto resolved = pool->BatchResolve(keys, include_descs, allow_file_refs); + const bool busy = std::any_of(resolved.begin(), resolved.end(), [](const auto& entry) { + return EffectiveResolveOutcome(entry.resolved) == ResolveOutcome::kBusy; + }); + if (!busy) return resolved; + ++attempts; + const auto remaining = deadline - std::chrono::steady_clock::now(); + if (remaining <= std::chrono::steady_clock::duration::zero()) { + MORI_UMBP_WARN("[PoolClient] local BatchResolve BUSY timeout after {} attempts", attempts); + return resolved; + } + const auto sleep_for = + std::min(backoff, std::chrono::duration_cast(remaining)); + if (sleep_for.count() > 0) std::this_thread::sleep_for(sleep_for); + backoff = std::min(backoff * 2, std::chrono::milliseconds{50}); + } } -bool SizeMatchesAllocation(uint64_t size, size_t num_pages, uint64_t page_size) { - if (page_size == 0 || num_pages == 0 || size == 0) return false; - if (size > num_pages * page_size) return false; - if (size <= (num_pages - 1) * page_size) return false; - return true; +// Classify a caller's buffer so engine selection can see what it really is. +// +// An unregistered pointer used to be described unconditionally as host bytes, +// which is what CanHandle dispatches on: a device pointer would therefore be +// claimed by LocalCopyEngine and memcpy'd, or staged through MoriIoEngine's +// host bounce buffer. Neither fails at the call site — they corrupt or fault +// somewhere else — so the classification has to happen before any engine sees +// the pair. Registered buffers never come through here; their ref already +// carries loc/device from RegisterMemory. +TransferRef ClassifiedUserBytes(void* ptr, uint64_t size) { + const PointerLocation location = DetectPointerLocation(ptr); + if (!location.IsDevice()) return TransferRef::HostBytes(ptr, size); + return TransferRef::HostBytes(ptr, size, mori::io::MemoryLocationType::GPU, location.device_id); } -uint64_t AlignUp(uint64_t value, uint64_t alignment) { - if (alignment == 0) return value; - uint64_t remainder = value % alignment; - return remainder == 0 ? value : (value + alignment - remainder); +// A file ref can only be READ into device memory: GdsEngine claims the +// (file, GPU) pair and no engine claims (file, host). Callers use this to +// decide whether asking a medium for a file ref is even useful. +bool DestinationIsDevice(void* ptr, uint64_t size) { + return ClassifiedUserBytes(ptr, size).loc == mori::io::MemoryLocationType::GPU; } -// Group `pages` by buffer_index, preserving first-seen ordering so -// IOEngine BatchWrite/BatchRead pair indexing stays predictable. -struct ScatterGroup { - uint32_t buffer_index; - std::vector src_page_indices; -}; -std::vector GroupPagesByBuffer(const std::vector& pages) { - std::vector groups; - std::unordered_map buf_to_group; - groups.reserve(pages.size()); - for (size_t i = 0; i < pages.size(); ++i) { - uint32_t bi = pages[i].buffer_index; - auto it = buf_to_group.find(bi); - if (it == buf_to_group.end()) { - buf_to_group.emplace(bi, groups.size()); - groups.push_back({bi, {}}); - it = buf_to_group.find(bi); +// Read-side range validity: inside the object and non-overlapping. Weaker than +// RangesTileObject, which the write side needs — a read may take any subset of +// an object, but a write must account for every byte it commits. +bool RangesAreDisjointAndInBounds(size_t object_size, const std::vector& sizes, + const std::vector& offsets) { + if (object_size == 0 || sizes.empty() || sizes.size() != offsets.size()) return false; + std::vector order(sizes.size()); + std::iota(order.begin(), order.end(), size_t{0}); + std::sort(order.begin(), order.end(), + [&](size_t a, size_t b) { return offsets[a] < offsets[b]; }); + size_t previous_end = 0; + bool first = true; + for (size_t index : order) { + if (sizes[index] == 0 || IsObjectRangeOverflow(offsets[index], sizes[index], object_size)) { + return false; } - groups[it->second].src_page_indices.push_back(i); + if (!first && offsets[index] < previous_end) return false; + previous_end = offsets[index] + sizes[index]; + first = false; } - return groups; -} - -// --------------------------------------------------------------------------- -// SSD / lease env knobs -// --------------------------------------------------------------------------- - -uint32_t ReleaseLeaseMaxRetries() { - static const uint32_t v = GetEnvUint32("UMBP_RELEASE_LEASE_MAX_RETRIES", 2, /*min_allowed=*/1); - return v; -} - -// Crash-restart SSD leftover policy (discard): wipe physical SSD bytes a -// previous crashed process left behind at startup. Default on; set -// UMBP_SSD_STARTUP_DISCARD=0/false to keep leftover (opt-out hook for a future -// rebuild-from-backend policy — there is no rebuild path yet). -bool SsdStartupDiscardEnabled() { - const char* v = std::getenv("UMBP_SSD_STARTUP_DISCARD"); - if (v == nullptr) return true; - const std::string s(v); - return !(s == "0" || s == "false" || s == "FALSE" || s == "off"); -} - -// Total attempts for a remote SSD get. Retryable outcomes (kRetry) are -// NO_SLOT (transient slot exhaustion, no slot claimed) and a reader-local lease -// expiry. Defaults to 1 (no retry); operators opt into bounded retry to absorb -// slot contention. A slow/failed PrepareSsdRead (DEADLINE_EXCEEDED etc.) is a -// hard failure, NOT retried — the peer may already hold a claimed slot, so a -// retry would only pile up more staging-slot occupation. NOT_FOUND is always a -// definitive miss (no retry). -uint32_t SsdGetTransientMaxAttempts() { - static const uint32_t v = GetEnvUint32("UMBP_SSD_GET_MAX_ATTEMPTS", 1, /*min_allowed=*/1); - return v; + return true; } -std::chrono::milliseconds SsdGetRetryBackoff() { - static const auto v = std::chrono::milliseconds( - GetEnvUint32("UMBP_SSD_GET_RETRY_BACKOFF_MS", 2, /*min_allowed=*/1)); - return v; -} +// Scratch-arena slice alignment. Objects are packed back-to-back into the +// arena; 64 B keeps each slice on a cache line so two concurrently-filled +// slices never share one. +constexpr size_t kRangedScratchAlignment = 64; -// Bounded client deadline for a single PrepareSsdRead RPC so a hung/slow peer -// cannot block the serial per-key remote SSD read loop indefinitely. A read -// that cannot even return within this budget is past any useful lease window -// and is treated as a hard failure (NOT retried: the peer may already hold a -// claimed slot, and a retry would only claim another). When the env is unset -// the caller falls back to the configured lease timeout (cluster-homogeneous -// assumption); 0 here means "use the fallback". -std::chrono::milliseconds SsdPrepareRpcTimeoutOverride() { - static const auto v = GetEnvMilliseconds("UMBP_SSD_PREPARE_TIMEOUT_MS", - std::chrono::milliseconds(0), /*min_allowed=*/0); - return v; +bool AlignUpChecked(size_t value, size_t alignment, size_t* out) { + if (!out || alignment == 0) return false; + const size_t remainder = value % alignment; + if (remainder == 0) { + *out = value; + return true; + } + const size_t add = alignment - remainder; + if (value > std::numeric_limits::max() - add) return false; + *out = value + add; + return true; } -// Bounded client deadline for a best-effort ReleaseSsdLease RPC. Release is -// never on the critical path (the slot is also reclaimed by TTL), so cap it -// tightly to avoid blocking on a slow peer. -std::chrono::milliseconds ReleaseLeaseRpcTimeout() { - static const auto v = GetEnvMilliseconds("UMBP_RELEASE_LEASE_TIMEOUT_MS", - std::chrono::milliseconds(1000), /*min_allowed=*/1); - return v; -} +// --------------------------------------------------------------------------- +// Lease env knobs +// --------------------------------------------------------------------------- // Peer-side DRAM/HBM read lease: how long a single Resolve protects its key's // pages from concurrent local Evict, covering one RDMA read. Only needs to @@ -334,13 +622,10 @@ std::chrono::milliseconds DramReadLeaseTtl() { return v; } -// Peer-side SSD read-staging slot lease: how long a claimed staging slot is -// reserved for a reader (peer reclaims by TTL if the best-effort ReleaseSsdLease -// is lost) and, mirrored back to the reader, the validity window it anchors at -// t_send. Must exceed one SSD read + RDMA (slower than DRAM), but too long -// pins one of only ~16 slots on a lost release, so 3 s balances both. Also the -// fallback for the PrepareSsdRead RPC deadline when UMBP_SSD_PREPARE_TIMEOUT_MS -// is unset. +// The SSD equivalent, which is a different number for a different reason: it has +// to cover an SSD read plus the RDMA, not just the RDMA. It is also the window a +// staging page stays claimed, so it sets this backend's read concurrency against +// a fixed arena -- shortening it frees pages and costs read coalescing. std::chrono::milliseconds SsdReadLeaseTtl() { static const auto v = GetEnvMilliseconds("UMBP_SSD_READ_LEASE_MS", std::chrono::milliseconds(3000), /*min_allowed=*/1); @@ -351,44 +636,160 @@ std::chrono::milliseconds SsdReadLeaseTtl() { // Config / proto translation // --------------------------------------------------------------------------- -// Build a TierConfig from the PoolClientConfig (DRAM only — HBM -// support requires per-tier buffer plumbing the upper layers don't -// currently provide). Returns an empty config when the engine has -// no DRAM buffers, signalling that no DRAM allocator should be built. -PeerDramAllocator::TierConfig BuildDramTierConfig(const std::vector& bufs, - const std::vector& mems) { - PeerDramAllocator::TierConfig cfg; - if (bufs.size() != mems.size()) return cfg; - for (size_t i = 0; i < bufs.size(); ++i) { - cfg.buffer_sizes.push_back(bufs[i].size); - msgpack::sbuffer sbuf; - msgpack::pack(sbuf, mems[i]); - cfg.buffer_descs.emplace_back(sbuf.data(), sbuf.data() + sbuf.size()); - // Local host base pointer, so PeerDramAllocator can resolve a committed - // key's pages to readable segments for the SSD copy worker. - cfg.buffer_bases.push_back(bufs[i].buffer); - } - return cfg; -} - // Translate a peer-side ::umbp::AllocateSlotResponse / ResolveKeyResponse // into the C++ shapes our code consumes. PoolClient::SlotPlan FromAllocateSlotResponse(const ::umbp::AllocateSlotResponse& resp) { PoolClient::SlotPlan p; p.slot_id = resp.slot_id(); p.page_size = resp.page_size(); + p.backend_id = resp.backend_id(); p.pages.reserve(resp.pages_size()); for (const auto& pp : resp.pages()) p.pages.push_back({pp.buffer_index(), pp.page_index()}); p.descs.reserve(resp.descs_size()); for (const auto& d : resp.descs()) { BufferMemoryDescBytes b; b.buffer_index = d.buffer_index(); + b.backend_id = d.backend_id(); b.desc_bytes.assign(d.desc().begin(), d.desc().end()); p.descs.push_back(std::move(b)); } return p; } +// --------------------------------------------------------------------------- +// Engine outcome -> per-entry failure +// +// A TransferItem's `tag` IS the index of the entry that produced it, so +// mapping an engine failure back to keys is an index lookup. A failed plan +// fails EVERY key that contributed a segment to it (per-item AND) — the same +// granularity the pre-refactor per-(localMR, remoteMR) group had, for the same +// reason: the wire reports success per transfer, not per key. +// +// Free templates rather than PoolClient members so both entry types share one +// definition; deduction means neither ever has to be named here. +// --------------------------------------------------------------------------- + +template +void ApplyTransferFailures(std::vector& entries, + const std::vector& failures, const char* what) { + for (const auto& f : failures) { + for (size_t tag : f.tags) { + if (tag >= entries.size()) continue; + auto& entry = entries[tag]; + MORI_UMBP_ERROR("{} transfer failed: code={} msg='{}' peer_engine='{}' key='{}'", what, + f.code, f.message, f.endpoint, + (entry.item && entry.item->key) ? *entry.item->key : std::string{""}); + entry.failed = true; + } + } +} + +// Items no engine could carry. Distinct from a failed transfer: these never +// reached the wire, which usually means a peer descriptor was missing, the item +// was larger than the engine's bounce pool, or the caller's buffer is +// unregistered GPU memory going to a remote peer — staging that would mean a +// host memcpy from device memory, so MoriIoEngine rejects it instead. +template +void ApplyRejectedTags(std::vector& entries, const std::vector& tags, + const char* what) { + for (size_t tag : tags) { + if (tag >= entries.size()) continue; + auto& entry = entries[tag]; + MORI_UMBP_WARN("{} transfer unplannable (no engine claimed the endpoints), key='{}'", what, + (entry.item && entry.item->key) ? *entry.item->key : std::string{""}); + entry.failed = true; + } +} + +std::vector EffectiveBackendConfigs(const PoolClientConfig& config) { + if (!config.backends.empty()) return config.backends; + + BackendInstanceConfig legacy; + legacy.name = DefaultBackendInstanceName(config.medium); + legacy.tier = config.medium; + legacy.dram = config.dram; + legacy.hbm = config.hbm; + legacy.ssd = config.ssd; + legacy.ssd_staging_buffer_slots = config.ssd_staging_buffer_slots; + return {std::move(legacy)}; +} + +std::unique_ptr MakeConfiguredBackend(const BackendInstanceConfig& config, + uint64_t page_size, bool staging_use_hugepages, + uint64_t staging_hugepage_size) { + switch (config.tier) { + case TierType::DRAM: { + PageBackend::OwnershipConfig ownership; + ownership.buffer_sizes = config.dram.buffer_sizes; + ownership.use_hugepages = config.dram.use_hugepages; + ownership.hugepage_size = config.dram.hugepage_size; + ownership.numa_node = config.dram.numa_node; + ownership.prefault = config.dram.prefault; + return MakePageBackend(TierType::DRAM, page_size, std::move(ownership), + /*pending_ttl=*/std::chrono::milliseconds{30000}, + /*read_lease_ttl=*/DramReadLeaseTtl()); + } + case TierType::HBM: + return MakeHbmBackend(page_size, config.hbm.device, config.hbm.buffer_sizes, + /*pending_ttl=*/std::chrono::milliseconds{30000}, + /*read_lease_ttl=*/DramReadLeaseTtl()); + case TierType::SSD: { + SsdBackend::Config ssd_cfg; + ssd_cfg.page_size = page_size; + ssd_cfg.staging_pages = config.ssd_staging_buffer_slots > 0 + ? static_cast(config.ssd_staging_buffer_slots) + : 16; + ssd_cfg.staging_use_hugepages = staging_use_hugepages; + ssd_cfg.staging_hugepage_size = staging_hugepage_size; + ssd_cfg.ssd = config.ssd; + ssd_cfg.ssd.enabled = true; + ssd_cfg.read_lease_ttl = SsdReadLeaseTtl(); + return MakeSsdBackend(std::move(ssd_cfg)); + } + case TierType::UNKNOWN: + return nullptr; + } + return nullptr; +} + +std::unique_ptr MakeConfiguredPoolPolicy( + const PoolClientConfig& config, const std::vector& backends, + const BackendRegistry& registry, std::shared_ptr* tier_graph) { + switch (config.placement_policy) { + case PoolPlacementPolicy::SINGLE_BACKEND: + return MakeSingleBackendPolicy(); + case PoolPlacementPolicy::WEIGHTED: { + std::vector weights; + weights.reserve(backends.size()); + for (const auto& backend : backends) { + if (backend.placement_weight == 0) { + MORI_UMBP_ERROR("[PoolClient] weighted backend '{}' has zero placement weight", + backend.name); + return nullptr; + } + weights.push_back(BackendPlacementWeight{backend.name, backend.placement_weight}); + } + return MakeWeightedPlacementPolicy(std::move(weights)); + } + case PoolPlacementPolicy::TIERED: + if (config.logical_tiers.empty()) { + MORI_UMBP_ERROR("[PoolClient] tiered placement has no logical tiers"); + return nullptr; + } + if (tier_graph == nullptr) return nullptr; + { + auto compiled = LogicalTierGraph::Compile(config.logical_tiers, registry); + if (!compiled.ok()) { + MORI_UMBP_ERROR("[PoolClient] invalid logical tier graph: {}", compiled.error); + return nullptr; + } + *tier_graph = std::move(compiled.graph); + return MakeTieredPlacementPolicy(*tier_graph); + } + } + return nullptr; +} + } // namespace // --------------------------------------------------------------------------- @@ -402,144 +803,261 @@ bool PoolClient::Init() { bool expected = false; if (!initialized_.compare_exchange_strong(expected, true)) return true; - master_client_ = std::make_unique(config_.master_config); + if (config_.workload_trace_path.empty()) { + if (const char* path = std::getenv("UMBP_WORKLOAD_TRACE_PATH")) { + config_.workload_trace_path = path; + } + } + if (const char* value = std::getenv("UMBP_WORKLOAD_TRACE_CLIENT_ID")) { + config_.workload_trace_client_id = static_cast(std::strtoull(value, nullptr, 10)); + } + if (const char* value = std::getenv("UMBP_WORKLOAD_TRACE_SEED")) { + config_.workload_trace_seed = std::strtoull(value, nullptr, 10); + } + + if (!config_.policy_config_path.empty()) { + auto loaded = LoadBackendPolicyFile(config_.policy_config_path); + std::string error; + if (!loaded.ok() || !ApplyBackendPolicy(*loaded.config, &config_, &error)) { + MORI_UMBP_ERROR("[PoolClient] failed to load backend policy '{}': {}", + config_.policy_config_path, loaded.ok() ? error : loaded.error); + initialized_.store(false); + return false; + } + } - // IO Engine setup (RDMA data plane). + // No address, no master. Everything below -- backends, transfer engine, peer + // service -- is built either way; only routing, registration and the + // heartbeat are skipped. That is what lets one binary be a single node or a + // cluster member depending on config alone. + if (!config_.master_config.master_address.empty()) { + master_client_ = std::make_unique(config_.master_config); + } else { + MORI_UMBP_INFO("[PoolClient] no master configured; running single-node"); + } + + // The one byte-moving path (design doc §4). Order is preference order: + // a pair both of whose endpoints are host-addressable never reaches the wire. + // + // This is the composition root, and the only place any concrete engine is + // named — the same rule Phase 5 applies to backends. Everything downstream + // holds TransferEngine / MemoryRegistrar / PeerDirectory. + auto composite = std::make_unique(); + composite->AddEngine(std::make_unique()); + // Both-local pairs with a GPU endpoint. Disjoint from LocalCopyEngine (which + // requires both sides CPU) and from MoriIoEngine (which requires exactly one + // side remote), so this is registration order as documentation, not as a + // tie-break — but it must come before the wire engine on principle, since an + // HBM pair that reached mori-io would be refused rather than served slowly. + auto hbm = std::make_unique(); + // Kept as a raw observer so the live medium's host buffers can be declared + // kernel-addressable once it has allocated them, below. This is the one + // engine-specific reference PoolClient holds, and it carries no data-plane + // call: everything that moves bytes still goes through transfer_engine_. + hbm_engine_ = hbm.get(); + composite->AddEngine(std::move(hbm)); +#ifdef UMBP_ENABLE_GDS + // The file->GPU engine for SsdBackend's FileRefs. It claims only (file, + // device) pairs, which no memory engine touches, so registration order is + // documentation. Present only when the build found hipfile. + composite->AddEngine(std::make_unique()); +#endif if (!config_.io_engine.host.empty()) { mori::io::IOEngineConfig io_cfg; io_cfg.host = config_.io_engine.host; io_cfg.port = config_.io_engine.port; - io_engine_ = std::make_unique(config_.master_config.node_id, io_cfg); - mori::io::RdmaBackendConfig rdma_cfg; - - rdma_cfg.qpPerTransfer = 4; - rdma_cfg.enableTransferChunking = true; - rdma_cfg.numNicsPerTransfer = 4; - // UMBP only sets the config defaults above. All RDMA knobs (qpPerTransfer - // / postBatchSize / numWorkerThreads / pollCqMode / chunking / numNics / - // ...) are overridable via MORI_IO_* env in the RDMA backend ctor, shared - // by every IO backend user; no UMBP-specific entry points. - io_engine_->CreateBackend(mori::io::BackendType::RDMA, rdma_cfg); - - staging_buffer_ = std::make_unique(config_.staging_buffer_size); - std::memset(staging_buffer_.get(), 0, config_.staging_buffer_size); - staging_mem_ = io_engine_->RegisterMemory(staging_buffer_.get(), config_.staging_buffer_size, - -1, mori::io::MemoryLocationType::CPU); - - for (const auto& dram : config_.dram_buffers) { - if (dram.buffer && dram.size > 0) { - auto mem = io_engine_->RegisterMemory(dram.buffer, dram.size, -1, - mori::io::MemoryLocationType::CPU); - export_dram_mems_.push_back(mem); - } + auto rdma = std::make_unique(config_.master_config.node_id, io_cfg, + config_.staging_buffer_size); + if (!rdma->Init()) { + MORI_UMBP_ERROR("[PoolClient] MoriIoEngine init failed on {}:{}", config_.io_engine.host, + config_.io_engine.port); + Shutdown(); + return false; } - MORI_UMBP_INFO("[PoolClient] IOEngine initialized on {}:{} ({} DRAM buffers)", - config_.io_engine.host, config_.io_engine.port, export_dram_mems_.size()); + peer_directory_ = rdma.get(); + composite->AddEngine(std::move(rdma)); } + transfer_engine_ = std::move(composite); - // Peer-side allocator: even SSD-only deployments build one (with - // empty DRAM/HBM tiers) so the SSD CommitSsdWrite path has an event - // outbox to push ADD events into. + // Peer-side backend instances. The legacy single-medium selector is + // synthesized into one entry here; topology-aware callers may + // supply several named instances, including several of the same TierType. + // This composition root remains the only place concrete backend factories + // are named. The configured PoolPolicy below owns placement among them. const uint64_t page_size = config_.dram_page_size > 0 ? config_.dram_page_size : 2ULL * 1024 * 1024; - PeerDramAllocator::TierConfig dram_cfg = - io_engine_ ? BuildDramTierConfig(config_.dram_buffers, export_dram_mems_) - : PeerDramAllocator::TierConfig{}; - PeerDramAllocator::TierConfig hbm_cfg; // HBM not currently exposed via PoolClientConfig - peer_alloc_ = - std::make_unique(page_size, std::move(dram_cfg), std::move(hbm_cfg), - /*pending_ttl=*/std::chrono::milliseconds{30000}, - /*read_lease_ttl=*/DramReadLeaseTtl()); - peer_alloc_->StartReaper(); - master_client_->SetPeerDramAllocator(peer_alloc_.get()); - - // Peer-side SSD tier owner. Built only when SSD is enabled; registered as an - // owned-location event source + SSD capacity provider. The copy-on-commit - // pipeline below populates it. Clear() quiesces the pipeline and clears - // peer_ssd_ alongside peer_alloc_. - if (config_.ssd.enabled) { - peer_ssd_ = std::make_unique(config_.ssd); - // Crash-restart leftover (discard): wipe stale SSD bytes before the - // pipeline starts so used capacity and the empty owned_ map start consistent - // (env-gated; see SsdStartupDiscardEnabled / DiscardLeftoverOnStartup). - if (SsdStartupDiscardEnabled()) peer_ssd_->DiscardLeftoverOnStartup(); - master_client_->SetPeerSsdManager(peer_ssd_.get()); - // Copy-on-commit pipeline: commits enqueue here, a worker pins the DRAM - // pages and writes them to the local SSD tier. Started below. - ssd_copy_pipeline_ = std::make_unique(peer_alloc_.get(), peer_ssd_.get(), - config_.copy_pipeline.queue_depth, - config_.copy_pipeline.worker_threads); - ssd_copy_pipeline_->Start(); + + auto backend_configs = EffectiveBackendConfigs(config_); + if (backend_configs.empty()) { + MORI_UMBP_ERROR("[PoolClient] no backend instances configured"); + Shutdown(); + return false; + } + medium_ = backend_configs.front().tier; // legacy default for paths without a routed tier + for (const auto& backend_config : backend_configs) { + if (backend_config.name.empty() || registry_.Get(backend_config.name) != nullptr) { + MORI_UMBP_ERROR("[PoolClient] backend instance name '{}' is empty or duplicated", + backend_config.name); + Shutdown(); + return false; + } + auto backend = + MakeConfiguredBackend(backend_config, page_size, config_.ssd_staging_use_hugepages, + config_.ssd_staging_hugepage_size); + if (backend == nullptr) { + MORI_UMBP_ERROR("[PoolClient] backend '{}' has unknown medium {}", backend_config.name, + static_cast(backend_config.tier)); + Shutdown(); + return false; + } + + // Apply the generic operation/byte/latency instrumentation uniformly to + // every named backend, including multiple instances of the same medium. + backend = MakeInstrumentedBackend(std::move(backend)); + + // A node with no master owns its whole storage policy, and both halves of + // that follow from the same fact: nobody will ever call Evict() on this + // backend, and nobody will ever drain its heartbeat outbox. So the backend + // evicts for itself, and stops recording events addressed to a master that + // does not exist. Both are no-ops in a cluster -- master-driven eviction + // stays the only policy there, and the outbox keeps feeding the heartbeat. + // + // Before Init(), which is the contract EnableLocalEviction states. + if (!master_client_) { + backend->SetEventPublishing(false); + backend->EnableLocalEviction(config_.local_evict_high_watermark, + config_.local_evict_low_watermark); + } + + // Narrowed to MemoryRegistrar: a backend publishes endpoints, it does not + // move bytes, and that is a compile-time fact (design doc §5 Rule C). + if (!backend->Init(static_cast(transfer_engine_.get()))) { + MORI_UMBP_ERROR("[PoolClient] backend '{}' ({}) Init failed", backend_config.name, + TierTypeName(backend_config.tier)); + backend->Shutdown(); + Shutdown(); + return false; + } + if (!registry_.Register(backend_config.name, std::move(backend))) { + Shutdown(); + return false; + } + } + + std::shared_ptr tier_graph; + auto placement_policy = + MakeConfiguredPoolPolicy(config_, backend_configs, registry_, &tier_graph); + if (placement_policy == nullptr) { + MORI_UMBP_ERROR("[PoolClient] invalid placement policy configuration"); + Shutdown(); + return false; + } + default_pool_ = + std::make_unique(®istry_, std::move(placement_policy), transfer_engine_.get()); + const bool weighted_placement = config_.placement_policy == PoolPlacementPolicy::WEIGHTED || + config_.placement_policy == PoolPlacementPolicy::TIERED; + if (master_client_) { + master_client_->SetAggregateBackendCapacities(weighted_placement); + master_client_->SetBackendRegistry(®istry_); + master_client_->SetPeerPool(default_pool_.get()); + } + + // Declare every host region the gather kernel may dereference: every live + // backend's buffers and the ranged scratch arenas. These are copied + // fragment-wise against GPU callers, which is the shape the kernel exists + // for. Registration is best-effort and asynchronous for large pools, so + // until it lands the engine just takes hipMemcpy — no ordering requirement + // here beyond the buffers existing. + if (hbm_engine_ != nullptr) { + for (auto* live : registry_.All()) { + if (live == nullptr) continue; + for (uint32_t i = 0; i < live->BufferCount(); ++i) { + const TransferRef ref = live->BufferRef(i); + // Only host buffers: a device-resident medium is already addressable + // from a kernel, and hipHostRegister on it would fail. + if (ref.HasHostPtr() && ref.loc == mori::io::MemoryLocationType::CPU) { + hbm_engine_->AddHostGatherRegion(ref.host_ptr, ref.size); + } + } + } + if (config_.ranged_get_scratch_buffer != nullptr && config_.ranged_get_scratch_size > 0) { + hbm_engine_->AddHostGatherRegion(config_.ranged_get_scratch_buffer, + config_.ranged_get_scratch_size); + } + if (config_.ranged_put_scratch_buffer != nullptr && config_.ranged_put_scratch_size > 0) { + hbm_engine_->AddHostGatherRegion(config_.ranged_put_scratch_buffer, + config_.ranged_put_scratch_size); + } } + if (master_client_) master_client_->SetBackendRegistry(®istry_); + + // Every instrumented component on this node rides the existing metrics tick: + // the storage backend (generic slot-lifecycle series from the decorator, plus + // whatever the medium says about its own internals) and the transfer engine + // (per-engine bytes, plans and in-flight time). Backend-agnostic by + // construction — PoolClient forwards what each component samples and never + // learns which medium or transport produced it. + if (master_client_) master_client_->AddMetricsProvider([this] { PublishComponentMetrics(); }); + // Pack engine_desc for master registration. std::vector engine_desc_bytes; - if (io_engine_) { - msgpack::sbuffer sbuf; - msgpack::pack(sbuf, io_engine_->GetEngineDesc()); - engine_desc_bytes.assign(sbuf.data(), sbuf.data() + sbuf.size()); - } - - // SSD staging buffer (one per process; not part of DRAM exports). Remote SSD - // reads are served out of it; allocated up front when SSD is enabled. - if (config_.ssd.enabled) { - ssd_staging_buffer_ = std::make_unique(config_.ssd_staging_buffer_size); - std::memset(ssd_staging_buffer_.get(), 0, config_.ssd_staging_buffer_size); - if (io_engine_) { - ssd_staging_mem_ = - io_engine_->RegisterMemory(ssd_staging_buffer_.get(), config_.ssd_staging_buffer_size, -1, - mori::io::MemoryLocationType::CPU); - msgpack::sbuffer sbuf; - msgpack::pack(sbuf, ssd_staging_mem_); - ssd_staging_mem_desc_bytes_.assign(sbuf.data(), sbuf.data() + sbuf.size()); - } - } - - if (config_.peer_service_port > 0) { - // Wire the SSD read path: peer_ssd_ + the RDMA-registered SSD staging buffer - // (both null/empty when SSD is disabled, leaving SsdRpcAvailable() false). - peer_service_ = std::make_unique( - peer_alloc_.get(), peer_ssd_.get(), ssd_staging_buffer_.get(), - ssd_staging_buffer_ ? config_.ssd_staging_buffer_size : 0, ssd_staging_mem_desc_bytes_, - config_.ssd_staging_buffer_slots, SsdReadLeaseTtl(), engine_desc_bytes, - master_client_.get(), ssd_copy_pipeline_.get()); + if (peer_directory_ != nullptr) engine_desc_bytes = peer_directory_->PackedLocalEngineDesc(); + + if (config_.peer_service_port > 0 || config_.auto_peer_service_port) { + peer_service_ = std::make_unique(*default_pool_, engine_desc_bytes, + master_client_.get()); if (!peer_service_->Start(config_.peer_service_port)) { MORI_UMBP_ERROR("[PoolClient] PeerService failed to start on port {}", config_.peer_service_port); peer_service_.reset(); - initialized_ = false; + Shutdown(); return false; } + config_.peer_service_port = peer_service_->BoundPort(); } std::string peer_address; - if (config_.peer_service_port > 0) { + if (peer_service_ != nullptr) { std::string host = config_.master_config.node_address; peer_address = host + ":" + std::to_string(config_.peer_service_port); } - // Register the SSD metrics provider before RegisterSelf starts the metrics - // thread (the list is read lock-free afterwards). See PublishSsdMetrics. - if (config_.ssd.enabled) { - master_client_->AddMetricsProvider([this] { PublishSsdMetrics(); }); + // Register conservatively with one instance per tier. Once the Master + // confirms max_allocatable_bytes support, weighted heartbeats advertise the + // full aggregate without making rolling upgrades over-admit large values. + std::map tier_caps; + for (auto* backend : registry_.All()) { + if (registry_.Get(backend->Tier()) != backend) continue; + auto cap = backend->Capacity(); + if (cap.total_bytes == 0) continue; + tier_caps[backend->Tier()] = cap; } - - // Master register. In the new design master holds no DRAM-side - // metadata; only membership + capacity-snapshot (SSD capacity rides in - // tier_capacities as TierType::SSD). - auto status = - master_client_->RegisterSelf(config_.tier_capacities, peer_address, engine_desc_bytes); - if (!status.ok()) { - MORI_UMBP_ERROR("[PoolClient] RegisterSelf failed: {}", status.error_message()); - initialized_ = false; - return false; + if (master_client_) { + const auto logical_caps = default_pool_ == nullptr + ? std::map{} + : default_pool_->LogicalTierCapacities(); + auto status = + master_client_->RegisterSelf(tier_caps, peer_address, engine_desc_bytes, logical_caps); + if (!status.ok()) { + MORI_UMBP_ERROR("[PoolClient] RegisterSelf failed: {}", status.error_message()); + Shutdown(); + return false; + } + if (config_.master_config.auto_heartbeat) master_client_->StartHeartbeat(); + if (weighted_placement && !master_client_->SupportsMaxAllocatableCapacity()) { + MORI_UMBP_WARN( + "[PoolClient] Master lacks max_allocatable_bytes support; weighted placement remains " + "enabled but capacity advertisement is limited to the first instance per tier"); + } } - if (config_.master_config.auto_heartbeat) master_client_->StartHeartbeat(); - - // Start the async re-cache worker only when the feature is on and this node has - // an exportable local DRAM tier to install into. - if (config_.cache_remote_fetches && peer_alloc_) { + // Start the async re-cache worker only when something can feed it and this + // node has an exportable local medium to install into. Two independent + // producers share it: cache_remote_fetches (whole-object BatchGet) and + // ranged_locality_prefetch (remote ranged reads). + if ((config_.cache_remote_fetches || config_.ranged_locality_prefetch) && + registry_.Get(medium_) != nullptr) { { std::lock_guard lk(recache_mutex_); recache_stop_ = false; @@ -547,6 +1065,23 @@ bool PoolClient::Init() { recache_worker_ = std::thread([this] { ReCacheWorkerLoop(); }); } + if (!config_.workload_trace_path.empty()) { + try { + benchmark::WorkloadTraceRecorderOptions options; + options.path = config_.workload_trace_path; + options.client_id = config_.workload_trace_client_id; + options.seed = config_.workload_trace_seed; + options.node_id = config_.master_config.node_id; + options.backend_policy = config_.policy_config_path; + workload_recorder_ = std::make_unique(std::move(options)); + } catch (const std::exception& exception) { + MORI_UMBP_ERROR("[PoolClient] failed to open workload trace '{}': {}", + config_.workload_trace_path, exception.what()); + Shutdown(); + return false; + } + } + MORI_UMBP_INFO("[PoolClient] Initialized node_id='{}'", config_.master_config.node_id); return true; } @@ -556,21 +1091,28 @@ void PoolClient::Shutdown() { initialized_ = false; // Stop the async re-cache worker first: it calls ExecuteLocalPut (which uses - // peer_alloc_ + master_client_), so it must be joined before those are torn + // the registry + master_client_), so it must be joined before those are torn // down below. { std::lock_guard lk(recache_mutex_); recache_stop_ = true; recache_queue_.clear(); + prefetch_inflight_.clear(); } recache_cv_.notify_all(); if (recache_worker_.joinable()) recache_worker_.join(); + if (workload_recorder_ != nullptr) { + try { + workload_recorder_->Close(); + } catch (const std::exception& exception) { + MORI_UMBP_ERROR("[PoolClient] workload trace is incomplete: {}", exception.what()); + } + workload_recorder_.reset(); + } if (master_client_) { master_client_->StopHeartbeat(); - // Stop the periodic metrics thread before tearing down the SSD components - // its provider reads (peer_service_ / ssd_copy_pipeline_ / peer_ssd_), so no - // provider callback runs against freed state. Idempotent with ~MasterClient. + // Idempotent with ~MasterClient. master_client_->StopMetricsReporting(); auto status = master_client_->UnregisterSelf(); if (!status.ok()) { @@ -585,40 +1127,29 @@ void PoolClient::Shutdown() { peer_service_.reset(); - // Stop the copy pipeline AFTER the peer service (RPC commit path) is gone so - // no commit can enqueue into a stopping pipeline. Stop() drops queued tasks - // and waits for the in-flight copy to finish + release its DRAM pin before - // we tear down peer_alloc_ / peer_ssd_ below. - if (ssd_copy_pipeline_) { - ssd_copy_pipeline_->Stop(); - ssd_copy_pipeline_.reset(); - } + // hipHostUnregister the gather regions before their backing memory goes + // away: the medium's buffers die with the registry below, and the scratch + // arena is caller-owned and freed after Shutdown returns. + if (hbm_engine_ != nullptr) hbm_engine_->ClearHostGatherRegions(); - if (peer_alloc_) { - peer_alloc_->StopReaper(); - peer_alloc_.reset(); + // Every backend deregisters its memory through transfer_engine_ inside its + // own destructor — this MUST run before transfer_engine_ is torn down below. + // MasterClient borrows the registry, so unbind it first. + if (master_client_) { + master_client_->SetPeerPool(nullptr); + master_client_->SetBackendRegistry(nullptr); } + default_pool_.reset(); + registry_ = BackendRegistry{}; - // Heartbeat thread is already stopped above, so MasterClient no longer reads - // ssd_manager_; safe to drop the SSD tier owner here. - peer_ssd_.reset(); - - if (io_engine_) { - { - std::lock_guard lock(registered_mem_mutex_); - for (auto& reg : registered_regions_) io_engine_->DeregisterMemory(reg.mem_desc); - registered_regions_.clear(); - } - if (staging_buffer_) io_engine_->DeregisterMemory(staging_mem_); - if (ssd_staging_buffer_) { - io_engine_->DeregisterMemory(ssd_staging_mem_); - ssd_staging_buffer_.reset(); - } - for (auto& mem : export_dram_mems_) io_engine_->DeregisterMemory(mem); - export_dram_mems_.clear(); - io_engine_.reset(); - staging_buffer_.reset(); + if (transfer_engine_) { + std::unique_lock lock(registered_mem_mutex_); + for (auto& reg : registered_regions_) transfer_engine_->Deregister(reg.ref); + registered_regions_.clear(); } + peer_directory_ = nullptr; + hbm_engine_ = nullptr; + transfer_engine_.reset(); master_client_.reset(); } @@ -628,212 +1159,758 @@ bool PoolClient::Clear() { // clear and no master to notify. Treat as success so callers in // shutdown / teardown paths do not surface a spurious error. if (!initialized_.load()) return true; - // Quiesce the copy pipeline first: drop queued tasks and wait for any - // in-flight copy to finish + release its pin. Otherwise a copy completing - // after the clear below would re-record an SSD location and emit ADD SSD - // for a key we just collapsed. - if (ssd_copy_pipeline_) ssd_copy_pipeline_->Quiesce(); - if (peer_alloc_) peer_alloc_->ClearLocal(); - if (peer_ssd_) peer_ssd_->ClearLocal(); + // Clear physical state and the default Pool's logical placement index before + // the full-sync empty snapshot goes out. + if (default_pool_ != nullptr) default_pool_->ClearLocal(); bool ok = true; if (master_client_) { ok = master_client_->ClearFullSync(); if (!ok) MORI_UMBP_WARN("[PoolClient] Clear full-sync heartbeat failed"); } - - // Always resume — even if the full-sync RPC failed — so the pipeline is - // never left permanently paused (the heartbeat loop retries convergence). - if (ssd_copy_pipeline_) ssd_copy_pipeline_->Resume(); return ok; } bool PoolClient::IsInitialized() const { return initialized_; } MasterClient& PoolClient::Master() { return *master_client_; } -PeerDramAllocator* PoolClient::DramAllocator() { return peer_alloc_.get(); } + +void PoolClient::RouteAllPutsLocally(size_t count, + std::vector>* routes) const { + RoutePutResult local; + local.node_id = config_.master_config.node_id; + local.tier = medium_; + routes->assign(count, local); +} + +void PoolClient::CountMetric(std::string name, std::string help, MasterClient::Labels labels, + double delta) { + if (master_client_ == nullptr) return; + master_client_->AddCounter(std::move(name), std::move(help), std::move(labels), delta); +} +BackendRegistry& PoolClient::Backends() { return registry_; } +std::map PoolClient::TierReadHits() const { + return default_pool_ == nullptr ? std::map{} + : default_pool_->TierReadHits(); +} + +std::string PoolClient::LogicalTierForBackend(uint32_t backend_id) const { + return default_pool_ == nullptr ? std::string{} + : default_pool_->LogicalTierForBackend(backend_id); +} + +TierTransitionMetrics PoolClient::TransitionMetrics() const { + return default_pool_ == nullptr ? TierTransitionMetrics{} : default_pool_->TransitionMetrics(); +} // --------------------------------------------------------------------------- // Memory registration // --------------------------------------------------------------------------- -bool PoolClient::RegisterMemory(void* ptr, size_t size) { - if (!io_engine_) { - MORI_UMBP_ERROR("[PoolClient] RegisterMemory: IOEngine not available"); +bool PoolClient::RegisterMemory(void* ptr, size_t size, mori::io::MemoryLocationType loc, + int device, MemoryRegistration mode) { + if (!transfer_engine_ && mode == MemoryRegistration::kPinned) { + MORI_UMBP_ERROR("[PoolClient] RegisterMemory: transfer engine not available"); return false; } if (ptr == nullptr || size == 0) { MORI_UMBP_ERROR("[PoolClient] RegisterMemory: invalid args ptr={}, size={}", ptr, size); return false; } - std::lock_guard lock(registered_mem_mutex_); - for (auto& reg : registered_regions_) { - if (reg.base == ptr) return true; + std::unique_lock lock(registered_mem_mutex_); + const auto at = std::lower_bound( + registered_regions_.begin(), registered_regions_.end(), ptr, + [](const RegisteredRegion& r, const void* p) { return std::less{}(r.base, p); }); + if (at != registered_regions_.end() && at->base == ptr) { + const RegisteredRegion& reg = *at; + // Re-registering the same base with a larger size is not idempotent: the + // existing registration covers fewer bytes, so FindRegisteredMemory would + // start rejecting the tail and silently fall back to unregistered bytes. + if (size <= reg.size) return true; + MORI_UMBP_ERROR( + "[PoolClient] RegisterMemory: ptr={} already registered with smaller size {}<{}", ptr, + reg.size, size); + return false; + } + + if (mode == MemoryRegistration::kLocalCopyOnly) { + // Nothing to pin and nothing that can throw: the ref carries exactly what + // local engine selection reads off it, and every engine's Deregister + // no-ops on a ref without a descriptor. + TransferRef ref = TransferRef::HostBytes(ptr, size, loc, device); + registered_regions_.insert(at, RegisteredRegion{ptr, size, std::move(ref)}); + return true; + } + + // The engine may throw (mori-io pinning can fail on a region the NIC cannot + // map). A throw here would propagate out through the pybind boundary; the + // documented contract is a bool. + try { + TransferRef ref = transfer_engine_->RegisterMemory(ptr, size, loc, device); + // Validate what came back rather than trusting it. A ref whose location or + // device disagrees with what the caller allocated would send every later + // transfer to the wrong engine, and one that covers fewer bytes than asked + // would be used for the full range. + if (!ref.Valid() || ref.host_ptr != ptr || ref.size < size || ref.loc != loc || + (loc == mori::io::MemoryLocationType::GPU && device >= 0 && ref.device != device)) { + MORI_UMBP_ERROR( + "[PoolClient] RegisterMemory: engine returned an inconsistent ref for ptr={}, size={}, " + "loc={}, device={}", + ptr, size, static_cast(loc), device); + transfer_engine_->Deregister(ref); + return false; + } + // Inserted in place rather than appended: lookups binary-search this. + registered_regions_.insert(at, RegisteredRegion{ptr, size, std::move(ref)}); + } catch (const std::exception& error) { + MORI_UMBP_ERROR("[PoolClient] RegisterMemory failed for ptr={}, size={}: {}", ptr, size, + error.what()); + return false; + } catch (...) { + MORI_UMBP_ERROR("[PoolClient] RegisterMemory failed for ptr={}, size={}: unknown error", ptr, + size); + return false; } - auto mem_desc = io_engine_->RegisterMemory(ptr, size, -1, mori::io::MemoryLocationType::CPU); - registered_regions_.push_back({ptr, size, mem_desc}); return true; } void PoolClient::DeregisterMemory(void* ptr) { if (ptr == nullptr) return; - std::lock_guard lock(registered_mem_mutex_); - auto it = std::find_if(registered_regions_.begin(), registered_regions_.end(), - [ptr](const RegisteredRegion& r) { return r.base == ptr; }); - if (it != registered_regions_.end()) { - if (io_engine_) io_engine_->DeregisterMemory(it->mem_desc); + std::unique_lock lock(registered_mem_mutex_); + auto it = std::lower_bound( + registered_regions_.begin(), registered_regions_.end(), ptr, + [](const RegisteredRegion& r, const void* p) { return std::less{}(r.base, p); }); + if (it != registered_regions_.end() && it->base == ptr) { + if (transfer_engine_) transfer_engine_->Deregister(it->ref); registered_regions_.erase(it); } } -std::optional> PoolClient::FindRegisteredMemory( - const void* ptr, size_t size) { - auto addr = reinterpret_cast(ptr); - std::lock_guard lock(registered_mem_mutex_); - for (auto& reg : registered_regions_) { - auto base = reinterpret_cast(reg.base); - if (addr >= base && size <= reg.size && (addr - base) <= reg.size - size) { - return std::pair{reg.mem_desc, static_cast(addr - base)}; - } - } - return std::nullopt; +const PoolClient::RegisteredRegion* PoolClient::FindRegisteredRegionLocked(const void* ptr, + size_t size) const { + // Regions never overlap, so the only candidate is the last one whose base is + // <= ptr. upper_bound then step back finds it in O(log n). + auto at = std::upper_bound( + registered_regions_.begin(), registered_regions_.end(), ptr, + [](const void* p, const RegisteredRegion& r) { return std::less{}(p, r.base); }); + if (at == registered_regions_.begin()) return nullptr; + --at; + const auto addr = reinterpret_cast(ptr); + const auto base = reinterpret_cast(at->base); + if (size > at->size || (addr - base) > at->size - size) return nullptr; + return &*at; } -// --------------------------------------------------------------------------- -// Self-target fast paths -// --------------------------------------------------------------------------- +std::optional> PoolClient::FindRegisteredMemory(const void* ptr, + size_t size) const { + std::shared_lock lock(registered_mem_mutex_); + const RegisteredRegion* reg = FindRegisteredRegionLocked(ptr, size); + if (reg == nullptr) return std::nullopt; + const auto offset = reinterpret_cast(ptr) - reinterpret_cast(reg->base); + return std::pair{reg->ref, static_cast(offset)}; +} -bool PoolClient::LocalPutPages(const std::vector& pages, uint64_t page_size, - const void* src, size_t size) { - const char* src_bytes = static_cast(src); - for (size_t i = 0; i < pages.size(); ++i) { - const auto& p = pages[i]; - if (p.buffer_index >= config_.dram_buffers.size()) { - MORI_UMBP_ERROR("[PoolClient] local Put: invalid buffer_index {}", p.buffer_index); - return false; - } - auto& dram = config_.dram_buffers[p.buffer_index]; - const uint64_t off = static_cast(p.page_index) * page_size; - if (!dram.buffer || page_size > dram.size || off > dram.size - page_size) { - MORI_UMBP_ERROR("[PoolClient] local Put: OOB buf={} off={}", p.buffer_index, off); - return false; - } - const uint64_t bytes = LogicalPageBytes(i, pages.size(), page_size, size); - LocalCopyBlock(static_cast(dram.buffer) + off, src_bytes + i * page_size, bytes); - } - return true; +std::pair PoolClient::UserBufferRef(void* ptr, size_t size) const { + auto reg = FindRegisteredMemory(ptr, size); + if (reg.has_value()) return {reg->first, reg->second}; + return {ClassifiedUserBytes(ptr, size), 0}; } -bool PoolClient::LocalGetPages(const std::vector& pages, uint64_t page_size, - void* dst, size_t size) { - char* dst_bytes = static_cast(dst); +// --------------------------------------------------------------------------- +// Self-target paths +// +// Not a "fast path" any more, and that is the headline of Phase 6: a local +// access is just a transfer whose endpoints are both local, planned by the +// same engine as everything else. What used to make it special — a raw base +// pointer per DRAM buffer, held by PoolClient, memcpy'd through directly — is +// what made the local path host-DRAM-only and what forced PoolClient to name a +// concrete backend type to obtain those pointers. +// --------------------------------------------------------------------------- + +// Build one TransferItem per page between a caller buffer and a backend's own +// buffers. `to_backend` is Put (user -> pages); false is Get (pages -> user). +// Returns false when the backend publishes no endpoint for a referenced buffer, +// which means this medium cannot serve the access in-process. +bool PoolClient::BuildLocalPageTransfers(MediumBackend* backend, + const std::vector& pages, uint64_t page_size, + void* user, size_t size, bool to_backend, + std::vector* items) { + if (pages.empty() || page_size == 0) return false; + // Classified, not assumed host: this is the self-target path, where a device + // user buffer paired with a host DRAM page must reach HbmCopyEngine rather + // than LocalCopyEngine. One classification per batch, not per page. + const TransferRef user_ref = ClassifiedUserBytes(user, size); + items->reserve(pages.size()); for (size_t i = 0; i < pages.size(); ++i) { - const auto& p = pages[i]; - if (p.buffer_index >= config_.dram_buffers.size()) { - MORI_UMBP_ERROR("[PoolClient] local Get: invalid buffer_index {}", p.buffer_index); + TransferRef buf = backend->BufferRef(pages[i].buffer_index); + if (!buf.HasHostPtr()) { + MORI_UMBP_WARN("[PoolClient] local transfer: tier={} publishes no endpoint for buffer {}", + static_cast(backend->Tier()), pages[i].buffer_index); return false; } - auto& dram = config_.dram_buffers[p.buffer_index]; - const uint64_t off = static_cast(p.page_index) * page_size; - if (!dram.buffer || page_size > dram.size || off > dram.size - page_size) { - MORI_UMBP_ERROR("[PoolClient] local Get: OOB buf={} off={}", p.buffer_index, off); - return false; + TransferItem item; + item.size = LogicalPageBytes(i, pages.size(), page_size, size); + item.tag = i; + if (to_backend) { + item.src = user_ref; + item.src_offset = i * page_size; + item.dst = std::move(buf); + item.dst_offset = PageOffset(pages[i], page_size); + } else { + item.src = std::move(buf); + item.src_offset = PageOffset(pages[i], page_size); + item.dst = user_ref; + item.dst_offset = i * page_size; } - const uint64_t bytes = LogicalPageBytes(i, pages.size(), page_size, size); - LocalCopyBlock(dst_bytes + i * page_size, static_cast(dram.buffer) + off, bytes); + items->push_back(std::move(item)); } return true; } PoolClient::PutAttemptOutcome PoolClient::ExecuteLocalPut(const std::string& key, const void* src, size_t size, TierType tier, - bool enqueue_ssd_copy) { - if (!peer_alloc_) { - MORI_UMBP_ERROR("[PoolClient] Local Put requested but peer allocator unavailable"); + const std::string& logical_tier) { + if (default_pool_ == nullptr || registry_.Empty()) { + MORI_UMBP_ERROR("[PoolClient] Local Put requested but no default pool is initialized"); return PutAttemptOutcome::kFatal; } - auto alloc_res = peer_alloc_->Allocate(key, size, tier); + PoolPlacementRequest request; + request.key = key; + request.size = size; + request.tier = tier; + request.logical_tier = logical_tier; + auto pool_alloc = default_pool_->BatchAllocate({request}).front(); + auto* backend = registry_.Get(pool_alloc.backend_id); + // A medium that publishes no buffer endpoints cannot be reached in-process at + // all; route the key elsewhere rather than allocating a slot we cannot fill. + if (backend == nullptr || backend->BufferCount() == 0) { + if (pool_alloc.allocation.outcome == AllocateOutcome::kSuccessAllocated) { + default_pool_->BatchAbort( + {PoolSlotRef{pool_alloc.backend_id, pool_alloc.allocation.slot_id}}); + } + MORI_UMBP_WARN("[PoolClient] Local Put: tier={} has no in-process-addressable backend", + static_cast(tier)); + return PutAttemptOutcome::kRetry; + } + auto& alloc_res = pool_alloc.allocation; switch (alloc_res.outcome) { - case PeerDramAllocator::Outcome::kSuccessAlreadyExists: + case AllocateOutcome::kSuccessAlreadyExists: return PutAttemptOutcome::kSuccessAlreadyExists; - case PeerDramAllocator::Outcome::kFailed: - case PeerDramAllocator::Outcome::kFailedNoSpace: - // Peer allocator already logged the specific reason. + case AllocateOutcome::kFailed: + case AllocateOutcome::kFailedNoSpace: + // Backend already logged the specific reason. return PutAttemptOutcome::kRetry; - case PeerDramAllocator::Outcome::kSuccessAllocated: + case AllocateOutcome::kSuccessAllocated: break; } - const auto& pending = *alloc_res.slot; - if (!LocalPutPages(pending.pages, peer_alloc_->PageSize(), src, size)) { - peer_alloc_->Abort(pending.slot_id); + std::vector items; + if (!BuildLocalPageTransfers(backend, alloc_res.pages, alloc_res.page_size, + const_cast(src), size, /*to_backend=*/true, &items) || + !transfer_engine_->Transfer(items, /*failed_tags=*/nullptr)) { + default_pool_->BatchAbort({PoolSlotRef{pool_alloc.backend_id, alloc_res.slot_id}}); return PutAttemptOutcome::kFatal; } - uint64_t committed_bytes = 0; - if (!peer_alloc_->Commit(pending.slot_id, key, committed_bytes)) { - peer_alloc_->Abort(pending.slot_id); + PoolCommitRequest commit; + commit.slot = PoolSlotRef{pool_alloc.backend_id, alloc_res.slot_id}; + commit.key = key; + if (!default_pool_->BatchCommit({commit}).front().commit.success) { + default_pool_->BatchAbort({commit.slot}); return PutAttemptOutcome::kFatal; } - // Owner-side commit succeeded: best-effort async copy to local SSD. - if (enqueue_ssd_copy && ssd_copy_pipeline_) { - ssd_copy_pipeline_->Enqueue(SsdCopyTask{key, tier, size}); - } - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL_HELP, - {{"traffic", "local"}}, static_cast(size)); - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL_HELP, - {{"traffic", "local"}}, static_cast(size)); + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + static_cast(size)); + CountMetric(MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + static_cast(size)); return PutAttemptOutcome::kSuccess; } PoolClient::GetAttemptOutcome PoolClient::ExecuteLocalGet(const std::string& key, void* dst, size_t size) { - if (!peer_alloc_) { - MORI_UMBP_ERROR("[PoolClient] Local Get requested but peer allocator unavailable"); - return GetAttemptOutcome::kFatal; - } - auto resolved = peer_alloc_->Resolve(key); - if (!resolved.found) { - return GetAttemptOutcome::kRetry; - } - if (!LocalGetPages(resolved.pages, peer_alloc_->PageSize(), dst, size)) { + if (default_pool_ == nullptr || registry_.Empty()) { + MORI_UMBP_ERROR("[PoolClient] Local Get requested but no default pool is initialized"); return GetAttemptOutcome::kFatal; } - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, - {{"traffic", "local"}}, static_cast(size)); - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, - {{"traffic", "local"}}, static_cast(size)); + auto pool_resolved = + ResolveLocalBatchWithBusyRetry(default_pool_.get(), {key}, /*include_descs=*/false, + /*allow_file_refs=*/true) + .front(); + auto* backend = registry_.Get(pool_resolved.backend_id); + bool served = pool_resolved.resolved.found && backend != nullptr; + if (served) { + auto& resolved = pool_resolved.resolved; + // Same guard the remote path applies after BatchResolveKeys: a stored size + // that disagrees with the requested one is a different object, and copying + // `size` bytes out of a slot sized for something else would read past it. + if (resolved.size != size) { + MORI_UMBP_WARN("[PoolClient] local Get: size mismatch for key='{}' (wanted {}, got {})", key, + size, resolved.size); + return GetAttemptOutcome::kRetry; + } + std::vector items; + if (resolved.file_ref.IsFile()) { + // Zero-copy GDS path: the backend published a file range, not a staged + // host page. Read it straight into the caller's (device) buffer; the + // planner routes the (file, GPU) pair to GdsEngine. + TransferItem item; + item.src = resolved.file_ref; + item.dst = ClassifiedUserBytes(dst, size); + item.size = size; + items.push_back(std::move(item)); + } else if (!BuildLocalPageTransfers(backend, resolved.pages, resolved.page_size, dst, size, + /*to_backend=*/false, &items)) { + // This medium holds the key but cannot be read in-process (no published + // endpoint for its buffers). Route elsewhere rather than reporting a + // miss, which would make the client exclude a node that does hold it. + return GetAttemptOutcome::kRetry; + } + if (!transfer_engine_->Transfer(items, /*failed_tags=*/nullptr)) { + return GetAttemptOutcome::kFatal; + } + served = true; + } + if (!served) return GetAttemptOutcome::kRetry; + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + static_cast(size)); + CountMetric(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + static_cast(size)); return GetAttemptOutcome::kSuccess; } -void PoolClient::MaybeReCacheAfterRemote(const std::string& key, const void* src, size_t size) { - if (!peer_alloc_) return; // no exportable local DRAM tier on this node - // Admission gate (cache_remote_fetches / size==0 / NEVER / SIZE cap): shared - // pure predicate, unit-tested in test_cache_remote_admission.cpp. - if (!ShouldAdmitReCache(config_.cache_remote_fetches, config_.cache_remote_admission, - config_.admission_max_block_bytes, size)) { - MORI_UMBP_DEBUG("[PoolClient] MaybeReCacheAfterRemote: key='{}' size={} not admitted", key, - size); - return; +void PoolClient::ResolveLocalBatch(const std::vector& keys, + const std::vector& candidates, + std::vector* holders, + std::vector* resolutions, + const std::function& dst_is_device) { + holders->assign(keys.size(), nullptr); + if (resolutions != nullptr) resolutions->assign(keys.size(), ResolvedEntry{}); + if (candidates.empty() || default_pool_ == nullptr || registry_.Empty()) return; + + std::vector batch; + batch.reserve(candidates.size()); + for (size_t i : candidates) batch.push_back(keys[i]); + + // Both readers fed from here — ServeLocalGets and BatchGetRanges — can + // consume a file ref, and this resolve never leaves the process, so ask for + // one whenever the caller can say where the bytes land. + const bool want_file_refs = static_cast(dst_is_device); + auto found = ResolveLocalBatchWithBusyRetry(default_pool_.get(), batch, + /*include_descs=*/false, + /*allow_file_refs=*/want_file_refs); + + // A file ref is unreadable into HOST memory, and the medium cannot know the + // destination at resolve time. Re-resolve those keys with file refs off so + // they come back on staged pages -- correct, just not zero-copy. + // + // Deliberately a second pass rather than classifying every destination up + // front: with GDS off nothing publishes a file ref, this loop finds nothing, + // and the ordinary local read pays no extra HIP call at all. + if (want_file_refs) { + std::vector restage; + std::vector restage_at; + for (size_t j = 0; j < candidates.size() && j < found.size(); ++j) { + if (!found[j].resolved.file_ref.IsFile()) continue; + if (dst_is_device(candidates[j])) continue; + restage.push_back(batch[j]); + restage_at.push_back(j); + } + if (!restage.empty()) { + auto staged = ResolveLocalBatchWithBusyRetry(default_pool_.get(), restage, + /*include_descs=*/false, + /*allow_file_refs=*/false); + for (size_t k = 0; k < restage_at.size() && k < staged.size(); ++k) { + found[restage_at[k]] = std::move(staged[k]); + } + } } - // ALWAYS and SIZE both delegate DRAM-capacity enforcement to the peer - // allocator: Allocate returns kFailedNoSpace when the tier is full, which we - // treat as a best-effort miss (the remote read result is unaffected). - - // Prepare the job outside the queue lock: the source buffer is valid for this - // call, but copying a multi-MiB block while holding recache_mutex_ would - // serialize unrelated Get finalizers behind this memcpy. - ReCacheJob job; - job.key = key; - job.bytes = std::unique_ptr(new (std::nothrow) char[size]); - if (!job.bytes) { - MORI_UMBP_DEBUG("[PoolClient] MaybeReCacheAfterRemote: allocation failed for key='{}' size={}", - key, size); - return; + for (size_t j = 0; j < candidates.size() && j < found.size(); ++j) { + if (!found[j].resolved.found) continue; + auto* backend = registry_.Get(found[j].backend_id); + if (backend == nullptr) continue; + (*holders)[candidates[j]] = backend; + if (resolutions != nullptr) { + (*resolutions)[candidates[j]] = std::move(found[j].resolved); + } } - job.size = size; - LocalCopyBlock(job.bytes.get(), src, size); +} + +// --------------------------------------------------------------------------- +// Ranged I/O — object ranges onto tier pages +// --------------------------------------------------------------------------- + +// The file-ref counterpart of BuildLocalRangeTransfers: the medium published a +// range of its own storage instead of staged pages, so each caller range is one +// read at (value start + range offset) -- GdsEngine adds src_offset to the +// ref's file_offset. No page walk, because there are no pages to walk. +// +// The ranges a layer-wise reader asks for are not 4 KiB aligned in general, +// while the file ref's own start and length are. That is not a correctness +// problem: hipFile serves an unaligned range through its compat path, which +// gives up the DMA fastpath for that read but returns the right bytes. +bool PoolClient::BuildLocalFileRangeTransfers(const TransferRef& file_ref, + const std::vector& ranges, size_t tag, + std::vector* items) { + if (!file_ref.IsFile() || ranges.empty()) return false; + items->reserve(items->size() + ranges.size()); + for (const ObjectRange& r : ranges) { + if (r.size == 0) continue; + // Same reason BuildLocalRangeTransfers prefers the registration table: a + // registered ref names the region BASE, so every range landing in one + // caller buffer shares a (src, dst) pair and folds into a single plan. + const auto [dst_ref, dst_base_offset] = UserBufferRef(r.user, r.size); + TransferItem item; + item.src = file_ref; + item.src_offset = r.object_offset; + item.dst = dst_ref; + item.dst_offset = dst_base_offset; + item.size = r.size; + item.tag = tag; + items->push_back(std::move(item)); + } + return true; +} + +bool PoolClient::BuildLocalRangeTransfers(MediumBackend* backend, + const std::vector& pages, + uint64_t page_size, uint64_t stored_size, + const std::vector& ranges, bool to_backend, + size_t tag, std::vector* items, + double* classify_sink) { + if (pages.empty() || page_size == 0) return false; + + // Describe each range's caller buffer once, ahead of the page walk, because a + // range can produce several page fragments and the description belongs to the + // range. Callers already truncate `items` when this returns false, so bailing + // before anything is emitted is equivalent to bailing partway through. + // + // Prefer the registration table. A registered ref already carries loc and + // device, so it answers what ClassifiedUserBytes asks hipPointerGetAttributes + // for -- and it answers it for the whole region, not per range. That matters + // twice over here: + // + // * The HIP call is per RANGE, and a layer-wise reader carries thousands. + // * A registered ref names the REGION base, so every range landing in one + // caller buffer shares a (src, dst) pair and LocalCopyEngine::Plan folds + // them into a single plan. Describing them by their own addresses made + // each range its own plan, with its own vectors, and defeated the + // coalescing that Plan exists to do. + // + // In standalone-process mode the server registers every client region it + // imports, so this is the path a real deployment takes; ClassifiedUserBytes + // stays as the answer for genuinely unregistered memory, where the pointer + // really does have to be interrogated. + std::vector user_refs; + std::vector user_offsets; + { + PhaseTimer classify_timer(classify_sink); + user_refs.reserve(ranges.size()); + user_offsets.reserve(ranges.size()); + // One shared lock for the whole batch, not one per range. + std::shared_lock lock(registered_mem_mutex_); + for (const auto& range : ranges) { + if (range.user == nullptr) return false; + if (const RegisteredRegion* reg = FindRegisteredRegionLocked(range.user, range.size)) { + user_refs.push_back(reg->ref); + user_offsets.push_back(reinterpret_cast(range.user) - + reinterpret_cast(reg->base)); + } else { + user_refs.push_back(ClassifiedUserBytes(range.user, range.size)); + user_offsets.push_back(0); + } + } + } + + for (size_t r = 0; r < ranges.size(); ++r) { + const auto& range = ranges[r]; + const TransferRef& user_ref = user_refs[r]; + const uint64_t user_base = user_offsets[r]; + + // The walk owns every bound: object overflow, page-list overrun, and the + // short last page a range must not run off the end of. + const bool walked = ForEachRangePageFragment( + pages, page_size, stored_size, range.object_offset, range.size, + [&](size_t page_index, uint64_t tier_offset, size_t fragment, size_t copied) { + TransferRef buf = backend->BufferRef(pages[page_index].buffer_index); + if (!buf.HasHostPtr()) { + MORI_UMBP_WARN( + "[PoolClient] ranged transfer: tier={} publishes no endpoint for buffer {}", + static_cast(backend->Tier()), pages[page_index].buffer_index); + return false; + } + TransferItem item; + item.size = fragment; + item.tag = tag; + // `copied` is relative to the range; the ref may describe the whole + // registered region the range sits inside, so bias by where the range + // starts in it. user_base is 0 for an unregistered pointer, whose ref + // describes exactly the range. + const uint64_t user_offset = user_base + copied; + if (to_backend) { + item.src = user_ref; + item.src_offset = user_offset; + item.dst = std::move(buf); + item.dst_offset = tier_offset; + } else { + item.src = std::move(buf); + item.src_offset = tier_offset; + item.dst = user_ref; + item.dst_offset = user_offset; + } + items->push_back(std::move(item)); + return true; + }); + if (!walked) return false; + } + return true; +} + +bool PoolClient::CopyContiguousToRanges(const void* src, size_t object_size, + const std::vector& ranges) { + // The arena is plain host memory this client owns, so it is described + // directly rather than resolved through a backend's published endpoints. + return CopyContiguousToRanges(TransferRef::HostBytes(const_cast(src), object_size), + /*src_base=*/0, object_size, ranges); +} + +bool PoolClient::BuildContiguousToRangesItems(const TransferRef& src, uint64_t src_base, + size_t object_size, + const std::vector& ranges, size_t tag, + std::vector* items) { + const size_t before = items->size(); + for (const auto& range : ranges) { + // All-or-nothing per object: a partially appended object would be + // transferred as a torn read, so roll back to the batch's prior end. + if (range.size == 0 || range.user == nullptr || + IsObjectRangeOverflow(range.object_offset, range.size, object_size)) { + items->resize(before); + return false; + } + TransferItem item; + item.size = range.size; + item.tag = tag; + item.src = src; + item.src_offset = src_base + range.object_offset; + item.dst = ClassifiedUserBytes(range.user, range.size); + item.dst_offset = 0; + items->push_back(std::move(item)); + } + return items->size() > before; +} + +bool PoolClient::CopyContiguousToRanges(const TransferRef& src, uint64_t src_base, + size_t object_size, + const std::vector& ranges) { + std::vector items; + items.reserve(ranges.size()); + if (!BuildContiguousToRangesItems(src, src_base, object_size, ranges, /*tag=*/0, &items)) { + return false; + } + return transfer_engine_->Transfer(items, /*failed_tags=*/nullptr); +} + +bool PoolClient::BuildRangesToContiguousItems(const std::vector& ranges, void* dst, + size_t object_size, size_t tag, + std::vector* items) { + const TransferRef object_ref = TransferRef::HostBytes(dst, object_size); + const size_t before = items->size(); + for (const auto& range : ranges) { + if (range.size == 0 || range.user == nullptr || + IsObjectRangeOverflow(range.object_offset, range.size, object_size)) { + items->resize(before); + return false; + } + TransferItem item; + item.size = range.size; + item.tag = tag; + item.src = ClassifiedUserBytes(range.user, range.size); + item.src_offset = 0; + item.dst = object_ref; + item.dst_offset = range.object_offset; + items->push_back(std::move(item)); + } + return items->size() > before; +} + +bool PoolClient::CopyRangesToContiguous(const std::vector& ranges, void* dst, + size_t object_size) { + std::vector items; + items.reserve(ranges.size()); + if (!BuildRangesToContiguousItems(ranges, dst, object_size, /*tag=*/0, &items)) return false; + return transfer_engine_->Transfer(items, /*failed_tags=*/nullptr); +} + +void PoolClient::ExecuteLocalPutRangesBatch(const std::vector& requests, + std::vector* results, double* committed_bytes, + const RangedPhaseSinks* sinks_or_null) { + static const RangedPhaseSinks kNoSinks; + const RangedPhaseSinks& sinks = sinks_or_null != nullptr ? *sinks_or_null : kNoSinks; + if (requests.empty()) return; + if (default_pool_ == nullptr || registry_.Empty()) { + MORI_UMBP_ERROR("[PoolClient] Local ranged Put requested but no pool is initialized"); + return; + } + + // Keep the optimized one-batch allocation shape, but let PeerPool choose + // each backend so weighted/logical-tier placement and pending state remain + // authoritative. + std::vector allocations; + std::vector backend_of(requests.size(), nullptr); + { + PhaseTimer alloc_timer(sinks.resolve); + std::vector asks; + asks.reserve(requests.size()); + for (const auto& request : requests) { + asks.push_back(PoolPlacementRequest{*request.key, request.object_size, request.tier, + /*backend_name=*/{}, request.logical_tier}); + } + allocations = default_pool_->BatchAllocate(asks); + allocations.resize(requests.size()); + for (size_t r = 0; r < requests.size(); ++r) { + backend_of[r] = registry_.Get(allocations[r].backend_id); + } + } + + struct PendingWrite { + size_t request_index = 0; + PoolSlotRef slot; + }; + std::vector to_abort; + std::vector pending; + std::vector items; + pending.reserve(requests.size()); + { + size_t range_total = 0; + for (const auto& request : requests) range_total += request.ranges.size(); + items.reserve(range_total); + } + + for (size_t r = 0; r < requests.size(); ++r) { + MediumBackend* const backend = backend_of[r]; + const auto& request = requests[r]; + const AllocateResult& alloc_res = allocations[r].allocation; + if (alloc_res.outcome == AllocateOutcome::kSuccessAlreadyExists) { + (*results)[request.result_index] = true; + continue; + } + if (alloc_res.outcome != AllocateOutcome::kSuccessAllocated) continue; + const PoolSlotRef slot{allocations[r].backend_id, alloc_res.slot_id}; + if (backend == nullptr || backend->BufferCount() == 0) { + MORI_UMBP_WARN("[PoolClient] Local ranged Put: selected backend has no local endpoint"); + to_abort.push_back(slot); + continue; + } + + // tag = index into `requests`, so the engine's failed tags map straight + // back to the slot that has to be aborted. + const size_t before = items.size(); + bool built = false; + { + PhaseTimer build_timer(sinks.build); + built = BuildLocalRangeTransfers(backend, alloc_res.pages, alloc_res.page_size, + request.object_size, request.ranges, /*to_backend=*/true, + /*tag=*/r, &items, sinks.classify); + } + if (!built) { + items.resize(before); + to_abort.push_back(slot); + continue; + } + pending.push_back({r, slot}); + } + // Classification is timed inside the build window; subtract it so the phases + // stay disjoint. + if (sinks.build != nullptr && sinks.classify != nullptr) *sinks.build -= *sinks.classify; + if (sinks.items != nullptr) *sinks.items += items.size(); + auto flush_aborts = [this, &to_abort] { + if (to_abort.empty()) return; + default_pool_->BatchAbort(to_abort); + to_abort.clear(); + }; + + if (pending.empty()) { + flush_aborts(); + return; + } + + std::vector failed_tags; + transfer_engine_->Transfer(items, &failed_tags, sinks.steps); + const std::unordered_set failed(failed_tags.begin(), failed_tags.end()); + + // Commit is still strictly after the transfer, and a slot that fails either + // step is still aborted -- only the number of calls changes. + std::vector commits; + std::vector commit_owner; + commits.reserve(pending.size()); + commit_owner.reserve(pending.size()); + for (const auto& write : pending) { + if (failed.count(write.request_index) != 0) { + to_abort.push_back(write.slot); + continue; + } + commits.push_back(PoolCommitRequest{write.slot, *requests[write.request_index].key}); + commit_owner.push_back(write.request_index); + } + + double put_bytes = 0.0; + { + PhaseTimer commit_timer(sinks.commit); + auto outcomes = default_pool_->BatchCommit(commits); + for (size_t j = 0; j < commit_owner.size(); ++j) { + const auto& request = requests[commit_owner[j]]; + if (j >= outcomes.size() || !outcomes[j].commit.success) { + to_abort.push_back(commits[j].slot); + continue; + } + (*results)[request.result_index] = true; + put_bytes += static_cast(request.object_size); + } + } + flush_aborts(); + + if (committed_bytes != nullptr) *committed_bytes += put_bytes; + if (put_bytes > 0.0) { + // One update for the batch rather than two per key; the totals are the + // same. AddCounter locks and materializes its help text every call, which + // is not something to do once per object. + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + put_bytes); + CountMetric(MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + put_bytes); + } +} + +void PoolClient::MaybeReCacheAfterRemote(const std::string& key, const void* src, size_t size) { + auto* local = registry_.Get(medium_); + if (local == nullptr || local->BufferCount() == 0) return; // no exportable local medium here + // Admission gate (cache_remote_fetches / size==0 / NEVER / SIZE cap): shared + // pure predicate, unit-tested in test_cache_remote_admission.cpp. + if (!ShouldAdmitReCache(config_.cache_remote_fetches, config_.cache_remote_admission, + config_.admission_max_block_bytes, size)) { + MORI_UMBP_DEBUG("[PoolClient] MaybeReCacheAfterRemote: key='{}' size={} not admitted", key, + size); + return; + } + // ALWAYS and SIZE both delegate capacity enforcement to the peer allocator: + // Allocate returns kFailedNoSpace when the medium is full, which we treat as + // a best-effort miss (the remote read result is unaffected). + + // Prepare the job outside the queue lock: the source buffer is valid for this + // call, but copying a multi-MiB block while holding recache_mutex_ would + // serialize unrelated Get finalizers behind this memcpy. + ReCacheJob job; + job.key = key; + job.bytes = std::unique_ptr(new (std::nothrow) char[size]); + if (!job.bytes) { + MORI_UMBP_DEBUG("[PoolClient] MaybeReCacheAfterRemote: allocation failed for key='{}' size={}", + key, size); + return; + } + job.size = size; + HostCopyBlock(job.bytes.get(), src, size); // Enqueue for asynchronous install. The actual DRAM Allocate + copy + // Commit→KvEvent::ADD publish is performed by ReCacheWorkerLoop OFF the Get @@ -852,7 +1929,190 @@ void PoolClient::MaybeReCacheAfterRemote(const std::string& key, const void* src recache_cv_.notify_one(); } +// How many whole-object prefetches one worker pass may take at once. Large +// enough that the per-batch resolve RPC is amortised over a realistic page +// batch, small enough that one batch's slot allocations and wire time stay +// bounded when objects are multi-MiB. +constexpr size_t kPrefetchBatchMax = 256; + +bool PoolClient::CanInstallLocally() const { + auto* local = registry_.Get(medium_); + return local != nullptr && local->BufferCount() != 0; +} + +bool PoolClient::LocalityPrefetchAdmits(size_t object_size) const { + if (!config_.ranged_locality_prefetch) return false; + if (!CanInstallLocally()) return false; + // Same admission policy as the non-ranged re-cache, minus its enable flag: + // ranged_locality_prefetch is this path's switch and was checked above, so + // the predicate is asked only about the policy and the block size. + return ShouldAdmitReCache(/*cache_remote_fetches=*/true, config_.cache_remote_admission, + config_.admission_max_block_bytes, object_size); +} + +void PoolClient::MaybeInstallCompleteArenaObject(const std::string& key, const void* arena_slice, + size_t object_size) { + if (!CanInstallLocally()) return; + // Synchronous on purpose, and cheap relative to the alternative: the slice is + // live only until the next sub-batch reuses it, and a local copy of the + // object beats a second RDMA of the same bytes. This is also exactly what + // the pre-ranged code did on every remote ranged read -- the difference is + // that it now only happens when the object really is complete here. + const auto installed = ExecuteLocalPut(key, arena_slice, object_size, medium_); + if (installed != PutAttemptOutcome::kSuccess && + installed != PutAttemptOutcome::kSuccessAlreadyExists) { + NoteRangedInstallFailure(); + } +} + +void PoolClient::MaybePrefetchWholeObject(const std::string& key, size_t object_size, + const RouteGetResult& route) { + if (!LocalityPrefetchAdmits(object_size)) return; + + ReCacheJob job; + job.key = key; + job.size = object_size; + job.route = route; // bytes stays null: the worker pulls them itself + + { + std::lock_guard lk(recache_mutex_); + if (recache_stop_) return; + if (recache_queue_.size() >= recache_queue_max_) { + MORI_UMBP_DEBUG("[PoolClient] MaybePrefetchWholeObject: queue full, dropping key='{}'", key); + return; + } + if (!prefetch_inflight_.insert(key).second) return; // already queued or running + recache_queue_.push_back(std::move(job)); + } + recache_cv_.notify_one(); +} + +void PoolClient::FetchWholeObjectsIntoMedium(std::vector& jobs) { + auto* backend = registry_.Get(medium_); + if (backend == nullptr || backend->BufferCount() == 0 || jobs.empty()) return; + + // Drop anything that landed while it waited in the queue — a concurrent + // request, or the offload path writing it back. One batched resolve, before + // any allocation: the alternative wastes a whole object of wire per key. + std::vector keys; + keys.reserve(jobs.size()); + for (const auto& job : jobs) keys.push_back(job.key); + const auto resolved = backend->BatchResolve(keys, /*include_descs=*/false); + + std::vector wanted; + std::vector requests; + wanted.reserve(jobs.size()); + requests.reserve(jobs.size()); + for (size_t i = 0; i < jobs.size(); ++i) { + if (i < resolved.size() && resolved[i].found) continue; + if (!jobs[i].route.has_value()) continue; + wanted.push_back(i); + requests.push_back(AllocateRequest{jobs[i].key, jobs[i].size}); + } + if (requests.empty()) return; + + auto allocs = backend->BatchAllocate(requests); + + // Peer pages and medium pages are both registered host memory, so the objects + // move peer -> medium as plain RDMA: no arena, no bounce buffer, no memcpy, + // and the reader's destination is never touched. + // + // The remote get path describes its destination as one contiguous span, so + // the slot has to be one too. The page allocator already prefers a + // same-buffer continuous run for exactly this reason, and distributed mode + // stores one key per page, so the common cases are covered; a scattered slot + // is left uncached rather than served by a path that would misplace it. + // + // slot_refs is reserved up front and never grows: the plan holds pointers + // into it. + BatchGetPlan plan; + std::vector slot_refs; + std::vector planned; // index into wanted/allocs + slot_refs.reserve(wanted.size()); + planned.reserve(wanted.size()); + + for (size_t w = 0; w < wanted.size(); ++w) { + auto& alloc = allocs[w]; + if (alloc.outcome != AllocateOutcome::kSuccessAllocated) continue; + const auto& job = jobs[wanted[w]]; + + bool contiguous = !alloc.pages.empty(); + for (size_t p = 1; p < alloc.pages.size() && contiguous; ++p) { + contiguous = alloc.pages[p].buffer_index == alloc.pages.front().buffer_index && + alloc.pages[p].page_index == alloc.pages.front().page_index + p; + } + TransferRef slot_buffer = + contiguous ? backend->BufferRef(alloc.pages.front().buffer_index) : TransferRef{}; + if (!contiguous || !slot_buffer.Valid()) { + MORI_UMBP_DEBUG("[PoolClient] prefetch: slot for key='{}' is not one addressable run", + job.key); + // Aborted here and never revisited: the commit/abort pass below walks + // `planned`, which this key does not enter. + backend->BatchAbort({alloc.slot_id}); + continue; + } + + slot_refs.push_back(std::move(slot_buffer)); + // Named by the backend's own ref rather than a raw pointer: the medium pool + // is RDMA-reachable but was never handed to RegisterMemory, so + // UserBufferRef would not find it and the transfer would fall through to a + // bounce the staging pool may be too small for. + plan.remote_groups[job.route->node_id].push_back( + BatchGetItem{.index = planned.size(), + .key = &job.key, + .dst = nullptr, + .size = job.size, + .dst_ref = &slot_refs.back(), + .dst_ref_offset = PageOffset(alloc.pages.front(), alloc.page_size), + .route = *job.route}); + planned.push_back(w); + } + if (planned.empty()) return; + + std::vector fetched(planned.size(), false); + ExecuteRemoteBatchGetPlan(plan, &fetched, /*recache_remote=*/false); + + std::vector commits; + std::vector aborts; + commits.reserve(planned.size()); + // Everything this pass wanted but could not plan -- no slot, or a slot that + // is not one addressable run -- already failed to become local. + size_t not_installed = wanted.size() - planned.size(); + for (size_t i = 0; i < planned.size(); ++i) { + auto& alloc = allocs[planned[i]]; + if (!fetched[i]) { + aborts.push_back(alloc.slot_id); + ++not_installed; + continue; + } + commits.push_back(CommitRequest{alloc.slot_id, jobs[wanted[planned[i]]].key}); + } + if (!commits.empty()) { + const auto results = backend->BatchCommit(commits); + for (size_t c = 0; c < results.size(); ++c) { + if (!results[c].success) { + aborts.push_back(commits[c].slot_id); + ++not_installed; + } + } + } + if (!aborts.empty()) backend->BatchAbort(aborts); + // Reported in one go rather than per key: this is a background pass, and the + // counter is there to say how much locality is being lost, not to trace it. + NoteRangedInstallFailure(not_installed); +} + void PoolClient::ReCacheWorkerLoop() { + // Nobody is blocked on this thread. Its whole purpose is to make LATER calls + // faster, so it must not borrow cores from the calls happening NOW -- doing + // so cost 63% of TTFL on a remote layer-wise restore at 8 ranks. Set once + // here rather than around each copy: everything this loop does is background + // by construction. + MarkThreadBackgroundCopies(); + + // Reused across iterations so a steady stream of prefetches does not + // reallocate this vector on every batch. + std::vector prefetch_batch; for (;;) { ReCacheJob job; { @@ -861,13 +2121,42 @@ void PoolClient::ReCacheWorkerLoop() { if (recache_stop_ && recache_queue_.empty()) return; job = std::move(recache_queue_.front()); recache_queue_.pop_front(); + // A prefetch is worth batching: its cost per key is a peer resolve RPC + // plus an RDMA round trip, and this is the only thread paying it. Take + // every prefetch already waiting, up to a cap that keeps one batch's + // allocation and wire time bounded. + if (job.bytes == nullptr) { + prefetch_batch.clear(); + prefetch_batch.push_back(std::move(job)); + for (auto it = recache_queue_.begin(); + it != recache_queue_.end() && prefetch_batch.size() < kPrefetchBatchMax;) { + if (it->bytes != nullptr) { + ++it; // an install job; leave it in order + continue; + } + prefetch_batch.push_back(std::move(*it)); + it = recache_queue_.erase(it); + } + } } - // Install into local DRAM. ExecuteLocalPut allocates a slot in dram_buffers, - // copies the bytes, and Commit queues a KvEvent::ADD that reaches the master - // via heartbeat — mirroring the local Put publish path. kSuccessAlreadyExists - // makes this idempotent for a repeat remote read of the same key. - switch (ExecuteLocalPut(job.key, job.bytes.get(), job.size, TierType::DRAM, - /*enqueue_ssd_copy=*/false)) { + if (!prefetch_batch.empty()) { + // Whole-object locality prefetch. The dedup slots are released + // afterwards, so a later miss can try again for whatever did not land. + FetchWholeObjectsIntoMedium(prefetch_batch); + { + std::lock_guard lk(recache_mutex_); + for (const auto& done : prefetch_batch) prefetch_inflight_.erase(done.key); + } + MORI_UMBP_DEBUG("[PoolClient] ReCacheWorker: prefetched {} objects", prefetch_batch.size()); + prefetch_batch.clear(); + continue; + } + // Install into this node's medium. ExecuteLocalPut allocates a slot on that + // backend, copies the bytes, and Commit queues a KvEvent::ADD that reaches + // the master via heartbeat — mirroring the local Put publish path. + // kSuccessAlreadyExists makes this idempotent for a repeat remote read of + // the same key. + switch (ExecuteLocalPut(job.key, job.bytes.get(), job.size, medium_)) { case PutAttemptOutcome::kSuccess: MORI_UMBP_DEBUG("[PoolClient] ReCacheWorker: re-cached key='{}' size={}", job.key, job.size); @@ -882,42 +2171,6 @@ void PoolClient::ReCacheWorkerLoop() { } } -PoolClient::GetAttemptOutcome PoolClient::ExecuteLocalSsdGet(const std::string& key, void* dst, - size_t size) { - // reader == owner: read straight into the user buffer (no staging / RDMA / - // lease — those serve remote readers). kNotFound -> kRetry (miss, not error). - if (!peer_ssd_) { - MORI_UMBP_ERROR("[PoolClient] Local SSD Get requested but PeerSsdManager unavailable"); - return GetAttemptOutcome::kFatal; - } - SsdReadOutcome outcome = peer_ssd_->PrepareRead(key, dst, size); - switch (outcome.status) { - case SsdReadStatus::kOk: - break; - case SsdReadStatus::kNotFound: - return GetAttemptOutcome::kRetry; - case SsdReadStatus::kSizeTooLarge: - case SsdReadStatus::kError: - MORI_UMBP_WARN("[PoolClient] Local SSD Get key='{}' failed (status={}, size={})", key, - static_cast(outcome.status), outcome.size); - return GetAttemptOutcome::kFatal; - } - // Guard against a short read filling only part of the user buffer (mirrors the - // remote path's size check). - if (outcome.size != size) { - MORI_UMBP_WARN("[PoolClient] Local SSD Get key='{}' size mismatch (wanted {}, got {})", key, - size, outcome.size); - return GetAttemptOutcome::kFatal; - } - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, - {{"traffic", "local"}}, static_cast(size)); - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, - {{"traffic", "local"}}, static_cast(size)); - return GetAttemptOutcome::kSuccess; -} - // --------------------------------------------------------------------------- // BatchPut // --------------------------------------------------------------------------- @@ -933,10 +2186,6 @@ bool PoolClient::Put(const std::string& key, const void* src, size_t size) { std::vector PoolClient::BatchPut(const std::vector& keys, const std::vector& srcs, const std::vector& sizes) { - // NOTE: BatchPut only lands data in the DRAM/HBM tier. The SSD tier copy is - // asynchronous and owner-driven: a successful slot Commit (local in - // ExecuteLocalPut, remote on the peer's CommitSlot) enqueues the copy onto the - // SsdCopyPipeline. There is no explicit SSD write on this path. const auto call_start = std::chrono::steady_clock::now(); if (keys.size() != srcs.size() || keys.size() != sizes.size()) { MORI_UMBP_ERROR("[PoolClient] BatchPut: vector length mismatch"); @@ -954,12 +2203,16 @@ std::vector PoolClient::BatchPut(const std::vector& keys, for (size_t i = 0; i < sizes.size(); ++i) block_sizes[i] = static_cast(sizes[i]); std::vector> routes; std::unordered_set excludes; - auto status = master_client_->BatchRoutePut(keys, block_sizes, excludes, &routes); - if (!status.ok()) { - MORI_UMBP_ERROR("[PoolClient] BatchPut: BatchRoutePut failed: {}", status.error_message()); - return std::vector(keys.size(), false); + if (!HasMaster()) { + RouteAllPutsLocally(keys.size(), &routes); + } else { + auto status = master_client_->BatchRoutePut(keys, block_sizes, excludes, &routes); + if (!status.ok()) { + MORI_UMBP_ERROR("[PoolClient] BatchPut: BatchRoutePut failed: {}", status.error_message()); + return std::vector(keys.size(), false); + } + if (routes.size() < keys.size()) routes.resize(keys.size()); } - if (routes.size() < keys.size()) routes.resize(keys.size()); BatchPutPlan plan = PartitionBatchPutTargets(keys, srcs, sizes, routes, &outcomes); ExecuteBatchPutPlan(plan, &outcomes); @@ -969,10 +2222,10 @@ std::vector PoolClient::BatchPut(const std::vector& keys, std::chrono::duration_cast>(call_end - call_start).count(); if (seconds > 0.0) { auto split = ComputeBatchBandwidthBytes(outcomes, sizes, routes, config_.master_config.node_id); - ObserveBatchBandwidth(*master_client_, split.local, seconds, + ObserveBatchBandwidth(master_client_.get(), split.local, seconds, MORI_UMBP_METRIC_CLIENT_BATCH_PUT_BANDWIDTH, MORI_UMBP_METRIC_CLIENT_BATCH_PUT_BANDWIDTH_HELP, "local"); - ObserveBatchBandwidth(*master_client_, split.remote, seconds, + ObserveBatchBandwidth(master_client_.get(), split.remote, seconds, MORI_UMBP_METRIC_CLIENT_BATCH_PUT_BANDWIDTH, MORI_UMBP_METRIC_CLIENT_BATCH_PUT_BANDWIDTH_HELP, "remote"); } @@ -981,6 +2234,15 @@ std::vector PoolClient::BatchPut(const std::vector& keys, for (size_t i = 0; i < outcomes.size(); ++i) { results[i] = (outcomes[i] != PutEntryOutcome::kFailed); } + if (workload_recorder_ != nullptr) { + std::vector recorded(outcomes.size()); + for (size_t i = 0; i < outcomes.size(); ++i) { + recorded[i] = outcomes[i] != PutEntryOutcome::kFailed + ? benchmark::WorkloadTraceOutcome::kSuccess + : benchmark::WorkloadTraceOutcome::kFailure; + } + workload_recorder_->RecordBatchPut(keys, sizes, recorded); + } return results; } @@ -1010,7 +2272,10 @@ PoolClient::BatchPutPlan PoolClient::PartitionBatchPutTargets( .index = i, .key = &keys[i], .src = srcs[i], .size = sizes[i], .route = route}); continue; } - if (route.tier != TierType::DRAM && route.tier != TierType::HBM) continue; + // No tier filter: every medium a peer advertises publishes registered + // pages (SSD's are its staging arena — see ssd_backend.h), so the remote + // put path is the same for all of them. The old DRAM/HBM allowlist here + // silently dropped puts master routed to a peer's SSD. plan.remote_groups[route.node_id].push_back(BatchPutItem{ .index = i, .key = &keys[i], .src = srcs[i], .size = sizes[i], .route = route}); } @@ -1019,28 +2284,111 @@ PoolClient::BatchPutPlan PoolClient::PartitionBatchPutTargets( void PoolClient::ExecuteBatchPutPlan(const BatchPutPlan& plan, std::vector* results) { - // Deferred local puts, parallel: per-key memcpy is lock-free (the allocator - // serializes Allocate/Commit); results is not vector-bit-packed, so - // workers write distinct indices directly. AddCounter / timing stay here. + // Deferred local puts. Allocate and commit are one pool call each for the + // whole batch: PeerPool takes its operation lock, consults placement and + // pending state and builds its per-round containers once per call, so + // driving it a key at a time made that fixed cost -- not the copy -- the + // dominant term of a local put. Only the copy is inherently per-key, and + // that is what stays parallel: distinct slots, distinct result indices, so + // workers still write directly. Same shape as ExecuteLocalPutRangesBatch, + // which is why the ranged path never paid this. auto run_local_put = [&]() { const auto& local = plan.local_items; if (local.empty()) return; + if (default_pool_ == nullptr || registry_.Empty()) { + MORI_UMBP_ERROR("[PoolClient] Local Put requested but no default pool is initialized"); + return; + } const int nthr = LocalCopyThreads("UMBP_DRAM_WRITE_THREADS"); const auto t0 = std::chrono::steady_clock::now(); - LocalParallelFor(local.size(), nthr, [&](size_t k) { + + std::vector asks; + asks.reserve(local.size()); + for (const auto& item : local) { + asks.push_back(PoolPlacementRequest{*item.key, item.size, item.route.tier, + /*backend_name=*/{}, item.route.logical_tier}); + } + auto allocations = default_pool_->BatchAllocate(asks); + allocations.resize(local.size()); + + // `staged` carries the keys that reached the copy phase; every slot that + // does not survive placement, copy or commit lands in `to_abort`. + std::vector to_abort; + std::vector staged; + std::vector backend_of(local.size(), nullptr); + staged.reserve(local.size()); + for (size_t k = 0; k < local.size(); ++k) { + const auto& alloc = allocations[k].allocation; + if (alloc.outcome == AllocateOutcome::kSuccessAlreadyExists) { + (*results)[local[k].index] = PutEntryOutcome::kAlreadyExists; + continue; + } + if (alloc.outcome != AllocateOutcome::kSuccessAllocated) continue; + auto* backend = registry_.Get(allocations[k].backend_id); + // A medium that publishes no buffer endpoints cannot be reached in-process + // at all; drop the reservation rather than fill a slot we cannot write. + if (backend == nullptr || backend->BufferCount() == 0) { + MORI_UMBP_WARN("[PoolClient] Local Put: tier={} has no in-process-addressable backend", + static_cast(local[k].route.tier)); + to_abort.push_back(PoolSlotRef{allocations[k].backend_id, alloc.slot_id}); + continue; + } + backend_of[k] = backend; + staged.push_back(k); + } + + // vector, not vector: workers write distinct indices. + std::vector copied(local.size(), 0); + ParallelFor(staged.size(), nthr, [&](size_t s) { + const size_t k = staged[s]; const auto& item = local[k]; - switch (ExecuteLocalPut(*item.key, item.src, item.size, item.route.tier)) { - case PutAttemptOutcome::kSuccess: - (*results)[item.index] = PutEntryOutcome::kSucceeded; - break; - case PutAttemptOutcome::kSuccessAlreadyExists: - (*results)[item.index] = PutEntryOutcome::kAlreadyExists; - break; - case PutAttemptOutcome::kRetry: - case PutAttemptOutcome::kFatal: - break; + const auto& alloc = allocations[k].allocation; + std::vector items; + if (BuildLocalPageTransfers(backend_of[k], alloc.pages, alloc.page_size, + const_cast(item.src), item.size, /*to_backend=*/true, + &items) && + transfer_engine_->Transfer(items, /*failed_tags=*/nullptr)) { + copied[k] = 1; } }); + + std::vector commits; + std::vector commit_owner; + commits.reserve(staged.size()); + commit_owner.reserve(staged.size()); + for (size_t k : staged) { + const PoolSlotRef slot{allocations[k].backend_id, allocations[k].allocation.slot_id}; + if (copied[k] == 0) { + to_abort.push_back(slot); + continue; + } + commits.push_back(PoolCommitRequest{slot, *local[k].key}); + commit_owner.push_back(k); + } + + double put_bytes = 0.0; + auto outcomes = default_pool_->BatchCommit(commits); + for (size_t j = 0; j < commit_owner.size(); ++j) { + if (j >= outcomes.size() || !outcomes[j].commit.success) { + to_abort.push_back(commits[j].slot); + continue; + } + const size_t k = commit_owner[j]; + (*results)[local[k].index] = PutEntryOutcome::kSucceeded; + put_bytes += static_cast(local[k].size); + } + if (!to_abort.empty()) default_pool_->BatchAbort(to_abort); + + if (put_bytes > 0.0) { + // One update for the batch rather than two per key; the totals are the + // same, and AddCounter locks and materializes its help text every call. + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + put_bytes); + CountMetric(MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_INBOUND_PUT_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + put_bytes); + } if (std::getenv("UMBP_LOCAL_COPY_TIMING")) { double sec = std::chrono::duration_cast>( std::chrono::steady_clock::now() - t0) @@ -1053,50 +2401,37 @@ void PoolClient::ExecuteBatchPutPlan(const BatchPutPlan& plan, } }; - const auto& remote_groups = plan.remote_groups; - - // All-ZC or all-staging (upper-layer invariant): probe one item's src. - const bool is_zc = - !remote_groups.empty() && FindRegisteredMemory(remote_groups.begin()->second.front().src, - remote_groups.begin()->second.front().size) - .has_value(); - - if (!is_zc) { - // Staging / no remote: each peer submits then IMMEDIATELY waits. The - // in-flight holds staging_mutex_ submit..wait, so a submit-all over multiple - // staging peers would deadlock on the single lock. - run_local_put(); - for (const auto& [node_id, items] : remote_groups) { - if (auto f = SubmitRemoteBatchPut(items, results, /*permit_staging=*/true)) { - WaitRemoteBatchPut(*f, results); - } - } - return; - } - - // Zero-copy: submit every peer (not waited) to overlap the wire across peers, - // run local puts in that window, then wait all + commit. On early exit - // ~RemoteDramPutInFlight drains statuses; the wait does mapping + commit/abort. - std::vector> inflights; - inflights.reserve(remote_groups.size()); - for (const auto& [node_id, items] : remote_groups) { - if (auto f = SubmitRemoteBatchPut(items, results, /*permit_staging=*/false)) { - inflights.push_back(std::move(f)); + // Submit every peer (not waited) to overlap the wire across peers, run the + // local puts in that window, then wait all + commit. On early exit the + // engine handle's destructor drains; the wait does mapping + commit/abort. + // + // There is no longer an all-zero-copy / all-staging fork here. Staging is + // the engine's bounce pool, and a plan that needs it settles inside Submit, + // so submit-all is unconditionally safe — a batch that mixes registered and + // unregistered buffers, which used to be a contract violation the client + // failed, now just works. + std::vector> inflights; + inflights.reserve(plan.remote_groups.size()); + { + // NOTE: SubmitRemoteBatchPut's BatchAllocateSlots is a SYNCHRONOUS RPC, so + // this loop serializes one allocate round trip per peer even though the + // comment above promises overlap -- only the RDMA that follows overlaps. + for (const auto& [node_id, items] : plan.remote_groups) { + if (auto f = SubmitRemoteBatchPut(items, results)) inflights.push_back(std::move(f)); } } run_local_put(); for (auto& f : inflights) WaitRemoteBatchPut(*f, results); } -std::unique_ptr PoolClient::SubmitRemoteBatchPut( - const std::vector& items, std::vector* results, - bool permit_staging) { +std::unique_ptr PoolClient::SubmitRemoteBatchPut( + const std::vector& items, std::vector* results) { if (items.empty()) return nullptr; auto fail_all = [&] { for (const auto& item : items) (*results)[item.index] = PutEntryOutcome::kFailed; }; - if (!io_engine_) { - MORI_UMBP_ERROR("[PoolClient] SubmitRemoteBatchPut: io_engine_ not initialized (items={})", + if (peer_directory_ == nullptr) { + MORI_UMBP_ERROR("[PoolClient] SubmitRemoteBatchPut: no RDMA engine configured (items={})", items.size()); fail_all(); return nullptr; @@ -1114,7 +2449,7 @@ std::unique_ptr PoolClient::SubmitRemoteBatch } auto* stub = static_cast<::umbp::UMBPPeer::Stub*>(peer.peer_stub.get()); - auto inflight = std::make_unique(); + auto inflight = std::make_unique(); inflight->peer = &peer; inflight->stub = stub; @@ -1143,139 +2478,62 @@ std::unique_ptr PoolClient::SubmitRemoteBatch return nullptr; } - std::vector transfers; - uint64_t staging_bytes = 0; - if (!BuildRemotePutTransfers(inflight->entries, peer, &transfers, &staging_bytes)) { - MORI_UMBP_WARN( - "[PoolClient] SubmitRemoteBatchPut: BuildRemotePutTransfers failed, node='{}' entries={} " - "-> aborting all slots", - first.route.node_id, inflight->entries.size()); - // Build failed wholesale: abort everything allocated. + // Abort everything allocated and fail every key: used on the paths that + // return nullptr, where no WaitRemoteBatchPut/finalize will run. + auto abort_everything = [&] { std::vector all = std::move(inflight->abort_slots); for (auto& entry : inflight->entries) { all.push_back(entry.slot_id); (*results)[entry.result_index] = PutEntryOutcome::kFailed; } abort_now(std::move(all)); - return nullptr; - } - inflight->staging_bytes = staging_bytes; + }; - // Staging uses the shared staging_buffer_, so a cross-peer submit-all cannot - // serialize on the single staging_mutex_ without deadlocking: the zero-copy - // path passes permit_staging=false and fails such a batch (contract: a batch - // is all-zc or all-staging). The serial staging path acquires staging_mutex_, - // copies src->staging here (BEFORE BatchWrite, opposite of Get), and parks the - // lock in the in-flight until WaitRemoteBatchPut releases it. - if (staging_bytes > 0) { - if (!permit_staging) { - MORI_UMBP_ERROR( - "[PoolClient] SubmitRemoteBatchPut: unexpected staging_bytes={} on zero-copy path " - "node='{}' -> failing batch", - staging_bytes, first.route.node_id); - std::vector all = std::move(inflight->abort_slots); - for (auto& entry : inflight->entries) { - all.push_back(entry.slot_id); - (*results)[entry.result_index] = PutEntryOutcome::kFailed; - } - abort_now(std::move(all)); - return nullptr; - } - inflight->staging_lock = std::unique_lock(staging_mutex_); - for (auto& entry : inflight->entries) { - if (entry.failed || !entry.use_staging) continue; - // Defensive re-check before the memcpy (build already validated the cursor). - if (entry.staging_offset + entry.item->size > config_.staging_buffer_size) { - MORI_UMBP_ERROR( - "[PoolClient] SubmitRemoteBatchPut: staging offset overflow (should not happen), " - "key='{}' staging_offset={} size={} cap={}", - (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, - entry.staging_offset, entry.item ? entry.item->size : 0, config_.staging_buffer_size); - entry.failed = true; - continue; - } - std::memcpy(staging_buffer_.get() + entry.staging_offset, entry.item->src, entry.item->size); - } + std::vector transfer_items; + if (!BuildRemotePutTransfers(inflight->entries, first.route.node_id, &transfer_items)) { + MORI_UMBP_WARN( + "[PoolClient] SubmitRemoteBatchPut: BuildRemotePutTransfers failed, node='{}' entries={} " + "-> aborting all slots", + first.route.node_id, inflight->entries.size()); + abort_everything(); + return nullptr; } - // Drop transfers whose entry failed during build. Those failed entries ride in + // Drop items whose entry failed during build. Those failed entries ride in // inflight->entries and are aborted by FinalizeRemotePutEntries at wait time — - // do NOT early-abort them here (they're not a whole-batch failure). - std::vector active_transfers; - active_transfers.reserve(transfers.size()); - for (const auto& t : transfers) { - if (!inflight->entries[t.entry_index].failed) active_transfers.push_back(t); - } - if (active_transfers.empty()) { - // Nothing to post: no in-flight returned, so finalize never runs. Release - // staging_mutex_ before the abort RPC (no need to hold it over the wire), - // then abort every allocated slot and fail every key. - if (inflight->staging_lock.owns_lock()) inflight->staging_lock.unlock(); - std::vector all = std::move(inflight->abort_slots); - for (auto& entry : inflight->entries) { - all.push_back(entry.slot_id); - (*results)[entry.result_index] = PutEntryOutcome::kFailed; - } - abort_now(std::move(all)); + // do NOT early-abort them here (they are not a whole-batch failure). + std::vector active; + active.reserve(transfer_items.size()); + for (auto& item : transfer_items) { + if (!inflight->entries[item.tag].failed) active.push_back(std::move(item)); + } + if (active.empty()) { + abort_everything(); return nullptr; } - // Collapse same-pair pages into one outer transfer, materialised INTO the - // in-flight so the args outlive the post (see RemoteDramGetInFlight). - inflight->groups = GroupTransfersByPair(active_transfers); - const size_t G = inflight->groups.size(); - inflight->local_descs.resize(G); - inflight->remote_descs.resize(G); - inflight->local_offsets.resize(G); - inflight->remote_offsets.resize(G); - inflight->sizes_v.resize(G); - inflight->statuses = std::vector(G); // built once, never resized/moved - inflight->status_ptrs.resize(G); - inflight->ids.resize(G); - for (size_t g = 0; g < G; ++g) { - inflight->local_descs[g] = inflight->groups[g].local_desc; - inflight->remote_descs[g] = inflight->groups[g].remote_desc; - inflight->local_offsets[g] = std::move(inflight->groups[g].local_offsets); - inflight->remote_offsets[g] = std::move(inflight->groups[g].remote_offsets); - inflight->sizes_v[g] = std::move(inflight->groups[g].sizes); - inflight->status_ptrs[g] = &inflight->statuses[g]; - inflight->ids[g] = io_engine_->AllocateTransferUniqueId(); - } - - // POST the writes; do NOT wait. All args are owned by *inflight; for staging - // the src bytes are already in staging_buffer_ under staging_lock. - io_engine_->BatchWrite(inflight->local_descs, inflight->local_offsets, inflight->remote_descs, - inflight->remote_offsets, inflight->sizes_v, inflight->status_ptrs, - inflight->ids); + TransferPlanSet planned = transfer_engine_->Plan(active); + ApplyRejectedTags(inflight->entries, planned.rejected_tags, "RemotePut"); + if (planned.plans.empty()) { + abort_everything(); + return nullptr; + } + // POST; do NOT wait. Everything the post references — including any bytes + // staged through the engine's bounce pool — is owned by the returned handle. + inflight->handle = transfer_engine_->Submit(std::move(planned.plans)); + if (inflight->handle == nullptr) { + abort_everything(); + return nullptr; + } return inflight; } -void PoolClient::WaitRemoteBatchPut(RemoteDramPutInFlight& f, - std::vector* results) { +void PoolClient::WaitRemoteBatchPut(RemotePutInFlight& f, std::vector* results) { if (f.drained) return; - f.drained = true; // set early (mirror Get) so the destructor never re-waits. - // Wait every group (never break early); a non-success group fails all of its - // contributing keys (per-item AND). - const size_t G = f.groups.size(); - for (size_t g = 0; g < G; ++g) { - f.statuses[g].Wait(); - if (f.statuses[g].Succeeded()) continue; - for (size_t ei : f.groups[g].entry_indices) { - auto& entry = f.entries[ei]; - MORI_UMBP_ERROR( - "RemotePut BatchWrite failed: code={} msg='{}' peer_engine='{}' key='{}' slot_id={} " - "use_staging={}", - f.statuses[g].CodeUint32(), f.statuses[g].Message(), f.groups[g].remote_desc.engineKey, - (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, entry.slot_id, - entry.use_staging); - entry.failed = true; - } - } - - // RDMA has consumed staging_buffer_; release staging_mutex_ (Put has no - // staging->dst copy-out — the data went to the peer). - if (f.staging_lock.owns_lock()) f.staging_lock.unlock(); - + f.drained = true; + std::vector failures; + if (f.handle != nullptr) f.handle->Wait(&failures); + ApplyTransferFailures(f.entries, failures, "RemotePut"); FinalizeRemotePutEntries(f.entries, f.abort_slots, results, f.stub); } @@ -1291,6 +2549,7 @@ bool PoolClient::AllocateRemotePutEntries(const std::vector& items entry->set_size(item.size); entry->set_tier(static_cast<::umbp::TierType>(item.route.tier)); entry->set_key(*item.key); + entry->set_logical_tier(item.route.logical_tier); } ::umbp::BatchAllocateSlotsResponse alloc_resp; @@ -1369,88 +2628,71 @@ bool PoolClient::AllocateRemotePutEntries(const std::vector& items return true; } -bool PoolClient::BuildRemotePutTransfers(std::vector& entries, PeerConnection& peer, - std::vector* transfers, - uint64_t* staging_bytes) { - transfers->clear(); - uint64_t cursor = 0; +bool PoolClient::BuildRemotePutTransfers(std::vector& entries, + const std::string& node_id, + std::vector* items) { + items->clear(); + + // Hydrate every entry's descs first, then snapshot the peer's buffers once + // PER BACKEND the batch touches: the loop below indexes a snapshot instead of + // taking the engine's remote lock per page. A batch may span media (each + // item carries its own route.tier), and buffer_index is backend-local, so one + // snapshot per peer would index the wrong medium's buffers. A concurrent + // hydrate can only add buffers, so a snapshot taken here is never stale for + // the indices these entries reference. + for (const auto& entry : entries) { + if (!entry.plan.descs.empty()) peer_directory_->CacheRemoteBuffers(node_id, entry.plan.descs); + } + std::array, kMaxBackendsPerPeer> snapshots; + std::array snapped{}; + auto buffers_for = [&](uint32_t backend_id) -> const std::vector& { + static const std::vector kNone; + if (backend_id >= kMaxBackendsPerPeer) return kNone; + if (!snapped[backend_id]) { + snapshots[backend_id] = peer_directory_->RemoteBufferSnapshot(node_id, backend_id); + snapped[backend_id] = true; + } + return snapshots[backend_id]; + }; for (size_t idx = 0; idx < entries.size(); ++idx) { auto& entry = entries[idx]; - - auto zero_copy = FindRegisteredMemory(entry.item->src, entry.item->size); - if (zero_copy.has_value()) { - entry.zero_copy = zero_copy; - entry.use_staging = false; - entry.staging_offset = 0; - } else { - entry.zero_copy.reset(); - entry.use_staging = true; - cursor = AlignUp(cursor, entry.plan.page_size); - const uint64_t aligned_size = AlignUp(entry.item->size, entry.plan.page_size); - if (cursor + aligned_size > config_.staging_buffer_size) { - MORI_UMBP_WARN( - "[PoolClient] BuildRemotePutTransfers: staging buffer overflow at entry {}/{}, " - "key='{}' size={} aligned={} cursor={} cap={}", - idx, entries.size(), + // Whether this goes zero-copy or through the engine's bounce pool is the + // engine's decision; this only names the endpoint. + const auto [src, src_base] = + UserBufferRef(const_cast(entry.item->src), entry.item->size); + const std::vector& remote = buffers_for(entry.plan.backend_id); + + std::vector entry_items; + entry_items.reserve(entry.plan.pages.size()); + for (size_t p = 0; p < entry.plan.pages.size(); ++p) { + const auto& page = entry.plan.pages[p]; + if (page.buffer_index >= remote.size() || !remote[page.buffer_index].HasMemoryDesc()) { + MORI_UMBP_ERROR( + "[PoolClient] BuildRemotePutTransfers: peer published no buffer, " + "key='{}' backend={} buffer_index={} peer_buffers={} page_index={}", (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, - entry.item ? entry.item->size : 0, aligned_size, cursor, config_.staging_buffer_size); - return false; - } - entry.staging_offset = cursor; - cursor += aligned_size; - } - - mori::io::MemoryDesc local_desc{}; - uint64_t base_offset = 0; - if (entry.use_staging) { - local_desc = staging_mem_; - base_offset = entry.staging_offset; - } else { - local_desc = entry.zero_copy->first; - base_offset = entry.zero_copy->second; - } - - // Hydrate + snapshot remote descs inside ONE peers_mutex_ window: a - // concurrent EnsureBufferDescsCached may resize dram_memories, so the - // reads below must not race it. zero_copy/staging above touch no peer - // state and stay outside the lock. - std::vector entry_transfers; - entry_transfers.reserve(entry.plan.pages.size()); - { - std::lock_guard lock(peers_mutex_); - EnsureBufferDescsCachedLocked(peer, entry.plan.descs); - for (size_t p = 0; p < entry.plan.pages.size(); ++p) { - const auto& page = entry.plan.pages[p]; - if (page.buffer_index >= peer.dram_memories.size() || - !IsValidMemoryDesc(peer.dram_memories[page.buffer_index])) { - MORI_UMBP_ERROR( - "[PoolClient] BuildRemotePutTransfers: invalid peer dram_memories slot, " - "key='{}' buffer_index={} peer_dram_size={} page_index={}", - (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, - page.buffer_index, peer.dram_memories.size(), page.page_index); - entry.failed = true; - entry_transfers.clear(); - break; - } - TransferInstruction instr; - instr.entry_index = idx; - instr.local_desc = local_desc; - instr.local_offset = base_offset + static_cast(p) * entry.plan.page_size; - instr.remote_desc = peer.dram_memories[page.buffer_index]; - instr.remote_offset = static_cast(page.page_index) * entry.plan.page_size; - instr.size = - LogicalPageBytes(p, entry.plan.pages.size(), entry.plan.page_size, entry.item->size); - entry_transfers.push_back(std::move(instr)); + entry.plan.backend_id, page.buffer_index, remote.size(), page.page_index); + entry.failed = true; + entry_items.clear(); + break; } + TransferItem item; + item.tag = idx; + item.src = src; + item.src_offset = src_base + static_cast(p) * entry.plan.page_size; + item.dst = remote[page.buffer_index]; + item.dst_offset = static_cast(page.page_index) * entry.plan.page_size; + item.size = + LogicalPageBytes(p, entry.plan.pages.size(), entry.plan.page_size, entry.item->size); + entry_items.push_back(std::move(item)); } - if (!entry_transfers.empty()) { - transfers->insert(transfers->end(), entry_transfers.begin(), entry_transfers.end()); + if (!entry_items.empty()) { + items->insert(items->end(), std::make_move_iterator(entry_items.begin()), + std::make_move_iterator(entry_items.end())); } } - - *staging_bytes = cursor; return true; } @@ -1494,10 +2736,9 @@ void PoolClient::FinalizeRemotePutEntries(std::vector& entries, auto idx = commit_indices[i]; auto& entry = entries[idx]; if (commit_resp.success(static_cast(i))) { - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL_HELP, - {{"traffic", "remote"}}, - static_cast(entry.item->size)); + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_PUT_BYTES_TOTAL_HELP, + {{"traffic", "remote"}}, static_cast(entry.item->size)); (*results)[entry.result_index] = PutEntryOutcome::kSucceeded; } else { // Peer allocator already logged the reason (SLOT_GONE / PRE_CLEAR). @@ -1537,13 +2778,189 @@ bool PoolClient::Get(const std::string& key, void* dst, size_t size) { return !results.empty() && results[0]; } -std::vector PoolClient::BatchGet(const std::vector& keys, - const std::vector& dsts, - const std::vector& sizes) { - const auto call_start = std::chrono::steady_clock::now(); - std::vector results(keys.size(), false); - if (keys.size() != dsts.size() || keys.size() != sizes.size()) { - MORI_UMBP_ERROR("[PoolClient] BatchGet: vector length mismatch"); +// Serve `indices` from this node's own media in one resolve and one transfer. +// Hits are written into *results; anything this node could not serve is +// appended to *missed when the caller wants it back (null when there is no +// fallback left to try). +// +// Batched on both axes deliberately. One BatchResolve per BACKEND, not per +// key: that mutex is shared with Allocate / Commit / Evict, and in +// standalone-process mode every rank on the node shares this client. Then ONE +// Transfer for every key, tagged by key index so the engine attributes a +// failure to its key instead of failing the batch. The per-key alternative +// measured 1.4-1.6x SLOWER than route-first at batch >= 32 despite issuing no +// RPC at all; batching it made local-first win at every batch size. +void PoolClient::ServeLocalGets(const std::vector& keys, + const std::vector& dsts, const std::vector& sizes, + const std::vector& indices, std::vector* results, + std::vector* missed) { + const auto miss = [&](size_t i) { + if (missed != nullptr) missed->push_back(i); + }; + if (indices.empty()) return; + + std::vector holders; + std::vector resolutions; + ResolveLocalBatch(keys, indices, &holders, &resolutions, + [&](size_t i) { return DestinationIsDevice(dsts[i], sizes[i]); }); + + const auto t0 = std::chrono::steady_clock::now(); + std::vector items; + std::vector local_keys; + local_keys.reserve(indices.size()); + for (size_t i : indices) { + MediumBackend* const holder = holders[i]; + if (holder == nullptr) { + miss(i); + continue; + } + const ResolvedEntry& resolved = resolutions[i]; + // A stored size that disagrees with the requested one is a different + // object; copying sizes[i] out of a slot sized for something else would + // read past it. Look for it elsewhere rather than report a hit. + if (resolved.size != sizes[i]) { + MORI_UMBP_WARN("[PoolClient] local Get: size mismatch for key='{}' (wanted {}, got {})", + keys[i], sizes[i], resolved.size); + miss(i); + continue; + } + const size_t before = items.size(); + if (resolved.file_ref.IsFile()) { + // Zero-copy GDS path, the same one ExecuteLocalGet takes: the backend + // published a file range rather than staged pages, so read it straight + // into the caller's device buffer and let the planner pick GdsEngine. + TransferItem item; + item.src = resolved.file_ref; + item.dst = ClassifiedUserBytes(dsts[i], sizes[i]); + item.size = sizes[i]; + items.push_back(std::move(item)); + } else if (!BuildLocalPageTransfers(holder, resolved.pages, resolved.page_size, dsts[i], + sizes[i], /*to_backend=*/false, &items)) { + // This medium holds the key but publishes no in-process endpoint for its + // buffers. Drop what was appended for it -- a half-built key must not + // ride along in the batch -- and route it instead. + items.resize(before); + miss(i); + continue; + } + for (size_t k = before; k < items.size(); ++k) items[k].tag = i; + local_keys.push_back(i); + } + if (items.empty()) return; + + std::vector failed_tags; + transfer_engine_->Transfer(items, &failed_tags); + const std::unordered_set failed(failed_tags.begin(), failed_tags.end()); + + double local_get_bytes = 0.0; + for (size_t i : local_keys) { + if (failed.count(i) != 0) { + // Any failed fragment fails the whole key: the destination now holds a + // partial object, and reporting a hit would hand the caller bytes it + // cannot tell are incomplete. + miss(i); + continue; + } + (*results)[i] = true; + local_get_bytes += static_cast(sizes[i]); + } + if (local_get_bytes > 0.0) { + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + local_get_bytes); + CountMetric(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + local_get_bytes); + } + if (std::getenv("UMBP_LOCAL_COPY_TIMING")) { + const double sec = std::chrono::duration_cast>( + std::chrono::steady_clock::now() - t0) + .count(); + MORI_UMBP_INFO("[LocalCopy] GET keys={} items={} bytes={} elapsed_ms={:.3f} GiB_s={:.2f}", + local_keys.size(), items.size(), local_get_bytes, sec * 1000.0, + local_get_bytes / (sec > 0 ? sec : 1e-12) / (1024.0 * 1024 * 1024)); + } +} + +// Whole-batch entry point for the local-first arm: filter what cannot be served +// at all, serve the rest, and hand back the keys the master still has to route. +std::vector PoolClient::ServeLocalBatchGet(const std::vector& keys, + const std::vector& dsts, + const std::vector& sizes, + std::vector* results) { + std::vector candidates; + candidates.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + // Silently skipped here; PartitionBatchGetTargets sees every key and owns + // the warning, so routing it through both would log each one twice. + if (sizes[i] == 0) continue; + candidates.push_back(i); + } + std::vector missed; + missed.reserve(candidates.size()); + ServeLocalGets(keys, dsts, sizes, candidates, results, &missed); + return missed; +} + +// Route the keys the local phase missed and group them by peer. `routes` is +// indexed by ORIGINAL key index and left nullopt for anything served locally, +// which is how ComputeBatchBandwidthBytes already reads a missing route. +// +// Self is excluded: ServeLocalBatchGet already asked every medium here, so a +// route back to this node could only repeat that miss. The returned plan has +// no local half for the same reason. False means the routing RPC failed. +bool PoolClient::RouteGetsInto(const std::vector& keys, + const std::vector& indices, bool exclude_self, + std::vector>* routes) { + if (indices.empty()) return true; + // No master means no peers: this node already looked, so the miss is final + // and every slot stays nullopt. Not a failure -- there was nothing to ask. + if (!HasMaster()) return true; + + std::vector route_keys; + route_keys.reserve(indices.size()); + for (size_t i : indices) route_keys.push_back(keys[i]); + + std::unordered_set excludes; + if (exclude_self) excludes.insert(config_.master_config.node_id); + + std::vector> answers; + auto status = master_client_->BatchRouteGet(route_keys, excludes, &answers); + if (!status.ok()) { + MORI_UMBP_ERROR("[PoolClient] BatchGet: BatchRouteGet failed: {}", status.error_message()); + return false; + } + // A short reply is not an error; the unanswered keys stay nullopt, which the + // partition reads as "nowhere". + answers.resize(indices.size()); + for (size_t r = 0; r < indices.size(); ++r) (*routes)[indices[r]] = answers[r]; + return true; +} + +void PoolClient::ObserveBatchGetBandwidth(const std::vector& results, + const std::vector& sizes, + const std::vector>& routes, + std::chrono::steady_clock::time_point call_start) { + const double seconds = std::chrono::duration_cast>( + std::chrono::steady_clock::now() - call_start) + .count(); + if (seconds <= 0.0) return; + auto split = ComputeBatchBandwidthBytes(results, sizes, routes, config_.master_config.node_id); + ObserveBatchBandwidth(master_client_.get(), split.local, seconds, + MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH, + MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH_HELP, "local"); + ObserveBatchBandwidth(master_client_.get(), split.remote, seconds, + MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH, + MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH_HELP, "remote"); +} + +std::vector PoolClient::BatchGet(const std::vector& keys, + const std::vector& dsts, + const std::vector& sizes) { + const auto call_start = std::chrono::steady_clock::now(); + std::vector results(keys.size(), false); + if (keys.size() != dsts.size() || keys.size() != sizes.size()) { + MORI_UMBP_ERROR("[PoolClient] BatchGet: vector length mismatch"); return results; } if (!initialized_) { @@ -1551,38 +2968,830 @@ std::vector PoolClient::BatchGet(const std::vector& keys, return results; } + // Both arms are route-then-partition; they differ only in WHICH keys reach + // the master and whether an unroutable key may still fall back to a local + // read. Left nullopt for anything served locally, which is how + // ComputeBatchBandwidthBytes already reads a missing route. + std::vector> routes(keys.size()); + const bool local_first = config_.local_first && !registry_.Empty(); + + std::vector to_route; + if (local_first) { + to_route = ServeLocalBatchGet(keys, dsts, sizes, &results); + // The whole point of the flag: a batch this node could satisfy by itself + // never touches the master. + if (to_route.empty()) { + ObserveBatchGetBandwidth(results, sizes, routes, call_start); + return results; + } + } else { + to_route.resize(keys.size()); + std::iota(to_route.begin(), to_route.end(), 0); + } + + // A routing failure is not fatal to what is already in hand: under local-first + // the phase-1 hits stand, since they never depended on this RPC. + if (RouteGetsInto(keys, to_route, /*exclude_self=*/local_first, &routes)) { + const BatchGetPlan plan = PartitionBatchGetTargets( + keys, dsts, sizes, routes, local_first ? LocalFallback::kSkip : LocalFallback::kAllow); + ExecuteBatchGetPlan(plan, keys, dsts, sizes, &results); + } + + ObserveBatchGetBandwidth(results, sizes, routes, call_start); + return results; +} + +// --------------------------------------------------------------------------- +// Ranged batch entry points +// --------------------------------------------------------------------------- + +namespace { + +// Fills `out` rather than returning a fresh vector: the caller loops over keys +// and one allocation per key adds up on a call carrying thousands of ranges. +void FillReadRanges(const std::vector& dsts, const std::vector& sizes, + const std::vector& offsets, std::vector* out) { + out->clear(); + out->reserve(dsts.size()); + for (size_t r = 0; r < dsts.size(); ++r) out->push_back({dsts[r], sizes[r], offsets[r]}); +} + +std::vector MakeWriteRanges(const std::vector& srcs, + const std::vector& sizes, + const std::vector& offsets) { + std::vector out; + out.reserve(srcs.size()); + for (size_t r = 0; r < srcs.size(); ++r) { + out.push_back({const_cast(srcs[r]), sizes[r], offsets[r]}); + } + return out; +} + +} // namespace + +// The route-first arm of BatchGetRanges: one BatchRouteGet over EVERY key, +// issued before anything is served. +// +// Turning local_first off cannot mean "do not look locally" here -- the local +// medium is the only thing that can serve a local key on this path. What it +// means is "ask the master the way the whole-object path used to", so the +// difference between the arms is purely WHEN and over HOW MANY keys the routing +// RPC is issued: all of them up front, versus just the misses (or none at all +// when the batch is entirely local). +// +// Uses the same exclude set as the missed-key routing below, so both arms send +// an identical request shape and a self-route cannot reach the remote planner +// from either one. False means the RPC failed. +bool PoolClient::RouteAllRangesUpFront(const std::vector& keys, double* route_sink, + std::vector>* preroutes) { + if (!HasMaster()) { + // Nothing to pre-route; phase 2 will find every miss unroutable, which on a + // single node is the correct answer rather than a degraded one. + preroutes->assign(keys.size(), std::nullopt); + return true; + } + std::unordered_set excludes{config_.master_config.node_id}; + PhaseTimer route_timer(route_sink); + auto status = master_client_->BatchRouteGet(keys, excludes, preroutes); + if (!status.ok()) { + MORI_UMBP_ERROR("[PoolClient] BatchGetRanges: BatchRouteGet failed: {}", + status.error_message()); + return false; + } + preroutes->resize(keys.size()); + return true; +} + +// Routes for the keys the local phase missed, parallel to `missed`. +// +// Under local_first this is the only routing RPC the call makes. Under the +// route-first arm phase 0 already routed every key, so this projects those +// answers onto the misses instead of paying a second round trip -- that arm +// must cost ONE routing RPC, not two, or a comparison against it measures the +// harness rather than the flag. False means the RPC failed. +bool PoolClient::RouteMissedRanges(const std::vector& route_keys, + const std::vector& missed, + const std::vector>& preroutes, + double* route_sink, + std::vector>* routes) { + if (!HasMaster()) { + routes->assign(route_keys.size(), std::nullopt); + return true; + } + if (!config_.local_first) { + routes->reserve(missed.size()); + for (size_t index : missed) routes->push_back(preroutes[index]); + } else { + std::unordered_set excludes{config_.master_config.node_id}; + PhaseTimer route_timer(route_sink); + auto status = master_client_->BatchRouteGet(route_keys, excludes, routes); + if (!status.ok()) { + MORI_UMBP_ERROR("[PoolClient] BatchGetRanges: BatchRouteGet failed: {}", + status.error_message()); + return false; + } + } + routes->resize(route_keys.size()); + return true; +} + +std::vector PoolClient::BatchGetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& src_offsets) { + const size_t n = keys.size(); + std::vector results(n, false); + if (!initialized_ || n == 0 || !RangeBatchShapeValid(n, dsts, sizes, src_offsets)) { + return results; + } + + // Observed at scope exit for whichever phases ran; see ScopedBatchBandwidth. + double local_bytes = 0.0; + double remote_bytes = 0.0; + ScopedBatchBandwidth bandwidth( + master_client_.get(), MORI_UMBP_METRIC_CLIENT_BATCH_GET_RANGES_BANDWIDTH, + MORI_UMBP_METRIC_CLIENT_BATCH_GET_RANGES_BANDWIDTH_HELP, local_bytes, remote_bytes); + + // Sub-timers. `dbg` is null when disabled, which makes every PhaseTimer inert. + const bool ranged_debug = RangedDebugEnabled(); + RangedPhases phases; + phases.keys = n; + if (ranged_debug && n > 0) phases.key0 = &keys[0]; + RangedPhases* dbg = ranged_debug ? &phases : nullptr; + ScopedRangedReport ranged_report("get", phases, ranged_debug, local_bytes, remote_bytes); + + // Accumulated over the batch and reported once at the end rather than per + // key. AddCounter takes a lock and builds a metric key out of the name and + // labels, and the help text is a hundred-odd bytes that has to be + // materialized on every call -- affordable once, not a thousand times. The + // counter totals are identical either way; only the number of updates + // changes. + double local_get_bytes = 0.0; + auto record_local_get = [&](size_t index) { + local_get_bytes += + static_cast(std::accumulate(sizes[index].begin(), sizes[index].end(), size_t{0})); + }; + auto flush_local_get_bytes = [&] { + if (local_get_bytes == 0.0) return; + local_bytes += local_get_bytes; + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + local_get_bytes); + CountMetric(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "local"}}, + local_get_bytes); + local_get_bytes = 0.0; + }; + // Declared after ScopedBatchBandwidth and ScopedRangedReport, so it runs + // BEFORE either of them reads local_bytes. + ScopeExit flush_guard{flush_local_get_bytes}; + + // ---- Phase 0: the route-first arm (local_first = false) ---- + std::vector> preroutes; + if (!config_.local_first && + !RouteAllRangesUpFront(keys, dbg ? &dbg->route : nullptr, &preroutes)) { + return results; + } + + // ---- Phase 1: everything this node already holds, in one submit ---- + // + // Every key's items carry tag = key index, so the engine reports failures per + // key and one bad key does not fail the batch. This is why the local half + // needs no per-key threading the way whole-object BatchGet does: the copies + // are already batched into a single plan set. + std::vector local_items; + std::vector local_keys; + std::vector missed; + missed.reserve(n); + // A TransferItem carries two TransferRefs, each with a MemoryDesc, so growing + // this vector by doubling relocates a lot of bytes. One range is at least one + // item -- more only when it straddles a page -- so the range count is a tight + // lower bound and covers the common case exactly. + { + size_t range_total = 0; + for (const auto& per_key : sizes) range_total += per_key.size(); + local_items.reserve(range_total); + } + std::vector key_ranges; + + // A key that asked for no ranges is never resolved -- there is nothing to + // serve it from, and it stays false either way. + std::vector candidates; + candidates.reserve(n); + for (size_t i = 0; i < n; ++i) { + if (!sizes[i].empty()) candidates.push_back(i); + } + std::vector holders; + std::vector resolutions; + { + PhaseTimer resolve_timer(dbg ? &dbg->resolve : nullptr); + ResolveLocalBatch(keys, candidates, &holders, &resolutions, [&](size_t i) { + // One key's ranges normally share a caller allocation, but nothing + // guarantees it; a single host range makes the whole key unservable + // from a file ref, so require all of them. + const std::vector& key_dsts = dsts[i]; + for (size_t r = 0; r < key_dsts.size(); ++r) { + if (!DestinationIsDevice(key_dsts[r], sizes[i][r])) return false; + } + return !key_dsts.empty(); + }); + } + + for (size_t i = 0; i < n; ++i) { + if (sizes[i].empty()) continue; // asked for nothing; stays false + + MediumBackend* const holder = holders[i]; + const ResolvedEntry& resolved = resolutions[i]; + if (holder == nullptr) { + missed.push_back(i); + continue; + } + // The get API takes no object_sizes; the resolved entry is where the true + // stored size becomes known. First hit labels the call. + if (dbg != nullptr && dbg->object_size == 0) { + dbg->object_size = static_cast(resolved.size); + } + if (!RangesAreDisjointAndInBounds(static_cast(resolved.size), sizes[i], + src_offsets[i])) { + MORI_UMBP_ERROR("[PoolClient] BatchGetRanges: invalid local ranges for key='{}' size={}", + keys[i], resolved.size); + continue; // a caller bug, not a miss — do not go looking for a replica + } + const size_t before = local_items.size(); + bool built = false; + { + PhaseTimer build_timer(dbg ? &dbg->build : nullptr); + // Refilled per key rather than returned fresh: one allocation for the + // call instead of one per key. + FillReadRanges(dsts[i], sizes[i], src_offsets[i], &key_ranges); + built = + resolved.file_ref.IsFile() + ? BuildLocalFileRangeTransfers(resolved.file_ref, key_ranges, /*tag=*/i, &local_items) + : BuildLocalRangeTransfers(holder, resolved.pages, resolved.page_size, resolved.size, + key_ranges, /*to_backend=*/false, + /*tag=*/i, &local_items, dbg ? &dbg->classify : nullptr); + } + if (!built) { + // The medium holds the key but cannot be read in-process. Same call as + // Same disposition the local get path takes on a size mismatch: look + // reporting a miss the caller cannot distinguish from absence. + local_items.resize(before); + missed.push_back(i); + continue; + } + local_keys.push_back(i); + } + if (dbg != nullptr) { + // Classification is timed inside the build window, so subtract it out to + // keep the phases disjoint and `other` meaningful. + dbg->build -= dbg->classify; + dbg->items += local_items.size(); + } + if (!local_items.empty()) { + std::vector failed_tags; + { + PhaseTimer xfer_timer(dbg ? &dbg->xfer : nullptr); + TransferEngine::StepTiming steps; + if (dbg != nullptr) steps = {&dbg->xfer_plan, &dbg->xfer_submit, &dbg->xfer_wait}; + transfer_engine_->Transfer(local_items, &failed_tags, dbg != nullptr ? &steps : nullptr); + } + const std::unordered_set failed(failed_tags.begin(), failed_tags.end()); + for (size_t i : local_keys) { + if (failed.count(i) != 0) continue; + results[i] = true; + record_local_get(i); + if (dbg != nullptr) dbg->local_keys += 1; + } + } + if (dbg != nullptr) dbg->remote_keys = missed.size(); + if (missed.empty()) return results; + + // No master, no remote phase: nothing can route these keys anywhere, so a + // local miss is the final answer (the same disposition BatchGet takes). + // Returning here rather than falling into the arena check matters because an + // embedded deployment allocates no arena -- without this, every ordinary + // cache miss would report itself as a scratch-arena error. + if (!HasMaster()) return results; + + // ---- Phase 2: the rest, through the GET scratch arena ---- + // + // Only the requested spans cross the wire, packed back-to-back into the + // arena. A layer-wise reader asks for one layer group out of many, so + // fetching the whole object would move an order of magnitude more bytes than + // it uses, and would do it on the first group -- the one the caller cannot + // overlap with anything. Packing spans instead also multiplies how many keys + // fit in one arena round by the same factor. + // + // GET has its own arena and mutex, so it overlaps a concurrent remote PUT + // (which uses the separate PUT arena/mutex). + char* const scratch_base = static_cast(config_.ranged_get_scratch_buffer); + const size_t scratch_size = config_.ranged_get_scratch_size; + if (scratch_base == nullptr || scratch_size == 0 || + !FindRegisteredMemory(scratch_base, scratch_size).has_value()) { + MORI_UMBP_ERROR("[PoolClient] BatchGetRanges: ranged scratch is absent or not registered"); + return results; + } + + // Routing is a master RPC that never touches the arena, so it runs before the + // lock is taken rather than holding every other arena user behind it. + std::vector route_keys; + route_keys.reserve(missed.size()); + for (size_t index : missed) route_keys.push_back(keys[index]); std::vector> routes; + if (!RouteMissedRanges(route_keys, missed, preroutes, dbg ? &dbg->route : nullptr, &routes)) { + return results; + } + + // No tier filter here. Upstream restricts remote ranged reads to DRAM/HBM + // because its remote path is a hand-written DRAM RDMA read; ours goes through + // ExecuteRemoteBatchGetPlan, which reaches any medium the owning peer + // publishes — including SSD, whose staging page is ordinary registered host + // memory and so is readable at range granularity like any other. + // + // Each eligible key becomes one fetch unit: its spans, the bytes they total, + // and where in the arena they will land. Splitting by span prefix rather + // than by key means an object larger than the arena is no longer unservable — + // only a single span larger than the arena is. + std::vector units; + units.reserve(missed.size()); + + for (size_t r = 0; r < missed.size(); ++r) { + const size_t original = missed[r]; + if (!routes[r].has_value()) continue; // genuinely nowhere in the cluster + const auto& route = *routes[r]; + const size_t object_size = static_cast(route.size); + if (route.node_id == config_.master_config.node_id || route.size == 0 || + route.size > std::numeric_limits::max() || + !RangesAreDisjointAndInBounds(object_size, sizes[original], src_offsets[original])) { + MORI_UMBP_ERROR("[PoolClient] BatchGetRanges: invalid route/ranges for key='{}' size={}", + keys[original], route.size); + continue; + } + + // Built aside and only spliced in once the WHOLE key has been laid out. A + // key can need several units, and a later range that cannot be served must + // discard the earlier ones too -- committing them would fetch part of the + // key and still report it whole, which the caller cannot detect. + std::vector key_units; + RangeFetchUnit unit; + unit.original = original; + unit.object_size = object_size; + unit.route = routes[r]; + // Ascending and gapless from byte 0 is what makes the arena slice equal the + // object; a split breaks it because no single slice then holds everything. + bool ascending_from_zero = true; + bool split = false; + bool servable = true; + size_t cursor = 0; + for (size_t j = 0; j < sizes[original].size(); ++j) { + const size_t span_bytes = sizes[original][j]; + if (span_bytes > scratch_size) { + MORI_UMBP_ERROR( + "[PoolClient] BatchGetRanges: range exceeds ranged scratch key='{}' range={} " + "scratch={}", + keys[original], span_bytes, scratch_size); + servable = false; + break; + } + if (src_offsets[original][j] != cursor) ascending_from_zero = false; + cursor += span_bytes; + if (unit.bytes + span_bytes > scratch_size) { + split = true; + key_units.push_back(std::move(unit)); + unit = RangeFetchUnit{}; + unit.original = original; + unit.object_size = object_size; + unit.route = routes[r]; + } + // The copy-out list stays one entry per caller range -- each has its own + // buffer -- but the FETCH list does not have to. Consecutive layers are + // consecutive bytes of the object, and the arena packs them in the same + // order, so a whole layer group is one contiguous read. + // + // The engine's planner would coalesce these too, but only after sorting + // and bucketing every one of them: a 61-range whole-object read would + // build 61 TransferItems per key just to merge them back into one. That + // CPU is invisible next to a multi-MiB object and very visible next to a + // small one, so the merge happens here instead. + unit.packed.push_back( + ObjectRange{.user = dsts[original][j], .size = span_bytes, .object_offset = unit.bytes}); + if (!unit.spans.empty() && + unit.spans.back().object_offset + unit.spans.back().size == src_offsets[original][j]) { + unit.spans.back().size += span_bytes; + } else { + unit.spans.push_back( + ByteSpan{.object_offset = src_offsets[original][j], .size = span_bytes}); + } + unit.bytes += span_bytes; + } + if (!servable) continue; // key stays false; none of its units are queued + if (!unit.spans.empty()) { + unit.holds_whole_object = !split && ascending_from_zero && unit.bytes == object_size; + key_units.push_back(std::move(unit)); + } + units.insert(units.end(), std::make_move_iterator(key_units.begin()), + std::make_move_iterator(key_units.end())); + } + if (units.empty()) return results; + + // A key served by several units only succeeds if every one of them does. + std::unordered_map unit_ok; + for (const auto& unit : units) unit_ok.emplace(unit.original, true); + + // Whole-object units skip the arena entirely when a slot can be had; what + // comes back is everything that still needs one. + units = ServeWholeObjectUnitsFromMedium(keys, std::move(units), &unit_ok, &remote_bytes); + if (units.empty()) { + for (const auto& [original, ok] : unit_ok) results[original] = ok; + return results; + } + + // Only the arena users serialize; the local hits above were fully concurrent, + // and so were the routing RPC and the slot-served units. + std::unique_lock scratch_lock(ranged_get_scratch_mutex_, std::defer_lock); + { + PhaseTimer lock_timer(dbg ? &dbg->lock : nullptr); + scratch_lock.lock(); + } + + // Pack as many units into the arena as fit, fetch that group, then repeat. + // Units are 64 B apart so two concurrently-filled slices never share a cache + // line; spans WITHIN a unit are exactly packed, which is what lets a group of + // adjacent layer ranges coalesce into a single wire segment. + PhaseTimer remote_timer(dbg ? &dbg->remote : nullptr); + size_t pos = 0; + while (pos < units.size()) { + std::vector scratch_offsets; + size_t cursor = 0; + size_t end = pos; + for (; end < units.size(); ++end) { + size_t aligned = 0; + if (!AlignUpChecked(cursor, kRangedScratchAlignment, &aligned) || + aligned > scratch_size - units[end].bytes) { + break; + } + scratch_offsets.push_back(aligned); + cursor = aligned + units[end].bytes; + } + // Unit sizing already caps every unit at scratch_size, so the first one + // always fits and this loop always advances. + const size_t count = end - pos; + + std::vector sub_keys; + std::vector sub_dsts; + std::vector sub_object_sizes; + std::vector sub_bytes; + std::vector*> sub_spans; + std::vector> sub_routes; + sub_keys.reserve(count); + sub_dsts.reserve(count); + sub_object_sizes.reserve(count); + sub_bytes.reserve(count); + sub_spans.reserve(count); + sub_routes.reserve(count); + for (size_t j = 0; j < count; ++j) { + const auto& unit = units[pos + j]; + sub_keys.push_back(keys[unit.original]); + sub_dsts.push_back(scratch_base + scratch_offsets[j]); + sub_object_sizes.push_back(unit.object_size); + sub_bytes.push_back(unit.bytes); + sub_spans.push_back(&unit.spans); + sub_routes.push_back(unit.route); + } + + std::vector fetched(count, false); + BatchGetPlan plan = PartitionBatchGetRangeTargets(sub_keys, sub_dsts, sub_object_sizes, + sub_bytes, sub_spans, sub_routes); + // recache_remote=false: a ranged entry never holds the whole object, so + // there is nothing the generic re-cache could legally install. Locality is + // restored by MaybePrefetchWholeObject below instead. + ExecuteRemoteBatchGetPlan(plan, &fetched, /*recache_remote=*/false); + + // Scatter EVERY fetched unit out of the arena with ONE transfer, for the + // same reason the put side assembles with one (see BatchPutRanges): a + // per-unit CopyContiguousToRanges is a blocking Transfer each, and the + // launch + synchronize dominates the copy at these object sizes. + std::vector scatter_items; + std::vector scatter_tag_to_j; + scatter_items.reserve(count); + scatter_tag_to_j.reserve(count); + for (size_t j = 0; j < count; ++j) { + if (!fetched[j]) continue; + const auto& unit = units[pos + j]; + const size_t tag = scatter_tag_to_j.size(); + if (!BuildContiguousToRangesItems(TransferRef::HostBytes(sub_dsts[j], unit.bytes), + /*src_base=*/0, unit.bytes, unit.packed, tag, + &scatter_items)) { + continue; + } + scatter_tag_to_j.push_back(j); + } + std::unordered_set scatter_failed; + if (!scatter_items.empty()) { + std::vector failed_tags; + transfer_engine_->Transfer(scatter_items, &failed_tags); + for (size_t tag : failed_tags) { + if (tag < scatter_tag_to_j.size()) scatter_failed.insert(scatter_tag_to_j[tag]); + } + } + std::unordered_set scattered(scatter_tag_to_j.begin(), scatter_tag_to_j.end()); + + for (size_t j = 0; j < count; ++j) { + const auto& unit = units[pos + j]; + if (!fetched[j]) { + unit_ok[unit.original] = false; + continue; + } + if (scattered.count(j) == 0 || scatter_failed.count(j) != 0) { + unit_ok[unit.original] = false; + continue; + } + // unit.bytes is both what the caller asked for and what crossed the wire: + // the fetch is range-granular now, so the two no longer differ. + remote_bytes += static_cast(unit.bytes); + + // Make the NEXT read of this key local. Two ways, and exactly one of + // them applies: + // * the arena slice already IS the whole object -> install it from + // there, which costs one local copy and nothing on the wire. Only + // units the medium could not give a slot to reach this; + // * it is not -> ask the background worker to pull the object. + // Installing from the arena has to happen now, before the next sub-batch + // reuses this slice. + if (unit.holds_whole_object) { + MaybeInstallCompleteArenaObject(sub_keys[j], sub_dsts[j], unit.object_size); + } else { + MaybePrefetchWholeObject(sub_keys[j], unit.object_size, *unit.route); + } + } + pos = end; + } + + for (const auto& [original, ok] : unit_ok) results[original] = ok; + return results; +} + +std::vector PoolClient::BatchPutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& dst_offsets) { + const size_t n = keys.size(); + std::vector results(n, false); + if (!initialized_ || n == 0 || object_sizes.size() != n || + !RangeBatchShapeValid(n, srcs, sizes, dst_offsets)) { + return results; + } + + // Observed at scope exit for whichever phases ran; see ScopedBatchBandwidth. + double local_bytes = 0.0; + double remote_bytes = 0.0; + ScopedBatchBandwidth bandwidth( + master_client_.get(), MORI_UMBP_METRIC_CLIENT_BATCH_PUT_RANGES_BANDWIDTH, + MORI_UMBP_METRIC_CLIENT_BATCH_PUT_RANGES_BANDWIDTH_HELP, local_bytes, remote_bytes); + + const bool ranged_debug = RangedDebugEnabled(); + RangedPhases phases; + phases.keys = n; + if (ranged_debug && n > 0) { + // One call is always one pool (the tree connector groups by PoolName), so + // key/size from the first entry label the whole batch. + phases.key0 = &keys[0]; + phases.object_size = object_sizes[0]; + } + RangedPhases* dbg = ranged_debug ? &phases : nullptr; + ScopedRangedReport ranged_report("put", phases, ranged_debug, local_bytes, remote_bytes); + + // A put must account for every byte of the object it commits, so the ranges + // have to tile it — being merely disjoint and in bounds would leave gaps of + // whatever the slot happened to hold. + std::vector valid; + std::vector route_keys; + std::vector route_sizes; + valid.reserve(n); + route_keys.reserve(n); + route_sizes.reserve(n); + auto validate_timer = std::make_unique(dbg ? &dbg->validate : nullptr); + for (size_t i = 0; i < n; ++i) { + if (!RangesTileObject(object_sizes[i], sizes[i], dst_offsets[i])) { + MORI_UMBP_ERROR("[PoolClient] BatchPutRanges: ranges do not tile key='{}' size={}", keys[i], + object_sizes[i]); + continue; + } + valid.push_back(i); + route_keys.push_back(keys[i]); + route_sizes.push_back(object_sizes[i]); + } + validate_timer.reset(); + if (valid.empty()) return results; + + std::vector> routes; std::unordered_set excludes; - auto status = master_client_->BatchRouteGet(keys, excludes, &routes); - if (!status.ok()) { - MORI_UMBP_ERROR("[PoolClient] BatchGet: BatchRouteGet failed: {}", status.error_message()); + // Unconditional, unlike the get path, which routes only what it missed + // locally. Every put pays this RPC even when every key lands on this node. + grpc::Status route_status; + { + PhaseTimer route_timer(dbg ? &dbg->route : nullptr); + if (!HasMaster()) { + RouteAllPutsLocally(route_keys.size(), &routes); + } else { + route_status = master_client_->BatchRoutePut(route_keys, route_sizes, excludes, &routes); + } + } + if (!route_status.ok()) { + MORI_UMBP_ERROR("[PoolClient] BatchPutRanges: BatchRoutePut failed: {}", + route_status.error_message()); return results; } - if (routes.size() < keys.size()) { - routes.resize(keys.size()); + routes.resize(valid.size()); + + std::vector remote; + std::vector local_writes; + remote.reserve(valid.size()); + local_writes.reserve(valid.size()); + for (size_t r = 0; r < valid.size(); ++r) { + const size_t original = valid[r]; + if (!routes[r].has_value()) continue; + const auto& route = *routes[r]; + if (route.outcome == RoutePutOutcome::kAlreadyExists) { + results[original] = true; + continue; + } + if (route.node_id != config_.master_config.node_id) { + remote.push_back(r); + continue; + } + // Routed here: write the ranges straight into the slot's pages. No + // assembly buffer — a TransferItem is a range, so the scattered sources + // land where they belong. + local_writes.push_back({original, &keys[original], object_sizes[original], + MakeWriteRanges(srcs[original], sizes[original], dst_offsets[original]), + route.tier, route.logical_tier}); + } + // One transfer for every locally-routed key in the batch, not one per key. + if (dbg != nullptr) { + dbg->local_keys = local_writes.size(); + dbg->remote_keys = remote.size(); } + { + PhaseTimer xfer_timer(dbg ? &dbg->xfer : nullptr); + TransferEngine::StepTiming steps; + RangedPhaseSinks sinks; + if (dbg != nullptr) { + steps = {&dbg->xfer_plan, &dbg->xfer_submit, &dbg->xfer_wait}; + sinks = {&dbg->resolve, &dbg->classify, &dbg->build, &dbg->commit, &dbg->items, &steps}; + } + ExecuteLocalPutRangesBatch(local_writes, &results, &local_bytes, + dbg != nullptr ? &sinks : nullptr); + } + if (dbg != nullptr) { + // The sub-phases above are all timed inside the xfer window, so xfer alone + // is left holding the engine's copy. Subtract rather than re-time: nesting + // is what makes each sub-phase attributable to the call it belongs to. + dbg->xfer -= dbg->resolve + dbg->classify + dbg->build + dbg->commit; + } + if (remote.empty()) return results; + + // A key routed to another node has to be one contiguous object on the wire, + // so this is the one direction that needs the arena. PUT has its own arena + // and mutex, so it overlaps a concurrent remote GET (which uses the separate + // GET arena/mutex). + char* const scratch_base = static_cast(config_.ranged_put_scratch_buffer); + const size_t scratch_size = config_.ranged_put_scratch_size; + if (scratch_base == nullptr || scratch_size == 0 || + !FindRegisteredMemory(scratch_base, scratch_size).has_value()) { + MORI_UMBP_ERROR("[PoolClient] BatchPutRanges: ranged scratch is absent or not registered"); + return results; + } + std::unique_lock scratch_lock(ranged_put_scratch_mutex_, std::defer_lock); + { + PhaseTimer lock_timer(dbg ? &dbg->lock : nullptr); + scratch_lock.lock(); + } + + PhaseTimer remote_timer(dbg ? &dbg->remote : nullptr); + size_t pos = 0; + while (pos < remote.size()) { + std::vector scratch_offsets; + size_t cursor = 0; + size_t end = pos; + for (; end < remote.size(); ++end) { + const size_t object_size = object_sizes[valid[remote[end]]]; + size_t aligned = 0; + if (!AlignUpChecked(cursor, kRangedScratchAlignment, &aligned) || + object_size > scratch_size || aligned > scratch_size - object_size) { + break; + } + scratch_offsets.push_back(aligned); + cursor = aligned + object_size; + } + if (end == pos) { + const size_t original = valid[remote[pos]]; + MORI_UMBP_ERROR( + "[PoolClient] BatchPutRanges: object exceeds ranged scratch key='{}' size={} scratch={}", + keys[original], object_sizes[original], scratch_size); + ++pos; + continue; + } - BatchGetPlan plan = PartitionBatchGetTargets(keys, dsts, sizes, routes); - ExecuteBatchGetPlan(plan, keys, dsts, sizes, &results); + const size_t count = end - pos; + std::vector sub_keys; + std::vector sub_srcs; + std::vector sub_sizes; + std::vector> sub_routes; + std::vector sub_originals; + sub_keys.reserve(count); + sub_srcs.reserve(count); + sub_sizes.reserve(count); + sub_routes.reserve(count); + sub_originals.reserve(count); + // Assemble EVERY object in this arena round with ONE transfer. + // + // This used to call CopyRangesToContiguous per key, and each of those is a + // blocking Transfer -- its own kernel launch and hipStreamSynchronize. On + // an 8-rank embedded deployment that showed up as 369 transfers carrying + // 369 segments, 14.5 us apiece for ~5 us of actual copy: ~65% of remote + // transfer time was per-call overhead. The local path + // (ExecuteLocalPutRangesBatch) already accumulates across requests and + // transfers once; this makes the remote path match it. + // + // Tagged by j so a single unassemblable object fails alone, which is what + // the per-key call gave us for free. + std::vector assembly_items; + std::vector assembly_tag_to_j; + assembly_items.reserve(count); + assembly_tag_to_j.reserve(count); + for (size_t j = 0; j < count; ++j) { + const size_t original = valid[remote[pos + j]]; + void* slice = scratch_base + scratch_offsets[j]; + const size_t tag = assembly_tag_to_j.size(); + if (!BuildRangesToContiguousItems( + MakeWriteRanges(srcs[original], sizes[original], dst_offsets[original]), slice, + object_sizes[original], tag, &assembly_items)) { + MORI_UMBP_ERROR("[PoolClient] BatchPutRanges: assembly plan failed for key='{}'", + keys[original]); + continue; + } + assembly_tag_to_j.push_back(j); + } + std::unordered_set assembly_failed; + if (!assembly_items.empty()) { + std::vector failed_tags; + transfer_engine_->Transfer(assembly_items, &failed_tags); + for (size_t tag : failed_tags) { + if (tag < assembly_tag_to_j.size()) assembly_failed.insert(assembly_tag_to_j[tag]); + } + } + std::unordered_set assembled(assembly_tag_to_j.begin(), assembly_tag_to_j.end()); + + for (size_t j = 0; j < count; ++j) { + const size_t route_index = remote[pos + j]; + const size_t original = valid[route_index]; + void* slice = scratch_base + scratch_offsets[j]; + if (assembled.count(j) == 0 || assembly_failed.count(j) != 0) { + MORI_UMBP_ERROR("[PoolClient] BatchPutRanges: assembly failed for key='{}'", + keys[original]); + continue; + } + sub_keys.push_back(keys[original]); + sub_srcs.push_back(slice); + sub_sizes.push_back(object_sizes[original]); + sub_routes.push_back(routes[route_index]); + sub_originals.push_back(original); + } - const auto call_end = std::chrono::steady_clock::now(); - const double seconds = - std::chrono::duration_cast>(call_end - call_start).count(); - if (seconds > 0.0) { - auto split = ComputeBatchBandwidthBytes(results, sizes, routes, config_.master_config.node_id); - ObserveBatchBandwidth(*master_client_, split.local, seconds, - MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH, - MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH_HELP, "local"); - ObserveBatchBandwidth(*master_client_, split.remote, seconds, - MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH, - MORI_UMBP_METRIC_CLIENT_BATCH_GET_BANDWIDTH_HELP, "remote"); + if (!sub_keys.empty()) { + std::vector outcomes(sub_keys.size(), PutEntryOutcome::kFailed); + BatchPutPlan plan = + PartitionBatchPutTargets(sub_keys, sub_srcs, sub_sizes, sub_routes, &outcomes); + ExecuteBatchPutPlan(plan, &outcomes); + for (size_t j = 0; j < outcomes.size(); ++j) { + results[sub_originals[j]] = outcomes[j] != PutEntryOutcome::kFailed; + // kAlreadyExists is success to the caller but moved nothing, matching + // IsCountedForBandwidth's rule for the whole-object BatchPut family. + if (outcomes[j] == PutEntryOutcome::kSucceeded) { + remote_bytes += static_cast(sub_sizes[j]); + } + } + } + pos = end; + } + if (workload_recorder_ != nullptr) { + std::vector recorded(results.size()); + for (size_t i = 0; i < results.size(); ++i) { + recorded[i] = results[i] ? benchmark::WorkloadTraceOutcome::kSuccess + : benchmark::WorkloadTraceOutcome::kFailure; + } + workload_recorder_->RecordBatchPut(keys, object_sizes, recorded); } return results; } PoolClient::BatchGetPlan PoolClient::PartitionBatchGetTargets( const std::vector& keys, const std::vector& dsts, - const std::vector& sizes, const std::vector>& routes) { + const std::vector& sizes, const std::vector>& routes, + LocalFallback fallback) { + const bool allow_local = fallback == LocalFallback::kAllow && !registry_.Empty(); BatchGetPlan plan; for (size_t i = 0; i < keys.size(); ++i) { // Zero-size gets are rejected before local fallback or remote read: an @@ -1592,145 +3801,241 @@ PoolClient::BatchGetPlan PoolClient::PartitionBatchGetTargets( MORI_UMBP_WARN("[PoolClient] BatchGet: skipping zero-size get for key='{}'", keys[i]); continue; } - if (i >= routes.size() || !routes[i].has_value()) { - if (peer_alloc_) plan.local_dram_indices.push_back(i); + // No route, or one pointing back at this node. Under kAllow that is a + // local read the caller has not attempted yet -- deferred as an index so + // ExecuteBatchGetPlan can run it inside the remote in-flight window. Under + // kSkip the local media were already asked, and the key is either served + // (routes[i] stayed nullopt) or genuinely absent; either way it is done. + if (i >= routes.size() || !routes[i].has_value() || + routes[i]->node_id == config_.master_config.node_id) { + if (allow_local) plan.local_indices.push_back(i); continue; } - const auto& route = routes[i].value(); - if (route.node_id == config_.master_config.node_id) { - // Self-target: both DRAM/HBM and SSD are deferred (collected as indices) - // so ExecuteBatchGetPlan can place them inside the remote-DRAM in-flight - // window in the overlap path. - if (route.tier == TierType::SSD) { - plan.local_ssd_indices.push_back(i); - } else { - plan.local_dram_indices.push_back(i); - } - continue; - } - BatchGetItem item{.index = i, - .key = &keys[i], - .dst = const_cast(dsts[i]), - .size = sizes[i], - .route = route}; - if (route.tier == TierType::SSD) { - plan.remote_ssd_groups[route.node_id].push_back(std::move(item)); - } else { - plan.remote_dram_groups[route.node_id].push_back(std::move(item)); - } + const auto& route = *routes[i]; + plan.remote_groups[route.node_id].push_back(BatchGetItem{ + .index = i, .key = &keys[i], .dst = dsts[i], .size = sizes[i], .route = route}); } return plan; } -void PoolClient::ExecuteBatchGetPlan(const BatchGetPlan& plan, const std::vector& keys, - const std::vector& dsts, - const std::vector& sizes, std::vector* results) { - // Parallel local DRAM reads: different threads handle different keys. Resolve - // is mutex-serialized in the allocator; the per-key memcpy in - // ExecuteLocalGet->LocalGetPages runs lock-free in parallel. results is - // std::vector (bit-packed) so threads write a temp buffer; merge serially. - auto run_local_dram = [&]() { - const auto& idx = plan.local_dram_indices; - if (idx.empty()) return; - const int nthr = LocalCopyThreads("UMBP_DRAM_READ_THREADS"); - const auto t0 = std::chrono::steady_clock::now(); - std::vector ok(idx.size(), 0); - LocalParallelFor(idx.size(), nthr, [&](size_t k) { - const size_t i = idx[k]; - if (ExecuteLocalGet(keys[i], const_cast(dsts[i]), sizes[i]) == - GetAttemptOutcome::kSuccess) { - ok[k] = 1; - } - }); - size_t tot = 0; - for (size_t k = 0; k < idx.size(); ++k) { - if (ok[k]) { - (*results)[idx[k]] = true; - tot += sizes[idx[k]]; - } +std::vector PoolClient::ServeWholeObjectUnitsFromMedium( + const std::vector& keys, std::vector units, + std::unordered_map* unit_ok, double* remote_bytes) { + auto* backend = registry_.Get(medium_); + std::vector leftover; + leftover.reserve(units.size()); + + // Only a unit that is the whole object can go straight into a slot. No + // switch and no size cap: the bytes are crossing the wire either way, so + // landing them in a slot instead of the arena costs nothing extra and is what + // the pre-ranged code did on every remote ranged read. + const bool can_install = CanInstallLocally(); + std::vector candidates; + for (size_t i = 0; i < units.size(); ++i) { + if (can_install && units[i].holds_whole_object) { + candidates.push_back(i); + } else { + leftover.push_back(std::move(units[i])); } - if (std::getenv("UMBP_LOCAL_COPY_TIMING")) { - double sec = std::chrono::duration_cast>( - std::chrono::steady_clock::now() - t0) - .count(); - MORI_UMBP_INFO("[LocalCopy] GET keys={} bytes={} threads={} elapsed_ms={:.3f} GiB_s={:.2f}", - idx.size(), tot, nthr, sec * 1000.0, - tot / (sec > 0 ? sec : 1e-12) / (1024.0 * 1024 * 1024)); + } + if (candidates.empty()) return leftover; + + // Already here? Then phase 1 raced us to it; hand the unit back rather than + // allocate a second slot for a key the medium already holds. + std::vector candidate_keys; + candidate_keys.reserve(candidates.size()); + for (size_t i : candidates) candidate_keys.push_back(keys[units[i].original]); + const auto resolved = backend->BatchResolve(candidate_keys, /*include_descs=*/false); + + std::vector wanted; + std::vector requests; + wanted.reserve(candidates.size()); + requests.reserve(candidates.size()); + for (size_t c = 0; c < candidates.size(); ++c) { + if (c < resolved.size() && resolved[c].found) { + leftover.push_back(std::move(units[candidates[c]])); + continue; } - }; + wanted.push_back(candidates[c]); + requests.push_back(AllocateRequest{candidate_keys[c], units[candidates[c]].object_size}); + } + if (requests.empty()) return leftover; - // Local SSD self-target reads (deferred from partition): serial on this - // thread, reading straight into the user buffer (no staging / RDMA / lease). - // Independent per key, so its position in the schedule is correctness-neutral. - auto run_local_ssd = [&]() { - for (size_t i : plan.local_ssd_indices) { - if (ExecuteLocalSsdGet(keys[i], const_cast(dsts[i]), sizes[i]) == - GetAttemptOutcome::kSuccess) { - (*results)[i] = true; - } + auto allocs = backend->BatchAllocate(requests); + + // slot_refs and slot_offsets are reserved up front and never grow: the plan + // holds pointers into the first. + BatchGetPlan plan; + std::vector slot_refs; + std::vector slot_offsets; + std::vector planned; // index into wanted/allocs + slot_refs.reserve(wanted.size()); + slot_offsets.reserve(wanted.size()); + planned.reserve(wanted.size()); + + for (size_t w = 0; w < wanted.size(); ++w) { + auto& alloc = allocs[w]; + RangeFetchUnit& unit = units[wanted[w]]; + if (alloc.outcome != AllocateOutcome::kSuccessAllocated) { + leftover.push_back(std::move(unit)); // medium is full; the arena still works + continue; } - }; - auto run_remote_ssd = [&]() { - for (const auto& [node_id, items] : plan.remote_ssd_groups) { - ProcessRemoteSsdBatchGet(items, results); + // The remote get path names its destination as one contiguous span, so the + // slot has to be one. The page allocator already prefers a same-buffer + // continuous run, and distributed mode stores one key per page, so the + // common cases are covered. + bool contiguous = !alloc.pages.empty(); + for (size_t q = 1; q < alloc.pages.size() && contiguous; ++q) { + contiguous = alloc.pages[q].buffer_index == alloc.pages.front().buffer_index && + alloc.pages[q].page_index == alloc.pages.front().page_index + q; + } + TransferRef slot_buffer = + contiguous ? backend->BufferRef(alloc.pages.front().buffer_index) : TransferRef{}; + if (!contiguous || !slot_buffer.Valid() || !slot_buffer.HasHostPtr()) { + backend->BatchAbort({alloc.slot_id}); + NoteRangedInstallFailure(); // the arena will still serve it, but not cache it + leftover.push_back(std::move(unit)); + continue; } - }; - const auto& remote_dram_groups = plan.remote_dram_groups; - - // Zero-copy? The upper layer guarantees a batch is all-ZC or all-staging, so - // probe one remote-DRAM item's dst registration (unordered_map order is - // unspecified, but any item answers the all-or-nothing question). - const bool is_zc = !remote_dram_groups.empty() && - FindRegisteredMemory(remote_dram_groups.begin()->second.front().dst, - remote_dram_groups.begin()->second.front().size) - .has_value(); - - if (!is_zc) { - // Staging (non-zero-copy), or no remote DRAM at all. Order: local SSD, local - // DRAM, staging remote DRAM (per peer), remote SSD. Staging must submit then - // IMMEDIATELY wait per peer: the in-flight holds staging_mutex_ from submit - // until the wait copies out of the shared staging_buffer_, so accumulating - // staging in-flights into a submit-all list would deadlock on that lock. - run_local_ssd(); - run_local_dram(); - for (const auto& [node_id, items] : remote_dram_groups) { - if (auto f = SubmitRemoteBatchGet(items, results, /*permit_staging=*/true)) { - WaitRemoteBatchGet(*f, results); // immediate; releases staging_mutex_ + const uint64_t slot_offset = PageOffset(alloc.pages.front(), alloc.page_size); + slot_offsets.push_back(slot_offset); + slot_refs.push_back(std::move(slot_buffer)); + // Named by the backend's own ref rather than a raw pointer: the medium pool + // is RDMA-reachable but was never handed to RegisterMemory, so + // UserBufferRef would not find it. + plan.remote_groups[unit.route->node_id].push_back(BatchGetItem{.index = planned.size(), + .key = &keys[unit.original], + .dst = nullptr, + .size = unit.object_size, + .dst_bytes = unit.bytes, + .spans = &unit.spans, + .dst_ref = &slot_refs.back(), + .dst_ref_offset = slot_offset, + .route = *unit.route}); + planned.push_back(w); + } + if (planned.empty()) return leftover; + + std::vector fetched(planned.size(), false); + ExecuteRemoteBatchGetPlan(plan, &fetched, /*recache_remote=*/false); + + std::vector commits; + std::vector aborts; + commits.reserve(planned.size()); + for (size_t i = 0; i < planned.size(); ++i) { + auto& alloc = allocs[planned[i]]; + RangeFetchUnit& unit = units[wanted[planned[i]]]; + if (!fetched[i]) { + // The bytes never arrived, so there is nothing worth keeping. Not a + // fallback to the arena: see the note on the declaration. + aborts.push_back(alloc.slot_id); + (*unit_ok)[unit.original] = false; + continue; + } + // Copy out BEFORE committing. A pending slot cannot be reclaimed; a + // committed one is an ordinary evictable object, and reading it after that + // races the allocator. + // The backend's own ref, not a bare pointer: an HBM pool hands back a + // device pointer, and re-wrapping it as host bytes makes the engine copy + // in the wrong direction (or memcpy device memory). + if (!CopyContiguousToRanges(slot_refs[i], slot_offsets[i], unit.object_size, unit.packed)) { + (*unit_ok)[unit.original] = false; + } + // Commit either way: the slot holds the object whether or not the caller's + // copy-out worked, and caching it still helps the next reader. + commits.push_back(CommitRequest{alloc.slot_id, keys[unit.original]}); + } + if (!commits.empty()) { + const auto committed = backend->BatchCommit(commits); + size_t refused = 0; + for (size_t c = 0; c < committed.size(); ++c) { + if (!committed[c].success) { + aborts.push_back(commits[c].slot_id); + ++refused; } } - run_remote_ssd(); - return; + NoteRangedInstallFailure(refused); } + if (!aborts.empty()) backend->BatchAbort(aborts); + return leftover; +} - // Zero-copy remote DRAM: submit every peer (posted, not waited) to overlap - // wire time across peers, run local DRAM/SSD + remote SSD in that window, then - // wait all. On early/exceptional exit ~RemoteDramGetInFlight drains in-flight - // statuses (lifetime safety); the wait loop does failure mapping + backfill. - std::vector> inflights; - inflights.reserve(remote_dram_groups.size()); - - for (const auto& [node_id, items] : remote_dram_groups) { - if (auto f = SubmitRemoteBatchGet(items, results, /*permit_staging=*/false)) { - inflights.push_back(std::move(f)); +PoolClient::BatchGetPlan PoolClient::PartitionBatchGetRangeTargets( + const std::vector& keys, const std::vector& arena_slices, + const std::vector& object_sizes, const std::vector& packed_bytes, + const std::vector*>& spans, + const std::vector>& routes) { + BatchGetPlan plan; + for (size_t i = 0; i < keys.size(); ++i) { + // Unroutable and self-routed keys were filtered by the caller, which is + // what makes this plan remote-only. Anything that slips through is a + // caller bug, not a local read to fall back to. + if (i >= routes.size() || !routes[i].has_value() || packed_bytes[i] == 0 || + spans[i] == nullptr || spans[i]->empty()) { + continue; } + plan.remote_groups[routes[i]->node_id].push_back(BatchGetItem{.index = i, + .key = &keys[i], + .dst = arena_slices[i], + .size = object_sizes[i], + .dst_bytes = packed_bytes[i], + .spans = spans[i], + .route = *routes[i]}); } - run_local_dram(); - run_local_ssd(); - // TODO(perf): remote SSD is still per-key serial (not yet optimized). - run_remote_ssd(); - for (auto& f : inflights) WaitRemoteBatchGet(*f, results); + return plan; } -std::unique_ptr PoolClient::SubmitRemoteBatchGet( - const std::vector& items, std::vector* results, bool permit_staging) { +void PoolClient::ExecuteBatchGetPlan(const BatchGetPlan& plan, const std::vector& keys, + const std::vector& dsts, + const std::vector& sizes, std::vector* results, + bool recache_remote) { + // The local half, run inside the remote submit..wait window so it overlaps + // the wire. One batched resolve and one Transfer -- the same core the + // local-first arm uses; nothing left to route, so misses simply stay false. + auto run_local = [&]() { + ServeLocalGets(keys, dsts, sizes, plan.local_indices, results, /*missed=*/nullptr); + }; + + ExecuteRemoteBatchGetPlan(plan, results, recache_remote, run_local); +} + +void PoolClient::ExecuteRemoteBatchGetPlan(const BatchGetPlan& plan, std::vector* results, + bool recache_remote, + const std::function& in_flight_window) { + // Submit every peer (posted, not waited) to overlap wire time across peers, + // run whatever the caller wants overlapped in that window, then wait all. On + // early/exceptional exit the engine handle's destructor drains in-flight + // statuses (lifetime safety); the wait loop does failure mapping + backfill. + // + // As in Put, there is no all-zero-copy / all-staging fork any more: staging + // lives in the engine's bounce pool and a plan that needs it settles inside + // Submit, so submit-all is unconditionally safe. + std::vector> inflights; + inflights.reserve(plan.remote_groups.size()); + for (const auto& [node_id, items] : plan.remote_groups) { + if (auto f = SubmitRemoteBatchGet(items, results)) inflights.push_back(std::move(f)); + } + if (in_flight_window) in_flight_window(); + for (auto& f : inflights) WaitRemoteBatchGet(*f, results, recache_remote); +} + +void PoolClient::NoteRangedInstallFailure(size_t count) { + if (count == 0) return; + CountMetric(MORI_UMBP_METRIC_RANGED_REMOTE_INSTALL_FAILURES_TOTAL, + MORI_UMBP_METRIC_RANGED_REMOTE_INSTALL_FAILURES_TOTAL_HELP, {}, + static_cast(count)); +} + +std::unique_ptr PoolClient::SubmitRemoteBatchGet( + const std::vector& items, std::vector* results) { if (items.empty()) return nullptr; auto fail_all = [&] { for (const auto& item : items) (*results)[item.index] = false; }; - if (!io_engine_) { - MORI_UMBP_ERROR("[PoolClient] SubmitRemoteBatchGet: io_engine_ not initialized (items={})", + if (peer_directory_ == nullptr) { + MORI_UMBP_ERROR("[PoolClient] SubmitRemoteBatchGet: no RDMA engine configured (items={})", items.size()); fail_all(); return nullptr; @@ -1748,7 +4053,7 @@ std::unique_ptr PoolClient::SubmitRemoteBatch } auto* stub = static_cast<::umbp::UMBPPeer::Stub*>(peer.peer_stub.get()); - auto inflight = std::make_unique(); + auto inflight = std::make_unique(); inflight->peer = &peer; // resolve RPC + per-key validation; failed keys already written to *results. @@ -1756,124 +4061,56 @@ std::unique_ptr PoolClient::SubmitRemoteBatch return nullptr; } - std::vector transfers; - uint64_t staging_bytes = 0; - if (!BuildRemoteGetTransfers(inflight->entries, peer, &transfers, &staging_bytes)) { + std::vector transfer_items; + if (!BuildRemoteGetTransfers(inflight->entries, first.route.node_id, &transfer_items)) { MORI_UMBP_WARN( "[PoolClient] SubmitRemoteBatchGet: BuildRemoteGetTransfers failed, node='{}' entries={}", first.route.node_id, inflight->entries.size()); for (auto& entry : inflight->entries) (*results)[entry.result_index] = false; return nullptr; } - inflight->staging_bytes = staging_bytes; - // Staging needs the shared staging_buffer_, which a cross-peer submit-all - // window cannot serialize on the single staging_mutex_ without deadlocking. - // So the zero-copy submit-all caller passes permit_staging=false: a batch that - // unexpectedly needs staging (upper layer guarantees all-zc or all-staging) is - // failed rather than risking a deadlock. The serial staging caller passes - // true and we park staging_mutex_ in the in-flight until the wait copies out. - if (staging_bytes > 0) { - if (!permit_staging) { - MORI_UMBP_ERROR( - "[PoolClient] SubmitRemoteBatchGet: unexpected staging_bytes={} on zero-copy path " - "node='{}' -> failing batch", - staging_bytes, first.route.node_id); - for (auto& entry : inflight->entries) (*results)[entry.result_index] = false; - return nullptr; + + // Drop items whose entry failed during build (peer published no buffer etc.). + std::vector active; + active.reserve(transfer_items.size()); + for (auto& item : transfer_items) { + if (!inflight->entries[item.tag].failed) active.push_back(std::move(item)); + } + if (active.empty()) { + for (auto& entry : inflight->entries) { + if (entry.failed) (*results)[entry.result_index] = false; } - inflight->staging_lock = std::unique_lock(staging_mutex_); + return nullptr; } - // Drop transfers whose entry failed during build (invalid remote desc etc.). - std::vector active_transfers; - active_transfers.reserve(transfers.size()); - for (const auto& t : transfers) { - if (!inflight->entries[t.entry_index].failed) active_transfers.push_back(t); - } - if (active_transfers.empty()) { + TransferPlanSet planned = transfer_engine_->Plan(active); + ApplyRejectedTags(inflight->entries, planned.rejected_tags, "RemoteGet"); + if (planned.plans.empty()) { for (auto& entry : inflight->entries) { if (entry.failed) (*results)[entry.result_index] = false; } - return nullptr; // staging_lock (if held) released by the in-flight's destruction - } - - // Collapse same-(localMR, remoteMR) pages into one outer transfer per pair, - // then materialise the BatchRead args INTO the in-flight struct so they outlive - // the post. - inflight->groups = GroupTransfersByPair(active_transfers); - const size_t G = inflight->groups.size(); - inflight->local_descs.resize(G); - inflight->remote_descs.resize(G); - inflight->local_offsets.resize(G); - inflight->remote_offsets.resize(G); - inflight->sizes_v.resize(G); - // Build statuses at final size once and never resize/move (backend holds raw - // TransferStatus*); status_ptrs alias into it. vector move-assign steals the - // buffer (default allocator), so elements are not moved and addresses are - // stable for the in-flight lifetime. - inflight->statuses = std::vector(G); - inflight->status_ptrs.resize(G); - inflight->ids.resize(G); - for (size_t g = 0; g < G; ++g) { - inflight->local_descs[g] = inflight->groups[g].local_desc; - inflight->remote_descs[g] = inflight->groups[g].remote_desc; - inflight->local_offsets[g] = std::move(inflight->groups[g].local_offsets); - inflight->remote_offsets[g] = std::move(inflight->groups[g].remote_offsets); - inflight->sizes_v[g] = std::move(inflight->groups[g].sizes); - inflight->status_ptrs[g] = &inflight->statuses[g]; - inflight->ids[g] = io_engine_->AllocateTransferUniqueId(); - } - - // POST the reads; do NOT wait. Everything BatchRead may reference is owned by - // *inflight; no allocation happens after this point in the submit window. For - // staging, the RDMA lands in staging_buffer_ under the held staging_lock. - io_engine_->BatchRead(inflight->local_descs, inflight->local_offsets, inflight->remote_descs, - inflight->remote_offsets, inflight->sizes_v, inflight->status_ptrs, - inflight->ids); + return nullptr; + } + // POST; do NOT wait. Everything the post references is owned by the returned + // handle, including any bytes staged through the engine's bounce pool. + inflight->handle = transfer_engine_->Submit(std::move(planned.plans)); + if (inflight->handle == nullptr) { + for (auto& entry : inflight->entries) (*results)[entry.result_index] = false; + return nullptr; + } return inflight; } -void PoolClient::WaitRemoteBatchGet(RemoteDramGetInFlight& f, std::vector* results) { +void PoolClient::WaitRemoteBatchGet(RemoteGetInFlight& f, std::vector* results, + bool recache_remote) { if (f.drained) return; f.drained = true; - // Wait every group (never break early): IN_PROGRESS groups complete here; - // INIT/terminal return immediately. A non-success group fails all of its - // contributing keys (per-item AND). - const size_t G = f.groups.size(); - for (size_t g = 0; g < G; ++g) { - f.statuses[g].Wait(); - if (f.statuses[g].Succeeded()) continue; - for (size_t ei : f.groups[g].entry_indices) { - auto& entry = f.entries[ei]; - MORI_UMBP_ERROR( - "RemoteGet BatchRead failed: code={} msg='{}' peer_engine='{}' key='{}' use_staging={}", - f.statuses[g].CodeUint32(), f.statuses[g].Message(), f.groups[g].remote_desc.engineKey, - (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, - entry.use_staging); - entry.failed = true; - } - } - - // Staging (non-zero-copy): RDMA landed in staging_buffer_; copy each surviving - // entry out to its user dst, then release staging_mutex_ (held since submit). - if (f.staging_lock.owns_lock()) { - for (auto& entry : f.entries) { - if (entry.failed || !entry.use_staging) continue; - if (entry.staging_offset + entry.item->size > config_.staging_buffer_size) { - MORI_UMBP_ERROR( - "[PoolClient] WaitRemoteBatchGet: staging offset overflow (should not happen), " - "key='{}' staging_offset={} size={} cap={}", - (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, - entry.staging_offset, entry.item ? entry.item->size : 0, config_.staging_buffer_size); - entry.failed = true; - continue; - } - std::memcpy(entry.item->dst, staging_buffer_.get() + entry.staging_offset, entry.item->size); - } - f.staging_lock.unlock(); - } - - FinalizeRemoteGetEntries(f.entries, results); + std::vector failures; + if (f.handle != nullptr) f.handle->Wait(&failures); + ApplyTransferFailures(f.entries, failures, "RemoteGet"); + // Nothing to copy out here even for a staged read: the engine owns its bounce + // pool and lands the bytes in the user's dst before its Wait returns. + FinalizeRemoteGetEntries(f.entries, results, recache_remote); } bool PoolClient::PrepareRemoteGetEntries(const std::vector& items, @@ -1886,49 +4123,77 @@ bool PoolClient::PrepareRemoteGetEntries(const std::vector& items, // them (from the GetPeerInfo handshake, or a prior resolve). A wrong guess // is safe: a missing descriptor is caught by the transfer-build guard and the // entry degrades to a miss, never a corrupt read. - bool have_descs = false; - { - std::lock_guard lock(peers_mutex_); - have_descs = !peer.dram_memories.empty(); - } + const bool have_descs = + peer_directory_ != nullptr && peer_directory_->HasRemoteBuffers(peer.node_id); ::umbp::BatchResolveKeysRequest resolve_req; for (const auto& item : items) resolve_req.add_keys(*item.key); resolve_req.set_omit_descs(have_descs); - ::umbp::BatchResolveKeysResponse resolve_resp; - grpc::ClientContext resolve_ctx; - auto resolve_status = stub->BatchResolveKeys(&resolve_ctx, resolve_req, &resolve_resp); - if (!resolve_status.ok() || - BatchResolveKeyCount(resolve_resp) != static_cast(items.size())) { - MORI_UMBP_WARN("[PoolClient] BatchResolveKeys failed on {}: {}", items.front().route.node_id, - resolve_status.error_message()); - for (const auto& item : items) { - (*results)[item.index] = false; + DecodedBatchResolve decoded; + const auto retry_timeout = ResolveBusyRetryTimeout(); + const auto retry_deadline = std::chrono::steady_clock::now() + retry_timeout; + std::chrono::milliseconds backoff{1}; + size_t busy_attempts = 0; + while (true) { + ::umbp::BatchResolveKeysResponse resolve_resp; + grpc::ClientContext resolve_ctx; + const auto remaining = retry_deadline - std::chrono::steady_clock::now(); + if (remaining <= std::chrono::steady_clock::duration::zero()) { + MORI_UMBP_WARN("[PoolClient] BatchResolveKeys BUSY timeout on {} after {} attempts", + items.front().route.node_id, busy_attempts); + for (const auto& item : items) (*results)[item.index] = false; + return false; + } + resolve_ctx.set_deadline(std::chrono::system_clock::now() + remaining); + auto resolve_status = stub->BatchResolveKeys(&resolve_ctx, resolve_req, &resolve_resp); + if (!resolve_status.ok() || + BatchResolveKeyCount(resolve_resp) != static_cast(items.size())) { + MORI_UMBP_WARN("[PoolClient] BatchResolveKeys failed on {}: {}", items.front().route.node_id, + resolve_status.error_message()); + for (const auto& item : items) (*results)[item.index] = false; + return false; } - return false; - } - DecodedBatchResolve decoded = DecodeBatchResolveResponse(resolve_resp); - if (decoded.keys.size() != items.size()) { - // Malformed (mismatched parallel arrays); fail the whole batch rather than - // partially-read it. - MORI_UMBP_WARN("[PoolClient] BatchResolveKeys malformed response on {}: {} keys for {} items", - items.front().route.node_id, decoded.keys.size(), items.size()); - for (const auto& item : items) { - (*results)[item.index] = false; + decoded = DecodeBatchResolveResponse(resolve_resp); + if (decoded.keys.size() != items.size()) { + // Malformed (mismatched parallel arrays); fail the whole batch rather + // than partially reading it. + MORI_UMBP_WARN("[PoolClient] BatchResolveKeys malformed response on {}: {} keys for {} items", + items.front().route.node_id, decoded.keys.size(), items.size()); + for (const auto& item : items) (*results)[item.index] = false; + return false; } - return false; + + const bool busy = std::any_of(decoded.keys.begin(), decoded.keys.end(), [](const auto& key) { + return key.outcome == ResolveOutcome::kBusy; + }); + if (!busy) break; + + // Discard the ENTIRE response. Keeping its successful SSD entries while + // retrying only BUSY keys would pin their leases and can make forward + // progress impossible; it can also leave stale page locations if a retry + // outlives the lease. + ++busy_attempts; + const auto sleep_for = + std::min(backoff, std::chrono::duration_cast( + retry_deadline - std::chrono::steady_clock::now())); + if (sleep_for.count() > 0) std::this_thread::sleep_for(sleep_for); + backoff = std::min(backoff * 2, std::chrono::milliseconds{50}); } // Hydrate the batch-level descriptors once (skipped when the peer honored // omit_descs and sent none). - if (!decoded.descs.empty()) EnsureBufferDescsCached(peer, decoded.descs); + if (!decoded.descs.empty()) peer_directory_->CacheRemoteBuffers(peer.node_id, decoded.descs); entries->reserve(items.size()); for (size_t i = 0; i < items.size(); ++i) { const auto& item = items[i]; const auto& key = decoded.keys[i]; if (!key.found) { + if (key.outcome == ResolveOutcome::kFailed) { + MORI_UMBP_ERROR("[PoolClient] BatchGet: peer reported permanent resolve failure key='{}'", + *item.key); + } (*results)[item.index] = false; continue; } @@ -1948,6 +4213,9 @@ bool PoolClient::PrepareRemoteGetEntries(const std::vector& items, entry.result_index = item.index; entry.item = &item; entry.plan.page_size = decoded.page_size; + // Per key, because a resolve batch may now be served from several media at + // once; the pages below are indices into THIS backend's buffers. + entry.plan.backend_id = key.backend_id; entry.plan.pages = std::move(decoded.keys[i].pages); // Descriptors were hydrated batch-level above; the per-entry plan carries // none (BuildRemoteGetTransfers' EnsureBufferDescsCached call is a no-op on @@ -1958,174 +4226,167 @@ bool PoolClient::PrepareRemoteGetEntries(const std::vector& items, return !entries->empty(); } -bool PoolClient::BuildRemoteGetTransfers(std::vector& entries, PeerConnection& peer, - std::vector* transfers, - uint64_t* staging_bytes) { - transfers->clear(); - uint64_t cursor = 0; +bool PoolClient::BuildRemoteGetTransfers(std::vector& entries, + const std::string& node_id, + std::vector* items) { + items->clear(); + + // Batch-level descs were already hydrated by PrepareRemoteGetEntries; any + // per-entry ones are folded in here before the snapshots (see + // BuildRemotePutTransfers for why one snapshot per backend beats a lock per + // page — and why one snapshot per PEER would index the wrong medium now that + // a resolve batch can span media). + for (const auto& entry : entries) { + if (!entry.plan.descs.empty()) peer_directory_->CacheRemoteBuffers(node_id, entry.plan.descs); + } + std::array, kMaxBackendsPerPeer> snapshots; + std::array snapped{}; + auto buffers_for = [&](uint32_t backend_id) -> const std::vector& { + static const std::vector kNone; + if (backend_id >= kMaxBackendsPerPeer) return kNone; + if (!snapped[backend_id]) { + snapshots[backend_id] = peer_directory_->RemoteBufferSnapshot(node_id, backend_id); + snapped[backend_id] = true; + } + return snapshots[backend_id]; + }; for (size_t idx = 0; idx < entries.size(); ++idx) { auto& entry = entries[idx]; - - auto zero_copy = FindRegisteredMemory(entry.item->dst, entry.item->size); - if (zero_copy.has_value()) { - entry.zero_copy = zero_copy; - entry.use_staging = false; - entry.staging_offset = 0; + // DstBytes(), not size: a ranged entry's destination is an arena slice + // holding only the spans. Passing the object size would make + // FindRegisteredMemory's containment check reject a slice that sits near + // the end of the arena, silently downgrading a zero-copy read to a staged + // one instead of failing. + TransferRef dst; + uint64_t dst_base = 0; + if (entry.item->dst_ref != nullptr) { + dst = *entry.item->dst_ref; + dst_base = entry.item->dst_ref_offset; } else { - entry.zero_copy.reset(); - entry.use_staging = true; - cursor = AlignUp(cursor, entry.plan.page_size); - const uint64_t aligned_size = AlignUp(entry.item->size, entry.plan.page_size); - if (cursor + aligned_size > config_.staging_buffer_size) { - MORI_UMBP_WARN( - "[PoolClient] BuildRemoteGetTransfers: staging buffer overflow at entry {}/{}, " - "key='{}' size={} aligned={} cursor={} cap={}", - idx, entries.size(), - (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, - entry.item ? entry.item->size : 0, aligned_size, cursor, config_.staging_buffer_size); - return false; - } - entry.staging_offset = cursor; - cursor += aligned_size; - } - - mori::io::MemoryDesc local_desc{}; - uint64_t base_offset = 0; - if (entry.use_staging) { - local_desc = staging_mem_; - base_offset = entry.staging_offset; - } else if (entry.zero_copy.has_value()) { - local_desc = entry.zero_copy->first; - base_offset = entry.zero_copy->second; - } else { - entry.failed = true; - continue; + std::tie(dst, dst_base) = UserBufferRef(entry.item->dst, entry.item->DstBytes()); } + const std::vector& remote = buffers_for(entry.plan.backend_id); - // Hydrate + snapshot remote descs inside ONE peers_mutex_ window (see - // BuildRemotePutTransfers). - std::vector entry_transfers; - entry_transfers.reserve(entry.plan.pages.size()); - { - std::lock_guard lock(peers_mutex_); - EnsureBufferDescsCachedLocked(peer, entry.plan.descs); - for (size_t p = 0; p < entry.plan.pages.size(); ++p) { - const auto& page = entry.plan.pages[p]; - if (page.buffer_index >= peer.dram_memories.size() || - !IsValidMemoryDesc(peer.dram_memories[page.buffer_index])) { + // Shared by both shapes: a page the peer never published means this key + // cannot be read at all, so the entry fails whole rather than partially. + auto peer_buffer = [&](const PageLocation& page) -> const TransferRef* { + if (page.buffer_index >= remote.size() || !remote[page.buffer_index].HasMemoryDesc()) { + MORI_UMBP_ERROR( + "[PoolClient] BuildRemoteGetTransfers: peer published no buffer, " + "key='{}' backend={} buffer_index={} peer_buffers={} page_index={}", + (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, + entry.plan.backend_id, page.buffer_index, remote.size(), page.page_index); + return nullptr; + } + return &remote[page.buffer_index]; + }; + + std::vector entry_items; + + if (entry.item->Ranged()) { + // Move only the requested spans, laid down back-to-back at the + // destination in the caller's span order. The spans are NOT assumed + // ascending — a packed multi-component KV pool emits them interleaved — + // so each is walked independently and any merging is left to the engine's + // planner, which coalesces adjacent segments per MR pair anyway. + // + // One item per (span x page it touches), so the page count alone is not + // the bound. + entry_items.reserve(entry.item->spans->size() + entry.plan.pages.size()); + size_t packed_cursor = 0; + for (const auto& span : *entry.item->spans) { + const bool walked = ForEachRangePageFragment( + entry.plan.pages, entry.plan.page_size, entry.item->size, span.object_offset, span.size, + [&](size_t page_index, uint64_t tier_offset, size_t fragment, size_t copied) { + const TransferRef* buf = peer_buffer(entry.plan.pages[page_index]); + if (buf == nullptr) return false; + TransferItem item; + item.tag = idx; + item.src = *buf; + item.src_offset = tier_offset; + item.dst = dst; + item.dst_offset = dst_base + packed_cursor + copied; + item.size = fragment; + entry_items.push_back(std::move(item)); + return true; + }); + if (!walked) { MORI_UMBP_ERROR( - "[PoolClient] BuildRemoteGetTransfers: invalid peer dram_memories slot, " - "key='{}' buffer_index={} peer_dram_size={} page_index={}", + "[PoolClient] BuildRemoteGetTransfers: span [{},{}) outside object, key='{}' size={}", + span.object_offset, span.object_offset + span.size, (entry.item && entry.item->key) ? *entry.item->key : std::string{""}, - page.buffer_index, peer.dram_memories.size(), page.page_index); + entry.item->size); entry.failed = true; - entry_transfers.clear(); + entry_items.clear(); break; } - TransferInstruction instr; - instr.entry_index = idx; - instr.local_desc = local_desc; - instr.local_offset = base_offset + static_cast(p) * entry.plan.page_size; - instr.remote_desc = peer.dram_memories[page.buffer_index]; - instr.remote_offset = static_cast(page.page_index) * entry.plan.page_size; - instr.size = + packed_cursor += span.size; + } + } else { + entry_items.reserve(entry.plan.pages.size()); + for (size_t p = 0; p < entry.plan.pages.size(); ++p) { + const TransferRef* buf = peer_buffer(entry.plan.pages[p]); + if (buf == nullptr) { + entry.failed = true; + entry_items.clear(); + break; + } + TransferItem item; + item.tag = idx; + item.src = *buf; + item.src_offset = + static_cast(entry.plan.pages[p].page_index) * entry.plan.page_size; + item.dst = dst; + item.dst_offset = dst_base + static_cast(p) * entry.plan.page_size; + item.size = LogicalPageBytes(p, entry.plan.pages.size(), entry.plan.page_size, entry.item->size); - entry_transfers.push_back(std::move(instr)); + entry_items.push_back(std::move(item)); } } - if (!entry_transfers.empty()) { - transfers->insert(transfers->end(), entry_transfers.begin(), entry_transfers.end()); + if (!entry_items.empty()) { + items->insert(items->end(), std::make_move_iterator(entry_items.begin()), + std::make_move_iterator(entry_items.end())); } } - - *staging_bytes = cursor; return true; } void PoolClient::FinalizeRemoteGetEntries(std::vector& entries, - std::vector* results) { + std::vector* results, bool recache_remote) { for (auto& entry : entries) { if (entry.failed) { (*results)[entry.result_index] = false; continue; } - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, - {{"traffic", "remote"}}, static_cast(entry.item->size)); - master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, - MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, - {{"traffic", "remote"}}, static_cast(entry.item->size)); + // Bytes that actually crossed the wire, which for a ranged entry is the + // span total and not the object size. Reporting the object size here would + // overstate remote GET bandwidth by the whole read-amplification factor the + // ranged path exists to remove. + const double moved = static_cast(entry.item->DstBytes()); + CountMetric(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "remote"}}, + moved); + CountMetric(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, + MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "remote"}}, + moved); (*results)[entry.result_index] = true; // Re-cache the remotely-fetched block into local DRAM (best-effort): the // user dst is already populated (staging copy-out and zero-copy both land // before Finalize runs), so subsequent reads of this key route local. - if (entry.item) { + // + // Never for a ranged entry: the destination holds only the requested spans, + // and a medium slot is whole-object, so installing it would publish a key + // whose remaining bytes are undefined. The ranged path schedules its own + // whole-object prefetch instead (MaybePrefetchWholeObject). + if (recache_remote && entry.item && !entry.item->Ranged()) { MaybeReCacheAfterRemote(*entry.item->key, entry.item->dst, entry.item->size); } } } -// --------------------------------------------------------------------------- -// Remote DRAM transfer grouping (shared by Put and Get) -// --------------------------------------------------------------------------- - -std::vector PoolClient::GroupTransfersByPair( - const std::vector& active) { - std::vector groups; - groups.reserve(active.size()); - // Group by (local MR id, remote MR id); first appearance defines order. - std::unordered_map pair_to_group; - pair_to_group.reserve(active.size() * 2); - - for (const auto& t : active) { - const PairKey key{t.local_desc.id, t.remote_desc.id}; - auto it = pair_to_group.find(key); - size_t gi; - if (it == pair_to_group.end()) { - gi = groups.size(); - pair_to_group.emplace(key, gi); - PairGroup g; - g.local_desc = t.local_desc; - g.remote_desc = t.remote_desc; - groups.push_back(std::move(g)); - } else { - gi = it->second; - } - PairGroup& g = groups[gi]; - const size_t lo = static_cast(t.local_offset); - const size_t ro = static_cast(t.remote_offset); - const size_t sz = static_cast(t.size); - // Coalesce with the previous segment when BOTH local and remote are - // exactly contiguous. This is NOT redundant with the backend's WR - // merging: the backend would fold these into the same WR either way, but - // doing it here shrinks the inner SG vector it has to sort/merge - // (O(M log M)) and allocate, which matters when M is large (big batch x - // pages). Same bytes, so per-pair failure granularity is unchanged. - // The merged size must stay within uint32_t since the backend stores it - // in ibv_sge.length (matches its own WR-merge cap in common.cpp). - const bool can_coalesce = - !g.sizes.empty() && AdjacentNoOverflow(g.local_offsets.back(), g.sizes.back(), lo) && - AdjacentNoOverflow(g.remote_offsets.back(), g.sizes.back(), ro) && - g.sizes.back() <= static_cast(std::numeric_limits::max()) - sz; - if (can_coalesce) { - g.sizes.back() += sz; - } else { - g.local_offsets.push_back(lo); - g.remote_offsets.push_back(ro); - g.sizes.push_back(sz); - } - // De-dup contributing entries: a single entry's pages arrive - // consecutively within a pair, so a back() check suffices. - if (g.entry_indices.empty() || g.entry_indices.back() != t.entry_index) { - g.entry_indices.push_back(t.entry_index); - } - } - return groups; -} - // --------------------------------------------------------------------------- // Cluster-wide existence check // --------------------------------------------------------------------------- @@ -2138,9 +4399,48 @@ bool PoolClient::Exists(const std::string& key) { std::vector PoolClient::BatchExists(const std::vector& keys) { if (!initialized_ || keys.empty()) return std::vector(keys.size(), false); - std::vector out; - auto status = master_client_->BatchLookup(keys, &out); - if (!status.ok() || out.size() != keys.size()) return std::vector(keys.size(), false); + // Local-first: this node's own backends answer conclusively for the keys they + // hold, so a batch that is entirely local needs no master round trip at all. + // The misses are NOT conclusive -- a peer may hold them -- so they still go to + // the master, and only they do. + std::vector out(keys.size(), false); + std::vector unknown; + std::vector unknown_indices; + if (config_.local_first) { + // Containment, not resolution: this only needs to know whether the key is + // here. A resolve would heap-copy the page vector of every key it finds and + // renew a read lease on it, and every byte of that is discarded one line + // below. + if (default_pool_ != nullptr) { + const auto local = default_pool_->BatchContains(keys); + for (size_t i = 0; i < keys.size() && i < local.size(); ++i) out[i] = local[i]; + } + for (size_t i = 0; i < keys.size(); ++i) { + if (!out[i]) { + unknown.push_back(keys[i]); + unknown_indices.push_back(i); + } + } + if (unknown.empty()) return out; + } else { + unknown = keys; + unknown_indices.resize(keys.size()); + std::iota(unknown_indices.begin(), unknown_indices.end(), 0); + } + + // With no master this node is the whole cluster, so what it does not hold + // does not exist. `out` already carries that answer. + if (!HasMaster()) return out; + + std::vector remote; + auto status = master_client_->BatchLookup(unknown, &remote); + // A lookup failure must not downgrade a local hit to a miss: the local half of + // the answer did not depend on the RPC. Only the keys we could not answer + // ourselves fall back to false. + if (!status.ok() || remote.size() != unknown.size()) return out; + for (size_t j = 0; j < unknown_indices.size(); ++j) { + if (remote[j]) out[unknown_indices[j]] = true; + } return out; } @@ -2151,17 +4451,26 @@ std::vector PoolClient::BatchExists(const std::vector& keys) bool PoolClient::ReportExternalKvBlocks(const std::vector& hashes, TierType tier) { if (!initialized_) return false; if (hashes.empty()) return true; + // External KV placement lives only in the master's index; with no master + // there is no such index and nothing to report to. + if (!HasMaster()) return true; return master_client_->ReportExternalKvBlocks(config_.master_config.node_id, hashes, tier).ok(); } bool PoolClient::RevokeExternalKvBlocks(const std::vector& hashes, TierType tier) { if (!initialized_) return false; if (hashes.empty()) return true; + // External KV placement lives only in the master's index; with no master + // there is no such index and nothing to report to. + if (!HasMaster()) return true; return master_client_->RevokeExternalKvBlocks(config_.master_config.node_id, hashes, tier).ok(); } bool PoolClient::RevokeAllExternalKvBlocksAtTier(TierType tier) { if (!initialized_) return false; + // External KV placement lives only in the master's index; with no master + // there is no such index and nothing to report to. + if (!HasMaster()) return true; return master_client_->RevokeAllExternalKvBlocksAtTier(config_.master_config.node_id, tier).ok(); } @@ -2169,6 +4478,9 @@ bool PoolClient::MatchExternalKv(const std::vector& hashes, std::vector* out_matches, bool count_as_hit) { if (!initialized_) return false; + // External KV placement lives only in the master's index; with no master + // there is no such index and nothing to report to. + if (!HasMaster()) return true; return master_client_->MatchExternalKv(hashes, out_matches, count_as_hit).ok(); } @@ -2176,6 +4488,9 @@ bool PoolClient::GetExternalKvHitCounts( const std::vector& hashes, std::vector* out_entries) { if (!initialized_) return false; + // External KV placement lives only in the master's index; with no master + // there is no such index and nothing to report to. + if (!HasMaster()) return true; return master_client_->GetExternalKvHitCounts(hashes, out_entries).ok(); } @@ -2190,200 +4505,29 @@ PoolClient::PeerConnection& PoolClient::GetOrConnectPeer(const std::string& node if (it != peers_.end()) return *it->second; auto conn = std::make_unique(); + conn->node_id = node_id; conn->peer_address = peer_address; - // engine_desc is hydrated lazily in EnsurePeerServiceConnection from - // the peer's GetPeerInfo response. + // The peer's engine desc and buffer descriptors are hydrated lazily in + // EnsurePeerServiceConnection, into the transfer engine rather than here. auto& ref = *conn; peers_[node_id] = std::move(conn); return ref; } -void PoolClient::EnsureBufferDescsCachedLocked(PeerConnection& peer, - const std::vector& descs) { - // Caller holds peers_mutex_. - if (!io_engine_) return; - for (const auto& d : descs) { - if (peer.dram_memories.size() <= d.buffer_index) { - peer.dram_memories.resize(d.buffer_index + 1); - } - if (IsValidMemoryDesc(peer.dram_memories[d.buffer_index])) continue; - if (d.desc_bytes.empty()) continue; - auto handle = - msgpack::unpack(reinterpret_cast(d.desc_bytes.data()), d.desc_bytes.size()); - peer.dram_memories[d.buffer_index] = handle.get().as(); - } -} - -void PoolClient::EnsureBufferDescsCached(PeerConnection& peer, - const std::vector& descs) { - std::lock_guard lock(peers_mutex_); - EnsureBufferDescsCachedLocked(peer, descs); -} - -// --------------------------------------------------------------------------- -// RDMA scatter helpers (unchanged from prior impl) -// --------------------------------------------------------------------------- - -bool PoolClient::RemoteDramScatterWrite(PeerConnection& peer, - const std::vector& pages, uint64_t page_size, - const void* src, size_t size, bool zero_copy) { - if (!io_engine_) return false; - if (pages.empty() || page_size == 0) return false; - if (!SizeMatchesAllocation(size, pages.size(), page_size)) return false; - - mori::io::MemoryDesc local_mem; - size_t local_base_offset = 0; - bool used_zero_copy = false; - std::unique_lock staging_lock(staging_mutex_, std::defer_lock); - if (zero_copy) { - auto reg = FindRegisteredMemory(src, size); - if (reg) { - local_mem = reg->first; - local_base_offset = reg->second; - used_zero_copy = true; - } - } - if (!used_zero_copy) { - if (size > config_.staging_buffer_size) return false; - staging_lock.lock(); - std::memcpy(staging_buffer_.get(), src, size); - local_mem = staging_mem_; - local_base_offset = 0; - } - - auto groups = GroupPagesByBuffer(pages); - const size_t N = groups.size(); - - mori::io::MemDescVec remote_descs; - remote_descs.reserve(N); - { - std::lock_guard lock(peers_mutex_); - for (size_t k = 0; k < N; ++k) { - const auto& g = groups[k]; - if (g.buffer_index >= peer.dram_memories.size() || - !IsValidMemoryDesc(peer.dram_memories[g.buffer_index])) { - return false; - } - remote_descs.push_back(peer.dram_memories[g.buffer_index]); - } - } - - mori::io::MemDescVec local_descs(N, local_mem); - mori::io::BatchSizeVec local_offsets(N), remote_offsets(N), sizes_v(N); - for (size_t k = 0; k < N; ++k) { - const auto& g = groups[k]; - local_offsets[k].reserve(g.src_page_indices.size()); - remote_offsets[k].reserve(g.src_page_indices.size()); - sizes_v[k].reserve(g.src_page_indices.size()); - for (size_t spi : g.src_page_indices) { - local_offsets[k].push_back(local_base_offset + spi * page_size); - remote_offsets[k].push_back(static_cast(pages[spi].page_index) * page_size); - sizes_v[k].push_back(LogicalPageBytes(spi, pages.size(), page_size, size)); - } - } - - std::vector statuses(N); - mori::io::TransferStatusPtrVec status_ptrs(N); - mori::io::TransferUniqueIdVec ids(N); - for (size_t k = 0; k < N; ++k) { - status_ptrs[k] = &statuses[k]; - ids[k] = io_engine_->AllocateTransferUniqueId(); - } - io_engine_->BatchWrite(local_descs, local_offsets, remote_descs, remote_offsets, sizes_v, - status_ptrs, ids); - bool all_ok = true; - for (auto& s : statuses) { - s.Wait(); - if (!s.Succeeded()) all_ok = false; - } - return all_ok; -} - -bool PoolClient::RemoteDramScatterRead(PeerConnection& peer, const std::vector& pages, - uint64_t page_size, void* dst, size_t size, bool zero_copy) { - if (!io_engine_) return false; - if (pages.empty() || page_size == 0) return false; - if (!SizeMatchesAllocation(size, pages.size(), page_size)) return false; - - mori::io::MemoryDesc local_mem; - size_t local_base_offset = 0; - bool used_zero_copy = false; - std::unique_lock staging_lock(staging_mutex_, std::defer_lock); - if (zero_copy) { - auto reg = FindRegisteredMemory(dst, size); - if (reg) { - local_mem = reg->first; - local_base_offset = reg->second; - used_zero_copy = true; - } - } - if (!used_zero_copy) { - if (size > config_.staging_buffer_size) return false; - staging_lock.lock(); - local_mem = staging_mem_; - local_base_offset = 0; - } - - auto groups = GroupPagesByBuffer(pages); - const size_t N = groups.size(); - - mori::io::MemDescVec remote_descs; - remote_descs.reserve(N); - { - std::lock_guard lock(peers_mutex_); - for (size_t k = 0; k < N; ++k) { - const auto& g = groups[k]; - if (g.buffer_index >= peer.dram_memories.size() || - !IsValidMemoryDesc(peer.dram_memories[g.buffer_index])) { - return false; - } - remote_descs.push_back(peer.dram_memories[g.buffer_index]); - } - } - - mori::io::MemDescVec local_descs(N, local_mem); - mori::io::BatchSizeVec local_offsets(N), remote_offsets(N), sizes_v(N); - for (size_t k = 0; k < N; ++k) { - const auto& g = groups[k]; - local_offsets[k].reserve(g.src_page_indices.size()); - remote_offsets[k].reserve(g.src_page_indices.size()); - sizes_v[k].reserve(g.src_page_indices.size()); - for (size_t spi : g.src_page_indices) { - local_offsets[k].push_back(local_base_offset + spi * page_size); - remote_offsets[k].push_back(static_cast(pages[spi].page_index) * page_size); - sizes_v[k].push_back(LogicalPageBytes(spi, pages.size(), page_size, size)); - } - } - - std::vector statuses(N); - mori::io::TransferStatusPtrVec status_ptrs(N); - mori::io::TransferUniqueIdVec ids(N); - for (size_t k = 0; k < N; ++k) { - status_ptrs[k] = &statuses[k]; - ids[k] = io_engine_->AllocateTransferUniqueId(); - } - io_engine_->BatchRead(local_descs, local_offsets, remote_descs, remote_offsets, sizes_v, - status_ptrs, ids); - bool all_ok = true; - for (auto& s : statuses) { - s.Wait(); - if (!s.Succeeded()) all_ok = false; - } - if (!all_ok) return false; - if (!used_zero_copy) std::memcpy(dst, staging_buffer_.get(), size); - return true; -} - // --------------------------------------------------------------------------- -// SSD path (preserved from prior impl) +// Peer connection setup // --------------------------------------------------------------------------- bool PoolClient::EnsurePeerServiceConnection(PeerConnection& peer) { - std::lock_guard lock(peer.ssd_op_mutex); + std::lock_guard lock(peer.conn_mutex); if (peer.peer_address.empty()) { return false; } + // GetPeerInfo returns two different kinds of fact and they now go to two + // different owners: the stub is a control-plane connection kept here, while + // the peer's engine desc and buffer descriptors are transfer-layer facts and + // go straight into the engine's remote cache. auto hydrate_from_peer = [&](::umbp::UMBPPeer::Stub* stub) -> bool { ::umbp::GetPeerInfoRequest req; ::umbp::GetPeerInfoResponse resp; @@ -2395,49 +4539,47 @@ bool PoolClient::EnsurePeerServiceConnection(PeerConnection& peer) { return false; } - if (!resp.engine_desc().empty()) { - auto handle = msgpack::unpack(resp.engine_desc().data(), resp.engine_desc().size()); - peer.engine_desc = handle.get().as(); - if (io_engine_) { - io_engine_->RegisterRemoteEngine(peer.engine_desc); - peer.engine_registered = true; + if (peer_directory_ != nullptr) { + if (!peer_directory_->EnsureRemoteEngine(peer.node_id, resp.engine_desc())) return false; + + // Every backend's buffers arrive in one list, each entry naming the + // backend its buffer_index belongs to. Before backend_id, this list held + // one medium's buffers and the rest were unreachable — yet HasRemoteBuffers + // still went true, so the next resolve asked the peer to omit descriptors + // and the missing media's pages were read against the published one's + // memory. + std::vector descs; + descs.reserve(resp.buffer_descs_size()); + for (const auto& d : resp.buffer_descs()) { + if (d.desc().empty()) continue; + BufferMemoryDescBytes b; + b.buffer_index = d.buffer_index(); + b.backend_id = d.backend_id(); + b.desc_bytes.assign(d.desc().begin(), d.desc().end()); + descs.push_back(std::move(b)); } - } else if (io_engine_ && !peer.engine_registered) { - return false; - } - - if (!resp.ssd_staging_mem_desc().empty()) { - auto handle = - msgpack::unpack(resp.ssd_staging_mem_desc().data(), resp.ssd_staging_mem_desc().size()); - peer.ssd_staging_mem = handle.get().as(); - peer.ssd_staging_size = resp.ssd_staging_size(); - } - - for (const auto& d : resp.dram_memory_descs()) { - if (peer.dram_memories.size() <= d.buffer_index()) { - peer.dram_memories.resize(d.buffer_index() + 1); - } - if (IsValidMemoryDesc(peer.dram_memories[d.buffer_index()])) continue; - if (d.desc().empty()) continue; - auto h = msgpack::unpack(d.desc().data(), d.desc().size()); - peer.dram_memories[d.buffer_index()] = h.get().as(); + peer_directory_->CacheRemoteBuffers(peer.node_id, descs); } return true; }; + const bool engine_known = + peer_directory_ == nullptr || peer_directory_->HasRemoteEngine(peer.node_id); + if (peer.peer_stub) { - if (!peer.engine_registered && io_engine_) { + if (!engine_known) { auto* stub = static_cast<::umbp::UMBPPeer::Stub*>(peer.peer_stub.get()); if (!hydrate_from_peer(stub)) { peer.peer_stub.reset(); - peer.engine_registered = false; + if (peer_directory_ != nullptr) peer_directory_->ForgetRemote(peer.node_id); return false; } } return true; } - auto channel = grpc::CreateChannel(peer.peer_address, grpc::InsecureChannelCredentials()); + auto channel = grpc::CreateCustomChannel(peer.peer_address, grpc::InsecureChannelCredentials(), + GrpcChannelArgs()); auto stub = ::umbp::UMBPPeer::NewStub(channel); if (!hydrate_from_peer(stub.get())) { return false; @@ -2448,276 +4590,67 @@ bool PoolClient::EnsurePeerServiceConnection(PeerConnection& peer) { return true; } -// Remote SSD get for one key (reader != owner): key-based PrepareSsdRead on the -// peer reads the bytes into its serving staging slot; we RDMA them out of the -// published staging buffer, then issue a best-effort ReleaseSsdLease. Outcomes: -// kMiss (NOT_FOUND, definitive miss); kRetry (retryable: NO_SLOT or a -// reader-local lease expiry); kError (not-served, not retried: rpc failure -// incl. DEADLINE_EXCEEDED, size mismatch, RDMA failure). -// -// Lease gating: the deadline is anchored at t_send (before the RPC) so it stays -// conservative against the peer, which counts the same TTL from request receipt -// (see ssd_read_lease.h). A read is reported successful only if RDMA finished -// before that deadline; once expired we return a transient retry and do NOT -// release (the peer reclaims the slot by TTL). A late-returning PrepareSsdRead -// is caught by the pre-RDMA expiry check; a hung one by the RPC deadline. -PoolClient::SsdGetOutcome PoolClient::RemoteSsdReadOnce(PeerConnection& peer, - const std::string& key, void* dst, - size_t size) { - namespace lease = ssd_read_lease; - if (!io_engine_) return SsdGetOutcome::kError; - if (!EnsurePeerServiceConnection(peer)) return SsdGetOutcome::kError; - if (!IsValidMemoryDesc(peer.ssd_staging_mem)) return SsdGetOutcome::kError; - auto* stub = static_cast<::umbp::UMBPPeer::Stub*>(peer.peer_stub.get()); - - // Anchor the lease deadline before sending; also bound the RPC itself so a - // hung peer can't stall the serial batch. Fall back to the configured lease - // timeout when the dedicated timeout env is unset (cluster-homogeneous). - const auto t_send = std::chrono::steady_clock::now(); - auto rpc_timeout = SsdPrepareRpcTimeoutOverride(); - if (rpc_timeout.count() == 0) { - rpc_timeout = SsdReadLeaseTtl(); - } - - ::umbp::PrepareSsdReadRequest req; - req.set_key(key); - req.set_max_size(size); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - // gRPC deadlines are specialized for system_clock; the lease gating below - // uses the monotonic steady_clock t_send for the actual validity window. - ctx.set_deadline(std::chrono::system_clock::now() + rpc_timeout); - const grpc::Status rpc = stub->PrepareSsdRead(&ctx, req, &resp); - if (!rpc.ok()) { - // Includes DEADLINE_EXCEEDED (peer slow / may already hold a claimed slot) - // and UNAVAILABLE: hard failure, not retried — see SsdGetTransientMaxAttempts. - MORI_UMBP_DEBUG("[PoolClient] RemoteSsdRead key='{}' PrepareSsdRead rpc failed (code={})", key, - static_cast(rpc.error_code())); - return SsdGetOutcome::kError; - } - - switch (resp.status()) { - case ::umbp::SSD_READ_OK: - break; - case ::umbp::SSD_READ_NOT_FOUND: - return SsdGetOutcome::kMiss; - case ::umbp::SSD_READ_NO_SLOT: - // Transient slot exhaustion (no slot was claimed): retryable, not a miss. - return SsdGetOutcome::kRetry; - default: // SIZE_TOO_LARGE / ERROR / unexpected - MORI_UMBP_WARN("[PoolClient] RemoteSsdRead key='{}' status={}", key, - static_cast(resp.status())); - return SsdGetOutcome::kError; - } - - const auto deadline_ttl = resp.lease_ttl_ms(); - - // If the lease already elapsed (slow/late PrepareSsdRead return), don't even - // RDMA — the staging bytes may already be getting recycled. Transient, not a - // miss; leave the slot for the peer's TTL reclaim (no release). - if (lease::LeaseExpired(t_send, deadline_ttl, std::chrono::steady_clock::now())) { - MORI_UMBP_DEBUG("[PoolClient] RemoteSsdRead key='{}' lease expired before RDMA (ttl_ms={})", - key, deadline_ttl); - return SsdGetOutcome::kRetry; - } - - if (resp.size() != size) { - MORI_UMBP_WARN("[PoolClient] RemoteSsdRead key='{}' size mismatch (wanted {}, got {})", key, - size, resp.size()); - ReleaseSsdLeaseBestEffort(stub, resp.lease_id()); - return SsdGetOutcome::kError; - } - - // RDMA the staged bytes home: zero-copy if the dst is registered, else stage - // through our own buffer and memcpy. - bool rdma_ok = false; - auto reg = FindRegisteredMemory(dst, size); - if (reg) { - auto uid = io_engine_->AllocateTransferUniqueId(); - mori::io::TransferStatus status; - io_engine_->Read(reg->first, reg->second, peer.ssd_staging_mem, resp.staging_offset(), size, - &status, uid); - status.Wait(); - rdma_ok = status.Succeeded(); - } - if (!rdma_ok && size <= config_.staging_buffer_size) { - std::lock_guard lock(staging_mutex_); - auto uid = io_engine_->AllocateTransferUniqueId(); - mori::io::TransferStatus status; - io_engine_->Read(staging_mem_, 0, peer.ssd_staging_mem, resp.staging_offset(), size, &status, - uid); - status.Wait(); - if (status.Succeeded()) { - std::memcpy(dst, staging_buffer_.get(), size); - rdma_ok = true; - } - } - - // Decide against the lease deadline: a read that finished after the deadline - // is untrusted (the peer may have recycled the slot mid-RDMA), so it becomes - // a transient retry and we skip release. The caller uses the return value; - // any bytes already written to dst on the expired path are not consumed. - const bool expired = lease::LeaseExpired(t_send, deadline_ttl, std::chrono::steady_clock::now()); - const auto decision = lease::DecideSsdReadOutcome(expired, rdma_ok); - if (decision.release) ReleaseSsdLeaseBestEffort(stub, resp.lease_id()); - if (expired && rdma_ok) { - MORI_UMBP_DEBUG("[PoolClient] RemoteSsdRead key='{}' lease expired after RDMA (ttl_ms={})", key, - deadline_ttl); - } - switch (decision.outcome) { - case lease::GateOutcome::kSuccess: - return SsdGetOutcome::kSuccess; - case lease::GateOutcome::kRetry: - return SsdGetOutcome::kRetry; - case lease::GateOutcome::kError: - return SsdGetOutcome::kError; - } - return SsdGetOutcome::kError; -} - -// Best-effort lease release. Correctness does not depend on it: if it fails or -// is never called, the peer reclaims the slot when the lease TTL expires. Each -// attempt is bounded by a short deadline so a slow peer can't stall the caller. -void PoolClient::ReleaseSsdLeaseBestEffort(::umbp::UMBPPeer::Stub* stub, uint64_t lease_id) { - if (lease_id == 0) return; - const uint32_t max_retries = ReleaseLeaseMaxRetries(); - const auto timeout = ReleaseLeaseRpcTimeout(); - for (uint32_t attempt = 0; attempt < max_retries; ++attempt) { - ::umbp::ReleaseSsdLeaseRequest rel_req; - rel_req.set_lease_id(lease_id); - ::umbp::ReleaseSsdLeaseResponse rel_resp; - grpc::ClientContext rel_ctx; - rel_ctx.set_deadline(std::chrono::system_clock::now() + timeout); - if (stub->ReleaseSsdLease(&rel_ctx, rel_req, &rel_resp).ok()) break; - } -} - -void PoolClient::ProcessRemoteSsdBatchGet(const std::vector& items, - std::vector* results) { - if (items.empty()) return; - const auto& first = items.front(); - auto& peer = GetOrConnectPeer(first.route.node_id, first.route.peer_address); - - const uint32_t max_attempts = SsdGetTransientMaxAttempts(); - uint64_t transient_not_served = 0; // aggregated per batch → one AddCounter below - for (const auto& item : items) { - bool done = false; - for (uint32_t attempt = 0; attempt < max_attempts && !done; ++attempt) { - SsdGetOutcome outcome = RemoteSsdReadOnce(peer, *item.key, item.dst, item.size); - switch (outcome) { - case SsdGetOutcome::kSuccess: - (*results)[item.index] = true; - done = true; - break; - case SsdGetOutcome::kMiss: // definitive miss - case SsdGetOutcome::kError: // hard failure (already logged) - done = true; - break; - case SsdGetOutcome::kRetry: // transient NO_SLOT / reader-local lease expiry; not a miss - if (attempt + 1 < max_attempts) std::this_thread::sleep_for(SsdGetRetryBackoff()); - break; - } - } - if (!done) { - ++transient_not_served; - MORI_UMBP_WARN( - "[PoolClient] Remote SSD get key='{}' still transient-failing (NO_SLOT/lease-expired) " - "after {} attempts; reporting as not-served this round (not a definitive miss)", - *item.key, max_attempts); - } - } - // Optional reader-side diagnostic: lease expiry is reader-local and never - // shows up in the peer's ssd_read_total, so surface transient not-served here. - // Aggregated to a single AddCounter per batch (cheap, no per-key lock churn). - if (transient_not_served > 0 && master_client_) { - master_client_->AddCounter(MORI_UMBP_METRIC_SSD_READ_CLIENT_TRANSIENT_TOTAL, - MORI_UMBP_METRIC_SSD_READ_CLIENT_TRANSIENT_TOTAL_HELP, {}, - static_cast(transient_not_served)); - } -} - -void PoolClient::PublishSsdMetrics() { +void PoolClient::PublishComponentMetrics() { if (!master_client_) return; - // Ship a monotonic peer-side counter as the delta since the last tick. `last` - // is updated even on a zero delta (or a defensive counter reset) so the next - // tick stays correct. Runs once per metrics flush in the metrics thread. - auto ship_counter = [&](const char* name, const char* help, MasterClient::Labels labels, - uint64_t current, uint64_t& last) { - if (current > last) { - master_client_->AddCounter(name, help, std::move(labels), - static_cast(current - last)); + MetricPublisher::Sink sink{ + [this](const char* name, const char* help, const MetricLabels& labels, double delta) { + CountMetric(name, help, labels, delta); + }, + [this](const char* name, const char* help, const MetricLabels& labels, double value) { + master_client_->SetGauge(name, help, labels, value); + }}; + + // Storage backends. tier= and backend= come from the component's own + // identity, so this loop never names a medium and a fourth one needs no edit + // here. Both halves of a backend's metrics — the generic series + // InstrumentedBackend measured and whatever the medium publishes about its + // own internals — arrive through the one SampleMetrics() call. + for (MediumBackend* backend : registry_.All()) { + if (backend == nullptr) continue; + const auto* entry = registry_.GetEntry(backend); + const std::string backend_name = entry == nullptr ? std::string(backend->Name()) : entry->name; + const MetricLabels labels = {{"tier", TierTypeName(backend->Tier())}, + {"backend", backend_name}}; + metric_publisher_.Publish(std::string("backend:") + backend_name, labels, *backend, sink); + } + + // The transfer layer. Its samples already carry engine=, stamped by the + // composite that dispatched them, so there is nothing to add at this level. + if (transfer_engine_ != nullptr) { + metric_publisher_.Publish("transfer", {}, *transfer_engine_, sink); + } + + // Logical tier transitions. PeerPool accumulates these and nothing outside + // the tier benchmark reads them, so without this a served workload cannot + // distinguish a tier graph that is migrating from one whose every migration + // fails. Read hits carry the logical tier as a label because the per-backend + // labels above collapse every instance of a medium into one series, which + // hides which tier actually served a read. + if (default_pool_ != nullptr) { + std::vector samples; + const auto sample = [&samples](const char* name, const char* help, MetricLabels labels, + uint64_t value) { + samples.push_back(MetricSample{name, help, std::move(labels), value}); + }; + + const TierTransitionMetrics tiers = default_pool_->TransitionMetrics(); + sample(MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS, MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS_HELP, + {{"outcome", "attempted"}}, tiers.attempted); + sample(MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS, MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS_HELP, + {{"outcome", "succeeded"}}, tiers.succeeded); + sample(MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS, MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS_HELP, + {{"outcome", "failed"}}, tiers.failed); + sample(MORI_UMBP_METRIC_CLIENT_TIER_OFFLOADED_BYTES, + MORI_UMBP_METRIC_CLIENT_TIER_OFFLOADED_BYTES_HELP, {}, tiers.offloaded_bytes); + sample(MORI_UMBP_METRIC_CLIENT_TIER_PROMOTED_BYTES, + MORI_UMBP_METRIC_CLIENT_TIER_PROMOTED_BYTES_HELP, {}, tiers.promoted_bytes); + + for (const auto& [tier, hits] : default_pool_->TierReadHits()) { + sample(MORI_UMBP_METRIC_CLIENT_TIER_READ_HITS, MORI_UMBP_METRIC_CLIENT_TIER_READ_HITS_HELP, + {{"tier", tier}}, hits); } - last = current; - }; - - if (ssd_copy_pipeline_) { - ship_counter(MORI_UMBP_METRIC_SSD_COPY_ENQUEUED_TOTAL, - MORI_UMBP_METRIC_SSD_COPY_ENQUEUED_TOTAL_HELP, {}, ssd_copy_pipeline_->Enqueued(), - ssd_metrics_last_.copy_enqueued); - ship_counter(MORI_UMBP_METRIC_SSD_COPY_SUCCEEDED_TOTAL, - MORI_UMBP_METRIC_SSD_COPY_SUCCEEDED_TOTAL_HELP, {}, ssd_copy_pipeline_->CopiedOk(), - ssd_metrics_last_.copy_succeeded); - ship_counter(MORI_UMBP_METRIC_SSD_COPY_FAILED_TOTAL, - MORI_UMBP_METRIC_SSD_COPY_FAILED_TOTAL_HELP, {}, ssd_copy_pipeline_->Failed(), - ssd_metrics_last_.copy_failed); - ship_counter(MORI_UMBP_METRIC_SSD_COPY_DROPPED_TOTAL, - MORI_UMBP_METRIC_SSD_COPY_DROPPED_TOTAL_HELP, {{"reason", "queue_full"}}, - ssd_copy_pipeline_->Dropped(), ssd_metrics_last_.copy_dropped_queue_full); - ship_counter(MORI_UMBP_METRIC_SSD_COPY_DROPPED_TOTAL, - MORI_UMBP_METRIC_SSD_COPY_DROPPED_TOTAL_HELP, {{"reason", "stopped"}}, - ssd_copy_pipeline_->DroppedStopped(), ssd_metrics_last_.copy_dropped_stopped); - } - - if (peer_ssd_) { - ship_counter(MORI_UMBP_METRIC_SSD_READ_TOTAL, MORI_UMBP_METRIC_SSD_READ_TOTAL_HELP, - {{"status", "ok"}}, peer_ssd_->ReadOk(), ssd_metrics_last_.read_ok); - ship_counter(MORI_UMBP_METRIC_SSD_READ_TOTAL, MORI_UMBP_METRIC_SSD_READ_TOTAL_HELP, - {{"status", "not_found"}}, peer_ssd_->ReadNotFound(), - ssd_metrics_last_.read_not_found); - ship_counter(MORI_UMBP_METRIC_SSD_READ_TOTAL, MORI_UMBP_METRIC_SSD_READ_TOTAL_HELP, - {{"status", "size_too_large"}}, peer_ssd_->ReadSizeTooLarge(), - ssd_metrics_last_.read_size_too_large); - ship_counter(MORI_UMBP_METRIC_SSD_READ_TOTAL, MORI_UMBP_METRIC_SSD_READ_TOTAL_HELP, - {{"status", "error"}}, peer_ssd_->ReadError(), ssd_metrics_last_.read_error); - - // SSD IO byte counters -> bandwidth via rate() in Grafana. - ship_counter(MORI_UMBP_METRIC_SSD_COPY_BYTES_TOTAL, MORI_UMBP_METRIC_SSD_COPY_BYTES_TOTAL_HELP, - {}, peer_ssd_->CopyBytes(), ssd_metrics_last_.copy_bytes); - ship_counter(MORI_UMBP_METRIC_SSD_READ_BYTES_TOTAL, MORI_UMBP_METRIC_SSD_READ_BYTES_TOTAL_HELP, - {}, peer_ssd_->ReadBytes(), ssd_metrics_last_.read_bytes); - - ship_counter(MORI_UMBP_METRIC_SSD_EVICTION_ROUNDS_TOTAL, - MORI_UMBP_METRIC_SSD_EVICTION_ROUNDS_TOTAL_HELP, {}, peer_ssd_->EvictionRounds(), - ssd_metrics_last_.evict_rounds); - ship_counter(MORI_UMBP_METRIC_SSD_EVICTION_VICTIMS_TOTAL, - MORI_UMBP_METRIC_SSD_EVICTION_VICTIMS_TOTAL_HELP, {}, peer_ssd_->EvictionVictims(), - ssd_metrics_last_.evict_victims); - ship_counter(MORI_UMBP_METRIC_SSD_EVICTION_BYTES_FREED_TOTAL, - MORI_UMBP_METRIC_SSD_EVICTION_BYTES_FREED_TOTAL_HELP, {}, - peer_ssd_->EvictionBytesFreed(), ssd_metrics_last_.evict_bytes_freed); - ship_counter(MORI_UMBP_METRIC_SSD_EVICTION_BACKEND_FAILED_TOTAL, - MORI_UMBP_METRIC_SSD_EVICTION_BACKEND_FAILED_TOTAL_HELP, {}, - peer_ssd_->EvictionBackendFailures(), ssd_metrics_last_.evict_backend_failed); - } - - if (peer_service_) { - const auto& m = peer_service_->Metrics(); - const uint64_t expired = m.expired_reclaims.load(std::memory_order_relaxed); - const uint64_t slot_full = m.slot_full_rejects.load(std::memory_order_relaxed); - ship_counter(MORI_UMBP_METRIC_SSD_STAGING_EXPIRED_RECLAIMS_TOTAL, - MORI_UMBP_METRIC_SSD_STAGING_EXPIRED_RECLAIMS_TOTAL_HELP, {}, expired, - ssd_metrics_last_.staging_expired_reclaims); - ship_counter(MORI_UMBP_METRIC_SSD_STAGING_SLOT_FULL_REJECTS_TOTAL, - MORI_UMBP_METRIC_SSD_STAGING_SLOT_FULL_REJECTS_TOTAL_HELP, {}, slot_full, - ssd_metrics_last_.staging_slot_full_rejects); - // A NO_SLOT read outcome IS a slot-full reject; surface it under the unified - // read_total{status=no_slot} too (same event, peer-service view of reads). - ship_counter(MORI_UMBP_METRIC_SSD_READ_TOTAL, MORI_UMBP_METRIC_SSD_READ_TOTAL_HELP, - {{"status", "no_slot"}}, slot_full, ssd_metrics_last_.read_no_slot); - master_client_->SetGauge(MORI_UMBP_METRIC_SSD_STAGING_SLOTS_IN_USE, - MORI_UMBP_METRIC_SSD_STAGING_SLOTS_IN_USE_HELP, {}, - static_cast(peer_service_->SnapshotReadSlotsInUse())); + metric_publisher_.Publish("pool", {}, samples, sink); } } diff --git a/src/umbp/distributed/proto/umbp.proto b/src/umbp/distributed/proto/umbp.proto index 8bc2805ab..aceb5d5ff 100644 --- a/src/umbp/distributed/proto/umbp.proto +++ b/src/umbp/distributed/proto/umbp.proto @@ -19,6 +19,22 @@ message TierCapacity { TierType tier = 1; uint64 total_capacity_bytes = 2; // Total capacity on this tier uint64 available_capacity_bytes = 3; // Current free capacity on this tier + uint64 max_allocatable_bytes = 4; // Largest value one backend can admit +} + +message LogicalTierCapacity { + string name = 1; + TierType representative_tier = 2; + uint64 total_capacity_bytes = 3; + uint64 available_capacity_bytes = 4; + uint64 max_allocatable_bytes = 5; + bool put_eligible = 6; + // The tier's own pressure as the offload decision measures it: the highest + // utilization among its backends, not the aggregate above. The two differ + // whenever a tier mixes backend sizes, and for the entry tier the aggregate + // covers every tier it can spill into, so this is the only field a reader can + // compare against the tier's watermarks. + double peak_member_utilization = 7; } // One page slot inside one of a node's DRAM/HBM buffers. Carried by the @@ -29,12 +45,26 @@ message PageLocation { uint32 page_index = 2; } -// One DRAM/HBM buffer's RDMA MemoryDesc plus the buffer_index it belongs -// to. Peer responses ship a deduplicated, ascending list so the writer -// can hydrate its peer-side buffer_index → MemoryDesc cache in one batch. +// One buffer's RDMA MemoryDesc plus the address it lives at. Peer responses +// ship a deduplicated, ascending list so the writer can hydrate its peer-side +// buffer cache in one batch. +// +// THE ADDRESS IS THE PAIR, not buffer_index alone. Every backend numbers its +// own buffers from 0 (PageBackend from its allocator, SsdBackend's staging +// arena is hardcoded to 0), so a bare buffer_index is ambiguous the moment a +// second medium is live on a peer — and in practice every backend publishes +// exactly one buffer, so index 0 collides on EVERY mixed-media node. Before +// backend_id existed, one medium's descriptors silently overwrote another's in +// the reader's cache and reads were served from the wrong memory: a hit, with +// garbage bytes. +// +// backend_id is assigned by the peer's BackendRegistry and is peer-local and +// opaque to the client, which only ever echoes it back. buffer_index stays +// backend-local and unchanged — no backend knows this field exists. message BufferMemoryDesc { uint32 buffer_index = 1; bytes desc = 2; // msgpack-packed mori::io::MemoryDesc + uint32 backend_id = 3; } // ============================================================ @@ -58,10 +88,13 @@ message RegisterClientRequest { string peer_address = 4; // PeerService gRPC address bytes engine_desc = 5; // packed EngineDesc repeated string tags = 13; // opaque key=value labels (e.g. "sgl_role=prefill") + repeated LogicalTierCapacity logical_tier_capacities = 14; } message RegisterClientResponse { uint64 heartbeat_interval_ms = 1; // Recommended heartbeat interval uint64 ack_seq = 2; // Initial heartbeat seq peer should send (= 0) + bool supports_max_allocatable_bytes = 3; + bool supports_logical_tiers = 4; } message UnregisterClientRequest { @@ -86,6 +119,7 @@ message KvEvent { string key = 2; TierType tier = 3; uint64 size = 4; // ADD only + string logical_tier = 5; } message EventBundle { @@ -107,6 +141,7 @@ message HeartbeatRequest { repeated TierCapacity tier_capacities = 4; // ground truth from peer's allocator uint64 delta_seq_baseline = 6; // valid for full-sync; master adopts this cursor repeated TierKvCount tier_kv_counts = 7; // live owned-key count per tier (client-reported) + repeated LogicalTierCapacity logical_tier_capacities = 8; } message HeartbeatResponse { ClientStatus status = 1; @@ -131,6 +166,7 @@ message RoutePutResponse { string node_id = 2; TierType tier = 3; string peer_address = 4; + string logical_tier = 5; } message RouteGetRequest { @@ -144,6 +180,7 @@ message RouteGetResponse { TierType tier = 3; uint64 size = 4; string peer_address = 5; + string logical_tier = 6; } // ============================================================ @@ -169,6 +206,7 @@ message BatchRoutePutEntry { string node_id = 2; TierType tier = 3; string peer_address = 4; + string logical_tier = 5; } message BatchRoutePutResponse { repeated BatchRoutePutEntry entries = 1; @@ -198,6 +236,7 @@ message BatchRouteGetResponse { repeated uint32 node_ref = 2; // per key; 1-based index into nodes, 0 = not found [packed] repeated TierType tier = 3; // per key [packed] repeated uint64 size = 4; // per key [packed] + repeated string logical_tier = 5; // per key } // Read-only batched existence probe. Unlike BatchRouteGet, this does diff --git a/src/umbp/distributed/proto/umbp_peer.proto b/src/umbp/distributed/proto/umbp_peer.proto index ff37fda7a..54680d9b3 100644 --- a/src/umbp/distributed/proto/umbp_peer.proto +++ b/src/umbp/distributed/proto/umbp_peer.proto @@ -6,14 +6,12 @@ import "umbp.proto"; service UMBPPeer { rpc GetPeerInfo(GetPeerInfoRequest) returns (GetPeerInfoResponse); - // SSD read staging. The direct-SSD-put RPCs (AllocateWriteSlot / - // CommitSsdWrite) were removed in the SSD-tier redesign: SSD is now an - // async copy-on-commit cold tier, not a writer-driven put target. - // PrepareSsdRead is key-based (the peer looks up the SSD location by key via - // PeerSsdManager) and returns a typed status + the actual size; ReleaseSsdLease - // is a best-effort fast release (slots are also reclaimed by lease TTL). - rpc PrepareSsdRead (PrepareSsdReadRequest) returns (PrepareSsdReadResponse); - rpc ReleaseSsdLease (ReleaseSsdLeaseRequest) returns (ReleaseSsdLeaseResponse); + // PrepareSsdRead / ReleaseSsdLease (SSD read-staging lease protocol) were + // removed in the backend-agnostic refactor Phase 0: SSD is unwired from the + // distributed data plane (see design-backend-agnostic-refactor.md). SSD + // returns post-refactor as an SsdBackend : MediumBackend adapter that will + // not reuse this lease shape — do not reuse these RPC names for anything + // else in the meantime. // DRAM/HBM allocator + key map (new — peer is now the canonical owner). rpc AllocateSlot(AllocateSlotRequest) returns (AllocateSlotResponse); @@ -36,52 +34,26 @@ message GetPeerInfoRequest {} message GetPeerInfoResponse { bytes engine_desc = 1; // SSD capacity is reported via heartbeat tier_capacities, not here. - bytes ssd_staging_mem_desc = 6; // packed MemoryDesc for dedicated SSD staging buffer - uint64 ssd_staging_size = 7; // total staging buffer size - - // DRAM/HBM buffer descriptors — peer-owned in the new design, exposed - // here so writers can hydrate first-contact peers without going through - // AllocateSlot or ResolveKey first. - repeated BufferMemoryDesc dram_memory_descs = 8; - uint64 dram_page_size = 9; -} - -// ============================================================ -// SSD read staging (key-based; see the UMBPPeer service comment above) -// ============================================================ - -// Per-key SSD read outcome. Callers MUST distinguish a transient failure -// (NO_SLOT) from a definitive miss (NOT_FOUND): a transient failure must never -// be surfaced to hicache as a cache miss. Lease expiry is NOT a wire status: -// the reader decides it locally against lease_ttl_ms and the peer reclaims the -// slot by TTL (see PrepareSsdReadResponse). -enum SsdReadStatus { - SSD_READ_OK = 0; - SSD_READ_NOT_FOUND = 1; // key not on this peer's SSD tier → caller treats as miss - SSD_READ_NO_SLOT = 2; // staging slots exhausted → transient/retryable, not a miss - SSD_READ_SIZE_TOO_LARGE = 3; // actual size > reader cap / staging slot size - SSD_READ_ERROR = 4; // backend read error (e.g. concurrent evict / corruption) -} - -message PrepareSsdReadRequest { - string key = 1; - // Reader's destination buffer size (capacity upper bound). The peer rejects - // with SSD_READ_SIZE_TOO_LARGE when the key's actual size exceeds it. - uint64 max_size = 2; -} -message PrepareSsdReadResponse { - SsdReadStatus status = 1; - uint64 staging_offset = 2; // offset into the published ssd_staging_mem_desc - uint64 size = 3; // actual readable bytes - uint64 lease_id = 4; - uint64 lease_ttl_ms = 5; -} - -message ReleaseSsdLeaseRequest { - uint64 lease_id = 1; -} -message ReleaseSsdLeaseResponse { - bool success = 1; + reserved 6, 7; // ssd_staging_mem_desc, ssd_staging_size — SSD read-staging + // buffer publication, removed with the lease protocol + // (Phase 0 of design-backend-agnostic-refactor.md) + reserved "ssd_staging_mem_desc", "ssd_staging_size"; + + reserved 8, 9; // dram_memory_descs, dram_page_size — a flat, single-medium + // buffer_index space that could advertise only ONE backend's + // buffers. A peer with two live media dropped the rest, and + // the reader then resolved the survivors' indices against the + // wrong memory. Replaced by buffer_descs, whose entries carry + // backend_id. + reserved "dram_memory_descs", "dram_page_size"; + + // Buffer descriptors for EVERY live backend, so a first-contact writer can + // hydrate without going through AllocateSlot or ResolveKey first. Each entry + // carries the backend_id its buffer_index belongs to; see BufferMemoryDesc. + repeated BufferMemoryDesc buffer_descs = 10; + // Uniform across backends — BackendRegistry refuses to register a backend + // that disagrees, because every wire carrying pages carries ONE page_size. + uint64 page_size = 11; } // ============================================================ @@ -96,6 +68,10 @@ message AllocateSlotRequest { uint64 size = 1; TierType tier = 2; string key = 3; // owned_ dedup (master-index-lag fallback) + // Optional explicit instance. When absent, the peer selects the first + // backend registered for `tier`. + optional string backend_name = 4; + string logical_tier = 5; } // AllocateSlot / BatchAllocateSlots per-entry outcome. // FAILED_NO_SPACE is split out so the writer can keep retrying with @@ -115,8 +91,11 @@ message AllocateSlotResponse { uint64 slot_id = 2; repeated PageLocation pages = 3; uint64 page_size = 4; - repeated BufferMemoryDesc descs = 5; // dedup'd by buffer_index + repeated BufferMemoryDesc descs = 5; // dedup'd by (backend_id, buffer_index) uint64 pending_ttl_ms = 6; // single TTL in the system + // Which backend `pages` index into. One per response, not one per page: a + // slot lives entirely in one medium. + uint32 backend_id = 7; } // Commit moves pending → owned. After a successful commit the peer @@ -150,6 +129,8 @@ message ResolveKeyResponse { uint64 page_size = 3; repeated BufferMemoryDesc descs = 4; uint64 size = 5; + // Which backend `pages` index into; a key lives entirely in one medium. + uint32 backend_id = 6; } // Master-driven eviction. Idempotent: already-evicted keys produce a @@ -198,12 +179,27 @@ message BatchAbortSlotsResponse { message BatchResolveKeysRequest { repeated string keys = 1; // When set, the peer omits the per-buffer descriptors from the response: the - // client has already hydrated buffer_index -> MemoryDesc from GetPeerInfo - // (dram_memory_descs) and the read path looks descriptors up in that cache. - // A wrongly-omitted descriptor degrades to a cache miss on the client (the - // transfer-build guard rejects the entry), never to a corrupt read. + // client has already hydrated (backend_id, buffer_index) -> MemoryDesc from + // GetPeerInfo (buffer_descs) and the read path looks descriptors up in that + // cache. A wrongly-omitted descriptor degrades to a cache miss on the client + // (the transfer-build guard rejects the entry), never to a corrupt read — + // which holds only because GetPeerInfo now publishes EVERY backend's buffers. + // While it published one medium's, this flag suppressed the descriptors that + // would have corrected the reader, and the miss became a wrong-memory read. bool omit_descs = 2; } + +enum ResolveStatus { + RESOLVE_STATUS_UNSPECIFIED = 0; + RESOLVE_STATUS_FOUND = 1; + RESOLVE_STATUS_MISSING = 2; + // The peer owns the key but cannot publish it until transient staging + // pressure clears. Retry the same peer; do not treat this as a cache miss. + RESOLVE_STATUS_BUSY = 3; + // Permanent backend/shape failure. Retrying unchanged input cannot help. + RESOLVE_STATUS_FAILED = 4; +} + // Struct-of-arrays layout. The previous shape repeated a full ResolveKeyResponse // per key, each carrying its own copy of the peer's buffer descriptors and // page_size — the descriptors and page_size are batch-invariant, so that @@ -214,8 +210,10 @@ message BatchResolveKeysRequest { // buffer_index/page_index with page_count[i] delimiting key i's slice. This // removes the per-key sub-message framing and the descriptor duplication. message BatchResolveKeysResponse { - // Deduplicated buffer descriptors, once per batch (empty when the request set - // omit_descs, or when no resolved key was found). + // Deduplicated buffer descriptors, once per batch, by (backend_id, + // buffer_index) (empty when the request set omit_descs, or when no resolved + // key was found). A batch may now span media, so these are NOT all from one + // backend. repeated BufferMemoryDesc descs = 1; // Peer-global page size. uint64 page_size = 2; @@ -224,6 +222,14 @@ message BatchResolveKeysResponse { repeated TierType tier = 4; repeated uint64 size = 5; repeated uint32 page_count = 6; + // Which backend key i's pages index into. Per KEY, not per page: a key lives + // entirely in one medium, so this costs one uint32 per key rather than one + // per page on the hottest array in the protocol. + repeated uint32 backend_id = 9; + // Optional for wire compatibility. New peers emit one status per key; new + // clients derive FOUND/MISSING from `found` when talking to an old peer that + // omits this field. + repeated ResolveStatus resolve_status = 10; // Flattened page locations across all keys: key i owns the next page_count[i] // (buffer_index, page_index) pairs, walking keys in order. repeated uint32 buffer_index = 7; diff --git a/src/umbp/distributed/proto/umbp_standalone.proto b/src/umbp/distributed/proto/umbp_standalone.proto index c094cb737..3dfb02a97 100644 --- a/src/umbp/distributed/proto/umbp_standalone.proto +++ b/src/umbp/distributed/proto/umbp_standalone.proto @@ -9,8 +9,10 @@ service UMBPStandalone { rpc Put(PutRequest) returns (BoolResponse); rpc Get(GetRequest) returns (BoolResponse); rpc BatchPut(BatchDataRequest) returns (BatchBoolResponse); + rpc BatchPutRanges(BatchRangeDataRequest) returns (BatchBoolResponse); rpc BatchPutWithDepth(BatchDataWithDepthRequest) returns (BatchBoolResponse); rpc BatchGet(BatchDataRequest) returns (BatchBoolResponse); + rpc BatchGetRanges(BatchRangeDataRequest) returns (BatchBoolResponse); rpc Exists(KeyRequest) returns (BoolResponse); rpc BatchExists(BatchKeysRequest) returns (BatchBoolResponse); rpc BatchExistsConsecutive(BatchKeysRequest) returns (CountResponse); @@ -33,6 +35,7 @@ message Empty {} message PingResponse { bool ready = 1; StandaloneBackendMode deployment_mode = 2; + bool supports_ranged_io = 3; } message BoolResponse { @@ -86,6 +89,20 @@ message BatchDataRequest { repeated uint64 region_bases = 5; } +message BatchRangeDataRequest { + repeated string keys = 1; + string client_id = 2; + // The remaining fields are flattened over entries. range_counts[i] records + // how many ranges belong to keys[i]. + repeated uint32 range_counts = 3; + repeated uint64 shm_offsets = 4; + repeated uint64 region_bases = 5; + repeated uint64 sizes = 6; + repeated uint64 object_offsets = 7; + // Put only: the expected full object size, one per key. Empty for get. + repeated uint64 object_sizes = 8; +} + message BatchDataWithDepthRequest { repeated string keys = 1; string client_id = 2; @@ -99,6 +116,11 @@ message BatchBoolResponse { repeated bool ok = 1; } +enum MemoryKind { + MEMORY_KIND_HOST_SHM = 0; + MEMORY_KIND_GPU_IPC = 1; +} + message RegisterMemoryRequest { string client_id = 1; uint64 worker_base = 2; @@ -106,6 +128,11 @@ message RegisterMemoryRequest { string worker_node_id = 4; string worker_node_address = 5; repeated string tags = 6; + MemoryKind kind = 7; + int32 device_id = 8; + bytes ipc_handle = 9; + uint64 ipc_offset = 10; + uint64 alloc_base = 11; } message DeregisterMemoryRequest { diff --git a/src/umbp/distributed/proto/umbp_workload.proto b/src/umbp/distributed/proto/umbp_workload.proto new file mode 100644 index 000000000..23574a60a --- /dev/null +++ b/src/umbp/distributed/proto/umbp_workload.proto @@ -0,0 +1,48 @@ +syntax = "proto3"; + +package umbp.benchmark; + +// This schema is persisted on disk. Existing field numbers and meanings must +// never be changed; add new fields when the format evolves. +message WorkloadTraceHeader { + enum TimeUnit { + TIME_UNIT_UNSPECIFIED = 0; + TIME_UNIT_NANOSECONDS = 1; + } + + uint32 schema_version = 1; + TimeUnit time_unit = 2; + uint64 seed = 3; + map metadata = 4; +} + +message WorkloadEvent { + enum Operation { + OPERATION_UNSPECIFIED = 0; + PUT = 1; + GET = 2; + } + + // What the recorded system answered. This is provenance, not an instruction: + // a replay reissues the operation regardless and reports its own result, so + // the two can be compared. Traces written before this field existed leave it + // unspecified, which readers must treat as "not recorded" rather than as a + // failure. + enum Outcome { + OUTCOME_UNSPECIFIED = 0; + OUTCOME_SUCCESS = 1; + // A GET that found nothing. Kept in the trace because dropping misses + // rewrites the workload into one with a perfect hit rate. + OUTCOME_MISS = 2; + OUTCOME_FAILURE = 3; + } + + uint64 relative_timestamp_ns = 1; + uint32 client_id = 2; + uint64 operation_id = 3; + Operation operation = 4; + string key = 5; + uint64 value_size = 6; + uint64 batch_id = 7; + Outcome outcome = 8; +} diff --git a/src/umbp/distributed/routing/route_get_strategy.cpp b/src/umbp/distributed/routing/route_get_strategy.cpp index c210902d4..f0f83bfc8 100644 --- a/src/umbp/distributed/routing/route_get_strategy.cpp +++ b/src/umbp/distributed/routing/route_get_strategy.cpp @@ -43,30 +43,6 @@ namespace { return oss.str(); } -// Lower rank = higher read priority. SSD is the slow cold tier, so it ranks -// last among the real tiers; UNKNOWN sorts after everything so a malformed -// location never wins over a usable one. -int TierReadRank(TierType tier) { - switch (tier) { - case TierType::HBM: - return 0; - case TierType::DRAM: - return 1; - case TierType::SSD: - return 2; - default: - return 3; - } -} - -// Pick a uniformly random element of a non-empty index list using the shared -// thread_local RNG. -size_t PickRandomIndex(const std::vector& indices) { - thread_local std::mt19937 rng{std::random_device{}()}; - std::uniform_int_distribution dist(0, indices.size() - 1); - return indices[dist(rng)]; -} - // Core selection cores shared by the single-key Select and the batched // BatchSelect. Precondition: `locations` is non-empty (callers guarantee it). Location PickRandomReplica(const std::vector& locations) { @@ -91,46 +67,32 @@ Location PickRandomReplica(const std::vector& locations) { return selected; } -Location PickTierPriorityReplica(const std::vector& locations, - const std::string& node_id) { - // Find the best (lowest-rank) tier present, then collect every replica on it - // and choose one at random so load still spreads within a tier. - int best_rank = TierReadRank(locations[0].tier); - for (const auto& loc : locations) { - best_rank = std::min(best_rank, TierReadRank(loc.tier)); - } - std::vector best_tier_indices; - for (size_t i = 0; i < locations.size(); ++i) { - if (TierReadRank(locations[i].tier) == best_rank) best_tier_indices.push_back(i); - } - - // Requester-local preference: if the asking node itself holds a replica in the - // best tier, serve it locally instead of a random peer. This is what makes - // cache_remote_fetches pay off — a node that re-cached a remotely-fetched block - // reads its own copy on the next Get (no RDMA), matching the local-first - // behavior of the pre-dual-scheme UMBPClient. Non-replica requesters still - // spread randomly. Empty node_id (unknown caller) falls through to random. +Location PickLocalPreferringReplica(const std::vector& locations, + const std::string& node_id) { + // No tier ranking: every medium is equivalent, so any replica is as good as + // any other and the choice is made on locality and load alone (backend- + // agnostic refactor Phase 4 — the HBM > DRAM > SSD order is deleted, not + // replaced). Re-introducing a preference means advertising it from the + // backend, not hardcoding a tier list here. + // + // Requester-local preference: if the asking node itself holds a replica, serve + // it locally instead of a random peer. This is what makes cache_remote_fetches + // pay off — a node that re-cached a remotely-fetched block reads its own copy + // on the next Get (no RDMA), matching the local-first behavior of the + // pre-dual-scheme UMBPClient. Non-replica requesters still spread randomly. + // Empty node_id (unknown caller) falls through to random. if (!node_id.empty()) { - for (size_t i : best_tier_indices) { - if (locations[i].node_id == node_id) { + for (const auto& loc : locations) { + if (loc.node_id == node_id) { MORI_UMBP_DEBUG( - "[TierPriorityRouteGetStrategy] requester-local hit node={} tier={} size={}", - locations[i].node_id, TierTypeName(locations[i].tier), locations[i].size); - return locations[i]; + "[LocalPreferringRouteGetStrategy] requester-local hit node={} tier={} size={}", + loc.node_id, TierTypeName(loc.tier), loc.size); + return loc; } } } - size_t choice = PickRandomIndex(best_tier_indices); - const auto& selected = locations[choice]; - // NOTE: temporarily disabled — the macro args are evaluated even when the log - // level suppresses output; SummarizeLocations() dominated the hot path. - // MORI_UMBP_DEBUG( - // "[TierPriorityRouteGetStrategy] {} candidates -> best_tier={} ({} replicas) choice node={} - // " "tier={} size={}, candidates=[{}]", locations.size(), TierTypeName(selected.tier), - // best_tier_indices.size(), selected.node_id, TierTypeName(selected.tier), selected.size, - // SummarizeLocations(locations)); - return selected; + return PickRandomReplica(locations); } } // namespace @@ -169,22 +131,23 @@ std::vector RandomRouteGetStrategy::BatchSelect( return out; } -Location TierPriorityRouteGetStrategy::Select(const std::vector& locations, - const std::string& node_id) { +Location LocalPreferringRouteGetStrategy::Select(const std::vector& locations, + const std::string& node_id) { if (locations.empty()) { MORI_UMBP_WARN( - "[TierPriorityRouteGetStrategy] received empty location set; returning default Location"); + "[LocalPreferringRouteGetStrategy] received empty location set; returning default " + "Location"); return {}; } - return PickTierPriorityReplica(locations, node_id); + return PickLocalPreferringReplica(locations, node_id); } -std::vector TierPriorityRouteGetStrategy::BatchSelect( +std::vector LocalPreferringRouteGetStrategy::BatchSelect( const std::vector>& per_key_locations, const std::string& node_id) { std::vector out(per_key_locations.size()); for (size_t i = 0; i < per_key_locations.size(); ++i) { if (per_key_locations[i].empty()) continue; // not routed; leave default - out[i] = PickTierPriorityReplica(per_key_locations[i], node_id); + out[i] = PickLocalPreferringReplica(per_key_locations[i], node_id); } return out; } diff --git a/src/umbp/distributed/routing/route_put_strategy.cpp b/src/umbp/distributed/routing/route_put_strategy.cpp index 1e47ccd45..4af071a1c 100644 --- a/src/umbp/distributed/routing/route_put_strategy.cpp +++ b/src/umbp/distributed/routing/route_put_strategy.cpp @@ -35,11 +35,26 @@ namespace mori::umbp { // --------------------------------------------------------------------------- namespace { -// SSD is intentionally excluded: there is no direct SSD put — the SSD copy is -// filled asynchronously by copy-on-commit. RoutePut must never steer a put at -// a tier with no direct-put semantics, even though SSD capacity is reported via -// heartbeat. -constexpr std::array kPutTierOrder = {TierType::HBM, TierType::DRAM}; +// One placement candidate: a (node, tier) pair with room for the block. +// +// There is no tier order. The former kPutTierOrder = {HBM, DRAM} (with SSD +// excluded as a non-direct-put target) was deleted in the backend-agnostic +// refactor Phase 4 rather than reimplemented: every medium is currently +// equivalent, so a tier a node advertises capacity for is a tier it accepts +// puts on. Re-introducing "this medium takes no direct puts" — which SSD will +// need when it returns as an SsdBackend — means the backend advertising it +// (design doc §3 BackendProperties::put_eligible), not a tier list here. +struct PlacementSlot { + size_t index; // into `candidates` + TierType tier; + std::string logical_tier; + uint64_t available; +}; + +uint64_t AllocatableBytes(const TierCapacity& capacity) { + if (capacity.max_allocatable_bytes == 0) return capacity.available_bytes; + return std::min(capacity.available_bytes, capacity.max_allocatable_bytes); +} std::string JoinStrings(const std::vector& items) { if (items.empty()) return ""; @@ -82,27 +97,41 @@ std::string SummarizeClientTiers(const std::vector& alive_clients) } // Indices of candidates that can fit block_size on a single @p tier. -std::vector CollectEligibleOnTier(const std::vector& candidates, - TierType tier, uint64_t block_size, - const std::unordered_set& exclude_nodes) { - std::vector indices; +// Every (node, tier) pair that can hold `block_size`. tier_capacities is a +// std::map, so the pairs come out in ascending TierType per node — the order is +// deterministic, which keeps most-available tie-breaks reproducible. +std::vector CollectEligibleSlots( + const std::vector& candidates, uint64_t block_size, + const std::unordered_set& exclude_nodes) { + std::vector slots; for (size_t i = 0; i < candidates.size(); ++i) { const auto& client = candidates[i]; if (exclude_nodes.count(client.node_id)) continue; - auto it = client.tier_capacities.find(tier); - if (it == client.tier_capacities.end()) continue; - if (it->second.available_bytes < block_size) continue; - indices.push_back(i); + if (!client.logical_tier_capacities.empty()) { + for (const auto& [name, logical] : client.logical_tier_capacities) { + if (!logical.put_eligible || AllocatableBytes(logical.capacity) < block_size) { + continue; + } + slots.push_back({i, logical.representative_tier, name, logical.capacity.available_bytes}); + } + } else { + for (const auto& [tier, cap] : client.tier_capacities) { + if (AllocatableBytes(cap) < block_size) continue; + slots.push_back({i, tier, {}, cap.available_bytes}); + } + } } - return indices; + return slots; } -RoutePutResult MakeRouted(const ClientRecord& client, TierType tier) { +RoutePutResult MakeRouted(const ClientRecord& client, TierType tier, + std::string logical_tier = {}) { return RoutePutResult{ .outcome = RoutePutOutcome::kRouted, .node_id = client.node_id, .peer_address = client.peer_address, .tier = tier, + .logical_tier = std::move(logical_tier), }; } @@ -110,10 +139,17 @@ RoutePutResult MakeRouted(const ClientRecord& client, TierType tier) { // 0 when the node or tier is absent. Used only to report the pre-deduction // figure in the routed INFO log. uint64_t LookupAvailableBytes(const std::vector& candidates, - const std::string& node_id, TierType tier) { + const std::string& node_id, TierType tier, + const std::string& logical_tier) { auto client_it = std::find_if(candidates.begin(), candidates.end(), [&](const ClientRecord& c) { return c.node_id == node_id; }); if (client_it == candidates.end()) return 0; + if (!logical_tier.empty()) { + auto logical_it = client_it->logical_tier_capacities.find(logical_tier); + return logical_it == client_it->logical_tier_capacities.end() + ? 0 + : logical_it->second.capacity.available_bytes; + } auto tier_it = client_it->tier_capacities.find(tier); if (tier_it == client_it->tier_capacities.end()) return 0; return tier_it->second.available_bytes; @@ -134,6 +170,18 @@ bool ApplyProjectedDeduction(std::vector& candidates, const RouteP result.node_id); return false; } + if (!result.logical_tier.empty()) { + auto logical_it = client_it->logical_tier_capacities.find(result.logical_tier); + if (logical_it == client_it->logical_tier_capacities.end() || + logical_it->second.capacity.available_bytes < block_size) { + MORI_UMBP_ERROR( + "[RoutePutStrategy] projected-deduction: selected logical tier absent/full on node {}", + result.node_id); + return false; + } + logical_it->second.capacity.available_bytes -= block_size; + return true; + } auto tier_it = client_it->tier_capacities.find(result.tier); if (tier_it == client_it->tier_capacities.end()) { MORI_UMBP_ERROR("[RoutePutStrategy] projected-deduction: selected tier absent on node {}", @@ -191,13 +239,22 @@ size_t ConfigurableRoutePutStrategy::PickWeighted(const std::vector& w std::optional ConfigurableRoutePutStrategy::TrySelectOnNodeTier( const std::vector& candidates, const std::string& node_id, TierType tier, - uint64_t block_size, const std::unordered_set& exclude_nodes) const { + uint64_t block_size, const std::unordered_set& exclude_nodes, + const std::string& logical_tier) const { if (node_id.empty() || exclude_nodes.count(node_id)) return std::nullopt; auto it = std::find_if(candidates.begin(), candidates.end(), [&](const ClientRecord& c) { return c.node_id == node_id; }); if (it == candidates.end()) return std::nullopt; + if (!logical_tier.empty()) { + auto cap = it->logical_tier_capacities.find(logical_tier); + if (cap == it->logical_tier_capacities.end() || !cap->second.put_eligible || + AllocatableBytes(cap->second.capacity) < block_size) { + return std::nullopt; + } + return MakeRouted(*it, cap->second.representative_tier, logical_tier); + } auto cap = it->tier_capacities.find(tier); - if (cap == it->tier_capacities.end() || cap->second.available_bytes < block_size) { + if (cap == it->tier_capacities.end() || AllocatableBytes(cap->second) < block_size) { return std::nullopt; } return MakeRouted(*it, tier); @@ -206,49 +263,68 @@ std::optional ConfigurableRoutePutStrategy::TrySelectOnNodeTier( std::optional ConfigurableRoutePutStrategy::TrySelectOnNode( const std::vector& candidates, const std::string& node_id, uint64_t block_size, const std::unordered_set& exclude_nodes) const { - for (TierType tier : kPutTierOrder) { - if (auto r = TrySelectOnNodeTier(candidates, node_id, tier, block_size, exclude_nodes)) { - return r; + if (node_id.empty() || exclude_nodes.count(node_id)) return std::nullopt; + auto it = std::find_if(candidates.begin(), candidates.end(), + [&](const ClientRecord& c) { return c.node_id == node_id; }); + if (it == candidates.end()) return std::nullopt; + + // With every medium equivalent, "the best tier on this node" is just the one + // with the most room; ascending-TierType map order breaks ties. + const TierCapacity* best = nullptr; + TierType best_tier = TierType::UNKNOWN; + std::string best_logical; + if (!it->logical_tier_capacities.empty()) { + for (const auto& [name, logical] : it->logical_tier_capacities) { + if (!logical.put_eligible) continue; + const auto& cap = logical.capacity; + if (AllocatableBytes(cap) < block_size) continue; + if (best == nullptr || cap.available_bytes > best->available_bytes) { + best = ∩ + best_tier = logical.representative_tier; + best_logical = name; + } + } + } else { + for (const auto& [tier, cap] : it->tier_capacities) { + if (AllocatableBytes(cap) < block_size) continue; + if (best == nullptr || cap.available_bytes > best->available_bytes) { + best = ∩ + best_tier = tier; + } } } - return std::nullopt; + if (best == nullptr) return std::nullopt; + return MakeRouted(*it, best_tier, best_logical); } std::optional ConfigurableRoutePutStrategy::SelectByAlgo( const std::vector& candidates, uint64_t block_size, const std::unordered_set& exclude_nodes, const std::optional& preferred_node) { - for (TierType tier : kPutTierOrder) { - // Node preference applies only within this tier, so a preferred node that is - // full on the faster tier never preempts a remote node that still has room - // there: tier priority is preserved. - if (preferred_node) { - if (auto r = - TrySelectOnNodeTier(candidates, *preferred_node, tier, block_size, exclude_nodes)) { - return r; - } + // A preferred node now wins outright when it has room anywhere: with no tier + // order left, there is no "remote HBM beats local DRAM" case for it to break. + if (preferred_node) { + if (auto r = TrySelectOnNode(candidates, *preferred_node, block_size, exclude_nodes)) { + return r; } - std::vector eligible = - CollectEligibleOnTier(candidates, tier, block_size, exclude_nodes); - if (eligible.empty()) continue; - - auto available = [&](size_t idx) { - return candidates[idx].tier_capacities.at(tier).available_bytes; - }; - size_t chosen = eligible.front(); - if (algo_ == SelectAlgo::kRandom) { - std::vector weights; - weights.reserve(eligible.size()); - for (size_t idx : eligible) weights.push_back(available(idx)); - chosen = eligible[PickWeighted(weights)]; - } else { - for (size_t idx : eligible) { - if (available(idx) > available(chosen)) chosen = idx; - } + } + + std::vector slots = CollectEligibleSlots(candidates, block_size, exclude_nodes); + if (slots.empty()) return std::nullopt; + + size_t chosen = 0; + if (algo_ == SelectAlgo::kRandom) { + std::vector weights; + weights.reserve(slots.size()); + for (const auto& slot : slots) weights.push_back(slot.available); + chosen = PickWeighted(weights); + } else { + for (size_t k = 1; k < slots.size(); ++k) { + if (slots[k].available > slots[chosen].available) chosen = k; } - return MakeRouted(candidates[chosen], tier); } - return std::nullopt; + return MakeRouted(candidates[slots[chosen].index], slots[chosen].tier, + slots[chosen].logical_tier); } std::vector> ConfigurableRoutePutStrategy::SelectBatch( @@ -283,6 +359,7 @@ std::vector> ConfigurableRoutePutStrategy::SelectB // pick as nodes fill. std::optional anchor_node; std::optional anchor_tier; // pinned only for the kSame whole-batch hit + std::string anchor_logical_tier; if (affinity_ == NodeAffinity::kLocal) { if (!requester_node_id.empty()) anchor_node = requester_node_id; } else if (affinity_ == NodeAffinity::kSame) { @@ -296,6 +373,7 @@ std::vector> ConfigurableRoutePutStrategy::SelectB if (auto whole = SelectByAlgo(candidates, total, exclude_nodes)) { anchor_node = whole->node_id; anchor_tier = whole->tier; + anchor_logical_tier = whole->logical_tier; } } } @@ -330,7 +408,7 @@ std::vector> ConfigurableRoutePutStrategy::SelectB std::optional selected; if (affinity_ == NodeAffinity::kLocal) { - // Node-first: local HBM -> local DRAM, then global HBM -> DRAM fallback. + // Node-first: the local node's roomiest tier, then a global fallback. if (anchor_node) { selected = TrySelectOnNode(candidates, *anchor_node, block_size, exclude_nodes); } @@ -338,15 +416,15 @@ std::vector> ConfigurableRoutePutStrategy::SelectB } else if (affinity_ == NodeAffinity::kSame) { // Whole-batch hit: keep every key on the pinned node AND tier. if (anchor_tier) { - selected = - TrySelectOnNodeTier(candidates, *anchor_node, *anchor_tier, block_size, exclude_nodes); + selected = TrySelectOnNodeTier(candidates, *anchor_node, *anchor_tier, block_size, + exclude_nodes, anchor_logical_tier); } if (!selected) { - // Tier-first with the sticky node only preferred within each tier, so a - // spill never prefers the anchor's DRAM over a remote node's HBM. Drop - // the tier pin and re-anchor to wherever this key actually landed. + // The pinned tier is full. Drop the tier pin, keep preferring the + // sticky node, and re-anchor to wherever this key actually landed. selected = SelectByAlgo(candidates, block_size, exclude_nodes, anchor_node); anchor_tier = std::nullopt; + anchor_logical_tier.clear(); if (selected && selected->outcome == RoutePutOutcome::kRouted) { anchor_node = selected->node_id; } @@ -360,8 +438,8 @@ std::vector> ConfigurableRoutePutStrategy::SelectB // failure, never as a route (capture available_bytes before the deduction // so the INFO reflects the snapshot the decision was made on). if (selected && selected->outcome == RoutePutOutcome::kRouted) { - const uint64_t available = - LookupAvailableBytes(candidates, selected->node_id, selected->tier); + const uint64_t available = LookupAvailableBytes(candidates, selected->node_id, selected->tier, + selected->logical_tier); if (ApplyProjectedDeduction(candidates, *selected, block_size)) { MORI_UMBP_INFO( "[RoutePutStrategy] block_size={} tier={} selected node={} available_bytes={} " diff --git a/src/umbp/distributed/routing/router.cpp b/src/umbp/distributed/routing/router.cpp index 7f69e85e7..11f5edd9f 100644 --- a/src/umbp/distributed/routing/router.cpp +++ b/src/umbp/distributed/routing/router.cpp @@ -30,12 +30,11 @@ namespace mori::umbp { Router::Router(IMasterMetadataStore& store, std::unique_ptr get_strategy, std::unique_ptr put_strategy) : store_(store) { - // Default to tier-priority (HBM > DRAM > SSD): with the SSD cold tier live, a - // random pick could route a key that also has a DRAM/HBM copy to the slow - // SSD. Callers can still inject RandomRouteGetStrategy (or any other) via - // config_.get_strategy. + // Default to locality-preferring: a requester that already holds a replica + // reads its own copy instead of RDMA-ing a peer's. Callers can still inject + // RandomRouteGetStrategy (or any other) via config_.get_strategy. get_strategy_ = - get_strategy ? std::move(get_strategy) : std::make_unique(); + get_strategy ? std::move(get_strategy) : std::make_unique(); // Default to most-available / no-affinity: the single built-in put strategy. // Callers can still inject any RoutePutStrategy (e.g. an env-configured // ConfigurableRoutePutStrategy from master startup) via config_.put_strategy. diff --git a/src/umbp/distributed/transfer/composite_transfer_engine.cpp b/src/umbp/distributed/transfer/composite_transfer_engine.cpp new file mode 100644 index 000000000..c6434687e --- /dev/null +++ b/src/umbp/distributed/transfer/composite_transfer_engine.cpp @@ -0,0 +1,405 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/transfer/composite_transfer_engine.h" + +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" + +namespace mori::umbp { + +// Wraps the per-engine handles a single Submit fanned out into. +// +// Also the transfer layer's measurement point. Each sub-handle is tagged with +// the engine that produced it and the shape of what it carries, so Wait can +// charge bytes, failures and in-flight time to that engine — the composite is +// the only object that knows the mapping, and Wait is the only moment the +// outcome is known. +class CompositeTransferEngine::FanOutHandle final : public TransferHandle { + public: + // One sub-handle plus what it costs. `bytes` and `plans` are known at submit + // time; how many plans failed is known only after Wait. + struct Charge { + size_t engine_index = 0; + TransferDirection dir = TransferDirection::kLocal; + uint64_t bytes = 0; + uint64_t plans = 0; + }; + + explicit FanOutHandle(CompositeTransferEngine* owner) : owner_(owner) {} + + void Add(std::unique_ptr handle, std::vector charges) { + if (handle == nullptr) { + // Nothing posted: the caller fails these tags, and so do the metrics. + for (auto& c : charges) { + owner_->RecordSettled(c.engine_index, c.dir, /*bytes=*/0, /*nanos=*/0, c.plans); + } + return; + } + handles_.push_back({std::move(handle), std::move(charges)}); + } + void AddFailure(TransferFailure f) { failures_.push_back(std::move(f)); } + bool Empty() const { return handles_.empty() && failures_.empty(); } + + void Wait(std::vector* failures) override { + if (drained_) return; + drained_ = true; + if (failures != nullptr) { + for (auto& f : failures_) failures->push_back(std::move(f)); + } + failures_.clear(); + // Wait every sub-handle, never breaking early: an unwaited handle would + // leave live statuses behind. + for (auto& entry : handles_) { + std::vector mine; + const auto start = std::chrono::steady_clock::now(); + entry.handle->Wait(&mine); + const uint64_t nanos = + static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count()); + + // One TransferFailure per failed plan (see TransferHandle::Wait), so the + // count maps onto the plans this engine was charged. A handle carries at + // most one charge per direction, and a submit is usually uniform in + // direction, so the fill order below is exact in practice; when one + // engine posted both a push and a pull in the same submit, the failure + // total is right and its split across the two directions can be off. + // Wait does not say which plan failed, and adding a plan id to + // TransferFailure to sharpen a metric would be paying in the interface + // for something no caller needs. + uint64_t failed = static_cast(mine.size()); + for (auto& c : entry.charges) { + const uint64_t attributed = failed >= c.plans ? c.plans : failed; + failed -= attributed; + // Time is charged once per (engine, direction) group, not per plan: it + // is wall time in flight, and the plans inside a group overlap. + owner_->RecordSettled(c.engine_index, c.dir, c.bytes, nanos, attributed); + } + if (failures != nullptr) { + for (auto& f : mine) failures->push_back(std::move(f)); + } + } + } + + private: + struct Entry { + std::unique_ptr handle; + std::vector charges; + }; + + CompositeTransferEngine* owner_ = nullptr; + std::vector handles_; + std::vector failures_; + bool drained_ = false; +}; + +void CompositeTransferEngine::AddEngine(std::unique_ptr engine) { + if (engine == nullptr) return; + engines_.push_back(std::move(engine)); + // Grown in lockstep so an engine's index is its counters' index. This is the + // whole registration cost of instrumenting a new transport. + counters_.push_back(std::make_unique()); +} + +void CompositeTransferEngine::RecordSettled(size_t engine_index, TransferDirection dir, + uint64_t bytes, uint64_t nanos, uint64_t failed_plans) { + if (engine_index >= counters_.size()) return; + DirectionCounters& c = counters_[engine_index]->by_direction[DirectionIndex(dir)]; + if (bytes > 0) c.bytes.fetch_add(bytes, std::memory_order_relaxed); + if (nanos > 0) c.nanos.fetch_add(nanos, std::memory_order_relaxed); + if (failed_plans > 0) c.failed_plans.fetch_add(failed_plans, std::memory_order_relaxed); +} + +TransferRef CompositeTransferEngine::RegisterMemory(void* base, size_t size, + mori::io::MemoryLocationType loc, int device) { + TransferRef merged; + for (auto& engine : engines_) { + TransferRef part = engine->RegisterMemory(base, size, loc, device); + if (part.HasHostPtr() && !merged.HasHostPtr()) { + merged.host_ptr = part.host_ptr; + merged.size = part.size; + merged.loc = part.loc; + merged.device = part.device; + } + if (part.HasMemoryDesc()) { + if (!merged.HasMemoryDesc()) { + merged.mem = part.mem; + } else { + // TransferRef currently carries ONE mori::io::MemoryDesc, so a second + // engine's registration has nowhere to go — and Deregister would never + // release it. Unreachable with today's engine set (only MoriIoEngine + // registers); it becomes real with a second registering transport, at + // which point TransferRef needs a handle per transport the way + // mori::io::MemoryDesc already keeps ipcHandle beside fabricHandle. + MORI_UMBP_WARN( + "[CompositeTransferEngine] '{}' also registered {} bytes but TransferRef holds only " + "one descriptor; that registration will leak", + engine->Name(), size); + engine->Deregister(part); + } + } + } + return merged; +} + +TransferRef CompositeTransferEngine::RegisterFile(int fd, uint64_t offset, uint64_t size) { + // Unlike RegisterMemory, a file range fans out to exactly one taker: only a + // file-capable engine (GdsEngine) claims it, and it owns the fd registration + // outright. First match wins; an invalid ref means no file engine is + // configured and the caller must fall back to a memory path. + for (auto& engine : engines_) { + TransferRef part = engine->RegisterFile(fd, offset, size); + if (part.IsFile()) return part; + } + return TransferRef{}; +} + +void CompositeTransferEngine::Deregister(const TransferRef& ref) { + for (auto& engine : engines_) engine->Deregister(ref); +} + +TransferEngine* CompositeTransferEngine::SelectEngine(const TransferRef& src, + const TransferRef& dst) const { + for (const auto& engine : engines_) { + if (engine->CanHandle(src, dst)) return engine.get(); + } + return nullptr; +} + +bool CompositeTransferEngine::CanHandle(const TransferRef& src, const TransferRef& dst) const { + return SelectEngine(src, dst) != nullptr; +} + +TransferPlanSet CompositeTransferEngine::Plan(const std::vector& items) const { + TransferPlanSet out; + if (items.empty()) return out; + + auto adopt = [&out](TransferPlanSet sub, TransferEngine* engine) { + for (auto& plan : sub.plans) { + plan.engine = engine; + out.plans.push_back(std::move(plan)); + } + out.rejected_tags.insert(out.rejected_tags.end(), sub.rejected_tags.begin(), + sub.rejected_tags.end()); + }; + + // Fast path: every item goes to the same engine, so delegate the caller's own + // vector. Worth a separate path because a TransferItem holds two + // TransferRefs and each carries a MemoryDesc with three std::strings — + // partitioning by copy would allocate per page on the hot path, and in + // practice a batch is uniform (one peer group is all-remote, a self-target + // batch is all-local). + TransferEngine* first = SelectEngine(items.front().src, items.front().dst); + if (first != nullptr) { + bool uniform = true; + for (size_t i = 1; i < items.size() && uniform; ++i) { + uniform = SelectEngine(items[i].src, items[i].dst) == first; + } + if (uniform) { + adopt(first->Plan(items), first); + return out; + } + } + + // Mixed batch: partition, then plan per engine in ONE call each — grouping is + // what Plan is for, and planning item-by-item would defeat it. + std::unordered_map> by_engine; + for (const auto& item : items) { + TransferEngine* engine = SelectEngine(item.src, item.dst); + if (engine == nullptr) { + out.rejected_tags.push_back(item.tag); + rejected_items_.fetch_add(1, std::memory_order_relaxed); + continue; + } + by_engine[engine].push_back(item); + } + + for (auto& [engine, engine_items] : by_engine) { + adopt(engine->Plan(engine_items), engine); + } + return out; +} + +size_t CompositeTransferEngine::IndexOf(const TransferEngine* engine) const { + for (size_t i = 0; i < engines_.size(); ++i) { + if (engines_[i].get() == engine) return i; + } + return engines_.size(); // not ours; RecordSettled drops it +} + +std::unique_ptr CompositeTransferEngine::Submit(std::vector plans) { + if (plans.empty()) return nullptr; + + std::unordered_map> by_engine; + auto handle = std::make_unique(this); + for (auto& plan : plans) { + TransferEngine* engine = plan.engine; + if (engine == nullptr) { + // A plan this composite did not produce. Fail its keys rather than + // guessing an engine: the endpoints alone cannot tell us which engine + // built the offsets (a bounce plan's local offsets mean nothing outside + // the engine that owns that pool). + TransferFailure f; + f.tags = plan.tags; + f.message = "plan has no owning engine"; + handle->AddFailure(std::move(f)); + continue; + } + by_engine[engine].push_back(std::move(plan)); + } + + for (auto& [engine, engine_plans] : by_engine) { + const size_t engine_index = IndexOf(engine); + + // Snapshot the tags BEFORE the move so a Submit that posts nothing can + // still fail exactly the keys it dropped. The metric charges are built in + // the same pass, for the same reason. + std::vector tags; + FanOutHandle::Charge charges[kDirectionCount]; + for (size_t d = 0; d < kDirectionCount; ++d) { + charges[d].engine_index = engine_index; + charges[d].dir = static_cast(d); + } + for (const auto& plan : engine_plans) { + tags.insert(tags.end(), plan.tags.begin(), plan.tags.end()); + FanOutHandle::Charge& c = charges[DirectionIndex(plan.dir)]; + ++c.plans; + for (size_t size : plan.sizes) c.bytes += size; + } + + std::vector live; + for (auto& c : charges) { + if (c.plans == 0) continue; + // Posted count goes in now; how many of them failed is decided in Wait. + if (engine_index < counters_.size()) { + counters_[engine_index]->by_direction[DirectionIndex(c.dir)].plans.fetch_add( + c.plans, std::memory_order_relaxed); + } + live.push_back(c); + } + + auto sub = engine->Submit(std::move(engine_plans)); + if (sub != nullptr) { + handle->Add(std::move(sub), std::move(live)); + continue; + } + // Submit posted nothing: the caller will never get a completion for these. + handle->Add(nullptr, std::move(live)); + TransferFailure f; + f.message = std::string(engine->Name()) + " submitted nothing"; + f.tags = std::move(tags); + handle->AddFailure(std::move(f)); + } + + if (handle->Empty()) return nullptr; + return handle; +} + +namespace { + +const char* DirectionName(TransferDirection dir) { + switch (dir) { + case TransferDirection::kPush: + return "push"; + case TransferDirection::kPull: + return "pull"; + case TransferDirection::kLocal: + return "local"; + default: + return "unknown"; + } +} + +constexpr double kNanosPerSecond = 1e9; + +} // namespace + +std::vector CompositeTransferEngine::SampleMetrics() const { + std::vector out; + + for (size_t i = 0; i < engines_.size() && i < counters_.size(); ++i) { + const char* engine_name = engines_[i]->Name(); + + for (size_t d = 0; d < kDirectionCount; ++d) { + const DirectionCounters& c = counters_[i]->by_direction[d]; + const uint64_t plans = c.plans.load(std::memory_order_relaxed); + if (plans == 0) continue; + const char* dir = DirectionName(static_cast(d)); + const uint64_t failed = c.failed_plans.load(std::memory_order_relaxed); + + out.push_back(MetricSample{MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL, + MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL_HELP, + {{"engine", engine_name}, {"direction", dir}, {"status", "ok"}}, + // Plans posted but not yet settled count as ok until Wait says + // otherwise; the correction lands on the next tick. + plans >= failed ? plans - failed : 0}); + if (failed > 0) { + out.push_back( + MetricSample{MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL, + MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL_HELP, + {{"engine", engine_name}, {"direction", dir}, {"status", "failed"}}, + failed}); + } + + const uint64_t bytes = c.bytes.load(std::memory_order_relaxed); + if (bytes > 0) { + out.push_back(MetricSample{MORI_UMBP_METRIC_TRANSFER_BYTES_TOTAL, + MORI_UMBP_METRIC_TRANSFER_BYTES_TOTAL_HELP, + {{"engine", engine_name}, {"direction", dir}}, + bytes}); + } + const uint64_t nanos = c.nanos.load(std::memory_order_relaxed); + if (nanos > 0) { + out.push_back(MetricSample{MORI_UMBP_METRIC_TRANSFER_SECONDS_TOTAL, + MORI_UMBP_METRIC_TRANSFER_SECONDS_TOTAL_HELP, + {{"engine", engine_name}, {"direction", dir}}, + nanos, + MetricKind::kCounter, + 1.0 / kNanosPerSecond}); + } + } + + // Whatever the engine publishes about its own internals, stamped with the + // engine that produced it so it shares the dashboard's engine= dimension. + for (MetricSample& s : engines_[i]->SampleMetrics()) { + s.labels.insert(s.labels.begin(), {"engine", engine_name}); + out.push_back(std::move(s)); + } + } + + // Items no engine would take. Not an engine failure — a routing one — so it + // gets its own engine value rather than being blamed on whoever ran last. + const uint64_t rejected = rejected_items_.load(std::memory_order_relaxed); + if (rejected > 0) { + out.push_back(MetricSample{MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL, + MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL_HELP, + {{"engine", "none"}, {"direction", "none"}, {"status", "rejected"}}, + rejected}); + } + return out; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/transfer/gds_engine.cpp b/src/umbp/distributed/transfer/gds_engine.cpp new file mode 100644 index 000000000..313489ac3 --- /dev/null +++ b/src/umbp/distributed/transfer/gds_engine.cpp @@ -0,0 +1,180 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/transfer/gds_engine.h" + +#include + +#include +#include +#include + +#include "mori/utils/mori_log.hpp" + +namespace mori::umbp { +namespace { + +inline bool IsGpu(const TransferRef& r) { return r.loc == mori::io::MemoryLocationType::GPU; } + +// Inline-completion handle, same shape as the copy engines': every read runs +// before Submit returns, so the handle only replays the outcome. +class SettledHandle final : public TransferHandle { + public: + explicit SettledHandle(std::vector failures) : failures_(std::move(failures)) {} + + void Wait(std::vector* failures) override { + if (reported_) return; + reported_ = true; + if (failures != nullptr) { + for (auto& f : failures_) failures->push_back(std::move(f)); + } + failures_.clear(); + } + + private: + std::vector failures_; + bool reported_ = false; +}; + +// Turn a hipFileRead return code into a human-readable reason. >= 0 is a byte +// count (a short read here); -1 is a POSIX errno; any other negative is the +// negated hipFileOpError_t. +std::string ReadFailureReason(ssize_t n, size_t want) { + if (n == -1) return std::string("errno=") + std::to_string(errno); + if (n < 0) return hipFileGetOpErrorString(static_cast(-n)); + return std::string("short read ") + std::to_string(n) + "/" + std::to_string(want); +} + +} // namespace + +GdsEngine::~GdsEngine() { + std::lock_guard lock(handles_mutex_); + for (auto& [fd, entry] : handles_) { + if (entry.handle != nullptr) { + hipFileHandleDeregister(static_cast(entry.handle)); + } + } + handles_.clear(); +} + +TransferRef GdsEngine::RegisterFile(int fd, uint64_t offset, uint64_t size) { + if (fd < 0) return TransferRef{}; + std::lock_guard lock(handles_mutex_); + + auto it = handles_.find(fd); + if (it != handles_.end()) { + ++it->second.refcount; + return TransferRef::File(fd, offset, size, it->second.handle); + } + + hipFileDescr_t descr{}; + descr.type = hipFileHandleTypeOpaqueFD; + descr.handle.fd = fd; + hipFileHandle_t handle = nullptr; + hipFileError_t err = hipFileHandleRegister(&handle, &descr); + if (err.err != hipFileSuccess || handle == nullptr) { + MORI_UMBP_ERROR("[GdsEngine] hipFileHandleRegister(fd={}) failed: {}", fd, + hipFileGetOpErrorString(err.err)); + return TransferRef{}; + } + handles_.emplace(fd, HandleEntry{static_cast(handle), 1}); + return TransferRef::File(fd, offset, size, static_cast(handle)); +} + +void GdsEngine::Deregister(const TransferRef& ref) { + if (!ref.IsFile()) return; + std::lock_guard lock(handles_mutex_); + auto it = handles_.find(ref.file_fd); + if (it == handles_.end()) return; + if (--it->second.refcount <= 0) { + if (it->second.handle != nullptr) { + hipFileHandleDeregister(static_cast(it->second.handle)); + } + handles_.erase(it); + } +} + +bool GdsEngine::CanHandle(const TransferRef& src, const TransferRef& dst) const { + // Read path only: a file source into a device buffer. A file destination or + // a host destination is not ours. + return src.IsFile() && !dst.IsFile() && dst.HasHostPtr() && IsGpu(dst); +} + +TransferPlanSet GdsEngine::Plan(const std::vector& items) const { + TransferPlanSet out; + out.plans.reserve(items.size()); + for (const auto& item : items) { + if (item.size == 0) continue; + if (!CanHandle(item.src, item.dst)) { + out.rejected_tags.push_back(item.tag); + continue; + } + // Bound the read to the destination buffer: an overrun into a device pool + // is silent corruption, not a segfault, so reject it here rather than let + // hipFileRead scribble past the slot. + if (item.size > item.dst.size || item.dst_offset > item.dst.size - item.size) { + out.rejected_tags.push_back(item.tag); + continue; + } + TransferPlan plan; + plan.src = item.src; + plan.dst = item.dst; + plan.dir = TransferDirection::kLocal; + plan.src_offsets.push_back(static_cast(item.src_offset)); + plan.dst_offsets.push_back(static_cast(item.dst_offset)); + plan.sizes.push_back(static_cast(item.size)); + plan.tags.push_back(item.tag); + out.plans.push_back(std::move(plan)); + } + return out; +} + +std::unique_ptr GdsEngine::Submit(std::vector plans) { + if (plans.empty()) return nullptr; + std::vector failures; + + for (const auto& plan : plans) { + auto fh = static_cast(plan.src.gds_handle); + if (fh == nullptr) { + failures.push_back( + TransferFailure{plan.tags, 0, "GdsEngine: file ref has no hipFile handle", "gds"}); + continue; + } + // buffer_base is the GPU allocation base; buffer_offset selects the slot + // inside it, and file_offset is the range's absolute position on disk. + void* dst_base = plan.dst.host_ptr; + for (size_t i = 0; i < plan.sizes.size(); ++i) { + const uint64_t file_off = plan.src.file_offset + plan.src_offsets[i]; + const ssize_t n = hipFileRead(fh, dst_base, plan.sizes[i], static_cast(file_off), + static_cast(plan.dst_offsets[i])); + if (n < 0 || static_cast(n) != plan.sizes[i]) { + const std::string why = ReadFailureReason(n, plan.sizes[i]); + MORI_UMBP_ERROR("[GdsEngine] hipFileRead file_off={} size={} failed: {}", file_off, + plan.sizes[i], why); + failures.push_back(TransferFailure{plan.tags, 0, "GdsEngine: hipFileRead: " + why, "gds"}); + break; // one failure per plan; a partially-read key is failed wholesale + } + } + } + return std::make_unique(std::move(failures)); +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/transfer/hbm_copy_engine.cpp b/src/umbp/distributed/transfer/hbm_copy_engine.cpp new file mode 100644 index 000000000..7edf9f38c --- /dev/null +++ b/src/umbp/distributed/transfer/hbm_copy_engine.cpp @@ -0,0 +1,738 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/transfer/hbm_copy_engine.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/common/device_gather.h" + +namespace mori::umbp { + +namespace { + +// True iff [next, ...) is exactly adjacent after [base, base+len). Same guard +// LocalCopyEngine uses; duplicated rather than shared because the two engines +// are otherwise independent and this is three lines. +inline bool AdjacentNoOverflow(size_t base, size_t len, size_t next) { + return len <= std::numeric_limits::max() - base && base + len == next; +} + +struct PtrPairKey { + const void* src; + const void* dst; + bool operator==(const PtrPairKey& o) const noexcept { return src == o.src && dst == o.dst; } +}; +struct PtrPairKeyHash { + size_t operator()(const PtrPairKey& k) const noexcept { + size_t h = std::hash{}(k.src); + h ^= std::hash{}(k.dst) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); + return h; + } +}; + +inline bool IsGpu(const TransferRef& r) { return r.loc == mori::io::MemoryLocationType::GPU; } + +// The copy kind is a function of the PAIR, like engine selection itself. +// Spelled explicitly rather than relying on hipMemcpyDefault + unified +// addressing so a mis-registered endpoint (loc=GPU on host memory, say) fails +// loudly at the copy instead of silently taking a slower inferred path. +inline hipMemcpyKind KindFor(const TransferRef& src, const TransferRef& dst) { + const bool s = IsGpu(src); + const bool d = IsGpu(dst); + if (s && d) return hipMemcpyDeviceToDevice; + if (s) return hipMemcpyDeviceToHost; + if (d) return hipMemcpyHostToDevice; + return hipMemcpyHostToHost; // unreachable: CanHandle requires one GPU side. +} + +// --------------------------------------------------------------------------- +// Gather fast path +// +// A plan that decomposed into many small segments is the case hipMemcpy is +// worst at: each call costs a fixed submission (~5.4 us with 8 ranks active) +// that a 4 KiB copy cannot amortize. One kernel doing all of them measures +// ~0.57 us per fragment and is flat in fragment size — see the rationale on +// LaunchDeviceGather. Above roughly 128 KiB per fragment the copy engine wins +// again, because one large hipMemcpyAsync already amortizes its own +// submission. +// +// This is what makes ranged I/O worth having on a GPU caller: reading one +// layer out of page-granular objects is exactly "dozens of small strided +// fragments", and it is why device_gather moved out of the DRAM tier into +// common/ (upstream 93f2998a). +// --------------------------------------------------------------------------- + +constexpr size_t kGatherFragmentThreshold = 128ULL << 10; + +// Mirrors device_gather.hip's kVectorBytes. Used ONLY to report, in debug +// mode, whether a segment would take the kernel's uint4 fast path or fall to +// its byte loop; it never affects a copy. If the kernel's vector width ever +// changes, this becomes a reporting inaccuracy, not a correctness bug. +constexpr size_t kVectorBytes = 16; + +// --------------------------------------------------------------------------- +// Debug mode (UMBP_HBM_COPY_DEBUG=1) +// +// Turns this engine into an instrument. For every batch it reports which of +// the two paths carried it, why the other one declined, the exact segment +// geometry needed to rebuild the case in a unit test, and the bandwidth +// ACTUALLY achieved. +// +// On measurement honesty: the per-fragment costs quoted in device_gather.h +// (~0.57 us, "flat across fragment sizes from 8 KiB to 221 KiB") are host-side +// submission costs, not data-movement costs -- a figure that is flat across a +// 27x size range cannot be a bandwidth. So dividing bytes by it overstates +// throughput badly. What is timed here instead is launch-to-completion, and +// both paths are already synchronous at this point -- the gather bucket does +// its own hipStreamSynchronize, and the fallback loop uses blocking hipMemcpy +// -- so enabling debug adds clock reads and logging but imposes NO extra +// synchronization and does not itself change which path runs. +// +// Knobs (all read once, on first use): +// UMBP_HBM_COPY_DEBUG=1 enable +// UMBP_HBM_COPY_DEBUG_SAMPLE=N emit a per-call line every Nth call (1) +// UMBP_HBM_COPY_DEBUG_FRAGMENTS=N per-call segment detail lines (16; 0=off) +// UMBP_HBM_COPY_DEBUG_SUMMARY_SEC=N cumulative rollup cadence (10; 0=off) +// +// Sampling exists because a serving run calls this millions of times; the +// rollup is what you read for a bandwidth answer, the per-call lines are what +// you read to build the reproducer. +// --------------------------------------------------------------------------- + +bool DebugEnvFlag(const char* name, bool fallback) { + const char* value = std::getenv(name); + if (value == nullptr) return fallback; + const std::string text(value); + return text != "0" && text != "off" && text != "false"; +} + +size_t DebugEnvSize(const char* name, size_t fallback) { + const char* value = std::getenv(name); + if (value == nullptr) return fallback; + char* end = nullptr; + const unsigned long long parsed = std::strtoull(value, &end, 10); + if (end == value || parsed > std::numeric_limits::max()) return fallback; + return static_cast(parsed); +} + +bool DebugEnabled() { + static const bool enabled = DebugEnvFlag("UMBP_HBM_COPY_DEBUG", false); + return enabled; +} + +size_t DebugSampleEvery() { + static const size_t n = std::max(DebugEnvSize("UMBP_HBM_COPY_DEBUG_SAMPLE", 1), 1); + return n; +} + +size_t DebugMaxSegments() { + static const size_t n = DebugEnvSize("UMBP_HBM_COPY_DEBUG_FRAGMENTS", 16); + return n; +} + +size_t DebugSummarySeconds() { + static const size_t n = DebugEnvSize("UMBP_HBM_COPY_DEBUG_SUMMARY_SEC", 10); + return n; +} + +// One counter for the whole engine, so sampling thins the combined stream +// rather than letting a rarely-taken path go unseen behind a common one. +bool DebugSampleThisCall() { + static std::atomic counter{0}; + return (counter.fetch_add(1, std::memory_order_relaxed) % DebugSampleEvery()) == 0; +} + +const char* KindName(hipMemcpyKind kind) { + switch (kind) { + case hipMemcpyHostToDevice: + return "H2D"; + case hipMemcpyDeviceToHost: + return "D2H"; + case hipMemcpyDeviceToDevice: + return "D2D"; + case hipMemcpyHostToHost: + return "H2H"; + default: + return "?"; + } +} + +inline double Seconds(std::chrono::steady_clock::time_point start, + std::chrono::steady_clock::time_point end) { + return std::chrono::duration_cast>(end - start).count(); +} + +inline double GiBps(double bytes, double seconds) { + if (seconds <= 0.0) return 0.0; + return (bytes / seconds) / static_cast(1ULL << 30); +} + +// Cumulative per-path totals. A serving run's per-call numbers are noisy and +// bimodal (a cold page fault dwarfs a warm copy), so the aggregate is the one +// worth quoting. +class DebugStats { + public: + void Record(bool gather, uint64_t segments, double bytes, double seconds) { + std::lock_guard lock(mutex_); + Totals& t = gather ? gather_ : memcpy_; + t.calls += 1; + t.segments += segments; + t.bytes += bytes; + t.seconds += seconds; + ReportIfDueLocked(); + } + + private: + struct Totals { + uint64_t calls = 0; + uint64_t segments = 0; + double bytes = 0.0; + double seconds = 0.0; + }; + + void ReportIfDueLocked() { + const size_t cadence = DebugSummarySeconds(); + if (cadence == 0) return; + const auto now = std::chrono::steady_clock::now(); + if (Seconds(last_report_, now) < static_cast(cadence)) return; + last_report_ = now; + MORI_UMBP_INFO( + "[HbmCopyEngine][dbg] SUMMARY gather[calls={} segs={} bytes={:.3f}GiB busy={:.3f}s " + "mean_seg={}B agg={:.3f}GiB/s] hipMemcpy[calls={} segs={} bytes={:.3f}GiB busy={:.3f}s " + "mean_seg={}B agg={:.3f}GiB/s]", + gather_.calls, gather_.segments, gather_.bytes / static_cast(1ULL << 30), + gather_.seconds, gather_.segments ? gather_.bytes / gather_.segments : 0.0, + GiBps(gather_.bytes, gather_.seconds), memcpy_.calls, memcpy_.segments, + memcpy_.bytes / static_cast(1ULL << 30), memcpy_.seconds, + memcpy_.segments ? memcpy_.bytes / memcpy_.segments : 0.0, + GiBps(memcpy_.bytes, memcpy_.seconds)); + } + + std::mutex mutex_; + Totals gather_; + Totals memcpy_; + std::chrono::steady_clock::time_point last_report_ = std::chrono::steady_clock::now(); +}; + +DebugStats& Stats() { + // Never destroyed: worker threads can still be copying during static + // teardown, same reasoning as the gather streams. + static DebugStats* stats = new DebugStats(); + return *stats; +} + +// Everything a unit test needs to rebuild this exact copy: the direction, the +// device, the two base pointers (for alignment), and the segment geometry. +void LogPlanGeometry(const char* path, size_t plan_index, const TransferPlan& plan, int device_id) { + const size_t limit = DebugMaxSegments(); + if (limit == 0) return; + const auto src_addr = reinterpret_cast(plan.src.host_ptr); + const auto dst_addr = reinterpret_cast(plan.dst.host_ptr); + MORI_UMBP_INFO( + "[HbmCopyEngine][dbg] {} plan={} kind={} dev={} src_base={:#x} dst_base={:#x} " + "src_loc={} dst_loc={} src_dev={} dst_dev={} src_size={} dst_size={} segs={} tags={}", + path, plan_index, KindName(KindFor(plan.src, plan.dst)), device_id, src_addr, dst_addr, + static_cast(plan.src.loc), static_cast(plan.dst.loc), plan.src.device, + plan.dst.device, plan.src.size, plan.dst.size, plan.sizes.size(), plan.tags.size()); + const size_t shown = std::min(limit, plan.sizes.size()); + for (size_t i = 0; i < shown; ++i) { + // vec16 reports whether the kernel's uint4 fast path applies to this + // segment; a single unaligned segment drops it to the byte loop. + const uintptr_t mask = (src_addr + plan.src_offsets[i]) | (dst_addr + plan.dst_offsets[i]) | + static_cast(plan.sizes[i]); + MORI_UMBP_INFO("[HbmCopyEngine][dbg] seg[{}] src_off={} dst_off={} bytes={} vec16={}", i, + plan.src_offsets[i], plan.dst_offsets[i], plan.sizes[i], + (mask & (kVectorBytes - 1)) == 0 ? 1 : 0); + } + if (plan.sizes.size() > shown) { + MORI_UMBP_INFO("[HbmCopyEngine][dbg] ... {} more segs", plan.sizes.size() - shown); + } +} + +// One non-blocking stream per device, per thread. Submit runs on PoolClient's +// executor threads, so a shared stream would serialize copies that the +// hipMemcpy path ran independently. +std::map& GatherStreams() { + static thread_local std::map* streams = new std::map(); + return *streams; +} + +hipStream_t GatherStream(int device_id) { + auto& streams = GatherStreams(); + auto it = streams.find(device_id); + if (it != streams.end()) return it->second; + hipStream_t stream = nullptr; + if (hipStreamCreateWithFlags(&stream, hipStreamNonBlocking) != hipSuccess) { + (void)hipGetLastError(); + return nullptr; + } + streams.emplace(device_id, stream); + return stream; +} + +// Latches on the first gather failure and never tries again. A kernel that +// cannot launch here will not start working later, and the fallback is always +// available, so one failed launch should not cost every later batch a retry. +std::atomic& GatherDisabled() { + static std::atomic disabled{false}; + return disabled; +} + +// Which device a plan's copies should run on. A D2D pair between two devices +// needs a current device with peer access to both; we pick the destination's, +// which is the one hipMemcpy attributes the copy to. +inline int DeviceFor(const TransferPlan& plan) { + if (IsGpu(plan.dst) && plan.dst.device >= 0) return plan.dst.device; + if (IsGpu(plan.src) && plan.src.device >= 0) return plan.src.device; + return -1; +} + +// Every plan this engine produces is executed before Submit returns, so the +// handle only has to replay the outcome. (Same shape as LocalCopyEngine's; +// kept file-local in both because it is an implementation detail of "this +// engine completes inline", not a shared contract.) +class SettledHandle final : public TransferHandle { + public: + explicit SettledHandle(std::vector failures) : failures_(std::move(failures)) {} + + void Wait(std::vector* failures) override { + if (reported_) return; + reported_ = true; + if (failures != nullptr) { + for (auto& f : failures_) failures->push_back(std::move(f)); + } + failures_.clear(); + } + + private: + std::vector failures_; + bool reported_ = false; +}; + +} // namespace + +// --------------------------------------------------------------------------- +// TransferEngine +// --------------------------------------------------------------------------- + +bool HbmCopyEngine::CanHandle(const TransferRef& src, const TransferRef& dst) const { + if (!src.HasHostPtr() || !dst.HasHostPtr()) return false; + // At least one GPU endpoint. The both-CPU case is LocalCopyEngine's and is + // left to it: hipMemcpy would be correct but strictly slower than the + // NT-AVX2 path, and overlapping the two engines would make composite + // selection order load-bearing for performance. + return IsGpu(src) || IsGpu(dst); +} + +TransferPlanSet HbmCopyEngine::Plan(const std::vector& items) const { + TransferPlanSet out; + std::unordered_map pair_to_plan; + pair_to_plan.reserve(items.size() * 2); + + for (const auto& item : items) { + if (item.size == 0) continue; + if (!CanHandle(item.src, item.dst)) { + out.rejected_tags.push_back(item.tag); + continue; + } + // Bounds are checked here, once, rather than per copy: a page index that + // does not fit its buffer is a bug in the backend's slot bookkeeping, and + // the key must be failed rather than copied past the end. On a GPU + // endpoint that matters more than on a host one — an overrun is not a + // segfault but silent corruption of whatever else lives in the pool. + 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; + } + + const PtrPairKey key{item.src.host_ptr, item.dst.host_ptr}; + auto it = pair_to_plan.find(key); + size_t pi; + if (it == pair_to_plan.end()) { + pi = out.plans.size(); + pair_to_plan.emplace(key, pi); + TransferPlan plan; + plan.src = item.src; + plan.dst = item.dst; + plan.dir = TransferDirection::kLocal; + out.plans.push_back(std::move(plan)); + } else { + pi = it->second; + } + TransferPlan& plan = out.plans[pi]; + const size_t so = static_cast(item.src_offset); + const size_t dof = static_cast(item.dst_offset); + const size_t sz = static_cast(item.size); + const bool can_coalesce = !plan.sizes.empty() && + AdjacentNoOverflow(plan.src_offsets.back(), plan.sizes.back(), so) && + AdjacentNoOverflow(plan.dst_offsets.back(), plan.sizes.back(), dof); + if (can_coalesce) { + plan.sizes.back() += sz; + } else { + plan.src_offsets.push_back(so); + plan.dst_offsets.push_back(dof); + plan.sizes.push_back(sz); + } + if (plan.tags.empty() || plan.tags.back() != item.tag) plan.tags.push_back(item.tag); + } + return out; +} + +void HbmCopyEngine::AddHostGatherRegion(void* base, size_t bytes) { + if (base == nullptr || bytes == 0) return; + std::lock_guard lock(host_regions_mutex_); + host_regions_.push_back(std::make_unique(base, bytes)); +} + +void HbmCopyEngine::ClearHostGatherRegions() { + std::lock_guard lock(host_regions_mutex_); + host_regions_.clear(); +} + +void* HbmCopyEngine::HostRegionDeviceAddress(const void* ptr, size_t size, int device_id) const { + std::lock_guard lock(host_regions_mutex_); + for (const auto& region : host_regions_) { + if (region == nullptr) continue; + if (void* alias = region->DeviceAddress(ptr, size, device_id)) return alias; + } + return nullptr; +} + +// Run the gather kernel over every eligible segment in the batch, and report +// which plans it completed. +// +// Bucketing is per DEVICE, not per plan, and that is the whole point. A plan +// is one (src base, dst base) pair, so a caller reading three ranges into three +// separate GPU allocations produces three single-segment plans — nothing to +// gather within any of them, while together they are exactly the scattered +// small-fragment batch the kernel wins on. Direction does not partition the +// buckets either: once the host side is registered both endpoints are +// dereferenceable from the device, so one kernel can carry H2D and D2H +// fragments at once. +// +// Never reports a plan it did not fully complete. Anything declined or failed +// is left for the hipMemcpy loop, which reaches the same result whether or not +// some fragments already landed. +const char* HbmCopyEngine::GatherSkipName(GatherSkip reason) { + switch (reason) { + case GatherSkip::kTaken: + return "taken"; + case GatherSkip::kUnrecorded: + return "unrecorded"; + case GatherSkip::kDisabled: + return "gather-disabled"; + case GatherSkip::kKindNotHostDevice: + return "kind-not-h2d-or-d2h"; + case GatherSkip::kNoDevice: + return "no-device"; + case GatherSkip::kHostNotGatherable: + return "host-not-gatherable"; + case GatherSkip::kNoFragments: + return "no-fragments"; + case GatherSkip::kTooFewFragments: + return "too-few-fragments"; + case GatherSkip::kFragmentAtOrAboveThreshold: + return "mean-seg-at-or-above-threshold"; + case GatherSkip::kSetDeviceFailed: + return "hipSetDevice-failed"; + case GatherSkip::kNoStream: + return "no-stream"; + case GatherSkip::kLaunchFailed: + return "launch-failed"; + case GatherSkip::kSyncFailed: + return "sync-failed"; + } + return "?"; +} + +std::vector HbmCopyEngine::GatherEligiblePlans(const std::vector& plans, + std::vector* skip_reasons) { + std::vector done(plans.size(), 0); + if (skip_reasons != nullptr) skip_reasons->assign(plans.size(), GatherSkip::kUnrecorded); + // Submit passes nullptr unless debug mode is on, so the hot path pays one + // predicted null check per verdict and stores nothing. + auto note = [skip_reasons](size_t plan_index, GatherSkip reason) { + if (skip_reasons != nullptr) (*skip_reasons)[plan_index] = reason; + }; + if (GatherDisabled().load(std::memory_order_relaxed) || !DeviceGatherEnabled()) { + for (size_t p = 0; p < plans.size(); ++p) note(p, GatherSkip::kDisabled); + return done; + } + + struct Bucket { + std::vector fragments; + std::vector plan_indices; + size_t total_bytes = 0; + // Debug-only shape. A bucket may legitimately mix directions: once the + // host side is registered both endpoints are device-dereferenceable, so + // one kernel can carry H2D and D2H fragments at once. + size_t min_fragment = std::numeric_limits::max(); + size_t max_fragment = 0; + size_t h2d_fragments = 0; + size_t d2h_fragments = 0; + }; + std::map buckets; + + for (size_t p = 0; p < plans.size(); ++p) { + const auto& plan = plans[p]; + const hipMemcpyKind kind = KindFor(plan.src, plan.dst); + // D2D has no host side to register, and H2H never reaches this engine. + if (kind != hipMemcpyHostToDevice && kind != hipMemcpyDeviceToHost) { + note(p, GatherSkip::kKindNotHostDevice); + continue; + } + const int device_id = DeviceFor(plan); + if (device_id < 0) { + note(p, GatherSkip::kNoDevice); + continue; + } + + const char* src = static_cast(plan.src.host_ptr); + char* dst = static_cast(plan.dst.host_ptr); + const bool host_is_src = kind == hipMemcpyHostToDevice; + + std::vector fragments; + fragments.reserve(plan.sizes.size()); + size_t plan_bytes = 0; + bool eligible = true; + for (size_t i = 0; i < plan.sizes.size() && eligible; ++i) { + const void* fragment_src = src + plan.src_offsets[i]; + void* fragment_dst = dst + plan.dst_offsets[i]; + const void* host_side = host_is_src ? fragment_src : fragment_dst; + // A kernel cannot dereference plain mmap memory — it faults the GPU — so + // an uncovered host side disqualifies the whole plan. Nor can it use the + // host address: only the host side is replaced by its per-device alias. + void* const host_alias = HostRegionDeviceAddress(host_side, plan.sizes[i], device_id); + if (host_alias == nullptr) { + eligible = false; + break; + } + if (host_is_src) { + fragments.push_back({host_alias, fragment_dst, plan.sizes[i]}); + } else { + fragments.push_back({fragment_src, host_alias, plan.sizes[i]}); + } + plan_bytes += plan.sizes[i]; + } + if (!eligible) { + note(p, GatherSkip::kHostNotGatherable); + continue; + } + if (fragments.empty()) { + note(p, GatherSkip::kNoFragments); + continue; + } + + auto& bucket = buckets[device_id]; + for (const auto& fragment : fragments) { + bucket.min_fragment = std::min(bucket.min_fragment, fragment.bytes); + bucket.max_fragment = std::max(bucket.max_fragment, fragment.bytes); + } + (host_is_src ? bucket.h2d_fragments : bucket.d2h_fragments) += fragments.size(); + bucket.fragments.insert(bucket.fragments.end(), fragments.begin(), fragments.end()); + bucket.plan_indices.push_back(p); + bucket.total_bytes += plan_bytes; + } + + const bool debug = DebugEnabled(); + for (auto& [device_id, bucket] : buckets) { + // Bound to a plain reference first: capturing the structured binding + // `bucket` directly is a C++20 extension and warns under this build. + auto& bucket_plans = bucket.plan_indices; + auto note_bucket = [&bucket_plans, ¬e](GatherSkip reason) { + for (size_t p : bucket_plans) note(p, reason); + }; + // One fragment is a plain hipMemcpy's best case; above the threshold the + // copy engine wins because a single large async copy amortizes its own + // submission. + if (bucket.fragments.size() < 2) { + note_bucket(GatherSkip::kTooFewFragments); + continue; + } + if (bucket.total_bytes / bucket.fragments.size() >= kGatherFragmentThreshold) { + note_bucket(GatherSkip::kFragmentAtOrAboveThreshold); + continue; + } + if (hipSetDevice(device_id) != hipSuccess) { + (void)hipGetLastError(); + note_bucket(GatherSkip::kSetDeviceFailed); + continue; + } + hipStream_t stream = GatherStream(device_id); + if (stream == nullptr) { + note_bucket(GatherSkip::kNoStream); + continue; + } + // Timed launch-to-completion. The synchronize below is NOT added by debug + // mode -- it is what makes the gather path synchronous in the first place + // -- so the interval is the kernel's real cost, and enabling debug does not + // perturb what is being measured. + const auto started = std::chrono::steady_clock::now(); + if (!LaunchDeviceGather(bucket.fragments.data(), bucket.fragments.size(), device_id, stream)) { + note_bucket(GatherSkip::kLaunchFailed); + continue; + } + const hipError_t sync = hipStreamSynchronize(stream); + if (sync != hipSuccess) { + // Latch off: a kernel that cannot complete here will not start working + // later, and every caller still has the fallback. + (void)hipGetLastError(); + GatherDisabled().store(true, std::memory_order_relaxed); + GatherStreams().erase(device_id); + MORI_UMBP_ERROR("[HbmCopyEngine] gather sync dev={} failed: {}", device_id, + hipGetErrorString(sync)); + note_bucket(GatherSkip::kSyncFailed); + continue; + } + for (size_t p : bucket.plan_indices) done[p] = 1; + if (!debug) continue; + + const double seconds = Seconds(started, std::chrono::steady_clock::now()); + const double bytes = static_cast(bucket.total_bytes); + Stats().Record(/*gather=*/true, bucket.fragments.size(), bytes, seconds); + note_bucket(GatherSkip::kTaken); + if (!DebugSampleThisCall()) continue; + MORI_UMBP_INFO( + "[HbmCopyEngine][dbg] path=gather dev={} plans={} segs={} h2d={} d2h={} bytes={} " + "seg_min={} seg_mean={} seg_max={} seconds={:.6f} GiBps={:.3f}", + device_id, bucket.plan_indices.size(), bucket.fragments.size(), bucket.h2d_fragments, + bucket.d2h_fragments, bucket.total_bytes, bucket.min_fragment, + bucket.total_bytes / bucket.fragments.size(), bucket.max_fragment, seconds, + GiBps(bytes, seconds)); + for (size_t p : bucket.plan_indices) LogPlanGeometry("gather", p, plans[p], device_id); + } + return done; +} + +std::unique_ptr HbmCopyEngine::Submit(std::vector plans) { + if (plans.empty()) return nullptr; + + std::vector failures; + + // Restore the caller's current device on the way out. Submit runs on + // PoolClient's executor threads, which are not ours to leave re-pointed at + // another device. + int entry_device = -1; + bool device_touched = false; + int current_device = -1; + + // The gather pass sets the device itself, so the save has to happen before + // it rather than at the first plan that needs one. + if (hipGetDevice(&entry_device) != hipSuccess) { + entry_device = -1; + } else { + device_touched = true; + } + const bool debug = DebugEnabled(); + std::vector skips; + const std::vector gathered = GatherEligiblePlans(plans, debug ? &skips : nullptr); + // The pass may have left the current device anywhere. + current_device = -1; + + for (size_t plan_index = 0; plan_index < plans.size(); ++plan_index) { + const auto& plan = plans[plan_index]; + if (gathered[plan_index] != 0) continue; + const int want_device = DeviceFor(plan); + if (want_device >= 0 && want_device != current_device) { + if (!device_touched) { + if (hipGetDevice(&entry_device) != hipSuccess) entry_device = -1; + device_touched = true; + } + const hipError_t derr = hipSetDevice(want_device); + if (derr != hipSuccess) { + failures.push_back(TransferFailure{plan.tags, static_cast(derr), + std::string("hipSetDevice(") + + std::to_string(want_device) + + "): " + hipGetErrorString(derr), + std::string("hbm:dev") + std::to_string(want_device)}); + continue; + } + current_device = want_device; + } + + const hipMemcpyKind kind = KindFor(plan.src, plan.dst); + char* dst = static_cast(plan.dst.host_ptr); + const char* src = static_cast(plan.src.host_ptr); + + // One failure 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. + // + // hipMemcpy is BLOCKING, so this loop is already launch-to-completion and + // the debug clock below needs no synchronize of its own. + size_t plan_bytes = 0; + const auto started = std::chrono::steady_clock::now(); + hipError_t err = hipSuccess; + for (size_t i = 0; i < plan.sizes.size(); ++i) { + err = hipMemcpy(dst + plan.dst_offsets[i], src + plan.src_offsets[i], plan.sizes[i], kind); + if (err != hipSuccess) break; + plan_bytes += plan.sizes[i]; + } + if (debug && err == hipSuccess) { + const double seconds = Seconds(started, std::chrono::steady_clock::now()); + Stats().Record(/*gather=*/false, plan.sizes.size(), static_cast(plan_bytes), seconds); + if (DebugSampleThisCall()) { + const size_t mean = plan.sizes.empty() ? 0 : plan_bytes / plan.sizes.size(); + MORI_UMBP_INFO( + "[HbmCopyEngine][dbg] path=hipMemcpy dev={} kind={} segs={} bytes={} seg_mean={} " + "seconds={:.6f} GiBps={:.3f} why_not_gather={}", + want_device, KindName(kind), plan.sizes.size(), plan_bytes, mean, seconds, + GiBps(static_cast(plan_bytes), seconds), + GatherSkipName(plan_index < skips.size() ? skips[plan_index] + : GatherSkip::kUnrecorded)); + LogPlanGeometry("hipMemcpy", plan_index, plan, want_device); + } + } + if (err != hipSuccess) { + MORI_UMBP_ERROR("[HbmCopyEngine] hipMemcpy kind={} dev={} failed: {}", static_cast(kind), + want_device, hipGetErrorString(err)); + failures.push_back(TransferFailure{plan.tags, static_cast(err), + std::string("hipMemcpy: ") + hipGetErrorString(err), + std::string("hbm:dev") + std::to_string(want_device)}); + } + } + + if (device_touched && entry_device >= 0) { + // Best effort: a failure here cannot be attributed to any one plan, and the + // transfers themselves already settled. + (void)hipSetDevice(entry_device); + } + + return std::make_unique(std::move(failures)); +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/transfer/local_copy_engine.cpp b/src/umbp/distributed/transfer/local_copy_engine.cpp new file mode 100644 index 000000000..a8defac73 --- /dev/null +++ b/src/umbp/distributed/transfer/local_copy_engine.cpp @@ -0,0 +1,480 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/transfer/local_copy_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__x86_64__) || defined(__i386__) +#include +#endif + +namespace mori::umbp { + +// --------------------------------------------------------------------------- +// Block copy +// +// In distributed mode 1 key == 1 page (master page_size == KV block size), so +// a self-target copy moves one ~MiB-scale block per call. Local mode's +// DRAMTier optimization parallelizes WITHIN a call's pages (always 1 here) and +// so never applied; the cross-key parallelism lives in PoolClient's batch +// executors instead, and what is left to win per key is the cache bypass. +// --------------------------------------------------------------------------- + +namespace { + +#if defined(__x86_64__) || defined(__i386__) +__attribute__((target("avx2"))) inline void NtCopyAvx2(char* d, const char* s, size_t n) { + size_t head = (32 - (reinterpret_cast(d) & 31)) & 31; + if (head > n) head = n; + std::memcpy(d, s, head); + size_t i = head; + for (; i + 128 <= n; i += 128) { + __m256i a = _mm256_loadu_si256(reinterpret_cast(s + i)); + __m256i b = _mm256_loadu_si256(reinterpret_cast(s + i + 32)); + __m256i c = _mm256_loadu_si256(reinterpret_cast(s + i + 64)); + __m256i e = _mm256_loadu_si256(reinterpret_cast(s + i + 96)); + _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i), a); + _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 32), b); + _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 64), c); + _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 96), e); + } + for (; i + 32 <= n; i += 32) { + __m256i a = _mm256_loadu_si256(reinterpret_cast(s + i)); + _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i), a); + } + if (i < n) std::memcpy(d + i, s + i, n - i); + _mm_sfence(); +} +inline bool NtSupported() { return __builtin_cpu_supports("avx2"); } +#else +inline void NtCopyAvx2(char* d, const char* s, size_t n) { std::memcpy(d, s, n); } +inline bool NtSupported() { return false; } +#endif + +// True iff [next, ...) is exactly adjacent after [base, base+len). +inline bool AdjacentNoOverflow(size_t base, size_t len, size_t next) { + return len <= std::numeric_limits::max() - base && base + len == next; +} + +struct PtrPairKey { + const void* src; + const void* dst; + bool operator==(const PtrPairKey& o) const noexcept { return src == o.src && dst == o.dst; } +}; +struct PtrPairKeyHash { + size_t operator()(const PtrPairKey& k) const noexcept { + size_t h = std::hash{}(k.src); + h ^= std::hash{}(k.dst) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); + return h; + } +}; + +// Every plan this engine produces is executed before Submit returns, so the +// handle only has to replay the outcome. +class SettledHandle final : public TransferHandle { + public: + explicit SettledHandle(std::vector failures) : failures_(std::move(failures)) {} + + void Wait(std::vector* failures) override { + if (reported_) return; + reported_ = true; + if (failures != nullptr) { + for (auto& f : failures_) failures->push_back(std::move(f)); + } + failures_.clear(); + } + + private: + std::vector failures_; + bool reported_ = false; +}; + +// Tuning knobs for the copy itself. Kept internal: they describe how this +// engine schedules work, not anything a caller can act on. + +// Smallest block the non-temporal path is used for. +// +// This was 256 KiB, a figure that came from whole-object copies where one key +// is one ~MiB page. A RANGED copy moves one LAYER of one page, and layers are +// small: DeepSeek-V4-Pro's three pools are 37,440 B, 8,448 B and 1,728 B. So +// every ranged copy fell under the old threshold and took the ordinary memcpy, +// paying read-for-ownership on a destination it completely overwrites. +// +// The threshold has to be set against a real model's fragment sizes, not a +// synthetic one: a uniform 73,728 B shape is larger than anything the model +// actually has, so a threshold that looks right there never fires in practice. +// +// 2 KiB, measured on DeepSeek-V4-Pro's geometry (8 ranks, 32 pages, served +// entirely from the local medium, medians of 3): restore 12.04 -> 11.43 ms and +// TTFL 1,529 -> 1,353 us. Below 2 KiB the remaining movement is inside this +// machine's run-to-run spread. Tunable because the right answer is a property +// of the machine's cache, not of the code. +size_t NtMinBytes() { + static const size_t kDefault = size_t{2} << 10; + static const size_t n = [] { + const char* raw = std::getenv("UMBP_DRAM_NT_COPY_MIN_BYTES"); + if (raw == nullptr) return kDefault; + char* end = nullptr; + const unsigned long long v = std::strtoull(raw, &end, 10); + if (end == raw || v == 0) return kDefault; + return static_cast(v); + }(); + return n; +} + +// Backing store for MarkThreadBackgroundCopies(); see the header for why a +// thread would declare itself background. +bool& ThreadBackgroundFlag() { + static thread_local bool background = false; + return background; +} + +// A persistent set of helper threads that Submit can lend its copies to. +// +// Persistent, not spawned per call, because spawning was measured to cost more +// than the parallelism gained on any workload with a remote half. Creating and +// joining threads maps and unmaps their stacks, which serializes on the +// process's address space -- so the damage is not proportional to how many +// helpers are used, and indeed capping the whole process to FOUR concurrent +// spawned helpers still cost 25% of TTFL on a remote layer-wise restore. The +// pool makes fan-out cost a condition-variable wake instead. +class CopyPool { + public: + struct Batch { + const std::function* fn = nullptr; + size_t count = 0; + std::atomic next{0}; + std::atomic done{0}; + // Completion is waited on, not spun on. A spin is only free when the box + // has a spare core, and the process this runs in does not: it is also + // running the model's forward pass, the HIP runtime and every other rank. + std::mutex mutex; + std::condition_variable cv; + }; + + static CopyPool& Instance() { + // Never destroyed: the helpers outlive any translation unit's static + // destructor order, and tearing them down during exit buys nothing. + static CopyPool* pool = new CopyPool(); + return *pool; + } + + // Publishes `batch`, works on it from the calling thread, and returns once + // every item has run. Helpers are best-effort: if they are all busy the + // caller simply does the whole batch itself. + void Run(const std::shared_ptr& batch, int helpers) { + { + std::lock_guard lock(mutex_); + GrowLocked(helpers); + queue_.push_back(batch); + } + // One wake per helper this batch wants. notify_all would wake every idle + // thread in the process to serve one batch. + for (int i = 0; i < helpers; ++i) cv_.notify_one(); + + Drain(*batch); + + // Draining only means THIS thread found nothing left to claim; a helper may + // still be inside the last item. The batch is shared_ptr-owned so waiting + // for it is safe. + { + std::unique_lock lock(batch->mutex); + batch->cv.wait( + lock, [&batch] { return batch->done.load(std::memory_order_acquire) >= batch->count; }); + } + { + std::lock_guard lock(mutex_); + auto it = std::find(queue_.begin(), queue_.end(), batch); + if (it != queue_.end()) queue_.erase(it); + } + } + + private: + // Upper bound on helper threads for the whole process. They are shared by + // every caller, so this is not per call. + static int PoolCap() { + static const int cap = [] { + const unsigned hc = std::thread::hardware_concurrency(); + const int n = hc == 0 ? 4 : static_cast(hc); + return std::max(1, std::min(n, 32)); + }(); + return cap; + } + + static void Drain(Batch& batch) { + for (size_t i = batch.next.fetch_add(1, std::memory_order_relaxed); i < batch.count; + i = batch.next.fetch_add(1, std::memory_order_relaxed)) { + (*batch.fn)(i); + if (batch.done.fetch_add(1, std::memory_order_acq_rel) + 1 == batch.count) { + // Notify under the batch mutex the waiter checks its predicate under, + // so a waiter that has not yet slept cannot miss this. + std::lock_guard lock(batch.mutex); + batch.cv.notify_all(); + } + } + } + + // Adds threads only for the helpers this batch cannot get from the idle ones. + // + // The previous version sized the pool to `helpers` outright, which with the + // default of one helper per call meant ONE thread for the whole process: the + // second concurrent caller onwards got no help at all and still paid the + // queue, the wake and the wait. + void GrowLocked(int helpers) { + const int deficit = helpers - idle_; + for (int i = 0; i < deficit && static_cast(threads_.size()) < PoolCap(); ++i) { + threads_.emplace_back([this] { WorkerLoop(); }); + } + } + + // The oldest batch that still has items to claim. Exhausted batches are left + // for their publisher to remove, so they must be skipped rather than blocking + // the queue. + std::shared_ptr PickLocked() { + for (auto& batch : queue_) { + if (batch->next.load(std::memory_order_relaxed) < batch->count) return batch; + } + return nullptr; + } + + void WorkerLoop() { + // A helper exists to serve other threads' copies, so its own copies -- of + // which it has none -- must never recurse into fanning out. + MarkThreadBackgroundCopies(); + for (;;) { + std::shared_ptr batch; + { + std::unique_lock lock(mutex_); + ++idle_; + cv_.wait(lock, [this, &batch] { return (batch = PickLocked()) != nullptr; }); + --idle_; + } + // Left queued while draining: several helpers working the same batch is + // the point, and they share its cursor. + Drain(*batch); + } + } + + std::mutex mutex_; + std::condition_variable cv_; + std::deque> queue_; + std::vector threads_; + int idle_ = 0; +}; + +// How many threads to spread one Submit's copies over, counting the caller. +// +// A single memcpy stream tops out well below what the memory system can do -- +// measured at ~6 GB/s here, and the same ~6 GB/s whether 1 rank or 8 are +// copying concurrently, so the ceiling is the stream and not the DRAM. Local +// mode already knew this: DRAMTier fans its batch reads out over +// UMBP_DRAM_READ_THREADS (default 4), and its comment says ">1 breaks the +// single-core memcpy ceiling". The distributed path never did, because +// LocalCopyEngine was written when one key was one whole page and the batch +// executors above it supplied the parallelism. Ranged I/O broke that +// assumption: a layer-wise reader hands one thread thousands of small copies. +// +// 4, matching local mode's UMBP_DRAM_READ_THREADS. On DeepSeek-V4-Pro's real +// pool geometry it is best or tied-best on every shape measured (medians of 3, +// restore ms / TTFL us): +// +// threads=1 threads=2 threads=4 threads=8 +// local 1 rank 9.88 / 690 9.71 / 518 9.61 / 425 9.57 / 381 +// local 8 ranks 10.06 / 853 10.03 / 723 9.83 / 611 9.94 / 682 +// remote 1 rank 13.23 / 2294 12.06 / 2229 11.50 / 2304 11.51 / 2284 +// remote 8 ranks 84.44 / 12399 90.55 / 12860 82.58 / 11384 79.78 / 11600 +// +// An earlier revision defaulted to 2, because 4 appeared to cost a concurrent +// remote half its first layer group. That was not bandwidth: it was this +// file's own pool handing out one helper for the whole process and making +// everyone else spin waiting for it. With that fixed the effect is gone, and +// the value that looked reckless is simply the right one. 8 buys a little more +// at one rank and gives it back at eight, so 4 is where this stops. +// +// Below the byte floor the threads cost more than they save, so a small batch +// stays on the calling thread and pays nothing. +int CopyThreads(size_t total_bytes, size_t copy_count) { + if (ThreadBackgroundFlag()) return 1; + static const int kMax = [] { + int n = 4; + if (const char* raw = std::getenv("UMBP_DRAM_COPY_THREADS")) { + const int v = std::atoi(raw); + if (v >= 1) n = v; + } + const unsigned hc = std::thread::hardware_concurrency(); + if (hc > 0 && n > static_cast(hc)) n = static_cast(hc); + return n < 1 ? 1 : n; + }(); + static const size_t kMinBytes = [] { + if (const char* raw = std::getenv("UMBP_DRAM_COPY_THREADS_MIN_BYTES")) { + char* end = nullptr; + const unsigned long long v = std::strtoull(raw, &end, 10); + if (end != raw) return static_cast(v); + } + return size_t{1} << 20; + }(); + if (kMax <= 1 || copy_count < 2 || total_bytes < kMinBytes) return 1; + return std::min(kMax, static_cast(copy_count)); +} + +} // namespace + +// Set on a worker thread at start-up; read by CopyThreads above. +void MarkThreadBackgroundCopies() { ThreadBackgroundFlag() = true; } + +bool ThreadDoesBackgroundCopies() { return ThreadBackgroundFlag(); } + +void HostCopyBlock(void* dst, const void* src, size_t size) { + static const bool kNt = NtSupported() && !(std::getenv("UMBP_DRAM_NT_COPY") && + std::getenv("UMBP_DRAM_NT_COPY")[0] == '0'); + static const size_t kNtMinBytes = NtMinBytes(); + if (kNt && size >= kNtMinBytes) { + NtCopyAvx2(static_cast(dst), static_cast(src), size); + } else { + std::memcpy(dst, src, size); + } +} + +// --------------------------------------------------------------------------- +// TransferEngine +// --------------------------------------------------------------------------- + +bool LocalCopyEngine::CanHandle(const TransferRef& src, const TransferRef& dst) const { + return src.HasHostPtr() && dst.HasHostPtr() && src.loc == mori::io::MemoryLocationType::CPU && + dst.loc == mori::io::MemoryLocationType::CPU; +} + +TransferPlanSet LocalCopyEngine::Plan(const std::vector& items) const { + TransferPlanSet out; + std::unordered_map pair_to_plan; + pair_to_plan.reserve(items.size() * 2); + + for (const auto& item : items) { + if (item.size == 0) continue; + if (!CanHandle(item.src, item.dst)) { + out.rejected_tags.push_back(item.tag); + continue; + } + // Bounds are checked here, once, rather than per copy: a page index that + // does not fit its buffer is a bug in the backend's slot bookkeeping, and + // the key must be failed rather than copied past the end. + 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; + } + + const PtrPairKey key{item.src.host_ptr, item.dst.host_ptr}; + auto it = pair_to_plan.find(key); + size_t pi; + if (it == pair_to_plan.end()) { + pi = out.plans.size(); + pair_to_plan.emplace(key, pi); + TransferPlan plan; + plan.src = item.src; + plan.dst = item.dst; + plan.dir = TransferDirection::kLocal; + out.plans.push_back(std::move(plan)); + } else { + pi = it->second; + } + TransferPlan& plan = out.plans[pi]; + const size_t so = static_cast(item.src_offset); + const size_t dof = static_cast(item.dst_offset); + const size_t sz = static_cast(item.size); + const bool can_coalesce = !plan.sizes.empty() && + AdjacentNoOverflow(plan.src_offsets.back(), plan.sizes.back(), so) && + AdjacentNoOverflow(plan.dst_offsets.back(), plan.sizes.back(), dof); + if (can_coalesce) { + plan.sizes.back() += sz; + } else { + plan.src_offsets.push_back(so); + plan.dst_offsets.push_back(dof); + plan.sizes.push_back(sz); + } + if (plan.tags.empty() || plan.tags.back() != item.tag) plan.tags.push_back(item.tag); + } + return out; +} + +std::unique_ptr LocalCopyEngine::Submit(std::vector plans) { + if (plans.empty()) return nullptr; + + // Flatten first. A ranged call arrives as many plans of one copy each, a + // whole-object one as a single plan of many; spreading work over threads has + // to see past that shape difference or it parallelizes only one of them. + struct Copy { + char* dst; + const char* src; + size_t size; + }; + std::vector copies; + size_t total_copies = 0; + for (const auto& plan : plans) total_copies += plan.sizes.size(); + copies.reserve(total_copies); + size_t total_bytes = 0; + for (const auto& plan : plans) { + char* dst = static_cast(plan.dst.host_ptr); + const char* src = static_cast(plan.src.host_ptr); + for (size_t i = 0; i < plan.sizes.size(); ++i) { + copies.push_back({dst + plan.dst_offsets[i], src + plan.src_offsets[i], plan.sizes[i]}); + total_bytes += plan.sizes[i]; + } + } + + const int threads = CopyThreads(total_bytes, copies.size()); + if (threads <= 1) { + for (const Copy& c : copies) HostCopyBlock(c.dst, c.src, c.size); + } else { + // Items are claimed off one shared cursor rather than split up front: the + // copies in a batch are not the same size (a range that straddles a page + // boundary is emitted as fragments), so equal counts would not be equal + // work. + const std::function fn = [&copies](size_t i) { + HostCopyBlock(copies[i].dst, copies[i].src, copies[i].size); + }; + auto batch = std::make_shared(); + batch->fn = &fn; + batch->count = copies.size(); + CopyPool::Instance().Run(batch, threads - 1); + } + + // No failure mode: bounds were validated at plan time and a memcpy cannot + // fail afterwards. + return std::make_unique(std::vector{}); +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/transfer/mori_io_engine.cpp b/src/umbp/distributed/transfer/mori_io_engine.cpp new file mode 100644 index 000000000..ff11b9609 --- /dev/null +++ b/src/umbp/distributed/transfer/mori_io_engine.cpp @@ -0,0 +1,664 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/transfer/mori_io_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/io/backend.hpp" +#include "mori/utils/mori_log.hpp" + +namespace mori::umbp { + +namespace { + +// Bucket key for grouping page transfers by their endpoint pair. +// +// MemoryUniqueId is allocated per IOEngine, so two DIFFERENT peers can both +// publish buffer id 1 — the id alone does NOT identify an endpoint. The map +// therefore buckets on ids (cheap, no string in the key) and each bucket holds +// candidate plan indices that the caller disambiguates by comparing engine +// keys. Collisions are rare, so the string compare is off the common path. +struct PairKey { + mori::io::MemoryUniqueId src; + mori::io::MemoryUniqueId dst; + bool operator==(const PairKey& o) const noexcept { return src == o.src && dst == o.dst; } +}; +struct PairKeyHash { + size_t operator()(const PairKey& k) const noexcept { + // hash_combine (boost-style): independent of size_t width. + size_t h = std::hash{}(k.src); + h ^= std::hash{}(k.dst) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); + return h; + } +}; + +// Do `plan` and `item` name the same two endpoints? Compares engine keys, not +// just ids, for the reason above. +bool SameEndpoints(const TransferPlan& plan, const TransferRef& src, const TransferRef& dst) { + return plan.src.mem.id == src.mem.id && plan.dst.mem.id == dst.mem.id && + plan.src.mem.engineKey == src.mem.engineKey && plan.dst.mem.engineKey == dst.mem.engineKey; +} + +// True iff [next, ...) is exactly adjacent after [base, base+len), with no +// size_t overflow in base+len. Used to coalesce contiguous SG segments. +inline bool AdjacentNoOverflow(size_t base, size_t len, size_t next) { + return len <= std::numeric_limits::max() - base && base + len == next; +} + +// Append one segment to `plan`, coalescing with the previous when BOTH sides +// are exactly contiguous. +// +// This is NOT redundant with the RDMA backend's own WR merging: the backend +// would fold these into the same WR either way, but doing it here shrinks the +// inner SG vector it has to sort/merge (O(M log M)) and allocate, which matters +// when M is large (big batch x pages). Same bytes, so per-plan failure +// granularity is unchanged. The merged size must stay within uint32_t since +// the backend stores it in ibv_sge.length. +void AppendSegment(TransferPlan* plan, size_t src_off, size_t dst_off, size_t size, size_t tag) { + const bool can_coalesce = + !plan->sizes.empty() && + AdjacentNoOverflow(plan->src_offsets.back(), plan->sizes.back(), src_off) && + AdjacentNoOverflow(plan->dst_offsets.back(), plan->sizes.back(), dst_off) && + plan->sizes.back() <= static_cast(std::numeric_limits::max()) - size; + if (can_coalesce) { + plan->sizes.back() += size; + } else { + plan->src_offsets.push_back(src_off); + plan->dst_offsets.push_back(dst_off); + plan->sizes.push_back(size); + } + // De-dup contributing tags: a single item's pages arrive consecutively + // within a pair, so a back() check suffices. + if (plan->tags.empty() || plan->tags.back() != tag) plan->tags.push_back(tag); +} + +} // namespace + +// --------------------------------------------------------------------------- +// Handles +// --------------------------------------------------------------------------- + +// A handle over posted-but-not-waited mori-io work. +// +// Lifetime contract, inherited verbatim from the pre-refactor +// RemoteDram{Get,Put}InFlight: the RDMA backend keeps a RAW TransferStatus* +// into `statuses`, so each PostSet is heap-owned and its status vector is built +// once at final size and never resized or moved. +class MoriIoEngine::RdmaHandle final : public TransferHandle { + public: + // One BatchRead or one BatchWrite. mori-io's batch calls are direction-typed + // (local-dest + remote-src vs local-src + remote-dest), so a plan set that + // mixes directions posts one PostSet per direction. + struct PostSet { + mori::io::MemDescVec local_descs; + mori::io::MemDescVec remote_descs; + mori::io::BatchSizeVec local_offsets; + mori::io::BatchSizeVec remote_offsets; + mori::io::BatchSizeVec sizes; + std::vector statuses; // built once, never resized/moved + mori::io::TransferStatusPtrVec status_ptrs; + mori::io::TransferUniqueIdVec ids; + std::vector> tags; // per group + std::vector endpoints; // per group, remote engine key + bool is_pull = false; + }; + + ~RdmaHandle() override { + // Safety net: on an early/exceptional destroy (before Wait) drain so the + // backend's completion callback never writes freed memory. No failure + // mapping here — that is Wait's job. + if (drained_) return; + for (auto& set : sets_) { + for (auto& s : set->statuses) s.Wait(); + } + } + + void Wait(std::vector* failures) override { + if (drained_) return; + drained_ = true; + if (failures != nullptr) { + failures->insert(failures->end(), std::make_move_iterator(presettled_.begin()), + std::make_move_iterator(presettled_.end())); + } + presettled_.clear(); + // Wait every group; never break early, so no status is left live. + for (auto& set : sets_) { + for (size_t g = 0; g < set->statuses.size(); ++g) { + set->statuses[g].Wait(); + if (set->statuses[g].Succeeded()) continue; + if (failures == nullptr) continue; + TransferFailure f; + f.tags = set->tags[g]; + f.code = set->statuses[g].CodeUint32(); + f.message = set->statuses[g].Message(); + f.endpoint = set->endpoints[g]; + failures->push_back(std::move(f)); + } + } + } + + PostSet* AddPostSet() { + sets_.push_back(std::make_unique()); + return sets_.back().get(); + } + bool Empty() const { return sets_.empty(); } + void AddPresettled(TransferFailure f) { presettled_.push_back(std::move(f)); } + + private: + std::vector> sets_; + // Failures from bounce plans that Submit already ran to completion. + std::vector presettled_; + bool drained_ = false; +}; + +// --------------------------------------------------------------------------- +// Lifecycle +// --------------------------------------------------------------------------- + +MoriIoEngine::MoriIoEngine(std::string engine_key, mori::io::IOEngineConfig io_config, + uint64_t bounce_bytes) + : local_engine_key_(std::move(engine_key)), + io_config_(std::move(io_config)), + bounce_size_(bounce_bytes) {} + +MoriIoEngine::~MoriIoEngine() { Shutdown(); } + +bool MoriIoEngine::Init() { + if (io_engine_ != nullptr) return true; + if (io_config_.host.empty()) return false; + + io_engine_ = std::make_unique(local_engine_key_, io_config_); + // Adopt whatever key the engine actually published: IsRemoteMemory compares + // a descriptor's engineKey against this, so a transformed key here would make + // every local buffer look remote. + local_engine_key_ = io_engine_->GetEngineDesc().key; + + mori::io::RdmaBackendConfig rdma_cfg; + rdma_cfg.qpPerTransfer = 4; + rdma_cfg.enableTransferChunking = true; + rdma_cfg.numNicsPerTransfer = 4; + // UMBP only sets the config defaults above. All RDMA knobs (qpPerTransfer / + // postBatchSize / numWorkerThreads / pollCqMode / chunking / numNics / ...) + // are overridable via MORI_IO_* env in the RDMA backend ctor, shared by every + // IO backend user; no UMBP-specific entry points. + io_engine_->CreateBackend(mori::io::BackendType::RDMA, rdma_cfg); + + if (bounce_size_ > 0) { + bounce_buffer_ = std::make_unique(bounce_size_); + std::memset(bounce_buffer_.get(), 0, bounce_size_); + bounce_ref_ = RegisterMemory(bounce_buffer_.get(), bounce_size_, + mori::io::MemoryLocationType::CPU, /*device=*/-1); + if (!bounce_ref_.HasMemoryDesc()) { + MORI_UMBP_WARN("[MoriIoEngine] bounce buffer could not be registered; staging disabled"); + bounce_buffer_.reset(); + bounce_ref_ = TransferRef{}; + bounce_size_ = 0; + } + } + + MORI_UMBP_INFO("[MoriIoEngine] initialized on {}:{} bounce_bytes={}", io_config_.host, + io_config_.port, bounce_size_); + return true; +} + +void MoriIoEngine::Shutdown() { + { + std::lock_guard lock(remotes_mutex_); + remotes_.clear(); + } + if (io_engine_ != nullptr && bounce_buffer_ != nullptr) { + io_engine_->DeregisterMemory(bounce_ref_.mem); + } + bounce_ref_ = TransferRef{}; + bounce_buffer_.reset(); + bounce_size_ = 0; + io_engine_.reset(); +} + +// --------------------------------------------------------------------------- +// Registration +// --------------------------------------------------------------------------- + +TransferRef MoriIoEngine::RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) { + TransferRef ref = TransferRef::HostBytes(base, size, loc, device); + if (io_engine_ == nullptr || base == nullptr || size == 0) return ref; + ref.mem = io_engine_->RegisterMemory(base, size, device, loc); + return ref; +} + +void MoriIoEngine::Deregister(const TransferRef& ref) { + if (io_engine_ == nullptr || !ref.HasMemoryDesc()) return; + io_engine_->DeregisterMemory(ref.mem); +} + +std::vector MoriIoEngine::PackedLocalEngineDesc() const { + if (io_engine_ == nullptr) return {}; + msgpack::sbuffer sbuf; + msgpack::pack(sbuf, io_engine_->GetEngineDesc()); + return std::vector(sbuf.data(), sbuf.data() + sbuf.size()); +} + +// --------------------------------------------------------------------------- +// Remote endpoints +// --------------------------------------------------------------------------- + +bool MoriIoEngine::EnsureRemoteEngine(const std::string& node_id, + const std::string& packed_engine_desc) { + if (packed_engine_desc.empty()) { + // Nothing new to learn: success only if we already know this peer's engine + // (or if there is no engine at all, in which case there is nothing to + // register and the caller's transfers will be rejected by CanHandle). + if (io_engine_ == nullptr) return true; + return HasRemoteEngine(node_id); + } + + mori::io::EngineDesc desc; + try { + auto handle = msgpack::unpack(packed_engine_desc.data(), packed_engine_desc.size()); + desc = handle.get().as(); + } catch (const std::exception& e) { + MORI_UMBP_ERROR("[MoriIoEngine] EnsureRemoteEngine: bad engine desc from '{}': {}", node_id, + e.what()); + return false; + } + + if (io_engine_ != nullptr) io_engine_->RegisterRemoteEngine(desc); + + std::lock_guard lock(remotes_mutex_); + auto& remote = remotes_[node_id]; + remote.engine_desc = std::move(desc); + remote.engine_registered = (io_engine_ != nullptr); + return true; +} + +bool MoriIoEngine::HasRemoteEngine(const std::string& node_id) const { + std::lock_guard lock(remotes_mutex_); + auto it = remotes_.find(node_id); + return it != remotes_.end() && it->second.engine_registered; +} + +void MoriIoEngine::ForgetRemote(const std::string& node_id) { + std::lock_guard lock(remotes_mutex_); + remotes_.erase(node_id); +} + +void MoriIoEngine::CacheRemoteBuffers(const std::string& node_id, + const std::vector& descs) { + if (descs.empty()) return; + std::lock_guard lock(remotes_mutex_); + auto& remote = remotes_[node_id]; + for (const auto& d : descs) { + if (d.backend_id >= kMaxBackendsPerPeer) { + MORI_UMBP_ERROR("[MoriIoEngine] CacheRemoteBuffers: bad backend_id={} from '{}' (max {})", + d.backend_id, node_id, kMaxBackendsPerPeer); + continue; + } + auto& shelf = remote.buffers[d.backend_id]; + if (shelf.size() <= d.buffer_index) shelf.resize(d.buffer_index + 1); + if (d.desc_bytes.empty()) continue; + + mori::io::MemoryDesc mem; + try { + auto handle = + msgpack::unpack(reinterpret_cast(d.desc_bytes.data()), d.desc_bytes.size()); + mem = handle.get().as(); + } catch (const std::exception& e) { + MORI_UMBP_ERROR( + "[MoriIoEngine] CacheRemoteBuffers: bad desc from '{}' backend={} buffer_index={}: {}", + node_id, d.backend_id, d.buffer_index, e.what()); + continue; + } + + if (shelf[d.buffer_index].HasMemoryDesc()) { + // Already hydrated. First-write-wins is only safe while a descriptor is + // immutable for a given address — a CHANGED descriptor means two buffers + // are claiming one address, which does not fail loudly on its own: the + // reader would keep planning RDMA against valid-looking memory and return + // the wrong bytes as a hit. That is exactly how the pre-backend_id + // collision stayed invisible, so say so and keep the incumbent. + const auto& held = shelf[d.buffer_index].mem; + if (!(held == mem)) { + MORI_UMBP_ERROR( + "[MoriIoEngine] CacheRemoteBuffers: conflicting desc from '{}' backend={} " + "buffer_index={} (held id={} data={:#x} size={}, offered id={} data={:#x} size={}); " + "keeping the held one — the peer is publishing two buffers at one address", + node_id, d.backend_id, d.buffer_index, held.id, held.data, held.size, mem.id, mem.data, + mem.size); + } + continue; + } + shelf[d.buffer_index] = TransferRef::Remote(std::move(mem)); + } +} + +bool MoriIoEngine::HasRemoteBuffers(const std::string& node_id) const { + std::lock_guard lock(remotes_mutex_); + auto it = remotes_.find(node_id); + if (it == remotes_.end()) return false; + // True means "this peer's buffers are hydrated", which is what the caller + // turns into BatchResolveKeysRequest.omit_descs. GetPeerInfo publishes every + // backend at once, so any non-empty shelf means the whole handshake landed. + for (const auto& shelf : it->second.buffers) { + if (!shelf.empty()) return true; + } + return false; +} + +std::vector MoriIoEngine::RemoteBufferSnapshot(const std::string& node_id, + uint32_t backend_id) const { + if (backend_id >= kMaxBackendsPerPeer) return {}; + std::lock_guard lock(remotes_mutex_); + auto it = remotes_.find(node_id); + if (it == remotes_.end()) return {}; + return it->second.buffers[backend_id]; +} + +// --------------------------------------------------------------------------- +// Planning +// --------------------------------------------------------------------------- + +namespace { + +// A ref names memory this engine could reach over the wire. +bool IsRemoteMemory(const TransferRef& ref, const std::string& local_key) { + return ref.HasMemoryDesc() && ref.mem.engineKey != local_key; +} + +// Whether the local endpoint can be staged through the host bounce pool. +// +// Addressability is not enough: staging is a std::memcpy in and out of a +// host-allocated region, so a GPU endpoint cannot take this path at all. An +// unregistered GPU buffer is therefore rejected by Plan and surfaces as a +// per-key failure, rather than faulting inside the memcpy. This is the single +// place the decision is made, which is why it belongs here and not at each of +// PoolClient's call sites — the transfer layer is the only layer that knows +// whether staging is involved. +bool IsBounceable(const TransferRef& local) { + return local.HasHostPtr() && local.loc == mori::io::MemoryLocationType::CPU; +} + +} // namespace + +bool MoriIoEngine::CanHandle(const TransferRef& src, const TransferRef& dst) const { + if (io_engine_ == nullptr) return false; + const bool src_remote = IsRemoteMemory(src, local_engine_key_); + const bool dst_remote = IsRemoteMemory(dst, local_engine_key_); + if (src_remote == dst_remote) return false; // both local or both remote + const TransferRef& local = src_remote ? dst : src; + if (local.HasMemoryDesc()) return true; // zero copy + return IsBounceable(local) && bounce_size_ > 0; // staged through the pool +} + +TransferPlanSet MoriIoEngine::Plan(const std::vector& items) const { + TransferPlanSet out; + if (items.empty()) return out; + + // Direct (zero-copy) groups, bucketed by the endpoint MR id pair; a bucket + // holds candidate plan indices disambiguated by engine key (see PairKey). + // First-appearance order is preserved so submit order stays deterministic. + std::unordered_map, PairKeyHash> pair_to_plans; + pair_to_plans.reserve(items.size() * 2); + + // Bounce groups, bucketed by (remote MR id, direction). Each accumulates + // until the pool is full, then a fresh plan starts — several plans over one + // pool are fine because Submit runs each to completion before the next. + std::unordered_map, PairKeyHash> bounce_to_plans; + std::vector bounce_cursor; // parallel to out.plans, only for bounce plans + + for (const auto& item : items) { + if (item.size == 0) continue; + const bool src_remote = IsRemoteMemory(item.src, local_engine_key_); + const bool dst_remote = IsRemoteMemory(item.dst, local_engine_key_); + if (src_remote == dst_remote) { + out.rejected_tags.push_back(item.tag); + continue; + } + const TransferDirection dir = dst_remote ? TransferDirection::kPush : TransferDirection::kPull; + const TransferRef& local = src_remote ? item.dst : item.src; + const TransferRef& remote = src_remote ? item.src : item.dst; + const uint64_t local_off = src_remote ? item.dst_offset : item.src_offset; + const uint64_t remote_off = src_remote ? item.src_offset : item.dst_offset; + + if (local.HasMemoryDesc()) { + const PairKey key{item.src.mem.id, item.dst.mem.id}; + auto& candidates = pair_to_plans[key]; + size_t pi = out.plans.size(); + bool found = false; + for (size_t cand : candidates) { + if (SameEndpoints(out.plans[cand], item.src, item.dst)) { + pi = cand; + found = true; + break; + } + } + if (!found) { + candidates.push_back(pi); + TransferPlan plan; + plan.src = item.src; + plan.dst = item.dst; + plan.dir = dir; + out.plans.push_back(std::move(plan)); + bounce_cursor.push_back(0); + } + AppendSegment(&out.plans[pi], item.src_offset, item.dst_offset, item.size, item.tag); + continue; + } + + // Staged: the local endpoint is addressable host memory but not registered. + if (bounce_size_ == 0 || !IsBounceable(local) || item.size > bounce_size_) { + out.rejected_tags.push_back(item.tag); + continue; + } + const PairKey key{remote.mem.id, static_cast(dir)}; + auto& bounce_candidates = bounce_to_plans[key]; + size_t pi = out.plans.size(); + bool fresh = true; + for (size_t& cand : bounce_candidates) { + const TransferPlan& p = out.plans[cand]; + const TransferRef& p_remote = (dir == TransferDirection::kPush) ? p.dst : p.src; + if (p.dir != dir || p_remote.mem.engineKey != remote.mem.engineKey) continue; + // Right endpoint; reuse it only while the pool still has room, otherwise + // retire it and open a fresh plan in its place. + if (bounce_cursor[cand] + item.size <= bounce_size_) { + pi = cand; + fresh = false; + } else { + cand = pi; // this bucket entry now names the plan we are about to add + } + break; + } + if (fresh) { + if (bounce_candidates.empty() || std::find(bounce_candidates.begin(), bounce_candidates.end(), + pi) == bounce_candidates.end()) { + bounce_candidates.push_back(pi); + } + TransferPlan plan; + plan.dir = dir; + plan.uses_bounce = true; + // The local endpoint IS the bounce region; the remote side is unchanged. + if (dir == TransferDirection::kPush) { + plan.src = bounce_ref_; + plan.dst = remote; + } else { + plan.src = remote; + plan.dst = bounce_ref_; + } + out.plans.push_back(std::move(plan)); + bounce_cursor.push_back(0); + } + TransferPlan& plan = out.plans[pi]; + const uint64_t bounce_off = bounce_cursor[pi]; + bounce_cursor[pi] += item.size; + if (dir == TransferDirection::kPush) { + AppendSegment(&plan, bounce_off, remote_off, item.size, item.tag); + } else { + AppendSegment(&plan, remote_off, bounce_off, item.size, item.tag); + } + plan.bounce_copies.push_back(TransferPlan::BounceCopy{ + static_cast(local.host_ptr) + local_off, bounce_off, item.size}); + } + + return out; +} + +// --------------------------------------------------------------------------- +// Submission +// --------------------------------------------------------------------------- + +bool MoriIoEngine::PostPlans(const std::vector& plans, RdmaHandle* handle) { + // One PostSet per direction: mori-io's batch calls are direction-typed. + RdmaHandle::PostSet* pull = nullptr; + RdmaHandle::PostSet* push = nullptr; + auto set_for = [&](TransferDirection dir) -> RdmaHandle::PostSet* { + RdmaHandle::PostSet*& slot = (dir == TransferDirection::kPull) ? pull : push; + if (slot == nullptr) { + slot = handle->AddPostSet(); + slot->is_pull = (dir == TransferDirection::kPull); + } + return slot; + }; + + for (const auto& plan : plans) { + RdmaHandle::PostSet* set = set_for(plan.dir); + const bool pulling = (plan.dir == TransferDirection::kPull); + // local = dst when pulling, src when pushing. + set->local_descs.push_back(pulling ? plan.dst.mem : plan.src.mem); + set->remote_descs.push_back(pulling ? plan.src.mem : plan.dst.mem); + set->local_offsets.push_back(pulling ? plan.dst_offsets : plan.src_offsets); + set->remote_offsets.push_back(pulling ? plan.src_offsets : plan.dst_offsets); + set->sizes.push_back(plan.sizes); + set->tags.push_back(plan.tags); + set->endpoints.push_back(pulling ? plan.src.mem.engineKey : plan.dst.mem.engineKey); + } + + bool posted = false; + for (RdmaHandle::PostSet* set : {pull, push}) { + if (set == nullptr) continue; + const size_t g = set->local_descs.size(); + // Built once at final size and never resized/moved: the backend holds raw + // TransferStatus* into this vector for the life of the post. + set->statuses = std::vector(g); + set->status_ptrs.resize(g); + set->ids.resize(g); + for (size_t i = 0; i < g; ++i) { + set->status_ptrs[i] = &set->statuses[i]; + set->ids[i] = io_engine_->AllocateTransferUniqueId(); + } + if (set->is_pull) { + io_engine_->BatchRead(set->local_descs, set->local_offsets, set->remote_descs, + set->remote_offsets, set->sizes, set->status_ptrs, set->ids); + } else { + io_engine_->BatchWrite(set->local_descs, set->local_offsets, set->remote_descs, + set->remote_offsets, set->sizes, set->status_ptrs, set->ids); + } + posted = true; + } + return posted; +} + +bool MoriIoEngine::RunBouncePlanInline(const TransferPlan& plan, TransferFailure* failure) { + auto fail = [&](uint32_t code, std::string msg) { + failure->tags = plan.tags; + failure->code = code; + failure->message = std::move(msg); + failure->endpoint = + (plan.dir == TransferDirection::kPull) ? plan.src.mem.engineKey : plan.dst.mem.engineKey; + return false; + }; + + // Held for the whole staged round trip. Never held across a return, which is + // why a submit-all loop over several peers cannot deadlock on it. + std::lock_guard lock(bounce_mutex_); + if (bounce_buffer_ == nullptr) return fail(0, "bounce buffer not configured"); + + const bool pushing = (plan.dir == TransferDirection::kPush); + if (pushing) { + for (const auto& c : plan.bounce_copies) { + if (c.bounce_offset + c.size > bounce_size_) return fail(0, "bounce offset overflow"); + std::memcpy(bounce_buffer_.get() + c.bounce_offset, c.host, c.size); + } + } + + RdmaHandle inflight; + if (!PostPlans({plan}, &inflight)) return fail(0, "post failed"); + std::vector failures; + inflight.Wait(&failures); + if (!failures.empty()) { + *failure = std::move(failures.front()); + return false; + } + + if (!pushing) { + for (const auto& c : plan.bounce_copies) { + if (c.bounce_offset + c.size > bounce_size_) return fail(0, "bounce offset overflow"); + std::memcpy(c.host, bounce_buffer_.get() + c.bounce_offset, c.size); + } + } + return true; +} + +std::unique_ptr MoriIoEngine::Submit(std::vector plans) { + if (plans.empty() || io_engine_ == nullptr) return nullptr; + + // Post the zero-copy plans FIRST (they do not block), then run the staged + // ones inline so the staged round trips overlap the posted wire rather than + // serializing behind it. + std::vector direct; + std::vector bounced; + direct.reserve(plans.size()); + for (size_t i = 0; i < plans.size(); ++i) { + if (plans[i].uses_bounce) { + bounced.push_back(i); + } else { + direct.push_back(std::move(plans[i])); + } + } + + auto handle = std::make_unique(); + bool any = false; + if (!direct.empty()) any = PostPlans(direct, handle.get()); + + for (size_t i : bounced) { + TransferFailure failure; + if (RunBouncePlanInline(plans[i], &failure)) { + any = true; + continue; + } + MORI_UMBP_ERROR("[MoriIoEngine] staged transfer failed: code={} msg='{}' peer_engine='{}'", + failure.code, failure.message, failure.endpoint); + handle->AddPresettled(std::move(failure)); + any = true; + } + + if (!any) return nullptr; + return handle; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/transfer/transfer_engine.cpp b/src/umbp/distributed/transfer/transfer_engine.cpp new file mode 100644 index 000000000..6869fc309 --- /dev/null +++ b/src/umbp/distributed/transfer/transfer_engine.cpp @@ -0,0 +1,107 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/transfer/transfer_engine.h" + +#include +#include + +namespace mori::umbp { + +namespace { + +// Adds its lifetime to a sink; a null sink makes it inert, so an untimed +// Transfer does not even read the clock. +class StepTimer { + public: + explicit StepTimer(double* sink) + : sink_(sink), + start_(sink ? std::chrono::steady_clock::now() : std::chrono::steady_clock::time_point{}) {} + ~StepTimer() { Stop(); } + StepTimer(const StepTimer&) = delete; + StepTimer& operator=(const StepTimer&) = delete; + void Stop() { + if (sink_ == nullptr) return; + *sink_ += std::chrono::duration_cast>( + std::chrono::steady_clock::now() - start_) + .count(); + sink_ = nullptr; + } + + private: + double* sink_; + std::chrono::steady_clock::time_point start_; +}; + +} // namespace + +bool TransferEngine::Transfer(const std::vector& items, + std::vector* failed_tags, const StepTiming* timing) { + if (items.empty()) return true; + + TransferPlanSet planned; + { + StepTimer plan_timer(timing ? timing->plan : nullptr); + planned = Plan(items); + } + bool ok = planned.rejected_tags.empty(); + if (failed_tags != nullptr) { + failed_tags->insert(failed_tags->end(), planned.rejected_tags.begin(), + planned.rejected_tags.end()); + } + if (planned.plans.empty()) return ok; + + // Snapshot every tag that made it into a plan BEFORE the move, so a Submit + // that posts nothing can still fail exactly those keys. + std::vector planned_tags; + for (const auto& plan : planned.plans) { + planned_tags.insert(planned_tags.end(), plan.tags.begin(), plan.tags.end()); + } + + std::unique_ptr handle; + { + StepTimer submit_timer(timing ? timing->submit : nullptr); + handle = Submit(std::move(planned.plans)); + } + if (handle == nullptr) { + if (failed_tags != nullptr) { + failed_tags->insert(failed_tags->end(), planned_tags.begin(), planned_tags.end()); + } + return false; + } + + std::vector failures; + { + StepTimer wait_timer(timing ? timing->wait : nullptr); + handle->Wait(&failures); + } + if (failures.empty()) return ok; + if (failed_tags != nullptr) { + for (const auto& f : failures) { + failed_tags->insert(failed_tags->end(), f.tags.begin(), f.tags.end()); + } + std::sort(failed_tags->begin(), failed_tags->end()); + failed_tags->erase(std::unique(failed_tags->begin(), failed_tags->end()), failed_tags->end()); + } + return false; +} + +} // namespace mori::umbp diff --git a/src/umbp/doc/design-backend-agnostic-refactor.md b/src/umbp/doc/design-backend-agnostic-refactor.md new file mode 100644 index 000000000..676621cb4 --- /dev/null +++ b/src/umbp/doc/design-backend-agnostic-refactor.md @@ -0,0 +1,849 @@ +# UMBP Backend-Agnostic Refactor — Plan + +**Goal:** make UMBP's data plane, control plane, and routing plane backend-agnostic. +Adding a new storage medium becomes: implement one interface + register it. No +call-site edits in `PoolClient`, `PeerServiceServer`, or the routing strategies. + +**Method:** remove SSD **from distributed mode only**, then generalize what +remains. Local (standalone) mode keeps its SSD tier, and the whole storage +implementation stack underneath it — SPDK, the io_uring/posix drivers, the +offset allocator, the proxy protocol — stays in tree, built and tested, +untouched. + +This is deliberate and temporary: SSD is expected back in distributed mode +*after* the refactor, as an `SsdBackend : MediumBackend`. Taking it out first is +purely to make the refactor simple — the distributed data plane collapses to one +shape (async RDMA page slots) while it is being generalized, and SSD re-enters +through the finished abstraction instead of being carried through every phase. +So the rule for Phase 0 is: **delete the coupling, keep the implementation.** +Anything reusable by a future SSD backend is preserved. + +During Phases 1–5 the live distributed media are DRAM and HBM, which are already +two instances of the same page-slot model. + +**Scope:** `src/umbp/` — `distributed/` (pool client, peer service, routing), +`local/` (standalone stack), `storage/`. + +--- + +## 1. What is actually coupled + +Four specific things, not a diffuse mess. + +**1. `PoolClient` hardcodes two different data-plane shapes.** +DRAM is async — `SubmitRemoteBatchGet` / `WaitRemoteBatchGet`, RDMA straight +into the peer's pages. SSD is blocking-only — `RemoteSsdReadOnce`, which does a +key-based `PrepareSsdRead` into a peer-side staging slot, takes a lease, RDMAs +out of the staging buffer, then releases. `BatchGetPlan`'s 4-way partition +(`remote_dram_groups` / `remote_ssd_groups` / `local_dram_indices` / +`local_ssd_indices`) exists only to reconcile these two shapes inside one batch +while preserving overlap. + +This asymmetry is roughly 90% of the refactoring cost, and it disappears the +moment SSD leaves the *distributed* path. Nothing about it is inherent to the +storage medium — it is inherent to `PrepareSsdRead`'s key-based, lease-held, +staging-slot shape. + +**2. `PeerServiceServer` takes four typed pointers.** +`PeerDramAllocator*`, `PeerSsdManager*`, `SsdCopyPipeline*`, plus the SSD +staging base/size/`MemoryDesc` args. The SSD-only RPCs (`PrepareSsdRead`, +`ReleaseSsdLease`) and the read-slot lease state machine live in the handlers. + +**3. Routing is nearly agnostic already.** *(resolved — Phase 4)* +`RoutePutStrategy` and `RouteGetStrategy` only ever see +`ClientRecord.tier_capacities` and `Location.tier`. The single hardcode is two +orderings: `TierPriorityRouteGetStrategy`'s `HBM > DRAM > SSD`, and put's +"HBM then DRAM, SSD is never a direct-put target." Both were **deleted** in +Phase 4 rather than re-expressed as advertised properties — see §5 Phase 4. + +**4. Memory ownership lives entirely outside the backend.** +This is the one the original plan missed, and it is the one that makes the +Phase 5 acceptance test unreachable. + +| Concern | Where it lives today | Medium-specific? | +|---|---|---| +| Pool allocation | `distributed_client.cpp:41-53` — `HostMemAllocator`, hugepages, NUMA, prefault | yes | +| Pool free | `distributed_client.cpp:271-276` | yes | +| RDMA registration | `pool_client.cpp:360-366` — `MemoryLocationType::CPU` hardcoded | yes | +| Desc/base marshalling into the allocator | `BuildDramTierConfig`, `pool_client.cpp:286-303` | yes | +| Local put/get byte copy | `pool_client.cpp:559-599` + `LocalCopyBlock` `:181-190` (AVX2-NT / `memcpy`) | yes | +| Buffer indexing | `config_.dram_buffers[p.buffer_index]` `:564`, `:585` — **no tier dimension** | yes | +| Copy parallelism | `UMBP_DRAM_{WRITE,READ}_THREADS` `:836`, `:1435` | yes | +| Capacity advertisement | `distributed_client.cpp:70-71` — `{DRAM, {size,size}}` literal | yes | +| Re-cache install tier | `pool_client.cpp:718` — `TierType::DRAM` literal | yes | +| Put tier filter | `pool_client.cpp:821` — `!= DRAM && != HBM` | yes | + +`PeerDramAllocator` owns only bookkeeping: the page bitmap, the key→slot map, +pending slots and their reaper, read leases, copy pins, the event outbox, the +clear generation. Everything about the memory itself arrives pre-cooked through +`BuildDramTierConfig`. + +So §3's premise — that a backend owns the slots it publishes descriptors for — +is aspirational. Today the client allocates, registers, and copies; the backend +is a bookkeeper over memory it does not own. Adding HBM under this ownership needs +a new config buffer list, a second `RegisterMemory` with `GPU`, a tier branch in +`LocalCopyBlock`, and a fix for the missing tier dimension: four call-site edits +before the registry line. + +**Not part of this** — a separate axis that must not be moved into a backend: +`PoolClient::RegisterMemory` (`:526`) registers the **caller's** buffer as +`MemoryLocationType::CPU` unconditionally. That is wrong for a GPU `src`/`dst` +(SGLang HiCache KV pages), but it is a property of the caller's allocation, not +of any storage medium. It needs a location parameter on the public API, not a +backend. + +**Key observation:** `PeerDramAllocator` is *shaped* for multiple media — an +internal `map` plus parallel `tier_descs_` / +`tier_bases_` maps, with a `TierType` parameter threaded through every method — +but it serves exactly **one** live tier. `PoolClient::Init` default-constructs an +empty `hbm_cfg` (`pool_client.cpp:377`, "HBM not currently exposed via +PoolClientConfig"), so no HBM buffer is ever passed in. `TierType::HBM`, its +routing rank, and the tier map are unexercised scaffolding. + +The refactor *inverts* that map: one backend instance per medium, instead of one +object holding a map of media. This deletes generality that was never live and +relocates it to the registry, where it is exercised. + +**Second key observation — the distributed SSD code is a thin adapter, not a +storage stack.** `PeerSsdManager` (665 LOC) forward-declares `TierBackend` and +holds one (`peer_ssd_manager.h:43,65`); the actual bytes are moved by local +mode's `SsdTier` / `SpdkSsdTier` over `storage/io/**` and `storage/spdk/**` +(~5.5k LOC). Distributed SSD therefore sits *on top of* local SSD. Unwiring it +from distributed mode is a cut at the adapter seam: it touches none of the +storage implementation, and — because the adapter's own dependencies are just +`config.h` and `types.h` — not even the adapter itself has to go. What is +actually deleted is the *lease-shaped call sites*, not SSD code. That is why +"keep it, just stop calling it" is cheap rather than a compromise. + +--- + +## 2. The boundary + +Everything below follows from one property of `mori::io::MemoryDesc` +(`include/mori/io/common.hpp:100-127`): it is **self-describing about the +medium**. + +```cpp +MemoryLocationType loc; // CPU | GPU +int deviceId; std::string deviceBusId; int numaNode; +std::array ipcHandle; +std::array fabricHandle; int vpodId; +``` + +Hence the criterion for deciding where any piece of code belongs: + +> **Bytes addressed by a descriptor are already medium-agnostic. Bytes addressed +> by a raw pointer are medium-specific. Registration is the one-way conversion +> between the two, and it belongs to whoever allocated the memory.** + +Three consequences, each checkable against the tree: + +**The remote path already works for a medium it was never written for.** A peer +whose pages are HBM can be read today, unmodified, because that peer registered +them and the desc carries `loc=GPU`. `RemoteDramScatterRead`, +`GroupTransfersByPair`, `SubmitRemoteBatchGet` and the peer desc cache never +dereference anything — they only pass descs to the IO engine. None of them +needed this refactor and none of them change. + +**The local path breaks because it contains no descriptor at all.** +`LocalCopyBlock(static_cast(dram.buffer) + off, ...)` (`:575`) operates on +a raw `void*` out of `config_.dram_buffers`, bypassing the descriptor layer. + +**That is a deliberate performance choice, not an oversight.** The local fast +path exists because a `memcpy` beats an engine round trip. Medium-awareness is +the price already being paid for it. This matters for §4: the fast path folds +back into the transfer layer only if a local transfer through the abstraction is +genuinely as cheap, which is a measurement, not an argument. + +Practical form of the rule, usable in review: **if adding a new medium would +force you to edit it, it belongs to the backend.** + +--- + +## 3. The abstraction + +The system has three components. The refactor's job is to make the boundaries +between them real. + +| Component | Owns | Knows about media? | +|---|---|---| +| **Control plane** — master, routing, heartbeat | who holds what, capacity, eviction, placement | only via advertised `BackendProperties` | +| **Storage backends** — DRAM, HBM, SSD, … | bytes, the key→slot map, capacity, eviction | yes — this is the only layer that does | +| **Transfer engine** (§4) | moving bytes between descriptors | no — descriptors are self-describing | + +New header: `include/umbp/distributed/peer/medium_backend.h`. + +A backend's job is to **publish descriptors** for the slots it owns, and to keep +the key→slot bookkeeping around them. It does not move bytes. + +```cpp +class MediumBackend { + public: + virtual ~MediumBackend() = default; + + // ---- identity ---- + virtual TierType Tier() const = 0; + virtual const char* Name() const = 0; + + // ---- ownership (new; see §1 item 4) ---- + // The backend allocates its own pool with its own policy (hugepage/NUMA vs + // hipMalloc vs file extents) and registers it, choosing the location type + // itself. PoolClient hands it the engine; it does not hand it memory. + virtual bool Init(TransferEngine*) = 0; + virtual void Shutdown() = 0; + + // ---- control plane (shipped by the heartbeat) ---- + virtual TierCapacity Capacity() const = 0; + virtual uint64_t OwnedKeyCount() const = 0; + virtual std::vector DrainPendingEvents() = 0; + virtual std::vector SnapshotOwnedKeys() const = 0; + virtual std::vector SnapshotOwnedKeysForFullSync() = 0; + virtual void ClearLocal() = 0; + virtual void ClearFullSyncAcked() = 0; + virtual bool IsClearFullSyncPending() const = 0; + + // ---- slot lifecycle (peer-side; driven by PeerServiceServer handlers) ---- + virtual std::vector BatchAllocate(const std::vector&) = 0; + virtual std::vector BatchCommit(const std::vector&) = 0; + virtual std::vector BatchAbort(const std::vector& slot_ids) = 0; + virtual std::vector BatchResolve(const std::vector& keys, + bool include_descs) = 0; + virtual std::vector Evict(const std::vector& keys) = 0; + + // ---- bootstrap (GetPeerInfo) ---- + virtual uint64_t PageSize() const = 0; + virtual std::vector AllBufferDescs() const = 0; +}; +``` + +`AllocateResult` and `ResolvedEntry` carry `{slot_id, pages, size, page_size, +descs, pending_ttl_ms}`. No method takes a `TierType`: one instance is one +medium, and the caller has already dispatched on it via the registry. + +**`BackendProperties` was proposed here and then NOT adopted.** The idea was to +make read order and put-eligibility *advertised* facts rather than a hand-written +tier list in the strategies. Phase 4 found the simpler answer: every medium in +the system today is equivalent, so the tier lists were **deleted outright** and +nothing replaced them. An advertised order would have been scaffolding nothing +exercises — exactly the mistake §1's "key observation" calls out about +`PeerDramAllocator`'s unexercised tier map. + +The trait comes back when a medium that genuinely differs does. That is SSD: +it takes no direct puts. `SsdBackend` reintroduces `put_eligible` *with* the +backend that needs it, which is the same "1 file + 1 registry line" test Phase 5 +already applies. + +**What is deliberately NOT on this interface**, and why — each of these was +proposed and rejected once §4 was settled: + +- **`LocalCopyIn` / `LocalCopyOut`.** Local access is a transfer whose endpoints + are both local; §4 selects a memcpy or pread engine for it. Putting a copy + method on the backend would re-admit a second byte-moving path. +- **A staging/bounce pool.** Belongs to the transfer layer, which is the only + layer that can observe transfer completion (§4). +- **A "not ready / retry here" resolve state.** Only needed to express + staging-pool exhaustion, which is no longer a backend resource. + +All three were artifacts of a transfer engine that speaks only memory. They are +listed here so they are not rediscovered and re-added. + +This interface absorbs the existing `OwnedLocationSource` — that type goes away. +The dormant `PeerSsdManager` (Phase 0) keeps its `DrainPendingEvents` / +`SnapshotOwnedKeys` methods but drops the `: public OwnedLocationSource` base and +the `override`s — a ~4-line change. + +Ownership: a `BackendRegistry` (`map>`) held +by `PoolClient` is the single place in the tree where a concrete backend type is +named. It exposes `Get(tier)`, `All()`, and `ByReadRank()` — the last being the +peer-local counterpart of `RouteGetStrategy`'s cross-peer ranking, used when a +key is mirrored across media. + +--- + +## 4. The transfer engine + +`mori-io` is a **memory-to-memory** engine: `BackendType` is +`{XGMI, RDMA, TCP, FABRIC}` and `MemoryLocationType` is `{CPU, GPU}` +(`include/mori/io/enum.hpp:27-41`). There is no file or object endpoint. + +**Decision: do not change mori-io.** Abstract the transfer engine inside UMBP +and make mori-io one implementation. UMBP defines its own descriptor variant, so +mori-io never has to learn about files: + +```cpp +struct TransferRef { + std::variant // bucket + key + target; +}; + +class TransferEngine { + public: + virtual ~TransferEngine() = default; + + // ---- registration ---- + // The backend supplies the facts a descriptor cannot recover — location type, + // device, NUMA node — because it is the allocator (§6). + virtual TransferRef RegisterMemory(void* base, size_t size, MemoryLocationType, int device) = 0; + virtual TransferRef RegisterFile(int fd, uint64_t offset, uint64_t size) = 0; + virtual void Deregister(const TransferRef&) = 0; + + // ---- transfer ---- + virtual bool CanHandle(const TransferRef& src, const TransferRef& dst) const = 0; + virtual std::vector Plan(const std::vector&) = 0; // see below + virtual Handle Submit(const std::vector&) = 0; + virtual bool Wait(Handle&) = 0; +}; +``` + +`MoriIoEngine` accepts only the `MemoryRef` variant and unwraps straight to +`mori::io::MemoryDesc`. A `PosixEngine` / `GdsEngine` accepts a `FileRef` on one +side. Engine selection is a function of the `(src, dst)` pair. + +**One type, including registration.** Registration and transfer are different +shapes — registering fans *out* across every transport that can reach an +endpoint, while a transfer dispatches to exactly *one*. mori-io already carries +that fan-out in its descriptor: `MemoryDesc` holds `ipcHandle` and +`fabricHandle` side by side, "so XGMI and Fabric backends can register the same +memory without clobbering each other's handle." + +That difference does not need a second class. `CompositeTransferEngine : +TransferEngine` holds the concrete engines, implements `RegisterMemory` / +`RegisterFile` by fanning out and merging the per-transport handles into one +`TransferRef`, and implements `Plan` / `Submit` by selecting or chaining. The +composite *is* a `TransferEngine`, so the rest of the system holds exactly one +pointer to exactly one type: `PoolClient` owns it, and backends receive the same +pointer at `Init`. + +The cost of folding registration in: a backend now holds something it *could* +call `Submit` on. That invariant — backends do not move bytes — moves from the +type system to Phase 5's lint, as one rule: no file under a backend directory may +call `Submit` / `Plan` / `Wait`. + +**`Plan()` is virtual on purpose.** `GroupTransfersByPair` — collapsing a +scatter into one transfer per `(localMR, remoteMR)` — is an RDMA optimization +that cuts CQE and post count. A memcpy engine gains nothing from it; a file +engine would rather group by file and merge adjacent offset ranges. Submit-all- +then-wait *is* universal and belongs in the base class; the grouping strategy is +not. Note the current implementation assumes one page size per MR pair, so +pushing it down requires generalizing it first. + +**Bounce buffers belong here, not in a backend.** The transfer layer is the only +layer with a completion signal (`TransferStatus::Wait()`). A staging pool inside +a storage backend must guess with a TTL, because the backend cannot observe when +a remote reader's RDMA finishes — that is exactly the shape of the +`PrepareSsdRead` lease this refactor exists to delete. One pool, owned by the +transfer layer, released on completion rather than on a timer. + +**Chaining.** When no single engine spans the pair — a 3FS file on node A, a +reader on node B — the layer composes: a file engine fills a bounce buffer, the +RDMA engine ships it. The extra hop is real, but it is an explicit plan in one +place, and if mori-io ever grows a storage backend the planner simply stops +chaining. Nothing above notices. + +**What stays in the client**, because it needs semantics the engine must not +have: which node and tier to target; retry-elsewhere with an exclude set (needs +to know replicas exist); the slot lifecycle, where `Commit` is what publishes a +key; capacity, eviction, and leases. + +--- + +## 5. Phases + +### Phase 0 — unwire SSD from distributed mode *(done — `ef081fd2`)* + +Cut at the adapter seam. Local mode is not touched; no storage-layer code is +deleted. + +**Deleted — the coupling only:** + +- `include/umbp/distributed/ssd_read_lease.h` (76 LOC) and + `test_ssd_read_lease_gating` — the lease *is* the coupling, and the re-added + backend will not use it +- The construction and wiring of `PeerSsdManager` / `SsdCopyPipeline` in + `distributed_client.cpp`, `pool_client.{h,cpp}`, `peer_service.{h,cpp}`, and + the `AddOwnedLocationSource(ssd_manager)` arm in `master_client.cpp` +- Every SSD member and method in `pool_client.{h,cpp}` — `ExecuteLocalSsdGet`, + `RemoteSsdReadOnce`, `ReleaseSsdLeaseBestEffort`, `ProcessRemoteSsdBatchGet`, + `PublishSsdMetrics`, `SsdMetricsLastShipped`, `SsdGetOutcome`, the SSD staging + buffer and its `MemoryDesc` (~470 SSD-touching lines across the pair) +- The SSD read-slot state machine, `StagingMetrics`, and the staging ctor args in + `peer_service.{h,cpp}` +- `PrepareSsdRead` / `ReleaseSsdLease` from `distributed/proto/umbp_peer.proto` + — `reserved` the field/method tags rather than freeing them +- SSD counters in `master/master_metrics.h` +- The 3 tests that exercise the *coupling* rather than the adapter: + `test_peer_ssd_read_rpc`, `test_ssd_read_lease_gating`, `test_ssd_reliability`, + plus the SSD arms of `test_peer_service` + +`BatchGetPlan` collapsed from 4 buckets to 2 (`remote_groups`, `local_indices`). + +`TierType::SSD = 3` stays reserved in the enum so wire tags stay stable; every +*live* distributed code path that produces or consumes it is gone. + +**Kept — dormant but still compiled and still tested:** + +- `distributed/peer/ssd/peer_ssd_manager.{h,cpp}` (665 LOC) and + `distributed/peer/ssd/ssd_copy_pipeline.{h,cpp}` (~250 LOC). Nothing + constructs them; they remain in the build. Moved under + `distributed/peer/ssd/` so "dormant" is legible in the tree and the Phase 5 + lint exemption can be a directory rather than a filename list. +- `PeerSsdConfig` (`distributed/config.h`) — `peer_ssd_manager.h` needs it. +- Their 3 unit tests: `test_peer_ssd_manager`, `test_peer_ssd_eviction`, + `test_ssd_copy_pipeline`. These test the adapter standalone and stay green. + +**Kept — untouched, still built, still tested:** `storage/io/**`, +`storage/spdk/**` and the `3rdparty/spdk` submodule, `local/tiers/*` including +`LocalStorageManager`'s demote/promote, `UMBPSsdConfig` and its Python binding, +and the local SSD test suites. + +*Gate:* ctest green — **including the local SSD suites and the 3 dormant adapter +tests**, unmodified. That is the check that the seam was cut in the right place. + +### Phase 1 — write the interfaces + +Add `medium_backend.h` (§3). Nothing implements it yet. + +*Gate:* tree compiles. Note this gate is weak by construction — a header nothing +includes is not in any translation unit. Verify self-containment explicitly +(`-fsyntax-only` against the project's real flags) rather than relying on the +build. + +### Phase 2 — DRAM becomes a backend, and takes ownership of its memory + +Two pieces, and the second is the one §1 item 4 added: + +**(a) Split the map.** `PeerDramAllocator`'s internal `map` +becomes one instance per medium. Rename to `PageBackend : MediumBackend`. Every +method loses its `TierType` parameter — this makes the file *smaller*. + +**(b) Move ownership down.** The backend allocates its own pool and registers it, +choosing `MemoryLocationType` itself. Consequences up the stack: +`PoolClientConfig::dram_buffers` is deleted; `DistributedClient` stops calling +`HostMemAllocator`; `BuildDramTierConfig` is deleted; the +`{DRAM, {size,size}}` capacity literal is deleted in favour of `Capacity()` +aggregated over the registry. + +This does **not** wait for Phase 6. The `TransferEngine` a backend receives at +`Init` is, in this phase, a ~30-line `PoolClient`-owned shim whose registration +methods forward to `io_engine_->RegisterMemory` and whose transfer methods are +unimplemented — nothing calls them yet, since the remote path still runs through +`PoolClient` until Phase 6. Phase 6 replaces the shim with +`CompositeTransferEngine` and no backend changes. Phase 1's header only +forward-declares the type, so it compiles with nothing behind it. + +Consequence to accept deliberately: `owned_`, the read lease, and the clear +generation become per-medium, so a key may exist in both DRAM and HBM. The +design already allows this (see the `NodeMatch` mirror comment in `types.h`). +Peer-local tier selection at `Resolve` time becomes a small policy function +driven by `BackendRegistry::ByReadRank()`. + +**The local fast path is NOT touched in this phase.** `LocalPutPages` / +`LocalGetPages` stay in `PoolClient`, host-DRAM-only. That is sound as long as +only one *real* medium is live, which is why option (b) below is the default. + +**HBM is not a free second backend.** It does not work today — no HBM buffer is +ever allocated or registered (§1). Two options: + +- **(a) Wire HBM up for real.** This is the genuine proof of the abstraction, + but it is *blocked on Phase 6*: a second real medium makes the host-only local + copy path wrong, and `config_.dram_buffers[buffer_index]` has no tier + dimension. Doing it before Phase 6 means interim scaffolding that Phase 6 then + deletes. +- **(b) `MockBackend` as the second registered backend.** Cheap, proves the + registry and dispatch paths, keeps the local path honestly single-medium. + +Default to (b). + +*Gate:* DRAM behavior unchanged with a second backend registered and dispatched +to; no buffer pointer crosses into `PoolClientConfig`. + +### Phase 3 — peer server and pool client go agnostic + +`PeerServiceServer` takes a `BackendRegistry*` instead of typed pointers; +handlers dispatch on the request's tier tag. Single-key RPCs are served by +one-element batches. + +`PoolClient`'s local paths go through `registry_.Get(tier)`. The remote paths are +already tier-tagged on the wire, so once the SSD special case is gone they become +generic with no further work. + +*Gate:* `bench_pool_client_batch_get` throughput unchanged. The submit/wait +overlap is a throughput property that no assertion catches — it must be measured, +not asserted. + +### Phase 4 — routing plane *(done)* + +Delete both hardcoded tier orders. **Scope decision: delete, do not replace.** +The original plan derived read order and put-eligibility from an advertised +`BackendProperties`; with every live medium equivalent there is nothing to rank, +so the orders were removed and no mechanism took their place (see §3). + +What went: + +- `TierReadRank` and the `HBM > DRAM > SSD` read order. + `TierPriorityRouteGetStrategy` → `LocalPreferringRouteGetStrategy`: it keeps + the requester-local preference (the thing that makes `cache_remote_fetches` + pay off) and drops the ranking. The old name described the deleted half. +- `kPutTierOrder = {HBM, DRAM}` and its "SSD is never a direct-put target" + exclusion. `SelectByAlgo` now scores every `(node, tier)` pair that has room, + on free space alone. +- `BackendProperties`, `MediumBackend::Properties()`, and + `BackendRegistry::ByReadRank()` — the last became redundant with `All()`, + which already iterates in a deterministic (ascending-`TierType`) order. +- `EvictionManager`'s `if (tier == SSD) continue`. It existed because `EvictKey` + only ever reached the peer's DRAM allocator; since Phase 3 it fans out to + every backend by key, so an overloaded medium now evicts from itself. + +**Behavior changes to know about** (all invisible while DRAM is the only live +medium, all asserted by rewritten tests): + +| Situation | Before | After | +|---|---|---| +| Requester holds a replica on a "slower" tier | remote faster tier wins | local replica wins (no RDMA) | +| Node advertises HBM 10G / DRAM 400G | HBM | DRAM (more room) | +| Only SSD has room | unroutable | routed to SSD | +| Same-node affinity, anchor's tier fills | spill to a remote node's faster tier | spill stays on the anchor node, changes tier | + +*Gate (revised):* the original gate — "`test_tier_priority_route_get` and +`test_route_put_strategy` pass unmodified" — cannot hold, because those tests +*are* the hardcode. `test_tier_priority_route_get` was replaced by +`test_local_preferring_route_get`, and 6 assertions across the put/get suites +were inverted; each inverted test names the expectation it replaces. Everything +else passes untouched. + +### Phase 5 — lock it in *(Rules A and C done in Phase 6; Rule B lint outstanding)* + +The original plan was one lint script enforcing three rules. Two of the three +are better enforced by the type system, so the lint shrinks to the one rule +types cannot cheaply express. + +**Rule A — only one place may name a concrete backend.** *Types — done.* +Phase 6 deleted `LocalBufferViews()`, the one call that forced `PoolClient` to +hold a `PageBackend*`, so `PoolClient::Init` now builds through +`MakePageBackend() -> unique_ptr`. The class stays in its header +because `test_page_backend` drives the allocator bookkeeping through the +`TierConfig` constructor directly; the rule as stated — one place names it — is +satisfied, and that place is the composition root. + +The same rule now holds for the transfer layer, which did not exist when this +was written: `PoolClient::Init` is also the only file naming `LocalCopyEngine`, +`MoriIoEngine` or `CompositeTransferEngine`. Reaching mori-io's peer handshake +(`EnsureRemoteEngine`, `CacheRemoteBuffers`, …) needed a `MoriIoEngine*` at +first; those calls became the `PeerDirectory` interface instead, so a second +remote transport implements an interface rather than editing `pool_client.cpp`. + +**Rule B — no file under `distributed/` may include `local/tiers/*`.** *Lint.* +The obvious structural fix — split `umbp_common` into `umbp_local` and +`umbp_distributed` so the include cannot resolve — does not work yet: +`page_backend.h` needs `HostMemAllocator`, which lives under `local/`. Cutting +that edge means relocating the allocator to `common/` first, a separate refactor. +Until then this rule stays a lint rule, scoped to `local/tiers/*` (not all of +`local/`), with `distributed/peer/ssd/` exempted as the dormant adapter — its +three `local/tiers/*` includes are the only violations in tree. + +**Rule C — no backend may call `TransferEngine::{Plan,Submit,Wait}`.** *Types — +done in Phase 6.* This rule existed only because §4 folded registration into +`TransferEngine`, handing backends an object they could transfer with. That +argument conflated *one class* with *one interface*; keeping one object and +giving it two views turns the rule into a compile error at zero runtime cost: + +```cpp +class MemoryRegistrar { /* RegisterMemory, Deregister */ }; +class TransferEngine : public MemoryRegistrar { /* + CanHandle, Plan, Submit */ }; +``` + +`MediumBackend::Init` takes a `MemoryRegistrar*`. There is no `Submit` on it to +call, and no accessor that hands back the engine. + +So what is left of Phase 5 is: cherry-pick +`tools/umbp_backend_abstraction_lint.py` and its pre-commit hook, scoped to +Rule B. **That is still outstanding** — the tool is not in tree. `MockBackend` +already exists (Phase 2, test-only since Phase 3). + +**On the budget table.** The lint carries a per-file table of tolerated +violations that "may only shrink". Prefer starting with none: the only Rule B +violations are in `distributed/peer/ssd/`, which is exempted by directory +anyway. A rule with no exceptions survives; a rule with a table of them drifts, +because the escape hatch becomes the habit. + +*Acceptance test:* adding a backend = 1 new file + 1 registry line + 0 call-site +edits. **No linter proves this — adding `SsdBackend` does**, and it is the cheap +test, since it mostly forwards to the `PeerSsdManager` that has been dormant and +under test since Phase 0. Phase 6 removed the blocker this paragraph named (the +local path is no longer host-DRAM-only, and `BufferRef` is how a medium +publishes local endpoints), so the test is now runnable; it has not been run. + +Two backend abstractions coexist by design, at different layers: +`MediumBackend` is the *distributed storage* contract (publishes descriptors, +registry-dispatched); `TierBackend` is *local mode's* storage-medium contract +(blocking read/write of bytes). The original coupling came from these two being +tangled inside `PoolClient`, not from both existing. + +### Phase 6 — abstract the transfer engine *(done)* + +`TransferEngine` and `TransferRef` landed as specified (§4). `MoriIoEngine` is +the first implementation and mori-io itself is not modified; +`CompositeTransferEngine` replaced the Phase 2 registration shim. Moved from +`PoolClient` into the engine: `GroupTransfersByPair` (now `MoriIoEngine::Plan`), +the bounce buffer and its mutex, peer engine registration and desc caching. +`RemoteDramScatterWrite` / `RemoteDramScatterRead` / `GroupPagesByBuffer` were +deleted outright — dead since the batch path landed, and a second byte-moving +path is exactly what §6 forbids. + +The local fast path **did** fold in: `LocalPutPages` / `LocalGetPages` / +`LocalCopyBlock` are gone, and `ExecuteLocalPut` / `ExecuteLocalGet` now build +`TransferItem`s against `MediumBackend::BufferRef()` and hand them to the same +planner as everything else. See "the measurement" below. + +**Four deviations from §4, each deliberate:** + +- **`TransferRef` is a struct of merged handles, not a `std::variant`.** + Registration fans *out*: the same DRAM buffer is a raw pointer (memcpy-able) + and an RDMA MR (peer-readable) *at the same time*, and which one is used is a + property of the (src, dst) PAIR. A variant forces a buffer to be one or the + other and makes the local path — memcpy between two buffers that are also + registered — inexpressible. `mori::io::MemoryDesc` already resolves this the + same way, holding `ipcHandle` and `fabricHandle` side by side. +- **No `FileRef` / `ObjectRef`, no `RegisterFile`.** No engine in tree consumes + one, and this plan has already refused an abstraction nothing exercises once + (§3, `BackendProperties`). The same test applies: `SsdBackend` is what needs a + file endpoint, so `SsdBackend` brings it — a `kind` tag plus a second handle + set in one header, with `CanHandle` already in place to route it. +- **`Wait` lives on the handle, not on the engine.** `Submit` returns a + `unique_ptr` whose destructor drains. That is not cosmetic: + the RDMA backend holds raw `TransferStatus*` into the handle's status vector, + so the drain-on-early-destroy safety net that used to live in + `~RemoteDramGetInFlight` has to sit with the statuses themselves. +- **The schedulers stayed in `PoolClient`.** §5 put submit-all-then-wait in the + base class, but the overlap that matters spans *several* `Submit` calls (one + per peer) with the local reads run in the gap — that is the client's schedule + over a batch, not one engine's over one call. Same for + `UMBP_DRAM_{WRITE,READ}_THREADS`, which §5 also listed as folding in: they + parallelize ACROSS the keys of one batch, and a batch is not a concept the + engine has. What is in the base class is `Transfer()`, the plan+submit+wait + convenience for callers with nothing to overlap. + +**Two behavior changes, both improvements, both asserted by tests:** + +| Situation | Before | After | +|---|---|---| +| Batch mixes registered and unregistered caller buffers | contract violation, batch failed | works; the engine stages the unregistered ones | +| Staged batch needs more than the bounce buffer holds | whole peer batch failed | chunked into pool-sized round trips | + +Both fall out of moving staging into the engine. The old code reserved the +whole batch's staging up front and held one mutex from submit to wait, so a +submit-all over several staging peers would deadlock — hence `permit_staging`, +the all-zero-copy-or-all-staging contract, and the two-armed fork in +`ExecuteBatch{Put,Get}Plan`. All of it is gone: a plan that needs the pool +completes *inside* `Submit`, so the lock is never held across a return. +`test_cross_node_smoke`'s `PutStagingOverflowFailsBatchCleanly` asserted the old +behavior and was replaced by `PutStagingLargerThanPoolIsChunkedNotFailed` plus +`PutPageLargerThanStagingPoolFailsBatchCleanly` — the failure that is still a +failure is a single *page* larger than the entire pool, which cannot be chunked. + +**Phase 5's two type-enforced rules closed here**, as §5 predicted: + +- *Rule C* — `MemoryRegistrar` (register/deregister) and + `TransferEngine : MemoryRegistrar` (+ `CanHandle`/`Plan`/`Submit`). + `MediumBackend::Init` takes a `MemoryRegistrar*`, so "a backend must not move + bytes" is a compile error, not a lint rule. +- *Rule A* — deleting `LocalBufferViews()` removed the last concrete-typed call, + so `PoolClient::Init` builds through `MakePageBackend() -> unique_ptr< + MediumBackend>`. `PageBackend`'s class definition stays in its header because + its unit tests drive the allocator bookkeeping directly; what matters for the + acceptance test is that the production path names no concrete backend, and it + no longer does. + +`LocalBufferViews()` was replaced on the interface by `BufferRef(buffer_index)` ++ `BufferCount()`. The difference is the point: a `TransferRef` is +medium-agnostic (§2), a raw base pointer is not. Publishing refs is what let +`PoolClient` drop `local_copy_backend_` / `local_buffers_` / `CanCopyLocally`, +and it is why a second medium's local access needs no tier branch in a copy loop. + +**Layout.** §3's three components are now three directories, so the boundary is +visible in the tree rather than only in the prose — and so the Phase 5 lint +rules can be scoped by directory, which is what Phase 0 already did for the +dormant SSD adapter: + +``` +distributed/ + master/ routing/ control plane + transfer/ transfer_engine mori_io_engine local_copy_engine composite_* + peer/ peer_service.{h,cpp} batch_resolve_codec.h (the RPC surface) + backend/ medium_backend.h page_backend mock_backend peer_page_allocator.h + ssd/ peer_ssd_manager ssd_copy_pipeline (dormant since Phase 0) +``` + +`transfer/` is a sibling of `peer/`, not a child of it, and the include graph is +why: `peer_service` — the reason `peer/` exists — references the transfer layer +**zero** times, because a peer hands out descriptors and the *initiator* moves +the bytes. `transfer/` depends on nothing but `types.h`; `backend/` and +`pool_client` depend on it. It is the lowest layer, so nesting it inside a +higher-level sibling would invert the layering. `LocalCopyEngine` settles it +from the other direction: a memcpy between two of this node's own buffers has no +peer in it at all. + +`backend/` does belong under `peer/`: `PeerServiceServer` dispatches every +Allocate/Commit/Resolve/Evict into `BackendRegistry`, and capacity, eviction and +the heartbeat outbox are all statements about what this node holds *for the +cluster*. + +Rule C is a type now, but "no file under `backend/` may name a transfer type +other than `MemoryRegistrar`/`TransferRef`" is a one-line grep against a +directory, and so is Rule B's `local/tiers/*` exemption. + +*Gate:* `bench_pool_client_batch_get` / `bench_pool_client_batch_put` unchanged; +see §9. + +Only after this phase does Phase 2 option (a), real HBM, become a clean +1-file change. + +--- + +## 6. Decisions + +**Branch from `main`, not from `refactor/umbp-storage-backend-api`.** +Most of that branch's remaining work — the `PeerTierBackend` async Submit/Wait +API and the 4-way `BatchGetPlan` partition — exists specifically to preserve +SSD's blocking data-plane shape alongside DRAM's async one. That shape is exactly +what Phase 0 removes. Cherry-pick only the lint tool. + +**Backends publish descriptors; the transfer layer decides how to move them.** +This supersedes the earlier "backends must be RDMA-addressable page slots." +The weaker requirement is the correct one: a backend's contract is to produce a +`TransferRef`, not to produce registered memory. A medium whose bytes are not +registerable produces a `FileRef` or `ObjectRef` and the transfer layer picks or +composes an engine. The invariant that must not be broken is that there is one +byte-moving path, chosen by the planner — admitting a second one at the backend +level is what produced the original coupling. + +**The transfer engine is abstracted inside UMBP; mori-io is not modified.** +UMBP owns `TransferRef`, so mori-io never needs a file or object endpoint type. +This keeps the blast radius off mori-io's other consumers. + +**Staging belongs to the transfer layer.** It is the only layer that can observe +completion. A staging pool in a storage backend must fall back to a TTL, which +is the `PrepareSsdRead` lease under another name. + +**Allocation and registration belong to the backend.** Whoever allocates memory +must supply the facts a descriptor cannot recover — hugepage/NUMA policy for +DRAM, device and location type for HBM. Note this is why the caller-buffer +registration at `pool_client.cpp:526` is a *different* bug on a different axis: +same seam, different allocator, so the fix is a location parameter on the public +API rather than a move into a backend. + +**The `local/` standalone stack is kept, not folded in.** +It has its own `TierBackend`, is reachable through `UMBPClient`'s standalone mode +and live in the Python bindings, and it is the only remaining consumer of the SSD +storage stack after Phase 0. Phase 5 fences it off with lint rather than +deleting it. + +**SSD is unwired from distributed mode, not deleted from it.** +Carrying SSD live through the refactor was estimated at 2.5 weeks for Phases 2–3 +alone, because every phase would have to design around the blocking lease shape. +Since neither the storage stack nor `PeerSsdManager` leaves the tree, and both +stay under test, the re-add is an adapter, not a rewrite. + +The one thing not preserved is the lease: `SsdReadLease`, `PrepareSsdRead`, +`ReleaseSsdLease`, and the peer-side slot state machine. + +**Re-add path (post-Phase 6).** `SsdBackend : MediumBackend` mostly *forwards* to +the still-tested `PeerSsdManager`: `BatchResolve` → `PrepareRead`, `Evict` → +`Evict`, `DrainPendingEvents` / `SnapshotOwnedKeys` → the same methods, +`Capacity` → `Capacity`. The piece the original plan deferred — owning a pool of +registered staging pages — is **no longer the backend's problem**: it publishes a +`FileRef` and the transfer layer supplies the bounce buffer if the chosen engine +needs one. It also brings back `put_eligible=false` (§3): Phase 4 deleted the +router's "SSD is not a direct-put target" rule, so SSD must re-assert it itself +or the router will place direct puts on it. Drivers, SPDK env, allocator, proxy protocol, segment format, +eviction policy, the peer-side ownership map and the copy-on-commit pipeline are +reused as-is. + +Concretely, `SsdBackend` is now three things, none of which touch a call site: +add a `kind` tag + file handle set to `TransferRef`; implement +`BufferRef`/`BufferCount` over its extents; register it in `PoolClient::Init`. +The engine that reads those extents (`PosixEngine` / `GdsEngine`) is a fourth +file, and `CompositeTransferEngine::SelectEngine` routes to it with no edit. + +**One place still names a concrete type, deliberately: `PoolClient::Init`.** +It constructs `LocalCopyEngine`, `MoriIoEngine`, `CompositeTransferEngine` and +(via the `Make*Backend` factories) the node's one medium backend. That is a +composition root, not a leak — the test the acceptance criterion actually states +is that *no other* file names one, and none does. Everything downstream holds +`MediumBackend`, `TransferEngine`, `MemoryRegistrar`, or `PeerDirectory`. + +**One medium per node — `UMBPDistributedConfig::medium` selects it.** +Registering DRAM unconditionally and adding HBM/SSD beside it (as Phase 2b–7 +did) reads like a tier stack and is not one: since Phase 4 the master treats +every advertised tier as an equally valid put target, so the second medium is +*mirrored into by free capacity*, not demoted to. That costs capacity, splits a +node's keys across two pools for no gain, and makes a "--tier ssd" measurement +actually measure DRAM. Local tiering would be a routing-plane feature (a +per-node tier order plus promote/demote), and nothing has asked for one: +heterogeneity comes from different **nodes** picking different media, which the +routing plane already handles. So the selector is a single enum, adding a medium +is a `case` rather than an `if`, and the two data-plane leftovers §1 listed — +the `TierType::DRAM` re-cache literal and the `!= DRAM && != HBM` put filter +(which silently dropped every put master routed to a peer's SSD) — are gone +rather than generalized to a list. + +--- + +## 7. Effort + +| Phase | Work | +|---|---| +| 0 — unwire SSD from distributed | 1–1.5 days *(done)* | +| 1 — interfaces | 0.5 day | +| 2 — DRAM backend **+ ownership move** | 4–5 days | +| 3 — peer + pool client | 2–3 days | +| 4 — routing | 1 day | +| 5 — lint + cleanup | 1 day | +| 6 — transfer engine | 4–6 days | + +Phase 2 grew from the original 2–3 days: splitting the tier map is the small +half, and relocating allocation, registration and capacity reaches up into +`PoolClientConfig`, `DistributedClient` and `UMBPConfig`. + +Phase 2 assumes option (b) — `MockBackend` as the second registered backend. +Real HBM is a feature on top, and is cleanest after Phase 6. + +Not counted: re-adding `SsdBackend`, ~1–1.5 days now that the staging pool moved +to the transfer layer. It is the real acceptance test for Phase 5's +*1 file + 1 registry line* claim. + +Keeping SSD alive carries a standing cost worth naming: the SPDK submodule, +`storage/**`, the local SSD suites, and the dormant adapter plus its 3 tests all +stay in build and CI while nothing in distributed mode calls them. Dormant +compiled code also drifts — if `KvEvent` or `types.h` changes under it, someone +must fix a file no one is using. Those 3 tests are what turns "preserved source" +into "preserved working behavior." + +--- + +## 8. Open — shared media and node-scoped routing + +Not solved by anything above, and worth knowing before S3 or 3FS is attempted. + +The whole routing plane assumes a key lives on a *machine*: + +```cpp +struct Location { std::string node_id; uint64_t size; TierType tier; }; +``` + +True for DRAM, HBM and local SSD. Meaningless for S3 or 3FS, where every node +can reach the same objects. Three concrete failures: + +- The master records whichever node happened to write an object and routes all + readers there, creating a hot spot for data any node could have fetched + directly. +- Per-node `Capacity()` is wrong: the capacity is shared, so N nodes each + reporting the whole bucket lets the master over-commit it N times. +- Eviction becomes a coordination problem — N nodes believe they own the same + object. + +Fixing this needs `Location` to express "reachable from anywhere" plus a +shared-capacity rule in the master. It is control-plane work, orthogonal to +`MediumBackend` and to §4, and it does not block any phase above. diff --git a/src/umbp/doc/design-master-control-plane.md b/src/umbp/doc/design-master-control-plane.md index 288ce565b..3b88096d1 100644 --- a/src/umbp/doc/design-master-control-plane.md +++ b/src/umbp/doc/design-master-control-plane.md @@ -173,8 +173,7 @@ include/umbp/ ├── distributed/ │ ├── config.h # MasterServerConfig, PoolClientConfig, ClientRegistryConfig │ ├── distributed_client.h # DistributedClient (IUMBPClient impl) -│ ├── pool_client.h # PoolClient (master + peer + IO engine glue) -│ ├── pool_allocator.h +│ ├── pool_client.h # PoolClient (master + peer + transfer engine glue) │ ├── obs_counters.h # MORI_UMBP_OBS_* / test-seam build switch │ ├── types.h # TierType, Location, KvEvent, ClientRecord, ... │ ├── master/ @@ -185,12 +184,23 @@ include/umbp/ │ │ ├── external_kv_block_index.h │ │ ├── eviction_manager.h │ │ └── master_metrics.h # MORI_UMBP_METRIC_* name/help strings +│ ├── transfer/ # the one byte-moving path; a sibling of +│ │ │ # peer/ because peer_service never uses it +│ │ ├── transfer_engine.h # TransferRef, MemoryRegistrar, PeerDirectory +│ │ ├── mori_io_engine.h # RDMA, via mori-io +│ │ ├── local_copy_engine.h # both endpoints local (no peer involved) +│ │ └── composite_transfer_engine.h │ ├── peer/ -│ │ ├── peer_service.h # UMBPPeer gRPC server -│ │ ├── peer_dram_allocator.h # canonical per-node DRAM/HBM owner -│ │ ├── peer_ssd_manager.h # canonical per-node SSD owner -│ │ ├── ssd_copy_pipeline.h # async copy-on-commit DRAM→SSD -│ │ └── peer_page_allocator.h # PageBitmapAllocator +│ │ ├── peer_service.h # UMBPPeer gRPC server (dispatches via BackendRegistry) +│ │ ├── batch_resolve_codec.h # BatchResolveKeys wire codec +│ │ ├── backend/ # storage media: own bytes, publish endpoints +│ │ │ ├── medium_backend.h # MediumBackend contract + BackendRegistry +│ │ │ ├── page_backend.h # canonical per-node DRAM/HBM owner +│ │ │ ├── mock_backend.h # second registered medium, test-only +│ │ │ └── peer_page_allocator.h # PageBitmapAllocator +│ │ └── ssd/ # dormant since Phase 0 of the backend refactor +│ │ ├── peer_ssd_manager.h # canonical per-node SSD owner +│ │ └── ssd_copy_pipeline.h # async copy-on-commit DRAM→SSD │ └── routing/ │ ├── router.h │ ├── route_get_strategy.h # TierPriorityRouteGetStrategy (default), RandomRouteGetStrategy @@ -972,25 +982,40 @@ Built-in metric families include: `mori_umbp_heartbeat_seq_gap_total` master counters, - `mori_umbp_external_kv_*` counters/gauges for the external-KV index. -SSD-tier families are reported by the owner peer (`node=`) via -`ReportMetrics` — names per `master_metrics.h`: - -- copy-on-commit: `mori_umbp_ssd_copy_enqueued_total`, - `mori_umbp_ssd_copy_succeeded_total`, - `mori_umbp_ssd_copy_failed_total`, - `mori_umbp_ssd_copy_dropped_total` (`reason=queue_full|stopped`), - and `mori_umbp_ssd_copy_bytes_total`, -- reads: `mori_umbp_ssd_read_total` - (`status=ok|not_found|no_slot|size_too_large|error`), - `mori_umbp_ssd_read_bytes_total`, and the reader-side - `mori_umbp_ssd_read_client_transient_total`, -- peer-local eviction: `mori_umbp_ssd_eviction_rounds_total`, - `mori_umbp_ssd_eviction_victims_total`, - `mori_umbp_ssd_eviction_bytes_freed_total`, - `mori_umbp_ssd_eviction_backend_failed_total`, -- read staging: `mori_umbp_ssd_staging_slots_in_use`, - `mori_umbp_ssd_staging_expired_reclaims_total`, - `mori_umbp_ssd_staging_slot_full_rejects_total`. +Peer-side storage and transfer families are reported by the owning peer +(`node=`) via `ReportMetrics`. They are named by the COMPONENT, not by +the medium — see `umbp/distributed/metrics/component_metrics.h`: + +- storage backends, emitted for every medium by the instrumentation + decorator that wraps each `MediumBackend` + (`instrumented_backend.h`), labeled `tier=, + backend=`: + `mori_umbp_backend_ops_total` (`op=allocate|commit|abort|resolve|evict`, + `status=ok|exists|no_space|miss|failed`), + `mori_umbp_backend_batches_total`, `mori_umbp_backend_bytes_total` + (`op=commit|resolve|evict`) and `mori_umbp_backend_op_seconds_total`, +- medium-internal detail a backend publishes for itself, under generic + names with the specifics in a label: + `mori_umbp_backend_medium_events_total{event=...}`, + `mori_umbp_backend_medium_bytes_total{event=...}` and the + `mori_umbp_backend_medium_state{state=...}` gauge. SSD uses these for + drive read outcomes, single-flight coalescing, eviction rounds and + staging-arena pressure, +- the transfer layer, emitted by `CompositeTransferEngine` for whichever + engine carried each plan, labeled `engine=, + direction=push|pull|local`: `mori_umbp_transfer_ops_total` + (`status=ok|failed`, plus `engine="none", status="rejected"` for items + no engine would take), `mori_umbp_transfer_bytes_total` and + `mori_umbp_transfer_seconds_total`. + +**A new backend or transfer engine emits all of the above with no metrics +code of its own** — it is instrumented by being composed in, which is why +`examples/monitoring/grafana/dashboards/umbp_backends.json` is a single +dashboard whose panels group by `tier` / `engine` rather than one dashboard +per medium. A component publishes for itself only what its interface cannot +show from outside, and it names an *event*, never a metric. The dashboards +beside it — client data rate, master RPC latency, RPC call rates, external +KV — are backend-agnostic and stayed as they were. Set `MORI_UMBP_LOG_LEVEL=DEBUG` (or `UMBP_LOG_LEVEL=0`) to surface per-RPC traces from `MORI_UMBP_INFO` / `MORI_UMBP_DEBUG`. diff --git a/src/umbp/doc/design-ssd-ranged-io.md b/src/umbp/doc/design-ssd-ranged-io.md new file mode 100644 index 000000000..3f384979b --- /dev/null +++ b/src/umbp/doc/design-ssd-ranged-io.md @@ -0,0 +1,557 @@ +# Ranged I/O on the SSD Medium — Standalone UMBP Server and Distributed Master + +Ranged (sub-object) I/O is unavailable on any node whose medium is SSD. One +expression turns it off: + +```cpp +// distributed_client.cpp:388 +return config_.distributed.has_value() && config_.distributed->medium != UMBPMedium::SSD; +``` + +`pure-ssd-mode.md` describes the deployment this blocks: **every** node sets +`medium: SSD`, so on that deployment `SupportsRangedIO()` is false cluster-wide, +and the sglang tree connector falls back to whole-object I/O for every request. + +**Status: D0 landed** on `feat/umbp-ssd-ranged-io`. §3 is implemented and +tested; D1, D1c, X1 and D2 remain open. §3.1 records one thing the +implementation corrected about the shape of the fix. + +This document covers the distributed data plane only — a client talking to a +master, whether that client is a `DistributedClient` in-process or a standalone +UMBP server forwarding for one. Written against `refactor/umbp-backend-agnostic` +at `84adf761`. + +**Out of scope: local mode.** `StandaloneClient`, `LocalStorageManager` and the +DRAM tier are not touched. Where this design reaches `SSDTier` / +`ShardedSsdTier` it does so through `PeerSsdManager`, which deliberately owns a +bare `TierBackend` and must not pull in `LocalStorageManager` or +`LocalBlockIndex` (`peer_ssd_manager.h:55`); that is a peer-side dependency of +the distributed path, not local mode leaking in. See §9 for the rest of the +non-goals. + +Companion reading: [`design-backend-agnostic-refactor.md`](design-backend-agnostic-refactor.md) +for the `MediumBackend` split, [`design-tree-connector-port.md`](design-tree-connector-port.md) +§10.3/§12 for how ranged I/O entered the distributed data plane, and +[`pure-ssd-mode.md`](pure-ssd-mode.md) for the deployment this work serves. + +--- + +## 1. Why SSD is different, and where it is not + +`SsdBackend`'s founding decision is that **the staging arena is the medium** +(`ssd_backend.h`, class comment). SSD bytes are not addressable by the transfer +layer, so a Resolve reads the object off disk into a registered host DRAM page +and publishes *that* page. Everything downstream — +`PoolClient::BuildLocalRangeTransfers`, the RDMA path, `HbmCopyEngine` — sees +ordinary host memory and needs no SSD awareness at all. + +The consequence for ranged I/O is the crux of this document: + +> Once the object is staged, ranged access **already works** — the ranges are +> just offsets into a published page. What does *not* work is the saving. The +> disk read happened at Resolve time, at full object size, before anyone said +> which bytes they wanted. + +So the gate at `distributed_client.cpp:388` is not protecting a correctness +hole. Lifting it yields a *correct* ranged path immediately, with **zero +disk-side savings**. That is exactly the follow-up +`design-tree-connector-port.md` §8 item 4 predicted, and it is why this work +splits into "make it work" (§3) and "make it worth it" (§4, §5). + +--- + +## 2. Invariants this design must not break + +1. **A backend does not move bytes.** `MediumBackend::Init` receives a + `MemoryRegistrar`, not memory, and that is a compile-time fence + (`medium_backend.h`; §5 Rule C of the refactor doc). Nothing here adds a copy + method to a backend. +2. **Pages tile the object, in order.** `PageLocation` carries + `(buffer_index, page_index)` and nothing else (`types.h:149`); + `BuildLocalRangeTransfers` maps an object offset to a page by division + (`pool_client.cpp:851`). Any "publish only part of an object" scheme has to + say what it does about this — §4.1. +3. **A ranged put tiles its object.** `RangesTileObject` is checked before + routing (`pool_client.cpp:1897`). Ranged puts are a *scatter-gather write of + a whole new object*, never a partial update of an existing one. +4. **A ranged get is a subset, not a tiling.** `RangesAreDisjointAndInBounds` is + deliberately weaker than `RangesTileObject` (`pool_client.cpp:170`). +5. **Unsupported is "the wrong door", not a miss.** The pattern `TierBackend` + established for `ranged_read` (`tier_backend.h:39-43`) — an all-false default + is indistinguishable from "every key missed", so capability is advertised + separately — applies to every capability added below. +6. **Per-key results.** One bad key fails that key, never the batch; shape + errors fail the whole call. +7. **No advertised tier order.** `medium_backend.h` rejected `BackendProperties` + and read priority outright. Nothing below reintroduces one; §6's `Contains` + is a side-effect-free existence predicate, which is a different thing. + +--- + +## 3. D0 — lift the gate, and be honest about what it buys + +Trace what actually happens with the medium test removed. + +**Ranged get, key held on this node.** `PoolClient::BatchGetRanges` Phase 1 +resolves the key against each registered backend (`pool_client.cpp:1715`). +`SsdBackend::BatchResolve` stages the whole object into a DRAM page and +publishes it under a read lease (`ssd_backend.cpp:380`). `BuildLocalRangeTransfers` +copies only the requested ranges out of that page. **Correct today.** + +**Ranged get, key held remotely.** Phase 2 fetches the whole object into the +ranged scratch arena via `ExecuteBatchGetPlan` and copies the ranges out +(`pool_client.cpp:1758`). The owning peer's `SsdBackend` stages it; the wire +carries the whole object. **Correct today.** + +**Ranged put, routed to this node.** `ExecuteLocalPutRangesBatch` allocates a +slot — for `SsdBackend`, a staging page (`ssd_backend.cpp:243`) — writes the +ranges straight into it, and commits, which spills the page to disk +(`ssd_backend.cpp:295`). Because a ranged put *tiles* its object (invariant 3), +every byte of the page is written before the spill. **Correct today, and with no +penalty**: the disk sees exactly the one whole-object write it would have seen +anyway. + +**Ranged put, routed remotely.** Assembled in the scratch arena and sent as one +contiguous object. **Correct today.** + +So D0 is: delete the medium test and let `SupportsRangedIO()` mean what its name +says — the arena check above it already carries the opt-in. + +```diff +- return config_.distributed.has_value() && config_.distributed->medium != UMBPMedium::SSD; ++ return true; // ranged_scratch_size_ != 0 was already checked above +``` + +What D0 buys, and what it does not: + +| | before D0 | after D0 | +|---|---|---| +| ranged put to an SSD node | unavailable | full speed — one whole-object disk write, as before | +| ranged get, disk bytes | n/a | still the whole object | +| ranged get, wire bytes | n/a | still the whole object | +| ranged get, host↔device copy bytes | n/a | **only the requested ranges** | + +### 3.1 The flag is a declaration, not a guard + +Worth stating plainly because it changes how the rest of this document reads, +and the implementation is what made it obvious: +`DistributedClient::BatchGetRanges` and `BatchPutRanges` forward to `PoolClient` +**unconditionally** — neither consults `SupportsRangedIO()`. The flag is +something callers *ask*, and the sglang tree connector is the caller that asks. + +So D0 does not unblock a code path. It stops the client misreporting a +capability it always had. Two consequences: + +- **The characterization tests had to be written to the conjunction.** A test + that only drove the data path would pass against the pre-D0 gate and pin + nothing, because the gate never stood in the data path's way. Every case in + §10 therefore opens by asserting `SupportsRangedIO()`, so what is pinned is + "the client says yes, *and* the bytes are right" — and all five fail against + the old gate. +- **Anything reached by bypassing the flag was already live.** A caller that + never asked has been issuing correct ranged I/O against SSD nodes all along. + That is not an argument for having left the gate — the connector *did* ask, + which is the whole reason the deployment lost the feature. + +### 3.2 What D0 buys + +The put half is a genuine, complete win. The get half is a correctness +enablement whose only saving is the final copy — worth having on a GPU +destination, where `HbmCopyEngine`'s gather is the expensive part, and worth +nothing on the drive. **Do not ship D0 announced as "SSD ranged reads work" +without D1**; that is how a feature gets benchmarked, found flat, and blamed. + +--- + +## 4. D1 — extent-scoped staging + +D1 is where the disk saving comes from: read only the requested extent off disk, +stage only that, publish only that. + +### 4.1 The tiling assumption, and `ResolvedEntry::covered_offset` + +The obstruction is invariant 2. A published page set is assumed to tile +`[0, object_size)` in order, and `BuildLocalRangeTransfers` finds a byte by +`object_offset / page_size`. A staging page holding object bytes `[a, b)` breaks +that assumption **silently** — every offset would be wrong, and the result is +plausible bytes at the wrong place, the worst failure shape available. + +Three ways to close it: + +**(a) Add an offset to `PageLocation`.** One `uint64_t object_offset` field, and +`BuildLocalRangeTransfers` intersects ranges against page extents instead of +dividing. *Rejected*: `PageLocation` is on the wire in `AllocateSlotResponse` / +`ResolveKeyResponse` and their batch forms, and in the master's metadata. It is +the most widely shared value type in the system, and every DRAM and HBM path +would carry a field only SSD ever sets. + +**(b) A per-resolve covered extent on `ResolvedEntry`.** Add +`uint64_t covered_offset = 0` alongside the existing `pages` / `size` / +`page_size`. The pages still tile — they tile +`[covered_offset, covered_offset + covered_size)` rather than the whole object. +`size` continues to mean the *object's* size, so range validation is unchanged. +`BuildLocalRangeTransfers` subtracts `covered_offset` from each range's +`object_offset` before its page arithmetic, and rejects any range not contained +in the covered window. **Chosen.** It is local to the resolve result (21 call +sites, all in-process at D1), DRAM and HBM leave it at zero and stay +byte-identical, and it states the true thing: a resolve answers *for an extent*, +and whole-object is the degenerate case where the extent is everything. + +**(c) Resolve-with-ranges as a distinct `MediumBackend` entry point.** +*Rejected*: it forks the backend interface for one medium, which is exactly the +per-medium branching the refactor removed. + +### 4.2 The change, end to end + +1. `MediumBackend::BatchResolve` gains a requested-extent argument — one + `std::optional` per key, defaulted to "everything". DRAM and + HBM ignore it: their pages already hold the object, they answer + `covered_offset = 0`, and they pay nothing. +2. `PeerSsdManager` gains `PrepareReadRangesBatch`, delegating to + `TierBackend::ReadBatchRangesIntoPtr` — see §4.3 for what that requires. +3. `SsdBackend::BatchResolve` computes the **enclosing** extent of the requested + ranges (one span, `[min offset, max end)`), reads that, stages it, and reports + `covered_offset`. One span rather than the exact set because a staging page is + one contiguous allocation; the ranges within it are the caller's problem + again, exactly as they are for DRAM. +4. `PoolClient::BatchGetRanges` Phase 1 passes the enclosing extent down; + `BuildLocalRangeTransfers` honours `covered_offset`. + +The one-key-one-page limit (`ssd_backend.cpp:258`, `:420`) is *relaxed* by this, +not tightened: an extent is by construction no larger than its object, so a key +whose object exceeds `page_size` becomes servable ranged even though it can never +be served whole. Worth a test; not a goal — in distributed mode master's +`page_size` is the KV block size, so the case is rare. + +**Staging-arena pressure improves too.** A resolve currently borrows a full +`page_size` page regardless of how much of it the reader wants; with D1 occupancy +tracks demand. The arena remains sized for read *concurrency* (`ssd_backend.h`), +so this does not change the sizing rule, only the headroom. + +### 4.3 The peer-side SSD tier dependency — multi-drive + +`PeerSsdManager` constructs its `TierBackend` through `CreateSsdBackend` +(`peer_ssd_manager.cpp:35`): one storage dir gives an `SSDTier`, several +comma-separated dirs give a `ShardedSsdTier` (`ssd_backend_factory.h:33`). +Multi-drive is the production topology — it is the whole point of +`pure-ssd-mode.md` — so D1's `ReadBatchRangesIntoPtr` call lands on +`ShardedSsdTier` in every real deployment. + +- `SSDTier::ReadBatchRangesIntoPtr` **exists** (`ssd_tier.cpp:710`, from + `abf38df7`) and is the reference implementation: it resolves keys under `mu_`, + plans outside it, merges object-contiguous ranges into runs, issues them as one + `io_driver_->ReadBatch` with per-run fallback, and reads straight into the + caller's buffer when a single host range is O_DIRECT-compatible. +- `ShardedSsdTier` **does not forward it**, and `Capabilities()` never sets + `ranged_read` (`sharded_ssd_tier.cpp:237`). Both are omissions, not decisions. + +So D1 requires, on `ShardedSsdTier`: + +- Override `ReadBatchRangesIntoPtr`: bucket keys by `ShardOf(key)`, issue **one + call per shard**, scatter results back — the same shard-and-regroup the class + already does for `ReadBatchIntoPtr`. +- Fold `ranged_read` into the existing capability intersection: ranged-capable + only if every shard is. + +Do not fall back to a per-key loop. One call per shard is what lets each drive +run its own `io_uring` batch concurrently, which is the aggregate-bandwidth +property the mode exists for. + +This is scoped strictly to the peer-side read path. `ShardedSsdTier` is reached +here as `PeerSsdManager`'s bare `TierBackend`; nothing about `LocalStorageManager` +routing, DRAM tiers or `StandaloneClient` is in scope. + +### 4.4 Single-flight coalescing must key on the extent, not the key + +This is the one genuine correctness hazard D1 introduces, and it is easy to miss. + +`PeerSsdManager` coalesces concurrent same-key reads when +`UMBPSsdConfig::single_flight_reads` is on — the default +(`config.h:116`, `UMBP_SSD_SINGLE_FLIGHT`). One requester becomes the leader, +others attach to its `ReadEpisode` as `FollowerTarget{dst, cap}`, and the leader +copies its staging buffer into every follower's destination +(`peer_ssd_manager.h:256-286`). + +Today that is sound because every read of a key is a whole-object read, so the +leader's bytes are exactly what each follower wanted. Under D1 they are not: two +requesters can want **different extents** of the same key, and a follower +attached to a leader reading `[0, 4096)` would silently receive those bytes as if +they were its own `[8192, 12288)`. + +Three admissible answers; pick per what the workload actually does: + +- **Attach only on an identical extent.** Key the episode on + `(key, covered_offset, covered_size)`. Simplest, and exactly right for the + workload that motivates single-flight in the first place — MLA keys identical + across TP ranks, where every rank asks for the same thing. +- **Attach on containment.** A follower whose extent is inside the leader's takes + a sub-slice of the leader's buffer. Strictly better coverage, one offset + subtraction more code, and it needs `FollowerTarget` to carry the follower's own + offset. +- **Do not coalesce ranged reads.** Correct and trivial; gives up the `read_dup` + saving on the exact access pattern that has the most duplication. + +**Containment is the recommendation**, with identical-extent as the fallback if +the episode bookkeeping proves awkward. Whichever is chosen, the wrong answer — +attaching on key alone — produces correct-looking bytes and no error, so this +needs a test that asserts *content*, not success. + +--- + +## 5. D2 — ranged resolve on the peer wire + +After D1 a locally-held key reads only its extent off disk. A remotely-held one +still crosses the wire whole, because Phase 2 fetches the object into the scratch +arena and slices it client-side (`pool_client.cpp:1758`). + +Closing that means the requested extent must travel to the peer, and +`umbp_peer.proto` has no notion of ranges today. Concretely: `ResolveKeyRequest` +and `BatchResolveKeysRequest` gain an optional extent per key; +`ResolveKeyResponse` and `BatchResolveKeysResponse` gain the matching +`covered_offset`. `peer_service.cpp` passes both through to `BatchResolve`, which +D1 already taught to take them. + +The client then also stops needing the scratch arena for a remote ranged *get*: +with the peer publishing an extent-scoped page, the reader can RDMA directly into +its own scattered destinations, which is what `BuildLocalRangeTransfers` does for +the local case. The arena stays required for remote ranged *put* — the object +must be contiguous on the wire. + +D2 is deliberately last. It is the only piece with a wire-compatibility story: an +old peer ignores the optional field and answers whole-object, which a new client +must detect via `covered_offset == 0 && covered_size == size` and handle by +slicing — i.e. exactly D0's behaviour, so the fallback is code that already +exists. It is also the only piece that helps DRAM and HBM as much as SSD, so it +can be justified on its own and does not need to wait for an SSD argument. + +--- + +## 6. X1 — resolve-as-probe, which D0 exposes + +`PoolClient::BatchGetRanges` Phase 1 finds a key's holder by calling +`BatchResolve` on **every registered backend until one answers** +(`pool_client.cpp:1715`). For DRAM that is a map lookup. For `SsdBackend` a +resolve is a *disk read, a staging-page allocation and a read lease* — so using +it as a probe is expensive, and on a node with more than one medium a probe that +reaches SSD first pays all of that for a key another backend may also hold. + +Today a node registers exactly one medium (`PoolClientConfig::medium`), so the +loop degenerates and this is latent. D0 is what makes the SSD case reachable at +all. The fix: + +> Give `MediumBackend` a `Contains(keys)` predicate that every backend answers +> without side effects. `SsdBackend` forwards it to `PeerSsdManager::SizeOf` +> (`peer_ssd_manager.h:83`), already an index lookup. Phase 1 probes with +> `Contains`, then resolves once against the winner. + +Ordering the registry cheapest-medium-first would merely make the bad case rarer, +and would edge toward the advertised tier order invariant 7 forbids. `Contains` +is three lines per backend and is the honest fix. + +Whole-object `BatchGet` has the same structure and gets the same benefit; this is +a shared fix that ranged work happens to force. + +--- + +## 7. The standalone UMBP server + +The server has no ranged logic of its own — it forwards to the client it was +configured with, so it inherits D0–D2 whole. Three things must be checked rather +than assumed: + +- **Capability propagation.** `supports_ranged_io` rides on the `Ping` response + (`umbp_standalone.proto:38`, `standalone_server.cpp:265`) and is cached in + `StandaloneProcessClient` (`standalone_process_client.h:80`). It is read from + the *server's* client, so a server over a distributed backend on an SSD-medium + node starts reporting true at D0 with no code change. The flag is cached at + connect time; a client that connected before the server's capability changed + keeps the stale answer for its lifetime. Pre-existing and acceptable + (capabilities do not change under a running server), recorded so it is not + later mistaken for a bug. +- **The opt-in is an env var here.** The server parses + `UMBP_DISTRIBUTED_RANGED_SCRATCH_BYTES` into `dist.ranged_scratch_size` + (`standalone_server_main.cpp:160`), which defaults to zero and therefore + defaults ranged I/O off. D0 does not change that — a pure-SSD deployment must + set it, and `pure-ssd-mode.md`'s env list should gain it when D0 lands. +- **The wire format needs nothing.** `BatchRangeDataRequest` flattens ranges + behind a `range_counts` prefix (`umbp_standalone.proto:92`) and is + medium-agnostic, through D2 included. + +--- + +## 8. Phasing + +| Phase | What | Depends on | Value | +|---|---|---|---| +| **D0** | lift the medium gate (§3) — **LANDED** | — | **ranged put on SSD nodes, at full speed**; ranged get correct, no disk saving | +| **D1a** | `ShardedSsdTier` forwards `ReadBatchRangesIntoPtr`, capability intersected (§4.3) | — | prerequisite; no user-visible change alone | +| **D1b** | extent-scoped staging via `ResolvedEntry::covered_offset` (§4.1, §4.2) | D0, D1a | **ranged get reads only its extent off disk** | +| **D1c** | extent-aware single-flight coalescing (§4.4) | D1b | correctness — must land with D1b, not after | +| **X1** | `MediumBackend::Contains` probe (§6) | — | removes a resolve-as-probe cost D0 exposes | +| **D2** | ranged resolve on the peer wire (§5) | D1b | remote ranged get stops shipping whole objects | + +Recommended order: **D1a → D0 → D1b+D1c → X1 → D2.** + +D0 went first in practice, because it is two lines and unblocks the deployment +immediately; D1a is inert on its own and unblocks the only phase that delivers +a disk-side saving, so it is next. D1b and D1c ship together — D1b without D1c +is a silent-wrong-bytes bug on the default configuration. + +**Landed for D0** (`feat/umbp-ssd-ranged-io`): +`distributed_client.cpp` `SupportsRangedIO()` reduced to the arena check; +`distributed_client.h`'s comment corrected (it justified the gate by saying +`SsdBackend` publishes storage refs — it does not, and never did: `ssd_backend.h` +weighs a file endpoint against staging and picks staging, so the gate rested on +a backend nobody built); `tests/cpp/umbp/distributed/test_ssd_ranged_io.cpp` +added and registered as the `umbp_ssd_ranged_io` integration case. + +--- + +## 9. Non-goals + +1. **Local mode.** `StandaloneClient`, `LocalStorageManager` ranged-write + routing, the DRAM tier and the `UMBPRole::SharedSSDLeader/Follower` shared-SSD + roles are all out of scope. Ranged writes there remain DRAM-only + (`local_storage_manager.cpp:721`, `:770`) and the follower role remains + excluded (`standalone_client.h:78`). §4.3 touches `ShardedSsdTier` only as + `PeerSsdManager`'s `TierBackend`, and adds nothing local mode consumes. +2. **Read-modify-write.** A ranged put tiles its object (invariant 3). Updating + part of a stored object in place is not added; the segmented-log layout is + append-only, so a partial update means a new record plus a rewrite of the + rest — strictly worse than a whole-object put. +3. **Per-block checksums / record format v4.** A record CRC covers the whole + value (`segment_format.h:91`), so a partial read has nothing to check itself + against, and `SSDTier`'s ranged read path verifies nothing. This design + inherits that narrowing rather than fixing it. Note it costs nothing on the + deployment in question: `pure-ssd-mode.md` recommends + `UMBP_SSD_VERIFY_CRC=0`, so nothing was verifying anyway. Per-block checksums + are the right answer only if ranged reads ever become the default access + shape rather than an opt-in. +4. **SPDK backends.** `SpdkSsdTier` and `SpdkProxyTier` get no ranged support. + `SpdkProxyTier` moves bytes through a shared-memory ring with its own + caching, so a "range" is not a device offset there at all; it would mean a + `spdk_proxy_protocol.h` change. Both correctly inherit the all-false default + and, under §4.3's intersection, correctly advertise no ranged capability — + so they are routed around rather than reporting phantom misses. +5. **A `FileRef` transfer endpoint / GDS path.** `ssd_backend.h` argues at length + for staging over a file endpoint, and ranged I/O does not weaken that + argument — if anything D1 strengthens it, since an extent-scoped stage is a + smaller bounce than a whole-object one. A `FileRef` backend remains the right + answer for zero-copy and remains a strictly larger change. +6. **Collapsing whole-object I/O into the degenerate ranged case.** Still open + from `design-tree-connector-port.md` §8 item 2; still not this work. +7. **Making ranged I/O the default.** It stays opt-in behind the scratch arena. + Below some fetched fraction it wins and above it it loses; nothing here + changes that trade-off, only where it applies. + +--- + +## 10. Verification + +**Unit.** `tests/cpp/umbp/distributed/test_pool_client_ranges.cpp` is the +fixture — it drives `BackendRegistry` / `MediumBackend` directly, so it can host +an `SsdBackend` without a master. + +- **D0** — *written, in `test_ssd_ranged_io.cpp`; 5 cases, all passing.* + `SupportsRangedIO()` true for an SSD medium with an arena and false without; + local ranged put assembling an object from out-of-order tiling ranges + (verified by a whole-object read-back, since a ranged read-back would agree + with a consistently wrong mapping); local ranged get of a non-tiling subset; + a guard-byte case proving a ranged get writes only what was asked for; and a + remote round trip through an SSD-only peer. Each drives a real + `IUMBPClient` and asserts the advertised capability first — see §3.1 for why + that assertion is load-bearing rather than decorative. All five fail against + the pre-D0 gate. + + A sixth case covers §7's hop, which the other five do not touch: a standalone + server holding the `DistributedClient` while a worker forwards to it over + gRPC, asserting `GetBackendMode() == Distributed`, the capability surviving + the `Ping`, and a ranged round trip. Writing it surfaced a trap worth + repeating: the standalone-process data plane is **shared memory**, so caller + buffers must be shm-backed and registered. An ordinary heap pointer fails + `ResolveRanges` server-side and returns a per-key `false` — correct, but at + the call site indistinguishable from the store refusing the write. +- **D1a**: a new `test_sharded_ssd_tier_ranges.cpp` asserting keys land on the + shard `ShardOf` names, that a batch spanning shards issues **one call per + shard** (counted via a mock `TierBackend`), and that `ranged_read` is the AND + over shards — including a mixed fixture where one shard is not ranged-capable. +- **D1b**: assert the **disk read size**, not just the bytes returned. An + instrumented `TierBackend` counting bytes read is what distinguishes D1b from + D0; without it the D1b test passes on D0's code. Also: a range outside the + covered window is rejected; DRAM and HBM report `covered_offset == 0` and are + byte-identical to before; an object larger than `page_size` is servable ranged. +- **D1c**: two concurrent readers of the same key requesting **disjoint** + extents, asserting each receives its own bytes. This test must compare + content — the bug it guards against returns success with wrong data. Run it + with `UMBP_SSD_SINGLE_FLIGHT=1` explicitly, since that is the default the bug + lives on. +- **X1**: a two-backend registry asserting `SsdBackend` performs no disk read for + a key it does not hold. + +**Measured, 2026-08-19, n06-21** (`umbp_bench.py ranged_perf`, one node, +`umbp_master` + two standalone servers, object 32x36KiB, batch 32, median of +3-7). Same drive, same object, same access shape; only the read path differs: + +| deployment | G=1 (3.1% fetched) | G=32 (100%) | +|---|---|---| +| `umbp-local`, 1 drive (true partial pread) | **6.96-9.83x faster** | 0.63-0.92x | +| `umbp-server` over RDMA (stages whole object) | 0.23-0.30x | 0.27x | +| `umbp-local`, 3 drives sharded | **every ranged read misses** | — | + +Three things this settles: + +1. **D1 is worth about an order of magnitude at low fetched fractions, and it is + not a projection.** Local mode already reads only the requested extent off + the same drive and is ~7-10x faster there; distributed cannot, because the + resolve stages before the ranges are known. The gap between the two rows IS + the D1 opportunity, measured. +2. **The distributed cost is per key and does not amortize.** A batch sweep at + G=1 gives 0.71 / 0.70 / 0.69 ms per key at batch 8 / 32 / 128, while + whole-object improves with batch (0.31 -> 0.215 -> 0.018 ms/key). Something + in the ranged path is serialized per key; §6's key-by-key Phase 1 probe is + the prime suspect, not yet profiled. That makes X1 a prerequisite for D1 + paying off, not an optional cleanup. +3. **D1a is not merely a missing optimisation — multi-drive is broken today.** + `--ssd-dir a,b,c` misses every ranged read where one drive succeeds on the + identical config, and `supports_ranged_io()` still returns true on that node + (`StandaloneClient` only checks the role), so the client advertises a + capability the sharded tier cannot serve. Fixing the forwarding must also fix + the advertisement. + +Caveat on the distributed whole-object baseline: it is partly lease-warmed. Read +leases keep staged pages in DRAM across passes, and whole-object at batch=128 +reported ~62 GB/s, impossible from one NVMe. Ranged's flatness and its per-key +constant are unaffected. The node was also heavily loaded (load ~43), which is +why the ratio, not the absolute, is the reported figure. + +**Performance.** Two things to hold onto when reading numbers: + +- The existing finding is the baseline to beat: ranged **loses** to whole-object + above roughly an eighth of the object fetched, and the crossover is what D1 + moves. Report the crossover point, not a single ratio. +- Page cache hides everything. `pure-ssd-mode.md`'s direct-I/O settings plus + `UMBP_SSD_TIMING=1` (diagnostic only — it prints several `[SsdPerf/*]` lines + per batch) are how device-side numbers become visible. A run showing no + `dev_GB_s` movement between D0 and D1 is measuring the page cache, not the + drive. + +**End to end.** The sglang tree connector with `ssd_enabled=true` on a +`pure-ssd-mode.md` deployment, plus `UMBP_DISTRIBUTED_RANGED_SCRATCH_BYTES` set: +before D0 it declines ranged I/O outright; after, it should negotiate ranged and +produce identical output. Equality of generated tokens against a whole-object run +is the acceptance test. + +--- + +## 11. Support matrix, after + +| Client | Backend | Ranged get | Ranged put | +|---|---|---|---| +| `DistributedClient` | — | opt-in, all media incl. SSD (D0); extent-scoped disk reads (D1); extent-scoped wire (D2) | opt-in, all media incl. SSD (D0) | +| `StandaloneProcessClient` | Distributed | same, inherited | same, inherited | +| `StandaloneProcessClient` | Local | unchanged — out of scope (§9.1) | unchanged — out of scope (§9.1) | +| `StandaloneClient` | — | unchanged — out of scope (§9.1) | unchanged — out of scope (§9.1) | + +`design-tree-connector-port.md` §9's matrix predates `abf38df7` and is stale on +the local-mode rows independently of this work; it should be corrected to point +here for the distributed rows once D0 lands. diff --git a/src/umbp/doc/design-standalone-process-mode.md b/src/umbp/doc/design-standalone-process-mode.md index f9d03a146..3c7aac90d 100644 --- a/src/umbp/doc/design-standalone-process-mode.md +++ b/src/umbp/doc/design-standalone-process-mode.md @@ -102,9 +102,8 @@ Two new roles, named after Mooncake's split for consistency: │ Fd-handoff listener (raw AF_UNIX, SEPARATE socket path from │ │ the gRPC UDS — gRPC cannot carry an fd) │ │ │ - │ client_ : IUMBPClient — either │ - │ StandaloneClient (local-backend, §3.1) │ - │ DistributedClient (distributed-backend, §6) │ + │ client_ : IUMBPClient — a DistributedClient, deployed either │ + │ with no master (embedded, §3.1) or with one (§6) │ │ │ │ Per-worker shm registry: client_id → {fd, base, size, offset} │ └─────────────────────────────────────────────────────────────────┘ @@ -120,18 +119,23 @@ separate raw `AF_UNIX` socket alongside the gRPC control-plane UDS. constructed through the existing `CreateUMBPClient(config)` factory (`umbp_client_factory.cpp`). This yields three deployment shapes: +> **Superseded.** `StandaloneClient` no longer exists. The server's +> backend is always a `DistributedClient`; what the shapes below called +> "local-backend" is now that same client deployed with no master, which +> the factory synthesizes when the config names no deployment. The +> distinction the table draws is real, but it is a property of the +> configuration, not of which class is constructed. + | Shape | `client_` backend | Cross-node | Server's own storage | |---|---|---|---| -| **local-backend** (same-host, fallback/dev-convenience) | `StandaloneClient` | no | `LocalStorageManager` (DRAM+SSD tiers, private anonymous/hugetlb memory — reused unmodified from `local/`) | -| **distributed-backend, backend enabled** (Mooncake-parity shape) | `DistributedClient` | yes, via Master + RDMA | `DistributedClient`'s own DRAM pool (a distinct concept from `LocalStorageManager` — see §6.5) | -| **distributed-backend, backend disabled/unconfigured** | `StandaloneClient` | no | same as local-backend | - -**distributed-backend is the primary target this feature builds -toward**, since it is the shape that actually matches Mooncake parity -(goal 2). **local-backend is a same-host fallback / dev-convenience -shape**, not a co-equal alternative — it remains fully supported and -is the default when the server has no distributed-backend -configuration. +| **embedded** (same-host, no master) | `DistributedClient`, `master_address` empty | no | its own DRAM pool; nothing routed, registered or heartbeated, and it evicts on its own watermarks | +| **master-backed** (Mooncake-parity shape) | `DistributedClient`, `master_address` set | yes, via Master + RDMA | its own DRAM pool, with eviction driven by master | + +**master-backed is the primary target this feature builds toward**, +since it is the shape that actually matches Mooncake parity (goal 2). +**embedded is a same-host fallback / dev-convenience shape**, not a +co-equal alternative — it remains fully supported and is what the +server runs when given no distributed configuration. Existing data-path RPC handlers (Put/Get/BatchPut/BatchGet/Exists/ Clear/Flush) require no logic changes across these shapes: they diff --git a/src/umbp/doc/design-tree-connector-port.md b/src/umbp/doc/design-tree-connector-port.md new file mode 100644 index 000000000..56ad3a156 --- /dev/null +++ b/src/umbp/doc/design-tree-connector-port.md @@ -0,0 +1,715 @@ +# Porting the Tree-Connector Work onto the Backend-Agnostic Branch — Plan + +**Goal:** bring the GPU-buffer and ranged-multi-buffer-I/O work from +`origin/feat/umbp-tree-connector` onto the backend-agnostic branch, so a +GPU-resident sglang HiCache connector can drive UMBP without a host bounce. + +**Provenance.** Upstream branch `feat/umbp-tree-connector`, originally at +`1ab7e751` with nine UMBP commits on the common base `ae9635f7`. Our side is +`integration/umbp-ssd-verify` at `ee4a5861`, ten UMBP commits on the same base +(the Phase 0–6 backend-agnostic refactor, the HBM/SSD `MediumBackend`s, one +medium per node, and the segment CRC work). This port lands on +`feat/umbp-tree-connector-port`, branched from `ee4a5861`. + +Upstream has since been rebased onto a newer main and grown three more UMBP +commits; it is now at `875d4ac3` with twelve. The rebase left the original nine +byte-identical (`git range-diff 9d5e30f5..1ab7e751 f185c5f0..362cd22c` reports +`=` for all nine), so the second round of porting is exactly those three +commits and nothing else — see §8. + +**Round three.** Upstream is now at `c0044f87` with thirteen, and was *not* +rebased this time (`git range-diff f185c5f0..362cd22c 7f6effc5..c0044f87` +reports `=` for all twelve carried commits), so the third round is exactly one +commit — see §12. This branch was at the same time rebased off `ee4a5861` onto +`c12c68e8`, the backend-agnostic branch's current head, which added the +pure-SSD/multi-drive work; see §12.1 for what that cost. + +**Round four.** Both branches moved again within hours of round three landing. +Upstream `c0044f87` → `0a174274`, two test-quality commits; backend-agnostic +`c12c68e8` → `9febaeaa`, one commit touching only `tests/python/umbp/ +umbp_bench.py`. See §14. + +**Method:** not a merge, and not a series of cherry-picks. The two branches +rewrote almost disjoint halves of `src/umbp`, so most upstream files can be +taken verbatim; the small shared set is hand-merged, and the files where the +two designs actually collide are *reimplemented* rather than ported. + +**Scope:** `src/umbp/` (local tier stack, standalone server/client, distributed +client, common helpers, proto), `src/pybind/`, `tests/`. + +--- + +## 1. Why a merge is the wrong tool + +A trial merge of the upstream branch into ours reports exactly one conflicted +file. That number is misleading in both directions. + +`src/umbp/distributed/pool_client.cpp` conflicts, and it is not resolvable by +picking hunks: upstream's five edit sites are `LocalPutPages`, `LocalGetPages`, +`RemoteDramScatterWrite`, `RemoteDramScatterRead` and `RemoteSsdReadOnce`, all +deleted by Phase 6. `FindRegisteredMemory` survives but now returns +`TransferRef`, not `mori::io::MemoryDesc`. + +Three files auto-merge *cleanly but wrongly*. The worst is +`standalone/standalone_process_client.cpp`: our `RegisterMemory` throws when +`loc != CPU`, and upstream's `7af295f6` implements exactly that case over GPU +IPC. Git keeps both, the throw wins, and the feature is dead with no conflict +marker to notice. + +So the merge would be silently wrong in the places that matter most. + +## 2. What the port does not need to do + +Three pieces of upstream work are already satisfied by the backend-agnostic +refactor. Porting them would duplicate the abstraction they generalize. + +**GPU endpoint dispatch.** `HbmCopyEngine::CanHandle` claims any pair with at +least one GPU endpoint, derives the `hipMemcpyKind` from the pair, and selects +the device the copy runs on (`hbm_copy_engine.cpp:68–81, 111–118`). That is what +upstream's `DeviceCopy` + `ScopedHipDevice` do at each call site, except ours is +selected by `CompositeTransferEngine` instead of branched on inline. We still +take `common/device_copy.{h,cpp}` — the *local* tier stack needs it, and +`DetectPointerLocation` is needed at the API boundary (§4.1) — but no call site +in `pool_client.cpp` gains a manual `hipMemcpy`. + +**Run coalescing.** Both local engines merge adjacent segments while building +the plan (`local_copy_engine.cpp:177–186`, `hbm_copy_engine.cpp:160–163`). This +is upstream's `BuildTierPageRuns` merge generalized: per endpoint pair, across +keys, on host and GPU paths alike, and with no `UMBP_LOCAL_PAGE_MERGE` escape +hatch because it is the only path. The distributed half of `500e186f` is +therefore dropped; its local half is taken. + +**Pointer location in the API.** `IUMBPClient::RegisterMemory` already takes +`loc` and `device` and plumbs them to `transfer_engine_->RegisterMemory` +(`umbp_client.h:131–140`, `pool_client.cpp:509–526`). Upstream had to infer +location from the pointer because the base signature could not express it. We +keep our explicit parameters *and* add inference, because a connector written +against upstream calls `register_memory(ptr, size)` with no location — see §4.1. + +## 3. File classification + +The rule: for any file where `git diff ae9635f7 HEAD -- ` is empty, our +branch never touched it, so upstream's final state *is* the correct merge +result. Take it verbatim. Everything else is hand-merged. + +**Take verbatim (25 files).** Local tier stack and its new source-private +headers (`dram_tier.{h,cpp}`, `tier_backend.{h,cpp}`, `local_storage_manager.{h,cpp}`, +`device_copy_run.h`, `device_gather.{h,hip}`, `host_registration.{h,cpp}`), +`local/standalone_client.{h,cpp}`, `standalone/standalone_server.cpp`, +`common/device_copy.{h,cpp}`, `common/range_utils.h`, +`proto/umbp_standalone.proto`, and the tests +(`test_dram_tier_{ranges,gather,concurrency}.cpp`, `test_umbp_client.cpp`, +`test_standalone_shm_ipc.cpp`, `test_umbp_client_ptr.py`). + +`test_standalone_shm_ipc.cpp` already exists at the base and is already wired in +`tests/cpp/umbp/distributed/CMakeLists.txt`; our branch rewrote that CMakeLists +but kept the entry, so no build wiring is needed for it. + +**Hand-merge (9 files).** + +| File | Reconciliation | +|---|---| +| `include/umbp/umbp_client.h` | Add upstream's two pure virtuals; keep our four-argument `RegisterMemory`. | +| `pybind/pybind_umbp.cpp` | Add the two ranged bindings next to ours. | +| `src/umbp/CMakeLists.txt` | Add `common/device_copy.cpp`, `local/tiers/host_registration.cpp`, `local/tiers/device_gather.hip` to `UMBP_CORE_SRCS`; add `hip::host` as a PUBLIC link on `umbp_core`. | +| `distributed/distributed_client.{h,cpp}` | Add the two warn-once stubs (see §5). | +| `standalone/standalone_process_client.{h,cpp}` | Take upstream's GPU-IPC registration and ranged forwarding; delete our `loc != CPU` throw. | +| `distributed/pool_client.cpp` | **Reimplemented, not ported** — §4. | +| `tests/.../test_pool_client_batch_put.cpp` | Append upstream's two multi-page byte-exact cases; they drive the public API and exercise our coalescing. | +| `tests/cpp/umbp/local/CMakeLists.txt` | Union of their three dram-tier test targets and our `test_segment_crc`. | + +## 4. The four gaps — the whole distributed-side port + +Upstream's 200-line `pool_client.cpp` diff exists to close these. In our +architecture they close much smaller, because the transfer layer already carries +pointer location and offsets. + +### 4.1 An unregistered device pointer is labelled host memory + +`PoolClient::UserBufferRef` (`pool_client.cpp:552–556`) falls back to +`TransferRef::HostBytes(ptr, size)`, which stamps `loc = CPU`. An unregistered +device pointer therefore routes to `LocalCopyEngine`, which `memcpy`s from +device memory. + +Fix: classify the pointer in the fallback with `DetectPointerLocation` and build +a GPU `TransferRef` when it is device memory, so `HbmCopyEngine` claims the +pair. This is one function, and it is what makes an unmodified upstream-era +connector — `register_memory(ptr, size)` with no location — correct on our +branch instead of merely accepted. + +### 4.2 The remote bounce path stages GPU memory through a host buffer + +`MoriIoEngine` decides bounce eligibility on `local.HasHostPtr()` alone +(`mori_io_engine.cpp:403`, `:462`) and then `std::memcpy`s in and out of the +host staging region (`:592`, `:608`). A GPU local endpoint that misses zero-copy +is therefore memcpy'd from device memory. + +Fix: require `local.loc == mori::io::MemoryLocationType::CPU` in both +predicates, so a GPU endpoint that cannot be zero-copied is *rejected by Plan* +and surfaces as a per-key failure. One predicate change at each of two sites is +the single-point equivalent of upstream's five scattered "refusing host staging +fallback" guards — and it is strictly better placed: the transfer layer is the +only layer that knows whether staging is involved. + +### 4.3 `RegisterMemory` lacks upstream's hardening + +Ours returns `true` when a pointer is already registered regardless of the new +size, does not validate the descriptor it got back, and does not catch +(`pool_client.cpp:509–526`). + +Fix: port upstream's logic against `TransferRef` — reject a re-registration that +grows the region, validate base/size/loc/device on the returned ref and +deregister on mismatch, and convert throws into `false`. + +### 4.4 Standalone-process mode throws on the case upstream implements + +`StandaloneProcessClient::RegisterMemory` raises "only CPU (AnonymousShm-backed +host) buffers are supported". Replace the throw with upstream's +`RegisterDeviceMemory` path, keeping our explicit `loc`/`device` parameters and +inferring them when the caller leaves them at the default. + +**The server-side half, which is not visible in the client diff:** +`standalone_server.cpp:757` registers the mapped worker region with the +two-argument `client_->RegisterMemory(base, size)`, which on our branch defaults +to `loc = CPU, device = -1`. On the GPU-IPC path the server maps a *device* +handle there, so that call must pass `loc = GPU` and the device ordinal — +otherwise the region is registered as host memory and lands straight back in +§4.1 and §4.2. + +## 5. Deliberate non-goals + +**Distributed ranged I/O stayed a stub — in round one only.** ~~Upstream ships +`DistributedClient::BatchGetRanges` / `BatchPutRanges` as warn-once all-false, +and this port keeps that.~~ **Superseded by §10.3:** upstream implemented it in +`875d4ac3` and this branch has now ported it. The reasoning below is kept +because it was the design note the implementation was written against, and it +held up — the prediction that "generalizing it is turning a constant into a +parameter" is what the ported code actually does. + +The original note: implementing it is design work (object-range → page-range +mapping, and whether the master's metadata needs to describe sub-object +extents), and doing it inside the first round of the port would make both +unreviewable. + +Worth recording for whoever picks it up: on our branch this is *much* closer +than upstream's shape suggests. `MediumBackend` has no data-movement virtual at +all — allocate, commit, resolve, evict, publish `TransferRef`s — and every byte +moves as a `TransferItem`, which already carries `src_offset`, `dst_offset` and +`size`. A `TransferItem` *is* a range. So no engine, backend or wire-format +change is needed; what is needed is the mapping in `PoolClient`, and the +whole-object path already computes it with the range pinned to `[0, size)` +(`LogicalPageBytes`, consumed by `BuildLocalPageTransfers`). Generalizing it is +turning a constant into a parameter. + +Note also that the gating factor is the *backend*, not the client mode. The +standalone server forwards ranged RPCs straight to its inner client, so +standalone-process over a Local backend works while standalone-process over a +Distributed backend inherited the stub. Implementing this lit up both +distributed rows at once, as predicted. + +**Whole-object I/O is not collapsed into ranged I/O.** Upstream added ranged ops +as a parallel path at every layer — `ReadBatchRangesIntoPtr` beside +`ReadBatchIntoPtr`, `BatchWriteRanges` beside `BatchWrite`, a second RPC pair, a +second client method pair. Whole-object I/O *is* the degenerate single-range +case, so those pairs could collapse into one primitive plus a thin adapter, +halving the surface instead of doubling it. + +This port keeps upstream's parallel structure. Two reasons: a port that also +redesigns is not reviewable against its source, and the collapse has a real +cost — the whole-object path takes shortcuts a general ranged path cannot (the +DRAM tier's parallel non-temporal write of one contiguous payload per key, and +coalescing that assumes dense user offsets), which become "detect the +single-full-range case and take the fast branch". Recommended as a follow-up, +deliberately out of scope here. + +**The three unrelated upstream commits are not taken.** `01a5eaa6`, `3552d318` +and `9d5e30f5` are EPv2 and build fixes that account for every `pyproject.toml` +and `python/mori/ops/` line in the upstream branch diff. They come from main on +their own schedule. + +## 6. Stages — as landed + +| Commit | Contents | +|---|---| +| `doc(umbp): plan the tree-connector port…` | this document | +| `feat(umbp): take the tree-connector local and standalone half verbatim` | the 24-file take-verbatim set (§3); touches nothing the refactor modified, so it cannot regress the distributed data plane | +| `feat(umbp): wire the ranged I/O interface and GPU-capable standalone registration` | `CMakeLists.txt`, `umbp_client.h`, `pybind_umbp.cpp`, `DistributedClient` stubs, local test CMake, and §4.4 both halves | +| `feat(umbp): accept GPU user buffers in the distributed data plane` | §4.1, §4.2, §4.3 | +| `test(umbp): cover multi-page runs and GPU user buffers through PoolClient` | upstream's two multi-page cases plus two GPU cases neither branch had | +| `test(umbp): assert why the GPU remote-path put fails…` | strengthens the rejection case after measuring the fixture's real routing | +| `doc(umbp): record what the port actually built, ran and found` | §7, §8 | +| `doc(umbp): record the mode support matrix and the local SSD device-buffer hole` | §9 | + +§4.4 was folded into the interface commit rather than kept separate: the ranged +pure virtuals and the `RegisterMemory` signature live in the same two files, so +splitting them would have produced a commit that does not compile. + +**Round two** (§10), one commit per upstream commit, in upstream order: + +| Commit | Contents | +|---|---| +| `refactor(umbp): share host registration and gather helpers` | §10.1 — upstream `93f2998a` | +| `fix(umbp): bound heartbeat event batches` | §10.2 — upstream `1d3c859e` | +| `feat(umbp): support ranged I/O in the distributed data plane` | §10.3 — upstream `875d4ac3`, plus the `HbmCopyEngine` gather path it needs | + +## 7. Verification — what was actually run + +Built and run in `rocm/mori-dev:…-rdmapush` with `/apps/ditian12` bind-mounted at +the same path, so the existing `build_check/` cache stays valid. The image has +`ENTRYPOINT ["bash"]`, so `--entrypoint bash` is required or the command is +passed to bash as a filename. + +- **Build:** `umbp_core` (including `device_gather.hip`, compiled for gfx950 and + host), `umbp_client`, `umbp_standalone_server`, and every touched test target. + Clean apart from the pre-existing spdlog `-Wdeprecated-literal-operator` + warnings. +- **pybind:** there is no cmake target for `src/pybind/pybind_umbp.cpp` — it is + built by `pip install`, so a plain `cmake --build` never compiles it and a + typo in a binding would surface only at install time. Checked directly with + `hipcc -fsyntax-only -std=c++17 -I . -I src/umbp/include -I include -I + 3rdparty/spdlog/include -I build_check/src/umbp/proto_gen -I -I + `; note `-I .` and the generated `proto_gen` directory are both + required. +- **ctest:** 22 umbp tests pass (`-E 'cross_node|e2e|medium_selection'`, which + need a second node). This host has `/dev/kfd` and `/dev/dri` but **no** + `/dev/infiniband`; the live `PoolClient` tests nonetheless pass here, so the + known infiniband constraint applies to the multi-node suites rather than to + these. +- **The gather suite runs twice by design** — `umbp_dram_tier_gather` and + `umbp_dram_tier_gather_kernel_off` — because the kernel and the copy engine + must be indistinguishable in their results. Both pass. +- **The four new `PoolClient` cases were confirmed to run, not skip.** The GPU + cases skip when no HIP device is present, and a skip is not a pass, so they + were re-run under `--gtest_filter` and observed as `[ OK ]`. +- **The rejection case was measured, not assumed.** Its first version allowed + either disposition per key; instrumenting it showed all four keys rejected, + because the fixture pins `caller_` to a single page to force remote routing. + It now asserts that, plus the `transfer unplannable` WARN (which pins the + reason to the transfer layer refusing the GPU endpoint) and that no key was + published. `UmbpLogCapture` forces the module to WARN, which is why that + message is observable in the test and invisible in a default-level run. +- **Python:** `BUILD_UMBP=ON pip3 install --no-build-isolation .` from a copy of + the tree, then `pytest tests/python/umbp/test_umbp_client_ptr.py` — 5 passed. + This is the only check that exercises the two new pybind bindings end to end, + and `test_batch_put_get_ranges_round_trip_and_shape_failure` is the closest + thing to a contract test for what the connector calls: a 12-byte object + assembled from three scattered 4-byte buffers in shuffled offset order, read + back as two out-of-order sub-ranges, plus a ragged inner vector rejected as a + whole-call shape error. It needs a full install, so it is not part of the + ctest run above. + +**Not verified here.** The multi-node suites (`cross_node`, `e2e`, +`medium_selection`) need a second node and were excluded. Nothing exercises +GPU-IPC registration in standalone-process mode against a live server — §4.4 is +covered by construction and by review, not by a test — and no measurement was +taken of the GPU path's throughput. + +## 8. Follow-ups this port deliberately leaves open + +1. Implement `BatchGetRanges`/`BatchPutRanges` for distributed (§5). +2. Collapse whole-object I/O into the degenerate ranged case (§5). +3. `HbmCopyEngine` issues blocking `hipMemcpy` on the null stream. Upstream's + `549c14d5` found that sharing one stream per device serialized every batch in + the DRAM tier (1.7 → 4.9 GB/s once split); our engine has the same coupling + in a different shape. Measure once GPU buffers are live end-to-end. +4. SSD ranged reads will work through the staging buffer as soon as §5 lands, + but without disk-side savings until `SsdCopyPipeline` learns sub-extent + reads. +5. **GPU buffers are not supported against the local SSD tier.** Inherited from + upstream, not introduced here. `SSDTier::ReadIntoPtr` passes the caller's + pointer straight into `ReadRecordLocked` (`ssd_tier.cpp:301`), which preads + at that address; neither `ssd_tier.cpp` nor `local_storage_manager.cpp` has + any device awareness, and `LocalStorageManager::ReadIntoPtrNoPromote` reads + from the owning tier directly rather than promoting to DRAM first. So in + Local mode (and standalone-process over a Local backend) a device + destination works for a DRAM hit and not for an SSD hit. It should fail + rather than corrupt, but that is untested and the behaviour may differ under + large-BAR host visibility. + + Note the asymmetry, because it is the refactor paying off rather than an + accident: distributed mode has no equivalent hole. There every byte moves + through the transfer layer, so an SSD-resident object reaches a device buffer + via `SsdBackend`'s staging refs and `HbmCopyEngine` — device-awareness is one + property of `TransferRef` instead of a per-tier concern. Fixing the local + side means teaching `SSDTier` to stage, or promoting to DRAM before serving a + device destination. + +## 9. Mode support matrix + +Both features, as this branch leaves them (updated after §12). + +| Client | Backend | GPU buffers | Ranged I/O | +|---|---|---|---| +| `StandaloneClient` | — | DRAM tier only (§8.5) | DRAM tier only | +| `StandaloneProcessClient` | Local | DRAM tier only (§8.5) | DRAM tier only | +| `StandaloneProcessClient` | Distributed | yes, all media | opt-in (§12), except SSD (§10.3) | +| `DistributedClient` | — | yes, all media | opt-in (§12), except SSD (§10.3) | + +Ranged I/O is DRAM-only within Local mode: `TierBackend`'s default ranged +methods return all-false (`tier_backend.cpp:52`, `:69`) and only `DRAMTier` +overrides them. + +In distributed mode two independent conditions must both hold, and +`SupportsRangedIO()` tests both, so a client reports the truth rather than the +caller discovering it as a failed batch. The medium must be able to serve it — +ranged access maps object ranges onto pages a backend publishes as *in-process* +endpoints, and `SsdBackend` publishes storage refs — and the deployment must +have opted in by configuring a ranged scratch arena (§12), which defaults to +zero. + +A client can now ask instead of assuming: `GetBackendMode()` reports what is +behind a forwarding client, and `SupportsRangedIO()` reports the capability. +Both are exposed to Python (`get_backend_mode`, `supports_ranged_io`). + +--- + +## 10. Round two — upstream's three later commits + +Ported after upstream's rebase. In upstream order: + +### 10.1 `93f2998a` — share host registration and gather helpers + +Moves `HostTierRegistration` and `LaunchDeviceGather` from DRAM-tier-private +headers into `src/umbp/include/umbp/common/`, leaving forwarding shims at the +old paths. Our branch had never touched those four files, so this cherry-picked +clean. + +One deliberate deviation: upstream's move rewrote both headers' comments down to +a few lines. This branch keeps the original rationale — the measured +per-fragment costs that justify a kernel over `hipMemcpy` / `hipMemcpy2DAsync`, +the ~128 KiB crossover, and why registration is single-shot and all-or-nothing. +Those numbers are why the code has the shape it has, and §10.3 depends on the +crossover being written down. + +### 10.2 `1d3c859e` — bound heartbeat event batches + +A real bug fix, not a feature: a large offload burst shipped every unacked +bundle in each heartbeat, so once a request exceeded the RPC deadline nothing +was acked and the next heartbeat resent the backlog plus new events. Upstream +saw 921 `DEADLINE_EXCEEDED` failures on an 8-rank run. + +`master_client.cpp` took upstream's change as-is; the only reconciliation is +that the drain reads `DrainAllBackends(Backends())` rather than the deleted +`DrainAllSources(owned_sources_)`. + +The **test** is the interesting half, and a good example of the §1 hazard. +Upstream's fake implements `OwnedLocationSource` and attaches via +`AddOwnedLocationSource` — both deleted by Phase 3. The cherry-pick auto-merges +without a conflict marker and does not compile. It was rebuilt to drive a +`MockBackend` (a full `MediumBackend` already in the tree) through a +`BackendRegistry`, which exercises the same `DrainAllBackends` path production +uses. The registry is attached *after* the commits that build the backlog, since +`SetBackendRegistry` installs the auto-flush hook that would otherwise drain it +before it accumulates. + +### 10.3 `875d4ac3` — ranged I/O in the distributed data plane + +The one that retires §5. §5's prediction held: the object-range → page-range +mapping is one function, `BuildLocalRangeTransfers`, and no backend, engine or +wire-format change was needed. + +What upstream's 780-line `pool_client.cpp` diff spends its length on, and why +none of it is here: a `RangeCopyRun` planner with adjacency merging, a +per-endpoint `hipMemcpyKind`, `ScopedHipDevice` at each call site, and a manual +gather/`hipMemcpyAsync` fork. Every one of those is a thing the transfer layer +already does — §2's three "already ours" items, applied to a second feature. + +Four differences worth reviewing against the source: + +1. **A locally-routed ranged put needs no assembly buffer.** Upstream assembles + scattered ranges into the scratch arena and writes the object; here the + ranges are written straight into the slot's pages, because a `TransferItem` + carries offsets. The arena is needed only for the remote direction, where the + object must be contiguous on the wire. +2. **A remote object is served from the arena, and the local install is a + caching side effect.** Upstream installs into the tier first and serves + through it, so an install failure loses a fetch it already paid for. The + install-failure metric is kept and now counts exactly what its name says: a + missed caching opportunity, not a failed read. +3. **No DRAM/HBM tier filter on remote ranged reads.** Upstream needs one + because its remote path is a hand-written DRAM RDMA read. Ours fetches the + whole object through `ExecuteBatchGetPlan`, which reaches any medium the + owning peer publishes. +4. **Per-key results instead of all-or-nothing.** Items are tagged with the + caller's key index, so the engine's failed tags map back per key. + +**Where the gather kernel went, and why.** Upstream calls `LaunchDeviceGather` +from `pool_client.cpp`. Here it lives in `HbmCopyEngine`, for the §4.2 reason: +the engine is the only layer that sees the segment shape a copy decomposes +into, so it is the only one that can decide when a kernel beats `hipMemcpy`. + +The load-bearing detail is that bucketing is **per device across the whole +batch, not per plan**. A plan is one `(src base, dst base)` pair, so a caller +reading three ranges into three separate GPU allocations produces three +single-segment plans — nothing to gather within any of them, and together +exactly the scattered small-fragment batch the kernel wins on. Getting this +wrong is not a correctness bug, which is why it is recorded: the first version +bucketed per plan, produced correct bytes, and silently never launched a kernel. +The same insight is why the local ranged put batches across keys — one kernel +for the batch rather than one per key. + +Consequence beyond ranged I/O: `PoolClient` declares the live medium's host +buffers and the scratch arena as gather regions, so *every* GPU transfer in the +tree gets the fast path, not just ranged ones. + +**The test** is upstream's `test_pool_client_ranges.cpp` rebuilt against this +architecture: the fixture drives `BackendRegistry` / `MediumBackend` instead of +the deleted `PeerDramAllocator`, and the backend self-allocates its pool +(Phase 2b) so the test passes a size rather than a buffer. All four cases pass, +including the two that assert the gather kernel actually ran — which is what +caught the per-plan bucketing mistake. + +--- + +## 11. Verification — round two + +Same container and constraints as §7. + +- **Build:** whole tree clean, plus the `pybind_umbp.cpp` syntax check (still + not covered by any cmake target — see §7 for the exact invocation). +- **ctest:** `-R umbp -E 'cross_node|e2e|medium_selection'`, 22 of 23 pass. + `umbp_pool_client_ranges` is new and passes all four cases; this host needs + `--privileged --device=/dev/infiniband` for it, unlike the suites §7 ran. +- **One failure, pre-existing and unrelated:** `umbp_local_client` aborts on + `test_gpu_put_get`'s `batch_a == host_a` when `umbp_host_mem_allocator` has + run first in the same ctest invocation. It passes standalone, and it fails + identically at `b1176b61` — the commit *before* any round-two data-plane work + — so it is not a regression from this round. It is a cross-test interaction + through shared host state, and it is worth chasing separately: the two tests + are separate processes, so the coupling has to be the filesystem or hugepages + rather than anything in the library. + +--- + +## 12. Round three — rebase onto `c12c68e8`, and upstream's one later commit + +Two independent movements, done together because the second lands on top of the +first. + +### 12.1 The rebase — `ee4a5861` → `c12c68e8` + +The backend-agnostic branch grew one commit, `c12c68e8` ("port the pure-SSD / +multi-drive work onto the backend-agnostic SSD path", #554): `ShardedSSDTier` +and its factory, O_DIRECT with the 4 KiB-aligned record format v3, hardware +CRC-32C, `PeerSsdManager` single-flight, and a `parallel_for` helper. + +Replaying the twelve port commits over it produced **one conflict, in one file, +of one line**: `pool_client.cpp`'s include block, where `c12c68e8` added +`umbp/common/parallel_for.h` and the ranged commit added +`umbp/common/range_utils.h`. Both were kept. `git range-diff` over the twelve +shows only context shifts — no hunk was rewritten and no resolution changed +anyone's semantics. + +That is worth recording rather than just being lucky, because it is §1's +argument holding a second time from the other side: the SSD work is a +`MediumBackend` and a local tier, the tree-connector work is the client-side +data plane and the DRAM tier, and the two barely intersect. + +Two interactions were checked rather than assumed: + +- **`ShardedSSDTier` inherits the ranged stubs, not the DRAM overrides.** It + derives from `TierBackend`, whose default ranged methods return all-false + (§9), so the new sharded tier reports exactly what the single-drive tier + reports. §8.5's local-SSD device-buffer hole is unchanged by it — wider, in + that it now covers N drives instead of one, but not different. +- **The new SSD tests build and pass under the port's CMake wiring.** + `test_sharded_ssd_tier`, `test_ssd_direct_io`, `test_segment_crc` and + `test_peer_ssd_single_flight` were the risk, because the port rewrote both + test `CMakeLists.txt` files. See §13. + +### 12.2 `c0044f87` — make the distributed ranged scratch opt-in + +The arena defaulted to 256 MiB, so every distributed client allocated and +RDMA-registered it whether or not it ever issued a ranged operation. Zero is now +the default and means "this deployment does not do ranged I/O": no allocation, +no registration, `SupportsRangedIO()` false. `UMBPConfig::Validate` no longer +rejects zero, and neither does `UMBP_DISTRIBUTED_RANGED_SCRATCH_BYTES`. + +Two deviations from upstream, both because our side already differs at the site: + +1. **`SupportsRangedIO()` gains the arena test rather than being replaced by + it.** Upstream's is now exactly `ranged_scratch_size_ > 0`; ours keeps the + medium test as well (§9). The conditions are independent and neither implies + the other — the arena is what the *remote* direction lands in, and the medium + is what decides whether object ranges can map onto published in-process + endpoints at all. Upstream has no medium selector, so it cannot express the + second half. + +2. **`StandaloneServer`'s `shared_reads_` keeps our predicate.** Upstream now + gates concurrent standalone reads on `SupportsRangedIO()`. Taken verbatim + here that would serialize *every* read of a distributed SSD-medium + deployment — the pure-SSD mode §12.1 just merged in — because our + `SupportsRangedIO()` also excludes that medium. The gate buys nothing + either: the RPCs under that lock are whole-object reads, and the only two + paths that touch the arena (`BatchGetRanges` / `BatchPutRanges` in + `pool_client.cpp`) already serialize on `ranged_scratch_mutex_` and bail out + cleanly when it is absent. + +Upstream's cleanup paths null the DRAM pool alongside the arena and guard the +frees with `if (ranged_scratch_)`. Neither carries over: the DRAM backend has +owned its own pool since Phase 2b, and `HostMemAllocator::Free` is already a +no-op on an unallocated handle (`host_mem_allocator.cpp:452`), so `Close()` and +the constructor's `release_scratch` are correct unchanged. + +The test is upstream's, adapted to drive a real `DistributedClient` through +`CreateUMBPClient` rather than the file's bare `PoolClient` fixture — the opt-in +lives in the constructor and in `SupportsRangedIO()`, neither of which +`PoolClient` sees. + +**Consequence for a connector.** Ranged I/O is now off by default in +distributed mode. A GPU-resident HiCache connector that wants it must set +`distributed.ranged_scratch_size` (Python: `ranged_scratch_size` on +`UMBPDistributedConfig`; `UMBP_DISTRIBUTED_RANGED_SCRATCH_BYTES` for the +standalone server) and should read `supports_ranged_io()` back rather than +assume it took effect. + +--- + +## 13. Verification — round three + +Same container as §7 (`rocm/mori-dev:…-rdmapush`, `--entrypoint bash`, +`/apps/ditian12` bind-mounted at the same path), on a host with `/dev/kfd`, +`/dev/dri` and `/dev/infiniband`. + +- **Build:** whole tree clean, incremental over the existing `build_check/`. +- **pybind:** `pybind_umbp.cpp` still has no cmake target, so it was + syntax-checked directly again — see §7 for the invocation. Clean apart from + the pre-existing spdlog `-Wdeprecated-literal-operator` warnings. +- **ctest: the whole suite this time, not the `-R umbp` subset.** That filter, + used in §7 and §11, silently skips 18 targets whose names begin `test_` — + including every SSD target `c12c68e8` added, which are exactly the ones this + round needed to see. With `-E 'cross_node|e2e|medium_selection'` only, + **46 of 47 pass**: `test_sharded_ssd_tier`, `test_ssd_direct_io`, + `test_segment_crc` and `test_peer_ssd_single_flight` among them, and + `umbp_pool_client_ranges` passing all five cases including the new opt-in one. + +**The one failure is `umbp_local_client`, and §11's account of it was wrong.** +§11 recorded it as an ordering artefact that "passes standalone". Measured +properly — 12 runs of the binary from a neutral working directory — it aborts on +`test_gpu_put_get`'s `batch_a == host_a` in 4 of 12 runs at HEAD, with +`UMBP_DRAM_GATHER_KERNEL=0` as well as with the kernel on. So it is **flaky, not +order-dependent, and not the gather kernel**; the single-item `Put`/`Get` +immediately above it never fails, only the batched device path does. + +It is not a regression from this round. The same measurement against +`99fa37d5` — the pre-rebase branch head, built from its own worktree — gives 4 +of 12 on the identical assertion. Same test, same rate, before any of this +round's work. + +(That comparison build also fails `test_batch_get_follower` in the runs that get +past the GPU test, which HEAD's `build_check` never does. That is a property of +the throwaway build directory, not of the commit: follower mode depends on +on-disk state the fresh tree does not have. It is noted only so the raw counts +above are not mistaken for a second finding.) + +Chasing the `test_gpu_put_get` flake belongs with the round-one verbatim import +that introduced it, and wants a look at the DRAM tier's batched device path +rather than at anything ranged. + +--- + +## 14. Round four — both branches moved again + +Recorded because the cost is the interesting part: rounds one to three each +needed real design work, and this one needed almost none. That is the port +converging, not luck. + +### 14.1 Rebase `c12c68e8` → `9febaeaa` + +One commit, "Bench/umbp staging autosize (#564)", touching only +`tests/python/umbp/umbp_bench.py` — a file no port commit has ever touched. +Replaying the sixteen produced no conflict and `git range-diff` reports `=` for +all sixteen: not merely resolvable, but byte-identical after the move. + +### 14.2 Upstream `ec72ec17` + `0a174274` — gather assertions state their precondition + +Taken as one commit here, because the second replaces a temporary CI step the +first adds; landing them separately would add and remove the same block. + +The launch-counter assertions (§10.3) need the DRAM tier's host memory +registered for GPU access. That registration is best-effort by design: where +`hipHostRegister` does not take, `Covers()` stays false, the planner leaves the +batch on the copy engine, and the counter never moves while every byte +expectation still holds. Asserting it is asserting a machine capability — which +is how three `DeviceGatherTest` cases and +`PoolClientRangesTest.GpuRangesUseGatherKernel` came to fail upstream's MI355X +runner on `0 vs 0` with nothing actually wrong. + +Both files now probe the registration once and gate only the counter checks on +it. `test_dram_tier_gather.cpp` gains its own `main()` so a skip exits 77 and +lands in the ctest summary rather than reading as a pass — gtest exits 0 on +`GTEST_SKIP`, so without it ctest saw a pass wherever the path was unreachable. + +Both applied without conflict, including into this branch's *rebuilt* +`test_pool_client_ranges.cpp` (§10.3): the probe sits beside `FreePort()` and +the two guards land on the assertions our version kept from upstream's. + +**The check that mattered.** A commit whose entire content is "skip this +assertion sometimes" is §1's "auto-merges cleanly but wrongly" hazard wearing +its friendliest disguise: if the probe returned false here, the port would +still build, the suite would still be green, and the gather coverage would be +silently gone. Verified under `ctest -V` that it is not: all five +`DeviceGatherTest` / `HostTierRegistrationTest` cases and all five +`PoolClientRangesTest` cases report `[ OK ]` and not `SKIPPED`, on both the +kernel and `kernel_off` variants. The registration takes on this host, so the +launch counters are still enforced. + +### 14.3 Verification + +Same container and host as §13 (n06-21). Incremental build clean; the umbp test +project **46 of 46, five consecutive runs**, with the port-flake work of §15 +already in. No new failures, and `umbp_local_client` — flaky on n09-29 (§13) — +has still never reproduced here. + +--- + +## 15. Test ports: guessed, raced, and fixed + +Not a port concern, but it surfaced running §13's suite on a loaded node and +the fix landed on this branch, so it belongs in the record. + +`umbp_master_client_flush_heartbeat` failed once in 20 runs and passed in +isolation. Its fixture guessed a port from `55500 + rand() % 1500`. These tests +run in a `--network host` container sharing the host's port space with every +other container on the machine, so the guess collides — and the collision +surfaces as `ASSERT_NE(server_, nullptr)` with the real `EADDRINUSE` visible +only in gRPC's log noise above it. + +The five affected tests did not have one problem, so they did not get one fix: + +| Test | Was | Now | +|---|---|---| +| `flush_heartbeat`, `batch_route_get_columnar`, `master_client_lifecycle` | guessed from a time-seeded range | `AddListeningPort(":0", &selected_port)` | +| `umbp_tags` (E2E suite) | `AllocPort()` — bind :0, read, close | same, via the builder's out-param | +| `umbp_tags`, `client_metrics` (MasterServer) | `AllocPort()` for the gRPC port | `listen_address ":0"` + `GetBoundPort()` | +| `client_metrics` (IO engine) | `AllocPort()`, plus a retry-until-unique helper | `io_engine.port = 0` | +| `umbp_tags`, `client_metrics` (metrics port) | `AllocPort()` — unavoidable | retry the server on a fresh port | + +The last row is the one that cannot be fixed by asking the OS. +`MasterServerConfig::metrics_port` is a number, `0` already means "no metrics +server", and nothing reports what it bound, so it must be chosen before +`MasterServer::Run()` binds it. Losing that race is not a test failure but a +`std::terminate` — `MetricsServer`'s constructor throws on the `Run()` thread: + +``` +terminate called after throwing an instance of 'std::runtime_error' + what(): [metrics] bind() on port 46077 failed: Address already in use +``` + +which is exactly how it presented: `umbp_tags` aborting in roughly one +full-suite run in six while passing 30 of 30 standalone. It needs the rest of +the suite competing for ephemeral ports to reproduce. Both fixtures now run +`Run()` under a `try`/`catch` and retry the whole server on a fresh metrics +port. + +The clean fix is a bound-metrics-port accessor on `MetricsServer` and +`MasterServer`, mirroring `GetBoundPort()`. Deliberately not taken: that is a +production API change for a test problem, and this branch is a port. + +Also: `test_umbp_tags`' `AllocPort` returned a hardcoded `50400` on failure — +with two callers, precisely the collision the helper exists to prevent. It +throws now. + +Measured on n06-21 alongside a dozen other containers: 30 standalone runs each +of the four fast targets, 10 of `client_metrics`, then 10 consecutive +full-suite runs, all clean. diff --git a/src/umbp/doc/pure-ssd-mode.md b/src/umbp/doc/pure-ssd-mode.md new file mode 100644 index 000000000..9db76a07c --- /dev/null +++ b/src/umbp/doc/pure-ssd-mode.md @@ -0,0 +1,255 @@ +# Pure-SSD mode and multi-drive SSD tiers + +Run a UMBP node with **no host-DRAM tier**, SSD carrying the whole cache, with +one node's SSD tier spanning **several drives** at once. + +The point of the mode is aggregate NVMe bandwidth: N nodes × M drives all +writing and reading in parallel. Most of this page is about not accidentally +giving that up. + +## Quick start + +Two pieces: the `extra_config` your engine passes to `UMBPStore`, and a handful +of env vars. This is a validated 2-drive-per-node deployment under SGLang +HiCache. + +**`extra_config`, per engine:** + +```json +{ + "medium": "SSD", + "ssd_enabled": true, + "ssd_storage_dir": "/umbp_ssd/drive1,/umbp_ssd/drive2", + "ssd_capacity_bytes": 68719476736, + "ssd_backend": "file", + "ssd_io_backend": "io_uring", + + "master_address": "", + "node_address": "", + "io_engine_port": "19600", + "peer_service_port": "19700", + "cache_remote_fetches": false, + "kv_events_subscriber": true, + "kv_events_endpoint": "tcp://localhost:6557", + + "staging_buffer_size": 2147483648, + "ssd_staging_buffer_slots": 512, + "ssd_staging_buffer_size": 4294967296, + "ssd_staging_use_hugepages": true, + "ssd_staging_hugepage_size": 2097152 +} +``` + +**Env vars, on each engine:** + +```bash +export UMBP_SSD_VERIFY_CRC=0 # skip checksums, matching the DRAM tier +export UMBP_SSD_TIER_IO_THREADS=4 +export UMBP_SSD_DURABILITY=strict # ~free once direct I/O is on +export UMBP_SSD_READ_LEASE_MS=30000 +export UMBP_DRAM_USE_HUGEPAGES=1 +export UMBP_DISTRIBUTED_SSD_STAGING_USE_HUGEPAGES=1 + +# Only if the engine issues ranged (sub-object) I/O -- the SGLang tree +# connector does. Zero is the default and turns the feature off, and on a +# pure-SSD deployment every node is SSD, so leaving it unset costs ranged I/O +# cluster-wide. Size it above the largest object a remote ranged batch moves. +export UMBP_DISTRIBUTED_RANGED_SCRATCH_BYTES=268435456 +``` + +Ranged I/O on an SSD medium used to be refused outright regardless of this +variable; that exclusion is gone. Note what it does and does not buy: a ranged +**put** runs at full speed, while a ranged **get** still reads the whole object +off the drive and saves only the final copy into the caller's buffers, because +a resolve stages the object before the requested ranges are known. See +[design-ssd-ranged-io.md](design-ssd-ranged-io.md). + +Deliberately **not** in that list: `UMBP_SSD_TIMING=1`. It is a diagnostic, not a +production setting — it prints several `[SsdPerf/*]` lines per batch, which on a +busy node is a lot of log. Turn it on when investigating throughput, off again +after; see [runtime-env-vars.md](runtime-env-vars.md) for what it reports. + +**Env var, on `umbp_master` — not the engine:** + +```bash +export UMBP_ROUTE_PUT_SELECT_ALGO=random # see "Spreading writes" +``` + +**Prerequisite, on every node before launching anything:** reserve hugepages — +`sudo sysctl -w vm.nr_hugepages=350000` — or engine/worker startup fails with +`RegisterRdmaMemoryRegionAuto failed ... errno:12 (Cannot allocate memory)`. + +## Enabling the mode + +`medium: SSD` **is** the switch. A node serves exactly one medium, and naming +SSD is the whole opt-in: the node registers an `SsdBackend`, advertises SSD +capacity, and takes puts directly. No separate "pure-SSD" flag and no +`UMBP_ROUTE_PUT_SSD_MODE` — the routing plane does not rank tiers, so every +advertised medium is an equally valid put target and an SSD node needs no +special case (see `route_put_strategy.h`). + +This also means the mode is per node, not per cluster: SSD, DRAM and HBM nodes +coexist in one pool, and which medium a key lands on is a property of the node +master chose, not of the key. + +Selecting SSD with `ssd.capacity_bytes` or `ssd.segment_size_bytes` unset is +rejected at startup rather than silently serving a zero-capacity tier. + +**How the selector reaches the node.** `UMBPDistributedConfig::medium` is +exposed three ways: the pybind field (`config.distributed.medium`), and +`UMBP_DISTRIBUTED_MEDIUM=SSD` for the standalone server binary. The `"medium"` +key shown in the `extra_config` above needs a matching entry in sglang's own +`umbp_store.py` parser — the same gap `single_flight_reads` has. Until that +lands, an SGLang engine selects SSD through the pybind field or the env var, not +through `extra_config` JSON. + +## Spreading writes + +**Set `UMBP_ROUTE_PUT_SELECT_ALGO=random`.** This is the single most important +knob in the mode, and the default is wrong for it. + +The default, `most_available`, sends each key to the emptiest node. Node capacity +only refreshes on a heartbeat, so every batch in between sees the same "emptiest" +node and piles onto it — and concurrent writers herd onto it together. Aggregate +write bandwidth collapses to one node's drives. `random` is capacity-weighted +(a node that cannot fit a block is never chosen), so it still fills evenly +without the herding. + +Leave `UMBP_ROUTE_PUT_NODE_AFFINITY` at `none`. `same` and `local` trade fan-out +for locality, which is the opposite of what this mode is for. + +Both are read by the **master process**. Setting them on an engine does nothing. +The master logs its resolved strategy at startup — check that line if a change +seems to have no effect. + +## Sizing + +| Setting | Meaning | +|---|---| +| `ssd_storage_dir` | Comma-separated, **one directory per physical drive**. More than one turns on multi-drive: keys are balanced across drives and batch I/O runs on all of them at once. | +| `ssd_capacity_bytes` | The **total** across every drive, split evenly — not per drive. | +| `ssd_staging_buffer_slots` / `_size` | Read staging. Per-slot size (`size / slots`) must stay **≥ your largest single value**. 512 slots over 4 GiB = 8 MiB/slot. | +| (writes) | Puts land in the same staging arena, borrowed for the duration of the spill — there is no separate write-staging region to size. | + +Too few staging slots and a large batch spends rounds waiting for one. With the +default `UMBP_SSD_GET_MAX_ATTEMPTS=1` that surfaces as a **cache miss**, not a +retry — raise the slot count, or that setting, if your read staging is small +relative to your batches. + +Each key lives whole on one drive, so a single-key read costs one drive's I/O. +Capacity is charged in **padded on-disk bytes**: values are padded to 4 KiB, so +a 512 B value occupies 8 KiB. Negligible for multi-MiB KV pages, but size +`ssd_capacity_bytes` with your real value size in mind if yours are small. + +## Direct I/O + +**Direct I/O is the default.** Keep it on for any deployment whose numbers you intend to +believe.** Without it the tier opens buffered, and reads are served from the +page cache — a whole working set can be served while moving *zero* bytes to the +device, which makes drive-count and DRAM-vs-SSD comparisons meaningless. + +Pair it with `ssd_io_backend: io_uring`; the posix driver on the direct path is +far slower, and the tier warns about it. + +| Env var | Default | Meaning | +|---|---|---| +| `UMBP_SSD_DIRECT_IO` | `1` | `O_DIRECT`, bypassing the page cache. On by default; `0` restores buffered | +| `UMBP_SSD_VERIFY_CRC` | `1` | `0` = skip checksum verification. Checksumming is 51–67% of GET time and the DRAM tier does none, so turn it off to compare the two without that confound. | +| `UMBP_SSD_TIER_IO_THREADS` | `4` | Parallelism **within one drive's** batch. Worth ~42% on reads versus 1. | +| `UMBP_SSD_SHARD_IO_THREADS` | `0` | Parallelism **across drives**. `0` = one per drive; leave it there unless CPU-bound. | +| `UMBP_SSD_DURABILITY` | `strict` | With direct I/O on, strict costs ~1.4% instead of 42%. Little reason to relax it. | + +**Confirm it is on**: run `iostat -x 1` during a read-heavy phase. Buffered shows +~1% device utilisation while reporting absurd bandwidth; direct shows the drives +busy at something near their rating. + +## Storage-only nodes (no engine, no GPU) + +To add SSD capacity to the pool from a node that isn't running an SGLang +engine — a dedicated storage box, or extra drives on a node already busy with +compute — join it as a bare `UMBPStore` client instead: + +```bash +python tests/python/umbp/umbp_store_node.py \ + --hicache-storage-backend-extra-config \ + '{"master_address": "", "node_id": "ssd-node-1", + "node_address": "", "medium": "SSD", + "ssd_enabled": true, "ssd_storage_dir": "/mnt/nvme0,/mnt/nvme1", + "ssd_capacity_bytes": 1000000000000}' +``` + +It registers with `umbp_master`, heartbeats, and idles — same config surface +(JSON `extra_config` or the matching `UMBP_*` env vars) as an engine's +`UMBPStore`, just with `mem_pool_host=None`. Run it in the same +image/container as the engine (needs sglang + mori importable). Ctrl-C / +`SIGTERM` flushes and exits cleanly; there is no remote-clear signal wired up, +so clearing one means restarting the process. + +If you want compute nodes to contribute *no* cache capacity of their own — +all SSD capacity coming from dedicated storage nodes instead — set the +engine's own `ssd_capacity_bytes` to something unusably tiny (e.g. `1048576`, +1 MB) rather than disabling its SSD tier: `ssd_enabled: true` is still needed +for HiCache's L3 path to initialize, but `random` routing skips any node that +"cannot fit a block," so a 1 MB node is never chosen as a put target. + +## Troubleshooting + +| Symptom | Likely cause | +|---|---| +| Absurd read bandwidth, `iostat` shows ~1% device utilisation | You are reading the page cache: either `UMBP_SSD_DIRECT_IO=0`, or the startup probe fell back to buffered (the filesystem rejects `O_DIRECT`). Check for `direct_io=true` on the `[SSDTier]` init line — requested is not the same as active | +| Adding drives or nodes doesn't add bandwidth | Writes are skewed. Set `UMBP_ROUTE_PUT_SELECT_ALGO=random`, check `UMBP_ROUTE_PUT_NODE_AFFINITY=none` | +| `random` set but still skewed | It was set on the engine instead of `umbp_master`. Check the master's startup strategy line | +| Reads report misses under load | Read staging slots exhausted. Raise `ssd_staging_buffer_slots` or `UMBP_SSD_GET_MAX_ATTEMPTS` | +| Puts fail on an SSD node | The staging arena is exhausted, so `BatchAllocate` had no page to hand out. Check `mori_umbp_backend_medium_events_total{tier="SSD", event="staging_slot_full_reject"}` and raise `ssd_staging_buffer_slots` | +| Node comes up with an empty cache after an upgrade | On-disk record format changed; older records are dropped on version mismatch | +| Tier fills far earlier than expected | Capacity is charged in padded bytes — small values pay up to 16× | +| Startup rejects `medium: SSD` | `ssd.capacity_bytes` or `ssd.segment_size_bytes` is 0 — the selected medium must be sized | +| A config key seems to do nothing | `extra_config` is an allow-list; unknown or wrong-mode keys are ignored with a warning. Grep the engine log for `UMBPStore:` | +| Reads keep missing after killing/relaunching a node several times | Master registry doesn't expire dead clients — stale entries stay capacity-weighted and `random` can still pick them. Restart `umbp_master` along with every engine/worker whenever node topology (count, capacity, drives) changes | +| Near-100% misses (`NO_SLOT`/lease-expired in the log) right after adding capacity per node | Staging slots didn't scale with the new concurrent load. `ssd_staging_buffer_slots` needs to grow with how many concurrent callers can hit one node, not just with data volume | +| Resend right after a write misses instead of hitting | Write-through acks drain asynchronously; a flush/resend that lands before the SSD write actually completes evicts the key with no L3 trace. Give it a moment (scales with value size) before flushing/resending | + +## Measured: L1 vs L2 vs L3 + +sglang HiCache, Kimi-K2.6-MXFP4, TP8, one fixed-length request per round. "Request +length" is the **prefix**; every request appends and generates exactly 1 token on +top, so the number **is** TTFT. Each tier isolated by forcing a resend to hit only +that tier, verified against `/metrics` per-tier hit counters rather than timing. +8 physical drives / 4 shards, 2 local + 6 remote across 3 standalone workers on a +second node. The engine's own SSD capacity was shrunk to 1 MB so it can never be +capacity-selected as a put target and leak reads back to its own idle disk. + +| Request length | Cold recompute | L1 hit | L2 hit | L3 hit | +|---|---|---|---|---| +| 4K | ≈0.18s | ≈0.10s | ≈0.20s | ≈0.14s | +| 16K | ≈0.73s | ≈0.12s | ≈0.20s | ≈0.20s | +| 32K | ≈1.42s | ≈0.14s | ≈0.21s | ≈0.27s | +| 128K | ≈7.7s | ≈0.29s | ≈0.45s | ≈0.84s | + +L1 always wins. L3 closes almost all the way to L2 at small sizes — beating it at +4K — but the gap widens with length, since an L3 hit still pays an SSD round trip. +L3 is ~1.3x faster than cold recompute at 4K and ~9x by 128K. + +## Notes + +- **Scope traps.** `UMBP_ROUTE_PUT_SELECT_ALGO` is read by the master. + `UMBP_DISTRIBUTED_SSD_STAGING_USE_HUGEPAGES` is read only by the standalone + server binary — on the SGLang path the `ssd_staging_use_hugepages` JSON key is + the one that takes effect; setting both covers either launch mode. + `UMBP_DRAM_USE_HUGEPAGES` does nothing on a node whose medium is SSD — that + node registers no DRAM pool. +- `cache_remote_fetches: false` keeps remote-fetch cost honest by not re-caching + whole-object remote reads. It does **not** cover remote *ranged* reads: those + land the object in the local medium whenever it is already in hand (the fetch + brought all of it), which costs nothing on the wire. The extra background + pull, which does cost a second copy on the wire, is + `ranged_locality_prefetch: false` + locally. Turn it on for a production hit-rate setup. +- **SPDK.** `UMBP_SPDK_REACTOR_MASK` defaults to `0x1` — a single core polling + all drives. Give it at least one core per drive (`0x3`, `0xf`, …) before + expecting multi-drive scaling on that path. Multi-drive tiers and SPDK RAID0 + compose; a drive in the list may itself be a RAID0 bdev. +- Every setting above is also a plain field on `UMBPConfig` / `UMBPDistributedConfig` + for callers driving the client directly, without SGLang. + +Full env-var reference: [runtime-env-vars.md](runtime-env-vars.md). diff --git a/src/umbp/doc/runtime-env-vars.md b/src/umbp/doc/runtime-env-vars.md index 9b0d6c8de..384a58eb6 100644 --- a/src/umbp/doc/runtime-env-vars.md +++ b/src/umbp/doc/runtime-env-vars.md @@ -63,7 +63,7 @@ Read by the **master process** (`bin/master_main.cpp` via | `UMBP_HIT_INDEX_TTL_SEC` | `7200` | sec | External KV hit-count entry TTL. A hash with no counted match for longer than this is removed from the hit index. | | `UMBP_HIT_INDEX_GC_INTERVAL_SEC` | `60` | sec | External KV hit-count GC sweep interval. | | `UMBP_HIT_QUERY_MAX_BATCH` | `4096` | count | Maximum hashes accepted by one `GetExternalKvHitCounts` request. Oversized requests return gRPC `INVALID_ARGUMENT`; the server does not truncate. | -| `UMBP_ROUTE_PUT_SELECT_ALGO` | `most_available` | enum | Base RoutePut placement algorithm over eligible nodes (HBM before DRAM; SSD never a direct-put target). `most_available` = pick the node with the most projected free space; `random` = capacity-weighted random (probability proportional to projected `available_bytes`, never picks a node that cannot fit). Unknown value → default + one WARN. | +| `UMBP_ROUTE_PUT_SELECT_ALGO` | `most_available` | enum | Base RoutePut placement algorithm over eligible node entry pools. `most_available` = pick the node with the most projected free space; `random` = capacity-weighted random (probability proportional to projected `available_bytes`, never picks a node that cannot fit). Unknown value → default + one WARN. | | `UMBP_ROUTE_PUT_NODE_AFFINITY` | `none` | enum | Node-affinity bias layered on top of the base algorithm. `none` = pure base algorithm; `same` = try to place the whole batch on one node that fits the non-dedup total, else per-key sticky to the first picked node; `local` = per-key prefer the requester's local node. All three fall back to the base algorithm so affinity never makes a key fail that the base algorithm could route. Unknown value → default + one WARN. | | `UMBP_MASTER_INDEX_SHARDS` | `32` | count | Number of independently-locked, key-hashed shards backing the block-location index inside `InMemoryMasterMetadataStore` (the block lock domain, separate from the single `meta_mutex_` that guards client records + external-KV). A heartbeat's event batch only takes the exclusive lock on the shards its keys hash into, so unrelated `RoutePut` / `BatchLookup` readers on other shards don't block behind a large apply; full-sync likewise becomes N small critical sections instead of one giant one. Read once at store construction via `std::strtol`; unset / unparseable / `< 1` → default `32` (a WARN is logged on unparseable input), clamped to a max of `4096`. `1` reproduces the old single-lock block index. Production guidance for heavy heartbeat fan-out (hundreds of clients): `64`. | @@ -75,7 +75,8 @@ that has loaded `libmori_pybinds.so`). | Env var | Default | Unit | Description | |---|---|---|---| | `UMBP_DRAM_READ_LEASE_MS` | `500` | ms | Peer-side DRAM/HBM read lease: how long a single `PeerDramAllocator::Resolve` protects its key's pages from concurrent local `Evict`, covering one RDMA read of those pages. Only needs to exceed one DRAM RDMA round trip (sub-ms), so 500 ms is ~100x margin. Read once at `PoolClient::Init`; `min_allowed=1`. | -| `UMBP_SSD_READ_LEASE_MS` | `3000` | ms | Peer-side SSD read-staging slot lease: how long a claimed staging slot is reserved before the peer reclaims it by TTL (the fallback when the reader's best-effort `ReleaseSsdLease` is lost), and, echoed back in `PrepareSsdReadResponse.lease_ttl_ms`, the reader's validity window anchored at `t_send`. Must exceed one SSD read + RDMA (slower than DRAM), but too long pins one of only ~16 slots on a lost release. Also the fallback for the `PrepareSsdRead` RPC deadline when `UMBP_SSD_PREPARE_TIMEOUT_MS` is unset. Read once at `PoolClient::Init`; `min_allowed=1`. | +| `UMBP_SSD_READ_LEASE_MS` | `3000` | ms | Peer-side SSD read-staging slot lease: how long a claimed staging slot is reserved before the peer reclaims it by TTL (the fallback when the reader's best-effort `ReleaseSsdLease` is lost), and, echoed back in `PrepareSsdReadResponse.lease_ttl_ms`, the reader's validity window anchored at `t_send`. Must exceed one SSD read + RDMA (slower than DRAM), but too long pins a configured staging slot on a lost release. Also the fallback for the `PrepareSsdRead` RPC deadline when `UMBP_SSD_PREPARE_TIMEOUT_MS` is unset. Read once at `PoolClient::Init`; `min_allowed=1`. Also the read lease of a policy-declared `SsdBackend`, where it governs how long a contiguous staging span stays claimed and so sets that backend's read concurrency against a fixed arena — see `staging_slots`. | +| `UMBP_RESOLVE_BUSY_TIMEOUT_MS` | `30000` | ms | Overall deadline for retrying a peer-local or remote batch resolve that returned `BUSY` because SSD staging was temporarily full. Retries discard the whole response and use exponential backoff capped at 50 ms. A batch whose own working set exceeds the arena returns a permanent failure immediately and is not retried. Values above 300000 ms are clamped. | | `UMBP_RPC_SHUTDOWN_TIMEOUT_MS` | `3000` | ms | Deadline for `UnregisterClient` and the last `Heartbeat` in `~MasterClient`. Bounds `~MasterClient` worst-case at ≤ 2 × this value. | | `UMBP_GRPC_SHUTDOWN_DEADLINE_SEC` | `3` | sec | `server_->Shutdown(deadline)` budget, shared by master and peer service. | | `UMBP_METRICS_REPORT_INTERVAL_MS` | `1000` | ms | Cadence at which the pool client's `MasterClient` flushes buffered counters/gauges/histograms via `ReportMetrics`. | @@ -85,6 +86,15 @@ that has loaded `libmori_pybinds.so`). | `UMBP_RELEASE_LEASE_TIMEOUT_MS` | `1000` | ms | Per-attempt gRPC deadline for the best-effort `ReleaseSsdLease` RPC so a slow peer can't stall the reader. `min_allowed=1`. | | `UMBP_SSD_PREPARE_TIMEOUT_MS` | `0` | ms | Per-call gRPC deadline for `PrepareSsdRead` so a hung/slow peer can't stall the serial batch. `0` = fall back to `UMBP_SSD_READ_LEASE_MS` (cluster-homogeneous). A timed-out / failed prepare is a hard not-served outcome (NOT retried, and never a miss). `min_allowed=0`. | | `UMBP_AUTO_FLUSH_EVENT_THRESHOLD` | `128` | count | Peer-side unshipped `KvEvent` outbox size at which a completed batch of puts auto-triggers a heartbeat flush (`FlushHeartbeat`), so the ADDs become visible at the master without waiting for the heartbeat interval or an explicit `Flush()`. Counted on `PeerDramAllocator` only (SSD events still wait for the interval). Parsed via `std::strtoull` (no WARN on bad input); unset / unparseable -> default `128`; `0` disables size-based auto-flush entirely (ADDs then ship only on the heartbeat interval or an explicit `Flush()`); set to a very large value to keep auto-flush armed but effectively never fire on size. Cached on first use in `MasterClient::SetPeerDramAllocator`. | +| `UMBP_WORKLOAD_TRACE_PATH` | empty | path | Records successful production PUT/GET calls in the benchmark trace format. Empty disables recording. | +| `UMBP_WORKLOAD_TRACE_CLIENT_ID` | `0` | id | Logical client stream written into production trace events. | +| `UMBP_WORKLOAD_TRACE_SEED` | `0` | seed | Payload seed stored in the trace header for deterministic replay. | +| `UMBP_PEER_MIN_POLLERS` | `8` | count | Minimum gRPC sync poller threads on the peer service. The default pool is small, and every peer RPC that stages SSD bytes holds one of those threads for the whole device read — on a loopback read the pool, not the drive, is the binding constraint. Same head-of-line-blocking fix as `UMBP_MASTER_MIN_POLLERS`. | +| `UMBP_PEER_MAX_POLLERS` | `64` | count | Maximum gRPC sync poller threads on the peer service. Clamped up to `UMBP_PEER_MIN_POLLERS` if set lower. | +| `UMBP_DISTRIBUTED_SSD_STAGING_USE_HUGEPAGES` | `0` | bool | Back the SSD staging arena with hugetlbfs pages. Every byte the SSD backend moves crosses that arena twice (device <-> arena, arena <-> wire), so its TLB behaviour is on the critical path in a way an ordinary buffer's is not. Falls back to 4 KiB pages when no hugetlb pages are free — a node must still come up. Standalone server binary only; the SGLang path uses the pybind `ssd_staging_use_hugepages` field. | +| `UMBP_DISTRIBUTED_SSD_STAGING_HUGEPAGE_SIZE` | 2 MiB | bytes | Hugepage size requested for the SSD staging arena. | +| `UMBP_DISTRIBUTED_RANGED_SCRATCH_BYTES` | `0` | bytes | Size of the registered host arena backing ranged (sub-object) multi-buffer I/O, and **the entire opt-in for the feature**: at `0` the client allocates and registers nothing and `SupportsRangedIO()` reports false, so a caller that asks — the SGLang tree connector does — falls back to whole-object I/O. Purely a remote-path resource: remote gets are staged in it and remote puts assembled in it, while ranged I/O served by this node's own medium never touches it. Must exceed the largest object a ranged batch will move remotely; an object bigger than the arena fails that key with an error rather than the batch. Applies on **every** medium including SSD (that exclusion was lifted — see `doc/design-ssd-ranged-io.md`). Standalone server binary only; the SGLang path uses the pybind `ranged_scratch_size` field. | +| `UMBP_DISTRIBUTED_MEDIUM` | `DRAM` | enum | The one medium this node serves when no backend policy is given: `DRAM`, `HBM` or `SSD`. Selecting `SSD` is the whole opt-in for a storage node — there is no separate pure-SSD flag. An unknown value is rejected at startup rather than silently defaulting. Ignored once `UMBP_BACKEND_POLICY` names a policy, which describes the backends directly. Standalone server binary only; the SGLang path uses the pybind `medium` field. | ## SPDK proxy @@ -119,9 +129,17 @@ Python wrapper construct one) — they are read once. |---|---|---| | `UMBP_DRAM_CAPACITY` | 4 GiB | `dram.capacity_bytes`. | | `UMBP_DRAM_HIGH_WM` / `UMBP_DRAM_LOW_WM` | `0.9` / `0.7` | DRAM tier eviction watermarks. | -| `UMBP_SSD_ENABLED` | `1` | `0` to disable the SSD tier entirely. | -| `UMBP_SSD_DIR` | `/tmp/umbp_ssd` | POSIX backend root. | -| `UMBP_SSD_CAPACITY` | 32 GiB | `ssd.capacity_bytes`. | +| `UMBP_SSD_ENABLED` | `1` | `0` to disable the SSD tier entirely. **This no longer selects a medium**: a single-medium SSD node is chosen with `UMBP_DISTRIBUTED_MEDIUM=SSD` (or `distributed.medium`), or named as a backend in `UMBP_BACKEND_POLICY`, and `umbp_standalone_server` refuses to start if this is set while neither named one, rather than quietly serving DRAM. | +| `UMBP_EMBEDDED_DRAM_PAGE_SIZE` | 2 MiB | Allocation granularity of the pool an **embedded** deployment gets when the config names no deployment at all. The pool is paged, so a value smaller than a page still occupies one: a caller storing KV pages should set this to the page's exact byte size, and one storing small values wants it smaller. Ignored once `distributed` is configured — set `distributed.dram_page_size` there instead. | +| `UMBP_SSD_DIR` | `/tmp/umbp_ssd` | POSIX backend root(s). **Comma-separated for multi-drive**, one directory per physical drive (e.g. `/mnt/nvme0,/mnt/nvme1`). More than one turns the tier into a `ShardedSsdTier`: keys are placed on the drive with the most free space and batch IO runs on every drive at once, so the tier delivers their aggregate bandwidth. | +| `UMBP_SSD_CAPACITY` | 32 GiB | `ssd.capacity_bytes`. With multiple `UMBP_SSD_DIR` entries this is the TOTAL budget, split evenly across the drives. Charged in **padded** on-disk bytes (see `UMBP_SSD_DIRECT_IO`), not raw value bytes. | +| `UMBP_SSD_DIRECT_IO` | `1` | `ssd.direct_io` — `1` opens segments `O_DIRECT`, bypassing the page cache. **On by default**; set `0` only to deliberately measure or exploit the page cache. Buffered can serve an entire working set from page cache and move **zero** bytes to the device, which makes drive-count and DRAM-vs-SSD comparisons meaningless — that is why the default is direct. Requires record format v3 (already unconditional), so a directory reads back either way. | +| `UMBP_SSD_VERIFY_CRC` | `1` | Compute checksums on write, verify on read. Off isolates how much of the SSD path is integrity work rather than storage — the DRAM tier does no checksumming, so an unqualified comparison conflates the two. Records written with it off carry `kFlagNoCrc` and stay readable either way. | +| `UMBP_SSD_TIER_IO_THREADS` | `4` | Worker threads for the CPU-bound phases *inside one drive's* batch (checksum verify on read, checksum + record assembly on write). Matches the DRAM tier's default so a DRAM-vs-SSD comparison is not silently 4 threads against 1. | +| `UMBP_SSD_SHARD_IO_THREADS` | `0` | Worker threads for the multi-drive fan-out. `0` = one per drive, which is what saturates N drives. Ignored with a single `UMBP_SSD_DIR`. | +| `UMBP_SSD_SINGLE_FLIGHT` | `1` | Coalesce concurrent reads of the same key: the first requester touches the drive, the rest are served by memcpy from its buffer. Exists for MLA + TP, where every attention-TP rank GETs a byte-identical key and one page is otherwise read `tp_size` times. MHA keys carry a per-rank suffix and never collide, so this is a no-op there. Reaches only configs built by `UMBPConfig::FromEnvironment()` — sglang's `UMBPStore` needs the pybind `single_flight_reads` field instead. | +| `UMBP_SSD_DURABILITY` | `strict` | `strict` fdatasync()s every batch write; `relaxed` leaves it to the kernel. Relaxed is safe for a pure cache — the bytes are re-fetchable and the peer discards leftover segments at startup — so the flush buys nothing there while costing a large share of write time. Unknown values leave the configured mode untouched. | +| `UMBP_SSD_TIMING` | unset | Diagnostic. Prints a stage-by-stage `[SsdPerf/tier]`, `[SsdPerf/shard]` and `[SsdPerf/peer]` breakdown per batch (index lookup / device IO / CRC / fan-out, with implied GB/s), which is what answers "device or CPU?" without a profiler. Several lines per batch — turn it off again after investigating. | | `UMBP_SSD_BACKEND` | `file` | `file` or `spdk`. Implicitly upgraded to `spdk` if `UMBP_SPDK_NVME_PCI` is set. | | `UMBP_EVICTION_POLICY` | `lru` | Forwarded to `eviction.policy`. | | `UMBP_ROLE` | (empty) | `leader` / `follower` / `standalone`. If unset, falls back to `LOCAL_RANK` / `OMPI_COMM_WORLD_LOCAL_RANK` / `SLURM_LOCALID` / `MPI_LOCALRANKID`: rank 0 → leader, others → follower. | @@ -153,11 +171,12 @@ operators can find them in one place. | Env var | Description | |---|---| -| `UMBP_MASTER_ADDRESS` | `host:port` of the master to connect to (e.g. `10.0.0.1:15558`). | +| `UMBP_MASTER_ADDRESS` | `host:port` of the master to connect to (e.g. `10.0.0.1:15558`). **Optional.** Without it the same client is an embedded deployment: same backends and transfer engine, nothing routed, registered or heartbeated, and a local miss is the final answer. It is also what decides whether the backend evicts for itself and whether it records heartbeat events at all. | +| `UMBP_BACKEND_POLICY` | JSON backend/tier policy used by a distributed-backed standalone server. | | `UMBP_MASTER_LISTEN` | `host:port` the master should listen on (when starting it locally). | | `UMBP_MASTER_AUTO_START` | `true`/`false`: auto-spawn `umbp_master` on this node before connecting. | | `UMBP_MASTER_BIN` | Path to the `umbp_master` binary. The Python `mori.umbp` package auto-fills this from the packaged binary; override to point at a custom build. | -| `UMBP_NODE_ADDRESS` | This node's address as advertised to peers. Must be reachable from every other node. | +| `UMBP_NODE_ADDRESS` | This node's address as advertised to peers. Must be reachable from every other node. Required only **with** a master: with none, nothing registers and no peer can dial in, so `umbp_standalone_server` synthesizes it (as it does `UMBP_NODE_ID`) and omitting `UMBP_IO_ENGINE_HOST` simply builds no RDMA engine. | | `UMBP_IO_ENGINE_HOST` | `mori::io::IOEngine` listener host (typically `127.0.0.1`). | | `UMBP_IO_ENGINE_PORT` / `UMBP_IO_ENGINE_PORTS` | IO engine port (single port, or comma-separated list for multi-engine deployments). | | `UMBP_PEER_SERVICE_PORT` | Port `PeerServiceServer` should bind. | diff --git a/src/umbp/include/umbp/common/aligned_buffer.h b/src/umbp/include/umbp/common/aligned_buffer.h new file mode 100644 index 000000000..7c4b4f8ff --- /dev/null +++ b/src/umbp/include/umbp/common/aligned_buffer.h @@ -0,0 +1,102 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include + +#include "umbp/local/tiers/segment/segment_format.h" + +namespace mori::umbp { + +// Heap buffer aligned to segment::kRecordAlign, with the size rounded up to the +// same boundary. O_DIRECT requires buffer address, file offset and length to all +// be alignment multiples; std::vector only guarantees max_align_t (16 B), +// so it cannot back a direct read or write. +class AlignedBuffer { + public: + AlignedBuffer() = default; + explicit AlignedBuffer(size_t size) { Resize(size); } + + ~AlignedBuffer() { std::free(data_); } + + AlignedBuffer(const AlignedBuffer&) = delete; + AlignedBuffer& operator=(const AlignedBuffer&) = delete; + + AlignedBuffer(AlignedBuffer&& other) noexcept + : data_(std::exchange(other.data_, nullptr)), + size_(std::exchange(other.size_, 0)), + capacity_(std::exchange(other.capacity_, 0)) {} + + AlignedBuffer& operator=(AlignedBuffer&& other) noexcept { + if (this != &other) { + std::free(data_); + data_ = std::exchange(other.data_, nullptr); + size_ = std::exchange(other.size_, 0); + capacity_ = std::exchange(other.capacity_, 0); + } + return *this; + } + + // Grows (never shrinks) the allocation and sets the logical size. Bytes + // between `size` and the padded capacity are zeroed so record padding never + // leaks heap contents to disk. + void Resize(size_t size) { + const size_t padded = static_cast(segment::AlignUp(size)); + if (padded > capacity_) { + void* p = nullptr; + if (posix_memalign(&p, static_cast(segment::kRecordAlign), padded) != 0) { + throw std::bad_alloc(); + } + std::free(data_); + data_ = static_cast(p); + capacity_ = padded; + } + size_ = size; + if (padded > size_) std::memset(data_ + size_, 0, padded - size_); + } + + char* data() { return data_; } + const char* data() const { return data_; } + size_t size() const { return size_; } + // Allocation rounded up to kRecordAlign — the length to hand a direct I/O op. + size_t padded_size() const { return static_cast(segment::AlignUp(size_)); } + bool empty() const { return size_ == 0; } + + private: + char* data_ = nullptr; + size_t size_ = 0; + size_t capacity_ = 0; +}; + +// True when `ptr` and `size` both satisfy the direct-I/O alignment rules. File +// offsets are aligned by construction in the v3 layout, so a caller that passes +// this can issue the read straight into the user buffer with no bounce. +inline bool IsDirectIoCompatible(const void* ptr, size_t size) { + return (reinterpret_cast(ptr) % segment::kRecordAlign) == 0 && + (size % segment::kRecordAlign) == 0; +} + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/common/config.h b/src/umbp/include/umbp/common/config.h index 21727f887..147aabee4 100644 --- a/src/umbp/include/umbp/common/config.h +++ b/src/umbp/include/umbp/common/config.h @@ -78,6 +78,10 @@ struct UMBPDurabilityConfig { struct UMBPSsdConfig { bool enabled = true; + // Comma-separated (same convention as spdk_nvme_pci_addr). Naming >1 directory + // — normally one mount point per drive — fans the tier across them via + // ShardedSsdTier for their aggregate bandwidth. capacity_bytes is the TOTAL, + // split evenly among the directories. std::string storage_dir = "/tmp/umbp_ssd"; size_t capacity_bytes = 32ULL * 1024 * 1024 * 1024; UMBPSsdLayoutMode layout_mode = UMBPSsdLayoutMode::SegmentedLog; @@ -85,6 +89,57 @@ struct UMBPSsdConfig { UMBPIoConfig io; UMBPDurabilityConfig durability; + // Threads for the sharded tier's batch paths, i.e. ACROSS drives. 0 = one per + // shard, which is what saturates N drives. Ignored with a single directory. + int shard_io_threads = 0; + + // Threads for the CPU-bound phases WITHIN one SSDTier (CRC verify on read, CRC + // + record assembly on write). Defaults to 4 to match the DRAM tier, so a + // DRAM-vs-SSD comparison is not silently 4 threads against 1. + int tier_io_threads = 4; + + // Bypass the page cache (O_DIRECT). On by default: the tier is itself a cache, + // so buffered I/O adds a second unmanaged DRAM cache that steals host memory + // from the DRAM tier and makes any drive measurement meaningless (a run can be + // served entirely from RAM, moving zero bytes to the device). Safe as a default + // because the tier probes O_DIRECT at startup and falls back to buffered with a + // warning on tmpfs/overlayfs. Set false to deliberately exploit the page cache. + bool direct_io = true; + + // Coalesce concurrent reads of the same key: only the first touches the drive, + // the rest are served by memcpy from its buffer. For MLA + TP, where every + // attention-TP rank GETs a byte-identical key; MHA keys carry a per-rank suffix + // and never collide, so it is a no-op there. NOTE: UMBP_SSD_SINGLE_FLIGHT only + // reaches configs built by FromEnvironment() (the standalone server); sglang + // constructs UMBPConfig directly and needs the pybind field. [SsdPerf/peer] + // reports merged=/merged_total=, which is what tells you it fired. + bool single_flight_reads = true; + + // Checksum on write, verify on read. Turning it off isolates how much of the + // SSD path's cost is integrity work (the DRAM tier does none). Records written + // with it off are marked kFlagNoCrc and stay readable either way. + bool verify_crc = true; + + // Split storage_dir on ',' — one entry per drive. Always returns at least + // one element (falls back to the default when the string is empty). + std::vector StorageDirs() const { + std::vector dirs; + size_t start = 0; + while (start <= storage_dir.size()) { + size_t comma = storage_dir.find(',', start); + if (comma == std::string::npos) comma = storage_dir.size(); + std::string part = storage_dir.substr(start, comma - start); + // Trim surrounding whitespace so "a, b" works as well as "a,b". + size_t b = part.find_first_not_of(" \t"); + size_t e = part.find_last_not_of(" \t"); + if (b != std::string::npos) dirs.push_back(part.substr(b, e - b + 1)); + if (comma == storage_dir.size()) break; + start = comma + 1; + } + if (dirs.empty()) dirs.push_back("/tmp/umbp_ssd"); + return dirs; + } + // Local SSD-tier capacity watermarks for the distributed PeerSsdManager's // local eviction. When used/total crosses high_watermark the owner // peer evicts its oldest keys down to low_watermark. Mirrors the DRAM tier's @@ -94,10 +149,10 @@ struct UMBPSsdConfig { double low_watermark = 0.7; // SSD backend selection. "file" uses the segmented-log SSDTier; "spdk" / - // "spdk_proxy" use the SPDK NVMe path (direct SpdkSsdTier in standalone, or + // "spdk_proxy" use the SPDK NVMe path (direct SpdkSsdTier for one process, or // SpdkProxyTier when sharing the device across processes). Kept here (rather - // than at UMBPConfig top level) so both the standalone LocalStorageManager and - // the distributed PeerSsdManager select the backend from the same config. + // than at UMBPConfig top level) so PeerSsdManager and any direct SSDTier user + // select the backend from the same config. std::string ssd_backend = "file"; // "file", "spdk" or "spdk_proxy" std::string spdk_bdev_name; // e.g. "Malloc0" or "NVMe0n1" std::string spdk_reactor_mask = "0x1"; // CPU core mask for SPDK reactors @@ -141,6 +196,13 @@ struct UMBPSsdConfig { if (error_message) *error_message = "ssd.segment_size_bytes must be > 0"; return false; } + // Records are padded to segment::kRecordAlign (4096), so a segment whose + // size is not a multiple of it would leave the append cursor unaligned at + // the roll-over and break direct I/O on the next segment. + if (segment_size_bytes % 4096 != 0) { + if (error_message) *error_message = "ssd.segment_size_bytes must be a multiple of 4096"; + return false; + } // Watermarks must satisfy 0 < low < high <= 1. Fail fast on a misconfigured // value rather than silently clamping (a clamp would hide the config error). if (!(high_watermark > 0.0 && high_watermark <= 1.0 && low_watermark > 0.0 && @@ -207,6 +269,37 @@ inline bool ShouldAdmitReCache(bool cache_remote_fetches, CacheRemoteAdmission p return true; // ALWAYS, or SIZE within cap } +// User-facing HBM-tier knobs for distributed mode. Unlike dram/ssd, HBM has +// no local (non-distributed) mode, so it lives only here rather than on the +// top-level UMBPConfig. Mirrors HbmOwnershipConfig's shape (device + +// buffer_sizes) deliberately: hipMalloc has no hugepage/NUMA/prefault +// dimension the way host memory does, so there is nothing else to expose. +// No `enabled` flag: UMBPDistributedConfig::medium is the single selector, +// and these knobs are read only when it names HBM. +struct UMBPHbmConfig { + int device = 0; // GPU ordinal the pool is allocated on + uint64_t capacity_bytes = 0; // single-buffer pool size +}; + +// The one storage medium a legacy distributed node serves. +// +// UMBP's routing plane does not tier: every advertised medium is an equally +// valid put target (see medium_backend.h and Phase 4 of the backend-agnostic +// refactor), so a node registering two backends MIRRORS across them rather +// than promoting/demoting between them — which is not what "DRAM + SSD" reads +// like, and costs capacity for nothing. Without backend_policy_path, a node +// picks exactly one medium and the cluster gets its heterogeneity from having +// different nodes pick differently. A JSON backend policy explicitly opts +// into peer-local heterogeneous tiering. +// +// The medium each key lands on is therefore a property of the node master +// routed it to, not of a per-node tier order. +enum class UMBPMedium { + DRAM, // host memory (default; the pre-selector behaviour) + HBM, // device memory on UMBPHbmConfig::device + SSD, // local NVMe/file storage, staged through registered host memory +}; + // User-facing distributed configuration. Set UMBPConfig::distributed to enable // distributed mode. Internally translated to PoolClientConfig by DistributedClient. struct UMBPDistributedConfig { @@ -215,7 +308,24 @@ struct UMBPDistributedConfig { size_t staging_buffer_size = 64ULL * 1024 * 1024; // 64 MB - // Dedicated SSD read staging, allocated only when ssd.enabled. Per-slot + // Registered host arena used by ranged multi-buffer I/O. Remote objects are + // fetched into disjoint slices here before being installed into the local + // medium; ranged puts routed to another node assemble their scattered ranges + // into a matching arena. Purely a remote-path resource — ranged I/O served by + // this node's own medium never touches it. + // + // This sizes EACH of the two arenas UMBP allocates: a separate GET arena and + // PUT arena, each under its own mutex, so a remote ranged get and a remote + // ranged put run concurrently instead of serializing on one lock. Each must + // hold at least one whole object. + // + // Zero keeps ranged remote I/O disabled without allocating or registering + // additional host memory; callers that need it must opt in explicitly. An + // existing distributed deployment that never issues ranged I/O therefore + // stops paying for an arena it does not use. + size_t ranged_scratch_size = 0; + + // Dedicated SSD read staging, allocated only when medium == SSD. Per-slot // (this / ssd_staging_buffer_slots) must be >= the largest single-key page KV // (61-layer MLA page ~= 4.5 MB). size_t ssd_staging_buffer_size = 268435456; // 256 MiB @@ -223,10 +333,64 @@ struct UMBPDistributedConfig { // Remote SSD read staging slots; per-slot = ssd_staging_buffer_size / this. int ssd_staging_buffer_slots = 16; + // Back the SSD staging arena with hugetlbfs pages. Every byte the SSD + // backend moves crosses that arena twice (device <-> arena, arena <-> wire), + // so its TLB behaviour sits on the critical path in a way an ordinary + // buffer's does not. Falls back to 4 KiB pages when no hugetlb pages are + // free — a node must still come up. + bool ssd_staging_use_hugepages = false; + size_t ssd_staging_hugepage_size = 2ULL * 1024 * 1024; // 2 MiB + uint16_t peer_service_port = 0; // gRPC peer service port bool cache_remote_fetches = true; // cache remotely-fetched blocks locally + // After a remote RANGED read, pull the whole object into this node's medium + // in the background so the next read of the key is a local hit. + // + // Separate from cache_remote_fetches on purpose. That flag gates a re-cache + // that copies out of the caller's destination buffer, which a GPU-destination + // deployment cannot use and therefore has to turn off; this one never touches + // the caller's buffer — it reads the peer straight into a freshly allocated + // local slot — so the same deployment can keep locality. + // + // The traffic is duplicate: the layer-wise reader will pull the same object a + // slice at a time regardless, so this costs up to one extra copy of the + // object on the wire, in exchange for the object being local for the NEXT + // request. Best-effort throughout — bounded queue, drop on full, admission + // gated by cache_remote_admission / admission_max_block_bytes. + bool ranged_locality_prefetch = true; + + // Ask this node's own media before asking the master. + // + // A read is a two-part question: "who holds this key" (master) and "give me + // the bytes" (whoever holds it). When this node holds the key itself, the + // first part is answered by its own backends, and the routing RPC only + // confirms what a local resolve already knew. A single-node deployment — + // one peer, master on localhost — never gets a different answer, so every + // routing round trip on that path is pure latency. + // + // With this on, BatchGet and BatchExists resolve locally first and contact + // the master ONLY for the keys this node missed; a fully-local batch issues + // no RPC at all. A local hit is conclusive (the backends own the bytes), so + // this cannot invent a hit; a local MISS is not conclusive — another node may + // hold the key — which is why the misses still go to the master. + // + // The cost is on multi-node deployments: local and remote reads stop + // overlapping (the local half is served before the routing RPC is issued + // rather than inside its in-flight window), and a batch whose keys are mostly + // remote pays a local resolve that mostly misses. Set false to restore + // route-first ordering. + // + // One visible consequence: on a node that holds the key, Exists() now answers + // as soon as the put commits, rather than a heartbeat publication interval + // later once the master's index has it. That is a more current answer about + // the key, but it means Exists() on the HOLDER is no longer a barrier for + // "the rest of the cluster can route to this". Poll the node that will + // actually read it -- which has to go to the master -- when that is what is + // being waited for. + bool local_first = true; + // Admission gate for re-caching. Only consulted when cache_remote_fetches is true. CacheRemoteAdmission cache_remote_admission = CacheRemoteAdmission::SIZE; @@ -234,13 +398,39 @@ struct UMBPDistributedConfig { // 0 means unlimited (no size gate). Default 16 MB. size_t admission_max_block_bytes = 16ULL * 1024 * 1024; - // Page size used by Master's PageBitmapAllocator for this node's DRAM/HBM - // tier. Reported via RegisterClient. Same value applies to both DRAM - // and HBM. Forwarded to PoolClientConfig::dram_page_size by - // DistributedClient unmodified. + // Page size used by Master's PageBitmapAllocator for this node's medium. + // Reported via RegisterClient. Forwarded to PoolClientConfig::dram_page_size + // by DistributedClient unmodified. (Named for DRAM because the wire field + // is; it applies to whichever medium this node serves.) // 0 = delegate to Master's ClientRegistryConfig::default_dram_page_size // (2 MiB by default). Set to an explicit byte count to override. uint64_t dram_page_size = 0; + + // Which medium the legacy distributed data plane serves. Ignored when + // backend_policy_path is set. Defaults to DRAM, so existing deployments are + // bit-identical. + // + // The selected medium's sizing knobs come from the config it names: + // DRAM -> UMBPConfig::dram (capacity, hugepages, NUMA, prefault) + // HBM -> distributed.hbm (device, capacity) + // SSD -> UMBPConfig::ssd (storage_dir, capacity, layout, io, backend) + // The other two are ignored rather than validated, so a config that carries + // all three (a common deployment template) selects by this field alone. + UMBPMedium medium = UMBPMedium::DRAM; + + UMBPHbmConfig hbm; + + // Optional JSON backend/tier policy. When set, it replaces the legacy + // single-medium ownership settings with named heterogeneous backends, + // weighted logical tiers, and peer-local offload rules. Appended for + // aggregate-init compatibility. + std::string backend_policy_path; + + // Optional production PUT/GET trace. Payload bytes are regenerated during + // replay, so production traces should disable payload validation. + std::string workload_trace_path; + uint32_t workload_trace_client_id = 0; + uint64_t workload_trace_seed = 0; }; // User-facing same-host standalone-process configuration. Set @@ -265,11 +455,17 @@ struct UMBPConfig { UMBPEvictionConfig eviction; UMBPCopyPipelineConfig copy_pipeline; - // Role is the source of truth for runtime behavior. + // Compatibility only. These selected the shared-SSD leader/follower + // behaviour of the local storage manager, which is gone: a shared pool is a + // deployment now (a master, or a standalone server), not a role a private + // client plays over a shared directory. Nothing reads them any more; they + // stay so existing Python and C++ callers still compile and so + // FromEnvironment keeps accepting UMBP_ROLE. + // + // The one place the distinction survives is where it always belonged: an + // SSDTier opened with SSDAccessMode::ReadOnlyShared reads a segment log + // another process writes. UMBPRole role = UMBPRole::Standalone; - - // Backward compatibility fields for older Python/C++ callers. - // New code should set `role` instead. bool follower_mode = false; bool force_ssd_copy_on_write = false; @@ -334,11 +530,12 @@ struct UMBPConfig { } if (distributed.has_value()) { const auto& d = distributed.value(); - if (d.master_config.master_address.empty()) { - if (error_message) - *error_message = "distributed.master_config.master_address must not be empty"; - return false; - } + // master_address MAY be empty. That is the single-node deployment: the + // client keeps its whole data plane -- backends, transfer engine, peer + // service -- and simply never consults a master. Every read already + // resolves locally first (see local_first), and on one node a local miss + // is the final answer, so there is nothing a master could add. Setting + // the address later is what turns the same process into a cluster member. if (d.master_config.node_id.empty()) { if (error_message) *error_message = "distributed.master_config.node_id must not be empty"; return false; @@ -348,6 +545,30 @@ struct UMBPConfig { *error_message = "distributed.master_config.node_address must not be empty"; return false; } + // Only the selected medium's sizing is checked: a node that serves HBM + // still carries a defaulted dram/ssd block it never allocates from. + if (d.backend_policy_path.empty() && d.medium == UMBPMedium::HBM && + d.hbm.capacity_bytes == 0) { + if (error_message) + *error_message = + "distributed.hbm.capacity_bytes must be > 0 when distributed.medium is HBM"; + return false; + } + // Not ssd.Validate(): that returns early on ssd.enabled == false, and + // selecting SSD here IS the opt-in (DistributedClient enables the tier + // from `medium`, so an unset ssd.enabled must not skip the sizing check). + if (d.backend_policy_path.empty() && d.medium == UMBPMedium::SSD) { + if (ssd.capacity_bytes == 0) { + if (error_message) + *error_message = "ssd.capacity_bytes must be > 0 when distributed.medium is SSD"; + return false; + } + if (ssd.segment_size_bytes == 0) { + if (error_message) + *error_message = "ssd.segment_size_bytes must be > 0 when distributed.medium is SSD"; + return false; + } + } } if (distributed.has_value() && standalone_process.has_value()) { if (error_message) @@ -391,6 +612,30 @@ struct UMBPConfig { cfg.ssd.enabled = getenv_int("UMBP_SSD_ENABLED", cfg.ssd.enabled ? 1 : 0) != 0; cfg.ssd.storage_dir = getenv_str("UMBP_SSD_DIR", cfg.ssd.storage_dir); cfg.ssd.capacity_bytes = getenv_size("UMBP_SSD_CAPACITY", cfg.ssd.capacity_bytes); + cfg.ssd.shard_io_threads = getenv_int("UMBP_SSD_SHARD_IO_THREADS", cfg.ssd.shard_io_threads); + cfg.ssd.tier_io_threads = getenv_int("UMBP_SSD_TIER_IO_THREADS", cfg.ssd.tier_io_threads); + cfg.ssd.direct_io = getenv_int("UMBP_SSD_DIRECT_IO", cfg.ssd.direct_io ? 1 : 0) != 0; + cfg.ssd.single_flight_reads = + getenv_int("UMBP_SSD_SINGLE_FLIGHT", cfg.ssd.single_flight_reads ? 1 : 0) != 0; + cfg.ssd.verify_crc = getenv_int("UMBP_SSD_VERIFY_CRC", cfg.ssd.verify_crc ? 1 : 0) != 0; + // Durability of the SSD cache tier. "strict" (default) fdatasync()s every + // batch write; "relaxed" leaves the data in the page cache and lets the + // kernel flush it. Relaxed is safe for a pure cache — the bytes are + // re-fetchable, and the distributed peer discards leftover segments at + // startup anyway (see PeerSsdManager::DiscardLeftoverOnStartup) — so the + // flush buys nothing there while costing a large share of write time. + { + const std::string durability = + getenv_str("UMBP_SSD_DURABILITY", + cfg.ssd.durability.mode == UMBPDurabilityMode::Relaxed ? "relaxed" : "strict"); + if (durability == "relaxed" || durability == "RELAXED") { + cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; + } else if (durability == "strict" || durability == "STRICT") { + cfg.ssd.durability.mode = UMBPDurabilityMode::Strict; + } + // Any other value leaves the configured mode untouched; Validate() does + // not police this field, so silently keeping the default beats guessing. + } cfg.eviction.policy = getenv_str("UMBP_EVICTION_POLICY", cfg.eviction.policy); cfg.dram.high_watermark = getenv_double("UMBP_DRAM_HIGH_WM", cfg.dram.high_watermark); cfg.dram.low_watermark = getenv_double("UMBP_DRAM_LOW_WM", cfg.dram.low_watermark); diff --git a/src/umbp/include/umbp/common/device_copy.h b/src/umbp/include/umbp/common/device_copy.h new file mode 100644 index 000000000..18da9477c --- /dev/null +++ b/src/umbp/include/umbp/common/device_copy.h @@ -0,0 +1,62 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#pragma once + +#include + +#include + +namespace mori::umbp { + +struct PointerLocation { + bool is_device{false}; + int device_id{-1}; + + bool IsDevice() const { return is_device; } +}; + +// Managed memory is intentionally reported as host memory: it is not a stable +// GPU registration target because its backing pages may migrate. +PointerLocation DetectPointerLocation(const void* ptr); + +class ScopedHipDevice { + public: + explicit ScopedHipDevice(int device_id); + ~ScopedHipDevice(); + + ScopedHipDevice(const ScopedHipDevice&) = delete; + ScopedHipDevice& operator=(const ScopedHipDevice&) = delete; + + bool IsValid() const { return valid_; } + + private: + int previous_device_{-1}; + bool changed_{false}; + bool valid_{false}; +}; + +bool DeviceCopy(void* dst, const void* src, size_t size, hipMemcpyKind kind, int device_id); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/common/device_gather.h b/src/umbp/include/umbp/common/device_gather.h new file mode 100644 index 000000000..79a57d662 --- /dev/null +++ b/src/umbp/include/umbp/common/device_gather.h @@ -0,0 +1,67 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include + +#include +#include + +namespace mori::umbp { + +struct DeviceGatherFragment { + const void* src; + void* dst; + size_t bytes; +}; + +// Copies `count` independent fragments between a HostTierRegistration-covered +// region and one GPU in a single kernel launch, in either direction. +// +// This exists because the DRAM tier's copy path is scatter-gather by nature: +// reading one layer out of page-granular objects means dozens of small strided +// fragments per call, and hipMemcpy has no scatter-gather form. Submitting them +// one at a time costs ~5.4 us each with 8 ranks active; hipMemcpy2DAsync can +// express the uniform-stride case but collapses to ~194 us per fragment +// whenever the source pages are cold, which for a live tier is always. +// A kernel is insensitive to both: measured ~0.57 us per fragment, and flat +// across fragment sizes from 8 KiB to 221 KiB. +// +// Only worth it below roughly 128 KiB per fragment. Above that the copy engine +// wins, because one large hipMemcpyAsync already amortizes its submission. +// +// The host side must be inside a HostTierRegistration-covered range: a kernel +// cannot dereference plain mmap memory at all, it faults. Returns false before +// enqueueing any work when the launch could not be issued, so the caller can +// fall back to hipMemcpy. +bool LaunchDeviceGather(const DeviceGatherFragment* fragments, size_t count, int device_id, + hipStream_t stream); + +// Whether the gather kernel path is compiled in and not disabled via +// UMBP_DRAM_GATHER_KERNEL=0. +bool DeviceGatherEnabled(); + +// Kernels launched since process start. Lets a test assert which path a batch +// actually took, rather than passing whether or not the kernel ran. +uint64_t DeviceGatherLaunchCount(); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/common/grpc_limits.h b/src/umbp/include/umbp/common/grpc_limits.h new file mode 100644 index 000000000..940822b5c --- /dev/null +++ b/src/umbp/include/umbp/common/grpc_limits.h @@ -0,0 +1,142 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +// The ONE gRPC message-size limit for every UMBP surface. +// +// gRPC defaults to 4 MiB on RECEIVE and unlimited on send, and it enforces that +// at the receiver: an oversized message is rejected whole, with +// "Received message larger than max (N vs. 4194304)", after the sender has +// already serialized and shipped it. There is no partial delivery. +// +// Setting it per call site is how UMBP ended up with the limit configured on +// every surface that talks to the MASTER and on none of the surfaces that talk +// to a PEER -- so master-side traffic could grow to 64 MB while peer-side +// traffic silently kept a 4 MiB ceiling. A single EvictKey carrying 155,157 +// victim keys (18.7 MB) was refused on every one of 81 consecutive eviction +// rounds, the tier never reclaimed a byte, and writes then failed with +// NO_SPACE until the deployment was restarted. Nothing logged the connection +// between the two. +// +// So the limit lives here, one value, applied through the same two helpers by +// every channel and every server. A surface that forgets to call them is a +// surface running on gRPC's default, which is the bug above. +// +// This header needs gRPC, so it is only includable from umbp_common sources. +// It sits in common/ rather than distributed/ because both distributed/ and +// standalone/ dial and serve. + +#include + +#include +#include +#include +#include + +#include "umbp/common/env_time.h" + +namespace mori::umbp { + +// 64 MiB. Chosen to match what the master surfaces already used, so this +// change lowers no existing limit. It is a ceiling for a control-plane +// message, not a data path: bulk bytes move over RDMA or shared memory, and +// anything approaching this size is a batch that should be chunked +// (see GrpcMaxItemsPerBatch). +inline constexpr uint32_t kDefaultGrpcMaxMessageBytes = 64u * 1024u * 1024u; + +// Floor of 1 MiB: below that the handshake and registration messages +// themselves start failing, which presents as a node that cannot join rather +// than as a size limit, and is a far worse thing to debug than a slow one. +inline constexpr uint32_t kMinGrpcMaxMessageBytes = 1u * 1024u * 1024u; + +// Resolved once per process. Static so every channel and server in this +// process agrees even if the environment is mutated later -- two surfaces +// disagreeing about the limit is exactly the failure this header exists to +// prevent. +inline uint32_t GrpcMaxMessageBytes() { + static const uint32_t value = GetEnvUint32("UMBP_GRPC_MAX_MESSAGE_BYTES", + kDefaultGrpcMaxMessageBytes, kMinGrpcMaxMessageBytes); + return value; +} + +// Apply to a client channel. Both directions: a client that can send 64 MB +// but only receive 4 MiB fails on the response instead of the request, which +// is the same bug wearing a different hat. +inline void ApplyGrpcLimits(grpc::ChannelArguments* args) { + const int bytes = static_cast(GrpcMaxMessageBytes()); + args->SetMaxReceiveMessageSize(bytes); + args->SetMaxSendMessageSize(bytes); +} + +// Sugar for the common case of a channel with no other arguments. +inline grpc::ChannelArguments GrpcChannelArgs() { + grpc::ChannelArguments args; + ApplyGrpcLimits(&args); + return args; +} + +// Apply to a server. The receive limit is the one that actually rejects +// traffic, so a server that skips this is the site that produces the +// "Received message larger than max" above. +inline void ApplyGrpcLimits(grpc::ServerBuilder* builder) { + const int bytes = static_cast(GrpcMaxMessageBytes()); + builder->SetMaxReceiveMessageSize(bytes); + builder->SetMaxSendMessageSize(bytes); +} + +// How many strings fit in one message, given the limit above. +// +// Raising the ceiling moves the wall; it does not remove it, because the +// batches that approach it grow with the deployment (an eviction round frees a +// FRACTION of the tier, so a bigger tier means more victims) while the limit is +// fixed. A caller shipping a `repeated string` therefore has to split, and +// this is the split point. +// +// Measured rather than assumed: each key costs its own bytes plus a field tag +// and a length varint. kPerItemOverhead is 8 rather than the 3 a <16 KiB +// string actually needs, and kFrameReserve holds back room for the rest of the +// message, because being one byte over costs the WHOLE batch -- the asymmetry +// between "slightly smaller chunks" and "nothing is evicted, ever" is not +// close. +// +// Returns at least 1: a single item larger than the whole budget still has to +// be attempted, and failing on it is more honest than silently sending an +// empty batch forever. +inline size_t GrpcMaxItemsPerBatch(const std::vector& items, size_t start) { + constexpr size_t kPerItemOverhead = 8; + constexpr size_t kFrameReserve = 64 * 1024; + + const size_t limit = static_cast(GrpcMaxMessageBytes()); + const size_t budget = limit > kFrameReserve ? limit - kFrameReserve : limit / 2; + + size_t bytes = 0; + size_t count = 0; + for (size_t i = start; i < items.size(); ++i) { + const size_t cost = items[i].size() + kPerItemOverhead; + if (count > 0 && bytes + cost > budget) break; + bytes += cost; + ++count; + } + return count == 0 ? 1 : count; +} + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/common/host_registration.h b/src/umbp/include/umbp/common/host_registration.h new file mode 100644 index 000000000..61a41f29c --- /dev/null +++ b/src/umbp/include/umbp/common/host_registration.h @@ -0,0 +1,103 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include + +namespace mori::umbp { + +// Makes a host region addressable from GPU kernels via one whole-region +// hipHostRegister call. +// +// Two things depend on this. A gather kernel cannot dereference plain mmap +// memory at all -- it faults the GPU -- so registration is a hard prerequisite +// for the scatter-gather copy path. Independently, registering takes +// hipMemcpyAsync off the pageable staging path, which alone is worth ~6x on the +// per-fragment copies that remain. +// +// Registration runs at roughly 13 GiB/s and does not parallelize: splitting it +// across threads measured slower, because it serializes inside the driver. A +// 512 GiB tier would therefore add ~40 s to server startup, so anything larger +// than the synchronous threshold is registered on a background thread and the +// caller simply behaves as it did before until that finishes. It must be a +// single registration covering the whole region -- see Register() -- so this is +// all-or-nothing rather than a progressive watermark. +// +// The device need not see the region at the host's address -- hipHostRegister's +// own documentation says the two "will have a different value" on many systems +// and that device code must use the device pointer. Equal addresses are a CUDA +// UVA property ROCm never promised. So the alias is recorded per device and +// DeviceAddress() translates by offset, which is sound because the whole region +// is one registration of one contiguous range. +class HostTierRegistration { + public: + // Registers [base, base + bytes). `bytes` must be the mapped (page-aligned) + // length, not the usable capacity. Never throws and never fails hard: if + // registration is unavailable the object simply reports Covers() == false + // forever. + HostTierRegistration(void* base, size_t bytes); + ~HostTierRegistration(); + + HostTierRegistration(const HostTierRegistration&) = delete; + HostTierRegistration& operator=(const HostTierRegistration&) = delete; + + // True when [ptr, ptr + size) is registered, and therefore safe to copy from + // or to without the pageable staging path. False is always the safe answer. + // Weaker than "safe to hand to a kernel" -- that needs DeviceAddress(). + bool Covers(const void* ptr, size_t size) const; + + // Whether a kernel on `device_id` can reach this region at all, i.e. whether + // DeviceAddress() will return non-null for a covered pointer. + bool GatherableOn(int device_id) const; + + // The address `device_id` must use to reach [host_ptr, host_ptr + size), or + // null when that range is outside the region or the device has no alias, in + // which case the caller must stay on the copy-engine path. Subsumes Covers(), + // so a caller that needs the address should not test both. + void* DeviceAddress(const void* host_ptr, size_t size, int device_id) const; + + // Zero until registration succeeds, then the full region. For logging and + // tests. + size_t RegisteredBytes() const { return registered_bytes_.load(std::memory_order_acquire); } + + private: + void Register(); + // Fills alias_bases_, one entry per device. False only when the devices + // cannot be enumerated; a device without an alias gets a null entry. + bool RecordDeviceAliases(); + // Bounds check against an already-loaded registered_bytes_, so an entry point + // that needs the value anyway reads the atomic once. + bool CoversRegistered(const void* ptr, size_t size, size_t registered) const; + + char* base_{nullptr}; + size_t bytes_{0}; + std::atomic registered_bytes_{0}; + // Indexed by device ordinal. Written before registered_bytes_ is published + // and read-only after, so that release/acquire pair is what publishes it. + std::vector alias_bases_; + std::thread worker_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/local/tiers/copy_pipeline.h b/src/umbp/include/umbp/common/parallel_for.h similarity index 55% rename from src/umbp/include/umbp/local/tiers/copy_pipeline.h rename to src/umbp/include/umbp/common/parallel_for.h index 8f045736e..6760e7987 100644 --- a/src/umbp/include/umbp/local/tiers/copy_pipeline.h +++ b/src/umbp/include/umbp/common/parallel_for.h @@ -21,50 +21,38 @@ // SOFTWARE. #pragma once +#include #include -#include -#include #include -#include -#include -#include #include #include -#include "umbp/common/config.h" -#include "umbp/local/tiers/local_storage_manager.h" - namespace mori::umbp { -class CopyPipeline { - public: - CopyPipeline(LocalStorageManager& storage, const UMBPCopyPipelineConfig& config, UMBPRole role); - ~CopyPipeline(); - - bool MaybeCopyToSharedSSD(const std::string& key); - void MaybeBatchCopyToSharedSSD(const std::vector& keys); - bool Drain(std::chrono::milliseconds timeout); - - private: - struct CopyTask { - std::string key; +// Run fn(i) for i in [0, n), on up to `threads` workers, work-stealing so an +// uneven index cost cannot strand a worker. Each index runs exactly once, so +// fn may write result[i] without locking. The caller's thread takes part and +// the call returns only once every index is done. +template +void ParallelFor(size_t n, int threads, Fn&& fn) { + if (n == 0) return; + const size_t workers = std::min(n, static_cast(std::max(threads, 1))); + if (workers <= 1) { + for (size_t i = 0; i < n; ++i) fn(i); + return; + } + std::atomic next{0}; + auto worker = [&]() { + for (size_t i = next.fetch_add(1, std::memory_order_relaxed); i < n; + i = next.fetch_add(1, std::memory_order_relaxed)) { + fn(i); + } }; - - bool EnqueueCopyToSSD(const std::string& key); - size_t EnqueueCopyToSSDBatch(const std::vector& keys); - void CopyWorkerLoop(); - - LocalStorageManager& storage_; - UMBPCopyPipelineConfig config_; - UMBPRole role_; - - std::atomic stop_copy_worker_{false}; - std::vector copy_workers_; - std::mutex copy_mu_; - std::condition_variable copy_cv_; - std::condition_variable drain_cv_; - std::deque copy_queue_; - size_t in_flight_copies_ = 0; -}; + std::vector pool; + pool.reserve(workers - 1); + for (size_t w = 0; w + 1 < workers; ++w) pool.emplace_back(worker); + worker(); + for (auto& t : pool) t.join(); +} } // namespace mori::umbp diff --git a/src/umbp/include/umbp/common/range_utils.h b/src/umbp/include/umbp/common/range_utils.h new file mode 100644 index 000000000..4138c26e9 --- /dev/null +++ b/src/umbp/include/umbp/common/range_utils.h @@ -0,0 +1,71 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include + +namespace mori::umbp { + +inline bool IsObjectRangeOverflow(size_t offset, size_t size, size_t limit) { + return size > limit || offset > limit - size; +} + +template +inline bool RangeBatchShapeValid(size_t entry_count, const std::vector>& ptrs, + const std::vector>& sizes, + const std::vector>& offsets) { + if (ptrs.size() != entry_count || sizes.size() != entry_count || offsets.size() != entry_count) { + return false; + } + for (size_t i = 0; i < entry_count; ++i) { + if (ptrs[i].size() != sizes[i].size() || ptrs[i].size() != offsets[i].size()) { + return false; + } + } + return true; +} + +inline bool RangesTileObject(size_t object_size, const std::vector& sizes, + const std::vector& offsets) { + if (object_size == 0 || sizes.empty() || sizes.size() != offsets.size()) return false; + + std::vector order(sizes.size()); + std::iota(order.begin(), order.end(), size_t{0}); + std::sort(order.begin(), order.end(), + [&](size_t a, size_t b) { return offsets[a] < offsets[b]; }); + + size_t cursor = 0; + for (size_t index : order) { + const size_t size = sizes[index]; + const size_t offset = offsets[index]; + if (size == 0 || offset != cursor || IsObjectRangeOverflow(offset, size, object_size)) { + return false; + } + cursor = offset + size; // Safe after IsObjectRangeOverflow. + } + return cursor == object_size; +} + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/common/ssd_perf.h b/src/umbp/include/umbp/common/ssd_perf.h new file mode 100644 index 000000000..83af2a2fc --- /dev/null +++ b/src/umbp/include/umbp/common/ssd_perf.h @@ -0,0 +1,69 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +// Stage-by-stage timing for the SSD batch put/get path, enabled by UMBP_SSD_TIMING. +// When off every helper is a cached bool test and no clock is read, so this can +// stay in the hot path. Emitted at INFO (module UMBP) as [SsdPerf/tier] (one +// drive: lookup | IO | CRC), [SsdPerf/shard] (per-drive fan-out), [SsdPerf/peer] +// and [SsdPerf/remote], each with keys, bytes, ms and the implied GB/s. + +#include +#include +#include +#include + +namespace mori::umbp::ssdperf { + +using Clock = std::chrono::steady_clock; +using TimePoint = Clock::time_point; + +// Env is read once; the hot-path check is a load of a cached bool. +inline bool Enabled() { + static const bool enabled = [] { + const char* v = std::getenv("UMBP_SSD_TIMING"); + if (v == nullptr || *v == '\0') return false; + const std::string s(v); + return !(s == "0" || s == "false" || s == "FALSE" || s == "off" || s == "OFF"); + }(); + return enabled; +} + +// Meaningless when timing is off, but only ever consumed inside an Enabled() branch. +inline TimePoint Now() { return Enabled() ? Clock::now() : TimePoint{}; } + +inline double MsBetween(TimePoint a, TimePoint b) { + return std::chrono::duration_cast>(b - a).count() * 1000.0; +} + +inline double MsSince(TimePoint a) { return Enabled() ? MsBetween(a, Clock::now()) : 0.0; } + +// Guarded against a zero denominator so a log line never prints inf. +inline double GbPerSec(uint64_t bytes, double ms) { + return ms > 0.0 ? static_cast(bytes) / (ms * 1.0e6) : 0.0; +} + +inline double Pct(double part_ms, double total_ms) { + return total_ms > 0.0 ? part_ms * 100.0 / total_ms : 0.0; +} + +} // namespace mori::umbp::ssdperf diff --git a/src/umbp/include/umbp/distributed/benchmark/payload.h b/src/umbp/include/umbp/distributed/benchmark/payload.h new file mode 100644 index 000000000..0445cb33d --- /dev/null +++ b/src/umbp/include/umbp/distributed/benchmark/payload.h @@ -0,0 +1,42 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#pragma once + +#include +#include +#include +#include + +namespace mori::umbp::benchmark { + +void FillDeterministicPayload(std::string_view key, uint64_t operation_id, uint64_t seed, + void* destination, size_t size); + +std::vector GenerateDeterministicPayload(std::string_view key, uint64_t operation_id, + uint64_t seed, size_t size); + +bool ValidateDeterministicPayload(std::string_view key, uint64_t operation_id, uint64_t seed, + const void* data, size_t size); + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/include/umbp/distributed/benchmark/pool_workload_client.h b/src/umbp/include/umbp/distributed/benchmark/pool_workload_client.h new file mode 100644 index 000000000..22e9cb705 --- /dev/null +++ b/src/umbp/include/umbp/distributed/benchmark/pool_workload_client.h @@ -0,0 +1,183 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/workload_runner.h" +#include "umbp/distributed/pool_client.h" + +namespace mori::umbp::benchmark { + +struct PoolWorkloadClientOptions { + // A key becomes visible to other clients only once the writer's next + // heartbeat reaches the master, so a GET that arrives first has to wait + // rather than count as a miss. The window applies from the PUT this adapter + // performed. + std::chrono::milliseconds publication_timeout{5000}; + // Window for keys this adapter never wrote, e.g. keys a separate writer put + // there. Zero means a missing key is reported immediately. + std::chrono::milliseconds unpublished_timeout{0}; + std::chrono::milliseconds retry_interval{5}; +}; + +// Bridges the transport-neutral WorkloadClient the runner drives onto real +// PoolClients, mapping each workload client id to its own PoolClient so a +// multi-client run exercises routing and remote fetch rather than one process +// talking to itself. +class PoolWorkloadClient final : public WorkloadClient { + public: + explicit PoolWorkloadClient(std::vector>* clients, + PoolWorkloadClientOptions options = {}) + : clients_(clients), + options_(options), + batch_put_calls_(clients == nullptr ? 0 : clients->size()), + batch_get_calls_(clients == nullptr ? 0 : clients->size()) {} + + ClientResult Put(uint32_t id, const std::string& key, const uint8_t* data, size_t size) override { + PoolClient* client = Select(id); + if (client == nullptr || !client->Put(key, data, size)) return ClientResult::kFailed; + Published(key); + client->Master().FlushHeartbeat(); + return ClientResult::kSuccess; + } + + ClientResult Get(uint32_t id, const std::string& key, uint8_t* data, size_t size) override { + PoolClient* client = Select(id); + if (client == nullptr) return ClientResult::kFailed; + const auto deadline = Deadline(key); + do { + if (client->Get(key, data, size)) { + Clear(key); + return ClientResult::kSuccess; + } + std::this_thread::sleep_for(options_.retry_interval); + } while (std::chrono::steady_clock::now() < deadline); + Clear(key); + return client->Exists(key) ? ClientResult::kFailed : ClientResult::kNotFound; + } + + std::vector BatchPut(uint32_t id, const std::vector& keys, + const std::vector>& values) override { + PoolClient* client = Select(id); + if (client == nullptr || keys.size() != values.size()) { + return std::vector(keys.size(), ClientResult::kFailed); + } + Count(&batch_put_calls_, id); + std::vector pointers; + std::vector sizes; + pointers.reserve(values.size()); + sizes.reserve(values.size()); + for (const auto& value : values) { + pointers.push_back(value.data()); + sizes.push_back(value.size()); + } + const auto results = client->BatchPut(keys, pointers, sizes); + std::vector converted(keys.size(), ClientResult::kFailed); + for (size_t i = 0; i < keys.size() && i < results.size(); ++i) { + if (results[i]) { + converted[i] = ClientResult::kSuccess; + Published(keys[i]); + } + } + client->Master().FlushHeartbeat(); + return converted; + } + + std::vector BatchGet(uint32_t id, const std::vector& keys, + const std::vector& sizes, + std::vector>* values) override { + PoolClient* client = Select(id); + if (client == nullptr || keys.size() != sizes.size() || values == nullptr) { + return std::vector(keys.size(), ClientResult::kFailed); + } + Count(&batch_get_calls_, id); + values->resize(keys.size()); + std::vector pointers; + pointers.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + (*values)[i].assign(sizes[i], 0); + pointers.push_back((*values)[i].data()); + } + const auto results = client->BatchGet(keys, pointers, sizes); + std::vector converted(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + if (i < results.size() && results[i]) { + Clear(keys[i]); + converted[i] = ClientResult::kSuccess; + } else { + // Retry the stragglers one key at a time so a single unpublished key + // does not re-fetch the whole batch. + converted[i] = Get(id, keys[i], (*values)[i].data(), sizes[i]); + } + } + return converted; + } + + uint64_t BatchPutCalls(uint32_t id) const { return Calls(batch_put_calls_, id); } + uint64_t BatchGetCalls(uint32_t id) const { return Calls(batch_get_calls_, id); } + + private: + PoolClient* Select(uint32_t id) const { + return clients_ != nullptr && id < clients_->size() ? (*clients_)[id].get() : nullptr; + } + static void Count(std::vector>* counters, uint32_t id) { + if (id < counters->size()) ++(*counters)[id]; + } + static uint64_t Calls(const std::vector>& counters, uint32_t id) { + return id < counters.size() ? counters[id].load() : 0; + } + void Published(const std::string& key) { + std::lock_guard lock(mutex_); + pending_[key] = std::chrono::steady_clock::now() + options_.publication_timeout; + } + std::chrono::steady_clock::time_point Deadline(const std::string& key) { + std::lock_guard lock(mutex_); + const auto found = pending_.find(key); + if (found != pending_.end()) return found->second; + return std::chrono::steady_clock::now() + options_.unpublished_timeout; + } + void Clear(const std::string& key) { + std::lock_guard lock(mutex_); + pending_.erase(key); + } + + std::vector>* clients_; + PoolWorkloadClientOptions options_; + std::vector> batch_put_calls_; + std::vector> batch_get_calls_; + std::mutex mutex_; + std::unordered_map pending_; +}; + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/include/umbp/distributed/benchmark/workload_runner.h b/src/umbp/include/umbp/distributed/benchmark/workload_runner.h new file mode 100644 index 000000000..8399ef155 --- /dev/null +++ b/src/umbp/include/umbp/distributed/benchmark/workload_runner.h @@ -0,0 +1,112 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#pragma once + +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/workload_source.h" + +namespace mori::umbp::benchmark { + +enum class ClientResult { + kSuccess, + kNotFound, + kFailed, +}; + +// A transport-neutral interface suitable for a real UMBP adapter or a test fake. +// Implementations used with multiple clients must permit concurrent calls. +class WorkloadClient { + public: + virtual ~WorkloadClient() = default; + virtual ClientResult Put(uint32_t client_id, const std::string& key, const uint8_t* data, + size_t size) = 0; + virtual ClientResult Get(uint32_t client_id, const std::string& key, uint8_t* data, + size_t size) = 0; + + virtual std::vector BatchPut(uint32_t client_id, + const std::vector& keys, + const std::vector>& values); + virtual std::vector BatchGet(uint32_t client_id, + const std::vector& keys, + const std::vector& sizes, + std::vector>* values); +}; + +struct OperationMetrics { + uint64_t attempted = 0; + uint64_t succeeded = 0; + uint64_t failed = 0; + uint64_t succeeded_bytes = 0; +}; + +struct DistributionMetrics { + std::vector samples_ns; + uint64_t p50_ns = 0; + uint64_t p95_ns = 0; + uint64_t p99_ns = 0; + uint64_t max_ns = 0; +}; + +struct WorkloadMetrics { + OperationMetrics total; + OperationMetrics puts; + OperationMetrics gets; + uint64_t get_misses = 0; + uint64_t get_validation_failures = 0; + // What the trace says the recorded system answered, so a replay can be read + // against the run it came from instead of being assumed to start from a + // perfect hit rate. + uint64_t recorded_misses = 0; + uint64_t recorded_failures = 0; + DistributionMetrics latency; + DistributionMetrics schedule_lag; +}; + +struct WorkloadRunnerOptions { + // Ignore trace timestamps and issue each client's operations immediately. + bool max_throughput = false; + // Trace intervals are multiplied by this value. Must be finite and positive. + double time_scale = 1.0; + bool validate_get_payloads = true; +}; + +class WorkloadRunner { + public: + explicit WorkloadRunner(WorkloadClient* client, WorkloadRunnerOptions options = {}); + + WorkloadMetrics Run(WorkloadSource* source); + + private: + using Clock = std::chrono::steady_clock; + + WorkloadClient* client_; + WorkloadRunnerOptions options_; +}; + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/include/umbp/distributed/benchmark/workload_source.h b/src/umbp/include/umbp/distributed/benchmark/workload_source.h new file mode 100644 index 000000000..c2d720202 --- /dev/null +++ b/src/umbp/include/umbp/distributed/benchmark/workload_source.h @@ -0,0 +1,110 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/workload_trace.h" + +namespace mori::umbp::benchmark { + +class WorkloadSource { + public: + virtual ~WorkloadSource() = default; + virtual uint64_t seed() const = 0; + virtual bool Next(::umbp::benchmark::WorkloadEvent* event) = 0; +}; + +class TraceWorkloadSource final : public WorkloadSource { + public: + explicit TraceWorkloadSource(const std::string& path); + + uint64_t seed() const override { return reader_.header().seed(); } + bool Next(::umbp::benchmark::WorkloadEvent* event) override; + const ::umbp::benchmark::WorkloadTraceHeader& header() const { return reader_.header(); } + + private: + TraceReader reader_; +}; + +enum class SyntheticProfile { + kSequential, + kUniform, + kHotsetZipf, + kReadAfterWrite, + kMixed, + kCapacityPressure, +}; + +enum class ValueSizeDistribution { + kFixed, + kUniform, + kLogUniform, +}; + +struct SyntheticWorkloadConfig { + SyntheticProfile profile = SyntheticProfile::kMixed; + uint64_t seed = 1; + uint64_t operation_count = 1000; + uint64_t key_count = 100; + size_t min_value_size = 4096; + size_t max_value_size = 4096; + ValueSizeDistribution value_size_distribution = ValueSizeDistribution::kFixed; + double read_ratio = 0.5; + uint32_t client_count = 1; + uint32_t batch_size = 1; + double qps = 0.0; + double hotset_fraction = 0.1; + double zipf_exponent = 1.1; + std::string key_prefix = "umbp-bench-"; +}; + +class SyntheticWorkloadSource final : public WorkloadSource { + public: + explicit SyntheticWorkloadSource(SyntheticWorkloadConfig config); + + uint64_t seed() const override { return config_.seed; } + bool Next(::umbp::benchmark::WorkloadEvent* event) override; + const SyntheticWorkloadConfig& config() const { return config_; } + + private: + uint64_t SelectKeyIndex(uint64_t sequence); + uint64_t SelectValueSize(); + bool SelectRead(uint64_t sequence); + + SyntheticWorkloadConfig config_; + std::mt19937_64 random_; + std::vector zipf_cdf_; + std::vector key_versions_; + std::vector last_put_operation_; + std::vector last_put_size_; + uint64_t generated_ = 0; +}; + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/include/umbp/distributed/benchmark/workload_trace.h b/src/umbp/include/umbp/distributed/benchmark/workload_trace.h new file mode 100644 index 000000000..08ed5e6ef --- /dev/null +++ b/src/umbp/include/umbp/distributed/benchmark/workload_trace.h @@ -0,0 +1,80 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include "umbp_workload.pb.h" + +namespace mori::umbp::benchmark { + +inline constexpr uint32_t kWorkloadTraceSchemaVersion = 1; + +// Errors caused by I/O, a malformed trace, or an unsupported version. +class TraceError : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +class TraceWriter { + public: + TraceWriter(const std::string& path, const ::umbp::benchmark::WorkloadTraceHeader& header); + ~TraceWriter(); + + TraceWriter(const TraceWriter&) = delete; + TraceWriter& operator=(const TraceWriter&) = delete; + + void Write(const ::umbp::benchmark::WorkloadEvent& event); + void Close(); + + private: + void WriteRecord(const google::protobuf::MessageLite& record, size_t max_bytes); + + std::ofstream stream_; + bool closed_ = false; +}; + +class TraceReader { + public: + explicit TraceReader(const std::string& path); + + const ::umbp::benchmark::WorkloadTraceHeader& header() const { return header_; } + + // Returns false only at a clean record boundary at end-of-file. + bool ReadNext(::umbp::benchmark::WorkloadEvent* event); + + private: + bool ReadRecord(google::protobuf::MessageLite* record, size_t max_bytes, bool allow_eof); + + std::ifstream stream_; + ::umbp::benchmark::WorkloadTraceHeader header_; +}; + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/include/umbp/distributed/benchmark/workload_trace_recorder.h b/src/umbp/include/umbp/distributed/benchmark/workload_trace_recorder.h new file mode 100644 index 000000000..697c571e3 --- /dev/null +++ b/src/umbp/include/umbp/distributed/benchmark/workload_trace_recorder.h @@ -0,0 +1,101 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/workload_trace.h" + +namespace mori::umbp::benchmark { + +struct WorkloadTraceRecorderOptions { + std::string path; + uint32_t client_id = 0; + uint64_t seed = 0; + std::string node_id; + std::string backend_policy; +}; + +// How the recorded system answered an operation. A miss is distinct from a +// failure because only one of them says something about the cache. +enum class WorkloadTraceOutcome { + kSuccess, + kMiss, + kFailure, +}; + +class WorkloadTraceRecorder { + public: + using Options = WorkloadTraceRecorderOptions; + + explicit WorkloadTraceRecorder(Options options); + ~WorkloadTraceRecorder(); + + WorkloadTraceRecorder(const WorkloadTraceRecorder&) = delete; + WorkloadTraceRecorder& operator=(const WorkloadTraceRecorder&) = delete; + + // Records every attempted operation, keys verbatim. A trace that keeps only + // the operations that succeeded, or that rewrites keys to make each write + // unique, is no longer the workload that ran: the reuse and overwrite + // patterns that decide what a tier policy does are exactly what those + // transformations remove. + void RecordBatchPut(const std::vector& keys, const std::vector& sizes, + const std::vector& outcomes) { + RecordBatch(::umbp::benchmark::WorkloadEvent::PUT, keys, sizes, outcomes); + } + void RecordBatchGet(const std::vector& keys, const std::vector& sizes, + const std::vector& outcomes) { + RecordBatch(::umbp::benchmark::WorkloadEvent::GET, keys, sizes, outcomes); + } + void Close(); + uint64_t event_count() const; + + private: + using Clock = std::chrono::steady_clock; + + void RecordBatch(::umbp::benchmark::WorkloadEvent::Operation operation, + const std::vector& keys, const std::vector& sizes, + const std::vector& outcomes); + + uint64_t RelativeTimestampNs() const; + + bool WriteLocked(const ::umbp::benchmark::WorkloadEvent& event); + + const uint32_t client_id_; + mutable std::mutex mutex_; + TraceWriter writer_; + const Clock::time_point start_; + uint64_t next_operation_id_ = 0; + uint64_t next_batch_id_ = 0; + uint64_t event_count_ = 0; + bool failed_ = false; + std::string error_; +}; + +} // namespace mori::umbp::benchmark diff --git a/src/umbp/include/umbp/distributed/config.h b/src/umbp/include/umbp/distributed/config.h index 1887d5663..0c57f5e3c 100644 --- a/src/umbp/include/umbp/distributed/config.h +++ b/src/umbp/include/umbp/distributed/config.h @@ -32,6 +32,7 @@ #include "umbp/common/config.h" #include "umbp/common/env_time.h" +#include "umbp/distributed/pool/policy_config.h" #include "umbp/distributed/types.h" namespace mori::umbp { @@ -41,6 +42,23 @@ class RouteGetStrategy; class RoutePutStrategy; class MasterEvictStrategy; +// The user-facing medium selector, lowered to the tier id the routing plane +// and the wire already speak. Two enums rather than one because the +// dependency is one-directional (distributed/config.h -> common/config.h) and +// UMBPMedium must be nameable by callers that never include the distributed +// headers. +inline TierType ToTierType(UMBPMedium medium) { + switch (medium) { + case UMBPMedium::HBM: + return TierType::HBM; + case UMBPMedium::SSD: + return TierType::SSD; + case UMBPMedium::DRAM: + break; + } + return TierType::DRAM; +} + struct ClientRegistryConfig { std::chrono::seconds heartbeat_ttl{10}; std::chrono::seconds reaper_interval{5}; @@ -66,6 +84,14 @@ struct ClientRegistryConfig { // (UMBPDistributedConfig / PoolClientConfig) default their // `dram_page_size` to 0 and rely on this value to materialize. uint64_t default_dram_page_size = 2ULL * 1024 * 1024; // 2 MiB + + // Current Masters advertise this so weighted peers may heartbeat aggregate + // same-tier capacity together with max_allocatable_bytes. A pre-weighted + // Master leaves it false (proto3 default); peers then keep heartbeats at the + // first instance per tier so rolling upgrades cannot over-admit a value no + // single backend can hold. Tests set this false to emulate that Master. + bool advertise_max_allocatable_bytes = true; + bool advertise_logical_tiers = true; }; struct EvictionConfig { @@ -131,9 +157,34 @@ struct MasterServerConfig { MasterServerConfig& operator=(MasterServerConfig&&) noexcept; }; -struct ExportableDram { - void* buffer = nullptr; - size_t size = 0; +// Ownership config for this node's DRAM/HBM backend pool(s) — sizes only. +// PageBackend::Init(TransferEngine*) self-allocates one HostMemAllocator +// buffer per entry in `buffer_sizes` and registers it with the transfer +// engine; PoolClientConfig never sees a buffer pointer (backend-agnostic +// refactor Phase 2b — see design-backend-agnostic-refactor.md §1 item 4). +// `buffer_sizes` usually holds exactly one entry (one pool); more than one +// exercises PageBitmapAllocator's cross-buffer scatter/gather strategy. +struct DramOwnershipConfig { + std::vector buffer_sizes; + bool use_hugepages = false; + uint64_t hugepage_size = 2ULL * 1024 * 1024; + int numa_node = -1; + bool prefault = true; +}; + +// HBM-tier ownership knobs. Deliberately NOT a superset of the DRAM ones: +// hipMalloc has no hugepage, NUMA or prefault dimension, and the device ordinal +// it does have is meaningless for host memory. That asymmetry is the reason +// each medium brings its own PageMemorySource rather than sharing one options +// struct with per-medium fields nobody else reads. +// +// No `enabled` flag: PoolClientConfig::medium selects the one live medium, and +// these knobs are read only when it names HBM. +struct HbmOwnershipConfig { + // GPU ordinal the pool is allocated on. Fixed at Init, not inherited from + // whichever thread allocates first (see HbmPageMemorySource). + int device = 0; + std::vector buffer_sizes; }; // SSD-tier construction parameters lowered from the user-facing UMBPConfig. @@ -142,22 +193,61 @@ struct ExportableDram { // peer only needs that subset — not the whole global config. ssd_backend // (posix / spdk / spdk_proxy) lives inside UMBPSsdConfig, so PeerSsdManager // picks the backend from cfg.ssd directly. +// +// `enabled` stays because PeerSsdManager's own contract is written against it +// (it refuses to open a tier that is off); PoolClient sets it from +// PoolClientConfig::medium rather than from user config. struct PeerSsdConfig { bool enabled = false; UMBPSsdConfig ssd; }; +enum class PoolPlacementPolicy { + SINGLE_BACKEND = 0, + WEIGHTED = 1, + TIERED = 2, +}; + +// One named backend instance on a peer. Only the ownership block selected by +// `tier` is read. Page size remains peer-global because the RPC wire carries a +// single page_size for batches that may span instances. +struct BackendInstanceConfig { + std::string name; + TierType tier = TierType::UNKNOWN; + DramOwnershipConfig dram; + HbmOwnershipConfig hbm; + PeerSsdConfig ssd; + int ssd_staging_buffer_slots = 16; + // Relative admission share within this tier when WEIGHTED placement is + // enabled. Must be positive; legacy and SINGLE_BACKEND paths ignore it. + // Appended to preserve existing positional aggregate initializers. + uint32_t placement_weight = 1; +}; + struct PoolClientConfig { UMBPMasterClientConfig master_config; UMBPIoEngineConfig io_engine; size_t staging_buffer_size = 64ULL * 1024 * 1024; - // SSD read-staging tuning (peer side). More slots reduce NO_SLOT under large - // concurrent prefetch batches, but shrink per-slot size (= staging_buffer_size - // / slots), which must stay >= the largest single SSD block. The slot lease - // TTL (the primary slot-reclaim mechanism; ReleaseSsdLease is best-effort) is - // resolved from UMBP_SSD_READ_LEASE_MS at PeerService construction, not here. + // Caller-owned, RDMA-registered host arenas for ranged I/O. DistributedClient + // owns the backing mappings and frees them only after PoolClient::Shutdown has + // deregistered the regions. Direct PoolClient tests may provide their own. + // + // Two separate arenas — one for remote ranged GET, one for remote ranged PUT + // — each under its own mutex in PoolClient, so a remote get and a remote put + // run concurrently instead of serializing on one lock (the load/offload + // overlap sglang's direct linker wants). Each must hold at least one whole + // object. Both zero keeps ranged remote I/O disabled with no host memory + // registered; SupportsRangedIO() requires both to be set. + void* ranged_get_scratch_buffer = nullptr; + size_t ranged_get_scratch_size = 0; + void* ranged_put_scratch_buffer = nullptr; + size_t ranged_put_scratch_size = 0; + + // SSD read-staging tuning. More slots increase concurrent reads and writes; + // each policy-created SsdBackend owns `slots * page_size` staging bytes. Its + // read lease is resolved from UMBP_SSD_READ_LEASE_MS during PoolClient::Init. int ssd_staging_buffer_slots = 16; // Backs ssd_staging_buffer_, allocated only when ssd.enabled. A remote SSD @@ -165,10 +255,24 @@ struct PoolClientConfig { // must be >= the largest single-key page KV (61-layer MLA page ~= 4.5 MB). size_t ssd_staging_buffer_size = 268435456; // 256 MiB - std::vector dram_buffers; + // Back the SSD staging arena with hugetlbfs pages. Every SSD backend this + // client registers uses the same node-wide setting. Falls back to 4 KiB + // pages when no hugetlb pages are free so the node can still start. + bool ssd_staging_use_hugepages = false; + size_t ssd_staging_hugepage_size = 2ULL * 1024 * 1024; // 2 MiB + + // Legacy one-medium selector. Used only when `backends` is empty; in that + // mode PoolClient::Init reads exactly one ownership block below. + TierType medium = TierType::DRAM; + + DramOwnershipConfig dram; + HbmOwnershipConfig hbm; PeerSsdConfig ssd; - std::map tier_capacities; + // Named multi-backend configuration. Empty selects the legacy fields above; + // PoolClient lowers them at Init time so mutations made after + // ToPoolClientConfig() remain authoritative. + std::vector backends; uint16_t peer_service_port = 0; @@ -178,6 +282,17 @@ struct PoolClientConfig { CacheRemoteAdmission cache_remote_admission = CacheRemoteAdmission::SIZE; size_t admission_max_block_bytes = 16ULL * 1024 * 1024; + // Background whole-object pull after a remote ranged read; see + // UMBPDistributedConfig::ranged_locality_prefetch for why it is not folded + // into cache_remote_fetches. + bool ranged_locality_prefetch = true; + + // Resolve on this node's own media before asking the master to route, and + // skip the master entirely when the whole batch is local. See + // UMBPDistributedConfig::local_first for the reasoning and the multi-node + // trade-off. + bool local_first = true; + // Page size used by Master's PageBitmapAllocator for this node's DRAM/HBM // tier. Reported via RegisterClient. Same value applies to both DRAM // and HBM. Forwarded unmodified to MasterClient::RegisterSelf by @@ -186,40 +301,100 @@ struct PoolClientConfig { // (2 MiB by default). Set to an explicit byte count to override. uint64_t dram_page_size = 0; + // Watermarks for the paged medium's own eviction, used ONLY when this node + // has no master (PoolClient::Init enables it exactly then). With a master, + // eviction is a cluster-wide decision shipped as Evict RPCs and these are + // ignored; without one, nothing else would ever free a page. Lowered from + // UMBPDramConfig::high_watermark / low_watermark, which is where the + // deleted local backend read the same policy from. The SSD medium has its + // own pair inside UMBPSsdConfig and does not consult these. + double local_evict_high_watermark = 0.9; + double local_evict_low_watermark = 0.7; + UMBPCopyPipelineConfig copy_pipeline = [] { UMBPCopyPipelineConfig c; c.worker_threads = 1; return c; }(); + + // Compatibility is the default. WEIGHTED distributes new keys among all + // configured same-tier instances according to placement_weight. Appended to + // preserve existing positional aggregate initializers. + PoolPlacementPolicy placement_policy = PoolPlacementPolicy::SINGLE_BACKEND; + + // Opt-in ephemeral peer service for test/benchmark harnesses. When true, + // PoolClient starts the service even when peer_service_port is zero and + // advertises the port selected atomically by gRPC, avoiding probe/bind races. + bool auto_peer_service_port = false; + + // Optional declarative policy. policy_config_path is loaded at Init and + // lowered into named backends plus logical_tiers. Callers that load JSON + // themselves may populate logical_tiers directly. Empty fields preserve the + // legacy SINGLE_BACKEND/WEIGHTED behavior. Appended for aggregate-init + // compatibility. + std::string policy_config_path; + std::vector logical_tiers; + + // Optional production traffic recorder. Empty path disables recording. + std::string workload_trace_path; + uint32_t workload_trace_client_id = 0; + uint64_t workload_trace_seed = 0; }; // Lower a user-facing UMBPDistributedConfig to the internal PoolClientConfig. // Kept as a free function (not a member of UMBPDistributedConfig) so that // common/config.h does not need to include distributed/config.h — the // dependency is one-directional: distributed/config.h -> common/config.h. -// DRAM buffers and tier capacities are caller-supplied because they live in -// DistributedClient (pool mmap'd memory), not in the user-facing config. +// `dram` is caller-supplied (sizes only, no pointer — see DramOwnershipConfig) +// because DistributedClient is the one place that knows the top-level +// UMBPConfig::dram ownership knobs (hugepages/numa/prefault), which sit +// beside — not inside — UMBPDistributedConfig. PoolClient's DRAM PageBackend +// self-allocates from this at Init(); tier capacities are no longer +// caller-supplied at all — they are derived from BackendRegistry::Capacity() +// after Init (backend-agnostic refactor Phase 2). inline PoolClientConfig ToPoolClientConfig(const UMBPDistributedConfig& dc, - std::vector dram_buffers, - std::map tier_capacities, - PeerSsdConfig ssd = {}) { + DramOwnershipConfig dram, PeerSsdConfig ssd = {}, + double dram_high_watermark = 0.9, + double dram_low_watermark = 0.7) { PoolClientConfig pc; pc.master_config = dc.master_config; pc.io_engine = dc.io_engine; pc.staging_buffer_size = dc.staging_buffer_size; + // The ranged scratch buffers/sizes are set by DistributedClient after it + // allocates and registers the two arenas (see DistributedClient ctor). pc.ssd_staging_buffer_size = dc.ssd_staging_buffer_size; pc.ssd_staging_buffer_slots = dc.ssd_staging_buffer_slots; + pc.ssd_staging_use_hugepages = dc.ssd_staging_use_hugepages; + pc.ssd_staging_hugepage_size = dc.ssd_staging_hugepage_size; pc.peer_service_port = dc.peer_service_port; pc.cache_remote_fetches = dc.cache_remote_fetches; pc.cache_remote_admission = dc.cache_remote_admission; pc.admission_max_block_bytes = dc.admission_max_block_bytes; + pc.ranged_locality_prefetch = dc.ranged_locality_prefetch; + pc.local_first = dc.local_first; // 0 propagates through PoolClient -> MasterClient::RegisterSelf -> // proto -> ClientRegistry, where it is interpreted as "use the // registry-wide default_dram_page_size". pc.dram_page_size = dc.dram_page_size; - pc.dram_buffers = std::move(dram_buffers); - pc.tier_capacities = std::move(tier_capacities); + pc.policy_config_path = dc.backend_policy_path; + pc.workload_trace_path = dc.workload_trace_path; + pc.workload_trace_client_id = dc.workload_trace_client_id; + pc.workload_trace_seed = dc.workload_trace_seed; + pc.dram = std::move(dram); pc.ssd = std::move(ssd); + // Unlike dram/ssd, UMBPHbmConfig carries no ownership knobs that live + // outside UMBPDistributedConfig (no hugepages/NUMA/prefault dimension for + // hipMalloc'd memory), so it can be lowered directly here instead of via a + // caller-supplied parameter. + pc.hbm.device = dc.hbm.device; + if (dc.hbm.capacity_bytes > 0) pc.hbm.buffer_sizes = {dc.hbm.capacity_bytes}; + pc.local_evict_high_watermark = dram_high_watermark; + pc.local_evict_low_watermark = dram_low_watermark; + pc.medium = ToTierType(dc.medium); + // The medium is what opts the SSD tier in; PeerSsdManager keys off this. + pc.ssd.enabled = (pc.medium == TierType::SSD); + // Deliberately leave pc.backends empty. PoolClient synthesizes the legacy + // instance at Init so callers may still mutate medium/dram/hbm/ssd afterward. return pc; } diff --git a/src/umbp/include/umbp/distributed/distributed_client.h b/src/umbp/include/umbp/distributed/distributed_client.h index 2d39b50be..d5d6f8821 100644 --- a/src/umbp/include/umbp/distributed/distributed_client.h +++ b/src/umbp/include/umbp/distributed/distributed_client.h @@ -35,9 +35,13 @@ namespace mori::umbp { -/// Distributed IUMBPClient implementation — master-led global routing -/// with RDMA/MORI-IO data plane. All routing decisions go through the -/// Master; this client does not use LocalStorageManager or LocalBlockIndex. +/// The IUMBPClient implementation: this node's medium, plus master-led global +/// routing over an RDMA/MORI-IO data plane when a master address is configured. +/// +/// With none it is an embedded, single-process store — the same backends and +/// the same transfer engine, with nothing routed, registered or heartbeated, +/// and a local miss as the final answer. Configuration alone decides which, +/// which is why there is no separate local implementation. class DistributedClient : public IUMBPClient { public: explicit DistributedClient(const UMBPConfig& config); @@ -58,6 +62,19 @@ class DistributedClient : public IUMBPClient { std::vector BatchGet(const std::vector& keys, const std::vector& dsts, const std::vector& sizes) override; + // Not implemented — see src/umbp/doc/design-tree-connector-port.md §5. The + // transfer layer can already express a range (TransferItem carries + // src_offset/dst_offset/size); what is missing is the object-range to + // page-range mapping and the master-side metadata question. + std::vector BatchGetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& src_offsets) override; + std::vector BatchPutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& dst_offsets) override; std::vector BatchExists(const std::vector& keys) const override; size_t BatchExistsConsecutive(const std::vector& keys) const override; @@ -66,8 +83,24 @@ class DistributedClient : public IUMBPClient { void Close() override; bool IsDistributed() const override; UMBPDeploymentMode GetDeploymentMode() const override { return UMBPDeploymentMode::Distributed; } + // One condition: the scratch arena must exist, because the remote direction + // has nowhere to land otherwise — that is upstream's opt-in rule, and it + // defaults to off. + // + // There is deliberately no medium condition. This used to also exclude SSD, + // on the reasoning that ranged I/O maps object ranges onto pages a backend + // publishes as in-process endpoints and SSD publishes storage refs instead. + // That describes a backend nobody built: ssd_backend.h weighs a file endpoint + // against staging and picks staging, so SsdBackend publishes ordinary + // registered host pages like every other medium. Every ranged path works on + // an SSD node. See the definition for what the medium does still change (the + // device read is whole-object until D1) and doc/design-ssd-ranged-io.md. + bool SupportsRangedIO() const override; - bool RegisterMemory(uintptr_t ptr, size_t size) override; + bool RegisterMemory(uintptr_t ptr, size_t size, + mori::io::MemoryLocationType loc = mori::io::MemoryLocationType::CPU, + int device = -1, + MemoryRegistration mode = MemoryRegistration::kPinned) override; void DeregisterMemory(uintptr_t ptr) override; bool ReportExternalKvBlocks(const std::vector& hashes, TierType tier) override; @@ -80,9 +113,27 @@ class DistributedClient : public IUMBPClient { private: UMBPConfig config_; - void* dram_pool_ = nullptr; - size_t dram_pool_size_ = 0; - HostBufferHandle dram_pool_handle_; + // The only buffers this class still allocates. Phase 2b moved medium pools + // into the backends, but the ranged scratch arenas are not a medium: they are + // client-side staging regions for objects fetched from, or assembled for, + // another node, so they belong to whoever owns the PoolClient. + // + // Two separate arenas — one for remote ranged GET, one for remote ranged PUT + // — each RDMA-registered and each under its own mutex in PoolClient, so a + // remote get and a remote put run concurrently instead of serializing on one + // lock (the load/offload overlap sglang's direct linker wants). Each is sized + // to config.distributed.ranged_scratch_size; allocated only when that is > 0. + void* ranged_get_scratch_ = nullptr; + size_t ranged_get_scratch_size_ = 0; + HostBufferHandle ranged_get_scratch_handle_; + void* ranged_put_scratch_ = nullptr; + size_t ranged_put_scratch_size_ = 0; + HostBufferHandle ranged_put_scratch_handle_; + // No master address => no routing => this node can never hold, fetch or + // serve a remote key. Set once at construction; read by SupportsRangedIO, + // which is the one place the distinction changes an answer rather than just + // a code path. + bool local_only_ = false; std::unique_ptr pool_client_; std::atomic closing_{false}; mutable std::shared_mutex op_mutex_; diff --git a/src/umbp/include/umbp/distributed/master/in_memory_master_metadata_store.h b/src/umbp/include/umbp/distributed/master/in_memory_master_metadata_store.h index 20f1118be..95b7f2199 100644 --- a/src/umbp/include/umbp/distributed/master/in_memory_master_metadata_store.h +++ b/src/umbp/include/umbp/distributed/master/in_memory_master_metadata_store.h @@ -85,10 +85,11 @@ class InMemoryMasterMetadataStore : public IMasterMetadataStore { std::chrono::system_clock::time_point now, std::chrono::system_clock::duration stale_after) override; void UnregisterClient(const std::string& node_id) override; - HeartbeatResult ApplyHeartbeat(const std::string& node_id, uint64_t seq, - std::chrono::system_clock::time_point now, - const std::map& caps, - const std::vector& events, bool is_full_sync) override; + HeartbeatResult ApplyHeartbeat( + const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, + const std::map& caps, const std::vector& events, + bool is_full_sync, + const std::map& logical_caps = {}) override; std::vector ExpireStaleClients( std::chrono::system_clock::time_point cutoff) override; diff --git a/src/umbp/include/umbp/distributed/master/master_client.h b/src/umbp/include/umbp/distributed/master/master_client.h index 4d329e05b..0b5eb3d7d 100644 --- a/src/umbp/include/umbp/distributed/master/master_client.h +++ b/src/umbp/include/umbp/distributed/master/master_client.h @@ -42,7 +42,7 @@ #include #include "umbp/distributed/config.h" -#include "umbp/distributed/peer/owned_location_source.h" +#include "umbp/distributed/peer/backend/medium_backend.h" #include "umbp/distributed/routing/route_put_strategy.h" #include "umbp/distributed/types.h" @@ -54,8 +54,7 @@ class HeartbeatResponse; namespace mori::umbp { -class PeerDramAllocator; -class PeerSsdManager; +class PeerPool; inline constexpr std::size_t kMasterClientMaxPendingHistograms = 15000; @@ -68,6 +67,7 @@ struct RouteGetResult { TierType tier = TierType::UNKNOWN; uint64_t size = 0; std::string peer_address; + std::string logical_tier; }; class MasterClient { @@ -86,9 +86,10 @@ class MasterClient { // membership + capacity-snapshot metadata is shipped — DRAM/HBM // descriptors are peer-internal now. `tier_capacities` is the single // source of per-tier capacity, including SSD (TierType::SSD). - grpc::Status RegisterSelf(const std::map& tier_capacities, - const std::string& peer_address = "", - const std::vector& engine_desc_bytes = {}); + grpc::Status RegisterSelf( + const std::map& tier_capacities, const std::string& peer_address = "", + const std::vector& engine_desc_bytes = {}, + const std::map& logical_tier_capacities = {}); grpc::Status UnregisterSelf(); // --- Router --- @@ -122,24 +123,26 @@ class MasterClient { grpc::Status BatchLookup(const std::vector& keys, std::vector* out); // --- Heartbeat (event-driven) --- - // Bind a PeerDramAllocator whose outbox the heartbeat thread will - // drain. Pass nullptr for SSD-only peers (skipped, not registered). - // Also keeps the concrete pointer for DRAM-specific duties the - // OwnedLocationSource interface does not cover (capacity snapshot, - // owned-key counts, distributed-clear write gate). Must be set before - // StartHeartbeat() — the heartbeat thread reads sources once per tick. - void SetPeerDramAllocator(PeerDramAllocator* dram_alloc); - - // Bind the SSD manager: registers it as an owned-location event source - // AND keeps the concrete pointer so heartbeat can merge live SSD - // capacity into tier_capacities. Pass nullptr to skip. Must be set + // Bind every storage medium live on this peer. The heartbeat drains each + // backend's outbox and concatenates them into ONE bundle under a single + // monotonic seq, and derives tier_capacities / tier_kv_counts by aggregating + // over the same set — so a backend added to the registry participates with + // no change here (backend-agnostic refactor Phase 3). Also installs the + // auto-flush hook on each backend. Pass nullptr to unbind. Must be set + // before StartHeartbeat() — the heartbeat thread reads the registry once per + // tick. + void SetBackendRegistry(BackendRegistry* registry); + void SetPeerPool(PeerPool* pool) { peer_pool_ = pool; } + + // Weighted placement makes every same-tier instance reachable through one + // tier route, so capacity heartbeats may aggregate them. Must be configured // before StartHeartbeat(). - void SetPeerSsdManager(PeerSsdManager* ssd_manager); - - // Register an additional owned-location event source whose events are - // drained/snapshotted into the same heartbeat bundle (single monotonic - // seq). Null is ignored. Must be set before StartHeartbeat(). - void AddOwnedLocationSource(OwnedLocationSource* source); + void SetAggregateBackendCapacities(bool enabled) { + aggregate_backend_capacities_requested_ = enabled; + aggregate_backend_capacities_ = enabled && supports_max_allocatable_capacity_; + } + bool SupportsMaxAllocatableCapacity() const { return supports_max_allocatable_capacity_; } + bool SupportsLogicalTiers() const { return supports_logical_tiers_; } void StartHeartbeat(); void StopHeartbeat(); @@ -147,13 +150,12 @@ class MasterClient { void FlushHeartbeat(); // Synchronously clear this node's UMBP-owned and external HiCache - // placement from master. When peer_alloc_ is bound, caller MUST have - // already invoked peer_alloc_->ClearLocal() (PoolClient::Clear() - // handles this); SSD-only peers skip that step. Returns true only - // after both the full-sync heartbeat and the external KV revoke RPC - // are acknowledged by master; on any failure returns false and leaves - // PeerDramAllocator::clear_full_sync_pending_ closed for the caller - // to retry. + // placement from master. When a registry is bound, caller MUST have already + // invoked ClearLocal() on every backend (PoolClient::Clear() handles this); + // SSD-only peers skip that step. Returns true only after both the full-sync + // heartbeat and the external KV revoke RPC are acknowledged by master; on any + // failure returns false and leaves each backend's clear-full-sync gate closed + // for the caller to retry. bool ClearFullSync(); // --- Client-side metrics --- @@ -226,27 +228,26 @@ class MasterClient { std::mutex hb_cv_mutex_; std::condition_variable hb_cv_; - // Cached tier capacities — heartbeat reports the latest peer - // allocator snapshot when the bound PeerDramAllocator is non-null, - // else falls back to whatever was last set here. + // Cached tier capacities — heartbeat reports the live per-backend snapshot + // when a registry is bound, else falls back to whatever was last set here. std::mutex caps_mutex_; std::map current_capacities_; - // Peer-event source for the heartbeat thread. Non-owning; lifetime - // is managed by PoolClient. Kept as a concrete pointer (in addition to - // its slot in owned_sources_) for DRAM-specific duties the - // OwnedLocationSource interface does not cover: TierCapacitiesSnapshot, - // OwnedKeyCountByTier, ClearLocal/ClearFullSyncAcked. - PeerDramAllocator* peer_alloc_ = nullptr; - - // Non-owning. Kept concrete (in addition to owned_sources_) so heartbeat - // can merge live SSD capacity into tier_capacities. - PeerSsdManager* ssd_manager_ = nullptr; - - // All owned-location event sources (DRAM allocator + SSD manager). Drained - // and snapshotted into a single heartbeat bundle under one monotonic seq. - // Populated before StartHeartbeat(); read-only afterwards (no lock needed). - std::vector owned_sources_; + // Every storage medium on this peer: the heartbeat's event sources AND the + // source of tier_capacities / tier_kv_counts / the distributed-clear write + // gate. Non-owning; lifetime is managed by PoolClient. Bound before + // StartHeartbeat(); read-only afterwards (no lock needed). + BackendRegistry* registry_ = nullptr; + PeerPool* peer_pool_ = nullptr; + bool aggregate_backend_capacities_requested_ = false; + bool aggregate_backend_capacities_ = false; + bool supports_max_allocatable_capacity_ = false; + bool supports_logical_tiers_ = false; + // Initial conservative registration payload, retained so an UNKNOWN response + // can rebuild the full peer record after Master state loss. + std::map registration_capacities_; + std::string registered_peer_address_; + std::vector registered_engine_desc_bytes_; // Serializes the actual Heartbeat RPC: at most one full-sync or // delta heartbeat is on the wire at a time. ClearFullSync() takes @@ -274,13 +275,18 @@ class MasterClient { ::umbp::HeartbeatResponse* resp); // Snapshot the freshest tier capacities and refresh the - // current_capacities_ fallback cache. When peer_alloc_ is bound, - // returns its bitmap-derived snapshot (DRAM/HBM) merged with cached - // entries for tiers it does not manage (e.g. SSD). When peer_alloc_ - // is null, returns the cached snapshot unchanged. Acquires - // caps_mutex_ internally. + // current_capacities_ fallback cache. When a registry is bound, returns each + // backend's live (bitmap-derived) capacity merged with cached entries for + // tiers no backend manages (e.g. SSD). When it is null, returns the cached + // snapshot unchanged. Acquires caps_mutex_ internally. std::map SnapshotAndCacheTierCapacities(); + // Owned-key count per tier, aggregated over every registered backend. + std::map OwnedKeyCountsByTier() const; + + // Every live backend, or empty when no registry is bound. + std::vector Backends() const; + // --- Metrics buffering --- struct PendingSample { std::string name; diff --git a/src/umbp/include/umbp/distributed/master/master_metadata_store.h b/src/umbp/include/umbp/distributed/master/master_metadata_store.h index ba770e764..8ea1ff32d 100644 --- a/src/umbp/include/umbp/distributed/master/master_metadata_store.h +++ b/src/umbp/include/umbp/distributed/master/master_metadata_store.h @@ -319,10 +319,10 @@ class IMasterMetadataStore { // (e.g. 100k events) here, and have the peer fragment larger // resyncs into multiple full_sync calls or shift to a snapshot- // then-delta protocol before the Redis backend ships. - virtual HeartbeatResult ApplyHeartbeat(const std::string& node_id, uint64_t seq, - std::chrono::system_clock::time_point now, - const std::map& caps, - const std::vector& events, bool is_full_sync) = 0; + virtual HeartbeatResult ApplyHeartbeat( + const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, + const std::map& caps, const std::vector& events, + bool is_full_sync, const std::map& logical_caps = {}) = 0; // Reaper pass. Atomically: // - Flips status ALIVE → EXPIRED for every record with diff --git a/src/umbp/include/umbp/distributed/master/master_metrics.h b/src/umbp/include/umbp/distributed/master/master_metrics.h index 0e58351fd..cd4f7b2dc 100644 --- a/src/umbp/include/umbp/distributed/master/master_metrics.h +++ b/src/umbp/include/umbp/distributed/master/master_metrics.h @@ -115,6 +115,72 @@ #define MORI_UMBP_METRIC_CLIENT_CAPACITY_UTILIZATION_HELP \ "Capacity utilization ratio (used / total) in [0,1]" +// --- Per-client logical tier capacity gauges ------------------------------- +// The capacity gauges above key on the medium, so a policy that places two +// DRAM backends in different logical tiers reports both as a single DRAM +// series and the hierarchy stops being readable: per-tier occupancy, whether a +// watermark drain reached its low watermark, and which tiers accept new PUTs +// all disappear into the sum. These carry the tier name the policy declared. +// Labels: node=, logical_tier=, tier= + +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_TOTAL \ + "mori_umbp_client_logical_tier_capacity_total_bytes" +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_TOTAL_HELP \ + "Total capacity bytes of a logical tier on this client's pool" + +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_AVAIL \ + "mori_umbp_client_logical_tier_capacity_available_bytes" +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_AVAIL_HELP \ + "Available capacity bytes of a logical tier on this client's pool" + +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_USED \ + "mori_umbp_client_logical_tier_capacity_used_bytes" +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_USED_HELP \ + "Used capacity bytes of a logical tier (total - available)" + +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_UTILIZATION \ + "mori_umbp_client_logical_tier_capacity_utilization_ratio" +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_CAPACITY_UTILIZATION_HELP \ + "Utilization ratio (used / total) of a logical tier, in [0,1]" + +// A tier that stopped accepting PUTs is indistinguishable from one nothing +// happens to be writing to unless the eligibility itself is published. +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_PUT_ELIGIBLE "mori_umbp_client_logical_tier_put_eligible" +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_PUT_ELIGIBLE_HELP \ + "1 when a logical tier currently accepts new PUTs, 0 when it does not" + +// The ratio the offload decision reads. utilization_ratio above aggregates the +// tier's backends, and on the entry tier it aggregates every tier reachable by +// offload, so comparing either against a watermark is meaningless. This is the +// series to plot against high_watermark and low_watermark. +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_PEAK_UTILIZATION \ + "mori_umbp_client_logical_tier_peak_member_utilization_ratio" +#define MORI_UMBP_METRIC_CLIENT_LOGICAL_PEAK_UTILIZATION_HELP \ + "Highest utilization among a logical tier's backends, in [0,1]; what the watermarks compare" + +// --- Logical tier transitions ---------------------------------------------- +// Whether an offload or a promotion ever ran is otherwise only visible to a +// process that links the tier benchmark, so a served workload cannot tell a +// working tier graph from one whose every migration fails. + +#define MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS "mori_umbp_client_tier_transitions_total" +#define MORI_UMBP_METRIC_CLIENT_TIER_TRANSITIONS_HELP \ + "Logical tier transitions on this client's pool, by outcome" + +#define MORI_UMBP_METRIC_CLIENT_TIER_OFFLOADED_BYTES "mori_umbp_client_tier_offloaded_bytes_total" +#define MORI_UMBP_METRIC_CLIENT_TIER_OFFLOADED_BYTES_HELP "Bytes moved to a later tier by offload" + +#define MORI_UMBP_METRIC_CLIENT_TIER_PROMOTED_BYTES "mori_umbp_client_tier_promoted_bytes_total" +#define MORI_UMBP_METRIC_CLIENT_TIER_PROMOTED_BYTES_HELP \ + "Bytes moved to an earlier tier by promote-on-read" + +// Labelled by logical tier, so which tier served a read stays visible even +// though the per-backend labels collapse every instance of a medium into one +// series. +#define MORI_UMBP_METRIC_CLIENT_TIER_READ_HITS "mori_umbp_client_tier_read_hits_total" +#define MORI_UMBP_METRIC_CLIENT_TIER_READ_HITS_HELP \ + "Reads served by each logical tier of this client's pool" + // --- Per-client RPC call counters ------------------------------------------ // Full name: prefix + sanitized_node_id @@ -166,6 +232,11 @@ #define MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP \ "Total bytes delivered to this client (inbound reads) split by local/remote traffic" +#define MORI_UMBP_METRIC_RANGED_REMOTE_INSTALL_FAILURES_TOTAL \ + "mori_umbp_ranged_remote_install_failures_total" +#define MORI_UMBP_METRIC_RANGED_REMOTE_INSTALL_FAILURES_TOTAL_HELP \ + "Remote ranged objects that could not be synchronously installed in the local medium" + // --- Heartbeat / event-shipping counters (master-as-advisor) ---------------- #define MORI_UMBP_METRIC_HEARTBEAT_EVENTS_APPLIED_TOTAL "mori_umbp_heartbeat_events_applied_total" @@ -188,6 +259,27 @@ "BatchGet e2e call bandwidth in GiB/s (successful bytes only, split by client and local/remote " \ "traffic)" +// Ranged siblings. Kept as their own series rather than folded into the two +// above: a ranged call moves a SUBSET of each object, so mixing them would make +// the whole-object families' bytes-per-call meaningless. Bytes counted are the +// range bytes actually delivered to (or committed from) the caller's buffers, +// which is what the caller sees and what the sglang tree connector reports on +// its side -- note that a REMOTE ranged get still moves the whole object over +// the wire into the scratch arena, so its wire traffic exceeds what this +// histogram credits (mori_umbp_client_*bound_get_bytes_total covers that). + +#define MORI_UMBP_METRIC_CLIENT_BATCH_PUT_RANGES_BANDWIDTH \ + "mori_umbp_client_batch_put_ranges_bandwidth_gibps" +#define MORI_UMBP_METRIC_CLIENT_BATCH_PUT_RANGES_BANDWIDTH_HELP \ + "BatchPutRanges e2e call bandwidth in GiB/s (committed range bytes only, split by client and " \ + "local/remote traffic)" + +#define MORI_UMBP_METRIC_CLIENT_BATCH_GET_RANGES_BANDWIDTH \ + "mori_umbp_client_batch_get_ranges_bandwidth_gibps" +#define MORI_UMBP_METRIC_CLIENT_BATCH_GET_RANGES_BANDWIDTH_HELP \ + "BatchGetRanges e2e call bandwidth in GiB/s (delivered range bytes only, split by client and " \ + "local/remote traffic)" + // --- MasterClient -> MasterServer RPC latency (client-perceived) ----------- // Histogram of round-trip latency for every RPC method on the // MasterClient channel, reported by clients via ReportMetrics. Labels @@ -208,99 +300,14 @@ "Number of histogram observations dropped client-side because the pending buffer hit its cap " \ "(see kMasterClientMaxPendingHistograms in master_client.h)" -// --- SSD tier: copy / read / eviction / staging (peer-reported) ------------ -// All shipped via ReportMetrics from the owner peer (node=). Labels -// are low-cardinality fixed enums only — status / reason — never key / path / -// error string / lease id. Counters are accumulated peer-side as cheap relaxed -// atomics at the event point and converted to ReportMetrics deltas once per -// metrics flush tick (no per-key AddCounter on the commit / read hot paths). -// SSD capacity is NOT re-exported here: it rides heartbeat tier_capacities as -// mori_umbp_client_capacity_{used,total}_bytes{tier="SSD"} (master emits the -// tier label upper-cased, e.g. tier="SSD" — match that casing in queries). - -// copy-on-commit pipeline (ssd_copy_pipeline.cpp) -#define MORI_UMBP_METRIC_SSD_COPY_ENQUEUED_TOTAL "mori_umbp_ssd_copy_enqueued_total" -#define MORI_UMBP_METRIC_SSD_COPY_ENQUEUED_TOTAL_HELP \ - "SSD copy-on-commit tasks accepted into the async copy queue" - -#define MORI_UMBP_METRIC_SSD_COPY_SUCCEEDED_TOTAL "mori_umbp_ssd_copy_succeeded_total" -#define MORI_UMBP_METRIC_SSD_COPY_SUCCEEDED_TOTAL_HELP \ - "SSD copy-on-commit tasks that completed successfully: either the bytes were written to the " \ - "backend (emits ADD SSD) OR the content-addressed key was already resident (dedup fast path, " \ - "no write, no event). Not a count of physical writes — see " \ - "mori_umbp_ssd_copy_bytes_total for physical SSD write bytes." - -#define MORI_UMBP_METRIC_SSD_COPY_FAILED_TOTAL "mori_umbp_ssd_copy_failed_total" -#define MORI_UMBP_METRIC_SSD_COPY_FAILED_TOTAL_HELP \ - "SSD copy-on-commit tasks whose backend Write failed (no SSD copy, no event)" - -// label: reason=queue_full|stopped -#define MORI_UMBP_METRIC_SSD_COPY_DROPPED_TOTAL "mori_umbp_ssd_copy_dropped_total" -#define MORI_UMBP_METRIC_SSD_COPY_DROPPED_TOTAL_HELP \ - "SSD copy tasks dropped at enqueue (reason=queue_full when the bounded queue " \ - "was full, reason=stopped when the pipeline was stopped/quiescing)" - -// label: status=ok|not_found|no_slot|size_too_large|error -#define MORI_UMBP_METRIC_SSD_READ_TOTAL "mori_umbp_ssd_read_total" -#define MORI_UMBP_METRIC_SSD_READ_TOTAL_HELP \ - "SSD reads served by this peer's SSD tier, by outcome (covers local owner " \ - "reads and remote PrepareSsdRead). not_found = stale-route miss; no_slot = " \ - "staging slots exhausted (transient, not a miss)" - -// SSD IO byte counters. Bandwidth is derived in Grafana/Prometheus via -// rate([]) (bytes/s) — same convention as the client -// inbound/outbound byte counters. copy_bytes = bytes landed on SSD by a -// successful copy-on-commit Write; read_bytes = bytes served by a successful -// SSD read (local owner read or remote PrepareSsdRead). -#define MORI_UMBP_METRIC_SSD_COPY_BYTES_TOTAL "mori_umbp_ssd_copy_bytes_total" -#define MORI_UMBP_METRIC_SSD_COPY_BYTES_TOTAL_HELP \ - "Bytes written to the SSD tier by successful copy-on-commit. rate() = offered " \ - "SSD write throughput (bytes / wall-clock; a load signal, not device per-IO " \ - "bandwidth — it averages in idle gaps)" - -#define MORI_UMBP_METRIC_SSD_READ_BYTES_TOTAL "mori_umbp_ssd_read_bytes_total" -#define MORI_UMBP_METRIC_SSD_READ_BYTES_TOTAL_HELP \ - "Bytes read from the SSD tier by successful reads. rate() = offered SSD read " \ - "throughput (bytes / wall-clock; a load signal, not device per-IO bandwidth)" - -// eviction (peer_ssd_manager.cpp, local watermark + LRU) -#define MORI_UMBP_METRIC_SSD_EVICTION_ROUNDS_TOTAL "mori_umbp_ssd_eviction_rounds_total" -#define MORI_UMBP_METRIC_SSD_EVICTION_ROUNDS_TOTAL_HELP \ - "Local SSD eviction rounds that actually ran (used above high watermark)" - -#define MORI_UMBP_METRIC_SSD_EVICTION_VICTIMS_TOTAL "mori_umbp_ssd_eviction_victims_total" -#define MORI_UMBP_METRIC_SSD_EVICTION_VICTIMS_TOTAL_HELP \ - "SSD keys evicted locally (REMOVE SSD emitted)" - -#define MORI_UMBP_METRIC_SSD_EVICTION_BYTES_FREED_TOTAL "mori_umbp_ssd_eviction_bytes_freed_total" -#define MORI_UMBP_METRIC_SSD_EVICTION_BYTES_FREED_TOTAL_HELP "Bytes freed by local SSD eviction" - -#define MORI_UMBP_METRIC_SSD_EVICTION_BACKEND_FAILED_TOTAL \ - "mori_umbp_ssd_eviction_backend_failed_total" -#define MORI_UMBP_METRIC_SSD_EVICTION_BACKEND_FAILED_TOTAL_HELP \ - "Local SSD evictions where the backend Evict failed (key kept for retry)" - -// staging (peer_service.cpp SSD read slots) -#define MORI_UMBP_METRIC_SSD_STAGING_SLOTS_IN_USE "mori_umbp_ssd_staging_slots_in_use" -#define MORI_UMBP_METRIC_SSD_STAGING_SLOTS_IN_USE_HELP \ - "SSD read staging slots currently in use (Preparing or Leased), sampled once per flush" - -#define MORI_UMBP_METRIC_SSD_STAGING_EXPIRED_RECLAIMS_TOTAL \ - "mori_umbp_ssd_staging_expired_reclaims_total" -#define MORI_UMBP_METRIC_SSD_STAGING_EXPIRED_RECLAIMS_TOTAL_HELP \ - "SSD read staging slots reclaimed after lease TTL expiry" - -#define MORI_UMBP_METRIC_SSD_STAGING_SLOT_FULL_REJECTS_TOTAL \ - "mori_umbp_ssd_staging_slot_full_rejects_total" -#define MORI_UMBP_METRIC_SSD_STAGING_SLOT_FULL_REJECTS_TOTAL_HELP \ - "PrepareSsdRead calls rejected with NO_SLOT because all staging slots were busy" - -// Optional reader-side diagnostic: per-batch count of remote SSD reads that -// stayed transient-not-served (NO_SLOT / reader-local lease expiry) after the -// configured attempts. Reported by the reader node; no reason label (kept -// cheap — one AddCounter per batch). lease expiry is reader-local and never -// visible in the peer-side ssd_read_total, so this is its only observability. -#define MORI_UMBP_METRIC_SSD_READ_CLIENT_TRANSIENT_TOTAL "mori_umbp_ssd_read_client_transient_total" -#define MORI_UMBP_METRIC_SSD_READ_CLIENT_TRANSIENT_TOTAL_HELP \ - "Remote SSD reads reported not-served this round due to transient NO_SLOT / " \ - "reader-local lease expiry (not a definitive miss)" +// --- Storage-backend and transfer-layer metrics ----------------------------- +// +// NOT here. Everything a medium or a transport reports now lives in +// umbp/distributed/metrics/component_metrics.h under names that carry the +// medium in a LABEL rather than in the identifier, because the previous +// arrangement — mori_umbp_ssd_* beside an implicit DRAM set — is what forced a +// separate dashboard per medium and left an SSD panel wired to counters that +// no longer had a publisher after the backend-agnostic refactor. +// +// This header keeps what the MASTER itself measures. A peer-side component's +// metrics arrive through ReportMetrics and are named by the component. diff --git a/src/umbp/include/umbp/distributed/master/master_server.h b/src/umbp/include/umbp/distributed/master/master_server.h index 1874a300d..39ea97cca 100644 --- a/src/umbp/include/umbp/distributed/master/master_server.h +++ b/src/umbp/include/umbp/distributed/master/master_server.h @@ -59,6 +59,12 @@ class MasterServer { // has called BuildAndStart(). uint16_t GetBoundPort() const { return bound_port_.load(); } + // Latest registered/heartbeat view of one peer. Used by tests to observe + // whether capacity heartbeats aggregated across same-tier backends. + std::optional GetClient(const std::string& node_id) const { + return store_ ? store_->GetClient(node_id) : std::nullopt; + } + private: MasterServerConfig config_; // Single owner of all master metadata state (block locations, client diff --git a/src/umbp/include/umbp/distributed/master/redis/lua_scripts.h b/src/umbp/include/umbp/distributed/master/redis/lua_scripts.h index b814cd9e7..07cd9d653 100644 --- a/src/umbp/include/umbp/distributed/master/redis/lua_scripts.h +++ b/src/umbp/include/umbp/distributed/master/redis/lua_scripts.h @@ -67,7 +67,7 @@ namespace mori::umbp::redis { // register_client: // KEYS[1] = node key -// ARGV = [tag, node_id, now_ms, stale_after_ms, addr, peer, caps, engine, tags] +// ARGV = [tag, node_id, now_ms, stale_after_ms, addr, peer, caps, engine, tags, lcaps] // Returns 1 if registered/revived, 0 if rejected (ALIVE and not stale). inline const std::string kRegisterClientLua = R"LUA(--!df flags=allow-undeclared-keys local nodeKey = KEYS[1] @@ -87,7 +87,7 @@ redis.call('DEL', nodeKey) redis.call('HSET', nodeKey, 'status', 1, 'last_hb', now, 'reg_at', now, 'seq', 0, 'addr', ARGV[5], 'peer', ARGV[6], 'caps', ARGV[7], - 'engine', ARGV[8], 'tags', ARGV[9]) + 'engine', ARGV[8], 'tags', ARGV[9], 'lcaps', ARGV[10]) redis.call('SADD', tag .. ':nodes:alive', nodeId) redis.call('HSET', tag .. ':alive_peers', nodeId, ARGV[6]) return 1 @@ -95,8 +95,8 @@ return 1 // apply_heartbeat: // KEYS[1] = node key -// ARGV = [tag, node_id, seq, now_ms, is_full_sync, caps_blob, n_events, -// then per event: kind, block_key, tier, size] +// ARGV = [tag, node_id, seq, now_ms, is_full_sync, caps_blob, lcaps_blob, n_events, +// then per event: kind, block_key, tier, size, logical_tier] // `block_key` is the FULL (already sharded) block key composed by the caller // (KeySchema::Block), so this script never has to know the shard layout. The // node's reverse-index set stores these full block keys as members. @@ -110,7 +110,8 @@ local seq = tonumber(ARGV[3]) local now = tonumber(ARGV[4]) local full = tonumber(ARGV[5]) local caps = ARGV[6] -local nev = tonumber(ARGV[7]) +local lcaps = ARGV[7] +local nev = tonumber(ARGV[8]) if redis.call('EXISTS', nodeKey) == 0 then return { 'UNKNOWN', '0' } @@ -127,7 +128,8 @@ if full == 0 and seq ~= last + 1 then return { 'SEQ_GAP', tostring(last) } end -redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1, 'seq', seq, 'caps', caps) +redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1, 'seq', seq, + 'caps', caps, 'lcaps', lcaps) redis.call('SADD', aliveSet, nodeId) redis.call('HSET', peers, nodeId, peer) @@ -143,19 +145,25 @@ local function cleanupEmpty(bk) if not anyLoc then redis.call('DEL', bk) end end -local function addLoc(blockKey, tier, size) +local function locationField(tier, logical) + local field = 'l|' .. nodeId .. '|' .. tier + if logical ~= '' then field = field .. '|' .. logical end + return field +end + +local function addLoc(blockKey, tier, size, logical) if redis.call('EXISTS', blockKey) == 0 then redis.call('HSET', blockKey, '_created', now, '_lacc', now, '_acnt', 0) end - local field = 'l|' .. nodeId .. '|' .. tier + local field = locationField(tier, logical) if redis.call('HEXISTS', blockKey, field) == 0 then redis.call('HSET', blockKey, field, size) redis.call('SADD', blocksSet, blockKey) end end -local function removeLoc(blockKey, tier) - local field = 'l|' .. nodeId .. '|' .. tier +local function removeLoc(blockKey, tier, logical) + local field = locationField(tier, logical) if redis.call('HDEL', blockKey, field) == 1 then local flds = redis.call('HKEYS', blockKey) local nodeStill = false @@ -180,18 +188,18 @@ if full == 1 then end redis.call('DEL', blocksSet) for i = 0, nev - 1 do - local base = 8 + i * 4 + local base = 9 + i * 5 if tonumber(ARGV[base]) == 0 then - addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) + addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3], ARGV[base + 4]) end end else for i = 0, nev - 1 do - local base = 8 + i * 4 + local base = 9 + i * 5 if tonumber(ARGV[base]) == 0 then - addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) + addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3], ARGV[base + 4]) else - removeLoc(ARGV[base + 1], ARGV[base + 2]) + removeLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 4]) end end end @@ -204,7 +212,7 @@ return { 'APPLIED', tostring(seq) } // KEYS stays single-slot on Redis Cluster / one proactor on Dragonfly). // ARGV = [now_ms, lease_ms, n_exclude, exclude_node_1..exclude_node_k] // Returns an array of n elements; element i is a flat array -// [node, size, tier, node, size, tier, ...] of the surviving locations. +// [node, size, tier, logical_tier, ...] of the surviving locations. // Only keys with >=1 surviving location get a lease + access bump. // // The per-key HGETALL already returns _acnt, so the lease/access bump folds @@ -234,11 +242,15 @@ for i = 1, #KEYS do local sep = string.find(rest, '|', 1, true) if sep ~= nil then local node = string.sub(rest, 1, sep - 1) - local tier = string.sub(rest, sep + 1) + local tail = string.sub(rest, sep + 1) + local tierSep = string.find(tail, '|', 1, true) + local tier = tierSep and string.sub(tail, 1, tierSep - 1) or tail + local logical = tierSep and string.sub(tail, tierSep + 1) or '' if not excl[node] then locs[#locs + 1] = node locs[#locs + 1] = v locs[#locs + 1] = tier + locs[#locs + 1] = logical touched = true end end @@ -402,7 +414,7 @@ return dead // apply_heartbeat_control (control instance): // KEYS[1] = node key -// ARGV = [tag, node_id, seq, now_ms, is_full_sync, caps_blob] +// ARGV = [tag, node_id, seq, now_ms, is_full_sync, caps_blob, lcaps_blob] // seq-CAS + record + nodes:alive/alive_peers ONLY (no block work). // Returns { status_string, acked_seq_string } like apply_heartbeat. inline const std::string kApplyHeartbeatControlLua = R"LUA(--!df flags=allow-undeclared-keys @@ -413,6 +425,7 @@ local seq = tonumber(ARGV[3]) local now = tonumber(ARGV[4]) local full = tonumber(ARGV[5]) local caps = ARGV[6] +local lcaps = ARGV[7] if redis.call('EXISTS', nodeKey) == 0 then return { 'UNKNOWN', '0' } end @@ -426,7 +439,8 @@ if full == 0 and seq ~= last + 1 then redis.call('HSET', peers, nodeId, peer) return { 'SEQ_GAP', tostring(last) } end -redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1, 'seq', seq, 'caps', caps) +redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1, 'seq', seq, + 'caps', caps, 'lcaps', lcaps) redis.call('SADD', aliveSet, nodeId) redis.call('HSET', peers, nodeId, peer) return { 'APPLIED', tostring(seq) } @@ -434,7 +448,7 @@ return { 'APPLIED', tostring(seq) } // apply_block_events (one shard's instance): // ARGV = [revidx_key, node_prefix, is_full_sync, now_ms, n_events, -// then per event: kind, block_key, tier, size] +// then per event: kind, block_key, tier, size, logical_tier] // `revidx_key` is this (node, shard) reverse-index set; `node_prefix` is // 'l||'. All block keys passed in ARGV are on this instance/slot. // Idempotent: ADD overwrites, REMOVE of a missing loc is a no-op, full_sync @@ -455,19 +469,25 @@ local function cleanupEmpty(bk) redis.call('DEL', bk) end -local function addLoc(bk, tier, size) +local function locationField(tier, logical) + local field = nodePfx .. tier + if logical ~= '' then field = field .. '|' .. logical end + return field +end + +local function addLoc(bk, tier, size, logical) if redis.call('EXISTS', bk) == 0 then redis.call('HSET', bk, '_created', now, '_lacc', now, '_acnt', 0) end - local field = nodePfx .. tier + local field = locationField(tier, logical) if redis.call('HEXISTS', bk, field) == 0 then redis.call('HSET', bk, field, size) redis.call('SADD', revidx, bk) end end -local function removeLoc(bk, tier) - local field = nodePfx .. tier +local function removeLoc(bk, tier, logical) + local field = locationField(tier, logical) if redis.call('HDEL', bk, field) == 1 then local flds = redis.call('HKEYS', bk) local still = false @@ -490,16 +510,16 @@ if full == 1 then end redis.call('DEL', revidx) for i = 0, nev - 1 do - local base = 6 + i * 4 - if tonumber(ARGV[base]) == 0 then addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) end + local base = 6 + i * 5 + if tonumber(ARGV[base]) == 0 then addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3], ARGV[base + 4]) end end else for i = 0, nev - 1 do - local base = 6 + i * 4 + local base = 6 + i * 5 if tonumber(ARGV[base]) == 0 then - addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) + addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3], ARGV[base + 4]) else - removeLoc(ARGV[base + 1], ARGV[base + 2]) + removeLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 4]) end end end @@ -871,7 +891,9 @@ for ki = 1, #KEYS do local sep = string.find(rest, '|', 1, true) if sep ~= nil then local node = string.sub(rest, 1, sep - 1) - local tier = string.sub(rest, sep + 1) + local tail = string.sub(rest, sep + 1) + local tierSep = string.find(tail, '|', 1, true) + local tier = tierSep and string.sub(tail, 1, tierSep - 1) or tail if wanted[node] and wanted[node][tier] then cands[#cands + 1] = bk cands[#cands + 1] = node diff --git a/src/umbp/include/umbp/distributed/master/redis_master_metadata_store.h b/src/umbp/include/umbp/distributed/master/redis_master_metadata_store.h index ec1589d1f..c93a0922e 100644 --- a/src/umbp/include/umbp/distributed/master/redis_master_metadata_store.h +++ b/src/umbp/include/umbp/distributed/master/redis_master_metadata_store.h @@ -98,10 +98,11 @@ class RedisMasterMetadataStore : public IMasterMetadataStore { std::chrono::system_clock::time_point now, std::chrono::system_clock::duration stale_after) override; void UnregisterClient(const std::string& node_id) override; - HeartbeatResult ApplyHeartbeat(const std::string& node_id, uint64_t seq, - std::chrono::system_clock::time_point now, - const std::map& caps, - const std::vector& events, bool is_full_sync) override; + HeartbeatResult ApplyHeartbeat( + const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, + const std::map& caps, const std::vector& events, + bool is_full_sync, + const std::map& logical_caps = {}) override; std::vector ExpireStaleClients( std::chrono::system_clock::time_point cutoff) override; diff --git a/src/umbp/include/umbp/distributed/metrics/component_metrics.h b/src/umbp/include/umbp/distributed/metrics/component_metrics.h new file mode 100644 index 000000000..031b6c493 --- /dev/null +++ b/src/umbp/include/umbp/distributed/metrics/component_metrics.h @@ -0,0 +1,221 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include + +// --------------------------------------------------------------------------- +// The one observability seam the refactored data plane publishes through. +// +// The rule this header exists to enforce: A COMPONENT IS OBSERVED, IT DOES NOT +// OBSERVE ITSELF. A storage medium and a transfer engine each get their +// generic metrics from a decorator that sits on their interface, not from +// calls they make — so a NEW backend or a NEW engine is instrumented the +// moment it is composed in, with no metrics code of its own and no dashboard +// edit. See instrumented_backend.h and CompositeTransferEngine. +// +// What a component may still publish for itself is what its interface cannot +// show from outside — the drive's own read outcomes, a staging arena's +// occupancy. Those ride SampleMetrics() below, and they go out under a +// GENERIC metric name with the specifics in a label (event=/state=), never +// under a name that spells the medium. A metric named mori_umbp_ssd_* forces +// a per-medium panel; mori_umbp_backend_medium_events_total{tier="SSD"} shows +// up in the panel that was already there. +// +// Nothing here knows about gRPC or the master: a component fills in samples, +// MetricPublisher diffs them and hands deltas to a sink. That is what makes +// the whole path unit-testable without standing up a server. +// --------------------------------------------------------------------------- + +namespace mori::umbp { + +// --------------------------------------------------------------------------- +// Metric-name vocabulary +// +// Backend-agnostic by construction: every name below is shared by every +// medium, and the medium is a LABEL. A dashboard panel written as +// `sum by (tier) (rate(mori_umbp_backend_bytes_total{op="commit"}[$__rate_interval]))` +// gains a new series the day a new backend registers and needs no edit. +// --------------------------------------------------------------------------- + +// --- Generic storage-backend metrics (emitted by InstrumentedBackend) ------- +// Labels on all of them: tier=, backend= + +// op=, status= +// Counted per ENTRY (one key), not per batch call — the batch shape is the +// wire's business, the key count is the workload. +#define MORI_UMBP_METRIC_BACKEND_OPS_TOTAL "mori_umbp_backend_ops_total" +#define MORI_UMBP_METRIC_BACKEND_OPS_TOTAL_HELP \ + "Storage-backend slot-lifecycle operations by op and status, counted per key. Emitted by " \ + "the instrumentation decorator for every medium, so a new backend appears here with no " \ + "code of its own" + +// op=: one batch call, whatever its size. +// Pair with _ops_total for average batch depth, which is the knob that decides +// whether a medium is being driven efficiently. +#define MORI_UMBP_METRIC_BACKEND_BATCHES_TOTAL "mori_umbp_backend_batches_total" +#define MORI_UMBP_METRIC_BACKEND_BATCHES_TOTAL_HELP \ + "Storage-backend batch calls by op (one per call regardless of batch size)" + +// op=. commit = bytes landed in the medium, resolve = +// bytes handed out to readers, evict = bytes freed. rate() gives per-medium +// write / read / reclaim bandwidth in one panel. +#define MORI_UMBP_METRIC_BACKEND_BYTES_TOTAL "mori_umbp_backend_bytes_total" +#define MORI_UMBP_METRIC_BACKEND_BYTES_TOTAL_HELP \ + "Bytes committed to / resolved from / freed in a storage backend, by op" + +// op=. Seconds spent INSIDE the backend +// call. rate(seconds_total)/rate(batches_total) is mean call latency; a +// histogram would need a second transport and this is the number that moves. +#define MORI_UMBP_METRIC_BACKEND_OP_SECONDS_TOTAL "mori_umbp_backend_op_seconds_total" +#define MORI_UMBP_METRIC_BACKEND_OP_SECONDS_TOTAL_HELP \ + "Cumulative seconds spent inside storage-backend calls, by op. Divide by " \ + "mori_umbp_backend_batches_total for mean call latency, or by " \ + "mori_umbp_backend_ops_total for mean per-key latency" + +// --- Medium-specific detail, under generic names --------------------------- +// A medium reports what its interface cannot show from outside. The medium +// names the EVENT, never the metric, so every medium lands in the same panel. +// Labels: tier, backend, event= + +#define MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL "mori_umbp_backend_medium_events_total" +#define MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL_HELP \ + "Medium-internal events a backend chose to publish (drive read outcomes, single-flight " \ + "coalescing, eviction rounds, staging pressure), keyed by event" + +#define MORI_UMBP_METRIC_BACKEND_MEDIUM_BYTES_TOTAL "mori_umbp_backend_medium_bytes_total" +#define MORI_UMBP_METRIC_BACKEND_MEDIUM_BYTES_TOTAL_HELP \ + "Medium-internal byte counters (e.g. bytes that actually reached the device, as opposed " \ + "to the logical bytes in mori_umbp_backend_bytes_total), keyed by event" + +// Gauge. Labels: tier, backend, state= +#define MORI_UMBP_METRIC_BACKEND_MEDIUM_STATE "mori_umbp_backend_medium_state" +#define MORI_UMBP_METRIC_BACKEND_MEDIUM_STATE_HELP \ + "Medium-internal live state a backend chose to publish (arena occupancy, queue depth)" + +// --- Generic transfer-layer metrics (emitted by CompositeTransferEngine) ---- +// Labels: engine=, direction= + +// status=. One count per PLAN — the unit an engine actually posts. +#define MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL "mori_umbp_transfer_ops_total" +#define MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL_HELP \ + "Transfer plans posted by engine, direction and outcome. engine=\"none\" with " \ + "status=\"rejected\" counts items no engine could carry, which is a routing bug and not " \ + "a transfer failure" + +#define MORI_UMBP_METRIC_TRANSFER_BYTES_TOTAL "mori_umbp_transfer_bytes_total" +#define MORI_UMBP_METRIC_TRANSFER_BYTES_TOTAL_HELP \ + "Bytes posted through a transfer engine, by engine and direction" + +// Wall time from Submit to the completion of Wait, charged to the engine that +// carried the plans. rate(bytes)/rate(seconds) is achieved bandwidth. +#define MORI_UMBP_METRIC_TRANSFER_SECONDS_TOTAL "mori_umbp_transfer_seconds_total" +#define MORI_UMBP_METRIC_TRANSFER_SECONDS_TOTAL_HELP \ + "Cumulative seconds transfers were in flight (submit to settled), by engine and " \ + "direction. Divide mori_umbp_transfer_bytes_total by this for achieved bandwidth" + +// --------------------------------------------------------------------------- +// Sample types +// --------------------------------------------------------------------------- + +enum class MetricKind { + // Value is MONOTONIC and the publisher ships `value - last`. A decrease is + // read as a rebuilt source and rebases rather than shipping a negative delta. + kCounter, + // Value is the current reading and is shipped as-is every tick. + kGauge, +}; + +using MetricLabels = std::vector>; + +// One sample from a component. `name`/`help` are the Prometheus identity and +// MUST be string literals with static storage duration — the same metric +// reported under different label sets shares both. +struct MetricSample { + const char* name = nullptr; + const char* help = nullptr; + MetricLabels labels; + uint64_t value = 0; + MetricKind kind = MetricKind::kCounter; + // Applied by the publisher to the delta (counter) or the reading (gauge). + // Lets a component accumulate in the integer unit it can add atomically — + // nanoseconds, pages — while the metric keeps the unit its name promises. + double scale = 1.0; +}; + +// Anything that can be sampled once per metrics tick. Implemented by +// MediumBackend and TransferEngine so the publisher walks both through one +// pointer type and neither has to know how the samples leave the process. +// +// Defaulted to empty: a component with nothing of its own to say (and every +// test double) needs no code, and still gets the generic metrics its decorator +// emits. Never called under a component lock — treat it as a plain read of +// relaxed atomics. +class MetricSource { + public: + virtual ~MetricSource() = default; + virtual std::vector SampleMetrics() const { return {}; } +}; + +// --------------------------------------------------------------------------- +// MetricPublisher — turns snapshots into deltas and hands them to a sink +// +// One instance holds the baselines for every source it publishes, keyed by +// (source id, metric name, labels). Two sources reporting the same metric +// name must not share a baseline or one would cancel the other's progress +// out, which is what `source_id` separates. +// --------------------------------------------------------------------------- +class MetricPublisher { + public: + struct Sink { + // delta >= 0, already differenced. Not called for a zero delta. + std::function + counter; + // Current reading, shipped every tick. + std::function + gauge; + }; + + // Sample `source`, prepend `source_labels` to every sample's own labels, and + // emit through `sink`. `source_id` identifies the source's baselines and + // must be stable across ticks for the same object. + void Publish(const std::string& source_id, const MetricLabels& source_labels, + const MetricSource& source, const Sink& sink); + + // Same, for samples a caller already has in hand. + void Publish(const std::string& source_id, const MetricLabels& source_labels, + const std::vector& samples, const Sink& sink); + + // Drop every baseline. The next Publish treats each sample as a fresh + // source, which is what a re-Init wants. + void Reset() { last_.clear(); } + + private: + std::unordered_map last_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/obs_counters.h b/src/umbp/include/umbp/distributed/obs_counters.h deleted file mode 100644 index ee6ba9ade..000000000 --- a/src/umbp/include/umbp/distributed/obs_counters.h +++ /dev/null @@ -1,56 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -// Test/bench-only build switch (renamed from MORI_UMBP_OBS_COUNTERS in v2.1.2). -// -// MORI_UMBP_TESTING covers TWO orthogonal responsibilities, gated by the same -// macro because both are exercised only by the same downstream targets: -// -// 1. Observability counter increments. Atomic counter members and their -// public getters are declared unconditionally (ABI stable across -// test/release builds); only the increment call sites are gated, so -// release builds pay zero CPU cost and getters return 0. -// -// 2. Test seams. Selected private members (e.g. PoolClient::IssueBatchWrite) -// are declared `virtual` only when the macro is defined, allowing test -// subclasses to inject failures. Release builds keep them as plain -// members for inlining. -// -// Build implication: test/release object files MUST NOT be mixed-linked. -// Toggling MORI_UMBP_TESTING changes class layout (vtable presence) and ABI. -// The CI test target builds with -DMORI_UMBP_TESTING=ON unconditionally; -// production release leaves it OFF. -// -// Usage: -// MORI_UMBP_OBS_INC(some_counter_); -// MORI_UMBP_OBS_ADD(some_counter_, n); -// MORI_UMBP_TEST_VIRTUAL void Foo(); // virtual under -DMORI_UMBP_TESTING -#ifdef MORI_UMBP_TESTING -#define MORI_UMBP_OBS_INC(counter) (counter).fetch_add(1, std::memory_order_relaxed) -#define MORI_UMBP_OBS_ADD(counter, n) (counter).fetch_add((n), std::memory_order_relaxed) -#define MORI_UMBP_TEST_VIRTUAL virtual -#else -#define MORI_UMBP_OBS_INC(counter) ((void)0) -#define MORI_UMBP_OBS_ADD(counter, n) ((void)0) -#define MORI_UMBP_TEST_VIRTUAL -#endif diff --git a/src/umbp/include/umbp/distributed/peer/backend/hbm_backend.h b/src/umbp/include/umbp/distributed/peer/backend/hbm_backend.h new file mode 100644 index 000000000..5541fbf03 --- /dev/null +++ b/src/umbp/include/umbp/distributed/peer/backend/hbm_backend.h @@ -0,0 +1,80 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/peer/backend/page_backend.h" + +namespace mori::umbp { + +// GPU HBM, allocated with hipMalloc on a fixed device. +// +// THE WHOLE HBM BACKEND IS THIS CLASS. There is no HbmBackend type, and that +// is the finding rather than a shortcut: PageBackend's slot lifecycle, bitmap +// allocator, event outbox, reaper, read leases and copy pins never consult the +// tier, so a second paged medium needs no second copy of any of it. What HBM +// actually changes is three facts — hipMalloc instead of mmap, loc=GPU instead +// of CPU, and a real device ordinal instead of -1 — which is exactly the +// surface PageMemorySource exposes. +// +// The device is fixed at construction rather than read from the ambient current +// device at Allocate time. A backend is called from the peer service's gRPC +// handler threads and the heartbeat thread (MediumBackend's threading note), and +// none of those has a meaningful current device — inheriting one would make the +// pool land on whichever thread happened to allocate first. +class HbmPageMemorySource final : public PageMemorySource { + public: + explicit HbmPageMemorySource(int device) : device_(device) {} + ~HbmPageMemorySource() override { Release(); } + + // hipMalloc one buffer per size, on device_. All-or-nothing. + bool Allocate(const std::vector& sizes, std::vector* out) override; + void Release() override; + + mori::io::MemoryLocationType LocationType() const override { + return mori::io::MemoryLocationType::GPU; + } + int Device() const override { return device_; } + const char* Name() const override { return "HbmPageMemorySource"; } + + private: + int device_ = 0; + std::vector ptrs_; +}; + +// The one place a caller outside this file obtains an HBM backend — the +// counterpart of MakePageBackend, and subject to the same Phase 5 Rule A: it +// returns MediumBackend, so PoolClient::Init names no concrete backend type. +// +// `device` is the GPU ordinal the pool is allocated on. `buffer_sizes` is one +// hipMalloc per entry, in bytes. +std::unique_ptr MakeHbmBackend(uint64_t page_size, int device, + std::vector buffer_sizes, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/backend/instrumented_backend.h b/src/umbp/include/umbp/distributed/peer/backend/instrumented_backend.h new file mode 100644 index 000000000..05eafbd77 --- /dev/null +++ b/src/umbp/include/umbp/distributed/peer/backend/instrumented_backend.h @@ -0,0 +1,156 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include + +#include "umbp/distributed/metrics/component_metrics.h" +#include "umbp/distributed/peer/backend/medium_backend.h" + +namespace mori::umbp { + +// --------------------------------------------------------------------------- +// InstrumentedBackend — generic storage metrics for EVERY medium, for free +// +// A MediumBackend decorator that forwards every call and measures the slot +// lifecycle on the way through: entries by outcome, bytes, batch count, and +// time spent inside the medium. It is the answer to "how do we keep metrics +// working when a backend is added" — the answer is that adding a backend is +// not a metrics event at all. PoolClient::Init wraps whatever the medium +// switch produced, so a medium written next year emits the same series as +// DRAM does today, lands in the same dashboard panels, and its author never +// opens this file. +// +// Why a decorator rather than counters at the call sites: the backends are +// driven from three places (PeerServiceServer for remote callers, PoolClient +// for same-node access, the eviction path), and instrumenting callers means +// every new caller is a new hole. There is exactly one interface, so there is +// exactly one place to stand. +// +// What it deliberately does NOT do: name a medium, look at a TierType, or +// branch on the wrapped type. Tier() and Name() are pass-through, and they +// become the tier= / backend= labels on everything published. +// +// Cost: one steady_clock pair and a handful of relaxed atomic adds per BATCH +// call (not per key). Against a call that touches a medium — a hipMemcpy, a +// drive read, a bitmap allocation under a mutex — that is not measurable. +// --------------------------------------------------------------------------- +class InstrumentedBackend final : public MediumBackend { + public: + explicit InstrumentedBackend(std::unique_ptr inner); + + // The wrapped backend, for the few call sites that legitimately need the + // concrete object (tests asserting medium internals). Never used to bypass + // the decorator on a data path. + MediumBackend* Inner() const { return inner_.get(); } + + // ---- identity: pass-through, and the source of the metric labels ---- + TierType Tier() const override { return inner_->Tier(); } + const char* Name() const override { return inner_->Name(); } + + // ---- ownership ---- + bool Init(MemoryRegistrar* registrar) override; + void Shutdown() override { inner_->Shutdown(); } + + // ---- control plane ---- + TierCapacity Capacity() const override { return inner_->Capacity(); } + uint64_t OwnedKeyCount() const override { return inner_->OwnedKeyCount(); } + bool Contains(const std::string& key) const override { return inner_->Contains(key); } + bool AcquireMigrationRead(const std::string& key, ResolvedEntry* resolved) override { + return inner_->AcquireMigrationRead(key, resolved); + } + void ReleaseMigrationRead(const std::string& key) override { inner_->ReleaseMigrationRead(key); } + std::vector DrainPendingEvents() override { return inner_->DrainPendingEvents(); } + std::vector SnapshotOwnedKeys() const override { return inner_->SnapshotOwnedKeys(); } + std::vector SnapshotOwnedKeysForFullSync() override { + return inner_->SnapshotOwnedKeysForFullSync(); + } + void ClearLocal() override { inner_->ClearLocal(); } + void ClearFullSyncAcked() override { inner_->ClearFullSyncAcked(); } + bool IsClearFullSyncPending() const override { return inner_->IsClearFullSyncPending(); } + void SetEventPublishing(bool enabled) override { inner_->SetEventPublishing(enabled); } + void EnableLocalEviction(double high_watermark, double low_watermark) override { + inner_->EnableLocalEviction(high_watermark, low_watermark); + } + + void SetAutoFlushHook(size_t threshold, std::function cb) override { + inner_->SetAutoFlushHook(threshold, std::move(cb)); + } + + // ---- observability ---- + // The generic series this decorator measured, followed by whatever the + // wrapped medium publishes about its own internals. One source, so + // PoolClient publishes a backend with a single call and cannot forget the + // medium's half. + std::vector SampleMetrics() const override; + + // ---- slot lifecycle: measured, then forwarded ---- + std::vector BatchAllocate(const std::vector& reqs) override; + std::vector BatchCommit(const std::vector& reqs) override; + std::vector BatchAbort(const std::vector& slot_ids) override; + std::vector BatchResolve(const std::vector& keys, bool include_descs, + bool allow_file_refs = false) override; + std::vector Evict(const std::vector& keys) override; + + // ---- bootstrap / local endpoints: pass-through ---- + uint64_t PageSize() const override { return inner_->PageSize(); } + std::vector AllBufferDescs() const override { + return inner_->AllBufferDescs(); + } + size_t BufferCount() const override { return inner_->BufferCount(); } + TransferRef BufferRef(uint32_t buffer_index) const override { + return inner_->BufferRef(buffer_index); + } + + private: + // One op's accumulators. Relaxed atomics: these are read once per metrics + // tick from another thread and are never correctness state, so ordering + // against the data plane buys nothing. + struct OpCounters { + std::atomic batches{0}; + std::atomic nanos{0}; + std::atomic bytes{0}; + }; + + // Per (op, status) entry counts. A fixed array rather than a map because the + // set is closed and this is touched per call: kOpCount * kStatusCount is 25 + // atomics, and the indices are the label values. + enum Op { kAllocate, kCommit, kAbort, kResolve, kEvict, kOpCount }; + enum Status { kOk, kExists, kNoSpace, kMiss, kFailed, kStatusCount }; + + void Add(Op op, Status status, uint64_t n) { + entries_[op][status].fetch_add(n, std::memory_order_relaxed); + } + + std::unique_ptr inner_; + OpCounters ops_[kOpCount]; + std::atomic entries_[kOpCount][kStatusCount]; +}; + +// Wrap `inner` so it reports the generic backend metrics. Null in, null out — +// a failed medium construction stays a failed medium construction. +std::unique_ptr MakeInstrumentedBackend(std::unique_ptr inner); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/backend/medium_backend.h b/src/umbp/include/umbp/distributed/peer/backend/medium_backend.h new file mode 100644 index 000000000..3d3caa395 --- /dev/null +++ b/src/umbp/include/umbp/distributed/peer/backend/medium_backend.h @@ -0,0 +1,623 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/distributed/metrics/component_metrics.h" +#include "umbp/distributed/transfer/transfer_engine.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +// One storage medium on a peer. +// +// A backend OWNS bytes and PUBLISHES DESCRIPTORS for them. It does not move +// them. That split follows from one property of mori::io::MemoryDesc: it is +// self-describing about the medium (loc, deviceId, numaNode, ipcHandle, +// fabricHandle). So — +// +// bytes addressed by a descriptor are already medium-agnostic; +// bytes addressed by a raw pointer are medium-specific; +// registration is the one-way conversion, and belongs to whoever allocated. +// +// A backend is therefore the only layer that knows what kind of memory it is +// holding, because it is the one that allocated it. Everything downstream of +// registration works off the descriptor and needs no per-medium branch (see +// design doc §2 for the rule, §3 for the three-component split). +// +// One instance == one medium. That is why no method takes a TierType: the +// tier is the identity of the object, and the caller has already dispatched on +// it via BackendRegistry. This inverts PeerDramAllocator's internal +// map, relocating the multi-medium generality to the registry +// where it is exercised. +// +// This interface absorbs OwnedLocationSource: the heartbeat drains/snapshots +// every registered backend and concatenates the events into ONE bundle under a +// single monotonic seq — never one seq per medium (that breaks the ack / +// seq-gap full-sync recovery). +// +// The backend contract treats media as equivalent: it has no built-in read +// priority or put-eligibility flag. Legacy policies retain that behavior. The +// optional JSON TieredPlacementPolicy supplies ordering and offload topology +// above this interface, keeping physical storage implementations policy-free. +// +// Threading: implementations must be safe to call from the peer service's gRPC +// handler threads and the heartbeat thread concurrently. +// +// NOT part of this contract, deliberately. The first three were each proposed +// and then rejected once the transfer engine (§4) was settled; they are listed +// so they are not rediscovered and re-added: +// * LocalCopyIn / LocalCopyOut — local access is a transfer whose endpoints +// are both local; the engine's planner picks a memcpy or pread path for it. +// A copy method here would re-admit a second byte-moving path, which is the +// coupling this whole refactor exists to remove. +// * a staging / bounce pool — belongs to the transfer engine, the only layer +// that can observe completion. A pool here would have to fall back to a +// TTL, which is the deleted PrepareSsdRead lease under another name. +// * byte-moving or staging methods — staging ownership stays an +// implementation detail. ResolveOutcome::kBusy is intentionally part of +// the RESULT, however: a staged backend must distinguish transient pressure +// from kMissing or the client excludes a node that does hold the key. +// * local mode's TierBackend (blocking read/write of bytes) — a different +// medium contract at a different layer, fenced off by lint in Phase 5 +// * single-key Allocate/Commit/Abort/Resolve — the single-key RPCs are +// served by one-element batches +// +// A backend also must not move bytes itself, and since Phase 6 that is a +// COMPILE ERROR rather than a lint rule: Init receives a MemoryRegistrar, which +// has only RegisterMemory/Deregister on it. The object behind that pointer is +// the same TransferEngine everything else holds — one object, two views (design +// doc §5 Rule C). + +// --------------------------------------------------------------------------- +// Data-plane value types +// +// These are the medium-agnostic forms of PeerDramAllocator's nested request / +// result structs. Two differences from those, both intentional: +// * no `tier` field — see the "one instance == one medium" note above +// * no allocator-internal state (generation, deadline) — the wire carries +// pending_ttl_ms, and generation is a backend implementation detail +// --------------------------------------------------------------------------- + +// Outcome of one slot allocation. kFailedNoSpace is split out from kFailed so +// the writer can keep retrying on other peers; every other failure collapses +// into kFailed and is diagnosed from the backend's own log. +enum class AllocateOutcome { + kSuccessAllocated, + kSuccessAlreadyExists, // dedup hit against the backend's owned map + kFailed, // generic — see backend log for reason + kFailedNoSpace, // medium exhausted; retry on another peer +}; + +struct AllocateRequest { + std::string key; + uint64_t size = 0; +}; + +// Populated fields below `outcome` are meaningful ONLY when +// outcome == kSuccessAllocated. +struct AllocateResult { + AllocateOutcome outcome = AllocateOutcome::kFailed; + uint64_t slot_id = 0; + std::vector pages; + uint64_t size = 0; + uint64_t page_size = 0; + // Descs for exactly the buffers referenced by `pages`, deduplicated and + // ascending by buffer_index — enough for the writer to RDMA without a + // follow-up GetPeerInfo. + std::vector descs; + // How long the slot survives without a matching BatchCommit / BatchAbort. + uint64_t pending_ttl_ms = 0; +}; + +struct CommitRequest { + uint64_t slot_id = 0; + std::string key; +}; + +struct CommitResult { + bool success = false; // false = slot unknown: reaped, aborted, or never existed + uint64_t bytes_committed = 0; +}; + +// A resolve needs one more state than a cache lookup. kBusy means this backend +// owns the key but cannot publish its bytes until a transient resource (SSD +// staging) becomes available; callers must retry THIS peer rather than exclude +// it as they do for kMissing. kFailed is a permanent shape/backend error and +// must not be retried indefinitely. +enum class ResolveOutcome { + kMissing, + kFound, + kBusy, + kFailed, +}; + +struct ResolvedEntry { + ResolveOutcome outcome = ResolveOutcome::kMissing; + // Kept for source and wire compatibility. New code should set outcome too; + // readers accept found=true as kFound for older backend implementations. + bool found = false; + std::vector pages; + uint64_t size = 0; + uint64_t page_size = 0; + // Empty when the caller passed include_descs=false (it already hydrated them + // from GetPeerInfo), or when found=false. + std::vector descs; + + // Set instead of `pages` when the backend serves this key by a zero-copy file + // range (SsdBackend + GdsEngine): a File TransferRef the reader DMAs straight + // into device memory, with `pages` left empty. A memory-backed medium leaves + // this invalid and publishes pages as before. + TransferRef file_ref{}; +}; + +inline ResolveOutcome EffectiveResolveOutcome(const ResolvedEntry& entry) { + return entry.found ? ResolveOutcome::kFound : entry.outcome; +} + +struct EvictResult { + std::string key; + // 0 if the key was unknown, already freed, or protected (read lease / pin). + // The master retries protected keys on a later round. + uint64_t bytes_freed = 0; +}; + +class MediumBackend : public MetricSource { + public: + virtual ~MediumBackend() = default; + + MediumBackend(const MediumBackend&) = delete; + MediumBackend& operator=(const MediumBackend&) = delete; + + // ---- identity + routing advertisement ---- + + virtual TierType Tier() const = 0; + virtual const char* Name() const = 0; + + // ---- ownership ---- + + // Allocate this medium's pool with its own policy (hugepage/NUMA for DRAM, + // hipMalloc for HBM, file extents for SSD) and register it with `registrar`, + // supplying the facts a descriptor cannot recover: location type, device, + // NUMA node. PoolClient hands over the registrar; it does NOT hand over + // memory. + // + // This is the inversion the refactor turns on. Before Phase 2 allocation + // lived in DistributedClient and registration in PoolClient::Init with + // MemoryLocationType::CPU hardcoded, so a "backend" was a bookkeeper over + // memory it did not own (design doc §1 item 4). + virtual bool Init(MemoryRegistrar* registrar) = 0; + + // Deregister and release the pool. Must tolerate being called after a failed + // Init, and must not run concurrently with any other method. + virtual void Shutdown() = 0; + + // ---- control plane (shipped by the heartbeat) ---- + + // Live capacity, derived from the backend's own allocator so it always + // reflects pending + owned. Cheap enough to call every heartbeat. + virtual TierCapacity Capacity() const = 0; + virtual uint64_t OwnedKeyCount() const = 0; + + // Events queued since the last call (delta heartbeat). Clears the outbox. + virtual std::vector DrainPendingEvents() = 0; + + // Every owned key as ADD events. const: a snapshot must not mutate state. + virtual std::vector SnapshotOwnedKeys() const = 0; + + // Full-sync snapshot that ALSO atomically drops the outbox, in the SAME + // critical section. After a full sync the snapshot is authoritative, so a + // still-queued event is already reflected in it and must not be re-shipped + // as a redundant delta. Two separate locks would drop events committed in + // between. + virtual std::vector SnapshotOwnedKeysForFullSync() = 0; + + // Distributed Clear: drop owned state, invalidate outstanding slots so they + // fail at commit, and clear the outbox. Allocation stays gated until + // ClearFullSyncAcked(). + virtual void ClearLocal() = 0; + + // Called by the heartbeat thread once master has acked the first full-sync + // empty snapshot. Re-enables allocation. + virtual void ClearFullSyncAcked() = 0; + + // True between ClearLocal() and ClearFullSyncAcked(). + virtual bool IsClearFullSyncPending() const = 0; + + // Register a hook the backend invokes once its unshipped-event outbox + // reaches `threshold`, so a batch of puts flushes the heartbeat instead of + // waiting for the next tick. On the interface (rather than on the concrete + // backend) because every medium's outbox feeds the same single-bundle + // heartbeat: a newly registered backend must participate without MasterClient + // learning its type. `cb` MUST be cheap and MUST NOT re-enter the backend — + // it should only signal the heartbeat thread. Pass a very large threshold to + // disable. + virtual void SetAutoFlushHook(size_t threshold, std::function cb) = 0; + + // Turn the heartbeat outbox off. Every ADD/REMOVE a backend queues is + // addressed to a master, and the ONLY consumer is MasterClient's heartbeat + // (DrainAllBackends below). A node with no master therefore queues an event + // per put that nothing will ever drain — an unbounded outbox on the one + // deployment that has no reason to keep one. PoolClient::Init turns + // publishing off when it builds no MasterClient; a backend that ignores the + // call (the default) simply has no outbox worth gating. + virtual void SetEventPublishing(bool /*enabled*/) {} + + // Turn on the backend's OWN eviction, for a node with no master. + // + // Evict(keys) above is a master decision: it picks victims with a + // cluster-wide view of who else holds the key. Nothing in the system calls + // it otherwise, so a masterless node needs some medium-local policy or it + // simply fills up. A backend that already self-evicts (SsdBackend, whose + // PeerSsdManager runs its own watermark loop from UMBPSsdConfig) ignores + // this; the default is to ignore it. + virtual void EnableLocalEviction(double /*high_watermark*/, double /*low_watermark*/) {} + + // ---- observability ---- + // + // A backend does NOT report its own operation counts, byte volumes or + // latencies. InstrumentedBackend sits on this interface and derives all of + // those from the calls below, which is why a new medium is fully observable + // the moment PoolClient composes it in — see instrumented_backend.h. + // + // MetricSource::SampleMetrics(), inherited here, is for the one thing a + // decorator cannot see: state INSIDE the medium (the drive's own read + // outcomes, single-flight coalescing, a staging arena's occupancy). Publish + // those under the generic MORI_UMBP_METRIC_BACKEND_MEDIUM_* names with the + // specifics in an event=/state= label, never under a name that spells the + // medium — a metric called mori_umbp_ssd_something forces a second dashboard, + // which is exactly what the single-dashboard layout removed. + // + // tier= and backend= are added by the publisher; do not set them here. + + // ---- slot lifecycle (peer-side; driven by PeerServiceServer handlers) ---- + // + // Deliberately not called "data plane": these calls hand out and retire slots + // and publish descriptors for them. The bytes move through the transfer + // engine, not through here. + + // One result per request, in request order. + virtual std::vector BatchAllocate(const std::vector&) = 0; + virtual std::vector BatchCommit(const std::vector&) = 0; + + // Idempotent: a slot already reaped or never seen still reports true. False + // is reserved for a state not currently produced. + virtual std::vector BatchAbort(const std::vector& slot_ids) = 0; + + // Look up keys and take a read lease on each hit, protecting it from + // concurrent Evict for the lease window. `include_descs=false` skips + // building descs entirely for the whole batch. A medium whose own bytes are + // not registerable publishes a descriptor naming its storage (a file/object + // ref) rather than a memory one — the engine decides how to reach it, and + // supplies a bounce buffer if the chosen path needs one. + // + // `allow_file_refs` says the caller can consume a TransferRef::File -- it + // will read the medium's own storage directly (GdsEngine DMAs the range into + // the destination) instead of a staged page. It defaults to FALSE because a + // file ref only means something IN THIS PROCESS: an fd plus a locally + // registered hipFile handle is meaningless to a peer, and a resolve that put + // one on the wire would ship a slot carrying no pages at all. Callers that + // serve a remote reader, and callers that only test `found`, leave it off. + // Every declaration of this method repeats the same default deliberately -- + // a default argument on a virtual binds statically, so they must not drift. + virtual std::vector BatchResolve(const std::vector& keys, + bool include_descs, + bool allow_file_refs = false) = 0; + + // Acquire a peer-local migration pin and resolve the key without changing or + // cancelling normal RPC read leases. The caller must release a successful + // acquisition after its synchronous transfer completes. + virtual bool AcquireMigrationRead(const std::string& key, ResolvedEntry* resolved) { + (void)key; + (void)resolved; + return false; + } + virtual void ReleaseMigrationRead(const std::string& key) { (void)key; } + + // Authoritative metadata lookup with no staging, descriptor construction, or + // read-lease side effects. Pool placement validation must use this instead of + // interpreting a temporarily unresolvable object as absent. + virtual bool Contains(const std::string& key) const = 0; + + // Containment for a whole batch, in request order. Exists answers from this + // rather than from BatchResolve because a resolve builds -- and heap-copies -- + // the page vector of every key it finds, which an existence check throws away. + // The default is the honest per-key loop; a backend that can answer the batch + // under one lock should override it. + virtual std::vector BatchContains(const std::vector& keys) const { + std::vector out(keys.size(), false); + for (size_t i = 0; i < keys.size(); ++i) out[i] = Contains(keys[i]); + return out; + } + + // Master-driven eviction. Idempotent; see EvictResult::bytes_freed. One + // result per key, in request order — the peer service relies on that to sum + // freed bytes for a key mirrored across media. + virtual std::vector Evict(const std::vector& keys) = 0; + + // ---- bootstrap (GetPeerInfo) ---- + + // Page size of this medium. Also mirrored into every AllocateResult / + // ResolvedEntry so a batch spanning backends stays self-describing. + virtual uint64_t PageSize() const = 0; + + // Descs for ALL configured buffers, ascending by buffer_index, so a + // first-contact writer can hydrate its peer-side cache in one round trip + // instead of learning buffers one allocation at a time. + virtual std::vector AllBufferDescs() const = 0; + + // ---- local (same-process) endpoints ---- + + // How many buffers this backend published; PageLocation::buffer_index is an + // index into that range. Zero for a backend with no page-addressable + // buffers. + virtual size_t BufferCount() const = 0; + + // The endpoint for one buffer, in the form a same-process transfer plans + // against — i.e. the UNPACKED, in-process counterpart of + // AllBufferDescs()[i], with no msgpack round trip on the local hot path. + // Returns an invalid ref for an out-of-range index. + // + // This is what replaced PageBackend::LocalBufferViews() in Phase 6, and the + // difference is the point: a TransferRef is medium-agnostic (§2: "bytes + // addressed by a descriptor are already medium-agnostic"), a raw base pointer + // was not. Publishing refs instead of pointers is what let PoolClient stop + // naming a concrete backend type, and what makes a second medium's local + // access work with no tier branch in a copy loop. + virtual TransferRef BufferRef(uint32_t buffer_index) const = 0; + + protected: + MediumBackend() = default; +}; + +// Owns every backend instance live on this peer. Names are stable +// configuration identities; backend ids are dense peer-local wire identities. +class BackendRegistry { + public: + // Defined in types.h so the transfer layer can size its per-backend buffer + // shelves without depending on this header (see kMaxBackendsPerPeer). + static constexpr uint32_t kMaxBackends = kMaxBackendsPerPeer; + + struct Entry { + uint32_t backend_id = 0; + std::string name; + TierType tier = TierType::UNKNOWN; + std::unique_ptr backend; + }; + + // Legacy registration keeps one conventional name per tier. Registering the + // same name again replaces that instance and preserves its backend id. + bool Register(std::unique_ptr backend) { + if (backend == nullptr) return true; + return Register(DefaultBackendInstanceName(backend->Tier()), std::move(backend)); + } + + bool Register(std::string name, std::unique_ptr backend) { + if (backend == nullptr) return true; + if (name.empty()) { + MORI_UMBP_ERROR("[BackendRegistry] backend instance name must not be empty"); + return false; + } + + const TierType tier = backend->Tier(); + auto existing = id_by_name_.find(name); + const uint32_t replacing_id = existing == id_by_name_.end() ? kMaxBackends : existing->second; + + // Every wire carrying pages has one peer-global page size. When replacing + // an entry, compare against the other live instances so replacing the sole + // backend may legitimately establish a new size. + uint64_t expected_page_size = 0; + for (const auto& entry : entries_) { + if (entry->backend_id == replacing_id) continue; + expected_page_size = entry->backend->PageSize(); + break; + } + if (expected_page_size != 0 && backend->PageSize() != expected_page_size) { + MORI_UMBP_ERROR( + "[BackendRegistry] backend '{}' tier={} page_size={} disagrees with the peer's " + "page_size={}; every instance on a node must agree", + name, static_cast(tier), backend->PageSize(), expected_page_size); + return false; + } + + if (existing != id_by_name_.end()) { + auto& entry = entries_[existing->second]; + entry->tier = tier; + entry->backend = std::move(backend); + page_size_ = entries_.front()->backend->PageSize(); + return true; + } + + if (entries_.size() >= kMaxBackends) { + MORI_UMBP_ERROR("[BackendRegistry] too many backends (max {})", kMaxBackends); + return false; + } + const uint32_t id = static_cast(entries_.size()); + auto entry = std::make_unique(); + entry->backend_id = id; + entry->name = std::move(name); + entry->tier = tier; + entry->backend = std::move(backend); + id_by_name_[entry->name] = id; + entries_.push_back(std::move(entry)); + page_size_ = entries_.front()->backend->PageSize(); + return true; + } + + // The backend's peer-local id, as stamped onto every buffer descriptor and + // page set it produces. kMaxBackends for an unregistered backend. + uint32_t BackendId(const MediumBackend* backend) const { + if (backend == nullptr) return kMaxBackends; + for (const auto& entry : entries_) { + if (entry->backend.get() == backend) return entry->backend_id; + } + return kMaxBackends; + } + + // Uniform across every registered backend (see Register). Zero when empty. + uint64_t PageSize() const { return page_size_; } + + // Null if the tier has no live backend on this peer — a normal condition + // (e.g. a node configured with DRAM only), not an error. Callers respond to + // a request for an absent tier with found=false / success=false. + MediumBackend* Get(TierType tier) const { + for (const auto& entry : entries_) { + if (entry->tier == tier) return entry->backend.get(); + } + return nullptr; + } + + MediumBackend* Get(const std::string& name) const { + const Entry* entry = GetEntry(name); + return entry == nullptr ? nullptr : entry->backend.get(); + } + + MediumBackend* Get(uint32_t backend_id) const { + const Entry* entry = GetEntry(backend_id); + return entry == nullptr ? nullptr : entry->backend.get(); + } + + const Entry* GetEntry(const std::string& name) const { + auto it = id_by_name_.find(name); + return it == id_by_name_.end() ? nullptr : GetEntry(it->second); + } + + Entry* GetEntry(const std::string& name) { + return const_cast(std::as_const(*this).GetEntry(name)); + } + + const Entry* GetEntry(uint32_t backend_id) const { + return backend_id < entries_.size() ? entries_[backend_id].get() : nullptr; + } + + Entry* GetEntry(uint32_t backend_id) { + return const_cast(std::as_const(*this).GetEntry(backend_id)); + } + + const Entry* GetEntry(const MediumBackend* backend) const { + const uint32_t id = BackendId(backend); + return id == kMaxBackends ? nullptr : GetEntry(id); + } + + bool Empty() const { return entries_.empty(); } + size_t Size() const { return entries_.size(); } + + // Every live backend in dense backend-id order. + std::vector All() const { + std::vector out; + out.reserve(entries_.size()); + for (const auto& entry : entries_) out.push_back(entry->backend.get()); + return out; + } + + private: + std::vector> entries_; + std::unordered_map id_by_name_; + uint64_t page_size_ = 0; +}; + +// --------------------------------------------------------------------------- +// Heartbeat event aggregation +// +// These absorb the deleted OwnedLocationSource free functions (design doc §3: +// "This interface absorbs the existing OwnedLocationSource — that type goes +// away"). Exposed rather than private to MasterClient so they can be unit +// tested against backends without standing up a master RPC. +// --------------------------------------------------------------------------- + +// Drain every backend and concat into ONE event list, in the given order. The +// heartbeat shipper wraps the result in a SINGLE EventBundle under one +// monotonic seq — concat here, never one bundle/seq per medium (that breaks the +// ack / seq-gap full-sync recovery). Null entries are skipped. +inline std::vector DrainAllBackends(const std::vector& backends) { + std::vector raw; + for (auto* backend : backends) { + if (backend == nullptr) continue; + auto events = backend->DrainPendingEvents(); + raw.insert(raw.end(), std::make_move_iterator(events.begin()), + std::make_move_iterator(events.end())); + } + + // Master locations are (node, physical tier), not backend instances. A + // same-tier migration therefore must not publish the source REMOVE after the + // target ADD as if the node had lost the key. Coalesce by (tier,key) against + // authoritative current ownership across all instances. + struct EventState { + std::optional add; + bool saw_remove = false; + }; + std::map, EventState> states; + for (auto& event : raw) { + auto& state = states[{event.tier, event.key}]; + if (event.kind == KvEvent::Kind::ADD) { + state.add = std::move(event); + } else { + state.saw_remove = true; + } + } + std::vector merged; + merged.reserve(states.size()); + for (auto& [identity, state] : states) { + const auto& [tier, key] = identity; + const bool still_owned = + std::any_of(backends.begin(), backends.end(), [&](MediumBackend* backend) { + return backend != nullptr && backend->Tier() == tier && backend->Contains(key); + }); + if (state.add.has_value() && (still_owned || !state.saw_remove)) { + merged.push_back(std::move(*state.add)); + } else if (state.saw_remove && !still_owned) { + merged.push_back(KvEvent{KvEvent::Kind::REMOVE, key, tier, 0}); + } + } + return merged; +} + +// Full-sync counterpart: snapshot every backend and concat, each backend ALSO +// atomically clearing its outbox in the same critical section (the snapshot is +// authoritative, so the queued delta is redundant and must not be re-shipped). +// Null entries are skipped. Non-const because it mutates the backends. +inline std::vector SnapshotAllBackendsForFullSync( + const std::vector& backends) { + std::vector merged; + for (auto* backend : backends) { + if (backend == nullptr) continue; + auto events = backend->SnapshotOwnedKeysForFullSync(); + merged.insert(merged.end(), std::make_move_iterator(events.begin()), + std::make_move_iterator(events.end())); + } + return merged; +} + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/backend/mock_backend.h b/src/umbp/include/umbp/distributed/peer/backend/mock_backend.h new file mode 100644 index 000000000..05ec0a43b --- /dev/null +++ b/src/umbp/include/umbp/distributed/peer/backend/mock_backend.h @@ -0,0 +1,283 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +// Second registered backend proving BackendRegistry dispatch (design doc +// Phase 2, default option (b)): "MockBackend as the second registered +// backend. Cheap, proves the registry and dispatch paths, keeps the local +// path honestly single-medium." +// +// Deliberately NOT wired into live routing: Capacity() always reports zero, +// so master never picks this node's MockBackend-fronted tier as a Put/Get +// target — the DRAM backend's behavior stays byte-for-byte unchanged with +// this second backend registered alongside it. The point of this class is +// to be a real, in-memory, fully-functional MediumBackend implementation +// that a test can drive through BackendRegistry — not a no-op stub. +class MockBackend : public MediumBackend { + public: + // The registry requires every backend on a peer to agree on a page size, so + // a mock sharing a registry with a real page-based backend must be told which + // one to report. + explicit MockBackend(TierType tier, uint64_t page_size = 1) + : tier_(tier), page_size_(page_size) {} + + TierType Tier() const override { return tier_; } + const char* Name() const override { return "MockBackend"; } + + bool Init(MemoryRegistrar* /*registrar*/) override { + initialized_ = true; + return true; + } + void Shutdown() override { + std::lock_guard lock(mutex_); + owned_.clear(); + pending_.clear(); + initialized_ = false; + } + + // Always zero: this backend holds real bytes (for dispatch tests) but + // never advertises capacity, so master-side routing never selects it. + TierCapacity Capacity() const override { return TierCapacity{}; } + + uint64_t OwnedKeyCount() const override { + std::lock_guard lock(mutex_); + return owned_.size(); + } + + std::vector DrainPendingEvents() override { + std::lock_guard lock(mutex_); + std::vector out; + out.swap(pending_events_); + return out; + } + std::vector SnapshotOwnedKeys() const override { + std::lock_guard lock(mutex_); + return SnapshotLocked(); + } + std::vector SnapshotOwnedKeysForFullSync() override { + std::lock_guard lock(mutex_); + auto out = SnapshotLocked(); + pending_events_.clear(); + return out; + } + + void ClearLocal() override { + std::lock_guard lock(mutex_); + owned_.clear(); + pending_.clear(); + pending_events_.clear(); + } + void ClearFullSyncAcked() override {} + bool IsClearFullSyncPending() const override { return false; } + + void SetAutoFlushHook(size_t threshold, std::function cb) override { + std::lock_guard lock(mutex_); + auto_flush_threshold_ = threshold; + auto_flush_cb_ = std::move(cb); + } + + std::vector BatchAllocate(const std::vector& entries) override { + std::vector out(entries.size()); + std::lock_guard lock(mutex_); + for (size_t i = 0; i < entries.size(); ++i) { + const auto& req = entries[i]; + if (req.size == 0) { + out[i].outcome = AllocateOutcome::kFailed; + continue; + } + if (owned_.find(req.key) != owned_.end()) { + out[i].outcome = AllocateOutcome::kSuccessAlreadyExists; + continue; + } + const uint64_t slot_id = next_slot_id_++; + pending_[slot_id] = {req.key, std::vector(req.size)}; + out[i].outcome = AllocateOutcome::kSuccessAllocated; + out[i].slot_id = slot_id; + out[i].size = req.size; + out[i].page_size = req.size; + out[i].pending_ttl_ms = 0; + } + return out; + } + + std::vector BatchCommit(const std::vector& entries) override { + std::vector out(entries.size()); + std::lock_guard lock(mutex_); + for (size_t i = 0; i < entries.size(); ++i) { + auto it = pending_.find(entries[i].slot_id); + if (it == pending_.end()) continue; + const uint64_t committed_size = it->second.bytes.size(); + owned_[entries[i].key] = std::move(it->second.bytes); + QueueEventLocked(KvEvent{KvEvent::Kind::ADD, entries[i].key, tier_, committed_size}); + out[i].success = true; + out[i].bytes_committed = committed_size; + pending_.erase(it); + } + return out; + } + + std::vector BatchAbort(const std::vector& slot_ids) override { + std::lock_guard lock(mutex_); + std::vector out(slot_ids.size(), true); + for (uint64_t id : slot_ids) pending_.erase(id); + return out; + } + + std::vector BatchResolve(const std::vector& keys, bool include_descs, + bool allow_file_refs = false) override { + (void)allow_file_refs; // this mock stages everything; it publishes no file refs + std::vector out(keys.size()); + std::lock_guard lock(mutex_); + for (size_t i = 0; i < keys.size(); ++i) { + auto it = owned_.find(keys[i]); + if (it == owned_.end()) continue; + out[i].outcome = ResolveOutcome::kFound; + out[i].found = true; + out[i].size = it->second.size(); + if (published_buffers_ == 0) continue; + // Every published mock hands out buffer 0, which is the collision a real + // peer has too (each backend publishes exactly one buffer today). + out[i].pages = {PageLocation{0, 0}}; + out[i].page_size = PageSize(); + if (include_descs) out[i].descs = AllBufferDescsLocked(); + } + return out; + } + + bool Contains(const std::string& key) const override { + std::lock_guard lock(mutex_); + return owned_.find(key) != owned_.end(); + } + + std::vector Evict(const std::vector& keys) override { + std::vector out; + out.reserve(keys.size()); + std::lock_guard lock(mutex_); + for (const auto& key : keys) { + EvictResult r; + r.key = key; + auto it = owned_.find(key); + if (it != owned_.end()) { + r.bytes_freed = it->second.size(); + QueueEventLocked(KvEvent{KvEvent::Kind::REMOVE, key, tier_, 0}); + owned_.erase(it); + } + out.push_back(std::move(r)); + } + return out; + } + + uint64_t PageSize() const override { return page_size_; } + + // Publish `count` fake buffers, numbered from 0 — like every real backend, + // and the reason a bare buffer_index is not an address: two mocks that both + // publish will both claim buffer 0. `desc_tag` is the first descriptor byte + // so a test can tell WHOSE buffer answered. Off by default, so a mock that + // does not call this behaves exactly as before. + void PublishBuffers(uint32_t count, uint8_t desc_tag) { + std::lock_guard lock(mutex_); + published_buffers_ = count; + desc_tag_ = desc_tag; + } + + std::vector AllBufferDescs() const override { + std::lock_guard lock(mutex_); + return AllBufferDescsLocked(); + } + + // Zero unless PublishBuffers() was called; then each key resolves to a page in + // buffer 0. A backend keeps each key in its own std::vector rather than in + // paged buffers, so with no published buffers a caller that tries to transfer + // against it gets an invalid ref and must treat the key as unservable here — + // which is exactly the shape a real medium takes when a tier has no live + // backend on this node. + size_t BufferCount() const override { + std::lock_guard lock(mutex_); + return published_buffers_; + } + TransferRef BufferRef(uint32_t /*buffer_index*/) const override { return TransferRef{}; } + + private: + // Backends leave backend_id at 0 — the peer service stamps the owner at the + // wire boundary (see BufferMemoryDescBytes). The mock does the same so a + // test exercises the real stamping path rather than pre-tagging for it. + std::vector AllBufferDescsLocked() const { + std::vector out; + out.reserve(published_buffers_); + for (uint32_t i = 0; i < published_buffers_; ++i) { + out.push_back( + {i, std::vector{desc_tag_, static_cast(i)}, /*backend_id=*/0}); + } + return out; + } + + uint32_t published_buffers_ = 0; + uint8_t desc_tag_ = 0; + + struct PendingSlot { + std::string key; + std::vector bytes; + }; + + // Mirrors PageBackend::QueueEventLocked: the hook fires under the lock, so + // the callback must not re-enter. Caller MUST hold `mutex_`. + void QueueEventLocked(KvEvent event) { + pending_events_.push_back(std::move(event)); + if (auto_flush_cb_ && pending_events_.size() >= auto_flush_threshold_) auto_flush_cb_(); + } + + std::vector SnapshotLocked() const { + std::vector out; + out.reserve(owned_.size()); + for (const auto& [key, bytes] : owned_) { + out.push_back(KvEvent{KvEvent::Kind::ADD, key, tier_, bytes.size()}); + } + return out; + } + + TierType tier_; + uint64_t page_size_ = 1; + bool initialized_ = false; + mutable std::mutex mutex_; + uint64_t next_slot_id_ = 1; + std::unordered_map pending_; + std::unordered_map> owned_; + std::vector pending_events_; + size_t auto_flush_threshold_ = SIZE_MAX; // events before a flush; default = never + std::function auto_flush_cb_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/backend/page_backend.h b/src/umbp/include/umbp/distributed/peer/backend/page_backend.h new file mode 100644 index 000000000..d0d8840c0 --- /dev/null +++ b/src/umbp/include/umbp/distributed/peer/backend/page_backend.h @@ -0,0 +1,633 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/io/engine.hpp" +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/peer/backend/page_pool.h" +#include "umbp/distributed/peer/backend/peer_page_allocator.h" +#include "umbp/distributed/types.h" +#include "umbp/local/host_mem_allocator.h" + +namespace mori::umbp { + +// --------------------------------------------------------------------------- +// PageMemorySource — the ONE medium-specific part of a paged backend +// +// PageBackend's class comment has always claimed it serves "DRAM or HBM", and +// everything below the ownership line genuinely was medium-agnostic: the +// bitmap allocator, slot lifecycle, event outbox, reaper, read leases and copy +// pins never look at the tier. But Init/Shutdown/BuildBufferRefs hardcoded +// HostMemAllocator, MemoryLocationType::CPU and device=-1, so the claim did not +// survive contact with an actual second medium. +// +// This interface is that hardcoding, lifted. It holds exactly the three things +// that differ per medium, and nothing else: +// * how the pool is allocated and released +// * the location type a descriptor cannot recover +// * the device a descriptor cannot recover +// +// This is the same rule the MediumBackend header states for the backend/ +// transfer split, applied one level down: registration is the conversion from +// medium-specific to medium-agnostic, and belongs to whoever allocated. A +// PageMemorySource is "whoever allocated"; PageBackend is everything after. +// +// Adding a paged medium is therefore implementing FIVE methods, not +// reimplementing MediumBackend's twenty-three. +class PageMemorySource { + public: + virtual ~PageMemorySource() = default; + + PageMemorySource(const PageMemorySource&) = delete; + PageMemorySource& operator=(const PageMemorySource&) = delete; + + struct Buffer { + void* base = nullptr; + // Usable size. May EXCEED the requested size (host hugepage rounding) — + // PageBackend publishes this, not the request, so the extra is allocatable. + uint64_t size = 0; + }; + + // Allocate one buffer per entry of `sizes`, appending to `out`. All or + // nothing: on failure the implementation must release whatever it took in + // this call and return false, leaving `out` untouched. Zero-sized entries + // are skipped, so `out` may be shorter than `sizes`. + virtual bool Allocate(const std::vector& sizes, std::vector* out) = 0; + + // Free everything Allocate handed out. Idempotent; must tolerate being + // called after a failed Allocate and without one. + virtual void Release() = 0; + + // The facts a descriptor cannot recover, supplied to RegisterMemory and + // mirrored into every TransferRef this backend publishes. + virtual mori::io::MemoryLocationType LocationType() const = 0; + virtual int Device() const = 0; + + virtual const char* Name() const = 0; + + protected: + PageMemorySource() = default; +}; + +// Host DRAM: anonymous or hugetlb pages via HostMemAllocator, optionally +// NUMA-bound and prefaulted. loc=CPU, device=-1. This is verbatim what +// PageBackend::Init used to do inline. +class HostPageMemorySource final : public PageMemorySource { + public: + struct Options { + bool use_hugepages = false; + uint64_t hugepage_size = 2ULL * 1024 * 1024; + int numa_node = -1; + bool prefault = true; + }; + + explicit HostPageMemorySource(Options opts) : opts_(opts) {} + ~HostPageMemorySource() override { Release(); } + + bool Allocate(const std::vector& sizes, std::vector* out) override; + void Release() override; + + mori::io::MemoryLocationType LocationType() const override { + return mori::io::MemoryLocationType::CPU; + } + int Device() const override { return -1; } + const char* Name() const override { return "HostPageMemorySource"; } + + private: + Options opts_; + std::vector handles_; +}; + +// One medium's (DRAM or HBM) peer-owned allocator + key map. One instance == +// one medium: the tier is the identity of the object (design doc §3), which is +// why no method below takes a TierType — PeerDramAllocator's internal +// map is gone, relocated to BackendRegistry (one PageBackend +// instance per registry entry). +// +// In the master-as-advisor design, this object is the canonical owner of +// every per-key page set on the local node — master's GlobalBlockIndex is a +// downstream projection of the events this object emits. +// +// Lifecycle of a Put on this peer: +// 1. Allocate(size) -> AllocateResult { slot_id, pages, ... } +// 2. (Writer RDMAs into the slot's pages.) +// 3. Commit(slot_id, key) -> moves pending -> owned, queues an ADD event. +// OR Abort(slot_id) -> drops pending; no event. +// If the writer dies between (1) and (3), the reaper drops the pending +// slot at `pending_ttl` and no event is ever emitted. +// +// Lifecycle of an Evict (master-driven): +// 1. Master sends EvictKey(keys[]) — handled by the peer service. +// 2. Peer service calls Evict(keys) here. +// 3. For each key actually freed, a REMOVE event is queued. +// 4. Heartbeat ships the REMOVE events and master drops the index entry. +// +// MediumBackend is the only base: the Phase 2 transitional second base +// (OwnedLocationSource) is gone, along with the type itself — MasterClient now +// consumes BackendRegistry directly and MediumBackend carries the heartbeat +// event contract (design doc §3, "this interface absorbs the existing +// OwnedLocationSource"). +class PageBackend : public MediumBackend { + public: + // Buffers already known (bases/descs supplied by the caller at + // construction). Used directly by tests that want to exercise the + // allocator's bookkeeping without a real TransferEngine/IOEngine — + // Init() on an instance built this way is a no-op success (already + // configured). Production backends use the OwnershipConfig constructor + // below instead, so no buffer pointer ever crosses PoolClientConfig + // (design doc §1 item 4 / Phase 2b). + struct TierConfig { + std::vector buffer_sizes; + std::vector> buffer_descs; // packed mori::io::MemoryDesc bytes + // Local host base pointer of each buffer (same order as buffer_sizes). + // Used by AcquireDramCopyPin to resolve a (buffer_index, page_index) + // page to a directly-readable local pointer for the SSD copy worker. + // Empty is allowed for tiers/deployments with no DRAM copy (e.g. unit + // tests that never pin); pin acquisition then yields no segments. + std::vector buffer_bases; + }; + + // Ownership config: sizes only. Init(TransferEngine*) self-allocates one + // HostMemAllocator buffer per entry and registers it with the engine — + // PageBackend is the allocator now, not a bookkeeper over memory handed to + // it from outside (design doc §1 item 4). + struct OwnershipConfig { + std::vector buffer_sizes; + bool use_hugepages = false; + uint64_t hugepage_size = 2ULL * 1024 * 1024; + int numa_node = -1; + bool prefault = true; + }; + + // `pending_ttl` is the only TTL in the system. After this elapses + // without a matching Commit / Abort, the reaper frees the slot's + // pages. `read_lease_ttl` is how long a single Resolve protects its + // key from concurrent Evict. `reaper_interval` is the wakeup cadence + // for sweeping expired pendings and read leases. + PageBackend(TierType tier, uint64_t page_size, TierConfig cfg, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl = std::chrono::milliseconds{500}, + std::chrono::milliseconds reaper_interval = std::chrono::milliseconds{200}); + + // Production constructor: no memory yet. Init(TransferEngine*) performs + // the actual self-allocation + registration. Sugar for the + // PageMemorySource constructor below with a HostPageMemorySource — kept so + // existing DRAM call sites and their tests read unchanged. + PageBackend(TierType tier, uint64_t page_size, OwnershipConfig ownership, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl = std::chrono::milliseconds{500}, + std::chrono::milliseconds reaper_interval = std::chrono::milliseconds{200}); + + // General production constructor: any medium, supplied as a PageMemorySource. + // This is the one Init() actually works against; the OwnershipConfig overload + // funnels into it, so DRAM and HBM share a single allocation/registration + // path rather than two that must be kept in step. + PageBackend(TierType tier, uint64_t page_size, std::unique_ptr source, + std::vector buffer_sizes, std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl = std::chrono::milliseconds{500}, + std::chrono::milliseconds reaper_interval = std::chrono::milliseconds{200}); + + ~PageBackend() override; + + PageBackend(const PageBackend&) = delete; + PageBackend& operator=(const PageBackend&) = delete; + + // ==================== MediumBackend ==================== + + TierType Tier() const override { return tier_; } + const char* Name() const override { return "PageBackend"; } + + // No-op success if constructed with a pre-built TierConfig (already + // configured; the caller drives the reaper itself via RunReaperOnceForTest). + // Otherwise self-allocates OwnershipConfig::buffer_sizes via HostMemAllocator, + // registers each buffer with `registrar`, and starts the reaper — a backend + // is fully live when Init returns, so no caller has to know it has a reaper. + bool Init(MemoryRegistrar* registrar) override; + void Shutdown() override; + + TierCapacity Capacity() const override; + uint64_t OwnedKeyCount() const override; + + std::vector DrainPendingEvents() override; + std::vector SnapshotOwnedKeys() const override; + std::vector SnapshotOwnedKeysForFullSync() override; + + void ClearLocal() override; + void ClearFullSyncAcked() override; + bool IsClearFullSyncPending() const override { + return clear_full_sync_pending_.load(std::memory_order_acquire); + } + + std::vector BatchAllocate(const std::vector& entries) override; + std::vector BatchCommit(const std::vector& entries) override; + std::vector BatchAbort(const std::vector& slot_ids) override; + std::vector BatchResolve(const std::vector& keys, bool include_descs, + bool allow_file_refs = false) override; + bool AcquireMigrationRead(const std::string& key, ResolvedEntry* resolved) override; + void ReleaseMigrationRead(const std::string& key) override; + bool Contains(const std::string& key) const override; + std::vector BatchContains(const std::vector& keys) const override; + std::vector Evict(const std::vector& keys) override; + + uint64_t PageSize() const override { return page_size_; } + std::vector AllBufferDescs() const override; + + size_t BufferCount() const override { return buffer_refs_.size(); } + TransferRef BufferRef(uint32_t buffer_index) const override; + + // ==================== Single-key convenience ==================== + // NOT part of MediumBackend (design doc §3: "single-key Allocate/Commit/ + // Abort/Resolve — the single-key RPCs are served by one-element batches"). + // Kept as dedicated (non-wrapping) methods for the hot single-key path in + // PoolClient/PeerServiceServer — same shape as before, minus the tier + // parameter. `descs` is left empty on the AllocateResult/ResolvedEntry + // returned here; callers that need descs call BufferDescsForPages + // separately (mirrors pre-refactor behavior). + + AllocateResult Allocate(const std::string& key, uint64_t size); + bool Commit(uint64_t slot_id, const std::string& key, uint64_t& bytes_committed); + bool Abort(uint64_t slot_id); + ResolvedEntry Resolve(const std::string& key); + + uint64_t PendingTtlMs() const { return static_cast(pending_ttl_.count()); } + + // Returns the deduplicated, ascending-by-buffer_index list of descs + // referenced by `pages`. + std::vector BufferDescsForPages( + const std::vector& pages) const; + + // ==================== DRAM copy pin (SSD copy-on-commit) ==================== + + // A copy pin protects a committed key's pages from eviction while the SSD + // copy worker reads them. It is an in-process lifetime guard (NOT a + // master lease): pages stay readable until ReleaseDramCopyPin. `segments` + // are directly-readable local pointers + lengths (pages may be + // non-contiguous across buffers). + struct DramCopyPin { + std::vector> segments; + size_t total_size = 0; + uint64_t pin_token = 0; // release-time validation + }; + + // Atomically (under mutex_) confirm `key` is owned, mark it pinned, and + // resolve its pages to local segments. Returns nullopt if the key is + // not owned (already evicted -> worker drops the task) or is already + // pinned (duplicate task). While pinned, Evict(key) reports + // bytes_freed=0 and emits no REMOVE (master retries next round). There + // is NO TTL: the pin lives until ReleaseDramCopyPin. The caller MUST + // release (use a RAII guard) so a worker exit always frees the pin. + std::optional AcquireDramCopyPin(const std::string& key); + + // Release a pin. No-op if the key is not pinned or the token does not + // match (tolerates duplicate / late release). + void ReleaseDramCopyPin(const std::string& key, uint64_t pin_token); + + // ==================== Heartbeat auto-flush ==================== + + // MediumBackend contract. The hook is invoked while the allocator lock is + // held, so `cb` must not re-enter this object. `threshold` should be >= 1. + void SetAutoFlushHook(size_t threshold, std::function cb) override; + + // MediumBackend contract. False stops queueing ADD/REMOVE and drops + // whatever is already queued: with no master nothing drains the outbox, so + // keeping one is a leak of one event per put. + void SetEventPublishing(bool enabled) override; + + // ==================== Local (masterless) eviction ==================== + // + // Eviction here is normally the MASTER's policy: it chooses victims with a + // cluster-wide view and calls Evict(keys) through the peer service. Nothing + // else ever calls Evict, so a node with no master fills its pool and + // Allocate starts answering NO_SPACE forever — where the deleted local + // backend's LocalStorageManager would have evicted on a watermark. This is + // that watermark loop, restored on the backend that now serves embedded + // deployments. + // + // Opt-in so the cluster path is bit-for-bit unchanged: with a master, a + // second local policy would free keys the master still indexes and hands out + // routes for. PoolClient::Init enables it exactly when it builds no + // MasterClient. + // + // Enabling it turns on an LRU ordering (inserted by Commit, renewed by + // Resolve) that is not maintained otherwise, so the read path pays nothing + // in a cluster. A round runs on the committing thread — there is no + // eviction thread — and frees oldest-first down to `low` once usage reaches + // `high`. Leased and copy-pinned keys are skipped, exactly as master-driven + // Evict skips them; a round that can only find protected keys frees what it + // can and stops rather than spinning. + // + // Must be called before Init(). Watermarks must satisfy + // 0 < low < high <= 1; anything else is rejected with a warning and leaves + // local eviction off. + void EnableLocalEviction(double high_watermark, double low_watermark) override; + + // Test seam: run one local-eviction round synchronously. Returns the number + // of keys freed (0 when local eviction is off or usage is below the high + // watermark). + size_t RunLocalEvictionOnceForTest() { return MaybeEvictToLowWatermark(); } + + // ==================== Reaper ==================== + + // Test seam: run one reaper sweep synchronously without the thread. + void RunReaperOnceForTest() { ReaperSweep(); } + + private: + // Lifecycle is owned by Init/Shutdown, not by the caller. Both are + // idempotent. + void StartReaper(); + void StopReaper(); + + struct PendingSlot { + uint64_t slot_id = 0; + std::vector pages; + uint64_t size = 0; + std::chrono::steady_clock::time_point deadline; + // Snapshot of allocator_generation_ at Allocate(). Commit() rejects + // slots whose generation no longer matches the current. + uint64_t generation = 0; + }; + + struct OwnedSlot { + std::vector pages; + uint64_t size = 0; + // Eviction fence, renewed by every Resolve. Lives here rather than in a + // map keyed by the same string, which cost a second hash of a ~128-byte + // key per resolve. Zero means no lease. + // + // Atomic, and stored as a steady_clock nanosecond count rather than a + // time_point, because the resolve paths renew it under a SHARED lock: + // many readers may renew the same key at once. Renewal is a CAS-max + // (RenewLeaseAtomic) so a slower thread can never shorten a lease another + // thread already extended. + std::atomic read_lease_ns{0}; + // Position in lru_, valid only while in_lru is true (local eviction on). + // Held here for the same reason as the lease: the alternative is a second + // map keyed by the same string. + // + // GUARDED BY lru_mutex_, not mutex_ -- the resolve paths touch the LRU + // while holding only a shared lock on mutex_. + std::list::iterator lru_it{}; + bool in_lru = false; + + // std::atomic is neither copyable nor movable, but OwnedSlot is stored by + // value in owned_ and moved into place on commit, so the moves are spelled + // out. Relaxed is enough: the map slot is exclusively owned at this point. + OwnedSlot() = default; + OwnedSlot(OwnedSlot&& other) noexcept + : pages(std::move(other.pages)), + size(other.size), + read_lease_ns(other.read_lease_ns.load(std::memory_order_relaxed)), + lru_it(other.lru_it), + in_lru(other.in_lru) {} + OwnedSlot& operator=(OwnedSlot&& other) noexcept { + if (this == &other) return *this; + pages = std::move(other.pages); + size = other.size; + read_lease_ns.store(other.read_lease_ns.load(std::memory_order_relaxed), + std::memory_order_relaxed); + lru_it = other.lru_it; + in_lru = other.in_lru; + return *this; + } + OwnedSlot(const OwnedSlot&) = delete; + OwnedSlot& operator=(const OwnedSlot&) = delete; + }; + + // steady_clock <-> the atomic lease representation. + static int64_t SteadyNs(std::chrono::steady_clock::time_point t) { + return std::chrono::duration_cast(t.time_since_epoch()).count(); + } + // Extend `slot`'s lease to `deadline_ns` without ever shortening it. Safe + // under a shared lock; concurrent renewals settle on the latest deadline. + static void RenewLeaseAtomic(OwnedSlot& slot, int64_t deadline_ns) { + int64_t current = slot.read_lease_ns.load(std::memory_order_relaxed); + while (current < deadline_ns && + !slot.read_lease_ns.compare_exchange_weak( + current, deadline_ns, std::memory_order_relaxed, std::memory_order_relaxed)) { + } + } + + AllocateResult AllocateLocked(const std::string& key, uint64_t size); + bool CommitLocked(uint64_t slot_id, const std::string& key, uint64_t& bytes_committed); + bool AbortLocked(uint64_t slot_id); + + // Single entry point for appending to the event outbox; when the outbox + // reaches auto_flush_threshold_ it invokes auto_flush_cb_ (under the lock) + // to wake the heartbeat. Caller MUST hold `mutex_`. + void QueueEventLocked(KvEvent event); + + // Build the full ADD list for every owned key. Caller MUST hold `mutex_`. + std::vector SnapshotOwnedKeysLocked() const; + + // True iff `slot`'s read-lease deadline is still in the future. + static bool HasActiveReadLeaseLocked(const OwnedSlot& slot, + std::chrono::steady_clock::time_point now); + + // Caller MUST hold `mutex_`. True iff `key` currently has an active copy + // pin (protects its pages from eviction). + bool HasActivePinLocked(const std::string& key) const; + + // Caller MUST hold `mutex_`. Resolve `pages` to directly-readable local + // segments via buffer_bases_. Empty if no bases are configured. + std::vector> BuildCopySegmentsLocked( + const std::vector& pages, uint64_t total_size) const; + + // Caller MUST hold `mutex_`. + std::vector BuildBufferDescsLocked( + const std::vector& pages) const; + + // Build buffer_refs_ from the installed bases + descs. Called once, from + // InstallTierConfig. + void BuildBufferRefs(); + + // Install `cfg` as this backend's live buffers (used by both constructors: + // directly for the TierConfig ctor, and from Init() after self-allocation + // for the OwnershipConfig ctor). + void InstallTierConfig(TierConfig cfg); + + void ReaperLoop(); + void ReaperSweep(); + + // ---- local eviction (no-master deployments; see EnableLocalEviction) ---- + + // Caller MUST hold `lru_mutex_`. Move `key` to the MRU end, inserting it if + // it is not yet tracked. No-op when local eviction is off. + void TouchLruLocked(const std::string& key, OwnedSlot& slot); + + // Caller MUST hold `lru_mutex_`. Drop `key` from the LRU if tracked. + void RemoveFromLruLocked(OwnedSlot& slot); + + // Self-locking wrappers for single-key callers. Lock order is always + // mutex_ -> lru_mutex_, never the reverse. + void TouchLru(const std::string& key, OwnedSlot& slot); + void RemoveFromLru(OwnedSlot& slot); + + // Run one round if local eviction is on and usage has reached the high + // watermark. Takes `mutex_` itself, so callers must NOT hold it. Returns + // the number of keys freed. + size_t MaybeEvictToLowWatermark(); + + // Owned pages held back at ClearLocal() because of an active read lease. + // Released by ReaperSweep() when release_at <= now. + struct DeferredFree { + std::string key; // for debug log only. + std::vector pages; + std::chrono::steady_clock::time_point release_at; + }; + + TierType tier_; + // Reader/writer, not exclusive: Resolve, BatchResolve and Contains are the + // hot path of a 100%-hit restore and only read owned_/pending_, so they take + // a shared lock and run concurrently. Everything that mutates the index + // (Allocate, Commit, Abort, Evict, pins, ClearLocal) takes it exclusively. + // The two things the resolve paths still write -- the read lease and the LRU + // position -- were moved off this lock: the lease is an atomic on the slot, + // the LRU has lru_mutex_. + mutable std::shared_mutex mutex_; + uint64_t page_size_; + std::chrono::milliseconds pending_ttl_; + std::chrono::milliseconds read_lease_ttl_; + std::chrono::milliseconds reaper_interval_; + + // Held as the interface, not the concrete bitmap: swapping the page- + // allocation strategy is a one-line change at the two construction sites in + // Init/InstallTierConfig, and nothing else in this class learns which pool + // it got. See page_pool.h. + std::unique_ptr allocator_; + std::vector> buffer_descs_; + // Local host base pointer per buffer_index. Source of truth for page -> + // local pointer, used by AcquireDramCopyPin. + std::vector buffer_bases_; + // The same buffers as transfer endpoints, built once at InstallTierConfig so + // BufferRef() costs an index rather than an msgpack unpack. Immutable after + // Init, hence readable without mutex_. + std::vector buffer_refs_; + + std::unordered_map pending_; + std::unordered_map owned_; + std::vector pending_events_; + + // Auto-flush state (see QueueEventLocked): + size_t auto_flush_threshold_ = SIZE_MAX; // events before a flush; default = never + std::function auto_flush_cb_; // set once at startup; wakes the heartbeat + + // Active copy pins. An entry means `key`'s owned pages are protected + // from eviction until ReleaseDramCopyPin. No TTL: lifetime is bound to + // the worker via a RAII guard, not a deadline (force-freeing under an + // in-flight backend Write would be a use-after-free). + struct PinState { + uint64_t token = 0; + std::chrono::steady_clock::time_point acquired_at; // for long-running warning only + }; + std::unordered_map pins_; + uint64_t next_pin_token_ = 1; + + std::vector deferred_frees_; + + // ---- local eviction state (see EnableLocalEviction) ---- + // Front = most recently committed/resolved, back = eviction candidate. + // Empty and unmaintained unless local_evict_enabled_. + std::list lru_; + // Guards lru_ and the per-slot lru_it/in_lru fields. Separate from mutex_ + // so a resolve can record an access while holding only a shared lock; the + // critical section is a list splice, not a hash lookup plus a page-vector + // copy. Always taken after mutex_ when both are held. + mutable std::mutex lru_mutex_; + bool local_evict_enabled_ = false; + double local_evict_high_wm_ = 0.9; + double local_evict_low_wm_ = 0.7; + // Serializes rounds without holding mutex_ across one: a second committing + // thread backs off instead of over-evicting (same shape as PeerSsdManager). + std::mutex local_evict_round_mutex_; + // Latches the "nothing reclaimable" warning so a wedged pool reports once + // per episode rather than once per commit. Cleared by the next round that + // frees something. + bool local_evict_warned_ = false; + + // False stops QueueEventLocked from recording anything (no master to ship + // to). True is the cluster default. + bool event_publishing_ = true; + + // Bumped by ClearLocal(); snapshotted into each PendingSlot. Commit() + // rejects pre-clear slots via mismatch. Local-only. + uint64_t allocator_generation_ = 0; + + std::atomic next_slot_id_{1}; + + // Set by ClearLocal(), cleared by ClearFullSyncAcked(). Gates Allocate() + // so no new owned key appears between local clear and the first acked + // full-sync empty heartbeat. + std::atomic clear_full_sync_pending_{false}; + + std::thread reaper_thread_; + std::atomic reaper_running_{false}; + std::mutex reaper_cv_mutex_; + std::condition_variable reaper_cv_; + + // ---- Ownership (only set by the self-allocating constructors) ---- + // Null for the pre-configured TierConfig constructor, which is handed buffers + // it does not own. Non-null for both production constructors, so Init() has + // exactly one path regardless of medium. + std::unique_ptr source_; + std::vector source_buffer_sizes_; + + // Mirrored from source_ at construction (defaults suit the TierConfig ctor, + // which is host memory by convention). Read by BuildBufferRefs, which runs + // for both constructors and so cannot reach through source_ unconditionally. + mori::io::MemoryLocationType buffer_loc_ = mori::io::MemoryLocationType::CPU; + int buffer_device_ = -1; + + bool owns_memory_ = false; // true once Init() has self-allocated + std::vector owned_refs_; + MemoryRegistrar* registrar_ = nullptr; // non-owning; set at Init(), used at Shutdown() +}; + +// The one place a caller outside this file obtains a PageBackend. +// +// Phase 5 Rule A ("only one place may name a concrete backend") was blocked on +// PoolClient::Init needing LocalBufferViews() for the local memcpy path; Phase 6 +// deleted that call, so construction can go through a factory returning the +// interface. The class itself stays in this header because its unit tests +// drive the allocator bookkeeping directly through the TierConfig constructor — +// what matters for the acceptance test is that the PRODUCTION path names no +// concrete type, and it no longer does. +std::unique_ptr MakePageBackend(TierType tier, uint64_t page_size, + PageBackend::OwnershipConfig ownership, + std::chrono::milliseconds pending_ttl, + std::chrono::milliseconds read_lease_ttl); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/backend/page_pool.h b/src/umbp/include/umbp/distributed/peer/backend/page_pool.h new file mode 100644 index 000000000..b80d65800 --- /dev/null +++ b/src/umbp/include/umbp/distributed/peer/backend/page_pool.h @@ -0,0 +1,99 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include + +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +// The page-allocation seam for the PAGED media -- DRAM and HBM. +// +// PageBackend is medium-agnostic above PageMemorySource: the two media differ +// in how the pool is obtained and in the loc/device stamped on published refs, +// never in how pages are handed out. So one pool implementation serves both, +// and swapping it swaps it for both. +// +// WHY ONLY HERE. The SSD side already has its plug-in seam one level up -- +// TierBackend, with five implementations (SSDTier, ShardedSsdTier, +// SpdkSsdTier, SpdkProxyTier, DummySsdTier) -- and its allocation models do +// not fit this contract anyway: the SPDK OffsetAllocator hands out +// variable-size CONTIGUOUS byte extents in a single space, and the segment log +// is an append cursor whose reclamation is whole-segment GC, not per-record +// free. A single interface over all three would be their union, which is not +// a contract. This one exists because the opposite is true here: one +// implementation today, several credible ones (bitmap, binned free-list, slab, +// hierarchical bitmap), and exactly one call site to switch them at. +// +// THREAD SAFETY: implementations hold NO internal lock. The caller +// serializes -- PageBackend::mutex_ on the peer, ClientRegistry::mutex_ on +// master. Stated here rather than left to accident: a lock-free +// implementation would want to relax this, and that has to be a deliberate +// decision rather than something an implementation quietly assumes. +// +// ALL-OR-NOTHING: a failed Allocate must leave the pool exactly as it found +// it. Callers treat nullopt as "no space" and do not clean up after it. +class PagePool { + public: + virtual ~PagePool() = default; + + PagePool(const PagePool&) = delete; + PagePool& operator=(const PagePool&) = delete; + + // Hand out `num_pages` pages, or nullopt if the pool cannot. + // + // The returned pages need NOT be contiguous and need not share a buffer -- + // callers address them through PageLocation and never assume adjacency. + // Whether an implementation prefers contiguous runs (better for a single + // RDMA), packs within one buffer, or scatters across buffers is POLICY that + // belongs to the implementation; it is deliberately not in this contract, so + // a new pool can change placement without touching a caller. + virtual std::optional> Allocate(uint32_t num_pages) = 0; + + // Return pages to the pool. Must tolerate an empty list, and must ignore + // out-of-range entries rather than trap -- teardown paths can present pages + // from a pool generation that has already been reset. + virtual void Deallocate(const std::vector& pages) = 0; + + // Capacity, in bytes. Called every heartbeat, so keep it O(buffers). + virtual uint64_t TotalBytes() const = 0; + virtual uint64_t AvailableBytes() const = 0; + uint64_t UsedBytes() const { return TotalBytes() - AvailableBytes(); } + + virtual uint64_t PageSize() const = 0; + + // Geometry, for publishing transfer endpoints. Deliberately NOT a handle to + // internal per-buffer state: PageBackend needs only each buffer's extent to + // build its TransferRefs, and exposing more would tie this interface to one + // implementation's bookkeeping (a bitmap's BufferState, say) and make a + // free-list or slab implementation awkward to write. + virtual size_t NumBuffers() const = 0; + virtual uint32_t BufferPageCount(size_t buffer_index) const = 0; + + protected: + PagePool() = default; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/peer_page_allocator.h b/src/umbp/include/umbp/distributed/peer/backend/peer_page_allocator.h similarity index 68% rename from src/umbp/include/umbp/distributed/peer/peer_page_allocator.h rename to src/umbp/include/umbp/distributed/peer/backend/peer_page_allocator.h index 65b21a94c..765b9ec2d 100644 --- a/src/umbp/include/umbp/distributed/peer/peer_page_allocator.h +++ b/src/umbp/include/umbp/distributed/peer/backend/peer_page_allocator.h @@ -27,12 +27,13 @@ #include #include +#include "umbp/distributed/peer/backend/page_pool.h" #include "umbp/distributed/types.h" namespace mori::umbp { // Page-granularity bitmap allocator, one instance per (node_id, tier) pair -// (owned by Master on the control plane, or by PeerDramAllocator on the peer). +// (owned by Master on the control plane, or by PageBackend on the peer). // Internally holds N BufferState entries (one per physical buffer registered // by the Client at RegisterClient time); each BufferState carries an // independent `vector` bitmap. @@ -47,8 +48,13 @@ namespace mori::umbp { // // THREAD SAFETY: PageBitmapAllocator holds NO internal mutex. Callers MUST // serialize every method via the owning object's mutex (e.g. -// `ClientRegistry::mutex_` on master, `PeerDramAllocator::mutex_` on peer). -class PageBitmapAllocator { +// `ClientRegistry::mutex_` on master, `PageBackend::mutex_` on peer). +// +// This is the default PagePool implementation. It is a bitmap with a next-fit +// cursor: cheap to reason about and allocation-free, but its search is still +// proportional to the run it examines rather than O(1). See page_pool.h for +// why that seam exists and what a faster implementation would look like. +class PageBitmapAllocator : public PagePool { public: struct BufferState { uint32_t buffer_index = 0; @@ -56,6 +62,23 @@ class PageBitmapAllocator { uint32_t total_pages = 0; std::vector bitmap; // true = page is currently allocated uint32_t free_count = 0; + // Next-fit cursor: where the last successful search STOPPED. + // + // Both searches used to restart at index 0 on every Allocate, which is + // O(total_pages) per call regardless of how much was already allocated. + // That is invisible on a coarse pool and crippling on a fine one: a 64 GiB + // pool of 1728 B pages -- what the sglang direct linker mandates, since its + // page must not exceed the smallest per-layer object -- has 39.8 MILLION + // pages, so every allocation rescanned tens of millions of bits and got + // slower as the pool filled. Measured on an 8-rank embedded deployment: + // BatchAllocateSlots was 93.8% of remote put time at ~280 us per page. + // + // The cursor makes the common case proportional to the run actually + // examined rather than to the pool. Correctness is unaffected: both + // searches wrap and still cover every page before reporting failure, so + // they find a run exactly when the old scan would have. Only the CHOSEN + // run differs, and no caller depends on which free pages it gets. + uint32_t cursor = 0; }; // Construct an allocator from a list of per-buffer byte sizes. @@ -86,7 +109,7 @@ class PageBitmapAllocator { // All-or-nothing allocate. See class doc for strategy ordering; nullopt // leaves every bitmap bit untouched. - std::optional> Allocate(uint32_t num_pages) { + std::optional> Allocate(uint32_t num_pages) override { if (num_pages == 0) return std::nullopt; // Strategy 1: same-buffer continuous run. @@ -102,6 +125,8 @@ class PageBitmapAllocator { pages.push_back({buf.buffer_index, idx}); } buf.free_count -= num_pages; + // Resume the next search just past this run. + buf.cursor = (*run_start + num_pages) % buf.total_pages; return pages; } } @@ -118,6 +143,7 @@ class PageBitmapAllocator { pages.push_back({buf.buffer_index, idx}); } buf.free_count -= num_pages; + if (!idxs.empty()) buf.cursor = (idxs.back() + 1) % buf.total_pages; return pages; } @@ -149,7 +175,7 @@ class PageBitmapAllocator { // Idempotent free: for each entry, only flip true -> false. Out-of-range // buffer_index / page_index and already-free pages are silently skipped // (do NOT throw, do NOT underflow free_count). - void Deallocate(const std::vector& pages) { + void Deallocate(const std::vector& pages) override { for (const auto& p : pages) { if (p.buffer_index >= buffers_.size()) continue; auto& buf = buffers_[p.buffer_index]; @@ -160,7 +186,7 @@ class PageBitmapAllocator { } } - uint64_t TotalBytes() const { + uint64_t TotalBytes() const override { uint64_t sum = 0; for (const auto& b : buffers_) { sum += static_cast(b.total_pages) * page_size_; @@ -168,16 +194,17 @@ class PageBitmapAllocator { return sum; } - uint64_t AvailableBytes() const { + uint64_t AvailableBytes() const override { uint64_t free_pages = 0; for (const auto& b : buffers_) free_pages += b.free_count; return free_pages * page_size_; } - uint64_t UsedBytes() const { return TotalBytes() - AvailableBytes(); } - - uint64_t PageSize() const { return page_size_; } - size_t NumBuffers() const { return buffers_.size(); } + uint64_t PageSize() const override { return page_size_; } + size_t NumBuffers() const override { return buffers_.size(); } + uint32_t BufferPageCount(size_t buffer_index) const override { + return buffer_index < buffers_.size() ? buffers_[buffer_index].total_pages : 0; + } // Read-only access to per-buffer state (e.g. for diagnostics / Heartbeat // status reporting). Returned reference is invalidated by any subsequent @@ -187,13 +214,26 @@ class PageBitmapAllocator { private: // Find a contiguous run of `n` free pages in `buf`. Returns the starting // page_index on success, nullopt if no such run exists. O(total_pages). + // Next-fit: start where the last search stopped and wrap once. A run may + // straddle the wrap point only if it starts before it, so the second pass + // scans [0, cursor + n) -- the overlap is what keeps this exhaustive. static std::optional FindContinuousFreeRun(const BufferState& buf, uint32_t n) { if (n == 0 || n > buf.total_pages) return std::nullopt; + const uint32_t start = buf.cursor < buf.total_pages ? buf.cursor : 0; uint32_t run = 0; - for (uint32_t i = 0; i < buf.total_pages; ++i) { + for (uint32_t i = start; i < buf.total_pages; ++i) { + if (!buf.bitmap[i]) { + if (++run == n) return i + 1 - n; + } else { + run = 0; + } + } + run = 0; + const uint32_t wrap_end = std::min(static_cast(start) + n, + static_cast(buf.total_pages)); + for (uint32_t i = 0; i < wrap_end; ++i) { if (!buf.bitmap[i]) { - ++run; - if (run == n) return i + 1 - n; + if (++run == n) return i + 1 - n; } else { run = 0; } @@ -206,7 +246,14 @@ class PageBitmapAllocator { static std::vector CollectFirstNFree(const BufferState& buf, uint32_t n) { std::vector result; result.reserve(n); - for (uint32_t i = 0; i < buf.total_pages && result.size() < n; ++i) { + // Same next-fit start as the run search, and the same exhaustiveness + // argument: the wrap pass covers [0, start), so together they visit every + // page. free_count >= n is the caller's guarantee that this succeeds. + const uint32_t start = buf.cursor < buf.total_pages ? buf.cursor : 0; + for (uint32_t i = start; i < buf.total_pages && result.size() < n; ++i) { + if (!buf.bitmap[i]) result.push_back(i); + } + for (uint32_t i = 0; i < start && result.size() < n; ++i) { if (!buf.bitmap[i]) result.push_back(i); } return result; diff --git a/src/umbp/include/umbp/distributed/peer/backend/ssd_backend.h b/src/umbp/include/umbp/distributed/peer/backend/ssd_backend.h new file mode 100644 index 000000000..75c3e7fd1 --- /dev/null +++ b/src/umbp/include/umbp/distributed/peer/backend/ssd_backend.h @@ -0,0 +1,296 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/config.h" +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/peer/backend/page_backend.h" +#include "umbp/distributed/peer/ssd/peer_ssd_manager.h" + +namespace mori::umbp { + +// SSD as a MediumBackend, staged through registered host DRAM. +// +// WHY STAGED, AND NOT A FileRef. transfer_engine.h reserves a file endpoint +// for exactly this backend ("SsdBackend is what needs a FileRef, so SsdBackend +// brings it"). This implementation deliberately does NOT take that option, and +// the reason is worth stating because it is the whole design: +// +// SSD bytes are not addressable. Every endpoint UMBP moves bytes between is +// a memory range — a raw pointer locally, an RDMA MR remotely. A file offset +// is neither. Two ways to close that gap: +// +// (a) teach the transfer layer a file endpoint — a new TransferRef kind, a +// PosixFileEngine, and CHAINING in CompositeTransferEngine for the +// remote reader (file -> bounce -> wire), which that class explicitly +// does not implement today; +// (b) make the SSD's published endpoint ordinary registered host memory, +// and move bytes between it and the device inside the backend. +// +// (b) is what this class does. The cost is a host-DRAM copy on each side; +// the benefit is that SSD reaches the data plane with ZERO new transfer-layer +// concepts — no new endpoint kind, no new engine, no chaining — and a remote +// peer reading from this node's SSD sees a perfectly ordinary registered +// buffer and needs no code at all. A FileRef backend remains the right +// answer for a zero-copy / GDS path; it is a strictly larger change and this +// one does not block it. +// +// SO: THE STAGING ARENA IS THE MEDIUM, as far as everything outside this class +// is concerned. The pages this backend publishes are host DRAM. What makes it +// SSD is where the bytes live BETWEEN a Commit and the next Resolve, which is +// the only window in which a cache tier's capacity actually means anything. +// +// LIFECYCLE OF A PUT +// 1. BatchAllocate(key, size) reserves a staging page and publishes it. +// 2. The writer RDMAs (or memcpys) into that page — normal data plane, no +// SSD awareness anywhere. +// 3. BatchCommit spills the staging page to the SSD via PeerSsdManager::Write, +// frees the page, and queues an ADD event. The staging page is NOT the +// key's home; it is borrowed for the duration of the write. +// +// LIFECYCLE OF A GET +// 1. BatchResolve fills a staging page from the SSD (PeerSsdManager:: +// PrepareRead) and publishes that page under a READ LEASE. +// 2. The reader moves bytes out of it, exactly as it would from DRAM. +// 3. The reaper frees the staging page when the lease expires. +// +// This is the one place SSD's asymmetry with DRAM shows through the +// interface: a Resolve does real IO and can fail, and it consumes a scarce +// resource (staging pages) that a DRAM resolve does not. A transient +// shortfall is reported as kBusy, after rolling back every reservation made +// by that backend batch. PoolClient retries the whole response so no page +// location survives across attempts. A batch whose own working set exceeds +// the arena is kFailed instead: unchanged retries can never make it fit. +// +// A key may span several staging pages. Runs are contiguous inside the one +// registered arena because PeerSsdManager::PrepareRead takes one (ptr, +// capacity), while AllocateResult / ResolvedEntry publish every page so the +// ordinary transfer path can address the object at page granularity. +class SsdBackend : public MediumBackend { + public: + struct Config { + // Must match the DRAM/HBM page size the node registered with master: a + // batch spanning media stays self-describing only if every medium agrees. + uint64_t page_size = 2ULL * 1024 * 1024; + + // Staging pages held in host DRAM. This is the backend's aggregate + // in-flight capacity: an object of N pages consumes N entries until its + // write commits or its read lease expires. + uint32_t staging_pages = 64; + + // Back the staging arena with hugetlbfs pages. Every byte this backend + // moves crosses the arena twice (device <-> arena, arena <-> wire), so its + // TLB behaviour is on the critical path in a way an ordinary buffer's is + // not. Falls back to 4 KiB pages with a warning when hugepages cannot be + // reserved — a failure here must not stop the node coming up. + bool staging_use_hugepages = false; + uint64_t staging_hugepage_size = 2ULL * 1024 * 1024; + + PeerSsdConfig ssd; + + std::chrono::milliseconds pending_ttl{30000}; + std::chrono::milliseconds read_lease_ttl{3000}; + std::chrono::milliseconds reaper_interval{200}; + }; + + explicit SsdBackend(Config cfg); + ~SsdBackend() override; + + SsdBackend(const SsdBackend&) = delete; + SsdBackend& operator=(const SsdBackend&) = delete; + + // ==================== MediumBackend ==================== + + TierType Tier() const override { return TierType::SSD; } + const char* Name() const override { return "SsdBackend"; } + + // Allocates + registers the staging arena, discards crash leftovers, and + // starts the reaper. Fails if the SSD tier itself could not be opened. + bool Init(MemoryRegistrar* registrar) override; + void Shutdown() override; + + // From the SSD tier, NOT from the staging arena: capacity is a statement + // about what this node can hold for the cluster, and the staging arena holds + // nothing across a Commit. + TierCapacity Capacity() const override; + uint64_t OwnedKeyCount() const override; + + // Forwarded to PeerSsdManager, which already emits TierType::SSD events and + // whose header anticipated exactly this adapter. + std::vector DrainPendingEvents() override; + std::vector SnapshotOwnedKeys() const override; + std::vector SnapshotOwnedKeysForFullSync() override; + + void ClearLocal() override; + void ClearFullSyncAcked() override; + bool IsClearFullSyncPending() const override { + return clear_full_sync_pending_.load(std::memory_order_acquire); + } + + std::vector BatchAllocate(const std::vector& entries) override; + std::vector BatchCommit(const std::vector& entries) override; + std::vector BatchAbort(const std::vector& slot_ids) override; + std::vector BatchResolve(const std::vector& keys, bool include_descs, + bool allow_file_refs = false) override; + bool AcquireMigrationRead(const std::string& key, ResolvedEntry* resolved) override; + void ReleaseMigrationRead(const std::string& key) override; + bool Contains(const std::string& key) const override; + std::vector Evict(const std::vector& keys) override; + + void SetAutoFlushHook(size_t threshold, std::function cb) override; + + // MediumBackend contract; forwarded to the manager that owns the outbox. + void SetEventPublishing(bool enabled) override; + + // PeerSsdManager's read/eviction counters plus this class's own staging-arena + // counters, as monotonic values. PoolClient ships the deltas. + std::vector SampleMetrics() const override; + + uint64_t PageSize() const override { return cfg_.page_size; } + std::vector AllBufferDescs() const override; + + size_t BufferCount() const override { return buffer_ref_.Valid() ? 1 : 0; } + TransferRef BufferRef(uint32_t buffer_index) const override; + + // Test seam: run one reaper sweep synchronously without the thread. + void RunReaperOnceForTest() { ReaperSweep(); } + + private: + struct StagingSpan { + uint32_t first_page = 0; + uint32_t page_count = 0; + }; + + // A contiguous staging span borrowed for a write (pending slot) or read. + struct PendingSlot { + uint64_t slot_id = 0; + std::string key; + StagingSpan span; + uint64_t size = 0; + std::chrono::steady_clock::time_point deadline; + }; + + struct ReadLease { + StagingSpan span; + uint64_t size = 0; + std::chrono::steady_clock::time_point expires_at; + }; + + // Caller MUST hold mutex_. A zero-length result means no contiguous run of + // the requested size is currently available. + StagingSpan AcquireStagingSpanLocked(uint32_t page_count); + void ReleaseStagingSpanLocked(StagingSpan span); + + // Local address of a staging page. Valid after Init. + void* StagingPagePtr(uint32_t page_index) const; + + // hipFile handle for a segment fd, obtained lazily from the registrar's file + // engine (GdsEngine) and shared by every key on that segment. Returns nullptr + // when no file engine is configured; that result is cached so a fd is probed + // at most once. Deregistered in Shutdown. + void* GdsHandleForFd(int fd); + + // Caller MUST hold mutex_. Bumps the unshipped-event count and fires the + // auto-flush hook once it crosses the threshold. + void NoteEventQueuedLocked(); + + void StartReaper(); + void StopReaper(); + void ReaperLoop(); + void ReaperSweep(); + + Config cfg_; + + mutable std::mutex mutex_; + + std::unique_ptr ssd_; + + // The staging arena: ONE host buffer of staging_pages * page_size, so every + // page is contiguous (PrepareRead's requirement) and one registration covers + // the lot. buffer_index is therefore always 0. + std::unique_ptr staging_source_; + void* staging_base_ = nullptr; + uint64_t staging_size_ = 0; + TransferRef buffer_ref_; + std::vector buffer_desc_; + std::vector staging_page_used_; + + std::unordered_map pending_; + std::unordered_map read_leases_; + std::unordered_map migration_reads_; + + // Opt-in switch (UMBP_ENABLE_GDS=1): route SSD reads of O_DIRECT records + // through the file->GPU GdsEngine instead of the staging arena. Off by + // default even when the build has hipfile and a GdsEngine is registered; + // correctness never depends on it (hipfile's own compat mode already covers + // fastpath misses), so an unset switch keeps every read on the staging path. + // A conservative default; read once at Init. + bool gds_enabled_ = false; + + // GDS file handles by segment fd (see GdsHandleForFd). On its own lock, off + // mutex_, so a resolve holding mutex_ can register a handle without widening + // what mutex_ guards. + std::mutex gds_mutex_; + std::unordered_map gds_handles_; + + size_t unshipped_events_ = 0; + size_t auto_flush_threshold_ = SIZE_MAX; + std::function auto_flush_cb_; + + std::atomic next_slot_id_{1}; + std::atomic clear_full_sync_pending_{false}; + + // Staging-arena observability. Relaxed atomics, never correctness state. + // slot_full_rejects_ is the one to watch: it counts resolve batches that + // returned BUSY because the arena was temporarily full. + std::atomic slot_full_rejects_{0}; + std::atomic staging_expired_reclaims_{0}; + // Live arena occupancy, mirrored out of staging_page_used_ so SampleMetrics + // can read it without taking mutex_ (the metrics tick must never contend with + // a read). + std::atomic staging_pages_in_use_{0}; + bool initialized_ = false; + MemoryRegistrar* registrar_ = nullptr; + + std::thread reaper_thread_; + std::atomic reaper_running_{false}; + std::mutex reaper_cv_mutex_; + std::condition_variable reaper_cv_; +}; + +// The one place a caller outside this file obtains an SSD backend — same Phase 5 +// Rule A shape as MakePageBackend / MakeHbmBackend. +std::unique_ptr MakeSsdBackend(SsdBackend::Config cfg); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h b/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h index cbcc4ae48..63a763ecf 100644 --- a/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h +++ b/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h @@ -31,6 +31,7 @@ #include #include +#include "umbp/distributed/peer/backend/medium_backend.h" #include "umbp/distributed/types.h" #include "umbp_peer.pb.h" @@ -39,19 +40,26 @@ namespace mori::umbp { // Host-side view of one resolved key, independent of the generated proto. // Used as the encoder input on the peer side. struct ResolvedKeyEntry { + ResolveOutcome outcome = ResolveOutcome::kMissing; bool found = false; TierType tier = TierType::UNKNOWN; uint64_t size = 0; std::vector pages; + // Which backend `pages` index into. A batch may span media, so this is per + // key; `pages` carries backend-local buffer indices and is meaningless + // without it. + uint32_t backend_id = 0; }; // One key decoded out of the SoA response (page_size and descs are batch-level // and live on DecodedBatchResolve, not here). struct DecodedResolveKey { + ResolveOutcome outcome = ResolveOutcome::kMissing; bool found = false; TierType tier = TierType::UNKNOWN; uint64_t size = 0; std::vector pages; + uint32_t backend_id = 0; }; // Whole-batch decode result. @@ -88,6 +96,35 @@ inline ::umbp::TierType BatchResolveTierToProto(TierType t) { } } +inline ::umbp::ResolveStatus ResolveOutcomeToProto(ResolveOutcome outcome) { + switch (outcome) { + case ResolveOutcome::kFound: + return ::umbp::RESOLVE_STATUS_FOUND; + case ResolveOutcome::kBusy: + return ::umbp::RESOLVE_STATUS_BUSY; + case ResolveOutcome::kFailed: + return ::umbp::RESOLVE_STATUS_FAILED; + case ResolveOutcome::kMissing: + default: + return ::umbp::RESOLVE_STATUS_MISSING; + } +} + +inline ResolveOutcome ResolveOutcomeFromProto(::umbp::ResolveStatus status) { + switch (status) { + case ::umbp::RESOLVE_STATUS_FOUND: + return ResolveOutcome::kFound; + case ::umbp::RESOLVE_STATUS_BUSY: + return ResolveOutcome::kBusy; + case ::umbp::RESOLVE_STATUS_FAILED: + return ResolveOutcome::kFailed; + case ::umbp::RESOLVE_STATUS_MISSING: + case ::umbp::RESOLVE_STATUS_UNSPECIFIED: + default: + return ResolveOutcome::kMissing; + } +} + // Encode resolved keys into the SoA response. `descs` are the batch-level // deduplicated buffer descriptors (by buffer_index); pass an empty vector to // omit them (honoring BatchResolveKeysRequest.omit_descs). Clears `resp` @@ -102,6 +139,7 @@ inline void EncodeBatchResolveResponse(const std::vector& keys for (const auto& d : descs) { auto* desc = resp->add_descs(); desc->set_buffer_index(d.buffer_index); + desc->set_backend_id(d.backend_id); desc->set_desc(std::string(d.desc_bytes.begin(), d.desc_bytes.end())); } @@ -111,13 +149,18 @@ inline void EncodeBatchResolveResponse(const std::vector& keys resp->mutable_tier()->Reserve(static_cast(keys.size())); resp->mutable_size()->Reserve(static_cast(keys.size())); resp->mutable_page_count()->Reserve(static_cast(keys.size())); + resp->mutable_backend_id()->Reserve(static_cast(keys.size())); + resp->mutable_resolve_status()->Reserve(static_cast(keys.size())); resp->mutable_buffer_index()->Reserve(static_cast(total_pages)); resp->mutable_page_index()->Reserve(static_cast(total_pages)); for (const auto& k : keys) { - resp->add_found(k.found); + const ResolveOutcome outcome = k.found ? ResolveOutcome::kFound : k.outcome; + resp->add_found(outcome == ResolveOutcome::kFound); + resp->add_resolve_status(ResolveOutcomeToProto(outcome)); resp->add_tier(BatchResolveTierToProto(k.tier)); resp->add_size(k.size); + resp->add_backend_id(k.backend_id); resp->add_page_count(static_cast(k.pages.size())); for (const auto& p : k.pages) { resp->add_buffer_index(p.buffer_index); @@ -146,6 +189,7 @@ inline DecodedBatchResolve DecodeBatchResolveResponse( for (const auto& d : resp.descs()) { BufferMemoryDescBytes b; b.buffer_index = d.buffer_index(); + b.backend_id = d.backend_id(); b.desc_bytes.assign(d.desc().begin(), d.desc().end()); out.descs.push_back(std::move(b)); } @@ -155,9 +199,10 @@ inline DecodedBatchResolve DecodeBatchResolveResponse( // hold at least the summed page_count. Anything else is a malformed peer // response; refuse to decode it. if (resp.tier_size() != n || resp.size_size() != n || resp.page_count_size() != n || - resp.buffer_index_size() != resp.page_index_size()) { + resp.backend_id_size() != n || resp.buffer_index_size() != resp.page_index_size()) { return out; } + if (resp.resolve_status_size() != 0 && resp.resolve_status_size() != n) return out; size_t total_pages = 0; for (int i = 0; i < n; ++i) total_pages += resp.page_count(i); @@ -169,9 +214,13 @@ inline DecodedBatchResolve DecodeBatchResolveResponse( size_t cursor = 0; for (int i = 0; i < n; ++i) { DecodedResolveKey k; - k.found = resp.found(i); + k.outcome = resp.resolve_status_size() == n + ? ResolveOutcomeFromProto(resp.resolve_status(i)) + : (resp.found(i) ? ResolveOutcome::kFound : ResolveOutcome::kMissing); + k.found = k.outcome == ResolveOutcome::kFound; k.tier = BatchResolveTierFromProto(resp.tier(i)); k.size = resp.size(i); + k.backend_id = resp.backend_id(i); const uint32_t pc = resp.page_count(i); k.pages.reserve(pc); for (uint32_t p = 0; p < pc; ++p) { diff --git a/src/umbp/include/umbp/distributed/peer/owned_location_source.h b/src/umbp/include/umbp/distributed/peer/owned_location_source.h deleted file mode 100644 index 2903704e8..000000000 --- a/src/umbp/include/umbp/distributed/peer/owned_location_source.h +++ /dev/null @@ -1,96 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include - -#include "umbp/distributed/types.h" - -namespace mori::umbp { - -// A peer-side source of owned-location mutations (ADD/REMOVE KvEvents) that the -// heartbeat shipper drains and snapshots. Each tier-owning component on a peer -// (PeerDramAllocator for DRAM/HBM, PeerSsdManager for SSD) implements this so -// MasterClient can aggregate every source into one heartbeat bundle under a -// single monotonic sequence number — never one seq per tier (that would break -// the ack / seq-gap full-sync recovery). -// -// This interface deliberately covers ONLY events. Per-tier capacity, owned-key -// counts and the distributed-clear write gate stay on the concrete owners -// (PeerDramAllocator / PeerSsdManager) because they are tier-specific and not -// shared across all sources. -class OwnedLocationSource { - public: - virtual ~OwnedLocationSource() = default; - - // Drain the events queued since the last call (delta heartbeat). Clears the - // source's outbox. - virtual std::vector DrainPendingEvents() = 0; - - // Full snapshot of every owned key as ADD events (full sync on seq gap or - // master restart). const: a snapshot must not mutate the source's state. - virtual std::vector SnapshotOwnedKeys() const = 0; - - // Snapshot for a full sync that ALSO atomically drops the source's outbox. - // After a full sync the snapshot is the authoritative state, so any event - // still queued at snapshot time is already reflected in it (ADDs are in the - // snapshot; REMOVEs are implied by the master's full replace) and must not be - // re-shipped as a redundant post-full-sync delta. The snapshot and the outbox - // clear MUST happen in the SAME critical section — doing it in two separate - // locks would drop events committed in between. - virtual std::vector SnapshotOwnedKeysForFullSync() = 0; -}; - -// Drain every source and concat into one event list, in source order. The -// heartbeat shipper wraps the result in a SINGLE EventBundle under one -// monotonic seq — concat here, never one bundle/seq per source. Null sources -// are skipped. Exposed (vs. private to MasterClient) so it can be unit-tested -// against mock sources without standing up a master RPC. -inline std::vector DrainAllSources(const std::vector& sources) { - std::vector merged; - for (auto* src : sources) { - if (src == nullptr) continue; - auto events = src->DrainPendingEvents(); - merged.insert(merged.end(), std::make_move_iterator(events.begin()), - std::make_move_iterator(events.end())); - } - return merged; -} - -// Snapshot every source and concat into one event list, in source order (full- -// sync path). Each source ALSO atomically clears its outbox: the snapshot is -// authoritative, so the queued delta is redundant and must not be re-shipped -// afterwards. Null sources are skipped. Non-const because it mutates sources. -inline std::vector SnapshotAllSourcesForFullSync( - const std::vector& sources) { - std::vector merged; - for (auto* src : sources) { - if (src == nullptr) continue; - auto events = src->SnapshotOwnedKeysForFullSync(); - merged.insert(merged.end(), std::make_move_iterator(events.begin()), - std::make_move_iterator(events.end())); - } - return merged; -} - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/peer_dram_allocator.h b/src/umbp/include/umbp/distributed/peer/peer_dram_allocator.h deleted file mode 100644 index 46383a71d..000000000 --- a/src/umbp/include/umbp/distributed/peer/peer_dram_allocator.h +++ /dev/null @@ -1,413 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/peer/owned_location_source.h" -#include "umbp/distributed/peer/peer_page_allocator.h" -#include "umbp/distributed/types.h" - -namespace mori::umbp { - -// Peer-owned DRAM/HBM allocator + key map. In the master-as-advisor -// design, this object is the canonical owner of every per-key page set -// on the local node — master's GlobalBlockIndex is a downstream -// projection of the events this object emits. -// -// Lifecycle of a Put on this peer: -// 1. Allocate(size, tier) -> PendingSlot { slot_id, pages, ... } -// 2. (Writer RDMAs into the slot's pages.) -// 3. Commit(slot_id, key) -> moves pending -> owned, queues an ADD event. -// OR Abort(slot_id) -> drops pending; no event. -// If the writer dies between (1) and (3), the reaper drops the pending -// slot at `pending_ttl` and no event is ever emitted. -// -// Lifecycle of an Evict (master-driven): -// 1. Master sends EvictKey(keys[]) — handled by the peer service. -// 2. Peer service calls Evict(keys) here. -// 3. For each key actually freed, a REMOVE event is queued. -// 4. Heartbeat ships the REMOVE events and master drops the index entry. -class PeerDramAllocator : public OwnedLocationSource { - public: - // Per-tier inputs: each i'th buffer's size in bytes, and the matching - // packed mori::io::MemoryDesc that the writer will use to RDMA into - // that buffer. `descs.size()` must equal `buffer_sizes.size()` (or - // both be empty, meaning "this tier is not configured"). - struct TierConfig { - std::vector buffer_sizes; - std::vector> buffer_descs; // packed mori::io::MemoryDesc bytes - // Local host base pointer of each buffer (same order as buffer_sizes). - // Used by AcquireDramCopyPin to resolve a (buffer_index, page_index) - // page to a directly-readable local pointer for the SSD copy worker. - // Empty is allowed for tiers/deployments with no DRAM copy (e.g. unit - // tests that never pin); pin acquisition then yields no segments. - std::vector buffer_bases; - }; - - // `pending_ttl` is the only TTL in the system. After this elapses - // without a matching Commit / Abort, the reaper frees the slot's - // pages. `read_lease_ttl` is how long a single Resolve protects its - // key from concurrent Evict. `reaper_interval` - // is the wakeup cadence for sweeping expired pendings and read - // leases. - PeerDramAllocator(uint64_t page_size, TierConfig dram, TierConfig hbm, - std::chrono::milliseconds pending_ttl, - std::chrono::milliseconds read_lease_ttl = std::chrono::milliseconds{500}, - std::chrono::milliseconds reaper_interval = std::chrono::milliseconds{200}); - ~PeerDramAllocator(); - - PeerDramAllocator(const PeerDramAllocator&) = delete; - PeerDramAllocator& operator=(const PeerDramAllocator&) = delete; - - // -------- RPC entry points -------- - - struct PendingSlot { - uint64_t slot_id = 0; - TierType tier = TierType::UNKNOWN; - std::vector pages; - uint64_t size = 0; - std::chrono::steady_clock::time_point deadline; - // Snapshot of allocator_generation_ at Allocate(). Commit() - // rejects slots whose generation no longer matches the current. - uint64_t generation = 0; - }; - - struct OwnedSlot { - TierType tier = TierType::UNKNOWN; - std::vector pages; - uint64_t size = 0; - }; - - struct ResolveResult { - bool found = false; - TierType tier = TierType::UNKNOWN; - std::vector pages; - uint64_t size = 0; - }; - - // ResolveResult + descs built under the same lock. - struct ResolvedEntry { - bool found = false; - TierType tier = TierType::UNKNOWN; - std::vector pages; - uint64_t size = 0; - std::vector descs; - }; - - struct EvictResult { - std::string key; - uint64_t bytes_freed = 0; // 0 if key was unknown / already freed / read-leased - }; - - // Allocate() outcome. Only kSuccessAllocated populates slot. - // kFailedNoSpace is split out so the writer can keep retrying on - // other peers; every other failure collapses into kFailed and is - // diagnosed via the WARN log emitted inside Allocate(). - enum class Outcome { - kSuccessAllocated, - kSuccessAlreadyExists, // owned_[key] dedup hit - kFailed, // generic — see allocator log for reason - kFailedNoSpace, // tier exhausted; retry on another peer - }; - - struct AllocateResult { - Outcome outcome = Outcome::kFailed; - std::optional slot; // populated iff outcome == kSuccessAllocated - }; - - struct AllocateRequest { - std::string key; - uint64_t size = 0; - TierType tier = TierType::UNKNOWN; - }; - - struct BatchAllocateResult { - Outcome outcome = Outcome::kFailed; - std::optional slot; - std::vector descs; - }; - - struct CommitRequest { - uint64_t slot_id = 0; - std::string key; - }; - - struct CommitResult { - bool success = false; - uint64_t bytes_committed = 0; - }; - - // Reserve `size` bytes on `tier` for `key`. `key` enables owned_ - // dedup (master-index-lag fallback; primary dedup is at BatchRoutePut). - AllocateResult Allocate(const std::string& key, uint64_t size, TierType tier); - - std::vector BatchAllocate(const std::vector& entries); - - // Move pending -> owned and queue an ADD event. Returns false if - // slot_id is unknown (already reaped, already aborted, or never - // existed) — the writer treats false as a Put failure. - bool Commit(uint64_t slot_id, const std::string& key, uint64_t& bytes_committed); - - std::vector BatchCommit(const std::vector& entries); - - // Drop a pending slot. Idempotent: returns true if the slot was - // dropped here OR was already gone (reaped or never existed). False - // is reserved for a state we don't currently produce (kept for future - // contract changes). - bool Abort(uint64_t slot_id); - - std::vector BatchAbort(const std::vector& slot_ids); - - // Look up a key the writer was just routed to. Extends the - // read-lease deadline for `key` to now + `read_lease_ttl_`; - // concurrent Evict requests for that key during the lease window - // report bytes_freed=0. - ResolveResult Resolve(const std::string& key); - - // Batched Resolve + BufferDescsForPages under a single mutex_ hold. - // Per-key behavior byte-identical to Resolve() + BufferDescsForPages(). - // `include_descs=false` skips BuildBufferDescsLocked entirely (leaving - // descs empty) — for callers that already have the buffer descs cached - // (e.g. from GetPeerInfo) and would otherwise discard the result. - std::vector BatchResolve(const std::vector& keys, - bool include_descs = true); - - // Master-driven eviction. Idempotent: keys that are unknown or - // already gone produce zero-bytes entries; keys with active read - // leases produce zero-bytes entries (master will retry next round). - // For every key actually freed, a REMOVE event is queued. - std::vector Evict(const std::vector& keys); - - // -------- DRAM copy pin (SSD copy-on-commit) -------- - - // A copy pin protects a committed key's DRAM pages from eviction while - // the SSD copy worker reads them. It is an in-process lifetime guard - // (NOT a master lease): pages stay readable until ReleaseDramCopyPin. - // `segments` are directly-readable local pointers + lengths (pages may - // be non-contiguous across buffers). - struct DramCopyPin { - std::vector> segments; - size_t total_size = 0; - uint64_t pin_token = 0; // release-time validation - }; - - // Atomically (under mutex_) confirm `key` is owned, mark it pinned, and - // resolve its pages to local segments. Returns nullopt if the key is - // not owned (already evicted -> worker drops the task) or is already - // pinned (duplicate task). While pinned, Evict(key) reports - // bytes_freed=0 and emits no REMOVE DRAM (master retries next round). - // There is NO TTL: the pin lives until ReleaseDramCopyPin. The caller - // MUST release (use a RAII guard) so a worker exit always frees the pin. - std::optional AcquireDramCopyPin(const std::string& key); - - // Release a pin. No-op if the key is not pinned or the token does not - // match (tolerates duplicate / late release). - void ReleaseDramCopyPin(const std::string& key, uint64_t pin_token); - - // -------- Distributed Clear -------- - - // Drop owned/lease state, bump allocator_generation_ so pre-clear - // pending slots fail at Commit(), and clear the event outbox. Owned - // pages with an active read lease are deferred to the reaper; others - // are freed immediately. Allocate() returns nullopt until - // ClearFullSyncAcked(). - void ClearLocal(); - - // Called by the heartbeat thread after the first full-sync empty - // snapshot is acked by master. Re-enables Allocate(). - void ClearFullSyncAcked(); - - bool IsClearFullSyncPending() const { - return clear_full_sync_pending_.load(std::memory_order_acquire); - } - - // -------- Heartbeat helpers -------- - - // Drain the outbox of events queued since the last call. Called by - // the heartbeat shipper; clears the buffer. OwnedLocationSource. - std::vector DrainPendingEvents() override; - - // Register a hook invoked (while the allocator lock is held) once the - // unshipped-event outbox reaches `threshold`, so the heartbeat is flushed - // automatically after a batch of puts instead of relying on the application to - // call FlushHeartbeat(). `cb` MUST be cheap and MUST NOT re-enter the - // allocator — it should only signal the heartbeat thread (set a flag + - // notify). `threshold` should be >= 1; pass a very large value to disable. - void SetAutoFlushHook(size_t threshold, std::function cb); - - // Full snapshot of every owned key as ADD events. Used when master - // requests a full sync (seq gap or master restart). OwnedLocationSource. - std::vector SnapshotOwnedKeys() const override; - - // Full-sync snapshot that also atomically drops the event outbox (and disarms - // auto-flush) under the same lock. See OwnedLocationSource. - std::vector SnapshotOwnedKeysForFullSync() override; - - // Live owned-key count per tier. O(tiers) — cheap to call every - // heartbeat. Used by the heartbeat shipper for per-client metrics. - std::map OwnedKeyCountByTier() const; - - // Live capacity per tier — derived directly from the underlying - // bitmap allocators, so always reflects pending+owned correctly. - std::map TierCapacitiesSnapshot() const; - - // -------- Wire-side helpers (used by peer service handlers) -------- - - uint64_t PageSize() const { return page_size_; } - uint64_t PendingTtlMs() const { return static_cast(pending_ttl_.count()); } - - // Returns descs for all configured buffers on `tier` (sorted by - // buffer_index ascending), ready to drop into a peer-service RPC - // response. Empty if the tier is not configured. Used by - // GetPeerInfo for first-contact bootstrap; AllocateSlot / ResolveKey - // need only the descs that actually appear in `pages` and use the - // overload below. - std::vector AllBufferDescs(TierType tier) const; - - // Returns the deduplicated, ascending-by-buffer_index list of descs - // referenced by `pages`. - std::vector BufferDescsForPages( - TierType tier, const std::vector& pages) const; - - // -------- Reaper -------- - - void StartReaper(); - void StopReaper(); - - // Test seam: run one reaper sweep synchronously without the thread. - // Public so unit tests can drive deterministic TTL expiry without - // racing the background thread. - void RunReaperOnceForTest() { ReaperSweep(); } - - private: - // Caller MUST hold `mutex_`. - PageBitmapAllocator* AllocatorForLocked(TierType tier); - const PageBitmapAllocator* AllocatorForLocked(TierType tier) const; - - AllocateResult AllocateLocked(const std::string& key, uint64_t size, TierType tier); - bool CommitLocked(uint64_t slot_id, const std::string& key, uint64_t& bytes_committed); - bool AbortLocked(uint64_t slot_id); - - // Single entry point for appending to the event outbox; when the outbox - // reaches auto_flush_threshold_ it invokes auto_flush_cb_ (under the lock) to - // wake the heartbeat. Caller MUST hold `mutex_`. - void QueueEventLocked(KvEvent event); - - // Build the full ADD list for every owned key. Caller MUST hold `mutex_`. - std::vector SnapshotOwnedKeysLocked() const; - - // Caller MUST hold `mutex_`. True iff `key`'s read-lease deadline - // is still in the future. Drops the entry if it has expired. - bool HasActiveReadLeaseLocked(const std::string& key); - - // Caller MUST hold `mutex_`. True iff `key` currently has an active - // copy pin (protects its pages from eviction). - bool HasActivePinLocked(const std::string& key) const; - - // Caller MUST hold `mutex_`. Resolve `pages` to directly-readable - // local segments via the per-tier buffer bases. Empty if the tier has - // no bases configured. - std::vector> BuildCopySegmentsLocked( - TierType tier, const std::vector& pages, uint64_t total_size) const; - - // Caller MUST hold `mutex_`. - std::vector BuildBufferDescsLocked( - TierType tier, const std::vector& pages) const; - - void ReaperLoop(); - void ReaperSweep(); - - // Owned pages held back at ClearLocal() because of an active read - // lease. Released by ReaperSweep() when release_at <= now. - struct DeferredFree { - std::string key; // for debug log only. - TierType tier = TierType::UNKNOWN; - std::vector pages; - std::chrono::steady_clock::time_point release_at; - }; - - mutable std::mutex mutex_; - uint64_t page_size_; - std::chrono::milliseconds pending_ttl_; - std::chrono::milliseconds read_lease_ttl_; - std::chrono::milliseconds reaper_interval_; - - std::map> allocators_; - std::map>> tier_descs_; - // Per-tier local host base pointer per buffer_index (parallel to the - // allocator's buffers). Source of truth for page -> local pointer used - // by AcquireDramCopyPin. - std::map> tier_bases_; - - std::unordered_map pending_; - std::unordered_map owned_; - std::unordered_map read_lease_until_; - std::vector pending_events_; - - // Auto-flush state (see QueueEventLocked): - size_t auto_flush_threshold_ = SIZE_MAX; // events before a flush; default = never - std::function auto_flush_cb_; // set once at startup; wakes the heartbeat - - // Active copy pins. An entry means `key`'s owned pages are protected - // from eviction until ReleaseDramCopyPin. No TTL: lifetime is bound to - // the worker via a RAII guard, not a deadline (force-freeing under an - // in-flight backend Write would be a use-after-free). - struct PinState { - uint64_t token = 0; - std::chrono::steady_clock::time_point acquired_at; // for long-running warning only - }; - std::unordered_map pins_; - uint64_t next_pin_token_ = 1; - - std::vector deferred_frees_; - - // Bumped by ClearLocal(); snapshotted into each PendingSlot. - // Commit() rejects pre-clear slots via mismatch. Local-only. - uint64_t allocator_generation_ = 0; - - std::atomic next_slot_id_{1}; - - // Set by ClearLocal(), cleared by ClearFullSyncAcked(). Gates - // Allocate() so no new owned key appears between local clear and the - // first acked full-sync empty heartbeat. - std::atomic clear_full_sync_pending_{false}; - - std::thread reaper_thread_; - std::atomic reaper_running_{false}; - std::mutex reaper_cv_mutex_; - std::condition_variable reaper_cv_; -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/peer/peer_service.h b/src/umbp/include/umbp/distributed/peer/peer_service.h index 128f85ac5..184ab514b 100644 --- a/src/umbp/include/umbp/distributed/peer/peer_service.h +++ b/src/umbp/include/umbp/distributed/peer/peer_service.h @@ -23,85 +23,46 @@ #include -#include -#include #include #include -#include -#include #include namespace mori::umbp { -class PeerDramAllocator; -class PeerSsdManager; +class BackendRegistry; class MasterClient; -class SsdCopyPipeline; - -// Prometheus-only observability counters for the SSD read-staging slots. NOT -// correctness state (the slot state machine in peer_service.cpp is the source -// of truth) — relaxed atomics incremented at discrete events and read once per -// metrics tick by PoolClient::PublishSsdMetrics. -struct StagingMetrics { - std::atomic expired_reclaims{0}; // leased slot reclaimed past TTL - std::atomic slot_full_rejects{0}; // PrepareSsdRead -> NO_SLOT -}; +class PeerPool; class PeerServiceServer { public: - // `dram_alloc` is non-owning and may be null when the host process has - // no DRAM/HBM tier (SSD-only deployments). When null, the - // AllocateSlot/CommitSlot/AbortSlot/ResolveKey/EvictKey handlers - // respond with success=false / found=false. + // `pool` is non-owning (PoolClient outlives this server). Handlers delegate + // logical placement and backend dispatch to it; no concrete backend type is + // named here. // - // `peer_ssd` + the SSD staging region (base / size / packed MemoryDesc) drive - // the SSD read RPCs: when all are present the peer serves PrepareSsdRead out - // of `peer_ssd` into the staging buffer. When `peer_ssd` is null (SSD - // disabled) the staging args are typically null/0 and the SSD read RPCs - // report SSD_READ_ERROR (SsdRpcAvailable() == false). The staging buffer - // must be RDMA-registered by the caller; its MemoryDesc is published via - // GetPeerInfo so readers can RDMA out of it. - // - // `copy_pipeline` (non-owning, may be null when SSD is disabled) receives an - // SsdCopyTask after each successful DRAM commit so the owner peer copies the - // committed bytes to its local SSD tier asynchronously. - PeerServiceServer(PeerDramAllocator* dram_alloc, PeerSsdManager* peer_ssd = nullptr, - void* ssd_staging_base = nullptr, size_t ssd_staging_size = 0, - std::vector ssd_staging_mem_desc_bytes = {}, int num_read_slots = 16, - std::chrono::milliseconds lease_timeout = std::chrono::milliseconds{3000}, - std::vector engine_desc_bytes = {}, - MasterClient* master_client = nullptr, - SsdCopyPipeline* copy_pipeline = nullptr); + // SSD read staging (PrepareSsdRead/ReleaseSsdLease) was removed in the + // backend-agnostic refactor Phase 0 — SSD is unwired from the distributed + // data plane (see design-backend-agnostic-refactor.md). + PeerServiceServer(PeerPool& pool, std::vector engine_desc_bytes = {}, + MasterClient* master_client = nullptr); + // Source-compatible adapter for embedders that have not created a Pool yet. + // The server owns an implicit SingleBackendPolicy PeerPool over `registry`. + PeerServiceServer(BackendRegistry* registry, std::vector engine_desc_bytes = {}, + MasterClient* master_client = nullptr); ~PeerServiceServer(); bool Start(uint16_t port); void Stop(); - - const StagingMetrics& Metrics() const { return metrics_; } - - // SSD read staging slots currently in use (Preparing or Leased). Sampled - // once per metrics flush by PoolClient's metrics provider for a gauge. - size_t SnapshotReadSlotsInUse() const; - - // Read-only access for the heartbeat shipper (lives in MasterClient - // / PoolClient). Never null after construction with a non-null - // allocator argument. - PeerDramAllocator* DramAllocator() const { return dram_alloc_; } + uint16_t BoundPort() const { return bound_port_; } private: - void* ssd_staging_base_; - size_t ssd_staging_size_; - PeerSsdManager* peer_ssd_; - PeerDramAllocator* dram_alloc_; + std::unique_ptr owned_pool_; + PeerPool* pool_; MasterClient* master_client_; - SsdCopyPipeline* copy_pipeline_ = nullptr; - - StagingMetrics metrics_; - std::vector ssd_staging_mem_desc_bytes_; std::vector engine_desc_bytes_; std::unique_ptr server_; + uint16_t bound_port_ = 0; class UMBPPeerServiceImpl; std::unique_ptr service_; diff --git a/src/umbp/include/umbp/distributed/peer/peer_ssd_manager.h b/src/umbp/include/umbp/distributed/peer/ssd/peer_ssd_manager.h similarity index 52% rename from src/umbp/include/umbp/distributed/peer/peer_ssd_manager.h rename to src/umbp/include/umbp/distributed/peer/ssd/peer_ssd_manager.h index 852861dac..3810c4d25 100644 --- a/src/umbp/include/umbp/distributed/peer/peer_ssd_manager.h +++ b/src/umbp/include/umbp/distributed/peer/ssd/peer_ssd_manager.h @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -35,36 +36,41 @@ #include #include "umbp/distributed/config.h" // PeerSsdConfig -#include "umbp/distributed/peer/owned_location_source.h" #include "umbp/distributed/types.h" +// RecordLocation is returned by value from LocateRecord, so the tier interface +// is included here rather than forward-declared. peer/ssd is the SSD adapter +// that already builds on the local tier (Phase 5 lint exempts this directory), +// so this crosses no new boundary. +#include "umbp/local/tiers/tier_backend.h" namespace mori::umbp { -class TierBackend; // umbp/local/tiers/tier_backend.h — kept out of this header. - enum class SsdReadStatus { kOk, kNotFound, kSizeTooLarge, kError }; struct SsdReadOutcome { SsdReadStatus status = SsdReadStatus::kError; size_t size = 0; }; -// Peer-side owner of the local SSD tier in the master-as-advisor design. -// Single responsibility: manage one SSD TierBackend + the key->SSD-location -// map + capacity + the owned-location event outbox + the read-prepare and -// local-eviction paths. It deliberately reuses ONLY the low-level TierBackend -// (SSDTier); it must NOT pull in LocalStorageManager / LocalBlockIndex (which -// carry their own DRAM tier + demote/promote) — peer DRAM is owned by -// PeerDramAllocator and two DRAM concepts would scramble ownership. -class PeerSsdManager : public OwnedLocationSource { +// Peer-side owner of the local SSD tier: one SSD TierBackend + the +// key->SSD-location map, capacity, the owned-location event outbox, and the +// read-prepare and local-eviction paths. Reached from the distributed data +// plane through SsdBackend : MediumBackend, which forwards to it. +// +// It reuses ONLY the low-level TierBackend (SSDTier). Peer DRAM is owned by +// PageBackend, and a second thing here that also held DRAM and demoted between +// the two would scramble that ownership — which is what the deleted +// LocalStorageManager did, and part of why it went. +class PeerSsdManager { public: explicit PeerSsdManager(const PeerSsdConfig& cfg); // Test-only: inject a ready-made backend and explicit watermarks so unit // tests can drive eviction with a controllable (e.g. blocking) fake backend. // Production code must use the config constructor. - PeerSsdManager(std::unique_ptr backend, double high_watermark, double low_watermark); + PeerSsdManager(std::unique_ptr backend, double high_watermark, double low_watermark, + bool single_flight = true); - ~PeerSsdManager() override; + ~PeerSsdManager(); PeerSsdManager(const PeerSsdManager&) = delete; PeerSsdManager& operator=(const PeerSsdManager&) = delete; @@ -74,6 +80,22 @@ class PeerSsdManager : public OwnedLocationSource { bool Exists(const std::string& key) const; + // Stored size of `key`, or 0 if unknown. Needed by SsdBackend::BatchResolve: + // PrepareRead wants a staging buffer chosen BEFORE the read, so the reader has + // to know how big the key is while it still only holds the key's name. + // `owned_` has carried this all along; only the accessor was missing. + uint64_t SizeOf(const std::string& key) const; + + // Hold a key against local watermark eviction across a peer-local migration. + // Uses the same refcount as active reads so eviction has one protection rule. + bool PinForMigration(const std::string& key); + void UnpinForMigration(const std::string& key); + + // Physical location of |key|'s value for a zero-copy GDS read, or nullopt when + // the key is unknown or being evicted (mirroring PrepareRead's kNotFound). + // Forwards to the SSD tier; does no device IO. + std::optional LocateRecord(const std::string& key) const; + // Write the key's bytes (assembled from possibly non-contiguous DRAM source // segments) to the SSD backend. On success records the SSD location and // queues an ADD SSD event; on failure records nothing and queues nothing @@ -81,6 +103,21 @@ class PeerSsdManager : public OwnedLocationSource { bool Write(const std::string& key, const std::vector>& segments, size_t total_size); + // Batched form of Write for contiguous sources — the SsdBackend::BatchCommit + // path. One dedup pass, ONE backend BatchWrite (which a multi-drive + // ShardedSsdTier fans across every drive in parallel), then one recording + // pass. This is what makes an SSD-target batch put scale with drive count + // instead of serializing key-by-key; a per-key Write loop would idle every + // drive but one. + // + // Result length and order match @p keys. An already-owned key reports true + // without any device IO (LRU refreshed). On a partial failure the failed keys + // are retried once after a single eviction round, so a batch that trips the + // high watermark mid-way still lands. + std::vector WriteBatch(const std::vector& keys, + const std::vector& srcs, + const std::vector& sizes); + // Local eviction of a single key. Read priority: a key with an // in-flight PrepareRead (inflight_reads_ > 0) is NOT evicted (returns false). // Concurrency: marks the key in evicting_ under the lock, runs the backend @@ -125,13 +162,34 @@ class PeerSsdManager : public OwnedLocationSource { // in-flight (inflight_reads_) across that window so eviction skips it. SsdReadOutcome PrepareRead(const std::string& key, void* staging_ptr, size_t staging_cap); - // OwnedLocationSource — all events carry TierType::SSD. - std::vector DrainPendingEvents() override; - std::vector SnapshotOwnedKeys() const override; + // Batched form of PrepareRead: resolve + mark every key in flight under one + // lock, issue ONE backend BatchReadIntoPtr (fanned across drives by + // ShardedSsdTier), then release the marks in one pass. Per-key semantics are + // identical to PrepareRead — including kNotFound for an evicting key and + // kSizeTooLarge rejected before any device IO. + // + // This is the throughput path: the per-key form costs one device read per + // key and leaves every drive but one idle, so a multi-drive tier can only be + // saturated through here. + // + // @p dsts and @p caps are parallel to @p keys; the result is too. + std::vector PrepareReadBatch(const std::vector& keys, + const std::vector& dsts, + const std::vector& caps); + + // Same shape as OwnedLocationSource (see the class comment above) — all + // events carry TierType::SSD. + std::vector DrainPendingEvents(); + + // False stops recording ADD/REMOVE and drops what is queued. The outbox has + // exactly one consumer -- MasterClient's heartbeat -- so a node with no + // master would otherwise grow it by one entry per written key forever. + void SetEventPublishing(bool enabled); + std::vector SnapshotOwnedKeys() const; // Full-sync snapshot that also atomically drops the event outbox under the - // same lock. See OwnedLocationSource. - std::vector SnapshotOwnedKeysForFullSync() override; + // same lock. + std::vector SnapshotOwnedKeysForFullSync(); // Crash-restart leftover policy (discard): after a crash owned_ is empty but // physical SSD bytes may remain, diverging used capacity from owned_. This @@ -140,14 +198,22 @@ class PeerSsdManager : public OwnedLocationSource { // synchronized against Write/PrepareRead). No-op when SSD is disabled. void DiscardLeftoverOnStartup(); - // Prometheus-only observability snapshots (see metrics_ below); sampled once - // per metrics tick by PublishSsdMetrics(), never drive correctness. + // Prometheus-only observability snapshots (see metrics_ below); never drive + // correctness. Unread while dormant (PoolClient::PublishSsdMetrics, the + // former sampler, was removed with the rest of Phase 0's distributed-mode + // SSD wiring — see design-backend-agnostic-refactor.md). uint64_t ReadOk() const { return metrics_.read_ok.load(std::memory_order_relaxed); } uint64_t ReadNotFound() const { return metrics_.read_not_found.load(std::memory_order_relaxed); } uint64_t ReadSizeTooLarge() const { return metrics_.read_size_too_large.load(std::memory_order_relaxed); } uint64_t ReadError() const { return metrics_.read_error.load(std::memory_order_relaxed); } + // Single-flight coalescing: leads issue device IO, dups are the concurrent + // same-key requests, merged is the subset of dups actually served from a + // leader's buffer. dup - merged is the residue that still hit the drive. + uint64_t ReadLead() const { return metrics_.read_lead.load(std::memory_order_relaxed); } + uint64_t ReadDup() const { return metrics_.read_dup.load(std::memory_order_relaxed); } + uint64_t ReadMerged() const { return metrics_.read_merged.load(std::memory_order_relaxed); } // Byte counters for SSD IO bandwidth (rate() in Grafana = bytes/s). uint64_t CopyBytes() const { return metrics_.copy_bytes.load(std::memory_order_relaxed); } uint64_t ReadBytes() const { return metrics_.read_bytes.load(std::memory_order_relaxed); } @@ -195,16 +261,71 @@ class PeerSsdManager : public OwnedLocationSource { std::unordered_map owned_; std::list lru_; // front = most-recently-used, back = LRU std::vector pending_events_; + // Guarded by mutex_ alongside pending_events_. See SetEventPublishing. + bool event_publishing_ = true; + + // One follower attached to a leader's read. Delivery is a memcpy into + // {dst, cap} out of the leader's staging slot. + // + // Held by shared_ptr on both sides: the leader swaps the follower list out of + // the episode when it seals, so a follower cannot read its own outcome out of + // that list afterwards — it keeps its own handle instead. + // + // Buffer lifetime: both leader and followers are blocked inside + // PrepareRead/PrepareReadBatch for the whole copy, so neither side's staging + // slot can be recycled underneath the memcpy. + struct FollowerTarget { + void* dst = nullptr; + size_t cap = 0; + // Written by the leader before it sets episode->done, read by the follower + // after it wakes. + bool ok = false; + }; + + // One leader's read, from claim to fan-out. Followers capture a shared_ptr + // at attach time and wait on *this* object, never on mutable per-key state: + // otherwise a leader that finishes and is immediately replaced by a new + // leader on the same key would reset the flag a not-yet-woken follower is + // waiting on, and that follower would silently start tracking a read whose + // fan-out list it is not in. Holding the shared_ptr also lets the episode + // outlive its slot, so the last follower to wake still finds valid state. + // + // Note there is deliberately no episode-level `ok`: the leader's read outcome + // reaches each follower through that follower's own FollowerTarget::ok. + // `done` is the only thing followers wait on. + struct ReadEpisode { + bool sealed = false; // snapshot taken; late arrivals must not attach + bool done = false; // leader finished (read + fan-out) + std::vector> followers; + std::condition_variable cv; // waited on under mutex_ + }; + + struct InflightRead { + int refs = 0; // leader + followers + independents holding this key + // Non-null exactly while a leader is mid-read; the leader clears it on + // completion so the next arrival starts a fresh episode. + std::shared_ptr episode; + }; // Read priority + eviction coordination (all guarded by mutex_): - // inflight_reads_: key -> active PrepareRead count (entry exists only while - // > 0). Eviction skips keys with a live read. + // inflight_reads_: key -> InflightRead (entry exists only while refs > 0). + // Eviction skips keys with a live read. // evicting_: keys currently inside Evict's backend-evict window; new reads // of these miss (kNotFound) and SelectVictims skips them. - std::unordered_map inflight_reads_; + std::unordered_map> inflight_reads_; std::unordered_set evicting_; std::condition_variable reads_drained_cv_; // notified when inflight_reads_ empties + // Coalesce concurrent same-key reads (UMBPSsdConfig::single_flight_reads). + // Off => every requester reads the device, and read_dup still reports the + // headroom that turning it on would recover. + bool single_flight_ = true; + + // Cap on how long a follower waits for its leader. A wedged or crashed + // leader must surface as a read error rather than hanging the RPC; the caller + // treats it as a miss and refetches from source. + static constexpr int kFollowerWaitSeconds = 30; + // Prometheus-only observability counters: relaxed atomics bumped at discrete // events, read once per metrics tick. NOT correctness state // (owned_/lru_/inflight_reads_ are authoritative); deletable with the provider. @@ -213,6 +334,18 @@ class PeerSsdManager : public OwnedLocationSource { std::atomic read_not_found{0}; std::atomic read_size_too_large{0}; std::atomic read_error{0}; + // Single-flight accounting. A requester that finds no in-flight read of + // the key is a "lead" (it is the one that issues device IO); one that finds + // an in-flight read is a "dup" — a concurrent request for the identical key + // that could attach to it instead. MLA + TP is the case that produces them + // (all attention-TP ranks GET a byte-identical key); MHA keys carry a + // per-rank suffix and should never dup. + std::atomic read_lead{0}; + std::atomic read_dup{0}; + // Of the read_dup requests, those actually served by memcpy from a leader + // (single_flight_ on and the merge window not missed). read_dup minus this + // is the residue that still hit the drive. + std::atomic read_merged{0}; std::atomic copy_bytes{0}; // bytes written to SSD (write IO) std::atomic read_bytes{0}; // bytes read from SSD (read IO) std::atomic evict_rounds{0}; diff --git a/src/umbp/include/umbp/distributed/peer/ssd_copy_pipeline.h b/src/umbp/include/umbp/distributed/peer/ssd/ssd_copy_pipeline.h similarity index 93% rename from src/umbp/include/umbp/distributed/peer/ssd_copy_pipeline.h rename to src/umbp/include/umbp/distributed/peer/ssd/ssd_copy_pipeline.h index a6c0b6770..75040de68 100644 --- a/src/umbp/include/umbp/distributed/peer/ssd_copy_pipeline.h +++ b/src/umbp/include/umbp/distributed/peer/ssd/ssd_copy_pipeline.h @@ -35,12 +35,12 @@ namespace mori::umbp { -class PeerDramAllocator; +class PageBackend; class PeerSsdManager; // One async SSD copy request, enqueued after a DRAM commit succeeds. It // deliberately carries NO source pointer: the worker re-derives the readable -// segments from `key` via PeerDramAllocator::AcquireDramCopyPin, so an +// segments from `key` via PageBackend::AcquireDramCopyPin, so an // already-evicted key is simply dropped (the pin acquisition fails). struct SsdCopyTask { std::string key; @@ -62,7 +62,7 @@ struct SsdCopyTask { // emits ADD SSD. class SsdCopyPipeline { public: - SsdCopyPipeline(PeerDramAllocator* dram, PeerSsdManager* ssd, size_t queue_depth = 4096, + SsdCopyPipeline(PageBackend* dram, PeerSsdManager* ssd, size_t queue_depth = 4096, size_t worker_threads = 1); ~SsdCopyPipeline(); @@ -91,7 +91,9 @@ class SsdCopyPipeline { bool Enqueue(SsdCopyTask task); // Prometheus-only observability snapshots (see metrics_ below). Not - // correctness state; sampled once per metrics tick by PublishSsdMetrics(). + // correctness state. Unread while dormant (PoolClient::PublishSsdMetrics, + // the former sampler, was removed with the rest of Phase 0's + // distributed-mode SSD wiring — see design-backend-agnostic-refactor.md). uint64_t Enqueued() const { return metrics_.enqueued.load(std::memory_order_relaxed); } // CopiedOk counts completed copies: a real backend Write OR a content-addressed // dedup hit (already resident) — both report success, only the former emits ADD SSD. @@ -108,7 +110,7 @@ class SsdCopyPipeline { void WorkerLoop(); void RunTask(const SsdCopyTask& task); - PeerDramAllocator* dram_; + PageBackend* dram_; PeerSsdManager* ssd_; const size_t queue_depth_; const size_t worker_threads_; diff --git a/src/umbp/include/umbp/distributed/pool/peer_pool.h b/src/umbp/include/umbp/distributed/pool/peer_pool.h new file mode 100644 index 000000000..c8e18ff74 --- /dev/null +++ b/src/umbp/include/umbp/distributed/pool/peer_pool.h @@ -0,0 +1,235 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/pool/pool_policy.h" +#include "umbp/distributed/pool/tier_transition.h" + +namespace mori::umbp { + +struct PoolAllocateResult { + uint32_t backend_id = BackendRegistry::kMaxBackends; + AllocateResult allocation; +}; + +struct PoolSlotRef { + uint32_t backend_id = BackendRegistry::kMaxBackends; + uint64_t local_slot_id = 0; +}; + +struct PoolCommitRequest { + PoolSlotRef slot; + std::string key; +}; + +struct PoolCommitResult { + uint32_t backend_id = BackendRegistry::kMaxBackends; + CommitResult commit; +}; + +struct PoolResolvedEntry { + uint32_t backend_id = BackendRegistry::kMaxBackends; + TierType tier = TierType::UNKNOWN; + ResolvedEntry resolved; +}; + +// Why a key is being removed. A tiered pool answers the two intents +// differently, so the caller has to say which one it means instead of leaving +// it to the tier configuration. +enum class PoolEvictMode { + // Free bytes on the tier that holds the key. Where a tier configures + // on_evict offload the key is demoted, not dropped: the master budgets per + // (node, tier), so moving the bytes downstream does relieve the pressure it + // measured, and the data stays reachable. + kReclaim, + // Drop the key from every backend on the peer regardless of tier + // configuration, for when the value itself must stop existing. + kDiscard, +}; + +// The implicit peer-local default Pool. It owns logical key placement and +// delegates backend choice to PoolPolicy; BackendRegistry continues to own the +// physical storage instances. This keeps the first Pool layer behaviorally +// compatible while creating the seam for weighted and tiered policies. +class PeerPool { + public: + PeerPool(BackendRegistry* backends, std::unique_ptr policy, + TransferEngine* transfer_engine = nullptr); + ~PeerPool(); + + BackendRegistry* Backends() const { return backends_; } + + std::vector BatchAllocate(const std::vector& requests); + std::vector BatchCommit(const std::vector& requests); + std::vector BatchAbort(const std::vector& slots); + std::vector BatchResolve(const std::vector& keys, + bool include_descs, bool allow_file_refs = false); + // Existence only, in request order. Cheaper than BatchResolve, which + // heap-copies the page vector of every key it finds -- work an existence + // check discards. Takes no read lease and does not touch recency: asking + // whether a key is here is not a read of it. + std::vector BatchContains(const std::vector& keys); + std::vector Evict(const std::vector& keys, PoolEvictMode mode); + void ClearLocal(); + std::vector DrainPendingEvents(); + std::vector SnapshotOwnedKeysForFullSync(); + std::map LogicalTierCapacities() const; + std::string LogicalTierForBackend(uint32_t backend_id) const; + + std::optional PlacementBackend(const std::string& key) const; + size_t PlacementCount() const; + TierTransitionMetrics TransitionMetrics() const; + // Reads this peer served, by the logical tier that held the key. Says whether + // a tier policy is actually keeping hot data where it was supposed to. + std::map TierReadHits() const; + + private: + struct PendingPlacement { + uint32_t backend_id = BackendRegistry::kMaxBackends; + uint64_t local_slot_id = 0; + std::chrono::steady_clock::time_point expires_at = std::chrono::steady_clock::time_point::max(); + }; + + struct TransitionJob { + TierTransitionKind kind = TierTransitionKind::kOffload; + std::string key; + uint32_t source_backend_id = BackendRegistry::kMaxBackends; + LogicalTierGraph::TierIndex source_tier = 0; + uint8_t attempts = 0; + }; + + // What the copy phase needs, resolved while the pool lock is still held. + struct TransitionPlan { + bool valid = false; + bool remove_source = false; + std::vector targets; + }; + + void EnqueueTransition(TransitionJob job); + void TransitionWorkerLoop(); + void StopTransitionWorker(); + void TouchLocked(const std::string& key); + void ForgetAccessLocked(const std::string& key); + // Whether this read should promote the key out of `source_tier`, per that + // tier's trigger. Counts the read for an kOnHits tier, so it must only be + // called once per served read and only when a promotion could actually run. + bool PromoteOnReadLocked(const std::string& key, LogicalTierGraph::TierIndex source_tier); + // Backend that currently holds the key, or kMaxBackends. The placement index + // is process-local and can be stale or empty (after a restart, or once + // another path evicted the key), so the backends are the authority. + uint32_t FindOwnerLocked(const std::string& key) const; + // Queues at most `max_count` offload candidates for one tier, oldest first, + // and reports how many it queued. Bounded on purpose: a peer can hold + // millions of placements, so neither the scan nor the queue may be + // proportional to that. + size_t EnqueueWatermarkCandidatesLocked(LogicalTierGraph::TierIndex tier, size_t max_count); + void MaybeEnqueueWatermarkOffloadLocked(); + bool WatermarkDrivenLocked(const TransitionJob& job) const; + TransitionPlan PlanTransitionLocked(const TransitionJob& job); + void FinishTransitionLocked(const TransitionJob& job, const TierTransitionResult& result); + // Plans under `lock`, releases it for the byte copy, then reacquires it to + // publish the outcome. Holding it across the copy would stall every + // allocate / commit / resolve on the peer for the duration of an SSD write. + TierTransitionResult RunTransition(std::unique_lock& lock, const TransitionJob& job); + + BackendRegistry* backends_; + std::unique_ptr policy_; + TransferEngine* transfer_engine_; + std::shared_ptr tier_graph_; + TierTransitionExecutor transition_executor_; + // Only watermark-driven offload consumes recency order. On-evict policies + // name their victim explicitly, and non-tiered pools never read the LRU. + bool track_access_order_ = false; + + // Resolve dispatch deliberately drops operation_mutex_ while thread-safe + // backends do their work. ClearLocal takes this exclusively so it cannot + // invalidate an SSD staging result while a resolve is between its plan and + // publish phases. + mutable std::shared_mutex lifecycle_mutex_; + // Serializes policy decisions with backend lifecycle operations. The first + // Pool still uses one metadata lock, but slow resolve dispatch does not hold + // it; later policies can shard the protected state by key without changing + // the public interface. + mutable std::mutex operation_mutex_; + // Every use is a point lookup, insert or erase. Keeping this unordered avoids + // string comparisons in BatchResolve's publish critical section. + std::unordered_map placements_; + // A migration may install its target while an independent read lease delays + // source deletion. Eviction retries must drain only that source, never fan + // out and delete the durable target. + std::unordered_map> draining_sources_; + // Keys whose bytes are being copied right now, with the pool lock released. + // Eviction defers to the copy instead of racing it, and ClearLocal waits for + // the set to drain so a late target commit cannot resurrect a wiped key. + std::unordered_set migrating_keys_; + std::condition_variable transition_idle_cv_; + std::unordered_map pending_keys_; + std::map, std::string> pending_slots_; + std::unordered_map::iterator> last_access_; + // Reads served for a key from a tier whose trigger is kOnHits. Separate from + // last_access_, which is an LRU position and carries no count. Cleared when the + // promotion is queued and whenever the key leaves the access index, so it + // cannot outlive the key it counts for. + std::unordered_map promote_hit_counts_; + // Hottest key is at the front. Stable list iterators let a touch splice an + // existing key without allocating a tree node or copying the key. + std::list access_order_; + // Earliest next scan per tier, moved forward only when a scan found nothing + // to queue. A pressured tier with no candidates must not make every commit + // batch pay for another walk. + std::vector tier_scan_backoff_; + TierTransitionMetrics transition_metrics_; + std::map tier_read_hits_; + // Invalidates a resolve plan if a clear completed while it was dispatched. + uint64_t clear_epoch_ = 0; + + std::mutex transition_mutex_; + std::condition_variable transition_cv_; + std::deque transition_queue_; + std::unordered_set queued_transition_keys_; + // Outstanding jobs per source tier. A tier that already has work queued needs + // no rescan, which is what keeps repeated commits off the placement map. + std::vector queued_by_tier_; + bool stop_transition_worker_ = false; + std::thread transition_worker_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/pool/policy_config.h b/src/umbp/include/umbp/distributed/pool/policy_config.h new file mode 100644 index 000000000..e815f9dd2 --- /dev/null +++ b/src/umbp/include/umbp/distributed/pool/policy_config.h @@ -0,0 +1,132 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +enum class PoolOffloadTrigger { + kOnEvict, + kWatermark, +}; + +// Mirrors PoolOffloadTrigger so a tier states when to promote the same way it +// states when to offload, and a new rule is a new enumerator rather than +// another boolean beside the last one. +enum class PoolPromoteTrigger { + kNever, + kOnRead, + kOnHits, +}; + +enum class PoolTransitionMode { + kMove, + kCopy, +}; + +struct PolicyBackendSpec { + std::string name; + TierType tier = TierType::UNKNOWN; + uint64_t capacity_bytes = 0; + std::vector devices; + int numa_node = -1; + std::string path; + // SSD only. The backend's staging arena is staging_slots * page_size, and a + // read that cannot claim a slot is reported as a miss rather than as + // backpressure, so this is a hard read-concurrency ceiling and not a tuning + // hint. 0 leaves BackendInstanceConfig's default. + int staging_slots = 0; +}; + +struct LogicalTierBackendConfig { + std::string backend_name; + uint32_t weight = 0; +}; + +struct LogicalTierConfig { + std::vector members; + std::vector offload_to; + PoolOffloadTrigger trigger = PoolOffloadTrigger::kOnEvict; + double high_watermark = 0.9; + double low_watermark = 0.7; + std::string name; + bool entry = false; + PoolPromoteTrigger promote_trigger = PoolPromoteTrigger::kNever; + // Reads on this tier before a key is promoted. Only kOnHits reads it, and a + // threshold of 1 is kOnRead spelled differently, so the parser requires >= 2. + uint32_t promote_hits = 0; + PoolTransitionMode promote_mode = PoolTransitionMode::kCopy; +}; + +struct BackendPolicyConfig { + std::vector backends; + std::vector logical_tiers; + uint32_t schema_version = 1; + std::string entry_tier; +}; + +// Everything a caller can learn about a tier list by looking at the list +// alone: resolved names, which tier owns each backend name, the entry tier, +// and offload edges resolved to tier indices with duplicates dropped. +struct LogicalTierIndex { + std::vector names; + std::unordered_map by_name; + std::unordered_map by_backend; + std::vector> offload_to; + size_t entry_tier = 0; + bool entry_flagged = false; +}; + +// The single statement of what makes a tier list well formed, shared by the +// JSON front end, ApplyBackendPolicy, and LogicalTierGraph::Compile so a rule +// is written once and every entry point holds configs to the same bar. Returns +// an empty string and fills the index when the list is well formed, otherwise +// the reason it is not and an index no caller may read. +std::string IndexLogicalTiers(const std::vector& tiers, LogicalTierIndex* index); + +struct BackendPolicyLoadResult { + std::optional config; + std::string error; + + bool ok() const { return config.has_value(); } +}; + +BackendPolicyLoadResult LoadBackendPolicyJson(std::string_view json); +BackendPolicyLoadResult LoadBackendPolicyFile(std::string_view path); + +struct PoolClientConfig; + +bool ApplyBackendPolicy(const BackendPolicyConfig& policy, PoolClientConfig* config, + std::string* error, std::string storage_path_suffix = {}); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/pool/pool_policy.h b/src/umbp/include/umbp/distributed/pool/pool_policy.h new file mode 100644 index 000000000..40fa57838 --- /dev/null +++ b/src/umbp/include/umbp/distributed/pool/pool_policy.h @@ -0,0 +1,188 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/pool/tier_graph.h" + +namespace mori::umbp { + +// Policy input for one new logical block. Backend names are an explicit +// override used by topology-aware callers; tier-only requests preserve the +// legacy route contract. +struct PoolPlacementRequest { + std::string key; + uint64_t size = 0; + TierType tier = TierType::UNKNOWN; + std::string backend_name; + std::string logical_tier; +}; + +// Decision-only interface. A policy returns backend ids; PeerPool owns slot +// lifecycle and data movement. Policy implementations must not call a backend +// or TransferEngine. +class PoolPolicy { + public: + virtual ~PoolPolicy() = default; + + // Ordered candidates for one placement attempt, best first. A single-choice + // policy returns one id; others append deterministic no-space fallbacks. + virtual std::vector PutOrder(const BackendRegistry& backends, + const PoolPlacementRequest& request) const = 0; + + virtual std::vector ReadOrder(const BackendRegistry& backends) const = 0; + virtual std::shared_ptr TierGraph() const { return {}; } +}; + +// Every backend the peer holds, in registry order. Shared by the policies that +// have no read preference of their own. +inline std::vector AllBackendIds(const BackendRegistry& backends) { + std::vector order; + order.reserve(backends.Size()); + for (auto* backend : backends.All()) { + const uint32_t id = backends.BackendId(backend); + if (id < BackendRegistry::kMaxBackends) order.push_back(id); + } + return order; +} + +// Compatibility policy for the implicit default pool. An explicit backend name +// selects that instance; otherwise the first instance of the routed TierType is +// used, exactly matching the pre-Pool behavior. +class SingleBackendPolicy final : public PoolPolicy { + public: + std::vector PutOrder(const BackendRegistry& backends, + const PoolPlacementRequest& request) const override { + MediumBackend* backend = nullptr; + if (!request.backend_name.empty()) { + backend = backends.Get(request.backend_name); + if (backend != nullptr && backend->Tier() != request.tier) backend = nullptr; + } else { + backend = backends.Get(request.tier); + } + if (backend == nullptr) return {}; + const uint32_t id = backends.BackendId(backend); + return id < BackendRegistry::kMaxBackends ? std::vector{id} : std::vector{}; + } + + std::vector ReadOrder(const BackendRegistry& backends) const override { + return AllBackendIds(backends); + } +}; + +inline std::unique_ptr MakeSingleBackendPolicy() { + return std::make_unique(); +} + +struct BackendPlacementWeight { + std::string backend_name; + uint32_t weight = 1; +}; + +// Deterministic weighted placement among configured instances of the requested +// tier. Hashing the logical key makes retries select the same backend without +// shared counters and gives stable distribution independent of request order. +class WeightedPlacementPolicy final : public PoolPolicy { + public: + explicit WeightedPlacementPolicy(std::vector weights) + : weights_(std::move(weights)) {} + + std::vector PutOrder(const BackendRegistry& backends, + const PoolPlacementRequest& request) const override { + if (!request.backend_name.empty()) { + auto* backend = backends.Get(request.backend_name); + if (backend == nullptr || backend->Tier() != request.tier) return {}; + const uint32_t id = backends.BackendId(backend); + return id < BackendRegistry::kMaxBackends ? std::vector{id} + : std::vector{}; + } + + std::vector candidates; + candidates.reserve(weights_.size()); + for (const auto& configured : weights_) { + auto* backend = backends.Get(configured.backend_name); + if (backend == nullptr || backend->Tier() != request.tier) continue; + candidates.push_back({configured.backend_name, configured.weight}); + } + return WeightedBackendOrder(candidates, backends, request.key); + } + + std::vector ReadOrder(const BackendRegistry& backends) const override { + return AllBackendIds(backends); + } + + private: + std::vector weights_; +}; + +inline std::unique_ptr MakeWeightedPlacementPolicy( + std::vector weights) { + return std::make_unique(std::move(weights)); +} + +class TieredPlacementPolicy final : public PoolPolicy { + public: + explicit TieredPlacementPolicy(std::shared_ptr graph) + : graph_(std::move(graph)) {} + + std::vector PutOrder(const BackendRegistry& backends, + const PoolPlacementRequest& request) const override { + if (!request.backend_name.empty()) { + auto* backend = backends.Get(request.backend_name); + if (backend == nullptr) return {}; + const uint32_t id = backends.BackendId(backend); + return id < BackendRegistry::kMaxBackends ? std::vector{id} + : std::vector{}; + } + if (graph_ == nullptr) return {}; + return request.logical_tier.empty() + ? graph_->PutOrder(request.key) + : graph_->PutOrderFromTier(request.logical_tier, request.key); + } + + std::vector ReadOrder(const BackendRegistry&) const override { + return graph_ == nullptr ? std::vector{} : graph_->ReadOrder(); + } + std::shared_ptr TierGraph() const override { return graph_; } + + private: + std::shared_ptr graph_; +}; + +inline std::unique_ptr MakeTieredPlacementPolicy( + std::shared_ptr graph) { + return std::make_unique(std::move(graph)); +} + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/pool/tier_graph.h b/src/umbp/include/umbp/distributed/pool/tier_graph.h new file mode 100644 index 000000000..11d92815a --- /dev/null +++ b/src/umbp/include/umbp/distributed/pool/tier_graph.h @@ -0,0 +1,119 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/pool/policy_config.h" + +namespace mori::umbp { + +uint64_t StablePlacementHash(std::string_view key); + +std::vector WeightedBackendOrder(const std::vector& members, + const BackendRegistry& backends, std::string_view key); + +class LogicalTierGraph { + public: + using TierIndex = size_t; + static constexpr TierIndex kEntryTierIndex = 0; + + struct Node { + std::string name; + std::vector members; + std::vector offload_to; + PoolOffloadTrigger trigger = PoolOffloadTrigger::kOnEvict; + double high_watermark = 0.9; + double low_watermark = 0.7; + PoolPromoteTrigger promote_trigger = PoolPromoteTrigger::kNever; + uint32_t promote_hits = 0; + PoolTransitionMode promote_mode = PoolTransitionMode::kCopy; + }; + + struct CompileResult; + + static CompileResult Compile(const std::vector& tiers, + const BackendRegistry& backends); + + size_t TierCount() const { return nodes_.size(); } + TierIndex EntryTierIndex() const { return entry_tier_; } + const Node& NodeAt(TierIndex tier) const { return nodes_.at(tier); } + + std::optional TierIndexForBackendId(uint32_t backend_id) const; + std::string NameForBackend(uint32_t backend_id) const; + std::map CapacitySnapshot() const; + + double MemberUtilization(uint32_t backend_id) const; + // Watermarks describe pressure, and pressure does not average: a tier that + // mixes a 512 GiB DRAM member with a 1 TiB SSD member would report only ~33% + // aggregate use with DRAM completely full, so nothing would ever offload and + // the DRAM-weighted share of the traffic would silently spill to SSD. The + // fullest member therefore speaks for the tier. + double PeakMemberUtilization(TierIndex tier) const; + bool AtOrAbove(TierIndex tier, double watermark) const; + bool MemberAtOrAbove(uint32_t backend_id, double watermark) const; + + std::vector WeightedMemberOrder(TierIndex tier, std::string_view key) const; + std::vector PutOrder(std::string_view key) const; + std::vector PutOrderFromTier(std::string_view name, std::string_view key) const; + std::vector TransitionTargetOrder(TierIndex source, std::string_view key) const; + std::vector PromoteTargetOrder(TierIndex source, std::string_view key) const; + std::vector ReadOrder() const; + + private: + LogicalTierGraph(const BackendRegistry& backends, std::vector nodes, + std::vector> upstream, + std::unordered_map tier_by_backend_id, + TierIndex entry_tier); + + std::vector ReachableTierOrder(TierIndex source, bool upstream, + bool include_source) const; + std::vector BackendOrder(const std::vector& tiers, + std::string_view key) const; + + const BackendRegistry* backends_; + std::vector nodes_; + std::vector> upstream_; + std::unordered_map tier_by_backend_id_; + TierIndex entry_tier_ = kEntryTierIndex; +}; + +struct LogicalTierGraph::CompileResult { + std::shared_ptr graph; + std::string error; + + bool ok() const { return graph != nullptr; } +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/pool/tier_transition.h b/src/umbp/include/umbp/distributed/pool/tier_transition.h new file mode 100644 index 000000000..dd36e0de9 --- /dev/null +++ b/src/umbp/include/umbp/distributed/pool/tier_transition.h @@ -0,0 +1,68 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#pragma once + +#include +#include +#include + +#include "umbp/distributed/peer/backend/medium_backend.h" + +namespace mori::umbp { + +enum class TierTransitionKind { + kOffload, + kPromotion, +}; + +struct TierTransitionResult { + bool success = false; + uint32_t target_backend_id = BackendRegistry::kMaxBackends; + uint64_t bytes_moved = 0; + uint64_t bytes_freed = 0; + bool source_draining = false; +}; + +// Executes one peer-local copy or move. Placement policy and scheduling remain +// in PeerPool; this class owns only the transactional byte-moving sequence. +class TierTransitionExecutor { + public: + TierTransitionExecutor(BackendRegistry* backends, TransferEngine* transfer_engine) + : backends_(backends), transfer_engine_(transfer_engine) {} + + TierTransitionResult Execute(const std::string& key, uint32_t source_backend_id, + const std::vector& target_backend_ids, + bool remove_source) const; + + private: + static bool BuildTransfers(MediumBackend* source, const ResolvedEntry& resolved, + MediumBackend* target, const AllocateResult& allocation, + std::vector* items); + + BackendRegistry* backends_; + TransferEngine* transfer_engine_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/pool_allocator.h b/src/umbp/include/umbp/distributed/pool_allocator.h deleted file mode 100644 index eaf1dc05b..000000000 --- a/src/umbp/include/umbp/distributed/pool_allocator.h +++ /dev/null @@ -1,110 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include - -namespace mori::umbp { - -struct PoolAllocator { - uint64_t total_size = 0; - uint64_t used_size = 0; - - struct OffsetTracker { - uint64_t bump = 0; - std::vector> free_list; // {offset, size} - }; - std::optional offset_tracker; - - std::optional Allocate(uint64_t size) { - if (size == 0 || size > AvailableBytes()) { - return std::nullopt; - } - - if (!offset_tracker) { - used_size += size; - return uint64_t{0}; - } - - auto& tracker = *offset_tracker; - - for (auto it = tracker.free_list.begin(); it != tracker.free_list.end(); ++it) { - if (it->second >= size) { - uint64_t offset = it->first; - if (it->second == size) { - tracker.free_list.erase(it); - } else { - it->first += size; - it->second -= size; - } - used_size += size; - return offset; - } - } - - if (tracker.bump + size <= total_size) { - uint64_t offset = tracker.bump; - tracker.bump += size; - used_size += size; - return offset; - } - - return std::nullopt; - } - - void Deallocate(uint64_t offset, uint64_t size) { - if (size == 0) return; - if (size > used_size) { - used_size = 0; - return; - } - used_size -= size; - - if (!offset_tracker) return; - - auto& fl = offset_tracker->free_list; - auto pos = std::lower_bound(fl.begin(), fl.end(), std::make_pair(offset, uint64_t{0})); - pos = fl.insert(pos, {offset, size}); - - auto next = std::next(pos); - if (next != fl.end() && pos->first + pos->second == next->first) { - pos->second += next->second; - fl.erase(next); - } - - if (pos != fl.begin()) { - auto prev = std::prev(pos); - if (prev->first + prev->second == pos->first) { - prev->second += pos->second; - fl.erase(pos); - } - } - } - - uint64_t AvailableBytes() const { return total_size - used_size; } -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/pool_client.h b/src/umbp/include/umbp/distributed/pool_client.h index 09bc413e4..cb678a7d6 100644 --- a/src/umbp/include/umbp/distributed/pool_client.h +++ b/src/umbp/include/umbp/distributed/pool_client.h @@ -22,13 +22,16 @@ #pragma once #include +#include #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -39,15 +42,21 @@ #include "mori/io/engine.hpp" #include "umbp/distributed/config.h" #include "umbp/distributed/master/master_client.h" +#include "umbp/distributed/metrics/component_metrics.h" +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/transfer/transfer_engine.h" #include "umbp/distributed/types.h" #include "umbp_peer.grpc.pb.h" namespace mori::umbp { -class PeerDramAllocator; +namespace benchmark { +class WorkloadTraceRecorder; +} + class PeerServiceServer; -class PeerSsdManager; -class SsdCopyPipeline; +class PeerPool; +class HbmCopyEngine; // Short name for log output. Generic FAILED maps to "FAILED" — the // detailed reason for that case lives in the peer's allocator log. @@ -98,10 +107,20 @@ class PoolClient { const std::string& NodeId() const { return config_.master_config.node_id; } - // Pin a caller-owned region for zero-copy RDMA. Calls into the IO - // engine's RegisterMemory; the descriptor is cached and looked up by - // (ptr, size) on the Put/Get hot paths. - bool RegisterMemory(void* ptr, size_t size); + // Record a caller-owned region. The descriptor is cached and looked up by + // (ptr, size) on the Put/Get hot paths; `loc`/`device` describe the caller's + // allocation (CPU, or a GPU ordinal) so the transfer layer routes to + // HbmCopyEngine instead of assuming host memory. + // + // kPinned also hands the region to the IO engine for zero-copy RDMA. + // kLocalCopyOnly skips that, for a region that is only ever a local copy + // endpoint -- it still has to be recorded, because a range that misses this + // table is described by its own address rather than its region base and so + // becomes its own transfer plan. + bool RegisterMemory(void* ptr, size_t size, + mori::io::MemoryLocationType loc = mori::io::MemoryLocationType::CPU, + int device = -1, MemoryRegistration mode = MemoryRegistration::kPinned); + void DeregisterMemory(void* ptr); // Hot paths. Both retry up to `max_route_retries` times when the @@ -117,17 +136,84 @@ class PoolClient { std::vector BatchGet(const std::vector& keys, const std::vector& dsts, const std::vector& sizes); + // Ranged multi-buffer I/O. One stored object is backed by scattered tier + // pages while the caller supplies several disjoint object-relative ranges, + // each with its own buffer. This is what a KV connector wants: read layer k + // of every block without materializing the blocks. + // + // Both directions map a range onto pages the same way the whole-object path + // does, because a TransferItem already carries src_offset/dst_offset/size — + // it *is* a range. Whole-object I/O is the degenerate case with the range + // pinned to [0, size); nothing in the backend, the engine or the wire format + // changes to support this. + // + // Get: ranges must be disjoint and inside the stored object. Keys held by + // this node are served straight out of its medium; the rest are fetched whole + // into the registered scratch arena, installed locally, and served from + // there. + // + // Put: ranges must *tile* the object — disjoint and covering [0, object_size) + // exactly — since a partial write would commit a slot with undefined gaps. + // A key routed to this node is written range-by-range directly into its + // pages, with no assembly step; a key routed elsewhere is assembled in the + // scratch arena and sent as one object. + // One caller range against one stored object. `user` is the caller's buffer + // for this range (the whole of it — the range's own base), `object_offset` + // where the range starts inside the stored object. + struct ObjectRange { + void* user = nullptr; + size_t size = 0; + size_t object_offset = 0; + }; + + // Where a ranged call's internal helpers report the time they spent, so the + // phases can be attributed without the helpers knowing about the reporting + // machinery. Split this finely because the costs scale differently: `resolve` + // is per key, `classify` is per RANGE, and one call carries thousands of + // ranges per key batch. Every member is nullable; null is inert. + struct RangedPhaseSinks { + double* resolve = nullptr; // medium index lookup / slot allocation + double* classify = nullptr; // classifying the caller's range pointers + double* build = nullptr; // TransferItem assembly (classify excluded) + double* commit = nullptr; // slot commit / abort + size_t* items = nullptr; // TransferItems handed to the engine + // Forwarded straight to TransferEngine::Transfer so the engine's own three + // steps land in the same report as the phases around them. + const TransferEngine::StepTiming* steps = nullptr; + }; + + std::vector BatchGetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& src_offsets); + std::vector BatchPutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& dst_offsets); + // Cluster-wide existence check — issues a RouteGet and reports // whether master surfaced any replica. No RDMA, no lease bump. bool Exists(const std::string& key); std::vector BatchExists(const std::vector& keys); - void* SsdStagingPtr() const { return ssd_staging_buffer_.get(); } - size_t SsdStagingSize() const { return config_.ssd_staging_buffer_size; } - const std::vector& SsdStagingMemDescBytes() const { return ssd_staging_mem_desc_bytes_; } - MasterClient& Master(); - PeerDramAllocator* DramAllocator(); + + // Every named storage backend live on this node. Callers use them through + // MediumBackend — no concrete backend type is named outside PoolClient::Init. + BackendRegistry& Backends(); + TierTransitionMetrics TransitionMetrics() const; + std::map TierReadHits() const; + std::string LogicalTierForBackend(uint32_t backend_id) const; + + // Legacy default medium: the first configured backend's tier. Valid after + // Init and retained for callers that have not adopted named instances. + TierType Medium() const { return medium_; } + + // Whether this node is part of a master-coordinated cluster. False is the + // single-node deployment: the data plane is unchanged, but nothing is routed, + // registered, or heartbeated, and a local miss is the final answer. + bool HasMaster() const { return master_client_ != nullptr; } bool IsInitialized() const; @@ -146,6 +232,11 @@ class PoolClient { std::vector pages; uint64_t page_size = 0; std::vector descs; + // Which of the peer's backends `pages` index into. A slot lives entirely + // in one medium, so one id covers the whole plan; without it the + // backend-local buffer_index does not name a buffer (see + // BufferMemoryDescBytes). + uint32_t backend_id = 0; }; // Per-entry outcome inside the Put pipeline; projected to `bool` at @@ -153,7 +244,7 @@ class PoolClient { // `kAlreadyExists` (master- or peer-side dedup) is success-to-caller // but excluded from bandwidth metrics — no bytes on the wire. // Aggregates all SUCCESS_* / FAILED_* variants of - // PeerDramAllocator::Outcome / proto AllocateSlotOutcome into 3 buckets; + // AllocateOutcome / proto AllocateSlotOutcome into 3 buckets; // the specific failure reason is logged at the call site, not propagated. enum class PutEntryOutcome { kFailed, kSucceeded, kAlreadyExists }; @@ -161,100 +252,110 @@ class PoolClient { PoolClientConfig config_; std::atomic initialized_{false}; + // Sample every instrumented component on this node — each registered storage + // backend, plus the transfer engine — and ship the result. Registered as a + // MasterClient metrics provider, so it runs on the metrics thread and never + // on a data-plane path. + // + // There is no per-component code here and no place to add any: a component is + // a MetricSource, the labels that identify it come from Tier() plus the + // registry instance name, and MetricPublisher does the differencing. That makes a new backend + // or a new transfer engine visible in Grafana without touching this file. + void PublishComponentMetrics(); + + // Holds the delta baselines for every component published above, keyed by + // (component, metric, labels). Touched once per tick; the counters + // themselves live in the components. + MetricPublisher metric_publisher_; + std::unique_ptr master_client_; - // Peer-side DRAM/HBM allocator. Owned here because PoolClient is - // the natural lifetime anchor for the per-process IO engine + DRAM - // buffers; PeerServiceServer borrows it. - std::unique_ptr peer_alloc_; - // Peer-side SSD tier owner. Built only when config_.ssd.enabled; registered - // with MasterClient as an owned-location source + SSD capacity provider. - std::unique_ptr peer_ssd_; - // Async copy-on-commit pipeline. Built only when SSD is enabled; - // borrows peer_alloc_ (pin source) + peer_ssd_ (write target). Started in - // Init, stopped in Shutdown (after the peer service so no commit can enqueue - // into a stopped pipeline). - std::unique_ptr ssd_copy_pipeline_; - std::unique_ptr peer_service_; + // Every storage medium live on this node. Owned here because PoolClient is + // the natural lifetime anchor for the per-process IO engine + backend pools. + // MasterClient borrows it for heartbeat aggregation; default_pool_ borrows it + // for logical placement and peer/local dispatch. + BackendRegistry registry_; - std::unique_ptr io_engine_; - mori::io::MemoryDesc staging_mem_{}; - std::vector export_dram_mems_; - std::unique_ptr staging_buffer_; - std::mutex staging_mutex_; + // The implicit default logical pool. It borrows registry_, owns placement + // state and policy, and is the dispatch surface for local and peer RPC paths. + std::unique_ptr default_pool_; + std::unique_ptr workload_recorder_; - std::unique_ptr ssd_staging_buffer_; - mori::io::MemoryDesc ssd_staging_mem_{}; - std::vector ssd_staging_mem_desc_bytes_; + // Legacy default tier, cached from the first configured backend. Read by + // re-cache paths that do not yet carry a PoolPolicy decision. + TierType medium_ = TierType::DRAM; + + std::unique_ptr peer_service_; - // Lazy peer connections (one per remote node). Engine descs cached - // here; DRAM memory descs hydrated on first AllocateSlot/ResolveKey - // response that references their buffer_index. + // The one byte-moving path (design doc §4). PoolClient owns the engine; + // backends receive it narrowed to MemoryRegistrar at Init, so they can + // publish endpoints but cannot transfer. + // + // `peers_` (below) is the control-plane half of talking to another node — + // the gRPC stub. `peer_directory_` is the TRANSPORT half, and it is a + // non-owning pointer to the sub-engine that implements PeerDirectory, not to + // a concrete engine type: adding a second remote transport means implementing + // that interface, not editing this file. Null when this node has no remote + // transport configured, in which case only local transfers are servicable. + std::unique_ptr transfer_engine_; + PeerDirectory* peer_directory_ = nullptr; + // Observer into transfer_engine_'s composite, used only to declare which + // host regions the gather kernel may dereference. Not a second data path. + HbmCopyEngine* hbm_engine_ = nullptr; + + // Lazy peer connections (one per remote node). Purely control plane since + // Phase 6: the peer's engine desc and buffer descriptors moved into + // MoriIoEngine, which is the layer that uses them. struct PeerConnection { + std::string node_id; std::string peer_address; - mori::io::EngineDesc engine_desc; - std::vector dram_memories; // indexed by buffer_index - bool engine_registered = false; std::unique_ptr peer_stub{nullptr, +[](void*) {}}; - std::mutex ssd_op_mutex; - mori::io::MemoryDesc ssd_staging_mem{}; - size_t ssd_staging_size = 0; + // Guards first-contact connection setup (stub creation + engine/buffer-desc + // hydration via GetPeerInfo) in EnsurePeerServiceConnection. + std::mutex conn_mutex; }; std::mutex peers_mutex_; std::unordered_map> peers_; // Caller MUST NOT hold peers_mutex_; this helper acquires it. PeerConnection& GetOrConnectPeer(const std::string& node_id, const std::string& peer_address); - // Hydrate peer.dram_memories[bd.buffer_index] for every entry in - // `descs`. Idempotent. Acquires peers_mutex_ internally. - void EnsureBufferDescsCached(PeerConnection& peer, - const std::vector& descs); - // Same as above, but the caller MUST already hold peers_mutex_. Lets the - // Build* helpers hydrate AND snapshot remote descs inside a single lock - // window so a concurrent hydrate (which may resize dram_memories) cannot - // race the reads. - void EnsureBufferDescsCachedLocked(PeerConnection& peer, - const std::vector& descs); - - // Same-tier RDMA scatter helpers (keep as much of the prior impl as - // possible — the IO engine call shape is unchanged). - bool RemoteDramScatterWrite(PeerConnection& peer, const std::vector& pages, - uint64_t page_size, const void* src, size_t size, bool zero_copy); - bool RemoteDramScatterRead(PeerConnection& peer, const std::vector& pages, - uint64_t page_size, void* dst, size_t size, bool zero_copy); - - // Self-target fast paths (no RDMA, no peer RPC). - bool LocalPutPages(const std::vector& pages, uint64_t page_size, const void* src, - size_t size); - bool LocalGetPages(const std::vector& pages, uint64_t page_size, void* dst, - size_t size); bool EnsurePeerServiceConnection(PeerConnection& peer); - // Outcome of one remote SSD get attempt. kRetry is a retryable transient - // condition (NO_SLOT or a reader-local lease expiry); kMiss (NOT_FOUND) is a - // definitive miss; kError is the not-served, not-retried catch-all (rpc - // failure incl. DEADLINE_EXCEEDED, size mismatch, RDMA failure) — not strictly - // a hard error. Keeping kMiss distinct lets BatchGet avoid surfacing a - // non-served key as a cache miss. - enum class SsdGetOutcome { kSuccess, kMiss, kRetry, kError }; - - // Remote SSD get path (reader != owner): key-based PrepareSsdRead -> RDMA out - // of the peer's published SSD staging buffer -> best-effort ReleaseSsdLease. - SsdGetOutcome RemoteSsdReadOnce(PeerConnection& peer, const std::string& key, void* dst, - size_t size); - void ReleaseSsdLeaseBestEffort(::umbp::UMBPPeer::Stub* stub, uint64_t lease_id); - - // Zero-copy registered memory regions. + // Endpoint for a caller-supplied buffer, plus the offset of `ptr` within it: + // the registered region when RegisterMemory pinned it (zero copy, ref covers + // the whole region), otherwise a plain host-bytes ref at offset 0 that the + // engine will stage through its own bounce buffer. + // + // One ref instead of the old optional + use_staging + staging + // offset triple, because whether a transfer needs staging is the transfer + // layer's decision, not the client's. + std::pair UserBufferRef(void* ptr, size_t size) const; + + // Zero-copy registered memory regions, kept sorted by `base`. + // + // Sorted because a ranged call looks one up PER RANGE, and a layer-wise + // reader carries thousands of ranges against a caller that registered one + // buffer per layer -- a linear scan makes that quadratic. Shared, because + // those lookups are reads and every rank on the node shares this client. struct RegisteredRegion { void* base; size_t size; - mori::io::MemoryDesc mem_desc; + TransferRef ref; }; - std::mutex registered_mem_mutex_; + mutable std::shared_mutex registered_mem_mutex_; std::vector registered_regions_; - std::optional> FindRegisteredMemory(const void* ptr, - size_t size); + + // The region covering [ptr, ptr+size), or null. Caller holds + // registered_mem_mutex_; the returned pointer is valid only under that lock. + // Exposed separately from FindRegisteredMemory so a batch can take the lock + // once and then resolve every range under it. + const RegisteredRegion* FindRegisteredRegionLocked(const void* ptr, size_t size) const; + + // The registered ref covering [ptr, ptr+size) plus the offset of ptr within + // it, or nullopt when the region was never registered. + std::optional> FindRegisteredMemory(const void* ptr, + size_t size) const; // Single-attempt outcome from a peer call; mapped to PutEntryOutcome // by the caller (Partition / Allocate). @@ -262,8 +363,151 @@ class PoolClient { enum class GetAttemptOutcome { kSuccess, kRetry, kFatal }; PutAttemptOutcome ExecuteLocalPut(const std::string& key, const void* src, size_t size, - TierType tier, bool enqueue_ssd_copy = true); + TierType tier, const std::string& logical_tier = {}); GetAttemptOutcome ExecuteLocalGet(const std::string& key, void* dst, size_t size); + + // Resolve `keys` through the PeerPool. On return holders[i] is the + // backend that owns keys[i] (nullptr when nothing here does) and, when it is + // non-null, resolutions[i] is that backend's entry. Both are sized to + // keys.size(); `candidates` selects which indices are asked about, so a + // caller can skip entries it has already rejected. + // + // `resolutions` may be null for a caller that only needs to know WHETHER a + // key is here -- an Exists does, and materializing an entry per key costs it + // a pages vector and a descs vector it never reads (measured 6-15%). + // + // One pool batch preserves placement/read-order, tier touches, and promotion + // semantics while avoiding one resolve call per key. + // `dst_is_device` answers, for an ORIGINAL key index, whether that key's + // destination is device memory -- the only place a file ref can be read. + // Empty (the default) means the caller cannot consume a file ref at all, so + // none is requested and every key resolves onto staged pages. + void ResolveLocalBatch(const std::vector& keys, + const std::vector& candidates, + std::vector* holders, + std::vector* resolutions, + const std::function& dst_is_device = {}); + + // One TransferItem per page between a caller buffer and `backend`'s own + // buffers. `to_backend` is Put (user -> pages), false is Get (pages -> user). + // False when the backend publishes no in-process endpoint for a referenced + // buffer — that medium cannot serve the access here, and the caller routes + // elsewhere rather than reporting a miss. + bool BuildLocalPageTransfers(MediumBackend* backend, const std::vector& pages, + uint64_t page_size, void* user, size_t size, bool to_backend, + std::vector* items); + + // ---- ranged I/O internals ---- + + // The route-first arm of BatchGetRanges: one BatchRouteGet over every key, + // before anything is served. `route_sink` is the optional phase timer. + // False means the routing RPC failed. + bool RouteAllRangesUpFront(const std::vector& keys, double* route_sink, + std::vector>* preroutes); + + // Routes for the keys BatchGetRanges' local phase missed, parallel to + // `missed`. Issues the RPC under local_first; under the route-first arm it + // projects phase 0's `preroutes` instead, so that arm costs one routing RPC + // rather than two. False means the routing RPC failed. + bool RouteMissedRanges(const std::vector& route_keys, + const std::vector& missed, + const std::vector>& preroutes, + double* route_sink, std::vector>* routes); + + // The ranged counterpart of BuildLocalPageTransfers, and the only place the + // object-range -> page-range mapping lives. Emits one TransferItem per + // (range x page) intersection; the engines coalesce adjacent ones back into + // single copies while planning, so a range spanning N contiguous pages costs + // one copy, not N. Every item carries `tag` = the caller's key index, so a + // partial failure comes back per key rather than failing the batch. + // + // Appends to `items`; returns false if the medium publishes no endpoint for a + // referenced buffer, or a range falls outside the stored object. + // + // `classify_sink`, when non-null, accumulates the seconds spent classifying + // the caller's range pointers. That cost scales with ranges rather than + // keys, so it is reported as its own phase; null makes it inert. + bool BuildLocalRangeTransfers(MediumBackend* backend, const std::vector& pages, + uint64_t page_size, uint64_t stored_size, + const std::vector& ranges, bool to_backend, size_t tag, + std::vector* items, double* classify_sink = nullptr); + + // The file-ref counterpart of BuildLocalRangeTransfers, for a medium that + // published a range of its own storage instead of staged pages. + bool BuildLocalFileRangeTransfers(const TransferRef& file_ref, + const std::vector& ranges, size_t tag, + std::vector* items); + + // Copy between one contiguous host object and a set of caller ranges, through + // the transfer engine so a device-resident caller buffer is handled by + // HbmCopyEngine rather than memcpy'd. Used for the scratch arena, which is + // not a medium and so has no PageLocations to map. + bool CopyContiguousToRanges(const void* src, size_t object_size, + const std::vector& ranges); + // Same, but for a source the backend already described. A medium slot is + // NOT necessarily host memory -- an HBM pool's BufferRef carries a device + // pointer -- so it has to keep the ref's loc/device rather than be re-wrapped + // as host bytes, or the engine picks the wrong copy direction. + bool CopyContiguousToRanges(const TransferRef& src, uint64_t src_base, size_t object_size, + const std::vector& ranges); + bool CopyRangesToContiguous(const std::vector& ranges, void* dst, + size_t object_size); + + // Item-building halves of the two copies above. They append to `items` + // instead of transferring, so a caller with MANY objects can issue ONE + // Transfer for all of them. + // + // This is the difference between the local and remote ranged paths, and it + // was worth ~65% of remote transfer time: ExecuteLocalPutRangesBatch + // accumulates every request's items and transfers once, while the remote + // arena loop called CopyRangesToContiguous per key -- one blocking transfer, + // one kernel launch and one hipStreamSynchronize EACH. Measured on an + // 8-rank embedded deployment: 369 transfers carrying 369 segments, 14.5 us + // apiece, of which only ~5 us was the copy. + // + // `tag` is the caller's key index; a failure comes back through + // Transfer's failed_tags so one bad object does not fail the batch. + bool BuildRangesToContiguousItems(const std::vector& ranges, void* dst, + size_t object_size, size_t tag, + std::vector* items); + bool BuildContiguousToRangesItems(const TransferRef& src, uint64_t src_base, size_t object_size, + const std::vector& ranges, size_t tag, + std::vector* items); + + // One locally-routed ranged put. `result_index` is the caller's key index, + // which is what gets written back into the results vector. + struct LocalRangeWriteRequest { + size_t result_index = 0; + const std::string* key = nullptr; + size_t object_size = 0; + std::vector ranges; + TierType tier = TierType::UNKNOWN; + std::string logical_tier; + }; + + // Allocate a slot per request, write the ranges straight into its pages, and + // commit. No assembly buffer: the ranges are written where they belong. + // + // Batched across keys deliberately. Every request's items go into ONE + // transfer, so a batch of GPU-sourced puts becomes a single gather kernel + // instead of one per key — the difference the kernel exists to make. Slots + // are held allocated across the transfer and committed or aborted per key + // afterwards, keyed off the engine's failed tags. + // + // Ranges must already have been validated as tiling their object. + // + // `committed_bytes`, when non-null, accumulates the bytes this call actually + // wrote. Deliberately not derivable from `results`: a key already present in + // the medium reports success without moving anything, and crediting it would + // inflate the bandwidth histogram. + // + // `sinks`, when non-null, attributes the call's time to phases; see + // RangedPhaseSinks. Its members are independently nullable too, and a null + // one costs nothing -- that is what keeps this off the hot path when the + // ranged debug switch is off. + void ExecuteLocalPutRangesBatch(const std::vector& requests, + std::vector* results, double* committed_bytes = nullptr, + const RangedPhaseSinks* sinks = nullptr); // After a successful remote DRAM fetch, if cache_remote_fetches is enabled and // the admission gate admits the block, enqueue it for asynchronous install into // this node's local DRAM tier (see ReCacheWorkerLoop). The install (DRAM @@ -272,6 +516,46 @@ class PoolClient { // (kSuccessAlreadyExists) and best-effort: the queue is bounded and drops on // full; failure does not affect the returned Get result. void MaybeReCacheAfterRemote(const std::string& key, const void* src, size_t size); + + // After a remote RANGED read served `fetched_bytes` of `object_size`, queue a + // background pull of the WHOLE object into this node's medium, so the next + // read of the key is a local hit instead of another trip to the peer. + // + // Why this exists at all: a ranged read holds only the caller's spans, so + // MaybeReCacheAfterRemote — which installs the buffer it is handed — has + // nothing legal to install. Locality has to be rebuilt from the peer. + // + // Skipped when this call already covered the object (`fetched_bytes == + // object_size`): the whole-object reader would otherwise pull every object + // twice, once for itself and once for a cache it just filled. + // + // Best-effort in every direction: gated by ranged_locality_prefetch and the + // re-cache admission policy, deduplicated against both the queue and this + // node's medium, bounded, and silent on failure. It never runs on, and never + // blocks, the read that scheduled it. + void MaybePrefetchWholeObject(const std::string& key, size_t object_size, + const RouteGetResult& route); + // The other half of the same decision: the arena slice already holds the + // whole object (the caller's ranges tiled it in order), so locality costs one + // local copy instead of a second trip to the peer. Synchronous because the + // slice is live only until the next sub-batch reuses it. Mutually exclusive + // with MaybePrefetchWholeObject — exactly one runs per fetched key. + void MaybeInstallCompleteArenaObject(const std::string& key, const void* arena_slice, + size_t object_size); + // The two locality paths are gated separately, because they do not cost the + // same thing. + // + // Landing the object in this node's medium when it is ALREADY IN HAND -- the + // arena slice is the whole object, or the peer can write straight into a slot + // -- costs nothing on the wire. The pre-ranged code did it on every remote + // ranged read, unconditionally, so gating it would be a silent regression + // against a switch whose documented job is something else. The only question + // worth asking is whether there is a medium to install into. + bool CanInstallLocally() const; + // Fetching the object a SECOND time in the background costs up to one extra + // copy of it on the wire. That is what ranged_locality_prefetch switches off, + // and what the re-cache admission policy sizes. + bool LocalityPrefetchAdmits(size_t object_size) const; // Background worker that drains recache_queue_ and performs the DRAM install // via ExecuteLocalPut. Started in Init (when cache_remote_fetches), stopped in // Shutdown before peer_alloc_ is torn down. @@ -280,20 +564,57 @@ class PoolClient { // Async re-cache install queue. MaybeReCacheAfterRemote copies the block bytes // into a job here (the source buffer is not owned past the Get return), and the // worker thread installs it. Bounded to keep memory + publish churn in check. + // Two shapes share one queue and one worker so they also share its lifecycle, + // its bound and its drain-on-shutdown: + // bytes != nullptr -> install this buffer (MaybeReCacheAfterRemote) + // bytes == nullptr -> pull the object from `route` first + // (MaybePrefetchWholeObject) struct ReCacheJob { std::string key; std::unique_ptr bytes; size_t size = 0; + std::optional route; }; + // Pull whole objects from their peers straight into freshly allocated local + // slots and commit them. Peer pages and medium pages are both registered + // host memory, so this is RDMA with no bounce buffer and no memcpy — unlike + // the ReCacheJob path, which has to carry a heap copy of the bytes. + // + // BATCHED deliberately, and it has to be. Per key the fixed cost is a peer + // resolve RPC plus an RDMA submit/wait; done one at a time on this single + // worker thread that cost is paid once per key, and for small objects it + // dominates the bytes so thoroughly that the worker cannot keep up with the + // reader — the prefetch then loses more races than it wins and the whole + // feature turns negative. Batching amortises the resolve over every key + // routed to the same peer and gives the engine one large scatter-gather to + // post. + // + // Best-effort per key: already-local keys are dropped before allocating (the + // alternative wastes a whole object of wire), and any key whose allocation or + // transfer fails is aborted without affecting the others. + void FetchWholeObjectsIntoMedium(std::vector& jobs); + std::deque recache_queue_; std::mutex recache_mutex_; + // Keys with a whole-object pull already queued or running. A layer-wise + // reader names the same key once per layer group, so without this every group + // that misses locally would queue another pull of bytes already in flight. + std::unordered_set prefetch_inflight_; + + // Serializes users of each caller-owned ranged scratch arena. Only the remote + // half of a ranged operation takes one — keys served by this node's own medium + // never touch an arena and stay fully concurrent. + // + // Separate GET and PUT arenas each get their own mutex, so a remote ranged GET + // and a remote ranged PUT run concurrently instead of serializing on one lock + // — the load/offload overlap sglang's direct linker wants. (Two same-kind ops + // still serialize on their arena's mutex.) + std::mutex ranged_get_scratch_mutex_; + std::mutex ranged_put_scratch_mutex_; std::condition_variable recache_cv_; std::thread recache_worker_; bool recache_stop_ = false; size_t recache_queue_max_ = 1024; - // Self-target SSD get: read straight from the local SSD tier into the user - // buffer (no staging / RDMA / lease). - GetAttemptOutcome ExecuteLocalSsdGet(const std::string& key, void* dst, size_t size); struct BatchPutItem { size_t index; const std::string* key; @@ -301,30 +622,57 @@ class PoolClient { size_t size; RoutePutResult route; }; + // One byte range of a stored object, named by the caller. + struct ByteSpan { + size_t object_offset = 0; + size_t size = 0; + }; + struct BatchGetItem { size_t index; const std::string* key; void* dst; - size_t size; + size_t size; // the OBJECT's stored size — what the peer must report + // Ranged reads move only `spans` and land them packed at `dst`, so the + // bytes written are the span total rather than the object size. Both + // fields default to the whole-object behaviour, which is what plain + // BatchGet wants and why it needs no changes. + size_t dst_bytes = 0; // 0 => size + const std::vector* spans = nullptr; // null/empty => whole object + // Pre-resolved destination, used instead of `dst` when set. UserBufferRef + // only knows regions registered through PoolClient::RegisterMemory, so a + // destination inside a medium backend's own pool — which is RDMA-reachable + // but was never handed to RegisterMemory — has to be named by the ref the + // backend publishes. Locality prefetch writes straight into a medium slot + // and is the only user. + const TransferRef* dst_ref = nullptr; + uint64_t dst_ref_offset = 0; RouteGetResult route; + + size_t DstBytes() const { return dst_bytes != 0 ? dst_bytes : size; } + bool Ranged() const { return spans != nullptr && !spans->empty(); } }; - // Routing plan for one BatchGet: which keys go to which tier/target. Pure - // grouping — no IO is issued here. Remote DRAM/HBM reads go through the - // batched RDMA path (remote_dram_groups); remote SSD reads through a per-key - // staging+lease path (remote_ssd_groups); self-target DRAM and SSD reads are - // deferred (collected as indices) so ExecuteBatchGetPlan can place them inside - // the remote-DRAM in-flight window when overlapping. + // Routing plan for one BatchGet: which keys go to which target. Pure + // grouping — no IO is issued here. Remote reads go through the batched RDMA + // path (remote_groups, keyed by peer node_id); self-target reads are + // deferred (collected as indices) so ExecuteBatchGetPlan can place them + // inside the remote-DRAM in-flight window when overlapping. + // What a key that no PEER can serve should fall back to. Route-first has not + // looked locally yet, so an unroutable or self-routed key still deserves a + // local read. Local-first has already asked every medium on this node, so the + // same key must be dropped -- sending it back would re-resolve a key just + // proven absent, which is the per-key cost ServeLocalBatchGet exists to avoid. + enum class LocalFallback { kAllow, kSkip }; + struct BatchGetPlan { - std::unordered_map> remote_dram_groups; - std::unordered_map> remote_ssd_groups; - std::vector local_dram_indices; - std::vector local_ssd_indices; + std::unordered_map> remote_groups; + std::vector local_indices; }; // Pure grouping for one BatchPut (no IO, no local put executed; mirrors - // BatchGetPlan). Put has no remote-SSD target. local_items hold full items - // (not bare indices) so the deferred local memcpy keeps its route tier. + // BatchGetPlan). local_items hold full items (not bare indices) so the + // deferred local memcpy keeps its route tier. struct BatchPutPlan { std::unordered_map> remote_groups; std::vector local_items; @@ -349,53 +697,132 @@ class PoolClient { BatchGetPlan PartitionBatchGetTargets(const std::vector& keys, const std::vector& dsts, const std::vector& sizes, - const std::vector>& routes); - // Execute a BatchGetPlan: local DRAM/SSD, remote SSD, and the remote-DRAM - // submit/wait arrangement. Zero-copy remote DRAM submits all peers, runs local - // DRAM/SSD + remote SSD INSIDE that submit..wait gap (so they overlap the DRAM - // wire), then waits all peers. Staging (non-zero-copy) runs per peer serially - // (submit -> wait). Reads the plan; writes per-key outcomes into *results. + const std::vector>& routes, + LocalFallback fallback); + + // Route `indices` of `keys` and scatter the answers into *routes at their + // ORIGINAL key index, leaving every other slot nullopt -- which is how + // ComputeBatchBandwidthBytes already reads a missing route, and what lets one + // partition pass serve both arms. `exclude_self` is set by the caller that + // has already tried this node's own media. False means the RPC failed. + bool RouteGetsInto(const std::vector& keys, const std::vector& indices, + bool exclude_self, std::vector>* routes); + + // ---- local-first get (PoolClientConfig::local_first) ---- + + // Serve `indices` from this node's own media in one resolve and one transfer. + // Hits go into *results; unservable keys are appended to *missed, which may be + // null when the caller has no fallback left. Shared by both BatchGet arms. + void ServeLocalGets(const std::vector& keys, const std::vector& dsts, + const std::vector& sizes, const std::vector& indices, + std::vector* results, std::vector* missed); + + // Whole-batch wrapper for the local-first arm: drops what cannot be served at + // all, serves the rest, and returns the keys the master still has to route. + std::vector ServeLocalBatchGet(const std::vector& keys, + const std::vector& dsts, + const std::vector& sizes, + std::vector* results); + + // Counter sink that tolerates a node running without a master: metrics + // accumulate on the MasterClient and ride its flush tick, so with no master + // there is simply nowhere to put them. + void CountMetric(std::string name, std::string help, MasterClient::Labels labels, double delta); + + // Placement when there is no master to ask: this node is the only candidate, + // so every key is routed to it on the medium it serves. The put paths then + // proceed exactly as they would for a master-issued self-route. + void RouteAllPutsLocally(size_t count, std::vector>* routes) const; + + // Local/remote bandwidth histograms for one BatchGet. A method rather than + // an inline tail because local-first gives the call several exit points. + void ObserveBatchGetBandwidth(const std::vector& results, const std::vector& sizes, + const std::vector>& routes, + std::chrono::steady_clock::time_point call_start); + // Execute a BatchGetPlan: local reads and the remote-DRAM submit/wait + // arrangement. Zero-copy remote DRAM submits all peers, runs local reads + // INSIDE that submit..wait gap (so they overlap the DRAM wire), then waits + // all peers. Staging (non-zero-copy) runs per peer serially (submit -> + // wait). Reads the plan; writes per-key outcomes into *results. + // `recache_remote=false` suppresses the asynchronous re-cache of remotely + // fetched blocks. void ExecuteBatchGetPlan(const BatchGetPlan& plan, const std::vector& keys, const std::vector& dsts, const std::vector& sizes, - std::vector* results); - - struct TransferInstruction { - size_t entry_index; - mori::io::MemoryDesc local_desc; - uint64_t local_offset; - mori::io::MemoryDesc remote_desc; - uint64_t remote_offset; - uint64_t size; + std::vector* results, bool recache_remote = true); + + // The remote half of ExecuteBatchGetPlan: submit every peer, then wait every + // peer. Split out for the ranged path, whose plan has no local half by + // construction (BatchGetRanges filters unroutable and self-routed keys out + // before building it), so threading the keys/dsts/sizes the local half needs + // would only pass three vectors nothing reads. + // `in_flight_window`, when set, runs after every peer has been posted and + // before any of them is waited on -- which is where local reads belong, so + // they overlap the wire instead of following it. A ranged plan has no local + // half by construction and passes nothing. + void ExecuteRemoteBatchGetPlan(const BatchGetPlan& plan, std::vector* results, + bool recache_remote, + const std::function& in_flight_window = {}); + // One place to say "this key was fetched but did not become local". Every + // way that can happen -- no slot, a slot that cannot be addressed as one run, + // a refused commit, a background pull that never landed -- reports here, so + // the counter measures the feature rather than one branch of it. + void NoteRangedInstallFailure(size_t count = 1); + + // One key's share of one arena round: the spans to fetch, where they land in + // the slice, and the caller buffers they are copied out to. A key needs more + // than one when its spans do not all fit at once. + struct RangeFetchUnit { + size_t original = 0; // caller key index + size_t object_size = 0; + std::optional route; + std::vector spans; + std::vector packed; // caller pointer + slice-relative offset + size_t bytes = 0; // == sum of spans + // The slice is byte-for-byte the whole object, which happens when the + // caller's ranges tile it in ascending order and it all fits in one unit. + // The whole-object reader (one call for every layer) hits this; the + // layer-wise reader never does. + bool holds_whole_object = false; }; - // A set of page transfers sharing the same (localMR, remoteMR) pair. All - // its pages collapse into ONE outer IO transfer whose inner offset/size - // vectors are scatter-gather segments, cutting CQE/status/post count vs - // one-transfer-per-page. `entry_indices` is the de-duplicated list of - // entries contributing pages to this group (for per-key failure mapping). - struct PairGroup { - mori::io::MemoryDesc local_desc; - mori::io::MemoryDesc remote_desc; - std::vector local_offsets; - std::vector remote_offsets; - std::vector sizes; - std::vector entry_indices; - }; - // Group `active` page transfers by (local_desc.id, remote_desc.id), - // preserving first-appearance order (stable). Assumes a fixed page size - // per (localMR, remoteMR) pair (1 buffer == 1 page size in the current - // model); mixed page sizes would require folding page size into the key. - static std::vector GroupTransfersByPair( - const std::vector& active); + // Serve whole-object ranged reads out of a fresh medium slot instead of the + // arena. Their span layout already equals the object, so the peer can write + // straight into the slot: one copy instead of two, the arena stays free for + // the readers that actually need it, and the arena mutex is never taken. + // Committing the slot afterwards is the local install, for nothing. + // + // Returns the units it could not take -- no slot to be had, or a slot that is + // not one contiguous run -- for the arena path to serve the ordinary way. + // That fallback is decided here, on a cheap local allocation, and never after + // a failed transfer: a transfer does not fail because of where it was + // pointed, and retrying it elsewhere would only double the latency of a + // failure the caller treats as fatal. + // `remote_bytes` accumulates what this path pulled over the wire. It has to + // be reported here rather than by the arena loop, because a batch whose units + // are all slot-served never reaches that loop -- and that is exactly the + // shape this path exists to make fast, so leaving it out would blank the + // bandwidth metric precisely where it matters. + std::vector ServeWholeObjectUnitsFromMedium( + const std::vector& keys, std::vector units, + std::unordered_map* unit_ok, double* remote_bytes); + + // Remote-only sibling of PartitionBatchGetTargets for ranged reads. Each key + // contributes one item carrying its arena slice, its span list, and the span + // total. Spans are passed by pointer rather than by value: the caller + // already owns them for the whole call, and copying one small vector per key + // per sub-batch is a heap allocation per key that buys nothing. Both the + // pointees and `keys` must outlive the returned plan. + BatchGetPlan PartitionBatchGetRangeTargets( + const std::vector& keys, const std::vector& arena_slices, + const std::vector& object_sizes, const std::vector& packed_bytes, + const std::vector*>& spans, + const std::vector>& routes); struct RemotePutEntry { size_t result_index; const BatchPutItem* item; SlotPlan plan; uint64_t slot_id; - std::optional> zero_copy; - bool use_staging = false; - uint64_t staging_offset = 0; bool failed = false; }; @@ -403,41 +830,18 @@ class PoolClient { size_t result_index; const BatchGetItem* item; SlotPlan plan; - std::optional> zero_copy; - bool use_staging = false; - uint64_t staging_offset = 0; bool failed = false; }; - // Remote SSD reads (one staging slot + lease per key) with bounded retry - // (default off) on transient NO_SLOT / reader-local lease expiry; an rpc - // failure is hard not-served, and a NOT_FOUND short-circuits as a miss. - void ProcessRemoteSsdBatchGet(const std::vector& items, std::vector* results); - - // Periodic SSD metrics provider (registered in Init() when SSD is enabled, run - // once per metrics flush tick in the metrics thread). Reads the cheap atomics - // on the pipeline / PeerSsdManager / PeerService and ships counter deltas + a - // staging gauge, keeping AddCounter off the commit/read hot paths. The - // last-shipped values below make each tick report only the delta. - void PublishSsdMetrics(); - struct SsdMetricsLastShipped { - uint64_t copy_enqueued = 0, copy_succeeded = 0, copy_failed = 0; - uint64_t copy_dropped_queue_full = 0, copy_dropped_stopped = 0; - uint64_t read_ok = 0, read_not_found = 0, read_size_too_large = 0, read_error = 0; - uint64_t read_no_slot = 0; - uint64_t copy_bytes = 0, read_bytes = 0; - uint64_t evict_rounds = 0, evict_victims = 0, evict_bytes_freed = 0, evict_backend_failed = 0; - uint64_t staging_expired_reclaims = 0, staging_slot_full_rejects = 0; - }; - SsdMetricsLastShipped ssd_metrics_last_; - bool AllocateRemotePutEntries(const std::vector& items, ::umbp::UMBPPeer::Stub* stub, std::vector* entries, std::vector* abort_slots, std::vector* results); - bool BuildRemotePutTransfers(std::vector& entries, PeerConnection& peer, - std::vector* transfers, - uint64_t* staging_bytes); + // Build one TransferItem per page, tagged with the entry index. Whether an + // item ends up zero-copy or staged, and how items group into wire transfers, + // are both the engine's business now — this only names endpoints. + bool BuildRemotePutTransfers(std::vector& entries, const std::string& node_id, + std::vector* items); void FinalizeRemotePutEntries(std::vector& entries, std::vector& abort_slots, std::vector* results, @@ -446,110 +850,67 @@ class PoolClient { bool PrepareRemoteGetEntries(const std::vector& items, PeerConnection& peer, ::umbp::UMBPPeer::Stub* stub, std::vector* entries, std::vector* results); - bool BuildRemoteGetTransfers(std::vector& entries, PeerConnection& peer, - std::vector* transfers, - uint64_t* staging_bytes); - void FinalizeRemoteGetEntries(std::vector& entries, std::vector* results); - - // One posted-but-not-yet-waited remote-DRAM read for a single peer; the - // scheduler waits it later. Owned by a unique_ptr and never moved after - // submit: the RDMA backend keeps a RAW TransferStatus* into `statuses`, so it - // must outlive BatchRead and never move (sized once to G, never resized). - struct RemoteDramGetInFlight { + bool BuildRemoteGetTransfers(std::vector& entries, const std::string& node_id, + std::vector* items); + void FinalizeRemoteGetEntries(std::vector& entries, std::vector* results, + bool recache_remote); + + // One posted-but-not-yet-waited remote read for a single peer; the scheduler + // waits it later. The lifetime contract that used to live here — statuses + // sized once and never moved, drained by the destructor — moved into the + // engine's TransferHandle, which is where the raw TransferStatus* the RDMA + // backend holds actually live. + struct RemoteGetInFlight { PeerConnection* peer = nullptr; std::vector entries; - std::vector groups; - mori::io::MemDescVec local_descs; - mori::io::MemDescVec remote_descs; - mori::io::BatchSizeVec local_offsets; - mori::io::BatchSizeVec remote_offsets; - mori::io::BatchSizeVec sizes_v; - std::vector statuses; // size G; built once, never resized/moved - mori::io::TransferStatusPtrVec status_ptrs; // &statuses[g] - mori::io::TransferUniqueIdVec ids; - bool drained = false; // set once statuses have been Wait()ed (drain or WaitRemoteBatchGet) - // Staging (non-zero-copy) state; zero / unlocked for the zero-copy path. - uint64_t staging_bytes = 0; - std::unique_lock staging_lock; // holds staging_mutex_ submit..memcpy - - // Safety net: BatchRead is posted (not waited) and the backend CQ callback - // holds raw TransferStatus* into `statuses`. On an early/exceptional destroy - // (before WaitRemoteBatchGet), drain so the callback never writes freed - // memory. No failure mapping / result backfill — that's WaitRemoteBatchGet. - ~RemoteDramGetInFlight() { - if (drained) return; - for (auto& s : statuses) s.Wait(); - } + std::unique_ptr handle; + bool drained = false; }; // Submit half: GetOrConnectPeer + EnsurePeerServiceConnection + - // PrepareRemoteGetEntries + BuildRemoteGetTransfers + GroupTransfersByPair + - // BatchRead (NOT waited). Returns the in-flight handle, or nullptr if nothing - // is in flight (peer unreachable / resolve / build failure — failed keys - // already written to *results). When `permit_staging` is false (the zero-copy - // submit-all path), a batch that needs staging is treated as a contract - // violation and failed rather than acquiring staging_mutex_ (a submit-all over - // multiple staging peers would deadlock on the single lock). When true (the - // serial staging path), staging_mutex_ is acquired here and parked in the - // in-flight until WaitRemoteBatchGet copies staging -> dst. - std::unique_ptr SubmitRemoteBatchGet( - const std::vector& items, std::vector* results, bool permit_staging); - // Wait half: wait every group (never break early), aggregate per-pair failure - // back to per-key (per-item AND); for a staging in-flight, copy staging_buffer_ - // -> user dst and release staging_mutex_; then FinalizeRemoteGetEntries. - void WaitRemoteBatchGet(RemoteDramGetInFlight& inflight, std::vector* results); - - // One posted-but-not-yet-waited remote-DRAM write for a single peer (same - // lifetime contract as RemoteDramGetInFlight: unique_ptr-owned, never moved - // after submit; the backend holds raw TransferStatus* into `statuses`, sized - // once to G). Put also carries the slot lifecycle: `entries`/`abort_slots` - // feed FinalizeRemotePutEntries and `stub` issues its commit/abort RPCs. - struct RemoteDramPutInFlight { + // PrepareRemoteGetEntries + BuildRemoteGetTransfers + engine Plan/Submit (NOT + // waited). Returns the in-flight handle, or nullptr if nothing is in flight + // (peer unreachable / resolve / build failure — failed keys already written + // to *results). + // + // There is no longer a permit_staging flag: a batch whose dst is unregistered + // is staged INSIDE the engine's Submit and comes back already settled, so a + // submit-all loop over several peers cannot deadlock on the bounce pool and a + // batch that mixes registered and unregistered buffers is no longer a + // contract violation — it just works. + std::unique_ptr SubmitRemoteBatchGet(const std::vector& items, + std::vector* results); + // Wait half: wait the handle (never breaks early), map per-plan failure back + // to per-key (per-item AND), then FinalizeRemoteGetEntries. + void WaitRemoteBatchGet(RemoteGetInFlight& inflight, std::vector* results, + bool recache_remote); + + // Put's counterpart. Also carries the slot lifecycle: `entries` / + // `abort_slots` feed FinalizeRemotePutEntries and `stub` issues its + // commit/abort RPCs. + struct RemotePutInFlight { PeerConnection* peer = nullptr; ::umbp::UMBPPeer::Stub* stub = nullptr; std::vector entries; // Malformed slots from Allocate (not in `entries`); Finalize appends // entry.failed slots and aborts the union (peer Abort is idempotent). std::vector abort_slots; - std::vector groups; - mori::io::MemDescVec local_descs; - mori::io::MemDescVec remote_descs; - mori::io::BatchSizeVec local_offsets; - mori::io::BatchSizeVec remote_offsets; - mori::io::BatchSizeVec sizes_v; - std::vector statuses; // size G; built once, never resized/moved - mori::io::TransferStatusPtrVec status_ptrs; // &statuses[g] - mori::io::TransferUniqueIdVec ids; + std::unique_ptr handle; bool drained = false; - uint64_t staging_bytes = 0; - std::unique_lock staging_lock; // holds staging_mutex_ submit(memcpy)..wait - - // Safety net (mirror Get): drain posted statuses on an early destroy so the - // backend CQ callback never writes freed memory. Slots are NOT committed/ - // aborted here — an un-committed slot never enters the master index, so it is - // correctness-neutral and the peer reaper reclaims it at pending_ttl. - ~RemoteDramPutInFlight() { - if (drained) return; - for (auto& s : statuses) s.Wait(); - } }; - // Submit half: allocate + build + (staging: memcpy src -> staging_buffer_ under - // staging_mutex_) + group + BatchWrite (NOT waited). Returns the in-flight, or - // nullptr if nothing is posted — in which case any allocated slots are aborted - // and the keys written kFailed here. Entries that fail during build but still - // post ride in the in-flight and are aborted by FinalizeRemotePutEntries at - // wait time (no early abort, avoids double-abort). permit_staging=false on the - // zero-copy submit-all path (staging would deadlock the single lock); true on - // the serial staging path, which parks staging_mutex_ in the in-flight. - std::unique_ptr SubmitRemoteBatchPut( - const std::vector& items, std::vector* results, - bool permit_staging); - // Wait half: wait every group (never break early), aggregate per-pair failure - // back to per-key (per-item AND), release staging_mutex_ (Put has no - // staging->dst copy-out), then FinalizeRemotePutEntries (commit survivors, - // abort failures + malformed slots). - void WaitRemoteBatchPut(RemoteDramPutInFlight& inflight, std::vector* results); + // Submit half: allocate + build + engine Plan/Submit (NOT waited). Returns + // the in-flight, or nullptr if nothing is posted — in which case any + // allocated slots are aborted and the keys written kFailed here. Entries + // that fail during build but still post ride in the in-flight and are aborted + // by FinalizeRemotePutEntries at wait time (no early abort, avoids + // double-abort). + std::unique_ptr SubmitRemoteBatchPut(const std::vector& items, + std::vector* results); + // Wait half: wait the handle, map per-plan failure back to per-key, then + // FinalizeRemotePutEntries (commit survivors, abort failures + malformed + // slots). + void WaitRemoteBatchPut(RemotePutInFlight& inflight, std::vector* results); }; } // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/range_map.h b/src/umbp/include/umbp/distributed/range_map.h new file mode 100644 index 000000000..e18ce9202 --- /dev/null +++ b/src/umbp/include/umbp/distributed/range_map.h @@ -0,0 +1,116 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include + +#include "umbp/common/range_utils.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +// The object-range -> page-range mapping, in one place. +// +// An object is stored as a page list; a caller names byte ranges of the object. +// Turning one into the other is the same arithmetic whether the pages live in +// this process's medium (PoolClient::BuildLocalRangeTransfers) or on a peer +// (PoolClient::BuildRemoteGetTransfers) — only the TransferRef the page index +// resolves to differs. It lives here so both callers share one implementation +// and so it can be unit-tested directly: a mistake here does not fail, it +// silently reads or writes the wrong bytes of a KV block. + +// Offset of a page within its buffer, as a plain byte offset. +inline uint64_t PageOffset(const PageLocation& page, uint64_t page_size) { + return static_cast(page.page_index) * page_size; +} + +// Bytes of page `i` that the object actually owns. Every page is full except +// the last, which is short whenever the object is not a page multiple. +inline uint64_t LogicalPageBytes(size_t i, size_t num_pages, uint64_t page_size, + size_t total_size) { + return (i + 1 == num_pages) ? (total_size - i * page_size) : page_size; +} + +// Does `size` plausibly occupy exactly `num_pages` pages of `page_size`? +// One page fewer would not hold it; one more would not have been allocated. +inline bool SizeMatchesAllocation(uint64_t size, size_t num_pages, uint64_t page_size) { + if (page_size == 0 || num_pages == 0 || size == 0) return false; + if (size > num_pages * page_size) return false; + if (size <= (num_pages - 1) * page_size) return false; + return true; +} + +// Walk [object_offset, object_offset + range_size) across `pages`, calling +// `emit` once per page the range touches. +// +// emit(page_index, tier_offset, fragment, copied_so_far) -> bool +// +// `tier_offset` is the byte offset of the fragment within pages[page_index]'s +// BUFFER (so the caller only has to supply the TransferRef for +// pages[page_index].buffer_index); `fragment` is how many bytes of this page +// the range covers; `copied_so_far` is how far into the range that fragment +// starts, which is what a contiguous destination wants as its own offset. +// Returning false from `emit` aborts the walk and makes this return false. +// +// Returns false — without emitting a partial result the caller could mistake +// for success — when the range runs off the end of the object, off the end of +// the page list, or into the part of the short last page the object does not +// own. Adjacent fragments are NOT merged here; the transfer engines coalesce +// while planning, and merging early would hide the page structure from them. +template +bool ForEachRangePageFragment(const std::vector& pages, uint64_t page_size, + uint64_t stored_size, size_t object_offset, size_t range_size, + Emit&& emit) { + if (range_size == 0) return false; + if (!SizeMatchesAllocation(stored_size, pages.size(), page_size)) return false; + if (IsObjectRangeOverflow(object_offset, range_size, static_cast(stored_size))) { + return false; + } + + size_t offset = object_offset; + size_t copied = 0; + while (copied < range_size) { + const size_t page_index = static_cast(offset / page_size); + const size_t within_page = static_cast(offset % page_size); + if (page_index >= pages.size()) return false; + + const uint64_t page_bytes = + LogicalPageBytes(page_index, pages.size(), page_size, static_cast(stored_size)); + if (within_page >= page_bytes) return false; + + const size_t fragment = + std::min(range_size - copied, static_cast(page_bytes - within_page)); + if (!emit(page_index, PageOffset(pages[page_index], page_size) + within_page, fragment, + copied)) { + return false; + } + + offset += fragment; + copied += fragment; + } + return true; +} + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/routing/route_get_strategy.h b/src/umbp/include/umbp/distributed/routing/route_get_strategy.h index 2d9c55bd2..0c43c424c 100644 --- a/src/umbp/include/umbp/distributed/routing/route_get_strategy.h +++ b/src/umbp/include/umbp/distributed/routing/route_get_strategy.h @@ -65,11 +65,16 @@ class RandomRouteGetStrategy : public RouteGetStrategy { const std::string& node_id) override; }; -/// Tier-priority strategy: prefer the fastest tier present (HBM > DRAM > SSD), -/// then pick a random replica within that tier. This is the distributed read -/// path's default — without it a key with both a DRAM and an SSD copy could be -/// routed to the slow SSD at random. Unknown-tier locations rank last. -class TierPriorityRouteGetStrategy : public RouteGetStrategy { +/// Locality-preferring strategy and the distributed read path's default: if the +/// requesting node holds a replica, read it locally (no RDMA); otherwise pick a +/// random replica. +/// +/// It does NOT rank tiers. The former HBM > DRAM > SSD order was deleted in the +/// backend-agnostic refactor Phase 4 rather than reimplemented, because every +/// medium is currently equivalent. Restoring a read preference is a matter of +/// the backend advertising one (design doc §3 BackendProperties), not of adding +/// a tier list back here. +class LocalPreferringRouteGetStrategy : public RouteGetStrategy { public: Location Select(const std::vector& locations, const std::string& node_id) override; std::vector BatchSelect(const std::vector>& per_key_locations, diff --git a/src/umbp/include/umbp/distributed/routing/route_put_strategy.h b/src/umbp/include/umbp/distributed/routing/route_put_strategy.h index 6a6ea101f..5f34533bf 100644 --- a/src/umbp/include/umbp/distributed/routing/route_put_strategy.h +++ b/src/umbp/include/umbp/distributed/routing/route_put_strategy.h @@ -49,6 +49,7 @@ struct RoutePutResult { std::string node_id; std::string peer_address; TierType tier = TierType::UNKNOWN; + std::string logical_tier; }; /// Abstract interface for batch RoutePut node placement. Implement this to @@ -89,10 +90,11 @@ class RoutePutStrategy { /// - node affinity: none / same-node / local-node. /// /// Both knobs are wired from env vars at master startup -/// (UMBP_ROUTE_PUT_SELECT_ALGO / UMBP_ROUTE_PUT_NODE_AFFINITY). Tier order is -/// always HBM -> DRAM; SSD is never a direct-put target. Projected capacity is -/// deducted on the batch-local candidates copy within SelectBatch; nothing is -/// written back to the registry. +/// (UMBP_ROUTE_PUT_SELECT_ALGO / UMBP_ROUTE_PUT_NODE_AFFINITY). There is no +/// tier order: every medium is treated as equivalent, so any tier a node +/// advertises capacity for is a valid put target (backend-agnostic refactor +/// Phase 4). Projected capacity is deducted on the batch-local candidates copy +/// within SelectBatch; nothing is written back to the registry. class ConfigurableRoutePutStrategy : public RoutePutStrategy { public: enum class SelectAlgo { kMostAvailable, kRandom }; @@ -113,24 +115,25 @@ class ConfigurableRoutePutStrategy : public RoutePutStrategy { private: /// Try only @p node_id on exactly @p tier; nullopt if it cannot fit - /// @p block_size. No cross-node, no cross-tier fallback. + /// @p block_size. No cross-node, no cross-tier fallback. Used by same-node + /// affinity to pin a whole batch to one (node, tier) — stickiness, not a tier + /// preference. std::optional TrySelectOnNodeTier( const std::vector& candidates, const std::string& node_id, TierType tier, - uint64_t block_size, const std::unordered_set& exclude_nodes) const; + uint64_t block_size, const std::unordered_set& exclude_nodes, + const std::string& logical_tier = {}) const; - /// Try only @p node_id, HBM then DRAM; nullopt if it cannot fit @p block_size. - /// Performs no cross-node fallback — that is the caller's explicit job. + /// Try only @p node_id, on whichever of its tiers has the most room; nullopt + /// if none can fit @p block_size. Performs no cross-node fallback — that is + /// the caller's explicit job. std::optional TrySelectOnNode( const std::vector& candidates, const std::string& node_id, uint64_t block_size, const std::unordered_set& exclude_nodes) const; - /// Base algorithm, tier priority paramount (HBM before DRAM): walk tiers in - /// order and route on the first tier that has room. Within a tier, most- - /// available picks the largest free space; random draws weighted by it. When - /// @p preferred_node is set and has room on the current tier, it wins over the - /// algorithm — but only within that tier, so tier priority is never broken - /// (a remote HBM node still beats the preferred node's DRAM). This is the - /// explicit global-fallback entry point. + /// Base algorithm over every (node, tier) pair with room: most-available picks + /// the largest free space, random draws weighted by it. When @p preferred_node + /// is set and has room on any of its tiers it wins over the algorithm outright. + /// This is the explicit global-fallback entry point. std::optional SelectByAlgo( const std::vector& candidates, uint64_t block_size, const std::unordered_set& exclude_nodes, diff --git a/src/umbp/include/umbp/distributed/ssd_read_lease.h b/src/umbp/include/umbp/distributed/ssd_read_lease.h deleted file mode 100644 index 10005d5ff..000000000 --- a/src/umbp/include/umbp/distributed/ssd_read_lease.h +++ /dev/null @@ -1,76 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -// Pure decision helpers for the reader-side remote SSD read lease path. -// -// These are intentionally free of PoolClient / RDMA / IO state so the lease -// gating policy can be unit-tested in isolation. The reader anchors its lease -// deadline at the moment BEFORE sending PrepareSsdRead (t_send) and treats the -// read as valid only while now <= t_send + lease_ttl_ms. This is strictly -// conservative against the peer: the peer starts the same TTL when it RECEIVES -// the request (received_at > t_send), so its reclaim point (received_at + ttl) -// is always later than the reader's deadline (t_send + ttl) — the reader gives -// up before the peer can reclaim and reassign the slot, so a read is never -// reported successful against a slot the peer has already recycled. -// -// Lease expiry is therefore a reader-local decision, not a wire status: there -// is no SSD_READ_LEASE_EXPIRED. On expiry the reader returns a transient -// retry (never a miss) and does NOT release — the peer reclaims by TTL. A -// reader-local lease expiry and a NO_SLOT response are the retryable outcomes; -// a failed/timed-out PrepareSsdRead is a hard failure (retrying a slow peer -// that may already hold a claimed slot just piles up more staging occupation), -// so it does not route through kRetry. - -#include -#include - -namespace mori::umbp::ssd_read_lease { - -// Reader-local lease outcome category, independent of PoolClient internals. -// Translated to PoolClient::SsdGetOutcome at the call site. -enum class GateOutcome { kSuccess, kRetry, kError }; - -struct GateDecision { - GateOutcome outcome; - bool release; // whether to issue a best-effort ReleaseSsdLease -}; - -// True once now is strictly past the deadline (t_send + lease_ttl_ms); now == -// deadline is still valid. ttl 0 expires as soon as now > t_send. -inline bool LeaseExpired(std::chrono::steady_clock::time_point t_send, uint64_t lease_ttl_ms, - std::chrono::steady_clock::time_point now) { - return now > t_send + std::chrono::milliseconds(lease_ttl_ms); -} - -// Decide the outcome of a prepared read. `expired` short-circuits to a -// transient retry with no release (the slot is left for the peer's TTL -// reclaim). Otherwise a successful RDMA serves the data and frees the slot -// fast via release; a failed RDMA is a hard error but still releases (the -// lease is still ours and matching by lease_id is safe). -inline GateDecision DecideSsdReadOutcome(bool expired, bool rdma_ok) { - if (expired) return {GateOutcome::kRetry, /*release=*/false}; - if (rdma_ok) return {GateOutcome::kSuccess, /*release=*/true}; - return {GateOutcome::kError, /*release=*/true}; -} - -} // namespace mori::umbp::ssd_read_lease diff --git a/src/umbp/include/umbp/distributed/transfer/composite_transfer_engine.h b/src/umbp/include/umbp/distributed/transfer/composite_transfer_engine.h new file mode 100644 index 000000000..4ee7e1149 --- /dev/null +++ b/src/umbp/include/umbp/distributed/transfer/composite_transfer_engine.h @@ -0,0 +1,133 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include + +#include "umbp/distributed/transfer/transfer_engine.h" + +namespace mori::umbp { + +// Holds the concrete engines and IS a TransferEngine, so the rest of the system +// holds exactly one pointer to exactly one type: PoolClient owns it, and +// backends receive the same object narrowed to MemoryRegistrar at Init. +// +// Registration FANS OUT — every engine that can reach a buffer gets to publish +// its handle, and the handles are merged into one TransferRef (see the comment +// on TransferRef for why that is a struct and not a variant). Transfer +// DISPATCHES — a (src, dst) pair goes to exactly one engine. That asymmetry is +// the whole reason this class exists; §4 argued it does not need a second class, +// only a second method shape. +// +// Selection is first-match in registration order, so order encodes preference: +// register LocalCopyEngine before MoriIoEngine and a both-endpoints-local pair +// never reaches the wire. In practice the two are disjoint (mori-io requires +// exactly one remote endpoint), so the order is documentation rather than a +// tie-break — but a future engine pair may genuinely overlap. +// +// NOT implemented: chaining, where no single engine spans the pair and the +// layer composes two hops through a bounce buffer (§4's "3FS file on node A, +// reader on node B"). It cannot arise while every endpoint is memory; it +// arrives with the first non-memory endpoint, which is SsdBackend's FileRef. +// It is also the transfer layer's OBSERVABILITY point, for the same reason it +// is the dispatch point: it is the one place that knows which engine carried +// which plan. Bytes, plan counts, failures and in-flight time are measured +// here and published under engine= / direction= labels, so an engine added +// with AddEngine() shows up in the existing dashboard panels without writing +// or wiring a single metric (see component_metrics.h). +class CompositeTransferEngine final : public TransferEngine { + public: + // Order is preference order. Null is ignored. + void AddEngine(std::unique_ptr engine); + + const char* Name() const override { return "CompositeTransferEngine"; } + + // Fan out: every sub-engine registers, and the per-transport handles are + // merged. Returns an invalid ref only when no engine claimed the memory. + TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) override; + // Fan out to the first file-capable sub-engine (GdsEngine): it owns the fd + // registration outright and returns a File ref carrying its handle. Invalid + // when no file engine is configured. + TransferRef RegisterFile(int fd, uint64_t offset, uint64_t size) override; + void Deregister(const TransferRef& ref) override; + + bool CanHandle(const TransferRef& src, const TransferRef& dst) const override; + + // Partition items by the engine that will carry them, delegate, and tag each + // returned plan with its engine so Submit can route it back. + TransferPlanSet Plan(const std::vector& items) const override; + + std::unique_ptr Submit(std::vector plans) override; + + // Which engine would carry this pair, or null. Exposed for tests and for + // callers that want to know whether a transfer is servicable before building + // a whole item list. + TransferEngine* SelectEngine(const TransferRef& src, const TransferRef& dst) const; + + // The generic per-engine series measured above, followed by whatever each + // sub-engine publishes about its own internals (engine= already stamped on + // both). One source for the whole transfer layer, so PoolClient publishes it + // with the same call it uses for a backend. + std::vector SampleMetrics() const override; + + private: + class FanOutHandle; + + // Per (engine, direction) accumulators. Indexed by the engine's position in + // engines_, which is fixed after composition, so the hot path is an array + // index and never a lookup. Relaxed atomics: read once per metrics tick, + // never correctness state. + struct DirectionCounters { + std::atomic plans{0}; + std::atomic failed_plans{0}; + std::atomic bytes{0}; + std::atomic nanos{0}; + }; + // kPush / kPull / kLocal, indexed by TransferDirection. + static constexpr size_t kDirectionCount = 3; + struct EngineCounters { + DirectionCounters by_direction[kDirectionCount]; + }; + + static size_t DirectionIndex(TransferDirection dir) { return static_cast(dir); } + // Position of `engine` in engines_, or engines_.size() when it is not ours. + size_t IndexOf(const TransferEngine* engine) const; + // Charge a settled plan back to the engine that carried it. Called from + // FanOutHandle::Wait, which is why the counters outlive any single submit. + void RecordSettled(size_t engine_index, TransferDirection dir, uint64_t bytes, uint64_t nanos, + uint64_t failed_plans); + + std::vector> engines_; + // One per engine, in engines_ order. Held by unique_ptr because the atomics + // inside are neither copyable nor movable and engines_ grows during + // composition. + std::vector> counters_; + // Items Plan() could not route to any engine. Not charged to an engine — + // there was none — so it rides engine="none", status="rejected". + mutable std::atomic rejected_items_{0}; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/transfer/gds_engine.h b/src/umbp/include/umbp/distributed/transfer/gds_engine.h new file mode 100644 index 000000000..1f86aa0d6 --- /dev/null +++ b/src/umbp/include/umbp/distributed/transfer/gds_engine.h @@ -0,0 +1,95 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/transfer/transfer_engine.h" + +namespace mori::umbp { + +// The SSD-file-to-GPU implementation of TransferEngine, over AMD hipFile (GDS). +// +// WHY THIS EXISTS. SsdBackend publishes a File TransferRef (an O_DIRECT fd + a +// byte range) instead of staging its bytes through host DRAM. No memory engine +// can move a file endpoint — their CanHandle keys on HasHostPtr() / +// HasMemoryDesc(), both false for a file ref — so a (file, GPU) pair was +// unservable until this engine. It reads the range straight into device memory +// with hipFileRead, which DMAs drive -> GPU with no host bounce (design §4/§6). +// +// READ ONLY, FOR NOW. CanHandle claims exactly (file source, device +// destination): the prefill KV-load path. The reverse (GPU -> file, via +// hipFileWrite) is a later change, and a host destination is left to the staged +// path, so this engine never widens what the memory engines already serve. +// +// The fd handle is registered once per fd and shared by every range on that +// segment (a segment file holds many keys): RegisterFile ref-counts, Deregister +// releases at zero. hipFileHandle_t is kept as an opaque void* so this header +// carries no hipfile dependency; only the .cpp includes . +class GdsEngine final : public TransferEngine { + public: + GdsEngine() = default; + ~GdsEngine() override; + + const char* Name() const override { return "GdsEngine"; } + + // Not a memory engine: an invalid ref tells the composite's memory fan-out to + // skip it. + TransferRef RegisterMemory(void*, size_t, mori::io::MemoryLocationType, int) override { + return TransferRef{}; + } + + // Register an O_DIRECT fd as a hipFile handle (once per fd, ref-counted) and + // return a File ref carrying it. Invalid ref if registration fails. + TransferRef RegisterFile(int fd, uint64_t offset, uint64_t size) override; + void Deregister(const TransferRef& ref) override; + + // A file source into a device destination — the read path — and nothing else. + bool CanHandle(const TransferRef& src, const TransferRef& dst) const override; + + // One plan per (file range -> device buffer). Grouping/coalescing is left + // out deliberately: hipFileRead already moves a contiguous range in one DMA, + // and a batch's ranges land in disjoint GPU buffers. + TransferPlanSet Plan(const std::vector& items) const override; + + // Performs the reads inline with hipFileRead; the returned handle is already + // settled (same contract as the copy engines). + std::unique_ptr Submit(std::vector plans) override; + + private: + // One registered fd: its opaque hipFileHandle_t and a ref count so a segment + // shared by many keys registers once and releases once. + struct HandleEntry { + void* handle = nullptr; // hipFileHandle_t + int refcount = 0; + }; + + mutable std::mutex handles_mutex_; + std::unordered_map handles_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/transfer/hbm_copy_engine.h b/src/umbp/include/umbp/distributed/transfer/hbm_copy_engine.h new file mode 100644 index 000000000..1e3c745dd --- /dev/null +++ b/src/umbp/include/umbp/distributed/transfer/hbm_copy_engine.h @@ -0,0 +1,172 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include + +#include "umbp/common/host_registration.h" +#include "umbp/distributed/transfer/transfer_engine.h" + +namespace mori::umbp { + +// The one-endpoint-is-GPU implementation of TransferEngine. +// +// WHY THIS EXISTS. Before this engine, a local pair with a GPU endpoint was +// servable by NO engine in tree, so an HBM backend's local Put/Get could not +// complete at all: +// +// LocalCopyEngine::CanHandle requires BOTH endpoints loc == CPU — a GPU +// endpoint is local but not memcpy-able, and the +// header says so ("falls to an engine that can +// reach it"). This is that engine. +// MoriIoEngine::CanHandle requires EXACTLY ONE remote endpoint +// (`src_remote == dst_remote` -> false), so a +// both-local pair is refused however it is typed. +// +// The REMOTE side of HBM already worked before this engine and still does not +// route here: a peer's HBM is reached through mori-io, whose MemoryDesc carries +// loc=GPU on the wire already (design doc §2, "this is why the remote path +// already works against an HBM peer unmodified, and why the local fast path +// does not"). This engine closes exactly the local half of that sentence. +// +// SELECTION IS DISJOINT, not preference-ordered. The three engines partition +// the (src, dst) space with no overlap, so CompositeTransferEngine's +// first-match order is documentation here rather than a tie-break: +// +// both local, both CPU -> LocalCopyEngine +// both local, either GPU -> HbmCopyEngine (this) +// exactly one remote -> MoriIoEngine +// +// That covers H2D (a caller's host Put src into an HBM slot), D2H (an HBM slot +// out to a caller's host Get dst), and D2D (HBM-to-HBM on one node, e.g. a +// mirror between two GPUs). +// +// NOTE ON `host_ptr`. TransferRef::host_ptr is documented as the "process-local +// view", not "host memory": for an hipMalloc'd buffer the DEVICE pointer is that +// view, since it is what hipMemcpy accepts and what is valid in this process. +// So HbmBackend registers its buffers with host_ptr = device pointer and +// loc = GPU, and HasHostPtr() reads as "addressable here" for both media. No +// new TransferRef field or `kind` tag was needed — which is the reverse of the +// SSD/FileRef case the base header reserves. +class HbmCopyEngine final : public TransferEngine { + public: + const char* Name() const override { return "HbmCopyEngine"; } + + // Nothing to pin. A device pointer is already a usable endpoint for + // hipMemcpy, exactly as a host pointer is for memcpy — the registration that + // DOES cost something (an RDMA MR) is MoriIoEngine's, and the composite fans + // registration out to it in the same call. + TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) override { + return TransferRef::HostBytes(base, size, loc, device); + } + void Deregister(const TransferRef&) override {} + + // Declare a host region as kernel-addressable, enabling the gather fast path + // for plans whose host side lies inside it. + // + // A gather kernel cannot dereference plain mmap memory — it faults the GPU — + // so the region must be hipHostRegister'd first, which is what + // HostTierRegistration does (asynchronously for large pools). Registration + // is best-effort and this is purely an optimisation hint: until it completes, + // and forever if it fails, Submit simply takes the hipMemcpy path it took + // before. + // + // PoolClient calls this once per buffer each medium publishes, after Init. + // It belongs on the engine rather than at the call sites because the engine + // is the only layer that sees the segment shape a copy actually decomposes + // into, which is what decides whether a kernel beats the copy engine. + void AddHostGatherRegion(void* base, size_t bytes); + void ClearHostGatherRegions(); + + // Both endpoints addressable in this process, at least one of them GPU. + bool CanHandle(const TransferRef& src, const TransferRef& dst) const override; + + // Group by (src base, dst base) and coalesce exactly-contiguous segments — + // same shape as LocalCopyEngine::Plan, and it pays MORE here. Each hipMemcpy + // carries a fixed launch/synchronization cost that a memcpy does not, so + // collapsing a run of adjacent pages into one call removes real overhead + // rather than just loop iterations. This is the concrete instance of the + // reason Plan() is virtual rather than a base-class helper. + TransferPlanSet Plan(const std::vector& items) const override; + + // Performs the copies inline; the returned handle is already settled. + // + // Synchronous hipMemcpy, deliberately, matching LocalCopyEngine's contract + // that a local transfer costs no round trip. An async/stream variant would + // overlap the segments within one plan, but the parallelism that actually + // matters here is ACROSS keys and already exists above this engine, in + // PoolClient's UMBP_DRAM_{READ,WRITE}_THREADS executors. A shared stream + // would also need its own synchronization, since Submit is called from those + // worker threads concurrently — cost with no matching win at KV-block sizes, + // where a multi-MiB copy dwarfs the per-call launch overhead. + std::unique_ptr Submit(std::vector plans) override; + + private: + // Why a plan did not take the gather kernel. The hot path never computes + // this: it is filled in only when debug mode is on (UMBP_HBM_COPY_DEBUG), so + // that the log can say which path carried a batch AND why the other one + // declined — the two questions are always asked together when a bandwidth + // number looks wrong. + enum class GatherSkip : uint8_t { + kTaken = 0, + kUnrecorded, // debug mode off, or a plan the pass never classified + kDisabled, // UMBP_DRAM_GATHER_KERNEL=0, or latched off by a failure + kKindNotHostDevice, // D2D: no host side to register, so no kernel form + kNoDevice, // neither endpoint named a device + kHostNotGatherable, // host side outside every HostTierRegistration, or + // registered but unreachable from this device + kNoFragments, // plan contributed no segments + kTooFewFragments, // < 2 in the bucket; one segment is hipMemcpy's best case + kFragmentAtOrAboveThreshold, // mean segment >= kGatherFragmentThreshold + kSetDeviceFailed, + kNoStream, + kLaunchFailed, + kSyncFailed, + }; + static const char* GatherSkipName(GatherSkip reason); + + // Run the gather kernel over the batch's eligible segments, bucketed by + // device. Returns a per-plan flag: 1 means fully copied by the kernel, 0 + // means "not taken" and never "failed" — Submit then runs its hipMemcpy loop + // for it, which reaches the same result whether or not fragments landed. + // May leave the current HIP device changed. + // + // `skip_reasons`, when non-null, is resized to plans.size() and filled with + // the per-plan verdict. Submit passes it only in debug mode. + std::vector GatherEligiblePlans(const std::vector& plans, + std::vector* skip_reasons); + + // The address `device_id` must use to reach `ptr`, or null when no registered + // region covers it or that device cannot gather from the one that does. + // Kernels take this, never the host pointer -- see HostTierRegistration. + void* HostRegionDeviceAddress(const void* ptr, size_t size, int device_id) const; + + mutable std::mutex host_regions_mutex_; + std::vector> host_regions_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/transfer/local_copy_engine.h b/src/umbp/include/umbp/distributed/transfer/local_copy_engine.h new file mode 100644 index 000000000..146e94584 --- /dev/null +++ b/src/umbp/include/umbp/distributed/transfer/local_copy_engine.h @@ -0,0 +1,99 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include + +#include "umbp/distributed/transfer/transfer_engine.h" + +namespace mori::umbp { + +// Non-temporal AVX2 block copy where it pays, plain memcpy otherwise. +// +// Exposed (rather than kept file-local to the engine) because the async +// re-cache path also copies a whole KV block off the Get critical path and +// wants the same cache-bypass behavior. Honors UMBP_DRAM_NT_COPY=0. +void HostCopyBlock(void* dst, const void* src, size_t size); + +// Declare that copies issued on THIS thread are not on anyone's critical path, +// so LocalCopyEngine::Submit should run them on the calling thread instead of +// fanning them out. +// +// Fanning out is a latency trade: it spends other cores to finish one caller's +// copy sooner. That is the right trade for a thread someone is blocked on and +// the wrong one for a background worker, which has no waiter and whose helpers +// land on whoever does. Measured: with the asynchronous locality install +// fanning out too, a remote layer-wise restore at 8 ranks saw TTFL go from +// 7,470 to 12,154 us while its own throughput barely moved -- and the whole +// regression vanished under --no-prefetch, which is what identified the +// installer as the thief. +// +// Thread-local and sticky, so a worker sets it once at start-up rather than +// around each call. +void MarkThreadBackgroundCopies(); +bool ThreadDoesBackgroundCopies(); + +// The both-endpoints-are-local implementation of TransferEngine. +// +// This is what makes the local fast path stop being a special case (design doc +// §5 Phase 6). Before this phase PoolClient held a raw base pointer per DRAM +// buffer and memcpy'd through it directly, which is why the local path was +// host-DRAM-only, why PoolClient had to name a concrete backend type to get +// those pointers, and why a second live medium would have needed a tier branch +// in the copy loop. Now a local access is just a transfer whose endpoints are +// both local, and the planner picks this engine for it. +// +// §2 warned that the fold-in is only correct if it is also free: "the local +// fast path exists because a memcpy beats an engine round trip... this is a +// measurement, not an argument." It is free here because nothing round trips — +// Submit performs the copies synchronously on the calling thread and returns an +// already-settled handle, so the cost over the old loop is one small plan +// vector per key, against a multi-MiB copy. The cross-key parallelism +// (UMBP_DRAM_{READ,WRITE}_THREADS) stays where it was, above this engine. +class LocalCopyEngine final : public TransferEngine { + public: + const char* Name() const override { return "LocalCopyEngine"; } + + // Nothing to pin: a process-local pointer is already a usable endpoint. + TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) override { + return TransferRef::HostBytes(base, size, loc, device); + } + void Deregister(const TransferRef&) override {} + + // Both endpoints must be host-addressable in this process. Note this is + // deliberately narrower than "both are local": a GPU endpoint is local but + // not memcpy-able, and falls to an engine that can reach it. + bool CanHandle(const TransferRef& src, const TransferRef& dst) const override; + + // Group by (src base, dst base) and coalesce exactly-contiguous segments, so + // a multi-page key whose pages happen to be adjacent collapses into one + // memcpy instead of one per page. + TransferPlanSet Plan(const std::vector& items) const override; + + // Performs the copies inline; the returned handle is already settled. + std::unique_ptr Submit(std::vector plans) override; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/transfer/mori_io_engine.h b/src/umbp/include/umbp/distributed/transfer/mori_io_engine.h new file mode 100644 index 000000000..acdf82af2 --- /dev/null +++ b/src/umbp/include/umbp/distributed/transfer/mori_io_engine.h @@ -0,0 +1,151 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "mori/io/engine.hpp" +#include "umbp/distributed/transfer/transfer_engine.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +// The mori-io (RDMA) implementation of TransferEngine, and the only place in +// UMBP that names mori::io::IOEngine. +// +// Everything moved here from PoolClient in Phase 6 is something that had to +// know about the wire: +// * the IOEngine itself, its RDMA backend config, and engine-desc packing +// * remote engine registration and the peer buffer-descriptor cache — a +// peer's MemoryDesc is a transfer-layer fact, not a routing fact +// * GroupTransfersByPair, now Plan(): collapsing a scatter into one transfer +// per (srcMR, dstMR) is an RDMA optimization, which is why Plan is virtual +// * the bounce buffer, because the transfer layer is the ONLY layer with a +// completion signal. A staging pool inside a storage backend has to guess +// with a TTL, which is the deleted PrepareSsdRead lease under another name. +// +// Threading: RegisterMemory/Deregister, the remote caches and the bounce pool +// each take their own lock; Plan() is pure and takes none. +class MoriIoEngine final : public TransferEngine, public PeerDirectory { + public: + // `bounce_bytes` sizes the single staging region used when a caller's buffer + // is not registered. Zero disables staging: an unregistered endpoint is then + // rejected by Plan rather than bounced. + MoriIoEngine(std::string engine_key, mori::io::IOEngineConfig io_config, uint64_t bounce_bytes); + ~MoriIoEngine() override; + + // Create the IOEngine + RDMA backend and register the bounce buffer. + // Returns false if the engine could not be created; the object is then + // inert (Active() == false) rather than unusable — a node with no RDMA still + // registers memory (host_ptr only) and serves purely local transfers through + // a different engine. + bool Init(); + void Shutdown(); + + bool Active() const { return io_engine_ != nullptr; } + + // ---- TransferEngine ---- + + const char* Name() const override { return "MoriIoEngine"; } + + TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) override; + void Deregister(const TransferRef& ref) override; + + // True iff exactly one endpoint is a registered REMOTE memory and the other + // is local — either registered (zero copy) or, when a bounce buffer is + // configured, merely addressable (staged). + bool CanHandle(const TransferRef& src, const TransferRef& dst) const override; + + // Group by (src MR id, dst MR id), first-appearance order, coalescing + // segments that are exactly contiguous on BOTH sides. Items whose local side + // is unregistered are grouped separately and assigned bounce offsets; a + // bounce group is chunked to fit the pool, and an item larger than the whole + // pool is rejected. + TransferPlanSet Plan(const std::vector& items) const override; + + std::unique_ptr Submit(std::vector plans) override; + + // ---- PeerDirectory ---- + // + // mori-io's half of "learning how to reach a node": the peer's EngineDesc has + // to be registered with the local engine and its published MemoryDescs + // unpacked before any transfer can name them. This state used to live in + // PoolClient::PeerConnection, which meant the client had to understand + // descriptors it never dereferenced. + + std::vector PackedLocalEngineDesc() const override; + bool EnsureRemoteEngine(const std::string& node_id, + const std::string& packed_engine_desc) override; + bool HasRemoteEngine(const std::string& node_id) const override; + void ForgetRemote(const std::string& node_id) override; + void CacheRemoteBuffers(const std::string& node_id, + const std::vector& descs) override; + bool HasRemoteBuffers(const std::string& node_id) const override; + std::vector RemoteBufferSnapshot(const std::string& node_id, + uint32_t backend_id) const override; + + const std::string& LocalEngineKey() const { return local_engine_key_; } + + private: + class RdmaHandle; + + struct RemoteNode { + mori::io::EngineDesc engine_desc; + bool engine_registered = false; + // One shelf per peer-side backend, each indexed by that backend's own + // buffer_index. A single flat vector here was the client half of the + // mixed-media corruption: two backends both publish a buffer 0, and + // whichever descriptor arrived first answered for both. + std::array, kMaxBackendsPerPeer> buffers; + }; + + // Post `plan` and hand the statuses to `handle`. Caller owns sequencing. + bool PostPlans(const std::vector& plans, RdmaHandle* handle); + + // Run one bounce plan to completion: stage in (kPush), post, wait, stage out + // (kPull), release. Returns false and fills `failure` on any error. + bool RunBouncePlanInline(const TransferPlan& plan, TransferFailure* failure); + + std::string local_engine_key_; + mori::io::IOEngineConfig io_config_; + std::unique_ptr io_engine_; + + mutable std::mutex remotes_mutex_; + std::unordered_map remotes_; + + // Bounce pool: one region, one mutex. A plan that needs it completes inside + // Submit (see TransferEngine::Submit), so the lock is never held across a + // return and a submit-all loop over several peers cannot deadlock on it. + uint64_t bounce_size_ = 0; + std::unique_ptr bounce_buffer_; + TransferRef bounce_ref_; + std::mutex bounce_mutex_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/transfer/transfer_engine.h b/src/umbp/include/umbp/distributed/transfer/transfer_engine.h new file mode 100644 index 000000000..eeaecb6a3 --- /dev/null +++ b/src/umbp/include/umbp/distributed/transfer/transfer_engine.h @@ -0,0 +1,410 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "mori/io/common.hpp" +#include "umbp/distributed/metrics/component_metrics.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +// --------------------------------------------------------------------------- +// The transfer layer (design doc §4) +// +// There is exactly ONE byte-moving path in the system and it is here. A +// storage backend owns bytes and publishes endpoints for them; the transfer +// layer decides how to move bytes between two endpoints. Admitting a second +// byte-moving path at the backend level is what produced the coupling this +// refactor exists to remove. +// +// mori-io is a MEMORY-to-memory engine (BackendType is {XGMI, RDMA, TCP, +// FABRIC}, MemoryLocationType is {CPU, GPU}); it has no file or object +// endpoint. Rather than teach it one, UMBP abstracts the transfer engine on +// its own side and makes mori-io one implementation — so mori-io's other +// consumers see no blast radius. +// --------------------------------------------------------------------------- + +// One endpoint of a transfer. +// +// NOT a std::variant, though §4 sketched one. Registration FANS OUT: the same +// buffer is reachable as a raw process-local pointer AND as an RDMA-registered +// MR at the same time, and which one is used is a property of the (src, dst) +// PAIR, not of the buffer. A variant would force a buffer to be one or the +// other, and would make the case this phase turns on — memcpy between two +// buffers that also happen to be RDMA-registered — inexpressible. So a +// TransferRef carries the per-transport handles side by side and each engine +// reads the one it needs. +// +// mori::io::MemoryDesc already works this way, and says so: ipcHandle and +// fabricHandle sit next to each other "so XGMI and Fabric backends can register +// the same memory without clobbering each other's handle." +// +// Empty handles are how an endpoint says "not reachable that way": +// * host_ptr == nullptr -> not addressable in this process (a peer's memory) +// * mem.size == 0 -> not registered with mori-io (a plain user buffer, +// or a node with no RDMA configured) +// An endpoint with neither is invalid and no engine will accept it. +// +// FILE ENDPOINT (kind == File): the "one-file change" the note below predicted. +// A `kind` tag plus a second handle set — an O_DIRECT fd, a file offset, and an +// opaque engine handle — lets a backend publish a byte range of a file that a +// file engine (GdsEngine) moves straight between the drive and device memory, +// with no host bounce. The memory-only engines ignore it for free: their +// CanHandle keys on HasHostPtr()/HasMemoryDesc(), both false for a file ref, so +// the existing per-pair selection already routes it to the only engine that +// claims it. SsdBackend is the backend that needed it (design doc §4/§6). +struct TransferRef { + enum class Kind : uint8_t { Memory, File }; + + // Which handle set below is authoritative. Memory refs (the default) carry a + // process-local pointer and/or a mori-io descriptor; File refs carry an + // O_DIRECT fd and offset. + Kind kind = Kind::Memory; + + // Process-local view. Set for anything this node allocated or a caller + // handed us; null for a peer's memory. + void* host_ptr = nullptr; + uint64_t size = 0; + mori::io::MemoryLocationType loc = mori::io::MemoryLocationType::CPU; + int device = -1; + + // mori-io registration. Valid iff an IOEngine registered these bytes + // (local), or they were hydrated from a peer's published descriptor (remote). + mori::io::MemoryDesc mem{}; + + // File view (kind == File). A byte range [file_offset, file_offset + size) of + // an O_DIRECT file. gds_handle is a file engine's opaque registration of the + // fd (a hipFileHandle_t), left null until RegisterFile fills it; `size` above + // is reused for the range length. + int file_fd = -1; + uint64_t file_offset = 0; + void* gds_handle = nullptr; + + bool HasHostPtr() const { return host_ptr != nullptr && size > 0; } + bool HasMemoryDesc() const { return mem.size > 0; } + bool IsFile() const { return kind == Kind::File && file_fd >= 0; } + bool Valid() const { return HasHostPtr() || HasMemoryDesc() || IsFile(); } + + // Unregistered, directly-addressable bytes: a caller's Put src / Get dst that + // was never passed to RegisterMemory. + static TransferRef HostBytes(void* p, uint64_t n, + mori::io::MemoryLocationType l = mori::io::MemoryLocationType::CPU, + int dev = -1) { + TransferRef r; + r.host_ptr = p; + r.size = n; + r.loc = l; + r.device = dev; + return r; + } + + // A peer's memory: reachable only through mori-io. + static TransferRef Remote(mori::io::MemoryDesc desc) { + TransferRef r; + r.size = desc.size; + r.loc = desc.loc; + r.device = desc.deviceId; + r.mem = std::move(desc); + return r; + } + + // A byte range of an O_DIRECT file, addressable only by a file engine. + // `handle` is the engine's registration of `fd`, filled in by RegisterFile. + static TransferRef File(int fd, uint64_t offset, uint64_t n, void* handle = nullptr) { + TransferRef r; + r.kind = Kind::File; + r.file_fd = fd; + r.file_offset = offset; + r.size = n; + r.gds_handle = handle; + return r; + } +}; + +// --------------------------------------------------------------------------- +// MemoryRegistrar — the ONLY view a storage backend gets +// +// Phase 5 Rule C ("no backend may call Plan/Submit/Wait") was going to be a +// lint rule because §4 folded registration into TransferEngine, handing +// backends an object they could transfer with. That conflated ONE CLASS with +// ONE INTERFACE. Keeping one object and giving it two views turns the rule +// into a compile error at zero runtime cost: MediumBackend::Init takes a +// MemoryRegistrar*, which has no Submit to call. +// --------------------------------------------------------------------------- +class MemoryRegistrar { + public: + virtual ~MemoryRegistrar() = default; + + MemoryRegistrar(const MemoryRegistrar&) = delete; + MemoryRegistrar& operator=(const MemoryRegistrar&) = delete; + + // Pin [base, base+size) for transfer and return the merged handle set. The + // CALLER supplies the facts a descriptor cannot recover — location type, + // device — because the caller is the allocator (design doc §6, "allocation + // and registration belong to the backend"). + // + // Always returns a ref with host_ptr set; whether it also carries an + // mori-io registration depends on what is configured. A node with no RDMA + // still gets a usable local-only endpoint rather than a failure. + virtual TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) = 0; + + // Register an O_DIRECT file range so a file engine can move it directly + // between the drive and device memory, returning a File TransferRef with the + // engine's handle filled in. Default: no file-capable engine, so the range + // is unservable and an invalid ref comes back. Only GdsEngine (and the + // composite that fans out to it) overrides this — the memory-only engines + // never touch a file. + virtual TransferRef RegisterFile(int fd, uint64_t offset, uint64_t size) { + (void)fd; + (void)offset; + (void)size; + return TransferRef{}; + } + + // Release every registration held in `ref`. Tolerates an invalid ref. + virtual void Deregister(const TransferRef& ref) = 0; + + protected: + MemoryRegistrar() = default; +}; + +// --------------------------------------------------------------------------- +// PeerDirectory — learning how to reach another node +// +// Some transports cannot address a peer until they have been TOLD about it: +// mori-io needs the peer's EngineDesc registered and the peer's published +// MemoryDescs unpacked before a single byte can move. That handshake is not a +// transfer, so it does not belong on TransferEngine — but it is also not the +// client's business HOW it happens, so it must not be reached by naming a +// concrete engine either. +// +// So it is its own interface, implemented by the engines that need it. +// PoolClient obtains one at composition time and drives peer setup through it; +// a second remote transport implements this and PoolClient does not change. +// A memcpy engine has no peers and implements nothing. +// +// Note "directory", not "cache": these calls are the authority on what this +// process knows about a peer's transport identity, and ForgetRemote is how a +// failed handshake is undone. +// --------------------------------------------------------------------------- +class PeerDirectory { + public: + virtual ~PeerDirectory() = default; + + PeerDirectory(const PeerDirectory&) = delete; + PeerDirectory& operator=(const PeerDirectory&) = delete; + + // This node's own transport identity, in the opaque form peers exchange. + // Empty when the transport is not active. + virtual std::vector PackedLocalEngineDesc() const = 0; + + // Learn a peer's transport identity from its packed descriptor. Idempotent. + // False only when the descriptor cannot be understood. An empty descriptor + // means "nothing new" and succeeds iff the peer is already known. + virtual bool EnsureRemoteEngine(const std::string& node_id, const std::string& packed) = 0; + virtual bool HasRemoteEngine(const std::string& node_id) const = 0; + + // Drop everything known about a peer, so the next access re-handshakes. + virtual void ForgetRemote(const std::string& node_id) = 0; + + // Learn the endpoints a peer published, by (backend_id, buffer_index). + // Idempotent: an address already known is left alone, so a stale descriptor + // cannot overwrite a live one. + // + // The address is the PAIR. buffer_index is backend-local — every backend on + // a peer numbers from 0, and each publishes exactly one buffer today — so + // caching by index alone let one medium's descriptor stand in for another's + // and served reads out of the wrong memory. See BufferMemoryDesc in + // umbp.proto. + virtual void CacheRemoteBuffers(const std::string& node_id, + const std::vector& descs) = 0; + virtual bool HasRemoteBuffers(const std::string& node_id) const = 0; + + // Every known endpoint for ONE of a peer's backends, indexed by that + // backend's own buffer_index. Taken ONCE per batch so a transfer-build loop + // never locks per page; a batch's pages may span backends, so it is taken per + // distinct backend_id in the batch rather than once per peer. An index that + // was never learned comes back invalid and the caller degrades that key to a + // miss. + virtual std::vector RemoteBufferSnapshot(const std::string& node_id, + uint32_t backend_id) const = 0; + + protected: + PeerDirectory() = default; +}; + +// --------------------------------------------------------------------------- +// Work +// --------------------------------------------------------------------------- + +// One contiguous copy request. `tag` is the caller's own per-entry index; it +// survives grouping so a per-plan failure maps back to the keys that +// contributed to it. +struct TransferItem { + TransferRef src; + uint64_t src_offset = 0; + TransferRef dst; + uint64_t dst_offset = 0; + uint64_t size = 0; + size_t tag = 0; +}; + +// Which endpoint of a plan is this node's. A plan is uniform in direction +// because it is grouped by (src endpoint, dst endpoint). +enum class TransferDirection { + kPush, // local -> remote (mori-io BatchWrite) + kPull, // remote -> local (mori-io BatchRead) + kLocal // local -> local (no wire) +}; + +// One grouped unit of work, produced by Plan() and consumed by Submit(). +// +// Collapsing a scatter into one plan per endpoint pair is an RDMA optimization +// — it cuts CQE and post count — which is exactly why Plan() is virtual and +// not a base-class helper: a memcpy engine gains nothing from pair grouping and +// a file engine would rather group by file and merge adjacent offset ranges. +// Submit-all-then-wait IS universal, and lives with the caller's scheduler. +struct TransferPlan { + TransferRef src; + TransferRef dst; + std::vector src_offsets; + std::vector dst_offsets; + std::vector sizes; + // De-duplicated tags of the items that contributed segments here. + std::vector tags; + TransferDirection dir = TransferDirection::kLocal; + + // Set when this plan's LOCAL endpoint is unregistered and the engine must + // stage through its own bounce buffer. The local-side offsets above are then + // offsets INTO the bounce buffer, and `bounce_copies` says which user bytes + // to stage in (kPush) or out (kPull) around the wire. See + // MoriIoEngine::Submit for why such a plan completes inline. + struct BounceCopy { + void* host = nullptr; + uint64_t bounce_offset = 0; + uint64_t size = 0; + }; + bool uses_bounce = false; + std::vector bounce_copies; + + // Set by CompositeTransferEngine so Submit can route a plan back to the + // engine that produced it. Null means "whoever I am submitted to". + class TransferEngine* engine = nullptr; +}; + +// What Plan() could not express. `rejected_tags` is not an error path the +// caller can ignore: those items were dropped and their keys must be failed. +struct TransferPlanSet { + std::vector plans; + std::vector rejected_tags; +}; + +// One failed plan, mapped back to the keys that rode in it. Carries enough to +// reproduce the pre-refactor per-transfer error log. +struct TransferFailure { + std::vector tags; + uint32_t code = 0; + std::string message; + std::string endpoint; // remote engine key, for diagnosis +}; + +// Posted, not-yet-settled work. OWNS everything Submit referenced — mori-io's +// RDMA backend keeps a raw TransferStatus* into the handle's status vector, so +// the handle must outlive the post and must never move its statuses. +// +// The destructor is a safety net, not the normal path: on an early or +// exceptional destroy it drains, so a completion callback never writes freed +// memory. It performs no failure mapping — that is Wait's job. +class TransferHandle { + public: + virtual ~TransferHandle() = default; + + TransferHandle(const TransferHandle&) = delete; + TransferHandle& operator=(const TransferHandle&) = delete; + + // Block until every posted plan settles, appending one TransferFailure per + // plan that did not succeed. Never breaks early — every plan is waited, so + // no status is left live. Idempotent: a second call appends nothing. + virtual void Wait(std::vector* failures) = 0; + + protected: + TransferHandle() = default; +}; + +// --------------------------------------------------------------------------- +// TransferEngine +// --------------------------------------------------------------------------- +class TransferEngine : public MemoryRegistrar, public MetricSource { + public: + virtual const char* Name() const = 0; + + // Observability, same split as MediumBackend: an engine does NOT count its + // own plans, bytes or in-flight time — CompositeTransferEngine derives those + // from the dispatch it already performs, so a new engine added with + // AddEngine() is measured without a line of metrics code. SampleMetrics() + // (from MetricSource) is only for transport-internal state a dispatcher + // cannot see, e.g. bounce-pool pressure. engine= is added by the publisher. + + // Can this engine move bytes from `src` to `dst`? Engine selection is a + // function of the PAIR, never of either endpoint alone. + virtual bool CanHandle(const TransferRef& src, const TransferRef& dst) const = 0; + + // Group `items` into the units this engine wants to submit. Pure: no IO is + // issued and no engine state changes, so a caller may plan on one thread and + // submit on another. + virtual TransferPlanSet Plan(const std::vector& items) const = 0; + + // Post `plans`. Returns a handle to wait on, or nullptr if nothing was + // posted (in which case every tag in `plans` must be treated as failed). + // + // 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 bounce pool's 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. + virtual std::unique_ptr Submit(std::vector plans) = 0; + + // Where Transfer() reports the seconds it spent in each of its three steps. + // Planning grows with the number of items while the move grows with bytes, so + // a caller that batches thousands of small ranges needs to tell them apart to + // know which one it is paying for. Null members are inert. + struct StepTiming { + double* plan = nullptr; + double* submit = nullptr; + double* wait = nullptr; + }; + + // Plan + Submit + Wait, for callers with nothing to overlap. Returns false + // and fills `failed_tags` if anything was rejected or failed. + bool Transfer(const std::vector& items, std::vector* failed_tags, + const StepTiming* timing = nullptr); +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/types.h b/src/umbp/include/umbp/distributed/types.h index ebc6a5da3..97ddf7227 100644 --- a/src/umbp/include/umbp/distributed/types.h +++ b/src/umbp/include/umbp/distributed/types.h @@ -39,9 +39,49 @@ enum class TierType : int { SSD = 3, }; +// Upper bound on live storage backends per peer. +// +// Lives here, not on BackendRegistry, because both layers that key on a +// backend id need it and neither may depend on the other: medium_backend.h +// already includes transfer_engine.h, so the transfer layer cannot include the +// backend layer back. types.h is the shared vocabulary both sides already use. +// +// A backend id is the other half of a buffer address (see BufferMemoryDescBytes) +// and the registry hands them out densely from 0, so this is also the width of +// the reader's per-backend buffer shelves. +// +// Sized for the topology a logical tier policy expands into rather than for one +// backend per TierType: an 8-GPU node whose HBM is declared once expands to +// eight instances, and a policy that also names hot and warm DRAM plus a pair of +// SSDs lands at twelve. The shelves are vectors, so unused ids cost a pointer +// triple each. +inline constexpr uint32_t kMaxBackendsPerPeer = 16; + struct TierCapacity { uint64_t total_bytes = 0; uint64_t available_bytes = 0; + // Largest single value the tier can currently admit. Zero means + // available_bytes for backward-compatible single-backend reports. + uint64_t max_allocatable_bytes = 0; +}; + +struct LogicalTierCapacity { + TierType representative_tier = TierType::UNKNOWN; + TierCapacity capacity; + bool put_eligible = false; + // What the watermark comparison actually reads: the highest utilization among + // this tier's own backends. `capacity` aggregates instead, and for the entry + // tier it aggregates every tier reachable by offload, so neither of its + // ratios can be checked against this tier's watermarks. + double peak_member_utilization = 0.0; +}; + +struct TierTransitionMetrics { + uint64_t attempted = 0; + uint64_t succeeded = 0; + uint64_t failed = 0; + uint64_t offloaded_bytes = 0; + uint64_t promoted_bytes = 0; }; struct ExternalKvHitCountEntry { @@ -57,9 +97,11 @@ struct Location { std::string node_id; uint64_t size = 0; TierType tier = TierType::UNKNOWN; + std::string logical_tier; bool operator==(const Location& other) const { - return node_id == other.node_id && size == other.size && tier == other.tier; + return node_id == other.node_id && size == other.size && tier == other.tier && + logical_tier == other.logical_tier; } }; @@ -128,6 +170,19 @@ enum class EvictionOrder : int { kLeastRecentlyAccessed = 1, // oldest last_accessed_at first }; +// What a RegisterMemory call is asking for. +// +// Registering is not only about pinning: it is also what lets a range be +// described by its region's base instead of its own address, which is what +// collapses a batch of ranges into one transfer plan rather than one per range. +// A region that will only ever be copied locally still wants that, but has no +// use for an RDMA MR — and for an IPC-imported ROCm mapping the dmabuf fallback +// is unsafe — so the two halves are separable. +enum class MemoryRegistration : int { + kPinned = 0, // record it and export it to the IO engine + kLocalCopyOnly = 1, // record it only; no MR, no dmabuf export +}; + // Structured form of one (buffer_index, page_index) slot. Used by the // peer DRAM/HBM allocator to describe which page slot a write should // land in, and by ResolveKey responses to tell readers where to RDMA @@ -146,14 +201,20 @@ struct PageLocation { } }; -// One peer-side buffer's RDMA MemoryDesc bytes plus the buffer_index it -// belongs to. Returned by PeerDramAllocator and the peer service in -// AllocateSlot / ResolveKey / GetPeerInfo responses so the Client can -// hydrate its peer-side buffer_index -> MemoryDesc cache in a single -// batch. +// One peer-side buffer's RDMA MemoryDesc bytes plus the address it lives at. +// Returned by the backends and the peer service in AllocateSlot / ResolveKey / +// GetPeerInfo responses so the Client can hydrate its peer-side +// (backend_id, buffer_index) -> MemoryDesc cache in a single batch. +// +// BACKENDS LEAVE backend_id AT 0 AND MUST NOT SET IT. Every backend numbers +// its buffers from 0 and knows nothing of its peers; the peer service stamps +// the owning backend on the way out (see BufferMemoryDesc in umbp.proto for why +// buffer_index alone is not an address). Keeping the stamp at the wire +// boundary is what lets a new backend be correct by construction. struct BufferMemoryDescBytes { uint32_t buffer_index = 0; std::vector desc_bytes; + uint32_t backend_id = 0; }; // One mutation in a peer's owned-key set, shipped via Heartbeat. Mirrors @@ -166,6 +227,7 @@ struct KvEvent { std::string key; TierType tier = TierType::UNKNOWN; uint64_t size = 0; // ADD only; REMOVE leaves this 0 + std::string logical_tier; }; struct EventBundle { @@ -185,6 +247,7 @@ struct ClientRecord { std::chrono::system_clock::time_point last_heartbeat; std::chrono::system_clock::time_point registered_at; std::map tier_capacities; + std::map logical_tier_capacities; std::string peer_address; std::vector engine_desc_bytes; @@ -206,6 +269,7 @@ struct ClientRegistration { std::string node_id; std::string node_address; std::map tier_capacities; + std::map logical_tier_capacities; std::string peer_address; std::vector engine_desc_bytes; std::vector tags; @@ -236,6 +300,21 @@ inline const char* TierTypeName(TierType t) { } } +// Stable name used when legacy one-backend configuration is lowered into the +// named multi-backend model. +inline const char* DefaultBackendInstanceName(TierType t) { + switch (t) { + case TierType::HBM: + return "hbm"; + case TierType::DRAM: + return "dram"; + case TierType::SSD: + return "ssd"; + default: + return "unknown"; + } +} + inline const char* ClientStatusName(ClientStatus s) { switch (s) { case ClientStatus::ALIVE: diff --git a/src/umbp/include/umbp/local/standalone_client.h b/src/umbp/include/umbp/local/standalone_client.h deleted file mode 100644 index b2c8906b5..000000000 --- a/src/umbp/include/umbp/local/standalone_client.h +++ /dev/null @@ -1,105 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include - -#include "umbp/local/block_index/local_block_index.h" -#include "umbp/local/tiers/copy_pipeline.h" -#include "umbp/local/tiers/local_storage_manager.h" -#include "umbp/umbp_client.h" - -namespace mori::umbp { - -/// Standalone IUMBPClient implementation — purely local DRAM + SSD storage -/// with no networking or master coordination. -class StandaloneClient : public IUMBPClient { - public: - explicit StandaloneClient(const UMBPConfig& config = UMBPConfig{}); - ~StandaloneClient() override; - - // ---- IUMBPClient interface ---- - bool Put(const std::string& key, uintptr_t src, size_t size) override; - bool Get(const std::string& key, uintptr_t dst, size_t size) override; - bool Exists(const std::string& key) const override; - - std::vector BatchPut(const std::vector& keys, - const std::vector& srcs, - const std::vector& sizes) override; - std::vector BatchPutWithDepth(const std::vector& keys, - const std::vector& srcs, - const std::vector& sizes, - const std::vector& depths) override; - std::vector BatchGet(const std::vector& keys, - const std::vector& dsts, - const std::vector& sizes) override; - std::vector BatchExists(const std::vector& keys) const override; - size_t BatchExistsConsecutive(const std::vector& keys) const override; - - bool Clear() override; - bool Flush() override; - void Close() override; - bool IsDistributed() const override; - UMBPDeploymentMode GetDeploymentMode() const override { return UMBPDeploymentMode::Local; } - - bool ReportExternalKvBlocks(const std::vector& /*hashes*/, - TierType /*tier*/) override { - return true; - } - bool RevokeExternalKvBlocks(const std::vector& /*hashes*/, - TierType /*tier*/) override { - return true; - } - bool RevokeAllExternalKvBlocksAtTier(TierType /*tier*/) override { return true; } - std::vector MatchExternalKv(const std::vector& /*hashes*/, - bool /*count_as_hit*/ = false) override { - return {}; - } - std::vector GetExternalKvHitCounts( - const std::vector& /*hashes*/) override { - return {}; - } - - // ---- Extra methods (not in IUMBPClient, for C++ tests and debugging) ---- - - bool Put(const std::string& key, const void* data, size_t size); - bool Remove(const std::string& key); - - mori::umbp::LocalBlockIndex& Index(); - LocalStorageManager& Storage(); - - private: - static UMBPConfig NormalizeConfig(const UMBPConfig& config); - - UMBPConfig config_; - UMBPRole role_; - mori::umbp::LocalBlockIndex index_; - LocalStorageManager storage_; - std::unique_ptr copy_pipeline_; - bool closed_ = false; -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/local/tiers/dram_tier.h b/src/umbp/include/umbp/local/tiers/dram_tier.h deleted file mode 100644 index 8df9c0f31..000000000 --- a/src/umbp/include/umbp/local/tiers/dram_tier.h +++ /dev/null @@ -1,147 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/local/host_mem_allocator.h" -#include "umbp/local/tiers/tier_backend.h" - -namespace mori::umbp { - -// DRAM Tier: mmap pre-allocated large memory block with offset allocator -class DRAMTier : public TierBackend { - public: - DRAMTier(size_t capacity, bool use_shm, const std::string& shm_name, bool use_hugepages, - size_t hugepage_size, int numa_node, bool prefault); - - // Backward-compatible overloads (no hugepages). - explicit DRAMTier(size_t capacity) - : DRAMTier(capacity, false, "/umbp_dram", false, 2ULL * 1024 * 1024, -1, true) {} - DRAMTier(size_t capacity, bool use_shm, const std::string& shm_name) - : DRAMTier(capacity, use_shm, shm_name, false, 2ULL * 1024 * 1024, -1, true) {} - - ~DRAMTier() override; - - // Non-copyable - DRAMTier(const DRAMTier&) = delete; - DRAMTier& operator=(const DRAMTier&) = delete; - - // TierBackend interface - // Write: allocate slot in pre-allocated memory, memcpy data. - // Does NOT self-evict on space pressure — returns false if no space. - // Upper layer (LocalStorageManager) is responsible for demoting keys. - bool Write(const std::string& key, const void* data, size_t size) override; - bool ReadIntoPtr(const std::string& key, uintptr_t dst_ptr, size_t size) override; - // Multi-threaded batch read: parallel CopyBlock of many KV blocks to break - // the single-core memcpy ceiling on cold DRAM. Holds mu_ for the whole batch - // (consistent with this tier's single-mutex model: serializes against other - // reads/writes), while the per-block copies run in parallel within the batch. - std::vector ReadBatchIntoPtr(const std::vector& keys, - const std::vector& dst_ptrs, - const std::vector& sizes) override; - // Multi-threaded batch write: serial slot allocation (mutates free_list_) - // followed by parallel non-temporal CopyBlock of each payload into its slot. - // Mirrors ReadBatchIntoPtr to break the single-core memcpy ceiling on the - // L2->L3 backup (PUT) path. Does NOT self-evict — keys that don't fit are - // left false so the upper layer can demote and retry per-key. - std::vector BatchWrite(const std::vector& keys, - const std::vector& data_ptrs, - const std::vector& sizes) override; - bool Exists(const std::string& key) const override; - bool Evict(const std::string& key) override; - std::pair Capacity() const override; - void Clear() override; - - // Extended interface overrides - TierCapabilities Capabilities() const override; - std::vector Read(const std::string& key) override; - std::string GetLRUKey() const override; - std::vector GetLRUCandidates(size_t max_candidates) const override; - std::optional GetLocationId(const std::string& key) const override; - - // DRAM-specific: zero-copy read returning internal pointer. - // Only safe for in-process mmap'd memory. Caller must not hold - // the returned pointer across Evict/Write calls. - const void* ReadPtr(const std::string& key, size_t* out_size) override; - - // Accessors for distributed integration. - // Returns the mmap'd base address for RDMA registration. - void* GetBasePtr() const { return base_ptr_; } - // Returns the byte offset of a key's slot, or nullopt if not found. - std::optional GetSlotOffset(const std::string& key) const; - - private: - void* base_ptr_; // mmap base address - size_t capacity_; // usable capacity (= requested size) - size_t mapped_size_; // actual mapped size (>= capacity_ with hugepages) - size_t used_; - int shm_fd_; // shm_open fd (-1 for anonymous) - bool use_shm_; - std::string shm_name_; - HostBufferHandle host_buf_handle_; // owned handle for non-shm path - - // Simple offset allocator: key -> (offset, size) - struct SlotInfo { - size_t offset; - size_t size; - }; - std::unordered_map slots_; - - // LRU linked list - std::list lru_list_; - std::unordered_map::iterator> lru_map_; - - // Free block management (simple free list) - struct FreeBlock { - size_t offset; - size_t size; - }; - std::list free_list_; - - mutable std::mutex mu_; - - // Threads used by ReadBatchIntoPtr for parallel CopyBlock. Default 8, override - // via env UMBP_DRAM_READ_THREADS, capped to hardware concurrency. >1 breaks - // the single-core memcpy ceiling on cold DRAM. - int read_threads_ = 4; - - // Threads used by BatchWrite for parallel CopyBlock. Default 8, override via - // env UMBP_DRAM_WRITE_THREADS, capped to hardware concurrency. >1 breaks the - // single-core memcpy ceiling on the PUT (backup) path. - int write_threads_ = 4; - - size_t Allocate(size_t size); // Allocate from free_list_ - void Deallocate(size_t offset, size_t size); // Return to free_list_ - void EvictLRU(); // Evict least recently used - void TouchLRU(const std::string& key); // Update LRU position -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/local/tiers/local_storage_manager.h b/src/umbp/include/umbp/local/tiers/local_storage_manager.h deleted file mode 100644 index e11d59452..000000000 --- a/src/umbp/include/umbp/local/tiers/local_storage_manager.h +++ /dev/null @@ -1,181 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/common/config.h" -#include "umbp/local/block_index/local_block_index.h" -#include "umbp/local/tiers/tier_backend.h" - -namespace mori::umbp { - -class LocalStorageManager { - public: - struct TierLocationInfo { - std::string location_id; - size_t size = 0; - }; - - // Callback fired after a key changes tier or is fully evicted. - // to_tier == std::nullopt means the key was evicted (no destination tier). - // WARNING: Must NOT call back into LocalStorageManager (deadlock risk). - using TierChangeCallback = std::function to_tier, - std::optional new_location)>; - - // index may be nullptr if index updates are not needed (testing). - explicit LocalStorageManager(const UMBPConfig& config, - mori::umbp::LocalBlockIndex* index = nullptr); - ~LocalStorageManager(); - - // Write to the specified tier. - // When writing to DRAM and space is insufficient, automatically demotes - // LRU keys to the next-slower tier to make room. - bool Write(const std::string& key, const void* data, size_t size, - StorageTier tier = StorageTier::CPU_DRAM); - bool WriteFromPtr(const std::string& key, uintptr_t src, size_t size, - StorageTier tier = StorageTier::CPU_DRAM); - // depth == -1 means no metadata (degenerates to plain LRU eviction). - bool WriteFromPtrWithDepth(const std::string& key, uintptr_t src, size_t size, int depth, - StorageTier tier = StorageTier::CPU_DRAM); - // Batch counterpart of WriteFromPtr. When the target tier advertises - // batch_write, payloads are copied in parallel via TierBackend::BatchWrite; - // any key that doesn't fit falls back to a per-key Write() with LRU demotion. - std::vector WriteBatchFromPtr(const std::vector& keys, - const std::vector& srcs, - const std::vector& sizes, - StorageTier tier = StorageTier::CPU_DRAM); - - bool ReadIntoPtr(const std::string& key, uintptr_t dst, size_t size); - bool ReadIntoPtrNoPromote(const std::string& key, uintptr_t dst, size_t size); - std::vector ReadBatchIntoPtr(const std::vector& keys, - const std::vector& dst_ptrs, - const std::vector& sizes); - bool Exists(const std::string& key) const; - bool Evict(const std::string& key); - std::pair Capacity(StorageTier tier) const; - - bool Demote(const std::string& key); // Move to next-slower tier - bool Promote(const std::string& key); // Move to next-faster tier - bool CopyToSSD(const std::string& key); // Non-destructive DRAM→SSD copy - bool CopyToSSDBatch(const std::vector& keys); // Batched DRAM→SSD copy - void Clear(); - - // Install a callback invoked after MoveKey() or Evict() completes. - // Must be called before any concurrent Put/Get/Evict (typically in constructor). - void SetOnTierChange(TierChangeCallback cb); - - // Flush all tiers — ensures pending write-back data is durable. - bool Flush(); - - // Batch operations — delegates to tier backend's batch methods. - std::vector BatchWrite(const std::vector& keys, - const std::vector& data_ptrs, - const std::vector& sizes, - StorageTier tier = StorageTier::CPU_DRAM); - std::vector BatchReadIntoPtr(const std::vector& keys, - const std::vector& dst_ptrs, - const std::vector& sizes); - - // Access tiers generically - TierBackend* GetTier(StorageTier tier); - const TierBackend* GetTier(StorageTier tier) const; - - // Typed access (returns nullptr if tier not present or wrong type) - template - T* GetTierAs(StorageTier tier) { - return dynamic_cast(GetTier(tier)); - } - - private: - UMBPConfig config_; - UMBPRole role_; - mori::umbp::LocalBlockIndex* index_; // non-owning, may be nullptr - TierChangeCallback on_tier_change_; // set once before concurrent access - - // Ordered fastest-to-slowest: [{CPU_DRAM, dram}, {LOCAL_SSD, ssd}, ...] - struct TierEntry { - StorageTier id; - std::unique_ptr backend; - }; - std::vector tiers_; - - // ----------------------------------------------------------------------- - // Depth and group metadata (protected by depth_mu_) - // - // LocalStorageManager has no class-level mutex; tier backends use their own - // internal locks. depth_map_ and group_map_ are manager-owned state accessed - // from concurrent Put and eviction paths, so they require their own lock. - // ----------------------------------------------------------------------- - mutable std::shared_mutex depth_mu_; - std::unordered_map depth_map_; // key → chain depth - std::unordered_map> group_map_; // base_hash → keys - - // depth_map_ + group_map_ helpers (all acquire depth_mu_ internally) - void RecordDepth(const std::string& key, int depth); - int GetDepth(const std::string& key) const; // returns -1 if unknown - void RemoveDepthAndGroup(const std::string& key); - void RecordGroup(const std::string& key); - std::vector GetGroup(const std::string& key) const; - - // Strip _k/_v and rank suffixes to recover the SHA256 base hash. - // Returns key unchanged if parsing fails (safe fallback). - static std::string ExtractBaseHash(const std::string& key); - - // Select a victim from |tier| using the configured eviction policy. - // For "prefix_aware_lru": inspect up to eviction_candidate_window candidates, - // score by depth, pick the deepest (suffix block). Falls back to plain LRU - // if no metadata is available. - // Returns empty string if the tier is empty. - std::string SelectVictim(TierBackend* tier); - - // ----------------------------------------------------------------------- - // Existing helpers - // ----------------------------------------------------------------------- - TierBackend* FindTierHolding(const std::string& key); - const TierBackend* FindTierHolding(const std::string& key) const; - TierBackend* NextSlowerTier(StorageTier current); - TierBackend* NextFasterTier(StorageTier current); - bool MoveKey(const std::string& key, TierBackend* from, TierBackend* to); - bool DemoteLRUForSpace(TierBackend* tier); - bool InsertReadCacheNoWriteback(const std::string& key); - void UpsertIndexTier(const std::string& key, StorageTier tier, size_t size_hint); - static std::optional BuildTierLocationInfo(TierBackend* tier, - const std::string& key, size_t size); - - void MaybeAutoPromote(const std::string& key); - -#ifdef __linux__ - int EnsureProxyDaemon(const std::string& shm_name); - int SpawnProxyDaemon(const std::string& shm_name); -#endif -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/local/tiers/segment/segment_format.h b/src/umbp/include/umbp/local/tiers/segment/segment_format.h index fd2336646..a38efc297 100644 --- a/src/umbp/include/umbp/local/tiers/segment/segment_format.h +++ b/src/umbp/include/umbp/local/tiers/segment/segment_format.h @@ -28,8 +28,32 @@ namespace mori::umbp::segment { constexpr uint32_t kRecordMagic = 0x554D4250; // "UMBP" -constexpr uint16_t kRecordVersion = 1; +// v2 switched the record checksum from CRC-32/ISO-HDLC to CRC-32C so it can use +// the SSE4.2 crc32 instruction. The checksum values differ, so v1 records must +// not be verified with the v2 routine; the scanner drops them on version +// mismatch and the segment is refilled. +// v3 padded the layout to kRecordAlign so every record starts, and every value +// starts and ends, on an alignment boundary — the precondition for O_DIRECT. +// The padding is unconditional (not gated on direct_io) so one directory can be +// read back with either setting; only the reader's buffer alignment differs. +// v2 records are dropped by the scanner on version mismatch, same as v1. +constexpr uint16_t kRecordVersion = 3; constexpr uint16_t kFlagCommitted = 1; +// Set when the writer skipped checksumming (UMBP_SSD_VERIFY_CRC=0). Recorded +// per record so a store written with checksums off is still readable by a +// process with them on, instead of reporting every key as corrupt. +constexpr uint16_t kFlagNoCrc = 2; + +// On-disk alignment for record and value boundaries. Deliberately a fixed +// constant rather than the device's reported DIO requirement: the layout must +// not change meaning when the same directory moves between devices. 4096 +// covers every realistic requirement (512 and 4096 both divide it), and the +// O_DIRECT probe only has to confirm the device needs no more than this. +constexpr uint64_t kRecordAlign = 4096; + +constexpr uint64_t AlignUp(uint64_t v, uint64_t align = kRecordAlign) { + return (v + align - 1) & ~(align - 1); +} struct RecordHeader { uint32_t magic = 0; @@ -43,6 +67,26 @@ struct RecordHeader { }; static_assert(sizeof(RecordHeader) == 32, "unexpected padding in RecordHeader"); +// v3 record layout, all quantities multiples of kRecordAlign: +// +// [ prefix: RecordHeader | key | pad ][ value | pad ] +// ^ record_offset ^ value_offset +// +// A record therefore occupies PrefixBytes() + PaddedValueBytes() and the next +// record starts on an aligned boundary by construction. +constexpr uint64_t PrefixBytes(size_t key_len) { + return AlignUp(sizeof(RecordHeader) + static_cast(key_len)); +} +constexpr uint64_t PaddedValueBytes(size_t value_size) { + return AlignUp(static_cast(value_size)); +} +constexpr uint64_t RecordBytes(size_t key_len, size_t value_size) { + return PrefixBytes(key_len) + PaddedValueBytes(value_size); +} + +// Streaming CRC-32C update; hardware-accelerated where SSE4.2 is available and +// bit-identical on the portable fallback. Not complemented -- callers finish a +// digest by inverting the returned value, as ComputeRecordCrc32 does. uint32_t CrcUpdate(const void* data, size_t size, uint32_t crc = 0xFFFFFFFFu); uint32_t ComputeRecordCrc32(const std::string& key, const void* value, size_t value_size); std::string BuildFileName(uint64_t segment_id); diff --git a/src/umbp/include/umbp/local/tiers/segment/segment_index.h b/src/umbp/include/umbp/local/tiers/segment/segment_index.h index 3dcf8501c..92deb85a8 100644 --- a/src/umbp/include/umbp/local/tiers/segment/segment_index.h +++ b/src/umbp/include/umbp/local/tiers/segment/segment_index.h @@ -83,8 +83,10 @@ class Index { void TouchLRU(const std::string& key); void RemoveLRU(const std::string& key); + // Reserves aligned space for one record. `crc_valid` false records that the + // writer skipped checksumming, so readers know not to verify `crc32`. bool PrepareWrite(const std::string& key, size_t size, size_t key_len, uint32_t crc32, Meta* seg, - WriteReservation* out); + WriteReservation* out, bool crc_valid = true); void RollbackWrite(const WriteReservation& reservation); void RecordRecoveredEntry(const std::string& key, const KeyMeta& meta); bool EraseKey(const std::string& key); diff --git a/src/umbp/include/umbp/local/tiers/segment/segment_scanner.h b/src/umbp/include/umbp/local/tiers/segment/segment_scanner.h index cdbdc4c1f..9f97a3c7c 100644 --- a/src/umbp/include/umbp/local/tiers/segment/segment_scanner.h +++ b/src/umbp/include/umbp/local/tiers/segment/segment_scanner.h @@ -30,9 +30,13 @@ namespace mori::umbp::segment { class Scanner { public: + // `extra_open_flags` is OR'd into open() for segments the scanner discovers — + // O_DIRECT when the tier runs unbuffered, so the fds it installs in the index + // match the ones SSDTier opens itself. All of the scanner's own reads are + // issued in aligned kRecordAlign blocks so they are legal on such an fd. bool RefreshFromDisk(const std::string& dir, StorageIoDriver& io_driver, Index& index, - bool read_only_shared, bool force_full_rescan, - std::string* error_message) const; + bool read_only_shared, bool force_full_rescan, std::string* error_message, + int extra_open_flags = 0) const; }; } // namespace mori::umbp::segment diff --git a/src/umbp/include/umbp/local/tiers/segment/segment_types.h b/src/umbp/include/umbp/local/tiers/segment/segment_types.h index 193430d17..5eb0fbf59 100644 --- a/src/umbp/include/umbp/local/tiers/segment/segment_types.h +++ b/src/umbp/include/umbp/local/tiers/segment/segment_types.h @@ -36,6 +36,17 @@ struct KeyMeta { uint32_t size = 0; uint32_t crc32 = 0; uint64_t generation = 0; + // Bytes this record actually occupies in the segment, i.e. `size` plus the + // header, key and v3 alignment padding. Capacity is accounted in these rather + // than in `size` so a store of small values cannot overrun capacity_bytes on + // disk: padding rounds every record up to a kRecordAlign multiple, which for a + // 4 KiB value is 2x its size (and for the multi-MiB KV pages this tier + // normally holds, under 0.1%). + uint32_t disk_bytes = 0; + // False for records written with checksumming disabled (kFlagNoCrc). Readers + // must skip verification for those regardless of their own verify_crc setting, + // otherwise every such key reads back as corrupt. + bool crc_valid = true; }; struct Meta { diff --git a/src/umbp/include/umbp/local/tiers/segment/segment_writer.h b/src/umbp/include/umbp/local/tiers/segment/segment_writer.h index 8d371b50b..2031533e8 100644 --- a/src/umbp/include/umbp/local/tiers/segment/segment_writer.h +++ b/src/umbp/include/umbp/local/tiers/segment/segment_writer.h @@ -25,6 +25,7 @@ #include #include +#include "umbp/common/aligned_buffer.h" #include "umbp/local/tiers/segment/segment_format.h" #include "umbp/local/tiers/segment/segment_index.h" #include "umbp/storage/io/storage_io_driver.h" @@ -32,27 +33,56 @@ namespace mori::umbp::segment { struct PreparedRecord { - std::vector record; + // Aligned rather than std::vector: with O_DIRECT the source buffer of a + // write must sit on a kRecordAlign boundary, which the default allocator does + // not guarantee. Build() sizes it to RecordBytes(), already an alignment + // multiple, so the whole buffer can go to the device in one op. + AlignedBuffer record; WriteReservation reservation; + uint32_t crc32 = 0; // set by Build, consumed by Reserve + bool crc_valid = true; // false when checksumming was skipped }; class Writer { public: - explicit Writer(StorageIoDriver& io_driver) : io_driver_(io_driver) {} + // `compute_crc` false skips checksumming on this store's writes and stamps + // kFlagNoCrc, so readers know to skip verification for these records. + explicit Writer(StorageIoDriver& io_driver, bool compute_crc = true) + : io_driver_(io_driver), compute_crc_(compute_crc) {} - // Phase 1 (caller holds mu_): prepare record buffer and reserve index space. - // Returns false if capacity is exhausted. + // Phase 1a (NO lock held): checksum the record and assemble its on-disk bytes. + // Pure CPU over caller-owned memory — it touches no Index or Meta state, so it + // must run outside the tier mutex. This is the expensive half (a CRC and a + // full copy of the value); keeping it under the lock would block every + // concurrent reader on the same drive for the duration of the batch. + // `generation` is left zero and stamped by Reserve. + void Build(const std::string& key, const void* data, size_t size, PreparedRecord* out) const; + + // Phase 1b (caller holds mu_): reserve index/segment space for an already-built + // record and stamp the header field that depends on the reservation. Cheap: + // no checksum, no payload copy. Returns false if capacity is exhausted. + bool Reserve(const std::string& key, size_t size, Meta* segment_meta, Index& index, + PreparedRecord* out) const; + + // Build + Reserve in one call, for callers already holding mu_. Prefer the + // split form on any path where the lock is contended. bool Prepare(const std::string& key, const void* data, size_t size, Meta* segment_meta, Index& index, PreparedRecord* out) const; // Phase 2 (caller holds io_mu_ only): write the prepared record to disk. - IoStatus WriteRecord(int fd, const PreparedRecord& pr, bool should_sync) const; + // `sync_ms_out`, when non-null, receives the milliseconds spent in the + // durability flush alone, so a caller can bill the flush separately from the + // device write (see ssd_perf.h); it is left untouched when should_sync=false. + IoStatus WriteRecord(int fd, const PreparedRecord& pr, bool should_sync, + double* sync_ms_out = nullptr) const; // Phase 2 batch variant: write multiple prepared records to disk. - IoStatus WriteRecords(int fd, const std::vector& records, bool should_sync) const; + IoStatus WriteRecords(int fd, const std::vector& records, bool should_sync, + double* sync_ms_out = nullptr) const; private: StorageIoDriver& io_driver_; + bool compute_crc_ = true; }; } // namespace mori::umbp::segment diff --git a/src/umbp/include/umbp/local/tiers/sharded_ssd_tier.h b/src/umbp/include/umbp/local/tiers/sharded_ssd_tier.h new file mode 100644 index 000000000..46506a13b --- /dev/null +++ b/src/umbp/include/umbp/local/tiers/sharded_ssd_tier.h @@ -0,0 +1,133 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/local/tiers/tier_backend.h" + +namespace mori::umbp { + +// Fan one logical SSD tier across N drives (one sub-backend per mount point), so +// a batch is written to and read from every drive at once for their AGGREGATE +// bandwidth. +// +// Not SPDK RAID0: that stripes at the block layer and needs the whole set behind +// SPDK, uniform drives, and a device-wide rebuild to change membership. This +// sits one level up - works over ordinary mounts, tolerates differently-sized +// drives, and keeps each key whole on one drive so a single-key read costs one +// drive's IO instead of a strip-split. The two compose; a shard may be a RAID0. +// +// Placement is balanced, not hashed: each write goes to the shard with the most +// free space, degenerating to round-robin for uniform values on uniform drives. +// key_shard_ records the choice, so a re-put lands back on the same drive. +// +// Concurrency: key_shard_ is guarded by a shared_mutex; sub-backends are called +// WITHOUT it held, so shard IO runs fully in parallel. +class ShardedSsdTier : public TierBackend { + public: + // @p shards must be non-empty; every element non-null. @p io_threads caps + // the worker threads used by the batch paths (<= shard count); 0 means "one + // per shard", which is what saturates N drives. + explicit ShardedSsdTier(std::vector> shards, int io_threads = 0); + ~ShardedSsdTier() override; + + ShardedSsdTier(const ShardedSsdTier&) = delete; + ShardedSsdTier& operator=(const ShardedSsdTier&) = delete; + + size_t ShardCount() const { return shards_.size(); } + + // Test/observability: which shard currently holds @p key, or -1 if unknown. + int ShardOf(const std::string& key) const; + + // Test/observability: per-shard (used, total) — the balance check. + std::vector> PerShardCapacity() const; + + bool Write(const std::string& key, const void* data, size_t size) override; + bool ReadIntoPtr(const std::string& key, uintptr_t dst_ptr, size_t size) override; + bool Exists(const std::string& key) const override; + bool Evict(const std::string& key) override; + std::pair Capacity() const override; + void Clear() override; + std::vector Read(const std::string& key) override; + TierCapabilities Capabilities() const override; + std::string GetLRUKey() const override; + std::vector GetLRUCandidates(size_t max_candidates) const override; + std::optional GetLocationId(const std::string& key) const override; + std::optional LocateRecord(const std::string& key) const override; + bool Flush() override; + void SetColdRead(bool enable) override; + + // Parallel across shards: keys are bucketed by their (assigned or resolved) + // shard and each bucket is handed to its own worker. This is the whole point + // of the class — N drives move bytes concurrently. + std::vector BatchWrite(const std::vector& keys, + const std::vector& data_ptrs, + const std::vector& sizes) override; + std::vector BatchReadIntoPtr(const std::vector& keys, + const std::vector& dst_ptrs, + const std::vector& sizes) override; + bool WriteBatch(const std::vector& keys, const std::vector& data_ptrs, + const std::vector& sizes) override; + std::vector ReadBatchIntoPtr(const std::vector& keys, + const std::vector& dst_ptrs, + const std::vector& sizes) override; + + private: + // Shard for a new key of @p size: the one with the most projected free space + // that can still fit it, or -1 when none can. Caller holds map_mu_ exclusively + // and passes the projected reservation through free_hint_ (see ReserveLocked). + int PickShardLocked(size_t size) const; + + // Resolve an existing key's shard; -1 when the key is unknown here. Caller + // holds at least a shared lock on map_mu_. + int LookupLocked(const std::string& key) const; + + // Record/undo a placement decision. used_hint_ is a placement hint only: + // Capacity() always reports the sub-backends' own numbers, so drift costs + // balance, never correctness. + void ReserveLocked(int shard, size_t size); + void ReleaseLocked(int shard, size_t size); + + // Shard assignment + the byte count needed to release the placement hint on + // evict, so eviction never has to re-read a key just to learn its size. + struct KeyLoc { + uint32_t shard = 0; + uint64_t size = 0; + }; + + std::vector> shards_; + std::vector shard_total_; // immutable per-shard capacity + int io_threads_; + + mutable std::shared_mutex map_mu_; + std::unordered_map key_shard_; + std::vector used_hint_; // guarded by map_mu_ +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/local/block_index/local_block_index.h b/src/umbp/include/umbp/local/tiers/ssd_backend_factory.h similarity index 63% rename from src/umbp/include/umbp/local/block_index/local_block_index.h rename to src/umbp/include/umbp/local/tiers/ssd_backend_factory.h index 0e053b9cc..afa9d279d 100644 --- a/src/umbp/include/umbp/local/block_index/local_block_index.h +++ b/src/umbp/include/umbp/local/tiers/ssd_backend_factory.h @@ -21,31 +21,21 @@ // SOFTWARE. #pragma once -#include -#include -#include -#include -#include -#include +#include -#include "umbp/local/storage_tier.h" +#include "umbp/common/config.h" +#include "umbp/local/tiers/ssd_tier.h" +#include "umbp/local/tiers/tier_backend.h" namespace mori::umbp { -class LocalBlockIndex { - public: - // Index operations - bool MayExist(const std::string& key) const; - std::optional Lookup(const std::string& key) const; - void Insert(const std::string& key, const LocalLocation& loc); - std::optional Remove(const std::string& key); - bool UpdateTier(const std::string& key, StorageTier new_tier); - size_t Count() const; - void Clear(); - - private: - mutable std::shared_mutex mu_; - std::unordered_map index_; -}; +// Build the `file`-backend SSD tier: one storage directory yields a plain +// SSDTier, several (comma-separated, one mount per drive) a ShardedSsdTier. +// capacity_bytes is the total budget, split evenly across the directories. +// PeerSsdManager goes through here, so every SSD deployment gets identical +// multi-drive behaviour. SPDK backends are unaffected — +// they get multi-device support from SpdkEnv's RAID0 one level down. +std::unique_ptr MakeFileSsdBackend( + const UMBPSsdConfig& ssd_config, SSDAccessMode access_mode = SSDAccessMode::ReadWrite); } // namespace mori::umbp diff --git a/src/umbp/include/umbp/local/tiers/ssd_tier.h b/src/umbp/include/umbp/local/tiers/ssd_tier.h index e264ccb38..d88dcdb54 100644 --- a/src/umbp/include/umbp/local/tiers/ssd_tier.h +++ b/src/umbp/include/umbp/local/tiers/ssd_tier.h @@ -66,6 +66,20 @@ class SSDTier : public TierBackend { std::vector BatchReadIntoPtr(const std::vector& keys, const std::vector& dst_ptrs, const std::vector& sizes) override; + // Scattered sub-object reads: each key supplies several disjoint + // object-relative ranges, each with its own destination buffer. The v3 + // layout makes this cheap -- a value is one contiguous, alignment-anchored + // extent, so a range is just `value_offset + object_offset` handed to pread. + // + // CHECKSUMS ARE NOT VERIFIED HERE, whatever ssd.verify_crc says. A record's + // CRC covers the whole value, so a partial read has nothing to check itself + // against; the alternative (per-block checksums) is a record-format change. + // Whole-object reads keep verifying, so this narrows the guarantee only for + // callers that opted into ranged I/O. + std::vector ReadBatchRangesIntoPtr( + const std::vector& keys, const std::vector>& dst_ptrs, + const std::vector>& sizes, + const std::vector>& src_offsets) override; bool Exists(const std::string& key) const override; bool Evict(const std::string& key) override; std::pair Capacity() const override; @@ -76,28 +90,54 @@ class SSDTier : public TierBackend { std::vector GetLRUCandidates(size_t max_candidates) const override; const IoStatus& LastIoStatus() const { return last_io_status_; } std::optional GetLocationId(const std::string& key) const override; + std::optional LocateRecord(const std::string& key) const override; + // Enables O_DIRECT for segments opened from here on. Takes effect for + // segments opened after the call; existing fds keep their current mode, so + // the intended use is at construction (ssd.direct_io) rather than mid-run. + void SetColdRead(bool enable) override; + + // True when segment I/O is actually bypassing the page cache. False if + // direct I/O was never requested, or was requested and the probe failed. + bool direct_io_active() const { return direct_io_; } private: bool IsReadOnlyShared() const { return access_mode_ == SSDAccessMode::ReadOnlyShared; } bool ShouldSyncOnWrite() const { return ssd_config_.durability.mode == UMBPDurabilityMode::Strict; } + int SegmentOpenFlags() const; + // Empirical check that this directory's filesystem supports O_DIRECT at + // kRecordAlign: opens a probe file, writes and reads back one aligned block. + // Cheaper to trust than a reported capability, and it catches tmpfs/overlayfs + // (which reject the open) as well as devices needing coarser alignment. + bool ProbeDirectIo() const; bool EnsureActiveSegment(size_t need_bytes); bool RefreshFromDiskLocked(bool force_full_rescan); + // Startup repair: truncate each segment to the last record the scanner could + // parse. Removes records from an older kRecordVersion and torn tails, both of + // which would otherwise make every subsequently appended record unreadable + // after a restart. Owner-only; no-op for a fully-parsed segment. + void DropUnparsedTailsLocked(); bool OpenOrCreateSegmentLocked(uint64_t segment_id); bool RefreshFollowerLocked() const; segment::Meta* GetSegmentLocked(uint64_t segment_id); const segment::Meta* GetSegmentLocked(uint64_t segment_id) const; bool ReadRecordLocked(const std::string& key, void* dst, size_t size, uint32_t expected_crc, - uint64_t value_offset, int read_fd) const; + uint64_t value_offset, int read_fd, bool crc_valid) const; void RememberStatus(IoStatus status) const; + // Issue one value read, transparently bouncing through an aligned buffer when + // direct I/O is on and `dst`/`size` do not meet the alignment rules. + IoStatus ReadValueInto(int fd, void* dst, size_t size, uint64_t value_offset) const; + bool ShouldVerifyCrc(bool crc_valid) const { return ssd_config_.verify_crc && crc_valid; } + int TierThreads() const; std::string dir_; size_t capacity_; UMBPSsdConfig ssd_config_; SSDAccessMode access_mode_; + bool direct_io_ = false; mutable std::mutex mu_; mutable std::mutex io_mu_; diff --git a/src/umbp/include/umbp/local/tiers/tier_backend.h b/src/umbp/include/umbp/local/tiers/tier_backend.h index 17eac2684..c3b348051 100644 --- a/src/umbp/include/umbp/local/tiers/tier_backend.h +++ b/src/umbp/include/umbp/local/tiers/tier_backend.h @@ -36,11 +36,28 @@ struct TierCapabilities { bool zero_copy_read = false; bool batch_write = false; bool batch_read = false; + // Whether ReadBatchRangesIntoPtr is really implemented. The base class + // provides an all-false default, which is indistinguishable from "every key + // missed"; callers route on this flag instead so an unsupported tier falls + // through to another one rather than reporting a phantom miss. + bool ranged_read = false; }; -// Abstract base class for storage tier backends (DRAM, SSD, NVM, ...). +// Physical location of a value inside a segment file, for a zero-copy reader +// (GdsEngine) that DMAs the bytes itself with hipFileRead instead of asking the +// tier to copy them. Only the segment-file SSD tiers can answer; every other +// tier reports "no stable file location". +struct RecordLocation { + int fd = -1; // O_DIRECT segment descriptor + uint64_t value_offset = 0; // 4 KiB-aligned start of the value in the file + uint64_t readable_size = 0; // 4 KiB-aligned padded length safe to read + uint64_t value_size = 0; // logical value length (<= readable_size) + bool direct_io = false; // fd is O_DIRECT — a hipFile fastpath precondition +}; + +// Abstract base class for storage tier backends (SSD, SPDK, ...). // All tiers share a common interface for write/read/evict; tier-specific -// extensions (e.g., DRAMTier::ReadPtr) live in the concrete subclass. +// extensions (e.g. SSDTier::direct_io_active) live in the concrete subclass. class TierBackend { public: virtual ~TierBackend() = default; @@ -85,6 +102,10 @@ class TierBackend { virtual std::vector ReadBatchIntoPtr(const std::vector& keys, const std::vector& dst_ptrs, const std::vector& sizes); + virtual std::vector ReadBatchRangesIntoPtr( + const std::vector& keys, const std::vector>& dst_ptrs, + const std::vector>& sizes, + const std::vector>& src_offsets); // Return the LRU key, or empty string if empty. // Default returns "". Override in tiers with LRU tracking. @@ -99,6 +120,11 @@ class TierBackend { virtual std::vector BatchWrite(const std::vector& keys, const std::vector& data_ptrs, const std::vector& sizes); + virtual std::vector BatchWriteRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& src_ptrs, + const std::vector>& sizes, + const std::vector>& dst_offsets); // Batch read into user pointers. Default loops over ReadIntoPtr(). virtual std::vector BatchReadIntoPtr(const std::vector& keys, @@ -118,6 +144,15 @@ class TierBackend { // Callers prepend the store index before publishing to the Master. virtual std::optional GetLocationId(const std::string& key) const; + // Physical location of |key|'s value for a zero-copy reader that DMAs it + // itself (GdsEngine). Default: no stable file location — only the + // segment-file SSD tiers override it. Read-only: acquires no lock the caller + // must hold and does not touch LRU, so a Resolve can call it without + // perturbing recency. + virtual std::optional LocateRecord(const std::string& /*key*/) const { + return std::nullopt; + } + // Which StorageTier does this backend represent? StorageTier tier_id() const { return tier_id_; } diff --git a/src/umbp/include/umbp/standalone/standalone_process_client.h b/src/umbp/include/umbp/standalone/standalone_process_client.h index 0734b374f..98239357c 100644 --- a/src/umbp/include/umbp/standalone/standalone_process_client.h +++ b/src/umbp/include/umbp/standalone/standalone_process_client.h @@ -57,6 +57,15 @@ class StandaloneProcessClient : public IUMBPClient { std::vector BatchGet(const std::vector& keys, const std::vector& dsts, const std::vector& sizes) override; + std::vector BatchGetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& src_offsets) override; + std::vector BatchPutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& dst_offsets) override; std::vector BatchExists(const std::vector& keys) const override; size_t BatchExistsConsecutive(const std::vector& keys) const override; @@ -67,8 +76,15 @@ class StandaloneProcessClient : public IUMBPClient { UMBPDeploymentMode GetDeploymentMode() const override { return UMBPDeploymentMode::StandaloneProcess; } + UMBPDeploymentMode GetBackendMode() const override { return backend_mode_; } + bool SupportsRangedIO() const override { return supports_ranged_io_; } - bool RegisterMemory(uintptr_t ptr, size_t size) override; + // `mode` is accepted and ignored: this client owns nothing to pin. It forwards + // the region to the server, which decides how to declare it to the backend. + bool RegisterMemory(uintptr_t ptr, size_t size, + mori::io::MemoryLocationType loc = mori::io::MemoryLocationType::CPU, + int device = -1, + MemoryRegistration mode = MemoryRegistration::kPinned) override; void DeregisterMemory(uintptr_t ptr) override; bool ReportExternalKvBlocks(const std::vector& hashes, TierType tier) override; @@ -83,10 +99,18 @@ class StandaloneProcessClient : public IUMBPClient { // Resolves `ptr` against the registered host regions. On success writes the // region-relative `offset` and the matched region's worker VA `region_base`. bool OffsetFor(uintptr_t ptr, size_t size, uint64_t* offset, uint64_t* region_base) const; - bool WaitReady(int timeout_ms) const; + // The batch paths translate one pointer per RANGE -- a 1023-key ranged get + // carries roughly 5k of them -- so they take registration_mu_ once for the + // whole request and call this instead of a lock round trip per range. + bool OffsetForLocked(uintptr_t ptr, size_t size, uint64_t* offset, uint64_t* region_base) const; + // Not const: a successful Ping is where the server's backend mode and ranged + // capability become known, and they are cached on the client. + bool WaitReady(int timeout_ms); void MaybeAutoStart(); std::string ClientId(); void DeregisterMemoryLocked(); + bool RegisterDeviceMemory(uintptr_t ptr, size_t size, int device_id); + bool RegisterHostShmMemory(uintptr_t ptr, size_t size); UMBPConfig config_; UMBPStandaloneProcessConfig standalone_config_; @@ -98,15 +122,17 @@ class StandaloneProcessClient : public IUMBPClient { mutable std::shared_mutex op_mutex_; std::atomic closing_{false}; bool closed_ = false; + UMBPDeploymentMode backend_mode_ = UMBPDeploymentMode::StandaloneProcess; + bool supports_ranged_io_ = false; - // One registered host shared-memory region. A hybrid HiCache worker (e.g. - // DeepSeek-V4) registers several non-contiguous host pools per rank, so the - // client tracks N regions and resolves each data op to the one that owns its - // pointer. `base` is the worker VA base (== allocation->base), `size` its - // mapped size. + enum class RegionKind { kHostShm, kGpuIpc }; + + // A hybrid worker can register several non-contiguous host or GPU regions. + // `base` is the worker VA used as region_base in data requests. struct RegisteredRegion { uintptr_t base = 0; size_t size = 0; + RegionKind kind = RegionKind::kHostShm; }; mutable std::mutex registration_mu_; diff --git a/src/umbp/include/umbp/umbp_client.h b/src/umbp/include/umbp/umbp_client.h index fec4200d1..542bb5946 100644 --- a/src/umbp/include/umbp/umbp_client.h +++ b/src/umbp/include/umbp/umbp_client.h @@ -30,12 +30,20 @@ #include #include +#include "mori/io/enum.hpp" #include "umbp/common/config.h" #include "umbp/distributed/types.h" namespace mori::umbp { enum class UMBPDeploymentMode : int { + // No client reports Local any more: the local backend is gone and an + // embedded deployment is a DistributedClient with no master, which reports + // Distributed (config alone decides what a process is -- an embedded store + // is not a different implementation). The value is kept, and kept at 0, so + // an out-of-tree consumer that names it still compiles and so the + // STANDALONE_BACKEND_LOCAL wire value keeps its meaning; treat receiving it + // as "a server older than this change". Local = 0, StandaloneProcess = 1, Distributed = 2, @@ -44,8 +52,17 @@ enum class UMBPDeploymentMode : int { /// Abstract interface for UMBP storage clients. /// /// Two implementations exist behind this interface: -/// - StandaloneClient: purely local DRAM+SSD storage, no networking. -/// - DistributedClient: master-led global routing + RDMA data plane. +/// - DistributedClient: this node's medium, plus master-led global routing +/// and an RDMA data plane WHEN a master address is configured. With none, +/// the same class is an embedded, single-process store: same backends, +/// same transfer engine, no routing, no registration, no heartbeat. +/// - StandaloneProcessClient: gRPC forwarding to a standalone server, which +/// runs a DistributedClient of its own (with or without a master). +/// +/// There is deliberately no third, local-only implementation. "Local" is a +/// property of a deployment -- no master address -- not a separate client, so +/// embedded and distributed, standalone-process and in-process, are the same +/// code path with different configuration. /// /// Use CreateUMBPClient() to obtain the appropriate implementation based on /// UMBPConfig. All methods are zero-copy and pointer-based, designed for @@ -84,6 +101,21 @@ class IUMBPClient { const std::vector& dsts, const std::vector& sizes) = 0; + /// Read byte ranges of stored objects into scattered user buffers. Shape + /// errors fail the whole batch; data errors are reported per key. + virtual std::vector BatchGetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& src_offsets) = 0; + + /// Atomically publish objects assembled from scattered buffers. The ranges + /// for each key must exactly tile [0, object_sizes[i]). + virtual std::vector BatchPutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& dst_offsets) = 0; + virtual std::vector BatchExists(const std::vector& keys) const = 0; /// Returns the number of keys that exist consecutively from index 0. @@ -115,9 +147,18 @@ class IUMBPClient { /// Returns the concrete deployment mode behind this client. virtual UMBPDeploymentMode GetDeploymentMode() const { - return IsDistributed() ? UMBPDeploymentMode::Distributed : UMBPDeploymentMode::Local; + return IsDistributed() ? UMBPDeploymentMode::Distributed + : UMBPDeploymentMode::StandaloneProcess; } + /// Returns the backend behind a forwarding client. For direct clients this + /// equals GetDeploymentMode(); StandaloneProcessClient overrides it with the + /// mode reported by the server's Ping response. + virtual UMBPDeploymentMode GetBackendMode() const { return GetDeploymentMode(); } + + /// Whether this concrete backend implements ranged multi-buffer I/O. + virtual bool SupportsRangedIO() const { return false; } + // ---- Optional zero-copy hooks ---- // // Register a host buffer for zero-copy RDMA transfers. Standalone @@ -127,7 +168,19 @@ class IUMBPClient { // that *do* require registration MUST override; callers may treat a // `true` return as "registered or not-needed", and `false` as a hard // failure that must be surfaced. - virtual bool RegisterMemory(uintptr_t /*ptr*/, size_t /*size*/) { return true; } + /// `loc`/`device` describe the CALLER's allocation (CPU or a GPU ordinal), + /// not any storage medium — a GPU-resident src/dst (e.g. sglang HiCache + /// device-resident KV pages) must be registered as such so the transfer + /// layer picks HbmCopyEngine instead of assuming host memory. + /// `mode` picks whether the region is pinned and exported to the IO engine + /// or merely recorded for engine selection; see MemoryRegistration. A backend + /// with nothing to pin may ignore it. + virtual bool RegisterMemory( + uintptr_t /*ptr*/, size_t /*size*/, + mori::io::MemoryLocationType /*loc*/ = mori::io::MemoryLocationType::CPU, int /*device*/ = -1, + MemoryRegistration /*mode*/ = MemoryRegistration::kPinned) { + return true; + } virtual void DeregisterMemory(uintptr_t /*ptr*/) {} // ---- External KV Events (for unmanaged L1/L2 cache blocks) ---- @@ -181,9 +234,39 @@ class IUMBPClient { } }; +/// Fill in an embedded (in-process, no master, no networking) distributed +/// block for a config that names no deployment, and return the result. +/// +/// This is what makes "no configuration" mean "a private store in this +/// process" now that there is no separate local client: the medium comes from +/// the top-level dram/ssd blocks, identity is synthesized, and everything that +/// needs a peer -- master address, IO engine host, peer service port -- is left +/// empty so no socket is opened and no thread is started for a cluster this +/// process is not part of. +/// +/// Being embedded (in this process) and having a master are independent: a +/// `distributed` block does NOT mean "not embedded", it means "I have opinions +/// about some fields". Such a config keeps every value it sets and only has its +/// blanks filled -- identity, and `dram_page_size` when it is 0 -- so a caller +/// wanting one field need not hand-build the whole struct. The one value that +/// is adjusted rather than merely filled is that page size, and only for a pool +/// no peer shares (`peer_service_port == 0`): a page larger than the pool +/// yields zero pages and makes every put fail with NO_SPACE silently, so an +/// impossible size is shrunk to fit and warned about whether it was defaulted +/// or asked for. The full set of embedded choices above is applied only +/// when no `distributed` block was given at all, since only then is there no +/// choice of the caller's to overwrite. +/// +/// Two configs are returned unchanged: one naming `standalone_process` (a +/// different client entirely -- the pool lives in the server process) and one +/// whose `master_config.master_address` is non-empty (a cluster member, whose +/// settings belong to the deployment rather than to us). +UMBPConfig WithEmbeddedDefaults(const UMBPConfig& config); + /// Factory: creates the appropriate IUMBPClient implementation. -/// Creates StandaloneClient when config.distributed is not set, -/// DistributedClient when it is. +/// StandaloneProcessClient when config.standalone_process is set; otherwise a +/// DistributedClient over WithEmbeddedDefaults(config), which covers a cluster +/// member, a self-configured masterless node and a bare config alike. std::unique_ptr CreateUMBPClient(const UMBPConfig& config = UMBPConfig{}); } // namespace mori::umbp diff --git a/src/umbp/local/block_index/local_block_index.cpp b/src/umbp/local/block_index/local_block_index.cpp deleted file mode 100644 index 1f2faa743..000000000 --- a/src/umbp/local/block_index/local_block_index.cpp +++ /dev/null @@ -1,73 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/local/block_index/local_block_index.h" - -#include -#include - -namespace mori::umbp { - -bool LocalBlockIndex::MayExist(const std::string& key) const { - std::shared_lock lock(mu_); - return index_.count(key) > 0; -} - -std::optional LocalBlockIndex::Lookup(const std::string& key) const { - std::shared_lock lock(mu_); - auto it = index_.find(key); - if (it == index_.end()) return std::nullopt; - return it->second; -} - -void LocalBlockIndex::Insert(const std::string& key, const LocalLocation& loc) { - std::unique_lock lock(mu_); - index_[key] = loc; -} - -std::optional LocalBlockIndex::Remove(const std::string& key) { - std::unique_lock lock(mu_); - auto it = index_.find(key); - if (it == index_.end()) return std::nullopt; - auto loc = it->second; - index_.erase(it); - return loc; -} - -bool LocalBlockIndex::UpdateTier(const std::string& key, StorageTier new_tier) { - std::unique_lock lock(mu_); - auto it = index_.find(key); - if (it == index_.end()) return false; - it->second.tier = new_tier; - return true; -} - -size_t LocalBlockIndex::Count() const { - std::shared_lock lock(mu_); - return index_.size(); -} - -void LocalBlockIndex::Clear() { - std::unique_lock lock(mu_); - index_.clear(); -} - -} // namespace mori::umbp diff --git a/src/umbp/local/standalone_client.cpp b/src/umbp/local/standalone_client.cpp deleted file mode 100644 index a0d77d2f2..000000000 --- a/src/umbp/local/standalone_client.cpp +++ /dev/null @@ -1,344 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/local/standalone_client.h" - -#include -#include -#include - -#include "mori/utils/mori_log.hpp" -#include "umbp/local/tiers/dram_tier.h" - -namespace mori::umbp { - -UMBPConfig StandaloneClient::NormalizeConfig(const UMBPConfig& config) { - UMBPConfig normalized = config; - normalized.role = config.ResolveRole(); - normalized.follower_mode = (normalized.role == UMBPRole::SharedSSDFollower); - normalized.force_ssd_copy_on_write = (normalized.role == UMBPRole::SharedSSDLeader); - std::string error_message; - if (!normalized.Validate(&error_message)) { - throw std::runtime_error("invalid UMBP config: " + error_message); - } - return normalized; -} - -static const char* RoleStr(UMBPRole r) { - switch (r) { - case UMBPRole::SharedSSDLeader: - return "leader"; - case UMBPRole::SharedSSDFollower: - return "follower"; - default: - return "standalone"; - } -} - -StandaloneClient::StandaloneClient(const UMBPConfig& config) - : config_(NormalizeConfig(config)), role_(config_.ResolveRole()), storage_(config_, &index_) { - copy_pipeline_ = std::make_unique(storage_, config_.copy_pipeline, role_); - - MORI_UMBP_INFO( - "[StandaloneClient] initialized — role={} " - "dram={}MB hugepages={} hugepage_size={}MB numa_node={} " - "ssd={} ssd_backend={} ssd_capacity={}MB ssd_dir={} " - "eviction={} copy_pipeline_threads={} copy_pipeline_qdepth={}", - RoleStr(role_), config_.dram.capacity_bytes / (1024 * 1024), config_.dram.use_hugepages, - config_.dram.hugepage_size / (1024 * 1024), config_.dram.numa_node, config_.ssd.enabled, - config_.ssd.ssd_backend, config_.ssd.capacity_bytes / (1024 * 1024), config_.ssd.storage_dir, - config_.eviction.policy, config_.copy_pipeline.worker_threads, - config_.copy_pipeline.queue_depth); -} - -StandaloneClient::~StandaloneClient() { Close(); } - -// --------------------------------------------------------------------------- -// Core KV Operations (IUMBPClient interface) -// --------------------------------------------------------------------------- - -bool StandaloneClient::Put(const std::string& key, uintptr_t src, size_t size) { - if (role_ == UMBPRole::SharedSSDFollower) return false; - if (index_.MayExist(key)) return true; - - if (!storage_.WriteFromPtr(key, src, size)) return false; - - index_.Insert(key, {StorageTier::CPU_DRAM, 0, size}); - copy_pipeline_->MaybeCopyToSharedSSD(key); - return true; -} - -bool StandaloneClient::Get(const std::string& key, uintptr_t dst, size_t size) { - bool in_index = index_.MayExist(key); - - if (!in_index && role_ != UMBPRole::SharedSSDFollower) { - return false; - } - - bool ok = storage_.ReadIntoPtr(key, dst, size); - - if (role_ == UMBPRole::SharedSSDFollower) { - if (ok) { - StorageTier tier = StorageTier::LOCAL_SSD; - auto* dram = storage_.GetTier(StorageTier::CPU_DRAM); - if (dram && dram->Exists(key)) { - tier = StorageTier::CPU_DRAM; - } - if (!index_.UpdateTier(key, tier)) { - index_.Insert(key, {tier, 0, size}); - } - } else { - if (in_index && !storage_.Exists(key)) { - index_.Remove(key); - } - } - } else if (!ok && in_index && !storage_.Exists(key)) { - index_.Remove(key); - } - - if (!ok && role_ == UMBPRole::SharedSSDFollower && !in_index) { - if (!storage_.Exists(key)) { - index_.Remove(key); - } - } - - return ok; -} - -bool StandaloneClient::Exists(const std::string& key) const { - if (role_ == UMBPRole::SharedSSDFollower) { - return storage_.Exists(key); - } - return index_.MayExist(key); -} - -// --------------------------------------------------------------------------- -// Batch Operations -// --------------------------------------------------------------------------- - -std::vector StandaloneClient::BatchPut(const std::vector& keys, - const std::vector& srcs, - const std::vector& sizes) { - const size_t n = keys.size(); - std::vector results(n, false); - // Followers never write; keep results all-false (and skip leader SSD copy). - if (role_ == UMBPRole::SharedSSDFollower) return results; - - // Collect keys needing an actual write (skip ones already present), then issue - // a single batched write so the payload copies run in parallel inside the tier - // (mirrors BatchGet's single ReadBatchIntoPtr dispatch). - std::vector wmap; - std::vector wkeys; - std::vector wsrcs; - std::vector wsizes; - wmap.reserve(n); - wkeys.reserve(n); - wsrcs.reserve(n); - wsizes.reserve(n); - for (size_t i = 0; i < n; ++i) { - if (index_.MayExist(keys[i])) { - results[i] = true; - continue; - } - wmap.push_back(i); - wkeys.push_back(keys[i]); - wsrcs.push_back(srcs[i]); - wsizes.push_back(sizes[i]); - } - - if (!wkeys.empty()) { - auto wres = storage_.WriteBatchFromPtr(wkeys, wsrcs, wsizes); - for (size_t j = 0; j < wkeys.size(); ++j) { - if (!wres[j]) continue; - size_t i = wmap[j]; - index_.Insert(keys[i], {StorageTier::CPU_DRAM, 0, sizes[i]}); - results[i] = true; - } - } - - if (role_ == UMBPRole::SharedSSDLeader) { - std::vector ssd_keys; - for (size_t i = 0; i < n; ++i) { - if (results[i]) ssd_keys.push_back(keys[i]); - } - copy_pipeline_->MaybeBatchCopyToSharedSSD(ssd_keys); - } - return results; -} - -std::vector StandaloneClient::BatchPutWithDepth(const std::vector& keys, - const std::vector& srcs, - const std::vector& sizes, - const std::vector& depths) { - std::vector results(keys.size(), false); - - for (size_t i = 0; i < keys.size(); ++i) { - if (role_ == UMBPRole::SharedSSDFollower) continue; - if (index_.MayExist(keys[i])) { - results[i] = true; - continue; - } - int depth = (i < depths.size()) ? depths[i] : -1; - if (!storage_.WriteFromPtrWithDepth(keys[i], srcs[i], sizes[i], depth)) continue; - index_.Insert(keys[i], {StorageTier::CPU_DRAM, 0, sizes[i]}); - results[i] = true; - } - - if (role_ == UMBPRole::SharedSSDLeader) { - std::vector ssd_keys; - for (size_t i = 0; i < keys.size(); ++i) { - if (results[i]) ssd_keys.push_back(keys[i]); - } - copy_pipeline_->MaybeBatchCopyToSharedSSD(ssd_keys); - } - return results; -} - -std::vector StandaloneClient::BatchGet(const std::vector& keys, - const std::vector& dsts, - const std::vector& sizes) { - std::vector results(keys.size(), false); - if (keys.empty()) return results; - - std::vector read_indices; - std::vector was_in_index(keys.size(), false); - read_indices.reserve(keys.size()); - - for (size_t i = 0; i < keys.size(); ++i) { - was_in_index[i] = index_.MayExist(keys[i]); - if (!was_in_index[i] && role_ != UMBPRole::SharedSSDFollower) { - continue; - } - read_indices.push_back(i); - } - - if (read_indices.empty()) return results; - - std::vector batch_keys; - std::vector batch_ptrs; - std::vector batch_sizes; - batch_keys.reserve(read_indices.size()); - batch_ptrs.reserve(read_indices.size()); - batch_sizes.reserve(read_indices.size()); - for (size_t idx : read_indices) { - batch_keys.push_back(keys[idx]); - batch_ptrs.push_back(dsts[idx]); - batch_sizes.push_back(sizes[idx]); - } - - auto batch_results = storage_.ReadBatchIntoPtr(batch_keys, batch_ptrs, batch_sizes); - - for (size_t j = 0; j < read_indices.size(); ++j) { - size_t i = read_indices[j]; - bool local_hit = batch_results[j]; - - if (role_ == UMBPRole::SharedSSDFollower) { - if (local_hit) { - StorageTier tier = StorageTier::LOCAL_SSD; - auto* dram = storage_.GetTier(StorageTier::CPU_DRAM); - if (dram && dram->Exists(keys[i])) { - tier = StorageTier::CPU_DRAM; - } - if (!index_.UpdateTier(keys[i], tier)) { - index_.Insert(keys[i], {tier, 0, sizes[i]}); - } - } else if (!storage_.Exists(keys[i])) { - index_.Remove(keys[i]); - } - } else if (!local_hit && was_in_index[i] && !storage_.Exists(keys[i])) { - index_.Remove(keys[i]); - } - - results[i] = local_hit; - } - - return results; -} - -std::vector StandaloneClient::BatchExists(const std::vector& keys) const { - std::vector results(keys.size(), false); - for (size_t i = 0; i < keys.size(); ++i) { - results[i] = Exists(keys[i]); - } - return results; -} - -size_t StandaloneClient::BatchExistsConsecutive(const std::vector& keys) const { - for (size_t i = 0; i < keys.size(); ++i) { - if (!Exists(keys[i])) return i; - } - return keys.size(); -} - -// --------------------------------------------------------------------------- -// Lifecycle -// --------------------------------------------------------------------------- - -bool StandaloneClient::Clear() { - index_.Clear(); - storage_.Clear(); - return true; -} - -bool StandaloneClient::Flush() { - bool ok = true; - if (copy_pipeline_) { - ok = copy_pipeline_->Drain(std::chrono::seconds(30)); - } - return storage_.Flush() && ok; -} - -void StandaloneClient::Close() { - if (closed_) return; - closed_ = true; - // CopyPipeline and storage are cleaned up by their own destructors, - // which run after this in the member destruction order. -} - -bool StandaloneClient::IsDistributed() const { return false; } - -// --------------------------------------------------------------------------- -// Extra methods (not in IUMBPClient) -// --------------------------------------------------------------------------- - -bool StandaloneClient::Put(const std::string& key, const void* data, size_t size) { - if (role_ == UMBPRole::SharedSSDFollower) return false; - if (index_.MayExist(key)) return true; - - if (!storage_.Write(key, data, size)) return false; - - index_.Insert(key, {StorageTier::CPU_DRAM, 0, size}); - copy_pipeline_->MaybeCopyToSharedSSD(key); - return true; -} - -bool StandaloneClient::Remove(const std::string& key) { - auto loc = index_.Remove(key); - if (!loc) return false; - storage_.Evict(key); - return true; -} - -mori::umbp::LocalBlockIndex& StandaloneClient::Index() { return index_; } - -LocalStorageManager& StandaloneClient::Storage() { return storage_; } - -} // namespace mori::umbp diff --git a/src/umbp/local/tiers/copy_pipeline.cpp b/src/umbp/local/tiers/copy_pipeline.cpp deleted file mode 100644 index 620ce3aaa..000000000 --- a/src/umbp/local/tiers/copy_pipeline.cpp +++ /dev/null @@ -1,139 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/local/tiers/copy_pipeline.h" - -#include - -namespace mori::umbp { - -CopyPipeline::CopyPipeline(LocalStorageManager& storage, const UMBPCopyPipelineConfig& config, - UMBPRole role) - : storage_(storage), config_(config), role_(role) { - if (role_ == UMBPRole::SharedSSDLeader && config_.async_enabled) { - const size_t n_workers = std::max(1, config_.worker_threads); - copy_workers_.reserve(n_workers); - for (size_t i = 0; i < n_workers; ++i) { - copy_workers_.emplace_back(&CopyPipeline::CopyWorkerLoop, this); - } - } -} - -CopyPipeline::~CopyPipeline() { - if (copy_workers_.empty()) return; - { - std::lock_guard lock(copy_mu_); - stop_copy_worker_.store(true); - } - copy_cv_.notify_all(); - for (auto& worker : copy_workers_) { - if (worker.joinable()) worker.join(); - } -} - -bool CopyPipeline::MaybeCopyToSharedSSD(const std::string& key) { - if (role_ != UMBPRole::SharedSSDLeader) return true; - if (!config_.async_enabled) { - storage_.CopyToSSD(key); - return true; - } - if (!EnqueueCopyToSSD(key)) { - storage_.CopyToSSD(key); - } - return true; -} - -void CopyPipeline::MaybeBatchCopyToSharedSSD(const std::vector& keys) { - if (role_ != UMBPRole::SharedSSDLeader || keys.empty()) return; - - if (!config_.async_enabled) { - storage_.CopyToSSDBatch(keys); - return; - } - - size_t enqueued = EnqueueCopyToSSDBatch(keys); - if (enqueued < keys.size()) { - std::vector overflow(keys.begin() + enqueued, keys.end()); - storage_.CopyToSSDBatch(overflow); - } -} - -bool CopyPipeline::Drain(std::chrono::milliseconds timeout) { - if (copy_workers_.empty()) return true; - std::unique_lock lock(copy_mu_); - return drain_cv_.wait_for(lock, timeout, - [&]() { return copy_queue_.empty() && in_flight_copies_ == 0; }); -} - -bool CopyPipeline::EnqueueCopyToSSD(const std::string& key) { - std::unique_lock lock(copy_mu_); - if (copy_queue_.size() >= config_.queue_depth) return false; - copy_queue_.push_back({key}); - lock.unlock(); - copy_cv_.notify_one(); - return true; -} - -size_t CopyPipeline::EnqueueCopyToSSDBatch(const std::vector& keys) { - std::unique_lock lock(copy_mu_); - size_t remaining = - (config_.queue_depth > copy_queue_.size()) ? config_.queue_depth - copy_queue_.size() : 0; - size_t to_enqueue = std::min(keys.size(), remaining); - for (size_t i = 0; i < to_enqueue; ++i) { - copy_queue_.push_back({keys[i]}); - } - lock.unlock(); - if (to_enqueue > 0) copy_cv_.notify_all(); - return to_enqueue; -} - -void CopyPipeline::CopyWorkerLoop() { - const size_t batch_max = std::max(1, config_.batch_max_ops); - while (true) { - std::vector batch; - { - std::unique_lock lock(copy_mu_); - copy_cv_.wait(lock, [&]() { return stop_copy_worker_.load() || !copy_queue_.empty(); }); - if (stop_copy_worker_.load() && copy_queue_.empty()) return; - size_t n = std::min(batch_max, copy_queue_.size()); - batch.reserve(n); - for (size_t i = 0; i < n; ++i) { - batch.push_back(std::move(copy_queue_.front().key)); - copy_queue_.pop_front(); - } - ++in_flight_copies_; - } - - if (batch.size() == 1) { - storage_.CopyToSSD(batch[0]); - } else { - storage_.CopyToSSDBatch(batch); - } - - { - std::lock_guard lock(copy_mu_); - if (in_flight_copies_ > 0) --in_flight_copies_; - if (copy_queue_.empty() && in_flight_copies_ == 0) drain_cv_.notify_all(); - } - } -} - -} // namespace mori::umbp diff --git a/src/umbp/local/tiers/device_copy_run.h b/src/umbp/local/tiers/device_copy_run.h new file mode 100644 index 000000000..3a15d7999 --- /dev/null +++ b/src/umbp/local/tiers/device_copy_run.h @@ -0,0 +1,145 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +// Source-private run planner shared only by DRAMTier and its unit test. + +#include +#include +#include +#include +#include +#include + +namespace mori::umbp::detail { + +enum class DeviceCopyRunKind { kSingle, kLinear, kPitched }; + +struct DeviceCopyRun { + DeviceCopyRunKind kind{DeviceCopyRunKind::kSingle}; + size_t end{0}; // one past the last job + size_t bytes{0}; + size_t width{0}; + size_t spitch{0}; + size_t dpitch{0}; +}; + +// Off by default, and it should stay that way. Reading one layer out of +// page-sized objects gives hipMemcpy2DAsync a pitch/width ratio in the tens, +// and at that ratio it is bimodal: ~0.35 us per fragment once the source pages +// have been through it, ~194 us the first time. A live tier recycles slots +// constantly, so it only ever pays the cold price -- measured 133x slower than +// per-fragment 1D copies end to end. The gather kernel covers this shape +// instead. Set UMBP_DRAM_PITCHED_COPY=1 to measure the pitched path again. +inline bool PitchedCopyEnabled() { + static const bool enabled = [] { + const char* value = std::getenv("UMBP_DRAM_PITCHED_COPY"); + if (value == nullptr) return false; + const std::string text(value); + return text == "1" || text == "on" || text == "true"; + }(); + return enabled; +} + +inline uintptr_t PointerAddress(const void* ptr) { return reinterpret_cast(ptr); } + +inline bool Adjacent(uintptr_t base, size_t size, uintptr_t next) { + return size <= std::numeric_limits::max() - base && base + size == next; +} + +// Job must expose `src`, `dst`, and `size`. `ordered` must be sorted by source +// address. Kept as a small pure planner so path selection can be regression- +// tested without issuing HIP work. +// +// `allow_pitched` is false for the gather-kernel path, which wants only the +// flat merges: a pitched run describes several fragments in one submission, +// which is the copy engine's vocabulary, not the kernel's. Callers pass +// PitchedCopyEnabled() rather than having it read here, so the planner stays a +// pure function of its arguments and can be tested without the environment. +template +DeviceCopyRun FindDeviceCopyRun(const std::vector& ordered, size_t begin, + bool allow_pitched = true) { + const Job* first = ordered[begin]; + DeviceCopyRun run{DeviceCopyRunKind::kSingle, begin + 1, first->size, first->size, 0, 0}; + + // Preserve the existing fast path exactly: adjacent jobs may have different + // sizes and still form one flat memcpy. + size_t linear_end = begin + 1; + size_t linear_bytes = first->size; + while (linear_end < ordered.size()) { + const Job* prev = ordered[linear_end - 1]; + const Job* cur = ordered[linear_end]; + if (!Adjacent(PointerAddress(prev->src), prev->size, PointerAddress(cur->src)) || + !Adjacent(PointerAddress(prev->dst), prev->size, PointerAddress(cur->dst)) || + cur->size > std::numeric_limits::max() - linear_bytes) { + break; + } + linear_bytes += cur->size; + ++linear_end; + } + if (linear_end > begin + 1) { + run.kind = DeviceCopyRunKind::kLinear; + run.end = linear_end; + run.bytes = linear_bytes; + return run; + } + + if (!allow_pitched || begin + 1 >= ordered.size() || first->size == 0 || + ordered[begin + 1]->size != first->size) { + return run; + } + const uintptr_t src0 = PointerAddress(first->src); + const uintptr_t src1 = PointerAddress(ordered[begin + 1]->src); + const uintptr_t dst0 = PointerAddress(first->dst); + const uintptr_t dst1 = PointerAddress(ordered[begin + 1]->dst); + if (src1 <= src0 || dst1 <= dst0) return run; + + const size_t spitch = src1 - src0; + const size_t dpitch = dst1 - dst0; + if (spitch < first->size || dpitch < first->size) return run; + + size_t pitched_end = begin + 2; + while (pitched_end < ordered.size()) { + const Job* prev = ordered[pitched_end - 1]; + const Job* cur = ordered[pitched_end]; + const uintptr_t prev_src = PointerAddress(prev->src); + const uintptr_t cur_src = PointerAddress(cur->src); + const uintptr_t prev_dst = PointerAddress(prev->dst); + const uintptr_t cur_dst = PointerAddress(cur->dst); + if (cur->size != first->size || cur_src <= prev_src || cur_dst <= prev_dst || + cur_src - prev_src != spitch || cur_dst - prev_dst != dpitch) { + break; + } + ++pitched_end; + } + + const size_t rows = pitched_end - begin; + run.kind = DeviceCopyRunKind::kPitched; + run.end = pitched_end; + run.bytes = first->size <= std::numeric_limits::max() / rows ? first->size * rows : 0; + run.width = first->size; + run.spitch = spitch; + run.dpitch = dpitch; + return run; +} + +} // namespace mori::umbp::detail diff --git a/src/umbp/include/umbp/common/error_code.h b/src/umbp/local/tiers/device_gather.h similarity index 81% rename from src/umbp/include/umbp/common/error_code.h rename to src/umbp/local/tiers/device_gather.h index 328574384..63c32b052 100644 --- a/src/umbp/include/umbp/common/error_code.h +++ b/src/umbp/local/tiers/device_gather.h @@ -19,19 +19,8 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// MIT License #pragma once -namespace umbp { - -enum class ErrorCode : int { - OK = 0, - SPDK_INIT_FAIL = -1, - SPDK_IO_FAIL = -2, - DMA_ALLOC_FAIL = -3, - ALLOC_FAIL = -4, - INVALID_ARG = -5, -}; - -} // namespace umbp +// Backward-compatible source-tree include. New shared users include the common +// header directly. +#include "umbp/common/device_gather.h" diff --git a/src/umbp/local/tiers/device_gather.hip b/src/umbp/local/tiers/device_gather.hip new file mode 100644 index 000000000..5ba665d43 --- /dev/null +++ b/src/umbp/local/tiers/device_gather.hip @@ -0,0 +1,184 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/common/device_gather.h" + +namespace mori::umbp { +namespace { + +constexpr unsigned kThreadsPerBlock = 256; + +// Enough parallelism to keep the link busy on the largest fragment without +// launching blocks that immediately exit on the smallest. Fragments below one +// block's worth of work still get one block each. +constexpr unsigned kMaxBlocksPerFragment = 8; + +constexpr size_t kVectorBytes = 16; + +// A block copies its slice of one fragment. Fragments are independent, so the +// grid is flat: blockIdx.x / blocks_per_fragment picks the fragment and the +// remainder picks the slice. +__global__ void GatherFragmentsKernel(const DeviceGatherFragment* __restrict__ fragments, + unsigned blocks_per_fragment) { + const DeviceGatherFragment fragment = fragments[blockIdx.x / blocks_per_fragment]; + const char* __restrict__ src = static_cast(fragment.src); + char* __restrict__ dst = static_cast(fragment.dst); + + const size_t slice = blockIdx.x % blocks_per_fragment; + const size_t index = slice * blockDim.x + threadIdx.x; + const size_t stride = static_cast(blocks_per_fragment) * blockDim.x; + + // KV fragments are page-aligned in practice; the byte loop is correctness + // insurance for pools whose element size makes them odd. + const uintptr_t mask = reinterpret_cast(src) | reinterpret_cast(dst) | + static_cast(fragment.bytes); + if ((mask & (kVectorBytes - 1)) == 0) { + const uint4* __restrict__ src_vector = reinterpret_cast(src); + uint4* __restrict__ dst_vector = reinterpret_cast(dst); + const size_t count = fragment.bytes / kVectorBytes; + for (size_t i = index; i < count; i += stride) dst_vector[i] = src_vector[i]; + return; + } + for (size_t i = index; i < fragment.bytes; i += stride) dst[i] = src[i]; +} + +// Per-thread, per-device descriptor staging. +// +// Per-thread because two threads serving the same device share one stream (see +// DeviceStream), so a shared staging buffer would be a data race even though +// the stream itself is ordered. Never freed, for the same reason DeviceStream +// never destroys its streams: releasing HIP resources from a thread or static +// destructor races with runtime teardown. The count is bounded by the server's +// worker threads times the devices they touch. +struct DescriptorStaging { + DeviceGatherFragment* host = nullptr; + DeviceGatherFragment* device = nullptr; + size_t capacity = 0; +}; + +DescriptorStaging* StagingFor(int device_id, size_t count) { + static thread_local std::map* staging_by_device = + new std::map(); + + DescriptorStaging& staging = (*staging_by_device)[device_id]; + if (staging.capacity >= count) return &staging; + + // Grow with headroom so a slowly increasing batch size does not reallocate + // on every call. + const size_t capacity = std::max(count * 2, 256); + const size_t bytes = capacity * sizeof(DeviceGatherFragment); + DeviceGatherFragment* host = nullptr; + DeviceGatherFragment* device = nullptr; + hipError_t status = hipHostMalloc(reinterpret_cast(&host), bytes, hipHostMallocDefault); + if (status != hipSuccess) { + MORI_UMBP_ERROR("[DeviceGather] hipHostMalloc for {} gather descriptors failed on device {}: {}", + capacity, device_id, hipGetErrorString(status)); + (void)hipGetLastError(); + return nullptr; + } + status = hipMalloc(reinterpret_cast(&device), bytes); + if (status != hipSuccess) { + MORI_UMBP_ERROR("[DeviceGather] hipMalloc for {} gather descriptors failed on device {}: {}", + capacity, device_id, hipGetErrorString(status)); + (void)hipGetLastError(); + (void)hipHostFree(host); + return nullptr; + } + + if (staging.host != nullptr) (void)hipHostFree(staging.host); + if (staging.device != nullptr) (void)hipFree(staging.device); + staging.host = host; + staging.device = device; + staging.capacity = capacity; + return &staging; +} + +std::atomic& LaunchCounter() { + static std::atomic launches{0}; + return launches; +} + +} // namespace + +uint64_t DeviceGatherLaunchCount() { return LaunchCounter().load(std::memory_order_relaxed); } + +bool DeviceGatherEnabled() { + static const bool enabled = [] { + const char* value = std::getenv("UMBP_DRAM_GATHER_KERNEL"); + if (value == nullptr) return true; + const std::string text(value); + return text != "0" && text != "off" && text != "false"; + }(); + return enabled; +} + +bool LaunchDeviceGather(const DeviceGatherFragment* fragments, size_t count, int device_id, + hipStream_t stream) { + if (count == 0) return true; + + size_t max_bytes = 0; + for (size_t i = 0; i < count; ++i) max_bytes = std::max(max_bytes, fragments[i].bytes); + if (max_bytes == 0) return false; + + const size_t chunks = (max_bytes + kVectorBytes - 1) / kVectorBytes; + const unsigned blocks_per_fragment = static_cast(std::clamp( + (chunks + kThreadsPerBlock - 1) / kThreadsPerBlock, 1, kMaxBlocksPerFragment)); + if (count > std::numeric_limits::max() / blocks_per_fragment) return false; + + DescriptorStaging* staging = StagingFor(device_id, count); + if (staging == nullptr) return false; + + std::memcpy(staging->host, fragments, count * sizeof(DeviceGatherFragment)); + hipError_t status = + hipMemcpyAsync(staging->device, staging->host, count * sizeof(DeviceGatherFragment), + hipMemcpyHostToDevice, stream); + if (status != hipSuccess) { + MORI_UMBP_ERROR("[DeviceGather] gather descriptor upload failed on device {}: {}", device_id, + hipGetErrorString(status)); + (void)hipGetLastError(); + return false; + } + + hipLaunchKernelGGL(GatherFragmentsKernel, + dim3(static_cast(count) * blocks_per_fragment), + dim3(kThreadsPerBlock), 0, stream, staging->device, blocks_per_fragment); + status = hipGetLastError(); + if (status != hipSuccess) { + MORI_UMBP_ERROR("[DeviceGather] gather kernel launch failed on device {} (fragments={}): {}", + device_id, count, hipGetErrorString(status)); + return false; + } + LaunchCounter().fetch_add(1, std::memory_order_relaxed); + return true; +} + +} // namespace mori::umbp diff --git a/src/umbp/local/tiers/dram_tier.cpp b/src/umbp/local/tiers/dram_tier.cpp deleted file mode 100644 index 4beef9a89..000000000 --- a/src/umbp/local/tiers/dram_tier.cpp +++ /dev/null @@ -1,529 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/local/tiers/dram_tier.h" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#if defined(__x86_64__) || defined(__i386__) -#include -#endif - -namespace mori::umbp { - -namespace { - -#if defined(__x86_64__) || defined(__i386__) -// Non-temporal AVX2 (256-bit) copy: streaming stores bypass the cache and skip -// the read-for-ownership (RFO) on dst. This is the right choice for the real -// batch-get path, where each KV block is read ONCE from cold DRAM and the -// working set far exceeds L3: a cached copy moves 3x the block bytes through -// memory (read src + RFO dst + writeback dst), NT moves only 2x (read src + -// stream-write dst). -// -// Width: AVX2 (256-bit), not AVX-512. On Zen4 the 512-bit datapath is -// double-pumped over 256-bit units, so 512-bit stream stores give no real -// width advantage and can trip AVX-512 frequency throttling; the NT bottleneck -// is the write-combining buffer drain rate, which 256-bit already saturates. -// Measured cold 4 MiB blocks, 8 threads (no pinning) on Zen4 EPYC: -// avx2_nt ~134 > avx512_nt ~130 > glibc memcpy ~88 > cached storeu ~77. -// dst is a host pinned buffer; sfence orders the streaming stores before the -// subsequent host->device DMA reads it. -__attribute__((target("avx2"))) void NtCopyAvx2(char* d, const char* s, size_t n) { - size_t head = (32 - (reinterpret_cast(d) & 31)) & 31; - if (head > n) head = n; - std::memcpy(d, s, head); - size_t i = head; - for (; i + 128 <= n; i += 128) { - __m256i a = _mm256_loadu_si256(reinterpret_cast(s + i)); - __m256i b = _mm256_loadu_si256(reinterpret_cast(s + i + 32)); - __m256i c = _mm256_loadu_si256(reinterpret_cast(s + i + 64)); - __m256i e = _mm256_loadu_si256(reinterpret_cast(s + i + 96)); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i), a); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 32), b); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 64), c); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i + 96), e); - } - for (; i + 32 <= n; i += 32) { - __m256i a = _mm256_loadu_si256(reinterpret_cast(s + i)); - _mm256_stream_si256(reinterpret_cast<__m256i*>(d + i), a); - } - if (i < n) std::memcpy(d + i, s + i, n - i); - _mm_sfence(); -} -bool Avx2Supported() { return __builtin_cpu_supports("avx2"); } -#else -void NtCopyAvx2(char* d, const char* s, size_t n) { std::memcpy(d, s, n); } -bool Avx2Supported() { return false; } -#endif - -// Copy one KV block. Large blocks (>= 256 KiB, the real KV-page regime, always -// cold DRAM) use non-temporal stores (~1.5x over memcpy on Zen4). Tiny blocks -// fall back to glibc memcpy (its small-copy path is faster and they may be hot). -// Disable NT via UMBP_DRAM_NT_COPY=0. -inline void CopyBlock(void* dst, const void* src, size_t size) { - static const bool kNt = Avx2Supported() && !(std::getenv("UMBP_DRAM_NT_COPY") && - std::getenv("UMBP_DRAM_NT_COPY")[0] == '0'); - static const size_t kNtMinBytes = 256ull << 10; - if (kNt && size >= kNtMinBytes) { - NtCopyAvx2(static_cast(dst), static_cast(src), size); - } else { - std::memcpy(dst, src, size); - } -} - -} // namespace - -DRAMTier::DRAMTier(size_t capacity, bool use_shm, const std::string& shm_name, bool use_hugepages, - size_t hugepage_size, int numa_node, bool prefault) - : TierBackend(StorageTier::CPU_DRAM), - base_ptr_(nullptr), - capacity_(capacity), - mapped_size_(0), - used_(0), - shm_fd_(-1), - use_shm_(use_shm), - shm_name_(shm_name) { - if (use_shm_) { - shm_fd_ = shm_open(shm_name_.c_str(), O_CREAT | O_RDWR, 0666); - if (shm_fd_ < 0) { - throw std::runtime_error("shm_open failed: " + std::string(strerror(errno))); - } - if (ftruncate(shm_fd_, capacity_) < 0) { - close(shm_fd_); - shm_unlink(shm_name_.c_str()); - throw std::runtime_error("ftruncate failed: " + std::string(strerror(errno))); - } - base_ptr_ = mmap(nullptr, capacity_, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd_, 0); - if (base_ptr_ == MAP_FAILED) { - close(shm_fd_); - shm_unlink(shm_name_.c_str()); - throw std::runtime_error("mmap failed: " + std::string(strerror(errno))); - } - mapped_size_ = capacity_; - } else { - HostMemAllocator allocator; - HostBufferOptions opts; - opts.backing = - use_hugepages ? HostBufferBacking::kAnonymousHugetlb : HostBufferBacking::kAnonymous; - opts.hugepage_size = hugepage_size; - opts.numa_node = numa_node; - opts.prefault = prefault; - - host_buf_handle_ = allocator.Alloc(capacity_, opts); - if (!host_buf_handle_.valid()) { - throw std::runtime_error("DRAMTier: memory allocation failed for " + - std::to_string(capacity_) + " bytes"); - } - base_ptr_ = host_buf_handle_.ptr; - mapped_size_ = host_buf_handle_.mapped_size; - } - - // Initialize free list with entire capacity - free_list_.push_back({0, capacity_}); - - // Threads for parallel batch-read CopyBlock. Default 8, override via env, - // capped to hardware concurrency. >1 breaks the single-core memcpy ceiling. - if (const char* e = std::getenv("UMBP_DRAM_READ_THREADS")) { - int v = std::atoi(e); - if (v >= 1) read_threads_ = v; - } - if (const char* e = std::getenv("UMBP_DRAM_WRITE_THREADS")) { - int v = std::atoi(e); - if (v >= 1) write_threads_ = v; - } - unsigned hc = std::thread::hardware_concurrency(); - if (hc > 0 && read_threads_ > static_cast(hc)) read_threads_ = static_cast(hc); - if (read_threads_ < 1) read_threads_ = 1; - if (hc > 0 && write_threads_ > static_cast(hc)) write_threads_ = static_cast(hc); - if (write_threads_ < 1) write_threads_ = 1; -} - -DRAMTier::~DRAMTier() { - if (use_shm_) { - if (base_ptr_ && base_ptr_ != MAP_FAILED) { - munmap(base_ptr_, mapped_size_); - } - if (shm_fd_ >= 0) close(shm_fd_); - shm_unlink(shm_name_.c_str()); - } else { - HostMemAllocator allocator; - allocator.Free(host_buf_handle_); - } -} - -size_t DRAMTier::Allocate(size_t size) { - // First-fit allocation - for (auto it = free_list_.begin(); it != free_list_.end(); ++it) { - if (it->size >= size) { - size_t offset = it->offset; - if (it->size == size) { - free_list_.erase(it); - } else { - it->offset += size; - it->size -= size; - } - return offset; - } - } - return static_cast(-1); // Allocation failed -} - -void DRAMTier::Deallocate(size_t offset, size_t size) { - // Insert into sorted position and coalesce adjacent blocks - auto it = free_list_.begin(); - while (it != free_list_.end() && it->offset < offset) { - ++it; - } - - auto new_it = free_list_.insert(it, {offset, size}); - - // Coalesce with next block - auto next = std::next(new_it); - if (next != free_list_.end() && new_it->offset + new_it->size == next->offset) { - new_it->size += next->size; - free_list_.erase(next); - } - - // Coalesce with previous block - if (new_it != free_list_.begin()) { - auto prev = std::prev(new_it); - if (prev->offset + prev->size == new_it->offset) { - prev->size += new_it->size; - free_list_.erase(new_it); - } - } -} - -void DRAMTier::TouchLRU(const std::string& key) { - auto it = lru_map_.find(key); - if (it != lru_map_.end()) { - lru_list_.erase(it->second); - } - lru_list_.push_front(key); - lru_map_[key] = lru_list_.begin(); -} - -void DRAMTier::EvictLRU() { - if (lru_list_.empty()) return; - - const std::string& victim = lru_list_.back(); - auto slot_it = slots_.find(victim); - if (slot_it != slots_.end()) { - Deallocate(slot_it->second.offset, slot_it->second.size); - used_ -= slot_it->second.size; - slots_.erase(slot_it); - } - lru_map_.erase(victim); - lru_list_.pop_back(); -} - -bool DRAMTier::Write(const std::string& key, const void* data, size_t size) { - std::lock_guard lock(mu_); - - // If key already exists, free its old slot first - auto existing = slots_.find(key); - if (existing != slots_.end()) { - Deallocate(existing->second.offset, existing->second.size); - used_ -= existing->second.size; - slots_.erase(existing); - auto lru_it = lru_map_.find(key); - if (lru_it != lru_map_.end()) { - lru_list_.erase(lru_it->second); - lru_map_.erase(lru_it); - } - } - - // Try to allocate — do NOT self-evict. - // If no space, return false so upper layer can demote keys to SSD. - size_t offset = Allocate(size); - if (offset == static_cast(-1)) { - return false; - } - - std::memcpy(static_cast(base_ptr_) + offset, data, size); - slots_[key] = {offset, size}; - used_ += size; - TouchLRU(key); - return true; -} - -bool DRAMTier::ReadIntoPtr(const std::string& key, uintptr_t dst_ptr, size_t size) { - std::lock_guard lock(mu_); - - auto it = slots_.find(key); - if (it == slots_.end()) return false; - - // Reject if caller's buffer size does not match the stored block size. - // A mismatch indicates a caller bug (wrong page size); silently truncating - // would produce a partially-filled KV block with no error signal. - if (size != it->second.size) return false; - - std::memcpy(reinterpret_cast(dst_ptr), static_cast(base_ptr_) + it->second.offset, - size); - TouchLRU(key); - return true; -} - -std::vector DRAMTier::ReadBatchIntoPtr(const std::vector& keys, - const std::vector& dst_ptrs, - const std::vector& sizes) { - const size_t n = keys.size(); - std::vector results(n, false); - if (n == 0) return results; - - // Hold mu_ for the whole batch. This tier uses a single mutex (reads and - // writes are mutually exclusive), so holding it keeps slot offsets valid - // during the parallel copy without a use-after-free against Write/Evict/ - // Clear. The per-block copies still run in parallel within the batch, which - // is what breaks the single-core memcpy ceiling on cold DRAM. - std::lock_guard lock(mu_); - - struct Job { - void* dst; - const void* src; - size_t size; - size_t idx; - }; - std::vector jobs; - jobs.reserve(n); - for (size_t i = 0; i < n; ++i) { - auto it = slots_.find(keys[i]); - if (it == slots_.end()) continue; - if (sizes[i] != it->second.size) continue; - jobs.push_back({reinterpret_cast(dst_ptrs[i]), - static_cast(base_ptr_) + it->second.offset, sizes[i], i}); - } - - int num_threads = read_threads_; - if (num_threads > static_cast(jobs.size())) num_threads = static_cast(jobs.size()); - - if (num_threads <= 1) { - for (const auto& j : jobs) { - CopyBlock(j.dst, j.src, j.size); - results[j.idx] = true; - } - } else { - std::atomic next{0}; - auto worker = [&]() { - size_t i; - while ((i = next.fetch_add(1)) < jobs.size()) { - CopyBlock(jobs[i].dst, jobs[i].src, jobs[i].size); - } - }; - std::vector pool; - pool.reserve(num_threads); - for (int t = 0; t < num_threads; ++t) pool.emplace_back(worker); - for (auto& th : pool) th.join(); - for (const auto& j : jobs) results[j.idx] = true; - } - - for (const auto& j : jobs) TouchLRU(keys[j.idx]); - return results; -} - -std::vector DRAMTier::BatchWrite(const std::vector& keys, - const std::vector& data_ptrs, - const std::vector& sizes) { - const size_t n = keys.size(); - std::vector results(n, false); - if (n == 0) return results; - - // Hold mu_ for the whole batch (single-mutex model: serializes against other - // reads/writes). Slot allocation mutates free_list_/slots_ and must be serial; - // only the per-block payload copies run in parallel, which is what breaks the - // single-core memcpy ceiling on the backup path. - std::lock_guard lock(mu_); - - struct Job { - void* dst; - const void* src; - size_t size; - size_t idx; - size_t offset; - }; - std::vector jobs; - jobs.reserve(n); - - // Phase 1 (serial): free any existing slot for the key, then allocate. Does - // NOT self-evict — a key that doesn't fit is left false so the upper layer - // (LocalStorageManager) can demote LRU keys and retry per-key. - for (size_t i = 0; i < n; ++i) { - auto existing = slots_.find(keys[i]); - if (existing != slots_.end()) { - Deallocate(existing->second.offset, existing->second.size); - used_ -= existing->second.size; - slots_.erase(existing); - auto lru_it = lru_map_.find(keys[i]); - if (lru_it != lru_map_.end()) { - lru_list_.erase(lru_it->second); - lru_map_.erase(lru_it); - } - } - size_t offset = Allocate(sizes[i]); - if (offset == static_cast(-1)) continue; - jobs.push_back({static_cast(base_ptr_) + offset, data_ptrs[i], sizes[i], i, offset}); - } - - // Phase 2 (parallel): non-temporal CopyBlock each payload into its slot. - int num_threads = write_threads_; - if (num_threads > static_cast(jobs.size())) num_threads = static_cast(jobs.size()); - - if (num_threads <= 1) { - for (const auto& j : jobs) CopyBlock(j.dst, j.src, j.size); - } else { - std::atomic next{0}; - auto worker = [&]() { - size_t i; - while ((i = next.fetch_add(1)) < jobs.size()) { - CopyBlock(jobs[i].dst, jobs[i].src, jobs[i].size); - } - }; - std::vector pool; - pool.reserve(num_threads); - for (int t = 0; t < num_threads; ++t) pool.emplace_back(worker); - for (auto& th : pool) th.join(); - } - - // Phase 3 (serial): register slots + LRU, mark successes. - for (const auto& j : jobs) { - slots_[keys[j.idx]] = {j.offset, j.size}; - used_ += j.size; - TouchLRU(keys[j.idx]); - results[j.idx] = true; - } - return results; -} - -const void* DRAMTier::ReadPtr(const std::string& key, size_t* out_size) { - std::lock_guard lock(mu_); - - auto it = slots_.find(key); - if (it == slots_.end()) return nullptr; - - if (out_size) *out_size = it->second.size; - TouchLRU(key); - return static_cast(base_ptr_) + it->second.offset; -} - -std::vector DRAMTier::Read(const std::string& key) { - std::lock_guard lock(mu_); - - auto it = slots_.find(key); - if (it == slots_.end()) return {}; - - size_t sz = it->second.size; - std::vector buf(sz); - std::memcpy(buf.data(), static_cast(base_ptr_) + it->second.offset, sz); - TouchLRU(key); - return buf; -} - -TierCapabilities DRAMTier::Capabilities() const { - TierCapabilities caps; - caps.zero_copy_read = true; - caps.batch_read = true; // use the multi-threaded ReadBatchIntoPtr above - caps.batch_write = true; // use the multi-threaded BatchWrite below - return caps; -} - -bool DRAMTier::Exists(const std::string& key) const { - std::lock_guard lock(mu_); - return slots_.count(key) > 0; -} - -bool DRAMTier::Evict(const std::string& key) { - std::lock_guard lock(mu_); - - auto it = slots_.find(key); - if (it == slots_.end()) return false; - - Deallocate(it->second.offset, it->second.size); - used_ -= it->second.size; - slots_.erase(it); - - auto lru_it = lru_map_.find(key); - if (lru_it != lru_map_.end()) { - lru_list_.erase(lru_it->second); - lru_map_.erase(lru_it); - } - return true; -} - -std::pair DRAMTier::Capacity() const { - std::lock_guard lock(mu_); - return {used_, capacity_}; -} - -void DRAMTier::Clear() { - std::lock_guard lock(mu_); - slots_.clear(); - lru_list_.clear(); - lru_map_.clear(); - free_list_.clear(); - free_list_.push_back({0, capacity_}); - used_ = 0; -} - -std::vector DRAMTier::GetLRUCandidates(size_t max_candidates) const { - if (max_candidates == 0) max_candidates = 1; - std::lock_guard lock(mu_); - std::vector result; - result.reserve(std::min(max_candidates, lru_list_.size())); - // Walk from the back (LRU end) up to max_candidates entries. - auto it = lru_list_.rbegin(); - for (size_t i = 0; i < max_candidates && it != lru_list_.rend(); ++i, ++it) { - result.push_back(*it); - } - return result; -} - -std::string DRAMTier::GetLRUKey() const { - std::lock_guard lock(mu_); - if (lru_list_.empty()) return ""; - return lru_list_.back(); -} - -std::optional DRAMTier::GetSlotOffset(const std::string& key) const { - std::lock_guard lock(mu_); - auto it = slots_.find(key); - if (it == slots_.end()) return std::nullopt; - return it->second.offset; -} - -std::optional DRAMTier::GetLocationId(const std::string& key) const { - auto offset = GetSlotOffset(key); - if (!offset.has_value()) { - return std::nullopt; - } - return std::to_string(*offset); -} - -} // namespace mori::umbp diff --git a/src/umbp/local/tiers/host_registration.cpp b/src/umbp/local/tiers/host_registration.cpp new file mode 100644 index 000000000..9566fb54a --- /dev/null +++ b/src/umbp/local/tiers/host_registration.cpp @@ -0,0 +1,194 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/common/host_registration.h" + +#include + +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" + +namespace mori::umbp { +namespace { + +// Below this, registering inline costs under ~100 ms and keeps behaviour +// deterministic, which unit tests rely on. +constexpr size_t kSyncThresholdBytes = 2ull << 30; + +bool RegistrationEnabled() { + static const bool enabled = [] { + const char* value = std::getenv("UMBP_DRAM_HOST_REGISTER"); + if (value == nullptr) return true; + const std::string text(value); + return text != "0" && text != "off" && text != "false"; + }(); + return enabled; +} + +bool HasDevice() { + int count = 0; + if (hipGetDeviceCount(&count) != hipSuccess) { + (void)hipGetLastError(); + return false; + } + return count > 0; +} + +} // namespace + +HostTierRegistration::HostTierRegistration(void* base, size_t bytes) + : base_(static_cast(base)), bytes_(bytes) { + if (base_ == nullptr || bytes_ == 0 || !RegistrationEnabled() || !HasDevice()) { + bytes_ = 0; + return; + } + + if (bytes_ <= kSyncThresholdBytes) { + Register(); + return; + } + worker_ = std::thread([this]() { Register(); }); +} + +HostTierRegistration::~HostTierRegistration() { + if (worker_.joinable()) worker_.join(); + if (registered_bytes_.load(std::memory_order_acquire) == 0) return; + if (hipHostUnregister(base_) != hipSuccess) (void)hipGetLastError(); +} + +void HostTierRegistration::Register() { + const auto started = std::chrono::steady_clock::now(); + + // One call for the whole region, never several. hipMemcpy rejects any copy + // whose host range spans two separate registrations with hipErrorInvalidValue, + // so registering in chunks would break every object that happens to straddle + // a chunk boundary -- which is most of them once the tier is large. + // + // hipHostRegisterPortable so every rank's device can reach the tier; + // hipHostRegisterMapped so a kernel can dereference it directly. + const hipError_t status = + hipHostRegister(base_, bytes_, hipHostRegisterMapped | hipHostRegisterPortable); + if (status != hipSuccess) { + MORI_UMBP_WARN( + "[DRAMTier] hipHostRegister of {} MiB failed: {}; the GPU gather path stays off and " + "copies fall back to pageable hipMemcpy", + bytes_ >> 20, hipGetErrorString(status)); + (void)hipGetLastError(); + return; + } + + // Deliberately does not undo the registration on failure: pinning already + // succeeded and takes every hipMemcpy on this tier off the pageable staging + // path, so a device without an alias loses only the gather kernel. + if (!RecordDeviceAliases()) { + // Leaving alias_bases_ empty disables gather everywhere, which is the safe + // answer -- an unchecked device must not be published as gatherable. + MORI_UMBP_WARN( + "[DRAMTier] cannot enumerate devices to map the host region; the GPU gather path stays " + "off (the region stays registered, so copies keep the pinned path)"); + (void)hipGetLastError(); + } + + registered_bytes_.store(bytes_, std::memory_order_release); + const double seconds = + std::chrono::duration(std::chrono::steady_clock::now() - started).count(); + size_t gatherable = 0; + for (const char* alias : alias_bases_) { + if (alias != nullptr) ++gatherable; + } + MORI_UMBP_INFO( + "[DRAMTier] host memory registered for GPU access: {} MiB in {:.1f} s ({}/{} devices " + "gatherable)", + bytes_ >> 20, seconds, gatherable, alias_bases_.size()); +} + +bool HostTierRegistration::RecordDeviceAliases() { + int devices = 0; + if (hipGetDeviceCount(&devices) != hipSuccess || devices <= 0) return false; + + int previous = 0; + const bool have_previous = hipGetDevice(&previous) == hipSuccess; + if (!have_previous) (void)hipGetLastError(); + + alias_bases_.assign(static_cast(devices), nullptr); + for (int device = 0; device < devices; ++device) { + if (hipSetDevice(device) != hipSuccess) { + (void)hipGetLastError(); + continue; + } + void* alias = nullptr; + if (hipHostGetDevicePointer(&alias, base_, 0) != hipSuccess || alias == nullptr) { + (void)hipGetLastError(); + MORI_UMBP_WARN( + "[DRAMTier] device {} has no mapping for the registered host region; the " + "GPU gather path stays off for it", + device); + continue; + } + + // Only the base is queried: hipHostGetDevicePointer is documented in terms + // of the pointer an allocation starts at, so an interior offset is outside + // the contract and a rejection there would disable gather needlessly. + alias_bases_[static_cast(device)] = static_cast(alias); + } + + if (have_previous) (void)hipSetDevice(previous); + return true; +} + +bool HostTierRegistration::GatherableOn(int device_id) const { + if (registered_bytes_.load(std::memory_order_acquire) == 0) return false; + if (device_id < 0 || static_cast(device_id) >= alias_bases_.size()) return false; + return alias_bases_[static_cast(device_id)] != nullptr; +} + +void* HostTierRegistration::DeviceAddress(const void* host_ptr, size_t size, int device_id) const { + const size_t registered = registered_bytes_.load(std::memory_order_acquire); + if (registered == 0) return nullptr; + if (device_id < 0 || static_cast(device_id) >= alias_bases_.size()) return nullptr; + char* const alias = alias_bases_[static_cast(device_id)]; + if (alias == nullptr || !CoversRegistered(host_ptr, size, registered)) return nullptr; + const uintptr_t address = reinterpret_cast(host_ptr); + const uintptr_t base = reinterpret_cast(base_); + return alias + (address - base); +} + +bool HostTierRegistration::Covers(const void* ptr, size_t size) const { + return CoversRegistered(ptr, size, registered_bytes_.load(std::memory_order_acquire)); +} + +bool HostTierRegistration::CoversRegistered(const void* ptr, size_t size, size_t registered) const { + if (ptr == nullptr || bytes_ == 0) return false; + // Compared as integers, not pointers: `ptr` may be outside this region, and + // relational operators and subtraction on pointers into different objects are + // undefined behaviour even when the hardware would do the obvious thing. + const uintptr_t address = reinterpret_cast(ptr); + const uintptr_t base = reinterpret_cast(base_); + if (address < base) return false; + const uintptr_t offset = address - base; + return size <= registered && offset <= registered - size; +} + +} // namespace mori::umbp diff --git a/src/umbp/local/tiers/host_registration.h b/src/umbp/local/tiers/host_registration.h new file mode 100644 index 000000000..11ba6ad39 --- /dev/null +++ b/src/umbp/local/tiers/host_registration.h @@ -0,0 +1,26 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +// Backward-compatible source-tree include. New shared users include the common +// header directly. +#include "umbp/common/host_registration.h" diff --git a/src/umbp/local/tiers/local_storage_manager.cpp b/src/umbp/local/tiers/local_storage_manager.cpp deleted file mode 100644 index f25dcf2f1..000000000 --- a/src/umbp/local/tiers/local_storage_manager.cpp +++ /dev/null @@ -1,1128 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/local/tiers/local_storage_manager.h" - -#include -#include -#include -#include -#include -#include - -#include "mori/utils/mori_log.hpp" -#include "umbp/local/tiers/dram_tier.h" -#include "umbp/local/tiers/dummy_ssd_tier.h" -#include "umbp/local/tiers/ssd_tier.h" -#ifdef USE_SPDK -#include "umbp/local/tiers/spdk_ssd_tier.h" -#endif -#ifdef __linux__ -#include -#include -#include -#include - -#include "umbp/local/tiers/spdk_proxy_tier.h" -#include "umbp/storage/spdk/proxy/spdk_proxy_shm.h" -#endif - -namespace mori::umbp { - -// --------------------------------------------------------------------------- -// ExtractBaseHash -// -// Strips the SGLang-generated key suffix to recover the 64-hex-char SHA256 -// base hash. Key formats (rightmost fields stripped first): -// -// MHA, no PP: {hash}_{tp_rank}_{k|v} → strip _k/_v, strip _{digit} -// MHA, with PP: {hash}_{tp_rank}_{pp_rank}_{k|v} → strip _k/_v, strip _{digit}, strip _{digit} -// MLA, no PP: {hash}__k → strip _k, strip trailing _ -// MLA, with PP: {hash}_{pp_rank}_{k} → strip _k, strip _{digit} -// -// The SHA256 base hash is always exactly 64 lower-case hex characters. -// We use that invariant to anchor stripping: stop as soon as the remaining -// string length equals 64 and the remaining characters are all hex. -// --------------------------------------------------------------------------- -static bool IsHexChar(char c) { - return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'); -} - -static bool IsAllHex(const std::string& s, size_t len) { - for (size_t i = 0; i < len; ++i) { - if (!IsHexChar(s[i])) return false; - } - return true; -} - -std::string LocalStorageManager::ExtractBaseHash(const std::string& key) { - // Fast-path: already looks like a bare hash. - if (key.size() == 64 && IsAllHex(key, 64)) return key; - - std::string s = key; - - // Strip _k or _v suffix. - if (s.size() >= 2 && s[s.size() - 2] == '_' && (s.back() == 'k' || s.back() == 'v')) { - s.resize(s.size() - 2); - } else { - return key; // Unrecognised format — return unchanged. - } - - // Iteratively strip _{digit+} rank segments until the remainder is a 64-char - // hex string or no more strippable segments remain. - for (int pass = 0; pass < 2; ++pass) { - if (s.size() == 64 && IsAllHex(s, 64)) return s; - - // Find the last '_' and check that everything after it is digits. - size_t pos = s.rfind('_'); - if (pos == std::string::npos) break; - - bool all_digits = true; - for (size_t i = pos + 1; i < s.size(); ++i) { - if (s[i] < '0' || s[i] > '9') { - all_digits = false; - break; - } - } - - // Handle MLA no-PP double-underscore: after stripping _k we may have a - // trailing '_' with nothing after it (pos == s.size()-1). - if (pos == s.size() - 1) { - // Empty segment after last '_' — this is the double-underscore case. - s.resize(pos); - continue; - } - - if (!all_digits) break; - s.resize(pos); - } - - if (s.size() == 64 && IsAllHex(s, 64)) return s; - - // Fallback: return original key — group will contain only this key (size 1), - // which is safe (single-key eviction). - return key; -} - -// --------------------------------------------------------------------------- -// Depth and group map helpers -// --------------------------------------------------------------------------- - -void LocalStorageManager::RecordDepth(const std::string& key, int depth) { - std::unique_lock lock(depth_mu_); - depth_map_[key] = depth; -} - -int LocalStorageManager::GetDepth(const std::string& key) const { - std::shared_lock lock(depth_mu_); - auto it = depth_map_.find(key); - return (it != depth_map_.end()) ? it->second : -1; -} - -void LocalStorageManager::RemoveDepthAndGroup(const std::string& key) { - std::unique_lock lock(depth_mu_); - depth_map_.erase(key); - - std::string base = ExtractBaseHash(key); - auto git = group_map_.find(base); - if (git != group_map_.end()) { - auto& vec = git->second; - vec.erase(std::remove(vec.begin(), vec.end(), key), vec.end()); - if (vec.empty()) group_map_.erase(git); - } -} - -void LocalStorageManager::RecordGroup(const std::string& key) { - std::unique_lock lock(depth_mu_); - std::string base = ExtractBaseHash(key); - auto& vec = group_map_[base]; - // Only append if not already present (idempotent for re-puts). - if (std::find(vec.begin(), vec.end(), key) == vec.end()) { - vec.push_back(key); - } -} - -std::vector LocalStorageManager::GetGroup(const std::string& key) const { - std::shared_lock lock(depth_mu_); - std::string base = ExtractBaseHash(key); - auto it = group_map_.find(base); - if (it != group_map_.end() && !it->second.empty()) return it->second; - return {key}; -} - -// --------------------------------------------------------------------------- -// SelectVictim -// -// "lru" policy: O(1) — return plain LRU tail key. -// "prefix_aware_lru": scan up to eviction_candidate_window candidates from -// the LRU tail, score by depth (higher depth = preferred victim, deeper suffix -// block first). Tie-break by LRU position (earlier in the candidate list = -// older = preferred victim among equal-depth candidates). -// Falls back to plain LRU when no depth metadata is available. -// --------------------------------------------------------------------------- -std::string LocalStorageManager::SelectVictim(TierBackend* tier) { - if (config_.eviction.policy != "prefix_aware_lru") { - return tier->GetLRUKey(); - } - - size_t window = config_.eviction.candidate_window; - if (window == 0) window = 1; - - std::vector candidates = tier->GetLRUCandidates(window); - if (candidates.empty()) return ""; - - // Score: (depth, position) where higher depth wins, lower position (older) - // wins ties. depth == -1 is treated as 0 for scoring so metadata-free - // keys degrade to plain LRU ordering. - int best_depth = -2; // sentinel below any real score - size_t best_pos = 0; - std::string best_key; - - for (size_t i = 0; i < candidates.size(); ++i) { - int d = GetDepth(candidates[i]); // -1 if unknown - int score = (d >= 0) ? d : 0; - // Higher score wins; tie-break: smaller i (older LRU position) wins. - if (best_key.empty() || score > best_depth || (score == best_depth && i < best_pos)) { - best_depth = score; - best_pos = i; - best_key = candidates[i]; - } - } - - return best_key; -} - -// --------------------------------------------------------------------------- -// WriteFromPtrWithDepth -// --------------------------------------------------------------------------- -bool LocalStorageManager::WriteFromPtrWithDepth(const std::string& key, uintptr_t src, size_t size, - int depth, StorageTier tier) { - bool ok = WriteFromPtr(key, src, size, tier); - if (ok && depth >= 0) { - RecordDepth(key, depth); - RecordGroup(key); - } - return ok; -} - -LocalStorageManager::~LocalStorageManager() = default; - -// --------------------------------------------------------------------------- -// Auto-fork proxy daemon (Linux only) -// --------------------------------------------------------------------------- -#ifdef __linux__ -static std::string FindProxyBinary(const std::string& explicit_path) { - auto try_candidate = [](const std::string& candidate) -> std::string { - return (access(candidate.c_str(), X_OK) == 0) ? candidate : ""; - }; - - if (!explicit_path.empty()) { - if (auto hit = try_candidate(explicit_path); !hit.empty()) return hit; - MORI_UMBP_WARN("LSM: UMBP_SPDK_PROXY_BIN='{}' not executable", explicit_path); - } - - char exe_buf[4096]; - ssize_t len = readlink("/proc/self/exe", exe_buf, sizeof(exe_buf) - 1); - if (len > 0) { - exe_buf[len] = '\0'; - std::string dir(exe_buf); - size_t slash = dir.rfind('/'); - if (slash != std::string::npos) { - dir.resize(slash + 1); - if (auto hit = try_candidate(dir + "spdk_proxy"); !hit.empty()) return hit; - - // CMake build tree layout: test binaries live under - // build_umbp/tests/... while spdk_proxy lives under build_umbp/src/umbp. - std::string parent = dir; - for (int depth = 0; depth < 8; ++depth) { - if (auto hit = try_candidate(parent + "src/umbp/spdk_proxy"); !hit.empty()) { - return hit; - } - if (parent.empty() || parent == "/") break; - size_t end = parent.size(); - while (end > 0 && parent[end - 1] == '/') --end; - if (end == 0) break; - size_t prev = parent.rfind('/', end - 1); - if (prev == std::string::npos) break; - parent.resize(prev + 1); - } - } - } - - return "spdk_proxy"; -} - -class ScopedBootstrapLock { - public: - explicit ScopedBootstrapLock(const std::string& shm_name) { - path_ = ::umbp::proxy::ProxyShmRegion::BootstrapLockPath(shm_name); - fd_ = open(path_.c_str(), O_CREAT | O_RDWR, 0600); - if (fd_ < 0) { - MORI_UMBP_ERROR("LSM: open bootstrap lock '{}' failed: {}", path_, strerror(errno)); - return; - } - if (flock(fd_, LOCK_EX) != 0) { - MORI_UMBP_ERROR("LSM: flock bootstrap lock '{}' failed: {}", path_, strerror(errno)); - close(fd_); - fd_ = -1; - } - } - - ~ScopedBootstrapLock() { - if (fd_ >= 0) { - flock(fd_, LOCK_UN); - close(fd_); - } - } - - bool IsValid() const { return fd_ >= 0; } - - private: - std::string path_; - int fd_ = -1; -}; - -static void SetEnvVarFromConfig(const char* name, const std::string& value) { - if (!value.empty()) setenv(name, value.c_str(), 1); -} - -static void SetEnvVarFromConfig(const char* name, size_t value) { - setenv(name, std::to_string(value).c_str(), 1); -} - -static void SetEnvVarFromConfig(const char* name, int value) { - setenv(name, std::to_string(value).c_str(), 1); -} - -static void SetEnvVarFromConfig(const char* name, bool value) { - setenv(name, value ? "1" : "0", 1); -} - -int LocalStorageManager::SpawnProxyDaemon(const std::string& shm_name) { - std::string bin = FindProxyBinary(config_.ssd.spdk_proxy_bin); - MORI_UMBP_INFO("LSM: launching spdk_proxy binary '{}' for shm '{}'", bin, shm_name); - pid_t pid = fork(); - if (pid < 0) { - MORI_UMBP_ERROR("LSM: fork() failed: {}", strerror(errno)); - return -1; - } - - if (pid == 0) { - setsid(); - - SetEnvVarFromConfig("UMBP_SPDK_PROXY_SHM", shm_name); - SetEnvVarFromConfig("UMBP_SPDK_PROXY_MAX_CHANNELS", - static_cast(config_.ssd.spdk_proxy_max_channels)); - SetEnvVarFromConfig("UMBP_SPDK_PROXY_DATA_PER_CHANNEL_MB", - config_.ssd.spdk_proxy_data_per_channel_mb); - SetEnvVarFromConfig("UMBP_SPDK_PROXY_IDLE_EXIT_TIMEOUT_MS", - config_.ssd.spdk_proxy_idle_exit_timeout_ms); - SetEnvVarFromConfig("UMBP_SPDK_PROXY_ALLOW_BORROW", config_.ssd.spdk_proxy_allow_borrow); - SetEnvVarFromConfig("UMBP_SPDK_PROXY_RESERVED_SHARED_BYTES", - config_.ssd.spdk_proxy_reserved_shared_bytes); - SetEnvVarFromConfig("UMBP_SSD_CAPACITY", config_.ssd.capacity_bytes); - SetEnvVarFromConfig("UMBP_SPDK_NVME_PCI", config_.ssd.spdk_nvme_pci_addr); - SetEnvVarFromConfig("UMBP_SPDK_NVME_CTRL", config_.ssd.spdk_nvme_ctrl_name); - SetEnvVarFromConfig("UMBP_SPDK_BDEV", config_.ssd.spdk_bdev_name); - SetEnvVarFromConfig("UMBP_SPDK_REACTOR_MASK", config_.ssd.spdk_reactor_mask); - SetEnvVarFromConfig("UMBP_SPDK_MEM_MB", config_.ssd.spdk_mem_size_mb); - SetEnvVarFromConfig("UMBP_SPDK_IO_WORKERS", config_.ssd.spdk_io_workers); - - // Suppress child output unless verbose logging is requested. - // Uses direct getenv (fork-safe, no spdlog) to check the Mori log level. - { - const char* lv = getenv("MORI_UMBP_LOG_LEVEL"); - if (!lv) lv = getenv("MORI_GLOBAL_LOG_LEVEL"); - bool verbose = - lv && (strcmp(lv, "info") == 0 || strcmp(lv, "INFO") == 0 || strcmp(lv, "debug") == 0 || - strcmp(lv, "DEBUG") == 0 || strcmp(lv, "trace") == 0 || strcmp(lv, "TRACE") == 0); - if (!verbose) { - int devnull = open("/dev/null", O_WRONLY); - if (devnull >= 0) { - dup2(devnull, STDOUT_FILENO); - dup2(devnull, STDERR_FILENO); - close(devnull); - } - } - } - execlp(bin.c_str(), "spdk_proxy", static_cast(nullptr)); - fprintf(stderr, "[UMBP ERROR] execlp('%s') failed: %s\n", bin.c_str(), strerror(errno)); - _exit(127); - } - - MORI_UMBP_INFO("LSM: spawned spdk_proxy service pid={} shm='{}'", pid, shm_name); - return 0; -} - -int LocalStorageManager::EnsureProxyDaemon(const std::string& shm_name) { - int probe = ::umbp::proxy::ProxyShmRegion::ProbeExisting(shm_name); - if (probe == 1) return 0; - if (probe == -2) { - MORI_UMBP_ERROR("LSM: proxy protocol version mismatch on SHM '{}'", shm_name); - return -1; - } - if (!config_.ssd.spdk_proxy_auto_start) { - if (probe == -1) { - return SpdkProxyTier::WaitForProxy(shm_name, config_.ssd.spdk_proxy_startup_timeout_ms) ? 0 - : -1; - } - MORI_UMBP_ERROR("LSM: SPDK proxy absent on SHM '{}' and auto-start disabled", shm_name); - return -1; - } - - ScopedBootstrapLock lock(shm_name); - if (!lock.IsValid()) return -1; - - probe = ::umbp::proxy::ProxyShmRegion::ProbeExisting(shm_name); - if (probe == 1) return 0; - if (probe == -2) { - MORI_UMBP_ERROR("LSM: proxy protocol version mismatch on SHM '{}'", shm_name); - return -1; - } - if (probe == -1) { - if (SpdkProxyTier::WaitForProxy(shm_name, config_.ssd.spdk_proxy_startup_timeout_ms)) { - return 0; - } - - ::umbp::proxy::ProxyShmRegion existing; - if (existing.Attach(shm_name) == 0) { - auto* hdr = existing.Header(); - uint32_t st = hdr->state.load(std::memory_order_acquire); - uint32_t pid = hdr->proxy_pid.load(std::memory_order_acquire); - bool alive = false; -#ifdef __linux__ - alive = (pid > 0 && kill(static_cast(pid), 0) == 0); -#endif - if (alive && st == static_cast(::umbp::proxy::ProxyState::ERROR)) { - MORI_UMBP_ERROR("LSM: proxy on SHM '{}' is alive but in ERROR state", shm_name); - return -1; - } - if (alive && st != static_cast(::umbp::proxy::ProxyState::SHUTDOWN)) { - MORI_UMBP_ERROR("LSM: proxy on SHM '{}' did not become READY in time (state={})", shm_name, - st); - return -1; - } - } - } - - ::umbp::proxy::ProxyShmRegion::CleanupStale(shm_name); - if (SpawnProxyDaemon(shm_name) != 0) return -1; - - return SpdkProxyTier::WaitForProxy(shm_name, config_.ssd.spdk_proxy_startup_timeout_ms) ? 0 : -1; -} -#endif - -LocalStorageManager::LocalStorageManager(const UMBPConfig& config, - mori::umbp::LocalBlockIndex* index) - : config_(config), role_(config.ResolveRole()), index_(index) { - std::string error_message; - if (!config_.Validate(&error_message)) { - throw std::runtime_error("invalid UMBP config: " + error_message); - } - // DRAM tier is always present (fastest) - tiers_.push_back( - {StorageTier::CPU_DRAM, - std::make_unique(config_.dram.capacity_bytes, config_.dram.use_shared_memory, - config_.dram.shm_name, config_.dram.use_hugepages, - config_.dram.hugepage_size, config_.dram.numa_node, - config_.dram.prefault)}); - - // SSD tier is optional (slower) - if (config_.ssd.enabled) { - std::unique_ptr ssd_backend; - bool use_proxy = false; - - if (config_.ssd.ssd_backend == "spdk") { -#ifdef USE_SPDK - if (role_ == UMBPRole::Standalone) { - // Standalone with USE_SPDK compiled: try direct SPDK access first - // (best perf, single process, no proxy overhead). - auto spdk_tier = std::make_unique(config_.ssd); - if (spdk_tier->IsValid()) { - ssd_backend = std::unique_ptr(spdk_tier.release()); - } else { - throw std::runtime_error( - "UMBP_SSD_BACKEND=spdk: SpdkSsdTier init failed. " - "Run 'tools/umbp_spdk_preflight.sh --pci ' to diagnose."); - } - } else { - // Non-Standalone (Leader/Follower): always use proxy for - // multi-process coordination. - use_proxy = true; - } -#else - // USE_SPDK not set — SpdkSsdTier (direct) is unavailable. - // Fall back to proxy mode: SpdkProxyTier has zero SPDK dependency - // (pure POSIX SHM), and the separate spdk_proxy daemon handles NVMe. - use_proxy = true; -#endif - } else if (config_.ssd.ssd_backend == "spdk_proxy") { - use_proxy = true; - } - -#ifdef __linux__ - if (use_proxy && !ssd_backend) { - std::string proxy_shm_name = config_.ssd.spdk_proxy_shm_name; - if (proxy_shm_name.empty()) proxy_shm_name = ::umbp::proxy::kDefaultShmName; - - bool proxy_ready = false; - if (config_.ssd.spdk_proxy_auto_start) { - proxy_ready = (EnsureProxyDaemon(proxy_shm_name) == 0); - } else { - proxy_ready = - SpdkProxyTier::WaitForProxy(proxy_shm_name, config_.ssd.spdk_proxy_startup_timeout_ms); - } - - if (proxy_ready) { - auto proxy_tier = std::make_unique(config_.ssd); - if (proxy_tier->IsValid()) { - ssd_backend = std::unique_ptr(proxy_tier.release()); - } - } - - if (!ssd_backend) { - bool explicit_spdk = - (config_.ssd.ssd_backend == "spdk" || config_.ssd.ssd_backend == "spdk_proxy"); - if (explicit_spdk) { - throw std::runtime_error( - "UMBP_SSD_BACKEND=" + config_.ssd.ssd_backend + ": SPDK proxy connect failed (shm=" + - proxy_shm_name + "). Run 'tools/umbp_spdk_preflight.sh --pci ' to diagnose."); - } - fprintf(stderr, "[UMBP WARN] SPDK proxy connect failed. Falling back to POSIX SSD.\n"); - } - } -#endif - if (!ssd_backend && config_.ssd.ssd_backend == "dummy_storage") { - ssd_backend = std::make_unique(config_.ssd.capacity_bytes); - } - - if (!ssd_backend) { - SSDAccessMode segmented_access_mode = SSDAccessMode::ReadWrite; - if (role_ == UMBPRole::SharedSSDFollower) { - segmented_access_mode = SSDAccessMode::ReadOnlyShared; - } - ssd_backend = std::make_unique(config_.ssd.storage_dir, config_.ssd.capacity_bytes, - config_.ssd, segmented_access_mode); - } - tiers_.push_back({StorageTier::LOCAL_SSD, std::move(ssd_backend)}); - } -} - -TierBackend* LocalStorageManager::GetTier(StorageTier tier) { - for (auto& entry : tiers_) { - if (entry.id == tier) return entry.backend.get(); - } - return nullptr; -} - -const TierBackend* LocalStorageManager::GetTier(StorageTier tier) const { - for (const auto& entry : tiers_) { - if (entry.id == tier) return entry.backend.get(); - } - return nullptr; -} - -TierBackend* LocalStorageManager::FindTierHolding(const std::string& key) { - for (auto& entry : tiers_) { - if (entry.backend->Exists(key)) return entry.backend.get(); - } - return nullptr; -} - -const TierBackend* LocalStorageManager::FindTierHolding(const std::string& key) const { - for (const auto& entry : tiers_) { - if (entry.backend->Exists(key)) return entry.backend.get(); - } - return nullptr; -} - -TierBackend* LocalStorageManager::NextSlowerTier(StorageTier current) { - for (size_t i = 0; i + 1 < tiers_.size(); ++i) { - if (tiers_[i].id == current) return tiers_[i + 1].backend.get(); - } - return nullptr; -} - -TierBackend* LocalStorageManager::NextFasterTier(StorageTier current) { - for (size_t i = 1; i < tiers_.size(); ++i) { - if (tiers_[i].id == current) return tiers_[i - 1].backend.get(); - } - return nullptr; -} - -void LocalStorageManager::SetOnTierChange(TierChangeCallback cb) { - on_tier_change_ = std::move(cb); -} - -std::optional LocalStorageManager::BuildTierLocationInfo( - TierBackend* tier, const std::string& key, size_t size) { - if (!tier) { - return std::nullopt; - } - - auto raw_id = tier->GetLocationId(key); - if (!raw_id.has_value()) { - return std::nullopt; - } - - TierLocationInfo info; - info.location_id = "0:" + *raw_id; - info.size = size; - return info; -} - -bool LocalStorageManager::MoveKey(const std::string& key, TierBackend* from, TierBackend* to) { - // Fast path for DRAM->slower movement to avoid an intermediate vector copy. - if (from->Capabilities().zero_copy_read) { - size_t sz = 0; - const void* ptr = from->ReadPtr(key, &sz); - if (ptr && sz > 0 && to->Write(key, ptr, sz)) { - StorageTier from_id = from->tier_id(); - auto new_location = BuildTierLocationInfo(to, key, sz); - from->Evict(key); - UpsertIndexTier(key, to->tier_id(), sz); - if (on_tier_change_) on_tier_change_(key, from_id, to->tier_id(), new_location); - return true; - } - } - - auto data = from->Read(key); - if (data.empty()) return false; - if (!to->Write(key, data.data(), data.size())) return false; - StorageTier from_id = from->tier_id(); - auto new_location = BuildTierLocationInfo(to, key, data.size()); - from->Evict(key); - UpsertIndexTier(key, to->tier_id(), data.size()); - if (on_tier_change_) on_tier_change_(key, from_id, to->tier_id(), new_location); - return true; -} - -bool LocalStorageManager::DemoteLRUForSpace(TierBackend* tier) { - TierBackend* slower = NextSlowerTier(tier->tier_id()); - if (!slower) return false; - if (role_ == UMBPRole::SharedSSDFollower && slower->tier_id() == StorageTier::LOCAL_SSD) { - return false; - } - - std::string victim = SelectVictim(tier); - if (victim.empty()) return false; - - // Group-demote: move all keys sharing the same page together. - // depth/group metadata is preserved (key still exists in slower tier). - std::vector group = GetGroup(victim); - bool any_moved = false; - for (const auto& k : group) { - if (tier->Exists(k)) { - if (MoveKey(k, tier, slower)) { - any_moved = true; - } - } - } - return any_moved; -} - -bool LocalStorageManager::Write(const std::string& key, const void* data, size_t size, - StorageTier tier) { - TierBackend* target = GetTier(tier); - if (!target) return false; - - // Try direct write - if (target->Write(key, data, size)) return true; - - // Target full — try demoting LRU keys to next-slower tier - while (DemoteLRUForSpace(target)) { - if (target->Write(key, data, size)) return true; - } - TierBackend* faster_than_target = NextFasterTier(target->tier_id()); - while (true) { - std::string victim = SelectVictim(target); - if (victim.empty()) return false; - std::vector group = GetGroup(victim); - for (const auto& k : group) { - if (target->Exists(k)) { - bool only_on_target = !faster_than_target || !faster_than_target->Exists(k); - target->Evict(k); - if (only_on_target) { - if (on_tier_change_) on_tier_change_(k, target->tier_id(), std::nullopt, std::nullopt); - if (index_) index_->Remove(k); - RemoveDepthAndGroup(k); - } - } - } - if (target->Write(key, data, size)) return true; - } -} - -bool LocalStorageManager::WriteFromPtr(const std::string& key, uintptr_t src, size_t size, - StorageTier tier) { - return Write(key, reinterpret_cast(src), size, tier); -} - -std::vector LocalStorageManager::WriteBatchFromPtr(const std::vector& keys, - const std::vector& srcs, - const std::vector& sizes, - StorageTier tier) { - const size_t n = keys.size(); - std::vector results(n, false); - if (n == 0) return results; - - TierBackend* target = GetTier(tier); - if (!target) return results; - - // Fast path: parallel batch write when the tier supports it. BatchWrite does - // not self-evict, so keys that don't fit come back false; retry those per-key - // through Write() which demotes LRU victims to make room. - if (target->Capabilities().batch_write && n > 1) { - std::vector ptrs; - ptrs.reserve(n); - for (size_t i = 0; i < n; ++i) ptrs.push_back(reinterpret_cast(srcs[i])); - results = target->BatchWrite(keys, ptrs, sizes); - for (size_t i = 0; i < n; ++i) { - if (!results[i]) { - results[i] = Write(keys[i], reinterpret_cast(srcs[i]), sizes[i], tier); - } - } - return results; - } - - // Slow path: per-key (also covers single-element batches). - for (size_t i = 0; i < n; ++i) { - results[i] = Write(keys[i], reinterpret_cast(srcs[i]), sizes[i], tier); - } - return results; -} - -bool LocalStorageManager::ReadIntoPtr(const std::string& key, uintptr_t dst, size_t size) { - bool ok = ReadIntoPtrNoPromote(key, dst, size); - if (ok && config_.eviction.auto_promote_on_read) { - MaybeAutoPromote(key); - } - return ok; -} - -bool LocalStorageManager::ReadIntoPtrNoPromote(const std::string& key, uintptr_t dst, size_t size) { - if (index_) { - auto loc = index_->Lookup(key); - if (loc) { - TierBackend* hinted = GetTier(loc->tier); - if (hinted && hinted->Exists(key)) { - bool ok = hinted->ReadIntoPtr(key, dst, size); - if (ok) return true; - } - } - } - - for (size_t i = 0; i < tiers_.size(); ++i) { - if (tiers_[i].backend->Exists(key)) { - bool ok = tiers_[i].backend->ReadIntoPtr(key, dst, size); - return ok; - } - } - return false; -} - -std::vector LocalStorageManager::ReadBatchIntoPtr(const std::vector& keys, - const std::vector& dst_ptrs, - const std::vector& sizes) { - std::vector results(keys.size(), false); - if (keys.empty()) return results; - - // Partition keys by tier using index hints. - // tier_indices[t] holds original indices for keys hinted to tier t. - // Use the integer value of StorageTier as array index (0=DRAM, 1=SSD). - constexpr int kNumTiers = 2; - std::vector tier_indices[kNumTiers]; - std::vector no_hint_indices; - - for (size_t i = 0; i < keys.size(); ++i) { - bool has_hint = false; - if (index_) { - auto loc = index_->Lookup(keys[i]); - if (loc) { - int t = static_cast(loc->tier); - if (t >= 0 && t < kNumTiers) { - tier_indices[t].push_back(i); - has_hint = true; - } - } - } - if (!has_hint) { - no_hint_indices.push_back(i); - } - } - - // Dispatch per-tier batch reads. - for (int t = 0; t < kNumTiers; ++t) { - auto& indices = tier_indices[t]; - if (indices.empty()) continue; - - StorageTier tier_id = static_cast(t); - TierBackend* tier = GetTier(tier_id); - if (!tier) { - for (size_t idx : indices) no_hint_indices.push_back(idx); - continue; - } - - if (tier->Capabilities().batch_read && indices.size() > 1) { - std::vector batch_keys; - std::vector batch_ptrs; - std::vector batch_sizes; - batch_keys.reserve(indices.size()); - batch_ptrs.reserve(indices.size()); - batch_sizes.reserve(indices.size()); - for (size_t idx : indices) { - batch_keys.push_back(keys[idx]); - batch_ptrs.push_back(dst_ptrs[idx]); - batch_sizes.push_back(sizes[idx]); - } - auto batch_results = tier->ReadBatchIntoPtr(batch_keys, batch_ptrs, batch_sizes); - for (size_t j = 0; j < indices.size(); ++j) { - if (batch_results[j]) { - results[indices[j]] = true; - if (tier_id != StorageTier::CPU_DRAM && config_.eviction.auto_promote_on_read) { - MaybeAutoPromote(keys[indices[j]]); - } - } else { - no_hint_indices.push_back(indices[j]); - } - } - } else { - for (size_t idx : indices) { - if (tier->Exists(keys[idx])) { - bool ok = tier->ReadIntoPtr(keys[idx], dst_ptrs[idx], sizes[idx]); - if (ok) { - results[idx] = true; - if (tier_id != StorageTier::CPU_DRAM && config_.eviction.auto_promote_on_read) { - MaybeAutoPromote(keys[idx]); - } - } else { - no_hint_indices.push_back(idx); - } - } else { - no_hint_indices.push_back(idx); - } - } - } - } - - // Fallback: full tier scan for keys without hints or that failed hinted read. - for (size_t idx : no_hint_indices) { - if (results[idx]) continue; // already resolved - results[idx] = ReadIntoPtr(keys[idx], dst_ptrs[idx], sizes[idx]); - } - - return results; -} - -bool LocalStorageManager::Exists(const std::string& key) const { - return FindTierHolding(key) != nullptr; -} - -bool LocalStorageManager::Evict(const std::string& key) { - TierBackend* tier = FindTierHolding(key); - if (!tier) { - if (index_) index_->Remove(key); - RemoveDepthAndGroup(key); - return false; - } - bool ok = tier->Evict(key); - if (ok) { - if (on_tier_change_) on_tier_change_(key, tier->tier_id(), std::nullopt, std::nullopt); - if (index_) index_->Remove(key); - } - RemoveDepthAndGroup(key); - return ok; -} - -std::pair LocalStorageManager::Capacity(StorageTier tier) const { - const TierBackend* t = GetTier(tier); - if (!t) return {0, 0}; - return t->Capacity(); -} - -bool LocalStorageManager::Demote(const std::string& key) { - TierBackend* from = FindTierHolding(key); - if (!from) return false; - - TierBackend* to = NextSlowerTier(from->tier_id()); - if (!to) return false; - - return MoveKey(key, from, to); -} - -bool LocalStorageManager::Promote(const std::string& key) { - TierBackend* from = FindTierHolding(key); - if (!from) return false; - - TierBackend* to = NextFasterTier(from->tier_id()); - if (!to) return false; - - // Already in the faster tier - if (to->Exists(key)) return true; - - // Try direct move - auto data = from->Read(key); - if (data.empty()) return false; - - // Shared SSD follower mode is read-only against SSD. Promotion into DRAM - // must never demote DRAM victims into SSD. - if (role_ == UMBPRole::SharedSSDFollower && from->tier_id() == StorageTier::LOCAL_SSD && - to->tier_id() == StorageTier::CPU_DRAM) { - if (!to->Write(key, data.data(), data.size())) { - while (true) { - std::string victim = SelectVictim(to); - if (victim.empty()) return false; - std::vector group = GetGroup(victim); - for (const auto& k : group) { - if (to->Exists(k)) { - to->Evict(k); - if (index_) index_->Remove(k); - RemoveDepthAndGroup(k); - } - } - if (to->Write(key, data.data(), data.size())) break; - } - } - // Drop read-tracking metadata in source tier while preserving shared file. - auto new_location = BuildTierLocationInfo(to, key, data.size()); - StorageTier from_id = from->tier_id(); - from->Evict(key); - UpsertIndexTier(key, to->tier_id(), data.size()); - if (on_tier_change_) on_tier_change_(key, from_id, to->tier_id(), new_location); - return true; - } - - if (!to->Write(key, data.data(), data.size())) { - // Target full — demote LRU keys from target to make room - while (DemoteLRUForSpace(to)) { - if (to->Write(key, data.data(), data.size())) { - auto new_location = BuildTierLocationInfo(to, key, data.size()); - StorageTier from_id = from->tier_id(); - from->Evict(key); - UpsertIndexTier(key, to->tier_id(), data.size()); - if (on_tier_change_) on_tier_change_(key, from_id, to->tier_id(), new_location); - return true; - } - } - return false; // Cannot promote - } - - auto new_location = BuildTierLocationInfo(to, key, data.size()); - StorageTier from_id = from->tier_id(); - from->Evict(key); - UpsertIndexTier(key, to->tier_id(), data.size()); - if (on_tier_change_) on_tier_change_(key, from_id, to->tier_id(), new_location); - return true; -} - -bool LocalStorageManager::CopyToSSD(const std::string& key) { - if (role_ == UMBPRole::SharedSSDFollower) return false; - - TierBackend* dram = GetTier(StorageTier::CPU_DRAM); - TierBackend* ssd = GetTier(StorageTier::LOCAL_SSD); - if (!dram || !ssd) return false; - - // Already on SSD - if (ssd->Exists(key)) return true; - - // Zero-copy: get a raw pointer into mmap'd DRAM instead of allocating a copy. - size_t data_size = 0; - const void* ptr = dram->ReadPtr(key, &data_size); - if (!ptr || data_size == 0) return false; - - if (ssd->Write(key, ptr, data_size)) return true; - - // SSD full — evict SSD LRU entries to make room. - // If a victim is only on SSD (no DRAM copy), we must also update the index. - while (true) { - std::string victim = SelectVictim(ssd); - if (victim.empty()) return false; - - std::vector group = GetGroup(victim); - for (const auto& k : group) { - if (!ssd->Exists(k)) continue; - bool only_on_ssd = !dram->Exists(k); - // Keep the index consistent: if the victim is only on SSD, evicting it - // is data loss and the index must be updated. - if (index_) { - auto loc = index_->Lookup(k); - if (loc && (loc->tier == StorageTier::LOCAL_SSD || only_on_ssd)) { - index_->Remove(k); - } - } - ssd->Evict(k); - // Only clear metadata if key is gone from all tiers (data loss). - if (only_on_ssd) RemoveDepthAndGroup(k); - } - if (ssd->Write(key, ptr, data_size)) return true; - } -} - -bool LocalStorageManager::CopyToSSDBatch(const std::vector& keys) { - if (role_ == UMBPRole::SharedSSDFollower) return false; - - TierBackend* dram = GetTier(StorageTier::CPU_DRAM); - TierBackend* ssd = GetTier(StorageTier::LOCAL_SSD); - if (!dram || !ssd) return false; - - // Gather zero-copy pointers for keys not yet on SSD. - std::vector batch_keys; - std::vector batch_ptrs; - std::vector batch_sizes; - std::vector fallback_keys; - - for (const auto& key : keys) { - if (ssd->Exists(key)) continue; - size_t sz = 0; - const void* ptr = dram->ReadPtr(key, &sz); - if (!ptr || sz == 0) continue; - batch_keys.push_back(key); - batch_ptrs.push_back(ptr); - batch_sizes.push_back(sz); - } - - if (batch_keys.empty()) return true; - - // Try batch write if SSDTier is available. - if (ssd->Capabilities().batch_write && batch_keys.size() > 1) { - if (ssd->WriteBatch(batch_keys, batch_ptrs, batch_sizes)) return true; - // Batch failed (likely capacity) — fall through to per-key CopyToSSD. - } - - // Fallback: per-key CopyToSSD. - bool all_ok = true; - for (const auto& key : batch_keys) { - if (!CopyToSSD(key)) all_ok = false; - } - return all_ok; -} - -void LocalStorageManager::MaybeAutoPromote(const std::string& key) { - if (!config_.eviction.auto_promote_on_read) return; - - // Only promote if key is not in the fastest tier - if (tiers_.empty()) return; - TierBackend* fastest = tiers_.front().backend.get(); - if (fastest->Exists(key)) return; - - // Check fastest tier watermark before promoting - auto [used, total] = fastest->Capacity(); - double utilization = static_cast(used) / total; - if (utilization >= config_.dram.high_watermark) return; - - // Best-effort promote - if (role_ == UMBPRole::SharedSSDFollower) { - InsertReadCacheNoWriteback(key); - } else { - Promote(key); - } -} - -bool LocalStorageManager::InsertReadCacheNoWriteback(const std::string& key) { - TierBackend* source = GetTier(StorageTier::LOCAL_SSD); - TierBackend* dram = GetTier(StorageTier::CPU_DRAM); - if (!source || !dram) return false; - - if (dram->Exists(key)) { - UpsertIndexTier(key, StorageTier::CPU_DRAM, 0); - return true; - } - if (!source->Exists(key)) return false; - - auto data = source->Read(key); - if (data.empty()) return false; - - if (!dram->Write(key, data.data(), data.size())) { - // Follower mode never writes back to SSD. Evict DRAM-only victims. - while (true) { - std::string victim = SelectVictim(dram); - if (victim.empty()) return false; - std::vector group = GetGroup(victim); - for (const auto& k : group) { - if (dram->Exists(k)) { - dram->Evict(k); - if (index_) index_->Remove(k); - RemoveDepthAndGroup(k); - } - } - if (dram->Write(key, data.data(), data.size())) break; - } - } - - UpsertIndexTier(key, StorageTier::CPU_DRAM, data.size()); - return true; -} - -void LocalStorageManager::UpsertIndexTier(const std::string& key, StorageTier tier, - size_t size_hint) { - if (!index_) return; - if (!index_->UpdateTier(key, tier)) { - index_->Insert(key, {tier, 0, size_hint}); - } -} - -void LocalStorageManager::Clear() { - for (auto& entry : tiers_) { - entry.backend->Clear(); - } - if (index_) index_->Clear(); - std::unique_lock lock(depth_mu_); - depth_map_.clear(); - group_map_.clear(); -} - -bool LocalStorageManager::Flush() { - bool ok = true; - for (auto& entry : tiers_) { - if (!entry.backend->Flush()) ok = false; - } - return ok; -} - -std::vector LocalStorageManager::BatchWrite(const std::vector& keys, - const std::vector& data_ptrs, - const std::vector& sizes, - StorageTier tier) { - TierBackend* target = GetTier(tier); - if (!target) return std::vector(keys.size(), false); - - auto results = target->BatchWrite(keys, data_ptrs, sizes); - - // Fallback: retry failed items one-by-one with eviction support - for (size_t i = 0; i < results.size(); ++i) { - if (!results[i]) { - results[i] = Write(keys[i], data_ptrs[i], sizes[i], tier); - } - } - return results; -} - -std::vector LocalStorageManager::BatchReadIntoPtr(const std::vector& keys, - const std::vector& dst_ptrs, - const std::vector& sizes) { - std::vector results(keys.size(), false); - for (size_t i = 0; i < keys.size(); ++i) { - results[i] = ReadIntoPtr(keys[i], dst_ptrs[i], sizes[i]); - } - return results; -} - -} // namespace mori::umbp diff --git a/src/umbp/local/tiers/segment/segment_format.cpp b/src/umbp/local/tiers/segment/segment_format.cpp index 2aa30149f..7a25e6c49 100644 --- a/src/umbp/local/tiers/segment/segment_format.cpp +++ b/src/umbp/local/tiers/segment/segment_format.cpp @@ -21,22 +21,77 @@ // SOFTWARE. #include "umbp/local/tiers/segment/segment_format.h" +#include #include +#if defined(__x86_64__) || defined(__i386__) +#include +#define UMBP_SEGMENT_CRC32C_X86 1 +#endif + namespace mori::umbp::segment { -uint32_t CrcUpdate(const void* data, size_t size, uint32_t crc) { - const uint8_t* p = static_cast(data); - for (size_t i = 0; i < size; ++i) { - crc ^= static_cast(p[i]); - for (int j = 0; j < 8; ++j) { - const uint32_t mask = -(crc & 1u); - crc = (crc >> 1) ^ (0xEDB88320u & mask); +namespace { + +// CRC-32C (Castagnoli), reflected polynomial. This is the polynomial the SSE4.2 +// crc32 instruction implements, so the table path and the hardware path below +// produce identical checksums: a segment written on a host with hardware support +// still verifies on one without it. +constexpr uint32_t kCrc32cPoly = 0x82F63B78u; + +struct Crc32cTable { + uint32_t entry[256]; + + constexpr Crc32cTable() : entry() { + for (uint32_t i = 0; i < 256; ++i) { + uint32_t c = i; + for (int k = 0; k < 8; ++k) c = (c >> 1) ^ (kCrc32cPoly & (~(c & 1u) + 1u)); + entry[i] = c; } } +}; + +constexpr Crc32cTable kCrc32cTable{}; + +uint32_t CrcUpdateTable(const uint8_t* p, size_t size, uint32_t crc) { + for (size_t i = 0; i < size; ++i) crc = (crc >> 8) ^ kCrc32cTable.entry[(crc ^ p[i]) & 0xFFu]; return crc; } +#ifdef UMBP_SEGMENT_CRC32C_X86 +__attribute__((target("sse4.2"))) uint32_t CrcUpdateHw(const uint8_t* p, size_t size, + uint32_t crc) { + uint64_t acc = crc; + while (size >= sizeof(uint64_t)) { + uint64_t chunk; + std::memcpy(&chunk, p, sizeof(chunk)); + acc = _mm_crc32_u64(acc, chunk); + p += sizeof(uint64_t); + size -= sizeof(uint64_t); + } + uint32_t tail = static_cast(acc); + while (size--) tail = _mm_crc32_u8(tail, *p++); + return tail; +} + +bool DetectSse42() { + __builtin_cpu_init(); + return __builtin_cpu_supports("sse4.2"); +} + +const bool kHasSse42 = DetectSse42(); +#endif // UMBP_SEGMENT_CRC32C_X86 + +} // namespace + +uint32_t CrcUpdate(const void* data, size_t size, uint32_t crc) { + const uint8_t* p = static_cast(data); +#ifdef UMBP_SEGMENT_CRC32C_X86 + if (kHasSse42) return CrcUpdateHw(p, size, crc); +#endif + return CrcUpdateTable(p, size, crc); +} + uint32_t ComputeRecordCrc32(const std::string& key, const void* value, size_t value_size) { uint32_t crc = CrcUpdate(key.data(), key.size()); return ~CrcUpdate(value, value_size, crc); diff --git a/src/umbp/local/tiers/segment/segment_index.cpp b/src/umbp/local/tiers/segment/segment_index.cpp index 3f5d4bc35..626a9ee2f 100644 --- a/src/umbp/local/tiers/segment/segment_index.cpp +++ b/src/umbp/local/tiers/segment/segment_index.cpp @@ -42,41 +42,50 @@ void Index::RemoveLRU(const std::string& key) { } bool Index::PrepareWrite(const std::string& key, size_t size, size_t key_len, uint32_t crc32, - Meta* seg, WriteReservation* out) { + Meta* seg, WriteReservation* out, bool crc_valid) { if (!seg || !out) return false; + // Admission is decided on the padded on-disk footprint, which is what + // capacity_bytes is meant to bound; charging only `size` would let a store of + // small values overrun the drive by the padding ratio. + const size_t on_disk = static_cast(RecordBytes(key_len, size)); auto existing = key_meta_.find(key); - size_t existing_size = (existing == key_meta_.end()) ? 0 : existing->second.size; + size_t existing_size = (existing == key_meta_.end()) ? 0 : existing->second.disk_bytes; size_t effective_used = (used_ >= existing_size) ? (used_ - existing_size) : 0; - if (effective_used + size > capacity_) return false; + if (effective_used + on_disk > capacity_) return false; const uint64_t generation = generation_counter_++; out->key = key; out->had_previous = (existing != key_meta_.end()); if (out->had_previous) out->previous_meta = existing->second; - const size_t record_size = sizeof(RecordHeader) + key_len + size; + // v3 layout: the append cursor is always on a kRecordAlign boundary, so + // starting the record at write_offset already gives an aligned record_offset; + // the padded prefix then puts the value on an aligned boundary too. Both are + // preconditions for reading or writing this record with O_DIRECT. out->record_offset = seg->write_offset; out->previous_write_offset = seg->write_offset; - seg->write_offset += static_cast(record_size); + seg->write_offset += RecordBytes(key_len, size); out->meta.segment_id = seg->id; - out->meta.value_offset = out->record_offset + sizeof(RecordHeader) + key_len; + out->meta.value_offset = out->record_offset + PrefixBytes(key_len); out->meta.size = static_cast(size); out->meta.crc32 = crc32; out->meta.generation = generation; + out->meta.crc_valid = crc_valid; + out->meta.disk_bytes = static_cast(on_disk); if (existing != key_meta_.end()) { auto old_seg = segments_.find(existing->second.segment_id); - if (old_seg != segments_.end() && old_seg->second.live_bytes >= existing->second.size) { - old_seg->second.live_bytes -= existing->second.size; + if (old_seg != segments_.end() && old_seg->second.live_bytes >= existing->second.disk_bytes) { + old_seg->second.live_bytes -= existing->second.disk_bytes; } - used_ -= existing->second.size; + used_ -= existing->second.disk_bytes; } - seg->live_bytes += size; + seg->live_bytes += on_disk; key_meta_[key] = out->meta; - used_ += size; + used_ += on_disk; TouchLRU(key); return true; } @@ -86,12 +95,12 @@ void Index::RollbackWrite(const WriteReservation& reservation) { if (it != key_meta_.end() && it->second.generation == reservation.meta.generation) { auto seg = segments_.find(reservation.meta.segment_id); if (seg != segments_.end()) { - if (seg->second.live_bytes >= reservation.meta.size) { - seg->second.live_bytes -= reservation.meta.size; + if (seg->second.live_bytes >= reservation.meta.disk_bytes) { + seg->second.live_bytes -= reservation.meta.disk_bytes; } seg->second.write_offset = reservation.previous_write_offset; } - if (used_ >= reservation.meta.size) used_ -= reservation.meta.size; + if (used_ >= reservation.meta.disk_bytes) used_ -= reservation.meta.disk_bytes; key_meta_.erase(it); RemoveLRU(reservation.key); } @@ -99,8 +108,9 @@ void Index::RollbackWrite(const WriteReservation& reservation) { if (reservation.had_previous) { key_meta_[reservation.key] = reservation.previous_meta; auto old_seg = segments_.find(reservation.previous_meta.segment_id); - if (old_seg != segments_.end()) old_seg->second.live_bytes += reservation.previous_meta.size; - used_ += reservation.previous_meta.size; + if (old_seg != segments_.end()) + old_seg->second.live_bytes += reservation.previous_meta.disk_bytes; + used_ += reservation.previous_meta.disk_bytes; TouchLRU(reservation.key); } } @@ -110,16 +120,16 @@ void Index::RecordRecoveredEntry(const std::string& key, const KeyMeta& meta) { if (existing != key_meta_.end()) { if (existing->second.generation >= meta.generation) return; auto old_seg = segments_.find(existing->second.segment_id); - if (old_seg != segments_.end() && old_seg->second.live_bytes >= existing->second.size) { - old_seg->second.live_bytes -= existing->second.size; + if (old_seg != segments_.end() && old_seg->second.live_bytes >= existing->second.disk_bytes) { + old_seg->second.live_bytes -= existing->second.disk_bytes; } - if (used_ >= existing->second.size) used_ -= existing->second.size; + if (used_ >= existing->second.disk_bytes) used_ -= existing->second.disk_bytes; } auto seg = segments_.find(meta.segment_id); - if (seg != segments_.end()) seg->second.live_bytes += meta.size; + if (seg != segments_.end()) seg->second.live_bytes += meta.disk_bytes; key_meta_[key] = meta; - used_ += meta.size; + used_ += meta.disk_bytes; TouchLRU(key); ObserveGeneration(meta.generation); } @@ -129,10 +139,10 @@ bool Index::EraseKey(const std::string& key) { if (it == key_meta_.end()) return false; auto seg = segments_.find(it->second.segment_id); - if (seg != segments_.end() && seg->second.live_bytes >= it->second.size) { - seg->second.live_bytes -= it->second.size; + if (seg != segments_.end() && seg->second.live_bytes >= it->second.disk_bytes) { + seg->second.live_bytes -= it->second.disk_bytes; } - if (used_ >= it->second.size) used_ -= it->second.size; + if (used_ >= it->second.disk_bytes) used_ -= it->second.disk_bytes; key_meta_.erase(it); RemoveLRU(key); return true; diff --git a/src/umbp/local/tiers/segment/segment_scanner.cpp b/src/umbp/local/tiers/segment/segment_scanner.cpp index 3bf970378..aefbea1e5 100644 --- a/src/umbp/local/tiers/segment/segment_scanner.cpp +++ b/src/umbp/local/tiers/segment/segment_scanner.cpp @@ -27,10 +27,12 @@ #include #include +#include #include #include #include +#include "umbp/common/aligned_buffer.h" #include "umbp/local/tiers/segment/segment_format.h" namespace fs = std::filesystem; @@ -39,7 +41,7 @@ namespace mori::umbp::segment { bool Scanner::RefreshFromDisk(const std::string& dir, StorageIoDriver& io_driver, Index& index, bool read_only_shared, bool force_full_rescan, - std::string* error_message) const { + std::string* error_message, int extra_open_flags) const { if (force_full_rescan) index.ResetMetadata(); std::vector ids; @@ -64,7 +66,8 @@ bool Scanner::RefreshFromDisk(const std::string& dir, StorageIoDriver& io_driver Meta seg; seg.id = sid; seg.path = dir + "/" + BuildFileName(sid); - seg.fd = open(seg.path.c_str(), read_only_shared ? O_RDONLY : (O_RDWR | O_CREAT), 0644); + seg.fd = open(seg.path.c_str(), + (read_only_shared ? O_RDONLY : (O_RDWR | O_CREAT)) | extra_open_flags, 0644); if (seg.fd < 0) { if (error_message) *error_message = "failed to open segment " + seg.path; return false; @@ -92,30 +95,48 @@ bool Scanner::RefreshFromDisk(const std::string& dir, StorageIoDriver& io_driver uint64_t file_size = static_cast(st.st_size); uint64_t offset = force_full_rescan ? 0 : seg.scanned_offset; - while (offset + sizeof(RecordHeader) <= file_size) { + // Header and key live in the record's aligned prefix, so one kRecordAlign + // block covers both for any key short enough to fit — always true for the + // hashed keys this tier stores, with a re-read for the pathological case. + // Reading whole aligned blocks into an aligned buffer is what keeps this + // loop legal when the fd carries O_DIRECT. + AlignedBuffer prefix(static_cast(kRecordAlign)); + + while (offset + kRecordAlign <= file_size) { + IoStatus blk = io_driver.ReadAt(seg.fd, prefix.data(), prefix.padded_size(), offset); + if (!blk.ok()) break; + RecordHeader hdr; - IoStatus hdr_status = io_driver.ReadAt(seg.fd, &hdr, sizeof(hdr), offset); - if (!hdr_status.ok()) break; + std::memcpy(&hdr, prefix.data(), sizeof(hdr)); if (hdr.magic != kRecordMagic || hdr.version != kRecordVersion || hdr.key_len == 0) break; - const uint64_t rec_size = - sizeof(RecordHeader) + static_cast(hdr.key_len) + hdr.value_size; + + const uint64_t prefix_bytes = PrefixBytes(hdr.key_len); + const uint64_t rec_size = RecordBytes(hdr.key_len, hdr.value_size); if (offset + rec_size > file_size) break; + + // Key spills past the first block (key_len > kRecordAlign - 32): widen the + // buffer and re-read the whole prefix, still aligned on both ends. + if (prefix_bytes > prefix.padded_size()) { + prefix.Resize(static_cast(prefix_bytes)); + IoStatus wide = io_driver.ReadAt(seg.fd, prefix.data(), prefix.padded_size(), offset); + if (!wide.ok()) break; + } + if ((hdr.flags & kFlagCommitted) == 0) { offset += rec_size; continue; } - std::string key; - key.resize(hdr.key_len); - IoStatus key_status = io_driver.ReadAt(seg.fd, key.data(), hdr.key_len, offset + sizeof(hdr)); - if (!key_status.ok()) break; + const std::string key(prefix.data() + sizeof(hdr), hdr.key_len); KeyMeta meta; meta.segment_id = seg.id; - meta.value_offset = offset + sizeof(hdr) + hdr.key_len; + meta.value_offset = offset + prefix_bytes; meta.size = hdr.value_size; meta.crc32 = hdr.crc32; meta.generation = hdr.generation; + meta.crc_valid = (hdr.flags & kFlagNoCrc) == 0; + meta.disk_bytes = static_cast(rec_size); index.RecordRecoveredEntry(key, meta); offset += rec_size; } diff --git a/src/umbp/local/tiers/segment/segment_writer.cpp b/src/umbp/local/tiers/segment/segment_writer.cpp index 36a0d267f..f5e3431a6 100644 --- a/src/umbp/local/tiers/segment/segment_writer.cpp +++ b/src/umbp/local/tiers/segment/segment_writer.cpp @@ -21,54 +21,92 @@ // SOFTWARE. #include "umbp/local/tiers/segment/segment_writer.h" -namespace mori::umbp::segment { +#include -bool Writer::Prepare(const std::string& key, const void* data, size_t size, Meta* segment_meta, - Index& index, PreparedRecord* out) const { - if (!segment_meta || !out) return false; +#include "umbp/common/ssd_perf.h" - const size_t record_size = sizeof(RecordHeader) + key.size() + size; - const uint32_t crc32 = ComputeRecordCrc32(key, data, size); +namespace mori::umbp::segment { - if (!index.PrepareWrite(key, size, key.size(), crc32, segment_meta, &out->reservation)) - return false; +void Writer::Build(const std::string& key, const void* data, size_t size, + PreparedRecord* out) const { + if (!out) return; + + out->crc_valid = compute_crc_; + out->crc32 = compute_crc_ ? ComputeRecordCrc32(key, data, size) : 0; RecordHeader hdr; hdr.magic = kRecordMagic; hdr.version = kRecordVersion; - hdr.flags = kFlagCommitted; + hdr.flags = kFlagCommitted | (compute_crc_ ? 0 : kFlagNoCrc); hdr.key_len = static_cast(key.size()); hdr.value_size = static_cast(size); - hdr.crc32 = crc32; - hdr.generation = out->reservation.meta.generation; + hdr.crc32 = out->crc32; + hdr.generation = 0; // stamped by Reserve, which allocates the generation - out->record.resize(record_size); + // v3 layout: [header|key|pad] [value|pad], both parts alignment multiples, so + // the value lands on an aligned offset and the next record starts aligned. + // Resize() zeroes the padding, so no heap contents reach the device. + const uint64_t prefix = PrefixBytes(key.size()); + out->record.Resize(static_cast(RecordBytes(key.size(), size))); std::memcpy(out->record.data(), &hdr, sizeof(hdr)); std::memcpy(out->record.data() + sizeof(hdr), key.data(), key.size()); - std::memcpy(out->record.data() + sizeof(hdr) + key.size(), data, size); + std::memcpy(out->record.data() + prefix, data, size); +} + +bool Writer::Reserve(const std::string& key, size_t size, Meta* segment_meta, Index& index, + PreparedRecord* out) const { + if (!segment_meta || !out) return false; + if (out->record.size() < sizeof(RecordHeader)) return false; // Build() not run + + if (!index.PrepareWrite(key, size, key.size(), out->crc32, segment_meta, &out->reservation, + out->crc_valid)) + return false; + + // Patch the one header field the reservation produces, in place, instead of + // rebuilding the whole record. + const uint64_t generation = out->reservation.meta.generation; + std::memcpy(out->record.data() + offsetof(RecordHeader, generation), &generation, + sizeof(generation)); return true; } -IoStatus Writer::WriteRecord(int fd, const PreparedRecord& pr, bool should_sync) const { - IoStatus status = - io_driver_.WriteAt(fd, pr.record.data(), pr.record.size(), pr.reservation.record_offset); +bool Writer::Prepare(const std::string& key, const void* data, size_t size, Meta* segment_meta, + Index& index, PreparedRecord* out) const { + if (!segment_meta || !out) return false; + Build(key, data, size, out); + return Reserve(key, size, segment_meta, index, out); +} + +IoStatus Writer::WriteRecord(int fd, const PreparedRecord& pr, bool should_sync, + double* sync_ms_out) const { + // padded_size(), not size(): under O_DIRECT the length must be an alignment + // multiple. Build() already sized the record to one, so these agree — the + // explicit call documents the requirement and survives a future layout change. + IoStatus status = io_driver_.WriteAt(fd, pr.record.data(), pr.record.padded_size(), + pr.reservation.record_offset); if (!status.ok()) return status; - if (should_sync) return io_driver_.Sync(fd); - return IoStatus::Ok(); + if (!should_sync) return IoStatus::Ok(); + const auto t_sync = ssdperf::Now(); + status = io_driver_.Sync(fd); + if (sync_ms_out) *sync_ms_out += ssdperf::MsSince(t_sync); + return status; } -IoStatus Writer::WriteRecords(int fd, const std::vector& records, - bool should_sync) const { +IoStatus Writer::WriteRecords(int fd, const std::vector& records, bool should_sync, + double* sync_ms_out) const { std::vector ops; ops.reserve(records.size()); for (const auto& pr : records) { - ops.push_back({fd, pr.record.data(), pr.record.size(), pr.reservation.record_offset}); + ops.push_back({fd, pr.record.data(), pr.record.padded_size(), pr.reservation.record_offset}); } IoStatus status = io_driver_.WriteBatch(ops); if (!status.ok()) return status; - if (should_sync) return io_driver_.SyncMany({fd}); - return IoStatus::Ok(); + if (!should_sync) return IoStatus::Ok(); + const auto t_sync = ssdperf::Now(); + status = io_driver_.SyncMany({fd}); + if (sync_ms_out) *sync_ms_out += ssdperf::MsSince(t_sync); + return status; } } // namespace mori::umbp::segment diff --git a/src/umbp/local/tiers/sharded_ssd_tier.cpp b/src/umbp/local/tiers/sharded_ssd_tier.cpp new file mode 100644 index 000000000..b9e409059 --- /dev/null +++ b/src/umbp/local/tiers/sharded_ssd_tier.cpp @@ -0,0 +1,445 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/local/tiers/sharded_ssd_tier.h" + +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/common/parallel_for.h" +#include "umbp/common/ssd_perf.h" + +namespace mori::umbp { + +namespace { + +// Gather each shard's sub-batch, run `op` on it in parallel across drives, and +// scatter the per-key results back. Shared by the put and get paths, which +// differ only in the pointer type and the backend call. Buckets touch disjoint +// shards and disjoint result indices, so nothing here needs a lock. +template +void FanOutBatches(const std::vector>& buckets, const std::vector& busy, + int io_threads, const std::vector& keys, + const std::vector& ptrs, const std::vector& sizes, + std::vector& results, std::vector& shard_ms, Op&& op) { + ParallelFor(busy.size(), io_threads, [&](size_t bi) { + const int s = busy[bi]; + const auto t_shard = ssdperf::Now(); + const auto& idxs = buckets[s]; + std::vector k; + std::vector d; + std::vector z; + k.reserve(idxs.size()); + d.reserve(idxs.size()); + z.reserve(idxs.size()); + for (size_t i : idxs) { + k.push_back(keys[i]); + d.push_back(ptrs[i]); + z.push_back(sizes[i]); + } + const auto shard_res = op(s, k, d, z); + for (size_t j = 0; j < idxs.size(); ++j) { + results[idxs[j]] = j < shard_res.size() && shard_res[j]; + } + shard_ms[s] = ssdperf::MsSince(t_shard); // disjoint shards: no race on shard_ms + }); +} + +// Indices of the shards this batch actually touches. +std::vector BusyShards(const std::vector>& buckets) { + std::vector busy; + for (size_t s = 0; s < buckets.size(); ++s) { + if (!buckets[s].empty()) busy.push_back(static_cast(s)); + } + return busy; +} + +// "s0:keys=64/32MiB/12.4ms ...". Imbalanced keys = placement is not spreading +// the batch; per-shard times well below wall time = the drives are not +// overlapping (too few io_threads, or one straggler). +std::string FormatShardFanout(const std::vector& busy, + const std::vector>& buckets, + const std::vector& sizes, const std::vector& shard_ms, + size_t num_shards) { + std::ostringstream oss; + for (size_t s = 0; s < num_shards; ++s) { + if (s) oss << ' '; + uint64_t bytes = 0; + for (size_t i : buckets[s]) bytes += sizes[i]; + oss << 's' << s << ":keys=" << buckets[s].size() << '/' << (bytes >> 20) << "MiB/"; + oss.setf(std::ios::fixed); + oss.precision(2); + oss << shard_ms[s] << "ms"; + oss.unsetf(std::ios::fixed); + } + if (busy.empty()) oss << " "; + return oss.str(); +} + +} // namespace + +ShardedSsdTier::ShardedSsdTier(std::vector> shards, int io_threads) + : TierBackend(StorageTier::LOCAL_SSD), shards_(std::move(shards)) { + if (shards_.empty()) { + throw std::runtime_error("[ShardedSsdTier] requires at least one shard"); + } + for (const auto& s : shards_) { + if (!s) throw std::runtime_error("[ShardedSsdTier] null shard backend"); + } + shard_total_.reserve(shards_.size()); + used_hint_.assign(shards_.size(), 0); + for (size_t i = 0; i < shards_.size(); ++i) { + auto cap = shards_[i]->Capacity(); + used_hint_[i] = cap.first; + shard_total_.push_back(cap.second); + } + io_threads_ = io_threads > 0 ? io_threads : static_cast(shards_.size()); + MORI_UMBP_INFO("[ShardedSsdTier] {} shard(s), io_threads={}", shards_.size(), io_threads_); +} + +ShardedSsdTier::~ShardedSsdTier() = default; + +int ShardedSsdTier::LookupLocked(const std::string& key) const { + auto it = key_shard_.find(key); + return it == key_shard_.end() ? -1 : static_cast(it->second.shard); +} + +int ShardedSsdTier::PickShardLocked(size_t size) const { + int best = -1; + size_t best_free = 0; + for (size_t i = 0; i < shards_.size(); ++i) { + const size_t used = used_hint_[i]; + if (used >= shard_total_[i]) continue; + const size_t free_bytes = shard_total_[i] - used; + if (free_bytes < size) continue; + if (best < 0 || free_bytes > best_free) { + best = static_cast(i); + best_free = free_bytes; + } + } + return best; +} + +void ShardedSsdTier::ReserveLocked(int shard, size_t size) { used_hint_[shard] += size; } + +void ShardedSsdTier::ReleaseLocked(int shard, size_t size) { + used_hint_[shard] = used_hint_[shard] > size ? used_hint_[shard] - size : 0; +} + +int ShardedSsdTier::ShardOf(const std::string& key) const { + std::shared_lock lock(map_mu_); + return LookupLocked(key); +} + +std::vector> ShardedSsdTier::PerShardCapacity() const { + std::vector> out; + out.reserve(shards_.size()); + for (const auto& s : shards_) out.push_back(s->Capacity()); + return out; +} + +bool ShardedSsdTier::Write(const std::string& key, const void* data, size_t size) { + int shard; + bool is_new; + { + std::unique_lock lock(map_mu_); + shard = LookupLocked(key); + is_new = shard < 0; + if (is_new) { + shard = PickShardLocked(size); + if (shard < 0) { + MORI_UMBP_WARN("[ShardedSsdTier] Write key={} size={}: no shard has room", key, size); + return false; + } + // Reserve before dropping the lock so concurrent writers see the fill and + // spread out instead of all piling onto the same "most free" shard. + key_shard_[key] = KeyLoc{static_cast(shard), size}; + ReserveLocked(shard, size); + } + } + + const bool ok = shards_[shard]->Write(key, data, size); + if (!ok && is_new) { + std::unique_lock lock(map_mu_); + key_shard_.erase(key); + ReleaseLocked(shard, size); + } + return ok; +} + +bool ShardedSsdTier::ReadIntoPtr(const std::string& key, uintptr_t dst_ptr, size_t size) { + const int shard = ShardOf(key); + return shard >= 0 && shards_[shard]->ReadIntoPtr(key, dst_ptr, size); +} + +bool ShardedSsdTier::Exists(const std::string& key) const { return ShardOf(key) >= 0; } + +bool ShardedSsdTier::Evict(const std::string& key) { + int shard; + size_t freed = 0; + { + std::shared_lock lock(map_mu_); + auto it = key_shard_.find(key); + if (it == key_shard_.end()) return false; + shard = static_cast(it->second.shard); + freed = it->second.size; + } + if (!shards_[shard]->Evict(key)) return false; + std::unique_lock lock(map_mu_); + key_shard_.erase(key); + ReleaseLocked(shard, freed); + return true; +} + +std::pair ShardedSsdTier::Capacity() const { + size_t used = 0, total = 0; + for (const auto& s : shards_) { + auto cap = s->Capacity(); + used += cap.first; + total += cap.second; + } + return {used, total}; +} + +void ShardedSsdTier::Clear() { + for (auto& s : shards_) s->Clear(); + std::unique_lock lock(map_mu_); + key_shard_.clear(); + std::fill(used_hint_.begin(), used_hint_.end(), 0); +} + +std::vector ShardedSsdTier::Read(const std::string& key) { + const int shard = ShardOf(key); + return shard < 0 ? std::vector{} : shards_[shard]->Read(key); +} + +TierCapabilities ShardedSsdTier::Capabilities() const { + // Batch capability is this class's own; zero-copy is only advertised when + // every shard can do it. + TierCapabilities caps; + caps.batch_write = true; + caps.batch_read = true; + caps.zero_copy_read = true; + for (const auto& s : shards_) { + if (!s->Capabilities().zero_copy_read) caps.zero_copy_read = false; + } + return caps; +} + +std::string ShardedSsdTier::GetLRUKey() const { + // No globally exact LRU: shards keep independent recency lists. Take the + // oldest key of the fullest shard — the one eviction must relieve anyway. + // PeerSsdManager drives distributed eviction from its own global LRU. + int fullest = -1; + size_t most_used = 0; + for (size_t i = 0; i < shards_.size(); ++i) { + auto cap = shards_[i]->Capacity(); + if (fullest < 0 || cap.first > most_used) { + fullest = static_cast(i); + most_used = cap.first; + } + } + return fullest < 0 ? std::string() : shards_[fullest]->GetLRUKey(); +} + +std::vector ShardedSsdTier::GetLRUCandidates(size_t max_candidates) const { + // Round-robin merge of the per-shard tails, so victims are spread across + // drives instead of emptying one. + std::vector> per_shard; + per_shard.reserve(shards_.size()); + for (const auto& s : shards_) per_shard.push_back(s->GetLRUCandidates(max_candidates)); + + std::vector out; + out.reserve(max_candidates); + for (size_t round = 0; out.size() < max_candidates; ++round) { + bool progressed = false; + for (auto& cands : per_shard) { + if (round >= cands.size()) continue; + out.push_back(cands[round]); + progressed = true; + if (out.size() >= max_candidates) break; + } + if (!progressed) break; + } + return out; +} + +std::optional ShardedSsdTier::GetLocationId(const std::string& key) const { + const int shard = ShardOf(key); + if (shard < 0) return std::nullopt; + auto inner = shards_[shard]->GetLocationId(key); + if (!inner) return std::nullopt; + // Prefix the drive so a location string identifies which device holds the key. + return "s" + std::to_string(shard) + ":" + *inner; +} + +std::optional ShardedSsdTier::LocateRecord(const std::string& key) const { + const int shard = ShardOf(key); + if (shard < 0) return std::nullopt; + return shards_[shard]->LocateRecord(key); +} + +bool ShardedSsdTier::Flush() { + bool ok = true; + for (auto& s : shards_) ok = s->Flush() && ok; + return ok; +} + +void ShardedSsdTier::SetColdRead(bool enable) { + for (auto& s : shards_) s->SetColdRead(enable); +} + +std::vector ShardedSsdTier::BatchWrite(const std::vector& keys, + const std::vector& data_ptrs, + const std::vector& sizes) { + std::vector results(keys.size(), false); + if (keys.empty()) return results; + if (data_ptrs.size() != keys.size() || sizes.size() != keys.size()) { + MORI_UMBP_ERROR("[ShardedSsdTier] BatchWrite: mismatched input lengths ({} / {} / {})", + keys.size(), data_ptrs.size(), sizes.size()); + return results; + } + + // Assign every key to a shard under one lock, applying the projected + // reservation as we go so a batch of same-size values spreads evenly. + std::vector> buckets(shards_.size()); + std::vector is_new(keys.size(), 0); + { + std::unique_lock lock(map_mu_); + for (size_t i = 0; i < keys.size(); ++i) { + int shard = LookupLocked(keys[i]); + if (shard < 0) { + shard = PickShardLocked(sizes[i]); + if (shard < 0) continue; // leaves results[i] == false + key_shard_[keys[i]] = KeyLoc{static_cast(shard), sizes[i]}; + ReserveLocked(shard, sizes[i]); + is_new[i] = 1; + } + buckets[shard].push_back(i); + } + } + + const std::vector busy = BusyShards(buckets); + const auto t_fanout = ssdperf::Now(); + std::vector shard_ms(shards_.size(), 0.0); + + FanOutBatches(buckets, busy, io_threads_, keys, data_ptrs, sizes, results, shard_ms, + [&](int s, const auto& k, const auto& d, const auto& z) { + return shards_[s]->BatchWrite(k, d, z); + }); + + if (ssdperf::Enabled()) { + const double total_ms = ssdperf::MsSince(t_fanout); + uint64_t bytes = 0; + for (size_t z : sizes) bytes += z; + double slowest = 0.0; + for (double m : shard_ms) slowest = std::max(slowest, m); + MORI_UMBP_INFO( + "[SsdPerf/shard] PUT drives={} io_threads={} busy={} keys={} bytes={} total_ms={:.3f} " + "GB_s={:.2f} slowest_drive_ms={:.3f} overlap={:.2f}x | {}", + shards_.size(), io_threads_, busy.size(), keys.size(), bytes, total_ms, + ssdperf::GbPerSec(bytes, total_ms), slowest, + total_ms > 0.0 ? std::accumulate(shard_ms.begin(), shard_ms.end(), 0.0) / total_ms : 0.0, + FormatShardFanout(busy, buckets, sizes, shard_ms, shards_.size())); + } + + // Roll back the map/hint for newly-assigned keys whose write failed, so a + // retry is free to pick a different drive. + std::unique_lock lock(map_mu_); + for (size_t i = 0; i < keys.size(); ++i) { + if (!results[i] && is_new[i]) { + auto it = key_shard_.find(keys[i]); + if (it != key_shard_.end()) { + ReleaseLocked(static_cast(it->second.shard), it->second.size); + key_shard_.erase(it); + } + } + } + return results; +} + +std::vector ShardedSsdTier::BatchReadIntoPtr(const std::vector& keys, + const std::vector& dst_ptrs, + const std::vector& sizes) { + std::vector results(keys.size(), false); + if (keys.empty()) return results; + if (dst_ptrs.size() != keys.size() || sizes.size() != keys.size()) { + MORI_UMBP_ERROR("[ShardedSsdTier] BatchReadIntoPtr: mismatched input lengths ({} / {} / {})", + keys.size(), dst_ptrs.size(), sizes.size()); + return results; + } + + std::vector> buckets(shards_.size()); + { + std::shared_lock lock(map_mu_); + for (size_t i = 0; i < keys.size(); ++i) { + int shard = LookupLocked(keys[i]); + if (shard >= 0) buckets[shard].push_back(i); + } + } + + const std::vector busy = BusyShards(buckets); + const auto t_fanout = ssdperf::Now(); + std::vector shard_ms(shards_.size(), 0.0); + + FanOutBatches(buckets, busy, io_threads_, keys, dst_ptrs, sizes, results, shard_ms, + [&](int s, const auto& k, const auto& d, const auto& z) { + return shards_[s]->BatchReadIntoPtr(k, d, z); + }); + + if (ssdperf::Enabled()) { + const double total_ms = ssdperf::MsSince(t_fanout); + uint64_t bytes = 0, missing = 0; + for (size_t i = 0; i < keys.size(); ++i) bytes += sizes[i]; + for (const auto& b : buckets) missing += b.size(); + missing = keys.size() - missing; + double slowest = 0.0; + for (double m : shard_ms) slowest = std::max(slowest, m); + // overlap = sum(per-drive time) / wall time. ~1.0 means the drives ran one + // after another (no parallelism); ~N means all N drives were busy at once. + MORI_UMBP_INFO( + "[SsdPerf/shard] GET drives={} io_threads={} busy={} keys={} unmapped={} bytes={} " + "total_ms={:.3f} GB_s={:.2f} slowest_drive_ms={:.3f} overlap={:.2f}x | {}", + shards_.size(), io_threads_, busy.size(), keys.size(), missing, bytes, total_ms, + ssdperf::GbPerSec(bytes, total_ms), slowest, + total_ms > 0.0 ? std::accumulate(shard_ms.begin(), shard_ms.end(), 0.0) / total_ms : 0.0, + FormatShardFanout(busy, buckets, sizes, shard_ms, shards_.size())); + } + return results; +} + +bool ShardedSsdTier::WriteBatch(const std::vector& keys, + const std::vector& data_ptrs, + const std::vector& sizes) { + auto res = BatchWrite(keys, data_ptrs, sizes); + return std::all_of(res.begin(), res.end(), [](bool ok) { return ok; }); +} + +std::vector ShardedSsdTier::ReadBatchIntoPtr(const std::vector& keys, + const std::vector& dst_ptrs, + const std::vector& sizes) { + return BatchReadIntoPtr(keys, dst_ptrs, sizes); +} + +} // namespace mori::umbp diff --git a/src/umbp/local/tiers/ssd_backend_factory.cpp b/src/umbp/local/tiers/ssd_backend_factory.cpp new file mode 100644 index 000000000..69b9dcb12 --- /dev/null +++ b/src/umbp/local/tiers/ssd_backend_factory.cpp @@ -0,0 +1,65 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/local/tiers/ssd_backend_factory.h" + +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/local/tiers/sharded_ssd_tier.h" + +namespace mori::umbp { + +std::unique_ptr MakeFileSsdBackend(const UMBPSsdConfig& ssd_config, + SSDAccessMode access_mode) { + const std::vector dirs = ssd_config.StorageDirs(); + // Surface the resolved durability mode once: "strict" fdatasync()s every batch + // write and, on a real drive, dominates write time. UMBP_SSD_DURABILITY. + MORI_UMBP_INFO("[MakeFileSsdBackend] drives={} durability={} io_backend={} queue_depth={}", + dirs.size(), + ssd_config.durability.mode == UMBPDurabilityMode::Strict ? "strict" : "relaxed", + ssd_config.io.backend == UMBPIoBackend::IoUring ? "io_uring" : "posix", + ssd_config.io.queue_depth); + if (dirs.size() == 1) { + return std::make_unique(dirs[0], ssd_config.capacity_bytes, ssd_config, access_mode); + } + + // Split the budget evenly; the remainder goes to the first shard so the sum + // matches capacity_bytes exactly and the master's reported total is right. + const size_t per_shard = ssd_config.capacity_bytes / dirs.size(); + const size_t remainder = ssd_config.capacity_bytes % dirs.size(); + if (per_shard == 0) { + throw std::runtime_error("[MakeFileSsdBackend] ssd.capacity_bytes too small to split across " + + std::to_string(dirs.size()) + " directories"); + } + + std::vector> shards; + shards.reserve(dirs.size()); + for (size_t i = 0; i < dirs.size(); ++i) { + const size_t cap = per_shard + (i == 0 ? remainder : 0); + MORI_UMBP_INFO("[MakeFileSsdBackend] shard {} dir={} capacity={}B", i, dirs[i], cap); + shards.push_back(std::make_unique(dirs[i], cap, ssd_config, access_mode)); + } + return std::make_unique(std::move(shards), ssd_config.shard_io_threads); +} + +} // namespace mori::umbp diff --git a/src/umbp/local/tiers/ssd_tier.cpp b/src/umbp/local/tiers/ssd_tier.cpp index 90266acc4..b8160c808 100644 --- a/src/umbp/local/tiers/ssd_tier.cpp +++ b/src/umbp/local/tiers/ssd_tier.cpp @@ -25,12 +25,21 @@ #include #include +#include #include +#include #include +#include #include +#include #include #include "mori/utils/mori_log.hpp" +#include "umbp/common/aligned_buffer.h" +#include "umbp/common/device_copy.h" +#include "umbp/common/parallel_for.h" +#include "umbp/common/range_utils.h" +#include "umbp/common/ssd_perf.h" #include "umbp/local/tiers/segment/segment_format.h" namespace fs = std::filesystem; @@ -57,15 +66,126 @@ SSDTier::SSDTier(const std::string& dir, size_t capacity, const UMBPSsdConfig& s "SSDTier: io_uring backend requested but unavailable (kernel missing io_uring " "support or restricted by seccomp/capabilities); falling back to POSIX backend"); } - writer_ = std::make_unique(*io_driver_); + writer_ = std::make_unique(*io_driver_, ssd_config_.verify_crc); fs::create_directories(dir_); + + if (ssd_config_.direct_io) { + direct_io_ = ProbeDirectIo(); + if (!direct_io_) { + MORI_UMBP_WARN( + "[SSDTier] {}: direct I/O requested but this filesystem rejects O_DIRECT at {}B " + "alignment (tmpfs/overlayfs cannot do it); falling back to buffered I/O. Reads will " + "be served from the page cache, so reported tier bandwidth will NOT be the device's", + dir_, segment::kRecordAlign); + } + } + MORI_UMBP_INFO("[SSDTier] {}: direct_io={} verify_crc={} tier_io_threads={} io_backend={}", dir_, + direct_io_, ssd_config_.verify_crc, TierThreads(), + ssd_config_.io.backend == UMBPIoBackend::IoUring ? "io_uring" : "posix"); + if (direct_io_ && ssd_config_.io.backend != UMBPIoBackend::IoUring) { + MORI_UMBP_WARN( + "[SSDTier] {}: direct I/O on the POSIX backend issues one blocking pread per key " + "(queue depth 1). Every read is now a synchronous device round trip, which will look " + "far slower than it needs to — prefer ssd.io.backend=io_uring for direct I/O", + dir_); + } + std::lock_guard lock(mu_); RefreshFromDiskLocked(true); + DropUnparsedTailsLocked(); if (!IsReadOnlyShared() && index_.Segments().empty()) { OpenOrCreateSegmentLocked(0); } } +int SSDTier::TierThreads() const { + int t = ssd_config_.tier_io_threads; + if (t < 1) t = 1; + const unsigned hc = std::thread::hardware_concurrency(); + if (hc > 0 && t > static_cast(hc)) t = static_cast(hc); + return t; +} + +int SSDTier::SegmentOpenFlags() const { return direct_io_ ? O_DIRECT : 0; } + +bool SSDTier::ProbeDirectIo() const { + const std::string path = dir_ + "/.umbp_odirect_probe"; + // O_DIRECT is rejected at open() time on filesystems that cannot do it, but a + // successful open does not prove the alignment is accepted, so round-trip one + // real block before believing it. + int fd = open(path.c_str(), O_RDWR | O_CREAT | O_TRUNC | O_DIRECT, 0644); + if (fd < 0) return false; + + bool ok = false; + try { + AlignedBuffer block(static_cast(segment::kRecordAlign)); + std::memset(block.data(), 0xA5, block.padded_size()); + const ssize_t w = pwrite(fd, block.data(), block.padded_size(), 0); + if (w == static_cast(block.padded_size())) { + AlignedBuffer back(static_cast(segment::kRecordAlign)); + const ssize_t r = pread(fd, back.data(), back.padded_size(), 0); + ok = r == static_cast(back.padded_size()) && + std::memcmp(block.data(), back.data(), block.padded_size()) == 0; + } + } catch (const std::bad_alloc&) { + ok = false; + } + + close(fd); + unlink(path.c_str()); + return ok; +} + +void SSDTier::SetColdRead(bool enable) { + if (enable == direct_io_) return; + if (enable && !ProbeDirectIo()) { + MORI_UMBP_WARN("[SSDTier] {}: SetColdRead(true) ignored — filesystem rejects O_DIRECT", dir_); + return; + } + direct_io_ = enable; + MORI_UMBP_INFO("[SSDTier] {}: direct I/O {} for segments opened from here on", dir_, + enable ? "enabled" : "disabled"); +} + +void SSDTier::DropUnparsedTailsLocked() { + // A follower must never rewrite the shared log; the owner does this once at + // startup, before any write can land. + if (IsReadOnlyShared()) return; + + for (auto& kv : index_.MutableSegments()) { + auto& seg = kv.second; + if (seg.fd < 0) continue; + + struct stat st; + if (fstat(seg.fd, &st) != 0) continue; + const uint64_t file_size = static_cast(st.st_size); + if (seg.scanned_offset >= file_size) continue; // fully parsed, nothing to do + + // Bytes past the last parsed record boundary are permanently unreachable: + // the scanner stops at the first header it cannot read and every later + // restart stops at the same place. Two ways to get here -- + // * records written by an older kRecordVersion (e.g. the CRC-32/ISO-HDLC + // v1 format that preceded CRC-32C), and + // * a torn tail record from a crash mid-write. + // Either way, appending after them silently loses the new records on the + // next restart, so truncate instead. Safe by construction: the SSD tier is + // a cache whose contents are re-fetchable. + const uint64_t dropped = file_size - seg.scanned_offset; + MORI_UMBP_WARN( + "[SSDTier] {}: dropping {}B of unreadable tail (stale record version or torn write); " + "segment truncated to the last valid record at offset {}", + seg.path, dropped, seg.scanned_offset); + if (ftruncate(seg.fd, static_cast(seg.scanned_offset)) != 0) { + MORI_UMBP_ERROR("[SSDTier] {}: ftruncate to {} failed; segment left as-is", seg.path, + seg.scanned_offset); + continue; + } + // Reset the append cursor too, otherwise the next write leaves a hole that + // the scanner would stop at all over again. + seg.write_offset = seg.scanned_offset; + } +} + SSDTier::~SSDTier() { std::lock_guard lock(mu_); for (auto& kv : index_.MutableSegments()) { @@ -91,7 +211,7 @@ bool SSDTier::OpenOrCreateSegmentLocked(uint64_t segment_id) { seg.id = segment_id; seg.path = dir_ + "/" + segment::BuildFileName(segment_id); - int flags = IsReadOnlyShared() ? O_RDONLY : (O_RDWR | O_CREAT); + int flags = (IsReadOnlyShared() ? O_RDONLY : (O_RDWR | O_CREAT)) | SegmentOpenFlags(); seg.fd = open(seg.path.c_str(), flags, 0644); if (seg.fd < 0) return false; @@ -141,7 +261,7 @@ bool SSDTier::RefreshFromDiskLocked(bool force_full_rescan) { std::string error_message; bool ok = scanner_.RefreshFromDisk(dir_, *io_driver_, index_, IsReadOnlyShared(), - force_full_rescan, &error_message); + force_full_rescan, &error_message, SegmentOpenFlags()); if (!ok && !error_message.empty()) { RememberStatus(IoStatus::IoError(error_message)); } @@ -155,16 +275,20 @@ bool SSDTier::RefreshFollowerLocked() const { bool SSDTier::Write(const std::string& key, const void* data, size_t size) { if (IsReadOnlyShared()) return false; - const size_t record_size = sizeof(segment::RecordHeader) + key.size() + size; + const size_t record_size = static_cast(segment::RecordBytes(key.size(), size)); segment::PreparedRecord pr; + + // Phase 1a (no lock): checksum + assemble, same rationale as WriteBatch. + writer_->Build(key, data, size, &pr); + int write_fd = -1; { - // Phase 1: reserve index space and build record buffer under mu_ + // Phase 1b: reserve index space under mu_ std::lock_guard lock(mu_); if (!EnsureActiveSegment(record_size)) return false; auto* seg = GetSegmentLocked(index_.active_segment_id()); if (!seg) return false; - if (!writer_->Prepare(key, data, size, seg, index_, &pr)) return false; + if (!writer_->Reserve(key, size, seg, index_, &pr)) return false; write_fd = seg->fd; } @@ -193,9 +317,11 @@ bool SSDTier::WriteBatch(const std::vector& keys, if (keys.empty()) return true; if (IsReadOnlyShared()) return false; + // Padded size: what the batch will actually occupy on disk, and therefore what + // the segment-roll check has to be made against. size_t total_bytes = 0; for (size_t i = 0; i < keys.size(); ++i) { - total_bytes += sizeof(segment::RecordHeader) + keys[i].size() + sizes[i]; + total_bytes += static_cast(segment::RecordBytes(keys[i].size(), sizes[i])); } if (total_bytes > ssd_config_.segment_size_bytes) { bool all_ok = true; @@ -205,11 +331,29 @@ bool SSDTier::WriteBatch(const std::vector& keys, return all_ok; } + // Stage timers for the [SsdPerf/tier] PUT breakdown (no-ops unless + // UMBP_SSD_TIMING is set). `prepare` covers the CRC + record memcpy done + // under mu_, which is the stage that blocks concurrent reads on this drive. + const auto t_begin = ssdperf::Now(); + double build_ms = 0.0, lock_ms = 0.0, reserve_ms = 0.0, io_ms = 0.0, sync_ms = 0.0; + + // Phase 1a (NO lock): checksum + assemble every record. This is the expensive + // half of the old Prepare() and it touches no shared state, so keeping it out + // of mu_ stops a write batch from blocking concurrent reads on this drive for + // the whole of its CRC + copy time. + std::vector built(keys.size()); + ParallelFor(keys.size(), TierThreads(), + [&](size_t i) { writer_->Build(keys[i], data_ptrs[i], sizes[i], &built[i]); }); + build_ms = ssdperf::MsSince(t_begin); + std::vector prepared; int write_fd = -1; { - // Phase 1: reserve index space and build record buffers under mu_ + // Phase 1b: reserve index/segment space under mu_ (cheap: no CRC, no copy). + const auto t_lock = ssdperf::Now(); std::lock_guard lock(mu_); + const auto t_locked = ssdperf::Now(); + lock_ms = ssdperf::MsBetween(t_lock, t_locked); if (!EnsureActiveSegment(total_bytes)) return false; auto* seg = GetSegmentLocked(index_.active_segment_id()); if (!seg) return false; @@ -217,22 +361,38 @@ bool SSDTier::WriteBatch(const std::vector& keys, prepared.reserve(keys.size()); for (size_t i = 0; i < keys.size(); ++i) { - segment::PreparedRecord pr; - if (!writer_->Prepare(keys[i], data_ptrs[i], sizes[i], seg, index_, &pr)) continue; - prepared.push_back(std::move(pr)); + if (!writer_->Reserve(keys[i], sizes[i], seg, index_, &built[i])) continue; + prepared.push_back(std::move(built[i])); } + reserve_ms = ssdperf::MsSince(t_locked); } if (prepared.empty()) return true; // Phase 2: perform I/O outside mu_ const bool needs_io_lock = !io_driver_->Capabilities().thread_safe; + const auto t_io = ssdperf::Now(); IoStatus status; if (needs_io_lock) { std::lock_guard io_lock(io_mu_); - status = writer_->WriteRecords(write_fd, prepared, ShouldSyncOnWrite()); + status = writer_->WriteRecords(write_fd, prepared, ShouldSyncOnWrite(), &sync_ms); } else { - status = writer_->WriteRecords(write_fd, prepared, ShouldSyncOnWrite()); + status = writer_->WriteRecords(write_fd, prepared, ShouldSyncOnWrite(), &sync_ms); + } + io_ms = ssdperf::MsSince(t_io) - sync_ms; + + if (ssdperf::Enabled()) { + const double total_ms = ssdperf::MsSince(t_begin); + MORI_UMBP_INFO( + "[SsdPerf/tier] PUT dir={} keys={} bytes={} total_ms={:.3f} GB_s={:.2f} | " + "build_crc_memcpy={:.3f}ms ({:.0f}%, lock-free) mu_wait={:.3f}ms reserve={:.3f}ms " + "({:.0f}%, under mu_) dev_write={:.3f}ms ({:.0f}%) fsync={:.3f}ms ({:.0f}%) | " + "crc_GB_s={:.2f} dev_GB_s={:.2f} sync_on_write={}", + dir_, prepared.size(), total_bytes, total_ms, ssdperf::GbPerSec(total_bytes, total_ms), + build_ms, ssdperf::Pct(build_ms, total_ms), lock_ms, reserve_ms, + ssdperf::Pct(reserve_ms, total_ms), io_ms, ssdperf::Pct(io_ms, total_ms), sync_ms, + ssdperf::Pct(sync_ms, total_ms), ssdperf::GbPerSec(total_bytes, build_ms), + ssdperf::GbPerSec(total_bytes, io_ms), ShouldSyncOnWrite()); } if (!status.ok()) { @@ -244,22 +404,38 @@ bool SSDTier::WriteBatch(const std::vector& keys, return true; } +IoStatus SSDTier::ReadValueInto(int fd, void* dst, size_t size, uint64_t value_offset) const { + // Value offsets are alignment multiples by construction in the v3 layout, so + // only the caller's buffer and length can disqualify a direct read. When they + // do, read the padded value (always present on disk) into an aligned bounce + // buffer and hand back just the value bytes. + if (direct_io_ && !IsDirectIoCompatible(dst, size)) { + AlignedBuffer bounce(size); + IoStatus status = io_driver_->ReadAt(fd, bounce.data(), bounce.padded_size(), value_offset); + if (!status.ok()) return status; + std::memcpy(dst, bounce.data(), size); + return IoStatus::Ok(); + } + return io_driver_->ReadAt(fd, dst, size, value_offset); +} + bool SSDTier::ReadRecordLocked(const std::string& key, void* dst, size_t size, - uint32_t expected_crc, uint64_t value_offset, int read_fd) const { + uint32_t expected_crc, uint64_t value_offset, int read_fd, + bool crc_valid) const { const bool needs_external_lock = !io_driver_->Capabilities().thread_safe; IoStatus status; if (needs_external_lock) { std::lock_guard io_lock(io_mu_); - status = io_driver_->ReadAt(read_fd, dst, size, value_offset); + status = ReadValueInto(read_fd, dst, size, value_offset); } else { - status = io_driver_->ReadAt(read_fd, dst, size, value_offset); + status = ReadValueInto(read_fd, dst, size, value_offset); } if (!status.ok()) { RememberStatus(std::move(status)); return false; } - if (segment::ComputeRecordCrc32(key, dst, size) != expected_crc) { + if (ShouldVerifyCrc(crc_valid) && segment::ComputeRecordCrc32(key, dst, size) != expected_crc) { RememberStatus(IoStatus::Corruption("segment CRC mismatch")); return false; } @@ -270,6 +446,7 @@ bool SSDTier::ReadIntoPtr(const std::string& key, uintptr_t dst_ptr, size_t size int read_fd = -1; uint64_t value_offset = 0; uint32_t expected_crc = 0; + bool crc_valid = true; { std::lock_guard lock(mu_); auto* meta = index_.FindKey(key); @@ -284,10 +461,11 @@ bool SSDTier::ReadIntoPtr(const std::string& key, uintptr_t dst_ptr, size_t size read_fd = seg->fd; value_offset = meta->value_offset; expected_crc = meta->crc32; + crc_valid = meta->crc_valid; index_.TouchLRU(key); } return ReadRecordLocked(key, reinterpret_cast(dst_ptr), size, expected_crc, value_offset, - read_fd); + read_fd, crc_valid); } std::vector SSDTier::ReadBatchIntoPtr(const std::vector& keys, @@ -303,15 +481,27 @@ std::vector SSDTier::ReadBatchIntoPtr(const std::vector& keys uint64_t offset; uint32_t expected_crc; size_t size; - void* dst; + void* dst; // final destination (caller's buffer) + void* io_dst; // where the I/O lands: dst, or `bounce` when misaligned + size_t io_size; // bytes to transfer: size, or the padded value length + bool crc_valid; + AlignedBuffer bounce; // only allocated when a bounce is actually needed }; std::vector lookups; lookups.reserve(keys.size()); + // Stage timers for the [SsdPerf/tier] GET breakdown (no-ops unless + // UMBP_SSD_TIMING is set). The three stages are disjoint and cover the whole + // call, so their percentages answer "device or CPU?" directly. + const auto t_begin = ssdperf::Now(); + double lock_ms = 0.0, lookup_ms = 0.0, io_ms = 0.0, crc_ms = 0.0; + // Phase 1 (mu_ held): batch index lookup + metadata extraction. { std::lock_guard lock(mu_); + const auto t_locked = ssdperf::Now(); + lock_ms = ssdperf::MsBetween(t_begin, t_locked); // Follower: do a single refresh if any key is missing. if (IsReadOnlyShared()) { @@ -334,9 +524,31 @@ std::vector SSDTier::ReadBatchIntoPtr(const std::vector& keys auto* seg = GetSegmentLocked(meta->segment_id); if (!seg || seg->fd < 0) continue; index_.TouchLRU(keys[i]); - lookups.push_back({i, seg->fd, meta->value_offset, meta->crc32, sizes[i], - reinterpret_cast(dst_ptrs[i])}); + + ReadLookup lk; + lk.orig_idx = i; + lk.fd = seg->fd; + lk.offset = meta->value_offset; + lk.expected_crc = meta->crc32; + lk.size = sizes[i]; + lk.dst = reinterpret_cast(dst_ptrs[i]); + lk.crc_valid = meta->crc_valid; + lk.io_dst = lk.dst; + lk.io_size = lk.size; + // Direct I/O demands an aligned buffer and length. KV page buffers are + // normally both (page sizes here are exact 4 KiB multiples out of + // hugepage-backed allocations), so this bounce is the exception, not the + // rule — but it has to exist, because the failure mode without it is an + // EINVAL that the per-key fallback below would simply repeat, surfacing as + // a silent 100% miss rather than an error. + if (direct_io_ && !IsDirectIoCompatible(lk.dst, lk.size)) { + lk.bounce.Resize(lk.size); + lk.io_dst = lk.bounce.data(); + lk.io_size = lk.bounce.padded_size(); + } + lookups.push_back(std::move(lk)); } + lookup_ms = ssdperf::MsSince(t_locked); } if (lookups.empty()) return results; @@ -344,6 +556,7 @@ std::vector SSDTier::ReadBatchIntoPtr(const std::vector& keys // Phase 2 (io_mu_ if needed): batch I/O. const bool needs_io_lock = !io_driver_->Capabilities().thread_safe; const bool use_batch = io_driver_->Capabilities().batch_read && lookups.size() > 1; + const auto t_io = ssdperf::Now(); std::vector io_ok(lookups.size(), false); @@ -351,7 +564,7 @@ std::vector SSDTier::ReadBatchIntoPtr(const std::vector& keys std::vector ops; ops.reserve(lookups.size()); for (const auto& lk : lookups) { - ops.push_back({lk.fd, lk.dst, lk.size, lk.offset}); + ops.push_back({lk.fd, lk.io_dst, lk.io_size, lk.offset}); } IoStatus status; @@ -373,9 +586,9 @@ std::vector SSDTier::ReadBatchIntoPtr(const std::vector& keys IoStatus s; if (needs_io_lock) { std::lock_guard io_lock(io_mu_); - s = io_driver_->ReadAt(lk.fd, lk.dst, lk.size, lk.offset); + s = io_driver_->ReadAt(lk.fd, lk.io_dst, lk.io_size, lk.offset); } else { - s = io_driver_->ReadAt(lk.fd, lk.dst, lk.size, lk.offset); + s = io_driver_->ReadAt(lk.fd, lk.io_dst, lk.io_size, lk.offset); } io_ok[j] = s.ok(); if (!s.ok()) RememberStatus(std::move(s)); @@ -388,24 +601,61 @@ std::vector SSDTier::ReadBatchIntoPtr(const std::vector& keys IoStatus s; if (needs_io_lock) { std::lock_guard io_lock(io_mu_); - s = io_driver_->ReadAt(lk.fd, lk.dst, lk.size, lk.offset); + s = io_driver_->ReadAt(lk.fd, lk.io_dst, lk.io_size, lk.offset); } else { - s = io_driver_->ReadAt(lk.fd, lk.dst, lk.size, lk.offset); + s = io_driver_->ReadAt(lk.fd, lk.io_dst, lk.io_size, lk.offset); } io_ok[j] = s.ok(); if (!s.ok()) RememberStatus(std::move(s)); } } - // Phase 3 (no lock): per-key CRC verification. + io_ms = ssdperf::MsSince(t_io); + + // Phase 3 (no lock): copy out of any bounce buffers, then verify checksums. + // Both are pure CPU over disjoint indices, so they fan out across + // tier_io_threads — this is the phase that was single-threaded while the DRAM + // tier's equivalent memcpy ran on read_threads_ workers. + const auto t_crc = ssdperf::Now(); + uint64_t verified_bytes = 0; + for (size_t j = 0; j < lookups.size(); ++j) { + if (io_ok[j]) verified_bytes += lookups[j].size; + } + + std::vector key_ok(lookups.size(), 0); + ParallelFor(lookups.size(), TierThreads(), [&](size_t j) { + if (!io_ok[j]) return; + auto& lk = lookups[j]; + if (lk.io_dst != lk.dst) std::memcpy(lk.dst, lk.bounce.data(), lk.size); + // CRC covers exactly the value, never the record's alignment padding. + if (ShouldVerifyCrc(lk.crc_valid) && + segment::ComputeRecordCrc32(keys[lk.orig_idx], lk.dst, lk.size) != lk.expected_crc) { + return; // key_ok stays 0 + } + key_ok[j] = 1; + }); + for (size_t j = 0; j < lookups.size(); ++j) { if (!io_ok[j]) continue; - const auto& lk = lookups[j]; - if (segment::ComputeRecordCrc32(keys[lk.orig_idx], lk.dst, lk.size) != lk.expected_crc) { + if (!key_ok[j]) { RememberStatus(IoStatus::Corruption("segment CRC mismatch")); continue; } - results[lk.orig_idx] = true; + results[lookups[j].orig_idx] = true; + } + crc_ms = ssdperf::MsSince(t_crc); + + if (ssdperf::Enabled()) { + const double total_ms = lock_ms + lookup_ms + io_ms + crc_ms; + MORI_UMBP_INFO( + "[SsdPerf/tier] GET dir={} keys={} served={} bytes={} total_ms={:.3f} GB_s={:.2f} | " + "mu_wait={:.3f}ms index={:.3f}ms ({:.0f}%) dev_read={:.3f}ms ({:.0f}%) crc={:.3f}ms " + "({:.0f}%) | dev_GB_s={:.2f} crc_GB_s={:.2f} batched={}", + dir_, keys.size(), lookups.size(), verified_bytes, total_ms, + ssdperf::GbPerSec(verified_bytes, total_ms), lock_ms, lookup_ms, + ssdperf::Pct(lookup_ms, total_ms), io_ms, ssdperf::Pct(io_ms, total_ms), crc_ms, + ssdperf::Pct(crc_ms, total_ms), ssdperf::GbPerSec(verified_bytes, io_ms), + ssdperf::GbPerSec(verified_bytes, crc_ms), use_batch); } return results; @@ -424,6 +674,289 @@ std::vector SSDTier::BatchReadIntoPtr(const std::vector& keys return ReadBatchIntoPtr(keys, dst_ptrs, sizes); } +namespace { + +// One caller range to lift out of a run's bounce buffer once the read lands. +// `device_id` is -1 for ordinary host memory; anything else means the caller +// handed us GPU memory (the tree connector registers its device KV buffers over +// IPC, so the resolved destination really is a device pointer) and the lift has +// to be a hipMemcpy rather than a memcpy. +struct BounceCopy { + void* dst = nullptr; + size_t bounce_offset = 0; + size_t size = 0; + int device_id = -1; +}; + +// One device read. A run is a maximal set of ranges that are contiguous *in +// the object*, which is the shape a layer-group load produces: consecutive +// layers of one component sit next to each other in the stored object. +// Collapsing them into a single pread is the whole reason ranged SSD beats +// whole-object SSD -- it trades N device round trips for one read plus a host +// memcpy, and on this tier a device round trip costs orders of magnitude more +// than the memcpy does. +struct RunRead { + size_t key_index = 0; + int fd = -1; + uint64_t io_offset = 0; + size_t io_size = 0; + void* io_dst = nullptr; + AlignedBuffer bounce; // unused when reading straight into the caller + std::vector copies; // empty when reading straight into the caller +}; + +} // namespace + +std::vector SSDTier::ReadBatchRangesIntoPtr( + const std::vector& keys, const std::vector>& dst_ptrs, + const std::vector>& sizes, + const std::vector>& src_offsets) { + const size_t n = keys.size(); + std::vector results(n, false); + if (n == 0 || !RangeBatchShapeValid(n, dst_ptrs, sizes, src_offsets)) return results; + + const auto t_begin = ssdperf::Now(); + double lock_ms = 0.0, plan_ms = 0.0, io_ms = 0.0, copy_ms = 0.0; + + // What Phase 1 needs out of the index, so run planning can happen without mu_. + struct KeyLocation { + size_t key_index; + int fd; + uint64_t value_offset; + size_t object_size; + }; + std::vector located; + located.reserve(n); + + // Phase 1 (mu_ held): resolve every key to a segment fd + value offset. + { + std::lock_guard lock(mu_); + const auto t_locked = ssdperf::Now(); + lock_ms = ssdperf::MsBetween(t_begin, t_locked); + + // Follower: one refresh covers the whole batch, same as ReadBatchIntoPtr. + if (IsReadOnlyShared()) { + for (size_t i = 0; i < n; ++i) { + if (!sizes[i].empty() && !index_.FindKey(keys[i])) { + RefreshFromDiskLocked(false); + break; + } + } + } + + for (size_t i = 0; i < n; ++i) { + if (sizes[i].empty()) continue; // asked for nothing; stays false + auto* meta = index_.FindKey(keys[i]); + if (!meta) continue; + auto* seg = GetSegmentLocked(meta->segment_id); + if (!seg || seg->fd < 0) continue; + index_.TouchLRU(keys[i]); + located.push_back({i, seg->fd, meta->value_offset, meta->size}); + } + plan_ms = ssdperf::MsSince(t_locked); + } + + if (located.empty()) return results; + + // Phase 2 (no lock): validate ranges and plan the runs. + const auto t_plan = ssdperf::Now(); + std::vector runs; + std::vector expected(n, 0); + std::vector order; + for (const auto& loc : located) { + const size_t i = loc.key_index; + const size_t count = sizes[i].size(); + + bool valid = true; + for (size_t j = 0; j < count; ++j) { + if (sizes[i][j] == 0 || + IsObjectRangeOverflow(src_offsets[i][j], sizes[i][j], loc.object_size)) { + valid = false; + break; + } + } + if (!valid) continue; // a caller bug, reported as this key's failure only + + order.resize(count); + std::iota(order.begin(), order.end(), size_t{0}); + std::sort(order.begin(), order.end(), + [&](size_t a, size_t b) { return src_offsets[i][a] < src_offsets[i][b]; }); + + size_t start = 0; + while (start < count) { + // Extend while the next range begins exactly where this one ends. + // Overlapping ranges simply do not merge; they become separate reads. + size_t end = start + 1; + size_t span_end = src_offsets[i][order[start]] + sizes[i][order[start]]; + while (end < count && src_offsets[i][order[end]] == span_end) { + span_end += sizes[i][order[end]]; + ++end; + } + + const uint64_t file_begin = loc.value_offset + src_offsets[i][order[start]]; + const uint64_t file_end = loc.value_offset + span_end; + const size_t span = static_cast(file_end - file_begin); + void* first_dst = reinterpret_cast(dst_ptrs[i][order[start]]); + + // A device destination cannot receive a read() at all, so it always + // bounces through host memory. Classification costs one + // hipPointerGetAttributes per range, the same toll the DRAM tier pays. + bool any_device = false; + for (size_t j = start; j < end; ++j) { + if (DetectPointerLocation(reinterpret_cast(dst_ptrs[i][order[j]])) + .IsDevice()) { + any_device = true; + break; + } + } + + RunRead run; + run.key_index = i; + run.fd = loc.fd; + + // Straight into the caller's buffer when nothing forbids it: one host + // range, and (under O_DIRECT) an aligned file offset, buffer and length. + // Value offsets are aligned by construction, so only the range's own + // offset can break the file side. + const bool aligned_ok = !direct_io_ || ((file_begin % segment::kRecordAlign) == 0 && + IsDirectIoCompatible(first_dst, span)); + if (end - start == 1 && aligned_ok && !any_device) { + run.io_offset = file_begin; + run.io_size = span; + run.io_dst = first_dst; + } else { + // Bounce. Under O_DIRECT widen to the enclosing aligned window; the + // value is padded to kRecordAlign on disk, so the widened read cannot + // escape this record. + uint64_t io_begin = file_begin; + size_t io_len = span; + if (direct_io_) { + io_begin = file_begin & ~(segment::kRecordAlign - 1); + io_len = static_cast(segment::AlignUp(file_end - io_begin)); + } + run.bounce.Resize(io_len); + run.io_offset = io_begin; + run.io_size = io_len; + run.io_dst = run.bounce.data(); + run.copies.reserve(end - start); + for (size_t j = start; j < end; ++j) { + const size_t idx = order[j]; + void* dst = reinterpret_cast(dst_ptrs[i][idx]); + const auto where = DetectPointerLocation(dst); + run.copies.push_back( + {dst, static_cast(loc.value_offset + src_offsets[i][idx] - io_begin), + sizes[i][idx], where.IsDevice() ? where.device_id : -1}); + } + } + + runs.push_back(std::move(run)); + ++expected[i]; + start = end; + } + } + plan_ms += ssdperf::MsSince(t_plan); + + if (runs.empty()) return results; + + // Phase 3: issue every run as one batch, falling back per run if the driver + // rejects the batch as a whole. + const bool needs_io_lock = !io_driver_->Capabilities().thread_safe; + const bool use_batch = io_driver_->Capabilities().batch_read && runs.size() > 1; + const auto t_io = ssdperf::Now(); + std::vector run_ok(runs.size(), 0); + uint64_t io_bytes = 0; + + auto read_one = [&](size_t j) { + const auto& run = runs[j]; + IoStatus status; + if (needs_io_lock) { + std::lock_guard io_lock(io_mu_); + status = io_driver_->ReadAt(run.fd, run.io_dst, run.io_size, run.io_offset); + } else { + status = io_driver_->ReadAt(run.fd, run.io_dst, run.io_size, run.io_offset); + } + run_ok[j] = status.ok() ? 1 : 0; + if (!status.ok()) RememberStatus(std::move(status)); + }; + + if (use_batch) { + std::vector ops; + ops.reserve(runs.size()); + for (const auto& run : runs) ops.push_back({run.fd, run.io_dst, run.io_size, run.io_offset}); + + IoStatus status; + if (needs_io_lock) { + std::lock_guard io_lock(io_mu_); + status = io_driver_->ReadBatch(ops); + } else { + status = io_driver_->ReadBatch(ops); + } + if (status.ok()) { + std::fill(run_ok.begin(), run_ok.end(), 1); + } else { + RememberStatus(std::move(status)); + for (size_t j = 0; j < runs.size(); ++j) read_one(j); + } + } else { + for (size_t j = 0; j < runs.size(); ++j) read_one(j); + } + io_ms = ssdperf::MsSince(t_io); + for (size_t j = 0; j < runs.size(); ++j) { + if (run_ok[j]) io_bytes += runs[j].io_size; + } + + // Phase 4 (no lock): lift the bounced runs out to their caller buffers. + // Disjoint destinations, so this fans out across tier_io_threads. + const auto t_copy = ssdperf::Now(); + ParallelFor(runs.size(), TierThreads(), [&](size_t j) { + if (!run_ok[j]) return; + auto& run = runs[j]; + for (const auto& copy : run.copies) { + const void* src = run.bounce.data() + copy.bounce_offset; + if (copy.device_id < 0) { + std::memcpy(copy.dst, src, copy.size); + continue; + } + // One synchronous hipMemcpy per range. The DRAM tier merges and pipelines + // these; here the device copy sits behind a device read that costs orders + // of magnitude more, so the simple form is not the bottleneck yet. + if (!DeviceCopy(copy.dst, src, copy.size, hipMemcpyHostToDevice, copy.device_id)) { + run_ok[j] = 0; // distinct index per worker; no race + return; + } + } + }); + copy_ms = ssdperf::MsSince(t_copy); + + std::vector done(n, 0); + for (size_t j = 0; j < runs.size(); ++j) { + if (run_ok[j]) ++done[runs[j].key_index]; + } + uint64_t served_bytes = 0; + for (size_t i = 0; i < n; ++i) { + if (expected[i] == 0 || done[i] != expected[i]) continue; + results[i] = true; + for (size_t size : sizes[i]) served_bytes += size; + } + + if (ssdperf::Enabled()) { + const double total_ms = lock_ms + plan_ms + io_ms + copy_ms; + MORI_UMBP_INFO( + "[SsdPerf/tier] RANGED_GET dir={} keys={} located={} runs={} served_bytes={} " + "read_bytes={} total_ms={:.3f} GB_s={:.2f} | mu_wait={:.3f}ms plan={:.3f}ms ({:.0f}%) " + "dev_read={:.3f}ms ({:.0f}%) bounce_copy={:.3f}ms ({:.0f}%) | dev_GB_s={:.2f} " + "amplification={:.2f}x batched={}", + dir_, keys.size(), located.size(), runs.size(), served_bytes, io_bytes, total_ms, + ssdperf::GbPerSec(served_bytes, total_ms), lock_ms, plan_ms, + ssdperf::Pct(plan_ms, total_ms), io_ms, ssdperf::Pct(io_ms, total_ms), copy_ms, + ssdperf::Pct(copy_ms, total_ms), ssdperf::GbPerSec(io_bytes, io_ms), + served_bytes ? static_cast(io_bytes) / static_cast(served_bytes) : 0.0, + use_batch); + } + + return results; +} + bool SSDTier::Exists(const std::string& key) const { std::lock_guard lock(mu_); if (index_.HasKey(key)) return true; @@ -466,6 +999,7 @@ std::vector SSDTier::Read(const std::string& key) { uint64_t value_offset = 0; uint32_t read_size = 0; uint32_t expected_crc = 0; + bool crc_valid = true; { std::lock_guard lock(mu_); auto* meta = index_.FindKey(key); @@ -480,11 +1014,16 @@ std::vector SSDTier::Read(const std::string& key) { value_offset = meta->value_offset; read_size = meta->size; expected_crc = meta->crc32; + crc_valid = meta->crc_valid; index_.TouchLRU(key); } + // The returned vector is only max_align_t-aligned, so under direct I/O + // ReadValueInto bounces it; that is the whole cost of keeping this API's + // std::vector return type. std::vector out(read_size); - if (!ReadRecordLocked(key, out.data(), out.size(), expected_crc, value_offset, read_fd)) + if (!ReadRecordLocked(key, out.data(), out.size(), expected_crc, value_offset, read_fd, + crc_valid)) return {}; return out; } @@ -493,6 +1032,7 @@ TierCapabilities SSDTier::Capabilities() const { TierCapabilities caps; caps.batch_write = true; caps.batch_read = true; + caps.ranged_read = true; return caps; } @@ -519,4 +1059,25 @@ std::optional SSDTier::GetLocationId(const std::string& key) const return "seg" + std::to_string(meta->segment_id) + ":" + std::to_string(meta->value_offset); } +std::optional SSDTier::LocateRecord(const std::string& key) const { + std::lock_guard lock(mu_); + const auto* meta = index_.FindKey(key); + if (!meta && IsReadOnlyShared()) { + const_cast(this)->RefreshFromDiskLocked(false); + meta = index_.FindKey(key); + } + if (!meta) return std::nullopt; + const auto* seg = GetSegmentLocked(meta->segment_id); + if (!seg || seg->fd < 0) return std::nullopt; + RecordLocation loc; + loc.fd = seg->fd; + loc.value_offset = meta->value_offset; + // v3 pads every value to kRecordAlign, so the padded length is what a GDS read + // must move; the reader trims to value_size afterwards. + loc.readable_size = segment::AlignUp(meta->size); + loc.value_size = meta->size; + loc.direct_io = direct_io_; + return loc; +} + } // namespace mori::umbp diff --git a/src/umbp/local/tiers/tier_backend.cpp b/src/umbp/local/tiers/tier_backend.cpp index 114e5aa18..22c11b7cf 100644 --- a/src/umbp/local/tiers/tier_backend.cpp +++ b/src/umbp/local/tiers/tier_backend.cpp @@ -49,6 +49,13 @@ std::vector TierBackend::ReadBatchIntoPtr(const std::vector& return results; } +std::vector TierBackend::ReadBatchRangesIntoPtr( + const std::vector& keys, const std::vector>& /*dst_ptrs*/, + const std::vector>& /*sizes*/, + const std::vector>& /*src_offsets*/) { + return std::vector(keys.size(), false); +} + std::vector TierBackend::BatchWrite(const std::vector& keys, const std::vector& data_ptrs, const std::vector& sizes) { @@ -59,6 +66,14 @@ std::vector TierBackend::BatchWrite(const std::vector& keys, return results; } +std::vector TierBackend::BatchWriteRanges( + const std::vector& keys, const std::vector& /*object_sizes*/, + const std::vector>& /*src_ptrs*/, + const std::vector>& /*sizes*/, + const std::vector>& /*dst_offsets*/) { + return std::vector(keys.size(), false); +} + std::vector TierBackend::BatchReadIntoPtr(const std::vector& keys, const std::vector& dst_ptrs, const std::vector& sizes) { diff --git a/src/umbp/standalone/bin/standalone_server_main.cpp b/src/umbp/standalone/bin/standalone_server_main.cpp index 6cf402f77..b7440cdc7 100644 --- a/src/umbp/standalone/bin/standalone_server_main.cpp +++ b/src/umbp/standalone/bin/standalone_server_main.cpp @@ -23,11 +23,13 @@ // // MIT License #include +#include #include #include #include #include +#include #include #include #include @@ -108,6 +110,15 @@ bool ParseBoolEnv(const char* name, bool* out, std::string* error) { return false; } +// Identity for a masterless server: unique enough to tell two of them apart in +// a log, and never published anywhere. +std::string DefaultNodeId() { + char host[256] = {0}; + if (gethostname(host, sizeof(host) - 1) != 0) host[0] = '\0'; + const std::string hostname = host[0] != '\0' ? host : "localhost"; + return "umbp-server-" + hostname + "-" + std::to_string(static_cast(getpid())); +} + bool AnyEnv(const std::vector& names) { for (const char* name : names) { if (EnvString(name).has_value()) return true; @@ -117,15 +128,28 @@ bool AnyEnv(const std::vector& names) { bool ApplyDistributedBackendConfigFromEnv(mori::umbp::UMBPConfig* config, bool* distributed_requested, std::string* error) { - static const std::vector kRequiredDistributedEnv = { - "UMBP_MASTER_ADDRESS", - "UMBP_NODE_ADDRESS", - "UMBP_NODE_ID", - "UMBP_IO_ENGINE_HOST", + // Any one of these asks for an explicitly configured backend. Selecting the + // medium counts: a pure-SSD server is configured by naming SSD and nothing + // else, which is the whole command for a node that serves storage locally. + static const std::vector kDistributedSelectorEnv = { + "UMBP_MASTER_ADDRESS", "UMBP_NODE_ADDRESS", "UMBP_NODE_ID", + "UMBP_IO_ENGINE_HOST", "UMBP_DISTRIBUTED_MEDIUM", }; - *distributed_requested = AnyEnv(kRequiredDistributedEnv); + *distributed_requested = AnyEnv(kDistributedSelectorEnv); if (!*distributed_requested) { + // Nothing configured: the factory fills in an embedded DRAM deployment. + // Refuse the one case where that would silently do the wrong thing -- + // someone who explicitly asked for SSD and would otherwise get DRAM. + const char* ssd_enabled = std::getenv("UMBP_SSD_ENABLED"); + if (ssd_enabled != nullptr && ssd_enabled[0] != '\0' && ssd_enabled[0] != '0') { + *error = + "UMBP_SSD_ENABLED is set but no medium was selected. A node serves one " + "medium: set UMBP_DISTRIBUTED_MEDIUM=SSD (with UMBP_SSD_DIR / " + "UMBP_SSD_CAPACITY) to serve storage, or unset UMBP_SSD_ENABLED to " + "serve DRAM."; + return false; + } config->distributed.reset(); return true; } @@ -134,29 +158,40 @@ bool ApplyDistributedBackendConfigFromEnv(mori::umbp::UMBPConfig* config, auto node_address = EnvString("UMBP_NODE_ADDRESS"); auto node_id = EnvString("UMBP_NODE_ID"); auto io_engine_host = EnvString("UMBP_IO_ENGINE_HOST"); - std::vector missing; - if (!master_address.has_value()) missing.push_back("UMBP_MASTER_ADDRESS"); - if (!node_address.has_value()) missing.push_back("UMBP_NODE_ADDRESS"); - if (!node_id.has_value()) missing.push_back("UMBP_NODE_ID"); - if (!io_engine_host.has_value()) missing.push_back("UMBP_IO_ENGINE_HOST"); - if (!missing.empty()) { - std::ostringstream oss; - oss << "distributed-backed standalone server requested but missing required env:"; - for (const char* name : missing) oss << " " << name; - *error = oss.str(); - return false; + + // With a master this node is a cluster member: peers reach it at an address + // it must state, and identity is how master tells it apart, so all three are + // required. Without one, none of them has a consumer -- there is no + // registration, no routing and therefore no peer that could dial in -- so + // they are optional and an omitted IO engine host means no RDMA engine is + // constructed at all. + const bool has_master = master_address.has_value() && !master_address->empty(); + if (has_master) { + std::vector missing; + if (!node_address.has_value()) missing.push_back("UMBP_NODE_ADDRESS"); + if (!node_id.has_value()) missing.push_back("UMBP_NODE_ID"); + if (!io_engine_host.has_value()) missing.push_back("UMBP_IO_ENGINE_HOST"); + if (!missing.empty()) { + std::ostringstream oss; + oss << "master-backed standalone server requested but missing required env:"; + for (const char* name : missing) oss << " " << name; + *error = oss.str(); + return false; + } } mori::umbp::UMBPDistributedConfig dist; - dist.master_config.master_address = *master_address; - dist.master_config.node_address = *node_address; - dist.master_config.node_id = *node_id; - dist.io_engine.host = *io_engine_host; + dist.master_config.master_address = master_address.value_or(""); + dist.master_config.node_address = node_address.value_or("127.0.0.1"); + dist.master_config.node_id = node_id.value_or(DefaultNodeId()); + dist.io_engine.host = io_engine_host.value_or(""); if (!ParseUint16Env("UMBP_IO_ENGINE_PORT", &dist.io_engine.port, error)) return false; if (!ParseUint16Env("UMBP_PEER_SERVICE_PORT", &dist.peer_service_port, error)) return false; if (!ParseSizeEnv("UMBP_DISTRIBUTED_STAGING_BUFFER_SIZE", &dist.staging_buffer_size, error)) return false; + if (!ParseSizeEnv("UMBP_DISTRIBUTED_RANGED_SCRATCH_BYTES", &dist.ranged_scratch_size, error)) + return false; if (!ParseSizeEnv("UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SIZE", &dist.ssd_staging_buffer_size, error)) { return false; @@ -169,11 +204,50 @@ bool ApplyDistributedBackendConfigFromEnv(mori::umbp::UMBPConfig* config, *error = "UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SLOTS must be > 0"; return false; } + if (!ParseBoolEnv("UMBP_DISTRIBUTED_SSD_STAGING_USE_HUGEPAGES", &dist.ssd_staging_use_hugepages, + error)) { + return false; + } + if (!ParseSizeEnv("UMBP_DISTRIBUTED_SSD_STAGING_HUGEPAGE_SIZE", &dist.ssd_staging_hugepage_size, + error)) { + return false; + } + // The one medium this node serves. Selecting SSD here is the whole opt-in + // for a storage node — there is no separate "pure-SSD mode" flag. Unknown + // values are rejected rather than silently defaulting to DRAM, since serving + // the wrong medium is not something a node should discover at runtime. + if (const char* medium_env = std::getenv("UMBP_DISTRIBUTED_MEDIUM")) { + const std::string m(medium_env); + if (m == "DRAM" || m == "dram") { + dist.medium = mori::umbp::UMBPMedium::DRAM; + } else if (m == "HBM" || m == "hbm") { + dist.medium = mori::umbp::UMBPMedium::HBM; + } else if (m == "SSD" || m == "ssd") { + dist.medium = mori::umbp::UMBPMedium::SSD; + } else { + *error = "UMBP_DISTRIBUTED_MEDIUM must be one of: DRAM, HBM, SSD"; + return false; + } + } if (!ParseBoolEnv("UMBP_DISTRIBUTED_CACHE_REMOTE_FETCHES", &dist.cache_remote_fetches, error)) return false; + if (!ParseBoolEnv("UMBP_DISTRIBUTED_RANGED_LOCALITY_PREFETCH", &dist.ranged_locality_prefetch, + error)) { + return false; + } + if (!ParseBoolEnv("UMBP_DISTRIBUTED_LOCAL_FIRST", &dist.local_first, error)) { + return false; + } size_t dram_page_size = static_cast(dist.dram_page_size); if (!ParseSizeEnv("UMBP_DISTRIBUTED_DRAM_PAGE_SIZE", &dram_page_size, error)) return false; dist.dram_page_size = static_cast(dram_page_size); + if (dist.dram_page_size == 0) { + *error = "UMBP_DISTRIBUTED_DRAM_PAGE_SIZE must be explicitly set to > 0"; + return false; + } + if (auto policy = EnvString("UMBP_BACKEND_POLICY"); policy.has_value()) { + dist.backend_policy_path = *policy; + } config->distributed = dist; return true; @@ -252,7 +326,7 @@ int main(int argc, char** argv) { }); MORI_UMBP_INFO("[StandaloneServer] running on {} backend={}", address, - distributed_requested ? "distributed" : "local"); + distributed_requested ? "configured" : "embedded (DRAM, no master)"); server->Run(); stop_signal_waiter = true; diff --git a/src/umbp/standalone/external_kv_identity_client.cpp b/src/umbp/standalone/external_kv_identity_client.cpp index 715db360f..cac69174a 100644 --- a/src/umbp/standalone/external_kv_identity_client.cpp +++ b/src/umbp/standalone/external_kv_identity_client.cpp @@ -32,6 +32,7 @@ #include "mori/utils/mori_log.hpp" #include "umbp.grpc.pb.h" +#include "umbp/common/grpc_limits.h" namespace mori::umbp::standalone { namespace { @@ -82,8 +83,7 @@ ExternalKvIdentityClient::ExternalKvIdentityClient(Config config) : config_(std::move(config)), stub_(nullptr, [](void* p) { delete static_cast<::umbp::UMBPMaster::Stub*>(p); }) { grpc::ChannelArguments args; - args.SetMaxReceiveMessageSize(64 * 1024 * 1024); - args.SetMaxSendMessageSize(64 * 1024 * 1024); + ApplyGrpcLimits(&args); auto channel = grpc::CreateCustomChannel(config_.master_address, grpc::InsecureChannelCredentials(), args); channel_ = channel; diff --git a/src/umbp/standalone/standalone_process_client.cpp b/src/umbp/standalone/standalone_process_client.cpp index 23db61db7..52c4de5cc 100644 --- a/src/umbp/standalone/standalone_process_client.cpp +++ b/src/umbp/standalone/standalone_process_client.cpp @@ -41,6 +41,9 @@ #include #include "mori/utils/mori_log.hpp" +#include "umbp/common/device_copy.h" +#include "umbp/common/grpc_limits.h" +#include "umbp/common/range_utils.h" #include "umbp/local/host_mem_allocator.h" #include "umbp/standalone/ipc.h" @@ -75,6 +78,18 @@ TierType TierFromProto(::umbp::TierType tier) { } } +UMBPDeploymentMode BackendModeFromProto(::umbp::StandaloneBackendMode mode) { + switch (mode) { + case ::umbp::STANDALONE_BACKEND_LOCAL: + return UMBPDeploymentMode::Local; + case ::umbp::STANDALONE_BACKEND_DISTRIBUTED: + return UMBPDeploymentMode::Distributed; + case ::umbp::STANDALONE_BACKEND_UNKNOWN: + default: + return UMBPDeploymentMode::StandaloneProcess; + } +} + bool IsLocalRankZero() { for (const char* name : {"LOCAL_RANK", "OMPI_COMM_WORLD_LOCAL_RANK", "SLURM_LOCALID", "MPI_LOCALRANKID"}) { @@ -182,7 +197,8 @@ StandaloneProcessClient::StandaloneProcessClient(const UMBPConfig& config) : con address_ = standalone_config_.address; fd_socket_path_ = DeriveFdSocketPath(address_); - channel_ = grpc::CreateChannel(address_, grpc::InsecureChannelCredentials()); + channel_ = + grpc::CreateCustomChannel(address_, grpc::InsecureChannelCredentials(), GrpcChannelArgs()); stub_ = ::umbp::UMBPStandalone::NewStub(channel_); MaybeAutoStart(); @@ -205,7 +221,7 @@ std::string StandaloneProcessClient::ClientId() { return client_id_; } -bool StandaloneProcessClient::WaitReady(int timeout_ms) const { +bool StandaloneProcessClient::WaitReady(int timeout_ms) { const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); while (std::chrono::steady_clock::now() < deadline) { grpc::ClientContext ctx; @@ -213,7 +229,11 @@ bool StandaloneProcessClient::WaitReady(int timeout_ms) const { ::umbp::Empty req; ::umbp::PingResponse resp; grpc::Status status = stub_->Ping(&ctx, req, &resp); - if (status.ok() && resp.ready()) return true; + if (status.ok() && resp.ready()) { + backend_mode_ = BackendModeFromProto(resp.deployment_mode()); + supports_ranged_io_ = resp.supports_ranged_io(); + return true; + } std::this_thread::sleep_for(std::chrono::milliseconds(100)); } return false; @@ -252,6 +272,11 @@ void StandaloneProcessClient::MaybeAutoStart() { bool StandaloneProcessClient::OffsetFor(uintptr_t ptr, size_t size, uint64_t* offset, uint64_t* region_base) const { std::lock_guard lock(registration_mu_); + return OffsetForLocked(ptr, size, offset, region_base); +} + +bool StandaloneProcessClient::OffsetForLocked(uintptr_t ptr, size_t size, uint64_t* offset, + uint64_t* region_base) const { for (const auto& region : regions_) { if (ptr < region.base) continue; uintptr_t rel = ptr - region.base; @@ -405,6 +430,82 @@ std::vector StandaloneProcessClient::BatchGet(const std::vector(resp.ok().begin(), resp.ok().end()); } +std::vector StandaloneProcessClient::BatchGetRanges( + const std::vector& keys, const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& src_offsets) { + std::vector failed(keys.size(), false); + if (closing_) return failed; + std::shared_lock lk(op_mutex_); + if (closed_ || !RangeBatchShapeValid(keys.size(), dsts, sizes, src_offsets)) return failed; + + ::umbp::BatchRangeDataRequest req; + req.set_client_id(ClientId()); + { + std::lock_guard registration_lock(registration_mu_); + for (size_t i = 0; i < keys.size(); ++i) { + if (dsts[i].size() > std::numeric_limits::max()) return failed; + req.add_keys(keys[i]); + req.add_range_counts(static_cast(dsts[i].size())); + for (size_t j = 0; j < dsts[i].size(); ++j) { + uint64_t shm_offset = 0; + uint64_t region_base = 0; + if (!OffsetForLocked(dsts[i][j], sizes[i][j], &shm_offset, ®ion_base)) return failed; + req.add_shm_offsets(shm_offset); + req.add_region_bases(region_base); + req.add_sizes(sizes[i][j]); + req.add_object_offsets(src_offsets[i][j]); + } + } + } + + grpc::ClientContext ctx; + ::umbp::BatchBoolResponse resp; + const grpc::Status status = stub_->BatchGetRanges(&ctx, req, &resp); + if (!status.ok() || resp.ok_size() != static_cast(keys.size())) return failed; + return std::vector(resp.ok().begin(), resp.ok().end()); +} + +std::vector StandaloneProcessClient::BatchPutRanges( + const std::vector& keys, const std::vector& object_sizes, + const std::vector>& srcs, const std::vector>& sizes, + const std::vector>& dst_offsets) { + std::vector failed(keys.size(), false); + if (closing_) return failed; + std::shared_lock lk(op_mutex_); + if (closed_ || object_sizes.size() != keys.size() || + !RangeBatchShapeValid(keys.size(), srcs, sizes, dst_offsets)) { + return failed; + } + + ::umbp::BatchRangeDataRequest req; + req.set_client_id(ClientId()); + { + std::lock_guard registration_lock(registration_mu_); + for (size_t i = 0; i < keys.size(); ++i) { + if (srcs[i].size() > std::numeric_limits::max()) return failed; + req.add_keys(keys[i]); + req.add_object_sizes(object_sizes[i]); + req.add_range_counts(static_cast(srcs[i].size())); + for (size_t j = 0; j < srcs[i].size(); ++j) { + uint64_t shm_offset = 0; + uint64_t region_base = 0; + if (!OffsetForLocked(srcs[i][j], sizes[i][j], &shm_offset, ®ion_base)) return failed; + req.add_shm_offsets(shm_offset); + req.add_region_bases(region_base); + req.add_sizes(sizes[i][j]); + req.add_object_offsets(dst_offsets[i][j]); + } + } + } + + grpc::ClientContext ctx; + ::umbp::BatchBoolResponse resp; + const grpc::Status status = stub_->BatchPutRanges(&ctx, req, &resp); + if (!status.ok() || resp.ok_size() != static_cast(keys.size())) return failed; + return std::vector(resp.ok().begin(), resp.ok().end()); +} + std::vector StandaloneProcessClient::BatchExists(const std::vector& keys) const { if (closing_) return std::vector(keys.size(), false); std::shared_lock lk(op_mutex_); @@ -460,18 +561,122 @@ void StandaloneProcessClient::Close() { if (closed_) return; try { DeregisterMemoryLocked(); + } catch (const std::exception& error) { + MORI_UMBP_ERROR("[StandaloneProcessClient] deregistration during close failed: {}", + error.what()); } catch (...) { + MORI_UMBP_ERROR("[StandaloneProcessClient] deregistration during close failed"); } closed_ = true; stub_.reset(); channel_.reset(); } -bool StandaloneProcessClient::RegisterMemory(uintptr_t ptr, size_t size) { +bool StandaloneProcessClient::RegisterMemory(uintptr_t ptr, size_t size, + mori::io::MemoryLocationType loc, int device, + MemoryRegistration /*mode*/) { if (closing_) return false; std::unique_lock lk(op_mutex_); if (closed_) return false; + // Classification is authoritative, and `loc`/`device` only fill in what it + // cannot know. Two callers have to work here: ours, which states the + // location explicitly, and a connector written against the two-argument + // upstream API, which does not. Trusting the pointer over the argument + // means a caller who leaves `loc` at its CPU default cannot silently get a + // device buffer registered as host memory — the failure that would then show + // up as a memcpy from device memory, far from here. + PointerLocation location = DetectPointerLocation(reinterpret_cast(ptr)); + if (loc == mori::io::MemoryLocationType::GPU && !location.IsDevice()) { + MORI_UMBP_ERROR( + "[StandaloneProcessClient] RegisterMemory: caller declared GPU for ptr=0x{:x} but it is " + "not device memory", + ptr); + return false; + } + // A device ordinal from the caller is honoured when classification could not + // supply one (hipPointerGetAttributes reports -1 for some allocations). + if (location.IsDevice() && location.device_id < 0) location.device_id = device; + + if (location.IsDevice()) return RegisterDeviceMemory(ptr, size, location.device_id); + return RegisterHostShmMemory(ptr, size); +} + +bool StandaloneProcessClient::RegisterDeviceMemory(uintptr_t ptr, size_t size, int device_id) { + if (ptr == 0 || size == 0) return false; + ScopedHipDevice device_guard(device_id); + if (!device_guard.IsValid()) { + MORI_UMBP_ERROR("[StandaloneProcessClient] failed to select GPU device {}", device_id); + return false; + } + + void* allocation_base = nullptr; + size_t allocation_size = 0; + const hipError_t range_status = + hipMemGetAddressRange(reinterpret_cast(&allocation_base), &allocation_size, + reinterpret_cast(ptr)); + if (range_status != hipSuccess || allocation_base == nullptr) { + MORI_UMBP_ERROR("[StandaloneProcessClient] hipMemGetAddressRange failed for ptr=0x{:x}: {}", + ptr, hipGetErrorString(range_status)); + (void)hipGetLastError(); + return false; + } + + const uintptr_t allocation_address = reinterpret_cast(allocation_base); + if (ptr < allocation_address) return false; + const uint64_t ipc_offset = static_cast(ptr - allocation_address); + if (ipc_offset > allocation_size || size > allocation_size - ipc_offset) { + MORI_UMBP_ERROR( + "[StandaloneProcessClient] GPU registration range exceeds allocation: ptr=0x{:x} " + "size={} alloc_base=0x{:x} alloc_size={}", + ptr, size, allocation_address, allocation_size); + return false; + } + + hipIpcMemHandle_t handle{}; + const hipError_t handle_status = hipIpcGetMemHandle(&handle, allocation_base); + if (handle_status != hipSuccess) { + MORI_UMBP_ERROR("[StandaloneProcessClient] hipIpcGetMemHandle failed for ptr=0x{:x}: {}", ptr, + hipGetErrorString(handle_status)); + (void)hipGetLastError(); + return false; + } + + ::umbp::RegisterMemoryRequest req; + req.set_client_id(ClientId()); + req.set_worker_base(ptr); + req.set_size(size); + req.set_worker_node_id(standalone_config_.worker_node_id); + req.set_worker_node_address(standalone_config_.worker_node_address); + for (const auto& tag : standalone_config_.tags) req.add_tags(tag); + req.set_kind(::umbp::MEMORY_KIND_GPU_IPC); + req.set_device_id(device_id); + req.set_ipc_handle(reinterpret_cast(&handle), sizeof(handle)); + req.set_ipc_offset(ipc_offset); + req.set_alloc_base(allocation_address); + + grpc::ClientContext ctx; + ::umbp::BoolResponse resp; + const grpc::Status rpc_status = stub_->RegisterMemory(&ctx, req, &resp); + if (!rpc_status.ok() || !resp.ok()) { + MORI_UMBP_ERROR("[StandaloneProcessClient] GPU RegisterMemory failed: {}", + rpc_status.ok() ? resp.error() : rpc_status.error_message()); + return false; + } + + std::lock_guard lock(registration_mu_); + auto existing = std::find_if(regions_.begin(), regions_.end(), [&](const RegisteredRegion& r) { + return r.base == ptr && r.kind == RegionKind::kGpuIpc; + }); + if (existing != regions_.end()) { + existing->size = size; + } else { + regions_.push_back({ptr, size, RegionKind::kGpuIpc}); + } + return true; +} + +bool StandaloneProcessClient::RegisterHostShmMemory(uintptr_t ptr, size_t size) { auto allocation = HostMemAllocator::AcquireShmAllocation(ptr, size); if (!allocation.has_value()) { throw std::runtime_error( @@ -514,7 +719,7 @@ bool StandaloneProcessClient::RegisterMemory(uintptr_t ptr, size_t size) { existing->size = allocation->mapped_size; HostMemAllocator::ReleaseShmAllocation(base); } else { - regions_.push_back({base, allocation->mapped_size}); + regions_.push_back({base, allocation->mapped_size, RegionKind::kHostShm}); } acquired_kept = true; } catch (...) { @@ -533,7 +738,7 @@ void StandaloneProcessClient::DeregisterMemoryLocked() { std::lock_guard lock(registration_mu_); if (regions_.empty()) return; client_id = client_id_; - regions.swap(regions_); + regions = regions_; } // One RPC tears down all of this client's regions server-side (UnmapClient), @@ -542,8 +747,21 @@ void StandaloneProcessClient::DeregisterMemoryLocked() { ::umbp::DeregisterMemoryRequest req; req.set_client_id(client_id); ::umbp::Empty resp; - if (stub_) stub_->DeregisterMemory(&ctx, req, &resp); - for (const auto& region : regions) HostMemAllocator::ReleaseShmAllocation(region.base); + if (!stub_) throw std::runtime_error("standalone DeregisterMemory RPC has no active stub"); + const grpc::Status status = stub_->DeregisterMemory(&ctx, req, &resp); + if (!status.ok()) { + throw std::runtime_error("standalone DeregisterMemory RPC failed: " + status.error_message()); + } + + { + std::lock_guard lock(registration_mu_); + regions_.clear(); + } + for (const auto& region : regions) { + if (region.kind == RegionKind::kHostShm) { + HostMemAllocator::ReleaseShmAllocation(region.base); + } + } } void StandaloneProcessClient::DeregisterMemory(uintptr_t /*ptr*/) { diff --git a/src/umbp/standalone/standalone_server.cpp b/src/umbp/standalone/standalone_server.cpp index a1b9ae457..63a86ea71 100644 --- a/src/umbp/standalone/standalone_server.cpp +++ b/src/umbp/standalone/standalone_server.cpp @@ -37,15 +37,21 @@ #include #include #include +#include #include #include #include #include +#include #include #include +#include #include #include "mori/utils/mori_log.hpp" +#include "umbp/common/device_copy.h" +#include "umbp/common/grpc_limits.h" +#include "umbp/distributed/config.h" #include "umbp/standalone/external_kv_identity_client.h" #include "umbp/standalone/ipc.h" #include "umbp/umbp_client.h" @@ -116,7 +122,20 @@ UMBPConfig NormalizeBackendConfig(UMBPConfig config) { // server constructor. The server backend must never consume that field, // otherwise CreateUMBPClient would recursively create a client to itself. config.standalone_process.reset(); - return config; + // Resolve the deployment here rather than leaving it to CreateUMBPClient: + // the server keeps this config and answers questions about the backend from + // it (peer address, external-KV identity), so it has to describe the backend + // that was actually built. Without this, a server started with no + // distributed environment would hold distributed=nullopt while its client + // reported Distributed, and every external-identity registration would fail + // on a config field the client no longer agreed with. + return WithEmbeddedDefaults(config); +} + +// The live medium, or DRAM for a config that names none (which cannot happen +// after NormalizeBackendConfig, and is the harmless answer if it did). +UMBPMedium BackendMedium(const UMBPConfig& config) { + return config.distributed.has_value() ? config.distributed->medium : UMBPMedium::DRAM; } ::umbp::StandaloneBackendMode BackendModeToProto(UMBPDeploymentMode mode) { @@ -141,6 +160,27 @@ std::chrono::milliseconds FdHandshakeTimeout() { return std::chrono::milliseconds(value); } +// Keeps the read-handler bodies independent of the deployment-specific lock +// policy. Exactly one of the two deferred locks owns the mutex. +class ConditionalReadLock { + public: + ConditionalReadLock(std::shared_mutex& mutex, bool shared) + : shared_(mutex, std::defer_lock), exclusive_(mutex, std::defer_lock) { + if (shared) { + shared_.lock(); + } else { + exclusive_.lock(); + } + } + + ConditionalReadLock(const ConditionalReadLock&) = delete; + ConditionalReadLock& operator=(const ConditionalReadLock&) = delete; + + private: + std::shared_lock shared_; + std::unique_lock exclusive_; +}; + } // namespace class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { @@ -148,6 +188,12 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { Impl(const UMBPConfig& config, std::string address) : backend_config_(NormalizeBackendConfig(config)), client_(CreateUMBPClient(backend_config_)), + // Concurrent reads for every medium but SSD, whose manager serializes + // around its staging arena. Keyed on the LIVE medium rather than on + // ssd.enabled: that flag defaults to true and now describes a tier the + // backend may not serve at all, so reading it here would serialize a + // DRAM server for no reason. + shared_reads_(BackendMedium(backend_config_) != UMBPMedium::SSD), address_(std::move(address)), fd_socket_path_(DeriveFdSocketPath(address_)) {} @@ -171,8 +217,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { if (!StartFdListener()) return false; grpc::ServerBuilder builder; - builder.SetMaxReceiveMessageSize(64 * 1024 * 1024); - builder.SetMaxSendMessageSize(64 * 1024 * 1024); + ApplyGrpcLimits(&builder); builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::MIN_POLLERS, 4); builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::MAX_POLLERS, 32); @@ -190,6 +235,12 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { chmod(grpc_path.c_str(), 0600); MORI_UMBP_INFO("[StandaloneServer] listening grpc={} fd_socket={}", address_, fd_socket_path_); + if (shared_reads_) { + MORI_UMBP_INFO("[StandaloneServer] data plane: concurrent reads (medium={})", + TierTypeName(ToTierType(BackendMedium(backend_config_)))); + } else { + MORI_UMBP_INFO("[StandaloneServer] data plane: serialized reads (SSD medium)"); + } return true; } @@ -207,12 +258,12 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { } UnregisterAllExternalIdentities(); { - std::lock_guard lock(client_mu_); + std::unique_lock lock(client_mu_); client_->Flush(); } UnmapAll(); { - std::lock_guard lock(client_mu_); + std::unique_lock lock(client_mu_); client_->Close(); } unlink(UnixPathFromGrpcAddress(address_).c_str()); @@ -223,18 +274,21 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { ::umbp::PingResponse* response) override { response->set_ready(!shutdown_.load()); response->set_deployment_mode(BackendModeToProto(client_->GetDeploymentMode())); + // The inner client is the authority now that local SSD serves ranges; the + // extra ssd.enabled veto here would keep refusing a backend that can. + response->set_supports_ranged_io(client_->SupportsRangedIO()); return grpc::Status::OK; } grpc::Status Put(grpc::ServerContext*, const ::umbp::PutRequest* request, ::umbp::BoolResponse* response) override { + std::unique_lock lock(client_mu_); uintptr_t ptr = 0; if (!ResolveRange(request->client_id(), request->region_base(), request->shm_offset(), request->size(), &ptr)) { SetBool(response, false, "unregistered or out-of-range shm buffer"); return grpc::Status::OK; } - std::lock_guard lock(client_mu_); if (shutdown_.load()) { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; @@ -245,13 +299,13 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { grpc::Status Get(grpc::ServerContext*, const ::umbp::GetRequest* request, ::umbp::BoolResponse* response) override { + ConditionalReadLock lock(client_mu_, shared_reads_); uintptr_t ptr = 0; if (!ResolveRange(request->client_id(), request->region_base(), request->shm_offset(), request->size(), &ptr)) { SetBool(response, false, "unregistered or out-of-range shm buffer"); return grpc::Status::OK; } - std::lock_guard lock(client_mu_); if (shutdown_.load()) { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; @@ -262,14 +316,14 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { grpc::Status BatchPut(grpc::ServerContext*, const ::umbp::BatchDataRequest* request, ::umbp::BatchBoolResponse* response) override { + std::vector keys(request->keys().begin(), request->keys().end()); + std::vector sizes = Sizes(*request); + std::unique_lock lock(client_mu_); std::vector ptrs; if (!ResolveBatch(*request, &ptrs)) { FillFalse(request->keys_size(), response); return grpc::Status::OK; } - std::vector keys(request->keys().begin(), request->keys().end()); - std::vector sizes = Sizes(*request); - std::lock_guard lock(client_mu_); if (shutdown_.load()) { FillFalse(request->keys_size(), response); return grpc::Status::OK; @@ -278,12 +332,59 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { return grpc::Status::OK; } + grpc::Status BatchPutRanges(grpc::ServerContext*, const ::umbp::BatchRangeDataRequest* request, + ::umbp::BatchBoolResponse* response) override { + size_t total_ranges = 0; + if (!ValidateRangeRequest(*request, /*put=*/true, &total_ranges)) { + FillFalse(request->keys_size(), response); + return grpc::Status::OK; + } + + std::vector queries; + queries.reserve(total_ranges); + for (size_t i = 0; i < total_ranges; ++i) { + queries.push_back({request->region_bases(static_cast(i)), + request->shm_offsets(static_cast(i)), + request->sizes(static_cast(i))}); + } + + std::unique_lock lock(client_mu_); + std::vector flat_ptrs; + if (!ResolveRanges(request->client_id(), queries, &flat_ptrs, /*allow_zero=*/true) || + shutdown_.load()) { + FillFalse(request->keys_size(), response); + return grpc::Status::OK; + } + + std::vector keys(request->keys().begin(), request->keys().end()); + std::vector object_sizes; + object_sizes.reserve(keys.size()); + for (uint64_t size : request->object_sizes()) object_sizes.push_back(static_cast(size)); + std::vector> ptrs(keys.size()); + std::vector> sizes(keys.size()); + std::vector> offsets(keys.size()); + size_t cursor = 0; + for (size_t i = 0; i < keys.size(); ++i) { + const size_t count = request->range_counts(static_cast(i)); + ptrs[i].reserve(count); + sizes[i].reserve(count); + offsets[i].reserve(count); + for (size_t j = 0; j < count; ++j, ++cursor) { + ptrs[i].push_back(flat_ptrs[cursor]); + sizes[i].push_back(static_cast(request->sizes(static_cast(cursor)))); + offsets[i].push_back( + static_cast(request->object_offsets(static_cast(cursor)))); + } + } + FillResults(client_->BatchPutRanges(keys, object_sizes, ptrs, sizes, offsets), response); + return grpc::Status::OK; + } + grpc::Status BatchPutWithDepth(grpc::ServerContext*, const ::umbp::BatchDataWithDepthRequest* request, ::umbp::BatchBoolResponse* response) override { // region_bases is optional for legacy single-region callers; when present it - // must be parallel to keys (BatchPutWithDepth resolves inline, not through - // ResolveBatch). + // must be parallel to keys. const bool has_region_bases = request->region_bases_size() > 0; if (request->keys_size() != request->shm_offsets_size() || request->keys_size() != request->sizes_size() || @@ -291,24 +392,24 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { FillFalse(request->keys_size(), response); return grpc::Status::OK; } - std::vector ptrs; - ptrs.reserve(request->keys_size()); - for (int i = 0; i < request->keys_size(); ++i) { - uintptr_t ptr = 0; - uint64_t region_base = has_region_bases ? request->region_bases(i) : 0; - if (!ResolveRange(request->client_id(), region_base, request->shm_offsets(i), - request->sizes(i), &ptr)) { - FillFalse(request->keys_size(), response); - return grpc::Status::OK; - } - ptrs.push_back(ptr); - } std::vector keys(request->keys().begin(), request->keys().end()); std::vector sizes; sizes.reserve(request->sizes_size()); for (uint64_t size : request->sizes()) sizes.push_back(static_cast(size)); std::vector depths(request->depths().begin(), request->depths().end()); - std::lock_guard lock(client_mu_); + std::vector queries; + queries.reserve(request->keys_size()); + for (int i = 0; i < request->keys_size(); ++i) { + queries.push_back({has_region_bases ? request->region_bases(i) : 0, request->shm_offsets(i), + request->sizes(i)}); + } + + std::unique_lock lock(client_mu_); + std::vector ptrs; + if (!ResolveRanges(request->client_id(), queries, &ptrs, /*allow_zero=*/false)) { + FillFalse(request->keys_size(), response); + return grpc::Status::OK; + } if (shutdown_.load()) { FillFalse(request->keys_size(), response); return grpc::Status::OK; @@ -319,14 +420,14 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { grpc::Status BatchGet(grpc::ServerContext*, const ::umbp::BatchDataRequest* request, ::umbp::BatchBoolResponse* response) override { + std::vector keys(request->keys().begin(), request->keys().end()); + std::vector sizes = Sizes(*request); + ConditionalReadLock lock(client_mu_, shared_reads_); std::vector ptrs; if (!ResolveBatch(*request, &ptrs)) { FillFalse(request->keys_size(), response); return grpc::Status::OK; } - std::vector keys(request->keys().begin(), request->keys().end()); - std::vector sizes = Sizes(*request); - std::lock_guard lock(client_mu_); if (shutdown_.load()) { FillFalse(request->keys_size(), response); return grpc::Status::OK; @@ -335,9 +436,54 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { return grpc::Status::OK; } + grpc::Status BatchGetRanges(grpc::ServerContext*, const ::umbp::BatchRangeDataRequest* request, + ::umbp::BatchBoolResponse* response) override { + size_t total_ranges = 0; + if (!ValidateRangeRequest(*request, /*put=*/false, &total_ranges)) { + FillFalse(request->keys_size(), response); + return grpc::Status::OK; + } + + std::vector queries; + queries.reserve(total_ranges); + for (size_t i = 0; i < total_ranges; ++i) { + queries.push_back({request->region_bases(static_cast(i)), + request->shm_offsets(static_cast(i)), + request->sizes(static_cast(i))}); + } + + ConditionalReadLock lock(client_mu_, shared_reads_); + std::vector flat_ptrs; + if (!ResolveRanges(request->client_id(), queries, &flat_ptrs, /*allow_zero=*/true) || + shutdown_.load()) { + FillFalse(request->keys_size(), response); + return grpc::Status::OK; + } + + std::vector keys(request->keys().begin(), request->keys().end()); + std::vector> ptrs(keys.size()); + std::vector> sizes(keys.size()); + std::vector> offsets(keys.size()); + size_t cursor = 0; + for (size_t i = 0; i < keys.size(); ++i) { + const size_t count = request->range_counts(static_cast(i)); + ptrs[i].reserve(count); + sizes[i].reserve(count); + offsets[i].reserve(count); + for (size_t j = 0; j < count; ++j, ++cursor) { + ptrs[i].push_back(flat_ptrs[cursor]); + sizes[i].push_back(static_cast(request->sizes(static_cast(cursor)))); + offsets[i].push_back( + static_cast(request->object_offsets(static_cast(cursor)))); + } + } + FillResults(client_->BatchGetRanges(keys, ptrs, sizes, offsets), response); + return grpc::Status::OK; + } + grpc::Status Exists(grpc::ServerContext*, const ::umbp::KeyRequest* request, ::umbp::BoolResponse* response) override { - std::lock_guard lock(client_mu_); + ConditionalReadLock lock(client_mu_, shared_reads_); if (shutdown_.load()) { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; @@ -349,7 +495,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { grpc::Status BatchExists(grpc::ServerContext*, const ::umbp::BatchKeysRequest* request, ::umbp::BatchBoolResponse* response) override { std::vector keys(request->keys().begin(), request->keys().end()); - std::lock_guard lock(client_mu_); + ConditionalReadLock lock(client_mu_, shared_reads_); if (shutdown_.load()) { FillFalse(request->keys_size(), response); return grpc::Status::OK; @@ -361,7 +507,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { grpc::Status BatchExistsConsecutive(grpc::ServerContext*, const ::umbp::BatchKeysRequest* request, ::umbp::CountResponse* response) override { std::vector keys(request->keys().begin(), request->keys().end()); - std::lock_guard lock(client_mu_); + ConditionalReadLock lock(client_mu_, shared_reads_); if (shutdown_.load()) { response->set_count(0); return grpc::Status::OK; @@ -372,7 +518,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { grpc::Status Clear(grpc::ServerContext*, const ::umbp::Empty*, ::umbp::BoolResponse* response) override { - std::lock_guard lock(client_mu_); + std::unique_lock lock(client_mu_); if (shutdown_.load()) { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; @@ -383,7 +529,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { grpc::Status Flush(grpc::ServerContext*, const ::umbp::Empty*, ::umbp::BoolResponse* response) override { - std::lock_guard lock(client_mu_); + std::unique_lock lock(client_mu_); if (shutdown_.load()) { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; @@ -398,10 +544,21 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; } + if (request->kind() == ::umbp::MEMORY_KIND_GPU_IPC) { + std::lock_guard lifecycle_lock(external_identity_lifecycle_mu_); + RegisterGpuIpc(*request, response); + if (response->ok() && !EnsureExternalIdentity(*request)) { + MORI_UMBP_WARN( + "[StandaloneServer] external-KV identity registration failed for client_id={} " + "worker_node_id={}; continuing with core GPU registration", + request->client_id(), request->worker_node_id()); + } + return grpc::Status::OK; + } std::lock_guard lifecycle_lock(external_identity_lifecycle_mu_); bool ok = false; { - std::lock_guard lock(memory_mu_); + std::shared_lock lock(memory_mu_); auto it = memory_.find(request->client_id()); if (it != memory_.end()) { ok = std::any_of(it->second.begin(), it->second.end(), [&](const RegisteredMemory& mem) { @@ -492,7 +649,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { if (auto identity = GetExternalIdentity(request->client_id())) { matches = identity->MatchExternalKv(hashes, request->count_as_hit()); } else { - std::lock_guard lock(client_mu_); + ConditionalReadLock lock(client_mu_, shared_reads_); if (!shutdown_.load()) matches = client_->MatchExternalKv(hashes, request->count_as_hit()); } FillExternalKvMatches(matches, response); @@ -508,7 +665,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { if (auto identity = GetExternalIdentity(request->client_id())) { entries = identity->GetExternalKvHitCounts(hashes); } else { - std::lock_guard lock(client_mu_); + ConditionalReadLock lock(client_mu_, shared_reads_); if (!shutdown_.load()) entries = client_->GetExternalKvHitCounts(hashes); } FillExternalKvHitCounts(entries, response); @@ -516,10 +673,38 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { } private: + enum class MemoryKind { kHostShm, kGpuIpc }; + struct RegisteredMemory { void* base = nullptr; uint64_t worker_base = 0; uint64_t size = 0; + MemoryKind kind = MemoryKind::kHostShm; + int device_id = -1; + uint64_t alloc_base = 0; + std::string client_id; + }; + + struct RangeQuery { + uint64_t region_base = 0; + uint64_t offset = 0; + uint64_t size = 0; + }; + + struct IpcKey { + std::string client_id; + int device_id = -1; + uint64_t alloc_base = 0; + + bool operator<(const IpcKey& other) const { + return std::tie(client_id, device_id, alloc_base) < + std::tie(other.client_id, other.device_id, other.alloc_base); + } + }; + + struct IpcMapping { + void* base = nullptr; + size_t refcount = 0; }; bool BackendIsDistributed() const { @@ -643,8 +828,13 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { void StopFdListener() { if (!fd_running_.exchange(false)) return; if (listen_fd_ >= 0) shutdown(listen_fd_, SHUT_RDWR); - int active_fd = active_fd_client_.load(); - if (active_fd >= 0) shutdown(active_fd, SHUT_RDWR); + { + // The lock pairs with the accept loop's publish: either we interrupt a + // connection that is already blocked in the handshake, or the loop sees + // fd_running_ false and drops the connection without ever blocking. + std::lock_guard lock(active_fd_mu_); + if (active_fd_client_ >= 0) shutdown(active_fd_client_, SHUT_RDWR); + } if (fd_thread_.joinable()) fd_thread_.join(); if (listen_fd_ >= 0) { close(listen_fd_); @@ -661,13 +851,23 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { } continue; } - active_fd_client_.store(client_fd); + { + std::lock_guard lock(active_fd_mu_); + if (!fd_running_.load()) { + close(client_fd); + break; + } + active_fd_client_ = client_fd; + } if (!SetFdSocketTimeouts(client_fd, FdHandshakeTimeout())) { MORI_UMBP_WARN("[StandaloneServer] failed to set fd socket timeout: {}", std::strerror(errno)); } HandleFdConnection(client_fd); - active_fd_client_.store(-1); + { + std::lock_guard lock(active_fd_mu_); + active_fd_client_ = -1; + } close(client_fd); } } @@ -702,22 +902,12 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { } std::string client_id(msg.client_id); - RegisteredMemory entry{mapped, static_cast(msg.worker_base), msg.size}; - std::optional old_mem; - { - std::lock_guard lock(memory_mu_); - auto& regions = memory_[client_id]; - auto existing = std::find_if( - regions.begin(), regions.end(), - [&](const RegisteredMemory& mem) { return mem.worker_base == entry.worker_base; }); - if (existing != regions.end()) { - // Same region re-registered: replace the mapping, release the old one. - old_mem = *existing; - *existing = entry; - } else { - regions.push_back(entry); - } - } + RegisteredMemory entry; + entry.base = mapped; + entry.worker_base = static_cast(msg.worker_base); + entry.size = msg.size; + entry.client_id = client_id; + std::optional old_mem = InsertOrReplaceRegion(client_id, entry); // Release outside the lock, and via ReleaseRegisteredMemory so the backend // deregisters the region before its mapping is munmap'd (required by the // distributed backend; see design-standalone-process-mode.md §5.3/§6.3). @@ -727,22 +917,184 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { return true; } + bool HasIdenticalGpuRegion(const ::umbp::RegisterMemoryRequest& request) { + std::shared_lock lock(memory_mu_); + auto it = memory_.find(request.client_id()); + if (it == memory_.end()) return false; + return std::any_of(it->second.begin(), it->second.end(), [&](const RegisteredMemory& mem) { + return mem.kind == MemoryKind::kGpuIpc && mem.worker_base == request.worker_base() && + mem.size == request.size() && mem.device_id == request.device_id() && + mem.alloc_base == request.alloc_base(); + }); + } + + std::optional InsertOrReplaceRegion(const std::string& client_id, + const RegisteredMemory& entry) { + std::unique_lock lock(memory_mu_); + auto& regions = memory_[client_id]; + auto existing = std::find_if(regions.begin(), regions.end(), [&](const RegisteredMemory& mem) { + return mem.worker_base == entry.worker_base; + }); + if (existing == regions.end()) { + regions.push_back(entry); + return std::nullopt; + } + RegisteredMemory old = *existing; + *existing = entry; + return old; + } + + void RegisterGpuIpc(const ::umbp::RegisterMemoryRequest& request, + ::umbp::BoolResponse* response) { + const UMBPDeploymentMode mode = client_->GetDeploymentMode(); + if (mode != UMBPDeploymentMode::Distributed) { + SetBool(response, false, "GPU IPC registration requires a Distributed backend"); + return; + } + if (BackendMedium(backend_config_) == UMBPMedium::SSD) { + SetBool(response, false, + "GPU IPC registration is not supported on an SSD medium " + "(run this server on the DRAM medium)"); + return; + } + if (request.client_id().empty() || request.worker_base() == 0 || request.size() == 0 || + request.alloc_base() == 0 || request.ipc_handle().size() != sizeof(hipIpcMemHandle_t)) { + SetBool(response, false, "invalid GPU IPC registration payload"); + return; + } + if (HasIdenticalGpuRegion(request)) { + SetBool(response, true); + return; + } + + const IpcKey key{request.client_id(), request.device_id(), request.alloc_base()}; + ScopedHipDevice device_guard(request.device_id()); + if (!device_guard.IsValid()) { + SetBool(response, false, "hipSetDevice failed for the requested device_id"); + return; + } + + void* mapped_base = nullptr; + { + std::lock_guard lock(ipc_mu_); + auto existing = ipc_maps_.find(key); + if (existing != ipc_maps_.end()) { + ++existing->second.refcount; + mapped_base = existing->second.base; + } else { + hipIpcMemHandle_t handle{}; + std::memcpy(&handle, request.ipc_handle().data(), sizeof(handle)); + const hipError_t status = + hipIpcOpenMemHandle(&mapped_base, handle, hipIpcMemLazyEnablePeerAccess); + if (status != hipSuccess) { + const std::string error = hipGetErrorString(status); + (void)hipGetLastError(); + SetBool(response, false, "hipIpcOpenMemHandle failed: " + error); + return; + } + ipc_maps_.emplace(key, IpcMapping{mapped_base, 1}); + } + } + + const uintptr_t mapped_address = reinterpret_cast(mapped_base); + if (request.ipc_offset() > std::numeric_limits::max() - mapped_address) { + ReleaseIpcMapping(key); + SetBool(response, false, "GPU IPC offset overflows the mapped address"); + return; + } + + RegisteredMemory entry; + entry.base = reinterpret_cast(mapped_address + request.ipc_offset()); + entry.worker_base = request.worker_base(); + entry.size = request.size(); + entry.kind = MemoryKind::kGpuIpc; + entry.device_id = request.device_id(); + entry.alloc_base = request.alloc_base(); + entry.client_id = request.client_id(); + // Both modes declare the imported mapping to the inner backend; only Local + // pins it. Ranged distributed I/O stages remote objects through a registered + // host arena, so an RDMA MR over this mapping buys nothing, and the dmabuf + // fallback is unsafe for IPC-imported ROCm mappings. + // + // Declaring it is not optional in either mode. Leaving it out is correct -- + // an unregistered device pointer is classified rather than assumed host, so + // HbmCopyEngine still claims the pair -- but a range that misses the region + // table is described by its own address instead of its region base, so it + // becomes its own transfer plan instead of sharing one. Measured on + // DSv4-Pro/TP8 at 256K, that is plan = 23.2% of a ranged call against 1.4% + // once declared. + if (!RegisterBackendMemory(entry.base, static_cast(entry.size), + mori::io::MemoryLocationType::GPU, entry.device_id, + mode == UMBPDeploymentMode::Local + ? MemoryRegistration::kPinned + : MemoryRegistration::kLocalCopyOnly)) { + ReleaseIpcMapping(key); + SetBool(response, false, "inner backend RegisterMemory failed"); + return; + } + + std::optional old_mem = InsertOrReplaceRegion(request.client_id(), entry); + if (old_mem.has_value()) ReleaseRegisteredMemory(*old_mem); + MORI_UMBP_INFO( + "[StandaloneServer] registered GPU IPC client_id={} worker_base=0x{:x} size={}MB " + "device={}", + request.client_id(), request.worker_base(), request.size() / (1024 * 1024), + request.device_id()); + SetBool(response, true); + } + + void ReleaseIpcMapping(const IpcKey& key) { + void* mapped_base = nullptr; + { + std::lock_guard lock(ipc_mu_); + auto it = ipc_maps_.find(key); + if (it == ipc_maps_.end()) return; + if (--it->second.refcount != 0) return; + mapped_base = it->second.base; + ipc_maps_.erase(it); + } + + ScopedHipDevice device_guard(key.device_id); + if (!device_guard.IsValid()) { + MORI_UMBP_WARN("[StandaloneServer] failed to select GPU {} while closing IPC mapping", + key.device_id); + } + const hipError_t status = hipIpcCloseMemHandle(mapped_base); + if (status != hipSuccess) { + MORI_UMBP_WARN("[StandaloneServer] hipIpcCloseMemHandle failed for client_id={}: {}", + key.client_id, hipGetErrorString(status)); + (void)hipGetLastError(); + } + } + + size_t IpcMappingCount() { + std::lock_guard lock(ipc_mu_); + return ipc_maps_.size(); + } + void UnmapClient(const std::string& client_id) { std::vector mems; { - std::lock_guard lock(memory_mu_); + std::unique_lock lock(memory_mu_); auto it = memory_.find(client_id); if (it == memory_.end()) return; mems.swap(it->second); memory_.erase(it); } + const size_t mappings_before = IpcMappingCount(); for (const auto& mem : mems) ReleaseRegisteredMemory(mem); + const size_t mappings_after = IpcMappingCount(); + MORI_UMBP_INFO( + "[StandaloneServer] unmapped client_id={} regions={} ipc_allocations_released={} " + "ipc_allocations_remaining={}", + client_id, mems.size(), + mappings_before >= mappings_after ? mappings_before - mappings_after : 0, mappings_after); } void UnmapAll() { std::vector entries; { - std::lock_guard lock(memory_mu_); + std::unique_lock lock(memory_mu_); for (auto& kv : memory_) { for (const auto& mem : kv.second) entries.push_back(mem); } @@ -751,19 +1103,31 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { for (const auto& mem : entries) ReleaseRegisteredMemory(mem); } - bool RegisterBackendMemory(void* base, size_t size) { - std::lock_guard lock(client_mu_); + // `loc`/`device` describe the mapping this server made, not the worker's + // original allocation. They must be passed through: a GPU-IPC region is + // device memory in this process too, and registering it as host memory would + // send it down the inner client's host paths — a memcpy from device memory in + // the local case, and a host staging bounce in the distributed one. + bool RegisterBackendMemory(void* base, size_t size, + mori::io::MemoryLocationType loc = mori::io::MemoryLocationType::CPU, + int device = -1, + MemoryRegistration mode = MemoryRegistration::kPinned) { + std::unique_lock lock(client_mu_); if (shutdown_.load()) return false; - return client_->RegisterMemory(reinterpret_cast(base), size); + return client_->RegisterMemory(reinterpret_cast(base), size, loc, device, mode); } void ReleaseRegisteredMemory(const RegisteredMemory& mem) { if (!mem.base) return; { - std::lock_guard lock(client_mu_); + std::unique_lock lock(client_mu_); client_->DeregisterMemory(reinterpret_cast(mem.base)); } - munmap(mem.base, static_cast(mem.size)); + if (mem.kind == MemoryKind::kGpuIpc) { + ReleaseIpcMapping(IpcKey{mem.client_id, mem.device_id, mem.alloc_base}); + } else { + munmap(mem.base, static_cast(mem.size)); + } } // Resolves (client_id, region_base, offset) to a server-local pointer. @@ -772,10 +1136,18 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { bool ResolveRange(const std::string& client_id, uint64_t region_base, uint64_t offset, uint64_t size, uintptr_t* out_ptr) { if (!out_ptr || size == 0) return false; - std::lock_guard lock(memory_mu_); + std::shared_lock lock(memory_mu_); auto it = memory_.find(client_id); if (it == memory_.end()) return false; - for (const auto& mem : it->second) { + return ResolveRangeInRegions(it->second, region_base, offset, size, out_ptr, + /*allow_zero=*/false); + } + + static bool ResolveRangeInRegions(const std::vector& regions, + uint64_t region_base, uint64_t offset, uint64_t size, + uintptr_t* out_ptr, bool allow_zero) { + if (!out_ptr || (!allow_zero && size == 0)) return false; + for (const auto& mem : regions) { if (region_base != 0 && mem.worker_base != region_base) continue; if (offset > mem.size || size > mem.size - offset) { if (region_base != 0) return false; @@ -787,6 +1159,29 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { return false; } + bool ResolveRanges(const std::string& client_id, const std::vector& queries, + std::vector* ptrs, bool allow_zero) { + if (!ptrs) return false; + ptrs->clear(); + ptrs->reserve(queries.size()); + if (queries.empty()) return true; + + std::shared_lock lock(memory_mu_); + auto it = memory_.find(client_id); + if (it == memory_.end()) return false; + + for (const auto& query : queries) { + uintptr_t ptr = 0; + if (!ResolveRangeInRegions(it->second, query.region_base, query.offset, query.size, &ptr, + allow_zero)) { + ptrs->clear(); + return false; + } + ptrs->push_back(ptr); + } + return true; + } + bool ResolveBatch(const ::umbp::BatchDataRequest& request, std::vector* ptrs) { // region_bases is optional for legacy single-region callers; when present it // must be parallel to keys. @@ -796,17 +1191,47 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { (has_region_bases && request.keys_size() != request.region_bases_size())) { return false; } - ptrs->clear(); - ptrs->reserve(request.keys_size()); + std::vector queries; + queries.reserve(request.keys_size()); for (int i = 0; i < request.keys_size(); ++i) { - uintptr_t ptr = 0; uint64_t region_base = has_region_bases ? request.region_bases(i) : 0; - if (!ResolveRange(request.client_id(), region_base, request.shm_offsets(i), request.sizes(i), - &ptr)) { - return false; + queries.push_back({region_base, request.shm_offsets(i), request.sizes(i)}); + } + return ResolveRanges(request.client_id(), queries, ptrs, /*allow_zero=*/false); + } + + static bool ValidateRangeRequest(const ::umbp::BatchRangeDataRequest& request, bool put, + size_t* total_ranges) { + if (!total_ranges || request.range_counts_size() != request.keys_size()) return false; + if ((put && request.object_sizes_size() != request.keys_size()) || + (!put && request.object_sizes_size() != 0)) { + return false; + } + + size_t total = 0; + for (uint32_t count : request.range_counts()) { + if (count > std::numeric_limits::max() - total) return false; + total += count; + } + if (total != static_cast(request.shm_offsets_size()) || + total != static_cast(request.region_bases_size()) || + total != static_cast(request.sizes_size()) || + total != static_cast(request.object_offsets_size())) { + return false; + } + + if constexpr (sizeof(size_t) < sizeof(uint64_t)) { + for (uint64_t value : request.sizes()) { + if (value > std::numeric_limits::max()) return false; + } + for (uint64_t value : request.object_offsets()) { + if (value > std::numeric_limits::max()) return false; + } + for (uint64_t value : request.object_sizes()) { + if (value > std::numeric_limits::max()) return false; } - ptrs->push_back(ptr); } + *total_ranges = total; return true; } @@ -853,24 +1278,31 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { UMBPConfig backend_config_; std::unique_ptr client_; + const bool shared_reads_; std::string address_; std::string fd_socket_path_; std::unique_ptr server_; std::atomic shutdown_{false}; - std::mutex client_mu_; - std::mutex memory_mu_; + // Data operations acquire client_mu_ before memory_mu_ and keep the client + // lock through the backend call, making it the lifetime barrier for resolved + // host/GPU mappings. Lifecycle paths never hold the two locks at once. + std::shared_mutex client_mu_; + mutable std::shared_mutex memory_mu_; + std::mutex ipc_mu_; // A worker registers N non-contiguous host regions (e.g. DeepSeek-V4's KV // side pools), so each client_id maps to a list of regions, resolved by // worker_base at data-op time. std::map> memory_; + std::map ipc_maps_; std::mutex external_identity_lifecycle_mu_; std::mutex external_identity_mu_; std::map> external_identities_; std::atomic fd_running_{false}; int listen_fd_ = -1; - std::atomic active_fd_client_{-1}; + std::mutex active_fd_mu_; + int active_fd_client_ = -1; std::thread fd_thread_; }; diff --git a/src/umbp/umbp_client_factory.cpp b/src/umbp/umbp_client_factory.cpp index 87b152f4b..69dbc3355 100644 --- a/src/umbp/umbp_client_factory.cpp +++ b/src/umbp/umbp_client_factory.cpp @@ -19,21 +19,196 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. +#include + +#include +#include +#include + +#include "mori/utils/mori_log.hpp" #include "umbp/distributed/distributed_client.h" -#include "umbp/local/standalone_client.h" #include "umbp/standalone/standalone_process_client.h" #include "umbp/umbp_client.h" namespace mori::umbp { -std::unique_ptr CreateUMBPClient(const UMBPConfig& config) { - if (config.distributed.has_value()) { - return std::make_unique(config); +namespace { + +// Identity for a client that will never register with anyone. It still has to +// be unique and stable within the process, because it is what a route result is +// compared against to decide "this key is mine" -- and, more practically, +// because it is the string every log line from this client carries. Several +// clients per process is the normal case (one per rank), hence the counter. +std::string EmbeddedNodeId() { + static std::atomic counter{0}; + char host[256] = {0}; + if (gethostname(host, sizeof(host) - 1) != 0) host[0] = '\0'; + const std::string hostname = host[0] != '\0' ? host : "localhost"; + return "embedded-" + hostname + "-" + std::to_string(static_cast(getpid())) + "-" + + std::to_string(counter.fetch_add(1, std::memory_order_relaxed)); +} + +// Fallback allocation granularity of the embedded DRAM pool, used only when the +// caller expressed no preference. +// +// This matters more than it looks: the pool is paged, so a value smaller than +// one page still consumes a whole one. A caller storing KV pages sets the size +// to the page's exact byte size (which is what the distributed deployments do) +// and gets one page per key; a caller storing small values wants it small. The +// knob is UMBPDistributedConfig::dram_page_size, which WithEmbeddedDefaults +// honours for an embedded pool too, so this is reached only by a caller who set +// neither that field nor the environment override below. +uint64_t EmbeddedPageSize() { + // 64 KiB rather than the 2 MiB every other deployment uses. Those are told + // the KV page's exact byte size and genuinely get one page per key; a caller + // who named no deployment has told us nothing, and paying 2 MiB for a value + // that may be a few hundred bytes silently evicts a pool that had ample room. + // Small enough that a 4 MiB pool holds 64 values instead of 8, large enough + // that a multi-megabyte value still spans few enough pages for the allocator + // to find a contiguous run. Callers storing real KV pages should say so -- + // set distributed.dram_page_size, which this function no longer overrides. + constexpr uint64_t kDefault = 64ULL * 1024; + const char* raw = std::getenv("UMBP_EMBEDDED_DRAM_PAGE_SIZE"); + if (raw == nullptr || raw[0] == '\0') return kDefault; + char* end = nullptr; + const unsigned long long parsed = std::strtoull(raw, &end, 10); + if (end == raw || parsed == 0) { + MORI_UMBP_WARN("[UMBP] ignoring UMBP_EMBEDDED_DRAM_PAGE_SIZE='{}' (not a positive integer)", + raw); + return kDefault; } + return static_cast(parsed); +} + +// Shrink the page to something the pool can actually hold. +// +// The pool is paged and the allocator hands out whole pages, so a page larger +// than the pool means ZERO pages -- and then every put fails with NO_SPACE, +// silently, because a failed put is a legal answer rather than an error. The +// deleted local backend allocated exact sizes and had no such cliff, so a +// small pool that used to work would simply stop working. Fit instead, and +// say so. +uint64_t FitPageSizeToPool(uint64_t requested, uint64_t capacity_bytes) { + constexpr uint64_t kMinPages = 8; // a pool that holds fewer is not a cache + constexpr uint64_t kMinPageSize = 4096; + if (capacity_bytes == 0 || requested <= capacity_bytes / kMinPages) return requested; + + uint64_t fitted = kMinPageSize; + while (fitted * 2 <= capacity_bytes / kMinPages) fitted *= 2; + if (fitted >= requested) return requested; + MORI_UMBP_WARN( + "[UMBP] embedded page size {} KiB does not fit a {} MiB pool; using {} KiB " + "so the pool holds at least {} pages. Size the pool for the page (or set " + "UMBP_EMBEDDED_DRAM_PAGE_SIZE) to choose deliberately.", + requested / 1024, capacity_bytes / (1024 * 1024), fitted / 1024, kMinPages); + return fitted; +} + +} // namespace + +UMBPConfig WithEmbeddedDefaults(const UMBPConfig& config) { + // A standalone deployment is a different client entirely -- the pool lives in + // the server process -- so there is nothing here to fill in for it. + if (config.standalone_process.has_value()) return config; + + // Whether the store is embedded (in this process) and whether it has a master + // are independent axes, and DistributedClient already treats them that way: + // it derives local_only_ from master_address being empty and has no notion of + // "embedded" at all. So a caller who supplies distributed config is NOT + // opting out of in-process operation -- they are only opting out of having + // every field chosen for them. Fill the blanks they left rather than + // refusing to help, which is what forced a caller who wanted to set one field + // (dram_page_size, say) to hand-build the whole struct. + // + // A master address is the one thing that does change the answer: that node is + // a cluster member whose settings are the deployment's, not ours, so leave it + // exactly as given. + const bool synthesized = !config.distributed.has_value(); + UMBPConfig out = config; + UMBPDistributedConfig dist = synthesized ? UMBPDistributedConfig{} : config.distributed.value(); + if (!dist.master_config.master_address.empty()) return config; + + // Identity, which nothing but this function has any basis to choose. + if (dist.master_config.node_id.empty()) dist.master_config.node_id = EmbeddedNodeId(); + // Non-empty because Validate() requires it. Nothing reads it: it is only + // ever used to build the peer address this client would advertise, and with + // no master it advertises nothing. + if (dist.master_config.node_address.empty()) dist.master_config.node_address = "127.0.0.1"; + // Not a preference: with no master there is no MasterClient to beat to, so a + // heartbeat thread would be pure cost. The struct default is true because + // the common case is a cluster member. + dist.master_config.auto_heartbeat = false; + + // Page size, where 0 means "no opinion". Both the default and an explicit + // value are fitted to the pool: the cliff FitPageSizeToPool guards against -- + // a page bigger than the pool means ZERO pages and every put failing with + // NO_SPACE, silently -- does not care who chose the number, and honouring a + // deliberate-but-impossible value literally would be a worse answer than + // fitting it and saying so. Fitting is a no-op for anything that already + // leaves the pool at least kMinPages. + // + // Only for a pool no peer shares, though: peers in one tier must agree on + // page_size (see UMBPDistributedConfig::dram_page_size), and a value derived + // from THIS node's capacity would disagree with a peer sized differently, so + // an agreed size must survive untouched. peer_service_port == 0 means no + // gRPC peer service is bound, hence no peers to disagree with. + if (dist.peer_service_port == 0) { + const uint64_t requested = dist.dram_page_size > 0 ? dist.dram_page_size : EmbeddedPageSize(); + dist.dram_page_size = FitPageSizeToPool(requested, out.dram.capacity_bytes); + } + + if (synthesized) { + // Only reachable when the caller named no deployment at all, so there is no + // choice of theirs to overwrite. Everything that would need a peer is left + // empty, and each omission removes a whole subsystem rather than merely + // disabling it: no IO engine host means MoriIoEngine is never constructed; + // peer_service_port = 0 means no gRPC server is bound. What remains is one + // medium plus the local copy engines. + dist.io_engine.host.clear(); + dist.peer_service_port = 0; + + // ONE medium, and it is DRAM. This is the deliberate alignment with the + // distributed client: UMBP's routing plane does not tier within a node, so + // a node serves exactly one medium (see UMBPMedium in common/config.h). + // DRAM rather than a rule derived from ssd.enabled, because that flag + // defaults to true and would silently make every unconfigured caller an SSD + // node. Serving SSD is an explicit choice -- set distributed.medium. + dist.medium = UMBPMedium::DRAM; + + // Nothing can be remote, so every remote-path resource is off: no ranged + // scratch arenas to allocate or register, and no re-cache of fetches that + // cannot happen. + dist.ranged_scratch_size = 0; + dist.cache_remote_fetches = false; + dist.ranged_locality_prefetch = false; + dist.local_first = true; + } + + if (synthesized && out.ssd.enabled) { + MORI_UMBP_WARN( + "[UMBP] embedded deployment serves DRAM only; the configured SSD tier " + "(dir='{}', {} MiB) is NOT served. A node serves one medium: set " + "distributed.medium = SSD for an SSD node, or use a standalone server. " + "DRAM+SSD tiering on one node returns with the multi-backend work.", + out.ssd.storage_dir, out.ssd.capacity_bytes / (1024 * 1024)); + } + + MORI_UMBP_INFO("[UMBP] embedded deployment: node_id='{}' DRAM={} MiB page={} KiB (no master)", + dist.master_config.node_id, out.dram.capacity_bytes / (1024 * 1024), + dist.dram_page_size / 1024); + + out.distributed = dist; + return out; +} + +std::unique_ptr CreateUMBPClient(const UMBPConfig& config) { if (config.standalone_process.has_value()) { return std::make_unique(config); } - return std::make_unique(config); + // One class for every in-process store, whether it names a master, names none + // but configures itself, or names nothing at all. WithEmbeddedDefaults fills + // whatever was left blank and returns a cluster member's config untouched. + return std::make_unique(WithEmbeddedDefaults(config)); } } // namespace mori::umbp diff --git a/tests/cpp/cco/CMakeLists.txt b/tests/cpp/cco/CMakeLists.txt index 783f5accf..734215ac2 100644 --- a/tests/cpp/cco/CMakeLists.txt +++ b/tests/cpp/cco/CMakeLists.txt @@ -17,6 +17,10 @@ # link only mori_cco (+ pthread; HIP/MPI added per-target below). set(CCO_TEST_COMMON_LIBS mori_cco pthread) set(CCO_TEST_MPI_RANKS 8) +# These tests wait on real GPU/NIC completions, so a missing or busy device +# makes them block forever rather than fail. Without a cap one such test stalls +# the whole `ctest` run, which is how a hang here hides unrelated failures. +set(CCO_TEST_TIMEOUT_SECONDS 300) # MPI is a direct dependency of the cco tests (bootstrap + test_multiprocess). # Previously it came transitively via mori_application; find it here now that the @@ -81,11 +85,15 @@ function(add_cco_test src) INSTALL_RPATH_USE_LINK_PATH TRUE) add_test(NAME cco_${short} COMMAND ${target}) + set_tests_properties(cco_${short} PROPERTIES TIMEOUT + ${CCO_TEST_TIMEOUT_SECONDS}) if(uses_mpi AND WITH_MPI) add_test( NAME cco_${short}_mpi COMMAND ${MPIEXEC_EXECUTABLE} --allow-run-as-root ${MPIEXEC_NUMPROC_FLAG} ${CCO_TEST_MPI_RANKS} $) + set_tests_properties(cco_${short}_mpi + PROPERTIES TIMEOUT ${CCO_TEST_TIMEOUT_SECONDS}) endif() endfunction() diff --git a/tests/cpp/umbp/distributed/CMakeLists.txt b/tests/cpp/umbp/distributed/CMakeLists.txt index 2a56eadff..37d579c17 100644 --- a/tests/cpp/umbp/distributed/CMakeLists.txt +++ b/tests/cpp/umbp/distributed/CMakeLists.txt @@ -70,7 +70,16 @@ endif() target_link_libraries( test_umbp_client_metrics PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} gtest) -add_test(NAME umbp_client_metrics COMMAND test_umbp_client_metrics) +# Registered as two cases from one binary: the buffering suites talk to a mock +# service, while PoolClientLocalByteTracking stands up a real MasterServer and +# belongs to the integration tier below. The default case excludes by name, so a +# newly added suite lands there instead of silently going unrun. +add_test(NAME umbp_client_metrics + COMMAND test_umbp_client_metrics + --gtest_filter=-PoolClientLocalByteTrackingTest.*) +add_test(NAME umbp_client_metrics_pool_bytes + COMMAND test_umbp_client_metrics + --gtest_filter=PoolClientLocalByteTrackingTest.*) # ScopedRpcTimer / MasterClient RPC latency instrumentation. add_executable(test_umbp_master_client_rpc_latency @@ -110,7 +119,13 @@ endif() target_link_libraries( test_umbp_tags PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} gtest_main) -add_test(NAME umbp_tags COMMAND test_umbp_tags) +# Split like umbp_client_metrics: RealMasterServerTags needs a live MasterServer +# with a Prometheus port; the other suites run against the store or a mock +# service. +add_test(NAME umbp_tags + COMMAND test_umbp_tags --gtest_filter=-RealMasterServerTagsTest.*) +add_test(NAME umbp_tags_real_master + COMMAND test_umbp_tags --gtest_filter=RealMasterServerTagsTest.*) # -------------------------------------------------------------------------- # Unit tests (pure logic; no gRPC / RDMA). Linked against umbp_common only. @@ -120,6 +135,34 @@ add_executable(test_umbp_types test_types.cpp) target_link_libraries(test_umbp_types PRIVATE umbp_common gtest_main) add_test(NAME umbp_types COMMAND test_umbp_types) +add_executable(test_umbp_backend_policy_config test_backend_policy_config.cpp) +target_link_libraries(test_umbp_backend_policy_config PRIVATE umbp_common + gtest_main) +# The test loads the example policy shipped in configs/ instead of keeping its +# own copy, so the file and the parser cannot drift apart unnoticed. +target_compile_definitions( + test_umbp_backend_policy_config + PRIVATE + UMBP_EXAMPLE_POLICY_PATH="${CMAKE_SOURCE_DIR}/examples/umbp/multi_backend_policy.json" +) +add_test(NAME umbp_backend_policy_config + COMMAND test_umbp_backend_policy_config) + +# Versioned workload trace codec, deterministic payloads, synthetic generators, +# and the transport-neutral replay scheduler. These are hermetic unit tests; +# the real PoolClient adapter is covered by the integration smoke below. +add_executable(test_umbp_workload_trace test_umbp_workload_trace.cpp) +add_dependencies(test_umbp_workload_trace umbp_workload_proto) +target_link_libraries(test_umbp_workload_trace PRIVATE umbp_workload_bench + gtest_main) +add_test(NAME umbp_workload_trace COMMAND test_umbp_workload_trace) + +add_executable(test_umbp_workload_runner test_umbp_workload_runner.cpp) +add_dependencies(test_umbp_workload_runner umbp_workload_proto) +target_link_libraries(test_umbp_workload_runner PRIVATE umbp_workload_bench + gtest_main) +add_test(NAME umbp_workload_runner COMMAND test_umbp_workload_runner) + # env_time helpers are header-only; tests only exercise the parser path. NOTE: # Production callers cache results in function-local statics that cannot be # reset mid-process; any multi-value business test must fork. @@ -137,15 +180,26 @@ target_link_libraries(test_umbp_route_put_strategy PRIVATE umbp_common gtest_main) add_test(NAME umbp_route_put_strategy COMMAND test_umbp_route_put_strategy) -# Master eviction policy plug-in: default LRU behavior + the EvictionManager DI -# seam. The seam test calls EvictionManager::RunOnce() directly (a public -# single-pass trigger), so no special compile flags are needed. +# Master eviction policy plug-in: the default LRU victim selection in isolation, +# given candidates and a per-(node, tier) byte budget. add_executable(test_umbp_master_evict_strategy test_master_evict_strategy.cpp) target_link_libraries(test_umbp_master_evict_strategy PRIVATE umbp_common gtest_main) add_test(NAME umbp_master_evict_strategy COMMAND test_umbp_master_evict_strategy) +# The master-to-peer eviction chain, which the two tests above and test_peer_pool +# only cover from their own ends: watermark detection, victim selection, and the +# dispatch, wired into a real PeerPool so on_evict offload is driven the way a +# deployed master drives it. RunOnce() is private, so rounds are driven by +# Start() at the one-second floor rather than called directly. +add_executable(test_umbp_master_eviction_chain test_master_eviction_chain.cpp) +target_link_libraries(test_umbp_master_eviction_chain PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_umbp_master_eviction_chain PRIVATE cxx_std_17) +add_test(NAME umbp_master_eviction_chain + COMMAND test_umbp_master_eviction_chain) + # PageBitmapAllocator (peer-side per-tier page allocator primitive). add_executable(test_umbp_page_bitmap_allocator test_page_bitmap_allocator.cpp) target_link_libraries(test_umbp_page_bitmap_allocator PRIVATE umbp_common @@ -158,15 +212,6 @@ add_test(NAME umbp_page_bitmap_allocator # Registered as CTest cases under the "integration" label. # -------------------------------------------------------------------------- -add_executable(test_umbp_peer_service test_peer_service.cpp) -if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") - target_link_options(test_umbp_peer_service PRIVATE -Wl,--no-as-needed) -endif() -target_link_libraries( - test_umbp_peer_service PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} - ${_TEST_GRPCPP_LIB} gtest_main) -add_test(NAME umbp_peer_service COMMAND test_umbp_peer_service) - add_executable(test_umbp_cross_node_smoke test_cross_node_smoke.cpp) if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") target_link_options(test_umbp_cross_node_smoke PRIVATE -Wl,--no-as-needed) @@ -189,8 +234,103 @@ target_link_libraries( add_test(NAME umbp_pool_client_batch_put COMMAND test_umbp_pool_client_batch_put) -set_tests_properties(umbp_peer_service umbp_cross_node_smoke - umbp_pool_client_batch_put PROPERTIES LABELS "integration") +# End-to-end workload replay through a live in-process master, two PoolClients, +# and weighted same-tier backends. This needs the real gRPC + transfer path. +add_executable(test_umbp_tier_benchmark test_umbp_tier_benchmark.cpp) +add_dependencies(test_umbp_tier_benchmark umbp_workload_proto) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_umbp_tier_benchmark PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_umbp_tier_benchmark + PRIVATE umbp_core umbp_workload_bench ${_TEST_PROTOBUF_LIB} + ${_TEST_GRPCPP_LIB} gtest_main) +add_test(NAME umbp_tier_benchmark COMMAND test_umbp_tier_benchmark) + +# test_umbp_medium_selection — one medium per node: the UMBPMedium -> TierType +# lowering, the per-medium validation, and (live, 2-node) the registry holding +# exactly one backend plus a cross-node BatchPut onto an SSD-only peer. +add_executable(test_umbp_medium_selection test_medium_selection.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_umbp_medium_selection PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_umbp_medium_selection + PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} + gtest_main) +add_test(NAME umbp_medium_selection COMMAND test_umbp_medium_selection) + +# test_umbp_local_first — PoolClientConfig::local_first: a fully-local batch +# never contacts the master. Asserted by STOPPING the master and reading +# anyway, with a local_first=false twin that must fail on the same data. +add_executable(test_umbp_local_first test_local_first.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_umbp_local_first PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_umbp_local_first + PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} + gtest_main) +add_test(NAME umbp_local_first COMMAND test_umbp_local_first) + +# test_umbp_no_master — the distributed client with master_address empty: same +# data plane, no routing/registration/heartbeat, a local miss is final. No +# master is started in that file, so reaching for one hangs or crashes rather +# than passing quietly. +add_executable(test_umbp_no_master test_no_master.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_umbp_no_master PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_umbp_no_master + PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} + gtest_main) +add_test(NAME umbp_no_master COMMAND test_umbp_no_master) + +# test_umbp_pool_client_ranges — ranged multi-buffer I/O through PoolClient: +# local range reads/writes against the node's own medium, the remote path via +# the scratch arena, range validation, and GPU caller buffers. +add_executable(test_umbp_pool_client_ranges test_pool_client_ranges.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_umbp_pool_client_ranges PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_umbp_pool_client_ranges + PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} + gtest_main) +add_test(NAME umbp_pool_client_ranges COMMAND test_umbp_pool_client_ranges) + +# test_umbp_ssd_ranged_io — ranged multi-buffer I/O on an SSD-medium node +# (doc/design-ssd-ranged-io.md, D0): the capability flag now turning on the +# arena alone, plus local and remote ranged put/get against an SSD-only peer. +add_executable(test_umbp_ssd_ranged_io test_ssd_ranged_io.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_umbp_ssd_ranged_io PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_umbp_ssd_ranged_io + PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} + gtest_main) +add_test(NAME umbp_ssd_ranged_io COMMAND test_umbp_ssd_ranged_io) + +# test_umbp_range_map — the object-range -> page-range arithmetic on its own. +# Header-only and node-free, so it is a plain unit test rather than an +# integration one: the shapes that matter (short last page, range ending on the +# object's last byte, one byte past the end) are awkward to reach through a +# live two-node fixture, and a mistake there corrupts a KV block silently +# rather than failing. +add_executable(test_umbp_range_map test_range_map.cpp) +target_link_libraries(test_umbp_range_map PRIVATE umbp_common gtest_main) +add_test(NAME umbp_range_map COMMAND test_umbp_range_map) + +set_tests_properties(umbp_cross_node_smoke umbp_pool_client_batch_put + umbp_medium_selection umbp_tier_benchmark + umbp_client_metrics_pool_bytes umbp_tags_real_master + umbp_pool_client_ranges umbp_ssd_ranged_io + PROPERTIES + LABELS "integration" + ENVIRONMENT + "MORI_IO_QP_MAX_SEND_WR=1024;MORI_IO_QP_MAX_CQE=4096") # -------------------------------------------------------------------------- # Benchmarks (manual; built but NOT registered with add_test). @@ -223,6 +363,20 @@ target_link_libraries( PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} gtest_main) +# Layer-wise ranged-restore bench: drives BatchGetRanges the way the sglang +# tree connector does (N ranks walking layer groups over the same keys) and +# reports TTFL, wire bytes and repeat-remote alongside latency. Standalone +# executable, not a CTest target. +add_executable(bench_umbp_pool_client_ranges bench_pool_client_ranges.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(bench_umbp_pool_client_ranges PRIVATE + -Wl,--no-as-needed) +endif() +target_link_libraries( + bench_umbp_pool_client_ranges + PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} + gtest_main) + # RoutePut strategy bench: compares batch put / put-then-get performance and # placement distribution across UMBP_ROUTE_PUT_SELECT_ALGO x # UMBP_ROUTE_PUT_NODE_AFFINITY under roomy/tight capacity regimes. Standalone @@ -244,23 +398,21 @@ target_link_libraries(bench_umbp_route_strategy PRIVATE umbp_common) # =========================================================================== # Metadata-store / external-KV / SSD-tier suite (migrated from src/umbp/tests). -# These link umbp_common (pure logic) except test_peer_ssd_read_rpc, which -# exercises the gRPC peer service and links umbp_core + protobuf/gRPC. Each -# binary is registered as a single ctest via add_test (see the NOTE above). +# These link umbp_common (pure logic). Each binary is registered as a single +# ctest via add_test (see the NOTE above). # =========================================================================== -# test_master_metadata_store_interface — Phase 1 compile/instantiation gate for -# IMasterMetadataStore. Includes an isolated self-compile TU (proves the header -# is self-contained) and a GMock mock instantiation/signature-completeness test. -add_executable( - test_master_metadata_store_interface test_master_metadata_store_interface.cpp - master_metadata_store_self_compile.cpp) -target_include_directories(test_master_metadata_store_interface - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) -target_link_libraries(test_master_metadata_store_interface - PRIVATE umbp_common GTest::gmock GTest::gtest_main) -target_compile_features(test_master_metadata_store_interface PRIVATE cxx_std_17) -add_test(NAME test_master_metadata_store_interface COMMAND test_master_metadata_store_interface) +# IMasterMetadataStore header self-compile gate: this TU includes only the +# interface header, so building it proves the header is self-contained. No +# add_test — there is nothing to assert at runtime, and the interface's +# polymorphic use is exercised for real by test_in_memory_master_metadata_store, +# which drives every case through IMasterMetadataStore&. +add_library(umbp_master_metadata_store_self_compile OBJECT + master_metadata_store_self_compile.cpp) +target_link_libraries(umbp_master_metadata_store_self_compile + PRIVATE umbp_common) +target_compile_features(umbp_master_metadata_store_self_compile + PRIVATE cxx_std_17) # test_in_memory_master_metadata_store — Phase 2 behavioral suite for the # InMemoryMasterMetadataStore implementation of IMasterMetadataStore (§6a). @@ -271,12 +423,91 @@ target_link_libraries(test_in_memory_master_metadata_store target_compile_features(test_in_memory_master_metadata_store PRIVATE cxx_std_17) add_test(NAME test_in_memory_master_metadata_store COMMAND test_in_memory_master_metadata_store) -# test_peer_dram_allocator -add_executable(test_peer_dram_allocator test_peer_dram_allocator.cpp) -target_link_libraries(test_peer_dram_allocator PRIVATE umbp_common - GTest::gtest_main) -target_compile_features(test_peer_dram_allocator PRIVATE cxx_std_17) -add_test(NAME test_peer_dram_allocator COMMAND test_peer_dram_allocator) +# test_page_backend — PageBackend (formerly PeerDramAllocator), one instance +# == one medium (backend-agnostic refactor Phase 2) +add_executable(test_page_backend test_page_backend.cpp) +target_link_libraries(test_page_backend PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_page_backend PRIVATE cxx_std_17) +add_test(NAME test_page_backend COMMAND test_page_backend) + +# BackendRegistry dispatch and the heartbeat event aggregation MediumBackend +# absorbed from OwnedLocationSource. MockBackend is only the test double. +add_executable(test_backend_registry test_backend_registry.cpp) +target_link_libraries(test_backend_registry PRIVATE umbp_common GTest::gtest_main) +target_compile_features(test_backend_registry PRIVATE cxx_std_17) +add_test(NAME test_backend_registry COMMAND test_backend_registry) + +# test_peer_pool — implicit default Pool, SingleBackendPolicy compatibility, +# logical key placement, named same-tier backend dispatch, and recovery fallback. +# Registered under the umbp_ prefix so CI's ctest filter picks it up. +add_executable(test_peer_pool test_peer_pool.cpp) +target_link_libraries(test_peer_pool PRIVATE umbp_common GTest::gtest_main) +target_compile_features(test_peer_pool PRIVATE cxx_std_17) +add_test(NAME umbp_peer_pool COMMAND test_peer_pool) + +# test_transfer_engine — the Phase 6 transfer layer without gRPC or RDMA: +# LocalCopyEngine's planner/executor and CompositeTransferEngine's fan-out +# registration + per-pair engine selection. MoriIoEngine needs a real fabric +# and is covered end-to-end by test_cross_node_smoke instead. +add_executable(test_transfer_engine test_transfer_engine.cpp) +target_link_libraries(test_transfer_engine PRIVATE umbp_common GTest::gtest_main) +target_compile_features(test_transfer_engine PRIVATE cxx_std_17) +add_test(NAME test_transfer_engine COMMAND test_transfer_engine) + +# test_component_metrics — the observability seam: MetricPublisher's delta / +# gauge / baseline rules, InstrumentedBackend deriving the generic backend +# series from the MediumBackend interface alone (asserted against MockBackend, +# which carries no metrics code — that is the genericity guarantee), and the +# composite transfer engine's per-engine byte/plan accounting. +add_executable(test_component_metrics test_component_metrics.cpp) +target_link_libraries(test_component_metrics PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_component_metrics PRIVATE cxx_std_17) +add_test(NAME test_component_metrics COMMAND test_component_metrics) + +# test_gds_engine — GdsEngine routing/planning (no GPU) plus a GPU + hipFile +# read round trip (skipped without a device, or on a filesystem that rejects +# O_DIRECT). Built only when hipfile is available; umbp_common already carries +# hip::hipfile and the UMBP_ENABLE_GDS definition in that case. +if(TARGET hip::hipfile) + add_executable(test_gds_engine test_gds_engine.cpp) + target_link_libraries(test_gds_engine PRIVATE umbp_common hip::host GTest::gtest_main) + target_compile_features(test_gds_engine PRIVATE cxx_std_17) + add_test(NAME test_gds_engine COMMAND test_gds_engine) +endif() + +# test_hbm_backend — the HBM medium and the engine that had to exist for it: +# HbmCopyEngine's selection/planner (no GPU needed) plus real H2D/D2H/D2D bytes +# and a Put/Get round trip against an hipMalloc'd pool (skipped without a GPU). +add_executable(test_hbm_backend test_hbm_backend.cpp) +target_link_libraries(test_hbm_backend PRIVATE umbp_common hip::host + GTest::gtest_main) +target_compile_features(test_hbm_backend PRIVATE cxx_std_17) +add_test(NAME test_hbm_backend COMMAND test_hbm_backend) + +# test_ssd_backend — SSD as a MediumBackend, staged through registered host +# DRAM: Put spills to the tier, Get stages back under a read lease, and the +# staging arena's exhaustion / reaping / eviction interactions. +add_executable(test_ssd_backend test_ssd_backend.cpp) +target_link_libraries(test_ssd_backend PRIVATE umbp_common GTest::gtest_main) +target_compile_features(test_ssd_backend PRIVATE cxx_std_17) +add_test(NAME test_ssd_backend COMMAND test_ssd_backend) + +# test_peer_service_dispatch — PeerServiceServer dispatch through +# BackendRegistry: tier compatibility, named same-tier instances, the +# backend-id-tagged opaque slot_id that makes Commit/Abort unambiguous, +# mixed-backend batch ordering, and resolve/evict fan-out. Drives the real gRPC +# surface, so it needs the generated proto + protobuf. +add_executable(test_peer_service_dispatch test_peer_service_dispatch.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_peer_service_dispatch PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_peer_service_dispatch PRIVATE umbp_common ${_TEST_PROTOBUF_LIB} + GTest::gtest_main) +target_compile_features(test_peer_service_dispatch PRIVATE cxx_std_17) +add_test(NAME test_peer_service_dispatch COMMAND test_peer_service_dispatch) # test_router_dedup — master-side BatchRoutePut dedup via IMasterMetadataStore add_executable(test_router_dedup test_router_dedup.cpp) @@ -284,13 +515,21 @@ target_link_libraries(test_router_dedup PRIVATE umbp_common GTest::gtest_main) target_compile_features(test_router_dedup PRIVATE cxx_std_17) add_test(NAME test_router_dedup COMMAND test_router_dedup) -# test_peer_ssd_manager — SSD tier ownership + owned-location source +# test_peer_ssd_manager — SSD tier ownership (reachable again: SsdBackend is the +# MediumBackend adapter that drives it) add_executable(test_peer_ssd_manager test_peer_ssd_manager.cpp) target_link_libraries(test_peer_ssd_manager PRIVATE umbp_common GTest::gtest_main) target_compile_features(test_peer_ssd_manager PRIVATE cxx_std_17) add_test(NAME test_peer_ssd_manager COMMAND test_peer_ssd_manager) +# test_peer_ssd_single_flight — same-key read coalescing (leader/follower) +add_executable(test_peer_ssd_single_flight test_peer_ssd_single_flight.cpp) +target_link_libraries(test_peer_ssd_single_flight PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_peer_ssd_single_flight PRIVATE cxx_std_17) +add_test(NAME test_peer_ssd_single_flight COMMAND test_peer_ssd_single_flight) + # test_peer_ssd_eviction — LRU + watermark eviction + in-flight guard add_executable(test_peer_ssd_eviction test_peer_ssd_eviction.cpp) target_link_libraries(test_peer_ssd_eviction PRIVATE umbp_common @@ -305,42 +544,6 @@ target_link_libraries(test_ssd_copy_pipeline PRIVATE umbp_common target_compile_features(test_ssd_copy_pipeline PRIVATE cxx_std_17) add_test(NAME test_ssd_copy_pipeline COMMAND test_ssd_copy_pipeline) -# test_tier_priority_route_get — RouteGet tier-priority strategy -add_executable(test_tier_priority_route_get test_tier_priority_route_get.cpp) -target_link_libraries(test_tier_priority_route_get PRIVATE umbp_common - GTest::gtest_main) -target_compile_features(test_tier_priority_route_get PRIVATE cxx_std_17) -add_test(NAME test_tier_priority_route_get COMMAND test_tier_priority_route_get) - -# test_peer_ssd_read_rpc — SSD read RPC over a gRPC loopback. Needs the gRPC -# peer service + generated stubs (umbp_core) and protobuf/gRPC. -add_executable(test_peer_ssd_read_rpc test_peer_ssd_read_rpc.cpp) -if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") - target_link_options(test_peer_ssd_read_rpc PRIVATE -Wl,--no-as-needed) -endif() -target_link_libraries( - test_peer_ssd_read_rpc PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} - ${_TEST_GRPCPP_LIB} GTest::gtest_main) -target_compile_features(test_peer_ssd_read_rpc PRIVATE cxx_std_17) -add_test(NAME test_peer_ssd_read_rpc COMMAND test_peer_ssd_read_rpc) - -# test_ssd_read_lease_gating — pure reader-side lease gating decision logic -# (ssd_read_lease.h). Header-only under test (no gRPC / RDMA deps). -add_executable(test_ssd_read_lease_gating test_ssd_read_lease_gating.cpp) -target_link_libraries(test_ssd_read_lease_gating PRIVATE umbp_common - GTest::gtest_main) -target_compile_features(test_ssd_read_lease_gating PRIVATE cxx_std_17) -add_test(NAME test_ssd_read_lease_gating COMMAND test_ssd_read_lease_gating) - -# test_ssd_reliability — cross-component reliability: owned-source DRAM+SSD -# merge, SSD evict -> REMOVE -> master index convergence, tier-priority over the -# real index, crash-restart discard, and observability counters. -add_executable(test_ssd_reliability test_ssd_reliability.cpp) -target_link_libraries(test_ssd_reliability PRIVATE umbp_common - GTest::gtest_main) -target_compile_features(test_ssd_reliability PRIVATE cxx_std_17) -add_test(NAME test_ssd_reliability COMMAND test_ssd_reliability) - # test_cache_remote_admission — cache_remote_fetches admission-gate unit test # (hermetic; no GPU/RDMA/master). add_executable(test_cache_remote_admission test_cache_remote_admission.cpp) @@ -425,17 +628,15 @@ if(USE_REDIS_BACKEND) target_link_libraries(test_sharded_read PRIVATE umbp_common GTest::gtest_main) target_compile_features(test_sharded_read PRIVATE cxx_std_17) add_test(NAME test_sharded_read COMMAND test_sharded_read) +endif() - # redis-plus-plus link/build smoke test. Proves the submodule compiles with - # this toolchain and links (static redis++ + system hiredis) before any store - # code depends on it; at runtime it connects to UMBP_REDIS_URI if reachable, - # else skips. redis++_static carries its own include dirs; hiredis must be - # linked explicitly (the system hiredis ships no CMake config, so redis++ does - # not propagate it for the static lib). - add_executable(test_redispp_smoke test_redispp_smoke.cpp) - target_link_libraries(test_redispp_smoke - PRIVATE redis++_static ${HIREDIS_LIB} GTest::gtest_main) - target_include_directories(test_redispp_smoke PRIVATE ${HIREDIS_INCLUDE_DIR}) - target_compile_features(test_redispp_smoke PRIVATE cxx_std_17) - add_test(NAME test_redispp_smoke COMMAND test_redispp_smoke) +# GDS-vs-staging A/B for SSD reads into device memory. Run twice, once with +# UMBP_ENABLE_GDS=1, and diff the CSV. Standalone executable; not a CTest +# target, because it needs a real block-backed --dir to measure anything. +add_executable(bench_umbp_gds_ssd_get bench_gds_ssd_get.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(bench_umbp_gds_ssd_get PRIVATE -Wl,--no-as-needed) endif() +target_link_libraries( + bench_umbp_gds_ssd_get + PRIVATE umbp_core umbp_common hip::host ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB}) diff --git a/tests/cpp/umbp/distributed/bench_gds_ssd_get.cpp b/tests/cpp/umbp/distributed/bench_gds_ssd_get.cpp new file mode 100644 index 000000000..7eef8b9df --- /dev/null +++ b/tests/cpp/umbp/distributed/bench_gds_ssd_get.cpp @@ -0,0 +1,326 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// GDS-vs-staging A/B for SSD reads. Drives IUMBPClient::BatchGet against an +// SSD-only node and reports the two passes that actually differ: +// +// cold -- first touch of each key. The staging path must do a device read +// into the arena and then a second copy out to the caller; the GDS +// path DMAs the range straight into the destination. This is the +// only pass where zero-copy can win. +// warm -- the same keys again, now holding a read lease. BatchResolve +// serves a lease before it ever considers a file ref, so BOTH modes +// are an arena copy here and the numbers should converge. Reported +// because "GDS made no difference" is the expected answer for hot +// data, and it is worth being able to show that rather than assert +// it. +// +// The mode is not a flag: run the binary twice, once with UMBP_ENABLE_GDS=1 +// and once without, and diff the CSV. That keeps the two arms honest -- same +// binary, same node, same data, one environment variable apart. +// +// Usage: +// bench_gds_ssd_get [--dir PATH] [--value-bytes N] [--batch N] [--keys N] +// [--capacity-mb N] [--page-bytes N] [--slots N] +// [--host-dst] [--io-backend posix|uring] [--no-header] +// +// --dir must name a real block-backed filesystem. The SSD tier probes for +// O_DIRECT and silently falls back to buffered I/O when the probe fails, and +// a buffered fd is never eligible for the GDS fastpath (RecordLocation:: +// direct_io gates it), so pointing this at tmpfs measures nothing. +// +// CSV (stdout): +// mode,pass,dst,value_bytes,batch,keys,slots,wall_ms,gibps,us_per_key + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/common/config.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/umbp_client.h" + +namespace fs = std::filesystem; + +using mori::umbp::CreateUMBPClient; +using mori::umbp::IUMBPClient; +using mori::umbp::MasterServer; +using mori::umbp::MasterServerConfig; +using mori::umbp::UMBPConfig; +using mori::umbp::UMBPDistributedConfig; +using mori::umbp::UMBPIoBackend; +using mori::umbp::UMBPMedium; + +namespace { + +struct Options { + std::string dir; + size_t value_bytes = 1ULL << 20; // 1 MiB + size_t batch = 32; + size_t keys = 512; + size_t capacity_mb = 8192; + size_t page_bytes = 64ULL * 1024; + // Staging pages the SSD backend may hold in flight. 0 = size it from the + // batch. The arena is the staging arm's hard concurrency limit: an object of + // N pages holds N entries until its read lease expires, and a batch that does + // not fit resolves kBusy and backs off (up to 50 ms a try). Undersizing it + // measures the backoff, not the copy -- so size it, and expose the knob so + // arena pressure can be studied deliberately rather than stumbled into. + size_t slots = 0; + bool host_dst = false; + bool header = true; + UMBPIoBackend io_backend = UMBPIoBackend::Posix; +}; + +[[noreturn]] void Die(const std::string& why) { + std::fprintf(stderr, "bench_gds_ssd_get: %s\n", why.c_str()); + std::exit(2); +} + +#define HIP_OK(expr) \ + do { \ + hipError_t _e = (expr); \ + if (_e != hipSuccess) Die(std::string(#expr) + ": " + hipGetErrorString(_e)); \ + } while (0) + +// A device buffer, or a host one when --host-dst asks for the comparison arm. +// Host is worth measuring precisely because it CANNOT use GDS: the file engine +// claims only (file, GPU), so a host destination re-resolves onto staged pages. +class Destination { + public: + Destination(size_t bytes, bool host) : bytes_(bytes), host_(host) { + if (host_) { + ptr_ = std::malloc(bytes_); + if (ptr_ == nullptr) Die("malloc failed"); + } else { + HIP_OK(hipMalloc(&ptr_, bytes_)); + } + } + ~Destination() { + if (ptr_ == nullptr) return; + if (host_) { + std::free(ptr_); + } else { + (void)hipFree(ptr_); + } + } + Destination(const Destination&) = delete; + Destination& operator=(const Destination&) = delete; + + void* get() const { return ptr_; } + bool host() const { return host_; } + + private: + size_t bytes_ = 0; + bool host_ = false; + void* ptr_ = nullptr; +}; + +double Seconds(std::chrono::steady_clock::time_point a, std::chrono::steady_clock::time_point b) { + return std::chrono::duration_cast>(b - a).count(); +} + +// One pass over every key, in batches. Returns wall seconds; fails loudly, +// because a benchmark that silently measures misses is worse than no number. +double RunPass(IUMBPClient* client, const std::vector& keys, const Destination& dst, + const Options& opt) { + auto* base = static_cast(dst.get()); + const auto t0 = std::chrono::steady_clock::now(); + for (size_t start = 0; start < keys.size(); start += opt.batch) { + const size_t n = std::min(opt.batch, keys.size() - start); + std::vector batch_keys(keys.begin() + start, keys.begin() + start + n); + std::vector dsts(n); + std::vector sizes(n, opt.value_bytes); + for (size_t i = 0; i < n; ++i) { + dsts[i] = reinterpret_cast(base + i * opt.value_bytes); + } + auto ok = client->BatchGet(batch_keys, dsts, sizes); + for (size_t i = 0; i < n; ++i) { + if (i >= ok.size() || !ok[i]) Die("BatchGet miss for key " + batch_keys[i]); + } + } + return Seconds(t0, std::chrono::steady_clock::now()); +} + +void Report(const Options& opt, const char* pass, double sec) { + const double bytes = static_cast(opt.keys) * static_cast(opt.value_bytes); + const double gibps = bytes / (sec > 0 ? sec : 1e-12) / (1024.0 * 1024.0 * 1024.0); + const double us_per_key = sec * 1e6 / static_cast(opt.keys); + const char* mode = std::getenv("UMBP_ENABLE_GDS"); + const bool on = mode != nullptr && (std::string(mode) == "1" || std::string(mode) == "true" || + std::string(mode) == "on"); + std::printf("%s,%s,%s,%zu,%zu,%zu,%zu,%.3f,%.3f,%.2f\n", on ? "gds" : "staging", pass, + opt.host_dst ? "host" : "device", opt.value_bytes, opt.batch, opt.keys, opt.slots, + sec * 1000.0, gibps, us_per_key); + std::fflush(stdout); +} + +} // namespace + +int main(int argc, char** argv) { + Options opt; + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto next = [&]() -> std::string { + if (i + 1 >= argc) Die("missing value for " + a); + return argv[++i]; + }; + if (a == "--dir") + opt.dir = next(); + else if (a == "--value-bytes") + opt.value_bytes = std::stoull(next()); + else if (a == "--batch") + opt.batch = std::stoull(next()); + else if (a == "--keys") + opt.keys = std::stoull(next()); + else if (a == "--capacity-mb") + opt.capacity_mb = std::stoull(next()); + else if (a == "--page-bytes") + opt.page_bytes = std::stoull(next()); + else if (a == "--slots") + opt.slots = std::stoull(next()); + else if (a == "--host-dst") + opt.host_dst = true; + else if (a == "--no-header") + opt.header = false; + else if (a == "--io-backend") { + const std::string b = next(); + if (b == "uring" || b == "io_uring") + opt.io_backend = UMBPIoBackend::IoUring; + else if (b == "posix") + opt.io_backend = UMBPIoBackend::Posix; + else + Die("--io-backend must be posix or uring"); + } else { + Die("unknown flag: " + a); + } + } + if (opt.dir.empty()) Die("--dir is required and must be on a real block device"); + if (opt.batch == 0 || opt.keys == 0 || opt.value_bytes == 0) Die("sizes must be non-zero"); + if (opt.page_bytes == 0) Die("--page-bytes must be non-zero"); + const size_t pages_per_value = (opt.value_bytes + opt.page_bytes - 1) / opt.page_bytes; + // Default: the whole working set. A staged page is held until its read + // lease expires, not until the copy finishes, so an arena sized to one batch + // fills after a couple of batches and every later resolve returns kBusy and + // sleeps -- which measures the backoff, not the data path. Pass --slots + // explicitly to study that regime on purpose; it is a real one, and it is + // where zero-copy helps most, because a file ref occupies no arena at all. + if (opt.slots == 0) opt.slots = opt.keys * pages_per_value; + + const fs::path store = fs::path(opt.dir) / ("bench_gds_" + std::to_string(::getpid())); + fs::remove_all(store); + fs::create_directories(store); + + MasterServerConfig master_cfg; + master_cfg.listen_address = "0.0.0.0:0"; + master_cfg.registry_config.heartbeat_ttl = std::chrono::seconds(4); + auto master = std::make_unique(std::move(master_cfg)); + std::thread master_thread([&] { master->Run(); }); + for (int i = 0; i < 200 && master->GetBoundPort() == 0; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + if (master->GetBoundPort() == 0) Die("master failed to start"); + + UMBPConfig config; + config.dram.use_hugepages = false; + config.dram.capacity_bytes = 8 << 20; // ignored on an SSD node; see medium selection + config.ssd.storage_dir = store.string(); + config.ssd.capacity_bytes = opt.capacity_mb * 1024ULL * 1024ULL; + config.ssd.io.backend = opt.io_backend; + + UMBPDistributedConfig distributed; + distributed.master_config.node_id = "bench-gds-ssd"; + distributed.master_config.node_address = "127.0.0.1"; + distributed.master_config.master_address = "localhost:" + std::to_string(master->GetBoundPort()); + distributed.io_engine.host = "0.0.0.0"; + distributed.io_engine.port = 0; + distributed.peer_service_port = 0; + distributed.medium = UMBPMedium::SSD; + distributed.dram_page_size = opt.page_bytes; + // ssd_staging_buffer_slots is what PoolClient hands SsdBackend as + // Config::staging_pages -- the LOCAL arena, in pages, not the remote one. + distributed.ssd_staging_buffer_slots = static_cast(opt.slots); + distributed.ssd_staging_buffer_size = + std::max(256ULL << 20, opt.batch * opt.value_bytes * 4); + config.distributed = std::move(distributed); + + std::string err; + if (!config.Validate(&err)) Die("config invalid: " + err); + auto client = CreateUMBPClient(config); + if (client == nullptr) Die("CreateUMBPClient returned null"); + + // ---- populate ------------------------------------------------------------- + std::vector payload(opt.value_bytes, 0x5A); + std::vector keys; + keys.reserve(opt.keys); + for (size_t i = 0; i < opt.keys; ++i) keys.push_back("gdsbench/" + std::to_string(i)); + for (size_t i = 0; i < opt.keys; ++i) { + if (!client->Put(keys[i], reinterpret_cast(payload.data()), opt.value_bytes)) { + Die("Put failed for " + keys[i]); + } + } + + Destination dst(opt.batch * opt.value_bytes, opt.host_dst); + if (!opt.host_dst) { + client->RegisterMemory(reinterpret_cast(dst.get()), opt.batch * opt.value_bytes, + mori::io::MemoryLocationType::GPU, 0); + } + + if (opt.header) { + std::printf("mode,pass,dst,value_bytes,batch,keys,slots,wall_ms,gibps,us_per_key\n"); + } + Report(opt, "cold", RunPass(client.get(), keys, dst, opt)); + Report(opt, "warm", RunPass(client.get(), keys, dst, opt)); + + // Correctness, after the timed passes so it costs them nothing. The device + // arm is the one that matters: it is the only path that can have come from a + // hipFileRead, and nothing else in the tree checks those bytes end to end + // through the public client API. + { + std::vector back(opt.value_bytes, 0); + if (opt.host_dst) { + std::memcpy(back.data(), dst.get(), opt.value_bytes); + } else { + HIP_OK(hipMemcpy(back.data(), dst.get(), opt.value_bytes, hipMemcpyDeviceToHost)); + } + if (std::memcmp(back.data(), payload.data(), opt.value_bytes) != 0) { + Die("VERIFY FAILED: bytes read back do not match what was written"); + } + std::fprintf(stderr, "verify ok (%zu bytes)\n", opt.value_bytes); + } + + client->Close(); + master->Shutdown(); + if (master_thread.joinable()) master_thread.join(); + std::error_code ec; + fs::remove_all(store, ec); + return 0; +} diff --git a/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp b/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp index 8f17a3e9d..7753e0be1 100644 --- a/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp +++ b/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp @@ -806,14 +806,12 @@ int main(int argc, char** argv) { const size_t io_bytes = o.batch * o.page_bytes; // src / dst staging // ---- clients ---- - std::vector> dram_bufs(N); std::vector> src_bufs(N); std::vector> dst_bufs(N); std::vector> clients(N); const bool need_dst = (o.get_mode != GetMode::kExists); for (size_t id = 0; id < N; ++id) { - dram_bufs[id].assign(peer_buf_bytes, 0); src_bufs[id].assign(io_bytes, 0x5a); if (need_dst) dst_bufs[id].assign(io_bytes, 0); @@ -828,8 +826,7 @@ int main(int argc, char** argv) { cfg.io_engine.port = 0; cfg.peer_service_port = NextPeerServicePort(); cfg.dram_page_size = o.page_bytes; - cfg.dram_buffers = {{dram_bufs[id].data(), dram_bufs[id].size()}}; - cfg.tier_capacities = {{TierType::DRAM, {dram_bufs[id].size(), dram_bufs[id].size()}}}; + cfg.dram.buffer_sizes = {peer_buf_bytes}; clients[id] = std::make_unique(std::move(cfg)); if (!clients[id]->Init()) { std::fprintf(stderr, "client %zu init failed\n", id); diff --git a/tests/cpp/umbp/distributed/bench_pool_client_batch_get.cpp b/tests/cpp/umbp/distributed/bench_pool_client_batch_get.cpp index 0b71554dc..ed808eba1 100644 --- a/tests/cpp/umbp/distributed/bench_pool_client_batch_get.cpp +++ b/tests/cpp/umbp/distributed/bench_pool_client_batch_get.cpp @@ -145,7 +145,6 @@ bool ParseArgs(int argc, char** argv, BenchOpts* o) { // multi_peer scenarios the caller pre-registers caller_buf_ as dst; // in all_stg it doesn't, so REMOTE_STG path activates. struct PeerNode { - std::vector dram; // master-managed exportable DRAM (data lives here) std::vector seed; // src for the seed BatchPut (registered for ZC Put) std::unique_ptr client; }; @@ -154,7 +153,7 @@ class Cluster { public: Cluster(size_t page_bytes, size_t target_dram_bytes, size_t num_peers, size_t caller_buf_bytes, size_t caller_local_bytes, size_t seed_bytes_per_peer) - : page_bytes_(page_bytes), caller_buf_(caller_buf_bytes), caller_local_(caller_local_bytes) { + : page_bytes_(page_bytes), caller_buf_(caller_buf_bytes) { MasterServerConfig mcfg; mcfg.listen_address = "0.0.0.0:0"; // The master index is eventually consistent: a committed key becomes @@ -182,8 +181,7 @@ class Cluster { cc.io_engine.port = 0; cc.peer_service_port = NextPeerServicePort(); cc.dram_page_size = page_bytes; - cc.dram_buffers = {{caller_local_.data(), caller_local_.size()}}; - cc.tier_capacities = {{TierType::DRAM, {caller_local_.size(), caller_local_.size()}}}; + cc.dram.buffer_sizes = {caller_local_bytes}; caller_ = std::make_unique(std::move(cc)); if (!caller_->Init()) { std::cerr << "caller init failed\n"; @@ -192,7 +190,6 @@ class Cluster { peers_.resize(num_peers); for (size_t k = 0; k < num_peers; ++k) { - peers_[k].dram.assign(target_dram_bytes, 0); peers_[k].seed.assign(seed_bytes_per_peer, 0); PoolClientConfig tc; tc.master_config.node_id = "node-target-" + std::to_string(k); @@ -202,8 +199,7 @@ class Cluster { tc.io_engine.port = 0; tc.peer_service_port = NextPeerServicePort(); tc.dram_page_size = page_bytes; - tc.dram_buffers = {{peers_[k].dram.data(), peers_[k].dram.size()}}; - tc.tier_capacities = {{TierType::DRAM, {peers_[k].dram.size(), peers_[k].dram.size()}}}; + tc.dram.buffer_sizes = {target_dram_bytes}; peers_[k].client = std::make_unique(std::move(tc)); if (!peers_[k].client->Init()) { std::cerr << "target " << k << " init failed\n"; @@ -233,7 +229,6 @@ class Cluster { private: size_t page_bytes_; std::vector caller_buf_; - std::vector caller_local_; std::unique_ptr master_; std::thread server_thread_; std::unique_ptr caller_; diff --git a/tests/cpp/umbp/distributed/bench_pool_client_batch_put.cpp b/tests/cpp/umbp/distributed/bench_pool_client_batch_put.cpp index 22e79ae7b..53bd99767 100644 --- a/tests/cpp/umbp/distributed/bench_pool_client_batch_put.cpp +++ b/tests/cpp/umbp/distributed/bench_pool_client_batch_put.cpp @@ -132,7 +132,6 @@ bool ParseArgs(int argc, char** argv, BenchOpts* o) { // One peer = one PoolClient with a remote DRAM buffer. struct PeerNode { - std::vector buf; std::unique_ptr client; }; @@ -140,7 +139,7 @@ class Cluster { public: Cluster(size_t page_bytes, size_t target_dram_bytes, size_t num_peers, size_t caller_buf_bytes, size_t caller_local_bytes) - : page_bytes_(page_bytes), caller_buf_(caller_buf_bytes), caller_local_(caller_local_bytes) { + : page_bytes_(page_bytes), caller_buf_(caller_buf_bytes) { MasterServerConfig mcfg; mcfg.listen_address = "0.0.0.0:0"; master_ = std::make_unique(std::move(mcfg)); @@ -162,8 +161,7 @@ class Cluster { cc.io_engine.port = 0; cc.peer_service_port = NextPeerServicePort(); cc.dram_page_size = page_bytes; - cc.dram_buffers = {{caller_local_.data(), caller_local_.size()}}; - cc.tier_capacities = {{TierType::DRAM, {caller_local_.size(), caller_local_.size()}}}; + cc.dram.buffer_sizes = {caller_local_bytes}; caller_ = std::make_unique(std::move(cc)); if (!caller_->Init()) { std::cerr << "caller init failed\n"; @@ -172,7 +170,6 @@ class Cluster { peers_.resize(num_peers); for (size_t k = 0; k < num_peers; ++k) { - peers_[k].buf.assign(target_dram_bytes, 0); PoolClientConfig tc; tc.master_config.node_id = "node-target-" + std::to_string(k); tc.master_config.node_address = "127.0.0.1"; @@ -181,8 +178,7 @@ class Cluster { tc.io_engine.port = 0; tc.peer_service_port = NextPeerServicePort(); tc.dram_page_size = page_bytes; - tc.dram_buffers = {{peers_[k].buf.data(), peers_[k].buf.size()}}; - tc.tier_capacities = {{TierType::DRAM, {peers_[k].buf.size(), peers_[k].buf.size()}}}; + tc.dram.buffer_sizes = {target_dram_bytes}; peers_[k].client = std::make_unique(std::move(tc)); if (!peers_[k].client->Init()) { std::cerr << "target " << k << " init failed\n"; @@ -207,7 +203,6 @@ class Cluster { private: size_t page_bytes_; std::vector caller_buf_; - std::vector caller_local_; std::unique_ptr master_; std::thread server_thread_; std::unique_ptr caller_; diff --git a/tests/cpp/umbp/distributed/bench_pool_client_ranges.cpp b/tests/cpp/umbp/distributed/bench_pool_client_ranges.cpp new file mode 100644 index 000000000..940338f1f --- /dev/null +++ b/tests/cpp/umbp/distributed/bench_pool_client_ranges.cpp @@ -0,0 +1,1289 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Layer-wise KV restore bench for PoolClient's ranged I/O. +// +// WHAT IT MODELS +// +// A KV connector doing layer-wise loading against UMBP topology 3 (one node's +// standalone server whose inner backend is distributed, so one PoolClient is +// shared by every rank on that node). The shapes below are taken from +// sglang's umbp_tree_connector.py, not invented: +// +// * A page object holds ALL layers of one page; layer L occupies the byte +// range [L*layer_bytes, (L+1)*layer_bytes) inside it. +// * Loading walks the layer stack in groups of UMBP_LAYER_GROUP (default 8). +// One BatchGetRanges call carries a CHUNK OF PAGES, and each page carries +// one range per layer in the group -- not one key per call, and not one +// range per call. Chunk size is the RANGES_PER_CALL budget divided by the +// ranges each object actually carries. +// * Offload does NOT walk groups: one BatchPutRanges call writes every layer +// of a page at once, because a put's ranges must tile the object exactly. +// * Per rank there is exactly ONE load thread and ONE offload thread, so the +// concurrency reaching PoolClient is bounded by the rank count, not by an +// arbitrary thread count. +// * Forward runs concurrently with loading and waits per layer group: it can +// compute layer L while layer L+1 is still in flight. That overlap is the +// entire point of layer-wise, so this bench runs a real forward thread that +// blocks on each group and records how long it stalled. +// +// THE LOCALITY THAT MATTERS +// +// Ideally only the FIRST group touching a page pays a remote fetch, and groups +// 2..N are served by the local medium with no arena and no lock. How that +// locality gets built is exactly what this bench measures, and it is not free: +// a remote ranged get moves only the bytes asked for, so making the page local +// is a separate asynchronous pull that a later group races against. Whether +// it wins that race is visible as local_hit_pct and repeat_remote. +// +// A bench that forces every call remote would overstate how often the arena is +// reached, so this one lets the real behavior happen and reports the local-hit +// rate it observes (how that number is obtained, and what it cannot tell you, +// is on LocallyResidentCount below). +// +// WHERE "REMOTE" COMES FROM ON A SINGLE HOST +// +// UMBP decides local vs remote by comparing node ids, not machines: a key is +// remote when its route's node_id differs from this client's own +// (PartitionBatchGetTargets). This bench therefore builds TWO PoolClients in +// one process -- "node-local" and "node-peer" -- each with its own node id, +// peer-service port and io engine, and seeds every page onto the peer only. +// Reading them from the node is then genuinely routed as remote and really +// does go through BatchRouteGet, the MoriIoEngine RDMA path, the scratch arena +// and its mutex. +// +// So the CODE PATH is the real one, but both endpoints sit on one host and the +// bytes cross a local NIC. Treat the structural results (local-hit rate, +// whether the loader stays ahead of forward, how they scale with ranks) as +// meaningful, and the absolute latency/bandwidth as optimistic: a real +// two-machine deployment pays a longer wire. +// +// BUILDING. Like the other bench_* targets here it builds with the test tree +// but is not registered with ctest: it takes minutes and its output is numbers +// to read, not an assertion to pass. +// +// cmake -S . -B build_ranged_bench -GNinja -DBUILD_UMBP=ON -DBUILD_TESTS=ON \ +// -DCMAKE_BUILD_TYPE=Release +// cmake --build build_ranged_bench --target bench_umbp_pool_client_ranges -j +// +// On a NIC whose RC queues are smaller than mori-io's defaults the peer +// handshake fails with `ibv_create_qp ... Invalid argument`; cap the queues: +// MORI_IO_QP_MAX_SEND_WR=256 MORI_IO_QP_MAX_RECV_WR=256 MORI_IO_QP_MAX_CQE=1024 +// +// Usage: +// bench_pool_client_ranges [--ranks N] [--layers N] [--layer-group N] +// [--layer-bytes N] [--pools SPEC] [--pages N] +// [--requests N] [--compute-us-per-layer N] +// [--offload-ranks N] [--ranges-per-call N] +// [--scratch-mib N] [--page-bytes N] [--no-prefetch] +// [--interleave-groups] [--local-only] [--local-mode] +// [--verify] +// +// --local-only and --local-mode are different questions. --local-only keeps a +// master-backed client and puts every key on the reading node, so the local +// half of the routed path is what runs. --local-mode deploys the same client +// embedded -- no master, no peer, no arena at all -- the floor the distributed +// numbers are measured against. +// +// CSV (stdout): one row per rank count in --ranks +// ranks,layers,group,pages,object_kib,restore_ms_p50,restore_ms_p95, +// ttfl_us_p50,ttfl_us_p95,stall_ms_p50,get_gibps,put_gibps,local_hit_pct, +// remote_calls,remote_keys,repeat_remote,wire_whole_mib,wire_range_mib, +// local_write_mib +// +// The two wire columns are what the fetch WOULD move whole-object versus what +// a range-granular fetch actually moves, so their ratio is the read +// amplification this bench exists to measure. repeat_remote counts keys a +// single request had to fetch remotely more than once, which is what falls +// when locality is rebuilt between layer groups. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/config.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/pool_client.h" +#include "umbp/umbp_client.h" + +using mori::umbp::MasterServer; +using mori::umbp::MasterServerConfig; +using mori::umbp::PoolClient; +using mori::umbp::PoolClientConfig; + +namespace { + +// The bench drives one of two deployments through the same loop. +// +// Comparing them is the point: "distributed mode on one machine" and "local +// mode" run the same layer-wise restore against the same shapes, so any +// difference is the deployment and not the workload. That only holds if the +// measurement loop is literally the same code, which is why this indirection +// exists rather than a second copy of the loop. +// +// Pointers are void* here because that is PoolClient's signature; the local +// client takes uintptr_t and the adapter converts. +class RangedClient { + public: + virtual ~RangedClient() = default; + virtual bool Register(void* ptr, size_t size) = 0; + virtual std::vector Put(const std::vector& keys, + const std::vector& srcs, + const std::vector& sizes) = 0; + virtual std::vector GetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& offsets) = 0; + virtual std::vector PutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& offsets) = 0; + virtual std::vector Exists(const std::vector& keys) = 0; + // Whether every key can be served without leaving this node. For the + // distributed client that is a medium resolve; for the local one, where + // there is no other node, it is simply existence. + // How many of `keys` this deployment can serve from ITS OWN local + // storage right now, without reaching another node. A count rather than + // a bool because callers want both a per-call all-or-nothing check and a + // per-key credit -- deriving the former from a count keeps there from + // being two divergent notions of "resident" in this file. + virtual size_t ResidentCount(const std::vector& keys) = 0; +}; + +inline uint16_t NextPeerServicePort() { + static std::atomic next{ + static_cast(20000 + (static_cast(::getpid()) * 16u) % 40000u)}; + return next.fetch_add(1); +} + +// One KV pool: the layers it holds and how many bytes each occupies in a page. +// +// A model does not have one uniform pool. DeepSeek-V4 splits its layers across +// six, and their per-layer sizes span more than an order of magnitude +// (37,440 B for the c4 pool, 8,448 for its indexer, 1,728 for c128). That +// spread is the point: achieved bandwidth on this path is dominated by fragment +// size, measured at 5.4 GB/s at 1,728 B against 37.5 GB/s at 37,440 B, so a +// single --layer-bytes picks one point on a 7x curve and any change that helps +// small fragments specifically is mis-measured by it. +// +// `layers` are logical layer numbers, because pools interleave: DeepSeek-V4's +// c4 and c128 alternate down the stack, so which pools a layer group touches +// depends on their actual positions and not just their counts. +struct Pool { + std::string name; + std::vector layers; + size_t layer_bytes = 0; + + size_t object_bytes() const { return layers.size() * layer_bytes; } + // Position of a logical layer inside the object, or npos. The object packs + // the pool's own layers back to back, so this is an index into `layers`, not + // the logical number -- the twin of DevicePoolEntry.layer_mapping. + size_t IndexOf(size_t logical_layer) const { + for (size_t i = 0; i < layers.size(); ++i) { + if (layers[i] == logical_layer) return i; + } + return SIZE_MAX; + } +}; + +struct BenchOpts { + // Concurrency: one load thread + one offload thread per rank, all sharing one + // PoolClient (topology 3's node-level server). + std::vector rank_sweep; + size_t offload_ranks = SIZE_MAX; // default: same as ranks; 0 disables offload + + // Model shape. Defaults are DeepSeek-V3.2 / TP8 / page_size=64 / bf16, KV + // pool: 61 layers x 73,728 B per layer per page => a 4.29 MiB page object, + // read back 8 layers (576 KiB) at a time. The INDEXER pool of the same model + // is --layers 61 --layer-bytes 8448. + size_t layers = 61; + size_t layer_group = 8; // UMBP_LAYER_GROUP + size_t layer_bytes = 73728; // bytes of one layer of one page + size_t pages = 64; // pages (objects) restored per request + size_t requests = 3; // requests per rank + size_t compute_us_per_layer = 150; // simulated forward compute + + // Empty unless --pools was given, in which case it replaces the single + // implied pool above. Resolve() fills it either way so the rest of the + // bench only ever reads this. + std::vector pools; + + // Client-side budgets that shape the call pattern. + size_t ranges_per_call = 8192; // UMBP_RANGES_PER_CALL + size_t scratch_mib = 256; // ranged scratch arena, each direction + // Medium page size. 0 => one page per object, which is what distributed mode + // actually runs (master page_size == KV block size, see pool_client.cpp). + size_t page_bytes = 0; + // Background whole-object pull after a remote ranged read. Turning it off is + // how the wire cost of the duplicate traffic is measured. + bool locality_prefetch = true; + // Deal layers to groups round-robin instead of in contiguous blocks, so a + // group's ranges are strided through the object and cannot be merged into + // one read. Same byte count, same call count, same range count -- only + // adjacency changes. This is the adversarial case for any implementation + // that leans on contiguous layer groups. + bool interleave_groups = false; + // Seed onto the reading node instead of the peer, so every call is served by + // the local medium and Phase 2 never runs. This is the single-machine shape: + // a standalone process whose inner backend is distributed but whose keys all + // live here. It isolates the per-key and per-range CPU cost of the + // distributed path from anything to do with the wire. + bool local_only = false; + // Run the whole bench against an embedded deployment (no master, no peer, no + // arena) instead of a master-backed one. Same client class, same loop, same + // shapes, so the difference between the two runs is the deployment alone. + bool local_mode = false; + // Check every restored page byte-for-byte after each request. Off by + // default because the comparison lands inside the measured window; on, this + // is the only thing here that proves the numbers above came from correct + // reads and not from a fast wrong answer. + bool verify = false; + + // Turn --layers/--layer-bytes into a one-pool list when --pools was not + // given, so a command line that worked before this option existed still + // describes exactly the same workload. + void Resolve() { + if (pools.empty()) { + Pool single; + single.name = "kv"; + single.layer_bytes = layer_bytes; + single.layers.resize(layers); + for (size_t l = 0; l < layers; ++l) single.layers[l] = l; + pools.push_back(std::move(single)); + } + layers = 0; + for (const Pool& pool : pools) { + for (size_t l : pool.layers) layers = std::max(layers, l + 1); + } + if (layers == 0) layers = 1; + } + + // Bytes of one page across every pool -- the whole KV footprint of a token + // page, which is what a restore has to move. + size_t object_bytes() const { + size_t total = 0; + for (const Pool& pool : pools) total += pool.object_bytes(); + return total; + } + size_t group_count() const { return (layers + layer_group - 1) / layer_group; } + // Bytes one layer group moves, summed over the pools that group touches. + // With heterogeneous pools this varies by group, so it is the mean. + size_t group_bytes() const { + const size_t groups = group_count(); + return groups > 0 ? object_bytes() / groups : object_bytes(); + } + // Largest single object, which is what the scratch arena must be able to hold. + size_t max_object_bytes() const { + size_t most = 0; + for (const Pool& pool : pools) most = std::max(most, pool.object_bytes()); + return most; + } + size_t medium_page_bytes() const { return page_bytes != 0 ? page_bytes : max_object_bytes(); } +}; + +void Usage() { + std::cerr << "Usage: bench_pool_client_ranges [--ranks 1,2,4,8]\n" + << " [--layers N] [--layer-group N] [--layer-bytes N]\n" + << " [--pools SPEC] [--pages N] [--requests N]\n" + << " [--compute-us-per-layer N]\n" + << " [--offload-ranks N] [--ranges-per-call N] [--scratch-mib N]\n" + << " [--page-bytes N] [--no-prefetch] [--interleave-groups]\n" + << " [--local-only] [--local-mode] [--verify]\n" + << "\n" + << " --pools replaces --layers/--layer-bytes with a real model's several\n" + << " pools, whose per-layer sizes differ by more than 10x and decide the\n" + << " achieved bandwidth. Semicolon separated, each entry:\n" + << " name:layer_count:layer_bytes[:first/stride | :l1,l2,l3,...]\n" + << " The last field places the pool's layers in the stack (default 0,1,2..),\n" + << " which is what decides the layer groups a pool appears in. Use the\n" + << " explicit list when the layers are not an arithmetic progression.\n" + << " DeepSeek-V4-Pro, from its config's compress_ratios:\n" + << " --pools " + "'c4:30:37440:2/2;c4_indexer:30:8448:2/" + "2;c128:31:1728:0,1,3,5,7,9,11,13,15,17,19,21,23,25,27,29,31,33,35,37,39,41,43,45," + "47,49,51,53,55,57,59'\n"; +} + +std::vector ParseList(const std::string& s) { + std::vector out; + std::stringstream ss(s); + std::string tok; + while (std::getline(ss, tok, ',')) { + if (!tok.empty()) out.push_back(std::strtoull(tok.c_str(), nullptr, 10)); + } + return out; +} + +// "name:count:bytes[:first/stride]" entries separated by ';'. +std::vector ParsePools(const std::string& spec) { + std::vector pools; + std::stringstream entries(spec); + std::string entry; + while (std::getline(entries, entry, ';')) { + if (entry.empty()) continue; + std::vector field; + std::stringstream fields(entry); + std::string tok; + while (std::getline(fields, tok, ':')) field.push_back(tok); + if (field.size() < 3 || field.size() > 4) { + std::cerr << "--pools: expected name:count:bytes[:first/stride], got '" << entry << "'\n"; + std::exit(2); + } + Pool pool; + pool.name = field[0]; + const size_t count = std::strtoull(field[1].c_str(), nullptr, 10); + pool.layer_bytes = std::strtoull(field[2].c_str(), nullptr, 10); + if (count == 0 || pool.layer_bytes == 0) { + std::cerr << "--pools: count and bytes must be non-zero in '" << entry << "'\n"; + std::exit(2); + } + if (field.size() == 3) { + pool.layers.reserve(count); + for (size_t i = 0; i < count; ++i) pool.layers.push_back(i); + } else if (field[3].find('/') != std::string::npos) { + const size_t slash = field[3].find('/'); + const size_t first = std::strtoull(field[3].substr(0, slash).c_str(), nullptr, 10); + const size_t stride = std::strtoull(field[3].substr(slash + 1).c_str(), nullptr, 10); + if (stride == 0) { + std::cerr << "--pools: stride must be non-zero in '" << entry << "'\n"; + std::exit(2); + } + pool.layers.reserve(count); + for (size_t i = 0; i < count; ++i) pool.layers.push_back(first + i * stride); + } else { + // Explicit list. Needed because real pools are not always an arithmetic + // progression: DeepSeek-V4's c128 holds layers 0,1,3,5,...,59 -- the + // first two adjacent, the rest every other one. + pool.layers = ParseList(field[3]); + if (pool.layers.size() != count) { + std::cerr << "--pools: '" << pool.name << "' declares " << count << " layers but lists " + << pool.layers.size() << "\n"; + std::exit(2); + } + } + pools.push_back(std::move(pool)); + } + if (pools.empty()) { + std::cerr << "--pools: no pools parsed from '" << spec << "'\n"; + std::exit(2); + } + return pools; +} + +bool ParseArgs(int argc, char** argv, BenchOpts* o) { + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto next = [&](const char* what) -> const char* { + if (i + 1 >= argc) { + std::cerr << "Missing value for " << what << "\n"; + std::exit(2); + } + return argv[++i]; + }; + if (a == "--ranks") { + o->rank_sweep = ParseList(next("--ranks")); + } else if (a == "--layers") { + o->layers = std::strtoull(next("--layers"), nullptr, 10); + } else if (a == "--layer-group") { + o->layer_group = std::strtoull(next("--layer-group"), nullptr, 10); + } else if (a == "--layer-bytes") { + o->layer_bytes = std::strtoull(next("--layer-bytes"), nullptr, 10); + } else if (a == "--pools") { + o->pools = ParsePools(next("--pools")); + } else if (a == "--pages") { + o->pages = std::strtoull(next("--pages"), nullptr, 10); + } else if (a == "--requests") { + o->requests = std::strtoull(next("--requests"), nullptr, 10); + } else if (a == "--compute-us-per-layer") { + o->compute_us_per_layer = std::strtoull(next("--compute-us-per-layer"), nullptr, 10); + } else if (a == "--offload-ranks") { + o->offload_ranks = std::strtoull(next("--offload-ranks"), nullptr, 10); + } else if (a == "--ranges-per-call") { + o->ranges_per_call = std::strtoull(next("--ranges-per-call"), nullptr, 10); + } else if (a == "--scratch-mib") { + o->scratch_mib = std::strtoull(next("--scratch-mib"), nullptr, 10); + } else if (a == "--page-bytes") { + o->page_bytes = std::strtoull(next("--page-bytes"), nullptr, 10); + } else if (a == "--no-prefetch") { + o->locality_prefetch = false; + } else if (a == "--interleave-groups") { + o->interleave_groups = true; + } else if (a == "--local-only") { + o->local_only = true; + } else if (a == "--local-mode") { + o->local_mode = true; + } else if (a == "--verify") { + o->verify = true; + } else if (a == "-h" || a == "--help") { + Usage(); + std::exit(0); + } else { + std::cerr << "Unknown arg: " << a << "\n"; + Usage(); + return false; + } + } + if (o->rank_sweep.empty()) o->rank_sweep = {1, 2, 4, 8}; + if (o->layer_group == 0) o->layer_group = 1; + if (o->layers == 0) o->layers = 1; + o->Resolve(); + return true; +} + +// Layers covered by group `g`, mirroring _layer_groups(). These are logical +// layer numbers spanning the whole stack; which pools they reach is a separate +// question answered by PoolGroupLayers. +// +// With --interleave-groups the same layers are dealt round-robin instead, so +// every group still covers its share of the object but its ranges are strided +// rather than contiguous. The partition stays exact either way. +std::vector GroupLayers(const BenchOpts& o, size_t g) { + std::vector out; + if (o.interleave_groups) { + for (size_t l = g; l < o.layers; l += o.group_count()) out.push_back(l); + return out; + } + const size_t start = g * o.layer_group; + for (size_t l = start; l < std::min(start + o.layer_group, o.layers); ++l) out.push_back(l); + return out; +} + +// Where this pool's own layers sit inside a group: object-relative indices, not +// logical numbers. Empty when the group misses the pool entirely, which is how +// interleaved pools end up issuing calls on different groups -- the twin of +// _plans_covering() skipping a plan. +std::vector PoolGroupLayers(const Pool& pool, const std::vector& group) { + std::vector out; + for (size_t logical : group) { + const size_t index = pool.IndexOf(logical); + if (index != SIZE_MAX) out.push_back(index); + } + return out; +} + +// Objects per RPC, budgeted by the ranges they actually carry -- the C++ twin +// of _entries_per_call(). +size_t EntriesPerCall(const BenchOpts& o, size_t ranges_per_object) { + return std::max(1, o.ranges_per_call / std::max(1, ranges_per_object)); +} + +double Percentile(std::vector v, double q) { + if (v.empty()) return 0.0; + std::sort(v.begin(), v.end()); + const size_t idx = std::min(v.size() - 1, static_cast(q * (v.size() - 1) + 0.5)); + return v[idx]; +} + +// Busy-wait: sleep_for's granularity is coarse next to a layer group's compute. +void SpinUs(uint64_t us) { + if (us == 0) return; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::microseconds(us); + while (std::chrono::steady_clock::now() < deadline) { + } +} + +// Two PoolClients in one process, distinct node ids: `node` is the one every +// rank drives (the node-level client of topology 3), `peer` exists only to own +// the seeded pages so that the node's first read of a page routes as remote. +// See "WHERE REMOTE COMES FROM ON A SINGLE HOST" at the top of this file. +class Cluster { + public: + Cluster(size_t node_capacity, size_t peer_capacity, size_t scratch_bytes, size_t page_bytes, + bool locality_prefetch) + : page_bytes_(page_bytes), + locality_prefetch_(locality_prefetch), + node_get_scratch_(scratch_bytes), + node_put_scratch_(scratch_bytes), + peer_get_scratch_(scratch_bytes), + peer_put_scratch_(scratch_bytes) { + MasterServerConfig mcfg; + mcfg.listen_address = "0.0.0.0:0"; + mcfg.registry_config.heartbeat_ttl = std::chrono::seconds{1}; + master_ = std::make_unique(std::move(mcfg)); + master_thread_ = std::thread([this] { master_->Run(); }); + for (int i = 0; i < 300 && master_->GetBoundPort() == 0; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + if (master_->GetBoundPort() == 0) { + std::cerr << "master failed to start\n"; + std::exit(2); + } + const std::string master_addr = "localhost:" + std::to_string(master_->GetBoundPort()); + + node_ = + MakeClient("node-local", master_addr, node_capacity, node_get_scratch_, node_put_scratch_); + peer_ = + MakeClient("node-peer", master_addr, peer_capacity, peer_get_scratch_, peer_put_scratch_); + } + + ~Cluster() { + if (node_) node_->Shutdown(); + if (peer_) peer_->Shutdown(); + if (master_) master_->Shutdown(); + if (master_thread_.joinable()) master_thread_.join(); + } + + PoolClient* node() { return node_.get(); } + PoolClient* peer() { return peer_.get(); } + + private: + std::unique_ptr MakeClient(const std::string& node_id, const std::string& master_addr, + size_t dram_capacity, std::vector& get_scratch, + std::vector& put_scratch) { + PoolClientConfig cfg; + cfg.master_config.node_id = node_id; + cfg.master_config.node_address = "127.0.0.1"; + cfg.master_config.master_address = master_addr; + cfg.io_engine.host = "0.0.0.0"; + cfg.io_engine.port = 0; + cfg.peer_service_port = NextPeerServicePort(); + cfg.dram_page_size = page_bytes_; + cfg.staging_buffer_size = 1024; // ranged remote must go through the arena + cfg.dram.buffer_sizes = {dram_capacity}; + cfg.ranged_get_scratch_buffer = get_scratch.data(); + cfg.ranged_get_scratch_size = get_scratch.size(); + cfg.ranged_put_scratch_buffer = put_scratch.data(); + cfg.ranged_put_scratch_size = put_scratch.size(); + cfg.cache_remote_fetches = false; + cfg.ranged_locality_prefetch = locality_prefetch_; + auto client = std::make_unique(std::move(cfg)); + if (!client->Init()) { + std::cerr << node_id << " init failed\n"; + std::exit(2); + } + if (!client->RegisterMemory(get_scratch.data(), get_scratch.size()) || + !client->RegisterMemory(put_scratch.data(), put_scratch.size())) { + std::cerr << node_id << " scratch registration failed\n"; + std::exit(2); + } + return client; + } + + size_t page_bytes_; + bool locality_prefetch_; + std::vector node_get_scratch_; + std::vector node_put_scratch_; + std::vector peer_get_scratch_; + std::vector peer_put_scratch_; + std::unique_ptr master_; + std::thread master_thread_; + std::unique_ptr node_; + std::unique_ptr peer_; +}; + +// How local_hit_pct is obtained, and what it is not. +// +// The bench cannot see which branch BatchGetRanges took -- PoolClient reports +// nothing about it. What it can do is apply the SAME criterion the call itself +// will apply: phase 1 resolves each key against this node's own medium and +// serves whatever it finds without ever reaching the arena, so a key that +// resolves here is a key that will not go remote. This runs that identical +// resolve just before the call. +// +// Two limits, so the number is not over-read: +// * It is a prediction from the same criterion, taken a moment earlier -- not +// an observation from inside mori. +// * It races: under concurrency the object can be evicted between this +// resolve and the call, which would then go remote after being counted +// local. Expect the reported rate to be slightly optimistic at high rank +// counts, which is also where eviction actually starts happening. +// +// Returns a per-key COUNT rather than a bool: local_hit_pct wants an +// all-or-nothing call classification (a chunk counts as local only when every +// key in it resolves), while local_write_mib wants partial credit (an object +// that landed counts even if others in its pool have not yet) -- two different +// questions the callers derive from the same count rather than two functions +// that could silently drift apart. +// +// Resolved for the whole chunk in one call: this sits inside the timed region, +// so it has to stay cheap. +size_t LocallyResidentCount(PoolClient* client, const std::vector& keys) { + auto* backend = client->Backends().Get(client->Medium()); + if (backend == nullptr) return 0; + auto resolved = backend->BatchResolve(keys, /*include_descs=*/false); + return static_cast(std::count_if(resolved.begin(), resolved.end(), + [](const auto& entry) { return entry.found; })); +} + +// Distributed mode: a PoolClient with a master behind it. +class DistributedRangedClient final : public RangedClient { + public: + explicit DistributedRangedClient(PoolClient* client) : client_(client) {} + + bool Register(void* ptr, size_t size) override { return client_->RegisterMemory(ptr, size); } + std::vector Put(const std::vector& keys, const std::vector& srcs, + const std::vector& sizes) override { + return client_->BatchPut(keys, srcs, sizes); + } + std::vector GetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& offsets) override { + return client_->BatchGetRanges(keys, dsts, sizes, offsets); + } + std::vector PutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& offsets) override { + return client_->BatchPutRanges(keys, object_sizes, srcs, sizes, offsets); + } + std::vector Exists(const std::vector& keys) override { + return client_->BatchExists(keys); + } + size_t ResidentCount(const std::vector& keys) override { + return LocallyResidentCount(client_, keys); + } + + private: + PoolClient* client_; +}; + +// Embedded mode: the same client class as the distributed arm, deployed with +// no master, no peer and no arena. The floor the distributed numbers are +// measured against -- and, since the convergence, a measurement of the +// deployment rather than of a second implementation. +class LocalRangedClient final : public RangedClient { + public: + LocalRangedClient(size_t dram_capacity, size_t page_bytes) { + mori::umbp::UMBPConfig cfg; + cfg.dram.capacity_bytes = dram_capacity; + // UMBPConfig enables an SSD tier by default; an embedded deployment serves + // one medium and would only warn about it. The distributed side of this + // comparison is DRAM, so say so. + cfg.ssd.enabled = false; + cfg = mori::umbp::WithEmbeddedDefaults(cfg); + // The pool is paged, so the page has to match the objects being stored or + // the two arms are not sizing the same way. + cfg.distributed->dram_page_size = page_bytes; + client_ = mori::umbp::CreateUMBPClient(cfg); + if (!client_->SupportsRangedIO()) { + std::cerr << "embedded client reports no ranged I/O support\n"; + std::exit(2); + } + } + + // Nothing to pin: local mode copies out of its own tier with no engine in + // the way, so a caller pointer is already usable. + bool Register(void*, size_t) override { return true; } + std::vector Put(const std::vector& keys, const std::vector& srcs, + const std::vector& sizes) override { + std::vector raw(srcs.size()); + for (size_t i = 0; i < srcs.size(); ++i) raw[i] = reinterpret_cast(srcs[i]); + return client_->BatchPut(keys, raw, sizes); + } + std::vector GetRanges(const std::vector& keys, + const std::vector>& dsts, + const std::vector>& sizes, + const std::vector>& offsets) override { + return client_->BatchGetRanges(keys, ToRaw(dsts), sizes, offsets); + } + std::vector PutRanges(const std::vector& keys, + const std::vector& object_sizes, + const std::vector>& srcs, + const std::vector>& sizes, + const std::vector>& offsets) override { + return client_->BatchPutRanges(keys, object_sizes, ToRaw(srcs), sizes, offsets); + } + std::vector Exists(const std::vector& keys) override { + return client_->BatchExists(keys); + } + // There is no other node, so anything that exists is here. + size_t ResidentCount(const std::vector& keys) override { + auto present = client_->BatchExists(keys); + return static_cast(std::count(present.begin(), present.end(), true)); + } + + private: + template + static std::vector> ToRaw(const std::vector>& in) { + std::vector> out(in.size()); + for (size_t i = 0; i < in.size(); ++i) { + out[i].reserve(in[i].size()); + for (P p : in[i]) out[i].push_back(reinterpret_cast(p)); + } + return out; + } + + std::unique_ptr client_; +}; + +void WaitAllVisible(RangedClient* client, const std::vector& keys, + std::chrono::milliseconds timeout = std::chrono::seconds{30}) { + if (keys.empty()) return; + const auto deadline = std::chrono::steady_clock::now() + timeout; + for (;;) { + auto present = client->Exists(keys); + if (std::all_of(present.begin(), present.end(), [](bool b) { return b; })) return; + if (std::chrono::steady_clock::now() >= deadline) { + std::cerr << "WaitAllVisible: keys not visible before timeout\n"; + std::exit(2); + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } +} + +// Publish whole page objects on the peer so the node must fetch them remotely. +// Byte a page's object should hold at `offset`. A pure function of the key's +// index so the loader can recompute the expectation without keeping a copy. +inline char SeedByte(size_t key_index, size_t offset) { + return static_cast((offset * 31 + key_index * 131 + 7) & 0xff); +} + +// `target` is the peer by default, which is what makes the reads remote; under +// --local-only it is the reading node itself, so the same keys resolve locally. +void SeedInto(RangedClient* target, const std::vector& keys, size_t object_bytes, + std::vector* storage) { + storage->assign(keys.size() * object_bytes, 0); + std::vector srcs(keys.size()); + std::vector sizes(keys.size(), object_bytes); + for (size_t i = 0; i < keys.size(); ++i) { + char* slot = storage->data() + i * object_bytes; + for (size_t b = 0; b < object_bytes; ++b) slot[b] = SeedByte(i, b); + srcs[i] = slot; + } + target->Register(storage->data(), storage->size()); + auto r = target->Put(keys, srcs, sizes); + for (size_t i = 0; i < keys.size(); ++i) { + if (!r[i]) { + std::cerr << "seed put failed for " << keys[i] << "\n"; + std::exit(2); + } + } +} + +// Per-request handoff between the loader thread and the forward thread, the +// role LayerWiseLoadCounter plays in the connector. +struct LayerGate { + std::mutex mu; + std::condition_variable cv; + size_t groups_done = 0; + bool failed = false; + + void Complete(size_t groups) { + { + std::lock_guard lk(mu); + groups_done = groups; + } + cv.notify_all(); + } + void Fail() { + { + std::lock_guard lk(mu); + failed = true; + groups_done = SIZE_MAX; + } + cv.notify_all(); + } + // Returns microseconds spent stalled waiting for this group. + double WaitFor(size_t group_index) { + const auto t0 = std::chrono::steady_clock::now(); + std::unique_lock lk(mu); + cv.wait(lk, [&] { return groups_done > group_index; }); + return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + } +}; + +struct RankStats { + std::vector restore_ms; // per request: first group issued -> last layer computed + std::vector ttfl_us; // per request: stall before group 0 is readable + std::vector stall_ms; // per request: total forward stall across groups + size_t get_calls = 0; + size_t get_calls_local = 0; // calls whose pages were all already local + size_t get_bytes = 0; // bytes the caller asked for + size_t put_bytes = 0; + // Wire accounting. These are PREDICTIONS from the same criterion + // LocallyResidentCount applies -- see its comment for what that does and does + // not mean. They are what the bench can know from outside mori. + // + // remote_keys is the useful one: it counts (key, call) pairs that had to reach + // the peer, so the reader can derive both wire figures and compare them + // directly across a before/after run: + // whole-object fetch -> remote_keys * object_bytes + // range-granular -> remote_keys * group_bytes + size_t remote_calls = 0; + size_t remote_keys = 0; + size_t repeat_remote = 0; // (key, call) pairs that went remote for a key + // this request had already fetched remotely + size_t local_write_bytes = 0; // objects that ended up installed on this node + size_t verify_mismatches = 0; // pages that came back wrong (--verify only) + // Per pool, indexed as BenchOpts::pools. Kept apart from the totals because + // the totals cannot show the thing heterogeneous pools exist to expose: a + // pool with a tenth the range size moves its bytes at a fraction of the rate, + // and only a per-pool split says which pool a change actually helped. + std::vector pool_bytes; + std::vector pool_ranges; + std::vector pool_calls; + std::vector pool_remote_keys; + + void Init(size_t pools) { + pool_bytes.assign(pools, 0); + pool_ranges.assign(pools, 0); + pool_calls.assign(pools, 0); + pool_remote_keys.assign(pools, 0); + } + // Measured window for this rank, so aggregate throughput excludes warm-up. + std::chrono::steady_clock::time_point window_start{}; + std::chrono::steady_clock::time_point window_end{}; +}; + +// One rank's layer-wise restore: a loader thread walking layer groups while a +// forward thread consumes them. +// `warmup_requests` leading requests run but are not recorded: the first +// restore on a fresh client also pays the peer handshake, which a long-lived +// connector has already paid. +void RunLoadRank(RangedClient* client, const BenchOpts& o, + const std::vector>>& request_keys, + size_t warmup_requests, std::atomic* warmed_ranks, RankStats* stats) { + // Destination for one request's pages, one buffer per pool because a pool's + // pages are only contiguous among themselves. Registered once for zero copy. + std::vector> kv(o.pools.size()); + for (size_t p = 0; p < o.pools.size(); ++p) { + kv[p].assign(o.pages * o.pools[p].object_bytes(), 0); + client->Register(kv[p].data(), kv[p].size()); + } + + for (size_t req = 0; req < request_keys.size(); ++req) { + const bool measured = req >= warmup_requests; + if (req == warmup_requests) { + stats->window_start = std::chrono::steady_clock::now(); + warmed_ranks->fetch_add(1, std::memory_order_relaxed); + } + const std::vector>& pool_keys = request_keys[req]; + LayerGate gate; + const size_t groups = o.group_count(); + // Keys this request has already pulled from the peer once. A second remote + // pull of the same key inside one load means the object did not become + // locally readable between two layer groups. + std::unordered_set pulled_remote; + + const auto req_start = std::chrono::steady_clock::now(); + + std::thread loader([&] { + for (size_t g = 0; g < groups; ++g) { + const std::vector group = GroupLayers(o, g); + bool ok = true; + // Each pool this group reaches issues its own calls, at its own range + // size -- the loop _run_layer_wise_batch runs over _plans_covering. + for (size_t pi = 0; pi < o.pools.size() && ok; ++pi) { + const Pool& pool = o.pools[pi]; + const std::vector layers = PoolGroupLayers(pool, group); + if (layers.empty()) continue; + const std::vector& keys = pool_keys[pi]; + // Every page carries one range per layer of this pool in this group. + const size_t step = EntriesPerCall(o, layers.size()); + for (size_t start = 0; start < keys.size() && ok; start += step) { + const size_t end = std::min(start + step, keys.size()); + const size_t count = end - start; + + std::vector chunk_keys(keys.begin() + start, keys.begin() + end); + std::vector> dsts(count); + std::vector> sizes(count); + std::vector> offsets(count); + for (size_t i = 0; i < count; ++i) { + char* page = kv[pi].data() + (start + i) * pool.object_bytes(); + for (size_t l : layers) { + dsts[i].push_back(page + l * pool.layer_bytes); + sizes[i].push_back(pool.layer_bytes); + offsets[i].push_back(l * pool.layer_bytes); + } + } + + // Whether this call can be served without the arena at all. + const bool all_local = client->ResidentCount(chunk_keys) == chunk_keys.size(); + + auto res = client->GetRanges(chunk_keys, dsts, sizes, offsets); + ok = res.size() == count && + std::all_of(res.begin(), res.end(), [](bool b) { return b; }); + + if (measured) { + stats->get_calls += 1; + if (all_local) stats->get_calls_local += 1; + if (ok) { + const size_t moved = count * layers.size() * pool.layer_bytes; + stats->get_bytes += moved; + stats->pool_bytes[pi] += moved; + stats->pool_ranges[pi] += count * layers.size(); + stats->pool_calls[pi] += 1; + } + } + if (!all_local) { + for (const auto& k : chunk_keys) { + const bool again = !pulled_remote.insert(k).second; + if (measured && again) stats->repeat_remote += 1; + } + if (measured) { + stats->remote_calls += 1; + stats->remote_keys += count; + stats->pool_remote_keys[pi] += count; + } + } + } + } + if (!ok) { + gate.Fail(); + return; + } + // Only now is every layer in the group readable, so they are released + // together -- the connector releases a whole group at once too. + gate.Complete(g + 1); + } + }); + + // Forward: wait for each group, then "compute" its layers while the loader + // is already fetching the next one. + double stall_us = 0.0; + double ttfl = 0.0; + for (size_t g = 0; g < groups; ++g) { + const double waited = gate.WaitFor(g); + if (g == 0) ttfl = waited; + stall_us += waited; + { + std::lock_guard lk(gate.mu); + if (gate.failed) break; + } + SpinUs(o.compute_us_per_layer * GroupLayers(o, g).size()); + } + loader.join(); + + const double restore_ms = + std::chrono::duration(std::chrono::steady_clock::now() - req_start) + .count(); + if (measured) { + stats->restore_ms.push_back(restore_ms); + stats->ttfl_us.push_back(ttfl); + stats->stall_ms.push_back(stall_us / 1000.0); + stats->window_end = std::chrono::steady_clock::now(); + // Byte-check the whole restored working set. Deliberately after the + // timing is recorded, so per-request latency is unaffected; it does sit + // inside the throughput window, so --verify runs are for correctness and + // the timing runs are separate. + if (o.verify) { + for (size_t pi = 0; pi < o.pools.size(); ++pi) { + const size_t object_bytes = o.pools[pi].object_bytes(); + for (size_t p = 0; p < o.pages; ++p) { + const char* page = kv[pi].data() + p * object_bytes; + // Seeded per pool, so the index must match SeedInto's ordering. + const size_t seed_index = req * o.pages + p; + for (size_t b = 0; b < object_bytes; ++b) { + if (page[b] != SeedByte(seed_index, b)) { + stats->verify_mismatches += 1; + break; + } + } + } + } + } + + // How much of this request's working set the remote path installed here. + // Counted after the load so an asynchronous install has had the whole + // request to land; it is still a lower bound if one lands later. + for (size_t pi = 0; pi < o.pools.size(); ++pi) { + stats->local_write_bytes += + client->ResidentCount(pool_keys[pi]) * o.pools[pi].object_bytes(); + } + } + } +} + +// One rank's offload stream: whole-object writes that overlap the loads above. +// This is the traffic that used to contend with loading on a single arena. +void RunOffloadRank(RangedClient* client, const BenchOpts& o, size_t rank, + const std::atomic& stop, const std::atomic& warmed_ranks, + size_t load_ranks, RankStats* stats) { + // One buffer per pool, mirroring the load side. + std::vector> kv(o.pools.size()); + for (size_t pi = 0; pi < o.pools.size(); ++pi) { + kv[pi].assign(o.pages * o.pools[pi].object_bytes(), 0); + for (size_t i = 0; i < kv[pi].size(); ++i) + kv[pi][i] = static_cast((i + rank + pi) & 0xff); + client->Register(kv[pi].data(), kv[pi].size()); + } + + for (size_t round = 0; !stop.load(std::memory_order_relaxed); ++round) { + for (size_t pi = 0; pi < o.pools.size() && !stop.load(std::memory_order_relaxed); ++pi) { + const Pool& pool = o.pools[pi]; + const size_t object_bytes = pool.object_bytes(); + // A put names every layer of the page at once: its ranges must tile the + // object exactly, so offload never walks groups. + const size_t step = EntriesPerCall(o, pool.layers.size()); + + for (size_t start = 0; start < o.pages && !stop.load(std::memory_order_relaxed); + start += step) { + const size_t end = std::min(start + step, o.pages); + const size_t count = end - start; + + std::vector chunk_keys(count); + std::vector object_sizes(count, object_bytes); + std::vector> srcs(count); + std::vector> sizes(count); + std::vector> offsets(count); + for (size_t i = 0; i < count; ++i) { + // Fresh key every time: a repeat would be short-circuited by + // RoutePut's dedup and never reach the data path. + chunk_keys[i] = "off-" + pool.name + "-r" + std::to_string(rank) + "-" + + std::to_string(round) + "-" + std::to_string(start + i); + char* page = kv[pi].data() + (start + i) * object_bytes; + for (size_t l = 0; l < pool.layers.size(); ++l) { + srcs[i].push_back(page + l * pool.layer_bytes); + sizes[i].push_back(pool.layer_bytes); + offsets[i].push_back(l * pool.layer_bytes); + } + } + auto res = client->PutRanges(chunk_keys, object_sizes, srcs, sizes, offsets); + // Offload keeps running through warm-up so the arena is loaded from the + // start, but only bytes inside the measured window are counted. + const bool in_window = warmed_ranks.load(std::memory_order_relaxed) >= load_ranks; + for (size_t i = 0; i < res.size(); ++i) { + if (res[i] && in_window) stats->put_bytes += object_bytes; + } + } + } + } +} + +void RunOne(const BenchOpts& o, size_t ranks) { + const size_t offload_ranks = (o.offload_ranks == SIZE_MAX) ? ranks : o.offload_ranks; + const size_t object_bytes = o.object_bytes(); + + // Whichever client is seeded holds every page for every rank and request + // (+1 warm-up); the other only needs a restore working set per rank, which is + // what makes groups 2..N local hits, plus slack for offloads routed there. + const size_t seeded_capacity = + std::max(size_t{512} << 20, + ranks * (o.requests + 1) * o.pages * object_bytes * 2 + (size_t{512} << 20)); + const size_t reader_capacity = + std::max(size_t{256} << 20, ranks * o.pages * object_bytes * 3); + const size_t node_capacity = o.local_only ? seeded_capacity : reader_capacity; + const size_t peer_capacity = o.local_only ? reader_capacity : seeded_capacity; + // The arena must hold at least one whole object: the remote path stages + // through it, and an object that does not fit is unservable. With several + // pools the largest one sets the floor. + const size_t scratch_bytes = std::max(o.max_object_bytes(), o.scratch_mib << 20); + + // Local mode needs no master, no peer and no arena, so the cluster is built + // only for the distributed side. Both then go through RangedClient. + std::unique_ptr cluster; + std::unique_ptr node_client; + std::unique_ptr peer_client; + if (o.local_mode) { + node_client = std::make_unique(seeded_capacity, o.medium_page_bytes()); + } else { + cluster = std::make_unique(node_capacity, peer_capacity, scratch_bytes, + o.medium_page_bytes(), o.locality_prefetch); + node_client = std::make_unique(cluster->node()); + peer_client = std::make_unique(cluster->peer()); + } + + // Seed: each (rank, request) gets its own pages, so the first group of every + // request is a real remote miss -- or, under --local-only, a local hit that + // no earlier group warmed. One extra leading request per rank is a warm-up + // that pays the peer handshake and is not recorded. + // Local mode has nowhere else to put anything; distributed mode seeds the + // peer unless --local-only asks for everything on the reading node. + RangedClient* const seed_target = + (o.local_mode || o.local_only) ? node_client.get() : peer_client.get(); + constexpr size_t kWarmupRequests = 1; + const size_t total_requests = o.requests + kWarmupRequests; + // [rank][request][pool] -> that pool's page keys. A pool's object holds only + // its own layers, so pools cannot share keys. + std::vector>>> keys(ranks); + std::vector>> seed_storage(ranks); + std::vector all_keys; + for (size_t r = 0; r < ranks; ++r) { + keys[r].resize(total_requests); + seed_storage[r].resize(o.pools.size()); + // Seeded per pool so SeedByte's index is (request, page) within the pool, + // which is what the verify pass recomputes. + std::vector> flat(o.pools.size()); + for (size_t q = 0; q < total_requests; ++q) { + keys[r][q].resize(o.pools.size()); + for (size_t pi = 0; pi < o.pools.size(); ++pi) { + for (size_t p = 0; p < o.pages; ++p) { + keys[r][q][pi].push_back(o.pools[pi].name + "-r" + std::to_string(r) + "-q" + + std::to_string(q) + "-p" + std::to_string(p)); + flat[pi].push_back(keys[r][q][pi].back()); + } + } + } + for (size_t pi = 0; pi < o.pools.size(); ++pi) { + SeedInto(seed_target, flat[pi], o.pools[pi].object_bytes(), &seed_storage[r][pi]); + all_keys.insert(all_keys.end(), flat[pi].begin(), flat[pi].end()); + } + } + if (cluster) { + // Distributed mode only: the seeded keys become visible to the reader once + // both nodes' outboxes have reached master. + cluster->peer()->Master().FlushHeartbeat(); + cluster->node()->Master().FlushHeartbeat(); + } + WaitAllVisible(node_client.get(), all_keys); + + std::vector load_stats(ranks); + std::vector offload_stats(std::max(offload_ranks, 1)); + for (auto& s : load_stats) s.Init(o.pools.size()); + for (auto& s : offload_stats) s.Init(o.pools.size()); + std::atomic stop{false}; + + // Offload streams run for the whole measurement window, the way a steady + // state has finished requests being written back while new ones load. + std::atomic warmed_ranks{0}; + std::vector offloaders; + for (size_t r = 0; r < offload_ranks; ++r) { + offloaders.emplace_back([&, r] { + RunOffloadRank(node_client.get(), o, r, stop, warmed_ranks, ranks, &offload_stats[r]); + }); + } + + std::vector loaders; + for (size_t r = 0; r < ranks; ++r) { + loaders.emplace_back([&, r] { + RunLoadRank(node_client.get(), o, keys[r], kWarmupRequests, &warmed_ranks, &load_stats[r]); + }); + } + for (auto& t : loaders) t.join(); + + stop.store(true, std::memory_order_relaxed); + for (auto& t : offloaders) t.join(); + + // Measured window: from the first rank leaving warm-up to the last rank + // finishing, so warm-up traffic is outside the throughput denominator. + auto window_start = std::chrono::steady_clock::time_point::max(); + auto window_end = std::chrono::steady_clock::time_point::min(); + for (const auto& s : load_stats) { + if (s.window_start != std::chrono::steady_clock::time_point{}) { + window_start = std::min(window_start, s.window_start); + window_end = std::max(window_end, s.window_end); + } + } + const double wall_s = window_end > window_start + ? std::chrono::duration(window_end - window_start).count() + : 0.0; + std::vector restore, ttfl, stall; + size_t get_bytes = 0, put_bytes = 0, calls = 0, calls_local = 0; + size_t remote_calls = 0, remote_keys = 0, repeat_remote = 0, local_write = 0, mismatches = 0; + std::vector pool_bytes(o.pools.size(), 0), pool_ranges(o.pools.size(), 0); + std::vector pool_calls(o.pools.size(), 0), pool_remote_keys(o.pools.size(), 0); + for (auto& s : load_stats) { + restore.insert(restore.end(), s.restore_ms.begin(), s.restore_ms.end()); + ttfl.insert(ttfl.end(), s.ttfl_us.begin(), s.ttfl_us.end()); + stall.insert(stall.end(), s.stall_ms.begin(), s.stall_ms.end()); + get_bytes += s.get_bytes; + calls += s.get_calls; + calls_local += s.get_calls_local; + remote_calls += s.remote_calls; + remote_keys += s.remote_keys; + repeat_remote += s.repeat_remote; + local_write += s.local_write_bytes; + mismatches += s.verify_mismatches; + for (size_t pi = 0; pi < o.pools.size(); ++pi) { + pool_bytes[pi] += s.pool_bytes[pi]; + pool_ranges[pi] += s.pool_ranges[pi]; + pool_calls[pi] += s.pool_calls[pi]; + pool_remote_keys[pi] += s.pool_remote_keys[pi]; + } + } + for (auto& s : offload_stats) put_bytes += s.put_bytes; + + constexpr double kGiB = 1024.0 * 1024.0 * 1024.0; + constexpr double kMiB = 1024.0 * 1024.0; + // Both wire figures, so a before/after pair can be compared without knowing + // which fetch granularity the binary under test used. + const double wire_whole_mib = remote_keys * object_bytes / kMiB; + const double wire_range_mib = remote_keys * o.group_bytes() / kMiB; + std::printf( + "%zu,%zu,%zu,%zu,%zu,%.2f,%.2f,%.1f,%.1f,%.2f,%.4f,%.4f,%.1f,%zu,%zu,%zu,%.1f,%.1f,%.1f\n", + ranks, o.layers, o.layer_group, o.pages, object_bytes / 1024, Percentile(restore, 0.50), + Percentile(restore, 0.95), Percentile(ttfl, 0.50), Percentile(ttfl, 0.95), + Percentile(stall, 0.50), wall_s > 0 ? (get_bytes / kGiB) / wall_s : 0.0, + wall_s > 0 ? (put_bytes / kGiB) / wall_s : 0.0, calls > 0 ? 100.0 * calls_local / calls : 0.0, + remote_calls, remote_keys, repeat_remote, wire_whole_mib, wire_range_mib, local_write / kMiB); + std::fflush(stdout); + + // Per-pool split. On stderr so the CSV schema above stays stable, and only + // when there is more than one pool to split. range_kib is the number that + // predicts the rate: this path is fragment-size bound, so two pools moving + // equal bytes at different range sizes will not move them at equal speed. + if (o.pools.size() > 1) { + std::fprintf(stderr, " %-16s %9s %10s %8s %10s %9s %10s\n", "pool", "range_kib", "ranges", + "calls", "bytes_mib", "share_pct", "remote_keys"); + for (size_t pi = 0; pi < o.pools.size(); ++pi) { + std::fprintf(stderr, " %-16s %9.2f %10zu %8zu %10.1f %9.1f %10zu\n", + o.pools[pi].name.c_str(), o.pools[pi].layer_bytes / 1024.0, pool_ranges[pi], + pool_calls[pi], pool_bytes[pi] / kMiB, + get_bytes > 0 ? 100.0 * pool_bytes[pi] / get_bytes : 0.0, pool_remote_keys[pi]); + } + std::fflush(stderr); + } + + if (o.verify) { + std::fprintf(stderr, "verify: %zu mismatched pages at %zu ranks\n", mismatches, ranks); + if (mismatches != 0) std::exit(1); + } +} + +} // namespace + +int main(int argc, char** argv) { + BenchOpts opts; + if (!ParseArgs(argc, argv, &opts)) return 2; + + std::fprintf(stderr, + "layer-wise restore: %zu layers in groups of %zu (%zu groups/request), " + "%zu pages/request, page=%zu KiB across %zu pool(s), mean group slice=%zu KiB, " + "medium page=%zu KiB, arena=%zu MiB, locality_prefetch=%s, layout=%s\n", + opts.layers, opts.layer_group, opts.group_count(), opts.pages, + opts.object_bytes() / 1024, opts.pools.size(), opts.group_bytes() / 1024, + opts.medium_page_bytes() / 1024, opts.scratch_mib, + opts.locality_prefetch ? "on" : "off", + opts.interleave_groups ? "interleaved" : "contiguous"); + for (const Pool& pool : opts.pools) { + std::fprintf(stderr, " pool %-16s %3zu layers x %8zu B = %8zu KiB/page\n", pool.name.c_str(), + pool.layers.size(), pool.layer_bytes, pool.object_bytes() / 1024); + } + + std::printf( + "ranks,layers,group,pages,object_kib,restore_ms_p50,restore_ms_p95,ttfl_us_p50," + "ttfl_us_p95,stall_ms_p50,get_gibps,put_gibps,local_hit_pct," + "remote_calls,remote_keys,repeat_remote,wire_whole_mib,wire_range_mib,local_write_mib\n"); + std::fflush(stdout); + + for (size_t ranks : opts.rank_sweep) RunOne(opts, ranks); + return 0; +} diff --git a/tests/cpp/umbp/distributed/bench_route_put_strategy.cpp b/tests/cpp/umbp/distributed/bench_route_put_strategy.cpp index ed3d7fb1d..7617ac12f 100644 --- a/tests/cpp/umbp/distributed/bench_route_put_strategy.cpp +++ b/tests/cpp/umbp/distributed/bench_route_put_strategy.cpp @@ -37,7 +37,7 @@ // the get. This wait is OUTSIDE the timed regions. // * Keys are never reclaimed and a remote BatchGet leaves a ~500ms read lease // that defers capacity free on Clear(). So every measured iter starts by -// Clear()-ing all nodes and polling each allocator's TierCapacitiesSnapshot +// Clear()-ing all nodes and polling each allocator's Capacity() // until available==total (lease-deferred frees released by the reaper). // * Placement is observed via reader->Master().BatchRouteGet (returns per-key // node_id). This bumps lease/access; the subsequent timed BatchGet RouteGets @@ -95,7 +95,7 @@ #include "umbp/distributed/config.h" #include "umbp/distributed/master/master_client.h" #include "umbp/distributed/master/master_server.h" -#include "umbp/distributed/peer/peer_dram_allocator.h" +#include "umbp/distributed/peer/backend/page_backend.h" #include "umbp/distributed/pool_client.h" using mori::umbp::MasterServer; @@ -264,8 +264,7 @@ uint64_t NodeStorageBytes(const std::string& regime, size_t peers, uint64_t batc // One symmetric node: exportable storage DRAM (receives writes) plus a private // io buffer used as the put src (when requester) and the get dst (when reader). struct Node { - std::vector storage; // master-managed exportable DRAM - std::vector io; // registered src/dst region + std::vector io; // registered src/dst region std::unique_ptr client; }; @@ -292,7 +291,6 @@ class Cluster { nodes_.resize(peers); for (size_t k = 0; k < peers; ++k) { - nodes_[k].storage.assign(node_storage_bytes, 0); nodes_[k].io.assign(io_bytes, 0); PoolClientConfig cc; cc.master_config.node_id = "node-" + std::to_string(k); @@ -302,8 +300,7 @@ class Cluster { cc.io_engine.port = 0; cc.peer_service_port = NextPeerServicePort(); cc.dram_page_size = page_bytes; - cc.dram_buffers = {{nodes_[k].storage.data(), nodes_[k].storage.size()}}; - cc.tier_capacities = {{TierType::DRAM, {nodes_[k].storage.size(), nodes_[k].storage.size()}}}; + cc.dram.buffer_sizes = {node_storage_bytes}; nodes_[k].client = std::make_unique(std::move(cc)); if (!nodes_[k].client->Init()) { std::cerr << "node " << k << " init failed\n"; @@ -333,9 +330,10 @@ class Cluster { const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); for (auto& n : nodes_) { while (true) { - auto caps = n.client->DramAllocator()->TierCapacitiesSnapshot(); - auto it = caps.find(TierType::DRAM); - const bool full = it != caps.end() && it->second.available_bytes == it->second.total_bytes; + auto* dram = n.client->Backends().Get(TierType::DRAM); + if (dram == nullptr) return false; + auto cap = dram->Capacity(); + const bool full = cap.available_bytes == cap.total_bytes; if (full) break; if (std::chrono::steady_clock::now() > deadline) return false; std::this_thread::sleep_for(std::chrono::milliseconds(10)); diff --git a/tests/cpp/umbp/distributed/bench_route_strategy.cpp b/tests/cpp/umbp/distributed/bench_route_strategy.cpp index 3fb1640d4..89b746499 100644 --- a/tests/cpp/umbp/distributed/bench_route_strategy.cpp +++ b/tests/cpp/umbp/distributed/bench_route_strategy.cpp @@ -58,10 +58,10 @@ #include "umbp/distributed/routing/route_get_strategy.h" #include "umbp/distributed/types.h" +using mori::umbp::LocalPreferringRouteGetStrategy; using mori::umbp::Location; using mori::umbp::RandomRouteGetStrategy; using mori::umbp::RouteGetStrategy; -using mori::umbp::TierPriorityRouteGetStrategy; using mori::umbp::TierType; namespace { @@ -83,7 +83,11 @@ void Usage() { // Resolve the --get-strategy name to a concrete RouteGetStrategy. std::unique_ptr MakeGetStrategy(const std::string& name) { if (name == "random") return std::make_unique(); - if (name == "tier-priority") return std::make_unique(); + // "tier-priority" kept as an alias so existing invocations keep working; the + // strategy no longer ranks tiers (backend-agnostic refactor Phase 4). + if (name == "local-preferring" || name == "tier-priority") { + return std::make_unique(); + } std::cerr << "Unknown --get-strategy: " << name << " (expected random|tier-priority)\n"; std::exit(2); } diff --git a/tests/cpp/umbp/distributed/mock_master_metadata_store.h b/tests/cpp/umbp/distributed/mock_master_metadata_store.h deleted file mode 100644 index 811e5919a..000000000 --- a/tests/cpp/umbp/distributed/mock_master_metadata_store.h +++ /dev/null @@ -1,126 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -// GMock mock for IMasterMetadataStore. -// -// Phase 1 use: instantiation gate. If this type compiles and instantiates, -// every pure-virtual on the interface is overridden with a well-typed -// signature — proving the contract has no orphaned/ill-typed methods. -// -// Reused in Phase 3 (consumer-integration) to assert that each rewired -// consumer (Router / EvictionManager / UMBPMasterServiceImpl handlers) calls -// the right store method with correctly-translated arguments. -#pragma once - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/master/master_metadata_store.h" -#include "umbp/distributed/types.h" - -namespace mori::umbp { - -class MockMasterMetadataStore : public IMasterMetadataStore { - public: - // Aliases for types whose commas would otherwise break the MOCK_METHOD macro - // parser (it splits the argument list on top-level commas). - using CapsMap = std::map; - using LocationBatch = std::vector>; - - // --- Cross-store writes --- - MOCK_METHOD(bool, RegisterClient, - (const ClientRegistration& registration, std::chrono::system_clock::time_point now, - std::chrono::system_clock::duration stale_after), - (override)); - MOCK_METHOD(void, UnregisterClient, (const std::string& node_id), (override)); - MOCK_METHOD(HeartbeatResult, ApplyHeartbeat, - (const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, - const CapsMap& caps, (const std::vector&)events, bool is_full_sync), - (override)); - MOCK_METHOD(std::vector, ExpireStaleClients, - (std::chrono::system_clock::time_point cutoff), (override)); - - // --- External-KV writes --- - MOCK_METHOD(bool, RegisterExternalKvIfAlive, - (const std::string& node_id, (const std::vector&)hashes, TierType tier), - (override)); - MOCK_METHOD(void, UnregisterExternalKv, - (const std::string& node_id, (const std::vector&)hashes, TierType tier), - (override)); - MOCK_METHOD(void, UnregisterExternalKvByTier, (const std::string& node_id, TierType tier), - (override)); - MOCK_METHOD(void, UnregisterExternalKvByNode, (const std::string& node_id), (override)); - MOCK_METHOD(std::size_t, GarbageCollectHits, (std::chrono::system_clock::time_point cutoff), - (override)); - - // --- Block reads --- - MOCK_METHOD(std::vector, LookupBlock, (const std::string& key), (const, override)); - MOCK_METHOD(std::vector, LookupBlockForRouteGet, - (const std::string& key, (const std::unordered_set&)exclude_nodes, - std::chrono::system_clock::time_point now, - std::chrono::system_clock::duration lease_duration), - (override)); - MOCK_METHOD(LocationBatch, BatchLookupBlockForRouteGet, - ((const std::vector&)keys, - (const std::unordered_set&)exclude_nodes, - std::chrono::system_clock::time_point now, - std::chrono::system_clock::duration lease_duration), - (override)); - MOCK_METHOD(std::vector, BatchExistsBlock, ((const std::vector&)keys), - (const, override)); - MOCK_METHOD((std::map>), EnumerateEvictionCandidates, - (const std::vector& buckets, EvictionOrder order, size_t max_per_bucket, - std::chrono::system_clock::time_point now), - (const, override)); - - // --- Client reads --- - MOCK_METHOD(std::optional, GetClient, (const std::string& node_id), - (const, override)); - MOCK_METHOD(bool, IsClientAlive, (const std::string& node_id), (const, override)); - MOCK_METHOD(std::optional, GetPeerAddress, (const std::string& node_id), - (const, override)); - MOCK_METHOD(std::vector, ListAliveClients, (), (const, override)); - MOCK_METHOD((std::unordered_map), GetAlivePeerView, (), - (const, override)); - MOCK_METHOD(std::size_t, AliveClientCount, (), (const, override)); - MOCK_METHOD(std::vector, GetClientTags, (const std::string& node_id), - (const, override)); - - // --- External-KV reads --- - MOCK_METHOD(std::vector, MatchExternalKv, - ((const std::vector&)hashes, bool count_as_hit, - std::chrono::system_clock::time_point now), - (override)); - MOCK_METHOD(std::vector, GetExternalKvHitCounts, - ((const std::vector&)hashes), (const, override)); - MOCK_METHOD(std::size_t, GetExternalKvCount, (const std::string& node_id), (const, override)); -}; - -} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_backend_policy_config.cpp b/tests/cpp/umbp/distributed/test_backend_policy_config.cpp new file mode 100644 index 000000000..34e6b337f --- /dev/null +++ b/tests/cpp/umbp/distributed/test_backend_policy_config.cpp @@ -0,0 +1,313 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include + +#include + +#include "umbp/distributed/config.h" +#include "umbp/distributed/pool/policy_config.h" + +namespace mori::umbp { +namespace { + +// The three-tier example we ship, read from disk rather than copied here, so +// that editing one without the other fails this test. +TEST(BackendPolicyConfig, ParsesAndLowersShippedExample) { + auto loaded = LoadBackendPolicyFile(UMBP_EXAMPLE_POLICY_PATH); + ASSERT_TRUE(loaded.ok()) << UMBP_EXAMPLE_POLICY_PATH << ": " << loaded.error; + EXPECT_EQ(loaded.config->schema_version, 1u); + EXPECT_EQ(loaded.config->entry_tier, "hot"); + ASSERT_EQ(loaded.config->backends.size(), 4u); + ASSERT_EQ(loaded.config->logical_tiers.size(), 3u); + + PoolClientConfig output; + output.dram_page_size = 2ULL << 20; + std::string error; + ASSERT_TRUE(ApplyBackendPolicy(*loaded.config, &output, &error, "-node-0")) << error; + EXPECT_EQ(output.placement_policy, PoolPlacementPolicy::TIERED); + ASSERT_EQ(output.backends.size(), 5u); + + const auto find_backend = [&](const std::string& name) -> const BackendInstanceConfig* { + for (const auto& backend : output.backends) { + if (backend.name == name) return &backend; + } + return nullptr; + }; + ASSERT_NE(find_backend("dram"), nullptr); + EXPECT_EQ(find_backend("dram")->dram.numa_node, 0); + ASSERT_NE(find_backend("hbm@0"), nullptr); + ASSERT_NE(find_backend("hbm@1"), nullptr); + EXPECT_EQ(find_backend("hbm@0")->hbm.buffer_sizes.front(), 40ULL << 30); + ASSERT_NE(find_backend("ssd_a"), nullptr); + EXPECT_EQ(find_backend("ssd_a")->ssd.ssd.storage_dir, "/mnt/kvcache/hot-node-0"); + + // One backend with two devices lowers to one member per device, splitting the + // declared capacity evenly, so the tier weights stay equal. + ASSERT_EQ(output.logical_tiers[0].members.size(), 2u); + EXPECT_EQ(output.logical_tiers[0].members[0].weight, output.logical_tiers[0].members[1].weight); + EXPECT_EQ(output.logical_tiers[0].offload_to, (std::vector{"dram", "ssd_a"})); + EXPECT_EQ(output.logical_tiers[0].trigger, PoolOffloadTrigger::kOnEvict); + EXPECT_EQ(output.logical_tiers[0].name, "hot"); + EXPECT_TRUE(output.logical_tiers[0].entry); + + ASSERT_EQ(output.logical_tiers[1].members.size(), 2u); + EXPECT_EQ(output.logical_tiers[1].members[0].backend_name, "dram"); + EXPECT_EQ(output.logical_tiers[1].members[0].weight, 70u); + EXPECT_EQ(output.logical_tiers[1].members[1].backend_name, "ssd_a"); + EXPECT_EQ(output.logical_tiers[1].members[1].weight, 30u); + EXPECT_EQ(output.logical_tiers[1].trigger, PoolOffloadTrigger::kWatermark); + EXPECT_DOUBLE_EQ(output.logical_tiers[1].high_watermark, 0.9); + EXPECT_DOUBLE_EQ(output.logical_tiers[1].low_watermark, 0.7); + + EXPECT_EQ(output.logical_tiers[2].promote_trigger, PoolPromoteTrigger::kOnRead); + EXPECT_EQ(output.logical_tiers[2].promote_mode, PoolTransitionMode::kCopy); + EXPECT_TRUE(output.logical_tiers[2].offload_to.empty()); +} + +TEST(BackendPolicyConfig, RejectsInvalidSchemaAndBackwardEdges) { + auto unknown = LoadBackendPolicyJson(R"json( + {"backends":{"d":{"type":"dram","capacity":"1GiB","extra":1}}, + "tiers":[{"backends":{"d":1}}]} + )json"); + EXPECT_FALSE(unknown.ok()); + EXPECT_NE(unknown.error.find("unknown field"), std::string::npos); + + auto backward = LoadBackendPolicyJson(R"json( + {"backends":{ + "a":{"type":"dram","capacity":"1GiB"}, + "b":{"type":"ssd","capacity":"1GiB","path":"/tmp/b"}}, + "tiers":[ + {"backends":{"a":1}}, + {"backends":{"b":1},"offload_to":["a"],"offload_trigger":"on_evict"}]} + )json"); + EXPECT_FALSE(backward.ok()); + EXPECT_NE(backward.error.find("strictly later tier"), std::string::npos); + + // Tier invariants live in the shared validator rather than in the JSON + // parser, so these two pin the wiring that runs it on the loading path. + auto watermarks = LoadBackendPolicyJson(R"json( + {"backends":{"d":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"d":1},"high_watermark":0.2,"low_watermark":0.9}]} + )json"); + EXPECT_FALSE(watermarks.ok()); + EXPECT_NE(watermarks.error.find("watermarks must satisfy"), std::string::npos); + + auto shared_backend = LoadBackendPolicyJson(R"json( + {"backends":{"d":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"d":1}},{"backends":{"d":1}}]} + )json"); + EXPECT_FALSE(shared_backend.ok()); + EXPECT_NE(shared_backend.error.find("more than one tier"), std::string::npos); + + // Loading and applying share one validator, so a topology that cannot be + // expanded is rejected by whichever entry point sees it first. + auto tiny_hbm = LoadBackendPolicyJson(R"json( + {"backends":{"h":{"type":"hbm","capacity":"1B","devices":[0,1]}}, + "tiers":[{"backends":{"h":1}}]} + )json"); + EXPECT_FALSE(tiny_hbm.ok()); + EXPECT_NE(tiny_hbm.error.find("per-device capacity"), std::string::npos); +} + +// The shipped example only exercises promote_mode "copy", so nothing pins the +// other branch of the enum: a policy asking for move would have parsed to copy +// and silently kept the source tier's copy alive. +TEST(BackendPolicyConfig, LowersMovePromotionMode) { + auto loaded = LoadBackendPolicyJson(R"json( + {"backends":{ + "h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "tiers":[ + {"backends":{"h":1},"offload_to":["cold"],"offload_trigger":"on_evict","name":"hot"}, + {"backends":{"c":1},"name":"cold","promote_trigger":"on_read","promote_mode":"move"}]} + )json"); + ASSERT_TRUE(loaded.ok()) << loaded.error; + + PoolClientConfig output; + output.dram_page_size = 64ULL << 10; + std::string error; + ASSERT_TRUE(ApplyBackendPolicy(*loaded.config, &output, &error)) << error; + ASSERT_EQ(output.logical_tiers.size(), 2u); + EXPECT_EQ(output.logical_tiers[1].promote_trigger, PoolPromoteTrigger::kOnRead); + EXPECT_EQ(output.logical_tiers[1].promote_mode, PoolTransitionMode::kMove); + // The entry tier keeps the default, so one tier asking for move does not + // change how the rest of the graph promotes. + EXPECT_EQ(output.logical_tiers[0].promote_mode, PoolTransitionMode::kCopy); + + auto bad = LoadBackendPolicyJson(R"json( + {"backends":{"d":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"d":1},"promote_mode":"relocate"}]} + )json"); + EXPECT_FALSE(bad.ok()); +} + +// promote_trigger is the extension point the promote policy is built on, so the +// rules that keep one policy from having two spellings are worth pinning: a +// threshold only where something reads it, never below 2, and never on the +// entry tier, which has no upstream tier to promote into. +TEST(BackendPolicyConfig, LowersPromoteTriggerAndRejectsIncoherentThresholds) { + auto loaded = LoadBackendPolicyJson(R"json( + {"backends":{ + "h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "tiers":[ + {"backends":{"h":1},"offload_to":["cold"],"offload_trigger":"watermark","name":"hot"}, + {"backends":{"c":1},"name":"cold","promote_trigger":"on_hits","promote_hits":4}]} + )json"); + ASSERT_TRUE(loaded.ok()) << loaded.error; + + PoolClientConfig output; + output.dram_page_size = 64ULL << 10; + std::string error; + ASSERT_TRUE(ApplyBackendPolicy(*loaded.config, &output, &error)) << error; + ASSERT_EQ(output.logical_tiers.size(), 2u); + EXPECT_EQ(output.logical_tiers[1].promote_trigger, PoolPromoteTrigger::kOnHits); + EXPECT_EQ(output.logical_tiers[1].promote_hits, 4u); + // Absent means never, so a tier says nothing about promotion by saying + // nothing rather than by carrying a threshold nothing reads. + EXPECT_EQ(output.logical_tiers[0].promote_trigger, PoolPromoteTrigger::kNever); + EXPECT_EQ(output.logical_tiers[0].promote_hits, 0u); + + const char* kRejected[] = { + // on_hits without a threshold: there is no default worth guessing. + R"json({"backends":{"h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"h":1},"offload_to":["cold"], + "offload_trigger":"on_evict","name":"hot"}, + {"backends":{"c":1},"name":"cold","promote_trigger":"on_hits"}]})json", + // A threshold where nothing reads it is a misconfiguration, not a hint. + R"json({"backends":{"h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"h":1},"offload_to":["cold"], + "offload_trigger":"on_evict","name":"hot"}, + {"backends":{"c":1},"name":"cold","promote_trigger":"on_read", + "promote_hits":3}]})json", + // 1 hit is on_read spelled a second way. + R"json({"backends":{"h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"h":1},"offload_to":["cold"], + "offload_trigger":"on_evict","name":"hot"}, + {"backends":{"c":1},"name":"cold","promote_trigger":"on_hits", + "promote_hits":1}]})json", + // The entry tier has nowhere to promote to; today this is silently dead + // config because the read path skips it. + R"json({"backends":{"h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "entry_tier":"hot", + "tiers":[{"backends":{"h":1},"offload_to":["cold"], + "offload_trigger":"on_evict","name":"hot", + "promote_trigger":"on_read"}, + {"backends":{"c":1},"name":"cold"}]})json", + // Unknown enumerator, so a typo cannot read as "never". + R"json({"backends":{"d":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"d":1},"promote_trigger":"on_every_read"}]})json", + // The old boolean must not keep working, or policies keep using it. + R"json({"backends":{"h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"h":1},"offload_to":["cold"], + "offload_trigger":"on_evict","name":"hot"}, + {"backends":{"c":1},"name":"cold","promote_on_read":true}]})json", + // Nor the pre-rename spelling of promote_mode, for the same reason. + R"json({"backends":{"h":{"type":"dram","capacity":"1GiB"}, + "c":{"type":"dram","capacity":"1GiB"}}, + "tiers":[{"backends":{"h":1},"offload_to":["cold"], + "offload_trigger":"on_evict","name":"hot"}, + {"backends":{"c":1},"name":"cold","promote_trigger":"on_read", + "promotion_mode":"move"}]})json", + }; + for (const char* json : kRejected) { + EXPECT_FALSE(LoadBackendPolicyJson(json).ok()) << json; + } +} + +// An SSD backend's staging arena is staging_slots * page_size, and a read that +// cannot claim a slot is reported as a miss rather than as backpressure, so the +// slot count is a hard read-concurrency ceiling. Left at the default of 16 a +// 64 KiB-page backend has a 1 MiB arena, which under a mixed workload fails the +// majority of its reads and, because the offload target then refuses pages, +// most of its writes too. Omitting the field must keep the previous default so +// existing policies are unaffected. +TEST(BackendPolicyConfig, LowersSsdStagingSlotsAndDefaultsWhenAbsent) { + auto loaded = LoadBackendPolicyJson(R"json( + {"backends":{ + "with_slots":{"type":"ssd","capacity":"1GiB","path":"/tmp/a","staging_slots":512}, + "without":{"type":"ssd","capacity":"1GiB","path":"/tmp/b"}}, + "tiers":[{"backends":{"with_slots":1}},{"backends":{"without":1}}]} + )json"); + ASSERT_TRUE(loaded.ok()) << loaded.error; + + PoolClientConfig output; + output.dram_page_size = 64ULL << 10; + std::string error; + ASSERT_TRUE(ApplyBackendPolicy(*loaded.config, &output, &error)) << error; + + const auto slots_of = [&](const std::string& name) { + for (const auto& backend : output.backends) { + if (backend.name == name) return backend.ssd_staging_buffer_slots; + } + ADD_FAILURE() << "backend '" << name << "' missing"; + return -1; + }; + EXPECT_EQ(slots_of("with_slots"), 512); + EXPECT_EQ(slots_of("without"), BackendInstanceConfig{}.ssd_staging_buffer_slots); + + // A DRAM backend has no staging arena, so the field must not be silently + // accepted there and then ignored. + auto on_dram = LoadBackendPolicyJson(R"json( + {"backends":{"d":{"type":"dram","capacity":"1GiB","staging_slots":512}}, + "tiers":[{"backends":{"d":1}}]} + )json"); + EXPECT_FALSE(on_dram.ok()); + EXPECT_NE(on_dram.error.find("unknown field"), std::string::npos); + + auto negative = LoadBackendPolicyJson(R"json( + {"backends":{"s":{"type":"ssd","capacity":"1GiB","path":"/tmp/c","staging_slots":-1}}, + "tiers":[{"backends":{"s":1}}]} + )json"); + EXPECT_FALSE(negative.ok()); +} + +TEST(BackendPolicyConfig, ApplyIsAtomicWhenLoweringFails) { + auto loaded = LoadBackendPolicyJson(R"json( + {"backends":{"d":{"type":"dram","capacity":"1MiB"}}, + "tiers":[{"backends":{"d":1}}]} + )json"); + ASSERT_TRUE(loaded.ok()) << loaded.error; + + PoolClientConfig output; + // Only Apply can reject this: whether a capacity is usable depends on the + // page size of the client the policy is being applied to. + output.dram_page_size = 2ULL << 20; + output.backends.push_back(BackendInstanceConfig{}); + output.backends.front().name = "unchanged"; + std::string error; + EXPECT_FALSE(ApplyBackendPolicy(*loaded.config, &output, &error)); + EXPECT_NE(error.find("smaller than page size"), std::string::npos); + ASSERT_EQ(output.backends.size(), 1u); + EXPECT_EQ(output.backends.front().name, "unchanged"); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_backend_registry.cpp b/tests/cpp/umbp/distributed/test_backend_registry.cpp new file mode 100644 index 000000000..4e749ebfc --- /dev/null +++ b/tests/cpp/umbp/distributed/test_backend_registry.cpp @@ -0,0 +1,174 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// BackendRegistry dispatch and cross-backend event aggregation. MockBackend is +// only the test double; its own behavior is exercised transitively here and by +// the PeerPool / PeerService suites. + +#include + +#include +#include +#include + +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/peer/backend/mock_backend.h" + +namespace mori::umbp { + +// The registry dispatches to whichever concrete backend is registered for a +// tier, without the caller naming MockBackend anywhere below this line — this +// is the acceptance property the design doc calls out for Phase 2. +TEST(BackendRegistryDispatch, RoutesByTierThroughTheInterfaceOnly) { + BackendRegistry registry; + registry.Register(std::make_unique(TierType::HBM)); + + MediumBackend* hbm = registry.Get(TierType::HBM); + ASSERT_NE(hbm, nullptr); + EXPECT_EQ(hbm->Tier(), TierType::HBM); + EXPECT_EQ(registry.Get(TierType::DRAM), nullptr); + + auto allocated = hbm->BatchAllocate({{"key", 64}}); + ASSERT_EQ(allocated.size(), 1u); + ASSERT_EQ(allocated[0].outcome, AllocateOutcome::kSuccessAllocated); + hbm->BatchCommit({{allocated[0].slot_id, "key"}}); + EXPECT_EQ(hbm->OwnedKeyCount(), 1u); + + ASSERT_EQ(registry.All().size(), 1u); + EXPECT_EQ(registry.All()[0], hbm); +} + +TEST(BackendRegistryDispatch, NamedSameTierInstancesKeepDenseStableIds) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram-primary", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram-secondary", std::make_unique(TierType::DRAM))); + + ASSERT_EQ(registry.Size(), 2u); + EXPECT_EQ(registry.Get(TierType::DRAM), registry.Get("dram-primary")); + EXPECT_EQ(registry.Get(0u), registry.Get("dram-primary")); + EXPECT_EQ(registry.Get(1u), registry.Get("dram-secondary")); + EXPECT_EQ(registry.BackendId(registry.Get("dram-secondary")), 1u); + + auto allocated = registry.Get("dram-primary")->BatchAllocate({{"old-state", 8}}); + registry.Get("dram-primary")->BatchCommit({{allocated.front().slot_id, "old-state"}}); + ASSERT_EQ(registry.Get("dram-primary")->OwnedKeyCount(), 1u); + ASSERT_TRUE(registry.Register("dram-primary", std::make_unique(TierType::DRAM))); + EXPECT_EQ(registry.Get("dram-primary")->OwnedKeyCount(), 0u); + EXPECT_EQ(registry.Size(), 2u); + EXPECT_EQ(registry.BackendId(registry.Get("dram-primary")), 0u); + ASSERT_EQ(registry.All().size(), 2u); + EXPECT_EQ(registry.All()[0], registry.Get("dram-primary")); + EXPECT_EQ(registry.All()[1], registry.Get("dram-secondary")); +} + +// ---- Heartbeat event aggregation ------------------------------------------- +// These replace the OwnedLocationSourceAgg tests that lived in +// test_peer_ssd_manager.cpp: OwnedLocationSource is gone and MediumBackend +// carries the event contract (design doc §3, backend-agnostic refactor +// Phase 3). What matters is unchanged — every medium's events concat into ONE +// list, in order, so the shipper can wrap them in a single bundle under one +// monotonic seq. + +namespace { + +// Commit `key` on `b` so exactly one ADD lands in its outbox. +void CommitOne(MediumBackend* b, const std::string& key, uint64_t size) { + auto allocated = b->BatchAllocate({{key, size}}); + b->BatchCommit({{allocated[0].slot_id, key}}); +} + +} // namespace + +TEST(BackendEventAgg, DrainAndSnapshotConcatAcrossBackendsInOrder) { + MockBackend hbm(TierType::HBM); + MockBackend ssd(TierType::SSD); + CommitOne(&hbm, "h1", 10); + CommitOne(&ssd, "s1", 30); + + std::vector backends = {&hbm, &ssd}; + + auto drained = DrainAllBackends(backends); + ASSERT_EQ(drained.size(), 2u); + EXPECT_EQ(drained[0].key, "h1"); + EXPECT_EQ(drained[0].tier, TierType::HBM); + EXPECT_EQ(drained[1].key, "s1"); + EXPECT_EQ(drained[1].tier, TierType::SSD); + // Drain cleared both outboxes. + EXPECT_TRUE(DrainAllBackends(backends).empty()); + + // The full-sync snapshot reports owned keys regardless of the drained outbox. + auto snap = SnapshotAllBackendsForFullSync(backends); + ASSERT_EQ(snap.size(), 2u); + EXPECT_EQ(snap[1].tier, TierType::SSD); +} + +TEST(BackendEventAgg, FullSyncSnapshotAlsoDropsTheOutbox) { + MockBackend b(TierType::HBM); + CommitOne(&b, "x", 1); + + // Snapshot is authoritative, so the still-queued ADD must not be re-shipped + // as a redundant delta afterwards. + auto snap = SnapshotAllBackendsForFullSync({&b}); + ASSERT_EQ(snap.size(), 1u); + EXPECT_TRUE(DrainAllBackends({&b}).empty()); +} + +TEST(BackendEventAgg, SameTierRelocationKeepsMasterLocation) { + MockBackend source(TierType::SSD); + MockBackend target(TierType::SSD); + CommitOne(&source, "moved", 8); + ASSERT_EQ(DrainAllBackends({&source, &target}).size(), 1u); + + CommitOne(&target, "moved", 8); + source.Evict({"moved"}); + auto relocated = DrainAllBackends({&source, &target}); + ASSERT_EQ(relocated.size(), 1u); + EXPECT_EQ(relocated.front().kind, KvEvent::Kind::ADD); + EXPECT_EQ(relocated.front().tier, TierType::SSD); + + target.Evict({"moved"}); + auto removed = DrainAllBackends({&source, &target}); + ASSERT_EQ(removed.size(), 1u); + EXPECT_EQ(removed.front().kind, KvEvent::Kind::REMOVE); +} + +TEST(BackendEventAgg, NullBackendsAreSkipped) { + MockBackend only(TierType::SSD); + CommitOne(&only, "x", 1); + std::vector backends = {nullptr, &only, nullptr}; + auto drained = DrainAllBackends(backends); + ASSERT_EQ(drained.size(), 1u); + EXPECT_EQ(drained[0].key, "x"); + EXPECT_TRUE(SnapshotAllBackendsForFullSync({nullptr}).empty()); +} + +TEST(BackendEventAgg, AutoFlushHookFiresAtThreshold) { + MockBackend b(TierType::HBM); + int fires = 0; + b.SetAutoFlushHook(2, [&] { ++fires; }); + CommitOne(&b, "a", 1); + EXPECT_EQ(fires, 0); + CommitOne(&b, "b", 1); + EXPECT_EQ(fires, 1); +} + +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_batch_resolve_codec.cpp b/tests/cpp/umbp/distributed/test_batch_resolve_codec.cpp index 29df4b893..3ea6b3618 100644 --- a/tests/cpp/umbp/distributed/test_batch_resolve_codec.cpp +++ b/tests/cpp/umbp/distributed/test_batch_resolve_codec.cpp @@ -49,10 +49,10 @@ BufferMemoryDescBytes MakeDesc(uint32_t buffer_index, const std::string& bytes) return d; } -// A representative batch: 4 keys mixing found / not-found, varying page counts +// A representative batch mixing found / not-found / busy, varying page counts // (including a found-but-zero-page key), and several tiers. std::vector SampleEntries() { - std::vector entries(4); + std::vector entries(5); entries[0].found = true; entries[0].tier = TierType::HBM; @@ -71,11 +71,15 @@ std::vector SampleEntries() { entries[3].size = 0; // found but zero pages (degenerate but legal) entries[3].pages = {}; + entries[4].outcome = ResolveOutcome::kBusy; + entries[4].tier = TierType::SSD; + return entries; } void ExpectKeyMatches(const ResolvedKeyEntry& src, const DecodedResolveKey& dec) { EXPECT_EQ(src.found, dec.found); + EXPECT_EQ(src.found ? ResolveOutcome::kFound : src.outcome, dec.outcome); if (!src.found) return; EXPECT_EQ(src.tier, dec.tier); EXPECT_EQ(src.size, dec.size); @@ -179,6 +183,23 @@ TEST(BatchResolveCodec, EmptyBatch) { EXPECT_EQ(decoded.page_size, 4096u); } +TEST(BatchResolveCodec, DecodesLegacyFoundOnlyResponse) { + ::umbp::BatchResolveKeysResponse resp; + resp.set_page_size(4096); + resp.add_found(true); + resp.add_tier(::umbp::TIER_DRAM); + resp.add_size(4096); + resp.add_page_count(1); + resp.add_backend_id(0); + resp.add_buffer_index(0); + resp.add_page_index(7); + + DecodedBatchResolve decoded = DecodeBatchResolveResponse(resp); + ASSERT_EQ(decoded.keys.size(), 1u); + EXPECT_TRUE(decoded.keys[0].found); + EXPECT_EQ(decoded.keys[0].outcome, ResolveOutcome::kFound); +} + TEST(BatchResolveCodec, MalformedMismatchedArraysRejected) { // found has 2 entries but the per-key arrays disagree -> decoder must refuse. ::umbp::BatchResolveKeysResponse resp; diff --git a/tests/cpp/umbp/distributed/test_batch_route_get_columnar.cpp b/tests/cpp/umbp/distributed/test_batch_route_get_columnar.cpp index de98cc478..37e0f0f0a 100644 --- a/tests/cpp/umbp/distributed/test_batch_route_get_columnar.cpp +++ b/tests/cpp/umbp/distributed/test_batch_route_get_columnar.cpp @@ -45,7 +45,6 @@ #include #include -#include #include #include #include @@ -61,15 +60,6 @@ namespace mori::umbp { namespace { -static uint16_t AllocPort() { - static std::atomic next{0}; - if (next.load() == 0) { - std::srand(static_cast(std::time(nullptr))); - next.store(static_cast(52000 + (std::rand() % 1500))); - } - return next.fetch_add(7); -} - // Canned per-key routing answer the mock server hands back. A nullopt entry // (or a key absent from the table) is "not found". struct Route { @@ -136,13 +126,19 @@ class MockMasterService final : public ::umbp::UMBPMaster::Service { class BatchRouteGetColumnarTest : public ::testing::Test { protected: void SetUp() override { - address_ = "127.0.0.1:" + std::to_string(AllocPort()); service_ = std::make_unique(); + + // Port 0 asks the OS for a free one and reports it back; a guessed port + // collides on a busy host and BuildAndStart signals that only by + // returning nullptr. + int selected_port = 0; grpc::ServerBuilder builder; - builder.AddListeningPort(address_, grpc::InsecureServerCredentials()); + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &selected_port); builder.RegisterService(service_.get()); server_ = builder.BuildAndStart(); ASSERT_NE(server_, nullptr); + ASSERT_NE(selected_port, 0); + address_ = "127.0.0.1:" + std::to_string(selected_port); UMBPMasterClientConfig cfg; cfg.node_id = "caller"; diff --git a/tests/cpp/umbp/distributed/test_component_metrics.cpp b/tests/cpp/umbp/distributed/test_component_metrics.cpp new file mode 100644 index 000000000..b45f93d2b --- /dev/null +++ b/tests/cpp/umbp/distributed/test_component_metrics.cpp @@ -0,0 +1,473 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// --------------------------------------------------------------------------- +// The metrics seam: MetricPublisher, InstrumentedBackend, and the transfer +// layer's per-engine accounting. +// +// The load-bearing test in this file is InstrumentedBackendMetrics. +// MetricsAreGenericOverBackends: MockBackend contains NOT ONE LINE of metrics +// code, and the assertions still hold. That is the property the refactor is +// supposed to guarantee — a medium added later reports the same series and +// lands in the same dashboard panels — and it is the property that silently +// rots if someone reintroduces per-medium metric names or moves counting back +// into the call sites. Keep this test backend-agnostic: if it ever needs to +// know which medium it wrapped, the property is gone. +// --------------------------------------------------------------------------- + +#include + +#include +#include +#include +#include + +#include "umbp/distributed/metrics/component_metrics.h" +#include "umbp/distributed/peer/backend/instrumented_backend.h" +#include "umbp/distributed/peer/backend/mock_backend.h" +#include "umbp/distributed/transfer/composite_transfer_engine.h" +#include "umbp/distributed/transfer/local_copy_engine.h" + +namespace mori::umbp { +namespace { + +// A metric identity as a dashboard sees it: name plus its sorted label set. +std::string Identity(const char* name, const MetricLabels& labels) { + std::map sorted(labels.begin(), labels.end()); + std::string out = name; + for (const auto& [k, v] : sorted) { + out += ','; + out += k; + out += '='; + out += v; + } + return out; +} + +// Collects what a publisher emitted, so a test can assert on the wire form +// rather than on a component's internals. +struct Collector { + std::map counters; // identity -> summed delta + std::map gauges; // identity -> last value + + MetricPublisher::Sink Sink() { + return MetricPublisher::Sink{ + [this](const char* name, const char*, const MetricLabels& labels, double delta) { + counters[Identity(name, labels)] += delta; + }, + [this](const char* name, const char*, const MetricLabels& labels, double value) { + gauges[Identity(name, labels)] = value; + }}; + } + + bool Has(const std::string& identity) const { return counters.count(identity) > 0; } + double Counter(const std::string& identity) const { + auto it = counters.find(identity); + return it == counters.end() ? -1.0 : it->second; + } +}; + +// A MetricSource whose samples a test controls outright. +class FakeSource : public MetricSource { + public: + std::vector samples; + std::vector SampleMetrics() const override { return samples; } +}; + +constexpr const char* kName = "test_metric_total"; +constexpr const char* kHelp = "help"; + +// --------------------------------------------------------------------------- +// MetricPublisher +// --------------------------------------------------------------------------- + +TEST(MetricPublisher, ShipsDeltasNotAbsoluteValues) { + FakeSource src; + src.samples = {MetricSample{kName, kHelp, {}, 10}}; + + MetricPublisher pub; + Collector c; + pub.Publish("src", {}, src, c.Sink()); + EXPECT_DOUBLE_EQ(c.Counter(kName), 10.0); + + // Same absolute counter, one tick later: nothing new happened. + c.counters.clear(); + pub.Publish("src", {}, src, c.Sink()); + EXPECT_FALSE(c.Has(kName)) << "an unchanged counter must not re-ship its total"; + + src.samples[0].value = 25; + pub.Publish("src", {}, src, c.Sink()); + EXPECT_DOUBLE_EQ(c.Counter(kName), 15.0); +} + +TEST(MetricPublisher, RebasesInsteadOfShippingNegativeDeltas) { + FakeSource src; + src.samples = {MetricSample{kName, kHelp, {}, 100}}; + MetricPublisher pub; + Collector c; + pub.Publish("src", {}, src, c.Sink()); + + // A component that was torn down and rebuilt restarts its counter. The + // publisher must absorb that, not emit a negative delta. + c.counters.clear(); + src.samples[0].value = 5; + pub.Publish("src", {}, src, c.Sink()); + EXPECT_FALSE(c.Has(kName)); + + // ...and the next real increment is measured against the NEW baseline. + src.samples[0].value = 8; + pub.Publish("src", {}, src, c.Sink()); + EXPECT_DOUBLE_EQ(c.Counter(kName), 3.0); +} + +TEST(MetricPublisher, SourcesDoNotShareABaseline) { + // Two components reporting the same metric name: if they shared a baseline, + // one's progress would cancel the other's out. This is the bug the source id + // exists to prevent, and it is exactly the shape two backends take. + FakeSource a; + FakeSource b; + a.samples = {MetricSample{kName, kHelp, {}, 10}}; + b.samples = {MetricSample{kName, kHelp, {}, 1000}}; + + MetricPublisher pub; + Collector c; + pub.Publish("a", {{"tier", "DRAM"}}, a, c.Sink()); + pub.Publish("b", {{"tier", "SSD"}}, b, c.Sink()); + + EXPECT_DOUBLE_EQ(c.Counter(Identity(kName, {{"tier", "DRAM"}})), 10.0); + EXPECT_DOUBLE_EQ(c.Counter(Identity(kName, {{"tier", "SSD"}})), 1000.0); + + c.counters.clear(); + a.samples[0].value = 11; + pub.Publish("a", {{"tier", "DRAM"}}, a, c.Sink()); + pub.Publish("b", {{"tier", "SSD"}}, b, c.Sink()); + EXPECT_DOUBLE_EQ(c.Counter(Identity(kName, {{"tier", "DRAM"}})), 1.0); + EXPECT_FALSE(c.Has(Identity(kName, {{"tier", "SSD"}}))); +} + +TEST(MetricPublisher, GaugesShipEveryTickAndIgnoreBaselines) { + FakeSource src; + src.samples = {MetricSample{"test_gauge", kHelp, {}, 7, MetricKind::kGauge}}; + MetricPublisher pub; + Collector c; + + pub.Publish("src", {}, src, c.Sink()); + EXPECT_DOUBLE_EQ(c.gauges["test_gauge"], 7.0); + + // A gauge that goes DOWN is normal (an arena drains), and it must be + // reported as-is rather than swallowed the way a counter regression is. + src.samples[0].value = 2; + pub.Publish("src", {}, src, c.Sink()); + EXPECT_DOUBLE_EQ(c.gauges["test_gauge"], 2.0); +} + +TEST(MetricPublisher, AppliesScaleSoNamedUnitsStayHonest) { + // Components accumulate in whatever integer unit an atomic add can carry; + // the metric still has to mean what its name says. + FakeSource src; + src.samples = { + MetricSample{"test_seconds_total", kHelp, {}, 2'500'000'000ULL, MetricKind::kCounter, 1e-9}}; + MetricPublisher pub; + Collector c; + pub.Publish("src", {}, src, c.Sink()); + EXPECT_DOUBLE_EQ(c.Counter("test_seconds_total"), 2.5); +} + +TEST(MetricPublisher, PrependsSourceLabelsToEverySample) { + FakeSource src; + src.samples = {MetricSample{kName, kHelp, {{"op", "commit"}}, 4}}; + MetricPublisher pub; + Collector c; + pub.Publish("src", {{"tier", "HBM"}, {"backend", "Mock"}}, src, c.Sink()); + EXPECT_DOUBLE_EQ( + c.Counter(Identity(kName, {{"tier", "HBM"}, {"backend", "Mock"}, {"op", "commit"}})), 4.0); +} + +// --------------------------------------------------------------------------- +// InstrumentedBackend +// --------------------------------------------------------------------------- + +class InstrumentedBackendMetrics : public ::testing::Test { + protected: + void SetUp() override { + auto mock = std::make_unique(TierType::DRAM); + mock_ = mock.get(); + backend_ = MakeInstrumentedBackend(std::move(mock)); + ASSERT_TRUE(backend_->Init(nullptr)); + } + + // Publish with the same labels PoolClient uses, so the identities asserted + // below are byte-for-byte the ones a dashboard queries. + void Publish() { + collector_.counters.clear(); + collector_.gauges.clear(); + const MetricLabels labels = {{"tier", TierTypeName(backend_->Tier())}, + {"backend", backend_->Name()}}; + publisher_.Publish(std::string("backend:") + backend_->Name(), labels, *backend_, + collector_.Sink()); + } + + // mori_umbp_backend_ops_total{tier=...,backend=...,op=...,status=...} + double Ops(const char* op, const char* status) { + return collector_.Counter( + Identity(MORI_UMBP_METRIC_BACKEND_OPS_TOTAL, {{"tier", TierTypeName(backend_->Tier())}, + {"backend", backend_->Name()}, + {"op", op}, + {"status", status}})); + } + double Bytes(const char* op) { + return collector_.Counter(Identity( + MORI_UMBP_METRIC_BACKEND_BYTES_TOTAL, + {{"tier", TierTypeName(backend_->Tier())}, {"backend", backend_->Name()}, {"op", op}})); + } + double Batches(const char* op) { + return collector_.Counter(Identity( + MORI_UMBP_METRIC_BACKEND_BATCHES_TOTAL, + {{"tier", TierTypeName(backend_->Tier())}, {"backend", backend_->Name()}, {"op", op}})); + } + + // Allocate + commit one key of `size` bytes, the way a writer does. + void PutKey(const std::string& key, uint64_t size) { + auto alloc = backend_->BatchAllocate({AllocateRequest{key, size}}); + ASSERT_EQ(alloc.size(), 1u); + if (alloc[0].outcome != AllocateOutcome::kSuccessAllocated) return; + backend_->BatchCommit({CommitRequest{alloc[0].slot_id, key}}); + } + + MockBackend* mock_ = nullptr; + std::unique_ptr backend_; + MetricPublisher publisher_; + Collector collector_; +}; + +TEST_F(InstrumentedBackendMetrics, MetricsAreGenericOverBackends) { + // MockBackend implements no observability of its own. Everything asserted + // here was derived from the MediumBackend interface by the decorator, which + // is the guarantee: a new medium is instrumented by being composed in. + PutKey("k1", 4096); + PutKey("k2", 2048); + backend_->BatchResolve({"k1", "missing"}, /*include_descs=*/false); + backend_->Evict({"k2"}); + Publish(); + + EXPECT_DOUBLE_EQ(Ops("allocate", "ok"), 2.0); + EXPECT_DOUBLE_EQ(Ops("commit", "ok"), 2.0); + EXPECT_DOUBLE_EQ(Bytes("commit"), 4096.0 + 2048.0); + + EXPECT_DOUBLE_EQ(Ops("resolve", "ok"), 1.0); + EXPECT_DOUBLE_EQ(Ops("resolve", "miss"), 1.0); + EXPECT_DOUBLE_EQ(Bytes("resolve"), 4096.0); + + EXPECT_DOUBLE_EQ(Ops("evict", "ok"), 1.0); + EXPECT_DOUBLE_EQ(Bytes("evict"), 2048.0); +} + +TEST_F(InstrumentedBackendMetrics, LabelsCarryTheMediumSoOnePanelShowsThemAll) { + // The medium must appear as a LABEL VALUE and never inside a metric name — + // that is what lets one panel render every backend. + PutKey("k", 128); + Publish(); + + bool found = false; + for (const auto& [identity, value] : collector_.counters) { + (void)value; + EXPECT_EQ(identity.find("dram"), std::string::npos) + << "medium leaked into a metric identifier: " << identity; + EXPECT_EQ(identity.find("_ssd_"), std::string::npos) + << "medium leaked into a metric identifier: " << identity; + if (identity.find(",tier=DRAM") != std::string::npos) found = true; + } + EXPECT_TRUE(found) << "no series carried the tier label"; +} + +TEST_F(InstrumentedBackendMetrics, CountsEntriesAndBatchesSeparately) { + // Batch depth is a workload property worth seeing: one call, three keys. + backend_->BatchAllocate( + {AllocateRequest{"a", 16}, AllocateRequest{"b", 16}, AllocateRequest{"c", 16}}); + Publish(); + EXPECT_DOUBLE_EQ(Ops("allocate", "ok"), 3.0); + EXPECT_DOUBLE_EQ(Batches("allocate"), 1.0); +} + +TEST_F(InstrumentedBackendMetrics, SeparatesDedupFromAFreshAllocation) { + // "already exists" is success to the caller but moves no bytes; a capacity + // panel that lumps it in with a real allocation reads as write amplification + // that is not happening. + PutKey("dup", 64); + backend_->BatchAllocate({AllocateRequest{"dup", 64}}); + Publish(); + EXPECT_DOUBLE_EQ(Ops("allocate", "exists"), 1.0); + EXPECT_DOUBLE_EQ(Ops("allocate", "ok"), 1.0); +} + +TEST_F(InstrumentedBackendMetrics, ReportsFailedAllocations) { + backend_->BatchAllocate({AllocateRequest{"zero", 0}}); // MockBackend fails size 0 + Publish(); + EXPECT_DOUBLE_EQ(Ops("allocate", "failed"), 1.0); +} + +TEST_F(InstrumentedBackendMetrics, CommitAgainstAReapedSlotCountsAsFailed) { + auto alloc = backend_->BatchAllocate({AllocateRequest{"k", 32}}); + ASSERT_EQ(alloc.size(), 1u); + backend_->BatchAbort({alloc[0].slot_id}); + backend_->BatchCommit({CommitRequest{alloc[0].slot_id, "k"}}); + Publish(); + EXPECT_DOUBLE_EQ(Ops("commit", "failed"), 1.0); + EXPECT_DOUBLE_EQ(Ops("abort", "ok"), 1.0); +} + +TEST_F(InstrumentedBackendMetrics, EvictingAnAbsentKeyIsAMissNotAFailure) { + // Master retries protected keys on a later round; an eviction that freed + // nothing is not an error and must not read as one. + backend_->Evict({"never-existed"}); + Publish(); + EXPECT_DOUBLE_EQ(Ops("evict", "miss"), 1.0); + EXPECT_DOUBLE_EQ(Ops("evict", "ok"), -1.0) << "no successful eviction should have been reported"; +} + +TEST_F(InstrumentedBackendMetrics, RecordsTimeSpentInsideTheMedium) { + PutKey("k", 64); + Publish(); + const double seconds = collector_.Counter(Identity( + MORI_UMBP_METRIC_BACKEND_OP_SECONDS_TOTAL, + {{"tier", TierTypeName(backend_->Tier())}, {"backend", backend_->Name()}, {"op", "commit"}})); + EXPECT_GT(seconds, 0.0); + // Scaled to seconds, not left in nanoseconds: a commit of one small key on + // an in-memory mock cannot plausibly take a second. + EXPECT_LT(seconds, 1.0); +} + +TEST_F(InstrumentedBackendMetrics, ForwardsEveryCallToTheWrappedBackend) { + // The decorator is on the live data path; behaviour must be untouched. + PutKey("k", 100); + EXPECT_EQ(backend_->OwnedKeyCount(), 1u); + EXPECT_EQ(mock_->OwnedKeyCount(), 1u); + + auto resolved = backend_->BatchResolve({"k"}, false); + ASSERT_EQ(resolved.size(), 1u); + EXPECT_TRUE(resolved[0].found); + EXPECT_EQ(resolved[0].size, 100u); + + EXPECT_EQ(backend_->Tier(), mock_->Tier()); + EXPECT_STREQ(backend_->Name(), mock_->Name()); + EXPECT_EQ(backend_->PageSize(), mock_->PageSize()); + + auto events = backend_->DrainPendingEvents(); + EXPECT_FALSE(events.empty()) << "the heartbeat's outbox must survive the decorator"; +} + +TEST_F(InstrumentedBackendMetrics, PassesThroughAMediumsOwnSamples) { + // A medium may still publish what the interface cannot show. Those samples + // must arrive alongside the generic ones, under the same tier/backend labels. + class ChattyBackend : public MockBackend { + public: + ChattyBackend() : MockBackend(TierType::SSD) {} + std::vector SampleMetrics() const override { + return {MetricSample{MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL, + MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL_HELP, + {{"event", "device_read_ok"}}, + 9}}; + } + }; + + auto wrapped = MakeInstrumentedBackend(std::make_unique()); + MetricPublisher pub; + Collector c; + const MetricLabels labels = {{"tier", TierTypeName(wrapped->Tier())}, + {"backend", wrapped->Name()}}; + pub.Publish("backend:chatty", labels, *wrapped, c.Sink()); + + EXPECT_DOUBLE_EQ(c.Counter(Identity( + MORI_UMBP_METRIC_BACKEND_MEDIUM_EVENTS_TOTAL, + {{"tier", "SSD"}, {"backend", "MockBackend"}, {"event", "device_read_ok"}})), + 9.0); +} + +// --------------------------------------------------------------------------- +// Transfer layer +// --------------------------------------------------------------------------- + +TEST(TransferEngineMetrics, ChargesBytesAndPlansToTheEngineThatCarriedThem) { + CompositeTransferEngine composite; + composite.AddEngine(std::make_unique()); + + std::vector src(1024, 7); + std::vector dst(1024, 0); + + TransferItem item; + item.src = TransferRef::HostBytes(src.data(), src.size()); + item.dst = TransferRef::HostBytes(dst.data(), dst.size()); + item.size = src.size(); + item.tag = 0; + + std::vector failed; + ASSERT_TRUE(composite.Transfer({item}, &failed)); + ASSERT_TRUE(failed.empty()); + EXPECT_EQ(dst, src) << "instrumentation must not disturb the transfer itself"; + + MetricPublisher pub; + Collector c; + pub.Publish("transfer", {}, composite, c.Sink()); + + const MetricLabels local = {{"engine", "LocalCopyEngine"}, {"direction", "local"}}; + EXPECT_DOUBLE_EQ(c.Counter(Identity(MORI_UMBP_METRIC_TRANSFER_BYTES_TOTAL, local)), 1024.0); + MetricLabels ok = local; + ok.push_back({"status", "ok"}); + EXPECT_DOUBLE_EQ(c.Counter(Identity(MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL, ok)), 1.0); +} + +TEST(TransferEngineMetrics, CountsItemsNoEngineWouldTake) { + // A rejected item is a routing failure, not a transport one, so it must not + // be blamed on whichever engine happened to be registered. + CompositeTransferEngine composite; + composite.AddEngine(std::make_unique()); + + TransferItem item; // both endpoints invalid: no engine can claim this pair + item.size = 64; + item.tag = 0; + TransferPlanSet planned = composite.Plan({item}); + EXPECT_EQ(planned.rejected_tags.size(), 1u); + + MetricPublisher pub; + Collector c; + pub.Publish("transfer", {}, composite, c.Sink()); + EXPECT_DOUBLE_EQ( + c.Counter(Identity(MORI_UMBP_METRIC_TRANSFER_OPS_TOTAL, + {{"engine", "none"}, {"direction", "none"}, {"status", "rejected"}})), + 1.0); +} + +TEST(TransferEngineMetrics, AnEngineNeedsNoMetricsCodeToBeMeasured) { + // The composite is the measurement point precisely so an engine can be added + // with AddEngine() alone. An engine registered but never used reports + // nothing rather than a row of zeros, keeping the series set honest. + CompositeTransferEngine composite; + composite.AddEngine(std::make_unique()); + + MetricPublisher pub; + Collector c; + pub.Publish("transfer", {}, composite, c.Sink()); + EXPECT_TRUE(c.counters.empty()); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_cross_node_smoke.cpp b/tests/cpp/umbp/distributed/test_cross_node_smoke.cpp index 64c00fa35..adadca859 100644 --- a/tests/cpp/umbp/distributed/test_cross_node_smoke.cpp +++ b/tests/cpp/umbp/distributed/test_cross_node_smoke.cpp @@ -92,16 +92,10 @@ inline bool WaitForExists(PoolClient* client, const std::string& key, class CrossNodeSmoke : public ::testing::Test { protected: void SetUp() override { - buf_a_ = std::malloc(kBufSize); - buf_b_ = std::malloc(kBufSize); caller_buf_ = std::malloc(kBufSize); read_buf_ = std::malloc(kBufSize); - ASSERT_NE(buf_a_, nullptr); - ASSERT_NE(buf_b_, nullptr); ASSERT_NE(caller_buf_, nullptr); ASSERT_NE(read_buf_, nullptr); - std::memset(buf_a_, 0, kBufSize); - std::memset(buf_b_, 0, kBufSize); std::memset(caller_buf_, 0, kBufSize); std::memset(read_buf_, 0, kBufSize); @@ -130,8 +124,7 @@ class CrossNodeSmoke : public ::testing::Test { cfg_a.io_engine.host = "0.0.0.0"; cfg_a.io_engine.port = 0; cfg_a.peer_service_port = NextPeerServicePort(); - cfg_a.dram_buffers = {{buf_a_, kBlockSize}}; - cfg_a.tier_capacities = {{TierType::DRAM, {kBlockSize, kBlockSize}}}; + cfg_a.dram.buffer_sizes = {kBlockSize}; cfg_a.dram_page_size = kBlockSize; client_a_ = std::make_unique(std::move(cfg_a)); ASSERT_TRUE(client_a_->Init()); @@ -143,8 +136,7 @@ class CrossNodeSmoke : public ::testing::Test { cfg_b.io_engine.host = "0.0.0.0"; cfg_b.io_engine.port = 0; cfg_b.peer_service_port = NextPeerServicePort(); - cfg_b.dram_buffers = {{buf_b_, kBufSize}}; - cfg_b.tier_capacities = {{TierType::DRAM, {kBufSize, kBufSize}}}; + cfg_b.dram.buffer_sizes = {kBufSize}; cfg_b.dram_page_size = kBlockSize; client_b_ = std::make_unique(std::move(cfg_b)); ASSERT_TRUE(client_b_->Init()); @@ -158,14 +150,10 @@ class CrossNodeSmoke : public ::testing::Test { if (client_a_) client_a_->Shutdown(); if (master_) master_->Shutdown(); if (server_thread_.joinable()) server_thread_.join(); - std::free(buf_a_); - std::free(buf_b_); std::free(caller_buf_); std::free(read_buf_); } - void* buf_a_ = nullptr; - void* buf_b_ = nullptr; void* caller_buf_ = nullptr; void* read_buf_ = nullptr; std::unique_ptr master_; @@ -272,8 +260,7 @@ class CrossNodeMultiPage : public ::testing::Test { void TearDown() override { TearDownClients(); } - std::unique_ptr MakeClient(const std::string& node_id, const NodeSetup& setup, - std::vector* owned_bufs) { + std::unique_ptr MakeClient(const std::string& node_id, const NodeSetup& setup) { PoolClientConfig cfg; cfg.master_config.node_id = node_id; cfg.master_config.node_address = "127.0.0.1"; @@ -282,16 +269,7 @@ class CrossNodeMultiPage : public ::testing::Test { cfg.io_engine.port = 0; cfg.peer_service_port = NextPeerServicePort(); cfg.dram_page_size = kPageSize; - uint64_t total = 0; - for (size_t sz : setup.buffer_sizes) { - void* p = std::malloc(sz); - EXPECT_NE(p, nullptr); - std::memset(p, 0, sz); - owned_bufs->push_back(p); - cfg.dram_buffers.push_back({p, sz}); - total += sz; - } - cfg.tier_capacities = {{TierType::DRAM, {total, total}}}; + cfg.dram.buffer_sizes = setup.buffer_sizes; auto cli = std::make_unique(std::move(cfg)); EXPECT_TRUE(cli->Init()); return cli; @@ -305,18 +283,12 @@ class CrossNodeMultiPage : public ::testing::Test { if (master_) master_->Shutdown(); if (server_thread_.joinable()) server_thread_.join(); master_.reset(); - for (void* p : owned_a_) std::free(p); - for (void* p : owned_b_) std::free(p); - owned_a_.clear(); - owned_b_.clear(); } std::unique_ptr master_; std::thread server_thread_; std::unique_ptr client_a_; std::unique_ptr client_b_; - std::vector owned_a_; - std::vector owned_b_; }; TEST_F(CrossNodeMultiPage, MultiPageSameBufferPutGet) { @@ -324,8 +296,8 @@ TEST_F(CrossNodeMultiPage, MultiPageSameBufferPutGet) { // page block. node-a is sized so node-b is the obvious most-available // target; we then PUT 3 pages from a and GET them back from a. StartMaster(); - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}, &owned_a_); - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 4}}, &owned_b_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 4}}); constexpr size_t kPayload = kPageSize * 3; std::vector src(kPayload); @@ -347,8 +319,8 @@ TEST_F(CrossNodeMultiPage, CrossBufferScatterPutGet) { StartMaster(); // node-a is the source; sized so it cannot accept the Put itself // (single page) and routes go to node-b. - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}, &owned_a_); - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize, kPageSize}}, &owned_b_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize, kPageSize}}); constexpr size_t kPayload = kPageSize * 2; std::vector src(kPayload); @@ -375,10 +347,10 @@ TEST_F(CrossNodeMultiPage, CrossBufferScatterPutGet) { // every get is zero-copy (not staging). TEST_F(CrossNodeMultiPage, BatchGetZeroCopyMultiKeyCrossBuffer) { StartMaster(); - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}, &owned_a_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}); // Two 2-page buffers: the 4 single-page keys distribute across both, so // their pages span two distinct remote buffer_index values. - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 2, kPageSize * 2}}, &owned_b_); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 2, kPageSize * 2}}); constexpr size_t kN = 4; std::vector keys; @@ -416,8 +388,8 @@ TEST_F(CrossNodeMultiPage, BatchGetZeroCopyMultiKeyCrossBuffer) { // per-pair merge feeds via entry.failed. TEST_F(CrossNodeMultiPage, BatchGetZeroCopyFailureIsolation) { StartMaster(); - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}, &owned_a_); - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 4}}, &owned_b_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize}}); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 4}}); std::vector keys = {"iso-0", "iso-1", "iso-2"}; std::vector> srcs(3); @@ -458,8 +430,8 @@ TEST_F(CrossNodeMultiPage, PartialTailSinglePage) { // size < page_size: single allocation, partial tail = size. Covers the // N=1 fast path through both Put/Get and the scatter helper. StartMaster(); - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}, &owned_a_); - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize}}, &owned_b_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize}}); constexpr size_t kPayload = 1234; // arbitrary, < kPageSize std::vector src(kPayload); @@ -482,8 +454,8 @@ TEST_F(CrossNodeMultiPage, PartialTailSinglePage) { TEST_F(CrossNodeMultiPage, PartialTailMultiPageSameBuffer) { // 2 full pages + a partial tail in a single contiguous buffer (Strategy 1). StartMaster(); - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}, &owned_a_); - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 4}}, &owned_b_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize * 4}}); constexpr size_t kTail = 333; constexpr size_t kPayload = kPageSize * 2 + kTail; @@ -509,8 +481,8 @@ TEST_F(CrossNodeMultiPage, PartialTailCrossBufferScatter) { // Regression guard: scatter helpers must identify "last page" by spi // (global page index), not by the position inside a group. StartMaster(); - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}, &owned_a_); - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize, kPageSize}}, &owned_b_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize, kPageSize}}); constexpr size_t kTail = 777; constexpr size_t kPayload = kPageSize + kTail; @@ -535,8 +507,8 @@ TEST_F(CrossNodeMultiPage, PartialTailGetSizeMismatchRejected) { // (size < stored) or pull stale bytes from the unused tail of the last // page (size > stored, still inside the page window). StartMaster(); - client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}, &owned_a_); - client_b_ = MakeClient("node-b", NodeSetup{{kPageSize}}, &owned_b_); + client_a_ = MakeClient("node-a", NodeSetup{{kPageSize / 2}}); + client_b_ = MakeClient("node-b", NodeSetup{{kPageSize}}); constexpr size_t kPayload = 999; std::vector src(kPayload, 'X'); @@ -597,16 +569,7 @@ class CrossNodeOverlap : public ::testing::Test { cfg.peer_service_port = NextPeerServicePort(); cfg.dram_page_size = kPageSize; if (staging_buffer_size > 0) cfg.staging_buffer_size = staging_buffer_size; - uint64_t total = 0; - for (size_t sz : buffer_sizes) { - void* p = std::malloc(sz); - EXPECT_NE(p, nullptr); - std::memset(p, 0, sz); - owned_bufs_.push_back(p); - cfg.dram_buffers.push_back({p, sz}); - total += sz; - } - cfg.tier_capacities = {{TierType::DRAM, {total, total}}}; + cfg.dram.buffer_sizes = buffer_sizes; auto cli = std::make_unique(std::move(cfg)); EXPECT_TRUE(cli->Init()); clients_.push_back(std::move(cli)); @@ -621,14 +584,11 @@ class CrossNodeOverlap : public ::testing::Test { if (master_) master_->Shutdown(); if (server_thread_.joinable()) server_thread_.join(); master_.reset(); - for (void* p : owned_bufs_) std::free(p); - owned_bufs_.clear(); } std::unique_ptr master_; std::thread server_thread_; std::vector> clients_; - std::vector owned_bufs_; }; TEST_F(CrossNodeOverlap, MultiPeerZeroCopyByteExact) { @@ -888,10 +848,14 @@ TEST_F(CrossNodeOverlap, PutStagingMultiPeerByteExact) { VerifyReadback(caller, keys, srcs, kPageSize); } -TEST_F(CrossNodeOverlap, PutStagingOverflowFailsBatchCleanly) { - // 1-page staging buffer + un-registered srcs -> BuildRemotePutTransfers - // overflow: the whole peer batch fails, every allocated slot is aborted (keys - // never visible), and the client stays usable for a later zero-copy Put. +TEST_F(CrossNodeOverlap, PutStagingLargerThanPoolIsChunkedNotFailed) { + // REPLACES PutStagingOverflowFailsBatchCleanly, whose expectation Phase 6 + // deliberately inverted. Before, a batch needing more staging than the buffer + // held was failed wholesale by BuildRemotePutTransfers, because the client + // reserved the whole batch's staging up front under one lock. The bounce pool + // now belongs to the transfer engine, which chunks a staged group into + // pool-sized plans and runs each to completion inside Submit — so 4 pages + // through a 1-page pool is 4 sequential round trips, not an error. StartMaster(); constexpr size_t kN = 4; PoolClient* caller = MakeClient("node-a", {kPageSize}, /*staging_buffer_size=*/kPageSize); @@ -909,7 +873,35 @@ TEST_F(CrossNodeOverlap, PutStagingOverflowFailsBatchCleanly) { auto put = caller->BatchPut(keys, psrcs, sizes); ASSERT_EQ(put.size(), kN); - for (size_t k = 0; k < kN; ++k) EXPECT_FALSE(put[k]) << "overflow key should fail " << keys[k]; + for (size_t k = 0; k < kN; ++k) EXPECT_TRUE(put[k]) << "chunked staged put failed " << keys[k]; + for (const auto& key : keys) ASSERT_TRUE(WaitForExists(caller, key)); + // Chunking must not scramble which bytes went where. + VerifyReadback(caller, keys, srcs, kPageSize); +} + +TEST_F(CrossNodeOverlap, PutPageLargerThanStagingPoolFailsBatchCleanly) { + // The failure that IS still a failure: one PAGE bigger than the whole bounce + // pool cannot be chunked, so the planner rejects it. Every allocated slot is + // aborted (keys never visible) and the client stays usable for a later + // zero-copy Put. + StartMaster(); + constexpr size_t kN = 4; + PoolClient* caller = MakeClient("node-a", {kPageSize}, /*staging_buffer_size=*/kPageSize / 2); + MakeClient("node-b", {kPageSize * 64}); + + std::vector keys; + std::vector> srcs(kN); + std::vector psrcs(kN); + std::vector sizes(kN, kPageSize); + for (size_t k = 0; k < kN; ++k) { + keys.push_back("ptoobig-" + std::to_string(k)); + srcs[k].assign(kPageSize, static_cast(0x41 + k)); // un-registered -> staging + psrcs[k] = srcs[k].data(); + } + + auto put = caller->BatchPut(keys, psrcs, sizes); + ASSERT_EQ(put.size(), kN); + for (size_t k = 0; k < kN; ++k) EXPECT_FALSE(put[k]) << "oversized key should fail " << keys[k]; // Slots were aborted, never committed -> keys must not be visible. auto present = caller->BatchExists(keys); ASSERT_EQ(present.size(), kN); @@ -918,7 +910,7 @@ TEST_F(CrossNodeOverlap, PutStagingOverflowFailsBatchCleanly) { // Client is still usable: a registered (zero-copy) single Put succeeds. std::vector ok_src(kPageSize, 0x7E); ASSERT_TRUE(caller->RegisterMemory(ok_src.data(), ok_src.size())); - EXPECT_TRUE(caller->Put("povf-ok", ok_src.data(), kPageSize)); + EXPECT_TRUE(caller->Put("ptoobig-ok", ok_src.data(), kPageSize)); caller->DeregisterMemory(ok_src.data()); } diff --git a/tests/cpp/umbp/distributed/test_gds_engine.cpp b/tests/cpp/umbp/distributed/test_gds_engine.cpp new file mode 100644 index 000000000..63e56c938 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_gds_engine.cpp @@ -0,0 +1,155 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// GdsEngine: routing/planning are pure and run anywhere; the read round trip +// needs a GPU and the hipFile driver and skips without them. Point TMPDIR at +// an ext4/xfs mount (e.g. TMPDIR=/mnt/gds) so the O_DIRECT open succeeds and the +// fastpath is actually exercised; on other filesystems it skips or falls back. +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "umbp/distributed/transfer/gds_engine.h" + +namespace mori::umbp { +namespace { + +TransferRef GpuBuf(void* p, uint64_t n, int dev = 0) { + return TransferRef::HostBytes(p, n, mori::io::MemoryLocationType::GPU, dev); +} + +TransferItem MakeItem(const TransferRef& src, uint64_t src_off, const TransferRef& dst, + uint64_t dst_off, uint64_t size, size_t tag) { + TransferItem it; + it.src = src; + it.src_offset = src_off; + it.dst = dst; + it.dst_offset = dst_off; + it.size = size; + it.tag = tag; + return it; +} + +// --------------------------------------------------------------------------- +// Routing / planning — no GPU, no hipFile driver +// --------------------------------------------------------------------------- + +TEST(GdsEngine, ClaimsFileToGpuOnly) { + GdsEngine engine; + TransferRef file = + TransferRef::File(/*fd=*/7, /*offset=*/0, /*n=*/4096, reinterpret_cast(0x1)); + std::vector stub(4096); + TransferRef gpu = GpuBuf(stub.data(), stub.size()); + TransferRef host = TransferRef::HostBytes(stub.data(), stub.size()); // CPU + + EXPECT_TRUE(engine.CanHandle(file, gpu)); // the read path + EXPECT_FALSE(engine.CanHandle(file, host)); // host dst -> staged path + EXPECT_FALSE(engine.CanHandle(gpu, file)); // write path, a later change + EXPECT_FALSE(engine.CanHandle(gpu, gpu)); // no file endpoint at all +} + +TEST(GdsEngine, PlansOneRangePerItemAndBoundsToBuffer) { + GdsEngine engine; + TransferRef file = + TransferRef::File(/*fd=*/7, /*offset=*/8192, /*n=*/1 << 20, reinterpret_cast(0x1)); + std::vector stub(4096); + TransferRef gpu = GpuBuf(stub.data(), stub.size()); + + std::vector items{ + MakeItem(file, 0, gpu, 0, 2048, 0), MakeItem(file, 2048, gpu, 2048, 2048, 1), + MakeItem(file, 0, gpu, 0, 8192, 2), // overruns the 4096-byte buffer + }; + TransferPlanSet planned = engine.Plan(items); + EXPECT_EQ(std::vector({2}), planned.rejected_tags); + ASSERT_EQ(2u, planned.plans.size()); + for (const auto& p : planned.plans) EXPECT_EQ(1u, p.sizes.size()); +} + +// --------------------------------------------------------------------------- +// End to end — needs a GPU and the hipFile driver +// --------------------------------------------------------------------------- + +bool HaveGpu() { + int n = 0; + return hipGetDeviceCount(&n) == hipSuccess && n > 0; +} + +TEST(GdsEngine, ReadsFileRangeIntoDeviceMemory) { + if (!HaveGpu()) GTEST_SKIP() << "no HIP device"; + + const size_t kSize = 8192; // 4 KiB-aligned so the fastpath can accept it + const std::string path = std::string(::testing::TempDir()) + "/gds_engine_rt.bin"; + std::vector payload(kSize); + for (size_t i = 0; i < kSize; ++i) payload[i] = static_cast(i * 131 + 7); + + { + int wfd = ::open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + ASSERT_GE(wfd, 0); + ASSERT_EQ(static_cast(kSize), ::write(wfd, payload.data(), kSize)); + ::close(wfd); + } + + int fd = ::open(path.c_str(), O_RDONLY | O_DIRECT); + if (fd < 0) { + ::unlink(path.c_str()); + GTEST_SKIP() << "filesystem rejects O_DIRECT (errno=" << errno << "); point TMPDIR at ext4/xfs"; + } + + GdsEngine engine; + TransferRef file = engine.RegisterFile(fd, /*offset=*/0, kSize); + if (!file.IsFile()) { + ::close(fd); + ::unlink(path.c_str()); + GTEST_SKIP() << "hipFileHandleRegister unavailable on this host"; + } + + void* dbuf = nullptr; + ASSERT_EQ(hipSuccess, hipMalloc(&dbuf, kSize)); + ASSERT_EQ(hipSuccess, hipMemset(dbuf, 0, kSize)); + TransferRef gpu = GpuBuf(dbuf, kSize); + + std::vector failed; + ASSERT_TRUE(engine.Transfer({MakeItem(file, 0, gpu, 0, kSize, 0)}, &failed)); + EXPECT_TRUE(failed.empty()); + + std::vector got(kSize, 0); + ASSERT_EQ(hipSuccess, hipMemcpy(got.data(), dbuf, kSize, hipMemcpyDeviceToHost)); + EXPECT_EQ(payload, got); + + hipFree(dbuf); + engine.Deregister(file); + ::close(fd); + ::unlink(path.c_str()); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_hbm_backend.cpp b/tests/cpp/umbp/distributed/test_hbm_backend.cpp new file mode 100644 index 000000000..c3732d88f --- /dev/null +++ b/tests/cpp/umbp/distributed/test_hbm_backend.cpp @@ -0,0 +1,395 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// HbmCopyEngine + the HBM backend: the second medium, and the engine that +// exists because no engine in tree could serve its local pairs. +// +// The planner/selection half needs no GPU and always runs. The byte-moving +// half needs a real device and SKIPS (rather than fails) without one, so the +// suite stays runnable on a CPU-only box — matching how the integration label +// is used for tests that need a fabric. + +#include +#include + +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/hbm_backend.h" +#include "umbp/distributed/peer/backend/page_backend.h" +#include "umbp/distributed/transfer/composite_transfer_engine.h" +#include "umbp/distributed/transfer/hbm_copy_engine.h" +#include "umbp/distributed/transfer/local_copy_engine.h" + +namespace mori::umbp { +namespace { + +bool HaveGpu() { + int count = 0; + return hipGetDeviceCount(&count) == hipSuccess && count > 0; +} + +TransferItem MakeItem(const TransferRef& src, uint64_t src_off, const TransferRef& dst, + uint64_t dst_off, uint64_t size, size_t tag) { + TransferItem item; + item.src = src; + item.src_offset = src_off; + item.dst = dst; + item.dst_offset = dst_off; + item.size = size; + item.tag = tag; + return item; +} + +TransferRef HostRef(void* p, uint64_t n) { + return TransferRef::HostBytes(p, n, mori::io::MemoryLocationType::CPU, -1); +} +TransferRef GpuRef(void* p, uint64_t n, int device = 0) { + return TransferRef::HostBytes(p, n, mori::io::MemoryLocationType::GPU, device); +} + +// RAII device buffer so a failed assertion cannot leak VRAM across cases. +class DeviceBuffer { + public: + explicit DeviceBuffer(size_t bytes) { + if (hipMalloc(&ptr_, bytes) != hipSuccess) ptr_ = nullptr; + size_ = ptr_ != nullptr ? bytes : 0; + } + ~DeviceBuffer() { + if (ptr_ != nullptr) (void)hipFree(ptr_); + } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + void* get() const { return ptr_; } + size_t size() const { return size_; } + bool valid() const { return ptr_ != nullptr; } + + private: + void* ptr_ = nullptr; + size_t size_ = 0; +}; + +// --------------------------------------------------------------------------- +// Selection — the reason this engine exists +// --------------------------------------------------------------------------- + +// The gap this engine was written to close: before it, a both-local pair with a +// GPU endpoint was claimed by NO engine, so an HBM backend's local Put/Get +// could not complete. Asserted against the real composite, not by inspection. +TEST(HbmCopyEngine, ClaimsTheLocalGpuPairNoOtherEngineWould) { + int host = 0; + int fake_device = 0; // no GPU needed: selection never dereferences. + const TransferRef h = HostRef(&host, sizeof(host)); + const TransferRef g = GpuRef(&fake_device, sizeof(fake_device)); + + LocalCopyEngine local; + HbmCopyEngine hbm; + + // The precondition: the host-only engine refuses every pair with a GPU side. + EXPECT_FALSE(local.CanHandle(h, g)); + EXPECT_FALSE(local.CanHandle(g, h)); + EXPECT_FALSE(local.CanHandle(g, g)); + + // And this engine picks up exactly those three. + EXPECT_TRUE(hbm.CanHandle(h, g)); // H2D + EXPECT_TRUE(hbm.CanHandle(g, h)); // D2H + EXPECT_TRUE(hbm.CanHandle(g, g)); // D2D +} + +// Disjoint, not merely ordered: the two local engines must never both claim a +// pair, or composite registration order would silently become a performance +// decision. +TEST(HbmCopyEngine, DoesNotOverlapLocalCopyEngine) { + int a = 0, b = 0; + const TransferRef h1 = HostRef(&a, sizeof(a)); + const TransferRef h2 = HostRef(&b, sizeof(b)); + + LocalCopyEngine local; + HbmCopyEngine hbm; + + EXPECT_TRUE(local.CanHandle(h1, h2)); + EXPECT_FALSE(hbm.CanHandle(h1, h2)); // both-CPU stays with the NT-AVX2 path +} + +TEST(HbmCopyEngine, CompositeRoutesGpuPairsHere) { + int host = 0; + int fake_device = 0; + CompositeTransferEngine composite; + composite.AddEngine(std::make_unique()); + composite.AddEngine(std::make_unique()); + + const TransferRef h = HostRef(&host, sizeof(host)); + const TransferRef g = GpuRef(&fake_device, sizeof(fake_device)); + + TransferEngine* for_host = composite.SelectEngine(h, h); + TransferEngine* for_gpu = composite.SelectEngine(h, g); + ASSERT_NE(for_host, nullptr); + ASSERT_NE(for_gpu, nullptr); + EXPECT_STREQ(for_host->Name(), "LocalCopyEngine"); + EXPECT_STREQ(for_gpu->Name(), "HbmCopyEngine"); +} + +// --------------------------------------------------------------------------- +// Planner +// --------------------------------------------------------------------------- + +TEST(HbmCopyEngine, CoalescesAdjacentSegments) { + int host = 0, dev = 0; + HbmCopyEngine engine; + const TransferRef h = HostRef(&host, 4096); + const TransferRef g = GpuRef(&dev, 4096); + + // Three adjacent pages on both sides collapse into one hipMemcpy — the win + // that matters more here than for memcpy, since each call has launch cost. + const auto set = engine.Plan({MakeItem(h, 0, g, 0, 1024, 0), MakeItem(h, 1024, g, 1024, 1024, 1), + MakeItem(h, 2048, g, 2048, 1024, 2)}); + + ASSERT_EQ(set.plans.size(), 1u); + EXPECT_TRUE(set.rejected_tags.empty()); + ASSERT_EQ(set.plans[0].sizes.size(), 1u); + EXPECT_EQ(set.plans[0].sizes[0], 3072u); + EXPECT_EQ(set.plans[0].tags.size(), 3u); +} + +TEST(HbmCopyEngine, RejectsOutOfBoundsRatherThanCorruptingThePool) { + int host = 0, dev = 0; + HbmCopyEngine engine; + const TransferRef h = HostRef(&host, 1024); + const TransferRef g = GpuRef(&dev, 1024); + + const auto set = engine.Plan({MakeItem(h, 0, g, 512, 1024, 7)}); + EXPECT_TRUE(set.plans.empty()); + ASSERT_EQ(set.rejected_tags.size(), 1u); + EXPECT_EQ(set.rejected_tags[0], 7u); +} + +TEST(HbmCopyEngine, RejectsPairsItCannotCarry) { + int a = 0, b = 0; + HbmCopyEngine engine; + const auto set = engine.Plan({MakeItem(HostRef(&a, 64), 0, HostRef(&b, 64), 0, 64, 3)}); + EXPECT_TRUE(set.plans.empty()); + ASSERT_EQ(set.rejected_tags.size(), 1u); + EXPECT_EQ(set.rejected_tags[0], 3u); +} + +// --------------------------------------------------------------------------- +// Real bytes (needs a GPU) +// --------------------------------------------------------------------------- + +TEST(HbmCopyEngine, RoundTripsHostToDeviceAndBack) { + if (!HaveGpu()) GTEST_SKIP() << "no GPU visible"; + + constexpr size_t kBytes = 256 * 1024; + DeviceBuffer device(kBytes); + ASSERT_TRUE(device.valid()); + + std::vector src(kBytes), dst(kBytes, 0); + std::iota(src.begin(), src.end(), 1); + + HbmCopyEngine engine; + const TransferRef s = HostRef(src.data(), kBytes); + const TransferRef g = GpuRef(device.get(), kBytes); + const TransferRef d = HostRef(dst.data(), kBytes); + + std::vector failed; + ASSERT_TRUE(engine.Transfer({MakeItem(s, 0, g, 0, kBytes, 0)}, &failed)) << "H2D failed"; + EXPECT_TRUE(failed.empty()); + + ASSERT_TRUE(engine.Transfer({MakeItem(g, 0, d, 0, kBytes, 0)}, &failed)) << "D2H failed"; + EXPECT_TRUE(failed.empty()); + + EXPECT_EQ(std::memcmp(src.data(), dst.data(), kBytes), 0); +} + +TEST(HbmCopyEngine, RoundTripsScatteredPagesThroughOneDeviceBuffer) { + if (!HaveGpu()) GTEST_SKIP() << "no GPU visible"; + + constexpr size_t kPage = 4096; + constexpr size_t kPages = 4; + DeviceBuffer device(kPage * kPages); + ASSERT_TRUE(device.valid()); + + std::vector src(kPage * kPages), dst(kPage * kPages, 0); + std::iota(src.begin(), src.end(), 7); + + HbmCopyEngine engine; + const TransferRef s = HostRef(src.data(), src.size()); + const TransferRef g = GpuRef(device.get(), kPage * kPages); + const TransferRef d = HostRef(dst.data(), dst.size()); + + // Deliberately non-adjacent order so the planner cannot coalesce everything. + // The stride must be coprime with kPages or this stops being a permutation + // and silently leaves pages uncopied (which is how this test first failed). + static_assert(kPages == 4, "stride 3 below is chosen coprime with kPages"); + std::vector up; + for (size_t i = 0; i < kPages; ++i) { + const size_t page = (i * 3) % kPages; // {0, 3, 2, 1} + up.push_back(MakeItem(s, page * kPage, g, page * kPage, kPage, i)); + } + std::vector failed; + ASSERT_TRUE(engine.Transfer(up, &failed)); + + std::vector down; + for (size_t i = 0; i < kPages; ++i) { + down.push_back(MakeItem(g, i * kPage, d, i * kPage, kPage, i)); + } + ASSERT_TRUE(engine.Transfer(down, &failed)); + + EXPECT_EQ(std::memcmp(src.data(), dst.data(), src.size()), 0); +} + +TEST(HbmCopyEngine, CopiesDeviceToDevice) { + if (!HaveGpu()) GTEST_SKIP() << "no GPU visible"; + + constexpr size_t kBytes = 64 * 1024; + DeviceBuffer a(kBytes), b(kBytes); + ASSERT_TRUE(a.valid()); + ASSERT_TRUE(b.valid()); + + std::vector src(kBytes), dst(kBytes, 0); + std::iota(src.begin(), src.end(), 3); + ASSERT_EQ(hipMemcpy(a.get(), src.data(), kBytes, hipMemcpyHostToDevice), hipSuccess); + + HbmCopyEngine engine; + std::vector failed; + ASSERT_TRUE(engine.Transfer( + {MakeItem(GpuRef(a.get(), kBytes), 0, GpuRef(b.get(), kBytes), 0, kBytes, 0)}, &failed)); + + ASSERT_EQ(hipMemcpy(dst.data(), b.get(), kBytes, hipMemcpyDeviceToHost), hipSuccess); + EXPECT_EQ(std::memcmp(src.data(), dst.data(), kBytes), 0); +} + +// --------------------------------------------------------------------------- +// The backend +// --------------------------------------------------------------------------- + +// A registrar that only hands back process-local refs — enough to exercise the +// backend's ownership path without an IOEngine, and it is what +// CompositeTransferEngine degrades to on a node with no RDMA configured. +class LocalOnlyRegistrar final : public MemoryRegistrar { + public: + TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) override { + ++registrations; + last_loc = loc; + last_device = device; + return TransferRef::HostBytes(base, size, loc, device); + } + void Deregister(const TransferRef&) override { ++deregistrations; } + + int registrations = 0; + int deregistrations = 0; + mori::io::MemoryLocationType last_loc = mori::io::MemoryLocationType::CPU; + int last_device = -1; +}; + +// The registration contract: the backend supplies the facts a descriptor cannot +// recover. Getting this wrong is invisible until a transfer picks the wrong +// engine, which is exactly what this asserts against. +TEST(HbmBackend, RegistersItsPoolAsGpuMemoryOnItsDevice) { + if (!HaveGpu()) GTEST_SKIP() << "no GPU visible"; + + constexpr uint64_t kPageSize = 64 * 1024; + auto backend = MakeHbmBackend(kPageSize, /*device=*/0, {kPageSize * 8}, + std::chrono::milliseconds{30000}, std::chrono::milliseconds{500}); + LocalOnlyRegistrar registrar; + ASSERT_TRUE(backend->Init(®istrar)); + + EXPECT_EQ(backend->Tier(), TierType::HBM); + EXPECT_EQ(registrar.registrations, 1); + EXPECT_EQ(registrar.last_loc, mori::io::MemoryLocationType::GPU); + EXPECT_EQ(registrar.last_device, 0); + + // And the endpoint it publishes says so too, which is what routes a local + // transfer to HbmCopyEngine instead of LocalCopyEngine. + ASSERT_EQ(backend->BufferCount(), 1u); + const TransferRef ref = backend->BufferRef(0); + EXPECT_TRUE(ref.Valid()); + EXPECT_EQ(ref.loc, mori::io::MemoryLocationType::GPU); + EXPECT_EQ(ref.device, 0); + + backend->Shutdown(); + EXPECT_EQ(registrar.deregistrations, 1); +} + +// The point of the whole exercise: a Put and a Get against HBM, moving real +// bytes, with the composite choosing the engine — no tier branch anywhere. +TEST(HbmBackend, PutsAndGetsThroughTheCompositeEngine) { + if (!HaveGpu()) GTEST_SKIP() << "no GPU visible"; + + constexpr uint64_t kPageSize = 64 * 1024; + auto backend = MakeHbmBackend(kPageSize, /*device=*/0, {kPageSize * 4}, + std::chrono::milliseconds{30000}, std::chrono::milliseconds{500}); + LocalOnlyRegistrar registrar; + ASSERT_TRUE(backend->Init(®istrar)); + + CompositeTransferEngine composite; + composite.AddEngine(std::make_unique()); + composite.AddEngine(std::make_unique()); + + // --- Put --- + auto allocated = backend->BatchAllocate({AllocateRequest{"key-a", kPageSize}}); + ASSERT_EQ(allocated.size(), 1u); + ASSERT_EQ(allocated[0].outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_EQ(allocated[0].pages.size(), 1u); + + std::vector payload(kPageSize); + std::iota(payload.begin(), payload.end(), 11); + + const TransferRef pool = backend->BufferRef(allocated[0].pages[0].buffer_index); + const uint64_t page_off = static_cast(allocated[0].pages[0].page_index) * kPageSize; + + std::vector failed; + ASSERT_TRUE(composite.Transfer( + {MakeItem(HostRef(payload.data(), kPageSize), 0, pool, page_off, kPageSize, 0)}, &failed)) + << "host -> HBM put failed"; + + auto committed = backend->BatchCommit({CommitRequest{allocated[0].slot_id, "key-a"}}); + ASSERT_EQ(committed.size(), 1u); + EXPECT_TRUE(committed[0].success); + + // --- Get --- + auto resolved = backend->BatchResolve({"key-a"}, /*include_descs=*/false); + ASSERT_EQ(resolved.size(), 1u); + ASSERT_TRUE(resolved[0].found); + ASSERT_EQ(resolved[0].pages.size(), 1u); + EXPECT_EQ(resolved[0].size, kPageSize); + + std::vector readback(kPageSize, 0); + const TransferRef read_pool = backend->BufferRef(resolved[0].pages[0].buffer_index); + const uint64_t read_off = static_cast(resolved[0].pages[0].page_index) * kPageSize; + ASSERT_TRUE(composite.Transfer( + {MakeItem(read_pool, read_off, HostRef(readback.data(), kPageSize), 0, kPageSize, 0)}, + &failed)) + << "HBM -> host get failed"; + + EXPECT_EQ(std::memcmp(payload.data(), readback.data(), kPageSize), 0); + + backend->Shutdown(); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_in_memory_master_metadata_store.cpp b/tests/cpp/umbp/distributed/test_in_memory_master_metadata_store.cpp index c5b47a67b..4d6206e65 100644 --- a/tests/cpp/umbp/distributed/test_in_memory_master_metadata_store.cpp +++ b/tests/cpp/umbp/distributed/test_in_memory_master_metadata_store.cpp @@ -213,6 +213,27 @@ TEST(InMemoryStore, HeartbeatDeltaAddRemove) { EXPECT_TRUE(store.LookupBlock("k1").empty()); } +TEST(InMemoryStore, NamedLogicalTiersAreDistinctLocations) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + auto hot = Add("k1", TierType::DRAM, 10); + hot.logical_tier = "hot"; + auto warm = Add("k1", TierType::DRAM, 10); + warm.logical_tier = "warm"; + ASSERT_EQ(Beat(store, "n1", 1, {hot, warm}, kT0).status, HeartbeatResult::APPLIED); + + auto locations = store.LookupBlock("k1"); + ASSERT_EQ(locations.size(), 2u); + EXPECT_NE(locations[0].logical_tier, locations[1].logical_tier); + + auto remove_hot = Remove("k1", TierType::DRAM); + remove_hot.logical_tier = "hot"; + ASSERT_EQ(Beat(store, "n1", 2, {remove_hot}, kT0).status, HeartbeatResult::APPLIED); + locations = store.LookupBlock("k1"); + ASSERT_EQ(locations.size(), 1u); + EXPECT_EQ(locations.front().logical_tier, "warm"); +} + TEST(InMemoryStore, HeartbeatFullSyncReplaces) { InMemoryMasterMetadataStore store; RegisterAlive(store, "n1"); diff --git a/tests/cpp/umbp/distributed/test_local_first.cpp b/tests/cpp/umbp/distributed/test_local_first.cpp new file mode 100644 index 000000000..0ede9ca0d --- /dev/null +++ b/tests/cpp/umbp/distributed/test_local_first.cpp @@ -0,0 +1,218 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// PoolClientConfig::local_first — resolve on this node's own media before +// asking the master, and skip the master entirely when the whole batch is here. +// +// The claim under test is not "fewer RPCs" but "NO RPC", so the assertions do +// not count calls: they STOP THE MASTER and then read. A read that still +// succeeds against a dead master provably never needed it, which is exactly the +// single-node case the flag exists for. The `local_first=false` twin is what +// keeps that honest -- it must fail on the same data, or the test is passing +// for some reason other than the flag. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/config.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/pool_client.h" + +namespace mori::umbp { +namespace { + +constexpr size_t kPageSize = 4096; +constexpr size_t kDramCap = 16ULL * 1024 * 1024; +constexpr size_t kObjectSize = 8192; + +uint16_t NextPeerServicePort() { + static std::atomic next{ + static_cast(51000 + (static_cast(::getpid()) % 3000))}; + return next.fetch_add(1); +} + +std::vector Pattern(size_t size, int seed) { + std::vector out(size); + for (size_t i = 0; i < size; ++i) out[i] = static_cast((i * 31 + seed) & 0xFF); + return out; +} + +class LocalFirstTest : public ::testing::Test { + protected: + void SetUp() override { + MasterServerConfig master_cfg; + master_cfg.listen_address = "0.0.0.0:0"; + master_ = std::make_unique(std::move(master_cfg)); + server_thread_ = std::thread([this] { master_->Run(); }); + for (int i = 0; i < 50 && master_->GetBoundPort() == 0; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_NE(master_->GetBoundPort(), 0) << "Master failed to start"; + master_addr_ = "localhost:" + std::to_string(master_->GetBoundPort()); + } + + void TearDown() override { + for (auto& c : clients_) { + if (c) c->Shutdown(); + } + clients_.clear(); + StopMaster(); + } + + // Idempotent: a test that stops the master mid-body must not have TearDown + // stop it again. + void StopMaster() { + if (master_) master_->Shutdown(); + if (server_thread_.joinable()) server_thread_.join(); + master_.reset(); + } + + PoolClientConfig DramConfig(const std::string& node_id, bool local_first) { + PoolClientConfig cfg; + cfg.master_config.node_id = node_id; + cfg.master_config.node_address = "127.0.0.1"; + cfg.master_config.master_address = master_addr_; + cfg.io_engine.host = "0.0.0.0"; + cfg.io_engine.port = 0; + cfg.peer_service_port = NextPeerServicePort(); + cfg.dram_page_size = kPageSize; + cfg.medium = TierType::DRAM; + cfg.dram.buffer_sizes = {kDramCap}; + cfg.cache_remote_fetches = false; + cfg.local_first = local_first; + return cfg; + } + + PoolClient* Start(PoolClientConfig cfg) { + auto client = std::make_unique(std::move(cfg)); + if (!client->Init()) return nullptr; + clients_.push_back(std::move(client)); + return clients_.back().get(); + } + + std::string master_addr_; + std::unique_ptr master_; + std::thread server_thread_; + std::vector> clients_; +}; + +// The single-node case the flag exists for: everything this process put is +// still readable once the master is gone. +TEST_F(LocalFirstTest, FullyLocalBatchGetSurvivesAMasterOutage) { + auto* node = Start(DramConfig("local-first-get", /*local_first=*/true)); + ASSERT_NE(node, nullptr); + + const std::vector keys = {"lf-a", "lf-b", "lf-c"}; + std::vector> objects; + std::vector srcs; + std::vector sizes; + for (size_t i = 0; i < keys.size(); ++i) { + objects.push_back(Pattern(kObjectSize, static_cast(i))); + srcs.push_back(objects.back().data()); + sizes.push_back(kObjectSize); + } + ASSERT_EQ(node->BatchPut(keys, srcs, sizes), std::vector(keys.size(), true)); + + StopMaster(); + + std::vector> reads(keys.size(), std::vector(kObjectSize, 0)); + std::vector dsts; + for (auto& r : reads) dsts.push_back(r.data()); + + EXPECT_EQ(node->BatchGet(keys, dsts, sizes), std::vector(keys.size(), true)); + for (size_t i = 0; i < keys.size(); ++i) { + EXPECT_EQ(std::memcmp(reads[i].data(), objects[i].data(), kObjectSize), 0) + << "wrong bytes for " << keys[i]; + } +} + +// Same data, same outage, flag off. If this passed, the test above would prove +// nothing about local_first. +TEST_F(LocalFirstTest, RouteFirstBatchGetDoesNotSurviveAMasterOutage) { + auto* node = Start(DramConfig("route-first-get", /*local_first=*/false)); + ASSERT_NE(node, nullptr); + + const std::vector keys = {"rf-a"}; + const auto object = Pattern(kObjectSize, 7); + const std::vector srcs = {object.data()}; + const std::vector sizes = {kObjectSize}; + ASSERT_EQ(node->BatchPut(keys, srcs, sizes), std::vector({true})); + + StopMaster(); + + std::vector read(kObjectSize, 0); + const std::vector dsts = {read.data()}; + EXPECT_EQ(node->BatchGet(keys, dsts, sizes), std::vector({false})) + << "route-first must go through the master, so a dead master is a miss"; +} + +TEST_F(LocalFirstTest, LocallyHeldExistsSurvivesAMasterOutage) { + auto* node = Start(DramConfig("local-first-exists", /*local_first=*/true)); + ASSERT_NE(node, nullptr); + + const std::vector keys = {"lfe-a", "lfe-b"}; + const auto object = Pattern(kObjectSize, 3); + const std::vector srcs = {object.data(), object.data()}; + const std::vector sizes = {kObjectSize, kObjectSize}; + ASSERT_EQ(node->BatchPut(keys, srcs, sizes), std::vector(keys.size(), true)); + + StopMaster(); + + EXPECT_EQ(node->BatchExists(keys), std::vector(keys.size(), true)); + EXPECT_TRUE(node->Exists(keys[0])); +} + +// A local MISS is not conclusive, so it must still be taken to the master -- +// and when the master cannot answer, the honest result is false rather than a +// fabricated hit. This is the half of the contract that keeps local-first from +// turning "I do not have it" into "nobody has it". +TEST_F(LocalFirstTest, UnknownKeyIsNotClaimedToExist) { + auto* node = Start(DramConfig("local-first-miss", /*local_first=*/true)); + ASSERT_NE(node, nullptr); + + const std::vector present = {"lfm-here"}; + const auto object = Pattern(kObjectSize, 11); + ASSERT_EQ(node->BatchPut(present, {object.data()}, {kObjectSize}), std::vector({true})); + + // Master alive: a key nobody ever put is absent, and the local hit beside it + // is unaffected by the lookup that answers for the miss. + const std::vector mixed = {"lfm-here", "lfm-never-put"}; + EXPECT_EQ(node->BatchExists(mixed), std::vector({true, false})); + + std::vector read(kObjectSize, 0); + std::vector unused(kObjectSize, 0); + const std::vector dsts = {read.data(), unused.data()}; + EXPECT_EQ(node->BatchGet(mixed, dsts, {kObjectSize, kObjectSize}), + std::vector({true, false})); + EXPECT_EQ(std::memcmp(read.data(), object.data(), kObjectSize), 0); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_master_client_flush_heartbeat.cpp b/tests/cpp/umbp/distributed/test_master_client_flush_heartbeat.cpp index 8997b4d81..0a49c3959 100644 --- a/tests/cpp/umbp/distributed/test_master_client_flush_heartbeat.cpp +++ b/tests/cpp/umbp/distributed/test_master_client_flush_heartbeat.cpp @@ -39,29 +39,20 @@ #include #include #include -#include -#include #include #include #include #include #include +#include #include "umbp.grpc.pb.h" #include "umbp/distributed/master/master_client.h" +#include "umbp/distributed/peer/backend/mock_backend.h" namespace mori::umbp { namespace { -static uint16_t AllocPort() { - static std::atomic next{0}; - if (next.load() == 0) { - std::srand(static_cast(std::time(nullptr))); - next.store(static_cast(55500 + (std::rand() % 1500))); - } - return next.fetch_add(7); -} - // -------------------------------------------------------------------------- // Fake master: records heartbeat arrival times; optionally blocks each // Heartbeat RPC until ReleaseAll() is called. @@ -77,12 +68,21 @@ class RecordingMasterService final : public ::umbp::UMBPMaster::Service { return grpc::Status::OK; } - grpc::Status Heartbeat(grpc::ServerContext* ctx, const ::umbp::HeartbeatRequest*, + grpc::Status Heartbeat(grpc::ServerContext* ctx, const ::umbp::HeartbeatRequest* req, ::umbp::HeartbeatResponse* resp) override { { std::lock_guard lock(mu_); ++count_; last_time_ = std::chrono::steady_clock::now(); + size_t event_count = 0; + uint64_t highest_seq = 0; + for (const auto& bundle : req->bundles()) { + event_count += static_cast(bundle.events_size()); + highest_seq = std::max(highest_seq, bundle.seq()); + } + heartbeat_event_counts_.push_back(event_count); + resp->set_acked_seq(highest_seq); + resp->set_status(::umbp::CLIENT_STATUS_ALIVE); entered_cv_.notify_all(); } if (block_heartbeats_) { @@ -95,7 +95,6 @@ class RecordingMasterService final : public ::umbp::UMBPMaster::Service { [&] { return released_ || ctx->IsCancelled(); }); } } - resp->set_acked_seq(0); return grpc::Status::OK; } @@ -133,6 +132,11 @@ class RecordingMasterService final : public ::umbp::UMBPMaster::Service { return last_time_; } + std::vector HeartbeatEventCounts() { + std::lock_guard lock(mu_); + return heartbeat_event_counts_; + } + private: const int interval_ms_; const bool block_heartbeats_; @@ -141,10 +145,43 @@ class RecordingMasterService final : public ::umbp::UMBPMaster::Service { std::condition_variable entered_cv_; std::condition_variable release_cv_; int count_ = 0; + std::vector heartbeat_event_counts_; std::chrono::steady_clock::time_point last_time_; bool released_ = false; }; +// Upstream drove this through a bespoke OwnedLocationSource holding a vector of +// events. That type is gone — MediumBackend absorbed it — so the backlog is +// built the way a real one arises: commit `count` keys into a backend and let +// its outbox fill. MockBackend is a full MediumBackend, so this exercises the +// same DrainAllBackends path production uses. +// +// The registry must be attached to the client only AFTER the commits, because +// SetBackendRegistry installs the auto-flush hook: with it installed, the +// commits themselves would flush and the backlog would never accumulate. +std::unique_ptr MakeRegistryWithBacklog(size_t count) { + auto backend = std::make_unique(TierType::DRAM); + + std::vector allocs; + allocs.reserve(count); + for (size_t i = 0; i < count; ++i) { + allocs.push_back(AllocateRequest{"heartbeat-key-" + std::to_string(i), /*size=*/4096}); + } + auto allocated = backend->BatchAllocate(allocs); + + std::vector commits; + commits.reserve(count); + for (size_t i = 0; i < count; ++i) { + EXPECT_EQ(allocated[i].outcome, AllocateOutcome::kSuccessAllocated); + commits.push_back(CommitRequest{allocated[i].slot_id, allocs[i].key}); + } + backend->BatchCommit(commits); + + auto registry = std::make_unique(); + registry->Register(std::move(backend)); + return registry; +} + // -------------------------------------------------------------------------- // Test fixture // -------------------------------------------------------------------------- @@ -152,14 +189,20 @@ class FlushHeartbeatTest : public ::testing::Test { protected: void BuildServer(int interval_ms, bool block = false) { service_ = std::make_unique(interval_ms, block); - uint16_t port = AllocPort(); - address_ = "127.0.0.1:" + std::to_string(port); + // Port 0 asks the OS for a free one and reports it back. A guessed port + // collides on a busy host — these tests run in a --network host container + // alongside everything else on the node — and BuildAndStart signals that + // by returning nullptr, so the collision surfaces as an unrelated + // assertion failure rather than as EADDRINUSE. + int selected_port = 0; grpc::ServerBuilder builder; - builder.AddListeningPort(address_, grpc::InsecureServerCredentials()); + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &selected_port); builder.RegisterService(service_.get()); server_ = builder.BuildAndStart(); ASSERT_NE(server_, nullptr); + ASSERT_NE(selected_port, 0); + address_ = "127.0.0.1:" + std::to_string(selected_port); } std::unique_ptr MakeRegisteredClient() { @@ -282,5 +325,27 @@ TEST_F(FlushHeartbeatTest, FlushWhileRPCInFlightFiresNextTickImmediately) { << " ms); flush_requested_ may not have been preserved across the in-flight RPC"; } +TEST_F(FlushHeartbeatTest, LargeEventBacklogIsSplitAcrossBoundedImmediateHeartbeats) { + constexpr size_t kEvents = 20'000; + constexpr size_t kDefaultMaxEventsPerRpc = 16 * 1024; + ASSERT_NO_FATAL_FAILURE(BuildServer(/*interval_ms=*/10'000)); + // Declared before the client so it outlives it: MasterClient holds the + // registry by raw pointer and touches it while shutting the heartbeat down. + auto registry = MakeRegistryWithBacklog(kEvents); + auto client = MakeRegisteredClient(); + client->SetBackendRegistry(registry.get()); + client->StartHeartbeat(); + + client->FlushHeartbeat(); + ASSERT_TRUE(service_->WaitForCount(2, std::chrono::milliseconds(2000))); + client->StopHeartbeat(); + + const auto counts = service_->HeartbeatEventCounts(); + ASSERT_GE(counts.size(), 2u); + EXPECT_LE(counts[0], kDefaultMaxEventsPerRpc); + EXPECT_LE(counts[1], kDefaultMaxEventsPerRpc); + EXPECT_EQ(counts[0] + counts[1], kEvents); +} + } // namespace } // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_master_client_lifecycle.cpp b/tests/cpp/umbp/distributed/test_master_client_lifecycle.cpp index c6862f3fa..be1c2367a 100644 --- a/tests/cpp/umbp/distributed/test_master_client_lifecycle.cpp +++ b/tests/cpp/umbp/distributed/test_master_client_lifecycle.cpp @@ -33,8 +33,6 @@ #include #include #include -#include -#include #include #include #include @@ -50,15 +48,6 @@ namespace { constexpr int kHeartbeatIntervalMs = 50; // fast loop so heartbeat enters quickly constexpr int kDestructorBudgetMs = 7000; -static uint16_t AllocPort() { - static std::atomic next{0}; - if (next.load() == 0) { - std::srand(static_cast(std::time(nullptr))); - next.store(static_cast(54000 + (std::rand() % 1500))); - } - return next.fetch_add(7); -} - // Fake master service: RegisterClient returns immediately so the client // can proceed to start its heartbeat loop. Heartbeat and UnregisterClient // park on a CV until ReleaseAll() is called from TearDown — this both @@ -121,15 +110,20 @@ class BlackholeMasterService final : public ::umbp::UMBPMaster::Service { class MasterClientLifecycleTest : public ::testing::Test { protected: void SetUp() override { - port_ = AllocPort(); - address_ = "127.0.0.1:" + std::to_string(port_); service_ = std::make_unique(kHeartbeatIntervalMs); + // Port 0 asks the OS for a free one and reports it back; a guessed port + // collides on a busy host and BuildAndStart signals that only by + // returning nullptr. + int selected_port = 0; grpc::ServerBuilder builder; - builder.AddListeningPort(address_, grpc::InsecureServerCredentials()); + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &selected_port); builder.RegisterService(service_.get()); server_ = builder.BuildAndStart(); ASSERT_NE(server_, nullptr); + ASSERT_NE(selected_port, 0); + port_ = static_cast(selected_port); + address_ = "127.0.0.1:" + std::to_string(port_); } void TearDown() override { diff --git a/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp b/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp index 9fed02443..a5bc7ac45 100644 --- a/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp +++ b/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp @@ -31,6 +31,7 @@ #include #include +#include "umbp/common/grpc_limits.h" #include "umbp/distributed/master/evict_strategy.h" #include "umbp/distributed/types.h" @@ -100,5 +101,63 @@ TEST(LruMasterEvictStrategy, SkipsTiersWithNoBudget) { EXPECT_TRUE(victims.empty()); } +// --------------------------------------------------------------------------- +// gRPC batch chunking (grpc_limits.h) +// +// The eviction round that exposed this selected 155,157 victims -- 18.7 MB of +// `repeated string` against a 4 MiB receive limit -- and protobuf refused the +// whole message, so 81 consecutive rounds freed nothing and the tier could not +// un-fill itself. These pin the split, not the transport. +// --------------------------------------------------------------------------- + +namespace { +std::vector MakeKeys(size_t n, size_t key_len) { + std::vector keys; + keys.reserve(n); + for (size_t i = 0; i < n; ++i) keys.emplace_back(key_len, 'k'); + return keys; +} +} // namespace + +TEST(GrpcBatchChunkingTest, SmallListGoesInOneChunk) { + const auto keys = MakeKeys(100, 117); + EXPECT_EQ(GrpcMaxItemsPerBatch(keys, 0), 100u); +} + +TEST(GrpcBatchChunkingTest, EveryChunkFitsUnderTheLimit) { + // 155,157 keys of the length the real deployment used. + const auto keys = MakeKeys(155157, 117); + const size_t limit = GrpcMaxMessageBytes(); + + size_t sent = 0, chunks = 0; + while (sent < keys.size()) { + const size_t take = GrpcMaxItemsPerBatch(keys, sent); + ASSERT_GT(take, 0u) << "a zero-sized chunk would loop forever"; + + size_t bytes = 0; + for (size_t i = sent; i < sent + take; ++i) bytes += keys[i].size() + 8; + EXPECT_LE(bytes, limit) << "chunk " << chunks << " would be refused"; + + sent += take; + ++chunks; + } + EXPECT_EQ(sent, keys.size()) << "chunking must cover every key exactly once"; + EXPECT_GE(chunks, 1u); +} + +TEST(GrpcBatchChunkingTest, OversizedSingleItemStillMakesProgress) { + // One key larger than the whole budget: returning 0 here would spin the + // dispatch loop forever, which is worse than letting the RPC fail. + const auto keys = MakeKeys(1, static_cast(GrpcMaxMessageBytes()) + 1024); + EXPECT_EQ(GrpcMaxItemsPerBatch(keys, 0), 1u); +} + +TEST(GrpcBatchChunkingTest, LimitHasAFloorAndADefault) { + // The default must not be below what the master surfaces already used, or + // this change would silently lower a working limit. + EXPECT_GE(GrpcMaxMessageBytes(), kMinGrpcMaxMessageBytes); + EXPECT_GE(kDefaultGrpcMaxMessageBytes, 64u * 1024u * 1024u); +} + } // namespace } // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_master_eviction_chain.cpp b/tests/cpp/umbp/distributed/test_master_eviction_chain.cpp new file mode 100644 index 000000000..ca5e0c70d --- /dev/null +++ b/tests/cpp/umbp/distributed/test_master_eviction_chain.cpp @@ -0,0 +1,515 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// The master decides what to evict and the peer executes it, and until now the +// two halves were only tested apart: test_master_evict_strategy covers victim +// selection given candidates, and test_peer_pool covers PeerPool::Evict given +// keys. Nothing joined them, so a break anywhere between "this tier is over its +// watermark" and "the key moved to the tier below" was invisible. +// +// That gap is not hypothetical. offload_trigger=on_evict and promote_mode=move +// both do their work inside PeerPool::Evict, whose only production caller is the +// EvictKey handler the master drives. Measuring them under umbp_tier_bench -- +// which starts no EvictionManager -- reports zero offloads and no reclamation, +// which reads exactly like the features being unimplemented. +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// eviction_manager.h only forward-declares MasterEvictStrategy, and the +// constructor's defaulted unique_ptr needs the complete type to destroy it. +#include "umbp/distributed/master/evict_strategy.h" +#include "umbp/distributed/master/eviction_manager.h" +#include "umbp/distributed/master/in_memory_master_metadata_store.h" +#include "umbp/distributed/peer/backend/page_backend.h" +#include "umbp/distributed/pool/peer_pool.h" +#include "umbp/distributed/transfer/local_copy_engine.h" + +namespace mori::umbp { +namespace { + +using Clock = std::chrono::system_clock; + +constexpr uint64_t kPageSize = 64; +constexpr auto kNoExpiry = std::chrono::milliseconds(30000); +constexpr char kNodeId[] = "node-a"; + +void RegisterPageBackend(BackendRegistry* registry, TransferEngine* engine, const std::string& name, + TierType tier, uint64_t pages) { + PageBackend::OwnershipConfig ownership; + ownership.buffer_sizes = {pages * kPageSize}; + auto backend = MakePageBackend(tier, kPageSize, ownership, kNoExpiry, kNoExpiry); + ASSERT_TRUE(backend->Init(engine)); + ASSERT_TRUE(registry->Register(name, std::move(backend))); +} + +// hot offloads to cold on evict, which is the trigger whose only driver is the +// master. Cold is a distinct medium so the two tiers land in different master +// buckets; a two-DRAM-tier policy would collapse into one, since the master +// tracks capacity per TierType and knows nothing of logical tiers. +std::vector HotColdOnEvict() { + LogicalTierConfig hot; + hot.name = "hot"; + hot.entry = true; + hot.members = {{"hot", 100}}; + hot.offload_to = {"cold"}; + hot.trigger = PoolOffloadTrigger::kOnEvict; + hot.high_watermark = 0.9; + hot.low_watermark = 0.7; + LogicalTierConfig cold; + cold.name = "cold"; + cold.members = {{"cold", 100}}; + return {hot, cold}; +} + +void CommitKey(PeerPool* pool, const std::string& key) { + PoolPlacementRequest request; + request.key = key; + request.size = kPageSize; + request.tier = TierType::DRAM; + auto allocation = pool->BatchAllocate({request}).front(); + ASSERT_EQ(allocation.allocation.outcome, AllocateOutcome::kSuccessAllocated) << key; + ASSERT_TRUE(pool->BatchCommit({PoolCommitRequest{ + {allocation.backend_id, allocation.allocation.slot_id}, key}}) + .front() + .commit.success) + << key; +} + +// Stands in for MasterPeerStubPool, minus gRPC: the handler in peer_service.cpp +// does exactly this one call, so routing the dispatch straight into the pool +// exercises the same decision-to-execution path a deployed master takes. +class PoolDispatcher final : public EvictKeyDispatcher { + public: + explicit PoolDispatcher(PeerPool* pool) : pool_(pool) {} + + void DispatchEvictKey(const std::string& node_id, const std::string& /*peer_address*/, + std::vector keys) override { + auto results = pool_->Evict(keys, PoolEvictMode::kReclaim); + std::lock_guard lock(mutex_); + ++rounds_; + for (size_t i = 0; i < keys.size(); ++i) { + dispatched_.push_back(keys[i]); + if (i < results.size()) freed_bytes_ += results[i].bytes_freed; + } + node_ids_.push_back(node_id); + } + + std::vector Dispatched() const { + std::lock_guard lock(mutex_); + return dispatched_; + } + size_t Rounds() const { + std::lock_guard lock(mutex_); + return rounds_; + } + uint64_t FreedBytes() const { + std::lock_guard lock(mutex_); + return freed_bytes_; + } + std::vector NodeIds() const { + std::lock_guard lock(mutex_); + return node_ids_; + } + + private: + PeerPool* pool_; + mutable std::mutex mutex_; + std::vector dispatched_; + std::vector node_ids_; + size_t rounds_ = 0; + uint64_t freed_bytes_ = 0; +}; + +void RegisterNode(InMemoryMasterMetadataStore* store, Clock::time_point now) { + ClientRegistration registration; + registration.node_id = kNodeId; + registration.node_address = "127.0.0.1:1"; + registration.peer_address = "127.0.0.1:2"; + ASSERT_TRUE(store->RegisterClient(registration, now, std::chrono::seconds(60))); +} + +// One heartbeat: `caps` is what makes a bucket look over or under its watermark, +// and `now` becomes the access time of any key this heartbeat introduces -- +// the store stamps it on first insertion only, so staggering heartbeats is how a +// test gets a victim order it can predict. +void ReportHeartbeat(InMemoryMasterMetadataStore* store, uint64_t seq, Clock::time_point now, + const std::map& caps, TierType key_tier, + const std::string& logical_tier, const std::vector& keys, + bool full_sync) { + std::vector events; + events.reserve(keys.size()); + for (const auto& key : keys) { + KvEvent event; + event.kind = KvEvent::Kind::ADD; + event.key = key; + event.tier = key_tier; + event.size = kPageSize; + event.logical_tier = logical_tier; + events.push_back(std::move(event)); + } + store->ApplyHeartbeat(kNodeId, seq, now, caps, events, full_sync); +} + +// Reports the node as holding `used` of `total` DRAM bytes, so RunOnce sees a +// bucket over its high watermark. +void PublishNode(InMemoryMasterMetadataStore* store, const std::vector& keys, + uint64_t total_bytes, uint64_t used_bytes) { + const auto now = Clock::now(); + RegisterNode(store, now); + std::map caps; + caps[TierType::DRAM] = TierCapacity{total_bytes, total_bytes - used_bytes, kPageSize}; + ReportHeartbeat(store, /*seq=*/1, now, caps, TierType::DRAM, "hot", keys, /*full_sync=*/true); +} + +EvictionConfig OneSecondRounds() { + EvictionConfig config; + config.high_watermark = 0.9; + config.low_watermark = 0.7; + // The loop sleeps before its first pass, and one second is the floor the + // env override enforces, so this is the shortest a test can wait for a + // round it did not trigger itself. + config.check_interval = std::chrono::seconds(1); + return config; +} + +bool WaitFor(const std::function& done, std::chrono::milliseconds budget) { + const auto deadline = std::chrono::steady_clock::now() + budget; + while (std::chrono::steady_clock::now() < deadline) { + if (done()) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + return done(); +} + +// The whole chain: a tier over its high watermark, through victim selection and +// the dispatch the master performs, into the pool that owns the bytes. on_evict +// has to demote here -- dropping the key instead would free the same bytes and +// look identical to the master, which is why the assertion is on where the key +// ended up and not on how much was freed. +TEST(MasterEvictionChain, OverWatermarkTierDemotesIntoTheColdTier) { + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, /*pages=*/4); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, /*pages=*/8); + + auto compiled = LogicalTierGraph::Compile(HotColdOnEvict(), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + const std::vector keys = {"k0", "k1", "k2", "k3"}; + for (const auto& key : keys) CommitKey(&pool, key); + ASSERT_EQ(registry.Get("hot")->OwnedKeyCount(), keys.size()); + + InMemoryMasterMetadataStore store; + // 4 of 4 pages used is over the 0.9 watermark, and the drain target of 0.7 + // leaves a budget big enough for more than one key. + PublishNode(&store, keys, /*total_bytes=*/4 * kPageSize, /*used_bytes=*/4 * kPageSize); + + PoolDispatcher dispatcher(&pool); + EvictionManager manager(store, OneSecondRounds(), &dispatcher); + manager.Start(); + const bool dispatched = + WaitFor([&] { return !dispatcher.Dispatched().empty(); }, std::chrono::seconds(8)); + manager.Stop(); + + ASSERT_TRUE(dispatched) << "master never dispatched an EvictKey for an over-watermark tier"; + EXPECT_EQ(dispatcher.NodeIds().front(), kNodeId); + + // Every key the master named must have moved down a tier rather than vanished. + const auto evicted = dispatcher.Dispatched(); + ASSERT_FALSE(evicted.empty()); + for (const auto& key : evicted) { + EXPECT_FALSE(registry.Get("hot")->Contains(key)) << key << " still in hot"; + EXPECT_TRUE(registry.Get("cold")->Contains(key)) << key << " was dropped, not demoted"; + } + EXPECT_GT(dispatcher.FreedBytes(), 0u); + + // Keys the master did not name are untouched: eviction is bounded by the + // budget, so a round must not drain the whole tier. + EXPECT_EQ(registry.Get("hot")->OwnedKeyCount() + evicted.size(), keys.size()); +} + +// A tier under its watermark must produce no eviction at all. Without this the +// test above would still pass if RunOnce evicted unconditionally, which is the +// same observation -- keys moving to cold -- for an entirely different reason. +TEST(MasterEvictionChain, TierBelowWatermarkIsLeftAlone) { + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, /*pages=*/8); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, /*pages=*/8); + + auto compiled = LogicalTierGraph::Compile(HotColdOnEvict(), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + const std::vector keys = {"k0", "k1"}; + for (const auto& key : keys) CommitKey(&pool, key); + + InMemoryMasterMetadataStore store; + // 2 of 8 pages is a quarter full, far below the 0.9 watermark. + PublishNode(&store, keys, /*total_bytes=*/8 * kPageSize, /*used_bytes=*/2 * kPageSize); + + PoolDispatcher dispatcher(&pool); + EvictionManager manager(store, OneSecondRounds(), &dispatcher); + manager.Start(); + // Long enough for two rounds at the configured interval. + std::this_thread::sleep_for(std::chrono::milliseconds(2500)); + manager.Stop(); + + EXPECT_EQ(dispatcher.Rounds(), 0u) << "evicted from a tier that was under its watermark"; + for (const auto& key : keys) { + EXPECT_TRUE(registry.Get("hot")->Contains(key)) << key; + EXPECT_FALSE(registry.Get("cold")->Contains(key)) << key; + } +} + +// promote_mode=move deletes the cold copy once a key is promoted, but the +// delete races the read lease the triggering read still holds, so it defers: the +// source parks in draining_sources_ and the promotion reports success. Nothing +// requeues it, and the one place that drains it is PeerPool::Evict -- so whether +// `move` ever reclaims anything at all depends entirely on this chain, and a +// benchmark with no eviction driver reports move and copy as byte-identical. +// +// This pins the reclamation down: a master round that names the key frees the +// parked source and leaves the promoted copy alone. Losing the target instead +// would be silent data loss, which is the one way move can be worse than copy. +TEST(MasterEvictionChain, MasterRoundReclaimsTheParkedSourceOfAMovePromotion) { + // Short enough that the lease which defers the source delete has expired by + // the time the master's round arrives a second later. + constexpr auto kShortLease = std::chrono::milliseconds(50); + LocalCopyEngine engine; + BackendRegistry registry; + PageBackend::OwnershipConfig hot_ownership; + hot_ownership.buffer_sizes = {4 * kPageSize}; + auto hot = MakePageBackend(TierType::DRAM, kPageSize, hot_ownership, kNoExpiry, kShortLease); + ASSERT_TRUE(hot->Init(&engine)); + ASSERT_TRUE(registry.Register("hot", std::move(hot))); + PageBackend::OwnershipConfig cold_ownership; + cold_ownership.buffer_sizes = {4 * kPageSize}; + auto cold = MakePageBackend(TierType::SSD, kPageSize, cold_ownership, kNoExpiry, kShortLease); + ASSERT_TRUE(cold->Init(&engine)); + ASSERT_TRUE(registry.Register("cold", std::move(cold))); + + auto tiers = HotColdOnEvict(); + tiers.back().promote_trigger = PoolPromoteTrigger::kOnRead; + tiers.back().promote_mode = PoolTransitionMode::kMove; + auto compiled = LogicalTierGraph::Compile(tiers, registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + // Put the key in cold directly, then read it so the promotion fires. + PoolPlacementRequest request; + request.key = "moved"; + request.size = kPageSize; + request.tier = TierType::DRAM; + request.logical_tier = "cold"; + auto allocation = pool.BatchAllocate({request}).front(); + ASSERT_EQ(allocation.allocation.outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_TRUE( + pool.BatchCommit( + {PoolCommitRequest{{allocation.backend_id, allocation.allocation.slot_id}, "moved"}}) + .front() + .commit.success); + ASSERT_TRUE(registry.Get("cold")->Contains("moved")); + + ASSERT_TRUE(pool.BatchResolve({"moved"}, false).front().resolved.found); + ASSERT_TRUE( + WaitFor([&] { return registry.Get("hot")->Contains("moved"); }, std::chrono::seconds(2))) + << "promotion never reached the hot tier"; + // The source is still there: this is the deferral, not a failure. + EXPECT_TRUE(registry.Get("cold")->Contains("moved")) + << "source was deleted during the promotion, so there is nothing to reclaim"; + + InMemoryMasterMetadataStore store; + PublishNode(&store, {"moved"}, /*total_bytes=*/4 * kPageSize, /*used_bytes=*/4 * kPageSize); + + PoolDispatcher dispatcher(&pool); + EvictionManager manager(store, OneSecondRounds(), &dispatcher); + manager.Start(); + const bool reclaimed = + WaitFor([&] { return !registry.Get("cold")->Contains("moved"); }, std::chrono::seconds(8)); + manager.Stop(); + + EXPECT_TRUE(reclaimed) << "the parked source outlived a master eviction round that named it"; + // The promoted copy is the only one left, so the key is still servable. + EXPECT_TRUE(registry.Get("hot")->Contains("moved")) + << "the evict drained the promoted target instead of the parked source"; +} + +// Master-side eviction reads one capacity per TierType, and a client with a +// weighted multi-backend policy sums its backends into that one figure +// (MasterClient::SnapshotAndCacheTierCapacities). So a hot tier sharing DRAM with +// a larger cold tier is reported against their combined size, and a hot tier that +// is completely full reads as a medium that is nearly empty. +// +// The pool side already rejects this reasoning for its own watermarks: +// PeakMemberUtilization exists because "pressure does not average", and lets the +// fullest member speak for the tier. RunOnce averages anyway, being one level up +// and seeing only the sum. +// +// Both halves below hold the hot tier at 4 of 4 pages. The only difference is +// whether a sibling logical tier shares its medium, which is not a property of +// the hot tier at all -- and it decides whether the master acts. +TEST(MasterEvictionChain, SiblingTierInTheSameMediumHidesTheOverload) { + const std::vector keys = {"k0", "k1", "k2", "k3"}; + + const auto run_round = [&keys](uint64_t reported_total_pages) { + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, /*pages=*/4); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, /*pages=*/8); + auto compiled = LogicalTierGraph::Compile(HotColdOnEvict(), registry); + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + for (const auto& key : keys) CommitKey(&pool, key); + + InMemoryMasterMetadataStore store; + PublishNode(&store, keys, reported_total_pages * kPageSize, /*used_bytes=*/4 * kPageSize); + PoolDispatcher dispatcher(&pool); + EvictionManager manager(store, OneSecondRounds(), &dispatcher); + manager.Start(); + WaitFor([&] { return !dispatcher.Dispatched().empty(); }, std::chrono::seconds(4)); + manager.Stop(); + return dispatcher.Dispatched().size(); + }; + + // Hot alone in its medium: 4 of 4 pages, over the watermark, master acts. + EXPECT_GT(run_round(/*reported_total_pages=*/4), 0u) + << "a full tier alone in its medium was not evicted from"; + + // Same full hot tier, now summed with a 32-page DRAM sibling: 4 of 36 pages + // reads as 11% and the master sees nothing to do. This is the bug, and the + // assertion is deliberately on the current behaviour -- when the decision moves + // to per-tier utilisation this expectation is what should fail. + EXPECT_EQ(run_round(/*reported_total_pages=*/36), 0u) + << "master evicted despite the aggregate being under the watermark -- if " + "this now fails, the decision has moved off the per-medium sum and this " + "test should be inverted"; +} + +// The timing half of the finding above, which the single-key test cannot reach: +// with one candidate the master has no ordering decision to make. Victims are +// named least-recently-accessed first, and a key is promoted precisely because +// it was just read, so its parked source should be the last thing its tier gives +// up. That is what makes `move` reclaim late rather than not at all -- a +// distinction worth a test, because the two look identical over any window +// shorter than the tier's turnover. +TEST(MasterEvictionChain, ParkedMoveSourceIsNamedLastAmongColderKeys) { + constexpr auto kShortLease = std::chrono::milliseconds(50); + LocalCopyEngine engine; + BackendRegistry registry; + PageBackend::OwnershipConfig hot_ownership; + hot_ownership.buffer_sizes = {4 * kPageSize}; + auto hot = MakePageBackend(TierType::DRAM, kPageSize, hot_ownership, kNoExpiry, kShortLease); + ASSERT_TRUE(hot->Init(&engine)); + ASSERT_TRUE(registry.Register("hot", std::move(hot))); + PageBackend::OwnershipConfig cold_ownership; + cold_ownership.buffer_sizes = {10 * kPageSize}; + auto cold = MakePageBackend(TierType::SSD, kPageSize, cold_ownership, kNoExpiry, kShortLease); + ASSERT_TRUE(cold->Init(&engine)); + ASSERT_TRUE(registry.Register("cold", std::move(cold))); + + auto tiers = HotColdOnEvict(); + tiers.back().promote_trigger = PoolPromoteTrigger::kOnRead; + tiers.back().promote_mode = PoolTransitionMode::kMove; + auto compiled = LogicalTierGraph::Compile(tiers, registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + const std::vector colder = {"cold0", "cold1", "cold2"}; + const auto place_in_cold = [&](const std::string& key) { + PoolPlacementRequest request; + request.key = key; + request.size = kPageSize; + request.tier = TierType::DRAM; + request.logical_tier = "cold"; + auto allocation = pool.BatchAllocate({request}).front(); + ASSERT_EQ(allocation.allocation.outcome, AllocateOutcome::kSuccessAllocated) << key; + ASSERT_TRUE(pool.BatchCommit({PoolCommitRequest{ + {allocation.backend_id, allocation.allocation.slot_id}, key}}) + .front() + .commit.success) + << key; + }; + for (const auto& key : colder) place_in_cold(key); + place_in_cold("moved"); + + // Read "moved" so it promotes; move defers the source delete behind the read + // lease, leaving the source parked in cold alongside the three colder keys. + ASSERT_TRUE(pool.BatchResolve({"moved"}, false).front().resolved.found); + ASSERT_TRUE( + WaitFor([&] { return registry.Get("hot")->Contains("moved"); }, std::chrono::seconds(2))); + ASSERT_TRUE(registry.Get("cold")->Contains("moved")); + + InMemoryMasterMetadataStore store; + const auto now = Clock::now(); + RegisterNode(&store, now); + // 10 of 10 SSD pages used, draining to 0.7, so the budget covers three keys of + // the four present. DRAM is left comfortably under its watermark so the hot + // tier contributes no bucket of its own. + std::map caps; + caps[TierType::SSD] = TierCapacity{10 * kPageSize, 0, kPageSize}; + caps[TierType::DRAM] = TierCapacity{10 * kPageSize, 9 * kPageSize, kPageSize}; + // The three colder keys enter ten seconds earlier, so their access stamps are + // unambiguously older than the promoted key's. + ReportHeartbeat(&store, /*seq=*/1, now - std::chrono::seconds(10), caps, TierType::SSD, "cold", + colder, /*full_sync=*/true); + ReportHeartbeat(&store, /*seq=*/2, now, caps, TierType::SSD, "cold", {"moved"}, + /*full_sync=*/false); + + PoolDispatcher dispatcher(&pool); + EvictionManager manager(store, OneSecondRounds(), &dispatcher); + manager.Start(); + const bool dispatched = WaitFor([&] { return dispatcher.Dispatched().size() >= colder.size(); }, + std::chrono::seconds(8)); + manager.Stop(); + + ASSERT_TRUE(dispatched) << "master never named the colder keys"; + const auto named = dispatcher.Dispatched(); + + // The parked source is the newest entry in its tier, so a budget that covers + // three of four keys must spend it on the other three. + EXPECT_EQ(std::count(named.begin(), named.end(), "moved"), 0) + << "the just-promoted key was named ahead of colder entries"; + for (const auto& key : colder) { + EXPECT_EQ(std::count(named.begin(), named.end(), key), 1) << key << " was not named"; + } + // And it really is still occupying the tier: this is the lag, stated as an + // observation rather than inferred from the ordering. + EXPECT_TRUE(registry.Get("cold")->Contains("moved")) + << "parked source was reclaimed after all, so the ordering concern is moot"; + EXPECT_TRUE(registry.Get("hot")->Contains("moved")); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_master_metadata_store_interface.cpp b/tests/cpp/umbp/distributed/test_master_metadata_store_interface.cpp deleted file mode 100644 index 9fcd952cd..000000000 --- a/tests/cpp/umbp/distributed/test_master_metadata_store_interface.cpp +++ /dev/null @@ -1,134 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -// Phase 1 compile/instantiation gate for IMasterMetadataStore. -// -// The interface is abstract with no implementation yet, so there is no runtime -// behavior to exercise. The bar for Phase 1 is that the contract is well-formed -// and instantiable: -// 1. MockMasterMetadataStore overrides every pure-virtual (a missing or -// ill-typed override makes the mock abstract → fails to instantiate). -// 2. A MockMasterMetadataStore is usable through an IMasterMetadataStore&, -// proving the override set is complete. -// Behavioral assertions arrive with InMemoryMasterMetadataStore in Phase 2. - -#include - -#include - -#include "mock_master_metadata_store.h" -#include "umbp/distributed/master/master_metadata_store.h" - -namespace mori::umbp { -namespace { - -// Instantiation gate: if the interface had an orphaned or ill-typed pure -// virtual, MockMasterMetadataStore would stay abstract and this would not -// compile. -TEST(MasterMetadataStoreInterface, MockIsInstantiableThroughInterface) { - MockMasterMetadataStore mock; - IMasterMetadataStore& store = mock; - (void)store; - SUCCEED(); -} - -// Signature-completeness spot check: name every interface method once through -// the base-class pointer with a default ON_CALL, mirroring the §1b delta table -// plus the two added hit-count methods (GetExternalKvHitCounts, -// GarbageCollectHits) and the `now` parameter on MatchExternalKv. This guards -// against silently dropping the live GetExternalKvHitCounts RPC path. -TEST(MasterMetadataStoreInterface, EveryMethodIsCallableThroughInterface) { - using ::testing::_; - using ::testing::NiceMock; - using ::testing::Return; - using namespace std::chrono_literals; - - // NiceMock: these are default-action calls, not behavior under test, so the - // "uninteresting call" warnings would just be noise. - NiceMock mock; - const auto now = std::chrono::system_clock::now(); - - ON_CALL(mock, RegisterClient(_, _, _)).WillByDefault(Return(true)); - ON_CALL(mock, ApplyHeartbeat(_, _, _, _, _, _)) - .WillByDefault(Return(HeartbeatResult{HeartbeatResult::APPLIED, 0})); - ON_CALL(mock, ExpireStaleClients(_)).WillByDefault(Return(std::vector{})); - ON_CALL(mock, RegisterExternalKvIfAlive(_, _, _)).WillByDefault(Return(true)); - ON_CALL(mock, GarbageCollectHits(_)).WillByDefault(Return(0)); - ON_CALL(mock, LookupBlock(_)).WillByDefault(Return(std::vector{})); - ON_CALL(mock, LookupBlockForRouteGet(_, _, _, _)).WillByDefault(Return(std::vector{})); - ON_CALL(mock, BatchLookupBlockForRouteGet(_, _, _, _)) - .WillByDefault(Return(std::vector>{})); - ON_CALL(mock, BatchExistsBlock(_)).WillByDefault(Return(std::vector{})); - ON_CALL(mock, EnumerateEvictionCandidates(_, _, _, _)) - .WillByDefault(Return(std::map>{})); - ON_CALL(mock, GetClient(_)).WillByDefault(Return(std::nullopt)); - ON_CALL(mock, IsClientAlive(_)).WillByDefault(Return(false)); - ON_CALL(mock, GetPeerAddress(_)).WillByDefault(Return(std::nullopt)); - ON_CALL(mock, ListAliveClients()).WillByDefault(Return(std::vector{})); - ON_CALL(mock, AliveClientCount()).WillByDefault(Return(0)); - ON_CALL(mock, GetClientTags(_)).WillByDefault(Return(std::vector{})); - ON_CALL(mock, MatchExternalKv(_, _, _)).WillByDefault(Return(std::vector{})); - ON_CALL(mock, GetExternalKvHitCounts(_)) - .WillByDefault(Return(std::vector{})); - ON_CALL(mock, GetExternalKvCount(_)).WillByDefault(Return(0)); - - IMasterMetadataStore& store = mock; - - // Cross-store writes. - ClientRegistration reg; - reg.node_id = "node-a"; - EXPECT_TRUE(store.RegisterClient(reg, now, 30s)); - store.UnregisterClient("node-a"); - EXPECT_EQ(store.ApplyHeartbeat("node-a", 1, now, {}, {}, false).status, HeartbeatResult::APPLIED); - EXPECT_TRUE(store.ExpireStaleClients(now).empty()); - - // External-KV writes. - EXPECT_TRUE(store.RegisterExternalKvIfAlive("node-a", {"h0"}, TierType::HBM)); - store.UnregisterExternalKv("node-a", {"h0"}, TierType::HBM); - store.UnregisterExternalKvByTier("node-a", TierType::HBM); - store.UnregisterExternalKvByNode("node-a"); - EXPECT_EQ(store.GarbageCollectHits(now), 0u); - - // Block reads. - EXPECT_TRUE(store.LookupBlock("k0").empty()); - EXPECT_TRUE(store.LookupBlockForRouteGet("k0", {}, now, 5s).empty()); - EXPECT_TRUE(store.BatchLookupBlockForRouteGet({"k0"}, {}, now, 5s).empty()); - EXPECT_TRUE(store.BatchExistsBlock({"k0"}).empty()); - EXPECT_TRUE( - store.EnumerateEvictionCandidates({}, EvictionOrder::kLeastRecentlyAccessed, 0, now).empty()); - - // Client reads. - EXPECT_FALSE(store.GetClient("node-a").has_value()); - EXPECT_FALSE(store.IsClientAlive("node-a")); - EXPECT_FALSE(store.GetPeerAddress("node-a").has_value()); - EXPECT_TRUE(store.ListAliveClients().empty()); - EXPECT_EQ(store.AliveClientCount(), 0u); - EXPECT_TRUE(store.GetClientTags("node-a").empty()); - - // External-KV reads, incl. the two added hit-count methods + `now` param. - EXPECT_TRUE(store.MatchExternalKv({"h0"}, /*count_as_hit=*/true, now).empty()); - EXPECT_TRUE(store.GetExternalKvHitCounts({"h0"}).empty()); - EXPECT_EQ(store.GetExternalKvCount("node-a"), 0u); -} - -} // namespace -} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_medium_selection.cpp b/tests/cpp/umbp/distributed/test_medium_selection.cpp new file mode 100644 index 000000000..fb4aa5569 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_medium_selection.cpp @@ -0,0 +1,349 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// ONE MEDIUM PER NODE. +// +// PoolClient::Init used to register DRAM unconditionally and then add HBM/SSD +// beside it, which made "--tier ssd" mean "DRAM *and* SSD". That is not a tier +// stack: the routing plane treats every advertised tier as an equally valid put +// target (Phase 4 deleted the hardcoded tier orders), so the second medium +// mirrors the first rather than sitting behind it. UMBPDistributedConfig:: +// medium now selects exactly one. +// +// Covered here: +// 1. The lowering — UMBPMedium -> PoolClientConfig::medium, and the SSD +// opt-in coming from the selector rather than UMBPConfig::ssd.enabled. +// 2. Validation — the selected medium's sizing is checked, the others' +// are not. +// 3. The registry — a live PoolClient holds exactly one backend, and an +// SSD node has NO DRAM tier to fall back to. +// 4. The wire — a cross-node BatchPut lands on a peer whose only medium +// is SSD. PartitionBatchPutTargets used to drop any route that was not +// DRAM/HBM, so this is the case that silently failed. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/common/config.h" +#include "umbp/distributed/config.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/pool_client.h" + +namespace mori::umbp { +namespace { + +namespace fs = std::filesystem; + +constexpr size_t kPageSize = 4096; +constexpr size_t kDramCap = 8 << 20; + +// --------------------------------------------------------------------------- +// 1. Lowering +// --------------------------------------------------------------------------- + +UMBPDistributedConfig MinimalDistributedConfig() { + UMBPDistributedConfig dc; + dc.master_config.master_address = "localhost:1"; + dc.master_config.node_id = "n0"; + dc.master_config.node_address = "127.0.0.1"; + return dc; +} + +TEST(MediumLoweringTest, DefaultsToDram) { + auto pc = ToPoolClientConfig(MinimalDistributedConfig(), DramOwnershipConfig{}, PeerSsdConfig{}); + EXPECT_EQ(pc.medium, TierType::DRAM); + EXPECT_EQ(pc.placement_policy, PoolPlacementPolicy::SINGLE_BACKEND); + EXPECT_TRUE(pc.backends.empty()) + << "legacy lowering must defer backend synthesis until PoolClient::Init"; + // The SSD tier stays shut unless the selector names it, whatever the SSD + // block says — UMBPSsdConfig::enabled defaults to true and describes the + // LOCAL-mode tier, so keying off it would opt every deployment in. + PeerSsdConfig ssd; + ssd.enabled = true; + ssd.ssd.enabled = true; + auto pc2 = ToPoolClientConfig(MinimalDistributedConfig(), DramOwnershipConfig{}, ssd); + EXPECT_EQ(pc2.medium, TierType::DRAM); + EXPECT_FALSE(pc2.ssd.enabled); +} + +TEST(MediumLoweringTest, SelectorPicksTierAndOptsSsdIn) { + auto dc = MinimalDistributedConfig(); + + dc.medium = UMBPMedium::HBM; + dc.hbm.device = 3; + dc.hbm.capacity_bytes = 1 << 20; + auto hbm = ToPoolClientConfig(dc, DramOwnershipConfig{}, PeerSsdConfig{}); + EXPECT_EQ(hbm.medium, TierType::HBM); + EXPECT_TRUE(hbm.backends.empty()); + EXPECT_EQ(hbm.hbm.device, 3); + ASSERT_EQ(hbm.hbm.buffer_sizes.size(), 1u); + EXPECT_EQ(hbm.hbm.buffer_sizes[0], 1u << 20); + EXPECT_FALSE(hbm.ssd.enabled); + + dc.medium = UMBPMedium::SSD; + auto ssd = ToPoolClientConfig(dc, DramOwnershipConfig{}, PeerSsdConfig{}); + EXPECT_EQ(ssd.medium, TierType::SSD); + EXPECT_TRUE(ssd.ssd.enabled) << "selecting SSD is what opts the tier in"; +} + +// --------------------------------------------------------------------------- +// 2. Validation — only the selected medium is sized-checked +// --------------------------------------------------------------------------- + +UMBPConfig MinimalUmbpConfig() { + UMBPConfig cfg; + cfg.distributed = MinimalDistributedConfig(); + return cfg; +} + +TEST(MediumValidationTest, UnselectedMediaAreNotChecked) { + auto cfg = MinimalUmbpConfig(); + // hbm.capacity_bytes is 0 here, and that is fine: this is a DRAM node. + std::string err; + EXPECT_TRUE(cfg.Validate(&err)) << err; +} + +TEST(MediumValidationTest, SelectedHbmNeedsCapacity) { + auto cfg = MinimalUmbpConfig(); + cfg.distributed->medium = UMBPMedium::HBM; + std::string err; + EXPECT_FALSE(cfg.Validate(&err)); + EXPECT_NE(err.find("distributed.hbm.capacity_bytes"), std::string::npos) << err; + + cfg.distributed->hbm.capacity_bytes = 1 << 20; + EXPECT_TRUE(cfg.Validate(&err)) << err; +} + +TEST(MediumValidationTest, SelectedSsdNeedsCapacityEvenWhenTierFlagIsOff) { + auto cfg = MinimalUmbpConfig(); + cfg.distributed->medium = UMBPMedium::SSD; + cfg.ssd.enabled = false; // must NOT short-circuit the sizing check + cfg.ssd.capacity_bytes = 0; + std::string err; + EXPECT_FALSE(cfg.Validate(&err)); + EXPECT_NE(err.find("ssd.capacity_bytes"), std::string::npos) << err; +} + +// --------------------------------------------------------------------------- +// 3 + 4. Live PoolClient: one registered backend, and a remote put onto it +// --------------------------------------------------------------------------- + +// Free ephemeral port; PoolClient binds and advertises this verbatim, so a +// hardcoded base would collide with concurrent test processes. +uint16_t NextPeerServicePort() { + int fd = ::socket(AF_INET, SOCK_STREAM, 0); + if (fd >= 0) { + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_ANY); + addr.sin_port = 0; + socklen_t len = sizeof(addr); + if (::bind(fd, reinterpret_cast(&addr), sizeof(addr)) == 0 && + ::getsockname(fd, reinterpret_cast(&addr), &len) == 0) { + uint16_t port = ntohs(addr.sin_port); + ::close(fd); + return port; + } + ::close(fd); + } + static std::atomic next{ + static_cast(54000 + (static_cast(::getpid()) % 4000))}; + return next.fetch_add(1); +} + +class MediumSelectionTest : public ::testing::Test { + protected: + void SetUp() override { + static std::atomic counter{0}; + ssd_dir_ = fs::temp_directory_path() / ("umbp_medium_sel_" + std::to_string(::getpid()) + "_" + + std::to_string(counter.fetch_add(1))); + fs::remove_all(ssd_dir_); + + MasterServerConfig master_cfg; + master_cfg.listen_address = "0.0.0.0:0"; + master_ = std::make_unique(std::move(master_cfg)); + server_thread_ = std::thread([this] { master_->Run(); }); + for (int i = 0; i < 50 && master_->GetBoundPort() == 0; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_NE(master_->GetBoundPort(), 0) << "Master failed to start"; + master_addr_ = "localhost:" + std::to_string(master_->GetBoundPort()); + } + + void TearDown() override { + for (auto& c : clients_) { + if (c) c->Shutdown(); + } + clients_.clear(); + if (master_) master_->Shutdown(); + if (server_thread_.joinable()) server_thread_.join(); + std::error_code ec; + fs::remove_all(ssd_dir_, ec); + } + + PoolClientConfig BaseConfig(const std::string& node_id) { + PoolClientConfig cfg; + cfg.master_config.node_id = node_id; + cfg.master_config.node_address = "127.0.0.1"; + cfg.master_config.master_address = master_addr_; + cfg.io_engine.host = "0.0.0.0"; + cfg.io_engine.port = 0; + cfg.peer_service_port = NextPeerServicePort(); + cfg.dram_page_size = kPageSize; + // Off so a remote read stays remote and the assertions below describe the + // medium under test, not a re-cached copy on the reader. + cfg.cache_remote_fetches = false; + return cfg; + } + + PoolClientConfig DramConfig(const std::string& node_id, size_t capacity) { + auto cfg = BaseConfig(node_id); + cfg.medium = TierType::DRAM; + cfg.dram.buffer_sizes = {capacity}; + return cfg; + } + + PoolClientConfig SsdConfig(const std::string& node_id) { + auto cfg = BaseConfig(node_id); + cfg.medium = TierType::SSD; + // DRAM sizing is still present and must be ignored — a node that selected + // SSD gets no host pool, whatever this says. + cfg.dram.buffer_sizes = {kDramCap}; + cfg.ssd_staging_buffer_slots = 8; + cfg.ssd.ssd.storage_dir = ssd_dir_.string(); + cfg.ssd.ssd.capacity_bytes = 64ULL * 1024 * 1024; + cfg.ssd.ssd.io.backend = UMBPIoBackend::Posix; // avoid io_uring container flakiness + return cfg; + } + + PoolClient* Start(PoolClientConfig cfg) { + auto client = std::make_unique(std::move(cfg)); + if (!client->Init()) return nullptr; + clients_.push_back(std::move(client)); + return clients_.back().get(); + } + + fs::path ssd_dir_; + std::string master_addr_; + std::unique_ptr master_; + std::thread server_thread_; + std::vector> clients_; +}; + +TEST_F(MediumSelectionTest, DramNodeRegistersOnlyDram) { + auto* node = Start(DramConfig("node-dram", kDramCap)); + ASSERT_NE(node, nullptr); + EXPECT_EQ(node->Medium(), TierType::DRAM); + EXPECT_EQ(node->Backends().All().size(), 1u); + EXPECT_NE(node->Backends().Get(TierType::DRAM), nullptr); + EXPECT_EQ(node->Backends().Get(TierType::SSD), nullptr); + EXPECT_EQ(node->Backends().Get(TierType::HBM), nullptr); +} + +TEST_F(MediumSelectionTest, SsdNodeHasNoDramTier) { + auto* node = Start(SsdConfig("node-ssd")); + ASSERT_NE(node, nullptr); + EXPECT_EQ(node->Medium(), TierType::SSD); + ASSERT_EQ(node->Backends().All().size(), 1u) << "an SSD node must not also carry a DRAM pool"; + EXPECT_EQ(node->Backends().All()[0]->Tier(), TierType::SSD); + EXPECT_EQ(node->Backends().Get(TierType::DRAM), nullptr); +} + +TEST_F(MediumSelectionTest, SsdOnlyNodeRoundTripsLocally) { + auto* node = Start(SsdConfig("node-ssd-local")); + ASSERT_NE(node, nullptr); + + std::vector payload(kPageSize, 0x5A); + ASSERT_TRUE(node->Put("ssd-local-key", payload.data(), payload.size())); + + std::vector readback(kPageSize, 0); + ASSERT_TRUE(node->Get("ssd-local-key", readback.data(), readback.size())); + EXPECT_EQ(std::memcmp(payload.data(), readback.data(), payload.size()), 0); +} + +// The regression that motivated the tier filter's removal: master routes the +// batch to the only node with capacity, whose only medium is SSD. The old +// DRAM/HBM allowlist in PartitionBatchPutTargets dropped exactly these items, +// so every key came back false with nothing on the wire. +TEST_F(MediumSelectionTest, CrossNodeBatchPutLandsOnSsdPeer) { + // One page of DRAM ~= no capacity, so the batch must route to the peer. + auto* caller = Start(DramConfig("node-caller", kPageSize)); + ASSERT_NE(caller, nullptr); + auto* target = Start(SsdConfig("node-ssd-target")); + ASSERT_NE(target, nullptr); + + constexpr size_t kKeys = 4; + std::vector keys; + std::vector> buffers; + std::vector srcs; + std::vector sizes; + for (size_t i = 0; i < kKeys; ++i) { + keys.push_back("ssd-remote-" + std::to_string(i)); + buffers.emplace_back(kPageSize, static_cast(0x20 + i)); + sizes.push_back(kPageSize); + } + for (auto& b : buffers) srcs.push_back(b.data()); + + auto put_results = caller->BatchPut(keys, srcs, sizes); + ASSERT_EQ(put_results.size(), kKeys); + for (size_t i = 0; i < kKeys; ++i) { + EXPECT_TRUE(put_results[i]) << "key " << keys[i] << " was not placed on the SSD peer"; + } + + // And it really is on the peer's SSD, readable back over the wire. Retried: + // the peer publishes its committed keys to master through the heartbeat, so a + // read issued immediately after the put can legitimately find no route yet. + std::vector> reads(kKeys, std::vector(kPageSize, 0)); + std::vector dsts; + for (auto& r : reads) dsts.push_back(r.data()); + std::vector get_results; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(20); + do { + get_results = caller->BatchGet(keys, dsts, sizes); + ASSERT_EQ(get_results.size(), kKeys); + bool all = true; + for (size_t i = 0; i < kKeys; ++i) all = all && get_results[i]; + if (all) break; + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } while (std::chrono::steady_clock::now() < deadline); + + for (size_t i = 0; i < kKeys; ++i) { + EXPECT_TRUE(get_results[i]) << "key " << keys[i] << " unreadable from the SSD peer"; + if (get_results[i]) { + EXPECT_EQ(std::memcmp(buffers[i].data(), reads[i].data(), kPageSize), 0); + } + } +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_no_master.cpp b/tests/cpp/umbp/distributed/test_no_master.cpp new file mode 100644 index 000000000..f1317117e --- /dev/null +++ b/tests/cpp/umbp/distributed/test_no_master.cpp @@ -0,0 +1,290 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// The distributed client with NO master configured. +// +// An empty master_address is the single-node deployment: the whole data plane +// -- backends, transfer engine, peer service -- is built exactly as it would be +// in a cluster, but nothing is routed, registered or heartbeated, and a local +// miss is the final answer rather than a question for someone else. +// +// The point of these tests is that this is the SAME client, not a fallback to +// local mode: GetDeploymentMode() still reports Distributed, and the config +// alone decides which it is. No master is started anywhere in this file, so +// any attempt to reach one shows up as a hang or a crash rather than passing +// quietly. + +#include +#include + +#include +#include +#include +#include +#include + +#include "umbp/distributed/config.h" +#include "umbp/distributed/pool_client.h" +#include "umbp/umbp_client.h" + +namespace mori::umbp { +namespace { + +constexpr size_t kPageSize = 4096; +constexpr size_t kDramCap = 32ULL * 1024 * 1024; +constexpr size_t kObjectSize = 8192; + +uint16_t NextPeerServicePort() { + static std::atomic next{ + static_cast(47000 + (static_cast(::getpid()) % 3000))}; + return next.fetch_add(1); +} + +std::vector Pattern(size_t size, int seed) { + std::vector out(size); + for (size_t i = 0; i < size; ++i) out[i] = static_cast((i * 37 + seed) & 0xFF); + return out; +} + +class NoMasterTest : public ::testing::Test { + protected: + // Deliberately identical to a clustered config except for the one empty + // field, so a failure here cannot be blamed on some other difference. + PoolClientConfig Config(const std::string& node_id) { + PoolClientConfig cfg; + cfg.master_config.node_id = node_id; + cfg.master_config.node_address = "127.0.0.1"; + cfg.master_config.master_address = ""; // <-- the whole difference + cfg.io_engine.host = "0.0.0.0"; + cfg.io_engine.port = 0; + cfg.peer_service_port = NextPeerServicePort(); + cfg.dram_page_size = kPageSize; + cfg.medium = TierType::DRAM; + cfg.dram.buffer_sizes = {kDramCap}; + cfg.cache_remote_fetches = false; + return cfg; + } + + PoolClient* Start(PoolClientConfig cfg) { + auto client = std::make_unique(std::move(cfg)); + if (!client->Init()) return nullptr; + clients_.push_back(std::move(client)); + return clients_.back().get(); + } + + void TearDown() override { + for (auto& c : clients_) { + if (c) c->Shutdown(); + } + clients_.clear(); + } + + std::vector> clients_; +}; + +TEST_F(NoMasterTest, InitSucceedsAndReportsNoMaster) { + auto* node = Start(Config("no-master-init")); + ASSERT_NE(node, nullptr) << "Init must not require a master"; + EXPECT_FALSE(node->HasMaster()); +} + +TEST_F(NoMasterTest, PutGetRoundTrip) { + auto* node = Start(Config("no-master-rt")); + ASSERT_NE(node, nullptr); + + const std::vector keys = {"nm-a", "nm-b", "nm-c"}; + std::vector> objects; + std::vector srcs; + std::vector sizes; + for (size_t i = 0; i < keys.size(); ++i) { + objects.push_back(Pattern(kObjectSize, static_cast(i))); + srcs.push_back(objects.back().data()); + sizes.push_back(kObjectSize); + } + ASSERT_EQ(node->BatchPut(keys, srcs, sizes), std::vector(keys.size(), true)) + << "a put with no master must place on this node rather than fail to route"; + + std::vector> reads(keys.size(), std::vector(kObjectSize, 0)); + std::vector dsts; + for (auto& r : reads) dsts.push_back(r.data()); + EXPECT_EQ(node->BatchGet(keys, dsts, sizes), std::vector(keys.size(), true)); + for (size_t i = 0; i < keys.size(); ++i) { + EXPECT_EQ(std::memcmp(reads[i].data(), objects[i].data(), kObjectSize), 0) + << "wrong bytes for " << keys[i]; + } +} + +TEST_F(NoMasterTest, ExistsAnswersFromTheLocalMedia) { + auto* node = Start(Config("no-master-exists")); + ASSERT_NE(node, nullptr); + + const std::vector present = {"nme-a", "nme-b"}; + const auto object = Pattern(kObjectSize, 5); + ASSERT_EQ(node->BatchPut(present, {object.data(), object.data()}, {kObjectSize, kObjectSize}), + std::vector(present.size(), true)); + + EXPECT_EQ(node->BatchExists(present), std::vector(present.size(), true)); + EXPECT_TRUE(node->Exists(present[0])); + + // The half that would otherwise have gone to the master. With no cluster + // behind it, "this node does not have it" IS the final answer -- and it must + // be false, not a fabricated hit and not a hang. + const std::vector mixed = {"nme-a", "nme-never-put"}; + EXPECT_EQ(node->BatchExists(mixed), std::vector({true, false})); +} + +TEST_F(NoMasterTest, MissingKeyIsAMissNotAnError) { + auto* node = Start(Config("no-master-miss")); + ASSERT_NE(node, nullptr); + + const auto object = Pattern(kObjectSize, 9); + ASSERT_EQ(node->BatchPut({"nmm-here"}, {object.data()}, {kObjectSize}), + std::vector({true})); + + std::vector hit(kObjectSize, 0); + std::vector miss(kObjectSize, 0); + const std::vector dsts = {hit.data(), miss.data()}; + EXPECT_EQ(node->BatchGet({"nmm-here", "nmm-absent"}, dsts, {kObjectSize, kObjectSize}), + std::vector({true, false})); + EXPECT_EQ(std::memcmp(hit.data(), object.data(), kObjectSize), 0); +} + +TEST_F(NoMasterTest, RangedGetServesFromTheLocalMedium) { + auto* node = Start(Config("no-master-ranged")); + ASSERT_NE(node, nullptr); + + const std::string key = "nmr-obj"; + const auto object = Pattern(kObjectSize, 13); + ASSERT_EQ(node->BatchPut({key}, {object.data()}, {kObjectSize}), std::vector({true})); + + // Two disjoint windows, one at each end of the object. + const size_t kFront = 512; + const size_t kTail = 256; + std::vector front(kFront, 0); + std::vector tail(kTail, 0); + const std::vector> dsts = {{front.data(), tail.data()}}; + const std::vector> sizes = {{kFront, kTail}}; + const std::vector> offsets = {{0, kObjectSize - kTail}}; + + EXPECT_EQ(node->BatchGetRanges({key}, dsts, sizes, offsets), std::vector({true})); + EXPECT_EQ(std::memcmp(front.data(), object.data(), kFront), 0); + EXPECT_EQ(std::memcmp(tail.data(), object.data() + kObjectSize - kTail, kTail), 0); + + // A key nothing here holds cannot be served and has nowhere else to come + // from; the batch must report it missing rather than route or hang. + std::vector unused(kFront, 0); + const std::vector> miss_dsts = {{unused.data()}}; + EXPECT_EQ(node->BatchGetRanges({"nmr-absent"}, miss_dsts, {{kFront}}, {{0}}), + std::vector({false})); +} + +TEST_F(NoMasterTest, RangedPutAssemblesLocally) { + auto* node = Start(Config("no-master-ranged-put")); + ASSERT_NE(node, nullptr); + + const std::string key = "nmrp-obj"; + const auto object = Pattern(kObjectSize, 17); + const size_t half = kObjectSize / 2; + const std::vector> srcs = {{object.data(), object.data() + half}}; + const std::vector> sizes = {{half, kObjectSize - half}}; + const std::vector> offsets = {{0, half}}; + + ASSERT_EQ(node->BatchPutRanges({key}, {kObjectSize}, srcs, sizes, offsets), + std::vector({true})) + << "a ranged put with no master must place on this node"; + + std::vector read(kObjectSize, 0); + const std::vector dsts = {read.data()}; + EXPECT_EQ(node->BatchGet({key}, dsts, {kObjectSize}), std::vector({true})); + EXPECT_EQ(std::memcmp(read.data(), object.data(), kObjectSize), 0); +} + +// External-KV state lives only in the master's index. With no master these +// must degrade to a clean no-op: reporting succeeds vacuously and a match finds +// nothing. A crash or a false positive here would break the sglang tree +// connector, which calls them unconditionally. +TEST_F(NoMasterTest, ExternalKvDegradesToNoOp) { + auto* node = Start(Config("no-master-extkv")); + ASSERT_NE(node, nullptr); + + const std::vector hashes = {"h1", "h2"}; + EXPECT_TRUE(node->ReportExternalKvBlocks(hashes, TierType::DRAM)); + EXPECT_TRUE(node->RevokeExternalKvBlocks(hashes, TierType::DRAM)); + EXPECT_TRUE(node->RevokeAllExternalKvBlocksAtTier(TierType::DRAM)); + + std::vector matches; + EXPECT_TRUE(node->MatchExternalKv(hashes, &matches, /*count_as_hit=*/false)); + EXPECT_TRUE(matches.empty()); +} + +TEST_F(NoMasterTest, ClearAndFlushSucceed) { + auto* node = Start(Config("no-master-lifecycle")); + ASSERT_NE(node, nullptr); + + const auto object = Pattern(kObjectSize, 21); + ASSERT_EQ(node->BatchPut({"nml-a"}, {object.data()}, {kObjectSize}), std::vector({true})); + EXPECT_TRUE(node->Clear()) << "Clear must converge vacuously with no master to sync to"; + EXPECT_EQ(node->BatchExists({"nml-a"}), std::vector({false})); +} + +// The same process, the same client class, both deployments -- which is the +// property that makes single-node and clustered a config choice rather than a +// code path. Runs them in sequence so a leak of one into the other shows up. +TEST_F(NoMasterTest, SameClientClassServesBothDeployments) { + auto* solo = Start(Config("no-master-mode")); + ASSERT_NE(solo, nullptr); + EXPECT_FALSE(solo->HasMaster()); + + UMBPConfig cfg; + cfg.dram.capacity_bytes = kDramCap; + UMBPDistributedConfig dc; + dc.master_config.node_id = "no-master-client"; + dc.master_config.node_address = "127.0.0.1"; + dc.master_config.master_address = ""; + dc.io_engine.host = "0.0.0.0"; + dc.io_engine.port = 0; + dc.peer_service_port = NextPeerServicePort(); + dc.dram_page_size = kPageSize; + cfg.distributed = dc; + + std::string error; + ASSERT_TRUE(cfg.Validate(&error)) << "an empty master address must validate: " << error; + + auto client = CreateUMBPClient(cfg); + ASSERT_NE(client, nullptr); + EXPECT_EQ(client->GetDeploymentMode(), UMBPDeploymentMode::Distributed) + << "no master must not silently downgrade to the Local backend"; + EXPECT_TRUE(client->IsDistributed()); + + const auto object = Pattern(kObjectSize, 23); + std::vector read(kObjectSize, 0); + const auto src = reinterpret_cast(object.data()); + const auto dst = reinterpret_cast(read.data()); + ASSERT_EQ(client->BatchPut({"nmc-a"}, {src}, {kObjectSize}), std::vector({true})); + EXPECT_EQ(client->BatchGet({"nmc-a"}, {dst}, {kObjectSize}), std::vector({true})); + EXPECT_EQ(std::memcmp(read.data(), object.data(), kObjectSize), 0); + EXPECT_EQ(client->BatchExistsConsecutive({"nmc-a", "nmc-absent"}), 1u); + client->Close(); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_peer_dram_allocator.cpp b/tests/cpp/umbp/distributed/test_page_backend.cpp similarity index 62% rename from tests/cpp/umbp/distributed/test_peer_dram_allocator.cpp rename to tests/cpp/umbp/distributed/test_page_backend.cpp index da2ac1a66..bbe4d18ba 100644 --- a/tests/cpp/umbp/distributed/test_peer_dram_allocator.cpp +++ b/tests/cpp/umbp/distributed/test_page_backend.cpp @@ -29,7 +29,7 @@ #include #include -#include "umbp/distributed/peer/peer_dram_allocator.h" +#include "umbp/distributed/peer/backend/page_backend.h" namespace mori::umbp { @@ -38,37 +38,40 @@ namespace { // 3 buffers x 4 pages of 1 KiB = 12 KiB total DRAM. constexpr uint64_t kPageSize = 1024; -PeerDramAllocator::TierConfig MakeDramCfg() { - PeerDramAllocator::TierConfig cfg; +PageBackend::TierConfig MakeDramCfg() { + PageBackend::TierConfig cfg; cfg.buffer_sizes = {kPageSize * 4, kPageSize * 4, kPageSize * 4}; cfg.buffer_descs = {{0xA0, 0xA1}, {0xB0, 0xB1}, {0xC0, 0xC1}}; return cfg; } -PeerDramAllocator::TierConfig EmptyCfg() { return {}; } - -std::unique_ptr MakeAllocator( +std::unique_ptr MakeAllocator( std::chrono::milliseconds pending_ttl = std::chrono::milliseconds{5000}, std::chrono::milliseconds read_lease_ttl = std::chrono::milliseconds{500}) { - return std::make_unique(kPageSize, MakeDramCfg(), EmptyCfg(), pending_ttl, - read_lease_ttl); + return std::make_unique(TierType::DRAM, kPageSize, MakeDramCfg(), pending_ttl, + read_lease_ttl); } -// Strip AllocateResult down to its slot for tests that don't exercise -// the dedup outcome. -std::optional AllocateOk(PeerDramAllocator& a, - const std::string& key, uint64_t size, - TierType tier) { - return a.Allocate(key, size, tier).slot; +// True iff `r` carries a live allocated slot (mirrors the pre-refactor +// AllocateResult::slot.has_value(), now folded into AllocateResult's own +// fields — see medium_backend.h). +bool HasSlot(const AllocateResult& r) { return r.outcome == AllocateOutcome::kSuccessAllocated; } + +// Strip AllocateResult down to itself for tests that don't exercise the +// dedup outcome — nullopt unless the allocation actually landed a slot. +std::optional AllocateOk(PageBackend& a, const std::string& key, uint64_t size) { + auto r = a.Allocate(key, size); + if (!HasSlot(r)) return std::nullopt; + return r; } } // namespace // ---- Allocate / Commit / Resolve happy path --------------------------------- -TEST(PeerDramAllocator, CommitMakesKeyResolvable) { +TEST(PageBackend, CommitMakesKeyResolvable) { auto a = MakeAllocator(); - auto pending = AllocateOk(*a, "key-1", kPageSize, TierType::DRAM); + auto pending = AllocateOk(*a, "key-1", kPageSize); ASSERT_TRUE(pending.has_value()); EXPECT_EQ(pending->size, kPageSize); EXPECT_EQ(pending->pages.size(), 1u); @@ -79,7 +82,6 @@ TEST(PeerDramAllocator, CommitMakesKeyResolvable) { auto r = a->Resolve("key-1"); EXPECT_TRUE(r.found); EXPECT_EQ(r.size, kPageSize); - EXPECT_EQ(r.tier, TierType::DRAM); EXPECT_EQ(r.pages, pending->pages); auto events = a->DrainPendingEvents(); @@ -93,54 +95,54 @@ TEST(PeerDramAllocator, CommitMakesKeyResolvable) { // ---- Allocate-side dedup ---------------------------------------------------- // Defensive layer for master-index lag (primary dedup is at BatchRoutePut). -TEST(PeerDramAllocator, AllocateRejectsAlreadyOwnedKey) { +TEST(PageBackend, AllocateRejectsAlreadyOwnedKey) { auto a = MakeAllocator(); - auto first = AllocateOk(*a, "A", kPageSize, TierType::DRAM); + auto first = AllocateOk(*a, "A", kPageSize); ASSERT_TRUE(first.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(first->slot_id, "A", committed_bytes)); a->DrainPendingEvents(); - const auto cap_after_commit = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_after_commit = a->Capacity(); - auto second = a->Allocate("A", kPageSize, TierType::DRAM); - EXPECT_EQ(second.outcome, PeerDramAllocator::Outcome::kSuccessAlreadyExists); - EXPECT_FALSE(second.slot.has_value()); + auto second = a->Allocate("A", kPageSize); + EXPECT_EQ(second.outcome, AllocateOutcome::kSuccessAlreadyExists); + EXPECT_FALSE(HasSlot(second)); // No pages reserved -> capacity unchanged. - const auto cap_after_dedup = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_after_dedup = a->Capacity(); EXPECT_EQ(cap_after_dedup.available_bytes, cap_after_commit.available_bytes); } -TEST(PeerDramAllocator, AllocateAllowsDifferentKey) { +TEST(PageBackend, AllocateAllowsDifferentKey) { auto a = MakeAllocator(); - auto first = AllocateOk(*a, "A", kPageSize, TierType::DRAM); + auto first = AllocateOk(*a, "A", kPageSize); ASSERT_TRUE(first.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(first->slot_id, "A", committed_bytes)); - auto second = a->Allocate("B", kPageSize, TierType::DRAM); - EXPECT_EQ(second.outcome, PeerDramAllocator::Outcome::kSuccessAllocated); - ASSERT_TRUE(second.slot.has_value()); - EXPECT_TRUE(a->Commit(second.slot->slot_id, "B", committed_bytes)); + auto second = a->Allocate("B", kPageSize); + EXPECT_EQ(second.outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_TRUE(HasSlot(second)); + EXPECT_TRUE(a->Commit(second.slot_id, "B", committed_bytes)); } // Lax mode: pending_ not checked. Two same-key Allocates before any // Commit both succeed; race absorbed by Commit() (see // DuplicateCommitIsIdempotentAndKeepsFirst). -TEST(PeerDramAllocator, AllocateDoesNotRejectOnPendingDuplicate) { +TEST(PageBackend, AllocateDoesNotRejectOnPendingDuplicate) { auto a = MakeAllocator(); - auto first = a->Allocate("A", kPageSize, TierType::DRAM); - EXPECT_EQ(first.outcome, PeerDramAllocator::Outcome::kSuccessAllocated); - ASSERT_TRUE(first.slot.has_value()); + auto first = a->Allocate("A", kPageSize); + EXPECT_EQ(first.outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_TRUE(HasSlot(first)); - auto second = a->Allocate("A", kPageSize, TierType::DRAM); - EXPECT_EQ(second.outcome, PeerDramAllocator::Outcome::kSuccessAllocated); - ASSERT_TRUE(second.slot.has_value()); - ASSERT_NE(second.slot->slot_id, first.slot->slot_id); + auto second = a->Allocate("A", kPageSize); + EXPECT_EQ(second.outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_TRUE(HasSlot(second)); + ASSERT_NE(second.slot_id, first.slot_id); } // ---- Duplicate Commit idempotency ------------------------------------------- @@ -148,12 +150,12 @@ TEST(PeerDramAllocator, AllocateDoesNotRejectOnPendingDuplicate) { // Commit — once owned_["dup-key"] is set, the new owned_-check in // Allocate would reject the second slot before it could reach Commit. -TEST(PeerDramAllocator, DuplicateCommitIsIdempotentAndKeepsFirst) { +TEST(PageBackend, DuplicateCommitIsIdempotentAndKeepsFirst) { auto a = MakeAllocator(); - auto first = AllocateOk(*a, "dup-key", kPageSize, TierType::DRAM); + auto first = AllocateOk(*a, "dup-key", kPageSize); ASSERT_TRUE(first.has_value()); - auto second = AllocateOk(*a, "dup-key", kPageSize, TierType::DRAM); + auto second = AllocateOk(*a, "dup-key", kPageSize); ASSERT_TRUE(second.has_value()); ASSERT_NE(second->slot_id, first->slot_id); @@ -169,7 +171,7 @@ TEST(PeerDramAllocator, DuplicateCommitIsIdempotentAndKeepsFirst) { EXPECT_EQ(events[0].key, "dup-key"); // First owned (1 page) + second still pending (1 page) = 2 occupied. - const auto cap_after_first_commit = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_after_first_commit = a->Capacity(); EXPECT_EQ(cap_after_first_commit.available_bytes, cap_after_first_commit.total_bytes - 2 * kPageSize); @@ -189,7 +191,7 @@ TEST(PeerDramAllocator, DuplicateCommitIsIdempotentAndKeepsFirst) { EXPECT_EQ(r.size, kPageSize); // Second slot's pages freed -> only first occupies (1 page). - const auto cap_after_dup = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_after_dup = a->Capacity(); EXPECT_EQ(cap_after_dup.available_bytes, cap_after_dup.total_bytes - kPageSize); EXPECT_EQ(cap_after_dup.total_bytes, cap_after_first_commit.total_bytes); @@ -200,31 +202,26 @@ TEST(PeerDramAllocator, DuplicateCommitIsIdempotentAndKeepsFirst) { // ---- ENOSPC ----------------------------------------------------------------- -TEST(PeerDramAllocator, AllocateReturnsNulloptWhenFull) { +TEST(PageBackend, AllocateReturnsNulloptWhenFull) { auto a = MakeAllocator(); std::vector slot_ids; for (int i = 0; i < 12; ++i) { - auto p = AllocateOk(*a, "k-" + std::to_string(i), kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "k-" + std::to_string(i), kPageSize); ASSERT_TRUE(p.has_value()) << "i=" << i; slot_ids.push_back(p->slot_id); } - EXPECT_FALSE(AllocateOk(*a, "k-overflow", kPageSize, TierType::DRAM).has_value()); + EXPECT_FALSE(AllocateOk(*a, "k-overflow", kPageSize).has_value()); EXPECT_TRUE(a->Abort(slot_ids.back())); - EXPECT_TRUE(AllocateOk(*a, "k-recovered", kPageSize, TierType::DRAM).has_value()); -} - -TEST(PeerDramAllocator, UnconfiguredTierReturnsNullopt) { - auto a = MakeAllocator(); - EXPECT_FALSE(AllocateOk(*a, "k", kPageSize, TierType::HBM).has_value()); + EXPECT_TRUE(AllocateOk(*a, "k-recovered", kPageSize).has_value()); } // ---- Pending TTL ------------------------------------------------------------ -TEST(PeerDramAllocator, PendingSlotExpiresAfterTtl) { - auto a = std::make_unique(kPageSize, MakeDramCfg(), EmptyCfg(), - /*pending_ttl=*/std::chrono::milliseconds{1}); - auto pending = AllocateOk(*a, "key-late", kPageSize, TierType::DRAM); +TEST(PageBackend, PendingSlotExpiresAfterTtl) { + auto a = std::make_unique(TierType::DRAM, kPageSize, MakeDramCfg(), + /*pending_ttl=*/std::chrono::milliseconds{1}); + auto pending = AllocateOk(*a, "key-late", kPageSize); ASSERT_TRUE(pending.has_value()); std::this_thread::sleep_for(std::chrono::milliseconds{20}); @@ -235,15 +232,15 @@ TEST(PeerDramAllocator, PendingSlotExpiresAfterTtl) { EXPECT_EQ(committed_bytes, 0u); EXPECT_TRUE(a->DrainPendingEvents().empty()); - auto cap = a->TierCapacitiesSnapshot(); - EXPECT_EQ(cap[TierType::DRAM].available_bytes, cap[TierType::DRAM].total_bytes); + auto cap = a->Capacity(); + EXPECT_EQ(cap.available_bytes, cap.total_bytes); } // ---- Abort idempotency ------------------------------------------------------ -TEST(PeerDramAllocator, AbortIsIdempotent) { +TEST(PageBackend, AbortIsIdempotent) { auto a = MakeAllocator(); - auto pending = AllocateOk(*a, "k", kPageSize, TierType::DRAM); + auto pending = AllocateOk(*a, "k", kPageSize); ASSERT_TRUE(pending.has_value()); EXPECT_TRUE(a->Abort(pending->slot_id)); EXPECT_TRUE(a->Abort(pending->slot_id)); @@ -253,9 +250,9 @@ TEST(PeerDramAllocator, AbortIsIdempotent) { // ---- Evict idempotency + REMOVE event --------------------------------------- -TEST(PeerDramAllocator, EvictRemovesKeyAndQueuesEvent) { +TEST(PageBackend, EvictRemovesKeyAndQueuesEvent) { auto a = MakeAllocator(); - auto p = AllocateOk(*a, "k", kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "k", kPageSize); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, "k", committed_bytes)); EXPECT_EQ(committed_bytes, p->size); @@ -282,11 +279,11 @@ TEST(PeerDramAllocator, EvictRemovesKeyAndQueuesEvent) { // ---- Resolve-during-Evict race --------------------------------------------- -TEST(PeerDramAllocator, EvictDefersWhenReadLeaseActive) { - auto a = std::make_unique(kPageSize, MakeDramCfg(), EmptyCfg(), - /*pending_ttl=*/std::chrono::milliseconds{5000}, - /*read_lease_ttl=*/std::chrono::milliseconds{200}); - auto p = AllocateOk(*a, "k", kPageSize, TierType::DRAM); +TEST(PageBackend, EvictDefersWhenReadLeaseActive) { + auto a = std::make_unique(TierType::DRAM, kPageSize, MakeDramCfg(), + /*pending_ttl=*/std::chrono::milliseconds{5000}, + /*read_lease_ttl=*/std::chrono::milliseconds{200}); + auto p = AllocateOk(*a, "k", kPageSize); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, "k", committed_bytes)); EXPECT_EQ(committed_bytes, p->size); @@ -311,11 +308,11 @@ TEST(PeerDramAllocator, EvictDefersWhenReadLeaseActive) { // ---- Full-sync snapshot ----------------------------------------------------- -TEST(PeerDramAllocator, SnapshotOwnedKeysReturnsEveryAdd) { +TEST(PageBackend, SnapshotOwnedKeysReturnsEveryAdd) { auto a = MakeAllocator(); for (int i = 0; i < 5; ++i) { const std::string k = "k-" + std::to_string(i); - auto p = AllocateOk(*a, k, kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, k, kPageSize); ASSERT_TRUE(p.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, k, committed_bytes)); @@ -335,13 +332,13 @@ TEST(PeerDramAllocator, SnapshotOwnedKeysReturnsEveryAdd) { // ---- Buffer descs filtered to the page set --------------------------------- -TEST(PeerDramAllocator, BufferDescsForPagesDedupAndOrder) { +TEST(PageBackend, BufferDescsForPagesDedupAndOrder) { auto a = MakeAllocator(); - auto p = AllocateOk(*a, "k", kPageSize * 5, TierType::DRAM); + auto p = AllocateOk(*a, "k", kPageSize * 5); ASSERT_TRUE(p.has_value()); ASSERT_EQ(p->pages.size(), 5u); - auto descs = a->BufferDescsForPages(TierType::DRAM, p->pages); + auto descs = a->BufferDescsForPages(p->pages); ASSERT_EQ(descs.size(), 2u); EXPECT_EQ(descs[0].buffer_index, 0u); EXPECT_EQ(descs[1].buffer_index, 1u); @@ -351,66 +348,63 @@ TEST(PeerDramAllocator, BufferDescsForPagesDedupAndOrder) { // ---- BatchAllocate / BatchCommit / BatchAbort ------------------------------- -TEST(PeerDramAllocator, BatchAllocateEmptyInputReturnsEmpty) { +TEST(PageBackend, BatchAllocateEmptyInputReturnsEmpty) { auto a = MakeAllocator(); EXPECT_TRUE(a->BatchAllocate({}).empty()); } -TEST(PeerDramAllocator, BatchAllocateMixedOutcomesAndDescs) { +TEST(PageBackend, BatchAllocateMixedOutcomesAndDescs) { auto a = MakeAllocator(); - auto owned = AllocateOk(*a, "owned", kPageSize, TierType::DRAM); + auto owned = AllocateOk(*a, "owned", kPageSize); ASSERT_TRUE(owned.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(owned->slot_id, "owned", committed_bytes)); a->DrainPendingEvents(); - std::vector requests; - requests.push_back({"owned", kPageSize, TierType::DRAM}); - requests.push_back({"ok", kPageSize * 5, TierType::DRAM}); - requests.push_back({"bad-tier", kPageSize, TierType::HBM}); - requests.push_back({"zero", 0, TierType::DRAM}); - requests.push_back({"too-big", kPageSize * 20, TierType::DRAM}); + std::vector requests; + requests.push_back({"owned", kPageSize}); + requests.push_back({"ok", kPageSize * 5}); + requests.push_back({"zero", 0}); + requests.push_back({"too-big", kPageSize * 20}); auto results = a->BatchAllocate(requests); ASSERT_EQ(results.size(), requests.size()); - EXPECT_EQ(results[0].outcome, PeerDramAllocator::Outcome::kSuccessAlreadyExists); - EXPECT_FALSE(results[0].slot.has_value()); + EXPECT_EQ(results[0].outcome, AllocateOutcome::kSuccessAlreadyExists); + EXPECT_FALSE(HasSlot(results[0])); EXPECT_TRUE(results[0].descs.empty()); - EXPECT_EQ(results[1].outcome, PeerDramAllocator::Outcome::kSuccessAllocated); - ASSERT_TRUE(results[1].slot.has_value()); - EXPECT_EQ(results[1].slot->size, kPageSize * 5); - EXPECT_EQ(results[1].slot->pages.size(), 5u); + EXPECT_EQ(results[1].outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_TRUE(HasSlot(results[1])); + EXPECT_EQ(results[1].size, kPageSize * 5); + EXPECT_EQ(results[1].pages.size(), 5u); ASSERT_EQ(results[1].descs.size(), 2u); EXPECT_EQ(results[1].descs[0].buffer_index, 0u); EXPECT_EQ(results[1].descs[1].buffer_index, 1u); - EXPECT_EQ(results[2].outcome, PeerDramAllocator::Outcome::kFailed); - EXPECT_FALSE(results[2].slot.has_value()); - EXPECT_EQ(results[3].outcome, PeerDramAllocator::Outcome::kFailed); - EXPECT_FALSE(results[3].slot.has_value()); - EXPECT_EQ(results[4].outcome, PeerDramAllocator::Outcome::kFailedNoSpace); - EXPECT_FALSE(results[4].slot.has_value()); + EXPECT_EQ(results[2].outcome, AllocateOutcome::kFailed); + EXPECT_FALSE(HasSlot(results[2])); + EXPECT_EQ(results[3].outcome, AllocateOutcome::kFailedNoSpace); + EXPECT_FALSE(HasSlot(results[3])); } -TEST(PeerDramAllocator, BatchCommitMixedSuccessAndFailure) { +TEST(PageBackend, BatchCommitMixedSuccessAndFailure) { auto a = MakeAllocator(); auto allocated = a->BatchAllocate({ - {"dup", kPageSize, TierType::DRAM}, - {"dup", kPageSize * 2, TierType::DRAM}, - {"unique", kPageSize, TierType::DRAM}, + {"dup", kPageSize}, + {"dup", kPageSize * 2}, + {"unique", kPageSize}, }); ASSERT_EQ(allocated.size(), 3u); - ASSERT_TRUE(allocated[0].slot.has_value()); - ASSERT_TRUE(allocated[1].slot.has_value()); - ASSERT_TRUE(allocated[2].slot.has_value()); + ASSERT_TRUE(HasSlot(allocated[0])); + ASSERT_TRUE(HasSlot(allocated[1])); + ASSERT_TRUE(HasSlot(allocated[2])); auto committed = a->BatchCommit({ - {allocated[0].slot->slot_id, "dup"}, + {allocated[0].slot_id, "dup"}, {999999, "missing"}, - {allocated[1].slot->slot_id, "dup"}, - {allocated[2].slot->slot_id, "unique"}, + {allocated[1].slot_id, "dup"}, + {allocated[2].slot_id, "unique"}, }); ASSERT_EQ(committed.size(), 4u); EXPECT_TRUE(committed[0].success); @@ -424,7 +418,7 @@ TEST(PeerDramAllocator, BatchCommitMixedSuccessAndFailure) { auto dup = a->Resolve("dup"); ASSERT_TRUE(dup.found); - EXPECT_EQ(dup.pages, allocated[0].slot->pages); + EXPECT_EQ(dup.pages, allocated[0].pages); EXPECT_EQ(dup.size, kPageSize); auto unique = a->Resolve("unique"); ASSERT_TRUE(unique.found); @@ -438,59 +432,58 @@ TEST(PeerDramAllocator, BatchCommitMixedSuccessAndFailure) { EXPECT_EQ(events[1].key, "unique"); } -TEST(PeerDramAllocator, BatchAbortMixedSlotsIsIdempotent) { +TEST(PageBackend, BatchAbortMixedSlotsIsIdempotent) { auto a = MakeAllocator(); auto allocated = a->BatchAllocate({ - {"drop", kPageSize, TierType::DRAM}, - {"keep", kPageSize, TierType::DRAM}, + {"drop", kPageSize}, + {"keep", kPageSize}, }); ASSERT_EQ(allocated.size(), 2u); - ASSERT_TRUE(allocated[0].slot.has_value()); - ASSERT_TRUE(allocated[1].slot.has_value()); + ASSERT_TRUE(HasSlot(allocated[0])); + ASSERT_TRUE(HasSlot(allocated[1])); - auto aborted = a->BatchAbort({allocated[0].slot->slot_id, 999999}); + auto aborted = a->BatchAbort({allocated[0].slot_id, 999999}); ASSERT_EQ(aborted.size(), 2u); EXPECT_TRUE(aborted[0]); EXPECT_TRUE(aborted[1]); uint64_t committed_bytes = 0; - EXPECT_FALSE(a->Commit(allocated[0].slot->slot_id, "drop", committed_bytes)); - EXPECT_TRUE(a->Commit(allocated[1].slot->slot_id, "keep", committed_bytes)); + EXPECT_FALSE(a->Commit(allocated[0].slot_id, "drop", committed_bytes)); + EXPECT_TRUE(a->Commit(allocated[1].slot_id, "keep", committed_bytes)); EXPECT_EQ(committed_bytes, kPageSize); EXPECT_TRUE(a->Resolve("keep").found); } // ---- BatchResolve ---------------------------------------------------------- -TEST(PeerDramAllocator, BatchResolveEmptyInputReturnsEmpty) { +TEST(PageBackend, BatchResolveEmptyInputReturnsEmpty) { auto a = MakeAllocator(); - EXPECT_TRUE(a->BatchResolve({}).empty()); + EXPECT_TRUE(a->BatchResolve({}, /*include_descs=*/true).empty()); } -TEST(PeerDramAllocator, BatchResolveMixedHitsAndMisses) { +TEST(PageBackend, BatchResolveMixedHitsAndMisses) { auto a = MakeAllocator(); // 5 pages over 4-pages-per-buffer config -> exercises dedup'd descs. - auto p_hit = AllocateOk(*a, "hit", kPageSize * 5, TierType::DRAM); + auto p_hit = AllocateOk(*a, "hit", kPageSize * 5); ASSERT_TRUE(p_hit.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p_hit->slot_id, "hit", committed_bytes)); - auto p_small = AllocateOk(*a, "small", kPageSize, TierType::DRAM); + auto p_small = AllocateOk(*a, "small", kPageSize); ASSERT_TRUE(p_small.has_value()); ASSERT_TRUE(a->Commit(p_small->slot_id, "small", committed_bytes)); a->DrainPendingEvents(); auto ref_hit = a->Resolve("hit"); - auto ref_descs_hit = a->BufferDescsForPages(ref_hit.tier, ref_hit.pages); + auto ref_descs_hit = a->BufferDescsForPages(ref_hit.pages); auto ref_small = a->Resolve("small"); - auto ref_descs_small = a->BufferDescsForPages(ref_small.tier, ref_small.pages); + auto ref_descs_small = a->BufferDescsForPages(ref_small.pages); ASSERT_TRUE(ref_hit.found); ASSERT_TRUE(ref_small.found); - auto results = a->BatchResolve({"hit", "ghost-a", "small", "ghost-b"}); + auto results = a->BatchResolve({"hit", "ghost-a", "small", "ghost-b"}, /*include_descs=*/true); ASSERT_EQ(results.size(), 4u); EXPECT_TRUE(results[0].found); - EXPECT_EQ(results[0].tier, ref_hit.tier); EXPECT_EQ(results[0].pages, ref_hit.pages); EXPECT_EQ(results[0].size, ref_hit.size); ASSERT_EQ(results[0].descs.size(), ref_descs_hit.size()); @@ -505,7 +498,6 @@ TEST(PeerDramAllocator, BatchResolveMixedHitsAndMisses) { EXPECT_TRUE(results[1].descs.empty()); EXPECT_TRUE(results[2].found); - EXPECT_EQ(results[2].tier, ref_small.tier); EXPECT_EQ(results[2].pages, ref_small.pages); EXPECT_EQ(results[2].size, ref_small.size); ASSERT_EQ(results[2].descs.size(), ref_descs_small.size()); @@ -517,20 +509,20 @@ TEST(PeerDramAllocator, BatchResolveMixedHitsAndMisses) { EXPECT_FALSE(results[3].found); } -TEST(PeerDramAllocator, BatchResolveExtendsLeaseForHitsOnly) { - auto a = std::make_unique(kPageSize, MakeDramCfg(), EmptyCfg(), - /*pending_ttl=*/std::chrono::milliseconds{5000}, - /*read_lease_ttl=*/std::chrono::milliseconds{500}); - auto p_x = AllocateOk(*a, "x", kPageSize, TierType::DRAM); +TEST(PageBackend, BatchResolveExtendsLeaseForHitsOnly) { + auto a = std::make_unique(TierType::DRAM, kPageSize, MakeDramCfg(), + /*pending_ttl=*/std::chrono::milliseconds{5000}, + /*read_lease_ttl=*/std::chrono::milliseconds{500}); + auto p_x = AllocateOk(*a, "x", kPageSize); ASSERT_TRUE(p_x.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p_x->slot_id, "x", committed_bytes)); - auto p_y = AllocateOk(*a, "y", kPageSize, TierType::DRAM); + auto p_y = AllocateOk(*a, "y", kPageSize); ASSERT_TRUE(p_y.has_value()); ASSERT_TRUE(a->Commit(p_y->slot_id, "y", committed_bytes)); a->DrainPendingEvents(); - auto results = a->BatchResolve({"x", "missing", "y"}); + auto results = a->BatchResolve({"x", "missing", "y"}, /*include_descs=*/true); ASSERT_EQ(results.size(), 3u); ASSERT_TRUE(results[0].found); ASSERT_FALSE(results[1].found); @@ -546,7 +538,7 @@ TEST(PeerDramAllocator, BatchResolveExtendsLeaseForHitsOnly) { // Miss must not poison read_lease_until_: a subsequent // Allocate+Commit+Evict on the same key must free as if never touched. - auto p_miss = AllocateOk(*a, "missing", kPageSize, TierType::DRAM); + auto p_miss = AllocateOk(*a, "missing", kPageSize); ASSERT_TRUE(p_miss.has_value()); ASSERT_TRUE(a->Commit(p_miss->slot_id, "missing", committed_bytes)); a->DrainPendingEvents(); @@ -555,17 +547,17 @@ TEST(PeerDramAllocator, BatchResolveExtendsLeaseForHitsOnly) { EXPECT_EQ(evict_missing[0].bytes_freed, kPageSize); } -TEST(PeerDramAllocator, BatchResolveLeaseExpiresLikeSingleKeyResolve) { - auto a = std::make_unique(kPageSize, MakeDramCfg(), EmptyCfg(), - /*pending_ttl=*/std::chrono::milliseconds{5000}, - /*read_lease_ttl=*/std::chrono::milliseconds{50}); - auto p = AllocateOk(*a, "k", kPageSize, TierType::DRAM); +TEST(PageBackend, BatchResolveLeaseExpiresLikeSingleKeyResolve) { + auto a = std::make_unique(TierType::DRAM, kPageSize, MakeDramCfg(), + /*pending_ttl=*/std::chrono::milliseconds{5000}, + /*read_lease_ttl=*/std::chrono::milliseconds{50}); + auto p = AllocateOk(*a, "k", kPageSize); ASSERT_TRUE(p.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, "k", committed_bytes)); a->DrainPendingEvents(); - auto results = a->BatchResolve({"k"}); + auto results = a->BatchResolve({"k"}, /*include_descs=*/true); ASSERT_EQ(results.size(), 1u); ASSERT_TRUE(results[0].found); @@ -577,37 +569,36 @@ TEST(PeerDramAllocator, BatchResolveLeaseExpiresLikeSingleKeyResolve) { EXPECT_EQ(evicted[0].bytes_freed, kPageSize); } -// ---- Capacities snapshot ---------------------------------------------------- +// ---- Capacity snapshot ---------------------------------------------------- -TEST(PeerDramAllocator, TierCapacitiesReflectAllocations) { +TEST(PageBackend, CapacityReflectsAllocations) { auto a = MakeAllocator(); - auto cap0 = a->TierCapacitiesSnapshot(); - ASSERT_EQ(cap0.count(TierType::DRAM), 1u); - const uint64_t total = cap0[TierType::DRAM].total_bytes; - EXPECT_EQ(cap0[TierType::DRAM].available_bytes, total); + auto cap0 = a->Capacity(); + const uint64_t total = cap0.total_bytes; + EXPECT_EQ(cap0.available_bytes, total); - auto p = AllocateOk(*a, "k", kPageSize * 3, TierType::DRAM); + auto p = AllocateOk(*a, "k", kPageSize * 3); ASSERT_TRUE(p.has_value()); - auto cap1 = a->TierCapacitiesSnapshot(); - EXPECT_EQ(cap1[TierType::DRAM].available_bytes, total - 3 * kPageSize); + auto cap1 = a->Capacity(); + EXPECT_EQ(cap1.available_bytes, total - 3 * kPageSize); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, "k", committed_bytes)); EXPECT_EQ(committed_bytes, p->size); - auto cap2 = a->TierCapacitiesSnapshot(); - EXPECT_EQ(cap2[TierType::DRAM].available_bytes, total - 3 * kPageSize); + auto cap2 = a->Capacity(); + EXPECT_EQ(cap2.available_bytes, total - 3 * kPageSize); ASSERT_EQ(a->Evict({"k"})[0].bytes_freed, 3 * kPageSize); - auto cap3 = a->TierCapacitiesSnapshot(); - EXPECT_EQ(cap3[TierType::DRAM].available_bytes, total); + auto cap3 = a->Capacity(); + EXPECT_EQ(cap3.available_bytes, total); } // ---- Commit after reap ------------------------------------------------------ -TEST(PeerDramAllocator, CommitAfterReapReturnsFalse) { - auto a = std::make_unique(kPageSize, MakeDramCfg(), EmptyCfg(), - std::chrono::milliseconds{1}); - auto p = AllocateOk(*a, "doomed", kPageSize, TierType::DRAM); +TEST(PageBackend, CommitAfterReapReturnsFalse) { + auto a = std::make_unique(TierType::DRAM, kPageSize, MakeDramCfg(), + std::chrono::milliseconds{1}); + auto p = AllocateOk(*a, "doomed", kPageSize); ASSERT_TRUE(p.has_value()); std::this_thread::sleep_for(std::chrono::milliseconds{20}); a->RunReaperOnceForTest(); @@ -619,19 +610,19 @@ TEST(PeerDramAllocator, CommitAfterReapReturnsFalse) { // ---- Distributed Clear ------------------------------------------------------ -TEST(PeerDramAllocator, ClearLocalReleasesOwnedAndCancelsPending) { +TEST(PageBackend, ClearLocalReleasesOwnedAndCancelsPending) { auto a = MakeAllocator(); - auto pA = AllocateOk(*a, "A", kPageSize, TierType::DRAM); + auto pA = AllocateOk(*a, "A", kPageSize); ASSERT_TRUE(pA.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(pA->slot_id, "A", committed_bytes)); - auto pB = AllocateOk(*a, "B", kPageSize * 2, TierType::DRAM); + auto pB = AllocateOk(*a, "B", kPageSize * 2); ASSERT_TRUE(pB.has_value()); a->DrainPendingEvents(); // discard the A ADD - const auto cap_before = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_before = a->Capacity(); ASSERT_LT(cap_before.available_bytes, cap_before.total_bytes); a->ClearLocal(); @@ -642,7 +633,7 @@ TEST(PeerDramAllocator, ClearLocalReleasesOwnedAndCancelsPending) { EXPECT_TRUE(a->DrainPendingEvents().empty()); // Owned pages (A) returned immediately; pending pages (B) still held. - auto cap_after_clear = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_after_clear = a->Capacity(); EXPECT_EQ(cap_after_clear.available_bytes, cap_before.total_bytes - 2 * kPageSize); // Committing the cancelled pending fails AND releases its pages @@ -652,24 +643,24 @@ TEST(PeerDramAllocator, ClearLocalReleasesOwnedAndCancelsPending) { EXPECT_TRUE(a->DrainPendingEvents().empty()); EXPECT_FALSE(a->Resolve("B").found); - auto cap_final = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_final = a->Capacity(); EXPECT_EQ(cap_final.available_bytes, cap_final.total_bytes); } -TEST(PeerDramAllocator, ClearLocalGatesAllocateUntilAcked) { +TEST(PageBackend, ClearLocalGatesAllocateUntilAcked) { auto a = MakeAllocator(); a->ClearLocal(); - EXPECT_FALSE(AllocateOk(*a, "blocked", kPageSize, TierType::DRAM).has_value()); + EXPECT_FALSE(AllocateOk(*a, "blocked", kPageSize).has_value()); a->ClearFullSyncAcked(); EXPECT_FALSE(a->IsClearFullSyncPending()); - EXPECT_TRUE(AllocateOk(*a, "ok-after-ack", kPageSize, TierType::DRAM).has_value()); + EXPECT_TRUE(AllocateOk(*a, "ok-after-ack", kPageSize).has_value()); } -TEST(PeerDramAllocator, ClearLocalDropsQueuedAdds) { +TEST(PageBackend, ClearLocalDropsQueuedAdds) { auto a = MakeAllocator(); - auto p = AllocateOk(*a, "k", kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "k", kPageSize); ASSERT_TRUE(p.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, "k", committed_bytes)); @@ -681,9 +672,9 @@ TEST(PeerDramAllocator, ClearLocalDropsQueuedAdds) { EXPECT_TRUE(a->SnapshotOwnedKeys().empty()); } -TEST(PeerDramAllocator, AbortReleasesCancelledPending) { +TEST(PageBackend, AbortReleasesCancelledPending) { auto a = MakeAllocator(); - auto p = AllocateOk(*a, "p1", kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "p1", kPageSize); ASSERT_TRUE(p.has_value()); a->ClearLocal(); @@ -691,17 +682,17 @@ TEST(PeerDramAllocator, AbortReleasesCancelledPending) { EXPECT_TRUE(a->Abort(p->slot_id)); a->ClearFullSyncAcked(); - auto cap = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap = a->Capacity(); EXPECT_EQ(cap.available_bytes, cap.total_bytes); } // Pre-clear pending Commit fails; post-ack new Allocate+Commit succeeds. -TEST(PeerDramAllocator, PendingGenerationRejectsPreClearCommit) { +TEST(PageBackend, PendingGenerationRejectsPreClearCommit) { auto a = MakeAllocator(); - auto pB = AllocateOk(*a, "B", kPageSize * 2, TierType::DRAM); + auto pB = AllocateOk(*a, "B", kPageSize * 2); ASSERT_TRUE(pB.has_value()); - const auto cap_before = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_before = a->Capacity(); EXPECT_EQ(cap_before.available_bytes, cap_before.total_bytes - 2 * kPageSize); a->ClearLocal(); @@ -710,12 +701,12 @@ TEST(PeerDramAllocator, PendingGenerationRejectsPreClearCommit) { EXPECT_FALSE(a->Commit(pB->slot_id, "B", committed_bytes)); EXPECT_EQ(committed_bytes, 0u); EXPECT_TRUE(a->DrainPendingEvents().empty()); - auto cap_after_reject = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_after_reject = a->Capacity(); EXPECT_EQ(cap_after_reject.available_bytes, cap_after_reject.total_bytes); a->ClearFullSyncAcked(); - auto pC = AllocateOk(*a, "C", kPageSize, TierType::DRAM); + auto pC = AllocateOk(*a, "C", kPageSize); ASSERT_TRUE(pC.has_value()); ASSERT_TRUE(a->Commit(pC->slot_id, "C", committed_bytes)); EXPECT_EQ(committed_bytes, kPageSize); @@ -723,10 +714,10 @@ TEST(PeerDramAllocator, PendingGenerationRejectsPreClearCommit) { } // Repeated Clears still reject the original pre-clear pending Commit. -TEST(PeerDramAllocator, PendingGenerationSurvivesDoubleClear) { +TEST(PageBackend, PendingGenerationSurvivesDoubleClear) { auto a = MakeAllocator(); - auto pB = AllocateOk(*a, "B", kPageSize, TierType::DRAM); + auto pB = AllocateOk(*a, "B", kPageSize); ASSERT_TRUE(pB.has_value()); a->ClearLocal(); @@ -734,28 +725,28 @@ TEST(PeerDramAllocator, PendingGenerationSurvivesDoubleClear) { uint64_t committed_bytes = 0; EXPECT_FALSE(a->Commit(pB->slot_id, "B", committed_bytes)); - auto cap_after_reject = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_after_reject = a->Capacity(); EXPECT_EQ(cap_after_reject.available_bytes, cap_after_reject.total_bytes); a->ClearFullSyncAcked(); - auto pC = AllocateOk(*a, "C", kPageSize, TierType::DRAM); + auto pC = AllocateOk(*a, "C", kPageSize); ASSERT_TRUE(pC.has_value()); EXPECT_TRUE(a->Commit(pC->slot_id, "C", committed_bytes)); } // Leased owned key: logically gone at Clear, pages freed by reaper after // the lease expires. -TEST(PeerDramAllocator, ClearLocalDefersLeasedOwnedPages) { +TEST(PageBackend, ClearLocalDefersLeasedOwnedPages) { auto a = MakeAllocator(/*pending_ttl=*/std::chrono::milliseconds{5000}, /*read_lease_ttl=*/std::chrono::milliseconds{200}); - auto p = AllocateOk(*a, "A", kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "A", kPageSize); ASSERT_TRUE(p.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, "A", committed_bytes)); a->DrainPendingEvents(); - const auto cap_committed = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_committed = a->Capacity(); EXPECT_EQ(cap_committed.available_bytes, cap_committed.total_bytes - kPageSize); ASSERT_TRUE(a->Resolve("A").found); // lease. @@ -765,37 +756,37 @@ TEST(PeerDramAllocator, ClearLocalDefersLeasedOwnedPages) { EXPECT_FALSE(a->Resolve("A").found); EXPECT_TRUE(a->SnapshotOwnedKeys().empty()); - auto cap_after_clear = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_after_clear = a->Capacity(); EXPECT_EQ(cap_after_clear.available_bytes, cap_committed.total_bytes - kPageSize); // Pre-TTL sweep: no-op. a->RunReaperOnceForTest(); - auto cap_no_op_sweep = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_no_op_sweep = a->Capacity(); EXPECT_EQ(cap_no_op_sweep.available_bytes, cap_committed.total_bytes - kPageSize); // Past TTL: pages return to bitmap. std::this_thread::sleep_for(std::chrono::milliseconds{300}); a->RunReaperOnceForTest(); - auto cap_swept = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_swept = a->Capacity(); EXPECT_EQ(cap_swept.available_bytes, cap_swept.total_bytes); } // Leased owned A defers; pending B rejects via generation. -TEST(PeerDramAllocator, ClearLocalMixedPendingAndLeased) { +TEST(PageBackend, ClearLocalMixedPendingAndLeased) { auto a = MakeAllocator(/*pending_ttl=*/std::chrono::milliseconds{5000}, /*read_lease_ttl=*/std::chrono::milliseconds{200}); - auto pA = AllocateOk(*a, "A", kPageSize, TierType::DRAM); + auto pA = AllocateOk(*a, "A", kPageSize); ASSERT_TRUE(pA.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(pA->slot_id, "A", committed_bytes)); ASSERT_TRUE(a->Resolve("A").found); // lease. - auto pB = AllocateOk(*a, "B", kPageSize * 2, TierType::DRAM); + auto pB = AllocateOk(*a, "B", kPageSize * 2); ASSERT_TRUE(pB.has_value()); a->DrainPendingEvents(); - const auto total = a->TierCapacitiesSnapshot()[TierType::DRAM].total_bytes; + const auto total = a->Capacity().total_bytes; a->ClearLocal(); @@ -803,115 +794,115 @@ TEST(PeerDramAllocator, ClearLocalMixedPendingAndLeased) { EXPECT_TRUE(a->SnapshotOwnedKeys().empty()); // A deferred + B pending: 3 pages occupied. - auto cap_after_clear = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_after_clear = a->Capacity(); EXPECT_EQ(cap_after_clear.available_bytes, total - 3 * kPageSize); // Commit(B) fails on generation mismatch, releases B. EXPECT_FALSE(a->Commit(pB->slot_id, "B", committed_bytes)); - auto cap_after_reject = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_after_reject = a->Capacity(); EXPECT_EQ(cap_after_reject.available_bytes, total - kPageSize); // only A. // Past lease + sweep: A released. std::this_thread::sleep_for(std::chrono::milliseconds{300}); a->RunReaperOnceForTest(); - auto cap_final = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap_final = a->Capacity(); EXPECT_EQ(cap_final.available_bytes, total); } // Sweeps are no-ops while the deferred lease is still active. -TEST(PeerDramAllocator, ClearLocalSweepRespectsTtl) { +TEST(PageBackend, ClearLocalSweepRespectsTtl) { auto a = MakeAllocator(/*pending_ttl=*/std::chrono::milliseconds{5000}, /*read_lease_ttl=*/std::chrono::milliseconds{10000}); - auto p = AllocateOk(*a, "A", kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "A", kPageSize); ASSERT_TRUE(p.has_value()); uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, "A", committed_bytes)); ASSERT_TRUE(a->Resolve("A").found); - const auto cap_committed = a->TierCapacitiesSnapshot()[TierType::DRAM]; + const auto cap_committed = a->Capacity(); a->ClearLocal(); // Lease still live: every sweep is a no-op. for (int i = 0; i < 3; ++i) { a->RunReaperOnceForTest(); - auto cap = a->TierCapacitiesSnapshot()[TierType::DRAM]; + auto cap = a->Capacity(); EXPECT_EQ(cap.available_bytes, cap_committed.available_bytes) << "sweep i=" << i; } } -// ---- OwnedKeyCountByTier ---------------------------------------------------- +// ---- OwnedKeyCount ----------------------------------------------------------- -TEST(PeerDramAllocator, OwnedKeyCountByTierTracksCommitsAndEvicts) { +TEST(PageBackend, OwnedKeyCountTracksCommitsAndEvicts) { auto a = MakeAllocator(); - auto counts0 = a->OwnedKeyCountByTier(); - EXPECT_EQ(counts0[TierType::DRAM], 0u); - EXPECT_EQ(counts0[TierType::HBM], 0u); - EXPECT_EQ(counts0[TierType::SSD], 0u); + EXPECT_EQ(a->OwnedKeyCount(), 0u); for (int i = 0; i < 3; ++i) { const std::string k = "key-dram-" + std::to_string(i); - auto p = AllocateOk(*a, k, kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, k, kPageSize); ASSERT_TRUE(p.has_value()) << "i=" << i; uint64_t committed_bytes = 0; ASSERT_TRUE(a->Commit(p->slot_id, k, committed_bytes)); } - auto counts1 = a->OwnedKeyCountByTier(); - EXPECT_EQ(counts1[TierType::DRAM], 3u); - EXPECT_EQ(counts1[TierType::HBM], 0u); - EXPECT_EQ(counts1[TierType::SSD], 0u); + EXPECT_EQ(a->OwnedKeyCount(), 3u); a->Evict({"key-dram-0"}); - auto counts2 = a->OwnedKeyCountByTier(); - EXPECT_EQ(counts2[TierType::DRAM], 2u); - EXPECT_EQ(counts2[TierType::HBM], 0u); + EXPECT_EQ(a->OwnedKeyCount(), 2u); } -TEST(PeerDramAllocator, OwnedKeyCountByTierMultiTier) { - PeerDramAllocator::TierConfig hbm_cfg; +// A DRAM instance and an HBM instance are fully independent objects — one +// instance == one medium (backend-agnostic refactor Phase 2), so there is no +// shared multi-tier state to cross-check beyond each Tier() staying fixed. +TEST(PageBackend, SeparateInstancesPerMediumAreIndependent) { + auto dram = MakeAllocator(); + PageBackend::TierConfig hbm_cfg; hbm_cfg.buffer_sizes = {kPageSize * 4}; hbm_cfg.buffer_descs = {{0xD0, 0xD1}}; - auto a = std::make_unique(kPageSize, MakeDramCfg(), hbm_cfg, - std::chrono::milliseconds{5000}); + auto hbm = std::make_unique(TierType::HBM, kPageSize, hbm_cfg, + std::chrono::milliseconds{5000}); + + EXPECT_EQ(dram->Tier(), TierType::DRAM); + EXPECT_EQ(hbm->Tier(), TierType::HBM); for (int i = 0; i < 2; ++i) { const std::string k = "d-" + std::to_string(i); - auto p = AllocateOk(*a, k, kPageSize, TierType::DRAM); + auto p = AllocateOk(*dram, k, kPageSize); ASSERT_TRUE(p.has_value()); uint64_t committed_bytes = 0; - ASSERT_TRUE(a->Commit(p->slot_id, k, committed_bytes)); + ASSERT_TRUE(dram->Commit(p->slot_id, k, committed_bytes)); } { - auto p = AllocateOk(*a, "h-0", kPageSize, TierType::HBM); + auto p = AllocateOk(*hbm, "h-0", kPageSize); ASSERT_TRUE(p.has_value()); uint64_t committed_bytes = 0; - ASSERT_TRUE(a->Commit(p->slot_id, "h-0", committed_bytes)); + ASSERT_TRUE(hbm->Commit(p->slot_id, "h-0", committed_bytes)); } - auto counts = a->OwnedKeyCountByTier(); - EXPECT_EQ(counts[TierType::DRAM], 2u); - EXPECT_EQ(counts[TierType::HBM], 1u); - EXPECT_EQ(counts[TierType::SSD], 0u); + EXPECT_EQ(dram->OwnedKeyCount(), 2u); + EXPECT_EQ(hbm->OwnedKeyCount(), 1u); + // Neither instance sees the other's keys. + EXPECT_FALSE(dram->Resolve("h-0").found); + EXPECT_FALSE(hbm->Resolve("d-0").found); } // ---- Auto-flush: cb fires exactly when the outbox first reaches threshold ---- -TEST(PeerDramAllocatorAutoFlush, FiresAtThreshold) { +TEST(PageBackendAutoFlush, FiresAtThreshold) { auto a = MakeAllocator(); int fires = 0; a->SetAutoFlushHook(3, [&] { ++fires; }); for (int i = 0; i < 2; ++i) { const std::string k = "af-" + std::to_string(i); - auto p = AllocateOk(*a, k, kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, k, kPageSize); ASSERT_TRUE(p.has_value()); uint64_t b = 0; ASSERT_TRUE(a->Commit(p->slot_id, k, b)); } EXPECT_EQ(fires, 0); // below threshold -> no fire - auto p = AllocateOk(*a, "af-2", kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "af-2", kPageSize); ASSERT_TRUE(p.has_value()); uint64_t b = 0; ASSERT_TRUE(a->Commit(p->slot_id, "af-2", b)); @@ -919,11 +910,11 @@ TEST(PeerDramAllocatorAutoFlush, FiresAtThreshold) { } // ---- No hook registered -> auto-flush disabled by default (never fires) ---- -TEST(PeerDramAllocatorAutoFlush, NoHookNoFire) { +TEST(PageBackendAutoFlush, NoHookNoFire) { auto a = MakeAllocator(); for (int i = 0; i < 10; ++i) { const std::string k = "nh-" + std::to_string(i); - auto p = AllocateOk(*a, k, kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, k, kPageSize); ASSERT_TRUE(p.has_value()); uint64_t b = 0; ASSERT_TRUE(a->Commit(p->slot_id, k, b)); @@ -932,11 +923,11 @@ TEST(PeerDramAllocatorAutoFlush, NoHookNoFire) { } // ---- Full-sync snapshot returns all owned AND atomically clears the outbox ---- -TEST(PeerDramAllocatorAutoFlush, FullSyncSnapshotClearsOutbox) { +TEST(PageBackendAutoFlush, FullSyncSnapshotClearsOutbox) { auto a = MakeAllocator(); for (int i = 0; i < 3; ++i) { const std::string k = "fs-" + std::to_string(i); - auto p = AllocateOk(*a, k, kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, k, kPageSize); ASSERT_TRUE(p.has_value()); uint64_t b = 0; ASSERT_TRUE(a->Commit(p->slot_id, k, b)); @@ -949,7 +940,7 @@ TEST(PeerDramAllocatorAutoFlush, FullSyncSnapshotClearsOutbox) { EXPECT_TRUE(a->DrainPendingEvents().empty()); // A commit AFTER the snapshot is a fresh delta (only genuinely new events). - auto p = AllocateOk(*a, "fs-new", kPageSize, TierType::DRAM); + auto p = AllocateOk(*a, "fs-new", kPageSize); ASSERT_TRUE(p.has_value()); uint64_t b = 0; ASSERT_TRUE(a->Commit(p->slot_id, "fs-new", b)); diff --git a/tests/cpp/umbp/distributed/test_page_bitmap_allocator.cpp b/tests/cpp/umbp/distributed/test_page_bitmap_allocator.cpp index 1afb5928e..8aa3bed8f 100644 --- a/tests/cpp/umbp/distributed/test_page_bitmap_allocator.cpp +++ b/tests/cpp/umbp/distributed/test_page_bitmap_allocator.cpp @@ -25,7 +25,7 @@ #include #include -#include "umbp/distributed/peer/peer_page_allocator.h" +#include "umbp/distributed/peer/backend/peer_page_allocator.h" #include "umbp/distributed/types.h" namespace { diff --git a/tests/cpp/umbp/distributed/test_peer_pool.cpp b/tests/cpp/umbp/distributed/test_peer_pool.cpp new file mode 100644 index 000000000..b181e6ec6 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_peer_pool.cpp @@ -0,0 +1,938 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/instrumented_backend.h" +#include "umbp/distributed/peer/backend/mock_backend.h" +#include "umbp/distributed/peer/backend/page_backend.h" +#include "umbp/distributed/pool/peer_pool.h" +#include "umbp/distributed/transfer/local_copy_engine.h" + +namespace mori::umbp { +namespace { + +PoolPlacementRequest PutRequest(std::string key, std::string backend_name = {}) { + PoolPlacementRequest request; + request.key = std::move(key); + request.size = 64; + request.tier = TierType::DRAM; + request.backend_name = std::move(backend_name); + return request; +} + +constexpr auto kNoExpiry = std::chrono::milliseconds(30000); + +// A page backend the tier tests can move real bytes through. +void RegisterPageBackend(BackendRegistry* registry, TransferEngine* engine, const std::string& name, + TierType tier, uint64_t page_size, uint64_t pages, + std::chrono::milliseconds read_lease_ttl = kNoExpiry) { + PageBackend::OwnershipConfig ownership; + ownership.buffer_sizes = {pages * page_size}; + auto backend = MakePageBackend(tier, page_size, ownership, kNoExpiry, read_lease_ttl); + backend = MakeInstrumentedBackend(std::move(backend)); + ASSERT_TRUE(backend->Init(engine)); + ASSERT_TRUE(registry->Register(name, std::move(backend))); +} + +// The topology most tier tests need: a 'hot' tier that offloads to a 'cold' +// one. Watermarks are only consulted for the watermark trigger, so the same +// pair of values suits both triggers. +std::vector HotColdTiers(PoolOffloadTrigger trigger) { + LogicalTierConfig hot; + hot.name = "hot"; + hot.entry = true; + hot.members = {{"hot", 100}}; + hot.offload_to = {"cold"}; + hot.trigger = trigger; + hot.high_watermark = 0.4; + hot.low_watermark = 0.2; + LogicalTierConfig cold; + cold.name = "cold"; + cold.members = {{"cold", 100}}; + return {hot, cold}; +} + +// Puts a key through the pool and returns where it landed. +PoolAllocateResult CommitKey(PeerPool* pool, const std::string& key, + const std::string& backend_name = {}, + const std::string& logical_tier = {}) { + auto request = PutRequest(key, backend_name); + request.logical_tier = logical_tier; + auto allocation = pool->BatchAllocate({request}).front(); + EXPECT_EQ(allocation.allocation.outcome, AllocateOutcome::kSuccessAllocated) << key; + EXPECT_TRUE(pool->BatchCommit({PoolCommitRequest{ + {allocation.backend_id, allocation.allocation.slot_id}, key}}) + .front() + .commit.success) + << key; + return allocation; +} + +// Transitions run on a worker thread, so arrival is observed rather than +// awaited. +bool WaitForKey(BackendRegistry* registry, const std::string& backend, const std::string& key) { + for (int attempt = 0; attempt < 100 && !registry->Get(backend)->Contains(key); ++attempt) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + return registry->Get(backend)->Contains(key); +} + +class NoSpaceBackend final : public MockBackend { + public: + explicit NoSpaceBackend(TierType tier) : MockBackend(tier) {} + + std::vector BatchAllocate(const std::vector& entries) override { + std::vector results(entries.size()); + for (auto& result : results) result.outcome = AllocateOutcome::kFailedNoSpace; + return results; + } +}; + +class BusyResolveBackend final : public MockBackend { + public: + explicit BusyResolveBackend(TierType tier) : MockBackend(tier) {} + + std::vector BatchResolve(const std::vector& keys, bool include_descs, + bool allow_file_refs = false) override { + if (!busy_) return MockBackend::BatchResolve(keys, include_descs, allow_file_refs); + std::vector results(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + if (Contains(keys[i])) results[i].outcome = ResolveOutcome::kBusy; + } + return results; + } + + void SetBusy(bool busy) { busy_ = busy; } + + private: + bool busy_ = true; +}; + +// Stands in for a slow migration target: parks inside the allocation a +// transition performs, so a test can observe the pool while bytes are moving. +class BlockingAllocateBackend final : public MockBackend { + public: + BlockingAllocateBackend(TierType tier, uint64_t page_size) : MockBackend(tier, page_size) {} + + std::vector BatchAllocate(const std::vector& entries) override { + { + std::lock_guard lock(mutex_); + entered_ = true; + } + cv_.notify_all(); + { + std::unique_lock lock(mutex_); + cv_.wait_for(lock, std::chrono::seconds(5), [this] { return released_; }); + } + std::vector results(entries.size()); + for (auto& result : results) result.outcome = AllocateOutcome::kFailedNoSpace; + return results; + } + + bool WaitUntilEntered() { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, std::chrono::seconds(5), [this] { return entered_; }); + } + + void Release() { + { + std::lock_guard lock(mutex_); + released_ = true; + } + cv_.notify_all(); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + bool entered_ = false; + bool released_ = false; +}; + +// Parks one backend resolve so the test can prove PeerPool does not retain its +// peer-wide metadata lock while a medium performs slow IO. +class BlockingResolveBackend final : public MockBackend { + public: + explicit BlockingResolveBackend(TierType tier) : MockBackend(tier) {} + + std::vector BatchResolve(const std::vector& keys, bool include_descs, + bool allow_file_refs = false) override { + { + std::lock_guard lock(mutex_); + entered_ = true; + } + cv_.notify_all(); + { + std::unique_lock lock(mutex_); + cv_.wait_for(lock, std::chrono::seconds(5), [this] { return released_; }); + } + return MockBackend::BatchResolve(keys, include_descs, allow_file_refs); + } + + bool WaitUntilEntered() { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, std::chrono::seconds(5), [this] { return entered_; }); + } + + void Release() { + { + std::lock_guard lock(mutex_); + released_ = true; + } + cv_.notify_all(); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + bool entered_ = false; + bool released_ = false; +}; + +TEST(PeerPool, DefaultPolicyPreservesTierOnlySelection) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto allocated = pool.BatchAllocate({PutRequest("k")}).front(); + ASSERT_EQ(allocated.allocation.outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(allocated.backend_id, registry.BackendId(registry.Get("dram_a"))); +} + +TEST(PeerPool, PreservesBusyOwnerInsteadOfTurningItIntoAMiss) { + BackendRegistry registry; + auto backend = std::make_unique(TierType::DRAM); + auto* busy = backend.get(); + ASSERT_TRUE(registry.Register(std::move(backend))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + CommitKey(&pool, "busy-key"); + + auto first = pool.BatchResolve({"busy-key"}, false); + ASSERT_EQ(first.size(), 1u); + EXPECT_FALSE(first[0].resolved.found); + EXPECT_EQ(first[0].resolved.outcome, ResolveOutcome::kBusy); + + busy->SetBusy(false); + auto second = pool.BatchResolve({"busy-key"}, false); + ASSERT_TRUE(second[0].resolved.found); + EXPECT_EQ(EffectiveResolveOutcome(second[0].resolved), ResolveOutcome::kFound); +} + +TEST(PeerPool, NamedBackendCommitBuildsLogicalPlacement) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto allocated = pool.BatchAllocate({PutRequest("k", "dram_b")}).front(); + ASSERT_EQ(allocated.allocation.outcome, AllocateOutcome::kSuccessAllocated); + const uint32_t second_id = registry.BackendId(registry.Get("dram_b")); + ASSERT_EQ(allocated.backend_id, second_id); + + PoolCommitRequest commit; + commit.slot = PoolSlotRef{allocated.backend_id, allocated.allocation.slot_id}; + commit.key = "k"; + ASSERT_TRUE(pool.BatchCommit({commit}).front().commit.success); + EXPECT_EQ(pool.PlacementBackend("k"), std::optional{second_id}); + EXPECT_EQ(pool.PlacementCount(), 1u); + + auto resolved = pool.BatchResolve({"k"}, /*include_descs=*/false).front(); + ASSERT_TRUE(resolved.resolved.found); + EXPECT_EQ(resolved.backend_id, second_id); + + // Logical dedup happens above the physical backends: asking another instance + // for the same committed key must not create a second copy. + auto duplicate = pool.BatchAllocate({PutRequest("k", "dram_a")}).front(); + EXPECT_EQ(duplicate.allocation.outcome, AllocateOutcome::kSuccessAlreadyExists); + EXPECT_EQ(duplicate.backend_id, second_id); + EXPECT_EQ(registry.Get("dram_a")->OwnedKeyCount(), 0u); + + auto evicted = pool.Evict({"k"}, PoolEvictMode::kReclaim); + ASSERT_EQ(evicted.size(), 1u); + EXPECT_EQ(evicted[0].bytes_freed, 64u); + EXPECT_EQ(pool.PlacementCount(), 0u); +} + +TEST(PeerPool, ResolveFallbackRepairsPlacementIndex) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto* second = registry.Get("dram_b"); + auto allocated = second->BatchAllocate({AllocateRequest{"recovered", 32}}).front(); + ASSERT_TRUE(second->BatchCommit({CommitRequest{allocated.slot_id, "recovered"}}).front().success); + + auto resolved = pool.BatchResolve({"recovered"}, /*include_descs=*/false).front(); + ASSERT_TRUE(resolved.resolved.found); + const uint32_t second_id = registry.BackendId(second); + EXPECT_EQ(resolved.backend_id, second_id); + EXPECT_EQ(pool.PlacementBackend("recovered"), std::optional{second_id}); +} + +TEST(PeerPool, ColdIndexDoesNotDuplicatePersistentKey) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + auto* second = registry.Get("dram_b"); + auto physical = second->BatchAllocate({AllocateRequest{"persistent", 32}}).front(); + ASSERT_TRUE(second->BatchCommit({CommitRequest{physical.slot_id, "persistent"}}).front().success); + + PeerPool pool(®istry, MakeSingleBackendPolicy()); + auto result = pool.BatchAllocate({PutRequest("persistent", "dram_a")}).front(); + EXPECT_EQ(result.allocation.outcome, AllocateOutcome::kSuccessAlreadyExists); + EXPECT_EQ(result.backend_id, registry.BackendId(second)); + EXPECT_EQ(registry.Get("dram_a")->OwnedKeyCount(), 0u); +} + +TEST(PeerPool, ClearDropsPhysicalAndLogicalState) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto allocated = pool.BatchAllocate({PutRequest("k")}).front(); + PoolCommitRequest commit; + commit.slot = PoolSlotRef{allocated.backend_id, allocated.allocation.slot_id}; + commit.key = "k"; + ASSERT_TRUE(pool.BatchCommit({commit}).front().commit.success); + ASSERT_EQ(pool.PlacementCount(), 1u); + + pool.ClearLocal(); + EXPECT_EQ(pool.PlacementCount(), 0u); + EXPECT_EQ(registry.Get("dram")->OwnedKeyCount(), 0u); +} + +TEST(PeerPool, BatchResultsRemainPositionalAcrossBackends) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto results = pool.BatchAllocate( + {PutRequest("a", "dram_b"), PutRequest("bad", "missing"), PutRequest("b", "dram_a")}); + ASSERT_EQ(results.size(), 3u); + EXPECT_EQ(results[0].backend_id, registry.BackendId(registry.Get("dram_b"))); + EXPECT_EQ(results[0].allocation.outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(results[1].allocation.outcome, AllocateOutcome::kFailed); + EXPECT_EQ(results[2].backend_id, registry.BackendId(registry.Get("dram_a"))); + EXPECT_EQ(results[2].allocation.outcome, AllocateOutcome::kSuccessAllocated); +} + +TEST(PeerPool, ReservesDuplicateKeysUntilCommitOrAbort) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto results = pool.BatchAllocate({PutRequest("same", "dram_a"), PutRequest("same", "dram_b")}); + ASSERT_EQ(results.size(), 2u); + EXPECT_EQ(results[0].allocation.outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(results[1].allocation.outcome, AllocateOutcome::kFailed); + + ASSERT_TRUE( + pool.BatchAbort({PoolSlotRef{results[0].backend_id, results[0].allocation.slot_id}}).front()); + auto retried = pool.BatchAllocate({PutRequest("same", "dram_b")}).front(); + EXPECT_EQ(retried.allocation.outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(retried.backend_id, registry.BackendId(registry.Get("dram_b"))); +} + +TEST(PeerPool, CommitMustMatchReservedSlotAndKey) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto allocated = pool.BatchAllocate({PutRequest("reserved")}).front(); + PoolCommitRequest wrong; + wrong.slot = PoolSlotRef{allocated.backend_id, allocated.allocation.slot_id}; + wrong.key = "different"; + EXPECT_FALSE(pool.BatchCommit({wrong}).front().commit.success); + + PoolCommitRequest correct = wrong; + correct.key = "reserved"; + EXPECT_TRUE(pool.BatchCommit({correct}).front().commit.success); + EXPECT_EQ(pool.PlacementBackend("reserved"), std::optional{allocated.backend_id}); + EXPECT_FALSE(pool.PlacementBackend("different").has_value()); +} + +TEST(PeerPool, StalePlacementDoesNotSuppressNewAllocation) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto allocated = pool.BatchAllocate({PutRequest("stale", "dram_a")}).front(); + PoolCommitRequest commit; + commit.slot = PoolSlotRef{allocated.backend_id, allocated.allocation.slot_id}; + commit.key = "stale"; + ASSERT_TRUE(pool.BatchCommit({commit}).front().commit.success); + + // Simulate physical state being cleared/recovered outside this process-local + // index. Allocate validates the indexed backend before claiming a duplicate. + registry.Get("dram_a")->ClearLocal(); + auto replacement = pool.BatchAllocate({PutRequest("stale", "dram_b")}).front(); + EXPECT_EQ(replacement.allocation.outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(replacement.backend_id, registry.BackendId(registry.Get("dram_b"))); +} + +TEST(WeightedPlacementPolicy, DistributesDeterministicallyWithinRequestedTier) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("ssd_a", std::make_unique(TierType::SSD))); + auto policy = MakeWeightedPlacementPolicy({{"dram_a", 3}, {"dram_b", 7}, {"ssd_a", 100}}); + + const uint32_t dram_a = registry.BackendId(registry.Get("dram_a")); + const uint32_t dram_b = registry.BackendId(registry.Get("dram_b")); + size_t on_a = 0; + size_t on_b = 0; + for (size_t i = 0; i < 10000; ++i) { + auto request = PutRequest("weighted-" + std::to_string(i * 2654435761ULL)); + auto order = policy->PutOrder(registry, request); + ASSERT_FALSE(order.empty()); + if (order.front() == dram_a) ++on_a; + if (order.front() == dram_b) ++on_b; + } + + EXPECT_EQ(on_a + on_b, 10000u); + EXPECT_NEAR(static_cast(on_a) / 10000.0, 0.3, 0.03); + + auto repeated = PutRequest("stable-key"); + auto first = policy->PutOrder(registry, repeated); + for (int i = 0; i < 10; ++i) { + EXPECT_EQ(policy->PutOrder(registry, repeated), first); + } +} + +TEST(WeightedPlacementPolicy, ExplicitBackendNameOverridesWeightedSelection) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + auto policy = MakeWeightedPlacementPolicy({{"dram_a", 1}, {"dram_b", 100}}); + + auto order = policy->PutOrder(registry, PutRequest("forced", "dram_a")); + ASSERT_FALSE(order.empty()); + EXPECT_EQ(order.front(), registry.BackendId(registry.Get("dram_a"))); +} + +TEST(PeerPool, WeightedPolicyPlacesNewKeysOnMultipleBackends) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeWeightedPlacementPolicy({{"dram_a", 1}, {"dram_b", 1}})); + + std::vector requests; + for (size_t i = 0; i < 128; ++i) { + requests.push_back(PutRequest("pool-weighted-" + std::to_string(i))); + } + auto results = pool.BatchAllocate(requests); + + const uint32_t dram_a = registry.BackendId(registry.Get("dram_a")); + const uint32_t dram_b = registry.BackendId(registry.Get("dram_b")); + size_t on_a = 0; + size_t on_b = 0; + std::vector slots; + for (const auto& result : results) { + ASSERT_EQ(result.allocation.outcome, AllocateOutcome::kSuccessAllocated); + if (result.backend_id == dram_a) ++on_a; + if (result.backend_id == dram_b) ++on_b; + slots.push_back(PoolSlotRef{result.backend_id, result.allocation.slot_id}); + } + EXPECT_GT(on_a, 0u); + EXPECT_GT(on_b, 0u); + pool.BatchAbort(slots); +} + +TEST(PeerPool, WeightedPolicyFallsBackWhenPrimaryHasNoSpace) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + auto policy = MakeWeightedPlacementPolicy({{"dram_a", 1}, {"dram_b", 1}}); + + PoolPlacementRequest request; + bool found_primary = false; + for (size_t i = 0; i < 100; ++i) { + request = PutRequest("fallback-" + std::to_string(i)); + auto order = policy->PutOrder(registry, request); + if (!order.empty() && order.front() == registry.BackendId(registry.Get("dram_a"))) { + found_primary = true; + break; + } + } + ASSERT_TRUE(found_primary); + ASSERT_EQ(policy->PutOrder(registry, request).front(), + registry.BackendId(registry.Get("dram_a"))); + + PeerPool pool(®istry, std::move(policy)); + auto result = pool.BatchAllocate({request}).front(); + EXPECT_EQ(result.allocation.outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(result.backend_id, registry.BackendId(registry.Get("dram_b"))); +} + +TEST(TieredPlacementPolicy, RejectsTierGraphWithoutTransferEngine) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("only", std::make_unique(TierType::DRAM))); + auto compiled = LogicalTierGraph::Compile({LogicalTierConfig{{{"only", 1}}}}, registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + EXPECT_THROW(PeerPool(®istry, MakeTieredPlacementPolicy(compiled.graph)), + std::invalid_argument); +} + +TEST(TieredPlacementPolicy, IgnoresPhysicalRouteAndSpillsToConfiguredTarget) { + LocalCopyEngine engine; + BackendRegistry registry; + ASSERT_TRUE(registry.Register("hot", std::make_unique(TierType::HBM))); + ASSERT_TRUE(registry.Register("cold", std::make_unique(TierType::SSD))); + auto compiled = LogicalTierGraph::Compile(HotColdTiers(PoolOffloadTrigger::kOnEvict), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + auto request = PutRequest("spill"); + request.tier = TierType::DRAM; + auto result = pool.BatchAllocate({request}).front(); + EXPECT_EQ(result.allocation.outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(result.backend_id, registry.BackendId(registry.Get("cold"))); +} + +TEST(TieredPlacementPolicy, HonorsExplicitLogicalTierAndTagsEvents) { + LocalCopyEngine engine; + BackendRegistry registry; + ASSERT_TRUE(registry.Register("hot", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("cold", std::make_unique(TierType::DRAM))); + auto compiled = LogicalTierGraph::Compile(HotColdTiers(PoolOffloadTrigger::kOnEvict), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + auto allocation = CommitKey(&pool, "named", {}, "cold"); + ASSERT_EQ(allocation.backend_id, registry.BackendId(registry.Get("cold"))); + + auto events = pool.DrainPendingEvents(); + ASSERT_EQ(events.size(), 1u); + EXPECT_EQ(events.front().logical_tier, "cold"); + auto capacities = pool.LogicalTierCapacities(); + EXPECT_TRUE(capacities.count("hot")); + EXPECT_TRUE(capacities.count("cold")); +} + +TEST(PeerPool, OnEvictMigratesBytesToConfiguredBackend) { + constexpr uint64_t kPageSize = 64; + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 2); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 2); + auto compiled = LogicalTierGraph::Compile(HotColdTiers(PoolOffloadTrigger::kOnEvict), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + auto allocation = pool.BatchAllocate({PutRequest("move")}).front(); + ASSERT_EQ(allocation.allocation.outcome, AllocateOutcome::kSuccessAllocated); + auto source_ref = registry.Get("hot")->BufferRef(allocation.allocation.pages[0].buffer_index); + ASSERT_TRUE(source_ref.HasHostPtr()); + const std::string payload(64, 'x'); + std::memcpy(static_cast(source_ref.host_ptr) + + allocation.allocation.pages[0].page_index * kPageSize, + payload.data(), payload.size()); + ASSERT_TRUE(pool.BatchCommit({PoolCommitRequest{ + {allocation.backend_id, allocation.allocation.slot_id}, "move"}}) + .front() + .commit.success); + + auto evicted = pool.Evict({"move"}, PoolEvictMode::kReclaim); + ASSERT_EQ(evicted.size(), 1u); + EXPECT_EQ(evicted.front().bytes_freed, payload.size()); + EXPECT_FALSE(registry.Get("hot")->Contains("move")); + EXPECT_TRUE(registry.Get("cold")->Contains("move")); + EXPECT_EQ(pool.PlacementBackend("move"), + std::optional{registry.BackendId(registry.Get("cold"))}); + + auto resolved = pool.BatchResolve({"move"}, false).front(); + ASSERT_TRUE(resolved.resolved.found); + auto target_ref = registry.Get("cold")->BufferRef(resolved.resolved.pages[0].buffer_index); + ASSERT_TRUE(target_ref.HasHostPtr()); + EXPECT_EQ(std::string(static_cast(target_ref.host_ptr) + + resolved.resolved.pages[0].page_index * kPageSize, + payload.size()), + payload); +} + +TEST(PeerPool, EvictionRetryDrainsLeasedSourceWithoutDeletingTarget) { + constexpr uint64_t kPageSize = 64; + constexpr auto kShortLease = std::chrono::milliseconds(1); + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 2, kShortLease); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 2, kShortLease); + auto compiled = LogicalTierGraph::Compile(HotColdTiers(PoolOffloadTrigger::kOnEvict), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + CommitKey(&pool, "leased"); + ASSERT_TRUE(registry.Get("hot")->BatchResolve({"leased"}, false).front().found); + + auto first = pool.Evict({"leased"}, PoolEvictMode::kReclaim); + ASSERT_EQ(first.front().bytes_freed, 0u); + EXPECT_TRUE(registry.Get("hot")->Contains("leased")); + EXPECT_TRUE(registry.Get("cold")->Contains("leased")); + + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + auto retry = pool.Evict({"leased"}, PoolEvictMode::kReclaim); + EXPECT_EQ(retry.front().bytes_freed, 64u); + EXPECT_FALSE(registry.Get("hot")->Contains("leased")); + EXPECT_TRUE(registry.Get("cold")->Contains("leased")); +} + +TEST(PeerPool, DiscardDeletesInsteadOfDemoting) { + constexpr uint64_t kPageSize = 64; + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 2); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 2); + // The same on_evict topology that demotes under kReclaim. + auto compiled = LogicalTierGraph::Compile(HotColdTiers(PoolOffloadTrigger::kOnEvict), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + CommitKey(&pool, "gone"); + + auto discarded = pool.Evict({"gone"}, PoolEvictMode::kDiscard); + ASSERT_EQ(discarded.size(), 1u); + EXPECT_EQ(discarded.front().bytes_freed, 64u); + EXPECT_FALSE(registry.Get("hot")->Contains("gone")); + EXPECT_FALSE(registry.Get("cold")->Contains("gone")); + EXPECT_EQ(pool.PlacementCount(), 0u); +} + +TEST(PeerPool, WatermarkMigratesCommittedKey) { + // Unlike the other tier tests, this one asserts on utilization, so the owned + // buffer must be a whole number of host allocation granules: a smaller + // request is rounded up by the memory source, and PageBackend adopts the + // size it actually got, which would put utilization far below the watermark. + constexpr uint64_t kPageSize = 4096; + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 2); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 2); + // Fails loudly if the granularity assumption above ever stops holding. + ASSERT_EQ(registry.Get("hot")->Capacity().total_bytes, 2 * kPageSize); + + auto compiled = LogicalTierGraph::Compile(HotColdTiers(PoolOffloadTrigger::kWatermark), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + CommitKey(&pool, "watermark"); + EXPECT_TRUE(WaitForKey(®istry, "cold", "watermark")); + EXPECT_FALSE(registry.Get("hot")->Contains("watermark")); +} + +TEST(PeerPool, WatermarkFollowsTheFullestMemberOfAMixedTier) { + constexpr uint64_t kPageSize = 4096; + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "small", TierType::DRAM, kPageSize, 2); + RegisterPageBackend(®istry, &engine, "large", TierType::DRAM, kPageSize, 64); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 64); + + LogicalTierConfig parallel; + parallel.members = {{"small", 1}, {"large", 1}}; + parallel.offload_to = {"cold"}; + parallel.trigger = PoolOffloadTrigger::kWatermark; + parallel.high_watermark = 0.4; + parallel.low_watermark = 0.2; + auto compiled = LogicalTierGraph::Compile( + {parallel, LogicalTierConfig{{{"cold", 1}}, {}, PoolOffloadTrigger::kOnEvict}}, registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + // One page on each member. Aggregated over the tier that is ~3% used, so an + // averaged watermark would never fire, yet "small" is half full. + CommitKey(&pool, "on_large", "large"); + CommitKey(&pool, "on_small", "small"); + + EXPECT_TRUE(WaitForKey(®istry, "cold", "on_small")); + EXPECT_FALSE(registry.Get("small")->Contains("on_small")); + // The roomy member is not under pressure, so its key must stay put. + EXPECT_TRUE(registry.Get("large")->Contains("on_large")); + EXPECT_FALSE(registry.Get("cold")->Contains("on_large")); +} + +TEST(PeerPool, MigrationDoesNotBlockConcurrentPoolOperations) { + constexpr uint64_t kPageSize = 4096; + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 2); + auto cold = std::make_unique(TierType::SSD, kPageSize); + auto* blocking = cold.get(); + ASSERT_TRUE(registry.Register("cold", std::move(cold))); + + auto compiled = LogicalTierGraph::Compile(HotColdTiers(PoolOffloadTrigger::kWatermark), registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + CommitKey(&pool, "blocked"); + ASSERT_TRUE(blocking->WaitUntilEntered()); + + // The transition is parked mid-copy. An unrelated resolve must not wait for + // it, which is only true while the copy runs without the pool lock. + auto concurrent = std::async(std::launch::async, + [&] { return pool.BatchResolve({"unrelated"}, false).size(); }); + EXPECT_EQ(concurrent.wait_for(std::chrono::seconds(2)), std::future_status::ready); + blocking->Release(); + EXPECT_EQ(concurrent.get(), 1u); +} + +TEST(PeerPool, SlowResolveDoesNotBlockUnrelatedPreferredBackend) { + BackendRegistry registry; + auto slow = std::make_unique(TierType::SSD); + auto* blocking = slow.get(); + ASSERT_TRUE(registry.Register("slow", std::move(slow))); + ASSERT_TRUE(registry.Register("fast", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + CommitKey(&pool, "fast-key", "fast"); + + auto slow_resolve = + std::async(std::launch::async, [&] { return pool.BatchResolve({"missing"}, false); }); + ASSERT_TRUE(blocking->WaitUntilEntered()); + + // The preferred placement resolves entirely on the fast backend. It must + // not queue behind unrelated IO already running on another medium. + auto fast_resolve = + std::async(std::launch::async, [&] { return pool.BatchResolve({"fast-key"}, false); }); + const auto status = fast_resolve.wait_for(std::chrono::seconds(2)); + blocking->Release(); + + EXPECT_EQ(status, std::future_status::ready); + auto fast_result = fast_resolve.get(); + ASSERT_EQ(fast_result.size(), 1u); + EXPECT_TRUE(fast_result.front().resolved.found); + EXPECT_EQ(slow_resolve.get().size(), 1u); +} + +TEST(PeerPool, ResolveMissDoesNotEraseConcurrentCommit) { + BackendRegistry registry; + auto slow = std::make_unique(TierType::SSD); + auto* blocking = slow.get(); + ASSERT_TRUE(registry.Register("slow", std::move(slow))); + ASSERT_TRUE(registry.Register("fast", std::make_unique(TierType::DRAM))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + + auto stale_resolve = + std::async(std::launch::async, [&] { return pool.BatchResolve({"raced"}, false); }); + ASSERT_TRUE(blocking->WaitUntilEntered()); + + auto concurrent_commit = std::async(std::launch::async, [&] { + auto allocation = pool.BatchAllocate({PutRequest("raced", "fast")}).front(); + bool committed = false; + if (allocation.allocation.outcome == AllocateOutcome::kSuccessAllocated) { + committed = + pool.BatchCommit({PoolCommitRequest{ + {allocation.backend_id, allocation.allocation.slot_id}, "raced"}}) + .front() + .commit.success; + } + return std::make_pair(allocation, committed); + }); + const auto status = concurrent_commit.wait_for(std::chrono::seconds(2)); + blocking->Release(); + + EXPECT_EQ(status, std::future_status::ready); + const auto [allocation, committed] = concurrent_commit.get(); + EXPECT_TRUE(committed); + EXPECT_EQ(allocation.backend_id, registry.BackendId(registry.Get("fast"))); + EXPECT_EQ(stale_resolve.get().size(), 1u); + EXPECT_EQ(pool.PlacementBackend("raced"), + std::optional{registry.BackendId(registry.Get("fast"))}); +} + +TEST(PeerPool, ClearWaitsForInFlightResolveAndLeavesNoStalePlacement) { + BackendRegistry registry; + auto backend = std::make_unique(TierType::DRAM); + auto* blocking = backend.get(); + ASSERT_TRUE(registry.Register("dram", std::move(backend))); + PeerPool pool(®istry, MakeSingleBackendPolicy()); + CommitKey(&pool, "clear-race", "dram"); + + auto resolve = + std::async(std::launch::async, [&] { return pool.BatchResolve({"clear-race"}, false); }); + ASSERT_TRUE(blocking->WaitUntilEntered()); + + auto clear = std::async(std::launch::async, [&] { pool.ClearLocal(); }); + EXPECT_EQ(clear.wait_for(std::chrono::milliseconds(100)), std::future_status::timeout); + blocking->Release(); + + ASSERT_TRUE(resolve.get().front().resolved.found); + clear.get(); + EXPECT_EQ(pool.PlacementBackend("clear-race"), std::nullopt); + EXPECT_FALSE(registry.Get("dram")->Contains("clear-race")); +} + +TEST(PeerPool, ReadPromotesColdKeyWithoutDeletingSource) { + constexpr uint64_t kPageSize = 64; + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 2); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 2); + + auto tiers = HotColdTiers(PoolOffloadTrigger::kOnEvict); + tiers.back().promote_trigger = PoolPromoteTrigger::kOnRead; + auto compiled = LogicalTierGraph::Compile(tiers, registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + CommitKey(&pool, "promote", {}, "cold"); + ASSERT_TRUE(pool.BatchResolve({"promote"}, false).front().resolved.found); + + EXPECT_TRUE(WaitForKey(®istry, "hot", "promote")); + EXPECT_TRUE(registry.Get("cold")->Contains("promote")); +} + +// The on_hits trigger, which is what made promote_trigger a key-value policy +// rather than a boolean. The threshold has to be counted per key, so one key +// crossing it must not carry another one along, and the count must not survive +// the promotion it caused. +TEST(PeerPool, ReadPromotesOnlyAfterConfiguredHitCount) { + constexpr uint64_t kPageSize = 64; + constexpr uint32_t kHits = 3; + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 4); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 4); + + auto tiers = HotColdTiers(PoolOffloadTrigger::kOnEvict); + tiers.back().promote_trigger = PoolPromoteTrigger::kOnHits; + tiers.back().promote_hits = kHits; + auto compiled = LogicalTierGraph::Compile(tiers, registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + CommitKey(&pool, "warm", {}, "cold"); + CommitKey(&pool, "bystander", {}, "cold"); + + // Below the threshold the reads still have to be served, from the cold tier, + // while the count builds up. + for (uint32_t read = 1; read < kHits; ++read) { + ASSERT_TRUE(pool.BatchResolve({"warm"}, false).front().resolved.found) << read; + } + // Settled first, because promotion is asynchronous: checking straight after + // the read would also pass while a wrongly queued promotion was still in + // flight, which is the bug this is meant to catch. + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + EXPECT_FALSE(registry.Get("hot")->Contains("warm")) + << "promoted after " << kHits - 1 << " reads, below the threshold of " << kHits; + + ASSERT_TRUE(pool.BatchResolve({"warm"}, false).front().resolved.found); + EXPECT_TRUE(WaitForKey(®istry, "hot", "warm")); + + // The other cold key was read zero times and must not have been dragged along + // by a count kept per tier instead of per key. + EXPECT_FALSE(registry.Get("hot")->Contains("bystander")); + EXPECT_TRUE(registry.Get("cold")->Contains("bystander")); +} + +// The move half of promote_mode, which nothing else covers. Copy leaves the +// cold copy behind, so a lost promotion costs nothing; move deletes it, and the +// promoted copy becomes the only one. That makes two things load-bearing that +// copy mode never has to get right: the source has to go away (or the tier keeps +// paying for a key it no longer serves), and the bytes have to survive the hop +// (or the delete has destroyed the only copy). +// +// The read that triggers the promotion is still holding a lease on the source +// page while the promotion runs, so the delete cannot land inside the transition +// itself; the source is parked as draining and the next evict of that key +// finishes it. Hence the short lease and the explicit evict below. +TEST(PeerPool, ReadPromotesWithMoveDrainsSourceAndKeepsBytes) { + constexpr uint64_t kPageSize = 64; + constexpr auto kShortLease = std::chrono::milliseconds(1); + LocalCopyEngine engine; + BackendRegistry registry; + RegisterPageBackend(®istry, &engine, "hot", TierType::DRAM, kPageSize, 2, kShortLease); + RegisterPageBackend(®istry, &engine, "cold", TierType::SSD, kPageSize, 2, kShortLease); + + auto tiers = HotColdTiers(PoolOffloadTrigger::kOnEvict); + tiers.back().promote_trigger = PoolPromoteTrigger::kOnRead; + tiers.back().promote_mode = PoolTransitionMode::kMove; + auto compiled = LogicalTierGraph::Compile(tiers, registry); + ASSERT_TRUE(compiled.ok()) << compiled.error; + PeerPool pool(®istry, MakeTieredPlacementPolicy(compiled.graph), &engine); + + // Placed straight onto cold and written through its own page, so the payload + // the promotion has to carry is known rather than incidental. + auto request = PutRequest("moved"); + request.logical_tier = "cold"; + auto allocation = pool.BatchAllocate({request}).front(); + ASSERT_EQ(allocation.allocation.outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_EQ(allocation.backend_id, registry.BackendId(registry.Get("cold"))); + auto source_ref = registry.Get("cold")->BufferRef(allocation.allocation.pages[0].buffer_index); + ASSERT_TRUE(source_ref.HasHostPtr()); + const std::string payload(kPageSize, 'm'); + std::memcpy(static_cast(source_ref.host_ptr) + + allocation.allocation.pages[0].page_index * kPageSize, + payload.data(), payload.size()); + ASSERT_TRUE( + pool.BatchCommit( + {PoolCommitRequest{{allocation.backend_id, allocation.allocation.slot_id}, "moved"}}) + .front() + .commit.success); + + ASSERT_TRUE(pool.BatchResolve({"moved"}, false).front().resolved.found); + EXPECT_TRUE(WaitForKey(®istry, "hot", "moved")); + EXPECT_EQ(pool.PlacementBackend("moved"), + std::optional{registry.BackendId(registry.Get("hot"))}); + + // Past the lease, so the parked source can actually be freed. The evict has to + // take the draining source and leave the promoted copy alone, which is the one + // way move can lose data outright. + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + pool.Evict({"moved"}, PoolEvictMode::kReclaim); + EXPECT_FALSE(registry.Get("cold")->Contains("moved")); + EXPECT_EQ(registry.Get("cold")->OwnedKeyCount(), 0u); + EXPECT_TRUE(registry.Get("hot")->Contains("moved")); + EXPECT_EQ(registry.Get("hot")->OwnedKeyCount(), 1u); + EXPECT_EQ(pool.PlacementBackend("moved"), + std::optional{registry.BackendId(registry.Get("hot"))}); + + auto resolved = pool.BatchResolve({"moved"}, false).front(); + ASSERT_TRUE(resolved.resolved.found); + auto target_ref = registry.Get("hot")->BufferRef(resolved.resolved.pages[0].buffer_index); + ASSERT_TRUE(target_ref.HasHostPtr()); + EXPECT_EQ(std::string(static_cast(target_ref.host_ptr) + + resolved.resolved.pages[0].page_index * kPageSize, + payload.size()), + payload); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_peer_service.cpp b/tests/cpp/umbp/distributed/test_peer_service.cpp deleted file mode 100644 index ed7506f13..000000000 --- a/tests/cpp/umbp/distributed/test_peer_service.cpp +++ /dev/null @@ -1,357 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/common/config.h" -#include "umbp/distributed/config.h" -#include "umbp/distributed/peer/peer_dram_allocator.h" -#include "umbp/distributed/peer/peer_service.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" -#include "umbp_peer.grpc.pb.h" - -namespace mori::umbp { -namespace { - -constexpr size_t kStagingSize = 4096; -constexpr uint64_t kPageSize = 4096; -constexpr uint16_t kBasePort = 50200; -constexpr int kNumReadSlots = 4; -constexpr int kLeaseTimeoutS = 2; - -// Ask the kernel for a currently-free ephemeral port. PeerServiceServer binds -// this port directly, so a hard-coded base (kBasePort) collides with concurrent -// test processes / leftover servers on a shared (self-hosted CI) host and makes -// the suite flaky. -static uint16_t AllocPort() { - int fd = ::socket(AF_INET, SOCK_STREAM, 0); - if (fd >= 0) { - sockaddr_in addr{}; - addr.sin_family = AF_INET; - addr.sin_addr.s_addr = htonl(INADDR_ANY); - addr.sin_port = 0; - socklen_t len = sizeof(addr); - if (::bind(fd, reinterpret_cast(&addr), sizeof(addr)) == 0 && - ::getsockname(fd, reinterpret_cast(&addr), &len) == 0) { - uint16_t port = ntohs(addr.sin_port); - ::close(fd); - return port; - } - ::close(fd); - } - static std::atomic next{ - static_cast(kBasePort + (static_cast(::getpid()) % 4000))}; - return next.fetch_add(1); -} - -class PeerServiceSlotTest : public ::testing::Test { - protected: - void SetUp() override { - staging_buffer_ = std::malloc(kStagingSize); - ASSERT_NE(staging_buffer_, nullptr); - std::memset(staging_buffer_, 0, kStagingSize); - - ssd_dir_ = std::filesystem::temp_directory_path() / - ("umbp_test_ssd_" + std::to_string(getpid()) + "_" + std::to_string(AllocPort())); - std::filesystem::create_directories(ssd_dir_); - - ssd_staging_mem_desc_ = {0xD0, 0xE0, 0xF0}; - - // Peer-side SSD tier owner (file SSDTier on a temp dir; Posix I/O backend to - // avoid io_uring availability differences inside the build container). - PeerSsdConfig ssd_cfg; - ssd_cfg.enabled = true; - ssd_cfg.ssd.enabled = true; - ssd_cfg.ssd.storage_dir = ssd_dir_.string(); - ssd_cfg.ssd.capacity_bytes = 1 << 20; - ssd_cfg.ssd.io.backend = UMBPIoBackend::Posix; - peer_ssd_ = std::make_unique(ssd_cfg); - - // Standalone DRAM allocator: the SSD-path RPCs exercised here never - // allocate DRAM, but PeerServiceServer requires a non-null allocator. - // An empty TierConfig (no DRAM/HBM buffers) is sufficient and avoids - // standing up a PoolClient / connecting to a master. - dram_alloc_ = std::make_unique(kPageSize, PeerDramAllocator::TierConfig{}, - PeerDramAllocator::TierConfig{}, - std::chrono::milliseconds{1000}); - - port_ = AllocPort(); - - server_ = std::make_unique( - dram_alloc_.get(), peer_ssd_.get(), staging_buffer_, kStagingSize, ssd_staging_mem_desc_, - kNumReadSlots, std::chrono::seconds(kLeaseTimeoutS), std::vector{}, - /*master_client=*/nullptr); - server_->Start(port_); - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - - auto channel = grpc::CreateChannel("localhost:" + std::to_string(port_), - grpc::InsecureChannelCredentials()); - stub_ = ::umbp::UMBPPeer::NewStub(channel); - } - - void TearDown() override { - server_->Stop(); - server_.reset(); - dram_alloc_.reset(); - peer_ssd_.reset(); - std::free(staging_buffer_); - std::filesystem::remove_all(ssd_dir_); - } - - void WriteTestDataToSsd(const std::string& key, const std::string& data) { - peer_ssd_->Write(key, {{data.data(), data.size()}}, data.size()); - } - - void* staging_buffer_ = nullptr; - std::filesystem::path ssd_dir_; - std::vector ssd_staging_mem_desc_; - uint16_t port_ = 0; - std::unique_ptr peer_ssd_; - std::unique_ptr dram_alloc_; - std::unique_ptr server_; - std::unique_ptr<::umbp::UMBPPeer::Stub> stub_; -}; - -// --- GetPeerInfo --- - -TEST_F(PeerServiceSlotTest, GetPeerInfoReturnsStagingInfo) { - ::umbp::GetPeerInfoRequest request; - ::umbp::GetPeerInfoResponse response; - grpc::ClientContext context; - - auto status = stub_->GetPeerInfo(&context, request, &response); - ASSERT_TRUE(status.ok()) << status.error_message(); - EXPECT_EQ(response.ssd_staging_mem_desc(), - std::string(ssd_staging_mem_desc_.begin(), ssd_staging_mem_desc_.end())); - EXPECT_EQ(response.ssd_staging_size(), kStagingSize); -} - -// --- PrepareSsdRead --- - -TEST_F(PeerServiceSlotTest, PrepareSsdReadSuccess) { - const std::string data = "read me from ssd"; - WriteTestDataToSsd("block_read", data); - - ::umbp::PrepareSsdReadRequest req; - req.set_key("block_read"); - req.set_max_size(data.size()); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - - auto status = stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_TRUE(status.ok()) << status.error_message(); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK); - EXPECT_LT(resp.staging_offset(), kStagingSize); - EXPECT_EQ(resp.size(), data.size()); - EXPECT_GT(resp.lease_id(), 0u); - EXPECT_GT(resp.lease_ttl_ms(), 0u); - - std::string loaded(static_cast(staging_buffer_) + resp.staging_offset(), - resp.size()); - EXPECT_EQ(loaded, data); -} - -TEST_F(PeerServiceSlotTest, PrepareSsdReadNotFound) { - ::umbp::PrepareSsdReadRequest req; - req.set_key("nonexistent"); - req.set_max_size(64); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - - auto status = stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - EXPECT_EQ(resp.status(), ::umbp::SSD_READ_NOT_FOUND); -} - -TEST_F(PeerServiceSlotTest, PrepareSsdReadTooLarge) { - // A key whose actual size exceeds the per-slot size must report - // SIZE_TOO_LARGE (distinct from NOT_FOUND), not be silently served. - const std::string big(kStagingSize, 'x'); // larger than one slot (kStagingSize/kNumReadSlots) - WriteTestDataToSsd("block_big", big); - - ::umbp::PrepareSsdReadRequest req; - req.set_key("block_big"); - req.set_max_size(kStagingSize); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - - auto status = stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - EXPECT_EQ(resp.status(), ::umbp::SSD_READ_SIZE_TOO_LARGE); -} - -TEST_F(PeerServiceSlotTest, PrepareSsdReadExhaustSlots) { - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "block_" + std::to_string(i); - const std::string data = "data_" + std::to_string(i); - WriteTestDataToSsd(key, data); - } - - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "block_" + std::to_string(i); - ::umbp::PrepareSsdReadRequest req; - req.set_key(key); - req.set_max_size(64); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - auto status = stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK) << "Slot " << i << " should succeed"; - } - - WriteTestDataToSsd("block_extra", "extra"); - ::umbp::PrepareSsdReadRequest req; - req.set_key("block_extra"); - req.set_max_size(64); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - auto status = stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - // Transient slot exhaustion is NO_SLOT (retryable) — never a NOT_FOUND miss. - EXPECT_EQ(resp.status(), ::umbp::SSD_READ_NO_SLOT) << "Should be NO_SLOT when slots exhausted"; -} - -// --- ReleaseSsdLease --- - -TEST_F(PeerServiceSlotTest, ReleaseSsdLeaseSuccess) { - const std::string data = "release test"; - WriteTestDataToSsd("block_rel", data); - - uint64_t lease_id; - { - ::umbp::PrepareSsdReadRequest req; - req.set_key("block_rel"); - req.set_max_size(data.size()); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK); - lease_id = resp.lease_id(); - } - - { - ::umbp::ReleaseSsdLeaseRequest req; - req.set_lease_id(lease_id); - ::umbp::ReleaseSsdLeaseResponse resp; - grpc::ClientContext ctx; - auto status = stub_->ReleaseSsdLease(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - EXPECT_TRUE(resp.success()); - } - - { - ::umbp::ReleaseSsdLeaseRequest req; - req.set_lease_id(lease_id); - ::umbp::ReleaseSsdLeaseResponse resp; - grpc::ClientContext ctx; - auto status = stub_->ReleaseSsdLease(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - EXPECT_FALSE(resp.success()) << "Double release should fail"; - } -} - -TEST_F(PeerServiceSlotTest, ReleaseSsdLeaseInvalid) { - ::umbp::ReleaseSsdLeaseRequest req; - req.set_lease_id(9999); - ::umbp::ReleaseSsdLeaseResponse resp; - grpc::ClientContext ctx; - auto status = stub_->ReleaseSsdLease(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - EXPECT_FALSE(resp.success()); -} - -// --- Slot isolation: different slots get different offsets --- - -TEST_F(PeerServiceSlotTest, MultipleReadSlotsDifferentOffsets) { - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "iso_" + std::to_string(i); - const std::string data(16, 'a' + i); - WriteTestDataToSsd(key, data); - } - - std::vector offsets; - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "iso_" + std::to_string(i); - ::umbp::PrepareSsdReadRequest req; - req.set_key(key); - req.set_max_size(16); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK); - offsets.push_back(resp.staging_offset()); - } - - for (size_t i = 0; i < offsets.size(); ++i) { - for (size_t j = i + 1; j < offsets.size(); ++j) { - EXPECT_NE(offsets[i], offsets[j]) << "Slots " << i << " and " << j << " overlap"; - } - } -} - -// --- TTL reclaim for read slots --- - -TEST_F(PeerServiceSlotTest, ReadSlotTtlReclaim) { - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "ttl_" + std::to_string(i); - const std::string data = "ttl_data_" + std::to_string(i); - WriteTestDataToSsd(key, data); - } - - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "ttl_" + std::to_string(i); - ::umbp::PrepareSsdReadRequest req; - req.set_key(key); - req.set_max_size(10); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK); - } - - std::this_thread::sleep_for(std::chrono::seconds(kLeaseTimeoutS + 1)); - - const std::string key = "ttl_0"; - ::umbp::PrepareSsdReadRequest req; - req.set_key(key); - req.set_max_size(10); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - auto status = stub_->PrepareSsdRead(&ctx, req, &resp); - ASSERT_TRUE(status.ok()); - EXPECT_EQ(resp.status(), ::umbp::SSD_READ_OK) << "Should succeed after TTL reclaim"; -} - -} // namespace -} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_peer_service_dispatch.cpp b/tests/cpp/umbp/distributed/test_peer_service_dispatch.cpp new file mode 100644 index 000000000..80c33989d --- /dev/null +++ b/tests/cpp/umbp/distributed/test_peer_service_dispatch.cpp @@ -0,0 +1,583 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// PeerServiceServer tier dispatch (backend-agnostic refactor Phase 3). +// +// The peer service used to hold one typed PageBackend* and serve every RPC from +// it. It now holds a PeerPool*, which owns logical placement and dispatches to +// the backend the policy picked. These tests drive the real gRPC surface with +// two MockBackends registered, so they assert the wire-visible contract rather +// than the handlers' internals: +// +// * a request routes to the backend registered for its tier, and a tier with +// no backend fails cleanly instead of hitting whichever one happens to be +// there +// * Commit / Abort carry NO backend selector on the wire, yet still reach the +// backend that handed the slot out — the peer tags backend_id into the opaque +// slot_id (this is what makes two backends numbering their slots from 1 +// unambiguous) +// * a mixed-tier batch is grouped per backend but answered in request order +// * Resolve / Evict, which carry no tier either, walk the peer-local +// read-rank order and fan out respectively + +#include +#include + +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/peer/backend/mock_backend.h" +#include "umbp/distributed/peer/peer_service.h" +#include "umbp/distributed/pool/peer_pool.h" +#include "umbp_peer.grpc.pb.h" + +namespace mori::umbp { +namespace { + +// Every medium is equivalent (Phase 4), so the registry walks them in dense +// backend-id order, i.e. the order SetUp registers them — HBM before DRAM here. +// The Resolve test depends only on that order being deterministic, not on it +// meaning "faster". +constexpr TierType kFirstByTier = TierType::HBM; +constexpr TierType kSecondByTier = TierType::DRAM; + +class PeerServiceDispatchTest : public ::testing::Test { + protected: + void SetUp() override { + registry_.Register(std::make_unique(kFirstByTier)); + registry_.Register(std::make_unique(kSecondByTier)); + pool_ = std::make_unique(®istry_, MakeSingleBackendPolicy()); + + // OS-assigned ports are not reachable through PeerServiceServer::Start + // (it does not report the bound port back), so probe a random high range + // rather than hardcoding one port that another process may already hold. + std::mt19937 rng{std::random_device{}()}; + std::uniform_int_distribution pick(20000, 60000); + for (int attempt = 0; attempt < 50 && !started_; ++attempt) { + const uint16_t port = static_cast(pick(rng)); + server_ = std::make_unique(*pool_); + if (server_->Start(port)) { + started_ = true; + auto channel = grpc::CreateChannel("127.0.0.1:" + std::to_string(port), + grpc::InsecureChannelCredentials()); + stub_ = ::umbp::UMBPPeer::NewStub(channel); + } else { + server_.reset(); + } + } + ASSERT_TRUE(started_) << "could not bind a peer service port"; + } + + void TearDown() override { + stub_.reset(); + if (server_) server_->Stop(); + server_.reset(); + } + + MediumBackend* Backend(TierType tier) { return registry_.Get(tier); } + + // The registry hands out MediumBackend; PublishBuffers is test scaffolding + // that only the mock has. Safe because SetUp registers nothing else. + MockBackend* Mock(TierType tier) { return static_cast(registry_.Get(tier)); } + + static ::umbp::TierType Proto(TierType t) { return static_cast<::umbp::TierType>(t); } + + // Allocate one slot and return the opaque, backend-tagged slot_id. + ::umbp::AllocateSlotResponse Allocate(const std::string& key, uint64_t size, TierType tier) { + ::umbp::AllocateSlotRequest req; + req.set_key(key); + req.set_size(size); + req.set_tier(Proto(tier)); + ::umbp::AllocateSlotResponse resp; + grpc::ClientContext ctx; + EXPECT_TRUE(stub_->AllocateSlot(&ctx, req, &resp).ok()); + return resp; + } + + bool Commit(uint64_t slot_id, const std::string& key) { + ::umbp::CommitSlotRequest req; + req.set_slot_id(slot_id); + req.set_key(key); + ::umbp::CommitSlotResponse resp; + grpc::ClientContext ctx; + EXPECT_TRUE(stub_->CommitSlot(&ctx, req, &resp).ok()); + return resp.success(); + } + + // Put `key` straight into one backend, bypassing the RPC surface, so a test + // can set up peer-side state without asserting on the path it is testing. + void SeedKey(TierType tier, const std::string& key, uint64_t size) { + auto* backend = Backend(tier); + ASSERT_NE(backend, nullptr); + auto allocated = backend->BatchAllocate({{key, size}}); + ASSERT_EQ(allocated[0].outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_TRUE(backend->BatchCommit({{allocated[0].slot_id, key}})[0].success); + } + + BackendRegistry registry_; + std::unique_ptr pool_; + std::unique_ptr server_; + std::unique_ptr<::umbp::UMBPPeer::Stub> stub_; + bool started_ = false; +}; + +// ---- Allocate --------------------------------------------------------------- + +TEST_F(PeerServiceDispatchTest, AllocateRoutesToTheBackendForItsTier) { + auto resp = Allocate("k", 128, kFirstByTier); + ASSERT_EQ(resp.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + ASSERT_TRUE(Commit(resp.slot_id(), "k")); + + EXPECT_EQ(Backend(kFirstByTier)->OwnedKeyCount(), 1u); + EXPECT_EQ(Backend(kSecondByTier)->OwnedKeyCount(), 0u); +} + +TEST_F(PeerServiceDispatchTest, AllocateForATierWithNoBackendFailsCleanly) { + // SSD is a valid wire tier with no backend registered here — a normal + // deployment shape, so it must be a FAILED outcome and not fall through to + // whichever backend happens to exist. + auto resp = Allocate("k", 128, TierType::SSD); + EXPECT_EQ(resp.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); + EXPECT_EQ(Backend(kFirstByTier)->OwnedKeyCount(), 0u); + EXPECT_EQ(Backend(kSecondByTier)->OwnedKeyCount(), 0u); +} + +// ---- The slot_id carries the backend id ------------------------------------- + +TEST_F(PeerServiceDispatchTest, ConcurrentSlotsOnTwoBackendsStayDistinct) { + // Each backend numbers its own slots from 1, so without the backend tag these + // two allocations would come back with the SAME opaque id and Commit could + // not tell them apart. + auto first = Allocate("a", 8, kFirstByTier); + auto second = Allocate("b", 8, kSecondByTier); + ASSERT_EQ(first.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + ASSERT_EQ(second.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + EXPECT_NE(first.slot_id(), second.slot_id()); + + EXPECT_TRUE(Commit(first.slot_id(), "a")); + EXPECT_TRUE(Commit(second.slot_id(), "b")); + EXPECT_EQ(Backend(kFirstByTier)->OwnedKeyCount(), 1u); + EXPECT_EQ(Backend(kSecondByTier)->OwnedKeyCount(), 1u); +} + +TEST_F(PeerServiceDispatchTest, CommitOfAnUnknownSlotFails) { + // Slot 0 decodes to tier UNKNOWN, which has no backend: "slot unknown". + EXPECT_FALSE(Commit(0, "nope")); +} + +// ---- Two live media share one buffer_index space ---------------------------- +// +// buffer_index is BACKEND-local: every backend numbers its buffers from 0, and +// every real backend publishes exactly one buffer, so on a peer with two live +// media both claim index 0. These tests pin the two wire surfaces where that +// used to lose data. Both fail without BufferMemoryDesc.backend_id. + +TEST_F(PeerServiceDispatchTest, GetPeerInfoPublishesEveryBackendsBuffers) { + // Distinct first descriptor byte per backend so we can tell whose is whose. + constexpr uint8_t kFirstTag = 0xA1; + constexpr uint8_t kSecondTag = 0xB2; + Mock(kFirstByTier)->PublishBuffers(1, kFirstTag); + Mock(kSecondByTier)->PublishBuffers(1, kSecondTag); + + ::umbp::GetPeerInfoRequest req; + ::umbp::GetPeerInfoResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->GetPeerInfo(&ctx, req, &resp).ok()); + + // Both media must be advertised. The bootstrap wire used to carry one, and + // dropping the other was not merely incomplete: the reader caches by index + // and then asks the peer to omit descriptors it believes it holds, so the + // unadvertised medium's pages were read against the advertised medium's + // memory — a hit, with the wrong bytes. + ASSERT_EQ(resp.buffer_descs_size(), 2); + + std::map backend_id_by_tag; + for (const auto& d : resp.buffer_descs()) { + ASSERT_FALSE(d.desc().empty()); + // Every backend numbers from 0, so the indices collide by design; only the + // backend_id separates them. + EXPECT_EQ(d.buffer_index(), 0u); + backend_id_by_tag[static_cast(d.desc()[0])] = d.backend_id(); + } + ASSERT_EQ(backend_id_by_tag.count(kFirstTag), 1u) << "first medium was not advertised"; + ASSERT_EQ(backend_id_by_tag.count(kSecondTag), 1u) << "second medium was not advertised"; + EXPECT_NE(backend_id_by_tag[kFirstTag], backend_id_by_tag[kSecondTag]) + << "two media sharing a backend_id collapses back into one address space"; + + // Uniform across backends, so one field still describes the whole peer. + EXPECT_EQ(resp.page_size(), Backend(kFirstByTier)->PageSize()); +} + +TEST_F(PeerServiceDispatchTest, BatchResolveAnswersKeysHeldByDifferentMedia) { + Mock(kFirstByTier)->PublishBuffers(1, 0xA1); + Mock(kSecondByTier)->PublishBuffers(1, 0xB2); + SeedKey(kFirstByTier, "in-first", 8); + SeedKey(kSecondByTier, "in-second", 8); + + ::umbp::BatchResolveKeysRequest req; + req.add_keys("in-first"); + req.add_keys("in-second"); + ::umbp::BatchResolveKeysResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->BatchResolveKeys(&ctx, req, &resp).ok()); + + // One response used to be served entirely by the FIRST medium with any hit, + // so "in-second" came back found=false purely because another key in the same + // batch hit earlier — a silent hit-rate loss on every mixed-media node. + ASSERT_EQ(resp.found_size(), 2); + ASSERT_EQ(resp.backend_id_size(), 2); + EXPECT_TRUE(resp.found(0)) << "key held by the first medium"; + EXPECT_TRUE(resp.found(1)) << "key held by the second medium was reported missing"; + + // Each key names the medium that served it, so its backend-local pages can be + // resolved against the right buffers. + EXPECT_NE(resp.backend_id(0), resp.backend_id(1)); + EXPECT_EQ(resp.tier(0), Proto(kFirstByTier)); + EXPECT_EQ(resp.tier(1), Proto(kSecondByTier)); +} + +TEST_F(PeerServiceDispatchTest, AllocateNamesTheBackendItsPagesBelongTo) { + Mock(kFirstByTier)->PublishBuffers(1, 0xA1); + Mock(kSecondByTier)->PublishBuffers(1, 0xB2); + + auto first = Allocate("a", 8, kFirstByTier); + auto second = Allocate("b", 8, kSecondByTier); + ASSERT_EQ(first.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + ASSERT_EQ(second.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + + // Same reason Commit needs the backend tagged into slot_id: the writer has to + // know which buffers its pages index into before it can RDMA into them. + EXPECT_NE(first.backend_id(), second.backend_id()); + for (const auto& d : first.descs()) EXPECT_EQ(d.backend_id(), first.backend_id()); + for (const auto& d : second.descs()) EXPECT_EQ(d.backend_id(), second.backend_id()); +} + +TEST_F(PeerServiceDispatchTest, AbortOfAnUnknownSlotIsIdempotentlyTrue) { + ::umbp::AbortSlotRequest req; + req.set_slot_id(0); + ::umbp::AbortSlotResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->AbortSlot(&ctx, req, &resp).ok()); + EXPECT_TRUE(resp.success()); +} + +TEST_F(PeerServiceDispatchTest, AbortReachesTheBackendThatOwnsTheSlot) { + auto allocated = Allocate("a", 8, kFirstByTier); + ASSERT_EQ(allocated.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + + ::umbp::AbortSlotRequest req; + req.set_slot_id(allocated.slot_id()); + ::umbp::AbortSlotResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->AbortSlot(&ctx, req, &resp).ok()); + EXPECT_TRUE(resp.success()); + + // The slot is gone, so the commit that would have followed now fails. + EXPECT_FALSE(Commit(allocated.slot_id(), "a")); + EXPECT_EQ(Backend(kFirstByTier)->OwnedKeyCount(), 0u); +} + +// ---- Mixed-tier batches ----------------------------------------------------- + +TEST_F(PeerServiceDispatchTest, BatchAllocateGroupsByTierButAnswersInRequestOrder) { + ::umbp::BatchAllocateSlotsRequest req; + const std::vector> wanted = { + {"e0", kFirstByTier}, + {"e1", TierType::SSD}, // no backend -> FAILED, but keeps its slot in the answer + {"e2", kSecondByTier}, + {"e3", kFirstByTier}, + }; + for (const auto& [key, tier] : wanted) { + auto* entry = req.add_entries(); + entry->set_key(key); + entry->set_size(16); + entry->set_tier(Proto(tier)); + } + + ::umbp::BatchAllocateSlotsResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->BatchAllocateSlots(&ctx, req, &resp).ok()); + ASSERT_EQ(resp.entries_size(), 4); + EXPECT_EQ(resp.entries(0).outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + EXPECT_EQ(resp.entries(1).outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); + EXPECT_EQ(resp.entries(2).outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + EXPECT_EQ(resp.entries(3).outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + + // Commit them back through the batch RPC — also mixed-tier, also positional. + ::umbp::BatchCommitSlotsRequest commit_req; + for (int i : {0, 2, 3}) { + auto* entry = commit_req.add_entries(); + entry->set_slot_id(resp.entries(i).slot_id()); + entry->set_key(wanted[i].first); + } + ::umbp::BatchCommitSlotsResponse commit_resp; + grpc::ClientContext commit_ctx; + ASSERT_TRUE(stub_->BatchCommitSlots(&commit_ctx, commit_req, &commit_resp).ok()); + ASSERT_EQ(commit_resp.success_size(), 3); + for (int i = 0; i < 3; ++i) EXPECT_TRUE(commit_resp.success(i)) << "entry " << i; + + // e0 and e3 landed on the first backend, e2 on the second. + EXPECT_EQ(Backend(kFirstByTier)->OwnedKeyCount(), 2u); + EXPECT_EQ(Backend(kSecondByTier)->OwnedKeyCount(), 1u); +} + +TEST_F(PeerServiceDispatchTest, BatchCommitOfAnUnroutableSlotReportsFalseInPlace) { + auto allocated = Allocate("a", 8, kSecondByTier); + ::umbp::BatchCommitSlotsRequest req; + auto* bad = req.add_entries(); + bad->set_slot_id(0); // backend 0, local slot 0: unknown + bad->set_key("bad"); + auto* good = req.add_entries(); + good->set_slot_id(allocated.slot_id()); + good->set_key("a"); + + ::umbp::BatchCommitSlotsResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->BatchCommitSlots(&ctx, req, &resp).ok()); + ASSERT_EQ(resp.success_size(), 2); + EXPECT_FALSE(resp.success(0)); + EXPECT_TRUE(resp.success(1)); +} + +// ---- Resolve / Evict carry no tier ------------------------------------------ + +TEST_F(PeerServiceDispatchTest, ResolveTakesTheFirstHitInMediaOrder) { + // Same key mirrored across both media with different sizes, so the response + // says which one served it. + SeedKey(kFirstByTier, "mirrored", 111); + SeedKey(kSecondByTier, "mirrored", 222); + + ::umbp::ResolveKeyRequest req; + req.set_key("mirrored"); + ::umbp::ResolveKeyResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->ResolveKey(&ctx, req, &resp).ok()); + ASSERT_TRUE(resp.found()); + EXPECT_EQ(resp.size(), 111u); +} + +TEST_F(PeerServiceDispatchTest, ResolveFindsAKeyHeldOnlyByALowerRankedMedium) { + SeedKey(kSecondByTier, "only-second", 77); + + ::umbp::ResolveKeyRequest req; + req.set_key("only-second"); + ::umbp::ResolveKeyResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->ResolveKey(&ctx, req, &resp).ok()); + ASSERT_TRUE(resp.found()); + EXPECT_EQ(resp.size(), 77u); +} + +TEST_F(PeerServiceDispatchTest, ResolveMissIsNotAnError) { + ::umbp::ResolveKeyRequest req; + req.set_key("absent"); + ::umbp::ResolveKeyResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->ResolveKey(&ctx, req, &resp).ok()); + EXPECT_FALSE(resp.found()); +} + +TEST_F(PeerServiceDispatchTest, BatchResolveReportsTheServingTier) { + SeedKey(kSecondByTier, "b-only", 55); + + ::umbp::BatchResolveKeysRequest req; + req.add_keys("absent"); + req.add_keys("b-only"); + ::umbp::BatchResolveKeysResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->BatchResolveKeys(&ctx, req, &resp).ok()); + ASSERT_EQ(resp.found_size(), 2); + ASSERT_EQ(resp.resolve_status_size(), 2); + EXPECT_FALSE(resp.found(0)); + EXPECT_TRUE(resp.found(1)); + EXPECT_EQ(resp.resolve_status(0), ::umbp::RESOLVE_STATUS_MISSING); + EXPECT_EQ(resp.resolve_status(1), ::umbp::RESOLVE_STATUS_FOUND); + ASSERT_EQ(resp.tier_size(), 2); + EXPECT_EQ(resp.tier(1), Proto(kSecondByTier)); + EXPECT_EQ(resp.size(1), 55u); +} + +TEST_F(PeerServiceDispatchTest, EvictFansOutAcrossEveryMedium) { + // A key mirrored across media must disappear from ALL of them, and the freed + // bytes master sizes its next round from are the sum. + SeedKey(kFirstByTier, "mirrored", 111); + SeedKey(kSecondByTier, "mirrored", 222); + + ::umbp::EvictKeyRequest req; + req.add_keys("mirrored"); + ::umbp::EvictKeyResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->EvictKey(&ctx, req, &resp).ok()); + ASSERT_EQ(resp.evicted_size(), 1); + EXPECT_EQ(resp.evicted(0).key(), "mirrored"); + EXPECT_EQ(resp.evicted(0).bytes_freed(), 333u); + EXPECT_EQ(Backend(kFirstByTier)->OwnedKeyCount(), 0u); + EXPECT_EQ(Backend(kSecondByTier)->OwnedKeyCount(), 0u); +} + +TEST_F(PeerServiceDispatchTest, EvictOfAnAbsentKeyReportsZeroBytes) { + ::umbp::EvictKeyRequest req; + req.add_keys("absent"); + ::umbp::EvictKeyResponse resp; + grpc::ClientContext ctx; + ASSERT_TRUE(stub_->EvictKey(&ctx, req, &resp).ok()); + ASSERT_EQ(resp.evicted_size(), 1); + EXPECT_EQ(resp.evicted(0).bytes_freed(), 0u); +} + +TEST(PeerServiceSameTierDispatch, NamedInstancesAllocateAndCommitIndependently) { + BackendRegistry registry; + ASSERT_TRUE(registry.Register("dram_a", std::make_unique(TierType::DRAM))); + ASSERT_TRUE(registry.Register("dram_b", std::make_unique(TierType::DRAM))); + + PeerPool pool(®istry, MakeSingleBackendPolicy()); + PeerServiceServer server(pool); + std::mt19937 rng{std::random_device{}()}; + std::uniform_int_distribution pick(20000, 60000); + uint16_t bound = 0; + for (int attempt = 0; attempt < 50 && bound == 0; ++attempt) { + const uint16_t candidate = static_cast(pick(rng)); + if (server.Start(candidate)) bound = candidate; + } + ASSERT_NE(bound, 0); + + auto channel = + grpc::CreateChannel("127.0.0.1:" + std::to_string(bound), grpc::InsecureChannelCredentials()); + auto stub = ::umbp::UMBPPeer::NewStub(channel); + const auto allocate = [&](const std::string& key, const std::string& backend_name) { + ::umbp::AllocateSlotRequest request; + request.set_key(key); + request.set_size(8); + request.set_tier(::umbp::TIER_DRAM); + request.set_backend_name(backend_name); + ::umbp::AllocateSlotResponse response; + grpc::ClientContext context; + EXPECT_TRUE(stub->AllocateSlot(&context, request, &response).ok()); + return response; + }; + const auto commit = [&](uint64_t slot_id, const std::string& key) { + ::umbp::CommitSlotRequest request; + request.set_slot_id(slot_id); + request.set_key(key); + ::umbp::CommitSlotResponse response; + grpc::ClientContext context; + EXPECT_TRUE(stub->CommitSlot(&context, request, &response).ok()); + return response.success(); + }; + + auto first = allocate("a", "dram_a"); + auto second = allocate("b", "dram_b"); + ASSERT_EQ(first.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + ASSERT_EQ(second.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALLOCATED); + EXPECT_NE(first.backend_id(), second.backend_id()); + EXPECT_NE(first.slot_id(), second.slot_id()); + EXPECT_TRUE(commit(first.slot_id(), "a")); + EXPECT_TRUE(commit(second.slot_id(), "b")); + EXPECT_EQ(registry.Get("dram_a")->OwnedKeyCount(), 1u); + EXPECT_EQ(registry.Get("dram_b")->OwnedKeyCount(), 1u); + + auto missing_backend = allocate("c", "missing"); + EXPECT_EQ(missing_backend.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); + server.Stop(); +} + +// ---- A registry-less peer still answers ------------------------------------- + +TEST(PeerServiceNoRegistry, AnswersWithoutCrashing) { + PeerServiceServer server(nullptr); + std::mt19937 rng{std::random_device{}()}; + std::uniform_int_distribution pick(20000, 60000); + uint16_t bound = 0; + for (int attempt = 0; attempt < 50 && bound == 0; ++attempt) { + const uint16_t port = static_cast(pick(rng)); + if (server.Start(port)) bound = port; + } + ASSERT_NE(bound, 0); + + auto channel = + grpc::CreateChannel("127.0.0.1:" + std::to_string(bound), grpc::InsecureChannelCredentials()); + auto stub = ::umbp::UMBPPeer::NewStub(channel); + + ::umbp::AllocateSlotRequest alloc_req; + alloc_req.set_key("k"); + alloc_req.set_size(8); + alloc_req.set_tier(::umbp::TIER_DRAM); + ::umbp::AllocateSlotResponse alloc_resp; + grpc::ClientContext alloc_ctx; + ASSERT_TRUE(stub->AllocateSlot(&alloc_ctx, alloc_req, &alloc_resp).ok()); + EXPECT_EQ(alloc_resp.outcome(), ::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); + + ::umbp::ResolveKeyRequest resolve_req; + resolve_req.set_key("k"); + ::umbp::ResolveKeyResponse resolve_resp; + grpc::ClientContext resolve_ctx; + ASSERT_TRUE(stub->ResolveKey(&resolve_ctx, resolve_req, &resolve_resp).ok()); + EXPECT_FALSE(resolve_resp.found()); + + server.Stop(); +} + +// ---- One peer service owns a port ------------------------------------------- +// +// gRPC enables SO_REUSEPORT by default for TCP servers on Linux, so before +// PeerServiceServer::Start passed GRPC_ARG_ALLOW_REUSEPORT=0 a SECOND server +// bound the same port successfully and the kernel split incoming connections +// between the two. A client dialing that address was then answered, some +// fraction of the time, by a peer service with entirely different backends +// registered — a wrong answer, not a connection error. +// +// That is a production hazard (two UMBP processes sharing a peer port silently +// mis-serve each other's traffic) and it was also the cause of the umbp suite's +// random cross-test flakiness: PeerServiceDispatchTest::SetUp picks a random +// port in [20000,60000) and relies on Start() failing to detect a collision, so +// with reuseport on it never retried and instead talked to the wrong server. +TEST(PeerServicePortExclusivity, SecondServerCannotBindTheSamePort) { + BackendRegistry registry_a; + BackendRegistry registry_b; + PeerPool pool_a(®istry_a, MakeSingleBackendPolicy()); + PeerPool pool_b(®istry_b, MakeSingleBackendPolicy()); + PeerServiceServer first(pool_a); + PeerServiceServer second(pool_b); + + std::mt19937 rng{std::random_device{}()}; + std::uniform_int_distribution pick(20000, 60000); + uint16_t port = 0; + for (int attempt = 0; attempt < 50 && port == 0; ++attempt) { + const uint16_t candidate = static_cast(pick(rng)); + if (first.Start(candidate)) port = candidate; + } + ASSERT_NE(port, 0) << "could not bind any port for the first server"; + + EXPECT_FALSE(second.Start(port)) + << "a second peer service bound port " << port + << " — SO_REUSEPORT is on, so clients would be split between two servers"; + + second.Stop(); + first.Stop(); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_peer_ssd_eviction.cpp b/tests/cpp/umbp/distributed/test_peer_ssd_eviction.cpp index 591d3b1d3..be19c8f93 100644 --- a/tests/cpp/umbp/distributed/test_peer_ssd_eviction.cpp +++ b/tests/cpp/umbp/distributed/test_peer_ssd_eviction.cpp @@ -38,7 +38,7 @@ #include #include -#include "umbp/distributed/peer/peer_ssd_manager.h" +#include "umbp/distributed/peer/ssd/peer_ssd_manager.h" #include "umbp/local/tiers/tier_backend.h" namespace mori::umbp { @@ -223,6 +223,13 @@ TEST(PeerSsdEviction, WatermarkTriggersEvictionDownToLow) { EXPECT_EQ(CountKind(events, KvEvent::Kind::REMOVE), 2); } +TEST(PeerSsdEviction, SuccessfulWriteDoesNotEvictItself) { + auto h = MakeHarness(/*capacity=*/100, /*high=*/0.9, /*low=*/0.7); + std::string value(90, 'x'); + ASSERT_TRUE(h.mgr->Write("only", OneSeg(value), value.size())); + EXPECT_TRUE(h.mgr->Exists("only")); +} + TEST(PeerSsdEviction, EnospcTriggersEvictThenRetry) { // Fill to 800/1000 (below the 0.9 high watermark, so no watermark eviction // fires during the fill), then write a 300-byte value that overflows the diff --git a/tests/cpp/umbp/distributed/test_peer_ssd_manager.cpp b/tests/cpp/umbp/distributed/test_peer_ssd_manager.cpp index d75dacf5a..0cea25a13 100644 --- a/tests/cpp/umbp/distributed/test_peer_ssd_manager.cpp +++ b/tests/cpp/umbp/distributed/test_peer_ssd_manager.cpp @@ -30,8 +30,7 @@ #include #include -#include "umbp/distributed/peer/owned_location_source.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" +#include "umbp/distributed/peer/ssd/peer_ssd_manager.h" namespace mori::umbp { namespace { @@ -179,59 +178,5 @@ TEST_F(PeerSsdManagerTest, PrepareReadRejectsOverCapBeforeIo) { EXPECT_EQ(out.size, value.size()); } -// ---- Unified owned-location source aggregation ------------------------------ - -// Minimal OwnedLocationSource that replays a fixed event list, used to verify -// MasterClient's multi-source concat logic without a live master. -class FakeSource : public OwnedLocationSource { - public: - explicit FakeSource(std::vector events) : events_(std::move(events)) {} - std::vector DrainPendingEvents() override { - auto out = events_; - drained_ = true; - return out; - } - std::vector SnapshotOwnedKeys() const override { return events_; } - std::vector SnapshotOwnedKeysForFullSync() override { - drained_ = true; // outbox dropped by the authoritative full-sync snapshot - return events_; - } - bool drained_ = false; - - private: - std::vector events_; -}; - -TEST(OwnedLocationSourceAgg, DrainAndSnapshotConcatAcrossSourcesInOrder) { - FakeSource dram({{KvEvent::Kind::ADD, "d1", TierType::DRAM, 10}, - {KvEvent::Kind::ADD, "d2", TierType::DRAM, 20}}); - FakeSource ssd({{KvEvent::Kind::ADD, "s1", TierType::SSD, 30}}); - - std::vector sources = {&dram, &ssd}; - - auto drained = DrainAllSources(sources); - ASSERT_EQ(drained.size(), 3u); - EXPECT_EQ(drained[0].key, "d1"); - EXPECT_EQ(drained[0].tier, TierType::DRAM); - EXPECT_EQ(drained[1].key, "d2"); - EXPECT_EQ(drained[2].key, "s1"); - EXPECT_EQ(drained[2].tier, TierType::SSD); - EXPECT_TRUE(dram.drained_); - EXPECT_TRUE(ssd.drained_); - - auto snap = SnapshotAllSourcesForFullSync(sources); - ASSERT_EQ(snap.size(), 3u); - EXPECT_EQ(snap[2].tier, TierType::SSD); -} - -TEST(OwnedLocationSourceAgg, NullSourcesAreSkipped) { - FakeSource only({{KvEvent::Kind::ADD, "x", TierType::SSD, 1}}); - std::vector sources = {nullptr, &only, nullptr}; - auto drained = DrainAllSources(sources); - ASSERT_EQ(drained.size(), 1u); - EXPECT_EQ(drained[0].key, "x"); - EXPECT_TRUE(SnapshotAllSourcesForFullSync({nullptr}).empty()); -} - } // namespace } // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_peer_ssd_read_rpc.cpp b/tests/cpp/umbp/distributed/test_peer_ssd_read_rpc.cpp deleted file mode 100644 index 8ffd3b2f5..000000000 --- a/tests/cpp/umbp/distributed/test_peer_ssd_read_rpc.cpp +++ /dev/null @@ -1,258 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -// -// RPC-level integration test for the SSD read path: a real PeerServiceServer -// (backed by a real PeerSsdManager / POSIX SSDTier) served over a gRPC loopback -// channel. It exercises prepare -> read-from-staging -> release / TTL and, -// crucially, asserts that OK / NOT_FOUND / NO_SLOT / SIZE_TOO_LARGE are each -// reported as distinct statuses so a transient failure is never collapsed into -// a NOT_FOUND miss. RDMA is intentionally out of scope here (the staging buffer -// is read directly); the full BatchGet -> RDMA path needs a live cluster. -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/config.h" -#include "umbp/distributed/peer/peer_service.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" -#include "umbp_peer.grpc.pb.h" - -namespace mori::umbp { -namespace { - -namespace fs = std::filesystem; - -constexpr size_t kStagingSize = 4096; -constexpr int kNumReadSlots = 4; // -> 1024 B per slot -constexpr int kLeaseTimeoutS = 2; - -uint16_t AllocPort() { - static std::atomic next{51300}; - return next.fetch_add(1); -} - -class PeerSsdReadRpcTest : public ::testing::Test { - protected: - void SetUp() override { - staging_buffer_ = std::malloc(kStagingSize); - ASSERT_NE(staging_buffer_, nullptr); - std::memset(staging_buffer_, 0, kStagingSize); - - dir_ = fs::temp_directory_path() / - ("umbp_ssd_rpc_" + std::to_string(::getpid()) + "_" + std::to_string(AllocPort())); - fs::remove_all(dir_); - - PeerSsdConfig cfg; - cfg.enabled = true; - cfg.ssd.enabled = true; - cfg.ssd.storage_dir = dir_.string(); - cfg.ssd.capacity_bytes = 1 << 20; - cfg.ssd.io.backend = UMBPIoBackend::Posix; // avoid io_uring container flakiness - peer_ssd_ = std::make_unique(cfg); - - // Fake staging MemoryDesc bytes — GetPeerInfo just echoes them; this test - // reads the staging buffer directly rather than RDMA-ing it. - staging_desc_ = {0xAB, 0xCD}; - - port_ = AllocPort(); - server_ = std::make_unique( - /*dram_alloc=*/nullptr, peer_ssd_.get(), staging_buffer_, kStagingSize, staging_desc_, - kNumReadSlots, std::chrono::seconds(kLeaseTimeoutS)); - ASSERT_TRUE(server_->Start(port_)); - std::this_thread::sleep_for(std::chrono::milliseconds(150)); - - auto channel = grpc::CreateChannel("localhost:" + std::to_string(port_), - grpc::InsecureChannelCredentials()); - stub_ = ::umbp::UMBPPeer::NewStub(channel); - } - - void TearDown() override { - server_->Stop(); - server_.reset(); - peer_ssd_.reset(); - std::free(staging_buffer_); - fs::remove_all(dir_); - } - - void WriteSsd(const std::string& key, const std::string& data) { - ASSERT_TRUE(peer_ssd_->Write(key, {{data.data(), data.size()}}, data.size())); - } - - ::umbp::PrepareSsdReadResponse Prepare(const std::string& key, uint64_t max_size) { - ::umbp::PrepareSsdReadRequest req; - req.set_key(key); - req.set_max_size(max_size); - ::umbp::PrepareSsdReadResponse resp; - grpc::ClientContext ctx; - EXPECT_TRUE(stub_->PrepareSsdRead(&ctx, req, &resp).ok()); - return resp; - } - - void* staging_buffer_ = nullptr; - fs::path dir_; - std::vector staging_desc_; - uint16_t port_ = 0; - std::unique_ptr peer_ssd_; - std::unique_ptr server_; - std::unique_ptr<::umbp::UMBPPeer::Stub> stub_; -}; - -TEST_F(PeerSsdReadRpcTest, OkReadsBytesIntoStaging) { - const std::string data = "ssd-read-rpc-ok"; - WriteSsd("k-ok", data); - - auto resp = Prepare("k-ok", data.size()); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK); - EXPECT_EQ(resp.size(), data.size()); - EXPECT_LT(resp.staging_offset(), kStagingSize); - EXPECT_GT(resp.lease_id(), 0u); - std::string loaded(static_cast(staging_buffer_) + resp.staging_offset(), - resp.size()); - EXPECT_EQ(loaded, data); -} - -TEST_F(PeerSsdReadRpcTest, NotFoundIsADistinctMiss) { - auto resp = Prepare("absent", 64); - EXPECT_EQ(resp.status(), ::umbp::SSD_READ_NOT_FOUND); -} - -TEST_F(PeerSsdReadRpcTest, SizeTooLargeIsDistinct) { - // A key bigger than one slot (1024 B) must report SIZE_TOO_LARGE, not OK and - // not NOT_FOUND. - const std::string big(2048, 'q'); - WriteSsd("k-big", big); - auto resp = Prepare("k-big", kStagingSize); - EXPECT_EQ(resp.status(), ::umbp::SSD_READ_SIZE_TOO_LARGE); -} - -// The key assertion the review asked for: slot exhaustion is NO_SLOT -// (retryable), never collapsed into NOT_FOUND. A present key and an absent key -// under exhaustion are distinguishable. -TEST_F(PeerSsdReadRpcTest, NoSlotIsDistinctFromNotFound) { - std::vector<::umbp::PrepareSsdReadResponse> held; - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "hold-" + std::to_string(i); - WriteSsd(key, "payload"); - auto resp = Prepare(key, 64); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK); - held.push_back(resp); // keep the lease held (no release) so slots stay busy - } - - // A present key with all slots busy -> NO_SLOT (retryable), NOT a miss. - WriteSsd("present-extra", "payload"); - EXPECT_EQ(Prepare("present-extra", 64).status(), ::umbp::SSD_READ_NO_SLOT); - - // An absent key under the same exhaustion is still NO_SLOT (slot check - // precedes the key lookup), so the caller cannot mistake exhaustion for a - // definitive miss. - EXPECT_EQ(Prepare("absent-extra", 64).status(), ::umbp::SSD_READ_NO_SLOT); -} - -// Many concurrent readers contend for a fixed pool of staging slots: at most -// kNumReadSlots win OK, the rest get NO_SLOT (a retryable transient), and NEVER -// NOT_FOUND for a present key. Also exercises the staging observability -// (slot_full_rejects counter + the in-use gauge accessor). -TEST_F(PeerSsdReadRpcTest, ConcurrentReadersExhaustSlotsWithoutFalseMiss) { - const std::string data = "concurrent-payload"; - for (int i = 0; i < 32; ++i) WriteSsd("ck-" + std::to_string(i), data); - - const uint64_t slot_full_before = server_->Metrics().slot_full_rejects.load(); - - constexpr int kReaders = 24; // >> kNumReadSlots, leases held (never released) - std::atomic ok{0}, no_slot{0}, other{0}; - std::vector threads; - for (int i = 0; i < kReaders; ++i) { - threads.emplace_back([&, i] { - auto resp = Prepare("ck-" + std::to_string(i), data.size()); - switch (resp.status()) { - case ::umbp::SSD_READ_OK: - ok.fetch_add(1); - break; - case ::umbp::SSD_READ_NO_SLOT: - no_slot.fetch_add(1); - break; - default: - other.fetch_add(1); // NOT_FOUND / SIZE_TOO_LARGE / ERROR must never happen here - break; - } - }); - } - for (auto& t : threads) t.join(); - - EXPECT_EQ(other.load(), 0) << "present keys never report a false miss under contention"; - EXPECT_LE(ok.load(), kNumReadSlots) << "at most one OK per staging slot"; - EXPECT_EQ(ok.load() + no_slot.load(), kReaders); - EXPECT_GT(no_slot.load(), 0) << "with more readers than slots, some must see NO_SLOT"; - - // The NO_SLOT rejections were counted, and the gauge sees the held leases. - EXPECT_GE(server_->Metrics().slot_full_rejects.load() - slot_full_before, - static_cast(no_slot.load())); - EXPECT_EQ(server_->SnapshotReadSlotsInUse(), static_cast(ok.load())); -} - -// A best-effort release frees the slot; double release reports false. -TEST_F(PeerSsdReadRpcTest, ReleaseFreesSlotAndIsBestEffort) { - WriteSsd("k-rel", "payload"); - auto resp = Prepare("k-rel", 64); - ASSERT_EQ(resp.status(), ::umbp::SSD_READ_OK); - - ::umbp::ReleaseSsdLeaseRequest rel; - rel.set_lease_id(resp.lease_id()); - ::umbp::ReleaseSsdLeaseResponse rel_resp; - grpc::ClientContext ctx; - ASSERT_TRUE(stub_->ReleaseSsdLease(&ctx, rel, &rel_resp).ok()); - EXPECT_TRUE(rel_resp.success()); - - ::umbp::ReleaseSsdLeaseResponse rel_resp2; - grpc::ClientContext ctx2; - ASSERT_TRUE(stub_->ReleaseSsdLease(&ctx2, rel, &rel_resp2).ok()); - EXPECT_FALSE(rel_resp2.success()) << "double release is a no-op"; -} - -// Leased slots are reclaimed by TTL even without a release, so a fresh prepare -// succeeds after the lease elapses (slot lifecycle: Leased -> reclaimed). -TEST_F(PeerSsdReadRpcTest, LeasedSlotsReclaimedByTtl) { - for (int i = 0; i < kNumReadSlots; ++i) { - const std::string key = "ttl-" + std::to_string(i); - WriteSsd(key, "payload"); - ASSERT_EQ(Prepare(key, 64).status(), ::umbp::SSD_READ_OK); // never released - } - EXPECT_EQ(Prepare("ttl-0", 64).status(), ::umbp::SSD_READ_NO_SLOT); // all busy - - std::this_thread::sleep_for(std::chrono::seconds(kLeaseTimeoutS + 1)); - - WriteSsd("ttl-after", "payload"); - EXPECT_EQ(Prepare("ttl-after", 64).status(), ::umbp::SSD_READ_OK) << "TTL should reclaim a slot"; -} - -} // namespace -} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_peer_ssd_single_flight.cpp b/tests/cpp/umbp/distributed/test_peer_ssd_single_flight.cpp new file mode 100644 index 000000000..20e382115 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_peer_ssd_single_flight.cpp @@ -0,0 +1,302 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Single-flight read coalescing in PeerSsdManager. +// +// Motivation: with MLA + TP every attention-TP rank GETs a byte-identical key, +// so one logical page was being read from the drive tp_size times over. A +// production TP8 trace measured 85.6% of reads redundant, and per-key device +// latency rising ~linearly with same-key fan-out (3.3ms solo -> 27.7ms at 6 +// concurrent). These tests pin the behaviour that removes it. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/ssd/peer_ssd_manager.h" +#include "umbp/local/tiers/tier_backend.h" + +namespace mori::umbp { +namespace { + +constexpr int kFanout = 8; // TP8: one leader + 7 followers + +// In-memory backend whose reads can be held open, so every reader is guaranteed +// to be inside the merge window at the same time. +class GatedBackend : public TierBackend { + public: + explicit GatedBackend(size_t capacity) + : TierBackend(StorageTier::LOCAL_SSD), capacity_(capacity) {} + + bool Write(const std::string& key, const void* data, size_t size) override { + std::lock_guard lk(mu_); + if (used_ + size > capacity_) return false; + store_[key].assign(static_cast(data), static_cast(data) + size); + used_ += size; + return true; + } + + bool ReadIntoPtr(const std::string& key, uintptr_t dst, size_t size) override { + { + std::unique_lock lk(gate_mu_); + ++reads_started_; + started_cv_.notify_all(); + gate_cv_.wait(lk, [this] { return !blocked_; }); + if (fail_) return false; + } + std::lock_guard lk(mu_); + auto it = store_.find(key); + if (it == store_.end() || it->second.size() != size) return false; + std::memcpy(reinterpret_cast(dst), it->second.data(), size); + return true; + } + + bool Exists(const std::string& key) const override { + std::lock_guard lk(mu_); + return store_.count(key) != 0; + } + bool Evict(const std::string& key) override { + std::lock_guard lk(mu_); + auto it = store_.find(key); + if (it == store_.end()) return false; + used_ -= it->second.size(); + store_.erase(it); + return true; + } + std::pair Capacity() const override { + std::lock_guard lk(mu_); + return {used_, capacity_}; + } + void Clear() override { + std::lock_guard lk(mu_); + store_.clear(); + used_ = 0; + } + + void BlockReads() { + std::lock_guard lk(gate_mu_); + blocked_ = true; + } + void UnblockReads() { + { + std::lock_guard lk(gate_mu_); + blocked_ = false; + } + gate_cv_.notify_all(); + } + void SetFailReads(bool f) { + std::lock_guard lk(gate_mu_); + fail_ = f; + } + void WaitReadsStarted(int n) { + std::unique_lock lk(gate_mu_); + started_cv_.wait(lk, [&] { return reads_started_ >= n; }); + } + int reads_started() const { + std::lock_guard lk(gate_mu_); + return reads_started_; + } + + private: + mutable std::mutex mu_; + std::unordered_map> store_; + size_t used_ = 0; + size_t capacity_; + + mutable std::mutex gate_mu_; + std::condition_variable gate_cv_; + std::condition_variable started_cv_; + bool blocked_ = false; + bool fail_ = false; + int reads_started_ = 0; +}; + +struct Harness { + GatedBackend* backend; + std::unique_ptr mgr; +}; + +Harness MakeHarness(bool single_flight = true) { + auto be = std::make_unique(1'000'000); + GatedBackend* raw = be.get(); + return Harness{raw, std::make_unique(std::move(be), 0.9, 0.7, single_flight)}; +} + +std::vector> OneSeg(const std::string& s) { + return {{s.data(), s.size()}}; +} + +// Spin until `n` requesters have found a read already in flight. Deterministic +// stand-in for "all followers have attached" — polling a public counter beats +// sleeping, which would make the merge window a race. +void WaitForDups(const PeerSsdManager& mgr, uint64_t n) { + while (mgr.ReadDup() < n) std::this_thread::yield(); +} + +// All fan-out readers of one key collapse to a single device read, and every +// one of them still receives the correct bytes. +TEST(PeerSsdSingleFlight, ConcurrentSameKeyReadsIssueOneDeviceRead) { + auto h = MakeHarness(); + const std::string payload = "kv-page-payload"; + ASSERT_TRUE(h.mgr->Write("K", OneSeg(payload), payload.size())); + + h.backend->BlockReads(); + std::vector> bufs(kFanout, std::vector(payload.size(), '\0')); + std::vector outs(kFanout); + std::vector readers; + for (int i = 0; i < kFanout; ++i) { + readers.emplace_back( + [&, i] { outs[i] = h.mgr->PrepareRead("K", bufs[i].data(), bufs[i].size()); }); + } + h.backend->WaitReadsStarted(1); // leader is on the device + WaitForDups(*h.mgr, kFanout - 1); // all 7 followers have attached + h.backend->UnblockReads(); + for (auto& t : readers) t.join(); + + EXPECT_EQ(h.backend->reads_started(), 1) << "followers must not touch the drive"; + EXPECT_EQ(h.mgr->ReadMerged(), static_cast(kFanout - 1)); + for (int i = 0; i < kFanout; ++i) { + EXPECT_EQ(outs[i].status, SsdReadStatus::kOk) << "reader " << i; + EXPECT_EQ(std::string(bufs[i].data(), outs[i].size), payload) << "reader " << i; + } +} + +// With coalescing off every requester reads the drive — the pre-single-flight +// behaviour, and the A/B baseline the UMBP_SSD_SINGLE_FLIGHT knob selects. +TEST(PeerSsdSingleFlight, DisabledIssuesOneDeviceReadPerRequester) { + auto h = MakeHarness(/*single_flight=*/false); + const std::string payload = "kv-page-payload"; + ASSERT_TRUE(h.mgr->Write("K", OneSeg(payload), payload.size())); + + h.backend->BlockReads(); + std::vector> bufs(kFanout, std::vector(payload.size(), '\0')); + std::vector outs(kFanout); + std::vector readers; + for (int i = 0; i < kFanout; ++i) { + readers.emplace_back( + [&, i] { outs[i] = h.mgr->PrepareRead("K", bufs[i].data(), bufs[i].size()); }); + } + h.backend->WaitReadsStarted(kFanout); // all of them reach the device + h.backend->UnblockReads(); + for (auto& t : readers) t.join(); + + EXPECT_EQ(h.backend->reads_started(), kFanout); + EXPECT_EQ(h.mgr->ReadMerged(), 0u); + for (int i = 0; i < kFanout; ++i) { + EXPECT_EQ(outs[i].status, SsdReadStatus::kOk); + EXPECT_EQ(std::string(bufs[i].data(), outs[i].size), payload); + } +} + +// A leader whose read fails must release its followers with an error rather +// than leaving them to hit the 30s timeout. +TEST(PeerSsdSingleFlight, LeaderFailureReleasesFollowersWithError) { + auto h = MakeHarness(); + const std::string payload = "kv-page-payload"; + ASSERT_TRUE(h.mgr->Write("K", OneSeg(payload), payload.size())); + + h.backend->BlockReads(); + h.backend->SetFailReads(true); + std::vector> bufs(kFanout, std::vector(payload.size(), '\0')); + std::vector outs(kFanout); + std::vector readers; + for (int i = 0; i < kFanout; ++i) { + readers.emplace_back( + [&, i] { outs[i] = h.mgr->PrepareRead("K", bufs[i].data(), bufs[i].size()); }); + } + h.backend->WaitReadsStarted(1); + WaitForDups(*h.mgr, kFanout - 1); + h.backend->UnblockReads(); + for (auto& t : readers) t.join(); // hangs here if the failure is not published + + for (int i = 0; i < kFanout; ++i) { + EXPECT_EQ(outs[i].status, SsdReadStatus::kError) << "reader " << i; + } +} + +// Distinct keys share no episode, so nothing merges — this is the MHA case, +// where keys carry a per-rank suffix and there is no redundancy to remove. +TEST(PeerSsdSingleFlight, DistinctKeysDoNotMerge) { + auto h = MakeHarness(); + for (int i = 0; i < kFanout; ++i) { + ASSERT_TRUE(h.mgr->Write("K" + std::to_string(i), OneSeg("payload!"), 8)); + } + + h.backend->BlockReads(); + std::vector> bufs(kFanout, std::vector(8, '\0')); + std::vector outs(kFanout); + std::vector readers; + for (int i = 0; i < kFanout; ++i) { + readers.emplace_back([&, i] { + outs[i] = h.mgr->PrepareRead("K" + std::to_string(i), bufs[i].data(), bufs[i].size()); + }); + } + h.backend->WaitReadsStarted(kFanout); + h.backend->UnblockReads(); + for (auto& t : readers) t.join(); + + EXPECT_EQ(h.backend->reads_started(), kFanout); + EXPECT_EQ(h.mgr->ReadMerged(), 0u); + for (int i = 0; i < kFanout; ++i) { + EXPECT_EQ(outs[i].status, SsdReadStatus::kOk); + EXPECT_EQ(std::string(bufs[i].data(), outs[i].size), "payload!"); + } +} + +// Back-to-back rounds on one key: a follower that has not yet woken when the +// next leader claims the key must still report its own leader's result. This +// is the generation hazard that per-episode state exists to prevent. +TEST(PeerSsdSingleFlight, SequentialRoundsStayIndependent) { + auto h = MakeHarness(); + const std::string payload = "kv-page-payload"; + ASSERT_TRUE(h.mgr->Write("K", OneSeg(payload), payload.size())); + + for (int round = 0; round < 20; ++round) { + h.backend->BlockReads(); + std::vector> bufs(kFanout, std::vector(payload.size(), '\0')); + std::vector outs(kFanout); + std::vector readers; + for (int i = 0; i < kFanout; ++i) { + readers.emplace_back( + [&, i] { outs[i] = h.mgr->PrepareRead("K", bufs[i].data(), bufs[i].size()); }); + } + h.backend->WaitReadsStarted(round + 1); // reads_started() is cumulative + h.backend->UnblockReads(); + for (auto& t : readers) t.join(); + for (int i = 0; i < kFanout; ++i) { + EXPECT_EQ(outs[i].status, SsdReadStatus::kOk) << "round " << round << " reader " << i; + EXPECT_EQ(std::string(bufs[i].data(), outs[i].size), payload); + } + } +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_pool_client_batch_put.cpp b/tests/cpp/umbp/distributed/test_pool_client_batch_put.cpp index 9a92ddc30..0595b000c 100644 --- a/tests/cpp/umbp/distributed/test_pool_client_batch_put.cpp +++ b/tests/cpp/umbp/distributed/test_pool_client_batch_put.cpp @@ -32,6 +32,7 @@ // spdlog output by the substring "BatchPut: src not registered for key=". #include +#include #include #include #include @@ -39,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -160,16 +162,10 @@ class BatchPutWarnTest : public ::testing::Test { void SetUp() override { caller_buf_ = std::malloc(kCallerBuf); registered_buf_ = std::malloc(kCallerBuf); - target_buf_ = std::malloc(kRemoteCap); - caller_local_ = std::malloc(kPageSize); ASSERT_NE(caller_buf_, nullptr); ASSERT_NE(registered_buf_, nullptr); - ASSERT_NE(target_buf_, nullptr); - ASSERT_NE(caller_local_, nullptr); std::memset(caller_buf_, 0, kCallerBuf); std::memset(registered_buf_, 0, kCallerBuf); - std::memset(target_buf_, 0, kRemoteCap); - std::memset(caller_local_, 0, kPageSize); MasterServerConfig master_cfg; master_cfg.listen_address = "0.0.0.0:0"; @@ -193,8 +189,7 @@ class BatchPutWarnTest : public ::testing::Test { cfg_caller.io_engine.port = 0; cfg_caller.peer_service_port = NextPeerServicePort(); cfg_caller.dram_page_size = kPageSize; - cfg_caller.dram_buffers = {{caller_local_, kPageSize}}; - cfg_caller.tier_capacities = {{TierType::DRAM, {kPageSize, kPageSize}}}; + cfg_caller.dram.buffer_sizes = {kPageSize}; caller_ = std::make_unique(std::move(cfg_caller)); ASSERT_TRUE(caller_->Init()); @@ -206,8 +201,7 @@ class BatchPutWarnTest : public ::testing::Test { cfg_target.io_engine.port = 0; cfg_target.peer_service_port = NextPeerServicePort(); cfg_target.dram_page_size = kPageSize; - cfg_target.dram_buffers = {{target_buf_, kRemoteCap}}; - cfg_target.tier_capacities = {{TierType::DRAM, {kRemoteCap, kRemoteCap}}}; + cfg_target.dram.buffer_sizes = {kRemoteCap}; target_ = std::make_unique(std::move(cfg_target)); ASSERT_TRUE(target_->Init()); } @@ -219,8 +213,6 @@ class BatchPutWarnTest : public ::testing::Test { if (server_thread_.joinable()) server_thread_.join(); std::free(caller_buf_); std::free(registered_buf_); - std::free(caller_local_); - std::free(target_buf_); } // Build a batch backed by `base` (one slot per page). Caller chooses @@ -243,8 +235,6 @@ class BatchPutWarnTest : public ::testing::Test { void* caller_buf_ = nullptr; void* registered_buf_ = nullptr; - void* target_buf_ = nullptr; - void* caller_local_ = nullptr; std::unique_ptr master_; std::thread server_thread_; std::unique_ptr caller_; @@ -375,5 +365,201 @@ TEST_F(BatchPutWarnTest, ZeroSizeGetEntriesFailOthersSucceed) { EXPECT_TRUE(gres[2]) << "key=" << gkeys[2]; } +// --------------------------------------------------------------------------- +// Multi-page objects exercise run coalescing in LocalCopyEngine. Every other +// test here uses exactly one page per key, so the merge path never runs there. +// +// A wrong run boundary or a miscomputed run length still returns true, so these +// assert bytes rather than status. The fill pattern varies per byte so a +// shifted, duplicated or truncated run is visible. +// --------------------------------------------------------------------------- + +namespace { +void FillPattern(char* buf, size_t bytes, uint8_t seed) { + for (size_t i = 0; i < bytes; ++i) { + buf[i] = static_cast((i * 31u + seed) & 0xFFu); + } +} +} // namespace + +TEST_F(BatchPutWarnTest, MultiPageObjectRoundTripsByteExact) { + constexpr size_t kPages = 4; + constexpr size_t kSize = kPages * kPageSize; // whole pages, no partial tail + + std::vector src(kSize); + FillPattern(src.data(), kSize, 0x5A); + std::vector keys = {"mp-full"}; + std::vector srcs = {src.data()}; + std::vector sizes = {kSize}; + auto pres = target_->BatchPut(keys, srcs, sizes); + ASSERT_EQ(pres.size(), 1u); + ASSERT_TRUE(pres[0]); + + std::vector dst(kSize, 0); + std::vector dsts = {dst.data()}; + auto gres = target_->BatchGet(keys, dsts, sizes); + ASSERT_EQ(gres.size(), 1u); + ASSERT_TRUE(gres[0]); + EXPECT_EQ(std::memcmp(src.data(), dst.data(), kSize), 0) + << "multi-page object did not round trip byte-exact"; +} + +TEST_F(BatchPutWarnTest, MultiPageObjectWithPartialTailRoundTripsByteExact) { + // Deliberately not a page multiple: the final page is short, which is the + // one place LogicalPageBytes differs and the easiest boundary to get wrong + // when accumulating a run length. + constexpr size_t kSize = 3 * kPageSize + 1234; + + std::vector src(kSize); + FillPattern(src.data(), kSize, 0xA5); + std::vector keys = {"mp-tail"}; + std::vector srcs = {src.data()}; + std::vector sizes = {kSize}; + auto pres = target_->BatchPut(keys, srcs, sizes); + ASSERT_EQ(pres.size(), 1u); + ASSERT_TRUE(pres[0]); + + // One extra byte of guard on each side catches a run that copies too much. + std::vector dst(kSize + 2, 0x7E); + std::vector dsts = {dst.data() + 1}; + auto gres = target_->BatchGet(keys, dsts, sizes); + ASSERT_EQ(gres.size(), 1u); + ASSERT_TRUE(gres[0]); + EXPECT_EQ(std::memcmp(src.data(), dst.data() + 1, kSize), 0) + << "partial-tail object did not round trip byte-exact"; + EXPECT_EQ(dst.front(), static_cast(0x7E)) << "run underran the destination"; + EXPECT_EQ(dst.back(), static_cast(0x7E)) << "run overran the destination"; +} + +// --------------------------------------------------------------------------- +// GPU user buffers. Neither branch had coverage for a device pointer driven +// through PoolClient, which is the whole point of the port: an unregistered +// device pointer used to be described as host bytes and memcpy'd. +// --------------------------------------------------------------------------- + +namespace { + +// Device memory is not available on every machine that runs this suite (and is +// not available at all in a container without /dev/kfd), so the GPU cases skip +// rather than fail. A skip here is not a pass — it means the case did not run. +bool DeviceMemoryAvailable() { + int count = 0; + if (hipGetDeviceCount(&count) != hipSuccess) { + (void)hipGetLastError(); + return false; + } + return count > 0; +} + +class DeviceBuffer { + public: + explicit DeviceBuffer(size_t bytes) { + if (hipMalloc(&ptr_, bytes) != hipSuccess) { + (void)hipGetLastError(); + ptr_ = nullptr; + } + } + ~DeviceBuffer() { + if (ptr_ != nullptr) (void)hipFree(ptr_); + } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + void* get() const { return ptr_; } + bool Valid() const { return ptr_ != nullptr; } + + private: + void* ptr_ = nullptr; +}; + +} // namespace + +// The self-target path with a device source and destination. This is the case +// BuildLocalPageTransfers used to mislabel: the user ref was built with +// TransferRef::HostBytes directly, so LocalCopyEngine claimed the pair and +// memcpy'd from device memory. Multi-page so it also covers coalescing on the +// GPU side. +TEST_F(BatchPutWarnTest, DeviceBufferRoundTripsByteExactOnLocalPath) { + if (!DeviceMemoryAvailable()) GTEST_SKIP() << "no HIP device available"; + + constexpr size_t kSize = 3 * kPageSize + 512; // multi-page with a short tail + DeviceBuffer device_src(kSize); + DeviceBuffer device_dst(kSize); + if (!device_src.Valid() || !device_dst.Valid()) GTEST_SKIP() << "hipMalloc failed"; + + std::vector host(kSize); + FillPattern(host.data(), kSize, 0x3C); + ASSERT_EQ(hipMemcpy(device_src.get(), host.data(), kSize, hipMemcpyHostToDevice), hipSuccess); + ASSERT_EQ(hipMemset(device_dst.get(), 0, kSize), hipSuccess); + + std::vector keys = {"gpu-local"}; + std::vector srcs = {device_src.get()}; + std::vector sizes = {kSize}; + auto pres = target_->BatchPut(keys, srcs, sizes); + ASSERT_EQ(pres.size(), 1u); + ASSERT_TRUE(pres[0]) << "device-source put must be handled by HbmCopyEngine"; + + std::vector dsts = {device_dst.get()}; + auto gres = target_->BatchGet(keys, dsts, sizes); + ASSERT_EQ(gres.size(), 1u); + ASSERT_TRUE(gres[0]) << "device-destination get must be handled by HbmCopyEngine"; + + std::vector back(kSize, 0); + ASSERT_EQ(hipMemcpy(back.data(), device_dst.get(), kSize, hipMemcpyDeviceToHost), hipSuccess); + EXPECT_EQ(std::memcmp(host.data(), back.data(), kSize), 0) + << "device round trip through the local path was not byte-exact"; +} + +// An UNREGISTERED device source on a remote-bound batch must be rejected, not +// host-staged. Before the port the bounce predicate tested only HasHostPtr(), +// so this std::memcpy'd out of hipMalloc'd memory — which faults rather than +// corrupting, so the old behaviour would take the whole suite down here. +// +// The fixture pins caller_ to a single page precisely to force remote routing, +// which is what StagingFallbackSucceedsWithoutWarn depends on: the same shape +// with a HOST buffer succeeds for every key by staging. So the contrast is the +// assertion — identical batch, device memory instead of host memory, and every +// key must now fail. +// +// The log check is what makes the failure meaningful rather than incidental: it +// pins the reason to MoriIoEngine finding no engine for the pair. UmbpLogCapture +// forces the module to WARN, which is why the message is observable here and +// not in a default-level run. +TEST_F(BatchPutWarnTest, UnregisteredDeviceSourceIsRejectedNotHostStaged) { + if (!DeviceMemoryAvailable()) GTEST_SKIP() << "no HIP device available"; + + constexpr size_t kKeys = 4; + DeviceBuffer device_src(kKeys * kPerKey); + if (!device_src.Valid()) GTEST_SKIP() << "hipMalloc failed"; + ASSERT_EQ(hipMemset(device_src.get(), 0x6B, kKeys * kPerKey), hipSuccess); + + std::vector keys; + std::vector srcs; + std::vector sizes; + for (size_t i = 0; i < kKeys; ++i) { + keys.push_back("gpu-remote-" + std::to_string(i)); + srcs.push_back(static_cast(device_src.get()) + i * kPerKey); + sizes.push_back(kPerKey); + } + + UmbpLogCapture cap; + auto results = caller_->BatchPut(keys, srcs, sizes); + ASSERT_EQ(results.size(), keys.size()) << "result vector length must be preserved"; + for (size_t i = 0; i < results.size(); ++i) { + EXPECT_FALSE(results[i]) << "key=" << keys[i] + << ": unregistered device memory must not be host-staged"; + } + EXPECT_GE(cap.CountSubstring("transfer unplannable"), 1u) + << "the batch failed, but not because the transfer layer refused the GPU endpoint"; + + // None of the keys may have been published either, or a later Get would hand + // back whatever the staging buffer happened to hold. + auto exists = caller_->BatchExists(keys); + ASSERT_EQ(exists.size(), keys.size()); + for (size_t i = 0; i < exists.size(); ++i) { + EXPECT_FALSE(exists[i]) << "key=" << keys[i] << " was published despite a failed put"; + } +} + } // namespace } // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_pool_client_ranges.cpp b/tests/cpp/umbp/distributed/test_pool_client_ranges.cpp new file mode 100644 index 000000000..b1bf98697 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_pool_client_ranges.cpp @@ -0,0 +1,1396 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/common/device_gather.h" +#include "umbp/common/host_registration.h" +#include "umbp/distributed/config.h" +#include "umbp/distributed/master/master_client.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/pool_client.h" +#include "umbp/umbp_client.h" + +namespace mori::umbp { +namespace { + +constexpr size_t kPageSize = 4096; +constexpr size_t kObjectSize = 2 * kPageSize; +constexpr size_t kCallerCapacity = 16 * kPageSize; +constexpr size_t kTargetCapacity = 256 * kPageSize; +constexpr size_t kScratchSize = kObjectSize; // Forces one object per sub-batch. + +// The kernel dereferences the tier directly, so it needs the host region +// registered for GPU access; without it the client uses the copy engine. +bool GatherPathAvailable() { + static const bool available = [] { + std::vector probe(64 * 1024); + return HostTierRegistration(probe.data(), probe.size()).GatherableOn(0); + }(); + return available; +} + +uint16_t FreePort() { + int fd = ::socket(AF_INET, SOCK_STREAM, 0); + if (fd >= 0) { + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_ANY); + addr.sin_port = 0; + socklen_t len = sizeof(addr); + if (::bind(fd, reinterpret_cast(&addr), sizeof(addr)) == 0 && + ::getsockname(fd, reinterpret_cast(&addr), &len) == 0) { + const uint16_t port = ntohs(addr.sin_port); + ::close(fd); + return port; + } + ::close(fd); + } + static std::atomic next{56000}; + return next.fetch_add(1); +} + +bool WaitForExists(PoolClient* client, const std::string& key) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < deadline) { + if (client->Exists(key)) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + return client->Exists(key); +} + +bool WaitForExists(IUMBPClient* client, const std::string& key) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < deadline) { + const auto found = client->BatchExists({key}); + if (found.size() == 1 && found[0]) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + const auto found = client->BatchExists({key}); + return found.size() == 1 && found[0]; +} + +class PoolClientRangesTest : public ::testing::Test { + protected: + void SetUp() override { + MasterServerConfig master_cfg; + master_cfg.listen_address = "0.0.0.0:0"; + // The registry derives a 2 s client heartbeat from this TTL: short enough + // for normal publication tests, while the 12 s expiry window keeps a stopped + // caller alive through the stale-self assertion. + master_cfg.registry_config.heartbeat_ttl = std::chrono::seconds(4); + master_ = std::make_unique(std::move(master_cfg)); + master_thread_ = std::thread([this] { master_->Run(); }); + for (int i = 0; i < 100 && master_->GetBoundPort() == 0; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_NE(master_->GetBoundPort(), 0); + master_address_ = "localhost:" + std::to_string(master_->GetBoundPort()); + + caller_get_scratch_.resize(kScratchSize); + caller_put_scratch_.resize(kScratchSize); + target_get_scratch_.resize(kScratchSize); + target_put_scratch_.resize(kScratchSize); + caller_ = + MakeClient("ranges-caller", kCallerCapacity, caller_get_scratch_, caller_put_scratch_); + target_ = + MakeClient("ranges-target", kTargetCapacity, target_get_scratch_, target_put_scratch_); + } + + void TearDown() override { + if (caller_) caller_->Shutdown(); + if (target_) target_->Shutdown(); + caller_.reset(); + target_.reset(); + if (master_) master_->Shutdown(); + if (master_thread_.joinable()) master_thread_.join(); + master_.reset(); + } + + std::unique_ptr MakeClient(const std::string& node_id, size_t dram_capacity, + std::vector& get_scratch, + std::vector& put_scratch, + bool locality_prefetch = true) { + PoolClientConfig cfg; + cfg.master_config.node_id = node_id; + cfg.master_config.node_address = "127.0.0.1"; + cfg.master_config.master_address = master_address_; + cfg.io_engine.host = "0.0.0.0"; + cfg.io_engine.port = 0; + cfg.peer_service_port = FreePort(); + cfg.dram_page_size = kPageSize; + // Smaller than one object: ranged remote I/O can pass only through the + // registered arena's zero-copy descriptor, never the legacy staging path. + cfg.staging_buffer_size = 1024; + // The DRAM backend self-allocates its own pool (refactor Phase 2b), so the + // test hands it a size rather than a buffer. + cfg.dram.buffer_sizes = {dram_capacity}; + // Separate GET and PUT arenas so remote ranged get/put run concurrently. + cfg.ranged_get_scratch_buffer = get_scratch.data(); + cfg.ranged_get_scratch_size = get_scratch.size(); + cfg.ranged_put_scratch_buffer = put_scratch.data(); + cfg.ranged_put_scratch_size = put_scratch.size(); + cfg.cache_remote_fetches = false; + cfg.ranged_locality_prefetch = locality_prefetch; + auto client = std::make_unique(std::move(cfg)); + EXPECT_TRUE(client->Init()); + EXPECT_TRUE(client->RegisterMemory(get_scratch.data(), get_scratch.size())); + EXPECT_TRUE(client->RegisterMemory(put_scratch.data(), put_scratch.size())); + return client; + } + + // A full DistributedClient (not the bare PoolClient the other cases drive), + // because the opt-in this asserts lives in DistributedClient's constructor + // and in SupportsRangedIO(), neither of which PoolClient sees. + UMBPConfig DistributedConfig(const std::string& node_id, size_t ranged_scratch_size) { + UMBPConfig config; + config.dram.capacity_bytes = 4 * kTargetCapacity; + config.dram.use_hugepages = false; + config.ssd.enabled = false; + + UMBPDistributedConfig distributed; + distributed.master_config.node_id = node_id; + distributed.master_config.node_address = "127.0.0.1"; + distributed.master_config.master_address = master_address_; + distributed.io_engine.host = "0.0.0.0"; + distributed.io_engine.port = 0; + distributed.peer_service_port = FreePort(); + distributed.dram_page_size = kPageSize; + distributed.staging_buffer_size = 4 * kObjectSize; + distributed.ranged_scratch_size = ranged_scratch_size; + distributed.cache_remote_fetches = false; + config.distributed = std::move(distributed); + return config; + } + + // Publish a replica on one specific node, bypassing RoutePut's de-duplication + // so two nodes can hold the same key. Writes through the backend's published + // endpoint rather than a caller-owned buffer, since the medium owns its pool. + void PutLocalReplica(PoolClient* client, const std::string& key, + const std::vector& object) { + auto* backend = client->Backends().Get(TierType::DRAM); + ASSERT_NE(backend, nullptr); + auto allocated = backend->BatchAllocate({AllocateRequest{key, object.size()}}).front(); + ASSERT_EQ(allocated.outcome, AllocateOutcome::kSuccessAllocated); + + size_t copied = 0; + for (const auto& page : allocated.pages) { + TransferRef buf = backend->BufferRef(page.buffer_index); + ASSERT_TRUE(buf.HasHostPtr()); + const size_t bytes = std::min(kPageSize, object.size() - copied); + const size_t offset = static_cast(page.page_index) * kPageSize; + ASSERT_LE(offset + bytes, buf.size); + std::memcpy(static_cast(buf.host_ptr) + offset, object.data() + copied, bytes); + copied += bytes; + } + ASSERT_EQ(copied, object.size()); + + auto committed = backend->BatchCommit({CommitRequest{allocated.slot_id, key}}).front(); + ASSERT_TRUE(committed.success); + ASSERT_EQ(committed.bytes_committed, object.size()); + } + + // Publish `object` on the target only, so the caller's sole route to it is + // remote. PutLocalReplica writes the target's medium directly, which is what + // keeps a copy off the caller. + void SeedRemoteObject(const std::string& key, const std::vector& object) { + PutLocalReplica(target_.get(), key, object); + target_->Master().FlushHeartbeat(); + caller_->Master().FlushHeartbeat(); + ASSERT_TRUE(WaitForExists(caller_.get(), key)); + } + + // Is `key` readable from this node's own medium? The same question phase 1 + // of BatchGetRanges asks, and therefore the thing that decides whether a + // later read goes remote. + static bool LocallyResident(PoolClient* client, const std::string& key) { + auto* backend = client->Backends().Get(client->Medium()); + return backend != nullptr && + backend->BatchResolve({key}, /*include_descs=*/false).front().found; + } + + static bool WaitLocallyResident(PoolClient* client, const std::string& key, + std::chrono::milliseconds timeout = std::chrono::seconds(5)) { + const auto deadline = std::chrono::steady_clock::now() + timeout; + while (std::chrono::steady_clock::now() < deadline) { + if (LocallyResident(client, key)) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + return LocallyResident(client, key); + } + + // Two DRAM instances under a tiered policy whose entry tier is the SECOND + // one registered, so selecting by TierType and asking the policy disagree. + std::unique_ptr MakeTieredClient(const std::string& node_id) { + PoolClientConfig cfg; + cfg.master_config.node_id = node_id; + cfg.master_config.node_address = "127.0.0.1"; + cfg.master_config.master_address = master_address_; + cfg.io_engine.host = "0.0.0.0"; + cfg.io_engine.port = 0; + cfg.peer_service_port = FreePort(); + cfg.dram_page_size = kPageSize; + cfg.staging_buffer_size = 1024; + tiered_get_scratch_.resize(kScratchSize); + tiered_put_scratch_.resize(kScratchSize); + cfg.ranged_get_scratch_buffer = tiered_get_scratch_.data(); + cfg.ranged_get_scratch_size = tiered_get_scratch_.size(); + cfg.ranged_put_scratch_buffer = tiered_put_scratch_.data(); + cfg.ranged_put_scratch_size = tiered_put_scratch_.size(); + cfg.cache_remote_fetches = false; + + BackendInstanceConfig cold; + cold.name = "cold"; + cold.tier = TierType::DRAM; + cold.dram.buffer_sizes = {kTargetCapacity}; + BackendInstanceConfig hot = cold; + hot.name = "hot"; + cfg.backends = {cold, hot}; + + cfg.placement_policy = PoolPlacementPolicy::TIERED; + LogicalTierConfig hot_tier; + hot_tier.name = "hot"; + hot_tier.members = {{"hot", 1}}; + hot_tier.entry = true; + hot_tier.offload_to = {"cold"}; + LogicalTierConfig cold_tier; + cold_tier.name = "cold"; + cold_tier.members = {{"cold", 1}}; + cfg.logical_tiers = {hot_tier, cold_tier}; + + auto client = std::make_unique(std::move(cfg)); + EXPECT_TRUE(client->Init()); + EXPECT_TRUE(client->RegisterMemory(tiered_get_scratch_.data(), tiered_get_scratch_.size())); + EXPECT_TRUE(client->RegisterMemory(tiered_put_scratch_.data(), tiered_put_scratch_.size())); + return client; + } + + bool WaitForRoute(PoolClient* client, const std::string& key, + const std::unordered_set& excludes, + const std::string& expected_node) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < deadline) { + std::vector> routes; + const auto status = client->Master().BatchRouteGet({key}, excludes, &routes); + if (status.ok() && routes.size() == 1 && routes[0].has_value() && + routes[0]->node_id == expected_node) { + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + return false; + } + + std::unique_ptr master_; + std::thread master_thread_; + std::string master_address_; + std::unique_ptr caller_; + std::unique_ptr target_; + std::vector caller_get_scratch_; + std::vector caller_put_scratch_; + std::vector target_get_scratch_; + std::vector target_put_scratch_; + std::vector tiered_get_scratch_; + std::vector tiered_put_scratch_; +}; + +// --------------------------------------------------------------------------- +// Ranged I/O must go through the pool, not straight at the backend registry. +// A tiered policy shows both directions at once: the entry tier is +// deliberately not the first-registered DRAM instance, so a write that picked +// a backend by TierType lands on the wrong one, and a read that scanned the +// registry serves the bytes but leaves the pool's placement index and read +// accounting empty -- which is what makes watermark offload and +// promote-on-read inert. +// --------------------------------------------------------------------------- +TEST_F(PoolClientRangesTest, LocalRangedIoIsPlacedAndAccountedByThePool) { + caller_->Shutdown(); + target_->Shutdown(); + auto tiered = MakeTieredClient("ranges-tiered"); + ASSERT_NE(tiered, nullptr); + + auto* cold = tiered->Backends().Get("cold"); + auto* hot = tiered->Backends().Get("hot"); + ASSERT_NE(cold, nullptr); + ASSERT_NE(hot, nullptr); + // Selecting by TierType would hand back "cold": it is the first DRAM entry. + ASSERT_EQ(tiered->Backends().Get(TierType::DRAM), cold); + + const std::string key = "tiered-range-object"; + std::vector head(1024, 0x41); + std::vector tail(kObjectSize - head.size(), 0x42); + const std::vector srcs = {head.data(), tail.data()}; + const std::vector sizes = {head.size(), tail.size()}; + const std::vector offsets = {0, head.size()}; + ASSERT_EQ(tiered->BatchPutRanges({key}, {kObjectSize}, {srcs}, {sizes}, {offsets}), + std::vector({true})); + + // Placed where the policy's entry tier says, not where the routed TierType + // would have put it. + EXPECT_TRUE(hot->Contains(key)); + EXPECT_FALSE(cold->Contains(key)); + + std::vector readback(kObjectSize, 0); + const std::vector dsts = {readback.data()}; + ASSERT_EQ(tiered->BatchGetRanges({key}, {dsts}, {{kObjectSize}}, {{0}}), + std::vector({true})); + EXPECT_EQ(std::memcmp(readback.data(), head.data(), head.size()), 0); + EXPECT_EQ(std::memcmp(readback.data() + head.size(), tail.data(), tail.size()), 0); + + // Served through the pool, so the read is attributed to the logical tier + // that held the key. A registry scan would leave this map empty. + const auto hits = tiered->TierReadHits(); + const auto hot_hits = hits.find("hot"); + ASSERT_NE(hot_hits, hits.end()); + EXPECT_GE(hot_hits->second, 1u); + + tiered->Shutdown(); +} + +TEST_F(PoolClientRangesTest, RemoteRoundTripSubBatchesAndInstallsLocally) { + constexpr size_t kKeys = 2; + std::vector keys = {"range-object-0", "range-object-1"}; + std::vector> objects(kKeys, std::vector(kObjectSize)); + for (size_t k = 0; k < kKeys; ++k) { + for (size_t i = 0; i < kObjectSize; ++i) { + objects[k][i] = static_cast((i * 13 + k * 37) & 0xff); + } + } + + std::vector object_sizes(kKeys, kObjectSize); + std::vector> put_ptrs(kKeys); + std::vector> put_sizes(kKeys); + std::vector> put_offsets(kKeys); + for (size_t k = 0; k < kKeys; ++k) { + put_ptrs[k] = {objects[k].data(), objects[k].data() + 1500, objects[k].data() + 6000}; + put_sizes[k] = {1500, 4500, kObjectSize - 6000}; + put_offsets[k] = {0, 1500, 6000}; + } + + auto put = caller_->BatchPutRanges(keys, object_sizes, put_ptrs, put_sizes, put_offsets); + ASSERT_EQ(put.size(), kKeys); + EXPECT_TRUE(put[0]); + EXPECT_TRUE(put[1]); + caller_->Master().FlushHeartbeat(); + target_->Master().FlushHeartbeat(); + ASSERT_TRUE(WaitForExists(caller_.get(), keys[0])); + ASSERT_TRUE(WaitForExists(caller_.get(), keys[1])); + + std::vector> out_a(kKeys, std::vector(900, 0)); + std::vector> out_b(kKeys, std::vector(700, 0)); + std::vector> get_ptrs(kKeys); + std::vector> get_sizes(kKeys, {900, 700}); + std::vector> get_offsets(kKeys, {200, 5000}); + for (size_t k = 0; k < kKeys; ++k) get_ptrs[k] = {out_a[k].data(), out_b[k].data()}; + + auto first_get = caller_->BatchGetRanges(keys, get_ptrs, get_sizes, get_offsets); + ASSERT_EQ(first_get.size(), kKeys); + for (size_t k = 0; k < kKeys; ++k) { + ASSERT_TRUE(first_get[k]) << "key=" << keys[k]; + EXPECT_EQ(std::memcmp(out_a[k].data(), objects[k].data() + 200, 900), 0); + EXPECT_EQ(std::memcmp(out_b[k].data(), objects[k].data() + 5000, 700), 0); + } + + // Reading the same ranges again must give the same bytes whichever way it is + // served: still remotely, or locally once the background prefetch has landed. + // Both are correct, so this asserts the payload rather than the route -- + // which route was taken is what LocalityPrefetch* below pins down. + for (auto& v : out_a) std::fill(v.begin(), v.end(), 0); + for (auto& v : out_b) std::fill(v.begin(), v.end(), 0); + auto second_get = caller_->BatchGetRanges(keys, get_ptrs, get_sizes, get_offsets); + EXPECT_EQ(second_get, std::vector({true, true})); + for (size_t k = 0; k < kKeys; ++k) { + EXPECT_EQ(std::memcmp(out_a[k].data(), objects[k].data() + 200, 900), 0); + EXPECT_EQ(std::memcmp(out_b[k].data(), objects[k].data() + 5000, 700), 0); + } +} + +TEST_F(PoolClientRangesTest, LocalCopyOnlyRegistrationServesRangesAndSharesOnePlan) { + // kLocalCopyOnly records a region without handing it to the IO engine. It + // exists because the two halves of "register" are separable: a region that is + // only ever a local copy endpoint has no use for an RDMA MR, but it still has + // to be in the table -- a range that misses it is described by its own + // address instead of its region base, so instead of every range sharing one + // (src, dst) pair and collapsing into a single transfer plan, each range + // becomes its own plan. That is a latency bug with no wrong answer attached, + // which is exactly why it needs a test rather than a reviewer. + const std::string key = "local-copy-only-ranges"; + std::vector object(kObjectSize); + for (size_t i = 0; i < object.size(); ++i) object[i] = static_cast((i * 31 + 7) & 0xff); + PutLocalReplica(caller_.get(), key, object); + + // One buffer, several disjoint ranges landing in it -- the layer-wise restore + // shape, and the case the pair has to collapse for. + std::vector dst(kObjectSize, 0); + ASSERT_TRUE(caller_->RegisterMemory(dst.data(), dst.size(), mori::io::MemoryLocationType::CPU, -1, + MemoryRegistration::kLocalCopyOnly)); + + // Disjoint and inside the 2-page object; the middle one straddles the page + // boundary so the walk emits more than one fragment for a single range. + const std::vector sizes = {2048, 2048, 1024}; + const std::vector offsets = {0, 3072, 6144}; + std::vector dsts; + size_t cursor = 0; + for (size_t bytes : sizes) { + dsts.push_back(dst.data() + cursor); + cursor += bytes; + } + ASSERT_EQ(caller_->BatchGetRanges({key}, {dsts}, {sizes}, {offsets}), std::vector({true})); + + cursor = 0; + for (size_t i = 0; i < sizes.size(); ++i) { + EXPECT_EQ(std::memcmp(dst.data() + cursor, object.data() + offsets[i], sizes[i]), 0) + << "range " << i << " restored the wrong bytes"; + cursor += sizes[i]; + } + + // Idempotent for a size already covered, and still refuses to silently shrink + // a region -- same contract as the pinned path. + EXPECT_TRUE(caller_->RegisterMemory(dst.data(), dst.size() / 2, mori::io::MemoryLocationType::CPU, + -1, MemoryRegistration::kLocalCopyOnly)); + EXPECT_FALSE(caller_->RegisterMemory(dst.data(), dst.size() * 2, + mori::io::MemoryLocationType::CPU, -1, + MemoryRegistration::kLocalCopyOnly)); + caller_->DeregisterMemory(dst.data()); +} + +TEST_F(PoolClientRangesTest, InvalidRangesFailWithoutPublishing) { + std::vector src(kObjectSize, 0x5a); + std::vector keys = {"range-gap", "range-overlap"}; + std::vector object_sizes = {kObjectSize, kObjectSize}; + std::vector> ptrs = {{src.data(), src.data() + 4096}, + {src.data(), src.data() + 2048}}; + std::vector> sizes = {{4096, 2048}, {4096, 4096}}; + std::vector> offsets = {{0, 6144}, {0, 2048}}; + auto result = caller_->BatchPutRanges(keys, object_sizes, ptrs, sizes, offsets); + EXPECT_EQ(result, std::vector({false, false})); + EXPECT_FALSE(caller_->Exists(keys[0])); + EXPECT_FALSE(caller_->Exists(keys[1])); +} + +TEST_F(PoolClientRangesTest, DistributedScratchIsExplicitlyOptedIn) { + EXPECT_EQ(UMBPDistributedConfig{}.ranged_scratch_size, 0u); + + auto config = DistributedConfig("ranges-no-scratch", /*ranged_scratch_size=*/0); + std::string validation_error; + ASSERT_TRUE(config.Validate(&validation_error)) << validation_error; + auto client = CreateUMBPClient(config); + ASSERT_NE(client, nullptr); + EXPECT_EQ(client->GetDeploymentMode(), UMBPDeploymentMode::Distributed); + EXPECT_FALSE(client->SupportsRangedIO()); + + const std::string key = "ordinary-io-without-ranged-scratch"; + std::vector source(kObjectSize); + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast((i * 17 + 3) & 0xff); + } + auto put = client->BatchPut({key}, {reinterpret_cast(source.data())}, {source.size()}); + ASSERT_EQ(put, std::vector({true})); + ASSERT_TRUE(client->Flush()); + ASSERT_TRUE(WaitForExists(client.get(), key)); + + std::vector restored(source.size(), 0); + auto get = + client->BatchGet({key}, {reinterpret_cast(restored.data())}, {restored.size()}); + ASSERT_EQ(get, std::vector({true})); + EXPECT_EQ(restored, source); + + // Bypassing the advertised capability remains a safe per-key failure. Use a + // missing key so the call reaches the remote-scratch guard after its local probe. + std::vector ranged_out(16, 0); + auto ranged = client->BatchGetRanges({"missing-without-ranged-scratch"}, + {{reinterpret_cast(ranged_out.data())}}, + {{ranged_out.size()}}, {{0}}); + EXPECT_EQ(ranged, std::vector({false})); + client->Close(); + client.reset(); + + auto opted_in_config = DistributedConfig("ranges-with-scratch", kScratchSize); + ASSERT_TRUE(opted_in_config.Validate(&validation_error)) << validation_error; + auto opted_in = CreateUMBPClient(opted_in_config); + ASSERT_NE(opted_in, nullptr); + EXPECT_TRUE(opted_in->SupportsRangedIO()); + opted_in->Close(); +} + +TEST_F(PoolClientRangesTest, StaleSelfLocationIsExcludedBeforeRemoteFetch) { + const std::string key = "range-stale-self"; + std::vector stale_object(kObjectSize, 0x11); + std::vector remote_object(kObjectSize); + for (size_t i = 0; i < remote_object.size(); ++i) { + remote_object[i] = static_cast((i * 29 + 7) & 0xff); + } + + // Publish two replicas without going through RoutePut's de-duplication. The + // caller heartbeat is then stopped before local eviction, deliberately + // leaving its location stale at master while the target remains readable. + PutLocalReplica(caller_.get(), key, stale_object); + PutLocalReplica(target_.get(), key, remote_object); + caller_->Master().FlushHeartbeat(); + target_->Master().FlushHeartbeat(); + ASSERT_TRUE(WaitForRoute(caller_.get(), key, {"ranges-target"}, "ranges-caller")); + ASSERT_TRUE(WaitForRoute(caller_.get(), key, {"ranges-caller"}, "ranges-target")); + + caller_->Master().StopHeartbeat(); + auto evicted = caller_->Backends().Get(TierType::DRAM)->Evict({key}); + ASSERT_EQ(evicted.size(), 1u); + ASSERT_EQ(evicted[0].bytes_freed, kObjectSize); + ASSERT_TRUE(WaitForRoute(caller_.get(), key, {}, "ranges-caller")) + << "master must still prefer the caller's deliberately stale location"; + + std::vector out_a(777, 0); + std::vector out_b(999, 0); + auto result = + caller_->BatchGetRanges({key}, {{out_a.data(), out_b.data()}}, {{777, 999}}, {{123, 5000}}); + + ASSERT_EQ(result, std::vector({true})); + EXPECT_EQ(std::memcmp(out_a.data(), remote_object.data() + 123, out_a.size()), 0); + EXPECT_EQ(std::memcmp(out_b.data(), remote_object.data() + 5000, out_b.size()), 0); +} + +TEST_F(PoolClientRangesTest, GpuRangesUseGatherKernel) { + struct HipAllocation { + void* ptr = nullptr; + ~HipAllocation() { + if (ptr) (void)hipFree(ptr); + } + }; + + const std::vector part_sizes = {1024, 2048, kObjectSize - 3072}; + std::vector gpu_src(part_sizes.size()); + std::vector> host_src(part_sizes.size()); + std::vector src_ptrs; + for (size_t i = 0; i < part_sizes.size(); ++i) { + host_src[i].resize(part_sizes[i], static_cast(0x31 + i)); + ASSERT_EQ(hipMalloc(&gpu_src[i].ptr, part_sizes[i]), hipSuccess); + ASSERT_EQ(hipMemcpy(gpu_src[i].ptr, host_src[i].data(), part_sizes[i], hipMemcpyHostToDevice), + hipSuccess); + src_ptrs.push_back(gpu_src[i].ptr); + } + + const bool gather_available = DeviceGatherEnabled() && GatherPathAvailable(); + const uint64_t launches_before = DeviceGatherLaunchCount(); + const std::vector gpu_keys = {"gpu-range-object-0", "gpu-range-object-1"}; + auto put = target_->BatchPutRanges(gpu_keys, {kObjectSize, kObjectSize}, {src_ptrs, src_ptrs}, + {part_sizes, part_sizes}, {{0, 1024, 3072}, {0, 1024, 3072}}); + ASSERT_EQ(put, std::vector({true, true})); + if (gather_available) EXPECT_EQ(DeviceGatherLaunchCount(), launches_before + 1); + + const std::vector get_sizes = {600, 700, 800}; + const std::vector get_offsets = {100, 2200, 6500}; + std::vector gpu_dst(2 * get_sizes.size()); + std::vector> dst_ptrs(2); + for (size_t key = 0; key < 2; ++key) { + for (size_t i = 0; i < get_sizes.size(); ++i) { + auto& allocation = gpu_dst[key * get_sizes.size() + i]; + ASSERT_EQ(hipMalloc(&allocation.ptr, get_sizes[i]), hipSuccess); + ASSERT_EQ(hipMemset(allocation.ptr, 0, get_sizes[i]), hipSuccess); + dst_ptrs[key].push_back(allocation.ptr); + } + } + auto get = target_->BatchGetRanges(gpu_keys, dst_ptrs, {get_sizes, get_sizes}, + {get_offsets, get_offsets}); + ASSERT_EQ(get, std::vector({true, true})); + if (gather_available) EXPECT_EQ(DeviceGatherLaunchCount(), launches_before + 2); + + std::vector full(kObjectSize); + std::memset(full.data(), 0x31, 1024); + std::memset(full.data() + 1024, 0x32, 2048); + std::memset(full.data() + 3072, 0x33, kObjectSize - 3072); + for (size_t key = 0; key < 2; ++key) { + for (size_t i = 0; i < get_sizes.size(); ++i) { + std::vector actual(get_sizes[i]); + ASSERT_EQ(hipMemcpy(actual.data(), gpu_dst[key * get_sizes.size() + i].ptr, get_sizes[i], + hipMemcpyDeviceToHost), + hipSuccess); + EXPECT_EQ(std::memcmp(actual.data(), full.data() + get_offsets[i], get_sizes[i]), 0); + } + } +} + +// --------------------------------------------------------------------------- +// Separate GET / PUT scratch arenas: a remote ranged get and put use different +// buffers under different mutexes, so they overlap and must never clobber each +// other. The fixture already gives every client a separate get arena and put +// arena (each one object); these exercise the two together. +// --------------------------------------------------------------------------- + +// A remote ranged put (PUT arena) then a remote ranged get (GET arena) of the +// same keys round-trips the exact bytes. +TEST_F(PoolClientRangesTest, SplitRemotePutThenGetRoundTrips) { + constexpr size_t kKeys = 3; + std::vector keys = {"split-0", "split-1", "split-2"}; + std::vector> objects(kKeys, std::vector(kObjectSize)); + for (size_t k = 0; k < kKeys; ++k) { + for (size_t i = 0; i < kObjectSize; ++i) { + objects[k][i] = static_cast((i * 7 + k * 101) & 0xff); + } + } + + std::vector object_sizes(kKeys, kObjectSize); + std::vector> put_ptrs(kKeys); + std::vector> put_sizes(kKeys); + std::vector> put_offsets(kKeys); + for (size_t k = 0; k < kKeys; ++k) { + put_ptrs[k] = {objects[k].data(), objects[k].data() + 2048}; + put_sizes[k] = {2048, kObjectSize - 2048}; + put_offsets[k] = {0, 2048}; + } + auto put = caller_->BatchPutRanges(keys, object_sizes, put_ptrs, put_sizes, put_offsets); + ASSERT_EQ(put, std::vector(kKeys, true)); + caller_->Master().FlushHeartbeat(); + target_->Master().FlushHeartbeat(); + for (const auto& key : keys) ASSERT_TRUE(WaitForExists(caller_.get(), key)); + + std::vector> out(kKeys, std::vector(kObjectSize, 0)); + std::vector> get_ptrs(kKeys); + std::vector> get_sizes(kKeys, {kObjectSize}); + std::vector> get_offsets(kKeys, {0}); + for (size_t k = 0; k < kKeys; ++k) get_ptrs[k] = {out[k].data()}; + + auto got = caller_->BatchGetRanges(keys, get_ptrs, get_sizes, get_offsets); + ASSERT_EQ(got, std::vector(kKeys, true)); + for (size_t k = 0; k < kKeys; ++k) { + EXPECT_EQ(std::memcmp(out[k].data(), objects[k].data(), kObjectSize), 0) << "key=" << keys[k]; + } +} + +// Concurrent remote gets (GET arena) and remote puts (PUT arena) on disjoint key +// sets: every payload must survive byte-for-byte, proving the two arenas never +// overwrite each other under real contention. +TEST_F(PoolClientRangesTest, SplitConcurrentGetPutDoNotClobber) { + constexpr size_t kSeed = 8; // keys pre-seeded on target for the get threads + constexpr size_t kPutPerThread = 8; + constexpr size_t kGetThreads = 3; + constexpr size_t kPutThreads = 3; + + // Seed keys the get threads will fetch remotely; each has a distinct fill. + std::vector seed_keys(kSeed); + std::vector> seed_objs(kSeed, std::vector(kObjectSize)); + for (size_t s = 0; s < kSeed; ++s) { + for (size_t i = 0; i < kObjectSize; ++i) + seed_objs[s][i] = static_cast((i + s * 17) & 0xff); + seed_keys[s] = "seed-" + std::to_string(s); + PutLocalReplica(target_.get(), seed_keys[s], seed_objs[s]); + } + target_->Master().FlushHeartbeat(); + caller_->Master().FlushHeartbeat(); + for (const auto& k : seed_keys) ASSERT_TRUE(WaitForExists(caller_.get(), k)); + + std::atomic ok{true}; + std::vector threads; + + // GET threads: read whole seeded objects, verify bytes (GET arena). + for (size_t t = 0; t < kGetThreads; ++t) { + threads.emplace_back([&, t] { + for (size_t it = 0; it < 12; ++it) { + const size_t s = (t + it) % kSeed; + std::vector out(kObjectSize, 0); + std::vector> ptrs = {{out.data()}}; + std::vector> sizes = {{kObjectSize}}; + std::vector> offsets = {{0}}; + auto r = caller_->BatchGetRanges({seed_keys[s]}, ptrs, sizes, offsets); + if (r.size() != 1 || !r[0] || std::memcmp(out.data(), seed_objs[s].data(), kObjectSize)) { + ok.store(false); + } + } + }); + } + + // PUT threads: write unique objects (PUT arena) concurrently with the gets. + // Read-back verification happens after join to avoid racing master's + // ADD-event propagation. + auto put_key = [](size_t t, size_t p) { + return "cput-" + std::to_string(t) + "-" + std::to_string(p); + }; + auto put_fill = [](size_t t, size_t p) { return static_cast((t * 31 + p * 3 + 1) & 0xff); }; + for (size_t t = 0; t < kPutThreads; ++t) { + threads.emplace_back([&, t] { + for (size_t p = 0; p < kPutPerThread; ++p) { + std::vector obj(kObjectSize); + std::memset(obj.data(), put_fill(t, p), kObjectSize); + std::vector> src = {{obj.data(), obj.data() + 4096}}; + std::vector> sizes = {{4096, kObjectSize - 4096}}; + std::vector> offs = {{0, 4096}}; + auto pr = caller_->BatchPutRanges({put_key(t, p)}, {kObjectSize}, src, sizes, offs); + if (pr.size() != 1 || !pr[0]) ok.store(false); + } + }); + } + + for (auto& th : threads) th.join(); + EXPECT_TRUE(ok.load()); + + // After the concurrent phase, every put's bytes must be intact — a clobber + // between the two arenas would corrupt these. + caller_->Master().FlushHeartbeat(); + target_->Master().FlushHeartbeat(); + for (size_t t = 0; t < kPutThreads; ++t) { + for (size_t p = 0; p < kPutPerThread; ++p) { + const std::string key = put_key(t, p); + ASSERT_TRUE(WaitForExists(caller_.get(), key)); + std::vector back(kObjectSize, 0xAB); + std::vector> gp = {{back.data()}}; + std::vector> gs = {{kObjectSize}}; + std::vector> go = {{0}}; + auto gr = caller_->BatchGetRanges({key}, gp, gs, go); + ASSERT_EQ(gr.size(), 1u); + ASSERT_TRUE(gr[0]) << "key=" << key; + std::vector expect(kObjectSize, static_cast(put_fill(t, p))); + EXPECT_EQ(std::memcmp(back.data(), expect.data(), kObjectSize), 0) << "key=" << key; + } + } +} + +// --------------------------------------------------------------------------- +// Range-granular remote fetch +// --------------------------------------------------------------------------- + +// The arena is the whole evidence: whatever the wire brought back is sitting in +// it, and whatever it did not bring back still holds the poison. So filling it +// with a known byte and inspecting it afterwards answers both "did only the +// requested ranges cross" and "were they packed the way the copy-out expects" +// in one assertion, with no instrumentation inside PoolClient. +TEST_F(PoolClientRangesTest, RemoteRangedFetchMovesOnlyRequestedBytes) { + const std::string key = "ranged-only-requested"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 7 + 3) & 0xff); + SeedRemoteObject(key, object); + + constexpr size_t kSpanA = 900; + constexpr size_t kSpanB = 700; + constexpr size_t kOffsetA = 200; + constexpr size_t kOffsetB = 5000; + constexpr size_t kPacked = kSpanA + kSpanB; + constexpr char kPoison = static_cast(0xAA); + std::fill(caller_get_scratch_.begin(), caller_get_scratch_.end(), kPoison); + + std::vector out_a(kSpanA, 0), out_b(kSpanB, 0); + std::vector> ptrs = {{out_a.data(), out_b.data()}}; + std::vector> sizes = {{kSpanA, kSpanB}}; + std::vector> offsets = {{kOffsetA, kOffsetB}}; + auto got = caller_->BatchGetRanges({key}, ptrs, sizes, offsets); + ASSERT_EQ(got.size(), 1u); + ASSERT_TRUE(got[0]); + + EXPECT_EQ(std::memcmp(out_a.data(), object.data() + kOffsetA, kSpanA), 0); + EXPECT_EQ(std::memcmp(out_b.data(), object.data() + kOffsetB, kSpanB), 0); + + // Packed in caller order, back to back. + EXPECT_EQ(std::memcmp(caller_get_scratch_.data(), object.data() + kOffsetA, kSpanA), 0); + EXPECT_EQ(std::memcmp(caller_get_scratch_.data() + kSpanA, object.data() + kOffsetB, kSpanB), 0); + // ...and nothing else moved. A whole-object fetch would have overwritten + // every byte of the arena up to kObjectSize. + for (size_t i = kPacked; i < caller_get_scratch_.size(); ++i) { + ASSERT_EQ(caller_get_scratch_[i], kPoison) << "arena byte " << i << " was overwritten"; + } +} + +// A range that starts mid-page, crosses into the next, and a range that ends +// exactly at the end of a SHORT last page. The last-page clamp is the one +// piece of the remote builder's arithmetic that nothing else reaches, and +// getting it wrong reads bytes the object does not own. +TEST_F(PoolClientRangesTest, RemoteRangedFetchCrossesPageBoundary) { + constexpr size_t kShortObject = kPageSize + 1000; // last page holds 1000 B + const std::string key = "ranged-short-tail"; + std::vector object(kShortObject); + for (size_t i = 0; i < kShortObject; ++i) object[i] = static_cast((i * 11 + 5) & 0xff); + SeedRemoteObject(key, object); + + constexpr size_t kCrossOffset = kPageSize - 300; // 300 B in page 0, 500 B in page 1 + constexpr size_t kCrossSize = 800; + constexpr size_t kTailSize = 400; + constexpr size_t kTailOffset = kShortObject - kTailSize; // ends on the last byte + + std::vector cross(kCrossSize, 0), tail(kTailSize, 0); + std::vector> ptrs = {{cross.data(), tail.data()}}; + std::vector> sizes = {{kCrossSize, kTailSize}}; + std::vector> offsets = {{kCrossOffset, kTailOffset}}; + auto got = caller_->BatchGetRanges({key}, ptrs, sizes, offsets); + ASSERT_EQ(got.size(), 1u); + ASSERT_TRUE(got[0]); + EXPECT_EQ(std::memcmp(cross.data(), object.data() + kCrossOffset, kCrossSize), 0); + EXPECT_EQ(std::memcmp(tail.data(), object.data() + kTailOffset, kTailSize), 0); +} + +// What the sglang direct linker actually does: the same key set, once per layer +// group, each call taking a different slice of every object. +TEST_F(PoolClientRangesTest, MultiLayerGroupsOverSameKeys) { + constexpr size_t kKeys = 3; + constexpr size_t kGroups = 8; + constexpr size_t kSliceSize = kObjectSize / kGroups; + + std::vector keys; + std::vector> objects; + for (size_t k = 0; k < kKeys; ++k) { + keys.push_back("layerwise-" + std::to_string(k)); + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) { + object[i] = static_cast((i * 3 + k * 101) & 0xff); + } + SeedRemoteObject(keys.back(), object); + objects.push_back(std::move(object)); + } + + std::vector> restored(kKeys, std::vector(kObjectSize, 0)); + for (size_t g = 0; g < kGroups; ++g) { + const size_t slice_offset = g * kSliceSize; + std::vector> ptrs(kKeys); + std::vector> sizes(kKeys, {kSliceSize}); + std::vector> offsets(kKeys, {slice_offset}); + for (size_t k = 0; k < kKeys; ++k) ptrs[k] = {restored[k].data() + slice_offset}; + + auto got = caller_->BatchGetRanges(keys, ptrs, sizes, offsets); + ASSERT_EQ(got.size(), kKeys); + for (size_t k = 0; k < kKeys; ++k) ASSERT_TRUE(got[k]) << "group=" << g << " key=" << keys[k]; + } + + // Every group, wherever it was served from, contributed the right bytes. + for (size_t k = 0; k < kKeys; ++k) { + EXPECT_EQ(std::memcmp(restored[k].data(), objects[k].data(), kObjectSize), 0) + << "key=" << keys[k]; + } +} + +// Spans totalling more than the arena are split across sub-batches rather than +// failing the key -- the packing unit is a span prefix, not an object. +TEST_F(PoolClientRangesTest, PartialReadSplitsWhenSpansExceedArena) { + // The spans have to total MORE than the arena to split, which a single object + // the size of the arena can never do -- so the object is bigger, and the read + // is deliberately NOT a tiling one. That keeps holds_whole_object false, so + // this exercises the split on the arena path rather than the slot path + // ObjectLargerThanArenaIsServedAcrossSubBatches already covers. + constexpr size_t kBigObject = 5 * kPageSize; // 20480, arena is 8192 + const std::string key = "ranged-split"; + std::vector object(kBigObject); + for (size_t i = 0; i < kBigObject; ++i) object[i] = static_cast((i * 17 + 9) & 0xff); + SeedRemoteObject(key, object); + + // 3 x 4000 = 12000 > 8192, with gaps between them so they cannot tile. + constexpr size_t kSpan = 4000; + const size_t starts[3] = {0, 6000, 13000}; + std::vector> out(3, std::vector(kSpan, 0)); + std::vector> ptrs(1); + std::vector> sizes(1), offsets(1); + for (size_t i = 0; i < 3; ++i) { + ptrs[0].push_back(out[i].data()); + sizes[0].push_back(kSpan); + offsets[0].push_back(starts[i]); + } + + auto got = caller_->BatchGetRanges({key}, ptrs, sizes, offsets); + ASSERT_EQ(got.size(), 1u); + ASSERT_TRUE(got[0]); + for (size_t i = 0; i < 3; ++i) { + EXPECT_EQ(std::memcmp(out[i].data(), object.data() + starts[i], kSpan), 0) << "span " << i; + } +} + +// One span bigger than the whole arena is the only remaining unservable shape, +// and it must fail that key alone. +// A single range bigger than the whole arena is the one shape that is still +// unservable. It must fail that key ALONE, and -- the part worth testing -- +// it must not leave the key half-delivered: the object here is larger than the +// arena, so the earlier ranges have already been split into their own fetch +// unit by the time the oversized one is seen. Committing those units would +// fetch part of the key and still report it whole, which the caller has no way +// to detect. +TEST_F(PoolClientRangesTest, OversizedRangeFailsItsKeyWholeAndAlone) { + constexpr size_t kBigObject = 5 * kPageSize; // > kScratchSize + static_assert(kBigObject > kScratchSize, "the split path needs an oversized object"); + const std::string bad_key = "oversize-bad"; + const std::string ok_key = "oversize-good"; + std::vector big(kBigObject), small(kObjectSize); + for (size_t i = 0; i < kBigObject; ++i) big[i] = static_cast((i * 59 + 4) & 0xff); + for (size_t i = 0; i < kObjectSize; ++i) small[i] = static_cast((i * 61 + 9) & 0xff); + SeedRemoteObject(bad_key, big); + SeedRemoteObject(ok_key, small); + + // 5000 + 5000 forces a split (the arena is 8192), then 10480 cannot be served + // at all. Every one of the three has to be refused. + constexpr char kPoison = static_cast(0x5a); + std::vector a(5000, kPoison), b(5000, kPoison), c(kBigObject - 10000, kPoison); + std::vector good(600, 0); + std::vector> ptrs = {{a.data(), b.data(), c.data()}, {good.data()}}; + std::vector> sizes = {{a.size(), b.size(), c.size()}, {600}}; + std::vector> offsets = {{0, 5000, 10000}, {900}}; + + auto got = caller_->BatchGetRanges({bad_key, ok_key}, ptrs, sizes, offsets); + ASSERT_EQ(got.size(), 2u); + EXPECT_FALSE(got[0]); + EXPECT_TRUE(got[1]); + + // Nothing partially delivered: the split-off prefix must not have been + // fetched either. + EXPECT_EQ(a, std::vector(a.size(), kPoison)) << "first range was delivered anyway"; + EXPECT_EQ(b, std::vector(b.size(), kPoison)); + EXPECT_EQ(c, std::vector(c.size(), kPoison)); + EXPECT_EQ(std::memcmp(good.data(), small.data() + 900, 600), 0); +} + +// The same split path when every range IS servable: an object larger than the +// arena, read whole, has to come back correct across several sub-batches. +TEST_F(PoolClientRangesTest, ObjectLargerThanArenaIsServedAcrossSubBatches) { + constexpr size_t kBigObject = 5 * kPageSize; + const std::string key = "oversize-object"; + std::vector object(kBigObject); + for (size_t i = 0; i < kBigObject; ++i) object[i] = static_cast((i * 67 + 5) & 0xff); + SeedRemoteObject(key, object); + + constexpr size_t kChunk = 4096; + std::vector out(kBigObject, 0); + std::vector> ptrs(1); + std::vector> sizes(1), offsets(1); + for (size_t off = 0; off < kBigObject; off += kChunk) { + const size_t len = std::min(kChunk, kBigObject - off); + ptrs[0].push_back(out.data() + off); + sizes[0].push_back(len); + offsets[0].push_back(off); + } + ASSERT_TRUE(caller_->BatchGetRanges({key}, ptrs, sizes, offsets).front()); + EXPECT_EQ(std::memcmp(out.data(), object.data(), kBigObject), 0); +} + +// --------------------------------------------------------------------------- +// Locality prefetch +// --------------------------------------------------------------------------- + +TEST_F(PoolClientRangesTest, LocalityPrefetchInstallsWholeObjectAfterRangedRead) { + const std::string key = "prefetch-on"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 23 + 7) & 0xff); + SeedRemoteObject(key, object); + ASSERT_FALSE(LocallyResident(caller_.get(), key)); + + std::vector out(512, 0); + std::vector> ptrs = {{out.data()}}; + std::vector> sizes = {{512}}; + std::vector> offsets = {{1024}}; + ASSERT_TRUE(caller_->BatchGetRanges({key}, ptrs, sizes, offsets).front()); + + // Asynchronous by design: the read that scheduled it must not have waited. + ASSERT_TRUE(WaitLocallyResident(caller_.get(), key)); + + // And what landed is the WHOLE object, not just the slice that was read. + std::vector whole(kObjectSize, 0); + std::vector> whole_ptrs = {{whole.data()}}; + std::vector> whole_sizes = {{kObjectSize}}; + std::vector> whole_offsets = {{0}}; + ASSERT_TRUE(caller_->BatchGetRanges({key}, whole_ptrs, whole_sizes, whole_offsets).front()); + EXPECT_EQ(std::memcmp(whole.data(), object.data(), kObjectSize), 0); +} + +TEST_F(PoolClientRangesTest, LocalityPrefetchOffLeavesNothingLocal) { + std::vector no_prefetch_get(kScratchSize), no_prefetch_put(kScratchSize); + auto quiet = MakeClient("ranges-no-prefetch", kCallerCapacity, no_prefetch_get, no_prefetch_put, + /*locality_prefetch=*/false); + + const std::string key = "prefetch-off"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 29 + 11) & 0xff); + PutLocalReplica(target_.get(), key, object); + target_->Master().FlushHeartbeat(); + quiet->Master().FlushHeartbeat(); + ASSERT_TRUE(WaitForExists(quiet.get(), key)); + + std::vector out(512, 0); + std::vector> ptrs = {{out.data()}}; + std::vector> sizes = {{512}}; + std::vector> offsets = {{1024}}; + ASSERT_TRUE(quiet->BatchGetRanges({key}, ptrs, sizes, offsets).front()); + EXPECT_EQ(std::memcmp(out.data(), object.data() + 1024, 512), 0); + + // Give a prefetch every chance to appear, then assert it did not. + EXPECT_FALSE(WaitLocallyResident(quiet.get(), key, std::chrono::milliseconds(750))); + quiet->Shutdown(); +} + +// Eight layer groups over one cold key must schedule ONE whole-object pull, not +// one per group: the dedup set is what keeps the duplicate traffic bounded at +// 1x the object instead of 8x. +// The switch buys off the SECOND fetch, not locality itself. A whole-object +// read already has every byte in hand, so landing it in a slot costs nothing on +// the wire and must still happen -- gating it would silently undo what the +// pre-ranged code did on every remote ranged read. +TEST_F(PoolClientRangesTest, WholeObjectReadStillInstallsWithPrefetchOff) { + std::vector quiet_get(kScratchSize), quiet_put(kScratchSize); + auto quiet = MakeClient("ranges-quiet-whole", kCallerCapacity, quiet_get, quiet_put, + /*locality_prefetch=*/false); + + const std::string key = "prefetch-off-whole"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 83 + 7) & 0xff); + PutLocalReplica(target_.get(), key, object); + target_->Master().FlushHeartbeat(); + quiet->Master().FlushHeartbeat(); + ASSERT_TRUE(WaitForExists(quiet.get(), key)); + ASSERT_FALSE(LocallyResident(quiet.get(), key)); + + // Four ascending gapless quarters: the ranges tile the object, so the fetch + // brings all of it and the slot path applies. + constexpr size_t kQuarter = kObjectSize / 4; + std::vector out(kObjectSize, 0); + std::vector> ptrs(1); + std::vector> sizes(1), offsets(1); + for (size_t q = 0; q < 4; ++q) { + ptrs[0].push_back(out.data() + q * kQuarter); + sizes[0].push_back(kQuarter); + offsets[0].push_back(q * kQuarter); + } + ASSERT_TRUE(quiet->BatchGetRanges({key}, ptrs, sizes, offsets).front()); + EXPECT_EQ(std::memcmp(out.data(), object.data(), kObjectSize), 0); + + // Synchronous: the slot is committed before the call returns, so no waiting. + EXPECT_TRUE(LocallyResident(quiet.get(), key)) + << "a switch named for the background pull must not disable the free install"; + quiet->Shutdown(); +} + +// The sglang shape against a cold key: eight layer groups, each naming the same +// key, with the object only ever arriving from the peer. Every group must +// return its own bytes and the key must end up local. +// +// This does NOT pin the in-flight dedup (prefetch_inflight_). Dedup has no +// external observable -- a duplicate pull is absorbed by the worker's +// BatchResolve check and by kSuccessAlreadyExists, so removing the set changes +// nothing a caller can see. Naming it here would claim coverage that does not +// exist; it is a throughput guard, and the bench is where it shows up. +TEST_F(PoolClientRangesTest, LayerWiseGroupsOverOneColdKeyAllReturnTheirBytes) { + const std::string key = "prefetch-dedup"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 31 + 13) & 0xff); + SeedRemoteObject(key, object); + + constexpr size_t kGroups = 8; + constexpr size_t kSliceSize = kObjectSize / kGroups; + std::vector restored(kObjectSize, 0); + for (size_t g = 0; g < kGroups; ++g) { + const size_t slice_offset = g * kSliceSize; + std::vector> ptrs = {{restored.data() + slice_offset}}; + std::vector> sizes = {{kSliceSize}}; + std::vector> offsets = {{slice_offset}}; + ASSERT_TRUE(caller_->BatchGetRanges({key}, ptrs, sizes, offsets).front()) << "group=" << g; + } + EXPECT_EQ(std::memcmp(restored.data(), object.data(), kObjectSize), 0); + + ASSERT_TRUE(WaitLocallyResident(caller_.get(), key)); + // A second pull would have had to allocate a second slot for the same key; + // the medium reports one, and the bytes are intact. + std::vector whole(kObjectSize, 0); + std::vector> whole_ptrs = {{whole.data()}}; + std::vector> whole_sizes = {{kObjectSize}}; + std::vector> whole_offsets = {{0}}; + ASSERT_TRUE(caller_->BatchGetRanges({key}, whole_ptrs, whole_sizes, whole_offsets).front()); + EXPECT_EQ(std::memcmp(whole.data(), object.data(), kObjectSize), 0); +} + +// When the caller's ranges tile the object in order, the arena slice IS the +// object, so locality costs one local copy and nothing on the wire. This is +// the whole-object reader's shape (one call naming every layer). +TEST_F(PoolClientRangesTest, CompleteArenaObjectIsInstalledWithoutRefetching) { + const std::string key = "arena-complete"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 47 + 3) & 0xff); + SeedRemoteObject(key, object); + ASSERT_FALSE(LocallyResident(caller_.get(), key)); + + // Four ascending, gapless ranges covering [0, kObjectSize). + constexpr size_t kQuarter = kObjectSize / 4; + std::vector out(kObjectSize, 0); + std::vector> ptrs(1); + std::vector> sizes(1), offsets(1); + for (size_t q = 0; q < 4; ++q) { + ptrs[0].push_back(out.data() + q * kQuarter); + sizes[0].push_back(kQuarter); + offsets[0].push_back(q * kQuarter); + } + ASSERT_TRUE(caller_->BatchGetRanges({key}, ptrs, sizes, offsets).front()); + EXPECT_EQ(std::memcmp(out.data(), object.data(), kObjectSize), 0); + + // Installed from the arena, synchronously -- no waiting for a worker, and no + // second trip to the peer. + EXPECT_TRUE(LocallyResident(caller_.get(), key)); + + std::vector again(kObjectSize, 0); + std::vector> again_ptrs = {{again.data()}}; + std::vector> again_sizes = {{kObjectSize}}; + std::vector> again_offsets = {{0}}; + ASSERT_TRUE(caller_->BatchGetRanges({key}, again_ptrs, again_sizes, again_offsets).front()); + EXPECT_EQ(std::memcmp(again.data(), object.data(), kObjectSize), 0); +} + +// A whole-object read does not need the arena at all: the slot it is going to +// be cached in has the same layout, so the peer writes straight into it. The +// arena is the witness -- poison it, and if a single byte changed the read went +// the long way round. +TEST_F(PoolClientRangesTest, WholeObjectReadBypassesTheArena) { + const std::string key = "slot-direct"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 71 + 13) & 0xff); + SeedRemoteObject(key, object); + ASSERT_FALSE(LocallyResident(caller_.get(), key)); + + constexpr char kPoison = static_cast(0xC3); + std::fill(caller_get_scratch_.begin(), caller_get_scratch_.end(), kPoison); + + constexpr size_t kQuarter = kObjectSize / 4; + std::vector out(kObjectSize, 0); + std::vector> ptrs(1); + std::vector> sizes(1), offsets(1); + for (size_t q = 0; q < 4; ++q) { + ptrs[0].push_back(out.data() + q * kQuarter); + sizes[0].push_back(kQuarter); + offsets[0].push_back(q * kQuarter); + } + ASSERT_TRUE(caller_->BatchGetRanges({key}, ptrs, sizes, offsets).front()); + + EXPECT_EQ(std::memcmp(out.data(), object.data(), kObjectSize), 0); + EXPECT_TRUE(LocallyResident(caller_.get(), key)) << "committing the slot IS the install"; + EXPECT_EQ(caller_get_scratch_, std::vector(kScratchSize, kPoison)) + << "the arena was used for a read that did not need it"; +} + +// ...and when the medium cannot give it a slot, the same read still has to +// work. The fallback is the arena, decided on the failed allocation rather +// than after a failed transfer. +TEST_F(PoolClientRangesTest, WholeObjectReadFallsBackToArenaWithoutASlot) { + // One page of DRAM, objects are two: BatchAllocate can never satisfy one. + std::vector tiny_get(kScratchSize), tiny_put(kScratchSize); + auto cramped = MakeClient("ranges-cramped", kPageSize, tiny_get, tiny_put); + + const std::string key = "slot-unavailable"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 73 + 19) & 0xff); + PutLocalReplica(target_.get(), key, object); + target_->Master().FlushHeartbeat(); + cramped->Master().FlushHeartbeat(); + ASSERT_TRUE(WaitForExists(cramped.get(), key)); + + std::vector out(kObjectSize, 0); + std::vector> ptrs = {{out.data()}}; + std::vector> sizes = {{kObjectSize}}; + std::vector> offsets = {{0}}; + ASSERT_TRUE(cramped->BatchGetRanges({key}, ptrs, sizes, offsets).front()) + << "no slot must not mean no read"; + EXPECT_EQ(std::memcmp(out.data(), object.data(), kObjectSize), 0); + EXPECT_FALSE(LocallyResident(cramped.get(), key)) << "nowhere to cache it, and that is fine"; + cramped->Shutdown(); +} + +// Ranges that cover every byte but arrive out of order do NOT make the arena +// equal the object -- it is packed in caller order -- so this must fall back to +// the background pull rather than install a scrambled slice. +TEST_F(PoolClientRangesTest, OutOfOrderFullCoverageDoesNotInstallFromArena) { + const std::string key = "arena-scrambled"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 53 + 17) & 0xff); + SeedRemoteObject(key, object); + + constexpr size_t kHalf = kObjectSize / 2; + std::vector out(kObjectSize, 0); + // Second half named first. + std::vector> ptrs = {{out.data() + kHalf, out.data()}}; + std::vector> sizes = {{kHalf, kHalf}}; + std::vector> offsets = {{kHalf, 0}}; + ASSERT_TRUE(caller_->BatchGetRanges({key}, ptrs, sizes, offsets).front()); + EXPECT_EQ(std::memcmp(out.data(), object.data(), kObjectSize), 0); + + // It still becomes local, but via the pull -- so what lands must be the real + // object and not the arena's caller-order layout. + ASSERT_TRUE(WaitLocallyResident(caller_.get(), key)); + std::vector again(kObjectSize, 0); + std::vector> again_ptrs = {{again.data()}}; + std::vector> again_sizes = {{kObjectSize}}; + std::vector> again_offsets = {{0}}; + ASSERT_TRUE(caller_->BatchGetRanges({key}, again_ptrs, again_sizes, again_offsets).front()); + EXPECT_EQ(std::memcmp(again.data(), object.data(), kObjectSize), 0); +} + +TEST_F(PoolClientRangesTest, MixedLocalAndRemoteRangedKeys) { + const std::string local_key = "mixed-local"; + const std::string remote_key = "mixed-remote"; + std::vector local_object(kObjectSize), remote_object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) { + local_object[i] = static_cast((i * 37 + 2) & 0xff); + remote_object[i] = static_cast((i * 41 + 6) & 0xff); + } + PutLocalReplica(caller_.get(), local_key, local_object); + SeedRemoteObject(remote_key, remote_object); + + std::vector out_local(600, 0), out_remote(600, 0); + std::vector> ptrs = {{out_local.data()}, {out_remote.data()}}; + std::vector> sizes = {{600}, {600}}; + std::vector> offsets = {{300}, {4500}}; + + auto got = caller_->BatchGetRanges({local_key, remote_key}, ptrs, sizes, offsets); + ASSERT_EQ(got.size(), 2u); + EXPECT_TRUE(got[0]); + EXPECT_TRUE(got[1]); + EXPECT_EQ(std::memcmp(out_local.data(), local_object.data() + 300, 600), 0); + EXPECT_EQ(std::memcmp(out_remote.data(), remote_object.data() + 4500, 600), 0); +} + +// A key whose ranges do not fit its object must fail alone, and must not +// disturb a healthy sibling in the same call. +TEST_F(PoolClientRangesTest, PartialFailureAttributionAcrossKeys) { + const std::string bad_key = "attrib-bad"; + const std::string good_key = "attrib-good"; + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) object[i] = static_cast((i * 43 + 8) & 0xff); + SeedRemoteObject(bad_key, object); + SeedRemoteObject(good_key, object); + + std::vector bad_out(600, 0x5a), good_out(600, 0); + std::vector> ptrs = {{bad_out.data()}, {good_out.data()}}; + std::vector> sizes = {{600}, {600}}; + // Runs 300 B past the end of the object. + std::vector> offsets = {{kObjectSize - 300}, {1200}}; + + auto got = caller_->BatchGetRanges({bad_key, good_key}, ptrs, sizes, offsets); + ASSERT_EQ(got.size(), 2u); + EXPECT_FALSE(got[0]); + EXPECT_TRUE(got[1]); + EXPECT_EQ(std::memcmp(good_out.data(), object.data() + 1200, 600), 0); + EXPECT_EQ(bad_out, std::vector(600, 0x5a)) << "failed key's buffer was written"; +} + +// Everything above reads one or two keys from one thread. The paths this +// change added only diverge under load: which keys are local when a call +// starts, whether a slot can be had, whether the background pull lands before +// the next group asks -- all of it decided by a race. So drive every shape at +// once over an overlapping key set and check the bytes afterwards. +// +// The caller's medium deliberately cannot hold all the keys, so slot +// allocation fails part of the time and the arena fallback is exercised too. +TEST_F(PoolClientRangesTest, ConcurrentMixedShapesAllReturnCorrectBytes) { + constexpr size_t kKeys = 12; // > kCallerCapacity / kObjectSize, so slots run out + constexpr size_t kThreads = 4; + constexpr size_t kRounds = 3; + constexpr size_t kGroups = 8; + constexpr size_t kSlice = kObjectSize / kGroups; + + std::vector keys; + std::vector> objects; + for (size_t k = 0; k < kKeys; ++k) { + keys.push_back("stress-" + std::to_string(k)); + std::vector object(kObjectSize); + for (size_t i = 0; i < kObjectSize; ++i) { + object[i] = static_cast((i * 97 + k * 211 + 5) & 0xff); + } + SeedRemoteObject(keys.back(), object); + objects.push_back(std::move(object)); + } + + std::atomic mismatches{0}; + std::atomic failures{0}; + std::vector threads; + for (size_t t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t] { + // Half the threads read layer-wise (a slice per call, never the whole + // object); half read the whole object in one call (the slot path). Both + // touch the same keys. + const bool layer_wise = (t % 2) == 0; + std::vector> restored(kKeys, std::vector(kObjectSize, 0)); + std::vector served(kKeys, 1); + for (size_t round = 0; round < kRounds; ++round) { + for (auto& buffer : restored) std::fill(buffer.begin(), buffer.end(), 0); + std::fill(served.begin(), served.end(), 1); + + if (layer_wise) { + for (size_t g = 0; g < kGroups; ++g) { + const size_t offset = g * kSlice; + std::vector> ptrs(kKeys); + std::vector> sizes(kKeys, {kSlice}); + std::vector> offsets(kKeys, {offset}); + for (size_t k = 0; k < kKeys; ++k) ptrs[k] = {restored[k].data() + offset}; + const auto got = caller_->BatchGetRanges(keys, ptrs, sizes, offsets); + for (size_t k = 0; k < kKeys; ++k) { + if (!got[k]) { + served[k] = false; + failures.fetch_add(1, std::memory_order_relaxed); + } + } + } + } else { + std::vector> ptrs(kKeys); + std::vector> sizes(kKeys), offsets(kKeys); + for (size_t k = 0; k < kKeys; ++k) { + for (size_t g = 0; g < kGroups; ++g) { + ptrs[k].push_back(restored[k].data() + g * kSlice); + sizes[k].push_back(kSlice); + offsets[k].push_back(g * kSlice); + } + } + const auto got = caller_->BatchGetRanges(keys, ptrs, sizes, offsets); + for (size_t k = 0; k < kKeys; ++k) { + if (!got[k]) { + served[k] = false; + failures.fetch_add(1, std::memory_order_relaxed); + } + } + } + + // Only a key the client SAID it served is checked. A refused read + // leaves its buffer untouched, and counting that as corruption would + // turn every capacity refusal into a false alarm. + for (size_t k = 0; k < kKeys; ++k) { + if (!served[k]) continue; + if (std::memcmp(restored[k].data(), objects[k].data(), kObjectSize) != 0) { + mismatches.fetch_add(1, std::memory_order_relaxed); + } + } + } + }); + } + for (auto& thread : threads) thread.join(); + + // A read may legitimately fail here -- the medium is intentionally too small, + // so a slot or an arena round can be refused -- but it must never return the + // wrong bytes. + EXPECT_EQ(mismatches.load(), 0u) << "a caller buffer came back wrong"; + // Failures are reported, not asserted away: the medium really is too small + // here, so a refusal is a legitimate outcome and pinning it to zero would + // make this test fail for the one reason it is allowed to. + if (failures.load() != 0) { + GTEST_LOG_(INFO) << failures.load() + << " reads were refused (capacity), none returned bad bytes"; + } +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_range_map.cpp b/tests/cpp/umbp/distributed/test_range_map.cpp new file mode 100644 index 000000000..38a43037b --- /dev/null +++ b/tests/cpp/umbp/distributed/test_range_map.cpp @@ -0,0 +1,204 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// ForEachRangePageFragment is the only place the object-range -> page-range +// mapping lives, for both the local medium and a peer's pages. An error here +// does not fail: it reads or writes the wrong bytes of a KV block and the +// caller sees a successful transfer. Until this file existed the arithmetic +// was only reachable through a live two-node fixture, so the awkward shapes +// (short last page, range ending exactly on the object's last byte, one byte +// past the end) went untested. + +#include + +#include +#include +#include + +#include "umbp/distributed/range_map.h" + +namespace mori::umbp { +namespace { + +struct Fragment { + size_t page_index; + uint64_t tier_offset; + size_t bytes; + size_t copied; + + bool operator==(const Fragment& other) const { + return page_index == other.page_index && tier_offset == other.tier_offset && + bytes == other.bytes && copied == other.copied; + } +}; + +// Pages laid out as a contiguous run starting at `first_page` of buffer 0, +// which is what the allocator produces for a single-slot object. +std::vector PageRun(size_t count, uint32_t first_page = 0, uint32_t buffer = 0) { + std::vector pages; + pages.reserve(count); + for (size_t i = 0; i < count; ++i) { + PageLocation page; + page.buffer_index = buffer; + page.page_index = first_page + static_cast(i); + pages.push_back(page); + } + return pages; +} + +std::vector Collect(const std::vector& pages, uint64_t page_size, + uint64_t stored_size, size_t offset, size_t size, bool* ok) { + std::vector out; + *ok = ForEachRangePageFragment( + pages, page_size, stored_size, offset, size, + [&](size_t page_index, uint64_t tier_offset, size_t bytes, size_t copied) { + out.push_back(Fragment{page_index, tier_offset, bytes, copied}); + return true; + }); + return out; +} + +constexpr uint64_t kPage = 4096; + +TEST(RangeMapTest, RangeInsideOnePage) { + bool ok = false; + const auto frags = Collect(PageRun(4), kPage, 4 * kPage, 100, 500, &ok); + EXPECT_TRUE(ok); + ASSERT_EQ(frags.size(), 1u); + EXPECT_EQ(frags[0], (Fragment{0, 100, 500, 0})); +} + +TEST(RangeMapTest, RangeStartsMidPageAndSpansTwo) { + bool ok = false; + const auto frags = Collect(PageRun(4), kPage, 4 * kPage, kPage - 300, 800, &ok); + EXPECT_TRUE(ok); + ASSERT_EQ(frags.size(), 2u); + EXPECT_EQ(frags[0], (Fragment{0, kPage - 300, 300, 0})); + EXPECT_EQ(frags[1], (Fragment{1, kPage, 500, 300})); +} + +TEST(RangeMapTest, RangeSpansEveryPage) { + bool ok = false; + const auto frags = Collect(PageRun(3), kPage, 3 * kPage, 0, 3 * kPage, &ok); + EXPECT_TRUE(ok); + ASSERT_EQ(frags.size(), 3u); + for (size_t p = 0; p < 3; ++p) { + EXPECT_EQ(frags[p], (Fragment{p, p * kPage, kPage, p * kPage})) << "page " << p; + } +} + +// tier_offset is an offset into the BUFFER, so a slot that does not start at +// page 0 must be reported relative to the buffer, not to the object. +TEST(RangeMapTest, TierOffsetIsBufferRelative) { + bool ok = false; + const auto frags = Collect(PageRun(2, /*first_page=*/7), kPage, 2 * kPage, 64, 128, &ok); + EXPECT_TRUE(ok); + ASSERT_EQ(frags.size(), 1u); + EXPECT_EQ(frags[0], (Fragment{0, 7 * kPage + 64, 128, 0})); +} + +// The last page of an object is short whenever the object is not a page +// multiple, and the range must be clamped to the part the object owns. +TEST(RangeMapTest, ShortLastPageIsClamped) { + constexpr uint64_t kStored = kPage + 1000; + bool ok = false; + const auto frags = Collect(PageRun(2), kPage, kStored, kPage - 200, 1200, &ok); + EXPECT_TRUE(ok); + ASSERT_EQ(frags.size(), 2u); + EXPECT_EQ(frags[0], (Fragment{0, kPage - 200, 200, 0})); + EXPECT_EQ(frags[1], (Fragment{1, kPage, 1000, 200})); +} + +TEST(RangeMapTest, RangeEndingOnTheObjectsLastByte) { + constexpr uint64_t kStored = kPage + 1000; + bool ok = false; + const auto frags = Collect(PageRun(2), kPage, kStored, kStored - 400, 400, &ok); + EXPECT_TRUE(ok); + ASSERT_EQ(frags.size(), 1u); + EXPECT_EQ(frags[0], (Fragment{1, kPage + 600, 400, 0})); +} + +TEST(RangeMapTest, WholeObjectWithShortLastPage) { + constexpr uint64_t kStored = 2 * kPage + 1; + bool ok = false; + const auto frags = Collect(PageRun(3), kPage, kStored, 0, kStored, &ok); + EXPECT_TRUE(ok); + ASSERT_EQ(frags.size(), 3u); + EXPECT_EQ(frags[2], (Fragment{2, 2 * kPage, 1, 2 * kPage})); +} + +TEST(RangeMapTest, OneBytePastTheEndIsRejected) { + constexpr uint64_t kStored = kPage + 1000; + bool ok = true; + const auto frags = Collect(PageRun(2), kPage, kStored, kStored - 400, 401, &ok); + EXPECT_FALSE(ok); + EXPECT_TRUE(frags.empty()) << "a rejected range must not emit fragments"; +} + +TEST(RangeMapTest, OffsetPastTheEndIsRejected) { + bool ok = true; + Collect(PageRun(2), kPage, 2 * kPage, 2 * kPage, 1, &ok); + EXPECT_FALSE(ok); +} + +TEST(RangeMapTest, ZeroSizeIsRejected) { + bool ok = true; + Collect(PageRun(2), kPage, 2 * kPage, 0, 0, &ok); + EXPECT_FALSE(ok); +} + +// A page count that cannot hold the object, or that is one page more than the +// object needs, means the caller paired a size with the wrong slot. +TEST(RangeMapTest, MismatchedPageCountIsRejected) { + bool ok = true; + Collect(PageRun(1), kPage, 2 * kPage, 0, 100, &ok); + EXPECT_FALSE(ok) << "too few pages for the object"; + + ok = true; + Collect(PageRun(3), kPage, 2 * kPage, 0, 100, &ok); + EXPECT_FALSE(ok) << "one more page than the object needs"; +} + +TEST(RangeMapTest, EmptyPagesOrZeroPageSizeIsRejected) { + bool ok = true; + Collect({}, kPage, kPage, 0, 100, &ok); + EXPECT_FALSE(ok); + + ok = true; + Collect(PageRun(1), 0, kPage, 0, 100, &ok); + EXPECT_FALSE(ok); +} + +// The walk stops at the first refusal instead of running the range to its end. +TEST(RangeMapTest, EmitRefusalAbortsTheWalk) { + size_t calls = 0; + const bool ok = ForEachRangePageFragment(PageRun(4), kPage, 4 * kPage, 0, 4 * kPage, + [&](size_t, uint64_t, size_t, size_t) { + ++calls; + return calls < 2; + }); + EXPECT_FALSE(ok); + EXPECT_EQ(calls, 2u); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_redis_master_metadata_store.cpp b/tests/cpp/umbp/distributed/test_redis_master_metadata_store.cpp index 245e3d59d..cc8f91fbf 100644 --- a/tests/cpp/umbp/distributed/test_redis_master_metadata_store.cpp +++ b/tests/cpp/umbp/distributed/test_redis_master_metadata_store.cpp @@ -268,7 +268,7 @@ class RedisStoreTest : public ::testing::TestWithParam { r.node_id = id; r.node_address = id + ".addr"; r.peer_address = id + ".peer:1234"; - r.tier_capacities[TierType::DRAM] = TierCapacity{1u << 30, 1u << 30}; + r.tier_capacities[TierType::DRAM] = TierCapacity{1u << 30, 1u << 30, 1u << 29}; r.tags = {"sgl_role=prefill", "zone=a"}; return r; } @@ -325,9 +325,26 @@ TEST_P(RedisStoreTest, RegisterMakesClientAlive) { EXPECT_EQ(rec->status, ClientStatus::ALIVE); ASSERT_EQ(rec->tier_capacities.count(TierType::DRAM), 1u); EXPECT_EQ(rec->tier_capacities[TierType::DRAM].total_bytes, 1u << 30); + EXPECT_EQ(rec->tier_capacities[TierType::DRAM].max_allocatable_bytes, 1u << 29); EXPECT_EQ(store_->GetClientTags("n1").size(), 2u); } +TEST_P(RedisStoreTest, PreservesLogicalTierPeakUtilization) { + auto registration = MakeReg("n1"); + LogicalTierCapacity hot; + hot.representative_tier = TierType::DRAM; + hot.capacity = TierCapacity{1024, 256, 128}; + hot.put_eligible = true; + hot.peak_member_utilization = 0.875; + registration.logical_tier_capacities["hot"] = hot; + + ASSERT_TRUE(store_->RegisterClient(registration, now_, 30s)); + auto record = store_->GetClient("n1"); + ASSERT_TRUE(record.has_value()); + ASSERT_EQ(record->logical_tier_capacities.count("hot"), 1u); + EXPECT_DOUBLE_EQ(record->logical_tier_capacities.at("hot").peak_member_utilization, 0.875); +} + TEST_P(RedisStoreTest, RegisterRejectsAliveDuplicateButAllowsStale) { EXPECT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); // Alive and not stale -> rejected. diff --git a/tests/cpp/umbp/distributed/test_redispp_smoke.cpp b/tests/cpp/umbp/distributed/test_redispp_smoke.cpp deleted file mode 100644 index bd6e5e48a..000000000 --- a/tests/cpp/umbp/distributed/test_redispp_smoke.cpp +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -// redis-plus-plus build/link smoke test. -// -// This does NOT test the UMBP store — it only proves that the redis++ submodule -// compiles with this toolchain and links (static redis++ + system hiredis) -// before any store code depends on it. It is the de-risk gate for adopting -// redis++ as the Redis-backend client library (cluster / sentinel / pooling): -// if this builds and links, the later seam migration rests on a proven base. -// -// Runtime behaviour: connects to UMBP_REDIS_URI (default tcp://127.0.0.1:6379) -// and, if a cluster seed is given via UMBP_REDIS_CLUSTER_SEEDS, to that cluster; -// both skip cleanly when nothing is reachable, so BUILD_TESTS on a host without -// a store never fails. - -#include -#include - -#include -#include - -namespace { - -std::string EnvOr(const char* key, const std::string& def) { - const char* v = std::getenv(key); - return (v != nullptr && *v != '\0') ? std::string(v) : def; -} - -// Compile check: constructing ConnectionOptions exercises the redis++ headers. -// Actual linkage of the static archive is proven by the ping tests below, which -// reference symbols defined in redis++'s .cpp files (Redis::ping / -// RedisCluster) regardless of whether they run or skip at runtime. -TEST(RedisPlusPlusSmoke, HeadersCompile) { - sw::redis::ConnectionOptions opts; - opts.host = "127.0.0.1"; - opts.port = 6379; - EXPECT_EQ(opts.port, 6379); -} - -TEST(RedisPlusPlusSmoke, PingSingleIfReachable) { - const std::string uri = EnvOr("UMBP_REDIS_URI", "tcp://127.0.0.1:6379"); - try { - sw::redis::Redis redis(uri); - const std::string pong = redis.ping(); - EXPECT_EQ(pong, "PONG"); - } catch (const sw::redis::Error& e) { - GTEST_SKIP() << "no RESP store reachable at " << uri << " (" << e.what() << "); skipping"; - } -} - -TEST(RedisPlusPlusSmoke, PingClusterIfSeedGiven) { - const char* seeds = std::getenv("UMBP_REDIS_CLUSTER_SEEDS"); - if (seeds == nullptr || *seeds == '\0') { - GTEST_SKIP() << "UMBP_REDIS_CLUSTER_SEEDS not set; skipping cluster smoke"; - } - try { - // RedisCluster fetches CLUSTER SLOTS on construction, proving cluster-path - // symbols link and a real cluster is reachable. - sw::redis::RedisCluster cluster{std::string(seeds)}; - const std::string pong = cluster.redis("smoke", false).ping(); - EXPECT_EQ(pong, "PONG"); - } catch (const sw::redis::Error& e) { - GTEST_SKIP() << "no Redis Cluster reachable at " << seeds << " (" << e.what() << "); skipping"; - } -} - -} // namespace diff --git a/tests/cpp/umbp/distributed/test_route_get_strategy.cpp b/tests/cpp/umbp/distributed/test_route_get_strategy.cpp index bdea94b93..70d9f1b59 100644 --- a/tests/cpp/umbp/distributed/test_route_get_strategy.cpp +++ b/tests/cpp/umbp/distributed/test_route_get_strategy.cpp @@ -22,6 +22,7 @@ #include #include +#include #include #include @@ -107,33 +108,10 @@ class LocalityAwareGetStrategy : public RouteGetStrategy { } }; -TEST(CustomRouteGetStrategyTest, PrefersLocalReplica) { - LocalityAwareGetStrategy strategy; - std::vector locations = { - MakeLoc("node-a", "loc-1"), - MakeLoc("node-b", "loc-2"), - MakeLoc("node-c", "loc-3"), - }; - - auto selected = strategy.Select(locations, "node-b"); - EXPECT_EQ(selected.node_id, "node-b"); -} - -TEST(CustomRouteGetStrategyTest, FallsBackToFirstWhenNoLocalReplica) { - LocalityAwareGetStrategy strategy; - std::vector locations = { - MakeLoc("node-a", "loc-1"), - MakeLoc("node-b", "loc-2"), - }; - - auto selected = strategy.Select(locations, "node-x"); - EXPECT_EQ(selected.node_id, "node-a"); -} - // ---- BatchSelect tests ---- TEST(RouteGetBatchSelectTest, ReturnsOneResultPerKeyInOrder) { - TierPriorityRouteGetStrategy strategy; + LocalPreferringRouteGetStrategy strategy; std::vector> per_key = { {MakeLoc("node-a", "loc-1")}, {MakeLoc("node-b", "loc-2")}, @@ -162,17 +140,20 @@ TEST(RouteGetBatchSelectTest, EmptyCandidateListLeftAsDefault) { EXPECT_EQ(out[2].node_id, "node-c"); } -TEST(RouteGetBatchSelectTest, TierPriorityPicksFastestTierPerKey) { - TierPriorityRouteGetStrategy strategy; - Location hbm = MakeLoc("node-a", "loc-1"); // HBM by default in MakeLoc - Location ssd = MakeLoc("node-b", "loc-2"); - ssd.tier = TierType::SSD; +// Was TierPriorityPicksFastestTierPerKey. Phase 4 deleted the tier order, so +// what BatchSelect applies per key is the locality rule: the requester's own +// replica wins even though it sits on the "slower" medium. +TEST(RouteGetBatchSelectTest, AppliesLocalPreferencePerKey) { + LocalPreferringRouteGetStrategy strategy; + Location remote_hbm = MakeLoc("node-a", "loc-1"); // HBM by default in MakeLoc + Location local_ssd = MakeLoc("requester", "loc-2"); + local_ssd.tier = TierType::SSD; - std::vector> per_key = {{ssd, hbm}}; + std::vector> per_key = {{remote_hbm, local_ssd}}; auto out = strategy.BatchSelect(per_key, "requester"); ASSERT_EQ(out.size(), 1u); - EXPECT_EQ(out[0].tier, TierType::HBM); - EXPECT_EQ(out[0].node_id, "node-a"); + EXPECT_EQ(out[0].node_id, "requester"); + EXPECT_EQ(out[0].tier, TierType::SSD); } TEST(RouteGetBatchSelectTest, CustomStrategyUsesBaseDefaultLoop) { @@ -192,5 +173,65 @@ TEST(RouteGetBatchSelectTest, CustomStrategyUsesBaseDefaultLoop) { EXPECT_EQ(out[2].node_id, "node-c"); // local replica preferred } +// ---- LocalPreferringRouteGetStrategy tests ---- + +Location MakeLoc(const std::string& node_id, TierType tier) { + Location loc; + loc.node_id = node_id; + loc.size = 4096; + loc.tier = tier; + return loc; +} + +TEST(LocalPreferringRouteGetStrategyTest, LocalReplicaBeatsRemotePeersOnTheSameTier) { + LocalPreferringRouteGetStrategy strategy; + std::vector locations = { + MakeLoc("dram-a", TierType::DRAM), + MakeLoc("requester", TierType::DRAM), + MakeLoc("dram-c", TierType::DRAM), + }; + + for (int i = 0; i < 100; ++i) { + EXPECT_EQ(strategy.Select(locations, "requester").node_id, "requester"); + } +} + +// With every medium equivalent, a non-local request can reach every replica +// rather than being confined to a preferred tier. +TEST(LocalPreferringRouteGetStrategyTest, NonLocalRequesterSpreadsAcrossEveryTier) { + LocalPreferringRouteGetStrategy strategy; + std::vector locations = { + MakeLoc("dram-a", TierType::DRAM), + MakeLoc("dram-b", TierType::DRAM), + MakeLoc("hbm-x", TierType::HBM), + MakeLoc("ssd-y", TierType::SSD), + }; + + std::set seen; + for (int i = 0; i < 2000; ++i) { + seen.insert(strategy.Select(locations, "requester").node_id); + } + EXPECT_EQ(seen.size(), 4u) << "every replica should be reachable, regardless of tier"; +} + +TEST(LocalPreferringRouteGetStrategyTest, EmptyRequesterFallsBackToRandom) { + LocalPreferringRouteGetStrategy strategy; + std::vector locations = { + MakeLoc("node-a", TierType::DRAM), + MakeLoc("node-b", TierType::DRAM), + }; + + std::set seen; + for (int i = 0; i < 500; ++i) seen.insert(strategy.Select(locations, "").node_id); + EXPECT_EQ(seen.size(), 2u); +} + +TEST(LocalPreferringRouteGetStrategyTest, EmptyReturnsDefault) { + LocalPreferringRouteGetStrategy strategy; + auto selected = strategy.Select({}, "requester"); + EXPECT_EQ(selected.tier, TierType::UNKNOWN); + EXPECT_TRUE(selected.node_id.empty()); +} + } // namespace } // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_route_put_strategy.cpp b/tests/cpp/umbp/distributed/test_route_put_strategy.cpp index 61c81456c..9af2861f8 100644 --- a/tests/cpp/umbp/distributed/test_route_put_strategy.cpp +++ b/tests/cpp/umbp/distributed/test_route_put_strategy.cpp @@ -67,7 +67,10 @@ std::optional SelectOne(ConfigurableRoutePutStrategy& strat, // ---- most_available / none: single-key placement (migrated from the old // TierAwareMostAvailableStrategy::Select coverage) ---- -TEST(MostAvailableNoneTest, PrefersHBMOverDRAM) { +// INVERTS the old PrefersHBMOverDRAM. Phase 4 deleted the HBM -> DRAM put +// order, so on a node advertising both, most-available picks whichever tier +// actually has more room — here DRAM's 400GB over HBM's 10GB. +TEST(MostAvailableNoneTest, PicksTheRoomiestTierOnANodeRegardlessOfMedium) { ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); std::vector clients = { @@ -78,10 +81,25 @@ TEST(MostAvailableNoneTest, PrefersHBMOverDRAM) { auto result = SelectOne(strategy, clients, 4096, /*exclude=*/{}); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result->node_id, "node-a"); + EXPECT_EQ(result->tier, TierType::DRAM); +} + +TEST(MostAvailableNoneTest, RoutesOnlyToNamedEntryPool) { + ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); + auto client = MakeClient("node-a", "addr-a", {{TierType::DRAM, TierCapacity{100 * GB, 90 * GB}}}); + client.logical_tier_capacities["hot"] = + LogicalTierCapacity{TierType::HBM, TierCapacity{20 * GB, 10 * GB, 10 * GB}, true}; + client.logical_tier_capacities["warm"] = + LogicalTierCapacity{TierType::DRAM, TierCapacity{80 * GB, 40 * GB, 40 * GB}, false}; + + auto result = SelectOne(strategy, {client}, 4096, {}); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result->logical_tier, "hot"); EXPECT_EQ(result->tier, TierType::HBM); } -TEST(MostAvailableNoneTest, FallsThroughToDRAMWhenHBMFull) { +// Unchanged by Phase 4: HBM has no room, so DRAM is the only eligible tier. +TEST(MostAvailableNoneTest, PicksTheOnlyTierWithRoom) { ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); std::vector clients = { @@ -94,10 +112,12 @@ TEST(MostAvailableNoneTest, FallsThroughToDRAMWhenHBMFull) { EXPECT_EQ(result->tier, TierType::DRAM); } -TEST(MostAvailableNoneTest, SsdIsNotADirectPutTarget) { - // SSD capacity is reported via heartbeat but RoutePut never steers a put at - // SSD: the SSD copy is filled asynchronously by copy-on-commit. With HBM and - // DRAM full, placement must return nullopt even when SSD has ample space. +// INVERTS the old SsdIsNotADirectPutTarget. Phase 4 treats every medium as +// equivalent, so a tier a node advertises capacity for is a tier it accepts +// puts on. When SSD returns as an SsdBackend that genuinely takes no direct +// puts, it advertises that itself (design doc §3 put_eligible) — the router +// does not assume it. +TEST(MostAvailableNoneTest, EveryAdvertisedTierIsAValidPutTarget) { ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); std::vector clients = { @@ -108,7 +128,8 @@ TEST(MostAvailableNoneTest, SsdIsNotADirectPutTarget) { }; auto result = SelectOne(strategy, clients, 4096, /*exclude=*/{}); - EXPECT_FALSE(result.has_value()); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result->tier, TierType::SSD); } TEST(MostAvailableNoneTest, ReturnsNulloptWhenAllFull) { @@ -136,6 +157,25 @@ TEST(MostAvailableNoneTest, ReturnsNulloptWhenBlockTooLarge) { EXPECT_FALSE(result.has_value()); } +TEST(MostAvailableNoneTest, RejectsValueLargerThanAnyBackendInAggregatedTier) { + ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); + std::vector clients = { + MakeClient("node-a", "addr-a", {{TierType::DRAM, TierCapacity{20 * GB, 20 * GB, 10 * GB}}}), + }; + + EXPECT_FALSE(SelectOne(strategy, clients, 11 * GB, /*exclude=*/{}).has_value()); + EXPECT_TRUE(SelectOne(strategy, clients, 10 * GB, /*exclude=*/{}).has_value()); +} + +TEST(MostAvailableNoneTest, ZeroMaxAllocatablePreservesLegacyCapacitySemantics) { + ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); + std::vector clients = { + MakeClient("node-a", "addr-a", {{TierType::DRAM, TierCapacity{20 * GB, 20 * GB}}}), + }; + + EXPECT_TRUE(SelectOne(strategy, clients, 15 * GB, /*exclude=*/{}).has_value()); +} + TEST(MostAvailableNoneTest, PicksMostAvailableOnSameTier) { ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); @@ -152,7 +192,9 @@ TEST(MostAvailableNoneTest, PicksMostAvailableOnSameTier) { EXPECT_EQ(result->tier, TierType::HBM); } -TEST(MostAvailableNoneTest, HBMPreferredEvenIfDRAMHasMoreSpace) { +// INVERTS the old HBMPreferredEvenIfDRAMHasMoreSpace: free space is now the +// only signal, so the tier with more of it wins. +TEST(MostAvailableNoneTest, MoreFreeSpaceWinsAcrossTiersOnTheSameNode) { ConfigurableRoutePutStrategy strategy(Algo::kMostAvailable, Affinity::kNone); std::vector clients = { @@ -162,7 +204,7 @@ TEST(MostAvailableNoneTest, HBMPreferredEvenIfDRAMHasMoreSpace) { auto result = SelectOne(strategy, clients, 4096, /*exclude=*/{}); ASSERT_TRUE(result.has_value()); - EXPECT_EQ(result->tier, TierType::HBM); + EXPECT_EQ(result->tier, TierType::DRAM); } TEST(MostAvailableNoneTest, EmptyClientListReturnsNullopt) { @@ -339,19 +381,24 @@ TEST(ConfigurableRoutePut, OrderAndDedupPreserved) { EXPECT_EQ(out[2]->outcome, RoutePutOutcome::kAlreadyExists); } -TEST(ConfigurableRoutePut, NeverRoutesToSsd) { +// INVERTS the old NeverRoutesToSsd. Phase 4 removed the "SSD is not a direct +// put target" rule from the router; a medium that cannot take direct puts will +// advertise that itself when one exists (design doc §3 put_eligible). +TEST(ConfigurableRoutePut, RoutesToSsdLikeAnyOtherAdvertisedTier) { ConfigurableRoutePutStrategy strat(Algo::kMostAvailable, Affinity::kNone); std::vector only_ssd{ MakeClient("a", "a-addr", {{TierType::SSD, {100 * GB, 100 * GB}}})}; auto out = strat.SelectBatch("req", {1 * GB}, NoDedup(1), only_ssd, {}); ASSERT_EQ(out.size(), 1u); - EXPECT_FALSE(out[0].has_value()); + ASSERT_TRUE(out[0].has_value()); + EXPECT_EQ(out[0]->tier, TierType::SSD); + // Alongside DRAM it competes purely on free space, and wins here. std::vector ssd_and_dram{MakeClient( "a", "a-addr", {{TierType::SSD, {100 * GB, 100 * GB}}, {TierType::DRAM, {2 * GB, 2 * GB}}})}; auto out2 = strat.SelectBatch("req", {1 * GB}, NoDedup(1), ssd_and_dram, {}); ASSERT_TRUE(out2[0].has_value()); - EXPECT_EQ(out2[0]->tier, TierType::DRAM); + EXPECT_EQ(out2[0]->tier, TierType::SSD); } // Strong, seed-independent: a node that cannot fit the block is never selected. @@ -444,9 +491,11 @@ TEST(ConfigurableRoutePut, SameWholeBatchPinsNodeAndTier) { } } -// same per-key fallback must not break tier priority: when the sticky anchor's -// HBM is full, a remote node's HBM beats the anchor's own DRAM. -TEST(ConfigurableRoutePut, SameFallbackPrefersRemoteHbmOverAnchorDram) { +// INVERTS the old SameFallbackPrefersRemoteHbmOverAnchorDram. With no tier +// order there is nothing for node stickiness to break, so when the anchor's +// first tier fills, the spill stays on the anchor node and changes tier rather +// than jumping to a remote node's "faster" medium. +TEST(ConfigurableRoutePut, SameFallbackKeepsTheAnchorNodeAndChangesTier) { ConfigurableRoutePutStrategy strat(Algo::kMostAvailable, Affinity::kSame); std::vector clients{ MakeClient("a", "a-addr", {{TierType::HBM, {2 * GB, 2 * GB}}}), @@ -458,13 +507,13 @@ TEST(ConfigurableRoutePut, SameFallbackPrefersRemoteHbmOverAnchorDram) { ASSERT_EQ(out.size(), 2u); ASSERT_TRUE(out[0].has_value()); ASSERT_TRUE(out[1].has_value()); - // key 0: most-available HBM -> b (3G > a 2G); anchor becomes b, b HBM -> 1G. + // key 0: most-available -> b (3G > a's 2G); anchor becomes b, b HBM -> 1G. EXPECT_EQ(out[0]->node_id, "b"); EXPECT_EQ(out[0]->tier, TierType::HBM); - // key 1: anchor b's HBM (1G) cannot fit; instead of falling to b's DRAM, tier - // priority routes to a's HBM. - EXPECT_EQ(out[1]->node_id, "a"); - EXPECT_EQ(out[1]->tier, TierType::HBM); + // key 1: anchor b's HBM (1G) cannot fit, so the spill takes b's DRAM (3G) — + // the sticky node is kept, which is the point of same-node affinity. + EXPECT_EQ(out[1]->node_id, "b"); + EXPECT_EQ(out[1]->tier, TierType::DRAM); } // dedup keys consume no projected capacity and do not count toward the same-node diff --git a/tests/cpp/umbp/distributed/test_ssd_backend.cpp b/tests/cpp/umbp/distributed/test_ssd_backend.cpp new file mode 100644 index 000000000..dd155d304 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_ssd_backend.cpp @@ -0,0 +1,647 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// SsdBackend: SSD reaching the data plane as an ordinary MediumBackend by +// staging through registered host DRAM. +// +// What is worth asserting here is the part that differs from DRAM — the spill +// on Commit, the fill on Resolve, and the staging arena that both borrow from — +// plus the invariant that makes the whole design work: what the backend +// PUBLISHES is plain host memory, so no transfer-layer change was needed. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/peer/backend/ssd_backend.h" +#include "umbp/distributed/transfer/composite_transfer_engine.h" +#include "umbp/distributed/transfer/local_copy_engine.h" +#ifdef UMBP_ENABLE_GDS +#include + +#include "umbp/distributed/transfer/gds_engine.h" +#include "umbp/distributed/transfer/hbm_copy_engine.h" +#endif + +namespace mori::umbp { +namespace { + +namespace fs = std::filesystem; + +constexpr uint64_t kPageSize = 64 * 1024; + +// Hands back process-local refs only — what CompositeTransferEngine degrades to +// on a node with no RDMA, and enough to exercise ownership without an IOEngine. +class LocalOnlyRegistrar final : public MemoryRegistrar { + public: + TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) override { + ++registrations; + last_loc = loc; + return TransferRef::HostBytes(base, size, loc, device); + } + void Deregister(const TransferRef&) override { ++deregistrations; } + + int registrations = 0; + int deregistrations = 0; + mori::io::MemoryLocationType last_loc = mori::io::MemoryLocationType::GPU; +}; + +class SsdBackendTest : public ::testing::Test { + protected: + void SetUp() override { + static std::atomic counter{0}; + dir_ = fs::temp_directory_path() / ("umbp_ssd_backend_test_" + std::to_string(::getpid()) + + "_" + std::to_string(counter.fetch_add(1))); + fs::remove_all(dir_); + } + + void TearDown() override { + if (backend_ != nullptr) backend_->Shutdown(); + backend_.reset(); + std::error_code ec; + fs::remove_all(dir_, ec); + } + + // `staging_pages` is the backend's real concurrency limit, so tests that care + // about exhaustion set it small. + SsdBackend* Start(uint32_t staging_pages = 8, + std::chrono::milliseconds read_lease_ttl = std::chrono::milliseconds{500}, + uint64_t ssd_capacity = 64ULL * 1024 * 1024) { + SsdBackend::Config cfg; + cfg.page_size = kPageSize; + cfg.staging_pages = staging_pages; + cfg.read_lease_ttl = read_lease_ttl; + cfg.ssd.enabled = true; + cfg.ssd.ssd.enabled = true; + cfg.ssd.ssd.storage_dir = dir_.string(); + cfg.ssd.ssd.capacity_bytes = ssd_capacity; + cfg.ssd.ssd.io.backend = UMBPIoBackend::Posix; // avoid io_uring container flakiness + + auto owned = std::make_unique(std::move(cfg)); + backend_ = std::move(owned); + EXPECT_TRUE(backend_->Init(®istrar_)); + return backend_.get(); + } + + // The full Put path a real writer takes: reserve a page, move bytes into it + // through the transfer layer, commit. + bool Put(SsdBackend* backend, const std::string& key, const std::vector& payload) { + auto allocated = backend->BatchAllocate({AllocateRequest{key, payload.size()}}); + if (allocated.size() != 1 || allocated[0].outcome != AllocateOutcome::kSuccessAllocated) { + return false; + } + const TransferRef pool = backend->BufferRef(allocated[0].pages[0].buffer_index); + const uint64_t off = static_cast(allocated[0].pages[0].page_index) * kPageSize; + + TransferItem item; + item.src = TransferRef::HostBytes(const_cast(payload.data()), payload.size()); + item.dst = pool; + item.dst_offset = off; + item.size = payload.size(); + std::vector failed; + if (!engine_.Transfer({item}, &failed)) return false; + + auto committed = backend->BatchCommit({CommitRequest{allocated[0].slot_id, key}}); + return committed.size() == 1 && committed[0].success; + } + + // The full Get path: resolve (which stages off the SSD), then move bytes out. + bool Get(SsdBackend* backend, const std::string& key, std::vector* out) { + auto resolved = backend->BatchResolve({key}, /*include_descs=*/false); + if (resolved.size() != 1 || !resolved[0].found) return false; + out->assign(resolved[0].size, 0); + + const TransferRef pool = backend->BufferRef(resolved[0].pages[0].buffer_index); + const uint64_t off = static_cast(resolved[0].pages[0].page_index) * kPageSize; + + TransferItem item; + item.src = pool; + item.src_offset = off; + item.dst = TransferRef::HostBytes(out->data(), out->size()); + item.size = resolved[0].size; + std::vector failed; + return engine_.Transfer({item}, &failed); + } + + static std::vector Payload(size_t n, char seed) { + std::vector v(n); + std::iota(v.begin(), v.end(), seed); + return v; + } + + SsdBackendTest() { engine_.AddEngine(std::make_unique()); } + + fs::path dir_; + LocalOnlyRegistrar registrar_; + CompositeTransferEngine engine_; + std::unique_ptr backend_; +}; + +// --------------------------------------------------------------------------- +// The design invariant +// --------------------------------------------------------------------------- + +// The reason SSD needed no new engine, no new TransferRef kind and no chaining: +// what it publishes is ordinary registered HOST memory. If this ever changes, +// every claim in ssd_backend.h's header comment stops holding. +TEST_F(SsdBackendTest, PublishesPlainHostMemorySoNoNewEngineIsNeeded) { + SsdBackend* backend = Start(); + + EXPECT_EQ(backend->Tier(), TierType::SSD); + EXPECT_EQ(registrar_.registrations, 1); // one arena, one registration + EXPECT_EQ(registrar_.last_loc, mori::io::MemoryLocationType::CPU); + + ASSERT_EQ(backend->BufferCount(), 1u); + const TransferRef ref = backend->BufferRef(0); + ASSERT_TRUE(ref.Valid()); + EXPECT_TRUE(ref.HasHostPtr()); + EXPECT_EQ(ref.loc, mori::io::MemoryLocationType::CPU); + + // Which means the plain host engine claims its pairs, with nothing new added. + LocalCopyEngine local; + EXPECT_TRUE(local.CanHandle(ref, ref)); +} + +TEST_F(SsdBackendTest, ReportsAnOutOfRangeBufferAsInvalid) { + SsdBackend* backend = Start(); + EXPECT_FALSE(backend->BufferRef(1).Valid()); +} + +// --------------------------------------------------------------------------- +// Put / Get +// --------------------------------------------------------------------------- + +TEST_F(SsdBackendTest, PutThenGetReturnsTheSameBytes) { + SsdBackend* backend = Start(); + const auto payload = Payload(4096, 5); + + ASSERT_TRUE(Put(backend, "key-a", payload)); + + std::vector readback; + ASSERT_TRUE(Get(backend, "key-a", &readback)); + ASSERT_EQ(readback.size(), payload.size()); + EXPECT_EQ(std::memcmp(readback.data(), payload.data(), payload.size()), 0); +} + +TEST_F(SsdBackendTest, MultiPageObjectRoundTripsWithoutChangingPageSize) { + SsdBackend* backend = Start(/*staging_pages=*/4); + const auto payload = Payload(2 * kPageSize + 123, 17); + + auto allocated = backend->BatchAllocate({AllocateRequest{"large", payload.size()}}); + ASSERT_EQ(allocated.size(), 1u); + ASSERT_EQ(allocated[0].outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_EQ(allocated[0].pages.size(), 3u); + EXPECT_EQ(allocated[0].page_size, kPageSize); + EXPECT_EQ(allocated[0].pages[1].page_index, allocated[0].pages[0].page_index + 1); + EXPECT_EQ(allocated[0].pages[2].page_index, allocated[0].pages[1].page_index + 1); + ASSERT_TRUE(backend->BatchAbort({allocated[0].slot_id})[0]); + + ASSERT_TRUE(Put(backend, "large", payload)); + auto resolved = backend->BatchResolve({"large"}, false); + ASSERT_TRUE(resolved[0].found); + EXPECT_EQ(resolved[0].pages.size(), 3u); + std::vector readback; + ASSERT_TRUE(Get(backend, "large", &readback)); + EXPECT_EQ(readback, payload); +} + +TEST_F(SsdBackendTest, DsV4SizedObjectRoundTripsWith64KiBPages) { + constexpr size_t kDsV4SwaObjectSize = 9135360; + SsdBackend* backend = Start(/*staging_pages=*/160); + const auto payload = Payload(kDsV4SwaObjectSize, 23); + + ASSERT_TRUE(Put(backend, "deepseek-v4-swa", payload)); + auto resolved = backend->BatchResolve({"deepseek-v4-swa"}, false); + ASSERT_TRUE(resolved[0].found); + EXPECT_EQ(resolved[0].pages.size(), (kDsV4SwaObjectSize + kPageSize - 1) / kPageSize); + + std::vector readback; + ASSERT_TRUE(Get(backend, "deepseek-v4-swa", &readback)); + EXPECT_EQ(readback, payload); +} + +TEST_F(SsdBackendTest, GetOfAnUnknownKeyMisses) { + SsdBackend* backend = Start(); + auto resolved = backend->BatchResolve({"never-written"}, false); + ASSERT_EQ(resolved.size(), 1u); + EXPECT_FALSE(resolved[0].found); +} + +TEST_F(SsdBackendTest, CommitQueuesAnSsdAddEvent) { + SsdBackend* backend = Start(); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 1))); + + const auto events = backend->DrainPendingEvents(); + ASSERT_EQ(events.size(), 1u); + EXPECT_EQ(events[0].kind, KvEvent::Kind::ADD); + EXPECT_EQ(events[0].key, "key-a"); + EXPECT_EQ(events[0].tier, TierType::SSD); + + EXPECT_TRUE(backend->DrainPendingEvents().empty()); // outbox cleared +} + +TEST_F(SsdBackendTest, MultipleKeysRoundTripIndependently) { + SsdBackend* backend = Start(); + const auto a = Payload(2048, 1); + const auto b = Payload(4096, 64); + + ASSERT_TRUE(Put(backend, "key-a", a)); + ASSERT_TRUE(Put(backend, "key-b", b)); + EXPECT_EQ(backend->OwnedKeyCount(), 2u); + + std::vector got_a, got_b; + ASSERT_TRUE(Get(backend, "key-a", &got_a)); + ASSERT_TRUE(Get(backend, "key-b", &got_b)); + EXPECT_EQ(got_a, a); + EXPECT_EQ(got_b, b); +} + +// --------------------------------------------------------------------------- +// The staging arena — what actually differs from DRAM +// --------------------------------------------------------------------------- + +// A staging page is BORROWED for a write, not the key's home. If Commit leaked +// it the backend would wedge after `staging_pages` puts, which is the bug this +// guards: with 2 pages, 10 sequential puts must all succeed. +TEST_F(SsdBackendTest, CommitReturnsTheStagingPageToTheArena) { + SsdBackend* backend = Start(/*staging_pages=*/2); + for (int i = 0; i < 10; ++i) { + ASSERT_TRUE(Put(backend, "key-" + std::to_string(i), Payload(512, static_cast(i)))) + << "put " << i << " failed — staging page leaked?"; + } + EXPECT_EQ(backend->OwnedKeyCount(), 10u); +} + +TEST_F(SsdBackendTest, AbortReturnsTheStagingPage) { + SsdBackend* backend = Start(/*staging_pages=*/1); + + auto first = backend->BatchAllocate({AllocateRequest{"key-a", 1024}}); + ASSERT_EQ(first[0].outcome, AllocateOutcome::kSuccessAllocated); + + // The only page is taken, so a second reservation must report exhaustion. + auto blocked = backend->BatchAllocate({AllocateRequest{"key-b", 1024}}); + EXPECT_EQ(blocked[0].outcome, AllocateOutcome::kFailedNoSpace); + + ASSERT_TRUE(backend->BatchAbort({first[0].slot_id})[0]); + + auto after = backend->BatchAllocate({AllocateRequest{"key-b", 1024}}); + EXPECT_EQ(after[0].outcome, AllocateOutcome::kSuccessAllocated); +} + +TEST_F(SsdBackendTest, AbortIsIdempotent) { + SsdBackend* backend = Start(); + auto allocated = backend->BatchAllocate({AllocateRequest{"key-a", 1024}}); + ASSERT_TRUE(backend->BatchAbort({allocated[0].slot_id})[0]); + EXPECT_TRUE(backend->BatchAbort({allocated[0].slot_id})[0]); // already gone + EXPECT_TRUE(backend->BatchAbort({99999})[0]); // never existed +} + +TEST_F(SsdBackendTest, CommitOfAnAbortedSlotFails) { + SsdBackend* backend = Start(); + auto allocated = backend->BatchAllocate({AllocateRequest{"key-a", 1024}}); + ASSERT_TRUE(backend->BatchAbort({allocated[0].slot_id})[0]); + + auto committed = backend->BatchCommit({CommitRequest{allocated[0].slot_id, "key-a"}}); + EXPECT_FALSE(committed[0].success); +} + +TEST_F(SsdBackendTest, AbortReturnsEveryPageOfAMultiPageSpan) { + SsdBackend* backend = Start(/*staging_pages=*/3); + auto allocated = backend->BatchAllocate({AllocateRequest{"large", 2 * kPageSize}}); + ASSERT_EQ(allocated[0].outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_EQ(allocated[0].pages.size(), 2u); + + EXPECT_EQ(backend->BatchAllocate({AllocateRequest{"also-large", 2 * kPageSize}})[0].outcome, + AllocateOutcome::kFailedNoSpace); + ASSERT_TRUE(backend->BatchAbort({allocated[0].slot_id})[0]); + + auto all_pages = backend->BatchAllocate({AllocateRequest{"largest", 3 * kPageSize}}); + ASSERT_EQ(all_pages[0].outcome, AllocateOutcome::kSuccessAllocated); + EXPECT_EQ(all_pages[0].pages.size(), 3u); +} + +// An object larger than the entire arena is a permanent shape failure, not +// transient pressure that another allocation retry could solve. +TEST_F(SsdBackendTest, RejectsAKeyLargerThanTheStagingArena) { + SsdBackend* backend = Start(/*staging_pages=*/2); + auto allocated = backend->BatchAllocate({AllocateRequest{"too-big", 2 * kPageSize + 1}}); + ASSERT_EQ(allocated.size(), 1u); + EXPECT_EQ(allocated[0].outcome, AllocateOutcome::kFailed); +} + +TEST_F(SsdBackendTest, RejectsAKeyLargerThanSsdEvenWhenItFitsStaging) { + SsdBackend* backend = + Start(/*staging_pages=*/4, std::chrono::milliseconds{500}, /*ssd_capacity=*/2 * kPageSize); + auto allocated = backend->BatchAllocate({AllocateRequest{"too-big", 2 * kPageSize + 1}}); + ASSERT_EQ(allocated.size(), 1u); + EXPECT_EQ(allocated[0].outcome, AllocateOutcome::kFailed); + EXPECT_EQ(backend->Capacity().max_allocatable_bytes, 2 * kPageSize); +} + +TEST_F(SsdBackendTest, RejectsAnOverflowingStagingArenaConfiguration) { + SsdBackend::Config cfg; + cfg.page_size = std::numeric_limits::max(); + cfg.staging_pages = 2; + EXPECT_THROW((void)SsdBackend(std::move(cfg)), std::invalid_argument); +} + +TEST_F(SsdBackendTest, DedupsAKeyAlreadyOnTheTier) { + SsdBackend* backend = Start(); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 1))); + + auto again = backend->BatchAllocate({AllocateRequest{"key-a", 1024}}); + EXPECT_EQ(again[0].outcome, AllocateOutcome::kSuccessAlreadyExists); +} + +// The read lease is what keeps staged bytes alive while a reader consumes them; +// the reaper is what stops that from being a leak. +TEST_F(SsdBackendTest, ReaperReclaimsStagingPagesFromExpiredReadLeases) { + SsdBackend* backend = Start(/*staging_pages=*/1, /*read_lease_ttl=*/std::chrono::milliseconds{1}); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 1))); + + auto resolved = backend->BatchResolve({"key-a"}, false); + ASSERT_TRUE(resolved[0].found); + + // The single page is now leased, so a write cannot get one. + EXPECT_EQ(backend->BatchAllocate({AllocateRequest{"key-b", 512}})[0].outcome, + AllocateOutcome::kFailedNoSpace); + + std::this_thread::sleep_for(std::chrono::milliseconds{10}); + backend->RunReaperOnceForTest(); + + EXPECT_EQ(backend->BatchAllocate({AllocateRequest{"key-b", 512}})[0].outcome, + AllocateOutcome::kSuccessAllocated); +} + +// A second reader of an already-staged key reuses the page rather than issuing +// a redundant SSD read — and, with one page configured, could not get a second +// page even if it tried. +TEST_F(SsdBackendTest, ConcurrentResolvesShareOneStagedPage) { + SsdBackend* backend = Start(/*staging_pages=*/1); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 9))); + + auto first = backend->BatchResolve({"key-a"}, false); + ASSERT_TRUE(first[0].found); + auto second = backend->BatchResolve({"key-a"}, false); + ASSERT_TRUE(second[0].found); + + EXPECT_EQ(first[0].pages[0].page_index, second[0].pages[0].page_index); + EXPECT_EQ(first[0].size, second[0].size); +} + +TEST_F(SsdBackendTest, DuplicateKeysInOneBatchShareOneMultiPageSpan) { + SsdBackend* backend = Start(/*staging_pages=*/3); + ASSERT_TRUE(Put(backend, "large", Payload(2 * kPageSize + 1, 9))); + + auto resolved = backend->BatchResolve({"large", "large"}, false); + ASSERT_EQ(resolved.size(), 2u); + ASSERT_TRUE(resolved[0].found); + ASSERT_TRUE(resolved[1].found); + ASSERT_EQ(resolved[0].pages.size(), 3u); + EXPECT_EQ(resolved[0].pages, resolved[1].pages); + EXPECT_EQ(resolved[0].size, resolved[1].size); +} + +TEST_F(SsdBackendTest, ResolveReportsBusyWhenStagingIsExhausted) { + SsdBackend* backend = Start(/*staging_pages=*/1); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 1))); + ASSERT_TRUE(Put(backend, "key-b", Payload(1024, 2))); + + auto a = backend->BatchResolve({"key-a"}, false); + ASSERT_TRUE(a[0].found); // takes the only page and holds it under a lease + + auto b = backend->BatchResolve({"key-b"}, false); + EXPECT_FALSE(b[0].found); + EXPECT_EQ(b[0].outcome, ResolveOutcome::kBusy); +} + +TEST_F(SsdBackendTest, BusyBatchRollsBackEveryNewReservation) { + SsdBackend* backend = Start(/*staging_pages=*/4); + ASSERT_TRUE(Put(backend, "blocker", Payload(1024, 1))); + ASSERT_TRUE(Put(backend, "large-a", Payload(2 * kPageSize, 2))); + ASSERT_TRUE(Put(backend, "large-b", Payload(2 * kPageSize, 3))); + ASSERT_TRUE(backend->BatchResolve({"blocker"}, false)[0].found); // one page remains leased + + auto resolved = backend->BatchResolve({"large-a", "large-b"}, false); + ASSERT_EQ(resolved.size(), 2u); + EXPECT_EQ(resolved[0].outcome, ResolveOutcome::kBusy); + EXPECT_EQ(resolved[1].outcome, ResolveOutcome::kBusy); + + // The first two-page reservation was returned when the second could not fit: + // all three pages not held by blocker are available to one writer. + auto allocated = backend->BatchAllocate({AllocateRequest{"three-pages", 3 * kPageSize}}); + EXPECT_EQ(allocated[0].outcome, AllocateOutcome::kSuccessAllocated); +} + +TEST_F(SsdBackendTest, BatchLargerThanArenaFailsWithoutRetryLoop) { + SsdBackend* backend = Start(/*staging_pages=*/3); + ASSERT_TRUE(Put(backend, "large-a", Payload(2 * kPageSize, 2))); + ASSERT_TRUE(Put(backend, "large-b", Payload(2 * kPageSize, 3))); + + auto resolved = backend->BatchResolve({"large-a", "large-b"}, false); + ASSERT_EQ(resolved.size(), 2u); + EXPECT_EQ(resolved[0].outcome, ResolveOutcome::kFailed); + EXPECT_EQ(resolved[1].outcome, ResolveOutcome::kFailed); +} + +// --------------------------------------------------------------------------- +// Eviction and Clear +// --------------------------------------------------------------------------- + +TEST_F(SsdBackendTest, EvictFreesTheKeyAndReportsItsBytes) { + SsdBackend* backend = Start(); + const auto payload = Payload(2048, 3); + ASSERT_TRUE(Put(backend, "key-a", payload)); + backend->DrainPendingEvents(); // discard the ADD + + auto evicted = backend->Evict({"key-a"}); + ASSERT_EQ(evicted.size(), 1u); + EXPECT_EQ(evicted[0].key, "key-a"); + EXPECT_EQ(evicted[0].bytes_freed, payload.size()); + + EXPECT_FALSE(backend->BatchResolve({"key-a"}, false)[0].found); + + const auto events = backend->DrainPendingEvents(); + ASSERT_EQ(events.size(), 1u); + EXPECT_EQ(events[0].kind, KvEvent::Kind::REMOVE); + EXPECT_EQ(events[0].tier, TierType::SSD); +} + +TEST_F(SsdBackendTest, EvictOfAnUnknownKeyFreesNothing) { + SsdBackend* backend = Start(); + auto evicted = backend->Evict({"nope"}); + ASSERT_EQ(evicted.size(), 1u); + EXPECT_EQ(evicted[0].bytes_freed, 0u); +} + +// A staged read holds the bytes live; master retries on a later round. +TEST_F(SsdBackendTest, EvictSkipsAKeyWithALiveReadLease) { + SsdBackend* backend = Start(); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 1))); + ASSERT_TRUE(backend->BatchResolve({"key-a"}, false)[0].found); + + auto evicted = backend->Evict({"key-a"}); + EXPECT_EQ(evicted[0].bytes_freed, 0u); + EXPECT_TRUE(backend->BatchResolve({"key-a"}, false)[0].found); // still there +} + +TEST_F(SsdBackendTest, EvictResultsAreOnePerKeyInRequestOrder) { + SsdBackend* backend = Start(); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 1))); + + auto evicted = backend->Evict({"missing-1", "key-a", "missing-2"}); + ASSERT_EQ(evicted.size(), 3u); + EXPECT_EQ(evicted[0].key, "missing-1"); + EXPECT_EQ(evicted[1].key, "key-a"); + EXPECT_EQ(evicted[2].key, "missing-2"); + EXPECT_EQ(evicted[1].bytes_freed, 1024u); +} + +TEST_F(SsdBackendTest, ClearLocalDropsEverythingAndGatesAllocation) { + SsdBackend* backend = Start(); + ASSERT_TRUE(Put(backend, "key-a", Payload(1024, 1))); + + backend->ClearLocal(); + EXPECT_TRUE(backend->IsClearFullSyncPending()); + EXPECT_EQ(backend->OwnedKeyCount(), 0u); + + // Gated: no new owned key may appear before master acks the empty full sync. + EXPECT_EQ(backend->BatchAllocate({AllocateRequest{"key-b", 512}})[0].outcome, + AllocateOutcome::kFailed); + + backend->ClearFullSyncAcked(); + EXPECT_FALSE(backend->IsClearFullSyncPending()); + EXPECT_EQ(backend->BatchAllocate({AllocateRequest{"key-b", 512}})[0].outcome, + AllocateOutcome::kSuccessAllocated); +} + +TEST_F(SsdBackendTest, ShutdownDeregistersTheArena) { + SsdBackend* backend = Start(); + backend->Shutdown(); + EXPECT_EQ(registrar_.deregistrations, 1); + backend->Shutdown(); // idempotent + EXPECT_EQ(registrar_.deregistrations, 1); +} + +#ifdef UMBP_ENABLE_GDS +// End to end: with a GdsEngine in the transfer stack and an O_DIRECT-capable +// filesystem, Resolve publishes a FileRef (not a staging page) and the engine +// reads the segment range straight into GPU memory. Point TMPDIR at ext4/xfs +// (e.g. /mnt/gds); skips without a GPU or where the filesystem rejects O_DIRECT. +TEST(SsdBackendGds, ResolvePublishesFileRefAndReadsIntoDeviceMemory) { + // GDS is opt-in at runtime (default off). Enable it for this test only and + // restore on every exit path (including GTEST_SKIP), so sibling staging tests + // in this binary keep their default-off behaviour. + struct GdsEnvGuard { + GdsEnvGuard() { ::setenv("UMBP_ENABLE_GDS", "1", /*overwrite=*/1); } + ~GdsEnvGuard() { ::unsetenv("UMBP_ENABLE_GDS"); } + } gds_env_guard; + + int ndev = 0; + if (hipGetDeviceCount(&ndev) != hipSuccess || ndev == 0) GTEST_SKIP() << "no HIP device"; + + const auto dir = fs::temp_directory_path() / ("umbp_ssd_gds_" + std::to_string(::getpid())); + fs::remove_all(dir); + + // A transfer stack that can move a file into GPU memory: the memory engines + // for the Put, the GdsEngine for the zero-copy Get. + CompositeTransferEngine composite; + composite.AddEngine(std::make_unique()); + composite.AddEngine(std::make_unique()); + composite.AddEngine(std::make_unique()); + + SsdBackend::Config cfg; + cfg.page_size = kPageSize; + cfg.staging_pages = 8; + cfg.ssd.enabled = true; + cfg.ssd.ssd.enabled = true; + cfg.ssd.ssd.storage_dir = dir.string(); + cfg.ssd.ssd.capacity_bytes = 64ULL * 1024 * 1024; + cfg.ssd.ssd.io.backend = UMBPIoBackend::Posix; + SsdBackend backend(std::move(cfg)); + ASSERT_TRUE(backend.Init(&composite)); + + // Put: allocate a staging page, move the payload in, commit (spills to SSD). + std::vector payload(8192); + std::iota(payload.begin(), payload.end(), 3); + auto alloc = backend.BatchAllocate({AllocateRequest{"k/gds", payload.size()}}); + ASSERT_EQ(1u, alloc.size()); + ASSERT_EQ(AllocateOutcome::kSuccessAllocated, alloc[0].outcome); + { + TransferRef pool = backend.BufferRef(alloc[0].pages[0].buffer_index); + TransferItem in; + in.src = TransferRef::HostBytes(payload.data(), payload.size()); + in.dst = pool; + in.dst_offset = static_cast(alloc[0].pages[0].page_index) * kPageSize; + in.size = payload.size(); + ASSERT_TRUE(composite.Transfer({in}, nullptr)); + } + auto commit = backend.BatchCommit({CommitRequest{alloc[0].slot_id, "k/gds"}}); + ASSERT_EQ(1u, commit.size()); + ASSERT_TRUE(commit[0].success); + + auto resolved = backend.BatchResolve({"k/gds"}, /*include_descs=*/false); + ASSERT_EQ(1u, resolved.size()); + ASSERT_TRUE(resolved[0].found); + if (!resolved[0].file_ref.IsFile()) { + backend.Shutdown(); + fs::remove_all(dir); + GTEST_SKIP() << "GDS branch inactive (filesystem rejects O_DIRECT?)"; + } + EXPECT_EQ(resolved[0].size, payload.size()); + + // Read the file range straight into device memory through GdsEngine. + void* dbuf = nullptr; + ASSERT_EQ(hipSuccess, hipMalloc(&dbuf, payload.size())); + ASSERT_EQ(hipSuccess, hipMemset(dbuf, 0, payload.size())); + TransferItem out; + out.src = resolved[0].file_ref; + out.dst = TransferRef::HostBytes(dbuf, payload.size(), mori::io::MemoryLocationType::GPU, 0); + out.size = resolved[0].size; + std::vector failed; + ASSERT_TRUE(composite.Transfer({out}, &failed)); + EXPECT_TRUE(failed.empty()); + + std::vector got(payload.size(), 0); + ASSERT_EQ(hipSuccess, hipMemcpy(got.data(), dbuf, payload.size(), hipMemcpyDeviceToHost)); + EXPECT_EQ(payload, got); + + (void)hipFree(dbuf); + backend.Shutdown(); + fs::remove_all(dir); +} +#endif // UMBP_ENABLE_GDS + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_ssd_copy_pipeline.cpp b/tests/cpp/umbp/distributed/test_ssd_copy_pipeline.cpp index 1ee21b71c..2aa7edcd9 100644 --- a/tests/cpp/umbp/distributed/test_ssd_copy_pipeline.cpp +++ b/tests/cpp/umbp/distributed/test_ssd_copy_pipeline.cpp @@ -32,12 +32,9 @@ #include #include -#include "umbp/distributed/peer/peer_dram_allocator.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" -#include "umbp/distributed/peer/ssd_copy_pipeline.h" -#include "umbp/local/block_index/local_block_index.h" -#include "umbp/local/tiers/copy_pipeline.h" -#include "umbp/local/tiers/local_storage_manager.h" +#include "umbp/distributed/peer/backend/page_backend.h" +#include "umbp/distributed/peer/ssd/peer_ssd_manager.h" +#include "umbp/distributed/peer/ssd/ssd_copy_pipeline.h" namespace mori::umbp { namespace { @@ -47,7 +44,7 @@ namespace fs = std::filesystem; constexpr uint64_t kPageSize = 1024; // Concatenate a pin's segments into one buffer for content comparison. -std::string Concat(const PeerDramAllocator::DramCopyPin& pin) { +std::string Concat(const PageBackend::DramCopyPin& pin) { std::string out; for (const auto& [ptr, len] : pin.segments) { out.append(static_cast(ptr), len); @@ -55,42 +52,40 @@ std::string Concat(const PeerDramAllocator::DramCopyPin& pin) { return out; } -// ---- DramCopyPin unit tests (direct on PeerDramAllocator) ------------------- +// ---- DramCopyPin unit tests (direct on PageBackend) ------------------- class DramCopyPinTest : public ::testing::Test { protected: void SetUp() override { backing_.assign(kPageSize * 8, 0); - PeerDramAllocator::TierConfig dram; + PageBackend::TierConfig dram; dram.buffer_sizes = {kPageSize * 8}; dram.buffer_descs = {{0x01, 0x02}}; dram.buffer_bases = {backing_.data()}; - dram_ = std::make_unique(kPageSize, std::move(dram), - PeerDramAllocator::TierConfig{}, - /*pending_ttl=*/std::chrono::milliseconds{5000}, - /*read_lease_ttl=*/std::chrono::milliseconds{0}); + dram_ = std::make_unique(TierType::DRAM, kPageSize, std::move(dram), + /*pending_ttl=*/std::chrono::milliseconds{5000}, + /*read_lease_ttl=*/std::chrono::milliseconds{0}); } // Allocate, write `value` into its pages, commit. Backing memory is owned // by the test, so we resolve pages -> offset exactly like the real writer. void PutLocal(const std::string& key, const std::string& value) { - auto res = dram_->Allocate(key, value.size(), TierType::DRAM); - ASSERT_EQ(res.outcome, PeerDramAllocator::Outcome::kSuccessAllocated); - const auto& slot = *res.slot; + auto res = dram_->Allocate(key, value.size()); + ASSERT_EQ(res.outcome, AllocateOutcome::kSuccessAllocated); size_t off = 0; - for (const auto& p : slot.pages) { + for (const auto& p : res.pages) { const size_t bytes = std::min(kPageSize, value.size() - off); std::memcpy(backing_.data() + static_cast(p.page_index) * kPageSize, value.data() + off, bytes); off += bytes; } uint64_t committed = 0; - ASSERT_TRUE(dram_->Commit(slot.slot_id, key, committed)); + ASSERT_TRUE(dram_->Commit(res.slot_id, key, committed)); ASSERT_EQ(committed, value.size()); } std::vector backing_; - std::unique_ptr dram_; + std::unique_ptr dram_; }; TEST_F(DramCopyPinTest, AcquireResolvesSegmentsToCommittedBytes) { @@ -143,28 +138,27 @@ TEST_F(DramCopyPinTest, EvictBlockedWhilePinnedThenAllowedAfterRelease) { TEST(DramCopyPinNonContiguous, SegmentsSpanMultipleBuffers) { // Two 1-page buffers force a cross-buffer page set for a 2-page key. std::vector b0(kPageSize, 0), b1(kPageSize, 0); - PeerDramAllocator::TierConfig dram; + PageBackend::TierConfig dram; dram.buffer_sizes = {kPageSize, kPageSize}; dram.buffer_descs = {{0x01}, {0x02}}; dram.buffer_bases = {b0.data(), b1.data()}; - PeerDramAllocator alloc(kPageSize, std::move(dram), PeerDramAllocator::TierConfig{}, - std::chrono::milliseconds{5000}, std::chrono::milliseconds{0}); + PageBackend alloc(TierType::DRAM, kPageSize, std::move(dram), std::chrono::milliseconds{5000}, + std::chrono::milliseconds{0}); const std::string value(kPageSize + 5, 'Q'); - auto res = alloc.Allocate("k", value.size(), TierType::DRAM); - ASSERT_EQ(res.outcome, PeerDramAllocator::Outcome::kSuccessAllocated); - const auto& slot = *res.slot; - ASSERT_EQ(slot.pages.size(), 2u); + auto res = alloc.Allocate("k", value.size()); + ASSERT_EQ(res.outcome, AllocateOutcome::kSuccessAllocated); + ASSERT_EQ(res.pages.size(), 2u); std::vector bases = {b0.data(), b1.data()}; size_t off = 0; - for (const auto& p : slot.pages) { + for (const auto& p : res.pages) { const size_t bytes = std::min(kPageSize, value.size() - off); std::memcpy(bases[p.buffer_index] + static_cast(p.page_index) * kPageSize, value.data() + off, bytes); off += bytes; } uint64_t committed = 0; - ASSERT_TRUE(alloc.Commit(slot.slot_id, "k", committed)); + ASSERT_TRUE(alloc.Commit(res.slot_id, "k", committed)); auto pin = alloc.AcquireDramCopyPin("k"); ASSERT_TRUE(pin.has_value()); @@ -184,13 +178,13 @@ class SsdCopyPipelineTest : public ::testing::Test { fs::remove_all(dir_); backing_.assign(kPageSize * 16, 0); - PeerDramAllocator::TierConfig dram; + PageBackend::TierConfig dram; dram.buffer_sizes = {kPageSize * 16}; dram.buffer_descs = {{0x01}}; dram.buffer_bases = {backing_.data()}; - dram_ = std::make_unique( - kPageSize, std::move(dram), PeerDramAllocator::TierConfig{}, - std::chrono::milliseconds{5000}, std::chrono::milliseconds{0}); + dram_ = std::make_unique(TierType::DRAM, kPageSize, std::move(dram), + std::chrono::milliseconds{5000}, + std::chrono::milliseconds{0}); PeerSsdConfig ssd_cfg; ssd_cfg.enabled = true; @@ -207,18 +201,17 @@ class SsdCopyPipelineTest : public ::testing::Test { } void PutLocal(const std::string& key, const std::string& value) { - auto res = dram_->Allocate(key, value.size(), TierType::DRAM); - ASSERT_EQ(res.outcome, PeerDramAllocator::Outcome::kSuccessAllocated); - const auto& slot = *res.slot; + auto res = dram_->Allocate(key, value.size()); + ASSERT_EQ(res.outcome, AllocateOutcome::kSuccessAllocated); size_t off = 0; - for (const auto& p : slot.pages) { + for (const auto& p : res.pages) { const size_t bytes = std::min(kPageSize, value.size() - off); std::memcpy(backing_.data() + static_cast(p.page_index) * kPageSize, value.data() + off, bytes); off += bytes; } uint64_t committed = 0; - ASSERT_TRUE(dram_->Commit(slot.slot_id, key, committed)); + ASSERT_TRUE(dram_->Commit(res.slot_id, key, committed)); } bool WaitForSsd(const std::string& key, std::chrono::milliseconds timeout) { @@ -232,7 +225,7 @@ class SsdCopyPipelineTest : public ::testing::Test { fs::path dir_; std::vector backing_; - std::unique_ptr dram_; + std::unique_ptr dram_; std::unique_ptr ssd_; }; @@ -339,34 +332,5 @@ TEST_F(SsdCopyPipelineTest, QuiesceThenClearLeavesNoStaleSsdState) { pipeline.Stop(); } - -TEST(LocalCopyPipelineDrainTest, DrainWaitsForQueuedCopies) { - UMBPConfig cfg; - cfg.role = UMBPRole::SharedSSDLeader; - cfg.force_ssd_copy_on_write = true; - cfg.dram.capacity_bytes = 1 << 20; - cfg.ssd.enabled = true; - cfg.ssd.ssd_backend = "dummy_storage"; - cfg.ssd.capacity_bytes = 1 << 20; - cfg.copy_pipeline.async_enabled = true; - cfg.copy_pipeline.worker_threads = 1; - cfg.copy_pipeline.queue_depth = 8; - - LocalBlockIndex index; - LocalStorageManager storage(cfg, &index); - CopyPipeline pipeline(storage, cfg.copy_pipeline, cfg.ResolveRole()); - - std::string payload = "copy-pipeline-drain"; - ASSERT_TRUE( - storage.WriteFromPtr("k", reinterpret_cast(payload.data()), payload.size())); - index.Insert("k", {StorageTier::CPU_DRAM, 0, payload.size()}); - - ASSERT_TRUE(pipeline.MaybeCopyToSharedSSD("k")); - EXPECT_TRUE(pipeline.Drain(std::chrono::seconds(2))); - auto* ssd = storage.GetTier(StorageTier::LOCAL_SSD); - ASSERT_NE(ssd, nullptr); - EXPECT_TRUE(ssd->Exists("k")); -} - } // namespace } // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_ssd_ranged_io.cpp b/tests/cpp/umbp/distributed/test_ssd_ranged_io.cpp new file mode 100644 index 000000000..923c61b31 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_ssd_ranged_io.cpp @@ -0,0 +1,479 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// RANGED I/O ON AN SSD-MEDIUM NODE (doc/design-ssd-ranged-io.md, D0) +// +// DistributedClient::SupportsRangedIO() used to answer false whenever the +// node's medium was SSD, on the reasoning that ranged access maps object +// ranges onto pages a backend publishes as in-process endpoints and SSD +// publishes storage refs instead. SsdBackend does not do that: it weighs a +// file endpoint against staging and picks staging (ssd_backend.h), so it +// publishes ordinary registered host pages and a resolved SSD key reaches +// BuildLocalRangeTransfers looking exactly like a DRAM one. +// +// The gate was therefore refusing paths that already worked. This suite pins +// all four of them so they cannot regress: +// +// 1. The capability itself — the arena is the only condition, on SSD as on +// DRAM. +// 2. Local ranged put+get — scattered sources tile an object onto an SSD +// node's own medium; a subset reads back. +// 3. Remote ranged put+get — the same across the wire to an SSD-only peer, +// through the scratch arena. +// 4. Subset semantics — a ranged get writes the requested bytes and +// nothing else, which is what distinguishes it from a whole-object read. +// +// WHAT THIS SUITE DELIBERATELY DOES NOT ASSERT: that any of it is faster. A +// resolve stages the whole object before the ranges are known, so a ranged get +// off SSD still reads the full object from the device. Making the device read +// follow the requested extent is D1, and the test that belongs with it counts +// bytes read at the TierBackend — an assertion this suite would pass without +// D1's code, which is exactly why it is not written here. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/common/config.h" +#include "umbp/distributed/config.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/pool_client.h" +#include "umbp/local/host_mem_allocator.h" +#include "umbp/standalone/standalone_server.h" +#include "umbp/umbp_client.h" + +namespace mori::umbp { +namespace { + +namespace fs = std::filesystem; + +constexpr size_t kPageSize = 4096; +// One key is one page on an SSD node (ssd_backend.h: a key larger than +// page_size is refused at BatchAllocate), so the object is a page. +constexpr size_t kObjectSize = kPageSize; +constexpr size_t kScratchSize = 8 * kObjectSize; +// A single page of DRAM ~= no capacity, so a put from this node must route out. +constexpr size_t kNoCapacity = kPageSize; + +uint16_t NextPeerServicePort() { + static std::atomic next{ + static_cast(58000 + (static_cast(::getpid()) % 4000))}; + return next.fetch_add(1); +} + +std::vector Pattern(size_t size, unsigned seed) { + std::vector out(size); + for (size_t i = 0; i < size; ++i) { + out[i] = static_cast(((i * 31) + seed * 7 + 11) & 0xff); + } + return out; +} + +bool WaitForExists(IUMBPClient* client, const std::string& key) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < deadline) { + const auto found = client->BatchExists({key}); + if (found.size() == 1 && found[0]) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + const auto found = client->BatchExists({key}); + return found.size() == 1 && found[0]; +} + +uintptr_t Addr(void* p) { return reinterpret_cast(p); } +uintptr_t Addr(const void* p) { return reinterpret_cast(p); } + +class SsdRangedIoTest : public ::testing::Test { + protected: + void SetUp() override { + static std::atomic counter{0}; + ssd_dir_ = fs::temp_directory_path() / ("umbp_ssd_ranged_" + std::to_string(::getpid()) + "_" + + std::to_string(counter.fetch_add(1))); + fs::remove_all(ssd_dir_); + + MasterServerConfig master_cfg; + master_cfg.listen_address = "0.0.0.0:0"; + // The registry derives a 2 s client heartbeat from this TTL. Left at the + // default, every WaitForExists below sits out a full publication interval. + master_cfg.registry_config.heartbeat_ttl = std::chrono::seconds(4); + master_ = std::make_unique(std::move(master_cfg)); + server_thread_ = std::thread([this] { master_->Run(); }); + for (int i = 0; i < 50 && master_->GetBoundPort() == 0; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_NE(master_->GetBoundPort(), 0) << "Master failed to start"; + master_addr_ = "localhost:" + std::to_string(master_->GetBoundPort()); + } + + void TearDown() override { + for (auto& c : clients_) { + if (c) c->Close(); + } + clients_.clear(); + if (master_) master_->Shutdown(); + if (server_thread_.joinable()) server_thread_.join(); + std::error_code ec; + fs::remove_all(ssd_dir_, ec); + } + + // Deliberately the full IUMBPClient, not the bare PoolClient the sibling + // ranges suite drives. SupportsRangedIO() lives on DistributedClient, and + // the point of every test here is to tie the data path to what the client + // ADVERTISES -- see the note above TEST 2. + UMBPConfig BaseConfig(const std::string& node_id, size_t ranged_scratch_size) { + UMBPConfig config; + config.dram.use_hugepages = false; + // The SSD tier's own knobs live on UMBPConfig, not on the distributed + // block: DistributedClient lowers config.ssd into PeerSsdConfig::ssd, and + // only the *enabled* flag comes from the medium selector. + config.ssd.storage_dir = ssd_dir_.string(); + config.ssd.capacity_bytes = 64ULL * 1024 * 1024; + config.ssd.io.backend = UMBPIoBackend::Posix; // avoid io_uring container flakiness + + UMBPDistributedConfig distributed; + distributed.master_config.node_id = node_id; + distributed.master_config.node_address = "127.0.0.1"; + distributed.master_config.master_address = master_addr_; + distributed.io_engine.host = "0.0.0.0"; + distributed.io_engine.port = 0; + distributed.peer_service_port = NextPeerServicePort(); + distributed.dram_page_size = kPageSize; + // Smaller than one object, so a remote ranged transfer can only pass + // through the registered arena and never the legacy staging path. + distributed.staging_buffer_size = 1024; + distributed.ranged_scratch_size = ranged_scratch_size; + // Off, so a remote read stays remote and the assertions describe the SSD + // peer rather than a re-cached copy on the reader. + distributed.cache_remote_fetches = false; + distributed.ssd_staging_buffer_slots = 8; + config.distributed = std::move(distributed); + return config; + } + + UMBPConfig SsdConfig(const std::string& node_id, size_t ranged_scratch_size = kScratchSize) { + auto config = BaseConfig(node_id, ranged_scratch_size); + // The selector IS the opt-in: this is what sets PeerSsdConfig::enabled. + config.distributed->medium = UMBPMedium::SSD; + // Present and to be ignored -- an SSD node gets no host pool whatever this + // says (see test_medium_selection.cpp). + config.dram.capacity_bytes = 8 << 20; + return config; + } + + UMBPConfig DramConfig(const std::string& node_id, size_t capacity, + size_t ranged_scratch_size = kScratchSize) { + auto config = BaseConfig(node_id, ranged_scratch_size); + config.distributed->medium = UMBPMedium::DRAM; + config.dram.capacity_bytes = capacity; + return config; + } + + IUMBPClient* Start(UMBPConfig config) { + std::string validation_error; + EXPECT_TRUE(config.Validate(&validation_error)) << validation_error; + auto client = CreateUMBPClient(config); + if (client == nullptr) return nullptr; + clients_.push_back(std::move(client)); + return clients_.back().get(); + } + + fs::path ssd_dir_; + std::string master_addr_; + std::unique_ptr master_; + std::thread server_thread_; + std::vector> clients_; +}; + +// --------------------------------------------------------------------------- +// 1. The capability +// --------------------------------------------------------------------------- + +// The whole of D0. Before it, the second EXPECT_TRUE was false: an SSD node +// could configure an arena and still be told it had no ranged I/O. +TEST_F(SsdRangedIoTest, SsdMediumSupportsRangedIoOnceTheArenaExists) { + auto* without = Start(SsdConfig("ssd-no-scratch", /*ranged_scratch_size=*/0)); + ASSERT_NE(without, nullptr); + EXPECT_EQ(without->GetDeploymentMode(), UMBPDeploymentMode::Distributed); + EXPECT_FALSE(without->SupportsRangedIO()) << "the arena is still the opt-in on SSD"; + + auto* with = Start(SsdConfig("ssd-with-scratch")); + ASSERT_NE(with, nullptr); + EXPECT_TRUE(with->SupportsRangedIO()) << "medium must no longer be a condition"; +} + +// --------------------------------------------------------------------------- +// 2-4. Local: the node's own SSD medium +// --------------------------------------------------------------------------- +// +// Each of these opens by asserting SupportsRangedIO(), and that is not +// ceremony. The ranged entry points forward to PoolClient unconditionally +// (distributed_client.cpp) -- the flag is a DECLARATION callers consult, not a +// guard the data path enforces, which is why the sglang tree connector simply +// declined to ask. So the contract worth pinning is the conjunction: the +// client says yes, AND the bytes are right. A test that skipped the assertion +// would pass against the pre-D0 gate and pin nothing. + +// A ranged put tiles its object out of scattered sources; the whole-object Get +// then proves the bytes were assembled in the right order on the device, which +// a ranged read-back alone would not (a consistently wrong mapping would agree +// with itself). +TEST_F(SsdRangedIoTest, LocalRangedPutAssemblesTheObjectOnSsd) { + auto* node = Start(SsdConfig("ssd-local-put")); + ASSERT_NE(node, nullptr); + ASSERT_TRUE(node->SupportsRangedIO()); + + const std::string key = "ssd-ranged-local-put"; + const auto object = Pattern(kObjectSize, 1); + const size_t quarter = kObjectSize / 4; + + // Deliberately out of order: the ranges tile the object, but the caller + // supplies them middle-first. + const std::vector srcs = {Addr(object.data() + quarter), Addr(object.data()), + Addr(object.data() + 2 * quarter)}; + const std::vector sizes = {quarter, quarter, kObjectSize - 2 * quarter}; + const std::vector offsets = {quarter, 0, 2 * quarter}; + + auto put = node->BatchPutRanges({key}, {kObjectSize}, {srcs}, {sizes}, {offsets}); + ASSERT_EQ(put, std::vector({true})); + ASSERT_TRUE(WaitForExists(node, key)); + + std::vector whole(kObjectSize, 0); + auto got = node->BatchGet({key}, {Addr(whole.data())}, {whole.size()}); + ASSERT_EQ(got, std::vector({true})); + EXPECT_EQ(whole, object) << "ranges landed at the wrong offsets on the device"; +} + +// The read direction, against an object written the ordinary way -- so a +// failure here is the ranged read, not a ranged write that happened to agree. +TEST_F(SsdRangedIoTest, LocalRangedGetReadsASubsetFromSsd) { + auto* node = Start(SsdConfig("ssd-local-get")); + ASSERT_NE(node, nullptr); + ASSERT_TRUE(node->SupportsRangedIO()); + + const std::string key = "ssd-ranged-local-get"; + const auto object = Pattern(kObjectSize, 2); + ASSERT_EQ(node->BatchPut({key}, {Addr(object.data())}, {object.size()}), + std::vector({true})); + ASSERT_TRUE(WaitForExists(node, key)); + + // Two disjoint windows that do NOT tile the object -- the read side is a + // subset, which is what separates it from a whole-object read. + std::vector head(64, 0); + std::vector middle(128, 0); + const std::vector dsts = {Addr(head.data()), Addr(middle.data())}; + const std::vector sizes = {head.size(), middle.size()}; + const std::vector offsets = {0, kObjectSize / 2}; + + auto got = node->BatchGetRanges({key}, {dsts}, {sizes}, {offsets}); + ASSERT_EQ(got, std::vector({true})); + EXPECT_EQ(std::memcmp(head.data(), object.data(), head.size()), 0); + EXPECT_EQ(std::memcmp(middle.data(), object.data() + kObjectSize / 2, middle.size()), 0); +} + +// A ranged get must touch the requested bytes and nothing else. Guard bytes +// on either side of the destination catch an implementation that quietly +// copies the whole staged object out -- which, on SSD, is precisely the shape +// the staging arena makes easy to do by accident. +TEST_F(SsdRangedIoTest, LocalRangedGetWritesOnlyTheRequestedBytes) { + auto* node = Start(SsdConfig("ssd-local-subset")); + ASSERT_NE(node, nullptr); + ASSERT_TRUE(node->SupportsRangedIO()); + + const std::string key = "ssd-ranged-subset"; + const auto object = Pattern(kObjectSize, 3); + ASSERT_EQ(node->BatchPut({key}, {Addr(object.data())}, {object.size()}), + std::vector({true})); + ASSERT_TRUE(WaitForExists(node, key)); + + constexpr size_t kGuard = 32; + constexpr size_t kWanted = 96; + std::vector buffer(kGuard + kWanted + kGuard, 0x7E); + const std::vector dsts = {Addr(buffer.data() + kGuard)}; + const std::vector sizes = {kWanted}; + const std::vector offsets = {kObjectSize - kWanted}; // the object's tail + + auto got = node->BatchGetRanges({key}, {dsts}, {sizes}, {offsets}); + ASSERT_EQ(got, std::vector({true})); + EXPECT_EQ(std::memcmp(buffer.data() + kGuard, object.data() + kObjectSize - kWanted, kWanted), 0); + for (size_t i = 0; i < kGuard; ++i) { + EXPECT_EQ(buffer[i], 0x7E) << "leading guard clobbered at " << i; + EXPECT_EQ(buffer[kGuard + kWanted + i], 0x7E) << "trailing guard clobbered at " << i; + } +} + +// --------------------------------------------------------------------------- +// 5. Remote: across the wire to an SSD-only peer +// --------------------------------------------------------------------------- + +// The caller has a page of DRAM and therefore no capacity, so master routes +// both operations to the SSD peer and they travel through the scratch arena. +// This is the pair the medium gate cost most: a caller honouring the flag +// declined before ever consulting master, on a deployment where EVERY node is +// SSD (pure-ssd-mode.md) and so every key is remote. +TEST_F(SsdRangedIoTest, RemoteRangedRoundTripThroughAnSsdPeer) { + auto* caller = Start(DramConfig("ssd-remote-caller", kNoCapacity)); + ASSERT_NE(caller, nullptr); + ASSERT_TRUE(caller->SupportsRangedIO()); + auto* target = Start(SsdConfig("ssd-remote-target")); + ASSERT_NE(target, nullptr); + ASSERT_TRUE(target->SupportsRangedIO()); + + const std::string key = "ssd-ranged-remote"; + const auto object = Pattern(kObjectSize, 4); + const size_t half = kObjectSize / 2; + + const std::vector srcs = {Addr(object.data()), Addr(object.data() + half)}; + const std::vector put_sizes = {half, kObjectSize - half}; + const std::vector put_offsets = {0, half}; + + auto put = caller->BatchPutRanges({key}, {kObjectSize}, {srcs}, {put_sizes}, {put_offsets}); + ASSERT_EQ(put, std::vector({true})) << "ranged put did not reach the SSD peer"; + // Polled on the CALLER, not the holder: what the ranged get below needs is + // not "the target has the bytes" but "the master can route the caller to + // them", and only the caller's lookup waits for that. Under local_first the + // holder answers Exists out of its own backend, which is true as soon as the + // put commits -- a publication interval before the master knows -- so polling + // `target` here would stop waiting for the very thing this test needs. + ASSERT_TRUE(WaitForExists(caller, key)) << "the object never landed on the SSD node"; + + std::vector front(200, 0); + std::vector tail(100, 0); + const std::vector dsts = {Addr(front.data()), Addr(tail.data())}; + const std::vector get_sizes = {front.size(), tail.size()}; + const std::vector get_offsets = {0, kObjectSize - tail.size()}; + + auto got = caller->BatchGetRanges({key}, {dsts}, {get_sizes}, {get_offsets}); + ASSERT_EQ(got, std::vector({true})) << "ranged get did not reach the SSD peer"; + EXPECT_EQ(std::memcmp(front.data(), object.data(), front.size()), 0); + EXPECT_EQ(std::memcmp(tail.data(), object.data() + kObjectSize - tail.size(), tail.size()), 0); +} + +// --------------------------------------------------------------------------- +// 6. Through the standalone UMBP server +// --------------------------------------------------------------------------- + +// The topology the feature is actually deployed in: an engine talks gRPC to a +// standalone server on its node, and the SERVER is the one holding the +// DistributedClient against master. Everything above drives DistributedClient +// in-process, so none of it exercises the hop that matters here -- +// supports_ranged_io is answered by the server's client and carried on the +// Ping response (umbp_standalone.proto), then CACHED in the worker for its +// lifetime. A capability that propagated wrongly, or not at all, would leave +// the deployment exactly where the medium gate left it: correct code that +// every caller declines to use. +TEST_F(SsdRangedIoTest, StandaloneServerOverAnSsdBackendAdvertisesAndServesRangedIo) { + const std::string address = + "unix:///tmp/umbp_ssd_ranged_srv_" + std::to_string(::getpid()) + ".sock"; + + // The server's own config: distributed + SSD. standalone_process is left + // unset, so CreateUMBPClient inside the server builds a DistributedClient + // rather than a client pointed at itself. + auto server_cfg = SsdConfig("ssd-standalone-backend"); + standalone::StandaloneServer server(server_cfg, address); + ASSERT_TRUE(server.Start()); + std::thread server_thread([&server]() { server.Run(); }); + + { + // The worker's config: standalone_process ONLY. The factory tests + // `distributed` first, so a config carrying both would silently produce a + // second DistributedClient instead of a forwarding one. + UMBPConfig worker_cfg; + UMBPStandaloneProcessConfig sp_cfg; + sp_cfg.address = address; + sp_cfg.startup_timeout_ms = 5000; + worker_cfg.standalone_process = sp_cfg; + + auto worker = CreateUMBPClient(worker_cfg); + ASSERT_NE(worker, nullptr); + EXPECT_EQ(worker->GetDeploymentMode(), UMBPDeploymentMode::StandaloneProcess); + EXPECT_EQ(worker->GetBackendMode(), UMBPDeploymentMode::Distributed) + << "the server must be reporting its DistributedClient, not a local one"; + ASSERT_TRUE(worker->SupportsRangedIO()) + << "the SSD backend's capability did not survive the Ping hop"; + + // Standalone-process mode moves bytes over SHARED MEMORY, not over the + // gRPC message: the worker registers an shm-backed region, the server maps + // it, and the wire carries (region_base, offset, size) triples that the + // server resolves back to its own mapping. Ordinary heap pointers fail + // ResolveRanges and come back as a per-key false -- which is what an + // unregistered buffer SHOULD do, and is worth knowing is the failure shape. + HostMemAllocator allocator; + HostBufferOptions opts; + opts.backing = HostBufferBacking::kAnonymousShm; + opts.prefault = false; + HostBufferHandle region = allocator.Alloc(4 * kObjectSize, opts); + ASSERT_TRUE(region.valid()); + ASSERT_TRUE( + worker->RegisterMemory(reinterpret_cast(region.ptr), region.mapped_size)); + + auto* base = static_cast(region.ptr); + char* src_region = base; // [0, kObjectSize) + char* dst_region = base + 2 * kObjectSize; // read destinations + + const std::string key = "ssd-ranged-standalone"; + const auto object = Pattern(kObjectSize, 5); + std::memcpy(src_region, object.data(), kObjectSize); + const size_t half = kObjectSize / 2; + + const std::vector srcs = {Addr(src_region), Addr(src_region + half)}; + const std::vector put_sizes = {half, kObjectSize - half}; + const std::vector put_offsets = {0, half}; + auto put = worker->BatchPutRanges({key}, {kObjectSize}, {srcs}, {put_sizes}, {put_offsets}); + ASSERT_EQ(put, std::vector({true})); + ASSERT_TRUE(WaitForExists(worker.get(), key)); + + constexpr size_t kHead = 96; + constexpr size_t kTail = 64; + char* head = dst_region; + char* tail = dst_region + kObjectSize; // a second, non-adjacent destination + std::memset(head, 0, kHead); + std::memset(tail, 0, kTail); + const std::vector dsts = {Addr(head), Addr(tail)}; + const std::vector get_sizes = {kHead, kTail}; + const std::vector get_offsets = {0, kObjectSize - kTail}; + auto got = worker->BatchGetRanges({key}, {dsts}, {get_sizes}, {get_offsets}); + ASSERT_EQ(got, std::vector({true})); + EXPECT_EQ(std::memcmp(head, object.data(), kHead), 0); + EXPECT_EQ(std::memcmp(tail, object.data() + kObjectSize - kTail, kTail), 0); + + worker->DeregisterMemory(reinterpret_cast(region.ptr)); + allocator.Free(region); + + worker->Close(); + } + + server.Shutdown(); + if (server_thread.joinable()) server_thread.join(); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_ssd_read_lease_gating.cpp b/tests/cpp/umbp/distributed/test_ssd_read_lease_gating.cpp deleted file mode 100644 index 1b105d0d1..000000000 --- a/tests/cpp/umbp/distributed/test_ssd_read_lease_gating.cpp +++ /dev/null @@ -1,97 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -// -// Pure-logic unit tests for the reader-side remote SSD read lease gating -// (umbp/distributed/ssd_read_lease.h). These cover the decision policy without -// a cluster / RDMA: the full PrepareSsdRead -> RDMA path is exercised at the -// RPC level in test_peer_ssd_read_rpc.cpp. Retryable outcomes are NO_SLOT and -// a reader-local lease expiry; rpc failures are not-served (RPC-test covered). -#include - -#include - -#include "umbp/distributed/ssd_read_lease.h" - -namespace mori::umbp::ssd_read_lease { -namespace { - -using std::chrono::milliseconds; -using std::chrono::steady_clock; - -// ---- LeaseExpired ---- - -TEST(SsdReadLeaseGating, NotExpiredBeforeDeadline) { - const auto t_send = steady_clock::now(); - EXPECT_FALSE(LeaseExpired(t_send, /*lease_ttl_ms=*/1000, t_send + milliseconds(500))); -} - -TEST(SsdReadLeaseGating, ExpiredAfterDeadline) { - const auto t_send = steady_clock::now(); - EXPECT_TRUE(LeaseExpired(t_send, /*lease_ttl_ms=*/1000, t_send + milliseconds(1001))); -} - -TEST(SsdReadLeaseGating, ExactlyAtDeadlineIsNotExpired) { - // Boundary: now == t_send + ttl uses '>' so it is still valid. - const auto t_send = steady_clock::now(); - EXPECT_FALSE(LeaseExpired(t_send, /*lease_ttl_ms=*/1000, t_send + milliseconds(1000))); -} - -TEST(SsdReadLeaseGating, ZeroTtlIsBornExpired) { - const auto t_send = steady_clock::now(); - EXPECT_FALSE(LeaseExpired(t_send, /*lease_ttl_ms=*/0, t_send)); // exactly t_send: valid - EXPECT_TRUE(LeaseExpired(t_send, /*lease_ttl_ms=*/0, t_send + milliseconds(1))); -} - -// ---- DecideSsdReadOutcome ---- -// Situation A (not expired): a good RDMA serves + releases; a failed RDMA is a -// hard error but still releases (the lease is still ours). -// Situation B (expired): always a transient retry, and NEVER release (the slot -// is left for the peer's TTL reclaim), regardless of whether the RDMA "worked". - -TEST(SsdReadLeaseGating, ValidAndRdmaOk_ServesAndReleases) { - const auto d = DecideSsdReadOutcome(/*expired=*/false, /*rdma_ok=*/true); - EXPECT_EQ(d.outcome, GateOutcome::kSuccess); - EXPECT_TRUE(d.release); -} - -TEST(SsdReadLeaseGating, ValidAndRdmaFailed_ErrorButReleases) { - const auto d = DecideSsdReadOutcome(/*expired=*/false, /*rdma_ok=*/false); - EXPECT_EQ(d.outcome, GateOutcome::kError); - EXPECT_TRUE(d.release); -} - -TEST(SsdReadLeaseGating, ExpiredWithRdmaOk_RetryNoRelease) { - // The dangerous case: RDMA "succeeded" but the lease elapsed, so the bytes - // are untrusted (the peer may have recycled the slot). Must NOT be success. - const auto d = DecideSsdReadOutcome(/*expired=*/true, /*rdma_ok=*/true); - EXPECT_EQ(d.outcome, GateOutcome::kRetry); - EXPECT_FALSE(d.release); -} - -TEST(SsdReadLeaseGating, ExpiredWithRdmaFailed_RetryNoRelease) { - const auto d = DecideSsdReadOutcome(/*expired=*/true, /*rdma_ok=*/false); - EXPECT_EQ(d.outcome, GateOutcome::kRetry); - EXPECT_FALSE(d.release); -} - -} // namespace -} // namespace mori::umbp::ssd_read_lease diff --git a/tests/cpp/umbp/distributed/test_ssd_reliability.cpp b/tests/cpp/umbp/distributed/test_ssd_reliability.cpp deleted file mode 100644 index 783a988a9..000000000 --- a/tests/cpp/umbp/distributed/test_ssd_reliability.cpp +++ /dev/null @@ -1,403 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -// -// Cross-component reliability tests: combinations no single-component test -// covers. All deterministic, no real disk / RDMA / master RPC: -// * the unified owned-location event source merges DRAM + SSD into one -// snapshot/delta (so a heartbeat full-sync ships SSD owned keys too); -// * a local SSD eviction's REMOVE SSD event converges the master -// metadata store while leaving the DRAM bucket intact; -// * tier-priority RouteGet over the real index picks DRAM, then SSD once the -// DRAM replica is removed; -// * crash-restart leftover is discarded at startup; -// * the SSD observability counters increment at the right events. -// -// (copy-pin vs DRAM evict is covered by test_ssd_copy_pipeline's -// EvictBlockedWhilePinnedThenAllowedAfterRelease; seq-gap -> full-sync by -// test_in_memory_master_metadata_store's heartbeat SeqGap cases.) -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/master/in_memory_master_metadata_store.h" -#include "umbp/distributed/peer/owned_location_source.h" -#include "umbp/distributed/peer/peer_ssd_manager.h" -#include "umbp/distributed/routing/route_get_strategy.h" -#include "umbp/distributed/types.h" -#include "umbp/local/tiers/tier_backend.h" - -namespace mori::umbp { -namespace { - -// Minimal in-memory TierBackend (mirrors the one in test_peer_ssd_eviction) so -// PeerSsdManager runs without real disk IO. Exposes a forced evict failure for -// the backend-failure counter and lets the test pre-seed bytes (crash leftover). -class FakeBackend : public TierBackend { - public: - explicit FakeBackend(size_t capacity) - : TierBackend(StorageTier::LOCAL_SSD), capacity_(capacity) {} - - bool Write(const std::string& key, const void* data, size_t size) override { - std::lock_guard lk(mu_); - auto it = store_.find(key); - size_t prev = (it == store_.end()) ? 0 : it->second.size(); - if (used_ - prev + size > capacity_) return false; - store_[key].assign(static_cast(data), static_cast(data) + size); - used_ = used_ - prev + size; - return true; - } - bool ReadIntoPtr(const std::string& key, uintptr_t dst, size_t size) override { - std::lock_guard lk(mu_); - auto it = store_.find(key); - if (it == store_.end() || it->second.size() != size) return false; - std::memcpy(reinterpret_cast(dst), it->second.data(), size); - return true; - } - bool Exists(const std::string& key) const override { - std::lock_guard lk(mu_); - return store_.count(key) != 0; - } - bool Evict(const std::string& key) override { - std::lock_guard lk(mu_); - if (fail_evict_) return false; - auto it = store_.find(key); - if (it == store_.end()) return false; - used_ -= it->second.size(); - store_.erase(it); - return true; - } - std::pair Capacity() const override { - std::lock_guard lk(mu_); - return {used_, capacity_}; - } - void Clear() override { - std::lock_guard lk(mu_); - ++clear_calls_; - store_.clear(); - used_ = 0; - } - void SetFailEvict(bool f) { - std::lock_guard lk(mu_); - fail_evict_ = f; - } - int clear_calls() const { - std::lock_guard lk(mu_); - return clear_calls_; - } - - private: - mutable std::mutex mu_; - std::unordered_map> store_; - size_t used_ = 0; - size_t capacity_; - bool fail_evict_ = false; - int clear_calls_ = 0; -}; - -std::vector> OneSeg(const std::string& s) { - return {{s.data(), s.size()}}; -} - -bool HasLoc(const std::vector& locs, const std::string& node, TierType tier) { - for (const auto& l : locs) { - if (l.node_id == node && l.tier == tier) return true; - } - return false; -} - -int CountTier(const std::vector& events, KvEvent::Kind kind, TierType tier) { - int n = 0; - for (const auto& e : events) { - if (e.kind == kind && e.tier == tier) ++n; - } - return n; -} - -// Under the merged store a block location can only be created through an -// ApplyHeartbeat from a registered (alive) node — locations no longer exist -// independently of a client record the way the old GlobalBlockIndex allowed. -// These helpers register a node once and apply heartbeat deltas with an -// ascending seq, standing in for the old GlobalBlockIndex::ApplyEvents. -constexpr uint64_t kGB = 1024ULL * 1024 * 1024; - -std::map MakeCaps() { - std::map caps; - caps[TierType::DRAM] = {8 * kGB, 8 * kGB}; - caps[TierType::SSD] = {8 * kGB, 8 * kGB}; - return caps; -} - -ClientRegistration MakeReg(const std::string& node_id) { - ClientRegistration reg; - reg.node_id = node_id; - reg.node_address = node_id + ":1"; - reg.peer_address = node_id + ":peer"; - reg.tier_capacities = MakeCaps(); - return reg; -} - -// A canned owned-location source standing in for PeerDramAllocator so the -// aggregation can be tested without standing up a DRAM allocator. -class FakeOwnedSource : public OwnedLocationSource { - public: - std::vector delta; - std::vector snapshot; - std::vector DrainPendingEvents() override { - std::vector out; - out.swap(delta); - return out; - } - std::vector SnapshotOwnedKeys() const override { return snapshot; } - std::vector SnapshotOwnedKeysForFullSync() override { - delta.clear(); // outbox dropped by the authoritative full-sync snapshot - return snapshot; - } -}; - -// --------------------------------------------------------------------------- -// Unified owned-location source: DRAM + SSD merge into one bundle. -// --------------------------------------------------------------------------- - -// A heartbeat full-sync snapshots ALL sources; SSD owned keys must be present -// alongside DRAM in the merged snapshot (otherwise master would drop the SSD -// tier on a seq-gap recovery). -TEST(SsdReliability, FullSyncSnapshotMergesDramAndSsdOwnedKeys) { - FakeOwnedSource dram; - dram.snapshot = {KvEvent{KvEvent::Kind::ADD, "d-key", TierType::DRAM, 10}}; - - auto be = std::make_unique(1'000'000); - PeerSsdManager ssd(std::move(be), 0.9, 0.7); - ASSERT_TRUE(ssd.Write("s-key", OneSeg("ssddata"), 7)); - ssd.DrainPendingEvents(); // the ADD SSD delta; snapshot is independent - - std::vector sources = {&dram, &ssd}; - auto snap = SnapshotAllSourcesForFullSync(sources); - - EXPECT_EQ(CountTier(snap, KvEvent::Kind::ADD, TierType::DRAM), 1); - EXPECT_EQ(CountTier(snap, KvEvent::Kind::ADD, TierType::SSD), 1); -} - -// A delta heartbeat drains ALL sources and concatenates into one list. -TEST(SsdReliability, DeltaDrainMergesDramAndSsdEvents) { - FakeOwnedSource dram; - dram.delta = {KvEvent{KvEvent::Kind::REMOVE, "d-key", TierType::DRAM, 0}}; - - auto be = std::make_unique(1'000'000); - PeerSsdManager ssd(std::move(be), 0.9, 0.7); - ASSERT_TRUE(ssd.Write("s-key", OneSeg("ssddata"), 7)); // queues ADD SSD delta - - std::vector sources = {&dram, &ssd}; - auto merged = DrainAllSources(sources); - - EXPECT_EQ(CountTier(merged, KvEvent::Kind::REMOVE, TierType::DRAM), 1); - EXPECT_EQ(CountTier(merged, KvEvent::Kind::ADD, TierType::SSD), 1); - // Draining again yields nothing (outbox cleared on both sources). - EXPECT_TRUE(DrainAllSources(sources).empty()); -} - -// --------------------------------------------------------------------------- -// SSD local eviction -> REMOVE SSD -> master metadata store converges. -// --------------------------------------------------------------------------- - -// A key mirrored on DRAM + SSD of one owner: a local SSD eviction emits -// REMOVE SSD, and applying that to the master store drops only the SSD bucket -// (the DRAM replica, owned independently, stays routable). -TEST(SsdReliability, LocalSsdEvictionRemoveConvergesMasterIndex) { - InMemoryMasterMetadataStore store; - const auto now = std::chrono::system_clock::now(); - ASSERT_TRUE(store.RegisterClient(MakeReg("owner"), now, std::chrono::seconds{30})); - - auto be = std::make_unique(1'000'000); - PeerSsdManager ssd(std::move(be), 0.9, 0.7); - - // DRAM replica added independently (a DRAM owner would emit this). - ASSERT_EQ(store - .ApplyHeartbeat("owner", /*seq=*/1, now, MakeCaps(), - {KvEvent{KvEvent::Kind::ADD, "k", TierType::DRAM, 100}}, - /*is_full_sync=*/false) - .status, - HeartbeatResult::APPLIED); - // SSD copy lands -> ADD SSD drained into the store. - ASSERT_TRUE(ssd.Write("k", OneSeg(std::string(100, 'x')), 100)); - ASSERT_EQ(store - .ApplyHeartbeat("owner", /*seq=*/2, now, MakeCaps(), ssd.DrainPendingEvents(), - /*is_full_sync=*/false) - .status, - HeartbeatResult::APPLIED); - - auto both = store.LookupBlock("k"); - ASSERT_TRUE(HasLoc(both, "owner", TierType::DRAM)); - ASSERT_TRUE(HasLoc(both, "owner", TierType::SSD)); - - // Local SSD eviction -> REMOVE SSD -> store drops only the SSD bucket. - ASSERT_TRUE(ssd.Evict("k")); - auto ssd_events = ssd.DrainPendingEvents(); - EXPECT_EQ(CountTier(ssd_events, KvEvent::Kind::REMOVE, TierType::SSD), 1); - ASSERT_EQ(store - .ApplyHeartbeat("owner", /*seq=*/3, now, MakeCaps(), ssd_events, - /*is_full_sync=*/false) - .status, - HeartbeatResult::APPLIED); - - auto after = store.LookupBlock("k"); - EXPECT_TRUE(HasLoc(after, "owner", TierType::DRAM)); // DRAM replica still routable - EXPECT_FALSE(HasLoc(after, "owner", TierType::SSD)); // SSD bucket converged away -} - -// --------------------------------------------------------------------------- -// Tier-priority RouteGet over the real store: DRAM first, SSD after evict. -// --------------------------------------------------------------------------- - -TEST(SsdReliability, TierPriorityRoutesDramThenSsdAfterDramRemoved) { - InMemoryMasterMetadataStore store; - const auto now = std::chrono::system_clock::now(); - const std::unordered_set kNoExclude; - ASSERT_TRUE(store.RegisterClient(MakeReg("owner"), now, std::chrono::seconds{30})); - ASSERT_EQ(store - .ApplyHeartbeat("owner", /*seq=*/1, now, MakeCaps(), - {KvEvent{KvEvent::Kind::ADD, "k", TierType::DRAM, 100}, - KvEvent{KvEvent::Kind::ADD, "k", TierType::SSD, 100}}, - /*is_full_sync=*/false) - .status, - HeartbeatResult::APPLIED); - - TierPriorityRouteGetStrategy strategy; - - auto locs = store.BatchLookupBlockForRouteGet({"k"}, kNoExclude, now, std::chrono::seconds{10}); - ASSERT_EQ(locs.size(), 1u); - auto dram_pick = strategy.Select(locs[0], "reader"); - EXPECT_EQ(dram_pick.tier, TierType::DRAM) << "prefers the fast DRAM replica"; - - // DRAM evicted -> only the SSD bucket remains -> RouteGet must serve from SSD. - ASSERT_EQ(store - .ApplyHeartbeat("owner", /*seq=*/2, now, MakeCaps(), - {KvEvent{KvEvent::Kind::REMOVE, "k", TierType::DRAM, 0}}, - /*is_full_sync=*/false) - .status, - HeartbeatResult::APPLIED); - auto locs2 = store.BatchLookupBlockForRouteGet({"k"}, kNoExclude, now, std::chrono::seconds{10}); - ASSERT_EQ(locs2.size(), 1u); - auto ssd_pick = strategy.Select(locs2[0], "reader"); - EXPECT_EQ(ssd_pick.tier, TierType::SSD) << "falls back to the surviving SSD replica"; - EXPECT_EQ(ssd_pick.node_id, "owner"); -} - -// --------------------------------------------------------------------------- -// Crash-restart leftover handling (discard). -// --------------------------------------------------------------------------- - -// After a crash owned_ is empty but the backend still holds bytes from the -// previous run. DiscardLeftoverOnStartup wipes them so used capacity starts -// at 0 (no divergence between the empty owned_ map and the physical device). -TEST(SsdReliability, StartupDiscardWipesLeftoverBytes) { - auto be = std::make_unique(1'000'000); - FakeBackend* raw = be.get(); - // Simulate a previous process's bytes left on the device. - ASSERT_TRUE(raw->Write("orphan-1", "leftover-a", 10)); - ASSERT_TRUE(raw->Write("orphan-2", "leftover-b", 10)); - ASSERT_GT(raw->Capacity().first, 0u); - - // Fresh manager: owned_ is empty, but the backend reports used > 0. - PeerSsdManager ssd(std::move(be), 0.9, 0.7); - EXPECT_TRUE(ssd.SnapshotOwnedKeys().empty()); - ASSERT_GT(ssd.Capacity().first, 0u); - - ssd.DiscardLeftoverOnStartup(); - - EXPECT_EQ(raw->clear_calls(), 1); - EXPECT_EQ(ssd.Capacity().first, 0u); // leftover gone -> consistent with empty owned_ -} - -TEST(SsdReliability, StartupDiscardOnCleanTierIsNoop) { - auto be = std::make_unique(1'000'000); - FakeBackend* raw = be.get(); - PeerSsdManager ssd(std::move(be), 0.9, 0.7); - - ssd.DiscardLeftoverOnStartup(); // used == 0 -> skip the wipe entirely - EXPECT_EQ(raw->clear_calls(), 0); -} - -// --------------------------------------------------------------------------- -// Observability counters increment at the right events. -// --------------------------------------------------------------------------- - -TEST(SsdReliability, ReadCountersTrackOutcomes) { - auto be = std::make_unique(1'000'000); - PeerSsdManager ssd(std::move(be), 0.9, 0.7); - ASSERT_TRUE(ssd.Write("k", OneSeg("0123456789"), 10)); - - std::vector buf(10); - EXPECT_EQ(ssd.PrepareRead("k", buf.data(), buf.size()).status, SsdReadStatus::kOk); - EXPECT_EQ(ssd.PrepareRead("absent", buf.data(), buf.size()).status, SsdReadStatus::kNotFound); - EXPECT_EQ(ssd.PrepareRead("k", buf.data(), /*cap=*/1).status, SsdReadStatus::kSizeTooLarge); - - EXPECT_EQ(ssd.ReadOk(), 1u); - EXPECT_EQ(ssd.ReadNotFound(), 1u); - EXPECT_EQ(ssd.ReadSizeTooLarge(), 1u); - EXPECT_EQ(ssd.ReadError(), 0u); -} - -TEST(SsdReliability, EvictionCountersTrackVictimsBytesAndBackendFailures) { - auto be = std::make_unique(1'000'000); - FakeBackend* raw = be.get(); - PeerSsdManager ssd(std::move(be), 0.9, 0.7); - ASSERT_TRUE(ssd.Write("a", OneSeg(std::string(40, 'a')), 40)); - ASSERT_TRUE(ssd.Write("b", OneSeg(std::string(60, 'b')), 60)); - - ASSERT_TRUE(ssd.Evict("a")); - EXPECT_EQ(ssd.EvictionVictims(), 1u); - EXPECT_EQ(ssd.EvictionBytesFreed(), 40u); - EXPECT_EQ(ssd.EvictionBackendFailures(), 0u); - - // Backend refuses the next evict -> the failure is counted, the key kept. - raw->SetFailEvict(true); - EXPECT_FALSE(ssd.Evict("b")); - EXPECT_EQ(ssd.EvictionBackendFailures(), 1u); - EXPECT_EQ(ssd.EvictionVictims(), 1u); // unchanged - EXPECT_TRUE(ssd.Exists("b")); -} - -TEST(SsdReliability, WatermarkEvictionCountsARound) { - // capacity 1000, high 0.9 (=>900), low 0.5 (=>500); 100-byte values. After - // the 9th write used hits 900 -> one eviction round runs. - auto be = std::make_unique(1000); - PeerSsdManager ssd(std::move(be), 0.9, 0.5); - std::string val(100, 'x'); - for (int i = 1; i <= 9; ++i) { - ASSERT_TRUE(ssd.Write("k" + std::to_string(i), OneSeg(val), val.size())); - } - EXPECT_GE(ssd.EvictionRounds(), 1u); - EXPECT_GE(ssd.EvictionVictims(), 1u); - EXPECT_LE(ssd.Capacity().first, 500u); -} - -} // namespace -} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_standalone_shm_ipc.cpp b/tests/cpp/umbp/distributed/test_standalone_shm_ipc.cpp index a1043c85b..c356c1b1e 100644 --- a/tests/cpp/umbp/distributed/test_standalone_shm_ipc.cpp +++ b/tests/cpp/umbp/distributed/test_standalone_shm_ipc.cpp @@ -22,6 +22,7 @@ // Copyright © Advanced Micro Devices, Inc. All rights reserved. // // MIT License +#include #include #include #include @@ -32,15 +33,18 @@ #include #include #include +#include #include #include #include +#include #include #include "umbp/local/host_mem_allocator.h" #include "umbp/standalone/ipc.h" #include "umbp/standalone/standalone_server.h" #include "umbp/umbp_client.h" +#include "umbp_standalone.grpc.pb.h" namespace mori::umbp { namespace { @@ -179,7 +183,53 @@ TEST(StandaloneShmIpcTest, RawUdsFdRegistrationTransfersFd) { EXPECT_TRUE(receiver_ok.load()); } -TEST(StandaloneShmIpcTest, StandaloneClientUsesNonZeroOffsetsAndCanReregister) { +TEST(StandaloneShmIpcTest, GpuRegistrationRejectsSsdBackedServer) { + const std::string suffix = std::to_string(getpid()); + const std::string address = "unix:///tmp/umbp_standalone_gpu_reject_" + suffix + ".sock"; + const std::string grpc_path = standalone::UnixPathFromGrpcAddress(address); + const std::string fd_path = standalone::DeriveFdSocketPath(address); + const std::string ssd_path = "/tmp/umbp_standalone_gpu_reject_ssd_" + suffix; + unlink(grpc_path.c_str()); + unlink(fd_path.c_str()); + std::filesystem::remove_all(ssd_path); + + UMBPConfig config; + config.dram.capacity_bytes = 1 << 20; + config.ssd.enabled = true; + config.ssd.storage_dir = ssd_path; + config.ssd.capacity_bytes = 4 << 20; + config.ssd.segment_size_bytes = 1 << 20; + // Serving SSD is now a MEDIUM, not the ssd.enabled flag: a server can carry + // SSD sizing and still serve DRAM, and GPU IPC is fine on DRAM. Name the + // medium, or this asserts against a server that has no reason to refuse. + config = WithEmbeddedDefaults(config); + config.distributed->medium = UMBPMedium::SSD; + standalone::StandaloneServer server(config, address); + ASSERT_TRUE(server.Start()); + std::thread server_thread([&]() { server.Run(); }); + + auto channel = grpc::CreateChannel(address, grpc::InsecureChannelCredentials()); + auto stub = ::umbp::UMBPStandalone::NewStub(channel); + grpc::ClientContext context; + ::umbp::RegisterMemoryRequest request; + request.set_kind(::umbp::MEMORY_KIND_GPU_IPC); + request.set_client_id("gpu-client"); + request.set_worker_base(0x1000); + request.set_size(4096); + ::umbp::BoolResponse response; + const grpc::Status status = stub->RegisterMemory(&context, request, &response); + ASSERT_TRUE(status.ok()); + EXPECT_FALSE(response.ok()); + EXPECT_NE(response.error().find("SSD"), std::string::npos); + + server.Shutdown(); + server_thread.join(); + unlink(grpc_path.c_str()); + unlink(fd_path.c_str()); + std::filesystem::remove_all(ssd_path); +} + +TEST(StandaloneShmIpcTest, WorkerRegistrationUsesNonZeroOffsetsAndCanReregister) { const std::string address = "unix:///tmp/umbp_standalone_e2e_" + std::to_string(getpid()) + ".grpc.sock"; const std::string grpc_path = standalone::UnixPathFromGrpcAddress(address); @@ -202,6 +252,10 @@ TEST(StandaloneShmIpcTest, StandaloneClientUsesNonZeroOffsetsAndCanReregister) { UMBPConfig client_cfg = server_cfg; auto client = CreateUMBPClient(client_cfg); ASSERT_EQ(client->GetDeploymentMode(), UMBPDeploymentMode::StandaloneProcess); + // The server's backend is a DistributedClient with no master -- there is no + // Local backend any more, and an embedded deployment is not a separate mode. + EXPECT_EQ(client->GetBackendMode(), UMBPDeploymentMode::Distributed); + EXPECT_TRUE(client->SupportsRangedIO()); HostMemAllocator allocator; HostBufferOptions opts; @@ -236,7 +290,7 @@ TEST(StandaloneShmIpcTest, StandaloneClientUsesNonZeroOffsetsAndCanReregister) { unlink(fd_path.c_str()); } -TEST(StandaloneShmIpcTest, StandaloneClientResolvesAcrossMultipleRegions) { +TEST(StandaloneShmIpcTest, WorkerRegistrationResolvesAcrossMultipleRegions) { const std::string address = "unix:///tmp/umbp_standalone_multiregion_" + std::to_string(getpid()) + ".grpc.sock"; const std::string grpc_path = standalone::UnixPathFromGrpcAddress(address); @@ -315,6 +369,47 @@ TEST(StandaloneShmIpcTest, StandaloneClientResolvesAcrossMultipleRegions) { for (int i = 0; i < 8; ++i) EXPECT_EQ(bytes_a[300 + i], static_cast(0xa0 + i)); for (int i = 0; i < 8; ++i) EXPECT_EQ(bytes_b[300 + i], static_cast(0xb0 + i)); + // One object can be assembled from, and read back into, ranges belonging to + // different registered regions. + for (int i = 0; i < 8; ++i) bytes_a[400 + i] = static_cast(0xc0 + i); + for (int i = 0; i < 8; ++i) bytes_b[400 + i] = static_cast(0xd0 + i); + auto range_put = client->BatchPutRanges( + {"range-key"}, {16}, + {{reinterpret_cast(bytes_b + 400), reinterpret_cast(bytes_a + 400)}}, + {{8, 8}}, {{8, 0}}); + ASSERT_EQ(range_put, std::vector({true})); + EXPECT_TRUE(client->Exists("range-key")); + + std::memset(bytes_a + 500, 0, 8); + std::memset(bytes_b + 500, 0, 8); + auto range_get = client->BatchGetRanges( + {"range-key"}, + {{reinterpret_cast(bytes_a + 500), reinterpret_cast(bytes_b + 500)}}, + {{8, 8}}, {{8, 0}}); + ASSERT_EQ(range_get, std::vector({true})); + for (int i = 0; i < 8; ++i) { + EXPECT_EQ(bytes_a[500 + i], static_cast(0xd0 + i)); + EXPECT_EQ(bytes_b[500 + i], static_cast(0xc0 + i)); + } + + // Malformed flattened range arrays are rejected before address resolution. + auto raw_stub = ::umbp::UMBPStandalone::NewStub( + grpc::CreateChannel(address, grpc::InsecureChannelCredentials())); + grpc::ClientContext malformed_context; + ::umbp::BatchRangeDataRequest malformed_request; + malformed_request.add_keys("malformed"); + malformed_request.add_range_counts(1); + malformed_request.add_shm_offsets(0); + // region_bases is deliberately missing. + malformed_request.add_sizes(8); + malformed_request.add_object_offsets(0); + malformed_request.add_object_sizes(8); + ::umbp::BatchBoolResponse malformed_response; + ASSERT_TRUE( + raw_stub->BatchPutRanges(&malformed_context, malformed_request, &malformed_response).ok()); + ASSERT_EQ(malformed_response.ok_size(), 1); + EXPECT_FALSE(malformed_response.ok(0)); + // A pointer outside every registered region fails cleanly (no crash, no hit). HostBufferHandle unregistered = allocator.Alloc(4096, opts); ASSERT_TRUE(unregistered.valid()); @@ -332,6 +427,246 @@ TEST(StandaloneShmIpcTest, StandaloneClientResolvesAcrossMultipleRegions) { unlink(fd_path.c_str()); } +TEST(StandaloneShmIpcTest, DeregistrationCannotUnmapAnInFlightDataOperation) { + constexpr size_t kValueSize = 32ULL << 20; + const std::string suffix = std::to_string(getpid()); + const std::string address = + "unix:///tmp/umbp_standalone_deregister_race_" + suffix + ".grpc.sock"; + const std::string grpc_path = standalone::UnixPathFromGrpcAddress(address); + const std::string fd_path = standalone::DeriveFdSocketPath(address); + unlink(grpc_path.c_str()); + unlink(fd_path.c_str()); + + UMBPConfig config; + config.dram.capacity_bytes = 2 * kValueSize; + config.ssd.enabled = false; + standalone::StandaloneServer server(config, address); + ASSERT_TRUE(server.Start()); + std::thread server_thread([&]() { server.Run(); }); + + HostMemAllocator allocator; + HostBufferOptions options; + options.backing = HostBufferBacking::kAnonymousShm; + options.prefault = false; + HostBufferHandle handle = allocator.Alloc(2 * kValueSize, options); + ASSERT_TRUE(handle.valid()); + auto* bytes = static_cast(handle.ptr); + std::memset(bytes, 0x5a, kValueSize); + std::memset(bytes + kValueSize, 0, kValueSize); + + auto allocation = HostMemAllocator::LookupShmAllocation(reinterpret_cast(handle.ptr), + handle.mapped_size); + ASSERT_TRUE(allocation.has_value()); + constexpr char kClientId[] = "deregister-race-client"; + std::string registration_error; + ASSERT_EQ(standalone::SendFdRegistration(fd_path, allocation->fd, kClientId, + reinterpret_cast(handle.ptr), + allocation->mapped_size, 5000, ®istration_error), + 0) + << registration_error; + + auto channel = grpc::CreateChannel(address, grpc::InsecureChannelCredentials()); + auto stub = ::umbp::UMBPStandalone::NewStub(channel); + { + grpc::ClientContext context; + ::umbp::RegisterMemoryRequest request; + request.set_kind(::umbp::MEMORY_KIND_HOST_SHM); + request.set_client_id(kClientId); + request.set_worker_base(reinterpret_cast(handle.ptr)); + request.set_size(allocation->mapped_size); + ::umbp::BoolResponse response; + const grpc::Status status = stub->RegisterMemory(&context, request, &response); + ASSERT_TRUE(status.ok()); + ASSERT_TRUE(response.ok()) << response.error(); + } + + { + grpc::ClientContext context; + ::umbp::PutRequest request; + request.set_key("large-value"); + request.set_client_id(kClientId); + request.set_region_base(reinterpret_cast(handle.ptr)); + request.set_shm_offset(0); + request.set_size(kValueSize); + ::umbp::BoolResponse response; + const grpc::Status status = stub->Put(&context, request, &response); + ASSERT_TRUE(status.ok()); + ASSERT_TRUE(response.ok()) << response.error(); + } + + std::atomic get_started{false}; + grpc::Status get_status; + ::umbp::BoolResponse get_response; + std::thread getter([&]() { + grpc::ClientContext context; + ::umbp::GetRequest request; + request.set_key("large-value"); + request.set_client_id(kClientId); + request.set_region_base(reinterpret_cast(handle.ptr)); + request.set_shm_offset(kValueSize); + request.set_size(kValueSize); + get_started.store(true, std::memory_order_release); + get_status = stub->Get(&context, request, &get_response); + }); + + while (!get_started.load(std::memory_order_acquire)) std::this_thread::yield(); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + grpc::Status deregister_status; + { + grpc::ClientContext context; + ::umbp::DeregisterMemoryRequest request; + request.set_client_id(kClientId); + ::umbp::Empty response; + deregister_status = stub->DeregisterMemory(&context, request, &response); + } + getter.join(); + + ASSERT_TRUE(deregister_status.ok()); + ASSERT_TRUE(get_status.ok()); + if (get_response.ok()) { + EXPECT_EQ(std::memcmp(bytes, bytes + kValueSize, kValueSize), 0); + } + + // Once deregistration returns, new operations must fail resolution rather + // than reaching a stale server-side mapping. + { + grpc::ClientContext context; + ::umbp::GetRequest request; + request.set_key("large-value"); + request.set_client_id(kClientId); + request.set_region_base(reinterpret_cast(handle.ptr)); + request.set_shm_offset(kValueSize); + request.set_size(kValueSize); + ::umbp::BoolResponse response; + const grpc::Status status = stub->Get(&context, request, &response); + EXPECT_TRUE(status.ok()); + EXPECT_FALSE(response.ok()); + } + + allocator.Free(handle); + server.Shutdown(); + server_thread.join(); + unlink(grpc_path.c_str()); + unlink(fd_path.c_str()); +} + +TEST(StandaloneShmIpcTest, WritersCompleteUnderContinuousReaderLoad) { + constexpr int kReaderCount = 8; + constexpr size_t kValueSize = 2ULL << 20; + const std::string address = + "unix:///tmp/umbp_standalone_writer_liveness_" + std::to_string(getpid()) + ".grpc.sock"; + const std::string grpc_path = standalone::UnixPathFromGrpcAddress(address); + const std::string fd_path = standalone::DeriveFdSocketPath(address); + unlink(grpc_path.c_str()); + unlink(fd_path.c_str()); + + UMBPConfig config; + config.dram.capacity_bytes = 8 * kValueSize; + config.ssd.enabled = false; + UMBPStandaloneProcessConfig standalone_config; + standalone_config.address = address; + standalone_config.startup_timeout_ms = 5000; + config.standalone_process = standalone_config; + + standalone::StandaloneServer server(config, address); + ASSERT_TRUE(server.Start()); + std::thread server_thread([&]() { server.Run(); }); + + HostMemAllocator allocator; + HostBufferOptions options; + options.backing = HostBufferBacking::kAnonymousShm; + options.prefault = false; + + std::vector> readers; + std::vector reader_buffers; + readers.reserve(kReaderCount); + reader_buffers.reserve(kReaderCount); + for (int i = 0; i < kReaderCount; ++i) { + reader_buffers.push_back(allocator.Alloc(kValueSize, options)); + ASSERT_TRUE(reader_buffers.back().valid()); + auto client = CreateUMBPClient(config); + ASSERT_TRUE(client->RegisterMemory(reinterpret_cast(reader_buffers.back().ptr), + reader_buffers.back().mapped_size)); + readers.push_back(std::move(client)); + } + + HostBufferHandle writer_buffer = allocator.Alloc(2 * kValueSize, options); + ASSERT_TRUE(writer_buffer.valid()); + auto writer = CreateUMBPClient(config); + ASSERT_TRUE(writer->RegisterMemory(reinterpret_cast(writer_buffer.ptr), + writer_buffer.mapped_size)); + auto* writer_bytes = static_cast(writer_buffer.ptr); + std::memset(writer_bytes, 0x31, kValueSize); + std::memset(writer_bytes + kValueSize, 0x72, kValueSize); + ASSERT_TRUE(writer->Put("read-hot-key", reinterpret_cast(writer_bytes), kValueSize)); + + std::atomic stop{false}; + std::atomic readers_started{0}; + std::vector reader_threads; + reader_threads.reserve(kReaderCount); + for (int i = 0; i < kReaderCount; ++i) { + reader_threads.emplace_back([&, i]() { + readers_started.fetch_add(1, std::memory_order_release); + while (!stop.load(std::memory_order_acquire)) { + readers[i]->Get("read-hot-key", reinterpret_cast(reader_buffers[i].ptr), + kValueSize); + } + }); + } + while (readers_started.load(std::memory_order_acquire) != kReaderCount) { + std::this_thread::yield(); + } + // Let every synchronous reader enter its steady RPC loop before introducing + // a writer; otherwise this could accidentally test an uncontended lock. + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + + auto run_with_deadline = [&](auto operation) { + std::atomic done{false}; + std::atomic ok{false}; + std::thread operation_thread([&]() { + ok.store(operation(), std::memory_order_relaxed); + done.store(true, std::memory_order_release); + }); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (!done.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + const bool completed_in_time = done.load(std::memory_order_acquire); + if (!completed_in_time) stop.store(true, std::memory_order_release); + operation_thread.join(); + return std::pair{completed_in_time, ok.load(std::memory_order_relaxed)}; + }; + + const auto put_result = run_with_deadline([&]() { + return writer->Put("writer-liveness-key", + reinterpret_cast(writer_bytes + kValueSize), kValueSize); + }); + EXPECT_TRUE(put_result.first) << "Put starved behind continuous readers"; + EXPECT_TRUE(put_result.second); + + std::pair clear_result{false, false}; + if (put_result.first) clear_result = run_with_deadline([&]() { return writer->Clear(); }); + EXPECT_TRUE(clear_result.first) << "Clear starved behind continuous readers"; + EXPECT_TRUE(clear_result.second); + + stop.store(true, std::memory_order_release); + for (auto& thread : reader_threads) thread.join(); + + for (size_t i = 0; i < readers.size(); ++i) { + readers[i]->DeregisterMemory(reinterpret_cast(reader_buffers[i].ptr)); + readers[i]->Close(); + allocator.Free(reader_buffers[i]); + } + writer->DeregisterMemory(reinterpret_cast(writer_buffer.ptr)); + writer->Close(); + allocator.Free(writer_buffer); + + server.Shutdown(); + server_thread.join(); + unlink(grpc_path.c_str()); + unlink(fd_path.c_str()); +} + TEST(StandaloneShmIpcTest, ShutdownDoesNotHangOnHalfOpenFdConnection) { const std::string address = "unix:///tmp/umbp_standalone_halfopen_" + std::to_string(getpid()) + ".grpc.sock"; diff --git a/tests/cpp/umbp/distributed/test_tier_priority_route_get.cpp b/tests/cpp/umbp/distributed/test_tier_priority_route_get.cpp deleted file mode 100644 index c910e1228..000000000 --- a/tests/cpp/umbp/distributed/test_tier_priority_route_get.cpp +++ /dev/null @@ -1,142 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include - -#include -#include -#include - -#include "umbp/distributed/routing/route_get_strategy.h" - -namespace mori::umbp { -namespace { - -Location MakeLoc(const std::string& node_id, TierType tier) { - Location loc; - loc.node_id = node_id; - loc.size = 4096; - loc.tier = tier; - return loc; -} - -// With a DRAM (or HBM) replica present alongside an SSD one, the strategy must -// never route to the slow SSD tier. -TEST(TierPriorityRouteGetStrategyTest, PrefersDramOverSsd) { - TierPriorityRouteGetStrategy strategy; - std::vector locations = { - MakeLoc("ssd-node", TierType::SSD), - MakeLoc("dram-node", TierType::DRAM), - }; - for (int i = 0; i < 100; ++i) { - auto selected = strategy.Select(locations, "requester"); - EXPECT_EQ(selected.tier, TierType::DRAM); - EXPECT_EQ(selected.node_id, "dram-node"); - } -} - -// HBM beats both DRAM and SSD. -TEST(TierPriorityRouteGetStrategyTest, PrefersHbmOverDramAndSsd) { - TierPriorityRouteGetStrategy strategy; - std::vector locations = { - MakeLoc("ssd-node", TierType::SSD), - MakeLoc("dram-node", TierType::DRAM), - MakeLoc("hbm-node", TierType::HBM), - }; - for (int i = 0; i < 100; ++i) { - auto selected = strategy.Select(locations, "requester"); - EXPECT_EQ(selected.tier, TierType::HBM); - } -} - -// When SSD is the only tier present it is selected (read-from-SSD is valid). -TEST(TierPriorityRouteGetStrategyTest, FallsBackToSsdWhenOnlyTier) { - TierPriorityRouteGetStrategy strategy; - std::vector locations = { - MakeLoc("ssd-a", TierType::SSD), - MakeLoc("ssd-b", TierType::SSD), - }; - for (int i = 0; i < 50; ++i) { - auto selected = strategy.Select(locations, "requester"); - EXPECT_EQ(selected.tier, TierType::SSD); - } -} - -// Within the winning tier, selection spreads across all replicas on that tier -// and never leaks to a lower tier. -TEST(TierPriorityRouteGetStrategyTest, RandomWithinBestTierOnly) { - TierPriorityRouteGetStrategy strategy; - std::vector locations = { - MakeLoc("dram-a", TierType::DRAM), - MakeLoc("dram-b", TierType::DRAM), - MakeLoc("dram-c", TierType::DRAM), - MakeLoc("ssd-x", TierType::SSD), - }; - std::set seen; - for (int i = 0; i < 2000; ++i) { - auto selected = strategy.Select(locations, "requester"); - ASSERT_EQ(selected.tier, TierType::DRAM) << "must never pick the SSD replica"; - seen.insert(selected.node_id); - } - EXPECT_EQ(seen.size(), 3u) << "all three DRAM replicas should be reachable"; - EXPECT_EQ(seen.count("ssd-x"), 0u); -} - -TEST(TierPriorityRouteGetStrategyTest, PrefersRequesterLocalReplicaWithinBestTier) { - TierPriorityRouteGetStrategy strategy; - std::vector locations = { - MakeLoc("dram-a", TierType::DRAM), - MakeLoc("requester", TierType::DRAM), - MakeLoc("dram-c", TierType::DRAM), - MakeLoc("ssd-x", TierType::SSD), - }; - - for (int i = 0; i < 100; ++i) { - auto selected = strategy.Select(locations, "requester"); - EXPECT_EQ(selected.tier, TierType::DRAM); - EXPECT_EQ(selected.node_id, "requester"); - } -} - -TEST(TierPriorityRouteGetStrategyTest, RequesterLocalLowerTierDoesNotBeatBestTier) { - TierPriorityRouteGetStrategy strategy; - std::vector locations = { - MakeLoc("dram-a", TierType::DRAM), - MakeLoc("requester", TierType::SSD), - }; - - for (int i = 0; i < 100; ++i) { - auto selected = strategy.Select(locations, "requester"); - EXPECT_EQ(selected.tier, TierType::DRAM); - EXPECT_EQ(selected.node_id, "dram-a"); - } -} - -TEST(TierPriorityRouteGetStrategyTest, EmptyReturnsDefault) { - TierPriorityRouteGetStrategy strategy; - std::vector locations; - auto selected = strategy.Select(locations, "requester"); - EXPECT_EQ(selected.tier, TierType::UNKNOWN); - EXPECT_TRUE(selected.node_id.empty()); -} - -} // namespace -} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_transfer_engine.cpp b/tests/cpp/umbp/distributed/test_transfer_engine.cpp new file mode 100644 index 000000000..73bbc8a6d --- /dev/null +++ b/tests/cpp/umbp/distributed/test_transfer_engine.cpp @@ -0,0 +1,401 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Phase 6 transfer layer, without gRPC or RDMA: LocalCopyEngine's planner and +// executor, and CompositeTransferEngine's fan-out registration + per-pair +// engine selection. +// +// The mori-io engine is exercised end-to-end by test_cross_node_smoke (which +// needs a real fabric); what is unit-testable here is everything that decides +// WHICH engine runs and WHAT it is asked to run. + +#include + +#include +#include +#include +#include +#include + +#include "umbp/distributed/transfer/composite_transfer_engine.h" +#include "umbp/distributed/transfer/local_copy_engine.h" +#include "umbp/distributed/transfer/transfer_engine.h" + +namespace mori::umbp { +namespace { + +TransferItem MakeItem(const TransferRef& src, uint64_t src_off, const TransferRef& dst, + uint64_t dst_off, uint64_t size, size_t tag) { + TransferItem item; + item.src = src; + item.src_offset = src_off; + item.dst = dst; + item.dst_offset = dst_off; + item.size = size; + item.tag = tag; + return item; +} + +// A ref that looks like a peer's memory: reachable only through mori-io, with +// an engine key that is not ours. +TransferRef RemoteRef(const std::string& engine_key, mori::io::MemoryUniqueId id, uint64_t size) { + mori::io::MemoryDesc desc; + desc.engineKey = engine_key; + desc.id = id; + desc.size = size; + desc.loc = mori::io::MemoryLocationType::CPU; + return TransferRef::Remote(std::move(desc)); +} + +// --------------------------------------------------------------------------- +// LocalCopyEngine +// --------------------------------------------------------------------------- + +TEST(LocalCopyEngine, CopiesEveryScatteredPage) { + LocalCopyEngine engine; + std::vector src(4096), dst(4096, 0); + std::iota(src.begin(), src.end(), 1); + + TransferRef s = + engine.RegisterMemory(src.data(), src.size(), mori::io::MemoryLocationType::CPU, -1); + TransferRef d = + engine.RegisterMemory(dst.data(), dst.size(), mori::io::MemoryLocationType::CPU, -1); + + // Deliberately out of order and non-adjacent, so nothing coalesces. + std::vector items{ + MakeItem(s, 2048, d, 0, 1024, 0), + MakeItem(s, 0, d, 2048, 1024, 1), + }; + std::vector failed; + ASSERT_TRUE(engine.Transfer(items, &failed)); + EXPECT_TRUE(failed.empty()); + EXPECT_EQ(0, std::memcmp(dst.data(), src.data() + 2048, 1024)); + EXPECT_EQ(0, std::memcmp(dst.data() + 2048, src.data(), 1024)); +} + +TEST(LocalCopyEngine, CoalescesContiguousSegmentsIntoOnePlan) { + LocalCopyEngine engine; + std::vector src(4096), dst(4096, 0); + std::iota(src.begin(), src.end(), 7); + TransferRef s = TransferRef::HostBytes(src.data(), src.size()); + TransferRef d = TransferRef::HostBytes(dst.data(), dst.size()); + + // Four adjacent pages, same buffer pair: one plan, one segment. + std::vector items; + for (size_t i = 0; i < 4; ++i) { + items.push_back(MakeItem(s, i * 1024, d, i * 1024, 1024, i)); + } + TransferPlanSet planned = engine.Plan(items); + ASSERT_EQ(1u, planned.plans.size()); + EXPECT_TRUE(planned.rejected_tags.empty()); + EXPECT_EQ(1u, planned.plans[0].sizes.size()); + EXPECT_EQ(4096u, planned.plans[0].sizes[0]); + // Every contributing key is still named, so a failure maps back to all four. + EXPECT_EQ(std::vector({0, 1, 2, 3}), planned.plans[0].tags); + + auto handle = engine.Submit(std::move(planned.plans)); + ASSERT_NE(nullptr, handle); + std::vector failures; + handle->Wait(&failures); + EXPECT_TRUE(failures.empty()); + EXPECT_EQ(src, dst); +} + +TEST(LocalCopyEngine, SeparatePlanPerBufferPair) { + LocalCopyEngine engine; + std::vector a(1024, 'a'), b(1024, 'b'), out1(1024, 0), out2(1024, 0); + TransferRef ra = TransferRef::HostBytes(a.data(), a.size()); + TransferRef rb = TransferRef::HostBytes(b.data(), b.size()); + TransferRef r1 = TransferRef::HostBytes(out1.data(), out1.size()); + TransferRef r2 = TransferRef::HostBytes(out2.data(), out2.size()); + + std::vector items{MakeItem(ra, 0, r1, 0, 1024, 0), MakeItem(rb, 0, r2, 0, 1024, 1)}; + TransferPlanSet planned = engine.Plan(items); + EXPECT_EQ(2u, planned.plans.size()); +} + +TEST(LocalCopyEngine, RejectsOutOfBoundsWithoutCopying) { + LocalCopyEngine engine; + std::vector src(1024, 'x'), dst(1024, 0); + TransferRef s = TransferRef::HostBytes(src.data(), src.size()); + TransferRef d = TransferRef::HostBytes(dst.data(), dst.size()); + + std::vector items{ + MakeItem(s, 0, d, 0, 512, 0), // fine + MakeItem(s, 512, d, 512, 1024, 1) // runs past the end of both + }; + TransferPlanSet planned = engine.Plan(items); + EXPECT_EQ(std::vector({1}), planned.rejected_tags); + ASSERT_EQ(1u, planned.plans.size()); + EXPECT_EQ(std::vector({0}), planned.plans[0].tags); + + // Transfer() surfaces the rejection as a failure without touching the bytes + // the rejected item would have written. + std::vector failed; + EXPECT_FALSE(engine.Transfer(items, &failed)); + EXPECT_EQ(std::vector({1}), failed); + for (size_t i = 512; i < dst.size(); ++i) EXPECT_EQ(0, dst[i]) << "byte " << i; +} + +TEST(LocalCopyEngine, RefusesAnEndpointItCannotAddress) { + LocalCopyEngine engine; + std::vector local(1024, 'x'); + TransferRef host = TransferRef::HostBytes(local.data(), local.size()); + TransferRef remote = RemoteRef("some-peer", /*id=*/7, 1024); + + EXPECT_TRUE(engine.CanHandle(host, host)); + EXPECT_FALSE(engine.CanHandle(host, remote)); + EXPECT_FALSE(engine.CanHandle(remote, host)); + EXPECT_FALSE(engine.CanHandle(remote, remote)); +} + +// --------------------------------------------------------------------------- +// CompositeTransferEngine +// --------------------------------------------------------------------------- + +// Records what it was asked to do; claims exactly the pairs it is told to. +class SpyEngine final : public TransferEngine { + public: + explicit SpyEngine(bool claims_remote) : claims_remote_(claims_remote) {} + + const char* Name() const override { return "SpyEngine"; } + + TransferRef RegisterMemory(void* base, size_t size, mori::io::MemoryLocationType loc, + int device) override { + ++registrations; + TransferRef ref = TransferRef::HostBytes(base, size, loc, device); + // Pretend to be a transport that also produces a mori-io registration. + ref.mem.engineKey = "spy"; + ref.mem.id = next_id_++; + ref.mem.size = size; + return ref; + } + void Deregister(const TransferRef&) override { ++deregistrations; } + + bool CanHandle(const TransferRef& src, const TransferRef& dst) const override { + const bool any_remote = (src.host_ptr == nullptr) || (dst.host_ptr == nullptr); + return claims_remote_ == any_remote; + } + + TransferPlanSet Plan(const std::vector& items) const override { + planned_batches.push_back(items.size()); + TransferPlanSet out; + TransferPlan plan; + for (const auto& item : items) { + plan.src = item.src; + plan.dst = item.dst; + plan.sizes.push_back(item.size); + plan.src_offsets.push_back(item.src_offset); + plan.dst_offsets.push_back(item.dst_offset); + plan.tags.push_back(item.tag); + } + if (!plan.sizes.empty()) out.plans.push_back(std::move(plan)); + return out; + } + + std::unique_ptr Submit(std::vector plans) override { + submitted += plans.size(); + return nullptr; // "posted nothing" — the composite must fail these tags + } + + int registrations = 0; + int deregistrations = 0; + size_t submitted = 0; + mutable std::vector planned_batches; + + private: + bool claims_remote_; + mori::io::MemoryUniqueId next_id_ = 1; +}; + +TEST(CompositeTransferEngine, RegistrationFansOutAndMergesHandles) { + auto composite = std::make_unique(); + // LocalCopyEngine contributes only a host pointer; the spy contributes a + // mori-io registration. A merged ref must carry BOTH — that is the whole + // reason TransferRef is a struct of per-transport handles and not a variant. + composite->AddEngine(std::make_unique()); + auto spy_owned = std::make_unique(/*claims_remote=*/true); + SpyEngine* spy = spy_owned.get(); + composite->AddEngine(std::move(spy_owned)); + + std::vector buf(256, 'z'); + TransferRef ref = + composite->RegisterMemory(buf.data(), buf.size(), mori::io::MemoryLocationType::CPU, -1); + EXPECT_EQ(1, spy->registrations); + EXPECT_TRUE(ref.HasHostPtr()); + EXPECT_EQ(buf.data(), ref.host_ptr); + EXPECT_TRUE(ref.HasMemoryDesc()); + EXPECT_EQ("spy", ref.mem.engineKey); + + composite->Deregister(ref); + EXPECT_EQ(1, spy->deregistrations); +} + +TEST(CompositeTransferEngine, SelectsPerPairAndPlansOncePerEngine) { + auto composite = std::make_unique(); + composite->AddEngine(std::make_unique()); + auto spy_owned = std::make_unique(/*claims_remote=*/true); + SpyEngine* spy = spy_owned.get(); + composite->AddEngine(std::move(spy_owned)); + + std::vector src(2048, 'q'), dst(2048, 0); + TransferRef s = TransferRef::HostBytes(src.data(), src.size()); + TransferRef d = TransferRef::HostBytes(dst.data(), dst.size()); + TransferRef remote = RemoteRef("peer", /*id=*/3, 2048); + + std::vector items{ + MakeItem(s, 0, d, 0, 512, 0), // local -> local: LocalCopyEngine + MakeItem(s, 512, remote, 0, 512, 1), // local -> remote: SpyEngine + MakeItem(s, 1024, d, 1024, 512, 2), // local -> local + }; + + TransferPlanSet planned = composite->Plan(items); + EXPECT_TRUE(planned.rejected_tags.empty()); + // The two local items were planned in ONE call to the local engine, not one + // call per item: partition-then-plan is what keeps grouping meaningful. + ASSERT_EQ(1u, spy->planned_batches.size()); + EXPECT_EQ(1u, spy->planned_batches[0]); + // Every plan is tagged with the engine that produced it. + ASSERT_FALSE(planned.plans.empty()); + for (const auto& plan : planned.plans) EXPECT_NE(nullptr, plan.engine); +} + +TEST(CompositeTransferEngine, RejectsAPairNoEngineClaims) { + auto composite = std::make_unique(); + composite->AddEngine(std::make_unique()); + + std::vector src(512, 'r'); + TransferRef s = TransferRef::HostBytes(src.data(), src.size()); + TransferRef remote = RemoteRef("peer", /*id=*/9, 512); + + std::vector items{MakeItem(s, 0, remote, 0, 512, 42)}; + TransferPlanSet planned = composite->Plan(items); + EXPECT_TRUE(planned.plans.empty()); + EXPECT_EQ(std::vector({42}), planned.rejected_tags); +} + +TEST(CompositeTransferEngine, FailsTagsWhenASubEngineSubmitsNothing) { + // A sub-engine that posts nothing must not silently drop keys: the caller + // would wait for a completion that never comes and report success. + auto composite = std::make_unique(); + auto spy_owned = std::make_unique(/*claims_remote=*/true); + composite->AddEngine(std::move(spy_owned)); + + std::vector src(512, 'r'); + TransferRef s = TransferRef::HostBytes(src.data(), src.size()); + TransferRef remote = RemoteRef("peer", /*id=*/9, 512); + + std::vector failed; + EXPECT_FALSE(composite->Transfer({MakeItem(s, 0, remote, 0, 512, 5)}, &failed)); + EXPECT_EQ(std::vector({5}), failed); +} + +// --------------------------------------------------------------------------- +// File endpoints (GdsEngine's shape, without HIP) +// --------------------------------------------------------------------------- + +// Stands in for GdsEngine: the only engine that registers and claims files. +class FakeFileEngine final : public TransferEngine { + public: + const char* Name() const override { return "FakeFileEngine"; } + TransferRef RegisterMemory(void*, size_t, mori::io::MemoryLocationType, int) override { + return TransferRef{}; // not a memory engine + } + TransferRef RegisterFile(int fd, uint64_t offset, uint64_t size) override { + ++file_registrations; + return TransferRef::File(fd, offset, size, reinterpret_cast(0xF11E)); + } + void Deregister(const TransferRef& ref) override { + if (ref.IsFile()) ++file_deregistrations; + } + bool CanHandle(const TransferRef& src, const TransferRef& dst) const override { + return src.IsFile() && dst.loc == mori::io::MemoryLocationType::GPU; + } + TransferPlanSet Plan(const std::vector&) const override { return {}; } + std::unique_ptr Submit(std::vector) override { return nullptr; } + + int file_registrations = 0; + int file_deregistrations = 0; +}; + +TEST(TransferRef, FileEndpointIsValidAndIgnoredByMemoryEngines) { + TransferRef f = TransferRef::File(/*fd=*/7, /*offset=*/4096, /*n=*/65536); + EXPECT_TRUE(f.IsFile()); + EXPECT_TRUE(f.Valid()); + EXPECT_FALSE(f.HasHostPtr()); + EXPECT_FALSE(f.HasMemoryDesc()); + EXPECT_EQ(7, f.file_fd); + EXPECT_EQ(4096u, f.file_offset); + EXPECT_EQ(65536u, f.size); + + // A memory engine never claims a file endpoint, so per-pair selection leaves + // it for the file engine. + LocalCopyEngine mem; + std::vector stub(64); + TransferRef gpu = + TransferRef::HostBytes(stub.data(), stub.size(), mori::io::MemoryLocationType::GPU, 0); + EXPECT_FALSE(mem.CanHandle(f, gpu)); +} + +TEST(CompositeTransferEngine, RegisterFileFansOutToTheFileEngine) { + auto composite = std::make_unique(); + composite->AddEngine(std::make_unique()); + auto file_owned = std::make_unique(); + FakeFileEngine* file = file_owned.get(); + composite->AddEngine(std::move(file_owned)); + + TransferRef ref = composite->RegisterFile(/*fd=*/9, /*offset=*/8192, /*size=*/4096); + EXPECT_EQ(1, file->file_registrations); + ASSERT_TRUE(ref.IsFile()); + EXPECT_EQ(9, ref.file_fd); + EXPECT_EQ(8192u, ref.file_offset); + EXPECT_EQ(reinterpret_cast(0xF11E), ref.gds_handle); + + composite->Deregister(ref); + EXPECT_EQ(1, file->file_deregistrations); +} + +TEST(CompositeTransferEngine, RegisterFileIsInvalidWithoutAFileEngine) { + auto composite = std::make_unique(); + composite->AddEngine(std::make_unique()); + TransferRef ref = composite->RegisterFile(/*fd=*/3, /*offset=*/0, /*size=*/4096); + EXPECT_FALSE(ref.IsFile()); + EXPECT_FALSE(ref.Valid()); +} + +TEST(CompositeTransferEngine, SelectsTheFileEngineForAFileToGpuPair) { + auto composite = std::make_unique(); + composite->AddEngine(std::make_unique()); + composite->AddEngine(std::make_unique()); + + TransferRef src = + TransferRef::File(/*fd=*/5, /*offset=*/0, /*n=*/4096, reinterpret_cast(0x1)); + std::vector gpu(4096); + TransferRef dst = + TransferRef::HostBytes(gpu.data(), gpu.size(), mori::io::MemoryLocationType::GPU, 0); + TransferEngine* sel = composite->SelectEngine(src, dst); + ASSERT_NE(nullptr, sel); + EXPECT_STREQ("FakeFileEngine", sel->Name()); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_umbp_client_metrics.cpp b/tests/cpp/umbp/distributed/test_umbp_client_metrics.cpp index d53f53a87..e01231f5a 100644 --- a/tests/cpp/umbp/distributed/test_umbp_client_metrics.cpp +++ b/tests/cpp/umbp/distributed/test_umbp_client_metrics.cpp @@ -24,16 +24,16 @@ // // 1. MasterClientMetricsTest — unit tests for MasterClient::AddCounter, // SetGauge, and Observe buffering. Uses a fake recording server that -// returns a 100ms heartbeat interval, allowing the flush to be verified -// within a short wall-clock window. +// captures ReportMetrics calls. main() sets the metrics reporting interval +// to 10ms so buffering semantics can be verified without one-second waits. // // 2. PoolClientLocalByteTrackingTest — integration test verifying that a // single-node PoolClient reports correct Put/Get byte counts (traffic=local) // through the full pipeline: PoolClient → MasterClient buffer → // ReportMetrics RPC → MasterServer → Prometheus HTTP endpoint. // -// main() sets UMBP_HEARTBEAT_TTL_SEC=1 before any test runs so that the real -// MasterServer in suite 2 returns ≈500ms heartbeat/metrics flush intervals. +// main() shortens the independently configured heartbeat and metrics intervals +// before either suite constructs a client. #include #include @@ -60,6 +60,7 @@ #include "umbp/distributed/config.h" #include "umbp/distributed/master/master_client.h" #include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/pool/policy_config.h" #include "umbp/distributed/pool_client.h" namespace mori::umbp { @@ -90,17 +91,14 @@ static uint16_t AllocPort() { // ============================================================ // Fake gRPC server that captures ReportMetrics calls. -// RegisterClient returns a configurable heartbeat interval so -// the metrics flush cadence is under test control. +// RegisterClient returns a short heartbeat interval; the independent +// metrics cadence is controlled by UMBP_METRICS_REPORT_INTERVAL_MS in main(). // ============================================================ class RecordingMasterService final : public ::umbp::UMBPMaster::Service { public: - explicit RecordingMasterService(int heartbeat_interval_ms) - : heartbeat_interval_ms_(heartbeat_interval_ms) {} - grpc::Status RegisterClient(grpc::ServerContext*, const ::umbp::RegisterClientRequest*, ::umbp::RegisterClientResponse* resp) override { - resp->set_heartbeat_interval_ms(heartbeat_interval_ms_); + resp->set_heartbeat_interval_ms(100); return grpc::Status::OK; } @@ -139,7 +137,6 @@ class RecordingMasterService final : public ::umbp::UMBPMaster::Service { } private: - int heartbeat_interval_ms_; std::mutex mu_; std::condition_variable cv_; std::vector<::umbp::ReportMetricsRequest> requests_; @@ -180,10 +177,8 @@ static double SumCounterDelta(const std::vector<::umbp::MetricSample>& samples, // ============================================================ class MasterClientMetricsTest : public ::testing::Test { protected: - static constexpr int kFlushIntervalMs = 100; - void SetUp() override { - service_ = std::make_unique(kFlushIntervalMs); + service_ = std::make_unique(); grpc::ServerBuilder builder; int selected_port = 0; @@ -372,77 +367,95 @@ constexpr size_t kLocalBufSize = 8 << 20; class PoolClientLocalByteTrackingTest : public ::testing::Test { protected: void SetUp() override { - // Kernel-assigned ephemeral ports are expected to be plentiful in CI; - // bounded retries protect against pathological duplicate picks. - constexpr int kMaxPortAllocAttempts = 32; // far above expected collision rate in CI - auto alloc_unique_port = [&](std::initializer_list used) { - for (int attempt = 0; attempt < kMaxPortAllocAttempts; ++attempt) { - const uint16_t candidate = AllocPort(); - bool duplicate = false; - for (const uint16_t port : used) { - if (candidate == port) { - duplicate = true; - break; - } - } - if (!duplicate) return candidate; - } - throw std::runtime_error("Failed to allocate unique test port"); - }; - - master_port_ = AllocPort(); - metrics_port_ = alloc_unique_port({master_port_}); - io_port_ = alloc_unique_port({master_port_, metrics_port_}); - - buf_ = std::malloc(kLocalBufSize); src_ = std::malloc(kLocalPageSize); dst_ = std::malloc(kLocalPageSize); - ASSERT_NE(buf_, nullptr); ASSERT_NE(src_, nullptr); ASSERT_NE(dst_, nullptr); - std::memset(buf_, 0, kLocalBufSize); std::memset(src_, 0xAB, kLocalPageSize); std::memset(dst_, 0, kLocalPageSize); - MasterServerConfig master_cfg; - // Short TTL so the recommended heartbeat/metrics interval is ≈100ms. - master_cfg.registry_config.heartbeat_ttl = std::chrono::seconds(1); - master_cfg.listen_address = "0.0.0.0:" + std::to_string(master_port_); - master_cfg.metrics_port = metrics_port_; - master_ = std::make_unique(std::move(master_cfg)); - server_thread_ = std::thread([this] { master_->Run(); }); - std::this_thread::sleep_for(std::chrono::milliseconds(500)); + // The master's gRPC port and the IO engine's port are requested as 0 and + // read back, so neither is reserved and released. metrics_port_ cannot be: + // MasterServerConfig takes it as a number, treats 0 as "no metrics server", + // and exposes no accessor for what it bound. Losing AllocPort's + // reserve-close-rebind race throws inside MetricsServer's constructor on + // the Run() thread, which is std::terminate rather than a test failure, so + // retry the whole server on a fresh port instead. + constexpr int kMaxAttempts = 8; + for (int attempt = 0; attempt < kMaxAttempts; ++attempt) { + metrics_port_ = AllocPort(); + + MasterServerConfig master_cfg; + // Short TTL so the recommended heartbeat/metrics interval is ≈100ms. + master_cfg.registry_config.heartbeat_ttl = std::chrono::seconds(1); + master_cfg.listen_address = "0.0.0.0:0"; + master_cfg.metrics_port = metrics_port_; + master_ = std::make_unique(std::move(master_cfg)); + run_failed_.store(false); + server_thread_ = std::thread([this] { + try { + master_->Run(); + } catch (const std::exception&) { + run_failed_.store(true); + } + }); + + // Wait for the assigned port rather than sleeping a fixed interval and + // hoping: it has to be read back anyway, and it is only non-zero once the + // server is actually listening. + bool ready = false; + for (int i = 0; i < 500; ++i) { + if (master_->GetBoundPort() != 0) { + ready = true; + break; + } + if (run_failed_.load()) break; + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + if (ready && !run_failed_.load()) break; + + master_->Shutdown(); + if (server_thread_.joinable()) server_thread_.join(); + master_.reset(); + } + ASSERT_NE(master_, nullptr); + ASSERT_FALSE(run_failed_.load()) << "MasterServer::Run() kept failing to bind a metrics port"; + ASSERT_NE(master_->GetBoundPort(), 0); + master_port_ = static_cast(master_->GetBoundPort()); PoolClientConfig cfg; cfg.master_config.node_id = "node-local"; cfg.master_config.node_address = "127.0.0.1"; cfg.master_config.master_address = "localhost:" + std::to_string(master_port_); cfg.io_engine.host = "0.0.0.0"; - cfg.io_engine.port = io_port_; + cfg.io_engine.port = 0; cfg.dram_page_size = kLocalPageSize; - cfg.dram_buffers = {{buf_, kLocalBufSize}}; - cfg.tier_capacities = {{TierType::DRAM, {kLocalBufSize, kLocalBufSize}}}; + ConfigureClient(&cfg); client_ = std::make_unique(std::move(cfg)); ASSERT_TRUE(client_->Init()); } + // Lets a fixture put a richer topology behind the same master bring-up. The + // base case is the one flat DRAM backend the byte-tracking tests expect. + virtual void ConfigureClient(PoolClientConfig* config) { + config->dram.buffer_sizes = {kLocalBufSize}; + } + void TearDown() override { if (client_) client_->Shutdown(); if (master_) master_->Shutdown(); if (server_thread_.joinable()) server_thread_.join(); - std::free(buf_); std::free(src_); std::free(dst_); } uint16_t master_port_ = 0; uint16_t metrics_port_ = 0; - uint16_t io_port_ = 0; - void* buf_ = nullptr; void* src_ = nullptr; void* dst_ = nullptr; std::unique_ptr master_; std::thread server_thread_; + std::atomic run_failed_{false}; std::unique_ptr client_; }; @@ -524,6 +537,148 @@ TEST_F(PoolClientLocalByteTrackingTest, LocalPutGetBytesCounted) { << "Unexpected remote batch-get bandwidth series in single-node setup"; } +// The ranged siblings get their OWN histograms, and the reason they are worth a +// test is that they are the only bandwidth series with early returns between +// the first byte and the observation: a batch served entirely from the local +// medium returns before the remote phase exists. If ScopedBatchBandwidth ever +// stops covering that path, the tree-connector panels silently render empty +// while ranged I/O runs at full rate. +TEST_F(PoolClientLocalByteTrackingTest, LocalRangedPutGetBandwidthCounted) { + const std::string key = "ranged-metric-key"; + constexpr size_t kHalf = kLocalPageSize / 2; + + // A put must TILE its object; a get may take any subset. + const std::vector> put_srcs = { + {src_, static_cast(src_) + kHalf}}; + const std::vector> put_sizes = {{kHalf, kHalf}}; + const std::vector> put_offsets = {{0, kHalf}}; + ASSERT_TRUE( + client_->BatchPutRanges({key}, {kLocalPageSize}, put_srcs, put_sizes, put_offsets).front()); + + const std::vector> get_dsts = {{dst_}}; + const std::vector> get_sizes = {{kHalf}}; + const std::vector> get_offsets = {{kHalf}}; + ASSERT_TRUE(client_->BatchGetRanges({key}, get_dsts, get_sizes, get_offsets).front()); + EXPECT_EQ(std::memcmp(dst_, static_cast(src_) + kHalf, kHalf), 0) + << "Ranged get did not return the bytes the ranged put committed"; + + std::this_thread::sleep_for(std::chrono::milliseconds(1500)); + + const std::string body = FetchPrometheusMetrics(metrics_port_); + ASSERT_FALSE(body.empty()) << "Could not fetch Prometheus metrics from port " << metrics_port_; + + auto hist = [&](const std::string& name, const std::string& suffix, const std::string& traffic) { + return ParseMetricValue(body, name + suffix, "traffic=\"" + traffic + "\""); + }; + + for (const char* name : {"mori_umbp_client_batch_put_ranges_bandwidth_gibps", + "mori_umbp_client_batch_get_ranges_bandwidth_gibps"}) { + EXPECT_GT(hist(name, "_sum", "local"), 0.0) + << "Expected local " << name << " sum > 0; Prometheus shows " + << hist(name, "_sum", "local"); + EXPECT_GT(hist(name, "_count", "local"), 0.0) + << "Expected local " << name << " count > 0; Prometheus shows " + << hist(name, "_count", "local"); + EXPECT_EQ(hist(name, "_sum", "remote"), -1.0) + << "Unexpected remote " << name << " series in single-node setup"; + } + + // Distinct series from the whole-object families: a ranged call moves a + // subset, so folding them together would make bytes-per-call meaningless. + EXPECT_EQ(hist("mori_umbp_client_batch_get_bandwidth_gibps", "_count", "local"), -1.0) + << "Ranged get leaked into the whole-object BatchGet bandwidth histogram"; +} + +// End-to-end proof that the refactored data plane is still wired to Prometheus: +// a Put and a Get on the live path must show up as the GENERIC backend and +// transfer series, carrying the medium and the engine as labels. +// +// This is the test that fails if someone unwires the instrumentation decorator +// in PoolClient::Init, or reintroduces medium-specific metric names — either of +// which leaves the dashboard rendering an empty panel while the system runs +// perfectly, which is exactly how the SSD panels ended up dead. +TEST_F(PoolClientLocalByteTrackingTest, BackendAndTransferSeriesReachPrometheus) { + ASSERT_TRUE(client_->Put("generic-metrics-key", src_, kLocalPageSize)); + ASSERT_TRUE(client_->Get("generic-metrics-key", dst_, kLocalPageSize)); + + std::this_thread::sleep_for(std::chrono::milliseconds(1500)); + + const std::string body = FetchPrometheusMetrics(metrics_port_); + ASSERT_FALSE(body.empty()) << "Could not fetch Prometheus metrics from port " << metrics_port_; + + // ParseMetricValue matches ONE substring, and label order in the rendered + // line follows insertion order (node, tags, source labels, sample labels), so + // match every label independently rather than betting on that order. + auto value_with = [&body](const std::string& name, + const std::vector& label_substrs) -> double { + size_t pos = 0; + while (pos < body.size()) { + const size_t nl = body.find('\n', pos); + const std::string line = + body.substr(pos, nl == std::string::npos ? std::string::npos : nl - pos); + pos = (nl == std::string::npos) ? body.size() : nl + 1; + if (line.empty() || line.front() == '#') continue; + if (line.rfind(name, 0) != 0) continue; // metric name starts the line + bool all = true; + for (const auto& sub : label_substrs) { + if (line.find(sub) == std::string::npos) { + all = false; + break; + } + } + if (!all) continue; + const size_t sp = line.rfind(' '); + if (sp == std::string::npos) continue; + try { + return std::stod(line.substr(sp + 1)); + } catch (...) { + } + } + return -1.0; + }; + + // The medium and named backend instance are labels on shared metric names. + // Legacy one-backend configuration lowers to the stable instance name + // "dram"; explicit multi-backend policies use their configured names. + const std::vector dram = {"tier=\"DRAM\"", "backend=\"dram\""}; + + auto with = [&dram](std::vector extra) { + std::vector out = dram; + out.insert(out.end(), extra.begin(), extra.end()); + return out; + }; + + EXPECT_GE(value_with("mori_umbp_backend_ops_total", with({"op=\"allocate\"", "status=\"ok\""})), + 1.0) + << "no allocate recorded for the live medium"; + EXPECT_GE(value_with("mori_umbp_backend_ops_total", with({"op=\"commit\"", "status=\"ok\""})), + 1.0); + EXPECT_GE(value_with("mori_umbp_backend_ops_total", with({"op=\"resolve\"", "status=\"ok\""})), + 1.0); + EXPECT_GE(value_with("mori_umbp_backend_bytes_total", with({"op=\"commit\""})), + static_cast(kLocalPageSize)); + EXPECT_GE(value_with("mori_umbp_backend_bytes_total", with({"op=\"resolve\""})), + static_cast(kLocalPageSize)); + EXPECT_GE(value_with("mori_umbp_backend_batches_total", with({"op=\"commit\""})), 1.0); + EXPECT_GT(value_with("mori_umbp_backend_op_seconds_total", with({"op=\"commit\""})), 0.0); + + // The transfer layer, charged to whichever engine carried the plans. Both + // endpoints are on this node, so the local engine took them. + EXPECT_GE(value_with("mori_umbp_transfer_bytes_total", + {"engine=\"LocalCopyEngine\"", "direction=\"local\""}), + static_cast(kLocalPageSize)); + EXPECT_GE(value_with("mori_umbp_transfer_ops_total", + {"engine=\"LocalCopyEngine\"", "direction=\"local\"", "status=\"ok\""}), + 1.0); + + // Nothing may be published under a name that spells the medium: such a metric + // needs its own panel, and one panel per medium is the arrangement the single + // dashboard replaced. + EXPECT_EQ(body.find("mori_umbp_dram_"), std::string::npos); + EXPECT_EQ(body.find("mori_umbp_ssd_"), std::string::npos); + EXPECT_EQ(body.find("mori_umbp_hbm_"), std::string::npos); +} + // Verifies that the four per-(node,tier) capacity gauges (total / available / // used / utilization_ratio) reach Prometheus via the heartbeat path, and that // available/used/utilization react to an allocation. @@ -583,14 +738,108 @@ TEST_F(PoolClientLocalByteTrackingTest, TierCapacityGaugesPublished) { EXPECT_NEAR(util1, used1 / total1, 1e-3) << "utilization must equal used / total"; } +// Two DRAM backends in different logical tiers, which is the case the +// per-medium gauges cannot describe: both are tier="DRAM", so they land on one +// series and the hierarchy is unreadable. +class PoolClientLogicalTierMetricsTest : public PoolClientLocalByteTrackingTest { + protected: + static constexpr uint64_t kHotBytes = 1ULL << 20; + static constexpr uint64_t kColdBytes = 4ULL << 20; + + void ConfigureClient(PoolClientConfig* config) override { + static constexpr char kPolicy[] = R"json({ + "schema_version": 1, + "entry_tier": "hot", + "backends": { + "dram_hot": {"type": "dram", "capacity": "1MiB"}, + "dram_cold": {"type": "dram", "capacity": "4MiB"} + }, + "tiers": [ + { + "name": "hot", + "backends": {"dram_hot": 100}, + "offload_to": ["cold"], + "offload_trigger": "watermark", + "high_watermark": 0.9, + "low_watermark": 0.7 + }, + {"name": "cold", "backends": {"dram_cold": 100}, "promote_trigger": "on_read"} + ] + })json"; + auto loaded = LoadBackendPolicyJson(kPolicy); + ASSERT_TRUE(loaded.ok()) << loaded.error; + std::string error; + ASSERT_TRUE(ApplyBackendPolicy(*loaded.config, config, &error)) << error; + } +}; + +// Each logical tier has to reach Prometheus as its own series, carrying the +// name the policy gave it. Without this a two-DRAM-tier policy publishes one +// tier="DRAM" series holding the sum, which cannot answer where a key sits, or +// whether a drain reached its low watermark, or which tiers take new PUTs. +TEST_F(PoolClientLogicalTierMetricsTest, PublishesCapacityPerLogicalTier) { + std::this_thread::sleep_for(std::chrono::milliseconds(1500)); + const std::string body = FetchPrometheusMetrics(metrics_port_); + ASSERT_FALSE(body.empty()) << "Could not fetch Prometheus metrics from port " << metrics_port_; + + const std::string hot = "logical_tier=\"hot\""; + const std::string cold = "logical_tier=\"cold\""; + const double hot_total = + ParseMetricValue(body, "mori_umbp_client_logical_tier_capacity_total_bytes", hot); + const double cold_total = + ParseMetricValue(body, "mori_umbp_client_logical_tier_capacity_total_bytes", cold); + ASSERT_GE(hot_total, 0.0) << "hot tier missing from /metrics"; + ASSERT_GE(cold_total, 0.0) << "cold tier missing from /metrics"; + + // Cold is its own backends' sum. The entry tier is not: its capacity covers + // everything a PUT addressed to it can spill into, because that is what the + // router needs. Pinning both shapes keeps the difference from being read as + // one tier's occupancy. + EXPECT_NEAR(cold_total, static_cast(kColdBytes), 16.0); + EXPECT_NEAR(hot_total, static_cast(kHotBytes + kColdBytes), 16.0) + << "entry tier capacity should aggregate every offload-reachable tier"; + + for (const std::string& tier : {hot, cold}) { + const double avail = + ParseMetricValue(body, "mori_umbp_client_logical_tier_capacity_available_bytes", tier); + const double used = + ParseMetricValue(body, "mori_umbp_client_logical_tier_capacity_used_bytes", tier); + const double total = + ParseMetricValue(body, "mori_umbp_client_logical_tier_capacity_total_bytes", tier); + const double util = + ParseMetricValue(body, "mori_umbp_client_logical_tier_capacity_utilization_ratio", tier); + const double eligible = + ParseMetricValue(body, "mori_umbp_client_logical_tier_put_eligible", tier); + const double peak = + ParseMetricValue(body, "mori_umbp_client_logical_tier_peak_member_utilization_ratio", tier); + ASSERT_GE(avail, 0.0) << tier << " available missing"; + ASSERT_GE(used, 0.0) << tier << " used missing"; + ASSERT_GE(util, 0.0) << tier << " utilization missing"; + ASSERT_GE(eligible, 0.0) << tier << " put_eligible missing"; + // The watermark comparison reads this one, so its absence is what made a + // drain unobservable even with the aggregate gauges present. + ASSERT_GE(peak, 0.0) << tier << " peak_member_utilization missing"; + EXPECT_LE(peak, 1.0); + EXPECT_NEAR(used + avail, total, 16.0); + EXPECT_NEAR(util, used / total, 1e-3); + } + + // Only the entry tier takes new PUTs; cold is reachable by offload alone. + EXPECT_EQ(ParseMetricValue(body, "mori_umbp_client_logical_tier_put_eligible", hot), 1.0); + EXPECT_EQ(ParseMetricValue(body, "mori_umbp_client_logical_tier_put_eligible", cold), 0.0); +} + } // namespace } // namespace mori::umbp int main(int argc, char** argv) { // Must be set before any GetEnvSeconds("UMBP_HEARTBEAT_TTL_SEC") static - // is initialized. TTL=1s → recommended interval ≈ 500ms, which the - // metrics flush thread reuses. + // is initialized. TTL=1s keeps the real-server heartbeat tests short. ::setenv("UMBP_HEARTBEAT_TTL_SEC", "1", 1); + // MetricsReportIntervalMs() caches this on the first MasterClient + // construction. A short interval keeps the mock-server buffering tests + // deterministic without sleeping for the production default of one second. + ::setenv("UMBP_METRICS_REPORT_INTERVAL_MS", "10", 1); ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); } diff --git a/tests/cpp/umbp/distributed/test_umbp_tags.cpp b/tests/cpp/umbp/distributed/test_umbp_tags.cpp index 3a0084c32..8e58d54b8 100644 --- a/tests/cpp/umbp/distributed/test_umbp_tags.cpp +++ b/tests/cpp/umbp/distributed/test_umbp_tags.cpp @@ -47,6 +47,7 @@ #include #include #include +#include #include #include #include @@ -65,20 +66,30 @@ namespace { // Helpers // --------------------------------------------------------------------------- +// Only for ports that must be chosen before the thing that binds them exists +// (MasterServerConfig::metrics_port). Prefer letting the server bind :0 and +// report back where an accessor exists — this reserve-and-release sequence has +// a window in which the port can be taken. +// +// Failure throws rather than returning a fixed fallback: a hardcoded port is +// what every caller then collides on, which is the failure this helper exists +// to avoid. static uint16_t AllocPort() { - // Bind to :0 and let the kernel pick a free port, then close immediately. int fd = ::socket(AF_INET, SOCK_STREAM, 0); - if (fd < 0) return 50400; + if (fd < 0) throw std::runtime_error("AllocPort socket() failed"); struct sockaddr_in addr{}; addr.sin_family = AF_INET; addr.sin_port = 0; addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); if (::bind(fd, reinterpret_cast(&addr), sizeof(addr)) != 0) { ::close(fd); - return 50400; + throw std::runtime_error("AllocPort bind() failed"); } socklen_t len = sizeof(addr); - ::getsockname(fd, reinterpret_cast(&addr), &len); + if (::getsockname(fd, reinterpret_cast(&addr), &len) != 0) { + ::close(fd); + throw std::runtime_error("AllocPort getsockname() failed"); + } ::close(fd); return ntohs(addr.sin_port); } @@ -246,15 +257,20 @@ class CapturingMasterService final : public ::umbp::UMBPMaster::Service { class MasterClientTagsE2ETest : public ::testing::Test { protected: void SetUp() override { - port_ = AllocPort(); - addr_ = "127.0.0.1:" + std::to_string(port_); - svc_ = std::make_unique(); + + // Bind the port through the builder rather than reserving one with + // AllocPort first: the reserve-close-rebind sequence leaves a window in + // which anything else on the host can take the port back. + int selected_port = 0; grpc::ServerBuilder builder; - builder.AddListeningPort(addr_, grpc::InsecureServerCredentials()); + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &selected_port); builder.RegisterService(svc_.get()); server_ = builder.BuildAndStart(); ASSERT_NE(server_, nullptr); + ASSERT_NE(selected_port, 0); + port_ = static_cast(selected_port); + addr_ = "127.0.0.1:" + std::to_string(port_); } void TearDown() override { @@ -334,18 +350,45 @@ class RealMasterServerTagsTest : public ::testing::Test { setenv("UMBP_HEARTBEAT_TTL_SEC", "2", 1); setenv("UMBP_METRICS_REPORT_INTERVAL_MS", "50", 1); - port_ = AllocPort(); - metrics_port_ = AllocPort(); + // The gRPC port is bound as :0 and read back, so it is never reserved and + // released. metrics_port_ cannot be: MasterServerConfig takes it as a + // number, treats 0 as "no metrics server", and exposes no accessor for + // what it bound — so AllocPort's reserve-close-rebind window is + // unavoidable there. Losing that race throws inside MetricsServer's + // constructor on the Run() thread, which is std::terminate rather than a + // test failure, so retry the whole server on a fresh port instead. + constexpr int kMaxAttempts = 8; + for (int attempt = 0; attempt < kMaxAttempts; ++attempt) { + metrics_port_ = AllocPort(); + + MasterServerConfig cfg = MasterServerConfig::FromEnvironment(); + cfg.listen_address = "127.0.0.1:0"; + cfg.metrics_port = metrics_port_; + server_ = std::make_unique(std::move(cfg)); + run_failed_.store(false); + server_thread_ = std::thread([this] { + try { + server_->Run(); + } catch (const std::exception&) { + run_failed_.store(true); + } + }); + + if (WaitFor([this] { return server_->GetBoundPort() != 0 || run_failed_.load(); }, + std::chrono::seconds(5)) && + !run_failed_.load()) { + break; + } + + server_->Shutdown(); + if (server_thread_.joinable()) server_thread_.join(); + server_.reset(); + } + ASSERT_NE(server_, nullptr); + ASSERT_FALSE(run_failed_.load()) << "MasterServer::Run() kept failing to bind a metrics port"; + ASSERT_NE(server_->GetBoundPort(), 0); + port_ = static_cast(server_->GetBoundPort()); addr_ = "127.0.0.1:" + std::to_string(port_); - - MasterServerConfig cfg = MasterServerConfig::FromEnvironment(); - cfg.listen_address = addr_; - cfg.metrics_port = metrics_port_; - server_ = std::make_unique(std::move(cfg)); - server_thread_ = std::thread([this] { server_->Run(); }); - - // Wait for the gRPC port to be bound - ASSERT_TRUE(WaitFor([this] { return server_->GetBoundPort() != 0; }, std::chrono::seconds(5))); } void TearDown() override { @@ -395,6 +438,7 @@ class RealMasterServerTagsTest : public ::testing::Test { std::string addr_; std::unique_ptr server_; std::thread server_thread_; + std::atomic run_failed_{false}; std::unique_ptr client_; }; diff --git a/tests/cpp/umbp/distributed/test_umbp_tier_benchmark.cpp b/tests/cpp/umbp/distributed/test_umbp_tier_benchmark.cpp new file mode 100644 index 000000000..5f1794219 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_umbp_tier_benchmark.cpp @@ -0,0 +1,489 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/payload.h" +#include "umbp/distributed/benchmark/pool_workload_client.h" +#include "umbp/distributed/benchmark/workload_runner.h" +#include "umbp/distributed/config.h" +#include "umbp/distributed/master/master_server.h" +#include "umbp/distributed/peer/backend/medium_backend.h" +#include "umbp/distributed/pool_client.h" +#include "umbp/distributed/routing/route_put_strategy.h" + +namespace mori::umbp { +namespace { + +namespace bench = benchmark; +using Event = ::umbp::benchmark::WorkloadEvent; + +constexpr size_t kPageSize = 4096; +constexpr size_t kKeysPerClient = 12; +constexpr size_t kBackendCapacity = kPageSize * 64; + +uint64_t AdvertisedDramTotal(MasterServer* master, const std::string& node_id) { + if (master == nullptr) return 0; + const auto record = master->GetClient(node_id); + if (!record.has_value()) return 0; + const auto found = record->tier_capacities.find(TierType::DRAM); + return found == record->tier_capacities.end() ? 0 : found->second.total_bytes; +} + +std::string AdvertisedDramDebug(MasterServer* master, const std::string& node_id) { + if (master == nullptr) return "master=null"; + const auto record = master->GetClient(node_id); + if (!record.has_value()) return "no ClientRecord"; + const auto found = record->tier_capacities.find(TierType::DRAM); + if (found == record->tier_capacities.end()) return "no DRAM tier"; + return "total=" + std::to_string(found->second.total_bytes) + + " available=" + std::to_string(found->second.available_bytes) + + " max_alloc=" + std::to_string(found->second.max_allocatable_bytes); +} + +bool WaitForAdvertisedDram(MasterServer* master, const std::string& node_id, + uint64_t expected_total, + std::optional expected_available = std::nullopt) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + do { + const auto record = master == nullptr ? std::nullopt : master->GetClient(node_id); + if (record.has_value()) { + const auto found = record->tier_capacities.find(TierType::DRAM); + if (found != record->tier_capacities.end() && found->second.total_bytes == expected_total && + (!expected_available.has_value() || + found->second.available_bytes == *expected_available)) { + return true; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } while (std::chrono::steady_clock::now() < deadline); + ADD_FAILURE() << node_id << " " << AdvertisedDramDebug(master, node_id) + << " expected_total=" << expected_total + << (expected_available.has_value() + ? " expected_available=" + std::to_string(*expected_available) + : ""); + return false; +} + +bool WaitForAdvertisedDramTotal(MasterServer* master, const std::string& node_id, + uint64_t expected) { + return WaitForAdvertisedDram(master, node_id, expected); +} + +class VectorSource final : public bench::WorkloadSource { + public: + explicit VectorSource(std::vector events, uint64_t seed) + : events_(std::move(events)), seed_(seed) {} + + uint64_t seed() const override { return seed_; } + + bool Next(Event* event) override { + if (next_ == events_.size()) return false; + *event = events_[next_++]; + return true; + } + + private: + std::vector events_; + uint64_t seed_; + size_t next_ = 0; +}; + +Event MakeEvent(uint32_t client_id, uint64_t operation_id, Event::Operation operation, + std::string key, uint64_t batch_id) { + Event event; + event.set_client_id(client_id); + event.set_operation_id(operation_id); + event.set_operation(operation); + event.set_key(std::move(key)); + event.set_value_size(kPageSize); + event.set_batch_id(batch_id); + return event; +} + +std::vector MakeReplay() { + std::vector events; + events.reserve(2 * 2 * kKeysPerClient); + for (uint32_t client_id = 0; client_id < 2; ++client_id) { + for (size_t key_index = 0; key_index < kKeysPerClient; ++key_index) { + const uint64_t operation_id = client_id * kKeysPerClient + key_index + 1; + events.push_back( + MakeEvent(client_id, operation_id, Event::PUT, + "tier-smoke-" + std::to_string(client_id) + "-" + std::to_string(key_index), + 100 + client_id)); + } + for (size_t key_index = 0; key_index < kKeysPerClient; ++key_index) { + const uint64_t operation_id = client_id * kKeysPerClient + key_index + 1; + events.push_back( + MakeEvent(client_id, operation_id, Event::GET, + "tier-smoke-" + std::to_string(client_id) + "-" + std::to_string(key_index), + 200 + client_id)); + } + } + return events; +} + +// Every case here needs the same thing before it can say anything about +// placement: one in-process master on an ephemeral port, and PoolClients that +// really register with it. Only the routing affinity, the capacity contract, and +// the backend shapes differ, so those are the parameters. +class InProcessClusterTest : public ::testing::Test { + protected: + using Affinity = ConfigurableRoutePutStrategy::NodeAffinity; + + void StartMaster(Affinity affinity = Affinity::kNone, + bool advertise_max_allocatable_bytes = false) { + MasterServerConfig config; + config.listen_address = "127.0.0.1:0"; + config.metrics_port = 0; + config.registry_config.default_dram_page_size = kPageSize; + config.registry_config.heartbeat_ttl = std::chrono::seconds(1); + config.registry_config.advertise_max_allocatable_bytes = advertise_max_allocatable_bytes; + config.put_strategy = std::make_unique( + ConfigurableRoutePutStrategy::SelectAlgo::kMostAvailable, affinity, 17); + master_ = std::make_unique(std::move(config)); + master_thread_ = std::thread([this] { master_->Run(); }); + + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (master_->GetBoundPort() == 0 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_NE(master_->GetBoundPort(), 0) << "in-process master failed to start"; + } + + // Weighted placement across named DRAM backends, which is the configuration + // under test in every case; callers supply the backends. + void StartClient(const std::string& node_id, const std::vector& backends) { + PoolClientConfig config; + config.master_config.master_address = "127.0.0.1:" + std::to_string(master_->GetBoundPort()); + config.master_config.node_id = node_id; + config.master_config.node_address = "127.0.0.1"; + config.master_config.auto_heartbeat = true; + config.io_engine.host = "127.0.0.1"; + config.io_engine.port = 0; + config.auto_peer_service_port = true; + config.dram_page_size = kPageSize; + config.cache_remote_fetches = false; + config.placement_policy = PoolPlacementPolicy::WEIGHTED; + config.backends = backends; + + auto client = std::make_unique(std::move(config)); + ASSERT_TRUE(client->Init()) << "failed to initialize PoolClient " << node_id; + clients_.push_back(std::move(client)); + } + + static BackendInstanceConfig DramBackend(const std::string& name, uint64_t capacity, + uint32_t weight) { + BackendInstanceConfig backend; + backend.name = name; + backend.tier = TierType::DRAM; + backend.dram.buffer_sizes = {capacity}; + backend.placement_weight = weight; + return backend; + } + + void Settle(std::chrono::milliseconds wait) { + for (const auto& client : clients_) client->Master().FlushHeartbeat(); + std::this_thread::sleep_for(wait); + } + + PoolClient* client() const { return clients_.empty() ? nullptr : clients_.front().get(); } + + void TearDown() override { + for (auto& client : clients_) { + if (client) client->Clear(); + } + for (auto it = clients_.rbegin(); it != clients_.rend(); ++it) { + if (*it) (*it)->Shutdown(); + } + clients_.clear(); + if (master_) master_->Shutdown(); + if (master_thread_.joinable()) master_thread_.join(); + master_.reset(); + } + + std::unique_ptr master_; + std::thread master_thread_; + std::vector> clients_; +}; + +class TierBenchmarkSmokeTest : public InProcessClusterTest { + protected: + void SetUp() override { + ASSERT_NO_FATAL_FAILURE(StartMaster()); + for (uint32_t client_id = 0; client_id < 2; ++client_id) { + ASSERT_NO_FATAL_FAILURE(StartClient("tier-smoke-node-" + std::to_string(client_id), + {DramBackend("dram-0", kBackendCapacity, 1), + DramBackend("dram-1", kBackendCapacity, 2)})); + } + ASSERT_EQ(clients_.size(), 2u); + Settle(std::chrono::milliseconds(100)); + } +}; + +TEST_F(TierBenchmarkSmokeTest, ReplaysValidatedBatchesAcrossWeightedDramBackends) { + constexpr uint64_t kPayloadSeed = 0x5eed; + VectorSource source(MakeReplay(), kPayloadSeed); + // Client B may read a key client A only just wrote, so the reader waits out + // the peer-event to master-index publication window instead of missing. + bench::PoolWorkloadClient adapter( + &clients_, {std::chrono::seconds(5), std::chrono::seconds(5), std::chrono::milliseconds(25)}); + bench::WorkloadRunnerOptions options; + options.max_throughput = true; + options.validate_get_payloads = true; + bench::WorkloadRunner runner(&adapter, options); + + const bench::WorkloadMetrics metrics = runner.Run(&source); + + EXPECT_EQ(metrics.total.attempted, 4 * kKeysPerClient); + EXPECT_EQ(metrics.total.succeeded, metrics.total.attempted); + EXPECT_EQ(metrics.total.failed, 0u); + EXPECT_EQ(metrics.puts.succeeded, 2 * kKeysPerClient); + EXPECT_EQ(metrics.gets.succeeded, 2 * kKeysPerClient); + EXPECT_EQ(metrics.get_misses, 0u); + EXPECT_EQ(metrics.get_validation_failures, 0u); + for (uint32_t client_id = 0; client_id < 2; ++client_id) { + EXPECT_EQ(adapter.BatchPutCalls(client_id), 1u); + EXPECT_EQ(adapter.BatchGetCalls(client_id), 1u); + } + + size_t populated_backends = 0; + size_t owned_keys = 0; + for (const auto& client : clients_) { + ASSERT_EQ(client->Backends().All().size(), 2u); + for (MediumBackend* backend : client->Backends().All()) { + EXPECT_EQ(backend->Tier(), TierType::DRAM); + const size_t backend_keys = backend->OwnedKeyCount(); + owned_keys += backend_keys; + if (backend_keys != 0) ++populated_backends; + } + } + EXPECT_EQ(owned_keys, 2 * kKeysPerClient); + EXPECT_GE(populated_backends, 2u) + << "the deterministic 24-key sample should exercise weighted placement"; + + for (auto& client : clients_) ASSERT_TRUE(client->Clear()); + for (const auto& client : clients_) { + for (MediumBackend* backend : client->Backends().All()) { + EXPECT_EQ(backend->OwnedKeyCount(), 0u); + } + } +} + +TEST_F(TierBenchmarkSmokeTest, ClientBReadsClientAKeysWithPayloadValidation) { + constexpr size_t kRemoteKeys = 8; + constexpr uint64_t kPayloadSeed = 0xA11CE; + PoolClient* writer = clients_[0].get(); + PoolClient* reader = clients_[1].get(); + ASSERT_NE(writer, nullptr); + ASSERT_NE(reader, nullptr); + + std::vector keys; + keys.reserve(kRemoteKeys); + for (size_t i = 0; i < kRemoteKeys; ++i) { + const std::string key = "tier-remote-" + std::to_string(i); + auto payload = bench::GenerateDeterministicPayload(key, i + 1, kPayloadSeed, kPageSize); + ASSERT_TRUE(writer->Put(key, payload.data(), payload.size())) << key; + keys.push_back(key); + } + writer->Master().FlushHeartbeat(); + + for (size_t i = 0; i < kRemoteKeys; ++i) { + std::vector got(kPageSize, 0); + bool ok = false; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + do { + if (reader->Get(keys[i], got.data(), got.size())) { + ok = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } while (std::chrono::steady_clock::now() < deadline); + ASSERT_TRUE(ok) << "client B missed client A's key " << keys[i]; + EXPECT_TRUE( + bench::ValidateDeterministicPayload(keys[i], i + 1, kPayloadSeed, got.data(), got.size())); + } +} + +// The weight of 100 to 1 makes dram-0 the primary for practically every key, so +// a fallback onto dram-1 is a decision the pool made and not a coin flip. +class TinyWeightedPoolTest : public InProcessClusterTest { + protected: + static constexpr size_t kTinyPages = 4; + static constexpr size_t kTinyCapacity = kPageSize * kTinyPages; + + void SetUp() override { + ASSERT_NO_FATAL_FAILURE(StartMaster(Affinity::kLocal)); + ASSERT_NO_FATAL_FAILURE(StartClient( + "tier-tiny-node-0", + {DramBackend("dram-0", kTinyCapacity, 100), DramBackend("dram-1", kTinyCapacity, 1)})); + Settle(std::chrono::milliseconds(100)); + } +}; + +TEST_F(TinyWeightedPoolTest, FallsBackThenFailsWithoutLeakingSlots) { + MediumBackend* primary = client()->Backends().Get("dram-0"); + MediumBackend* fallback = client()->Backends().Get("dram-1"); + ASSERT_NE(primary, nullptr); + ASSERT_NE(fallback, nullptr); + + std::vector value(kPageSize, 0x5A); + size_t succeeded = 0; + size_t failed = 0; + constexpr size_t kAttempts = kTinyPages * 2 + 4; + for (size_t i = 0; i < kAttempts; ++i) { + const std::string key = "tiny-" + std::to_string(i); + if (client()->Put(key, value.data(), value.size())) { + ++succeeded; + } else { + ++failed; + } + } + + EXPECT_EQ(succeeded, 2 * kTinyPages); + EXPECT_EQ(failed, 4u); + EXPECT_EQ(primary->OwnedKeyCount(), kTinyPages) << "weight 100:1 should fill dram-0 first"; + EXPECT_EQ(fallback->OwnedKeyCount(), kTinyPages) + << "no-space on the primary must fall back to dram-1"; + EXPECT_EQ(primary->OwnedKeyCount() + fallback->OwnedKeyCount(), succeeded); + + const auto primary_cap = primary->Capacity(); + const auto fallback_cap = fallback->Capacity(); + EXPECT_EQ(primary_cap.available_bytes, 0u); + EXPECT_EQ(fallback_cap.available_bytes, 0u); + EXPECT_EQ(primary_cap.max_allocatable_bytes, 0u); + EXPECT_EQ(fallback_cap.max_allocatable_bytes, 0u); + EXPECT_EQ(primary_cap.total_bytes, kTinyCapacity); + EXPECT_EQ(fallback_cap.total_bytes, kTinyCapacity); + EXPECT_EQ(primary_cap.available_bytes + primary->OwnedKeyCount() * kPageSize, + primary_cap.total_bytes); + EXPECT_EQ(fallback_cap.available_bytes + fallback->OwnedKeyCount() * kPageSize, + fallback_cap.total_bytes); + + ASSERT_TRUE(client()->Clear()); + EXPECT_EQ(primary->OwnedKeyCount(), 0u); + EXPECT_EQ(fallback->OwnedKeyCount(), 0u); + EXPECT_EQ(primary->Capacity().available_bytes, kTinyCapacity); + EXPECT_EQ(fallback->Capacity().available_bytes, kTinyCapacity); +} + +// Two DRAM backends of very different sizes: the aggregate is three times the +// largest single allocation the pool can serve, which is exactly where a master +// that only understands totals overcommits the node. +class UnequalDramWeightedPoolTest : public InProcessClusterTest { + protected: + static constexpr size_t kPrimaryPages = 4; + static constexpr size_t kSecondaryPages = 12; + static constexpr size_t kPrimaryCapacity = kPageSize * kPrimaryPages; + static constexpr size_t kSecondaryCapacity = kPageSize * kSecondaryPages; + static constexpr size_t kPooledCapacity = kPrimaryCapacity + kSecondaryCapacity; + static constexpr const char* kNodeId = "tier-cap-node-0"; + + void StartCluster(bool advertise_max_allocatable_bytes) { + ASSERT_NO_FATAL_FAILURE(StartMaster(Affinity::kLocal, advertise_max_allocatable_bytes)); + ASSERT_NO_FATAL_FAILURE(StartClient(kNodeId, {DramBackend("dram-0", kPrimaryCapacity, 1), + DramBackend("dram-1", kSecondaryCapacity, 1)})); + Settle(std::chrono::milliseconds(150)); + client()->Master().FlushHeartbeat(); + } + + size_t FillLocalPool(const std::string& key_prefix) { + std::vector value(kPageSize, 0x3C); + size_t succeeded = 0; + const size_t attempts = kPrimaryPages + kSecondaryPages; + for (size_t i = 0; i < attempts; ++i) { + if (client()->Put(key_prefix + std::to_string(i), value.data(), value.size())) ++succeeded; + } + return succeeded; + } +}; + +class NewMasterWeightedPoolTest : public UnequalDramWeightedPoolTest { + protected: + void SetUp() override { StartCluster(/*advertise_max_allocatable_bytes=*/true); } +}; + +class LegacyMasterWeightedPoolTest : public UnequalDramWeightedPoolTest { + protected: + void SetUp() override { StartCluster(/*advertise_max_allocatable_bytes=*/false); } +}; + +TEST_F(NewMasterWeightedPoolTest, NegotiatesMaxAllocatableAndSumsSameTierCapacity) { + EXPECT_TRUE(client()->Master().SupportsMaxAllocatableCapacity()) + << "current Master must advertise supports_max_allocatable_bytes so weighted " + "peers can aggregate same-tier capacity"; + + MediumBackend* primary = client()->Backends().Get("dram-0"); + MediumBackend* secondary = client()->Backends().Get("dram-1"); + ASSERT_NE(primary, nullptr); + ASSERT_NE(secondary, nullptr); + EXPECT_EQ(primary->Capacity().total_bytes + secondary->Capacity().total_bytes, kPooledCapacity); + + EXPECT_EQ(FillLocalPool("new-master-"), kPrimaryPages + kSecondaryPages); + EXPECT_GT(primary->OwnedKeyCount(), 0u); + EXPECT_GT(secondary->OwnedKeyCount(), 0u); + + client()->Master().FlushHeartbeat(); + EXPECT_TRUE(WaitForAdvertisedDram(master_.get(), kNodeId, kPooledCapacity, /*available=*/0)) + << "new Master must store the aggregated same-tier pool, not the first instance"; + EXPECT_NE(AdvertisedDramTotal(master_.get(), kNodeId), kPrimaryCapacity); +} + +TEST_F(LegacyMasterWeightedPoolTest, KeepsFirstInstanceCapacityAndStillPlacesOnBothBackends) { + EXPECT_FALSE(client()->Master().SupportsMaxAllocatableCapacity()) + << "omitting supports_max_allocatable_bytes must disable aggregate heartbeats"; + + MediumBackend* primary = client()->Backends().Get("dram-0"); + MediumBackend* secondary = client()->Backends().Get("dram-1"); + ASSERT_NE(primary, nullptr); + ASSERT_NE(secondary, nullptr); + + EXPECT_EQ(FillLocalPool("legacy-"), kPrimaryPages + kSecondaryPages); + EXPECT_GT(primary->OwnedKeyCount(), 0u); + EXPECT_GT(secondary->OwnedKeyCount(), 0u); + EXPECT_EQ(primary->OwnedKeyCount() + secondary->OwnedKeyCount(), kPrimaryPages + kSecondaryPages) + << "weighted placement remains enabled under a legacy Master; only advertisement shrinks"; + + client()->Master().FlushHeartbeat(); + EXPECT_TRUE(WaitForAdvertisedDram(master_.get(), kNodeId, kPrimaryCapacity, /*available=*/0)) + << "heartbeats stay at the first same-tier instance so the old Master cannot " + "over-admit a value larger than any single backend"; + EXPECT_EQ(AdvertisedDramTotal(master_.get(), kNodeId), kPrimaryCapacity); + EXPECT_NE(AdvertisedDramTotal(master_.get(), kNodeId), kPooledCapacity); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_umbp_workload_runner.cpp b/tests/cpp/umbp/distributed/test_umbp_workload_runner.cpp new file mode 100644 index 000000000..d259efd81 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_umbp_workload_runner.cpp @@ -0,0 +1,268 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/workload_runner.h" + +namespace mori::umbp::benchmark { +namespace { + +using Event = ::umbp::benchmark::WorkloadEvent; + +Event MakeEvent(uint32_t client, uint64_t id, Event::Operation operation, const std::string& key, + uint64_t timestamp_ns = 0, uint64_t batch_id = 0) { + Event event; + event.set_client_id(client); + event.set_operation_id(id); + event.set_operation(operation); + event.set_key(key); + event.set_value_size(32); + event.set_relative_timestamp_ns(timestamp_ns); + event.set_batch_id(batch_id); + return event; +} + +WorkloadRunnerOptions MaxThroughputOptions() { + WorkloadRunnerOptions options; + options.max_throughput = true; + return options; +} + +class VectorSource final : public WorkloadSource { + public: + explicit VectorSource(std::vector events, uint64_t seed = 7) + : events_(std::move(events)), seed_(seed) {} + + uint64_t seed() const override { return seed_; } + bool Next(Event* event) override { + if (next_ == events_.size()) return false; + *event = events_[next_++]; + return true; + } + + private: + std::vector events_; + uint64_t seed_; + size_t next_ = 0; +}; + +struct ClientCall { + uint32_t client_id; + std::string key; +}; + +class FakeClient final : public WorkloadClient { + public: + ClientResult Put(uint32_t client_id, const std::string& key, const uint8_t* data, + size_t size) override { + std::lock_guard lock(mutex_); + calls_.push_back({client_id, key}); + if (key == "put-fail") return ClientResult::kFailed; + values_[client_id][key] = std::vector(data, data + size); + return ClientResult::kSuccess; + } + + ClientResult Get(uint32_t client_id, const std::string& key, uint8_t* data, + size_t size) override { + std::lock_guard lock(mutex_); + calls_.push_back({client_id, key}); + const auto client = values_.find(client_id); + if (client == values_.end()) return ClientResult::kNotFound; + const auto found = client->second.find(key); + if (found == client->second.end()) return ClientResult::kNotFound; + if (found->second.size() != size) return ClientResult::kFailed; + std::copy(found->second.begin(), found->second.end(), data); + if (key == "corrupt" && size != 0) data[0] ^= 1; + return ClientResult::kSuccess; + } + + std::vector BatchPut(uint32_t client_id, const std::vector& keys, + const std::vector>& values) override { + ++batch_put_calls; + return WorkloadClient::BatchPut(client_id, keys, values); + } + + std::vector BatchGet(uint32_t client_id, const std::vector& keys, + const std::vector& sizes, + std::vector>* values) override { + ++batch_get_calls; + return WorkloadClient::BatchGet(client_id, keys, sizes, values); + } + + std::vector Calls() const { + std::lock_guard lock(mutex_); + return calls_; + } + + std::atomic batch_put_calls{0}; + std::atomic batch_get_calls{0}; + + private: + mutable std::mutex mutex_; + std::unordered_map>> values_; + std::vector calls_; +}; + +TEST(WorkloadRunnerTest, PreservesOrderingWithinEachClient) { + std::vector events; + for (uint64_t i = 0; i < 20; ++i) { + events.push_back(MakeEvent(0, i + 1, Event::PUT, "a-" + std::to_string(i))); + events.push_back(MakeEvent(1, i + 101, Event::PUT, "b-" + std::to_string(i))); + } + VectorSource source(std::move(events)); + FakeClient client; + WorkloadRunner runner(&client, MaxThroughputOptions()); + const WorkloadMetrics metrics = runner.Run(&source); + EXPECT_EQ(metrics.puts.succeeded, 40); + + uint64_t next_a = 0; + uint64_t next_b = 0; + for (const ClientCall& call : client.Calls()) { + if (call.key[0] == 'a') { + EXPECT_EQ(call.key, "a-" + std::to_string(next_a++)); + } else { + EXPECT_EQ(call.key, "b-" + std::to_string(next_b++)); + } + } + EXPECT_EQ(next_a, 20); + EXPECT_EQ(next_b, 20); +} + +TEST(WorkloadRunnerTest, DeliversDistinctClientIdsInPerClientOrder) { + VectorSource source({ + MakeEvent(11, 1, Event::PUT, "first-11"), + MakeEvent(22, 2, Event::PUT, "first-22"), + MakeEvent(11, 3, Event::PUT, "second-11"), + MakeEvent(22, 4, Event::PUT, "second-22"), + }); + FakeClient client; + WorkloadRunner runner(&client, MaxThroughputOptions()); + EXPECT_EQ(runner.Run(&source).total.succeeded, 4); + + std::vector client_11_keys; + std::vector client_22_keys; + for (const ClientCall& call : client.Calls()) { + if (call.client_id == 11) { + client_11_keys.push_back(call.key); + } else if (call.client_id == 22) { + client_22_keys.push_back(call.key); + } else { + ADD_FAILURE() << "unexpected client id " << call.client_id; + } + } + EXPECT_EQ(client_11_keys, (std::vector{"first-11", "second-11"})); + EXPECT_EQ(client_22_keys, (std::vector{"first-22", "second-22"})); +} + +TEST(WorkloadRunnerTest, HonorsSchedulingAndMaxThroughputModes) { + constexpr uint64_t kDelayNs = 30'000'000; + FakeClient scheduled_client; + VectorSource scheduled_source({MakeEvent(0, 1, Event::PUT, "scheduled", kDelayNs)}); + const auto scheduled_start = std::chrono::steady_clock::now(); + WorkloadRunner scheduled_runner(&scheduled_client); + WorkloadMetrics scheduled_metrics = scheduled_runner.Run(&scheduled_source); + const auto scheduled_elapsed = std::chrono::steady_clock::now() - scheduled_start; + EXPECT_GE(scheduled_elapsed, std::chrono::milliseconds(20)); + ASSERT_EQ(scheduled_metrics.schedule_lag.samples_ns.size(), 1); + + FakeClient fast_client; + VectorSource fast_source({MakeEvent(0, 1, Event::PUT, "fast", kDelayNs)}); + const auto fast_start = std::chrono::steady_clock::now(); + WorkloadRunner fast_runner(&fast_client, MaxThroughputOptions()); + WorkloadMetrics fast_metrics = fast_runner.Run(&fast_source); + const auto fast_elapsed = std::chrono::steady_clock::now() - fast_start; + EXPECT_LT(fast_elapsed, scheduled_elapsed); + EXPECT_TRUE(fast_metrics.schedule_lag.samples_ns.empty()); +} + +TEST(WorkloadRunnerTest, GroupsCompatibleBatchIds) { + std::vector events; + for (uint64_t i = 0; i < 4; ++i) { + events.push_back(MakeEvent(0, i + 1, Event::PUT, "key-" + std::to_string(i), 0, 7)); + } + for (uint64_t i = 0; i < 4; ++i) { + events.push_back(MakeEvent(0, i + 1, Event::GET, "key-" + std::to_string(i), 0, 8)); + } + VectorSource source(std::move(events)); + FakeClient client; + WorkloadRunner runner(&client, MaxThroughputOptions()); + const WorkloadMetrics metrics = runner.Run(&source); + EXPECT_EQ(client.batch_put_calls, 1); + EXPECT_EQ(client.batch_get_calls, 1); + EXPECT_EQ(metrics.total.succeeded, 8); +} + +TEST(WorkloadRunnerTest, AccountsForFailuresAndComputesPercentiles) { + std::vector events = { + MakeEvent(0, 1, Event::PUT, "good"), MakeEvent(0, 1, Event::GET, "good"), + MakeEvent(0, 2, Event::GET, "missing"), MakeEvent(0, 3, Event::PUT, "put-fail"), + MakeEvent(0, 4, Event::PUT, "corrupt"), MakeEvent(0, 4, Event::GET, "corrupt"), + }; + VectorSource source(std::move(events)); + FakeClient client; + WorkloadRunner runner(&client, MaxThroughputOptions()); + const WorkloadMetrics metrics = runner.Run(&source); + + EXPECT_EQ(metrics.total.attempted, 6); + EXPECT_EQ(metrics.total.succeeded, 3); + EXPECT_EQ(metrics.total.failed, 3); + EXPECT_EQ(metrics.get_misses, 1); + EXPECT_EQ(metrics.get_validation_failures, 1); + ASSERT_EQ(metrics.latency.samples_ns.size(), 6); + EXPECT_LE(metrics.latency.p50_ns, metrics.latency.p95_ns); + EXPECT_LE(metrics.latency.p95_ns, metrics.latency.p99_ns); + EXPECT_LE(metrics.latency.p99_ns, metrics.latency.max_ns); +} + +TEST(WorkloadRunnerTest, RejectsConflictingPutsForImmutableKeys) { + VectorSource source({ + MakeEvent(0, 1, Event::PUT, "immutable"), + MakeEvent(0, 2, Event::PUT, "immutable"), + }); + FakeClient client; + WorkloadRunner runner(&client, MaxThroughputOptions()); + EXPECT_THROW(runner.Run(&source), std::invalid_argument); +} + +TEST(WorkloadRunnerTest, RejectsDuplicatePutsSoTransferredBytesStayExact) { + VectorSource source({ + MakeEvent(0, 1, Event::PUT, "immutable"), + MakeEvent(0, 1, Event::PUT, "immutable"), + }); + FakeClient client; + WorkloadRunner runner(&client, MaxThroughputOptions()); + EXPECT_THROW(runner.Run(&source), std::invalid_argument); +} + +} // namespace +} // namespace mori::umbp::benchmark diff --git a/tests/cpp/umbp/distributed/test_umbp_workload_trace.cpp b/tests/cpp/umbp/distributed/test_umbp_workload_trace.cpp new file mode 100644 index 000000000..7c796ee76 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_umbp_workload_trace.cpp @@ -0,0 +1,281 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/benchmark/payload.h" +#include "umbp/distributed/benchmark/workload_source.h" +#include "umbp/distributed/benchmark/workload_trace.h" +#include "umbp/distributed/benchmark/workload_trace_recorder.h" + +namespace mori::umbp::benchmark { +namespace { + +std::string TempPath(const std::string& suffix) { + static std::atomic sequence{0}; + const auto id = std::chrono::steady_clock::now().time_since_epoch().count(); + return (std::filesystem::temp_directory_path() / + ("umbp-workload-" + std::to_string(getpid()) + "-" + std::to_string(id) + "-" + + std::to_string(sequence.fetch_add(1)) + suffix)) + .string(); +} + +::umbp::benchmark::WorkloadTraceHeader Header() { + ::umbp::benchmark::WorkloadTraceHeader header; + header.set_schema_version(kWorkloadTraceSchemaVersion); + header.set_time_unit(::umbp::benchmark::WorkloadTraceHeader::TIME_UNIT_NANOSECONDS); + header.set_seed(42); + (*header.mutable_metadata())["profile"] = "test"; + return header; +} + +::umbp::benchmark::WorkloadEvent Event() { + ::umbp::benchmark::WorkloadEvent event; + event.set_relative_timestamp_ns(1234); + event.set_client_id(3); + event.set_operation_id(9); + event.set_operation(::umbp::benchmark::WorkloadEvent::PUT); + event.set_key("key"); + event.set_value_size(257); + event.set_batch_id(7); + return event; +} + +std::vector ReadBytes(const std::string& path) { + std::ifstream input(path, std::ios::binary); + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; +} + +void WriteBytes(const std::string& path, const std::vector& bytes) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + output.write(bytes.data(), static_cast(bytes.size())); +} + +TEST(WorkloadTraceTest, RoundTripsHeaderAndEvents) { + const std::string path = TempPath(".trace"); + { + TraceWriter writer(path, Header()); + writer.Write(Event()); + writer.Close(); + } + + TraceReader reader(path); + EXPECT_EQ(reader.header().seed(), 42); + EXPECT_EQ(reader.header().metadata().at("profile"), "test"); + ::umbp::benchmark::WorkloadEvent actual; + ASSERT_TRUE(reader.ReadNext(&actual)); + EXPECT_EQ(actual.SerializeAsString(), Event().SerializeAsString()); + EXPECT_FALSE(reader.ReadNext(&actual)); + std::filesystem::remove(path); +} + +TEST(WorkloadTraceTest, RejectsTruncationCorruptionAndUnsupportedEnvelope) { + const std::string valid = TempPath("-valid.trace"); + { + TraceWriter writer(valid, Header()); + writer.Write(Event()); + writer.Close(); + } + const std::vector bytes = ReadBytes(valid); + + const std::string truncated = TempPath("-truncated.trace"); + WriteBytes(truncated, std::vector(bytes.begin(), bytes.end() - 1)); + TraceReader truncated_reader(truncated); + ::umbp::benchmark::WorkloadEvent event; + EXPECT_THROW(truncated_reader.ReadNext(&event), TraceError); + + const std::string corrupt = TempPath("-corrupt.trace"); + auto corrupt_bytes = bytes; + corrupt_bytes.back() = static_cast(0xff); + WriteBytes(corrupt, corrupt_bytes); + TraceReader corrupt_reader(corrupt); + EXPECT_THROW(corrupt_reader.ReadNext(&event), TraceError); + + const std::string unsupported = TempPath("-unsupported.trace"); + auto unsupported_bytes = bytes; + unsupported_bytes[8] = 2; + WriteBytes(unsupported, unsupported_bytes); + EXPECT_THROW(TraceReader ignored(unsupported), TraceError); + + const std::string unsupported_schema = TempPath("-unsupported-schema.trace"); + auto unsupported_schema_bytes = bytes; + ASSERT_GT(unsupported_schema_bytes.size(), 17); + ASSERT_EQ(static_cast(unsupported_schema_bytes[16]), 0x08); + unsupported_schema_bytes[17] = 2; + WriteBytes(unsupported_schema, unsupported_schema_bytes); + EXPECT_THROW(TraceReader ignored(unsupported_schema), TraceError); + + std::filesystem::remove(valid); + std::filesystem::remove(truncated); + std::filesystem::remove(corrupt); + std::filesystem::remove(unsupported); + std::filesystem::remove(unsupported_schema); +} + +TEST(WorkloadTraceTest, PayloadIsDeterministicAndValidatable) { + const auto first = GenerateDeterministicPayload("alpha", 17, 99, 37); + const auto second = GenerateDeterministicPayload("alpha", 17, 99, 37); + EXPECT_EQ(first, second); + EXPECT_NE(first, GenerateDeterministicPayload("alpha", 18, 99, 37)); + EXPECT_TRUE(ValidateDeterministicPayload("alpha", 17, 99, first.data(), first.size())); + auto damaged = first; + damaged[13] ^= 1; + EXPECT_FALSE(ValidateDeterministicPayload("alpha", 17, 99, damaged.data(), damaged.size())); +} + +TEST(WorkloadTraceTest, SyntheticGenerationIsDeterministicAcrossProfiles) { + for (SyntheticProfile profile : {SyntheticProfile::kSequential, SyntheticProfile::kUniform, + SyntheticProfile::kHotsetZipf, SyntheticProfile::kReadAfterWrite, + SyntheticProfile::kMixed, SyntheticProfile::kCapacityPressure}) { + SyntheticWorkloadConfig config; + config.profile = profile; + config.seed = 123; + config.operation_count = 32; + config.key_count = 8; + config.min_value_size = 8; + config.max_value_size = 64; + config.value_size_distribution = ValueSizeDistribution::kUniform; + config.client_count = 3; + config.batch_size = 2; + config.qps = 1000; + + SyntheticWorkloadSource first(config); + SyntheticWorkloadSource second(config); + ::umbp::benchmark::WorkloadEvent left; + ::umbp::benchmark::WorkloadEvent right; + for (uint64_t i = 0; i < config.operation_count; ++i) { + ASSERT_TRUE(first.Next(&left)); + ASSERT_TRUE(second.Next(&right)); + EXPECT_EQ(left.SerializeAsString(), right.SerializeAsString()); + } + EXPECT_FALSE(first.Next(&left)); + } +} + +TEST(WorkloadTraceTest, SyntheticKeysAreImmutableAndDependenciesStayOnOneClient) { + SyntheticWorkloadConfig config; + config.profile = SyntheticProfile::kMixed; + config.seed = 9; + config.operation_count = 64; + config.key_count = 4; + config.client_count = 3; + config.batch_size = 2; + config.qps = 1000; + config.read_ratio = 0.5; + + SyntheticWorkloadSource source(config); + std::map put_ids; + ::umbp::benchmark::WorkloadEvent event; + while (source.Next(&event)) { + const size_t version_separator = event.key().find("-v", config.key_prefix.size()); + ASSERT_NE(version_separator, std::string::npos); + const uint64_t base_key = std::stoull( + event.key().substr(config.key_prefix.size(), version_separator - config.key_prefix.size())); + EXPECT_EQ(event.client_id(), base_key % config.client_count); + if (event.operation() == ::umbp::benchmark::WorkloadEvent::PUT) { + EXPECT_TRUE(put_ids.emplace(event.key(), event.operation_id()).second); + } else { + ASSERT_TRUE(put_ids.count(event.key())); + EXPECT_EQ(event.operation_id(), put_ids[event.key()]); + } + } + EXPECT_GE(put_ids.size(), config.key_count); +} + +TEST(WorkloadTraceTest, CapacityPressureUsesUniqueWritesAndBatchTimestamps) { + SyntheticWorkloadConfig config; + config.profile = SyntheticProfile::kCapacityPressure; + config.operation_count = 12; + config.key_count = 2; + config.client_count = 2; + config.batch_size = 3; + config.qps = 600; + config.read_ratio = 1.0; // pressure overrides this to remain write-only + + SyntheticWorkloadSource source(config); + std::set keys; + std::map timestamp_by_batch; + ::umbp::benchmark::WorkloadEvent event; + while (source.Next(&event)) { + EXPECT_EQ(event.operation(), ::umbp::benchmark::WorkloadEvent::PUT); + EXPECT_TRUE(keys.insert(event.key()).second); + const auto [it, inserted] = + timestamp_by_batch.emplace(event.batch_id(), event.relative_timestamp_ns()); + if (!inserted) EXPECT_EQ(event.relative_timestamp_ns(), it->second); + } + EXPECT_EQ(keys.size(), config.operation_count); +} + +TEST(WorkloadTraceRecorderTest, KeepsKeysVerbatimAndRecordsEveryOutcome) { + const std::string path = TempPath(".trace"); + { + WorkloadTraceRecorderOptions options; + options.path = path; + options.client_id = 7; + options.seed = 42; + WorkloadTraceRecorder recorder(std::move(options)); + recorder.RecordBatchPut({"key"}, {64}, {WorkloadTraceOutcome::kSuccess}); + recorder.RecordBatchPut({"key"}, {64}, {WorkloadTraceOutcome::kSuccess}); + recorder.RecordBatchPut({"full"}, {64}, {WorkloadTraceOutcome::kFailure}); + recorder.RecordBatchGet({"key", "missing"}, {64, 64}, + {WorkloadTraceOutcome::kSuccess, WorkloadTraceOutcome::kMiss}); + EXPECT_EQ(recorder.event_count(), 5u); + recorder.Close(); + } + + TraceReader reader(path); + std::vector<::umbp::benchmark::WorkloadEvent> events; + ::umbp::benchmark::WorkloadEvent event; + while (reader.ReadNext(&event)) events.push_back(event); + ASSERT_EQ(events.size(), 5u); + // The overwrite is two writes of one key, which is what makes it an + // overwrite; a versioned key would replay as two unrelated objects. + EXPECT_EQ(events[0].key(), "key"); + EXPECT_EQ(events[1].key(), "key"); + EXPECT_EQ(events[2].key(), "full"); + EXPECT_EQ(events[2].outcome(), ::umbp::benchmark::WorkloadEvent::OUTCOME_FAILURE); + EXPECT_EQ(events[3].key(), "key"); + EXPECT_EQ(events[3].outcome(), ::umbp::benchmark::WorkloadEvent::OUTCOME_SUCCESS); + // A GET for a key this client never wrote is still traffic the tiers saw. + EXPECT_EQ(events[4].key(), "missing"); + EXPECT_EQ(events[4].outcome(), ::umbp::benchmark::WorkloadEvent::OUTCOME_MISS); + EXPECT_EQ(events[3].batch_id(), events[4].batch_id()); + EXPECT_LE(events[0].relative_timestamp_ns(), events[1].relative_timestamp_ns()); + EXPECT_LE(events[1].relative_timestamp_ns(), events[3].relative_timestamp_ns()); + std::filesystem::remove(path); +} + +} // namespace +} // namespace mori::umbp::benchmark diff --git a/tests/cpp/umbp/local/CMakeLists.txt b/tests/cpp/umbp/local/CMakeLists.txt index 301b38055..22494e540 100644 --- a/tests/cpp/umbp/local/CMakeLists.txt +++ b/tests/cpp/umbp/local/CMakeLists.txt @@ -1,45 +1,40 @@ -# C++ unit tests -add_executable(test_umbp_local_block_index test_local_block_index.cpp) -target_link_libraries(test_umbp_local_block_index PRIVATE umbp_core) -add_test(NAME umbp_local_block_index COMMAND test_umbp_local_block_index) - -add_executable(test_umbp_local_storage_manager test_local_storage_manager.cpp) -target_link_libraries(test_umbp_local_storage_manager PRIVATE umbp_core) -add_test(NAME umbp_local_storage_manager - COMMAND test_umbp_local_storage_manager) - +# Storage-tier unit tests. +# +# These cover umbp_core's tier layer -- the SSD/SPDK/segment backends and the +# host allocator -- which outlived the local client that used to sit on top of +# them: PeerSsdManager builds the distributed SSD medium on exactly these +# classes. The tests that drove StandaloneClient, LocalStorageManager, +# LocalBlockIndex and DRAMTier went with those types; DRAM is now covered by the +# PageBackend tests under ../distributed. add_executable(test_umbp_host_mem_allocator test_host_mem_allocator.cpp) target_link_libraries(test_umbp_host_mem_allocator PRIVATE umbp_core) add_test(NAME umbp_host_mem_allocator COMMAND test_umbp_host_mem_allocator) -add_executable(test_umbp_local_client test_umbp_client.cpp) -target_link_libraries(test_umbp_local_client PRIVATE umbp_core) -add_test(NAME umbp_local_client COMMAND test_umbp_local_client) +add_executable(test_ssd_tier test_ssd_tier.cpp) +target_link_libraries(test_ssd_tier PRIVATE umbp_core) +add_test(NAME test_ssd_tier COMMAND test_ssd_tier) -add_executable(test_umbp_local_follower_mode test_follower_mode.cpp) -target_link_libraries(test_umbp_local_follower_mode PRIVATE umbp_core) -add_test(NAME umbp_local_follower_mode COMMAND test_umbp_local_follower_mode) +add_executable(test_umbp_ssd_tier_ranges test_ssd_tier_ranges.cpp) +target_link_libraries(test_umbp_ssd_tier_ranges PRIVATE umbp_core GTest::gtest_main) +target_compile_features(test_umbp_ssd_tier_ranges PRIVATE cxx_std_17) +add_test(NAME umbp_ssd_tier_ranges COMMAND test_umbp_ssd_tier_ranges) -add_executable(test_umbp_local_eviction test_prefix_aware_eviction.cpp) -target_link_libraries(test_umbp_local_eviction PRIVATE umbp_core) -add_test(NAME umbp_local_eviction COMMAND test_umbp_local_eviction) +add_executable(test_segment_crc test_segment_crc.cpp) +target_link_libraries(test_segment_crc PRIVATE umbp_core) +add_test(NAME test_segment_crc COMMAND test_segment_crc) -add_executable(test_prefix_aware_eviction test_prefix_aware_eviction.cpp) -target_link_libraries(test_prefix_aware_eviction PRIVATE umbp_core) -add_test(NAME test_prefix_aware_eviction COMMAND test_prefix_aware_eviction) +add_executable(test_sharded_ssd_tier test_sharded_ssd_tier.cpp) +target_link_libraries(test_sharded_ssd_tier PRIVATE umbp_core) +add_test(NAME test_sharded_ssd_tier COMMAND test_sharded_ssd_tier) -add_executable(test_ssd_tier test_ssd_tier.cpp) -target_link_libraries(test_ssd_tier PRIVATE umbp_core) -add_test(NAME test_ssd_tier COMMAND test_ssd_tier) +add_executable(test_ssd_direct_io test_ssd_direct_io.cpp) +target_link_libraries(test_ssd_direct_io PRIVATE umbp_core) +add_test(NAME test_ssd_direct_io COMMAND test_ssd_direct_io) add_executable(test_dummy_ssd_tier test_dummy_ssd_tier.cpp) target_link_libraries(test_dummy_ssd_tier PRIVATE umbp_core) add_test(NAME test_dummy_ssd_tier COMMAND test_dummy_ssd_tier) -# Micro-benchmark (manual run only, not registered as a test) -add_executable(bench_umbp_micro bench_umbp_micro.cpp) -target_link_libraries(bench_umbp_micro PRIVATE umbp_core) - # --------------------------------------------------------------------------- # Helper: symlink spdk_proxy next to a target so SpawnProxyDaemon() finds it # --------------------------------------------------------------------------- @@ -55,8 +50,6 @@ endfunction() # When spdk_proxy is available, symlink it next to targets that need it if(TARGET spdk_proxy) - umbp_symlink_spdk_proxy(bench_umbp_micro) - # SPDK test/benchmark targets (require SPDK libraries at link time) add_executable(test_spdk_env test_spdk_env.cpp) target_link_libraries(test_spdk_env PRIVATE umbp_core) @@ -73,15 +66,7 @@ if(UNIX) add_executable(bench_proxy_ipc bench_proxy_ipc.cpp) target_link_libraries(bench_proxy_ipc PRIVATE umbp_core) - add_executable(bench_umbp_e2e bench_umbp_e2e.cpp) - target_link_libraries(bench_umbp_e2e PRIVATE umbp_core) - - add_executable(test_umbp_e2e test_umbp_e2e.cpp) - target_link_libraries(test_umbp_e2e PRIVATE umbp_core) - if(TARGET spdk_proxy) umbp_symlink_spdk_proxy(bench_proxy_ipc) - umbp_symlink_spdk_proxy(bench_umbp_e2e) - umbp_symlink_spdk_proxy(test_umbp_e2e) endif() endif() diff --git a/tests/cpp/umbp/local/bench_umbp_e2e.cpp b/tests/cpp/umbp/local/bench_umbp_e2e.cpp deleted file mode 100644 index 32095d6e6..000000000 --- a/tests/cpp/umbp/local/bench_umbp_e2e.cpp +++ /dev/null @@ -1,744 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -// -// umbp_bench: End-to-end benchmark using StandaloneClient (full stack). -// -// Modes: -// Throughput (default): -// ./umbp_bench # POSIX single -// UMBP_SPDK_NVME_PCI=... ./umbp_bench --ranks=8 # SPDK multi -// -// Latency (single-thread QD=1, per-op timing): -// ./umbp_bench --latency # POSIX -// UMBP_SPDK_NVME_PCI=... ./umbp_bench --latency # SPDK -// -// Cold reads use O_DIRECT (POSIX) or bypass ring cache (SPDK). - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef __linux__ -#include -#include -#include -#include -#endif - -#include "umbp/local/standalone_client.h" - -using namespace mori::umbp; - -static double NowSec() { - auto tp = std::chrono::steady_clock::now(); - return std::chrono::duration(tp.time_since_epoch()).count(); -} - -static std::string MakeSessionId() { - std::random_device rd; - char buf[16]; - snprintf(buf, sizeof(buf), "%08x", rd()); - return buf; -} - -struct SizeSpec { - size_t size; - int count; -}; - -static const SizeSpec kSpecs[] = { - {4 * 1024, 2000}, {32 * 1024, 2000}, {128 * 1024, 2000}, - {512 * 1024, 1024}, {1024 * 1024, 512}, {2ULL * 1024 * 1024, 256}, - {8ULL * 1024 * 1024, 64}, {16ULL * 1024 * 1024, 32}, {32ULL * 1024 * 1024, 16}, - {64ULL * 1024 * 1024, 8}, {128ULL * 1024 * 1024, 4}, {256ULL * 1024 * 1024, 2}, - {288ULL * 1024 * 1024, 2}, {512ULL * 1024 * 1024, 2}, -}; - -struct BenchResult { - size_t value_size; - int count; - double write_mbps; - double read_mbps; // first-read (cold, real business scenario) - double cache_read_mbps; // best of subsequent iterations (heap cache) - int read_ok; - int corrupt; - double cold_t0, cold_t1; // wall-clock timestamps for cold read - double cache_t0, cache_t1; // wall-clock timestamps for best cache read -}; - -// Coordination structure for multi-rank phased benchmark. -// Shared between forked processes via anonymous mmap. -struct alignas(64) BenchCoord { - std::atomic write_gen; // Leader sets to size_idx+1 after writing - std::atomic reads_complete; // All ranks increment after reading - std::atomic iter_barrier; // Sync ranks between read iterations - int num_ranks; - char session[16]; // Shared session ID -}; - -// --------------------------------------------------------------------------- -// Single-rank benchmark: same process writes and reads (original behavior) -// --------------------------------------------------------------------------- -static BenchResult RunBatch(StandaloneClient& client, int rank_id, const std::string& session, - size_t value_size, int count, int iterations) { - std::string prefix = - "bench_r" + std::to_string(rank_id) + "_" + session + "_" + std::to_string(value_size) + "_"; - - uint32_t seed = 0; - for (char c : session) seed = seed * 31 + static_cast(c); - - std::vector> bufs(count); - std::vector sizes(count, value_size); - for (int i = 0; i < count; ++i) - bufs[i].resize(value_size, static_cast((seed + i + 1) & 0xFF)); - - double total_bytes = static_cast(value_size) * count; - - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) return {value_size, count, 0, 0, 0, 0, 0, 0, 0, 0, 0}; - - std::vector cptrs(count); - for (int i = 0; i < count; ++i) cptrs[i] = bufs[i].data(); - - double best_write = 0; - for (int iter = 0; iter < iterations; ++iter) { - std::vector wkeys(count); - for (int i = 0; i < count; ++i) - wkeys[i] = prefix + "w" + std::to_string(iter) + "_" + std::to_string(i); - - double t0 = NowSec(); - auto wr = ssd->BatchWrite(wkeys, cptrs, sizes); - double t1 = NowSec(); - - int ok = 0; - for (auto b : wr) ok += b; - double mbps = (total_bytes / (1024.0 * 1024.0)) / (t1 - t0); - if (ok == count && mbps > best_write) best_write = mbps; - } - - std::vector rkeys(count); - for (int i = 0; i < count; ++i) rkeys[i] = prefix + "r_" + std::to_string(i); - - int write_ok = 0; - for (int attempt = 0; attempt < 3 && write_ok < count; ++attempt) { - auto wr = ssd->BatchWrite(rkeys, cptrs, sizes); - write_ok = 0; - for (auto b : wr) write_ok += b; - if (write_ok < count && attempt < 2) - fprintf(stderr, " [rank %d] read-key write %d/%d, retrying...\n", rank_id, write_ok, count); - } - if (write_ok < count) - fprintf(stderr, " [rank %d] WARNING: only wrote %d/%d read-keys for %zuKB\n", rank_id, - write_ok, count, value_size / 1024); - - std::vector> read_bufs(count, std::vector(value_size, 0)); - std::vector dst_ptrs(count); - for (int i = 0; i < count; ++i) dst_ptrs[i] = reinterpret_cast(read_bufs[i].data()); - - double first_read = 0; - double cache_read = 0; - double cold_t0 = 0, cold_t1 = 0; - double cache_last_t0 = 0, cache_last_t1 = 0; - int best_read_ok = 0; - int corrupt = 0; - for (int iter = 0; iter < iterations; ++iter) { - ssd->SetColdRead(iter == 0); - - for (auto& b : read_bufs) std::memset(b.data(), 0, b.size()); - - double t0 = NowSec(); - auto rr = ssd->BatchReadIntoPtr(rkeys, dst_ptrs, sizes); - double t1 = NowSec(); - - int ok = 0; - for (auto b : rr) ok += b; - double mbps = (total_bytes / (1024.0 * 1024.0)) / (t1 - t0); - - if (iter == 0) { - if (ok > 0) { - first_read = mbps; - best_read_ok = ok; - } - cold_t0 = t0; - cold_t1 = t1; - for (int i = 0; i < count; ++i) { - if (!rr[i]) continue; - char expected = static_cast((seed + i + 1) & 0xFF); - const char* buf = read_bufs[i].data(); - if (buf[0] != expected || buf[value_size / 2] != expected || - buf[value_size - 1] != expected) { - ++corrupt; - if (corrupt <= 3) - fprintf(stderr, - " [rank %d] CORRUPT: key %d, expected 0x%02x, " - "got first=0x%02x last=0x%02x\n", - rank_id, i, (unsigned char)expected, (unsigned char)buf[0], - (unsigned char)buf[value_size - 1]); - } - } - } else { - if (ok > 0) { - cache_read = mbps; - cache_last_t0 = t0; - cache_last_t1 = t1; - } - } - } - ssd->SetColdRead(false); - - return {value_size, count, best_write, first_read, cache_read, best_read_ok, - corrupt, cold_t0, cold_t1, cache_last_t0, cache_last_t1}; -} - -// --------------------------------------------------------------------------- -// Multi-rank phased benchmark (real Prefill/Decode scenario) -// Phase 1: Leader writes data (measure write bandwidth) -// Phase 2: All ranks concurrently read the SAME data (measure read bandwidth) -// Phase 3: Barrier — sync before next size -// --------------------------------------------------------------------------- -#ifdef __linux__ -static BenchResult RunBatchPhased(StandaloneClient& client, int rank_id, size_t value_size, - int count, int iterations, int size_idx, BenchCoord* coord) { - std::string session(coord->session); - std::string prefix = "bench_" + session + "_" + std::to_string(value_size) + "_"; - - uint32_t seed = 0; - for (char c : session) seed = seed * 31 + static_cast(c); - - double total_bytes = static_cast(value_size) * count; - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) return {value_size, count, 0, 0, 0, 0, 0, 0, 0, 0, 0}; - - double best_write = 0; - - // === WRITE PHASE (Leader only, others wait) === - if (rank_id == 0) { - std::vector> bufs(count); - std::vector sizes(count, value_size); - for (int i = 0; i < count; ++i) - bufs[i].resize(value_size, static_cast((seed + i + 1) & 0xFF)); - - std::vector cptrs(count); - for (int i = 0; i < count; ++i) cptrs[i] = bufs[i].data(); - - for (int iter = 0; iter < iterations; ++iter) { - std::vector wkeys(count); - for (int i = 0; i < count; ++i) - wkeys[i] = prefix + "w" + std::to_string(iter) + "_" + std::to_string(i); - - double t0 = NowSec(); - auto wr = ssd->BatchWrite(wkeys, cptrs, sizes); - double t1 = NowSec(); - - int ok = 0; - for (auto b : wr) ok += b; - double mbps = (total_bytes / (1024.0 * 1024.0)) / (t1 - t0); - if (ok == count && mbps > best_write) best_write = mbps; - } - - // Write the read-benchmark keys (data all ranks will read) - std::vector rkeys(count); - for (int i = 0; i < count; ++i) rkeys[i] = prefix + "r_" + std::to_string(i); - - int write_ok = 0; - for (int attempt = 0; attempt < 3 && write_ok < count; ++attempt) { - auto wr = ssd->BatchWrite(rkeys, cptrs, sizes); - write_ok = 0; - for (auto b : wr) write_ok += b; - } - if (write_ok < count) - fprintf(stderr, " [Leader] WARNING: wrote %d/%d read-keys for %zuKB\n", write_ok, count, - value_size / 1024); - - client.Flush(); - - coord->write_gen.store(size_idx + 1, std::memory_order_release); - } - - // === WAIT for Leader to finish writing === - while (coord->write_gen.load(std::memory_order_acquire) < size_idx + 1) sched_yield(); - - // === READ PHASE (All ranks read the SAME data concurrently) === - std::vector rkeys(count); - for (int i = 0; i < count; ++i) rkeys[i] = prefix + "r_" + std::to_string(i); - - std::vector> read_bufs(count, std::vector(value_size, 0)); - std::vector dst_ptrs(count); - std::vector sizes(count, value_size); - for (int i = 0; i < count; ++i) dst_ptrs[i] = reinterpret_cast(read_bufs[i].data()); - - double first_read = 0; // iteration 0: cold read (O_DIRECT / bypass cache) - double cache_read = 0; // last cache iteration (steady-state) - double cold_t0 = 0, cold_t1 = 0; - double cache_last_t0 = 0, cache_last_t1 = 0; - int best_read_ok = 0; - int corrupt = 0; - for (int iter = 0; iter < iterations; ++iter) { - ssd->SetColdRead(iter == 0); - - if (iter > 0) { - coord->iter_barrier.fetch_add(1, std::memory_order_acq_rel); - int sync = (size_idx * (iterations - 1) + iter) * coord->num_ranks; - while (coord->iter_barrier.load(std::memory_order_acquire) < sync) sched_yield(); - } - - for (auto& b : read_bufs) std::memset(b.data(), 0, b.size()); - - double t0 = NowSec(); - auto rr = ssd->BatchReadIntoPtr(rkeys, dst_ptrs, sizes); - double t1 = NowSec(); - - int ok = 0; - for (auto b : rr) ok += b; - double mbps = (total_bytes / (1024.0 * 1024.0)) / (t1 - t0); - - if (iter == 0) { - if (ok > 0) { - first_read = mbps; - best_read_ok = ok; - } - cold_t0 = t0; - cold_t1 = t1; - for (int i = 0; i < count; ++i) { - if (!rr[i]) continue; - char expected = static_cast((seed + i + 1) & 0xFF); - const char* buf = read_bufs[i].data(); - if (buf[0] != expected || buf[value_size / 2] != expected || - buf[value_size - 1] != expected) { - ++corrupt; - if (corrupt <= 3) - fprintf(stderr, - " [rank %d] CORRUPT: key %d, expected 0x%02x, " - "got first=0x%02x last=0x%02x\n", - rank_id, i, (unsigned char)expected, (unsigned char)buf[0], - (unsigned char)buf[value_size - 1]); - } - } - } else { - if (ok > 0) { - cache_read = mbps; - cache_last_t0 = t0; - cache_last_t1 = t1; - } - } - } - - ssd->SetColdRead(false); - - // === BARRIER: all ranks done reading this size === - coord->reads_complete.fetch_add(1, std::memory_order_acq_rel); - int target = (size_idx + 1) * coord->num_ranks; - while (coord->reads_complete.load(std::memory_order_acquire) < target) sched_yield(); - - return {value_size, count, best_write, first_read, cache_read, best_read_ok, - corrupt, cold_t0, cold_t1, cache_last_t0, cache_last_t1}; -} -#endif - -// --------------------------------------------------------------------------- -// Output formatting -// --------------------------------------------------------------------------- -static void PrintResults(int rank_id, int num_ranks, const char* role_str, const char* backend, - const std::vector& results) { - printf("\n"); - printf("========================================================\n"); - printf(" StandaloneClient E2E Benchmark — rank=%d/%d role=%s backend=%s\n", rank_id, num_ranks, - role_str, backend); - printf("========================================================\n"); - printf(" %8s %6s %10s %10s %12s\n", "ValSize", "Count", "Write MB/s", "ColdRd MB/s", - "HotRd MB/s"); - - for (const auto& r : results) { - char sz_label[16]; - if (r.value_size >= 1024 * 1024) - snprintf(sz_label, sizeof(sz_label), "%zuMB", r.value_size / (1024 * 1024)); - else - snprintf(sz_label, sizeof(sz_label), "%zuKB", r.value_size / 1024); - - if (r.write_mbps > 0) - printf(" %8s %6d %10.0f %10.0f %12.0f", sz_label, r.count, r.write_mbps, r.read_mbps, - r.cache_read_mbps); - else - printf(" %8s %6d %10s %10.0f %12.0f", sz_label, r.count, "-", r.read_mbps, - r.cache_read_mbps); - if (r.read_ok != r.count) printf(" *** READ %d/%d", r.read_ok, r.count); - if (r.corrupt > 0) printf(" *** CORRUPT %d", r.corrupt); - printf("\n"); - } - fflush(stdout); -} - -#ifdef __linux__ -static void WriteResultsToPipe(int fd, const std::vector& results) { - uint32_t n = static_cast(results.size()); - (void)!write(fd, &n, sizeof(n)); - for (const auto& r : results) (void)!write(fd, &r, sizeof(r)); -} - -static std::vector ReadResultsFromPipe(int fd) { - uint32_t n = 0; - if (read(fd, &n, sizeof(n)) != sizeof(n)) return {}; - std::vector results(n); - for (uint32_t i = 0; i < n; ++i) (void)!read(fd, &results[i], sizeof(BenchResult)); - return results; -} -#endif - -static int RunBenchmarkProcess(int rank_id, int num_ranks, int pipe_fd, void* coord_ptr) { - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - - StandaloneClient client(cfg); - - UMBPRole role = cfg.ResolveRole(); - const char* role_str = (role == UMBPRole::Standalone) ? "Standalone" - : (role == UMBPRole::SharedSSDLeader) ? "Leader" - : "Follower"; - const char* backend = cfg.ssd.ssd_backend.c_str(); - - const int iterations = 2; // 0=cold (O_DIRECT/bypass cache), 1=hot - - std::vector results; - -#ifdef __linux__ - auto* coord = static_cast(coord_ptr); - // SPDK Leader/Follower: phased mode (Leader writes → barrier → all read same data) - // POSIX / Standalone: per-rank mode (each rank writes and reads independently) - bool use_phased = - coord && (role == UMBPRole::SharedSSDLeader || role == UMBPRole::SharedSSDFollower); - if (use_phased) { - for (size_t s = 0; s < sizeof(kSpecs) / sizeof(kSpecs[0]); ++s) - results.push_back(RunBatchPhased(client, rank_id, kSpecs[s].size, kSpecs[s].count, iterations, - static_cast(s), coord)); - } else -#endif - { - (void)coord_ptr; - std::string session = MakeSessionId(); - for (const auto& s : kSpecs) - results.push_back(RunBatch(client, rank_id, session, s.size, s.count, iterations)); - } - - if (pipe_fd >= 0) { -#ifdef __linux__ - WriteResultsToPipe(pipe_fd, results); -#endif - } else { - PrintResults(rank_id, num_ranks, role_str, backend, results); - } - - (void)role_str; - (void)backend; - return 0; -} - -// --------------------------------------------------------------------------- -// Latency benchmark: single-thread QD=1, per-op timing -// --------------------------------------------------------------------------- -static const SizeSpec kLatencySpecs[] = { - {4 * 1024, 500}, {32 * 1024, 500}, {128 * 1024, 200}, - {512 * 1024, 100}, {1024 * 1024, 50}, {2ULL * 1024 * 1024, 50}, - {8ULL * 1024 * 1024, 20}, {16ULL * 1024 * 1024, 10}, {32ULL * 1024 * 1024, 8}, - {64ULL * 1024 * 1024, 4}, {128ULL * 1024 * 1024, 4}, {256ULL * 1024 * 1024, 2}, - {288ULL * 1024 * 1024, 2}, {512ULL * 1024 * 1024, 2}, -}; - -struct LatencyStats { - double avg_us, p50_us, p99_us, max_us; -}; - -static LatencyStats ComputeStats(std::vector& samples_us) { - if (samples_us.empty()) return {0, 0, 0, 0}; - std::sort(samples_us.begin(), samples_us.end()); - size_t n = samples_us.size(); - double sum = 0; - for (double v : samples_us) sum += v; - return {sum / static_cast(n), samples_us[n / 2], - samples_us[std::min(n - 1, static_cast(n * 0.99))], samples_us[n - 1]}; -} - -static int RunLatencyBench() { - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - StandaloneClient client(cfg); - - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) { - fprintf(stderr, "ERROR: no SSD tier available\n"); - return 1; - } - - const char* backend = cfg.ssd.ssd_backend.c_str(); - std::string session = MakeSessionId(); - - ssd->SetColdRead(true); - - printf("\n================================================================\n"); - printf(" WRITE LATENCY — backend=%s (single-thread QD=1, O_DIRECT)\n", backend); - printf("================================================================\n"); - printf(" %8s %5s %10s %10s %10s %10s\n", "ValSize", "Count", "Avg(us)", "P50(us)", - "P99(us)", "Max(us)"); - - for (const auto& spec : kLatencySpecs) { - size_t vsize = spec.size; - int count = spec.count; - - char sz_label[16]; - if (vsize >= 1024 * 1024) - snprintf(sz_label, sizeof(sz_label), "%zuMB", vsize / (1024 * 1024)); - else - snprintf(sz_label, sizeof(sz_label), "%zuKB", vsize / 1024); - - std::string prefix = "lat_" + session + "_" + std::to_string(vsize) + "_"; - - std::vector data(vsize); - for (size_t j = 0; j < vsize; ++j) data[j] = static_cast((j + 1) & 0xFF); - - std::vector samples; - samples.reserve(count); - for (int i = 0; i < count; ++i) { - std::string key = prefix + std::to_string(i); - double t0 = NowSec(); - ssd->Write(key, data.data(), vsize); - double t1 = NowSec(); - samples.push_back((t1 - t0) * 1e6); - } - ssd->Flush(); - - auto st = ComputeStats(samples); - printf(" %8s %5d %10.1f %10.1f %10.1f %10.1f\n", sz_label, count, st.avg_us, st.p50_us, - st.p99_us, st.max_us); - } - - printf("\n================================================================\n"); - printf(" READ LATENCY — backend=%s (single-thread QD=1, O_DIRECT)\n", backend); - printf("================================================================\n"); - printf(" %8s %5s %10s %10s %10s %10s\n", "ValSize", "Count", "Avg(us)", "P50(us)", - "P99(us)", "Max(us)"); - - for (const auto& spec : kLatencySpecs) { - size_t vsize = spec.size; - int count = spec.count; - - char sz_label[16]; - if (vsize >= 1024 * 1024) - snprintf(sz_label, sizeof(sz_label), "%zuMB", vsize / (1024 * 1024)); - else - snprintf(sz_label, sizeof(sz_label), "%zuKB", vsize / 1024); - - std::string prefix = "lat_" + session + "_" + std::to_string(vsize) + "_"; - - std::vector read_buf(vsize, 0); - uintptr_t dst = reinterpret_cast(read_buf.data()); - - std::vector samples; - samples.reserve(count); - for (int i = 0; i < count; ++i) { - std::string key = prefix + std::to_string(i); - double t0 = NowSec(); - ssd->ReadIntoPtr(key, dst, vsize); - double t1 = NowSec(); - samples.push_back((t1 - t0) * 1e6); - } - - auto st = ComputeStats(samples); - printf(" %8s %5d %10.1f %10.1f %10.1f %10.1f\n", sz_label, count, st.avg_us, st.p50_us, - st.p99_us, st.max_us); - } - ssd->SetColdRead(false); - - printf("================================================================\n"); - fflush(stdout); - return 0; -} - -int main(int argc, char** argv) { - int num_ranks = 1; - bool latency_mode = false; - for (int i = 1; i < argc; ++i) { - if (strncmp(argv[i], "--ranks=", 8) == 0) { - num_ranks = std::atoi(argv[i] + 8); - } else if (strcmp(argv[i], "--latency") == 0) { - latency_mode = true; - } - } - - if (latency_mode) return RunLatencyBench(); - - if (num_ranks <= 1) { - return RunBenchmarkProcess(0, 1, -1, nullptr); - } - -#ifdef __linux__ - printf("Launching %d ranks (phased: Leader writes → barrier → all read)...\n", num_ranks); - fflush(stdout); - - // Shared coordination via anonymous mmap (inherited by fork children) - auto* coord = static_cast( - mmap(nullptr, sizeof(BenchCoord), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0)); - coord->write_gen.store(0, std::memory_order_relaxed); - coord->reads_complete.store(0, std::memory_order_relaxed); - coord->iter_barrier.store(0, std::memory_order_relaxed); - coord->num_ranks = num_ranks; - std::string session = MakeSessionId(); - strncpy(coord->session, session.c_str(), sizeof(coord->session) - 1); - coord->session[sizeof(coord->session) - 1] = '\0'; - - struct RankInfo { - pid_t pid; - int pipe_fd; - int rank_id; - }; - std::vector ranks; - - for (int r = 0; r < num_ranks; ++r) { - int pipefd[2]; - if (pipe(pipefd) != 0) { - fprintf(stderr, "pipe() failed for rank %d\n", r); - continue; - } - - pid_t pid = fork(); - if (pid == 0) { - close(pipefd[0]); - char rank_str[16]; - snprintf(rank_str, sizeof(rank_str), "%d", r); - setenv("LOCAL_RANK", rank_str, 1); - int rc = RunBenchmarkProcess(r, num_ranks, pipefd[1], coord); - close(pipefd[1]); - _exit(rc); - } - if (pid < 0) { - fprintf(stderr, "fork() failed for rank %d: %s\n", r, strerror(errno)); - close(pipefd[0]); - close(pipefd[1]); - continue; - } - close(pipefd[1]); - ranks.push_back({pid, pipefd[0], r}); - } - - int failures = 0; - struct RankResult { - int rank_id; - std::vector results; - bool ok; - }; - std::vector all_results; - - for (auto& ri : ranks) { - int status; - waitpid(ri.pid, &status, 0); - bool ok = WIFEXITED(status) && WEXITSTATUS(status) == 0; - auto results = ReadResultsFromPipe(ri.pipe_fd); - close(ri.pipe_fd); - all_results.push_back({ri.rank_id, std::move(results), ok}); - if (!ok) failures++; - } - - munmap(coord, sizeof(BenchCoord)); - - printf("\n=== %d/%d ranks completed successfully ===\n", num_ranks - failures, num_ranks); - - auto cfg = UMBPConfig::FromEnvironment(); - const char* backend = cfg.ssd.ssd_backend.c_str(); - bool is_phased = (num_ranks > 1); - - for (auto& rr : all_results) { - const char* role_str; - if (is_phased) - role_str = (rr.rank_id == 0) ? "Leader" : "Follower"; - else - role_str = "Standalone"; - if (rr.ok && !rr.results.empty()) { - PrintResults(rr.rank_id, num_ranks, role_str, backend, rr.results); - } else { - printf("\n[rank %d] FAILED (no results)\n", rr.rank_id); - } - } - - size_t num_specs = sizeof(kSpecs) / sizeof(kSpecs[0]); - if (is_phased) { - printf("\n========================================================\n"); - printf(" AGGREGATE (Leader write / %d-rank concurrent read)\n", num_ranks); - printf("========================================================\n"); - } else { - printf("\n========================================================\n"); - printf(" AGGREGATE (sum of %d independent ranks)\n", num_ranks); - printf("========================================================\n"); - } - printf(" %8s %10s %10s %12s\n", "ValSize", "Write MB/s", "ColdRd MB/s", "HotRd MB/s"); - for (size_t s = 0; s < num_specs; ++s) { - double sum_w = 0; - double min_cold_t0 = 1e18, max_cold_t1 = 0; - double min_cache_t0 = 1e18, max_cache_t1 = 0; - double total_cold_bytes = 0, total_cache_bytes = 0; - for (auto& rr : all_results) { - if (!rr.ok || s >= rr.results.size()) continue; - auto& r = rr.results[s]; - if (is_phased) { - if (rr.rank_id == 0) sum_w = r.write_mbps; - } else { - sum_w += r.write_mbps; - } - - double bytes = static_cast(r.value_size) * r.count; - if (r.cold_t0 > 0 && r.cold_t1 > r.cold_t0) { - if (r.cold_t0 < min_cold_t0) min_cold_t0 = r.cold_t0; - if (r.cold_t1 > max_cold_t1) max_cold_t1 = r.cold_t1; - total_cold_bytes += bytes; - } - if (r.cache_t0 > 0 && r.cache_t1 > r.cache_t0) { - if (r.cache_t0 < min_cache_t0) min_cache_t0 = r.cache_t0; - if (r.cache_t1 > max_cache_t1) max_cache_t1 = r.cache_t1; - total_cache_bytes += bytes; - } - } - - double wall_read = (max_cold_t1 > min_cold_t0) - ? (total_cold_bytes / (1024.0 * 1024.0)) / (max_cold_t1 - min_cold_t0) - : 0; - double wall_cache = (max_cache_t1 > min_cache_t0) ? (total_cache_bytes / (1024.0 * 1024.0)) / - (max_cache_t1 - min_cache_t0) - : 0; - - char sz_label[16]; - if (kSpecs[s].size >= 1024 * 1024) - snprintf(sz_label, sizeof(sz_label), "%zuMB", kSpecs[s].size / (1024 * 1024)); - else - snprintf(sz_label, sizeof(sz_label), "%zuKB", kSpecs[s].size / 1024); - printf(" %8s %10.0f %10.0f %12.0f\n", sz_label, sum_w, wall_read, wall_cache); - } - fflush(stdout); - - return failures > 0 ? 1 : 0; -#else - fprintf(stderr, "Multi-rank mode requires Linux\n"); - return 1; -#endif -} diff --git a/tests/cpp/umbp/local/bench_umbp_micro.cpp b/tests/cpp/umbp/local/bench_umbp_micro.cpp deleted file mode 100644 index c962677e6..000000000 --- a/tests/cpp/umbp/local/bench_umbp_micro.cpp +++ /dev/null @@ -1,3656 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -// UMBP Benchmark Tool -// -// Standardized benchmark for measuring UMBP write-path optimizations: -// zero-copy CopyToSSD, lock splitting, multi-threaded async copy, batch io_uring. -// -// Usage: -// bench_umbp [OPTIONS] -// --profile Preset config (default: medium) -// --num-keys N Keys per scenario -// --value-size N Value size in bytes -// --batch-size N Batch size -// --iters N Measurement iterations (default: 10) -// --ssd-backend SSD backend (default: file) -// --filter SUBSTRING Run only matching scenarios -// --dir PATH Temp directory path -// -h, --help Help - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/common/config.h" -#include "umbp/local/standalone_client.h" -#include "umbp/local/storage_tier.h" -#include "umbp/local/tiers/dram_tier.h" -#include "umbp/local/tiers/local_storage_manager.h" -#include "umbp/local/tiers/segment/segment_format.h" -#include "umbp/local/tiers/spdk_proxy_tier.h" -#include "umbp/local/tiers/ssd_tier.h" -#include "umbp/storage/io/storage_io_driver.h" - -using namespace mori::umbp; - -namespace fs = std::filesystem; - -// --------------------------------------------------------------------------- -// Clock alias -// --------------------------------------------------------------------------- -using Clock = std::chrono::steady_clock; - -// --------------------------------------------------------------------------- -// BenchConfig -// --------------------------------------------------------------------------- -struct BenchConfig { - size_t num_keys = 1000; - size_t value_size = 4096; - size_t batch_size = 64; - size_t warmup_iters = 1; - size_t measure_iters = 10; - - size_t dram_capacity = 64ULL * 1024 * 1024; - size_t ssd_capacity = 256ULL * 1024 * 1024; - size_t segment_size = 64ULL * 1024 * 1024; - UMBPIoBackend ssd_io_backend = UMBPIoBackend::Posix; - size_t ssd_io_queue_depth = 4096; - UMBPDurabilityMode ssd_durability_mode = UMBPDurabilityMode::Relaxed; - - std::vector thread_counts = {1, 2, 4, 8}; - - std::string ssd_backend = "file"; // "file" or "spdk" - - std::string base_dir = "/tmp/umbp_bench"; - std::string filter; - bool list_scenarios = false; -}; - -// --------------------------------------------------------------------------- -// E2E Config — sglang-aligned parameters -// --------------------------------------------------------------------------- -enum class E2EModelMode { MHA, MLA }; - -struct E2EConfig { - E2EModelMode mode = E2EModelMode::MLA; - - // Model architecture (defaults: DeepSeek-V3/R1) - size_t num_layers = 61; - size_t num_kv_heads = 8; // MHA only - size_t head_dim = 128; // MHA only - size_t kv_lora_rank = 512; // MLA only - size_t qk_rope_head_dim = 64; // MLA only - - // Common - std::string kv_cache_dtype = "bf16"; // bf16, fp16, fp8_e4m3, fp8_e5m2 - size_t page_size = 1; // tokens per page - - // Bench parameters - size_t num_pages = 512; - size_t batch_pages = 128; // sglang storage_batch_size - double dedup_ratio = 0.5; - int prefix_depth_base = 10; - - size_t DtypeSize() const { - if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8_e5m2" || kv_cache_dtype == "fp8") - return 1; - return 2; // bf16, fp16 - } - - size_t KvCacheDim() const { return kv_lora_rank + qk_rope_head_dim; } - - size_t ValueSizePerKey() const { - if (mode == E2EModelMode::MLA) - return num_layers * page_size * KvCacheDim() * DtypeSize(); - else - return num_layers * page_size * num_kv_heads * head_dim * DtypeSize(); - } - - size_t KeysPerPage() const { return (mode == E2EModelMode::MLA) ? 1 : 2; } -}; - -static void ApplyModelPreset(E2EConfig& e2e, const std::string& preset) { - if (preset == "deepseek-v3" || preset == "deepseek-r1") { - e2e.mode = E2EModelMode::MLA; - e2e.num_layers = 61; - e2e.kv_lora_rank = 512; - e2e.qk_rope_head_dim = 64; - e2e.kv_cache_dtype = "bf16"; - e2e.page_size = 1; - } else if (preset == "deepseek-v2") { - e2e.mode = E2EModelMode::MLA; - e2e.num_layers = 60; - e2e.kv_lora_rank = 512; - e2e.qk_rope_head_dim = 64; - e2e.kv_cache_dtype = "bf16"; - e2e.page_size = 1; - } else if (preset == "llama-70b") { - e2e.mode = E2EModelMode::MHA; - e2e.num_layers = 80; - e2e.num_kv_heads = 8; - e2e.head_dim = 128; - e2e.kv_cache_dtype = "bf16"; - e2e.page_size = 1; - } else if (preset == "llama-8b") { - e2e.mode = E2EModelMode::MHA; - e2e.num_layers = 32; - e2e.num_kv_heads = 8; - e2e.head_dim = 128; - e2e.kv_cache_dtype = "bf16"; - e2e.page_size = 1; - } else { - std::cerr << "Unknown model preset: " << preset << std::endl; - std::exit(1); - } -} - -static std::string ToLower(std::string value) { - std::transform(value.begin(), value.end(), value.begin(), - [](unsigned char c) { return static_cast(std::tolower(c)); }); - return value; -} - -namespace mori::umbp::bench { - -inline std::string ToLowerCopy(std::string value) { - std::transform(value.begin(), value.end(), value.begin(), - [](unsigned char c) { return static_cast(std::tolower(c)); }); - return value; -} - -inline bool FilterMatches(const std::string& filter, const std::string& scenario_name) { - if (filter.empty()) return true; - return ToLowerCopy(scenario_name).find(ToLowerCopy(filter)) != std::string::npos; -} - -inline size_t CountTrue(const std::vector& values) { - return static_cast(std::count(values.begin(), values.end(), true)); -} - -inline size_t CountSuccessfulPages(const std::vector& key_results, size_t keys_per_page) { - if (keys_per_page == 0) return 0; - size_t ok_pages = 0; - for (size_t i = 0; i < key_results.size(); i += keys_per_page) { - bool page_ok = true; - for (size_t j = 0; j < keys_per_page && i + j < key_results.size(); ++j) { - page_ok = page_ok && key_results[i + j]; - } - if (page_ok) ++ok_pages; - } - return ok_pages; -} - -inline double ThroughputOpsPerSec(size_t successful_ops, double elapsed_sec) { - return elapsed_sec > 0.0 ? static_cast(successful_ops) / elapsed_sec : 0.0; -} - -inline double ThroughputMbPerSec(size_t successful_bytes, double elapsed_sec) { - return elapsed_sec > 0.0 ? static_cast(successful_bytes) / (1024.0 * 1024.0) / elapsed_sec - : 0.0; -} - -struct ResultTally { - size_t requested_ops = 0; - size_t successful_ops = 0; - size_t requested_bytes = 0; - size_t successful_bytes = 0; - std::vector latencies_us; - - void ReserveLatencySamples(size_t n) { latencies_us.reserve(n); } - - size_t failed_ops() const { - return (requested_ops >= successful_ops) ? (requested_ops - successful_ops) : 0; - } - - size_t sample_count() const { return latencies_us.size(); } - - void AddSample(size_t requested, size_t succeeded, size_t req_bytes, size_t ok_bytes, - double total_us) { - requested_ops += requested; - successful_ops += succeeded; - requested_bytes += req_bytes; - successful_bytes += ok_bytes; - if (requested == 0) return; - latencies_us.push_back(total_us); - } - - void AddCall(size_t requested, size_t succeeded, size_t req_bytes, size_t ok_bytes, - double total_us) { - requested_ops += requested; - successful_ops += succeeded; - requested_bytes += req_bytes; - successful_bytes += ok_bytes; - latencies_us.push_back(total_us); - } - - void AddOp(bool ok, size_t bytes, double total_us) { - AddSample(1, ok ? 1 : 0, bytes, ok ? bytes : 0, total_us); - } - - void Merge(const ResultTally& other) { - requested_ops += other.requested_ops; - successful_ops += other.successful_ops; - requested_bytes += other.requested_bytes; - successful_bytes += other.successful_bytes; - latencies_us.insert(latencies_us.end(), other.latencies_us.begin(), other.latencies_us.end()); - } -}; - -enum class BatchWriteMode { - Fused, - Fallback, -}; - -inline size_t EstimateRecordBytes(size_t key_size, size_t value_size, size_t record_header_size) { - return record_header_size + key_size + value_size; -} - -inline size_t EstimateBatchRecordBytes(size_t keys_in_batch, size_t key_size, size_t value_size, - size_t record_header_size) { - return keys_in_batch * EstimateRecordBytes(key_size, value_size, record_header_size); -} - -inline BatchWriteMode ClassifyBatchWriteMode(size_t keys_in_batch, size_t key_size, - size_t value_size, size_t segment_size, - size_t record_header_size) { - return EstimateBatchRecordBytes(keys_in_batch, key_size, value_size, record_header_size) <= - segment_size - ? BatchWriteMode::Fused - : BatchWriteMode::Fallback; -} - -struct WorkloadSummary { - size_t num_keys = 0; - size_t value_size = 0; - size_t batch_size = 0; - size_t dram_capacity = 0; - size_t ssd_capacity = 0; - size_t segment_size = 0; - size_t key_size_hint = 0; - size_t record_header_size = 0; -}; - -struct ConfigWarning { - std::string message; -}; - -inline std::vector CollectConfigWarnings(const WorkloadSummary& summary) { - std::vector warnings; - if (summary.value_size > summary.dram_capacity) { - warnings.push_back( - {"value_size exceeds DRAM capacity; DRAM-resident and cold CopyToSSD scenarios will " - "partially fail or be skipped."}); - } - if (summary.value_size > summary.ssd_capacity) { - warnings.push_back( - {"value_size exceeds SSD capacity; single-value SSD writes cannot succeed."}); - } - if (summary.num_keys * summary.value_size > summary.ssd_capacity) { - warnings.push_back( - {"working set exceeds SSD capacity; warm-cache SSD read scenarios will report mixed " - "residency rather than full-hit throughput."}); - } - if (summary.batch_size > 0 && - ClassifyBatchWriteMode(summary.batch_size, summary.key_size_hint, summary.value_size, - summary.segment_size, - summary.record_header_size) == BatchWriteMode::Fallback) { - warnings.push_back( - {"batch payload exceeds segment_size; WriteBatch-style scenarios will " - "exercise fallback/per-key behavior instead of fused writes."}); - } - return warnings; -} - -} // namespace mori::umbp::bench - -namespace bench = mori::umbp::bench; - -struct IoBackendSpec { - UMBPIoBackend backend; - const char* cli_name; - const char* display_name; - const char* path_suffix; -}; - -static const std::vector& IoBackendSpecs() { - static const std::vector specs = { - {UMBPIoBackend::Posix, "posix", "POSIX", "posix"}, - {UMBPIoBackend::IoUring, "io_uring", "io_uring", "iouring"}, - }; - return specs; -} - -static const IoBackendSpec& GetIoBackendSpec(UMBPIoBackend backend) { - for (const auto& spec : IoBackendSpecs()) { - if (spec.backend == backend) return spec; - } - throw std::invalid_argument("Unknown I/O backend enum value"); -} - -static bool ParseIoBackend(const std::string& text, UMBPIoBackend& backend) { - std::string lower = ToLower(text); - if (lower == "posix") { - backend = UMBPIoBackend::Posix; - return true; - } - if (lower == "io_uring" || lower == "iouring" || lower == "uring") { - backend = UMBPIoBackend::IoUring; - return true; - } - return false; -} - -static const char* DurabilityLabel(UMBPDurabilityMode mode) { - return (mode == UMBPDurabilityMode::Strict) ? "Strict" : "Relaxed"; -} - -static bool ParseDurabilityMode(const std::string& text, UMBPDurabilityMode& mode) { - std::string lower = ToLower(text); - if (lower == "strict" || lower == "sync") { - mode = UMBPDurabilityMode::Strict; - return true; - } - if (lower == "relaxed" || lower == "async") { - mode = UMBPDurabilityMode::Relaxed; - return true; - } - return false; -} - -static std::string BackendVariantLabel(UMBPIoBackend backend, size_t queue_depth) { - const auto& spec = GetIoBackendSpec(backend); - if (backend != UMBPIoBackend::IoUring) return "backend=" + std::string(spec.display_name); - return "backend=" + std::string(spec.display_name) + "/qd=" + std::to_string(queue_depth); -} - -static bool IsSpdk(const BenchConfig& cfg) { return cfg.ssd_backend == "spdk"; } - -static UMBPConfig MakeBaseSsdConfig(const BenchConfig& cfg) { - UMBPConfig ucfg = IsSpdk(cfg) ? UMBPConfig::FromEnvironment() : UMBPConfig(); - ucfg.ssd.enabled = true; - ucfg.ssd.ssd_backend = cfg.ssd_backend; - ucfg.ssd.capacity_bytes = cfg.ssd_capacity; - ucfg.ssd.io.backend = cfg.ssd_io_backend; - ucfg.ssd.io.queue_depth = cfg.ssd_io_queue_depth; - ucfg.ssd.durability.mode = cfg.ssd_durability_mode; - ucfg.ssd.segment_size_bytes = cfg.segment_size; - return ucfg; -} - -static UMBPConfig MakeBaseSsdConfig(const BenchConfig& cfg, UMBPIoBackend backend, - UMBPDurabilityMode durability, size_t queue_depth) { - UMBPConfig ucfg = MakeBaseSsdConfig(cfg); - ucfg.ssd.io.backend = backend; - ucfg.ssd.io.queue_depth = queue_depth; - ucfg.ssd.durability.mode = durability; - ucfg.ssd.ssd_backend = cfg.ssd_backend; - return ucfg; -} - -static UMBPConfig MakeStandaloneClientConfig(const BenchConfig& cfg, const std::string& storage_dir, - size_t dram_capacity, - size_t ssd_capacity_override = 0) { - UMBPConfig ucfg = MakeBaseSsdConfig(cfg); - ucfg.dram.capacity_bytes = dram_capacity; - ucfg.ssd.enabled = true; - ucfg.ssd.storage_dir = storage_dir; - if (ssd_capacity_override > 0) ucfg.ssd.capacity_bytes = ssd_capacity_override; - ucfg.role = IsSpdk(cfg) ? UMBPRole::SharedSSDLeader : UMBPRole::Standalone; - ucfg.copy_pipeline.async_enabled = false; - ucfg.eviction.auto_promote_on_read = false; - return ucfg; -} - -static UMBPConfig MakeLeaderClientConfig(const BenchConfig& cfg, const std::string& storage_dir, - size_t dram_capacity, bool async_copy, - size_t ssd_capacity_override = 0) { - UMBPConfig ucfg = - MakeStandaloneClientConfig(cfg, storage_dir, dram_capacity, ssd_capacity_override); - ucfg.role = UMBPRole::SharedSSDLeader; - ucfg.copy_pipeline.async_enabled = async_copy; - return ucfg; -} - -static UMBPConfig MakeFollowerClientConfig(const BenchConfig& cfg, const std::string& storage_dir, - size_t dram_capacity, size_t ssd_capacity_override = 0) { - UMBPConfig ucfg = - MakeStandaloneClientConfig(cfg, storage_dir, dram_capacity, ssd_capacity_override); - ucfg.role = UMBPRole::SharedSSDFollower; - return ucfg; -} - -struct ScopedFd { - int fd = -1; - - ScopedFd() = default; - explicit ScopedFd(int fd) : fd(fd) {} - ~ScopedFd() { - if (fd >= 0) close(fd); - } - - ScopedFd(const ScopedFd&) = delete; - ScopedFd& operator=(const ScopedFd&) = delete; - - ScopedFd(ScopedFd&& other) noexcept : fd(other.fd) { other.fd = -1; } - ScopedFd& operator=(ScopedFd&& other) noexcept { - if (this != &other) { - if (fd >= 0) close(fd); - fd = other.fd; - other.fd = -1; - } - return *this; - } -}; - -static ScopedFd OpenBenchFile(const std::string& path) { - int fd = open(path.c_str(), O_CREAT | O_TRUNC | O_RDWR, 0644); - if (fd < 0) { - throw std::runtime_error("Failed to open benchmark file '" + path + - "': " + std::string(std::strerror(errno))); - } - return ScopedFd(fd); -} - -static void EnsureIoOk(const IoStatus& status, const std::string& context) { - if (!status.ok()) { - throw std::runtime_error(context + ": " + status.message()); - } -} - -// --------------------------------------------------------------------------- -// BenchResult -// --------------------------------------------------------------------------- -struct DisplayTag { - std::string key; - std::string value; -}; - -struct BenchDisplay { - std::string benchmark; - std::string workload; - std::vector scenario_tags; - std::string latency_scope; -}; - -struct BenchResult { - BenchDisplay display; - size_t requested_ops = 0; - size_t successful_ops = 0; - size_t failed_ops = 0; - size_t requested_bytes = 0; - size_t successful_bytes = 0; - size_t sample_count = 0; - double elapsed_sec = 0.0; - - double lat_min_us = 0.0; - double lat_avg_us = 0.0; - double lat_p50_us = 0.0; - double lat_p95_us = 0.0; - double lat_p99_us = 0.0; - double lat_max_us = 0.0; - - double throughput_ops_sec() const { - return bench::ThroughputOpsPerSec(successful_ops, elapsed_sec); - } - double throughput_mb_sec() const { - return bench::ThroughputMbPerSec(successful_bytes, elapsed_sec); - } -}; - -static DisplayTag Tag(std::string key, std::string value) { - return DisplayTag{std::move(key), std::move(value)}; -} - -static int DisplayTagOrder(const std::string& key) { - static const std::vector order = { - "backend", "qd", "durability", "path", "source", "phase", "resident", - "dram", "ssd_seed", "dram_seed", "batch", "bs", "op", "threads", - "copy", "fds", "hit", "dedup", "bytes", - }; - for (size_t i = 0; i < order.size(); ++i) { - if (order[i] == key) return static_cast(i); - } - return static_cast(order.size()); -} - -static BenchDisplay MakeDisplay(std::string benchmark, std::string workload = {}, - std::initializer_list scenario_tags = {}, - std::string latency_scope = {}) { - BenchDisplay display; - display.benchmark = std::move(benchmark); - display.workload = std::move(workload); - display.scenario_tags.assign(scenario_tags.begin(), scenario_tags.end()); - std::sort(display.scenario_tags.begin(), display.scenario_tags.end(), - [](const DisplayTag& lhs, const DisplayTag& rhs) { - int lhs_order = DisplayTagOrder(lhs.key); - int rhs_order = DisplayTagOrder(rhs.key); - if (lhs_order != rhs_order) return lhs_order < rhs_order; - return lhs.key < rhs.key; - }); - display.latency_scope = std::move(latency_scope); - return display; -} - -static bool ContainsTagKey(const std::vector& tags, const std::string& key) { - return std::any_of(tags.begin(), tags.end(), - [&](const DisplayTag& tag) { return tag.key == key; }); -} - -static BenchDisplay ParseLegacyDisplay(const std::string& benchmark, const std::string& variant) { - std::vector parts; - std::stringstream ss(variant); - std::string part; - while (std::getline(ss, part, '/')) { - if (!part.empty()) parts.push_back(part); - } - - std::string workload; - size_t i = 0; - if (parts.size() >= 3 && (parts[0] == "MLA" || parts[0] == "MHA") && - parts[1].find("pg") != std::string::npos && parts[2].find("KB") != std::string::npos) { - workload = parts[0] + "/" + parts[1] + "/" + parts[2]; - i = 3; - } - - std::vector tags; - std::string explicit_latency_scope; - for (; i < parts.size(); ++i) { - const std::string& token = parts[i]; - size_t eq = token.find('='); - if (eq == std::string::npos) continue; - - std::string key = token.substr(0, eq); - std::string value = token.substr(eq + 1); - if ((key == "resident" || key == "dram_seed" || key == "ssd_seed") && i + 1 < parts.size() && - parts[i + 1].find('=') == std::string::npos) { - value += "/" + parts[++i]; - } - - if (key == "timing") { - explicit_latency_scope = value; - continue; - } - tags.push_back(Tag(key, value)); - } - - if (explicit_latency_scope.empty()) { - if (benchmark.find("ExistsScan") != std::string::npos || - benchmark.find("PrefetchCycle") != std::string::npos) { - explicit_latency_scope = "per-cycle"; - } else if (benchmark.rfind("E2E ", 0) == 0 && benchmark.find("Leader") == std::string::npos) { - explicit_latency_scope = "per-batch"; - } else if (ContainsTagKey(tags, "batch") || ContainsTagKey(tags, "bs") || - ContainsTagKey(tags, "fds")) { - explicit_latency_scope = "per-batch"; - } else { - explicit_latency_scope = "per-op"; - } - } - - BenchDisplay display; - display.benchmark = benchmark; - display.workload = workload; - display.scenario_tags = std::move(tags); - std::sort(display.scenario_tags.begin(), display.scenario_tags.end(), - [](const DisplayTag& lhs, const DisplayTag& rhs) { - int lhs_order = DisplayTagOrder(lhs.key); - int rhs_order = DisplayTagOrder(rhs.key); - if (lhs_order != rhs_order) return lhs_order < rhs_order; - return lhs.key < rhs.key; - }); - display.latency_scope = explicit_latency_scope; - return display; -} - -static std::string FormatScenario(const BenchDisplay& display) { - if (display.scenario_tags.empty()) return "-"; - - auto residency_tier_label = [&]() -> std::string { - if (display.benchmark.rfind("DRAM ", 0) == 0) return "DRAM"; - if (display.benchmark.rfind("SSD ", 0) == 0) return "SSD"; - if (display.benchmark.rfind("IO Backend Read", 0) == 0) return "SSD"; - return "storage"; - }; - - auto format_tag_key = [](const std::string& key) { - if (key == "qd") return std::string("queue depth"); - if (key == "source") return std::string("read source"); - if (key == "dram") return std::string("DRAM fit"); - if (key == "ssd_seed") return std::string("SSD preloaded"); - if (key == "dram_seed") return std::string("DRAM preloaded"); - if (key == "batch") return std::string("batch mode"); - if (key == "bs") return std::string("batch size"); - if (key == "op") return std::string("operation"); - if (key == "copy") return std::string("copy mode"); - if (key == "fds") return std::string("files synced"); - if (key == "hit") return std::string("existing pages"); - if (key == "dedup") return std::string("reused pages"); - if (key == "bytes") return std::string("throughput"); - if (key == "resident") return std::string("residency"); - return key; - }; - - auto format_tag_value = [&](const std::string& key, const std::string& value) { - if (key == "path") { - if (value == "ssd-miss") return std::string("needs SSD copy"); - if (value == "ssd-hit") return std::string("already on SSD"); - } - if (key == "resident") { - if (value == "all") return "all data on " + residency_tier_label(); - return value + " on " + residency_tier_label(); - } - if (key == "batch") { - if (value == "fused") return std::string("fused write"); - if (value == "fallback") return std::string("per-key fallback"); - } - if (key == "op") { - if (value == "single") return std::string("single call"); - if (value == "batch") return std::string("batch call"); - } - if (key == "copy") { - if (value == "sync") return std::string("sync"); - if (value == "async") return std::string("async"); - } - if (key == "source" && value == "ssd-only") return std::string("SSD only"); - if (key == "phase") { - if (value == "fits-in-dram") return std::string("fits in DRAM"); - if (value == "spills-to-ssd") return std::string("spills to SSD"); - if (value == "read-after-spill") return std::string("read after spill"); - if (value == "pressure") return std::string("under pressure"); - } - if (key == "hit" || key == "dedup" || key == "dram") { - return value; - } - if (key == "bytes" && value == "physical") return std::string("physical writes only"); - return value; - }; - - std::ostringstream oss; - for (size_t i = 0; i < display.scenario_tags.size(); ++i) { - if (i > 0) oss << ", "; - oss << format_tag_key(display.scenario_tags[i].key) << "=" - << format_tag_value(display.scenario_tags[i].key, display.scenario_tags[i].value); - } - return oss.str(); -} - -static std::string FormatScopeLabel(const std::string& scope) { - if (scope == "per-op") return "single call"; - if (scope == "per-batch") return "batch call"; - if (scope == "per-cycle") return "full cycle"; - if (scope == "end-to-end") return "full request"; - return scope; -} - -static std::vector WrapCell(std::string text, size_t width) { - if (width == 0) return {text}; - if (text.empty()) return {"-"}; - - std::vector lines; - while (text.size() > width) { - size_t split = text.rfind(' ', width); - if (split == std::string::npos || split == 0) { - split = text.rfind(',', width); - if (split != std::string::npos && split + 1 < text.size()) ++split; - } - if (split == std::string::npos || split == 0) split = width; - - lines.push_back(text.substr(0, split)); - text.erase(0, split); - while (!text.empty() && text.front() == ' ') text.erase(text.begin()); - } - if (!text.empty()) lines.push_back(text); - if (lines.empty()) lines.push_back("-"); - return lines; -} - -enum class OutputTableKind { - Detail, - Summary, -}; - -constexpr int kBenchmarkColWidth = 30; -constexpr int kWorkloadColWidth = 20; -constexpr int kScenarioColWidth = 58; -constexpr int kScopeColWidth = 12; - -static void PrintSectionTitle(const std::string& section) { - std::cout << "\n=== " << section << " ===" << std::endl; -} - -static void PrintTableDivider() { - constexpr size_t kDividerWidth = kBenchmarkColWidth + kWorkloadColWidth + kScenarioColWidth + - kScopeColWidth + 8 * 4 + 10 + 9 * 6 + 12 + 18; - std::printf("%s\n", std::string(kDividerWidth, '-').c_str()); -} - -static void PrintTableHeader(OutputTableKind kind) { - if (kind == OutputTableKind::Detail) { - std::printf("%-*s %-*s %-*s %-*s %8s %8s %8s %8s %10s %9s %9s %9s %9s %9s %9s %12s\n", - kBenchmarkColWidth, "Benchmark", kWorkloadColWidth, "Workload", kScenarioColWidth, - "Scenario", kScopeColWidth, "Scope", "req", "ok", "fail", "samples", "MB/s", - "min(us)", "avg(us)", "p50(us)", "p95(us)", "p99(us)", "max(us)", "ops/s"); - } else { - std::printf("%-*s %-*s %-*s %-*s %8s %8s %12s %10s\n", kBenchmarkColWidth, "Benchmark", - kWorkloadColWidth, "Workload", kScenarioColWidth, "Scenario", kScopeColWidth, - "Scope", "ok", "fail", "ops/s", "MB/s"); - } - PrintTableDivider(); -} - -static void PrintSectionHeader(const std::string& section) { - PrintSectionTitle(section); - PrintTableHeader(OutputTableKind::Detail); -} - -static void PrintDetailResult(const BenchResult& r) { - auto benchmark_lines = WrapCell(r.display.benchmark, kBenchmarkColWidth); - auto workload_lines = - WrapCell(r.display.workload.empty() ? "-" : r.display.workload, kWorkloadColWidth); - auto scenario_lines = WrapCell(FormatScenario(r.display), kScenarioColWidth); - auto scope_lines = - WrapCell(r.display.latency_scope.empty() ? "-" : FormatScopeLabel(r.display.latency_scope), - kScopeColWidth); - size_t line_count = std::max(std::max(benchmark_lines.size(), workload_lines.size()), - std::max(scenario_lines.size(), scope_lines.size())); - - for (size_t i = 0; i < line_count; ++i) { - const char* benchmark = (i < benchmark_lines.size()) ? benchmark_lines[i].c_str() : ""; - const char* workload = (i < workload_lines.size()) ? workload_lines[i].c_str() : ""; - const char* scenario = (i < scenario_lines.size()) ? scenario_lines[i].c_str() : ""; - const char* scope = (i < scope_lines.size()) ? scope_lines[i].c_str() : ""; - if (i == 0) { - std::printf( - "%-*s %-*s %-*s %-*s %8zu %8zu %8zu %8zu %10.1f %9.1f %9.1f %9.1f %9.1f %9.1f %9.1f " - "%12.0f\n", - kBenchmarkColWidth, benchmark, kWorkloadColWidth, workload, kScenarioColWidth, scenario, - kScopeColWidth, scope, r.requested_ops, r.successful_ops, r.failed_ops, r.sample_count, - r.throughput_mb_sec(), r.lat_min_us, r.lat_avg_us, r.lat_p50_us, r.lat_p95_us, - r.lat_p99_us, r.lat_max_us, r.throughput_ops_sec()); - } else { - std::printf("%-*s %-*s %-*s %-*s\n", kBenchmarkColWidth, benchmark, kWorkloadColWidth, - workload, kScenarioColWidth, scenario, kScopeColWidth, scope); - } - } -} - -static void PrintSummaryResult(const BenchResult& r) { - auto benchmark_lines = WrapCell(r.display.benchmark, kBenchmarkColWidth); - auto workload_lines = - WrapCell(r.display.workload.empty() ? "-" : r.display.workload, kWorkloadColWidth); - auto scenario_lines = WrapCell(FormatScenario(r.display), kScenarioColWidth); - auto scope_lines = - WrapCell(r.display.latency_scope.empty() ? "-" : FormatScopeLabel(r.display.latency_scope), - kScopeColWidth); - size_t line_count = std::max(std::max(benchmark_lines.size(), workload_lines.size()), - std::max(scenario_lines.size(), scope_lines.size())); - - for (size_t i = 0; i < line_count; ++i) { - const char* benchmark = (i < benchmark_lines.size()) ? benchmark_lines[i].c_str() : ""; - const char* workload = (i < workload_lines.size()) ? workload_lines[i].c_str() : ""; - const char* scenario = (i < scenario_lines.size()) ? scenario_lines[i].c_str() : ""; - const char* scope = (i < scope_lines.size()) ? scope_lines[i].c_str() : ""; - if (i == 0) { - std::printf("%-*s %-*s %-*s %-*s %8zu %8zu %12.0f %10.1f\n", kBenchmarkColWidth, benchmark, - kWorkloadColWidth, workload, kScenarioColWidth, scenario, kScopeColWidth, scope, - r.successful_ops, r.failed_ops, r.throughput_ops_sec(), r.throughput_mb_sec()); - } else { - std::printf("%-*s %-*s %-*s %-*s\n", kBenchmarkColWidth, benchmark, kWorkloadColWidth, - workload, kScenarioColWidth, scenario, kScopeColWidth, scope); - } - } -} - -// --------------------------------------------------------------------------- -// ScopedTempDir — RAII directory cleanup -// --------------------------------------------------------------------------- -struct ScopedTempDir { - std::string path; - explicit ScopedTempDir(const std::string& p) : path(p) { - fs::remove_all(path); - fs::create_directories(path); - } - ~ScopedTempDir() { - std::error_code ec; - fs::remove_all(path, ec); - } - ScopedTempDir(const ScopedTempDir&) = delete; - ScopedTempDir& operator=(const ScopedTempDir&) = delete; -}; - -// --------------------------------------------------------------------------- -// TierScope — RAII for SSDTier (POSIX) or external TierBackend* (SPDK) -// --------------------------------------------------------------------------- -struct TierScope { - std::unique_ptr tmp; - std::unique_ptr local_tier; - TierBackend* tier; - - // POSIX: create local SSDTier - TierScope(const std::string& dir, size_t capacity, const UMBPConfig& ucfg) - : tmp(std::make_unique(dir)), - local_tier(std::make_unique(tmp->path, capacity, ucfg.ssd)), - tier(local_tier.get()) {} - - // SPDK: borrow external tier - explicit TierScope(TierBackend* ext) : tier(ext) {} - - TierScope(const TierScope&) = delete; - TierScope& operator=(const TierScope&) = delete; -}; - -// --------------------------------------------------------------------------- -// Data generation -// --------------------------------------------------------------------------- -static std::string MakeKey(size_t idx) { - char buf[32]; - std::snprintf(buf, sizeof(buf), "bench_%08zu", idx); - return std::string(buf); -} - -static std::vector GenerateKeys(size_t n) { - std::vector keys(n); - for (size_t i = 0; i < n; ++i) { - keys[i] = MakeKey(i); - } - return keys; -} - -static std::vector> GenerateValues(size_t n, size_t value_size) { - std::mt19937 rng(42); - std::uniform_int_distribution dist(0, 255); - std::vector> values(n); - for (size_t i = 0; i < n; ++i) { - values[i].resize(value_size); - for (size_t j = 0; j < value_size; ++j) { - values[i][j] = static_cast(dist(rng)); - } - } - return values; -} - -// --------------------------------------------------------------------------- -// E2E helpers — key generation, host buffer, depth computation -// --------------------------------------------------------------------------- - -// Generates UMBP keys matching umbp_store.py's _batch_preprocess format. -// Key suffix rules (from umbp_store.py lines 154-160): -// MHA pp=1: "{hash}_{tp_rank}_k/v" -// MHA pp>1: "{hash}_{tp_rank}_{pp_rank}_k/v" -// MLA pp=1: "{hash}__k" (mla_suffix="" → double underscore) -// MLA pp>1: "{hash}_{pp_rank}_k" -struct E2EKeyGenerator { - E2EModelMode mode; - size_t tp_rank = 0; - size_t pp_rank = 0; - size_t pp_size = 1; - - // Pre-compute suffix string matching umbp_store.py's __init__. - std::string MhaSuffix() const { - if (pp_size > 1) return std::to_string(tp_rank) + "_" + std::to_string(pp_rank); - return std::to_string(tp_rank); - } - - std::string MlaSuffix() const { - if (pp_size > 1) return std::to_string(pp_rank); - return ""; // pp_size==1: mla_suffix = "" - } - - void KeysForPage(size_t page_idx, std::vector& out) const { - char hash[32]; - std::snprintf(hash, sizeof(hash), "e2e_%08zu", page_idx); - if (mode == E2EModelMode::MLA) { - // "{hash}_{mla_suffix}_k" — when mla_suffix="" this becomes "{hash}__k" - out.push_back(std::string(hash) + "_" + MlaSuffix() + "_k"); - } else { - std::string suffix = MhaSuffix(); - out.push_back(std::string(hash) + "_" + suffix + "_k"); - out.push_back(std::string(hash) + "_" + suffix + "_v"); - } - } - - std::vector KeysForPages(size_t start, size_t count) const { - std::vector keys; - size_t keys_per_page = (mode == E2EModelMode::MLA) ? 1 : 2; - keys.reserve(count * keys_per_page); - for (size_t i = 0; i < count; ++i) { - KeysForPage(start + i, keys); - } - return keys; - } -}; - -// Contiguous host buffer simulating sglang's pinned host KV pool. -// -// Matches sglang memory_pool_host.py get_page_buffer_meta() for page_first layout: -// MHA: [K0 K1 ... Kn | V0 V1 ... Vn] — K region then V region, separated by v_offset. -// k_ptr[i] = base + i * value_size -// v_ptr[i] = k_ptr[i] + v_offset (v_offset = num_pages * value_size) -// MLA: [K0 K1 ... Kn] — single contiguous K region, no V. -struct E2EHostBuffer { - std::vector data; - size_t value_size; - size_t keys_per_page; // 2 for MHA, 1 for MLA - size_t num_pages; - size_t v_offset; // byte offset from K region to V region (MHA only) - - E2EHostBuffer(size_t num_pages, size_t value_size, size_t keys_per_page) - : value_size(value_size), keys_per_page(keys_per_page), num_pages(num_pages) { - // MHA: v_offset = num_pages * value_size (all K pages, then all V pages). - v_offset = num_pages * value_size; - data.resize(num_pages * keys_per_page * value_size); - std::mt19937 rng(12345); - std::uniform_int_distribution dist(0, 255); - for (auto& b : data) b = static_cast(dist(rng)); - } - - // Fill ptrs/sizes for pages [start, start+count), matching get_page_buffer_meta. - // MHA output order: [k_ptr0, v_ptr0, k_ptr1, v_ptr1, ...] - // MLA output order: [k_ptr0, k_ptr1, ...] - void GetBatchMeta(size_t start, size_t count, std::vector& ptrs, - std::vector& sizes) const { - ptrs.clear(); - sizes.clear(); - size_t total_keys = count * keys_per_page; - ptrs.reserve(total_keys); - sizes.reserve(total_keys); - auto base = reinterpret_cast(data.data()); - for (size_t i = 0; i < count; ++i) { - size_t page = start + i; - uintptr_t k_ptr = base + page * value_size; - ptrs.push_back(k_ptr); - sizes.push_back(value_size); - if (keys_per_page == 2) { - ptrs.push_back(k_ptr + v_offset); // V region at fixed offset - sizes.push_back(value_size); - } - } - } - - // Separate read buffer with the same layout for BatchGet. - static void MakeReadMeta(size_t count, size_t keys_per_page, size_t value_size, - std::vector& buf, std::vector& ptrs, - std::vector& sizes) { - size_t total_bytes = count * keys_per_page * value_size; - buf.assign(total_bytes, 0); - ptrs.clear(); - sizes.clear(); - size_t total = count * keys_per_page; - ptrs.reserve(total); - sizes.reserve(total); - auto base = reinterpret_cast(buf.data()); - size_t read_v_offset = count * value_size; - for (size_t i = 0; i < count; ++i) { - uintptr_t k_ptr = base + i * value_size; - ptrs.push_back(k_ptr); - sizes.push_back(value_size); - if (keys_per_page == 2) { - ptrs.push_back(k_ptr + read_v_offset); - sizes.push_back(value_size); - } - } - } -}; - -// Generate depth values matching umbp_store.py _compute_expanded_depths. -// depth = prefix_depth_base + page_index; duplicated for K/V in MHA. -static std::vector GenerateDepths(const E2EConfig& e2e, size_t start_page, size_t count) { - std::vector depths; - depths.reserve(count * e2e.KeysPerPage()); - for (size_t i = 0; i < count; ++i) { - int depth = e2e.prefix_depth_base + static_cast(start_page + i); - depths.push_back(depth); - if (e2e.mode == E2EModelMode::MHA) depths.push_back(depth); // V same depth as K - } - return depths; -} - -// --------------------------------------------------------------------------- -// Latency statistics -// --------------------------------------------------------------------------- -static double PercentileValue(const std::vector& latencies, double pct) { - if (latencies.empty()) return 0.0; - if (latencies.size() == 1) return latencies.front(); - double index = pct * static_cast(latencies.size() - 1); - size_t lo = static_cast(index); - size_t hi = std::min(lo + 1, latencies.size() - 1); - double frac = index - static_cast(lo); - return latencies[lo] + (latencies[hi] - latencies[lo]) * frac; -} - -static void ComputeLatencyStats(std::vector& latencies, BenchResult& result) { - if (latencies.empty()) return; - std::sort(latencies.begin(), latencies.end()); - size_t n = latencies.size(); - result.lat_min_us = latencies.front(); - result.lat_max_us = latencies.back(); - result.lat_avg_us = - std::accumulate(latencies.begin(), latencies.end(), 0.0) / static_cast(n); - result.lat_p50_us = PercentileValue(latencies, 0.50); - result.lat_p95_us = PercentileValue(latencies, 0.95); - result.lat_p99_us = PercentileValue(latencies, 0.99); -} - -// --------------------------------------------------------------------------- -// Result formatting -// --------------------------------------------------------------------------- -static double WallSeconds(const Clock::time_point& start, const Clock::time_point& end) { - return std::chrono::duration(end - start).count(); -} - -static void PrintHeader(const std::string& section) { PrintSectionHeader(section); } - -// Record a benchmark result: compute stats, print, and append to results vector. -static void RecordResult(BenchDisplay display, const bench::ResultTally& tally, double elapsed_sec, - std::vector& results); -static void RecordResult(BenchDisplay display, size_t requested_ops, size_t requested_bytes, - double elapsed_sec, std::vector& latencies, - std::vector& results); - -static void RecordResult(const std::string& name, const std::string& variant, - const bench::ResultTally& tally, double elapsed_sec, - std::vector& results) { - RecordResult(ParseLegacyDisplay(name, variant), tally, elapsed_sec, results); -} - -static void RecordResult(BenchDisplay display, const bench::ResultTally& tally, double elapsed_sec, - std::vector& results) { - BenchResult r; - r.display = std::move(display); - r.requested_ops = tally.requested_ops; - r.successful_ops = tally.successful_ops; - r.failed_ops = tally.failed_ops(); - r.requested_bytes = tally.requested_bytes; - r.successful_bytes = tally.successful_bytes; - r.sample_count = tally.sample_count(); - r.elapsed_sec = elapsed_sec; - auto latencies = tally.latencies_us; - ComputeLatencyStats(latencies, r); - PrintDetailResult(r); - results.push_back(r); -} - -static void RecordResult(const std::string& name, const std::string& variant, size_t requested_ops, - size_t requested_bytes, double elapsed_sec, std::vector& latencies, - std::vector& results) { - RecordResult(ParseLegacyDisplay(name, variant), requested_ops, requested_bytes, elapsed_sec, - latencies, results); -} - -static void RecordResult(BenchDisplay display, size_t requested_ops, size_t requested_bytes, - double elapsed_sec, std::vector& latencies, - std::vector& results) { - bench::ResultTally tally; - tally.requested_ops = requested_ops; - tally.successful_ops = requested_ops; - tally.requested_bytes = requested_bytes; - tally.successful_bytes = requested_bytes; - tally.latencies_us = latencies; - RecordResult(std::move(display), tally, elapsed_sec, results); -} - -// --------------------------------------------------------------------------- -// Scenario filter check -// --------------------------------------------------------------------------- -static bool ShouldRun(const BenchConfig& cfg, const std::string& name) { - return bench::FilterMatches(cfg.filter, name); -} - -static const std::vector& ScenarioFilters() { - static const std::vector filters = { - "DRAM", "SSD Tier", "Batch", "CopyToSSD", "IO Backend", "Durability", - "StorageIo", "Concurrent", "Leader", "Capacity", "E2E", - }; - return filters; -} - -static void PrintScenarioList() { - std::printf("Available scenario filters:\n"); - for (const auto& name : ScenarioFilters()) { - std::printf(" %s\n", name.c_str()); - } -} - -static void PrintSkipMessage(const std::string& section, const std::string& reason) { - std::cout << "\n=== " << section << " ===" << std::endl; - std::printf("[skipped] %s\n", reason.c_str()); -} - -// --------------------------------------------------------------------------- -// Pre-generated batch descriptors -// -// Construct once before warmup/measure so that vector slicing and allocation -// stay outside the timed path, keeping throughput numbers pure. -// --------------------------------------------------------------------------- -struct WriteBatchDesc { - std::vector keys; - std::vector data_ptrs; - std::vector sizes; -}; - -struct ReadBatchDesc { - std::vector keys; - std::vector dst_ptrs; - std::vector sizes; -}; - -static std::vector BuildWriteBatches(const std::vector& all_keys, - const std::vector>& values, - size_t batch_size) { - std::vector descs; - for (size_t i = 0; i < all_keys.size(); i += batch_size) { - size_t end = std::min(i + batch_size, all_keys.size()); - WriteBatchDesc d; - d.keys.assign(all_keys.begin() + i, all_keys.begin() + end); - d.data_ptrs.reserve(end - i); - d.sizes.reserve(end - i); - for (size_t j = i; j < end; ++j) { - d.data_ptrs.push_back(values[j].data()); - d.sizes.push_back(values[j].size()); - } - descs.push_back(std::move(d)); - } - return descs; -} - -static std::vector BuildReadBatches(const std::vector& all_keys, - const std::vector& all_ptrs, - const std::vector& all_sizes, - size_t batch_size) { - std::vector descs; - for (size_t i = 0; i < all_keys.size(); i += batch_size) { - size_t end = std::min(i + batch_size, all_keys.size()); - ReadBatchDesc d; - d.keys.assign(all_keys.begin() + i, all_keys.begin() + end); - d.dst_ptrs.assign(all_ptrs.begin() + i, all_ptrs.begin() + end); - d.sizes.assign(all_sizes.begin() + i, all_sizes.begin() + end); - descs.push_back(std::move(d)); - } - return descs; -} - -// Build key-only batch slices (e.g. for CopyToSSDBatch). -static std::vector> BuildKeyBatches( - const std::vector& all_keys, size_t batch_size) { - std::vector> descs; - for (size_t i = 0; i < all_keys.size(); i += batch_size) { - size_t end = std::min(i + batch_size, all_keys.size()); - descs.emplace_back(all_keys.begin() + i, all_keys.begin() + end); - } - return descs; -} - -static size_t CountExistingKeys(const TierBackend* tier, const std::vector& keys) { - size_t existing = 0; - for (const auto& key : keys) { - if (tier->Exists(key)) ++existing; - } - return existing; -} - -static std::string ResidentValue(size_t resident_keys, size_t total_keys) { - if (resident_keys >= total_keys) return "all"; - std::ostringstream oss; - oss << resident_keys << "/" << total_keys; - return oss.str(); -} - -static std::string ResidencyVariantPrefix(size_t resident_keys, size_t total_keys) { - return "resident=" + ResidentValue(resident_keys, total_keys); -} - -static std::string JoinVariant(std::initializer_list parts) { - std::ostringstream oss; - bool first = true; - for (const auto& value : parts) { - if (value.empty()) continue; - if (!first) oss << "/"; - oss << value; - first = false; - } - return oss.str(); -} - -static bool BatchFitsSegment(const WriteBatchDesc& desc, size_t segment_size) { - size_t total_bytes = 0; - for (size_t i = 0; i < desc.keys.size(); ++i) { - total_bytes += sizeof(segment::RecordHeader) + desc.keys[i].size() + desc.sizes[i]; - } - return total_bytes <= segment_size; -} - -static std::string BatchModeValue(const std::vector& descs, size_t segment_size) { - for (const auto& desc : descs) { - if (!BatchFitsSegment(desc, segment_size)) return "fallback"; - } - return "fused"; -} - -static std::string BatchModeLabel(const std::vector& descs, size_t segment_size) { - return "batch=" + BatchModeValue(descs, segment_size); -} - -static std::string BatchSizeValue(size_t batch_size) { return std::to_string(batch_size); } - -static std::string BackendDisplayValue(UMBPIoBackend backend) { - return GetIoBackendSpec(backend).display_name; -} - -static std::string E2EWorkloadLabel(const E2EConfig& e2e) { - std::string mode = (e2e.mode == E2EModelMode::MLA) ? "MLA" : "MHA"; - return mode + "/" + std::to_string(e2e.batch_pages) + "pg/" + - std::to_string(e2e.ValueSizePerKey() / 1024) + "KB"; -} - -static void PrintConfigWarnings(const BenchConfig& cfg) { - bench::WorkloadSummary summary; - summary.num_keys = cfg.num_keys; - summary.value_size = cfg.value_size; - summary.batch_size = cfg.batch_size; - summary.dram_capacity = cfg.dram_capacity; - summary.ssd_capacity = cfg.ssd_capacity; - summary.segment_size = cfg.segment_size; - summary.key_size_hint = MakeKey(0).size(); - summary.record_header_size = sizeof(segment::RecordHeader); - - auto warnings = bench::CollectConfigWarnings(summary); - for (const auto& warning : warnings) { - std::printf(" warning = %s\n", warning.message.c_str()); - } -} - -// --------------------------------------------------------------------------- -// A. DRAM Tier Benchmarks -// --------------------------------------------------------------------------- -static void BenchDRAMTier(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "DRAM")) return; - if (cfg.value_size > cfg.dram_capacity) { - PrintSkipMessage("DRAM Tier", - "value_size exceeds dram_capacity; a single payload cannot fit in DRAM."); - return; - } - - PrintHeader("DRAM Tier"); - - DRAMTier tier(cfg.dram_capacity); - - // Pre-fill for read benchmarks - size_t resident_keys = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (tier.Write(keys[i], values[i].data(), values[i].size())) ++resident_keys; - } - std::string read_variant = - JoinVariant({ResidencyVariantPrefix(resident_keys, keys.size()), "op=single"}); - - // 1. DRAM Write - { - // Clear and re-fill for write benchmark - tier.Clear(); - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - tier.Write(keys[i], values[i].data(), values[i].size()); - } - } - tier.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier.Write(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - double us = std::chrono::duration(t1 - t0).count(); - tally.AddOp(ok, cfg.value_size, us); - } - } - auto run_end = Clock::now(); - - RecordResult("DRAM Write", "op=single", tally, WallSeconds(run_start, run_end), results); - } - - // Re-fill for reads - tier.Clear(); - resident_keys = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (tier.Write(keys[i], values[i].data(), values[i].size())) ++resident_keys; - } - read_variant = JoinVariant({ResidencyVariantPrefix(resident_keys, keys.size()), "op=single"}); - - // 2. DRAM Read (ReadIntoPtr) - { - std::vector buf(cfg.value_size); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - tier.ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier.ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("DRAM Read", read_variant, tally, WallSeconds(run_start, run_end), results); - } - - // 3. DRAM ReadPtr (zero-copy) - { - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - size_t sz = 0; - tier.ReadPtr(keys[i], &sz); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < keys.size(); ++i) { - size_t sz = 0; - auto t0 = Clock::now(); - const void* ptr = tier.ReadPtr(keys[i], &sz); - auto t1 = Clock::now(); - tally.AddOp(ptr != nullptr && sz == cfg.value_size, cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("DRAM ReadPtr (zero-copy)", read_variant, tally, WallSeconds(run_start, run_end), - results); - } -} - -// --------------------------------------------------------------------------- -// B. SSD Tier Benchmarks -// --------------------------------------------------------------------------- -static void BenchSSDTier(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results, TierBackend* ext_tier = nullptr) { - if (!ShouldRun(cfg, "SSD Tier")) return; - if (!ext_tier && cfg.value_size > cfg.ssd_capacity) { - PrintSkipMessage("SSD Tier", - "value_size exceeds ssd_capacity; a single payload cannot fit in SSD."); - return; - } - - PrintHeader("SSD Tier"); - - UMBPConfig ucfg = MakeBaseSsdConfig(cfg); - - // Write - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/ssd_tier_write", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - tier.Write(keys[i], values[i].data(), values[i].size()); - } - } - tier.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier.Write(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Tier Write", "op=single", tally, WallSeconds(run_start, run_end), results); - } - - // BatchWrite - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/ssd_tier_write_batch", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - auto wdescs = BuildWriteBatches(keys, values, cfg.batch_size); - std::string batch_variant = JoinVariant( - {BatchModeLabel(wdescs, cfg.segment_size), "bs=" + std::to_string(cfg.batch_size)}); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - tier.Clear(); - for (const auto& d : wdescs) tier.BatchWrite(d.keys, d.data_ptrs, d.sizes); - } - tier.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - tier.Clear(); - for (const auto& d : wdescs) { - auto t0 = Clock::now(); - auto write_results = tier.BatchWrite(d.keys, d.data_ptrs, d.sizes); - auto t1 = Clock::now(); - size_t wrote = bench::CountTrue(write_results); - tally.AddSample(d.keys.size(), wrote, d.keys.size() * cfg.value_size, - wrote * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Tier WriteBatch", batch_variant, tally, WallSeconds(run_start, run_end), - results); - } - - // Read - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/ssd_tier_read", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - // Pre-fill - size_t resident_keys = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (tier.Write(keys[i], values[i].data(), values[i].size())) ++resident_keys; - } - std::string read_variant = - JoinVariant({ResidencyVariantPrefix(resident_keys, keys.size()), "op=single"}); - - std::vector buf(cfg.value_size); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - tier.ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier.ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Tier Read", read_variant, tally, WallSeconds(run_start, run_end), results); - } - - // BatchReadIntoPtr - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/ssd_tier_read_batch", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - // Pre-fill - size_t resident_keys = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (tier.Write(keys[i], values[i].data(), values[i].size())) ++resident_keys; - } - std::string read_variant = JoinVariant({ResidencyVariantPrefix(resident_keys, keys.size()), - "bs=" + std::to_string(cfg.batch_size)}); - - // Prepare per-key read buffers and pre-build batch descriptors. - std::vector> bufs(keys.size(), std::vector(cfg.value_size)); - std::vector ptrs(keys.size()); - std::vector sizes(keys.size(), cfg.value_size); - for (size_t i = 0; i < keys.size(); ++i) { - ptrs[i] = reinterpret_cast(bufs[i].data()); - } - auto rdescs = BuildReadBatches(keys, ptrs, sizes, cfg.batch_size); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (const auto& d : rdescs) tier.BatchReadIntoPtr(d.keys, d.dst_ptrs, d.sizes); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (const auto& d : rdescs) { - auto t0 = Clock::now(); - auto batch_results = tier.BatchReadIntoPtr(d.keys, d.dst_ptrs, d.sizes); - auto t1 = Clock::now(); - size_t ok = bench::CountTrue(batch_results); - tally.AddSample(d.keys.size(), ok, d.keys.size() * cfg.value_size, ok * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Tier ReadBatch", read_variant, tally, WallSeconds(run_start, run_end), - results); - } -} - -// --------------------------------------------------------------------------- -// C. Batch vs Single Write -// --------------------------------------------------------------------------- -static void BenchBatchWrite(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results, TierBackend* ext_tier = nullptr) { - if (!ShouldRun(cfg, "Batch")) return; - if (!ext_tier && cfg.value_size > cfg.ssd_capacity) { - PrintSkipMessage("Batch vs Single Write", - "value_size exceeds ssd_capacity; SSD write comparisons cannot succeed."); - return; - } - - PrintHeader("Batch vs Single Write"); - - UMBPConfig ucfg = MakeBaseSsdConfig(cfg); - - // Sequential single-key write - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/batch_seq", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - tier.Write(keys[i], values[i].data(), values[i].size()); - } - } - tier.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier.Write(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Write Compare", "op=single", tally, WallSeconds(run_start, run_end), results); - } - - // BatchWrite - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/batch_batch", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - auto wdescs = BuildWriteBatches(keys, values, cfg.batch_size); - std::string batch_variant = JoinVariant({"op=batch", BatchModeLabel(wdescs, cfg.segment_size), - "bs=" + std::to_string(cfg.batch_size)}); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - tier.Clear(); - for (const auto& d : wdescs) tier.BatchWrite(d.keys, d.data_ptrs, d.sizes); - } - tier.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - tier.Clear(); - for (const auto& d : wdescs) { - auto t0 = Clock::now(); - auto write_results = tier.BatchWrite(d.keys, d.data_ptrs, d.sizes); - auto t1 = Clock::now(); - size_t wrote = bench::CountTrue(write_results); - tally.AddSample(d.keys.size(), wrote, d.keys.size() * cfg.value_size, - wrote * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Write Compare", batch_variant, tally, WallSeconds(run_start, run_end), - results); - } -} - -// --------------------------------------------------------------------------- -// C2. Batch vs Single Read -// --------------------------------------------------------------------------- -static void BenchBatchRead(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results, TierBackend* ext_tier = nullptr) { - if (!ShouldRun(cfg, "Batch")) return; - if (!ext_tier && cfg.value_size > cfg.ssd_capacity) { - PrintSkipMessage("Batch vs Single Read", - "value_size exceeds ssd_capacity; SSD read comparisons cannot succeed."); - return; - } - - PrintHeader("Batch vs Single Read"); - - UMBPConfig ucfg = MakeBaseSsdConfig(cfg); - - // Sequential single-key read - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/batch_read_seq", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - size_t resident_keys = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (tier.Write(keys[i], values[i].data(), values[i].size())) ++resident_keys; - } - std::string seq_variant = - JoinVariant({ResidencyVariantPrefix(resident_keys, keys.size()), "op=single"}); - - std::vector buf(cfg.value_size); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - tier.ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier.ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Read Compare", seq_variant, tally, WallSeconds(run_start, run_end), results); - } - - // BatchReadIntoPtr - { - auto ts = ext_tier ? TierScope(ext_tier) - : TierScope(cfg.base_dir + "/batch_read_batch", cfg.ssd_capacity, ucfg); - TierBackend& tier = *ts.tier; - - size_t resident_keys = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (tier.Write(keys[i], values[i].data(), values[i].size())) ++resident_keys; - } - - std::vector> bufs(keys.size(), std::vector(cfg.value_size)); - std::vector ptrs(keys.size()); - std::vector sizes(keys.size(), cfg.value_size); - for (size_t i = 0; i < keys.size(); ++i) { - ptrs[i] = reinterpret_cast(bufs[i].data()); - } - auto rdescs = BuildReadBatches(keys, ptrs, sizes, cfg.batch_size); - std::string batch_variant = JoinVariant({ResidencyVariantPrefix(resident_keys, keys.size()), - "op=batch", "bs=" + std::to_string(cfg.batch_size)}); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (const auto& d : rdescs) tier.BatchReadIntoPtr(d.keys, d.dst_ptrs, d.sizes); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (const auto& d : rdescs) { - auto t0 = Clock::now(); - auto batch_results = tier.BatchReadIntoPtr(d.keys, d.dst_ptrs, d.sizes); - auto t1 = Clock::now(); - size_t ok = bench::CountTrue(batch_results); - tally.AddSample(d.keys.size(), ok, d.keys.size() * cfg.value_size, ok * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Read Compare", batch_variant, tally, WallSeconds(run_start, run_end), - results); - } -} - -// --------------------------------------------------------------------------- -// D. CopyToSSD vs CopyToSSDBatch -// --------------------------------------------------------------------------- -static void BenchCopyToSSD(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "CopyToSSD")) return; - - PrintHeader("CopyToSSD Paths"); - - auto seed_dram = [&](LocalStorageManager& mgr) { - size_t ready = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (mgr.Write(keys[i], values[i].data(), values[i].size(), StorageTier::CPU_DRAM)) ++ready; - } - return ready; - }; - - auto seed_ssd = [&](TierBackend* ssd) { - size_t ready = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (ssd->Write(keys[i], values[i].data(), values[i].size())) ++ready; - } - return ready; - }; - - auto hot_variant = [&](size_t seeded, const std::string& suffix) { - std::ostringstream oss; - oss << "path=ssd-hit/ssd_seed=" << seeded << "/" << keys.size() << "/" << suffix; - return oss.str(); - }; - - auto cold_variant = [&](size_t seeded, const std::string& suffix) { - std::ostringstream oss; - oss << "path=ssd-miss/dram_seed=" << seeded << "/" << keys.size() << "/" << suffix; - return oss.str(); - }; - - // Single CopyToSSD, cold path. - { - ScopedTempDir tmp(cfg.base_dir + "/copy_single_cold"); - UMBPConfig ucfg = MakeStandaloneClientConfig(cfg, tmp.path + "/ssd", cfg.dram_capacity); - fs::create_directories(ucfg.ssd.storage_dir); - LocalStorageManager mgr(ucfg); - TierBackend* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - size_t dram_ready = seed_dram(mgr); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - mgr.CopyToSSD(keys[i]); - if (ssd->Exists(keys[i])) ssd->Evict(keys[i]); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = mgr.CopyToSSD(keys[i]); - auto t1 = Clock::now(); - bool copied = ok && ssd->Exists(keys[i]); - tally.AddOp(copied, cfg.value_size, - std::chrono::duration(t1 - t0).count()); - if (copied) ssd->Evict(keys[i]); - } - } - auto run_end = Clock::now(); - - RecordResult("CopyToSSD", cold_variant(dram_ready, "op=single"), tally, - WallSeconds(run_start, run_end), results); - } - - // Single CopyToSSD, already-present fast path. - { - ScopedTempDir tmp(cfg.base_dir + "/copy_single_hot"); - UMBPConfig ucfg = MakeStandaloneClientConfig(cfg, tmp.path + "/ssd", cfg.dram_capacity); - fs::create_directories(ucfg.ssd.storage_dir); - LocalStorageManager mgr(ucfg); - TierBackend* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - size_t ssd_ready = seed_ssd(ssd); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - mgr.CopyToSSD(keys[i]); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = mgr.CopyToSSD(keys[i]); - auto t1 = Clock::now(); - tally.AddOp(ok && ssd->Exists(keys[i]), cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("CopyToSSD", hot_variant(ssd_ready, "op=single"), tally, - WallSeconds(run_start, run_end), results); - } - - auto kdescs = BuildKeyBatches(keys, cfg.batch_size); - auto wdescs = BuildWriteBatches(keys, values, cfg.batch_size); - std::string batch_mode = BatchModeLabel(wdescs, cfg.segment_size); - - // Batch CopyToSSDBatch, cold path. - { - ScopedTempDir tmp(cfg.base_dir + "/copy_batch_cold"); - UMBPConfig ucfg = MakeStandaloneClientConfig(cfg, tmp.path + "/ssd", cfg.dram_capacity); - fs::create_directories(ucfg.ssd.storage_dir); - LocalStorageManager mgr(ucfg); - TierBackend* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - size_t dram_ready = seed_dram(mgr); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (const auto& batch : kdescs) { - mgr.CopyToSSDBatch(batch); - for (const auto& key : batch) { - if (ssd->Exists(key)) ssd->Evict(key); - } - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - double measured_sec = 0.0; - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (const auto& batch : kdescs) { - auto t0 = Clock::now(); - mgr.CopyToSSDBatch(batch); - auto t1 = Clock::now(); - measured_sec += WallSeconds(t0, t1); - tally.AddSample(batch.size(), batch.size(), batch.size() * cfg.value_size, - batch.size() * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - // Reset SSD state for next iteration (untimed). - for (const auto& key : batch) { - if (ssd->Exists(key)) ssd->Evict(key); - } - } - } - - RecordResult( - "CopyToSSD", - cold_variant(dram_ready, JoinVariant({batch_mode, "bs=" + std::to_string(cfg.batch_size)})), - tally, measured_sec, results); - } - - // Batch CopyToSSDBatch, already-present fast path. - { - ScopedTempDir tmp(cfg.base_dir + "/copy_batch_hot"); - UMBPConfig ucfg = MakeStandaloneClientConfig(cfg, tmp.path + "/ssd", cfg.dram_capacity); - fs::create_directories(ucfg.ssd.storage_dir); - LocalStorageManager mgr(ucfg); - TierBackend* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - size_t ssd_ready = seed_ssd(ssd); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (const auto& batch : kdescs) mgr.CopyToSSDBatch(batch); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (const auto& batch : kdescs) { - auto t0 = Clock::now(); - mgr.CopyToSSDBatch(batch); - auto t1 = Clock::now(); - // All keys pre-seeded on SSD; no need for CountExistingKeys here. - tally.AddSample(batch.size(), batch.size(), batch.size() * cfg.value_size, - batch.size() * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult( - "CopyToSSD", - hot_variant(ssd_ready, JoinVariant({batch_mode, "bs=" + std::to_string(cfg.batch_size)})), - tally, WallSeconds(run_start, run_end), results); - } -} - -// --------------------------------------------------------------------------- -// E. IO Backend: POSIX vs io_uring -// --------------------------------------------------------------------------- -static void BenchIOBackend(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "IO Backend")) return; - if (IsSpdk(cfg)) { - PrintSkipMessage("IO Backend Sweep", "not applicable for SPDK backend"); - return; - } - if (cfg.value_size > cfg.ssd_capacity) { - PrintSkipMessage("IO Backend Sweep", - "value_size exceeds ssd_capacity; backend SSD comparisons cannot succeed."); - return; - } - - PrintHeader("IO Backend Sweep"); - - auto run_backend = [&](UMBPIoBackend backend) { - const auto& spec = GetIoBackendSpec(backend); - UMBPConfig ucfg = - MakeBaseSsdConfig(cfg, backend, cfg.ssd_durability_mode, cfg.ssd_io_queue_depth); - std::string variant = BackendVariantLabel(backend, cfg.ssd_io_queue_depth); - - // Write - { - ScopedTempDir tmp(cfg.base_dir + "/io_" + std::string(spec.path_suffix) + "_w"); - std::unique_ptr tier; - try { - tier = std::make_unique(tmp.path, cfg.ssd_capacity, ucfg.ssd); - } catch (const std::exception& e) { - std::printf("[SKIP] %s not available: %s\n", variant.c_str(), e.what()); - return; - } - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - tier->Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - tier->Write(keys[i], values[i].data(), values[i].size()); - } - } - tier->Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - tier->Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier->Write(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("IO Backend Write", variant, tally, WallSeconds(run_start, run_end), results); - } - - // Read - { - ScopedTempDir tmp(cfg.base_dir + "/io_" + std::string(spec.path_suffix) + "_r"); - std::unique_ptr tier; - try { - tier = std::make_unique(tmp.path, cfg.ssd_capacity, ucfg.ssd); - } catch (const std::exception& e) { - std::printf("[SKIP] %s not available: %s\n", variant.c_str(), e.what()); - return; - } - - // Pre-fill - size_t resident_keys = 0; - for (size_t i = 0; i < keys.size(); ++i) { - if (tier->Write(keys[i], values[i].data(), values[i].size())) ++resident_keys; - } - std::string read_variant = - JoinVariant({variant, ResidencyVariantPrefix(resident_keys, keys.size())}); - - std::vector buf(cfg.value_size); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - tier->ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier->ReadIntoPtr(keys[i], reinterpret_cast(buf.data()), buf.size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("IO Backend Read", read_variant, tally, WallSeconds(run_start, run_end), - results); - } - }; - - for (const auto& spec : IoBackendSpecs()) { - run_backend(spec.backend); - } -} - -// --------------------------------------------------------------------------- -// F. Durability: Strict vs Relaxed -// --------------------------------------------------------------------------- -static void BenchDurability(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "Durability")) return; - if (IsSpdk(cfg)) { - PrintSkipMessage("Durability (Strict vs Relaxed)", "not applicable for SPDK backend"); - return; - } - if (cfg.value_size > cfg.ssd_capacity) { - PrintSkipMessage( - "Durability (Strict vs Relaxed)", - "value_size exceeds ssd_capacity; durability write comparisons cannot succeed."); - return; - } - - PrintHeader("Durability (Strict vs Relaxed)"); - - auto run_mode = [&](UMBPDurabilityMode mode, const std::string& label) { - UMBPConfig ucfg = MakeBaseSsdConfig(cfg, cfg.ssd_io_backend, mode, cfg.ssd_io_queue_depth); - - ScopedTempDir tmp(cfg.base_dir + "/dur_" + label); - SSDTier tier(tmp.path, cfg.ssd_capacity, ucfg.ssd); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - tier.Write(keys[i], values[i].data(), values[i].size()); - } - } - tier.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - tier.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = tier.Write(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("SSD Write Durability", "durability=" + label, tally, - WallSeconds(run_start, run_end), results); - }; - - run_mode(UMBPDurabilityMode::Strict, "strict"); - run_mode(UMBPDurabilityMode::Relaxed, "relaxed"); -} - -// --------------------------------------------------------------------------- -// G. StorageIoDriver microbench -// --------------------------------------------------------------------------- -static void BenchStorageIoDriver(const BenchConfig& cfg, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "StorageIoDriver")) return; - if (IsSpdk(cfg)) { - PrintSkipMessage("StorageIoDriver", "not applicable for SPDK backend"); - return; - } - - PrintHeader("StorageIoDriver"); - std::string backend_variant = BackendVariantLabel(cfg.ssd_io_backend, cfg.ssd_io_queue_depth); - - std::unique_ptr driver; - try { - driver = - CreateStorageIoDriver(cfg.ssd_io_backend, static_cast(cfg.ssd_io_queue_depth)); - if (cfg.ssd_io_backend == UMBPIoBackend::IoUring && !driver->Capabilities().native_async) { - std::printf("[SKIP] %s not available: native async initialization failed\n", - backend_variant.c_str()); - return; - } - } catch (const std::exception& e) { - std::printf("[SKIP] %s not available: %s\n", backend_variant.c_str(), e.what()); - return; - } - - ScopedTempDir tmp(cfg.base_dir + "/driver_" + - std::string(GetIoBackendSpec(cfg.ssd_io_backend).path_suffix)); - ScopedFd rw_fd = OpenBenchFile(tmp.path + "/rw.bin"); - - auto prefill_rw_file = [&]() { - for (size_t i = 0; i < values.size(); ++i) { - EnsureIoOk(driver->WriteAt(rw_fd.fd, values[i].data(), values[i].size(), i * cfg.value_size), - "prefill rw.bin"); - } - EnsureIoOk(driver->Sync(rw_fd.fd), "sync rw.bin"); - }; - - // --- WriteAt --- - { - std::vector latencies; - latencies.reserve(values.size() * cfg.measure_iters); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < values.size(); ++i) { - EnsureIoOk( - driver->WriteAt(rw_fd.fd, values[i].data(), values[i].size(), i * cfg.value_size), - "warmup WriteAt"); - } - } - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < values.size(); ++i) { - auto t0 = Clock::now(); - EnsureIoOk( - driver->WriteAt(rw_fd.fd, values[i].data(), values[i].size(), i * cfg.value_size), - "WriteAt"); - auto t1 = Clock::now(); - latencies.push_back(std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("Driver WriteAt", backend_variant, values.size() * cfg.measure_iters, - values.size() * cfg.measure_iters * cfg.value_size, - WallSeconds(run_start, run_end), latencies, results); - } - - // --- ReadAt --- - { - prefill_rw_file(); - std::vector buf(cfg.value_size); - std::vector latencies; - latencies.reserve(values.size() * cfg.measure_iters); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < values.size(); ++i) { - EnsureIoOk(driver->ReadAt(rw_fd.fd, buf.data(), buf.size(), i * cfg.value_size), - "warmup ReadAt"); - } - } - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < values.size(); ++i) { - auto t0 = Clock::now(); - EnsureIoOk(driver->ReadAt(rw_fd.fd, buf.data(), buf.size(), i * cfg.value_size), "ReadAt"); - auto t1 = Clock::now(); - latencies.push_back(std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("Driver ReadAt", backend_variant, values.size() * cfg.measure_iters, - values.size() * cfg.measure_iters * cfg.value_size, - WallSeconds(run_start, run_end), latencies, results); - } - - // --- WriteBatch --- - { - ScopedFd batch_fd = OpenBenchFile(tmp.path + "/batch.bin"); - size_t total_batches = (values.size() + cfg.batch_size - 1) / cfg.batch_size; - std::vector> write_batches; - write_batches.reserve(total_batches); - for (size_t i = 0; i < values.size(); i += cfg.batch_size) { - size_t end = std::min(i + cfg.batch_size, values.size()); - std::vector ops; - ops.reserve(end - i); - for (size_t j = i; j < end; ++j) { - ops.push_back({batch_fd.fd, values[j].data(), values[j].size(), j * cfg.value_size}); - } - write_batches.push_back(std::move(ops)); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(values.size() * cfg.measure_iters); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (const auto& ops : write_batches) { - EnsureIoOk(driver->WriteBatch(ops), "warmup WriteBatch"); - } - } - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (const auto& ops : write_batches) { - auto t0 = Clock::now(); - EnsureIoOk(driver->WriteBatch(ops), "WriteBatch"); - auto t1 = Clock::now(); - tally.AddSample(ops.size(), ops.size(), ops.size() * cfg.value_size, - ops.size() * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("Driver WriteBatch", - JoinVariant({backend_variant, "bs=" + std::to_string(cfg.batch_size)}), tally, - WallSeconds(run_start, run_end), results); - } - - // --- ReadBatch --- - { - ScopedFd batch_fd = OpenBenchFile(tmp.path + "/read_batch.bin"); - for (size_t i = 0; i < values.size(); ++i) { - EnsureIoOk( - driver->WriteAt(batch_fd.fd, values[i].data(), values[i].size(), i * cfg.value_size), - "prefill ReadBatch"); - } - EnsureIoOk(driver->Sync(batch_fd.fd), "sync ReadBatch"); - - size_t total_batches = (values.size() + cfg.batch_size - 1) / cfg.batch_size; - std::vector> read_buffers(values.size(), std::vector(cfg.value_size)); - std::vector> read_batches; - read_batches.reserve(total_batches); - for (size_t i = 0; i < values.size(); i += cfg.batch_size) { - size_t end = std::min(i + cfg.batch_size, values.size()); - std::vector ops; - ops.reserve(end - i); - for (size_t j = i; j < end; ++j) { - ops.push_back( - {batch_fd.fd, read_buffers[j].data(), read_buffers[j].size(), j * cfg.value_size}); - } - read_batches.push_back(std::move(ops)); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(values.size() * cfg.measure_iters); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (const auto& ops : read_batches) { - EnsureIoOk(driver->ReadBatch(ops), "warmup ReadBatch"); - } - } - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (const auto& ops : read_batches) { - auto t0 = Clock::now(); - EnsureIoOk(driver->ReadBatch(ops), "ReadBatch"); - auto t1 = Clock::now(); - tally.AddSample(ops.size(), ops.size(), ops.size() * cfg.value_size, - ops.size() * cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("Driver ReadBatch", - JoinVariant({backend_variant, "bs=" + std::to_string(cfg.batch_size)}), tally, - WallSeconds(run_start, run_end), results); - } - - // --- Sync --- - { - ScopedFd sync_fd = OpenBenchFile(tmp.path + "/sync.bin"); - std::vector latencies; - latencies.reserve(values.size() * cfg.measure_iters); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < values.size(); ++i) { - EnsureIoOk(driver->WriteAt(sync_fd.fd, values[i].data(), values[i].size(), 0), - "warmup Sync write"); - EnsureIoOk(driver->Sync(sync_fd.fd), "warmup Sync"); - } - } - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t i = 0; i < values.size(); ++i) { - EnsureIoOk(driver->WriteAt(sync_fd.fd, values[i].data(), values[i].size(), 0), - "prepare Sync"); - auto t0 = Clock::now(); - EnsureIoOk(driver->Sync(sync_fd.fd), "Sync"); - auto t1 = Clock::now(); - latencies.push_back(std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("Driver Sync", backend_variant, values.size() * cfg.measure_iters, 0, - WallSeconds(run_start, run_end), latencies, results); - } - - // --- SyncMany --- - { - size_t sync_group_size = std::min(cfg.batch_size, values.size()); - std::vector sync_fds; - sync_fds.reserve(sync_group_size); - for (size_t i = 0; i < sync_group_size; ++i) { - sync_fds.push_back(OpenBenchFile(tmp.path + "/sync_many_" + std::to_string(i) + ".bin")); - } - - size_t total_groups = (values.size() + sync_group_size - 1) / sync_group_size; - bench::ResultTally tally; - tally.ReserveLatencySamples(values.size() * cfg.measure_iters); - - auto write_group = [&](size_t start_idx) { - size_t count = std::min(sync_group_size, values.size() - start_idx); - for (size_t i = 0; i < count; ++i) { - EnsureIoOk(driver->WriteAt(sync_fds[i].fd, values[start_idx + i].data(), - values[start_idx + i].size(), 0), - "prepare SyncMany"); - } - return count; - }; - - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t start = 0; start < values.size(); start += sync_group_size) { - size_t count = write_group(start); - std::vector fds; - fds.reserve(count); - for (size_t i = 0; i < count; ++i) fds.push_back(sync_fds[i].fd); - EnsureIoOk(driver->SyncMany(fds), "warmup SyncMany"); - } - } - - auto run_start = Clock::now(); - for (size_t m = 0; m < cfg.measure_iters; ++m) { - for (size_t start = 0; start < values.size(); start += sync_group_size) { - size_t count = write_group(start); - std::vector fds; - fds.reserve(count); - for (size_t i = 0; i < count; ++i) fds.push_back(sync_fds[i].fd); - auto t0 = Clock::now(); - EnsureIoOk(driver->SyncMany(fds), "SyncMany"); - auto t1 = Clock::now(); - tally.AddSample(count, count, 0, 0, - std::chrono::duration(t1 - t0).count()); - } - } - auto run_end = Clock::now(); - - RecordResult("Driver SyncMany", - JoinVariant({backend_variant, "fds=" + std::to_string(sync_group_size)}), tally, - WallSeconds(run_start, run_end), results); - } -} - -// --------------------------------------------------------------------------- -// H. Concurrent Scaling (StandaloneClient Put + Get) -// --------------------------------------------------------------------------- -static void BenchConcurrent(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "Concurrent")) return; - - PrintHeader("Concurrent Scaling"); - - for (int nthreads : cfg.thread_counts) { - ScopedTempDir tmp(cfg.base_dir + "/concurrent_" + std::to_string(nthreads)); - - UMBPConfig ucfg = MakeStandaloneClientConfig(cfg, tmp.path + "/ssd", cfg.dram_capacity); - - fs::create_directories(ucfg.ssd.storage_dir); - StandaloneClient client(ucfg); - - size_t keys_per_thread = keys.size() / static_cast(nthreads); - if (keys_per_thread == 0) continue; - - std::string variant = "threads=" + std::to_string(nthreads); - - // --- Put --- - { - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - client.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - client.Put(keys[i], values[i].data(), values[i].size()); - } - } - client.Clear(); - - std::vector thread_tallies(nthreads); - for (int t = 0; t < nthreads; ++t) { - thread_tallies[t].ReserveLatencySamples(keys_per_thread * cfg.measure_iters); - } - double elapsed_sec = 0.0; - - for (size_t m = 0; m < cfg.measure_iters; ++m) { - client.Clear(); - auto iter_start = Clock::now(); - std::vector threads; - for (int t = 0; t < nthreads; ++t) { - threads.emplace_back([&, t]() { - size_t start = t * keys_per_thread; - size_t end = start + keys_per_thread; - for (size_t i = start; i < end; ++i) { - auto t0 = Clock::now(); - bool ok = client.Put(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - thread_tallies[t].AddOp(ok, cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - }); - } - for (auto& th : threads) th.join(); - elapsed_sec += WallSeconds(iter_start, Clock::now()); - } - - bench::ResultTally tally; - for (const auto& thread_tally : thread_tallies) tally.Merge(thread_tally); - - RecordResult("Concurrent Put", variant, tally, elapsed_sec, results); - } - - // --- Get --- - { - // Ensure data is present - client.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - client.Put(keys[i], values[i].data(), values[i].size()); - } - - // Warmup - std::vector wbuf(cfg.value_size); - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - for (size_t i = 0; i < keys.size(); ++i) { - client.Get(keys[i], reinterpret_cast(wbuf.data()), wbuf.size()); - } - } - - std::vector thread_tallies(nthreads); - for (int t = 0; t < nthreads; ++t) { - thread_tallies[t].ReserveLatencySamples(keys_per_thread * cfg.measure_iters); - } - double elapsed_sec = 0.0; - - for (size_t m = 0; m < cfg.measure_iters; ++m) { - auto iter_start = Clock::now(); - std::vector threads; - for (int t = 0; t < nthreads; ++t) { - threads.emplace_back([&, t]() { - std::vector buf(cfg.value_size); - size_t start = t * keys_per_thread; - size_t end = start + keys_per_thread; - for (size_t i = start; i < end; ++i) { - auto t0 = Clock::now(); - bool ok = client.Get(keys[i], reinterpret_cast(buf.data()), buf.size()); - auto t1 = Clock::now(); - thread_tallies[t].AddOp(ok, cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - }); - } - for (auto& th : threads) th.join(); - elapsed_sec += WallSeconds(iter_start, Clock::now()); - } - - bench::ResultTally tally; - for (const auto& thread_tally : thread_tallies) tally.Merge(thread_tally); - - RecordResult("Concurrent Get", variant, tally, elapsed_sec, results); - } - } -} - -// --------------------------------------------------------------------------- -// I. Leader Mode: sync vs async copy -// --------------------------------------------------------------------------- -static void BenchLeaderMode(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "Leader")) return; - - PrintHeader("Leader Mode (sync vs async copy)"); - - auto run_mode = [&](bool async_copy, const std::string& label) { - ScopedTempDir tmp(cfg.base_dir + "/leader_" + label); - - UMBPConfig ucfg = MakeLeaderClientConfig(cfg, tmp.path + "/ssd", cfg.dram_capacity, async_copy); - - fs::create_directories(ucfg.ssd.storage_dir); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - StandaloneClient client(ucfg); - for (size_t i = 0; i < keys.size(); ++i) { - client.Put(keys[i], values[i].data(), values[i].size()); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(keys.size() * cfg.measure_iters); - double elapsed_sec = 0.0; - - for (size_t m = 0; m < cfg.measure_iters; ++m) { - auto iter_start = Clock::now(); - { - StandaloneClient client(ucfg); - for (size_t i = 0; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = client.Put(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, - std::chrono::duration(t1 - t0).count()); - } - } - elapsed_sec += WallSeconds(iter_start, Clock::now()); - } - - RecordResult("Leader Put", JoinVariant({label, "timing=end-to-end"}), tally, elapsed_sec, - results); - }; - - run_mode(false, "copy=sync"); - run_mode(true, "copy=async"); -} - -// --------------------------------------------------------------------------- -// J. Capacity Pressure -// --------------------------------------------------------------------------- -static void BenchCapacityPressure(const BenchConfig& cfg, const std::vector& keys, - const std::vector>& values, - std::vector& results) { - if (!ShouldRun(cfg, "Capacity")) return; - - PrintHeader("Capacity Pressure"); - - // DRAM capacity = 50% of total data => second half triggers eviction - size_t pressure_dram = keys.size() * cfg.value_size / 2; - size_t half = keys.size() / 2; - if (half == 0) return; - - ScopedTempDir tmp(cfg.base_dir + "/pressure"); - - UMBPConfig ucfg = MakeStandaloneClientConfig(cfg, tmp.path + "/ssd", pressure_dram); - - fs::create_directories(ucfg.ssd.storage_dir); - - // No pressure: first half, DRAM not full - { - StandaloneClient client(ucfg); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - client.Clear(); - for (size_t i = 0; i < half; ++i) { - client.Put(keys[i], values[i].data(), values[i].size()); - } - } - client.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(half * cfg.measure_iters); - double elapsed_sec = 0.0; - - for (size_t m = 0; m < cfg.measure_iters; ++m) { - client.Clear(); - auto iter_start = Clock::now(); - for (size_t i = 0; i < half; ++i) { - auto t0 = Clock::now(); - bool ok = client.Put(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - elapsed_sec += WallSeconds(iter_start, Clock::now()); - } - - RecordResult("Capacity Put", "phase=fits-in-dram/dram=50%", tally, elapsed_sec, results); - } - - // Under pressure: write all keys, second half triggers eviction - { - StandaloneClient client(ucfg); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - client.Clear(); - for (size_t i = 0; i < keys.size(); ++i) { - client.Put(keys[i], values[i].data(), values[i].size()); - } - } - client.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(half * cfg.measure_iters); - double elapsed_sec = 0.0; - - for (size_t m = 0; m < cfg.measure_iters; ++m) { - client.Clear(); - // Write first half to fill DRAM - for (size_t i = 0; i < half; ++i) { - client.Put(keys[i], values[i].data(), values[i].size()); - } - // Write second half — triggers eviction + demotion - auto iter_start = Clock::now(); - for (size_t i = half; i < keys.size(); ++i) { - auto t0 = Clock::now(); - bool ok = client.Put(keys[i], values[i].data(), values[i].size()); - auto t1 = Clock::now(); - tally.AddOp(ok, cfg.value_size, std::chrono::duration(t1 - t0).count()); - } - elapsed_sec += WallSeconds(iter_start, Clock::now()); - } - - RecordResult("Capacity Put", "phase=spills-to-ssd/dram=50%", tally, elapsed_sec, results); - } -} - -// --------------------------------------------------------------------------- -// K. E2E StandaloneClient Benchmark (sglang connector simulation) -// --------------------------------------------------------------------------- -static void BenchE2E(const BenchConfig& cfg, const E2EConfig& e2e, - std::vector& results) { - if (!ShouldRun(cfg, "E2E")) return; - - size_t value_size = e2e.ValueSizePerKey(); - size_t keys_per_page = e2e.KeysPerPage(); - std::string mode_str = (e2e.mode == E2EModelMode::MLA) ? "MLA" : "MHA"; - std::string variant_label = mode_str + "/" + std::to_string(e2e.batch_pages) + "pg/" + - std::to_string(value_size / 1024) + "KB"; - - PrintSectionTitle("E2E StandaloneClient (" + variant_label + ")"); - std::printf(" mode = %s\n", mode_str.c_str()); - std::printf(" num_layers = %zu\n", e2e.num_layers); - if (e2e.mode == E2EModelMode::MLA) { - std::printf(" kv_cache_dim = %zu (lora=%zu + rope=%zu)\n", e2e.KvCacheDim(), e2e.kv_lora_rank, - e2e.qk_rope_head_dim); - } else { - std::printf(" num_kv_heads = %zu, head_dim = %zu\n", e2e.num_kv_heads, e2e.head_dim); - } - std::printf(" kv_cache_dtype= %s (%zu bytes)\n", e2e.kv_cache_dtype.c_str(), e2e.DtypeSize()); - std::printf(" value_size = %zu bytes (%zu KB)\n", value_size, value_size / 1024); - std::printf(" keys_per_page = %zu\n", keys_per_page); - std::printf(" num_pages = %zu\n", e2e.num_pages); - std::printf(" batch_pages = %zu\n", e2e.batch_pages); - std::printf(" dedup_ratio = %.0f%%\n", e2e.dedup_ratio * 100); - - E2EKeyGenerator keygen{e2e.mode, 0, 0}; - E2EHostBuffer host_buf(e2e.num_pages, value_size, keys_per_page); - - // DRAM sized to hold all data with headroom. - size_t total_data = e2e.num_pages * keys_per_page * value_size; - - // E2E DRAM-only config: reuses cfg's SSD base settings (io_backend, durability, etc.) - // but disables SSD. If an SSD scenario is needed, MakeStandaloneClientConfig enables it. - auto MakeDramOnlyConfig = [&]() -> UMBPConfig { - UMBPConfig ucfg = MakeBaseSsdConfig(cfg); // inherit SSD params for consistency - ucfg.dram.capacity_bytes = total_data * 2; - ucfg.ssd.enabled = false; - ucfg.role = UMBPRole::Standalone; - ucfg.copy_pipeline.async_enabled = false; - ucfg.eviction.auto_promote_on_read = false; - return ucfg; - }; - - size_t batches_per_iter = (e2e.num_pages + e2e.batch_pages - 1) / e2e.batch_pages; - - // Auto-scale iters to keep percentile rows reasonably stable. - // Batch-style rows target >=100 batch samples; whole-cycle rows need a higher floor - // because p99 from only ~100 calls is still noisy. - // User --iters is treated as a minimum; we scale up if batches_per_iter is small. - constexpr size_t kMinBatchSamples = 100; - constexpr size_t kMinWholeCycleSamples = 200; - size_t e2e_iters = cfg.measure_iters; - if (batches_per_iter > 0 && batches_per_iter * e2e_iters < kMinBatchSamples) { - e2e_iters = (kMinBatchSamples + batches_per_iter - 1) / batches_per_iter; - } - size_t total_samples = batches_per_iter * e2e_iters; - size_t e2e_call_iters = std::max(e2e_iters, kMinWholeCycleSamples); - std::printf(" e2e_iters = %zu (auto-scaled from %zu; %zu batches/iter, %zu total samples)\n", - e2e_iters, cfg.measure_iters, batches_per_iter, total_samples); - if (e2e_call_iters != e2e_iters) { - std::printf(" e2e_call_iters= %zu (whole-cycle latency rows)\n", e2e_call_iters); - } - PrintTableHeader(OutputTableKind::Detail); - - // Helper: fill all pages into a client (untimed). - auto FillAll = [&](StandaloneClient& client) { - for (size_t b = 0; b < e2e.num_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, e2e.num_pages - b); - auto keys = keygen.KeysForPages(b, count); - std::vector ptrs; - std::vector sizes; - host_buf.GetBatchMeta(b, count, ptrs, sizes); - auto depths = GenerateDepths(e2e, b, count); - client.BatchPutWithDepth(keys, ptrs, sizes, depths); - } - }; - - // Helper: fill pages [0, num_pages) via BatchPut, collecting per-page metrics. - auto FillAllTimed = [&](StandaloneClient& client, size_t num_pages, bench::ResultTally& tally) { - for (size_t b = 0; b < num_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, num_pages - b); - auto keys = keygen.KeysForPages(b, count); - std::vector ptrs; - std::vector sizes; - host_buf.GetBatchMeta(b, count, ptrs, sizes); - auto depths = GenerateDepths(e2e, b, count); - auto t0 = Clock::now(); - auto batch_results = client.BatchPutWithDepth(keys, ptrs, sizes, depths); - auto t1 = Clock::now(); - size_t ok_pages = bench::CountSuccessfulPages(batch_results, keys_per_page); - tally.AddSample(count, ok_pages, count * keys_per_page * value_size, - ok_pages * keys_per_page * value_size, - std::chrono::duration(t1 - t0).count()); - } - }; - - auto FillAllTimedDedup = [&](StandaloneClient& client, size_t num_pages, size_t prefill_pages, - bench::ResultTally& tally) { - for (size_t b = 0; b < num_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, num_pages - b); - auto keys = keygen.KeysForPages(b, count); - std::vector ptrs; - std::vector sizes; - host_buf.GetBatchMeta(b, count, ptrs, sizes); - auto depths = GenerateDepths(e2e, b, count); - auto t0 = Clock::now(); - auto batch_results = client.BatchPutWithDepth(keys, ptrs, sizes, depths); - auto t1 = Clock::now(); - - size_t ok_pages = bench::CountSuccessfulPages(batch_results, keys_per_page); - size_t dedup_pages = 0; - if (b < prefill_pages) { - dedup_pages = std::min(count, prefill_pages - b); - } - size_t physical_write_pages = (ok_pages > dedup_pages) ? (ok_pages - dedup_pages) : 0; - tally.AddSample(count, ok_pages, count * keys_per_page * value_size, - physical_write_pages * keys_per_page * value_size, - std::chrono::duration(t1 - t0).count()); - } - }; - - // Helper: read all pages [0, num_pages) via BatchGet, collecting per-batch latencies. - std::vector read_buf; - std::vector read_ptrs; - std::vector read_sizes; - - auto ReadAllTimed = [&](StandaloneClient& client, size_t num_pages, bench::ResultTally& tally) { - for (size_t b = 0; b < num_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, num_pages - b); - auto keys = keygen.KeysForPages(b, count); - E2EHostBuffer::MakeReadMeta(count, keys_per_page, value_size, read_buf, read_ptrs, - read_sizes); - auto t0 = Clock::now(); - auto batch_results = client.BatchGet(keys, read_ptrs, read_sizes); - auto t1 = Clock::now(); - size_t ok_pages = bench::CountSuccessfulPages(batch_results, keys_per_page); - tally.AddSample(count, ok_pages, count * keys_per_page * value_size, - ok_pages * keys_per_page * value_size, - std::chrono::duration(t1 - t0).count()); - } - }; - - // Helper: read all pages [0, num_pages) without timing (for warmup). - auto ReadAll = [&](StandaloneClient& client, size_t num_pages) { - for (size_t b = 0; b < num_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, num_pages - b); - auto keys = keygen.KeysForPages(b, count); - E2EHostBuffer::MakeReadMeta(count, keys_per_page, value_size, read_buf, read_ptrs, - read_sizes); - client.BatchGet(keys, read_ptrs, read_sizes); - } - }; - - // --------------------------------------------------------------- - // (a) E2E BatchSet — fresh writes via BatchPutWithDepth - // --------------------------------------------------------------- - { - ScopedTempDir tmp(cfg.base_dir + "/e2e_batchset"); - UMBPConfig ucfg = MakeDramOnlyConfig(); - StandaloneClient client(ucfg); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - client.Clear(); - FillAll(client); - } - client.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e.num_pages * e2e_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_iters; ++m) { - client.Clear(); - FillAllTimed(client, e2e.num_pages, tally); - } - auto run_end = Clock::now(); - RecordResult("E2E BatchSet", variant_label, tally, WallSeconds(run_start, run_end), results); - } - - // --------------------------------------------------------------- - // (b) E2E BatchGet — read-back via BatchGet - // --------------------------------------------------------------- - { - ScopedTempDir tmp(cfg.base_dir + "/e2e_batchget"); - UMBPConfig ucfg = MakeDramOnlyConfig(); - StandaloneClient client(ucfg); - FillAll(client); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) ReadAll(client, e2e.num_pages); - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e.num_pages * e2e_iters); - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_iters; ++m) ReadAllTimed(client, e2e.num_pages, tally); - auto run_end = Clock::now(); - RecordResult("E2E BatchGet", variant_label, tally, WallSeconds(run_start, run_end), results); - } - - // --------------------------------------------------------------- - // (c) E2E ExistsScan — BatchExistsConsecutive with partial fill - // --------------------------------------------------------------- - { - ScopedTempDir tmp(cfg.base_dir + "/e2e_exists"); - UMBPConfig ucfg = MakeDramOnlyConfig(); - StandaloneClient client(ucfg); - - // Fill first half of pages to test early-stop behavior. - size_t fill_pages = e2e.num_pages / 2; - for (size_t b = 0; b < fill_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, fill_pages - b); - auto keys = keygen.KeysForPages(b, count); - std::vector ptrs; - std::vector sizes; - host_buf.GetBatchMeta(b, count, ptrs, sizes); - auto depths = GenerateDepths(e2e, b, count); - client.BatchPutWithDepth(keys, ptrs, sizes, depths); - } - - // Query all pages — consecutive hits stop at fill_pages boundary. - auto all_keys = keygen.KeysForPages(0, e2e.num_pages); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - client.BatchExistsConsecutive(all_keys); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e_call_iters); - size_t hit_pct = (e2e.num_pages == 0) ? 0 : ((fill_pages * 100) / e2e.num_pages); - - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_call_iters; ++m) { - auto t0 = Clock::now(); - client.BatchExistsConsecutive(all_keys); - auto t1 = Clock::now(); - tally.AddCall(/*requested=*/1, /*succeeded=*/1, /*req_bytes=*/0, /*ok_bytes=*/0, - std::chrono::duration(t1 - t0).count()); - } - auto run_end = Clock::now(); - RecordResult("E2E ExistsScan", - JoinVariant({variant_label, "hit=" + std::to_string(hit_pct) + "%"}), tally, - WallSeconds(run_start, run_end), results); - } - - // --------------------------------------------------------------- - // (d) E2E Dedup — batch_set with prefix reuse (dedup_ratio pre-filled) - // - // Each iteration: clear → re-fill only prefix pages → batch_set all pages. - // Pages [0, prefill) hit MayExist dedup; pages [prefill, num_pages) are fresh writes. - // --------------------------------------------------------------- - { - ScopedTempDir tmp(cfg.base_dir + "/e2e_dedup"); - UMBPConfig ucfg = MakeDramOnlyConfig(); - - size_t prefill_pages = static_cast(e2e.num_pages * e2e.dedup_ratio); - - // Helper: clear and re-seed prefix pages in a client. - auto SeedPrefix = [&](StandaloneClient& c) { - c.Clear(); - // FillAll variant with partial page count (untimed). - for (size_t b = 0; b < prefill_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, prefill_pages - b); - auto keys = keygen.KeysForPages(b, count); - std::vector ptrs; - std::vector sizes; - host_buf.GetBatchMeta(b, count, ptrs, sizes); - auto depths = GenerateDepths(e2e, b, count); - c.BatchPutWithDepth(keys, ptrs, sizes, depths); - } - }; - - // Warmup - { - StandaloneClient client(ucfg); - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - SeedPrefix(client); - FillAll(client); - } - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e.num_pages * e2e_iters); - - StandaloneClient client(ucfg); - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_iters; ++m) { - SeedPrefix(client); - FillAllTimedDedup(client, e2e.num_pages, prefill_pages, tally); - } - auto run_end = Clock::now(); - RecordResult( - "E2E Dedup", - JoinVariant({variant_label, - "dedup=" + std::to_string(static_cast(e2e.dedup_ratio * 100)) + "%", - "bytes=physical"}), - tally, WallSeconds(run_start, run_end), results); - } - - // --------------------------------------------------------------- - // (e) E2E PrefetchCycle — exists → get pipeline (sglang prefetch flow) - // --------------------------------------------------------------- - { - ScopedTempDir tmp(cfg.base_dir + "/e2e_prefetch"); - UMBPConfig ucfg = MakeDramOnlyConfig(); - StandaloneClient client(ucfg); - FillAll(client); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - auto all_keys = keygen.KeysForPages(0, e2e.num_pages); - client.BatchExistsConsecutive(all_keys); - ReadAll(client, e2e.num_pages); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e_call_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_call_iters; ++m) { - auto t0 = Clock::now(); - - // Step 1: exists check (determines how many pages to fetch). - auto all_keys = keygen.KeysForPages(0, e2e.num_pages); - size_t hit = client.BatchExistsConsecutive(all_keys); - size_t hit_pages = hit / keys_per_page; - - // Step 2: batch_get the hit pages in chunks (matching sglang flow). - // Note: untimed inner loop — entire prefetch cycle is timed as one op. - for (size_t b = 0; b < hit_pages; b += e2e.batch_pages) { - size_t count = std::min(e2e.batch_pages, hit_pages - b); - auto keys = keygen.KeysForPages(b, count); - E2EHostBuffer::MakeReadMeta(count, keys_per_page, value_size, read_buf, read_ptrs, - read_sizes); - client.BatchGet(keys, read_ptrs, read_sizes); - } - - auto t1 = Clock::now(); - tally.AddCall(/*requested=*/1, /*succeeded=*/1, - /*req_bytes=*/e2e.num_pages * keys_per_page * value_size, - /*ok_bytes=*/hit_pages * keys_per_page * value_size, - std::chrono::duration(t1 - t0).count()); - } - auto run_end = Clock::now(); - RecordResult("E2E PrefetchCycle", variant_label, tally, WallSeconds(run_start, run_end), - results); - } - - // --------------------------------------------------------------- - // (f) E2E Capacity Pressure — DRAM too small, eviction spills to SSD, - // then BatchGet reads back from SSD. - // - // DRAM = 50% of total data → second half triggers eviction. - // Measures: write under pressure (set), then read-back (get from SSD). - // --------------------------------------------------------------- - // SSD scenarios — always run when BenchE2E is entered (gated by "E2E" filter above). - { - ScopedTempDir tmp(cfg.base_dir + "/e2e_capacity"); - - UMBPConfig ucfg = - MakeStandaloneClientConfig(cfg, tmp.path + "/ssd", total_data / 2, total_data * 2); - fs::create_directories(ucfg.ssd.storage_dir); - - // --- Write under capacity pressure --- - { - StandaloneClient client(ucfg); - - // Warmup - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - client.Clear(); - FillAll(client); - } - client.Clear(); - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e.num_pages * e2e_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_iters; ++m) { - client.Clear(); - FillAllTimed(client, e2e.num_pages, tally); - } - auto run_end = Clock::now(); - RecordResult("E2E Capacity Put", - JoinVariant({variant_label, "phase=spills-to-ssd", "dram=50%"}), tally, - WallSeconds(run_start, run_end), results); - } - - // --- Read-back (early pages evicted to SSD) --- - { - StandaloneClient client(ucfg); - FillAll(client); // first half evicted to SSD, second half in DRAM - - for (size_t w = 0; w < cfg.warmup_iters; ++w) ReadAll(client, e2e.num_pages); - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e.num_pages * e2e_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_iters; ++m) ReadAllTimed(client, e2e.num_pages, tally); - auto run_end = Clock::now(); - RecordResult("E2E Capacity Get", - JoinVariant({variant_label, "phase=read-after-spill", "dram=50%"}), tally, - WallSeconds(run_start, run_end), results); - } - } - - // --------------------------------------------------------------- - // (g) E2E Leader Mode — SharedSSDLeader with async copy pipeline. - // - // Mirrors sglang MLA + TP>1 deployment: - // BatchPutWithDepth writes to DRAM, CopyPipeline async-copies to SSD. - // Compares sync vs async copy throughput. - // --------------------------------------------------------------- - { - auto run_leader = [&](bool async_copy, const std::string& label) { - ScopedTempDir tmp(cfg.base_dir + "/e2e_leader_" + label); - - UMBPConfig ucfg = MakeLeaderClientConfig(cfg, tmp.path + "/ssd", total_data * 2, async_copy, - total_data * 2); - fs::create_directories(ucfg.ssd.storage_dir); - - // Warmup — construct+destroy client so destructor drains async queue. - for (size_t w = 0; w < cfg.warmup_iters; ++w) { - StandaloneClient client(ucfg); - FillAll(client); - } - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e.num_pages * e2e_iters); - - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_iters; ++m) { - bench::ResultTally iter_tally; - auto iter_start = Clock::now(); - { - StandaloneClient client(ucfg); // fresh client per iter (destructor drains async) - FillAllTimed(client, e2e.num_pages, iter_tally); - } - auto iter_end = Clock::now(); - tally.AddSample(iter_tally.requested_ops, iter_tally.successful_ops, - iter_tally.requested_bytes, iter_tally.successful_bytes, - std::chrono::duration(iter_end - iter_start).count()); - } - auto run_end = Clock::now(); - RecordResult("E2E Leader Set", JoinVariant({variant_label, label, "timing=end-to-end"}), - tally, WallSeconds(run_start, run_end), results); - }; - - run_leader(false, "copy=sync"); - run_leader(true, "copy=async"); - } - - // --------------------------------------------------------------- - // (h) E2E Follower Get — pure SSD read via SharedSSDFollower. - // - // Mirrors sglang MLA + TP>1, rank>0: - // Leader writes all pages to shared SSD, Follower reads from SSD only. - // Measures pure SSD read throughput (no DRAM hits). - // --------------------------------------------------------------- - { - ScopedTempDir tmp(cfg.base_dir + "/e2e_follower"); - std::string ssd_dir = tmp.path + "/ssd"; - - // Step 1: Leader writes all pages to SSD. - { - UMBPConfig leader_cfg = - MakeLeaderClientConfig(cfg, ssd_dir, total_data * 2, false, total_data * 2); - fs::create_directories(leader_cfg.ssd.storage_dir); - StandaloneClient leader(leader_cfg); - FillAll(leader); - // Leader destructor drains copy pipeline, ensuring all data is on SSD. - } - - // Step 2: Follower reads from SSD. - UMBPConfig follower_cfg = - MakeFollowerClientConfig(cfg, ssd_dir, total_data * 2, total_data * 2); - StandaloneClient follower(follower_cfg); - - for (size_t w = 0; w < cfg.warmup_iters; ++w) ReadAll(follower, e2e.num_pages); - - bench::ResultTally tally; - tally.ReserveLatencySamples(e2e.num_pages * e2e_iters); - auto run_start = Clock::now(); - for (size_t m = 0; m < e2e_iters; ++m) ReadAllTimed(follower, e2e.num_pages, tally); - auto run_end = Clock::now(); - RecordResult("E2E Follower Get", JoinVariant({variant_label, "source=ssd-only"}), tally, - WallSeconds(run_start, run_end), results); - } -} - -// --------------------------------------------------------------------------- -// CLI parsing & profiles -// --------------------------------------------------------------------------- -static void ApplyProfile(BenchConfig& cfg, E2EConfig& e2e, const std::string& profile) { - if (profile == "small") { - cfg.num_keys = 200; - cfg.value_size = 1024; - cfg.batch_size = 16; - cfg.dram_capacity = 4ULL * 1024 * 1024; - cfg.ssd_capacity = 16ULL * 1024 * 1024; - cfg.segment_size = 4ULL * 1024 * 1024; - cfg.thread_counts = {1, 2}; - // E2E: DeepSeek-V3, small scale - ApplyModelPreset(e2e, "deepseek-v3"); - e2e.num_pages = 64; - e2e.batch_pages = 16; - } else if (profile == "medium") { - cfg.num_keys = 1000; - cfg.value_size = 4096; - cfg.batch_size = 64; - cfg.dram_capacity = 64ULL * 1024 * 1024; - cfg.ssd_capacity = 256ULL * 1024 * 1024; - cfg.segment_size = 64ULL * 1024 * 1024; - cfg.thread_counts = {1, 2, 4, 8}; - // E2E: DeepSeek-V3, default scale - ApplyModelPreset(e2e, "deepseek-v3"); - e2e.num_pages = 512; - e2e.batch_pages = 128; - } else if (profile == "large") { - cfg.num_keys = 10000; - cfg.value_size = 64 * 1024; - cfg.batch_size = 128; - cfg.dram_capacity = 512ULL * 1024 * 1024; - cfg.ssd_capacity = 2ULL * 1024 * 1024 * 1024; - cfg.segment_size = 256ULL * 1024 * 1024; - cfg.thread_counts = {1, 2, 4, 8}; - // E2E: DeepSeek-V3, large scale - ApplyModelPreset(e2e, "deepseek-v3"); - e2e.num_pages = 2048; - e2e.batch_pages = 128; - } else { - std::cerr << "Unknown profile: " << profile << std::endl; - std::exit(1); - } -} - -static void PrintUsage(const char* argv0) { - std::printf( - "Usage: %s [OPTIONS]\n" - "\n" - "General:\n" - " --profile Preset config (default: medium)\n" - " --list-scenarios Print available scenario filters and exit\n" - " --num-keys N Keys per scenario\n" - " --value-size N Value size in bytes\n" - " --batch-size N Batch size\n" - " --warmup-iters N Warmup iterations\n" - " --iters N Measurement iterations (default: 10)\n" - " --filter SUBSTRING Run only matching scenarios\n" - " --dir PATH Temp directory path\n" - " -h, --help Help\n" - "\n" - "SSD / driver:\n" - " --dram-capacity N DRAM capacity in bytes\n" - " --ssd-capacity N SSD capacity in bytes\n" - " --segment-size N SSD segment size in bytes\n" - " --ssd-io-backend \n" - " --ssd-queue-depth N Storage I/O queue depth\n" - " --ssd-durability \n" - " --ssd-backend SSD backend (default: file)\n" - " spdk requires UMBP_SPDK_NVME_PCI env var\n" - "\n" - "E2E (sglang connector simulation):\n" - " --model \n" - " --mode Override model attention mode\n" - " --num-layers N Transformer layers\n" - " --num-kv-heads N KV heads (MHA only)\n" - " --head-dim N Head dimension (MHA only)\n" - " --kv-lora-rank N LoRA rank (MLA only)\n" - " --qk-rope-head-dim N RoPE head dim (MLA only)\n" - " --kv-cache-dtype \n" - " --page-size N Tokens per page\n" - " --num-pages N Total pages for E2E test\n" - " --batch-pages N Pages per batch call\n" - " --dedup-ratio F Prefix reuse ratio (0.0-1.0)\n", - argv0); -} - -struct ParsedArgs { - BenchConfig cfg; - E2EConfig e2e; -}; - -static ParsedArgs ParseArgs(int argc, char* argv[]) { - BenchConfig cfg; - E2EConfig e2e; - std::string profile = "medium"; - std::string model_preset; - - // Track which fields the user explicitly overrides - bool override_num_keys = false; - bool override_value_size = false; - bool override_batch_size = false; - bool override_warmup_iters = false; - bool override_dram_capacity = false; - bool override_ssd_capacity = false; - bool override_segment_size = false; - bool override_ssd_io_backend = false; - bool override_ssd_io_queue_depth = false; - bool override_ssd_durability = false; - - size_t user_num_keys = 0; - size_t user_value_size = 0; - size_t user_batch_size = 0; - size_t user_warmup_iters = 0; - size_t user_dram_capacity = 0; - size_t user_ssd_capacity = 0; - size_t user_segment_size = 0; - UMBPIoBackend user_ssd_io_backend = UMBPIoBackend::Posix; - size_t user_ssd_io_queue_depth = 0; - UMBPDurabilityMode user_ssd_durability = UMBPDurabilityMode::Relaxed; - - // E2E overrides - bool override_mode = false; - bool override_num_layers = false; - bool override_num_kv_heads = false; - bool override_head_dim = false; - bool override_kv_lora_rank = false; - bool override_qk_rope = false; - bool override_kv_cache_dtype = false; - bool override_page_size = false; - bool override_num_pages = false; - bool override_batch_pages = false; - bool override_dedup_ratio = false; - - E2EModelMode user_mode{}; - size_t user_num_layers = 0, user_num_kv_heads = 0, user_head_dim = 0; - size_t user_kv_lora_rank = 0, user_qk_rope = 0; - std::string user_kv_cache_dtype; - size_t user_page_size = 0, user_num_pages = 0, user_batch_pages = 0; - double user_dedup_ratio = 0.0; - - for (int i = 1; i < argc; ++i) { - std::string arg(argv[i]); - if (arg == "-h" || arg == "--help") { - PrintUsage(argv[0]); - std::exit(0); - } else if (arg == "--list-scenarios") { - cfg.list_scenarios = true; - } else if (arg == "--profile" && i + 1 < argc) { - profile = argv[++i]; - } else if (arg == "--num-keys" && i + 1 < argc) { - user_num_keys = std::stoull(argv[++i]); - override_num_keys = true; - } else if (arg == "--value-size" && i + 1 < argc) { - user_value_size = std::stoull(argv[++i]); - override_value_size = true; - } else if (arg == "--batch-size" && i + 1 < argc) { - user_batch_size = std::stoull(argv[++i]); - override_batch_size = true; - } else if (arg == "--warmup-iters" && i + 1 < argc) { - user_warmup_iters = std::stoull(argv[++i]); - override_warmup_iters = true; - } else if (arg == "--iters" && i + 1 < argc) { - cfg.measure_iters = std::stoull(argv[++i]); - } else if (arg == "--filter" && i + 1 < argc) { - cfg.filter = argv[++i]; - } else if (arg == "--dir" && i + 1 < argc) { - cfg.base_dir = argv[++i]; - } else if (arg == "--dram-capacity" && i + 1 < argc) { - user_dram_capacity = std::stoull(argv[++i]); - override_dram_capacity = true; - } else if (arg == "--ssd-capacity" && i + 1 < argc) { - user_ssd_capacity = std::stoull(argv[++i]); - override_ssd_capacity = true; - } else if (arg == "--segment-size" && i + 1 < argc) { - user_segment_size = std::stoull(argv[++i]); - override_segment_size = true; - } else if (arg == "--ssd-io-backend" && i + 1 < argc) { - std::string backend = argv[++i]; - if (!ParseIoBackend(backend, user_ssd_io_backend)) { - std::cerr << "Error: --ssd-io-backend must be one of: posix, io_uring" - << " (got '" << backend << "')\n"; - std::exit(1); - } - override_ssd_io_backend = true; - } else if (arg == "--ssd-queue-depth" && i + 1 < argc) { - user_ssd_io_queue_depth = std::stoull(argv[++i]); - override_ssd_io_queue_depth = true; - } else if (arg == "--ssd-durability" && i + 1 < argc) { - std::string durability = argv[++i]; - if (!ParseDurabilityMode(durability, user_ssd_durability)) { - std::cerr << "Error: --ssd-durability must be 'strict' or 'relaxed', got '" << durability - << "'\n"; - std::exit(1); - } - override_ssd_durability = true; - } else if (arg == "--ssd-backend" && i + 1 < argc) { - std::string val = argv[++i]; - std::string lower = ToLower(val); - if (lower != "file" && lower != "spdk") { - std::cerr << "Error: --ssd-backend must be 'file' or 'spdk'\n"; - std::exit(1); - } - cfg.ssd_backend = lower; - // E2E flags - } else if (arg == "--model" && i + 1 < argc) { - model_preset = argv[++i]; - } else if (arg == "--mode" && i + 1 < argc) { - std::string m = argv[++i]; - if (m == "mla") { - user_mode = E2EModelMode::MLA; - } else if (m == "mha") { - user_mode = E2EModelMode::MHA; - } else { - std::cerr << "Error: --mode must be 'mha' or 'mla', got '" << m << "'\n"; - std::exit(1); - } - override_mode = true; - } else if (arg == "--num-layers" && i + 1 < argc) { - user_num_layers = std::stoull(argv[++i]); - override_num_layers = true; - } else if (arg == "--num-kv-heads" && i + 1 < argc) { - user_num_kv_heads = std::stoull(argv[++i]); - override_num_kv_heads = true; - } else if (arg == "--head-dim" && i + 1 < argc) { - user_head_dim = std::stoull(argv[++i]); - override_head_dim = true; - } else if (arg == "--kv-lora-rank" && i + 1 < argc) { - user_kv_lora_rank = std::stoull(argv[++i]); - override_kv_lora_rank = true; - } else if (arg == "--qk-rope-head-dim" && i + 1 < argc) { - user_qk_rope = std::stoull(argv[++i]); - override_qk_rope = true; - } else if (arg == "--kv-cache-dtype" && i + 1 < argc) { - user_kv_cache_dtype = argv[++i]; - override_kv_cache_dtype = true; - } else if (arg == "--page-size" && i + 1 < argc) { - user_page_size = std::stoull(argv[++i]); - override_page_size = true; - } else if (arg == "--num-pages" && i + 1 < argc) { - user_num_pages = std::stoull(argv[++i]); - override_num_pages = true; - } else if (arg == "--batch-pages" && i + 1 < argc) { - user_batch_pages = std::stoull(argv[++i]); - override_batch_pages = true; - } else if (arg == "--dedup-ratio" && i + 1 < argc) { - user_dedup_ratio = std::stod(argv[++i]); - override_dedup_ratio = true; - } else { - std::cerr << "Unknown option: " << arg << std::endl; - PrintUsage(argv[0]); - std::exit(1); - } - } - - // Apply profile first (sets both BenchConfig and E2EConfig defaults). - ApplyProfile(cfg, e2e, profile); - - // Override BenchConfig fields. - if (override_num_keys) cfg.num_keys = user_num_keys; - if (override_value_size) cfg.value_size = user_value_size; - if (override_batch_size) cfg.batch_size = user_batch_size; - if (override_warmup_iters) cfg.warmup_iters = user_warmup_iters; - if (override_dram_capacity) cfg.dram_capacity = user_dram_capacity; - if (override_ssd_capacity) cfg.ssd_capacity = user_ssd_capacity; - if (override_segment_size) cfg.segment_size = user_segment_size; - if (override_ssd_io_backend) cfg.ssd_io_backend = user_ssd_io_backend; - if (override_ssd_io_queue_depth) cfg.ssd_io_queue_depth = user_ssd_io_queue_depth; - if (override_ssd_durability) cfg.ssd_durability_mode = user_ssd_durability; - - // Apply model preset (overrides E2E model params from profile). - if (!model_preset.empty()) ApplyModelPreset(e2e, model_preset); - - // Override individual E2E fields. - if (override_mode) e2e.mode = user_mode; - if (override_num_layers) e2e.num_layers = user_num_layers; - if (override_num_kv_heads) e2e.num_kv_heads = user_num_kv_heads; - if (override_head_dim) e2e.head_dim = user_head_dim; - if (override_kv_lora_rank) e2e.kv_lora_rank = user_kv_lora_rank; - if (override_qk_rope) e2e.qk_rope_head_dim = user_qk_rope; - if (override_kv_cache_dtype) e2e.kv_cache_dtype = user_kv_cache_dtype; - if (override_page_size) e2e.page_size = user_page_size; - if (override_num_pages) e2e.num_pages = user_num_pages; - if (override_batch_pages) e2e.batch_pages = user_batch_pages; - if (override_dedup_ratio) e2e.dedup_ratio = user_dedup_ratio; - - // --- Input validation --- - if (cfg.num_keys == 0) { - std::cerr << "Error: --num-keys must be > 0\n"; - std::exit(1); - } - if (cfg.value_size == 0) { - std::cerr << "Error: --value-size must be > 0\n"; - std::exit(1); - } - if (cfg.batch_size == 0) { - std::cerr << "Error: --batch-size must be > 0\n"; - std::exit(1); - } - if (cfg.dram_capacity == 0) { - std::cerr << "Error: --dram-capacity must be > 0\n"; - std::exit(1); - } - if (cfg.ssd_capacity == 0) { - std::cerr << "Error: --ssd-capacity must be > 0\n"; - std::exit(1); - } - if (cfg.segment_size == 0) { - std::cerr << "Error: --segment-size must be > 0\n"; - std::exit(1); - } - if (cfg.ssd_io_queue_depth == 0) { - std::cerr << "Error: --ssd-queue-depth must be > 0\n"; - std::exit(1); - } - if (cfg.ssd_io_queue_depth > std::numeric_limits::max()) { - std::cerr << "Error: --ssd-queue-depth must fit in uint32_t\n"; - std::exit(1); - } - if (e2e.batch_pages == 0) { - std::cerr << "Error: --batch-pages must be > 0\n"; - std::exit(1); - } - if (e2e.num_pages == 0) { - std::cerr << "Error: --num-pages must be > 0\n"; - std::exit(1); - } - if (e2e.dedup_ratio < 0.0 || e2e.dedup_ratio > 1.0) { - std::cerr << "Error: --dedup-ratio must be in [0.0, 1.0]\n"; - std::exit(1); - } - if (e2e.page_size == 0) { - std::cerr << "Error: --page-size must be > 0\n"; - std::exit(1); - } - { - const auto& d = e2e.kv_cache_dtype; - if (d != "bf16" && d != "fp16" && d != "fp8" && d != "fp8_e4m3" && d != "fp8_e5m2") { - std::cerr << "Error: --kv-cache-dtype must be one of: bf16, fp16, fp8, fp8_e4m3, fp8_e5m2" - << " (got '" << d << "')\n"; - std::exit(1); - } - } - - return {cfg, e2e}; -} - -// --------------------------------------------------------------------------- -// main -// --------------------------------------------------------------------------- -int main(int argc, char* argv[]) { - auto [cfg, e2e] = ParseArgs(argc, argv); - if (cfg.list_scenarios) { - PrintScenarioList(); - return 0; - } - - std::printf("UMBP Benchmark\n"); - std::printf(" num_keys = %zu\n", cfg.num_keys); - std::printf(" value_size = %zu bytes\n", cfg.value_size); - std::printf(" batch_size = %zu\n", cfg.batch_size); - std::printf(" warmup_iters = %zu\n", cfg.warmup_iters); - std::printf(" measure_iters= %zu\n", cfg.measure_iters); - std::printf(" dram_capacity= %zu bytes\n", cfg.dram_capacity); - std::printf(" ssd_capacity = %zu bytes\n", cfg.ssd_capacity); - std::printf(" segment_size = %zu bytes\n", cfg.segment_size); - std::printf(" ssd_io_back = %s\n", GetIoBackendSpec(cfg.ssd_io_backend).display_name); - std::printf(" ssd_backend = %s\n", cfg.ssd_backend.c_str()); - std::printf(" ssd_queue_d = %zu\n", cfg.ssd_io_queue_depth); - std::printf(" ssd_durability = %s\n", DurabilityLabel(cfg.ssd_durability_mode)); - std::printf(" base_dir = %s\n", cfg.base_dir.c_str()); - if (!cfg.filter.empty()) { - std::printf(" filter = %s\n", cfg.filter.c_str()); - } - PrintConfigWarnings(cfg); - std::printf(" threads ="); - for (int t : cfg.thread_counts) std::printf(" %d", t); - std::printf("\n"); - - // Pre-generate data - std::printf("\nGenerating %zu keys x %zu bytes...\n", cfg.num_keys, cfg.value_size); - auto keys = GenerateKeys(cfg.num_keys); - auto values = GenerateValues(cfg.num_keys, cfg.value_size); - std::printf("Data generation complete.\n"); - - std::vector results; - - // SPDK anchor — keeps the proxy daemon alive for the entire benchmark run. - // All subsequent StandaloneClient instances probe the existing proxy and reuse it - // instead of spawning new ones (see LocalStorageManager proxy probe logic). - std::unique_ptr spdk_anchor; - TierBackend* ext_ssd = nullptr; - if (IsSpdk(cfg)) { - const char* pci = std::getenv("UMBP_SPDK_NVME_PCI"); - if (!pci || !pci[0]) { - std::fprintf(stderr, - "ERROR: --ssd-backend spdk requires UMBP_SPDK_NVME_PCI env var.\n" - " Example: UMBP_SPDK_NVME_PCI=0000:c1:00.0 ./bench_umbp --ssd-backend spdk\n" - " Find NVMe PCI addresses with: lspci | grep -i nvme\n"); - return 1; - } - std::printf("\nInitializing SPDK backend (PCI: %s)...\n", pci); - UMBPConfig acfg = MakeBaseSsdConfig(cfg); - acfg.dram.capacity_bytes = cfg.dram_capacity; - acfg.role = UMBPRole::SharedSSDLeader; - acfg.copy_pipeline.async_enabled = false; - acfg.eviction.auto_promote_on_read = false; - spdk_anchor = std::make_unique(acfg); - ext_ssd = spdk_anchor->Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ext_ssd || !dynamic_cast(ext_ssd)) { - std::fprintf(stderr, - "ERROR: SPDK SSD tier not available. " - "Check UMBP_SPDK_NVME_PCI env var.\n"); - return 1; - } - std::printf("SPDK backend initialized.\n"); - } - - // Phase 1: tier-level benchmarks - BenchDRAMTier(cfg, keys, values, results); - BenchSSDTier(cfg, keys, values, results, ext_ssd); - BenchBatchWrite(cfg, keys, values, results, ext_ssd); - BenchBatchRead(cfg, keys, values, results, ext_ssd); - - // Phase 2: POSIX-specific (auto-skipped for SPDK inside each function) - BenchIOBackend(cfg, keys, values, results); - BenchDurability(cfg, keys, values, results); - BenchStorageIoDriver(cfg, values, results); - - // Phase 3: client-level — each creates own StandaloneClient that reuses the - // anchor's proxy (probed as alive, no respawn needed). - BenchCopyToSSD(cfg, keys, values, results); - BenchConcurrent(cfg, keys, values, results); - BenchLeaderMode(cfg, keys, values, results); - BenchCapacityPressure(cfg, keys, values, results); - BenchE2E(cfg, e2e, results); - - // Anchor destroyed here (after all benchmarks) — proxy shuts down. - - if (results.empty()) { - if (!cfg.filter.empty()) { - std::fprintf(stderr, - "Warning: filter '%s' matched no benchmark scenarios. Use --list-scenarios.\n", - cfg.filter.c_str()); - } - return 1; - } - - // Summary - std::printf("\n=== Summary ===\n"); - PrintTableHeader(OutputTableKind::Summary); - for (const auto& r : results) { - PrintSummaryResult(r); - } - - std::printf("\nDone. %zu benchmarks completed.\n", results.size()); - return 0; -} diff --git a/tests/cpp/umbp/local/test_dummy_ssd_tier.cpp b/tests/cpp/umbp/local/test_dummy_ssd_tier.cpp index b54deece1..56fcae805 100644 --- a/tests/cpp/umbp/local/test_dummy_ssd_tier.cpp +++ b/tests/cpp/umbp/local/test_dummy_ssd_tier.cpp @@ -27,7 +27,6 @@ #include #include "umbp/local/tiers/dummy_ssd_tier.h" -#include "umbp/local/tiers/local_storage_manager.h" using namespace mori::umbp; @@ -365,215 +364,6 @@ void test_batch_read_partial_miss() { std::cout << "PASSED" << std::endl; } -// ============================================================ -// Integration tests: DummySsdTier via LocalStorageManager -// ============================================================ - -static UMBPConfig MakeDummyConfig(size_t dram_bytes, size_t ssd_bytes) { - UMBPConfig cfg; - cfg.dram.capacity_bytes = dram_bytes; - cfg.ssd.enabled = true; - cfg.ssd.ssd_backend = "dummy_storage"; - cfg.ssd.capacity_bytes = ssd_bytes; - return cfg; -} - -void test_lsm_dummy_write_read() { - std::cout << "test_lsm_dummy_write_read... "; - auto cfg = MakeDummyConfig(1024 * 1024, 10 * 1024 * 1024); - LocalStorageManager mgr(cfg); - - std::vector data(4096, 'S'); - assert(mgr.Write("lsm_dk", data.data(), data.size(), StorageTier::LOCAL_SSD)); - assert(mgr.Exists("lsm_dk")); - - // ReadIntoPtr succeeds (dummy returns true, does not write buffer) - std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("lsm_dk", reinterpret_cast(buf.data()), buf.size())); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_lsm_dummy_capacity() { - std::cout << "test_lsm_dummy_capacity... "; - auto cfg = MakeDummyConfig(1024 * 1024, 5 * 1024 * 1024); - LocalStorageManager mgr(cfg); - - auto [used, total] = mgr.Capacity(StorageTier::LOCAL_SSD); - assert(used == 0); - assert(total == 5 * 1024 * 1024); - - std::vector data(1024, 'C'); - mgr.Write("cap_dk", data.data(), data.size(), StorageTier::LOCAL_SSD); - - auto [used2, _] = mgr.Capacity(StorageTier::LOCAL_SSD); - assert(used2 == 1024); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_lsm_dummy_dram_auto_demote() { - std::cout << "test_lsm_dummy_dram_auto_demote... "; - auto cfg = MakeDummyConfig(1024, 10 * 1024 * 1024); - LocalBlockIndex index; - LocalStorageManager mgr(cfg, &index); - - std::vector d1(512, 'A'); - assert(mgr.Write("dk1", d1.data(), d1.size())); - index.Insert("dk1", {StorageTier::CPU_DRAM, 0, 512}); - - std::vector d2(512, 'B'); - assert(mgr.Write("dk2", d2.data(), d2.size())); - index.Insert("dk2", {StorageTier::CPU_DRAM, 0, 512}); - - // DRAM full → dk1 auto-demoted to dummy SSD - std::vector d3(512, 'C'); - assert(mgr.Write("dk3", d3.data(), d3.size())); - index.Insert("dk3", {StorageTier::CPU_DRAM, 0, 512}); - - // dk1 still exists (demoted to dummy SSD) - assert(mgr.Exists("dk1")); - - auto loc = index.Lookup("dk1"); - assert(loc.has_value()); - assert(loc->tier == StorageTier::LOCAL_SSD); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_lsm_dummy_ssd_full_eviction() { - std::cout << "test_lsm_dummy_ssd_full_eviction... "; - auto cfg = MakeDummyConfig(1024 * 1024, 1024); - LocalBlockIndex index; - LocalStorageManager mgr(cfg, &index); - - std::vector d1(512, 'A'), d2(512, 'B'); - assert(mgr.Write("se1", d1.data(), d1.size(), StorageTier::LOCAL_SSD)); - index.Insert("se1", {StorageTier::LOCAL_SSD, 0, 512}); - assert(mgr.Write("se2", d2.data(), d2.size(), StorageTier::LOCAL_SSD)); - index.Insert("se2", {StorageTier::LOCAL_SSD, 0, 512}); - - // SSD full → se1 (LRU) gets evicted to make room for se3 - std::vector d3(512, 'C'); - assert(mgr.Write("se3", d3.data(), d3.size(), StorageTier::LOCAL_SSD)); - - assert(!mgr.Exists("se1")); - assert(mgr.Exists("se2")); - assert(mgr.Exists("se3")); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_lsm_dummy_demote_and_promote() { - std::cout << "test_lsm_dummy_demote_and_promote... "; - auto cfg = MakeDummyConfig(1024 * 1024, 10 * 1024 * 1024); - cfg.eviction.auto_promote_on_read = false; - LocalBlockIndex index; - LocalStorageManager mgr(cfg, &index); - - std::vector data(4096, 'D'); - assert(mgr.Write("dp_dk", data.data(), data.size())); - index.Insert("dp_dk", {StorageTier::CPU_DRAM, 0, 4096}); - - // Demote DRAM → dummy SSD - assert(mgr.Demote("dp_dk")); - assert(mgr.Exists("dp_dk")); - auto loc1 = index.Lookup("dp_dk"); - assert(loc1.has_value()); - assert(loc1->tier == StorageTier::LOCAL_SSD); - - // Promote dummy SSD → DRAM - assert(mgr.Promote("dp_dk")); - auto loc2 = index.Lookup("dp_dk"); - assert(loc2.has_value()); - assert(loc2->tier == StorageTier::CPU_DRAM); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_lsm_dummy_batch_write_and_read() { - std::cout << "test_lsm_dummy_batch_write_and_read... "; - auto cfg = MakeDummyConfig(1024 * 1024, 10 * 1024 * 1024); - LocalStorageManager mgr(cfg); - - constexpr size_t kN = 8; - constexpr size_t kSize = 2048; - - std::vector keys(kN); - std::vector> payloads(kN); - std::vector ptrs(kN); - std::vector sizes(kN, kSize); - - for (size_t i = 0; i < kN; ++i) { - keys[i] = "lsm_batch_" + std::to_string(i); - payloads[i].assign(kSize, static_cast('a' + i)); - ptrs[i] = payloads[i].data(); - } - - auto wr = mgr.BatchWrite(keys, ptrs, sizes, StorageTier::LOCAL_SSD); - for (size_t i = 0; i < kN; ++i) { - assert(wr[i]); - } - - std::vector> bufs(kN, std::vector(kSize, 0)); - std::vector dst_ptrs(kN); - for (size_t i = 0; i < kN; ++i) { - dst_ptrs[i] = reinterpret_cast(bufs[i].data()); - } - - auto rd = mgr.ReadBatchIntoPtr(keys, dst_ptrs, sizes); - for (size_t i = 0; i < kN; ++i) { - assert(rd[i]); - } - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_lsm_dummy_copy_to_ssd() { - std::cout << "test_lsm_dummy_copy_to_ssd... "; - auto cfg = MakeDummyConfig(1024 * 1024, 10 * 1024 * 1024); - LocalStorageManager mgr(cfg); - - // Write to DRAM first - std::vector data(4096, 'T'); - assert(mgr.Write("copy_dk", data.data(), data.size(), StorageTier::CPU_DRAM)); - - // CopyToSSD should succeed with dummy tier - assert(mgr.CopyToSSD("copy_dk")); - - // Key should exist on SSD tier now - auto* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - assert(ssd != nullptr); - assert(ssd->Exists("copy_dk")); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_lsm_dummy_flush() { - std::cout << "test_lsm_dummy_flush... "; - auto cfg = MakeDummyConfig(1024 * 1024, 10 * 1024 * 1024); - LocalStorageManager mgr(cfg); - - std::vector data(1024, 'F'); - mgr.Write("fl_dk", data.data(), data.size(), StorageTier::LOCAL_SSD); - - assert(mgr.Flush()); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -// ============================================================ -// Main -// ============================================================ - int main() { std::cout << "=== DummySsdTier Unit Tests ===" << std::endl; test_write_and_exists(); @@ -593,16 +383,6 @@ int main() { test_batch_write_and_batch_read(); test_batch_read_partial_miss(); - std::cout << "\n=== DummySsdTier Integration Tests (via LocalStorageManager) ===" << std::endl; - test_lsm_dummy_write_read(); - test_lsm_dummy_capacity(); - test_lsm_dummy_dram_auto_demote(); - test_lsm_dummy_ssd_full_eviction(); - test_lsm_dummy_demote_and_promote(); - test_lsm_dummy_batch_write_and_read(); - test_lsm_dummy_copy_to_ssd(); - test_lsm_dummy_flush(); - std::cout << "\nAll DummySsdTier tests passed!" << std::endl; return 0; } diff --git a/tests/cpp/umbp/local/test_follower_mode.cpp b/tests/cpp/umbp/local/test_follower_mode.cpp deleted file mode 100644 index 194036ed8..000000000 --- a/tests/cpp/umbp/local/test_follower_mode.cpp +++ /dev/null @@ -1,458 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/local/standalone_client.h" -#include "umbp/local/tiers/ssd_tier.h" - -using namespace mori::umbp; - -namespace fs = std::filesystem; - -static const std::string SHARED_SSD_DIR = "/tmp/umbp_test_follower_ssd"; - -static void cleanup_dir(const std::string& dir) { - if (fs::exists(dir)) { - fs::remove_all(dir); - } -} - -static UMBPSsdConfig make_ssd_config() { - UMBPSsdConfig cfg; - cfg.io.backend = UMBPIoBackend::Posix; - cfg.durability.mode = UMBPDurabilityMode::Relaxed; - cfg.segment_size_bytes = 4 * 1024 * 1024; - return cfg; -} - -void test_ssd_follower_exists_fallback() { - std::cout << "test_ssd_follower_exists_fallback... "; - cleanup_dir(SHARED_SSD_DIR); - - auto cfg = make_ssd_config(); - - // Leader writes a key - SSDTier leader(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadWrite); - std::vector data(4096, 'X'); - assert(leader.Write("shared_k", data.data(), data.size())); - assert(leader.Exists("shared_k")); - - // Follower can discover it via segment scan - SSDTier follower(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadOnlyShared); - assert(follower.Exists("shared_k")); - assert(!follower.Exists("nonexistent")); - - // Non-follower SSDTier at same dir would see leader's key - // (segments are scanned on construction) - SSDTier normal(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadWrite); - assert(normal.Exists("shared_k")); - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -void test_ssd_follower_read_fallback() { - std::cout << "test_ssd_follower_read_fallback... "; - cleanup_dir(SHARED_SSD_DIR); - - auto cfg = make_ssd_config(); - - // Leader writes data - SSDTier leader(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadWrite); - std::vector data(4096); - for (size_t i = 0; i < data.size(); ++i) data[i] = static_cast(i % 256); - assert(leader.Write("read_k", data.data(), data.size())); - - // Follower reads via segment scan - SSDTier follower(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadOnlyShared); - std::vector buf(4096, 0); - assert(follower.ReadIntoPtr("read_k", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); - - // Second read should hit fast path (key now in follower's key_meta_) - std::vector buf2(4096, 0); - assert(follower.ReadIntoPtr("read_k", reinterpret_cast(buf2.data()), buf2.size())); - assert(buf2 == data); - - // Read with wrong size should fail - std::vector buf3(2048, 0); - assert(!follower.ReadIntoPtr("read_k", reinterpret_cast(buf3.data()), buf3.size())); - - // Read() (vector version) should also work - SSDTier follower2(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadOnlyShared); - auto read_data = follower2.Read("read_k"); - assert(read_data.size() == data.size()); - assert(read_data == data); - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -void test_ssd_follower_evict_no_delete() { - std::cout << "test_ssd_follower_evict_no_delete... "; - cleanup_dir(SHARED_SSD_DIR); - - auto cfg = make_ssd_config(); - - // Leader writes - SSDTier leader(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadWrite); - std::vector data(4096, 'Y'); - assert(leader.Write("evict_k", data.data(), data.size())); - - // Follower reads - SSDTier follower(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadOnlyShared); - assert(follower.Exists("evict_k")); - std::vector buf(4096, 0); - assert(follower.ReadIntoPtr("evict_k", reinterpret_cast(buf.data()), buf.size())); - - // Follower evicts — should NOT delete the segment file - assert(follower.Evict("evict_k")); - auto [used, cap] = follower.Capacity(); - assert(used <= cap); // guard against size_t underflow - - // Segment file should still exist on disk - bool found_segment = false; - for (const auto& entry : fs::directory_iterator(SHARED_SSD_DIR)) { - if (entry.path().filename().string().find("segment_") == 0) { - found_segment = true; - break; - } - } - assert(found_segment); - - // Key is evicted from follower's metadata. Since the segment was already - // fully scanned, an incremental rescan won't re-discover it. - assert(!follower.Exists("evict_k")); - - // But a fresh follower instance will find it via full initial scan. - SSDTier follower2(SHARED_SSD_DIR, 1 * 1024 * 1024, cfg, SSDAccessMode::ReadOnlyShared); - assert(follower2.Exists("evict_k")); - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -void test_follower_stale_index_after_evict() { - std::cout << "test_follower_stale_index_after_evict... "; - cleanup_dir(SHARED_SSD_DIR); - - // Leader writes to shared SSD (via copy-on-write) - UMBPConfig leader_cfg; - leader_cfg.dram.capacity_bytes = 1 * 1024 * 1024; - leader_cfg.ssd.enabled = true; - leader_cfg.ssd.storage_dir = SHARED_SSD_DIR; - leader_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - leader_cfg.ssd.io.backend = UMBPIoBackend::Posix; - leader_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - leader_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - leader_cfg.copy_pipeline.async_enabled = false; - leader_cfg.role = UMBPRole::SharedSSDLeader; - StandaloneClient leader(leader_cfg); - - std::vector data(4096, 'Q'); - assert(leader.Put("stale_k", data.data(), data.size())); - - // Follower reads from SSD but does NOT auto-promote into DRAM. - UMBPConfig follower_cfg; - follower_cfg.dram.capacity_bytes = 512 * 1024; - follower_cfg.ssd.enabled = true; - follower_cfg.ssd.storage_dir = SHARED_SSD_DIR; - follower_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - follower_cfg.ssd.io.backend = UMBPIoBackend::Posix; - follower_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - follower_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - follower_cfg.role = UMBPRole::SharedSSDFollower; - follower_cfg.eviction.auto_promote_on_read = false; - StandaloneClient follower(follower_cfg); - - std::vector buf(4096, 0); - assert(follower.Get("stale_k", reinterpret_cast(buf.data()), buf.size())); - assert(follower.Index().MayExist("stale_k")); - - // Leader evicts the key, clearing it from the segment metadata - assert(leader.Remove("stale_k")); - - // Follower's Exists() should reflect the removal after rescan - // Note: With segmented log, the data bytes remain on disk but the leader - // has evicted the key from its metadata. The follower's in-memory metadata - // may still have a stale entry. The follower can re-discover the key via - // rescan since bytes are still on disk in the segment file. - // This is acceptable behavior — the key hasn't been physically deleted. - // Verify that the follower can still read the data. - std::vector buf2(4096, 0); - bool read_ok = follower.Get("stale_k", reinterpret_cast(buf2.data()), buf2.size()); - // The follower has the key in its local metadata from the previous read, - // so this should succeed. - assert(read_ok); - assert(buf2 == data); - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -void test_leader_copy_to_ssd() { - std::cout << "test_leader_copy_to_ssd... "; - cleanup_dir(SHARED_SSD_DIR); - - UMBPConfig cfg; - cfg.dram.capacity_bytes = 1 * 1024 * 1024; - cfg.ssd.enabled = true; - cfg.ssd.storage_dir = SHARED_SSD_DIR; - cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - cfg.ssd.io.backend = UMBPIoBackend::Posix; - cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - cfg.copy_pipeline.async_enabled = false; - cfg.role = UMBPRole::SharedSSDLeader; - - StandaloneClient leader(cfg); - - std::vector data(4096, 'Z'); - assert(leader.Put("copy_k", data.data(), data.size())); - - // Data should be in DRAM (index says CPU_DRAM) - assert(leader.Exists("copy_k")); - auto loc = leader.Index().Lookup("copy_k"); - assert(loc.has_value()); - assert(loc->tier == StorageTier::CPU_DRAM); - - // Segment file should exist on SSD - bool found_segment = false; - for (const auto& entry : fs::directory_iterator(SHARED_SSD_DIR)) { - if (entry.path().filename().string().find("segment_") == 0) { - found_segment = true; - break; - } - } - assert(found_segment); - - // Verify data on SSD is correct via a follower read - auto ssd_cfg = make_ssd_config(); - SSDTier ssd_check(SHARED_SSD_DIR, 4 * 1024 * 1024, ssd_cfg, SSDAccessMode::ReadOnlyShared); - auto ssd_data = ssd_check.Read("copy_k"); - assert(ssd_data.size() == data.size()); - assert(ssd_data == data); - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -void test_e2e_leader_follower() { - std::cout << "test_e2e_leader_follower... "; - cleanup_dir(SHARED_SSD_DIR); - - // Leader config - UMBPConfig leader_cfg; - leader_cfg.dram.capacity_bytes = 1 * 1024 * 1024; - leader_cfg.ssd.enabled = true; - leader_cfg.ssd.storage_dir = SHARED_SSD_DIR; - leader_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - leader_cfg.ssd.io.backend = UMBPIoBackend::Posix; - leader_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - leader_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - leader_cfg.copy_pipeline.async_enabled = false; - leader_cfg.role = UMBPRole::SharedSSDLeader; - - StandaloneClient leader(leader_cfg); - - // Follower config — separate DRAM, shared SSD dir - UMBPConfig follower_cfg; - follower_cfg.dram.capacity_bytes = 512 * 1024; - follower_cfg.ssd.enabled = true; - follower_cfg.ssd.storage_dir = SHARED_SSD_DIR; - follower_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - follower_cfg.ssd.io.backend = UMBPIoBackend::Posix; - follower_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - follower_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - follower_cfg.role = UMBPRole::SharedSSDFollower; - follower_cfg.eviction.auto_promote_on_read = true; - - StandaloneClient follower(follower_cfg); - - // Leader writes 3 keys - for (int i = 0; i < 3; ++i) { - std::string key = "e2e_k" + std::to_string(i); - std::vector data(4096, 'A' + i); - assert(leader.Put(key.c_str(), data.data(), data.size())); - } - - // Follower should see all 3 keys - for (int i = 0; i < 3; ++i) { - std::string key = "e2e_k" + std::to_string(i); - assert(follower.Exists(key)); - } - - // Follower should read all 3 keys correctly - for (int i = 0; i < 3; ++i) { - std::string key = "e2e_k" + std::to_string(i); - std::vector buf(4096, 0); - assert(follower.Get(key, reinterpret_cast(buf.data()), buf.size())); - std::vector expected(4096, 'A' + i); - assert(buf == expected); - } - - // After reading with auto_promote, follower's index should have the keys - for (int i = 0; i < 3; ++i) { - std::string key = "e2e_k" + std::to_string(i); - assert(follower.Index().MayExist(key)); - } - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -void test_follower_batch_exists() { - std::cout << "test_follower_batch_exists... "; - cleanup_dir(SHARED_SSD_DIR); - - // Leader - UMBPConfig leader_cfg; - leader_cfg.dram.capacity_bytes = 1 * 1024 * 1024; - leader_cfg.ssd.enabled = true; - leader_cfg.ssd.storage_dir = SHARED_SSD_DIR; - leader_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - leader_cfg.ssd.io.backend = UMBPIoBackend::Posix; - leader_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - leader_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - leader_cfg.copy_pipeline.async_enabled = false; - leader_cfg.role = UMBPRole::SharedSSDLeader; - StandaloneClient leader(leader_cfg); - - // Follower - UMBPConfig follower_cfg; - follower_cfg.dram.capacity_bytes = 512 * 1024; - follower_cfg.ssd.enabled = true; - follower_cfg.ssd.storage_dir = SHARED_SSD_DIR; - follower_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - follower_cfg.ssd.io.backend = UMBPIoBackend::Posix; - follower_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - follower_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - follower_cfg.role = UMBPRole::SharedSSDFollower; - StandaloneClient follower(follower_cfg); - - // Leader writes 10 keys - std::vector keys; - for (int i = 0; i < 10; ++i) { - std::string key = "batch_k" + std::to_string(i); - keys.push_back(key); - std::vector data(1024, 'a' + (i % 26)); - assert(leader.Put(key.c_str(), data.data(), data.size())); - } - - // Follower BatchExists should find all - auto results = follower.BatchExists(keys); - assert(results.size() == 10); - for (size_t i = 0; i < results.size(); ++i) { - assert(results[i]); - } - - // BatchExists with a mix of existing and non-existing - std::vector mixed_keys = {"batch_k0", "batch_k1", "missing_key"}; - auto mixed_results = follower.BatchExists(mixed_keys); - assert(mixed_results[0] == true); - assert(mixed_results[1] == true); - assert(mixed_results[2] == false); - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -void test_follower_autopromote_no_writeback() { - std::cout << "test_follower_autopromote_no_writeback... "; - cleanup_dir(SHARED_SSD_DIR); - - UMBPConfig leader_cfg; - leader_cfg.dram.capacity_bytes = 1 * 1024 * 1024; - leader_cfg.ssd.enabled = true; - leader_cfg.ssd.storage_dir = SHARED_SSD_DIR; - leader_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - leader_cfg.ssd.io.backend = UMBPIoBackend::Posix; - leader_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - leader_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - leader_cfg.copy_pipeline.async_enabled = false; - leader_cfg.role = UMBPRole::SharedSSDLeader; - StandaloneClient leader(leader_cfg); - - std::vector d0(4096, 'M'); - std::vector d1(4096, 'N'); - assert(leader.Put("auto_k0", d0.data(), d0.size())); - assert(leader.Put("auto_k1", d1.data(), d1.size())); - - // Record segment file state before follower reads - std::string segment_path; - uintmax_t before_size = 0; - for (const auto& entry : fs::directory_iterator(SHARED_SSD_DIR)) { - if (entry.path().filename().string().find("segment_") == 0) { - segment_path = entry.path().string(); - before_size = fs::file_size(entry.path()); - break; - } - } - assert(!segment_path.empty()); - assert(before_size > 0); - - UMBPConfig follower_cfg; - follower_cfg.dram.capacity_bytes = 4096; // fit one block only - follower_cfg.ssd.enabled = true; - follower_cfg.ssd.storage_dir = SHARED_SSD_DIR; - follower_cfg.ssd.capacity_bytes = 4 * 1024 * 1024; - follower_cfg.ssd.io.backend = UMBPIoBackend::Posix; - follower_cfg.ssd.durability.mode = UMBPDurabilityMode::Relaxed; - follower_cfg.ssd.segment_size_bytes = 4 * 1024 * 1024; - follower_cfg.role = UMBPRole::SharedSSDFollower; - follower_cfg.eviction.auto_promote_on_read = true; - follower_cfg.dram.high_watermark = 2.0; // force promote path on every read - StandaloneClient follower(follower_cfg); - - std::vector buf(4096, 0); - assert(follower.Get("auto_k0", reinterpret_cast(buf.data()), buf.size())); - // Ensure filesystem mtime granularity won't mask accidental rewrites. - std::this_thread::sleep_for(std::chrono::milliseconds(1200)); - assert(follower.Get("auto_k1", reinterpret_cast(buf.data()), buf.size())); - - // Segment file size should not have grown — follower never writes back - uintmax_t after_size = fs::file_size(segment_path); - assert(after_size == before_size); - - cleanup_dir(SHARED_SSD_DIR); - std::cout << "PASSED" << std::endl; -} - -int main() { - test_ssd_follower_exists_fallback(); - test_ssd_follower_read_fallback(); - test_ssd_follower_evict_no_delete(); - test_follower_stale_index_after_evict(); - test_leader_copy_to_ssd(); - test_e2e_leader_follower(); - test_follower_batch_exists(); - test_follower_autopromote_no_writeback(); - - std::cout << "\nAll follower mode tests PASSED!" << std::endl; - return 0; -} diff --git a/tests/cpp/umbp/local/test_local_block_index.cpp b/tests/cpp/umbp/local/test_local_block_index.cpp deleted file mode 100644 index 98a522657..000000000 --- a/tests/cpp/umbp/local/test_local_block_index.cpp +++ /dev/null @@ -1,147 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include -#include -#include -#include - -#include "umbp/local/block_index/local_block_index.h" - -using namespace mori::umbp; -using mori::umbp::LocalBlockIndex; - -void test_insert_lookup() { - std::cout << "test_insert_lookup... "; - - LocalBlockIndex idx; - assert(idx.Count() == 0); - - // Insert - LocalLocation loc{StorageTier::CPU_DRAM, 100, 4096}; - idx.Insert("key1", loc); - assert(idx.Count() == 1); - - // Lookup existing - auto result = idx.Lookup("key1"); - assert(result.has_value()); - assert(result->tier == StorageTier::CPU_DRAM); - assert(result->offset == 100); - assert(result->size == 4096); - - // Lookup non-existing - auto result2 = idx.Lookup("key2"); - assert(!result2.has_value()); - - // MayExist - assert(idx.MayExist("key1") == true); - assert(idx.MayExist("key2") == false); - - std::cout << "PASSED" << std::endl; -} - -void test_remove() { - std::cout << "test_remove... "; - - LocalBlockIndex idx; - idx.Insert("key1", {StorageTier::CPU_DRAM, 0, 100}); - - auto removed = idx.Remove("key1"); - assert(removed.has_value()); - assert(removed->size == 100); - assert(idx.Count() == 0); - - // Remove non-existing - auto removed2 = idx.Remove("key1"); - assert(!removed2.has_value()); - - std::cout << "PASSED" << std::endl; -} - -void test_update_tier() { - std::cout << "test_update_tier... "; - - LocalBlockIndex idx; - idx.Insert("key1", {StorageTier::CPU_DRAM, 0, 100}); - - bool ok = idx.UpdateTier("key1", StorageTier::LOCAL_SSD); - assert(ok); - auto loc = idx.Lookup("key1"); - assert(loc->tier == StorageTier::LOCAL_SSD); - - // Update non-existing - bool fail = idx.UpdateTier("nonexist", StorageTier::CPU_DRAM); - assert(!fail); - - std::cout << "PASSED" << std::endl; -} - -void test_clear() { - std::cout << "test_clear... "; - - LocalBlockIndex idx; - idx.Insert("key1", {StorageTier::CPU_DRAM, 0, 100}); - idx.Insert("key2", {StorageTier::LOCAL_SSD, 0, 200}); - assert(idx.Count() == 2); - - idx.Clear(); - assert(idx.Count() == 0); - assert(!idx.MayExist("key1")); - - std::cout << "PASSED" << std::endl; -} - -void test_concurrent_access() { - std::cout << "test_concurrent_access... "; - - LocalBlockIndex idx; - const int num_threads = 8; - const int ops_per_thread = 1000; - - std::vector threads; - for (int t = 0; t < num_threads; ++t) { - threads.emplace_back([&idx, t, ops_per_thread]() { - for (int i = 0; i < ops_per_thread; ++i) { - std::string key = "t" + std::to_string(t) + "_k" + std::to_string(i); - idx.Insert(key, {StorageTier::CPU_DRAM, 0, 64}); - idx.MayExist(key); - idx.Lookup(key); - } - }); - } - - for (auto& th : threads) th.join(); - - assert(idx.Count() == num_threads * ops_per_thread); - - std::cout << "PASSED" << std::endl; -} - -int main() { - std::cout << "=== BlockIndex Tests ===" << std::endl; - test_insert_lookup(); - test_remove(); - test_update_tier(); - test_clear(); - test_concurrent_access(); - std::cout << "All BlockIndex tests passed!" << std::endl; - return 0; -} diff --git a/tests/cpp/umbp/local/test_local_storage_manager.cpp b/tests/cpp/umbp/local/test_local_storage_manager.cpp deleted file mode 100644 index d2cf56e7c..000000000 --- a/tests/cpp/umbp/local/test_local_storage_manager.cpp +++ /dev/null @@ -1,366 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include - -#include -#include -#include -#include - -#include "umbp/local/tiers/local_storage_manager.h" - -using namespace mori::umbp; - -void test_dram_write_read() { - std::cout << "test_dram_write_read... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; // 1 MB - config.ssd.enabled = false; - - LocalStorageManager mgr(config); - - std::vector data(4096, 'A'); - assert(mgr.Write("key1", data.data(), data.size())); - assert(mgr.Exists("key1")); - - std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("key1", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); - - std::cout << "PASSED" << std::endl; -} - -void test_dram_eviction_no_ssd() { - std::cout << "test_dram_eviction_no_ssd... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1024; // tiny: 1 KB - config.ssd.enabled = false; - - LocalStorageManager mgr(config); - - // Write 512 bytes — should fit - std::vector d1(512, 'A'); - assert(mgr.Write("key1", d1.data(), d1.size())); - - // Write another 512 bytes — should fit - std::vector d2(512, 'B'); - assert(mgr.Write("key2", d2.data(), d2.size())); - - // Write 512 more — DRAM full, no SSD → last-resort eviction of key1 (LRU) - std::vector d3(512, 'C'); - assert(mgr.Write("key3", d3.data(), d3.size())); - // key1 was LRU, should have been evicted (data lost, no SSD) - assert(!mgr.Exists("key1")); - assert(mgr.Exists("key3")); - - std::cout << "PASSED" << std::endl; -} - -void test_dram_full_auto_demote_to_ssd() { - std::cout << "test_dram_full_auto_demote_to_ssd... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1024; // tiny: 1 KB - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_ssd_autodemote"; - config.ssd.capacity_bytes = 10 * 1024 * 1024; - - mori::umbp::LocalBlockIndex index; - LocalStorageManager mgr(config, &index); - - // Write 512 bytes — fits in DRAM - std::vector d1(512, 'A'); - assert(mgr.Write("key1", d1.data(), d1.size())); - index.Insert("key1", {StorageTier::CPU_DRAM, 0, 512}); - - // Write another 512 bytes — fits in DRAM - std::vector d2(512, 'B'); - assert(mgr.Write("key2", d2.data(), d2.size())); - index.Insert("key2", {StorageTier::CPU_DRAM, 0, 512}); - - // Write 512 more — DRAM full, should auto-demote key1 (LRU) to SSD - std::vector d3(512, 'C'); - assert(mgr.Write("key3", d3.data(), d3.size())); - index.Insert("key3", {StorageTier::CPU_DRAM, 0, 512}); - - // key1 should still exist (demoted to SSD, not lost!) - assert(mgr.Exists("key1")); - assert(mgr.Exists("key3")); - - // key1's index should now say LOCAL_SSD - auto loc = index.Lookup("key1"); - assert(loc.has_value()); - assert(loc->tier == StorageTier::LOCAL_SSD); - - // Verify key1 data is intact (read from SSD) - std::vector buf(512, 0); - assert(mgr.ReadIntoPtr("key1", reinterpret_cast(buf.data()), buf.size())); - assert(buf == d1); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_ssd_write_read() { - std::cout << "test_ssd_write_read... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_ssd"; - config.ssd.capacity_bytes = 10 * 1024 * 1024; - - LocalStorageManager mgr(config); - - std::vector data(4096, 'X'); - assert(mgr.Write("ssd_key", data.data(), data.size(), StorageTier::LOCAL_SSD)); - assert(mgr.Exists("ssd_key")); - - std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("ssd_key", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_demote_promote_with_index() { - std::cout << "test_demote_promote_with_index... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_ssd_dp2"; - config.ssd.capacity_bytes = 10 * 1024 * 1024; - config.eviction.auto_promote_on_read = false; - - mori::umbp::LocalBlockIndex index; - LocalStorageManager mgr(config, &index); - - // Write to DRAM - std::vector data(4096, 'D'); - assert(mgr.Write("dp_key", data.data(), data.size())); - index.Insert("dp_key", {StorageTier::CPU_DRAM, 0, 4096}); - - // Verify index says DRAM - auto loc1 = index.Lookup("dp_key"); - assert(loc1->tier == StorageTier::CPU_DRAM); - - // Demote DRAM → SSD - assert(mgr.Demote("dp_key")); - assert(mgr.Exists("dp_key")); - - // Index should now say SSD - auto loc2 = index.Lookup("dp_key"); - assert(loc2.has_value()); - assert(loc2->tier == StorageTier::LOCAL_SSD); - - // Verify data integrity after demote - std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("dp_key", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); - - // Promote SSD → DRAM - assert(mgr.Promote("dp_key")); - - // Index should now say DRAM again - auto loc3 = index.Lookup("dp_key"); - assert(loc3.has_value()); - assert(loc3->tier == StorageTier::CPU_DRAM); - - std::vector buf2(4096, 0); - assert(mgr.ReadIntoPtr("dp_key", reinterpret_cast(buf2.data()), buf2.size())); - assert(buf2 == data); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_capacity() { - std::cout << "test_capacity... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_ssd_cap"; - config.ssd.capacity_bytes = 10 * 1024 * 1024; - - LocalStorageManager mgr(config); - - auto [dram_used, dram_total] = mgr.Capacity(StorageTier::CPU_DRAM); - assert(dram_used == 0); - assert(dram_total == 1 * 1024 * 1024); - - auto [ssd_used, ssd_total] = mgr.Capacity(StorageTier::LOCAL_SSD); - assert(ssd_used == 0); - assert(ssd_total == 10 * 1024 * 1024); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_write_from_ptr() { - std::cout << "test_write_from_ptr... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = false; - - LocalStorageManager mgr(config); - - std::vector data(4096, 'P'); - uintptr_t src = reinterpret_cast(data.data()); - assert(mgr.WriteFromPtr("ptr_key", src, data.size())); - - std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("ptr_key", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); - - std::cout << "PASSED" << std::endl; -} - -// Regression test: SSD was previously self-evicting keys without notifying -// mori::umbp::LocalBlockIndex, leaving the index in a dirty state. After the fix, -// SSDTier never self-evicts; LocalStorageManager::Write() drives eviction -// with full index synchronization. -void test_ssd_full_evicts_with_index_sync() { - std::cout << "test_ssd_full_evicts_with_index_sync... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_ssd_idxsync"; - config.ssd.capacity_bytes = 1024; // tiny: 1 KB - - mori::umbp::LocalBlockIndex index; - LocalStorageManager mgr(config, &index); - - // Fill SSD with two 512-byte blocks (already demoted from DRAM) - std::vector d1(512, 'A'), d2(512, 'B'); - assert(mgr.Write("ssd1", d1.data(), d1.size(), StorageTier::LOCAL_SSD)); - index.Insert("ssd1", {StorageTier::LOCAL_SSD, 0, 512}); - assert(mgr.Write("ssd2", d2.data(), d2.size(), StorageTier::LOCAL_SSD)); - index.Insert("ssd2", {StorageTier::LOCAL_SSD, 0, 512}); - - // Writing a third block to SSD should evict ssd1 (LRU) with index removal - std::vector d3(512, 'C'); - assert(mgr.Write("ssd3", d3.data(), d3.size(), StorageTier::LOCAL_SSD)); - index.Insert("ssd3", {StorageTier::LOCAL_SSD, 0, 512}); - - // ssd1 must be gone from both storage and index - assert(!mgr.Exists("ssd1")); - assert(!index.Lookup("ssd1").has_value()); - - // ssd3 must be readable - std::vector buf(512, 0); - assert(mgr.ReadIntoPtr("ssd3", reinterpret_cast(buf.data()), buf.size())); - assert(buf == d3); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_read_into_ptr_size_mismatch() { - std::cout << "test_read_into_ptr_size_mismatch... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_ssd_sizemm"; - config.ssd.capacity_bytes = 10 * 1024 * 1024; - config.eviction.auto_promote_on_read = false; - - LocalStorageManager mgr(config); - - std::vector data(4096, 'Z'); - assert(mgr.Write("mm_key", data.data(), data.size())); - - // Exact size — should succeed - std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("mm_key", reinterpret_cast(buf.data()), 4096)); - - // Smaller size — must fail (not silently truncate) - std::vector small(2048, 0); - assert(!mgr.ReadIntoPtr("mm_key", reinterpret_cast(small.data()), 2048)); - - // Larger size — must fail - std::vector large(8192, 0); - assert(!mgr.ReadIntoPtr("mm_key", reinterpret_cast(large.data()), 8192)); - - // Same test for a key on SSD - assert(mgr.Write("ssd_mm", data.data(), data.size(), StorageTier::LOCAL_SSD)); - assert(mgr.ReadIntoPtr("ssd_mm", reinterpret_cast(buf.data()), 4096)); - assert(!mgr.ReadIntoPtr("ssd_mm", reinterpret_cast(small.data()), 2048)); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_shared_memory_mode() { - std::cout << "test_shared_memory_mode... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.dram.use_shared_memory = true; - config.dram.shm_name = "/umbp_test_shm_" + std::to_string(getpid()); - config.ssd.enabled = false; - - try { - LocalStorageManager mgr(config); - - std::vector data(4096, 'S'); - assert(mgr.Write("shm_key", data.data(), data.size())); - assert(mgr.Exists("shm_key")); - - std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("shm_key", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; - } catch (const std::runtime_error& e) { - const std::string message = e.what(); - if (message.find("shm_open failed") != std::string::npos) { - std::cout << "SKIPPED (" << message << ")" << std::endl; - return; - } - throw; - } -} - -int main() { - std::cout << "=== LocalStorageManager Tests ===" << std::endl; - test_dram_write_read(); - test_dram_eviction_no_ssd(); - test_dram_full_auto_demote_to_ssd(); - test_ssd_write_read(); - test_demote_promote_with_index(); - test_capacity(); - test_write_from_ptr(); - test_ssd_full_evicts_with_index_sync(); - test_read_into_ptr_size_mismatch(); - test_shared_memory_mode(); - std::cout << "All LocalStorageManager tests passed!" << std::endl; - return 0; -} diff --git a/tests/cpp/umbp/local/test_prefix_aware_eviction.cpp b/tests/cpp/umbp/local/test_prefix_aware_eviction.cpp deleted file mode 100644 index cdb87066d..000000000 --- a/tests/cpp/umbp/local/test_prefix_aware_eviction.cpp +++ /dev/null @@ -1,471 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include -#include -#include -#include -#include - -#include "umbp/local/block_index/local_block_index.h" -#include "umbp/local/standalone_client.h" -#include "umbp/local/tiers/local_storage_manager.h" - -using namespace mori::umbp; - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -// Build a fake 64-hex-char SHA256-like base hash. -// |index| is encoded in the first 2 hex chars so each call produces a unique -// but syntactically valid hex string (all chars in [0-9a-f]). -static std::string FakeHash(int index) { - char buf[3]; - snprintf(buf, sizeof(buf), "%02x", index & 0xff); - std::string h(64, 'a'); // 'a' is valid hex - h[0] = buf[0]; - h[1] = buf[1]; - return h; -} - -static UMBPConfig MakePrefixAwareConfig(size_t dram_bytes, size_t window = 16) { - UMBPConfig cfg; - cfg.dram.capacity_bytes = dram_bytes; - cfg.ssd.enabled = false; - cfg.eviction.policy = "prefix_aware_lru"; - cfg.eviction.candidate_window = window; - return cfg; -} - -// --------------------------------------------------------------------------- -// Test 1: Depth-biased eviction — deeper suffix block is evicted before -// shallower prefix block even when the prefix block is older (LRU). -// --------------------------------------------------------------------------- -void test_depth_biased_eviction() { - std::cout << "test_depth_biased_eviction... "; - - // DRAM holds exactly 2 × 512 B. Insert prefix (depth 0) then suffix (depth 5). - // The 3rd write cannot fit → must evict one. suffix (depth 5) should be evicted. - UMBPConfig cfg = MakePrefixAwareConfig(2 * 512); - mori::umbp::LocalBlockIndex index; - LocalStorageManager mgr(cfg, &index); - - std::string prefix_hash = FakeHash(1); - std::string suffix_hash = FakeHash(2); - std::string new_hash = FakeHash(3); - - std::vector data(512, 'X'); - - // Write prefix block (depth 0) — oldest in LRU - assert(mgr.WriteFromPtrWithDepth(prefix_hash + "_0_k", reinterpret_cast(data.data()), - data.size(), 0)); - index.Insert(prefix_hash + "_0_k", {StorageTier::CPU_DRAM, 0, 512}); - - // Write suffix block (depth 5) — newer in LRU - assert(mgr.WriteFromPtrWithDepth(suffix_hash + "_0_k", reinterpret_cast(data.data()), - data.size(), 5)); - index.Insert(suffix_hash + "_0_k", {StorageTier::CPU_DRAM, 0, 512}); - - // Write a third block — DRAM full, must evict. Prefix block is LRU-oldest, - // but suffix block has higher depth → suffix should be evicted. - std::vector new_data(512, 'Y'); - assert(mgr.WriteFromPtrWithDepth(new_hash + "_0_k", reinterpret_cast(new_data.data()), - new_data.size(), 3)); - - // Suffix (depth 5) must be gone; prefix (depth 0) must still exist. - assert(!mgr.Exists(suffix_hash + "_0_k")); - assert(mgr.Exists(prefix_hash + "_0_k")); - - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 2: Group eviction — both _k and _v variants of a page are evicted -// atomically when the group victim is selected. -// --------------------------------------------------------------------------- -void test_group_eviction() { - std::cout << "test_group_eviction... "; - - // DRAM = 2 × 512 B. Page A (depth 5) occupies both slots (A_k + A_v). - // Writing a third key must trigger group eviction of A (depth 5 is highest). - // Both A_k and A_v must disappear atomically. - UMBPConfig cfg = MakePrefixAwareConfig(2 * 512); - LocalStorageManager mgr(cfg); - - std::string hashA = FakeHash(1); - std::string hashC = FakeHash(3); - std::vector data(512, 'X'); - uintptr_t src = reinterpret_cast(data.data()); - - // Page A (depth 5): two keys fill DRAM - assert(mgr.WriteFromPtrWithDepth(hashA + "_0_k", src, 512, 5)); - assert(mgr.WriteFromPtrWithDepth(hashA + "_0_v", src, 512, 5)); - - // DRAM full. Writing one more key triggers eviction. - // Victim = A_k (first member of group A seen by SelectVictim). - // Group eviction: A_k and A_v both removed, freeing 1024 B → C_k fits. - std::vector new_data(512, 'Y'); - uintptr_t new_src = reinterpret_cast(new_data.data()); - assert(mgr.WriteFromPtrWithDepth(hashC + "_0_k", new_src, 512, 3)); - - // Both A variants must be gone. - assert(!mgr.Exists(hashA + "_0_k")); - assert(!mgr.Exists(hashA + "_0_v")); - - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 3: Fallback — keys with no depth metadata behave identically to LRU. -// --------------------------------------------------------------------------- -void test_fallback_to_lru() { - std::cout << "test_fallback_to_lru... "; - - // 2 slots. Insert key1 (no depth) then key2 (no depth). Fill cache → key1 - // (LRU-oldest) should be evicted. - UMBPConfig cfg = MakePrefixAwareConfig(2 * 512); - LocalStorageManager mgr(cfg); - - std::vector data(512, 'X'); - uintptr_t src = reinterpret_cast(data.data()); - - // Use plain WriteFromPtr (no depth) - assert(mgr.WriteFromPtr("key1_0_k", src, 512)); - assert(mgr.WriteFromPtr("key2_0_k", src, 512)); - - // Third write → evicts key1 (LRU-oldest, no depth metadata) - assert(mgr.WriteFromPtr("key3_0_k", src, 512)); - - assert(!mgr.Exists("key1_0_k")); - assert(mgr.Exists("key2_0_k")); - assert(mgr.Exists("key3_0_k")); - - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 4: Mixed metadata — depth-aware keys score over no-metadata keys. -// --------------------------------------------------------------------------- -void test_mixed_metadata() { - std::cout << "test_mixed_metadata... "; - - // DRAM = 3 × 512 B. Insert 3 keys, then write a 4th to force eviction. - // depth-5 key should be evicted (highest depth), not the LRU-oldest no-metadata key. - UMBPConfig cfg = MakePrefixAwareConfig(3 * 512); // fits exactly 3; 4th triggers eviction - LocalStorageManager mgr(cfg); - - std::string hashA = FakeHash(1); - std::string hashB = FakeHash(2); - std::string hashC = FakeHash(3); - std::string hashD = FakeHash(4); - std::vector data(512, 'X'); - uintptr_t src = reinterpret_cast(data.data()); - - // No-depth key (oldest in LRU) - assert(mgr.WriteFromPtr(hashA + "_0_k", src, 512)); - - // depth 3 - assert(mgr.WriteFromPtrWithDepth(hashB + "_0_k", src, 512, 3)); - - // depth 5 (deepest) - assert(mgr.WriteFromPtrWithDepth(hashC + "_0_k", src, 512, 5)); - - // Fourth write triggers eviction. depth-5 key should be victim. - assert(mgr.WriteFromPtrWithDepth(hashD + "_0_k", src, 512, 1)); - - assert(!mgr.Exists(hashC + "_0_k")); // depth 5 → evicted - assert(mgr.Exists(hashA + "_0_k")); // no depth → preserved - assert(mgr.Exists(hashB + "_0_k")); // depth 3 → preserved - - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 5: ExtractBaseHash — verify parsing for all key formats. -// --------------------------------------------------------------------------- -// Access the static method via a thin wrapper (it's private; test via a -// public path that exercises it: RecordGroup + GetGroup). -void test_extract_base_hash_via_group() { - std::cout << "test_extract_base_hash_via_group... "; - - UMBPConfig cfg = MakePrefixAwareConfig(16 * 512); - LocalStorageManager mgr(cfg); - - std::string hash = FakeHash(1); - std::vector data(512, 'X'); - uintptr_t src = reinterpret_cast(data.data()); - - // MHA no PP: {hash}_0_k and {hash}_0_v should group together. - assert(mgr.WriteFromPtrWithDepth(hash + "_0_k", src, 512, 2)); - assert(mgr.WriteFromPtrWithDepth(hash + "_0_v", src, 512, 2)); - - // MLA no PP (double underscore): {hash}__k - std::string hashB = FakeHash(2); - assert(mgr.WriteFromPtrWithDepth(hashB + "__k", src, 512, 3)); - - // MHA with PP: {hash}_0_0_k, {hash}_0_0_v should group. - std::string hashC = FakeHash(3); - assert(mgr.WriteFromPtrWithDepth(hashC + "_0_0_k", src, 512, 4)); - assert(mgr.WriteFromPtrWithDepth(hashC + "_0_0_v", src, 512, 4)); - - // Verify group eviction for MHA no PP: evict hash_0_k, hash_0_v should also go. - // Fill up to force eviction — we have 5 keys in 16-slot cache, so let's - // just call Evict directly on _0_k and check _0_v is NOT removed (group - // eviction only happens on space pressure paths; Evict() removes single key). - // Instead, verify via a small cache that forces group eviction. - - // Separate test: small cache with MHA pair, force group eviction on third write. - UMBPConfig cfg2 = MakePrefixAwareConfig(2 * 512); - LocalStorageManager mgr2(cfg2); - std::string hashX = FakeHash(10); - assert(mgr2.WriteFromPtrWithDepth(hashX + "_0_k", src, 512, 5)); - assert(mgr2.WriteFromPtrWithDepth(hashX + "_0_v", src, 512, 5)); - // Third write must evict both hashX_0_k and hashX_0_v (group). - std::string hashY = FakeHash(11); - assert(mgr2.WriteFromPtrWithDepth(hashY + "_0_k", src, 512, 1)); - assert(!mgr2.Exists(hashX + "_0_k")); - assert(!mgr2.Exists(hashX + "_0_v")); - - // MLA double-underscore: small cache, force eviction. - UMBPConfig cfg3 = MakePrefixAwareConfig(1 * 512); - LocalStorageManager mgr3(cfg3); - std::string hashZ = FakeHash(12); - std::string hashW = FakeHash(13); - assert(mgr3.WriteFromPtrWithDepth(hashZ + "__k", src, 512, 5)); - assert(mgr3.WriteFromPtrWithDepth(hashW + "__k", src, 512, 1)); - assert(!mgr3.Exists(hashZ + "__k")); // depth 5 evicted - assert(mgr3.Exists(hashW + "__k")); - - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 6: TOCTOU resilience — if the selected victim was concurrently evicted, -// Evict() returns false but does not crash or corrupt state. -// --------------------------------------------------------------------------- -void test_toctou_resilience() { - std::cout << "test_toctou_resilience... "; - - // DRAM = 1 slot (512 B). Write key1 (depth 5) — fills DRAM. - // Manually evict key1 to simulate a concurrent eviction that happens between - // SelectVictim choosing key1 and the eviction loop acting on it. - // Then write key2 — DRAM is already empty (key1 gone), so key2 fits. - // No crash or assertion failure = TOCTOU handled gracefully. - UMBPConfig cfg = MakePrefixAwareConfig(1 * 512); - LocalStorageManager mgr(cfg); - - std::vector data(512, 'X'); - uintptr_t src = reinterpret_cast(data.data()); - - // Fill the single slot. - assert(mgr.WriteFromPtrWithDepth("key1_0_k", src, 512, 5)); - - // Simulate concurrent eviction: key1 is removed before the Write loop acts. - assert(mgr.Evict("key1_0_k")); - - // Write key2 — DRAM is now empty, fits without eviction. - // SelectVictim would have returned key1 (highest depth in window) but the - // eviction attempt finds key1 already gone and succeeds on the next iteration. - assert(mgr.WriteFromPtrWithDepth("key2_0_k", src, 512, 1)); - assert(mgr.Exists("key2_0_k")); - assert(!mgr.Exists("key1_0_k")); - - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 7: Concurrent depth_map access — no data races under TSan. -// --------------------------------------------------------------------------- -void test_concurrent_depth_map_access() { - std::cout << "test_concurrent_depth_map_access... "; - - UMBPConfig cfg = MakePrefixAwareConfig(256 * 1024); // 256 KB - cfg.eviction.policy = "prefix_aware_lru"; - LocalStorageManager mgr(cfg); - - constexpr int kThreads = 4; - constexpr int kKeysPerThread = 20; - std::vector data(512, 'X'); - - std::vector threads; - threads.reserve(kThreads); - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&mgr, &data, t]() { - uintptr_t src = reinterpret_cast(data.data()); - for (int i = 0; i < kKeysPerThread; ++i) { - // Build a fake 64-char hex hash using thread/key index. - std::string hash(64, '0'); - // Encode thread and key index in hex chars. - char buf[16]; - snprintf(buf, sizeof(buf), "%02d%02d", t, i); - for (int c = 0; c < 4 && c < 64; ++c) hash[c] = buf[c]; - - std::string key = hash + "_0_k"; - int depth = i % 8; - mgr.WriteFromPtrWithDepth(key, src, 512, depth); - } - }); - } - for (auto& th : threads) th.join(); - - std::cout << "PASSED (run under ThreadSanitizer for full verification)" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 8: StandaloneClient depth-aware batch put — BatchPutWithDepth. -// --------------------------------------------------------------------------- -void test_umbp_client_batch_put_with_depth() { - std::cout << "test_umbp_client_batch_put_with_depth... "; - - UMBPConfig cfg; - cfg.dram.capacity_bytes = 4 * 1024; - cfg.ssd.enabled = false; - cfg.eviction.policy = "prefix_aware_lru"; - cfg.eviction.candidate_window = 16; - - StandaloneClient client(cfg); - - std::vector data(512, 'Z'); - uintptr_t src = reinterpret_cast(data.data()); - - std::vector keys = {"key_a_0_k", "key_b_0_k", "key_c_0_k"}; - std::vector ptrs = {src, src, src}; - std::vector sizes = {512, 512, 512}; - std::vector depths = {0, 3, 7}; - - auto results = client.BatchPutWithDepth(keys, ptrs, sizes, depths); - assert(results.size() == 3); - assert(results[0] && results[1] && results[2]); - - assert(client.Exists("key_a_0_k")); - assert(client.Exists("key_b_0_k")); - assert(client.Exists("key_c_0_k")); - - // Dedup: re-putting same keys should succeed (already exist). - auto results2 = client.BatchPutWithDepth(keys, ptrs, sizes, depths); - assert(results2[0] && results2[1] && results2[2]); - - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 9: SSD-path consistency — prefix_aware_lru still syncs index on -// SSD eviction. -// --------------------------------------------------------------------------- -void test_ssd_prefix_aware_index_sync() { - std::cout << "test_ssd_prefix_aware_index_sync... "; - - UMBPConfig cfg; - cfg.dram.capacity_bytes = 4 * 1024 * 1024; // large DRAM — no DRAM pressure - cfg.ssd.enabled = true; - cfg.ssd.storage_dir = "/tmp/umbp_test_prefix_ssd"; - cfg.ssd.capacity_bytes = 1024; // tiny SSD: 1 KB - cfg.eviction.policy = "prefix_aware_lru"; - cfg.eviction.candidate_window = 16; - - mori::umbp::LocalBlockIndex index; - LocalStorageManager mgr(cfg, &index); - - std::vector d1(512, 'A'), d2(512, 'B'), d3(512, 'C'); - - // Write two keys directly to SSD with depths. - assert(mgr.WriteFromPtrWithDepth("hash1_0_k", reinterpret_cast(d1.data()), 512, 2, - StorageTier::LOCAL_SSD)); - index.Insert("hash1_0_k", {StorageTier::LOCAL_SSD, 0, 512}); - - assert(mgr.WriteFromPtrWithDepth("hash2_0_k", reinterpret_cast(d2.data()), 512, 7, - StorageTier::LOCAL_SSD)); - index.Insert("hash2_0_k", {StorageTier::LOCAL_SSD, 0, 512}); - - // Writing a third 512-byte key to SSD forces eviction. - // hash2_0_k has depth 7 (higher) → should be evicted, index removed. - assert(mgr.WriteFromPtrWithDepth("hash3_0_k", reinterpret_cast(d3.data()), 512, 1, - StorageTier::LOCAL_SSD)); - index.Insert("hash3_0_k", {StorageTier::LOCAL_SSD, 0, 512}); - - assert(!mgr.Exists("hash2_0_k")); - assert(!index.Lookup("hash2_0_k").has_value()); - assert(mgr.Exists("hash1_0_k")); - assert(mgr.Exists("hash3_0_k")); - - mgr.Clear(); - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// Test 10: Follower safety — prefix_aware_lru never writes/deletes shared SSD. -// --------------------------------------------------------------------------- -void test_follower_no_ssd_write() { - std::cout << "test_follower_no_ssd_write... "; - - UMBPConfig cfg; - cfg.dram.capacity_bytes = 1 * 1024 * 1024; - cfg.ssd.enabled = true; - cfg.ssd.storage_dir = "/tmp/umbp_test_prefix_follower"; - cfg.ssd.capacity_bytes = 32 * 1024 * 1024; - cfg.eviction.policy = "prefix_aware_lru"; - cfg.eviction.candidate_window = 16; - cfg.role = UMBPRole::SharedSSDFollower; - - LocalStorageManager mgr(cfg); - - std::vector data(512, 'F'); - uintptr_t src = reinterpret_cast(data.data()); - - // Follower WriteFromPtrWithDepth should still write to DRAM. - assert(mgr.WriteFromPtrWithDepth("key1_0_k", src, 512, 0)); - assert(mgr.Exists("key1_0_k")); - - // Follower must not write to SSD directly. - // CopyToSSD returns false for followers. - assert(!mgr.CopyToSSD("key1_0_k")); - - // DRAM eviction in follower mode must not touch SSD. - // Fill DRAM with depth-aware keys. - for (int i = 0; i < 2048; ++i) { - std::string k = "loop_" + std::to_string(i) + "_0_k"; - mgr.WriteFromPtrWithDepth(k, src, 512, i % 10); - // No assert — DRAM will evict, but follower path must not write SSD. - } - - // If we get here without crash or SSD writes, follower safety holds. - std::cout << "PASSED" << std::endl; -} - -// --------------------------------------------------------------------------- -// main -// --------------------------------------------------------------------------- -int main() { - std::cout << "=== Prefix-Aware Eviction Tests ===" << std::endl; - test_depth_biased_eviction(); - test_group_eviction(); - test_fallback_to_lru(); - test_mixed_metadata(); - test_extract_base_hash_via_group(); - test_toctou_resilience(); - test_concurrent_depth_map_access(); - test_umbp_client_batch_put_with_depth(); - test_ssd_prefix_aware_index_sync(); - test_follower_no_ssd_write(); - std::cout << "All prefix-aware eviction tests passed!" << std::endl; - return 0; -} diff --git a/tests/cpp/umbp/local/test_segment_crc.cpp b/tests/cpp/umbp/local/test_segment_crc.cpp new file mode 100644 index 000000000..689966c23 --- /dev/null +++ b/tests/cpp/umbp/local/test_segment_crc.cpp @@ -0,0 +1,216 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Pins the segment record checksum to CRC-32C. +// +// CrcUpdate has two implementations selected at runtime (SSE4.2 hardware and a +// portable table). A segment written on a host with hardware support must +// verify on one without it, so the two MUST agree bit-for-bit -- this test +// compares whichever path this host selects against an independent bitwise +// reference, which is the shared ground truth for both. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/local/tiers/segment/segment_format.h" +#include "umbp/local/tiers/segment/segment_index.h" +#include "umbp/local/tiers/segment/segment_writer.h" +#include "umbp/storage/io/storage_io_driver.h" + +using namespace mori::umbp::segment; + +namespace { + +// Independent bitwise CRC-32C (Castagnoli, reflected). Deliberately naive and +// written from the polynomial rather than shared with the production code, so a +// change there cannot silently "fix" this reference too. +uint32_t ReferenceCrc32c(const void* data, size_t size, uint32_t crc = 0xFFFFFFFFu) { + const uint8_t* p = static_cast(data); + for (size_t i = 0; i < size; ++i) { + crc ^= p[i]; + for (int k = 0; k < 8; ++k) crc = (crc >> 1) ^ (0x82F63B78u & (~(crc & 1u) + 1u)); + } + return crc; +} + +void test_known_vector() { + std::cout << "test_known_vector... "; + // The standard CRC-32C check value: crc("123456789") == 0xE3069283. + // Catches an accidental revert to CRC-32/ISO-HDLC (0xEDB88320), whose check + // value is 0xCBF43926. + const std::string s = "123456789"; + assert(~CrcUpdate(s.data(), s.size()) == 0xE3069283u); + std::cout << "OK\n"; +} + +void test_matches_reference_all_tail_lengths() { + std::cout << "test_matches_reference_all_tail_lengths... "; + // The hardware path consumes 8 bytes at a time and finishes byte-wise, so + // every length mod 8 (and every length below one block) is a distinct case. + std::mt19937 rng(0xC0FFEE); + std::vector buf(8192); + for (auto& b : buf) b = static_cast(rng()); + + for (size_t n = 0; n <= 512; ++n) { + assert(CrcUpdate(buf.data(), n) == ReferenceCrc32c(buf.data(), n)); + } + for (size_t n : {1023u, 1024u, 1025u, 4095u, 4096u, 8191u, 8192u}) { + assert(CrcUpdate(buf.data(), n) == ReferenceCrc32c(buf.data(), n)); + } + // Unaligned starts: the hardware path must not assume 8-byte alignment. + for (size_t off = 1; off < 9; ++off) { + assert(CrcUpdate(buf.data() + off, 777) == ReferenceCrc32c(buf.data() + off, 777)); + } + std::cout << "OK\n"; +} + +void test_streaming_equals_one_shot() { + std::cout << "test_streaming_equals_one_shot... "; + // ComputeRecordCrc32 seeds the value digest with the key digest, so chunked + // updates must compose exactly like a single pass. + std::vector buf(10007); + for (size_t i = 0; i < buf.size(); ++i) buf[i] = static_cast(i * 31 + 7); + + const uint32_t one_shot = CrcUpdate(buf.data(), buf.size()); + uint32_t chunked = 0xFFFFFFFFu; + for (size_t off = 0; off < buf.size();) { + const size_t take = std::min(1 + (off % 97), buf.size() - off); + chunked = CrcUpdate(buf.data() + off, take, chunked); + off += take; + } + assert(one_shot == chunked); + std::cout << "OK\n"; +} + +void test_key_is_bound_to_record() { + std::cout << "test_key_is_bound_to_record... "; + // The read path compares against a checksum taken over key+value, which is + // what turns "bytes are intact" into "these are the bytes for THIS key" -- + // the guard against a stale index entry pointing at another record. + std::vector value(4096, 0xAB); + const uint32_t a = ComputeRecordCrc32("prefix/block/0001", value.data(), value.size()); + const uint32_t b = ComputeRecordCrc32("prefix/block/0002", value.data(), value.size()); + assert(a != b); + + // ...and a single flipped payload byte must change the digest. + std::vector corrupted = value; + corrupted[2048] ^= 0x01; + assert(ComputeRecordCrc32("prefix/block/0001", corrupted.data(), corrupted.size()) != a); + std::cout << "OK\n"; +} + +void test_empty_inputs() { + std::cout << "test_empty_inputs... "; + // Zero-length update is the identity on the running digest. + assert(CrcUpdate(nullptr, 0, 0x12345678u) == 0x12345678u); + const std::string key = "k"; + assert(ComputeRecordCrc32(key, nullptr, 0) == ~CrcUpdate(key.data(), key.size())); + std::cout << "OK\n"; +} + +void test_build_reserve_split() { + std::cout << "test_build_reserve_split... "; + // Build() runs outside the tier mutex and Reserve() inside it, so Build must + // produce the complete record except for `generation`, which only exists once + // the reservation is taken. Reserve patches it in place -- if that patch were + // lost, the scanner would recover every record with generation 0 and + // RecordRecoveredEntry's ordering check would keep the wrong version. + auto driver = mori::umbp::CreateStorageIoDriver(mori::umbp::UMBPIoBackend::Posix, 0); + Writer writer(*driver); + + const std::string key = "block/abc"; + std::vector value(256, 0x5A); + + PreparedRecord pr; + writer.Build(key, value.data(), value.size(), &pr); + + RecordHeader built{}; + std::memcpy(&built, pr.record.data(), sizeof(built)); + assert(built.magic == kRecordMagic); + assert(built.version == kRecordVersion); + assert(built.flags == kFlagCommitted); + assert(built.key_len == key.size()); + assert(built.value_size == value.size()); + assert(built.crc32 == ComputeRecordCrc32(key, value.data(), value.size())); + assert(pr.crc32 == built.crc32); + assert(built.generation == 0); // not yet reserved + // v3 pads the record so the value starts on a kRecordAlign boundary. + assert(pr.record.size() == RecordBytes(key.size(), value.size())); + assert(pr.record.size() % kRecordAlign == 0); + // Key and value bytes are already in place before the lock is taken. + assert(std::memcmp(pr.record.data() + sizeof(RecordHeader), key.data(), key.size()) == 0); + assert(std::memcmp(pr.record.data() + PrefixBytes(key.size()), value.data(), value.size()) == 0); + + // Reserve stamps a real generation into the already-built buffer. + Index index(1 << 20); + Meta seg; + seg.id = 0; + seg.write_offset = 0; + assert(writer.Reserve(key, value.size(), &seg, index, &pr)); + + RecordHeader reserved{}; + std::memcpy(&reserved, pr.record.data(), sizeof(reserved)); + assert(reserved.generation != 0); + assert(reserved.generation == pr.reservation.meta.generation); + // Reserve must touch nothing else. + assert(reserved.crc32 == built.crc32); + assert(reserved.key_len == built.key_len); + assert(reserved.value_size == built.value_size); + assert(std::memcmp(pr.record.data() + PrefixBytes(key.size()), value.data(), value.size()) == 0); + // The reservation must put the value on an aligned offset -- the precondition + // that makes the record readable and writable with O_DIRECT. + assert(pr.reservation.record_offset % kRecordAlign == 0); + assert(pr.reservation.meta.value_offset % kRecordAlign == 0); + assert(seg.write_offset % kRecordAlign == 0); + + // Generations strictly advance, so recovery can order two writes of one key. + PreparedRecord pr2; + writer.Build(key, value.data(), value.size(), &pr2); + assert(writer.Reserve(key, value.size(), &seg, index, &pr2)); + assert(pr2.reservation.meta.generation > pr.reservation.meta.generation); + + // Reserve on a record that was never built must fail rather than write a + // generation past the end of an empty buffer. + PreparedRecord empty; + assert(!writer.Reserve(key, value.size(), &seg, index, &empty)); + + std::cout << "OK\n"; +} + +} // namespace + +int main() { + test_known_vector(); + test_matches_reference_all_tail_lengths(); + test_streaming_equals_one_shot(); + test_key_is_bound_to_record(); + test_empty_inputs(); + test_build_reserve_split(); + std::cout << "All segment CRC tests passed\n"; + return 0; +} diff --git a/tests/cpp/umbp/local/test_sharded_ssd_tier.cpp b/tests/cpp/umbp/local/test_sharded_ssd_tier.cpp new file mode 100644 index 000000000..c6f20796b --- /dev/null +++ b/tests/cpp/umbp/local/test_sharded_ssd_tier.cpp @@ -0,0 +1,298 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// Multi-drive SSD tier: keys spread across N directories (one per drive), read +// back from whichever drive holds them, and batch IO runs on all drives at once. +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/local/tiers/sharded_ssd_tier.h" +#include "umbp/local/tiers/ssd_backend_factory.h" + +using namespace mori::umbp; + +namespace { + +// A fresh set of per-drive directories for one test case. +std::string MakeDirs(const std::string& stem, int n) { + std::string joined; + for (int i = 0; i < n; ++i) { + const std::string dir = "/tmp/umbp_test_" + stem + "_d" + std::to_string(i); + std::filesystem::remove_all(dir); + if (!joined.empty()) joined += ','; + joined += dir; + } + return joined; +} + +UMBPSsdConfig MakeConfig(const std::string& dirs, size_t capacity) { + UMBPSsdConfig cfg; + cfg.enabled = true; + cfg.storage_dir = dirs; + cfg.capacity_bytes = capacity; + cfg.segment_size_bytes = 4ULL * 1024 * 1024; + return cfg; +} + +void test_storage_dirs_parsing() { + std::cout << "test_storage_dirs_parsing... "; + UMBPSsdConfig cfg; + + cfg.storage_dir = "/mnt/nvme0"; + assert(cfg.StorageDirs().size() == 1); + assert(cfg.StorageDirs()[0] == "/mnt/nvme0"); + + cfg.storage_dir = "/mnt/nvme0,/mnt/nvme1,/mnt/nvme2"; + auto dirs = cfg.StorageDirs(); + assert(dirs.size() == 3); + assert(dirs[2] == "/mnt/nvme2"); + + // Whitespace around separators is tolerated. + cfg.storage_dir = "/mnt/a , /mnt/b"; + dirs = cfg.StorageDirs(); + assert(dirs.size() == 2); + assert(dirs[0] == "/mnt/a"); + assert(dirs[1] == "/mnt/b"); + + // Empty never yields an empty list. + cfg.storage_dir = ""; + assert(cfg.StorageDirs().size() == 1); + std::cout << "PASSED" << std::endl; +} + +// One directory must still produce a plain SSDTier (no sharding overhead, and +// no behavior change for every existing single-drive deployment). +void test_factory_single_dir_is_unsharded() { + std::cout << "test_factory_single_dir_is_unsharded... "; + auto cfg = MakeConfig(MakeDirs("factory_single", 1), 64ULL * 1024 * 1024); + auto backend = MakeFileSsdBackend(cfg); + assert(backend != nullptr); + assert(dynamic_cast(backend.get()) == nullptr); + std::cout << "PASSED" << std::endl; +} + +void test_factory_multi_dir_is_sharded() { + std::cout << "test_factory_multi_dir_is_sharded... "; + auto cfg = MakeConfig(MakeDirs("factory_multi", 3), 96ULL * 1024 * 1024); + auto backend = MakeFileSsdBackend(cfg); + auto* sharded = dynamic_cast(backend.get()); + assert(sharded != nullptr); + assert(sharded->ShardCount() == 3); + // capacity_bytes is the TOTAL budget, split across drives. + assert(backend->Capacity().second == 96ULL * 1024 * 1024); + std::cout << "PASSED" << std::endl; +} + +// The headline case: 1000 same-size keys over 2 drives land ~500/500 and every +// one reads back correctly from whichever drive got it. +void test_balanced_placement_and_readback() { + std::cout << "test_balanced_placement_and_readback... "; + auto cfg = MakeConfig(MakeDirs("balance", 2), 128ULL * 1024 * 1024); + auto backend = MakeFileSsdBackend(cfg); + auto* sharded = dynamic_cast(backend.get()); + assert(sharded != nullptr); + + constexpr int kKeys = 1000; + constexpr size_t kSize = 4096; + int per_shard[2] = {0, 0}; + for (int i = 0; i < kKeys; ++i) { + std::vector payload(kSize, static_cast('A' + (i % 26))); + const std::string key = "k" + std::to_string(i); + assert(backend->Write(key, payload.data(), payload.size())); + const int s = sharded->ShardOf(key); + assert(s == 0 || s == 1); + ++per_shard[s]; + } + // Exact round-robin for uniform sizes on uniform drives. + assert(per_shard[0] == kKeys / 2); + assert(per_shard[1] == kKeys / 2); + + for (int i = 0; i < kKeys; ++i) { + std::vector out(kSize, 0); + const std::string key = "k" + std::to_string(i); + assert(backend->ReadIntoPtr(key, reinterpret_cast(out.data()), out.size())); + assert(out == std::vector(kSize, static_cast('A' + (i % 26)))); + } + std::cout << "PASSED" << std::endl; +} + +// A re-put of a live key must go back to its own drive, never get duplicated +// onto a second one. +void test_rewrite_stays_on_same_shard() { + std::cout << "test_rewrite_stays_on_same_shard... "; + auto cfg = MakeConfig(MakeDirs("rewrite", 3), 96ULL * 1024 * 1024); + auto backend = MakeFileSsdBackend(cfg); + auto* sharded = dynamic_cast(backend.get()); + + std::vector a(2048, 'A'); + std::vector b(2048, 'B'); + assert(backend->Write("sticky", a.data(), a.size())); + const int first_shard = sharded->ShardOf("sticky"); + assert(first_shard >= 0); + assert(backend->Write("sticky", b.data(), b.size())); + assert(sharded->ShardOf("sticky") == first_shard); + + std::vector out(2048, 0); + assert(backend->ReadIntoPtr("sticky", reinterpret_cast(out.data()), out.size())); + assert(out == b); + std::cout << "PASSED" << std::endl; +} + +void test_batch_write_read_across_shards() { + std::cout << "test_batch_write_read_across_shards... "; + auto cfg = MakeConfig(MakeDirs("batch", 4), 128ULL * 1024 * 1024); + auto backend = MakeFileSsdBackend(cfg); + auto* sharded = dynamic_cast(backend.get()); + + constexpr int kKeys = 64; + constexpr size_t kSize = 8192; + std::vector keys; + std::vector> payloads; + std::vector srcs; + std::vector sizes; + for (int i = 0; i < kKeys; ++i) { + keys.push_back("bk" + std::to_string(i)); + payloads.emplace_back(kSize, static_cast(i)); + sizes.push_back(kSize); + } + for (auto& p : payloads) srcs.push_back(p.data()); + + auto wrote = backend->BatchWrite(keys, srcs, sizes); + assert(wrote.size() == static_cast(kKeys)); + for (bool ok : wrote) assert(ok); + + // Every drive should be carrying part of the batch — that is what makes the + // reads parallel. + for (const auto& cap : sharded->PerShardCapacity()) assert(cap.first > 0); + + std::vector> outs(kKeys, std::vector(kSize, 0)); + std::vector dsts; + for (auto& o : outs) dsts.push_back(reinterpret_cast(o.data())); + auto read = backend->BatchReadIntoPtr(keys, dsts, sizes); + assert(read.size() == static_cast(kKeys)); + for (int i = 0; i < kKeys; ++i) { + assert(read[i]); + assert(outs[i] == std::vector(kSize, static_cast(i))); + } + std::cout << "PASSED" << std::endl; +} + +// A key that was never written must miss, not read someone else's bytes, and a +// batch containing unknown keys must report exactly those as failures. +void test_missing_keys() { + std::cout << "test_missing_keys... "; + auto cfg = MakeConfig(MakeDirs("missing", 2), 64ULL * 1024 * 1024); + auto backend = MakeFileSsdBackend(cfg); + + std::vector payload(1024, 'X'); + assert(backend->Write("present", payload.data(), payload.size())); + assert(backend->Exists("present")); + assert(!backend->Exists("absent")); + + std::vector out(1024, 0); + assert(!backend->ReadIntoPtr("absent", reinterpret_cast(out.data()), out.size())); + + std::vector keys = {"present", "absent"}; + std::vector> outs(2, std::vector(1024, 0)); + std::vector dsts = {reinterpret_cast(outs[0].data()), + reinterpret_cast(outs[1].data())}; + std::vector sizes = {1024, 1024}; + auto read = backend->BatchReadIntoPtr(keys, dsts, sizes); + assert(read[0]); + assert(!read[1]); + std::cout << "PASSED" << std::endl; +} + +void test_evict_and_clear() { + std::cout << "test_evict_and_clear... "; + auto cfg = MakeConfig(MakeDirs("evict", 2), 64ULL * 1024 * 1024); + auto backend = MakeFileSsdBackend(cfg); + + std::vector payload(4096, 'E'); + for (int i = 0; i < 8; ++i) { + assert(backend->Write("ek" + std::to_string(i), payload.data(), payload.size())); + } + const size_t used_before = backend->Capacity().first; + assert(used_before > 0); + + assert(backend->Evict("ek3")); + assert(!backend->Exists("ek3")); + assert(!backend->Evict("ek3")); // idempotent-ish: already gone + assert(backend->Exists("ek4")); + + backend->Clear(); + assert(backend->Capacity().first == 0); + for (int i = 0; i < 8; ++i) assert(!backend->Exists("ek" + std::to_string(i))); + std::cout << "PASSED" << std::endl; +} + +// Uneven drive sizes: placement is capacity-aware, so the bigger drive takes +// more keys rather than the set failing once the small one fills. +void test_uneven_shards_place_by_free_space() { + std::cout << "test_uneven_shards_place_by_free_space... "; + const std::string dirs = MakeDirs("uneven", 2); + auto cfg = MakeConfig(dirs, 0); // capacity set per shard below + + std::vector parts = cfg.StorageDirs(); + std::vector> shards; + UMBPSsdConfig shard_cfg = cfg; + shard_cfg.capacity_bytes = 8ULL * 1024 * 1024; + shards.push_back(std::make_unique(parts[0], 8ULL * 1024 * 1024, shard_cfg)); + shard_cfg.capacity_bytes = 32ULL * 1024 * 1024; + shards.push_back(std::make_unique(parts[1], 32ULL * 1024 * 1024, shard_cfg)); + ShardedSsdTier tier(std::move(shards)); + + assert(tier.Capacity().second == 40ULL * 1024 * 1024); + + int per_shard[2] = {0, 0}; + std::vector payload(64 * 1024, 'U'); + for (int i = 0; i < 60; ++i) { + const std::string key = "uk" + std::to_string(i); + assert(tier.Write(key, payload.data(), payload.size())); + ++per_shard[tier.ShardOf(key)]; + } + // The 32MB drive must have absorbed strictly more than the 8MB one. + assert(per_shard[1] > per_shard[0]); + std::cout << "PASSED" << std::endl; +} + +} // namespace + +int main() { + std::cout << "=== Sharded (multi-drive) SSD Tier Tests ===" << std::endl; + test_storage_dirs_parsing(); + test_factory_single_dir_is_unsharded(); + test_factory_multi_dir_is_sharded(); + test_balanced_placement_and_readback(); + test_rewrite_stays_on_same_shard(); + test_batch_write_read_across_shards(); + test_missing_keys(); + test_evict_and_clear(); + test_uneven_shards_place_by_free_space(); + std::cout << "All sharded SSD tier tests passed!" << std::endl; + return 0; +} diff --git a/tests/cpp/umbp/local/test_ssd_direct_io.cpp b/tests/cpp/umbp/local/test_ssd_direct_io.cpp new file mode 100644 index 000000000..45c30b9ec --- /dev/null +++ b/tests/cpp/umbp/local/test_ssd_direct_io.cpp @@ -0,0 +1,343 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// Covers the direct-I/O SSD path and the two knobs added alongside it +// (verify_crc, tier_io_threads). The point of direct I/O here is measurement +// honesty: with buffered I/O the tier reads back out of the page cache at many +// times device bandwidth, so any drive-scaling or DRAM-vs-SSD number taken from +// it is meaningless. These tests pin the correctness properties that make the +// unbuffered path usable, not its speed. + +#include + +#include +#include +#include +#include +#include + +#include "umbp/common/aligned_buffer.h" +#include "umbp/local/tiers/segment/segment_format.h" +#include "umbp/local/tiers/ssd_tier.h" + +namespace fs = std::filesystem; +using mori::umbp::AlignedBuffer; +using mori::umbp::IsDirectIoCompatible; +using mori::umbp::SSDTier; +using mori::umbp::UMBPIoBackend; +using mori::umbp::UMBPSsdConfig; +using namespace mori::umbp::segment; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::cerr << "FAILED: " #cond " at " << __FILE__ << ":" << __LINE__ << std::endl; \ + std::exit(1); \ + } \ + } while (0) + +namespace { + +std::string MakeDir(const std::string& name) { + const std::string dir = "/tmp/umbp_test_directio_" + name; + fs::remove_all(dir); + fs::create_directories(dir); + return dir; +} + +UMBPSsdConfig BaseConfig(const std::string& dir, bool direct_io, bool verify_crc, + int tier_io_threads = 4) { + UMBPSsdConfig cfg; + cfg.enabled = true; + cfg.storage_dir = dir; + cfg.capacity_bytes = 64ULL * 1024 * 1024; + cfg.segment_size_bytes = 16ULL * 1024 * 1024; + cfg.io.backend = UMBPIoBackend::Posix; // no io_uring dependency in unit tests + cfg.direct_io = direct_io; + cfg.verify_crc = verify_crc; + cfg.tier_io_threads = tier_io_threads; + return cfg; +} + +// A value whose size is an exact kRecordAlign multiple, like the KV pages this +// tier actually stores — the case that reads straight into the caller's buffer +// with no bounce. +std::vector AlignedValue(char fill, size_t blocks = 2) { + return std::vector(static_cast(kRecordAlign) * blocks, fill); +} + +} // namespace + +// The v3 layout must put every record and every value on an alignment boundary, +// regardless of key length — that is the property O_DIRECT depends on, and it +// has to hold for the buffered path too so one directory works either way. +void test_layout_is_aligned() { + std::cout << "test_layout_is_aligned... "; + for (size_t key_len : {1u, 9u, 31u, 32u, 33u, 100u, 4063u, 4064u, 4065u, 9000u}) { + for (size_t value : {1u, 512u, 4095u, 4096u, 4097u, 1u << 20}) { + CHECK(PrefixBytes(key_len) % kRecordAlign == 0); + CHECK(PrefixBytes(key_len) >= sizeof(RecordHeader) + key_len); + CHECK(PaddedValueBytes(value) % kRecordAlign == 0); + CHECK(PaddedValueBytes(value) >= value); + CHECK(RecordBytes(key_len, value) % kRecordAlign == 0); + } + } + std::cout << "OK\n"; +} + +void test_aligned_buffer() { + std::cout << "test_aligned_buffer... "; + AlignedBuffer buf(100); + CHECK(reinterpret_cast(buf.data()) % kRecordAlign == 0); + CHECK(buf.size() == 100); + CHECK(buf.padded_size() == kRecordAlign); + // Padding must be zeroed, so record padding never leaks heap bytes to disk. + for (size_t i = buf.size(); i < buf.padded_size(); ++i) CHECK(buf.data()[i] == 0); + + CHECK(IsDirectIoCompatible(buf.data(), static_cast(kRecordAlign))); + CHECK(!IsDirectIoCompatible(buf.data(), 100)); // length unaligned + CHECK(!IsDirectIoCompatible(buf.data() + 1, kRecordAlign)); // address unaligned + + AlignedBuffer moved(std::move(buf)); + CHECK(reinterpret_cast(moved.data()) % kRecordAlign == 0); + std::cout << "OK\n"; +} + +// Round-trip through a tier, once buffered and once unbuffered. Both must +// return byte-identical data and survive a reopen (which exercises the scanner's +// aligned prefix reads on an O_DIRECT fd). +void test_round_trip(bool direct_io) { + std::cout << "test_round_trip(direct_io=" << direct_io << ")... "; + const std::string dir = MakeDir(direct_io ? "rt_direct" : "rt_buffered"); + auto cfg = BaseConfig(dir, direct_io, /*verify_crc=*/true); + + const std::vector keys = {"k/aligned", "k/odd", "k/big"}; + std::vector> values = {AlignedValue('A'), std::vector(1234, 'B'), + AlignedValue('C', 5)}; + + { + SSDTier tier(dir, cfg.capacity_bytes, cfg); + if (direct_io && !tier.direct_io_active()) { + std::cout << "SKIPPED (filesystem rejects O_DIRECT)\n"; + return; + } + std::vector ptrs; + std::vector sizes; + for (auto& v : values) { + ptrs.push_back(v.data()); + sizes.push_back(v.size()); + } + CHECK(tier.WriteBatch(keys, ptrs, sizes)); + + // Read back into an aligned destination (the no-bounce path) and into a + // deliberately misaligned one (the bounce path); both must be exact. + AlignedBuffer aligned_dst(values[0].size()); + CHECK(tier.ReadIntoPtr(keys[0], reinterpret_cast(aligned_dst.data()), + values[0].size())); + CHECK(std::memcmp(aligned_dst.data(), values[0].data(), values[0].size()) == 0); + + std::vector odd_dst(values[1].size()); + CHECK(tier.ReadIntoPtr(keys[1], reinterpret_cast(odd_dst.data()), values[1].size())); + CHECK(std::memcmp(odd_dst.data(), values[1].data(), values[1].size()) == 0); + + // Batch read, mixing aligned and unaligned destinations in one call. + AlignedBuffer b0(values[0].size()); + std::vector b1(values[1].size()); + AlignedBuffer b2(values[2].size()); + std::vector dsts = {reinterpret_cast(b0.data()), + reinterpret_cast(b1.data()), + reinterpret_cast(b2.data())}; + auto ok = + tier.ReadBatchIntoPtr(keys, dsts, {values[0].size(), values[1].size(), values[2].size()}); + CHECK(ok.size() == 3); + for (bool r : ok) CHECK(r); + CHECK(std::memcmp(b0.data(), values[0].data(), values[0].size()) == 0); + CHECK(std::memcmp(b1.data(), values[1].data(), values[1].size()) == 0); + CHECK(std::memcmp(b2.data(), values[2].data(), values[2].size()) == 0); + + CHECK(tier.Read(keys[1]) == values[1]); + } + + // Reopen: the scanner has to rebuild the index from the padded records, over + // an O_DIRECT fd when direct I/O is on. A regression here shows up as a + // silent empty tier, not an error, so assert every key came back. + { + SSDTier tier(dir, cfg.capacity_bytes, cfg); + for (size_t i = 0; i < keys.size(); ++i) { + CHECK(tier.Exists(keys[i])); + CHECK(tier.Read(keys[i]) == values[i]); + } + } + + fs::remove_all(dir); + std::cout << "OK\n"; +} + +// LocateRecord must hand back an fd and offset that, pread'd directly, yield the +// stored value — the contract GdsEngine relies on to hipFileRead a segment range +// straight into GPU memory. +void test_locate_record() { + std::cout << "test_locate_record... "; + const std::string dir = MakeDir("locate"); + auto cfg = BaseConfig(dir, /*direct_io=*/false, /*verify_crc=*/true); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + + const std::string key = "k/locate"; + const auto value = AlignedValue('L', 3); + CHECK(tier.Write(key, value.data(), value.size())); + + auto loc = tier.LocateRecord(key); + CHECK(loc.has_value()); + CHECK(loc->fd >= 0); + CHECK(loc->value_size == value.size()); + CHECK(loc->readable_size == PaddedValueBytes(value.size())); + CHECK(loc->value_offset % kRecordAlign == 0); + CHECK(loc->direct_io == tier.direct_io_active()); + + // The offset/size must point at the real bytes: pread the fd ourselves and + // compare — exactly what a GDS read does, but on the CPU. + std::vector got(loc->value_size, 0); + const ssize_t n = ::pread(loc->fd, got.data(), got.size(), static_cast(loc->value_offset)); + CHECK(n == static_cast(got.size())); + CHECK(std::memcmp(got.data(), value.data(), value.size()) == 0); + + // An unknown key has no location. + CHECK(!tier.LocateRecord("k/missing").has_value()); + + fs::remove_all(dir); + std::cout << "OK\n"; +} + +// verify_crc=0 must round-trip, and — the part that is easy to get wrong — a +// store written with checksums off must stay readable by a tier with checksums +// on. Without the kFlagNoCrc marker every such key would read back as corrupt. +void test_crc_disabled_round_trip() { + std::cout << "test_crc_disabled_round_trip... "; + const std::string dir = MakeDir("crc_off"); + const std::string key = "k/nocrc"; + const auto value = AlignedValue('Z'); + + { + auto cfg = BaseConfig(dir, /*direct_io=*/false, /*verify_crc=*/false); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + CHECK(tier.Write(key, value.data(), value.size())); + CHECK(tier.Read(key) == value); + } + + // Same directory, checksums back on. + { + auto cfg = BaseConfig(dir, /*direct_io=*/false, /*verify_crc=*/true); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + CHECK(tier.Exists(key)); + CHECK(tier.Read(key) == value); + + // A record written *with* a checksum is still verified by this tier. + const std::string key2 = "k/withcrc"; + const auto value2 = AlignedValue('Y'); + CHECK(tier.Write(key2, value2.data(), value2.size())); + CHECK(tier.Read(key2) == value2); + } + + fs::remove_all(dir); + std::cout << "OK\n"; +} + +// tier_io_threads only changes how the CPU phases are scheduled; results must be +// identical at any thread count. 1 vs 8 also exercises the ParallelFor +// boundary between the inline loop and the worker path. +void test_tier_threads_are_result_invariant() { + std::cout << "test_tier_threads_are_result_invariant... "; + std::vector keys; + std::vector> values; + for (int i = 0; i < 17; ++i) { // deliberately not a multiple of any thread count + keys.push_back("k/" + std::to_string(i)); + values.push_back(AlignedValue(static_cast('a' + i % 26))); + } + + for (int threads : {1, 2, 8}) { + const std::string dir = MakeDir("threads_" + std::to_string(threads)); + auto cfg = BaseConfig(dir, /*direct_io=*/false, /*verify_crc=*/true, threads); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + + std::vector ptrs; + std::vector sizes; + for (auto& v : values) { + ptrs.push_back(v.data()); + sizes.push_back(v.size()); + } + CHECK(tier.WriteBatch(keys, ptrs, sizes)); + + std::vector dst_bufs(keys.size()); + std::vector dsts; + for (size_t i = 0; i < keys.size(); ++i) { + dst_bufs[i].Resize(values[i].size()); + dsts.push_back(reinterpret_cast(dst_bufs[i].data())); + } + auto ok = tier.ReadBatchIntoPtr(keys, dsts, sizes); + for (size_t i = 0; i < keys.size(); ++i) { + CHECK(ok[i]); + CHECK(std::memcmp(dst_bufs[i].data(), values[i].data(), values[i].size()) == 0); + } + fs::remove_all(dir); + } + std::cout << "OK\n"; +} + +// Capacity is charged in padded on-disk bytes, so a store of small values cannot +// admit more than the drive can hold. Before v3 this was accounted in raw value +// bytes, which for a 512 B value under-counted the footprint by 16x. +void test_capacity_charges_padded_bytes() { + std::cout << "test_capacity_charges_padded_bytes... "; + const std::string dir = MakeDir("capacity"); + const std::string key_a = "a"; + const size_t value_size = 512; + const size_t per_record = static_cast(RecordBytes(key_a.size(), value_size)); + CHECK(per_record == 2 * kRecordAlign); // 4 KiB prefix + 4 KiB padded value + + auto cfg = BaseConfig(dir, /*direct_io=*/false, /*verify_crc=*/true); + cfg.capacity_bytes = 2 * per_record; + SSDTier tier(dir, cfg.capacity_bytes, cfg); + + const std::vector value(value_size, 'q'); + CHECK(tier.Write("a", value.data(), value.size())); + CHECK(tier.Write("b", value.data(), value.size())); + // Third write must be refused: the tier is full in on-disk terms even though + // only 1536 raw value bytes have been stored. + CHECK(!tier.Write("c", value.data(), value.size())); + CHECK(tier.Capacity().first == 2 * per_record); + + fs::remove_all(dir); + std::cout << "OK\n"; +} + +int main() { + test_layout_is_aligned(); + test_aligned_buffer(); + test_round_trip(/*direct_io=*/false); + test_round_trip(/*direct_io=*/true); + test_locate_record(); + test_crc_disabled_round_trip(); + test_tier_threads_are_result_invariant(); + test_capacity_charges_padded_bytes(); + std::cout << "All direct-IO SSD tests passed\n"; + return 0; +} diff --git a/tests/cpp/umbp/local/test_ssd_tier.cpp b/tests/cpp/umbp/local/test_ssd_tier.cpp index 8f5bef7e7..96c09abf7 100644 --- a/tests/cpp/umbp/local/test_ssd_tier.cpp +++ b/tests/cpp/umbp/local/test_ssd_tier.cpp @@ -24,7 +24,7 @@ #include #include -#include "umbp/local/tiers/local_storage_manager.h" +#include "umbp/local/tiers/ssd_tier.h" using namespace mori::umbp; @@ -33,23 +33,22 @@ void test_segmented_recovery() { const std::string dir = "/tmp/umbp_test_segmented_recovery"; UMBPConfig cfg; - cfg.dram.capacity_bytes = 1024 * 1024; cfg.ssd.enabled = true; cfg.ssd.storage_dir = dir; cfg.ssd.capacity_bytes = 64 * 1024 * 1024; { - LocalStorageManager mgr(cfg); + SSDTier tier(dir, cfg.ssd.capacity_bytes, cfg.ssd); std::vector payload(4096, 'R'); - assert(mgr.Write("recover_key", payload.data(), payload.size(), StorageTier::LOCAL_SSD)); + assert(tier.Write("recover_key", payload.data(), payload.size())); } { - LocalStorageManager mgr(cfg); + SSDTier tier(dir, cfg.ssd.capacity_bytes, cfg.ssd); std::vector buf(4096, 0); - assert(mgr.ReadIntoPtr("recover_key", reinterpret_cast(buf.data()), buf.size())); + assert(tier.ReadIntoPtr("recover_key", reinterpret_cast(buf.data()), buf.size())); assert(buf == std::vector(4096, 'R')); - mgr.Clear(); + tier.Clear(); } std::cout << "PASSED" << std::endl; } @@ -59,52 +58,47 @@ void test_segmented_overwrite_generation() { const std::string dir = "/tmp/umbp_test_segmented_overwrite"; UMBPConfig cfg; - cfg.dram.capacity_bytes = 1024 * 1024; cfg.ssd.enabled = true; cfg.ssd.storage_dir = dir; cfg.ssd.capacity_bytes = 64 * 1024 * 1024; - LocalStorageManager mgr(cfg); + SSDTier tier(dir, cfg.ssd.capacity_bytes, cfg.ssd); std::vector a(1024, 'A'); std::vector b(1024, 'B'); - assert(mgr.Write("gen_key", a.data(), a.size(), StorageTier::LOCAL_SSD)); - assert(mgr.Write("gen_key", b.data(), b.size(), StorageTier::LOCAL_SSD)); + assert(tier.Write("gen_key", a.data(), a.size())); + assert(tier.Write("gen_key", b.data(), b.size())); std::vector out(1024, 0); - assert(mgr.ReadIntoPtr("gen_key", reinterpret_cast(out.data()), out.size())); + assert(tier.ReadIntoPtr("gen_key", reinterpret_cast(out.data()), out.size())); assert(out == b); - mgr.Clear(); + tier.Clear(); std::cout << "PASSED" << std::endl; } -void test_segmented_follower_refresh() { - std::cout << "test_segmented_follower_refresh... "; - const std::string dir = "/tmp/umbp_test_segmented_follower"; +void test_segmented_shared_reader_refresh() { + std::cout << "test_segmented_shared_reader_refresh... "; + const std::string dir = "/tmp/umbp_test_segmented_shared_reader"; - UMBPConfig leader_cfg; - leader_cfg.dram.capacity_bytes = 1024 * 1024; - leader_cfg.ssd.enabled = true; - leader_cfg.ssd.storage_dir = dir; - leader_cfg.ssd.capacity_bytes = 64 * 1024 * 1024; - - leader_cfg.role = UMBPRole::SharedSSDLeader; - - UMBPConfig follower_cfg = leader_cfg; - follower_cfg.role = UMBPRole::SharedSSDFollower; - follower_cfg.follower_mode = true; + UMBPConfig cfg; + cfg.ssd.enabled = true; + cfg.ssd.storage_dir = dir; + cfg.ssd.capacity_bytes = 64 * 1024 * 1024; - LocalStorageManager leader(leader_cfg); - LocalStorageManager follower(follower_cfg); + // Two tiers over one directory: the writer owns the segment log, the reader + // opens it ReadOnlyShared and has to pick up records it did not write by + // re-scanning. This is the sharing the deleted leader/follower roles were + // built on, and it lives in SSDTier -- not in any client. + SSDTier writer(dir, cfg.ssd.capacity_bytes, cfg.ssd); + SSDTier reader(dir, cfg.ssd.capacity_bytes, cfg.ssd, SSDAccessMode::ReadOnlyShared); std::vector payload(2048, 'F'); - assert(leader.Write("follower_key", payload.data(), payload.size(), StorageTier::LOCAL_SSD)); + assert(writer.Write("shared_key", payload.data(), payload.size())); std::vector out(2048, 0); - assert(follower.ReadIntoPtr("follower_key", reinterpret_cast(out.data()), out.size())); + assert(reader.ReadIntoPtr("shared_key", reinterpret_cast(out.data()), out.size())); assert(out == payload); - leader.Clear(); - follower.Clear(); + writer.Clear(); std::cout << "PASSED" << std::endl; } @@ -113,7 +107,6 @@ void test_segmented_io_uring_backend() { const std::string dir = "/tmp/umbp_test_segmented_io_uring"; UMBPConfig cfg; - cfg.dram.capacity_bytes = 1024 * 1024; cfg.ssd.enabled = true; cfg.ssd.storage_dir = dir; cfg.ssd.capacity_bytes = 64 * 1024 * 1024; @@ -122,21 +115,21 @@ void test_segmented_io_uring_backend() { cfg.ssd.durability.mode = UMBPDurabilityMode::Strict; cfg.ssd.io.queue_depth = 128; - std::unique_ptr mgr_ptr; + std::unique_ptr tier_ptr; try { - mgr_ptr = std::make_unique(cfg); + tier_ptr = std::make_unique(dir, cfg.ssd.capacity_bytes, cfg.ssd); } catch (const std::runtime_error& e) { std::cout << "SKIPPED (io_uring unavailable: " << e.what() << ")" << std::endl; return; } - auto& mgr = *mgr_ptr; + auto& tier = *tier_ptr; std::vector payload(4096, 'U'); - assert(mgr.Write("uring_key", payload.data(), payload.size(), StorageTier::LOCAL_SSD)); + assert(tier.Write("uring_key", payload.data(), payload.size())); std::vector out(4096, 0); - assert(mgr.ReadIntoPtr("uring_key", reinterpret_cast(out.data()), out.size())); + assert(tier.ReadIntoPtr("uring_key", reinterpret_cast(out.data()), out.size())); assert(out == payload); - mgr.Clear(); + tier.Clear(); std::cout << "PASSED" << std::endl; } @@ -145,12 +138,11 @@ void test_ssd_batch_read() { const std::string dir = "/tmp/umbp_test_ssd_batch_read"; UMBPConfig cfg; - cfg.dram.capacity_bytes = 1024 * 1024; cfg.ssd.enabled = true; cfg.ssd.storage_dir = dir; cfg.ssd.capacity_bytes = 64 * 1024 * 1024; - LocalStorageManager mgr(cfg); + SSDTier tier(dir, cfg.ssd.capacity_bytes, cfg.ssd); // Write 10 keys to SSD. constexpr size_t kNumKeys = 10; @@ -160,12 +152,11 @@ void test_ssd_batch_read() { for (size_t i = 0; i < kNumKeys; ++i) { keys[i] = "batch_read_key_" + std::to_string(i); payloads[i].assign(kValueSize, static_cast('A' + i)); - assert(mgr.Write(keys[i], payloads[i].data(), payloads[i].size(), StorageTier::LOCAL_SSD)); + assert(tier.Write(keys[i], payloads[i].data(), payloads[i].size())); } // Batch read all 10 via ReadBatchIntoPtr. - auto* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - assert(ssd != nullptr); + auto* ssd = &tier; std::vector> buffers(kNumKeys, std::vector(kValueSize, 0)); std::vector ptrs(kNumKeys); @@ -180,7 +171,7 @@ void test_ssd_batch_read() { assert(buffers[i] == payloads[i]); } - mgr.Clear(); + tier.Clear(); std::cout << "PASSED" << std::endl; } @@ -189,12 +180,11 @@ void test_ssd_batch_read_partial() { const std::string dir = "/tmp/umbp_test_ssd_batch_read_partial"; UMBPConfig cfg; - cfg.dram.capacity_bytes = 1024 * 1024; cfg.ssd.enabled = true; cfg.ssd.storage_dir = dir; cfg.ssd.capacity_bytes = 64 * 1024 * 1024; - LocalStorageManager mgr(cfg); + SSDTier tier(dir, cfg.ssd.capacity_bytes, cfg.ssd); // Write only 5 keys to SSD. constexpr size_t kExist = 5; @@ -207,12 +197,11 @@ void test_ssd_batch_read_partial() { } for (size_t i = 0; i < kExist; ++i) { payloads[i].assign(kValueSize, static_cast('X' + i)); - assert(mgr.Write(keys[i], payloads[i].data(), payloads[i].size(), StorageTier::LOCAL_SSD)); + assert(tier.Write(keys[i], payloads[i].data(), payloads[i].size())); } // Batch read all 7 keys — first 5 should succeed, last 2 fail. - auto* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - assert(ssd != nullptr); + auto* ssd = &tier; std::vector> buffers(kTotal, std::vector(kValueSize, 0)); std::vector ptrs(kTotal); @@ -230,7 +219,7 @@ void test_ssd_batch_read_partial() { assert(!results[i]); } - mgr.Clear(); + tier.Clear(); std::cout << "PASSED" << std::endl; } @@ -239,21 +228,20 @@ void test_ssd_batch_read_io_uring() { const std::string dir = "/tmp/umbp_test_ssd_batch_read_uring"; UMBPConfig cfg; - cfg.dram.capacity_bytes = 1024 * 1024; cfg.ssd.enabled = true; cfg.ssd.storage_dir = dir; cfg.ssd.capacity_bytes = 64 * 1024 * 1024; cfg.ssd.io.backend = UMBPIoBackend::IoUring; cfg.ssd.io.queue_depth = 128; - std::unique_ptr mgr_ptr; + std::unique_ptr tier_ptr; try { - mgr_ptr = std::make_unique(cfg); + tier_ptr = std::make_unique(dir, cfg.ssd.capacity_bytes, cfg.ssd); } catch (const std::runtime_error& e) { std::cout << "SKIPPED (io_uring unavailable: " << e.what() << ")" << std::endl; return; } - auto& mgr = *mgr_ptr; + auto& tier = *tier_ptr; constexpr size_t kNumKeys = 8; constexpr size_t kValueSize = 8192; @@ -262,11 +250,10 @@ void test_ssd_batch_read_io_uring() { for (size_t i = 0; i < kNumKeys; ++i) { keys[i] = "uring_batch_" + std::to_string(i); payloads[i].assign(kValueSize, static_cast('0' + i)); - assert(mgr.Write(keys[i], payloads[i].data(), payloads[i].size(), StorageTier::LOCAL_SSD)); + assert(tier.Write(keys[i], payloads[i].data(), payloads[i].size())); } - auto* ssd = mgr.GetTier(StorageTier::LOCAL_SSD); - assert(ssd != nullptr); + auto* ssd = &tier; std::vector> buffers(kNumKeys, std::vector(kValueSize, 0)); std::vector ptrs(kNumKeys); @@ -281,7 +268,7 @@ void test_ssd_batch_read_io_uring() { assert(buffers[i] == payloads[i]); } - mgr.Clear(); + tier.Clear(); std::cout << "PASSED" << std::endl; } @@ -289,7 +276,7 @@ int main() { std::cout << "=== Segmented SSD Tier Tests ===" << std::endl; test_segmented_recovery(); test_segmented_overwrite_generation(); - test_segmented_follower_refresh(); + test_segmented_shared_reader_refresh(); test_segmented_io_uring_backend(); test_ssd_batch_read(); test_ssd_batch_read_partial(); diff --git a/tests/cpp/umbp/local/test_ssd_tier_ranges.cpp b/tests/cpp/umbp/local/test_ssd_tier_ranges.cpp new file mode 100644 index 000000000..bc911f723 --- /dev/null +++ b/tests/cpp/umbp/local/test_ssd_tier_ranges.cpp @@ -0,0 +1,403 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// Ranged (sub-object) reads against the local SSD tier. +// +// The shape under test is the one the sglang tree connector produces: one +// stored object holds a KV page across every layer, and a load asks for the +// byte slices belonging to one group of layers, each into its own buffer. +// Both the direct and buffered variants run, because O_DIRECT changes which +// reads can land in the caller's buffer and which have to bounce through an +// aligned window -- and the results must be indistinguishable. +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/local/tiers/segment/segment_format.h" +#include "umbp/local/tiers/ssd_tier.h" + +namespace mori::umbp { +namespace { + +namespace fs = std::filesystem; + +// Point this at a real block-backed filesystem to exercise the O_DIRECT path; +// the default /tmp is overlayfs or tmpfs in most containers, which rejects +// O_DIRECT outright and makes the direct variants degrade to buffered. +std::string TestRoot() { + const char* root = std::getenv("UMBP_TEST_SSD_RANGES_DIR"); + return root && *root ? root : "/tmp"; +} + +std::string MakeDir(const std::string& name) { + const std::string dir = TestRoot() + "/umbp_test_ssd_ranges_" + name; + fs::remove_all(dir); + fs::create_directories(dir); + return dir; +} + +UMBPSsdConfig BaseConfig(const std::string& dir, bool direct_io) { + UMBPSsdConfig cfg; + cfg.enabled = true; + cfg.storage_dir = dir; + cfg.capacity_bytes = 64ULL * 1024 * 1024; + cfg.segment_size_bytes = 16ULL * 1024 * 1024; + cfg.io.backend = UMBPIoBackend::Posix; // no io_uring dependency in unit tests + cfg.direct_io = direct_io; + cfg.verify_crc = true; // ranged reads must not verify regardless + cfg.tier_io_threads = 4; + return cfg; +} + +// Byte i of the object is a function of i, so any slice read back can be +// checked against the offset it claims to come from. +std::vector PatternedValue(size_t size, int salt = 0) { + std::vector value(size); + for (size_t i = 0; i < size; ++i) { + value[i] = static_cast((i * 31 + salt * 7) & 0xFF); + } + return value; +} + +void ExpectSlice(const std::vector& got, const std::vector& object, size_t offset) { + ASSERT_LE(offset + got.size(), object.size()); + EXPECT_TRUE(std::equal(got.begin(), got.end(), object.begin() + offset)) + << "slice at offset " << offset << " size " << got.size() << " does not match the object"; +} + +uintptr_t Ptr(std::vector& buffer) { return reinterpret_cast(buffer.data()); } + +// Skips rather than silently passing when a direct-I/O variant could not get +// direct I/O -- otherwise it re-runs the buffered path under a name that claims +// to cover alignment widening, which is the one part only O_DIRECT reaches. +#define SKIP_IF_DIRECT_UNAVAILABLE(tier, want_direct) \ + do { \ + if ((want_direct) && !(tier).direct_io_active()) { \ + GTEST_SKIP() << "filesystem rejects O_DIRECT; set UMBP_TEST_SSD_RANGES_DIR to a " \ + "block-backed path to cover the aligned-window read path"; \ + } \ + } while (0) + +class SsdTierRangesTest : public ::testing::TestWithParam {}; + +INSTANTIATE_TEST_SUITE_P(BufferedAndDirect, SsdTierRangesTest, ::testing::Values(false, true), + [](const ::testing::TestParamInfo& info) { + return info.param ? "direct" : "buffered"; + }); + +// The layer-group case: consecutive slices of one object, each into its own +// destination. These coalesce into one device read and scatter on the way out, +// so a byte error here means the bounce bookkeeping is wrong. +TEST_P(SsdTierRangesTest, ContiguousSlicesScatterToSeparateBuffers) { + const std::string dir = MakeDir(GetParam() ? "contig_direct" : "contig_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + constexpr size_t kLayer = 3000; // deliberately not a kRecordAlign multiple + constexpr size_t kLayers = 8; + const auto object = PatternedValue(kLayer * kLayers); + ASSERT_TRUE(tier.Write("page", object.data(), object.size())); + + std::vector> outs(kLayers, std::vector(kLayer, 0)); + std::vector ptrs; + std::vector sizes; + std::vector offsets; + for (size_t layer = 0; layer < kLayers; ++layer) { + ptrs.push_back(Ptr(outs[layer])); + sizes.push_back(kLayer); + offsets.push_back(layer * kLayer); + } + + ASSERT_EQ(tier.ReadBatchRangesIntoPtr({"page"}, {ptrs}, {sizes}, {offsets}), + std::vector({true})); + for (size_t layer = 0; layer < kLayers; ++layer) { + ExpectSlice(outs[layer], object, layer * kLayer); + } +} + +// A single group out of the middle of the object, which is what layer-wise +// loading actually issues -- and the case that must NOT read the whole object. +TEST_P(SsdTierRangesTest, MiddleGroupReadsOnlyItsOwnSlices) { + const std::string dir = MakeDir(GetParam() ? "mid_direct" : "mid_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + constexpr size_t kLayer = 4096; + constexpr size_t kLayers = 16; + const auto object = PatternedValue(kLayer * kLayers, 3); + ASSERT_TRUE(tier.Write("page", object.data(), object.size())); + + std::vector> outs(4, std::vector(kLayer, 0)); + std::vector ptrs; + std::vector sizes; + std::vector offsets; + for (size_t j = 0; j < outs.size(); ++j) { + ptrs.push_back(Ptr(outs[j])); + sizes.push_back(kLayer); + offsets.push_back((8 + j) * kLayer); + } + + ASSERT_EQ(tier.ReadBatchRangesIntoPtr({"page"}, {ptrs}, {sizes}, {offsets}), + std::vector({true})); + for (size_t j = 0; j < outs.size(); ++j) { + ExpectSlice(outs[j], object, (8 + j) * kLayer); + } +} + +// Ranges given out of order, non-adjacent, and overlapping. None of these can +// coalesce, so every one becomes its own read; correctness must not depend on +// the caller sorting anything. +TEST_P(SsdTierRangesTest, ShuffledSparseAndOverlappingRangesAreByteCorrect) { + const std::string dir = MakeDir(GetParam() ? "shuf_direct" : "shuf_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + const auto object = PatternedValue(10000, 5); + ASSERT_TRUE(tier.Write("page", object.data(), object.size())); + + std::vector tail(500, 0), head(500, 0), overlap(700, 0), lone(1, 0); + const std::vector offsets{9500, 0, 300, 4321}; + ASSERT_EQ(tier.ReadBatchRangesIntoPtr({"page"}, {{Ptr(tail), Ptr(head), Ptr(overlap), Ptr(lone)}}, + {{tail.size(), head.size(), overlap.size(), lone.size()}}, + {offsets}), + std::vector({true})); + ExpectSlice(tail, object, 9500); + ExpectSlice(head, object, 0); + ExpectSlice(overlap, object, 300); + ExpectSlice(lone, object, 4321); +} + +// Reading the whole object through the ranged path must agree with the +// whole-object path byte for byte -- this is the degenerate case the tree +// connector's prefetch strategy uses. +TEST_P(SsdTierRangesTest, FullCoverageMatchesWholeObjectRead) { + const std::string dir = MakeDir(GetParam() ? "full_direct" : "full_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + const auto object = PatternedValue(20000, 9); + ASSERT_TRUE(tier.Write("page", object.data(), object.size())); + + std::vector whole(object.size(), 0); + ASSERT_TRUE(tier.ReadIntoPtr("page", Ptr(whole), whole.size())); + + std::vector> parts{std::vector(7000, 0), std::vector(7000, 0), + std::vector(6000, 0)}; + ASSERT_EQ(tier.ReadBatchRangesIntoPtr({"page"}, {{Ptr(parts[0]), Ptr(parts[1]), Ptr(parts[2])}}, + {{7000, 7000, 6000}}, {{0, 7000, 14000}}), + std::vector({true})); + + std::vector stitched; + for (const auto& part : parts) stitched.insert(stitched.end(), part.begin(), part.end()); + EXPECT_EQ(stitched, whole); + EXPECT_EQ(stitched, object); +} + +TEST_P(SsdTierRangesTest, ManyKeysInOneBatch) { + const std::string dir = MakeDir(GetParam() ? "many_direct" : "many_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + constexpr size_t kKeys = 32; + constexpr size_t kLayer = 1024; + std::vector keys; + std::vector> objects; + for (size_t k = 0; k < kKeys; ++k) { + keys.push_back("page_" + std::to_string(k)); + objects.push_back(PatternedValue(kLayer * 4, static_cast(k))); + ASSERT_TRUE(tier.Write(keys.back(), objects.back().data(), objects.back().size())); + } + + std::vector> outs(kKeys * 2, std::vector(kLayer, 0)); + std::vector> ptrs(kKeys); + std::vector> sizes(kKeys), offsets(kKeys); + for (size_t k = 0; k < kKeys; ++k) { + ptrs[k] = {Ptr(outs[k * 2]), Ptr(outs[k * 2 + 1])}; + sizes[k] = {kLayer, kLayer}; + offsets[k] = {kLayer, kLayer * 2}; // layers 1 and 2: adjacent, so they coalesce + } + + EXPECT_EQ(tier.ReadBatchRangesIntoPtr(keys, ptrs, sizes, offsets), + std::vector(kKeys, true)); + for (size_t k = 0; k < kKeys; ++k) { + ExpectSlice(outs[k * 2], objects[k], kLayer); + ExpectSlice(outs[k * 2 + 1], objects[k], kLayer * 2); + } +} + +TEST_P(SsdTierRangesTest, ShapeErrorsFailWholeBatch) { + const std::string dir = MakeDir(GetParam() ? "shape_direct" : "shape_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + const auto object = PatternedValue(4096); + ASSERT_TRUE(tier.Write("a", object.data(), object.size())); + ASSERT_TRUE(tier.Write("b", object.data(), object.size())); + + std::vector out_a(64, 0), out_b(64, 0); + // Two sizes but one pointer for "a": a malformed request, not a miss. + EXPECT_EQ(tier.ReadBatchRangesIntoPtr({"a", "b"}, {{Ptr(out_a)}, {Ptr(out_b)}}, {{64, 64}, {64}}, + {{0}, {0}}), + std::vector({false, false})); +} + +TEST_P(SsdTierRangesTest, DataErrorsAreIsolatedPerEntry) { + const std::string dir = MakeDir(GetParam() ? "iso_direct" : "iso_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + const auto object = PatternedValue(4096); + ASSERT_TRUE(tier.Write("present", object.data(), object.size())); + + std::vector good(64, 0), missing(64, 0), overflow(64, 0); + EXPECT_EQ(tier.ReadBatchRangesIntoPtr({"present", "absent", "present"}, + {{Ptr(good)}, {Ptr(missing)}, {Ptr(overflow)}}, + {{64}, {64}, {64}}, {{128}, {0}, {4090}}), + std::vector({true, false, false})); + ExpectSlice(good, object, 128); +} + +TEST_P(SsdTierRangesTest, RejectsOverflowZeroAndEmptyEntries) { + const std::string dir = MakeDir(GetParam() ? "rej_direct" : "rej_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + const auto object = PatternedValue(4096); + ASSERT_TRUE(tier.Write("key", object.data(), object.size())); + std::vector out(64, 0); + + EXPECT_EQ(tier.ReadBatchRangesIntoPtr({"key"}, {{Ptr(out)}}, {{64}}, + {{std::numeric_limits::max() - 8}}), + std::vector({false})); + EXPECT_EQ(tier.ReadBatchRangesIntoPtr({"key"}, {{Ptr(out)}}, {{0}}, {{0}}), + std::vector({false})); + EXPECT_EQ(tier.ReadBatchRangesIntoPtr({"key"}, {{}}, {{}}, {{}}), std::vector({false})); + EXPECT_EQ(tier.ReadBatchRangesIntoPtr({}, {}, {}, {}), std::vector({})); +} + +// A value written with checksums on must still be readable through the ranged +// path, which does not (and cannot) verify them. +TEST_P(SsdTierRangesTest, RangedReadIgnoresRecordChecksum) { + const std::string dir = MakeDir(GetParam() ? "crc_direct" : "crc_buffered"); + auto cfg = BaseConfig(dir, GetParam()); + ASSERT_TRUE(cfg.verify_crc); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + SKIP_IF_DIRECT_UNAVAILABLE(tier, GetParam()); + + const auto object = PatternedValue(8192, 11); + ASSERT_TRUE(tier.Write("page", object.data(), object.size())); + + std::vector slice(1000, 0); + ASSERT_EQ(tier.ReadBatchRangesIntoPtr({"page"}, {{Ptr(slice)}}, {{1000}}, {{2000}}), + std::vector({true})); + ExpectSlice(slice, object, 2000); +} +// The tree connector loads straight into GPU buffers, so the tier has to accept +// a device destination. read() cannot write one, which is why these have to +// bounce through host memory -- a regression here shows up as a silent miss +// (EFAULT reported as "key not found"), not as a crash. +TEST(SsdTierRangesDeviceTest, ScattersIntoDeviceBuffers) { + int device_count = 0; + if (hipGetDeviceCount(&device_count) != hipSuccess || device_count == 0) { + GTEST_SKIP() << "no GPU available"; + } + + const std::string dir = MakeDir("device"); + auto cfg = BaseConfig(dir, /*direct_io=*/false); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + + constexpr size_t kLayer = 4096; + constexpr size_t kLayers = 4; + const auto object = PatternedValue(kLayer * kLayers, 29); + ASSERT_TRUE(tier.Write("page", object.data(), object.size())); + + void* device_buffer = nullptr; + ASSERT_EQ(hipMalloc(&device_buffer, kLayer * kLayers), hipSuccess); + ASSERT_EQ(hipMemset(device_buffer, 0, kLayer * kLayers), hipSuccess); + + std::vector ptrs; + std::vector sizes; + std::vector offsets; + for (size_t layer = 0; layer < kLayers; ++layer) { + ptrs.push_back(reinterpret_cast(device_buffer) + layer * kLayer); + sizes.push_back(kLayer); + offsets.push_back(layer * kLayer); + } + + ASSERT_EQ(tier.ReadBatchRangesIntoPtr({"page"}, {ptrs}, {sizes}, {offsets}), + std::vector({true})); + + std::vector host_result(kLayer * kLayers, 0); + ASSERT_EQ(hipMemcpy(host_result.data(), device_buffer, host_result.size(), hipMemcpyDeviceToHost), + hipSuccess); + EXPECT_EQ(host_result, object); + ASSERT_EQ(hipFree(device_buffer), hipSuccess); +} + +// A single non-contiguous slice, which would otherwise take the "read straight +// into the caller" fast path -- the path a device pointer must be kept off. +TEST(SsdTierRangesDeviceTest, SingleSliceIntoDeviceBufferBounces) { + int device_count = 0; + if (hipGetDeviceCount(&device_count) != hipSuccess || device_count == 0) { + GTEST_SKIP() << "no GPU available"; + } + + const std::string dir = MakeDir("device_single"); + auto cfg = BaseConfig(dir, /*direct_io=*/false); + SSDTier tier(dir, cfg.capacity_bytes, cfg); + + const auto object = PatternedValue(20000, 31); + ASSERT_TRUE(tier.Write("page", object.data(), object.size())); + + constexpr size_t kSlice = 4096; + void* device_buffer = nullptr; + ASSERT_EQ(hipMalloc(&device_buffer, kSlice), hipSuccess); + ASSERT_EQ(hipMemset(device_buffer, 0, kSlice), hipSuccess); + + ASSERT_EQ(tier.ReadBatchRangesIntoPtr({"page"}, {{reinterpret_cast(device_buffer)}}, + {{kSlice}}, {{8192}}), + std::vector({true})); + + std::vector host_result(kSlice, 0); + ASSERT_EQ(hipMemcpy(host_result.data(), device_buffer, kSlice, hipMemcpyDeviceToHost), + hipSuccess); + ExpectSlice(host_result, object, 8192); + ASSERT_EQ(hipFree(device_buffer), hipSuccess); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/local/test_umbp_client.cpp b/tests/cpp/umbp/local/test_umbp_client.cpp deleted file mode 100644 index 497191b7c..000000000 --- a/tests/cpp/umbp/local/test_umbp_client.cpp +++ /dev/null @@ -1,340 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include -#include -#include -#include - -#include "umbp/local/standalone_client.h" - -using namespace mori::umbp; - -void test_put_get() { - std::cout << "test_put_get... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = false; - - StandaloneClient client(config); - - std::vector data(4096, 'A'); - assert(client.Put("key1", data.data(), data.size())); - assert(client.Exists("key1")); - - std::vector buf(4096, 0); - assert(client.Get("key1", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); - - std::cout << "PASSED" << std::endl; -} - -void test_put_dedup() { - std::cout << "test_put_dedup... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = false; - - StandaloneClient client(config); - - std::vector data(4096, 'B'); - assert(client.Put("key1", data.data(), data.size())); - - // Second put with same key should return true (dedup) - std::vector data2(4096, 'C'); - assert(client.Put("key1", data2.data(), data2.size())); - - // Data should still be original (dedup skipped write) - std::vector buf(4096, 0); - assert(client.Get("key1", reinterpret_cast(buf.data()), buf.size())); - assert(buf == data); // Original data, not data2 - - std::cout << "PASSED" << std::endl; -} - -void test_remove() { - std::cout << "test_remove... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = false; - - StandaloneClient client(config); - - std::vector data(4096, 'D'); - assert(client.Put("key1", data.data(), data.size())); - assert(client.Exists("key1")); - - assert(client.Remove("key1")); - assert(!client.Exists("key1")); - assert(!client.Remove("key1")); // Already removed - - std::cout << "PASSED" << std::endl; -} - -void test_batch_put_get() { - std::cout << "test_batch_put_get... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = false; - - StandaloneClient client(config); - - const int N = 10; - std::vector keys; - std::vector> all_data; - std::vector ptrs; - std::vector sizes; - - for (int i = 0; i < N; ++i) { - keys.push_back("batch_key_" + std::to_string(i)); - all_data.emplace_back(1024, static_cast('0' + i)); - ptrs.push_back(reinterpret_cast(all_data.back().data())); - sizes.push_back(1024); - } - - auto put_results = client.BatchPut(keys, ptrs, sizes); - for (int i = 0; i < N; ++i) { - assert(put_results[i]); - } - - // Batch exists - auto exists_results = client.BatchExists(keys); - for (int i = 0; i < N; ++i) { - assert(exists_results[i]); - } - - // Batch get - std::vector> read_bufs(N, std::vector(1024, 0)); - std::vector dst_ptrs; - for (int i = 0; i < N; ++i) { - dst_ptrs.push_back(reinterpret_cast(read_bufs[i].data())); - } - - auto get_results = client.BatchGet(keys, dst_ptrs, sizes); - for (int i = 0; i < N; ++i) { - assert(get_results[i]); - assert(read_bufs[i] == all_data[i]); - } - - std::cout << "PASSED" << std::endl; -} - -void test_clear() { - std::cout << "test_clear... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = false; - - StandaloneClient client(config); - - std::vector data(4096, 'E'); - client.Put("key1", data.data(), data.size()); - client.Put("key2", data.data(), data.size()); - assert(client.Exists("key1")); - assert(client.Exists("key2")); - - client.Clear(); - assert(!client.Exists("key1")); - assert(!client.Exists("key2")); - - std::cout << "PASSED" << std::endl; -} - -void test_put_from_ptr_get_into_ptr() { - std::cout << "test_put_from_ptr_get_into_ptr... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1 * 1024 * 1024; - config.ssd.enabled = false; - - StandaloneClient client(config); - - std::vector data(8192, 'Z'); - uintptr_t src = reinterpret_cast(data.data()); - assert(client.Put("ptr_key", src, data.size())); - - std::vector buf(8192, 0); - uintptr_t dst = reinterpret_cast(buf.data()); - assert(client.Get("ptr_key", dst, buf.size())); - assert(buf == data); - - std::cout << "PASSED" << std::endl; -} - -void test_dram_full_demote_with_index() { - std::cout << "test_dram_full_demote_with_index... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1024; // 1 KB - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_client_demote"; - config.ssd.capacity_bytes = 10 * 1024 * 1024; - - StandaloneClient client(config); - - // Fill DRAM: 2 x 512 bytes - std::vector d1(512, 'A'); - assert(client.Put("k1", d1.data(), d1.size())); - std::vector d2(512, 'B'); - assert(client.Put("k2", d2.data(), d2.size())); - - // This write forces auto-demote of k1 (LRU) to SSD - std::vector d3(512, 'C'); - assert(client.Put("k3", d3.data(), d3.size())); - - // k1 should still be accessible (demoted to SSD, not lost) - assert(client.Exists("k1")); - - // k1's index should say LOCAL_SSD - auto loc = client.Index().Lookup("k1"); - assert(loc.has_value()); - assert(loc->tier == StorageTier::LOCAL_SSD); - - // k3's index should say CPU_DRAM - auto loc3 = client.Index().Lookup("k3"); - assert(loc3.has_value()); - assert(loc3->tier == StorageTier::CPU_DRAM); - - // Data integrity: read k1 back from SSD - std::vector buf(512, 0); - assert(client.Get("k1", reinterpret_cast(buf.data()), buf.size())); - assert(buf == d1); - - client.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_batch_get_ssd_tier() { - std::cout << "test_batch_get_ssd_tier... "; - - UMBPConfig config; - config.dram.capacity_bytes = 1024; // Tiny DRAM: forces SSD demotion - config.ssd.enabled = true; - config.ssd.storage_dir = "/tmp/umbp_test_batch_get_ssd"; - config.ssd.capacity_bytes = 10 * 1024 * 1024; - - StandaloneClient client(config); - - // Write 10 x 512-byte values; DRAM holds ~2, rest auto-demote to SSD. - constexpr int N = 10; - constexpr size_t kSize = 512; - std::vector keys(N); - std::vector> data(N); - for (int i = 0; i < N; ++i) { - keys[i] = "ssd_batch_" + std::to_string(i); - data[i].assign(kSize, static_cast('A' + i)); - assert(client.Put(keys[i], data[i].data(), data[i].size())); - } - - // Batch read all 10 — mix of DRAM and SSD. - std::vector> bufs(N, std::vector(kSize, 0)); - std::vector ptrs(N); - std::vector sizes(N, kSize); - for (int i = 0; i < N; ++i) { - ptrs[i] = reinterpret_cast(bufs[i].data()); - } - - auto results = client.BatchGet(keys, ptrs, sizes); - for (int i = 0; i < N; ++i) { - assert(results[i]); - assert(bufs[i] == data[i]); - } - - client.Clear(); - std::cout << "PASSED" << std::endl; -} - -void test_batch_get_follower() { - std::cout << "test_batch_get_follower... "; - - const std::string ssd_dir = "/tmp/umbp_test_batch_get_follower"; - - // Leader: writes to DRAM + copies to shared SSD. - UMBPConfig leader_cfg; - leader_cfg.dram.capacity_bytes = 1024 * 1024; - leader_cfg.ssd.enabled = true; - leader_cfg.ssd.storage_dir = ssd_dir; - leader_cfg.ssd.capacity_bytes = 10 * 1024 * 1024; - leader_cfg.role = UMBPRole::SharedSSDLeader; - leader_cfg.force_ssd_copy_on_write = true; - leader_cfg.copy_pipeline.async_enabled = false; // sync copy for determinism - - constexpr int N = 5; - constexpr size_t kSize = 2048; - std::vector keys(N); - std::vector> data(N); - - { - StandaloneClient leader(leader_cfg); - for (int i = 0; i < N; ++i) { - keys[i] = "follower_batch_" + std::to_string(i); - data[i].assign(kSize, static_cast('P' + i)); - assert(leader.Put(keys[i], data[i].data(), data[i].size())); - } - } // leader destructor drains copy pipeline - - // Follower: reads from shared SSD via BatchGet. - UMBPConfig follower_cfg; - follower_cfg.dram.capacity_bytes = 1024 * 1024; - follower_cfg.ssd.enabled = true; - follower_cfg.ssd.storage_dir = ssd_dir; - follower_cfg.ssd.capacity_bytes = 10 * 1024 * 1024; - follower_cfg.role = UMBPRole::SharedSSDFollower; - follower_cfg.follower_mode = true; - - StandaloneClient follower(follower_cfg); - - std::vector> bufs(N, std::vector(kSize, 0)); - std::vector ptrs(N); - std::vector sizes(N, kSize); - for (int i = 0; i < N; ++i) { - ptrs[i] = reinterpret_cast(bufs[i].data()); - } - - auto results = follower.BatchGet(keys, ptrs, sizes); - for (int i = 0; i < N; ++i) { - assert(results[i]); - assert(bufs[i] == data[i]); - } - - follower.Clear(); - std::cout << "PASSED" << std::endl; -} - -int main() { - std::cout << "=== StandaloneClient Tests ===" << std::endl; - test_put_get(); - test_put_dedup(); - test_remove(); - test_batch_put_get(); - test_clear(); - test_put_from_ptr_get_into_ptr(); - test_dram_full_demote_with_index(); - test_batch_get_ssd_tier(); - test_batch_get_follower(); - std::cout << "All StandaloneClient tests passed!" << std::endl; - return 0; -} diff --git a/tests/cpp/umbp/local/test_umbp_e2e.cpp b/tests/cpp/umbp/local/test_umbp_e2e.cpp deleted file mode 100644 index 9e4b6612b..000000000 --- a/tests/cpp/umbp/local/test_umbp_e2e.cpp +++ /dev/null @@ -1,879 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -// -// test_umbp_e2e: End-to-end integration tests for UMBP. -// Runs POSIX tests always. SPDK tests run only if UMBP_SPDK_NVME_PCI is set. -// All tests in one executable — run: ./test_umbp_e2e -// -// Test categories: -// [POSIX] Always available -// [ROLE] Role deduction logic (no I/O) -// [SPDK] Requires UMBP_SPDK_NVME_PCI (skipped if absent) -// [PROXY] Requires SPDK + Linux fork (skipped if absent) - -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef __unix__ -#include -#endif - -#include "umbp/common/config.h" -#include "umbp/local/standalone_client.h" - -#ifdef __linux__ -#include -#include -#include - -#include "umbp/local/tiers/spdk_proxy_tier.h" -#include "umbp/storage/spdk/proxy/spdk_proxy_protocol.h" -#include "umbp/storage/spdk/proxy/spdk_proxy_shm.h" -#endif - -using namespace mori::umbp; - -// --------------------------------------------------------------------------- -// Test infrastructure -// --------------------------------------------------------------------------- -static int g_passed = 0, g_failed = 0, g_skipped = 0; - -#define RUN_TEST(fn) \ - do { \ - printf(" %-50s", #fn "..."); \ - fflush(stdout); \ - bool _ok = fn(); \ - if (_ok) { \ - g_passed++; \ - printf("PASS\n"); \ - } else { \ - g_failed++; \ - printf("FAIL\n"); \ - } \ - } while (0) - -#define SKIP_TEST(name, reason) \ - do { \ - printf(" %-50s", name "..."); \ - g_skipped++; \ - printf("SKIP (%s)\n", reason); \ - } while (0) - -#define CHECK(cond) \ - do { \ - if (!(cond)) { \ - fprintf(stderr, "\n CHECK FAILED: %s (%s:%d)\n", #cond, __FILE__, __LINE__); \ - return false; \ - } \ - } while (0) - -// --------------------------------------------------------------------------- -// Helper: create a UMBPConfig for POSIX testing -// --------------------------------------------------------------------------- -static UMBPConfig MakePosixConfig(size_t dram_mb = 64, size_t ssd_mb = 256) { - UMBPConfig cfg; - cfg.dram.capacity_bytes = dram_mb * 1024 * 1024; - cfg.ssd.capacity_bytes = ssd_mb * 1024 * 1024; - cfg.ssd.io.backend = UMBPIoBackend::Posix; - cfg.ssd.ssd_backend = "file"; - cfg.ssd.storage_dir = "/tmp/umbp_e2e_test_" + std::to_string(getpid()); - cfg.role = UMBPRole::Standalone; - return cfg; -} - -// ========================================================================= -// [ROLE] Tests -// ========================================================================= - -static bool test_role_default_standalone() { - UMBPConfig cfg; - CHECK(cfg.ResolveRole() == UMBPRole::Standalone); - return true; -} - -static bool test_role_explicit_leader() { - UMBPConfig cfg; - cfg.role = UMBPRole::SharedSSDLeader; - CHECK(cfg.ResolveRole() == UMBPRole::SharedSSDLeader); - return true; -} - -static bool test_role_explicit_follower() { - UMBPConfig cfg; - cfg.role = UMBPRole::SharedSSDFollower; - CHECK(cfg.ResolveRole() == UMBPRole::SharedSSDFollower); - return true; -} - -static bool test_role_backward_compat_follower_mode() { - UMBPConfig cfg; - cfg.follower_mode = true; - CHECK(cfg.ResolveRole() == UMBPRole::SharedSSDFollower); - return true; -} - -static bool test_role_backward_compat_force_cow() { - UMBPConfig cfg; - cfg.force_ssd_copy_on_write = true; - CHECK(cfg.ResolveRole() == UMBPRole::SharedSSDLeader); - return true; -} - -static bool test_auto_rank_sentinel() { - CHECK(kAutoRankId == UINT32_MAX); - // Without UMBP_SPDK_PROXY_RANK env, FromEnvironment should give kAutoRankId - // (can't reliably unsetenv in test, so just verify the constant) - return true; -} - -#ifdef __linux__ -static bool test_role_from_local_rank_env() { - // Fork a child that sets LOCAL_RANK and checks role - int pipefd[2]; - CHECK(pipe(pipefd) == 0); - - pid_t pid = fork(); - if (pid == 0) { - close(pipefd[0]); - // Ensure no prior UMBP_ROLE - unsetenv("UMBP_ROLE"); - - // Test LOCAL_RANK=0 → Leader - setenv("LOCAL_RANK", "0", 1); - auto cfg0 = UMBPConfig::FromEnvironment(); - uint8_t r0 = (cfg0.role == UMBPRole::SharedSSDLeader) ? 1 : 0; - - // Test LOCAL_RANK=3 → Follower - setenv("LOCAL_RANK", "3", 1); - auto cfg3 = UMBPConfig::FromEnvironment(); - uint8_t r3 = (cfg3.role == UMBPRole::SharedSSDFollower) ? 1 : 0; - - uint8_t result = r0 & r3; - write(pipefd[1], &result, 1); - close(pipefd[1]); - _exit(0); - } - - close(pipefd[1]); - uint8_t result = 0; - read(pipefd[0], &result, 1); - close(pipefd[0]); - waitpid(pid, nullptr, 0); - - CHECK(result == 1); - return true; -} -#endif - -// ========================================================================= -// [POSIX] Tests -// ========================================================================= - -static bool test_posix_standalone_write_read() { - auto cfg = MakePosixConfig(); - StandaloneClient client(cfg); - - std::string key = "test_posix_wr_1"; - std::vector data(4096, 'A'); - CHECK(client.Put(key, data.data(), data.size())); - - std::vector buf(4096, 0); - CHECK(client.Get(key, reinterpret_cast(buf.data()), buf.size())); - CHECK(buf == data); - - client.Clear(); - return true; -} - -static bool test_posix_batch_write_read() { - auto cfg = MakePosixConfig(); - StandaloneClient client(cfg); - - const int N = 100; - const size_t sz = 8192; - std::vector keys(N); - std::vector> bufs(N); - std::vector ptrs(N); - std::vector sizes(N, sz); - - for (int i = 0; i < N; ++i) { - keys[i] = "posix_batch_" + std::to_string(i); - bufs[i].resize(sz, static_cast(i & 0xFF)); - ptrs[i] = reinterpret_cast(bufs[i].data()); - } - - auto wr = client.BatchPut(keys, ptrs, sizes); - int write_ok = 0; - for (auto b : wr) write_ok += b; - CHECK(write_ok == N); - - std::vector> read_bufs(N, std::vector(sz, 0)); - std::vector dst_ptrs(N); - for (int i = 0; i < N; ++i) dst_ptrs[i] = reinterpret_cast(read_bufs[i].data()); - - auto rr = client.BatchGet(keys, dst_ptrs, sizes); - int read_ok = 0; - for (auto b : rr) read_ok += b; - CHECK(read_ok == N); - - for (int i = 0; i < N; ++i) CHECK(read_bufs[i] == bufs[i]); - - client.Clear(); - return true; -} - -static bool test_posix_dedup() { - auto cfg = MakePosixConfig(); - StandaloneClient client(cfg); - - std::vector data(1024, 'X'); - CHECK(client.Put("dedup_key", data.data(), data.size())); - CHECK(client.Put("dedup_key", data.data(), data.size())); // should dedup - CHECK(client.Exists("dedup_key")); - - client.Clear(); - return true; -} - -static bool test_posix_evict() { - auto cfg = MakePosixConfig(); - StandaloneClient client(cfg); - - std::vector data(1024, 'E'); - CHECK(client.Put("evict_key", data.data(), data.size())); - CHECK(client.Exists("evict_key")); - CHECK(client.Remove("evict_key")); - CHECK(!client.Exists("evict_key")); - - client.Clear(); - return true; -} - -static bool test_posix_ssd_direct_write_read() { - auto cfg = MakePosixConfig(); - StandaloneClient client(cfg); - - std::vector data(32768, 'S'); - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - CHECK(ssd != nullptr); - CHECK(ssd->Write("ssd_direct_key", data.data(), data.size())); - - std::vector buf(32768, 0); - CHECK(ssd->ReadIntoPtr("ssd_direct_key", reinterpret_cast(buf.data()), buf.size())); - CHECK(buf == data); - - client.Clear(); - return true; -} - -static bool test_posix_dram_evict_to_ssd() { - // Small DRAM, force demotion to SSD - auto cfg = MakePosixConfig(1 /* 1MB DRAM */, 256); - StandaloneClient client(cfg); - - const size_t sz = 128 * 1024; // 128KB each - std::vector data(sz, 'D'); - - // Write 10 items → 1.25MB total, DRAM=1MB → some must demote to SSD - for (int i = 0; i < 10; ++i) { - std::string key = "dram_evict_" + std::to_string(i); - CHECK(client.Put(key, data.data(), data.size())); - } - - // All should still be readable (from DRAM or SSD) - for (int i = 0; i < 10; ++i) { - std::string key = "dram_evict_" + std::to_string(i); - CHECK(client.Exists(key)); - std::vector buf(sz, 0); - CHECK(client.Get(key, reinterpret_cast(buf.data()), sz)); - CHECK(buf == data); - } - - client.Clear(); - return true; -} - -static bool test_posix_capacity() { - auto cfg = MakePosixConfig(64, 256); - StandaloneClient client(cfg); - - auto* dram = client.Storage().GetTier(StorageTier::CPU_DRAM); - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - CHECK(dram != nullptr); - CHECK(ssd != nullptr); - - auto [dram_used, dram_total] = dram->Capacity(); - CHECK(dram_total > 0); - - auto [ssd_used, ssd_total] = ssd->Capacity(); - CHECK(ssd_total > 0); - - client.Clear(); - return true; -} - -static bool test_posix_empty_read() { - auto cfg = MakePosixConfig(); - StandaloneClient client(cfg); - - std::vector buf(1024, 0); - CHECK(!client.Get("nonexistent", reinterpret_cast(buf.data()), 1024)); - CHECK(!client.Exists("nonexistent")); - - return true; -} - -static bool test_posix_large_value() { - auto cfg = MakePosixConfig(16, 512); - StandaloneClient client(cfg); - - const size_t sz = 4 * 1024 * 1024; // 4MB - std::vector data(sz); - for (size_t i = 0; i < sz; ++i) data[i] = static_cast(i & 0xFF); - - // Write directly to SSD (4MB > 16MB DRAM) - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - CHECK(ssd != nullptr); - CHECK(ssd->Write("large_val", data.data(), data.size())); - - std::vector buf(sz, 0); - CHECK(ssd->ReadIntoPtr("large_val", reinterpret_cast(buf.data()), sz)); - CHECK(buf == data); - - client.Clear(); - return true; -} - -// ========================================================================= -// [SPDK] Tests — only if UMBP_SPDK_NVME_PCI is set -// ========================================================================= - -#ifdef __linux__ - -// Kill any residual spdk_proxy daemon and clean stale SHM before tests. -// SPDK takes an exclusive lock on NVMe cores — leftover processes from -// previous runs (manual daemon, crashed test, etc.) will block all SPDK init. -static void CleanupResidualSpdk() { - // Kill any lingering spdk_proxy processes - system("pkill -9 -x spdk_proxy 2>/dev/null"); - usleep(500000); // 500ms for process to die - - // Remove stale SHM files - umbp::proxy::ProxyShmRegion::CleanupStale(umbp::proxy::kDefaultShmName); - - // Also remove SPDK's per-pid lock files that may be stale - system("rm -f /var/tmp/spdk_pid*.lock 2>/dev/null"); -} - -// Run in a forked child so that SpdkEnv singleton (which holds NVMe core -// locks) is destroyed when the child exits — preventing it from blocking -// subsequent proxy tests in the parent process. -static bool test_spdk_standalone_write_read() { - pid_t pid = fork(); - if (pid == 0) { - UMBPConfig cfg; - cfg.ssd.ssd_backend = "spdk"; - cfg.role = UMBPRole::Standalone; - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - cfg.ssd.capacity_bytes = 1024ULL * 1024 * 1024; - - auto env_cfg = UMBPConfig::FromEnvironment(); - cfg.ssd.spdk_nvme_pci_addr = env_cfg.ssd.spdk_nvme_pci_addr; - cfg.ssd.spdk_reactor_mask = env_cfg.ssd.spdk_reactor_mask; - cfg.ssd.spdk_mem_size_mb = env_cfg.ssd.spdk_mem_size_mb; - cfg.ssd.spdk_io_workers = env_cfg.ssd.spdk_io_workers; - - StandaloneClient client(cfg); - - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) _exit(1); - - const size_t sz = 65536; - std::vector data(sz, 'Z'); - if (!ssd->Write("spdk_standalone_1", data.data(), data.size())) _exit(2); - - std::vector buf(sz, 0); - if (!ssd->ReadIntoPtr("spdk_standalone_1", reinterpret_cast(buf.data()), sz)) - _exit(3); - if (buf != data) _exit(4); - - client.Clear(); - _exit(0); - } - int status; - waitpid(pid, &status, 0); - CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 0); - return true; -} - -// --------------------------------------------------------------------------- -// [PROXY] Tests — Leader auto-fork + Follower connect -// --------------------------------------------------------------------------- - -// Helper: fork, set env, run function, return exit code -static int ForkAndRun(const char* local_rank, const char* backend, int (*fn)(const char*), - const char* arg = nullptr) { - pid_t pid = fork(); - if (pid == 0) { - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - setenv("LOCAL_RANK", local_rank, 1); - if (backend) setenv("UMBP_SSD_BACKEND", backend, 1); - _exit(fn(arg)); - } - int status; - waitpid(pid, &status, 0); - return WIFEXITED(status) ? WEXITSTATUS(status) : -1; -} - -// Signal file paths for leader-follower coordination. -// Two-phase handshake: leader signals "data ready", follower signals "done". -static constexpr const char* kLeaderReadyFile = "/tmp/umbp_e2e_leader_ready"; -static constexpr const char* kFollowerDoneFile = "/tmp/umbp_e2e_follower_done"; - -static bool WaitForFile(const char* path, int timeout_ms) { - for (int elapsed = 0; elapsed < timeout_ms; elapsed += 100) { - if (access(path, F_OK) == 0) return true; - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - } - return false; -} - -static void TouchFile(const char* path) { - int fd = open(path, O_CREAT | O_WRONLY, 0644); - if (fd >= 0) close(fd); -} - -static int leader_write_and_wait(const char*) { - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - StandaloneClient client(cfg); - - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) return 1; - - for (int i = 0; i < 10; ++i) { - std::string key = "proxy_shared_" + std::to_string(i); - std::vector data(4096, static_cast('A' + i)); - if (!ssd->Write(key, data.data(), data.size())) return 2; - } - - // Tell follower data is ready - TouchFile(kLeaderReadyFile); - - // Keep StandaloneClient (and proxy daemon) alive until follower finishes (60s) - if (!WaitForFile(kFollowerDoneFile, 60000)) { - fprintf(stderr, " leader: timed out waiting for follower\n"); - } - return 0; -} - -static int follower_read_and_verify(const char*) { - // Wait until leader signals that data is written (30s) - if (!WaitForFile(kLeaderReadyFile, 30000)) { - fprintf(stderr, " follower: timed out waiting for leader ready\n"); - return 5; - } - - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - StandaloneClient client(cfg); - - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) { - TouchFile(kFollowerDoneFile); - return 1; - } - - int ok = 0; - for (int i = 0; i < 10; ++i) { - std::string key = "proxy_shared_" + std::to_string(i); - std::vector buf(4096, 0); - if (ssd->ReadIntoPtr(key, reinterpret_cast(buf.data()), 4096)) { - std::vector expected(4096, static_cast('A' + i)); - if (buf == expected) ok++; - } - } - - TouchFile(kFollowerDoneFile); - return (ok == 10) ? 0 : 3; -} - -static bool test_proxy_auto_fork_leader_follower() { - unlink(kLeaderReadyFile); - unlink(kFollowerDoneFile); - - pid_t leader = fork(); - if (leader == 0) { - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - setenv("LOCAL_RANK", "0", 1); - setenv("UMBP_SSD_BACKEND", "spdk", 1); - _exit(leader_write_and_wait(nullptr)); - } - - pid_t follower = fork(); - if (follower == 0) { - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - setenv("LOCAL_RANK", "1", 1); - setenv("UMBP_SSD_BACKEND", "spdk", 1); - _exit(follower_read_and_verify(nullptr)); - } - - int lstatus, fstatus; - waitpid(leader, &lstatus, 0); - waitpid(follower, &fstatus, 0); - - unlink(kLeaderReadyFile); - unlink(kFollowerDoneFile); - - bool leader_ok = WIFEXITED(lstatus) && WEXITSTATUS(lstatus) == 0; - bool follower_ok = WIFEXITED(fstatus) && WEXITSTATUS(fstatus) == 0; - - if (!leader_ok) - fprintf(stderr, " leader exit=%d\n", WIFEXITED(lstatus) ? WEXITSTATUS(lstatus) : -1); - if (!follower_ok) - fprintf(stderr, " follower exit=%d\n", WIFEXITED(fstatus) ? WEXITSTATUS(fstatus) : -1); - - CHECK(leader_ok); - CHECK(follower_ok); - return true; -} - -static bool test_proxy_daemon_cleanup_after_leader_exit() { - // Service-mode proxy should survive the first client exiting, remain - // usable by a later client, then idle-exit after the last client leaves. - setenv("UMBP_SPDK_PROXY_IDLE_EXIT_TIMEOUT_MS", "2000", 1); - - pid_t leader = fork(); - if (leader == 0) { - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - setenv("LOCAL_RANK", "0", 1); - setenv("UMBP_SSD_BACKEND", "spdk", 1); - - int rc = 0; - { - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - StandaloneClient client(cfg); - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) rc = 1; - std::vector data(4096, 'Q'); - if (rc == 0 && !ssd->Write("proxy_survive_key", data.data(), data.size())) rc = 2; - } - _exit(rc); - } - - int status; - waitpid(leader, &status, 0); - CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 0); - - pid_t follower = fork(); - if (follower == 0) { - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - setenv("LOCAL_RANK", "1", 1); - setenv("UMBP_SSD_BACKEND", "spdk", 1); - - int rc = 0; - { - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - StandaloneClient client(cfg); - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) rc = 3; - - std::vector buf(4096, 0); - if (rc == 0 && !ssd->ReadIntoPtr("proxy_survive_key", reinterpret_cast(buf.data()), - buf.size())) { - rc = 4; - } - std::vector expected(4096, 'Q'); - if (rc == 0 && buf != expected) rc = 5; - } - _exit(rc); - } - - waitpid(follower, &status, 0); - CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 0); - - // Wait for idle-exit after the last client disconnects. - std::this_thread::sleep_for(std::chrono::seconds(4)); - int probe = umbp::proxy::ProxyShmRegion::ProbeExisting(umbp::proxy::kDefaultShmName); - CHECK(probe <= 0); - - unsetenv("UMBP_SPDK_PROXY_IDLE_EXIT_TIMEOUT_MS"); - umbp::proxy::ProxyShmRegion::CleanupStale(umbp::proxy::kDefaultShmName); - return true; -} - -static bool test_proxy_cas_rank_allocation() { - // Fork 3 "followers" that all try to auto-allocate rank slots. - // Each should get a unique rank. - - // First, spawn a leader to create the proxy - pid_t leader = fork(); - if (leader == 0) { - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - setenv("LOCAL_RANK", "0", 1); - setenv("UMBP_SSD_BACKEND", "spdk", 1); - - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - StandaloneClient client(cfg); - - // Wait for followers to finish (max 30s) - for (int i = 0; i < 300; ++i) { - if (access("/tmp/umbp_e2e_cas_done", F_OK) == 0) break; - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - } - unlink("/tmp/umbp_e2e_cas_done"); - _exit(0); - } - - // Wait for proxy to be ready - std::this_thread::sleep_for(std::chrono::seconds(3)); - - // Fork 3 followers - int pipe_fds[3][2]; - pid_t followers[3]; - for (int i = 0; i < 3; ++i) { - pipe(pipe_fds[i]); - followers[i] = fork(); - if (followers[i] == 0) { - close(pipe_fds[i][0]); - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - char lr[8]; - snprintf(lr, sizeof(lr), "%d", i + 1); - setenv("LOCAL_RANK", lr, 1); - setenv("UMBP_SSD_BACKEND", "spdk", 1); - - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - - // Wait for proxy, then auto-allocate rank via CAS - std::string shm_name = - cfg.ssd.spdk_proxy_shm_name.empty() ? "/umbp_spdk_proxy" : cfg.ssd.spdk_proxy_shm_name; - SpdkProxyTier::WaitForProxy(shm_name, 15000); - SpdkProxyTier tier(cfg.ssd); - uint8_t result = tier.IsValid() ? 1 : 0; - write(pipe_fds[i][1], &result, 1); - close(pipe_fds[i][1]); - _exit(0); - } - close(pipe_fds[i][1]); - } - - int ok_count = 0; - for (int i = 0; i < 3; ++i) { - uint8_t r = 0; - read(pipe_fds[i][0], &r, 1); - close(pipe_fds[i][0]); - if (r == 1) ok_count++; - int status; - waitpid(followers[i], &status, 0); - } - - // Signal leader to exit - int fd = open("/tmp/umbp_e2e_cas_done", O_CREAT | O_WRONLY, 0644); - if (fd >= 0) close(fd); - waitpid(leader, nullptr, 0); - unlink("/tmp/umbp_e2e_cas_done"); - - CHECK(ok_count == 3); - return true; -} - -static bool test_proxy_batch_write_read() { - pid_t leader = fork(); - if (leader == 0) { - unsetenv("UMBP_ROLE"); - unsetenv("UMBP_SPDK_PROXY_RANK"); - setenv("LOCAL_RANK", "0", 1); - setenv("UMBP_SSD_BACKEND", "spdk", 1); - - int rc = 0; - { - auto cfg = UMBPConfig::FromEnvironment(); - cfg.dram.capacity_bytes = 64ULL * 1024 * 1024; - StandaloneClient client(cfg); - - auto* ssd = client.Storage().GetTier(StorageTier::LOCAL_SSD); - if (!ssd) rc = 1; - - const int N = 50; - const size_t sz = 32768; - std::vector keys(N); - std::vector cptrs(N); - std::vector sizes(N, sz); - std::vector> bufs(N); - for (int i = 0; i < N; ++i) { - keys[i] = "proxy_batch_" + std::to_string(i); - bufs[i].resize(sz, static_cast(i & 0xFF)); - cptrs[i] = bufs[i].data(); - } - - if (rc == 0) { - auto wr = ssd->BatchWrite(keys, cptrs, sizes); - int wok = 0; - for (auto b : wr) wok += b; - if (wok != N) rc = 2; - } - - std::vector> rbufs(N, std::vector(sz, 0)); - std::vector dptrs(N); - for (int i = 0; i < N; ++i) dptrs[i] = reinterpret_cast(rbufs[i].data()); - - if (rc == 0) { - auto rr = ssd->BatchReadIntoPtr(keys, dptrs, sizes); - int rok = 0; - for (auto b : rr) rok += b; - if (rok != N) rc = 3; - } - - if (rc == 0) { - for (int i = 0; i < N; ++i) { - if (rbufs[i] != bufs[i]) { - fprintf(stderr, - " proxy batch mismatch idx=%d got0=%d exp0=%d got_last=%d exp_last=%d\n", i, - static_cast(rbufs[i][0]), static_cast(bufs[i][0]), - static_cast(rbufs[i][sz - 1]), - static_cast(bufs[i][sz - 1])); - rc = 4; - break; - } - } - } - } - _exit(rc); - } - - int status; - waitpid(leader, &status, 0); - if (!(WIFEXITED(status) && WEXITSTATUS(status) == 0)) { - fprintf(stderr, " proxy batch leader exit=%d\n", - WIFEXITED(status) ? WEXITSTATUS(status) : -1); - } - CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 0); - return true; -} -#endif // __linux__ - -// ========================================================================= -// main -// ========================================================================= - -int main() { - bool have_spdk = (std::getenv("UMBP_SPDK_NVME_PCI") != nullptr); - - printf("=== UMBP End-to-End Integration Tests ===\n"); - printf(" SPDK available: %s\n\n", have_spdk ? "YES" : "NO (set UMBP_SPDK_NVME_PCI to enable)"); - - // --- Role deduction tests --- - printf("[ROLE] Role deduction tests\n"); - RUN_TEST(test_role_default_standalone); - RUN_TEST(test_role_explicit_leader); - RUN_TEST(test_role_explicit_follower); - RUN_TEST(test_role_backward_compat_follower_mode); - RUN_TEST(test_role_backward_compat_force_cow); - RUN_TEST(test_auto_rank_sentinel); -#ifdef __linux__ - RUN_TEST(test_role_from_local_rank_env); -#endif - printf("\n"); - - // --- POSIX tests --- - printf("[POSIX] POSIX backend tests\n"); - RUN_TEST(test_posix_standalone_write_read); - RUN_TEST(test_posix_batch_write_read); - RUN_TEST(test_posix_dedup); - RUN_TEST(test_posix_evict); - RUN_TEST(test_posix_ssd_direct_write_read); - RUN_TEST(test_posix_dram_evict_to_ssd); - RUN_TEST(test_posix_capacity); - RUN_TEST(test_posix_empty_read); - RUN_TEST(test_posix_large_value); - printf("\n"); - - // --- SPDK tests --- -#ifdef __linux__ - if (have_spdk) { - printf("[CLEANUP] Killing residual SPDK processes...\n"); - CleanupResidualSpdk(); - - printf("[SPDK] SPDK standalone tests\n"); - RUN_TEST(test_spdk_standalone_write_read); - printf("\n"); - - // SPDK cleanup is not instant — wait for lock files to be released - printf("[CLEANUP] Waiting for SPDK lock release...\n"); - CleanupResidualSpdk(); - - printf("[PROXY] SPDK proxy auto-fork tests\n"); - RUN_TEST(test_proxy_batch_write_read); - CleanupResidualSpdk(); - RUN_TEST(test_proxy_auto_fork_leader_follower); - CleanupResidualSpdk(); - RUN_TEST(test_proxy_daemon_cleanup_after_leader_exit); - CleanupResidualSpdk(); - RUN_TEST(test_proxy_cas_rank_allocation); - CleanupResidualSpdk(); - printf("\n"); - } else { - printf("[SPDK] Skipped (UMBP_SPDK_NVME_PCI not set)\n"); - SKIP_TEST("test_spdk_standalone_write_read", "no SPDK"); - printf("\n"); - printf("[PROXY] Skipped (UMBP_SPDK_NVME_PCI not set)\n"); - SKIP_TEST("test_proxy_batch_write_read", "no SPDK"); - SKIP_TEST("test_proxy_auto_fork_leader_follower", "no SPDK"); - SKIP_TEST("test_proxy_daemon_cleanup_after_leader_exit", "no SPDK"); - SKIP_TEST("test_proxy_cas_rank_allocation", "no SPDK"); - printf("\n"); - } -#else - printf("[SPDK/PROXY] Skipped (Linux only)\n\n"); - g_skipped += 5; -#endif - - // --- Summary --- - printf("========================================\n"); - printf(" PASSED: %d FAILED: %d SKIPPED: %d\n", g_passed, g_failed, g_skipped); - printf("========================================\n"); - - if (g_failed > 0) { - printf("*** SOME TESTS FAILED ***\n"); - return 1; - } - - printf("=== ALL TESTS PASSED ===\n"); - return 0; -} diff --git a/tests/python/ops/bench_dispatch_combine.py b/tests/python/ops/bench_dispatch_combine.py index 3870c1e77..8dd08e51c 100644 --- a/tests/python/ops/bench_dispatch_combine.py +++ b/tests/python/ops/bench_dispatch_combine.py @@ -1656,7 +1656,7 @@ def bench_dispatch_combine( "'normal' is the legacy default (torch.randn). " "'uniform' samples on [-1, 1]. " "'lognormal' samples exp(N(0,1))-1 (heavy right tail). " - "'two_bucket' is 99% N(0,1) + 1% N(0, 16) for long-tail activations." + "'two_bucket' is 99%% N(0,1) + 1%% N(0, 16) for long-tail activations." ), ) parser.add_argument( diff --git a/tests/python/umbp/test_umbp_client_ptr.py b/tests/python/umbp/test_umbp_client_ptr.py index f429d3fff..95333d999 100644 --- a/tests/python/umbp/test_umbp_client_ptr.py +++ b/tests/python/umbp/test_umbp_client_ptr.py @@ -99,6 +99,42 @@ def test_put_from_ptr_get_into_ptr_round_trip(): assert client.clear() +def test_batch_put_get_ranges_round_trip_and_shape_failure(): + client = _make_client() + a = (ctypes.c_ubyte * 4)(*b"AAAA") + b = (ctypes.c_ubyte * 4)(*b"BBBB") + c = (ctypes.c_ubyte * 4)(*b"CCCC") + + assert client.batch_put_ranges_from_ptr( + ["range_key"], + [12], + [[ctypes.addressof(b), ctypes.addressof(c), ctypes.addressof(a)]], + [[4, 4, 4]], + [[4, 8, 0]], + ) == [True] + assert client.exists("range_key") + + prefix = (ctypes.c_ubyte * 4)() + suffix = (ctypes.c_ubyte * 4)() + assert client.batch_get_ranges_into_ptr( + ["range_key"], + [[ctypes.addressof(suffix), ctypes.addressof(prefix)]], + [[4, 4]], + [[8, 0]], + ) == [True] + assert bytes(prefix) == b"AAAA" + assert bytes(suffix) == b"CCCC" + + # A ragged inner vector is a whole-call shape error. + assert client.batch_get_ranges_into_ptr( + ["range_key", "missing"], + [[ctypes.addressof(prefix)], [ctypes.addressof(suffix)]], + [[4, 4], [4]], + [[0], [0]], + ) == [False, False] + assert client.clear() + + def test_batch_put_get_multiple_page_components(): client = _make_client() mem_pool = MockPageFirstKVCache(num_pages=4, page_size=1, element_size=256) diff --git a/tests/python/umbp/test_umbp_prim_bench.py b/tests/python/umbp/test_umbp_prim_bench.py new file mode 100644 index 000000000..e96c39c21 --- /dev/null +++ b/tests/python/umbp/test_umbp_prim_bench.py @@ -0,0 +1,164 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Key-space and payload logic of umbp_prim_bench. + +All of mori's imports in that module are function-local, so the parts tested +here -- which keys a call names, what each one must contain, and what answer +each is expected to give -- run without a store, a GPU or a server. +""" +import pytest + +bench = pytest.importorskip("umbp_prim_bench") + + +def _cfg(**over): + base = dict( + backend="umbp-server", + address="unix:///tmp/x.grpc.sock", + ops=("exists",), + rank_counts=(1,), + keys_per_rank=8, + batch=8, + passes=1, + ranges=4, + range_bytes=16, + loc="host", + register=False, + verify=True, + key_prefix="t", + dram_capacity=1 << 30, + exists_absent_frac=0.0, + ) + base.update(over) + return bench.Config(**base) + + +# --- payload --------------------------------------------------------------- + + +def test_payload_is_deterministic_key_derived_and_full_length(): + a = bench._payload("k1", 4096) + assert len(a) == 4096 + assert a == bench._payload("k1", 4096), "must reproduce across calls" + assert a != bench._payload("k2", 4096), "different keys, different bytes" + + +def test_payload_is_offset_sensitive(): + """A repeated-tag pattern would verify clean against a shifted read; this + is what makes a wrong range offset detectable at all.""" + p = bench._payload("k1", 4096) + assert p[0:64] != p[64:128] + assert p[0:64] != p[32:96], "a 32-byte shift must change the bytes" + + +def test_truncated_or_zero_filled_reads_are_detectable(): + """The keystream is a stream, so a short payload IS a prefix of a long + one -- which is harmless only because verification always compares the + full object length. These are the shapes that must not compare equal. + """ + full = bench._payload("k", 128) + assert full[:64] + bytes(64) != full, "a half-filled object must differ" + assert bytes(128) != full, "an untouched buffer must differ" + + +# --- key spaces ------------------------------------------------------------ + + +def test_absent_keys_are_disjoint_from_every_written_key(): + cfg = _cfg() + present = set() + for op in ("put", "get", "exists", "put_ranges", "get_ranges"): + for gen in range(0, 4): + present |= set(cfg.keys(op, 1, 0, gen)) + absent = set(cfg.keys("exists", 1, 0, 0, absent=True)) + assert absent and not (absent & present) + + +def test_absent_keys_have_the_same_shape_as_present_ones(): + """Same length and prefix, so mixing the two does not also mix two + different string-hashing costs into the measurement.""" + cfg = _cfg() + present = cfg.keys("exists", 1, 0, 0) + absent = cfg.keys("exists", 1, 0, 0, absent=True) + assert {len(k) for k in present} == {len(k) for k in absent} + + +# --- absent-key mixing ----------------------------------------------------- + + +def _exists_op(frac, keys_per_rank=8, batch=8): + cfg = _cfg(exists_absent_frac=frac, keys_per_rank=keys_per_rank, batch=batch) + return cfg, bench.ExistsOp(cfg, client=None, buf=None, ranks=1, rank=0) + + +@pytest.mark.parametrize("frac,expected", [(0.0, 0), (0.25, 2), (0.5, 4), (1.0, 8)]) +def test_absent_count_matches_the_requested_fraction(frac, expected): + _, op = _exists_op(frac) + assert len(op.absent_slots(8)) == expected + + +def test_absent_slots_are_spread_not_blocked(): + """Blocked at one end, whole calls would be single-population and the + per-call latency would depend on which chunk you looked at.""" + _, op = _exists_op(0.5) + slots = sorted(op.absent_slots(8)) + assert slots == [0, 2, 4, 6] + + +def test_expectation_tracks_the_absent_slots(): + cfg, op = _exists_op(0.5) + (batch,) = op.prepare(0) + absent = set(cfg.keys("exists", 1, 0, 0, absent=True)) + assert batch.expect is not None + for name, exp in zip(batch.key_names, batch.expect): + assert exp == (name not in absent), name + assert sum(batch.expect) == 4 + + +def test_all_present_probe_carries_no_absent_keys(): + cfg, op = _exists_op(0.0) + (batch,) = op.prepare(0) + assert all(batch.expect) + assert set(batch.key_names) <= set(cfg.keys("exists", 1, 0, 0)) + + +def test_all_absent_probe_expects_every_key_missing(): + cfg, op = _exists_op(1.0) + (batch,) = op.prepare(0) + assert not any(batch.expect) + assert set(batch.key_names) <= set(cfg.keys("exists", 1, 0, 0, absent=True)) + + +def test_absent_keys_do_not_shrink_the_call(): + """The probe must stay the same size at every fraction, or the sweep + compares calls of different widths.""" + for frac in (0.0, 0.5, 1.0): + _, op = _exists_op(frac, keys_per_rank=16, batch=4) + batches = op.prepare(0) + assert [b.keys for b in batches] == [4, 4, 4, 4] + assert all(len(b.key_names) == 4 for b in batches) + + +def test_fraction_is_rejected_outside_zero_to_one(): + for bad in ("-0.1", "1.5"): + with pytest.raises(SystemExit): + bench.parse_args(["umbp-local", "--exists-absent-frac", bad]) diff --git a/tests/python/umbp/test_umbp_prim_gate.py b/tests/python/umbp/test_umbp_prim_gate.py new file mode 100644 index 000000000..728c7551b --- /dev/null +++ b/tests/python/umbp/test_umbp_prim_gate.py @@ -0,0 +1,384 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Calibration for the primitive gate in umbp_prim_ci.py. + +The floor is only worth having if it separates the two builds it was drawn +from, so it is checked against both of them here rather than asserted to be +a sensible-looking number. The throughputs below are the measured +`umbp_prim_bench.py` results for the client-path regression that motivated +the gate: an exclusive lock over the resolve path left `batch_exists` +throughput flat in rank count, while the build before it scaled ~3.9x. + +Pure arithmetic over recorded numbers -- no mori, no server, no GPU -- so it +runs in the ordinary pytest step and fails loudly if someone retunes the +floor past the regression it exists to catch. +""" +import pytest + +umbp_prim_ci = pytest.importorskip("umbp_prim_ci") + +# Aggregate batch_exists keys/s by rank count, both arms, from one idle node +# with the same geometry and image pair. The 1- and 8-rank figures are the +# measured ones; 2 and 4 are reconstructed from the recorded scaling factors +# (1.91 / 3.27 and 1.37 / 0.91). The gate reads only the ends, so the middle +# is here to keep the fixture the shape of a real sweep. +GOOD_EXISTS = {1: 2_643_915, 2: 5_049_878, 4: 8_645_601, 8: 10_329_315} +BAD_EXISTS = {1: 964_177, 2: 1_320_922, 4: 877_401, 8: 720_637} + + +def _payload(exists_by_rank, failures=0, mismatches=0, absent_frac=0.0): + """A result file shaped like the one umbp_prim_bench.py writes.""" + rows = [] + for ranks, keys_per_s in exists_by_rank.items(): + for op in ("put", "get", "exists", "put_ranges", "get_ranges"): + rows.append( + { + "op": op, + "ranks": ranks, + "calls": 40, + "p50_ms": 1.0, + "p99_ms": 2.0, + "keys_per_s": keys_per_s if op == "exists" else 100_000.0, + "mib_per_s": 0.0 if op == "exists" else 1_000.0, + "failures": failures if op == "exists" else 0, + "mismatches": mismatches if op == "get" else 0, + } + ) + return { + "backend": "umbp-server", + "address": "unix:///tmp/x.grpc.sock", + "geometry": {"exists_absent_frac": absent_frac}, + "rows": rows, + } + + +def test_default_floor_separates_the_two_arms(): + """The whole point: the good build passes and the bad one does not.""" + floor = umbp_prim_ci.DEFAULT_EXISTS_SCALING_FLOOR + # These are idle-node numbers; the floor also has to clear the much + # lower scores a contended runner produces -- see the test below. + good = umbp_prim_ci.check_exists_scaling("good", _payload(GOOD_EXISTS), floor)[0] + bad = umbp_prim_ci.check_exists_scaling("bad", _payload(BAD_EXISTS), floor)[0] + assert good.ok, f"the known-good arm must pass: {good}" + assert not bad.ok, f"the known-regressed arm must fail: {bad}" + assert good.value == "3.91x" + assert bad.value == "0.75x" + + +def test_floor_has_margin_under_the_good_arm(): + """A floor set flush against the good arm turns a busy runner into a red + build, so keep the distance from it visible and enforced.""" + good_scale = GOOD_EXISTS[8] / GOOD_EXISTS[1] + bad_scale = BAD_EXISTS[8] / BAD_EXISTS[1] + floor = umbp_prim_ci.DEFAULT_EXISTS_SCALING_FLOOR + assert bad_scale < floor < good_scale + assert good_scale / floor >= 1.5, "less than 1.5x of headroom under the good arm" + + +def test_correctness_gate_catches_a_failed_operation(): + findings = umbp_prim_ci.check_correctness("arm", _payload(GOOD_EXISTS, failures=3)) + assert not findings[0].ok + assert "3 failed" in findings[0].detail + + +def test_correctness_gate_catches_a_readback_mismatch(): + findings = umbp_prim_ci.check_correctness( + "arm", _payload(GOOD_EXISTS, mismatches=1) + ) + assert not findings[0].ok + assert "1 mismatched" in findings[0].detail + + +def _arm(exists_by_rank, scaling_by_frac=None, **kwargs): + """An arm as run_arm returns it: one payload per pass. + + The sweep contributes one scaling pass per absent fraction; only the + 0.0 one is gated. + """ + arm = {"correctness": _payload(exists_by_rank, **kwargs)} + for frac in (0.0, 0.5, 1.0): + rows = (scaling_by_frac or {}).get(frac, exists_by_rank) + arm[f"scaling@{frac:g}"] = _payload(rows, absent_frac=frac) + return arm + + +def test_clean_run_passes_every_check(): + findings = umbp_prim_ci.gate({"nomaster": _arm(GOOD_EXISTS)}, 2.0) + assert all(f.ok for f in findings) + assert umbp_prim_ci.report_findings(findings) + + +def test_only_the_zero_absent_fraction_is_gated(): + """The other fractions have no measured baseline, so they are reported + and must not be able to fail the build.""" + findings = umbp_prim_ci.gate({"arm": _arm(GOOD_EXISTS)}, 2.0) + scaling = [f for f in findings if "exists scaling" in f.check] + assert len(scaling) == 3 + gated = {f.check.split("absent=")[1].split()[0] for f in scaling if f.gated} + assert gated == {"0"}, f"expected only absent=0 gated, got {gated}" + + +def test_a_regressed_absent_fraction_reports_but_does_not_fail(): + findings = umbp_prim_ci.gate( + {"arm": _arm(GOOD_EXISTS, scaling_by_frac={0.5: BAD_EXISTS, 1.0: BAD_EXISTS})}, + 2.0, + ) + assert umbp_prim_ci.report_findings(findings), "ungated fractions must not gate" + ungated = [f for f in findings if not f.gated] + assert ungated and all(f.value == "0.75x" for f in ungated) + + +def test_a_regressed_zero_fraction_does_fail(): + findings = umbp_prim_ci.gate( + {"arm": _arm(GOOD_EXISTS, scaling_by_frac={0.0: BAD_EXISTS})}, 2.0 + ) + assert not umbp_prim_ci.report_findings(findings) + + +def test_correctness_is_checked_on_the_absent_fraction_passes_too(): + """A False for a present key, or True for an absent one, lands in + `failures` on those passes -- a false positive being the worse bug.""" + arm = _arm(GOOD_EXISTS) + arm["scaling@1"]["rows"][0]["failures"] = 7 + findings = umbp_prim_ci.gate({"arm": arm}, 2.0) + bad = [f for f in findings if f.check.startswith("correctness") and not f.ok] + assert len(bad) == 1 and "scaling@1" in bad[0].check + assert not umbp_prim_ci.report_findings(findings) + + +def test_one_bad_arm_fails_the_whole_gate(): + findings = umbp_prim_ci.gate( + {"nomaster": _arm(GOOD_EXISTS), "master": _arm(BAD_EXISTS)}, 2.0 + ) + assert not umbp_prim_ci.report_findings(findings) + + +def test_correctness_reads_the_correctness_pass_not_the_scaling_one(): + """The scaling pass runs `exists` alone, so a correctness failure that + only the five-primitive pass can see must not be looked for in it.""" + findings = umbp_prim_ci.gate({"arm": _arm(GOOD_EXISTS, mismatches=1)}, 1.0) + correctness = [f for f in findings if f.check.startswith("correctness")] + assert not correctness[0].ok + + +#: Worst 1->8 `exists` scaling seen from an unmodified build on a contended +#: node, over three runs through the scaling pass's geometry (3.50 / 5.39 / +#: 4.33). The floor has to clear this, or a busy runner turns the gate red +#: on a healthy build. +WORST_HEALTHY_SCALING = 3.50 + + +def test_shipped_floor_sits_between_a_busy_healthy_run_and_the_regression(): + floor = umbp_prim_ci.DEFAULT_EXISTS_SCALING_FLOOR + assert BAD_EXISTS[8] / BAD_EXISTS[1] < floor < WORST_HEALTHY_SCALING + assert WORST_HEALTHY_SCALING / floor >= 1.4, "under 40% headroom on a busy runner" + assert floor / (BAD_EXISTS[8] / BAD_EXISTS[1]) >= 2.0, "too close to the regression" + + +def test_single_rank_count_is_not_a_pass(): + """A ratio needs two points; silently reporting PASS on one would make + the gate disappear the moment someone trims the sweep.""" + finding = umbp_prim_ci.check_exists_scaling("arm", _payload({8: 10_000.0}), 2.0)[0] + assert not finding.ok + assert "two rank counts" in finding.detail + + +def test_zero_baseline_is_a_failure_not_a_crash(): + finding = umbp_prim_ci.check_exists_scaling( + "arm", _payload({1: 0.0, 8: 10_000.0}), 2.0 + )[0] + assert not finding.ok + + +def test_pool_is_sized_from_the_whole_sweep(): + """The pool must cover every rank count's keys at once, not just the + largest -- each writes into its own namespace and none are freed.""" + from types import SimpleNamespace + + args = SimpleNamespace( + ranks=[1, 2, 4, 8], + ops=["put", "get", "exists", "put_ranges", "get_ranges"], + keys_per_rank=512, + batch=128, + passes=3, + ranges=4, + range_bytes=16384, + scaling_keys_per_rank=8192, + scaling_batch=256, + scaling_passes=100, + scaling_range_bytes=4096, + page_bytes=65536, + exists_absent_fracs=[0.0], + scaling_key_prefix="s00001", + ) + correctness, scaling = umbp_prim_ci.passes_for(args) + page = args.page_bytes + # 15 rank-slots x 512 keys x (3 seeded + 2 writers x 3 passes), one page each + assert correctness.resident_bytes(args.ranks, page) == 15 * 512 * 9 * page + # The scaling pass seeds exists once and never writes again, so its + # residency does not grow with its (large) pass count. + assert scaling.resident_bytes(args.ranks, page) == 15 * 8192 * 1 * page + + resident = umbp_prim_ci.pool_bytes(args, headroom=1.0) + assert resident == sum( + p.resident_bytes(args.ranks, page) for p in (correctness, scaling) + ) + assert umbp_prim_ci.pool_bytes(args, headroom=3.0) == 3 * resident + + largest_only = SimpleNamespace(**{**vars(args), "ranks": [8]}) + assert umbp_prim_ci.pool_bytes(largest_only, 1.0) < resident + + +def test_pool_is_charged_in_whole_pages_not_bytes(): + """The regression this guards cost a false A/B result. + + The DRAM medium is paged and hands out whole pages, so a 4 KiB value in a + 2 MiB page occupies 2 MiB. Sizing the pool by `keys * object_bytes` + understates demand by the padding ratio; the pool then evicts continuously + and, because the just-seeded keys are the least recently used, the run + dies in the visibility wait rather than reporting anything about space. + """ + from types import SimpleNamespace + + tiny = umbp_prim_ci.Pass( + "scaling", + ["exists"], + keys_per_rank=1000, + batch=100, + passes=1, + ranges=1, + range_bytes=4096, + ) + # One 4 KiB object still costs a whole 2 MiB page: 512x what bytes suggest. + assert tiny.resident_bytes([1], 2 << 20) == 1000 * (2 << 20) + assert tiny.resident_bytes([1], 2 << 20) == 512 * tiny.resident_bytes([1], 4096) + + # An object larger than a page spans pages rather than being rounded down. + wide = umbp_prim_ci.Pass( + "correctness", + ["exists"], + keys_per_rank=10, + batch=10, + passes=1, + ranges=1, + range_bytes=5 * 65536, + ) + assert wide.resident_bytes([1], 65536) == 10 * 5 * 65536 + + # And the pool the harness asks for tracks the page size it pins. + args = SimpleNamespace( + ranks=[1, 2], + ops=["exists"], + keys_per_rank=64, + batch=64, + passes=1, + ranges=1, + range_bytes=4096, + scaling_keys_per_rank=64, + scaling_batch=64, + scaling_passes=1, + scaling_range_bytes=4096, + page_bytes=65536, + exists_absent_fracs=[0.0], + scaling_key_prefix="s00001", + ) + small_page = umbp_prim_ci.pool_bytes(args, 1.0) + args.page_bytes = 2 << 20 + assert umbp_prim_ci.pool_bytes(args, 1.0) == 32 * small_page + + +# --------------------------------------------------------------------------- +# Media selection +# --------------------------------------------------------------------------- + + +def _media(spec, dirs=("/d0", "/d1"), devs=(0,)): + return umbp_prim_ci.parse_media(spec, list(dirs), list(devs)) + + +def test_single_medium_gets_the_whole_pool(): + for spec in ("dram", "ssd", "hbm"): + (b,) = _media(spec) + assert b.kind == spec and b.weight == 100.0 + + +def test_unweighted_media_split_the_remainder_evenly(): + assert [b.weight for b in _media("dram,ssd")] == [50.0, 50.0] + assert [round(b.weight) for b in _media("dram,ssd,hbm")] == [33, 33, 33] + + +def test_explicit_weights_are_normalised_to_a_hundred(): + got = {b.kind: b.weight for b in _media("dram:70,ssd:30")} + assert got == {"dram": 70.0, "ssd": 30.0} + # Relative weights, not percentages: 3:1 is the same as 75:25. + got = {b.kind: b.weight for b in _media("dram:3,ssd:1")} + assert got == {"dram": 75.0, "ssd": 25.0} + + +def test_a_partial_weight_leaves_the_rest_to_the_unweighted(): + got = {b.kind: b.weight for b in _media("dram:80,ssd")} + assert got == {"dram": 80.0, "ssd": 20.0} + + +@pytest.mark.parametrize( + "spec", ["", "tape", "dram:60,ssd:60,hbm", "dram,dram", "dram:0"] +) +def test_bad_media_specs_are_rejected(spec): + with pytest.raises(ValueError): + _media(spec) + + +def test_ssd_backends_consume_the_supplied_directories(): + with pytest.raises(ValueError): + _media("ssd", dirs=[]) + (b,) = _media("ssd", dirs=["/mnt/one"]) + assert b.path == "/mnt/one" + + +def test_policy_splits_capacity_by_weight_and_names_each_medium(): + backends = _media("dram:70,ssd:30", dirs=["/mnt/nvme"]) + pol = umbp_prim_ci.backend_policy(backends, 100 << 30) + assert pol["schema_version"] == 1 + assert pol["entry_tier"] in pol["tiers"][0]["name"] + # The unit suffix is mandatory; a bare byte count is rejected by the loader. + assert pol["backends"]["dram"]["capacity"].endswith("B") + assert int(pol["backends"]["dram"]["capacity"][:-1]) == int((100 << 30) * 0.7) + assert int(pol["backends"]["ssd"]["capacity"][:-1]) == int((100 << 30) * 0.3) + assert pol["backends"]["ssd"]["path"] == "/mnt/nvme" + # One logical tier: this measures placement across media, not promotion + # between tiers, which would put an eviction policy inside a latency number. + assert len(pol["tiers"]) == 1 + assert pol["tiers"][0]["backends"] == {"dram": 70.0, "ssd": 30.0} + + +def test_policy_gives_hbm_its_device_and_no_path(): + backends = _media("hbm", devs=[3]) + pol = umbp_prim_ci.backend_policy(backends, 8 << 30) + assert pol["backends"]["hbm3"]["devices"] == [3] + assert "path" not in pol["backends"]["hbm3"] + + +def test_tiny_shares_still_get_a_usable_backend(): + """A 1% share of a small pool must not round down to nothing.""" + backends = _media("dram:99,ssd:1", dirs=["/mnt/nvme"]) + pol = umbp_prim_ci.backend_policy(backends, 4 << 20) + assert int(pol["backends"]["ssd"]["capacity"][:-1]) >= (1 << 20) diff --git a/tests/python/umbp/umbp_bench.py b/tests/python/umbp/umbp_bench.py deleted file mode 100644 index 74d4616f0..000000000 --- a/tests/python/umbp/umbp_bench.py +++ /dev/null @@ -1,829 +0,0 @@ -# Copyright © Advanced Micro Devices, Inc. All rights reserved. -# -# MIT License -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -import argparse -import atexit -import ctypes -import os -import signal -import socket -import subprocess -import sys -import time -from abc import ABC, abstractmethod - -_DEFAULT_NR_OBJECTS = 2048 -_DEFAULT_VALUE_SIZE = 2097152 -_DEFAULT_SYNC_PORT = 18515 - - -# --- optional in-process server startup (--start-server) --- - - -def _port_in_use(port) -> bool: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.settimeout(0.5) - return s.connect_ex(("127.0.0.1", int(port))) == 0 - - -def _wait_for_port(port, timeout: float = 15.0) -> bool: - deadline = time.time() + timeout - while time.time() < deadline: - if _port_in_use(port): - return True - time.sleep(0.2) - return False - - -_PR_SET_PDEATHSIG = 1 - - -def _die_with_parent(): - # Linux only: ask the kernel to SIGKILL this child when the bench process - # dies -- including when the bench is itself killed with SIGKILL, which - # can't be caught so atexit never runs. Also start a new session so the - # child becomes a process-group leader; this lets us SIGKILL the whole - # group (including grandchildren that PDEATHSIG doesn't reach) on cleanup. - os.setsid() - ctypes.CDLL("libc.so.6", use_errno=True).prctl(_PR_SET_PDEATHSIG, signal.SIGKILL) - - -def _kill_group(proc): - # Kill the spawned process and any descendants it left behind (e.g. - # mooncake_master's metrics/admin server). PDEATHSIG only covers direct - # children and doesn't propagate to grandchildren, so reap the whole group. - if proc.poll() is not None: - return - try: - os.killpg(proc.pid, signal.SIGKILL) - except ProcessLookupError: - pass - except OSError: - proc.kill() - - -def _spawn(argv, log_path, env=None) -> subprocess.Popen: - log = open(log_path, "w") - try: - proc = subprocess.Popen( - argv, - stdout=log, - stderr=subprocess.STDOUT, - env=env, - preexec_fn=_die_with_parent, - ) - finally: - log.close() - atexit.register(lambda: _kill_group(proc)) - return proc - - -def _build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(prog="umbp_bench.py") - cmd_sub = parser.add_subparsers(dest="command", required=True, metavar="COMMAND") - - def _add_backend_sub(cmd_parser, *, batch_perf: bool = False): - be_sub = cmd_parser.add_subparsers( - dest="backend", required=True, metavar="BACKEND" - ) - for be_name in ("umbp", "mooncake"): - be_p = be_sub.add_parser(be_name) - if be_name == "umbp": - be_p.add_argument("master_address", help="e.g. ip:port") - else: - be_p.add_argument( - "metadata_server", help="e.g. http://ip:8080/metadata" - ) - be_p.add_argument("master_server", help="e.g. ip:50051") - be_p.add_argument("--nr-objects", type=int, default=_DEFAULT_NR_OBJECTS) - be_p.add_argument("--value-size", type=int, default=_DEFAULT_VALUE_SIZE) - be_p.add_argument( - "--role", - choices=["both", "writer", "reader"], - default="both", - help="Process role for split-node runs. 'both' (default) writes " - "then reads in a single process (original behavior). 'writer' " - "writes the dataset and stays alive holding it in DRAM until " - "the reader is done. 'reader' connects to the writer's " - "already-running master and reads. Run writer and reader on " - "different nodes, passing the reader the writer's master " - "address, and the same --key-prefix and --sync-port to both.", - ) - be_p.add_argument( - "--key-prefix", - default=None, - help="Shared key namespace (first 4 chars used). REQUIRED to match " - "between writer and reader in split runs so they compute the " - "same keys; defaults to a per-pid value when unset (fine for " - "role=both, broken for split).", - ) - be_p.add_argument( - "--sync-port", - type=int, - default=_DEFAULT_SYNC_PORT, - help="TCP port on the writer node used as a write-done/read-done " - "rendezvous between a split writer and reader (default: " - f"{_DEFAULT_SYNC_PORT}). Ignored for role=both.", - ) - be_p.add_argument( - "--numa-node", - type=int, - default=-1, - help="Pin this process to a NUMA node (default: -1 = no binding): " - "binds CPU affinity to the node's cpulist, allocates the bench " - "DMA buffers on it, and (UMBP) places the DRAM tier there. The " - "IO engine then auto-selects NUMA-local NICs. In a same-node " - "split, give the writer and reader different --numa-node values " - "to exercise a cross-NUMA transfer. Overridable per-allocation " - "via UMBP_DRAM_NUMA_NODE.", - ) - be_p.add_argument( - "--start-server", - action=argparse.BooleanOptionalAction, - default=True, - help="Launch the backend metadata/master server(s) for this run " - "and tear them down on exit (default: on; use --no-start-server " - "to reuse an already-running server).", - ) - if batch_perf: - be_p.add_argument("--passes", type=int, default=10) - be_p.add_argument("--batch-sizes", type=int, nargs="+", metavar="N") - be_p.add_argument( - "--dest-layout", - choices=["contiguous", "shuffled"], - default="shuffled", - help="Destination buffer layout for reads (default: shuffled). " - "'shuffled' maps keys to a random permutation of slots in the " - "same buffer so consecutive keys are non-adjacent in memory " - "(non-consecutive transfer). 'contiguous' packs objects " - "back-to-back (coalesceable), but a full-dataset batch then " - "coalesces into a single multi-GB SGE that the ionic provider " - "rejects (ibv_post_send EINVAL), so large batches all-miss.", - ) - be_p.add_argument( - "--shuffle-seed", - type=int, - default=0, - help="Seed for --dest-layout shuffled (default: 0).", - ) - - _add_backend_sub(cmd_sub.add_parser("correctness")) - _add_backend_sub(cmd_sub.add_parser("batch_perf"), batch_perf=True) - return parser - - -args = _build_parser().parse_args() -command = args.command -backend_name = args.backend -nr_objects = args.nr_objects -value_size = args.value_size -role = args.role -numa_node = args.numa_node - -if role != "both" and not args.key_prefix: - print( - "ERROR: split runs (--role writer/reader) require a shared --key-prefix on " - "both writer and reader so they compute the same keys", - file=sys.stderr, - ) - sys.exit(2) - -if backend_name == "umbp": - from mori.umbp import UMBPClient, UMBPConfig, UMBPDistributedConfig - - master_address = args.master_address - _resolve_host = master_address.split(":")[0] -else: - from mooncake.store import MooncakeDistributedStore, ReplicateConfig - - metadata_server = args.metadata_server - master_server = args.master_server - _resolve_host = master_server.split(":")[0] - -key_prefix = args.key_prefix if args.key_prefix else f"{os.getpid() % 10000:04d}" -key_size = 20 -pattern = b"\xab" - -with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as _s: - _s.connect((_resolve_host, 1)) - node_address = _s.getsockname()[0] - - -def make_key(object_id: int) -> str: - suffix = f"{object_id:016d}" - prefix_part = key_prefix[: key_size - len(suffix)].ljust( - key_size - len(suffix), "_" - ) - return f"{prefix_part}{suffix}" - - -def expected_payload(size: int) -> bytes: - reps = (size // len(pattern)) + 1 - return (pattern * reps)[:size] - - -# --- backend abstraction --- - - -class Backend(ABC): - """Uniform KV-store interface; subclass per backend.""" - - @classmethod - @abstractmethod - def start_server(cls, args): - """Launch this backend's metadata/master server(s) for the run; tear down on exit.""" - - @abstractmethod - def put(self, key: str, payload: bytes, ptr: int, size: int) -> bool: - """Write value. ptr is a pre-registered DMA buffer (may be ignored by some backends).""" - - @abstractmethod - def flush(self): ... - - @abstractmethod - def register_memory(self, ptr: int, size: int): - """Register a memory region for RDMA (no-op for backends that don't need it).""" - - @abstractmethod - def get_into_ptr(self, key: str, ptr: int, size: int) -> bool: - """Read value into ptr. Returns True on hit.""" - - @abstractmethod - def batch_get_into_ptr(self, keys: list, ptrs: list, sizes: list) -> list: - """Batch read into ptrs. Returns bytes_read per key (0 = miss).""" - - @abstractmethod - def make_reader(self) -> "Backend": - """Return a backend instance for reading (may be self).""" - - -class UMBPBackend(Backend): - @classmethod - def start_server(cls, args): - # In a split run the reader connects using the exact master port it was - # given on the CLI, so the writer must bind that port instead of an - # auto-picked one. role=both keeps auto-picking to avoid collisions on - # repeated single-node runs. - parts = args.master_address.split(":") - desired_port = int(parts[1]) if len(parts) > 1 and parts[1] else 0 - port = ( - desired_port if (args.role == "writer" and desired_port) else _free_port() - ) - mori_root = os.path.abspath( - os.path.join(os.path.dirname(__file__), "..", "..", "..") - ) - fqdn = subprocess.run( - ["hostname", "-f"], capture_output=True, text=True - ).stdout.strip() - binary = os.path.join(mori_root, f"build_{fqdn}", "src", "umbp", "umbp_master") - if not os.path.exists(binary): - print( - f"umbp_master binary not found: {binary} (run ./build.sh first)", - file=sys.stderr, - ) - sys.exit(1) - env = dict(os.environ) - env.setdefault("MORI_UMBP_LOG_LEVEL", "DEBUG") - proc = _spawn([binary, f"0.0.0.0:{port}", "0"], "/tmp/umbp_master.log", env) - if not _wait_for_port(port): - print( - "umbp_master failed to start; see /tmp/umbp_master.log", file=sys.stderr - ) - sys.exit(1) - print(f"started umbp_master pid={proc.pid} on 0.0.0.0:{port}", flush=True) - # Return the address the client should connect to (port may have been - # auto-picked, so it can differ from what was passed on the CLI). - return f"{args.master_address.split(':')[0]}:{port}" - - def __init__( - self, - node_id: str, - port: int, - reader_node_id: str = None, - reader_port: int = None, - ): - cfg = UMBPConfig() - cfg.dram.capacity_bytes = int( - os.environ.get("UMBP_DRAM_CAPACITY_BYTES", 8 * 1024 * 1024 * 1024) - ) - cfg.dram.use_hugepages = os.environ.get("UMBP_DRAM_USE_HUGEPAGES", "1") not in ( - "0", - "", - ) - cfg.dram.hugepage_size = int( - os.environ.get("UMBP_DRAM_HUGEPAGE_SIZE", 2 * 1024 * 1024) - ) - cfg.dram.numa_node = int(os.environ.get("UMBP_DRAM_NUMA_NODE", numa_node)) - dist = UMBPDistributedConfig() - dist.master_config.master_address = master_address - dist.master_config.node_id = node_id - dist.master_config.node_address = node_address - dist.peer_service_port = port - dist.io_engine.host = node_address - cfg.distributed = dist - self._client = UMBPClient(cfg) - self._reader_node_id = reader_node_id - self._reader_port = reader_port - - def put(self, key: str, payload: bytes, ptr: int, size: int) -> bool: - ctypes.memmove(ptr, payload, size) - return self._client.put_from_ptr(key, ptr, size) - - def flush(self): - self._client.flush() - - def register_memory(self, ptr: int, size: int): - self._client.register_memory(ptr, size) - - def get_into_ptr(self, key: str, ptr: int, size: int) -> bool: - return self._client.get_into_ptr(key, ptr, size) - - def batch_get_into_ptr(self, keys: list, ptrs: list, sizes: list) -> list: - raw = self._client.batch_get_into_ptr(keys, ptrs, sizes) - return [sz if hit else 0 for hit, sz in zip(raw, sizes)] - - def make_reader(self) -> "UMBPBackend": - return UMBPBackend(self._reader_node_id, self._reader_port) - - -class MooncakeBackend(Backend): - @classmethod - def start_server(cls, args): - scheme, _, rest = args.metadata_server.partition("://") - meta_hostport, _, meta_path = rest.partition("/") - meta_host = meta_hostport.partition(":")[0] - master_host = args.master_server.partition(":")[0] - # For a split run honor the CLI-provided meta/master ports (role=writer) - # so the reader on another node can connect using the same addresses. - cli_meta_port = int(meta_hostport.partition(":")[2] or 0) - cli_master_port = int(args.master_server.partition(":")[2] or 0) - meta_port = ( - cli_meta_port if (args.role == "writer" and cli_meta_port) else _free_port() - ) - master_port = ( - cli_master_port - if (args.role == "writer" and cli_master_port) - else _free_port() - ) - # mooncake_master also binds a metrics/admin HTTP server, hardcoded to - # 9003 by default. --enable_metric_reporting=false only disables - # reporting, not the bind, so a leaked master from a prior run holds - # 9003 and the next master collides. Give it a free port too. - metrics_port = _free_port() - meta = _spawn( - [ - "mooncake_http_metadata_server", - "--port", - str(meta_port), - "--host", - "0.0.0.0", - ], - "/tmp/mc_http_metadata.log", - ) - master = _spawn( - [ - "mooncake_master", - "--rpc_port", - str(master_port), - "--metrics_port", - str(metrics_port), - "--enable_metric_reporting=false", - ], - "/tmp/mc_master.log", - ) - if not _wait_for_port(meta_port) or not _wait_for_port(master_port): - print( - "mooncake servers failed to start; see /tmp/mc_*.log", file=sys.stderr - ) - sys.exit(1) - print( - f"started mooncake servers meta pid={meta.pid} port={meta_port} " - f"master pid={master.pid} port={master_port}", - flush=True, - ) - # Return the addresses the client should use (ports may have been auto-picked). - return ( - f"{scheme}://{meta_host}:{meta_port}/{meta_path}", - f"{master_host}:{master_port}", - ) - - def __init__(self): - # Enable hugepages (2MB) by default, matching the UMBP backend. These are - # read by the C++ client at store construction / setup() time. Presence of - # MC_STORE_USE_HUGEPAGE (any value) enables it; MC_STORE_HUGEPAGE_SIZE - # accepts "2MB" or "1GB". Both stay overridable via the environment. - os.environ.setdefault("MC_STORE_USE_HUGEPAGE", "1") - os.environ.setdefault("MC_STORE_HUGEPAGE_SIZE", "2MB") - segment_size = max(1 << 24, nr_objects * value_size * 2) - # local_buffer_size is the client-side RDMA staging pool. It is - # registered on EVERY NIC at setup, so on multi-NIC ionic an 8 GB pool - # blows up the device/IOMMU page-table footprint (8 GB x N NICs) and - # ibv_reg_mr returns ENOMEM at >=6 NICs. The zero-copy batch_get_into - # path never touches it; only sequential put()/single get() stage - # through it, one 2 MB object at a time, so a small pool is plenty. - # Override via MC_LOCAL_BUFFER_SIZE (bytes) if a workload needs more. - local_buffer_size = int( - os.environ.get("MC_LOCAL_BUFFER_SIZE", 256 * 1024 * 1024) - ) - self._store = MooncakeDistributedStore() - # rdma_devices is a comma-separated HCA list. Empty string => Mooncake - # auto-discovers all RDMA NICs (no longer pinned to a single device). - rdma_devices = os.environ.get("MC_RDMA_DEVICES", "") - # Pass the sizes by keyword: setup()'s positional order is - # (..., global_segment_size, local_buffer_size, ...) -- the reverse of - # how it reads -- so a positional swap silently mis-sizes the segment - # (segment too small => puts fail "insufficient space"). - ret = self._store.setup( - node_address, - metadata_server, - global_segment_size=segment_size, - local_buffer_size=local_buffer_size, - protocol="rdma", - rdma_devices=rdma_devices, - master_server_addr=master_server, - ) - if ret != 0: - print(f"mooncake setup failed: {ret}", file=sys.stderr) - sys.exit(1) - self._config = ReplicateConfig() - self._config.replica_num = 1 - - def put(self, key: str, payload: bytes, ptr: int, size: int) -> bool: - return self._store.put(key, payload, self._config) == 0 - - def flush(self): - pass - - def register_memory(self, ptr: int, size: int): - self._store.register_buffer(ptr, size) - - def get_into_ptr(self, key: str, ptr: int, size: int) -> bool: - data = self._store.get(key) - if data in (None, b""): - return False - ctypes.memmove(ptr, bytes(data), min(len(data), size)) - return True - - def batch_get_into_ptr(self, keys: list, ptrs: list, sizes: list) -> list: - return self._store.batch_get_into(keys, ptrs, sizes) - - def make_reader(self) -> "MooncakeBackend": - return self - - -# --- backend init --- - -_BACKENDS = {"umbp": UMBPBackend, "mooncake": MooncakeBackend} - - -def _free_port(): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as _s: - _s.bind(("", 0)) - return _s.getsockname()[1] - - -# --- split writer/reader rendezvous (TCP, hosted on the writer node) --- - - -def _sync_recv_until(conn, token: bytes, timeout: float = 600.0): - conn.settimeout(timeout) - buf = b"" - while token not in buf: - chunk = conn.recv(64) - if not chunk: - raise ConnectionError(f"sync peer closed before sending {token!r}") - buf += chunk - - -def _sync_listen(port: int): - srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - srv.bind(("0.0.0.0", port)) - srv.listen(1) - return srv - - -def _sync_connect(host: str, port: int, timeout: float = 300.0): - deadline = time.time() + timeout - last = None - while time.time() < deadline: - try: - return socket.create_connection((host, port), timeout=5) - except OSError as e: - last = e - time.sleep(0.5) - raise TimeoutError(f"could not reach writer sync {host}:{port}: {last}") - - -def _bind_numa_cpus(node: int): - # Pin CPU affinity to the given NUMA node's cpulist. Memory placement is - # handled separately (UMBP DRAM tier + bench buffers take an explicit - # numa_node), but binding CPUs first means any first-touch allocation and - # the IO engine's worker threads also land on this node. Call before - # constructing any client or allocating buffers so threads inherit it. - if node < 0: - return - try: - with open(f"/sys/devices/system/node/node{node}/cpulist") as f: - spec = f.read().strip() - except OSError as e: - print( - f"WARNING: NUMA node {node} has no cpulist ({e}); not binding CPUs", - file=sys.stderr, - ) - return - cpus = set() - for part in spec.split(","): - if not part: - continue - if "-" in part: - lo, hi = part.split("-") - cpus.update(range(int(lo), int(hi) + 1)) - else: - cpus.add(int(part)) - if not cpus: - return - os.sched_setaffinity(0, cpus) - print(f"bound to NUMA node {node}: {len(cpus)} cpus ({spec})", flush=True) - - -# Only the writer (or a combined both) owns the master/metadata server(s); the -# reader connects to the writer's already-running master, so it never starts one. -if args.start_server and role != "reader": - _resolved = _BACKENDS[backend_name].start_server(args) - # The server may have been moved to an auto-picked free port; point the - # client at the address(es) it actually bound. - if backend_name == "umbp": - master_address = _resolved - else: - metadata_server, master_server = _resolved - -# Pin to the requested NUMA node before any client/thread/buffer is created so -# affinity and first-touch placement are inherited. In a same-node split the -# writer and reader pass different --numa-node values (cross-NUMA transfer). -_bind_numa_cpus(numa_node) - -# Build the writer only for role in {both, writer}; build a standalone reader -# backend lazily (role=reader uses _build_reader; role=both uses -# writer.make_reader() so the two clients still share one process). -writer = None -if backend_name == "umbp": - - def _build_reader(): - rid = f"reader_{os.getpid()}_{int(time.time())}" - return UMBPBackend(rid, _free_port()) - - if role in ("both", "writer"): - ts = int(time.time()) - node_id = f"{command}_{os.getpid()}_{ts}" - reader_node_id = f"reader_{os.getpid()}_{ts}" - peer_service_port = _free_port() - reader_peer_port = _free_port() - writer = UMBPBackend( - node_id, peer_service_port, reader_node_id, reader_peer_port - ) - print( - f"backend=umbp role={role} writer_node_id={node_id} reader_node_id={reader_node_id} " - f"node_address={node_address} writer_port={peer_service_port} reader_port={reader_peer_port}", - flush=True, - ) -elif backend_name == "mooncake": - - def _build_reader(): - return MooncakeBackend() - - if role in ("both", "writer"): - print( - f"backend=mooncake role={role} node_address={node_address} metadata={metadata_server} master={master_server}", - flush=True, - ) - writer = MooncakeBackend() - -# --- host buffer allocation (hugepage-backed via mori's UMBP allocator) --- - - -class HostBuffer: - """Host DMA buffer allocated through mori's UMBPHostMemAllocator so the bench - buffers are hugepage-backed (matching the DRAM tier) instead of the 4 KiB-paged - ctypes arrays used previously. Keeps the handle alive for the object's lifetime - and exposes a raw integer `ptr`. Honors the same UMBP_DRAM_USE_HUGEPAGES / - UMBP_DRAM_HUGEPAGE_SIZE env knobs (default: on, 2MB).""" - - _allocator = None - - def __init__(self, size: int): - from mori.umbp import UMBPHostBufferBacking, UMBPHostMemAllocator - - if HostBuffer._allocator is None: - HostBuffer._allocator = UMBPHostMemAllocator() - use_hp = os.environ.get("UMBP_DRAM_USE_HUGEPAGES", "1") not in ("0", "") - hp_size = int(os.environ.get("UMBP_DRAM_HUGEPAGE_SIZE", 2 * 1024 * 1024)) - backing = ( - UMBPHostBufferBacking.AnonymousHugetlb - if use_hp - else UMBPHostBufferBacking.Anonymous - ) - numa = int(os.environ.get("UMBP_DRAM_NUMA_NODE", numa_node)) - self._handle = HostBuffer._allocator.alloc(size, backing, hp_size, numa, True) - if not self._handle: - raise RuntimeError(f"UMBPHostMemAllocator.alloc({size}) failed") - if use_hp and self._handle.actual_backing == UMBPHostBufferBacking.Anonymous: - print( - "WARNING: requested hugepages but kernel demoted to 4 KiB pages; " - "check vm.nr_hugepages and HugePages_Free in /proc/meminfo", - file=sys.stderr, - ) - self.ptr = int(self._handle.ptr) - self.size = size - - def __del__(self): - handle = getattr(self, "_handle", None) - if handle is not None and HostBuffer._allocator is not None: - try: - HostBuffer._allocator.free(handle) - except Exception: - pass - self._handle = None - - -if role in ("both", "writer"): - buf = HostBuffer(value_size) - ptr = buf.ptr - writer.register_memory(ptr, value_size) - - -# --- shared write helper --- - - -def write_all() -> bool: - payload = expected_payload(value_size) - ok = fail = 0 - for i in range(nr_objects): - if writer.put(make_key(i), payload, ptr, value_size): - ok += 1 - else: - fail += 1 - print(f"WRITE_DONE ok={ok} fail={fail} total={nr_objects}", flush=True) - writer.flush() - return fail == 0 - - -# --- write + reader setup (role-aware) --- -# -# role=writer : write the dataset, then host a TCP rendezvous on this node: -# signal READY (data flushed) -> wait for the reader's DONE -> release & exit. -# The writer must outlive the reads because UMBP serves them via RDMA out of -# this process's DRAM. -# role=reader : connect to the writer's rendezvous, wait for READY, read, then -# signal DONE. -# role=both : original single-process behavior, no rendezvous. - -if role == "writer": - srv = _sync_listen(args.sync_port) - write_all() - print( - f"WRITER_READY holding dataset; waiting for reader on :{args.sync_port}", - flush=True, - ) - conn, peer = srv.accept() - conn.sendall(b"READY\n") - print(f"reader connected from {peer[0]}:{peer[1]}; signalled READY", flush=True) - _sync_recv_until(conn, b"DONE") - print("READER_DONE; writer releasing dataset and exiting", flush=True) - conn.close() - srv.close() - sys.exit(0) - -sync_conn = None -if role == "reader": - sync_conn = _sync_connect(_resolve_host, args.sync_port) - print( - f"connected to writer sync {_resolve_host}:{args.sync_port}; waiting for READY", - flush=True, - ) - _sync_recv_until(sync_conn, b"READY") - print("writer signalled READY; starting reads", flush=True) - reader = _build_reader() -else: # both - write_all() - reader = writer.make_reader() - -# --- commands (reader side) --- - -exit_code = 0 - -if command == "correctness": - read_buf = HostBuffer(value_size) - read_ptr = read_buf.ptr - reader.register_memory(read_ptr, value_size) - hits = misses = mismatches = 0 - expected = expected_payload(value_size) - for i in range(nr_objects): - ctypes.memset(read_ptr, 0, value_size) - if not reader.get_into_ptr(make_key(i), read_ptr, value_size): - misses += 1 - continue - hits += 1 - if ctypes.string_at(read_ptr, value_size) != expected: - mismatches += 1 - print( - f"CORRECTNESS hits={hits} misses={misses} mismatches={mismatches} total={nr_objects}", - flush=True, - ) - if hits == nr_objects and mismatches == 0: - print("CORRECTNESS_OK", flush=True) - else: - print("CORRECTNESS_FAILED", flush=True) - exit_code = 1 - -elif command == "batch_perf": - passes = args.passes - batch_sizes = args.batch_sizes if args.batch_sizes else [1, 32, 128, nr_objects] - - keys = [make_key(i) for i in range(nr_objects)] - buf_total = value_size * nr_objects - whole = HostBuffer(buf_total) - base = whole.ptr - # Destination slot for each key. 'contiguous' packs them back-to-back so the - # IO engine can coalesce adjacent SGEs; 'shuffled' assigns each key a random - # slot in the same buffer so consecutive keys land at non-adjacent addresses - # (simulates a non-consecutive transfer; defeats adjacency coalescing). - if args.dest_layout == "shuffled": - import random - - slots = list(range(nr_objects)) - random.Random(args.shuffle_seed).shuffle(slots) - print(f"dest_layout=shuffled seed={args.shuffle_seed}", flush=True) - else: - slots = range(nr_objects) - ptrs = [base + slot * value_size for slot in slots] - sizes = [value_size] * nr_objects - - reader.register_memory(base, buf_total) - - for batch_size in batch_sizes: - batches = [keys[i : i + batch_size] for i in range(0, nr_objects, batch_size)] - ptr_batches = [ - ptrs[i : i + batch_size] for i in range(0, nr_objects, batch_size) - ] - size_batches = [ - sizes[i : i + batch_size] for i in range(0, nr_objects, batch_size) - ] - - # warmup - for kb, pb, sb in zip(batches, ptr_batches, size_batches): - reader.batch_get_into_ptr(kb, pb, sb) - - hits = misses = 0 - total_bytes = 0 - start = time.perf_counter() - for _ in range(passes): - for kb, pb, sb in zip(batches, ptr_batches, size_batches): - results = reader.batch_get_into_ptr(kb, pb, sb) - for r in results: - if r > 0: - hits += 1 - total_bytes += r - else: - misses += 1 - elapsed = time.perf_counter() - start - total_reqs = nr_objects * passes - total_batches = len(batches) * passes - print( - f"RESULT batch_size={batch_size} hits={hits} misses={misses} requests={total_reqs} " - f"batches={total_batches} duration={elapsed:.3f}s " - f"req_per_s={total_reqs / elapsed:.2f} batch_per_s={total_batches / elapsed:.2f} " - f"MiB_per_s={(total_bytes / 1024 / 1024) / elapsed:.2f} " - f"avg_latency_ms={(elapsed / total_reqs) * 1000:.3f}", - flush=True, - ) - -# In a split run, tell the writer we're finished so it can release the dataset -# and exit; the writer is blocked in _sync_recv_until(b"DONE") until this lands. -if sync_conn is not None: - try: - sync_conn.sendall(b"DONE\n") - sync_conn.close() - except OSError: - pass - -if exit_code: - sys.exit(exit_code) diff --git a/tests/python/umbp/umbp_prim_bench.py b/tests/python/umbp/umbp_prim_bench.py new file mode 100644 index 000000000..6d537675d --- /dev/null +++ b/tests/python/umbp/umbp_prim_bench.py @@ -0,0 +1,902 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Atomic UMBP primitive benchmark, swept over concurrency. + +One primitive per timed window and nothing else in it: every argument list is +built during an untimed prepare, so a measurement is the client call and not +the Python that assembles its arguments. Five primitives: + + put batch_put_from_ptr whole object write + get batch_get_into_ptr whole object read + exists batch_exists presence probe, moves no bytes + put_ranges batch_put_ranges_from_ptr write R ranges of one object + get_ranges batch_get_ranges_into_ptr read R ranges of one object + +Each runs at every rank count in --ranks, which is the axis that matters most +here: a shared-lock control plane and an exclusive-lock one look identical at +one rank and diverge at eight. Ranks are separate processes (the GIL makes +threads useless for this), started with multiprocessing's spawn so each builds +its own client, and held in lockstep by a Barrier around every timed window. + +Example +------- + python3 umbp_prim_bench.py umbp-server unix:///run/umbp/standalone/a.sock \\ + --ops put get exists put_ranges get_ranges --ranks 1 2 4 8 +""" +from __future__ import annotations + +import argparse +import ctypes +import multiprocessing as mp +import os +import random +import sys +import time +from dataclasses import asdict, dataclass, field +from typing import Callable + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +OPS = ("put", "get", "exists", "put_ranges", "get_ranges") + + +@dataclass(frozen=True) +class Config: + backend: str + address: str + ops: tuple[str, ...] + rank_counts: tuple[int, ...] + keys_per_rank: int + batch: int + passes: int + ranges: int + range_bytes: int + loc: str + register: bool + verify: bool + key_prefix: str + dram_capacity: int + #: Fraction of each `exists` call's keys that name something never + #: written. 0 probes only resident keys and never reaches the master. + exists_absent_frac: float = 0.0 + #: Where to write the machine-readable result, for a caller that gates on + #: these numbers rather than reading the table. Empty means don't. + json_path: str = "" + + @property + def object_bytes(self) -> int: + """One object is exactly its ranges, so whole and ranged ops move the + same bytes for the same key and stay comparable.""" + return self.ranges * self.range_bytes + + def keys( + self, op: str, ranks: int, rank: int, generation: int, absent: bool = False + ) -> list[str]: + # Disjoint per rank, as the linker's tp{r} suffix makes them: a shared + # keyspace would let one rank's put satisfy another's get and turn a + # concurrency sweep into a dedup measurement. + # + # The RANK COUNT is in the namespace too. Without it the ranks=2 step + # rewrites the keys the ranks=1 step already wrote, so every put after + # the first sweep step is an overwrite of a resident key -- which the + # store short-circuits, making the write primitive get faster the more + # concurrency it is under. + # The OP is in the namespace as well, so the two write primitives do + # not write each other's keys -- put and put_ranges both want a key + # that does not exist yet, and whichever ran second would otherwise be + # measuring an overwrite the store short-circuits. + # + # `absent` shifts the index past everything any op ever writes, so + # those keys are guaranteed missing by construction. Deriving them + # from the same namespace keeps the key SHAPE identical -- same + # length, same prefix -- so a probe that mixes the two is not + # secretly measuring two different string costs. + base = self.keys_per_rank if absent else 0 + return [ + f"{self.key_prefix}_{op}_n{ranks}_g{generation}_r{rank}_{base + i:07d}" + for i in range(self.keys_per_rank) + ] + + +# --------------------------------------------------------------------------- +# Buffers +# --------------------------------------------------------------------------- + + +class HostBuffer: + """Host DMA buffer from mori's allocator. + + Standalone-process mode resolves (region_base, offset) against the SERVER's + mapping of this memory, so it must be shm-backed or nothing can map it -- + and getting that wrong is quiet, surfacing as a per-key false that looks + exactly like the store refusing the operation. + """ + + def __init__(self, size: int, shm: bool): + from mori.umbp import UMBPHostBufferBacking, UMBPHostMemAllocator + + self._allocator = UMBPHostMemAllocator() + order = ( + [ + UMBPHostBufferBacking.AnonymousShmHugetlb, + UMBPHostBufferBacking.AnonymousShm, + ] + if shm + else [ + UMBPHostBufferBacking.AnonymousHugetlb, + UMBPHostBufferBacking.Anonymous, + ] + ) + self._handle = None + for backing in order: + self._handle = self._allocator.alloc(size, backing, 2 << 20, -1, True) + if self._handle: + break + if not self._handle: + raise RuntimeError(f"host alloc({size}) failed") + self.ptr = int(self._handle.ptr) + self.size = size + + def write(self, offset: int, payload: bytes) -> None: + ctypes.memmove(self.ptr + offset, payload, len(payload)) + + def read(self, offset: int, size: int) -> bytes: + return ctypes.string_at(self.ptr + offset, size) + + def close(self) -> None: + if getattr(self, "_handle", None) is not None: + self._allocator.free(self._handle) + self._handle = None + + +class DeviceBuffer: + """GPU buffer via raw HIP, so a read lands where the tree connector's reads + land. ctypes rather than torch keeps the bench usable in a bare container.""" + + _D2H, _H2D = 2, 1 + + def __init__(self, size: int, device: int): + self._hip = ctypes.CDLL("libamdhip64.so") + self._check(self._hip.hipSetDevice(ctypes.c_int(device)), "hipSetDevice") + raw = ctypes.c_void_p() + self._check( + self._hip.hipMalloc(ctypes.byref(raw), ctypes.c_size_t(size)), "hipMalloc" + ) + self._check(self._hip.hipMemset(raw, 0, ctypes.c_size_t(size)), "hipMemset") + self.ptr = raw.value + self.size = size + self.device = device + + @staticmethod + def device_count() -> int: + hip = ctypes.CDLL("libamdhip64.so") + n = ctypes.c_int(0) + if hip.hipGetDeviceCount(ctypes.byref(n)) != 0: + return 0 + return n.value + + def _check(self, rc: int, what: str) -> None: + if rc != 0: + raise RuntimeError(f"{what} failed: hip rc={rc}") + + def write(self, offset: int, payload: bytes) -> None: + self._check( + self._hip.hipMemcpy( + ctypes.c_void_p(self.ptr + offset), + payload, + ctypes.c_size_t(len(payload)), + ctypes.c_int(self._H2D), + ), + "hipMemcpy H2D", + ) + + def read(self, offset: int, size: int) -> bytes: + out = ctypes.create_string_buffer(size) + self._check( + self._hip.hipMemcpy( + out, + ctypes.c_void_p(self.ptr + offset), + ctypes.c_size_t(size), + ctypes.c_int(self._D2H), + ), + "hipMemcpy D2H", + ) + return out.raw + + def close(self) -> None: + ptr = getattr(self, "ptr", None) + if ptr: + self._hip.hipFree(ctypes.c_void_p(ptr)) + self.ptr = None + + +def make_buffer(cfg: Config, size: int, rank: int): + if cfg.loc == "gpu": + count = DeviceBuffer.device_count() + if count == 0: + raise RuntimeError("--loc gpu but no HIP device is visible") + return DeviceBuffer(size, rank % count) + return HostBuffer(size, shm=cfg.backend == "umbp-server") + + +# --------------------------------------------------------------------------- +# Client +# --------------------------------------------------------------------------- + + +def make_client(cfg: Config): + from mori.umbp import UMBPClient, UMBPConfig + + conf = UMBPConfig() + if cfg.backend == "umbp-server": + # Imported here, not at the top: a build without the standalone + # bindings can still run the embedded backend, and an unconditional + # import would turn that into an ImportError naming the wrong thing. + from mori.umbp import UMBPStandaloneProcessConfig + + sp = UMBPStandaloneProcessConfig() + sp.address = cfg.address + sp.startup_timeout_ms = int(os.environ.get("UMBP_STANDALONE_TIMEOUT_MS", 15000)) + # standalone_process ONLY: the factory tests `distributed` first, so a + # config carrying both quietly builds a DistributedClient instead. + conf.standalone_process = sp + else: # umbp-local: in-process client, no server and no master + conf.dram.capacity_bytes = cfg.dram_capacity + conf.dram.use_hugepages = False + return UMBPClient(conf) + + +def register(client, cfg: Config, buf) -> None: + """Register the whole buffer once, the way the tree connector does. + + The connector calls the two-argument form, which leaves `loc` at its CPU + default even for a device pointer -- mori detects the pointer itself. This + passes CPU for a GPU buffer too, deliberately: registering the way + production registers keeps the bench on the path that is actually shipped. + """ + if not cfg.register: + return + from mori.io import MemoryLocationType + + client.register_memory(buf.ptr, buf.size, MemoryLocationType.CPU, -1) + + +# --------------------------------------------------------------------------- +# Operations +# +# An Op is a prepare/invoke pair. Everything a call needs is built in +# prepare -- key lists, pointer lists, per-range size and offset tables -- so +# the timed window holds one client call and nothing else. Building those +# lists inside the loop is what made the old bench's per-call numbers include +# its own list comprehensions. +# --------------------------------------------------------------------------- + + +@dataclass +class Batch: + """One prepared primitive call.""" + + invoke: Callable[[], object] + keys: int + nbytes: int + #: The keys this call names, in slot order. Verification needs them to + #: know which payload belongs in which buffer slot. + key_names: tuple[str, ...] = () + #: Per-key expected truth, when it is not "all of them". A probe that + #: deliberately names absent keys expects False for those, and a False + #: there is the CORRECT answer -- scoring it as a failure would make the + #: absent-key sweep unusable. None means "expect every key to succeed". + expect: tuple[bool, ...] | None = None + + +class Op: + name: str + #: A put must land on keys that do not exist yet, or the store short-circuits + #: and the measurement is an overwrite of a resident key. + fresh_keys_per_pass = False + #: Reads need their objects present before the timed window opens. + needs_seed = True + + def __init__(self, cfg: Config, client, buf, ranks: int, rank: int): + self.cfg = cfg + self.client = client + self.buf = buf + self.ranks = ranks + self.rank = rank + + def slots(self, count: int) -> list[int]: + return [self.buf.ptr + i * self.cfg.object_bytes for i in range(count)] + + def prepare(self, generation: int) -> list[Batch]: + raise NotImplementedError + + +class PutOp(Op): + name = "put" + fresh_keys_per_pass = True + needs_seed = False + + def prepare(self, generation: int) -> list[Batch]: + cfg, out = self.cfg, [] + keys = cfg.keys(self.name, self.ranks, self.rank, generation) + for lo in range(0, len(keys), cfg.batch): + kb = keys[lo : lo + cfg.batch] + ptrs = self.slots(len(kb)) + sizes = [cfg.object_bytes] * len(kb) + out.append( + Batch( + lambda k=kb, p=ptrs, s=sizes: self.client.batch_put_from_ptr( + k, p, s + ), + len(kb), + len(kb) * cfg.object_bytes, + key_names=tuple(kb), + ) + ) + return out + + +class GetOp(Op): + name = "get" + + def prepare(self, generation: int) -> list[Batch]: + cfg, out = self.cfg, [] + keys = cfg.keys(self.name, self.ranks, self.rank, generation) + for lo in range(0, len(keys), cfg.batch): + kb = keys[lo : lo + cfg.batch] + ptrs = self.slots(len(kb)) + sizes = [cfg.object_bytes] * len(kb) + out.append( + Batch( + lambda k=kb, p=ptrs, s=sizes: self.client.batch_get_into_ptr( + k, p, s + ), + len(kb), + len(kb) * cfg.object_bytes, + key_names=tuple(kb), + ) + ) + return out + + +class ExistsOp(Op): + """Presence probe, optionally mixed with keys that were never written. + + The mix matters more than it looks. Under `local_first` a probe for a + key this node holds is answered locally and the master is never asked, so + an all-present probe measures only half the path. A local miss proves + nothing -- the key could live on a peer -- so every ABSENT key costs a + master round trip. That is the half a prefix cache actually pays on a + cold prefix, and the half where a control-plane regression shows up. + """ + + name = "exists" + + def absent_slots(self, count: int) -> set[int]: + """Which positions in a `count`-key call name an absent key. + + Spread evenly rather than blocked at one end, so no call is entirely + one population and every call pays a mix. + """ + n = int(round(count * self.cfg.exists_absent_frac)) + n = max(0, min(count, n)) + return {(i * count) // n for i in range(n)} if n else set() + + def prepare(self, generation: int) -> list[Batch]: + cfg, out = self.cfg, [] + present = cfg.keys(self.name, self.ranks, self.rank, generation) + # Indices at or past keys_per_rank are never written by any op, so + # these are absent by construction rather than by hoping for a miss. + absent = cfg.keys(self.name, self.ranks, self.rank, generation, absent=True) + for lo in range(0, len(present), cfg.batch): + kb = list(present[lo : lo + cfg.batch]) + holes = self.absent_slots(len(kb)) + for slot in holes: + kb[slot] = absent[lo + slot] + expect = tuple(slot not in holes for slot in range(len(kb))) + # No bytes at all: the whole call is the control plane. + out.append( + Batch( + lambda k=kb: self.client.batch_exists(k), + len(kb), + 0, + key_names=tuple(kb), + expect=expect, + ) + ) + return out + + +class RangedOp(Op): + """Shared geometry for the two ranged primitives. + + The ranges tile the object exactly, which the write side requires -- a put + that left a hole would be read back as garbage rather than as a miss. + """ + + def tables(self, count: int): + cfg = self.cfg + sizes = [[cfg.range_bytes] * cfg.ranges] * count + offsets = [[r * cfg.range_bytes for r in range(cfg.ranges)]] * count + ptrs = [ + [base + r * cfg.range_bytes for r in range(cfg.ranges)] + for base in self.slots(count) + ] + return ptrs, sizes, offsets + + +class PutRangesOp(RangedOp): + name = "put_ranges" + fresh_keys_per_pass = True + needs_seed = False + + def prepare(self, generation: int) -> list[Batch]: + cfg, out = self.cfg, [] + keys = cfg.keys(self.name, self.ranks, self.rank, generation) + for lo in range(0, len(keys), cfg.batch): + kb = keys[lo : lo + cfg.batch] + ptrs, sizes, offsets = self.tables(len(kb)) + objsz = [cfg.object_bytes] * len(kb) + out.append( + Batch( + lambda k=kb, o=objsz, p=ptrs, s=sizes, f=offsets: ( + self.client.batch_put_ranges_from_ptr(k, o, p, s, f) + ), + len(kb), + len(kb) * cfg.object_bytes, + key_names=tuple(kb), + ) + ) + return out + + +class GetRangesOp(RangedOp): + name = "get_ranges" + + def prepare(self, generation: int) -> list[Batch]: + cfg, out = self.cfg, [] + keys = cfg.keys(self.name, self.ranks, self.rank, generation) + for lo in range(0, len(keys), cfg.batch): + kb = keys[lo : lo + cfg.batch] + ptrs, sizes, offsets = self.tables(len(kb)) + out.append( + Batch( + lambda k=kb, p=ptrs, s=sizes, f=offsets: ( + self.client.batch_get_ranges_into_ptr(k, p, s, f) + ), + len(kb), + len(kb) * cfg.object_bytes, + key_names=tuple(kb), + ) + ) + return out + + +OP_TYPES: dict[str, type[Op]] = { + cls.name: cls for cls in (PutOp, GetOp, ExistsOp, PutRangesOp, GetRangesOp) +} + + +# --------------------------------------------------------------------------- +# Rank worker +# --------------------------------------------------------------------------- + + +@dataclass +class OpResult: + op: str + rank: int + calls: list[float] = field(default_factory=list) + keys: int = 0 + nbytes: int = 0 + wall: float = 0.0 + failures: int = 0 + mismatches: int = 0 + + +def _payload(key: str, size: int) -> bytes: + """The exact bytes `key` must hold, all `size` of them. + + A keyed pseudo-random stream rather than a repeated tag, because the + cheap patterns miss the bugs worth catching: with a 32-byte tag repeated + to fill the object, a range written or read at the wrong offset lands on + identical bytes and verifies clean. `random.Random(str)` seeds from a + SHA-512 of the string, so this is stable across processes and runs and + owes nothing to PYTHONHASHSEED, and `randbytes` is C-implemented so + filling a 64 KiB object costs microseconds. + """ + return random.Random(key).randbytes(size) + + +def _seed(cfg: Config, client, buf, op: str, ranks: int, rank: int) -> None: + """Untimed write of this op's read set (generation 0), plus the visibility wait. + + A put is not readable the instant it returns, so a read phase started too + early reports a total miss that is indistinguishable from a broken path. + """ + keys = cfg.keys(op, ranks, rank, 0) + for lo in range(0, len(keys), cfg.batch): + kb = keys[lo : lo + cfg.batch] + # Refill the slots for THIS chunk. The buffer only holds `batch` + # objects, so it is reused chunk after chunk; writing it once up front + # would leave every key past the first chunk holding the first + # chunk's bytes, and a full-payload check would then be verifying the + # wrong thing rather than nothing. + for i, key in enumerate(kb): + buf.write(i * cfg.object_bytes, _payload(key, cfg.object_bytes)) + ptrs = [buf.ptr + i * cfg.object_bytes for i in range(len(kb))] + res = client.batch_put_from_ptr(kb, ptrs, [cfg.object_bytes] * len(kb)) + if not all(res): + raise RuntimeError( + f"seed put failed for {sum(1 for r in res if not r)} keys" + ) + client.flush() + deadline = time.time() + float(os.environ.get("UMBP_VISIBLE_TIMEOUT_S", 120)) + while time.time() < deadline: + if all(client.batch_exists(keys[: cfg.batch])): + return + time.sleep(0.05) + raise RuntimeError("seeded keys never became visible") + + +def _run_op( + cfg: Config, client, buf, ranks: int, rank: int, name: str, barrier +) -> OpResult: + op = OP_TYPES[name](cfg, client, buf, ranks, rank) + result = OpResult(op=name, rank=rank) + + # Generation 0 is the read set. A write primitive gets its own generations + # so every pass is a genuine insert. + if op.needs_seed: + _seed(cfg, client, buf, name, ranks, rank) + + for p in range(cfg.passes): + generation = 1 + p if op.fresh_keys_per_pass else 0 + batches = op.prepare(generation) + barrier.wait() + start = time.perf_counter() + for batch in batches: + t0 = time.perf_counter() + res = batch.invoke() + result.calls.append(time.perf_counter() - t0) + result.keys += batch.keys + result.nbytes += batch.nbytes + if res is not None: + if batch.expect is None: + result.failures += sum(1 for r in res if not r) + else: + # A probe that names absent keys expects False for them; + # counting those as failures would make the sweep useless, + # and NOT checking them would miss a false positive -- + # which for a prefix-cache probe is the worse bug. + result.failures += sum( + 1 for r, e in zip(res, batch.expect) if bool(r) != e + ) + result.wall += time.perf_counter() - start + barrier.wait() + + if cfg.verify and name in ("get", "get_ranges"): + result.mismatches += _verify(cfg, op, buf) + return result + + +def _verify(cfg: Config, op: "Op", buf) -> int: + """Re-read every key through the op's OWN call and compare every byte. + + Untimed, and deliberately after the timed windows rather than inside + them. Three things make this stronger than sampling a tag: + + * it is every key, not the first two, so a fault that only affects later + chunks is reachable; + * it is the whole object, so truncation and partial writes show up; + * it goes through the op's own call, so `get_ranges` is verified with its + real offset and size tables -- which is the only way a wrong range + offset is ever caught. + + The buffer is zeroed first: without that a call which quietly moved + nothing leaves the previous chunk's correct bytes in place and verifies + clean. + """ + zeros = b"\x00" * (cfg.batch * cfg.object_bytes) + mismatches = 0 + for batch in op.prepare(0): + buf.write(0, zeros[: batch.keys * cfg.object_bytes]) + batch.invoke() + for slot, key in enumerate(batch.key_names): + got = buf.read(slot * cfg.object_bytes, cfg.object_bytes) + if got != _payload(key, cfg.object_bytes): + mismatches += 1 + return mismatches + + +def _worker(ranks: int, rank: int, cfg: Config, barrier, queue) -> None: + buf = client = None + try: + client = make_client(cfg) + buf = make_buffer(cfg, cfg.batch * cfg.object_bytes, rank) + register(client, cfg, buf) + for name in cfg.ops: + queue.put(_run_op(cfg, client, buf, ranks, rank, name, barrier)) + except BaseException as err: # noqa: BLE001 - the parent needs the reason + queue.put(("ERROR", rank, f"{type(err).__name__}: {err}")) + # Release every barrier the siblings are still waiting on, or they hang + # until the parent's join timeout instead of reporting this failure. + try: + barrier.abort() + except Exception: + pass + finally: + if buf is not None: + buf.close() + + +# --------------------------------------------------------------------------- +# Coordinator +# --------------------------------------------------------------------------- + + +def run_rank_count(cfg: Config, ranks: int) -> dict[str, list[OpResult]]: + ctx = mp.get_context("spawn") + barrier = ctx.Barrier(ranks) + queue = ctx.Queue() + procs = [ + ctx.Process(target=_worker, args=(ranks, r, cfg, barrier, queue), daemon=True) + for r in range(ranks) + ] + for p in procs: + p.start() + + expected = ranks * len(cfg.ops) + by_op: dict[str, list[OpResult]] = {name: [] for name in cfg.ops} + errors: list[str] = [] + for _ in range(expected): + try: + item = queue.get(timeout=1800) + except Exception: + errors.append("timed out waiting for a rank result") + break + if isinstance(item, tuple): + errors.append(f"rank {item[1]}: {item[2]}") + break + by_op[item.op].append(item) + + for p in procs: + p.join(timeout=30) + if p.is_alive(): + p.terminate() + if errors: + raise RuntimeError("; ".join(errors)) + return by_op + + +def _pct(values: list[float], q: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + return ordered[min(len(ordered) - 1, int(q * len(ordered)))] + + +@dataclass +class Row: + op: str + ranks: int + calls: int + p50_ms: float + p99_ms: float + keys_per_s: float + mib_per_s: float + failures: int + mismatches: int + + +def summarise(op: str, ranks: int, results: list[OpResult]) -> Row: + calls = [c for r in results for c in r.calls] + # Wall is the straggler's, not the sum: the ranks run inside one barrier, so + # the window is as long as the slowest of them. + wall = max((r.wall for r in results), default=0.0) + keys = sum(r.keys for r in results) + nbytes = sum(r.nbytes for r in results) + return Row( + op=op, + ranks=ranks, + calls=len(calls), + p50_ms=_pct(calls, 0.5) * 1e3, + p99_ms=_pct(calls, 0.99) * 1e3, + keys_per_s=keys / wall if wall else 0.0, + mib_per_s=(nbytes / 2**20) / wall if wall else 0.0, + failures=sum(r.failures for r in results), + mismatches=sum(r.mismatches for r in results), + ) + + +def report(cfg: Config, rows: list[Row]) -> None: + print( + f"\n{'op':<12}{'ranks':>6}{'calls':>7}{'p50_ms':>10}{'p99_ms':>10}" + f"{'keys/s':>12}{'MiB/s':>11}{'scale':>7}{'fail':>6}{'bad':>5}" + ) + base: dict[str, float] = {} + for row in rows: + base.setdefault(row.op, row.keys_per_s) + scale = row.keys_per_s / base[row.op] if base[row.op] else 0.0 + print( + f"{row.op:<12}{row.ranks:>6}{row.calls:>7}{row.p50_ms:>10.3f}" + f"{row.p99_ms:>10.3f}{row.keys_per_s:>12,.0f}{row.mib_per_s:>11,.1f}" + f"{scale:>6.2f}x{row.failures:>6}{row.mismatches:>5}" + ) + print( + f"\ngeometry: object={cfg.object_bytes}B ({cfg.ranges} x {cfg.range_bytes}B) " + f"keys/rank={cfg.keys_per_rank} batch={cfg.batch} passes={cfg.passes} " + f"loc={cfg.loc} register={cfg.register} " + f"exists_absent_frac={cfg.exists_absent_frac:g}", + flush=True, + ) + + +def write_json(cfg: Config, rows: list[Row], path: str) -> None: + """Dump the run so a gate can read numbers instead of parsing the table. + + The geometry travels with the rows on purpose: a scaling floor is only + meaningful against the shape that produced it, and a result file that + does not say what was measured cannot be compared to anything later. + """ + import json + + payload = { + "backend": cfg.backend, + "address": cfg.address, + "geometry": { + "object_bytes": cfg.object_bytes, + "ranges": cfg.ranges, + "range_bytes": cfg.range_bytes, + "keys_per_rank": cfg.keys_per_rank, + "batch": cfg.batch, + "passes": cfg.passes, + "loc": cfg.loc, + "register": cfg.register, + "exists_absent_frac": cfg.exists_absent_frac, + }, + "rows": [asdict(row) for row in rows], + } + with open(path, "w") as handle: + json.dump(payload, handle, indent=2) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def parse_args(argv: list[str] | None = None) -> Config: + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("backend", choices=["umbp-server", "umbp-local"]) + p.add_argument( + "address", + nargs="?", + default="", + help="standalone server address, e.g. unix:///run/umbp/standalone/a.sock " + "(umbp-server only)", + ) + p.add_argument("--ops", nargs="+", choices=OPS, default=list(OPS)) + p.add_argument( + "--ranks", + nargs="+", + type=int, + default=[1, 8], + help="rank counts to sweep (default: 1 8). Each is a separate set of " + "processes; a primitive is run at every one of them.", + ) + p.add_argument("--keys-per-rank", type=int, default=1024) + p.add_argument("--batch", type=int, default=256, help="keys per call") + p.add_argument("--passes", type=int, default=5) + p.add_argument("--ranges", type=int, default=8, help="ranges per object") + p.add_argument("--range-bytes", type=int, default=131072) + p.add_argument("--loc", choices=["host", "gpu"], default="gpu") + p.add_argument("--no-register", dest="register", action="store_false") + p.add_argument("--no-verify", dest="verify", action="store_false") + p.add_argument("--key-prefix", default=None) + p.add_argument( + "--exists-absent-frac", + type=float, + default=0.0, + help="fraction of each exists call naming keys that were never written " + "(0.0-1.0). Absent keys cannot be answered locally, so they are the " + "half of the probe that consults the master.", + ) + p.add_argument( + "--dram-capacity-gib", + type=int, + default=64, + help="umbp-local only; the server owns its own pool", + ) + p.add_argument( + "--json", + dest="json_path", + default="", + help="also write the rows to this path as JSON, for a gate to read", + ) + a = p.parse_args(argv) + if a.backend == "umbp-server" and not a.address: + p.error("umbp-server needs a server address") + if any(r < 1 for r in a.ranks): + p.error("--ranks must all be >= 1") + if not 0.0 <= a.exists_absent_frac <= 1.0: + p.error("--exists-absent-frac must be between 0.0 and 1.0") + return Config( + backend=a.backend, + address=a.address, + ops=tuple(a.ops), + rank_counts=tuple(a.ranks), + keys_per_rank=a.keys_per_rank, + batch=min(a.batch, a.keys_per_rank), + passes=a.passes, + ranges=a.ranges, + range_bytes=a.range_bytes, + loc=a.loc, + register=a.register, + verify=a.verify, + # Every rank count is a fresh set of processes writing the same keys, so + # the prefix has to differ per invocation or a later sweep step reads + # what an earlier one wrote and calls it a hit. + exists_absent_frac=a.exists_absent_frac, + key_prefix=a.key_prefix or f"p{os.getpid() % 100000:05d}", + dram_capacity=a.dram_capacity_gib << 30, + json_path=a.json_path, + ) + + +def main(argv: list[str] | None = None) -> int: + cfg = parse_args(argv) + print( + f"umbp_prim_bench backend={cfg.backend} address={cfg.address or '-'} " + f"ops={','.join(cfg.ops)} ranks={','.join(map(str, cfg.rank_counts))} " + f"key_prefix={cfg.key_prefix}", + flush=True, + ) + rows: list[Row] = [] + for ranks in cfg.rank_counts: + by_op = run_rank_count(cfg, ranks) + for name in cfg.ops: + rows.append(summarise(name, ranks, by_op[name])) + row = rows[-1] + print( + f" ranks={ranks:<3} {row.op:<11} p50={row.p50_ms:8.3f} ms " + f"keys/s={row.keys_per_s:>11,.0f} fail={row.failures}", + flush=True, + ) + # Report grouped by primitive so the rank sweep of one op reads down a + # column rather than being interleaved with the others. + rows.sort(key=lambda r: (cfg.ops.index(r.op), r.ranks)) + report(cfg, rows) + if cfg.json_path: + write_json(cfg, rows, cfg.json_path) + bad = sum(r.failures + r.mismatches for r in rows) + if bad: + print(f"FAILED: {bad} failed or mismatched operations", file=sys.stderr) + return 1 if bad else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/python/umbp/umbp_prim_ci.py b/tests/python/umbp/umbp_prim_ci.py new file mode 100644 index 000000000..171517aaf --- /dev/null +++ b/tests/python/umbp/umbp_prim_ci.py @@ -0,0 +1,1123 @@ +#!/usr/bin/env python3 +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""CI gate over the five UMBP primitives, swept over concurrency. + +Brings up standalone-process mode -- `umbp_standalone_server` plus, for the +second arm, a `umbp_master` in front of it -- runs `umbp_prim_bench.py` +against each, and turns the rows into a pass/fail verdict. + + python3 umbp_prim_ci.py --arms nomaster master + +Two things are gated, and the second is the reason this exists. + +**Correctness.** Every one of `put`, `get`, `exists`, `put_ranges` and +`get_ranges` must report zero failed operations and zero read-back +mismatches at every rank count. Eight concurrent clients writing and +reading disjoint keyspaces through one server is a real functional test on +its own. + +It is not, however, where a control-plane regression shows up: a build whose +resolve path serializes still returns every correct answer, just slowly. Do +not read a correctness pass as evidence about performance, or the reverse -- +and treat a correctness failure on ONE arm of an A/B as a claim about the +harness until the pool has been ruled out (see `Pass.resident_bytes`). + +**Concurrency scaling of `exists`.** Aggregate `batch_exists` throughput +must grow by at least `--exists-scaling-floor` from the smallest rank count +to the largest. This is the axis a whole benchmark suite missed once: an +exclusive lock over the resolve path is invisible at one rank and fatal at +eight, so an absolute latency number taken at a single concurrency cannot +see it, and comparing absolute numbers across CI runs on a shared machine +is noise. A ratio measured inside one run is not. + +The floor is 2.0x, and it was set by measuring rather than by borrowing a +number from an idle machine. On a busy node, three identical runs of one +unmodified build scored 3.50x, 5.39x and 4.33x, with the one-rank baseline +steady to within 3%; the regression this is aimed at scored 0.75x. So the +floor sits ~43% under the worst healthy run seen and 2.7x over the +regressed one. + +That margin exists because of the window, not because the threshold is +generous. Measured through the correctness pass's geometry -- a sub- +millisecond window -- the same three runs scattered 1.59x, 1.88x and 3.49x +and were not even monotonic in rank count. If this number ever reads +noisy, the fix is a longer window (`--scaling-passes`), not a lower floor. + +`exists` is the gated op deliberately -- it is the only one of the five +that moves no bytes, so its scaling is a property of the control plane and +not of how much DRAM bandwidth the runner had left. The other four are +measured and reported, but not gated: reads scale with the memory system, +and both a good and a bad build of the write path have been observed to +*anti*-scale under this geometry, so a floor there would be either +vacuous or permanently red. If a write-path scaling gate is ever wanted, +it needs its own baseline, not this one's. +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +import signal +import socket +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass +from pathlib import Path + +BENCH = Path(__file__).resolve().parent / "umbp_prim_bench.py" + +#: Aggregate `exists` throughput at the largest rank count, over the smallest. +#: See SCALING_* below for why this is measured in its own pass, and the +#: module docstring for how the floor was picked. +DEFAULT_EXISTS_SCALING_FLOOR = 2.0 + +#: Geometry for the scaling pass, which measures `exists` and nothing else. +#: +#: It is deliberately not the correctness pass's geometry. A `batch_exists` +#: call costs tens of microseconds, so the correctness pass's few dozen calls +#: add up to a timed window under a millisecond -- short enough that process +#: scheduling jitter, not the control plane, decides the throughput. Measured +#: on one loaded node, three identical runs of a single build scattered across +#: 1.59x, 1.88x and 3.49x on such a window; no floor survives that. So this +#: pass buys a window in the tens of milliseconds by issuing thousands of +#: calls, and buys it cheaply by shrinking the object to 4 KiB -- `exists` +#: moves no bytes, so object size costs only seed time and pool space. +SCALING_KEYS_PER_RANK = 8192 +SCALING_BATCH = 256 +SCALING_PASSES = 100 +SCALING_RANGE_BYTES = 4096 + +#: Absent-key fractions the scaling sweep runs at. 0.0 is the calibrated +#: gate (see the floor above); the rest are measured and reported so the +#: master-answered half of the probe is under the same concurrency sweep. +EXISTS_ABSENT_FRACS = (0.0, 0.5, 1.0) + +#: Storage media a backend can be. `hbm` is GPU memory -- the server's, not +#: the client buffer's, which is what --loc selects and is a separate axis. +MEDIA_KINDS = ("dram", "hbm", "ssd") + +#: Ops whose read set has to be written before the timed window, and so +#: occupy the pool for the whole run. +SEEDED_OPS = ("get", "exists", "get_ranges") +#: Ops that insert a fresh generation of keys on every pass. +WRITING_OPS = ("put", "put_ranges") + + +# --------------------------------------------------------------------------- +# Gate +# +# Pure: it takes parsed result payloads and returns findings. Nothing here +# imports mori or touches a process, so it is unit-testable against recorded +# numbers -- see test_umbp_prim_gate.py, which feeds it both arms of the +# regression this floor was calibrated on. +# --------------------------------------------------------------------------- + + +@dataclass +class Finding: + arm: str + check: str + value: str + threshold: str + ok: bool + detail: str = "" + #: False for a number that is measured and printed but does not decide + #: the build. Used for the absent-key fractions, which have no + #: calibrated baseline yet -- shipping a threshold nobody has measured is + #: how a gate becomes a flake. + gated: bool = True + + +def _rows_for(payload: dict, op: str) -> list[dict]: + return sorted( + (r for r in payload["rows"] if r["op"] == op), key=lambda r: r["ranks"] + ) + + +def check_correctness(arm: str, payload: dict, label: str = "") -> list[Finding]: + rows = payload["rows"] + bad = [r for r in rows if r["failures"] or r["mismatches"]] + detail = "; ".join( + f"{r['op']}@{r['ranks']}: {r['failures']} failed, {r['mismatches']} mismatched" + for r in bad + ) + return [ + Finding( + arm=arm, + check=f"correctness {label} ({len(rows)} rows)".replace(" ", " "), + value=f"{len(bad)} bad", + threshold="0 bad", + ok=not bad, + detail=detail, + ) + ] + + +def check_exists_scaling( + arm: str, payload: dict, floor: float, label: str = "", gated: bool = True +) -> list[Finding]: + rows = _rows_for(payload, "exists") + if len(rows) < 2: + return [ + Finding( + arm=arm, + check=f"exists scaling {label}", + gated=gated, + value="n/a", + threshold=f"{floor:.2f}x", + ok=False, + detail="needs at least two rank counts; the gate is a ratio", + ) + ] + low, high = rows[0], rows[-1] + # A zero at one rank is not a scaling result, it is a broken run, and + # dividing by it would hand back either inf or a crash. + if low["keys_per_s"] <= 0: + return [ + Finding( + arm=arm, + check=f"exists scaling {label} {low['ranks']}->{high['ranks']}", + value="0 keys/s at the base rank count", + threshold=f"{floor:.2f}x", + ok=False, + gated=gated, + ) + ] + scale = high["keys_per_s"] / low["keys_per_s"] + return [ + Finding( + arm=arm, + check=f"exists scaling {label} {low['ranks']}->{high['ranks']}", + value=f"{scale:.2f}x", + threshold=(f">= {floor:.2f}x" if gated else "(not gated)"), + ok=(scale >= floor) if gated else True, + gated=gated, + detail=( + f"{low['keys_per_s']:,.0f} -> {high['keys_per_s']:,.0f} keys/s" + if scale < floor + else "" + ), + ) + ] + + +def _absent_frac(payload: dict) -> float: + return float(payload.get("geometry", {}).get("exists_absent_frac", 0.0)) + + +def gate(results: dict[str, dict], floor: float) -> list[Finding]: + """Correctness over every pass; the scaling floor over the calibrated one. + + Correctness is checked on the absent-key passes too, and that is not + ceremony: those passes assert the probe answers False for a key that was + never written, so a false POSITIVE lands here. For a prefix-cache probe + that is the more dangerous direction -- it makes the engine read a key + that does not exist. + + The scaling floor applies only at absent_frac = 0, which is the one + measured against both arms of a known regression. The other fractions + put the master-answered half of the probe under the same sweep and are + reported, but they gate nothing until someone baselines them. + """ + findings: list[Finding] = [] + for arm, arm_results in results.items(): + for name, payload in sorted(arm_results.items()): + findings += check_correctness(arm, payload, label=name) + for name, payload in sorted(arm_results.items()): + if not name.startswith("scaling"): + continue + frac = _absent_frac(payload) + findings += check_exists_scaling( + arm, + payload, + floor, + label=f"absent={frac:g}", + gated=(frac == 0.0), + ) + return findings + + +def scaling_table(results: dict[str, dict]) -> list[str]: + """Every op at every rank count, including the four that are not gated. + + Reported rather than gated: reads scale with the memory system and both a + good and a regressed build have been seen to anti-scale on the writes, so + a floor there would be either vacuous or permanently red. Printing them + still makes a change visible to whoever reads the log. + """ + lines = [ + f"\n{'arm':<12}{'pass':<12}{'op':<12}{'ranks':>7}{'keys/s':>14}{'scale':>8}" + ] + for arm, arm_results in results.items(): + for which, payload in arm_results.items(): + for op in ("put", "get", "exists", "put_ranges", "get_ranges"): + rows = _rows_for(payload, op) + base = rows[0]["keys_per_s"] if rows else 0.0 + for row in rows: + scale = row["keys_per_s"] / base if base else 0.0 + lines.append( + f"{arm:<12}{which:<12}{op:<12}{row['ranks']:>7}" + f"{row['keys_per_s']:>14,.0f}{scale:>7.2f}x" + ) + return lines + + +def report_findings(findings: list[Finding]) -> bool: + print(f"\n{'arm':<10}{'check':<42}{'value':>16}{'threshold':>14}{'':>8}") + for f in findings: + verdict = " PASS" if f.ok else " FAIL" + if not f.gated: + verdict = " info" + print(f"{f.arm:<10}{f.check:<42}{f.value:>16}{f.threshold:>14}{verdict:>8}") + if f.detail: + print(f"{'':<10} {f.detail}") + return all(f.ok for f in findings if f.gated) + + +# --------------------------------------------------------------------------- +# Media +# +# What the SERVER stores in. Three ways to say it, and which one applies is +# forced by what the server can actually be told: +# +# * one DRAM backend -- the implicit embedded deployment, no selector env. +# Kept as the default because it is the configuration the scaling floor was +# calibrated against; switching the default would silently re-baseline it. +# * one SSD backend -- UMBP_DISTRIBUTED_MEDIUM=SSD plus the SSD dir and +# capacity. Naming a medium makes the server build an explicit distributed +# config, which is also why the page size stops being optional there. +# * anything else -- a backend-policy file. HBM has NO capacity env +# var at all (Validate() demands hbm.capacity_bytes > 0 and nothing parses +# one), so GPU memory is reachable ONLY this way; and mixtures need +# per-backend weights, which only the policy expresses. +# +# STATUS, 2026-09-03. Only `--media dram` is verified end to end against a +# live server. The rest is wired and unit-tested but does NOT yet work: +# +# ssd every put fails with "[SsdBackend] Resolve: staging arena busy; +# rolled back 0 reservations". Rolled back ZERO means the very first +# AcquireStagingSpanLocked failed, and it failed with 256 free pages +# against a 16-key batch -- so this is not arena sizing, which is +# already computed from the widest call. AcquireStagingSpanLocked +# returns empty for page_count == 0, so the likely cause is a +# per-key page count of zero; that is inside the backend, not here. +# combo the policy file now loads (its capacity needs a unit suffix, which +# cost a round), but the server then reports +# `medium=[DRAM pool=...]` with no SSD backend at all, so the policy +# is not visibly taking effect on the media. Unresolved. +# hbm untested. +# +# Do not read a passing run at another --media as evidence until those are +# settled: `dram` is the only one that has produced a real measurement. +# --------------------------------------------------------------------------- + + +@dataclass +class Backend: + """One medium in the server's storage, and its share of the pool.""" + + kind: str + weight: float + path: str = "" # ssd only + device: int = 0 # hbm only + + @property + def name(self) -> str: + return f"{self.kind}{self.device}" if self.kind == "hbm" else self.kind + + +def parse_media( + spec: str, ssd_dirs: list[str], hbm_devices: list[int] +) -> list[Backend]: + """`dram`, `ssd`, `hbm`, `dram:70,ssd:30`, `dram:50,ssd:30,hbm:20`. + + Weights are relative and normalised to 100; omitted weights split what is + left evenly, so `dram,ssd` is 50/50 without having to say so. + """ + entries = [part.strip() for part in spec.split(",") if part.strip()] + if not entries: + raise ValueError("--media is empty") + parsed: list[tuple[str, float | None]] = [] + for entry in entries: + kind, _, weight = entry.partition(":") + kind = kind.strip().lower() + if kind not in MEDIA_KINDS: + raise ValueError( + f"unknown medium {kind!r}; pick from {', '.join(MEDIA_KINDS)}" + ) + parsed.append((kind, float(weight) if weight else None)) + if len({k for k, _ in parsed}) != len(parsed): + raise ValueError("--media names a medium twice; give one weight instead") + + named = sum(w for _, w in parsed if w is not None) + unweighted = [i for i, (_, w) in enumerate(parsed) if w is None] + if unweighted: + if named >= 100: + raise ValueError(f"weights already total {named:g}; nothing left to split") + share = (100.0 - named) / len(unweighted) + parsed = [(k, share if w is None else w) for k, w in parsed] + total = sum(w for _, w in parsed) + if total <= 0: + raise ValueError("--media weights must be positive") + + out, ssd_i = [], 0 + for kind, weight in parsed: + weight = weight * 100.0 / total + if kind == "ssd": + if ssd_i >= len(ssd_dirs): + raise ValueError("more ssd backends than --ssd-dirs entries") + out.append(Backend(kind, weight, path=ssd_dirs[ssd_i])) + ssd_i += 1 + elif kind == "hbm": + out.append(Backend(kind, weight, device=hbm_devices[0])) + else: + out.append(Backend(kind, weight)) + return out + + +def backend_policy(backends: list[Backend], capacity: int) -> dict: + """One logical tier over every backend, split by weight. + + A single tier on purpose: this measures placement across media on one + node, not promotion and demotion between tiers, which would put an + eviction policy inside a latency measurement. + """ + entry = "bench" + spec: dict = {} + for b in backends: + share = max(1 << 20, int(capacity * b.weight / 100.0)) + # The unit is not optional: the loader rejects a bare byte count with + # "capacity unit must be B, KiB, MiB, GiB, or TiB". + one: dict = {"type": b.kind, "capacity": f"{share}B"} + if b.kind == "ssd": + one["path"] = b.path + elif b.kind == "hbm": + one["devices"] = [b.device] + spec[b.name] = one + return { + "schema_version": 1, + "entry_tier": entry, + "backends": spec, + "tiers": [ + {"name": entry, "backends": {b.name: round(b.weight, 3) for b in backends}} + ], + } + + +# --------------------------------------------------------------------------- +# Deployment +# --------------------------------------------------------------------------- + + +def _outbound_ip() -> str: + """This host's address on the route out, which is what the master hands + peers and what the IO engine binds. 127.0.0.1 does not work: the address + is published, not merely dialled.""" + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + sock.connect(("1.1.1.1", 53)) + return sock.getsockname()[0] + except OSError: + return "127.0.0.1" + finally: + sock.close() + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + return sock.getsockname()[1] + + +def _require_gpu() -> None: + """The master arm needs a GPU; say so before the server dies proving it. + + A master-backed server builds a MoRI IO engine, which initialises ROCm. + In a container without /dev/kfd that fails inside engine construction and + the process exits 255 after printing only "rsmi_num_monitor_devices + reported 0 GPUs" -- which reads like an unrelated telemetry warning. The + masterless arm has no engine and does not care. + """ + if not os.path.exists("/dev/kfd"): + raise RuntimeError( + "the master arm needs a GPU (/dev/kfd is not present): a " + "master-backed standalone server builds a MoRI IO engine, which " + "initialises ROCm. Run the container with --device=/dev/kfd " + "--device=/dev/dri (docker/ci_run.sh already does), or use " + "--arms nomaster." + ) + + +def _binary(env_name: str, override: str) -> str: + """Resolve a UMBP binary. + + Importing mori.umbp is what publishes UMBP_STANDALONE_BIN and + UMBP_MASTER_BIN, pointing at the binaries packaged beside the extension + module -- so the server under test is always the one built from the tree + that produced the installed wheel, never a stray one on PATH. + """ + if override: + return override + try: + import mori.umbp # noqa: F401 (imported for its import-time side effect) + except ImportError as err: + raise RuntimeError( + f"cannot import mori.umbp to locate {env_name}: {err}" + ) from err + path = os.environ.get(env_name, "") + if not path or not os.access(path, os.X_OK): + raise RuntimeError( + f"{env_name} is unset or not executable ({path!r}). Build mori with " + "BUILD_UMBP=ON, or pass the path explicitly." + ) + return path + + +class Deployment: + """A standalone server, with or without a master in front of it. + + Context manager: everything it starts is torn down on the way out, + including on an exception, because a leaked server holds both the unix + socket and the pool and the next arm would fail to bind. + """ + + def __init__( + self, + arm: str, + run_dir: Path, + dram_bytes: int, + page_bytes: int, + server_bin: str, + master_bin: str, + ready_timeout: float, + backends: list[Backend] | None = None, + ssd_staging_slots: int = 0, + ssd_staging_bytes: int = 0, + ): + self.ssd_staging_slots = ssd_staging_slots + self.ssd_staging_bytes = ssd_staging_bytes + self.backends = backends or [Backend("dram", 100.0)] + self.arm = arm + self.with_master = arm == "master" + self.run_dir = run_dir + self.dram_bytes = dram_bytes + self.page_bytes = page_bytes + self.server_bin = server_bin + self.master_bin = master_bin + self.ready_timeout = ready_timeout + self.procs: list[tuple[str, subprocess.Popen]] = [] + # The socket lives in a short path on purpose: an AF_UNIX path is + # capped at 108 bytes, and a workspace-relative one silently exceeds + # it on a CI checkout. The server derives the fd socket by replacing + # the .grpc.sock suffix, so the name has to end in exactly that. + self.sock_dir = Path(tempfile.mkdtemp(prefix=f"umbp-ci-{arm}-", dir="/tmp")) + self.address = f"unix://{self.sock_dir}/node.grpc.sock" + + # -- lifecycle --------------------------------------------------------- + + def __enter__(self) -> "Deployment": + try: + master_address = self._start_master() if self.with_master else "" + self._start_server(master_address) + except BaseException: + self.__exit__(*sys.exc_info()) + raise + return self + + def __exit__(self, *_exc) -> None: + for name, proc in reversed(self.procs): + if proc.poll() is not None: + continue + proc.send_signal(signal.SIGTERM) + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + print(f"[ci] {name} ignored SIGTERM, killing", flush=True) + proc.kill() + proc.wait(timeout=10) + self.procs.clear() + shutil.rmtree(self.sock_dir, ignore_errors=True) + + # -- processes --------------------------------------------------------- + + def _spawn(self, name: str, argv: list[str], env: dict[str, str]) -> Path: + log = self.run_dir / f"{self.arm}_{name}.log" + handle = open(log, "wb") + print(f"[ci] starting {name}: {' '.join(argv)}", flush=True) + proc = subprocess.Popen( + argv, stdout=handle, stderr=subprocess.STDOUT, env={**os.environ, **env} + ) + self.procs.append((name, proc)) + return log + + def _start_master(self) -> str: + ip = _outbound_ip() + port = _free_port() + log = self._spawn( + "master", + [self.master_bin, f"0.0.0.0:{port}", str(_free_port())], + {"MORI_UMBP_LOG_LEVEL": "info", "MORI_GLOBAL_LOG_LEVEL": "info"}, + ) + address = f"{ip}:{port}" + # The master logs its address before it serves, so the log is not the + # readiness signal here -- a completed TCP connect is. + self._await( + lambda: self._can_connect(ip, port), + what=f"master on {address}", + log=log, + proc_name="master", + ) + return address + + def _staging_env(self) -> dict: + """Size the SSD read-staging arena for the batches actually issued. + + A Resolve reserves one arena page per key IN THE WHOLE CALL, and if it + cannot get them all it rolls the reservation back and fails every key + with "staging arena busy" -- which surfaces as `seed put failed`, not + as anything mentioning staging. The default arena is 16 slots, so a + default 128-key batch cannot ever succeed. Slots must therefore scale + with the batch, and each slot must be at least one whole object (the + buffer is divided evenly among them). + """ + if not self.ssd_staging_slots: + return {} + return { + "UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SLOTS": str(self.ssd_staging_slots), + "UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SIZE": str(self.ssd_staging_bytes), + } + + def _media_env(self) -> dict: + """Translate the requested media into what the server can be told.""" + kinds = {b.kind for b in self.backends} + if len(self.backends) == 1 and kinds == {"dram"}: + # Implicit embedded DRAM: no selector, so the factory fills in the + # deployment. This is the shape the scaling floor was measured + # against; do not "tidy" it into an explicit medium without + # re-baselining, because they are not the same code path. + return {"UMBP_SSD_ENABLED": "0", "UMBP_DRAM_CAPACITY": str(self.dram_bytes)} + + if len(self.backends) == 1 and kinds == {"ssd"}: + ssd = self.backends[0] + os.makedirs(ssd.path, exist_ok=True) + return { + "UMBP_DISTRIBUTED_MEDIUM": "SSD", + "UMBP_SSD_ENABLED": "1", + "UMBP_SSD_DIR": ssd.path, + "UMBP_SSD_CAPACITY": str(self.dram_bytes), + # Explicit distributed config: the page size stops being + # optional, and 0 is rejected outright. + "UMBP_DISTRIBUTED_DRAM_PAGE_SIZE": str(self.page_bytes), + **self._staging_env(), + } + + # Everything else -- any mixture, and any HBM at all -- is a policy. + for b in self.backends: + if b.kind == "ssd": + os.makedirs(b.path, exist_ok=True) + path = self.run_dir / f"{self.arm}_backend_policy.json" + path.write_text( + json.dumps(backend_policy(self.backends, self.dram_bytes), indent=2) + ) + return { + "UMBP_BACKEND_POLICY": str(path), + # A medium must still be named and must still be one of the three + # words the server accepts -- it reads this with getenv() and an + # EMPTY string is a value, not an absence, so leaving it blank is + # rejected outright. The policy supplies the real backends; this + # only has to validate. It cannot be HBM: Validate() then demands + # hbm.capacity_bytes > 0 and no env var sets it. + "UMBP_DISTRIBUTED_MEDIUM": ("SSD" if kinds == {"ssd"} else "DRAM"), + # UMBP_BACKEND_POLICY is NOT one of the env vars that make the + # server build a distributed config, so on its own it is read and + # then thrown away with the rest of the deployment. Naming the + # node is the selector with the least meaning of its own -- the + # media all come from the policy. + "UMBP_NODE_ID": f"prim-ci-{self.arm}", + "UMBP_SSD_ENABLED": "1" if "ssd" in kinds else "0", + "UMBP_DRAM_CAPACITY": str(self.dram_bytes), + "UMBP_DISTRIBUTED_DRAM_PAGE_SIZE": str(self.page_bytes), + **(self._staging_env() if "ssd" in kinds else {}), + } + + def _start_server(self, master_address: str) -> None: + env = { + "UMBP_ROLE": "standalone", + # Not hugepages: CI has no reservation, and falling back is silent. + "UMBP_DRAM_USE_HUGEPAGES": "0", + # Pin the paged medium's page size instead of taking the 2 MiB + # default, so pool_bytes() computes against the size the server + # actually allocates in. Left implicit, a 64 KiB object pads out + # to a 2 MiB page and the pool runs 32x short. The masterless arm + # reads this name, the master-backed one reads + # UMBP_DISTRIBUTED_DRAM_PAGE_SIZE below, and they must agree or + # the two arms are not measuring the same store. + "UMBP_EMBEDDED_DRAM_PAGE_SIZE": str(self.page_bytes), + # The readiness line below is INFO, and the default level is WARN, + # so without this the wait can only ever time out. + "MORI_UMBP_LOG_LEVEL": "info", + "MORI_GLOBAL_LOG_LEVEL": "info", + } + env.update(self._media_env()) + if master_address: + _require_gpu() + ip = _outbound_ip() + env.update( + { + "UMBP_MASTER_ADDRESS": master_address, + # With a master all four are required: the node is a + # cluster member, so it must state where peers reach it. + "UMBP_NODE_ADDRESS": ip, + "UMBP_NODE_ID": f"{ip}:prim-ci", + "UMBP_IO_ENGINE_HOST": ip, + "UMBP_IO_ENGINE_PORT": str(_free_port()), + "UMBP_PEER_SERVICE_PORT": str(_free_port()), + # Defaults to 0, which the server rejects outright rather + # than delegating to the master's registry default. + "UMBP_DISTRIBUTED_DRAM_PAGE_SIZE": str(self.page_bytes), + } + ) + else: + # Naming any of the selector vars is what switches the server off + # its embedded backend, so a stale one in the environment would + # silently turn this arm into the other one. + for name in ( + "UMBP_MASTER_ADDRESS", + "UMBP_NODE_ADDRESS", + "UMBP_NODE_ID", + "UMBP_IO_ENGINE_HOST", + "UMBP_DISTRIBUTED_MEDIUM", + ): + # setdefault, not assignment: the media config above may have + # set one of these deliberately (the policy path needs a + # selector), and blanking it would drop the whole deployment. + env.setdefault(name, "") + + log = self._spawn("server", [self.server_bin, self.address], env) + # "data plane" is logged at the end of Start(), after both the gRPC + # port and the fd-handoff listener are up. The socket file appears + # well before that, so waiting on the path yields a client that + # connects and is then told the server is not ready. + self._await( + lambda: "data plane" in log.read_text(errors="replace"), + what=f"standalone server on {self.address}", + log=log, + proc_name="server", + ) + + # -- waiting ----------------------------------------------------------- + + @staticmethod + def _can_connect(host: str, port: int) -> bool: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1.0) + return sock.connect_ex((host, port)) == 0 + + def _await(self, ready, what: str, log: Path, proc_name: str) -> None: + deadline = time.time() + self.ready_timeout + while time.time() < deadline: + for name, proc in self.procs: + if name == proc_name and proc.poll() is not None: + raise RuntimeError( + f"{what} exited with {proc.returncode} before becoming " + f"ready:\n{self._tail(log)}" + ) + if ready(): + print(f"[ci] {what} ready", flush=True) + return + time.sleep(0.5) + raise RuntimeError( + f"{what} not ready after {self.ready_timeout:.0f}s:\n{self._tail(log)}" + ) + + @staticmethod + def _tail(log: Path, lines: int = 30) -> str: + try: + return "\n".join(log.read_text(errors="replace").splitlines()[-lines:]) + except OSError as err: + return f"(could not read {log}: {err})" + + +# --------------------------------------------------------------------------- +# Runner +# --------------------------------------------------------------------------- + + +@dataclass +class Pass: + """One bench invocation: which primitives, over what geometry.""" + + name: str + ops: list[str] + keys_per_rank: int + batch: int + passes: int + ranges: int + range_bytes: int + #: Fraction of each `exists` call naming a never-written key. Only the + #: scaling passes vary it; the correctness pass leaves it at 0. + exists_absent_frac: float = 0.0 + #: Key namespace. The scaling passes deliberately SHARE one: they probe + #: the same dataset and differ only in how much of each call is absent, + #: so seeding it once makes them comparable (identical store contents) + #: and stops the pool being charged for the same keys per fraction. + key_prefix: str = "" + + @property + def residency_key(self) -> tuple: + """Passes with the same value occupy the same pages, not more.""" + return ( + self.key_prefix, + tuple(self.ops), + self.keys_per_rank, + self.object_bytes, + self.passes, + ) + + @property + def object_bytes(self) -> int: + return self.ranges * self.range_bytes + + def resident_bytes(self, ranks: list[int], page_bytes: int) -> int: + """What this pass leaves in the pool for the rest of the run. + + Charged in PAGES, not in bytes, and that distinction is the whole + reason this method takes a page size. The DRAM medium is paged and + the allocator hands out whole pages, so a 4 KiB value in a 2 MiB page + occupies 2 MiB. Sizing this by `keys * object_bytes` understates + demand by the padding ratio -- 32x for a 64 KiB object in a 2 MiB + page -- and the pool then runs at a permanent eviction deficit. What + that looks like from the outside is NOT an out-of-space error: the + just-seeded keys are the least recently used, so they are evicted + first, and the run dies in the visibility wait with "seeded keys + never became visible". Measured, not deduced: it cost a whole + false A/B result. + + Every rank count also writes into its own key namespace and nothing + is deleted between them, so this sums over the whole sweep rather + than taking the largest rank count. + """ + seeded = sum(1 for op in SEEDED_OPS if op in self.ops) + written = sum(1 for op in WRITING_OPS if op in self.ops) + generations = seeded + written * self.passes + keys = sum(ranks) * self.keys_per_rank * generations + pages_per_key = max(1, -(-self.object_bytes // page_bytes)) + return keys * pages_per_key * page_bytes + + +def passes_for(args) -> list[Pass]: + """The measurements, in the order they run against one server. + + Split because they want opposite geometries: correctness wants all five + primitives over objects big enough for a transfer to mean something, + scaling wants `exists` alone over a window long enough to measure. + + Scaling is then one pass per absent-key fraction. The fraction decides + how much of the probe the master has to answer -- a resident key is + resolved locally and never reaches it, an absent one always does -- so + each fraction is a different amount of control plane under the same + concurrency sweep. Separate invocations rather than one compound run, + so each stays a single-purpose measurement. + """ + out = [ + Pass( + "correctness", + list(args.ops), + args.keys_per_rank, + args.batch, + args.passes, + args.ranges, + args.range_bytes, + ) + ] + for frac in args.exists_absent_fracs: + out.append( + Pass( + f"scaling@{frac:g}", + ["exists"], + args.scaling_keys_per_rank, + args.scaling_batch, + args.scaling_passes, + 1, + args.scaling_range_bytes, + exists_absent_frac=frac, + key_prefix=args.scaling_key_prefix, + ) + ) + return out + + +def ssd_staging(args) -> tuple[int, int]: + """(slots, bytes) for the SSD staging arena, from the widest call issued.""" + specs = passes_for(args) + widest = max(p.batch for p in specs) + biggest = max(p.object_bytes for p in specs) + slot = max(args.page_bytes, biggest) + # 4x the widest call, so the ranks running concurrently are not each + # waiting on the last one to release before their own batch can reserve. + slots = max(64, widest * 4) + return slots, slots * slot + + +def pool_bytes(args, headroom: float) -> int: + """Size the server pool from the geometry rather than a constant. + + Deduplicated by residency_key: the absent-fraction sweep is several + passes over ONE dataset, so charging the pool once per fraction would + triple a figure the run never occupies. + """ + by_keyspace = { + p.residency_key: p.resident_bytes(args.ranks, args.page_bytes) + for p in passes_for(args) + } + return int(sum(by_keyspace.values()) * headroom) + + +def run_arm(arm: str, args, run_dir: Path, dram: int) -> dict: + """Bring the arm up once and run both passes against it.""" + out: dict[str, dict] = {} + with Deployment( + arm=arm, + run_dir=run_dir, + dram_bytes=dram, + page_bytes=args.page_bytes, + server_bin=_binary("UMBP_STANDALONE_BIN", args.server_bin), + # Resolved only for the arm that runs one: the nomaster arm must not + # be blocked by a build that shipped no master binary. + master_bin=( + _binary("UMBP_MASTER_BIN", args.master_bin) if arm == "master" else "" + ), + ready_timeout=args.ready_timeout, + backends=args.backends, + **dict(zip(("ssd_staging_slots", "ssd_staging_bytes"), ssd_staging(args))), + ) as deploy: + for spec in passes_for(args): + path = run_dir / f"{arm}_{spec.name.replace('@', '_')}.json" + argv = [ + sys.executable, + str(BENCH), + "umbp-server", + deploy.address, + "--ops", + *spec.ops, + "--ranks", + *[str(r) for r in args.ranks], + "--keys-per-rank", + str(spec.keys_per_rank), + "--batch", + str(spec.batch), + "--passes", + str(spec.passes), + "--ranges", + str(spec.ranges), + "--range-bytes", + str(spec.range_bytes), + "--exists-absent-frac", + str(spec.exists_absent_frac), + *(("--key-prefix", spec.key_prefix) if spec.key_prefix else ()), + "--loc", + args.loc, + "--json", + str(path), + ] + print( + f"\n[ci] === arm {arm}, pass {spec.name} ===\n[ci] {' '.join(argv)}", + flush=True, + ) + rc = subprocess.run( + argv, env={**os.environ, "UMBP_STANDALONE_TIMEOUT_MS": "60000"} + ).returncode + if not path.exists(): + raise RuntimeError( + f"arm {arm}: the {spec.name} pass exited {rc} without writing {path}" + ) + payload = json.loads(path.read_text()) + payload["bench_returncode"] = rc + out[spec.name] = payload + return out + + +def parse_args(argv: list[str] | None = None): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument( + "--arms", + nargs="+", + choices=["nomaster", "master"], + default=["nomaster", "master"], + help="standalone server without a master, with one, or both (default: both)", + ) + p.add_argument("--ranks", nargs="+", type=int, default=[1, 2, 4, 8]) + p.add_argument( + "--ops", + nargs="+", + choices=["put", "get", "exists", "put_ranges", "get_ranges"], + default=["put", "get", "exists", "put_ranges", "get_ranges"], + ) + p.add_argument("--keys-per-rank", type=int, default=512) + p.add_argument("--batch", type=int, default=128) + p.add_argument("--passes", type=int, default=3) + p.add_argument("--ranges", type=int, default=4) + p.add_argument("--range-bytes", type=int, default=16384) + p.add_argument("--loc", choices=["host", "gpu"], default="host") + p.add_argument( + "--exists-absent-fracs", + nargs="+", + type=float, + default=list(EXISTS_ABSENT_FRACS), + help="absent-key fractions to sweep the exists probe over. 0 never " + "reaches the master, 1 always does.", + ) + p.add_argument("--scaling-keys-per-rank", type=int, default=SCALING_KEYS_PER_RANK) + p.add_argument("--scaling-batch", type=int, default=SCALING_BATCH) + p.add_argument( + "--scaling-passes", + type=int, + default=SCALING_PASSES, + help="raise this, not the floor, if the scaling number is noisy", + ) + p.add_argument("--scaling-range-bytes", type=int, default=SCALING_RANGE_BYTES) + p.add_argument( + "--page-bytes", + type=int, + default=65536, + help="page size of the server's DRAM medium, on both arms. The pool is " + "charged in whole pages per key, so this and the object size together " + "decide how much pool the sweep needs.", + ) + p.add_argument( + "--pool-headroom", + type=float, + default=3.0, + help="multiple of the run's resident set to give the pool", + ) + p.add_argument( + "--exists-scaling-floor", type=float, default=DEFAULT_EXISTS_SCALING_FLOOR + ) + p.add_argument( + "--media", + default="dram", + help="what the SERVER stores in: dram, hbm, ssd, or a weighted " + "mixture such as 'dram:70,ssd:30'. Weights are relative; omitted ones " + "split the remainder evenly. Note this is the server's storage, not " + "the client buffer's location, which is --loc.", + ) + p.add_argument( + "--ssd-dirs", + default="", + help="comma-separated directories, one per ssd backend in --media " + "(default: a per-run dir under $UMBP_PRIM_CI_SSD_ROOT or /data)", + ) + p.add_argument( + "--hbm-devices", + default="0", + help="comma-separated GPU ordinals for hbm backends", + ) + p.add_argument("--ready-timeout", type=float, default=300.0) + p.add_argument("--out-dir", default="") + p.add_argument("--server-bin", default="", help="override UMBP_STANDALONE_BIN") + p.add_argument("--master-bin", default="", help="override UMBP_MASTER_BIN") + a = p.parse_args(argv) + if len(a.ranks) < 2: + p.error("--ranks needs at least two counts: the gate is a ratio between them") + if "exists" not in a.ops: + p.error("--ops must include exists: it is the gated primitive") + if any(not 0.0 <= f <= 1.0 for f in a.exists_absent_fracs): + p.error("--exists-absent-fracs must all be between 0.0 and 1.0") + if 0.0 not in a.exists_absent_fracs: + p.error("--exists-absent-fracs must include 0.0: it is the gated point") + # Computed once, so every scaling pass in this run names the same keys. + a.scaling_key_prefix = f"s{os.getpid() % 100000:05d}" + + root = os.environ.get("UMBP_PRIM_CI_SSD_ROOT", "/data") + ssd_dirs = ( + [d.strip() for d in a.ssd_dirs.split(",") if d.strip()] + if a.ssd_dirs + else [f"{root}/umbp_prim_ci_{os.getpid()}_{i}" for i in range(4)] + ) + try: + a.backends = parse_media( + a.media, ssd_dirs, [int(d) for d in a.hbm_devices.split(",") if d.strip()] + ) + except ValueError as err: + p.error(f"--media: {err}") + return a + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + run_dir = Path(args.out_dir or tempfile.mkdtemp(prefix="umbp-prim-ci-")) + run_dir.mkdir(parents=True, exist_ok=True) + dram = pool_bytes(args, args.pool_headroom) + media = " + ".join( + f"{b.kind}{'@' + b.path if b.kind == 'ssd' else ''} {b.weight:.0f}%" + for b in args.backends + ) + lines = [ + f"[ci] arms={','.join(args.arms)} ranks={args.ranks} loc={args.loc}", + f"[ci] server media: {media}", + ] + for spec in passes_for(args): + lines.append( + f"[ci] pass {spec.name:<12} ops={','.join(spec.ops)} " + f"object={spec.object_bytes}B keys/rank={spec.keys_per_rank} " + f"batch={spec.batch} passes={spec.passes}" + ) + lines.append( + f"[ci] server pool: {dram / 2**30:.1f} GiB " + f"({args.pool_headroom:g}x both passes' resident set)" + ) + lines.append(f"[ci] logs and results: {run_dir}") + print("\n".join(lines), flush=True) + + results: dict[str, dict] = {} + failed_arms: list[str] = [] + for arm in args.arms: + try: + results[arm] = run_arm(arm, args, run_dir, dram) + except Exception as err: # noqa: BLE001 - one arm must not hide the other + print(f"\n[ci] arm {arm} FAILED TO RUN: {err}", file=sys.stderr, flush=True) + failed_arms.append(arm) + + if results: + print("\n".join(scaling_table(results)), flush=True) + ok = report_findings(gate(results, args.exists_scaling_floor)) + else: + ok = False + + if failed_arms: + print( + f"\n[ci] arms that never produced a result: {', '.join(failed_arms)}", + file=sys.stderr, + ) + verdict = ok and not failed_arms + print(f"\n[ci] {'PASS' if verdict else 'FAIL'}", flush=True) + return 0 if verdict else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/python/umbp/umbp_store_node.py b/tests/python/umbp/umbp_store_node.py new file mode 100644 index 000000000..6bb1facb8 --- /dev/null +++ b/tests/python/umbp/umbp_store_node.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Standalone UMBP storage-node launcher. + +Joins a UMBP distributed pool as a pure storage/capacity node (DRAM/SSD) +WITHOUT starting an SGLang inference engine. It does this by constructing +sglang's own `UMBPStore` (the exact same HiCache L3 backend class an sglang +engine instantiates) with `mem_pool_host=None`, so this process: + + - reads the exact same env vars sglang reads (UMBP_MASTER_ADDRESS, + UMBP_SSD_ENABLED, UMBP_SSD_DIR, UMBP_SSD_CAPACITY, UMBP_NODE_ID, + UMBP_NODE_ADDRESS, UMBP_IO_ENGINE_HOST, UMBP_IO_ENGINE_PORT, + UMBP_PEER_SERVICE_PORT, UMBP_SPDK_*, ...) via UMBPConfig.from_environment(), + - accepts the exact same `--hicache-storage-backend-extra-config` JSON + string / "@file.json|yaml|toml" syntax sglang's server_args.py takes, + - registers with the UMBP master and starts the heartbeat exactly as a + real sglang rank's UMBPStore would. + +There is no model, no scheduler, no GPU work — it just allocates DRAM/SSD +capacity, registers it with the master, and idles. + +Run this in the SAME container/image you use to launch the sglang engine +(needs sglang + mori built with BUILD_UMBP=ON importable). + +Example (pure-SSD storage node, joining an existing UMBP master): + + UMBP_MASTER_ADDRESS=10.0.0.5:50051 \\ + UMBP_SSD_ENABLED=1 \\ + UMBP_SSD_DIR=/mnt/nvme0,/mnt/nvme1 \\ + UMBP_SSD_CAPACITY=1000000000000 \\ + UMBP_NODE_ID=ssd-node-1 \\ + UMBP_NODE_ADDRESS=10.0.0.9 \\ + python umbp_store_node.py + +Or pass the same knobs via extra_config, exactly like an sglang launch +script's --hicache-storage-backend-extra-config: + + python umbp_store_node.py --hicache-storage-backend-extra-config \\ + '{"master_address": "10.0.0.5:50051", "node_id": "ssd-node-1", \\ + "node_address": "10.0.0.9", "ssd_enabled": true, \\ + "ssd_storage_dir": "/mnt/nvme0,/mnt/nvme1", \\ + "ssd_capacity_bytes": 1000000000000}' +""" + +import argparse +import json +import logging +import os +import signal +import socket +import sys +import time +from typing import Optional + +logger = logging.getLogger("umbp_store_node") + +# Keys that sglang's HiRadixCache pops out of extra_config before it ever +# reaches UMBPStore (see hiradix_cache.py::_parse_storage_backend_extra_config). +# Stripped here too so a config blob copy-pasted from an sglang launch script +# doesn't trigger UMBPStore's "unknown extra_config key" warning. +_HIRADIX_ONLY_EXTRA_KEYS = ( + "prefetch_threshold", + "prefetch_timeout_base", + "prefetch_timeout_per_ki_token", + "prefetch_timeout_max", + "hicache_storage_pass_prefix_keys", +) + + +def _parse_extra_config(raw: Optional[str]) -> dict: + """Same parsing as sglang's _parse_storage_backend_extra_config: a JSON + string, or "@path" to a .json/.toml/.yaml file.""" + extra_config = {} + if raw: + if raw.startswith("@"): + path = raw[1:] + ext = os.path.splitext(path)[1].lower() + with open(path, "rb" if ext == ".toml" else "r") as f: + if ext == ".json": + extra_config = json.load(f) + elif ext == ".toml": + import tomllib + + extra_config = tomllib.load(f) + elif ext in (".yaml", ".yml"): + import yaml + + extra_config = yaml.safe_load(f) + else: + raise ValueError(f"Unsupported config file {path} (format: {ext})") + else: + extra_config = json.loads(raw) + for key in _HIRADIX_ONLY_EXTRA_KEYS: + extra_config.pop(key, None) + return extra_config + + +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + description=( + "Join a UMBP distributed pool as a storage-only node " + "(no SGLang engine, no GPU)." + ), + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument( + "--hicache-storage-backend-extra-config", + dest="extra_config_raw", + default=None, + help=( + "Same flag name/format sglang uses: a JSON string, or '@path' to a " + "json/toml/yaml file, with UMBPStore extra_config keys (e.g. " + "master_address, node_id, node_address, ssd_enabled, " + "ssd_storage_dir, ssd_capacity_bytes, ...). Anything not set here " + "falls back to the matching UMBP_* env var, same as sglang." + ), + ) + # Rank/topology fields mirror sglang's HiCacheStorageConfig. A pure + # storage node normally leaves these at their defaults (rank 0 / size 1); + # they only matter if you intentionally run several store processes + # per node and want them to shard SSD dirs / node_ids by rank the same + # way sglang's per-rank UMBPStore instances do. + p.add_argument("--tp-rank", type=int, default=0) + p.add_argument("--tp-size", type=int, default=1) + p.add_argument("--pp-rank", type=int, default=0) + p.add_argument("--pp-size", type=int, default=1) + p.add_argument("--attn-cp-rank", type=int, default=0) + p.add_argument("--attn-cp-size", type=int, default=1) + p.add_argument("--is-mla-model", action="store_true") + p.add_argument("--enable-storage-metrics", action="store_true") + p.add_argument("--model-name", default=None) + p.add_argument( + "--log-level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"] + ) + return p + + +def main() -> None: + args = build_arg_parser().parse_args() + logging.basicConfig( + level=getattr(logging, args.log_level), + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + ) + + try: + from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig + from sglang.srt.mem_cache.storage.umbp.umbp_store import UMBPStore + except ImportError as exc: + logger.error( + "Could not import sglang's UMBPStore (%s). Run this inside the same " + "container/image used to launch the sglang UMBP engine — it needs " + "sglang and mori (built with BUILD_UMBP=ON) importable.", + exc, + ) + sys.exit(1) + + extra_config = _parse_extra_config(args.extra_config_raw) + + storage_config = HiCacheStorageConfig( + tp_rank=args.tp_rank, + tp_size=args.tp_size, + pp_rank=args.pp_rank, + pp_size=args.pp_size, + attn_cp_rank=args.attn_cp_rank, + attn_cp_size=args.attn_cp_size, + is_mla_model=args.is_mla_model, + enable_storage_metrics=args.enable_storage_metrics, + is_page_first_layout=True, + model_name=args.model_name, + extra_config=extra_config, + ) + + logger.info( + "Starting UMBP storage-only node on %s (no SGLang engine, no GPU). " + "extra_config keys=%s", + socket.gethostname(), + sorted(extra_config.keys()), + ) + + store = UMBPStore(storage_config=storage_config, mem_pool_host=None) + + is_distributed = False + try: + is_distributed = bool(store.client.is_distributed()) + except Exception: + pass + if not is_distributed: + logger.warning( + "UMBP client did NOT start in distributed mode — no master_address " + "was resolved from extra_config['master_address'] or the " + "UMBP_MASTER_ADDRESS env var. This process is only a local/standalone " + "pool, not a capacity contributor to a remote master." + ) + else: + logger.info( + "UMBP storage node registered with master and heartbeating. " + "Idling — press Ctrl-C to leave the pool and exit." + ) + + stop = False + + def _handle_signal(signum, _frame): + nonlocal stop + logger.info("Received signal %d, shutting down store node...", signum) + stop = True + + signal.signal(signal.SIGINT, _handle_signal) + signal.signal(signal.SIGTERM, _handle_signal) + + while not stop: + time.sleep(5) + + store.flush() + logger.info("UMBP storage node stopped.") + + +if __name__ == "__main__": + main()