diff --git a/runs/active-intervention-v0/extract_training.py b/runs/active-intervention-v0/extract_training.py new file mode 100644 index 0000000..8f7c8bf --- /dev/null +++ b/runs/active-intervention-v0/extract_training.py @@ -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() diff --git a/runs/active-intervention-v0/model.py b/runs/active-intervention-v0/model.py new file mode 100644 index 0000000..fe815e6 --- /dev/null +++ b/runs/active-intervention-v0/model.py @@ -0,0 +1,287 @@ +#!/usr/bin/env python3 +"""Small-data action-response model for the active intervention pilot. + +The model predicts the paired normalized SLO-goodput treatment effect from a +source measurement and a full MNS/MBBT action. Telemetry features are direct, +continuous engine measurements; there is no diagnosis-to-action rule here. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +PREFIX_FEATURES = ( + "normalized_slo_goodput", + "admitted_fraction", + "completed_over_admitted", + "completed_pass_rate", + "completed_fail_fraction_of_total", + "outstanding_over_admitted", + "ttft_max_over_slo_max", + "ttft_mean_over_slo_max", + "tpot_max_over_slo", + "tpot_mean_over_slo", + "admitted_input_tokens_mean_over_limit", +) + +TELEMETRY_FEATURES = ( + "scheduler_steps_per_s", + "batch_size_mean", + "batch_size_cv", + "batch_tokens_mean", + "batch_tokens_cv", + "decode_batch_size_mean", + "decode_batch_size_cv", + "prefill_token_fraction", + "queue_waiting_mean", + "queue_running_mean", + "preemptions_per_step", + "kv_usage_mean", + "kv_usage_max", + "kv_usage_end_minus_start", + "graph_none_share", + "graph_full_share", + "graph_padding_fraction", +) + + +def _finite(value: Any, name: str) -> float: + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{name} must be finite") + return result + + +def feature_vector( + example: Mapping[str, Any], *, include_telemetry: bool +) -> tuple[list[str], np.ndarray]: + source = example["source"] + action = example["action"] + source_log_mns = math.log2(_finite(source["mns"], "source MNS")) + source_log_mbbt = math.log2(_finite(source["mbbt"], "source MBBT")) + target_log_mns = math.log2(_finite(action["target_mns"], "target MNS")) + target_log_mbbt = math.log2(_finite(action["target_mbbt"], "target MBBT")) + delta_mns = target_log_mns - source_log_mns + delta_mbbt = target_log_mbbt - source_log_mbbt + names = [ + "source_log2_mns", + "source_log2_mbbt", + "target_log2_mns", + "target_log2_mbbt", + "delta_log2_mns", + "delta_log2_mbbt", + "delta_product", + "offered_rate_per_gpu", + ] + values = [ + source_log_mns, + source_log_mbbt, + target_log_mns, + target_log_mbbt, + delta_mns, + delta_mbbt, + delta_mns * delta_mbbt, + _finite(source["offered_rate_per_gpu"], "offered rate"), + ] + for name in PREFIX_FEATURES: + names.append(f"outcome.{name}") + values.append(_finite(source["outcome"][name], name)) + if include_telemetry: + for name in TELEMETRY_FEATURES: + value = _finite(source["telemetry"][name], name) + names.extend( + ( + f"telemetry.{name}", + f"telemetry.{name}*delta_mns", + f"telemetry.{name}*delta_mbbt", + ) + ) + values.extend((value, value * delta_mns, value * delta_mbbt)) + vector = np.asarray(values, dtype=np.float64) + if not np.all(np.isfinite(vector)): + raise ValueError("feature vector contains a non-finite value") + return names, vector + + +@dataclass(frozen=True) +class RidgeModel: + feature_names: tuple[str, ...] + mean: np.ndarray + scale: np.ndarray + weights: np.ndarray + intercept: float + regularization: float + + def predict(self, values: np.ndarray) -> float: + if values.shape != self.mean.shape: + raise ValueError("ridge prediction feature shape mismatch") + normalized = (values - self.mean) / self.scale + return float(self.intercept + normalized @ self.weights) + + def to_json(self) -> dict[str, Any]: + return { + "feature_names": list(self.feature_names), + "mean": self.mean.tolist(), + "scale": self.scale.tolist(), + "weights": self.weights.tolist(), + "intercept": self.intercept, + "regularization": self.regularization, + } + + @classmethod + def from_json(cls, payload: Mapping[str, Any]) -> "RidgeModel": + return cls( + feature_names=tuple(str(value) for value in payload["feature_names"]), + mean=np.asarray(payload["mean"], dtype=np.float64), + scale=np.asarray(payload["scale"], dtype=np.float64), + weights=np.asarray(payload["weights"], dtype=np.float64), + intercept=float(payload["intercept"]), + regularization=float(payload["regularization"]), + ) + + +def fit_ridge( + examples: Sequence[Mapping[str, Any]], + *, + include_telemetry: bool, + regularization: float, +) -> RidgeModel: + if not examples: + raise ValueError("ridge fit requires examples") + if regularization <= 0: + raise ValueError("ridge regularization must be positive") + encoded = [ + feature_vector(example, include_telemetry=include_telemetry) + for example in examples + ] + names = encoded[0][0] + if any(item[0] != names for item in encoded): + raise ValueError("feature names changed across examples") + x = np.stack([item[1] for item in encoded]) + y = np.asarray( + [ + _finite(example["target_delta_normalized_goodput"], "target effect") + for example in examples + ], + dtype=np.float64, + ) + mean = x.mean(axis=0) + scale = x.std(axis=0) + scale[scale < 1e-12] = 1.0 + normalized = (x - mean) / scale + intercept = float(y.mean()) + centered = y - intercept + system = normalized.T @ normalized + regularization * np.eye(x.shape[1]) + weights = np.linalg.solve(system, normalized.T @ centered) + return RidgeModel( + feature_names=tuple(names), + mean=mean, + scale=scale, + weights=weights, + intercept=intercept, + regularization=regularization, + ) + + +def fit_jackknife_ensemble( + examples: Sequence[Mapping[str, Any]], + *, + include_telemetry: bool, + regularization: float, + group_key: str = "decision_id", +) -> list[RidgeModel]: + groups = sorted({str(example[group_key]) for example in examples}) + if len(groups) < 3: + raise ValueError("jackknife ensemble requires at least three groups") + models = [] + for held_out in groups: + training = [ + example for example in examples if str(example[group_key]) != held_out + ] + models.append( + fit_ridge( + training, + include_telemetry=include_telemetry, + regularization=regularization, + ) + ) + return models + + +def ensemble_predict( + models: Sequence[RidgeModel], + example: Mapping[str, Any], + *, + include_telemetry: bool, +) -> dict[str, float]: + if not models: + raise ValueError("ensemble prediction requires models") + source = example["source"] + action = example["action"] + if ( + int(action["target_mns"]) == int(source["mns"]) + and int(action["target_mbbt"]) == int(source["mbbt"]) + ): + return {"mean": 0.0, "std": 0.0, "min": 0.0, "max": 0.0, "distinct_n": 1} + names, values = feature_vector(example, include_telemetry=include_telemetry) + if any(model.feature_names != tuple(names) for model in models): + raise ValueError("ensemble feature schema mismatch") + raw = np.asarray([model.predict(values) for model in models], dtype=np.float64) + clipped = np.clip(raw, -1.0, 1.0) + return { + "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)), + } + + +def select_action( + models: Sequence[RidgeModel], + candidates: Sequence[Mapping[str, Any]], + *, + include_telemetry: bool, + confidence_z: float = 1.0, + minimum_margin: float = 0.02, +) -> dict[str, Any]: + if len(candidates) < 2: + raise ValueError("action selection requires at least two candidates") + rows = [] + for example in candidates: + prediction = ensemble_predict( + models, example, include_telemetry=include_telemetry + ) + rows.append( + { + "action_id": str(example["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 models_to_json(models: Iterable[RidgeModel]) -> list[dict[str, Any]]: + return [model.to_json() for model in models] + + +def models_from_json(payload: Iterable[Mapping[str, Any]]) -> list[RidgeModel]: + return [RidgeModel.from_json(item) for item in payload] diff --git a/runs/active-intervention-v0/test_model.py b/runs/active-intervention-v0/test_model.py new file mode 100644 index 0000000..c49401a --- /dev/null +++ b/runs/active-intervention-v0/test_model.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + + +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 + + +def example(model, decision: str, action: str, pressure: float, target: float): + outcome = { + name: 0.5 for name in model.PREFIX_FEATURES + } + telemetry = {name: 0.0 for name in model.TELEMETRY_FEATURES} + telemetry["queue_waiting_mean"] = pressure + telemetry["batch_size_mean"] = pressure + return { + "decision_id": decision, + "source": { + "mns": 16, + "mbbt": 8192, + "offered_rate_per_gpu": 2.0, + "outcome": outcome, + "telemetry": telemetry, + }, + "action": { + "id": action, + "target_mns": 64 if action == "mns" else 16, + "target_mbbt": 8192 if action == "mns" else 16384, + }, + "target_normalized_goodput": target, + "target_delta_normalized_goodput": target - 0.5, + } + + +def main() -> None: + model = load_model() + examples = [] + for index, pressure in enumerate((0.2, 0.5, 0.8), 1): + examples.extend( + ( + example(model, f"d{index}", "mns", pressure, 0.5 + pressure / 2), + example(model, f"d{index}", "mbbt", pressure, 0.6 - pressure / 4), + ) + ) + fitted = model.fit_ridge( + examples, include_telemetry=True, regularization=1.0 + ) + encoded = fitted.to_json() + restored = model.RidgeModel.from_json(encoded) + names, values = model.feature_vector(examples[-2], include_telemetry=True) + assert tuple(names) == restored.feature_names + assert abs(fitted.predict(values) - restored.predict(values)) < 1e-12 + ensemble = model.fit_jackknife_ensemble( + examples, include_telemetry=True, regularization=1.0 + ) + decision = model.select_action( + ensemble, examples[-2:], include_telemetry=True, minimum_margin=0.0 + ) + assert decision["selected_action"] == "mns" + assert all(-1.0 <= row["prediction"]["mean"] <= 1.0 for row in decision["candidates"]) + noop = example(model, "noop", "noop", 0.8, 0.5) + noop["action"]["target_mbbt"] = 8192 + prediction = model.ensemble_predict( + ensemble, noop, include_telemetry=True + ) + assert prediction == { + "mean": 0.0, + "std": 0.0, + "min": 0.0, + "max": 0.0, + "distinct_n": 1, + } + print("active intervention model: PASS") + + +if __name__ == "__main__": + main() diff --git a/runs/active-intervention-v0/test_pipeline.py b/runs/active-intervention-v0/test_pipeline.py new file mode 100644 index 0000000..ee7805c --- /dev/null +++ b/runs/active-intervention-v0/test_pipeline.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import importlib.util +import json +import sys +import tempfile +from pathlib import Path + + +HERE = Path(__file__).resolve().parent + + +def load(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 + + +def write_json(path: Path, payload) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload) + "\n", encoding="utf-8") + + +def write_jsonl(path: Path, rows) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" + ) + + +def engine_record(index: int, timestamp_ns: int) -> dict: + alternate = index % 2 + return { + "step_index": index, + "submit_mono_ns": timestamp_ns, + "model_executed": True, + "scheduled_requests": 1 + alternate, + "decode_batch_size": alternate, + "prefill_tokens": 8 + alternate, + "decode_tokens": alternate, + "preemptions": 0, + "queues": {"waiting": alternate, "running": 1 + alternate}, + "kv": {"usage": 0.1 + 0.01 * alternate}, + "cudagraph": { + "runtime_mode": "FULL" if alternate else "NONE", + "bucket_tokens": 16, + "padding_tokens": alternate, + }, + "dropped_records_before": 0, + } + + +def main() -> None: + extractor = load("active_intervention_extract_test", HERE / "extract_training.py") + trainer = load("active_intervention_train_test", HERE / "train_policy.py") + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + run_root = root / "runs" + configs = [ + {"id": "a_base", "mns": 16, "mbbt": 8192}, + {"id": "a_mns", "mns": 64, "mbbt": 8192}, + {"id": "a_mbbt", "mns": 16, "mbbt": 16384}, + {"id": "b_base", "mns": 64, "mbbt": 2048}, + {"id": "b_mns", "mns": 128, "mbbt": 2048}, + {"id": "b_mbbt", "mns": 64, "mbbt": 8192}, + ] + manifest = { + "engine": {"duration_s": 300.0, "tp": 4}, + "configs": configs, + "repetitions": { + str(rep): {"selection": {"offered_req_s_per_gpu": 0.01}} + for rep in (1, 2, 3) + }, + "regimes": { + "A": { + "source": "a_base", + "actions": {"mns": "a_mns", "mbbt": "a_mbbt"}, + }, + "B": { + "source": "b_base", + "actions": {"mns": "b_mns", "mbbt": "b_mbbt"}, + }, + }, + } + manifest_path = root / "manifest.json" + write_json(manifest_path, manifest) + + streams = [] + source_starts: dict[tuple[str, int], int] = {} + for source_index, source_id in enumerate(("a_base", "b_base")): + rows = [] + index = 0 + for repetition in (1, 2, 3): + start_ns = int((source_index * 2000 + repetition * 400) * 1e9) + source_starts[(source_id, repetition)] = start_ns + for second in (1, 30, 76, 105, 151, 180, 226, 255): + rows.append(engine_record(index, start_ns + int(second * 1e9))) + index += 1 + stream_path = root / f"{source_id}-stream.jsonl" + write_jsonl(stream_path, rows) + streams.append( + { + "config_id": source_id, + "stream": str(stream_path), + "stream_sha256": extractor.sha256_file(stream_path), + } + ) + + request_rows = [ + { + "request_id": f"r{index}", + "arrival_s": arrival, + "completed_elapsed_s": arrival + 10, + "slo_pass": index != 3, + "ttft_ms": 1000 + index * 100, + "tpot_ms": 20 + index, + "raw_input_tokens": 1000 + index * 100, + } + for index, arrival in enumerate((5.0, 80.0, 155.0, 230.0), 1) + ] + for source_id in ("a_base", "b_base"): + for repetition in (1, 2, 3): + write_jsonl( + run_root + / "sessions" + / source_id + / f"rep{repetition}" + / "requests.jsonl", + request_rows, + ) + + goodput = { + "a_base": 0.020, + "a_mns": 0.036, + "a_mbbt": 0.028, + "b_base": 0.032, + "b_mns": 0.030, + "b_mbbt": 0.038, + } + runs = [] + for config in configs: + for repetition in (1, 2, 3): + item = { + "config_id": config["id"], + "repetition": repetition, + "outcome": { + "slo_goodput_req_s": goodput[config["id"]] + + repetition * 0.0001 + }, + } + if config["id"] in ("a_base", "b_base"): + start_ns = source_starts[(config["id"], repetition)] + item["state"] = { + "interval": { + "start_ns": start_ns, + "end_ns": start_ns + int(300 * 1e9), + } + } + runs.append(item) + audit = { + "schema": "action-aware-constraint-pilot-audit-v0", + "sanity": {"red_flags": []}, + "streams": streams, + "runs": runs, + } + audit_path = root / "audit.json" + write_json(audit_path, audit) + + dataset = extractor.build_dataset( + audit_path=audit_path, manifest_path=manifest_path, run_root=run_root + ) + assert dataset["status"] == "VALID" + assert len(dataset["examples"]) == 72 + assert not dataset["sanity"]["red_flags"] + assert all( + "exclusive" not in feature + for example in dataset["examples"] + for feature in example["source"]["telemetry"] + ) + dataset_path = root / "dataset.json" + write_json(dataset_path, dataset) + policy = trainer.build_policy(dataset_path) + assert policy["status"] in { + "RETROSPECTIVE_INCREMENTAL_SIGNAL", + "NO_RETROSPECTIVE_INCREMENTAL_SIGNAL", + } + assert not policy["sanity"]["red_flags"] + print("active intervention pipeline: PASS") + + +if __name__ == "__main__": + main() diff --git a/runs/active-intervention-v0/train_policy.py b/runs/active-intervention-v0/train_policy.py new file mode 100644 index 0000000..8bbae02 --- /dev/null +++ b/runs/active-intervention-v0/train_policy.py @@ -0,0 +1,318 @@ +#!/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 + + +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) + } + decision_rows.append( + { + "holdout": held_out, + "decision_id": decision_id, + "selected_action": selected["action_id"], + "best_actions": sorted(best_actions), + "correct": selected["action_id"] in best_actions, + "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, + "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 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 + ) + 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, + } + selected_phase = incremental_candidates[0] if incremental_candidates else phases[-1] + status = ( + "RETROSPECTIVE_INCREMENTAL_SIGNAL" + if incremental_candidates + else "NO_RETROSPECTIVE_INCREMENTAL_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, + }, + "measurement_policy": { + "selected_phase": selected_phase, + "selected_cutoff_s": phase_results[selected_phase]["cutoff_s"], + "selection_reason": ( + "earliest phase passing the frozen incremental gate" + if incremental_candidates + else "no incremental phase; retain full measurement for exploratory held-out test" + ), + }, + "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()