364 lines
13 KiB
Python
364 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Freeze the unseen-trace 2x2 active intervention development surface."""
|
|
|
|
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
|
|
|
|
|
|
SCHEMA = "active-intervention-prospective-manifest-v0"
|
|
TP = 4
|
|
REPETITIONS = (1, 2, 3)
|
|
DURATION_S = 300.0
|
|
REPLAY_TIME_SCALE = 0.5
|
|
OFFERED_RATE_PER_GPU = 2.75
|
|
TARGET_COUNT = round(OFFERED_RATE_PER_GPU * DURATION_S * TP)
|
|
WINDOW_ID = "chat_w20260313_1000"
|
|
ENGINE_VERSION = "0.24.1.dev3+g668cfb7e2"
|
|
|
|
|
|
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 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 configs() -> list[dict[str, Any]]:
|
|
return [
|
|
{
|
|
"id": "source_mns32_mbbt4096",
|
|
"mns": 32,
|
|
"mbbt": 4096,
|
|
"repetition_order": [1, 2, 3],
|
|
},
|
|
{
|
|
"id": "mns64_mbbt4096",
|
|
"mns": 64,
|
|
"mbbt": 4096,
|
|
"repetition_order": [2, 3, 1],
|
|
},
|
|
{
|
|
"id": "mns32_mbbt8192",
|
|
"mns": 32,
|
|
"mbbt": 8192,
|
|
"repetition_order": [3, 1, 2],
|
|
},
|
|
{
|
|
"id": "joint_mns64_mbbt8192",
|
|
"mns": 64,
|
|
"mbbt": 8192,
|
|
"repetition_order": [1, 3, 2],
|
|
},
|
|
]
|
|
|
|
|
|
def partition_trace(source: Path, output_root: Path) -> dict[str, Any]:
|
|
source_sha = sha256_file(source)
|
|
output_root.mkdir(parents=True, exist_ok=True)
|
|
paths = {rep: output_root / f"rep{rep}.jsonl" for rep in REPETITIONS}
|
|
temporary = {rep: path.with_suffix(".jsonl.tmp") for rep, path in paths.items()}
|
|
handles = {rep: temporary[rep].open("w", encoding="utf-8") for rep in REPETITIONS}
|
|
counts = {rep: 0 for rep in REPETITIONS}
|
|
id_digests = {rep: hashlib.sha256() for rep in REPETITIONS}
|
|
total = 0
|
|
try:
|
|
with source.open(encoding="utf-8") as input_file:
|
|
for line_number, line in enumerate(input_file, start=1):
|
|
if not line.strip():
|
|
continue
|
|
row = json.loads(line)
|
|
original_id = str(row.get("request_id") or row.get("id") or line_number)
|
|
digest = hashlib.sha256(
|
|
f"{source_sha}:{line_number}:{original_id}".encode()
|
|
).hexdigest()
|
|
repetition = int(digest[:16], 16) % len(REPETITIONS) + 1
|
|
row["request_id"] = f"active-r{repetition}-{digest}"
|
|
handles[repetition].write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
counts[repetition] += 1
|
|
total += 1
|
|
id_digests[repetition].update(row["request_id"].encode() + b"\n")
|
|
finally:
|
|
for handle in handles.values():
|
|
handle.close()
|
|
for repetition in REPETITIONS:
|
|
os.replace(temporary[repetition], paths[repetition])
|
|
partitions = {
|
|
str(rep): {
|
|
"path": str(paths[rep]),
|
|
"rows": counts[rep],
|
|
"bytes": paths[rep].stat().st_size,
|
|
"sha256": sha256_file(paths[rep]),
|
|
"request_id_order_sha256": id_digests[rep].hexdigest(),
|
|
}
|
|
for rep in REPETITIONS
|
|
}
|
|
return {
|
|
"source": str(source),
|
|
"source_sha256": source_sha,
|
|
"source_rows": total,
|
|
"partition_rule": "sha256(source_sha:line_number:original_id) modulo 3",
|
|
"partitions": partitions,
|
|
}
|
|
|
|
|
|
def materialize_study(
|
|
base_study: Path,
|
|
target: Path,
|
|
*,
|
|
repetition: int,
|
|
trace_path: Path,
|
|
windows_path: Path,
|
|
) -> None:
|
|
payload = json.loads(base_study.read_text(encoding="utf-8"))
|
|
payload["study_id"] = f"active-intervention-trace13-rep{repetition}"
|
|
payload["hardware"]["host_candidates"] = ["dash0"]
|
|
payload["engine"]["engine_version"] = ENGINE_VERSION
|
|
trace = payload["trace"]
|
|
trace.update(
|
|
{
|
|
"windows_path": str(windows_path),
|
|
"window_id": WINDOW_ID,
|
|
"trace_file_override": str(trace_path),
|
|
"completion_tokens_override": 128,
|
|
"replay_time_scale": REPLAY_TIME_SCALE,
|
|
"early_stop_max_lag_s": None,
|
|
"early_stop_max_elapsed_s": 360.0,
|
|
"restart_engine_after_early_stop": False,
|
|
"adaptive_stop": {"enabled": False},
|
|
}
|
|
)
|
|
atomic_json(target, payload)
|
|
|
|
|
|
def attainable_anchor(requests: list[Any], target_count: int) -> tuple[float, list[Any]]:
|
|
ordered = sorted(float(request.sampling_u) for request in requests)
|
|
if target_count <= 0 or target_count > len(ordered):
|
|
raise ValueError(
|
|
f"target count {target_count} is outside available range 1..{len(ordered)}"
|
|
)
|
|
candidates = []
|
|
for index in sorted({target_count - 1, min(target_count, len(ordered) - 1)}):
|
|
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], item[2])
|
|
)
|
|
return anchor, selected
|
|
|
|
|
|
def selection_record(selected: list[Any]) -> dict[str, Any]:
|
|
return {
|
|
"anchor": max(float(request.sampling_u) for request in selected),
|
|
"selected_count": len(selected),
|
|
"target_count": TARGET_COUNT,
|
|
"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(
|
|
*,
|
|
base_study: Path,
|
|
base_action_manifest: Path,
|
|
source_trace: Path,
|
|
windows_path: Path,
|
|
private_root: Path,
|
|
policy_path: Path,
|
|
) -> dict[str, Any]:
|
|
base_manifest = json.loads(base_action_manifest.read_text(encoding="utf-8"))
|
|
if base_manifest.get("status") != "PASS":
|
|
raise ValueError("base action-aware manifest did not pass")
|
|
policy = json.loads(policy_path.read_text(encoding="utf-8"))
|
|
if policy.get("schema") != "active-intervention-policy-v0":
|
|
raise ValueError("unexpected frozen policy schema")
|
|
if policy.get("sanity", {}).get("red_flags"):
|
|
raise ValueError("frozen policy contains red flags")
|
|
|
|
partition = partition_trace(source_trace, private_root / "traces")
|
|
repetitions = {}
|
|
selected_sets: list[set[str]] = []
|
|
for repetition in REPETITIONS:
|
|
trace_path = Path(partition["partitions"][str(repetition)]["path"])
|
|
study_path = private_root / "studies" / f"rep{repetition}-tp4.json"
|
|
materialize_study(
|
|
base_study,
|
|
study_path,
|
|
repetition=repetition,
|
|
trace_path=trace_path,
|
|
windows_path=windows_path,
|
|
)
|
|
study = load_study_spec(study_path)
|
|
window, requests = load_trace_requests(study, study_spec_path=study_path)
|
|
duration_s = float(window.window_end - window.window_start)
|
|
if not math.isclose(duration_s, DURATION_S, abs_tol=1e-9):
|
|
raise ValueError(f"rep{repetition}: duration {duration_s} != {DURATION_S}")
|
|
_anchor, selected = attainable_anchor(requests, TARGET_COUNT)
|
|
record = selection_record(selected)
|
|
selected_sets.append({request.row_id for request in selected})
|
|
repetitions[str(repetition)] = {
|
|
"study": str(study_path),
|
|
"study_sha256": sha256_file(study_path),
|
|
"trace": partition["partitions"][str(repetition)],
|
|
"available_filtered_requests": len(requests),
|
|
"selection": record,
|
|
}
|
|
|
|
frozen_configs = configs()
|
|
config_ids = {str(config["id"]) for config in frozen_configs}
|
|
invariants = {
|
|
"three_nonempty_trace_partitions": all(
|
|
int(item["rows"]) > 0 for item in partition["partitions"].values()
|
|
),
|
|
"partition_rows_conserved": sum(
|
|
int(item["rows"]) for item in partition["partitions"].values()
|
|
)
|
|
== int(partition["source_rows"]),
|
|
"selected_sets_disjoint": all(
|
|
not selected_sets[left] & selected_sets[right]
|
|
for left in range(len(selected_sets))
|
|
for right in range(left + 1, len(selected_sets))
|
|
),
|
|
"target_count_attained": all(
|
|
abs(int(item["selection"]["selected_count"]) - TARGET_COUNT) <= 1
|
|
for item in repetitions.values()
|
|
),
|
|
"four_unique_configs": len(config_ids) == 4,
|
|
"two_by_two_surface": {
|
|
(int(config["mns"]), int(config["mbbt"]))
|
|
for config in frozen_configs
|
|
}
|
|
== {(32, 4096), (64, 4096), (32, 8192), (64, 8192)},
|
|
"repetition_orders_are_permutations": all(
|
|
sorted(config["repetition_order"]) == list(REPETITIONS)
|
|
for config in frozen_configs
|
|
),
|
|
}
|
|
red_flags = [name for name, passed in invariants.items() if not passed]
|
|
return {
|
|
"schema": SCHEMA,
|
|
"status": "PASS" if not red_flags else "STOP",
|
|
"source": {
|
|
"window_id": WINDOW_ID,
|
|
"source_trace": str(source_trace),
|
|
"source_trace_sha256": partition["source_sha256"],
|
|
"windows_path": str(windows_path),
|
|
"base_study": str(base_study),
|
|
"base_study_sha256": sha256_file(base_study),
|
|
"base_action_manifest": str(base_action_manifest),
|
|
"base_action_manifest_sha256": sha256_file(base_action_manifest),
|
|
},
|
|
"policy": {
|
|
"path": str(policy_path),
|
|
"sha256": sha256_file(policy_path),
|
|
"status": policy["status"],
|
|
"training": policy["training"],
|
|
"measurement_policy": policy["measurement_policy"],
|
|
"launch_reason": (
|
|
"bounded unseen-trace joint-action test after a negative narrow "
|
|
"retrospective replay"
|
|
),
|
|
},
|
|
"engine": {
|
|
"tp": TP,
|
|
"duration_s": DURATION_S,
|
|
"client_timeout_s": 450.0,
|
|
"burnin_max_elapsed_s": 90.0,
|
|
"disable_slo_early_stop": True,
|
|
},
|
|
"burnin": base_manifest["burnin"],
|
|
"private": {"trace_partition": partition},
|
|
"repetitions": repetitions,
|
|
"configs": frozen_configs,
|
|
"source_config_id": "source_mns32_mbbt4096",
|
|
"actions": {
|
|
"noop": "source_mns32_mbbt4096",
|
|
"mns": "mns64_mbbt4096",
|
|
"mbbt": "mns32_mbbt8192",
|
|
"joint": "joint_mns64_mbbt8192",
|
|
},
|
|
"checkpoints": {
|
|
"fractions": [0.25, 0.50, 0.75, 1.0],
|
|
"seconds": [75.0, 150.0, 225.0, 300.0],
|
|
},
|
|
"gates": {
|
|
"acceptable_regret": 0.02,
|
|
"source_ceiling_normalized_goodput": 0.98,
|
|
"confirmation_trigger_gpu_cost_reduction": 0.10,
|
|
"contribution_gpu_cost_reduction": 0.20,
|
|
"maximum_task_regret": 0.05,
|
|
},
|
|
"budget": {
|
|
"hard_cap_h20_hours": 6.0,
|
|
"session_estimate_h20_hours": 1.3,
|
|
"safety_h20_hours": 0.3,
|
|
"expected_h20_hours": [4.6, 5.5],
|
|
"expected_wall_minutes": [75, 100],
|
|
},
|
|
"sanity": {"invariants": invariants, "red_flags": red_flags},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-study", type=Path, required=True)
|
|
parser.add_argument("--base-action-manifest", type=Path, required=True)
|
|
parser.add_argument("--source-trace", type=Path, required=True)
|
|
parser.add_argument("--windows-path", type=Path, required=True)
|
|
parser.add_argument("--private-root", type=Path, required=True)
|
|
parser.add_argument("--policy", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
payload = build(
|
|
base_study=args.base_study,
|
|
base_action_manifest=args.base_action_manifest,
|
|
source_trace=args.source_trace,
|
|
windows_path=args.windows_path,
|
|
private_root=args.private_root,
|
|
policy_path=args.policy,
|
|
)
|
|
atomic_json(args.output, payload)
|
|
print(json.dumps({"status": payload["status"], "sanity": payload["sanity"]}))
|
|
if payload["status"] != "PASS":
|
|
raise SystemExit("prospective manifest preflight failed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|