From 989df9ad674b939c08889efb966e9c7459622a08 Mon Sep 17 00:00:00 2001
From: Tiago Reis <tiago_reis@hey.com>
Date: Mon, 7 Jun 2021 19:39:52 +0100
Subject: [PATCH 1/2] Skip download and copy of runfile if already present in
 tmp path

---
 tasks/install_runfile.yml | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tasks/install_runfile.yml b/tasks/install_runfile.yml
index 4ede7fc..12bfe81 100644
--- a/tasks/install_runfile.yml
+++ b/tasks/install_runfile.yml
@@ -54,14 +54,19 @@
       - "Driver installed: {{ cuda_driver_installed }}"
       - "Toolkit installed: {{ cuda_toolkit_installed }}"
 
-- name: "Create temporary directory for runfile"
-  file:
-    path: /tmp/cuda_runfile
-    state: directory
+- name: Check runfile already exists
+  stat:
+    path: /tmp/cuda_runfile/{{ cuda_runfile_sh }}
+  register: cuda_runtime_file
 
 - name: "Obtain runfile"
   block:
 
+  - name: "Create temporary directory for runfile"
+    file:
+      path: /tmp/cuda_runfile
+      state: directory
+
   - name: "Copy pre-downloaded runfile"
     copy:
       src: "{{ cuda_runfile_sh }}"
@@ -74,8 +79,9 @@
       dest: "/tmp/cuda_runfile/{{ cuda_runfile_sh }}"
     when: cuda_runfile_download
 
-  when: (cuda_runfile_toolkit and not cuda_toolkit_installed) or
-        (cuda_runfile_driver and not cuda_driver_installed)
+  when: not cuda_runtime_file.stat.exists 
+    and ((cuda_runfile_toolkit and not cuda_toolkit_installed)
+    or (cuda_runfile_driver and not cuda_driver_installed))
 
 - name: "Run installer for toolkit"
   command: bash /tmp/cuda_runfile/{{ cuda_runfile_sh }} --silent --toolkit

From 7cebc506c778bb3ee2a8927a34fd26486543fab9 Mon Sep 17 00:00:00 2001
From: Tiago Reis <tiago_reis@hey.com>
Date: Sun, 20 Jun 2021 19:01:23 +0100
Subject: [PATCH 2/2] Add checksum validation

---
 defaults/main.yml         |  2 ++
 tasks/install_runfile.yml | 27 +++++++++++++++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/defaults/main.yml b/defaults/main.yml
index 9d6f6bb..a0914e3 100644
--- a/defaults/main.yml
+++ b/defaults/main.yml
@@ -9,6 +9,8 @@ cuda_packages:
  - cuda
 cuda_use_runfile: False
 cuda_runfile_url: "https://developer.download.nvidia.com/compute/cuda/11.2.2/local_installers/cuda_11.2.2_460.32.03_linux.run"
+cuda_runfile_checksum: True
+cuda_runfile_checksum_url: "https://developer.download.nvidia.com/compute/cuda/11.2.2/docs/sidebar/md5sum.txt"
 cuda_runfile_driver: True
 cuda_runfile_toolkit: True
 cuda_runfile_download: True
diff --git a/tasks/install_runfile.yml b/tasks/install_runfile.yml
index 12bfe81..310a06c 100644
--- a/tasks/install_runfile.yml
+++ b/tasks/install_runfile.yml
@@ -7,7 +7,7 @@
   when: ansible_pkg_mgr in ["yum", "dnf"]
 
 - name: "Ensure kernel headers are installed (apt)"
-  yum:
+  apt:
     name:
       - linux-headers-generic
       - build-essential
@@ -57,8 +57,30 @@
 - name: Check runfile already exists
   stat:
     path: /tmp/cuda_runfile/{{ cuda_runfile_sh }}
+    checksum_algorithm: md5
   register: cuda_runtime_file
 
+- name: Download checksum file
+  get_url:
+    url: "{{ cuda_runfile_checksum_url }}"
+    dest: /tmp/cuda_checksum.tsv
+  when: cuda_runfile_checksum
+
+- name: Get checksums from file
+  community.general.read_csv:
+      path: /tmp/cuda_checksum.tsv
+      delimiter: " "
+      fieldnames: checksum, file
+      key: file
+  register: checksums
+  when: cuda_runfile_checksum
+
+- name: Get checksums from dict
+  set_fact:
+    cuda_runfile_valid_checksum: "{{checksums[ cuda_runfile_sh ].checksum}}" == cuda_runtime_file.stat.checksum
+  when: cuda_runtime_file.stat.exists
+
+
 - name: "Obtain runfile"
   block:
 
@@ -79,10 +101,11 @@
       dest: "/tmp/cuda_runfile/{{ cuda_runfile_sh }}"
     when: cuda_runfile_download
 
-  when: not cuda_runtime_file.stat.exists 
+  when: not (cuda_runtime_file.stat.exists and cuda_runfile_valid_checksum)
     and ((cuda_runfile_toolkit and not cuda_toolkit_installed)
     or (cuda_runfile_driver and not cuda_driver_installed))
 
+
 - name: "Run installer for toolkit"
   command: bash /tmp/cuda_runfile/{{ cuda_runfile_sh }} --silent --toolkit
   register: cuda_toolkit_install_out