From 5b391a491eafc9b08d1d6024d1ea58202853f8b1 Mon Sep 17 00:00:00 2001
From: Kenji Miyake <31987104+kenji-miyake@users.noreply.github.com>
Date: Mon, 23 May 2022 20:07:22 +0900
Subject: [PATCH] feat(setup, ansible): add an option for not installing
 cuda-drivers (#338)

Signed-off-by: Kenji Miyake <kenji.miyake@tier4.jp>
---
 ansible/roles/cuda/README.md          |  7 ++++---
 ansible/roles/cuda/defaults/main.yaml |  1 +
 ansible/roles/cuda/tasks/main.yaml    | 13 +++++++++++--
 setup-dev-env.sh                      | 10 +++++++++-
 4 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/ansible/roles/cuda/README.md b/ansible/roles/cuda/README.md
index d45b1d1d..f011866a 100644
--- a/ansible/roles/cuda/README.md
+++ b/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
 
diff --git a/ansible/roles/cuda/defaults/main.yaml b/ansible/roles/cuda/defaults/main.yaml
index e69de29b..3488c01e 100644
--- a/ansible/roles/cuda/defaults/main.yaml
+++ b/ansible/roles/cuda/defaults/main.yaml
@@ -0,0 +1 @@
+install_cuda_drivers: true
diff --git a/ansible/roles/cuda/tasks/main.yaml b/ansible/roles/cuda/tasks/main.yaml
index 7b9d2960..e282f264 100644
--- a/ansible/roles/cuda/tasks/main.yaml
+++ b/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
diff --git a/setup-dev-env.sh b/setup-dev-env.sh
index 63215586..a3510788 100755
--- a/setup-dev-env.sh
+++ b/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