import logging import threading import time import uuid from collections import deque from enum import Enum from typing import Any, Callable, Deque, Dict, List, Optional from tokenizers.decoders import DecodeStream from astrai.inference.core.metrics import MetricsCollector from astrai.tokenize.tokenizer import AutoTokenizer logger = logging.getLogger(__name__) 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, ): 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.status = TaskStatus.PENDING self.output_ids: List[int] = [] self.output_logprobs: List[float] = [] self.input_tokens: int = 0 self.output_tokens: int = 0 self._decoder: Optional[StreamDecoder] = None 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]) def flush_remaining(self, tokenizer: AutoTokenizer) -> str: """Emit any text still buffered in the decoder. With the Rust-native DecodeStream, the stream is always in a correct state — any completed text was already emitted by the last ``push``. A trailing incomplete multi-byte sequence has no valid text to emit, so this is a no-op. """ return "" @property def next_pos(self) -> int: # The first output is sampled from prefill and enters KV on the next step. return self.input_tokens + max(0, len(self.output_ids) - 1) 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, 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 len(prompt_ids) > self.max_seq_len: if stream_callback: stream_callback(STOP) return task_id 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, ) 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()