From 70b219b630b53faedc8161abec4896eec0149902 Mon Sep 17 00:00:00 2001 From: TTThanos Date: Fri, 11 Sep 2026 13:42:29 +0800 Subject: [PATCH] [Bugfix][TENT] Forward classic RDMA device filters Preserve the device whitelist passed through the classic TransferEngine API when MC_USE_TENT is enabled, and inject it into TENT's `topology/rdma_whitelist` before topology discovery. Keep explicit caller filters ahead of MC_TENT_CONF and MC_TE_FILTERS defaults, retain filters across move operations and the pre-init setter, and leave empty filters to the existing environment configuration. Add compatibility, configuration-precedence, move-semantics, and real multi-HCA regression coverage. --- .../transfer_engine/transfer_engine_py.cpp | 13 ++- .../include/transfer_engine.h | 8 ++ .../src/transfer_engine.cpp | 47 ++++++--- .../tent/src/runtime/transfer_engine_impl.cpp | 5 + .../transfer_engine_config_override_test.cpp | 39 ++++++++ mooncake-transfer-engine/tests/CMakeLists.txt | 2 + .../tests/transport_uint_test.cpp | 98 +++++++++++++++++++ 7 files changed, 192 insertions(+), 20 deletions(-) diff --git a/mooncake-integration/transfer_engine/transfer_engine_py.cpp b/mooncake-integration/transfer_engine/transfer_engine_py.cpp index 3b87e59e4a..0b9ae65f14 100644 --- a/mooncake-integration/transfer_engine/transfer_engine_py.cpp +++ b/mooncake-integration/transfer_engine/transfer_engine_py.cpp @@ -1158,13 +1158,12 @@ std::string TransferEnginePy::getLocalTopology(const char* device_name) { getenv("MC_USE_TENT") != nullptr || getenv("MC_USE_TEV1") != nullptr; #ifdef USE_TENT if (use_tent) { - // The classic shim (TransferEngine(true, filter)) silently drops the - // filter on the TENT path and builds its own Config in init(), so - // inject the whitelist via the per-instance Config that TENT's public - // constructor already accepts. Avoids touching the process-global - // MC_TE_FILTERS env var (racey under concurrent callers, leaked on - // throw). Note: if MC_TE_FILTERS is also set in env, loadFromEnv() - // inside TransferEngineImpl will override this — env takes priority. + // This helper only needs topology, so use the native TENT Config path + // directly instead of constructing the classic compatibility shim. + // Keep the filter per-instance and avoid the process-global + // MC_TE_FILTERS environment variable, which is unsafe for concurrent + // callers. Explicit Config values take precedence over environment + // defaults inside TransferEngineImpl. auto conf = std::make_shared(); conf->set("metadata_type", "p2p"); if (!device_name_safe.empty()) { diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index d8e4350bfa..b4342ce893 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -33,6 +33,7 @@ namespace mooncake { class ShutdownToken; class TransferEngineImpl; namespace tent { +class Config; class TransferEngine; }; #if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ @@ -294,9 +295,16 @@ class TransferEngine { std::string showLinks(bool json = false) const; private: + std::shared_ptr buildTentConfig( + const std::string& metadata_conn_string, + const std::string& local_server_name) const; + std::shared_ptr impl_; std::shared_ptr impl_tent_; std::shared_ptr shutdown_token_; + // Classic callers provide this through TransferEngine(auto_discover, + // filter) before init() creates the native TENT engine. + std::vector tent_device_filter_; bool use_tent_{false}; friend class TransferEngineImplTestPeer; }; diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 56be21f692..2ada330931 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -93,6 +93,7 @@ TransferEngine& TransferEngine::operator=(TransferEngine&& other) noexcept { freeEngine(); impl_ = std::move(other.impl_); impl_tent_ = std::move(other.impl_tent_); + tent_device_filter_ = std::move(other.tent_device_filter_); use_tent_ = other.use_tent_; const bool shutdown_enabled = static_cast(other.shutdown_token_); detachShutdownToken(other.shutdown_token_); @@ -421,7 +422,9 @@ TransferEngine::TransferEngine(bool auto_discover, if (getenv("MC_USE_TENT") || getenv("MC_USE_TEV1")) { use_tent_ = true; } - if (!use_tent_) { + if (use_tent_) { + tent_device_filter_ = filter; + } else { impl_ = std::make_shared(auto_discover, filter); } } @@ -430,6 +433,7 @@ TransferEngine::TransferEngine(TransferEngine&& other) noexcept : impl_(std::move(other.impl_)), impl_tent_(std::move(other.impl_tent_)), shutdown_token_(nullptr), + tent_device_filter_(std::move(other.tent_device_filter_)), use_tent_(other.use_tent_) { const bool shutdown_enabled = static_cast(other.shutdown_token_); detachShutdownToken(other.shutdown_token_); @@ -444,6 +448,7 @@ TransferEngine& TransferEngine::operator=(TransferEngine&& other) noexcept { freeEngine(); impl_ = std::move(other.impl_); impl_tent_ = std::move(other.impl_tent_); + tent_device_filter_ = std::move(other.tent_device_filter_); use_tent_ = other.use_tent_; const bool shutdown_enabled = static_cast(other.shutdown_token_); detachShutdownToken(other.shutdown_token_); @@ -478,6 +483,26 @@ static std::pair parseConnectionStringInternal( return result; } +std::shared_ptr TransferEngine::buildTentConfig( + const std::string& metadata_conn_string, + const std::string& local_server_name) const { + auto config = std::make_shared(); + if (!local_server_name.empty()) + config->set("local_segment_name", local_server_name); + if (metadata_conn_string == P2PHANDSHAKE) { + config->set("metadata_type", "p2p"); + } else { + auto [type, servers] = + parseConnectionStringInternal(metadata_conn_string); + if (!type.empty()) config->set("metadata_type", type); + if (!servers.empty()) config->set("metadata_servers", servers); + } + if (!tent_device_filter_.empty()) { + config->set("topology/rdma_whitelist", tent_device_filter_); + } + return config; +} + int TransferEngine::init(const std::string& metadata_conn_string, const std::string& local_server_name, const std::string& ip_or_host_name, @@ -494,17 +519,7 @@ int TransferEngine::init(const std::string& metadata_conn_string, return impl_->init(metadata_conn_string, local_server_name, ip_or_host_name, rpc_port); } else { - auto config = std::make_shared(); - if (!local_server_name.empty()) - config->set("local_segment_name", local_server_name); - if (metadata_conn_string == P2PHANDSHAKE) { - config->set("metadata_type", "p2p"); - } else { - auto [type, servers] = - parseConnectionStringInternal(metadata_conn_string); - if (!type.empty()) config->set("metadata_type", type); - if (!servers.empty()) config->set("metadata_servers", servers); - } + auto config = buildTentConfig(metadata_conn_string, local_server_name); if (protocol == "tcp") { mooncake::tent::ConfigHelper::forceTcp(*config); if (!std::getenv("MC_FORCE_TCP")) { @@ -908,7 +923,13 @@ void TransferEngine::setAutoDiscover(const AutoDiscoverConfig& config) { } void TransferEngine::setWhitelistFilters(std::vector&& filters) { - if (!use_tent_) impl_->setWhitelistFilters(std::move(filters)); + if (!use_tent_) { + impl_->setWhitelistFilters(std::move(filters)); + } else if (!impl_tent_) { + tent_device_filter_ = std::move(filters); + } else { + LOG(WARNING) << "Cannot change the TENT RDMA device filter after init"; + } } int TransferEngine::numContexts() const { diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index ea04769ef5..10014b031d 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -95,6 +95,7 @@ struct PreservedTentConfigOverrides { std::optional rpc_server_hostname; std::optional rpc_server_port; bool force_tcp{false}; + std::optional> rdma_whitelist; }; template @@ -219,6 +220,8 @@ PreservedTentConfigOverrides captureExplicitTransferEngineConfig( preserved.force_tcp = config.get("transports/force_tcp", false); preserved.rpc_server_port = captureExplicitConfigValue(config, "rpc_server_port", json()); + preserved.rdma_whitelist = captureExplicitConfigValue( + config, "topology/rdma_whitelist", std::vector()); return preserved; } @@ -245,6 +248,8 @@ void restoreExplicitTransferEngineConfig( if (preserved.force_tcp) { ConfigHelper::forceTcp(config); } + restoreExplicitConfigValue(config, "topology/rdma_whitelist", + preserved.rdma_whitelist); } TransferEngineImpl::TransferEngineImpl() diff --git a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp index 28b66f1a5d..30829297b6 100644 --- a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp @@ -411,6 +411,45 @@ TEST(TransferEngineConfigOverrideTest, CustomTopoJsonEnvLoadsPath) { "/tmp/mooncake-nic-priority-matrix.json"); } +TEST(TransferEngineConfigOverrideTest, + ExplicitRdmaWhitelistOverridesLegacyFilterEnv) { + EnvVarGuard guard("MC_TE_FILTERS", "mlx5_from_env_0,mlx5_from_env_1"); + + auto config = std::make_shared(); + const std::vector explicit_filter{"mlx5_requested"}; + config->set("topology/rdma_whitelist", explicit_filter); + config->set("rpc_server_hostname", kInvalidHostname); + + TransferEngineImpl engine(config); + + EXPECT_FALSE(engine.available()); + EXPECT_EQ(config->getArray("topology/rdma_whitelist"), + explicit_filter); +} + +TEST(TransferEngineConfigOverrideTest, + ExplicitRdmaWhitelistOverridesMcTentConf) { + TempConfigFile conf_file(R"({ + "topology": { + "rdma_whitelist": ["mlx5_from_env_0", "mlx5_from_env_1"] + } + })"); + EnvVarGuard guard("MC_TENT_CONF", conf_file.path()); + + auto config = std::make_shared(); + const std::vector explicit_filter{"mlx5_requested"}; + config->set("topology/rdma_whitelist", explicit_filter); + // Stop construction before platform probing; this test only needs the + // constructor's config merge and remains hardware-independent. + config->set("rpc_server_hostname", kInvalidHostname); + + TransferEngineImpl engine(config); + + EXPECT_FALSE(engine.available()); + EXPECT_EQ(config->getArray("topology/rdma_whitelist"), + explicit_filter); +} + TEST(TransferEngineConfigOverrideTest, ExplicitMetadataOverridesDriveSuccessfulHttpInitialization) { #ifdef _WIN32 diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index bff35ab2a8..10f47c2cf0 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -18,6 +18,8 @@ add_executable(transport_uint_test ${WORKSPACE}/transport_uint_test.cpp) target_link_libraries(transport_uint_test PUBLIC transfer_engine gtest gtest_main) if(USE_TENT) + # Match the compatibility shim public class definition and enable the + # MC_USE_TENT regression cases in transport_uint_test.cpp. target_compile_definitions(transport_uint_test PRIVATE USE_TENT) endif() add_test(NAME transport_uint_test COMMAND transport_uint_test) diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index 1d122c6ba5..a657cd8525 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -16,6 +16,9 @@ #include #include #include +#ifdef USE_TENT +#include +#endif #include #include @@ -119,9 +122,104 @@ class TransferEngineImplTestPeer { if (!engine.impl_tent_ || !engine.impl_tent_->impl_) return nullptr; return engine.impl_tent_->impl_->conf_.get(); } + + static std::shared_ptr buildTentConfig( + const TransferEngine& engine, const std::string& metadata, + const std::string& segment) { + return engine.buildTentConfig(metadata, segment); + } #endif }; +#ifdef USE_TENT +class ScopedUnsetEnvVar { + public: + explicit ScopedUnsetEnvVar(const char* name) : name_(name) { + if (const char* old = std::getenv(name)) old_value_ = old; + unsetenv(name); + } + + ~ScopedUnsetEnvVar() { + if (old_value_.has_value()) + setenv(name_.c_str(), old_value_->c_str(), 1); + } + + private: + std::string name_; + std::optional old_value_; +}; + +TEST(TransferEngineTentCompatibilityTest, + ConstructorDeviceFilterIsForwardedToTentConfig) { + ScopedEnvVar use_tent("MC_USE_TENT", "1"); + const std::vector filter{"mlx5_0", "mlx5_2"}; + TransferEngine engine(/*auto_discover=*/true, filter); + ASSERT_TRUE(engine.isUsingTent()); + + auto config = TransferEngineImplTestPeer::buildTentConfig( + engine, P2PHANDSHAKE, "local-segment"); + + EXPECT_EQ(config->getArray("topology/rdma_whitelist"), filter); +} + +TEST(TransferEngineTentCompatibilityTest, + SetterDeviceFilterIsForwardedBeforeInit) { + ScopedEnvVar use_tent("MC_USE_TENT", "1"); + TransferEngine engine(/*auto_discover=*/true); + engine.setWhitelistFilters({"mlx5_1"}); + + auto config = TransferEngineImplTestPeer::buildTentConfig( + engine, P2PHANDSHAKE, "local-segment"); + + EXPECT_EQ(config->getArray("topology/rdma_whitelist"), + (std::vector{"mlx5_1"})); +} + +TEST(TransferEngineTentCompatibilityTest, + DeviceFilterSurvivesMoveConstructionAndAssignment) { + ScopedEnvVar use_tent("MC_USE_TENT", "1"); + const std::vector filter{"mlx5_move"}; + TransferEngine source(/*auto_discover=*/true, filter); + TransferEngine moved(std::move(source)); + + auto moved_config = TransferEngineImplTestPeer::buildTentConfig( + moved, P2PHANDSHAKE, "local-segment"); + EXPECT_EQ(moved_config->getArray("topology/rdma_whitelist"), + filter); + + TransferEngine assigned(/*auto_discover=*/true); + assigned = std::move(moved); + auto assigned_config = TransferEngineImplTestPeer::buildTentConfig( + assigned, P2PHANDSHAKE, "local-segment"); + EXPECT_EQ(assigned_config->getArray("topology/rdma_whitelist"), + filter); +} + +TEST(TransferEngineTentCompatibilityTest, + ConstructorDeviceFilterRestrictsDiscoveredTopology) { + int count = 0; + ibv_device** devices = ibv_get_device_list(&count); + if (!devices || count < 2) { + if (devices) ibv_free_device_list(devices); + GTEST_SKIP() << "Requires at least two RDMA devices"; + } + const std::string selected = ibv_get_device_name(devices[0]); + ibv_free_device_list(devices); + + ScopedEnvVar use_tent("MC_USE_TENT", "1"); + ScopedEnvVar hostname("MOONCAKE_LOCAL_HOSTNAME", "127.0.0.1"); + ScopedUnsetEnvVar tent_conf("MC_TENT_CONF"); + ScopedUnsetEnvVar custom_topology("MC_CUSTOM_TOPO_JSON"); + TransferEngine engine(/*auto_discover=*/true, {selected}); + ASSERT_EQ(engine.init(P2PHANDSHAKE, ""), 0); + + const auto topology = engine.getLocalTopology(); + ASSERT_NE(topology, nullptr); + EXPECT_EQ(topology->getHcaList(), (std::vector{selected})); +} + +#endif + TEST(TransferEngineAutoDiscoverTest, SelectsEfaForEfaProtocol) { TransferEngineImpl engine(false); engine.setAutoDiscover({.enabled = true, .protocol = "efa"});