Tsume Shogi Generator

generate_tsume_parallel.py

import asyncio
import json
import os
from pathlib import Path

CLI = str(Path("build/shogiboardq-cli").resolve())
ENGINE = "KomoringHeights  1.1.0 64ZEN2"  # Use the exact name from list-engines
PLIES = 3
TARGET = 100
WORKERS = 4
TIMEOUT_MS = 5000
MAX_ATTACK = 4
MAX_DEFEND = 1
ATTACK_RANGE = 3
ALLOW_FINAL_ALTERNATIVES = True
OUTPUT = Path(f"tsume-work/tsume_{PLIES}ply_candidates.txt")
ENV = dict(os.environ)
# Optional: use a separately prepared configuration on Linux
# ENV["XDG_CONFIG_HOME"] = "/path/to/tsume-config"

positions = {}
active = set()


def remember(sfen, pv):
    fields = sfen.split()
    if len(fields) != 4 or len(pv) != PLIES:
        raise ValueError("Invalid SFEN or PV length")
    key = " ".join(fields[:3])  # Board, turn, and hands; ignore the move number
    if key not in positions and len(positions) < TARGET:
        positions[key] = sfen + " moves " + " ".join(pv)


def checkpoint():
    temporary = OUTPUT.with_suffix(".tmp")
    text = "# CLI candidates (before independent collection validation)\n"
    text += "\n".join(positions.values()) + "\n"
    temporary.write_text(text, encoding="utf-8")
    temporary.replace(OUTPUT)


async def stop(process):
    if process.returncode is None:
        try:
            process.stdin.write(b"stop\n")
            await process.stdin.drain()
        except (BrokenPipeError, ConnectionResetError):
            pass


async def close_process(process):
    await stop(process)
    try:
        await asyncio.wait_for(process.wait(), timeout=15)
    except asyncio.TimeoutError:
        process.kill()
        await process.wait()
    active.discard(process)


async def worker(number):
    while len(positions) < TARGET:
        args = [
            CLI, "generate-tsume", "--engine", ENGINE,
            "--target-moves", str(PLIES),
            "--max-positions", str(min(100, TARGET - len(positions))),
            "--timeout-ms", str(TIMEOUT_MS),
            "--max-attack", str(MAX_ATTACK), "--max-defend", str(MAX_DEFEND),
            "--attack-range", str(ATTACK_RANGE), "--stdin-control",
        ]
        if not ALLOW_FINAL_ALTERNATIVES:
            args.append("--no-final-alternatives")
        log_path = OUTPUT.parent / f"worker-{PLIES}ply-{number}.jsonl"
        err_path = OUTPUT.parent / f"worker-{PLIES}ply-{number}.stderr.log"
        with log_path.open("a", encoding="utf-8") as log, err_path.open("ab") as err:
            process = await asyncio.create_subprocess_exec(
                *args, env=ENV,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE, stderr=err,
            )
            active.add(process)
            try:
                # Another worker may have reached the target during startup.
                if len(positions) >= TARGET:
                    await stop(process)
                async for raw in process.stdout:
                    line = raw.decode("utf-8")
                    log.write(line)
                    log.flush()
                    event = json.loads(line)
                    if event["event"] == "error":
                        raise RuntimeError(event["message"])
                    if event["event"] != "position":
                        continue
                    remember(event["sfen"], event["pv"])
                    checkpoint()
                    print(f"{len(positions)}/{TARGET}", flush=True)
                    if len(positions) == TARGET:
                        await asyncio.gather(*(stop(p) for p in list(active)))
                if await process.wait() != 0:
                    raise RuntimeError(f"CLI failed; see {err_path}")
            finally:
                await close_process(process)


async def main():
    if TARGET < 1 or WORKERS < 1:
        raise ValueError("TARGET and WORKERS must be positive")
    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    if OUTPUT.exists():
        for line in OUTPUT.read_text(encoding="utf-8").splitlines():
            if line and not line.startswith("#"):
                sfen, moves = line.split(" moves ", 1)
                remember(sfen, moves.split())
    tasks = [asyncio.create_task(worker(i)) for i in range(WORKERS)]
    try:
        await asyncio.gather(*tasks)
    finally:
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        await asyncio.gather(*(close_process(p) for p in list(active)))
        checkpoint()
    print(f"Collected: {OUTPUT} (validate the collection separately)")


if __name__ == "__main__":
    asyncio.run(main())