Expand test matrix to also install different computation backends (#90)

* also install torchvision and torchaudio

* test cu116

* fix candidate selection

* revert cu116

* try computation backend ci

* fix rocm install

* fix cu102

* disable rocm for now

Author: pmeier

.github/workflows/install.yml

@@ -9,7 +9,7 @@ on:
   pull_request:
 
 jobs:
-  torch_cpu:
+  channel:
     strategy:
       matrix:
         os:
@@ -49,14 +49,96 @@ jobs:
       - name: Collect environment information
         run: python scripts/collect_env.py
 
-      - name: Install torch
-        run: ltt install --cpuonly --pytorch-channel=${{ matrix.pytorch-channel }} torch
+      - name: Install PyTorch distributions
+        run:
+          ltt install --pytorch-channel=${{ matrix.pytorch-channel }} torch torchvision
+          torchaudio
 
       - name: Check if CPU only
         shell: python
         run: |
           import sys
           import torch
 
-          compiled_with_cuda = hasattr(torch._C, "_cuda_getDeviceCount")
-          sys.exit(compiled_with_cuda)
+          cuda = torch.version.cuda
+          print(f"cuda = {cuda}")
+
+          hip = torch.version.hip
+          print(f"hip = {hip}")
+
+          sys.exit(cuda or hip)
+
+  computation-backend:
+    strategy:
+      matrix:
+        os:
+          - ubuntu-latest
+          - windows-latest
+          - macos-latest
+        python-version:
+          - "3.7"
+          - "3.8"
+          - "3.9"
+          - "3.10"
+        pytorch-computation-backend:
+          - cpu
+          - cu113
+          - cu116
+        exclude:
+          - os: macos-latest
+            pytorch-computation-backend: cu113
+          - os: macos-latest
+            pytorch-computation-backend: cu116
+        include:
+          - os: ubuntu-latest
+            pytorch-computation-backend: cu102
+      # TODO: find a way to test this
+      #          - os: ubuntu-latest
+      #            pytorch-computation-backend: rocm5.1.1
+      fail-fast: false
+
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - name: Setup development environment
+        uses: ./.github/actions/setup-dev-env
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Collect environment information
+        run: python scripts/collect_env.py
+
+      - name: Install torch
+        run:
+          ltt install --pytorch-computation-backend=${{
+          matrix.pytorch-computation-backend }} torch==1.12.1
+
+      - name: Check computation backend
+        shell: python
+        run: |
+          import sys
+          import torch
+          from light_the_torch._cb import ComputationBackend, CUDABackend, ROCmBackend, CPUBackend
+
+          expected = ComputationBackend.from_str("${{ matrix.pytorch-computation-backend }}")
+
+          cuda = torch.version.cuda
+          print(f"cuda = {cuda}")
+
+          hip = torch.version.hip
+          print(f"hip = {hip}")
+
+          if cuda:
+              actual = CUDABackend.from_str(f"cu{torch.version.cuda}")
+          elif hip:
+              rocm = ".".join(torch.version.hip.split(".")[:2])
+              actual = ROCmBackend.from_str(f"rocm{rocm}")
+          else:
+              actual = CPUBackend()
+
+          sys.exit(actual != expected)

light_the_torch/_patch.py

@@ -278,6 +278,9 @@ def extract_local_specifier(candidate):
         return local
 
     def preprocessing(input):
+        if not input.candidates:
+            return
+
         candidates = iter(input.candidates)
         candidate = next(candidates)