name: BioNeMo Sub-Package Workflow

on:
  # 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. 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
      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:
    name: "[Configure Workflow Packages] Identify sub-packages for testing and publishing."
    runs-on: ubuntu-latest
    outputs:
      workflow_packages: ${{ steps.parse-dispatch-packages.outputs.dispatch_packages }}
    steps:
      - 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 "[BioNeMo Sub-Package CI] Sub-packages to stage: $dispatch_packages"

  install-and-test:
    needs: configure-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."
    # 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 -qyy git curl lsb-release build-essential
      - 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: ${{ 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 uv maturin
          # Install required core & optional [test] dependencies.
          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.
        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: [configure-workflow-packages, install-and-test]
    if: ${{ needs.configure-workflow-packages.outputs.workflow_packages != '[]' && github.event.inputs.publish == 'true' }}
    outputs:
      staged_packages: ${{ 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 }}] Build the sub-package."
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
        with:
          persist-credentials: false
      - uses: actions/setup-python@v5
        with:
          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" && "${{ 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
          # 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 }}-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]
    # 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.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
    steps:
      - id: download-distribution
        name: Download the built distribution.
        uses: actions/download-artifact@v4
        with:
          name: ${{ matrix.package }}-build-artifacts
          path: sub-packages/${{ matrix.package }}/dist
      - id: publish-to-testpypi
        name: Publish distribution 📦 to Test PyPI for PR.
        if: ${{ github.event.inputs.pypi == 'false' }}
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          verbose: true
          packages-dir: sub-packages/${{ matrix.package }}/dist
          repository-url: https://test.pypi.org/legacy/
          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'
        # 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
          packages-dir: sub-packages/${{ matrix.package }}/dist
          skip-existing: ${{ github.event.inputs.version_overwrite }}
