325 lines
14 KiB
Python
325 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Extract paired source/action examples from the accepted action-aware run."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from statistics import fmean
|
|
from typing import Any, Mapping
|
|
|
|
|
|
PHASES = ("0.25", "0.50", "0.75", "1.00")
|
|
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
|
|
|
|
|
|
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 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):
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
records.append(json.loads(line))
|
|
except json.JSONDecodeError as error:
|
|
raise ValueError(f"{path}:{line_number}: invalid JSON") from error
|
|
if not records:
|
|
raise ValueError(f"{path}: no request records")
|
|
return records
|
|
|
|
|
|
def prefix_outcome(
|
|
requests: list[Mapping[str, Any]], *, cutoff_s: float, offered_total: float
|
|
) -> dict[str, float]:
|
|
admitted = [request for request in requests if float(request["arrival_s"]) <= cutoff_s]
|
|
completed = [
|
|
request
|
|
for request in requests
|
|
if request.get("completed_elapsed_s") is not None
|
|
and float(request["completed_elapsed_s"]) <= cutoff_s
|
|
]
|
|
if not admitted:
|
|
raise ValueError("prefix has no admitted requests")
|
|
admitted_ids = {str(request["request_id"]) for request in admitted}
|
|
if any(str(request["request_id"]) not in admitted_ids for request in completed):
|
|
raise ValueError("prefix completion precedes admission")
|
|
passed = sum(bool(request["slo_pass"]) for request in completed)
|
|
ttft = [float(request["ttft_ms"]) for request in completed]
|
|
tpot = [float(request["tpot_ms"]) for request in completed]
|
|
total = len(requests)
|
|
return {
|
|
"normalized_slo_goodput": passed / cutoff_s / offered_total,
|
|
"admitted_fraction": len(admitted) / total,
|
|
"completed_over_admitted": len(completed) / len(admitted),
|
|
"completed_pass_rate": passed / max(1, len(completed)),
|
|
"completed_fail_fraction_of_total": (len(completed) - passed) / total,
|
|
"outstanding_over_admitted": (len(admitted) - len(completed)) / len(admitted),
|
|
"ttft_max_over_slo_max": max(ttft, default=0.0) / 6000.0,
|
|
"ttft_mean_over_slo_max": fmean(ttft) / 6000.0 if ttft else 0.0,
|
|
"tpot_max_over_slo": max(tpot, default=0.0) / 50.0,
|
|
"tpot_mean_over_slo": fmean(tpot) / 50.0 if tpot else 0.0,
|
|
"admitted_input_tokens_mean_over_limit": fmean(
|
|
float(request["raw_input_tokens"]) for request in admitted
|
|
)
|
|
/ 8192.0,
|
|
}
|
|
|
|
|
|
def telemetry_record(state: Mapping[str, Any]) -> dict[str, float]:
|
|
common = state["common"]
|
|
engine = state["engine_only"]
|
|
executed_steps = int(state["sanity"]["executed_steps"])
|
|
if executed_steps <= 0:
|
|
raise ValueError("telemetry phase contains no executed engine steps")
|
|
return {
|
|
"scheduler_steps_per_s": float(common["scheduler_steps_per_s"]),
|
|
"batch_size_mean": float(common["batch_size"]["mean"]),
|
|
"batch_size_cv": float(common["batch_size"]["cv"]),
|
|
"batch_tokens_mean": float(common["batch_tokens"]["mean"]),
|
|
"batch_tokens_cv": float(common["batch_tokens"]["cv"]),
|
|
"decode_batch_size_mean": float(common["decode_batch_size"]["mean"]),
|
|
"decode_batch_size_cv": float(common["decode_batch_size"]["cv"]),
|
|
"prefill_token_fraction": float(common["prefill_token_fraction"]),
|
|
"queue_waiting_mean": float(common["queue_waiting_mean"]),
|
|
"queue_running_mean": float(common["queue_running_mean"]),
|
|
"preemptions_per_step": float(common["preemptions"]) / executed_steps,
|
|
"kv_usage_mean": float(engine["kv_usage_mean"]),
|
|
"kv_usage_max": float(engine["kv_usage_max"]),
|
|
"kv_usage_end_minus_start": float(engine["kv_usage_end_minus_start"]),
|
|
"graph_none_share": float(engine["graph_none_share"]),
|
|
"graph_full_share": float(engine["graph_full_share"]),
|
|
"graph_padding_fraction": float(engine["graph_padding_fraction"]),
|
|
}
|
|
|
|
|
|
def load_stream(path: Path, *, expected_sha256: str) -> list[dict[str, Any]]:
|
|
if sha256_file(path) != expected_sha256:
|
|
raise ValueError(f"engine stream hash mismatch: {path}")
|
|
decoded = load_jsonl(path)
|
|
records = [row for row in decoded if "step_index" in row]
|
|
if not records:
|
|
raise ValueError(f"engine stream has no Layer-1 records: {path}")
|
|
return records
|
|
|
|
|
|
def build_dataset(
|
|
*, audit_path: Path, manifest_path: Path, run_root: Path
|
|
) -> dict[str, Any]:
|
|
audit = json.loads(audit_path.read_text(encoding="utf-8"))
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
if audit.get("schema") != "action-aware-constraint-pilot-audit-v0":
|
|
raise ValueError("unexpected action-aware audit schema")
|
|
if audit["sanity"]["red_flags"]:
|
|
raise ValueError(f"action-aware audit red flags: {audit['sanity']['red_flags']}")
|
|
configs = {str(item["id"]): item for item in manifest["configs"]}
|
|
runs = {
|
|
(str(run["config_id"]), int(run["repetition"])): run
|
|
for run in audit["runs"]
|
|
}
|
|
source_ids = {str(regime["source"]) for regime in manifest["regimes"].values()}
|
|
stream_entries = {
|
|
str(item["config_id"]): item
|
|
for item in audit["streams"]
|
|
if str(item["config_id"]) in source_ids
|
|
}
|
|
if set(stream_entries) != source_ids:
|
|
raise ValueError("audit is missing a source config engine stream")
|
|
streams = {
|
|
config_id: load_stream(
|
|
Path(item["stream"]), expected_sha256=str(item["stream_sha256"])
|
|
)
|
|
for config_id, item in stream_entries.items()
|
|
}
|
|
examples = []
|
|
request_hashes = []
|
|
for regime_name, regime in sorted(manifest["regimes"].items()):
|
|
source_id = str(regime["source"])
|
|
for repetition in sorted(int(value) for value in manifest["repetitions"]):
|
|
source_run = runs[(source_id, repetition)]
|
|
source_config = configs[source_id]
|
|
request_path = run_root / "sessions" / source_id / f"rep{repetition}" / "requests.jsonl"
|
|
requests = load_jsonl(request_path)
|
|
request_hashes.append(sha256_file(request_path))
|
|
offered_rate_per_gpu = float(
|
|
manifest["repetitions"][str(repetition)]["selection"][
|
|
"offered_req_s_per_gpu"
|
|
]
|
|
)
|
|
offered_total = offered_rate_per_gpu * int(manifest["engine"]["tp"])
|
|
source_goodput = float(source_run["outcome"]["slo_goodput_req_s"])
|
|
source_normalized = min(1.0, source_goodput / offered_total)
|
|
decision_id = f"{regime_name}-rep{repetition}"
|
|
for phase in PHASES:
|
|
cutoff_s = float(manifest["engine"]["duration_s"]) * float(phase)
|
|
outcome = prefix_outcome(
|
|
requests, cutoff_s=cutoff_s, offered_total=offered_total
|
|
)
|
|
admitted_count = sum(
|
|
float(request["arrival_s"]) <= cutoff_s for request in requests
|
|
)
|
|
start_ns = int(source_run["state"]["interval"]["start_ns"])
|
|
phase_state = summarize_engine(
|
|
streams[source_id],
|
|
start_ns=start_ns,
|
|
end_ns=start_ns + round(cutoff_s * 1e9),
|
|
request_count=admitted_count,
|
|
)
|
|
if not all(phase_state["sanity"]["invariants"].values()):
|
|
raise ValueError(
|
|
f"engine state invariant failed: {decision_id} phase {phase}"
|
|
)
|
|
telemetry = telemetry_record(phase_state)
|
|
actions = {"noop": source_id, **regime["actions"]}
|
|
for action_name, target_id in sorted(actions.items()):
|
|
target_run = runs[(str(target_id), repetition)]
|
|
target_config = configs[str(target_id)]
|
|
target_goodput = float(target_run["outcome"]["slo_goodput_req_s"])
|
|
normalized = target_goodput / offered_total
|
|
if not 0.0 <= normalized <= 1.0 + 1e-12:
|
|
raise ValueError("target normalized goodput is outside [0, 1]")
|
|
examples.append(
|
|
{
|
|
"phase": phase,
|
|
"cutoff_s": cutoff_s,
|
|
"decision_id": decision_id,
|
|
"regime": regime_name,
|
|
"repetition": repetition,
|
|
"source": {
|
|
"config_id": source_id,
|
|
"mns": int(source_config["mns"]),
|
|
"mbbt": int(source_config["mbbt"]),
|
|
"offered_rate_per_gpu": offered_rate_per_gpu,
|
|
"outcome": outcome,
|
|
"telemetry": telemetry,
|
|
},
|
|
"action": {
|
|
"id": action_name,
|
|
"target_config_id": str(target_id),
|
|
"target_mns": int(target_config["mns"]),
|
|
"target_mbbt": int(target_config["mbbt"]),
|
|
},
|
|
"target_slo_goodput_req_s": target_goodput,
|
|
"target_normalized_goodput": min(1.0, normalized),
|
|
"source_normalized_goodput": source_normalized,
|
|
"target_delta_normalized_goodput": min(1.0, normalized)
|
|
- source_normalized,
|
|
}
|
|
)
|
|
invariants = {
|
|
"expected_examples": len(examples) == len(PHASES) * 2 * 3 * 3,
|
|
"four_phases": sorted({example["phase"] for example in examples})
|
|
== sorted(PHASES),
|
|
"six_decisions": len({example["decision_id"] for example in examples}) == 6,
|
|
"three_actions_per_decision_phase": all(
|
|
sum(
|
|
item["decision_id"] == decision
|
|
and item["phase"] == phase
|
|
for item in examples
|
|
)
|
|
== 3
|
|
for decision in {item["decision_id"] for item in examples}
|
|
for phase in PHASES
|
|
),
|
|
"targets_not_all_identical": len(
|
|
{example["target_normalized_goodput"] for example in examples}
|
|
)
|
|
> 1,
|
|
"bounded_prefix_ratios": all(
|
|
0.0 <= float(value) <= 1.0
|
|
for example in examples
|
|
for key, value in example["source"]["outcome"].items()
|
|
if key
|
|
in {
|
|
"admitted_fraction",
|
|
"completed_over_admitted",
|
|
"completed_pass_rate",
|
|
"completed_fail_fraction_of_total",
|
|
"outstanding_over_admitted",
|
|
}
|
|
),
|
|
"direct_telemetry_without_binding_labels": all(
|
|
not any(token in key for token in ("exclusive", "unresolved", "both"))
|
|
for example in examples
|
|
for key in example["source"]["telemetry"]
|
|
),
|
|
"treatment_effects_bounded": all(
|
|
-1.0 <= float(example["target_delta_normalized_goodput"]) <= 1.0
|
|
for example in examples
|
|
),
|
|
}
|
|
red_flags = [name for name, passed in invariants.items() if not passed]
|
|
if red_flags:
|
|
raise RuntimeError(f"training dataset sanity failed: {red_flags}")
|
|
return {
|
|
"schema": "active-intervention-training-v0",
|
|
"status": "VALID",
|
|
"provenance": {
|
|
"audit": str(audit_path),
|
|
"audit_sha256": sha256_file(audit_path),
|
|
"manifest": str(manifest_path),
|
|
"manifest_sha256": sha256_file(manifest_path),
|
|
"run_root": str(run_root),
|
|
"source_request_sha256": sorted(set(request_hashes)),
|
|
"source_stream_sha256": sorted(
|
|
str(item["stream_sha256"]) for item in stream_entries.values()
|
|
),
|
|
},
|
|
"examples": examples,
|
|
"sanity": {"invariants": invariants, "red_flags": red_flags},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--audit", type=Path, required=True)
|
|
parser.add_argument("--manifest", type=Path, required=True)
|
|
parser.add_argument("--run-root", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
dataset = build_dataset(
|
|
audit_path=args.audit,
|
|
manifest_path=args.manifest,
|
|
run_root=args.run_root,
|
|
)
|
|
atomic_json(args.output, dataset)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"status": dataset["status"],
|
|
"examples": len(dataset["examples"]),
|
|
"sanity": dataset["sanity"],
|
|
},
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|