evo2-sae serve: API under /api + optional static frontend mount

Move the API routes under an /api prefix (one APIRouter + include_router) and, when a built
frontend is configured (build_app(static_dir=...) or DASHBOARD_DIST env), mount it at / via
StaticFiles(html=True). This lets a single container serve both the dashboard and the API on one
origin: the frontend always calls /api/* (in dev via the Vite proxy, in prod from the same server).

The static mount is generic — it serves whatever dir it's pointed at and knows nothing about the
dashboard; the dashboard recipe (#1623) supplies the dir + the Docker build. With no frontend
configured the server is API-only and / 404s (never crashes). Startup already tolerates a load
failure (stays not-ready -> 503), so a frontend+API smoke needs no GPU/checkpoints.

Tests: re-point existing contract tests to /api/*, add SPA-index/asset served, API-reachable-under-
prefix, unknown-/api-is-404-not-SPA, and API-only-when-no-frontend.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Polina Binder <pbinder@nvidia.com>

Author: polinabinder1

interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py

@@ -15,21 +15,25 @@
 
 """FastAPI server over the Evo2SAE engine — the live backend the viz talks to.
 
-Endpoints: /health, /features, /annotate (per-base activations for a pasted
-sequence), /generate (autoregressive generation + optional SAE-feature clamp).
-This is a thin layer; all model work lives in `core.Evo2SAE`.
+Endpoints (under /api): /api/health, /api/features, /api/annotate (per-base activations for a
+pasted sequence), /api/generate (autoregressive generation + optional SAE-feature clamp). The
+/api prefix lets a prebuilt frontend be served from "/" on the same origin (single-container
+deploy); when no frontend is configured the server is API-only. This is a thin layer; all model
+work lives in `core.Evo2SAE`.
 """
 
 from __future__ import annotations
 
 import logging
 import os
 from contextlib import asynccontextmanager
+from pathlib import Path
 from typing import Optional
 
 import anyio
-from fastapi import FastAPI, HTTPException, Response
+from fastapi import APIRouter, FastAPI, HTTPException, Response
 from fastapi.responses import JSONResponse
+from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 
 from . import core
@@ -39,6 +43,19 @@
 logger = logging.getLogger("evo2_sae_infer.server")
 
 
+def _resolve_static_dir(static_dir: Optional[str]) -> Optional[str]:
+    """Directory of a prebuilt frontend to serve at ``/``, or None for an API-only server.
+
+    Explicit ``static_dir`` arg wins, else the ``DASHBOARD_DIST`` env var. Returns None unless the
+    path is an existing directory — so a server with no built frontend (dev, or an image built
+    without one) just serves the API and ``/`` 404s, instead of crashing. This layer is generic:
+    it serves whatever static dir it's pointed at and knows nothing about the dashboard (the
+    dashboard recipe supplies the dir via DASHBOARD_DIST / the Docker build).
+    """
+    cand = static_dir or os.getenv("DASHBOARD_DIST")
+    return cand if (cand and Path(cand).is_dir()) else None
+
+
 class AnnotateRequest(BaseModel):
     """Request body for /annotate (top-k feature scan or an explicit feature pick)."""
 
@@ -70,8 +87,14 @@ class GenerateRequest(BaseModel):
     compare_baseline: bool = False
 
 
-def build_app(engine: Evo2SAE) -> FastAPI:
-    """Build the FastAPI app; the engine is loaded once in the lifespan handler."""
+def build_app(engine: Evo2SAE, static_dir: Optional[str] = None) -> FastAPI:
+    """Build the FastAPI app; the engine is loaded once in the lifespan handler.
+
+    API routes live under ``/api`` (so the dashboard and the API can be served from one origin:
+    the frontend always calls ``/api/*``, in dev via the Vite proxy and in production from the
+    same server). If a built frontend is found (``static_dir`` / ``DASHBOARD_DIST``), it is mounted
+    at ``/`` so a single container serves both the UI and the API; otherwise the server is API-only.
+    """
     # One GPU (the engine serializes model calls with a lock), so cap how many sync requests run
     # at once: excess requests wait for a worker instead of piling up dozens of parked threads.
     # NOTE: generation is bounded only by the context window now, so a single /generate can run
@@ -113,7 +136,10 @@ def _require_ready():
         if not engine.ready:
             raise HTTPException(503, "Backend not ready")
 
-    @app.get("/health")
+    # All endpoints under /api (mounted below) so the SPA at "/" never collides with them.
+    api = APIRouter()
+
+    @api.get("/health")
     def health(response: Response):
         if not engine.ready:
             response.status_code = 503  # readiness probes shed this pod until load finishes (body still informative)
@@ -128,7 +154,7 @@ def health(response: Response):
             "max_seq_len": engine.max_seq_len,  # context budget — UI caps generation length to this
         }
 
-    @app.get("/features")
+    @api.get("/features")
     def features():
         _require_ready()
         rows = [
@@ -137,7 +163,7 @@ def features():
         rows.sort(key=lambda r: r["id"])
         return rows
 
-    @app.post("/annotate")
+    @api.post("/annotate")
     def annotate(req: AnnotateRequest):
         _require_ready()
         try:
@@ -183,7 +209,7 @@ def annotate(req: AnnotateRequest):
             "features": feats,
         }
 
-    @app.post("/generate")
+    @api.post("/generate")
     def generate(req: GenerateRequest):
         _require_ready()
         try:
@@ -200,4 +226,16 @@ def generate(req: GenerateRequest):
         except ValueError as e:
             raise HTTPException(413 if "too long" in str(e) else 400, str(e))
 
+    app.include_router(api, prefix="/api")
+
+    # Serve a prebuilt frontend at "/" when present, so one container serves UI + API. Mounted
+    # AFTER the API router, so /api/* always resolves to the API; unknown /api/* paths 404 here
+    # (StaticFiles only serves index.html for "/", not as a SPA catch-all — the UI is tabs at "/").
+    resolved = _resolve_static_dir(static_dir)
+    if resolved:
+        app.mount("/", StaticFiles(directory=resolved, html=True), name="dashboard")
+        logger.info("serving dashboard from %s", resolved)
+    else:
+        logger.info("no frontend mounted (API-only)")
+
     return app

interpretability/sparse_autoencoders/recipes/evo2/tests/test_server.py

@@ -34,8 +34,25 @@ def client(fake_engine):
         yield c
 
 
+@pytest.fixture
+def dist_dir(tmp_path):
+    """A minimal stand-in for a built frontend (feature_explorer/dist) — index + one asset."""
+    d = tmp_path / "dist"
+    (d / "assets").mkdir(parents=True)
+    (d / "index.html").write_text('<!doctype html><html><body><div id="root"></div></body></html>')
+    (d / "assets" / "app.js").write_text("console.log('dashboard')")
+    return d
+
+
+@pytest.fixture
+def static_client(fake_engine, dist_dir):
+    """A client for the single-container shape: API under /api + the frontend served at /."""
+    with TestClient(build_app(fake_engine, static_dir=str(dist_dir))) as c:
+        yield c
+
+
 def test_health(client):
-    r = client.get("/health")
+    r = client.get("/api/health")
     assert r.status_code == 200  # 200 only when ready
     b = r.json()
     assert b["ready"] is True and b["layer"] == 19
@@ -44,72 +61,74 @@ def test_health(client):
 
 def test_annotate_rejects_too_long(client, fake_engine):
     seq = "A" * (fake_engine.max_seq_len + 1)  # exceeds the context budget
-    assert client.post("/annotate", json={"sequence": seq}).status_code == 413
+    assert client.post("/api/annotate", json={"sequence": seq}).status_code == 413
 
 
 def test_features(client):
-    rows = client.get("/features").json()
+    rows = client.get("/api/features").json()
     assert {"id", "label", "natural_peak"} <= set(rows[0])
 
 
 def test_annotate_returns_per_base_activations(client):
-    b = client.post("/annotate", json={"sequence": "ACGTACGT", "organism": "None (raw DNA)"}).json()
+    b = client.post("/api/annotate", json={"sequence": "ACGTACGT", "organism": "None (raw DNA)"}).json()
     assert {"sequence", "features", "bases", "tag_len", "layer", "n_tokens"} <= set(b)
     assert b["features"][0]["activations"]  # the per-base track the viz plots
 
 
 def test_annotate_rejects_non_dna(client):
-    assert client.post("/annotate", json={"sequence": "ZZZZ"}).status_code == 400
+    assert client.post("/api/annotate", json={"sequence": "ZZZZ"}).status_code == 400
 
 
 def test_annotate_pick_mode(client):
-    b = client.post("/annotate", json={"sequence": "ACGTACGT", "mode": "pick", "feature_ids": [1]}).json()
+    b = client.post("/api/annotate", json={"sequence": "ACGTACGT", "mode": "pick", "feature_ids": [1]}).json()
     assert [f["feature_id"] for f in b["features"]] == [1]
     assert b["features"][0]["activations"]  # per-base track returned for the picked feature
 
 
 def test_annotate_pick_requires_ids(client):
-    assert client.post("/annotate", json={"sequence": "ACGT", "mode": "pick"}).status_code == 400
+    assert client.post("/api/annotate", json={"sequence": "ACGT", "mode": "pick"}).status_code == 400
 
 
 def test_annotate_pick_rejects_out_of_range_id(client, fake_engine):
     # user-supplied pick ids: an over-range id must 400 (not 500/IndexError), a negative one must
     # 400 (not silently index the wrong feature via torch negative-indexing).
-    over = client.post("/annotate", json={"sequence": "ACGT", "mode": "pick", "feature_ids": [fake_engine.n_features]})
+    over = client.post(
+        "/api/annotate", json={"sequence": "ACGT", "mode": "pick", "feature_ids": [fake_engine.n_features]}
+    )
     assert over.status_code == 400
-    neg = client.post("/annotate", json={"sequence": "ACGT", "mode": "pick", "feature_ids": [-1]})
+    neg = client.post("/api/annotate", json={"sequence": "ACGT", "mode": "pick", "feature_ids": [-1]})
     assert neg.status_code == 400
 
 
 def test_annotate_rejects_invalid_mode(client):
-    assert client.post("/annotate", json={"sequence": "ACGT", "mode": "bogus"}).status_code == 400
+    assert client.post("/api/annotate", json={"sequence": "ACGT", "mode": "bogus"}).status_code == 400
 
 
 def test_annotate_clamps_k_into_range(client, fake_engine):
-    client.post("/annotate", json={"sequence": "ACGT", "k": 999})
+    client.post("/api/annotate", json={"sequence": "ACGT", "k": 999})
     assert fake_engine.last_k == 64  # upper bound
-    client.post("/annotate", json={"sequence": "ACGT", "k": 0})
+    client.post("/api/annotate", json={"sequence": "ACGT", "k": 0})
     assert fake_engine.last_k == 1  # lower bound
 
 
 def test_generate_returns_sequence(client):
-    b = client.post("/generate", json={"prompt": "ACGT", "organism": "None (raw DNA)"}).json()
+    b = client.post("/api/generate", json={"prompt": "ACGT", "organism": "None (raw DNA)"}).json()
     assert b["generation"]["sequence"]
 
 
 def test_generate_rejects_out_of_range_feature(client):
-    r = client.post("/generate", json={"prompt": "ACGT", "features": [{"feature_id": 999}]})
+    r = client.post("/api/generate", json={"prompt": "ACGT", "features": [{"feature_id": 999}]})
     assert r.status_code == 400  # the wedge guard, surfaced to the client
 
 
 def test_generate_rejects_too_long(client, fake_engine):
     seq = "A" * (fake_engine.max_seq_len + 1)  # exceeds the context budget -> 413 (parity w/ annotate)
-    assert client.post("/generate", json={"prompt": seq}).status_code == 413
+    assert client.post("/api/generate", json={"prompt": seq}).status_code == 413
 
 
 def test_generate_compare_baseline(client):
     b = client.post(
-        "/generate",
+        "/api/generate",
         json={"prompt": "ACGT", "features": [{"feature_id": 0, "strength": 5.0}], "compare_baseline": True},
     ).json()
     assert b["baseline"]["sequence"]  # unsteered comparison returned alongside the steered one
@@ -118,14 +137,54 @@ def test_generate_compare_baseline(client):
 def test_rejects_oversized_body(monkeypatch, fake_engine):
     monkeypatch.setenv("MAX_BODY_BYTES", "100")
     with TestClient(build_app(fake_engine)) as c:
-        assert c.post("/annotate", json={"sequence": "ACGT" * 100}).status_code == 413
+        assert c.post("/api/annotate", json={"sequence": "ACGT" * 100}).status_code == 413
 
 
 def test_endpoints_503_until_ready(fake_engine):
     fake_engine.ready = False
     fake_engine.load = lambda: None  # startup leaves it not-ready
     with TestClient(build_app(fake_engine)) as c:
-        assert c.get("/health").status_code == 503  # readiness probe sheds the pod
-        assert c.get("/features").status_code == 503
-        assert c.post("/annotate", json={"sequence": "ACGT"}).status_code == 503
-        assert c.post("/generate", json={"prompt": "ACGT", "organism": "None (raw DNA)"}).status_code == 503
+        assert c.get("/api/health").status_code == 503  # readiness probe sheds the pod
+        assert c.get("/api/features").status_code == 503
+        assert c.post("/api/annotate", json={"sequence": "ACGT"}).status_code == 503
+        assert c.post("/api/generate", json={"prompt": "ACGT", "organism": "None (raw DNA)"}).status_code == 503
+
+
+# ---------------------------------------------------------- single-container: SPA at / + API at /api
+def test_serves_spa_index(static_client):
+    """With a built frontend mounted, GET / returns the SPA index (HTML), not an API response."""
+    r = static_client.get("/")
+    assert r.status_code == 200
+    assert "text/html" in r.headers["content-type"]
+    assert 'id="root"' in r.text
+
+
+def test_serves_static_asset(static_client):
+    """Frontend assets are served from the mount (so the SPA's JS/CSS load same-origin)."""
+    r = static_client.get("/assets/app.js")
+    assert r.status_code == 200
+    assert "dashboard" in r.text
+
+
+def test_api_reachable_under_prefix_with_frontend(static_client):
+    """The API stays fully reachable under /api even with the SPA mounted at / (no shadowing)."""
+    b = static_client.get("/api/health").json()
+    assert b["ready"] is True and b["layer"] == 19
+
+
+def test_unknown_api_path_is_404_not_spa(static_client):
+    """An unknown /api/* path 404s — the /api namespace never falls through to the SPA index."""
+    r = static_client.get("/api/does-not-exist")
+    assert r.status_code == 404
+    assert 'id="root"' not in r.text
+
+
+def test_api_only_when_no_frontend(fake_engine, monkeypatch):
+    """No static_dir and no DASHBOARD_DIST -> API-only: / 404s (nothing mounted), /api still works.
+
+    Guards dev / a frontend-less image: a missing build degrades to API-only, never a crash.
+    """
+    monkeypatch.delenv("DASHBOARD_DIST", raising=False)
+    with TestClient(build_app(fake_engine)) as c:
+        assert c.get("/").status_code == 404
+        assert c.get("/api/health").status_code == 200