Add prospective active intervention experiment
This commit is contained in:
362
runs/active-intervention-v0/prepare_prospective.py
Normal file
362
runs/active-intervention-v0/prepare_prospective.py
Normal file
@@ -0,0 +1,362 @@
|
||||
#!/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"
|
||||
|
||||
|
||||
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"] = "0.24.1.dev3+opprof"
|
||||
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()
|
||||
Reference in New Issue
Block a user