xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-25 08:04:26 +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:39:45 +00:00
parent 6cd8dec23f
commit c3f7afa6a8
1 changed files with 7 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -42,7 +42,9 @@ 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 = 4
import os
VLLM_DEBUG_NIXL_XFER_TIME = os.getenv("VLLM_DEBUG_NIXL_XFER_TIME", "0") == "1"
logger = init_logger(__name__) logger = init_logger(__name__)
@ -50,16 +52,11 @@ logger = init_logger(__name__)
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_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_VERSION = None
NIXL_NUM_WORKERS
class NixlAgentMetadata( class NixlAgentMetadata(
msgspec.Struct, msgspec.Struct,
omit_defaults=True, # type: ignore[call-arg] omit_defaults=True, # type: ignore[call-arg]
@ -368,7 +365,7 @@ class NixlConnectorWorker:
# Agent. # Agent.
import os import os
NIXL_NUM_WORKERS = int(os.getenv("VLLM_NIXL_NUM_WORKERS", "1")) NIXL_NUM_WORKERS = int(os.getenv("VLLM_NIXL_NUM_WORKERS", "8"))
logger.info(f"Using NIXL_NUM_WORKERS={NIXL_NUM_WORKERS} for NIXL agent.") logger.info(f"Using NIXL_NUM_WORKERS={NIXL_NUM_WORKERS} for NIXL agent.")
config = nixl_agent_config(enable_prog_thread=False, num_threads=NIXL_NUM_WORKERS) config = nixl_agent_config(enable_prog_thread=False, num_threads=NIXL_NUM_WORKERS)
@ -1038,7 +1035,9 @@ class NixlConnectorWorker:
start = time.perf_counter() start = time.perf_counter()
self.nixl_wrapper.transfer(handle) self.nixl_wrapper.transfer(handle)
end = time.perf_counter() end = time.perf_counter()
logger.info(f"TIME: {end - start}") if VLLM_DEBUG_NIXL_XFER_TIME:
# Log the time taken for the transfer.
logger.info(f"TIME: {end - start}")
# Use handle to check completion in future step(). # Use handle to check completion in future step().
# TODO (NickLucche) surface xfer elapsed time # TODO (NickLucche) surface xfer elapsed time