tpu_unit_tests.yaml

name: CI with GCP TPU

on: [pull_request]

jobs:
  test:
    runs-on: ubuntu-latest
    env:
      TPU_ZONE: "us-central2-b"

    steps:
      - name: Checkout code
        uses: actions/checkout@v2

      - name: Set up Google Cloud SDK
        uses: google-github-actions/setup-gcloud@v1
        with:
          project_id: ${{ secrets.GCP_PROJECT_ID }}

      - name: Authenticate to Google Cloud
        uses: google-github-actions/auth@v1
        with:
          credentials_json: ${{ secrets.GCP_SA_KEY }}

      - name: Configure Google Cloud
        run: |
          gcloud config set project ${{ secrets.GCP_PROJECT_ID }}

      - name: Create VM
        run: |
          export TPU_NAME=ci-run-${{ github.run_id }}
          eval "$(ssh-agent -s)"
          TRUE_SHA=${{ github.event.pull_request.head.sha }}
          bash infra/spin-up-vm.sh $TPU_NAME -z ${TPU_ZONE} -t v4-8 --preemptible --retries 1

      - name: Run most tests
        run: |
          export TPU_NAME=ci-run-${{ github.run_id }}
          python infra/launch.py --foreground --tpu=$TPU_NAME --zone=$TPU_ZONE -- /opt/levanter/.venv/bin/pytest tests -m "not entry"
# Something's wrong with these
#
#      - name: Run forked tests
#        run: |
#          export TPU_NAME=ci-run-${{ github.run_id }}
#          gcloud compute tpus tpu-vm ssh $TPU_NAME --zone ${TPU_ZONE} --command "PYTHONPATH=$PYTHONPATH:levanter/tests bash levanter/infra/run.sh pytest --forked levanter/tests -m 'entry'"
#
      - name: Cleanup
        if: ${{ always() }}
        run: |
          export TPU_NAME=ci-run-${{ github.run_id }}
          echo gcloud compute tpus tpu-vm delete $TPU_NAME --zone ${TPU_ZONE} --quiet
          gcloud compute tpus tpu-vm delete $TPU_NAME --zone ${TPU_ZONE} --quiet