ViperEkura
47b3ed4e44
- InferenceEngine/Scheduler accept an explicit backend - capture request-level attn_backend context onto Task - split prefill/decode batches by backend instance - ASTR_BACKEND env overrides ContextVar as process-wide policy - report resolved backend and CUDA-graph state in benchmark
291 lines
9.1 KiB
Python
291 lines
9.1 KiB
Python
import threading
|
|
import time
|
|
import uuid
|
|
from collections import deque
|
|
from enum import Enum
|
|
from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional
|
|
|
|
from tokenizers.decoders import DecodeStream
|
|
|
|
from astrai.inference.metrics import MetricsCollector
|
|
from astrai.tokenize.tokenizer import AutoTokenizer
|
|
|
|
if TYPE_CHECKING:
|
|
from astrai.extension import AttentionBackend
|
|
|
|
STOP = object()
|
|
|
|
|
|
class StreamDecoder:
|
|
"""Incremental decoder backed by the tokenizers library's DecodeStream.
|
|
|
|
Delegates to the Rust-native streaming decoder which maintains an
|
|
O(1) bounded token buffer internally (via prefix drain), avoiding
|
|
the O(n²) cost of re-decoding the full history on each step.
|
|
|
|
Multi-byte UTF-8 sequences split across token boundaries are
|
|
buffered until complete; ``push`` returns "" while the trailing
|
|
sequence is still incomplete.
|
|
"""
|
|
|
|
__slots__ = ("_stream", "_tok")
|
|
|
|
def __init__(self, tokenizer: AutoTokenizer):
|
|
self._tok = tokenizer._tokenizer
|
|
self._stream = DecodeStream(skip_special_tokens=True)
|
|
|
|
def push(self, token_id: int) -> str:
|
|
"""Append a token ID and return newly completed text.
|
|
|
|
Returns "" while a multi-byte character is still incomplete.
|
|
"""
|
|
chunk = self._stream.step(self._tok, token_id)
|
|
return chunk or ""
|
|
|
|
|
|
class TaskStatus(Enum):
|
|
"""Task lifecycle states."""
|
|
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
FINISHED = "finished"
|
|
ABORTED = "aborted"
|
|
|
|
|
|
class Task:
|
|
"""Single generation request: prompt, sampling params, output state."""
|
|
|
|
def __init__(
|
|
self,
|
|
task_id: str,
|
|
prompt_ids: 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,
|
|
backend: Optional["AttentionBackend"] = None,
|
|
):
|
|
self.task_id = task_id
|
|
self.prompt_ids = prompt_ids
|
|
self.max_tokens = max_tokens
|
|
self.temperature = temperature
|
|
self.top_p = top_p
|
|
self.top_k = top_k
|
|
self.frequency_penalty = frequency_penalty
|
|
self.rep_window = rep_window
|
|
self.backend = backend
|
|
|
|
self.status = TaskStatus.PENDING
|
|
self.output_ids: List[int] = []
|
|
self.output_logprobs: List[float] = []
|
|
self.input_tokens: int = 0
|
|
self.output_tokens: int = 0
|
|
self._kv_len: int = 0
|
|
self._decoder: Optional[StreamDecoder] = None
|
|
|
|
def mark_prefill_done(self):
|
|
"""Prompt KV is materialized by prefill; first output sampled but
|
|
not yet written to KV."""
|
|
self._kv_len = self.input_tokens
|
|
|
|
def advance_kv(self):
|
|
"""One more position written to KV (after a decode forward)."""
|
|
self._kv_len += 1
|
|
|
|
def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
|
|
"""Decode the last appended output token, buffering incomplete
|
|
multi-byte sequences across calls.
|
|
|
|
Lazily creates a :class:`StreamDecoder` on first use.
|
|
"""
|
|
if self._decoder is None:
|
|
self._decoder = StreamDecoder(tokenizer)
|
|
return self._decoder.push(self.output_ids[-1])
|
|
|
|
@property
|
|
def next_pos(self) -> int:
|
|
"""KV position where the next decode step will write."""
|
|
return self._kv_len
|
|
|
|
@property
|
|
def prefill_done(self) -> bool:
|
|
"""True when all prompt KV entries are materialized."""
|
|
return self._kv_len >= self.input_tokens > 0
|
|
|
|
def is_finished(self, stop_ids: List[int]) -> bool:
|
|
if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
|
|
return True
|
|
if self.output_ids and self.output_ids[-1] in stop_ids:
|
|
return True
|
|
return False
|
|
|
|
|
|
class TaskManager:
|
|
"""Thread-safe task queues and lifecycle transitions (no page ops)."""
|
|
|
|
def __init__(
|
|
self,
|
|
tokenizer: AutoTokenizer,
|
|
max_batch_size: int = 16,
|
|
max_seq_len: int = 8192,
|
|
metrics: Optional["MetricsCollector"] = None,
|
|
):
|
|
self.tokenizer = tokenizer
|
|
self.max_batch_size = max_batch_size
|
|
self.max_seq_len = max_seq_len
|
|
|
|
self.waiting_queue: Deque[Task] = deque()
|
|
self.active_tasks: List[Task] = []
|
|
self._callbacks: Dict[str, Callable[[str], None]] = {}
|
|
|
|
self._task_event = threading.Event()
|
|
self._lock = threading.Lock()
|
|
|
|
self._total_tasks = 0
|
|
self._total_tokens = 0
|
|
|
|
self._metrics = metrics
|
|
|
|
def add_task(
|
|
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,
|
|
backend: Optional["AttentionBackend"] = None,
|
|
stream_callback: Optional[Callable[[str], None]] = None,
|
|
) -> str:
|
|
task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
|
|
prompt_ids = self.tokenizer.encode(prompt)
|
|
if len(prompt_ids) > self.max_seq_len:
|
|
prompt_ids = prompt_ids[-self.max_seq_len :]
|
|
|
|
if max_tokens is None:
|
|
max_tokens = self.max_seq_len - len(prompt_ids)
|
|
else:
|
|
max_tokens = min(max_tokens, self.max_seq_len - len(prompt_ids))
|
|
|
|
task = Task(
|
|
task_id=task_id,
|
|
prompt_ids=prompt_ids,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
frequency_penalty=frequency_penalty,
|
|
rep_window=rep_window,
|
|
backend=backend,
|
|
)
|
|
|
|
with self._lock:
|
|
self.waiting_queue.append(task)
|
|
self._total_tasks += 1
|
|
if stream_callback:
|
|
self._callbacks[task_id] = stream_callback
|
|
|
|
if self._metrics is not None:
|
|
self._metrics.register(task_id)
|
|
|
|
self._task_event.set()
|
|
return task_id
|
|
|
|
def remove_task(self, task_id: str) -> List[Task]:
|
|
with self._lock:
|
|
removed_active = [t for t in self.active_tasks if t.task_id == task_id]
|
|
self.waiting_queue = deque(
|
|
t for t in self.waiting_queue if t.task_id != task_id
|
|
)
|
|
self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
|
|
self._callbacks.pop(task_id, None)
|
|
return removed_active
|
|
|
|
def invoke_callback(self, task_id: str, token: str):
|
|
cb = self._callbacks.get(task_id)
|
|
if cb:
|
|
cb(token)
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
stats: Dict[str, Any] = {
|
|
"total_tasks": self._total_tasks,
|
|
"total_tokens": self._total_tokens,
|
|
"active_tasks": len(self.active_tasks),
|
|
"waiting_queue": len(self.waiting_queue),
|
|
}
|
|
if self._metrics is not None:
|
|
stats.update(self._metrics.get_stats())
|
|
return stats
|
|
|
|
def remove_finished_tasks(self, stop_ids: List[int]) -> List[Task]:
|
|
with self._lock:
|
|
finished = []
|
|
for task in self.active_tasks:
|
|
if task.status == TaskStatus.ABORTED:
|
|
finished.append(task)
|
|
elif task.is_finished(stop_ids):
|
|
task.status = TaskStatus.FINISHED
|
|
finished.append(task)
|
|
self._total_tokens += task.output_tokens
|
|
|
|
if self._metrics is not None:
|
|
for task in finished:
|
|
self._metrics.mark_finished(
|
|
task.task_id, task.input_tokens, task.output_tokens
|
|
)
|
|
|
|
self.active_tasks = [
|
|
t
|
|
for t in self.active_tasks
|
|
if t.status not in (TaskStatus.FINISHED, TaskStatus.ABORTED)
|
|
]
|
|
return finished
|
|
|
|
def pull_candidates(self, n: int) -> List[Task]:
|
|
to_add: List[Task] = []
|
|
with self._lock:
|
|
take = min(n, len(self.waiting_queue))
|
|
for _ in range(take):
|
|
to_add.append(self.waiting_queue.popleft())
|
|
return to_add
|
|
|
|
def activate(self, task: Task):
|
|
task.status = TaskStatus.RUNNING
|
|
with self._lock:
|
|
self.active_tasks.append(task)
|
|
|
|
def return_to_waiting(self, tasks: List[Task]):
|
|
with self._lock:
|
|
for task in reversed(tasks):
|
|
self.waiting_queue.appendleft(task)
|
|
|
|
def has_work(self) -> bool:
|
|
return bool(self.active_tasks or self.waiting_queue)
|
|
|
|
def wait_for_tasks(self, timeout: float = 1.0):
|
|
with self._lock:
|
|
if self.waiting_queue or self.active_tasks:
|
|
return
|
|
self._task_event.clear()
|
|
self._task_event.wait(timeout=timeout)
|
|
|
|
def get_active_tasks(self) -> List[Task]:
|
|
with self._lock:
|
|
return list(self.active_tasks)
|
|
|
|
def get_waiting_tasks(self) -> List[Task]:
|
|
with self._lock:
|
|
return list(self.waiting_queue)
|
|
|
|
def clear_queues(self):
|
|
with self._lock:
|
|
self.waiting_queue.clear()
|
|
self.active_tasks.clear()
|
|
self._callbacks.clear()
|
|
|
|
def wake(self):
|
|
self._task_event.set()
|