Setup Docker for TPU execution and update infra scripts. #104
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |