3 files changed, 29 insertions, 0 deletions
diff --git a/roles/cuda/templates/cuda.sh.j2 b/roles/cuda/templates/cuda.sh.j2
new file mode 100755
index 0000000..78393a1
--- /dev/null
+++ b/roles/cuda/templates/cuda.sh.j2
@@ -0,0 +1,7 @@
+#!/bin/bash
+# {{ ansible_managed }}
+export PATH={{ cuda_bash_path }}:$PATH
+export LD_LIBRARY_PATH={{ cuda_bash_ld_lib_path }}:$LD_LIBRARY_PATH
+export CUDA_ROOT={{ cuda_bash_cuda_root }}
+export CUDA_INC_DIR={{ cuda_bash_cuda_inc_dir }}:$CUDA_INC_DIR
+export CPATH={{ cuda_bash_cpath }}:$CPATH
diff --git a/roles/cuda/templates/cuda_init.service.j2 b/roles/cuda/templates/cuda_init.service.j2
new file mode 100644
index 0000000..c61cbad
--- /dev/null
+++ b/roles/cuda/templates/cuda_init.service.j2
@@ -0,0 +1,13 @@
+[Unit]
+Description=Initialize nvidia/cuda devices
+Before=slurm.service
+DefaultDependencies=no
+Conflicts=shutdown.target
+
+[Service]
+ExecStart=/usr/local/bin/cuda_init.sh
+Type=oneshot
+RemainAfterExit=yes
+
+[Install]
+WantedBy=multi-user.target
diff --git a/roles/cuda/templates/cuda_init.sh.j2 b/roles/cuda/templates/cuda_init.sh.j2
new file mode 100644
index 0000000..3c85ea2
--- /dev/null
+++ b/roles/cuda/templates/cuda_init.sh.j2
@@ -0,0 +1,9 @@
+#!/bin/bash
+# During boot we initialize the GPUs (creates /dev/nvidia*)
+
+if [ ! -f "/usr/bin/nvidia-smi" ]; then
+        logger -s -t nvidia-smi "Script $0 could not find /usr/bin/nvidia-smi"
+else
+        /usr/bin/nvidia-smi --compute-mode={{ cuda_init_compute_mode }}
+        /usr/bin/nvidia-smi --persistence-mode={{ cuda_init_persistence_mode }}
+fi