Enhance Dockerfile and application startup process with prewarm functionality

- Added environment variables for model warmup control in Dockerfile and README.md.
- Implemented startup prewarm logic in FastAPI application to load models before serving traffic, with options to skip prewarm.
- Introduced warm methods in user modeling and recommendation services to ensure models are ready for inference.
- Updated docker_build_assets.py to include model warming during the build process, with a flag to skip if necessary.
- Enhanced example environment file to include new configuration options for startup behavior.

Dockerfile

@@ -35,13 +35,18 @@ ENV OMP_NUM_THREADS=2 \
 
 ARG HF_TOKEN=
 ARG HUGGING_FACE_HUB_TOKEN=
+ARG DOCKER_BUILD_SKIP_LLM_WARM=
 ENV HF_TOKEN=${HF_TOKEN}
 ENV HUGGING_FACE_HUB_TOKEN=${HUGGING_FACE_HUB_TOKEN}
+ENV DOCKER_BUILD_SKIP_LLM_WARM=${DOCKER_BUILD_SKIP_LLM_WARM}
 
 COPY . .
 
 RUN python scripts/docker_build_assets.py
 
+ENV HF_HUB_OFFLINE=1 \
+    TRANSFORMERS_OFFLINE=1
+
 EXPOSE 7860
 
 CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1", "--timeout-keep-alive", "30"]

README.md

@@ -103,6 +103,10 @@ docker compose up --build -d
 
 Default compose maps **`7860:7860`**. The image bakes **`/code/data/business_catalog_embedded.jsonl`** and **`/code/data/task_a_reviews_embedded.jsonl`** at build time (or stubs if Yelp JSON is missing). Override with a bind mount, e.g. `./data:/code/data`, if you rebuild those files locally.
 
+The Docker image sets **`HF_HUB_OFFLINE=1`** and **`TRANSFORMERS_OFFLINE=1`** so the running container does not call the Hugging Face Hub (models must be fully cached during `docker build`). `scripts/docker_build_assets.py` runs **`warm_runtime_models()`** after data JSONL: one SentenceTransformer forward and one causal LM forward on CPU (set build-arg **`DOCKER_BUILD_SKIP_LLM_WARM=1`** if the builder OOMs).
+
+On startup, **`STARTUP_PREWARM`** (default **`user_modeling`**) loads that task’s embedder + optional RAG index + LLM before serving traffic (`all` = Task A and Task B, uses ~2× LLM RAM). Disable with **`SKIP_STARTUP_PREWARM=1`**.
+
 ### Smoke checks
 
 OpenAPI: `http://localhost:7860/docs` when using Docker (port **7860**). Local `uvicorn` defaults to **8080** unless you set `PORT`.

app/main.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import logging
 import os
+from contextlib import asynccontextmanager
 from pathlib import Path
 
 from dotenv import load_dotenv
