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
459 lines
18 KiB
Python
459 lines
18 KiB
Python
import logging
|
|
import threading
|
|
import uuid
|
|
from contextlib import nullcontext
|
|
from functools import wraps
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
|
|
|
|
import torch
|
|
|
|
from astrai.extension import (
|
|
ATTN_BACKEND,
|
|
AttentionBackend,
|
|
attn_backend,
|
|
get_backend,
|
|
)
|
|
from astrai.inference.cache import PagePool, TaskCacheManager
|
|
from astrai.inference.metrics import MetricsCollector
|
|
from astrai.inference.runtime.executor import Executor
|
|
from astrai.inference.runtime.stepper import Stepper
|
|
from astrai.inference.task import (
|
|
STOP,
|
|
GenerationResult,
|
|
Task,
|
|
TaskManager,
|
|
TaskStatus,
|
|
)
|
|
from astrai.inference.versioning import PolicyVersionGuard
|
|
from astrai.model.automodel import AutoModel
|
|
from astrai.tokenize.tokenizer import AutoTokenizer
|
|
|
|
logger = logging.getLogger(__name__)
|
|
T = TypeVar("T")
|
|
|
|
|
|
def _with_weight_lock(method):
|
|
@wraps(method)
|
|
def synchronized(self, *args, **kwargs):
|
|
with self._weight_lock:
|
|
return method(self, *args, **kwargs)
|
|
|
|
return synchronized
|
|
|
|
|
|
class InferenceScheduler:
|
|
"""Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: AutoModel,
|
|
tokenizer: AutoTokenizer,
|
|
max_batch_size: int = 16,
|
|
max_seq_len: Optional[int] = None,
|
|
device: Optional[str] = None,
|
|
dtype: Optional[torch.dtype] = None,
|
|
cache: Optional[PagePool] = None,
|
|
enable_cuda_graph: bool = True,
|
|
backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
|
|
policy_version: int = 0,
|
|
):
|
|
if (
|
|
isinstance(policy_version, bool)
|
|
or not isinstance(policy_version, int)
|
|
or policy_version < 0
|
|
):
|
|
raise ValueError("policy_version must be a non-negative integer")
|
|
config = model.config
|
|
|
|
if max_seq_len is not None:
|
|
self.max_seq_len = max_seq_len
|
|
elif config.max_position_embeddings is not None:
|
|
self.max_seq_len = config.max_position_embeddings
|
|
else:
|
|
raise ValueError(
|
|
"max_seq_len must be provided either as argument "
|
|
"or in model config (config.max_position_embeddings)"
|
|
)
|
|
self.device = device or next(model.parameters()).device
|
|
self.dtype = dtype or next(model.parameters()).dtype
|
|
|
|
head_dim = config.hidden_size // config.num_attention_heads
|
|
|
|
if cache is not None:
|
|
self._cache = cache
|
|
else:
|
|
self._cache = PagePool(
|
|
n_layers=config.num_hidden_layers,
|
|
n_kv_heads=config.num_key_value_heads,
|
|
head_dim=head_dim,
|
|
max_batch_size=max_batch_size,
|
|
max_seq_len=self.max_seq_len,
|
|
device=self.device,
|
|
dtype=self.dtype,
|
|
)
|
|
|
|
self._metrics = MetricsCollector()
|
|
|
|
self._task_cache = TaskCacheManager(self._cache)
|
|
|
|
self._task_mgr = TaskManager(
|
|
tokenizer=tokenizer,
|
|
max_batch_size=max_batch_size,
|
|
max_seq_len=self.max_seq_len,
|
|
metrics=self._metrics,
|
|
)
|
|
|
|
if backend is None:
|
|
self._backend = None
|
|
active_backend = get_backend()
|
|
else:
|
|
active_backend = backend
|
|
with attn_backend(active_backend):
|
|
if backend is not None:
|
|
self._backend = get_backend()
|
|
self._backend_name = type(get_backend()).__name__
|
|
self._executor = Executor(
|
|
model=model,
|
|
kv_cache=self._cache,
|
|
task_cache=self._task_cache,
|
|
device=self.device,
|
|
dtype=self.dtype,
|
|
enable_cuda_graph=enable_cuda_graph,
|
|
)
|
|
|
|
self._stepper = Stepper(
|
|
self._cache, self._task_cache, self._executor, self._metrics
|
|
)
|
|
|
|
self._stop_event = threading.Event()
|
|
self._loop_thread: Optional[threading.Thread] = None
|
|
self._policy_guard = PolicyVersionGuard(
|
|
policy_version,
|
|
ensure_ready=self._ensure_weight_update_ready,
|
|
on_commit=self._task_cache.invalidate_cache,
|
|
)
|
|
# Synchronous generation shares the guard's generation/weight mutex.
|
|
self._weight_lock = self._policy_guard.lock
|
|
|
|
@property
|
|
def policy_version(self) -> int:
|
|
"""Version of the model weights used for subsequent generations."""
|
|
return self._policy_guard.policy_version
|
|
|
|
def _ensure_weight_update_ready(self) -> None:
|
|
"""Check weight update preconditions. Must be called under the lock."""
|
|
if self._loop_thread is not None and self._loop_thread.is_alive():
|
|
raise RuntimeError("Stop the scheduler before updating model weights")
|
|
if self._task_mgr.get_active_tasks() or self._task_mgr.get_waiting_tasks():
|
|
raise RuntimeError("Cannot update model weights while tasks are queued")
|
|
|
|
def update_weights(self, policy_version: int) -> int:
|
|
"""Acknowledge an in-place weight update and invalidate stale KV state.
|
|
|
|
The scheduler owns the same model object as the in-process trainer, so
|
|
weights have already changed when this method is called. The explicit
|
|
version update makes that lifecycle visible and prevents prefix KV
|
|
entries produced by older weights from being reused.
|
|
"""
|
|
return self._policy_guard.update_weights(policy_version)
|
|
|
|
def apply_weight_update(
|
|
self, policy_version: Optional[int], update: Callable[[], T]
|
|
) -> T:
|
|
"""Mutate shared weights and publish their version without generation.
|
|
|
|
``policy_version=None`` derives ``live + 1`` under the same lock, for
|
|
callers that only need "advance by one" (e.g. ``optimizer.step()``)
|
|
without a read-compute-write race on the current version.
|
|
"""
|
|
return self._policy_guard.apply_weight_update(policy_version, update)
|
|
|
|
def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T:
|
|
"""Inspect state while the scheduler's policy version remains stable."""
|
|
return self._policy_guard.with_policy_snapshot(inspect)
|
|
|
|
def add_task(self, prompt: str, **kwargs) -> str:
|
|
return self._task_mgr.add_task(prompt, **kwargs)
|
|
|
|
def cancel_task(self, task_id: str) -> bool:
|
|
"""Cancel a waiting or active task without freeing in-use KV state."""
|
|
immediate, cancelled = self._task_mgr.cancel_task(task_id)
|
|
for task in immediate:
|
|
self._metrics.mark_finished(
|
|
task.task_id, task.input_tokens, task.output_tokens
|
|
)
|
|
if cancelled:
|
|
self._task_mgr.wake()
|
|
return cancelled
|
|
|
|
def remove_task(self, task_id: str) -> bool:
|
|
"""Backward-compatible alias for cancellation."""
|
|
return self.cancel_task(task_id)
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
stats = self._task_mgr.get_stats()
|
|
stats["kv_cache_tasks"] = self._task_cache.task_count
|
|
stats["policy_version"] = self._policy_guard.policy_version
|
|
return stats
|
|
|
|
@property
|
|
def backend_name(self) -> str:
|
|
return self._backend_name
|
|
|
|
@property
|
|
def cuda_graph_enabled(self) -> bool:
|
|
return self._executor.cuda_graph_enabled
|
|
|
|
def _backend_context(self):
|
|
if self._backend is None:
|
|
return nullcontext()
|
|
return attn_backend(self._backend)
|
|
|
|
def _step(
|
|
self, tasks: List[Task], return_logprobs: bool = False
|
|
) -> Tuple[List[Task], List[Task]]:
|
|
"""Advance every active task by one token; see :class:`Stepper`."""
|
|
return self._stepper.step(tasks, return_logprobs=return_logprobs)
|
|
|
|
def _run_generation_loop(self):
|
|
stop_ids = self._task_mgr.tokenizer.stop_ids
|
|
try:
|
|
with self._backend_context():
|
|
while not self._stop_event.is_set():
|
|
finished = self._task_mgr.remove_finished_tasks(stop_ids)
|
|
for task in finished:
|
|
if task.status == TaskStatus.FINISHED:
|
|
self._task_cache.task_record_hashes(
|
|
task.task_id,
|
|
self._task_cache.task_cacheable_ids(
|
|
task.task_id, task.prompt_ids, task.output_ids
|
|
),
|
|
)
|
|
self._task_cache.task_free(task.task_id)
|
|
|
|
active = self._task_mgr.get_active_tasks()
|
|
available = self._task_mgr.max_batch_size - len(active)
|
|
if available > 0:
|
|
candidates = self._task_mgr.pull_candidates(available)
|
|
failed = []
|
|
for task in candidates:
|
|
if self._task_cache.task_alloc(
|
|
task.task_id, task.prompt_ids
|
|
):
|
|
if not self._task_mgr.activate(task):
|
|
self._task_cache.task_free(task.task_id)
|
|
self._metrics.mark_finished(
|
|
task.task_id,
|
|
task.input_tokens,
|
|
task.output_tokens,
|
|
)
|
|
else:
|
|
failed.append(task)
|
|
if failed:
|
|
self._task_mgr.return_to_waiting(failed)
|
|
|
|
if not self._task_mgr.has_work():
|
|
self._task_mgr.wait_for_tasks(timeout=1.0)
|
|
continue
|
|
|
|
active = [
|
|
task
|
|
for task in self._task_mgr.get_active_tasks()
|
|
if task.status != TaskStatus.ABORTED
|
|
]
|
|
|
|
decoded, aborted = self._stepper.step(active)
|
|
|
|
# One dispatch per step: batch-aware sinks take their
|
|
# lock (and wake waiters) once instead of once per token.
|
|
events: List[Tuple[str, Any]] = [(t.task_id, STOP) for t in aborted]
|
|
for t in decoded:
|
|
if t.status == TaskStatus.ABORTED:
|
|
continue
|
|
new_text = t.decode_new_token(self._task_mgr.tokenizer)
|
|
if new_text:
|
|
events.append((t.task_id, new_text))
|
|
if t.is_finished(stop_ids):
|
|
events.append((t.task_id, STOP))
|
|
if events:
|
|
self._task_mgr.invoke_callbacks(events)
|
|
|
|
except Exception as e:
|
|
self._stop_event.set()
|
|
logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
|
|
self._abort_and_clear(free_waiting=False)
|
|
|
|
def start(self):
|
|
if self._loop_thread is not None and self._loop_thread.is_alive():
|
|
return
|
|
self._stop_event.clear()
|
|
t = threading.Thread(target=self._run_generation_loop, daemon=True)
|
|
t.start()
|
|
self._loop_thread = t
|
|
|
|
def stop(self):
|
|
self._stop_event.set()
|
|
self._task_mgr.wake()
|
|
if self._loop_thread is not None:
|
|
self._loop_thread.join(timeout=2.0)
|
|
self._loop_thread = None
|
|
self._abort_and_clear(free_waiting=True)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
def _abort_and_clear(self, free_waiting: bool):
|
|
"""Invoke STOP callbacks, release cache slots, and clear task queues."""
|
|
active = self._task_mgr.get_active_tasks()
|
|
waiting = self._task_mgr.get_waiting_tasks()
|
|
for task in active:
|
|
self._task_mgr.invoke_callback(task.task_id, STOP)
|
|
self._task_cache.task_free(task.task_id)
|
|
self._metrics.mark_finished(
|
|
task.task_id, task.input_tokens, task.output_tokens
|
|
)
|
|
for task in waiting:
|
|
self._task_mgr.invoke_callback(task.task_id, STOP)
|
|
if free_waiting:
|
|
self._task_cache.task_free(task.task_id)
|
|
self._metrics.mark_finished(
|
|
task.task_id, task.input_tokens, task.output_tokens
|
|
)
|
|
self._task_mgr.clear_queues()
|
|
|
|
@_with_weight_lock
|
|
def run_batch(
|
|
self,
|
|
prompt_ids_list: List[List[int]],
|
|
*,
|
|
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,
|
|
return_logprobs: bool = False,
|
|
return_details: bool = False,
|
|
) -> List[Any]:
|
|
"""Synchronous batch generation without the scheduler thread.
|
|
|
|
Accepts already-tokenized prompts (no string round-trip) and runs
|
|
prefill + decode to completion on the calling thread. Designed for
|
|
RL rollout, where logprobs of the behaviour policy must be collected
|
|
alongside generated tokens.
|
|
|
|
Args:
|
|
prompt_ids_list: ``B`` prompts, each a list of token IDs.
|
|
max_tokens: Maximum tokens to generate per prompt. ``None``
|
|
uses ``self.max_seq_len - len(prompt_ids)``.
|
|
temperature/top_p/top_k/frequency_penalty/rep_window: Sampling
|
|
parameters (uniform across the batch).
|
|
return_logprobs: If ``True``, return ``(token_ids, logprobs)``
|
|
tuples per prompt (logprobs aligned 1-to-1 with token_ids).
|
|
return_details: If ``True``, return a structured result per prompt
|
|
with terminal and error reasons. Logprobs are populated when
|
|
``return_logprobs`` is also ``True``.
|
|
|
|
Returns:
|
|
Structured results when ``return_details`` is ``True``;
|
|
otherwise generated token IDs per prompt, or token/logprob tuples
|
|
when ``return_logprobs`` is ``True``.
|
|
"""
|
|
stop_ids = self._task_mgr.tokenizer.stop_ids
|
|
seq_cap = self.max_seq_len
|
|
request_backend = get_backend(use_default=False)
|
|
|
|
tasks: List[Optional[Task]] = []
|
|
error_reasons: List[Optional[str]] = []
|
|
for ids in prompt_ids_list:
|
|
if not ids:
|
|
tasks.append(None)
|
|
error_reasons.append("prompt_empty")
|
|
continue
|
|
if len(ids) >= seq_cap:
|
|
tasks.append(None)
|
|
error_reasons.append("prompt_too_long")
|
|
continue
|
|
t_max = max_tokens
|
|
if t_max is None:
|
|
t_max = seq_cap - len(ids)
|
|
else:
|
|
t_max = min(t_max, seq_cap - len(ids))
|
|
if t_max <= 0:
|
|
tasks.append(None)
|
|
error_reasons.append("max_tokens_non_positive")
|
|
continue
|
|
task = Task(
|
|
task_id=f"batch_{uuid.uuid4().hex[:8]}",
|
|
prompt_ids=list(ids),
|
|
max_tokens=t_max,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
frequency_penalty=frequency_penalty,
|
|
rep_window=rep_window,
|
|
backend=request_backend,
|
|
)
|
|
if not self._task_cache.task_alloc(task.task_id, task.prompt_ids):
|
|
tasks.append(None)
|
|
error_reasons.append("kv_cache_allocation_failed")
|
|
continue
|
|
task.input_tokens = len(task.prompt_ids)
|
|
self._metrics.register(task.task_id)
|
|
tasks.append(task)
|
|
error_reasons.append(None)
|
|
|
|
runtime_errors: Dict[str, str] = {}
|
|
try:
|
|
live = [t for t in tasks if t is not None]
|
|
|
|
with self._backend_context():
|
|
while live:
|
|
decoded, aborted = self._stepper.step(
|
|
live, return_logprobs=return_logprobs
|
|
)
|
|
for task in aborted:
|
|
runtime_errors[task.task_id] = "kv_cache_extension_failed"
|
|
live = [t for t in decoded if not t.is_finished(stop_ids)]
|
|
finally:
|
|
for t in tasks:
|
|
if t is not None:
|
|
self._metrics.mark_finished(
|
|
t.task_id, t.input_tokens, t.output_tokens
|
|
)
|
|
self._task_cache.task_free(t.task_id)
|
|
|
|
details: List[GenerationResult] = []
|
|
for t, setup_error in zip(tasks, error_reasons):
|
|
if t is None:
|
|
details.append(
|
|
GenerationResult(
|
|
token_ids=[],
|
|
logprobs=[],
|
|
finish_reason="rejected",
|
|
error_reason=setup_error,
|
|
)
|
|
)
|
|
else:
|
|
runtime_error = runtime_errors.get(t.task_id)
|
|
stopped = bool(t.output_ids and t.output_ids[-1] in stop_ids)
|
|
if runtime_error:
|
|
finish_reason = "rejected"
|
|
elif stopped:
|
|
finish_reason = "stop"
|
|
else:
|
|
finish_reason = "length"
|
|
details.append(
|
|
GenerationResult(
|
|
token_ids=list(t.output_ids),
|
|
logprobs=list(t.output_logprobs),
|
|
finish_reason=finish_reason,
|
|
error_reason=runtime_error,
|
|
)
|
|
)
|
|
|
|
if return_details:
|
|
return details
|
|
if return_logprobs:
|
|
return [(result.token_ids, result.logprobs) for result in details]
|
|
return [result.token_ids for result in details]
|