From 086f51d7db47db2aee11e43c1b9c17bd6cd5d3be Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:01:04 -0700 Subject: [PATCH 1/4] [https://nvbugs/6484986][fix] cancel pending UCX receive Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../ucx_utils/ucxCacheCommunicator.cpp | 2 +- cpp/tests/unit_tests/executor/ucxCommTest.cpp | 41 ++++++++++++++++++- tests/integration/test_lists/waives.txt | 1 - 3 files changed, 41 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp index 3cd388625640..7b817545b39c 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp +++ b/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"); diff --git a/cpp/tests/unit_tests/executor/ucxCommTest.cpp b/cpp/tests/unit_tests/executor/ucxCommTest.cpp index 5d2bfc6e772f..d3695144cc92 100644 --- a/cpp/tests/unit_tests/executor/ucxCommTest.cpp +++ b/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 +#include #include #include #include #include #include +#include #include #include #include @@ -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 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) { diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6dc4776bb633..170838aa6b2b 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -389,7 +389,6 @@ 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) From 444b4f5f1149edfc8dac8ed9ef66f31360b44a25 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:02:59 -0700 Subject: [PATCH 2/4] [https://nvbugs/6488811][test] unwaive exhausted KV timeout Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 170838aa6b2b..1659323574cf 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -390,7 +390,6 @@ unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_t 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_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) From 0c5ba5164196d7db396ec7c194b1e80e4b24634a Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 22 Jul 2026 22:55:01 -0700 Subject: [PATCH 3/4] [NVBUG 6484986][test] bound UCX cancellation test runtime Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- cpp/tests/CMakeLists.txt | 16 +++++++++++----- cpp/tests/unit_tests/executor/CMakeLists.txt | 2 +- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 99e16ea8b147..52336eb3629f 100644 --- a/cpp/tests/CMakeLists.txt +++ b/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,17 @@ 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 - PRE_TEST # WAR for DLL discovery on windows. - DISCOVERY_TIMEOUT 30) # Longer timeout needed because discovery - # can be slow on Windows + 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 add_dependencies(google-tests ${test_name}) endfunction() diff --git a/cpp/tests/unit_tests/executor/CMakeLists.txt b/cpp/tests/unit_tests/executor/CMakeLists.txt index 2275683879a4..6d851ebd4544 100644 --- a/cpp/tests/unit_tests/executor/CMakeLists.txt +++ b/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}) From 0fded855f041c1d6c9eee53b7c2d7eaf673318b5 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 23 Jul 2026 08:01:31 -0700 Subject: [PATCH 4/4] [https://nvbugs/6484986][fix] format CMake test timeout Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- cpp/tests/CMakeLists.txt | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 52336eb3629f..2c65dee55abe 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -70,9 +70,11 @@ function(add_gtest test_name test_src) gtest_discover_tests( ${test_name} 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_MODE + PRE_TEST # WAR for DLL discovery on windows. + DISCOVERY_TIMEOUT + 30) # Longer timeout needed because discovery can be slow on + # Windows add_dependencies(google-tests ${test_name}) endfunction()