benchmark.py

#!/usr/bin/env python3
"""Throughput sweep: measure games/min at several concurrency levels.

Usage: .venv/bin/python benchmark.py
Milestone evaluations are disabled during the sweep so pauses don't
contaminate the numbers. Training itself runs normally (the trajectories
still feed the trainer thread), so the result reflects real conditions.
"""
import time

from sockfish import config
config.EVAL_INTERVAL = 10 ** 9          # no milestone pauses mid-benchmark
from sockfish.training.pipeline import TrainingPipeline  # noqa: E402

SETTLE_S = 5
MEASURE_S = 15
LEVELS = [4, 8, 12, 16, 24, 32, 48, 64]


def main():
    pipeline = TrainingPipeline(config.find_stockfish())
    pipeline.start(LEVELS[0])
    results = {}
    try:
        for n in LEVELS:
            pipeline.set_concurrency(n)
            time.sleep(SETTLE_S)
            start = pipeline.status()["games_played"]
            t0 = time.time()
            time.sleep(MEASURE_S)
            elapsed = time.time() - t0
            delta = pipeline.status()["games_played"] - start
            gpm = delta / elapsed * 60
            results[n] = gpm
            print(f"{n:3d} workers -> {gpm:6,.0f} games/min")
    finally:
        pipeline.stop()
    best = max(results, key=results.get)
    print(f"\noptimal on this machine: {best} workers "
          f"({results[best]:,.0f} games/min)")


if __name__ == "__main__":
    main()