Skip to content

vllm.distributed.nixl_utils

Functions:

  • alias_nixl_for_ray –

    Let Ray's nixl._api import resolve to the ROCm implementation.

  • is_nixl_available –

    Lightweight check for the platform's NIXL package without importing it.

alias_nixl_for_ray()

Let Ray's nixl._api import resolve to the ROCm implementation.

Source code in vllm/distributed/nixl_utils.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def alias_nixl_for_ray() -> None:
    """Let Ray's ``nixl._api`` import resolve to the ROCm implementation."""
    if not current_platform.is_rocm() or "nixl._api" in sys.modules:
        return

    package_name = _get_nixl_package_name()
    if not is_nixl_available():
        raise ImportError(f"Ray's NIXL transport on ROCm requires {package_name}")

    _maybe_set_ucx_rcache_limit()
    package = importlib.import_module(package_name)
    api = importlib.import_module(f"{package_name}._api")
    sys.modules.setdefault("nixl", package)
    sys.modules.setdefault("nixl._api", api)

is_nixl_available()

Lightweight check for the platform's NIXL package without importing it.

Source code in vllm/distributed/nixl_utils.py 
86
87
88
89
90
91
def is_nixl_available() -> bool:
    """Lightweight check for the platform's NIXL package without importing it."""
    import importlib.util

    pkg = _get_nixl_package_name()
    return pkg in sys.modules or importlib.util.find_spec(pkg) is not None