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
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user