356 lines
14 KiB
Python
356 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Materialize session-disjoint pilot repeats and freeze attainable anchors.
|
|
|
|
The private outputs retain prompt text and stay on the experiment host. The
|
|
public manifest contains only aggregate counts, hashes, paths, and parameters.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
AITUNER_ROOT = Path(os.environ.get("AITUNER_ROOT", Path(__file__).resolve().parents[2]))
|
|
sys.path.insert(0, str(AITUNER_ROOT / "src"))
|
|
|
|
from aituner.spec import load_study_spec # noqa: E402
|
|
from aituner.trace import load_trace_requests, select_requests_for_threshold # noqa: E402
|
|
|
|
|
|
ROLES = ("burnin", "low1", "high1", "low2", "high2", "low3", "high3")
|
|
CELLS = {
|
|
"tp1_mns8": {"tp": 1, "mns": 8, "frontier_req_s_gpu": 2.3833333333333333},
|
|
"tp1_mns64": {"tp": 1, "mns": 64, "frontier_req_s_gpu": 2.3833333333333333},
|
|
"tp2_mns8": {"tp": 2, "mns": 8, "frontier_req_s_gpu": 2.2416666666666667},
|
|
"tp2_mns64": {"tp": 2, "mns": 64, "frontier_req_s_gpu": 2.3},
|
|
"tp4_mns16": {"tp": 4, "mns": 16, "frontier_req_s_gpu": 2.5},
|
|
"tp4_mns64": {"tp": 4, "mns": 64, "frontier_req_s_gpu": 2.5},
|
|
}
|
|
TARGET_MULTIPLIERS = {"low": 0.85, "high": 1.25}
|
|
|
|
|
|
def atomic_json(path: Path, payload: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_suffix(path.suffix + ".tmp")
|
|
tmp.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
|
os.replace(tmp, path)
|
|
|
|
|
|
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 order_hash(values: list[str]) -> str:
|
|
return hashlib.sha256("\n".join(values).encode()).hexdigest()
|
|
|
|
|
|
def resolve_source_trace(windows_path: Path, window_id: str) -> tuple[dict[str, Any], Path]:
|
|
payload = json.loads(windows_path.read_text(encoding="utf-8"))
|
|
for window in payload["windows"]:
|
|
if window["window_id"] != window_id:
|
|
continue
|
|
trace = Path(window["trace_file"])
|
|
if not trace.is_absolute():
|
|
trace = (windows_path.parent / trace).resolve()
|
|
return window, trace
|
|
raise ValueError(f"window not found: {window_id}")
|
|
|
|
|
|
def materialize_bands(
|
|
source_trace: Path,
|
|
source_window: dict[str, Any],
|
|
private_root: Path,
|
|
) -> tuple[Path, dict[str, Any]]:
|
|
traces_root = private_root / "traces"
|
|
traces_root.mkdir(parents=True, exist_ok=True)
|
|
temporary = {role: traces_root / f".{role}.jsonl.tmp" for role in ROLES}
|
|
final = {role: traces_root / f"{role}.jsonl" for role in ROLES}
|
|
handles = {role: temporary[role].open("w", encoding="utf-8") for role in ROLES}
|
|
stats = {
|
|
role: {
|
|
"rows": 0,
|
|
"sum_input_tokens": 0,
|
|
"min_timestamp": None,
|
|
"max_timestamp": None,
|
|
}
|
|
for role in ROLES
|
|
}
|
|
try:
|
|
with source_trace.open(encoding="utf-8") as source:
|
|
for line_number, line in enumerate(source):
|
|
row = json.loads(line)
|
|
value = float(row["sampling_u"])
|
|
if not 0.0 <= value <= 1.0:
|
|
raise ValueError(f"sampling_u outside [0,1] at line {line_number}")
|
|
band = min(len(ROLES) - 1, int(value * len(ROLES)))
|
|
role = ROLES[band]
|
|
remapped = value * len(ROLES) - band
|
|
row["sampling_u"] = min(remapped, math.nextafter(1.0, 0.0))
|
|
row["fidelity_pilot_band"] = role
|
|
handles[role].write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
timestamp = float(row["timestamp"])
|
|
item = stats[role]
|
|
item["rows"] += 1
|
|
item["sum_input_tokens"] += int(row.get("input_length") or 0)
|
|
item["min_timestamp"] = (
|
|
timestamp if item["min_timestamp"] is None
|
|
else min(float(item["min_timestamp"]), timestamp)
|
|
)
|
|
item["max_timestamp"] = (
|
|
timestamp if item["max_timestamp"] is None
|
|
else max(float(item["max_timestamp"]), timestamp)
|
|
)
|
|
finally:
|
|
for handle in handles.values():
|
|
handle.close()
|
|
for role in ROLES:
|
|
os.replace(temporary[role], final[role])
|
|
stats[role]["sha256"] = sha256_file(final[role])
|
|
stats[role]["bytes"] = final[role].stat().st_size
|
|
|
|
windows = []
|
|
for role in ROLES:
|
|
window = dict(source_window)
|
|
window["window_id"] = f"fidelity_pilot_{role}"
|
|
window["trace_file"] = f"traces/{role}.jsonl"
|
|
window["num_requests"] = stats[role]["rows"]
|
|
window["sum_input_length"] = stats[role]["sum_input_tokens"]
|
|
window["sampling_strategy"] = "session_uniform_seven_disjoint_bands_remapped"
|
|
window["fidelity_pilot_role"] = role
|
|
windows.append(window)
|
|
private_windows = private_root / "windows.json"
|
|
atomic_json(
|
|
private_windows,
|
|
{
|
|
"schema": "fidelity-pilot-private-windows-v1",
|
|
"roles": list(ROLES),
|
|
"windows": windows,
|
|
},
|
|
)
|
|
return private_windows, stats
|
|
|
|
|
|
def write_studies(
|
|
*,
|
|
base_primary: Path,
|
|
base_tp4: Path,
|
|
private_windows: Path,
|
|
private_root: Path,
|
|
) -> dict[str, dict[str, Path]]:
|
|
bases = {
|
|
"primary": json.loads(base_primary.read_text(encoding="utf-8")),
|
|
"tp4": json.loads(base_tp4.read_text(encoding="utf-8")),
|
|
}
|
|
result: dict[str, dict[str, Path]] = {}
|
|
for role in ROLES:
|
|
result[role] = {}
|
|
for tier, base in bases.items():
|
|
payload = json.loads(json.dumps(base))
|
|
payload["study_id"] = f"fidelity-prefix-pilot-{role}-{tier}"
|
|
payload["hardware"]["host_candidates"] = ["dash0"]
|
|
payload["engine"]["engine_version"] = "0.24.1.dev3+opprof"
|
|
payload["trace"]["windows_path"] = str(private_windows)
|
|
payload["trace"]["window_id"] = f"fidelity_pilot_{role}"
|
|
path = private_root / "studies" / f"{role}-{tier}.json"
|
|
atomic_json(path, payload)
|
|
result[role][tier] = path
|
|
return result
|
|
|
|
|
|
def attainable_anchor(requests: list[Any], target_count: int) -> tuple[float, list[Any]]:
|
|
ordered = sorted(float(request.sampling_u) for request in requests)
|
|
if not ordered:
|
|
raise ValueError("no requests after study filtering")
|
|
candidate_indices = sorted({
|
|
max(0, min(len(ordered) - 1, target_count - 1)),
|
|
max(0, min(len(ordered) - 1, target_count)),
|
|
})
|
|
candidates = []
|
|
for index in candidate_indices:
|
|
anchor = ordered[index]
|
|
selected = select_requests_for_threshold(requests, threshold=anchor)
|
|
candidates.append((abs(len(selected) - target_count), len(selected), anchor, selected))
|
|
_error, _count, anchor, selected = min(candidates, key=lambda item: (item[0], item[1]))
|
|
return anchor, selected
|
|
|
|
|
|
def selected_record(selected: list[Any], *, tp: int, duration_s: float) -> dict[str, Any]:
|
|
return {
|
|
"anchor": max(float(request.sampling_u) for request in selected),
|
|
"selected_count": len(selected),
|
|
"offered_req_s": len(selected) / duration_s,
|
|
"offered_req_s_per_gpu": len(selected) / duration_s / tp,
|
|
"request_id_order_sha256": order_hash([request.row_id for request in selected]),
|
|
"arrival_order_sha256": order_hash([f"{request.arrival_s:.12f}" for request in selected]),
|
|
"input_length_order_sha256": order_hash(
|
|
[str(request.prompt_tokens_hint) for request in selected]
|
|
),
|
|
}
|
|
|
|
|
|
def build_manifest(
|
|
*,
|
|
studies: dict[str, dict[str, Path]],
|
|
private_windows: Path,
|
|
band_stats: dict[str, Any],
|
|
source_trace: Path,
|
|
source_windows: Path,
|
|
source_window_id: str,
|
|
) -> dict[str, Any]:
|
|
loaded = {}
|
|
durations = {}
|
|
for role, tiers in studies.items():
|
|
loaded[role] = {}
|
|
for tier, path in tiers.items():
|
|
study = load_study_spec(path)
|
|
window, requests = load_trace_requests(study, study_spec_path=path)
|
|
loaded[role][tier] = requests
|
|
durations[role] = float(window.window_end - window.window_start)
|
|
|
|
cells = {}
|
|
all_hashes = []
|
|
for cell, config in CELLS.items():
|
|
tp = int(config["tp"])
|
|
tier = "tp4" if tp == 4 else "primary"
|
|
targets = {}
|
|
for level, multiplier in TARGET_MULTIPLIERS.items():
|
|
target_rate = float(config["frontier_req_s_gpu"]) * multiplier
|
|
target_count = round(target_rate * durations["low1"] * tp)
|
|
roles = [
|
|
role
|
|
for role in ROLES
|
|
if role.startswith(level) or (level == "low" and role == "burnin")
|
|
]
|
|
selections = {}
|
|
for role in roles:
|
|
anchor, selected = attainable_anchor(loaded[role][tier], target_count)
|
|
record = selected_record(selected, tp=tp, duration_s=durations[role])
|
|
record["anchor"] = anchor
|
|
record["study"] = str(studies[role][tier])
|
|
selections[role] = record
|
|
all_hashes.append(record["request_id_order_sha256"])
|
|
targets[level] = {
|
|
"multiplier": multiplier,
|
|
"target_req_s_per_gpu": target_rate,
|
|
"target_count": target_count,
|
|
"selections": selections,
|
|
}
|
|
cells[cell] = {**config, "targets": targets}
|
|
|
|
red_flags = []
|
|
for cell, config in cells.items():
|
|
for level, target in config["targets"].items():
|
|
if not target["selections"]:
|
|
red_flags.append(f"missing_{cell}_{level}")
|
|
for selection in target["selections"].values():
|
|
if selection["selected_count"] <= 0:
|
|
red_flags.append(f"empty_{cell}_{level}")
|
|
per_cell_distinct = {}
|
|
for cell, config in cells.items():
|
|
hashes = [
|
|
selection["request_id_order_sha256"]
|
|
for target in config["targets"].values()
|
|
for selection in target["selections"].values()
|
|
]
|
|
per_cell_distinct[cell] = len(hashes) == len(set(hashes))
|
|
if not per_cell_distinct[cell]:
|
|
red_flags.append(f"session_bands_overlap_{cell}")
|
|
return {
|
|
"schema": "fidelity-prefix-pilot-manifest-v1",
|
|
"status": "PASS" if not red_flags else "STOP",
|
|
"source": {
|
|
"windows": str(source_windows),
|
|
"window_id": source_window_id,
|
|
"trace": str(source_trace),
|
|
"trace_sha256": sha256_file(source_trace),
|
|
},
|
|
"private": {
|
|
"windows": str(private_windows),
|
|
"windows_sha256": sha256_file(private_windows),
|
|
"band_stats": band_stats,
|
|
"studies": {
|
|
role: {tier: str(path) for tier, path in tiers.items()}
|
|
for role, tiers in studies.items()
|
|
},
|
|
},
|
|
"roles": list(ROLES),
|
|
"cells": cells,
|
|
"execution": {
|
|
"cutoff_s": 5.0,
|
|
"replicates_per_level": 3,
|
|
"label": "2-of-3 session-disjoint repetitions",
|
|
"even_cell_order": ["low1", "high1", "high2", "low2", "low3", "high3"],
|
|
"odd_cell_order": ["high1", "low1", "low2", "high2", "high3", "low3"],
|
|
"hard_cap_h20_hours": 3.5,
|
|
},
|
|
"sanity": {
|
|
"red_flags": red_flags,
|
|
"n_cells": len(cells),
|
|
"n_roles": len(ROLES),
|
|
"selected_sets": len(all_hashes),
|
|
"distinct_selected_sets": len(set(all_hashes)),
|
|
"per_cell_selected_sets_distinct": per_cell_distinct,
|
|
"invariants": {
|
|
"cells_6": len(cells) == 6,
|
|
"roles_7": len(ROLES) == 7,
|
|
"band_rows_nonzero": all(stats["rows"] > 0 for stats in band_stats.values()),
|
|
"session_bands_disjoint_per_cell": all(per_cell_distinct.values()),
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--source-windows", type=Path, required=True)
|
|
parser.add_argument("--source-window-id", default="chat_w20260312_1000")
|
|
parser.add_argument("--base-primary-study", type=Path, required=True)
|
|
parser.add_argument("--base-tp4-study", type=Path, required=True)
|
|
parser.add_argument("--private-root", type=Path, required=True)
|
|
parser.add_argument("--public-manifest", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
|
|
source_window, source_trace = resolve_source_trace(
|
|
args.source_windows, args.source_window_id
|
|
)
|
|
private_windows, band_stats = materialize_bands(
|
|
source_trace, source_window, args.private_root
|
|
)
|
|
studies = write_studies(
|
|
base_primary=args.base_primary_study,
|
|
base_tp4=args.base_tp4_study,
|
|
private_windows=private_windows,
|
|
private_root=args.private_root,
|
|
)
|
|
manifest = build_manifest(
|
|
studies=studies,
|
|
private_windows=private_windows,
|
|
band_stats=band_stats,
|
|
source_trace=source_trace,
|
|
source_windows=args.source_windows,
|
|
source_window_id=args.source_window_id,
|
|
)
|
|
atomic_json(args.public_manifest, manifest)
|
|
print(json.dumps({
|
|
"status": manifest["status"],
|
|
"manifest": str(args.public_manifest),
|
|
"sanity": manifest["sanity"],
|
|
}, sort_keys=True))
|
|
if manifest["status"] != "PASS":
|
|
raise RuntimeError(f"pilot preflight failed: {manifest['sanity']['red_flags']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|