xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-09 10:47:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

updated

Browse Source 
Signed-off-by: Robert Shaw <robshaw@redhat.com>
This commit is contained in:
Robert Shaw 2025-07-16 20:29:24 +00:00
parent 723263fa23
commit 6cd8dec23f
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
View File

@ -10,6 +10,7 @@ from collections import defaultdict
from collections.abc import Iterator from collections.abc import Iterator
from concurrent.futures import Future, ThreadPoolExecutor from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass from dataclasses import dataclass
from importlib import metadata
from typing import TYPE_CHECKING, Any, Optional from typing import TYPE_CHECKING, Any, Optional
import msgspec import msgspec
@ -41,17 +42,23 @@ Transfer = tuple[int, float] # (xfer_handle, start_time)
EngineId = str EngineId = str
ReqId = str ReqId = str
GET_META_MSG = b"get_meta_msg" GET_META_MSG = b"get_meta_msg"
NIXL_NUM_WORKERS = 32 NIXL_NUM_WORKERS = 4
logger = init_logger(__name__) logger = init_logger(__name__)
# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used # Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used
try: try:
from nixl._api import nixl_agent as NixlWrapper, nixl_agent_config from nixl._api import nixl_agent as NixlWrapper, nixl_agent_config
NIXL_VERSION = metadata.version("nixl")
NIXL_major, NIXL_minor, NIXL_patch = map(int, NIXL_VERSION.split("."))
except ImportError: except ImportError:
logger.warning("NIXL is not available") logger.warning("NIXL is not available")
NixlWrapper = None NixlWrapper = None
NIXL_VERSION = None
NIXL_NUM_WORKERS
class NixlAgentMetadata( class NixlAgentMetadata(
msgspec.Struct, msgspec.Struct,
@ -352,9 +359,8 @@ class NixlConnectorWorker:
def __init__(self, vllm_config: VllmConfig, engine_id: str): def __init__(self, vllm_config: VllmConfig, engine_id: str):
if NixlWrapper is None: if NixlWrapper is None:
logger.error("NIXL is not available") logger.error("NIXL is not available")
raise RuntimeError("NIXL is not available") raise RuntimeError("NIXL is not available.")
logger.info("Initializing NIXL wrapper") logger.info("Initializing NIXL v%s: worker %s", NIXL_VERSION, engine_id)
logger.info("Initializing NIXL worker %s", engine_id)
# Config. # Config.
self.vllm_config = vllm_config self.vllm_config = vllm_config
@ -362,10 +368,10 @@ class NixlConnectorWorker:
# Agent. # Agent.
import os import os
if os.getenv("VLLM_USE_NIXL_WORKERS", "0") == "1": NIXL_NUM_WORKERS = int(os.getenv("VLLM_NIXL_NUM_WORKERS", "1"))
config = nixl_agent_config(num_threads=NIXL_NUM_WORKERS) logger.info(f"Using NIXL_NUM_WORKERS={NIXL_NUM_WORKERS} for NIXL agent.")
else:
config = None config = nixl_agent_config(enable_prog_thread=False, num_threads=NIXL_NUM_WORKERS)
self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), config) self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), config)
# Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}. # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.
self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict) self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict)
@ -454,7 +460,8 @@ class NixlConnectorWorker:
def __del__(self): def __del__(self):
"""Cleanup background threads on destruction.""" """Cleanup background threads on destruction."""
self._handshake_initiation_executor.shutdown(wait=False) if t_ := getattr(self, "_handshake_initiation_executor", None):
t_.shutdown(wait=False)
if self._nixl_handshake_listener_t: if self._nixl_handshake_listener_t:
self._nixl_handshake_listener_t.join(timeout=0) self._nixl_handshake_listener_t.join(timeout=0)