tpu_unit_tests.yaml

name: CI with GCP TPU

on: [push]

jobs:
  test:
    runs-on: ubuntu-latest
    strategy:
      matrix:
        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:
          service_account_key: ${{ secrets.GCP_SA_KEY }}
          project_id: ${{ secrets.GCP_PROJECT_ID }}

      - name: Authenticate Google Cloud
        run: |
          printf "%s" "${{ secrets.GCP_SA_KEY }}" > ${HOME}/gcloud-service-key.json
          gcloud auth activate-service-account --key-file=${HOME}/gcloud-service-key.json
          gcloud config set project ${{ secrets.GCP_PROJECT_ID }}

      - name: Create VM and Run Script
        run: |
          export TPU_NAME=tpu-${{ github.run_id }}
          cd levanter
          infra/babysit-tpu-vm.sh $TPU_NAME -z ${{ matrix.tpu-zone }} -t v4-8 --preemptible -s infra/helpers/setup-tpu-vm-tests.sh -- \
            bash levanter/infra/run.sh pytest tests