From 14b6b077cd820c1cb07b57a90c2592cd749d99b1 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  | 29 +++++++++++++++++------------
 .github/workflows/tests.yml | 29 +++++++++++++++++------------
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 474517a3..55dee661 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,23 @@ 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//./}/$(uname -m)/cuda-keyring_1.0-1_all.deb"
+            sudo dpkg -i cuda-keyring.deb
+          )
+          sudo apt-get update && sudo apt-get install "cuda-minimal-build-${CUDA_PKG_SUFFIX}" --yes
+          echo "PATH=/usr/local/cuda/bin${PATH:+:${PATH}}" >> "${GITHUB_ENV}"
+          echo "LD_LIBRARY_PATH=/usr/local/cuda/lib64${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 67732041..cbf8c350 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,23 @@ 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//./}/$(uname -m)/cuda-keyring_1.0-1_all.deb"
+            sudo dpkg -i cuda-keyring.deb
+          )
+          sudo apt-get update && sudo apt-get install "cuda-minimal-build-${CUDA_PKG_SUFFIX}" --yes
+          echo "PATH=/usr/local/cuda/bin${PATH:+:${PATH}}" >> "${GITHUB_ENV}"
+          echo "LD_LIBRARY_PATH=/usr/local/cuda/lib64${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