| from __future__ import annotations |
| import argparse |
| import os |
| import subprocess |
| import sys |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--checkpoint", default="/root/chess/ckpt/model.pt") |
| ap.add_argument("--out-dir", default="/root/chess/selfplay_data") |
| ap.add_argument("--games", type=int, default=4000, help="total games across all workers") |
| ap.add_argument("--workers", type=int, default=28) |
| ap.add_argument("--movetime", type=int, default=200) |
| ap.add_argument("--temperature-plies", type=int, default=10) |
| args = ap.parse_args() |
| per_worker = max(1, args.games // args.workers) |
| procs = [] |
| logs = [] |
| for i in range(args.workers): |
| worker_out = os.path.join(args.out_dir, f"worker{i}") |
| log_path = os.path.join(args.out_dir, f"worker{i}.log") |
| os.makedirs(args.out_dir, exist_ok=True) |
| log_f = open(log_path, "w") |
| cmd = [ |
| sys.executable, |
| "selfplay_finetune.py", |
| "--checkpoint", |
| args.checkpoint, |
| "--out-dir", |
| worker_out, |
| "--games", |
| str(per_worker), |
| "--movetime", |
| str(args.movetime), |
| "--temperature-plies", |
| str(args.temperature_plies), |
| ] |
| p = subprocess.Popen( |
| cmd, |
| stdout=log_f, |
| stderr=subprocess.STDOUT, |
| cwd=os.path.dirname(os.path.abspath(__file__)) or ".", |
| ) |
| procs.append(p) |
| logs.append(log_f) |
| print(f"[launcher] started worker {i} (pid {p.pid}), {per_worker} games -> {worker_out}") |
| exit_codes = [] |
| for i, p in enumerate(procs): |
| code = p.wait() |
| exit_codes.append(code) |
| print(f"[launcher] worker {i} exited with code {code}") |
| for f in logs: |
| f.close() |
| failed = [i for i, c in enumerate(exit_codes) if c != 0] |
| if failed: |
| print( |
| f"[launcher] WARNING: workers failed: {failed} -- check worker*.log in {args.out_dir}" |
| ) |
| print( |
| f"[launcher] DONE: {args.workers} workers, {per_worker * args.workers} games total requested" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
