From 47e9ac99379654fbcb33e9ea617c45e6585f003a Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Thu, 9 Jul 2026 06:11:47 +0000 Subject: [PATCH 1/3] feat(umbp): add standalone process mode --- python/mori/umbp/__init__.py | 25 + setup.py | 9 + src/pybind/pybind_umbp.cpp | 19 +- src/umbp/CMakeLists.txt | 45 +- .../distributed/proto/umbp_standalone.proto | 132 ++ .../doc/design-standalone-process-mode.md | 1156 +++++++++++++++++ src/umbp/include/umbp/common/config.h | 30 + .../umbp/distributed/distributed_client.h | 1 + .../include/umbp/local/host_mem_allocator.h | 20 + .../include/umbp/local/standalone_client.h | 1 + .../include/umbp/local/tiers/copy_pipeline.h | 4 + src/umbp/include/umbp/standalone/ipc.h | 58 + .../standalone/standalone_process_client.h | 107 ++ .../umbp/standalone/standalone_server.h | 53 + src/umbp/include/umbp/umbp_client.h | 11 + src/umbp/local/host_mem_allocator.cpp | 213 ++- src/umbp/local/standalone_client.cpp | 9 +- src/umbp/local/tiers/copy_pipeline.cpp | 14 + .../standalone/bin/standalone_server_main.cpp | 102 ++ src/umbp/standalone/ipc.cpp | 297 +++++ .../standalone/standalone_process_client.cpp | 600 +++++++++ src/umbp/standalone/standalone_server.cpp | 618 +++++++++ src/umbp/tests/CMakeLists.txt | 10 + src/umbp/tests/test_ssd_copy_pipeline.cpp | 31 + src/umbp/tests/test_standalone_shm_ipc.cpp | 285 ++++ src/umbp/umbp_client_factory.cpp | 4 + 26 files changed, 3842 insertions(+), 12 deletions(-) create mode 100644 src/umbp/distributed/proto/umbp_standalone.proto create mode 100644 src/umbp/doc/design-standalone-process-mode.md create mode 100644 src/umbp/include/umbp/standalone/ipc.h create mode 100644 src/umbp/include/umbp/standalone/standalone_process_client.h create mode 100644 src/umbp/include/umbp/standalone/standalone_server.h create mode 100644 src/umbp/standalone/bin/standalone_server_main.cpp create mode 100644 src/umbp/standalone/ipc.cpp create mode 100644 src/umbp/standalone/standalone_process_client.cpp create mode 100644 src/umbp/standalone/standalone_server.cpp create mode 100644 src/umbp/tests/test_standalone_shm_ipc.cpp 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..f1c607318 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,12 @@ 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); + py::class_(m, "UMBPConfig") .def(py::init<>()) .def_static("from_environment", &UMBPConfig::FromEnvironment) @@ -238,7 +251,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 +277,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..5afd1ae08 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,10 @@ 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/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 +398,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 +410,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..a002e0716 --- /dev/null +++ b/src/umbp/distributed/proto/umbp_standalone.proto @@ -0,0 +1,132 @@ +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; +} + +message BoolResponse { + bool ok = 1; + string error = 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; +} + +message DeregisterMemoryRequest { + string client_id = 1; +} + +message StandaloneExternalKvMutationRequest { + repeated string hashes = 1; + TierType tier = 2; +} + +message StandaloneExternalKvTierRequest { + TierType tier = 1; +} + +message StandaloneMatchExternalKvRequest { + repeated string hashes = 1; + bool count_as_hit = 2; +} + +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; +} + +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..5972a2ab7 --- /dev/null +++ b/src/umbp/doc/design-standalone-process-mode.md @@ -0,0 +1,1156 @@ +# UMBP Standalone Process Mode — Design (Draft v0.5) + +**Scope:** A new deployment mode for UMBP in which the DRAM/SSD tiers, +`LocalBlockIndex`, and `CopyPipeline` run inside a dedicated, long-lived +OS process on the same host, while one or more SGLang/vLLM worker +processes talk to it through `IUMBPClient` as if it were still +in-process. This is *not* the existing `distributed/` mode (separate +master + multi-node peers over gRPC + RDMA) — it targets the +**single-host, multi-worker-process** case that today forces every +worker to duplicate its own DRAM tier and lose the ability to share a +KV cache across sibling processes (e.g. TP ranks, or a restarted +worker resuming a warm cache). + +**Status:** design proposal, not yet implemented. Everything under +"Open questions" needs a decision before implementation starts. + +**Revision history:** this draft has gone through two source-grounded +code review passes so far. Each found concrete implementation gaps in +the prior draft — inline "(v0.2 fix)" / "(v0.3 fix)" markers throughout +mark exactly what changed and why, so the reasoning stays visible +rather than silently overwritten. + +- **v0.2** (§3.1, §3.2, §4.3, §5, §6, §7, §8): the `is_distributed()` + gate that would silently skip `register_memory` in this mode (§5); + the missing mechanism for getting a `memfd` out of `HostBufferHandle` + and into `StandaloneProcessClient::RegisterMemory` (§3.2, §5); an + internal inconsistency between "same UDS" and "second UDS" for fd + handoff (§3.2, §8); an incompatible TCP-fallback claim for an + `SCM_RIGHTS`-based data plane (§3.1, §6, §8); an overreach in + claiming the server's own `DRAMTier` shared-memory backing needs no + new code (§3.2, §7); plus a missing `CopyPipeline::Drain()` for + clean shutdown (§4.3). +- **v0.3** (§3.1, §4.1, §5, §6, §10): the still-unclosed activation + path for `cfg.standalone_process` — nothing in v0.2 specified who + actually sets it from Python, the same way `umbp_store.py` already + does for `cfg.distributed` (§5, §10); a leftover default-path + inconsistency between `.sock` and `.grpc.sock` + plus an unspecified fd-socket derivation rule (§3.1, §4.1, §6); an + fd-ownership contradiction between the registry design and + "`DeregisterMemory` closes the worker's fd" (§3.2); a missing `Ping` + RPC that §4.1's readiness probe depends on, plus explicit + client-local-vs-RPC method classification (§3.1); and small + corrections (a stale `batch_resolder_codec.h` typo, a table cell + claiming the server "owns the shared-memory DRAM region" that + contradicted v0.2's own fix). +- **v0.4** (§4.2, §5, §6): `UMBPStandaloneProcessConfig`'s + field/parsing ownership was still ambiguous — §6's table described + `fd_socket`/`transport`/`auto_start` with routing language that + conflicted with both the config struct definition (only + `address`/`auto_start`/`startup_timeout_ms`) and the §5 Python + snippet (which only ever set `address`). Resolved by picking one of + the two existing, mutually-inconsistent precedents in this codebase + (`cfg.distributed`: all fields Python-set, vs. + `spdk_proxy_auto_start`: read directly by C++ + `FromEnvironment()`) and applying it consistently: `address`, + `auto_start`, and `startup_timeout_ms` are now all Python-set + (mirroring `distributed`), while `fd_socket` and `transport` are + removed from the config surface entirely — `fd_socket` is always + mechanically derived from `address` (no knob needed), and + `transport` has no v0.1 config surface since `tcp_staging` isn't + built yet (§8 item 6b). +- **v0.5** (§5, §10): two blockers found by tracing SGLang's actual + construction order and `register_mem_pool_host`'s actual error + handling, both confirmed and resolved with the team rather than + picked unilaterally, since both are scope/policy calls, not + implementation details: (1) `extra_config["standalone_address"]` + cannot reliably activate standalone-process mode — verified, the + host memory pool is constructed and allocates its buffer in + `HiRadixCache.__init__` (`hiradix_cache.py:91,106`, + `pool_host/base.py:97,139`) *before* `extra_config` is parsed and + handed to `UMBPStore` (`hiradix_cache.py:175,182`), so + `umbp_host_allocator.py` can only ever see environment variables at + the point it must decide a buffer's backing. **Decision: v0.1 + activation is `UMBP_STANDALONE_ADDRESS`-only; + `extra_config["standalone_address"]` without the env var raises.** + (2) `register_mem_pool_host`'s existing exception/`False` handling + (`umbp_store.py:862-929`) silently downgrades to a warning — correct + for `distributed`, which has a real staging-buffer fallback, but + wrong for standalone-process mode, which has none in v0.1: a + swallowed registration failure there means every later `Put`/`Get` + fails with no correlated error. **Decision: under + `get_deployment_mode() == StandaloneProcess`, all three cases + (`disable_zero_copy_register` set, `register_memory` exception, + `register_memory() == False`) raise instead of warn; + `distributed`'s behavior is unchanged.** + +**Primary references (source-verified, see §9):** +- UMBP's own SPDK-proxy daemon (`storage/spdk/proxy/`, + `local/tiers/local_storage_manager.cpp:SpawnProxyDaemon/EnsureProxyDaemon`) + — the only existing "library spawns a separate process, backed by a + shared-memory ring protocol" precedent in this codebase. +- **Mooncake's Real Client / DummyClient split** + (`mooncake-store/src/real_client.cpp`, `real_client_main.cpp`, + `shm_helper.cpp`, `dummy_client.cpp`) — explicitly the primary model + for this design per team guidance ("Mooncake 是我们的老师"). +- LMCache's `lmcache/v1/multiprocess/` — reviewed and **not** used as + the primary template (its zero-copy path is CUDA-IPC/GPU-only and + has no host-memory equivalent; see §8.4). + +--- + +## 1. Why not just reuse the existing `distributed/` mode? + +UMBP already runs multi-process today (`umbp_master` + peers over +gRPC + RDMA). That mode solves a different problem: **cross-node** +sharing with RDMA as the bulk-transfer fabric, master-mediated routing, +and lease/eviction protocols built around remote peers that may be on +a different machine. Standalone-process mode is strictly **same-host**: +its entire job is to let N worker processes that already sit on one +NUMA machine share one DRAM tier without RDMA, without a master, and +without inventing new eviction protocols. Reusing `distributed/` +wholesale means dragging in a master, RDMA registration, and the +peer/lease machinery for a problem that doesn't need any of it — this +design intentionally sits *underneath* the existing `local/` +(`StandaloneClient`) code, not beside `distributed/`. + +Concretely: **the standalone-process server is `StandaloneClient`'s own +`LocalStorageManager`/`LocalBlockIndex`/`CopyPipeline`, unchanged, moved +into its own process and fronted by an RPC+shm shell.** No new tier +logic, no new eviction policy, no new durability logic — those are +100% reused from `local/`. + +--- + +## 2. Component overview + +Two new roles, named after Mooncake's split for consistency: + +| UMBP name | Mooncake analog | Lives in | Role | +|---|---|---|---| +| `umbp_standalone_server` | `mooncake_client` (Real Client) | its own process | Owns `LocalStorageManager` (DRAM/SSD tiers, private anon/hugetlb memory by default — see §3.2/§7, **not** a shared DRAM region), `LocalBlockIndex`, `CopyPipeline`. 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` | A new `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 — see §3.2/§8.3; gRPC cannot carry an fd) │ + │ │ + │ LocalStorageManager (REUSED, unmodified) │ + │ DRAMTier (own private anon/hugetlb memory, NOT shared — see │ + │ §3.2 note) ── SsdTier │ + │ LocalBlockIndex (REUSED, unmodified) │ + │ CopyPipeline (REUSED, unmodified; gains Drain(), see §4.3) │ + │ │ + │ Per-worker shm registry: client_id → {fd, base, size, offset} │ + └─────────────────────────────────────────────────────────────────┘ +``` + +Two logical channels, two socket paths, deliberately — see §3.2/§8.3 +for why gRPC cannot double as the fd-handoff channel. + +--- + +## 3. IPC mechanism + +### 3.1 Control plane: gRPC over a Unix domain socket + +**Decision: reuse UMBP's existing gRPC stack, not a new RPC framework, +and not Mooncake's `coro_rpc`.** + +Mooncake uses `ylt::coro_rpc` for its control plane. We explicitly do +**not** copy that choice: UMBP's `umbp_common` already depends on +gRPC/protobuf for the `distributed/` master/peer path +(`umbp/CMakeLists.txt:305-328,347-349`), and `pybind_umbp.cpp` already +has a proven, working convention for calling blocking C++ RPC methods +from Python with the GIL released +(`py::call_guard()`, +`pybind_umbp.cpp:248-279`). Adopting `coro_rpc` would require building +a new coroutine-to-pybind11 bridge with **no existing precedent in +either UMBP or the reports on Mooncake's own Python integration** — +pure added risk for zero benefit, since gRPC already does everything +Mooncake uses `coro_rpc` for (typed request/response, async server +loop, TCP or UDS transport). + +Design: + +- New proto service `UMBPStandalone` (new file, + `distributed/proto/umbp_standalone.proto`). **(v0.3 fix — this list + is now explicit about what is an RPC, what is proto-but-not-`IUMBPClient`, + and what stays entirely client-local, to avoid an implementer + guessing.)** + - 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`, **and + `GetExternalKvHitCounts`** (`umbp_client.h:166-170`; named + explicitly here since it's easy to drop — it has a default no-op + body in the interface, unlike its siblings, which makes it look + optional when it isn't for a real implementation). + - `Ping(Empty) -> PingResponse{ready: bool}`: new RPC, **no** + `IUMBPClient` analog, exists solely for the §4.1 readiness probe. + - **Not RPCs, stay entirely client-local** in + `StandaloneProcessClient`: `Close()` (tears down the local gRPC + channel + fd-handoff socket + registry state; no server-side call + needed beyond whatever `DeregisterMemory` already covers) and the + new `get_deployment_mode()` accessor from §5 (pure local constant, + never needs a round trip). + Reuse the existing struct-of-arrays batch codec approach from + `batch_resolve_codec.h` + (`distributed/peer/batch_resolve_codec.h:24-28`) for `BatchPut`/ + `BatchGet` so per-key protobuf submessage overhead doesn't reappear + here — this is exactly the pattern UMBP already validated for the + peer `BatchResolveKeys` RPC. +- **Default transport: gRPC over a Unix domain socket** + (`unix:///run/umbp/standalone/.grpc.sock`), not TCP. This is a + deliberate deviation from Mooncake (TCP, port 50052) and LMCache + (TCP, port 5555): both of those need to be reachable from a + different host or container in the general case, but standalone + mode by definition never leaves the host, so a UDS avoids picking a + port, avoids the network stack, and gets free filesystem-permission + based access control (§8.5). gRPC's UDS support + (`unix:` target scheme) is mature and already linked in via the same + gRPC dependency `umbp_common` uses today. + **(v0.2 fix) TCP is not a drop-in fallback for this design and must + not be presented as one.** §3.2's data plane depends on `SCM_RIGHTS` + fd passing, which is an `AF_UNIX`-only mechanism — there is no + TCP-socket equivalent for handing a `memfd` to another process, and + because the design deliberately uses an anonymous `memfd_create` + (not a named `shm_open` segment, precisely to avoid `/dev/shm` name + collisions, see §3.2), there is also no name the TCP-connected side + could use to reopen the segment out-of-band. Concretely, for v0.1: + the fd-handoff channel (§3.2, §8.3) **must** be a UDS reachable by + both processes (i.e. they must share the same mount/PID namespace + for `/dev/shm`-adjacent sockets, or a bind-mounted socket directory + across containers). If a deployment genuinely cannot provide that + (fully separate network namespaces with no shared filesystem), it + cannot use the zero-copy data plane at all; the only option is a + non-zero-copy fallback where `Put`/`Get` payload bytes travel inside + the gRPC message itself (i.e. the same "no RDMA, no shm" degraded + path `distributed/` already has via `staging_buffer_size`-capped + copies when `UMBP_DISABLE_ZERO_COPY_REGISTER=1`), not a variant of + the shm data plane over TCP. The gRPC *control* channel could in + principle still run over TCP in that scenario, but the two are no + longer both UDS by construction — this is a materially different, + slower deployment mode conceptually named `tcp_staging` to + distinguish it from the default `uds` path. **(v0.4 fix)** it has no + config/env surface in v0.1 at all, precisely because it isn't built + in v0.1 (§8 item 6b, §10 step 13) — see §6 for why a `transport` + knob is deliberately not added to `UMBPStandaloneProcessConfig` + ahead of the feature it would select. +- 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 today is + translated to `(client_id, shm_offset)` **inside** + `StandaloneProcessClient` before the RPC is issued (see §3.2) — this + is exactly Mooncake's `map_dummy_buffer_range_to_real` / + `shm_addr_offset` translation + (`real_client.cpp:1510-1583,3296-3318`), just carried over gRPC + instead of coro_rpc. + +### 3.2 Data plane: shared memory, modeled directly on Mooncake's `ShmHelper` + +This is the part of the design that should track Mooncake most +closely, because UMBP and Mooncake independently converged on the same +shape (control message carries a handle, bytes live in shared memory) +— UMBP's own SPDK-proxy already does this too +(`spdk_proxy_protocol.h:RingSlot.data_offset/data_size`). The novelty +here is *who allocates the shared segment*. + +**(v0.2 fix — this paragraph overreached in v0.1.)** `DRAMTier` does +already support POSIX shared memory — `use_shared_memory`/`shm_name` +in `UMBPDramConfig` (`common/config.h:57-58`) are wired through +`LocalStorageManager` (`local_storage_manager.cpp:443-444`) into +`DRAMTier`'s constructor, which does `shm_open`+`ftruncate`+`mmap`+ +`shm_unlink` (`dram_tier.cpp:104-175`). But two things in the v0.1 +draft were wrong about what this buys us: + +- The server's own DRAM cache **does not need to be visible to any + other process** — it is read/written exclusively by + `umbp_standalone_server` itself; no worker ever `mmap`s it directly + (workers only ever see their own registered host KV buffer, per + below). So `umbp_standalone_server` should default to the same + `kAnonymous`/`kAnonymousHugetlb` private backing + `StandaloneClient` already uses today — **not** + `use_shared_memory=true`. There is no reason to pay for a named + POSIX shm segment (and its cleanup/lifecycle concerns) for memory + nothing else needs to touch. +- Even where a future variant *did* want `use_shared_memory=true` (for + example, a warm-restart scenario that reattaches to a still-live + segment after the server process is replaced), the existing + `shm_open(shm_name_.c_str(), O_CREAT | O_RDWR, 0666)` call + (`dram_tier.cpp:114`) creates the segment world-read-write. That is + a pre-existing gap in `DRAMTier` itself (not introduced by this + design), and this design must not casually inherit it by pointing + more traffic at that code path while claiming "zero new C++ code" — + tightening `dram_tier.cpp`'s permissions is a separate, pre-existing + fix, out of scope here unless/until a future revision actually + chooses to share the server's DRAM tier across processes. + +The actual sharing requirement in this design is one-directional and +much narrower than "the whole DRAM tier": only the **worker's host KV +buffer** needs to become visible to the server process. That is the +problem the rest of this section solves. + +The problem this design actually has to solve is the **other** side: +today, `Put(key, src, size)`/`Get(key, dst, size)` take a pointer that +is only valid in the *caller's* address space +(`umbp_client.h:47-171`) — this is fine when `StandaloneClient` runs +in-process, and it is fine in `distributed/` mode because RDMA reads +directly from the registered remote memory without needing a shared +mapping. It is **not** fine when the server is a separate process on +the same host with no RDMA involved: the server cannot dereference a +pointer that only exists in the worker's page tables. + +Adopt Mooncake's answer exactly: + +1. **The worker process (not the server) owns the buffer**, exactly + like Mooncake's `DummyClient` owns its local buffer via + `setup_dummy()` → `shm_helper_->allocate(...)` + (`dummy_client.cpp:449-451`). Concretely: `UMBPHostMemAllocator` + (`host_mem_allocator.h/.cpp`, 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 (`shm_helper.cpp:102-123`), chosen over + UMBP's own named-`shm_open` DRAM-tier path specifically to avoid + `/dev/shm` name collisions across many SGLang instances on one + host (Mooncake's rationale for `memfd_create` over named + `shm_open` — anonymous, refcounted purely by open fds, nothing to + leak into `/dev/shm` if a process is killed uncleanly). +2. **One-time registration handshake**, triggered by the existing + `RegisterMemory(ptr, size)` call + (`umbp_store.py:912` → `register_mem_pool_host`). + + **(v0.2 fix — this step was underspecified in v0.1.)** + `IUMBPClient::RegisterMemory(uintptr_t ptr, size_t size)` + (`umbp_client.h:119`) and the pybind binding + `register_memory(ptr, size)` only carry a pointer and a size — there + is no fd parameter anywhere on this path, and `HostBufferHandle` + (`host_mem_allocator.h:41`) has no fd field either. The Python side + never sees a file descriptor (`kv_buffer.data_ptr()` is just an + integer, `umbp_store.py:896`), so it has nothing to pass through + pybind even if we wanted to add an argument. The fd has to be + recovered entirely on the C++ side of the process boundary that + already exists (`StandaloneProcessClient` and the allocator both + live in the same worker process), not threaded through Python: + + - `HostMemAllocator::Alloc` gains an internal, process-local + registry (a mutex-guarded `std::map` + where `AllocRecord` holds `{fd, size}`), populated only for + `kAnonymousShm` allocations at alloc time and erased on `Free`. + - `StandaloneProcessClient::RegisterMemory(ptr, size)` looks up + `ptr` in that registry (a range lookup, since `ptr` may be + `kv_buffer.data_ptr()` — the tensor's base address, which for the + host KV pool is normally exactly the allocation's base, but a + range/floor lookup is safer than requiring exact match) and, on a + hit, sends the recovered `fd` to the server via `SCM_RIGHTS` + ancillary data over the **separate raw-UDS fd-handoff socket** + (`.fd.sock`, distinct from the gRPC UDS — see §8.3), + mirroring `ipc_send_fd`/`ipc_recv_fd` (`shm_helper.cpp:181-230`) + and Mooncake's `IPC_SHM_REGISTER` handler + (`real_client.cpp:4046-4058`). On a miss (the pointer isn't + backed by a `kAnonymousShm` allocation — e.g. standalone-process + mode was configured but the host allocator wasn't switched to the + shm backing), `RegisterMemory` must fail 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` (mirrors `shm_contexts_`, + `real_client.cpp:4074-4082`). +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 + (mirrors `map_dummy_buffer_range_to_real`, + `real_client.cpp:1510-1583,3317`). No KV bytes ever cross the gRPC + channel; the gRPC message is a handful of integers per key. +4. `DeregisterMemory` is the inverse — **(v0.3 fix: fd ownership + clarified, since the v0.2 wording here conflicted with the §3.2 + step 2 registry design).** The allocator's ptr→fd registry (§3.2 + step 2), not `RegisterMemory`/`DeregisterMemory`, owns the fd: + `HostMemAllocator` opened it at `Alloc` time and is the only thing + that closes it, at `Free` time. `RegisterMemory` merely *borrows* + that fd to send over `SCM_RIGHTS`; sending a fd over `SCM_RIGHTS` + 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 of "this range + is registered" — it must **not** close the allocator-owned fd. + Closing happens exactly once, in `HostMemAllocator::Free`. Getting + this wrong either double-closes an fd number the kernel may have + already recycled for something else, or leaves a re-`RegisterMemory` + call after a `DeregisterMemory` trying to send an already-closed fd + — both are the kind of bug that only shows up under a + register/deregister/re-register cycle, so the unit test in §10 + step 11 must specifically exercise that cycle, not just a single + register→use→deregister pass. + +Net effect: the hot path becomes *exactly* the same shape as UMBP's +existing RDMA-registered `distributed/` path (`RegisterMemory` once, +then reference-by-handle on every call) — this is not a new interface +concept for UMBP, just a new transport underneath a pattern +`IUMBPClient` already exposes. + +### 3.3 Why not a busy-polled ring buffer (SPDK-proxy's own approach)? + +UMBP's own precedent for "separate process, shared-memory data plane" +(`spdk_proxy_protocol.h`) is a fixed-slot ring buffer with busy-polling +client-side (`spdk_proxy_tier.cpp:240,354`, `sleep_for(ProxyPollInterval())`) +and no RPC framework at all — every control operation (attach, admin +shutdown, stats) is also squeezed into the same 256-slot ring. This +works well for SPDK-proxy because it's low cardinality (few tenants, +few request types) and latency-critical at the microsecond level for +raw block I/O. It is a poor fit for the standalone-process control +plane because: + +- it has no natural way to express variable-shape requests (batch + registration, external-KV admission with variable-length hash + lists) without hand-rolling a second serialization format inside + the slot payload — gRPC already solves this, +- it offers no async multiplexing for slow control operations + without blocking the ring for the next request, +- 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 strictly for bulk bytes, which is the strategy this + design also adopts. + +The *data plane* still ends up looking like SPDK-proxy's idea (handle ++ offset/size, bytes in shm) — we're only rejecting the ring buffer +for the **control** plane, not the "shared memory for bytes" idea +itself. + +--- + +## 4. Lifecycle management + +This is the one dimension where UMBP's own precedent (SPDK-proxy: +library-spawned, idle-self-exit) and the Mooncake/LMCache precedent +(externally launched, long-lived, no self-exit) genuinely disagree, +and a specific choice has to be made for this codebase. + +**Decision: default to Mooncake's model (externally launched, +long-lived), with UMBP's SPDK-proxy auto-start pattern available as an +opt-in convenience, and idle-self-exit *disabled by default*.** + +Rationale: unlike the SPDK-proxy (a stateless passthrough to a shared +NVMe device — nothing is lost if it exits and respawns), the standalone +server **holds the DRAM cache tier's actual data**. Self-exiting on +idle silently discards the cache, defeating the purpose of the whole +feature. Neither Mooncake's Real Client nor LMCache's multiprocess +server self-terminate in the reports we reviewed — that absence is +itself evidence, not just a gap. + +### 4.1 Primary path: externally launched + +- Launched the same way `run_umbp_single_node_hicache.sh` already + launches `umbp_master` today (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 exact + naming convention of `UMBP_MASTER_ADDRESS` + (`runtime-env-vars.md:146-152`) for consistency. +- **fd-socket path derivation, made explicit (v0.3), confirmed to have + no separate config/env knob (v0.4 — see §6).** The fd-handoff socket + path is derived from `cfg.standalone_process.address` mechanically, + by both processes independently: strip the `unix://` scheme prefix + to get a filesystem path (raw `AF_UNIX` `connect()`/`bind()` need a + path, not a URI), then replace a trailing `.grpc.sock` with + `.fd.sock`; if the configured address doesn't end in `.grpc.sock` + (a custom path), append `.fd.sock` to it instead of doing a + substring replace. E.g. + `unix:///run/umbp/standalone/node0.grpc.sock` → + `/run/umbp/standalone/node0.fd.sock`. This derivation must live in + one shared place (`StandaloneProcessClient`'s ctor and + `umbp_standalone_server`'s startup both call it) rather than being + reimplemented on each side, to guarantee they agree — and precisely + because it's a pure function of `address`, it deliberately has no + independent `UMBP_STANDALONE_FD_SOCKET` env var or config field to + keep in sync (§6). +- **(v0.3 fix) `Ping` must be added to the `UMBPStandalone` proto + service.** The method list in §3.1 mirrors `IUMBPClient` 1:1 and + `IUMBPClient` has no `Ping` method, so a reader implementing §3.1's + proto as written would have nothing for this readiness probe to + call. `Ping(Empty) -> PingResponse{ready: bool}` is a new, + UMBP-standalone-specific addition to the proto with no `IUMBPClient` + analog — call it out explicitly in §3.1's method list, not just + here in §4.1. +- Readiness probe: connect to the UDS and call `Ping` + (equivalent to `ProxyShmRegion::ProbeExisting`'s role for + SPDK-proxy, but over gRPC instead of a shm header flag, since there + is no shm header to probe until a client registers memory). + +### 4.2 Convenience path: auto-start (opt-in) + +- **(v0.4 fix — routing corrected to match §6's decision.)** When + `cfg.standalone_process.auto_start` is `true` (Python-set from + `UMBP_STANDALONE_AUTO_START`/`extra["standalone_auto_start"]`, §5, + §6 — `CreateUMBPClient` itself never reads that env var; by the time + it runs, the bool is already sitting on the config it was handed): + 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 + already in `UMBPConfig::FromEnvironment` + (`config.h:410-428`) — no new election protocol) probes the UDS at + `cfg.standalone_process.address`; if absent, it does the same + `fork()` + `setsid()` + `execlp()` sequence + `LocalStorageManager::SpawnProxyDaemon` already uses + (`local_storage_manager.cpp:319-372`), passing the spawned server's + own config down via env vars (this is a different hop than the one + above — see `UMBP_STANDALONE_BIN` in §6, which *is* read directly by + this C++ auto-start code, matching its deployment-only status), and + waits for readiness bounded by + `cfg.standalone_process.startup_timeout_ms` (also Python-set, §6 — + mirrors `spdk_proxy_startup_timeout_ms`'s *role* only; unlike its + SPDK-proxy analog it is not itself read by + `UMBPConfig::FromEnvironment`, per §6's Precedent A decision). A + bootstrap lock (reuse the existing `ScopedBootstrapLock` pattern, + `local_storage_manager.cpp:274-283`) prevents a thundering herd of + workers all trying to spawn the server simultaneously. +- **Not** the default, because auto-start from inside a worker process + means the worker's `fork()`'d child inherits CUDA/HIP context, open + gRPC channels, and GPU device state at the moment of fork — exactly + the fork-safety hazard `SpawnProxyDaemon` already has to be careful + about (`local_storage_manager.cpp:328` calls `setsid()` immediately, + and the child does nothing but `execlp()` before replacing itself — + **this rule must not be relaxed** for the standalone-server spawn + path either; no CUDA/HIP/gRPC call may happen between `fork()` and + `execlp()`). + +### 4.3 Shutdown + +- `SIGTERM`/`SIGINT` → drain in-flight RPCs (bounded deadline, + mirroring `UMBP_GRPC_SHUTDOWN_DEADLINE_SEC` already used by + `MasterServer::Shutdown()`, `master_server.cpp:809-816`) → flush + `CopyPipeline` (best-effort persist DRAM-tier dirty pages to SSD + before exit, if the SSD tier is enabled and + `force_ssd_copy_on_write` is set) → `munmap` all registered client + shm segments → exit. + + **(v0.2 fix — no API exists today for the "flush CopyPipeline" + step above.)** `CopyPipeline` only drains its queue in its + destructor (`copy_pipeline.cpp:40-52`, `stop_copy_worker_` + + worker-thread join), and `StandaloneClient::Flush()` + (`standalone_client.cpp:300`) only calls `storage_.Flush()` — it + never waits for the async SSD-copy queue to empty. A clean shutdown + that wants "no dirty DRAM pages lost" needs the queue drained + *before* the object is destroyed, while the server can still respond + to health checks. This requires a small, genuinely new piece of + work, not a reuse of an existing hook: add + `CopyPipeline::Drain(std::chrono::milliseconds timeout)` (blocks + until `copy_queue_` is empty and in-flight copies complete, or the + timeout elapses) and have both `umbp_standalone_server`'s shutdown + path and the `Flush` RPC call it. See §10 implementation order — + this is listed as its own task, not folded into "reuse + `LocalStorageManager` unmodified" from §1, because it genuinely + isn't unmodified once this exists. +- 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 recommended + worker-side behavior is to treat this the same as "cache miss / + cache unavailable" rather than crashing the inference process + (SGLang's `HiCacheStorage` abstraction already tolerates backend + failures returning `False`/`None`; this needs to be threaded through + explicitly — see Open Question in §8.2, since neither Mooncake nor + LMCache's reviewed code shows a supervised-restart or + reconnect-on-crash path to copy). + +### 4.4 Multi-tenant workers on one host + +- Multiple SGLang worker processes (e.g. one per TP/DP rank) may + attach to the same standalone server. Each gets a `client_id` + (UUID, assigned at first `RegisterMemory` call) and its own shm + registration entry — mirrors Mooncake's per-client `shm_contexts_` + map and UMBP's own SPDK-proxy per-tenant quota concept + (`TenantInfo.quota_bytes`, `spdk_proxy_protocol.h:227-240`). + Capacity accounting inside `LocalStorageManager`/`DRAMTier` is + already global (not per-caller) today — **this is an open question**, + not solved by this design (see §8.3): without per-tenant quotas, one + worker can starve another's share of the shared DRAM tier. + +--- + +## 5. Python / pybind interface adaptation + +**(v0.2 fix — the v0.1 headline claim "zero call-site changes in +`umbp_store.py`" is false and is retracted here.)** One call site does +need to change, and the reason is a real bug the v0.1 draft +introduced by combining two of its own recommendations +inconsistently: §5 (v0.1) said `is_distributed()` should keep meaning +"true cross-node distributed" and *not* be overloaded for +standalone-process mode, while `register_mem_pool_host` +(`umbp_store.py:870-881`) gates the entire `register_memory` call on +`if not is_distributed: return` (`umbp_store.py:880`). Taken together, +a `StandaloneProcessClient` — which correctly reports +`is_distributed() == False` per the v0.1 recommendation — would never +have its host KV buffer registered, and therefore never trigger the +fd-handoff handshake in §3.2 at all. The whole standalone-process data +plane silently never activates. This is not a minor wording issue; it +is a functional gap that must be fixed as part of this design, not +left as an implementation detail: + +- Add `get_deployment_mode()` (enum: `Local`/`StandaloneProcess`/ + `Distributed`) to `IUMBPClient`/pybind, as already proposed below. +- Change the gate at `umbp_store.py:880` from `if not is_distributed: + return` to also allow `StandaloneProcess`, e.g. + `if deployment_mode not in (Distributed, StandaloneProcess): return`. + This *is* a call-site change in `umbp_store.py`, and needs to be + listed explicitly in the implementation order (§10), not implied. + +With that correction, the remaining adaptation is small: + +- `CreateUMBPClient(const UMBPConfig&)` (`umbp_client_factory.cpp:28-33`) + gains a third branch: `config.standalone_process.has_value()` (new + optional field on `UMBPConfig`, alongside the existing `distributed` + optional) → construct `StandaloneProcessClient`. Precedence, if both + happen to be set, should be `distributed` first (cross-node case is + strictly more general) — but in practice these are mutually + exclusive deployment choices and `UMBPConfig::Validate()` should + reject setting both. +- **(v0.3 fix, superseded by v0.5 below — kept for history.)** Nothing + in v0.2 said *who sets* `cfg.standalone_process`; v0.3 filled that in + by mirroring how `umbp_store.py` populates `cfg.distributed` from + either `extra_config["master_address"]` or `UMBP_MASTER_ADDRESS` + (`umbp_store.py:458-465`). +- **(v0.5 fix — the v0.3 answer above is wrong for this specific field, + for a reason that doesn't apply to `master_address`.)** Mirroring + `master_address`'s two-source pattern (`extra_config` *or* env) + assumes both sources are read at the same point in the SGLang + startup sequence. They are not, and the difference matters here + specifically because of a new constraint standalone-process mode + introduces that distributed mode never had: **verified**, in + `HiRadixCache.__init__`, the host memory pool + (`MHATokenToKVPoolHost`, via `pool_host/base.py`) is constructed with + `allocator_type=server_args.hicache_storage_backend` + (`hiradix_cache.py:91,106`) and immediately allocates `kv_buffer` in + its own constructor (`pool_host/base.py:97` `get_allocator_from_storage`, + `pool_host/base.py:139` `self.kv_buffer = self.init_kv_buffer()`) — + this happens *before* `storage_backend_extra_config` is even parsed + and handed to `CacheController`/`UMBPStore` (`hiradix_cache.py:175,182`, + well after the pool already exists). `UMBPHostTensorAllocator` (used + by that pool) therefore can only see **process environment + variables** at the moment it decides a buffer's backing — never + `extra_config`, which doesn't exist yet at that point in the process. + This asymmetry is new to standalone-process mode: `distributed` + mode's `RegisterMemory` pins/registers *whatever pointer it's given* + for RDMA regardless of how it was allocated, so `master_address` + arriving late via `extra_config` was always fine — nothing upstream + needed to know in advance. Standalone-process mode's + `RegisterMemory` instead needs the buffer to already be a + `kAnonymousShm` allocation *before* `register_mem_pool_host` ever + runs, and that decision is made by `umbp_host_allocator.py` strictly + earlier than `UMBPStore.__init__` gets to see `extra_config` at all. + **Decision (confirmed with the team): v0.1 supports activation via + `UMBP_STANDALONE_ADDRESS` only.** `extra_config["standalone_address"]` + is not a valid activation path in v0.1 and `UMBPStore.__init__` must + reject it outright with an actionable error rather than silently + accepting a config that the host allocator already missed: + ```python + if extra.get("standalone_address") and not _optional_env_str( + "UMBP_STANDALONE_ADDRESS" + ): + raise ValueError( + "standalone_address in hicache-storage-backend-extra-config is " + "not supported: the host memory pool allocator decides its " + "backing (Anonymous vs. AnonymousShm) before extra_config is " + "parsed, so extra_config alone cannot activate standalone-process " + "mode reliably. Set the UMBP_STANDALONE_ADDRESS environment " + "variable instead (see design-standalone-process-mode.md §5)." + ) + standalone_address = _optional_env_str("UMBP_STANDALONE_ADDRESS") + if standalone_address and master_address: + raise ValueError( + "master_address and UMBP_STANDALONE_ADDRESS are mutually " + "exclusive (distributed vs. standalone-process mode)" + ) + if standalone_address and UMBPStandaloneProcessConfig is not None: + standalone_cfg = UMBPStandaloneProcessConfig() + standalone_cfg.address = str(standalone_address) + # auto_start/startup_timeout_ms may still come from extra_config -- + # unlike `address`, they carry no allocator-timing constraint, + # since they only affect what CreateUMBPClient does once cfg is + # already fully built. Only `address` (the activation switch) is + # env-only; these two may keep the extra/env fallback pattern. + standalone_cfg.auto_start = _strict_bool( + extra.get("standalone_auto_start", _optional_env_str("UMBP_STANDALONE_AUTO_START")), + "standalone_auto_start", + ) if extra.get("standalone_auto_start") or _optional_env_str("UMBP_STANDALONE_AUTO_START") else False + standalone_cfg.startup_timeout_ms = int( + extra.get( + "standalone_startup_timeout_ms", + _optional_env_str("UMBP_STANDALONE_STARTUP_TIMEOUT_MS") or 30000, + ) + ) + cfg.standalone_process = standalone_cfg + ``` + This must run, and `cfg.standalone_process` must be populated, + **before** `UMBPClient(cfg)` is constructed (`umbp_store.py:790`) — + listed as its own step in §10 rather than folded silently into "wire + into `umbp_client_factory.cpp`", since that step is C++-only and + does not by itself make Python ever populate the field. The rejected + alternative — teaching `hiradix_cache.py`/`kv_cache_builder.py`/ + `memory_pool_host.py` to see `extra_config` before constructing the + host pool, so `extra_config` could activate standalone-process mode + too — was considered and explicitly deferred: it touches SGLang's + core pool-construction sequence outside this feature's natural + boundary, for a convenience (config-file-only activation) that + `UMBP_STANDALONE_ADDRESS` already covers. Revisit only if an env-var-only + activation switch turns out to be a real deployment blocker in + practice. +- **(v0.2 fix)** `pybind_umbp.cpp` is **not** change-free the way v0.1 + claimed. The `IUMBPClient` method binding table itself + (`put_from_ptr`/`get_into_ptr`/etc., `pybind_umbp.cpp:248-279`) + indeed needs no changes, since it dispatches through the vtable and + `StandaloneProcessClient` implements the same interface — that part + of the v0.1 claim was correct. But at least two *other* bindings in + the same file are missing and block this design from working at + all: + - `HostBufferBacking`/`UMBPHostBufferBacking` + (`pybind_umbp.cpp:42-45`) currently only exports `Anonymous` and + `AnonymousHugetlb` — it needs a third value, + `AnonymousShm`/`kAnonymousShm`, or `umbp_host_allocator.py` has no + way to ask for the new backing added in §3.2. + - `UMBPConfig`'s pybind class (`pybind_umbp.cpp:230-239`) only + exposes `dram`/`ssd`/`eviction`/`copy_pipeline`/`role`/ + `follower_mode`/`force_ssd_copy_on_write`/`distributed` — it needs + a new `standalone_process` property (and a new + `UMBPStandaloneProcessConfig` pybind class, mirroring how + `UMBPDistributedConfig` is bound at `pybind_umbp.cpp:218-228`) + before Python can construct or detect this config at all. + The GIL-release `call_guard` convention on the existing methods is + still correct as-is, since `StandaloneProcessClient`'s methods block + on a gRPC call and never call back into Python — same justification + as the existing comment (`pybind_umbp.cpp:245-246`, "block on RDMA, + SSD, or gRPC"). +- **(v0.5 fix — must read the raw env var directly, not `UMBPConfig`.)** + `umbp_host_allocator.py` needs one change: `UMBPHostTensorAllocator.__init__` + reads `os.environ["UMBP_STANDALONE_ADDRESS"]` directly (same style as + its existing `SGLANG_HICACHE_HOST_*` env reads, + `umbp_host_allocator.py:42-47`) — **not** anything derived from + `UMBPConfig` or `extra_config`, since neither exists yet at this + point in the process per the v0.5 fix above. If set, request the new + `AnonymousShm` backing from `UMBPHostMemAllocator.alloc(...)` instead + of the default anonymous/hugetlb backing. This is a small branch in + `UMBPHostTensorAllocator.allocate()` (`umbp_host_allocator.py:50`) — + the returned `handle.ptr` is used identically afterward (same + `ctypes.c_byte.from_address` + `torch.frombuffer` zero-copy wrap, + `umbp_host_allocator.py:86-87`), gated on the new `AnonymousShm` + pybind enum value existing (see above). +- **(v0.5 fix — the existing exception/false handling in + `register_mem_pool_host` was written for `distributed`'s optional, + gracefully-degradable RDMA registration and must not apply as-is to + standalone-process mode, which has no staging/inline fallback to + degrade to in v0.1.)** Verified: `register_mem_pool_host` + (`umbp_store.py:862-929`) currently (a) treats + `disable_zero_copy_register` as "log info, skip registration, stay + on staging path" (`umbp_store.py:~886-892`), (b) wraps the + `register_memory` call in `except Exception as exc: logger.warning(...); + return` (`umbp_store.py:~913`), and (c) treats a `False` return as + `logger.warning(...)` only, no raise (`umbp_store.py:~929`). All + three are correct for `distributed` — a failed/declined RDMA + registration there really does have a working fallback + (`distributed.staging_buffer_size`-capped copies). For + standalone-process mode none of these three can be a silent + downgrade, because there is nothing to downgrade *to*: a registry + miss means every subsequent `Put`/`Get` will fail against the + server's offset translation, just later and with a less + correlated-looking error. **Decision (confirmed with the team): when + `get_deployment_mode() == StandaloneProcess`, all three cases raise + instead of warn-and-return; `distributed`'s existing warn-and-fallback + behavior is unchanged.** Concretely, `register_mem_pool_host` branches + on `get_deployment_mode()` before each of the three points above: + ```python + is_standalone_process = ( + self.client.get_deployment_mode() == UMBPDeploymentMode.StandaloneProcess + ) + if is_standalone_process and getattr(self, "_disable_zero_copy_register", False): + raise RuntimeError( + "disable_zero_copy_register is not supported in standalone-process " + "mode: there is no staging-buffer fallback path to degrade to (v0.1)." + ) + ... + try: + ok = bool(self.client.register_memory(host_ptr, host_size)) + except Exception as exc: + if is_standalone_process: + raise RuntimeError( + f"register_memory failed in standalone-process mode: {exc}. " + "This is fatal, not a degradable condition -- see design-" + "standalone-process-mode.md §5." + ) from exc + logger.warning(...) # distributed: unchanged, falls back to staging + return + if not ok: + if is_standalone_process: + raise RuntimeError( + "register_memory returned false in standalone-process mode " + "(fatal -- no fallback path exists in v0.1)." + ) + logger.warning(...) # distributed: unchanged + ``` + `register_mem_pool_host`/`register_memory` call sites themselves + (`umbp_store.py:896-912`) still pass `(host_ptr, host_size)` + unchanged — the fd-passing handshake described in §3.2 happens + entirely inside `StandaloneProcessClient::RegisterMemory` in C++ via + the process-local fd registry (§3.2 step 2), invisible to Python. + Only the **gate** above it (`umbp_store.py:880`) and the + **error-handling branches** just described change, both listed as + their own items in §10. +- `is_distributed()` remains a pure bool getter, unchanged, and + continues to mean "true cross-node distributed" for the existing + external-KV branch points that already depend on it. The new + `get_deployment_mode()` accessor is not optional polish here — §5's + own fix above depends on it existing. + +--- + +## 6. New / reused configuration surface + +**(v0.4 fix — this section previously left it ambiguous whether each +knob is a `UMBPStandaloneProcessConfig` field set by Python or an env +var read directly by C++, and used both descriptions inconsistently +for the same knobs. Resolved below by picking one routing rule per +knob and sticking to it, using the two conflicting precedents already +in this codebase as the tie-breaker.)** + +Two existing precedents disagree on how "which deployment mode, and +with what parameters" config should flow, and this design has to pick +one per knob rather than blend them: + +- **Precedent A — `cfg.distributed`:** every sub-field + (`master_address`, `node_address`, `node_id`, ...) is read and + assigned **entirely in Python** (`umbp_store.py:458-495`), each via + the same `extra.get(key, _optional_env_str(ENV_VAR))` pattern. + `UMBPConfig::FromEnvironment()` never touches any of them. +- **Precedent B — `UMBPSsdConfig::spdk_proxy_auto_start`/ + `spdk_proxy_startup_timeout_ms`:** read **directly by + `UMBPConfig::FromEnvironment()`** in C++ (`config.h:399-402`, `UMBP_SPDK_PROXY_AUTO_START`/`UMBP_SPDK_PROXY_TIMEOUT_MS`), with no + Python involvement at all. + +**Decision: `standalone_process` follows Precedent A end to end** — +every field Python can set, it does, exactly like `distributed` +already works, rather than splitting the same config object across +two different parsing owners. Rationale: `standalone_process` is an +*activation* choice (like `distributed`), not a leaf tuning knob on an +already-active tier (like `spdk_proxy_auto_start` is, one level below +an already-enabled SSD tier) — it belongs with the precedent for the +thing it resembles, not the deepest-nested one that happened to be +read directly by C++. + +| Var | Consumed by | Purpose | Default | +|---|---|---|---| +| `UMBP_STANDALONE_ADDRESS` | **Python** (`UMBPStore.__init__`, §5) into `cfg.standalone_process.address` | UDS path (`unix:///run/umbp/standalone/.grpc.sock`) of the standalone server's **gRPC** socket. Presence enables standalone-process mode, mirroring how `UMBP_MASTER_ADDRESS` presence enables distributed mode today. | 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, §4.2). | `false` | +| `UMBP_STANDALONE_STARTUP_TIMEOUT_MS` | **Python** into `cfg.standalone_process.startup_timeout_ms` | Bound on waiting for readiness after spawn, mirrors `spdk_proxy_startup_timeout_ms`'s role (but is itself Python-set, per the decision above — it is not read the same way its SPDK-proxy analog is). | `30000` | +| `UMBP_STANDALONE_BIN` | **Deployment-only, not a config field** — read directly by the auto-start code path in C++ at spawn time, same status as `UMBP_MASTER_BIN`/`spdk_proxy_bin` per `runtime-env-vars.md`'s existing "deployment/launcher knobs, not parsed by `UMBPConfig`" category. | 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 (server-side only; never reaches a worker's `UMBPConfig` at all, so the Precedent A/B question doesn't apply to it) | `0` = never self-exit (default; see §4 rationale). Non-zero opts back into SPDK-proxy-style idle exit, only safe if SSD tier durability is also enabled. | `0` | +| `UMBP_STANDALONE_GRPC_SHUTDOWN_DEADLINE_SEC` | `umbp_standalone_server`'s own env read (server-side only) | Reuses the `MasterServer` shutdown-deadline convention. | same default as `UMBP_GRPC_SHUTDOWN_DEADLINE_SEC` | +| `UMBP_STANDALONE_SHM_DIR` | Both sides' own env read (deployment-only, not a config field — used only for the bootstrap-lock file path, mirrors `ScopedBootstrapLock`'s convention) | Directory for the bootstrap-lock file, not the shm segment itself (which is anonymous via `memfd_create`, never named on the filesystem — see §3.2). | `/tmp` | + +**(v0.4 fix) `fd_socket` and `transport` are explicitly *not* +`UMBPStandaloneProcessConfig` fields, and there is no +`UMBP_STANDALONE_FD_SOCKET`/`UMBP_STANDALONE_TRANSPORT` env var in +v0.1** — removing the ambiguity the v0.3 draft left open: + +- The fd-handoff socket path is **always mechanically derived** from + `cfg.standalone_process.address` using the deterministic rule + already given in §4.1 (strip `unix://`, replace trailing + `.grpc.sock` with `.fd.sock`, else append `.fd.sock`) — computed + identically inside `StandaloneProcessClient`'s ctor and + `umbp_standalone_server`'s startup, from the one `address` value + they already both have. No separate config surface, no env var, and + therefore nothing that can drift out of sync between the two + processes. If a real deployment later needs a non-derived override + (e.g. the fd socket must live on a different mount than the gRPC + one), add an explicit `fd_socket_override` field to + `UMBPStandaloneProcessConfig` *then*, as a v2 change — do not + pre-add it speculatively now. +- `transport` (`uds` vs. the deferred `tcp_staging`, §8 item 6b, §10 + step 13) has no config surface in v0.1 for the same reason + `tcp_staging` itself isn't built in v0.1: there is nothing yet for a + `transport` field to select between. When `tcp_staging` is actually + implemented in a v2 revision, add `cfg.standalone_process.transport` + (Python-set, Precedent A, for consistency with everything else in + this table) at that time — not before. + +`UMBPConfig` (C++) gains: `std::optional standalone_process;` +with exactly three fields — `address`, `auto_start`, +`startup_timeout_ms` — all three Python-set per the decision above, +same shape as the existing `UMBPDistributedConfig` +(`common/config.h:212-244`), so `ResolveRole()`/`Validate()` extend +naturally rather than needing a parallel code path. + +--- + +## 7. 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**; + standalone-process mode does not touch any file under + `distributed/master/`, `distributed/peer/`, or `distributed/routing/`. +- The shared code touched: `common/config.h` (new optional field), + `umbp_client_factory.cpp` (new branch), `local/host_mem_allocator.*` + (new `kAnonymousShm` backing plus the fd registry, additive), + `local/tiers/copy_pipeline.*` (new `Drain()`, see §4.3), + `pybind_umbp.cpp` (new `AnonymousShm` enum value + new + `standalone_process`/`UMBPStandaloneProcessConfig` bindings, see §5 + — **not** change-free, correcting the v0.1 claim), and + `umbp_store.py`'s `register_mem_pool_host` gate (one conditional, + see §5 — **not** call-site-free, correcting the v0.1 claim). +- A future extension (not in scope here, flagged for later): the + standalone server could itself become a `distributed/` peer/leader + (i.e. a host runs one standalone server that both serves local + workers over UDS *and* participates in the cross-node distributed + pool) — `LocalStorageManager`'s existing `SharedSSDLeader`/ + `SharedSSDFollower` roles already hint at this kind of layering, but + wiring it up is out of scope for v0.1. + +--- + +## 8. Open questions / risks + +1. **Crash / restart semantics are unsolved by every precedent + reviewed.** Neither LMCache's server, Mooncake's Real Client, nor + UMBP's own SPDK-proxy daemon has a supervised-restart or + reconnect-with-cache-rehydration story in the code we read. If the + standalone server dies, every attached worker loses its DRAM cache + simultaneously. Minimum bar for v0.1: workers must treat a broken + RPC channel as "cache unavailable," not a fatal error (falling back + to no-cache behavior, same as any other `HiCacheStorage` backend + failure) — this needs explicit handling in `StandaloneProcessClient` + and probably a small state machine (`CONNECTED` → `DISCONNECTED` → + optionally `RECONNECTING`), which does not exist in any of the + three reference systems today. +2. **Per-tenant DRAM quota is not implemented in `LocalStorageManager` + today.** Without it, multiple workers sharing one standalone + server's DRAM tier can starve each other. SPDK-proxy's + `TenantInfo.quota_bytes` concept (`spdk_proxy_protocol.h:227-240`) + is the natural template but would need to be added to + `DRAMTier`/`LocalBlockIndex`, which currently have no notion of + "owner" per key. +3. **(decided, not open as of v0.2, keeping the entry for traceability) + fd-passing handshake requires a second, separate socket path from + the gRPC UDS** (`.grpc.sock` vs `.fd.sock`, §2 diagram, + §6) since gRPC cannot itself carry a file descriptor. The v0.1 draft + said this but its own diagram contradicted it ("over the same + UDS") — that inconsistency is fixed in §2/§6/§9. What remains + genuinely open: whether the fd-handoff listener should be a + long-lived socket the server keeps open for the whole process + lifetime (simplest, matches Mooncake's `real_client_main.cpp:104`) + or a short-lived one opened only during an active registration + window (marginally more defensible against unauthorized local + connection attempts, at the cost of a startup race the client must + retry against). Proposal: long-lived, protected by filesystem + permissions (item 4 below) — simplicity over a marginal hardening + gain, revisit if a security review disagrees. +4. **Security/isolation**: the shm segment (memfd, `O_CLOEXEC` + recommended) and the UDS socket file must both be created with + `0600`/owner-only permissions, or any other local user can attach + to another tenant's KV cache. Neither Mooncake's nor UMBP's SPDK- + proxy code reviewed explicitly sets restrictive permissions on + their sockets/shm — this needs to be added deliberately, not copied. +5. **gRPC-over-UDS batch overhead vs. SPDK-proxy's ring buffer**: even + with struct-of-arrays batching, a gRPC call still costs a syscall + + protobuf encode/decode per `BatchGet`/`BatchPut`, which a busy-polled + shm ring buffer avoids entirely for the hottest path. If profiling + later shows this is a bottleneck for very small, very frequent Get + calls (e.g. single-block prefix-cache lookups), a v2 iteration could + add an optional shm-ring **fast path** for single-key Get/Put + alongside the gRPC path for everything else — deferred, not part of + v0.1, to avoid combining two IPC mechanisms before the simpler one + is proven to be insufficient. +6a. **(new in v0.2) The process-local ptr→fd registry (§3.2 step 2) + is new state with its own lifetime/thread-safety questions no + existing code needs to answer today.** `HostMemAllocator::Alloc`/ + `Free` (`host_mem_allocator.h:41-66`) are currently stateless + beyond the returned handle — adding a shared registry means: (a) it + must be safe for `Free` to run concurrently with a + `RegisterMemory` lookup from a different thread (mutex, as noted in + §3.2), (b) `DeregisterMemory` and `Free` ordering must be defined — + if a caller frees the underlying buffer before calling + `DeregisterMemory`, the registry entry and the server's `mmap` + would dangle; the design should make `Free` refuse (or at least + warn loudly) if an active registration still references that + pointer, rather than silently leaving the server with a stale + mapping. +6b. **(new in v0.2) The `tcp_staging` transport (§3.1, §6) is a real + second code path, not a config toggle on the same code.** Once + `Put`/`Get` payload bytes can travel two different ways (shm-offset + reference vs. inline bytes in the RPC message), `UMBPStandalone`'s + proto and `StandaloneProcessClient`'s implementation both need an + explicit branch, and testing needs to cover both — this roughly + doubles the data-plane test matrix for standalone-process mode + versus what v0.1 implied. Recommendation: **do not build + `tcp_staging` in v0.1** unless a concrete deployment already needs + it; ship UDS-only first and treat the inline-payload fallback as a + v2 addition once there's a real caller for it, rather than building + two data planes up front for a use case that may never materialize. +7. **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. This is the same + constraint `distributed/` already has between `umbp_master` and + worker processes, not a new problem — but standalone-process mode + is more likely to be deployed as a long-lived sidecar that outlives + several SGLang worker restarts/upgrades, which makes schema drift + more likely in practice than it is for `umbp_master` today. No + solution proposed here beyond "same binary/version pinning + discipline as the existing distributed mode." + +--- + +## 9. Source references used in this design + +All claims above trace to source read directly in this repo checkout, +not to general knowledge: + +- UMBP: `mori/src/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` + (`SpawnProxyDaemon`/`EnsureProxyDaemon`), + `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`, + `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`, local checkout): + `mooncake-store/src/real_client.cpp`, `real_client_main.cpp`, + `shm_helper.cpp`, `dummy_client.cpp`, `client_buffer.cpp`, + `docs/source/design/mooncake-store.md`, + `docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md`. +- LMCache (`/apps/nima/KVManager/LMCache`, local checkout, reviewed + for comparison only): `lmcache/v1/multiprocess/{mq,server,protocol, + custom_types,futures,config}.py`, `lmcache/v1/mp_observability/`. + +## 10. Suggested implementation order (for step 3, pending approval) + +**(Reordered/expanded across v0.2 and v0.3 to make each review-found +gap a concrete, individually-checkable task rather than an implicit +sub-step.)** + +1. `common/config.h`: add `UMBPStandaloneProcessConfig` + + `UMBPConfig::standalone_process`; extend `Validate()`/`ResolveRole()`. +2. `local/host_mem_allocator.{h,cpp}`: add `kAnonymousShm` backing + (`memfd_create`-based) **plus** the process-local ptr→fd registry + (§3.2 step 2, §8 item 6a) that `RegisterMemory` will query — these + belong in the same change since the registry is populated at + `Alloc`/`Free` time. +3. `local/tiers/copy_pipeline.{h,cpp}`: add `Drain(timeout)` (§4.3, + §8 item 7 fix) — small, independent of everything else, can land + first. +4. `src/pybind/pybind_umbp.cpp`: add `AnonymousShm` to + `UMBPHostBufferBacking`; bind `UMBPStandaloneProcessConfig` and + `UMBPConfig::standalone_process`; add `get_deployment_mode()` to + the `IUMBPClient` binding (§5). Do this **before** step 6 below — + the Python-side change depends on these bindings existing. +5. New proto `distributed/proto/umbp_standalone.proto` + + generated-code wiring in `CMakeLists.txt`; struct-of-arrays batch + codec reuse for `BatchPut`/`BatchGet` (§3.1). +6. New `umbp_standalone_server` binary (own `main.cpp`, thin wrapper + around a new `StandaloneServer` class that owns + `LocalStorageManager`/`LocalBlockIndex`/`CopyPipeline` — reused, not + reimplemented, with `DRAMTier` defaulting to private + anon/hugetlb per §3.2/§7, not `use_shared_memory=true`) — plus the + gRPC service on `.grpc.sock` and the **separate** raw-UDS + fd-handoff listener on `.fd.sock` (§2, §8 item 3). +7. New `StandaloneProcessClient : IUMBPClient` (worker-side): gRPC + stub, offset translation, the `SCM_RIGHTS` send side of the + fd-handoff handshake, wired into `umbp_client_factory.cpp`. + `RegisterMemory` must fail loudly (not silently degrade) when the + registry lookup misses (§3.2 step 2). +8. **(v0.3 addition, corrected in v0.5 — this step did not exist + before and is the activation path itself.)** `umbp_store.py`: in + `UMBPStore.__init__`, next to the existing `master_address` block + (`umbp_store.py:458-465`): **first**, raise if + `extra_config["standalone_address"]` is set without + `UMBP_STANDALONE_ADDRESS` also being set (v0.5 — extra_config alone + cannot activate this mode, see §5); **then** read + `UMBP_STANDALONE_ADDRESS`, construct `UMBPStandaloneProcessConfig`, + assign `cfg.standalone_process`, and raise if both `master_address` + and the standalone address are set (§5). Depends on step 4 (pybind + bindings for `UMBPStandaloneProcessConfig` must exist first). + Without this step, every other step in this list can be implemented + correctly and `UMBPStore` will still silently construct a plain + `StandaloneClient` forever — this is not optional polish, it is the + switch that turns the feature on. +9. **(v0.5 fix — must read the env var directly, at allocator + construction time, before `extra_config` exists — see §5.)** + `umbp_host_allocator.py`: `UMBPHostTensorAllocator` reads + `UMBP_STANDALONE_ADDRESS` directly from `os.environ` (not via + `UMBPConfig`/`extra_config`) and branches to request `AnonymousShm` + backing when set (depends on step 4 for the pybind enum value). + This step and step 8 both gate on the *same* env var by + construction, precisely so they can never disagree about whether + standalone-process mode is active. +10. **`umbp_store.py`: fix the `register_mem_pool_host` gate at line + 880** from `is_distributed()`-only to also accept + `get_deployment_mode() == StandaloneProcess` (§5) — called out as + its own task because it is a functional bug an implementation + could ship with even after step 8 above: the client would be + correctly constructed as `StandaloneProcessClient`, but still + never register its host buffer, so every `Put`/`Get` would fail at + the server's offset-translation step with no obvious cause. +10a. **(v0.5 addition.)** `umbp_store.py`: in the same method, make + `disable_zero_copy_register`-set / `register_memory` exception / + `register_memory() == False` all `raise` instead of + warn-and-return when `get_deployment_mode() == StandaloneProcess` + (§5); leave `distributed`'s existing warn-and-fallback behavior + untouched. Land this together with step 10 — both touch the same + function and the same review found both gaps together. +11. Lifecycle: auto-start path reusing `SpawnProxyDaemon`'s fork/exec + pattern; shell script analogous to `run_umbp_single_node_hicache.sh` + for the externally-launched path. +12. Tests: unit tests for the fd-handoff handshake and offset + translation, the ptr→fd registry's `Free`/`DeregisterMemory` + ordering including a register→deregister→re-register cycle (§3.2 + step 4, §8 item 6a), `CopyPipeline::Drain()` under load; an + integration test with 2+ worker processes sharing one standalone + server on one host, including a kill-the-server test that asserts + workers degrade to cache-unavailable rather than crashing (§8 + item 1), and an end-to-end test that `UMBP_STANDALONE_ADDRESS` + alone (via `UMBPStore` config, no other code changes) is sufficient + to activate the mode (guards against step 8 regressing silently). + **(v0.5 additions)** a Python test asserting + `extra_config["standalone_address"]` without the env var raises at + `UMBPStore.__init__` (guards step 8's rejection path); and Python + tests for each of the three `register_mem_pool_host` cases (§5, + step 10a) asserting they raise under `StandaloneProcess` and still + only warn under `Distributed` (guards against the two modes' + branches getting merged back together in a future edit). +13. **Deferred to v2, not v0.1** (§8 item 6b): `tcp_staging` transport + with inline RPC payload bytes. Build only if a concrete deployment + needs it. diff --git a/src/umbp/include/umbp/common/config.h b/src/umbp/include/umbp/common/config.h index f29ea1a06..b400be96c 100644 --- a/src/umbp/include/umbp/common/config.h +++ b/src/umbp/include/umbp/common/config.h @@ -243,6 +243,15 @@ 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 +}; + struct UMBPConfig { UMBPDramConfig dram; UMBPSsdConfig ssd; @@ -262,6 +271,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 +342,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/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..eed29bef9 --- /dev/null +++ b/src/umbp/standalone/bin/standalone_server_main.cpp @@ -0,0 +1,102 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#include + +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/standalone/ipc.h" +#include "umbp/standalone/standalone_server.h" + +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.distributed.reset(); + + 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(); + } + + mori::umbp::UMBPStandaloneProcessConfig standalone_cfg; + standalone_cfg.address = address; + config.standalone_process = standalone_cfg; + + mori::umbp::standalone::StandaloneServer server(config, address); + + 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, &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 {}", address); + server.Run(); + + stop_signal_waiter = true; + signal_waiter.join(); + MORI_UMBP_INFO("[StandaloneServer] exited cleanly"); + return 0; +} 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..a7c2a5ef0 --- /dev/null +++ b/src/umbp/standalone/standalone_process_client.cpp @@ -0,0 +1,600 @@ +// 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); + ::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)); + ::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)); + ::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)); + ::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); + ::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); + ::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..5af6e42e0 --- /dev/null +++ b/src/umbp/standalone/standalone_server.cpp @@ -0,0 +1,618 @@ +// 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 "mori/utils/mori_log.hpp" +#include "umbp/local/standalone_client.h" +#include "umbp/standalone/ipc.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; +} + +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) + : client_(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()); + } + { + std::lock_guard lock(client_mu_); + client_.Flush(); + client_.Close(); + } + UnmapAll(); + 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()); + 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 lock(memory_mu_); + auto it = memory_.find(request->client_id()); + bool ok = it != memory_.end() && it->second.worker_base == request->worker_base() && + it->second.size >= request->size(); + SetBool(response, ok, ok ? "" : "fd handoff registration was not found"); + return grpc::Status::OK; + } + + grpc::Status DeregisterMemory(grpc::ServerContext*, + const ::umbp::DeregisterMemoryRequest* request, + ::umbp::Empty*) override { + UnmapClient(request->client_id()); + return grpc::Status::OK; + } + + grpc::Status ReportExternalKvBlocks(grpc::ServerContext*, + const ::umbp::StandaloneExternalKvMutationRequest*, + ::umbp::BoolResponse* response) override { + SetBool(response, true); + return grpc::Status::OK; + } + + grpc::Status RevokeExternalKvBlocks(grpc::ServerContext*, + const ::umbp::StandaloneExternalKvMutationRequest*, + ::umbp::BoolResponse* response) override { + SetBool(response, true); + return grpc::Status::OK; + } + + grpc::Status RevokeAllExternalKvBlocksAtTier(grpc::ServerContext*, + const ::umbp::StandaloneExternalKvTierRequest*, + ::umbp::BoolResponse* response) override { + SetBool(response, true); + return grpc::Status::OK; + } + + grpc::Status MatchExternalKv(grpc::ServerContext*, + const ::umbp::StandaloneMatchExternalKvRequest*, + ::umbp::StandaloneMatchExternalKvResponse*) override { + return grpc::Status::OK; + } + + grpc::Status GetExternalKvHitCounts(grpc::ServerContext*, + const ::umbp::StandaloneExternalKvHitCountsRequest*, + ::umbp::StandaloneExternalKvHitCountsResponse*) override { + return grpc::Status::OK; + } + + private: + struct RegisteredMemory { + void* base = nullptr; + uint64_t worker_base = 0; + uint64_t size = 0; + }; + + 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; + } + + std::string client_id(msg.client_id); + std::lock_guard lock(memory_mu_); + auto old = memory_.find(client_id); + if (old != memory_.end() && old->second.base) { + munmap(old->second.base, static_cast(old->second.size)); + } + memory_[client_id] = RegisteredMemory{mapped, static_cast(msg.worker_base), msg.size}; + 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::lock_guard lock(memory_mu_); + auto it = memory_.find(client_id); + if (it == memory_.end()) return; + if (it->second.base) munmap(it->second.base, static_cast(it->second.size)); + memory_.erase(it); + } + + void UnmapAll() { + std::lock_guard lock(memory_mu_); + for (auto& kv : memory_) { + if (kv.second.base) munmap(kv.second.base, static_cast(kv.second.size)); + } + memory_.clear(); + } + + 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); + } + + StandaloneClient 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::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); } From c026dcb650a29c9ab3e3f3a0b5cf8f63fe1585aa Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Thu, 9 Jul 2026 09:33:22 +0000 Subject: [PATCH 2/3] Add distributed-backed UMBP standalone mode --- src/pybind/pybind_umbp.cpp | 5 +- src/umbp/CMakeLists.txt | 1 + .../distributed/proto/umbp_standalone.proto | 14 + .../doc/design-standalone-process-mode.md | 1436 ++++++++++++++++- src/umbp/include/umbp/common/config.h | 7 + .../standalone/external_kv_identity_client.h | 91 ++ .../standalone/bin/standalone_server_main.cpp | 182 ++- .../external_kv_identity_client.cpp | 300 ++++ .../standalone/standalone_process_client.cpp | 8 + src/umbp/standalone/standalone_server.cpp | 324 +++- 10 files changed, 2292 insertions(+), 76 deletions(-) create mode 100644 src/umbp/include/umbp/standalone/external_kv_identity_client.h create mode 100644 src/umbp/standalone/external_kv_identity_client.cpp diff --git a/src/pybind/pybind_umbp.cpp b/src/pybind/pybind_umbp.cpp index f1c607318..a4b28e17d 100644 --- a/src/pybind/pybind_umbp.cpp +++ b/src/pybind/pybind_umbp.cpp @@ -239,7 +239,10 @@ void RegisterMoriUmbp(py::module_& m) { .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("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<>()) diff --git a/src/umbp/CMakeLists.txt b/src/umbp/CMakeLists.txt index 5afd1ae08..5b02247aa 100644 --- a/src/umbp/CMakeLists.txt +++ b/src/umbp/CMakeLists.txt @@ -348,6 +348,7 @@ add_library( distributed/peer/peer_service.cpp distributed/pool_client.cpp distributed/distributed_client.cpp + standalone/external_kv_identity_client.cpp standalone/ipc.cpp standalone/standalone_process_client.cpp standalone/standalone_server.cpp) diff --git a/src/umbp/distributed/proto/umbp_standalone.proto b/src/umbp/distributed/proto/umbp_standalone.proto index a002e0716..4cb6580e2 100644 --- a/src/umbp/distributed/proto/umbp_standalone.proto +++ b/src/umbp/distributed/proto/umbp_standalone.proto @@ -32,6 +32,7 @@ message Empty {} message PingResponse { bool ready = 1; + StandaloneBackendMode deployment_mode = 2; } message BoolResponse { @@ -39,6 +40,12 @@ message BoolResponse { string error = 2; } +enum StandaloneBackendMode { + STANDALONE_BACKEND_UNKNOWN = 0; + STANDALONE_BACKEND_LOCAL = 1; + STANDALONE_BACKEND_DISTRIBUTED = 2; +} + message CountResponse { uint64 count = 1; } @@ -88,6 +95,9 @@ 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 { @@ -97,15 +107,18 @@ message DeregisterMemoryRequest { 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 { @@ -120,6 +133,7 @@ message StandaloneMatchExternalKvResponse { message StandaloneExternalKvHitCountsRequest { repeated string hashes = 1; + string client_id = 2; } message StandaloneExternalKvHitCountEntry { diff --git a/src/umbp/doc/design-standalone-process-mode.md b/src/umbp/doc/design-standalone-process-mode.md index 5972a2ab7..b07d89282 100644 --- a/src/umbp/doc/design-standalone-process-mode.md +++ b/src/umbp/doc/design-standalone-process-mode.md @@ -1,4 +1,218 @@ -# UMBP Standalone Process Mode — Design (Draft v0.5) +# UMBP Standalone Process Mode — Design (Draft v1.4) + +**v1.0 revision note:** §11.4.4 (external-KV) went through two rounds +of direct clarification and was substantially corrected, not just +touched up — worth reading in full even if earlier revisions of this +section were already reviewed. Summary of the arc: v0.7 treated +external-KV as a data-plane feature and concluded the server can't +physically serve HBM-tier reads, so it decided to disable external-KV +entirely for a distributed-backed server. That premise was wrong: +external-KV is a pure query/registry service (`Report`/ +`MatchExternalKv` just record and answer "who holds hash H," a caller +— typically a custom SGLang router — uses the answer to route a *new* +request elsewhere; UMBP itself never moves the reported bytes, in +`distributed/` mode or here). With that corrected, the real question +became identity granularity, and a second clarification established +that the project's actual deployment shape is 8-GPU hosts running +SGLang TP=8, **and** that the *baseline* non-standalone implementation +already supports multiple independent replicas sharing one host for +free (`node_id` is built from per-process rank coordinates with no +awareness of logical TP/DP group boundaries, `umbp_store.py:512-519`). +Per explicit instruction to match the baseline rather than settle for +a narrower capability, §11.4.4 now specifies **per-worker distributed +sub-identities** (the server runs one independent, external-KV-only +identity per connected worker, sharing one physical `peer_address` — +verified this pass that `peer_address` uniqueness is not required +anywhere in `ClientRegistry`/`MasterServer` — see §11.4.4 for +citations) rather than either disabling the feature or collapsing all +workers into one shared identity. **(Corrected in v1.2: this is +implemented by a new, purpose-built `ExternalKvIdentityClient` class — +not `MasterClient` — specifically so each sub-identity reports empty +`tier_capacities`, publishes no owned-KV events, and is therefore +structurally ineligible for ordinary routing/eviction; see §11.4.4 and +§11.4.7 for the full reasoning and interface. The paragraph above +originally said "`MasterClient` registration," which is no longer +accurate — kept here only so the revision history stays honest about +what changed and why.)** This is real, non-trivial, correctly-scoped +implementation work (up to N `ExternalKvIdentityClient` +registrations/heartbeats inside one process), not a simple flag flip — +it was **not** in the v0.7/v0.9 implementation-cost accounting and +should be weighed accordingly when this proposal is scheduled. + +**v0.9 revision note:** a fourth review pass raised two further +points, both accepted as valid (not over-engineering): (1) +`PingResponse.deployment_mode` was wrongly filed as a low-stakes, +deferrable nice-to-have — it is actually the only guard against the +single most dangerous silent-failure mode this whole section exists to +prevent (a parity-intended deployment silently ending up local-backed, +with every RPC still succeeding); promoted to required, with a new +assertion requirement for distributed-backend launch/test tooling. (2) +§11.4.6's env-var list for the server's own distributed backend config +was incomplete (listed 6 vars, the real `umbp_store.py` surface has ~12 +including SSD staging and cache-admission knobs) and, separately, +missed that `dram_page_size`'s normal auto-derivation depends on a +`mem_pool_host` the server never has — both fixed, with a proposed +`BuildDistributedBackendConfigFromEnv()` helper shape and an explicit +required/optional split. See inline `(v0.9 fix)` markers in §11.4.5 and +§11.4.6. + +**v0.8 correction pass (citations only, no design changes):** every +Mooncake source citation in this document was originally verified +against `/apps/theresa/Mooncake` (commit `e1d6d6f6`, 2026-04-22). That +was the wrong path — the correct, current Mooncake checkout for this +project is `/apps/nima/KVManager/Mooncake` (commit `c9896684`, +2026-07-08, 75 commits ahead on the relevant files). All citations +were re-verified against the correct path this pass. Outcomes: +- Most citations were **directionally correct but line-stale** + (functions renamed/refactored, e.g. `map_shm_internal` is now a thin + wrapper around `map_shm_internal_with_device`; `ipc_send_fd`/ + `ipc_recv_fd` were renamed to `UdsConnection::sendFd`/`recvFd` and + moved into `uds_transport.cpp`) — corrected in place, marked + `(v0.7 citation fix)` inline (the fixes were made while the doc was + still at v0.7; the version bump to v0.8 here just reflects the + re-verification pass being complete). +- **One substantive error was found and corrected, not just a stale + line number**: this document had described Mooncake's primary shm + buffer as "the Real Client allocates it and hands the fd to the + application." The correct direction, confirmed against the current + checkout, is the opposite for the primary buffer: the **application** + (`DummyClient`) allocates and owns it, and sends the fd to the Real + Client, which is the passive second mapper. (A secondary, read-only + "hot cache" segment does go the other direction, but that is not + the buffer this design's citations were describing.) Reassuringly, + **UMBP's own already-shipped implementation never had this error** — + §3.2 already correctly has the worker (application) allocate and own + the buffer, matching Mooncake's real direction; only this document's + *prose describing* Mooncake's mechanism had drifted from the + primary-source direction. No code changes result from this + correction, only doc text (§11.3's "closed open item" section). +- The auto-start precedent claim ("zero matches" for fork/exec/spawn + patterns across `mooncake-integration/` and `mooncake-wheel/`) was + also corrected: `mooncake-integration/` (the actual application- + binding code) is still zero matches, but `mooncake-wheel/` (the + separate CLI/packaging layer) contains one real match — an explicit, + human-invoked CLI passthrough, not an implicit auto-start. The + underlying design decision (§11.4.6: distributed-backed + `umbp_standalone_server` supports external launch only, no + auto-start) is unaffected; only the supporting claim's precision was + corrected. + +**v1.1 revision note:** confirmed §11.4.4's per-worker `MasterClient` +sub-identity design (Option A) as final, after clarifying it involves +no worker→server→Master relay (each sub-identity heartbeats Master +directly from inside the server) and after explicitly weighing it +against a simpler Mooncake-aligned alternative (Option B: one shared +identity + `tenant_id`-style key namespacing, the pattern Mooncake +actually uses for its own multi-`DummyClient`-per-`RealClient` case). +Option B was rejected: preserving the `distributed/` baseline's +existing per-worker routing precision (verified as a real, already- +existing capability, not a nice-to-have) was judged more valuable than +the simplicity Option A gives up — accepted as genuinely new +engineering with no direct precedent in either UMBP or Mooncake, not a +known-safe pattern being reused. See the `(v1.1 —` marker in §11.4.4. + +**v1.2 revision note:** the v1.1 per-worker sub-identity design was +found to have three confirmed blockers, all fixed by replacing "N +independent `MasterClient` registrations" with a new, purpose-built +`ExternalKvIdentityClient` class: (A) a naive `MasterClient`-based +sub-identity with nonzero `tier_capacities` would have been eligible +for ordinary `RoutePut` target selection, silently polluting real KV +placement with fake storage nodes — closed structurally by giving the +new class no `tier_capacities`-reporting code path at all, rather than +relying on every future call site remembering to pass empty capacities +into `MasterClient`. (B) `MasterClient`'s existing re-register path has +a confirmed bug — it omits `peer_address`/`engine_desc` on re-register, +silently breaking `MatchExternalKv` responses after a Master +restart/reaper-expiry — closed by not reusing that code path at all; +the new class always includes both fields, by construction, rather +than patching a class shared by every real distributed peer. (C) the +registration-handshake field list was under-specified ("just +`node_id`") when the proto actually requires `node_address` too — +decided now: `worker_node_id` + `worker_node_address` + optional +`tags`, Python remains the single source of truth for deriving them, +no raw rank components cross the boundary. `ExternalKvIdentityClient` +itself is flagged as real, net-new implementation work not previously +accounted for in this proposal's cost estimate (§11.7 item 7). + +**v1.3 revision note:** a documentation-consistency and finalization +pass, no architecture changes from v1.2. Six items closed: (1) the +v1.0 revision note at the top of this document still said "`MasterClient` +registration per connected worker," contradicting §11.4.4's v1.2 +`ExternalKvIdentityClient` fix — corrected in place, kept visible +rather than silently rewritten. (2) `ExternalKvIdentityClient` was +still filed as an open item needing "a design pass" even though +§11.4.4 already made it a required component — it now has a full +interface/lifecycle spec in new §11.4.7 (Register/Heartbeat/Unregister/ +the five external-KV RPCs, the re-registration hard requirement, and +shutdown ordering). (3) confirmed via `eviction_manager.cpp:83-91` +(checked directly, not assumed) that the same empty-`tier_capacities` +invariant that hides a sub-identity from `RoutePut` also hides it from +`EvictionManager` — no separate mechanism needed, but stated +explicitly now rather than left to be inferred; also explicitly +decided `ExternalKvIdentityClient` heartbeats must never carry +`EventBundle`s, a separate `HeartbeatRequest` field from +`tier_capacities`. (4) the worker→server identity handshake now has a +concrete wire location: `worker_node_id`/`worker_node_address`/`tags` +are added to the existing `RegisterMemoryRequest` message +(`umbp_standalone.proto:87-91`), not a new RPC. (5) §11.4.6's five +new server-backend env vars are finalized as +`UMBP_DISTRIBUTED_STAGING_BUFFER_SIZE` and siblings (the +`UMBP_DISTRIBUTED_` prefix specifically because these five have no +existing worker-side env-var form to collide with, unlike the ones +reusing worker-side names verbatim) — no longer "proposed." (6) added +an explicit layering statement (new §11.4.7 lead-in): distributed- +backed server + UDS/shm fd handoff is the Mooncake-parity core; +`ExternalKvIdentityClient`/per-worker external-KV identity is a +separate UMBP-compatibility extension preserving a `distributed/`- +baseline capability, not required for declaring parity achieved, and +schedulable/testable independently. + +**v1.4 revision note:** two updates from the first implementation + +code-review round on `ExternalKvIdentityClient`. (1) A concurrency bug +found by review — `EnsureExternalIdentity`/`RemoveExternalIdentity` had +no serialization across their read-check/stop-old/construct-and-start- +new/insert sequence, so a `DeregisterMemory` racing an in-flight +`RegisterMemory` for the same `client_id` could silently fail to +remove an identity the register path was about to (re-)insert (a +permanent leak until full server restart), and two concurrent +`RegisterMemory` calls for the same `client_id` could briefly +double-register — was fixed with a single server-wide +`external_identity_lifecycle_mu_` serializing all three call sites; +§11.4.7 now documents this as an accepted trade-off (serializes +otherwise-unrelated workers' register/deregister calls when the lock +covers a slow/unavailable Master RPC, worst case approaching +`N * rpc_deadline` for the stated 8-worker deployment shape) rather +than a free fix, with an explicit note on when to revisit it +(per-`client_id` sharding, or decoupling external-KV registration from +the core memory-registration path). (2) §11.4.7's `Unregister()` +trigger list previously claimed worker crash/disconnect detection was +"already" handled by mirroring a `StandaloneProcessClient` state +machine — but §8 item 1 itself explicitly says that state machine +doesn't exist. This was an internal contradiction, not just a missing +callout: corrected to state plainly that no crash detection exists for +`ExternalKvIdentityClient` today, what the concrete consequence is (a +crashed worker's identity heartbeats forever, `ClientRegistry` never +expires it, only a full server restart clears it), and that this is a +deliberately deferred, tracked limitation rather than a silent gap. + +**Status (still true as of v1.4): §11 below is a NEW, UNAPPROVED +proposal. Nothing in §11 is implemented. Do not start coding against +it until it has been discussed — this revision exists to make the +proposal concrete enough to discuss, per explicit instruction not to +touch code first. v0.7 closed 4 of the 5 blockers a second review +pass found in the v0.6 draft (config/auto-start conflict, +external-KV routing correctness, backend registration +transaction/rollback + shutdown ordering, and the cross-mapping +memory-visibility open item); one naming/scope-framing question +remains genuinely open (§11.7 item 1).** §1-§10 +describe what is already implemented and shipped (v0.1-v0.5); §11 +revises the framing in §1 (see the callout there) to close a scope gap +identified after implementation: v0.1-v0.5 built a same-host-only +subsystem, whereas the stated goal was Mooncake-parity, and Mooncake's +"standalone" Real Client is cross-node-capable (it is a full +distributed client that happens to run in its own process, not a +narrower local-only subsystem). §11 proposes closing that gap without +discarding the v0.1-v0.5 implementation. **Scope:** A new deployment mode for UMBP in which the DRAM/SSD tiers, `LocalBlockIndex`, and `CopyPipeline` run inside a dedicated, long-lived @@ -101,6 +315,11 @@ rather than silently overwritten. ## 1. Why not just reuse the existing `distributed/` mode? +**(v0.6 note — this section's framing is revised by §11, not replaced. +Read this section for why v0.1-v0.5 built what they built; read §11 +for why that scope turned out to be narrower than "Mooncake parity" +requires, and how the two are reconciled without a rewrite.)** + UMBP already runs multi-process today (`umbp_master` + peers over gRPC + RDMA). That mode solves a different problem: **cross-node** sharing with RDMA as the bulk-transfer fabric, master-mediated routing, @@ -269,9 +488,13 @@ Design: translated to `(client_id, shm_offset)` **inside** `StandaloneProcessClient` before the RPC is issued (see §3.2) — this is exactly Mooncake's `map_dummy_buffer_range_to_real` / - `shm_addr_offset` translation - (`real_client.cpp:1510-1583,3296-3318`), just carried over gRPC - instead of coro_rpc. + `shm_addr_offset` translation **(v0.7 citation fix — re-verified + against the correct checkout; function confirmed still present, + only line numbers moved from a stale reference)**: + `real_client.cpp:369-379` (`map_dummy_buffer_range_to_real`), called + from the batch RPC handlers `get_into_range_shm_helper`/ + `get_into_ranges_shm_helper` (`real_client.cpp:4291`/`:4316`), just + carried over gRPC instead of coro_rpc. ### 3.2 Data plane: shared memory, modeled directly on Mooncake's `ShmHelper` @@ -331,13 +554,21 @@ Adopt Mooncake's answer exactly: 1. **The worker process (not the server) owns the buffer**, exactly like Mooncake's `DummyClient` owns its local buffer via - `setup_dummy()` → `shm_helper_->allocate(...)` - (`dummy_client.cpp:449-451`). Concretely: `UMBPHostMemAllocator` + `setup_dummy()` → `shm_helper_->allocate(...)` **(v0.7 citation fix + — re-verified against the correct Mooncake checkout, + `/apps/nima/KVManager/Mooncake`, not the stale one earlier research + in this doc used; line numbers below are corrected accordingly, the + underlying claim is unchanged and was directionally correct all + along)**: `dummy_client.cpp:469`, inside `setup_dummy()` + (`dummy_client.cpp:430-509`). Concretely: `UMBPHostMemAllocator` (`host_mem_allocator.h/.cpp`, 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 (`shm_helper.cpp:102-123`), chosen over + `ShmHelper::allocate()` uses (`shm_helper.cpp:74-140`: + `memfd_create_wrapper` at `:26-32`, `memfd_create` at `:108`, + `ftruncate` at `:117`, `mmap(MAP_SHARED|MAP_POPULATE)` at + `:123-124`), chosen over UMBP's own named-`shm_open` DRAM-tier path specifically to avoid `/dev/shm` name collisions across many SGLang instances on one host (Mooncake's rationale for `memfd_create` over named @@ -372,25 +603,40 @@ Adopt Mooncake's answer exactly: hit, sends the recovered `fd` to the server via `SCM_RIGHTS` ancillary data over the **separate raw-UDS fd-handoff socket** (`.fd.sock`, distinct from the gRPC UDS — see §8.3), - mirroring `ipc_send_fd`/`ipc_recv_fd` (`shm_helper.cpp:181-230`) - and Mooncake's `IPC_SHM_REGISTER` handler - (`real_client.cpp:4046-4058`). On a miss (the pointer isn't + mirroring Mooncake's fd-passing and registration handshake + **(v0.7 citation fix — the function names `ipc_send_fd`/ + `ipc_recv_fd` no longer exist in the current Mooncake checkout; + the mechanism was renamed/relocated, not removed)**: fd send/recv + is now `UdsConnection::sendFd`/`UdsConnection::recvFd` + (`uds_transport.cpp:167`/`:196`), and the registration entry point + is `RealClient::handle_ipc_shm_register` + (`real_client.cpp:5311`, dispatched from the `IPC_SHM_REGISTER` + case at `real_client.cpp:5286-5287`). On a miss (the pointer isn't backed by a `kAnonymousShm` allocation — e.g. standalone-process mode was configured but the host allocator wasn't switched to the shm backing), `RegisterMemory` must fail 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` (mirrors `shm_contexts_`, - `real_client.cpp:4074-4082`). + address space (`real_client.cpp:2136-2138`, inside + `map_shm_internal_with_device`) and stores `client_id → {base, + size}` in a registry keyed by `client_id` (mirrors + `shm_contexts_[client_id]`, e.g. `real_client.cpp:2158` — the map + itself is unchanged, only specific usage line numbers moved + across the file's many call sites). 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 - (mirrors `map_dummy_buffer_range_to_real`, - `real_client.cpp:1510-1583,3317`). No KV bytes ever cross the gRPC - channel; the gRPC message is a handful of integers per key. + and reads/writes there directly (mirrors + `RealClient::map_dummy_buffer_range_to_real`, now at + `real_client.cpp:369-379`, called from the batch RPC surface at + `get_into_range_shm_helper`/`get_into_ranges_shm_helper`, + `real_client.cpp:4291`/`:4316` — **(v0.7 citation fix: function + still exists exactly as described, only line numbers moved; the + RPC-handler line numbers were re-confirmed this pass, unlike some + other citations in this section which were only spot-checked)**). + No KV bytes ever cross the gRPC channel; the gRPC message is a + handful of integers per key. 4. `DeregisterMemory` is the inverse — **(v0.3 fix: fd ownership clarified, since the v0.2 wording here conflicted with the §3.2 step 2 registry design).** The allocator's ptr→fd registry (§3.2 @@ -1154,3 +1400,1159 @@ sub-step.)** 13. **Deferred to v2, not v0.1** (§8 item 6b): `tcp_staging` transport with inline RPC payload bytes. Build only if a concrete deployment needs it. + +--- + +## 11. Cross-Node Extension: making `umbp_standalone_server` optionally `DistributedClient`-backed (v0.7 proposal, UNAPPROVED) + +**v0.7 revision note:** a second review pass found 5 concrete gaps in +the v0.6 draft of this section — a config conflict with the already- +shipped v0.5 activation rule, an external-KV routing-correctness hole, +missing transaction/rollback semantics for backend memory +registration, an open correctness question left open instead of +closed, and an overstated "no logic change" claim. All five are +addressed below with inline `(v0.7 fix)` markers. **Guiding principle +adopted for this revision and going forward: where a design question +has a direct Mooncake precedent, resolve it by matching Mooncake, +rather than inventing a new answer** — per explicit direction ("我们与 +Mooncake 保持对标就好,类似的问题都一个思路"). Two of the five fixes +below were resolved this way, with the precedent re-verified against +the local Mooncake checkout (not recalled from memory): + +- **No *implicit* auto-start anywhere in Mooncake's application-side + binding code.** **(v0.7 citation fix, re-checked a second time + against the correct checkout `/apps/nima/KVManager/Mooncake` after + discovering the first verification pass had used a stale checkout at + a different path — the precise claim below is narrower than what was + first written, because the first, broader "zero matches across both + directories" claim turned out to be literally false.)** Re-checked + directly: `grep -rln "fork(|execvp|execve|posix_spawn|subprocess\.|Popen" mooncake-integration/` + returns zero matches — the `DummyClient` application-binding code + never spawns anything. `mooncake-wheel/` (the packaging/CLI layer, + not the application-binding layer) does contain one real match, + `mooncake-wheel/mooncake/cli_client.py:24` + (`subprocess.call([bin_path] + sys.argv[1:])`), but this is the + pass-through implementation of the explicitly-invoked `mooncake_client` + console-script command — functionally equivalent to a human running + the `mooncake_client` binary directly from a shell, not an implicit + auto-start triggered by normal `DummyClient`/integration-API usage. + (The other `mooncake-wheel/` hits — `cli.py`, `cli_bench.py`, + `setup.py`, `tests/test_cli.py` — wrap unrelated binaries, are + build-time only, or are an explicit human-run test harness, + respectively.) The Real Client is always started by something + outside the application (deployment system, operator, systemd unit, + or an explicit CLI invocation) — never implicitly, by any code path + a normal application integration would exercise. §11.4.6 below + adopts this exactly. +- **Mooncake's Real Client already does the "mmap the app's shm, then + register that same mapping with the distributed transport" pattern + this section proposes** **(v0.7 citation fix — line numbers below + corrected against the correct checkout; the function has been + renamed/refactored since the stale citation was written, but the + mechanism is unchanged)** — `real_client.cpp:2136-2137` + (the `mmap(nullptr, shm_size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0)` + call) and `real_client.cpp:2161-2162` + (`client_->RegisterLocalMemory(shm.shm_buffer, shm_size, ...)`, + munmap'd on failure), both inside + `RealClient::map_shm_internal_with_device` + (`real_client.cpp:2097-2186`), which the now-thin + `RealClient::map_shm_internal` (`real_client.cpp:2089-2095`) + delegates to. This is production code, not + a prototype — it directly validates §11.3's core mechanism (register + the *server's* local view of a shared mapping with the distributed + transport) as an already-shipped pattern, not a novel one this + design is the first to attempt. This substantially de-risks the + "cross-mapping visibility" open item in §11.3 (see the fix there). + +**Scope reframing (v0.7, requested explicitly):** §1's original prose +("Standalone-process mode is strictly same-host") describes the +history of what v0.1-v0.5 built, not the target end-state. Given the +project's actual goal is Mooncake parity, and Mooncake's "standalone" +Real Client is cross-node-capable by construction, this section's +distributed-backed shape — not the local-only shape §1 describes — is +the one that actually satisfies that goal. Going forward: +**`standalone-process/distributed-backend` (cross-node, Mooncake-parity +shape) is the primary target this feature is building toward; +`standalone-process/local-backend` (§1-§10, already shipped) is a +same-host fallback/dev-convenience shape, not a co-equal alternative.** +§1's prose is not rewritten in place (it remains an accurate +description of what v0.1-v0.5 actually built and why, and that +reasoning was sound for the narrower problem it solved), but it should +no longer be read as this feature's final scope statement — this +paragraph supersedes it for that purpose. See §11.7 item 1 for the +naming-convention decision that follows from this. + +### 11.0 Why this section exists + +After v0.1-v0.5 shipped, we compared what we built against what +Mooncake's Real Client actually does (source-verified earlier in this +project, see the LMCache/Mooncake research this design already cites) +and found a real divergence, not a rounding error: + +- **Mooncake's Real Client**: a full distributed `Client` (Master + + etcd + Transfer Engine/RDMA), just hosted in its own OS process. + **(v0.7 citation fix — corrected against the correct checkout; + construction pattern unchanged, only line numbers moved, and a + second, structurally identical call site was added since the stale + citation was written)** `real_client.cpp:672-675` (and duplicated at + `real_client.cpp:708-711`, an auto-port-binding-with-retry variant of + the same construction) constructs it via the exact same + `Client::Create(..., master_server_addr, transfer_engine, + {"client_mode":"real"})` an embedded client would use — "standalone" + only changes which process the library lives in, not which protocol + it speaks to the rest of the cluster. A cache miss can still be + served from a remote node over RDMA. +- **What v0.1-v0.5 built**: `umbp_standalone_server` wraps a concrete + `StandaloneClient` (`standalone_server.cpp:587`, + `client_(config)` at `standalone_server.cpp:125`), i.e. `local/`'s + `LocalStorageManager`/`LocalBlockIndex`/`CopyPipeline` moved into its + own process. It has **no** connection to `distributed/` — no master, + no RDMA, no peer protocol. A cache miss is just a miss; there is no + "somewhere else" to check. + +The UDS + `SCM_RIGHTS`-fd-handoff IPC layer between the worker and the +server (§2, §3) is *not* the divergence — that part already correctly +mirrors Mooncake's own app-process↔Real-Client IPC. The divergence is +narrower and specific: **what backs the server side of that IPC**, +local storage vs. a full distributed client. + +### 11.1 Proposal, in one sentence + +Make the concrete `StandaloneClient client_;` member in +`StandaloneServer::Impl` into an `IUMBPClient`-typed backend built +through the **already-existing** `CreateUMBPClient(config)` factory +(`umbp_client_factory.cpp:29-37`), and give `standalone_server_main.cpp` +a way to populate that config's `distributed` field (today it +unconditionally clears it — `standalone_server_main.cpp:43`, +`config.distributed.reset()`). Same server binary, same worker-facing +IPC protocol, same `umbp_standalone_server` deployment story — the +only thing that changes is which `IUMBPClient` implementation answers +Put/Get/etc, decided entirely by how the server operator starts it. + +This is deliberately **not** a rewrite: §11.5 below shows the existing +data-path RPC handlers (Put/Get/BatchPut/BatchGet/Exists/Clear/Flush) +need no logic changes at all, because they already resolve +`(client_id, shm_offset)` to a plain pointer before ever touching +`client_` (`standalone_server.cpp:541-568`) — `client_`'s concrete +type was never load-bearing for those handlers in the first place. + +### 11.2 Two independent `UMBPConfig` instances — this is the key structural point + +There are, and always were, two separate configuration surfaces in +play, and they must 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. **Unchanged by this proposal** — 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, new): + what `umbp_standalone_server` itself uses to decide whether *its* + `client_` backend should be a `StandaloneClient` (today's behavior) + or a `DistributedClient` (this proposal). This is a **second, + independent `UMBPConfig` instance** — never the same object as #1. + +**Verified this is not blocked by the existing mutual-exclusivity +check.** `UMBPConfig::Validate()` (`config.h:345-349`) rejects +`distributed.has_value() && standalone_process.has_value()` *within +one config instance* — it says nothing about two different `UMBPConfig` +objects existing in the same process. `standalone_server_main.cpp` +already builds one `UMBPConfig` for the worker-facing side (with only +`standalone_process` set); this proposal adds a second, separate +`UMBPConfig` for the server's own `client_` backend (with only +`distributed` set, when requested). Neither instance ever has both +fields set, so `Validate()`'s existing rule is satisfied unchanged — +no change needed to that check. + +``` + Worker process umbp_standalone_server process + ┌───────────────────────┐ ┌──────────────────────────────────┐ + │ UMBPConfig #1 │ UDS + │ UMBPConfig #1 (mirror of the │ + │ standalone_process = │───SCM_RIGHTS─│ worker's, address/fd-socket) │ + │ {address} │ │ -> used only to open sockets │ + │ distributed = null │ │ │ + └───────────────────────┘ │ UMBPConfig #2 (server's OWN, │ + │ new in this proposal) │ + │ 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 │ + │ (today's local-only │ + │ behavior, unchanged) │ + └──────────────────────────────────┘ + │ RDMA (if DistributedClient) + ▼ + ┌──────────────────────────────────┐ + │ umbp_master + other nodes' │ + │ peers (the existing distributed/│ + │ cluster, entirely unmodified) │ + └──────────────────────────────────┘ +``` + +### 11.3 Why the zero-copy story still holds with an RDMA backend + +This was the load-bearing technical question and it checks out against +the actual `PoolClient`/`DistributedClient` code (not assumed): + +- RDMA memory registration (`DistributedClient::RegisterMemory` → + `PoolClient::RegisterMemory`, `pool_client.cpp:635-651`) operates + purely on `(ptr, size)` **in the calling process's own virtual + address space** — it calls + `io_engine_->RegisterMemory(ptr, size, -1, MemoryLocationType::CPU)` + (`pool_client.cpp:648`) and caches `{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** + (`pool_client.cpp:635-675`, no ownership check anywhere). +- This means: `umbp_standalone_server`, after `mmap`-ing the worker's + `memfd`-backed shared region via the existing fd-handoff path (§3.2, + already built and working), 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 — full stop, regardless of which process's VA was used to + register them. +- **Get() zero-copy is conditional on this registration, and confirmed + in the actual RDMA read path**: `WaitRemoteBatchGet` + (`pool_client.cpp:1812-1853`) only skips the memcpy-out when the + destination pointer was pre-registered and `FindRegisteredMemory` + matches it (`pool_client.cpp:1661-1664`) — in that branch, the RDMA + read from the remote peer is posted directly into that registered + MR (`pool_client.cpp:1806-1808`) with **no memcpy** afterward. If + the pointer wasn't registered, RDMA lands in a shared staging buffer + and is then `memcpy`'d to the destination (`pool_client.cpp:1847`) — + the ordinary, already-existing non-zero-copy path. +- **Net effect**: register the server's local mapping once (at + fd-handoff time, alongside the existing mmap), and every subsequent + `Get` that targets an offset inside that region gets 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. This is *better* than Mooncake's shape in one + respect: the worker process never touches RDMA/IB verbs directly, + only the server does, entirely consistent with the "worker only + knows about UDS + shm" boundary the v0.1-v0.5 design already + established. + +**(v0.7 fix — closed, was left open in v0.6.)** Cross-mapping memory +visibility/ordering: does an RDMA-completed write into the physical +pages via the server's VA become visible to a subsequent read through +the worker's *independent* VA mapping of the same `memfd`? v0.6 left +this as an unresolved research question. It is now closed by two +independent legs, not a single hand-wave: + +1. **Mooncake precedent, re-verified against the correct checkout this + pass (an earlier pass had used a stale checkout and, separately, + mis-described the allocation direction — both corrected here).** + Mooncake's Real Client does the *identical* dual-mapping shape in + shipping production code, with this exact allocation/registration + direction: the **application** (`DummyClient`) is the one that + originally allocates the shared region, via + `ShmHelper::allocate()` (`dummy_client.cpp:469`, inside + `setup_dummy()`) — that allocation call itself performs the first + `mmap` (`shm_helper.cpp:123-124`), giving `DummyClient` its own + independent VA mapping of the physical pages from the moment of + allocation, not by "receiving and mapping a fd" the way the server + does. `DummyClient` then sends that fd to the Real Client + (`sendFd`, `dummy_client.cpp:408`). The **Real Client** receives it + (`recvFd` via `handle_ipc_shm_register`, `real_client.cpp:5311`) + and performs the *second*, independent `mmap` of the same physical + pages into its own address space (`real_client.cpp:2136-2137`), + then registers that server-local mapping with its own Transfer + Engine (`real_client.cpp:2161-2162`). This is the same shape + `umbp_standalone_server` already implements today (worker allocates + and owns the buffer, §3.2 step 1; server receives the fd and + independently mmaps it, §3.2 step 2) — UMBP's already-shipped + direction matches Mooncake's actual direction exactly; only this + design doc's *prose describing* Mooncake's mechanism had briefly + drifted from that (now corrected). Two independent processes, two + independent VAs, one physical region, the *receiving* side's VA + registered with a distributed transport, exactly this design's + shape — running in production today. This is strong evidence the + pattern is sound, not merely "should be fine" reasoning from first + principles. +2. **Explicit, adopted assumption for the implementation (decision, + not left implicit):** the gRPC `Get` response is defined as the + happens-before edge. `umbp_standalone_server` must only send the + gRPC response for a `Get` **after** the RDMA read completion is + observed server-side (already the natural control flow — there is + no reason to respond before the transfer completes). The worker + must not read the target offset until it has received that + response. This is not a new constraint on the implementation beyond + what it would already naturally do; it is being stated explicitly + here so it is a documented contract, not an accidental property of + the current code shape. + +Residual, non-blocking follow-up: add one integration test that +specifically exercises this path (RDMA `Get` into a `DistributedClient`- +backed server's mapped region, immediately followed by the worker +reading the result through its own independent mapping) once the +distributed-backend test infrastructure from §11.7 item 5 exists — this +is a good regression guard to have, not a prerequisite for the design +being considered closed. + +### 11.4 What changes, concretely, mapped against the current code + +Grounded in the actual current implementation (all citations verified +against source in this pass, not the earlier v0.1-v0.5 design intent): + +1. **`StandaloneServer::Impl`** (`standalone_server.h/.cpp`): + `StandaloneClient client_;` → `std::unique_ptr client_;`, + constructed via `CreateUMBPClient(backend_config)` instead of + `client_(config)` directly. `backend_config` is `UMBPConfig #2` + from §11.2, built by `standalone_server_main.cpp` (see #6 below). +2. **Put/Get/BatchPut/BatchPutWithDepth/BatchGet/Exists/BatchExists/ + BatchExistsConsecutive/Clear/Flush RPC handlers**: mechanically, + swapping `client_` from a value to a `unique_ptr` is a + `client_.` → `client_->` change and nothing more — they already + call `client_.Put(...)` etc. on plain resolved pointers + (`standalone_server.cpp:199-354`). **(v0.7 fix — "no logic change" + overstated in v0.6; three things must be documented even though no + code changes are required for them):** + - **Registration gating is already correct, verified this pass, no + new work needed**: `ResolveRange` (`standalone_server.cpp:541-549`) + already fails (`return false`) if `client_id` isn't present in + the `memory_` map, so `Put`/`Get`/etc. already cannot reach + `client_` at all before `RegisterMemory` has succeeded, for + either backend. The reviewer's concern here is already satisfied + by existing code — noted so a future reader doesn't think it's a + gap. + - **Failure-policy ambiguity (backend-unavailable vs. genuine + cache-miss both surface as `false`) is a pre-existing + characteristic of `IUMBPClient`, not something this proposal + introduces or worsens.** `DistributedClient::Get` already returns + `false` uniformly for "not found" and "RPC/transport failure" + today, for ordinary distributed-mode workers, independent of + whether it runs inside a standalone server. This proposal + inherits that behavior unchanged; fixing it (if ever needed) is + an `IUMBPClient` interface question orthogonal to this section + and explicitly out of scope here. + - **`Flush()`/`Clear()` mean different things per backend, and this + is an accepted semantic difference, not a bug to fix**: under + `StandaloneClient`, `Flush()` drains `CopyPipeline` to SSD + (§4.3); under `DistributedClient`, `Flush()`/`Clear()` operate on + the pool client's own state (heartbeat/registration bookkeeping, + not a local SSD tier — there is no local SSD tier in this backend + per §11.6). Both are the correct, existing behavior of the + respective `IUMBPClient` implementation for their backend; the + server passing the call through unchanged is correct, and this + should be documented as an accepted, backend-dependent semantic + rather than left for a future reader to wonder whether it's an + oversight. +3. **`RegisterMemory`/`DeregisterMemory` RPC handlers** + (`standalone_server.cpp:356-375`): **currently never touch `client_` + at all** — pure fd/mmap bookkeeping against the `memory_` map. Must + add, after a successful `mmap`: + `client_->RegisterMemory(server_local_ptr, size)`, and on + deregistration: `client_->DeregisterMemory(server_local_ptr)`. This + call is **safe to make unconditionally regardless of backend** — + `IUMBPClient::RegisterMemory`'s default (used by `StandaloneClient`, + which never overrides it) is already a no-op returning `true` + (`umbp_client.h`, documented as "Standalone mode needs no + registration (CPU-local memcpy)"). So no `IsDistributed()` branch + is needed here — the interface was already designed for exactly + this kind of backend-agnostic call. + + **(v0.7 fix — v0.6 said "add the call" but didn't specify the + transaction/rollback semantics around it; this is a genuine + correctness gap, not an implementation detail to leave for later.)** + + - **Rollback on backend registration failure.** `RegisterFd` + (`standalone_server.cpp:504-523`, confirmed this pass) is a + single-phase commit: on a successful `mmap`, it immediately writes + the entry into `memory_` and returns success — there is no + existing "pending" intermediate state to roll back from. The new + `client_->RegisterMemory(server_local_ptr, size)` call must be + placed **inside `RegisterFd`, immediately after `mmap` succeeds + and before the entry is written into `memory_`**. If it fails, + `RegisterFd` must `munmap` the just-created mapping and return + failure **without** writing to `memory_` — i.e. backend + registration failure is treated exactly like an `mmap` failure + already is, using the exact same single-phase-commit shape the + function already has. This requires no new state machine, only + ordering the existing calls correctly. + - **Deregistration ordering**: `DeregisterMemory` + (`standalone_server.cpp:370-375`) currently only calls + `UnmapClient()` (munmap + erase). Must call + `client_->DeregisterMemory(server_local_ptr)` **before** + `munmap`-ing (deregister the transport's view of the memory before + the memory itself disappears — the reverse order risks the + backend still holding a registered MR pointing at unmapped + memory, however briefly). + - **Shutdown ordering — `UnmapAll()` must deregister from the + backend before unmapping, and this must happen before + `client_->Close()`.** `UnmapAll()` (`standalone_server.cpp:533-538`) + currently only `munmap`s every entry in `memory_`. Once `client_` + can be a `DistributedClient` holding live RDMA registrations, this + must become: for each registered mapping, call + `client_->DeregisterMemory(base)`, *then* `munmap`. And this whole + sequence must run **before** `client_->Close()` in `Shutdown()` + (`standalone_server.cpp:154-169`, the ordering already fixed in + the v0.1-v0.5 review pass — see the earlier "shutdown ordering" + fix in this document's history) — deregistering after `Close()` + would call into a backend that has already torn down its RDMA/IO + engine state. + - **Explicitly out of scope for this proposal, tracked separately + (not a new problem introduced here):** if a worker sends its fd + via the raw fd-handoff socket and then crashes before ever issuing + the confirming gRPC `RegisterMemory` RPC, the mapping stays live in + `memory_` indefinitely (no timeout/reaper exists for this today). + This is a pre-existing characteristic of the v0.1-v0.5 fd-handoff + protocol, orthogonal to which backend `client_` is — it existed + before this proposal and isn't made worse by it. Worth a future + hardening pass (e.g. a liveness check tied to the worker's gRPC + channel), but not a blocker for this section. +4. **The five external-KV RPC handlers** + (`ReportExternalKvBlocks`/`RevokeExternalKvBlocks`/ + `RevokeAllExternalKvBlocksAtTier`/`MatchExternalKv`/ + `GetExternalKvHitCounts`, `standalone_server.cpp:377-408`): + currently hardcoded stubs (three always return `true`, two return + empty responses) that never touch `client_` at all. + + **(v1.0 fix — superseding both the v0.7 "disable entirely" decision + and the intermediate "naive single shared identity" idea; this is + the third and final revision of this item, corrected after two + rounds of clarification about what external-KV actually is and what + the baseline non-standalone implementation actually already + guarantees.)** + + **What external-KV actually is, clarified:** it 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 — it sends a NEW + request to whichever node the match points at, and that node's own + already-resident local cache (GPU HBM or otherwise) serves it + locally. **UMBP itself never moves the reported bytes for this + feature, in `distributed/` mode or here** — so the earlier v0.7 + concern "the server can't physically reach the worker's HBM memory" + does not apply: nothing ever needs to read that memory through + UMBP for this feature. That concern was based on treating + external-KV as a data-plane feature; it is not one. + + **What genuinely matters instead: does the reported/matched + *identity* (`node_id`) correctly reflect the actual deployment + topology, at the granularity a caller needs?** This is where the + real design work is, and it is grounded in what the *baseline*, + non-standalone implementation already does — per this project's + standing principle of matching existing capability rather than + inventing a narrower one: + + - **Verified against `umbp_store.py:512-519`**: a normal distributed + worker's `node_id` is built as + `f"{node_address}:dp{dp_rank}:pp{pp_rank}:tp{local_rank}"` — pure + per-process rank coordinates, with **zero awareness of which + logical TP/DP replica group the process belongs to**. This means + the baseline already, for free, supports multiple independent + replicas sharing one physical host (e.g. two TP=4 groups on one + 8-GPU box) — each of the 8 processes gets its own distinct + `node_id` regardless of logical grouping, because disambiguation + is by rank coordinate, not by group membership. This is a real, + already-existing capability of the system this proposal must not + silently drop. + - **This project's stated deployment shape is 8-GPU hosts running + SGLang TP=8 (one replica, 8 ranks, per host).** For that specific + shape, a single shared identity per host would actually have been + *semantically adequate* on its own (all 8 ranks in one TP group + process every request in lock-step, so "this host has hash H + cached" is a node-level fact, not a per-rank one) — but adopting a + single-shared-identity design would silently regress the + multi-replica-per-host case the baseline already supports, for no + reason other than convenience. Per explicit instruction, this + proposal should support that case too rather than quietly drop it. + + **Decision: the server maintains one independent distributed + sub-identity per connected worker (`client_id`), not one shared + identity for the whole host.** + + **(v1.2 fix — this used to say "N independent `MasterClient`-level + registrations." That was wrong in a way that would have broken + ordinary distributed routing, not just been imprecise. Corrected to + a purpose-built `ExternalKvIdentityClient` instead, per three + blockers found on review, all confirmed against source, not + speculative.)** + + **Blocker A — a naive `MasterClient`-based sub-identity would pollute + normal Put/Get routing.** Once `RegisterClient` succeeds, that + identity becomes eligible for `ClientRegistry::GetAliveClients()` + (`client_registry.cpp:50`), which ordinary `RoutePut` placement + selects targets from + (`route_put_strategy.cpp`'s `CollectEligibleOnTier`, iterating + *all* alive clients and matching on `tier_capacities`). If a + sub-identity reported any nonzero `tier_capacities`, real KV blocks + from unrelated Puts could get routed onto it — a node that isn't a + real storage participant at all. **Verified the fix works with the + existing code, no routing-code change needed**: `CollectEligibleOnTier` + already skips any candidate with no entry for the tier being routed + (`if (it == client.tier_capacities.end()) continue;`), and + `MasterServer::RegisterClient` does not reject empty/absent + `tier_capacities` (`master_server.cpp:218-240`) — so a sub-identity + registered with **empty `tier_capacities`** is already, structurally, + invisible to `RoutePut`/`RouteGet` selection, with zero changes to + the routing algorithm itself. + + **(v1.2 addition — the same invariant also covers `EvictionManager`, + checked directly rather than assumed to be covered by the RoutePut + fix alone.)** `EvictionManager::RunOnce` (`eviction_manager.cpp:83-91`) + also calls `registry_.GetAliveClients()` and iterates each client's + `tier_capacities` to detect overloaded node-tiers + (`for (const auto& [tier, cap] : client.tier_capacities) { ... }`); + a client with an empty `tier_capacities` map contributes zero + iterations and can never be flagged as overloaded or selected as an + eviction target — the same empty-capacities invariant that hides a + sub-identity from `RoutePut` also makes it structurally invisible to + `EvictionManager`, with no separate mechanism needed. **This must + hold by construction, not by convention that a future edit could + quietly break** (e.g. someone "helpfully" wiring up real capacity + reporting on `ExternalKvIdentityClient` to make monitoring output + look more complete) — which is exactly why Blocker B's fix is a + dedicated class with no `tier_capacities`-reporting code path at + all, rather than a `MasterClient` configured to behave this way by + convention. **Additionally decided: `ExternalKvIdentityClient`'s + heartbeat must never publish `EventBundle`s (owned-KV-block delta/ + full-sync events) — it has no owned index to report in the first + place, since it is not a real storage participant; this is stated + explicitly here, not left to be inferred from "empty capacities," + because the two are different fields in `HeartbeatRequest` + (`tier_capacities` vs. `bundles`, `distributed/proto/umbp.proto:103-109`) + and both must independently stay empty.** + + **Blocker B — reusing `MasterClient` directly is the wrong vehicle, + for two compounding reasons, so this is now a dedicated class:** + `MasterClient` is a full-featured class built for real distributed + peers — capacity heartbeats, KV-event-bundle sequencing, metrics + reporting, and (separately) a real, **confirmed** bug in its + re-register path: on receiving `CLIENT_STATUS_UNKNOWN`, the + re-register request sets only `node_id`/`node_address`/ + `tier_capacities`/`tags` (`master_client.cpp:~612`) and **omits + `peer_address`/`engine_desc`** — unlike the initial `RegisterSelf` + (`master_client.cpp:137-141`), which sets both. Since + `ClientRegistry::RegisterClient` unconditionally overwrites + `record.peer_address`/`record.engine_desc_bytes` with whatever the + request carried (empty, in the re-register case), any identity that + re-registers after Master forgets it (restart, reaper expiry) loses + its `peer_address` — and `MatchExternalKv`'s response, sourced from + `ClientRegistry::GetAlivePeerView()` (`master_server.cpp:590`), + would then return an empty `peer_address` for that node. This is a + pre-existing bug in `MasterClient` itself (it would affect real + distributed peers too, not just this proposal), but relying on it + un-fixed for a *new* use case is not acceptable, and patching the + shared `MasterClient` class carries real regression risk across + every existing distributed peer for the sake of a narrow new + feature. + + **Decision: introduce a new, purpose-built `ExternalKvIdentityClient` + class instead of instantiating N `MasterClient`s.** It implements + only `Register`/`Heartbeat`/`Unregister` plus the five external-KV + RPCs — no capacity heartbeats, no KV-event-bundle logic, no + `PeerDramAllocator`/`PeerSsdManager`/`PeerServiceServer` coupling. + This resolves both blockers structurally rather than by convention: + - **Blocker A is closed by construction, not configuration**: the + class has no `tier_capacities`-reporting code path at all, so + there is no field a future edit could accidentally populate with + real capacity — contrast with "reuse `MasterClient` but always + pass empty capacities," which depends on every future call site + remembering to keep passing empty. + - **Blocker B is closed by not repeating the bug**: since this is + new code, its `Register`/re-register path always includes + `peer_address`/`engine_desc` on every call, from the start — no + fix to shared `MasterClient` code required, no regression surface + on real distributed peers. + - Heartbeat's only job for this class is liveness (keep the + `ClientRegistry` entry from expiring) — it must **not** publish + owned-KV events or capacity snapshots, since it represents no real + storage. + - All N instances (N ≤ 8 for the stated deployment shape) **share + one common `peer_address`** (the server's own single physical + `PeerService` endpoint) — **verified this pass, not assumed**: + `peer_address` uniqueness is not required anywhere in the + registry/dispatch path — `ClientRegistry` stores a plain + `node_id → peer_address` map (`client_registry.cpp:217`), + `MasterServer::GetOrCreateStub` keys its gRPC stub cache by + `node_id` (`master_server.cpp:158-166`, tolerating multiple + `node_id`s pointing at the same `peer_address`), and the + `MatchExternalKv` response path returns `(node_id, peer_address)` + pairs per match (`master_server.cpp:595`) — nothing here assumes + or requires a 1:1 `node_id`↔`peer_address` mapping. + - Each of the five external-KV RPC handlers, on receiving a call + from a given `client_id`, dispatches it to *that worker's* + `ExternalKvIdentityClient` instance (`client_id → + ExternalKvIdentityClient` lookup, mirroring the existing + `client_id → {base, size}` lookup already used for offset + resolution, §3.2) rather than to any single, server-wide identity. + - The bulk Put/Get/RegisterMemory RDMA data path is **unaffected** + and continues to share the server's one real `DistributedClient` + backend (built via `CreateUMBPClient`, per §11.1/§11.4.1) for + efficiency — per §11.3/§11.4.3, RDMA registration is inherently + per-VA and identity-agnostic, so there is no reason to split it + per worker; only the identity-bound external-KV surface needs + per-worker handling, and now has a class whose only job is that. + + **Blocker C — the registration-handshake field list, decided now, + not left to implementation time.** `RegisterClientRequest` requires + both `node_id` **and** `node_address` + (`distributed/proto/umbp.proto:54-60`) — a v1.1 draft of this + section under-specified this as "just `node_id`." **Decided:** the + worker conveys `worker_node_id`, `worker_node_address`, and + optionally `tags` — **not raw rank components** (`dp_rank`/`pp_rank`/ + `tp_rank`). Python (`umbp_store.py`'s existing baseline node_id/ + node_address derivation logic, `umbp_store.py:496-519`) remains the + single source of truth for how these strings are computed; the + server never re-derives rank semantics itself. This avoids two + independent, potentially-diverging implementations of the same + derivation rule ever existing at once. + + **(v1.2 — wire-schema placement decided, not left as prose without + a concrete message.)** These fields are added to the **existing** + `RegisterMemoryRequest` message (`distributed/proto/umbp_standalone.proto:87-91`, + currently `{client_id, worker_base, size}`), not 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, + // mirrors RegisterClientRequest.tags + } + ``` + + Rationale for extending this message rather than adding a separate + `RegisterWorkerRequest`: `RegisterMemoryRequest` is already the gRPC + call that completes a worker's registration lifecycle (the + "confirm" step after the raw-UDS fd handoff, §3.2/§11.4.3) — a + worker's memory registration and its external-KV identity + registration should start and end together (the sub-identity should + not exist before the worker has a registered buffer, and should be + torn down at `DeregisterMemory`, per §11.4.7's shutdown/lifecycle + spec below), so piggybacking on the existing message avoids both a + new round trip and a second, independent lifecycle to keep in sync + with the first. `worker_node_id`/`worker_node_address` are plain + `string` fields (not `optional`) — a worker that never sets them + (e.g. a `StandaloneClient`-backed server has no use for them at all) + simply leaves them empty, and the server only constructs an + `ExternalKvIdentityClient` for a `client_id` whose + `RegisterMemoryRequest` carried a non-empty `worker_node_id`. + + **Real, bounded implementation cost, not hidden:** this means the + server runs up to N independent `ExternalKvIdentityClient` + registrations and heartbeat threads concurrently (N ≤ 8 for the + stated deployment + shape) — the same aggregate cost the baseline already pays today + (8 independent processes each running one), just consolidated into + one process. This is not free, but it is not new cost relative to + the baseline either; it is the honest price of preserving a + capability the baseline already has, not a discretionary addition. + + **(v1.1 — confirmed as the final decision, alternatives considered + and explicitly rejected, not left as an open trade-off.)** Two + points raised in review, both resolved: + - **No relay/double-hop exists in this design.** Each per-worker + `ExternalKvIdentityClient` sub-instance runs entirely inside + `umbp_standalone_server` and heartbeats `MasterServer` directly — + the worker is never in this path after the one-time `node_id` + handoff at registration; there is no "worker → server → Master" + forwarding of heartbeat traffic anywhere in this design. What *is* + real is N independent direct connections/threads originating from + the same host, not a relay chain. + - **Batching N identities onto one heartbeat connection is not + possible without a Master-protocol change, and is out of scope + here.** Verified: `RegisterClientRequest`/`HeartbeatRequest` + (`distributed/proto/umbp.proto:54-115`) both carry a single + `node_id` field, not `repeated` — the protocol has no batched- + identity shape today, and both `MasterClient` and the new + `ExternalKvIdentityClient` are 1-instance-per-`node_id` classes. + Adding batched multi-identity heartbeats would be + a `distributed/` master/protocol change affecting far more than + `umbp_standalone_server`, and is explicitly not attempted here. + - **No Mooncake precedent exists for this specific problem, checked + directly rather than assumed.** `real_client.cpp` constructs + exactly one `Client::Create(...)` for the Real Client's entire + lifetime (confirmed: exactly 2 call sites, both single-instance + construction — one primary, one port-retry variant — never one + per `DummyClient`). Multiple `DummyClient`s sharing one Real + Client are never given independent Master-facing identities in + Mooncake; multi-tenancy there is handled purely as **storage-key + namespacing** via `tenant_id` (`MakeTenantScopedStorageKey`) under + one shared identity — a fundamentally lighter-weight answer than + per-worker sub-identities, but one that only works because + Mooncake has no feature analogous to external-KV requiring + per-original-owner routing precision. This was raised explicitly + as an alternative (call it Option B: one shared identity + `client_id`- + based key/tenant namespacing, matching Mooncake's actual pattern, + at the cost of losing the multi-replica-per-host routing precision + §11.4.4 exists to preserve) and **explicitly rejected in favor of + the N-sub-identity design (Option A, this section)** — decision + confirmed: preserving parity with the `distributed/` baseline's + existing per-worker-identity capability outweighs the simplicity + Mooncake's narrower pattern would have offered here. This is + genuinely new engineering for UMBP with no precedent (in either + UMBP's own `distributed/` code or in Mooncake) to lean on for the + "N identities behind one process" mechanics specifically — flagged + honestly as such, not presented as a known-safe pattern being + reused. + + Independent of the above, add the same `client_mu_`-style locking + discipline these five handlers currently skip + (`standalone_server.cpp:206-213` pattern) around whichever + per-worker sub-identity they dispatch to, since each sub-identity + is now a stateful object with its own background threads. + + Under a `StandaloneClient` backend (no distributed identity at all), + behavior is unchanged from today: hardcoded stubs, matching + `StandaloneClient`'s own no-op external-KV methods + (`standalone_client.h:67-83`). +5. **`Ping` / capability signaling** (`standalone_server.cpp:193-197`, + `umbp_standalone.proto:33-35`): today `PingResponse` has only + `bool ready`, and nothing anywhere lets a worker learn whether the + server it's connected to is distributed-backed. + + **(v0.9 fix — upgraded from "open, low stakes" to a required part of + this proposal's implementation, not deferred.)** The reasoning in + the v0.7 draft — "the data plane doesn't care, only humans debugging + need it" — missed the actual failure mode this exists to catch: a + single `UMBP_STANDALONE_ADDRESS` can now resolve to *either* a + local-backed or a distributed-backed server (§11.6), and if the + whole point of this section is Mooncake parity (cross-node), the + most dangerous failure mode is exactly the quiet one — a worker (or + the deployment/test tooling that stood up the server) connects + successfully, every RPC succeeds, and the server has silently ended + up local-backed instead of distributed-backed (e.g. an operator's + deployment script forgot to set `UMBP_MASTER_ADDRESS` on the + *server's* side — a legitimate, non-erroring config per §11.2, not + the already-hard-failing misconfiguration case in item 6 below). + Nothing observes this today; the feature's entire purpose (cross- + node capability) can silently not exist while everything appears to + work. This is the same class of failure this document has rejected + everywhere else (§5, §8, item 6 below) — it should not be waved + through here just because the data path itself doesn't strictly + need it. + + **Decision: `PingResponse.deployment_mode` (new proto field, + enum `LOCAL`/`DISTRIBUTED`, additive/backward-compatible, + populated from `client_->GetDeploymentMode()`) is a required part + of this proposal's implementation, not an optional nice-to-have.** + Additionally: **any launch/health-check/integration-test tooling for + the distributed-backed shape must assert + `deployment_mode == DISTRIBUTED` after startup** — a distributed- + backed server that comes up as `LOCAL` (for whatever reason) must be + treated as a failed deployment, not a degraded-but-working one. The + local-backed shape itself remains a fully legitimate, intentional + deployment choice (§11.6) when that's what was actually configured + — this assertion only guards against the *parity* shape silently + collapsing into the *fallback* shape unnoticed. +6. **`standalone_server_main.cpp`**: currently unconditionally does + `config.distributed.reset()` (`standalone_server_main.cpp:43`) and + only ever builds `config.standalone_process`. Needs a **second, + independent** config-building path for `backend_config` (§11.2). + + **(v0.9 fix — the v0.7 env-var list was incomplete, not just + abbreviated; re-checked against the actual current + `umbp_store.py` code, not recalled from memory.)** The v0.7 draft + said "reuse the exact set `umbp_store.py` already reads" and then + listed only 6 variables. Re-reading `umbp_store.py:458-620` in full, + a real distributed worker's config surface is materially larger: + `master_address`, `node_address`, `node_id`, `auto_heartbeat` + (`umbp_store.py:524-527` — **note: `extra_config`-only, no env-var + fallback exists for this one today**), `io_engine.host`, + `io_engine.port`, `staging_buffer_size`, `ssd_staging_buffer_size`, + `ssd_staging_buffer_slots`, `peer_service_port`, + `cache_remote_fetches`, and `dram_page_size`. An implementer who + copied only the original 6-variable list would build a server-side + `DistributedClient` that silently behaves differently from an + ordinary distributed worker — e.g. no SSD staging buffer sizing, no + `cache_remote_fetches` admission control tuning — a real behavior + drift, not a cosmetic gap. + + **One field needs special handling, not a direct copy:** + `dram_page_size` in `umbp_store.py` (`umbp_store.py:598-620`) is + normally *auto-derived by probing `mem_pool_host`* (the worker's own + host KV tensor buffer) when not explicitly set — this derivation is + meaningless for `standalone_server_main.cpp`, which has no + `mem_pool_host` of its own (the server's `client_` backend isn't + attached to any particular worker's tensor layout — it may serve + many workers with potentially different layouts, per §4.4). The + server-side config builder must **only** support the explicit- + override path (an env var, e.g. `UMBP_DRAM_PAGE_SIZE`, mirroring + `extra_config["dram_page_size"]`'s role as "operator-controlled + escape hatch" per the comment at `umbp_store.py:590-598`) and leave + it at the `UMBPDistributedConfig` default (`0`, delegating to the + master's own `ClientRegistryConfig.default_dram_page_size`) when + unset — it must not attempt to replicate the `mem_pool_host`-probing + auto-derivation, which does not apply here. + + **Proposed shape for the implementation**: a single + `BuildDistributedBackendConfigFromEnv()` C++ helper (new, + `standalone_server_main.cpp` or a shared helper if reused elsewhere) + that mirrors `umbp_store.py`'s env-var-driven construction of + `UMBPDistributedConfig`, split explicitly into: + - **Required** (missing any of these means "distributed backend not + requested" — server falls back to local-backed, or hard-fails per + the decision later in this same item if some but not all are set, + which is itself a misconfiguration signal): `UMBP_MASTER_ADDRESS`, + `UMBP_NODE_ADDRESS`, `UMBP_NODE_ID`, `UMBP_IO_ENGINE_HOST`. + - **Optional, with the same defaults `UMBPDistributedConfig` already + has when unset**: `UMBP_IO_ENGINE_PORT`, `UMBP_PEER_SERVICE_PORT` + (both reuse the exact worker-side names, since those two already + have an established env-var convention — `umbp_store.py:526-543` — + to mirror), plus five genuinely new names, **finalized now, not + left as "proposed"**: `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` (per the special handling above). + **(v1.2 — decided, no longer "proposed.")** These five get the + `UMBP_DISTRIBUTED_` prefix specifically because, unlike + `UMBP_MASTER_ADDRESS`/`UMBP_NODE_ADDRESS`/`UMBP_NODE_ID`/ + `UMBP_IO_ENGINE_HOST`/`UMBP_IO_ENGINE_PORT`/`UMBP_PEER_SERVICE_PORT` + (which reuse worker-side names verbatim, referring to the exact + same concept in both places), these five have **no existing + env-var form at all today** — worker-side, they are + `extra_config`-only (`umbp_store.py:555-598`) — so there is no + established bare name to collide with, but there is a real future + risk of a bare `UMBP_STAGING_BUFFER_SIZE` etc. later being added + as a worker-side env var with different semantics (e.g. a + per-worker override rather than a server-backend-wide default); + the `_DISTRIBUTED_` infix makes clear these five specifically + configure the *server's own* `UMBPDistributedConfig`, not + anything worker-facing, closing off that ambiguity before it can + arise. `auto_heartbeat` has no worker-side env var today either + (only `extra_config`); **decided: the server-side helper keeps the + `UMBPMasterClientConfig` default (`true`) unconditionally** rather + than adding a dedicated env var for a field with no operational + reason to ever be `false` for these sub-identities-plus-shared- + backend registrations — simpler than inventing + `UMBP_DISTRIBUTED_AUTO_HEARTBEAT` for a knob nobody has asked to + turn off. + + **(v0.7 fix — decided, was flagged as a v0.6 blocker.) A + distributed-backed `umbp_standalone_server` supports external + launch ONLY — it does not support the §4.2 auto-start (worker + fork+exec) path. Local-backed auto-start (§4.2, already shipped in + v0.1-v0.5) is completely unaffected by this and continues to work + exactly as today.** + + The conflict this closes: v0.5 (§5) made + `UMBPStore.__init__` raise if a worker's own environment has both + `UMBP_STANDALONE_ADDRESS` and `UMBP_MASTER_ADDRESS` set (mutual + exclusivity, by design — a worker is either distributed or + standalone-process, never both). If a worker's own auto-start path + were to fork+exec a distributed-backed server, the env vars needed + to configure that child (`UMBP_MASTER_ADDRESS` etc.) would most + naturally live in the *parent's* (the worker's) environment too + (fork inherits environment by default) — which is exactly the + combination v0.5 already made illegal, on purpose, for the worker's + own config parsing. Two ways to resolve this were considered + (support auto-start via a separate `UMBP_STANDALONE_BACKEND_*` env + namespace the worker never reads, vs. drop auto-start for this mode + entirely) — **resolved by matching Mooncake rather than inventing + either option**: re-checked directly against the local Mooncake + checkout this pass (`grep -rln "fork\(|execvp|execve|posix_spawn|subprocess\.|Popen" mooncake-integration/ mooncake-wheel/` + → zero matches), Mooncake's application side (`DummyClient`) never + forks/execs the Real Client under any circumstance — the Real + Client is always started by something outside the application + (deployment system, operator, systemd unit, etc.), full stop, no + exceptions for any deployment shape. Adopting the same rule here is + simpler than a second env namespace, has zero conflict with the + already-shipped v0.5 rule, and matches the explicit goal of this + whole section (Mooncake parity) more directly than inventing a + UMBP-specific auto-start variant Mooncake itself doesn't have. + + In practice: `standalone_server_main.cpp` reads + `UMBP_MASTER_ADDRESS` etc. **from its own process environment**, + set by whatever external launcher starts it (a shell script, + Kubernetes sidecar spec, systemd unit — the same category of + launcher `run_umbp_single_node_hicache.sh` already represents for + `umbp_master` today). Workers never see or set these variables; + they only ever need `UMBP_STANDALONE_ADDRESS`, unchanged from + today. + + **Decision (also confirmed, unchanged from v0.6): if + `UMBP_MASTER_ADDRESS` is set but `PoolClient::Init` can't actually + stand up the RDMA IO engine** (e.g. `UMBP_IO_ENGINE_HOST` empty — + `PoolClient::Init` silently skips building the IO engine entirely + in that case, `pool_client.cpp:385-415`, rather than failing), + **server startup must hard-fail**, not fall back to a + `StandaloneClient` backend with a warning. Silently downgrading a + server an operator explicitly asked to be distributed-backed into a + local-only one is exactly the kind of silent-degradation failure + mode this whole design doc has repeatedly flagged as unacceptable + elsewhere (§5, §8). A misconfigured distributed-backed server + should refuse to start, not quietly become a different, smaller + feature. + +### 11.4.7 `ExternalKvIdentityClient`: interface, lifecycle, and layering (v1.2, new) + +**(This subsection exists because §11.7 previously filed +`ExternalKvIdentityClient` as an open item needing "its own design +pass" — that framing was too weak once §11.4.4 made it a required +component of the final design, not an optional extra. This is that +design pass, not a placeholder for a future one.)** + +**Layering — this is a UMBP compatibility extension, not part of +Mooncake-parity core.** Worth stating explicitly so it is never +mistaken for required minimum-viable scope: the Mooncake-parity core +of this whole section (§11) is `distributed-backed umbp_standalone_server ++ UDS/shm fd handoff` (§11.1-§11.3) — that alone is what makes this +design match Mooncake's Real Client shape. `ExternalKvIdentityClient` +exists purely to avoid *regressing* a capability the `distributed/` +baseline already has (per-worker external-KV routing precision, §11.4.4's +"why this matters" discussion) — it has no Mooncake analog and is not +required for Mooncake parity itself. Concretely, this means: +- It can be scheduled, implemented, and tested as its own, separable + unit of work after the parity core is working — a distributed-backed + server with `ExternalKvIdentityClient` not yet implemented is a + legitimate intermediate state (external-KV simply stays unavailable + for that server, exactly like the `StandaloneClient`-backed case + already handles it today), not a broken or non-conforming one. +- It must never be treated as sitting on the critical path for + declaring "Mooncake parity achieved" — conflating the two would + under-scope the parity-core milestone by tying it to a strictly + UMBP-specific feature Mooncake itself has no equivalent of. + +**Interface** (methods on the new class; not a proto service of its +own — it is a client-side C++ class inside `umbp_standalone_server`, +using the *existing* `UMBPMaster` gRPC service just like `MasterClient` +does): + +- `Register(worker_node_id, worker_node_address, peer_address, engine_desc, tags)` + — sends `RegisterClientRequest` with **empty `tier_capacities`** + (§11.4.4's Blocker A invariant) and the given `peer_address` + (shared across all sub-identities on this server, §11.4.4) / + `engine_desc` (may be empty/dummy — nothing ever RDMA-targets a + sub-identity's "engine," since external-KV never moves bytes through + UMBP, §11.4.4's opening clarification). Constructed the moment a + worker's `RegisterMemoryRequest` carries a non-empty + `worker_node_id` (§11.4.4 Blocker C's wire-schema decision). +- `Heartbeat()` — periodic, liveness-only. Sends `HeartbeatRequest` + with **empty `tier_capacities` and empty `bundles`** on every call + (§11.4.4's v1.2 addition) — never derives or forwards real capacity + or owned-KV-event data. 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 (worker-initiated + teardown), and (b) server shutdown (see below). + + **(v1.4 fix — this bullet previously claimed a third trigger, + "worker disconnect/crash detection (mirrors however + `StandaloneProcessClient` disconnect is already detected... §8 item + 1's connection state machine)," as if it already existed. That was + false, and worse than simply undocumented: §8 item 1 itself states + plainly that this exact mechanism is unsolved — "Crash / restart + semantics are unsolved by every precedent reviewed... this needs + explicit handling in `StandaloneProcessClient` and probably a small + state machine (`CONNECTED` → `DISCONNECTED` → optionally + `RECONNECTING`), which does not exist in any of the three reference + systems today." §11.4.7 was citing, as its basis for "already + detected," a mechanism §8 item 1 explicitly says does not exist — + an internal contradiction a reader of this subsection alone would + never catch.)** + + **Confirmed gap, stated plainly instead of implied as solved: there + is currently no worker crash/disconnect detection anywhere in + `standalone_server.cpp`.** The only two paths that ever tear down an + `ExternalKvIdentityClient` are `DeregisterMemory` (explicit, + worker-initiated) and full server `Shutdown()` + (`UnregisterAllExternalIdentities`, see below). **Consequence, stated + explicitly:** if a worker process crashes or is killed without ever + calling `DeregisterMemory` (SIGKILL, OOM-kill, or any other + non-graceful exit), its `ExternalKvIdentityClient` keeps running + inside `umbp_standalone_server` indefinitely — its heartbeat thread + (which lives in the *server* process, not the worker's, so the + worker's death has no effect on it at all) keeps sending + liveness-only heartbeats forever, so the Master's `ClientRegistry` + entry for that `node_id` never expires and stays `ALIVE`. + `MatchExternalKv`/`GetExternalKvHitCounts` will keep returning this + node_id/peer_address to callers for KV blocks that may no longer + exist (the crashed worker's shm segment is gone), and this is **not + cleaned up until the entire `umbp_standalone_server` process is + restarted** — there is no timeout, reaper, or liveness poll for this + case anywhere in the reviewed code. This is a known, deliberately + deferred limitation (fixing it properly needs a real + liveness/crash-detection mechanism — most plausibly something built + on top of the connection-state-machine work §8 item 1 already + flags as needed for `StandaloneProcessClient`'s own reconnect + handling, not a small patch specific to `ExternalKvIdentityClient` + alone), tracked here and in §8 item 1, not silently absent. +- `ReportExternalKvBlocks`/`RevokeExternalKvBlocks`/ + `RevokeAllExternalKvBlocksAtTier`/`MatchExternalKv`/ + `GetExternalKvHitCounts` — thin forwarding wrappers to the + corresponding `UMBPMaster` RPCs, using this sub-identity's own + `node_id`. These are what the five `standalone_server.cpp` external-KV + RPC handlers (§11.4.4, item 4) dispatch into once they've resolved + the calling `client_id` to its `ExternalKvIdentityClient` instance. + +**Re-registration must not repeat `MasterClient`'s bug (§11.4.4 +Blocker B) — stated as a hard requirement, not left implicit in "this +is new code so it won't have the bug":** every `RegisterClientRequest` +this class ever sends — the very first one and any sent after a +`CLIENT_STATUS_UNKNOWN` heartbeat response — **must** include +`peer_address` and `engine_desc`. Implementation-wise, the simplest way +to guarantee this is to store `peer_address_`/`engine_desc_bytes_` as +member fields set once at construction and have exactly one internal +`BuildRegisterRequest()` helper used for both the initial `Register()` +call and any re-register path, rather than two independently-written +request-construction code paths (which is how `MasterClient`'s bug was +introduced in the first place — the re-register path was written +separately from `RegisterSelf` and simply forgot two fields). + +**Lifecycle serialization trade-off (accepted v1.4 implementation +simplification):** the current implementation may use one server-wide +lifecycle mutex to serialize `ExternalKvIdentityClient` register, +deregister, and shutdown bookkeeping. This deliberately closes +register-vs-deregister races for the same `client_id`, but it also +means the lock can cover blocking Master RPCs (`RegisterClient` / +`UnregisterClient`, bounded by the RPC shutdown deadline). Therefore, +when the Master is slow or unavailable, otherwise unrelated workers can +queue behind each other during `RegisterMemory`/`DeregisterMemory`; in +the stated TP=8 deployment shape, the worst-case startup/shutdown delay +can be amplified toward `N * rpc_deadline` rather than one parallel +deadline. This is an accepted simplification because these are +lifecycle operations, not the Put/Get hot path, and because the current +target scale is small. If future deployments support many more workers, +frequent worker reconnects, or observe startup stalls caused by slow +Master RPCs, this lock should be the first place to revisit: either +replace it with per-`client_id` lifecycle locks or decouple +external-KV identity registration from core memory registration with an +async/retry path. + +**Shutdown ordering** (extends §4.3/§11.4.3's existing shutdown +sequence, does not replace it): `umbp_standalone_server`'s `Shutdown()` +must `Unregister()` all live `ExternalKvIdentityClient` instances +**before** tearing down the shared `DistributedClient` backend +(`client_->Close()`) or exiting — unregistering first ensures the +Master's `ClientRegistry` doesn't briefly carry stale `ALIVE` entries +for identities whose process is already gone, which would otherwise +sit around until heartbeat-timeout-driven reaping (§4, the same +"absent liveness" concern already motivated for the primary +`DistributedClient` identity). Order within `Shutdown()`, appended to +the sequence already specified in §4.3/§11.4.3: drain in-flight RPCs → +**unregister all `ExternalKvIdentityClient` instances** → flush +`CopyPipeline`/close the shared `DistributedClient` backend → unmap → +exit. + +### 11.5 What does *not* change + +- The worker-facing IPC layer: UDS control-plane socket, the separate + raw-UDS `.fd.sock` `SCM_RIGHTS` channel, offset translation, + `StandaloneProcessClient`'s locking discipline, the fd-ownership + registry in `HostMemAllocator` — **none of this is touched**. All of + it was built, reviewed, and bug-fixed independent of what backs the + server side, and stays exactly as-is. +- `local/` and today's `distributed/`-for-workers path — **both + unmodified**, per §7's existing compatibility statement, which this + proposal does not change. +- The `UMBPStandaloneProcessConfig` shape workers use (§6) — + unmodified. Workers still only ever set `address` (env-var-only + activation, per the v0.5 decision) and never need to know or care + what the server's own backend is. + +### 11.6 Consequence for "what does the local DRAM tier even mean now" + +Worth stating explicitly since it's a natural point of confusion: when +`client_` is a `DistributedClient`, the server's `LocalStorageManager`/ +`DRAMTier`/`SSDTier` concept (§3.2's earlier "server's own DRAM tier +should default to private anon/hugetlb" discussion) **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 +(`distributed_client.cpp:41-53`), used as its local slice of the +distributed pool, not as a passthrough cache for the registered worker +buffers. The registered worker buffers (via `RegisterMemory` in +§11.4.3) are a *separate* thing: they exist purely so RDMA can target +them directly, not so `DistributedClient` manages them as cache +storage. **There are now three, not two, deployment shapes for +`umbp_standalone_server`**, and the design should name them distinctly +to avoid ambiguity in docs/logs/config: + +| Shape | `client_` backend | Cross-node | Server's own storage | +|---|---|---|---| +| existing v0.1-v0.5 (unchanged) | `StandaloneClient` | no | `LocalStorageManager` (DRAM+SSD tiers, private memory) | +| this proposal, backend enabled | `DistributedClient` | **yes**, via Master+RDMA | `DistributedClient`'s own DRAM pool (distinct concept, not `LocalStorageManager`) | +| this proposal, backend disabled | `StandaloneClient` | no | same as row 1 — this is the fallback/default, identical to today | + +### 11.7 Open questions for discussion (before any code) + +**(v0.7: items 2 and 4 from the v0.6 draft are now resolved — see +§11.4.6 and §11.3 respectively — and are kept here only as a record of +what was decided and why. Items 1, 3, 5 remain genuinely open.)** + +1. **Naming/scope framing — still open, needs a decision.** The + "Scope reframing" paragraph near the top of §11 (added this pass) + already states that + distributed-backed is the shape that actually matches the stated + goal, and local-backed is a fallback/dev-convenience shape, not a + co-equal alternative. Given that, should the *document* (and + eventually config/logs) use a name like + `standalone-process/distributed-backend` (or "real-client mode") vs + `standalone-process/local-backend` per §11.6's table, rather than + presenting them as two unlabeled options? Recommend adopting this + secondary naming now, in the doc, ahead of implementation — it's a + documentation-only decision but should be settled before code + introduces its own inconsistent naming (env var prefixes, log + messages, etc). +2. ~~The hard-fail-on-misconfiguration decision in §11.4.6~~ — + **resolved**: hard-fail, confirmed in §11.4.6. +3. ~~Whether `PingResponse.deployment_mode` is worth the proto change + now~~ — **resolved (v0.9): required, not deferred.** See §11.4.5 — + this is now a mandatory guard against the parity shape silently + collapsing into the local-only fallback shape, not an optional + debugging aid. +4. ~~The cross-mapping memory-visibility question in §11.3~~ — + **resolved**: closed via the Mooncake production-precedent citation + plus the explicit gRPC-response-as-happens-before-edge contract, see + §11.3. +5. Testing implications, not yet scoped: the existing + `test_standalone_shm_ipc.cpp` suite only exercises the + `StandaloneClient`-backed path. A `DistributedClient`-backed test + needs at minimum a running `umbp_master` plus the RDMA/IO-engine + stack available in the test environment — likely a heavier + integration test than what exists today, possibly gated separately + (e.g. only run where RDMA hardware/loopback is available), which + needs its own test-infrastructure decision before it can be added + to `§10`'s implementation-order list. Should also cover: per-worker + external-KV sub-identity registration/dispatch under a + distributed-backed server with 2+ connected workers (§11.4.4 — this + is the item's own core mechanism now, not an edge case), and the + `RegisterMemory`/`DeregisterMemory` rollback/shutdown-ordering paths + added in §11.4.3. +6. ~~The exact registration-handshake field(s) for per-worker identity + conveyance~~ — **resolved (v1.2, Blocker C): `worker_node_id` + + `worker_node_address` + optional `tags`, not raw rank components.** + See §11.4.4. +7. ~~`ExternalKvIdentityClient` needs a design/implementation pass of + its own~~ — **resolved (v1.2): see §11.4.7** for the full interface + (`Register`/`Heartbeat`/`Unregister`/the five external-KV RPCs), + the re-registration hard requirement, and shutdown ordering. It is + still real, net-new implementation work — it doesn't exist anywhere + in the codebase today, unlike `MasterClient`, which at least could + have been (incorrectly) reused — so it still needs to be factored + into scheduling, not just the N-heartbeat-threads cost already noted + above; what changed is that its shape is now fully specified rather + than deferred. + +### 11.8 Source references for this section + +- `/apps/nima/KVManager/mori/src/umbp/standalone/standalone_server.{h,cpp}` +- `/apps/nima/KVManager/mori/src/umbp/standalone/bin/standalone_server_main.cpp` +- `/apps/nima/KVManager/mori/src/umbp/distributed/distributed_client.{h,cpp}` +- `/apps/nima/KVManager/mori/src/umbp/distributed/pool_client.{h,cpp}` +- `/apps/nima/KVManager/mori/src/umbp/include/umbp/distributed/master/master_client.h` +- `/apps/nima/KVManager/mori/src/umbp/include/umbp/common/config.h` + (`UMBPConfig::Validate()`, `UMBPDistributedConfig`, + `UMBPMasterClientConfig`, `UMBPIoEngineConfig`) +- `/apps/nima/KVManager/mori/src/umbp/umbp_client_factory.cpp` +- `/apps/nima/KVManager/mori/src/umbp/distributed/proto/umbp_standalone.proto` +- `/apps/nima/KVManager/mori/src/umbp/include/umbp/umbp_client.h` + (`IUMBPClient::RegisterMemory`'s default no-op-returns-true contract) diff --git a/src/umbp/include/umbp/common/config.h b/src/umbp/include/umbp/common/config.h index b400be96c..21727f887 100644 --- a/src/umbp/include/umbp/common/config.h +++ b/src/umbp/include/umbp/common/config.h @@ -250,6 +250,13 @@ 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 { 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/standalone/bin/standalone_server_main.cpp b/src/umbp/standalone/bin/standalone_server_main.cpp index eed29bef9..6cf402f77 100644 --- a/src/umbp/standalone/bin/standalone_server_main.cpp +++ b/src/umbp/standalone/bin/standalone_server_main.cpp @@ -25,22 +25,179 @@ #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.distributed.reset(); + 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') { @@ -51,11 +208,13 @@ int main(int argc, char** argv) { address = mori::umbp::standalone::DefaultStandaloneAddress(); } - mori::umbp::UMBPStandaloneProcessConfig standalone_cfg; - standalone_cfg.address = address; - config.standalone_process = standalone_cfg; - - mori::umbp::standalone::StandaloneServer server(config, address); + 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); @@ -68,13 +227,13 @@ int main(int argc, char** argv) { return 1; } - if (!server.Start()) { + 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, &signal_set, &stop_signal_waiter]() { + std::thread signal_waiter([server = server.get(), &signal_set, &stop_signal_waiter]() { while (!stop_signal_waiter.load()) { timespec timeout{}; timeout.tv_sec = 1; @@ -82,7 +241,7 @@ int main(int argc, char** argv) { const int signum = sigtimedwait(&signal_set, nullptr, &timeout); if (signum == SIGINT || signum == SIGTERM) { MORI_UMBP_INFO("[StandaloneServer] caught signal {}, shutting down", signum); - server.Shutdown(); + server->Shutdown(); return; } if (signum == -1 && errno != EAGAIN && errno != EINTR) { @@ -92,8 +251,9 @@ int main(int argc, char** argv) { } }); - MORI_UMBP_INFO("[StandaloneServer] running on {}", address); - server.Run(); + MORI_UMBP_INFO("[StandaloneServer] running on {} backend={}", address, + distributed_requested ? "distributed" : "local"); + server->Run(); stop_signal_waiter = true; signal_waiter.join(); 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/standalone_process_client.cpp b/src/umbp/standalone/standalone_process_client.cpp index a7c2a5ef0..52c64adba 100644 --- a/src/umbp/standalone/standalone_process_client.cpp +++ b/src/umbp/standalone/standalone_process_client.cpp @@ -472,6 +472,9 @@ bool StandaloneProcessClient::RegisterMemory(uintptr_t ptr, size_t size) { 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()) { @@ -528,6 +531,7 @@ bool StandaloneProcessClient::ReportExternalKvBlocks(const std::vectorReportExternalKvBlocks(&ctx, req, &resp); return status.ok() && resp.ok(); @@ -539,6 +543,7 @@ bool StandaloneProcessClient::RevokeExternalKvBlocks(const std::vectorRevokeExternalKvBlocks(&ctx, req, &resp); return status.ok() && resp.ok(); @@ -548,6 +553,7 @@ 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(); @@ -559,6 +565,7 @@ std::vector StandaloneProcessClient::MatchExternal ::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 {}; @@ -583,6 +590,7 @@ std::vector StandaloneProcessClient::GetEx 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 {}; diff --git a/src/umbp/standalone/standalone_server.cpp b/src/umbp/standalone/standalone_server.cpp index 5af6e42e0..6eb90b3a8 100644 --- a/src/umbp/standalone/standalone_server.cpp +++ b/src/umbp/standalone/standalone_server.cpp @@ -37,14 +37,17 @@ #include #include #include +#include #include +#include #include #include #include #include "mori/utils/mori_log.hpp" -#include "umbp/local/standalone_client.h" +#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 { @@ -107,6 +110,26 @@ bool SetFdSocketTimeouts(int fd, std::chrono::milliseconds timeout) { 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); @@ -122,7 +145,8 @@ std::chrono::milliseconds FdHandshakeTimeout() { class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { public: Impl(const UMBPConfig& config, std::string address) - : client_(config), + : backend_config_(NormalizeBackendConfig(config)), + client_(CreateUMBPClient(backend_config_)), address_(std::move(address)), fd_socket_path_(DeriveFdSocketPath(address_)) {} @@ -180,12 +204,16 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { if (server_) { server_->Shutdown(std::chrono::system_clock::now() + ShutdownDeadline()); } + UnregisterAllExternalIdentities(); { std::lock_guard lock(client_mu_); - client_.Flush(); - client_.Close(); + client_->Flush(); } UnmapAll(); + { + std::lock_guard lock(client_mu_); + client_->Close(); + } unlink(UnixPathFromGrpcAddress(address_).c_str()); unlink(fd_socket_path_.c_str()); } @@ -193,6 +221,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { 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; } @@ -208,7 +237,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; } - SetBool(response, client_.Put(request->key(), ptr, static_cast(request->size()))); + SetBool(response, client_->Put(request->key(), ptr, static_cast(request->size()))); return grpc::Status::OK; } @@ -224,7 +253,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; } - SetBool(response, client_.Get(request->key(), ptr, static_cast(request->size()))); + SetBool(response, client_->Get(request->key(), ptr, static_cast(request->size()))); return grpc::Status::OK; } @@ -242,7 +271,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { FillFalse(request->keys_size(), response); return grpc::Status::OK; } - FillResults(client_.BatchPut(keys, ptrs, sizes), response); + FillResults(client_->BatchPut(keys, ptrs, sizes), response); return grpc::Status::OK; } @@ -274,7 +303,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { FillFalse(request->keys_size(), response); return grpc::Status::OK; } - FillResults(client_.BatchPutWithDepth(keys, ptrs, sizes, depths), response); + FillResults(client_->BatchPutWithDepth(keys, ptrs, sizes, depths), response); return grpc::Status::OK; } @@ -292,7 +321,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { FillFalse(request->keys_size(), response); return grpc::Status::OK; } - FillResults(client_.BatchGet(keys, ptrs, sizes), response); + FillResults(client_->BatchGet(keys, ptrs, sizes), response); return grpc::Status::OK; } @@ -303,7 +332,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; } - SetBool(response, client_.Exists(request->key())); + SetBool(response, client_->Exists(request->key())); return grpc::Status::OK; } @@ -315,7 +344,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { FillFalse(request->keys_size(), response); return grpc::Status::OK; } - FillResults(client_.BatchExists(keys), response); + FillResults(client_->BatchExists(keys), response); return grpc::Status::OK; } @@ -327,7 +356,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { response->set_count(0); return grpc::Status::OK; } - response->set_count(client_.BatchExistsConsecutive(keys)); + response->set_count(client_->BatchExistsConsecutive(keys)); return grpc::Status::OK; } @@ -338,7 +367,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; } - SetBool(response, client_.Clear()); + SetBool(response, client_->Clear()); return grpc::Status::OK; } @@ -349,7 +378,7 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; } - SetBool(response, client_.Flush()); + SetBool(response, client_->Flush()); return grpc::Status::OK; } @@ -359,51 +388,117 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { SetBool(response, false, "server is shutting down"); return grpc::Status::OK; } - std::lock_guard lock(memory_mu_); - auto it = memory_.find(request->client_id()); - bool ok = it != memory_.end() && it->second.worker_base == request->worker_base() && - it->second.size >= request->size(); - SetBool(response, ok, ok ? "" : "fd handoff registration was not found"); + 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*, + const ::umbp::StandaloneExternalKvMutationRequest* request, ::umbp::BoolResponse* response) override { - SetBool(response, true); + 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*, + const ::umbp::StandaloneExternalKvMutationRequest* request, ::umbp::BoolResponse* response) override { - SetBool(response, true); + 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*, - ::umbp::BoolResponse* response) override { - SetBool(response, true); + 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*, - ::umbp::StandaloneMatchExternalKvResponse*) override { + 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*, - ::umbp::StandaloneExternalKvHitCountsResponse*) override { + 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; } @@ -414,6 +509,83 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { 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) { @@ -509,33 +681,63 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { 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::lock_guard lock(memory_mu_); - auto old = memory_.find(client_id); - if (old != memory_.end() && old->second.base) { - munmap(old->second.base, static_cast(old->second.size)); + 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}; } - 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::lock_guard lock(memory_mu_); - auto it = memory_.find(client_id); - if (it == memory_.end()) return; - if (it->second.base) munmap(it->second.base, static_cast(it->second.size)); - memory_.erase(it); + 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::lock_guard lock(memory_mu_); - for (auto& kv : memory_) { - if (kv.second.base) munmap(kv.second.base, static_cast(kv.second.size)); + 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)); } - memory_.clear(); + munmap(mem.base, static_cast(mem.size)); } bool ResolveRange(const std::string& client_id, uint64_t offset, uint64_t size, @@ -584,7 +786,32 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { for (int i = 0; i < n; ++i) response->add_ok(false); } - StandaloneClient client_; + 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_; @@ -593,6 +820,9 @@ class StandaloneServer::Impl final : public ::umbp::UMBPStandalone::Service { 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; From 4b4204fff919e29b285c34f7f413fea6e87f9669 Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Fri, 10 Jul 2026 02:26:56 +0000 Subject: [PATCH 3/3] rewrite doc --- .../doc/design-standalone-process-mode.md | 3477 +++++------------ 1 file changed, 983 insertions(+), 2494 deletions(-) diff --git a/src/umbp/doc/design-standalone-process-mode.md b/src/umbp/doc/design-standalone-process-mode.md index b07d89282..f9d03a146 100644 --- a/src/umbp/doc/design-standalone-process-mode.md +++ b/src/umbp/doc/design-standalone-process-mode.md @@ -1,372 +1,96 @@ -# UMBP Standalone Process Mode — Design (Draft v1.4) - -**v1.0 revision note:** §11.4.4 (external-KV) went through two rounds -of direct clarification and was substantially corrected, not just -touched up — worth reading in full even if earlier revisions of this -section were already reviewed. Summary of the arc: v0.7 treated -external-KV as a data-plane feature and concluded the server can't -physically serve HBM-tier reads, so it decided to disable external-KV -entirely for a distributed-backed server. That premise was wrong: -external-KV is a pure query/registry service (`Report`/ -`MatchExternalKv` just record and answer "who holds hash H," a caller -— typically a custom SGLang router — uses the answer to route a *new* -request elsewhere; UMBP itself never moves the reported bytes, in -`distributed/` mode or here). With that corrected, the real question -became identity granularity, and a second clarification established -that the project's actual deployment shape is 8-GPU hosts running -SGLang TP=8, **and** that the *baseline* non-standalone implementation -already supports multiple independent replicas sharing one host for -free (`node_id` is built from per-process rank coordinates with no -awareness of logical TP/DP group boundaries, `umbp_store.py:512-519`). -Per explicit instruction to match the baseline rather than settle for -a narrower capability, §11.4.4 now specifies **per-worker distributed -sub-identities** (the server runs one independent, external-KV-only -identity per connected worker, sharing one physical `peer_address` — -verified this pass that `peer_address` uniqueness is not required -anywhere in `ClientRegistry`/`MasterServer` — see §11.4.4 for -citations) rather than either disabling the feature or collapsing all -workers into one shared identity. **(Corrected in v1.2: this is -implemented by a new, purpose-built `ExternalKvIdentityClient` class — -not `MasterClient` — specifically so each sub-identity reports empty -`tier_capacities`, publishes no owned-KV events, and is therefore -structurally ineligible for ordinary routing/eviction; see §11.4.4 and -§11.4.7 for the full reasoning and interface. The paragraph above -originally said "`MasterClient` registration," which is no longer -accurate — kept here only so the revision history stays honest about -what changed and why.)** This is real, non-trivial, correctly-scoped -implementation work (up to N `ExternalKvIdentityClient` -registrations/heartbeats inside one process), not a simple flag flip — -it was **not** in the v0.7/v0.9 implementation-cost accounting and -should be weighed accordingly when this proposal is scheduled. - -**v0.9 revision note:** a fourth review pass raised two further -points, both accepted as valid (not over-engineering): (1) -`PingResponse.deployment_mode` was wrongly filed as a low-stakes, -deferrable nice-to-have — it is actually the only guard against the -single most dangerous silent-failure mode this whole section exists to -prevent (a parity-intended deployment silently ending up local-backed, -with every RPC still succeeding); promoted to required, with a new -assertion requirement for distributed-backend launch/test tooling. (2) -§11.4.6's env-var list for the server's own distributed backend config -was incomplete (listed 6 vars, the real `umbp_store.py` surface has ~12 -including SSD staging and cache-admission knobs) and, separately, -missed that `dram_page_size`'s normal auto-derivation depends on a -`mem_pool_host` the server never has — both fixed, with a proposed -`BuildDistributedBackendConfigFromEnv()` helper shape and an explicit -required/optional split. See inline `(v0.9 fix)` markers in §11.4.5 and -§11.4.6. - -**v0.8 correction pass (citations only, no design changes):** every -Mooncake source citation in this document was originally verified -against `/apps/theresa/Mooncake` (commit `e1d6d6f6`, 2026-04-22). That -was the wrong path — the correct, current Mooncake checkout for this -project is `/apps/nima/KVManager/Mooncake` (commit `c9896684`, -2026-07-08, 75 commits ahead on the relevant files). All citations -were re-verified against the correct path this pass. Outcomes: -- Most citations were **directionally correct but line-stale** - (functions renamed/refactored, e.g. `map_shm_internal` is now a thin - wrapper around `map_shm_internal_with_device`; `ipc_send_fd`/ - `ipc_recv_fd` were renamed to `UdsConnection::sendFd`/`recvFd` and - moved into `uds_transport.cpp`) — corrected in place, marked - `(v0.7 citation fix)` inline (the fixes were made while the doc was - still at v0.7; the version bump to v0.8 here just reflects the - re-verification pass being complete). -- **One substantive error was found and corrected, not just a stale - line number**: this document had described Mooncake's primary shm - buffer as "the Real Client allocates it and hands the fd to the - application." The correct direction, confirmed against the current - checkout, is the opposite for the primary buffer: the **application** - (`DummyClient`) allocates and owns it, and sends the fd to the Real - Client, which is the passive second mapper. (A secondary, read-only - "hot cache" segment does go the other direction, but that is not - the buffer this design's citations were describing.) Reassuringly, - **UMBP's own already-shipped implementation never had this error** — - §3.2 already correctly has the worker (application) allocate and own - the buffer, matching Mooncake's real direction; only this document's - *prose describing* Mooncake's mechanism had drifted from the - primary-source direction. No code changes result from this - correction, only doc text (§11.3's "closed open item" section). -- The auto-start precedent claim ("zero matches" for fork/exec/spawn - patterns across `mooncake-integration/` and `mooncake-wheel/`) was - also corrected: `mooncake-integration/` (the actual application- - binding code) is still zero matches, but `mooncake-wheel/` (the - separate CLI/packaging layer) contains one real match — an explicit, - human-invoked CLI passthrough, not an implicit auto-start. The - underlying design decision (§11.4.6: distributed-backed - `umbp_standalone_server` supports external launch only, no - auto-start) is unaffected; only the supporting claim's precision was - corrected. - -**v1.1 revision note:** confirmed §11.4.4's per-worker `MasterClient` -sub-identity design (Option A) as final, after clarifying it involves -no worker→server→Master relay (each sub-identity heartbeats Master -directly from inside the server) and after explicitly weighing it -against a simpler Mooncake-aligned alternative (Option B: one shared -identity + `tenant_id`-style key namespacing, the pattern Mooncake -actually uses for its own multi-`DummyClient`-per-`RealClient` case). -Option B was rejected: preserving the `distributed/` baseline's -existing per-worker routing precision (verified as a real, already- -existing capability, not a nice-to-have) was judged more valuable than -the simplicity Option A gives up — accepted as genuinely new -engineering with no direct precedent in either UMBP or Mooncake, not a -known-safe pattern being reused. See the `(v1.1 —` marker in §11.4.4. - -**v1.2 revision note:** the v1.1 per-worker sub-identity design was -found to have three confirmed blockers, all fixed by replacing "N -independent `MasterClient` registrations" with a new, purpose-built -`ExternalKvIdentityClient` class: (A) a naive `MasterClient`-based -sub-identity with nonzero `tier_capacities` would have been eligible -for ordinary `RoutePut` target selection, silently polluting real KV -placement with fake storage nodes — closed structurally by giving the -new class no `tier_capacities`-reporting code path at all, rather than -relying on every future call site remembering to pass empty capacities -into `MasterClient`. (B) `MasterClient`'s existing re-register path has -a confirmed bug — it omits `peer_address`/`engine_desc` on re-register, -silently breaking `MatchExternalKv` responses after a Master -restart/reaper-expiry — closed by not reusing that code path at all; -the new class always includes both fields, by construction, rather -than patching a class shared by every real distributed peer. (C) the -registration-handshake field list was under-specified ("just -`node_id`") when the proto actually requires `node_address` too — -decided now: `worker_node_id` + `worker_node_address` + optional -`tags`, Python remains the single source of truth for deriving them, -no raw rank components cross the boundary. `ExternalKvIdentityClient` -itself is flagged as real, net-new implementation work not previously -accounted for in this proposal's cost estimate (§11.7 item 7). - -**v1.3 revision note:** a documentation-consistency and finalization -pass, no architecture changes from v1.2. Six items closed: (1) the -v1.0 revision note at the top of this document still said "`MasterClient` -registration per connected worker," contradicting §11.4.4's v1.2 -`ExternalKvIdentityClient` fix — corrected in place, kept visible -rather than silently rewritten. (2) `ExternalKvIdentityClient` was -still filed as an open item needing "a design pass" even though -§11.4.4 already made it a required component — it now has a full -interface/lifecycle spec in new §11.4.7 (Register/Heartbeat/Unregister/ -the five external-KV RPCs, the re-registration hard requirement, and -shutdown ordering). (3) confirmed via `eviction_manager.cpp:83-91` -(checked directly, not assumed) that the same empty-`tier_capacities` -invariant that hides a sub-identity from `RoutePut` also hides it from -`EvictionManager` — no separate mechanism needed, but stated -explicitly now rather than left to be inferred; also explicitly -decided `ExternalKvIdentityClient` heartbeats must never carry -`EventBundle`s, a separate `HeartbeatRequest` field from -`tier_capacities`. (4) the worker→server identity handshake now has a -concrete wire location: `worker_node_id`/`worker_node_address`/`tags` -are added to the existing `RegisterMemoryRequest` message -(`umbp_standalone.proto:87-91`), not a new RPC. (5) §11.4.6's five -new server-backend env vars are finalized as -`UMBP_DISTRIBUTED_STAGING_BUFFER_SIZE` and siblings (the -`UMBP_DISTRIBUTED_` prefix specifically because these five have no -existing worker-side env-var form to collide with, unlike the ones -reusing worker-side names verbatim) — no longer "proposed." (6) added -an explicit layering statement (new §11.4.7 lead-in): distributed- -backed server + UDS/shm fd handoff is the Mooncake-parity core; -`ExternalKvIdentityClient`/per-worker external-KV identity is a -separate UMBP-compatibility extension preserving a `distributed/`- -baseline capability, not required for declaring parity achieved, and -schedulable/testable independently. - -**v1.4 revision note:** two updates from the first implementation + -code-review round on `ExternalKvIdentityClient`. (1) A concurrency bug -found by review — `EnsureExternalIdentity`/`RemoveExternalIdentity` had -no serialization across their read-check/stop-old/construct-and-start- -new/insert sequence, so a `DeregisterMemory` racing an in-flight -`RegisterMemory` for the same `client_id` could silently fail to -remove an identity the register path was about to (re-)insert (a -permanent leak until full server restart), and two concurrent -`RegisterMemory` calls for the same `client_id` could briefly -double-register — was fixed with a single server-wide -`external_identity_lifecycle_mu_` serializing all three call sites; -§11.4.7 now documents this as an accepted trade-off (serializes -otherwise-unrelated workers' register/deregister calls when the lock -covers a slow/unavailable Master RPC, worst case approaching -`N * rpc_deadline` for the stated 8-worker deployment shape) rather -than a free fix, with an explicit note on when to revisit it -(per-`client_id` sharding, or decoupling external-KV registration from -the core memory-registration path). (2) §11.4.7's `Unregister()` -trigger list previously claimed worker crash/disconnect detection was -"already" handled by mirroring a `StandaloneProcessClient` state -machine — but §8 item 1 itself explicitly says that state machine -doesn't exist. This was an internal contradiction, not just a missing -callout: corrected to state plainly that no crash detection exists for -`ExternalKvIdentityClient` today, what the concrete consequence is (a -crashed worker's identity heartbeats forever, `ClientRegistry` never -expires it, only a full server restart clears it), and that this is a -deliberately deferred, tracked limitation rather than a silent gap. - -**Status (still true as of v1.4): §11 below is a NEW, UNAPPROVED -proposal. Nothing in §11 is implemented. Do not start coding against -it until it has been discussed — this revision exists to make the -proposal concrete enough to discuss, per explicit instruction not to -touch code first. v0.7 closed 4 of the 5 blockers a second review -pass found in the v0.6 draft (config/auto-start conflict, -external-KV routing correctness, backend registration -transaction/rollback + shutdown ordering, and the cross-mapping -memory-visibility open item); one naming/scope-framing question -remains genuinely open (§11.7 item 1).** §1-§10 -describe what is already implemented and shipped (v0.1-v0.5); §11 -revises the framing in §1 (see the callout there) to close a scope gap -identified after implementation: v0.1-v0.5 built a same-host-only -subsystem, whereas the stated goal was Mooncake-parity, and Mooncake's -"standalone" Real Client is cross-node-capable (it is a full -distributed client that happens to run in its own process, not a -narrower local-only subsystem). §11 proposes closing that gap without -discarding the v0.1-v0.5 implementation. - -**Scope:** A new deployment mode for UMBP in which the DRAM/SSD tiers, -`LocalBlockIndex`, and `CopyPipeline` run inside a dedicated, long-lived -OS process on the same host, while one or more SGLang/vLLM worker -processes talk to it through `IUMBPClient` as if it were still -in-process. This is *not* the existing `distributed/` mode (separate -master + multi-node peers over gRPC + RDMA) — it targets the -**single-host, multi-worker-process** case that today forces every -worker to duplicate its own DRAM tier and lose the ability to share a -KV cache across sibling processes (e.g. TP ranks, or a restarted -worker resuming a warm cache). - -**Status:** design proposal, not yet implemented. Everything under -"Open questions" needs a decision before implementation starts. - -**Revision history:** this draft has gone through two source-grounded -code review passes so far. Each found concrete implementation gaps in -the prior draft — inline "(v0.2 fix)" / "(v0.3 fix)" markers throughout -mark exactly what changed and why, so the reasoning stays visible -rather than silently overwritten. - -- **v0.2** (§3.1, §3.2, §4.3, §5, §6, §7, §8): the `is_distributed()` - gate that would silently skip `register_memory` in this mode (§5); - the missing mechanism for getting a `memfd` out of `HostBufferHandle` - and into `StandaloneProcessClient::RegisterMemory` (§3.2, §5); an - internal inconsistency between "same UDS" and "second UDS" for fd - handoff (§3.2, §8); an incompatible TCP-fallback claim for an - `SCM_RIGHTS`-based data plane (§3.1, §6, §8); an overreach in - claiming the server's own `DRAMTier` shared-memory backing needs no - new code (§3.2, §7); plus a missing `CopyPipeline::Drain()` for - clean shutdown (§4.3). -- **v0.3** (§3.1, §4.1, §5, §6, §10): the still-unclosed activation - path for `cfg.standalone_process` — nothing in v0.2 specified who - actually sets it from Python, the same way `umbp_store.py` already - does for `cfg.distributed` (§5, §10); a leftover default-path - inconsistency between `.sock` and `.grpc.sock` - plus an unspecified fd-socket derivation rule (§3.1, §4.1, §6); an - fd-ownership contradiction between the registry design and - "`DeregisterMemory` closes the worker's fd" (§3.2); a missing `Ping` - RPC that §4.1's readiness probe depends on, plus explicit - client-local-vs-RPC method classification (§3.1); and small - corrections (a stale `batch_resolder_codec.h` typo, a table cell - claiming the server "owns the shared-memory DRAM region" that - contradicted v0.2's own fix). -- **v0.4** (§4.2, §5, §6): `UMBPStandaloneProcessConfig`'s - field/parsing ownership was still ambiguous — §6's table described - `fd_socket`/`transport`/`auto_start` with routing language that - conflicted with both the config struct definition (only - `address`/`auto_start`/`startup_timeout_ms`) and the §5 Python - snippet (which only ever set `address`). Resolved by picking one of - the two existing, mutually-inconsistent precedents in this codebase - (`cfg.distributed`: all fields Python-set, vs. - `spdk_proxy_auto_start`: read directly by C++ - `FromEnvironment()`) and applying it consistently: `address`, - `auto_start`, and `startup_timeout_ms` are now all Python-set - (mirroring `distributed`), while `fd_socket` and `transport` are - removed from the config surface entirely — `fd_socket` is always - mechanically derived from `address` (no knob needed), and - `transport` has no v0.1 config surface since `tcp_staging` isn't - built yet (§8 item 6b). -- **v0.5** (§5, §10): two blockers found by tracing SGLang's actual - construction order and `register_mem_pool_host`'s actual error - handling, both confirmed and resolved with the team rather than - picked unilaterally, since both are scope/policy calls, not - implementation details: (1) `extra_config["standalone_address"]` - cannot reliably activate standalone-process mode — verified, the - host memory pool is constructed and allocates its buffer in - `HiRadixCache.__init__` (`hiradix_cache.py:91,106`, - `pool_host/base.py:97,139`) *before* `extra_config` is parsed and - handed to `UMBPStore` (`hiradix_cache.py:175,182`), so - `umbp_host_allocator.py` can only ever see environment variables at - the point it must decide a buffer's backing. **Decision: v0.1 - activation is `UMBP_STANDALONE_ADDRESS`-only; - `extra_config["standalone_address"]` without the env var raises.** - (2) `register_mem_pool_host`'s existing exception/`False` handling - (`umbp_store.py:862-929`) silently downgrades to a warning — correct - for `distributed`, which has a real staging-buffer fallback, but - wrong for standalone-process mode, which has none in v0.1: a - swallowed registration failure there means every later `Put`/`Get` - fails with no correlated error. **Decision: under - `get_deployment_mode() == StandaloneProcess`, all three cases - (`disable_zero_copy_register` set, `register_memory` exception, - `register_memory() == False`) raise instead of warn; - `distributed`'s behavior is unchanged.** - -**Primary references (source-verified, see §9):** -- UMBP's own SPDK-proxy daemon (`storage/spdk/proxy/`, - `local/tiers/local_storage_manager.cpp:SpawnProxyDaemon/EnsureProxyDaemon`) - — the only existing "library spawns a separate process, backed by a - shared-memory ring protocol" precedent in this codebase. -- **Mooncake's Real Client / DummyClient split** - (`mooncake-store/src/real_client.cpp`, `real_client_main.cpp`, - `shm_helper.cpp`, `dummy_client.cpp`) — explicitly the primary model - for this design per team guidance ("Mooncake 是我们的老师"). -- LMCache's `lmcache/v1/multiprocess/` — reviewed and **not** used as - the primary template (its zero-copy path is CUDA-IPC/GPU-only and - has no host-memory equivalent; see §8.4). - ---- - -## 1. Why not just reuse the existing `distributed/` mode? - -**(v0.6 note — this section's framing is revised by §11, not replaced. -Read this section for why v0.1-v0.5 built what they built; read §11 -for why that scope turned out to be narrower than "Mooncake parity" -requires, and how the two are reconciled without a rewrite.)** - -UMBP already runs multi-process today (`umbp_master` + peers over -gRPC + RDMA). That mode solves a different problem: **cross-node** -sharing with RDMA as the bulk-transfer fabric, master-mediated routing, -and lease/eviction protocols built around remote peers that may be on -a different machine. Standalone-process mode is strictly **same-host**: -its entire job is to let N worker processes that already sit on one -NUMA machine share one DRAM tier without RDMA, without a master, and -without inventing new eviction protocols. Reusing `distributed/` -wholesale means dragging in a master, RDMA registration, and the -peer/lease machinery for a problem that doesn't need any of it — this -design intentionally sits *underneath* the existing `local/` -(`StandaloneClient`) code, not beside `distributed/`. - -Concretely: **the standalone-process server is `StandaloneClient`'s own -`LocalStorageManager`/`LocalBlockIndex`/`CopyPipeline`, unchanged, moved -into its own process and fronted by an RPC+shm shell.** No new tier -logic, no new eviction policy, no new durability logic — those are -100% reused from `local/`. - ---- - -## 2. Component overview +# 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` | `mooncake_client` (Real Client) | its own process | Owns `LocalStorageManager` (DRAM/SSD tiers, private anon/hugetlb memory by default — see §3.2/§7, **not** a shared DRAM region), `LocalBlockIndex`, `CopyPipeline`. 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` | A new `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. | +| `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 │ │ - ▼ ▼ ▼ ▼ + ┌────────────────────────────┐ ┌────────────────────────────┐ + │ 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 │ │ │ @@ -376,2183 +100,948 @@ Two new roles, named after Mooncake's split for consistency: │ instead of raw pointers │ │ │ │ Fd-handoff listener (raw AF_UNIX, SEPARATE socket path from │ - │ the gRPC UDS — see §3.2/§8.3; gRPC cannot carry an fd) │ + │ the gRPC UDS — gRPC cannot carry an fd) │ │ │ - │ LocalStorageManager (REUSED, unmodified) │ - │ DRAMTier (own private anon/hugetlb memory, NOT shared — see │ - │ §3.2 note) ── SsdTier │ - │ LocalBlockIndex (REUSED, unmodified) │ - │ CopyPipeline (REUSED, unmodified; gains Drain(), see §4.3) │ + │ 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 — see §3.2/§8.3 -for why gRPC cannot double as the fd-handoff channel. +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 -## 3. IPC mechanism +`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: -### 3.1 Control plane: gRPC over a Unix domain socket - -**Decision: reuse UMBP's existing gRPC stack, not a new RPC framework, -and not Mooncake's `coro_rpc`.** - -Mooncake uses `ylt::coro_rpc` for its control plane. We explicitly do -**not** copy that choice: UMBP's `umbp_common` already depends on -gRPC/protobuf for the `distributed/` master/peer path -(`umbp/CMakeLists.txt:305-328,347-349`), and `pybind_umbp.cpp` already -has a proven, working convention for calling blocking C++ RPC methods -from Python with the GIL released -(`py::call_guard()`, -`pybind_umbp.cpp:248-279`). Adopting `coro_rpc` would require building -a new coroutine-to-pybind11 bridge with **no existing precedent in -either UMBP or the reports on Mooncake's own Python integration** — -pure added risk for zero benefit, since gRPC already does everything -Mooncake uses `coro_rpc` for (typed request/response, async server -loop, TCP or UDS transport). +| 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: -- New proto service `UMBPStandalone` (new file, - `distributed/proto/umbp_standalone.proto`). **(v0.3 fix — this list - is now explicit about what is an RPC, what is proto-but-not-`IUMBPClient`, - and what stays entirely client-local, to avoid an implementer - guessing.)** +- 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`, **and - `GetExternalKvHitCounts`** (`umbp_client.h:166-170`; named - explicitly here since it's easy to drop — it has a default no-op - body in the interface, unlike its siblings, which makes it look - optional when it isn't for a real implementation). - - `Ping(Empty) -> PingResponse{ready: bool}`: new RPC, **no** - `IUMBPClient` analog, exists solely for the §4.1 readiness probe. - - **Not RPCs, stay entirely client-local** in - `StandaloneProcessClient`: `Close()` (tears down the local gRPC - channel + fd-handoff socket + registry state; no server-side call - needed beyond whatever `DeregisterMemory` already covers) and the - new `get_deployment_mode()` accessor from §5 (pure local constant, - never needs a round trip). - Reuse the existing struct-of-arrays batch codec approach from - `batch_resolve_codec.h` - (`distributed/peer/batch_resolve_codec.h:24-28`) for `BatchPut`/ - `BatchGet` so per-key protobuf submessage overhead doesn't reappear - here — this is exactly the pattern UMBP already validated for the - peer `BatchResolveKeys` RPC. -- **Default transport: gRPC over a Unix domain socket** - (`unix:///run/umbp/standalone/.grpc.sock`), not TCP. This is a - deliberate deviation from Mooncake (TCP, port 50052) and LMCache - (TCP, port 5555): both of those need to be reachable from a - different host or container in the general case, but standalone - mode by definition never leaves the host, so a UDS avoids picking a - port, avoids the network stack, and gets free filesystem-permission - based access control (§8.5). gRPC's UDS support - (`unix:` target scheme) is mature and already linked in via the same - gRPC dependency `umbp_common` uses today. - **(v0.2 fix) TCP is not a drop-in fallback for this design and must - not be presented as one.** §3.2's data plane depends on `SCM_RIGHTS` - fd passing, which is an `AF_UNIX`-only mechanism — there is no - TCP-socket equivalent for handing a `memfd` to another process, and - because the design deliberately uses an anonymous `memfd_create` - (not a named `shm_open` segment, precisely to avoid `/dev/shm` name - collisions, see §3.2), there is also no name the TCP-connected side - could use to reopen the segment out-of-band. Concretely, for v0.1: - the fd-handoff channel (§3.2, §8.3) **must** be a UDS reachable by - both processes (i.e. they must share the same mount/PID namespace - for `/dev/shm`-adjacent sockets, or a bind-mounted socket directory - across containers). If a deployment genuinely cannot provide that - (fully separate network namespaces with no shared filesystem), it - cannot use the zero-copy data plane at all; the only option is a - non-zero-copy fallback where `Put`/`Get` payload bytes travel inside - the gRPC message itself (i.e. the same "no RDMA, no shm" degraded - path `distributed/` already has via `staging_buffer_size`-capped - copies when `UMBP_DISABLE_ZERO_COPY_REGISTER=1`), not a variant of - the shm data plane over TCP. The gRPC *control* channel could in - principle still run over TCP in that scenario, but the two are no - longer both UDS by construction — this is a materially different, - slower deployment mode conceptually named `tcp_staging` to - distinguish it from the default `uds` path. **(v0.4 fix)** it has no - config/env surface in v0.1 at all, precisely because it isn't built - in v0.1 (§8 item 6b, §10 step 13) — see §6 for why a `transport` - knob is deliberately not added to `UMBPStandaloneProcessConfig` - ahead of the feature it would select. -- 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 today is - translated to `(client_id, shm_offset)` **inside** - `StandaloneProcessClient` before the RPC is issued (see §3.2) — this - is exactly Mooncake's `map_dummy_buffer_range_to_real` / - `shm_addr_offset` translation **(v0.7 citation fix — re-verified - against the correct checkout; function confirmed still present, - only line numbers moved from a stale reference)**: - `real_client.cpp:369-379` (`map_dummy_buffer_range_to_real`), called - from the batch RPC handlers `get_into_range_shm_helper`/ - `get_into_ranges_shm_helper` (`real_client.cpp:4291`/`:4316`), just - carried over gRPC instead of coro_rpc. - -### 3.2 Data plane: shared memory, modeled directly on Mooncake's `ShmHelper` - -This is the part of the design that should track Mooncake most -closely, because UMBP and Mooncake independently converged on the same -shape (control message carries a handle, bytes live in shared memory) -— UMBP's own SPDK-proxy already does this too -(`spdk_proxy_protocol.h:RingSlot.data_offset/data_size`). The novelty -here is *who allocates the shared segment*. - -**(v0.2 fix — this paragraph overreached in v0.1.)** `DRAMTier` does -already support POSIX shared memory — `use_shared_memory`/`shm_name` -in `UMBPDramConfig` (`common/config.h:57-58`) are wired through -`LocalStorageManager` (`local_storage_manager.cpp:443-444`) into -`DRAMTier`'s constructor, which does `shm_open`+`ftruncate`+`mmap`+ -`shm_unlink` (`dram_tier.cpp:104-175`). But two things in the v0.1 -draft were wrong about what this buys us: - -- The server's own DRAM cache **does not need to be visible to any - other process** — it is read/written exclusively by - `umbp_standalone_server` itself; no worker ever `mmap`s it directly - (workers only ever see their own registered host KV buffer, per - below). So `umbp_standalone_server` should default to the same - `kAnonymous`/`kAnonymousHugetlb` private backing - `StandaloneClient` already uses today — **not** - `use_shared_memory=true`. There is no reason to pay for a named - POSIX shm segment (and its cleanup/lifecycle concerns) for memory - nothing else needs to touch. -- Even where a future variant *did* want `use_shared_memory=true` (for - example, a warm-restart scenario that reattaches to a still-live - segment after the server process is replaced), the existing - `shm_open(shm_name_.c_str(), O_CREAT | O_RDWR, 0666)` call - (`dram_tier.cpp:114`) creates the segment world-read-write. That is - a pre-existing gap in `DRAMTier` itself (not introduced by this - design), and this design must not casually inherit it by pointing - more traffic at that code path while claiming "zero new C++ code" — - tightening `dram_tier.cpp`'s permissions is a separate, pre-existing - fix, out of scope here unless/until a future revision actually - chooses to share the server's DRAM tier across processes. - -The actual sharing requirement in this design is one-directional and -much narrower than "the whole DRAM tier": only the **worker's host KV -buffer** needs to become visible to the server process. That is the -problem the rest of this section solves. - -The problem this design actually has to solve is the **other** side: -today, `Put(key, src, size)`/`Get(key, dst, size)` take a pointer that -is only valid in the *caller's* address space -(`umbp_client.h:47-171`) — this is fine when `StandaloneClient` runs -in-process, and it is fine in `distributed/` mode because RDMA reads -directly from the registered remote memory without needing a shared -mapping. It is **not** fine when the server is a separate process on -the same host with no RDMA involved: the server cannot dereference a -pointer that only exists in the worker's page tables. - -Adopt Mooncake's answer exactly: + `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()` → `shm_helper_->allocate(...)` **(v0.7 citation fix - — re-verified against the correct Mooncake checkout, - `/apps/nima/KVManager/Mooncake`, not the stale one earlier research - in this doc used; line numbers below are corrected accordingly, the - underlying claim is unchanged and was directionally correct all - along)**: `dummy_client.cpp:469`, inside `setup_dummy()` - (`dummy_client.cpp:430-509`). Concretely: `UMBPHostMemAllocator` - (`host_mem_allocator.h/.cpp`, 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 (`shm_helper.cpp:74-140`: - `memfd_create_wrapper` at `:26-32`, `memfd_create` at `:108`, - `ftruncate` at `:117`, `mmap(MAP_SHARED|MAP_POPULATE)` at - `:123-124`), chosen over - UMBP's own named-`shm_open` DRAM-tier path specifically to avoid - `/dev/shm` name collisions across many SGLang instances on one - host (Mooncake's rationale for `memfd_create` over named - `shm_open` — anonymous, refcounted purely by open fds, nothing to - leak into `/dev/shm` if a process is killed uncleanly). + `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 - (`umbp_store.py:912` → `register_mem_pool_host`). - - **(v0.2 fix — this step was underspecified in v0.1.)** - `IUMBPClient::RegisterMemory(uintptr_t ptr, size_t size)` - (`umbp_client.h:119`) and the pybind binding - `register_memory(ptr, size)` only carry a pointer and a size — there - is no fd parameter anywhere on this path, and `HostBufferHandle` - (`host_mem_allocator.h:41`) has no fd field either. The Python side - never sees a file descriptor (`kv_buffer.data_ptr()` is just an - integer, `umbp_store.py:896`), so it has nothing to pass through - pybind even if we wanted to add an argument. The fd has to be - recovered entirely on the C++ side of the process boundary that - already exists (`StandaloneProcessClient` and the allocator both - live in the same worker process), not threaded through Python: - - - `HostMemAllocator::Alloc` gains an internal, process-local - registry (a mutex-guarded `std::map` - where `AllocRecord` holds `{fd, size}`), populated only for + `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 lookup, since `ptr` may be - `kv_buffer.data_ptr()` — the tensor's base address, which for the - host KV pool is normally exactly the allocation's base, but a - range/floor lookup is safer than requiring exact match) and, on a - hit, sends the recovered `fd` to the server via `SCM_RIGHTS` - ancillary data over the **separate raw-UDS fd-handoff socket** - (`.fd.sock`, distinct from the gRPC UDS — see §8.3), - mirroring Mooncake's fd-passing and registration handshake - **(v0.7 citation fix — the function names `ipc_send_fd`/ - `ipc_recv_fd` no longer exist in the current Mooncake checkout; - the mechanism was renamed/relocated, not removed)**: fd send/recv - is now `UdsConnection::sendFd`/`UdsConnection::recvFd` - (`uds_transport.cpp:167`/`:196`), and the registration entry point - is `RealClient::handle_ipc_shm_register` - (`real_client.cpp:5311`, dispatched from the `IPC_SHM_REGISTER` - case at `real_client.cpp:5286-5287`). On a miss (the pointer isn't - backed by a `kAnonymousShm` allocation — e.g. standalone-process - mode was configured but the host allocator wasn't switched to the - shm backing), `RegisterMemory` must fail loudly rather than - silently falling back to a broken pointer-based `Put`/`Get`. + `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 (`real_client.cpp:2136-2138`, inside - `map_shm_internal_with_device`) and stores `client_id → {base, - size}` in a registry keyed by `client_id` (mirrors - `shm_contexts_[client_id]`, e.g. `real_client.cpp:2158` — the map - itself is unchanged, only specific usage line numbers moved - across the file's many call sites). + 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 (mirrors - `RealClient::map_dummy_buffer_range_to_real`, now at - `real_client.cpp:369-379`, called from the batch RPC surface at - `get_into_range_shm_helper`/`get_into_ranges_shm_helper`, - `real_client.cpp:4291`/`:4316` — **(v0.7 citation fix: function - still exists exactly as described, only line numbers moved; the - RPC-handler line numbers were re-confirmed this pass, unlike some - other citations in this section which were only spot-checked)**). - No KV bytes ever cross the gRPC channel; the gRPC message is a - handful of integers per key. -4. `DeregisterMemory` is the inverse — **(v0.3 fix: fd ownership - clarified, since the v0.2 wording here conflicted with the §3.2 - step 2 registry design).** The allocator's ptr→fd registry (§3.2 - step 2), not `RegisterMemory`/`DeregisterMemory`, owns the fd: - `HostMemAllocator` opened it at `Alloc` time and is the only thing - that closes it, at `Free` time. `RegisterMemory` merely *borrows* - that fd to send over `SCM_RIGHTS`; sending a fd over `SCM_RIGHTS` - 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 of "this range - is registered" — it must **not** close the allocator-owned fd. - Closing happens exactly once, in `HostMemAllocator::Free`. Getting - this wrong either double-closes an fd number the kernel may have - already recycled for something else, or leaves a re-`RegisterMemory` - call after a `DeregisterMemory` trying to send an already-closed fd - — both are the kind of bug that only shows up under a - register/deregister/re-register cycle, so the unit test in §10 - step 11 must specifically exercise that cycle, not just a single - register→use→deregister pass. - -Net effect: the hot path becomes *exactly* the same shape as UMBP's -existing RDMA-registered `distributed/` path (`RegisterMemory` once, -then reference-by-handle on every call) — this is not a new interface -concept for UMBP, just a new transport underneath a pattern -`IUMBPClient` already exposes. - -### 3.3 Why not a busy-polled ring buffer (SPDK-proxy's own approach)? - -UMBP's own precedent for "separate process, shared-memory data plane" -(`spdk_proxy_protocol.h`) is a fixed-slot ring buffer with busy-polling -client-side (`spdk_proxy_tier.cpp:240,354`, `sleep_for(ProxyPollInterval())`) -and no RPC framework at all — every control operation (attach, admin -shutdown, stats) is also squeezed into the same 256-slot ring. This -works well for SPDK-proxy because it's low cardinality (few tenants, -few request types) and latency-critical at the microsecond level for -raw block I/O. It is a poor fit for the standalone-process control -plane because: - -- it has no natural way to express variable-shape requests (batch - registration, external-KV admission with variable-length hash - lists) without hand-rolling a second serialization format inside - the slot payload — gRPC already solves this, -- it offers no async multiplexing for slow control operations - without blocking the ring for the next request, -- 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 strictly for bulk bytes, which is the strategy this - design also adopts. - -The *data plane* still ends up looking like SPDK-proxy's idea (handle -+ offset/size, bytes in shm) — we're only rejecting the ring buffer -for the **control** plane, not the "shared memory for bytes" idea -itself. - ---- - -## 4. Lifecycle management - -This is the one dimension where UMBP's own precedent (SPDK-proxy: -library-spawned, idle-self-exit) and the Mooncake/LMCache precedent -(externally launched, long-lived, no self-exit) genuinely disagree, -and a specific choice has to be made for this codebase. - -**Decision: default to Mooncake's model (externally launched, -long-lived), with UMBP's SPDK-proxy auto-start pattern available as an -opt-in convenience, and idle-self-exit *disabled by default*.** - -Rationale: unlike the SPDK-proxy (a stateless passthrough to a shared -NVMe device — nothing is lost if it exits and respawns), the standalone -server **holds the DRAM cache tier's actual data**. Self-exiting on -idle silently discards the cache, defeating the purpose of the whole -feature. Neither Mooncake's Real Client nor LMCache's multiprocess -server self-terminate in the reports we reviewed — that absence is -itself evidence, not just a gap. - -### 4.1 Primary path: externally launched + 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` today (background process, log redirected, - `trap cleanup EXIT`) — no new orchestration concept, just a new - binary target. + 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 exact - naming convention of `UMBP_MASTER_ADDRESS` - (`runtime-env-vars.md:146-152`) for consistency. -- **fd-socket path derivation, made explicit (v0.3), confirmed to have - no separate config/env knob (v0.4 — see §6).** The fd-handoff socket - path is derived from `cfg.standalone_process.address` mechanically, - by both processes independently: strip the `unix://` scheme prefix - to get a filesystem path (raw `AF_UNIX` `connect()`/`bind()` need a - path, not a URI), then replace a trailing `.grpc.sock` with - `.fd.sock`; if the configured address doesn't end in `.grpc.sock` - (a custom path), append `.fd.sock` to it instead of doing a - substring replace. E.g. - `unix:///run/umbp/standalone/node0.grpc.sock` → - `/run/umbp/standalone/node0.fd.sock`. This derivation must live in - one shared place (`StandaloneProcessClient`'s ctor and - `umbp_standalone_server`'s startup both call it) rather than being - reimplemented on each side, to guarantee they agree — and precisely - because it's a pure function of `address`, it deliberately has no - independent `UMBP_STANDALONE_FD_SOCKET` env var or config field to - keep in sync (§6). -- **(v0.3 fix) `Ping` must be added to the `UMBPStandalone` proto - service.** The method list in §3.1 mirrors `IUMBPClient` 1:1 and - `IUMBPClient` has no `Ping` method, so a reader implementing §3.1's - proto as written would have nothing for this readiness probe to - call. `Ping(Empty) -> PingResponse{ready: bool}` is a new, - UMBP-standalone-specific addition to the proto with no `IUMBPClient` - analog — call it out explicitly in §3.1's method list, not just - here in §4.1. -- Readiness probe: connect to the UDS and call `Ping` - (equivalent to `ProxyShmRegion::ProbeExisting`'s role for - SPDK-proxy, but over gRPC instead of a shm header flag, since there - is no shm header to probe until a client registers memory). - -### 4.2 Convenience path: auto-start (opt-in) - -- **(v0.4 fix — routing corrected to match §6's decision.)** When - `cfg.standalone_process.auto_start` is `true` (Python-set from - `UMBP_STANDALONE_AUTO_START`/`extra["standalone_auto_start"]`, §5, - §6 — `CreateUMBPClient` itself never reads that env var; by the time - it runs, the bool is already sitting on the config it was handed): - 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 - already in `UMBPConfig::FromEnvironment` - (`config.h:410-428`) — no new election protocol) probes the UDS at - `cfg.standalone_process.address`; if absent, it does the same - `fork()` + `setsid()` + `execlp()` sequence - `LocalStorageManager::SpawnProxyDaemon` already uses - (`local_storage_manager.cpp:319-372`), passing the spawned server's - own config down via env vars (this is a different hop than the one - above — see `UMBP_STANDALONE_BIN` in §6, which *is* read directly by - this C++ auto-start code, matching its deployment-only status), and - waits for readiness bounded by - `cfg.standalone_process.startup_timeout_ms` (also Python-set, §6 — - mirrors `spdk_proxy_startup_timeout_ms`'s *role* only; unlike its - SPDK-proxy analog it is not itself read by - `UMBPConfig::FromEnvironment`, per §6's Precedent A decision). A - bootstrap lock (reuse the existing `ScopedBootstrapLock` pattern, - `local_storage_manager.cpp:274-283`) prevents a thundering herd of - workers all trying to spawn the server simultaneously. -- **Not** the default, because auto-start from inside a worker process - means the worker's `fork()`'d child inherits CUDA/HIP context, open - gRPC channels, and GPU device state at the moment of fork — exactly - the fork-safety hazard `SpawnProxyDaemon` already has to be careful - about (`local_storage_manager.cpp:328` calls `setsid()` immediately, - and the child does nothing but `execlp()` before replacing itself — - **this rule must not be relaxed** for the standalone-server spawn - path either; no CUDA/HIP/gRPC call may happen between `fork()` and - `execlp()`). - -### 4.3 Shutdown - -- `SIGTERM`/`SIGINT` → drain in-flight RPCs (bounded deadline, - mirroring `UMBP_GRPC_SHUTDOWN_DEADLINE_SEC` already used by - `MasterServer::Shutdown()`, `master_server.cpp:809-816`) → flush - `CopyPipeline` (best-effort persist DRAM-tier dirty pages to SSD - before exit, if the SSD tier is enabled and - `force_ssd_copy_on_write` is set) → `munmap` all registered client - shm segments → exit. - - **(v0.2 fix — no API exists today for the "flush CopyPipeline" - step above.)** `CopyPipeline` only drains its queue in its - destructor (`copy_pipeline.cpp:40-52`, `stop_copy_worker_` + - worker-thread join), and `StandaloneClient::Flush()` - (`standalone_client.cpp:300`) only calls `storage_.Flush()` — it - never waits for the async SSD-copy queue to empty. A clean shutdown - that wants "no dirty DRAM pages lost" needs the queue drained - *before* the object is destroyed, while the server can still respond - to health checks. This requires a small, genuinely new piece of - work, not a reuse of an existing hook: add - `CopyPipeline::Drain(std::chrono::milliseconds timeout)` (blocks - until `copy_queue_` is empty and in-flight copies complete, or the - timeout elapses) and have both `umbp_standalone_server`'s shutdown - path and the `Flush` RPC call it. See §10 implementation order — - this is listed as its own task, not folded into "reuse - `LocalStorageManager` unmodified" from §1, because it genuinely - isn't unmodified once this exists. -- 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 recommended - worker-side behavior is to treat this the same as "cache miss / - cache unavailable" rather than crashing the inference process - (SGLang's `HiCacheStorage` abstraction already tolerates backend - failures returning `False`/`None`; this needs to be threaded through - explicitly — see Open Question in §8.2, since neither Mooncake nor - LMCache's reviewed code shows a supervised-restart or - reconnect-on-crash path to copy). - -### 4.4 Multi-tenant workers on one host - -- Multiple SGLang worker processes (e.g. one per TP/DP rank) may - attach to the same standalone server. Each gets a `client_id` - (UUID, assigned at first `RegisterMemory` call) and its own shm - registration entry — mirrors Mooncake's per-client `shm_contexts_` - map and UMBP's own SPDK-proxy per-tenant quota concept - (`TenantInfo.quota_bytes`, `spdk_proxy_protocol.h:227-240`). - Capacity accounting inside `LocalStorageManager`/`DRAMTier` is - already global (not per-caller) today — **this is an open question**, - not solved by this design (see §8.3): without per-tenant quotas, one - worker can starve another's share of the shared DRAM tier. - ---- - -## 5. Python / pybind interface adaptation - -**(v0.2 fix — the v0.1 headline claim "zero call-site changes in -`umbp_store.py`" is false and is retracted here.)** One call site does -need to change, and the reason is a real bug the v0.1 draft -introduced by combining two of its own recommendations -inconsistently: §5 (v0.1) said `is_distributed()` should keep meaning -"true cross-node distributed" and *not* be overloaded for -standalone-process mode, while `register_mem_pool_host` -(`umbp_store.py:870-881`) gates the entire `register_memory` call on -`if not is_distributed: return` (`umbp_store.py:880`). Taken together, -a `StandaloneProcessClient` — which correctly reports -`is_distributed() == False` per the v0.1 recommendation — would never -have its host KV buffer registered, and therefore never trigger the -fd-handoff handshake in §3.2 at all. The whole standalone-process data -plane silently never activates. This is not a minor wording issue; it -is a functional gap that must be fixed as part of this design, not -left as an implementation detail: - -- Add `get_deployment_mode()` (enum: `Local`/`StandaloneProcess`/ - `Distributed`) to `IUMBPClient`/pybind, as already proposed below. -- Change the gate at `umbp_store.py:880` from `if not is_distributed: - return` to also allow `StandaloneProcess`, e.g. - `if deployment_mode not in (Distributed, StandaloneProcess): return`. - This *is* a call-site change in `umbp_store.py`, and needs to be - listed explicitly in the implementation order (§10), not implied. - -With that correction, the remaining adaptation is small: - -- `CreateUMBPClient(const UMBPConfig&)` (`umbp_client_factory.cpp:28-33`) - gains a third branch: `config.standalone_process.has_value()` (new - optional field on `UMBPConfig`, alongside the existing `distributed` - optional) → construct `StandaloneProcessClient`. Precedence, if both - happen to be set, should be `distributed` first (cross-node case is - strictly more general) — but in practice these are mutually - exclusive deployment choices and `UMBPConfig::Validate()` should - reject setting both. -- **(v0.3 fix, superseded by v0.5 below — kept for history.)** Nothing - in v0.2 said *who sets* `cfg.standalone_process`; v0.3 filled that in - by mirroring how `umbp_store.py` populates `cfg.distributed` from - either `extra_config["master_address"]` or `UMBP_MASTER_ADDRESS` - (`umbp_store.py:458-465`). -- **(v0.5 fix — the v0.3 answer above is wrong for this specific field, - for a reason that doesn't apply to `master_address`.)** Mirroring - `master_address`'s two-source pattern (`extra_config` *or* env) - assumes both sources are read at the same point in the SGLang - startup sequence. They are not, and the difference matters here - specifically because of a new constraint standalone-process mode - introduces that distributed mode never had: **verified**, in - `HiRadixCache.__init__`, the host memory pool - (`MHATokenToKVPoolHost`, via `pool_host/base.py`) is constructed with - `allocator_type=server_args.hicache_storage_backend` - (`hiradix_cache.py:91,106`) and immediately allocates `kv_buffer` in - its own constructor (`pool_host/base.py:97` `get_allocator_from_storage`, - `pool_host/base.py:139` `self.kv_buffer = self.init_kv_buffer()`) — - this happens *before* `storage_backend_extra_config` is even parsed - and handed to `CacheController`/`UMBPStore` (`hiradix_cache.py:175,182`, - well after the pool already exists). `UMBPHostTensorAllocator` (used - by that pool) therefore can only see **process environment - variables** at the moment it decides a buffer's backing — never - `extra_config`, which doesn't exist yet at that point in the process. - This asymmetry is new to standalone-process mode: `distributed` - mode's `RegisterMemory` pins/registers *whatever pointer it's given* - for RDMA regardless of how it was allocated, so `master_address` - arriving late via `extra_config` was always fine — nothing upstream - needed to know in advance. Standalone-process mode's - `RegisterMemory` instead needs the buffer to already be a - `kAnonymousShm` allocation *before* `register_mem_pool_host` ever - runs, and that decision is made by `umbp_host_allocator.py` strictly - earlier than `UMBPStore.__init__` gets to see `extra_config` at all. - **Decision (confirmed with the team): v0.1 supports activation via - `UMBP_STANDALONE_ADDRESS` only.** `extra_config["standalone_address"]` - is not a valid activation path in v0.1 and `UMBPStore.__init__` must - reject it outright with an actionable error rather than silently - accepting a config that the host allocator already missed: - ```python - if extra.get("standalone_address") and not _optional_env_str( - "UMBP_STANDALONE_ADDRESS" - ): - raise ValueError( - "standalone_address in hicache-storage-backend-extra-config is " - "not supported: the host memory pool allocator decides its " - "backing (Anonymous vs. AnonymousShm) before extra_config is " - "parsed, so extra_config alone cannot activate standalone-process " - "mode reliably. Set the UMBP_STANDALONE_ADDRESS environment " - "variable instead (see design-standalone-process-mode.md §5)." - ) - standalone_address = _optional_env_str("UMBP_STANDALONE_ADDRESS") - if standalone_address and master_address: - raise ValueError( - "master_address and UMBP_STANDALONE_ADDRESS are mutually " - "exclusive (distributed vs. standalone-process mode)" - ) - if standalone_address and UMBPStandaloneProcessConfig is not None: - standalone_cfg = UMBPStandaloneProcessConfig() - standalone_cfg.address = str(standalone_address) - # auto_start/startup_timeout_ms may still come from extra_config -- - # unlike `address`, they carry no allocator-timing constraint, - # since they only affect what CreateUMBPClient does once cfg is - # already fully built. Only `address` (the activation switch) is - # env-only; these two may keep the extra/env fallback pattern. - standalone_cfg.auto_start = _strict_bool( - extra.get("standalone_auto_start", _optional_env_str("UMBP_STANDALONE_AUTO_START")), - "standalone_auto_start", - ) if extra.get("standalone_auto_start") or _optional_env_str("UMBP_STANDALONE_AUTO_START") else False - standalone_cfg.startup_timeout_ms = int( - extra.get( - "standalone_startup_timeout_ms", - _optional_env_str("UMBP_STANDALONE_STARTUP_TIMEOUT_MS") or 30000, - ) - ) - cfg.standalone_process = standalone_cfg - ``` - This must run, and `cfg.standalone_process` must be populated, - **before** `UMBPClient(cfg)` is constructed (`umbp_store.py:790`) — - listed as its own step in §10 rather than folded silently into "wire - into `umbp_client_factory.cpp`", since that step is C++-only and - does not by itself make Python ever populate the field. The rejected - alternative — teaching `hiradix_cache.py`/`kv_cache_builder.py`/ - `memory_pool_host.py` to see `extra_config` before constructing the - host pool, so `extra_config` could activate standalone-process mode - too — was considered and explicitly deferred: it touches SGLang's - core pool-construction sequence outside this feature's natural - boundary, for a convenience (config-file-only activation) that - `UMBP_STANDALONE_ADDRESS` already covers. Revisit only if an env-var-only - activation switch turns out to be a real deployment blocker in - practice. -- **(v0.2 fix)** `pybind_umbp.cpp` is **not** change-free the way v0.1 - claimed. The `IUMBPClient` method binding table itself - (`put_from_ptr`/`get_into_ptr`/etc., `pybind_umbp.cpp:248-279`) - indeed needs no changes, since it dispatches through the vtable and - `StandaloneProcessClient` implements the same interface — that part - of the v0.1 claim was correct. But at least two *other* bindings in - the same file are missing and block this design from working at - all: - - `HostBufferBacking`/`UMBPHostBufferBacking` - (`pybind_umbp.cpp:42-45`) currently only exports `Anonymous` and - `AnonymousHugetlb` — it needs a third value, - `AnonymousShm`/`kAnonymousShm`, or `umbp_host_allocator.py` has no - way to ask for the new backing added in §3.2. - - `UMBPConfig`'s pybind class (`pybind_umbp.cpp:230-239`) only - exposes `dram`/`ssd`/`eviction`/`copy_pipeline`/`role`/ - `follower_mode`/`force_ssd_copy_on_write`/`distributed` — it needs - a new `standalone_process` property (and a new - `UMBPStandaloneProcessConfig` pybind class, mirroring how - `UMBPDistributedConfig` is bound at `pybind_umbp.cpp:218-228`) - before Python can construct or detect this config at all. - The GIL-release `call_guard` convention on the existing methods is - still correct as-is, since `StandaloneProcessClient`'s methods block - on a gRPC call and never call back into Python — same justification - as the existing comment (`pybind_umbp.cpp:245-246`, "block on RDMA, - SSD, or gRPC"). -- **(v0.5 fix — must read the raw env var directly, not `UMBPConfig`.)** - `umbp_host_allocator.py` needs one change: `UMBPHostTensorAllocator.__init__` - reads `os.environ["UMBP_STANDALONE_ADDRESS"]` directly (same style as - its existing `SGLANG_HICACHE_HOST_*` env reads, - `umbp_host_allocator.py:42-47`) — **not** anything derived from - `UMBPConfig` or `extra_config`, since neither exists yet at this - point in the process per the v0.5 fix above. If set, request the new - `AnonymousShm` backing from `UMBPHostMemAllocator.alloc(...)` instead - of the default anonymous/hugetlb backing. This is a small branch in - `UMBPHostTensorAllocator.allocate()` (`umbp_host_allocator.py:50`) — - the returned `handle.ptr` is used identically afterward (same - `ctypes.c_byte.from_address` + `torch.frombuffer` zero-copy wrap, - `umbp_host_allocator.py:86-87`), gated on the new `AnonymousShm` - pybind enum value existing (see above). -- **(v0.5 fix — the existing exception/false handling in - `register_mem_pool_host` was written for `distributed`'s optional, - gracefully-degradable RDMA registration and must not apply as-is to - standalone-process mode, which has no staging/inline fallback to - degrade to in v0.1.)** Verified: `register_mem_pool_host` - (`umbp_store.py:862-929`) currently (a) treats - `disable_zero_copy_register` as "log info, skip registration, stay - on staging path" (`umbp_store.py:~886-892`), (b) wraps the - `register_memory` call in `except Exception as exc: logger.warning(...); - return` (`umbp_store.py:~913`), and (c) treats a `False` return as - `logger.warning(...)` only, no raise (`umbp_store.py:~929`). All - three are correct for `distributed` — a failed/declined RDMA - registration there really does have a working fallback - (`distributed.staging_buffer_size`-capped copies). For - standalone-process mode none of these three can be a silent - downgrade, because there is nothing to downgrade *to*: a registry - miss means every subsequent `Put`/`Get` will fail against the - server's offset translation, just later and with a less - correlated-looking error. **Decision (confirmed with the team): when - `get_deployment_mode() == StandaloneProcess`, all three cases raise - instead of warn-and-return; `distributed`'s existing warn-and-fallback - behavior is unchanged.** Concretely, `register_mem_pool_host` branches - on `get_deployment_mode()` before each of the three points above: - ```python - is_standalone_process = ( - self.client.get_deployment_mode() == UMBPDeploymentMode.StandaloneProcess - ) - if is_standalone_process and getattr(self, "_disable_zero_copy_register", False): - raise RuntimeError( - "disable_zero_copy_register is not supported in standalone-process " - "mode: there is no staging-buffer fallback path to degrade to (v0.1)." - ) - ... - try: - ok = bool(self.client.register_memory(host_ptr, host_size)) - except Exception as exc: - if is_standalone_process: - raise RuntimeError( - f"register_memory failed in standalone-process mode: {exc}. " - "This is fatal, not a degradable condition -- see design-" - "standalone-process-mode.md §5." - ) from exc - logger.warning(...) # distributed: unchanged, falls back to staging - return - if not ok: - if is_standalone_process: - raise RuntimeError( - "register_memory returned false in standalone-process mode " - "(fatal -- no fallback path exists in v0.1)." - ) - logger.warning(...) # distributed: unchanged - ``` - `register_mem_pool_host`/`register_memory` call sites themselves - (`umbp_store.py:896-912`) still pass `(host_ptr, host_size)` - unchanged — the fd-passing handshake described in §3.2 happens - entirely inside `StandaloneProcessClient::RegisterMemory` in C++ via - the process-local fd registry (§3.2 step 2), invisible to Python. - Only the **gate** above it (`umbp_store.py:880`) and the - **error-handling branches** just described change, both listed as - their own items in §10. -- `is_distributed()` remains a pure bool getter, unchanged, and - continues to mean "true cross-node distributed" for the existing - external-KV branch points that already depend on it. The new - `get_deployment_mode()` accessor is not optional polish here — §5's - own fix above depends on it existing. - ---- - -## 6. New / reused configuration surface - -**(v0.4 fix — this section previously left it ambiguous whether each -knob is a `UMBPStandaloneProcessConfig` field set by Python or an env -var read directly by C++, and used both descriptions inconsistently -for the same knobs. Resolved below by picking one routing rule per -knob and sticking to it, using the two conflicting precedents already -in this codebase as the tie-breaker.)** - -Two existing precedents disagree on how "which deployment mode, and -with what parameters" config should flow, and this design has to pick -one per knob rather than blend them: - -- **Precedent A — `cfg.distributed`:** every sub-field - (`master_address`, `node_address`, `node_id`, ...) is read and - assigned **entirely in Python** (`umbp_store.py:458-495`), each via - the same `extra.get(key, _optional_env_str(ENV_VAR))` pattern. - `UMBPConfig::FromEnvironment()` never touches any of them. -- **Precedent B — `UMBPSsdConfig::spdk_proxy_auto_start`/ - `spdk_proxy_startup_timeout_ms`:** read **directly by - `UMBPConfig::FromEnvironment()`** in C++ (`config.h:399-402`, `UMBP_SPDK_PROXY_AUTO_START`/`UMBP_SPDK_PROXY_TIMEOUT_MS`), with no - Python involvement at all. - -**Decision: `standalone_process` follows Precedent A end to end** — -every field Python can set, it does, exactly like `distributed` -already works, rather than splitting the same config object across -two different parsing owners. Rationale: `standalone_process` is an -*activation* choice (like `distributed`), not a leaf tuning knob on an -already-active tier (like `spdk_proxy_auto_start` is, one level below -an already-enabled SSD tier) — it belongs with the precedent for the -thing it resembles, not the deepest-nested one that happened to be -read directly by C++. - -| Var | Consumed by | Purpose | Default | -|---|---|---|---| -| `UMBP_STANDALONE_ADDRESS` | **Python** (`UMBPStore.__init__`, §5) into `cfg.standalone_process.address` | UDS path (`unix:///run/umbp/standalone/.grpc.sock`) of the standalone server's **gRPC** socket. Presence enables standalone-process mode, mirroring how `UMBP_MASTER_ADDRESS` presence enables distributed mode today. | 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, §4.2). | `false` | -| `UMBP_STANDALONE_STARTUP_TIMEOUT_MS` | **Python** into `cfg.standalone_process.startup_timeout_ms` | Bound on waiting for readiness after spawn, mirrors `spdk_proxy_startup_timeout_ms`'s role (but is itself Python-set, per the decision above — it is not read the same way its SPDK-proxy analog is). | `30000` | -| `UMBP_STANDALONE_BIN` | **Deployment-only, not a config field** — read directly by the auto-start code path in C++ at spawn time, same status as `UMBP_MASTER_BIN`/`spdk_proxy_bin` per `runtime-env-vars.md`'s existing "deployment/launcher knobs, not parsed by `UMBPConfig`" category. | 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 (server-side only; never reaches a worker's `UMBPConfig` at all, so the Precedent A/B question doesn't apply to it) | `0` = never self-exit (default; see §4 rationale). Non-zero opts back into SPDK-proxy-style idle exit, only safe if SSD tier durability is also enabled. | `0` | -| `UMBP_STANDALONE_GRPC_SHUTDOWN_DEADLINE_SEC` | `umbp_standalone_server`'s own env read (server-side only) | Reuses the `MasterServer` shutdown-deadline convention. | same default as `UMBP_GRPC_SHUTDOWN_DEADLINE_SEC` | -| `UMBP_STANDALONE_SHM_DIR` | Both sides' own env read (deployment-only, not a config field — used only for the bootstrap-lock file path, mirrors `ScopedBootstrapLock`'s convention) | Directory for the bootstrap-lock file, not the shm segment itself (which is anonymous via `memfd_create`, never named on the filesystem — see §3.2). | `/tmp` | - -**(v0.4 fix) `fd_socket` and `transport` are explicitly *not* -`UMBPStandaloneProcessConfig` fields, and there is no -`UMBP_STANDALONE_FD_SOCKET`/`UMBP_STANDALONE_TRANSPORT` env var in -v0.1** — removing the ambiguity the v0.3 draft left open: - -- The fd-handoff socket path is **always mechanically derived** from - `cfg.standalone_process.address` using the deterministic rule - already given in §4.1 (strip `unix://`, replace trailing - `.grpc.sock` with `.fd.sock`, else append `.fd.sock`) — computed - identically inside `StandaloneProcessClient`'s ctor and - `umbp_standalone_server`'s startup, from the one `address` value - they already both have. No separate config surface, no env var, and - therefore nothing that can drift out of sync between the two - processes. If a real deployment later needs a non-derived override - (e.g. the fd socket must live on a different mount than the gRPC - one), add an explicit `fd_socket_override` field to - `UMBPStandaloneProcessConfig` *then*, as a v2 change — do not - pre-add it speculatively now. -- `transport` (`uds` vs. the deferred `tcp_staging`, §8 item 6b, §10 - step 13) has no config surface in v0.1 for the same reason - `tcp_staging` itself isn't built in v0.1: there is nothing yet for a - `transport` field to select between. When `tcp_staging` is actually - implemented in a v2 revision, add `cfg.standalone_process.transport` - (Python-set, Precedent A, for consistency with everything else in - this table) at that time — not before. - -`UMBPConfig` (C++) gains: `std::optional standalone_process;` -with exactly three fields — `address`, `auto_start`, -`startup_timeout_ms` — all three Python-set per the decision above, -same shape as the existing `UMBPDistributedConfig` -(`common/config.h:212-244`), so `ResolveRole()`/`Validate()` extend -naturally rather than needing a parallel code path. - ---- - -## 7. 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**; - standalone-process mode does not touch any file under - `distributed/master/`, `distributed/peer/`, or `distributed/routing/`. -- The shared code touched: `common/config.h` (new optional field), - `umbp_client_factory.cpp` (new branch), `local/host_mem_allocator.*` - (new `kAnonymousShm` backing plus the fd registry, additive), - `local/tiers/copy_pipeline.*` (new `Drain()`, see §4.3), - `pybind_umbp.cpp` (new `AnonymousShm` enum value + new - `standalone_process`/`UMBPStandaloneProcessConfig` bindings, see §5 - — **not** change-free, correcting the v0.1 claim), and - `umbp_store.py`'s `register_mem_pool_host` gate (one conditional, - see §5 — **not** call-site-free, correcting the v0.1 claim). -- A future extension (not in scope here, flagged for later): the - standalone server could itself become a `distributed/` peer/leader - (i.e. a host runs one standalone server that both serves local - workers over UDS *and* participates in the cross-node distributed - pool) — `LocalStorageManager`'s existing `SharedSSDLeader`/ - `SharedSSDFollower` roles already hint at this kind of layering, but - wiring it up is out of scope for v0.1. - ---- - -## 8. Open questions / risks - -1. **Crash / restart semantics are unsolved by every precedent - reviewed.** Neither LMCache's server, Mooncake's Real Client, nor - UMBP's own SPDK-proxy daemon has a supervised-restart or - reconnect-with-cache-rehydration story in the code we read. If the - standalone server dies, every attached worker loses its DRAM cache - simultaneously. Minimum bar for v0.1: workers must treat a broken - RPC channel as "cache unavailable," not a fatal error (falling back - to no-cache behavior, same as any other `HiCacheStorage` backend - failure) — this needs explicit handling in `StandaloneProcessClient` - and probably a small state machine (`CONNECTED` → `DISCONNECTED` → - optionally `RECONNECTING`), which does not exist in any of the - three reference systems today. -2. **Per-tenant DRAM quota is not implemented in `LocalStorageManager` - today.** Without it, multiple workers sharing one standalone - server's DRAM tier can starve each other. SPDK-proxy's - `TenantInfo.quota_bytes` concept (`spdk_proxy_protocol.h:227-240`) - is the natural template but would need to be added to - `DRAMTier`/`LocalBlockIndex`, which currently have no notion of - "owner" per key. -3. **(decided, not open as of v0.2, keeping the entry for traceability) - fd-passing handshake requires a second, separate socket path from - the gRPC UDS** (`.grpc.sock` vs `.fd.sock`, §2 diagram, - §6) since gRPC cannot itself carry a file descriptor. The v0.1 draft - said this but its own diagram contradicted it ("over the same - UDS") — that inconsistency is fixed in §2/§6/§9. What remains - genuinely open: whether the fd-handoff listener should be a - long-lived socket the server keeps open for the whole process - lifetime (simplest, matches Mooncake's `real_client_main.cpp:104`) - or a short-lived one opened only during an active registration - window (marginally more defensible against unauthorized local - connection attempts, at the cost of a startup race the client must - retry against). Proposal: long-lived, protected by filesystem - permissions (item 4 below) — simplicity over a marginal hardening - gain, revisit if a security review disagrees. -4. **Security/isolation**: the shm segment (memfd, `O_CLOEXEC` - recommended) and the UDS socket file must both be created with - `0600`/owner-only permissions, or any other local user can attach - to another tenant's KV cache. Neither Mooncake's nor UMBP's SPDK- - proxy code reviewed explicitly sets restrictive permissions on - their sockets/shm — this needs to be added deliberately, not copied. -5. **gRPC-over-UDS batch overhead vs. SPDK-proxy's ring buffer**: even - with struct-of-arrays batching, a gRPC call still costs a syscall + - protobuf encode/decode per `BatchGet`/`BatchPut`, which a busy-polled - shm ring buffer avoids entirely for the hottest path. If profiling - later shows this is a bottleneck for very small, very frequent Get - calls (e.g. single-block prefix-cache lookups), a v2 iteration could - add an optional shm-ring **fast path** for single-key Get/Put - alongside the gRPC path for everything else — deferred, not part of - v0.1, to avoid combining two IPC mechanisms before the simpler one - is proven to be insufficient. -6a. **(new in v0.2) The process-local ptr→fd registry (§3.2 step 2) - is new state with its own lifetime/thread-safety questions no - existing code needs to answer today.** `HostMemAllocator::Alloc`/ - `Free` (`host_mem_allocator.h:41-66`) are currently stateless - beyond the returned handle — adding a shared registry means: (a) it - must be safe for `Free` to run concurrently with a - `RegisterMemory` lookup from a different thread (mutex, as noted in - §3.2), (b) `DeregisterMemory` and `Free` ordering must be defined — - if a caller frees the underlying buffer before calling - `DeregisterMemory`, the registry entry and the server's `mmap` - would dangle; the design should make `Free` refuse (or at least - warn loudly) if an active registration still references that - pointer, rather than silently leaving the server with a stale - mapping. -6b. **(new in v0.2) The `tcp_staging` transport (§3.1, §6) is a real - second code path, not a config toggle on the same code.** Once - `Put`/`Get` payload bytes can travel two different ways (shm-offset - reference vs. inline bytes in the RPC message), `UMBPStandalone`'s - proto and `StandaloneProcessClient`'s implementation both need an - explicit branch, and testing needs to cover both — this roughly - doubles the data-plane test matrix for standalone-process mode - versus what v0.1 implied. Recommendation: **do not build - `tcp_staging` in v0.1** unless a concrete deployment already needs - it; ship UDS-only first and treat the inline-payload fallback as a - v2 addition once there's a real caller for it, rather than building - two data planes up front for a use case that may never materialize. -7. **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. This is the same - constraint `distributed/` already has between `umbp_master` and - worker processes, not a new problem — but standalone-process mode - is more likely to be deployed as a long-lived sidecar that outlives - several SGLang worker restarts/upgrades, which makes schema drift - more likely in practice than it is for `umbp_master` today. No - solution proposed here beyond "same binary/version pinning - discipline as the existing distributed mode." - ---- - -## 9. Source references used in this design - -All claims above trace to source read directly in this repo checkout, -not to general knowledge: - -- UMBP: `mori/src/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` - (`SpawnProxyDaemon`/`EnsureProxyDaemon`), - `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`, - `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`, local checkout): - `mooncake-store/src/real_client.cpp`, `real_client_main.cpp`, - `shm_helper.cpp`, `dummy_client.cpp`, `client_buffer.cpp`, - `docs/source/design/mooncake-store.md`, - `docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md`. -- LMCache (`/apps/nima/KVManager/LMCache`, local checkout, reviewed - for comparison only): `lmcache/v1/multiprocess/{mq,server,protocol, - custom_types,futures,config}.py`, `lmcache/v1/mp_observability/`. - -## 10. Suggested implementation order (for step 3, pending approval) - -**(Reordered/expanded across v0.2 and v0.3 to make each review-found -gap a concrete, individually-checkable task rather than an implicit -sub-step.)** - -1. `common/config.h`: add `UMBPStandaloneProcessConfig` + - `UMBPConfig::standalone_process`; extend `Validate()`/`ResolveRole()`. -2. `local/host_mem_allocator.{h,cpp}`: add `kAnonymousShm` backing - (`memfd_create`-based) **plus** the process-local ptr→fd registry - (§3.2 step 2, §8 item 6a) that `RegisterMemory` will query — these - belong in the same change since the registry is populated at - `Alloc`/`Free` time. -3. `local/tiers/copy_pipeline.{h,cpp}`: add `Drain(timeout)` (§4.3, - §8 item 7 fix) — small, independent of everything else, can land - first. -4. `src/pybind/pybind_umbp.cpp`: add `AnonymousShm` to - `UMBPHostBufferBacking`; bind `UMBPStandaloneProcessConfig` and - `UMBPConfig::standalone_process`; add `get_deployment_mode()` to - the `IUMBPClient` binding (§5). Do this **before** step 6 below — - the Python-side change depends on these bindings existing. -5. New proto `distributed/proto/umbp_standalone.proto` + - generated-code wiring in `CMakeLists.txt`; struct-of-arrays batch - codec reuse for `BatchPut`/`BatchGet` (§3.1). -6. New `umbp_standalone_server` binary (own `main.cpp`, thin wrapper - around a new `StandaloneServer` class that owns - `LocalStorageManager`/`LocalBlockIndex`/`CopyPipeline` — reused, not - reimplemented, with `DRAMTier` defaulting to private - anon/hugetlb per §3.2/§7, not `use_shared_memory=true`) — plus the - gRPC service on `.grpc.sock` and the **separate** raw-UDS - fd-handoff listener on `.fd.sock` (§2, §8 item 3). -7. New `StandaloneProcessClient : IUMBPClient` (worker-side): gRPC - stub, offset translation, the `SCM_RIGHTS` send side of the - fd-handoff handshake, wired into `umbp_client_factory.cpp`. - `RegisterMemory` must fail loudly (not silently degrade) when the - registry lookup misses (§3.2 step 2). -8. **(v0.3 addition, corrected in v0.5 — this step did not exist - before and is the activation path itself.)** `umbp_store.py`: in - `UMBPStore.__init__`, next to the existing `master_address` block - (`umbp_store.py:458-465`): **first**, raise if - `extra_config["standalone_address"]` is set without - `UMBP_STANDALONE_ADDRESS` also being set (v0.5 — extra_config alone - cannot activate this mode, see §5); **then** read - `UMBP_STANDALONE_ADDRESS`, construct `UMBPStandaloneProcessConfig`, - assign `cfg.standalone_process`, and raise if both `master_address` - and the standalone address are set (§5). Depends on step 4 (pybind - bindings for `UMBPStandaloneProcessConfig` must exist first). - Without this step, every other step in this list can be implemented - correctly and `UMBPStore` will still silently construct a plain - `StandaloneClient` forever — this is not optional polish, it is the - switch that turns the feature on. -9. **(v0.5 fix — must read the env var directly, at allocator - construction time, before `extra_config` exists — see §5.)** - `umbp_host_allocator.py`: `UMBPHostTensorAllocator` reads - `UMBP_STANDALONE_ADDRESS` directly from `os.environ` (not via - `UMBPConfig`/`extra_config`) and branches to request `AnonymousShm` - backing when set (depends on step 4 for the pybind enum value). - This step and step 8 both gate on the *same* env var by - construction, precisely so they can never disagree about whether - standalone-process mode is active. -10. **`umbp_store.py`: fix the `register_mem_pool_host` gate at line - 880** from `is_distributed()`-only to also accept - `get_deployment_mode() == StandaloneProcess` (§5) — called out as - its own task because it is a functional bug an implementation - could ship with even after step 8 above: the client would be - correctly constructed as `StandaloneProcessClient`, but still - never register its host buffer, so every `Put`/`Get` would fail at - the server's offset-translation step with no obvious cause. -10a. **(v0.5 addition.)** `umbp_store.py`: in the same method, make - `disable_zero_copy_register`-set / `register_memory` exception / - `register_memory() == False` all `raise` instead of - warn-and-return when `get_deployment_mode() == StandaloneProcess` - (§5); leave `distributed`'s existing warn-and-fallback behavior - untouched. Land this together with step 10 — both touch the same - function and the same review found both gaps together. -11. Lifecycle: auto-start path reusing `SpawnProxyDaemon`'s fork/exec - pattern; shell script analogous to `run_umbp_single_node_hicache.sh` - for the externally-launched path. -12. Tests: unit tests for the fd-handoff handshake and offset - translation, the ptr→fd registry's `Free`/`DeregisterMemory` - ordering including a register→deregister→re-register cycle (§3.2 - step 4, §8 item 6a), `CopyPipeline::Drain()` under load; an - integration test with 2+ worker processes sharing one standalone - server on one host, including a kill-the-server test that asserts - workers degrade to cache-unavailable rather than crashing (§8 - item 1), and an end-to-end test that `UMBP_STANDALONE_ADDRESS` - alone (via `UMBPStore` config, no other code changes) is sufficient - to activate the mode (guards against step 8 regressing silently). - **(v0.5 additions)** a Python test asserting - `extra_config["standalone_address"]` without the env var raises at - `UMBPStore.__init__` (guards step 8's rejection path); and Python - tests for each of the three `register_mem_pool_host` cases (§5, - step 10a) asserting they raise under `StandaloneProcess` and still - only warn under `Distributed` (guards against the two modes' - branches getting merged back together in a future edit). -13. **Deferred to v2, not v0.1** (§8 item 6b): `tcp_staging` transport - with inline RPC payload bytes. Build only if a concrete deployment - needs it. - ---- - -## 11. Cross-Node Extension: making `umbp_standalone_server` optionally `DistributedClient`-backed (v0.7 proposal, UNAPPROVED) - -**v0.7 revision note:** a second review pass found 5 concrete gaps in -the v0.6 draft of this section — a config conflict with the already- -shipped v0.5 activation rule, an external-KV routing-correctness hole, -missing transaction/rollback semantics for backend memory -registration, an open correctness question left open instead of -closed, and an overstated "no logic change" claim. All five are -addressed below with inline `(v0.7 fix)` markers. **Guiding principle -adopted for this revision and going forward: where a design question -has a direct Mooncake precedent, resolve it by matching Mooncake, -rather than inventing a new answer** — per explicit direction ("我们与 -Mooncake 保持对标就好,类似的问题都一个思路"). Two of the five fixes -below were resolved this way, with the precedent re-verified against -the local Mooncake checkout (not recalled from memory): - -- **No *implicit* auto-start anywhere in Mooncake's application-side - binding code.** **(v0.7 citation fix, re-checked a second time - against the correct checkout `/apps/nima/KVManager/Mooncake` after - discovering the first verification pass had used a stale checkout at - a different path — the precise claim below is narrower than what was - first written, because the first, broader "zero matches across both - directories" claim turned out to be literally false.)** Re-checked - directly: `grep -rln "fork(|execvp|execve|posix_spawn|subprocess\.|Popen" mooncake-integration/` - returns zero matches — the `DummyClient` application-binding code - never spawns anything. `mooncake-wheel/` (the packaging/CLI layer, - not the application-binding layer) does contain one real match, - `mooncake-wheel/mooncake/cli_client.py:24` - (`subprocess.call([bin_path] + sys.argv[1:])`), but this is the - pass-through implementation of the explicitly-invoked `mooncake_client` - console-script command — functionally equivalent to a human running - the `mooncake_client` binary directly from a shell, not an implicit - auto-start triggered by normal `DummyClient`/integration-API usage. - (The other `mooncake-wheel/` hits — `cli.py`, `cli_bench.py`, - `setup.py`, `tests/test_cli.py` — wrap unrelated binaries, are - build-time only, or are an explicit human-run test harness, - respectively.) The Real Client is always started by something - outside the application (deployment system, operator, systemd unit, - or an explicit CLI invocation) — never implicitly, by any code path - a normal application integration would exercise. §11.4.6 below - adopts this exactly. -- **Mooncake's Real Client already does the "mmap the app's shm, then - register that same mapping with the distributed transport" pattern - this section proposes** **(v0.7 citation fix — line numbers below - corrected against the correct checkout; the function has been - renamed/refactored since the stale citation was written, but the - mechanism is unchanged)** — `real_client.cpp:2136-2137` - (the `mmap(nullptr, shm_size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0)` - call) and `real_client.cpp:2161-2162` - (`client_->RegisterLocalMemory(shm.shm_buffer, shm_size, ...)`, - munmap'd on failure), both inside - `RealClient::map_shm_internal_with_device` - (`real_client.cpp:2097-2186`), which the now-thin - `RealClient::map_shm_internal` (`real_client.cpp:2089-2095`) - delegates to. This is production code, not - a prototype — it directly validates §11.3's core mechanism (register - the *server's* local view of a shared mapping with the distributed - transport) as an already-shipped pattern, not a novel one this - design is the first to attempt. This substantially de-risks the - "cross-mapping visibility" open item in §11.3 (see the fix there). - -**Scope reframing (v0.7, requested explicitly):** §1's original prose -("Standalone-process mode is strictly same-host") describes the -history of what v0.1-v0.5 built, not the target end-state. Given the -project's actual goal is Mooncake parity, and Mooncake's "standalone" -Real Client is cross-node-capable by construction, this section's -distributed-backed shape — not the local-only shape §1 describes — is -the one that actually satisfies that goal. Going forward: -**`standalone-process/distributed-backend` (cross-node, Mooncake-parity -shape) is the primary target this feature is building toward; -`standalone-process/local-backend` (§1-§10, already shipped) is a -same-host fallback/dev-convenience shape, not a co-equal alternative.** -§1's prose is not rewritten in place (it remains an accurate -description of what v0.1-v0.5 actually built and why, and that -reasoning was sound for the narrower problem it solved), but it should -no longer be read as this feature's final scope statement — this -paragraph supersedes it for that purpose. See §11.7 item 1 for the -naming-convention decision that follows from this. - -### 11.0 Why this section exists - -After v0.1-v0.5 shipped, we compared what we built against what -Mooncake's Real Client actually does (source-verified earlier in this -project, see the LMCache/Mooncake research this design already cites) -and found a real divergence, not a rounding error: - -- **Mooncake's Real Client**: a full distributed `Client` (Master + - etcd + Transfer Engine/RDMA), just hosted in its own OS process. - **(v0.7 citation fix — corrected against the correct checkout; - construction pattern unchanged, only line numbers moved, and a - second, structurally identical call site was added since the stale - citation was written)** `real_client.cpp:672-675` (and duplicated at - `real_client.cpp:708-711`, an auto-port-binding-with-retry variant of - the same construction) constructs it via the exact same - `Client::Create(..., master_server_addr, transfer_engine, - {"client_mode":"real"})` an embedded client would use — "standalone" - only changes which process the library lives in, not which protocol - it speaks to the rest of the cluster. A cache miss can still be - served from a remote node over RDMA. -- **What v0.1-v0.5 built**: `umbp_standalone_server` wraps a concrete - `StandaloneClient` (`standalone_server.cpp:587`, - `client_(config)` at `standalone_server.cpp:125`), i.e. `local/`'s - `LocalStorageManager`/`LocalBlockIndex`/`CopyPipeline` moved into its - own process. It has **no** connection to `distributed/` — no master, - no RDMA, no peer protocol. A cache miss is just a miss; there is no - "somewhere else" to check. - -The UDS + `SCM_RIGHTS`-fd-handoff IPC layer between the worker and the -server (§2, §3) is *not* the divergence — that part already correctly -mirrors Mooncake's own app-process↔Real-Client IPC. The divergence is -narrower and specific: **what backs the server side of that IPC**, -local storage vs. a full distributed client. - -### 11.1 Proposal, in one sentence - -Make the concrete `StandaloneClient client_;` member in -`StandaloneServer::Impl` into an `IUMBPClient`-typed backend built -through the **already-existing** `CreateUMBPClient(config)` factory -(`umbp_client_factory.cpp:29-37`), and give `standalone_server_main.cpp` -a way to populate that config's `distributed` field (today it -unconditionally clears it — `standalone_server_main.cpp:43`, -`config.distributed.reset()`). Same server binary, same worker-facing -IPC protocol, same `umbp_standalone_server` deployment story — the -only thing that changes is which `IUMBPClient` implementation answers -Put/Get/etc, decided entirely by how the server operator starts it. - -This is deliberately **not** a rewrite: §11.5 below shows the existing -data-path RPC handlers (Put/Get/BatchPut/BatchGet/Exists/Clear/Flush) -need no logic changes at all, because they already resolve -`(client_id, shm_offset)` to a plain pointer before ever touching -`client_` (`standalone_server.cpp:541-568`) — `client_`'s concrete -type was never load-bearing for those handlers in the first place. - -### 11.2 Two independent `UMBPConfig` instances — this is the key structural point - -There are, and always were, two separate configuration surfaces in -play, and they must stay separate: + `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. **Unchanged by this proposal** — 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, new): - what `umbp_standalone_server` itself uses to decide whether *its* - `client_` backend should be a `StandaloneClient` (today's behavior) - or a `DistributedClient` (this proposal). This is a **second, - independent `UMBPConfig` instance** — never the same object as #1. - -**Verified this is not blocked by the existing mutual-exclusivity -check.** `UMBPConfig::Validate()` (`config.h:345-349`) rejects -`distributed.has_value() && standalone_process.has_value()` *within -one config instance* — it says nothing about two different `UMBPConfig` -objects existing in the same process. `standalone_server_main.cpp` -already builds one `UMBPConfig` for the worker-facing side (with only -`standalone_process` set); this proposal adds a second, separate -`UMBPConfig` for the server's own `client_` backend (with only -`distributed` set, when requested). Neither instance ever has both -fields set, so `Validate()`'s existing rule is satisfied unchanged — -no change needed to that check. + 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_RIGHTS─│ worker's, address/fd-socket) │ - │ {address} │ │ -> used only to open sockets │ - │ distributed = null │ │ │ - └───────────────────────┘ │ UMBPConfig #2 (server's OWN, │ - │ new in this proposal) │ - │ 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 │ - │ (today's local-only │ - │ behavior, unchanged) │ - └──────────────────────────────────┘ - │ RDMA (if DistributedClient) - ▼ - ┌──────────────────────────────────┐ - │ umbp_master + other nodes' │ - │ peers (the existing distributed/│ - │ cluster, entirely unmodified) │ - └──────────────────────────────────┘ + 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) │ + └───────────────────────────────────┘ ``` -### 11.3 Why the zero-copy story still holds with an RDMA backend - -This was the load-bearing technical question and it checks out against -the actual `PoolClient`/`DistributedClient` code (not assumed): - -- RDMA memory registration (`DistributedClient::RegisterMemory` → - `PoolClient::RegisterMemory`, `pool_client.cpp:635-651`) operates - purely on `(ptr, size)` **in the calling process's own virtual - address space** — it calls - `io_engine_->RegisterMemory(ptr, size, -1, MemoryLocationType::CPU)` - (`pool_client.cpp:648`) and caches `{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** - (`pool_client.cpp:635-675`, no ownership check anywhere). -- This means: `umbp_standalone_server`, after `mmap`-ing the worker's - `memfd`-backed shared region via the existing fd-handoff path (§3.2, - already built and working), 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 — full stop, regardless of which process's VA was used to - register them. -- **Get() zero-copy is conditional on this registration, and confirmed - in the actual RDMA read path**: `WaitRemoteBatchGet` - (`pool_client.cpp:1812-1853`) only skips the memcpy-out when the - destination pointer was pre-registered and `FindRegisteredMemory` - matches it (`pool_client.cpp:1661-1664`) — in that branch, the RDMA - read from the remote peer is posted directly into that registered - MR (`pool_client.cpp:1806-1808`) with **no memcpy** afterward. If - the pointer wasn't registered, RDMA lands in a shared staging buffer - and is then `memcpy`'d to the destination (`pool_client.cpp:1847`) — - the ordinary, already-existing non-zero-copy path. -- **Net effect**: register the server's local mapping once (at - fd-handoff time, alongside the existing mmap), and every subsequent - `Get` that targets an offset inside that region gets 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. This is *better* than Mooncake's shape in one - respect: the worker process never touches RDMA/IB verbs directly, - only the server does, entirely consistent with the "worker only - knows about UDS + shm" boundary the v0.1-v0.5 design already - established. - -**(v0.7 fix — closed, was left open in v0.6.)** Cross-mapping memory -visibility/ordering: does an RDMA-completed write into the physical -pages via the server's VA become visible to a subsequent read through -the worker's *independent* VA mapping of the same `memfd`? v0.6 left -this as an unresolved research question. It is now closed by two -independent legs, not a single hand-wave: - -1. **Mooncake precedent, re-verified against the correct checkout this - pass (an earlier pass had used a stale checkout and, separately, - mis-described the allocation direction — both corrected here).** - Mooncake's Real Client does the *identical* dual-mapping shape in - shipping production code, with this exact allocation/registration - direction: the **application** (`DummyClient`) is the one that - originally allocates the shared region, via - `ShmHelper::allocate()` (`dummy_client.cpp:469`, inside - `setup_dummy()`) — that allocation call itself performs the first - `mmap` (`shm_helper.cpp:123-124`), giving `DummyClient` its own - independent VA mapping of the physical pages from the moment of - allocation, not by "receiving and mapping a fd" the way the server - does. `DummyClient` then sends that fd to the Real Client - (`sendFd`, `dummy_client.cpp:408`). The **Real Client** receives it - (`recvFd` via `handle_ipc_shm_register`, `real_client.cpp:5311`) - and performs the *second*, independent `mmap` of the same physical - pages into its own address space (`real_client.cpp:2136-2137`), - then registers that server-local mapping with its own Transfer - Engine (`real_client.cpp:2161-2162`). This is the same shape - `umbp_standalone_server` already implements today (worker allocates - and owns the buffer, §3.2 step 1; server receives the fd and - independently mmaps it, §3.2 step 2) — UMBP's already-shipped - direction matches Mooncake's actual direction exactly; only this - design doc's *prose describing* Mooncake's mechanism had briefly - drifted from that (now corrected). Two independent processes, two - independent VAs, one physical region, the *receiving* side's VA - registered with a distributed transport, exactly this design's - shape — running in production today. This is strong evidence the - pattern is sound, not merely "should be fine" reasoning from first - principles. -2. **Explicit, adopted assumption for the implementation (decision, - not left implicit):** the gRPC `Get` response is defined as the - happens-before edge. `umbp_standalone_server` must only send the - gRPC response for a `Get` **after** the RDMA read completion is - observed server-side (already the natural control flow — there is - no reason to respond before the transfer completes). The worker - must not read the target offset until it has received that - response. This is not a new constraint on the implementation beyond - what it would already naturally do; it is being stated explicitly - here so it is a documented contract, not an accidental property of - the current code shape. - -Residual, non-blocking follow-up: add one integration test that -specifically exercises this path (RDMA `Get` into a `DistributedClient`- -backed server's mapped region, immediately followed by the worker -reading the result through its own independent mapping) once the -distributed-backend test infrastructure from §11.7 item 5 exists — this -is a good regression guard to have, not a prerequisite for the design -being considered closed. - -### 11.4 What changes, concretely, mapped against the current code - -Grounded in the actual current implementation (all citations verified -against source in this pass, not the earlier v0.1-v0.5 design intent): - -1. **`StandaloneServer::Impl`** (`standalone_server.h/.cpp`): - `StandaloneClient client_;` → `std::unique_ptr client_;`, - constructed via `CreateUMBPClient(backend_config)` instead of - `client_(config)` directly. `backend_config` is `UMBPConfig #2` - from §11.2, built by `standalone_server_main.cpp` (see #6 below). -2. **Put/Get/BatchPut/BatchPutWithDepth/BatchGet/Exists/BatchExists/ - BatchExistsConsecutive/Clear/Flush RPC handlers**: mechanically, - swapping `client_` from a value to a `unique_ptr` is a - `client_.` → `client_->` change and nothing more — they already - call `client_.Put(...)` etc. on plain resolved pointers - (`standalone_server.cpp:199-354`). **(v0.7 fix — "no logic change" - overstated in v0.6; three things must be documented even though no - code changes are required for them):** - - **Registration gating is already correct, verified this pass, no - new work needed**: `ResolveRange` (`standalone_server.cpp:541-549`) - already fails (`return false`) if `client_id` isn't present in - the `memory_` map, so `Put`/`Get`/etc. already cannot reach - `client_` at all before `RegisterMemory` has succeeded, for - either backend. The reviewer's concern here is already satisfied - by existing code — noted so a future reader doesn't think it's a - gap. - - **Failure-policy ambiguity (backend-unavailable vs. genuine - cache-miss both surface as `false`) is a pre-existing - characteristic of `IUMBPClient`, not something this proposal - introduces or worsens.** `DistributedClient::Get` already returns - `false` uniformly for "not found" and "RPC/transport failure" - today, for ordinary distributed-mode workers, independent of - whether it runs inside a standalone server. This proposal - inherits that behavior unchanged; fixing it (if ever needed) is - an `IUMBPClient` interface question orthogonal to this section - and explicitly out of scope here. - - **`Flush()`/`Clear()` mean different things per backend, and this - is an accepted semantic difference, not a bug to fix**: under - `StandaloneClient`, `Flush()` drains `CopyPipeline` to SSD - (§4.3); under `DistributedClient`, `Flush()`/`Clear()` operate on - the pool client's own state (heartbeat/registration bookkeeping, - not a local SSD tier — there is no local SSD tier in this backend - per §11.6). Both are the correct, existing behavior of the - respective `IUMBPClient` implementation for their backend; the - server passing the call through unchanged is correct, and this - should be documented as an accepted, backend-dependent semantic - rather than left for a future reader to wonder whether it's an - oversight. -3. **`RegisterMemory`/`DeregisterMemory` RPC handlers** - (`standalone_server.cpp:356-375`): **currently never touch `client_` - at all** — pure fd/mmap bookkeeping against the `memory_` map. Must - add, after a successful `mmap`: - `client_->RegisterMemory(server_local_ptr, size)`, and on - deregistration: `client_->DeregisterMemory(server_local_ptr)`. This - call is **safe to make unconditionally regardless of backend** — - `IUMBPClient::RegisterMemory`'s default (used by `StandaloneClient`, - which never overrides it) is already a no-op returning `true` - (`umbp_client.h`, documented as "Standalone mode needs no - registration (CPU-local memcpy)"). So no `IsDistributed()` branch - is needed here — the interface was already designed for exactly - this kind of backend-agnostic call. - - **(v0.7 fix — v0.6 said "add the call" but didn't specify the - transaction/rollback semantics around it; this is a genuine - correctness gap, not an implementation detail to leave for later.)** - - - **Rollback on backend registration failure.** `RegisterFd` - (`standalone_server.cpp:504-523`, confirmed this pass) is a - single-phase commit: on a successful `mmap`, it immediately writes - the entry into `memory_` and returns success — there is no - existing "pending" intermediate state to roll back from. The new - `client_->RegisterMemory(server_local_ptr, size)` call must be - placed **inside `RegisterFd`, immediately after `mmap` succeeds - and before the entry is written into `memory_`**. If it fails, - `RegisterFd` must `munmap` the just-created mapping and return - failure **without** writing to `memory_` — i.e. backend - registration failure is treated exactly like an `mmap` failure - already is, using the exact same single-phase-commit shape the - function already has. This requires no new state machine, only - ordering the existing calls correctly. - - **Deregistration ordering**: `DeregisterMemory` - (`standalone_server.cpp:370-375`) currently only calls - `UnmapClient()` (munmap + erase). Must call - `client_->DeregisterMemory(server_local_ptr)` **before** - `munmap`-ing (deregister the transport's view of the memory before - the memory itself disappears — the reverse order risks the - backend still holding a registered MR pointing at unmapped - memory, however briefly). - - **Shutdown ordering — `UnmapAll()` must deregister from the - backend before unmapping, and this must happen before - `client_->Close()`.** `UnmapAll()` (`standalone_server.cpp:533-538`) - currently only `munmap`s every entry in `memory_`. Once `client_` - can be a `DistributedClient` holding live RDMA registrations, this - must become: for each registered mapping, call - `client_->DeregisterMemory(base)`, *then* `munmap`. And this whole - sequence must run **before** `client_->Close()` in `Shutdown()` - (`standalone_server.cpp:154-169`, the ordering already fixed in - the v0.1-v0.5 review pass — see the earlier "shutdown ordering" - fix in this document's history) — deregistering after `Close()` - would call into a backend that has already torn down its RDMA/IO - engine state. - - **Explicitly out of scope for this proposal, tracked separately - (not a new problem introduced here):** if a worker sends its fd - via the raw fd-handoff socket and then crashes before ever issuing - the confirming gRPC `RegisterMemory` RPC, the mapping stays live in - `memory_` indefinitely (no timeout/reaper exists for this today). - This is a pre-existing characteristic of the v0.1-v0.5 fd-handoff - protocol, orthogonal to which backend `client_` is — it existed - before this proposal and isn't made worse by it. Worth a future - hardening pass (e.g. a liveness check tied to the worker's gRPC - channel), but not a blocker for this section. -4. **The five external-KV RPC handlers** - (`ReportExternalKvBlocks`/`RevokeExternalKvBlocks`/ - `RevokeAllExternalKvBlocksAtTier`/`MatchExternalKv`/ - `GetExternalKvHitCounts`, `standalone_server.cpp:377-408`): - currently hardcoded stubs (three always return `true`, two return - empty responses) that never touch `client_` at all. - - **(v1.0 fix — superseding both the v0.7 "disable entirely" decision - and the intermediate "naive single shared identity" idea; this is - the third and final revision of this item, corrected after two - rounds of clarification about what external-KV actually is and what - the baseline non-standalone implementation actually already - guarantees.)** - - **What external-KV actually is, clarified:** it 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 — it sends a NEW - request to whichever node the match points at, and that node's own - already-resident local cache (GPU HBM or otherwise) serves it - locally. **UMBP itself never moves the reported bytes for this - feature, in `distributed/` mode or here** — so the earlier v0.7 - concern "the server can't physically reach the worker's HBM memory" - does not apply: nothing ever needs to read that memory through - UMBP for this feature. That concern was based on treating - external-KV as a data-plane feature; it is not one. - - **What genuinely matters instead: does the reported/matched - *identity* (`node_id`) correctly reflect the actual deployment - topology, at the granularity a caller needs?** This is where the - real design work is, and it is grounded in what the *baseline*, - non-standalone implementation already does — per this project's - standing principle of matching existing capability rather than - inventing a narrower one: - - - **Verified against `umbp_store.py:512-519`**: a normal distributed - worker's `node_id` is built as - `f"{node_address}:dp{dp_rank}:pp{pp_rank}:tp{local_rank}"` — pure - per-process rank coordinates, with **zero awareness of which - logical TP/DP replica group the process belongs to**. This means - the baseline already, for free, supports multiple independent - replicas sharing one physical host (e.g. two TP=4 groups on one - 8-GPU box) — each of the 8 processes gets its own distinct - `node_id` regardless of logical grouping, because disambiguation - is by rank coordinate, not by group membership. This is a real, - already-existing capability of the system this proposal must not - silently drop. - - **This project's stated deployment shape is 8-GPU hosts running - SGLang TP=8 (one replica, 8 ranks, per host).** For that specific - shape, a single shared identity per host would actually have been - *semantically adequate* on its own (all 8 ranks in one TP group - process every request in lock-step, so "this host has hash H - cached" is a node-level fact, not a per-rank one) — but adopting a - single-shared-identity design would silently regress the - multi-replica-per-host case the baseline already supports, for no - reason other than convenience. Per explicit instruction, this - proposal should support that case too rather than quietly drop it. - - **Decision: the server maintains one independent distributed - sub-identity per connected worker (`client_id`), not one shared - identity for the whole host.** - - **(v1.2 fix — this used to say "N independent `MasterClient`-level - registrations." That was wrong in a way that would have broken - ordinary distributed routing, not just been imprecise. Corrected to - a purpose-built `ExternalKvIdentityClient` instead, per three - blockers found on review, all confirmed against source, not - speculative.)** - - **Blocker A — a naive `MasterClient`-based sub-identity would pollute - normal Put/Get routing.** Once `RegisterClient` succeeds, that - identity becomes eligible for `ClientRegistry::GetAliveClients()` - (`client_registry.cpp:50`), which ordinary `RoutePut` placement - selects targets from - (`route_put_strategy.cpp`'s `CollectEligibleOnTier`, iterating - *all* alive clients and matching on `tier_capacities`). If a - sub-identity reported any nonzero `tier_capacities`, real KV blocks - from unrelated Puts could get routed onto it — a node that isn't a - real storage participant at all. **Verified the fix works with the - existing code, no routing-code change needed**: `CollectEligibleOnTier` - already skips any candidate with no entry for the tier being routed - (`if (it == client.tier_capacities.end()) continue;`), and - `MasterServer::RegisterClient` does not reject empty/absent - `tier_capacities` (`master_server.cpp:218-240`) — so a sub-identity - registered with **empty `tier_capacities`** is already, structurally, - invisible to `RoutePut`/`RouteGet` selection, with zero changes to - the routing algorithm itself. - - **(v1.2 addition — the same invariant also covers `EvictionManager`, - checked directly rather than assumed to be covered by the RoutePut - fix alone.)** `EvictionManager::RunOnce` (`eviction_manager.cpp:83-91`) - also calls `registry_.GetAliveClients()` and iterates each client's - `tier_capacities` to detect overloaded node-tiers - (`for (const auto& [tier, cap] : client.tier_capacities) { ... }`); - a client with an empty `tier_capacities` map contributes zero - iterations and can never be flagged as overloaded or selected as an - eviction target — the same empty-capacities invariant that hides a - sub-identity from `RoutePut` also makes it structurally invisible to - `EvictionManager`, with no separate mechanism needed. **This must - hold by construction, not by convention that a future edit could - quietly break** (e.g. someone "helpfully" wiring up real capacity - reporting on `ExternalKvIdentityClient` to make monitoring output - look more complete) — which is exactly why Blocker B's fix is a - dedicated class with no `tier_capacities`-reporting code path at - all, rather than a `MasterClient` configured to behave this way by - convention. **Additionally decided: `ExternalKvIdentityClient`'s - heartbeat must never publish `EventBundle`s (owned-KV-block delta/ - full-sync events) — it has no owned index to report in the first - place, since it is not a real storage participant; this is stated - explicitly here, not left to be inferred from "empty capacities," - because the two are different fields in `HeartbeatRequest` - (`tier_capacities` vs. `bundles`, `distributed/proto/umbp.proto:103-109`) - and both must independently stay empty.** - - **Blocker B — reusing `MasterClient` directly is the wrong vehicle, - for two compounding reasons, so this is now a dedicated class:** - `MasterClient` is a full-featured class built for real distributed - peers — capacity heartbeats, KV-event-bundle sequencing, metrics - reporting, and (separately) a real, **confirmed** bug in its - re-register path: on receiving `CLIENT_STATUS_UNKNOWN`, the - re-register request sets only `node_id`/`node_address`/ - `tier_capacities`/`tags` (`master_client.cpp:~612`) and **omits - `peer_address`/`engine_desc`** — unlike the initial `RegisterSelf` - (`master_client.cpp:137-141`), which sets both. Since - `ClientRegistry::RegisterClient` unconditionally overwrites - `record.peer_address`/`record.engine_desc_bytes` with whatever the - request carried (empty, in the re-register case), any identity that - re-registers after Master forgets it (restart, reaper expiry) loses - its `peer_address` — and `MatchExternalKv`'s response, sourced from - `ClientRegistry::GetAlivePeerView()` (`master_server.cpp:590`), - would then return an empty `peer_address` for that node. This is a - pre-existing bug in `MasterClient` itself (it would affect real - distributed peers too, not just this proposal), but relying on it - un-fixed for a *new* use case is not acceptable, and patching the - shared `MasterClient` class carries real regression risk across - every existing distributed peer for the sake of a narrow new - feature. - - **Decision: introduce a new, purpose-built `ExternalKvIdentityClient` - class instead of instantiating N `MasterClient`s.** It implements - only `Register`/`Heartbeat`/`Unregister` plus the five external-KV - RPCs — no capacity heartbeats, no KV-event-bundle logic, no - `PeerDramAllocator`/`PeerSsdManager`/`PeerServiceServer` coupling. - This resolves both blockers structurally rather than by convention: - - **Blocker A is closed by construction, not configuration**: the - class has no `tier_capacities`-reporting code path at all, so - there is no field a future edit could accidentally populate with - real capacity — contrast with "reuse `MasterClient` but always - pass empty capacities," which depends on every future call site - remembering to keep passing empty. - - **Blocker B is closed by not repeating the bug**: since this is - new code, its `Register`/re-register path always includes - `peer_address`/`engine_desc` on every call, from the start — no - fix to shared `MasterClient` code required, no regression surface - on real distributed peers. - - Heartbeat's only job for this class is liveness (keep the - `ClientRegistry` entry from expiring) — it must **not** publish - owned-KV events or capacity snapshots, since it represents no real - storage. - - All N instances (N ≤ 8 for the stated deployment shape) **share - one common `peer_address`** (the server's own single physical - `PeerService` endpoint) — **verified this pass, not assumed**: - `peer_address` uniqueness is not required anywhere in the - registry/dispatch path — `ClientRegistry` stores a plain - `node_id → peer_address` map (`client_registry.cpp:217`), - `MasterServer::GetOrCreateStub` keys its gRPC stub cache by - `node_id` (`master_server.cpp:158-166`, tolerating multiple - `node_id`s pointing at the same `peer_address`), and the - `MatchExternalKv` response path returns `(node_id, peer_address)` - pairs per match (`master_server.cpp:595`) — nothing here assumes - or requires a 1:1 `node_id`↔`peer_address` mapping. - - Each of the five external-KV RPC handlers, on receiving a call - from a given `client_id`, dispatches it to *that worker's* - `ExternalKvIdentityClient` instance (`client_id → - ExternalKvIdentityClient` lookup, mirroring the existing - `client_id → {base, size}` lookup already used for offset - resolution, §3.2) rather than to any single, server-wide identity. - - The bulk Put/Get/RegisterMemory RDMA data path is **unaffected** - and continues to share the server's one real `DistributedClient` - backend (built via `CreateUMBPClient`, per §11.1/§11.4.1) for - efficiency — per §11.3/§11.4.3, RDMA registration is inherently - per-VA and identity-agnostic, so there is no reason to split it - per worker; only the identity-bound external-KV surface needs - per-worker handling, and now has a class whose only job is that. - - **Blocker C — the registration-handshake field list, decided now, - not left to implementation time.** `RegisterClientRequest` requires - both `node_id` **and** `node_address` - (`distributed/proto/umbp.proto:54-60`) — a v1.1 draft of this - section under-specified this as "just `node_id`." **Decided:** the - worker conveys `worker_node_id`, `worker_node_address`, and - optionally `tags` — **not raw rank components** (`dp_rank`/`pp_rank`/ - `tp_rank`). Python (`umbp_store.py`'s existing baseline node_id/ - node_address derivation logic, `umbp_store.py:496-519`) remains the - single source of truth for how these strings are computed; the - server never re-derives rank semantics itself. This avoids two - independent, potentially-diverging implementations of the same - derivation rule ever existing at once. - - **(v1.2 — wire-schema placement decided, not left as prose without - a concrete message.)** These fields are added to the **existing** - `RegisterMemoryRequest` message (`distributed/proto/umbp_standalone.proto:87-91`, - currently `{client_id, worker_base, size}`), not 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, - // mirrors RegisterClientRequest.tags - } - ``` - - Rationale for extending this message rather than adding a separate - `RegisterWorkerRequest`: `RegisterMemoryRequest` is already the gRPC - call that completes a worker's registration lifecycle (the - "confirm" step after the raw-UDS fd handoff, §3.2/§11.4.3) — a - worker's memory registration and its external-KV identity - registration should start and end together (the sub-identity should - not exist before the worker has a registered buffer, and should be - torn down at `DeregisterMemory`, per §11.4.7's shutdown/lifecycle - spec below), so piggybacking on the existing message avoids both a - new round trip and a second, independent lifecycle to keep in sync - with the first. `worker_node_id`/`worker_node_address` are plain - `string` fields (not `optional`) — a worker that never sets them - (e.g. a `StandaloneClient`-backed server has no use for them at all) - simply leaves them empty, and the server only constructs an - `ExternalKvIdentityClient` for a `client_id` whose - `RegisterMemoryRequest` carried a non-empty `worker_node_id`. - - **Real, bounded implementation cost, not hidden:** this means the - server runs up to N independent `ExternalKvIdentityClient` - registrations and heartbeat threads concurrently (N ≤ 8 for the - stated deployment - shape) — the same aggregate cost the baseline already pays today - (8 independent processes each running one), just consolidated into - one process. This is not free, but it is not new cost relative to - the baseline either; it is the honest price of preserving a - capability the baseline already has, not a discretionary addition. - - **(v1.1 — confirmed as the final decision, alternatives considered - and explicitly rejected, not left as an open trade-off.)** Two - points raised in review, both resolved: - - **No relay/double-hop exists in this design.** Each per-worker - `ExternalKvIdentityClient` sub-instance runs entirely inside - `umbp_standalone_server` and heartbeats `MasterServer` directly — - the worker is never in this path after the one-time `node_id` - handoff at registration; there is no "worker → server → Master" - forwarding of heartbeat traffic anywhere in this design. What *is* - real is N independent direct connections/threads originating from - the same host, not a relay chain. - - **Batching N identities onto one heartbeat connection is not - possible without a Master-protocol change, and is out of scope - here.** Verified: `RegisterClientRequest`/`HeartbeatRequest` - (`distributed/proto/umbp.proto:54-115`) both carry a single - `node_id` field, not `repeated` — the protocol has no batched- - identity shape today, and both `MasterClient` and the new - `ExternalKvIdentityClient` are 1-instance-per-`node_id` classes. - Adding batched multi-identity heartbeats would be - a `distributed/` master/protocol change affecting far more than - `umbp_standalone_server`, and is explicitly not attempted here. - - **No Mooncake precedent exists for this specific problem, checked - directly rather than assumed.** `real_client.cpp` constructs - exactly one `Client::Create(...)` for the Real Client's entire - lifetime (confirmed: exactly 2 call sites, both single-instance - construction — one primary, one port-retry variant — never one - per `DummyClient`). Multiple `DummyClient`s sharing one Real - Client are never given independent Master-facing identities in - Mooncake; multi-tenancy there is handled purely as **storage-key - namespacing** via `tenant_id` (`MakeTenantScopedStorageKey`) under - one shared identity — a fundamentally lighter-weight answer than - per-worker sub-identities, but one that only works because - Mooncake has no feature analogous to external-KV requiring - per-original-owner routing precision. This was raised explicitly - as an alternative (call it Option B: one shared identity + `client_id`- - based key/tenant namespacing, matching Mooncake's actual pattern, - at the cost of losing the multi-replica-per-host routing precision - §11.4.4 exists to preserve) and **explicitly rejected in favor of - the N-sub-identity design (Option A, this section)** — decision - confirmed: preserving parity with the `distributed/` baseline's - existing per-worker-identity capability outweighs the simplicity - Mooncake's narrower pattern would have offered here. This is - genuinely new engineering for UMBP with no precedent (in either - UMBP's own `distributed/` code or in Mooncake) to lean on for the - "N identities behind one process" mechanics specifically — flagged - honestly as such, not presented as a known-safe pattern being - reused. - - Independent of the above, add the same `client_mu_`-style locking - discipline these five handlers currently skip - (`standalone_server.cpp:206-213` pattern) around whichever - per-worker sub-identity they dispatch to, since each sub-identity - is now a stateful object with its own background threads. - - Under a `StandaloneClient` backend (no distributed identity at all), - behavior is unchanged from today: hardcoded stubs, matching - `StandaloneClient`'s own no-op external-KV methods - (`standalone_client.h:67-83`). -5. **`Ping` / capability signaling** (`standalone_server.cpp:193-197`, - `umbp_standalone.proto:33-35`): today `PingResponse` has only - `bool ready`, and nothing anywhere lets a worker learn whether the - server it's connected to is distributed-backed. - - **(v0.9 fix — upgraded from "open, low stakes" to a required part of - this proposal's implementation, not deferred.)** The reasoning in - the v0.7 draft — "the data plane doesn't care, only humans debugging - need it" — missed the actual failure mode this exists to catch: a - single `UMBP_STANDALONE_ADDRESS` can now resolve to *either* a - local-backed or a distributed-backed server (§11.6), and if the - whole point of this section is Mooncake parity (cross-node), the - most dangerous failure mode is exactly the quiet one — a worker (or - the deployment/test tooling that stood up the server) connects - successfully, every RPC succeeds, and the server has silently ended - up local-backed instead of distributed-backed (e.g. an operator's - deployment script forgot to set `UMBP_MASTER_ADDRESS` on the - *server's* side — a legitimate, non-erroring config per §11.2, not - the already-hard-failing misconfiguration case in item 6 below). - Nothing observes this today; the feature's entire purpose (cross- - node capability) can silently not exist while everything appears to - work. This is the same class of failure this document has rejected - everywhere else (§5, §8, item 6 below) — it should not be waved - through here just because the data path itself doesn't strictly - need it. - - **Decision: `PingResponse.deployment_mode` (new proto field, - enum `LOCAL`/`DISTRIBUTED`, additive/backward-compatible, - populated from `client_->GetDeploymentMode()`) is a required part - of this proposal's implementation, not an optional nice-to-have.** - Additionally: **any launch/health-check/integration-test tooling for - the distributed-backed shape must assert - `deployment_mode == DISTRIBUTED` after startup** — a distributed- - backed server that comes up as `LOCAL` (for whatever reason) must be - treated as a failed deployment, not a degraded-but-working one. The - local-backed shape itself remains a fully legitimate, intentional - deployment choice (§11.6) when that's what was actually configured - — this assertion only guards against the *parity* shape silently - collapsing into the *fallback* shape unnoticed. -6. **`standalone_server_main.cpp`**: currently unconditionally does - `config.distributed.reset()` (`standalone_server_main.cpp:43`) and - only ever builds `config.standalone_process`. Needs a **second, - independent** config-building path for `backend_config` (§11.2). - - **(v0.9 fix — the v0.7 env-var list was incomplete, not just - abbreviated; re-checked against the actual current - `umbp_store.py` code, not recalled from memory.)** The v0.7 draft - said "reuse the exact set `umbp_store.py` already reads" and then - listed only 6 variables. Re-reading `umbp_store.py:458-620` in full, - a real distributed worker's config surface is materially larger: - `master_address`, `node_address`, `node_id`, `auto_heartbeat` - (`umbp_store.py:524-527` — **note: `extra_config`-only, no env-var - fallback exists for this one today**), `io_engine.host`, - `io_engine.port`, `staging_buffer_size`, `ssd_staging_buffer_size`, - `ssd_staging_buffer_slots`, `peer_service_port`, - `cache_remote_fetches`, and `dram_page_size`. An implementer who - copied only the original 6-variable list would build a server-side - `DistributedClient` that silently behaves differently from an - ordinary distributed worker — e.g. no SSD staging buffer sizing, no - `cache_remote_fetches` admission control tuning — a real behavior - drift, not a cosmetic gap. - - **One field needs special handling, not a direct copy:** - `dram_page_size` in `umbp_store.py` (`umbp_store.py:598-620`) is - normally *auto-derived by probing `mem_pool_host`* (the worker's own - host KV tensor buffer) when not explicitly set — this derivation is - meaningless for `standalone_server_main.cpp`, which has no - `mem_pool_host` of its own (the server's `client_` backend isn't - attached to any particular worker's tensor layout — it may serve - many workers with potentially different layouts, per §4.4). The - server-side config builder must **only** support the explicit- - override path (an env var, e.g. `UMBP_DRAM_PAGE_SIZE`, mirroring - `extra_config["dram_page_size"]`'s role as "operator-controlled - escape hatch" per the comment at `umbp_store.py:590-598`) and leave - it at the `UMBPDistributedConfig` default (`0`, delegating to the - master's own `ClientRegistryConfig.default_dram_page_size`) when - unset — it must not attempt to replicate the `mem_pool_host`-probing - auto-derivation, which does not apply here. - - **Proposed shape for the implementation**: a single - `BuildDistributedBackendConfigFromEnv()` C++ helper (new, - `standalone_server_main.cpp` or a shared helper if reused elsewhere) - that mirrors `umbp_store.py`'s env-var-driven construction of - `UMBPDistributedConfig`, split explicitly into: - - **Required** (missing any of these means "distributed backend not - requested" — server falls back to local-backed, or hard-fails per - the decision later in this same item if some but not all are set, - which is itself a misconfiguration signal): `UMBP_MASTER_ADDRESS`, - `UMBP_NODE_ADDRESS`, `UMBP_NODE_ID`, `UMBP_IO_ENGINE_HOST`. - - **Optional, with the same defaults `UMBPDistributedConfig` already - has when unset**: `UMBP_IO_ENGINE_PORT`, `UMBP_PEER_SERVICE_PORT` - (both reuse the exact worker-side names, since those two already - have an established env-var convention — `umbp_store.py:526-543` — - to mirror), plus five genuinely new names, **finalized now, not - left as "proposed"**: `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` (per the special handling above). - **(v1.2 — decided, no longer "proposed.")** These five get the - `UMBP_DISTRIBUTED_` prefix specifically because, unlike - `UMBP_MASTER_ADDRESS`/`UMBP_NODE_ADDRESS`/`UMBP_NODE_ID`/ - `UMBP_IO_ENGINE_HOST`/`UMBP_IO_ENGINE_PORT`/`UMBP_PEER_SERVICE_PORT` - (which reuse worker-side names verbatim, referring to the exact - same concept in both places), these five have **no existing - env-var form at all today** — worker-side, they are - `extra_config`-only (`umbp_store.py:555-598`) — so there is no - established bare name to collide with, but there is a real future - risk of a bare `UMBP_STAGING_BUFFER_SIZE` etc. later being added - as a worker-side env var with different semantics (e.g. a - per-worker override rather than a server-backend-wide default); - the `_DISTRIBUTED_` infix makes clear these five specifically - configure the *server's own* `UMBPDistributedConfig`, not - anything worker-facing, closing off that ambiguity before it can - arise. `auto_heartbeat` has no worker-side env var today either - (only `extra_config`); **decided: the server-side helper keeps the - `UMBPMasterClientConfig` default (`true`) unconditionally** rather - than adding a dedicated env var for a field with no operational - reason to ever be `false` for these sub-identities-plus-shared- - backend registrations — simpler than inventing - `UMBP_DISTRIBUTED_AUTO_HEARTBEAT` for a knob nobody has asked to - turn off. - - **(v0.7 fix — decided, was flagged as a v0.6 blocker.) A - distributed-backed `umbp_standalone_server` supports external - launch ONLY — it does not support the §4.2 auto-start (worker - fork+exec) path. Local-backed auto-start (§4.2, already shipped in - v0.1-v0.5) is completely unaffected by this and continues to work - exactly as today.** - - The conflict this closes: v0.5 (§5) made - `UMBPStore.__init__` raise if a worker's own environment has both - `UMBP_STANDALONE_ADDRESS` and `UMBP_MASTER_ADDRESS` set (mutual - exclusivity, by design — a worker is either distributed or - standalone-process, never both). If a worker's own auto-start path - were to fork+exec a distributed-backed server, the env vars needed - to configure that child (`UMBP_MASTER_ADDRESS` etc.) would most - naturally live in the *parent's* (the worker's) environment too - (fork inherits environment by default) — which is exactly the - combination v0.5 already made illegal, on purpose, for the worker's - own config parsing. Two ways to resolve this were considered - (support auto-start via a separate `UMBP_STANDALONE_BACKEND_*` env - namespace the worker never reads, vs. drop auto-start for this mode - entirely) — **resolved by matching Mooncake rather than inventing - either option**: re-checked directly against the local Mooncake - checkout this pass (`grep -rln "fork\(|execvp|execve|posix_spawn|subprocess\.|Popen" mooncake-integration/ mooncake-wheel/` - → zero matches), Mooncake's application side (`DummyClient`) never - forks/execs the Real Client under any circumstance — the Real - Client is always started by something outside the application - (deployment system, operator, systemd unit, etc.), full stop, no - exceptions for any deployment shape. Adopting the same rule here is - simpler than a second env namespace, has zero conflict with the - already-shipped v0.5 rule, and matches the explicit goal of this - whole section (Mooncake parity) more directly than inventing a - UMBP-specific auto-start variant Mooncake itself doesn't have. - - In practice: `standalone_server_main.cpp` reads - `UMBP_MASTER_ADDRESS` etc. **from its own process environment**, - set by whatever external launcher starts it (a shell script, - Kubernetes sidecar spec, systemd unit — the same category of - launcher `run_umbp_single_node_hicache.sh` already represents for - `umbp_master` today). Workers never see or set these variables; - they only ever need `UMBP_STANDALONE_ADDRESS`, unchanged from - today. +### 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 +} +``` - **Decision (also confirmed, unchanged from v0.6): if - `UMBP_MASTER_ADDRESS` is set but `PoolClient::Init` can't actually - stand up the RDMA IO engine** (e.g. `UMBP_IO_ENGINE_HOST` empty — - `PoolClient::Init` silently skips building the IO engine entirely - in that case, `pool_client.cpp:385-415`, rather than failing), - **server startup must hard-fail**, not fall back to a - `StandaloneClient` backend with a warning. Silently downgrading a - server an operator explicitly asked to be distributed-backed into a - local-only one is exactly the kind of silent-degradation failure - mode this whole design doc has repeatedly flagged as unacceptable - elsewhere (§5, §8). A misconfigured distributed-backed server - should refuse to start, not quietly become a different, smaller - feature. - -### 11.4.7 `ExternalKvIdentityClient`: interface, lifecycle, and layering (v1.2, new) - -**(This subsection exists because §11.7 previously filed -`ExternalKvIdentityClient` as an open item needing "its own design -pass" — that framing was too weak once §11.4.4 made it a required -component of the final design, not an optional extra. This is that -design pass, not a placeholder for a future one.)** - -**Layering — this is a UMBP compatibility extension, not part of -Mooncake-parity core.** Worth stating explicitly so it is never -mistaken for required minimum-viable scope: the Mooncake-parity core -of this whole section (§11) is `distributed-backed umbp_standalone_server -+ UDS/shm fd handoff` (§11.1-§11.3) — that alone is what makes this -design match Mooncake's Real Client shape. `ExternalKvIdentityClient` -exists purely to avoid *regressing* a capability the `distributed/` -baseline already has (per-worker external-KV routing precision, §11.4.4's -"why this matters" discussion) — it has no Mooncake analog and is not -required for Mooncake parity itself. Concretely, this means: -- It can be scheduled, implemented, and tested as its own, separable - unit of work after the parity core is working — a distributed-backed - server with `ExternalKvIdentityClient` not yet implemented is a - legitimate intermediate state (external-KV simply stays unavailable - for that server, exactly like the `StandaloneClient`-backed case - already handles it today), not a broken or non-conforming one. -- It must never be treated as sitting on the critical path for - declaring "Mooncake parity achieved" — conflating the two would - under-scope the parity-core milestone by tying it to a strictly - UMBP-specific feature Mooncake itself has no equivalent of. - -**Interface** (methods on the new class; not a proto service of its -own — it is a client-side C++ class inside `umbp_standalone_server`, -using the *existing* `UMBPMaster` gRPC service just like `MasterClient` -does): +`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`** - (§11.4.4's Blocker A invariant) and the given `peer_address` - (shared across all sub-identities on this server, §11.4.4) / - `engine_desc` (may be empty/dummy — nothing ever RDMA-targets a - sub-identity's "engine," since external-KV never moves bytes through - UMBP, §11.4.4's opening clarification). Constructed the moment a + — 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` (§11.4.4 Blocker C's wire-schema decision). + `worker_node_id`. - `Heartbeat()` — periodic, liveness-only. Sends `HeartbeatRequest` - with **empty `tier_capacities` and empty `bundles`** on every call - (§11.4.4's v1.2 addition) — never derives or forwards real capacity - or owned-KV-event data. 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 (worker-initiated - teardown), and (b) server shutdown (see below). - - **(v1.4 fix — this bullet previously claimed a third trigger, - "worker disconnect/crash detection (mirrors however - `StandaloneProcessClient` disconnect is already detected... §8 item - 1's connection state machine)," as if it already existed. That was - false, and worse than simply undocumented: §8 item 1 itself states - plainly that this exact mechanism is unsolved — "Crash / restart - semantics are unsolved by every precedent reviewed... this needs - explicit handling in `StandaloneProcessClient` and probably a small - state machine (`CONNECTED` → `DISCONNECTED` → optionally - `RECONNECTING`), which does not exist in any of the three reference - systems today." §11.4.7 was citing, as its basis for "already - detected," a mechanism §8 item 1 explicitly says does not exist — - an internal contradiction a reader of this subsection alone would - never catch.)** - - **Confirmed gap, stated plainly instead of implied as solved: there - is currently no worker crash/disconnect detection anywhere in - `standalone_server.cpp`.** The only two paths that ever tear down an - `ExternalKvIdentityClient` are `DeregisterMemory` (explicit, - worker-initiated) and full server `Shutdown()` - (`UnregisterAllExternalIdentities`, see below). **Consequence, stated - explicitly:** if a worker process crashes or is killed without ever - calling `DeregisterMemory` (SIGKILL, OOM-kill, or any other - non-graceful exit), its `ExternalKvIdentityClient` keeps running - inside `umbp_standalone_server` indefinitely — its heartbeat thread - (which lives in the *server* process, not the worker's, so the - worker's death has no effect on it at all) keeps sending - liveness-only heartbeats forever, so the Master's `ClientRegistry` - entry for that `node_id` never expires and stays `ALIVE`. - `MatchExternalKv`/`GetExternalKvHitCounts` will keep returning this - node_id/peer_address to callers for KV blocks that may no longer - exist (the crashed worker's shm segment is gone), and this is **not - cleaned up until the entire `umbp_standalone_server` process is - restarted** — there is no timeout, reaper, or liveness poll for this - case anywhere in the reviewed code. This is a known, deliberately - deferred limitation (fixing it properly needs a real - liveness/crash-detection mechanism — most plausibly something built - on top of the connection-state-machine work §8 item 1 already - flags as needed for `StandaloneProcessClient`'s own reconnect - handling, not a small patch specific to `ExternalKvIdentityClient` - alone), tracked here and in §8 item 1, not silently absent. + 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`. These are what the five `standalone_server.cpp` external-KV - RPC handlers (§11.4.4, item 4) dispatch into once they've resolved - the calling `client_id` to its `ExternalKvIdentityClient` instance. - -**Re-registration must not repeat `MasterClient`'s bug (§11.4.4 -Blocker B) — stated as a hard requirement, not left implicit in "this -is new code so it won't have the bug":** every `RegisterClientRequest` -this class ever sends — the very first one and any sent after a -`CLIENT_STATUS_UNKNOWN` heartbeat response — **must** include -`peer_address` and `engine_desc`. Implementation-wise, the simplest way -to guarantee this is to store `peer_address_`/`engine_desc_bytes_` as -member fields set once at construction and have exactly one internal -`BuildRegisterRequest()` helper used for both the initial `Register()` -call and any re-register path, rather than two independently-written -request-construction code paths (which is how `MasterClient`'s bug was -introduced in the first place — the re-register path was written -separately from `RegisterSelf` and simply forgot two fields). - -**Lifecycle serialization trade-off (accepted v1.4 implementation -simplification):** the current implementation may use one server-wide -lifecycle mutex to serialize `ExternalKvIdentityClient` register, -deregister, and shutdown bookkeeping. This deliberately closes -register-vs-deregister races for the same `client_id`, but it also -means the lock can cover blocking Master RPCs (`RegisterClient` / -`UnregisterClient`, bounded by the RPC shutdown deadline). Therefore, -when the Master is slow or unavailable, otherwise unrelated workers can -queue behind each other during `RegisterMemory`/`DeregisterMemory`; in -the stated TP=8 deployment shape, the worst-case startup/shutdown delay -can be amplified toward `N * rpc_deadline` rather than one parallel -deadline. This is an accepted simplification because these are -lifecycle operations, not the Put/Get hot path, and because the current -target scale is small. If future deployments support many more workers, -frequent worker reconnects, or observe startup stalls caused by slow -Master RPCs, this lock should be the first place to revisit: either -replace it with per-`client_id` lifecycle locks or decouple -external-KV identity registration from core memory registration with an -async/retry path. - -**Shutdown ordering** (extends §4.3/§11.4.3's existing shutdown -sequence, does not replace it): `umbp_standalone_server`'s `Shutdown()` -must `Unregister()` all live `ExternalKvIdentityClient` instances -**before** tearing down the shared `DistributedClient` backend -(`client_->Close()`) or exiting — unregistering first ensures the -Master's `ClientRegistry` doesn't briefly carry stale `ALIVE` entries -for identities whose process is already gone, which would otherwise -sit around until heartbeat-timeout-driven reaping (§4, the same -"absent liveness" concern already motivated for the primary -`DistributedClient` identity). Order within `Shutdown()`, appended to -the sequence already specified in §4.3/§11.4.3: drain in-flight RPCs → -**unregister all `ExternalKvIdentityClient` instances** → flush -`CopyPipeline`/close the shared `DistributedClient` backend → unmap → -exit. - -### 11.5 What does *not* change - -- The worker-facing IPC layer: UDS control-plane socket, the separate - raw-UDS `.fd.sock` `SCM_RIGHTS` channel, offset translation, - `StandaloneProcessClient`'s locking discipline, the fd-ownership - registry in `HostMemAllocator` — **none of this is touched**. All of - it was built, reviewed, and bug-fixed independent of what backs the - server side, and stays exactly as-is. -- `local/` and today's `distributed/`-for-workers path — **both - unmodified**, per §7's existing compatibility statement, which this - proposal does not change. -- The `UMBPStandaloneProcessConfig` shape workers use (§6) — - unmodified. Workers still only ever set `address` (env-var-only - activation, per the v0.5 decision) and never need to know or care - what the server's own backend is. - -### 11.6 Consequence for "what does the local DRAM tier even mean now" - -Worth stating explicitly since it's a natural point of confusion: when -`client_` is a `DistributedClient`, the server's `LocalStorageManager`/ -`DRAMTier`/`SSDTier` concept (§3.2's earlier "server's own DRAM tier -should default to private anon/hugetlb" discussion) **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 -(`distributed_client.cpp:41-53`), used as its local slice of the -distributed pool, not as a passthrough cache for the registered worker -buffers. The registered worker buffers (via `RegisterMemory` in -§11.4.3) are a *separate* thing: they exist purely so RDMA can target -them directly, not so `DistributedClient` manages them as cache -storage. **There are now three, not two, deployment shapes for -`umbp_standalone_server`**, and the design should name them distinctly -to avoid ambiguity in docs/logs/config: + `node_id`. -| Shape | `client_` backend | Cross-node | Server's own storage | +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 | |---|---|---|---| -| existing v0.1-v0.5 (unchanged) | `StandaloneClient` | no | `LocalStorageManager` (DRAM+SSD tiers, private memory) | -| this proposal, backend enabled | `DistributedClient` | **yes**, via Master+RDMA | `DistributedClient`'s own DRAM pool (distinct concept, not `LocalStorageManager`) | -| this proposal, backend disabled | `StandaloneClient` | no | same as row 1 — this is the fallback/default, identical to today | - -### 11.7 Open questions for discussion (before any code) - -**(v0.7: items 2 and 4 from the v0.6 draft are now resolved — see -§11.4.6 and §11.3 respectively — and are kept here only as a record of -what was decided and why. Items 1, 3, 5 remain genuinely open.)** - -1. **Naming/scope framing — still open, needs a decision.** The - "Scope reframing" paragraph near the top of §11 (added this pass) - already states that - distributed-backed is the shape that actually matches the stated - goal, and local-backed is a fallback/dev-convenience shape, not a - co-equal alternative. Given that, should the *document* (and - eventually config/logs) use a name like - `standalone-process/distributed-backend` (or "real-client mode") vs - `standalone-process/local-backend` per §11.6's table, rather than - presenting them as two unlabeled options? Recommend adopting this - secondary naming now, in the doc, ahead of implementation — it's a - documentation-only decision but should be settled before code - introduces its own inconsistent naming (env var prefixes, log - messages, etc). -2. ~~The hard-fail-on-misconfiguration decision in §11.4.6~~ — - **resolved**: hard-fail, confirmed in §11.4.6. -3. ~~Whether `PingResponse.deployment_mode` is worth the proto change - now~~ — **resolved (v0.9): required, not deferred.** See §11.4.5 — - this is now a mandatory guard against the parity shape silently - collapsing into the local-only fallback shape, not an optional - debugging aid. -4. ~~The cross-mapping memory-visibility question in §11.3~~ — - **resolved**: closed via the Mooncake production-precedent citation - plus the explicit gRPC-response-as-happens-before-edge contract, see - §11.3. -5. Testing implications, not yet scoped: the existing - `test_standalone_shm_ipc.cpp` suite only exercises the - `StandaloneClient`-backed path. A `DistributedClient`-backed test - needs at minimum a running `umbp_master` plus the RDMA/IO-engine - stack available in the test environment — likely a heavier - integration test than what exists today, possibly gated separately - (e.g. only run where RDMA hardware/loopback is available), which - needs its own test-infrastructure decision before it can be added - to `§10`'s implementation-order list. Should also cover: per-worker - external-KV sub-identity registration/dispatch under a - distributed-backed server with 2+ connected workers (§11.4.4 — this - is the item's own core mechanism now, not an edge case), and the - `RegisterMemory`/`DeregisterMemory` rollback/shutdown-ordering paths - added in §11.4.3. -6. ~~The exact registration-handshake field(s) for per-worker identity - conveyance~~ — **resolved (v1.2, Blocker C): `worker_node_id` + - `worker_node_address` + optional `tags`, not raw rank components.** - See §11.4.4. -7. ~~`ExternalKvIdentityClient` needs a design/implementation pass of - its own~~ — **resolved (v1.2): see §11.4.7** for the full interface - (`Register`/`Heartbeat`/`Unregister`/the five external-KV RPCs), - the re-registration hard requirement, and shutdown ordering. It is - still real, net-new implementation work — it doesn't exist anywhere - in the codebase today, unlike `MasterClient`, which at least could - have been (incorrectly) reused — so it still needs to be factored - into scheduling, not just the N-heartbeat-threads cost already noted - above; what changed is that its shape is now fully specified rather - than deferred. - -### 11.8 Source references for this section - -- `/apps/nima/KVManager/mori/src/umbp/standalone/standalone_server.{h,cpp}` -- `/apps/nima/KVManager/mori/src/umbp/standalone/bin/standalone_server_main.cpp` -- `/apps/nima/KVManager/mori/src/umbp/distributed/distributed_client.{h,cpp}` -- `/apps/nima/KVManager/mori/src/umbp/distributed/pool_client.{h,cpp}` -- `/apps/nima/KVManager/mori/src/umbp/include/umbp/distributed/master/master_client.h` -- `/apps/nima/KVManager/mori/src/umbp/include/umbp/common/config.h` - (`UMBPConfig::Validate()`, `UMBPDistributedConfig`, - `UMBPMasterClientConfig`, `UMBPIoEngineConfig`) -- `/apps/nima/KVManager/mori/src/umbp/umbp_client_factory.cpp` -- `/apps/nima/KVManager/mori/src/umbp/distributed/proto/umbp_standalone.proto` -- `/apps/nima/KVManager/mori/src/umbp/include/umbp/umbp_client.h` - (`IUMBPClient::RegisterMemory`'s default no-op-returns-true contract) +| `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/`.