| |
| """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() |
|
|
