362 lines
13 KiB
Python
362 lines
13 KiB
Python
#!/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()
|