Merge branch 'main' into eden_llama

Author: yzhang123

.github/workflows/unit-tests-framework.yml

@@ -30,10 +30,18 @@ jobs:
           fetch-depth: 0
           submodules: "recursive"
 
+      - name: Get merge-base commit
+        id: merge-base
+        run: |
+          # Get the merge-base between current branch and main
+          MERGE_BASE=$(git merge-base HEAD origin/main)
+          echo "merge-base=$MERGE_BASE" >> $GITHUB_OUTPUT
+          echo "Merge-base commit: $MERGE_BASE"
+
       - uses: step-security/changed-files@v46
         id: changed-files
         with:
-          base_sha: main
+          base_sha: ${{ steps.merge-base.outputs.merge-base }}
           files: |
             **
             !models/**
@@ -42,6 +50,7 @@ jobs:
             !.github/**
             !.gitignore
             !.devcontainer/**
+            !ci/scripts/recipes_local_test.py
             .github/workflows/unit-tests-framework.yml
 
       - name: Show output
@@ -239,8 +248,6 @@ jobs:
       - name: Run notebook tests
         env:
           BIONEMO_DATA_SOURCE: ngc
-          # this variable should be used in the notebooks to run a subset of the model layers or a smaller model/dataset
-          FAST_CI_MODE: true
         run: |
           chmod +x ./ci/scripts/run_pytest_notebooks.sh
           ./ci/scripts/run_pytest_notebooks.sh

.github/workflows/unit-tests-recipes.yml

@@ -35,13 +35,21 @@ jobs:
         with:
           fetch-depth: 0
 
+      - name: Get merge-base commit
+        id: merge-base
+        run: |
+          # Get the merge-base between current branch and main
+          MERGE_BASE=$(git merge-base HEAD origin/main)
+          echo "merge-base=$MERGE_BASE" >> $GITHUB_OUTPUT
+          echo "Merge-base commit: $MERGE_BASE"
+
       - name: Get changed files
         id: changed-files
         uses: step-security/changed-files@v46
         with:
           json: true
           matrix: true
-          base_sha: main
+          base_sha: ${{ steps.merge-base.outputs.merge-base }}
           dir_names: true
           dir_names_max_depth: 2
           files: |
@@ -84,21 +92,35 @@ jobs:
           # Assign Docker images to the selected directories
           # Currently, AMPLIFY is the only folder that needs a custom base image, since we have to support both TE and
           # xformers-based models for golden value testing. The rest of the models use the default pytorch image.
+
+          # This uses a squashed version of the pytorch:25.06-py3 image, generated with `docker-squash
+          # nvcr.io/nvidia/pytorch:25.06-py3 -t svcbionemo023/bionemo-framework:pytorch25.06-py3-squashed --output
+          # type=registry,compression=zstd,force-compression=true,oci-mediatypes=true,compression-level=15` and pushed
+          # to the dockerhub registry. Our github actions are able to cache image pulls from dockerhub but not nvcr, so
+          # hopefully this cuts down slightly on CI time at the expense of having a slightly in-directed image location.
+
           DIRS_WITH_IMAGES=$(echo "$DIRS" | jq -c '
             map({
               dir: .,
               image: (
                 if . == "models/amplify" then
                   "svcbionemo023/bionemo-framework:amplify-model-devcontainer-082025"
                 else
-                  "nvcr.io/nvidia/pytorch:25.06-py3"
+                  "svcbionemo023/bionemo-framework:pytorch25.06-py3-squashed-zstd"
                 end
               )
             })
           ')
           echo "dirs=$DIRS_WITH_IMAGES" >> $GITHUB_OUTPUT
+
       - name: Show output
         run: |
+          echo "=== Changed Files Analysis ==="
+          echo "Current branch: ${{ github.ref_name }}"
+          echo "Merge-base commit: ${{ steps.merge-base.outputs.merge-base }}"
+          echo "Changed files compared to merge-base:"
+          echo '${{ steps.changed-files.outputs.all_changed_files }}' | jq -r '.[]' | sed 's/^/  - /'
+          echo "Total changed files: $(echo '${{ steps.changed-files.outputs.all_changed_files }}' | jq '. | length')"
           echo '${{ toJSON(steps.changed-files.outputs) }}'
           echo '${{ toJSON(steps.set-dirs.outputs) }}'
         shell: bash
@@ -115,8 +137,12 @@ jobs:
       fail-fast: false
 
     steps:
+
+      - name: Show GPU info
+        run: nvidia-smi
       - name: Setup proxy cache
         uses: nv-gha-runners/setup-proxy-cache@main
+
       - name: Checkout repository
         uses: actions/checkout@v4
         with:
@@ -125,7 +151,6 @@ jobs:
 
       - name: Install dependencies
         working-directory: ${{ matrix.recipe.dir }}
-        #
         run: |
           if [ -f pyproject.toml ] || [ -f setup.py ]; then
             PIP_CONSTRAINT= pip install -e .

3rdparty/NeMo

@@ -49,6 +49,7 @@ apt-get install -qyy \
   curl \
   pre-commit \
   sudo \
+  emacs-nox \
   gnupg \
   unzip \
   libsqlite3-dev

Dockerfile

@@ -23,26 +23,39 @@
 from pathlib import Path
 from typing import List, Optional
 
+from platformdirs import user_cache_dir
+
+
+PIP_CACHE_DIR = user_cache_dir(appname="bionemo-pip-cache", appauthor="nvidia")
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 DOCKER_RUN_ARGS = [
     "--rm",
-    "-it",
     "--gpus",
     "all",
     "--ipc=host",
     "--ulimit",
     "memlock=-1",
     "--ulimit",
     "stack=67108864",
+    "-v",
+    f"{PIP_CACHE_DIR}:/workspace/.cache/pip",
 ]
 
 CUSTOM_CONTAINERS = {
     "models/amplify": "svcbionemo023/bionemo-framework:amplify-model-devcontainer-082025",
 }
-DEFAULT_CONTAINER = "nvcr.io/nvidia/pytorch:25.06-py3"
+
+# DEFAULT_CONTAINER = "nvcr.io/nvidia/pytorch:25.06-py3"
+
+# This is a squashed version of the pytorch:25.06-py3 image, generated with
+# docker-squash nvcr.io/nvidia/pytorch:25.06-py3 -t svcbionemo023/bionemo-framework:pytorch25.06-py3-squashed
+# --output type=registry,compression=zstd,force-compression=true,oci-mediatypes=true,compression-level=15
+# and pushed to the dockerhub registry. Our github actions are able to cache image pulls from dockerhub but not nvcr, so
+# hopefully this cuts down slightly on CI time at the expense of having a slightly in-directed image location.
+DEFAULT_CONTAINER = "svcbionemo023/bionemo-framework:pytorch25.06-py3-squashed"
 
 
 def get_git_root() -> str:
@@ -89,24 +102,26 @@ def run_tests_in_docker(work_dir: str) -> bool:
     install_and_test_script = textwrap.dedent("""
         set -e  # Exit on any error
 
-        echo "Checking for dependency files..."
+        # Ensure image-embedded constraints do not leak into local recipe installs
+        unset PIP_CONSTRAINT || true
 
+        echo "Checking for dependency files..."
         # Install dependencies based on available files
         if [ -f pyproject.toml ] || [ -f setup.py ]; then
             echo "Installing package in editable mode..."
-            PIP_CONSTRAINT= pip install -e .
+            PIP_CACHE_DIR=/workspace/.cache/pip pip install -e .
             echo "Installed package as editable package"
         elif [ -f requirements.txt ]; then
             echo "Installing from requirements.txt..."
-            PIP_CONSTRAINT= pip install -r requirements.txt
+            PIP_CACHE_DIR=/workspace/.cache/pip pip install -r requirements.txt
             echo "Installed from requirements.txt"
         else
             echo "No pyproject.toml, setup.py, or requirements.txt found"
             exit 1
         fi
 
         echo "Running tests..."
-        pytest -v .
+        python -m pytest -v .
         """)
 
     relative_path = Path(work_dir).relative_to(git_root).as_posix()
@@ -166,6 +181,8 @@ def main():
     if args.debug:
         logging.getLogger().setLevel(logging.DEBUG)
 
+    logger.info(f"Caching pip installations to: {PIP_CACHE_DIR}")
+
     # Get directories to test
     test_dirs = get_test_directories(args.directories)
 

ci/recipes_local_test.py ci/scripts/recipes_local_test.pyci/recipes_local_test.py renamed to ci/scripts/recipes_local_test.py

@@ -19,4 +19,4 @@
 # Enable strict mode with better error handling
 set -euox pipefail
 
-pytest -v --nbval-lax -x -p no:python docs/ sub-packages/
+FAST_CI_MODE=true pytest -v --nbval-lax -x -p no:python docs/ sub-packages/

ci/scripts/run_pytest_notebooks.sh

@@ -19,8 +19,8 @@ dependencies = [
     "nvidia_resiliency_ext",
     "omegaconf",
     "pytest",
-    "torch",
-    # "transformer_engine[pytorch]",
+    "torch==2.6.0a0+ecf3bae40a.nv25.01",
+    "transformer_engine[pytorch]",
     "transformers",
     "xformers",
 ]

models/amplify/pyproject.toml

@@ -86,14 +86,12 @@ def parse_args():
     )
     ap.add_argument(
         "--fp8",
-        type=bool,
         action="store_true",
         default=False,
         help="Whether to use vortex style FP8. Defaults to False.",
     )
     ap.add_argument(
         "--flash-decode",
-        type=bool,
         action="store_true",
         default=False,
         help="Whether to use flash decode. Defaults to True.",