Repository: allenai/tango (language: python)

Issue:
Use Beaker to run GPU tests
## Tasks
- [x] Set up a scalable, preemptible,  cheap cloud cluster on Beaker: https://beaker.org/cl/ai2/tango-gpu-tests/details
- [x] Set up a Beaker workspace: https://beaker.org/ws/ai2/tango-testing
- [ ] #183
- [x] #184 (merged into 183)


Discussion / hints:

Fix patch:
diff --git a/.dockerignore b/.dockerignore
index ac99a4f97..fccd1c6f0 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -7,7 +7,6 @@
 .mypy_cache
 docs
 examples
-scripts
 tests
 test_fixtures
 integration_tests
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 7a3d0e32c..b28dc719c 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -216,10 +216,88 @@ jobs:
           . .venv/bin/activate
           pip uninstall -y ai2-tango
 
+  gpu_tests:
+    name: GPU Tests
+    runs-on: [ubuntu-latest]
+    env:
+      BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
+      BEAKER_WORKSPACE: ai2/tango-testing
+      IMAGE_NAME: tango-testing
+    steps:
+      - uses: actions/checkout@v2
+
+      - 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: Install beaker client
+        shell: bash
+        run: |
+          mkdir -p "$HOME/bin"
+
+          # Download and install from latest GitHub release.
+          curl -s https://api.github.com/repos/allenai/beaker/releases/latest \
+            | grep 'browser_download_url.*linux' \
+            | cut -d '"' -f 4 \
+            | wget -qi - \
+          && tar -xvzf beaker_linux.tar.gz -C "$HOME/bin"
+
+          # Add to path.
+          echo "$HOME/bin" >> "$GITHUB_PATH"
+
+      - name: Verify beaker install
+        run: |
+          beaker account whoami
+
+      - name: Create beaker experiment config
+        run: |
+          cat >beaker_config.yml << EOL
+          version: v2-alpha
+          description: GPU Tests
+          tasks:
+            - name: tests
+              image:
+                beaker: petew/tango-testing
+              command: ["/entrypoint.sh", "pytest", "-v", "-m", "gpu", "tests/"]
+              envVars:
+                - name: COMMIT_SHA
+                  value: $COMMIT_SHA
+              result:
+                path: '/results'
+              resources:
+                gpuCount: 2
+              context:
+                cluster: ai2/tango-gpu-tests
+                priority: normal
+          EOL
+          cat beaker_config.yml
+
+      - name: Submit beaker job
+        run: |
+          EXPERIMENT=$(beaker experiment create beaker_config.yml --workspace $BEAKER_WORKSPACE --name gpu-tests-$COMMIT_SHA | awk '{print $2}')
+          echo "EXPERIMENT=$EXPERIMENT" >> $GITHUB_ENV
+          echo "Experiment $EXPERIMENT submitted. See progress at https://beaker.org/ex/$EXPERIMENT"
+
+      - name: Wait for job to finish
+        run: |
+          beaker experiment await $EXPERIMENT tests finalized --timeout 10m
+
+      - name: Get logs
+        if: always()
+        run: |
+          beaker experiment results $EXPERIMENT --prefix out.log --output results
+          cat results/tests/out.log
+
   release:
     name: Release
     runs-on: ubuntu-latest
-    needs: [checks]
+    needs: [checks, gpu_tests]
     if: startsWith(github.ref, 'refs/tags/')
     steps:
       - uses: actions/checkout@v1
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 005aa963c..75c89fb78 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added a '@requires_gpus' decorator for marking tests as needing GPUs. Tests marked with this will be run in the "GPU Tests" workflow
+  on dual k80 GPUs via Beaker.
+
+
 ## [v0.5.0](https://github.com/allenai/tango/releases/tag/v0.5.0) - 2022-02-09
 
 ### Added
diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh
new file mode 100644
index 000000000..b20ab0bcb
--- /dev/null
+++ b/scripts/entrypoint.sh
@@ -0,0 +1,22 @@
+#!/bin/sh
+
+# Exit script if any commands fail.
+set -e
+
+# Check that the environment variable has been set correctly
+if [ -z "$COMMIT_SHA" ]; then
+  echo >&2 'error: missing COMMIT_SHA environment variable'
+  exit 1
+fi
+
+# Clone and install tango.
+git clone https://github.com/allenai/tango.git
+cd tango
+git checkout "$COMMIT_SHA"
+/opt/conda/bin/pip install --no-cache-dir '.[dev,all]'
+
+# Create directory for results.
+mkdir -p /results
+
+# Execute the arguments to this script as commands themselves, piping output into a log file.
+exec "$@" | tee /results/out.log


Test patch:
diff --git a/.github/workflows/docker_testing.yml b/.github/workflows/docker_testing.yml
new file mode 100644
index 000000000..a49de2755
--- /dev/null
+++ b/.github/workflows/docker_testing.yml
@@ -0,0 +1,80 @@
+# This workflow is just for building our Docker image for GPU testing on Beaker,
+# and pushing it to Beaker. We only run it when the relevant Dockerfile (or .dockerignore) changes.
+name: Docker testing
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+on:
+  pull_request:
+    branches:
+      - main
+    paths:
+      - 'Dockerfile.test'
+      - '.dockerignore'
+  push:
+    branches:
+      - main
+    paths:
+      - 'Dockerfile.test'
+      - '.dockerignore'
+
+jobs:
+  build:
+    name: Build
+    runs-on: ubuntu-latest
+    env:
+      BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
+      BEAKER_WORKSPACE: ai2/tango-testing
+      IMAGE_NAME: tango-testing
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Build Docker image
+        run: |
+          docker build -t "$IMAGE_NAME" -f Dockerfile.test .
+
+      - 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: Test Docker image
+        run: |
+          docker run --rm --env COMMIT_SHA="$COMMIT_SHA" "$IMAGE_NAME" tango info
+
+      - name: Install beaker client
+        if: github.event_name != 'pull_request'
+        shell: bash
+        run: |
+          mkdir -p "$HOME/bin"
+
+          # Download and install from latest GitHub release.
+          curl -s https://api.github.com/repos/allenai/beaker/releases/latest \
+            | grep 'browser_download_url.*linux' \
+            | cut -d '"' -f 4 \
+            | wget -qi - \
+          && tar -xvzf beaker_linux.tar.gz -C "$HOME/bin"
+
+          # Add to path.
+          echo "$HOME/bin" >> "$GITHUB_PATH"
+
+      - name: Verify beaker install
+        if: github.event_name != 'pull_request'
+        run: |
+          beaker account whoami
+
+      - name: Push release to beaker
+        if: github.event_name != 'pull_request'
+        run: |
+          # We delete the previous version. This doesn't actually delete the backing Docker image, so
+          # we'll still benefit from layer caching when we push new version. The image might not exist
+          # yet, so it's ok if this fails.
+          beaker image delete ai2/"$IMAGE_NAME" || true
+          beaker image create --workspace "$BEAKER_WORKSPACE" --name "$IMAGE_NAME" "$IMAGE_NAME"
diff --git a/Dockerfile.test b/Dockerfile.test
new file mode 100644
index 000000000..9ee788320
--- /dev/null
+++ b/Dockerfile.test
@@ -0,0 +1,14 @@
+# This Dockerfile is for building an image suitable for running tango's GPU tests.
+# There are no instructions in the Dockerfile that install tango. Instead, the entrypoint
+# script handles installing tango from a particular commit at runtime, based on the environment
+# variable "COMMIT_SHA". That way we don't need to rebuild and push the image each time we run
+# GPU tests.
+
+FROM ghcr.io/allenai/pytorch:1.10.2-cuda11.3
+
+COPY scripts/entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
+WORKDIR /testing
+
+ENTRYPOINT ["/entrypoint.sh"]
diff --git a/pytest.ini b/pytest.ini
index a8a81cf0d..c6c219aa3 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -4,6 +4,7 @@ python_classes = Test* *Test
 log_format = %(asctime)s - %(levelname)s - %(name)s - %(message)s
 log_level = DEBUG
 markers =
+    gpu: marks tests that need GPUs
 filterwarnings =
     ignore:.*Consider increasing the value of the `num_workers` argument.*:UserWarning:pytorch_lightning\.trainer\.data_loading
     ignore:.*you defined a validation_step but have no val_dataloader.*:UserWarning:pytorch_lightning\.trainer\.configuration_validator
diff --git a/tango/common/testing.py b/tango/common/testing.py
index 98d28f502..da28718fe 100644
--- a/tango/common/testing.py
+++ b/tango/common/testing.py
@@ -105,3 +105,17 @@ def run_experiment(
     finally:
         test_case.teardown_method()
         teardown_logging()
+
+
+def requires_gpus(test_method):
+    """
+    Decorator to indicate that a test requires multiple GPU devices.
+    """
+    import pytest
+    import torch
+
+    return pytest.mark.gpu(
+        pytest.mark.skipif(torch.cuda.device_count() < 2, reason="2 or more GPUs required.")(
+            test_method
+        )
+    )
diff --git a/test_fixtures/beaker/nvidia_smi.yml b/test_fixtures/beaker/nvidia_smi.yml
new file mode 100644
index 000000000..49f16912a
--- /dev/null
+++ b/test_fixtures/beaker/nvidia_smi.yml
@@ -0,0 +1,16 @@
+# Used to test that GPUs in a cluster are available. Submit this to beaker with:
+# $ beaker experiment create test_fixtures/beaker/nvidia_smi.yml --workspace ai2/tango-testing --name tango-test-1
+version: v2-alpha
+description: NvidiaSMI
+tasks:
+  - name: nvidia-smi
+    image:
+      docker: nvidia/cuda:11.0-base
+    command: [nvidia-smi]
+    result:
+      path: '/unused'
+    resources:
+      gpuCount: 2
+    context:
+      cluster: ai2/tango-gpu-tests
+      priority: normal
diff --git a/tests/integrations/fairscale/train_test.py b/tests/integrations/fairscale/train_test.py
index 949d56478..522bbcec2 100644
--- a/tests/integrations/fairscale/train_test.py
+++ b/tests/integrations/fairscale/train_test.py
@@ -21,7 +21,12 @@ def teardown_method(self):
             pytest.param(
                 True,
                 id="fsdp=True",
-                marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA"),
+                marks=[
+                    pytest.mark.gpu,
+                    pytest.mark.skipif(
+                        torch.cuda.device_count() < 2, reason="Requires CUDA devices"
+                    ),
+                ],
             ),
             pytest.param(False, id="fsdp=False"),
         ),
@@ -39,7 +44,12 @@ def teardown_method(self):
             pytest.param(
                 True,
                 id="amp=True",
-                marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA"),
+                marks=[
+                    pytest.mark.gpu,
+                    pytest.mark.skipif(
+                        torch.cuda.device_count() < 2, reason="Requires CUDA devices"
+                    ),
+                ],
             ),
             pytest.param(False, id="amp=False"),
         ),