diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 80e697faf3be..c02464b7ecfb 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -363,9 +363,6 @@ unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py::test_capture_smoke S unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[fp8-2-1560] SKIP (https://nvbugs/6198760) unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[int8-2-1560] SKIP (https://nvbugs/6198760) unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) -unittest/disaggregated/test_kv_transfer.py SKIP (https://nvbugs/6403793) -unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp1_pp4_to_tp2_pp2] SKIP (https://nvbugs/6445316) -unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834) unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) diff --git a/tests/unittest/disaggregated/test_agent.py b/tests/unittest/disaggregated/test_agent.py index f2212a6c5f82..c8425ab49e18 100644 --- a/tests/unittest/disaggregated/test_agent.py +++ b/tests/unittest/disaggregated/test_agent.py @@ -6,8 +6,10 @@ import pytest import torch -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" from tensorrt_llm import logger from tensorrt_llm._torch.disaggregation.base.agent import ( diff --git a/tests/unittest/disaggregated/test_agent_multi_backends.py b/tests/unittest/disaggregated/test_agent_multi_backends.py index d983af152542..d82815c99a79 100644 --- a/tests/unittest/disaggregated/test_agent_multi_backends.py +++ b/tests/unittest/disaggregated/test_agent_multi_backends.py @@ -3,8 +3,11 @@ import pytest -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. The subprocesses +# spawned below inherit these via os.environ.copy(). +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" def test_load_agent_missing_module(): diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 651a347b2f08..26b9e2e9199d 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -28,7 +28,10 @@ # Exclude UCX IB transport (avoid NIXL setup hangs without IB) and gdr_copy # (avoid SIGSEGV at process exit from UCX rcache cleanup; gdr_copy disabled # falls back to cuda_ipc / cuda_copy without affecting correctness). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" from dataclasses import dataclass from typing import Dict, List, Optional @@ -55,9 +58,6 @@ AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType -# Reduce NIXL threads for unit test: default 8 threads per agent causes heavy -# contention when creating multiple agents on a single GPU in the same process. -os.environ.setdefault("TRTLLM_NIXL_NUM_THREADS", "0") # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 76a0d8db25b3..c2cfc1382e37 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -5,8 +5,19 @@ import time import uuid -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX config regardless of what the cluster/CI injects +# (the CI agent bootstrap exports UCX_TLS=tcp,cuda_copy,cuda_ipc before pytest +# starts, which a setdefault would leave in place): exclude IB (no fabric +# assumed) and gdr_copy (UCX rcache SIGABRT at teardown). +os.environ["UCX_TLS"] = "^ib,gdr_copy" +# Each NIXL agent spawns TRTLLM_NIXL_NUM_THREADS (default 8) busy-polling +# progress threads, and a single case builds up to 8 TransferWorkers (one per +# rank). On CI nodes shared with other single-GPU jobs the resulting CPU +# oversubscription inflates agent construction from ~3s to ~30s each, blowing +# the 120s per-test timeout intermittently (https://nvbugs/6426834). One +# progress thread is enough here: these tests verify transfer logic, not +# transfer-engine threading. +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" from dataclasses import dataclass from typing import List, Optional diff --git a/tests/unittest/disaggregated/test_kv_transfer_mp.py b/tests/unittest/disaggregated/test_kv_transfer_mp.py index 76e5f4aa7573..a52d95bde7a4 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_mp.py +++ b/tests/unittest/disaggregated/test_kv_transfer_mp.py @@ -8,7 +8,13 @@ import torch.multiprocessing as mp # Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX config regardless of what the cluster/CI injects; +# see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +# Limit NIXL busy-polling progress threads; see test_kv_transfer.py for the +# full rationale (intermittent 120s timeouts on shared CI nodes, +# https://nvbugs/6426834). +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" import tensorrt_llm import tensorrt_llm.bindings diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 0cf0b3209899..8688d75f9a53 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import os import threading import uuid from typing import Dict, List @@ -19,6 +20,11 @@ import pytest import torch +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" + import tensorrt_llm import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 diff --git a/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py b/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py index a81e05026cbb..d414d18fb48b 100644 --- a/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py +++ b/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py @@ -15,8 +15,10 @@ import torch.distributed as dist import torch.multiprocessing as mp -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" import tensorrt_llm import tensorrt_llm.bindings diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 0d6b31e367d1..241a4aee0d09 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -3,6 +3,7 @@ import gc import multiprocessing +import os import sys import time import uuid @@ -11,6 +12,11 @@ import pytest import torch +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see disaggregated/test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" + import tensorrt_llm import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm