Merge branch 'main' into farhadr/evo2_cleanup

Author: farhadrgh

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

@@ -1,115 +1,99 @@
 name: BioNeMo Sub-Package Workflow
 
 on:
-  # TODO(@cspades): Re-enable PR-based testing and publishing for debugging or when we have sufficient compute resources.
-  # pull_request:
-  #   types: [opened, synchronize, reopened, ready_for_review]
+  # To test or publish sub-packages or adjustments to this workflow that are branched in PR's, manually dispatch this workflow on the PR's branch here: https://github.com/NVIDIA/bionemo-framework/actions/workflows/bionemo-subpackage-ci.yml.
   workflow_dispatch:
     inputs:
       subpackages:
         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.
+        required: false
+        type: boolean
+        default: true
       publish:
-        description: Publish the built package to 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.
         required: false
         type: boolean
         default: false
-
-env:
-  # TODO(@cspades): Skip pytest for packages that require GPU's or special installation recipes until resolved.
-  SUBPACKAGE_SKIP_TEST: '["bionemo-moco","bionemo-noodles","bionemo-geometric","bionemo-webdatamodule","bionemo-size-aware-batching","bionemo-testing","bionemo-llm","bionemo-evo2","bionemo-esm2","bionemo-geneformer","bionemo-example_model","bionemo-fw"]'
-  # TODO(@cspades): Skip publishing for packages that require special installation recipes or should not be published to PyPI.
-  SUBPACKAGE_SKIP_PUBLISH: '["bionemo-noodles","bionemo-geometric","bionemo-testing","bionemo-example_model"]'
 
 jobs:
   configure-workflow-packages:
     name: "[Configure Workflow Packages] Identify sub-packages for testing and publishing."
-    # Sub-package CI override label in PR's.
-    if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'SKIP_SUBPACKAGE_CI') }}
     runs-on: ubuntu-latest
-    # Depending on the type of GitHub Event (Dispatch or PR), we stage different packages for testing and publishing in 'workflow_packages'.
     outputs:
-      workflow_packages: ${{ steps.collect-staged-packages.outputs.workflow_packages }}
+      workflow_packages: ${{ steps.parse-dispatch-packages.outputs.dispatch_packages }}
     steps:
-      - uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-          submodules: "recursive"
-      - id: collect-package-deltas
-        if: ${{ github.event_name == 'pull_request' }}
-        name: Check which sub-packages have been updated in the PR.
-        run: |
-          # Get list of changed files between PR base and head.
-          changed_files=$(git diff --name-only origin/${{ github.base_ref }})
-
-          # Find unique package directories that changed.
-          changed_packages=$(echo "$changed_files" | grep "^sub-packages/" | cut -d'/' -f2 | sort -u | jq -R -s -c 'split("\n")[:-1]')
-
-          # Send the list of changed packages to the next job.
-          echo "pr_packages=$changed_packages" >> "$GITHUB_OUTPUT"
-          echo "[PR Mode] Sub-packages to stage: $changed_packages"
       - id: parse-dispatch-packages
         if: ${{ github.event_name == 'workflow_dispatch' }}
         name: Parse the sub-packages specified in the workflow dispatch.
         run: |
           # Send the parsed list of sub-packages to the next job.
           dispatch_packages=$(echo '${{ github.event.inputs.subpackages }}' | jq -R -c 'split(",")')
           echo "dispatch_packages=$dispatch_packages" >> "$GITHUB_OUTPUT"
