Improve bot first-move latency via inference prewarm and UI flow/perf polish

Author: Dieg0Code

src/api/app.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 
 from api.config import Settings, get_settings
+from api.deps.inference import preload_inference_service
 from api.error_handling import register_error_handlers
 from api.modules.auth.rate_limit import AuthRateLimiter
 from api.modules.auth.router import router as auth_router
@@ -24,11 +28,20 @@
 def create_app(settings: Settings | None = None) -> FastAPI:
     cfg = settings or get_settings()
     configure_logging(cfg)
+
+    @asynccontextmanager
+    async def _lifespan(_app: FastAPI) -> AsyncIterator[None]:
+        # Preload once at process startup to remove cold-start lag from the
+        # first model move users see in live matches.
+        preload_inference_service()
+        yield
+
     app = FastAPI(
         title=cfg.app_name,
         debug=cfg.app_debug,
         docs_url=cfg.docs_url,
         redoc_url=cfg.redoc_url,
+        lifespan=_lifespan,
     )
     register_error_handlers(app)
     if cfg.app_log_requests:

src/api/deps/inference.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from functools import lru_cache
 
 from fastapi import HTTPException, status
@@ -8,6 +9,8 @@
 from api.inference_artifacts import resolve_artifact_uri
 from inference.service import InferenceService
 
+logger = logging.getLogger(__name__)
+
 
 @lru_cache(maxsize=1)
 def _build_inference_service(
@@ -54,3 +57,18 @@ def get_inference_service_dep() -> InferenceService:
             status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
             detail=f"Inference service unavailable: {exc}",
         ) from exc
+
+
+def preload_inference_service() -> InferenceService | None:
+    """
+    Best-effort preload to avoid first-turn latency spikes in bot matches.
+    """
+    try:
+        service = get_inference_service_dep()
+        service.warmup(mode="fast")
+        return service
+    except HTTPException as exc:
+        logger.warning("Inference preload skipped.", extra={"detail": str(exc.detail)})
+    except Exception:  # pragma: no cover - defensive path for runtime-specific failures.
+        logger.exception("Inference preload crashed unexpectedly.")
+    return None

src/api/modules/gameplay/router.py

@@ -20,11 +20,11 @@
 from agents.heuristic import heuristic_move
 from agents.random_agent import random_move
 from api.config import Settings, get_settings
-from api.db.enums import GameStatus
+from api.db.enums import AgentType, GameStatus
 from api.db.models import Game, User
 from api.deps.auth import get_auth_service_dep, get_current_user_dep
 from api.deps.gameplay import get_gameplay_service_dep
-from api.deps.inference import get_inference_service_dep
+from api.deps.inference import get_inference_service_dep, preload_inference_service
 from api.modules.auth.service import AuthService
 from api.modules.gameplay.schemas import (
     GameCreateRequest,
@@ -285,6 +285,8 @@ async def post_game(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail=str(exc),
         ) from exc
+    if game.player1_agent == AgentType.MODEL or game.player2_agent == AgentType.MODEL:
+        preload_inference_service()
     return await _to_game_response(gameplay_service, game)
 
 

src/api/modules/matches/model_runtime.py

@@ -70,4 +70,17 @@ def resolve_model_inference_service(
     )
 
 
-__all__ = ["resolve_model_inference_service"]
+def prewarm_model_inference_service(
+    *,
+    version: ModelVersion,
+    base_service: InferenceService | None,
+) -> InferenceService:
+    service = resolve_model_inference_service(
+        version=version,
+        base_service=base_service,
+    )
+    service.warmup(mode="fast")
+    return service
+
+
+__all__ = ["prewarm_model_inference_service", "resolve_model_inference_service"]

src/api/modules/matches/service.py

@@ -16,7 +16,10 @@
     WinnerSide,
 )
 from api.db.models import BotProfile, Game, GameMove, User
-from api.modules.matches.model_runtime import resolve_model_inference_service
+from api.modules.matches.model_runtime import (
+    prewarm_model_inference_service,
+    resolve_model_inference_service,
+)
 from api.modules.matches.repository import MatchesRepository
 from api.modules.matches.schemas import MatchCreateRequest, MatchMoveRequest
 from api.modules.ranking.service import RankingService
@@ -52,6 +55,11 @@ async def create_match(self, payload: MatchCreateRequest, actor_user_id: UUID) -
                 player2_agent = profile.agent_type
                 if profile.agent_type == AgentType.MODEL and model_version_id is None:
                     model_version_id = player2.model_version_id
