name: "BioNemo Image Build and Unit Tests"

on:
  push:
    branches:
      - main
      - "pull-request/[0-9]+"
      - "dependabot/**"
  merge_group:
    types: [checks_requested]
  schedule:
    - cron: "0 7 * * *" # Runs at 7 AM UTC daily (12 AM MST)

defaults:
  run:
    shell: bash -x -e -u -o pipefail {0}

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

jobs:
  pre-commit:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
        with:
          fetch-depth: 0
          submodules: "recursive"
      - uses: actions/setup-python@v5
        with:
          python-version: "3.12"
          cache: "pip"
      - run: pip install -r requirements-dev.txt
      - run: ./ci/scripts/static_checks.sh

  # With copy-pr-bot, we need to get the PR labels from the PR API rather than from the event metadata.
  get-pr-labels:
    runs-on: ubuntu-latest
    outputs:
      labels: ${{ steps.get-labels.outputs.labels }}
    steps:
      - name: Get PR number from branch
        if: startsWith(github.ref, 'refs/heads/pull-request/')
        id: get-pr-num
        run: |
          PR_NUM=$(echo ${{ github.ref_name }} | grep -oE '[0-9]+$')
          echo "pr_num=$PR_NUM" >> $GITHUB_OUTPUT

      - name: Get PR labels
        id: get-labels
        if: startsWith(github.ref, 'refs/heads/pull-request/')
        env:
          GH_TOKEN: ${{ github.token }}
        run: |
          LABELS=$(gh api repos/${{ github.repository }}/pulls/${{ steps.get-pr-num.outputs.pr_num }} --jq '[.labels[].name]' || echo "[]")
          echo "labels=$LABELS" >> $GITHUB_OUTPUT

      - name: Set empty labels for non-PR branches
        if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
        id: get-labels-empty
        run: echo "labels=[]" >> $GITHUB_OUTPUT

  build-bionemo-image:
    needs:
      - pre-commit
      - get-pr-labels
    runs-on: linux-amd64-cpu16
    if: ${{ !contains(fromJSON(needs.get-pr-labels.outputs.labels || '[]'), 'SKIP_CI') }}
    steps:
      - name: Login to Docker Hub
        uses: docker/login-action@v3
        with:
          username: ${{ vars.DOCKER_USERNAME }}
          password: ${{ secrets.DOCKER_PASSWORD }}

      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          submodules: "recursive"

      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3

      - name: Docker Metadata
        id: metadata
        uses: docker/metadata-action@v5
        with:
          images: svcbionemo023/bionemo-framework
          tags: |
            type=schedule
            type=ref,event=branch
            type=ref,event=tag
            type=ref,event=pr
            type=raw,value=${{ github.run_id }}

      # This action sets up our cache-from and cache-to flags appropriately; see the README of this action for more
      # info. It doesn't seem to cache correctly for merge_group events, so we need to add that as an extra argument in
      # the step below. There's probably a slight optimization to be had here by caching from the pr- caches for
      # merge_group events. See https://github.com/int128/docker-build-cache-config-action/issues/1005 for more info.
      - uses: int128/docker-build-cache-config-action@v1
        id: cache
        with:
          image: svcbionemo023/bionemo-build-cache

      - name: Build and push
        uses: docker/build-push-action@v5
        with:
          push: true
          tags: ${{ steps.metadata.outputs.tags }}
          labels: ${{ steps.metadata.outputs.labels }}
          cache-from: |
            ${{ steps.cache.outputs.cache-from }}
            type=registry,ref=svcbionemo023/bionemo-build-cache:main
          cache-to: ${{ steps.cache.outputs.cache-to }}


  run-tests:
    needs:
      - build-bionemo-image
      - get-pr-labels
    runs-on: linux-amd64-gpu-l4-latest-1
    container:
      image: svcbionemo023/bionemo-framework:${{ github.run_id }}
      credentials:
        username: ${{ vars.DOCKER_USERNAME }}
        password: ${{ secrets.DOCKER_PASSWORD }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Run tests
        # Tests in this stage generate code coverage metrics for the repository
        # Coverage data is uploaded to Codecov in subsequent stages
        env:
          BIONEMO_DATA_SOURCE: ngc
        run: ./ci/scripts/run_pytest.sh --no-nbval --skip-slow

      - name: Upload coverage to Codecov
        # Don't run coverage on merge queue or nightly CI to avoid duplicating reports
        # to codecov. See https://github.com/matplotlib/napari-matplotlib/issues/155
        if: github.event_name != 'merge_group' && github.event_name != 'schedule'
        uses: codecov/codecov-action@v5
        with:
          token: ${{ secrets.CODECOV_TOKEN }}

      - name: Upload test results to Codecov
        # Don't run coverage on merge queue or nightly CI to avoid duplicating reports
        # to codecov. See https://github.com/matplotlib/napari-matplotlib/issues/155
        if: ${{ !cancelled() && github.event_name != 'merge_group' && github.event_name != 'schedule' }}
        uses: codecov/test-results-action@v1
        with:
          token: ${{ secrets.CODECOV_TOKEN }}

  run-slow-tests:
    needs:
      - build-bionemo-image
      - get-pr-labels
    runs-on: linux-amd64-gpu-l4-latest-1
    container:
      image: svcbionemo023/bionemo-framework:${{ github.run_id }}
      credentials:
        username: ${{ vars.DOCKER_USERNAME }}
        password: ${{ secrets.DOCKER_PASSWORD }}
    if: |
      github.event_name == 'schedule' || github.event_name == 'merge_group' ||
      contains(fromJSON(needs.get-pr-labels.outputs.labels || '[]'), 'INCLUDE_SLOW_TESTS')
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Run slow tests
        env:
          BIONEMO_DATA_SOURCE: ngc
        # Not every sub-package has slow tests, and since some sub-packages have tests under the same name we need
        #  to run package by package like we do with the fast tests.
        run: ./ci/scripts/run_pytest.sh --no-nbval --only-slow --allow-no-tests


  run-notebooks-docs:
    needs:
      - build-bionemo-image
      - get-pr-labels
    runs-on: linux-amd64-gpu-l4-latest-1
    if: |
      github.event_name == 'schedule' || github.event_name == 'merge_group' ||
      contains(fromJSON(needs.get-pr-labels.outputs.labels || '[]'), 'INCLUDE_NOTEBOOKS_TESTS')
    container:
      image: svcbionemo023/bionemo-framework:${{ github.run_id }}
      credentials:
        username: ${{ vars.DOCKER_USERNAME }}
        password: ${{ secrets.DOCKER_PASSWORD }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - 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: |
          pytest -v --nbval-lax -p no:python docs/ sub-packages/

  verify-tests-status:
    # Base on the status of this job, the unit-tests workflow succeeds or fails
    # This steps checks the status of all test jobs and fails if any of them failed or were cancelled.
    # It is a workaround for the lack of a built-in feature to finalize a pipeline by checking the status of multiple jobs
    needs: # List all your run-*-test jobs
      - run-tests
      - run-slow-tests
      - run-notebooks-docs
      # Add all other run-*-test jobs
    runs-on: ubuntu-latest
    if: always() # This ensures the job runs even if previous jobs fail
    steps:
      - name: Check test job statuses
        run: |
          if [[ "${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') }}" == "true" ]]; then
            echo "Some test jobs have failed or been cancelled!"
            exit 1
          else
            echo "All test jobs have completed successfully or been skipped!"
            exit 0
          fi
