From 0a6a67360faeae00fa7eb261778dd63a8955597d Mon Sep 17 00:00:00 2001
From: kp992 <kpl992@outlook.com>
Date: Mon, 27 May 2024 21:55:08 +0530
Subject: [PATCH] use script

---
 .github/workflows/ci.yml | 14 ++++-------
 setup_cuda.sh            | 51 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 9 deletions(-)
 create mode 100644 setup_cuda.sh

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 9c781a1..c381df1 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -19,15 +19,11 @@ jobs:
           python-version: "3.11"
           environment-file: environment.yml
           activate-environment: quantecon   
-      - name: Install CUDA Toolkit
-        id: cuda-toolkit
-        uses: Jimver/cuda-toolkit@v0.2.15
-        with:
-          cuda: '12.4.1'
-      # - name: Install CUDA
-      #   shell: bash -l {0}
-      #   run: |
-      #     conda install cudatoolkit cuda-nvcc -c nvidia -c anaconda -y
+      - name: Install CUDA
+        shell: bash -l {0}
+        run: |
+          chmod +x setup_cuda.sh
+          ./setup_cuda.sh
       - name: Install JAX[CUDA] and Numpyro[CUDA]
         shell: bash -l {0}
         run: |
diff --git a/setup_cuda.sh b/setup_cuda.sh
new file mode 100644
index 0000000..8615c03
--- /dev/null
+++ b/setup_cuda.sh
@@ -0,0 +1,51 @@
+#!/bin/bash
+
+# Exit script on any error
+set -e
+
+# Define variables
+NV_CUDA_LIB_VERSION="12.4.1-1"
+NV_CUDA_CUDART_DEV_VERSION="12.4.127-1"
+NV_NVML_DEV_VERSION="12.4.127-1"
+NV_LIBCUSPARSE_DEV_VERSION="12.3.1.170-1"
+NV_LIBNPP_DEV_VERSION="12.2.5.30-1"
+NV_LIBNPP_DEV_PACKAGE="libnpp-dev-12-4=${NV_LIBNPP_DEV_VERSION}"
+NV_LIBCUBLAS_DEV_PACKAGE_NAME="libcublas-dev-12-4"
+NV_LIBCUBLAS_DEV_VERSION="12.4.5.8-1"
+NV_LIBCUBLAS_DEV_PACKAGE="${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}"
+NV_CUDA_NSIGHT_COMPUTE_VERSION="12.4.1-1"
+NV_CUDA_NSIGHT_COMPUTE_DEV_PACKAGE="cuda-nsight-compute-12-4=${NV_CUDA_NSIGHT_COMPUTE_VERSION}"
+NV_NVPROF_VERSION="12.4.127-1"
+NV_NVPROF_DEV_PACKAGE="cuda-nvprof-12-4=${NV_NVPROF_VERSION}"
+NV_LIBNCCL_DEV_PACKAGE_NAME="libnccl-dev"
+NV_LIBNCCL_DEV_PACKAGE_VERSION="2.21.5-1"
+NCCL_VERSION="2.21.5-1"
+NV_LIBNCCL_DEV_PACKAGE="${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda12.4"
+
+# Update package lists
+sudo apt-get update
+
+# Install CUDA development packages
+sudo apt-get install -y --no-install-recommends \
+    cuda-cudart-dev-12-4=${NV_CUDA_CUDART_DEV_VERSION} \
+    cuda-command-line-tools-12-4=${NV_CUDA_LIB_VERSION} \
+    cuda-minimal-build-12-4=${NV_CUDA_LIB_VERSION} \
+    cuda-libraries-dev-12-4=${NV_CUDA_LIB_VERSION} \
+    cuda-nvml-dev-12-4=${NV_NVML_DEV_VERSION} \
+    ${NV_NVPROF_DEV_PACKAGE} \
+    ${NV_LIBNPP_DEV_PACKAGE} \
+    libcusparse-dev-12-4=${NV_LIBCUSPARSE_DEV_VERSION} \
+    ${NV_LIBCUBLAS_DEV_PACKAGE} \
+    ${NV_LIBNCCL_DEV_PACKAGE} \
+    ${NV_CUDA_NSIGHT_COMPUTE_DEV_PACKAGE}
+
+# Clean up
+sudo rm -rf /var/lib/apt/lists/*
+
+# Prevent auto-upgrade of specific packages
+sudo apt-mark hold ${NV_LIBCUBLAS_DEV_PACKAGE_NAME} ${NV_LIBNCCL_DEV_PACKAGE_NAME}
+
+# Set environment variable
+export LIBRARY_PATH=/usr/local/cuda/lib64/stubs
+
+echo "CUDA development environment setup is complete."