main.yml

name: Main

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}

on:
  pull_request:
    branches:
      - "*"
  push:
    branches:
      - main
    tags:
      - "v*.*.*"

env:
  CACHE_PREFIX: v5 # Change this to invalidate existing cache.
  PYTHON_PATH: ./
  DEFAULT_PYTHON: 3.9
  WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
  BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
  BEAKER_WORKSPACE: ai2/tango-testing
  BEAKER_IMAGE: petew/tango-testing
  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

jobs:
  checks:
    name: python ${{ matrix.python }} - ${{ matrix.task.name }}
    runs-on: [ubuntu-latest]
    timeout-minutes: 30
    permissions:
      contents: "read"
      id-token: "write"
    strategy:
      fail-fast: false
      matrix:
        python: ["3.9"]
        task:
          - name: Lint
            extras: dev,all
            requires_torch: true
            run: |
              ruff check .

          - name: Type check
            extras: dev,all
            requires_torch: true
            run: |
              mypy --check-untyped-defs .

          - name: Build
            extras: dev,all
            requires_torch: true
            run: |
              tango --version
              python -m build

          - name: Style
            extras: dev
            requires_torch: false
            run: |
              isort --check .
              black --check .

          - name: Docs
            extras: dev,all
            requires_torch: true
            run: |
              cd docs && make html SPHINXOPTS="-W --keep-going"

          - name: Test
            extras: dev
            requires_torch: false
            run: |
              pytest -v --durations=10 --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/

          - name: Datasets integration
            extras: dev,datasets
            requires_torch: false
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/datasets tests/integrations/datasets

          - name: PyTorch integration
            extras: dev,torch
            requires_torch: true
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/torch tests/integrations/torch

          - name: Transformers integration
            extras: dev,flax,transformers
            requires_torch: true
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/transformers tests/integrations/transformers

          - name: FairScale integration
            extras: dev,fairscale
            requires_torch: true
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/fairscale tests/integrations/fairscale

          - name: W&B integration
            extras: dev,torch,flax,wandb
            requires_torch: true
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/wandb tests/integrations/wandb

          - name: Beaker integration
            extras: dev,beaker
            requires_torch: false
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/beaker tests/integrations/beaker

          - name: Flax integration
            extras: dev,flax,transformers
            requires_torch: false
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/flax tests/integrations/flax

          - name: GS integration
            extras: dev,gs
            requires_torch: false
            run: |
              pytest -v --color=yes --doctest-modules tango/integrations/gs tests/integrations/gs

          - name: Example - train_lm
            extras: dev,all
            requires_torch: true
            run: |
              cd examples/train_lm
              pytest -v --color=yes test.py

        include:
          # Run the core tests on other Python versions as well.
          - task:
              name: Test
              extras: dev
              requires_torch: false
              run: |
                pytest -v --durations=10 --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
            python: "3.8"

          - task:
              name: Test
              extras: dev
              requires_torch: false
              run: |
                pytest -v --durations=10 --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
            python: "3.10"

    steps:
      - uses: "actions/checkout@v3"
      - name: Checkout
        if: github.event_name != 'pull_request'
        uses: actions/checkout@v3

      # For pull requests we need to checkout the HEAD commit instead of the merge
      # commit since some tests depend on having an existing commit.
      - name: Checkout (pull request)
        if: github.event_name == 'pull_request'
        uses: actions/checkout@v3
        with:
          ref: ${{ github.event.pull_request.head.sha }}

      - name: Setup Python
        uses: actions/setup-python@v4
        with:
          python-version: ${{ matrix.python }}

      - name: Install prerequisites
        run: |
          pip install --upgrade pip setuptools wheel virtualenv

      - name: Set build variables
        shell: bash
        run: |
          set -e
          # Get the exact Python version to use in the cache key.
          echo "PYTHON_VERSION=$(python --version)" >> $GITHUB_ENV
          echo "RUNNER_ARCH=$(uname -m)" >> $GITHUB_ENV
          # Use week number in cache key so we can refresh the cache weekly.
          echo "WEEK_NUMBER=$(date +%V)" >> $GITHUB_ENV
          echo "EXTRAS_HASH=$(python scripts/hash_extras.py ${{ matrix.task.extras }})" >> $GITHUB_ENV

      - uses: actions/cache@v3
        id: virtualenv-cache
        with:
          path: .venv
          key: ${{ env.CACHE_PREFIX }}-${{ env.WEEK_NUMBER }}-${{ runner.os }}-${{ env.RUNNER_ARCH }}-${{ env.PYTHON_VERSION }}-${{ env.EXTRAS_HASH }}-${{ hashFiles('pyproject.toml') }}

      - name: Setup virtual environment (no cache hit)
        if: steps.virtualenv-cache.outputs.cache-hit != 'true'
        run: |
          test -d .venv || virtualenv -p $(which python) --copies --reset-app-data .venv

      # Reference: https://github.com/marketplace/actions/authenticate-to-google-cloud#setup
      - name: Authenticate to Google Cloud
        if: matrix.task.name == 'GS integration'
        uses: "google-github-actions/auth@v1"
        with:
          workload_identity_provider: "projects/10554368204/locations/global/workloadIdentityPools/tango-ci-pool/providers/tango-ci-provider"
          service_account: "tango-service@ai2-allennlp.iam.gserviceaccount.com"

      - name: Pre-install torch
        if: steps.virtualenv-cache.outputs.cache-hit != 'true' && (contains(matrix.task.extras, 'torch') || contains(matrix.task.extras, 'all') || matrix.task.requires_torch)
        run: |
          . .venv/bin/activate
          pip install torch==2.0.0 --extra-index-url https://download.pytorch.org/whl/cpu

      - name: Pre-install flax
        if: steps.virtualenv-cache.outputs.cache-hit != 'true' && (contains(matrix.task.extras, 'flax') || contains(matrix.task.extras, 'all'))
        run: |
          . .venv/bin/activate
          pip install flax jax jaxlib "tensorflow-cpu>=2.9.1" optax

      - name: Install editable (no cache hit)
        if: steps.virtualenv-cache.outputs.cache-hit != 'true'
        run: |
          . .venv/bin/activate
          pip install -e .[${{ matrix.task.extras }}]

      - name: Install editable (cache hit)
        if: steps.virtualenv-cache.outputs.cache-hit == 'true'
        run: |
          . .venv/bin/activate
          pip install --no-deps -e .[${{ matrix.task.extras }}]

      - name: Show environment info
        run: |
          . .venv/bin/activate
          echo "========= Python location ==========="
          which python
          echo "========= Python version ============"
          python --version
          echo "========= Python packages ==========="
          pip freeze
          echo "========= Tango installation ========"
          tango info

      - name: ${{ matrix.task.name }}
        run: |
          . .venv/bin/activate
          ${{ matrix.task.run }}

      - name: Upload package distribution files
        if: matrix.task.name == 'Build' && matrix.python == env.DEFAULT_PYTHON
        uses: actions/upload-artifact@v3
        with:
          name: package
          path: dist

      - name: Upload docs build
        if: matrix.task.name == 'Docs' && matrix.python == env.DEFAULT_PYTHON
        uses: actions/upload-artifact@v3
        with:
          name: docs
          path: docs/build

      - name: Clean up
        if: always()
        run: |
          . .venv/bin/activate
          pip uninstall -y ai2-tango

  gpu_tests:
    name: GPU Tests
    runs-on: ubuntu-latest
    steps:
      - name: Determine current commit SHA (pull request)
        if: github.event_name == 'pull_request'
        run: |
          echo "COMMIT_SHA=${{ github.event.pull_request.head.sha }}" >> $GITHUB_ENV

      - name: Determine current commit SHA (push)
        if: github.event_name != 'pull_request'
        run: |
          echo "COMMIT_SHA=$GITHUB_SHA" >> $GITHUB_ENV

      - name: GPU Tests
        uses: allenai/beaker-run-action@v1.2
        with:
          spec: |
            version: v2
            description: GPU Tests
            budget: ai2/oe-training
            tasks:
              - name: tests
                image:
                  beaker: ${{ env.BEAKER_IMAGE }}
                context:
                  preemptible: true 
                resources:
                  gpuCount: 2
                envVars:
                  - name: COMMIT_SHA
                    value: ${{ env.COMMIT_SHA }}
                command: ["/entrypoint.sh", "pytest", "-v", "-m", "gpu", "tests/"]
                result:
                  path: /unused
          token: ${{ secrets.BEAKER_TOKEN }}
          workspace: ${{ env.BEAKER_WORKSPACE }}
          clusters: ai2/general-cirrascale,ai2/allennlp-cirrascale,ai2/aristo-cirrascale,ai2/mosaic-cirrascale,ai2/s2-cirrascale

  release:
    name: Release
    runs-on: ubuntu-latest
    needs: [gpu_tests, checks]
    if: startsWith(github.ref, 'refs/tags/')
    steps:
      - uses: actions/checkout@v3
        with:
          fetch-depth: 0

      - name: Setup Python
        uses: actions/setup-python@v4
        with:
          python-version: ${{ env.DEFAULT_PYTHON }}

      - name: Install requirements
        run: |
          pip install -e .[dev]

      - name: Prepare environment
        run: |
          echo "RELEASE_VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_ENV
          echo "TAG=${GITHUB_REF#refs/tags/}" >> $GITHUB_ENV

      - name: Download package distribution files
        uses: actions/download-artifact@v3
        with:
          name: package
          path: dist

      - name: Generate release notes
        run: |
          python scripts/release_notes.py > ${{ github.workspace }}-RELEASE_NOTES.md

      - name: Publish package to PyPI
        run: |
          twine upload -u __token__ -p ${{ secrets.PYPI_PASSWORD }} dist/*

      - name: Publish GitHub release
        uses: softprops/action-gh-release@v1
        with:
          body_path: ${{ github.workspace }}-RELEASE_NOTES.md
          prerelease: ${{ contains(env.TAG, 'rc') }}
          files: |
            dist/*