#!/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())