ViperEkura
074642b6d2
- add BatchedStreamCallback sink type: TaskManager resolves a decode step's (task_id, token) events under one lock and delivers each sink a single list instead of one call per token - keep the plain Callable[[str]] callback contract: per-token callbacks still receive one call per event, and invoke_callback/cancel_task wrap single events for batched sinks - collect aborted, text, and finish STOP events in the scheduler decode loop and dispatch once per step instead of once per token - register one _ResultSink per generate call (replacing per-task closures) so GenerateResult takes its lock and wakes waiters once per step, with late-bind replay for tasks that start decoding before add_task returns their id - apply GenerateResult batches under a single condition hold via append_batch; append delegates to it - update engine test fakes to the batched contract and add coverage for event grouping, single-event dispatch, cancel STOP, and late-bind replay Benchmark: NVIDIA L20 (idle), CUDA 12.8, torch 2.11.0+cu128, 1.2B bf16 checkpoint, prompt 512, 256 greedy tokens, CUDA graph on, serving-level decode, 3 trials - batch 32: 7.808 -> 7.506 ms/token (4098 -> 4263 batch tok/s, +4.0%) - batch 1/8: unchanged within noise (3.768 -> 3.797 / 4.699 -> 4.607 ms/token) - full suite: 896 passed
366 lines
12 KiB
Python
366 lines
12 KiB
Python
"""Unified inference engine for continuous batching."""
|
|
|
|
import asyncio
|
|
import gc
|
|
import logging
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from astrai.extension import ATTN_BACKEND, AttentionBackend, get_backend
|
|
from astrai.inference.cache import PagePool
|
|
from astrai.inference.scheduler import InferenceScheduler
|
|
from astrai.inference.task import STOP, BatchedStreamCallback
|
|
from astrai.model import AutoModel
|
|
from astrai.tokenize import AutoTokenizer
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GenerateResult:
|
|
"""Thread-safe token accumulator for streaming and non-streaming modes."""
|
|
|
|
def __init__(self, count: int = 1):
|
|
self._cond = threading.Condition()
|
|
self._event = threading.Event()
|
|
self.tokens: List[Tuple[int, str]] = []
|
|
self.results: List[str] = [""] * count
|
|
self._done: List[bool] = [False] * count
|
|
self._completed = 0
|
|
self._total = count
|
|
|
|
def append(self, token: str, idx: int = 0):
|
|
self.append_batch([(idx, token)])
|
|
|
|
def append_batch(self, items: List[Tuple[int, Any]]) -> None:
|
|
"""Append multiple ``(idx, token)`` events under one lock/notify.
|
|
|
|
Batched counterpart to :meth:`append` for per-step delivery: state
|
|
updates for every event happen under a single condition hold and
|
|
waiters are woken once per batch instead of once per token.
|
|
"""
|
|
if not items:
|
|
return
|
|
with self._cond:
|
|
for idx, token in items:
|
|
self.tokens.append((idx, token))
|
|
if token is STOP:
|
|
if not self._done[idx]:
|
|
self._done[idx] = True
|
|
self._completed += 1
|
|
self._cond.notify_all()
|
|
else:
|
|
self.results[idx] += token
|
|
self._event.set()
|
|
|
|
def pop_all(self) -> List[Tuple[int, str]]:
|
|
with self._cond:
|
|
out = self.tokens.copy()
|
|
self.tokens.clear()
|
|
self._event.clear()
|
|
return out
|
|
|
|
def wait(self, timeout: Optional[float] = None) -> bool:
|
|
return self._event.wait(timeout=timeout)
|
|
|
|
def wait_completion(self, timeout: float = 300.0):
|
|
with self._cond:
|
|
if not self._cond.wait_for(
|
|
lambda: self._completed >= self._total, timeout=timeout
|
|
):
|
|
raise TimeoutError(
|
|
f"Generation timeout after {timeout}s "
|
|
f"({self._completed}/{self._total} completed)"
|
|
)
|
|
|
|
def get_results(self) -> List[str]:
|
|
with self._cond:
|
|
return self.results.copy()
|
|
|
|
|
|
class _ResultSink(BatchedStreamCallback):
|
|
"""Batched stream channel from the scheduler into one GenerateResult.
|
|
|
|
Registered as the ``stream_callback`` for every task of a single
|
|
``generate`` call, so the scheduler's one dispatch per decode step
|
|
maps to one ``append_batch`` (one lock, one waiter wake). A task can
|
|
start decoding the moment ``add_task`` returns — before the engine
|
|
learns its id — so events for ids not yet bound are buffered and
|
|
replayed on ``bind``.
|
|
"""
|
|
|
|
def __init__(self, result: GenerateResult):
|
|
self._result = result
|
|
self._lock = threading.Lock()
|
|
self._index_of: Dict[str, int] = {}
|
|
self._pending: List[Tuple[str, Any]] = []
|
|
|
|
def bind(self, task_id: str, idx: int) -> None:
|
|
with self._lock:
|
|
self._index_of[task_id] = idx
|
|
replay = [(idx, token) for tid, token in self._pending if tid == task_id]
|
|
if replay:
|
|
self._pending = [
|
|
(tid, token) for tid, token in self._pending if tid != task_id
|
|
]
|
|
if replay:
|
|
self._result.append_batch(replay)
|
|
|
|
def __call__(self, events: List[Tuple[str, Any]]) -> None:
|
|
with self._lock:
|
|
items: List[Tuple[int, Any]] = []
|
|
for tid, token in events:
|
|
idx = self._index_of.get(tid)
|
|
if idx is None:
|
|
self._pending.append((tid, token))
|
|
else:
|
|
items.append((idx, token))
|
|
if items:
|
|
self._result.append_batch(items)
|
|
|
|
|
|
class InferenceEngine:
|
|
"""Unified inference engine backed by continuous-batching scheduler."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: nn.Module,
|
|
tokenizer: AutoTokenizer,
|
|
max_batch_size: int = 1,
|
|
max_seq_len: Optional[int] = None,
|
|
cache: Optional[PagePool] = None,
|
|
enable_cuda_graph: bool = True,
|
|
backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
|
|
):
|
|
self.model = model
|
|
self.tokenizer = tokenizer
|
|
self.scheduler = InferenceScheduler(
|
|
model=self.model,
|
|
tokenizer=self.tokenizer,
|
|
max_batch_size=max_batch_size,
|
|
max_seq_len=max_seq_len,
|
|
cache=cache,
|
|
enable_cuda_graph=enable_cuda_graph,
|
|
backend=backend,
|
|
)
|
|
|
|
self.scheduler.start()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.shutdown()
|
|
return False
|
|
|
|
def generate(
|
|
self,
|
|
prompt: Union[str, List[str]],
|
|
stream: bool = False,
|
|
max_tokens: Optional[int] = None,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
top_k: int = 50,
|
|
frequency_penalty: float = 0.0,
|
|
rep_window: int = 64,
|
|
) -> Union[Generator, str, List[str]]:
|
|
is_batch = isinstance(prompt, list)
|
|
prompts = prompt if is_batch else [prompt]
|
|
|
|
if max_tokens is not None and max_tokens <= 0:
|
|
if stream:
|
|
return iter(())
|
|
results = [""] * len(prompts)
|
|
return results if is_batch else results[0]
|
|
|
|
return self._generate(
|
|
prompts,
|
|
is_batch,
|
|
stream,
|
|
max_tokens,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
frequency_penalty,
|
|
rep_window,
|
|
)
|
|
|
|
def generate_async(
|
|
self,
|
|
prompt: str,
|
|
max_tokens: Optional[int] = None,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
top_k: int = 50,
|
|
frequency_penalty: float = 0.0,
|
|
rep_window: int = 64,
|
|
) -> AsyncGenerator[str, None]:
|
|
request_backend = get_backend(use_default=False)
|
|
result = GenerateResult()
|
|
sink = _ResultSink(result)
|
|
task_id = self.scheduler.add_task(
|
|
prompt=prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
frequency_penalty=frequency_penalty,
|
|
rep_window=rep_window,
|
|
backend=request_backend,
|
|
stream_callback=sink,
|
|
)
|
|
sink.bind(task_id, 0)
|
|
|
|
async def _agen():
|
|
finished = False
|
|
try:
|
|
while not finished:
|
|
for _idx, token in result.pop_all():
|
|
if token is STOP:
|
|
finished = True
|
|
break
|
|
yield token
|
|
if not finished:
|
|
await asyncio.to_thread(result.wait, 0.05)
|
|
finally:
|
|
if not finished:
|
|
self.scheduler.cancel_task(task_id)
|
|
|
|
return _agen()
|
|
|
|
def _generate(
|
|
self,
|
|
prompts: List[str],
|
|
is_batch: bool,
|
|
stream: bool,
|
|
max_tokens: Optional[int],
|
|
temperature: float,
|
|
top_p: float,
|
|
top_k: int,
|
|
frequency_penalty: float,
|
|
rep_window: int,
|
|
) -> Union[Generator, str, List[str]]:
|
|
n = len(prompts)
|
|
request_backend = get_backend(use_default=False)
|
|
result = GenerateResult(count=n)
|
|
sink = _ResultSink(result)
|
|
task_ids = []
|
|
for i, p in enumerate(prompts):
|
|
task_id = self.scheduler.add_task(
|
|
prompt=p,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
frequency_penalty=frequency_penalty,
|
|
rep_window=rep_window,
|
|
backend=request_backend,
|
|
stream_callback=sink,
|
|
)
|
|
sink.bind(task_id, i)
|
|
task_ids.append(task_id)
|
|
|
|
if not stream:
|
|
try:
|
|
result.wait_completion()
|
|
except TimeoutError:
|
|
for tid in task_ids:
|
|
self.scheduler.cancel_task(tid)
|
|
raise
|
|
res = result.get_results()
|
|
return res if is_batch else res[0]
|
|
|
|
remaining = n
|
|
finished = [False] * n
|
|
|
|
def gen():
|
|
nonlocal remaining
|
|
try:
|
|
while remaining > 0:
|
|
items = result.pop_all()
|
|
for idx, token in items:
|
|
if token is STOP:
|
|
if not finished[idx]:
|
|
finished[idx] = True
|
|
remaining -= 1
|
|
else:
|
|
yield (idx, token) if is_batch else token
|
|
if remaining > 0:
|
|
result.wait(timeout=0.05)
|
|
finally:
|
|
for idx, task_id in enumerate(task_ids):
|
|
if not finished[idx]:
|
|
self.scheduler.cancel_task(task_id)
|
|
|
|
return gen()
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
return self.scheduler.get_stats()
|
|
|
|
@property
|
|
def backend_name(self) -> str:
|
|
return self.scheduler.backend_name
|
|
|
|
@property
|
|
def cuda_graph_enabled(self) -> bool:
|
|
return self.scheduler.cuda_graph_enabled
|
|
|
|
def shutdown(self):
|
|
self.scheduler.stop()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
|
|
def build_engine(
|
|
param_path: Optional[Union[str, Path]] = None,
|
|
*,
|
|
model: Optional[nn.Module] = None,
|
|
tokenizer: Optional[AutoTokenizer] = None,
|
|
device: Optional[str] = "cuda",
|
|
dtype: Optional[torch.dtype] = torch.bfloat16,
|
|
max_batch_size: int = 16,
|
|
max_seq_len: Optional[int] = None,
|
|
**engine_kwargs: Any,
|
|
) -> InferenceEngine:
|
|
"""Composition root for inference assembly.
|
|
|
|
Loads model and tokenizer from *param_path*, or accepts preloaded
|
|
objects, places the model, and returns a started InferenceEngine.
|
|
Extra *engine_kwargs* (cache, enable_cuda_graph, backend) pass
|
|
through to InferenceEngine. Placement parts left as None are skipped.
|
|
"""
|
|
if param_path is not None:
|
|
if model is not None or tokenizer is not None:
|
|
raise ValueError("pass either param_path or model+tokenizer, not both")
|
|
path = Path(param_path)
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Parameter directory not found: {path}")
|
|
tokenizer = AutoTokenizer.from_pretrained(path)
|
|
model = AutoModel.from_pretrained(path)
|
|
elif model is None or tokenizer is None:
|
|
raise ValueError("build_engine requires param_path or both model and tokenizer")
|
|
|
|
placement: Dict[str, Any] = {}
|
|
if device is not None:
|
|
placement["device"] = device
|
|
if dtype is not None:
|
|
placement["dtype"] = dtype
|
|
if placement:
|
|
model.to(**placement)
|
|
logger.info(
|
|
f"Model placed on {placement.get('device')} "
|
|
f"with dtype {placement.get('dtype')}"
|
|
)
|
|
|
|
return InferenceEngine(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
max_batch_size=max_batch_size,
|
|
max_seq_len=max_seq_len,
|
|
**engine_kwargs,
|
|
)
|