File size: 2,343 Bytes
23e02ba | 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 | #!/usr/bin/env python3
"""Create a node/chunk plan for full test-set CPU generation."""
from __future__ import annotations
import argparse
import json
import math
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/full_test"
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--chunk-size", type=int, required=True)
parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT))
parser.add_argument("--total-materials", type=int, default=0)
args = parser.parse_args()
if args.chunk_size <= 0:
raise ValueError("--chunk-size must be positive.")
total = args.total_materials or int(len(pd.read_pickle(TEST_PATH)))
output_root = Path(args.output_root)
output_root.mkdir(parents=True, exist_ok=True)
node_count = math.ceil(total / args.chunk_size)
rows = []
for node_index in range(node_count):
first_idx = node_index * args.chunk_size
num_materials = min(args.chunk_size, total - first_idx)
rows.append(
{
"node_index": node_index,
"first_idx": first_idx,
"num_materials": num_materials,
"output_dir": str(output_root / f"node_{node_index:05d}"),
}
)
payload = {
"model_path": str(MODEL_PATH),
"test_path": str(TEST_PATH),
"total_materials": total,
"chunk_size": args.chunk_size,
"node_count": node_count,
"output_root": str(output_root),
"rows": rows,
}
(output_root / "plan.json").write_text(json.dumps(payload, indent=2) + "\n")
with (output_root / "plan.tsv").open("w") as f:
f.write("node_index\tfirst_idx\tnum_materials\toutput_dir\n")
for row in rows:
f.write(
f"{row['node_index']}\t{row['first_idx']}\t{row['num_materials']}\t{row['output_dir']}\n"
)
print(json.dumps({k: payload[k] for k in ["total_materials", "chunk_size", "node_count", "output_root"]}, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
|