File size: 4,518 Bytes
bc29ee3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | #!/usr/bin/env python3
"""Merge independently evaluated two-block FPPP shards into one result set."""
from __future__ import annotations
import argparse
import csv
import json
import os
import shutil
from pathlib import Path
from typing import Any
PROMPT_IDS = list(range(80, 100))
SCHEDULE = "chunk0=FFFF; chunks1-6=FPPP"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--main_dir", type=Path, required=True)
parser.add_argument("--shard_dir", type=Path, required=True)
parser.add_argument("--sweep_dir", type=Path, required=True)
return parser.parse_args()
def atomic_json(path: Path, value: Any) -> None:
temporary = path.with_suffix(path.suffix + ".tmp")
temporary.write_text(
json.dumps(value, indent=2, ensure_ascii=False) + "\n",
encoding="utf-8",
)
os.replace(temporary, path)
def read_experiment_names(sweep_dir: Path) -> list[str]:
with (sweep_dir / "summary.csv").open(
"r", encoding="utf-8", newline=""
) as handle:
return [row["name"] for row in csv.DictReader(handle)]
def merge_pair_directories(main_dir: Path, shard_dir: Path) -> None:
for source in shard_dir.glob("pair_*"):
if source.is_dir():
shutil.copytree(source, main_dir / source.name, dirs_exist_ok=True)
def collect_rows(main_dir: Path, names: list[str]) -> list[dict[str, Any]]:
rows = []
for name in names:
metrics_path = main_dir / name / "metrics.json"
if not metrics_path.exists():
raise FileNotFoundError(f"Missing completed metrics: {metrics_path}")
metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
if metrics.get("status") != "complete":
raise RuntimeError(f"Incomplete metrics for {name}")
if metrics.get("schedule") != SCHEDULE:
raise RuntimeError(f"Unexpected schedule for {name}: {metrics.get('schedule')}")
if metrics.get("prompt_ids") != PROMPT_IDS:
raise RuntimeError(f"Unexpected prompt split for {name}")
rows.append(
{
"name": metrics["name"],
"source_layer_1": metrics["source_layers"][0],
"source_layer_2": metrics["source_layers"][1],
"pair_kind": metrics["pair_kind"],
"schedule": metrics["schedule"],
"num_prompts": metrics["num_prompts"],
"psnr": metrics["psnr"],
"ssim": metrics["ssim"],
"lpips": metrics["lpips"],
"rollout_psnr": metrics["rollout_psnr"],
"rollout_ssim": metrics["rollout_ssim"],
"rollout_lpips": metrics["rollout_lpips"],
"offline_final_val_flow_mse": metrics[
"offline_final_val_flow_mse"
],
"mean_generation_time_s": metrics["mean_generation_time_s"],
}
)
rows.sort(key=lambda row: float(row["lpips"]))
return rows
def write_summary(main_dir: Path, rows: list[dict[str, Any]]) -> None:
destination = main_dir / "summary.csv"
temporary = destination.with_suffix(".csv.tmp")
with temporary.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
writer.writeheader()
writer.writerows(rows)
os.replace(temporary, destination)
atomic_json(main_dir / "summary.json", rows)
def main() -> None:
args = parse_args()
main_dir = args.main_dir.resolve()
shard_dir = args.shard_dir.resolve()
sweep_dir = args.sweep_dir.resolve()
names = read_experiment_names(sweep_dir)
if len(names) != 43:
raise RuntimeError(f"Expected 43 experiments, found {len(names)}")
merge_pair_directories(main_dir, shard_dir)
rows = collect_rows(main_dir, names)
write_summary(main_dir, rows)
atomic_json(
main_dir / "manifest.json",
{
"status": "complete",
"architecture": "two_block_predictor",
"prompt_ids": PROMPT_IDS,
"experiments": names,
"rollout_schedule": "FPPP",
"rollout_definition": SCHEDULE,
"generation_seed_reset_per_prompt": 0,
"gpu_shards": {"2": 22, "4": 21},
"merged_shard_dir": str(shard_dir),
},
)
print(f"[merge] {len(rows)} pairs -> {main_dir / 'summary.csv'}", flush=True)
if __name__ == "__main__":
main()
|