feat(setup, ansible): add an option for not installing cuda-drivers (a…

…utowarefoundation#338)

Signed-off-by: Kenji Miyake <[email protected]>

ansible/roles/cuda/README.md

@@ -4,9 +4,10 @@ This role installs [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit) fol
 
 ## Inputs
 
-| Name         | Required | Description                  |
-| ------------ | -------- | ---------------------------- |
-| cuda_version | true     | The version of CUDA Toolkit. |
+| Name                 | Required | Description                      |
+| -------------------- | -------- | -------------------------------- |
+| cuda_version         | true     | The version of CUDA Toolkit.     |
+| install_cuda_drivers | false    | Whether to install cuda-drivers. |
 
 ## Manual Installation
 

ansible/roles/cuda/defaults/main.yaml

@@ -0,0 +1 @@
+install_cuda_drivers: true

ansible/roles/cuda/tasks/main.yaml

@@ -19,12 +19,21 @@
   register: dash_case_cuda_version
   changed_when: false
 
-- name: Install cuda-{{ dash_case_cuda_version.stdout }}
+- name: Install cuda-toolkit-{{ dash_case_cuda_version.stdout }}
   become: true
   ansible.builtin.apt:
-    name: cuda-{{ dash_case_cuda_version.stdout }}
+    name:
+      - cuda-toolkit-{{ dash_case_cuda_version.stdout }}
     update_cache: true
 
+- name: Install cuda-drivers
+  become: true
+  ansible.builtin.apt:
+    name:
+      - cuda-drivers
+    update_cache: true
+  when: install_cuda_drivers|bool
+
 - name: Add PATH to .bashrc
   ansible.builtin.lineinfile:
     dest: ~/.bashrc

setup-dev-env.sh

@@ -20,6 +20,9 @@ while [ "$1" != "" ]; do
     --no-nvidia)
         option_no_nvidia=true
         ;;
+    --no-cuda-drivers)
+        option_no_cuda_drivers=true
+        ;;
     *)
         args+=("$1")
         ;;
@@ -58,13 +61,18 @@ if [ "$option_verbose" = "true" ]; then
     ansible_args+=("-vvv")
 fi
 
-# Check NVIDIA Installation
+# Check installation of NVIDIA libraries
 if [ "$option_no_nvidia" = "true" ]; then
     ansible_args+=("--extra-vars" "install_nvidia=n")
 elif [ "$option_yes" = "true" ]; then
     ansible_args+=("--extra-vars" "install_nvidia=y")
 fi
 
+# Check installation of CUDA Drivers
+if [ "$option_no_cuda_drivers" = "true" ]; then
+    ansible_args+=("--extra-vars" "install_cuda_drivers=false")
+fi
+
 # Load env
 source "$SCRIPT_DIR/amd64.env"
 if [ "$(uname -m)" = "aarch64" ]; then