[TEST — DO NOT MERGE] verify #1667 interp GPU lane runs on the #1622 recipe

.github/workflows/unit-tests-interpretability-recipes.yaml

@@ -0,0 +1,146 @@
+name: "BioNeMo Interpretability Recipes CI"
+
+# CI for the SAE interpretability recipes under interpretability/sparse_autoencoders/recipes/*.
+# Generic matrix, modeled on the repo-wide unit-tests-recipes.yml but scoped to the interp subtree:
+#   * `changed-dirs` (cheap ubuntu) discovers which interp recipes to test and emits a matrix.
+#   * `unit-tests` runs each selected recipe on the L4: its own `.ci_build.sh` + `pytest tests/`.
+#
+# Eligible recipes = any interp recipe with its own `.ci_build.sh`. Today that's `evo2`; codonfm/esm2
+# have no `.ci_build.sh`/tests yet, so they are skipped until they add them. This also makes the lane
+# a green no-op before the evo2 SAE recipe lands (#1622) — the presence guard is "has a .ci_build.sh".
+#
+# What runs when:
+#   * change under a recipe's own dir            -> that recipe.
+#   * change to the shared `sae` lib, or to this workflow file, or the nightly schedule
+#                                                -> ALL eligible recipes (they all depend on sae).
+#   * change to an unrelated recipe / elsewhere  -> nothing (empty matrix, green no-op).
+# Each recipe's `.ci_build.sh` owns its own build (evo2 -> mbridge bionemo.evo2; esm2 -> HF; etc.),
+# so a codonfm/esm2 change never triggers the Evo2 megatron build, and vice-versa.
+
+on:
+  push:
+    branches:
+      - "pull-request/[0-9]+"
+      - "dependabot/**"
+  merge_group:
+    types: [checks_requested]
+  schedule:
+    - cron: "0 9 * * *" # Runs at 9 AM UTC daily (2 AM MST)
+
+defaults:
+  run:
+    shell: bash -x -e -u -o pipefail {0}
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  changed-dirs:
+    # Cheap ubuntu pre-job: decide which interp recipes to test and emit the matrix.
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    outputs:
+      dirs: ${{ steps.set-dirs.outputs.dirs }}
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5
+        with:
+          fetch-depth: 0
+          persist-credentials: false
+
+      - name: Determine which interp recipes to run
+        id: set-dirs
+        env:
+          EVENT_NAME: ${{ github.event_name }}
+        run: |
+          RROOT="interpretability/sparse_autoencoders/recipes"
+          SAE="interpretability/sparse_autoencoders/sae"
+          WF=".github/workflows/unit-tests-interpretability-recipes.yaml"
+
+          # Eligible recipes = those with their own .ci_build.sh (others are green no-ops until
+          # they add one; empty before #1622 lands -> whole lane is a no-op).
+          ALL=$(for d in "$RROOT"/*/; do [ -f "${d}.ci_build.sh" ] && echo "${d%/}"; done \
+                | jq -R -s -c 'split("\n") | map(select(length > 0))')
+          echo "Eligible recipes (have .ci_build.sh): $ALL"
+
+          MERGE_BASE=$(git merge-base HEAD origin/main)
+          CHANGED=$(git diff --name-only "$MERGE_BASE" HEAD)
+          echo "Changed files:"; echo "$CHANGED" | sed 's/^/  - /'
+
+          # sae lib / this workflow / nightly -> run EVERY eligible recipe (all depend on sae).
+          if [ "$EVENT_NAME" = "schedule" ] || echo "$CHANGED" | grep -qE "^(${SAE}/|${WF}$)"; then
+            DIRS="$ALL"
+          else
+            # otherwise -> only eligible recipes that have a changed file under them.
+            DIRS=$(echo "$ALL" | jq -c --arg changed "$CHANGED" '
+              ($changed | split("\n")) as $cf
+              | map(select(. as $d | $cf | any(startswith($d + "/"))))
+            ')
+          fi
+          echo "Recipes to run: $DIRS"
+
+          # Attach the runner image (the dockerhub-cached squashed pytorch, like the repo-wide lane).
+          IMG="svcbionemo023/bionemo-framework:pytorch26.04-py3-squashed"
+          MATRIX=$(echo "$DIRS" | jq -c --arg img "$IMG" \
+            'map({dir: ., name: (. | split("/") | last), image: $img})')
+          echo "dirs=$MATRIX" >> "$GITHUB_OUTPUT"
+          echo "matrix: $MATRIX"
+
+  unit-tests:
+    needs: changed-dirs
+    if: ${{ needs.changed-dirs.outputs.dirs != '' && needs.changed-dirs.outputs.dirs != '[]' }}
+    runs-on: linux-amd64-gpu-l4-latest-1
+    name: "interp-unit-tests (${{ matrix.recipe.name }})"
+    permissions:
+      contents: read
+    container:
+      image: ${{ matrix.recipe.image }}
+      options: --shm-size=16G
+      env:
+        CI: true
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        HF_HOME: /cache/huggingface
+    strategy:
+      fail-fast: false
+      matrix:
+        recipe: ${{ fromJson(needs.changed-dirs.outputs.dirs) }}
+
+    steps:
+      - name: Show GPU info
+        run: nvidia-smi
+
+      - name: Setup proxy cache
+        uses: nv-gha-runners/setup-proxy-cache@14229018fe157c83e03c008f27d183d8e99bc67c
+
+      - name: Checkout repository
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5
+        with:
+          # Interp recipes live under interpretability/sparse_autoencoders (recipe + shared sae);
+          # evo2 additionally builds on recipes/evo2_megatron. Check out both so any recipe's
+          # .ci_build.sh has what it needs.
+          sparse-checkout: |
+            interpretability/sparse_autoencoders
+            recipes/evo2_megatron
+          sparse-checkout-cone-mode: false
+          persist-credentials: false
+
+      - name: Cache Hugging Face models
+        uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830
+        with:
+          path: /cache/huggingface
+          key: ${{ runner.os }}-huggingface-interp-${{ matrix.recipe.name }}-${{ github.sha }}
+          restore-keys: |
+            ${{ runner.os }}-huggingface-interp-${{ matrix.recipe.name }}-
+            ${{ runner.os }}-huggingface-
+
+      - name: Install dependencies
+        working-directory: ${{ matrix.recipe.dir }}
+        run: bash .ci_build.sh
+
+      - name: Run tests
+        working-directory: ${{ matrix.recipe.dir }}
+        run: |
+          [ -f .ci_test_env.sh ] && source .ci_test_env.sh
+          pytest -v tests/

interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+# Build the Evo2-SAE recipe environment for CI.
+#
+# The inference engine imports `bionemo.evo2` — which lives ONLY in the mbridge
+# `recipes/evo2_megatron` recipe (there is no prebaked image that ships it). So we build
+# that environment with evo2_megatron's OWN CI build script, inheriting its pinned
+# megatron-bridge / causal-conv1d / TransformerEngine handling verbatim (no fork that
+# could drift), then install the SAE library + this recipe on top of the same venv.
+#
+# 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"
+
+# Fail loudly if the sibling recipe isn't checked out, rather than emitting a cryptic path
+# error deep inside the build. In CI this means the workflow's sparse-checkout MUST include
+# recipes/evo2_megatron.
+if [[ ! -f "$MEGATRON/.ci_build.sh" ]]; then
+  echo "ERROR: evo2_megatron not found at $MEGATRON" >&2
+  echo "       evo2-sae builds on the mbridge bionemo.evo2 recipe; ensure the checkout" >&2
+  echo "       includes recipes/evo2_megatron." >&2
+  exit 1
+fi
+
+# 1. Build the bionemo.evo2 (mbridge) environment. Creates $MEGATRON/.venv with the
+#    system-site torch/TE plus the full megatron stack.
+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).
+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/.ci_test_env.sh

@@ -0,0 +1,4 @@
+# Sourced by CI before running pytest (mirrors evo2_megatron/.ci_test_env.sh).
+# The recipe shares evo2_megatron's venv (see .ci_build.sh), which lives in the sibling
+# recipe dir rather than here.
+source ../../../../recipes/evo2_megatron/.venv/bin/activate

interpretability/sparse_autoencoders/recipes/evo2/Dockerfile

@@ -0,0 +1,36 @@
+# syntax=docker/dockerfile:1.4
+#
+# Runnable image for the Evo2 SAE recipe. Thin and non-forking: the engine imports bionemo.evo2
+# (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.
+#
+# 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):
+#   docker run --gpus all -it evo2-sae bash -lc "source .ci_test_env.sh && pytest tests/"
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3
+FROM ${BASE_IMAGE}
+
+WORKDIR /workspace
+# 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 install
+
+# 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/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/

interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml

@@ -12,14 +12,17 @@ dependencies = [
     "sae",
     "torch>=2.0",
     "numpy>=1.20",
+    "pandas>=2.0",  # cli.py `batch` -> parquet
     "pyarrow>=23.0.0",
+    "fastapi>=0.110",  # server.py
+    "uvicorn>=0.29",  # serve CLI
+    "anyio>=4.0",  # server.py: threadpool concurrency cap (ships with fastapi, declared explicitly)
 ]
 
-# No package code lives here yet — the recipe is just an entry-point for
-# scripts/ that depends on the shared `sae` workspace package. Declare no
-# packages so setuptools doesn't try to discover anything.
-[tool.setuptools]
-packages = []
+# The `evo2_sae` package (src/) holds the live inference engine + steering hook + the FastAPI
+# server and CLI; scripts/ (extract, train) are standalone entry points alongside it.
+[tool.setuptools.packages.find]
+where = ["src"]
 
 [tool.uv.sources]
 sae = { workspace = true }

interpretability/sparse_autoencoders/recipes/evo2/scripts/7b.sh

@@ -25,7 +25,7 @@ CHUNK_BP="${CHUNK_BP:-8192}"
 # An Evo2 7B MBridge checkpoint directory (see prerequisites above).
 CKPT_DIR="${CKPT_DIR:?Set CKPT_DIR to an Evo2 7B MBridge checkpoint directory (see header)}"
 FASTA="${FASTA:?Set FASTA to the (prok+euk) input sequences}"
-WORK_ROOT="${WORK_ROOT:-/data/interp/evo2}"
+WORK_ROOT="${WORK_ROOT:?Set WORK_ROOT to a (large, multi-TB) output dir for activations + SAE checkpoints}"
 
 NPROC="${NPROC:-8}"            # GPUs / DP ranks
 MAX_TOKENS="${MAX_TOKENS:-1000000000}"
@@ -72,7 +72,6 @@ torchrun --nproc_per_node="$NPROC" "${RECIPE_DIR}/scripts/train.py" \
     --cache-dir "$PARQUET_DIR" \
     --model-path "$CKPT_DIR" \
     --layer "$LAYER" \
-    --model-type topk \
     --expansion-factor 16 --top-k 128 \
     --normalize-input \
     --auxk 2048 --auxk-coef 0.03125 \

interpretability/sparse_autoencoders/recipes/evo2/scripts/launch_inference.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+# Launch the Evo2 SAE inference engine. One engine, four modes:
+#
+#   ./launch_inference.sh serve                       # live HTTP server on :8001 (viz backend)
+#   ./launch_inference.sh encode  --sequence ATGC...  # annotate ONE sequence -> top features
+#   ./launch_inference.sh batch   --fasta in.fa --out out.parquet   # MANY sequences -> parquet
+#   ./launch_inference.sh generate --prompt ATGC... --clamp 29244:300  # steer + generate DNA
+#
+# Steering loop: `encode` a sequence to find an active feature id, then
+# `generate --clamp ID:STRENGTH` (strength ~2-3x the feature's max_activation; repeat --clamp).
+#
+# Config via env. Required: EVO2_CKPT_DIR, SAE_CKPT_PATH. Optional (have defaults):
+# FEATURE_ANNOTATIONS, EMBEDDING_LAYER (26), DEVICE, PORT, CUDA_VISIBLE_DEVICES.
+#
+# Run INSIDE the evo2_megatron venv (it provides bionemo.evo2 + megatron) — same venv the tests
+# use. In the Docker image it's already active (on PATH), so just run this. On bare metal, activate
+# it first:  source <evo2_megatron>/.venv/bin/activate
+set -euo pipefail
+
+HERE="$(cd "$(dirname "$0")" && pwd)"
+RECIPE_DIR="$(cd "$HERE/.." && pwd)"  # recipes/evo2 — so the evo2_sae package imports
+
+# Required (no hardcoded defaults — supply your own paths via env):
+export EVO2_CKPT_DIR="${EVO2_CKPT_DIR:?Set EVO2_CKPT_DIR to an Evo2 MBridge checkpoint directory}"
+export SAE_CKPT_PATH="${SAE_CKPT_PATH:?Set SAE_CKPT_PATH to a trained SAE checkpoint (.pt)}"
+# Optional: feature-label parquet (empty = features are unlabeled). Layer defaults to 26.
+export FEATURE_ANNOTATIONS="${FEATURE_ANNOTATIONS:-}"
+export EMBEDDING_LAYER="${EMBEDDING_LAYER:-26}"
+
+export PYTHONPATH="$RECIPE_DIR/src${PYTHONPATH:+:$PYTHONPATH}"  # find evo2_sae without a pip install
+
+# We assume the evo2_megatron venv is already active (Docker: on PATH; bare metal: you sourced it).
+# Fail with a clear message rather than a deep ImportError if it isn't.
+if ! python -c "import bionemo.evo2" 2>/dev/null; then
+  echo "ERROR: bionemo.evo2 not importable — activate the evo2_megatron venv first" >&2
+  echo "       (source <evo2_megatron>/.venv/bin/activate) or run inside the Docker image." >&2
+  exit 1
+fi
+
+# One-shot modes (encode/batch/generate) just run once -> exec, signals pass straight through.
+# For `serve`, a bad request can trip a CUDA device-side assert that poisons the process context
+# (unrecoverable in-process — restart is the only cure). So tell the engine to exit the worker on
+# that fault (EXIT_ON_CUDA_WEDGE) and respawn it here, making recovery host-independent.
+#
+# The worker runs in the BACKGROUND and we `wait` on it so this supervisor stays responsive to
+# signals: a trap forwards SIGTERM/SIGINT to the worker (triggering uvicorn's graceful shutdown)
+# and stops the respawn loop — important when this script is PID 1 under `docker stop`/k8s, where
+# bash would otherwise not forward the signal and orphan the worker. Crash/wedge -> respawn with
+# backoff, capped so a persistent failure (e.g. port already bound) doesn't loop forever.
+if [[ "${1:-}" == "serve" ]]; then
+  export EXIT_ON_CUDA_WEDGE=1
+  child=""
+  stop=0
+  trap 'stop=1; [[ -n "$child" ]] && kill -TERM "$child" 2>/dev/null || true' TERM INT
+  rc=0
+  fails=0
+  while [[ "$stop" -eq 0 ]]; do
+    python -m evo2_sae.cli "$@" &
+    child=$!
+    wait "$child" && rc=0 || rc=$?
+    child=""
+    # stop on a clean exit or when a signal asked us to (143 SIGTERM / 130 SIGINT belt-and-suspenders).
+    [[ "$stop" -eq 1 || "$rc" -eq 0 || "$rc" -eq 143 || "$rc" -eq 130 ]] && break
+    fails=$((fails + 1))
+    if [[ "$fails" -ge 10 ]]; then
+      echo "[launch] serve exited ($rc) $fails times — giving up; fix the cause." >&2
+      break
+    fi
+    backoff=$([[ "$fails" -lt 5 ]] && echo 2 || echo 10)
+    echo "[launch] serve exited ($rc); restart $fails/10 in ${backoff}s…" >&2
+    sleep "$backoff" &
+    wait $! 2>/dev/null || true  # interruptible: a signal during backoff stops the loop promptly
+  done
+  exit "$rc"
+fi
+exec python -m evo2_sae.cli "$@"