wehub-resource-sync
827 lines
29 KiB
Python
827 lines
29 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import multiprocessing
|
|
import threading
|
|
import time
|
|
import weakref
|
|
from collections.abc import Callable, Sequence
|
|
from contextlib import AbstractContextManager
|
|
from dataclasses import dataclass
|
|
from multiprocessing import connection
|
|
from multiprocessing.process import BaseProcess
|
|
from multiprocessing.queues import Queue
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Generic,
|
|
TypeVar,
|
|
Union,
|
|
overload,
|
|
)
|
|
|
|
import torch
|
|
import uvloop
|
|
from torch.autograd.profiler import record_function
|
|
|
|
import vllm.envs as envs
|
|
from vllm.logger import init_logger
|
|
from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message
|
|
from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri
|
|
from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title
|
|
from vllm.utils.torch_utils import PIN_MEMORY
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
|
|
if TYPE_CHECKING:
|
|
import numpy as np
|
|
|
|
from vllm.v1.engine.coordinator import DPCoordinator
|
|
from vllm.v1.engine.utils import CoreEngineActorManager, CoreEngineProcManager
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class ConstantList(Generic[T], Sequence):
|
|
def __init__(self, x: list[T]) -> None:
|
|
self._x = x
|
|
|
|
def append(self, item):
|
|
raise TypeError("Cannot append to a constant list")
|
|
|
|
def extend(self, item):
|
|
raise TypeError("Cannot extend a constant list")
|
|
|
|
def insert(self, item):
|
|
raise TypeError("Cannot insert into a constant list")
|
|
|
|
def pop(self, item):
|
|
raise TypeError("Cannot pop from a constant list")
|
|
|
|
def remove(self, item):
|
|
raise TypeError("Cannot remove from a constant list")
|
|
|
|
def clear(self):
|
|
raise TypeError("Cannot clear a constant list")
|
|
|
|
def index(self, item: T, start: int = 0, stop: int | None = None) -> int:
|
|
return self._x.index(item, start, stop if stop is not None else len(self._x))
|
|
|
|
@overload
|
|
def __getitem__(self, item: int) -> T: ...
|
|
|
|
@overload
|
|
def __getitem__(self, s: slice, /) -> list[T]: ...
|
|
|
|
def __getitem__(self, item: int | slice) -> T | list[T]:
|
|
return self._x[item]
|
|
|
|
@overload
|
|
def __setitem__(self, item: int, value: T): ...
|
|
|
|
@overload
|
|
def __setitem__(self, s: slice, value: T, /): ...
|
|
|
|
def __setitem__(self, item: int | slice, value: T | list[T]):
|
|
raise TypeError("Cannot set item in a constant list")
|
|
|
|
def __delitem__(self, item):
|
|
raise TypeError("Cannot delete item from a constant list")
|
|
|
|
def __iter__(self):
|
|
return iter(self._x)
|
|
|
|
def __contains__(self, item):
|
|
return item in self._x
|
|
|
|
def __len__(self):
|
|
return len(self._x)
|
|
|
|
def __repr__(self):
|
|
return f"ConstantList({self._x})"
|
|
|
|
def copy(self) -> list[T]:
|
|
return self._x.copy()
|
|
|
|
|
|
class CpuGpuBuffer:
|
|
"""Buffer to easily copy tensors between CPU and GPU."""
|
|
|
|
def __init__(
|
|
self,
|
|
*size: int | torch.SymInt,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
pin_memory: bool = PIN_MEMORY,
|
|
with_numpy: bool = True,
|
|
) -> None:
|
|
# these buffers are mutable runtime state, so allocate them as normal
|
|
with torch.inference_mode(False):
|
|
self.cpu = torch.zeros(
|
|
*size, dtype=dtype, device="cpu", pin_memory=pin_memory
|
|
)
|
|
self.gpu = torch.zeros_like(self.cpu, device=device)
|
|
self.np: np.ndarray
|
|
# To keep type hints simple (avoiding generics and subclasses), we
|
|
# only conditionally create the numpy array attribute. This can cause
|
|
# AttributeError if `self.np` is accessed when `with_numpy=False`.
|
|
if with_numpy:
|
|
if dtype == torch.bfloat16:
|
|
raise ValueError(
|
|
"Bfloat16 torch tensors cannot be directly cast to a "
|
|
"numpy array, so call CpuGpuBuffer with with_numpy=False"
|
|
)
|
|
self.np = self.cpu.numpy()
|
|
|
|
def copy_to_gpu(self, n: int | None = None) -> torch.Tensor:
|
|
if n is None:
|
|
return self.gpu.copy_(self.cpu, non_blocking=True)
|
|
return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)
|
|
|
|
def copy_to_cpu(self, n: int | None = None) -> torch.Tensor:
|
|
"""NOTE: Because this method is non-blocking, explicit synchronization
|
|
is needed to ensure the data is copied to CPU."""
|
|
if n is None:
|
|
return self.cpu.copy_(self.gpu, non_blocking=True)
|
|
return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True)
|
|
|
|
|
|
def get_engine_client_zmq_addr(
|
|
local_only: bool,
|
|
host: str,
|
|
port: int = 0,
|
|
) -> str:
|
|
"""Return an IPC path (``local_only=True``) or ``tcp://host:port``.
|
|
|
|
``port=0`` lets the kernel assign the port at ``bind()`` time; the
|
|
caller must recover it via ``getsockopt(zmq.LAST_ENDPOINT)``."""
|
|
if local_only:
|
|
return get_open_zmq_ipc_path()
|
|
return get_tcp_uri(host, port)
|
|
|
|
|
|
class APIServerProcessManager:
|
|
"""Manages a group of API server processes.
|
|
|
|
Handles creation, monitoring, and termination of API server worker
|
|
processes. Also monitors extra processes to check if they are healthy.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
listen_address: str,
|
|
sock: Any,
|
|
args: argparse.Namespace,
|
|
num_servers: int,
|
|
input_addresses: list[str],
|
|
output_addresses: list[str],
|
|
target_server_fn: Callable | None = None,
|
|
stats_update_address: str | None = None,
|
|
tensor_queue: Queue | None = None,
|
|
):
|
|
"""Initialize and start API server worker processes.
|
|
|
|
``input_addresses``/``output_addresses`` may contain
|
|
``tcp://host:0`` placeholders; each child must report the actual
|
|
bound endpoint over its ``actual_address_pipe`` in ``client_config``
|
|
and the parent collects them via
|
|
:py:meth:`gather_actual_addresses`.
|
|
|
|
Args:
|
|
target_server_fn: Override function to call for each API server process
|
|
listen_address: Address to listen for client connections
|
|
sock: Socket for client connections
|
|
args: Command line arguments
|
|
num_servers: Number of API server processes to start
|
|
input_addresses: Input addresses for each API server
|
|
output_addresses: Output addresses for each API server
|
|
stats_update_address: Optional stats update address
|
|
tensor_queue: Optional tensor IPC queue for sharing MM tensors
|
|
"""
|
|
self.listen_address = listen_address
|
|
self.sock = sock
|
|
self.args = args
|
|
|
|
spawn_context = multiprocessing.get_context("spawn")
|
|
self.processes: list[BaseProcess] = []
|
|
self._address_pipes: list[connection.Connection] = []
|
|
|
|
for i, in_addr, out_addr in zip(
|
|
range(num_servers), input_addresses, output_addresses
|
|
):
|
|
client_config: dict[str, Any] = {
|
|
"input_address": in_addr,
|
|
"output_address": out_addr,
|
|
"client_count": num_servers,
|
|
"client_index": i,
|
|
}
|
|
if stats_update_address is not None:
|
|
client_config["stats_update_address"] = stats_update_address
|
|
if tensor_queue is not None:
|
|
client_config["tensor_queue"] = tensor_queue
|
|
|
|
parent_recv, child_send = spawn_context.Pipe(duplex=False)
|
|
self._address_pipes.append(parent_recv)
|
|
client_config["actual_address_pipe"] = child_send
|
|
|
|
proc = spawn_context.Process(
|
|
target=target_server_fn or run_api_server_worker_proc,
|
|
name=f"ApiServer_{i}",
|
|
args=(listen_address, sock, args, client_config),
|
|
)
|
|
self.processes.append(proc)
|
|
proc.start()
|
|
|
|
# Drop parent's write end so reader sees EOF on child death.
|
|
child_send.close()
|
|
|
|
logger.info("Started %d API server processes", len(self.processes))
|
|
|
|
# Shutdown only the API server processes on garbage collection
|
|
# The extra processes are managed by their owners
|
|
self._finalizer = weakref.finalize(self, shutdown, self.processes)
|
|
|
|
def gather_actual_addresses(
|
|
self,
|
|
timeout: float = envs.VLLM_ENGINE_READY_TIMEOUT_S,
|
|
) -> tuple[list[str], list[str]]:
|
|
"""Return (inputs, outputs) reported by each child, indexed by
|
|
``client_index``. Raises ``RuntimeError`` on timeout or premature
|
|
child exit."""
|
|
n = len(self._address_pipes)
|
|
inputs: list[str | None] = [None] * n
|
|
outputs: list[str | None] = [None] * n
|
|
pending: dict[connection.Connection, int] = {
|
|
pipe: i for i, pipe in enumerate(self._address_pipes)
|
|
}
|
|
sentinel_to_idx: dict[Any, int] = {
|
|
proc.sentinel: i for i, proc in enumerate(self.processes)
|
|
}
|
|
|
|
deadline = time.monotonic() + timeout
|
|
try:
|
|
while pending:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
missing = [self.processes[i].name for i in pending.values()]
|
|
raise RuntimeError(
|
|
f"Timed out after {timeout:.1f}s waiting for "
|
|
f"API server(s) to report bound ZMQ addresses: "
|
|
f"{missing}"
|
|
)
|
|
waitables: list[Any] = list(pending.keys()) + list(
|
|
sentinel_to_idx.keys()
|
|
)
|
|
ready = connection.wait(waitables, timeout=remaining)
|
|
# Drain pipes before checking sentinels: a child that sent
|
|
# its message and then exited can surface both events in
|
|
# the same poll, and we must record the success first.
|
|
for item in ready:
|
|
if isinstance(item, connection.Connection) and item in pending:
|
|
idx = pending.pop(item)
|
|
try:
|
|
msg: dict[str, str] = item.recv()
|
|
except EOFError as e:
|
|
raise RuntimeError(
|
|
f"API server {self.processes[idx].name} "
|
|
f"closed its address pipe without "
|
|
f"reporting its bound ZMQ addresses"
|
|
) from e
|
|
inputs[idx] = msg["input_address"]
|
|
outputs[idx] = msg["output_address"]
|
|
item.close()
|
|
for item in ready:
|
|
if item in sentinel_to_idx:
|
|
idx = sentinel_to_idx.pop(item)
|
|
pipe = self._address_pipes[idx]
|
|
if pipe in pending:
|
|
proc = self.processes[idx]
|
|
raise RuntimeError(
|
|
f"API server process {proc.name} exited "
|
|
f"(code={proc.exitcode}) before reporting "
|
|
f"its bound ZMQ addresses"
|
|
)
|
|
finally:
|
|
for pipe in pending:
|
|
with contextlib.suppress(Exception):
|
|
pipe.close()
|
|
|
|
return inputs, outputs # type: ignore[return-value]
|
|
|
|
def shutdown(self, timeout: float | None = None) -> None:
|
|
"""Shutdown API server processes with configurable timeout"""
|
|
for pipe in self._address_pipes:
|
|
with contextlib.suppress(Exception):
|
|
pipe.close()
|
|
self._address_pipes = []
|
|
|
|
if self._finalizer.detach() is not None:
|
|
shutdown(self.processes, timeout=timeout)
|
|
|
|
|
|
class RustFrontendProcessManager:
|
|
"""Manages a single Rust frontend subprocess.
|
|
|
|
Launches the Rust vllm-rs binary in 'frontend' mode, passing the
|
|
listening socket fd and ZMQ transport addresses. Provides the same
|
|
interface as APIServerProcessManager for process monitoring.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
binary_path: str,
|
|
sock: Any,
|
|
args: argparse.Namespace,
|
|
input_address: str,
|
|
output_address: str,
|
|
engine_start_index: int,
|
|
engine_count: int,
|
|
stats_update_address: str | None = None,
|
|
):
|
|
import os
|
|
import subprocess
|
|
|
|
fd = sock.fileno()
|
|
os.set_inheritable(fd, True)
|
|
|
|
cmd = [
|
|
binary_path,
|
|
"frontend",
|
|
"--listen-fd",
|
|
str(fd),
|
|
"--input-address",
|
|
input_address,
|
|
"--output-address",
|
|
output_address,
|
|
"--engine-start-index",
|
|
str(engine_start_index),
|
|
"--engine-count",
|
|
str(engine_count),
|
|
]
|
|
if stats_update_address is not None:
|
|
cmd.extend(["--coordinator-address", stats_update_address])
|
|
from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args
|
|
|
|
args_dict = jsonify_non_default_args(
|
|
args,
|
|
exclude={
|
|
"api_server_count",
|
|
# Python passes the bootstrapped engine range explicitly.
|
|
"data_parallel_rank",
|
|
"data_parallel_external_lb",
|
|
"data_parallel_hybrid_lb",
|
|
},
|
|
)
|
|
# The Rust `frontend` subcommand parses --args-json via serde_json,
|
|
# which bypasses clap and therefore ignores any `#[arg(env = ...)]`
|
|
# declarations on SharedRuntimeArgs fields. Forward the env-driven
|
|
# values explicitly so VLLM_ENGINE_READY_TIMEOUT_S and
|
|
# VLLM_HTTP_TIMEOUT_KEEP_ALIVE behave the same on both Python and Rust
|
|
# frontends.
|
|
args_dict["engine_ready_timeout_secs"] = envs.VLLM_ENGINE_READY_TIMEOUT_S
|
|
args_dict["http_timeout_keep_alive"] = envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE
|
|
args_json = json.dumps(args_dict, sort_keys=True)
|
|
cmd.extend(["--args-json", args_json])
|
|
|
|
logger.info("Launching Rust frontend: %s", " ".join(cmd))
|
|
self._proc = subprocess.Popen(cmd, pass_fds=(fd,))
|
|
|
|
# Create a process wrapper with a sentinel fd for monitoring
|
|
self.processes: list[_SubprocessWrapper] = [
|
|
_SubprocessWrapper(self._proc, "RustFrontend")
|
|
]
|
|
|
|
self._finalizer = weakref.finalize(self, _shutdown_subprocesses, self.processes)
|
|
|
|
def shutdown(self, timeout: float | None = None) -> None:
|
|
if self._finalizer.detach() is not None:
|
|
_shutdown_subprocesses(self.processes, timeout=timeout)
|
|
|
|
|
|
class _SubprocessWrapper:
|
|
"""Wraps subprocess.Popen to provide the BaseProcess-like interface
|
|
needed by wait_for_completion_or_failure."""
|
|
|
|
def __init__(self, proc, name: str):
|
|
self._proc = proc
|
|
self.name = name
|
|
self.pid = proc.pid
|
|
self._sentinel_conn: connection.Connection | None = None
|
|
self._sentinel_send: connection.Connection | None = None
|
|
|
|
# Use a Pipe-based sentinel so subprocess monitoring works uniformly
|
|
# across platforms with multiprocessing.connection.wait().
|
|
recv, send = connection.Pipe(duplex=False)
|
|
self._sentinel_conn = recv
|
|
self._sentinel_send = send
|
|
|
|
def monitor_subprocess() -> None:
|
|
try:
|
|
proc.wait()
|
|
finally:
|
|
with contextlib.suppress(Exception):
|
|
send.close()
|
|
|
|
threading.Thread(
|
|
target=monitor_subprocess, daemon=True, name=f"{name}Monitor"
|
|
).start()
|
|
|
|
@property
|
|
def sentinel(self):
|
|
return self._sentinel_conn
|
|
|
|
@property
|
|
def exitcode(self) -> int | None:
|
|
return self._proc.returncode if self._proc.poll() is not None else None
|
|
|
|
def is_alive(self) -> bool:
|
|
return self._proc.poll() is None
|
|
|
|
def terminate(self):
|
|
self._proc.terminate()
|
|
|
|
def join(self, timeout=None):
|
|
with contextlib.suppress(Exception):
|
|
self._proc.wait(timeout=timeout)
|
|
|
|
def __del__(self):
|
|
with contextlib.suppress(Exception):
|
|
if self._sentinel_conn is not None:
|
|
self._sentinel_conn.close()
|
|
if self._sentinel_send is not None:
|
|
self._sentinel_send.close()
|
|
|
|
|
|
def _shutdown_subprocesses(
|
|
procs: list[_SubprocessWrapper], timeout: float | None = None
|
|
) -> None:
|
|
"""Shutdown subprocess wrappers (mirrors the shutdown() function)."""
|
|
if timeout is None:
|
|
timeout = 0.0
|
|
timeout = max(timeout, 5.0)
|
|
|
|
logger.debug(
|
|
"[shutdown] Subprocess manager: start process_count=%d timeout=%ss",
|
|
len(procs),
|
|
timeout,
|
|
)
|
|
|
|
for proc in procs:
|
|
if proc.is_alive():
|
|
proc.terminate()
|
|
|
|
deadline = time.monotonic() + timeout
|
|
for proc in procs:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
if proc.is_alive():
|
|
proc.join(remaining)
|
|
|
|
remaining_pids = [
|
|
proc.pid for proc in procs if proc.is_alive() and proc.pid is not None
|
|
]
|
|
if remaining_pids:
|
|
logger.warning(
|
|
"[shutdown] Subprocess manager: force killing remaining processes count=%d",
|
|
len(remaining_pids),
|
|
)
|
|
for pid in remaining_pids:
|
|
kill_process_tree(pid)
|
|
|
|
logger.debug_once("[shutdown] Subprocess manager: complete")
|
|
|
|
|
|
def run_api_server_worker_proc(
|
|
listen_address, sock, args, client_config=None, **uvicorn_kwargs
|
|
) -> None:
|
|
"""Entrypoint for individual API server worker processes."""
|
|
|
|
from vllm.entrypoints.openai.api_server import run_server_worker
|
|
|
|
client_config = client_config or {}
|
|
server_index = client_config.get("client_index", 0)
|
|
|
|
# Set process title and add process-specific prefix to stdout and stderr.
|
|
set_process_title("APIServer", str(server_index))
|
|
decorate_logs()
|
|
|
|
uvloop.run(
|
|
run_server_worker(listen_address, sock, args, client_config, **uvicorn_kwargs)
|
|
)
|
|
|
|
|
|
def wait_for_completion_or_failure(
|
|
api_server_manager: "APIServerProcessManager | RustFrontendProcessManager",
|
|
engine_manager: Union["CoreEngineProcManager", "CoreEngineActorManager"]
|
|
| None = None,
|
|
coordinator: "DPCoordinator | None" = None,
|
|
) -> None:
|
|
"""Wait for all processes to complete or detect if any fail.
|
|
|
|
Raises an exception if any process exits with a non-zero status.
|
|
|
|
Args:
|
|
api_server_manager: The manager for API servers.
|
|
engine_manager: The manager for engine processes.
|
|
If CoreEngineProcManager, it manages local engines;
|
|
if CoreEngineActorManager, it manages all engines.
|
|
coordinator: The coordinator for data parallel.
|
|
"""
|
|
|
|
try:
|
|
logger.info("Waiting for API servers to complete ...")
|
|
# Create a mapping of sentinels to their corresponding processes
|
|
# for efficient lookup
|
|
sentinel_to_proc: dict[Any, BaseProcess | _SubprocessWrapper | None] = {
|
|
proc.sentinel: proc for proc in api_server_manager.processes
|
|
}
|
|
|
|
if coordinator:
|
|
sentinel_to_proc[coordinator.proc.sentinel] = coordinator.proc
|
|
|
|
if engine_manager:
|
|
core_shutdown_recv, core_shutdown_send = connection.Pipe(duplex=False)
|
|
|
|
def monitor_engines():
|
|
try:
|
|
engine_manager.monitor_engine_liveness()
|
|
finally:
|
|
core_shutdown_send.close()
|
|
core_shutdown_recv.close()
|
|
|
|
# start monitor for engine liveness
|
|
threading.Thread(target=monitor_engines, daemon=True).start()
|
|
sentinel_to_proc[core_shutdown_recv] = None # type: ignore[assignment]
|
|
|
|
# Check if any process terminates
|
|
while sentinel_to_proc:
|
|
# Wait for any process to terminate (or engine shutdown signal)
|
|
ready_sentinels: list[Any] = connection.wait(sentinel_to_proc)
|
|
|
|
# Process any terminated processes
|
|
for sentinel in ready_sentinels:
|
|
proc = sentinel_to_proc.pop(sentinel)
|
|
|
|
# Check if process exited with error
|
|
if proc is not None and proc.exitcode != 0:
|
|
raise RuntimeError(
|
|
f"Process {proc.name} (PID: {proc.pid}) "
|
|
f"died with exit code {proc.exitcode}"
|
|
)
|
|
if engine_manager and engine_manager.failed_proc_name is not None:
|
|
raise RuntimeError(
|
|
f"Engine core process {engine_manager.failed_proc_name} "
|
|
"died unexpectedly."
|
|
)
|
|
|
|
except KeyboardInterrupt:
|
|
logger.info("Received KeyboardInterrupt, shutting down API servers...")
|
|
except Exception as e:
|
|
logger.exception("Exception occurred while running API servers: %s", str(e))
|
|
raise
|
|
|
|
|
|
# Note(rob): shutdown function cannot be a bound method,
|
|
# else the gc cannot collect the object.
|
|
def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None:
|
|
"""Shutdown processes with timeout.
|
|
|
|
Args:
|
|
procs: List of processes to shutdown
|
|
timeout: Maximum time in seconds to wait for graceful shutdown
|
|
"""
|
|
if timeout is None:
|
|
# Keep a small grace period for best-effort cleanup paths that do not
|
|
# have a user-configured shutdown timeout.
|
|
timeout = 5.0
|
|
|
|
logger.debug(
|
|
"[shutdown] Process manager: start process_count=%d timeout=%ss",
|
|
len(procs),
|
|
timeout,
|
|
)
|
|
|
|
# Shutdown the process.
|
|
for proc in procs:
|
|
if proc.is_alive():
|
|
proc.terminate()
|
|
|
|
# Allow time for remaining procs to terminate.
|
|
deadline = time.monotonic() + timeout
|
|
for proc in procs:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
if proc.is_alive():
|
|
proc.join(remaining)
|
|
|
|
remaining_pids = [
|
|
proc.pid for proc in procs if proc.is_alive() and proc.pid is not None
|
|
]
|
|
if remaining_pids:
|
|
logger.warning(
|
|
"[shutdown] Process manager: force killing remaining processes count=%d",
|
|
len(remaining_pids),
|
|
)
|
|
for pid in remaining_pids:
|
|
kill_process_tree(pid)
|
|
|
|
logger.debug_once("[shutdown] Process manager: complete")
|
|
|
|
|
|
def copy_slice(
|
|
from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int
|
|
) -> torch.Tensor:
|
|
"""
|
|
Copy the first length elements of a tensor into another tensor in a
|
|
non-blocking manner.
|
|
|
|
Used to copy pinned CPU tensor data to pre-allocated GPU tensors.
|
|
|
|
Returns the sliced target tensor.
|
|
"""
|
|
return to_tensor[:length].copy_(from_tensor[:length], non_blocking=True)
|
|
|
|
|
|
def report_usage_stats(
|
|
vllm_config, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT
|
|
) -> None:
|
|
"""Report usage statistics if enabled."""
|
|
|
|
if not is_usage_stats_enabled():
|
|
return
|
|
|
|
from vllm.model_executor.model_loader import get_architecture_class_name
|
|
|
|
model_config = vllm_config.model_config
|
|
scheduler_config = vllm_config.scheduler_config
|
|
parallel_config = vllm_config.parallel_config
|
|
attention_config = vllm_config.attention_config
|
|
compilation_config = vllm_config.compilation_config
|
|
speculative_config = vllm_config.speculative_config
|
|
|
|
# Prepare KV connector string if applicable
|
|
kv_connector = None
|
|
if vllm_config.kv_transfer_config is not None:
|
|
kv_connector = vllm_config.kv_transfer_config.kv_connector
|
|
|
|
# Attention backend is None when set to "auto" (resolved at runtime per platform).
|
|
attention_backend = (
|
|
attention_config.backend.name if attention_config.backend is not None else None
|
|
)
|
|
|
|
# CompilationMode is an IntEnum; report the name for readability in dashboards.
|
|
compilation_mode = (
|
|
compilation_config.mode.name if compilation_config.mode is not None else None
|
|
)
|
|
|
|
# Speculative decoding fields default to None when spec decode is disabled.
|
|
spec_decode_method = (
|
|
speculative_config.method if speculative_config is not None else None
|
|
)
|
|
num_speculative_tokens = (
|
|
speculative_config.num_speculative_tokens
|
|
if speculative_config is not None
|
|
else None
|
|
)
|
|
|
|
if model_config.using_transformers_backend():
|
|
backend_cls = model_config._model_info.architecture
|
|
# Show what was wrapped e.g. TransformersForCausalLM(Starcoder2ForCausalLM)
|
|
architecture = f"{backend_cls}({model_config.architectures[0]})"
|
|
else:
|
|
architecture = get_architecture_class_name(model_config)
|
|
|
|
usage_message.report_usage(
|
|
architecture,
|
|
usage_context,
|
|
extra_kvs={
|
|
# Common configuration
|
|
"dtype": str(model_config.dtype),
|
|
"block_size": vllm_config.cache_config.block_size,
|
|
"gpu_memory_utilization": vllm_config.cache_config.gpu_memory_utilization,
|
|
"kv_cache_memory_bytes": vllm_config.cache_config.kv_cache_memory_bytes,
|
|
# Quantization
|
|
"quantization": model_config.quantization,
|
|
"kv_cache_dtype": str(vllm_config.cache_config.cache_dtype),
|
|
# Feature flags
|
|
"enable_lora": bool(vllm_config.lora_config),
|
|
"enable_prefix_caching": vllm_config.cache_config.enable_prefix_caching,
|
|
"enforce_eager": model_config.enforce_eager,
|
|
"disable_custom_all_reduce": parallel_config.disable_custom_all_reduce,
|
|
# Distributed parallelism settings
|
|
"tensor_parallel_size": parallel_config.tensor_parallel_size,
|
|
"data_parallel_size": parallel_config.data_parallel_size,
|
|
"pipeline_parallel_size": parallel_config.pipeline_parallel_size,
|
|
"enable_expert_parallel": parallel_config.enable_expert_parallel,
|
|
# All2All backend for MoE expert parallel
|
|
"all2all_backend": parallel_config.all2all_backend,
|
|
# KV connector used
|
|
"kv_connector": kv_connector,
|
|
# Batching limits — tuning knobs operators commonly override
|
|
"max_model_len": model_config.max_model_len,
|
|
"max_num_seqs": scheduler_config.max_num_seqs,
|
|
"max_num_batched_tokens": scheduler_config.max_num_batched_tokens,
|
|
# Attention backend (user-requested; None = auto-selected at runtime)
|
|
"attention_backend": attention_backend,
|
|
# torch.compile mode (e.g. NONE, STOCK_TORCH_COMPILE, VLLM_COMPILE)
|
|
"compilation_mode": compilation_mode,
|
|
# Speculative decoding configuration
|
|
"spec_decode_method": spec_decode_method,
|
|
"num_speculative_tokens": num_speculative_tokens,
|
|
# Wide expert parallel: load balancer + redundant/total expert counts
|
|
"enable_eplb": parallel_config.enable_eplb,
|
|
"num_redundant_experts": parallel_config.eplb_config.num_redundant_experts,
|
|
"num_experts": model_config.get_num_experts(),
|
|
},
|
|
)
|
|
|
|
|
|
_PROFILER_FUNC = None
|
|
|
|
|
|
def record_function_or_nullcontext(name: str) -> AbstractContextManager:
|
|
global _PROFILER_FUNC
|
|
|
|
# fast path assume it is set
|
|
if _PROFILER_FUNC is not None:
|
|
return _PROFILER_FUNC(name)
|
|
|
|
func = contextlib.nullcontext
|
|
if envs.VLLM_CUSTOM_SCOPES_FOR_PROFILING:
|
|
func = record_function
|
|
elif envs.VLLM_NVTX_SCOPES_FOR_PROFILING:
|
|
import nvtx
|
|
|
|
func = nvtx.annotate
|
|
|
|
_PROFILER_FUNC = func
|
|
return func(name)
|
|
|
|
|
|
def tensor_data(tensor: torch.Tensor) -> memoryview:
|
|
"""Get the raw data of a tensor as a uint8 memoryview, useful for
|
|
serializing and hashing.
|
|
|
|
Args:
|
|
tensor: The input tensor.
|
|
|
|
Returns:
|
|
A memoryview of the tensor data as uint8.
|
|
"""
|
|
return tensor.flatten().cpu().contiguous().view(torch.uint8).numpy().data
|
|
|
|
|
|
@dataclass
|
|
class IterationDetails:
|
|
num_ctx_requests: int
|
|
num_ctx_tokens: int
|
|
num_generation_requests: int
|
|
num_generation_tokens: int
|
|
|
|
def __repr__(self) -> str:
|
|
return f"IterationDetails(num_ctx_requests={self.num_ctx_requests},\
|
|
num_ctx_tokens={self.num_ctx_tokens}, \
|
|
num_generation_requests={self.num_generation_requests}, \
|
|
num_generation_tokens={self.num_generation_tokens})"
|
|
|
|
|
|
def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDetails:
|
|
"""
|
|
Compute the number of context/generation requests and tokens
|
|
for the current iteration's scheduler output. A requests is regarded
|
|
as a context request if its output tokens are still 0, an extended chunk
|
|
of chunked prefill falls into this category.
|
|
|
|
Args:
|
|
scheduler_output: The scheduler output for the current iteration.
|
|
|
|
Returns:
|
|
An IterationDetails object containing the number of
|
|
context/generation requests and tokens.
|
|
"""
|
|
num_context_requests = 0
|
|
num_context_tokens = 0
|
|
num_generation_requests = 0
|
|
num_generation_tokens = 0
|
|
new_req_ids = {new_req.req_id for new_req in scheduler_output.scheduled_new_reqs}
|
|
for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items():
|
|
if scheduler_output.scheduled_cached_reqs.is_context_phase(req_id) or (
|
|
req_id in new_req_ids
|
|
):
|
|
num_context_requests += 1
|
|
num_context_tokens += num_tokens
|
|
else:
|
|
num_generation_requests += 1
|
|
num_generation_tokens += num_tokens
|
|
return IterationDetails(
|
|
num_context_requests,
|
|
num_context_tokens,
|
|
num_generation_requests,
|
|
num_generation_tokens,
|
|
)
|