| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import importlib |
| 5 | import os |
| 6 | import sys |
| 7 | from typing import Any |
| 8 | |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.platforms import current_platform |
| 11 | |
| 12 | logger = init_logger(__name__) |
| 13 | |
| 14 | # declaration for static analyzers |
| 15 | NixlWrapper: Any |
| 16 | nixl_agent_config: Any |
| 17 | nixlXferTelemetry: Any |
| 18 | |
| 19 | |
| 20 | def _maybe_set_ucx_rcache_limit() -> None: |
| 21 | if "UCX_RCACHE_MAX_UNRELEASED" in os.environ: |
| 22 | return |
| 23 | |
| 24 | if "nixl" in sys.modules or "nixl_rocm" in sys.modules: |
| 25 | logger.warning_once( |
| 26 | "NIXL was already imported, we can't reset " |
| 27 | "UCX_RCACHE_MAX_UNRELEASED. " |
| 28 | "Please set it to '1024' manually." |
| 29 | ) |
| 30 | return |
| 31 | |
| 32 | logger.info_once( |
| 33 | "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare " |
| 34 | "memory leak in UCX when using NIXL." |
| 35 | ) |
| 36 | os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" |
| 37 | |
| 38 | |
| 39 | def _get_nixl_package_name() -> str: |
| 40 | return "nixl_rocm" if current_platform.is_rocm() else "nixl" |
| 41 | |
| 42 | |
| 43 | def _get_nixl_module_name(name: str) -> str: |
| 44 | package_name = _get_nixl_package_name() |
| 45 | if name == "nixlXferTelemetry": |
| 46 | return f"{package_name}._bindings" |
| 47 | return f"{package_name}._api" |
| 48 | |
| 49 | |
| 50 | def _load_nixl_attr(name: str) -> Any: |
| 51 | attr_name = { |
| 52 | "NixlWrapper": "nixl_agent", |
| 53 | "nixl_agent_config": "nixl_agent_config", |
| 54 | "nixlXferTelemetry": "nixlXferTelemetry", |
| 55 | }[name] |
| 56 | |
| 57 | _maybe_set_ucx_rcache_limit() |
| 58 | try: |
| 59 | module = importlib.import_module(_get_nixl_module_name(name)) |
| 60 | except ImportError: |
| 61 | if name == "NixlWrapper": |
| 62 | logger.warning_once("NIXL is not available") |
| 63 | elif name == "nixl_agent_config": |
| 64 | logger.warning_once("NIXL agent config is not available") |
| 65 | value = None |
| 66 | else: |
| 67 | value = getattr(module, attr_name, None) |
| 68 | if name == "NixlWrapper": |
| 69 | if value is None: |
| 70 | logger.warning_once("NIXL is not available") |
| 71 | else: |
| 72 | logger.info_once("NIXL is available") |
| 73 | elif name == "nixl_agent_config" and value is None: |
| 74 | logger.warning_once("NIXL agent config is not available") |
| 75 | |
| 76 | globals()[name] = value |
| 77 | return value |
| 78 | |
| 79 | |
| 80 | def __getattr__(name: str) -> Any: |
| 81 | if name in __all__: |
| 82 | return _load_nixl_attr(name) |
| 83 | raise AttributeError(f"module {__name__!r} has no attribute {name!r}") |
| 84 | |
| 85 | |
| 86 | def is_nixl_available() -> bool: |
| 87 | """Lightweight check for the platform's NIXL package without importing it.""" |
| 88 | import importlib.util |
| 89 | |
| 90 | pkg = _get_nixl_package_name() |
| 91 | return pkg in sys.modules or importlib.util.find_spec(pkg) is not None |
| 92 | |
| 93 | |
| 94 | def alias_nixl_for_ray() -> None: |
| 95 | """Let Ray's ``nixl._api`` import resolve to the ROCm implementation.""" |
| 96 | if not current_platform.is_rocm() or "nixl._api" in sys.modules: |
| 97 | return |
| 98 | |
| 99 | package_name = _get_nixl_package_name() |
| 100 | if not is_nixl_available(): |
| 101 | raise ImportError(f"Ray's NIXL transport on ROCm requires {package_name}") |
| 102 | |
| 103 | _maybe_set_ucx_rcache_limit() |
| 104 | package = importlib.import_module(package_name) |
| 105 | api = importlib.import_module(f"{package_name}._api") |
| 106 | sys.modules.setdefault("nixl", package) |
| 107 | sys.modules.setdefault("nixl._api", api) |
| 108 | |
| 109 | |
| 110 | __all__ = [ |
| 111 | "NixlWrapper", |
| 112 | "nixl_agent_config", |
| 113 | "nixlXferTelemetry", |
| 114 | "is_nixl_available", |
| 115 | "alias_nixl_for_ray", |
| 116 | ] |