mirror of https://github.com/razor-ai/soup.git
111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
"""v0.44.0 Part A — SSE training-stream payload schema.
|
|
|
|
Pure-Python: serialises a TrainEvent as a single SSE-frame string that the
|
|
FastAPI endpoint can write into a streaming response. No FastAPI/dep import.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, Optional
|
|
|
|
# Closed allowlist — drop any payload key not in this set.
|
|
_ALLOWED_KEYS = frozenset(
|
|
{
|
|
"type",
|
|
"ts",
|
|
"step",
|
|
"epoch",
|
|
"loss",
|
|
"lr",
|
|
"grad_norm",
|
|
"tokens_per_s",
|
|
"ema_loss",
|
|
"p95_step_ms",
|
|
"p99_step_ms",
|
|
"eta_seconds",
|
|
"message",
|
|
}
|
|
)
|
|
|
|
_VALID_TYPES = frozenset({"metric", "status", "log", "eval"})
|
|
|
|
_MAX_MESSAGE_LEN = 1024
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TrainEvent:
|
|
"""One SSE-streamed training event."""
|
|
|
|
type: str
|
|
ts: float = field(default_factory=lambda: time.time())
|
|
step: Optional[int] = None
|
|
epoch: Optional[float] = None
|
|
loss: Optional[float] = None
|
|
lr: Optional[float] = None
|
|
grad_norm: Optional[float] = None
|
|
tokens_per_s: Optional[float] = None
|
|
ema_loss: Optional[float] = None
|
|
p95_step_ms: Optional[float] = None
|
|
p99_step_ms: Optional[float] = None
|
|
eta_seconds: Optional[float] = None
|
|
message: Optional[str] = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.type not in _VALID_TYPES:
|
|
raise ValueError(
|
|
f"type must be one of {sorted(_VALID_TYPES)}; got {self.type!r}"
|
|
)
|
|
if isinstance(self.ts, bool) or not isinstance(self.ts, (int, float)):
|
|
raise ValueError("ts must be a number")
|
|
if not math.isfinite(float(self.ts)):
|
|
raise ValueError("ts must be finite")
|
|
if self.message is not None:
|
|
if not isinstance(self.message, str):
|
|
raise ValueError("message must be str or None")
|
|
if "\x00" in self.message:
|
|
raise ValueError("message contains NUL byte")
|
|
if len(self.message) > _MAX_MESSAGE_LEN:
|
|
raise ValueError(
|
|
f"message exceeds {_MAX_MESSAGE_LEN} chars"
|
|
)
|
|
|
|
|
|
def to_payload(event: TrainEvent) -> Dict[str, Any]:
|
|
"""Convert a TrainEvent into a JSON-serialisable dict, omitting None."""
|
|
raw: Dict[str, Any] = {
|
|
"type": event.type,
|
|
"ts": float(event.ts),
|
|
}
|
|
for key in (
|
|
"step",
|
|
"epoch",
|
|
"loss",
|
|
"lr",
|
|
"grad_norm",
|
|
"tokens_per_s",
|
|
"ema_loss",
|
|
"p95_step_ms",
|
|
"p99_step_ms",
|
|
"eta_seconds",
|
|
"message",
|
|
):
|
|
value = getattr(event, key)
|
|
if value is not None:
|
|
raw[key] = value
|
|
# Defence-in-depth: filter out any keys that drifted in from refactors.
|
|
return {key: value for key, value in raw.items() if key in _ALLOWED_KEYS}
|
|
|
|
|
|
def format_sse_frame(event: TrainEvent) -> str:
|
|
"""Serialise the event as a single Server-Sent-Events frame.
|
|
|
|
`data: {json}\\n\\n` — the standard W3C SSE wire format.
|
|
"""
|
|
payload = to_payload(event)
|
|
body = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
|
|
return f"data: {body}\n\n"
|