Zcode urma rpc

dependencies.sh

@@ -237,7 +237,17 @@ cd build
 check_success "Failed to change to build directory"
 
 echo "Configuring yalantinglibs..."
-cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF
+ylt_cmake_args=(
+    -DBUILD_EXAMPLES=OFF
+    -DBUILD_BENCHMARK=OFF
+    -DBUILD_UNIT_TESTS=OFF
+)
+
+if [ "${USE_YLT_URMA_RPC:-OFF}" = "ON" ] || [ "${USE_YLT_URMA_RPC:-OFF}" = "1" ]; then
+    ylt_cmake_args+=( -DYLT_ENABLE_URMA=ON )
+fi
+
+cmake .. "${ylt_cmake_args[@]}"
 check_success "Failed to configure yalantinglibs"
 
 echo "Building yalantinglibs (using $(nproc) cores)..."

docs/source/deployment/mooncake-store-deployment-guide.md

@@ -541,7 +541,12 @@ rpc_port: 50051
 
 | Variable | Default | Description |
 |----------|---------|-------------|
-| `MC_RPC_PROTOCOL` | `tcp` | RPC transport protocol between master and clients: `tcp` or `rdma` |
+| `MC_RPC_PROTOCOL` | `tcp` | RPC transport protocol between master and clients: `tcp`, `rdma`, or `urma`. `urma` requires `USE_YLT_URMA_RPC=ON`, yalantinglibs built with `YLT_ENABLE_URMA=ON`, and URMA/UMDK development and runtime libraries on the host. |
+| `MC_RPC_URMA_DEVICE` | `bonding_dev_0` | URMA device used when `MC_RPC_PROTOCOL=urma`. |
+| `MC_RPC_URMA_EID_INDEX` | `0` | URMA EID index used when `MC_RPC_PROTOCOL=urma`. |
+| `MC_RPC_URMA_QUEUE_DEPTH` | `64` | URMA send/receive queue depth used when `MC_RPC_PROTOCOL=urma`; also controls completion queue sizing. |
+| `MC_RPC_URMA_BUFFER_SIZE` | `4096` | URMA per-buffer size in bytes used when `MC_RPC_PROTOCOL=urma`. |
+| `MC_RPC_URMA_MAX_MEMORY_MIB` | `256` | URMA RPC buffer pool memory limit in MiB used when `MC_RPC_PROTOCOL=urma`. |
 | `MC_USE_TENT` / `MC_USE_TEV1` | unset | Set to any value to enable the TENT (next-gen) transfer engine |
 | `MC_STORE_CLUSTER_ID` | unset | Cluster ID label attached to client metrics |
 

docs/source/getting_started/supported-protocols.md

