[https://nvbugs/6426834][fix] Deflake test_kv_transfer: cap NIXL progress threads, pin UCX env

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)

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 (

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():

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
 # ---------------------------------------------------------------------------

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
 

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

tests/unittest/disaggregated/test_mamba_transfer.py

@@ -12,13 +12,19 @@
 # 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
 
 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

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

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