cpu_multinode_generation/sweep_single_node.py

#!/usr/bin/env python3
"""Sweep single-node CPU generation chunk sizes.

This script calls the existing cpu_generation_benchmark runner for material
counts 32, 48, 64, ... and stops when wall_seconds / num_materials first rises.
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path

import pandas as pd


ROOT = Path("/workspace/mp_20_pxrdnet")
MODEL_PATH = ROOT / "hydra/singlerun/2026-07-20/pxrdgen_raw512_run02"
TEST_PATH = ROOT / "data/mp_20_pxrdgen_xrd90_0p1_raw512/test.csv"
DEFAULT_OUTPUT_ROOT = (
    ROOT / "paper_results_pxrdgen_match_only/cpu_multinode/single_node_sweep"
)


def load_summary(output_dir: Path) -> dict:
    summary_path = output_dir / "timing" / "summary.json"
    if not summary_path.exists():
        raise FileNotFoundError(f"Missing summary: {summary_path}")
    return json.loads(summary_path.read_text())


def run_one(args: argparse.Namespace, n_materials: int) -> dict:
    output_dir = Path(args.output_root) / f"n_{n_materials:05d}"
    output_dir.mkdir(parents=True, exist_ok=True)
    command = [
        sys.executable,
        "-W",
        "ignore",
        str(ROOT / "cpu_generation_benchmark/run_cpu_generation_benchmark.py"),
        "--model-path",
        str(MODEL_PATH),
        "--output-dir",
        str(output_dir),
        "--first-idx",
        str(args.first_idx),
        "--num-materials",
        str(n_materials),
        "--num-starting-points",
        str(args.num_starting_points),
        "--workers",
        str(args.cpu_workers),
        "--torch-threads-per-worker",
        str(args.torch_threads_per_worker),
        "--num-gradient-steps",
        str(args.num_gradient_steps),
        "--n-step-each",
        str(args.n_step_each),
        "--progress-log-interval",
        str(args.progress_log_interval),
        "--progress-mininterval",
        str(args.progress_mininterval),
    ]
    if args.rebuild_limited_data:
        command.append("--rebuild-limited-data")

    env = os.environ.copy()
    env.update(
        {
            "PROJECT_ROOT": str(ROOT),
            "HYDRA_JOBS": str(ROOT / "hydra"),
            "WABDB_DIR": str(ROOT / "wabdb"),
            "PYTHONPATH": f"{ROOT}:{ROOT / 'scripts'}:{env.get('PYTHONPATH', '')}",
            "WANDB_MODE": "disabled",
            "CUDA_VISIBLE_DEVICES": "",
            "PIP_CACHE_DIR": env.get("PIP_CACHE_DIR", "/workspace/.cache/pip"),
            "TMPDIR": env.get("TMPDIR", "/workspace/tmp"),
        }
    )

    record = {
        "num_materials": n_materials,
        "output_dir": str(output_dir),
        "command": command,
        "started_at_unix": time.time(),
    }
    if args.dry_run:
        record.update(
            {
                "status": "dry_run",
                "wall_seconds": None,
                "avg_wall_seconds_per_material": None,
            }
        )
        return record

    log_path = output_dir / "benchmark.log"
    started = time.perf_counter()
    completed = subprocess.run(
        command,
        cwd=str(ROOT),
        env=env,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    elapsed = time.perf_counter() - started
    log_path.write_text(completed.stdout)
    record["returncode"] = completed.returncode
    record["elapsed_seconds_by_driver"] = elapsed
    if completed.returncode != 0:
        record.update({"status": "failed", "log_path": str(log_path)})
        return record

    summary = load_summary(output_dir)
    wall_seconds = float(summary["wall_seconds"])
    record.update(
        {
            "status": "ok",
            "summary": summary,
            "wall_seconds": wall_seconds,
            "avg_wall_seconds_per_material": wall_seconds / n_materials,
            "avg_wall_minutes_per_material": wall_seconds / n_materials / 60.0,
            "materials_per_wall_second": n_materials / wall_seconds,
            "log_path": str(log_path),
        }
    )
    return record


def write_outputs(output_root: Path, records: list[dict]) -> None:
    output_root.mkdir(parents=True, exist_ok=True)
    ok_records = [r for r in records if r.get("status") == "ok"]
    best = None
    if ok_records:
        best = min(ok_records, key=lambda r: r["avg_wall_seconds_per_material"])
    payload = {
        "model_path": str(MODEL_PATH),
        "test_path": str(TEST_PATH),
        "test_rows": int(len(pd.read_pickle(TEST_PATH))),
        "best_num_materials": best["num_materials"] if best else None,
        "best_avg_wall_seconds_per_material": best["avg_wall_seconds_per_material"]
        if best
        else None,
        "records": records,
    }
    (output_root / "sweep_summary.json").write_text(
        json.dumps(payload, indent=2) + "\n"
    )
    with (output_root / "sweep_summary.tsv").open("w") as f:
        f.write(
            "num_materials\tstatus\twall_seconds\tavg_wall_seconds_per_material\tavg_wall_minutes_per_material\tmaterials_per_wall_second\toutput_dir\n"
        )
        for r in records:
            f.write(
                f"{r.get('num_materials')}\t{r.get('status')}\t{r.get('wall_seconds')}\t"
                f"{r.get('avg_wall_seconds_per_material')}\t{r.get('avg_wall_minutes_per_material')}\t"
                f"{r.get('materials_per_wall_second')}\t{r.get('output_dir')}\n"
            )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--first-idx", type=int, default=int(os.environ.get("FIRST_IDX", 0)))
    parser.add_argument("--start-materials", type=int, default=int(os.environ.get("START_MATERIALS", 32)))
    parser.add_argument("--step-materials", type=int, default=int(os.environ.get("STEP_MATERIALS", 16)))
    parser.add_argument("--max-materials", type=int, default=int(os.environ.get("MAX_MATERIALS", 256)))
    parser.add_argument("--increase-tolerance", type=float, default=float(os.environ.get("INCREASE_TOLERANCE", 0.0)))
    parser.add_argument("--num-starting-points", type=int, default=int(os.environ.get("NUM_STARTING_POINTS", 1)))
    parser.add_argument("--cpu-workers", type=int, default=int(os.environ.get("CPU_WORKERS", 0)))
    parser.add_argument("--torch-threads-per-worker", type=int, default=int(os.environ.get("TORCH_THREADS_PER_WORKER", 0)))
    parser.add_argument("--num-gradient-steps", type=int, default=int(os.environ.get("NUM_GRADIENT_STEPS", 5000)))
    parser.add_argument("--n-step-each", type=int, default=int(os.environ.get("N_STEP_EACH", 100)))
    parser.add_argument("--progress-log-interval", type=int, default=int(os.environ.get("PROGRESS_LOG_INTERVAL", 500)))
    parser.add_argument("--progress-mininterval", type=float, default=float(os.environ.get("PROGRESS_MININTERVAL", 10)))
    parser.add_argument("--output-root", default=os.environ.get("OUTPUT_ROOT", str(DEFAULT_OUTPUT_ROOT)))
    parser.add_argument("--rebuild-limited-data", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    if args.start_materials <= 0 or args.step_materials <= 0:
        raise ValueError("start-materials and step-materials must be positive.")
    if args.max_materials < args.start_materials:
        raise ValueError("max-materials must be >= start-materials.")

    output_root = Path(args.output_root)
    records: list[dict] = []
    previous_avg = None
    for n_materials in range(
        args.start_materials, args.max_materials + 1, args.step_materials
    ):
        print(f"==> Running single-node CPU generation benchmark: N={n_materials}", flush=True)
        record = run_one(args, n_materials)
        records.append(record)
        write_outputs(output_root, records)
        print(json.dumps({k: record.get(k) for k in [
            "num_materials",
            "status",
            "wall_seconds",
            "avg_wall_seconds_per_material",
            "avg_wall_minutes_per_material",
            "materials_per_wall_second",
        ]}, indent=2), flush=True)

        if record.get("status") == "dry_run":
            previous_avg = previous_avg
            continue
        if record.get("status") != "ok":
            print("Stopping because benchmark failed.", flush=True)
            break
        current_avg = float(record["avg_wall_seconds_per_material"])
        if previous_avg is not None and current_avg > previous_avg * (1.0 + args.increase_tolerance):
            print(
                "Stopping because average time per material increased for the first time.",
                flush=True,
            )
            break
        previous_avg = current_avg

    return 0


if __name__ == "__main__":
    raise SystemExit(main())