Prepare Frontier code trace fidelity campaign
This commit is contained in:
361
runs/frontier-code-trace-v0/audit_code_trace.py
Normal file
361
runs/frontier-code-trace-v0/audit_code_trace.py
Normal file
@@ -0,0 +1,361 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Audit long code traces before choosing a replay window and max model length."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Sequence
|
||||
|
||||
|
||||
BLOCK_SIZE_CANDIDATES = (16, 32, 64, 128, 256, 512, 1024)
|
||||
MAX_MODEL_LEN_CANDIDATES = (40960, 65536, 98304, 131072, 262144)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--trace-root", type=Path)
|
||||
parser.add_argument("--source", type=Path, action="append")
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--min-minutes", type=int, default=60)
|
||||
parser.add_argument("--max-minutes", type=int, default=75)
|
||||
parser.add_argument("--bin-seconds", type=int, default=60)
|
||||
parser.add_argument("--max-acceptable-gap-s", type=float, default=5.0)
|
||||
parser.add_argument("--model-position-limit", type=int, default=262144)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def percentile(values: Sequence[int | float], fraction: float) -> float | None:
|
||||
if not values:
|
||||
return None
|
||||
ordered = sorted(float(value) for value in values)
|
||||
position = (len(ordered) - 1) * fraction
|
||||
lower = math.floor(position)
|
||||
upper = math.ceil(position)
|
||||
if lower == upper:
|
||||
return ordered[lower]
|
||||
return ordered[lower] * (upper - position) + ordered[upper] * (position - lower)
|
||||
|
||||
|
||||
def distribution(values: Sequence[int | float]) -> dict[str, int | float | None]:
|
||||
return {
|
||||
"count": len(values),
|
||||
"min": min(values) if values else None,
|
||||
"p50": percentile(values, 0.50),
|
||||
"p90": percentile(values, 0.90),
|
||||
"p95": percentile(values, 0.95),
|
||||
"p99": percentile(values, 0.99),
|
||||
"max": max(values) if values else None,
|
||||
"mean": statistics.fmean(values) if values else None,
|
||||
}
|
||||
|
||||
|
||||
def parse_hash_ids(value: Any) -> list[Any]:
|
||||
if isinstance(value, list):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
stripped = value.strip()
|
||||
if not stripped:
|
||||
return []
|
||||
if stripped.startswith("["):
|
||||
decoded = json.loads(stripped)
|
||||
if not isinstance(decoded, list):
|
||||
raise ValueError("hash_ids JSON must decode to a list")
|
||||
return decoded
|
||||
delimiter = "|" if "|" in stripped else ","
|
||||
return [part for part in stripped.split(delimiter) if part.strip()]
|
||||
if value is None:
|
||||
return []
|
||||
return [value]
|
||||
|
||||
|
||||
def iter_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]:
|
||||
with path.open() as stream:
|
||||
for line_number, line in enumerate(stream, 1):
|
||||
if not line.strip():
|
||||
continue
|
||||
row = json.loads(line)
|
||||
if not isinstance(row, dict):
|
||||
raise ValueError(f"{path}:{line_number}: row must be an object")
|
||||
yield line_number, row
|
||||
|
||||
|
||||
def choose_window(
|
||||
*,
|
||||
counts: Sequence[int],
|
||||
max_gaps: Sequence[float],
|
||||
first_timestamp: float,
|
||||
min_minutes: int,
|
||||
max_minutes: int,
|
||||
bin_seconds: int,
|
||||
max_acceptable_gap_s: float,
|
||||
) -> dict[str, Any] | None:
|
||||
candidates = []
|
||||
for minutes in range(max_minutes, min_minutes - 1, -1):
|
||||
bins = math.ceil(minutes * 60 / bin_seconds)
|
||||
for start_bin in range(0, len(counts) - bins + 1):
|
||||
selected = counts[start_bin : start_bin + bins]
|
||||
mean = statistics.fmean(selected)
|
||||
cv = statistics.pstdev(selected) / mean if mean else math.inf
|
||||
max_gap = max(max_gaps[start_bin : start_bin + bins], default=0.0)
|
||||
candidates.append(
|
||||
{
|
||||
"_score": (
|
||||
max_gap > max_acceptable_gap_s,
|
||||
cv,
|
||||
max_gap,
|
||||
-minutes,
|
||||
start_bin,
|
||||
),
|
||||
"start_bin": start_bin,
|
||||
"minutes": minutes,
|
||||
"count_mean_per_bin": mean,
|
||||
"count_cv": cv,
|
||||
"count_min_per_bin": min(selected),
|
||||
"count_max_per_bin": max(selected),
|
||||
"max_gap_s": max_gap,
|
||||
}
|
||||
)
|
||||
if not candidates:
|
||||
return None
|
||||
chosen = min(candidates, key=lambda item: item["_score"])
|
||||
chosen.pop("_score")
|
||||
chosen["start_timestamp"] = first_timestamp + chosen["start_bin"] * bin_seconds
|
||||
chosen["end_timestamp"] = chosen["start_timestamp"] + chosen["minutes"] * 60
|
||||
return chosen
|
||||
|
||||
|
||||
def scan_source(path: Path, args: argparse.Namespace) -> dict[str, Any]:
|
||||
rows = 0
|
||||
first_timestamp = None
|
||||
last_timestamp = None
|
||||
previous_timestamp = None
|
||||
counts: Counter[int] = Counter()
|
||||
max_gaps: dict[int, float] = {}
|
||||
input_lengths: list[int] = []
|
||||
output_lengths: list[int] = []
|
||||
total_lengths: list[int] = []
|
||||
hash_rows = 0
|
||||
hash_matches = Counter()
|
||||
prompt_rows = 0
|
||||
sampling_rows = 0
|
||||
schema_keys: Counter[str] = Counter()
|
||||
for line_number, row in iter_jsonl(path):
|
||||
missing = [
|
||||
key
|
||||
for key in ("timestamp", "input_length", "output_length")
|
||||
if key not in row
|
||||
]
|
||||
if missing:
|
||||
raise ValueError(f"{path}:{line_number}: missing required fields {missing}")
|
||||
timestamp = float(row["timestamp"])
|
||||
if first_timestamp is None:
|
||||
first_timestamp = timestamp
|
||||
if previous_timestamp is not None and timestamp < previous_timestamp:
|
||||
raise ValueError(
|
||||
f"{path}:{line_number}: timestamp {timestamp} < {previous_timestamp}"
|
||||
)
|
||||
bin_index = math.floor((timestamp - first_timestamp) / args.bin_seconds)
|
||||
counts[bin_index] += 1
|
||||
if previous_timestamp is not None:
|
||||
previous_bin = math.floor(
|
||||
(previous_timestamp - first_timestamp) / args.bin_seconds
|
||||
)
|
||||
max_gaps[previous_bin] = max(
|
||||
max_gaps.get(previous_bin, 0.0),
|
||||
timestamp - previous_timestamp,
|
||||
)
|
||||
input_tokens = int(row["input_length"])
|
||||
output_tokens = max(1, int(row["output_length"]))
|
||||
if input_tokens <= 0:
|
||||
raise ValueError(f"{path}:{line_number}: input_length must be positive")
|
||||
input_lengths.append(input_tokens)
|
||||
output_lengths.append(output_tokens)
|
||||
total_lengths.append(input_tokens + output_tokens)
|
||||
hashes = parse_hash_ids(row.get("hash_ids"))
|
||||
if hashes:
|
||||
hash_rows += 1
|
||||
for block_size in BLOCK_SIZE_CANDIDATES:
|
||||
if len(hashes) == math.ceil(input_tokens / block_size):
|
||||
hash_matches[block_size] += 1
|
||||
prompt_rows += int(
|
||||
isinstance(row.get("prompt"), (str, list)) and bool(row.get("prompt"))
|
||||
)
|
||||
sampling_rows += int("sampling_u" in row)
|
||||
schema_keys.update(row.keys())
|
||||
rows += 1
|
||||
previous_timestamp = timestamp
|
||||
last_timestamp = timestamp
|
||||
if not rows or first_timestamp is None or last_timestamp is None:
|
||||
raise ValueError(f"{path}: empty trace")
|
||||
total_bins = math.floor((last_timestamp - first_timestamp) / args.bin_seconds) + 1
|
||||
chosen = choose_window(
|
||||
counts=[counts[index] for index in range(total_bins)],
|
||||
max_gaps=[max_gaps.get(index, 0.0) for index in range(total_bins)],
|
||||
first_timestamp=first_timestamp,
|
||||
min_minutes=args.min_minutes,
|
||||
max_minutes=args.max_minutes,
|
||||
bin_seconds=args.bin_seconds,
|
||||
max_acceptable_gap_s=args.max_acceptable_gap_s,
|
||||
)
|
||||
return {
|
||||
"source": str(path.resolve()),
|
||||
"rows": rows,
|
||||
"first_timestamp": first_timestamp,
|
||||
"last_timestamp": last_timestamp,
|
||||
"span_s": last_timestamp - first_timestamp,
|
||||
"request_rate_per_s": rows / max(last_timestamp - first_timestamp, 1.0),
|
||||
"input_length": distribution(input_lengths),
|
||||
"output_length": distribution(output_lengths),
|
||||
"total_length": distribution(total_lengths),
|
||||
"over_max_model_len": {
|
||||
str(limit): {
|
||||
"requests": sum(value > limit for value in total_lengths),
|
||||
"fraction": sum(value > limit for value in total_lengths) / rows,
|
||||
}
|
||||
for limit in MAX_MODEL_LEN_CANDIDATES
|
||||
},
|
||||
"hash_contract": {
|
||||
"rows_with_hash_ids": hash_rows,
|
||||
"candidate_exact_match_rows": {
|
||||
str(size): hash_matches[size] for size in BLOCK_SIZE_CANDIDATES
|
||||
},
|
||||
"exact_source_block_size": next(
|
||||
(
|
||||
size
|
||||
for size in BLOCK_SIZE_CANDIDATES
|
||||
if hash_rows and hash_matches[size] == hash_rows
|
||||
),
|
||||
None,
|
||||
),
|
||||
},
|
||||
"prompt_rows": prompt_rows,
|
||||
"sampling_u_rows": sampling_rows,
|
||||
"schema_field_counts": dict(sorted(schema_keys.items())),
|
||||
"stable_window": chosen,
|
||||
}
|
||||
|
||||
|
||||
def scan_window(source: Path, window: dict[str, Any]) -> dict[str, Any]:
|
||||
start = float(window["start_timestamp"])
|
||||
end = float(window["end_timestamp"])
|
||||
inputs: list[int] = []
|
||||
outputs: list[int] = []
|
||||
totals: list[int] = []
|
||||
for _, row in iter_jsonl(source):
|
||||
timestamp = float(row["timestamp"])
|
||||
if timestamp < start:
|
||||
continue
|
||||
if timestamp >= end:
|
||||
break
|
||||
input_tokens = int(row["input_length"])
|
||||
output_tokens = max(1, int(row["output_length"]))
|
||||
inputs.append(input_tokens)
|
||||
outputs.append(output_tokens)
|
||||
totals.append(input_tokens + output_tokens)
|
||||
return {
|
||||
"requests": len(totals),
|
||||
"input_length": distribution(inputs),
|
||||
"output_length": distribution(outputs),
|
||||
"total_length": distribution(totals),
|
||||
"max_model_len_coverage": {
|
||||
str(limit): {
|
||||
"covered_requests": sum(value <= limit for value in totals),
|
||||
"excluded_requests": sum(value > limit for value in totals),
|
||||
"coverage": sum(value <= limit for value in totals) / len(totals),
|
||||
}
|
||||
for limit in MAX_MODEL_LEN_CANDIDATES
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def resolve_sources(args: argparse.Namespace) -> list[Path]:
|
||||
if args.source:
|
||||
return [path.resolve() for path in args.source]
|
||||
if args.trace_root is None:
|
||||
raise ValueError("provide --trace-root or one or more --source")
|
||||
sources = sorted(
|
||||
path.resolve()
|
||||
for path in args.trace_root.glob("*.jsonl")
|
||||
if "prompt" not in path.stem.lower()
|
||||
)
|
||||
if not sources:
|
||||
raise FileNotFoundError(f"no non-prompt JSONL files under {args.trace_root}")
|
||||
return sources
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if not 0 < args.min_minutes <= args.max_minutes:
|
||||
raise ValueError("require 0 < min_minutes <= max_minutes")
|
||||
sources = resolve_sources(args)
|
||||
files = [scan_source(path, args) for path in sources]
|
||||
eligible = [item for item in files if item["stable_window"] is not None]
|
||||
if not eligible:
|
||||
chosen = None
|
||||
data_gate = "BLOCKED_NO_1H_WINDOW"
|
||||
else:
|
||||
chosen = min(
|
||||
eligible,
|
||||
key=lambda item: (
|
||||
item["stable_window"]["max_gap_s"] > args.max_acceptable_gap_s,
|
||||
item["stable_window"]["count_cv"],
|
||||
-item["stable_window"]["minutes"],
|
||||
item["source"],
|
||||
),
|
||||
)
|
||||
chosen["selected_window_stats"] = scan_window(
|
||||
Path(chosen["source"]), chosen["stable_window"]
|
||||
)
|
||||
exact_block_size = chosen["hash_contract"]["exact_source_block_size"]
|
||||
max_total = chosen["selected_window_stats"]["total_length"]["max"]
|
||||
data_gate = (
|
||||
"PASS"
|
||||
if exact_block_size is not None
|
||||
and max_total is not None
|
||||
and max_total <= args.model_position_limit
|
||||
else "BLOCKED_HASH_OR_POSITION_CONTRACT"
|
||||
)
|
||||
recommendation = None
|
||||
if chosen is not None:
|
||||
maximum = chosen["selected_window_stats"]["total_length"]["max"]
|
||||
recommendation = next(
|
||||
(
|
||||
limit
|
||||
for limit in MAX_MODEL_LEN_CANDIDATES
|
||||
if maximum <= limit <= args.model_position_limit
|
||||
),
|
||||
None,
|
||||
)
|
||||
payload = {
|
||||
"schema": "frontier-code-trace-audit-v1",
|
||||
"trace_root": str(args.trace_root.resolve()) if args.trace_root else None,
|
||||
"sources": [str(path) for path in sources],
|
||||
"window_policy": {
|
||||
"min_minutes": args.min_minutes,
|
||||
"max_minutes": args.max_minutes,
|
||||
"bin_seconds": args.bin_seconds,
|
||||
"max_acceptable_gap_s": args.max_acceptable_gap_s,
|
||||
"selection": "lowest density CV after rejecting anomalous-gap windows",
|
||||
},
|
||||
"model_position_limit": args.model_position_limit,
|
||||
"files": files,
|
||||
"selected": chosen,
|
||||
"max_model_len_recommendation": recommendation,
|
||||
"data_gate": data_gate,
|
||||
"runtime_gate": (
|
||||
"PENDING: vLLM startup must prove enough KV blocks and nonzero "
|
||||
"max concurrency at the recommended max_model_len for each TP"
|
||||
),
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
||||
print(json.dumps({"data_gate": data_gate, "output": str(args.output)}))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user