Merge branch 'main' into farhadr/evo2_cleanup

Author: farhadrgh

.github/codecov.yml

@@ -0,0 +1,14 @@
+codecov:
+  require_ci_to_pass: false
+
+coverage:
+  status:
+    project:
+      default:
+        target: auto
+        threshold: 5
+
+comment:
+  layout: "diff, flags, files"
+  behavior: default
+  require_changes: false  # if true: only post the comment if coverage changes

.github/dependabot.yml

@@ -10,3 +10,13 @@ updates:
     reviewers:
       - "pstjohn"
       - "jstjohn"
+  - package-ecosystem: "docker"
+    directory: "/"
+    target-branch: "main"
+    open-pull-requests-limit: 1
+    schedule:
+      interval: "weekly"
+    reviewers:
+      - "pstjohn"
+      - "dorotat-nv"
+      - "trvachov"

.github/workflows/bionemo-subpackage-ci.yml

@@ -5,24 +5,52 @@ on:
   workflow_dispatch:
     inputs:
       subpackages:
-        description: BioNeMo sub-packages (comma-separated) to test or publish.
+        description: "BioNeMo sub-packages (comma-separated) to test or publish."
         required: true
         type: string
       test:
-        description: Test the sub-packages before publishing to PyPI. Strongly recommended for production releases to PyPI. Can be disabled when staging sub-packages on Test PyPI or publishing circular dependencies to PyPI.
+        description: "Test the sub-packages before publishing to PyPI. Strongly recommended for production releases to PyPI. Can be disabled when staging sub-packages on Test PyPI or publishing circular dependencies to PyPI."
         required: false
         type: boolean
         default: true
       publish:
-        description: Publish the built package to PyPI. If testing is specified, requires that all sub-package tests succeed based on dependencies published to Test PyPI or PyPI.
+        description: "Publish the built package to PyPI. If testing is specified, requires that all sub-package tests succeed based on dependencies published to Test PyPI or PyPI."
         required: false
         type: boolean
         default: false
       pypi:
-        description: Publish to PyPI instead of Test PyPI.
+        description: "Publish to PyPI instead of Test PyPI."
         required: false
         type: boolean
         default: false
+      version_overwrite:
+        description: "Overwrite the published version of the sub-package. (Sets skip-existing to False. Requires deleting existing wheels and other artifacts on PyPI.)"
+        required: false
+        type: boolean
+        default: false
+      build_framework:
+        description: "Build framework to use for building and publishing."
+        type: choice
+        options:
+          - "python"
+          - "rust_pyo3_maturin"
+        default: "python"
+        required: true
+      python_version:
+        description: "Python version to use for testing and publishing."
+        required: false
+        type: string
+        default: "3.12"
+      gpu_runner:
+        description: "Specify a GPU runner for testing on NVIDIA GitHub Actions. (For a list of available runners, refer to: https://docs.gha-runners.nvidia.com/runners/)"
+        required: false
+        type: string
+        default: "linux-amd64-gpu-l4-latest-1"
+      cuda_version:
+        description: "NVIDIA CUDA container version to use for testing."
+        required: false
+        type: string
+        default: "nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04"
 
 jobs:
   configure-workflow-packages:
@@ -49,18 +77,19 @@ jobs:
         package: ${{ fromJson(needs.configure-workflow-packages.outputs.workflow_packages) }}
       fail-fast: false  # Prevent all matrix jobs from failing if one fails.
     name: "[${{ matrix.package }}] Install and test sub-package."
-    runs-on: linux-amd64-gpu-l4-latest-1
-    container: # GPU jobs must run in a container. Use a fresh base container for package installation and testing.
-      image: nvidia/cuda:12.8.1-base-ubuntu22.04
+    # Use GPU runner only when testing, otherwise use a standard runner.
+    runs-on: ${{ github.event.inputs.test == 'true' && github.event.inputs.gpu_runner || 'ubuntu-latest' }}
+    container:
+      # GPU jobs must run in a container. Use a fresh CUDA base container for package installation and testing.
+      # If testing is disabled, use a lightweight container to quickly skip this job.
+      image: ${{ github.event.inputs.test == 'true' && github.event.inputs.cuda_version || 'ubuntu:latest' }}
     steps:
       # Silently skip all steps if testing is disabled, which does not block building or publishing.
       - name: Install git and system dependencies.
         if: ${{ github.event.inputs.test == 'true' }}
         run: |
           apt-get update
-          apt-get install -y git
-          apt-get install -y lsb-release      # No longer pre-installed in Ubuntu>=22.04.
-          apt-get install -y build-essential  # For installing C build tools, like GCC and make.
+          apt-get install -qyy git curl lsb-release build-essential
       - uses: actions/checkout@v4
         if: ${{ github.event.inputs.test == 'true' }}
         with:
