wehub-resource-sync
1282 lines
51 KiB
Python
1282 lines
51 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from vllm-mlx (https://github.com/vllm-project/vllm-mlx).
|
|
"""
|
|
Engine Core for oMLX continuous batching.
|
|
|
|
This module provides the EngineCore class that coordinates:
|
|
- Model loading and management
|
|
- Request scheduling via Scheduler
|
|
- Async request processing
|
|
- Output streaming
|
|
|
|
The design follows vLLM's engine architecture adapted for MLX.
|
|
"""
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import gc
|
|
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass, field
|
|
from typing import (
|
|
Any,
|
|
AsyncIterator,
|
|
Awaitable,
|
|
Callable,
|
|
Dict,
|
|
List,
|
|
Optional,
|
|
Tuple,
|
|
Union,
|
|
)
|
|
|
|
import mlx.core as mx
|
|
|
|
from .exceptions import PrefillMemoryExceededError
|
|
from .model_registry import get_registry
|
|
from .output_collector import RequestOutputCollector, RequestStreamState
|
|
from .request import Request, RequestOutput, SamplingParams
|
|
from .scheduler import Scheduler, SchedulerConfig, _sync_and_clear_cache
|
|
from .utils.compile_cache import (
|
|
clear_thread_compile_cache,
|
|
compile_cache_clear_available,
|
|
)
|
|
from .utils.fatal import FATAL_TEARDOWN_TIMEOUT_S, fatal_exit
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _raise_request_output_error(output: RequestOutput) -> None:
|
|
if output.error_code == "prefill_memory_exceeded":
|
|
metadata = output.error_metadata or {}
|
|
request_id = metadata.get("request_id")
|
|
estimated_bytes = metadata.get("estimated_bytes")
|
|
limit_bytes = metadata.get("limit_bytes")
|
|
raise PrefillMemoryExceededError(
|
|
message=output.error or "Prefill memory exceeded",
|
|
request_id=str(request_id) if request_id is not None else output.request_id,
|
|
estimated_bytes=(
|
|
int(estimated_bytes) if estimated_bytes is not None else None
|
|
),
|
|
limit_bytes=int(limit_bytes) if limit_bytes is not None else None,
|
|
)
|
|
raise RuntimeError(output.error)
|
|
|
|
|
|
_global_mlx_executor: concurrent.futures.ThreadPoolExecutor | None = None
|
|
|
|
# Fallback only: used when the MLX compile-cache clear symbol is unavailable
|
|
# (see omlx/utils/compile_cache.py). In that case a per-engine MLX worker
|
|
# thread cannot exit safely (its thread_local ~CompilerCache would free
|
|
# @mx.compile graphs' Python objects without the GIL -> crash), so close()
|
|
# keeps the executor + stream alive here for the process lifetime instead of
|
|
# shutting the thread down. With the clear symbol present (the normal path)
|
|
# these stay empty and the worker threads shut down normally.
|
|
_immortal_mlx_executors: list = []
|
|
_immortal_mlx_streams: list = []
|
|
|
|
|
|
def _final_engine_thread_reclaim(stream: Any) -> None:
|
|
"""Drop Python cycles and reclaim MLX buffers on the engine worker thread."""
|
|
gc.collect()
|
|
_sync_and_clear_cache(stream)
|
|
gc.collect()
|
|
|
|
|
|
def _init_mlx_thread() -> None:
|
|
"""Replace generation_stream with a thread-local stream on the executor thread.
|
|
|
|
mlx-lm's module-level ``generation_stream`` is created at import time in
|
|
whichever thread imported it first (the main thread at server startup).
|
|
Arrays produced inside ``with mx.stream(generation_stream):`` blocks carry
|
|
that stream reference. If the stream was created on the main thread,
|
|
subsequent ``.item()`` / ``mx.synchronize()`` calls from the executor
|
|
thread fail with "There is no Stream(gpu, 0) in current thread".
|
|
|
|
Fix: create a thread-local stream HERE and replace the module-level
|
|
``generation_stream`` in mlx_lm.generate and omlx.scheduler.
|
|
"""
|
|
import sys
|
|
|
|
import mlx.core as mx
|
|
|
|
stream = mx.new_thread_local_stream(mx.default_device())
|
|
|
|
gen_mod = sys.modules.get("mlx_lm.generate")
|
|
if gen_mod is not None:
|
|
gen_mod.generation_stream = stream
|
|
|
|
sched_mod = sys.modules.get("omlx.scheduler")
|
|
if sched_mod is not None:
|
|
sched_mod.generation_stream = stream
|
|
|
|
logger.info(f"MLX executor thread initialized: generation_stream = {stream}")
|
|
|
|
|
|
def get_mlx_executor() -> concurrent.futures.ThreadPoolExecutor:
|
|
"""Get or create the global MLX executor (lazy singleton).
|
|
|
|
mlx-lm's BatchGenerator uses a module-level Metal stream
|
|
(generation_stream), so ALL MLX GPU operations across all models
|
|
MUST be serialized onto one thread to prevent Metal command buffer
|
|
races that cause segfaults. See issue #85.
|
|
"""
|
|
global _global_mlx_executor
|
|
if _global_mlx_executor is None:
|
|
_global_mlx_executor = concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=1,
|
|
thread_name_prefix="mlx-global",
|
|
initializer=_init_mlx_thread,
|
|
)
|
|
return _global_mlx_executor
|
|
|
|
|
|
@dataclass
|
|
class EngineConfig:
|
|
"""Configuration for the engine."""
|
|
|
|
model_name: str = ""
|
|
scheduler_config: Optional[SchedulerConfig] = None
|
|
step_interval: float = 0.05 # Idle wait timeout; requests wake the loop
|
|
stream_interval: int = 1 # Tokens to batch before streaming (1=every token)
|
|
prefill_eviction_callback: Optional[Callable[[Any], Awaitable[bool]]] = None
|
|
# Decode burst: run several scheduler.step() calls per run_in_executor
|
|
# hand-off instead of one. Each decode token otherwise bounces back to the
|
|
# event loop, ping-ponging the GIL with the asyncio loop + uvicorn on the
|
|
# main thread; bursting keeps the MLX thread holding the GIL continuously.
|
|
# scheduler.step() services aborts/admission/finish every step, so
|
|
# correctness is unchanged and memory is identical (same tokens decoded,
|
|
# same KV cache; only a small list of K SchedulerOutputs is held per
|
|
# burst). The budget is a TIME ceiling so the event-loop pause (and thus
|
|
# new-request admission / abort / HTTP latency) is bounded consistently
|
|
# across hardware, and a slow prefill-chunk step ends the burst.
|
|
#
|
|
# Adaptive: with a single active request (the common local/single-user
|
|
# case) there is no concurrent request to stay responsive to, so we burst
|
|
# aggressively (decode_burst_budget_single_s). Once concurrent, we use the
|
|
# tight decode_burst_budget_s to keep admission/abort latency low.
|
|
# max_steps is a safety cap (bounds the host-side output list), NOT a
|
|
# memory knob. Set both budgets <= 0, or max_steps <= 1, to disable.
|
|
decode_burst_max_steps: int = field(
|
|
default_factory=lambda: int(os.environ.get("OMLX_DECODE_BURST_MAX_STEPS", "64"))
|
|
)
|
|
decode_burst_budget_single_s: float = field(
|
|
default_factory=lambda: float(
|
|
os.environ.get("OMLX_DECODE_BURST_BUDGET_SINGLE_S", "0.1")
|
|
)
|
|
)
|
|
decode_burst_budget_s: float = field(
|
|
default_factory=lambda: float(
|
|
os.environ.get("OMLX_DECODE_BURST_BUDGET_S", "0.03")
|
|
)
|
|
)
|
|
|
|
|
|
class EngineCore:
|
|
"""
|
|
Core engine for oMLX inference with continuous batching.
|
|
|
|
This engine runs the generation loop and manages request lifecycle.
|
|
It provides both sync and async interfaces for request handling.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: Any,
|
|
tokenizer: Any,
|
|
config: Optional[EngineConfig] = None,
|
|
engine_id: Optional[str] = None,
|
|
force_model_ownership: bool = True,
|
|
):
|
|
"""
|
|
Initialize the engine.
|
|
|
|
Args:
|
|
model: The MLX model
|
|
tokenizer: The tokenizer
|
|
config: Engine configuration
|
|
engine_id: Optional unique ID for this engine (auto-generated if None)
|
|
force_model_ownership: If True (default), forcibly take model ownership
|
|
from any existing engine. If False, raises
|
|
ModelOwnershipError if model is in use.
|
|
"""
|
|
self.model = model
|
|
self.tokenizer = tokenizer
|
|
self.config = config or EngineConfig()
|
|
self._engine_id = engine_id or str(uuid.uuid4())
|
|
self._owns_model = False
|
|
self._closed = False
|
|
|
|
# Acquire model ownership
|
|
registry = get_registry()
|
|
registry.acquire(
|
|
model=model,
|
|
engine=self,
|
|
engine_id=self._engine_id,
|
|
force=force_model_ownership,
|
|
)
|
|
self._owns_model = True
|
|
|
|
# Per-engine executor with dedicated mx.Stream (#1248).
|
|
# Each EngineCore gets its own thread + GPU stream so different
|
|
# models can run scheduler.step() concurrently.
|
|
self._mlx_stream = mx.new_thread_local_stream(mx.default_device())
|
|
self._mlx_executor = concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=1,
|
|
thread_name_prefix=f"mlx-engine-{self._engine_id[:8]}",
|
|
)
|
|
|
|
# Create scheduler with per-engine stream
|
|
scheduler_config = self.config.scheduler_config or SchedulerConfig()
|
|
self.scheduler = Scheduler(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
config=scheduler_config,
|
|
stream=self._mlx_stream,
|
|
)
|
|
|
|
# Output collectors for low-latency streaming (vLLM pattern)
|
|
self._output_collectors: Dict[str, RequestOutputCollector] = {}
|
|
self._stream_states: Dict[str, RequestStreamState] = {}
|
|
self._finished_events: Dict[str, asyncio.Event] = {}
|
|
# Finish timestamps for orphan-collector reaping (#1154).
|
|
# Normally a consumer drains and removes its own collector, but if the
|
|
# client disconnects mid-stream the SSE generator chain is abandoned and
|
|
# its cleanup finally only runs at GC time, so the collector lingers and
|
|
# the dashboard keeps showing the request as "Generating".
|
|
self._finished_at: Dict[str, float] = {}
|
|
self._last_reap = 0.0
|
|
|
|
# Engine state
|
|
self._running = False
|
|
self._task: Optional[asyncio.Task] = None
|
|
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
|
self._wake_event: Optional[asyncio.Event] = None
|
|
self._start_time: Optional[float] = None
|
|
self._steps_executed = 0
|
|
|
|
# Drop transient aliases after ownership moves to the engine/scheduler
|
|
# graph, so close()/deep_reset() can make that graph unreachable.
|
|
model = None
|
|
tokenizer = None
|
|
|
|
logger.debug(f"Engine {self._engine_id} initialized")
|
|
|
|
async def start(self) -> None:
|
|
"""Start the engine loop."""
|
|
if self._running:
|
|
return
|
|
|
|
self._loop = asyncio.get_running_loop()
|
|
self._wake_event = asyncio.Event()
|
|
self._running = True
|
|
self._start_time = time.time()
|
|
self._task = asyncio.create_task(self._engine_loop())
|
|
logger.info("Engine started")
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the engine loop."""
|
|
self._running = False
|
|
if self._wake_event is not None:
|
|
self._wake_event.set()
|
|
if self._task:
|
|
self._task.cancel()
|
|
with suppress(asyncio.CancelledError):
|
|
await self._task
|
|
self._task = None
|
|
self._wake_event = None
|
|
self._loop = None
|
|
logger.info("Engine stopped")
|
|
|
|
def is_running(self) -> bool:
|
|
"""Check if engine is running."""
|
|
return self._running
|
|
|
|
def _wake_engine_loop(self) -> None:
|
|
"""Wake the idle engine loop after scheduler-visible state changes."""
|
|
event = getattr(self, "_wake_event", None)
|
|
loop = getattr(self, "_loop", None)
|
|
if event is None or loop is None or loop.is_closed():
|
|
return
|
|
try:
|
|
running_loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
running_loop = None
|
|
if running_loop is loop:
|
|
event.set()
|
|
else:
|
|
loop.call_soon_threadsafe(event.set)
|
|
|
|
def _step_burst(self) -> list:
|
|
"""Run scheduler.step() several times in one executor hand-off.
|
|
|
|
Each decode token otherwise bounces back to the event loop, which
|
|
ping-pongs the GIL with the asyncio loop + uvicorn on the main thread
|
|
(~1ms/token of contention). Chaining a few steps lets the MLX thread
|
|
hold the GIL continuously (in-process sync loop hits ~80 tok/s vs ~74
|
|
through the per-token async hand-off).
|
|
|
|
scheduler.step() services aborts/admission/finish every step, so
|
|
correctness is unchanged; the only cost is event-loop responsiveness,
|
|
bounded by decode_burst_budget_s. Stops early when no work remains, a
|
|
prefill eviction needs the (async) callback, or the budget elapses —
|
|
the budget also ends the burst when a slow prefill-chunk step lands.
|
|
|
|
Runs on the MLX executor thread. Returns the SchedulerOutputs in order.
|
|
"""
|
|
max_steps = self.config.decode_burst_max_steps
|
|
outputs = [self.scheduler.step()]
|
|
if max_steps <= 1:
|
|
return outputs
|
|
# Adaptive budget: single active request -> aggressive (nothing else to
|
|
# stay responsive to); concurrent -> tight to keep admission/abort low.
|
|
running = getattr(self.scheduler, "running", None)
|
|
single = running is None or len(running) <= 1
|
|
budget = (
|
|
self.config.decode_burst_budget_single_s
|
|
if single
|
|
else self.config.decode_burst_budget_s
|
|
)
|
|
if budget <= 0:
|
|
return outputs
|
|
deadline = time.monotonic() + budget
|
|
while len(outputs) < max_steps:
|
|
last = outputs[-1]
|
|
if (
|
|
not last.has_work # throttled/idle: stop and let the loop wait
|
|
or not self.scheduler.has_requests()
|
|
or last.prefill_eviction_request is not None
|
|
or time.monotonic() >= deadline
|
|
):
|
|
break
|
|
outputs.append(self.scheduler.step())
|
|
return outputs
|
|
|
|
async def _engine_loop(self) -> None:
|
|
"""Main engine loop - runs scheduler steps on the MLX executor.
|
|
|
|
All scheduler steps run on _mlx_executor (single-worker thread) to
|
|
guarantee that MLX GPU operations are never concurrent. VLM vision
|
|
encoding also runs on the same executor, so inline scheduler.step()
|
|
on the event loop would race with vision mx.eval() and segfault.
|
|
"""
|
|
loop = asyncio.get_running_loop()
|
|
|
|
step_interval = self.config.step_interval
|
|
stream_interval = self.config.stream_interval
|
|
use_simple_streaming = stream_interval == 1
|
|
|
|
while self._running:
|
|
try:
|
|
# Sweep collectors orphaned by client disconnects (throttled).
|
|
now = time.monotonic()
|
|
if now - self._last_reap >= 1.0:
|
|
self._last_reap = now
|
|
self._reap_orphaned_collectors(now)
|
|
|
|
if self.scheduler.has_requests():
|
|
step_outputs = await loop.run_in_executor(
|
|
self._mlx_executor, self._step_burst
|
|
)
|
|
self._steps_executed += len(step_outputs)
|
|
|
|
# Distribute every step's outputs to collectors (one or
|
|
# more decode tokens per burst). collector.put() runs on the
|
|
# loop thread, keeping the asyncio.Event signalling
|
|
# thread-safe and per-token streaming intact.
|
|
collectors = self._output_collectors
|
|
states = self._stream_states
|
|
eviction_request = None
|
|
distributed = False
|
|
|
|
for output in step_outputs:
|
|
if (
|
|
eviction_request is None
|
|
and output.prefill_eviction_request is not None
|
|
):
|
|
eviction_request = output.prefill_eviction_request
|
|
|
|
outputs = output.outputs
|
|
if not outputs:
|
|
continue
|
|
distributed = True
|
|
|
|
for req_output in outputs:
|
|
rid = req_output.request_id
|
|
collector = collectors.get(rid)
|
|
|
|
if collector is not None:
|
|
# Optimized: skip stream_interval check when interval=1
|
|
if use_simple_streaming:
|
|
collector.put(req_output)
|
|
else:
|
|
state = states.get(rid)
|
|
if state and state.should_send(
|
|
req_output.completion_tokens,
|
|
req_output.finished,
|
|
):
|
|
collector.put(req_output)
|
|
state.mark_sent(req_output.completion_tokens)
|
|
|
|
if req_output.finished:
|
|
self._mark_request_finished(rid)
|
|
# Cleanup normally happens in the consumer
|
|
# (stream_outputs()/generate()); collectors left
|
|
# behind by a disconnected client are swept by
|
|
# _reap_orphaned_collectors() via _finished_at.
|
|
|
|
if distributed:
|
|
# Always yield to prevent event loop starvation.
|
|
# Without this, orphaned requests (client disconnected but
|
|
# request still in scheduler) block the entire event loop,
|
|
# making the server unresponsive to all HTTP requests.
|
|
await asyncio.sleep(0)
|
|
|
|
if eviction_request is not None:
|
|
callback = self.config.prefill_eviction_callback
|
|
if callback is not None:
|
|
logger.info(
|
|
"Running prefill LRU eviction for request %s",
|
|
eviction_request.request_id,
|
|
)
|
|
evicted = await callback(eviction_request)
|
|
if evicted:
|
|
logger.info(
|
|
"Prefill LRU eviction completed for request %s",
|
|
eviction_request.request_id,
|
|
)
|
|
else:
|
|
logger.info(
|
|
"No idle model evicted for request %s; "
|
|
"scheduler will fall back to throttling",
|
|
eviction_request.request_id,
|
|
)
|
|
else:
|
|
logger.debug(
|
|
"Prefill eviction requested for %s but no callback "
|
|
"is configured",
|
|
eviction_request.request_id,
|
|
)
|
|
continue
|
|
if not step_outputs[-1].has_work:
|
|
# Requests may be queued while scheduler admission is
|
|
# intentionally throttled by async cache cleanup. Avoid
|
|
# spinning the engine loop, but still let new requests
|
|
# wake the wait immediately.
|
|
event = self._wake_event
|
|
if event is None:
|
|
await asyncio.sleep(step_interval)
|
|
else:
|
|
event.clear()
|
|
with suppress(TimeoutError):
|
|
await asyncio.wait_for(
|
|
event.wait(), timeout=step_interval
|
|
)
|
|
else:
|
|
event = self._wake_event
|
|
if event is None:
|
|
await asyncio.sleep(step_interval)
|
|
else:
|
|
event.clear()
|
|
# Avoid losing a request that arrived between
|
|
# has_requests() and clear().
|
|
if self.scheduler.has_requests():
|
|
continue
|
|
with suppress(TimeoutError):
|
|
await asyncio.wait_for(event.wait(), timeout=step_interval)
|
|
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
logger.error(f"Engine loop error: {e}\n{traceback.format_exc()}")
|
|
# Fail all requests and remove from scheduler to prevent
|
|
# infinite loop (has_requests() must return False).
|
|
failed_ids = await loop.run_in_executor(
|
|
self._mlx_executor, self.scheduler.fail_all_requests
|
|
)
|
|
for rid in failed_ids:
|
|
collector = self._output_collectors.get(rid)
|
|
if collector is not None:
|
|
collector.put(
|
|
RequestOutput(
|
|
request_id=rid,
|
|
finished=True,
|
|
finish_reason="error",
|
|
error=str(e),
|
|
)
|
|
)
|
|
self._mark_request_finished(rid)
|
|
await asyncio.sleep(0.1)
|
|
|
|
async def add_request(
|
|
self,
|
|
prompt: Union[str, List[int]],
|
|
sampling_params: Optional[SamplingParams] = None,
|
|
request_id: Optional[str] = None,
|
|
images: Optional[List[Any]] = None,
|
|
videos: Optional[List[Any]] = None,
|
|
vlm_inputs_embeds: Optional[Any] = None,
|
|
vlm_extra_kwargs: Optional[Dict[str, Any]] = None,
|
|
vlm_image_hash: Optional[str] = None,
|
|
vlm_cache_key_start: int = 0,
|
|
vlm_cache_key_ranges: Optional[List[Tuple[int, str]]] = None,
|
|
specprefill: Optional[bool] = None,
|
|
specprefill_keep_pct: Optional[float] = None,
|
|
specprefill_threshold: Optional[int] = None,
|
|
specprefill_system_end: Optional[int] = None,
|
|
) -> str:
|
|
"""
|
|
Add a request for processing.
|
|
|
|
Args:
|
|
prompt: Input prompt (string or token IDs)
|
|
sampling_params: Generation parameters
|
|
request_id: Optional custom request ID
|
|
images: Optional images for multimodal
|
|
videos: Optional videos for multimodal
|
|
vlm_inputs_embeds: Pre-computed vision+text embeddings for VLM
|
|
vlm_extra_kwargs: Model-specific VLM kwargs (e.g., position_ids)
|
|
vlm_image_hash: SHA256 hash of images for prefix cache
|
|
specprefill: Per-request SpecPrefill override (True/False/None)
|
|
specprefill_keep_pct: Per-request keep rate override
|
|
specprefill_threshold: Per-request threshold override (min tokens)
|
|
|
|
Returns:
|
|
The request ID
|
|
"""
|
|
if request_id is None:
|
|
request_id = str(uuid.uuid4())
|
|
|
|
if sampling_params is None:
|
|
sampling_params = SamplingParams()
|
|
|
|
request = Request(
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
sampling_params=sampling_params,
|
|
images=images,
|
|
videos=videos,
|
|
vlm_inputs_embeds=vlm_inputs_embeds,
|
|
vlm_extra_kwargs=vlm_extra_kwargs,
|
|
vlm_image_hash=vlm_image_hash,
|
|
vlm_cache_key_start=vlm_cache_key_start,
|
|
vlm_cache_key_ranges=vlm_cache_key_ranges,
|
|
)
|
|
|
|
# SpecPrefill: resolve per-request settings.
|
|
# The scheduler checks _specprefill_enabled to decide whether to score.
|
|
if specprefill is not None:
|
|
request._specprefill_enabled = specprefill
|
|
elif self.scheduler._specprefill_draft_model is not None:
|
|
# Draft model is loaded → enable by default
|
|
request._specprefill_enabled = True
|
|
if specprefill_keep_pct is not None:
|
|
request._specprefill_keep_pct = specprefill_keep_pct
|
|
if specprefill_threshold is not None:
|
|
request._specprefill_threshold = specprefill_threshold
|
|
if specprefill_system_end is not None and specprefill_system_end > 0:
|
|
request.specprefill_system_end = specprefill_system_end
|
|
|
|
# Setup output collector with stream_interval from config
|
|
self._output_collectors[request_id] = RequestOutputCollector(aggregate=True)
|
|
self._stream_states[request_id] = RequestStreamState(
|
|
stream_interval=self.config.stream_interval
|
|
)
|
|
self._finished_events[request_id] = asyncio.Event()
|
|
|
|
# Add to scheduler — route through the MLX executor so that
|
|
# prefix cache reconstruction (mx.load, mx.concatenate) never
|
|
# races with scheduler.step() on the Metal stream. See #95.
|
|
#
|
|
# The scheduler may raise (PrefillMemoryExceededError, or other
|
|
# validation errors) before the request enters self.waiting. In
|
|
# that case the consumer in stream_outputs / generate never sees
|
|
# the request_id and its finally-block cleanup never fires —
|
|
# without the explicit cleanup below the per-rejection leak
|
|
# accumulates one collector + one stream_state + one
|
|
# asyncio.Event per refused request. Re-raise after cleanup so
|
|
# the typed exception still reaches the FastAPI 400 handler.
|
|
loop = asyncio.get_running_loop()
|
|
try:
|
|
await loop.run_in_executor(
|
|
self._mlx_executor, self.scheduler.add_request, request
|
|
)
|
|
except BaseException:
|
|
# If the caller is cancelled here (e.g. the client disconnected
|
|
# before the SSE stream began) — or the insert fails — the request
|
|
# never reaches stream_outputs()/generate()'s try/finally, so
|
|
# nothing would mark it finished or clean it up. The collector
|
|
# created above would then linger forever as a phantom the reaper
|
|
# cannot see (it was never stamped _finished_at), and the dashboard
|
|
# would show it as "Generating" indefinitely (#1154).
|
|
# Drop the tracking and abort any partial scheduler insert (the
|
|
# deferred abort is idempotent and harmless if it never landed).
|
|
try:
|
|
self.scheduler.abort_request(request_id)
|
|
except Exception as abort_exc: # noqa: BLE001
|
|
logger.debug(
|
|
f"Abort of partial insert for {request_id} failed: {abort_exc}"
|
|
)
|
|
self._cleanup_request(request_id)
|
|
raise
|
|
self._wake_engine_loop()
|
|
|
|
return request_id
|
|
|
|
async def abort_request(self, request_id: str) -> bool:
|
|
"""Abort a request.
|
|
|
|
Uses deferred abort pattern: scheduler.abort_request() just enqueues
|
|
the request ID into a thread-safe set. The actual abort is processed
|
|
at the start of the next scheduler.step() call, ensuring it runs in
|
|
the same execution context as generation (no race conditions).
|
|
|
|
Signals the consumer (stream_outputs/generate) with an abort error
|
|
so it can exit gracefully. Cleanup is handled by the consumer's
|
|
finally block, NOT here -- calling _cleanup_request() immediately
|
|
after put() would clear the output before the consumer can drain it.
|
|
"""
|
|
scheduler = getattr(self, "scheduler", None)
|
|
if getattr(self, "_closed", False) or scheduler is None:
|
|
logger.debug(
|
|
"Skipping abort for request %s because engine is already closed",
|
|
request_id,
|
|
)
|
|
return False
|
|
|
|
result = scheduler.abort_request(request_id)
|
|
|
|
# Signal consumer with abort error so any waiting
|
|
# stream_outputs() / generate() can exit gracefully.
|
|
# Matches abort_all_requests() pattern.
|
|
collector = self._output_collectors.get(request_id)
|
|
if collector is not None:
|
|
collector.put(
|
|
RequestOutput(
|
|
request_id=request_id,
|
|
finished=True,
|
|
finish_reason="abort",
|
|
error="Request aborted",
|
|
)
|
|
)
|
|
self._mark_request_finished(request_id)
|
|
self._wake_engine_loop()
|
|
|
|
return result
|
|
|
|
async def abort_all_requests(self) -> int:
|
|
"""Abort all active requests without stopping the engine.
|
|
|
|
Sends error output to all active collectors and marks requests
|
|
for deferred abort in the scheduler. Cleanup is handled by
|
|
the consumer (stream_outputs/generate).
|
|
"""
|
|
from .utils.proc_memory import get_phys_footprint
|
|
|
|
request_ids = list(self._output_collectors.keys())
|
|
ceiling = 0
|
|
sched = self.scheduler
|
|
if sched is not None:
|
|
ceiling = int(getattr(sched, "_memory_hard_limit_bytes", 0) or 0)
|
|
usage = get_phys_footprint()
|
|
usage_gb = usage / (1024**3)
|
|
ceiling_gb = ceiling / (1024**3) if ceiling > 0 else 0.0
|
|
for rid in request_ids:
|
|
self.scheduler.abort_request(rid)
|
|
collector = self._output_collectors.get(rid)
|
|
if collector is not None:
|
|
if ceiling > 0:
|
|
error_msg = (
|
|
f"Request aborted: process memory limit exceeded "
|
|
f"(usage {usage_gb:.1f} GB, ceiling {ceiling_gb:.1f} GB). "
|
|
"Reduce context size or lower memory_guard_tier."
|
|
)
|
|
else:
|
|
error_msg = (
|
|
f"Request aborted: process memory limit exceeded "
|
|
f"(usage {usage_gb:.1f} GB). "
|
|
"Reduce context size or lower memory_guard_tier."
|
|
)
|
|
collector.put(
|
|
RequestOutput(
|
|
request_id=rid,
|
|
finished=True,
|
|
finish_reason="error",
|
|
new_text=f"\n\n[Error: {error_msg}]",
|
|
error=error_msg,
|
|
)
|
|
)
|
|
self._mark_request_finished(rid)
|
|
if request_ids:
|
|
logger.warning(
|
|
f"Aborted {len(request_ids)} requests due to memory pressure"
|
|
)
|
|
self._wake_engine_loop()
|
|
return len(request_ids)
|
|
|
|
def _mark_request_finished(self, request_id: str) -> None:
|
|
"""Signal the consumer a request finished and stamp the finish time.
|
|
|
|
The timestamp lets _reap_orphaned_collectors() drop collectors whose
|
|
consumer never cleaned up (e.g. the client disconnected mid-stream and
|
|
the SSE generator chain was abandoned rather than closed).
|
|
"""
|
|
self._finished_at.setdefault(request_id, time.monotonic())
|
|
event = self._finished_events.get(request_id)
|
|
if event is not None:
|
|
event.set()
|
|
|
|
def _reap_orphaned_collectors(self, now: float, grace: float = 5.0) -> int:
|
|
"""Drop tracking for finished requests whose consumer never cleaned up.
|
|
|
|
stream_outputs()/generate() normally remove their own collector via
|
|
_cleanup_request() once the final output is drained. But when a client
|
|
disconnects mid-stream the SSE generator chain is abandoned instead of
|
|
closed, so that finally block only runs at non-deterministic GC time —
|
|
the collector lingers in _output_collectors and the request shows as
|
|
"Generating" forever on the dashboard (#1154).
|
|
|
|
This sweep removes the dict tracking for any request finished more than
|
|
``grace`` seconds ago. It is intentionally pop-only and never calls
|
|
``collector.clear()``: stream_outputs()/generate() hold their own
|
|
reference to the collector object, so dropping the dict entry cannot
|
|
truncate a slow-but-live consumer's output — only an over-eager clear()
|
|
could (which is what made the earlier _delayed_cleanup() approach race).
|
|
The grace period guarantees a live consumer (which drains in the same
|
|
event-loop turn the request finishes) has already self-cleaned.
|
|
"""
|
|
if not self._finished_at:
|
|
return 0
|
|
stale = [rid for rid, ts in self._finished_at.items() if now - ts >= grace]
|
|
for rid in stale:
|
|
# pop-only — see docstring; never clear() the collector object.
|
|
self._output_collectors.pop(rid, None)
|
|
self._stream_states.pop(rid, None)
|
|
self._finished_events.pop(rid, None)
|
|
self._finished_at.pop(rid, None)
|
|
if stale:
|
|
logger.debug(
|
|
"Reaped %d orphaned output collector(s) after disconnect: %s",
|
|
len(stale),
|
|
stale,
|
|
)
|
|
return len(stale)
|
|
|
|
def _cleanup_request(self, request_id: str) -> None:
|
|
"""Clean up request tracking.
|
|
|
|
Only cleans engine-core level state (collectors, events).
|
|
Scheduler state is cleaned by _do_abort_request (deferred abort)
|
|
or _cleanup_finished (normal completion).
|
|
"""
|
|
collector = self._output_collectors.pop(request_id, None)
|
|
if collector:
|
|
collector.clear()
|
|
self._stream_states.pop(request_id, None)
|
|
self._finished_events.pop(request_id, None)
|
|
self._finished_at.pop(request_id, None)
|
|
|
|
async def _delayed_cleanup(self, request_id: str, delay: float = 5.0) -> None:
|
|
"""
|
|
Cleanup request after delay if not already cleaned.
|
|
|
|
This handles the case where a client disconnects before consuming
|
|
the stream_outputs() generator, which would prevent the finally
|
|
block from running.
|
|
"""
|
|
await asyncio.sleep(delay)
|
|
if request_id in self._output_collectors:
|
|
logger.debug(f"Delayed cleanup for request {request_id}")
|
|
self._cleanup_request(request_id)
|
|
|
|
async def stream_outputs(
|
|
self,
|
|
request_id: str,
|
|
timeout: Optional[float] = None,
|
|
) -> AsyncIterator[RequestOutput]:
|
|
"""
|
|
Stream outputs for a request with low-latency non-blocking pattern.
|
|
|
|
Uses the vLLM pattern: get_nowait() or await get()
|
|
This avoids unnecessary task switches when output is available.
|
|
|
|
Args:
|
|
request_id: The request ID
|
|
timeout: Optional timeout in seconds
|
|
|
|
Yields:
|
|
RequestOutput objects as tokens are generated
|
|
"""
|
|
collector = self._output_collectors.get(request_id)
|
|
if collector is None:
|
|
# Request might not be added yet or already cleaned up
|
|
return
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
# Non-blocking drain pattern from vLLM
|
|
# Try get_nowait first to avoid task switch if output ready
|
|
if timeout:
|
|
output = collector.get_nowait()
|
|
if output is None:
|
|
output = await asyncio.wait_for(
|
|
collector.get(), timeout=timeout
|
|
)
|
|
else:
|
|
output = collector.get_nowait() or await collector.get()
|
|
|
|
yield output
|
|
|
|
if output.error:
|
|
_raise_request_output_error(output)
|
|
|
|
if output.finished:
|
|
break
|
|
|
|
except asyncio.TimeoutError:
|
|
logger.warning(f"Timeout waiting for request {request_id}")
|
|
break
|
|
|
|
finally:
|
|
self._cleanup_request(request_id)
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: Union[str, List[int]],
|
|
sampling_params: Optional[SamplingParams] = None,
|
|
request_id: Optional[str] = None,
|
|
**kwargs,
|
|
) -> RequestOutput:
|
|
"""
|
|
Generate a complete response (non-streaming).
|
|
|
|
This method is optimized to avoid streaming overhead when
|
|
you only need the final result.
|
|
|
|
Args:
|
|
prompt: Input prompt
|
|
sampling_params: Generation parameters
|
|
request_id: Optional request ID
|
|
|
|
Returns:
|
|
Final RequestOutput with complete text
|
|
"""
|
|
request_id = await self.add_request(
|
|
prompt=prompt,
|
|
sampling_params=sampling_params,
|
|
request_id=request_id,
|
|
**kwargs,
|
|
)
|
|
|
|
# Wait for completion using event instead of streaming
|
|
# This avoids the waiting_consumer tracking overhead
|
|
event = self._finished_events.get(request_id)
|
|
if event is None:
|
|
raise RuntimeError(f"No event for request {request_id}")
|
|
|
|
# Capture the collector reference BEFORE awaiting, mirroring
|
|
# stream_outputs(): the orphan reaper is pop-only and may drop the dict
|
|
# entry once the request is finished, but a held reference still drains.
|
|
# Re-fetching after the await would race the reaper if this coroutine is
|
|
# starved past the grace window under heavy load (#1154).
|
|
collector = self._output_collectors.get(request_id)
|
|
if collector is None:
|
|
raise RuntimeError(f"No collector for request {request_id}")
|
|
|
|
try:
|
|
# Wait for the request to finish
|
|
await event.wait()
|
|
except asyncio.CancelledError:
|
|
# Client disconnected or task was cancelled - abort the request
|
|
# to free scheduler/GPU resources (prevents orphaned requests)
|
|
logger.info(f"Request {request_id} cancelled, aborting")
|
|
await self.abort_request(request_id)
|
|
self._cleanup_request(request_id)
|
|
raise
|
|
|
|
# Drain all outputs and get the last one (using the captured reference)
|
|
final_output = None
|
|
while True:
|
|
output = collector.get_nowait()
|
|
if output is None:
|
|
break
|
|
final_output = output
|
|
|
|
# Cleanup
|
|
self._cleanup_request(request_id)
|
|
|
|
if final_output is None:
|
|
raise RuntimeError(f"No output for request {request_id}")
|
|
|
|
if final_output.error:
|
|
_raise_request_output_error(final_output)
|
|
|
|
return final_output
|
|
|
|
def generate_batch_sync(
|
|
self,
|
|
prompts: List[Union[str, List[int]]],
|
|
sampling_params: Optional[SamplingParams] = None,
|
|
) -> List[RequestOutput]:
|
|
"""
|
|
Generate responses synchronously for maximum throughput.
|
|
|
|
This bypasses the async engine loop entirely, running the scheduler
|
|
directly for optimal batching performance. Use this when you don't
|
|
need streaming and want maximum throughput.
|
|
|
|
Args:
|
|
prompts: List of input prompts
|
|
sampling_params: Generation parameters (same for all)
|
|
|
|
Returns:
|
|
List of RequestOutput in same order as prompts
|
|
"""
|
|
import uuid as uuid_module
|
|
|
|
from .request import Request
|
|
|
|
if sampling_params is None:
|
|
sampling_params = SamplingParams()
|
|
|
|
# Add all requests to scheduler
|
|
request_ids = []
|
|
for prompt in prompts:
|
|
request_id = str(uuid_module.uuid4())
|
|
request = Request(
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
sampling_params=sampling_params,
|
|
)
|
|
self.scheduler.add_request(request)
|
|
request_ids.append(request_id)
|
|
|
|
# Process until all done - direct scheduler access, no async overhead
|
|
results: Dict[str, RequestOutput] = {}
|
|
while self.scheduler.has_requests():
|
|
output = self.scheduler.step()
|
|
for req_output in output.outputs:
|
|
if req_output.finished:
|
|
results[req_output.request_id] = req_output
|
|
|
|
# Cleanup
|
|
for rid in request_ids:
|
|
self.scheduler.remove_finished_request(rid)
|
|
|
|
# Return in original order
|
|
return [results[rid] for rid in request_ids]
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Get engine statistics."""
|
|
scheduler_stats = self.scheduler.get_stats()
|
|
uptime = time.time() - self._start_time if self._start_time else 0
|
|
|
|
return {
|
|
"running": self._running,
|
|
"uptime_seconds": uptime,
|
|
"steps_executed": self._steps_executed,
|
|
"active_requests": len(self._output_collectors),
|
|
"stream_interval": self.config.stream_interval,
|
|
**scheduler_stats,
|
|
}
|
|
|
|
def get_cache_stats(self) -> Optional[Dict[str, Any]]:
|
|
"""Get prefix cache statistics."""
|
|
return self.scheduler.get_cache_stats()
|
|
|
|
def _release_model(self) -> None:
|
|
"""Release model ownership."""
|
|
if self._owns_model and not self._closed:
|
|
registry = get_registry()
|
|
registry.release(self.model, self._engine_id)
|
|
self._owns_model = False
|
|
logger.debug(f"Engine {self._engine_id} released model ownership")
|
|
|
|
def close(self) -> None:
|
|
"""
|
|
Explicitly close the engine and release resources.
|
|
|
|
This should be called when done using the engine, especially
|
|
if you plan to create another engine with the same model.
|
|
"""
|
|
if self._closed:
|
|
return
|
|
|
|
# Release model ownership BEFORE setting _closed
|
|
# (_release_model checks not self._closed)
|
|
if self._owns_model:
|
|
registry = get_registry()
|
|
registry.release(self.model, self._engine_id)
|
|
self._owns_model = False
|
|
logger.debug(f"Engine {self._engine_id} released model ownership")
|
|
|
|
self._closed = True
|
|
|
|
# Both shutdown() and deep_reset() touch the engine stream (directly
|
|
# or via _drain_pending_async_removes / _do_abort_request). The
|
|
# stream is bound to the engine's executor thread, so dispatch both
|
|
# through the executor; fall back to a direct call if the executor
|
|
# is already shut down.
|
|
for fn in (self.scheduler.shutdown, self.scheduler.deep_reset):
|
|
fn_name = getattr(fn, "__name__", repr(fn))
|
|
try:
|
|
self._mlx_executor.submit(fn).result(timeout=FATAL_TEARDOWN_TIMEOUT_S)
|
|
except concurrent.futures.TimeoutError:
|
|
fatal_exit(
|
|
f"Engine teardown timed out after "
|
|
f"{FATAL_TEARDOWN_TIMEOUT_S:.0f}s while running "
|
|
f"{fn_name} for engine {self._engine_id}"
|
|
)
|
|
except RuntimeError:
|
|
try:
|
|
fn()
|
|
except RuntimeError:
|
|
pass
|
|
except Exception:
|
|
logger.warning(
|
|
"Engine %s: %s raised during close() fallback",
|
|
self._engine_id,
|
|
getattr(fn, "__name__", fn),
|
|
exc_info=True,
|
|
)
|
|
except Exception:
|
|
# A failing shutdown/deep_reset must not abort close(), or the
|
|
# SSD cache manager below stays open and its writer thread keeps
|
|
# the manager (and its hot cache) alive until restart.
|
|
logger.warning(
|
|
"Engine %s: %s raised during close()",
|
|
self._engine_id,
|
|
getattr(fn, "__name__", fn),
|
|
exc_info=True,
|
|
)
|
|
|
|
# Drop the last bound-method reference from the teardown loop before
|
|
# the final GC/reclaim pass below.
|
|
fn = None
|
|
|
|
# Guarantee the SSD cache manager is released even if shutdown() did not
|
|
# reach its own close() above. The manager's writer thread holds a strong
|
|
# reference to it, so an unclosed manager leaks until restart.
|
|
manager = getattr(self.scheduler, "paged_ssd_cache_manager", None)
|
|
if manager is not None:
|
|
try:
|
|
manager.close()
|
|
except Exception:
|
|
logger.warning(
|
|
"Engine %s: SSD cache manager close() failed during teardown",
|
|
self._engine_id,
|
|
exc_info=True,
|
|
)
|
|
self.scheduler.paged_ssd_cache_manager = None
|
|
manager = None
|
|
|
|
# Clear output collectors before dropping model/scheduler references so
|
|
# any request-side caches they retain are eligible for the final reclaim.
|
|
for collector in self._output_collectors.values():
|
|
collector.clear()
|
|
self._output_collectors.clear()
|
|
self._stream_states.clear()
|
|
self._finished_events.clear()
|
|
self._finished_at.clear()
|
|
|
|
release_model_resources = getattr(self.model, "release_resources", None)
|
|
if callable(release_model_resources):
|
|
try:
|
|
release_model_resources()
|
|
except Exception:
|
|
logger.warning(
|
|
"Engine %s: model resource release failed during close()",
|
|
self._engine_id,
|
|
exc_info=True,
|
|
)
|
|
release_model_resources = None
|
|
|
|
# Release model, tokenizer, and scheduler references before the final
|
|
# MLX reclaim. The reclaim must run on this engine's worker thread and
|
|
# stream; clearing on the global executor cannot reliably return this
|
|
# thread/stream-local Metal memory to MLX.
|
|
self.model = None
|
|
self.tokenizer = None
|
|
self.scheduler = None
|
|
|
|
if self._mlx_executor is not None:
|
|
try:
|
|
self._mlx_executor.submit(
|
|
_final_engine_thread_reclaim, self._mlx_stream
|
|
).result(timeout=FATAL_TEARDOWN_TIMEOUT_S)
|
|
except concurrent.futures.TimeoutError:
|
|
fatal_exit(
|
|
f"Engine teardown timed out after "
|
|
f"{FATAL_TEARDOWN_TIMEOUT_S:.0f}s while reclaiming "
|
|
f"MLX memory for engine {self._engine_id}"
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
except Exception:
|
|
logger.warning(
|
|
"Engine %s: final MLX reclaim raised during close()",
|
|
self._engine_id,
|
|
exc_info=True,
|
|
)
|
|
|
|
# MLX's @mx.compile cache is a C++ thread_local CompilerCache. If
|
|
# this worker thread exits with a non-empty cache, ~CompilerCache
|
|
# frees the cached graphs' Python objects from a thread-exit handler
|
|
# WITHOUT the GIL -> "PyThreadState_Get: GIL is released" crash for
|
|
# models with module-scope @mx.compile graphs (DeepSeek V4 unload,
|
|
# ml-explore/mlx #3280). Clear the cache ON this worker thread (GIL
|
|
# held) before the thread is torn down so the destructor runs on an
|
|
# empty cache, then request shutdown without waiting indefinitely.
|
|
# See utils/compile_cache.py.
|
|
if compile_cache_clear_available():
|
|
try:
|
|
self._mlx_executor.submit(clear_thread_compile_cache).result(
|
|
timeout=FATAL_TEARDOWN_TIMEOUT_S
|
|
)
|
|
except concurrent.futures.TimeoutError:
|
|
fatal_exit(
|
|
f"Engine teardown timed out after "
|
|
f"{FATAL_TEARDOWN_TIMEOUT_S:.0f}s while clearing "
|
|
f"MLX compile cache for engine {self._engine_id}"
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
self._mlx_executor.shutdown(wait=False)
|
|
else:
|
|
# Fallback: the clear symbol is unavailable, so do NOT exit the
|
|
# worker thread (that would run the unsafe destructor). Keep it
|
|
# alive for the process lifetime via the module-global registry.
|
|
_immortal_mlx_executors.append(self._mlx_executor)
|
|
_immortal_mlx_streams.append(self._mlx_stream)
|
|
self._mlx_executor = None
|
|
|
|
logger.debug(f"Engine {self._engine_id} closed")
|
|
|
|
def __del__(self):
|
|
"""Cleanup on destruction."""
|
|
try:
|
|
self._release_model()
|
|
except Exception:
|
|
# Ignore errors during garbage collection
|
|
pass
|
|
|
|
@property
|
|
def engine_id(self) -> str:
|
|
"""Get the engine ID."""
|
|
return self._engine_id
|
|
|
|
|
|
class AsyncEngineCore:
|
|
"""
|
|
Async context manager wrapper for EngineCore.
|
|
|
|
Usage:
|
|
async with AsyncEngineCore(model, tokenizer) as engine:
|
|
request_id = await engine.add_request("Hello")
|
|
async for output in engine.stream_outputs(request_id):
|
|
print(output.new_text)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: Any,
|
|
tokenizer: Any,
|
|
config: Optional[EngineConfig] = None,
|
|
):
|
|
self.engine = EngineCore(model, tokenizer, config)
|
|
# Drop wrapper-local aliases after EngineCore takes ownership.
|
|
model = None
|
|
tokenizer = None
|
|
|
|
@property
|
|
def _mlx_executor(self):
|
|
"""Expose the MLX executor for VLM vision encoding."""
|
|
return self.engine._mlx_executor
|
|
|
|
async def __aenter__(self) -> "AsyncEngineCore":
|
|
await self.engine.start()
|
|
return self
|
|
|
|
async def __aexit__(self, *args) -> None:
|
|
await self.stop()
|
|
|
|
def start(self) -> None:
|
|
"""Start engine (creates task in current loop)."""
|
|
asyncio.create_task(self.engine.start())
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the engine."""
|
|
engine = getattr(self, "engine", None)
|
|
if engine is None:
|
|
return
|
|
await engine.stop()
|
|
|
|
async def add_request(
|
|
self,
|
|
prompt: Union[str, List[int]],
|
|
sampling_params: Optional[SamplingParams] = None,
|
|
request_id: Optional[str] = None,
|
|
**kwargs,
|
|
) -> str:
|
|
"""Add a request."""
|
|
return await self.engine.add_request(
|
|
prompt=prompt,
|
|
sampling_params=sampling_params,
|
|
request_id=request_id,
|
|
**kwargs,
|
|
)
|
|
|
|
async def abort_request(self, request_id: str) -> bool:
|
|
"""Abort a request."""
|
|
engine = getattr(self, "engine", None)
|
|
if engine is None:
|
|
logger.debug(
|
|
"Skipping abort for request %s because async engine is closed",
|
|
request_id,
|
|
)
|
|
return False
|
|
return await engine.abort_request(request_id)
|
|
|
|
async def abort_all_requests(self) -> int:
|
|
"""Abort all active requests without stopping the engine."""
|
|
engine = getattr(self, "engine", None)
|
|
if engine is None:
|
|
return 0
|
|
return await engine.abort_all_requests()
|
|
|
|
async def stream_outputs(
|
|
self,
|
|
request_id: str,
|
|
timeout: Optional[float] = None,
|
|
) -> AsyncIterator[RequestOutput]:
|
|
"""Stream outputs."""
|
|
async for output in self.engine.stream_outputs(request_id, timeout):
|
|
yield output
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: Union[str, List[int]],
|
|
sampling_params: Optional[SamplingParams] = None,
|
|
**kwargs,
|
|
) -> RequestOutput:
|
|
"""Generate complete response."""
|
|
return await self.engine.generate(
|
|
prompt=prompt,
|
|
sampling_params=sampling_params,
|
|
**kwargs,
|
|
)
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Get engine stats."""
|
|
return self.engine.get_stats()
|
|
|
|
def get_cache_stats(self) -> Optional[Dict[str, Any]]:
|
|
"""Get prefix cache statistics."""
|
|
return self.engine.get_cache_stats()
|