Files
aituner/runs/telemetry-residual/common_state.py

401 lines
15 KiB
Python

#!/usr/bin/env python3
"""Extract directly measurable engine/simulator state into one schema.
The schema deliberately keeps common, engine-only, and simulator-only fields
separate. Missing simulator mechanisms must remain missing; they must not be
filled by a heuristic bottleneck label.
"""
from __future__ import annotations
import csv
import json
import math
from pathlib import Path
from statistics import fmean
from typing import Any, Iterable, Mapping, Sequence
SCHEMA = "telemetry-common-state-v1"
def numeric(values: Iterable[float | int]) -> dict[str, Any]:
finite = [float(value) for value in values]
if not finite:
raise ValueError("numeric summary requires at least one value")
if any(not math.isfinite(value) for value in finite):
raise ValueError("numeric summary received a non-finite value")
mean = fmean(finite)
variance = fmean((value - mean) ** 2 for value in finite)
return {
"n": len(finite),
"min": min(finite),
"max": max(finite),
"mean": mean,
"cv": math.sqrt(variance) / abs(mean) if mean else 0.0,
"distinct_n": len(set(finite)),
}
def load_jsonl(path: Path) -> list[dict[str, Any]]:
rows = []
with path.open(encoding="utf-8") as source:
for line_number, line in enumerate(source, start=1):
if not line.strip():
continue
row = json.loads(line)
if not isinstance(row, dict):
raise ValueError(f"{path}:{line_number}: expected a JSON object")
rows.append(row)
if not rows:
raise ValueError(f"{path}: no JSONL records")
return rows
def _as_number(value: Any, *, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be numeric, got {value!r}")
result = float(value)
if not math.isfinite(result):
raise ValueError(f"{name} must be finite, got {value!r}")
return result
def _time_weighted_mean(
records: Sequence[Mapping[str, Any]],
*,
start_ns: int,
end_ns: int,
value,
) -> float:
if end_ns <= start_ns:
raise ValueError("time-weighted interval must be positive")
selected = [
record
for record in records
if start_ns <= int(record["submit_mono_ns"]) <= end_ns
]
if not selected:
raise ValueError("time-weighted interval contains no records")
selected.sort(key=lambda record: int(record["submit_mono_ns"]))
cursor = start_ns
total = 0.0
current = _as_number(value(selected[0]), name="time-weighted value")
for record in selected[1:]:
timestamp = int(record["submit_mono_ns"])
if timestamp < cursor:
raise ValueError("telemetry timestamps are not monotonic")
total += current * (timestamp - cursor)
cursor = timestamp
current = _as_number(value(record), name="time-weighted value")
total += current * (end_ns - cursor)
return total / (end_ns - start_ns)
def summarize_engine(
records: Sequence[Mapping[str, Any]],
*,
start_ns: int,
end_ns: int,
request_count: int,
) -> dict[str, Any]:
"""Summarize a measured engine interval from Layer-1 records."""
if request_count <= 0:
raise ValueError("request_count must be positive")
layer1 = [record for record in records if "step_index" in record]
if not layer1:
raise ValueError("engine stream has no Layer-1 records")
step_indexes = [int(record["step_index"]) for record in layer1]
if len(step_indexes) != len(set(step_indexes)):
raise ValueError("engine Layer-1 step indexes are not unique")
if step_indexes != sorted(step_indexes):
raise ValueError("engine Layer-1 step indexes are not ordered")
if any(int(record.get("dropped_records_before", 0)) != 0 for record in layer1):
raise ValueError("engine Layer-1 stream reports dropped records")
interval = [
record
for record in layer1
if start_ns <= int(record["submit_mono_ns"]) <= end_ns
]
if not interval:
raise ValueError("engine interval has no Layer-1 records")
executed = [record for record in interval if bool(record["model_executed"])]
if not executed:
raise ValueError("engine interval has no executed model steps")
duration_s = (end_ns - start_ns) / 1e9
if duration_s <= 0:
raise ValueError("engine interval duration must be positive")
batch_sizes = [int(record["scheduled_requests"]) for record in executed]
prefill_tokens = [int(record["prefill_tokens"]) for record in executed]
decode_tokens = [int(record["decode_tokens"]) for record in executed]
batch_tokens = [
prefill + decode
for prefill, decode in zip(prefill_tokens, decode_tokens, strict=True)
]
decode_batches = [int(record["decode_batch_size"]) for record in executed]
if any(value < 0 for value in batch_sizes + batch_tokens + decode_batches):
raise ValueError("engine batch counters must be non-negative")
if any(
int(record["prefill_tokens"]) + int(record["decode_tokens"])
<= 0
for record in executed
):
raise ValueError("executed engine step has no scheduled tokens")
waiting_mean = _time_weighted_mean(
interval,
start_ns=start_ns,
end_ns=end_ns,
value=lambda record: record["queues"]["waiting"],
)
running_mean = _time_weighted_mean(
interval,
start_ns=start_ns,
end_ns=end_ns,
value=lambda record: record["queues"]["running"],
)
kv_mean = _time_weighted_mean(
interval,
start_ns=start_ns,
end_ns=end_ns,
value=lambda record: record["kv"]["usage"],
)
kv_values = [float(record["kv"]["usage"]) for record in interval]
if any(not 0.0 <= value <= 1.0 for value in kv_values):
raise ValueError("engine KV usage must be in [0, 1]")
total_prefill = sum(prefill_tokens)
total_decode = sum(decode_tokens)
graph_modes = [str(record["cudagraph"]["runtime_mode"]) for record in executed]
bucket_tokens = sum(int(record["cudagraph"]["bucket_tokens"]) for record in executed)
padding_tokens = sum(int(record["cudagraph"]["padding_tokens"]) for record in executed)
common = {
"scheduler_steps_per_s": len(executed) / duration_s,
"batch_size": numeric(batch_sizes),
"batch_tokens": numeric(batch_tokens),
"decode_batch_size": numeric(decode_batches),
"prefill_token_fraction": total_prefill / (total_prefill + total_decode),
"queue_waiting_mean": waiting_mean,
"queue_running_mean": running_mean,
"queue_waiting_time_per_request_ms": waiting_mean * duration_s * 1000.0 / request_count,
"queue_running_time_per_request_ms": running_mean * duration_s * 1000.0 / request_count,
"preemptions": sum(int(record["preemptions"]) for record in executed),
}
result = {
"schema": SCHEMA,
"source": "engine_layer1",
"interval": {
"start_ns": start_ns,
"end_ns": end_ns,
"duration_s": duration_s,
"request_count": request_count,
},
"common": common,
"engine_only": {
"kv_usage_mean": kv_mean,
"kv_usage_max": max(kv_values),
"kv_usage_end_minus_start": kv_values[-1] - kv_values[0],
"graph_none_share": graph_modes.count("NONE") / len(graph_modes),
"graph_full_share": graph_modes.count("FULL") / len(graph_modes),
"graph_padding_fraction": padding_tokens / max(1, bucket_tokens),
},
"simulator_only": {},
"sanity": {
"records": len(interval),
"executed_steps": len(executed),
"step_index_min": min(int(record["step_index"]) for record in interval),
"step_index_max": max(int(record["step_index"]) for record in interval),
"invariants": {
"positive_duration": duration_s > 0,
"positive_request_count": request_count > 0,
"zero_drops": True,
"nonnegative_counters": True,
"kv_bounded": True,
"batch_values_not_all_identical": any(
summary["distinct_n"] > 1
for summary in (
common["batch_size"],
common["batch_tokens"],
common["decode_batch_size"],
)
),
},
},
}
return result
def _csv_rows(path: Path) -> list[dict[str, str]]:
with path.open(encoding="utf-8", newline="") as source:
rows = list(csv.DictReader(source))
if not rows:
raise ValueError(f"{path}: CSV contains no rows")
return rows
def _column(rows: Sequence[Mapping[str, str]], name: str) -> list[float]:
if name not in rows[0]:
raise ValueError(f"CSV is missing required column {name!r}")
values = []
for row in rows:
text = row.get(name, "")
if text == "":
raise ValueError(f"CSV column {name!r} contains an empty value")
values.append(_as_number(float(text), name=name))
return values
def summarize_frontier(
*,
system_metrics_path: Path,
request_metrics_path: Path,
batch_metrics_path: Path | None = None,
ledger_path: Path | None = None,
) -> dict[str, Any]:
"""Summarize a Frontier run, retaining unavailable state as null."""
system = json.loads(system_metrics_path.read_text(encoding="utf-8"))
throughput = system["throughput_metrics"]
duration_s = _as_number(
throughput["total_duration_seconds"], name="total_duration_seconds"
)
if duration_s <= 0:
raise ValueError("Frontier duration must be positive")
request_rows = _csv_rows(request_metrics_path)
waiting_ms = _column(request_rows, "request_waiting_time_total")
e2e_ms = _column(request_rows, "request_e2e_time")
running_ms = [max(0.0, e2e - waiting) for e2e, waiting in zip(e2e_ms, waiting_ms, strict=True)]
duration_ms = duration_s * 1000.0
request_count = len(request_rows)
common: dict[str, Any] = {
"scheduler_steps_per_s": None,
"batch_size": None,
"batch_tokens": None,
"decode_batch_size": None,
"prefill_token_fraction": None,
"queue_waiting_mean": sum(waiting_ms) / duration_ms,
"queue_running_mean": sum(running_ms) / duration_ms,
"queue_waiting_time_per_request_ms": fmean(waiting_ms),
"queue_running_time_per_request_ms": fmean(running_ms),
"preemptions": sum(
int(float(row.get("request_total_preemption_count") or 0))
for row in request_rows
),
}
batch_rows: list[dict[str, str]] = []
if batch_metrics_path is not None:
batch_rows = _csv_rows(batch_metrics_path)
batch_sizes = _column(batch_rows, "batch_size")
batch_tokens = _column(batch_rows, "batch_num_tokens")
prefill_tokens = _column(batch_rows, "batch_num_prefill_tokens")
decode_tokens = _column(batch_rows, "batch_num_decode_tokens")
if any(value < 0 for value in batch_sizes + batch_tokens + prefill_tokens + decode_tokens):
raise ValueError("Frontier batch counters must be non-negative")
common.update(
{
"scheduler_steps_per_s": len(batch_rows) / duration_s,
"batch_size": numeric(batch_sizes),
"batch_tokens": numeric(batch_tokens),
"decode_batch_size": numeric(decode_tokens),
"prefill_token_fraction": sum(prefill_tokens)
/ max(1.0, sum(prefill_tokens) + sum(decode_tokens)),
}
)
ledger_rows: list[dict[str, Any]] = []
if ledger_path is not None:
ledger_rows = load_jsonl(ledger_path)
for row in ledger_rows:
start = _as_number(row["stage_start_ts"], name="stage_start_ts")
end = _as_number(row["stage_end_ts"], name="stage_end_ts")
if end < start:
raise ValueError("Frontier ledger has a negative stage duration")
batch_distinct = (
max(
summary["distinct_n"]
for summary in (
common["batch_size"],
common["batch_tokens"],
common["decode_batch_size"],
)
)
if batch_rows
else None
)
return {
"schema": SCHEMA,
"source": "frontier",
"interval": {
"duration_s": duration_s,
"request_count": request_count,
},
"common": common,
"engine_only": {
"kv_usage_mean": None,
"kv_usage_max": None,
"kv_usage_end_minus_start": None,
"graph_none_share": None,
"graph_full_share": None,
"graph_padding_fraction": None,
},
"simulator_only": {
"request_waiting_time_ms": numeric(waiting_ms),
"request_running_time_ms": numeric(running_ms),
"ledger_rows": len(ledger_rows) if ledger_path is not None else None,
},
"sanity": {
"request_rows": request_count,
"batch_rows": len(batch_rows),
"ledger_rows": len(ledger_rows),
"invariants": {
"positive_duration": duration_s > 0,
"positive_request_count": request_count > 0,
"nonnegative_counters": True,
"request_values_not_all_identical": max(
numeric(waiting_ms)["distinct_n"],
numeric(running_ms)["distinct_n"],
)
> 1,
"batch_values_not_all_identical": (
batch_distinct > 1 if batch_distinct is not None else None
),
},
},
}
def residual(real: Mapping[str, Any], simulated: Mapping[str, Any]) -> dict[str, Any]:
if real.get("schema") != SCHEMA or simulated.get("schema") != SCHEMA:
raise ValueError("residual inputs must use the common-state schema")
values = {}
missing = []
for name, real_value in real["common"].items():
sim_value = simulated["common"].get(name)
if isinstance(real_value, dict):
if not isinstance(sim_value, dict):
missing.append(name)
continue
for statistic in ("mean", "max", "cv"):
key = f"{name}.{statistic}"
values[key] = float(real_value[statistic]) - float(sim_value[statistic])
continue
if real_value is None or sim_value is None:
missing.append(name)
continue
values[name] = float(real_value) - float(sim_value)
return {
"schema": "telemetry-state-residual-v1",
"values": values,
"missing_common_fields": sorted(missing),
"coverage": {
"available": len(values),
"missing": len(missing),
"common_field_count": len(real["common"]),
},
}