diff --git a/.github/workflows/unit-tests-interpretability-recipes.yaml b/.github/workflows/unit-tests-interpretability-recipes.yaml index c14f5b9f53..01a29c08e5 100644 --- a/.github/workflows/unit-tests-interpretability-recipes.yaml +++ b/.github/workflows/unit-tests-interpretability-recipes.yaml @@ -1,21 +1,24 @@ name: "BioNeMo Interpretability Recipes CI" -# CI for the SAE interpretability recipes under interpretability/sparse_autoencoders/recipes/*. +# CI for the SAE interpretability code under interpretability/sparse_autoencoders/* — both the +# recipes (recipes/*) and the shared `sae` library. # 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/`. +# * `changed-dirs` (cheap ubuntu) discovers which interp units to test and emits a matrix. +# * `unit-tests` runs each selected unit 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". +# Eligible units = any interp recipe OR the `sae` lib itself, each with its own `.ci_build.sh`. +# Today that's `sae` (its own lib tests) plus `evo2` once #1622 lands; codonfm/esm2 have no +# `.ci_build.sh`/tests yet, so they are skipped until they add them. 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 under the shared `sae` lib -> the `sae` lib's own tests AND every recipe +# (all recipes depend on sae). +# * change to this workflow file, or nightly -> ALL eligible units. # * 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. +# Each unit's `.ci_build.sh` owns its own build (evo2 -> mbridge bionemo.evo2; esm2 -> HF; sae -> a +# plain editable install), so a codonfm/esm2 change never triggers the Evo2 megatron build, etc. on: push: @@ -59,14 +62,16 @@ jobs: 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). + # Eligible units = interp recipes AND the shared `sae` lib, each gated on having its own + # .ci_build.sh (units without one are green no-ops until they add it). Including "$SAE" + # makes the lib a first-class matrix entry (build + `pytest tests/`), so sae/tests is + # tested directly instead of only triggering the recipes that depend on it. # Use `if` (not `[ -f ] && echo`): the short-circuit leaves exit 1 when the last dir # lacks a .ci_build.sh, which `set -e` would turn into a job failure instead of an # empty (green no-op) matrix. - ALL=$(for d in "$RROOT"/*/; do if [ -f "${d}.ci_build.sh" ]; then echo "${d%/}"; fi; done \ + ALL=$(for d in "$RROOT"/*/ "$SAE"/; do if [ -f "${d}.ci_build.sh" ]; then echo "${d%/}"; fi; done \ | jq -R -s -c 'split("\n") | map(select(length > 0))') - echo "Eligible recipes (have .ci_build.sh): $ALL" + echo "Eligible units (have .ci_build.sh): $ALL" MERGE_BASE=$(git merge-base HEAD origin/main) CHANGED=$(git diff --name-only "$MERGE_BASE" HEAD) diff --git a/interpretability/sparse_autoencoders/sae/.ci_build.sh b/interpretability/sparse_autoencoders/sae/.ci_build.sh new file mode 100755 index 0000000000..d0e24bdd45 --- /dev/null +++ b/interpretability/sparse_autoencoders/sae/.ci_build.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# Build the shared `sae` interpretability library for CI, mirroring each recipe's `.ci_build.sh` +# so the lib rides the same interp CI matrix (its own `.ci_build.sh` + `pytest tests/`) as a recipe. +# +# Unlike the evo2 recipe (which must build the bionemo.evo2 / mbridge megatron env), the `sae` lib +# is domain-agnostic — torch + triton, both already in the pytorch base image — so there is no venv +# to build. A plain editable install of the package plus its dev/test extras is all pytest needs: +# * test_kernels uses the image's Triton on the L4 GPU; +# * the test_tp_* suite is gloo/CPU multiprocessing, so it runs on the same box's CPUs. +set -euo pipefail + +pip install -e ".[dev]"