[https://nvbugs/6484986][test] add UCX cancellation regression coverage

cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");

cpp/tests/CMakeLists.txt

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION &
 # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
@@ -39,6 +39,7 @@ add_custom_target(google-tests)
 
 function(add_gtest test_name test_src)
   set(options NO_GTEST_MAIN NO_TLLM_LINKAGE)
+  set(oneValueArgs TIMEOUT)
   cmake_parse_arguments(ARGS "${options}" "${oneValueArgs}" "${multiValueArgs}"
                         ${ARGN})
   add_executable(${test_name} ${test_src})
@@ -61,12 +62,19 @@ function(add_gtest test_name test_src)
   target_compile_definitions(${test_name}
                              PUBLIC TOP_LEVEL_DIR="${TOP_LEVEL_DIR}")
 
+  set(test_properties ENVIRONMENT "CUDA_MODULE_LOADING=LAZY")
+  if(ARGS_TIMEOUT)
+    list(APPEND test_properties TIMEOUT "${ARGS_TIMEOUT}")
+  endif()
+
   gtest_discover_tests(
     ${test_name}
-    PROPERTIES ENVIRONMENT "CUDA_MODULE_LOADING=LAZY" DISCOVERY_MODE
+    PROPERTIES ${test_properties}
+               DISCOVERY_MODE
                PRE_TEST # WAR for DLL discovery on windows.
-               DISCOVERY_TIMEOUT 30) # Longer timeout needed because discovery
-                                     # can be slow on Windows
+               DISCOVERY_TIMEOUT
+               30) # Longer timeout needed because discovery can be slow on
+                   # Windows
   add_dependencies(google-tests ${test_name})
 endfunction()
 

cpp/tests/unit_tests/executor/CMakeLists.txt

@@ -27,7 +27,7 @@ add_gtest(loraConfigTest loraConfigTest.cpp)
 add_gtest(coalesceTest coalesceTest.cpp)
 add_gtest(genUniqueAgentNameTest genUniqueAgentNameTest.cpp)
 target_link_libraries(genUniqueAgentNameTest PRIVATE ${Python3_LIBRARIES})
-add_gtest(ucxCommTest ucxCommTest.cpp)
+add_gtest(ucxCommTest ucxCommTest.cpp TIMEOUT 60)
 target_link_libraries(ucxCommTest PRIVATE ${Python3_LIBRARIES})
 target_link_libraries(serializeUtilsTest PRIVATE ${Python3_LIBRARIES})
 

cpp/tests/unit_tests/executor/ucxCommTest.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -43,11 +43,14 @@
 #include "tensorrt_llm/runtime/common.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 #include "gtest/gtest.h"
+#include <atomic>
+#include <chrono>
 #include <csignal>
 #include <cstddef>
 #include <cstdint>
 #include <cstdio>
 #include <cstdlib>
+#include <future>
 #include <gmock/gmock.h>
 #include <memory>
 #include <random>
@@ -91,6 +94,42 @@ class UcxCommTest : public ::testing::Test
 using DataContext = tensorrt_llm::executor::kv_cache::DataContext;
 using TransceiverTag = tensorrt_llm::batch_manager::TransceiverTag;
 
+TEST_F(UcxCommTest, recvConnectCancellation)
+{
+    try
+    {
+        auto connectionManager = makeOneUcxConnectionManager();
+        ASSERT_NE(connectionManager, nullptr);
+
+        std::atomic<bool> transferTerminate{false};
+        TransceiverTag::Id id;
+        auto receiveFuture = std::async(std::launch::async,
+            [&]() {
+                return connectionManager->recvConnect(
+                    DataContext{TransceiverTag::kID_TAG, transferTerminate}, &id, sizeof(id));
+            });
+
+        constexpr auto kReceiveStartPeriod = std::chrono::milliseconds{100};
+        constexpr auto kCancellationTimeout = std::chrono::seconds{5};
+        ASSERT_EQ(receiveFuture.wait_for(kReceiveStartPeriod), std::future_status::timeout);
+
+        transferTerminate.store(true, std::memory_order_relaxed);
+        ASSERT_EQ(receiveFuture.wait_for(kCancellationTimeout), std::future_status::ready);
+        EXPECT_EQ(receiveFuture.get(), nullptr);
+    }
+    catch (std::exception const& e)
+    {
+        std::string error = e.what();
+        if (error.find("UCX wrapper library is not open correctly") != std::string::npos
+            || error.find("Unable to load UCX wrapper library symbol") != std::string::npos)
+        {
+            GTEST_SKIP() << "UCX wrapper library is not open correctly. Skip this test case.";
+        }
+
+        throw e;
+    }
+}
+
 TEST_F(UcxCommTest, Basic)
 {
 

tests/integration/test_lists/waives.txt

@@ -389,9 +389,7 @@ unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_2
 unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_tp2pp2 SKIP (https://nvbugs/6427411)
 unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504)
 unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504)
-unittest/llmapi/test_llm_pytorch.py::test_llm_context_only_timed_out[None] SKIP (https://nvbugs/6484986)
 unittest/llmapi/test_llm_pytorch.py::test_llm_context_only_timed_out_kv_cache_exhausted[None-UCX-1000] SKIP (https://nvbugs/6490004)
-unittest/llmapi/test_llm_pytorch.py::test_llm_context_only_timed_out_kv_cache_exhausted[None-UCX-100] SKIP (https://nvbugs/6488811)
 unittest/llmapi/test_memory_profiling.py::test_profile_kvcache SKIP (https://nvbugs/5580781)
 unittest/tools/test_layer_wise_benchmarks.py::test_performance_alignment[1] SKIP (https://nvbugs/6487836)
 unittest/tools/test_test_to_stage_mapping.py::test_cli_functionality[txt] SKIP (https://nvbugs/6482297)