feat: add per-task throughput and latency metrics

- extract TaskTiming + MetricsCollector out of Task/TaskManager
- unify prefill/decode timing into single record() context manager
- expose avg_ttft_ms, avg_decode_tps, avg_e2e_latency_ms via /stats

astrai/inference/core/__init__.py

@@ -10,6 +10,7 @@ from astrai.inference.core.cache import (
     page_hash,
 )
 from astrai.inference.core.executor import Executor
+from astrai.inference.core.metrics import MetricsCollector, TaskTiming
 from astrai.inference.core.scheduler import InferenceScheduler
 from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 
@@ -23,6 +24,8 @@ __all__ = [
     "page_hash",
     "Executor",
     "InferenceScheduler",
+    "MetricsCollector",
+    "TaskTiming",
     "STOP",
     "Task",
     "TaskManager",

astrai/inference/core/metrics.py

@@ -0,0 +1,225 @@
+"""Unified per-task perf/stats: timing records, context-manager scopes, aggregate reporting."""
+
+import time
+from collections import deque
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Deque, Dict, Generator, List, Literal, Optional
+
+
+@dataclass
+class TaskTiming:
+    """Timestamp snapshots and computed metrics for one generation task.
+
+    Created by :class:`MetricsCollector` at task-registration time;
+    updated via ``prefill_scope`` / ``mark_finished``.
+    """
+
+    task_id: str
+    arrival_time: float
+    prefill_start_time: Optional[float] = None
+    first_token_time: Optional[float] = None
+    finish_time: Optional[float] = None
+    input_tokens: int = 0
+    output_tokens: int = 0
+    _decode_steps: int = 0
+    _decode_total_s: float = 0.0
+
+    # derived metrics
+
+    # derived metrics
+
+    @property
+    def queue_wait_ms(self) -> Optional[float]:
+        if self.prefill_start_time is not None:
+            return (self.prefill_start_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def ttft_ms(self) -> Optional[float]:
+        if self.first_token_time is not None:
+            return (self.first_token_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def prefill_tps(self) -> Optional[float]:
+        if self.prefill_start_time is not None and self.first_token_time is not None:
+            d = self.first_token_time - self.prefill_start_time
+            if d > 0 and self.input_tokens > 0:
+                return self.input_tokens / d
+        return None
+
+    @property
+    def decode_tps(self) -> Optional[float]:
+        if self.first_token_time is not None and self.finish_time is not None:
+            d = self.finish_time - self.first_token_time
+            dt = self.output_tokens - 1
+            if dt > 0 and d > 0:
+                return dt / d
+        return None
+
+    @property
+    def decode_avg_ms(self) -> Optional[float]:
+        if self._decode_steps > 0 and self._decode_total_s > 0:
+            return (self._decode_total_s / self._decode_steps) * 1000
+        return None
+
+    @property
+    def e2e_latency_ms(self) -> Optional[float]:
+        if self.finish_time is not None:
+            return (self.finish_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def total_tps(self) -> Optional[float]:
+        if self.finish_time is not None:
+            total = self.input_tokens + self.output_tokens
+            d = self.finish_time - self.arrival_time
+            if total > 0 and d > 0:
+                return total / d
+        return None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "task_id": self.task_id,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "queue_wait_ms": (
+                round(self.queue_wait_ms, 2) if self.queue_wait_ms is not None else None
+            ),
+            "ttft_ms": (round(self.ttft_ms, 2) if self.ttft_ms is not None else None),
+            "prefill_tps": (
+                round(self.prefill_tps, 2) if self.prefill_tps is not None else None
+            ),
+            "decode_tps": (
+                round(self.decode_tps, 2) if self.decode_tps is not None else None
+            ),
+            "decode_avg_ms": (
+                round(self.decode_avg_ms, 2) if self.decode_avg_ms is not None else None
+            ),
+            "total_tps": (
+                round(self.total_tps, 2) if self.total_tps is not None else None
+            ),
+            "e2e_latency_ms": (
+                round(self.e2e_latency_ms, 2)
+                if self.e2e_latency_ms is not None
+                else None
+            ),
+        }
+
+
+class MetricsCollector:
+    """Single-owner perf/stats hub for all generation tasks.
+
+    Usage::
+
+        metrics = MetricsCollector()
+        metrics.register(task_id, arrival_time)
+
+        with metrics.prefill_scope(task_ids):
+            run_prefill(...)
+
+        metrics.mark_finished(task_id, input_tokens, output_tokens)
+
+        stats = metrics.get_stats()
+    """
+
+    def __init__(self, max_recent: int = 128):
+        self._timings: Dict[str, TaskTiming] = {}
+        self._completed: Deque[TaskTiming] = deque(maxlen=max_recent)
+
+        self._ttft_ms_sum = 0.0
+        self._ttft_ms_count = 0
+        self._decode_tps_sum = 0.0
+        self._decode_tps_count = 0
+        self._e2e_ms_sum = 0.0
+        self._e2e_ms_count = 0
+
+    def register(self, task_id: str):
+        """Create a timing record for a newly-created task."""
+        self._timings[task_id] = TaskTiming(task_id=task_id, arrival_time=time.time())
+
+    def mark_finished(self, task_id: str, input_tokens: int, output_tokens: int):
+        """Close timing for a finished/aborted task and move it to completed."""
+        timing = self._timings.pop(task_id, None)
+        if timing is None:
+            return
+        timing.finish_time = time.time()
+        timing.input_tokens = input_tokens
+        timing.output_tokens = output_tokens
+        self._completed.append(timing)
+        self._accumulate(timing)
+
+    def clear(self):
+        """Reset all state (e.g. on engine shutdown)."""
+        self._timings.clear()
+        self._completed.clear()
+        self._ttft_ms_sum = 0.0
+        self._ttft_ms_count = 0
+        self._decode_tps_sum = 0.0
+        self._decode_tps_count = 0
+        self._e2e_ms_sum = 0.0
+        self._e2e_ms_count = 0
+
+    # timing scopes
+
+    @contextmanager
+    def record(
+        self, task_ids: List[str], phase: Literal["prefill", "decode"]
+    ) -> Generator[None, None, None]:
+        tic = time.time()
+        yield
+        toc = time.time()
+        dt = toc - tic
+        for tid in task_ids:
+            t = self._timings.get(tid)
+            if t is None:
+                continue
+            if phase == "prefill":
+                t.prefill_start_time = tic
+                t.first_token_time = toc
+            elif phase == "decode":
+                t._decode_steps += 1
+                t._decode_total_s += dt
+
+    # access
+
+    def get_timing(self, task_id: str) -> Optional[TaskTiming]:
+        """Return the timing record for *task_id* (active or completed)."""
+        if task_id in self._timings:
+            return self._timings[task_id]
+        for t in self._completed:
+            if t.task_id == task_id:
+                return t
+        return None
+
+    # aggregate stats
+
+    def get_stats(self) -> Dict[str, Any]:
+        stats: Dict[str, Any] = {}
+        if self._ttft_ms_count > 0:
+            stats["avg_ttft_ms"] = round(self._ttft_ms_sum / self._ttft_ms_count, 2)
+        if self._decode_tps_count > 0:
+            stats["avg_decode_tps"] = round(
+                self._decode_tps_sum / self._decode_tps_count, 2
+            )
+        if self._e2e_ms_count > 0:
+            stats["avg_e2e_latency_ms"] = round(
+                self._e2e_ms_sum / self._e2e_ms_count, 2
+            )
+        if self._completed:
+            stats["recent_tasks"] = [t.to_dict() for t in self._completed]
+        return stats
+
+    # internal
+
+    def _accumulate(self, t: TaskTiming):
+        if t.ttft_ms is not None:
+            self._ttft_ms_sum += t.ttft_ms
+            self._ttft_ms_count += 1
+        if t.decode_tps is not None:
+            self._decode_tps_sum += t.decode_tps
+            self._decode_tps_count += 1
+        if t.e2e_latency_ms is not None:
+            self._e2e_ms_sum += t.e2e_latency_ms
+            self._e2e_ms_count += 1

astrai/inference/core/scheduler.py

@@ -7,6 +7,7 @@ import torch
 
 from astrai.inference.core.cache import PagePool
 from astrai.inference.core.executor import Executor
+from astrai.inference.core.metrics import MetricsCollector
 from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
@@ -56,10 +57,13 @@ class InferenceScheduler:
                 dtype=self.dtype,
             )
 
+        self._metrics = MetricsCollector()
+
         self._task_mgr = TaskManager(
             tokenizer=tokenizer,
             max_batch_size=max_batch_size,
             max_seq_len=self.max_seq_len,
+            metrics=self._metrics,
         )
 
         self._executor = Executor(
@@ -121,14 +125,17 @@ class InferenceScheduler:
                 groups.setdefault((len(t.prompt_ids), start_pos), []).append(t)
 
             for (prompt_len, start_pos), group in groups.items():
-                prefilled, step_out = self._executor.execute_prefill(
-                    group, prompt_len, start_pos, return_logprobs=return_logprobs
-                )
+                with self._metrics.record([t.task_id for t in group], "prefill"):
+                    prefilled, step_out = self._executor.execute_prefill(
+                        group, prompt_len, start_pos, return_logprobs=return_logprobs
+                    )
+
                 for t, out in zip(prefilled, step_out):
                     t.output_ids.append(out[0] if return_logprobs else out)
                     t.output_tokens += 1
                     prefilled_ids.add(t.task_id)
                     produced.append(t)
+
                 start_logical_page = start_pos // getattr(cache, "page_size", 64)
                 for t in group:
                     cache.task_record_hashes(
@@ -147,9 +154,10 @@ class InferenceScheduler:
                 aborted.append(t)
 
         if decoded:
-            step_out = self._executor.execute_decode(
-                decoded, return_logprobs=return_logprobs
-            )
+            with self._metrics.record([t.task_id for t in decoded], "decode"):
+                step_out = self._executor.execute_decode(
+                    decoded, return_logprobs=return_logprobs
+                )
             for t, out in zip(decoded, step_out):
                 t.output_ids.append(out[0] if return_logprobs else out)
                 t.output_tokens += 1
@@ -305,6 +313,7 @@ class InferenceScheduler:
                 tasks.append(None)
                 continue
             task.input_tokens = len(task.prompt_ids)
+            self._metrics.register(task.task_id)
             tasks.append(task)
 
         try:
@@ -316,6 +325,9 @@ class InferenceScheduler:
         finally:
             for t in tasks:
                 if t is not None:
+                    self._metrics.mark_finished(
+                        t.task_id, t.input_tokens, t.output_tokens
+                    )
                     cache.task_free(t.task_id)
 
         results: List[Any] = []

astrai/inference/core/task.py

@@ -8,6 +8,7 @@ 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__)
@@ -79,8 +80,6 @@ class Task:
         self.output_logprobs: List[float] = []
         self.input_tokens: int = 0
         self.output_tokens: int = 0
-        self.arrival_time = time.time()
-        self.finish_time: Optional[float] = None
         self._decoder: Optional[StreamDecoder] = None
 
     def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
@@ -124,6 +123,7 @@ class TaskManager:
         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
@@ -139,6 +139,8 @@ class TaskManager:
         self._total_tasks = 0
         self._total_tokens = 0
 
+        self._metrics = metrics
+
     def add_task(
         self,
         prompt: str,
@@ -182,6 +184,9 @@ class TaskManager:
             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
 
@@ -201,26 +206,33 @@ class TaskManager:
             cb(token)
 
     def get_stats(self) -> Dict[str, Any]:
-        return {
+        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:
-                    task.finish_time = time.time()
                     finished.append(task)
                 elif task.is_finished(stop_ids):
                     task.status = TaskStatus.FINISHED
-                    task.finish_time = time.time()
                     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