ci: switch to amazon runners

Author: d4l3k

.github/workflows/unittest.yaml

@@ -11,32 +11,24 @@ jobs:
     strategy:
       matrix:
         os:
-          - ubuntu-20.04
+          - "linux.2xlarge"
           - "linux.4xlarge.nvidia.gpu"
-    runs-on: ${{ matrix.os }}
-    steps:
-      - name: Setup Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: "3.10"
-          architecture: x64
-      - name: Checkout
-        uses: actions/checkout@v2
-      - name: Install Dependencies
-        run: |
-          set -eux
 
-          sudo apt-get install -y protobuf-compiler
+  uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      timeout: 120
+      runner: ${{ matrix.runs-on }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      script: |
+        conda create -n venv python=3.10 -y
+        conda activate venv
 
-          pip install -e .[dev] -v
-      - name: Run Python Tests
-        run: |
-          set -eux
+        yum install -y protobuf-compiler
 
-          pytest -v
-      - name: Run Rust Lint
-        run: |
-          set -eux
+        python -m pip install --upgrade pip
+        pip install -e .[dev] -v
 
-          cargo test -v
+        pytest -v
+        cargo test -v