Add prospective active intervention experiment
This commit is contained in:
441
runs/active-intervention-v0/prospective_decision.py
Normal file
441
runs/active-intervention-v0/prospective_decision.py
Normal file
@@ -0,0 +1,441 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user