Add action-conditioned intervention feasibility model
This commit is contained in:
324
runs/active-intervention-v0/extract_training.py
Normal file
324
runs/active-intervention-v0/extract_training.py
Normal file
@@ -0,0 +1,324 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user