---
# Part of role: nvidia
# Called by: ansible/playbooks/setup-gpu-node.yaml
# Description: Installs NVIDIA drivers, CUDA toolkit, and nvidia-container-toolkit.
#              Configures Docker and K3s containerd runtimes for GPU access.

- name: Add NVIDIA CUDA apt keyring
  ansible.builtin.shell:
    cmd: >
      curl -fsSL
      https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
      -o /tmp/cuda-keyring.deb && dpkg -i /tmp/cuda-keyring.deb
    creates: /usr/share/keyrings/cuda-archive-keyring.gpg

- name: Add NVIDIA container toolkit keyring
  ansible.builtin.shell:
    cmd: >
      curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey |
      gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
    creates: /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

- name: Add NVIDIA container toolkit repo
  ansible.builtin.shell:
    cmd: >
      curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list |
      sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' |
      tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
    creates: /etc/apt/sources.list.d/nvidia-container-toolkit.list

- name: Update apt cache
  ansible.builtin.apt:
    update_cache: true

- name: Install NVIDIA driver
  ansible.builtin.apt:
    name: "nvidia-driver-{{ nvidia_driver_version }}"
    state: present
  register: nvidia_driver_install

- name: Install CUDA toolkit
  ansible.builtin.apt:
    name: "cuda-toolkit-{{ cuda_version }}"
    state: present

- name: Install nvidia-container-toolkit
  ansible.builtin.apt:
    name: nvidia-container-toolkit
    state: present

- name: Add CUDA to system PATH
  ansible.builtin.copy:
    dest: /etc/profile.d/cuda.sh
    content: |
      export PATH=/usr/local/cuda/bin:$PATH
      export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
    mode: "0644"

- name: Check if Docker already has the NVIDIA runtime configured
  ansible.builtin.command:
    cmd: grep -q '"nvidia"' /etc/docker/daemon.json
  register: nvidia_docker_runtime_check
  changed_when: false
  failed_when: false

- name: Configure Docker runtime for NVIDIA
  ansible.builtin.command:
    cmd: nvidia-ctk runtime configure --runtime=docker
  when: nvidia_docker_runtime_check.rc != 0
  changed_when: true
  notify: Restart docker

- name: Reboot if driver was just installed
  ansible.builtin.reboot:
    reboot_timeout: 300
  when: nvidia_driver_install.changed

- name: Verify NVIDIA driver loaded
  ansible.builtin.command: nvidia-smi
  changed_when: false
  failed_when: false
  register: nvidia_smi_check

- name: Assert nvidia-smi succeeded
  ansible.builtin.assert:
    that: nvidia_smi_check.rc == 0
    fail_msg: "nvidia-smi failed — driver may not have loaded correctly"