diff --git a/configs/ismcts/default.yaml b/configs/ismcts/default.yaml index d76fc7b..4ed8153 100644 --- a/configs/ismcts/default.yaml +++ b/configs/ismcts/default.yaml @@ -36,6 +36,9 @@ training: replay_capacity: 100000 interleave_games: 8 interleave_max_batch: 64 + use_inference_server: true + inference_server_max_batch: 128 + inference_server_batch_timeout_ms: 10.0 optimization: learning_rate: 0.0003 grad_clip: 5.0 diff --git a/src/coolrl_lost_cities/games/classic/ismcts/config.py b/src/coolrl_lost_cities/games/classic/ismcts/config.py index 0b849fa..bafb0bf 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/config.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/config.py @@ -60,6 +60,9 @@ class TrainingConfig(StrictModel): interleave_max_batch: int = 64 num_workers: int = 1 worker_device: str = "cpu" + use_inference_server: bool = True + inference_server_max_batch: int = 128 + inference_server_batch_timeout_ms: float = 10.0 @field_validator( "games_per_iter", @@ -69,6 +72,7 @@ class TrainingConfig(StrictModel): "interleave_games", "interleave_max_batch", "num_workers", + "inference_server_max_batch", ) @classmethod def _positive_int(cls, value: int) -> int: @@ -76,6 +80,13 @@ class TrainingConfig(StrictModel): raise ValueError("must be positive") return value + @field_validator("inference_server_batch_timeout_ms") + @classmethod + def _positive_float(cls, value: float) -> float: + if value <= 0: + raise ValueError("must be positive") + return value + class IsMctsConfig(StrictModel): run: RunConfig = Field(default_factory=lambda: RunConfig(experiment_name="ismcts")) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/eval_worker.py b/src/coolrl_lost_cities/games/classic/ismcts/eval_worker.py index f4a3276..c9ee687 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/eval_worker.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/eval_worker.py @@ -13,22 +13,37 @@ from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import IsMctsConfig, config_from_dict +from .inference_server import InferenceClient +from .info_set import canonical_info_set_key from .mcts import IsMctsSearcher from .network import AlphaZeroNet +_INFERENCE_REQUEST_QUEUE: Any | None = None +_INFERENCE_RESPONSE_QUEUES: list[Any] | None = None + + +def init_eval_inference_queues(request_queue: Any, response_queues: list[Any]) -> None: + global _INFERENCE_REQUEST_QUEUE, _INFERENCE_RESPONSE_QUEUES + _INFERENCE_REQUEST_QUEUE = request_queue + _INFERENCE_RESPONSE_QUEUES = response_queues + @dataclass(frozen=True) class EvalWorkerBatch: worker_index: int config: dict[str, Any] game_config: dict[str, Any] - network_state: dict[str, Any] + network_state: dict[str, Any] | None mcts_config: dict[str, Any] opponent: str game_indices: list[int] seed: int device: str max_steps: int + tasks: list[tuple[str, int]] | None = None + use_inference_server: bool = False + request_queue: Any | None = None + response_queue: Any | None = None @dataclass @@ -41,6 +56,7 @@ class EvalWorkerResult: policy_turns: int play_actions: int timeouts: int + by_opponent: dict[str, dict[str, Any]] | None = None def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult: @@ -51,12 +67,30 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult: torch.set_num_threads(1) cfg: IsMctsConfig = config_from_dict(batch.config) game_config = LostCitiesConfig(**batch.game_config) - device = torch.device(batch.device) probe = GameState.new_game(game_config, seed=batch.seed) in_dim = input_dim(probe, cfg.encoding) - network = AlphaZeroNet.from_config(in_dim, probe.action_size, cfg).to(device) - network.load_state_dict(batch.network_state) - network.eval() + if batch.use_inference_server: + device = torch.device("cpu") + network = _NetworkShape(probe.action_size) + request_queue = batch.request_queue or _INFERENCE_REQUEST_QUEUE + response_queue = batch.response_queue + if response_queue is None and _INFERENCE_RESPONSE_QUEUES is not None: + response_queue = _INFERENCE_RESPONSE_QUEUES[batch.worker_index] + if request_queue is None or response_queue is None: + raise RuntimeError("inference server queues are required") + inference_client = InferenceClient( + batch.worker_index, + request_queue, + response_queue, + ) + else: + device = torch.device(batch.device) + network = AlphaZeroNet.from_config(in_dim, probe.action_size, cfg).to(device) + if batch.network_state is None: + raise RuntimeError("network_state is required without inference server") + network.load_state_dict(batch.network_state) + network.eval() + inference_client = None from .config import MctsConfig mcts_config = MctsConfig.model_validate(batch.mcts_config) @@ -67,11 +101,27 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult: policy_turns = 0 play_actions = 0 timeouts = 0 - for game_index in batch.game_indices: + tasks = batch.tasks or [(batch.opponent, game_index) for game_index in batch.game_indices] + by_opponent: dict[str, dict[str, Any]] = {} + for opponent_name, game_index in tasks: + bucket = by_opponent.setdefault( + opponent_name, + { + "score_diffs": [], + "wins0": 0, + "wins1": 0, + "draws": 0, + "policy_turns": 0, + "play_actions": 0, + "timeouts": 0, + }, + ) + game_policy_turns = 0 + game_play_actions = 0 policy_player = game_index % 2 opponents = [ - build_bot(batch.opponent, seed=batch.seed + game_index), - build_bot(batch.opponent, seed=batch.seed + game_index + 1), + build_bot(opponent_name, seed=batch.seed + game_index), + build_bot(opponent_name, seed=batch.seed + game_index + 1), ] state = GameState.new_game(game_config, seed=batch.seed + game_index) steps = 0 @@ -89,15 +139,21 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult: encoding=cfg.encoding, rng=random.Random(rng.randrange(2**31)), ) - visits = searcher.search(state, current) + visits = ( + _search_with_inference_server(searcher, state, current, inference_client) + if inference_client is not None + else searcher.search(state, current) + ) if visits: unified = max(visits, key=visits.get) else: unified = state.unified_legal_actions()[0] if state.phase == "card": policy_turns += 1 + game_policy_turns += 1 if unified % 2 == 0: play_actions += 1 + game_play_actions += 1 state.apply_unified_action(unified) else: action = opponents[current].act(state) @@ -105,14 +161,21 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult: steps += 1 if not terminated: timeouts += 1 + bucket["timeouts"] += 1 diff = float(state.score_diff(policy_player)) score_diffs.append(diff) + bucket["score_diffs"].append(diff) if diff > 0: wins0 += 1 + bucket["wins0"] += 1 elif diff < 0: wins1 += 1 + bucket["wins1"] += 1 else: draws += 1 + bucket["draws"] += 1 + bucket["policy_turns"] += game_policy_turns + bucket["play_actions"] += game_play_actions return EvalWorkerResult( worker_index=batch.worker_index, score_diffs=score_diffs, @@ -122,4 +185,47 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult: policy_turns=policy_turns, play_actions=play_actions, timeouts=timeouts, + by_opponent=by_opponent, ) + + +def _search_with_inference_server( + searcher: IsMctsSearcher, + state: GameState, + traverser: int, + inference_client: InferenceClient, +) -> dict[int, int]: + from .interleaved_self_play import _evaluate_global_batch + + root_key = canonical_info_set_key(state, state.current_player) + root = searcher.tree.get_or_create( + root_key, + player=state.current_player, + terminal=state.terminal, + ) + completed = 0 + sims = int(searcher.config.n_simulations) + while completed < sims: + quota = min(int(searcher.config.parallel_simulations), sims - completed) + pending = searcher.prepare_simulation_batch(state, traverser, quota) + if not pending: + break + jobs = [(_SearchProxy(searcher), item) for item in pending] + _evaluate_global_batch( + searcher.network, + jobs, + searcher.device, + inference_client=inference_client, + ) + completed += len(pending) + return {action: root.visits.get(action, 0) for action in state.unified_legal_actions()} + + +@dataclass +class _SearchProxy: + searcher: IsMctsSearcher + + +class _NetworkShape: + def __init__(self, action_size: int) -> None: + self.action_size = int(action_size) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py b/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py index a392330..19c32d9 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py @@ -11,7 +11,7 @@ from coolrl_lost_cities.games.classic.bots.registry import build_bot from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import IsMctsConfig, MctsConfig -from .eval_worker import EvalWorkerBatch, run_eval_worker +from .eval_worker import EvalWorkerBatch, init_eval_inference_queues, run_eval_worker from .mcts import IsMctsSearcher from .network import AlphaZeroNet @@ -174,6 +174,133 @@ def _evaluate_parallel( play_actions += res.play_actions timeouts += res.timeouts n = len(score_diffs) + return _evaluation_metrics( + score_diffs=score_diffs, + wins0=wins0, + wins1=wins1, + draws=draws, + policy_turns=policy_turns, + play_actions=play_actions, + timeouts=timeouts, + elapsed_seconds=time.perf_counter() - started, + n=n, + ) + + +def evaluate_opponents_with_mcts_parallel( + network: AlphaZeroNet, + game_config: LostCitiesConfig, + mcts_config: MctsConfig, + *, + config: IsMctsConfig, + opponents: tuple[str, ...], + games: int, + seed: int, + num_workers: int, + max_steps: int, + request_queue=None, + response_queues=None, +) -> dict[str, dict[str, float | int]]: + started = time.perf_counter() + tasks = [(opponent, game_index) for opponent in opponents for game_index in range(games)] + if not tasks: + return {} + effective_workers = min(max(1, int(num_workers)), len(tasks)) + tasks_per_worker = [tasks[i::effective_workers] for i in range(effective_workers)] + use_inference_server = request_queue is not None and response_queues is not None + cpu_state = ( + None + if use_inference_server + else {name: tensor.detach().cpu() for name, tensor in network.state_dict().items()} + ) + config_dict = config.to_dict() + game_snapshot = game_config.to_snapshot() + mcts_dict = mcts_config.model_dump(mode="json") + worker_device = str(config.training.worker_device) + batches = [ + EvalWorkerBatch( + worker_index=i, + config=config_dict, + game_config=game_snapshot, + network_state=cpu_state, + mcts_config=mcts_dict, + opponent=tasks_per_worker[i][0][0] if tasks_per_worker[i] else "", + game_indices=[], + seed=seed, + device=worker_device, + max_steps=max_steps, + tasks=tasks_per_worker[i], + use_inference_server=use_inference_server, + request_queue=None, + response_queue=None, + ) + for i in range(effective_workers) + ] + ctx = mp.get_context("spawn") + aggregate: dict[str, dict[str, object]] = { + opponent: { + "score_diffs": [], + "wins0": 0, + "wins1": 0, + "draws": 0, + "policy_turns": 0, + "play_actions": 0, + "timeouts": 0, + } + for opponent in opponents + } + executor_kwargs = ( + { + "initializer": init_eval_inference_queues, + "initargs": (request_queue, response_queues), + } + if use_inference_server + else {} + ) + with ProcessPoolExecutor( + max_workers=effective_workers, + mp_context=ctx, + **executor_kwargs, + ) as executor: + for result in executor.map(run_eval_worker, batches): + for opponent, bucket in (result.by_opponent or {}).items(): + dest = aggregate[opponent] + dest["score_diffs"].extend(bucket["score_diffs"]) + dest["wins0"] += int(bucket["wins0"]) + dest["wins1"] += int(bucket["wins1"]) + dest["draws"] += int(bucket["draws"]) + dest["policy_turns"] += int(bucket["policy_turns"]) + dest["play_actions"] += int(bucket["play_actions"]) + dest["timeouts"] += int(bucket["timeouts"]) + elapsed = time.perf_counter() - started + return { + opponent: _evaluation_metrics( + score_diffs=list(bucket["score_diffs"]), + wins0=int(bucket["wins0"]), + wins1=int(bucket["wins1"]), + draws=int(bucket["draws"]), + policy_turns=int(bucket["policy_turns"]), + play_actions=int(bucket["play_actions"]), + timeouts=int(bucket["timeouts"]), + elapsed_seconds=elapsed, + n=len(bucket["score_diffs"]), + ) + for opponent, bucket in aggregate.items() + } + + +def _evaluation_metrics( + *, + score_diffs: list[float], + wins0: int, + wins1: int, + draws: int, + policy_turns: int, + play_actions: int, + timeouts: int, + elapsed_seconds: float, + n: int, +) -> dict[str, float | int]: avg_diff = sum(score_diffs) / n if n else 0.0 return { "games": n, @@ -186,5 +313,5 @@ def _evaluate_parallel( "policy_turns": policy_turns, "play_action_rate": play_actions / policy_turns if policy_turns else 0.0, "max_step_timeouts": timeouts, - "elapsed_seconds": time.perf_counter() - started, + "elapsed_seconds": elapsed_seconds, } diff --git a/src/coolrl_lost_cities/games/classic/ismcts/inference_server.py b/src/coolrl_lost_cities/games/classic/ismcts/inference_server.py new file mode 100644 index 0000000..783a263 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/inference_server.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import itertools +import queue +import threading +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from .network import AlphaZeroNet + +InferenceRequest = tuple[int, int, np.ndarray, np.ndarray] | None +InferenceResponse = tuple[int, np.ndarray, np.ndarray] + + +class InferenceClient: + def __init__(self, worker_id: int, request_queue: Any, response_queue: Any) -> None: + self.worker_id = int(worker_id) + self.request_queue = request_queue + self.response_queue = response_queue + self._ids = itertools.count() + + def infer(self, infos: np.ndarray, masks: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + request_id = next(self._ids) + self.request_queue.put( + ( + self.worker_id, + request_id, + np.asarray(infos, dtype=np.float32), + np.asarray(masks, dtype=bool), + ) + ) + while True: + response_id, priors, values = self.response_queue.get() + if response_id == request_id: + return priors, values + raise RuntimeError( + f"inference response id mismatch: expected {request_id}, got {response_id}" + ) + + +@dataclass +class InferenceServer: + network: AlphaZeroNet + device: torch.device + request_queue: Any + response_queues: list[Any] + max_batch: int = 64 + batch_timeout_seconds: float = 0.001 + + def __post_init__(self) -> None: + self._thread: threading.Thread | None = None + self._stop = threading.Event() + self.forward_batches = 0 + self.forward_requests = 0 + self.forward_positions = 0 + + def start(self) -> None: + if self._thread is not None: + return + self.network.eval() + self._thread = threading.Thread(target=self._run, name="ismcts-inference-server") + self._thread.start() + + def stop(self) -> None: + self._stop.set() + self.request_queue.put(None) + if self._thread is not None: + self._thread.join() + self._thread = None + + def __enter__(self) -> InferenceServer: + self.start() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.stop() + + def _run(self) -> None: + while not self._stop.is_set(): + try: + first = self.request_queue.get(timeout=0.1) + except queue.Empty: + continue + if first is None: + break + batch: list[tuple[int, int, np.ndarray, np.ndarray]] = [first] + rows = _request_rows(first) + deadline = time.perf_counter() + self.batch_timeout_seconds + while rows < self.max_batch: + remaining = deadline - time.perf_counter() + if remaining <= 0: + break + try: + item = self.request_queue.get(timeout=remaining) + except queue.Empty: + break + if item is None: + self._stop.set() + break + batch.append(item) + rows += _request_rows(item) + self._serve(batch) + + def _serve(self, batch: list[tuple[int, int, np.ndarray, np.ndarray]]) -> None: + infos = np.concatenate([_ensure_2d(item[2]) for item in batch], axis=0) + masks = np.concatenate([_ensure_2d(item[3]) for item in batch], axis=0) + with torch.inference_mode(): + x = torch.as_tensor(infos, dtype=torch.float32, device=self.device) + legal = torch.as_tensor(masks, dtype=torch.bool, device=self.device) + logits, values = self.network(x, legal) + probs = torch.softmax(logits, dim=-1).masked_fill(~legal, 0.0) + normalizer = probs.sum(dim=-1, keepdim=True).clamp_min(1.0e-12) + priors = (probs / normalizer).detach().cpu().numpy() + values_np = values.detach().cpu().numpy() + cursor = 0 + for worker_id, request_id, request_infos, _request_masks in batch: + size = _ensure_2d(request_infos).shape[0] + self.response_queues[worker_id].put( + ( + request_id, + priors[cursor : cursor + size].astype(np.float32, copy=False), + values_np[cursor : cursor + size].astype(np.float32, copy=False), + ) + ) + cursor += size + self.forward_batches += 1 + self.forward_requests += len(batch) + self.forward_positions += int(infos.shape[0]) + + +def _ensure_2d(array: np.ndarray) -> np.ndarray: + array = np.asarray(array) + if array.ndim == 1: + return array[None, :] + return array + + +def _request_rows(item: tuple[int, int, np.ndarray, np.ndarray]) -> int: + return int(_ensure_2d(item[2]).shape[0]) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/interleaved_self_play.py b/src/coolrl_lost_cities/games/classic/ismcts/interleaved_self_play.py index 3e4b763..dc3445d 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/interleaved_self_play.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/interleaved_self_play.py @@ -10,6 +10,7 @@ from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import MctsConfig, TrainingConfig +from .inference_server import InferenceClient from .info_set import canonical_info_set_key from .mcts import IsMctsSearcher, PendingSimulation from .network import AlphaZeroNet @@ -55,6 +56,7 @@ def play_self_play_iteration( encoding=None, temperature: float = 1.0, max_steps: int = 10_000, + inference_client: InferenceClient | None = None, ) -> list[ReplaySample]: device = torch.device(device) completed: list[list[ReplaySample]] = [] @@ -102,7 +104,13 @@ def play_self_play_iteration( active = still_active if jobs: - _run_search_jobs(network, jobs, training_config.interleave_max_batch, device) + _run_search_jobs( + network, + jobs, + training_config.interleave_max_batch, + device, + inference_client=inference_client, + ) for job in jobs: _finish_decision(job, mcts_config, encoding, temperature) @@ -119,6 +127,8 @@ def _run_search_jobs( jobs: list[_SearchJob], max_batch: int, device: torch.device, + *, + inference_client: InferenceClient | None = None, ) -> None: while any(job.remaining > 0 for job in jobs): pending: list[tuple[_SearchJob, PendingSimulation]] = [] @@ -141,13 +151,15 @@ def _run_search_jobs( break if not pending: break - _evaluate_global_batch(network, pending, device) + _evaluate_global_batch(network, pending, device, inference_client=inference_client) def _evaluate_global_batch( network: AlphaZeroNet, pending: list[tuple[_SearchJob, PendingSimulation]], device: torch.device, + *, + inference_client: InferenceClient | None = None, ) -> None: network_pending = [(job, item) for job, item in pending if item.terminal_value is None] values_by_id: dict[int, float] = {} @@ -159,12 +171,17 @@ def _evaluate_global_batch( masks = np.stack( [item.legal_mask for _job, item in network_pending if item.legal_mask is not None] ) - with torch.inference_mode(): - x = torch.as_tensor(infos, dtype=torch.float32, device=device) - mask = torch.as_tensor(masks, dtype=torch.bool, device=device) - probs = network.policy_distribution(x, mask).detach().cpu().numpy() - _logits, values = network(x, mask) - values_np = values.detach().cpu().numpy() + if inference_client is None: + with torch.inference_mode(): + x = torch.as_tensor(infos, dtype=torch.float32, device=device) + mask = torch.as_tensor(masks, dtype=torch.bool, device=device) + logits, values = network(x, mask) + policy = torch.softmax(logits, dim=-1).masked_fill(~mask, 0.0) + normalizer = policy.sum(dim=-1, keepdim=True).clamp_min(1.0e-12) + probs = (policy / normalizer).detach().cpu().numpy() + values_np = values.detach().cpu().numpy() + else: + probs, values_np = inference_client.infer(infos, masks) for index, (_job, item) in enumerate(network_pending): priors_by_id[id(item)] = probs[index] values_by_id[id(item)] = float(values_np[index]) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/trainer.py b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py index 2f263b7..d2dfe4c 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/trainer.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py @@ -17,11 +17,12 @@ from coolrl_lost_cities.games.classic.deep_cfr.evaluate import evaluate_strategy from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import IsMctsConfig -from .evaluate import evaluate_with_mcts +from .evaluate import evaluate_opponents_with_mcts_parallel, evaluate_with_mcts +from .inference_server import InferenceServer from .interleaved_self_play import play_self_play_iteration from .network import AlphaZeroLogitsView, AlphaZeroNet from .replay_buffer import ReplayBuffer, ReplaySample -from .workers import SelfPlayWorkerBatch, run_self_play_worker +from .workers import SelfPlayWorkerBatch, init_inference_queues, run_self_play_worker @dataclass @@ -182,10 +183,14 @@ class IsMctsTrainer: base = total_games // effective_workers remainder = total_games % effective_workers per_worker = [base + (1 if i < remainder else 0) for i in range(effective_workers)] - # Move network state dict to CPU for cross-process transfer. - cpu_state = { - name: tensor.detach().cpu() for name, tensor in self.network.state_dict().items() - } + use_inference_server = bool(training_cfg.use_inference_server and effective_workers > 1) + # Move network state dict to CPU for cross-process transfer when workers + # run local inference. In server mode, workers never deserialize the model. + cpu_state = ( + None + if use_inference_server + else {name: tensor.detach().cpu() for name, tensor in self.network.state_dict().items()} + ) config_dict = self.config.to_dict() game_snapshot = self.game_config.to_snapshot() max_steps = self.config.evaluation.max_steps @@ -205,6 +210,7 @@ class IsMctsTrainer: temperature=temperature, max_steps=max_steps, device=worker_device, + use_inference_server=use_inference_server, ) ) samples: list[ReplaySample] = [] @@ -215,19 +221,60 @@ class IsMctsTrainer: flush=True, ) spawn_started = time.perf_counter() - with ProcessPoolExecutor(max_workers=effective_workers, mp_context=ctx) as executor: - futures = [executor.submit(run_self_play_worker, batch) for batch in batches] - print( - f" workers submitted in {time.perf_counter() - spawn_started:.1f}s, waiting for results...", - flush=True, + request_queue = ctx.Queue() if use_inference_server else None + response_queues = ( + [ctx.Queue() for _ in range(effective_workers)] if use_inference_server else None + ) + server = ( + InferenceServer( + self.network, + self.device, + request_queue, + response_queues, + max_batch=int(training_cfg.inference_server_max_batch), + batch_timeout_seconds=float(training_cfg.inference_server_batch_timeout_ms) + / 1000.0, ) - results = [] - for future in futures: - res = future.result() - results.append(res) + if use_inference_server + else None + ) + try: + if server is not None: + server.start() + executor_kwargs = ( + { + "initializer": init_inference_queues, + "initargs": (request_queue, response_queues), + } + if use_inference_server + else {} + ) + with ProcessPoolExecutor( + max_workers=effective_workers, + mp_context=ctx, + **executor_kwargs, + ) as executor: + futures = [executor.submit(run_self_play_worker, batch) for batch in batches] print( - f" worker {res.worker_index} done ({len(res.samples)} samples, " - f"elapsed {time.perf_counter() - spawn_started:.1f}s)", + f" workers submitted in {time.perf_counter() - spawn_started:.1f}s, waiting for results...", + flush=True, + ) + results = [] + for future in futures: + res = future.result() + results.append(res) + print( + f" worker {res.worker_index} done ({len(res.samples)} samples, " + f"elapsed {time.perf_counter() - spawn_started:.1f}s)", + flush=True, + ) + finally: + if server is not None: + server.stop() + print( + " inference server " + f"batches={server.forward_batches} requests={server.forward_requests} " + f"positions={server.forward_positions}", flush=True, ) for result in sorted(results, key=lambda item: item.worker_index): @@ -285,6 +332,87 @@ class IsMctsTrainer: return {} self.network.eval() results: dict[str, float | int] = {} + eval_workers = max( + 1, + int(self.config.evaluation.num_workers), + int(self.config.training.num_workers), + ) + use_eval_server = bool( + self.config.training.use_inference_server + and self.config.mcts.eval_with_mcts + and eval_workers > 1 + ) + if self.config.mcts.eval_with_mcts and eval_workers > 1: + print( + f" eval vs {', '.join(opponents)} (workers={eval_workers}, " + f"inference_server={use_eval_server})...", + flush=True, + ) + eval_mcts_cfg = self.config.mcts.model_copy() + if self.config.mcts.eval_n_simulations > 0: + eval_mcts_cfg = eval_mcts_cfg.model_copy( + update={"n_simulations": self.config.mcts.eval_n_simulations} + ) + ctx = mp.get_context("spawn") + request_queue = ctx.Queue() if use_eval_server else None + response_queues = ( + [ctx.Queue() for _ in range(eval_workers)] if use_eval_server else None + ) + server = ( + InferenceServer( + self.network, + self.device, + request_queue, + response_queues, + max_batch=int(self.config.training.inference_server_max_batch), + batch_timeout_seconds=float( + self.config.training.inference_server_batch_timeout_ms + ) + / 1000.0, + ) + if use_eval_server + else None + ) + started = time.perf_counter() + try: + if server is not None: + server.start() + eval_results = evaluate_opponents_with_mcts_parallel( + self.network, + self.game_config, + eval_mcts_cfg, + config=self.config, + opponents=tuple(opponents), + games=self.config.evaluation.games, + seed=self.config.run.seed + iteration * 1000, + num_workers=eval_workers, + max_steps=self.config.evaluation.max_steps, + request_queue=request_queue, + response_queues=response_queues, + ) + finally: + if server is not None: + server.stop() + print( + " eval inference server " + f"batches={server.forward_batches} requests={server.forward_requests} " + f"positions={server.forward_positions}", + flush=True, + ) + for opponent, result in eval_results.items(): + key = opponent.replace("-", "_") + for metric_key, value in result.items(): + results[f"eval/{key}/{metric_key}"] = value + par = result.get("play_action_rate", 0.0) + sd = result.get("avg_score_diff0", 0.0) + wr = result.get("win_rate0", 0.0) + print( + f" eval vs {opponent} done in {result.get('elapsed_seconds', 0.0):.1f}s " + f"PA={par:.2f} W={wr:.2f} S={sd:.1f}", + flush=True, + ) + print(f" eval opponents done in {time.perf_counter() - started:.1f}s", flush=True) + return results for opponent in opponents: print(f" eval vs {opponent}...", flush=True) opp_started = time.perf_counter() diff --git a/src/coolrl_lost_cities/games/classic/ismcts/workers.py b/src/coolrl_lost_cities/games/classic/ismcts/workers.py index 28d2ccc..4a80e47 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/workers.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/workers.py @@ -20,11 +20,20 @@ from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import IsMctsConfig, config_from_dict +from .inference_server import InferenceClient from .interleaved_self_play import play_self_play_iteration from .network import AlphaZeroNet from .replay_buffer import ReplaySample _TORCH_THREADS_CONFIGURED = False +_INFERENCE_REQUEST_QUEUE: Any | None = None +_INFERENCE_RESPONSE_QUEUES: list[Any] | None = None + + +def init_inference_queues(request_queue: Any, response_queues: list[Any]) -> None: + global _INFERENCE_REQUEST_QUEUE, _INFERENCE_RESPONSE_QUEUES + _INFERENCE_REQUEST_QUEUE = request_queue + _INFERENCE_RESPONSE_QUEUES = response_queues def _configure_worker_torch_threads() -> None: @@ -49,10 +58,13 @@ class SelfPlayWorkerBatch: base_seed: int config: dict[str, Any] game_config: dict[str, Any] - network_state: dict[str, Any] + network_state: dict[str, Any] | None temperature: float max_steps: int device: str + use_inference_server: bool = False + request_queue: Any | None = None + response_queue: Any | None = None @dataclass @@ -69,13 +81,31 @@ def run_self_play_worker(batch: SelfPlayWorkerBatch) -> SelfPlayWorkerResult: _configure_worker_torch_threads() cfg: IsMctsConfig = config_from_dict(batch.config) game_config = LostCitiesConfig(**batch.game_config) - device = torch.device(batch.device) probe = GameState.new_game(game_config, seed=batch.base_seed) in_dim = input_dim(probe, cfg.encoding) action_size = probe.action_size - network = AlphaZeroNet.from_config(in_dim, action_size, cfg).to(device) - network.load_state_dict(batch.network_state) - network.eval() + if batch.use_inference_server: + device = torch.device("cpu") + network = _NetworkShape(action_size) + request_queue = batch.request_queue or _INFERENCE_REQUEST_QUEUE + response_queue = batch.response_queue + if response_queue is None and _INFERENCE_RESPONSE_QUEUES is not None: + response_queue = _INFERENCE_RESPONSE_QUEUES[batch.worker_index] + if request_queue is None or response_queue is None: + raise RuntimeError("inference server queues are required") + inference_client = InferenceClient( + batch.worker_index, + request_queue, + response_queue, + ) + else: + device = torch.device(batch.device) + network = AlphaZeroNet.from_config(in_dim, action_size, cfg).to(device) + if batch.network_state is None: + raise RuntimeError("network_state is required without inference server") + network.load_state_dict(batch.network_state) + network.eval() + inference_client = None print( f" [worker {batch.worker_index}] init done in {_time.perf_counter() - _t0:.1f}s, self-play start", flush=True, @@ -96,9 +126,15 @@ def run_self_play_worker(batch: SelfPlayWorkerBatch) -> SelfPlayWorkerResult: encoding=cfg.encoding, temperature=batch.temperature, max_steps=batch.max_steps, + inference_client=inference_client, ) print( f" [worker {batch.worker_index}] self-play done in {_time.perf_counter() - _sp_t0:.1f}s ({len(samples)} samples)", flush=True, ) return SelfPlayWorkerResult(worker_index=batch.worker_index, samples=samples) + + +class _NetworkShape: + def __init__(self, action_size: int) -> None: + self.action_size = int(action_size) diff --git a/tests/games/classic/ismcts/test_ismcts.py b/tests/games/classic/ismcts/test_ismcts.py index a8b19da..a509e70 100644 --- a/tests/games/classic/ismcts/test_ismcts.py +++ b/tests/games/classic/ismcts/test_ismcts.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib.util +import multiprocessing as mp import random import sys from pathlib import Path @@ -16,6 +17,13 @@ from coolrl_lost_cities.games.classic.bots.heuristic_py import ( ) from coolrl_lost_cities.games.classic.ismcts.config import IsMctsConfig, MctsConfig from coolrl_lost_cities.games.classic.ismcts.determinization import sample_determinization +from coolrl_lost_cities.games.classic.ismcts.evaluate import ( + evaluate_opponents_with_mcts_parallel, +) +from coolrl_lost_cities.games.classic.ismcts.inference_server import ( + InferenceClient, + InferenceServer, +) from coolrl_lost_cities.games.classic.ismcts.info_set import canonical_info_set_key from coolrl_lost_cities.games.classic.ismcts.interleaved_self_play import ( play_self_play_iteration, @@ -450,3 +458,162 @@ def test_smoke_iter_with_batching(tmp_path) -> None: assert "mcts/avg_visit_entropy" in metrics assert "mcts/value_prediction_error" in metrics assert "mcts/policy_mcts_kl" in metrics + + +def test_inference_server_roundtrip_shapes() -> None: + state = GameState.new_game(mini_config(), seed=41) + net = AlphaZeroNet(input_dim(state), state.action_size, hidden_size=8, num_layers=1) + ctx = mp.get_context("spawn") + manager = ctx.Manager() + try: + request_queue = manager.Queue() + response_queues = [manager.Queue()] + server = InferenceServer( + net, + torch.device("cpu"), + request_queue, + response_queues, + max_batch=4, + ) + server.start() + try: + client = InferenceClient(0, request_queue, response_queues[0]) + infos = [] + masks = [] + for player in (0, 1, 0): + infos.append(encode_info_state(state, player)) + masks.append(np.asarray(state.unified_legal_mask(), dtype=bool)) + priors, values = client.infer(np.stack(infos), np.stack(masks)) + finally: + server.stop() + finally: + manager.shutdown() + + assert priors.shape == (3, state.action_size) + assert values.shape == (3,) + assert np.allclose(priors.sum(axis=1), 1.0) + assert np.all(priors[:, ~np.asarray(state.unified_legal_mask(), dtype=bool)] == 0.0) + + +def test_parallel_self_play_server_matches_sample_count(tmp_path) -> None: + base_config = { + "run": {"max_iterations": 1, "seed": 42, "device": "cpu"}, + "rules": { + "n_colors": 3, + "n_ranks": 5, + "n_handshakes": 1, + "hand_size": 4, + "bonus_threshold": 4, + }, + "network": {"hidden_size": 16, "num_layers": 1}, + "mcts": {"n_simulations": 2, "parallel_simulations": 2, "use_rollout_value": False}, + "training": { + "games_per_iter": 2, + "gradient_steps_per_iter": 1, + "batch_size": 8, + "num_workers": 2, + }, + "checkpoint": {"save_every": 0}, + "evaluation": {"eval_every": 0, "num_workers": 1, "max_steps": 80}, + } + off = IsMctsConfig.model_validate( + { + **base_config, + "training": {**base_config["training"], "use_inference_server": False}, + } + ) + on = IsMctsConfig.model_validate( + { + **base_config, + "training": {**base_config["training"], "use_inference_server": True}, + } + ) + torch.manual_seed(45) + off_trainer = IsMctsTrainer( + off, + off.rules.to_lost_cities_config(seed=off.run.seed), + run_dir=tmp_path / "off", + ) + torch.manual_seed(45) + on_trainer = IsMctsTrainer( + on, + on.rules.to_lost_cities_config(seed=on.run.seed), + run_dir=tmp_path / "on", + ) + + off_metrics = off_trainer.train()[0].to_dict() + on_metrics = on_trainer.train()[0].to_dict() + + assert on_metrics["samples/added"] == off_metrics["samples/added"] + + +def test_parallel_eval_server_matches_local_stats() -> None: + config = IsMctsConfig.model_validate( + { + "run": {"seed": 43, "device": "cpu"}, + "rules": { + "n_colors": 3, + "n_ranks": 5, + "n_handshakes": 1, + "hand_size": 4, + "bonus_threshold": 4, + }, + "network": {"hidden_size": 16, "num_layers": 1}, + "mcts": {"n_simulations": 2, "parallel_simulations": 2, "use_rollout_value": False}, + "training": {"num_workers": 2}, + "evaluation": {"games": 2, "opponents": ["random"], "num_workers": 2, "max_steps": 80}, + } + ) + game_config = config.rules.to_lost_cities_config(seed=config.run.seed) + state = GameState.new_game(game_config, seed=config.run.seed) + net = AlphaZeroNet(input_dim(state), state.action_size, hidden_size=16, num_layers=1) + local = evaluate_opponents_with_mcts_parallel( + net, + game_config, + config.mcts, + config=config, + opponents=("random",), + games=2, + seed=44, + num_workers=2, + max_steps=80, + ) + + ctx = mp.get_context("spawn") + manager = ctx.Manager() + try: + request_queue = manager.Queue() + response_queues = [manager.Queue() for _ in range(2)] + server = InferenceServer( + net, + torch.device("cpu"), + request_queue, + response_queues, + max_batch=8, + ) + server.start() + try: + server_result = evaluate_opponents_with_mcts_parallel( + net, + game_config, + config.mcts, + config=config, + opponents=("random",), + games=2, + seed=44, + num_workers=2, + max_steps=80, + request_queue=request_queue, + response_queues=response_queues, + ) + finally: + server.stop() + finally: + manager.shutdown() + + for key in ("games", "wins0", "wins1", "draws", "policy_turns", "max_step_timeouts"): + assert server_result["random"][key] == local["random"][key] + assert np.isclose( + server_result["random"]["avg_score_diff0"], + local["random"]["avg_score_diff0"], + )