@@ -16,6 +17,30 @@ logger = logging.getLogger(__name__)
 
 load_dotenv(Path(__file__).resolve().parents[1] / ".env")
 
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    if os.environ.get("SKIP_STARTUP_PREWARM", "").strip().lower() in ("1", "true", "yes"):
+        yield
+        return
+    mode = os.environ.get("STARTUP_PREWARM", "user_modeling").strip().lower()
+    logger.info("Startup prewarm (STARTUP_PREWARM=%s) …", mode)
+    try:
+        if mode in ("all", "both", "*"):
+            user_modeling_service().warm()
+            recommendation_service().warm()
+        elif mode in ("recommendation", "task_b", "task2", "2"):
+            recommendation_service().warm()
+        elif mode not in ("none", "off", "0", "skip"):
+            user_modeling_service().warm()
+        logger.info("Startup prewarm complete.")
+    except Exception:
+        logger.exception(
+            "Startup prewarm failed — first requests may be slow; set SKIP_STARTUP_PREWARM=1 to disable"
+        )
+    yield
+
+
 app = FastAPI(
     title="DSN X BCT — User modeling & Recommendation",
     description=(
@@ -23,6 +48,7 @@ app = FastAPI(
         "Task 2 (Recommendation): persona -> personalised ranked items."
     ),
     version="1.0",
+    lifespan=lifespan,
 )
 
 _um: UserModelingService | None = None

app/recommendation_pipeline.py

@@ -377,6 +377,11 @@ class RecommendationService:
         self.index.load()
         self._loaded = True
 
+    def warm(self) -> None:
+        self.ensure_catalog()
+        self._ensure_local_embedder()
+        self._ensure_local_rank_llm()
+
     def recommend(
         self,
         persona: str,

app/user_modeling.py

@@ -77,6 +77,12 @@ class UserModelingService:
             self._rag_index = TaskAReviewRagIndex(self._rag_path)
         return self._rag_index
 
+    def warm(self) -> None:
+        if self._rag_path.is_file():
+            _ = self._ensure_query_embedder()
+            self._rag().load()
+        self._ensure_local_llm()
+
     def _ensure_local_llm(self) -> tuple[Any, Any, Any]:
         if (
             self._local_llm_model is not None

env.example

@@ -8,4 +8,7 @@ TASK_B_LOCAL_EMBEDDING_MODEL=all-MiniLM-L6-v2
 TASK_B_LOCAL_LLM_MODEL=Qwen/Qwen2.5-1.5B-Instruct
 TASK_B_EMBEDDED_CATALOG=data/business_catalog_embedded.jsonl
 
+STARTUP_PREWARM=user_modeling
+SKIP_STARTUP_PREWARM=
+
 HF_TOKEN=

scripts/docker_build_assets.py

@@ -76,6 +76,50 @@ def prefetch_hub_files_only() -> None:
     print("docker_build_assets: Hub snapshots cached (LLM not loaded into RAM).")
 
 
+def warm_runtime_models() -> None:
+    print("docker_build_assets: warming models for runtime (CPU, one forward each)...")
+    import gc
+
+    emb_key = os.environ.get("TASK_B_LOCAL_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
+    from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]
+
+    st = SentenceTransformer(emb_key)
+    st.encode(["docker-build-warmup"], batch_size=1, show_progress_bar=False, convert_to_numpy=True)
+    del st
+    gc.collect()
+
+    if os.environ.get("DOCKER_BUILD_SKIP_LLM_WARM", "").strip().lower() in ("1", "true", "yes"):
+        print("docker_build_assets: DOCKER_BUILD_SKIP_LLM_WARM set — skipping causal LM warm.")
+        return
+
+    import torch  # type: ignore[import-untyped]
+    from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore[import-untyped]
+
+    llm_b = os.environ.get("TASK_B_LOCAL_LLM_MODEL", "Qwen/Qwen2.5-1.5B-Instruct").strip()
+    llm_a = os.environ.get("TASK_A_LOCAL_LLM_MODEL", "").strip()
+    to_load = [llm_b]
+    if llm_a and llm_a != llm_b:
+        to_load.append(llm_a)
+
+    for mid in to_load:
+        print(f"docker_build_assets: causal LM warm — {mid}")
+        tok = AutoTokenizer.from_pretrained(mid, trust_remote_code=True)
+        mdl = AutoModelForCausalLM.from_pretrained(
+            mid,
+            torch_dtype=torch.float32,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        mdl.eval()
+        with torch.no_grad():
+            batch = tok("warmup", return_tensors="pt")
+            mdl(**batch)
+        del mdl, tok
+        gc.collect()
+
+    print("docker_build_assets: model warm complete.")
+
+
 def yelp_review_path(rt: Path) -> Path:
     env_p = os.environ.get("YELP_REVIEW_JSON", "").strip()
     if env_p:
@@ -275,6 +319,8 @@ def main() -> None:
         build_stub_embedded(rt)
         build_stub_task_a_embedded(rt)
 
+    warm_runtime_models()
+
 
 if __name__ == "__main__":
     main()