-          echo "[Workflow Dispatch Mode] Sub-packages to stage: $dispatch_packages"
-      - id: collect-staged-packages
-        name: Collect the sub-packages to stage.
-        env:
-          PR_PACKAGES: ${{ steps.collect-package-deltas.outputs.pr_packages }}
-          DISPATCH_PACKAGES: ${{ steps.parse-dispatch-packages.outputs.dispatch_packages }}
-        run: |
-          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
-            echo "workflow_packages=$PR_PACKAGES" >> "$GITHUB_OUTPUT"
-          elif [[ "${{ github.event_name }}" == "workflow_dispatch" ]]; then
-            echo "workflow_packages=$DISPATCH_PACKAGES" >> "$GITHUB_OUTPUT"
-          else
-            echo "workflow_packages=[]" >> "$GITHUB_OUTPUT"
-          fi
+          echo "[BioNeMo Sub-Package CI] Sub-packages to stage: $dispatch_packages"
 
   install-and-test:
     needs: configure-workflow-packages
-    outputs:
-      test_packages: ${{ needs.configure-workflow-packages.outputs.workflow_packages }}
+    # Check if the previous job has any staged packages to test and publish.
+    if: ${{ needs.configure-workflow-packages.outputs.workflow_packages != '[]' }}
     strategy:
       matrix:
         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: ubuntu-latest
+    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
     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.
       - uses: actions/checkout@v4
+        if: ${{ github.event.inputs.test == 'true' }}
         with:
           fetch-depth: 0
           submodules: "recursive"
       - uses: actions/setup-python@v5
+        if: ${{ github.event.inputs.test == 'true' }}
         with:
           python-version: "3.12"
-          cache: "pip"
       - id: install-dispatch-subpackage
+        if: ${{ github.event.inputs.test == 'true' }}
         name: Install sub-package.
-        if: ${{ !contains(fromJson(env.SUBPACKAGE_SKIP_TEST), matrix.package) }}
         run: |
+          # Install sub-package and dependencies.
+          pip install --upgrade pip setuptools
           pip install pytest
-          pip install sub-packages/${{ matrix.package }}
+          # Install required core & optional [test] dependencies.
+          pip install --no-cache-dir 'sub-packages/${{ matrix.package }}[test]'
       - id: test-dispatch-subpackage
+        if: ${{ github.event.inputs.test == 'true' }}
         name: Test sub-package.
-        if: ${{ !contains(fromJson(env.SUBPACKAGE_SKIP_TEST), matrix.package) }}
         run: pytest -vv sub-packages/${{ matrix.package }}
 
   build-pypi:
     # Build distributions from either the workflow dispatch or PR.
     # Validate building before merging or publishing.
-    needs: install-and-test
+    needs: [configure-workflow-packages, install-and-test]
+    if: ${{ needs.configure-workflow-packages.outputs.workflow_packages != '[]' && github.event.inputs.publish == 'true' }}
     outputs:
-      staged_packages: ${{ needs.install-and-test.outputs.test_packages }}
+      staged_packages: ${{ needs.configure-workflow-packages.outputs.workflow_packages }}
     strategy:
       matrix:
-        package: ${{ fromJson(needs.install-and-test.outputs.test_packages) }}
+        package: ${{ fromJson(needs.configure-workflow-packages.outputs.workflow_packages) }}
       fail-fast: false  # Prevent all matrix jobs from failing if one fails.
     name: "[${{ matrix.package }}] Build the sub-package."
     runs-on: ubuntu-latest
@@ -120,48 +104,48 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
-          cache: "pip"
       - id: build-package
         name: Build a binary wheel and a source tarball for the sub-package.
-        if: ${{ !contains(fromJson(env.SUBPACKAGE_SKIP_PUBLISH), matrix.package) }}
         run: |
