#!/usr/bin/env python3 """Choose measurement horizon and next intervention from a completed source run.""" from __future__ import annotations import argparse import hashlib import importlib.util import json import math import os import statistics import sys from pathlib import Path from typing import Any, Mapping, Sequence import numpy as np 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 = "active-intervention-prospective-decision-v0" def load_module(name: str, path: Path): spec = importlib.util.spec_from_file_location(name, path) module = importlib.util.module_from_spec(spec) assert spec.loader is not None sys.modules[spec.name] = module spec.loader.exec_module(module) return module MODEL = load_module("active_intervention_prospective_model", HERE / "model.py") EXTRACT = load_module( "active_intervention_prospective_extract", HERE / "extract_training.py" ) 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: Sequence[float]) -> dict[str, Any]: finite = [float(value) for value in values] if not finite or any(not math.isfinite(value) for value in finite): raise ValueError("numeric summary requires finite values") return { "n": len(finite), "min": min(finite), "max": max(finite), "distinct_n": len(set(finite)), } def load_engine_records(source_root: Path) -> tuple[list[dict[str, Any]], Path]: streams = sorted((source_root / "opprof").glob("*.jsonl")) if len(streams) != 1: raise ValueError(f"expected one source engine stream, found {len(streams)}") records = [ row for row in EXTRACT.load_jsonl(streams[0]) if "step_index" in row ] if not records: raise ValueError("source engine stream has no Layer-1 records") return records, streams[0] def candidate_example( *, source_config: Mapping[str, Any], target_config: Mapping[str, Any], action_id: str, offered_rate_per_gpu: float, outcome: Mapping[str, float], telemetry: Mapping[str, float], ) -> dict[str, Any]: return { "source": { "mns": int(source_config["mns"]), "mbbt": int(source_config["mbbt"]), "offered_rate_per_gpu": float(offered_rate_per_gpu), "outcome": dict(outcome), "telemetry": dict(telemetry), }, "action": { "id": action_id, "target_mns": int(target_config["mns"]), "target_mbbt": int(target_config["mbbt"]), }, } def aggregate_checkpoint( *, models: Sequence[Any], examples_by_action: Mapping[str, Sequence[Mapping[str, Any]]], include_telemetry: bool, confidence_z: float, minimum_margin: float, ) -> dict[str, Any]: rows = [] for action_id, examples in sorted(examples_by_action.items()): raw = [] for example in examples: source = example["source"] action = example["action"] noop = ( int(source["mns"]) == int(action["target_mns"]) and int(source["mbbt"]) == int(action["target_mbbt"]) ) if noop: raw.extend(0.0 for _model in models) continue names, values = MODEL.feature_vector( example, include_telemetry=include_telemetry ) if any(model.feature_names != tuple(names) for model in models): raise ValueError("prospective feature schema does not match frozen model") raw.extend(model.predict(values) for model in models) clipped = np.clip(np.asarray(raw, dtype=np.float64), -1.0, 1.0) prediction = { "mean": float(clipped.mean()), "std": float(clipped.std(ddof=0)), "min": float(clipped.min()), "max": float(clipped.max()), "distinct_n": len(set(float(value) for value in clipped)), "sample_n": int(clipped.size), } rows.append( { "action_id": action_id, "prediction": prediction, "lower": prediction["mean"] - confidence_z * prediction["std"], "upper": prediction["mean"] + confidence_z * prediction["std"], } ) rows.sort(key=lambda row: (-row["prediction"]["mean"], row["action_id"])) best, second = rows[:2] margin = float(best["prediction"]["mean"] - second["prediction"]["mean"]) confident = bool( margin >= minimum_margin and best["lower"] > second["upper"] ) return { "selected_action": best["action_id"], "confident": confident, "predicted_margin": margin, "candidates": rows, } def apply_measurement_and_acquisition(checkpoints: list[dict[str, Any]]) -> dict[str, Any]: selected = checkpoints[-1] stop_reason = "full_measurement_fallback" for previous, current in zip(checkpoints, checkpoints[1:], strict=False): if ( previous["confident"] and current["confident"] and previous["selected_action"] == current["selected_action"] ): selected = current stop_reason = "two_consecutive_confident_checkpoints" break candidates = selected["candidates"] mean_best = candidates[0] non_noop = [row for row in candidates if row["action_id"] != "noop"] if selected["confident"]: chosen = mean_best decision_kind = "exploit" else: positive_ucb = [row for row in non_noop if float(row["upper"]) > 0.0] if positive_ucb: chosen = max( positive_ucb, key=lambda row: (float(row["upper"]), row["action_id"]), ) decision_kind = "diagnostic_ucb" else: chosen = next(row for row in candidates if row["action_id"] == "noop") decision_kind = "abstain_no_positive_ucb" remaining = [row for row in candidates if row["action_id"] != chosen["action_id"]] remaining.sort(key=lambda row: (-float(row["upper"]), row["action_id"])) order = [chosen["action_id"], *(row["action_id"] for row in remaining)] return { "selected_phase": selected["phase"], "selected_cutoff_s": selected["cutoff_s"], "measurement_stop_reason": stop_reason, "decision_kind": decision_kind, "selected_action": chosen["action_id"], "intervention_order": order, "selected_checkpoint": selected, "checkpoints": checkpoints, } def build_decision( *, manifest_path: Path, policy_path: Path, run_root: Path ) -> dict[str, Any]: manifest = json.loads(manifest_path.read_text(encoding="utf-8")) policy = json.loads(policy_path.read_text(encoding="utf-8")) if manifest.get("schema") != "active-intervention-prospective-manifest-v0": raise ValueError("unexpected prospective manifest schema") if policy.get("schema") != "active-intervention-policy-v0": raise ValueError("unexpected frozen policy schema") if sha256_file(policy_path) != manifest["policy"]["sha256"]: raise ValueError("frozen policy hash changed after manifest preparation") configs = {str(item["id"]): item for item in manifest["configs"]} source_id = str(manifest["source_config_id"]) source_config = configs[source_id] source_root = run_root / "sessions" / source_id engine_records, stream_path = load_engine_records(source_root) phases = [f"{fraction:.2f}" for fraction in manifest["checkpoints"]["fractions"]] confidence_z = float(policy["measurement_policy"]["confidence_z"]) minimum_margin = float(policy["measurement_policy"]["minimum_margin"]) examples: dict[str, dict[str, dict[str, Mapping[str, Any]]]] = {} source_measurements: dict[str, dict[str, Any]] = {} source_normalized = [] telemetry_values = [] for repetition in sorted(int(key) for key in manifest["repetitions"]): item = manifest["repetitions"][str(repetition)] result_root = source_root / f"rep{repetition}" result = json.loads((result_root / "result.json").read_text(encoding="utf-8")) if result["selection"]["request_id_order_sha256"] != item["selection"][ "request_id_order_sha256" ]: raise ValueError(f"source request hash mismatch: rep{repetition}") requests = EXTRACT.load_jsonl(result_root / "requests.jsonl") offered_rate = float(item["selection"]["offered_req_s_per_gpu"]) offered_total = offered_rate * int(manifest["engine"]["tp"]) source_normalized.append( float(result["slo_pass_count"]) / float(manifest["engine"]["duration_s"]) / offered_total ) start_ns = int(result["interval"]["start_mono_ns"]) examples[str(repetition)] = {} source_measurements[str(repetition)] = { "result": str(result_root / "result.json"), "result_sha256": sha256_file(result_root / "result.json"), "request_sha256": sha256_file(result_root / "requests.jsonl"), "phases": {}, } for phase, cutoff_s in zip( phases, manifest["checkpoints"]["seconds"], strict=True ): outcome = EXTRACT.prefix_outcome( requests, cutoff_s=float(cutoff_s), offered_total=offered_total ) admitted_count = sum( float(request["arrival_s"]) <= float(cutoff_s) for request in requests ) state = summarize_engine( engine_records, start_ns=start_ns, end_ns=start_ns + round(float(cutoff_s) * 1e9), request_count=admitted_count, ) if not all(state["sanity"]["invariants"].values()): raise ValueError( f"source engine state invariant failed: rep{repetition} {phase}" ) telemetry = EXTRACT.telemetry_record(state) telemetry_values.extend(float(value) for value in telemetry.values()) source_measurements[str(repetition)]["phases"][phase] = { "cutoff_s": float(cutoff_s), "outcome": outcome, "telemetry": telemetry, "engine_sanity": state["sanity"], } examples[str(repetition)][phase] = { action_id: candidate_example( source_config=source_config, target_config=configs[str(target_id)], action_id=action_id, offered_rate_per_gpu=offered_rate, outcome=outcome, telemetry=telemetry, ) for action_id, target_id in manifest["actions"].items() } decisions = {} for mode, include_telemetry in (("outcome_only", False), ("telemetry", True)): checkpoints = [] for phase, cutoff_s in zip( phases, manifest["checkpoints"]["seconds"], strict=True ): models = MODEL.models_from_json(policy["phases"][phase][mode]["models"]) examples_by_action = { action_id: [ examples[str(repetition)][phase][action_id] for repetition in sorted(int(key) for key in manifest["repetitions"]) ] for action_id in manifest["actions"] } checkpoint = aggregate_checkpoint( models=models, examples_by_action=examples_by_action, include_telemetry=include_telemetry, confidence_z=confidence_z, minimum_margin=minimum_margin, ) checkpoints.append( {"phase": phase, "cutoff_s": float(cutoff_s), **checkpoint} ) decisions[mode] = apply_measurement_and_acquisition(checkpoints) ceiling = float(manifest["gates"]["source_ceiling_normalized_goodput"]) source_median = float(statistics.median(source_normalized)) status = "STOP_SOURCE_CEILING" if source_median >= ceiling else "SELECTED" phase_admission_monotonic = all( all( left <= right + 1e-12 for left, right in zip(values, values[1:], strict=False) ) for repetition in source_measurements.values() for values in ( [ float(repetition["phases"][phase]["outcome"]["admitted_fraction"]) for phase in phases ], ) ) telemetry_ratio_keys = { "prefill_token_fraction", "kv_usage_mean", "kv_usage_max", "graph_none_share", "graph_full_share", "graph_padding_fraction", } telemetry_records = [ measurement["telemetry"] for repetition in source_measurements.values() for measurement in repetition["phases"].values() ] invariants = { "three_source_repetitions": len(source_normalized) == 3, "source_goodput_nonnegative": all(value >= 0.0 for value in source_normalized), "source_goodput_bounded": all( value <= 1.0 + 1e-12 for value in source_normalized ), "four_actions": set(manifest["actions"]) == {"noop", "mns", "mbbt", "joint"}, "four_checkpoints": len(phases) == 4, "finite_telemetry": all(math.isfinite(value) for value in telemetry_values), "nonnegative_telemetry": all( float(value) >= 0.0 for record in telemetry_records for key, value in record.items() if key != "kv_usage_end_minus_start" ), "telemetry_ratios_bounded": all( 0.0 <= float(record[key]) <= 1.0 + 1e-12 for record in telemetry_records for key in telemetry_ratio_keys ), "telemetry_not_all_identical": len(set(telemetry_values)) > 1, "phase_admission_monotonic": phase_admission_monotonic, "orders_are_permutations": all( set(decisions[mode]["intervention_order"]) == set(manifest["actions"]) for mode in decisions ), } red_flags = [name for name, passed in invariants.items() if not passed] if red_flags: status = "STOP_SANITY" return { "schema": SCHEMA, "status": status, "manifest": str(manifest_path), "manifest_sha256": sha256_file(manifest_path), "policy": str(policy_path), "policy_sha256": sha256_file(policy_path), "source_stream": str(stream_path), "source_stream_sha256": sha256_file(stream_path), "source_measurements": source_measurements, "source_normalized_goodput": { "values": source_normalized, "median": source_median, **numeric(source_normalized), }, "decisions": decisions, "sanity": { "invariants": invariants, "red_flags": red_flags, "telemetry_values": numeric(telemetry_values), }, } def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--policy", 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() decision = build_decision( manifest_path=args.manifest, policy_path=args.policy, run_root=args.run_root ) atomic_json(args.output, decision) print( json.dumps( { "status": decision["status"], "source_normalized_goodput": decision["source_normalized_goodput"], "outcome_only": { key: decision["decisions"]["outcome_only"][key] for key in ("selected_cutoff_s", "decision_kind", "selected_action") }, "telemetry": { key: decision["decisions"]["telemetry"][key] for key in ("selected_cutoff_s", "decision_kind", "selected_action") }, }, sort_keys=True, ) ) if __name__ == "__main__": main()