cleanup the duplicate logic of WITH_CUDA and ACCELERATOR

Author: stefdoerr

.github/workflows/docs.yml

@@ -26,10 +26,10 @@ jobs:
 
       - name: Install docs dependencies
         run: |
-          pip install -vv .
+          pip install -vv . --no-deps
           pip install -r docs/requirements.txt
         env:
-          WITH_CUDA: "0"
+          ACCELERATOR: "cpu"
 
       - name: Build Sphinx Documentation
         run: |

.github/workflows/publish.yml

@@ -59,15 +59,13 @@ jobs:
         env:
           ACCELERATOR: ${{ matrix.accelerator }}
           CPU_TRAIN: ${{ runner.os == 'macOS' && 'true' || 'false' }}
-          WITH_CUDA: ${{ matrix.accelerator != 'cpu' && '1' || '0' }}
 
       - name: Build wheels
         if: matrix.os == 'windows-2019'
         shell: cmd # Use cmd on Windows to avoid bash environment taking priority over MSVC variables
         run: python -m cibuildwheel --output-dir wheelhouse
         env:
           ACCELERATOR: ${{ matrix.accelerator }}
-          WITH_CUDA: ${{ matrix.accelerator != 'cpu' && '1' || '0' }}
           DISTUTILS_USE_SDK: "1" # Windows requires this to use vc for building
           SKIP_TORCH_COMPILE: "true"
 

.github/workflows/test.yml

@@ -73,7 +73,7 @@ jobs:
           python -m build            
           pip install dist/*.whl
         env:
-          WITH_CUDA: "0"
+          ACCELERATOR: "cpu"
 
       # - name: Install nnpops
       #   if: matrix.os == 'ubuntu-latest' || matrix.os == 'macos-latest'
@@ -85,7 +85,7 @@ jobs:
       - name: Run tests
         run: pytest -v -s --durations=10
         env:
-          WITH_CUDA: "0"
+          ACCELERATOR: "cpu"
           SKIP_TORCH_COMPILE: ${{ runner.os == 'Windows' && 'true' || 'false' }}
           OMP_PREFIX: ${{ runner.os == 'macOS' && '/Users/runner/miniconda3/envs/test' || '' }}
           CPU_TRAIN: ${{ runner.os == 'macOS' && 'true' || 'false' }}

cibuildwheel_support/before_all_linux.sh

@@ -30,8 +30,6 @@ if [ "$ACCELERATOR" == "cu118" ]; then
     ln -s /opt/rh/gcc-toolset-11/root/usr/bin/g++ /usr/local/cuda/bin/g++
 
     export CUDA_HOME="/usr/local/cuda"
-    export WITH_CUDA=1
-
 
     # Configure pip to use PyTorch extra-index-url for CUDA 11.8
     mkdir -p $HOME/.config/pip

setup.py

@@ -9,16 +9,9 @@
 import platform
 
 
-# If WITH_CUDA is defined
-env_with_cuda = os.getenv("WITH_CUDA")
-if env_with_cuda is not None:
-    if env_with_cuda not in ("0", "1"):
-        raise ValueError(
-            f"WITH_CUDA environment variable got invalid value {env_with_cuda}. Expected '0' or '1'"
-        )
-    use_cuda = env_with_cuda == "1"
-else:
-    use_cuda = torch.cuda._is_compiled()
+use_cuda = (
+    os.environ.get("ACCELERATOR", "").startswith("cu") or torch.cuda._is_compiled()
+)
 
 
 def set_torch_cuda_arch_list():
@@ -56,8 +49,8 @@ def set_torch_cuda_arch_list():
     ]
 
 extra_deps = []
-if use_cuda:
-    cuda_ver = os.getenv("ACCELERATOR", "")[2:4]
+if os.getenv("ACCELERATOR", "").startswith("cu"):
+    cuda_ver = os.getenv("ACCELERATOR")[2:4]
     extra_deps = [f"nvidia-cuda-runtime-cu{cuda_ver}"]
 
 ExtensionType = CppExtension if not use_cuda else CUDAExtension