tzhouamandClaude Fable 5.1
00510ce38b
One package that makes the POSIX assumptions measured in W0-W8 hold in every process vLLM / vLLM-Omni starts, and reports what it did: - provides fcntl (flock/lockf on msvcrt.locking; LOCK_SH served as exclusive) and uvloop (asyncio.run on a selector loop) when absent; - installs WindowsSelectorEventLoopPolicy and a signal.signal fallback for loop.add_signal_handler on the Windows loop classes; - replaces whole vLLM's get_open_zmq_ipc_path / find_loaded_library when the wheel lacks the branch patch and the upstream body hash is known; - verifies the wheel-only sites (sentinel poll, int64 seed) and reports wheel_required at the exact site instead of failing downstream; - shm_dir() replaces the /dev/shm literals; os.geteuid provided. Activated from vllm_omni/patch.py and as vllm.general_plugins / vllm_omni.general_plugins entry points (idempotent; no-op off Windows). prints the audit; the M0 placement report embeds it. 16 Linux tests through injected primitives. Design: analysis/design_windows_compat_plugin_20260921.md Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
830 lines
34 KiB
Python
830 lines
34 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""The M0 local text engine: run an admitted plan, and record what happened.
|
|
|
|
This is thin on purpose. It owns the plan, the session state references, the
|
|
bounded output stream, the cancellation epoch and the measurement; it owns no
|
|
kernels, no scheduler and no cache. vLLM keeps doing all of that inside the
|
|
Omni stage, which is the whole point of evolving the existing runtime instead
|
|
of writing a second one.
|
|
|
|
**On measuring memory here, honestly.** Three counters are sampled and they are
|
|
three different things, which the architecture record is emphatic about not
|
|
adding together:
|
|
|
|
* ``nvml_device_used`` is the whole card, every process on it. Under WSL2,
|
|
``nvmlDeviceGetComputeRunningProcesses`` returns an empty list, so this
|
|
engine *cannot* attribute GPU bytes to its own PIDs on this host. The delta
|
|
across a load is still evidence -- 7.8 GiB appearing while the host RSS does
|
|
not move says where the weights went -- but it is a delta on a shared
|
|
counter, and it is labelled that way.
|
|
* ``tree_rss`` / ``tree_pss`` are this process and its stage subprocesses.
|
|
PSS includes a share of file-backed pages, so on an mmap'd safetensors load
|
|
it counts pages that are page cache, not anonymous memory.
|
|
* ``host_available`` is the system's, and moves for reasons that are not us.
|
|
|
|
Sampling is at a fixed interval, so a peak shorter than the interval is missed.
|
|
That lower bound is reported with the numbers rather than assumed away.
|
|
|
|
**On cancellation.** ``cancel`` does three things in order: retire the stream's
|
|
epoch so nothing already in flight can be delivered, abort the request in the
|
|
backend so it stops producing, and release the session's reservation in the
|
|
ledger. Doing them in that order is what makes "取消后旧数据不能重新进入输出流"
|
|
true rather than likely.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import os
|
|
import threading
|
|
import time
|
|
from dataclasses import asdict, dataclass, field
|
|
from typing import Any
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.edge.local.plan import ExecutionPlan
|
|
from vllm_omni.edge.local.session import (
|
|
BoundedEventStream,
|
|
ChunkEvent,
|
|
StateHandle,
|
|
new_session_id,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
GiB = 2**30
|
|
SAMPLE_INTERVAL_S = 0.25
|
|
"""Memory sampling period. Anything shorter-lived than this is not seen; the
|
|
reported peaks are lower bounds and say so."""
|
|
|
|
RUNTIME_ENV: dict[str, str] = {
|
|
# FlashInfer's sampler JIT-compiles on first use and needs nvcc, which is
|
|
# not installed here -- nvidia-smi reporting a CUDA version is not a
|
|
# toolkit. Without this the *sampler* takes down the engine at the first
|
|
# token with "Could not find nvcc".
|
|
"VLLM_USE_FLASHINFER_SAMPLER": "0",
|
|
# Under WSL2 vLLM's own GPU model runner allocates a UvaBuffer for its
|
|
# input batch and refuses to start with "UVA is not available" unless
|
|
# pinned memory is enabled. Omni's AR model runner does not take that path,
|
|
# so this matters most when something compares the two.
|
|
"VLLM_WSL2_ENABLE_PIN_MEMORY": "1",
|
|
"TOKENIZERS_PARALLELISM": "false",
|
|
}
|
|
"""The environment the measured runs used, in one place.
|
|
|
|
It lives here rather than inline in ``start()`` because more than one entry
|
|
point has to apply it: a comparison that configures the reference differently
|
|
from the engine is measuring the configuration, not the engine. Applied with
|
|
``setdefault``, so a caller who set one of these keeps their value -- and the
|
|
run record reports what was actually in force."""
|
|
|
|
|
|
def apply_runtime_env() -> dict[str, str | None]:
|
|
"""Set :data:`RUNTIME_ENV` without overriding the caller, and report it."""
|
|
for key, value in RUNTIME_ENV.items():
|
|
os.environ.setdefault(key, value)
|
|
return {key: os.environ.get(key) for key in RUNTIME_ENV}
|
|
|
|
|
|
CANCEL_DRAIN_TIMEOUT_S = 10.0
|
|
"""How long ``cancel`` waits for the producer task to finish unwinding before
|
|
reporting ``producer_stopped: False`` and returning anyway. The consumer is
|
|
released before this wait starts, so exceeding this delays cleanup, not the
|
|
caller."""
|
|
|
|
|
|
# ---- measurement -----------------------------------------------------------
|
|
|
|
|
|
def _host_kind() -> str:
|
|
"""``windows``, ``wsl2`` or ``linux`` -- the three hosts this code has run on."""
|
|
import sys
|
|
|
|
if sys.platform == "win32":
|
|
return "windows"
|
|
try:
|
|
with open("/proc/version", encoding="utf-8", errors="replace") as f:
|
|
if "microsoft" in f.read().lower():
|
|
return "wsl2"
|
|
except OSError:
|
|
pass
|
|
return "linux"
|
|
|
|
|
|
def _host_compat_report() -> dict[str, Any] | None:
|
|
"""``vllm_omni.windows.report()`` on Windows, so a record says what made the run possible."""
|
|
from vllm_omni.windows import is_windows, report
|
|
|
|
if not is_windows():
|
|
return None
|
|
try:
|
|
return report()
|
|
except Exception as exc: # the audit must never take the engine down
|
|
return {"error": f"{type(exc).__name__}: {exc}"}
|
|
|
|
|
|
def gpu_attribution_for_tree() -> dict[str, Any]:
|
|
"""Try to attribute GPU bytes to this process tree, and say why when it fails.
|
|
|
|
[edge-infer W4] The earlier version of this report hardcoded a WSL2
|
|
explanation for an empty ``nvmlDeviceGetComputeRunningProcesses`` list.
|
|
The W1 acceptance run on *native Windows* then carried that WSL2 sentence
|
|
verbatim -- the conclusion (placement unverified) was right, the reason
|
|
offered to the operator was wrong. So this now measures instead of
|
|
assuming: it asks NVML, intersects with the process tree, and only when
|
|
that comes back empty does it name the host-specific mechanism (WSL2's
|
|
GPU-PV and Windows' WDDM driver model both hide per-process usage from
|
|
NVML; a bare Linux host normally reports it).
|
|
"""
|
|
host = _host_kind()
|
|
try:
|
|
import psutil
|
|
import pynvml
|
|
|
|
proc = psutil.Process()
|
|
tree = {proc.pid} | {c.pid for c in proc.children(recursive=True)}
|
|
pynvml.nvmlInit()
|
|
try:
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
|
running = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
|
|
finally:
|
|
with contextlib.suppress(Exception):
|
|
pynvml.nvmlShutdown()
|
|
except Exception as exc:
|
|
return {
|
|
"per_process_gpu_attribution": None,
|
|
"per_process_gpu_attribution_reason": (
|
|
f"NVML per-process query unavailable on this {host} host: {exc}. "
|
|
"The load delta is a whole-card measurement."
|
|
),
|
|
}
|
|
ours = [
|
|
{"pid": int(r.pid), "used_gpu_memory": getattr(r, "usedGpuMemory", None)}
|
|
for r in running
|
|
if int(r.pid) in tree
|
|
]
|
|
if ours:
|
|
return {
|
|
"per_process_gpu_attribution": ours,
|
|
"per_process_gpu_attribution_reason": None,
|
|
}
|
|
mechanism = {
|
|
"wsl2": "under WSL2 the GPU-PV driver reports no per-process compute list",
|
|
"windows": "under the Windows WDDM driver model NVML reports no per-process "
|
|
"compute list (the same NVML call that is empty under WSL2's GPU-PV)",
|
|
"linux": "no compute process from this tree appears in NVML's list",
|
|
}[host]
|
|
return {
|
|
"per_process_gpu_attribution": None,
|
|
"per_process_gpu_attribution_reason": (
|
|
f"nvmlDeviceGetComputeRunningProcesses returned {len(running)} process(es), "
|
|
f"none from this engine's tree of {len(tree)} PID(s): {mechanism}, so GPU "
|
|
"bytes cannot be attributed to this engine's PIDs on this host. The load "
|
|
"delta is a whole-card measurement."
|
|
),
|
|
}
|
|
|
|
|
|
|
|
@dataclass
|
|
class MemorySample:
|
|
unix: float
|
|
nvml_device_used: int | None
|
|
nvml_device_total: int | None
|
|
tree_rss: int
|
|
tree_pss: int
|
|
host_available: int
|
|
host_total: int
|
|
|
|
|
|
class MemorySampler:
|
|
"""Background sampler over the process tree, the host and the GPU.
|
|
|
|
A thread rather than a task: the engine's event loop is the thing being
|
|
measured, and a sampler that stops whenever the loop is busy would miss
|
|
exactly the moments worth sampling.
|
|
"""
|
|
|
|
def __init__(self, interval_s: float = SAMPLE_INTERVAL_S) -> None:
|
|
self.interval_s = interval_s
|
|
self.samples: list[MemorySample] = []
|
|
self._stop = threading.Event()
|
|
self._thread: threading.Thread | None = None
|
|
self._nvml = None
|
|
self._handle = None
|
|
|
|
def _init_nvml(self) -> None:
|
|
try:
|
|
import pynvml
|
|
|
|
pynvml.nvmlInit()
|
|
self._nvml = pynvml
|
|
self._handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
|
except Exception as exc: # no GPU, or no NVML in this venv
|
|
logger.debug("NVML unavailable, GPU memory will not be sampled: %s", exc)
|
|
self._nvml = None
|
|
|
|
def sample(self) -> MemorySample:
|
|
import psutil
|
|
|
|
proc = psutil.Process()
|
|
rss = pss = 0
|
|
for p in [proc] + proc.children(recursive=True):
|
|
try:
|
|
info = p.memory_full_info()
|
|
rss += info.rss
|
|
pss += getattr(info, "pss", 0) or 0
|
|
except (psutil.Error, OSError):
|
|
continue
|
|
vm = psutil.virtual_memory()
|
|
used = total = None
|
|
if self._nvml is not None:
|
|
try:
|
|
mem = self._nvml.nvmlDeviceGetMemoryInfo(self._handle)
|
|
used, total = int(mem.used), int(mem.total)
|
|
except Exception:
|
|
used = total = None
|
|
return MemorySample(
|
|
unix=time.time(), nvml_device_used=used, nvml_device_total=total,
|
|
tree_rss=rss, tree_pss=pss,
|
|
host_available=vm.available, host_total=vm.total,
|
|
)
|
|
|
|
def start(self) -> None:
|
|
self._init_nvml()
|
|
|
|
def loop() -> None:
|
|
while not self._stop.is_set():
|
|
try:
|
|
self.samples.append(self.sample())
|
|
except Exception as exc: # never take the engine down
|
|
logger.debug("memory sample failed: %s", exc)
|
|
self._stop.wait(self.interval_s)
|
|
|
|
self._thread = threading.Thread(target=loop, name="local-mem-sampler", daemon=True)
|
|
self._thread.start()
|
|
|
|
def stop(self) -> None:
|
|
self._stop.set()
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=2.0)
|
|
if self._nvml is not None:
|
|
with contextlib.suppress(Exception):
|
|
self._nvml.nvmlShutdown()
|
|
|
|
def peaks(self, since: float = 0.0) -> dict[str, Any]:
|
|
window = [s for s in self.samples if s.unix >= since]
|
|
if not window:
|
|
return {"samples": 0, "interval_s": self.interval_s}
|
|
gpu = [s.nvml_device_used for s in window if s.nvml_device_used is not None]
|
|
return {
|
|
"samples": len(window),
|
|
"interval_s": self.interval_s,
|
|
"note": (
|
|
"peaks are lower bounds: anything shorter than interval_s is missed. "
|
|
"nvml_device_used is the whole card including other processes; "
|
|
"tree_pss includes a share of file-backed pages. These three "
|
|
"counters are different accountings and must not be summed."
|
|
),
|
|
"nvml_device_used_peak": max(gpu) if gpu else None,
|
|
"nvml_device_used_min": min(gpu) if gpu else None,
|
|
"tree_rss_peak": max(s.tree_rss for s in window),
|
|
"tree_pss_peak": max(s.tree_pss for s in window),
|
|
"host_available_min": min(s.host_available for s in window),
|
|
}
|
|
|
|
|
|
# ---- per-request record ----------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class RequestRecord:
|
|
"""Timing for one request, in the units the validation plan asks for."""
|
|
|
|
request_id: str
|
|
session_id: str
|
|
epoch: int
|
|
prompt_chars: int
|
|
prompt_tokens: int | None = None
|
|
output_tokens: int = 0
|
|
submitted_unix: float = 0.0
|
|
first_token_unix: float | None = None
|
|
last_token_unix: float | None = None
|
|
finished: bool = False
|
|
cancelled: bool = False
|
|
error: str | None = None
|
|
token_unix: list[float] = field(default_factory=list)
|
|
text: str = ""
|
|
output_token_ids: list[int] = field(default_factory=list)
|
|
"""The sampled ids, kept so a run can be compared against a reference
|
|
decode token-for-token. Text is not enough for that: two different id
|
|
sequences can detokenize to the same string."""
|
|
|
|
@property
|
|
def ttft_s(self) -> float | None:
|
|
"""Submission to first token. Includes admission and scheduling, which
|
|
is what a caller experiences; it is not an isolated prefill kernel
|
|
time and is never reported as one."""
|
|
if self.first_token_unix is None:
|
|
return None
|
|
return self.first_token_unix - self.submitted_unix
|
|
|
|
@property
|
|
def decode_tok_per_s(self) -> float | None:
|
|
"""Tokens after the first, over the time between first and last.
|
|
|
|
``None`` below two tokens rather than a number computed from one
|
|
interval.
|
|
"""
|
|
if self.first_token_unix is None or self.last_token_unix is None:
|
|
return None
|
|
if self.output_tokens < 2:
|
|
return None
|
|
span = self.last_token_unix - self.first_token_unix
|
|
return (self.output_tokens - 1) / span if span > 0 else None
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
d = asdict(self)
|
|
d["ttft_s"] = self.ttft_s
|
|
d["decode_tok_per_s"] = self.decode_tok_per_s
|
|
return d
|
|
|
|
|
|
# ---- the engine ------------------------------------------------------------
|
|
|
|
|
|
class PlanNotAdmitted(RuntimeError):
|
|
"""Raised when someone tries to run a plan that refused."""
|
|
|
|
|
|
class LocalTextEngine:
|
|
"""One admitted plan, one Omni stage, streamed text with a cancel epoch."""
|
|
|
|
def __init__(self, plan: ExecutionPlan, *, extra_engine_kwargs: dict[str, Any] | None = None) -> None:
|
|
if not plan.admitted:
|
|
raise PlanNotAdmitted(plan.summary())
|
|
self.plan = plan
|
|
self.extra_engine_kwargs = dict(extra_engine_kwargs or {})
|
|
self.sampler = MemorySampler()
|
|
self.records: dict[str, RequestRecord] = {}
|
|
self._omni: Any = None
|
|
self._sessions: dict[str, StateHandle] = {}
|
|
self._streams: dict[str, BoundedEventStream] = {}
|
|
self._pumps: dict[str, asyncio.Task] = {}
|
|
self._inflight: dict[str, str] = {} # request_id -> session_id
|
|
self.load_seconds: float | None = None
|
|
self.load_memory: dict[str, Any] = {}
|
|
self._closed = False
|
|
|
|
# -- lifecycle ----------------------------------------------------------
|
|
async def start(self) -> None:
|
|
"""Load the model according to the plan, and price the load."""
|
|
from vllm_omni import AsyncOmni
|
|
|
|
apply_runtime_env()
|
|
self.sampler.start()
|
|
before = self.sampler.sample()
|
|
started = time.perf_counter()
|
|
kwargs = {**self.plan.engine_kwargs, **self.extra_engine_kwargs}
|
|
logger.info(
|
|
"[local] loading %s on %s with %s",
|
|
self.plan.manifest.model_dir, self.plan.selected.device_id, kwargs,
|
|
)
|
|
self._omni = AsyncOmni(model=self.plan.manifest.model_dir, **kwargs)
|
|
self.load_seconds = time.perf_counter() - started
|
|
after = self.sampler.sample()
|
|
|
|
self.load_memory = {
|
|
"before": asdict(before),
|
|
"after": asdict(after),
|
|
"nvml_device_used_delta": (
|
|
(after.nvml_device_used - before.nvml_device_used)
|
|
if (after.nvml_device_used is not None and before.nvml_device_used is not None)
|
|
else None
|
|
),
|
|
"tree_rss_delta": after.tree_rss - before.tree_rss,
|
|
"peaks_during_load": self.sampler.peaks(since=before.unix),
|
|
}
|
|
self._record_actuals()
|
|
|
|
def _record_actuals(self) -> None:
|
|
"""Turn the load into a bound on the reservations, without overclaiming.
|
|
|
|
The NVML delta across the load is **not** "the weights". It is the
|
|
weights plus the CUDA context, the allocator's arenas and whatever else
|
|
the backend set up, measured on a counter shared with every other
|
|
process on the card. So it is recorded as an *upper bound* on the
|
|
weights reservation rather than as its actual value: a ratio above 1.0
|
|
here means the budget was under the bound, not that the weights were
|
|
bigger than their files.
|
|
"""
|
|
pool = self.plan.selected.memory_pool if self.plan.selected else "vram"
|
|
if pool == "vram":
|
|
delta = self.load_memory.get("nvml_device_used_delta")
|
|
unit = "on the device (whole-card delta)"
|
|
else:
|
|
delta = self.load_memory.get("tree_rss_delta")
|
|
unit = "of process-tree RSS"
|
|
if not delta or delta <= 0:
|
|
return
|
|
for reservation in self.plan.reservations:
|
|
if reservation.purpose == "weights" and reservation.pool == pool:
|
|
reservation.actual_upper_bound_bytes = int(delta)
|
|
reservation.note += (
|
|
f"; the load moved +{delta / GiB:.2f} GiB {unit}, which bounds "
|
|
"weights + backend context together and is not an attribution "
|
|
"to the weights alone"
|
|
)
|
|
|
|
def measured_peak(self) -> dict[str, Any]:
|
|
"""Budgeted peak against the measured one, on the selected pool.
|
|
|
|
The measured number subtracts a pre-load baseline, because the counter
|
|
is shared: on this host the card already held 2.75 GiB before the
|
|
engine started, and reporting that as the engine's would be the same
|
|
mistake the architecture record calls out about summing WSL ``used``,
|
|
Windows ``used`` and PSS.
|
|
"""
|
|
if self.plan.selected is None:
|
|
return {}
|
|
pool = self.plan.selected.memory_pool
|
|
peaks = self.sampler.peaks()
|
|
before = self.load_memory.get("before") or {}
|
|
if pool == "vram":
|
|
baseline = before.get("nvml_device_used")
|
|
peak = peaks.get("nvml_device_used_peak")
|
|
counter = "nvml_device_used"
|
|
else:
|
|
baseline = before.get("tree_rss")
|
|
peak = peaks.get("tree_rss_peak")
|
|
counter = "tree_rss"
|
|
attributable = (peak - baseline) if (peak is not None and baseline is not None) else None
|
|
budget = self.plan.peak_bytes
|
|
return {
|
|
"pool": pool,
|
|
"counter": counter,
|
|
"baseline_bytes": baseline,
|
|
"peak_bytes": peak,
|
|
"attributable_peak_bytes": attributable,
|
|
"budgeted_peak_bytes": budget,
|
|
"budget_over_measured": (budget / attributable) if attributable else None,
|
|
"admitted_without_oom": True,
|
|
"caveats": (
|
|
"sampled every %.2f s, so the peak is a lower bound; the counter is "
|
|
"shared with other processes and the baseline subtraction only "
|
|
"removes what was there before the load, not what arrived during it. "
|
|
"A budget above the measured peak is the intended outcome -- this "
|
|
"budget is a safe upper bound for admission, not a prediction."
|
|
% self.sampler.interval_s
|
|
),
|
|
}
|
|
|
|
async def close(self) -> None:
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
for task in list(self._pumps.values()):
|
|
task.cancel()
|
|
for task in list(self._pumps.values()):
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
await task
|
|
self._pumps.clear()
|
|
for stream in list(self._streams.values()):
|
|
with contextlib.suppress(Exception):
|
|
await stream.close()
|
|
self._streams.clear()
|
|
if self._omni is not None:
|
|
for name in ("shutdown", "close"):
|
|
closer = getattr(self._omni, name, None)
|
|
if callable(closer):
|
|
with contextlib.suppress(Exception):
|
|
result = closer()
|
|
if asyncio.iscoroutine(result):
|
|
await result
|
|
break
|
|
self.sampler.stop()
|
|
|
|
async def __aenter__(self) -> LocalTextEngine:
|
|
await self.start()
|
|
return self
|
|
|
|
async def __aexit__(self, *exc: Any) -> None:
|
|
await self.close()
|
|
|
|
# -- sessions -----------------------------------------------------------
|
|
def open_session(self) -> StateHandle:
|
|
handle = StateHandle(
|
|
session_id=new_session_id(),
|
|
backend=self.plan.backend or "unknown",
|
|
artifact_id=self.plan.manifest.artifact_id,
|
|
)
|
|
self._sessions[handle.session_id] = handle
|
|
return handle
|
|
|
|
def session(self, session_id: str) -> StateHandle:
|
|
return self._sessions[session_id]
|
|
|
|
def close_session(self, session_id: str) -> None:
|
|
self._sessions.pop(session_id, None)
|
|
|
|
# -- generation ---------------------------------------------------------
|
|
async def submit(
|
|
self,
|
|
handle: StateHandle,
|
|
prompt: str,
|
|
*,
|
|
max_tokens: int = 128,
|
|
temperature: float = 0.0,
|
|
ignore_eos: bool = False,
|
|
request_id: str | None = None,
|
|
max_chunks: int = 64,
|
|
max_bytes: int = 1 << 20,
|
|
) -> tuple[str, BoundedEventStream]:
|
|
"""Start one request; return its id and its bounded event stream.
|
|
|
|
Returning the stream rather than yielding from here keeps cancellation
|
|
possible from outside the consumer: ``cancel`` can retire the epoch
|
|
while the caller is still awaiting an event.
|
|
"""
|
|
if self._omni is None:
|
|
raise RuntimeError("engine not started")
|
|
live = self._sessions.get(handle.session_id)
|
|
if live is None or not live.accepts(handle):
|
|
raise RuntimeError(
|
|
f"state handle for session {handle.session_id} is stale "
|
|
f"(epoch {handle.epoch}, live {live.epoch if live else 'closed'}); "
|
|
"open a new session rather than reusing retired state"
|
|
)
|
|
from vllm import SamplingParams
|
|
|
|
request_id = request_id or f"{handle.session_id}-{len(self.records)}"
|
|
stream = BoundedEventStream(max_chunks=max_chunks, max_bytes=max_bytes)
|
|
if handle.epoch:
|
|
# A session that has been cancelled before starts its new
|
|
# stream already fenced against everything older.
|
|
stream.retire_epoch(handle.epoch - 1)
|
|
record = RequestRecord(
|
|
request_id=request_id,
|
|
session_id=handle.session_id,
|
|
epoch=handle.epoch,
|
|
prompt_chars=len(prompt),
|
|
submitted_unix=time.time(),
|
|
)
|
|
self.records[request_id] = record
|
|
self._streams[request_id] = stream
|
|
self._inflight[request_id] = handle.session_id
|
|
|
|
params = SamplingParams(
|
|
temperature=temperature, max_tokens=max_tokens, ignore_eos=ignore_eos
|
|
)
|
|
self._pumps[request_id] = asyncio.create_task(
|
|
self._pump(request_id, handle, prompt, params, stream, record),
|
|
name=f"local-pump-{request_id}",
|
|
)
|
|
return request_id, stream
|
|
|
|
async def _pump(
|
|
self,
|
|
request_id: str,
|
|
handle: StateHandle,
|
|
prompt: str,
|
|
params: Any,
|
|
stream: BoundedEventStream,
|
|
record: RequestRecord,
|
|
) -> None:
|
|
"""Adapt Omni's cumulative outputs into incremental ChunkEvents.
|
|
|
|
Omni yields the output-so-far, like vLLM. Re-emitting that verbatim
|
|
would make a dropped chunk invisible to the consumer, so only the delta
|
|
goes on the wire and ``seq`` is what detects loss.
|
|
"""
|
|
seq = 0
|
|
emitted_text = ""
|
|
try:
|
|
async for out in self._omni.generate(
|
|
prompt, request_id=request_id, sampling_params_list=[params]
|
|
):
|
|
completion = _first_completion(out)
|
|
if completion is None:
|
|
continue
|
|
text = completion.get("text", "")
|
|
token_ids = completion.get("token_ids", ())
|
|
if record.prompt_tokens is None:
|
|
record.prompt_tokens = completion.get("prompt_tokens")
|
|
delta = text[len(emitted_text):] if text.startswith(emitted_text) else text
|
|
n_tokens = len(token_ids)
|
|
if n_tokens <= record.output_tokens and not delta:
|
|
continue
|
|
now = time.time()
|
|
if record.first_token_unix is None:
|
|
record.first_token_unix = now
|
|
record.last_token_unix = now
|
|
record.output_tokens = max(record.output_tokens, n_tokens)
|
|
record.token_unix.append(now)
|
|
emitted_text = text if text.startswith(emitted_text) else emitted_text + delta
|
|
if n_tokens >= len(record.output_token_ids):
|
|
record.output_token_ids = list(token_ids)
|
|
seq += 1
|
|
await stream.put(ChunkEvent(
|
|
request_id=request_id, stage_id=0, seq=seq, epoch=handle.epoch,
|
|
kind="token", payload={"text": delta, "n_tokens": n_tokens},
|
|
started_unix=record.submitted_unix, emitted_unix=now,
|
|
input_watermark=record.prompt_tokens or 0,
|
|
))
|
|
record.text = emitted_text
|
|
record.finished = True
|
|
await stream.put(ChunkEvent(
|
|
request_id=request_id, stage_id=0, seq=seq + 1, epoch=handle.epoch,
|
|
kind="done", payload={"output_tokens": record.output_tokens},
|
|
started_unix=record.submitted_unix, emitted_unix=time.time(), final=True,
|
|
))
|
|
except asyncio.CancelledError:
|
|
record.cancelled = True
|
|
record.text = emitted_text
|
|
raise
|
|
except Exception as exc:
|
|
record.error = str(exc)
|
|
record.text = emitted_text
|
|
with contextlib.suppress(Exception):
|
|
await stream.put(ChunkEvent(
|
|
request_id=request_id, stage_id=0, seq=seq + 1, epoch=handle.epoch,
|
|
kind="error", payload=None, started_unix=record.submitted_unix,
|
|
emitted_unix=time.time(), final=True, error=str(exc),
|
|
))
|
|
finally:
|
|
self._inflight.pop(request_id, None)
|
|
with contextlib.suppress(Exception):
|
|
await stream.close()
|
|
|
|
async def cancel(self, request_id: str) -> dict[str, Any]:
|
|
"""Retire the epoch, stop the backend, release the reservation.
|
|
|
|
The order matters and is the contract: by the time this returns, no
|
|
event produced before the cancel can still reach the consumer.
|
|
"""
|
|
record = self.records.get(request_id)
|
|
stream = self._streams.get(request_id)
|
|
session_id = self._inflight.get(request_id) or (record.session_id if record else None)
|
|
before = self.sampler.sample()
|
|
|
|
if stream is not None:
|
|
stream.retire_epoch(record.epoch if record else 0)
|
|
|
|
aborted = False
|
|
if self._omni is not None:
|
|
for name in ("abort", "abort_request", "abort_requests_async"):
|
|
fn = getattr(self._omni, name, None)
|
|
if callable(fn):
|
|
with contextlib.suppress(Exception):
|
|
result = fn(request_id)
|
|
if asyncio.iscoroutine(result):
|
|
await result
|
|
aborted = True
|
|
break
|
|
|
|
# Close the stream *before* waiting on the producer. The consumer is
|
|
# blocked in ``get()`` and the sentinel is what releases it; making
|
|
# that depend on the backend's unwind would mean a backend that hangs
|
|
# hangs the caller too, which is the opposite of what cancel is for.
|
|
if stream is not None:
|
|
with contextlib.suppress(Exception):
|
|
await stream.close()
|
|
|
|
task = self._pumps.pop(request_id, None)
|
|
pump_stopped = True
|
|
if task is not None:
|
|
task.cancel()
|
|
# Bounded: tearing down an async generator runs the backend's own
|
|
# cleanup, and that cleanup can wait on a request this method has
|
|
# already aborted. Waiting forever for it would turn a cancel into
|
|
# a deadlock, so the wait has a deadline and the outcome is
|
|
# reported rather than assumed.
|
|
done, pending = await asyncio.wait({task}, timeout=CANCEL_DRAIN_TIMEOUT_S)
|
|
pump_stopped = not pending
|
|
if done:
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
task.result()
|
|
|
|
if record is not None:
|
|
record.cancelled = True
|
|
if session_id and session_id in self._sessions:
|
|
# A cancelled session's state is retired, not reused: the next
|
|
# request on it must start from a fresh epoch.
|
|
self._sessions[session_id] = self._sessions[session_id].next_epoch()
|
|
|
|
# Give the backend a moment to free blocks, then look.
|
|
await asyncio.sleep(SAMPLE_INTERVAL_S * 2)
|
|
after = self.sampler.sample()
|
|
return {
|
|
"request_id": request_id,
|
|
"aborted_in_backend": aborted,
|
|
"producer_stopped": pump_stopped,
|
|
"epoch_retired": stream.epoch if stream is not None else None,
|
|
"dropped_stale_events": stream.dropped_stale if stream is not None else 0,
|
|
"output_tokens_at_cancel": record.output_tokens if record else 0,
|
|
"session_epoch_now": self._sessions[session_id].epoch if session_id in self._sessions else None,
|
|
"inflight_after": len(self._inflight),
|
|
"nvml_device_used_before": before.nvml_device_used,
|
|
"nvml_device_used_after": after.nvml_device_used,
|
|
"note": (
|
|
"a cancelled request's KV blocks return to vLLM's pool, which is "
|
|
"pre-allocated; the device counter is therefore expected NOT to "
|
|
"fall. What must fall is the in-flight count, and no retired-epoch "
|
|
"event may be delivered."
|
|
),
|
|
}
|
|
|
|
# -- reporting ----------------------------------------------------------
|
|
def report_placement(self) -> dict[str, Any]:
|
|
"""Where the stage is actually configured and running.
|
|
|
|
Two claims of different strength: the *configured* device, read back
|
|
from the stage's own ``VllmConfig`` rather than from the plan that
|
|
requested it, and the *observed* device memory delta across the load.
|
|
Neither is a per-process attribution, because NVML cannot give one
|
|
here; that limitation is in the output, not only in the docstring.
|
|
"""
|
|
configured: dict[str, Any] = {}
|
|
try:
|
|
configs = getattr(self._omni.engine, "stage_vllm_configs", None) or {}
|
|
cfg = configs[0] if isinstance(configs, (list, tuple)) else configs.get(0)
|
|
device_config = getattr(cfg, "device_config", None)
|
|
model_config = getattr(cfg, "model_config", None)
|
|
cache_config = getattr(cfg, "cache_config", None)
|
|
configured = {
|
|
"device": str(getattr(device_config, "device", "")),
|
|
"device_type": str(getattr(device_config, "device_type", "")),
|
|
"dtype": str(getattr(model_config, "dtype", "")),
|
|
"quantization": getattr(model_config, "quantization", None),
|
|
"max_model_len": getattr(model_config, "max_model_len", None),
|
|
"enforce_eager": getattr(model_config, "enforce_eager", None),
|
|
"kv_cache_memory_bytes": getattr(cache_config, "kv_cache_memory_bytes", None),
|
|
"num_gpu_blocks": getattr(cache_config, "num_gpu_blocks", None),
|
|
"block_size": getattr(cache_config, "block_size", None),
|
|
}
|
|
except Exception as exc:
|
|
configured = {"error": f"stage config unavailable: {exc}"}
|
|
|
|
return {
|
|
"planned_device": self.plan.selected.device_id if self.plan.selected else None,
|
|
"planned_backend": self.plan.backend,
|
|
"configured": configured,
|
|
"observed": {
|
|
"nvml_device_used_delta_at_load": self.load_memory.get("nvml_device_used_delta"),
|
|
"tree_rss_delta_at_load": self.load_memory.get("tree_rss_delta"),
|
|
**gpu_attribution_for_tree(),
|
|
},
|
|
# [edge-infer] On native Windows, which compatibility layer carried this
|
|
# process (shims, patches, wheel-provided sites). None elsewhere.
|
|
"host_compat": _host_compat_report(),
|
|
}
|
|
|
|
def report_usage(self) -> dict[str, Any]:
|
|
"""Budgeted vs measured, per reservation, plus the sampled peaks."""
|
|
return {
|
|
"load_seconds": self.load_seconds,
|
|
"reservations": [
|
|
{
|
|
**r.to_dict(),
|
|
"upper_bound_over_budget": (
|
|
r.actual_upper_bound_bytes / r.bytes
|
|
if (r.actual_upper_bound_bytes and r.bytes)
|
|
else None
|
|
),
|
|
}
|
|
for r in self.plan.reservations
|
|
],
|
|
"budgeted_peak_bytes": self.plan.peak_bytes,
|
|
"measured_peak": self.measured_peak(),
|
|
"load_memory": self.load_memory,
|
|
"peaks": self.sampler.peaks(),
|
|
"requests": [r.to_dict() for r in self.records.values()],
|
|
"streams": {rid: s.stats() for rid, s in self._streams.items()},
|
|
}
|
|
|
|
|
|
def _first_completion(out: Any) -> dict[str, Any] | None:
|
|
"""Pull ``(text, token_ids, prompt_tokens)`` out of an Omni output.
|
|
|
|
Kept in one function because the shape is the thing most likely to move
|
|
between vLLM/Omni versions, and a version bump should break here, loudly,
|
|
rather than silently start reporting zero tokens.
|
|
"""
|
|
outputs = getattr(out, "outputs", None)
|
|
if not outputs:
|
|
return None
|
|
first = outputs[0]
|
|
text = getattr(first, "text", None)
|
|
token_ids = getattr(first, "token_ids", None)
|
|
if text is None and token_ids is None:
|
|
return None
|
|
prompt_ids = getattr(out, "prompt_token_ids", None)
|
|
return {
|
|
"text": text or "",
|
|
"token_ids": tuple(token_ids or ()),
|
|
"prompt_tokens": len(prompt_ids) if prompt_ids is not None else None,
|
|
}
|