+          if [[ "${{ github.event.inputs.test }}" != "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 }}
       - id: upload-distribution
         name: Upload distribution packages to the workflow.
-        if: ${{ !contains(fromJson(env.SUBPACKAGE_SKIP_PUBLISH), matrix.package) }}
         uses: actions/upload-artifact@v4
         with:
           name: ${{ matrix.package }}-dist
           path: sub-packages/${{ matrix.package }}/dist
 
   publish-to-pypi:
-    needs: build-pypi
+    needs: [build-pypi, install-and-test]
+    # Require staged sub-package builds for publishing to PyPI.
+    if: ${{ needs.build-pypi.result == 'success' }}
     strategy:
       matrix:
         package: ${{ fromJson(needs.build-pypi.outputs.staged_packages) }}
       fail-fast: false  # Prevent all matrix jobs from failing if one fails.
     name: Publish ${{ matrix.package }} to PyPI.
     runs-on: ubuntu-latest
     environment:
-      name: ${{ github.event_name == 'pull_request' && 'testpypi' || 'pypi' }}
-      url: ${{ github.event_name == 'pull_request' && format('https://test.pypi.org/p/{0}', matrix.package) || format('https://pypi.org/p/{0}', matrix.package) }}
+      name: ${{ github.event.inputs.pypi && 'pypi' || 'testpypi' }}
+      url: ${{ github.event.inputs.pypi && format('https://pypi.org/p/{0}', matrix.package) || format('https://test.pypi.org/p/{0}', matrix.package) }}
     permissions:
       id-token: write
-    # Publish to Test PyPI for PRs, and publish to PyPI for (manually triggered) Workflow Dispatches.
-    # Note: Test PyPI is designed to run on each commit to the main branch and is often used to indicate a healthy release publishing pipeline.
     steps:
       - id: download-distribution
         name: Download the built distribution.
-        if: ${{ !contains(fromJson(env.SUBPACKAGE_SKIP_PUBLISH), matrix.package) }}
         uses: actions/download-artifact@v4
         with:
           name: ${{ matrix.package }}-dist
           path: sub-packages/${{ matrix.package }}/dist
-      # TODO(@cspades): Uncomment when developing on this workflow to test. Otherwise, no need to publish for every PR.
       - id: publish-to-testpypi
         name: Publish distribution 📦 to Test PyPI for PR.
-        if: ${{ github.event_name == 'pull_request' && !contains(fromJson(env.SUBPACKAGE_SKIP_PUBLISH), matrix.package) }}
+        if: ${{ github.event.inputs.pypi == 'false' }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           verbose: true
@@ -170,7 +154,10 @@ jobs:
           skip-existing: true
       - id: publish-to-pypi
         name: Publish distribution 📦 to PyPI for Workflow Dispatch.
-        if: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true' && !contains(fromJson(env.SUBPACKAGE_SKIP_PUBLISH), matrix.package) }}
+        # To require testing before publishing to PyPI, add: ... && needs.install-and-test.result == 'success'
+        # If testing is run but fails, the workflow will fail and not publish to PyPI (or Test PyPI).
+        # We strongly recommend testing when publishing to production PyPI.
+        if: ${{ github.event.inputs.pypi == 'true' }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           verbose: true

Dockerfile

@@ -35,12 +35,6 @@ apt-get upgrade -qyy \
 rm -rf /tmp/* /var/tmp/*
 EOF
 
-## BUMP TE from v1.14 to v1.13
- ARG TE_TAG=v1.13
- RUN NVTE_FRAMEWORK=pytorch NVTE_WITH_USERBUFFERS=1 MPI_HOME=/usr/local/mpi \
-    pip --disable-pip-version-check --no-cache-dir install \
-    git+https://github.com/NVIDIA/TransformerEngine.git@${TE_TAG}
-
 # Use a branch of causal_conv1d while the repository works on Blackwell support.
 RUN CAUSAL_CONV1D_FORCE_BUILD=TRUE pip --disable-pip-version-check --no-cache-dir install git+https://github.com/trvachov/causal-conv1d.git@52e06e3d5ca10af0c7eb94a520d768c48ef36f1f
 

docs/docs/user-guide/contributing/contributing.md

@@ -120,12 +120,6 @@ Key behaviors:
 - Skips entire CI pipeline
 - Use for documentation typos, README updates
 
-#### **SKIP_SUBPACKAGE_CI**
-
-- Skips installation, testing, and publication of individual sub-packages of BioNeMo.
-- For more granular controls on a per-package basis, such as skipping only testing or only publication to PyPI, you can modify hard-coded sub-package names (`bionemo-...`) listed within `SKIP...` variables in [`bionemo-framework/.github/workflows/bionemo-subpackage-ci.yml`](../../../../.github/workflows/bionemo-subpackage-ci.yml).
-- Sub-package CI is enabled (not skipped) by default. Utilized to test individual BioNeMo sub-packages without the environmental support of the BioNeMo Framework Container, which validates if the sub-package can be pusblished standalone to PyPI.
-
 #### **INCLUDE_NOTEBOOKS_TESTS**
 
 - Enables notebook validation tests
@@ -259,7 +253,7 @@ The resulting altered baseline files should then be committed.
   - Unit tests not associated with source code in BioNeMo can be placed anywhere reasonable under `tests/bionemo/<package-name-suffix>`.
 - Verify that the `pyproject.toml` is `pip install`-able (and `python -m build`-able).
   - If the sub-package is publishable, follow the instructions in [Publishing to PyPI](#publishing-to-pypi) to register or link your package to the sub-package workflow in BioNeMo Framework.
-  - If the sub-package cannot be installed, built, or published to PyPI, add the name of this sub-package to `SUBPACKAGE_SKIP_TEST` and `SUBPACKAGE_SKIP_PUBLISH` in [`bionemo-framework/.github/workflows/bionemo-subpackage-ci.yml`](../../../../.github/workflows/bionemo-subpackage-ci.yml).
+  - Add test dependencies to a `test` field under `[project.optional-dependencies]` for test-only dependencies.
 
 ### Publishing to PyPI
 
@@ -276,27 +270,33 @@ To publish your sub-package via "Trusted Publishing" to PyPI, you can follow the
   - Environment Name:
     - `pypi` for PyPI
     - `testpypi` for Test PyPI
-- **NVIDIA-Only**: Run the workflow!
-  - Create or update any PR with `git diff` changes to your `bionemo-framework/sub-packages`.
-    - Publishes to Test PyPI.
-  - Dispatch the `bionemo-subpackage-ci.yml` workflow from GitHub Actions.
-    - Publishes to PyPI.
-    - Required: Input a comma-separated list of sub-packages you want to test and/or publish into `subpackages`.
-    - Optional: Set `publish` to `true` if you want to publish to PyPI. (Default: `false`)
+- **NVIDIA-Only**: Run the workflow! For more information, refer to: [Sub-Package GitHub Actions Workflow](#sub-package-github-actions-workflow)
 - **Optional**: Add `bionemo` as an owner or maintainer of the PyPI package if you want help maintaining it.
   - **Disclaimer**: If this is not done, and the package becomes dysfunctional, then NVIDIA / BioNeMo are not responsible for the health of the package or the sub-package source code, because we will not have the ability to deprecate versions, etc.
 
-### Sub-Package CI Workflow Key Points
+### Sub-Package GitHub Actions Workflow
+
+- Dispatch the `bionemo-subpackage-ci.yml` workflow from GitHub Actions to test, build, and publish your sub-packages to PyPI!
+  - Required: Input a comma-separated list of sub-packages you want to test and/or publish into `subpackages`.
+    - For example, `bionemo-moco,bionemo-llm,bionemo-webdatamodule`. The sub-packages will be tested and published in separate parallel environments.
+  - Optional: Set `test` to `true` if you want to test your sub-package. (Default: `true`)
+    - Sub-packages that require pre- or post- installation steps may require modification of the `install-and-test` job in [`bionemo-framework/.github/workflows/bionemo-subpackage-ci.yml`](../../../../.github/workflows/bionemo-subpackage-ci.yml).
+  - Optional: Set `publish` to `true` if you want to publish to Test PyPI or PyPI. (Default: `false`)
+    - Pre-Requisite: [BioNeMo Publishing to PyPI](#publishing-to-pypi)
+    - Publishing requires package building, but does not require testing for flexibility of package management.
+  - Optional: Publishes to Test PyPI by default. To publish to PyPI, check `Publish to PyPI instead of TestPyPI`.
+
+### FAQ
 
-- Individually `pip install`, `pytest`, and `python -m build` every supported sub-package before publishing to PyPI.
-- Triggered by PR changes (which publish to Test PyPI) or manually-dispatched GitHub Actions Workflows.
-  - TODO(@cspades): PR workflows are deactivated pending sufficient compute resources. To reactivate them in your PR, uncomment the PR trigger in [`bionemo-framework/.github/workflows/bionemo-subpackage-ci.yml`](../../../../.github/workflows/bionemo-subpackage-ci.yml).
-- `SUBPACKAGE_SKIP_TEST` and `SUBPACKAGE_SKIP_PUBLISH` in [`bionemo-framework/.github/workflows/bionemo-subpackage-ci.yml`](../../../../.github/workflows/bionemo-subpackage-ci.yml) control whether your sub-package is installed/tested or published to PyPI.
-- BioNeMo Pull Request Archive
-  - Prototype: https://github.com/NVIDIA/bionemo-framework/pull/725
+- What do I do if I want to test and publish two updated sub-packages that depend on each other?
+  - To deal with circular dependencies, publish one package to PyPI first, followed by testing and publishing the other. `pip` installs dependencies in reverse topological order, and will resolve / break circular dependencies as long as dependency conflicts do not exist. (If dependency conflicts exist, resolve them!)
+  - For example, if `A` depends on `B`, and `B` depends on `A`...
+    - Publish `B` to PyPI without testing. Untested sub-packages will be published with the version suffix `*-dev`.
+    - Set `A` to depend on the latest version (i.e. the `*-dev` version) of `B`.
+    - Test and publish `A` to PyPI.
+    - Test and publish `B` (which depends on the now-released `A`) to PyPI.
 
 ### TODO
 
 - Support building packages that have installation dependencies, such as `bionemo-noodles` dependent on `maturin` or `bionemo-<model>` dependent on `transformer-engine`.
-- Support unit tests that require GPU.
 - Automatically cut a release tag for the sub-package via GHA.

sub-packages/bionemo-amplify/src/bionemo/amplify/train_amplify.py

@@ -307,7 +307,7 @@ def main(
         save_top_k=save_top_k,
         every_n_train_steps=val_check_interval,
         always_save_context=True,  # Enables the .nemo file-like checkpointing where all IOMixins are under SerDe
-        filename="{epoch}-{val_loss:.2f}-{step}-{consumed_samples}",  # Including step and consumed_samples in the checkpoint filename prevents duplicate filenames and bugs related to this.
+        filename="{epoch}-{step}-{consumed_samples}",  # Including step and consumed_samples in the checkpoint filename prevents duplicate filenames and bugs related to this.
     )
 
     # Setup the logger and train the model

sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py

@@ -333,7 +333,7 @@ def main(
             every_n_train_steps=val_check_interval,
             always_save_context=True,
             # Enables the .nemo file-like checkpointing where all IOMixins are under SerDe
-            filename="{epoch}-{val_loss:.2f}-{step}-{consumed_samples}",
+            filename="{epoch}-{step}-{consumed_samples}",
             # Including step and consumed_samples in the checkpoint filename prevents duplicate filenames and bugs related to this.
         )
         callbacks.append(checkpoint_callback)

sub-packages/bionemo-llm/pyproject.toml

@@ -21,6 +21,11 @@ dependencies = [
     # TODO: Figure out how (or whether) to include apex & TE as part of the dependencies.
 ]
 
+[project.optional-dependencies]
+test = [
+    "bionemo-testing"
+]
+
 [tool.setuptools.packages.find]
 where = ["src"]
 include = ["bionemo.*"]