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()