442 lines
16 KiB
Python
442 lines
16 KiB
Python
#!/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()
|