diff --git a/python/mori/umbp/__init__.py b/python/mori/umbp/__init__.py index 0706685c8..5332f4e3e 100644 --- a/python/mori/umbp/__init__.py +++ b/python/mori/umbp/__init__.py @@ -67,14 +67,36 @@ def _configure_packaged_umbp_master() -> None: os.environ["UMBP_MASTER_BIN"] = str(master_path) +def _configure_packaged_umbp_standalone_server() -> None: + if os.environ.get("UMBP_STANDALONE_BIN"): + return + + server_path = Path(__file__).resolve().parents[1] / "umbp_standalone_server" + if not server_path.is_file(): + return + + try: + mode = server_path.stat().st_mode + exec_bits = stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + if (mode & exec_bits) != exec_bits: + server_path.chmod(mode | exec_bits) + except OSError: + pass + + if os.access(server_path, os.X_OK): + os.environ["UMBP_STANDALONE_BIN"] = str(server_path) + + _configure_packaged_spdk_proxy() _configure_packaged_umbp_master() +_configure_packaged_umbp_standalone_server() from mori.cpp import ( CacheRemoteAdmission, UMBPClient, UMBPConfig, UMBPCopyPipelineConfig, + UMBPDeploymentMode, UMBPDistributedConfig, UMBPDramConfig, UMBPDurabilityMode, @@ -84,6 +106,7 @@ def _configure_packaged_umbp_master() -> None: UMBPIoBackend, UMBPIoConfig, UMBPRole, + UMBPStandaloneProcessConfig, UMBPSsdConfig, UMBPEvictionConfig, UMBPDurabilityConfig, @@ -97,6 +120,7 @@ def _configure_packaged_umbp_master() -> None: "UMBPClient", "UMBPConfig", "UMBPCopyPipelineConfig", + "UMBPDeploymentMode", "UMBPDistributedConfig", "UMBPDramConfig", "UMBPDurabilityConfig", @@ -107,6 +131,7 @@ def _configure_packaged_umbp_master() -> None: "UMBPIoBackend", "UMBPIoConfig", "UMBPRole", + "UMBPStandaloneProcessConfig", "UMBPSsdConfig", "UMBPEvictionConfig", "UMBPTierType", diff --git a/setup.py b/setup.py index 928e22762..a9e1cfc64 100644 --- a/setup.py +++ b/setup.py @@ -558,6 +558,14 @@ def build_extension(self, ext: Extension) -> None: elif umbp_master_dst.exists(): umbp_master_dst.unlink() + umbp_standalone_src = build_dir / "src/umbp/umbp_standalone_server" + umbp_standalone_dst = root_dir / "python/mori/umbp_standalone_server" + if umbp_standalone_src.exists(): + shutil.copyfile(umbp_standalone_src, umbp_standalone_dst) + os.chmod(umbp_standalone_dst, 0o700) + elif umbp_standalone_dst.exists(): + umbp_standalone_dst.unlink() + # CCO C++ examples: ship the built binaries when BUILD_EXAMPLES=ON. They # carry an $ORIGIN/../.. rpath (set in examples/CMakeLists.txt) so they # resolve libmori_*.so from site-packages/mori/ once installed here. @@ -723,6 +731,7 @@ def _cco_extension() -> list: "libmori_metrics.so", "libmori_collective.so", # optional: only present when BUILD_COLLECTIVE=ON "umbp_master", + "umbp_standalone_server", "_jit-sources/include/**/*.hpp", "_jit-sources/include/**/*.h", "_jit-sources/include/**/*.cuh", diff --git a/src/pybind/pybind_umbp.cpp b/src/pybind/pybind_umbp.cpp index fa06488bb..a4b28e17d 100644 --- a/src/pybind/pybind_umbp.cpp +++ b/src/pybind/pybind_umbp.cpp @@ -42,6 +42,7 @@ void RegisterMoriUmbp(py::module_& m) { py::enum_(m, "UMBPHostBufferBacking") .value("Anonymous", HostBufferBacking::kAnonymous) .value("AnonymousHugetlb", HostBufferBacking::kAnonymousHugetlb) + .value("AnonymousShm", HostBufferBacking::kAnonymousShm) .export_values(); py::class_(m, "UMBPHostBufferHandle") @@ -115,6 +116,12 @@ void RegisterMoriUmbp(py::module_& m) { .value("SharedSSDFollower", UMBPRole::SharedSSDFollower) .export_values(); + py::enum_(m, "UMBPDeploymentMode") + .value("Local", UMBPDeploymentMode::Local) + .value("StandaloneProcess", UMBPDeploymentMode::StandaloneProcess) + .value("Distributed", UMBPDeploymentMode::Distributed) + .export_values(); + py::enum_(m, "UMBPSsdLayoutMode") .value("SegmentedLog", UMBPSsdLayoutMode::SegmentedLog) .export_values(); @@ -228,6 +235,15 @@ void RegisterMoriUmbp(py::module_& m) { .def_readwrite("admission_max_block_bytes", &UMBPDistributedConfig::admission_max_block_bytes) .def_readwrite("dram_page_size", &UMBPDistributedConfig::dram_page_size); + py::class_(m, "UMBPStandaloneProcessConfig") + .def(py::init<>()) + .def_readwrite("address", &UMBPStandaloneProcessConfig::address) + .def_readwrite("auto_start", &UMBPStandaloneProcessConfig::auto_start) + .def_readwrite("startup_timeout_ms", &UMBPStandaloneProcessConfig::startup_timeout_ms) + .def_readwrite("worker_node_id", &UMBPStandaloneProcessConfig::worker_node_id) + .def_readwrite("worker_node_address", &UMBPStandaloneProcessConfig::worker_node_address) + .def_readwrite("tags", &UMBPStandaloneProcessConfig::tags); + py::class_(m, "UMBPConfig") .def(py::init<>()) .def_static("from_environment", &UMBPConfig::FromEnvironment) @@ -238,7 +254,8 @@ void RegisterMoriUmbp(py::module_& m) { .def_readwrite("role", &UMBPConfig::role) .def_readwrite("follower_mode", &UMBPConfig::follower_mode) .def_readwrite("force_ssd_copy_on_write", &UMBPConfig::force_ssd_copy_on_write) - .def_readwrite("distributed", &UMBPConfig::distributed); + .def_readwrite("distributed", &UMBPConfig::distributed) + .def_readwrite("standalone_process", &UMBPConfig::standalone_process); py::class_>(m, "UMBPClient") .def(py::init([](const UMBPConfig& cfg) { return CreateUMBPClient(cfg); }), @@ -263,7 +280,8 @@ void RegisterMoriUmbp(py::module_& m) { py::call_guard()) .def("clear", &IUMBPClient::Clear, py::call_guard()) .def("flush", &IUMBPClient::Flush, py::call_guard()) - .def("is_distributed", &IUMBPClient::IsDistributed) // pure getter, no I/O + .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("deregister_memory", &IUMBPClient::DeregisterMemory, py::arg("ptr"), diff --git a/src/umbp/CMakeLists.txt b/src/umbp/CMakeLists.txt index f5b355f2a..5b02247aa 100644 --- a/src/umbp/CMakeLists.txt +++ b/src/umbp/CMakeLists.txt @@ -155,6 +155,9 @@ if(HAVE_SPDK) # SPDK Proxy daemon (owns SPDK, serves ranks via SHM IPC) add_executable(spdk_proxy storage/spdk/proxy/spdk_proxy_daemon.cpp) target_link_libraries(spdk_proxy PRIVATE umbp_core) + set_target_properties(spdk_proxy PROPERTIES BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE) endif() # Python bindings are registered in src/pybind/pybind_umbp.cpp (compiled into @@ -260,6 +263,7 @@ message( 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_PROTO_GEN_DIR "${CMAKE_CURRENT_BINARY_DIR}/proto_gen") file(MAKE_DIRECTORY ${UMBP_PROTO_GEN_DIR}) @@ -273,6 +277,11 @@ set(UMBP_PEER_PROTO_HDRS "${UMBP_PROTO_GEN_DIR}/umbp_peer.pb.h") set(UMBP_PEER_GRPC_SRCS "${UMBP_PROTO_GEN_DIR}/umbp_peer.grpc.pb.cc") set(UMBP_PEER_GRPC_HDRS "${UMBP_PROTO_GEN_DIR}/umbp_peer.grpc.pb.h") +set(UMBP_STANDALONE_PROTO_SRCS "${UMBP_PROTO_GEN_DIR}/umbp_standalone.pb.cc") +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") + add_custom_command( OUTPUT ${UMBP_PROTO_SRCS} ${UMBP_PROTO_HDRS} ${UMBP_GRPC_SRCS} ${UMBP_GRPC_HDRS} @@ -299,6 +308,17 @@ add_custom_command( COMMENT "Generating UMBP Peer protobuf/gRPC C++ sources" VERBATIM) +add_custom_command( + OUTPUT ${UMBP_STANDALONE_PROTO_SRCS} ${UMBP_STANDALONE_PROTO_HDRS} + ${UMBP_STANDALONE_GRPC_SRCS} ${UMBP_STANDALONE_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_STANDALONE_PROTO_FILE} + DEPENDS ${UMBP_STANDALONE_PROTO_FILE} ${UMBP_PROTO_FILE} ${UMBP_PROTO_HDRS} + COMMENT "Generating UMBP Standalone protobuf/gRPC C++ sources" + VERBATIM) + # ---------- Static library: umbp_common --------------------------------------- # Shared between PoolClient integration and standalone control-plane tools. @@ -308,6 +328,8 @@ add_library( ${UMBP_GRPC_SRCS} ${UMBP_PEER_PROTO_SRCS} ${UMBP_PEER_GRPC_SRCS} + ${UMBP_STANDALONE_PROTO_SRCS} + ${UMBP_STANDALONE_GRPC_SRCS} distributed/master/global_block_index.cpp distributed/master/external_kv_block_index.cpp distributed/master/client_registry.cpp @@ -325,7 +347,11 @@ add_library( distributed/peer/ssd_copy_pipeline.cpp distributed/peer/peer_service.cpp distributed/pool_client.cpp - distributed/distributed_client.cpp) + distributed/distributed_client.cpp + standalone/external_kv_identity_client.cpp + standalone/ipc.cpp + standalone/standalone_process_client.cpp + standalone/standalone_server.cpp) target_include_directories( umbp_common PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${UMBP_PROTO_GEN_DIR} @@ -373,6 +399,9 @@ if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") endif() target_link_libraries(umbp_master PRIVATE umbp_common ${_PROTOBUF_LIBS} ${_GRPCPP_LIB}) +set_target_properties(umbp_master PROPERTIES BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE) # ---------- Client executable ------------------------------------------------- @@ -382,6 +411,21 @@ if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") endif() target_link_libraries(umbp_client PRIVATE umbp_common ${_PROTOBUF_LIBS} ${_GRPCPP_LIB}) +set_target_properties(umbp_client 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) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(umbp_standalone_server PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries(umbp_standalone_server PRIVATE umbp_common ${_PROTOBUF_LIBS} + ${_GRPCPP_LIB}) +set_target_properties(umbp_standalone_server PROPERTIES BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE) if(BUILD_TESTS) add_subdirectory(tests) diff --git a/src/umbp/distributed/proto/umbp_standalone.proto b/src/umbp/distributed/proto/umbp_standalone.proto new file mode 100644 index 000000000..4cb6580e2 --- /dev/null +++ b/src/umbp/distributed/proto/umbp_standalone.proto @@ -0,0 +1,146 @@ +syntax = "proto3"; +package umbp; + +import "umbp.proto"; + +service UMBPStandalone { + rpc Ping(Empty) returns (PingResponse); + + rpc Put(PutRequest) returns (BoolResponse); + rpc Get(GetRequest) returns (BoolResponse); + rpc BatchPut(BatchDataRequest) returns (BatchBoolResponse); + rpc BatchPutWithDepth(BatchDataWithDepthRequest) returns (BatchBoolResponse); + rpc BatchGet(BatchDataRequest) returns (BatchBoolResponse); + rpc Exists(KeyRequest) returns (BoolResponse); + rpc BatchExists(BatchKeysRequest) returns (BatchBoolResponse); + rpc BatchExistsConsecutive(BatchKeysRequest) returns (CountResponse); + + rpc Clear(Empty) returns (BoolResponse); + rpc Flush(Empty) returns (BoolResponse); + rpc RegisterMemory(RegisterMemoryRequest) returns (BoolResponse); + rpc DeregisterMemory(DeregisterMemoryRequest) returns (Empty); + + rpc ReportExternalKvBlocks(StandaloneExternalKvMutationRequest) returns (BoolResponse); + rpc RevokeExternalKvBlocks(StandaloneExternalKvMutationRequest) returns (BoolResponse); + rpc RevokeAllExternalKvBlocksAtTier(StandaloneExternalKvTierRequest) returns (BoolResponse); + rpc MatchExternalKv(StandaloneMatchExternalKvRequest) returns (StandaloneMatchExternalKvResponse); + rpc GetExternalKvHitCounts(StandaloneExternalKvHitCountsRequest) + returns (StandaloneExternalKvHitCountsResponse); +} + +message Empty {} + +message PingResponse { + bool ready = 1; + StandaloneBackendMode deployment_mode = 2; +} + +message BoolResponse { + bool ok = 1; + string error = 2; +} + +enum StandaloneBackendMode { + STANDALONE_BACKEND_UNKNOWN = 0; + STANDALONE_BACKEND_LOCAL = 1; + STANDALONE_BACKEND_DISTRIBUTED = 2; +} + +message CountResponse { + uint64 count = 1; +} + +message KeyRequest { + string key = 1; +} + +message BatchKeysRequest { + repeated string keys = 1; +} + +message PutRequest { + string key = 1; + string client_id = 2; + uint64 shm_offset = 3; + uint64 size = 4; +} + +message GetRequest { + string key = 1; + string client_id = 2; + uint64 shm_offset = 3; + uint64 size = 4; +} + +message BatchDataRequest { + repeated string keys = 1; + string client_id = 2; + repeated uint64 shm_offsets = 3; + repeated uint64 sizes = 4; +} + +message BatchDataWithDepthRequest { + repeated string keys = 1; + string client_id = 2; + repeated uint64 shm_offsets = 3; + repeated uint64 sizes = 4; + repeated int32 depths = 5; +} + +message BatchBoolResponse { + repeated bool ok = 1; +} + +message RegisterMemoryRequest { + string client_id = 1; + uint64 worker_base = 2; + uint64 size = 3; + string worker_node_id = 4; + string worker_node_address = 5; + repeated string tags = 6; +} + +message DeregisterMemoryRequest { + string client_id = 1; +} + +message StandaloneExternalKvMutationRequest { + repeated string hashes = 1; + TierType tier = 2; + string client_id = 3; +} + +message StandaloneExternalKvTierRequest { + TierType tier = 1; + string client_id = 2; +} + +message StandaloneMatchExternalKvRequest { + repeated string hashes = 1; + bool count_as_hit = 2; + string client_id = 3; +} + +message StandaloneExternalKvMatch { + string node_id = 1; + string peer_address = 2; + repeated TierHashes hashes_by_tier = 3; +} + +message StandaloneMatchExternalKvResponse { + repeated StandaloneExternalKvMatch matches = 1; +} + +message StandaloneExternalKvHitCountsRequest { + repeated string hashes = 1; + string client_id = 2; +} + +message StandaloneExternalKvHitCountEntry { + string hash = 1; + uint64 hit_count_total = 2; +} + +message StandaloneExternalKvHitCountsResponse { + repeated StandaloneExternalKvHitCountEntry entries = 1; +} diff --git a/src/umbp/doc/design-standalone-process-mode.md b/src/umbp/doc/design-standalone-process-mode.md new file mode 100644 index 000000000..f9d03a146 --- /dev/null +++ b/src/umbp/doc/design-standalone-process-mode.md @@ -0,0 +1,1047 @@ +# UMBP Standalone Process Mode — Design + +## 1. Background + +UMBP (mori's tiered KV-cache system) currently supports two deployment +shapes: + +- **`local/`** — an in-process `StandaloneClient` that owns a + `LocalStorageManager` (`DRAMTier` + `SSDTier`), a `LocalBlockIndex`, + and a `CopyPipeline`, all living inside the calling process. +- **`distributed/`** — a cross-node deployment (`umbp_master` + peers + over gRPC + RDMA) for sharing a KV cache across machines, with + master-mediated routing and lease/eviction protocols for remote + peers. + +Neither shape addresses a common single-host deployment pattern: an +8-GPU server running SGLang with tensor-parallel degree 8 spawns 8 +worker processes on one host, and today each process runs its own +`local/`-mode UMBP instance. This duplicates the DRAM tier 8x and gives +sibling processes (TP ranks, or a restarted worker resuming a warm +cache) no way to share a cache that physically fits on the same +machine. + +Mooncake solves the same problem with a **Real Client / DummyClient** +split: a long-lived Real Client process owns the actual distributed +storage client, and lightweight DummyClient shims inside each +application process forward calls to it over IPC. Mooncake is the +primary architectural reference for this design; LMCache's +`lmcache/v1/multiprocess/` was reviewed for comparison but not used as +a template (its zero-copy path is CUDA-IPC/GPU-only, with no +host-memory equivalent). + +A second, initially underappreciated fact about Mooncake's design +matters here: its "standalone" Real Client is not a same-host-only +cache-sharing shim — it is a **full distributed client** (Master + RDMA +transfer engine) that happens to run in its own OS process. A cache +miss on one host can still be served from a remote node over RDMA. +This makes Mooncake's standalone mode cross-node-capable by +construction, not just same-host-capable. + +## 2. Goals + +1. **Same-host sharing**: let N SGLang/vLLM worker processes on one + host share a single DRAM/SSD cache tier through a dedicated, + long-lived OS process, without each process duplicating its own + tier. +2. **Mooncake parity**: match Mooncake's Real Client capability, + including cross-node sharing — a standalone-process deployment + should be able to serve a cache miss from a remote node over RDMA, + not only from local storage. +3. **Preserve existing UMBP capabilities** that are not present in + Mooncake but are part of UMBP's baseline `distributed/` mode, + specifically per-worker external-KV routing precision (used by a + custom SGLang router to make placement decisions). +4. **No changes to `distributed/`'s cross-node machinery** — master, + RDMA registration, routing, and eviction are reused unmodified. + +### Non-goals (v1) + +- Per-tenant DRAM quota enforcement across workers sharing one server. +- A `tcp_staging` (non-UDS, inline-payload) transport — the zero-copy + shared-memory data plane requires `AF_UNIX`/`SCM_RIGHTS`; no + drop-in TCP equivalent exists. This is deferred until a concrete + deployment needs it. +- Worker crash/disconnect detection and cache rehydration after a + server restart. Tracked as a known limitation (§8). + +## 3. Architecture Overview + +Two new roles, named after Mooncake's split for consistency: + +| UMBP name | Mooncake analog | Lives in | Role | +|---|---|---|---| +| `umbp_standalone_server` | Real Client | its own process | Owns the storage/transport backend (see §3.1) and exposes `IUMBPClient` semantics over gRPC. `mmap`s each worker's registered host-buffer segment (via fd handoff) for zero-copy Put/Get access. | +| `StandaloneProcessClient` | DummyClient | inside the SGLang/vLLM worker process, behind `IUMBPClient` | An `IUMBPClient` implementation. Forwards every call to the server over gRPC. Registers the worker's host KV buffer with the server once via fd-passing, then references it by offset on every hot-path call. | + +``` + ┌────────────────────────────┐ ┌────────────────────────────┐ + │ SGLang worker process A │ │ SGLang worker process B │ + │ │ │ │ + │ UMBPStore (Python) │ │ UMBPStore (Python) │ + │ │ pybind11 │ │ │ pybind11 │ + │ ▼ │ │ ▼ │ + │ StandaloneProcessClient │ │ StandaloneProcessClient │ + │ (IUMBPClient impl) │ │ (IUMBPClient impl) │ + │ │ gRPC │ raw UDS │ │ │ gRPC │ raw UDS │ + │ │ (.grpc.sock)│(.fd.sock)│ │ │ (.grpc.sock)│(.fd.sock)│ + └────┼─────────────┼─────────┘ └────┼─────────────┼─────────┘ + │ │ │ │ + │ │ one-time fd handoff │ │ + │ │ (SCM_RIGHTS), at │ │ + │ │ RegisterMemory() time │ │ + ▼ ▼ ▼ ▼ + ┌─────────────────────────────────────────────────────────────────┐ + │ umbp_standalone_server │ + │ │ + │ UMBPStandaloneService (gRPC, control plane, own UDS socket) │ + │ Put/Get/BatchPut/BatchGet/Exists/... — same names as │ + │ IUMBPClient, arguments are (key, client_id, shm_offset, size) │ + │ instead of raw pointers │ + │ │ + │ 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) │ + │ │ + │ Per-worker shm registry: client_id → {fd, base, size, offset} │ + └─────────────────────────────────────────────────────────────────┘ +``` + +Two logical channels, two socket paths, deliberately: gRPC cannot +carry a file descriptor, so the one-time fd handoff (§4.1) must use a +separate raw `AF_UNIX` socket alongside the gRPC control-plane UDS. + +### 3.1 Deployment shapes + +`umbp_standalone_server`'s `client_` member is `IUMBPClient`-typed and +constructed through the existing `CreateUMBPClient(config)` factory +(`umbp_client_factory.cpp`). This yields three deployment shapes: + +| 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. + +Existing data-path RPC handlers (Put/Get/BatchPut/BatchGet/Exists/ +Clear/Flush) require no logic changes across these shapes: they +already resolve `(client_id, shm_offset)` to a plain pointer before +calling into `client_` (§4.2), so `client_`'s concrete type was never +load-bearing for them. + +## 4. IPC Mechanism + +### 4.1 Control plane: gRPC over a Unix domain socket + +UMBP reuses its existing gRPC stack for the control plane rather than +adopting a new RPC framework or Mooncake's `coro_rpc`: `umbp_common` +already depends on gRPC/protobuf for the `distributed/` master/peer +path, and `pybind_umbp.cpp` already has a working convention for +calling blocking C++ RPC methods from Python with the GIL released. +Adopting `coro_rpc` would add a new coroutine-to-pybind11 bridge with +no precedent in this codebase, for no benefit gRPC doesn't already +provide (typed request/response, an async server loop, UDS transport). + +Design: + +- A new proto service `UMBPStandalone` + (`distributed/proto/umbp_standalone.proto`) with: + - RPCs mirroring `IUMBPClient` 1:1: `Put`, `Get`, `BatchPut`, + `BatchPutWithDepth`, `BatchGet`, `Exists`, `BatchExists`, + `BatchExistsConsecutive`, `Clear`, `Flush`, `RegisterMemory`, + `DeregisterMemory`, and the full external-KV set — + `ReportExternalKvBlocks`, `RevokeExternalKvBlocks`, + `RevokeAllExternalKvBlocksAtTier`, `MatchExternalKv`, + `GetExternalKvHitCounts`. + - `Ping(Empty) -> PingResponse{ready: bool, deployment_mode: LOCAL|DISTRIBUTED}`: + a new RPC with no `IUMBPClient` analog, used for readiness probing + and deployment-mode signaling (§6.4). + - Methods that stay entirely client-local in `StandaloneProcessClient` + (no RPC): `Close()` (tears down the local gRPC channel + fd-handoff + socket + registry state) and `get_deployment_mode()` (a pure local + constant). + - `BatchPut`/`BatchGet` reuse the existing struct-of-arrays batch + codec (`distributed/peer/batch_resolve_codec.h`) so per-key + protobuf submessage overhead doesn't reappear here — the same + pattern already validated for the peer `BatchResolveKeys` RPC. +- **Transport: gRPC over a Unix domain socket** + (`unix:///run/umbp/standalone/.grpc.sock`), not TCP. This + is a deliberate deviation from Mooncake (TCP) and LMCache (TCP): + standalone mode never leaves the host, so a UDS avoids port + allocation, avoids the network stack, and gets filesystem-permission + based access control for free. TCP is not a drop-in fallback for + this design: the data plane (§4.2) depends on `SCM_RIGHTS` fd + passing, an `AF_UNIX`-only mechanism with no TCP equivalent. A + deployment that cannot share a UDS-reachable filesystem between + processes cannot use the zero-copy data plane at all; the only + fallback would be a non-zero-copy path where `Put`/`Get` payload + bytes travel inside the gRPC message itself. This mode + (`tcp_staging`) is not built in v1 (§8). +- **Argument shape**: control messages never carry KV bytes. `Put`/ + `Get` take `(key, client_id, shm_offset, size)` — the raw pointer + that `IUMBPClient::Put(key, uintptr_t src, size_t size)` takes is + translated to `(client_id, shm_offset)` inside + `StandaloneProcessClient` before the RPC is issued. This mirrors + Mooncake's `map_dummy_buffer_range_to_real`/`shm_addr_offset` + translation. + +### 4.2 Data plane: shared memory + +UMBP and Mooncake independently converged on the same shape for the +data plane: control messages carry a handle, KV bytes live in shared +memory (UMBP's own SPDK-proxy daemon uses the same idea for its ring +protocol). The design point that needs to be specified precisely is +**who allocates the shared segment**, and the answer is taken directly +from Mooncake: + +1. **The worker process (not the server) owns the buffer**, exactly + like Mooncake's `DummyClient` owns its local buffer via + `setup_dummy()` → `ShmHelper::allocate()`. Concretely, + `HostMemAllocator` (already used by `umbp_host_allocator.py` to + back SGLang's host KV pool) gains a new backing kind, + `kAnonymousShm`, using `memfd_create` + `ftruncate` + + `mmap(MAP_SHARED)` — the same primitives Mooncake's + `ShmHelper::allocate()` uses. `memfd_create` (anonymous, refcounted + purely by open fds) is chosen over UMBP's existing named-`shm_open` + `DRAMTier` path specifically to avoid `/dev/shm` name collisions + across many SGLang instances on one host, and to avoid leaking a + named segment into `/dev/shm` if a process is killed uncleanly. +2. **One-time registration handshake**, triggered by the existing + `RegisterMemory(ptr, size)` call. `IUMBPClient::RegisterMemory` + only carries a pointer and a size; the file descriptor is recovered + entirely on the C++ side of the worker process: + - `HostMemAllocator::Alloc` maintains a process-local, + mutex-guarded `ptr → {fd, size}` registry, populated only for + `kAnonymousShm` allocations at alloc time and erased on `Free`. + - `StandaloneProcessClient::RegisterMemory(ptr, size)` looks up + `ptr` in that registry (a range/floor lookup, since `ptr` may be + the tensor's base address rather than the allocation's exact + base) and, on a hit, sends the recovered fd to the server via + `SCM_RIGHTS` over the **separate raw-UDS fd-handoff socket** + (`.fd.sock`, distinct from the gRPC UDS) — mirroring + Mooncake's `UdsConnection::sendFd`/`recvFd` and + `RealClient::handle_ipc_shm_register`. On a registry miss (the + pointer isn't backed by a `kAnonymousShm` allocation), + `RegisterMemory` fails loudly rather than silently falling back + to a broken pointer-based `Put`/`Get`. + - The server `mmap`s the received fd read-write into its own + address space and stores `client_id → {base, size}` in a registry + keyed by `client_id` — mirroring Mooncake's `shm_contexts_`. +3. **Every subsequent Put/Get/BatchPut/BatchGet RPC carries only + `(client_id, shm_offset, size)`** — `StandaloneProcessClient` + computes `shm_offset = ptr - registered_base` locally before + issuing the RPC; the server computes `real_addr = base + + shm_offset` and reads/writes there directly. No KV bytes ever cross + the gRPC channel. +4. **fd ownership**: the allocator's ptr→fd registry, not + `RegisterMemory`/`DeregisterMemory`, owns the fd — `HostMemAllocator` + opens it at `Alloc` time and is the only thing that closes it, at + `Free` time. `RegisterMemory` only *borrows* the fd to send over + `SCM_RIGHTS` (which duplicates it into the receiving process, so + the sender's copy is unaffected by anything the server does with + its own copy). `DeregisterMemory` therefore only (a) tells the + server to `munmap` its copy and drop the `client_id` registry + entry, and (b) clears `StandaloneProcessClient`'s own bookkeeping — + it must **not** close the allocator-owned fd. Closing happens + exactly once, in `HostMemAllocator::Free`. + +The server's own DRAM cache tier (local-backend shape) does not need +to be visible to any other process — no worker ever `mmap`s it +directly — so it defaults to the same private +`kAnonymous`/`kAnonymousHugetlb` backing `StandaloneClient` already +uses today, not a named `use_shared_memory=true` POSIX shm segment. + +Net effect: the hot path becomes the same shape as UMBP's existing +RDMA-registered `distributed/` path (register once, then +reference-by-handle on every call) — not a new interface concept for +UMBP, just a new transport underneath a pattern `IUMBPClient` already +exposes. + +**Rejected alternative — a busy-polled ring buffer** (the SPDK-proxy +daemon's own approach): well suited to SPDK-proxy's low-cardinality, +microsecond-latency block I/O, but a poor fit for a control plane that +needs to express variable-shape requests (batch registration, +variable-length external-KV hash lists) and async multiplexing for +slow operations. Neither Mooncake nor LMCache use a pure ring buffer +for control; both back it with a real RPC/messaging library and +reserve raw shared memory for bulk bytes, which this design also does. +The data plane still uses the same "handle + offset/size, bytes in +shm" idea as SPDK-proxy — only the ring buffer itself is rejected, for +the control plane specifically. + +## 5. Lifecycle Management + +The standalone server holds the DRAM cache tier's actual data, so +unlike UMBP's own SPDK-proxy precedent (a stateless passthrough to a +shared NVMe device, safe to self-exit and respawn), it defaults to +Mooncake/LMCache's model: **externally launched, long-lived, no +self-exit**, with idle-self-exit available as an explicit opt-in. + +### 5.1 Primary path: externally launched + +- Launched the same way `run_umbp_single_node_hicache.sh` already + launches `umbp_master` (background process, log redirected, `trap + cleanup EXIT`) — no new orchestration concept, just a new binary + target. +- Discovery: `UMBP_STANDALONE_ADDRESS` (default + `unix:///run/umbp/standalone/.grpc.sock`), following the + naming convention of `UMBP_MASTER_ADDRESS`. +- The fd-handoff socket path is always **mechanically derived** from + `cfg.standalone_process.address` — by both processes independently: + strip the `unix://` scheme prefix, then replace a trailing + `.grpc.sock` with `.fd.sock` (or append `.fd.sock` for a custom + path that doesn't end in `.grpc.sock`). This derivation lives in one + shared place, computed identically on both sides, and deliberately + has no independent config field or env var, so it cannot drift out + of sync between the two processes. +- Readiness probe: connect to the UDS and call `Ping`. + +### 5.2 Convenience path: auto-start (opt-in, local-backend only) + +When `cfg.standalone_process.auto_start` is `true`: the first worker +process on a host to reach `CreateUMBPClient` (leader election reuses +the existing `UMBP_ROLE`/`LOCAL_RANK`/`OMPI_COMM_WORLD_LOCAL_RANK` +rank-0 logic in `UMBPConfig::FromEnvironment` — no new election +protocol) probes the UDS; if absent, it forks+execs the server binary +using the same `fork()` + `setsid()` + `execlp()` sequence +`LocalStorageManager::SpawnProxyDaemon` already uses, and waits for +readiness bounded by `startup_timeout_ms`. A bootstrap lock (the +existing `ScopedBootstrapLock` pattern) prevents a thundering herd of +workers all trying to spawn the server simultaneously. No CUDA/HIP/ +gRPC call may happen between `fork()` and `execlp()`, matching the +fork-safety discipline `SpawnProxyDaemon` already requires. + +Auto-start is opt-in, not default, because a worker's `fork()`'d child +inherits CUDA/HIP context and GPU device state at the moment of fork. + +**A distributed-backend server does not support auto-start — external +launch only.** This matches Mooncake, where the application side +(`DummyClient`) never forks/execs the Real Client under any +circumstance; the Real Client is always started by something outside +the application. It also avoids a real conflict: the env vars needed +to configure a distributed backend (`UMBP_MASTER_ADDRESS`, etc.) would +most naturally live in the worker's own environment if the worker +forked the child, but a worker's own config parsing treats +`UMBP_STANDALONE_ADDRESS` and `UMBP_MASTER_ADDRESS` both being set as +a hard error (mutual exclusivity between standalone-process and +distributed worker modes). Local-backend auto-start is unaffected by +this restriction. + +### 5.3 Shutdown + +`SIGTERM`/`SIGINT` triggers, in order: + +1. Drain in-flight RPCs (bounded deadline, mirroring + `UMBP_GRPC_SHUTDOWN_DEADLINE_SEC`/`MasterServer::Shutdown()`). +2. **(distributed-backend only)** Unregister all live + `ExternalKvIdentityClient` instances (§7) — before the shared + backend is closed, so the Master's `ClientRegistry` doesn't briefly + carry stale `ALIVE` entries for a process that is already exiting. +3. **(local-backend)** Flush `CopyPipeline` (best-effort persist + DRAM-tier dirty pages to SSD, if the SSD tier is enabled and + `force_ssd_copy_on_write` is set) — `CopyPipeline::Drain(timeout)` + blocks until the async copy queue is empty or the timeout elapses; + both the shutdown path and the `Flush` RPC call it. + **(distributed-backend)** Close the `DistributedClient` backend + (`client_->Close()`). +4. `munmap` all registered client shm segments. +5. Exit. + +Steps 2 and 3/4 for a given registered client must run in the order +"deregister from `client_` → `munmap`" (not the reverse), and the +whole backend-deregistration sequence must complete before +`client_->Close()` — deregistering after `Close()` would call into a +backend that has already torn down its RDMA/IO engine state. + +Client-side: if the RPC channel drops (server crashed or was killed), +`StandaloneProcessClient` must not silently return stale success — +every in-flight call fails, and the worker-side behavior is to treat +this as "cache miss / cache unavailable" (SGLang's `HiCacheStorage` +abstraction already tolerates backend failures returning +`False`/`None`), not to crash the inference process. There is +currently no supervised-restart or reconnect path (§8, item 1). + +### 5.4 Multi-tenant workers on one host + +Multiple SGLang worker processes (e.g. one per TP/DP rank) attach to +the same standalone server. Each gets a `client_id` (UUID, assigned at +first `RegisterMemory` call) and its own shm registration entry, +mirroring Mooncake's per-client `shm_contexts_` map. Capacity +accounting inside `LocalStorageManager`/`DRAMTier` is global, not +per-caller — per-tenant DRAM quota is not implemented (§8, item 2). + +### 5.5 Duplicate writes from sibling workers + +`local/` mode provides a Shared-SSD Leader/Follower model +(`UMBPRole::SharedSSDLeader`/`SharedSSDFollower`, `common/config.h`) +for the case where N sibling worker processes each own their own +`LocalStorageManager` but point at one shared on-disk SSD directory: +only the leader writes to SSD, followers open it read-only +(`standalone_client.cpp`, `local_storage_manager.cpp`, +`ssd_tier.cpp`), so sibling ranks do not issue duplicate or concurrent +writes to the shared segmented-log. Role assignment uses rank-0-elects- +leader logic in `UMBPConfig::FromEnvironment()`. + +Standalone-process mode addresses the same concern by construction +rather than by role coordination, because its topology is different: +there is a single backend instance (one `client_`, one +`LocalStorageManager` or one `DistributedClient`) shared by all workers +over IPC, and all data-path RPCs are serialized on the server's +`client_mu_` (`standalone_server.cpp`). There is therefore a single +writer to the backing storage, so the multi-writer-to-shared-storage +coordination the Leader/Follower model exists for does not arise. The +server accordingly constructs its backend with `role = +UMBPRole::Standalone` (`follower_mode = false`, +`standalone_server_main.cpp`); the Leader/Follower code path is not +exercised in this mode. + +Duplicate writes of the same key from different workers are collapsed +by the backend's existing content-addressed index: +`StandaloneClient::Put`/`BatchPut` return early for a key already +present (`index_.MayExist`, `standalone_client.cpp`), so a repeat write +of an already-stored key is a no-op rather than a second copy into the +tier. This is first-writer-wins with no content comparison: if two +workers write different bytes under the same key, writes after the +first are dropped (reported successful, not stored). Whether sibling +workers actually issue same-key writes is determined by the caller +(SGLang), which is outside this design's scope; this section only +describes the server's behavior when they do. + +## 6. Cross-Node Extension: `DistributedClient`-backed Server + +### 6.1 Two independent `UMBPConfig` instances + +There are two separate configuration surfaces, and they stay separate: + +1. **Worker-facing** (`standalone_process` field): what a + `StandaloneProcessClient` inside an SGLang worker uses to find and + talk to the server over UDS. A worker never knows or cares whether + the server it's talking to is itself distributed-backed. +2. **Server's own internal config** (`distributed` field): what + `umbp_standalone_server` itself uses to decide whether its + `client_` backend should be a `StandaloneClient` or a + `DistributedClient`. This is a second, independent `UMBPConfig` + instance, built by `standalone_server_main.cpp` from the server's + own process environment (§6.6) — never the same object as #1. + +`UMBPConfig::Validate()` rejects `distributed.has_value() && +standalone_process.has_value()` *within one config instance* — it does +not constrain two different `UMBPConfig` objects existing in the same +process, so this split requires no change to that check. + +``` + Worker process umbp_standalone_server process + ┌────────────────────────┐ ┌───────────────────────────────────┐ + │ UMBPConfig #1 │ UDS + │ UMBPConfig #1 (mirror of the │ + │ standalone_process = │─SCM_──▶│ worker's, address/fd-socket) │ + │ {address} │ RIGHTS │ -> used only to open sockets │ + │ distributed = null │ │ │ + └────────────────────────┘ │ UMBPConfig #2 (server's own) │ + │ distributed = {master_address, │ + │ node_id, node_address, │ + │ io_engine, ...} (if enabled) │ + │ standalone_process = null │ + │ │ │ + │ ▼ │ + │ client_ = CreateUMBPClient(#2) │ + │ -> DistributedClient │ + │ (Master + RDMA) │ + │ -- or, if #2.distributed │ + │ is unset -- │ + │ -> StandaloneClient │ + │ (local-backend, unchanged) │ + └───────────────────────────────────┘ + │ RDMA (if DistributedClient) + ▼ + ┌───────────────────────────────────┐ + │ umbp_master + other nodes' │ + │ peers (existing distributed/ │ + │ cluster, entirely unmodified) │ + └───────────────────────────────────┘ +``` + +### 6.2 Zero-copy RDMA registration of a shared mapping + +RDMA memory registration (`DistributedClient::RegisterMemory` → +`PoolClient::RegisterMemory`) operates purely on `(ptr, size)` in the +calling process's own virtual address space (`io_engine_->RegisterMemory +(ptr, size, ...)`), caching `{base, size, mem_desc}` keyed by that +local VA. Nothing checks who allocated the memory or whether another +process also maps the same physical pages. + +This means `umbp_standalone_server`, after `mmap`-ing the worker's +`memfd`-backed region via the fd-handoff path (§4.2), can call +`client_->RegisterMemory(server_local_va, size)` on its own local VA +for that mapping. Since the server's VA and the worker's VA back the +same physical pages (`MAP_SHARED` on the same `memfd`), registering +the server's VA for RDMA registers those physical pages for RDMA, +regardless of which process's VA was used to register them. + +`Get()` zero-copy is conditional on this registration: the RDMA read +path (`PoolClient`'s remote-get handling) only skips the memcpy-out +when the destination pointer was pre-registered and found in the +registered-memory index — in that branch, the RDMA read is posted +directly into the registered memory region with no memcpy afterward. +Registering the server's local mapping once, at fd-handoff time, +therefore gives every subsequent `Get` targeting that region true +zero-copy RDMA: the remote write lands in memory the worker can +already see, with no copy on the server side and no RDMA involvement +on the worker side at all — the worker process never touches RDMA/IB +verbs directly, only the server does. + +Mooncake's Real Client does the identical dual-mapping pattern in +production: the application (`DummyClient`) allocates the shared +region and performs the first `mmap` via `ShmHelper::allocate()`, then +sends the fd to the Real Client, which performs a second, independent +`mmap` of the same physical pages (`RealClient::map_shm_internal_with_device`) +and registers that server-local mapping with its own Transfer Engine. +This is the same allocation/registration direction +`umbp_standalone_server` already implements (worker allocates and owns +the buffer; server receives the fd and independently mmaps it) — a +shipped, production pattern, not a novel one this design introduces. + +**Memory-visibility contract**: the gRPC `Get` response is the +happens-before edge. `umbp_standalone_server` sends the gRPC response +for a `Get` only after the RDMA read completion is observed +server-side; the worker must not read the target offset until it has +received that response. + +### 6.3 Backend registration lifecycle + +`RegisterMemory`/`DeregisterMemory` currently only perform fd/mmap +bookkeeping and never touch `client_`. They are extended to also +register/deregister the mapped region with the backend: + +- `client_->RegisterMemory(server_local_ptr, size)` is safe to call + unconditionally regardless of backend — `IUMBPClient::RegisterMemory`'s + default (used by `StandaloneClient`, which never overrides it) is + already a no-op returning `true` ("standalone mode needs no + registration — CPU-local memcpy"). +- **Registration**: the call is placed inside the existing fd + registration path, immediately after `mmap` succeeds and before the + registry entry is written. If it fails, the just-created mapping is + `munmap`'d and the call returns failure without writing the registry + entry — the same single-phase-commit shape the registration path + already has for an `mmap` failure. +- **Deregistration**: `client_->DeregisterMemory(server_local_ptr)` is + called before `munmap`-ing, so the backend's view of the memory is + torn down before the memory itself disappears. +- **Shutdown**: every registered mapping is deregistered from `client_` + before being `munmap`'d, and this whole sequence runs before + `client_->Close()` (§5.3). + +Out of scope, tracked as a pre-existing characteristic orthogonal to +backend choice: if a worker sends its fd via the fd-handoff socket and +then crashes before issuing the confirming gRPC `RegisterMemory` call, +the mapping stays live indefinitely — there is no timeout/reaper for +this today (§8, item 9). + +`Flush()`/`Clear()` have backend-dependent meaning, by design: under +`StandaloneClient` they drain `CopyPipeline` to SSD; under +`DistributedClient` they operate on the pool client's own state +(heartbeat/registration bookkeeping — there is no local SSD tier in +this backend, §6.5). The server passes both calls through unchanged; +this is the correct, existing behavior of each `IUMBPClient` +implementation, not something this design needs to reconcile. + +### 6.4 Deployment-mode signaling + +`PingResponse` includes a `deployment_mode` field (`LOCAL`/ +`DISTRIBUTED`), populated from `client_->GetDeploymentMode()`. A +single `UMBP_STANDALONE_ADDRESS` can resolve to either a local-backed +or a distributed-backed server, and the failure mode this guards +against is silent: a distributed-backend deployment whose server +ends up local-backed anyway (e.g. an operator's launch script forgot +`UMBP_MASTER_ADDRESS`) would otherwise have every RPC succeed while +the feature's entire purpose — cross-node capability — silently does +not exist. Launch/health-check/integration-test tooling for the +distributed-backend shape must assert `deployment_mode == DISTRIBUTED` +after startup; a distributed-backend server that comes up as `LOCAL` +is a failed deployment, not a degraded-but-working one. + +### 6.5 The server's own storage, under each backend + +When `client_` is a `DistributedClient`, the server's +`LocalStorageManager`/`DRAMTier`/`SSDTier` concept does not apply at +all — `DistributedClient` doesn't use `LocalStorageManager`; it has +its own separately-owned DRAM pool (built via `HostMemAllocator` +inside `PoolClient`/`DistributedClient`'s own construction), used as +its local slice of the distributed pool, not as a passthrough cache +for registered worker buffers. The registered worker buffers (§6.3) +exist purely so RDMA can target them directly; `DistributedClient` +does not manage them as cache storage. See the deployment-shapes table +in §3.1. + +### 6.6 Server-side distributed configuration + +`standalone_server_main.cpp` builds the second, independent +`UMBPConfig` (§6.1) for the backend from its own process environment — +never from the worker's environment. A `BuildDistributedBackendConfigFromEnv()` +helper mirrors `umbp_store.py`'s env-var-driven construction of +`UMBPDistributedConfig`: + +- **Required** (absence means "distributed backend not requested" — + falls back to local-backend; some-but-not-all set is a + misconfiguration and hard-fails): `UMBP_MASTER_ADDRESS`, + `UMBP_NODE_ADDRESS`, `UMBP_NODE_ID`, `UMBP_IO_ENGINE_HOST`. +- **Optional**, defaulting the same way `UMBPDistributedConfig` + already does when unset: `UMBP_IO_ENGINE_PORT`, + `UMBP_PEER_SERVICE_PORT` (reusing the worker-side env-var names, + since those refer to the same concept in both places), and five + server-backend-only names with a `UMBP_DISTRIBUTED_` prefix (chosen + because these five have no existing worker-side env-var form to + collide with — worker-side they are `extra_config`-only): + `UMBP_DISTRIBUTED_STAGING_BUFFER_SIZE`, + `UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SIZE`, + `UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SLOTS`, + `UMBP_DISTRIBUTED_CACHE_REMOTE_FETCHES`, + `UMBP_DISTRIBUTED_DRAM_PAGE_SIZE`. +- `dram_page_size` is normally auto-derived by probing a worker's own + `mem_pool_host`; this does not apply server-side, since the server's + backend isn't attached to any one worker's tensor layout and may + serve workers with different layouts. The server-side builder only + supports the explicit-override env var above, defaulting to `0` + (delegating to the master's own default) when unset. +- `auto_heartbeat` has no dedicated env var; the server-side helper + keeps the `UMBPMasterClientConfig` default (`true`) unconditionally. + +If `UMBP_MASTER_ADDRESS` is set but the IO engine cannot actually be +stood up (e.g. `UMBP_IO_ENGINE_HOST` empty), **server startup hard-fails** +rather than silently falling back to a `StandaloneClient` backend — a +misconfigured distributed-backend server must refuse to start, not +quietly become a smaller feature. + +Distributed-backend supports external launch only (§5.2) — no +auto-start for any deployment shape. + +## 7. External-KV Per-Worker Identity (`ExternalKvIdentityClient`) + +### 7.1 Layering + +This is a UMBP-compatibility extension, not part of Mooncake-parity +core. The parity core (§6.1–§6.4) alone is what makes this design +match Mooncake's Real Client shape; `ExternalKvIdentityClient` exists +only to avoid regressing a capability the `distributed/` baseline +already has (per-worker external-KV routing precision). It has no +Mooncake analog, is not required for declaring Mooncake parity +achieved, and can be scheduled, implemented, and tested as a separable +unit of work — a distributed-backend server without it is a legitimate +intermediate state (external-KV simply stays unavailable, exactly as +the `StandaloneClient`-backed case already handles it). + +### 7.2 What external-KV is + +A pure query/registry service, not a data-transfer mechanism. +`Report`/`RevokeExternalKvBlocks` record advisory "node N holds hash H +at tier T" facts in the Master's `GlobalBlockIndex`; `MatchExternalKv` +answers "which node(s) hold these hashes" for a caller (typically a +custom SGLang router) making a routing decision — the caller sends a +new request to whichever node the match points at, and that node's own +already-resident local cache serves it. UMBP itself never moves the +reported bytes for this feature, in `distributed/` mode or here. + +### 7.3 Identity granularity + +What matters is whether the reported/matched identity (`node_id`) +reflects the actual deployment topology at the granularity a caller +needs. The `distributed/` baseline's `node_id` is built from +per-process rank coordinates +(`f"{node_address}:dp{dp_rank}:pp{pp_rank}:tp{local_rank}"`) with no +awareness of logical TP/DP group boundaries — this means the baseline +already supports, for free, multiple independent replicas sharing one +physical host (e.g. two TP=4 groups on one 8-GPU box), since each +process gets its own distinct `node_id` by rank coordinate regardless +of group membership. For this project's stated deployment shape +(8-GPU hosts running one SGLang TP=8 replica per host) a single shared +per-host identity would be semantically adequate on its own, but would +silently regress the multi-replica-per-host case the baseline already +supports. + +**Decision: the server maintains one independent distributed +sub-identity per connected worker (`client_id`)**, implemented by a +new, purpose-built `ExternalKvIdentityClient` class rather than N +`MasterClient` instances, for two reasons: + +- **Routing/eviction isolation.** A naive `MasterClient`-based + sub-identity with nonzero `tier_capacities` would become eligible + for ordinary `RoutePut`/`EvictionManager` target selection + (`ClientRegistry::GetAliveClients()`), risking real KV blocks being + routed onto a node that isn't a real storage participant. + `ExternalKvIdentityClient` has no `tier_capacities`-reporting code + path at all (always empty) and never publishes `EventBundle`s on + heartbeat — both `RoutePut`'s `CollectEligibleOnTier` and + `EvictionManager::RunOnce` already skip any client with no entry for + the tier being considered, so this closes the issue structurally, + with zero changes to the routing/eviction algorithms themselves. +- **Avoiding a known `MasterClient` bug.** `MasterClient`'s existing + re-register path (triggered by `CLIENT_STATUS_UNKNOWN`) omits + `peer_address`/`engine_desc`, silently breaking `MatchExternalKv` + responses for that identity after a Master restart or reaper expiry. + `ExternalKvIdentityClient` is new code with a single + `BuildRegisterRequest()` helper shared by the initial `Register()` + call and any re-register path, so both fields are always included by + construction. + +All N sub-identities (N ≤ 8 for the stated deployment shape) share one +common `peer_address` — the server's own single physical `PeerService` +endpoint. Neither `ClientRegistry` nor `MasterServer` assumes a 1:1 +`node_id` ↔ `peer_address` mapping (`ClientRegistry` stores a plain +`node_id → peer_address` map; `MasterServer::GetOrCreateStub` keys its +gRPC stub cache by `node_id`), so this is not blocked by any existing +invariant. Each of the five external-KV RPC handlers dispatches to the +calling `client_id`'s own `ExternalKvIdentityClient` instance. The +bulk Put/Get/RegisterMemory RDMA data path is unaffected and continues +to share the server's one `DistributedClient` backend — RDMA +registration is inherently per-VA and identity-agnostic, so only the +identity-bound external-KV surface needs per-worker handling. + +**Rejected alternative**: one shared identity + `tenant_id`-style key +namespacing under one shared identity — the pattern Mooncake itself +uses for its own multi-`DummyClient`-per-`RealClient` case. Rejected +because it would lose the `distributed/` baseline's existing +per-worker routing precision, and because Mooncake has no feature +analogous to external-KV that would validate the pattern for this use +case. This is genuinely new engineering for UMBP with no precedent in +either UMBP's own `distributed/` code or in Mooncake. + +### 7.4 Wire schema + +The worker conveys `worker_node_id`, `worker_node_address`, and +optionally `tags` — not raw rank components — added to the existing +`RegisterMemoryRequest` message rather than a new RPC: + +```protobuf +message RegisterMemoryRequest { + string client_id = 1; + uint64 worker_base = 2; + uint64 size = 3; + string worker_node_id = 4; // new + string worker_node_address = 5; // new + repeated string tags = 6; // new, opaque key=value labels +} +``` + +`RegisterMemoryRequest` is already the gRPC call that completes a +worker's registration lifecycle (the confirmation step after the raw +fd handoff), so piggybacking on it means a worker's memory +registration and its external-KV identity registration start and end +together, without a second round trip or a second lifecycle to keep in +sync. `worker_node_id`/`worker_node_address` are plain (not +`optional`) strings; a worker that never sets them (e.g. talking to a +`StandaloneClient`-backed server) leaves them empty, and the server +only constructs an `ExternalKvIdentityClient` for a `client_id` whose +`RegisterMemoryRequest` carried a non-empty `worker_node_id`. Python +(`umbp_store.py`'s existing node_id/node_address derivation) remains +the single source of truth for these strings; the server never +re-derives rank semantics. + +### 7.5 Interface and lifecycle + +`ExternalKvIdentityClient` is a client-side C++ class inside +`umbp_standalone_server`, using the existing `UMBPMaster` gRPC service +(the same one `MasterClient` uses) — not a proto service of its own: + +- `Register(worker_node_id, worker_node_address, peer_address, engine_desc, tags)` + — sends `RegisterClientRequest` with empty `tier_capacities` and the + given `peer_address`/`engine_desc`. Constructed the moment a + worker's `RegisterMemoryRequest` carries a non-empty + `worker_node_id`. +- `Heartbeat()` — periodic, liveness-only. Sends `HeartbeatRequest` + with empty `tier_capacities` and empty `bundles` on every call; its + only job is keeping the `ClientRegistry` entry from expiring so + `MatchExternalKv` continues to resolve this `node_id`. +- `Unregister()` — sends `UnregisterClientRequest`. Called on (a) the + corresponding worker's `DeregisterMemory` RPC, and (b) server + shutdown. +- `ReportExternalKvBlocks`/`RevokeExternalKvBlocks`/ + `RevokeAllExternalKvBlocksAtTier`/`MatchExternalKv`/ + `GetExternalKvHitCounts` — thin forwarding wrappers to the + corresponding `UMBPMaster` RPCs, using this sub-identity's own + `node_id`. + +Access to a per-`client_id` `ExternalKvIdentityClient` is guarded by +the same locking discipline used elsewhere in `standalone_server.cpp`, +since each sub-identity is a stateful object with its own background +heartbeat thread. + +Under a `StandaloneClient` backend, external-KV behavior is unchanged +from the local-backend shape: hardcoded stubs, matching +`StandaloneClient`'s own no-op external-KV methods. + +### 7.6 Known limitations + +**No crash/disconnect detection.** The only paths that tear down an +`ExternalKvIdentityClient` are `DeregisterMemory` (explicit, +worker-initiated) and full server shutdown. If a worker process +crashes or is killed without calling `DeregisterMemory`, its +`ExternalKvIdentityClient`'s heartbeat thread (which lives in the +server process, unaffected by the worker's death) keeps running +indefinitely — the Master's `ClientRegistry` entry for that `node_id` +never expires, and `MatchExternalKv`/`GetExternalKvHitCounts` keep +returning it for KV blocks that may no longer exist. This is not +cleaned up until the entire `umbp_standalone_server` process is +restarted. Fixing this needs a general connection-liveness/crash-detection +mechanism — the same kind of work needed for +`StandaloneProcessClient`'s own reconnect handling (§8, item 1) — and +is a deliberately deferred, tracked limitation. + +**Lifecycle lock serializes register/deregister across workers.** +Register, deregister, and shutdown bookkeeping for +`ExternalKvIdentityClient` are serialized by one server-wide lifecycle +lock, closing register-vs-deregister races for the same `client_id`. +This lock can be held across blocking Master RPCs +(`RegisterClient`/`UnregisterClient`, bounded by the RPC deadline), so +when the Master is slow or unavailable, otherwise-unrelated workers +can queue behind each other during `RegisterMemory`/`DeregisterMemory` +— worst case, startup/shutdown delay approaches `N * rpc_deadline` for +the stated 8-worker shape. This is an accepted trade-off at the +current scale (a lifecycle path, not the Put/Get hot path). If future +deployments need more workers, more frequent reconnects, or hit +startup stalls from slow Master RPCs, revisit with per-`client_id` +lifecycle locks or by decoupling external-KV identity registration +from core memory registration via an async/retry path. + +## 8. Known Limitations and Open Items + +1. **Crash/restart semantics.** No supervised-restart or + reconnect-with-cache-rehydration story exists — matches every + reference system reviewed (Mooncake, LMCache, UMBP's own SPDK-proxy + daemon). If the standalone server dies, every attached worker loses + its DRAM cache simultaneously. Workers must treat a broken RPC + channel as "cache unavailable," not fatal; this needs a connection + state machine (`CONNECTED` → `DISCONNECTED` → optionally + `RECONNECTING`) in `StandaloneProcessClient` that does not exist + today. +2. **Per-tenant DRAM quota** is not implemented in + `LocalStorageManager`/`DRAMTier` — no notion of "owner" per key + exists, so multiple workers sharing one server's DRAM tier can + starve each other. +3. **fd-handoff listener lifetime**: long-lived for the server + process's whole lifetime (matches Mooncake), chosen over a + short-lived per-registration-window socket for simplicity, relying + on filesystem permissions (item 4) rather than a narrower open + window for hardening. +4. **Security/isolation**: the shm segment (`memfd`, `O_CLOEXEC` + recommended) and both UDS socket files must be created with + `0600`/owner-only permissions, or another local user could attach + to a different tenant's KV cache. Neither Mooncake's nor UMBP's + SPDK-proxy code sets restrictive permissions on their sockets/shm + by default — this must be added deliberately. +5. **gRPC-over-UDS batch overhead**: a gRPC call costs a syscall + + protobuf encode/decode per `BatchGet`/`BatchPut` even with + struct-of-arrays batching, unlike a busy-polled shm ring. Acceptable + for v1; if profiling later shows this is a bottleneck for very + small, very frequent `Get` calls, a future iteration could add an + optional shm-ring fast path for single-key Get/Put alongside gRPC. +6. **Process-local ptr→fd registry** (`HostMemAllocator::Alloc`/`Free`) + must be safe for `Free` to run concurrently with a `RegisterMemory` + lookup from a different thread, and `Free`/`DeregisterMemory` + ordering must be defined — `Free` should refuse (or at least warn + loudly) if an active registration still references that pointer, + rather than silently leaving the server with a stale mapping. +7. **`tcp_staging` transport** is deferred; not built unless a + concrete deployment needs it (§2). +8. **Version/ABI skew**: the standalone server binary and the + `mori_pybinds`-linked proto definitions in each worker process must + agree on the `UMBPStandalone` proto schema — the same constraint + `distributed/` already has between `umbp_master` and workers, but + more likely in practice here since a standalone server is typically + a longer-lived sidecar that outlives several worker restarts/ + upgrades. No solution beyond the existing binary/version-pinning + discipline. +9. **No fd-handoff registration reaper**: if a worker sends its fd via + the raw fd-handoff socket and then crashes before issuing the + confirming gRPC `RegisterMemory` RPC, the mapping stays live in the + server's registry indefinitely — pre-existing, orthogonal to which + backend `client_` is. +10. **Distributed-backend test coverage**: the existing + `test_standalone_shm_ipc.cpp` suite only exercises the + `StandaloneClient`-backed path. A `DistributedClient`-backed test + needs a running `umbp_master` plus the RDMA/IO-engine stack, likely + gated separately (e.g. only where RDMA hardware/loopback is + available). Should also cover per-worker external-KV sub-identity + registration/dispatch with 2+ connected workers, and the + `RegisterMemory`/`DeregisterMemory` rollback/shutdown-ordering + paths (§6.3). +11. **Naming convention**: whether to formally distinguish + `standalone-process/distributed-backend` vs + `standalone-process/local-backend` in docs, config, and log output + (recommended, not yet adopted in code/env-var naming). + +## 9. Compatibility with Existing Modes + +- `local/` (in-process `StandaloneClient`) is unchanged and remains + the default when neither `distributed` nor `standalone_process` is + set. +- `distributed/` (cross-node master+peers+RDMA) is unchanged for + ordinary workers; this design does not touch any file under + `distributed/master/`, `distributed/peer/`, or `distributed/routing/`. +- Shared code touched: `common/config.h` (new `standalone_process` + optional field on `UMBPConfig`), `umbp_client_factory.cpp` (new + branch), `local/host_mem_allocator.*` (new `kAnonymousShm` backing + plus the fd registry, additive), `local/tiers/copy_pipeline.*` (new + `Drain()`), `pybind_umbp.cpp` (new `AnonymousShm` enum value and + `standalone_process`/`UMBPStandaloneProcessConfig` bindings), and + `umbp_store.py`'s `register_mem_pool_host` gate and error-handling + branches (§10). +- Future extension, out of scope here: the standalone server could + itself become a `distributed/` peer/leader (a host running one + standalone server that both serves local workers over UDS and + participates in the cross-node distributed pool directly) — + `LocalStorageManager`'s existing `SharedSSDLeader`/`SharedSSDFollower` + roles hint at this kind of layering. + +## 10. Python / Pybind Interface + +- `get_deployment_mode()` (enum `Local`/`StandaloneProcess`/ + `Distributed`) is added to `IUMBPClient`/pybind. `is_distributed()` + remains a pure bool getter, unchanged, continuing to mean "true + cross-node distributed" for existing external-KV branch points. +- `CreateUMBPClient(const UMBPConfig&)` gains a third branch: + `config.standalone_process.has_value()` (new optional field on + `UMBPConfig`, alongside the existing `distributed`) constructs + `StandaloneProcessClient`. `distributed` and `standalone_process` + are mutually exclusive within one `UMBPConfig`; + `UMBPConfig::Validate()` rejects setting both. +- **Activation is `UMBP_STANDALONE_ADDRESS`-only.** + `extra_config["standalone_address"]` without the env var raises at + `UMBPStore.__init__`. This is required because the host memory pool + (`MHATokenToKVPoolHost`) is constructed and allocates its `kv_buffer` + inside `HiRadixCache.__init__`, before `storage_backend_extra_config` + is parsed and handed to `UMBPStore` — `UMBPHostTensorAllocator` can + only see process environment variables at the point it must decide a + buffer's backing, so `extra_config` alone cannot reliably activate + this mode. `auto_start`/`startup_timeout_ms` (which only affect what + `CreateUMBPClient` does once the config is already built) may still + come from either `extra_config` or their env vars, since they carry + no allocator-timing constraint. +- `umbp_store.py`'s `register_mem_pool_host` gate is changed from `if + not is_distributed(): return` to also allow `StandaloneProcess` — + otherwise a correctly-constructed `StandaloneProcessClient` (which + reports `is_distributed() == False`) would never have its host KV + buffer registered, and the fd-handoff handshake (§4.2) would never + trigger. +- Under `get_deployment_mode() == StandaloneProcess`, all three of + `register_mem_pool_host`'s existing failure paths — a set + `disable_zero_copy_register`, a `register_memory` exception, and + `register_memory() == False` — raise instead of warn-and-return. + This differs from `distributed`'s behavior (unchanged: warn and fall + back to the staging-buffer path), because standalone-process mode + has no staging/inline fallback to degrade to: a swallowed + registration failure would otherwise mean every later `Put`/`Get` + fails later, with a less correlated-looking error. +- `pybind_umbp.cpp` additions: `AnonymousShm`/`kAnonymousShm` value on + `UMBPHostBufferBacking`; a new `standalone_process` property and + `UMBPStandaloneProcessConfig` pybind class on `UMBPConfig`. +- `umbp_host_allocator.py`: `UMBPHostTensorAllocator.__init__` reads + `UMBP_STANDALONE_ADDRESS` directly from `os.environ` (not via + `UMBPConfig`/`extra_config`, neither of which exists yet at this + point in the process) and requests the `AnonymousShm` backing when + set. + +## 11. Configuration Surface + +`UMBPConfig` gains `std::optional +standalone_process`, with exactly three fields, all Python-set +(mirroring `UMBPDistributedConfig`'s existing convention): `address`, +`auto_start`, `startup_timeout_ms`. + +| Var | Consumed by | Purpose | Default | +|---|---|---|---| +| `UMBP_STANDALONE_ADDRESS` | Python, into `cfg.standalone_process.address` | UDS path of the standalone server's gRPC socket. Presence enables standalone-process mode. | unset (disabled) | +| `UMBP_STANDALONE_AUTO_START` | Python, into `cfg.standalone_process.auto_start` | If `true` and no server found at `address`, `CreateUMBPClient` forks+execs `umbp_standalone_server` (rank-0-local only, local-backend only, §5.2). | `false` | +| `UMBP_STANDALONE_STARTUP_TIMEOUT_MS` | Python, into `cfg.standalone_process.startup_timeout_ms` | Bound on waiting for readiness after spawn. | `30000` | +| `UMBP_STANDALONE_BIN` | Deployment-only, read directly by the auto-start code at spawn time | Path override for the `umbp_standalone_server` binary. | resolved via `PATH`/build dir | +| `UMBP_STANDALONE_IDLE_EXIT_TIMEOUT_MS` | `umbp_standalone_server`'s own env read at startup | `0` = never self-exit (default). Non-zero opts into SPDK-proxy-style idle exit. | `0` | +| `UMBP_STANDALONE_GRPC_SHUTDOWN_DEADLINE_SEC` | `umbp_standalone_server`'s own env read | Reuses the `MasterServer` shutdown-deadline convention. | same as `UMBP_GRPC_SHUTDOWN_DEADLINE_SEC` | +| `UMBP_STANDALONE_SHM_DIR` | Both sides' own env read (deployment-only) | Directory for the bootstrap-lock file (not the shm segment itself, which is anonymous via `memfd_create`). | `/tmp` | + +Server-side distributed-backend variables are listed in §6.6. + +`fd_socket` and `transport` are deliberately not config fields: the +fd-handoff socket path is always derived from `address` (§5.1), and +`transport` (`uds` vs. the deferred `tcp_staging`) has no v1 surface +since `tcp_staging` isn't built. + +### 11.1 DRAM tier capacity: ownership moves from per-rank to the server + +The UMBP DRAM cache tier is sized by `UMBP_DRAM_CAPACITY` +(`UMBPDramConfig::capacity_bytes`, read by `UMBPConfig::FromEnvironment()` +in `common/config.h`; also settable in Python via +`extra_config["dram_capacity_bytes"]`). This tier is memory the backend +allocates for itself, distinct from SGLang's host KV pool +(the `umbp_host_allocator`-backed buffer that workers register for +zero-copy transfer, §4.2): `Put` copies bytes *from* the registered +host buffer *into* this tier. + +Where this tier lives, and therefore which process's `UMBP_DRAM_CAPACITY` +sizes it, differs by mode: + +- **`local/` (in-process)**: each worker constructs its own + `LocalStorageManager` and its own `DRAMTier`, each sized by that + worker process's `UMBP_DRAM_CAPACITY`. N workers on a host produce N + independent tiers. +- **Standalone-process**: `StandaloneProcessClient` allocates no tier — + it only forwards RPCs (`standalone_process_client.{h,cpp}`). The + single DRAM tier lives in `umbp_standalone_server` and is shared by + all connected workers (§5.4, §5.5). It is sized by the **server + process's** `UMBP_DRAM_CAPACITY` + (`standalone_server_main.cpp` builds the backend config via + `UMBPConfig::FromEnvironment()`). A worker's own `UMBP_DRAM_CAPACITY` + no longer sizes a local tier. + +The two launch paths (§5.1, §5.2) determine where the server reads that +value from: + +- **Externally launched (§5.1; also the only supported path for the + distributed backend)**: the server reads `UMBP_DRAM_CAPACITY` from its + own launch environment. Worker-side values are not consulted for tier + sizing. +- **Auto-start (§5.2; local-backend only)**: the bootstrap-winning + worker (local rank 0) forks the server and exports its own + `cfg.dram.capacity_bytes` into the child's environment as + `UMBP_DRAM_CAPACITY` (`standalone_process_client.cpp`, + `ExportServerEnv`). The server's tier is therefore sized from that one + worker's value, not an aggregate of all workers'. + +For the distributed backend, `UMBP_DRAM_CAPACITY` sizes the +`DistributedClient`/`PoolClient` DRAM pool +(`distributed/distributed_client.cpp`) rather than a +`LocalStorageManager` tier; the staging buffer is sized separately by +`UMBP_DISTRIBUTED_STAGING_BUFFER_SIZE` (§6.6). + +## 12. Source References + +- UMBP: `umbp_client_factory.cpp`, `include/umbp/umbp_client.h`, + `include/umbp/common/config.h`, `local/standalone_client.{h,cpp}`, + `local/host_mem_allocator.{h,cpp}`, `local/tiers/dram_tier.cpp`, + `local/tiers/local_storage_manager.cpp`, `local/tiers/copy_pipeline.{h,cpp}`, + `storage/spdk/proxy/spdk_proxy_protocol.h`, `storage/spdk/proxy/spdk_proxy_shm.cpp`, + `distributed/peer/batch_resolve_codec.h`, + `distributed/master/master_server.cpp`, `distributed/master/master_client.cpp`, + `distributed/master/client_registry.cpp`, `distributed/master/eviction_manager.cpp`, + `distributed/routing/router.cpp`, `distributed/routing/route_put_strategy.cpp`, + `distributed/distributed_client.{h,cpp}`, `distributed/pool_client.{h,cpp}`, + `distributed/proto/umbp.proto`, `distributed/proto/umbp_standalone.proto`, + `standalone/standalone_server.{h,cpp}`, `standalone/bin/standalone_server_main.cpp`, + `src/pybind/pybind_umbp.cpp`, `src/pybind/CMakeLists.txt`, + `doc/runtime-env-vars.md`, `scripts/run_umbp_single_node_hicache.sh`. +- SGLang: `sglang/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py`, + `umbp_host_allocator.py`. +- Mooncake (`/apps/nima/KVManager/Mooncake`): `mooncake-store/src/real_client.cpp`, + `real_client_main.cpp`, `shm_helper.cpp`, `dummy_client.cpp`, `client_buffer.cpp`, + `uds_transport.cpp`, `docs/source/design/mooncake-store.md`, + `docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md`. +- LMCache (`/apps/nima/KVManager/LMCache`, reviewed for comparison only): + `lmcache/v1/multiprocess/{mq,server,protocol,custom_types,futures,config}.py`, + `lmcache/v1/mp_observability/`. diff --git a/src/umbp/include/umbp/common/config.h b/src/umbp/include/umbp/common/config.h index f29ea1a06..21727f887 100644 --- a/src/umbp/include/umbp/common/config.h +++ b/src/umbp/include/umbp/common/config.h @@ -243,6 +243,22 @@ struct UMBPDistributedConfig { uint64_t dram_page_size = 0; }; +// User-facing same-host standalone-process configuration. Set +// UMBPConfig::standalone_process to make CreateUMBPClient construct a +// StandaloneProcessClient that talks to an umbp_standalone_server over UDS. +struct UMBPStandaloneProcessConfig { + std::string address; // e.g. unix:///run/umbp/standalone/node0.grpc.sock + bool auto_start = false; // opt-in fork+exec convenience path + int startup_timeout_ms = 30000; // readiness wait bound for auto_start + + // Optional distributed identity for external-KV reports routed through a + // distributed-backed standalone server. Empty means external-KV identity is + // not requested for this worker. + std::string worker_node_id; + std::string worker_node_address; + std::vector tags; +}; + struct UMBPConfig { UMBPDramConfig dram; UMBPSsdConfig ssd; @@ -262,6 +278,11 @@ struct UMBPConfig { // nullopt (default) = local-only mode with no network dependencies. std::optional distributed; + // Optional same-host standalone-process mode. Mutually exclusive with + // distributed: this mode uses one local server process and shm fd handoff, + // not the cross-node master/RDMA path. + std::optional standalone_process; + UMBPRole ResolveRole() const { if (role != UMBPRole::Standalone) { return role; @@ -328,6 +349,22 @@ struct UMBPConfig { return false; } } + if (distributed.has_value() && standalone_process.has_value()) { + if (error_message) + *error_message = "distributed and standalone_process are mutually exclusive"; + return false; + } + if (standalone_process.has_value()) { + const auto& sp = standalone_process.value(); + if (sp.address.empty()) { + if (error_message) *error_message = "standalone_process.address must not be empty"; + return false; + } + if (sp.startup_timeout_ms <= 0) { + if (error_message) *error_message = "standalone_process.startup_timeout_ms must be > 0"; + return false; + } + } return true; } diff --git a/src/umbp/include/umbp/distributed/distributed_client.h b/src/umbp/include/umbp/distributed/distributed_client.h index 689c19ca7..2d39b50be 100644 --- a/src/umbp/include/umbp/distributed/distributed_client.h +++ b/src/umbp/include/umbp/distributed/distributed_client.h @@ -65,6 +65,7 @@ class DistributedClient : public IUMBPClient { bool Flush() override; void Close() override; bool IsDistributed() const override; + UMBPDeploymentMode GetDeploymentMode() const override { return UMBPDeploymentMode::Distributed; } bool RegisterMemory(uintptr_t ptr, size_t size) override; void DeregisterMemory(uintptr_t ptr) override; diff --git a/src/umbp/include/umbp/local/host_mem_allocator.h b/src/umbp/include/umbp/local/host_mem_allocator.h index 476ffdaba..cfbdaf6fc 100644 --- a/src/umbp/include/umbp/local/host_mem_allocator.h +++ b/src/umbp/include/umbp/local/host_mem_allocator.h @@ -23,12 +23,14 @@ #include #include +#include namespace mori::umbp { enum class HostBufferBacking : int { kAnonymous = 0, kAnonymousHugetlb = 1, + kAnonymousShm = 2, }; struct HostBufferOptions { @@ -50,6 +52,12 @@ struct HostBufferHandle { class HostMemAllocator { public: + struct ShmAllocation { + void* base = nullptr; + size_t mapped_size = 0; + int fd = -1; + }; + HostMemAllocator() = default; ~HostMemAllocator() = default; @@ -66,6 +74,18 @@ class HostMemAllocator { // fails (a WARN is logged), to prevent a subsequent Free from munmap'ing // an address the kernel may have reused for a later mmap(). void Free(HostBufferHandle& handle); + + // Look up a live kAnonymousShm allocation containing [ptr, ptr + size). + // The returned fd is owned by the allocator registry; callers may pass it + // through SCM_RIGHTS but must not close it. + static std::optional LookupShmAllocation(uintptr_t ptr, size_t size); + + // Acquire/Release a live shm allocation while its fd is being handed to, or + // remains registered with, another process. Acquire increments the registry's + // active reference count before returning the fd, closing the lookup/use + // race with Free(). + static std::optional AcquireShmAllocation(uintptr_t ptr, size_t size); + static void ReleaseShmAllocation(uintptr_t base); }; } // namespace mori::umbp diff --git a/src/umbp/include/umbp/local/standalone_client.h b/src/umbp/include/umbp/local/standalone_client.h index 95f6a1415..b2c8906b5 100644 --- a/src/umbp/include/umbp/local/standalone_client.h +++ b/src/umbp/include/umbp/local/standalone_client.h @@ -63,6 +63,7 @@ class StandaloneClient : public IUMBPClient { 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 { diff --git a/src/umbp/include/umbp/local/tiers/copy_pipeline.h b/src/umbp/include/umbp/local/tiers/copy_pipeline.h index 546983573..8f045736e 100644 --- a/src/umbp/include/umbp/local/tiers/copy_pipeline.h +++ b/src/umbp/include/umbp/local/tiers/copy_pipeline.h @@ -22,6 +22,7 @@ #pragma once #include +#include #include #include #include @@ -42,6 +43,7 @@ class CopyPipeline { bool MaybeCopyToSharedSSD(const std::string& key); void MaybeBatchCopyToSharedSSD(const std::vector& keys); + bool Drain(std::chrono::milliseconds timeout); private: struct CopyTask { @@ -60,7 +62,9 @@ class CopyPipeline { 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; }; } // namespace mori::umbp diff --git a/src/umbp/include/umbp/standalone/external_kv_identity_client.h b/src/umbp/include/umbp/standalone/external_kv_identity_client.h new file mode 100644 index 000000000..78b336f16 --- /dev/null +++ b/src/umbp/include/umbp/standalone/external_kv_identity_client.h @@ -0,0 +1,91 @@ +// 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/umbp_client.h" + +namespace mori::umbp::standalone { + +// Lightweight Master identity used only for external-KV reports on behalf of +// one worker attached to a distributed-backed standalone server. It never +// reports storage capacities or owned-KV event bundles. +class ExternalKvIdentityClient { + public: + struct Config { + std::string master_address; + std::string node_id; + std::string node_address; + std::string peer_address; + std::vector engine_desc_bytes; + std::vector tags; + }; + + explicit ExternalKvIdentityClient(Config config); + ~ExternalKvIdentityClient(); + + ExternalKvIdentityClient(const ExternalKvIdentityClient&) = delete; + ExternalKvIdentityClient& operator=(const ExternalKvIdentityClient&) = delete; + + bool Start(); + void Stop(); + const std::string& node_id() const { return config_.node_id; } + + bool ReportExternalKvBlocks(const std::vector& hashes, TierType tier); + bool RevokeExternalKvBlocks(const std::vector& hashes, TierType tier); + bool RevokeAllExternalKvBlocksAtTier(TierType tier); + std::vector MatchExternalKv(const std::vector& hashes, + bool count_as_hit = false); + std::vector GetExternalKvHitCounts( + const std::vector& hashes); + + private: + bool RegisterLocked(); + void UnregisterLocked(); + bool SendHeartbeatOnceLocked(); + void HeartbeatLoop(); + + Config config_; + std::shared_ptr channel_; + std::unique_ptr stub_; + + std::mutex rpc_mu_; + std::mutex cv_mu_; + std::condition_variable cv_; + std::thread heartbeat_thread_; + std::atomic running_{false}; + bool registered_ = false; + uint64_t heartbeat_interval_ms_ = 1000; +}; + +} // namespace mori::umbp::standalone diff --git a/src/umbp/include/umbp/standalone/ipc.h b/src/umbp/include/umbp/standalone/ipc.h new file mode 100644 index 000000000..4a57485ae --- /dev/null +++ b/src/umbp/include/umbp/standalone/ipc.h @@ -0,0 +1,58 @@ +// 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 + +namespace mori::umbp::standalone { + +constexpr uint32_t kFdRegistrationMagic = 0x554d4250; // "UMBP" +constexpr uint32_t kFdRegistrationVersion = 1; +constexpr size_t kClientIdBytes = 64; + +struct FdRegistrationMessage { + uint32_t magic = kFdRegistrationMagic; + uint32_t version = kFdRegistrationVersion; + char client_id[kClientIdBytes] = {}; + uint64_t worker_base = 0; + uint64_t size = 0; +}; + +std::string UnixPathFromGrpcAddress(const std::string& address); +std::string DeriveFdSocketPath(const std::string& grpc_address); +std::string DefaultStandaloneAddress(); + +bool EnsureParentDirectory(const std::string& path, std::string* error = nullptr); + +int SendFdRegistration(const std::string& socket_path, int fd, const std::string& client_id, + uintptr_t worker_base, size_t size, int timeout_ms, std::string* error); +int RecvFdRegistration(int socket_fd, FdRegistrationMessage* message, std::string* error); + +bool SendStatus(int socket_fd, int32_t status, std::string* error = nullptr); +bool RecvStatus(int socket_fd, int32_t* status, std::string* error = nullptr); + +} // namespace mori::umbp::standalone diff --git a/src/umbp/include/umbp/standalone/standalone_process_client.h b/src/umbp/include/umbp/standalone/standalone_process_client.h new file mode 100644 index 000000000..4b24f9447 --- /dev/null +++ b/src/umbp/include/umbp/standalone/standalone_process_client.h @@ -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. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/umbp_client.h" +#include "umbp_standalone.grpc.pb.h" + +namespace mori::umbp::standalone { + +class StandaloneProcessClient : public IUMBPClient { + public: + explicit StandaloneProcessClient(const UMBPConfig& config); + ~StandaloneProcessClient() override; + + 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 { return false; } + UMBPDeploymentMode GetDeploymentMode() const override { + return UMBPDeploymentMode::StandaloneProcess; + } + + bool RegisterMemory(uintptr_t ptr, size_t size) override; + void DeregisterMemory(uintptr_t ptr) override; + + bool ReportExternalKvBlocks(const std::vector& hashes, TierType tier) override; + bool RevokeExternalKvBlocks(const std::vector& hashes, TierType tier) override; + bool RevokeAllExternalKvBlocksAtTier(TierType tier) override; + std::vector MatchExternalKv(const std::vector& hashes, + bool count_as_hit = false) override; + std::vector GetExternalKvHitCounts( + const std::vector& hashes) override; + + private: + bool OffsetFor(uintptr_t ptr, size_t size, uint64_t* offset) const; + bool WaitReady(int timeout_ms) const; + void MaybeAutoStart(); + std::string ClientId(); + void DeregisterMemoryLocked(); + + UMBPConfig config_; + UMBPStandaloneProcessConfig standalone_config_; + std::string address_; + std::string fd_socket_path_; + std::shared_ptr<::grpc::ChannelInterface> channel_; + std::unique_ptr<::umbp::UMBPStandalone::Stub> stub_; + + mutable std::shared_mutex op_mutex_; + std::atomic closing_{false}; + bool closed_ = false; + + mutable std::mutex registration_mu_; + std::string client_id_; + uintptr_t registered_base_ = 0; + size_t registered_size_ = 0; + bool registered_ = false; +}; + +} // namespace mori::umbp::standalone diff --git a/src/umbp/include/umbp/standalone/standalone_server.h b/src/umbp/include/umbp/standalone/standalone_server.h new file mode 100644 index 000000000..9eba001e1 --- /dev/null +++ b/src/umbp/include/umbp/standalone/standalone_server.h @@ -0,0 +1,53 @@ +// 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 "umbp/common/config.h" + +namespace mori::umbp::standalone { + +class StandaloneServer { + public: + StandaloneServer(UMBPConfig config, std::string address); + ~StandaloneServer(); + + bool Start(); + void Run(); + void Shutdown(); + + const std::string& address() const { return address_; } + + private: + class Impl; + + UMBPConfig config_; + std::string address_; + std::unique_ptr impl_; +}; + +} // namespace mori::umbp::standalone diff --git a/src/umbp/include/umbp/umbp_client.h b/src/umbp/include/umbp/umbp_client.h index 7221a64bc..fec4200d1 100644 --- a/src/umbp/include/umbp/umbp_client.h +++ b/src/umbp/include/umbp/umbp_client.h @@ -35,6 +35,12 @@ namespace mori::umbp { +enum class UMBPDeploymentMode : int { + Local = 0, + StandaloneProcess = 1, + Distributed = 2, +}; + /// Abstract interface for UMBP storage clients. /// /// Two implementations exist behind this interface: @@ -107,6 +113,11 @@ class IUMBPClient { /// Returns true when the client operates in distributed (master-led) mode. virtual bool IsDistributed() const = 0; + /// Returns the concrete deployment mode behind this client. + virtual UMBPDeploymentMode GetDeploymentMode() const { + return IsDistributed() ? UMBPDeploymentMode::Distributed : UMBPDeploymentMode::Local; + } + // ---- Optional zero-copy hooks ---- // // Register a host buffer for zero-copy RDMA transfers. Standalone diff --git a/src/umbp/local/host_mem_allocator.cpp b/src/umbp/local/host_mem_allocator.cpp index 3bfc49e4c..badddf666 100644 --- a/src/umbp/local/host_mem_allocator.cpp +++ b/src/umbp/local/host_mem_allocator.cpp @@ -21,6 +21,7 @@ // SOFTWARE. #include "umbp/local/host_mem_allocator.h" +#include #include #include @@ -33,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -45,6 +47,47 @@ namespace { constexpr size_t kDefaultPageSize = 4096; +#ifndef MFD_CLOEXEC +#define MFD_CLOEXEC 0x0001U +#endif + +struct ShmAllocRecord { + int fd = -1; + size_t size = 0; + size_t active_registrations = 0; + bool pending_free = false; +}; + +std::mutex& ShmRegistryMutex() { + static std::mutex mu; + return mu; +} + +std::map& ShmRegistry() { + static std::map registry; + return registry; +} + +#ifdef __linux__ +int MemfdCreate(const char* name, unsigned int flags) { +#if defined(__NR_memfd_create) + return static_cast(syscall(__NR_memfd_create, name, flags)); +#else + (void)name; + (void)flags; + errno = ENOSYS; + return -1; +#endif +} +#else +int MemfdCreate(const char* name, unsigned int flags) { + (void)name; + (void)flags; + errno = ENOSYS; + return -1; +} +#endif + size_t GetPageSize() { const long page = sysconf(_SC_PAGESIZE); return page > 0 ? static_cast(page) : kDefaultPageSize; @@ -203,6 +246,57 @@ HostBufferHandle AllocAnonymous(size_t size, const HostBufferOptions& opts) { return handle; } +HostBufferHandle AllocAnonymousShm(size_t size, const HostBufferOptions& opts) { + HostBufferHandle handle; + if (size == 0) return handle; + + const size_t page_size = GetPageSize(); + const std::optional mapped_size = AlignUpChecked(size, page_size); + if (!mapped_size.has_value()) return handle; + + const int fd = MemfdCreate("umbp_host_buffer", MFD_CLOEXEC); + if (fd < 0) { + const int err = errno; + MORI_UMBP_WARN("HostMemAllocator: memfd_create failed ({}: {})", err, std::strerror(err)); + return handle; + } + + if (ftruncate(fd, static_cast(*mapped_size)) != 0) { + const int err = errno; + MORI_UMBP_WARN("HostMemAllocator: ftruncate(memfd, {}) failed ({}: {})", *mapped_size, err, + std::strerror(err)); + close(fd); + return handle; + } + + void* ptr = mmap(nullptr, *mapped_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + if (ptr == MAP_FAILED) { + const int err = errno; + MORI_UMBP_WARN("HostMemAllocator: mmap(memfd, {}) failed ({}: {})", *mapped_size, err, + std::strerror(err)); + close(fd); + return handle; + } + + handle.ptr = ptr; + handle.requested_size = size; + handle.mapped_size = *mapped_size; + handle.actual_backing = HostBufferBacking::kAnonymousShm; + handle.actual_alignment = page_size; + + { + std::lock_guard lock(ShmRegistryMutex()); + ShmAllocRecord record; + record.fd = fd; + record.size = *mapped_size; + record.active_registrations = 0; + ShmRegistry()[ptr] = record; + } + + ApplyPostMappingPolicies(handle, opts); + return handle; +} + HostBufferHandle AllocAnonymousHugetlb(size_t size, const HostBufferOptions& opts) { HostBufferHandle handle; if (size == 0) return handle; @@ -261,6 +355,8 @@ HostBufferHandle HostMemAllocator::Alloc(size_t size, const HostBufferOptions& o return AllocAnonymous(size, opts); case HostBufferBacking::kAnonymousHugetlb: return AllocAnonymousHugetlb(size, opts); + case HostBufferBacking::kAnonymousShm: + return AllocAnonymousShm(size, opts); default: return {}; } @@ -269,21 +365,122 @@ HostBufferHandle HostMemAllocator::Alloc(size_t size, const HostBufferOptions& o void HostMemAllocator::Free(HostBufferHandle& handle) { if (!handle.valid()) return; - if (munmap(handle.ptr, handle.mapped_size) != 0) { + void* ptr_to_unmap = handle.ptr; + size_t size_to_unmap = handle.mapped_size; + bool defer_unmap = false; + int fd_to_close = -1; + if (handle.actual_backing == HostBufferBacking::kAnonymousShm) { + std::lock_guard lock(ShmRegistryMutex()); + auto it = ShmRegistry().find(handle.ptr); + if (it != ShmRegistry().end()) { + if (it->second.active_registrations > 0) { + MORI_UMBP_WARN( + "HostMemAllocator: freeing AnonymousShm buffer ptr={} while {} active registration(s) " + "still reference it; deferring munmap/close until deregistration", + handle.ptr, it->second.active_registrations); + it->second.pending_free = true; + defer_unmap = true; + } else { + fd_to_close = it->second.fd; + ShmRegistry().erase(it); + } + } + } + + handle.ptr = nullptr; + handle.requested_size = 0; + handle.mapped_size = 0; + + if (defer_unmap) return; + + if (munmap(ptr_to_unmap, size_to_unmap) != 0) { const int err = errno; MORI_UMBP_WARN( "HostMemAllocator: munmap failed ({}: {}); invalidating handle anyway " "to prevent a possible double-free of a reused VA range", err, std::strerror(err)); - // Fall through to invalidation below: keeping a stale-but-valid handle - // would let the next Free() munmap an address the kernel may have - // already handed back to a later mmap(). Accepting a small leak on - // the (already-rare) munmap-failure path is the lesser evil. + // Keeping a stale-but-valid handle would let the next Free() munmap an + // address the kernel may have already handed back to a later mmap(). } - handle.ptr = nullptr; - handle.requested_size = 0; - handle.mapped_size = 0; + if (fd_to_close >= 0) close(fd_to_close); +} + +std::optional HostMemAllocator::LookupShmAllocation(uintptr_t ptr, + size_t size) { + if (ptr == 0 || size == 0) return std::nullopt; + if (ptr > std::numeric_limits::max() - (size - 1)) return std::nullopt; + + std::lock_guard lock(ShmRegistryMutex()); + auto& registry = ShmRegistry(); + auto it = registry.upper_bound(reinterpret_cast(ptr)); + if (it == registry.begin()) return std::nullopt; + --it; + + uintptr_t base = reinterpret_cast(it->first); + uintptr_t end = base + it->second.size; + uintptr_t req_end = ptr + size; + if (ptr < base || req_end > end || it->second.fd < 0 || it->second.pending_free) { + return std::nullopt; + } + + ShmAllocation allocation; + allocation.base = it->first; + allocation.mapped_size = it->second.size; + allocation.fd = it->second.fd; + return allocation; +} + +std::optional HostMemAllocator::AcquireShmAllocation(uintptr_t ptr, + size_t size) { + if (ptr == 0 || size == 0) return std::nullopt; + if (ptr > std::numeric_limits::max() - (size - 1)) return std::nullopt; + + std::lock_guard lock(ShmRegistryMutex()); + auto& registry = ShmRegistry(); + auto it = registry.upper_bound(reinterpret_cast(ptr)); + if (it == registry.begin()) return std::nullopt; + --it; + + uintptr_t base = reinterpret_cast(it->first); + uintptr_t end = base + it->second.size; + uintptr_t req_end = ptr + size; + if (ptr < base || req_end > end || it->second.fd < 0 || it->second.pending_free) { + return std::nullopt; + } + + ++it->second.active_registrations; + ShmAllocation allocation; + allocation.base = it->first; + allocation.mapped_size = it->second.size; + allocation.fd = it->second.fd; + return allocation; +} + +void HostMemAllocator::ReleaseShmAllocation(uintptr_t base_ptr) { + void* ptr_to_unmap = nullptr; + size_t size_to_unmap = 0; + int fd_to_close = -1; + { + std::lock_guard lock(ShmRegistryMutex()); + auto it = ShmRegistry().find(reinterpret_cast(base_ptr)); + if (it == ShmRegistry().end()) return; + if (it->second.active_registrations > 0) --it->second.active_registrations; + if (it->second.active_registrations == 0 && it->second.pending_free) { + ptr_to_unmap = it->first; + size_to_unmap = it->second.size; + fd_to_close = it->second.fd; + ShmRegistry().erase(it); + } + } + + if (ptr_to_unmap) { + if (munmap(ptr_to_unmap, size_to_unmap) != 0) { + const int err = errno; + MORI_UMBP_WARN("HostMemAllocator: deferred munmap failed ({}: {})", err, std::strerror(err)); + } + if (fd_to_close >= 0) close(fd_to_close); + } } } // namespace mori::umbp diff --git a/src/umbp/local/standalone_client.cpp b/src/umbp/local/standalone_client.cpp index de979f994..a0d77d2f2 100644 --- a/src/umbp/local/standalone_client.cpp +++ b/src/umbp/local/standalone_client.cpp @@ -21,6 +21,7 @@ // SOFTWARE. #include "umbp/local/standalone_client.h" +#include #include #include @@ -297,7 +298,13 @@ bool StandaloneClient::Clear() { return true; } -bool StandaloneClient::Flush() { return storage_.Flush(); } +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; diff --git a/src/umbp/local/tiers/copy_pipeline.cpp b/src/umbp/local/tiers/copy_pipeline.cpp index cbf3762eb..620ce3aaa 100644 --- a/src/umbp/local/tiers/copy_pipeline.cpp +++ b/src/umbp/local/tiers/copy_pipeline.cpp @@ -76,6 +76,13 @@ void CopyPipeline::MaybeBatchCopyToSharedSSD(const std::vector& key } } +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; @@ -112,6 +119,7 @@ void CopyPipeline::CopyWorkerLoop() { batch.push_back(std::move(copy_queue_.front().key)); copy_queue_.pop_front(); } + ++in_flight_copies_; } if (batch.size() == 1) { @@ -119,6 +127,12 @@ void CopyPipeline::CopyWorkerLoop() { } 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(); + } } } diff --git a/src/umbp/standalone/bin/standalone_server_main.cpp b/src/umbp/standalone/bin/standalone_server_main.cpp new file mode 100644 index 000000000..6cf402f77 --- /dev/null +++ b/src/umbp/standalone/bin/standalone_server_main.cpp @@ -0,0 +1,262 @@ +// 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 +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/standalone/ipc.h" +#include "umbp/standalone/standalone_server.h" + +namespace { + +std::optional EnvString(const char* name) { + const char* value = std::getenv(name); + if (!value || value[0] == '\0') return std::nullopt; + return std::string(value); +} + +bool ParseSizeEnv(const char* name, size_t* out, std::string* error) { + auto raw = EnvString(name); + if (!raw.has_value()) return true; + try { + *out = static_cast(std::stoull(*raw)); + return true; + } catch (const std::exception& exc) { + *error = std::string("invalid ") + name + ": " + exc.what(); + return false; + } +} + +bool ParseUint16Env(const char* name, uint16_t* out, std::string* error) { + auto raw = EnvString(name); + if (!raw.has_value()) return true; + try { + unsigned long value = std::stoul(*raw); + if (value > 65535) { + *error = std::string(name) + " must be <= 65535"; + return false; + } + *out = static_cast(value); + return true; + } catch (const std::exception& exc) { + *error = std::string("invalid ") + name + ": " + exc.what(); + return false; + } +} + +bool ParseIntEnv(const char* name, int* out, std::string* error) { + auto raw = EnvString(name); + if (!raw.has_value()) return true; + try { + *out = std::stoi(*raw); + return true; + } catch (const std::exception& exc) { + *error = std::string("invalid ") + name + ": " + exc.what(); + return false; + } +} + +bool ParseBoolEnv(const char* name, bool* out, std::string* error) { + auto raw = EnvString(name); + if (!raw.has_value()) return true; + std::string value = *raw; + for (char& ch : value) ch = static_cast(std::tolower(ch)); + if (value == "1" || value == "true" || value == "yes" || value == "on") { + *out = true; + return true; + } + if (value == "0" || value == "false" || value == "no" || value == "off") { + *out = false; + return true; + } + *error = std::string("invalid boolean ") + name + "=" + *raw; + return false; +} + +bool AnyEnv(const std::vector& names) { + for (const char* name : names) { + if (EnvString(name).has_value()) return true; + } + return false; +} + +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", + }; + + *distributed_requested = AnyEnv(kRequiredDistributedEnv); + if (!*distributed_requested) { + config->distributed.reset(); + return true; + } + + auto master_address = EnvString("UMBP_MASTER_ADDRESS"); + 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; + } + + 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; + + 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_SSD_STAGING_BUFFER_SIZE", &dist.ssd_staging_buffer_size, + error)) { + return false; + } + if (!ParseIntEnv("UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SLOTS", &dist.ssd_staging_buffer_slots, + error)) { + return false; + } + if (dist.ssd_staging_buffer_slots <= 0) { + *error = "UMBP_DISTRIBUTED_SSD_STAGING_BUFFER_SLOTS must be > 0"; + return false; + } + if (!ParseBoolEnv("UMBP_DISTRIBUTED_CACHE_REMOTE_FETCHES", &dist.cache_remote_fetches, 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); + + config->distributed = dist; + return true; +} + +} // namespace + +int main(int argc, char** argv) { + mori::umbp::UMBPConfig config = mori::umbp::UMBPConfig::FromEnvironment(); + config.role = mori::umbp::UMBPRole::Standalone; + config.follower_mode = false; + config.force_ssd_copy_on_write = false; + config.standalone_process.reset(); + + bool distributed_requested = false; + std::string config_error; + if (!ApplyDistributedBackendConfigFromEnv(&config, &distributed_requested, &config_error)) { + MORI_UMBP_ERROR("[StandaloneServer] {}", config_error); + return 1; + } + if (!config.Validate(&config_error)) { + MORI_UMBP_ERROR("[StandaloneServer] invalid backend config: {}", config_error); + return 1; + } + + std::string address; + if (argc > 1 && argv[1] && argv[1][0] != '\0') { + address = argv[1]; + } else if (const char* env = std::getenv("UMBP_STANDALONE_ADDRESS")) { + address = env; + } else { + address = mori::umbp::standalone::DefaultStandaloneAddress(); + } + + std::unique_ptr server; + try { + server = std::make_unique(config, address); + } catch (const std::exception& exc) { + MORI_UMBP_ERROR("[StandaloneServer] backend initialization failed: {}", exc.what()); + return 1; + } + + sigset_t signal_set; + sigemptyset(&signal_set); + sigaddset(&signal_set, SIGINT); + sigaddset(&signal_set, SIGTERM); + + const int block_rc = pthread_sigmask(SIG_BLOCK, &signal_set, nullptr); + if (block_rc != 0) { + MORI_UMBP_ERROR("[StandaloneServer] failed to block signals: {}", std::strerror(block_rc)); + return 1; + } + + if (!server->Start()) { + MORI_UMBP_ERROR("[StandaloneServer] failed to start on {}", address); + return 1; + } + + std::atomic stop_signal_waiter{false}; + std::thread signal_waiter([server = server.get(), &signal_set, &stop_signal_waiter]() { + while (!stop_signal_waiter.load()) { + timespec timeout{}; + timeout.tv_sec = 1; + timeout.tv_nsec = 0; + const int signum = sigtimedwait(&signal_set, nullptr, &timeout); + if (signum == SIGINT || signum == SIGTERM) { + MORI_UMBP_INFO("[StandaloneServer] caught signal {}, shutting down", signum); + server->Shutdown(); + return; + } + if (signum == -1 && errno != EAGAIN && errno != EINTR) { + MORI_UMBP_ERROR("[StandaloneServer] sigtimedwait failed: {}", std::strerror(errno)); + return; + } + } + }); + + MORI_UMBP_INFO("[StandaloneServer] running on {} backend={}", address, + distributed_requested ? "distributed" : "local"); + server->Run(); + + stop_signal_waiter = true; + signal_waiter.join(); + MORI_UMBP_INFO("[StandaloneServer] exited cleanly"); + return 0; +} diff --git a/src/umbp/standalone/external_kv_identity_client.cpp b/src/umbp/standalone/external_kv_identity_client.cpp new file mode 100644 index 000000000..715db360f --- /dev/null +++ b/src/umbp/standalone/external_kv_identity_client.cpp @@ -0,0 +1,300 @@ +// 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/standalone/external_kv_identity_client.h" + +#include + +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp.grpc.pb.h" + +namespace mori::umbp::standalone { +namespace { + +::umbp::UMBPMaster::Stub* Stub(void* ptr) { return static_cast<::umbp::UMBPMaster::Stub*>(ptr); } + +::umbp::TierType TierToProto(TierType tier) { + switch (tier) { + 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; + } +} + +TierType TierFromProto(::umbp::TierType tier) { + switch (tier) { + case ::umbp::TIER_HBM: + return TierType::HBM; + case ::umbp::TIER_DRAM: + return TierType::DRAM; + case ::umbp::TIER_SSD: + return TierType::SSD; + default: + return TierType::UNKNOWN; + } +} + +int RpcShutdownTimeoutMs() { + const char* raw = std::getenv("UMBP_RPC_SHUTDOWN_TIMEOUT_MS"); + if (!raw || raw[0] == '\0') return 3000; + const int parsed = std::atoi(raw); + return parsed > 0 ? parsed : 3000; +} + +void SetDeadline(grpc::ClientContext* ctx) { + ctx->set_deadline(std::chrono::system_clock::now() + + std::chrono::milliseconds(RpcShutdownTimeoutMs())); +} + +} // namespace + +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); + auto channel = + grpc::CreateCustomChannel(config_.master_address, grpc::InsecureChannelCredentials(), args); + channel_ = channel; + stub_.reset(::umbp::UMBPMaster::NewStub(channel).release()); +} + +ExternalKvIdentityClient::~ExternalKvIdentityClient() { Stop(); } + +bool ExternalKvIdentityClient::Start() { + if (config_.master_address.empty() || config_.node_id.empty() || config_.node_address.empty()) { + MORI_UMBP_WARN("[ExternalKvIdentity] invalid config: master/node identity is empty"); + return false; + } + + { + std::lock_guard lock(rpc_mu_); + if (!RegisterLocked()) return false; + } + + running_.store(true); + heartbeat_thread_ = std::thread(&ExternalKvIdentityClient::HeartbeatLoop, this); + return true; +} + +void ExternalKvIdentityClient::Stop() { + bool was_running = running_.exchange(false); + if (was_running) { + cv_.notify_one(); + if (heartbeat_thread_.joinable()) heartbeat_thread_.join(); + } + + std::lock_guard lock(rpc_mu_); + UnregisterLocked(); +} + +bool ExternalKvIdentityClient::RegisterLocked() { + ::umbp::RegisterClientRequest req; + req.set_node_id(config_.node_id); + req.set_node_address(config_.node_address); + req.set_peer_address(config_.peer_address); + req.set_engine_desc(config_.engine_desc_bytes.data(), config_.engine_desc_bytes.size()); + for (const auto& tag : config_.tags) req.add_tags(tag); + + ::umbp::RegisterClientResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->RegisterClient(&ctx, req, &resp); + if (!status.ok() && status.error_code() != grpc::StatusCode::ALREADY_EXISTS) { + MORI_UMBP_WARN("[ExternalKvIdentity] RegisterClient failed node_id={} error={}", + config_.node_id, status.error_message()); + registered_ = false; + return false; + } + if (resp.heartbeat_interval_ms() > 0) heartbeat_interval_ms_ = resp.heartbeat_interval_ms(); + registered_ = true; + MORI_UMBP_INFO("[ExternalKvIdentity] registered node_id={} peer={}", config_.node_id, + config_.peer_address); + return true; +} + +void ExternalKvIdentityClient::UnregisterLocked() { + if (!registered_) return; + ::umbp::UnregisterClientRequest req; + req.set_node_id(config_.node_id); + ::umbp::UnregisterClientResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->UnregisterClient(&ctx, req, &resp); + if (!status.ok()) { + MORI_UMBP_WARN("[ExternalKvIdentity] UnregisterClient failed node_id={} error={}", + config_.node_id, status.error_message()); + } + registered_ = false; +} + +bool ExternalKvIdentityClient::SendHeartbeatOnceLocked() { + if (!registered_) return false; + + ::umbp::HeartbeatRequest req; + req.set_node_id(config_.node_id); + req.set_is_full_sync(false); + + ::umbp::HeartbeatResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->Heartbeat(&ctx, req, &resp); + if (!status.ok()) { + MORI_UMBP_WARN("[ExternalKvIdentity] Heartbeat failed node_id={} error={}", config_.node_id, + status.error_message()); + return false; + } + if (resp.status() == ::umbp::CLIENT_STATUS_UNKNOWN) { + MORI_UMBP_WARN("[ExternalKvIdentity] master forgot node_id={}, re-registering", + config_.node_id); + registered_ = false; + return RegisterLocked(); + } + return true; +} + +void ExternalKvIdentityClient::HeartbeatLoop() { + while (running_.load()) { + std::unique_lock lock(cv_mu_); + cv_.wait_for(lock, std::chrono::milliseconds(heartbeat_interval_ms_), + [this] { return !running_.load(); }); + lock.unlock(); + if (!running_.load()) break; + + std::lock_guard rpc_lock(rpc_mu_); + SendHeartbeatOnceLocked(); + } +} + +bool ExternalKvIdentityClient::ReportExternalKvBlocks(const std::vector& hashes, + TierType tier) { + if (hashes.empty()) return true; + std::lock_guard lock(rpc_mu_); + if (!registered_) return false; + ::umbp::ReportExternalKvBlocksRequest req; + req.set_node_id(config_.node_id); + req.set_tier(TierToProto(tier)); + for (const auto& hash : hashes) req.add_hashes(hash); + ::umbp::ReportExternalKvBlocksResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->ReportExternalKvBlocks(&ctx, req, &resp); + return status.ok(); +} + +bool ExternalKvIdentityClient::RevokeExternalKvBlocks(const std::vector& hashes, + TierType tier) { + if (hashes.empty()) return true; + std::lock_guard lock(rpc_mu_); + if (!registered_) return false; + ::umbp::RevokeExternalKvBlocksRequest req; + req.set_node_id(config_.node_id); + req.set_tier(TierToProto(tier)); + for (const auto& hash : hashes) req.add_hashes(hash); + ::umbp::RevokeExternalKvBlocksResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->RevokeExternalKvBlocks(&ctx, req, &resp); + return status.ok(); +} + +bool ExternalKvIdentityClient::RevokeAllExternalKvBlocksAtTier(TierType tier) { + std::lock_guard lock(rpc_mu_); + if (!registered_) return false; + ::umbp::RevokeAllExternalKvBlocksAtTierRequest req; + req.set_node_id(config_.node_id); + req.set_tier(TierToProto(tier)); + ::umbp::RevokeAllExternalKvBlocksAtTierResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->RevokeAllExternalKvBlocksAtTier(&ctx, req, &resp); + return status.ok(); +} + +std::vector ExternalKvIdentityClient::MatchExternalKv( + const std::vector& hashes, bool count_as_hit) { + std::vector out; + if (hashes.empty()) return out; + std::lock_guard lock(rpc_mu_); + if (!registered_) return out; + + ::umbp::MatchExternalKvRequest req; + for (const auto& hash : hashes) req.add_hashes(hash); + req.set_count_as_hit(count_as_hit); + ::umbp::MatchExternalKvResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->MatchExternalKv(&ctx, req, &resp); + if (!status.ok()) return out; + + out.reserve(resp.matches_size()); + for (const auto& m : resp.matches()) { + IUMBPClient::ExternalKvMatch match; + match.node_id = m.node_id(); + match.peer_address = m.peer_address(); + for (const auto& bucket : m.hashes_by_tier()) { + std::vector values(bucket.hashes().begin(), bucket.hashes().end()); + match.hashes_by_tier[TierFromProto(bucket.tier())] = std::move(values); + } + out.push_back(std::move(match)); + } + return out; +} + +std::vector ExternalKvIdentityClient::GetExternalKvHitCounts( + const std::vector& hashes) { + std::vector out; + if (hashes.empty()) return out; + std::lock_guard lock(rpc_mu_); + if (!registered_) return out; + + ::umbp::GetExternalKvHitCountsRequest req; + for (const auto& hash : hashes) req.add_hashes(hash); + ::umbp::GetExternalKvHitCountsResponse resp; + grpc::ClientContext ctx; + SetDeadline(&ctx); + auto status = Stub(stub_.get())->GetExternalKvHitCounts(&ctx, req, &resp); + if (!status.ok()) return out; + + out.reserve(resp.entries_size()); + for (const auto& e : resp.entries()) { + IUMBPClient::ExternalKvHitCountEntry entry; + entry.hash = e.hash(); + entry.hit_count_total = e.hit_count_total(); + out.push_back(std::move(entry)); + } + return out; +} + +} // namespace mori::umbp::standalone diff --git a/src/umbp/standalone/ipc.cpp b/src/umbp/standalone/ipc.cpp new file mode 100644 index 000000000..c18561aea --- /dev/null +++ b/src/umbp/standalone/ipc.cpp @@ -0,0 +1,297 @@ +// 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/standalone/ipc.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace mori::umbp::standalone { +namespace { + +bool StartsWith(const std::string& value, const char* prefix) { + return value.rfind(prefix, 0) == 0; +} + +void SetError(std::string* error, const std::string& message) { + if (error) *error = message; +} + +std::string ErrnoMessage(const std::string& op) { return op + ": " + std::strerror(errno); } + +bool FillSockaddr(const std::string& path, sockaddr_un* addr, socklen_t* addr_len, + std::string* error) { + if (!addr || !addr_len) return false; + if (path.empty()) { + SetError(error, "empty UDS path"); + return false; + } + if (path.size() >= sizeof(addr->sun_path)) { + SetError(error, "UDS path too long: " + path); + return false; + } + std::memset(addr, 0, sizeof(*addr)); + addr->sun_family = AF_UNIX; + std::strncpy(addr->sun_path, path.c_str(), sizeof(addr->sun_path) - 1); + *addr_len = static_cast(sizeof(sa_family_t) + path.size() + 1); + return true; +} + +bool SendAll(int fd, const void* data, size_t len, std::string* error) { + const char* pos = static_cast(data); + size_t left = len; + while (left > 0) { + ssize_t n = send(fd, pos, left, MSG_NOSIGNAL); + if (n < 0 && errno == EINTR) continue; + if (n <= 0) { + SetError(error, ErrnoMessage("send")); + return false; + } + pos += n; + left -= static_cast(n); + } + return true; +} + +bool RecvAll(int fd, void* data, size_t len, std::string* error) { + char* pos = static_cast(data); + size_t left = len; + while (left > 0) { + ssize_t n = recv(fd, pos, left, 0); + if (n < 0 && errno == EINTR) continue; + if (n <= 0) { + SetError(error, ErrnoMessage("recv")); + return false; + } + pos += n; + left -= static_cast(n); + } + return true; +} + +bool SetSocketTimeouts(int fd, int timeout_ms, std::string* error) { + if (timeout_ms <= 0) return true; + timeval tv; + tv.tv_sec = timeout_ms / 1000; + tv.tv_usec = (timeout_ms % 1000) * 1000; + if (setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)) != 0) { + SetError(error, ErrnoMessage("setsockopt(SO_SNDTIMEO)")); + return false; + } + if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) != 0) { + SetError(error, ErrnoMessage("setsockopt(SO_RCVTIMEO)")); + return false; + } + return true; +} + +} // namespace + +std::string UnixPathFromGrpcAddress(const std::string& address) { + if (StartsWith(address, "unix://")) return address.substr(7); + if (StartsWith(address, "unix:")) return address.substr(5); + return address; +} + +std::string DeriveFdSocketPath(const std::string& grpc_address) { + std::string path = UnixPathFromGrpcAddress(grpc_address); + const std::string suffix = ".grpc.sock"; + if (path.size() >= suffix.size() && + path.compare(path.size() - suffix.size(), suffix.size(), suffix) == 0) { + path.replace(path.size() - suffix.size(), suffix.size(), ".fd.sock"); + } else { + path += ".fd.sock"; + } + return path; +} + +std::string DefaultStandaloneAddress() { + const char* node = std::getenv("UMBP_NODE_ID"); + std::string node_id = (node && node[0] != '\0') ? node : "node0"; + return "unix:///run/umbp/standalone/" + node_id + ".grpc.sock"; +} + +bool EnsureParentDirectory(const std::string& path, std::string* error) { + size_t slash = path.rfind('/'); + if (slash == std::string::npos || slash == 0) return true; + + std::string parent = path.substr(0, slash); + size_t pos = 1; + while (pos <= parent.size()) { + size_t next = parent.find('/', pos); + std::string part = parent.substr(0, next == std::string::npos ? parent.size() : next); + if (!part.empty() && mkdir(part.c_str(), 0700) != 0 && errno != EEXIST) { + SetError(error, ErrnoMessage("mkdir(" + part + ")")); + return false; + } + if (next == std::string::npos) break; + pos = next + 1; + } + return true; +} + +int SendFdRegistration(const std::string& socket_path, int fd, const std::string& client_id, + uintptr_t worker_base, size_t size, int timeout_ms, std::string* error) { + if (client_id.empty() || client_id.size() >= kClientIdBytes) { + SetError(error, "client_id is empty or too long"); + return -1; + } + + int sock = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (sock < 0) { + SetError(error, ErrnoMessage("socket")); + return -1; + } + if (!SetSocketTimeouts(sock, timeout_ms, error)) { + close(sock); + return -1; + } + + sockaddr_un addr; + socklen_t addr_len = 0; + if (!FillSockaddr(socket_path, &addr, &addr_len, error)) { + close(sock); + return -1; + } + if (connect(sock, reinterpret_cast(&addr), addr_len) != 0) { + SetError(error, ErrnoMessage("connect(" + socket_path + ")")); + close(sock); + return -1; + } + + FdRegistrationMessage msg; + std::strncpy(msg.client_id, client_id.c_str(), sizeof(msg.client_id) - 1); + msg.worker_base = static_cast(worker_base); + msg.size = static_cast(size); + + msghdr hdr; + std::memset(&hdr, 0, sizeof(hdr)); + iovec iov; + iov.iov_base = &msg; + iov.iov_len = sizeof(msg); + hdr.msg_iov = &iov; + hdr.msg_iovlen = 1; + + char control[CMSG_SPACE(sizeof(int))]; + std::memset(control, 0, sizeof(control)); + hdr.msg_control = control; + hdr.msg_controllen = sizeof(control); + cmsghdr* cmsg = CMSG_FIRSTHDR(&hdr); + cmsg->cmsg_level = SOL_SOCKET; + cmsg->cmsg_type = SCM_RIGHTS; + cmsg->cmsg_len = CMSG_LEN(sizeof(int)); + std::memcpy(CMSG_DATA(cmsg), &fd, sizeof(int)); + + while (true) { + ssize_t sent = sendmsg(sock, &hdr, 0); + if (sent < 0 && errno == EINTR) continue; + if (sent != static_cast(sizeof(msg))) { + SetError(error, sent < 0 ? ErrnoMessage("sendmsg") : "short sendmsg"); + close(sock); + return -1; + } + break; + } + + int32_t status = -1; + if (!RecvStatus(sock, &status, error)) { + close(sock); + return -1; + } + close(sock); + return status; +} + +int RecvFdRegistration(int socket_fd, FdRegistrationMessage* message, std::string* error) { + if (!message) return -1; + msghdr hdr; + std::memset(&hdr, 0, sizeof(hdr)); + iovec iov; + iov.iov_base = message; + iov.iov_len = sizeof(*message); + hdr.msg_iov = &iov; + hdr.msg_iovlen = 1; + + char control[CMSG_SPACE(sizeof(int))]; + std::memset(control, 0, sizeof(control)); + hdr.msg_control = control; + hdr.msg_controllen = sizeof(control); + + ssize_t received = 0; + while (true) { + received = recvmsg(socket_fd, &hdr, 0); + if (received < 0 && errno == EINTR) continue; + break; + } + auto close_delivered_fd = [&]() { + cmsghdr* delivered = CMSG_FIRSTHDR(&hdr); + if (delivered && delivered->cmsg_level == SOL_SOCKET && delivered->cmsg_type == SCM_RIGHTS) { + int delivered_fd = -1; + std::memcpy(&delivered_fd, CMSG_DATA(delivered), sizeof(int)); + if (delivered_fd >= 0) close(delivered_fd); + } + }; + if (received != static_cast(sizeof(*message)) || + (hdr.msg_flags & (MSG_TRUNC | MSG_CTRUNC))) { + close_delivered_fd(); + SetError(error, received < 0 ? ErrnoMessage("recvmsg") : "malformed fd registration message"); + return -1; + } + + cmsghdr* cmsg = CMSG_FIRSTHDR(&hdr); + if (!cmsg || cmsg->cmsg_level != SOL_SOCKET || cmsg->cmsg_type != SCM_RIGHTS) { + close_delivered_fd(); + SetError(error, "fd registration message missing SCM_RIGHTS fd"); + return -1; + } + int received_fd = -1; + std::memcpy(&received_fd, CMSG_DATA(cmsg), sizeof(int)); + if (message->magic != kFdRegistrationMagic || message->version != kFdRegistrationVersion || + message->client_id[kClientIdBytes - 1] != '\0' || message->size == 0) { + close(received_fd); + SetError(error, "invalid fd registration payload"); + return -1; + } + return received_fd; +} + +bool SendStatus(int socket_fd, int32_t status, std::string* error) { + return SendAll(socket_fd, &status, sizeof(status), error); +} + +bool RecvStatus(int socket_fd, int32_t* status, std::string* error) { + if (!status) return false; + return RecvAll(socket_fd, status, sizeof(*status), error); +} + +} // namespace mori::umbp::standalone diff --git a/src/umbp/standalone/standalone_process_client.cpp b/src/umbp/standalone/standalone_process_client.cpp new file mode 100644 index 000000000..52c64adba --- /dev/null +++ b/src/umbp/standalone/standalone_process_client.cpp @@ -0,0 +1,608 @@ +// 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/standalone/standalone_process_client.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/local/host_mem_allocator.h" +#include "umbp/standalone/ipc.h" + +namespace mori::umbp::standalone { +namespace { + +std::atomic g_client_counter{0}; + +::umbp::TierType TierToProto(TierType tier) { + switch (tier) { + 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; + } +} + +TierType TierFromProto(::umbp::TierType tier) { + switch (tier) { + case ::umbp::TIER_HBM: + return TierType::HBM; + case ::umbp::TIER_DRAM: + return TierType::DRAM; + case ::umbp::TIER_SSD: + return TierType::SSD; + default: + return TierType::UNKNOWN; + } +} + +bool IsLocalRankZero() { + for (const char* name : + {"LOCAL_RANK", "OMPI_COMM_WORLD_LOCAL_RANK", "SLURM_LOCALID", "MPI_LOCALRANKID"}) { + const char* value = std::getenv(name); + if (value) return std::atoi(value) == 0; + } + return true; +} + +std::string BootstrapLockPath() { + const char* dir = std::getenv("UMBP_STANDALONE_SHM_DIR"); + std::string base = (dir && dir[0] != '\0') ? dir : "/tmp"; + if (!base.empty() && base.back() == '/') base.pop_back(); + return base + "/umbp_standalone_bootstrap.lock"; +} + +std::string FindStandaloneServerBinary() { + const char* env = std::getenv("UMBP_STANDALONE_BIN"); + return (env && env[0] != '\0') ? env : "umbp_standalone_server"; +} + +void SetEnv(const char* name, const std::string& value) { + if (!value.empty()) setenv(name, value.c_str(), 1); +} + +void SetEnv(const char* name, size_t value) { setenv(name, std::to_string(value).c_str(), 1); } + +void SetEnv(const char* name, int value) { setenv(name, std::to_string(value).c_str(), 1); } + +void SetEnv(const char* name, bool value) { setenv(name, value ? "1" : "0", 1); } + +void SetEnv(const char* name, double value) { setenv(name, std::to_string(value).c_str(), 1); } + +void ExportServerEnv(const UMBPConfig& config, const std::string& address) { + SetEnv("UMBP_STANDALONE_ADDRESS", address); + SetEnv("UMBP_ROLE", "standalone"); + SetEnv("UMBP_DRAM_CAPACITY", config.dram.capacity_bytes); + SetEnv("UMBP_DRAM_USE_HUGEPAGES", config.dram.use_hugepages); + SetEnv("UMBP_DRAM_HUGEPAGE_SIZE", config.dram.hugepage_size); + SetEnv("UMBP_DRAM_NUMA_NODE", config.dram.numa_node); + SetEnv("UMBP_DRAM_PREFAULT", config.dram.prefault); + SetEnv("UMBP_DRAM_HIGH_WM", config.dram.high_watermark); + SetEnv("UMBP_DRAM_LOW_WM", config.dram.low_watermark); + SetEnv("UMBP_SSD_ENABLED", config.ssd.enabled); + SetEnv("UMBP_SSD_DIR", config.ssd.storage_dir); + SetEnv("UMBP_SSD_CAPACITY", config.ssd.capacity_bytes); + SetEnv("UMBP_SSD_BACKEND", config.ssd.ssd_backend); + SetEnv("UMBP_SSD_HIGH_WM", config.ssd.high_watermark); + SetEnv("UMBP_SSD_LOW_WM", config.ssd.low_watermark); + SetEnv("UMBP_EVICTION_POLICY", config.eviction.policy); + SetEnv("UMBP_SPDK_BDEV", config.ssd.spdk_bdev_name); + SetEnv("UMBP_SPDK_REACTOR_MASK", config.ssd.spdk_reactor_mask); + SetEnv("UMBP_SPDK_MEM_MB", config.ssd.spdk_mem_size_mb); + SetEnv("UMBP_SPDK_NVME_PCI", config.ssd.spdk_nvme_pci_addr); + SetEnv("UMBP_SPDK_NVME_CTRL", config.ssd.spdk_nvme_ctrl_name); + SetEnv("UMBP_SPDK_IO_WORKERS", config.ssd.spdk_io_workers); + SetEnv("UMBP_SPDK_PROXY_SHM", config.ssd.spdk_proxy_shm_name); + SetEnv("UMBP_SPDK_PROXY_BIN", config.ssd.spdk_proxy_bin); + SetEnv("UMBP_SPDK_PROXY_TENANT_ID", static_cast(config.ssd.spdk_proxy_tenant_id)); + SetEnv("UMBP_SPDK_PROXY_TENANT_QUOTA_BYTES", config.ssd.spdk_proxy_tenant_quota_bytes); + SetEnv("UMBP_SPDK_PROXY_MAX_CHANNELS", static_cast(config.ssd.spdk_proxy_max_channels)); + SetEnv("UMBP_SPDK_PROXY_DATA_PER_CHANNEL_MB", config.ssd.spdk_proxy_data_per_channel_mb); + SetEnv("UMBP_SPDK_PROXY_TIMEOUT_MS", config.ssd.spdk_proxy_startup_timeout_ms); + SetEnv("UMBP_SPDK_PROXY_AUTO_START", config.ssd.spdk_proxy_auto_start); + SetEnv("UMBP_SPDK_PROXY_IDLE_EXIT_TIMEOUT_MS", config.ssd.spdk_proxy_idle_exit_timeout_ms); + SetEnv("UMBP_SPDK_PROXY_ALLOW_BORROW", config.ssd.spdk_proxy_allow_borrow); + SetEnv("UMBP_SPDK_PROXY_RESERVED_SHARED_BYTES", config.ssd.spdk_proxy_reserved_shared_bytes); +} + +class ScopedBootstrapLock { + public: + ScopedBootstrapLock() { + std::string path = BootstrapLockPath(); + fd_ = open(path.c_str(), O_CREAT | O_RDWR, 0600); + if (fd_ >= 0 && flock(fd_, LOCK_EX) != 0) { + close(fd_); + fd_ = -1; + } + } + + ~ScopedBootstrapLock() { + if (fd_ >= 0) { + flock(fd_, LOCK_UN); + close(fd_); + } + } + + bool valid() const { return fd_ >= 0; } + + private: + int fd_ = -1; +}; + +} // namespace + +StandaloneProcessClient::StandaloneProcessClient(const UMBPConfig& config) : config_(config) { + if (!config_.standalone_process.has_value()) { + throw std::runtime_error("StandaloneProcessClient requires UMBPConfig::standalone_process"); + } + standalone_config_ = config_.standalone_process.value(); + std::string error_message; + if (!config_.Validate(&error_message)) { + throw std::runtime_error("invalid UMBP config: " + error_message); + } + + address_ = standalone_config_.address; + fd_socket_path_ = DeriveFdSocketPath(address_); + channel_ = grpc::CreateChannel(address_, grpc::InsecureChannelCredentials()); + stub_ = ::umbp::UMBPStandalone::NewStub(channel_); + + MaybeAutoStart(); + if (!WaitReady(standalone_config_.startup_timeout_ms)) { + throw std::runtime_error("StandaloneProcessClient: server is not ready at " + address_); + } + + MORI_UMBP_INFO("[StandaloneProcessClient] connected address={} fd_socket={}", address_, + fd_socket_path_); +} + +StandaloneProcessClient::~StandaloneProcessClient() { Close(); } + +std::string StandaloneProcessClient::ClientId() { + std::lock_guard lock(registration_mu_); + if (!client_id_.empty()) return client_id_; + std::ostringstream oss; + oss << "umbp-" << getpid() << "-" << g_client_counter.fetch_add(1); + client_id_ = oss.str(); + return client_id_; +} + +bool StandaloneProcessClient::WaitReady(int timeout_ms) const { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); + while (std::chrono::steady_clock::now() < deadline) { + grpc::ClientContext ctx; + ctx.set_deadline(std::chrono::system_clock::now() + std::chrono::milliseconds(500)); + ::umbp::Empty req; + ::umbp::PingResponse resp; + grpc::Status status = stub_->Ping(&ctx, req, &resp); + if (status.ok() && resp.ready()) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + return false; +} + +void StandaloneProcessClient::MaybeAutoStart() { + if (WaitReady(200)) return; + if (!standalone_config_.auto_start) return; + + ScopedBootstrapLock lock; + if (!lock.valid()) { + MORI_UMBP_WARN("[StandaloneProcessClient] bootstrap lock unavailable; waiting for server"); + return; + } + if (WaitReady(200)) return; + if (!IsLocalRankZero()) return; + + std::string bin = FindStandaloneServerBinary(); + pid_t pid = fork(); + if (pid < 0) { + throw std::runtime_error("StandaloneProcessClient: fork() failed: " + + std::string(std::strerror(errno))); + } + if (pid == 0) { + setsid(); + ExportServerEnv(config_, address_); + execlp(bin.c_str(), "umbp_standalone_server", address_.c_str(), static_cast(nullptr)); + fprintf(stderr, "[UMBP ERROR] execlp('%s') failed: %s\n", bin.c_str(), std::strerror(errno)); + _exit(127); + } + MORI_UMBP_INFO( + "[StandaloneProcessClient] spawned umbp_standalone_server pid={} bin={} address={}", pid, bin, + address_); +} + +bool StandaloneProcessClient::OffsetFor(uintptr_t ptr, size_t size, uint64_t* offset) const { + std::lock_guard lock(registration_mu_); + if (!registered_ || ptr < registered_base_) return false; + uintptr_t rel = ptr - registered_base_; + if (rel > registered_size_ || size > registered_size_ - rel) return false; + *offset = static_cast(rel); + return true; +} + +bool StandaloneProcessClient::Put(const std::string& key, uintptr_t src, size_t size) { + if (closing_) return false; + std::shared_lock lk(op_mutex_); + if (closed_) return false; + uint64_t offset = 0; + if (!OffsetFor(src, size, &offset)) return false; + grpc::ClientContext ctx; + ::umbp::PutRequest req; + req.set_key(key); + req.set_client_id(ClientId()); + req.set_shm_offset(offset); + req.set_size(size); + ::umbp::BoolResponse resp; + grpc::Status status = stub_->Put(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +bool StandaloneProcessClient::Get(const std::string& key, uintptr_t dst, size_t size) { + if (closing_) return false; + std::shared_lock lk(op_mutex_); + if (closed_) return false; + uint64_t offset = 0; + if (!OffsetFor(dst, size, &offset)) return false; + grpc::ClientContext ctx; + ::umbp::GetRequest req; + req.set_key(key); + req.set_client_id(ClientId()); + req.set_shm_offset(offset); + req.set_size(size); + ::umbp::BoolResponse resp; + grpc::Status status = stub_->Get(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +bool StandaloneProcessClient::Exists(const std::string& key) const { + if (closing_) return false; + std::shared_lock lk(op_mutex_); + if (closed_) return false; + grpc::ClientContext ctx; + ::umbp::KeyRequest req; + req.set_key(key); + ::umbp::BoolResponse resp; + grpc::Status status = stub_->Exists(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +std::vector StandaloneProcessClient::BatchPut(const std::vector& keys, + const std::vector& srcs, + const std::vector& sizes) { + if (closing_) return std::vector(keys.size(), false); + std::shared_lock lk(op_mutex_); + if (closed_ || keys.size() != srcs.size() || keys.size() != sizes.size()) { + return std::vector(keys.size(), false); + } + ::umbp::BatchDataRequest req; + req.set_client_id(ClientId()); + for (size_t i = 0; i < keys.size(); ++i) { + uint64_t offset = 0; + if (!OffsetFor(srcs[i], sizes[i], &offset)) return std::vector(keys.size(), false); + req.add_keys(keys[i]); + req.add_shm_offsets(offset); + req.add_sizes(sizes[i]); + } + grpc::ClientContext ctx; + ::umbp::BatchBoolResponse resp; + grpc::Status status = stub_->BatchPut(&ctx, req, &resp); + if (!status.ok() || resp.ok_size() != static_cast(keys.size())) { + return std::vector(keys.size(), false); + } + return std::vector(resp.ok().begin(), resp.ok().end()); +} + +std::vector StandaloneProcessClient::BatchPutWithDepth(const std::vector& keys, + const std::vector& srcs, + const std::vector& sizes, + const std::vector& depths) { + if (closing_) return std::vector(keys.size(), false); + std::shared_lock lk(op_mutex_); + if (closed_ || keys.size() != srcs.size() || keys.size() != sizes.size()) { + return std::vector(keys.size(), false); + } + ::umbp::BatchDataWithDepthRequest req; + req.set_client_id(ClientId()); + for (size_t i = 0; i < keys.size(); ++i) { + uint64_t offset = 0; + if (!OffsetFor(srcs[i], sizes[i], &offset)) return std::vector(keys.size(), false); + req.add_keys(keys[i]); + req.add_shm_offsets(offset); + req.add_sizes(sizes[i]); + req.add_depths(i < depths.size() ? depths[i] : -1); + } + grpc::ClientContext ctx; + ::umbp::BatchBoolResponse resp; + grpc::Status status = stub_->BatchPutWithDepth(&ctx, req, &resp); + if (!status.ok() || resp.ok_size() != static_cast(keys.size())) { + return std::vector(keys.size(), false); + } + return std::vector(resp.ok().begin(), resp.ok().end()); +} + +std::vector StandaloneProcessClient::BatchGet(const std::vector& keys, + const std::vector& dsts, + const std::vector& sizes) { + if (closing_) return std::vector(keys.size(), false); + std::shared_lock lk(op_mutex_); + if (closed_ || keys.size() != dsts.size() || keys.size() != sizes.size()) { + return std::vector(keys.size(), false); + } + ::umbp::BatchDataRequest req; + req.set_client_id(ClientId()); + for (size_t i = 0; i < keys.size(); ++i) { + uint64_t offset = 0; + if (!OffsetFor(dsts[i], sizes[i], &offset)) return std::vector(keys.size(), false); + req.add_keys(keys[i]); + req.add_shm_offsets(offset); + req.add_sizes(sizes[i]); + } + grpc::ClientContext ctx; + ::umbp::BatchBoolResponse resp; + grpc::Status status = stub_->BatchGet(&ctx, req, &resp); + if (!status.ok() || resp.ok_size() != static_cast(keys.size())) { + return std::vector(keys.size(), false); + } + 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_); + if (closed_) return std::vector(keys.size(), false); + grpc::ClientContext ctx; + ::umbp::BatchKeysRequest req; + for (const auto& key : keys) req.add_keys(key); + ::umbp::BatchBoolResponse resp; + grpc::Status status = stub_->BatchExists(&ctx, req, &resp); + if (!status.ok() || resp.ok_size() != static_cast(keys.size())) { + return std::vector(keys.size(), false); + } + return std::vector(resp.ok().begin(), resp.ok().end()); +} + +size_t StandaloneProcessClient::BatchExistsConsecutive(const std::vector& keys) const { + if (closing_) return 0; + std::shared_lock lk(op_mutex_); + if (closed_) return 0; + grpc::ClientContext ctx; + ::umbp::BatchKeysRequest req; + for (const auto& key : keys) req.add_keys(key); + ::umbp::CountResponse resp; + grpc::Status status = stub_->BatchExistsConsecutive(&ctx, req, &resp); + return status.ok() ? static_cast(resp.count()) : 0; +} + +bool StandaloneProcessClient::Clear() { + if (closing_) return true; + std::unique_lock lk(op_mutex_); + if (closed_) return true; + grpc::ClientContext ctx; + ::umbp::Empty req; + ::umbp::BoolResponse resp; + grpc::Status status = stub_->Clear(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +bool StandaloneProcessClient::Flush() { + if (closing_) return true; + std::shared_lock lk(op_mutex_); + if (closed_) return true; + grpc::ClientContext ctx; + ::umbp::Empty req; + ::umbp::BoolResponse resp; + grpc::Status status = stub_->Flush(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +void StandaloneProcessClient::Close() { + closing_ = true; + std::unique_lock lk(op_mutex_); + if (closed_) return; + try { + DeregisterMemoryLocked(); + } catch (...) { + } + closed_ = true; + stub_.reset(); + channel_.reset(); +} + +bool StandaloneProcessClient::RegisterMemory(uintptr_t ptr, size_t size) { + if (closing_) return false; + std::unique_lock lk(op_mutex_); + if (closed_) return false; + + auto allocation = HostMemAllocator::AcquireShmAllocation(ptr, size); + if (!allocation.has_value()) { + throw std::runtime_error( + "StandaloneProcessClient::RegisterMemory requires an AnonymousShm-backed host buffer"); + } + + bool acquired_kept = false; + const std::string client_id = ClientId(); + try { + std::string error; + int status = SendFdRegistration( + fd_socket_path_, allocation->fd, client_id, reinterpret_cast(allocation->base), + allocation->mapped_size, standalone_config_.startup_timeout_ms, &error); + if (status != 0) { + throw std::runtime_error("fd handoff failed: " + error); + } + + grpc::ClientContext ctx; + ::umbp::RegisterMemoryRequest req; + req.set_client_id(client_id); + req.set_worker_base(reinterpret_cast(allocation->base)); + req.set_size(allocation->mapped_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); + ::umbp::BoolResponse resp; + grpc::Status rpc_status = stub_->RegisterMemory(&ctx, req, &resp); + if (!rpc_status.ok() || !resp.ok()) { + throw std::runtime_error("standalone RegisterMemory RPC failed: " + + (rpc_status.ok() ? resp.error() : rpc_status.error_message())); + } + + std::lock_guard lock(registration_mu_); + if (registered_) HostMemAllocator::ReleaseShmAllocation(registered_base_); + registered_base_ = reinterpret_cast(allocation->base); + registered_size_ = allocation->mapped_size; + registered_ = true; + acquired_kept = true; + } catch (...) { + if (!acquired_kept) { + HostMemAllocator::ReleaseShmAllocation(reinterpret_cast(allocation->base)); + } + throw; + } + return true; +} + +void StandaloneProcessClient::DeregisterMemoryLocked() { + std::string client_id; + uintptr_t base = 0; + { + std::lock_guard lock(registration_mu_); + if (!registered_) return; + client_id = client_id_; + base = registered_base_; + registered_ = false; + registered_base_ = 0; + registered_size_ = 0; + } + + grpc::ClientContext ctx; + ::umbp::DeregisterMemoryRequest req; + req.set_client_id(client_id); + ::umbp::Empty resp; + if (stub_) stub_->DeregisterMemory(&ctx, req, &resp); + HostMemAllocator::ReleaseShmAllocation(base); +} + +void StandaloneProcessClient::DeregisterMemory(uintptr_t /*ptr*/) { + if (closing_) return; + std::unique_lock lk(op_mutex_); + if (closed_) return; + DeregisterMemoryLocked(); +} + +bool StandaloneProcessClient::ReportExternalKvBlocks(const std::vector& hashes, + TierType tier) { + grpc::ClientContext ctx; + ::umbp::StandaloneExternalKvMutationRequest req; + for (const auto& hash : hashes) req.add_hashes(hash); + req.set_tier(TierToProto(tier)); + req.set_client_id(ClientId()); + ::umbp::BoolResponse resp; + grpc::Status status = stub_->ReportExternalKvBlocks(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +bool StandaloneProcessClient::RevokeExternalKvBlocks(const std::vector& hashes, + TierType tier) { + grpc::ClientContext ctx; + ::umbp::StandaloneExternalKvMutationRequest req; + for (const auto& hash : hashes) req.add_hashes(hash); + req.set_tier(TierToProto(tier)); + req.set_client_id(ClientId()); + ::umbp::BoolResponse resp; + grpc::Status status = stub_->RevokeExternalKvBlocks(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +bool StandaloneProcessClient::RevokeAllExternalKvBlocksAtTier(TierType tier) { + grpc::ClientContext ctx; + ::umbp::StandaloneExternalKvTierRequest req; + req.set_tier(TierToProto(tier)); + req.set_client_id(ClientId()); + ::umbp::BoolResponse resp; + grpc::Status status = stub_->RevokeAllExternalKvBlocksAtTier(&ctx, req, &resp); + return status.ok() && resp.ok(); +} + +std::vector StandaloneProcessClient::MatchExternalKv( + const std::vector& hashes, bool count_as_hit) { + grpc::ClientContext ctx; + ::umbp::StandaloneMatchExternalKvRequest req; + for (const auto& hash : hashes) req.add_hashes(hash); + req.set_count_as_hit(count_as_hit); + req.set_client_id(ClientId()); + ::umbp::StandaloneMatchExternalKvResponse resp; + grpc::Status status = stub_->MatchExternalKv(&ctx, req, &resp); + if (!status.ok()) return {}; + + std::vector out; + out.reserve(resp.matches_size()); + for (const auto& m : resp.matches()) { + IUMBPClient::ExternalKvMatch match; + match.node_id = m.node_id(); + match.peer_address = m.peer_address(); + for (const auto& bucket : m.hashes_by_tier()) { + std::vector values(bucket.hashes().begin(), bucket.hashes().end()); + match.hashes_by_tier[TierFromProto(bucket.tier())] = std::move(values); + } + out.push_back(std::move(match)); + } + return out; +} + +std::vector StandaloneProcessClient::GetExternalKvHitCounts( + const std::vector& hashes) { + grpc::ClientContext ctx; + ::umbp::StandaloneExternalKvHitCountsRequest req; + for (const auto& hash : hashes) req.add_hashes(hash); + req.set_client_id(ClientId()); + ::umbp::StandaloneExternalKvHitCountsResponse resp; + grpc::Status status = stub_->GetExternalKvHitCounts(&ctx, req, &resp); + if (!status.ok()) return {}; + std::vector out; + out.reserve(resp.entries_size()); + for (const auto& e : resp.entries()) { + IUMBPClient::ExternalKvHitCountEntry entry; + entry.hash = e.hash(); + entry.hit_count_total = e.hit_count_total(); + out.push_back(std::move(entry)); + } + return out; +} + +} // namespace mori::umbp::standalone diff --git a/src/umbp/standalone/standalone_server.cpp b/src/umbp/standalone/standalone_server.cpp new file mode 100644 index 000000000..6eb90b3a8 --- /dev/null +++ b/src/umbp/standalone/standalone_server.cpp @@ -0,0 +1,848 @@ +// 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/standalone/standalone_server.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/standalone/external_kv_identity_client.h" +#include "umbp/standalone/ipc.h" +#include "umbp/umbp_client.h" +#include "umbp_standalone.grpc.pb.h" + +namespace mori::umbp::standalone { +namespace { + +std::chrono::seconds ShutdownDeadline() { + const char* v = std::getenv("UMBP_STANDALONE_GRPC_SHUTDOWN_DEADLINE_SEC"); + if (!v) v = std::getenv("UMBP_GRPC_SHUTDOWN_DEADLINE_SEC"); + if (!v) return std::chrono::seconds(5); + int seconds = std::atoi(v); + return std::chrono::seconds(seconds > 0 ? seconds : 5); +} + +::umbp::TierType TierToProto(TierType tier) { + switch (tier) { + 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; + } +} + +TierType TierFromProto(::umbp::TierType tier) { + switch (tier) { + case ::umbp::TIER_HBM: + return TierType::HBM; + case ::umbp::TIER_DRAM: + return TierType::DRAM; + case ::umbp::TIER_SSD: + return TierType::SSD; + default: + return TierType::UNKNOWN; + } +} + +bool FillSockaddr(const std::string& path, sockaddr_un* addr, socklen_t* addr_len) { + if (path.empty() || path.size() >= sizeof(addr->sun_path)) return false; + std::memset(addr, 0, sizeof(*addr)); + addr->sun_family = AF_UNIX; + std::strncpy(addr->sun_path, path.c_str(), sizeof(addr->sun_path) - 1); + *addr_len = static_cast(sizeof(sa_family_t) + path.size() + 1); + return true; +} + +void SetBool(::umbp::BoolResponse* response, bool ok, const std::string& error = {}) { + response->set_ok(ok); + if (!error.empty()) response->set_error(error); +} + +bool SetFdSocketTimeouts(int fd, std::chrono::milliseconds timeout) { + if (fd < 0 || timeout.count() <= 0) return true; + timeval tv; + tv.tv_sec = static_cast(timeout.count() / 1000); + tv.tv_usec = static_cast((timeout.count() % 1000) * 1000); + return setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) == 0 && + setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)) == 0; +} + +UMBPConfig NormalizeBackendConfig(UMBPConfig config) { + // Older callers pass the worker-facing standalone_process field to the + // 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; +} + +::umbp::StandaloneBackendMode BackendModeToProto(UMBPDeploymentMode mode) { + switch (mode) { + case UMBPDeploymentMode::Distributed: + return ::umbp::STANDALONE_BACKEND_DISTRIBUTED; + case UMBPDeploymentMode::Local: + return ::umbp::STANDALONE_BACKEND_LOCAL; + case UMBPDeploymentMode::StandaloneProcess: + default: + return ::umbp::STANDALONE_BACKEND_UNKNOWN; + } +} + +std::chrono::milliseconds FdHandshakeTimeout() { + const char* raw = std::getenv("UMBP_STANDALONE_FD_HANDSHAKE_TIMEOUT_MS"); + if (!raw || raw[0] == '\0') return std::chrono::milliseconds(5000); + char* end = nullptr; + long value = std::strtol(raw, &end, 10); + if (end == raw || value <= 0) return std::chrono::milliseconds(5000); + if (value > INT_MAX) value = INT_MAX; + return std::chrono::milliseconds(value); +} + +} // namespace + +class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { + public: + Impl(const UMBPConfig& config, std::string address) + : backend_config_(NormalizeBackendConfig(config)), + client_(CreateUMBPClient(backend_config_)), + address_(std::move(address)), + fd_socket_path_(DeriveFdSocketPath(address_)) {} + + ~Impl() override { Shutdown(); } + + bool Start() { + std::string error; + const std::string grpc_path = UnixPathFromGrpcAddress(address_); + if (!EnsureParentDirectory(grpc_path, &error)) { + MORI_UMBP_ERROR("[StandaloneServer] {}", error); + return false; + } + if (!EnsureParentDirectory(fd_socket_path_, &error)) { + MORI_UMBP_ERROR("[StandaloneServer] {}", error); + return false; + } + + unlink(grpc_path.c_str()); + unlink(fd_socket_path_.c_str()); + + if (!StartFdListener()) return false; + + grpc::ServerBuilder builder; + builder.SetMaxReceiveMessageSize(64 * 1024 * 1024); + builder.SetMaxSendMessageSize(64 * 1024 * 1024); + builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::MIN_POLLERS, 4); + builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::MAX_POLLERS, 32); + + int selected_port = 0; + builder.AddListeningPort(address_, grpc::InsecureServerCredentials(), &selected_port); + builder.RegisterService(this); + mode_t old_umask = umask(0077); + server_ = builder.BuildAndStart(); + umask(old_umask); + if (!server_) { + MORI_UMBP_ERROR("[StandaloneServer] failed to start gRPC server on {}", address_); + StopFdListener(); + return false; + } + + chmod(grpc_path.c_str(), 0600); + MORI_UMBP_INFO("[StandaloneServer] listening grpc={} fd_socket={}", address_, fd_socket_path_); + return true; + } + + void Run() { + if (server_) server_->Wait(); + } + + void Shutdown() { + bool expected = false; + if (!shutdown_.compare_exchange_strong(expected, true)) return; + + StopFdListener(); + if (server_) { + server_->Shutdown(std::chrono::system_clock::now() + ShutdownDeadline()); + } + UnregisterAllExternalIdentities(); + { + std::lock_guard lock(client_mu_); + client_->Flush(); + } + UnmapAll(); + { + std::lock_guard lock(client_mu_); + client_->Close(); + } + unlink(UnixPathFromGrpcAddress(address_).c_str()); + unlink(fd_socket_path_.c_str()); + } + + grpc::Status Ping(grpc::ServerContext*, const ::umbp::Empty*, + ::umbp::PingResponse* response) override { + response->set_ready(!shutdown_.load()); + response->set_deployment_mode(BackendModeToProto(client_->GetDeploymentMode())); + return grpc::Status::OK; + } + + grpc::Status Put(grpc::ServerContext*, const ::umbp::PutRequest* request, + ::umbp::BoolResponse* response) override { + uintptr_t ptr = 0; + if (!ResolveRange(request->client_id(), 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; + } + SetBool(response, client_->Put(request->key(), ptr, static_cast(request->size()))); + return grpc::Status::OK; + } + + grpc::Status Get(grpc::ServerContext*, const ::umbp::GetRequest* request, + ::umbp::BoolResponse* response) override { + uintptr_t ptr = 0; + if (!ResolveRange(request->client_id(), 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; + } + SetBool(response, client_->Get(request->key(), ptr, static_cast(request->size()))); + return grpc::Status::OK; + } + + grpc::Status BatchPut(grpc::ServerContext*, const ::umbp::BatchDataRequest* request, + ::umbp::BatchBoolResponse* response) override { + 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; + } + FillResults(client_->BatchPut(keys, ptrs, sizes), response); + return grpc::Status::OK; + } + + grpc::Status BatchPutWithDepth(grpc::ServerContext*, + const ::umbp::BatchDataWithDepthRequest* request, + ::umbp::BatchBoolResponse* response) override { + if (request->keys_size() != request->shm_offsets_size() || + request->keys_size() != request->sizes_size()) { + 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; + if (!ResolveRange(request->client_id(), 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_); + if (shutdown_.load()) { + FillFalse(request->keys_size(), response); + return grpc::Status::OK; + } + FillResults(client_->BatchPutWithDepth(keys, ptrs, sizes, depths), response); + return grpc::Status::OK; + } + + grpc::Status BatchGet(grpc::ServerContext*, const ::umbp::BatchDataRequest* request, + ::umbp::BatchBoolResponse* response) override { + 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; + } + FillResults(client_->BatchGet(keys, ptrs, sizes), response); + return grpc::Status::OK; + } + + grpc::Status Exists(grpc::ServerContext*, const ::umbp::KeyRequest* request, + ::umbp::BoolResponse* response) override { + std::lock_guard lock(client_mu_); + if (shutdown_.load()) { + SetBool(response, false, "server is shutting down"); + return grpc::Status::OK; + } + SetBool(response, client_->Exists(request->key())); + return grpc::Status::OK; + } + + 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_); + if (shutdown_.load()) { + FillFalse(request->keys_size(), response); + return grpc::Status::OK; + } + FillResults(client_->BatchExists(keys), response); + return grpc::Status::OK; + } + + 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_); + if (shutdown_.load()) { + response->set_count(0); + return grpc::Status::OK; + } + response->set_count(client_->BatchExistsConsecutive(keys)); + return grpc::Status::OK; + } + + grpc::Status Clear(grpc::ServerContext*, const ::umbp::Empty*, + ::umbp::BoolResponse* response) override { + std::lock_guard lock(client_mu_); + if (shutdown_.load()) { + SetBool(response, false, "server is shutting down"); + return grpc::Status::OK; + } + SetBool(response, client_->Clear()); + return grpc::Status::OK; + } + + grpc::Status Flush(grpc::ServerContext*, const ::umbp::Empty*, + ::umbp::BoolResponse* response) override { + std::lock_guard lock(client_mu_); + if (shutdown_.load()) { + SetBool(response, false, "server is shutting down"); + return grpc::Status::OK; + } + SetBool(response, client_->Flush()); + return grpc::Status::OK; + } + + grpc::Status RegisterMemory(grpc::ServerContext*, const ::umbp::RegisterMemoryRequest* request, + ::umbp::BoolResponse* response) override { + if (shutdown_.load()) { + SetBool(response, false, "server is shutting down"); + return grpc::Status::OK; + } + std::lock_guard lifecycle_lock(external_identity_lifecycle_mu_); + bool ok = false; + { + std::lock_guard lock(memory_mu_); + auto it = memory_.find(request->client_id()); + ok = it != memory_.end() && it->second.worker_base == request->worker_base() && + it->second.size >= request->size(); + } + if (!ok) { + SetBool(response, false, "fd handoff registration was not found"); + return grpc::Status::OK; + } + + if (!EnsureExternalIdentity(*request)) { + MORI_UMBP_WARN( + "[StandaloneServer] external-KV identity registration failed for client_id={} " + "worker_node_id={}; continuing with core memory registration", + request->client_id(), request->worker_node_id()); + } + SetBool(response, true); + return grpc::Status::OK; + } + + grpc::Status DeregisterMemory(grpc::ServerContext*, + const ::umbp::DeregisterMemoryRequest* request, + ::umbp::Empty*) override { + std::lock_guard lifecycle_lock(external_identity_lifecycle_mu_); + RemoveExternalIdentity(request->client_id()); + UnmapClient(request->client_id()); + return grpc::Status::OK; + } + + grpc::Status ReportExternalKvBlocks(grpc::ServerContext*, + const ::umbp::StandaloneExternalKvMutationRequest* request, + ::umbp::BoolResponse* response) override { + if (!BackendIsDistributed()) { + SetBool(response, true); + return grpc::Status::OK; + } + auto identity = GetExternalIdentity(request->client_id()); + if (!identity) { + SetBool(response, false, "external-KV identity is not registered for client_id"); + return grpc::Status::OK; + } + std::vector hashes(request->hashes().begin(), request->hashes().end()); + SetBool(response, identity->ReportExternalKvBlocks(hashes, TierFromProto(request->tier()))); + return grpc::Status::OK; + } + + grpc::Status RevokeExternalKvBlocks(grpc::ServerContext*, + const ::umbp::StandaloneExternalKvMutationRequest* request, + ::umbp::BoolResponse* response) override { + if (!BackendIsDistributed()) { + SetBool(response, true); + return grpc::Status::OK; + } + auto identity = GetExternalIdentity(request->client_id()); + if (!identity) { + SetBool(response, false, "external-KV identity is not registered for client_id"); + return grpc::Status::OK; + } + std::vector hashes(request->hashes().begin(), request->hashes().end()); + SetBool(response, identity->RevokeExternalKvBlocks(hashes, TierFromProto(request->tier()))); + return grpc::Status::OK; + } + + grpc::Status RevokeAllExternalKvBlocksAtTier( + grpc::ServerContext*, const ::umbp::StandaloneExternalKvTierRequest* request, + ::umbp::BoolResponse* response) override { + if (!BackendIsDistributed()) { + SetBool(response, true); + return grpc::Status::OK; + } + auto identity = GetExternalIdentity(request->client_id()); + if (!identity) { + SetBool(response, false, "external-KV identity is not registered for client_id"); + return grpc::Status::OK; + } + SetBool(response, identity->RevokeAllExternalKvBlocksAtTier(TierFromProto(request->tier()))); + return grpc::Status::OK; + } + + grpc::Status MatchExternalKv(grpc::ServerContext*, + const ::umbp::StandaloneMatchExternalKvRequest* request, + ::umbp::StandaloneMatchExternalKvResponse* response) override { + if (!BackendIsDistributed()) return grpc::Status::OK; + std::vector hashes(request->hashes().begin(), request->hashes().end()); + std::vector matches; + if (auto identity = GetExternalIdentity(request->client_id())) { + matches = identity->MatchExternalKv(hashes, request->count_as_hit()); + } else { + std::lock_guard lock(client_mu_); + if (!shutdown_.load()) matches = client_->MatchExternalKv(hashes, request->count_as_hit()); + } + FillExternalKvMatches(matches, response); + return grpc::Status::OK; + } + + grpc::Status GetExternalKvHitCounts( + grpc::ServerContext*, const ::umbp::StandaloneExternalKvHitCountsRequest* request, + ::umbp::StandaloneExternalKvHitCountsResponse* response) override { + if (!BackendIsDistributed()) return grpc::Status::OK; + std::vector hashes(request->hashes().begin(), request->hashes().end()); + std::vector entries; + if (auto identity = GetExternalIdentity(request->client_id())) { + entries = identity->GetExternalKvHitCounts(hashes); + } else { + std::lock_guard lock(client_mu_); + if (!shutdown_.load()) entries = client_->GetExternalKvHitCounts(hashes); + } + FillExternalKvHitCounts(entries, response); + return grpc::Status::OK; + } + + private: + struct RegisteredMemory { + void* base = nullptr; + uint64_t worker_base = 0; + uint64_t size = 0; + }; + + bool BackendIsDistributed() const { + return client_ && client_->GetDeploymentMode() == UMBPDeploymentMode::Distributed; + } + + std::string BackendPeerAddress() const { + if (!backend_config_.distributed.has_value()) return ""; + const auto& dist = backend_config_.distributed.value(); + if (dist.peer_service_port == 0 || dist.master_config.node_address.empty()) return ""; + return dist.master_config.node_address + ":" + std::to_string(dist.peer_service_port); + } + + bool EnsureExternalIdentity(const ::umbp::RegisterMemoryRequest& request) { + if (!BackendIsDistributed()) return true; + if (request.worker_node_id().empty()) return true; + if (!backend_config_.distributed.has_value()) return false; + + std::shared_ptr old; + { + std::lock_guard lock(external_identity_mu_); + auto it = external_identities_.find(request.client_id()); + if (it != external_identities_.end()) { + if (it->second && it->second->node_id() == request.worker_node_id()) return true; + old = std::move(it->second); + external_identities_.erase(it); + } + } + if (old) old->Stop(); + + const auto& dist = backend_config_.distributed.value(); + ExternalKvIdentityClient::Config cfg; + cfg.master_address = dist.master_config.master_address; + cfg.node_id = request.worker_node_id(); + cfg.node_address = request.worker_node_address(); + cfg.peer_address = BackendPeerAddress(); + cfg.tags.assign(request.tags().begin(), request.tags().end()); + + auto identity = std::make_shared(std::move(cfg)); + if (!identity->Start()) return false; + + { + std::lock_guard lock(external_identity_mu_); + external_identities_[request.client_id()] = identity; + } + return true; + } + + std::shared_ptr GetExternalIdentity(const std::string& client_id) { + std::lock_guard lock(external_identity_mu_); + auto it = external_identities_.find(client_id); + return it == external_identities_.end() ? nullptr : it->second; + } + + void RemoveExternalIdentity(const std::string& client_id) { + std::shared_ptr identity; + { + std::lock_guard lock(external_identity_mu_); + auto it = external_identities_.find(client_id); + if (it == external_identities_.end()) return; + identity = std::move(it->second); + external_identities_.erase(it); + } + if (identity) identity->Stop(); + } + + void UnregisterAllExternalIdentities() { + std::lock_guard lifecycle_lock(external_identity_lifecycle_mu_); + std::vector> identities; + { + std::lock_guard lock(external_identity_mu_); + for (auto& kv : external_identities_) identities.push_back(std::move(kv.second)); + external_identities_.clear(); + } + for (auto& identity : identities) { + if (identity) identity->Stop(); + } + } + + bool StartFdListener() { + listen_fd_ = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (listen_fd_ < 0) { + MORI_UMBP_ERROR("[StandaloneServer] fd socket() failed: {}", std::strerror(errno)); + return false; + } + + sockaddr_un addr; + socklen_t addr_len = 0; + if (!FillSockaddr(fd_socket_path_, &addr, &addr_len)) { + MORI_UMBP_ERROR("[StandaloneServer] invalid fd socket path {}", fd_socket_path_); + close(listen_fd_); + listen_fd_ = -1; + return false; + } + + mode_t old_umask = umask(0077); + int bind_rc = bind(listen_fd_, reinterpret_cast(&addr), addr_len); + umask(old_umask); + if (bind_rc != 0) { + MORI_UMBP_ERROR("[StandaloneServer] bind('{}') failed: {}", fd_socket_path_, + std::strerror(errno)); + close(listen_fd_); + listen_fd_ = -1; + return false; + } + chmod(fd_socket_path_.c_str(), 0600); + + if (listen(listen_fd_, 16) != 0) { + MORI_UMBP_ERROR("[StandaloneServer] listen('{}') failed: {}", fd_socket_path_, + std::strerror(errno)); + close(listen_fd_); + listen_fd_ = -1; + return false; + } + + fd_running_.store(true); + fd_thread_ = std::thread([this]() { FdAcceptLoop(); }); + return true; + } + + 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); + if (fd_thread_.joinable()) fd_thread_.join(); + if (listen_fd_ >= 0) { + close(listen_fd_); + listen_fd_ = -1; + } + } + + void FdAcceptLoop() { + while (fd_running_.load()) { + int client_fd = accept4(listen_fd_, nullptr, nullptr, SOCK_CLOEXEC); + if (client_fd < 0) { + if (fd_running_.load()) { + MORI_UMBP_WARN("[StandaloneServer] accept fd socket failed: {}", std::strerror(errno)); + } + continue; + } + active_fd_client_.store(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); + close(client_fd); + } + } + + void HandleFdConnection(int client_fd) { + FdRegistrationMessage msg; + std::string error; + int fd = RecvFdRegistration(client_fd, &msg, &error); + if (fd < 0) { + MORI_UMBP_WARN("[StandaloneServer] fd registration receive failed: {}", error); + SendStatus(client_fd, -1); + return; + } + + int32_t status = RegisterFd(fd, msg) ? 0 : -1; + SendStatus(client_fd, status); + } + + bool RegisterFd(int fd, const FdRegistrationMessage& msg) { + void* mapped = + mmap(nullptr, static_cast(msg.size), PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + close(fd); + if (mapped == MAP_FAILED) { + MORI_UMBP_WARN("[StandaloneServer] mmap received fd failed: {}", std::strerror(errno)); + return false; + } + if (!RegisterBackendMemory(mapped, static_cast(msg.size))) { + MORI_UMBP_WARN("[StandaloneServer] backend RegisterMemory failed for client_id={}", + msg.client_id); + munmap(mapped, static_cast(msg.size)); + return false; + } + + std::string client_id(msg.client_id); + std::optional old_mem; + { + std::lock_guard lock(memory_mu_); + auto old = memory_.find(client_id); + if (old != memory_.end()) old_mem = old->second; + memory_[client_id] = + RegisteredMemory{mapped, static_cast(msg.worker_base), msg.size}; + } + if (old_mem.has_value()) ReleaseRegisteredMemory(*old_mem); + MORI_UMBP_INFO("[StandaloneServer] registered shm client_id={} worker_base=0x{:x} size={}MB", + client_id, msg.worker_base, msg.size / (1024 * 1024)); + return true; + } + + void UnmapClient(const std::string& client_id) { + std::optional mem; + { + std::lock_guard lock(memory_mu_); + auto it = memory_.find(client_id); + if (it == memory_.end()) return; + mem = it->second; + memory_.erase(it); + } + ReleaseRegisteredMemory(*mem); + } + + void UnmapAll() { + std::vector entries; + { + std::lock_guard lock(memory_mu_); + for (auto& kv : memory_) entries.push_back(kv.second); + memory_.clear(); + } + for (const auto& mem : entries) ReleaseRegisteredMemory(mem); + } + + bool RegisterBackendMemory(void* base, size_t size) { + std::lock_guard lock(client_mu_); + if (shutdown_.load()) return false; + return client_->RegisterMemory(reinterpret_cast(base), size); + } + + void ReleaseRegisteredMemory(const RegisteredMemory& mem) { + if (!mem.base) return; + { + std::lock_guard lock(client_mu_); + client_->DeregisterMemory(reinterpret_cast(mem.base)); + } + munmap(mem.base, static_cast(mem.size)); + } + + bool ResolveRange(const std::string& client_id, uint64_t offset, uint64_t size, + uintptr_t* out_ptr) { + if (!out_ptr || size == 0) return false; + std::lock_guard lock(memory_mu_); + auto it = memory_.find(client_id); + if (it == memory_.end()) return false; + const RegisteredMemory& mem = it->second; + if (offset > mem.size || size > mem.size - offset) return false; + *out_ptr = reinterpret_cast(mem.base) + static_cast(offset); + return true; + } + + bool ResolveBatch(const ::umbp::BatchDataRequest& request, std::vector* ptrs) { + if (!ptrs || request.keys_size() != request.shm_offsets_size() || + request.keys_size() != request.sizes_size()) { + return false; + } + ptrs->clear(); + ptrs->reserve(request.keys_size()); + for (int i = 0; i < request.keys_size(); ++i) { + uintptr_t ptr = 0; + if (!ResolveRange(request.client_id(), request.shm_offsets(i), request.sizes(i), &ptr)) { + return false; + } + ptrs->push_back(ptr); + } + return true; + } + + static std::vector Sizes(const ::umbp::BatchDataRequest& request) { + std::vector sizes; + sizes.reserve(request.sizes_size()); + for (uint64_t size : request.sizes()) sizes.push_back(static_cast(size)); + return sizes; + } + + static void FillResults(const std::vector& results, ::umbp::BatchBoolResponse* response) { + response->mutable_ok()->Reserve(static_cast(results.size())); + for (bool ok : results) response->add_ok(ok); + } + + static void FillFalse(int n, ::umbp::BatchBoolResponse* response) { + response->mutable_ok()->Reserve(n); + for (int i = 0; i < n; ++i) response->add_ok(false); + } + + static void FillExternalKvMatches(const std::vector& matches, + ::umbp::StandaloneMatchExternalKvResponse* response) { + for (const auto& match : matches) { + auto* out = response->add_matches(); + out->set_node_id(match.node_id); + out->set_peer_address(match.peer_address); + for (const auto& [tier, hashes] : match.hashes_by_tier) { + auto* bucket = out->add_hashes_by_tier(); + bucket->set_tier(TierToProto(tier)); + for (const auto& hash : hashes) bucket->add_hashes(hash); + } + } + } + + static void FillExternalKvHitCounts( + const std::vector& entries, + ::umbp::StandaloneExternalKvHitCountsResponse* response) { + for (const auto& entry : entries) { + auto* out = response->add_entries(); + out->set_hash(entry.hash); + out->set_hit_count_total(entry.hit_count_total); + } + } + + UMBPConfig backend_config_; + std::unique_ptr client_; + std::string address_; + std::string fd_socket_path_; + std::unique_ptr server_; + std::atomic shutdown_{false}; + + std::mutex client_mu_; + std::mutex memory_mu_; + std::map memory_; + 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::thread fd_thread_; +}; + +StandaloneServer::StandaloneServer(UMBPConfig config, std::string address) + : config_(std::move(config)), address_(std::move(address)) { + impl_ = std::make_unique(config_, address_); +} + +StandaloneServer::~StandaloneServer() { Shutdown(); } + +bool StandaloneServer::Start() { return impl_->Start(); } + +void StandaloneServer::Run() { impl_->Run(); } + +void StandaloneServer::Shutdown() { + if (impl_) impl_->Shutdown(); +} + +} // namespace mori::umbp::standalone diff --git a/src/umbp/tests/CMakeLists.txt b/src/umbp/tests/CMakeLists.txt index 40e431c15..a26258124 100644 --- a/src/umbp/tests/CMakeLists.txt +++ b/src/umbp/tests/CMakeLists.txt @@ -223,3 +223,13 @@ add_executable(test_cache_remote_admission test_cache_remote_admission.cpp) target_link_libraries(test_cache_remote_admission PRIVATE umbp_common GTest::gtest_main) gtest_discover_tests(test_cache_remote_admission) + +# --------------------------------------------------------------------------- +# test_standalone_shm_ipc — memfd-backed host allocator registry plus raw UDS +# fd handoff used by standalone-process mode. +# --------------------------------------------------------------------------- +add_executable(test_standalone_shm_ipc test_standalone_shm_ipc.cpp) +target_link_libraries(test_standalone_shm_ipc PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_standalone_shm_ipc PRIVATE cxx_std_17) +gtest_discover_tests(test_standalone_shm_ipc) diff --git a/src/umbp/tests/test_ssd_copy_pipeline.cpp b/src/umbp/tests/test_ssd_copy_pipeline.cpp index 1d144cbb6..1ee21b71c 100644 --- a/src/umbp/tests/test_ssd_copy_pipeline.cpp +++ b/src/umbp/tests/test_ssd_copy_pipeline.cpp @@ -35,6 +35,9 @@ #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" namespace mori::umbp { namespace { @@ -337,5 +340,33 @@ 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/src/umbp/tests/test_standalone_shm_ipc.cpp b/src/umbp/tests/test_standalone_shm_ipc.cpp new file mode 100644 index 000000000..76043be8e --- /dev/null +++ b/src/umbp/tests/test_standalone_shm_ipc.cpp @@ -0,0 +1,285 @@ +// 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 + +#include "umbp/local/host_mem_allocator.h" +#include "umbp/standalone/ipc.h" +#include "umbp/standalone/standalone_server.h" +#include "umbp/umbp_client.h" + +namespace mori::umbp { +namespace { + +TEST(StandaloneShmIpcTest, AnonymousShmRegistryLookupMapsSameMemory) { + HostMemAllocator allocator; + HostBufferOptions opts; + opts.backing = HostBufferBacking::kAnonymousShm; + opts.prefault = false; + + HostBufferHandle handle = allocator.Alloc(4096, opts); + ASSERT_TRUE(handle.valid()); + EXPECT_EQ(handle.actual_backing, HostBufferBacking::kAnonymousShm); + + auto allocation = + HostMemAllocator::LookupShmAllocation(reinterpret_cast(handle.ptr), 128); + ASSERT_TRUE(allocation.has_value()); + EXPECT_EQ(allocation->base, handle.ptr); + EXPECT_GE(allocation->mapped_size, handle.mapped_size); + ASSERT_GE(allocation->fd, 0); + + int dup_fd = dup(allocation->fd); + ASSERT_GE(dup_fd, 0); + void* mirror = + mmap(nullptr, allocation->mapped_size, PROT_READ | PROT_WRITE, MAP_SHARED, dup_fd, 0); + close(dup_fd); + ASSERT_NE(mirror, MAP_FAILED); + + static_cast(handle.ptr)[17] = 0x5a; + EXPECT_EQ(static_cast(mirror)[17], 0x5a); + munmap(mirror, allocation->mapped_size); + + allocator.Free(handle); + EXPECT_FALSE(handle.valid()); + EXPECT_FALSE( + HostMemAllocator::LookupShmAllocation(reinterpret_cast(allocation->base), 128) + .has_value()); +} + +TEST(StandaloneShmIpcTest, ActiveAnonymousShmFreeIsDeferredUntilRelease) { + HostMemAllocator allocator; + HostBufferOptions opts; + opts.backing = HostBufferBacking::kAnonymousShm; + opts.prefault = false; + + HostBufferHandle handle = allocator.Alloc(4096, opts); + ASSERT_TRUE(handle.valid()); + static_cast(handle.ptr)[9] = 0x33; + + auto held = HostMemAllocator::AcquireShmAllocation(reinterpret_cast(handle.ptr), 4096); + ASSERT_TRUE(held.has_value()); + int dup_fd = dup(held->fd); + ASSERT_GE(dup_fd, 0); + uintptr_t base = reinterpret_cast(held->base); + + allocator.Free(handle); + EXPECT_FALSE(handle.valid()); + EXPECT_FALSE(HostMemAllocator::LookupShmAllocation(base, 16).has_value()); + + HostMemAllocator::ReleaseShmAllocation(base); + void* mirror = mmap(nullptr, held->mapped_size, PROT_READ | PROT_WRITE, MAP_SHARED, dup_fd, 0); + close(dup_fd); + ASSERT_NE(mirror, MAP_FAILED); + EXPECT_EQ(static_cast(mirror)[9], 0x33); + munmap(mirror, held->mapped_size); +} + +bool FillSockaddr(const std::string& path, sockaddr_un* addr, socklen_t* addr_len) { + if (path.size() >= sizeof(addr->sun_path)) return false; + std::memset(addr, 0, sizeof(*addr)); + addr->sun_family = AF_UNIX; + std::strncpy(addr->sun_path, path.c_str(), sizeof(addr->sun_path) - 1); + *addr_len = static_cast(sizeof(sa_family_t) + path.size() + 1); + return true; +} + +TEST(StandaloneShmIpcTest, RawUdsFdRegistrationTransfersFd) { + HostMemAllocator allocator; + HostBufferOptions opts; + opts.backing = HostBufferBacking::kAnonymousShm; + opts.prefault = false; + HostBufferHandle handle = allocator.Alloc(4096, opts); + ASSERT_TRUE(handle.valid()); + static_cast(handle.ptr)[3] = 0x7b; + + auto allocation = + HostMemAllocator::LookupShmAllocation(reinterpret_cast(handle.ptr), 4096); + ASSERT_TRUE(allocation.has_value()); + + std::string path = "/tmp/umbp_standalone_ipc_test_" + std::to_string(getpid()) + ".sock"; + unlink(path.c_str()); + + int listen_fd = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + ASSERT_GE(listen_fd, 0); + sockaddr_un addr; + socklen_t addr_len = 0; + ASSERT_TRUE(FillSockaddr(path, &addr, &addr_len)); + ASSERT_EQ(bind(listen_fd, reinterpret_cast(&addr), addr_len), 0) + << std::strerror(errno); + ASSERT_EQ(listen(listen_fd, 1), 0) << std::strerror(errno); + + std::atomic receiver_ok{false}; + std::thread receiver([&]() { + int accepted = accept4(listen_fd, nullptr, nullptr, SOCK_CLOEXEC); + if (accepted < 0) return; + standalone::FdRegistrationMessage msg; + std::string error; + int received_fd = standalone::RecvFdRegistration(accepted, &msg, &error); + if (received_fd >= 0 && std::string(msg.client_id) == "client-a" && + msg.worker_base == reinterpret_cast(handle.ptr) && msg.size >= 4096) { + void* mirror = mmap(nullptr, static_cast(msg.size), PROT_READ | PROT_WRITE, + MAP_SHARED, received_fd, 0); + close(received_fd); + if (mirror != MAP_FAILED) { + receiver_ok.store(static_cast(mirror)[3] == 0x7b); + munmap(mirror, static_cast(msg.size)); + } + standalone::SendStatus(accepted, 0); + } else { + if (received_fd >= 0) close(received_fd); + standalone::SendStatus(accepted, -1); + } + close(accepted); + }); + + std::string error; + int status = standalone::SendFdRegistration(path, allocation->fd, "client-a", + reinterpret_cast(handle.ptr), + allocation->mapped_size, 1000, &error); + EXPECT_EQ(status, 0) << error; + receiver.join(); + close(listen_fd); + unlink(path.c_str()); + allocator.Free(handle); + + EXPECT_TRUE(receiver_ok.load()); +} + +TEST(StandaloneShmIpcTest, StandaloneClientUsesNonZeroOffsetsAndCanReregister) { + const std::string address = + "unix:///tmp/umbp_standalone_e2e_" + 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 server_cfg; + server_cfg.dram.capacity_bytes = 1 << 20; + server_cfg.ssd.enabled = false; + UMBPStandaloneProcessConfig sp_cfg; + sp_cfg.address = address; + sp_cfg.startup_timeout_ms = 5000; + server_cfg.standalone_process = sp_cfg; + + standalone::StandaloneServer server(server_cfg, address); + ASSERT_TRUE(server.Start()); + std::thread server_thread([&]() { server.Run(); }); + + UMBPConfig client_cfg = server_cfg; + auto client = CreateUMBPClient(client_cfg); + ASSERT_EQ(client->GetDeploymentMode(), UMBPDeploymentMode::StandaloneProcess); + + HostMemAllocator allocator; + HostBufferOptions opts; + opts.backing = HostBufferBacking::kAnonymousShm; + opts.prefault = false; + HostBufferHandle handle = allocator.Alloc(4096, opts); + ASSERT_TRUE(handle.valid()); + auto* bytes = static_cast(handle.ptr); + + ASSERT_TRUE(client->RegisterMemory(reinterpret_cast(handle.ptr), handle.mapped_size)); + + for (int i = 0; i < 16; ++i) bytes[32 + i] = static_cast(i + 1); + ASSERT_TRUE(client->Put("offset-key", reinterpret_cast(bytes + 32), 16)); + std::memset(bytes + 96, 0, 16); + ASSERT_TRUE(client->Get("offset-key", reinterpret_cast(bytes + 96), 16)); + for (int i = 0; i < 16; ++i) EXPECT_EQ(bytes[96 + i], static_cast(i + 1)); + + client->DeregisterMemory(reinterpret_cast(handle.ptr)); + ASSERT_TRUE(client->RegisterMemory(reinterpret_cast(handle.ptr), handle.mapped_size)); + for (int i = 0; i < 8; ++i) bytes[128 + i] = static_cast(0xa0 + i); + ASSERT_TRUE(client->Put("reregister-key", reinterpret_cast(bytes + 128), 8)); + std::memset(bytes + 192, 0, 8); + ASSERT_TRUE(client->Get("reregister-key", reinterpret_cast(bytes + 192), 8)); + for (int i = 0; i < 8; ++i) EXPECT_EQ(bytes[192 + i], static_cast(0xa0 + i)); + + client->DeregisterMemory(reinterpret_cast(handle.ptr)); + client->Close(); + allocator.Free(handle); + 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"; + 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 cfg; + cfg.dram.capacity_bytes = 1 << 20; + cfg.ssd.enabled = false; + UMBPStandaloneProcessConfig sp_cfg; + sp_cfg.address = address; + sp_cfg.startup_timeout_ms = 5000; + cfg.standalone_process = sp_cfg; + + standalone::StandaloneServer server(cfg, address); + ASSERT_TRUE(server.Start()); + std::thread server_thread([&]() { server.Run(); }); + + int sock = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + ASSERT_GE(sock, 0); + sockaddr_un addr; + socklen_t addr_len = 0; + ASSERT_TRUE(FillSockaddr(fd_path, &addr, &addr_len)); + ASSERT_EQ(connect(sock, reinterpret_cast(&addr), addr_len), 0) << std::strerror(errno); + + std::atomic done{false}; + std::thread shutdown_thread([&]() { + server.Shutdown(); + done.store(true); + }); + + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (!done.load() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_TRUE(done.load()); + close(sock); + shutdown_thread.join(); + server_thread.join(); + unlink(grpc_path.c_str()); + unlink(fd_path.c_str()); +} + +} // namespace +} // namespace mori::umbp diff --git a/src/umbp/umbp_client_factory.cpp b/src/umbp/umbp_client_factory.cpp index 8aa18fb0b..87b152f4b 100644 --- a/src/umbp/umbp_client_factory.cpp +++ b/src/umbp/umbp_client_factory.cpp @@ -21,6 +21,7 @@ // SOFTWARE. #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 { @@ -29,6 +30,9 @@ std::unique_ptr CreateUMBPClient(const UMBPConfig& config) { if (config.distributed.has_value()) { return std::make_unique(config); } + if (config.standalone_process.has_value()) { + return std::make_unique(config); + } return std::make_unique(config); }