Add ISMCTS inference server path

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

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"))

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)

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,
     }

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])

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])

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()

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)

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"],
+    )