coorl-lost-cities/docs/plans/batched_traversal_inference_server.md

Plan: Batched Traversal Inference Server (Priority #5, Option A)

Status: Ready for implementation Owner: Codex Background: See docs/performance.md → "Batched Traversal Inference: Design Decision (2026-05-07)" for the A/B/C analysis and rationale. This plan implements Option A.

Goal

Replace per-worker single-state CPU policy forwards in Deep CFR traversal with a central GPU inference server that batches policy requests across all workers. Targets the dominant phase (traversal_seconds ≈ 60% of iteration time).

Non-goals

• Do not modify the Cython traversal recursion structure.
• Do not modify the game engine, replay buffer, or training loop math.
• Do not implement Option C (single-process vectorized traversal). Keep it as future work.
• Do not require Cython nogil-cleanliness.
• Do not change the public CLI surface.

Success criteria

1. With traversal.inference_backend: server enabled on configs/deep_cfr/default.yaml, end-to-end training produces eval-winrate trajectories indistinguishable (within seed noise) from the current local backend over at least 50 iterations on home hardware.
2. On home (6-core + RTX 3090), traversal_seconds decreases by at least 30% compared to the current run profile in docs/performance.md.
3. On remote (32-core + weak GPU), traversal_seconds decreases or stays within 10% of current; if it regresses more, fall back to local is the operator's choice — the plan still ships.
4. With traversal.inference_backend: local, behavior is byte-identical to current main.
5. All existing tests pass. New unit tests for the inference client/server round-trip pass.

Key files (current)

• src/coolrl_lost_cities/games/classic/deep_cfr/workers.py — run_traversal_worker_batch. Spawns CPU networks per worker (device = torch.device("cpu") at line 62). This is the worker entry point that must learn about the inference server.
• src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx — recursive traversal. Two policy call sites:
  • Lines 473–475: networks[player](x).squeeze(0).detach().cpu().numpy() — advantage-net forward.
  • Lines 550–552: self.strategy_network(x).squeeze(0).detach().cpu().numpy() — strategy-net forward.
  • Both are the seams that route through the inference client when the backend is server.
• src/coolrl_lost_cities/games/classic/deep_cfr/networks.py — DeepCFRMLP.
• src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py — main loop; spawns the worker pool and owns the trainer-side networks. Must also start/stop the inference server and push weights periodically.
• src/coolrl_lost_cities/games/classic/deep_cfr/config.py — config schema. Add traversal.inference_backend and an inference_server block.

New files

• src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py — server process: owns models on GPU, drains the request queue, runs batched forward, writes responses.
• src/coolrl_lost_cities/games/classic/deep_cfr/inference_client.py — client used inside workers: claims a request slot, posts the encoded state, waits for the response. Provides a forward(network_id, player, state) API.
• src/coolrl_lost_cities/games/classic/deep_cfr/inference_buffers.py — shared-memory tensor pool: pre-allocated [num_slots, input_dim] request buffer and [num_slots, action_size] response buffer, plus per-slot ready events and a free-slot stack.
• tests/games/classic/deep_cfr/test_inference_server.py — unit/integration tests for the request path.

Architecture

Process layout

main process (trainer)
 ├─ inference server process (GPU)
 │   - owns advantage networks (per-player), strategy network, league snapshots
 │   - request buffer (shared mem): [num_slots, input_dim] float32
 │   - response buffer (shared mem): [num_slots, action_size] float32
 │   - control queue (mp.Queue): control messages only
 │   - per-slot ready event (mp.Event[num_slots])
 ├─ traversal worker 1 ─┐
 ├─ ...                 ├─ inference clients post requests, wait on per-slot event
 └─ traversal worker N ─┘

Request protocol

A worker policy call becomes:

1. Client pops a free slot id from a shared mp.Queue-backed free-slot stack (small int).
2. Client writes the encoded state into request_buffer[slot_id] (shared memory; no pickle).
3. Client puts a RequestMessage(slot_id, network_kind, player, network_index) onto the request control queue. network_kind ∈ {ADVANTAGE, STRATEGY, LEAGUE}. Pickle cost is negligible (small struct of ints).
4. Server drains the request queue with a short batch window (batch_window_us, default 200μs) up to max_batch (default 256). Empty drain blocks on the queue with a small timeout.
5. Server stacks the requested rows from request_buffer, runs forward, writes outputs back to response_buffer[slot_id] for each request.
6. Server fires the per-slot ready event for each completed request.
7. Client wakes on its slot's event, reads response_buffer[slot_id], copies to a local numpy array, returns the slot to the free-slot stack.

Weight sync

• Trainer holds master weights. On a fixed cadence (inference_server.weight_sync_every iterations, default 1 — push every iter), trainer sends a WeightUpdateMessage containing state_dicts via torch.multiprocessing (auto-shares tensor storage; cheap once and copied into server's GPU model).
• Server applies load_state_dict and signals weight_sync_complete. Trainer waits before kicking off the next traversal iteration.
• League snapshots: passed via the same channel when the league updates. League list is small relative to per-iter cost.

Backpressure & lifecycle

• Free-slot stack size = num_slots (default max(64, 4 * num_workers * worker_chunk_size)). Workers block on slot allocation if all slots in flight; this is the natural backpressure.
• Server shutdown: trainer puts a Shutdown sentinel on the control queue at training end. Server drains, exits.
• Crash isolation: if the server dies, workers will hang on their slot events. Trainer monitors the server process; on death, it raises and tears down the pool. No silent corruption.

IPC choice

• torch.multiprocessing for weight passing (auto-shares tensors).
• multiprocessing.shared_memory.SharedMemory (numpy view) for request/response buffers — manual slot management. Pre-allocated once at startup; no per-call allocation.
• multiprocessing.Queue for control messages only (slot ids and small structs).
• multiprocessing.Event array for per-slot wakeups.

Rationale: see docs/performance.md § "IPC mechanism: multiprocessing + shared memory".

Config schema

Extend config.py:

@dataclass
class InferenceServerConfig:
    enabled: bool = False               # if False, behave like current code
    device: str = "cuda"                # server-side device
    num_slots: int | None = None        # None → auto: max(64, 4 * num_workers * worker_chunk_size)
    max_batch: int = 256
    batch_window_us: int = 200
    weight_sync_every: int = 1          # iterations
    use_amp: bool = False               # eval-only AMP for forward (no grad)

@dataclass
class TraversalConfig:
    # ... existing fields ...
    inference_backend: Literal["local", "server"] = "local"

In default.yaml, leave inference_backend: local for now. Add an explicit configs/deep_cfr/default_server.yaml variant that flips it on for benchmarking.

Worker integration

run_traversal_worker_batch (workers.py) currently:

• Builds CPU DeepCFRMLP instances and loads state_dicts.
• Passes them as positional networks into the Traversal Cython object.

When inference_backend == "server":

• Skip building local networks. Instead, build an InferenceClient bound to the shared buffers and queues that the trainer wires in via TraversalWorkerBatch.
• Pass a small NetworkProxy object with the same call signature as the current network: proxy(x: torch.Tensor) -> torch.Tensor. Internally it converts to numpy, calls client.forward(...), returns a torch tensor.
• Critically: the Cython traversal sites at traversal.pyx:473-475 and :550-552 should not need source changes if NetworkProxy is a callable returning a 2-D tensor. The existing .squeeze(0).detach().cpu().numpy() chain still works on the proxy's returned tensor (which can just be a CPU tensor wrapping the numpy result). Verify this; if Cython has typed assumptions that reject a Python proxy, fall back to a thin Python helper invoked from the .pyx instead.

Add to TraversalWorkerBatch:

• inference_handles: InferenceClientHandles | None — shared-memory names, queue handles, event arrays. None when inference_backend == "local".

Trainer integration

trainer.py:

1. On run start, if inference_backend == "server": instantiate InferenceServer (spawns process), build InferenceClientHandles, push initial weights, wait for weight_sync_complete.
2. Per iteration: before pool.starmap(run_traversal_worker_batch, ...), push fresh weights if iter % weight_sync_every == 0. Pass inference_handles into each TraversalWorkerBatch.
3. After traversal: same as today.
4. On run end: send shutdown sentinel; join the server process.

Implementation steps (ordered, each independently mergeable)

Step 1: shared-memory buffer module

• Create inference_buffers.py with InferenceBuffers class: pre-allocates request/response numpy arrays via SharedMemory, exposes attach(name) for child processes, has release() cleanup.
• Free-slot management: mp.Queue of slot ids, populated at startup with range(num_slots).
• Per-slot ready events: [mp.Event() for _ in range(num_slots)].
• Unit test: parent creates buffers, child attaches, writes a row, parent reads. Verify zero-copy semantics.

Step 2: inference client

• Create inference_client.py. InferenceClient.forward(network_kind, player, network_index, state_np: np.ndarray) -> np.ndarray:
  • Pop free slot, write state, post request, wait on event, copy response, return slot.
• Add a NetworkProxy callable that wraps client.forward to look like a torch.nn.Module for traversal call sites.
• Unit test: stub a server thread that echoes state*2; assert client gets the expected output.

Step 3: inference server

• Create inference_server.py. Spawn-friendly entry function run_inference_server(handles, model_config, control_queue, weight_queue).
• Owns models on device. Sets eval() and inference_mode().
• Main loop: drain request queue with batch_window_us deadline, group by (network_kind, network_index), run batched forward per group, scatter outputs to response buffer slots, fire events.
• Handles WeightUpdateMessage and Shutdown.
• Optional use_amp: wrap forward in torch.autocast when configured.
• Unit test: send N requests across multiple kinds; verify outputs match model(stacked_input).

Step 4: config + trainer wiring

• Extend config.py with InferenceServerConfig and TraversalConfig.inference_backend.
• Update trainer.py to start/stop server, push weights, attach handles to TraversalWorkerBatch.
• Update workers.py: when inference_backend == "server", build proxies instead of CPU networks.
• Add configs/deep_cfr/default_server.yaml flipping traversal.inference_backend: server.

Step 5: traversal call-site verification

• Run with inference_backend: server and a tiny config (1 worker, 1 traversal). Confirm the proxy is callable from traversal.pyx:473-475 and :550-552 without Cython type errors.
• If Cython rejects the proxy: refactor those two call sites to invoke a Python helper that takes (networks, player, info_state) and returns a numpy array. The helper picks local or server path by inspecting the object. This is a 2-line change per site.

Step 6: integration tests

• tests/games/classic/deep_cfr/test_inference_server.py:
  • End-to-end smoke: 1 iteration of training with server backend on CPU device; assert no crash, replay buffer populated.
  • Equivalence: same seed, same initial weights, both backends → assert traversal samples match within numerical tolerance for at least 1 iteration. (May require weight_sync_every = 1 and deterministic GPU forward; if exact match is fragile, accept distributional equivalence over 10 iterations.)

Step 7: benchmarking

• Add scripts/bench_inference_backend.py (mirrors scripts/profile_gpu_forward.py style):
  • Run 10 iterations on default.yaml with inference_backend=local.
  • Run 10 iterations on default_server.yaml.
  • Print iteration-mean and per-phase mean times.
• Document results in a new docs/performance.md experiment subsection (date-stamped).

Risks and mitigations

• Cython proxy incompatibility (Step 5). Mitigation: pre-prototype with a 5-line Python script that imports Traversal and passes a stub callable in place of a DeepCFRMLP. If it works, the rest of the plan stands.
• Weight sync staleness invalidates CFR. Mitigation: default weight_sync_every: 1 (every iteration). Only loosen after measuring.
• GPU contention with eval. Eval already runs on the trainer's device; the inference server adds another tenant on the same GPU. Mitigation: serialize eval and traversal phases (they already are sequential in the iteration loop). Document this in inference_server.py.
• Deadlock on server crash. Mitigation: trainer monitors server process.is_alive() between iterations and raises if dead. Workers waiting on events will be torn down with the pool.
• Slot exhaustion under bursty load. Mitigation: default num_slots = 4 * num_workers * worker_chunk_size to absorb bursts. Operators can tune.
• Remote (weak GPU) regression. If traversal_seconds regresses on remote, the operator can flip inference_backend: local and ship without it. The plan succeeds either way.

What this plan explicitly does not do

• Does not implement async client (workers stay sync-blocking on slot events). Async client is a future optimization if IPC round-trip dominates after measurement.
• Does not implement nogil threading or Option C. See docs/performance.md for the deferred path.
• Does not optimize encoding (policy_encoding_seconds). If post-A measurements show encoding dominates, that is a separate work item.

Out-of-scope follow-ups (do not start)

• Option C (vectorized traversal) — only re-evaluate if A's measurements show GPU forward is no longer on the critical path.
• TensorRT / torch.compile on the inference server's models — defer until A baseline numbers are collected.
• Batched encoding inside the server (compute encoding from raw game state on GPU) — only if policy_encoding_seconds becomes the new bottleneck.