| |
| """Find the best (workers x threads) for self-play on this machine.""" |
| import argparse, json, os, sys, time |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| import torch |
| from tinychess.config import REGISTRY |
| from tinychess.model import TinyChess |
| from tinychess.selfplay import SelfPlayConfig |
| from tinychess.parallel import play_games_parallel |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--games", type=int, default=128) |
| ap.add_argument("--max-plies", type=int, default=80) |
| ap.add_argument("--grid", default="1x1,8x2,16x2,32x2,48x2,64x2,96x2,48x4") |
| ap.add_argument("--out", default="runs/bench_parallel.json") |
| a = ap.parse_args() |
|
|
| m = TinyChess(REGISTRY["baseline"]()); m.eval() |
| res = {"cpu_count": os.cpu_count(), "games": a.games, "runs": []} |
| for spec in a.grid.split(","): |
| w, t = (int(x) for x in spec.split("x")) |
| cfg = SelfPlayConfig(n_games=a.games, max_plies=a.max_plies, steps=4, |
| adjudicate="material", seed=0) |
| t0 = time.time() |
| recs, st = play_games_parallel(m, None, cfg, n_workers=w, threads_per_worker=t) |
| dt = time.time() - t0 |
| moves = sum(len(r.moves) for r in recs) |
| row = {"workers": w, "threads": t, "sec": round(dt, 2), |
| "moves_per_s": round(moves / dt, 1), "games": len(recs)} |
| res["runs"].append(row) |
| print(f"{w:3d}w x {t}t {dt:7.2f}s {row['moves_per_s']:8.1f} moves/s", flush=True) |
| json.dump(res, open(a.out, "w"), indent=2) |
| best = max(res["runs"], key=lambda r: r["moves_per_s"]) |
| res["best"] = best |
| json.dump(res, open(a.out, "w"), indent=2) |
| print(f"\nBEST: {best['workers']}w x {best['threads']}t = {best['moves_per_s']} moves/s") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
