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/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 99e16ea8b147..2c65dee55abe 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,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() 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}) 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..1659323574cf 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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)