diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py index dcfa28210f81..0b6a52d93de0 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py @@ -1,14 +1,33 @@ +import importlib import time from enum import Enum -from nixl import nixl_agent, nixl_agent_config, nixl_xfer_handle - from tensorrt_llm._utils import nvtx_range from tensorrt_llm.logger import logger # Import base classes for type compatibility from ..base.agent import BaseTransferAgent, RegMemoryDescs, TransferRequest, TransferStatus +# The PyPI ``nixl`` meta-package provides a top-level ``nixl`` shim that +# redirects to ``nixl_cu13`` or ``nixl_cu12``. When only ``nixl-cu13`` (or +# ``nixl-cu12``) is installed without the meta-package, ``import nixl`` fails. +# Try the CUDA-specific packages first, then fall back to the meta-package. +_nixl_mod = None +for _candidate in ("nixl_cu13", "nixl_cu12", "nixl"): + try: + _nixl_mod = importlib.import_module(_candidate) + break + except ImportError: + continue +if _nixl_mod is None: + raise ImportError( + "Could not find a NIXL Python package. " + "Install ``nixl-cu13`` (CUDA 13) or ``nixl-cu12`` (CUDA 12)." + ) +nixl_agent = _nixl_mod.nixl_agent +nixl_agent_config = _nixl_mod.nixl_agent_config +nixl_xfer_handle = _nixl_mod.nixl_xfer_handle + class TransferState(Enum): PENDING = "PENDING" diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 1c7cd90dc257..a2f897f2e29a 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -351,7 +351,6 @@ unittest/auto_deploy/singlegpu/models/test_qwen3_5_moe.py::test_vlm_wrapper_delt unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py::test_build_ad[deepseek-ai/DeepSeek-V3-llm_extra_args10] SKIP (https://nvbugs/5888827) unittest/auto_deploy/standalone/test_standalone_package.py::TestStandalonePackage::test_run_unit_tests SKIP (https://nvbugs/6160629) unittest/bindings/test_transfer_agent_bindings.py::TestNixlFunctionalTransfer::test_nixl_wait_in_progress_on_zero_timeout SKIP (https://nvbugs/6260897) -unittest/disaggregated/test_agent_multi_backends.py::test_run_with_different_env[1] SKIP (https://nvbugs/5979673) 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)