@@ -69,16 +98,41 @@ jobs:
       - uses: actions/setup-python@v5
         if: ${{ github.event.inputs.test == 'true' }}
         with:
-          python-version: "3.12"
-      - id: install-dispatch-subpackage
+          python-version: ${{ github.event.inputs.python_version }}
+      - id: install-rust
+        if: ${{ github.event.inputs.test == 'true' }}
+        name: Install Rust.
+        run: |
+          curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+          . $HOME/.cargo/env
+          rustc --version
+          cargo --version
+          rustup --version
+      - id: install-subpackage-core
         if: ${{ github.event.inputs.test == 'true' }}
         name: Install sub-package.
         run: |
+          # Setup environment, i.e. add Rust to PATH and silence pip root user warnings.
+          . $HOME/.cargo/env
           # Install sub-package and dependencies.
-          pip install --upgrade pip setuptools
-          pip install pytest
+          pip install --upgrade pip setuptools uv maturin
           # Install required core & optional [test] dependencies.
-          pip install --no-cache-dir 'sub-packages/${{ matrix.package }}[test]'
+          uv pip install --no-cache --system pytest sub-packages/${{ matrix.package }}[test]
+      - id: install-subpackage-post
+        if: ${{ github.event.inputs.test == 'true' }}
+        name: Install sub-package dependencies that need to be installed after the core dependencies.
+        run: |
+          # DEV: Post-install dependencies are configured in [project.optional-dependencies].
+          # `uv pip install --extra <optional-dependency> -r <pyproject.toml>` tracks
+          # post-dependencies in the pyproject.toml and avoids installing core dependencies
+          # redundantly, which causes errors with incompatible --config-setting.
+
+          # TransformerEngine
+          uv pip install --no-cache --no-build-isolation --system --extra te -r sub-packages/${{ matrix.package }}/pyproject.toml || echo "[BioNeMo Sub-Package CI] TE will not be installed."
+
+          # # Apex
+          # # NOTE: --cpp_ext and --cuda_ext are required for building fused Apex kernels.
+          # uv pip install --no-cache --no-build-isolation --system --config-setting="--build-option=--cpp_ext" --config-setting="--build-option=--cuda_ext" --extra apex -r sub-packages/${{ matrix.package }}/pyproject.toml || echo "[BioNeMo Sub-Package CI] Apex will not be installed."
       - id: test-dispatch-subpackage
         if: ${{ github.event.inputs.test == 'true' }}
         name: Test sub-package.
@@ -103,23 +157,30 @@ jobs:
           persist-credentials: false
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: ${{ github.event.inputs.python_version }}
       - id: build-package
         name: Build a binary wheel and a source tarball for the sub-package.
         run: |
-          if [[ "${{ github.event.inputs.test }}" != "true" ]]; then
+          if [[ "${{ github.event.inputs.test }}" != "true" && "${{ github.event.inputs.version_overwrite }}" != "true" ]]; then
             # For untested sub-packages, append '-dev' to the version for PyPI.
             sed -i 's/[[:space:]]*$//' sub-packages/${{ matrix.package }}/VERSION
             sed -i 's/$/-dev/' sub-packages/${{ matrix.package }}/VERSION
           fi
-          python -m pip install build
-          python -m build sub-packages/${{ matrix.package }}
+          # Build the sub-package.
+          if [[ "${{ github.event.inputs.build_framework }}" == "python" ]]; then
+            pip install build
+            python -m build sub-packages/${{ matrix.package }}
+          elif [[ "${{ github.event.inputs.build_framework }}" == "rust_pyo3_maturin" ]]; then
+            # Install maturin[zig] to build the Rust sub-package with compatibility for manylinux_X_Y using zig.
+            pip install maturin[zig]
+            maturin build --release --zig -m sub-packages/${{ matrix.package }}/Cargo.toml
+          fi
       - id: upload-distribution
         name: Upload distribution packages to the workflow.
         uses: actions/upload-artifact@v4
         with:
-          name: ${{ matrix.package }}-dist
-          path: sub-packages/${{ matrix.package }}/dist
+          name: ${{ matrix.package }}-build-artifacts
+          path: ${{ github.event.inputs.build_framework == 'rust_pyo3_maturin' && format('sub-packages/{0}/target/wheels', matrix.package) || format('sub-packages/{0}/dist', matrix.package) }}
 
   publish-to-pypi:
     needs: [build-pypi, install-and-test]
@@ -141,7 +202,7 @@ jobs:
         name: Download the built distribution.
         uses: actions/download-artifact@v4
         with:
