Files
aituner/runs/fidelity-headroom/analyze_prefixes.py

630 lines
24 KiB
Python

#!/usr/bin/env python3
"""Retrospective, leakage-bounded audit of short real-probe prefixes.
The outcome-only and instrumentation-aware models receive the same trial
prefix. The latter differs only by Layer-1 engine state. Existing Phase-6
request artifacts predate exact completion timestamps, so their completion
time is reconstructed from arrival + TTFT + token intervals and is explicitly
marked approximate. New artifacts use ``completed_elapsed_s`` directly.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable
import numpy as np
from analyze_existing import (
DEFAULT_REGULARIZATION,
REGULARIZATION_SENSITIVITY,
_classification_metrics,
_fit_logistic,
_group_bootstrap_delta,
_mcnemar_exact_p,
_sigmoid,
)
SCHEMA = "fidelity-prefix-v1"
DEFAULT_CUTOFFS = (5.0, 10.0, 15.0, 20.0)
POLICY_THRESHOLDS = (0.8, 0.9, 0.95)
OUTCOME_FEATURES = (
"log_offered_rate_per_gpu",
"log2_tp",
"log2_max_num_seqs",
"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",
)
INSTRUMENTATION_FEATURES = (
"model_steps_per_second",
"waiting_mean",
"waiting_max",
"waiting_nonzero_share",
"running_mean",
"running_max",
"decode_batch_mean",
"decode_batch_max",
"decode_batch_cv",
"kv_usage_mean",
"kv_usage_max",
"kv_usage_end_minus_start",
"graph_none_share",
"graph_full_share",
"padding_fraction",
"prefill_token_fraction",
"preemptions",
)
@dataclass(frozen=True)
class PrefixExample:
cell: str
anchor: float
cutoff_s: float
tp: int
full_elapsed_s: float
feasible: int
primary_feasible: int
outcome: tuple[float, ...]
instrumentation: tuple[float, ...]
completion_time_source: str
@property
def remaining_h20_hours(self) -> float:
return self.tp * max(0.0, self.full_elapsed_s - self.cutoff_s) / 3600.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 numeric(values: Iterable[float | int]) -> dict[str, Any]:
array = [float(value) for value in values]
return {
"n": len(array),
"min": min(array) if array else None,
"max": max(array) if array else None,
"distinct_n": len(set(array)),
}
def _cv(values: list[float]) -> float:
if not values:
return 0.0
array = np.asarray(values, dtype=np.float64)
mean = float(array.mean())
return float(array.std(ddof=0) / mean) if mean else 0.0
def completion_elapsed_s(request: dict[str, Any]) -> tuple[float | None, str]:
exact = request.get("completed_elapsed_s")
if exact is not None:
value = float(exact)
if value < 0 or not math.isfinite(value):
raise ValueError(f"invalid completed_elapsed_s={exact}")
return value, "exact_monotonic"
if not request.get("success"):
return None, "unobserved_failure"
required = (
request.get("arrival_s"),
request.get("ttft_ms"),
request.get("tpot_ms"),
request.get("completion_tokens"),
)
if any(value is None for value in required):
return None, "unobserved_failure"
arrival_s, ttft_ms, tpot_ms, completion_tokens = required
value = float(arrival_s) + (
float(ttft_ms) + max(int(completion_tokens) - 1, 0) * float(tpot_ms)
) / 1000.0
if value < 0 or not math.isfinite(value):
raise ValueError(f"invalid reconstructed completion time={value}")
return value, "reconstructed_from_latency"
def _load_jsonl(path: Path, *, require_key: str | None = None) -> list[dict[str, Any]]:
records = []
with path.open(encoding="utf-8") as source:
for line in source:
item = json.loads(line)
if require_key is None or require_key in item:
records.append(item)
return records
def _anchor_directory(cell_root: Path, anchor: float) -> Path:
matches = []
for result_path in cell_root.glob("anchor-*/result.json"):
payload = json.loads(result_path.read_text(encoding="utf-8"))
if math.isclose(float(payload["anchor"]), anchor, rel_tol=0.0, abs_tol=1e-15):
matches.append(result_path.parent)
if len(matches) != 1:
raise ValueError(f"expected one primary directory for anchor {anchor}: {matches}")
return matches[0]
def _prefix_features(
*,
primary: dict[str, Any],
tp: int,
max_num_seqs: int,
requests: list[dict[str, Any]],
records: list[dict[str, Any]],
cutoff_s: float,
) -> tuple[tuple[float, ...], tuple[float, ...], str]:
admitted = [request for request in requests if float(request["arrival_s"]) <= cutoff_s]
completed = []
sources = set()
for request in requests:
completed_s, source = completion_elapsed_s(request)
if completed_s is None or completed_s > cutoff_s:
continue
completed.append(request)
sources.add(source)
if not admitted or not records:
raise ValueError("prefix has no admitted requests or Layer-1 records")
if any(request not in admitted for request in completed):
raise ValueError("completed request was not admitted inside prefix")
total = len(requests)
passed = sum(bool(request["slo_pass"]) for request in completed)
ttft = [float(request["ttft_ms"]) for request in completed if request["ttft_ms"] is not None]
tpot = [float(request["tpot_ms"]) for request in completed if request["tpot_ms"] is not None]
offered_rate = float(primary["selection"]["offered_req_s_per_gpu"])
if offered_rate <= 0 or total <= 0:
raise ValueError("offered rate and selected request count must be positive")
outcome = (
math.log(offered_rate),
math.log2(float(tp)),
math.log2(float(max_num_seqs)),
len(admitted) / total,
len(completed) / len(admitted),
passed / max(1, len(completed)),
(len(completed) - passed) / total,
(len(admitted) - len(completed)) / len(admitted),
max(ttft, default=0.0) / 6000.0,
float(np.mean(ttft)) / 6000.0 if ttft else 0.0,
max(tpot, default=0.0) / 50.0,
float(np.mean(tpot)) / 50.0 if tpot else 0.0,
float(np.mean([float(request["raw_input_tokens"]) for request in admitted])) / 8192.0,
)
waiting = [float(record["queues"]["waiting"]) for record in records]
running = [float(record["queues"]["running"]) for record in records]
decode_batch = [float(record["decode_batch_size"]) for record in records]
kv_usage = [float(record["kv"]["usage"]) for record in records]
graph_modes = [str(record["cudagraph"]["runtime_mode"]) for record in records]
bucket_tokens = sum(int(record["cudagraph"]["bucket_tokens"]) for record in records)
padding_tokens = sum(int(record["cudagraph"]["padding_tokens"]) for record in records)
prefill_tokens = sum(int(record["prefill_tokens"]) for record in records)
decode_tokens = sum(int(record["decode_tokens"]) for record in records)
instrumentation = (
len(records) / cutoff_s,
float(np.mean(waiting)),
max(waiting),
sum(value > 0 for value in waiting) / len(waiting),
float(np.mean(running)),
max(running),
float(np.mean(decode_batch)),
max(decode_batch),
_cv(decode_batch),
float(np.mean(kv_usage)),
max(kv_usage),
kv_usage[-1] - kv_usage[0],
graph_modes.count("NONE") / len(graph_modes),
graph_modes.count("FULL") / len(graph_modes),
padding_tokens / max(1, bucket_tokens),
prefill_tokens / max(1, prefill_tokens + decode_tokens),
float(sum(int(record["preemptions"]) for record in records)),
)
completion_source = "+".join(sorted(sources)) if sources else "none_completed"
return outcome, instrumentation, completion_source
def build_examples(
phase6: dict[str, Any],
raw_root: Path,
cutoff_s: float,
) -> list[PrefixExample]:
examples = []
for cell, cell_result in sorted(phase6["cells"].items()):
cell_root = raw_root / cell
stream_path = next((cell_root / "opprof").glob("*.jsonl"))
stream = _load_jsonl(stream_path, require_key="submit_mono_ns")
for anchor in cell_result["anchors"]:
primary = anchor["primary"]
full_elapsed_s = float(primary["interval"]["elapsed_s"])
if full_elapsed_s + 1e-9 < cutoff_s:
continue
anchor_value = float(primary["anchor"])
anchor_root = _anchor_directory(cell_root, anchor_value)
requests = _load_jsonl(anchor_root / "requests.jsonl")
start_ns = int(primary["interval"]["start_mono_ns"])
end_ns = start_ns + int(cutoff_s * 1e9)
records = [
record
for record in stream
if record.get("model_executed")
and start_ns <= int(record["submit_mono_ns"]) <= end_ns
]
outcome, instrumentation, source = _prefix_features(
primary=primary,
tp=int(cell_result["tp"]),
max_num_seqs=int(cell_result["mns"]),
requests=requests,
records=records,
cutoff_s=cutoff_s,
)
examples.append(
PrefixExample(
cell=cell,
anchor=anchor_value,
cutoff_s=cutoff_s,
tp=int(cell_result["tp"]),
full_elapsed_s=full_elapsed_s,
feasible=int(bool(anchor["accepted_feasible"])),
primary_feasible=int(bool(primary["feasible"])),
outcome=outcome,
instrumentation=instrumentation,
completion_time_source=source,
)
)
return examples
def grouped_predictions(
examples: list[PrefixExample],
*,
instrumentation_aware: bool,
regularization: float,
) -> tuple[np.ndarray, np.ndarray, list[str]]:
probabilities = []
labels = []
groups = []
for held_out in sorted({example.cell for example in examples}):
train = [example for example in examples if example.cell != held_out]
test = [example for example in examples if example.cell == held_out]
def row(example: PrefixExample) -> np.ndarray:
values = example.outcome
if instrumentation_aware:
values += example.instrumentation
return np.asarray((1.0, *values), dtype=np.float64)
x_train = np.stack([row(example) for example in train])
x_test = np.stack([row(example) for example in test])
y_train = np.asarray([example.feasible for example in train], dtype=np.float64)
if len(set(y_train.tolist())) != 2:
raise ValueError(f"training fold for {held_out} has a single label")
mean = x_train[:, 1:].mean(axis=0)
standard_deviation = x_train[:, 1:].std(axis=0)
standard_deviation[standard_deviation < 1e-8] = 1.0
x_train[:, 1:] = (x_train[:, 1:] - mean) / standard_deviation
x_test[:, 1:] = (x_test[:, 1:] - mean) / standard_deviation
weights = _fit_logistic(x_train, y_train, regularization)
probabilities.extend(_sigmoid(x_test @ weights).tolist())
labels.extend(example.feasible for example in test)
groups.extend(held_out for _ in test)
return (
np.asarray(labels, dtype=np.int64),
np.asarray(probabilities, dtype=np.float64),
groups,
)
def fit_frozen_model(
examples: list[PrefixExample],
*,
instrumentation_aware: bool,
regularization: float,
) -> dict[str, Any]:
def row(example: PrefixExample) -> np.ndarray:
values = example.outcome
if instrumentation_aware:
values += example.instrumentation
return np.asarray((1.0, *values), dtype=np.float64)
matrix = np.stack([row(example) for example in examples])
labels = np.asarray([example.feasible for example in examples], dtype=np.float64)
if len(set(labels.tolist())) != 2:
raise ValueError("frozen model requires both feasibility labels")
mean = matrix[:, 1:].mean(axis=0)
standard_deviation = matrix[:, 1:].std(axis=0)
standard_deviation[standard_deviation < 1e-8] = 1.0
standardized = matrix.copy()
standardized[:, 1:] = (standardized[:, 1:] - mean) / standard_deviation
weights = _fit_logistic(standardized, labels, regularization)
probabilities = _sigmoid(standardized @ weights)
names = list(OUTCOME_FEATURES)
if instrumentation_aware:
names.extend(INSTRUMENTATION_FEATURES)
return {
"instrumentation_aware": instrumentation_aware,
"regularization": regularization,
"feature_names": names,
"feature_mean": mean.tolist(),
"feature_standard_deviation": standard_deviation.tolist(),
"weights_with_intercept_first": weights.tolist(),
"training_classification": _classification_metrics(labels, probabilities),
}
def predict_frozen_model(
model: dict[str, Any],
examples: list[PrefixExample],
) -> np.ndarray:
instrumentation_aware = bool(model["instrumentation_aware"])
rows = []
for example in examples:
values = example.outcome
if instrumentation_aware:
values += example.instrumentation
rows.append((1.0, *values))
matrix = np.asarray(rows, dtype=np.float64)
mean = np.asarray(model["feature_mean"], dtype=np.float64)
standard_deviation = np.asarray(
model["feature_standard_deviation"], dtype=np.float64
)
weights = np.asarray(model["weights_with_intercept_first"], dtype=np.float64)
if matrix.shape[1] != len(weights) or matrix.shape[1] - 1 != len(mean):
raise ValueError("frozen model feature dimensions do not match examples")
matrix[:, 1:] = (matrix[:, 1:] - mean) / standard_deviation
return _sigmoid(matrix @ weights)
def policy_metrics(
examples: list[PrefixExample],
labels: np.ndarray,
probabilities: np.ndarray,
threshold: float,
) -> dict[str, Any]:
accept = probabilities >= threshold
reject = probabilities <= 1.0 - threshold
decide = accept | reject
prediction = accept.astype(np.int64)
correct = prediction == labels
remaining = np.asarray(
[example.remaining_h20_hours for example in examples], dtype=np.float64
)
full_cost = sum(example.tp * example.full_elapsed_s / 3600.0 for example in examples)
saved = float(np.sum(remaining[decide]))
correct_saved = float(np.sum(remaining[decide & correct]))
invalid_saved = float(np.sum(remaining[decide & ~correct]))
def describe(mask: np.ndarray) -> list[dict[str, Any]]:
return [
{
"cell": example.cell,
"anchor": example.anchor,
"label_feasible": bool(label),
"probability_feasible": float(probability),
"remaining_h20_hours": example.remaining_h20_hours,
}
for example, label, probability, selected in zip(
examples, labels, probabilities, mask
)
if selected
]
return {
"threshold": threshold,
"early_accept": int(np.sum(accept)),
"early_reject": int(np.sum(reject)),
"abstain_continue_full": int(np.sum(~decide)),
"false_accept": int(np.sum(accept & (labels == 0))),
"false_reject": int(np.sum(reject & (labels == 1))),
"false_accept_examples": describe(accept & (labels == 0)),
"false_reject_examples": describe(reject & (labels == 1)),
"decision_coverage": float(np.mean(decide)),
"full_trial_h20_hours": float(full_cost),
"remaining_h20_hours_at_cutoff": float(np.sum(remaining)),
"saved_h20_hours_if_decisions_used": saved,
"correctly_saved_h20_hours": correct_saved,
"invalidly_saved_h20_hours": invalid_saved,
"valid_zero_error_policy": bool(np.all(correct[decide])),
"valid_cost_reduction_fraction": (
correct_saved / full_cost if invalid_saved == 0.0 and full_cost else None
),
}
def analyze_cutoff(examples: list[PrefixExample]) -> dict[str, Any]:
sensitivity = {}
headline = None
for regularization in REGULARIZATION_SENSITIVITY:
labels, outcome_probability, groups = grouped_predictions(
examples,
instrumentation_aware=False,
regularization=regularization,
)
instrument_labels, instrument_probability, instrument_groups = grouped_predictions(
examples,
instrumentation_aware=True,
regularization=regularization,
)
if not np.array_equal(labels, instrument_labels) or groups != instrument_groups:
raise AssertionError("paired folds or labels differ")
if groups != [example.cell for example in examples]:
raise AssertionError("prediction order differs from example order")
outcome_correct = (outcome_probability >= 0.5) == labels
instrument_correct = (instrument_probability >= 0.5) == labels
result = {
"outcome_only": {
"classification": _classification_metrics(labels, outcome_probability),
"policies": [
policy_metrics(examples, labels, outcome_probability, threshold)
for threshold in POLICY_THRESHOLDS
],
},
"instrumentation_aware": {
"classification": _classification_metrics(labels, instrument_probability),
"policies": [
policy_metrics(examples, labels, instrument_probability, threshold)
for threshold in POLICY_THRESHOLDS
],
},
"paired_correctness": {
"both_correct": int(np.sum(outcome_correct & instrument_correct)),
"outcome_only_correct": int(np.sum(outcome_correct & ~instrument_correct)),
"instrumentation_only_correct": int(np.sum(~outcome_correct & instrument_correct)),
"both_wrong": int(np.sum(~outcome_correct & ~instrument_correct)),
},
"bootstrap": _group_bootstrap_delta(
labels,
outcome_probability,
instrument_probability,
groups,
),
}
paired = result["paired_correctness"]
paired["mcnemar_exact_two_sided_p"] = _mcnemar_exact_p(
paired["outcome_only_correct"], paired["instrumentation_only_correct"]
)
sensitivity[str(regularization)] = result
if regularization == DEFAULT_REGULARIZATION:
headline = result
assert headline is not None
labels = [example.feasible for example in examples]
return {
"examples": len(examples),
"cells": len({example.cell for example in examples}),
"label_sanity": {
**numeric(labels),
"positive": sum(labels),
"negative": len(labels) - sum(labels),
"primary_adjudicated_disagreements": sum(
example.feasible != example.primary_feasible for example in examples
),
},
"completion_time_sources": {
source: sum(example.completion_time_source == source for example in examples)
for source in sorted({example.completion_time_source for example in examples})
},
"headline_regularization": DEFAULT_REGULARIZATION,
"headline": headline,
"regularization_sensitivity": sensitivity,
"remaining_h20_hours": numeric(
example.remaining_h20_hours for example in examples
),
}
def analyze(
phase6_path: Path,
raw_root: Path,
cutoffs: tuple[float, ...],
) -> dict[str, Any]:
phase6 = json.loads(phase6_path.read_text(encoding="utf-8"))
by_cutoff = {}
red_flags = []
for cutoff in cutoffs:
examples = build_examples(phase6, raw_root, cutoff)
if len({example.feasible for example in examples}) != 2:
red_flags.append(f"single_label_at_{cutoff:g}s")
continue
by_cutoff[f"{cutoff:g}"] = analyze_cutoff(examples)
if len({example.cell for example in examples}) != 12:
red_flags.append(f"incomplete_cells_at_{cutoff:g}s")
if not all(
math.isfinite(value)
for example in examples
for value in (*example.outcome, *example.instrumentation)
):
red_flags.append(f"nonfinite_features_at_{cutoff:g}s")
headline_deltas = {
cutoff: {
"accuracy": (
result["headline"]["instrumentation_aware"]["classification"]["accuracy"]
- result["headline"]["outcome_only"]["classification"]["accuracy"]
),
"brier": (
result["headline"]["instrumentation_aware"]["classification"]["brier"]
- result["headline"]["outcome_only"]["classification"]["brier"]
),
}
for cutoff, result in by_cutoff.items()
}
return {
"schema": SCHEMA,
"status": "PASS" if not red_flags else "STOP",
"scope": (
"retrospective single-workload prefix diagnostic; model selection, "
"threshold choice, and contribution claims require held-out prospective tasks"
),
"estimand": (
"2-of-3 adjudicated anchor feasibility from the first primary trial's "
"identical short real prefix"
),
"split": "leave-one-configuration-cell-out",
"model": "same L2 logistic model and folds; instrumentation model appends Layer-1 features",
"outcome_features": list(OUTCOME_FEATURES),
"instrumentation_features": list(INSTRUMENTATION_FEATURES),
"provenance": {
"phase6_metrics": str(phase6_path.resolve()),
"phase6_metrics_sha256": sha256_file(phase6_path),
"raw_root": str(raw_root.resolve()),
},
"cutoffs_s": list(cutoffs),
"cutoffs": by_cutoff,
"headline_incremental_deltas": headline_deltas,
"decision": {
"contribution_established": False,
"reason": (
"This dataset contains one workload and reconstructed rather than exact request "
"completion times. Three TP4 primary trials also disagree with their 2-of-3 "
"labels. It can reject a missing-signal premise but cannot establish "
"generalization or a paper-facing cost reduction."
),
},
"sanity": {
"red_flags": red_flags,
"cutoff_count": len(by_cutoff),
"invariants": {
"cutoffs_positive": all(cutoff > 0 for cutoff in cutoffs),
"paired_same_model_family": True,
"probabilities_checked_in_unit_interval": True,
"full_trial_label_not_used_as_feature": True,
"records_strictly_prefix_sliced": True,
},
},
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--phase6-metrics", type=Path, required=True)
parser.add_argument("--raw-root", type=Path, required=True)
parser.add_argument("--cutoffs", type=float, nargs="+", default=DEFAULT_CUTOFFS)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
result = analyze(args.phase6_metrics, args.raw_root, tuple(args.cutoffs))
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
print(json.dumps({"status": result["status"], "output": str(args.output)}, sort_keys=True))
if __name__ == "__main__":
main()