520 lines
21 KiB
Python
520 lines
21 KiB
Python
#!/usr/bin/env python3
|
|
"""Train and audit outcome-only versus telemetry action-response policies."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Any, Mapping, Sequence
|
|
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
|
|
|
|
def _load_model():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"active_intervention_model", HERE / "model.py"
|
|
)
|
|
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_model()
|
|
REGULARIZATION = 10.0
|
|
MINIMUM_MARGIN = 0.02
|
|
CONFIDENCE_Z = 1.0
|
|
ACCEPTABLE_REGRET = 0.02
|
|
|
|
|
|
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 grouped(
|
|
examples: Sequence[Mapping[str, Any]], key: str
|
|
) -> dict[str, list[Mapping[str, Any]]]:
|
|
result: dict[str, list[Mapping[str, Any]]] = defaultdict(list)
|
|
for example in examples:
|
|
result[str(example[key])].append(example)
|
|
return dict(result)
|
|
|
|
|
|
def evaluate_grouped_cv(
|
|
examples: Sequence[Mapping[str, Any]],
|
|
*,
|
|
include_telemetry: bool,
|
|
holdout_key: str,
|
|
) -> dict[str, Any]:
|
|
holdouts = grouped(examples, holdout_key)
|
|
decision_rows = []
|
|
for held_out, test_examples in sorted(holdouts.items()):
|
|
training = [example for example in examples if str(example[holdout_key]) != held_out]
|
|
if len({str(example["decision_id"]) for example in training}) < 2:
|
|
continue
|
|
model = MODEL.fit_ridge(
|
|
training,
|
|
include_telemetry=include_telemetry,
|
|
regularization=REGULARIZATION,
|
|
)
|
|
for decision_id, candidates in sorted(grouped(test_examples, "decision_id").items()):
|
|
predictions = []
|
|
for candidate in candidates:
|
|
source = candidate["source"]
|
|
action = candidate["action"]
|
|
noop = (
|
|
int(action["target_mns"]) == int(source["mns"])
|
|
and int(action["target_mbbt"]) == int(source["mbbt"])
|
|
)
|
|
if noop:
|
|
prediction = 0.0
|
|
else:
|
|
names, vector = MODEL.feature_vector(
|
|
candidate, include_telemetry=include_telemetry
|
|
)
|
|
if tuple(names) != model.feature_names:
|
|
raise ValueError("cross-validation feature schema mismatch")
|
|
prediction = max(-1.0, min(1.0, model.predict(vector)))
|
|
predictions.append(
|
|
{
|
|
"action_id": str(candidate["action"]["id"]),
|
|
"prediction": prediction,
|
|
"real": float(candidate["target_normalized_goodput"]),
|
|
}
|
|
)
|
|
predictions.sort(key=lambda row: (-row["prediction"], row["action_id"]))
|
|
selected = predictions[0]
|
|
oracle = max(row["real"] for row in predictions)
|
|
regret = 1.0 - selected["real"] / oracle if oracle > 0 else 0.0
|
|
best_actions = {
|
|
row["action_id"] for row in predictions if math.isclose(row["real"], oracle)
|
|
}
|
|
acceptable_actions = {
|
|
row["action_id"]
|
|
for row in predictions
|
|
if oracle <= 0
|
|
or 1.0 - float(row["real"]) / oracle <= ACCEPTABLE_REGRET + 1e-12
|
|
}
|
|
decision_rows.append(
|
|
{
|
|
"holdout": held_out,
|
|
"decision_id": decision_id,
|
|
"selected_action": selected["action_id"],
|
|
"best_actions": sorted(best_actions),
|
|
"acceptable_actions": sorted(acceptable_actions),
|
|
"correct": regret <= ACCEPTABLE_REGRET + 1e-12,
|
|
"selected_real": selected["real"],
|
|
"oracle_real": oracle,
|
|
"regret": regret,
|
|
"predictions": predictions,
|
|
}
|
|
)
|
|
if not decision_rows:
|
|
return {"status": "INSUFFICIENT_GROUPS", "decisions": []}
|
|
regrets = [float(row["regret"]) for row in decision_rows]
|
|
return {
|
|
"status": "VALID",
|
|
"holdout_key": holdout_key,
|
|
"acceptable_regret": ACCEPTABLE_REGRET,
|
|
"decision_n": len(decision_rows),
|
|
"correct_n": sum(bool(row["correct"]) for row in decision_rows),
|
|
"accuracy": sum(bool(row["correct"]) for row in decision_rows) / len(decision_rows),
|
|
"mean_regret": sum(regrets) / len(regrets),
|
|
"max_regret": max(regrets),
|
|
"decisions": decision_rows,
|
|
}
|
|
|
|
|
|
def paired_delta(outcome: Mapping[str, Any], telemetry: Mapping[str, Any]) -> dict[str, Any]:
|
|
if outcome.get("status") != "VALID" or telemetry.get("status") != "VALID":
|
|
return {"status": "INSUFFICIENT_GROUPS"}
|
|
outcome_by_id = {row["decision_id"]: row for row in outcome["decisions"]}
|
|
telemetry_by_id = {row["decision_id"]: row for row in telemetry["decisions"]}
|
|
common = sorted(set(outcome_by_id) & set(telemetry_by_id))
|
|
rows = []
|
|
for decision_id in common:
|
|
before = outcome_by_id[decision_id]
|
|
after = telemetry_by_id[decision_id]
|
|
rows.append(
|
|
{
|
|
"decision_id": decision_id,
|
|
"outcome_action": before["selected_action"],
|
|
"telemetry_action": after["selected_action"],
|
|
"action_changed": before["selected_action"] != after["selected_action"],
|
|
"regret_delta": float(after["regret"]) - float(before["regret"]),
|
|
"telemetry_corrected": (not before["correct"]) and bool(after["correct"]),
|
|
"telemetry_harmed": bool(before["correct"]) and (not after["correct"]),
|
|
}
|
|
)
|
|
return {
|
|
"status": "VALID",
|
|
"decision_n": len(rows),
|
|
"action_changed_n": sum(row["action_changed"] for row in rows),
|
|
"corrected_n": sum(row["telemetry_corrected"] for row in rows),
|
|
"harmed_n": sum(row["telemetry_harmed"] for row in rows),
|
|
"mean_regret_delta": (
|
|
sum(float(row["regret_delta"]) for row in rows) / len(rows) if rows else 0.0
|
|
),
|
|
"rows": rows,
|
|
}
|
|
|
|
|
|
def evaluate_sequential_measurement_cv(
|
|
examples: Sequence[Mapping[str, Any]],
|
|
*,
|
|
include_telemetry: bool,
|
|
holdout_key: str,
|
|
) -> dict[str, Any]:
|
|
"""Replay a two-consecutive-confident-checkpoint measurement policy."""
|
|
|
|
phases = sorted({str(example["phase"]) for example in examples}, key=float)
|
|
holdouts = grouped(examples, holdout_key)
|
|
rows = []
|
|
full_duration_s = max(float(example["cutoff_s"]) for example in examples)
|
|
for held_out, test_examples in sorted(holdouts.items()):
|
|
training = [
|
|
example for example in examples if str(example[holdout_key]) != held_out
|
|
]
|
|
if len({str(example["decision_id"]) for example in training}) < 3:
|
|
continue
|
|
phase_models = {}
|
|
for phase in phases:
|
|
phase_training = [
|
|
example for example in training if str(example["phase"]) == phase
|
|
]
|
|
phase_models[phase] = MODEL.fit_jackknife_ensemble(
|
|
phase_training,
|
|
include_telemetry=include_telemetry,
|
|
regularization=REGULARIZATION,
|
|
)
|
|
for decision_id, decision_examples in sorted(
|
|
grouped(test_examples, "decision_id").items()
|
|
):
|
|
checkpoints = []
|
|
by_phase = grouped(decision_examples, "phase")
|
|
for phase in phases:
|
|
candidates = by_phase[phase]
|
|
decision = MODEL.select_action(
|
|
phase_models[phase],
|
|
candidates,
|
|
include_telemetry=include_telemetry,
|
|
confidence_z=CONFIDENCE_Z,
|
|
minimum_margin=MINIMUM_MARGIN,
|
|
)
|
|
checkpoints.append(
|
|
{
|
|
"phase": phase,
|
|
"cutoff_s": float(candidates[0]["cutoff_s"]),
|
|
**decision,
|
|
}
|
|
)
|
|
selected_checkpoint = 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_checkpoint = current
|
|
stop_reason = "two_consecutive_confident_checkpoints"
|
|
break
|
|
candidates = by_phase[str(selected_checkpoint["phase"])]
|
|
real_by_action = {
|
|
str(candidate["action"]["id"]): float(
|
|
candidate["target_normalized_goodput"]
|
|
)
|
|
for candidate in candidates
|
|
}
|
|
target_by_action = {
|
|
str(candidate["action"]["id"]): str(
|
|
candidate["action"]["target_config_id"]
|
|
)
|
|
for candidate in candidates
|
|
}
|
|
selected_action = str(selected_checkpoint["selected_action"])
|
|
oracle = max(real_by_action.values())
|
|
selected_real = real_by_action[selected_action]
|
|
regret = 1.0 - selected_real / oracle if oracle > 0 else 0.0
|
|
source_tp = 4
|
|
target_s = 0.0 if selected_action == "noop" else full_duration_s
|
|
replay_gpu_seconds = source_tp * (
|
|
float(selected_checkpoint["cutoff_s"]) + target_s
|
|
)
|
|
rows.append(
|
|
{
|
|
"holdout": held_out,
|
|
"decision_id": decision_id,
|
|
"selected_phase": str(selected_checkpoint["phase"]),
|
|
"selected_cutoff_s": float(selected_checkpoint["cutoff_s"]),
|
|
"stop_reason": stop_reason,
|
|
"selected_action": selected_action,
|
|
"selected_target_config_id": target_by_action[selected_action],
|
|
"selected_real": selected_real,
|
|
"oracle_real": oracle,
|
|
"regret": regret,
|
|
"acceptable": regret <= ACCEPTABLE_REGRET + 1e-12,
|
|
"replay_gpu_seconds_lower_bound": replay_gpu_seconds,
|
|
"checkpoints": checkpoints,
|
|
}
|
|
)
|
|
if not rows:
|
|
return {"status": "INSUFFICIENT_GROUPS", "decisions": []}
|
|
regrets = [float(row["regret"]) for row in rows]
|
|
cutoffs = [float(row["selected_cutoff_s"]) for row in rows]
|
|
costs = [float(row["replay_gpu_seconds_lower_bound"]) for row in rows]
|
|
return {
|
|
"status": "VALID",
|
|
"holdout_key": holdout_key,
|
|
"measurement_rule": "earliest two consecutive confident checkpoints; otherwise full",
|
|
"acceptable_regret": ACCEPTABLE_REGRET,
|
|
"decision_n": len(rows),
|
|
"acceptable_n": sum(bool(row["acceptable"]) for row in rows),
|
|
"mean_regret": sum(regrets) / len(regrets),
|
|
"max_regret": max(regrets),
|
|
"mean_cutoff_s": sum(cutoffs) / len(cutoffs),
|
|
"total_replay_gpu_seconds_lower_bound": sum(costs),
|
|
"total_replay_h20_hours_lower_bound": sum(costs) / 3600.0,
|
|
"decisions": rows,
|
|
}
|
|
|
|
|
|
def paired_sequential_delta(
|
|
outcome: Mapping[str, Any], telemetry: Mapping[str, Any]
|
|
) -> dict[str, Any]:
|
|
if outcome.get("status") != "VALID" or telemetry.get("status") != "VALID":
|
|
return {"status": "INSUFFICIENT_GROUPS"}
|
|
before_by_id = {row["decision_id"]: row for row in outcome["decisions"]}
|
|
after_by_id = {row["decision_id"]: row for row in telemetry["decisions"]}
|
|
rows = []
|
|
for decision_id in sorted(set(before_by_id) & set(after_by_id)):
|
|
before = before_by_id[decision_id]
|
|
after = after_by_id[decision_id]
|
|
rows.append(
|
|
{
|
|
"decision_id": decision_id,
|
|
"outcome_action": before["selected_action"],
|
|
"telemetry_action": after["selected_action"],
|
|
"outcome_cutoff_s": before["selected_cutoff_s"],
|
|
"telemetry_cutoff_s": after["selected_cutoff_s"],
|
|
"outcome_regret": before["regret"],
|
|
"telemetry_regret": after["regret"],
|
|
"regret_delta": float(after["regret"]) - float(before["regret"]),
|
|
"gpu_seconds_delta": float(
|
|
after["replay_gpu_seconds_lower_bound"]
|
|
)
|
|
- float(before["replay_gpu_seconds_lower_bound"]),
|
|
"telemetry_corrected": (not before["acceptable"])
|
|
and bool(after["acceptable"]),
|
|
"telemetry_harmed": bool(before["acceptable"])
|
|
and (not after["acceptable"]),
|
|
}
|
|
)
|
|
outcome_cost = float(outcome["total_replay_gpu_seconds_lower_bound"])
|
|
telemetry_cost = float(telemetry["total_replay_gpu_seconds_lower_bound"])
|
|
return {
|
|
"status": "VALID",
|
|
"decision_n": len(rows),
|
|
"corrected_n": sum(row["telemetry_corrected"] for row in rows),
|
|
"harmed_n": sum(row["telemetry_harmed"] for row in rows),
|
|
"outcome_replay_gpu_seconds_lower_bound": outcome_cost,
|
|
"telemetry_replay_gpu_seconds_lower_bound": telemetry_cost,
|
|
"gpu_cost_reduction_fraction": (
|
|
1.0 - telemetry_cost / outcome_cost if outcome_cost > 0 else 0.0
|
|
),
|
|
"rows": rows,
|
|
}
|
|
|
|
|
|
def build_policy(dataset_path: Path) -> dict[str, Any]:
|
|
dataset = json.loads(dataset_path.read_text(encoding="utf-8"))
|
|
if dataset.get("status") != "VALID" or dataset["sanity"]["red_flags"]:
|
|
raise ValueError("training dataset is not valid")
|
|
examples = dataset["examples"]
|
|
phases = sorted({str(example["phase"]) for example in examples}, key=float)
|
|
phase_results = {}
|
|
incremental_candidates = []
|
|
for phase in phases:
|
|
selected = [example for example in examples if str(example["phase"]) == phase]
|
|
outcome_cv = evaluate_grouped_cv(
|
|
selected, include_telemetry=False, holdout_key="repetition"
|
|
)
|
|
telemetry_cv = evaluate_grouped_cv(
|
|
selected, include_telemetry=True, holdout_key="repetition"
|
|
)
|
|
outcome_regime = evaluate_grouped_cv(
|
|
selected, include_telemetry=False, holdout_key="regime"
|
|
)
|
|
telemetry_regime = evaluate_grouped_cv(
|
|
selected, include_telemetry=True, holdout_key="regime"
|
|
)
|
|
delta = paired_delta(outcome_cv, telemetry_cv)
|
|
outcome_models = MODEL.fit_jackknife_ensemble(
|
|
selected,
|
|
include_telemetry=False,
|
|
regularization=REGULARIZATION,
|
|
)
|
|
telemetry_models = MODEL.fit_jackknife_ensemble(
|
|
selected,
|
|
include_telemetry=True,
|
|
regularization=REGULARIZATION,
|
|
)
|
|
incremental = bool(
|
|
delta.get("status") == "VALID"
|
|
and int(delta["corrected_n"]) >= 1
|
|
and int(delta["harmed_n"]) == 0
|
|
and float(delta["mean_regret_delta"]) < -1e-12
|
|
and float(telemetry_cv["max_regret"]) <= 0.05
|
|
and telemetry_regime.get("status") == "VALID"
|
|
and float(telemetry_regime["mean_regret"])
|
|
<= float(outcome_regime["mean_regret"]) + 1e-12
|
|
and float(telemetry_regime["max_regret"]) <= 0.05
|
|
)
|
|
if incremental:
|
|
incremental_candidates.append(phase)
|
|
phase_results[phase] = {
|
|
"cutoff_s": float(selected[0]["cutoff_s"]),
|
|
"outcome_only": {
|
|
"leave_repetition_out": outcome_cv,
|
|
"leave_regime_out": outcome_regime,
|
|
"models": MODEL.models_to_json(outcome_models),
|
|
},
|
|
"telemetry": {
|
|
"leave_repetition_out": telemetry_cv,
|
|
"leave_regime_out": telemetry_regime,
|
|
"models": MODEL.models_to_json(telemetry_models),
|
|
},
|
|
"paired_incremental": delta,
|
|
"incremental_gate": incremental,
|
|
}
|
|
outcome_sequential = evaluate_sequential_measurement_cv(
|
|
examples, include_telemetry=False, holdout_key="repetition"
|
|
)
|
|
telemetry_sequential = evaluate_sequential_measurement_cv(
|
|
examples, include_telemetry=True, holdout_key="repetition"
|
|
)
|
|
sequential_delta = paired_sequential_delta(
|
|
outcome_sequential, telemetry_sequential
|
|
)
|
|
retrospective_cost_gate = bool(
|
|
sequential_delta.get("status") == "VALID"
|
|
and int(sequential_delta["harmed_n"]) == 0
|
|
and int(telemetry_sequential["acceptable_n"])
|
|
>= int(outcome_sequential["acceptable_n"])
|
|
and float(telemetry_sequential["max_regret"]) <= 0.05
|
|
and float(sequential_delta["gpu_cost_reduction_fraction"]) >= 0.10
|
|
)
|
|
status = (
|
|
"RETROSPECTIVE_GPU_COST_SIGNAL"
|
|
if retrospective_cost_gate
|
|
else "NO_RETROSPECTIVE_GPU_COST_SIGNAL"
|
|
)
|
|
target_values = [float(example["target_normalized_goodput"]) for example in examples]
|
|
effect_values = [
|
|
float(example["target_delta_normalized_goodput"]) for example in examples
|
|
]
|
|
invariants = {
|
|
"four_phases": len(phases) == 4,
|
|
"targets_bounded": all(0.0 <= value <= 1.0 for value in target_values),
|
|
"targets_not_all_identical": len(set(target_values)) > 1,
|
|
"effects_bounded": all(-1.0 <= value <= 1.0 for value in effect_values),
|
|
"effects_not_all_identical": len(set(effect_values)) > 1,
|
|
"models_present_every_phase": all(
|
|
phase_results[phase][mode]["models"]
|
|
for phase in phases
|
|
for mode in ("outcome_only", "telemetry")
|
|
),
|
|
}
|
|
red_flags = [name for name, passed in invariants.items() if not passed]
|
|
if red_flags:
|
|
raise RuntimeError(f"policy sanity failed: {red_flags}")
|
|
return {
|
|
"schema": "active-intervention-policy-v0",
|
|
"status": status,
|
|
"training": {
|
|
"dataset": str(dataset_path),
|
|
"dataset_sha256": sha256_file(dataset_path),
|
|
"examples": len(examples),
|
|
"decisions": len({example["decision_id"] for example in examples}),
|
|
"regularization": REGULARIZATION,
|
|
"confidence_z": CONFIDENCE_Z,
|
|
"minimum_margin": MINIMUM_MARGIN,
|
|
"acceptable_regret": ACCEPTABLE_REGRET,
|
|
},
|
|
"measurement_policy": {
|
|
"rule": "earliest two consecutive confident checkpoints; otherwise full",
|
|
"checkpoints": [phase_results[phase]["cutoff_s"] for phase in phases],
|
|
"confidence_z": CONFIDENCE_Z,
|
|
"minimum_margin": MINIMUM_MARGIN,
|
|
},
|
|
"sequential_replay": {
|
|
"outcome_only": outcome_sequential,
|
|
"telemetry": telemetry_sequential,
|
|
"paired_delta": sequential_delta,
|
|
"retrospective_gpu_cost_gate": retrospective_cost_gate,
|
|
"minimum_cost_reduction_fraction": 0.10,
|
|
},
|
|
"phases": phase_results,
|
|
"sanity": {
|
|
"invariants": invariants,
|
|
"red_flags": red_flags,
|
|
"target_normalized_goodput": {
|
|
"n": len(target_values),
|
|
"min": min(target_values),
|
|
"max": max(target_values),
|
|
"distinct_n": len(set(target_values)),
|
|
},
|
|
"target_delta_normalized_goodput": {
|
|
"n": len(effect_values),
|
|
"min": min(effect_values),
|
|
"max": max(effect_values),
|
|
"distinct_n": len(set(effect_values)),
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
policy = build_policy(args.dataset)
|
|
atomic_json(args.output, policy)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"status": policy["status"],
|
|
"measurement_policy": policy["measurement_policy"],
|
|
"sanity": policy["sanity"],
|
|
},
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|