ViperEkura
ae7fc3059a
- split KVCache into phase-specific PrefillKVCache/DecodeKVCache types selected by start_pos - unify steady-state detection in TaskCacheManager - guard decode steady-state reuse with the cached task signature so recycled req slots cannot replay a prior generation's tokens and positions - collapse attention backend fwd_decode/fwd_prefill into a single subclass-owned forward with a shared _check_fwd guard - fix thread-safety gap in weight update and validate prefill inputs before KV allocation - centralize magic constants in InferenceConfig and align docs with behavior
214 lines
7.1 KiB
Python
214 lines
7.1 KiB
Python
"""Unified per-task perf/stats: timing records, context-manager scopes, aggregate reporting."""
|
|
|
|
import threading
|
|
import time
|
|
from collections import deque
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Any, Deque, Dict, Generator, List, Literal, Optional
|
|
|
|
from astrai.config.inference_config import InferenceConfig
|
|
|
|
_config = InferenceConfig()
|
|
|
|
|
|
@dataclass
|
|
class TaskTiming:
|
|
"""Timestamp snapshots and computed metrics for one generation task.
|
|
|
|
Created by :class:`MetricsCollector` at task-registration time;
|
|
updated via ``record`` / ``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
|
|
|
|
@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.record(task_ids, "prefill"):
|
|
run_prefill(...)
|
|
|
|
metrics.mark_finished(task_id, input_tokens, output_tokens)
|
|
|
|
stats = metrics.get_stats()
|
|
"""
|
|
|
|
def __init__(self, max_recent: int = _config.max_recent_tasks):
|
|
self._timings: Dict[str, TaskTiming] = {}
|
|
self._completed: Deque[TaskTiming] = deque(maxlen=max_recent)
|
|
self._lock = threading.Lock()
|
|
|
|
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."""
|
|
with self._lock:
|
|
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."""
|
|
with self._lock:
|
|
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)
|
|
|
|
# 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
|
|
with self._lock:
|
|
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
|
|
|
|
# aggregate stats
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
with self._lock:
|
|
stats: Dict[str, Any] = {"in_flight_tasks": len(self._timings)}
|
|
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
|