Audit telemetry residual tuning premise
This commit is contained in:
316
runs/telemetry-residual/run_frontier_state.py
Normal file
316
runs/telemetry-residual/run_frontier_state.py
Normal file
@@ -0,0 +1,316 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Replay frozen Frontier fixtures with state-detail outputs enabled.
|
||||
|
||||
This runner does not modify the frozen fixture/config inputs. It reuses the
|
||||
audited SimFid command and only turns on existing Frontier batch/ledger output
|
||||
flags in a separate result root.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
AITUNER_ROOT = HERE.parents[1]
|
||||
sys.path.insert(0, str(HERE))
|
||||
|
||||
from common_state import numeric, summarize_frontier # noqa: E402, F401
|
||||
|
||||
|
||||
def load_module(name: str, path: Path):
|
||||
module_root = str(path.parent.resolve())
|
||||
if module_root not in sys.path:
|
||||
sys.path.insert(0, module_root)
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError(path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
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 git_capture(root: Path, *arguments: str) -> str:
|
||||
return subprocess.run(
|
||||
["git", "-C", str(root), *arguments],
|
||||
check=True,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
).stdout
|
||||
|
||||
|
||||
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")
|
||||
os.replace(temporary, path)
|
||||
|
||||
|
||||
def enable_state_outputs(command: list[str]) -> list[str]:
|
||||
result = list(command)
|
||||
disabled_ledger = "--no-metrics_config_store_frontier_stage_batch_ledger"
|
||||
enabled_ledger = "--metrics_config_store_frontier_stage_batch_ledger"
|
||||
enabled_batches = "--metrics_config_keep_individual_batch_metrics"
|
||||
disabled_batches = "--no-metrics_config_keep_individual_batch_metrics"
|
||||
if result.count(disabled_ledger) != 1 or enabled_ledger in result:
|
||||
raise ValueError("base command must explicitly disable one full Frontier ledger")
|
||||
result[result.index(disabled_ledger)] = enabled_ledger
|
||||
if disabled_batches in result:
|
||||
if result.count(disabled_batches) != 1:
|
||||
raise ValueError("duplicate disabled individual-batch flag")
|
||||
result[result.index(disabled_batches)] = enabled_batches
|
||||
elif enabled_batches not in result:
|
||||
result.append(enabled_batches)
|
||||
if result.count(enabled_ledger) != 1 or result.count(enabled_batches) != 1:
|
||||
raise ValueError("state-detail flags were not enabled exactly once")
|
||||
return result
|
||||
|
||||
|
||||
def find_state_metrics(run_root: Path) -> dict[str, Path]:
|
||||
patterns = {
|
||||
"system": "frontier_metrics/**/system_metrics.json",
|
||||
"requests": "frontier_metrics/**/request_metrics.csv",
|
||||
"batches": "frontier_metrics/**/monolithic_batch_metrics.csv",
|
||||
"ledger": "frontier_metrics/**/frontier_stage_batch_ledger.jsonl",
|
||||
}
|
||||
result = {}
|
||||
for name, pattern in patterns.items():
|
||||
matches = sorted(run_root.glob(pattern))
|
||||
if len(matches) != 1:
|
||||
raise ValueError(
|
||||
f"{run_root}: expected exactly one {name} artifact, found {len(matches)}"
|
||||
)
|
||||
result[name] = matches[0]
|
||||
if len({path.parent for path in result.values()}) != 1:
|
||||
raise ValueError("Frontier state artifacts are not colocated")
|
||||
return result
|
||||
|
||||
|
||||
def execute(args: argparse.Namespace) -> dict[str, Any]:
|
||||
prepared = json.loads(args.prepared_manifest.read_text(encoding="utf-8"))
|
||||
if prepared.get("status") != "PASS":
|
||||
raise RuntimeError("prepared simulator manifest did not pass")
|
||||
matches = [
|
||||
entry
|
||||
for entry in prepared["entries"]
|
||||
if entry["cell"] == args.cell and entry["role"] == args.role
|
||||
]
|
||||
if len(matches) != 1:
|
||||
raise ValueError(
|
||||
f"expected one prepared entry for {args.cell}/{args.role}, found {len(matches)}"
|
||||
)
|
||||
entry = matches[0]
|
||||
if args.output.exists() and any(args.output.iterdir()):
|
||||
raise FileExistsError(f"non-empty output already exists: {args.output}")
|
||||
args.output.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
driver = load_module(
|
||||
"telemetry_residual_execution_driver",
|
||||
args.replayserve_root / "runs/simfid_s2rb/results/execution_driver.py",
|
||||
)
|
||||
config_path = Path(entry["config"])
|
||||
config = json.loads(config_path.read_text(encoding="utf-8"))
|
||||
fixture_manifest_path = Path(entry["fixture_manifest"])
|
||||
fixture = json.loads(fixture_manifest_path.read_text(encoding="utf-8"))
|
||||
trace_path = Path(entry["frontier_csv"])
|
||||
sidecar_path = Path(entry["sidecar"])
|
||||
metrics_root = args.output / "frontier_metrics"
|
||||
run_id = f"telemetry_residual_r0_{args.cell}_{args.role}"
|
||||
knobs = config["frontier"]["knobs"]
|
||||
base_command = driver.build_command(
|
||||
trace_path=trace_path,
|
||||
metrics_root=metrics_root,
|
||||
run_id=run_id,
|
||||
knobs=knobs,
|
||||
)
|
||||
driver.audit_command(base_command, knobs)
|
||||
command = enable_state_outputs(base_command)
|
||||
row = {
|
||||
"hook_path": config["calibration"]["hook_path"],
|
||||
"applied_a_tp": config["calibration"]["a_tp"],
|
||||
"sidecar_path": str(sidecar_path),
|
||||
"request_count": int(fixture["request_count"]),
|
||||
"tensor_parallel_size": int(fixture["tensor_parallel_size"]),
|
||||
}
|
||||
environment = driver.environment_for(row)
|
||||
if environment.get("CUDA_VISIBLE_DEVICES") != "":
|
||||
raise ValueError("Frontier state replay must hide CUDA devices")
|
||||
manifest = {
|
||||
"schema": "telemetry-residual-frontier-state-run-v1",
|
||||
"entry": {
|
||||
"cell": args.cell,
|
||||
"role": args.role,
|
||||
"anchor": entry["anchor"],
|
||||
"request_count": entry["selected_count"],
|
||||
},
|
||||
"inputs": {
|
||||
"prepared_manifest": str(args.prepared_manifest.resolve()),
|
||||
"prepared_manifest_sha256": sha256_file(args.prepared_manifest),
|
||||
"config": str(config_path.resolve()),
|
||||
"config_sha256": sha256_file(config_path),
|
||||
"fixture_manifest": str(fixture_manifest_path.resolve()),
|
||||
"fixture_manifest_sha256": sha256_file(fixture_manifest_path),
|
||||
"frontier_csv": str(trace_path.resolve()),
|
||||
"frontier_csv_sha256": sha256_file(trace_path),
|
||||
"sidecar": str(sidecar_path.resolve()),
|
||||
"sidecar_sha256": sha256_file(sidecar_path),
|
||||
},
|
||||
"frontier": {
|
||||
"root": str(args.frontier_root.resolve()),
|
||||
"git_head": git_capture(args.frontier_root, "rev-parse", "HEAD").strip(),
|
||||
"git_status_short": git_capture(args.frontier_root, "status", "--short"),
|
||||
},
|
||||
"runner": {
|
||||
"script": str(Path(__file__).resolve()),
|
||||
"script_sha256": sha256_file(Path(__file__).resolve()),
|
||||
"aituner_git_head": git_capture(AITUNER_ROOT, "rev-parse", "HEAD").strip(),
|
||||
"aituner_git_status_short": git_capture(AITUNER_ROOT, "status", "--short"),
|
||||
},
|
||||
"environment": {
|
||||
key: environment[key]
|
||||
for key in (
|
||||
"PYTHONPATH",
|
||||
"FRONTIER_EXECUTION_TIME_SCALE",
|
||||
"CUDA_VISIBLE_DEVICES",
|
||||
"NVIDIA_VISIBLE_DEVICES",
|
||||
"FRONTIER_LOG_LEVEL",
|
||||
)
|
||||
},
|
||||
"command": command,
|
||||
"state_outputs": {
|
||||
"individual_batch_metrics": True,
|
||||
"full_stage_batch_ledger": True,
|
||||
},
|
||||
"contains_prompt_text": False,
|
||||
}
|
||||
atomic_json(args.output / "run_manifest.json", manifest)
|
||||
start = time.monotonic()
|
||||
with (args.output / "stdout.log").open("w", encoding="utf-8") as stdout, (
|
||||
args.output / "stderr.log"
|
||||
).open("w", encoding="utf-8") as stderr:
|
||||
try:
|
||||
process = subprocess.run(
|
||||
command,
|
||||
cwd=args.frontier_root,
|
||||
env=environment,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
timeout=args.timeout_s,
|
||||
)
|
||||
return_code = int(process.returncode)
|
||||
except subprocess.TimeoutExpired:
|
||||
return_code = 124
|
||||
runtime_s = time.monotonic() - start
|
||||
if return_code != 0:
|
||||
failure = {
|
||||
"status": "STOP",
|
||||
"return_code": return_code,
|
||||
"runtime_s": runtime_s,
|
||||
}
|
||||
atomic_json(args.output / "failure.json", failure)
|
||||
raise RuntimeError(f"Frontier state replay failed: {failure}")
|
||||
|
||||
paths = find_state_metrics(args.output)
|
||||
summary = summarize_frontier(
|
||||
system_metrics_path=paths["system"],
|
||||
request_metrics_path=paths["requests"],
|
||||
batch_metrics_path=paths["batches"],
|
||||
ledger_path=paths["ledger"],
|
||||
)
|
||||
atomic_json(args.output / "common-state.json", summary)
|
||||
scorer = driver.score_trial(row, paths["system"], paths["requests"])
|
||||
atomic_json(args.output / "scorer_output.json", scorer)
|
||||
sizes = {name: path.stat().st_size for name, path in paths.items()}
|
||||
red_flags = []
|
||||
if summary["interval"]["request_count"] != int(entry["selected_count"]):
|
||||
red_flags.append("request_count_mismatch")
|
||||
if summary["sanity"]["batch_rows"] <= 0:
|
||||
red_flags.append("no_batch_rows")
|
||||
if summary["sanity"]["ledger_rows"] <= 0:
|
||||
red_flags.append("no_ledger_rows")
|
||||
result = {
|
||||
"schema": "telemetry-residual-frontier-state-result-v1",
|
||||
"status": "PASS" if not red_flags else "STOP",
|
||||
"runtime_s": runtime_s,
|
||||
"return_code": return_code,
|
||||
"paths": {name: str(path.resolve()) for name, path in paths.items()},
|
||||
"bytes": sizes,
|
||||
"common_state": str((args.output / "common-state.json").resolve()),
|
||||
"red_flags": red_flags,
|
||||
"sanity": {
|
||||
"request_rows": summary["sanity"]["request_rows"],
|
||||
"batch_rows": summary["sanity"]["batch_rows"],
|
||||
"ledger_rows": summary["sanity"]["ledger_rows"],
|
||||
"artifact_bytes": {
|
||||
"n": len(sizes),
|
||||
"min": min(sizes.values()),
|
||||
"max": max(sizes.values()),
|
||||
"distinct_n": len(set(sizes.values())),
|
||||
},
|
||||
"invariants": {
|
||||
"zero_failures": return_code == 0,
|
||||
"gpu_visibility_disabled": True,
|
||||
"request_count_match": not red_flags,
|
||||
"nonnegative_file_sizes": all(size >= 0 for size in sizes.values()),
|
||||
"state_values_not_all_identical": summary["sanity"]["invariants"][
|
||||
"batch_values_not_all_identical"
|
||||
],
|
||||
},
|
||||
},
|
||||
}
|
||||
atomic_json(args.output / "result.json", result)
|
||||
if result["status"] != "PASS":
|
||||
raise RuntimeError(red_flags)
|
||||
return result
|
||||
|
||||
|
||||
def parser() -> argparse.ArgumentParser:
|
||||
result = argparse.ArgumentParser()
|
||||
result.add_argument("--prepared-manifest", type=Path, required=True)
|
||||
result.add_argument("--output", type=Path, required=True)
|
||||
result.add_argument("--replayserve-root", type=Path, required=True)
|
||||
result.add_argument("--frontier-root", type=Path, required=True)
|
||||
result.add_argument("--cell", required=True)
|
||||
result.add_argument("--role", choices=("low1", "high1"), required=True)
|
||||
result.add_argument("--timeout-s", type=float, default=300.0)
|
||||
return result
|
||||
|
||||
|
||||
def main() -> None:
|
||||
result = execute(parser().parse_args())
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"status": result["status"],
|
||||
"runtime_s": result["runtime_s"],
|
||||
"request_rows": result["sanity"]["request_rows"],
|
||||
"batch_rows": result["sanity"]["batch_rows"],
|
||||
"ledger_rows": result["sanity"]["ledger_rows"],
|
||||
"red_flags": result["red_flags"],
|
||||
},
|
||||
sort_keys=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user