diff --git a/dependencies.sh b/dependencies.sh index 8420bc32d1..763109f024 100755 --- a/dependencies.sh +++ b/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)..." diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index aa4799448b..3f0c7d2d99 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/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 | diff --git a/docs/source/getting_started/supported-protocols.md b/docs/source/getting_started/supported-protocols.md index 6ebbb9a796..e959e281ee 100644 --- a/docs/source/getting_started/supported-protocols.md +++ b/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 | diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index eefa33bdb4..8ca663a94a 100644 --- a/mooncake-common/common.cmake +++ b/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() diff --git a/mooncake-common/include/rpc_transport_config.h b/mooncake-common/include/rpc_transport_config.h new file mode 100644 index 0000000000..f3eb7ac1d2 --- /dev/null +++ b/mooncake-common/include/rpc_transport_config.h @@ -0,0 +1,103 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#ifdef YLT_ENABLE_URMA +#include +#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( + queue_depth, std::numeric_limits::max()); + + auto buffer_size = GetRpcUrmaUintFromEnv("MC_RPC_URMA_BUFFER_SIZE", 4096); + buffer_size = std::min( + buffer_size, std::numeric_limits::max()); + + auto max_memory_mib = + GetRpcUrmaUintFromEnv("MC_RPC_URMA_MAX_MEMORY_MIB", 256); + max_memory_mib = std::min( + max_memory_mib, std::numeric_limits::max() / 1024 / 1024); + + auto eid_index = GetRpcUrmaUintFromEnv("MC_RPC_URMA_EID_INDEX", 0); + eid_index = std::min( + eid_index, static_cast(std::numeric_limits::max())); + + coro_io::urma_socket_t::config_t config{}; + config.cq_size = static_cast(queue_depth * 2 + 8); + config.recv_buffer_cnt = static_cast(queue_depth); + config.send_buffer_cnt = static_cast(queue_depth); + config.buffer_size = static_cast(buffer_size); + config.max_memory_usage = max_memory_mib * 1024ull * 1024ull; + config.device_name = GetRpcUrmaDeviceFromEnv(); + config.eid_index = static_cast(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 diff --git a/mooncake-store/benchmarks/master_bench.cpp b/mooncake-store/benchmarks/master_bench.cpp index ad774a33b0..92b34cd4c0 100644 --- a/mooncake-store/benchmarks/master_bench.cpp +++ b/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); diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index cad4159225..d30d1394c0 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -1,16 +1,20 @@ #pragma once #include +#include #include #include #include -#include #include -#include +#include #include +#include #include #include #include +#ifdef YLT_ENABLE_URMA +#include +#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 +inline void MaybeEnableUrmaSocketConfig(SocketConfigVariant& socket_config) { + if constexpr (variant_contains_v) { + 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>( pool_conf); diff --git a/mooncake-store/src/dummy_client.cpp b/mooncake-store/src/dummy_client.cpp index c1a84a5c2a..6883d637ea 100644 --- a/mooncake-store/src/dummy_client.cpp +++ b/mooncake-store/src/dummy_client.cpp @@ -2,6 +2,9 @@ #include #include #include +#ifdef YLT_ENABLE_URMA +#include +#endif #include // For shm_open, mmap, munmap #include // 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::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>( pool_conf); diff --git a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp index 4f9ff54388..f40c724158 100644 --- a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp +++ b/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( mooncake::WrappedMasterServiceConfig( diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index 9e3c16facb..6c39e714ed 100644 --- a/mooncake-store/src/master.cpp +++ b/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::WrappedMasterServiceConfig(master_config, version)); diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 4e67ecc78f..73903bdb13 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -24,6 +24,9 @@ #include #include "mooncake_logging.h" +#ifdef YLT_ENABLE_URMA +#include +#endif #include "real_client.h" #include "client_buffer.hpp" @@ -33,6 +36,7 @@ #include "types.h" #include "utils.h" #include "rpc_types.h" +#include "rpc_transport_config.h" #include "file_storage.h" #include "gpu_staging_utils.h" #include "default_config.h" @@ -925,6 +929,15 @@ tl::expected RealClient::setup_internal( // Use port 0 to let the OS auto-allocate an available port. offload_rpc_server_ = std::make_unique(1, 0, "0.0.0.0"); +#ifdef YLT_ENABLE_URMA + const char *rpc_protocol = std::getenv("MC_RPC_PROTOCOL"); + if (rpc_protocol && std::string_view(rpc_protocol) == "urma") { + auto urma_config = MakeUrmaRpcConfigFromEnv(); + LOG(INFO) << "Offload RPC server using URMA RPC transport: " + << FormatUrmaRpcConfig(urma_config); + offload_rpc_server_->init_urma(urma_config); + } +#endif offload_rpc_server_ ->register_handler<&RealClient::batch_get_offload_object>(this); offload_rpc_server_ @@ -6323,6 +6336,14 @@ ClientRequester::ClientRequester() { 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) << "ClientRequester using URMA RPC transport: " + << FormatUrmaRpcConfig(urma_config); + } +#endif // Configure reasonable retry limits for SSD offload RPC connections. // - connect_retry_count: Maximum connection retry attempts (default: 3) // - reconnect_wait_time: Wait time between retries (default: 1000ms) diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index 8f0a4ba80c..1fe61cb037 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -1,11 +1,17 @@ #include #include +#include +#include #include +#ifdef YLT_ENABLE_URMA +#include +#endif #include "client_service.h" #include "config.h" #include "mooncake_logging.h" #include "real_client.h" +#include "rpc_transport_config.h" using namespace mooncake; @@ -129,6 +135,15 @@ int main(int argc, char *argv[]) { } coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, FLAGS_host); +#ifdef YLT_ENABLE_URMA + const char *rpc_protocol = std::getenv("MC_RPC_PROTOCOL"); + if (rpc_protocol && std::string_view(rpc_protocol) == "urma") { + auto urma_config = MakeUrmaRpcConfigFromEnv(); + LOG(INFO) << "Real client service using URMA RPC transport: " + << FormatUrmaRpcConfig(urma_config); + server.init_urma(urma_config); + } +#endif RegisterClientRpcService(server, *client_inst); LOG(INFO) << "Starting real client service on " << FLAGS_host << ":" diff --git a/mooncake-store/tests/test_server_helpers.h b/mooncake-store/tests/test_server_helpers.h index 5336e0e122..2d81b451e4 100644 --- a/mooncake-store/tests/test_server_helpers.h +++ b/mooncake-store/tests/test_server_helpers.h @@ -14,6 +14,7 @@ #include "http_metadata_server.h" #include "master_config.h" #include "rpc_service.h" +#include "rpc_transport_config.h" #include "types.h" #include #include "utils.h" @@ -93,6 +94,12 @@ class InProcMaster { 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 = MakeUrmaRpcConfigFromEnv(); + server_->init_urma(urma_config); + } +#endif uint64_t default_kv_lease_ttl = DEFAULT_DEFAULT_KV_LEASE_TTL; if (config.default_kv_lease_ttl.has_value()) { diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp index 158725f0ed..09e1800d1a 100644 --- a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp @@ -7,10 +7,14 @@ #include #include #include +#ifdef YLT_ENABLE_URMA +#include +#endif #include #include #include "async_simple/coro/SyncAwait.h" #include "default_config.h" +#include "rpc_transport_config.h" namespace mooncake { namespace py = pybind11; @@ -51,16 +55,23 @@ bool RpcCommunicator::initialize(const RpcCommunicatorConfig& config) { // Initialize client pools with proper configuration coro_io::client_pool::pool_config pool_conf{}; const char* value = std::getenv("MC_RPC_PROTOCOL"); +#ifdef YLT_ENABLE_URMA + bool use_urma_transport = false; +#endif if (value && std::string_view(value) == "rdma") { pool_conf.client_config.socket_config = coro_io::ib_socket_t::config_t{}; } - client_pools_ = - std::make_shared>( - pool_conf); +#ifdef YLT_ENABLE_URMA + else if (value && std::string_view(value) == "urma") { + auto urma_config = MakeUrmaRpcConfigFromEnv(); + pool_conf.client_config.socket_config = urma_config; + use_urma_transport = true; + LOG(INFO) << "RpcCommunicator client pool using URMA RPC transport: " + << FormatUrmaRpcConfig(urma_config); + } +#endif - LOG(INFO) << "create coro_rpc_client_pool with " << config.pool_size - << " threads"; if (!config.listen_address.empty()) { LOG(INFO) << "Initializing server on " << config.listen_address; @@ -88,15 +99,58 @@ bool RpcCommunicator::initialize(const RpcCommunicatorConfig& config) { LOG(WARNING) << "Falling back to TCP mode"; } } +#ifdef YLT_ENABLE_URMA + else if (value && std::string_view(value) == "urma") { + if (server_) { + try { + auto urma_config = MakeUrmaRpcConfigFromEnv(); + LOG(INFO) << "RpcCommunicator server using URMA RPC " + << "transport: " + << FormatUrmaRpcConfig(urma_config); + server_->init_urma(urma_config); + LOG(INFO) << "URMA initialized successfully"; + } catch (const std::exception& e) { + LOG(ERROR) << "URMA initialization failed: " << e.what(); + LOG(WARNING) << "Falling back to TCP mode"; + pool_conf = {}; + use_urma_transport = false; + } catch (...) { + LOG(ERROR) + << "URMA initialization failed with unknown error"; + LOG(WARNING) << "Falling back to TCP mode"; + pool_conf = {}; + use_urma_transport = false; + } + } else { + LOG(ERROR) << "Server pointer is null, cannot initialize URMA"; + LOG(WARNING) << "Falling back to TCP mode"; + pool_conf = {}; + use_urma_transport = false; + } + } +#endif server_->register_handler<&RpcCommunicator::handleDataTransfer, &RpcCommunicator::handleTensorTransfer>(this); } + + client_pools_ = + std::make_shared>( + pool_conf); + + LOG(INFO) << "create coro_rpc_client_pool with " << config.pool_size + << " threads"; LOG(INFO) << "Environment variable MC_RPC_PROTOCOL is set to " << (value ? value : "not set"); if (value && std::string_view(value) == "rdma") { LOG(INFO) << "Using RDMA transport for RPC communication"; - } else { + } +#ifdef YLT_ENABLE_URMA + else if (use_urma_transport) { + LOG(INFO) << "Using URMA transport for RPC communication"; + } +#endif + else { LOG(INFO) << "Using TCP transport for RPC communication"; } diff --git a/mooncake-wheel/tests/test_meta_server.py b/mooncake-wheel/tests/test_meta_server.py index 2a2f27b80a..49dfc23594 100644 --- a/mooncake-wheel/tests/test_meta_server.py +++ b/mooncake-wheel/tests/test_meta_server.py @@ -5,7 +5,8 @@ from mooncake.store import MooncakeDistributedStore # how to test: start mooncake_master and http_metadata_server. -# Default is tcp, if you want to use rdma, should set MC_RPC_PROTOCOL=rdma and DEVICE_NAME=rdma_xxx at first. +# Default is tcp. Set MC_RPC_PROTOCOL=rdma with DEVICE_NAME=rdma_xxx for RDMA, +# or set MC_RPC_PROTOCOL=urma on builds with USE_YLT_URMA_RPC=ON for URMA RPC. # ./mooncake_master # python mooncake-wheel/mooncake/http_metadata_server.py --port 8080 # python mooncake-wheel/tests/test_meta_server.py 127.0.0.1 8 @@ -21,8 +22,12 @@ def test_worker(store, num_req): num_thread = 1 if len(sys.argv) < 3 else int(sys.argv[2]) # Initialize the store store = MooncakeDistributedStore() -# Use TCP protocol by default for testing, also support rdma -protocol = os.getenv("MC_RPC_PROTOCOL", "tcp") +# Use TCP data-plane protocol by default. Keep the historical +# MC_RPC_PROTOCOL=rdma shortcut, but do not pass control-plane-only URMA to +# store.setup(). Set MOONCAKE_PROTOCOL to choose the data-plane explicitly. +protocol = os.getenv("MOONCAKE_PROTOCOL") +if protocol is None: + protocol = "rdma" if os.getenv("MC_RPC_PROTOCOL") == "rdma" else "tcp" device_name = os.getenv("DEVICE_NAME", "eth0") local_hostname = os.getenv("LOCAL_HOSTNAME", "127.0.0.1") metadata_server = os.getenv("METADATA_ADDR", f"http://{master_host}:8080/metadata")