[Bugfix][TENT] Forward classic RDMA device filters

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<mooncake::tent::Config>();
         conf->set("metadata_type", "p2p");
         if (!device_name_safe.empty()) {

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<mooncake::tent::Config> buildTentConfig(
+        const std::string& metadata_conn_string,
+        const std::string& local_server_name) const;
+
     std::shared_ptr<TransferEngineImpl> impl_;
     std::shared_ptr<mooncake::tent::TransferEngine> impl_tent_;
     std::shared_ptr<ShutdownToken> shutdown_token_;
+    // Classic callers provide this through TransferEngine(auto_discover,
+    // filter) before init() creates the native TENT engine.
+    std::vector<std::string> tent_device_filter_;
     bool use_tent_{false};
     friend class TransferEngineImplTestPeer;
 };

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<bool>(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<TransferEngineImpl>(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<bool>(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<bool>(other.shutdown_token_);
     detachShutdownToken(other.shutdown_token_);
@@ -478,6 +483,26 @@ static std::pair<std::string, std::string> parseConnectionStringInternal(
     return result;
 }
 
+std::shared_ptr<mooncake::tent::Config> TransferEngine::buildTentConfig(
+    const std::string& metadata_conn_string,
+    const std::string& local_server_name) const {
+    auto config = std::make_shared<mooncake::tent::Config>();
+    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<mooncake::tent::Config>();
-        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<std::string>&& 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 {

mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp

@@ -95,6 +95,7 @@ struct PreservedTentConfigOverrides {
     std::optional<std::string> rpc_server_hostname;
     std::optional<json> rpc_server_port;
     bool force_tcp{false};
+    std::optional<std::vector<std::string>> rdma_whitelist;
 };
 
 template <typename T>
@@ -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<std::string>());
     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()

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<Config>();
+    const std::vector<std::string> 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<std::string>("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<Config>();
+    const std::vector<std::string> 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<std::string>("topology/rdma_whitelist"),
+              explicit_filter);
+}
+
 TEST(TransferEngineConfigOverrideTest,
      ExplicitMetadataOverridesDriveSuccessfulHttpInitialization) {
 #ifdef _WIN32

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)

mooncake-transfer-engine/tests/transport_uint_test.cpp

@@ -16,6 +16,9 @@
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 #include <sys/time.h>
+#ifdef USE_TENT
+#include <infiniband/verbs.h>
+#endif
 
 #include <algorithm>
 #include <array>
@@ -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<tent::Config> 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<std::string> old_value_;
+};
+
+TEST(TransferEngineTentCompatibilityTest,
+     ConstructorDeviceFilterIsForwardedToTentConfig) {
+    ScopedEnvVar use_tent("MC_USE_TENT", "1");
+    const std::vector<std::string> 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<std::string>("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<std::string>("topology/rdma_whitelist"),
+              (std::vector<std::string>{"mlx5_1"}));
+}
+
+TEST(TransferEngineTentCompatibilityTest,
+     DeviceFilterSurvivesMoveConstructionAndAssignment) {
+    ScopedEnvVar use_tent("MC_USE_TENT", "1");
+    const std::vector<std::string> 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<std::string>("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<std::string>("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<std::string>{selected}));
+}
+
+#endif
+
 TEST(TransferEngineAutoDiscoverTest, SelectsEfaForEfaProtocol) {
     TransferEngineImpl engine(false);
     engine.setAutoDiscover({.enabled = true, .protocol = "efa"});