From cb60713a72869409cfa9fa64e0a67d648840e558 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 8 Aug 2026 12:07:08 +0800 Subject: [PATCH] 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 | 3 + astrai/inference/core/metrics.py | 225 +++++++++++++++++++++++++++++ astrai/inference/core/scheduler.py | 24 ++- astrai/inference/core/task.py | 22 ++- 4 files changed, 263 insertions(+), 11 deletions(-) create mode 100644 astrai/inference/core/metrics.py diff --git a/astrai/inference/core/__init__.py b/astrai/inference/core/__init__.py index 38137d6..673fefd 100644 --- a/astrai/inference/core/__init__.py +++ b/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", diff --git a/astrai/inference/core/metrics.py b/astrai/inference/core/metrics.py new file mode 100644 index 0000000..9996367 --- /dev/null +++ b/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 diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index 9236da2..ac2d204 100644 --- a/astrai/inference/core/scheduler.py +++ b/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] = [] diff --git a/astrai/inference/core/task.py b/astrai/inference/core/task.py index c3ba3a3..57a4e07 100644 --- a/astrai/inference/core/task.py +++ b/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