Files
aituner/runs/action-aware-v0/analyze_pilot.py

588 lines
22 KiB
Python

#!/usr/bin/env python3
"""Audit source-only constraint signals against crossed real interventions."""
from __future__ import annotations
import argparse
import hashlib
import json
import math
import os
import statistics
import sys
from pathlib import Path
from typing import Any, Iterable, Mapping
HERE = Path(__file__).resolve().parent
COMMON_STATE = HERE.parent / "telemetry-residual"
sys.path.insert(0, str(COMMON_STATE))
from common_state import summarize_engine # noqa: E402
SCHEMA = "action-aware-constraint-pilot-audit-v0"
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for chunk in iter(lambda: source.read(1 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def atomic_json(path: Path, payload: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
temporary.write_text(
json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
)
os.replace(temporary, path)
def numeric(values: Iterable[float]) -> dict[str, Any]:
finite = [float(value) for value in values]
if not finite:
raise ValueError("numeric summary requires values")
if any(not math.isfinite(value) for value in finite):
raise ValueError("numeric summary received non-finite values")
return {
"n": len(finite),
"min": min(finite),
"max": max(finite),
"distinct_n": len(set(finite)),
}
def quantile(values: Iterable[float], probability: float) -> float:
ordered = sorted(float(value) for value in values)
if not ordered:
raise ValueError("quantile requires values")
position = probability * (len(ordered) - 1)
lower = math.floor(position)
upper = math.ceil(position)
if lower == upper:
return ordered[lower]
weight = position - lower
return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
def load_jsonl(path: Path) -> list[dict[str, Any]]:
records = []
with path.open(encoding="utf-8") as source:
for line_number, line in enumerate(source, 1):
try:
records.append(json.loads(line))
except json.JSONDecodeError as error:
raise ValueError(f"{path}:{line_number}: invalid JSON") from error
return records
def binding_summary(
records: list[Mapping[str, Any]], *, mns: int, mbbt: int
) -> dict[str, Any]:
if not records:
raise ValueError("binding summary requires scheduler records")
counts = {
"mns_exclusive": 0,
"mbbt_exclusive": 0,
"both": 0,
"waiting_unresolved": 0,
"waiting": 0,
}
running_utilization = []
token_utilization = []
kv_usage = []
preemptions = 0
for record in records:
waiting = int(record["queues"]["waiting"]) + int(
record["queues"]["deferred"]
)
running = int(record["queues"]["running"])
scheduled_tokens = int(record["prefill_tokens"]) + int(
record["decode_tokens"]
)
if running > mns:
raise ValueError("running requests exceed configured MNS")
if scheduled_tokens > mbbt:
raise ValueError("scheduled tokens exceed configured MBBT")
mns_hit = waiting > 0 and running == mns
mbbt_hit = waiting > 0 and scheduled_tokens == mbbt
if waiting > 0:
counts["waiting"] += 1
if mns_hit and mbbt_hit:
counts["both"] += 1
elif mns_hit:
counts["mns_exclusive"] += 1
elif mbbt_hit:
counts["mbbt_exclusive"] += 1
else:
counts["waiting_unresolved"] += 1
running_utilization.append(running / mns)
token_utilization.append(scheduled_tokens / mbbt)
kv_usage.append(float(record["kv"]["usage"]))
preemptions += int(record["preemptions"])
count = len(records)
return {
"records": count,
**{f"{name}_count": value for name, value in counts.items()},
**{f"{name}_fraction": value / count for name, value in counts.items()},
"running_utilization_mean": statistics.fmean(running_utilization),
"running_utilization_max": max(running_utilization),
"token_utilization_mean": statistics.fmean(token_utilization),
"token_utilization_max": max(token_utilization),
"kv_usage_mean": statistics.fmean(kv_usage),
"kv_usage_max": max(kv_usage),
"preemptions": preemptions,
}
def request_summary(path: Path, expected_count: int) -> dict[str, Any]:
rows = load_jsonl(path)
if len(rows) != expected_count:
raise ValueError(f"request row count mismatch: {path}")
ttft = [float(row["ttft_ms"]) for row in rows if row["ttft_ms"] is not None]
tpot = [float(row["tpot_ms"]) for row in rows if row["tpot_ms"] is not None]
if not ttft or not tpot:
raise ValueError(f"missing request latency values: {path}")
return {
"ttft_ms": {f"p{int(p * 100)}": quantile(ttft, p) for p in (0.5, 0.95, 0.99)},
"tpot_ms": {f"p{int(p * 100)}": quantile(tpot, p) for p in (0.5, 0.95, 0.99)},
}
def load_stream(session_root: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
streams = sorted((session_root / "opprof").glob("*.jsonl"))
sidecars = sorted((session_root / "opprof").glob("*.jsonl.footer.json"))
if len(streams) != 1 or len(sidecars) != 1:
raise ValueError(f"expected one OpProf stream and sidecar: {session_root}")
decoded = load_jsonl(streams[0])
records = [row for row in decoded if "step_index" in row]
footers = [row for row in decoded if row.get("record_type") == "footer"]
sidecar = json.loads(sidecars[0].read_text(encoding="utf-8"))
indexes = [int(row["step_index"]) for row in records]
invariants = {
"one_footer_last": len(footers) == 1 and decoded[-1] is footers[0],
"sidecar_final": sidecar.get("final") is True,
"zero_drops": sidecar.get("dropped_records") == 0,
"written_matches_records": sidecar.get("written_records") == len(records),
"contiguous_step_indexes": indexes == list(range(len(indexes))),
"monotonic_timestamps": all(
int(right["submit_mono_ns"]) >= int(left["submit_mono_ns"])
for left, right in zip(records, records[1:], strict=False)
),
}
return records, {
"stream": str(streams[0]),
"stream_sha256": sha256_file(streams[0]),
"records": len(records),
"invariants": invariants,
}
def analyze_run(
*,
run_root: Path,
config: Mapping[str, Any],
repetition: int,
expected: Mapping[str, Any],
stream_records: list[Mapping[str, Any]],
duration_s: float,
phase_fractions: list[float],
) -> dict[str, Any]:
result_root = run_root / "sessions" / str(config["id"]) / f"rep{repetition}"
result_path = result_root / "result.json"
result = json.loads(result_path.read_text(encoding="utf-8"))
selection = result["selection"]
invariants = {
"result_schema": result.get("schema") == "action-aware-pilot-result-v0",
"config_id": result.get("config_id") == config["id"],
"tp": int(result.get("tp", -1)) == 4,
"mns": int(result.get("mns", -1)) == int(config["mns"]),
"mbbt": int(result.get("mbbt", -1)) == int(config["mbbt"]),
"uncensored": not bool(result.get("early_stopped", True)),
"slo_early_stop_disabled": result.get("slo_early_stop_disabled") is True,
"selection_count": int(selection["count"]) == int(expected["selected_count"]),
"request_accounting": int(result["observed_count"])
== int(expected["selected_count"]),
"request_hash": selection["request_id_order_sha256"]
== expected["request_id_order_sha256"],
"arrival_hash": selection["arrival_order_sha256"]
== expected["arrival_order_sha256"],
"length_hash": selection["raw_length_order_sha256"]
== expected["input_length_order_sha256"],
}
start_ns = int(result["interval"]["start_mono_ns"])
arrival_end_ns = start_ns + round(duration_s * 1e9)
full_records = [
record
for record in stream_records
if start_ns <= int(record["submit_mono_ns"]) <= arrival_end_ns
]
if not full_records:
raise ValueError(f"no telemetry records in measured window: {result_path}")
gaps = [
(int(right["submit_mono_ns"]) - int(left["submit_mono_ns"])) / 1e9
for left, right in zip(full_records, full_records[1:], strict=False)
]
coverage = {
"start_gap_s": (int(full_records[0]["submit_mono_ns"]) - start_ns) / 1e9,
"end_gap_s": (arrival_end_ns - int(full_records[-1]["submit_mono_ns"])) / 1e9,
"max_internal_gap_s": max(gaps, default=0.0),
}
invariants["telemetry_coverage"] = all(
0.0 <= value <= 1.0 for value in coverage.values()
)
binding = binding_summary(
full_records, mns=int(config["mns"]), mbbt=int(config["mbbt"])
)
phases = {}
for fraction in phase_fractions:
phase_end = start_ns + round(duration_s * fraction * 1e9)
phase_records = [
record
for record in full_records
if int(record["submit_mono_ns"]) <= phase_end
]
phases[f"{fraction:.2f}"] = binding_summary(
phase_records, mns=int(config["mns"]), mbbt=int(config["mbbt"])
)
state = summarize_engine(
full_records,
start_ns=start_ns,
end_ns=arrival_end_ns,
request_count=int(result["observed_count"]),
)
latency = request_summary(
result_root / "requests.jsonl", int(result["observed_count"])
)
return {
"config_id": config["id"],
"mns": int(config["mns"]),
"mbbt": int(config["mbbt"]),
"repetition": repetition,
"result_path": str(result_path),
"result_sha256": sha256_file(result_path),
"selection": {
"count": int(selection["count"]),
"request_id_order_sha256": selection["request_id_order_sha256"],
"arrival_order_sha256": selection["arrival_order_sha256"],
"raw_length_order_sha256": selection["raw_length_order_sha256"],
},
"outcome": {
"pass_rate": float(result["pass_rate"]),
"feasible": bool(result["feasible"]),
"slo_pass_count": int(result["slo_pass_count"]),
"slo_goodput_req_s": int(result["slo_pass_count"]) / duration_s,
"elapsed_s": float(result["interval"]["elapsed_s"]),
**latency,
},
"binding": binding,
"phases": phases,
"state": state,
"coverage": coverage,
"invariants": invariants,
}
def median(values: Iterable[float]) -> float:
return float(statistics.median(float(value) for value in values))
def evaluate_decisions(
runs: list[Mapping[str, Any]], manifest: Mapping[str, Any]
) -> dict[str, Any]:
by_key = {
(str(run["config_id"]), int(run["repetition"])): run for run in runs
}
repetitions = sorted(int(key) for key in manifest["repetitions"])
regime_results = {}
all_predictions = []
crossed_pass = True
binding_pass = True
material_ambiguity = False
for regime_name, regime in manifest["regimes"].items():
rows = []
source_runs = []
for repetition in repetitions:
source = by_key[(str(regime["source"]), repetition)]
mns_target = by_key[(str(regime["actions"]["mns"]), repetition)]
mbbt_target = by_key[(str(regime["actions"]["mbbt"]), repetition)]
source_runs.append(source)
source_goodput = float(source["outcome"]["slo_goodput_req_s"])
mns_goodput = float(mns_target["outcome"]["slo_goodput_req_s"])
mbbt_goodput = float(mbbt_target["outcome"]["slo_goodput_req_s"])
observed = (
"mns"
if mns_goodput > mbbt_goodput
else "mbbt"
if mbbt_goodput > mns_goodput
else "tie"
)
mns_score = float(source["binding"]["mns_exclusive_fraction"])
mbbt_score = float(source["binding"]["mbbt_exclusive_fraction"])
predicted = (
"mns"
if mns_score > mbbt_score
else "mbbt"
if mbbt_score > mns_score
else "tie"
)
phase_predictions = {}
for phase, summary in source["phases"].items():
left = float(summary["mns_exclusive_fraction"])
right = float(summary["mbbt_exclusive_fraction"])
phase_predictions[phase] = (
"mns" if left > right else "mbbt" if right > left else "tie"
)
margin = (
abs(mns_goodput - mbbt_goodput) / source_goodput
if source_goodput > 0
else None
)
row = {
"repetition": repetition,
"source_goodput_req_s": source_goodput,
"mns_target_goodput_req_s": mns_goodput,
"mbbt_target_goodput_req_s": mbbt_goodput,
"observed_winner": observed,
"predicted_winner": predicted,
"prediction_correct": predicted == observed,
"relative_winner_margin_over_source": margin,
"mns_exclusive_fraction": mns_score,
"mbbt_exclusive_fraction": mbbt_score,
"phase_predictions": phase_predictions,
"phase_stable": all(value == predicted for value in phase_predictions.values()),
}
rows.append(row)
all_predictions.append(row)
expected_winner = "mns" if regime_name == "A" else "mbbt"
minimum_margin = float(manifest["gates"]["minimum_relative_winner_margin"])
regime_crossed = all(
row["observed_winner"] == expected_winner
and row["relative_winner_margin_over_source"] is not None
and row["relative_winner_margin_over_source"] >= minimum_margin
for row in rows
)
crossed_pass &= regime_crossed
winning_key = f"{expected_winner}_exclusive_fraction"
losing_key = (
"mbbt_exclusive_fraction" if expected_winner == "mns" else "mns_exclusive_fraction"
)
winning_median = median(row[winning_key] for row in rows)
losing_median = median(row[losing_key] for row in rows)
ratio_pass = winning_median >= float(
manifest["gates"]["minimum_exclusive_ratio"]
) * losing_median
regime_binding = (
all(row["prediction_correct"] and row["phase_stable"] for row in rows)
and winning_median
>= float(manifest["gates"]["minimum_exclusive_fraction"])
and ratio_pass
)
binding_pass &= regime_binding
ambiguity_median = median(
float(run["binding"]["both_fraction"])
+ float(run["binding"]["waiting_unresolved_fraction"])
for run in source_runs
)
score_gap_median = median(
abs(
float(run["binding"]["mns_exclusive_fraction"])
- float(run["binding"]["mbbt_exclusive_fraction"])
)
for run in source_runs
)
kv_max_median = median(
float(run["binding"]["kv_usage_max"]) for run in source_runs
)
any_preemption = any(
int(run["binding"]["preemptions"]) > 0 for run in source_runs
)
regime_material = (
ambiguity_median >= score_gap_median
or kv_max_median >= float(manifest["gates"]["material_kv_usage"])
or any_preemption
)
material_ambiguity |= regime_material
regime_results[regime_name] = {
"source": regime["source"],
"actions": regime["actions"],
"expected_winner": expected_winner,
"crossed_response_pass": regime_crossed,
"binding_pass": regime_binding,
"winning_exclusive_median": winning_median,
"losing_exclusive_median": losing_median,
"exclusive_ratio_pass": ratio_pass,
"ambiguity_median": ambiguity_median,
"exclusive_gap_median": score_gap_median,
"kv_usage_max_median": kv_max_median,
"any_preemption": any_preemption,
"material_ambiguity": regime_material,
"repetitions": rows,
}
if not crossed_pass:
decision = "STOP_WORKLOAD_NOT_CROSSED"
elif not binding_pass:
decision = "STOP_BINDING_NOT_PREDICTIVE"
elif material_ambiguity:
decision = "OPEN_EXACT_ATTRIBUTION_ABLATION"
else:
decision = "STOP_NO_NEW_INSTRUMENTATION_NEEDED"
correct = sum(int(row["prediction_correct"]) for row in all_predictions)
return {
"decision": decision,
"crossed_response_pass": crossed_pass,
"binding_pass": binding_pass,
"material_ambiguity": material_ambiguity,
"regimes": regime_results,
"baselines": {
"always_mns_correct": sum(
int(row["observed_winner"] == "mns") for row in all_predictions
),
"always_mbbt_correct": sum(
int(row["observed_winner"] == "mbbt") for row in all_predictions
),
"binding_correct": correct,
"decision_count": len(all_predictions),
},
}
def analyze(run_root: Path, manifest_path: Path) -> dict[str, Any]:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
if manifest.get("schema") not in {
"action-aware-constraint-pilot-manifest-v0",
"action-aware-constraint-pilot-manifest-v1",
}:
raise ValueError("unexpected manifest schema")
duration_s = float(manifest["engine"]["duration_s"])
phase_fractions = [float(value) for value in manifest["gates"]["phase_fractions"]]
runs = []
stream_audits = []
for config in manifest["configs"]:
session_root = run_root / "sessions" / str(config["id"])
stream_records, stream_audit = load_stream(session_root)
stream_audit["config_id"] = config["id"]
stream_audits.append(stream_audit)
for repetition in sorted(int(key) for key in manifest["repetitions"]):
runs.append(
analyze_run(
run_root=run_root,
config=config,
repetition=repetition,
expected=manifest["repetitions"][str(repetition)]["selection"],
stream_records=stream_records,
duration_s=duration_s,
phase_fractions=phase_fractions,
)
)
invariants = {
"fifteen_runs": len(runs) == 15,
"five_streams": len(stream_audits) == 5,
"all_run_invariants": all(
all(bool(value) for value in run["invariants"].values()) for run in runs
),
"all_stream_invariants": all(
all(bool(value) for value in stream["invariants"].values())
for stream in stream_audits
),
"nonnegative_counters": all(
all(
float(run["binding"][key]) >= 0
for key in (
"mns_exclusive_count",
"mbbt_exclusive_count",
"both_count",
"waiting_unresolved_count",
"preemptions",
)
)
for run in runs
),
"ratios_bounded": all(
all(
0.0 <= float(run["binding"][key]) <= 1.0
for key in (
"mns_exclusive_fraction",
"mbbt_exclusive_fraction",
"both_fraction",
"waiting_unresolved_fraction",
"kv_usage_mean",
"kv_usage_max",
)
)
for run in runs
),
"per_config_results_not_all_identical": len(
{float(run["outcome"]["pass_rate"]) for run in runs}
)
> 1,
}
red_flags = [name for name, passed in invariants.items() if not passed]
decisions = (
evaluate_decisions(runs, manifest)
if not red_flags
else {
"decision": "STOP_DATA_INVALID",
"crossed_response_pass": False,
"binding_pass": False,
"material_ambiguity": False,
"regimes": {},
"baselines": {},
}
)
payload = {
"schema": SCHEMA,
"decision": decisions["decision"],
"manifest": str(manifest_path),
"manifest_sha256": sha256_file(manifest_path),
"run_root": str(run_root),
"runs": runs,
"streams": stream_audits,
"decision_audit": decisions,
"sanity": {
"runs": len(runs),
"pass_rate": numeric(run["outcome"]["pass_rate"] for run in runs),
"slo_goodput_req_s": numeric(
run["outcome"]["slo_goodput_req_s"] for run in runs
),
"telemetry_records_per_run": numeric(
run["binding"]["records"] for run in runs
),
"mns_values": numeric(run["mns"] for run in runs),
"mbbt_values": numeric(run["mbbt"] for run in runs),
"invariants": invariants,
"red_flags": red_flags,
},
}
return payload
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--run-root", type=Path, required=True)
parser.add_argument("--manifest", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
payload = analyze(args.run_root, args.manifest)
atomic_json(args.output, payload)
print(
json.dumps(
{
"decision": payload["decision"],
"sanity": payload["sanity"],
"decision_audit": payload["decision_audit"],
},
indent=2,
sort_keys=True,
)
)
if __name__ == "__main__":
main()