-          name: ${{ matrix.package }}-dist
+          name: ${{ matrix.package }}-build-artifacts
           path: sub-packages/${{ matrix.package }}/dist
       - id: publish-to-testpypi
         name: Publish distribution 📦 to Test PyPI for PR.
@@ -151,7 +212,7 @@ jobs:
           verbose: true
           packages-dir: sub-packages/${{ matrix.package }}/dist
           repository-url: https://test.pypi.org/legacy/
-          skip-existing: true
+          skip-existing: ${{ github.event.inputs.version_overwrite }}
       - id: publish-to-pypi
         name: Publish distribution 📦 to PyPI for Workflow Dispatch.
         # To require testing before publishing to PyPI, add: ... && needs.install-and-test.result == 'success'
@@ -162,4 +223,4 @@ jobs:
         with:
           verbose: true
           packages-dir: sub-packages/${{ matrix.package }}/dist
-          skip-existing: true
+          skip-existing: ${{ github.event.inputs.version_overwrite }}

Dockerfile

@@ -21,7 +21,7 @@
 #   training loss curves from NeMo.
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:25.01-py3
 
-FROM rust:1.82.0 AS rust-env
+FROM rust:1.86.0 AS rust-env
 
 RUN rustup set profile minimal && \
   rustup install 1.82.0 && \
@@ -114,6 +114,12 @@ RUN if [ "$TARGETARCH" = "arm64" ]; then \
     pip install .; \
 fi
 
+# On ARM, bits and bytes needs to be built from scratch
+RUN if [ "$TARGETARCH" = "arm64" ]; then \
+    cd / && pip uninstall bitsandbytes && \
+    git clone --single-branch --branch 0.45.5 https://github.com/bitsandbytes-foundation/bitsandbytes.git && \
+    cd bitsandbytes && pip install . && cd .. && rm -rf bitsandbytes; \
+fi
 ###############################################################################
 # /end ARM
 ###############################################################################
@@ -128,6 +134,9 @@ RUN pip install hatchling urllib3  # needed to install nemo-run
 ARG NEMU_RUN_TAG=v0.3.0
 RUN pip install nemo_run@git+https://github.com/NVIDIA/NeMo-Run.git@${NEMU_RUN_TAG} --use-deprecated=legacy-resolver
 
+# Rapids SingleCell Installation
+RUN pip install 'rapids-singlecell' --extra-index-url=https://pypi.nvidia.com
+
 RUN mkdir -p /workspace/bionemo2/
 
 WORKDIR /workspace
@@ -140,7 +149,7 @@ RUN rm -rf /opt/pytorch/pytorch/third_party/onnx
 # environment, and does not use the current uv.lock file. Note that with python 3.12, we now need to set
 # UV_BREAK_SYSTEM_PACKAGES, since the pytorch base image has made the decision not to use a virtual environment and UV
 # does not respect the PIP_BREAK_SYSTEM_PACKAGES environment variable set in the base dockerfile.
-COPY --from=ghcr.io/astral-sh/uv:0.4.25 /uv /usr/local/bin/uv
+COPY --from=ghcr.io/astral-sh/uv:0.6.13 /uv /usr/local/bin/uv
 ENV UV_LINK_MODE=copy \
   UV_COMPILE_BYTECODE=1 \
   UV_PYTHON_DOWNLOADS=never \
@@ -175,7 +184,7 @@ uv pip install maturin --no-build-isolation
 git clone https://github.com/NVIDIA/nvidia-resiliency-ext
 uv pip install nvidia-resiliency-ext/
 rm -rf nvidia-resiliency-ext/
-# ngcsdk causes strange dependency conflicts that we will resolve later
+# ngcsdk causes strange dependency conflicts (ngcsdk requires protobuf<4, but nemo_toolkit requires protobuf==4.24.4, deleting it from the uv pip install prevents installation conflicts)
 sed -i "/ngcsdk/d" ./sub-packages/bionemo-core/pyproject.toml
 # Remove llama-index because bionemo doesn't use it and it adds CVEs to container
 sed -i "/llama-index/d" ./3rdparty/NeMo/requirements/requirements_nlp.txt
@@ -185,7 +194,7 @@ uv pip install --no-build-isolation \
 -r /requirements-cve.txt \
 -r /requirements-test.txt
 
-# Install back ngcsdk. Somehow doing it here avoids a large dependency loop
+# Install back ngcsdk, as a WAR for the protobuf version conflict with nemo_toolkit.
 uv pip install ngcsdk
 
 # Addressing security scan issue - CVE vulnerability https://github.com/advisories/GHSA-g4r7-86gm-pgqc The package is a
