From 562381b3373d77cc683bb44e2bd4ccb640316982 Mon Sep 17 00:00:00 2001
From: Xuehai Pan <XuehaiPan@pku.edu.cn>
Date: Sun, 12 Feb 2023 12:57:54 +0000
Subject: [PATCH] refactor(workflows): rewrite setup CUDA Toolkit logic

---
 .github/workflows/lint.yml  | 27 +++++++++++++++------------
 .github/workflows/tests.yml | 27 +++++++++++++++------------
 2 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 474517a34..51e54dbd4 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -15,6 +15,9 @@ concurrency:
   group: "${{ github.workflow }}-${{ github.ref }}"
   cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
+env:
+  CUDA_VERSION: "11.7"
+
 jobs:
   lint:
     runs-on: ubuntu-latest
@@ -33,21 +36,21 @@ jobs:
           update-environment: true
 
       - name: Setup CUDA Toolkit
-        uses: Jimver/cuda-toolkit@v0.2.8
         id: cuda-toolkit
-        with:
-          cuda: "11.7.0"
-          method: network
-          sub-packages: '["nvcc"]'
-      - run: |
-          CUDA_VERSION="${{steps.cuda-toolkit.outputs.cuda}}"
-          echo "CUDA_VERSION=${CUDA_VERSION}" >> "${GITHUB_ENV}"
-          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
+        run: |
+          CUDA_PKG_SUFFIX="$(echo "${CUDA_VERSION}" | cut -d'.' -f-2 | tr '.' '-')"
+          sudo apt-get update && sudo apt-get install wget --yes
+          source /etc/os-release && \
+            wget -O cuda-keyring.deb "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${VERSION_ID//./}/x86_64/cuda-keyring_1.0-1_all.deb"
+          sudo dpkg -i cuda-keyring.deb
+          sudo apt-get update && sudo apt-get install "cuda-nvcc-${CUDA_PKG_SUFFIX}"
+          echo "PATH=/usr/local/cuda/bin${PATH:+:${PATH}}" >> "${GITHUB_ENV}"
+          echo "LD_LIBRARY_PATH=/usr/local/cuda/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" >> "${GITHUB_ENV}"
+
+          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_PKG_SUFFIX}" | tr -d '-')"
           echo "PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}" >> "${GITHUB_ENV}"
-
           echo "Installed CUDA version is: ${CUDA_VERSION}"
-          echo "CUDA install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
-          nvcc -V
+          /usr/local/cuda/bin/nvcc -V
           echo "Torch index URL: ${PIP_EXTRA_INDEX_URL}"
 
       - name: Upgrade pip
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 677320418..8f209109a 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -26,6 +26,9 @@ concurrency:
   group: "${{ github.workflow }}-${{ github.ref }}"
   cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
+env:
+  CUDA_VERSION: "11.7"
+
 jobs:
   test:
     name: Test with CXX/CUDA extensions on ubuntu-latest
@@ -45,21 +48,21 @@ jobs:
           update-environment: true
 
       - name: Setup CUDA Toolkit
-        uses: Jimver/cuda-toolkit@v0.2.8
         id: cuda-toolkit
-        with:
-          cuda: "11.7.0"
-          method: network
-          sub-packages: '["nvcc"]'
-      - run: |
-          CUDA_VERSION="${{steps.cuda-toolkit.outputs.cuda}}"
-          echo "CUDA_VERSION=${CUDA_VERSION}" >> "${GITHUB_ENV}"
-          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
+        run: |
+          CUDA_PKG_SUFFIX="$(echo "${CUDA_VERSION}" | cut -d'.' -f-2 | tr '.' '-')"
+          sudo apt-get update && sudo apt-get install wget --yes
+          source /etc/os-release && \
+            wget -O cuda-keyring.deb "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${VERSION_ID//./}/x86_64/cuda-keyring_1.0-1_all.deb"
+          sudo dpkg -i cuda-keyring.deb
+          sudo apt-get update && sudo apt-get install "cuda-nvcc-${CUDA_PKG_SUFFIX}"
+          echo "PATH=/usr/local/cuda/bin${PATH:+:${PATH}}" >> "${GITHUB_ENV}"
+          echo "LD_LIBRARY_PATH=/usr/local/cuda/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" >> "${GITHUB_ENV}"
+
+          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_PKG_SUFFIX}" | tr -d '-')"
           echo "PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}" >> "${GITHUB_ENV}"
-
           echo "Installed CUDA version is: ${CUDA_VERSION}"
-          echo "CUDA install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
-          nvcc -V
+          /usr/local/cuda/bin/nvcc -V
           echo "Torch index URL: ${PIP_EXTRA_INDEX_URL}"
 
       - name: Upgrade pip