@@ -314,6 +314,30 @@ export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE"
 export MOONCAKE_LOCAL_HOSTNAME="node1"
 ```
 
+### yalantinglibs RPC transport
+
+Mooncake's data-plane transport is configured separately from the yalantinglibs `coro_rpc` control-plane transport. The control-plane RPC transport is selected with `MC_RPC_PROTOCOL`:
+
+```bash
+# Default: TCP control-plane RPC
+unset MC_RPC_PROTOCOL
+
+# Existing IBV/RDMA control-plane RPC
+export MC_RPC_PROTOCOL=rdma
+
+# URMA control-plane RPC, requires Mooncake built with USE_YLT_URMA_RPC=ON,
+# yalantinglibs built with YLT_ENABLE_URMA=ON, and URMA/UMDK development
+# and runtime libraries available on the host.
+export MC_RPC_PROTOCOL=urma
+
+# Optional URMA RPC tuning. Defaults match the yalantinglibs URMA benchmark.
+export MC_RPC_URMA_DEVICE=bonding_dev_0
+export MC_RPC_URMA_EID_INDEX=0
+export MC_RPC_URMA_QUEUE_DEPTH=64
+export MC_RPC_URMA_BUFFER_SIZE=4096
+export MC_RPC_URMA_MAX_MEMORY_MIB=256
+```
+
 ## Choosing the Right Protocol
 
 | Scenario | Recommended Protocol | Notes |

mooncake-common/common.cmake

@@ -86,6 +86,7 @@ option(USE_MNNVL "option for using Multi-Node NVLink transport" OFF)
 option(USE_CXL "option for using CXL protocol" OFF)
 option(USE_EFA "option for using AWS EFA transport" OFF)
 option(USE_UB "option for using UB protocol transport" OFF)
+option(USE_YLT_URMA_RPC "option for enabling yalantinglibs URMA RPC transport" OFF)
 option(USE_SUNRISE "option for enabling gpu features for Sunrise GPU with Tang runtime" OFF)
 
 if (USE_UB)
@@ -429,5 +430,11 @@ endif()
 set(GFLAGS_USE_TARGET_NAMESPACE "true")
 find_package(yaml-cpp REQUIRED)
 find_package(gflags REQUIRED)
+if (USE_YLT_URMA_RPC)
+  set(YLT_ENABLE_URMA ON CACHE BOOL "Enable URMA support in yalantinglibs" FORCE)
+endif()
 find_package(yalantinglibs CONFIG REQUIRED)
 add_compile_definitions(YLT_ENABLE_IBV)
+if (USE_YLT_URMA_RPC)
+  add_compile_definitions(YLT_ENABLE_URMA)
+endif()

mooncake-common/include/rpc_transport_config.h

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+#include <limits>
+#include <sstream>
+#include <string>
+#include <string_view>
+
+#ifdef YLT_ENABLE_URMA
+#include <ylt/coro_io/urma/urma_socket.hpp>
+#endif
+
+namespace mooncake {
+
+inline std::string GetRpcProtocolFromEnv() {
+    const char* value = std::getenv("MC_RPC_PROTOCOL");
+    if (value && std::string_view(value) == "rdma") {
+        return "rdma";
+    }
+#ifdef YLT_ENABLE_URMA
+    if (value && std::string_view(value) == "urma") {
+        return "urma";
+    }
+#endif
+    return "tcp";
+}
+
+inline bool IsUrmaRpcProtocol() {
+    const char* value = std::getenv("MC_RPC_PROTOCOL");
+    return value && std::string_view(value) == "urma";
+}
+
+#ifdef YLT_ENABLE_URMA
+
+inline std::string GetRpcUrmaDeviceFromEnv() {
+    const char* value = std::getenv("MC_RPC_URMA_DEVICE");
+    return value ? std::string(value) : std::string("bonding_dev_0");
+}
+
+inline uint64_t GetRpcUrmaUintFromEnv(const char* name,
+                                      uint64_t default_value) {
+    const char* value = std::getenv(name);
+    if (!value || *value == '\0' || *value == '-') {
+        return default_value;
+    }
+
+    char* end = nullptr;
+    auto parsed = std::strtoull(value, &end, 10);
+    if (end == value || *end != '\0') {
+        return default_value;
+    }
+    return parsed;
+}
+
+inline coro_io::urma_socket_t::config_t MakeUrmaRpcConfigFromEnv() {
+    auto queue_depth = GetRpcUrmaUintFromEnv("MC_RPC_URMA_QUEUE_DEPTH", 64);
+    if (queue_depth == 0) {
+        queue_depth = 1;
+    }
+    queue_depth = std::min<uint64_t>(
+        queue_depth, std::numeric_limits<uint16_t>::max());
+
+    auto buffer_size = GetRpcUrmaUintFromEnv("MC_RPC_URMA_BUFFER_SIZE", 4096);
+    buffer_size = std::min<uint64_t>(
+        buffer_size, std::numeric_limits<uint32_t>::max());
+
+    auto max_memory_mib =
+        GetRpcUrmaUintFromEnv("MC_RPC_URMA_MAX_MEMORY_MIB", 256);
+    max_memory_mib = std::min<uint64_t>(
+        max_memory_mib, std::numeric_limits<uint64_t>::max() / 1024 / 1024);
+
+    auto eid_index = GetRpcUrmaUintFromEnv("MC_RPC_URMA_EID_INDEX", 0);
+    eid_index = std::min<uint64_t>(
+        eid_index, static_cast<uint64_t>(std::numeric_limits<int>::max()));
+
+    coro_io::urma_socket_t::config_t config{};
+    config.cq_size = static_cast<uint32_t>(queue_depth * 2 + 8);
+    config.recv_buffer_cnt = static_cast<uint16_t>(queue_depth);
+    config.send_buffer_cnt = static_cast<uint16_t>(queue_depth);
+    config.buffer_size = static_cast<uint32_t>(buffer_size);
+    config.max_memory_usage = max_memory_mib * 1024ull * 1024ull;
+    config.device_name = GetRpcUrmaDeviceFromEnv();
+    config.eid_index = static_cast<int>(eid_index);
+    config.tp_type = URMA_CTP;
+    return config;
+}
+
+inline std::string FormatUrmaRpcConfig(
+    const coro_io::urma_socket_t::config_t& config) {
+    std::ostringstream os;
+    os << "device=" << config.device_name << ", eid_index=" << config.eid_index
+       << ", queue_depth=" << config.recv_buffer_cnt
+       << ", cq_size=" << config.cq_size
+       << ", buffer_size=" << config.buffer_size
+       << ", max_memory_mib=" << config.max_memory_usage / 1024 / 1024;
+    return os.str();
+}
+
+#endif
+
+}  // namespace mooncake

mooncake-store/benchmarks/master_bench.cpp

@@ -25,6 +25,7 @@
 #include "glog/logging.h"
 
 #include "master_client.h"
+#include "rpc_transport_config.h"
 
 // Size units for better readability
 static constexpr size_t KiB = 1024;
@@ -370,6 +371,16 @@ int main(int argc, char** argv) {
 
     gflags::ParseCommandLineFlags(&argc, &argv, false);
 
+    LOG(INFO) << "master_bench RPC transport: "
+              << mooncake::GetRpcProtocolFromEnv();
+#ifdef YLT_ENABLE_URMA
+    if (mooncake::IsUrmaRpcProtocol()) {
+        auto urma_config = mooncake::MakeUrmaRpcConfigFromEnv();
+        LOG(INFO) << "master_bench URMA RPC config: "
+                  << mooncake::FormatUrmaRpcConfig(urma_config);
+    }
+#endif
+
     ping_thread = std::jthread([&](std::stop_token stop_token) {
         static const auto OneSecond = std::chrono::seconds(1);
 

mooncake-store/include/master_client.h

@@ -1,16 +1,20 @@
 #pragma once
 
 #include <csignal>
+#include <cstdlib>
 #include <memory>
 #include <string>
 #include <type_traits>
-#include <vector>
 #include <variant>
-#include <cstdlib>
+#include <vector>
 #include <boost/functional/hash.hpp>
+#include <glog/logging.h>
 #include <ylt/coro_rpc/coro_rpc_client.hpp>
 #include <ylt/coro_io/client_pool.hpp>
 #include <ylt/coro_io/ibverbs/ib_socket.hpp>
+#ifdef YLT_ENABLE_URMA
+#include <ylt/coro_io/urma/urma_socket.hpp>
+#endif
 
 #include "client_metric.h"
 #include "replica.h"
@@ -19,6 +23,7 @@
 #include "rpc_types.h"
 #include "master_metric_manager.h"
 #include "task_manager.h"
+#include "rpc_transport_config.h"
 
 namespace mooncake {
 
@@ -45,6 +50,16 @@ inline void MaybeEnableRdmaSocketConfig(SocketConfigVariant& socket_config) {
     }
 }
 
+#ifdef YLT_ENABLE_URMA
+template <typename SocketConfigVariant>
+inline void MaybeEnableUrmaSocketConfig(SocketConfigVariant& socket_config) {
+    if constexpr (variant_contains_v<SocketConfigVariant,
+                                     coro_io::urma_socket_t::config_t>) {
+        socket_config = MakeUrmaRpcConfigFromEnv();
+    }
+}
+#endif
+
 }  // namespace detail
 
 /**
@@ -70,6 +85,15 @@ class MasterClient {
             detail::MaybeEnableRdmaSocketConfig(
                 pool_conf.client_config.socket_config);
         }
+#ifdef YLT_ENABLE_URMA
+        else if (value && std::string_view(value) == "urma") {
+            detail::MaybeEnableUrmaSocketConfig(
+                pool_conf.client_config.socket_config);
+            auto urma_config = MakeUrmaRpcConfigFromEnv();
+            LOG(INFO) << "MasterClient using URMA RPC transport: "
+                      << FormatUrmaRpcConfig(urma_config);
+        }
+#endif
         client_pools_ =
             std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
                 pool_conf);

mooncake-store/src/dummy_client.cpp

@@ -2,6 +2,9 @@
 #include <csignal>
 #include <ylt/easylog/record.hpp>
 #include <ylt/coro_rpc/coro_rpc_client.hpp>
+#ifdef YLT_ENABLE_URMA
+#include <ylt/coro_io/urma/urma_socket.hpp>
+#endif
 
 #include <sys/mman.h>  // For shm_open, mmap, munmap
 #include <sys/stat.h>  // For S_IRUSR, S_IWUSR
@@ -18,6 +21,7 @@
 #include "types.h"
 #include "default_config.h"
 #include "config.h"
+#include "rpc_transport_config.h"
 #ifdef USE_ASCEND_DIRECT
 #include "acl/acl_rt.h"
 #include "ascend_allocator.h"
@@ -220,6 +224,19 @@ DummyClient::DummyClient()
     mooncake::init_ylt_log_level();
     // Initialize client pools
     coro_io::client_pool<coro_rpc::coro_rpc_client>::pool_config pool_conf{};
+    const char *value = std::getenv("MC_RPC_PROTOCOL");
+    if (value && std::string_view(value) == "rdma") {
+        pool_conf.client_config.socket_config =
+            coro_io::ib_socket_t::config_t{};
+    }
+#ifdef YLT_ENABLE_URMA
+    else if (value && std::string_view(value) == "urma") {
+        auto urma_config = MakeUrmaRpcConfigFromEnv();
+        pool_conf.client_config.socket_config = urma_config;
+        LOG(INFO) << "DummyClient using URMA RPC transport: "
+                  << FormatUrmaRpcConfig(urma_config);
+    }
+#endif
     client_pools_ =
         std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
             pool_conf);

mooncake-store/src/ha/leadership/master_service_supervisor.cpp

@@ -14,6 +14,7 @@
 #include "ha/leadership/leader_coordinator_factory.h"
 #include "ha/standby_controller.h"
 #include "rpc_service.h"
+#include "rpc_transport_config.h"
 
 namespace mooncake {
 namespace ha {
@@ -348,6 +349,14 @@ int RunSupervisorLoop(const HABackendSpec& spec,
         if (protocol && std::string_view(protocol) == "rdma") {
             server.init_ibv();
         }
+#ifdef YLT_ENABLE_URMA
+        else if (protocol && std::string_view(protocol) == "urma") {
+            auto urma_config = MakeUrmaRpcConfigFromEnv();
+            LOG(INFO) << "HA master service using URMA RPC transport: "
+                      << FormatUrmaRpcConfig(urma_config);
+            server.init_urma(urma_config);
+        }
+#endif
 
         auto wrapped_master_service = std::make_shared<WrappedMasterService>(
             mooncake::WrappedMasterServiceConfig(

mooncake-store/src/master.cpp

@@ -19,6 +19,7 @@
 #include "http_metadata_server.h"
 #include "mooncake_logging.h"
 #include "rpc_service.h"
+#include "rpc_transport_config.h"
 #include "types.h"
 #include "utils.h"
 
@@ -1076,11 +1077,7 @@ int main(int argc, char* argv[]) {
         return 1;
     }
 
-    const char* value = std::getenv("MC_RPC_PROTOCOL");
-    std::string protocol = "tcp";
-    if (value && std::string_view(value) == "rdma") {
-        protocol = "rdma";
-    }
+    auto protocol = mooncake::GetRpcProtocolFromEnv();
     LOG(INFO)
         << "Master service started on port " << master_config.rpc_port
         << ", max_threads=" << master_config.rpc_thread_num
@@ -1181,6 +1178,14 @@ int main(int argc, char* argv[]) {
         if (value && std::string_view(value) == "rdma") {
             server.init_ibv();
         }
+#ifdef YLT_ENABLE_URMA
+        else if (value && std::string_view(value) == "urma") {
+            auto urma_config = mooncake::MakeUrmaRpcConfigFromEnv();
+            LOG(INFO) << "Master service using URMA RPC transport: "
+                      << mooncake::FormatUrmaRpcConfig(urma_config);
+            server.init_urma(urma_config);
+        }
+#endif
         auto wrapped_master_service =
             std::make_shared<mooncake::WrappedMasterService>(
                 mooncake::WrappedMasterServiceConfig(master_config, version));