@@ -234,7 +243,7 @@ USER $USERNAME
 COPY --from=bionemo2-base --chown=$USERNAME:$USERNAME --chmod=777 \
   /usr/local/lib/python3.12/dist-packages /usr/local/lib/python3.12/dist-packages
 
-COPY --from=ghcr.io/astral-sh/uv:0.4.25 /uv /usr/local/bin/uv
+COPY --from=ghcr.io/astral-sh/uv:0.6.13 /uv /usr/local/bin/uv
 ENV UV_LINK_MODE=copy \
   UV_COMPILE_BYTECODE=0 \
   UV_PYTHON_DOWNLOADS=never \

README.md

@@ -13,6 +13,8 @@ fastest performance on the market. You can access BioNeMo Framework as a free co
 or learn more at <https://www.nvidia.com/en-us/clara/bionemo/> about getting an enterprise license for improved
 expert-level support.
 
+BioNeMo Framework is part of a larger ecosystem of NVIDIA Biopharma products. Get notified of new releases, bug fixes, critical security updates, and more for biopharma. [Subscribe.](https://www.nvidia.com/en-us/clara/biopharma/product-updates/)
+
 ## Structure of the Framework
 
 The `bionemo-framework` is organized into independently installable namespace packages. These are located under the

ci/benchmarks/partial-conv/esm2_pretrain.yaml

@@ -25,7 +25,8 @@ script_args:
   batch_size:
     value: 16
   max_steps:
-    value: 26500
+    value: 500000
+  stop_steps: 26500
 script: |-
   WANDB_API_KEY=$BIONEMO_WANDB_API_KEY ${variant}_${model} \
     --train-cluster-path=${data_path}/train_clusters.parquet \
@@ -38,6 +39,7 @@ script: |-
     --val-check-interval=1000 \
     --limit-val-batches=1 \
     --num-steps=${max_steps} \
+    --early-stop-on-step ${stop_steps} \
     --min-seq-length=1024 \
     --max-seq-length=1024 \
     --num-layers=33 \

ci/benchmarks/partial-conv/evo2_pretrain.yaml

@@ -68,7 +68,8 @@ script: |-
   --devices=${gpus} \
   --micro-batch-size=${batch_size} \
   --model-size=${config_name} \
-  --max-steps=${max_steps} --early-stop-on-step ${stop_steps} \
+  --max-steps=${max_steps} \
+  --early-stop-on-step ${stop_steps} \
   --limit-val-batches=20 \
   --log-every-n-steps=50 \
   --val-check-interval=500 \

ci/benchmarks/partial-conv/geneformer_pretrain.yaml

@@ -0,0 +1,60 @@
+scope: partial-conv
+time_limit: 14400
+script_args:
+  # All arguments referenced in the script string must be specified here.
+  # Arguments not referenced in the script string must have the 'arg' field specified.
+  # See jet/core/configs.py for the specification of the configuration class
+  workspace:
+    value: /workspace/bionemo2
+    key_segment: False
+  data_path:
+    value: /data/cellxgene_scdl
+    key_segment: False
+  model:
+    value: geneformer
+  variant:
+    value: train
+  config_name:
+    value: geneformer_config
+  precision:
+    value: [bf16-mixed]
+  nodes:
+    value: [2]
+  gpus:
+    value: 8
+  batch_size:
+    value: 32
+  max_steps:
+    value: 37000
+  lr:
+    value: 0.001
+  val_check_interval:
+    value: 500
+  acc:
+    value: 1
+
+script: |-
+   WANDB_API_KEY=$BIONEMO_WANDB_API_KEY ${variant}_${model} \
+    --data-dir ${data_path} \
+    --experiment-name ${batch_size}bs_${nodes}node_${gpus}gpu_${max_steps}s_${precision}prec \
+    --num-gpus ${gpus} \
+    --save-last-checkpoint \
+    --num-nodes ${nodes} \
+    --val-check-interval ${val_check_interval} \
+    --num-dataset-workers 8 \
+    --num-steps ${max_steps} \
+    --seq-length 2048 \
+    --limit-val-batches 8 \
+    --micro-batch-size ${batch_size} \
+    --resume-if-exists \
+    --log-every-n-steps 50 \
+    --lr ${lr} \
+    --create-tensorboard-logger \
+    --result-dir=${tensorboard_dir} \
+    --wandb-project ${wandb_project_name} \
+    --wandb-job-type=${pipeline_label} \
+    --cosine-rampup-frac 0.004331629559040111 \
+    --cosine-hold-frac 0.021658147795200554 \
+    --accumulate-grad-batches ${acc} \
+    --precision ${precision} \
+    --disable-checkpointing;