+                if profile.agent_type == AgentType.MODEL and model_version_id is not None:
+                    await self._prewarm_model_runtime(
+                        version_id=model_version_id,
+                        fallback_service=None,
+                    )
 
         game = Game(
             season_id=payload.season_id,
@@ -96,6 +104,11 @@ async def create_invitation(
 
         if opponent.is_bot:
             profile = await self._get_enabled_bot_profile(opponent.id)
+            if profile.agent_type == AgentType.MODEL and opponent.model_version_id is not None:
+                await self._prewarm_model_runtime(
+                    version_id=opponent.model_version_id,
+                    fallback_service=None,
+                )
             now = datetime.now(timezone.utc).replace(tzinfo=None)
             return await self.repository.create_game(
                 Game(
@@ -468,5 +481,22 @@ async def _resolve_model_bot_inference_service(
                 f"Model version '{version.name}' has no usable local artifact for inference."
             ) from None
 
+    async def _prewarm_model_runtime(
+        self,
+        *,
+        version_id: UUID,
+        fallback_service: InferenceService | None,
+    ) -> None:
+        version = await self.repository.get_model_version(version_id)
+        if version is None:
+            return
+        try:
+            prewarm_model_inference_service(
+                version=version,
+                base_service=fallback_service,
+            )
+        except (FileNotFoundError, RuntimeError, ValueError):
+            return
+
 
 

src/inference/service.py

@@ -154,6 +154,7 @@ def __init__(
                 "Inference initialization failed: neither torch checkpoint nor ONNX session is available."
             )
         self._mcts: MCTS | None = None
+        self._is_warmed_up = False
 
     @staticmethod
     def _resolve_device(device: str) -> str:
@@ -414,6 +415,15 @@ def _strong_result(self, board: AtaxxBoard) -> InferenceResult:
         value = float(value_tensor.item())
         return InferenceResult(move=move, action_idx=action_idx, value=value, mode="strong")
 
+    def warmup(self, *, mode: InferenceMode = "fast") -> None:
+        """
+        Prime inference runtime once so the first real bot turn avoids cold-start latency.
+        """
+        if self._is_warmed_up:
+            return
+        self.predict(board=AtaxxBoard(), mode=mode)
+        self._is_warmed_up = True
+
     def predict(self, board: AtaxxBoard, *, mode: InferenceMode = "fast") -> InferenceResult:
         if board.is_game_over():
             return InferenceResult(

tests/test_api_games.py

@@ -3,6 +3,7 @@
 import sys
 import unittest
 from pathlib import Path
+from unittest.mock import patch
 from uuid import UUID, uuid4
 
 from fastapi.testclient import TestClient
@@ -102,6 +103,34 @@ def test_create_game(self) -> None:
         self.assertEqual(payload["player2_agent"], "heuristic")
         self.assertIn("id", payload)
 
+    def test_create_game_prewarms_inference_when_model_agent_is_present(self) -> None:
+        client, _ = self._client_with_stub()
+        with patch("api.modules.gameplay.router.preload_inference_service") as preload:
+            response = client.post(
+                "/api/v1/gameplay/games",
+                json={
+                    "queue_type": "vs_ai",
+                    "player1_agent": "human",
+                    "player2_agent": "model",
+                },
+            )
+        self.assertEqual(response.status_code, 201)
+        preload.assert_called_once()
+
+    def test_create_game_skips_prewarm_for_non_model_agents(self) -> None:
+        client, _ = self._client_with_stub()
+        with patch("api.modules.gameplay.router.preload_inference_service") as preload:
+            response = client.post(
+                "/api/v1/gameplay/games",
+                json={
+                    "queue_type": "vs_ai",
+                    "player1_agent": "human",
+                    "player2_agent": "heuristic",
+                },
+            )
+        self.assertEqual(response.status_code, 201)
+        preload.assert_not_called()
+
     def test_get_game_by_id(self) -> None:
         client, _ = self._client_with_stub()
         created = client.post(

tests/test_api_inference_dep.py

@@ -3,13 +3,13 @@
 import sys
 import unittest
 from pathlib import Path
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 from fastapi import HTTPException
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
 
-from api.deps.inference import get_inference_service_dep
+from api.deps.inference import get_inference_service_dep, preload_inference_service
 
 
 class TestApiInferenceDep(unittest.TestCase):
@@ -22,7 +22,23 @@ def test_maps_module_not_found_to_http_503(self, *_: object) -> None:
         self.assertEqual(ctx.exception.status_code, 503)
         self.assertIn("Inference service unavailable", str(ctx.exception.detail))
 
+    @patch("api.deps.inference.get_inference_service_dep")
+    def test_preload_inference_service_warms_up_once(self, get_dep: Mock) -> None:
+        service = Mock()
+        get_dep.return_value = service
+
+        resolved = preload_inference_service()
+
+        self.assertIs(resolved, service)
+        service.warmup.assert_called_once_with(mode="fast")
+
+    @patch(
+        "api.deps.inference.get_inference_service_dep",
+        side_effect=HTTPException(status_code=503, detail="inference unavailable"),
+    )
+    def test_preload_inference_service_returns_none_when_unavailable(self, *_: object) -> None:
+        self.assertIsNone(preload_inference_service())
+
 
 if __name__ == "__main__":
     unittest.main()
-

tests/test_matches_service_model_inference.py

@@ -172,6 +172,50 @@ async def _run() -> None:
 
         asyncio.run(_run())
 
+    def test_create_invitation_prewarms_runtime_for_model_bot(self) -> None:
+        async def _run() -> None:
+            async with self.sessionmaker() as session:
+                service = MatchesService(repository=MatchesRepository(session=session))
+                human = User(username="human-c", email="human-c@example.com", is_active=True)
+                version = ModelVersion(
+                    name="ub_policy_spatial_v2",
+                    checkpoint_uri="checkpoints/policy_spatial_v2.ckpt",
+                    is_active=False,
+                )
+                bot = User(
+                    username="ub_bogonet_warmup",
+                    email="bogonet@example.com",
+                    is_active=True,
+                    is_bot=True,
+                    bot_kind=BotKind.MODEL,
+                    model_version_id=version.id,
+                )
+                session.add(human)
+                session.add(version)
+                session.add(bot)
+                session.add(
+                    BotProfile(
+                        user_id=bot.id,
+                        agent_type=AgentType.MODEL,
+                        model_mode="fast",
+                        enabled=True,
+                    )
+                )
+                await session.commit()
+
+                with patch("api.modules.matches.service.prewarm_model_inference_service") as prewarm:
+                    game = await service.create_invitation(
+                        actor_user_id=human.id,
+                        opponent_user_id=bot.id,
+                        rated=False,
+                    )
+
+                self.assertEqual(game.status, GameStatus.IN_PROGRESS)
+                self.assertEqual(game.player2_agent, AgentType.MODEL)
+                prewarm.assert_called_once()
+
+        asyncio.run(_run())
+
 
 if __name__ == "__main__":
     unittest.main()