feat: eliminate cold starts with build-time prefetch and startup warmup

Dockerfile

@@ -3,7 +3,7 @@ FROM python:3.11-slim
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
     PYTHONDONTWRITEBYTECODE=1 \
-    HF_HOME=/tmp/hf \
+    HF_HOME=/home/user/hf-cache \
     PORT=7860
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -12,9 +12,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     && rm -rf /var/lib/apt/lists/*
 
 RUN useradd -m -u 1000 user && \
-    mkdir -p /tmp/hf /tmp/p2 && \
-    chown -R user:user /tmp/hf /tmp/p2 && \
-    chmod 777 /tmp/hf /tmp/p2
+    mkdir -p /home/user/hf-cache /tmp/p2 && \
+    chown -R user:user /home/user/hf-cache /tmp/p2 && \
+    chmod 777 /tmp/p2
+
 USER user
 ENV HOME=/home/user \
     PATH=/home/user/.local/bin:$PATH
@@ -34,8 +35,20 @@ RUN --mount=type=secret,id=GITHUB_TOKEN,required=false,uid=1000 \
     pip install --no-cache-dir -r requirements.txt && \
     git config --global --unset-all url."https://${TOKEN}@github.com/".insteadOf 2>/dev/null || true
 
+# Copy registry and prefetch script before app code for optimal layer caching
+COPY --chown=user:user registry.py prefetch.py ./
+
+# Run build-time model prefetch to bake weights into HF_HOME
+RUN --mount=type=secret,id=GITHUB_TOKEN,required=false,uid=1000 \
+    python prefetch.py
+
+# Set offline mode for runtime only
+ENV HF_HUB_OFFLINE=1 \
+    TRANSFORMERS_OFFLINE=1
+
+# Copy application code
 COPY --chown=user:user . .
 
 EXPOSE 7860
 
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]

README.md

@@ -53,6 +53,7 @@ Set these in the Space settings under **Variables and secrets**:
 | `DISABLE_PROJECTS` | No | None | Comma-separated prefixes to skip loading (e.g. `project-2`). |
 | `MAX_BODY_SIZE_BYTES` | No | 10485760 (10 MB) | Max request body size before returning 413. |
 | `REQUEST_TIMEOUT_SECONDS` | No | 60.0 | Request timeout before returning 504. |
+| `WARMUP_TIMEOUT_SECONDS` | No | 300.0 | Warmup timeout per project in worker thread before timing out. |
 | `GITHUB_TOKEN` | Only for private repos | None | GitHub token passed as a build secret to clone private dependencies. |
 
 GitHub repository secrets (used by GitHub Actions):
@@ -71,14 +72,17 @@ Returns the status of the host and each registered project.
 {
   "status": "healthy",
   "loaded": {
-    "/project-1": "heartflow_backend.router:router",
-    "/project-2": "semantic_backend.router:router"
+    "/semantic": "semantic_backend.router:router"
+  },
+  "failed": {},
+  "warm_status": {
+    "/semantic": "warmed"
   },
-  "failed": {}
+  "total_warmup_seconds": 6.35
 }
 ```
 
-If any service fails to load, `status` changes to `"degraded"` or `"unhealthy"` and the error is listed under `"failed"`.
+If any service fails to load or warm up, `status` changes to `"degraded"` or `"unhealthy"` and the error is listed under `"failed"` or `"warm_status"`.
 
 Interactive documentation (`/docs`, `/redoc`, `/openapi.json`) is disabled.
 

main.py

@@ -2,6 +2,7 @@ import asyncio
 import importlib
 import logging
 import os
+import time
 from contextlib import asynccontextmanager
 from typing import Callable
 
@@ -16,9 +17,12 @@ logger = logging.getLogger("hf-backend")
 
 loaded_projects: dict[str, str] = {}
 failed_projects: dict[str, str] = {}
+warm_status: dict[str, str] = {}
+total_warmup_seconds: float = 0.0
 
 max_body_size = int(os.getenv("MAX_BODY_SIZE_BYTES", str(10 * 1024 * 1024)))
 request_timeout = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "60.0"))
+default_warmup_timeout = float(os.getenv("WARMUP_TIMEOUT_SECONDS", "300.0"))
 
 
 class RequestLimitMiddleware(BaseHTTPMiddleware):
@@ -52,11 +56,55 @@ class RequestLimitMiddleware(BaseHTTPMiddleware):
             return Response(content="Request Timeout", status_code=504)
 
 
+async def run_warmup_for_project(prefix: str, target: str, timeout: float):
+    logger.info(f"[{prefix}] Checking warmup for target '{target}'...")
+    try:
+        module_name = target.split(":")[0]
+        top_pkg_name = module_name.split(".")[0]
+        pkg = importlib.import_module(top_pkg_name)
+        warmup_fn = getattr(pkg, "warmup", None)
+        if warmup_fn is None or not callable(warmup_fn):
+            logger.info(f"[{prefix}] No callable warmup found on package '{top_pkg_name}'")
+            warm_status[prefix] = "no warmup"
+            return
+
+        logger.info(f"[{prefix}] Starting warmup with timeout {timeout}s...")
+        await asyncio.wait_for(asyncio.to_thread(warmup_fn), timeout=timeout)
+        warm_status[prefix] = "warmed"
+        logger.info(f"[{prefix}] Warmup completed successfully")
+    except asyncio.TimeoutError:
+        msg = f"failed: TimeoutError: warmup timed out after {timeout}s"
+        warm_status[prefix] = msg
+        logger.error(f"[{prefix}] Warmup timed out after {timeout}s")
+    except Exception as e:
+        msg = f"failed: {type(e).__name__}: {str(e)}"
+        warm_status[prefix] = msg
+        logger.error(f"[{prefix}] Warmup failed: {e}", exc_info=True)
+
+
+async def run_all_warmups():
+    global total_warmup_seconds
+    timeout = float(os.getenv("WARMUP_TIMEOUT_SECONDS", str(default_warmup_timeout)))
+    logger.info(f"Lifespan startup: beginning sequential project warmups (timeout={timeout}s)...")
+    t0 = time.perf_counter()
+    for prefix, target in list(loaded_projects.items()):
+        await run_warmup_for_project(prefix, target, timeout=timeout)
+    total_warmup_seconds = round(time.perf_counter() - t0, 3)
+    logger.info(f"Lifespan startup: all warmups finished in {total_warmup_seconds}s")
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    await run_all_warmups()
+    yield
+
+
 app = FastAPI(
     title="HF Backend",
     docs_url=None,
     redoc_url=None,
     openapi_url=None,
+    lifespan=lifespan,
 )
 
 app.add_middleware(RequestLimitMiddleware, max_size=max_body_size, timeout=request_timeout)
@@ -74,6 +122,10 @@ if allowed_origins:
 
 
 def load_routers():
+    loaded_projects.clear()
+    failed_projects.clear()
+    warm_status.clear()
+
     disabled_raw = [p.strip().strip("/") for p in os.getenv("DISABLE_PROJECTS", "").split(",") if p.strip()]
     disabled_set = {p for p in disabled_raw}
 
@@ -101,7 +153,10 @@ load_routers()
 
 @app.get("/health")
 def health():
-    if not failed_projects:
+    all_mounted = bool(loaded_projects) and len(failed_projects) == 0
+    all_warmed = bool(loaded_projects) and all(warm_status.get(p) == "warmed" for p in loaded_projects)
+
+    if all_mounted and all_warmed:
         overall_status = "healthy"
     elif loaded_projects:
         overall_status = "degraded"
@@ -112,10 +167,11 @@ def health():
         "status": overall_status,
         "loaded": loaded_projects,
         "failed": failed_projects,
+        "warm_status": warm_status,
+        "total_warmup_seconds": total_warmup_seconds,
     }
 
 
 @app.get("/")
 def root():
     return health()
-

prefetch.py

@@ -0,0 +1,34 @@
+import importlib
+import logging
+import sys
+
+from registry import PROJECT_REGISTRY
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("prefetch")
+
+
+def main():
+    logger.info("Starting model prefetch during build...")
+    for prefix, target in PROJECT_REGISTRY.items():
+        module_name = target.split(":")[0]
+        top_pkg_name = module_name.split(".")[0]
+        logger.info(f"[{prefix}] Checking warmup for package '{top_pkg_name}'...")
+        try:
+            pkg = importlib.import_module(top_pkg_name)
+            warmup_fn = getattr(pkg, "warmup", None)
+            if warmup_fn and callable(warmup_fn):
+                logger.info(f"[{prefix}] Running warmup() for '{top_pkg_name}'...")
+                warmup_fn()
+                logger.info(f"[{prefix}] Prefetch complete for '{top_pkg_name}'")
+            else:
+                logger.info(f"[{prefix}] No callable warmup() found for '{top_pkg_name}'")
+        except Exception as e:
+            logger.error(f"[{prefix}] Prefetch failed: {e}", exc_info=True)
+            sys.exit(1)
+
+    logger.info("All project models prefetched successfully.")
+
+
+if __name__ == "__main__":
+    main()

requirements.in

@@ -2,4 +2,4 @@ fastapi
 uvicorn[standard]
 pydantic
 starlette
-git+https://github.com/rachit-mangawa/Semantic-Comment-Analyze.git@v1.2#subdirectory=backend
+git+https://github.com/rachit-mangawa/Semantic-Comment-Analyze.git@v1.3#subdirectory=backend

requirements.txt

@@ -120,7 +120,7 @@ scipy==1.17.1
     # via
     #   scikit-learn
     #   sentence-transformers
-semantic-backend @ git+https://github.com/rachit-mangawa/Semantic-Comment-Analyze.git@480d2c3203f4773e805f8f58e63c1f1be315d6d1#subdirectory=backend
+semantic-backend @ git+https://github.com/rachit-mangawa/Semantic-Comment-Analyze.git@eb73c60a2d94fafa9d224d75d1ae8602a0fe372f#subdirectory=backend
     # via -r requirements.in
 sentence-transformers==6.2.0
     # via semantic-backend

tests/test_host.py

@@ -2,6 +2,7 @@ import asyncio
 import logging
 import os
 import sys
+import time
 import types
 import pytest
 from fastapi import APIRouter
@@ -32,6 +33,7 @@ def create_mock_modules():
 
     sem_router_mod.router = r
     sem.router = sem_router_mod
+    sem.warmup = lambda: True
     sys.modules["mock_semantic"] = sem
     sys.modules["mock_semantic.router"] = sem_router_mod
 
@@ -55,54 +57,76 @@ def test_host_with_mocks(monkeypatch, caplog):
 
     with caplog.at_level(logging.INFO):
         import main
-        client = TestClient(main.app)
+        with TestClient(main.app) as client:
+            assert "[/semantic] Successfully mounted at /semantic" in caplog.text
+            assert "[/broken] Failed to mount router" in caplog.text
+
+            res_health = client.get("/health")
+            assert res_health.status_code == 200
+            health_data = res_health.json()
+            assert health_data["status"] == "degraded"
+            assert "/semantic" in health_data["loaded"]
+            assert "/broken" in health_data["failed"]
+            assert health_data["warm_status"]["/semantic"] == "warmed"
+
+            # Verify GET / returns same info as /health
+            res_root = client.get("/")
+            assert res_root.status_code == 200
+            assert res_root.json() == health_data
+
+            res_sem = client.get("/semantic/ping")
+            assert res_sem.status_code == 200
+            assert res_sem.json() == {"project": "semantic", "status": "ok"}
+
+            for doc_url in ["/docs", "/redoc", "/openapi.json"]:
+                assert client.get(doc_url).status_code == 404
+
+            cors_allowed = client.get(
+                "/semantic/ping",
+                headers={"Origin": "https://allowed.com"}
+            )
+            assert cors_allowed.status_code == 200
+            assert cors_allowed.headers.get("access-control-allow-origin") == "https://allowed.com"
+
+            cors_rejected = client.get(
+                "/semantic/ping",
+                headers={"Origin": "https://disallowed.com"}
+            )
+            assert cors_rejected.headers.get("access-control-allow-origin") is None
+
+            oversized_data = {"text": "a" * 2000}
+            res_oversized = client.post("/semantic/echo", json=oversized_data)
+            assert res_oversized.status_code == 413
+
+            res_slow = client.get("/semantic/slow")
+            assert res_slow.status_code == 504
 
-    assert "[/semantic] Successfully mounted at /semantic" in caplog.text
-    assert "[/broken] Failed to mount router" in caplog.text
 
-    res_health = client.get("/health")
-    assert res_health.status_code == 200
-    health_data = res_health.json()
-    assert health_data["status"] == "degraded"
-    assert "/semantic" in health_data["loaded"]
-    assert "/broken" in health_data["failed"]
-
-    # Verify GET / returns same info as /health
-    res_root = client.get("/")
-    assert res_root.status_code == 200
-    assert res_root.json() == health_data
-
-    res_sem = client.get("/semantic/ping")
-    assert res_sem.status_code == 200
-    assert res_sem.json() == {"project": "semantic", "status": "ok"}
-
-    for doc_url in ["/docs", "/redoc", "/openapi.json"]:
-        assert client.get(doc_url).status_code == 404
+def test_kill_switch(monkeypatch):
+    create_mock_modules()
+    monkeypatch.setenv("DISABLE_PROJECTS", "semantic")
 
-    cors_allowed = client.get(
-        "/semantic/ping",
-        headers={"Origin": "https://allowed.com"}
-    )
-    assert cors_allowed.status_code == 200
-    assert cors_allowed.headers.get("access-control-allow-origin") == "https://allowed.com"
+    import registry
+    registry.PROJECT_REGISTRY = {
+        "/semantic": "mock_semantic.router:router",
+    }
 
-    cors_rejected = client.get(
-        "/semantic/ping",
-        headers={"Origin": "https://disallowed.com"}
-    )
-    assert cors_rejected.headers.get("access-control-allow-origin") is None
+    if "main" in sys.modules:
+        del sys.modules["main"]
 
-    oversized_data = {"text": "a" * 2000}
-    res_oversized = client.post("/semantic/echo", json=oversized_data)
-    assert res_oversized.status_code == 413
+    import main
+    with TestClient(main.app) as client:
+        res_health = client.get("/health")
+        health_data = res_health.json()
+        assert "/semantic" not in health_data["loaded"]
+        assert "/semantic" not in health_data["failed"]
 
-    res_slow = client.get("/semantic/slow")
-    assert res_slow.status_code == 504
+        assert client.get("/semantic/ping").status_code == 404
 
 
-def test_kill_switch(monkeypatch):
+def test_warmup_success(monkeypatch):
     create_mock_modules()
-    monkeypatch.setenv("DISABLE_PROJECTS", "semantic")
+    monkeypatch.setenv("DISABLE_PROJECTS", "")
 
     import registry
     registry.PROJECT_REGISTRY = {
@@ -113,39 +137,147 @@ def test_kill_switch(monkeypatch):
         del sys.modules["main"]
 
     import main
-    client = TestClient(main.app)
+    with TestClient(main.app) as client:
+        res_health = client.get("/health")
+        assert res_health.status_code == 200
+        data = res_health.json()
+        assert data["status"] == "healthy"
+        assert data["loaded"]["/semantic"] == "mock_semantic.router:router"
+        assert data["warm_status"]["/semantic"] == "warmed"
+        assert data["total_warmup_seconds"] >= 0
+
+        res_root = client.get("/")
+        assert res_root.status_code == 200
+        assert res_root.json() == data
+
+
+def test_warmup_exception(monkeypatch):
+    sem = types.ModuleType("err_semantic")
+    sem_router_mod = types.ModuleType("err_semantic.router")
+    r = APIRouter()
+
+    @r.get("/ping")
+    def ping():
+        return {"status": "ok"}
 
-    res_health = client.get("/health")
-    health_data = res_health.json()
-    assert "/semantic" not in health_data["loaded"]
-    assert "/semantic" not in health_data["failed"]
+    sem_router_mod.router = r
+    sem.router = sem_router_mod
 
-    assert client.get("/semantic/ping").status_code == 404
+    def fail_warmup():
+        raise RuntimeError("Weights file corrupted")
 
+    sem.warmup = fail_warmup
+    sys.modules["err_semantic"] = sem
+    sys.modules["err_semantic.router"] = sem_router_mod
 
-def test_healthy_status_when_all_load(monkeypatch):
-    create_mock_modules()
     monkeypatch.setenv("DISABLE_PROJECTS", "")
 
     import registry
     registry.PROJECT_REGISTRY = {
-        "/semantic": "mock_semantic.router:router",
+        "/err": "err_semantic.router:router",
     }
 
     if "main" in sys.modules:
         del sys.modules["main"]
 
     import main
-    client = TestClient(main.app)
+    with TestClient(main.app) as client:
+        res = client.get("/health")
+        assert res.status_code == 200
+        data = res.json()
+        assert data["status"] == "degraded"
+        assert data["loaded"]["/err"] == "err_semantic.router:router"
+        assert "failed: RuntimeError" in data["warm_status"]["/err"]
+
+        res_ping = client.get("/err/ping")
+        assert res_ping.status_code == 200
+
+
+def test_warmup_timeout(monkeypatch):
+    fast = types.ModuleType("fast_proj")
+    fast_router = types.ModuleType("fast_proj.router")
+    rf = APIRouter()
+    @rf.get("/ping")
+    def fast_ping():
+        return {"fast": "ok"}
+    fast_router.router = rf
+    fast.router = fast_router
+    fast.warmup = lambda: True
+    sys.modules["fast_proj"] = fast
+    sys.modules["fast_proj.router"] = fast_router
+
+    slow = types.ModuleType("slow_proj")
+    slow_router = types.ModuleType("slow_proj.router")
+    rs = APIRouter()
+    @rs.get("/ping")
+    def slow_ping():
+        return {"slow": "ok"}
+    slow_router.router = rs
+    slow.router = slow_router
+
+    def slow_warmup():
+        time.sleep(1.0)
+        return True
+
+    slow.warmup = slow_warmup
+    sys.modules["slow_proj"] = slow
+    sys.modules["slow_proj.router"] = slow_router
 
-    res_health = client.get("/health")
-    assert res_health.status_code == 200
-    assert res_health.json()["status"] == "healthy"
-    assert "/semantic" in res_health.json()["loaded"]
+    monkeypatch.setenv("DISABLE_PROJECTS", "")
+    monkeypatch.setenv("WARMUP_TIMEOUT_SECONDS", "0.1")
+
+    import registry
+    registry.PROJECT_REGISTRY = {
+        "/slow": "slow_proj.router:router",
+        "/fast": "fast_proj.router:router",
+    }
+
+    if "main" in sys.modules:
+        del sys.modules["main"]
 
-    res_root = client.get("/")
-    assert res_root.status_code == 200
-    assert res_root.json() == res_health.json()
+    import main
+    with TestClient(main.app) as client:
+        res = client.get("/health")
+        assert res.status_code == 200
+        data = res.json()
+        assert data["status"] == "degraded"
+        assert "failed: TimeoutError" in data["warm_status"]["/slow"]
+        assert data["warm_status"]["/fast"] == "warmed"
+        # Others still serve
+        assert client.get("/fast/ping").status_code == 200
+        assert client.get("/slow/ping").status_code == 200
+
+
+def test_project_without_warmup(monkeypatch):
+    nowarm = types.ModuleType("nowarm_proj")
+    nowarm_router = types.ModuleType("nowarm_proj.router")
+    r = APIRouter()
+    @r.get("/ping")
+    def ping():
+        return {"nowarm": "ok"}
+    nowarm_router.router = r
+    nowarm.router = nowarm_router
+    sys.modules["nowarm_proj"] = nowarm
+    sys.modules["nowarm_proj.router"] = nowarm_router
+
+    monkeypatch.setenv("DISABLE_PROJECTS", "")
+
+    import registry
+    registry.PROJECT_REGISTRY = {
+        "/nowarm": "nowarm_proj.router:router",
+    }
+
+    if "main" in sys.modules:
+        del sys.modules["main"]
+
+    import main
+    with TestClient(main.app) as client:
+        res = client.get("/health")
+        assert res.status_code == 200
+        data = res.json()
+        assert data["status"] == "degraded"
+        assert data["warm_status"]["/nowarm"] == "no warmup"
+        assert client.get("/nowarm/ping").status_code == 200
 
 
 def test_real_registry():
@@ -159,15 +291,26 @@ def test_real_registry():
         del sys.modules["main"]
 
     import main
-    client = TestClient(main.app)
-
-    res = client.get("/health")
-    assert res.status_code == 200
-    data = res.json()
-    assert data["status"] == "healthy"
-    assert "/semantic" in data["loaded"]
-    assert not data["failed"]
-
-    res_root = client.get("/")
-    assert res_root.status_code == 200
-    assert res_root.json() == data
+    with TestClient(main.app) as client:
+        res = client.get("/health")
+        assert res.status_code == 200
+        data = res.json()
+        assert data["status"] == "healthy"
+        assert "/semantic" in data["loaded"]
+        assert not data["failed"]
+        assert data["warm_status"]["/semantic"] == "warmed"
+        assert data["total_warmup_seconds"] >= 0
+
+        res_root = client.get("/")
+        assert res_root.status_code == 200
+        assert res_root.json() == data
+
+        # Check Semantic /health endpoint reports models_loaded
+        res_sem = client.get("/semantic/health")
+        assert res_sem.status_code == 200
+        sem_data = res_sem.json()
+        assert sem_data["status"] == "healthy"
+        assert sem_data["service"] == "semantic_backend"
+        assert sem_data["models_loaded"]["embedding"] is True
+        assert sem_data["models_loaded"]["classifier_head"] is True
+        assert sem_data["models_loaded"]["sentiment"] is True