evo2-sae dashboard: single-container serve (static frontend baked into image)

Make the dashboard + API deployable as ONE container so a public-repo user (or coworker) can
launch it with a single 'docker run', no Node and no second process:

  * Dockerfile: add a Node build stage that compiles feature_explorer -> static dist/, COPY it
    into the runtime image, and set DASHBOARD_DIST so server.py serves it at / (Node lives only
    in the build stage, never in the runtime GPU image).
  * vite.config.js: drop the /api rewrite — proxy /api straight through to :8001/api. Dev now hits
    the same paths as production (the server serves the API under /api), so there's no dev/prod drift.
  * move the /gene_embed endpoint onto the /api router (rebased onto the new server.py); frontend
    already calls ${BACKEND}/gene_embed (= /api/gene_embed), so no frontend code change.
  * re-point the gene_embed contract test to /api/gene_embed.

Dev mode (vite + launch_dashboard.py) is unchanged for UI iteration; the baked-in static build is
the deploy path. Frontend compile is exercised by the image build.

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/Dockerfile

@@ -4,11 +4,31 @@
 # (the mbridge evo2_megatron recipe), so .ci_build.sh delegates to evo2_megatron's OWN build
 # (inheriting its pinned megatron-bridge / causal-conv1d / TransformerEngine versions), then
 # installs the sae library + this recipe on top. We don't reimplement the megatron build here.
+# A Node build stage compiles the dashboard to static files baked into the image, so ONE
+# container serves both the UI (at /) and the API (at /api) — no Node at runtime.
 #
 # Build from the REPO ROOT (the context must include the recipes/evo2_megatron sibling):
 #   docker build -f interpretability/sparse_autoencoders/recipes/evo2/Dockerfile -t evo2-sae .
-# Run (needs a GPU + checkpoints; see the recipe README):
+# Serve the dashboard + API on one port (needs a GPU + checkpoints; see the recipe README):
+#   docker run --gpus all -p 8001:8001 \
+#     -e EVO2_CKPT_DIR=/ckpt/evo2 -e SAE_CKPT_PATH=/ckpt/sae.pt -e EMBEDDING_LAYER=26 \
+#     -v /my/checkpoints:/ckpt evo2-sae scripts/launch_inference.sh serve
+#   # then open http://localhost:8001
+# Or run the tests:
 #   docker run --gpus all -it evo2-sae bash -lc "source .ci_test_env.sh && pytest tests/"
+
+# Stage 1 — build the dashboard frontend to static files. Node lives ONLY in this stage; the
+# runtime image below gets just the compiled dist/, never npm/vite.
+FROM node:20-slim AS frontend
+WORKDIR /frontend
+# Copy manifests first so `npm ci` is cached unless dependencies change.
+COPY interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/package.json \
+     interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/package-lock.json ./
+RUN npm ci
+COPY interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/ ./
+RUN npm run build  # -> /frontend/dist
+
+# Stage 2 — the runnable GPU image.
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3
 FROM ${BASE_IMAGE}
 
@@ -31,6 +51,12 @@ COPY interpretability/sparse_autoencoders interpretability/sparse_autoencoders
 WORKDIR /workspace/interpretability/sparse_autoencoders/recipes/evo2
 RUN bash .ci_build.sh install
 
+# 3. Bake in the prebuilt dashboard and point the server at it (DASHBOARD_DIST), so the one
+#    container serves the UI at / and the API at /api. Only this layer changes on a frontend edit.
+COPY --from=frontend /frontend/dist \
+     /workspace/interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/dist
+ENV DASHBOARD_DIST=/workspace/interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/dist
+
 # Default to the built venv so `docker run … pytest tests/` (or importing Evo2SAE) just works.
 ENV VIRTUAL_ENV=/workspace/recipes/evo2_megatron/.venv
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"

interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/vite.config.js

@@ -11,14 +11,14 @@ export default defineConfig({
     host: '0.0.0.0',
     port: 5176,
     strictPort: true,
-    // The live backend (server.py) runs on :8001. Proxying it under
-    // /api means only the Vite port needs to be tunneled — the browser talks
-    // to Vite, Vite talks to the backend, both on the pod.
+    // The live backend (server.py) runs on :8001 and serves the API under /api. Proxy /api
+    // straight through (NO rewrite) so dev hits the same paths as production: in dev the browser
+    // talks to Vite and Vite forwards /api/* to :8001/api/*; in the single-container build the
+    // browser hits /api/* on the backend directly. Same frontend code, identical paths both ways.
     proxy: {
       '/api': {
         target: 'http://localhost:8001',
         changeOrigin: true,
-        rewrite: (path) => path.replace(/^\/api/, ''),
       },
     },
   },

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

@@ -122,7 +122,7 @@ def test_gene_embed_returns_decodable_matrix(client):
     import numpy as np
 
     genes = [{"symbol": "g1", "sequence": "ACGTACGT"}, {"symbol": "g2", "sequence": "TTTTGGGG"}]
-    b = client.post("/gene_embed", json={"genes": genes, "min_firing": 1}).json()
+    b = client.post("/api/gene_embed", json={"genes": genes, "min_firing": 1}).json()
     assert {"G_b64", "Gmax_b64", "n_features", "n_genes", "genes", "feature_ids"} <= set(b)
     assert b["n_genes"] == 2 and len(b["genes"]) == 2
     # only firing columns are shipped: feature_ids maps column -> real SAE feature id