build(evo2-sae): cache the megatron build separately from SAE code (Docker layers)

The Dockerfile copied the whole sparse_autoencoders tree before the single `RUN .ci_build.sh`,
so editing any SAE file invalidated that layer and re-ran the ~30-min megatron compile.

Split into two layers via a phase arg on .ci_build.sh (env | install | all):
- Layer 1: COPY recipes/evo2_megatron + the recipe's .ci_build.sh, `RUN .ci_build.sh env` — the
  expensive mbridge build, depends ONLY on evo2_megatron, so it stays cached across code edits.
- Layer 2: COPY the SAE source, `RUN .ci_build.sh install` — the two editable pip installs
  (seconds). A code change now invalidates only this layer.
CI still calls `.ci_build.sh` (no arg = all). Adds a per-Dockerfile .dockerignore to keep
node_modules/dist/__pycache__/.venv out of the build context.

Not docker-built here (no docker in this env) — bash syntax + phase routing verified; CI builds it.

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/.ci_build.sh

@@ -9,8 +9,16 @@
 #
 # The CI checkout must also provide the sibling recipe (recipes/evo2_megatron); it pulls its
 # own bionemo-core / bionemo-recipeutils deps from git, so nothing else is required locally.
+#
+# Phases (arg, default "all") — let the Dockerfile cache the expensive megatron build separately
+# from the cheap, frequently-changing SAE installs (so a code edit doesn't rebuild megatron):
+#   env      build the bionemo.evo2 (mbridge) env only       (depends on recipes/evo2_megatron)
+#   install  install the SAE library + this recipe only       (depends on the SAE source)
+#   all      both (CI / default)
 set -euo pipefail
 
+phase="${1:-all}"
+
 HERE="$(cd "$(dirname "$0")" && pwd)"
 SAE_LIB="$HERE/../../sae"                       # sparse_autoencoders/sae
 MEGATRON="$HERE/../../../../recipes/evo2_megatron"
@@ -27,10 +35,14 @@ fi
 
 # 1. Build the bionemo.evo2 (mbridge) environment. Creates $MEGATRON/.venv with the
 #    system-site torch/TE plus the full megatron stack.
-( cd "$MEGATRON" && bash .ci_build.sh )
+if [[ "$phase" == "env" || "$phase" == "all" ]]; then
+  ( cd "$MEGATRON" && bash .ci_build.sh )
+fi
 
 # 2. Add the generic SAE library + this recipe into that venv. PIP_CONSTRAINT= clears the
 #    TE pin constraint evo2_megatron sets (it must not block our pure-Python installs).
-source "$MEGATRON/.venv/bin/activate"
-PIP_CONSTRAINT= pip install -e "$SAE_LIB"
-PIP_CONSTRAINT= pip install -e "$HERE"
+if [[ "$phase" == "install" || "$phase" == "all" ]]; then
+  source "$MEGATRON/.venv/bin/activate"
+  PIP_CONSTRAINT= pip install -e "$SAE_LIB"
+  PIP_CONSTRAINT= pip install -e "$HERE"
+fi

interpretability/sparse_autoencoders/recipes/evo2/Dockerfile

@@ -13,13 +13,23 @@ ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3
 FROM ${BASE_IMAGE}
 
 WORKDIR /workspace
-# evo2_megatron pulls its bionemo-core / bionemo-recipeutils deps from git, so only these two
-# trees are needed in the build context.
-COPY interpretability/sparse_autoencoders interpretability/sparse_autoencoders
+# Two layers so editing SAE code doesn't rebuild the ~30-min megatron stack:
+#
+# 1. Expensive, rarely-changing: build the bionemo.evo2 (mbridge) env. Depends ONLY on
+#    recipes/evo2_megatron (which pulls its bionemo-core/recipeutils deps from git), so this layer
+#    stays cached across SAE code edits. Copy just the recipe's .ci_build.sh to drive the build.
 COPY recipes/evo2_megatron recipes/evo2_megatron
+COPY interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh \
+     interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh
+WORKDIR /workspace/interpretability/sparse_autoencoders/recipes/evo2
+RUN bash .ci_build.sh env
 
+# 2. Cheap, frequently-changing: add the SAE library + this recipe, install editable. A code
+#    change invalidates only this layer (the two pip installs, seconds), not the build above.
+WORKDIR /workspace
+COPY interpretability/sparse_autoencoders interpretability/sparse_autoencoders
 WORKDIR /workspace/interpretability/sparse_autoencoders/recipes/evo2
-RUN bash .ci_build.sh
+RUN bash .ci_build.sh install
 
 # Default to the built venv so `docker run … pytest tests/` (or importing Evo2SAE) just works.
 ENV VIRTUAL_ENV=/workspace/recipes/evo2_megatron/.venv

interpretability/sparse_autoencoders/recipes/evo2/Dockerfile.dockerignore

@@ -0,0 +1,11 @@
+# Per-Dockerfile ignore (BuildKit uses <dockerfile>.dockerignore when present). Paths are relative
+# to the build context (the repo root). Keep generated/heavy trees out of the context so they
+# neither bloat the upload nor bust the COPY cache when they change locally.
+**/node_modules/
+**/dist/
+**/.vite/
+**/__pycache__/
+**/*.pyc
+**/.venv/
+**/.pytest_cache/
+.git/