Nick HillandIlia Yastrebov
eaf25f675a
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Ilia Yastrebov <iyastrebov@nvidia.com>
137 lines
4.0 KiB
Python
137 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import functools
|
|
import importlib
|
|
import os
|
|
import sys
|
|
from typing import Any
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.platforms import current_platform
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# declaration for static analyzers
|
|
NixlWrapper: Any
|
|
nixl_agent_config: Any
|
|
nixlXferTelemetry: Any
|
|
|
|
|
|
def _maybe_set_ucx_rcache_limit() -> None:
|
|
if "UCX_RCACHE_MAX_UNRELEASED" in os.environ:
|
|
return
|
|
|
|
if "nixl" in sys.modules or "nixl_rocm" in sys.modules:
|
|
logger.warning_once(
|
|
"NIXL was already imported, we can't reset "
|
|
"UCX_RCACHE_MAX_UNRELEASED. "
|
|
"Please set it to '1024' manually."
|
|
)
|
|
return
|
|
|
|
logger.info_once(
|
|
"Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare "
|
|
"memory leak in UCX when using NIXL."
|
|
)
|
|
os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
|
|
|
|
|
|
def _get_nixl_package_name() -> str:
|
|
return "nixl_rocm" if current_platform.is_rocm() else "nixl"
|
|
|
|
|
|
def _get_nixl_module_name(name: str) -> str:
|
|
package_name = _get_nixl_package_name()
|
|
if name == "nixlXferTelemetry":
|
|
return f"{package_name}._bindings"
|
|
return f"{package_name}._api"
|
|
|
|
|
|
def _load_nixl_attr(name: str) -> Any:
|
|
attr_name = {
|
|
"NixlWrapper": "nixl_agent",
|
|
"nixl_agent_config": "nixl_agent_config",
|
|
"nixlXferTelemetry": "nixlXferTelemetry",
|
|
}[name]
|
|
|
|
_maybe_set_ucx_rcache_limit()
|
|
try:
|
|
module = importlib.import_module(_get_nixl_module_name(name))
|
|
except ImportError:
|
|
if name == "NixlWrapper":
|
|
logger.warning_once("NIXL is not available")
|
|
elif name == "nixl_agent_config":
|
|
logger.warning_once("NIXL agent config is not available")
|
|
value = None
|
|
else:
|
|
value = getattr(module, attr_name, None)
|
|
if name == "NixlWrapper":
|
|
if value is None:
|
|
logger.warning_once("NIXL is not available")
|
|
else:
|
|
logger.info_once("NIXL is available")
|
|
elif name == "nixl_agent_config" and value is None:
|
|
logger.warning_once("NIXL agent config is not available")
|
|
|
|
globals()[name] = value
|
|
return value
|
|
|
|
|
|
def __getattr__(name: str) -> Any:
|
|
if name in __all__:
|
|
return _load_nixl_attr(name)
|
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
|
|
|
|
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
|
|
|
|
|
|
@functools.cache
|
|
def check_nixl_version(min_version: str) -> None:
|
|
"""Raise if the installed NIXL package is older than ``min_version``."""
|
|
from importlib.metadata import packages_distributions, version
|
|
|
|
from packaging.version import Version
|
|
|
|
pkg = _get_nixl_package_name()
|
|
for dist in packages_distributions().get(pkg, ()):
|
|
installed = version(dist)
|
|
if Version(installed) < Version(min_version):
|
|
raise RuntimeError(
|
|
f"NIXL >= {min_version} is required, but {dist}=={installed} "
|
|
"is installed"
|
|
)
|
|
return
|
|
|
|
|
|
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)
|
|
|
|
|
|
__all__ = [
|
|
"NixlWrapper",
|
|
"nixl_agent_config",
|
|
"nixlXferTelemetry",
|
|
"is_nixl_available",
|
|
"check_nixl_version",
|
|
"alias_nixl_for_ray",
|
|
]
|