nvidia.yml

---
- name: Check and update Nvidia drivers based on environment
  hosts: "{{ target_hosts }}"  # Dynamically populated by Semaphore UI
  become: true

  tasks:
    - name: Display the current environment mode
      debug:
        msg: >
          {% if var_environment == "production" %}
            "Running in PRODUCTION mode: Will install the driver if newer."
          {% else %}
            "Running in TEST mode: Will only print driver versions."
          {% endif %}

    - name: Update apt cache
      apt:
        update_cache: yes

    - name: Get the currently installed Nvidia driver version
      shell: dpkg -l | grep -oP 'nvidia-driver-\d+(-server)?' | head -1
      register: current_driver_version
      changed_when: false

    - name: Debug output of dpkg for Nvidia drivers
      debug:
        msg: "Raw output from dpkg command: {{ current_driver_version.stdout }}"

    - name: Get the recommended Nvidia server driver version
      shell: ubuntu-drivers devices | grep -oP '(nvidia-driver-\d+-server)' | tail -1
      register: latest_driver_version
      changed_when: false

    - name: Debug output of ubuntu-drivers for Nvidia server drivers
      debug:
        msg: "Raw output from ubuntu-drivers command: {{ latest_driver_version.stdout }}"

    - name: Extract current driver version number
      set_fact:
        current_version_number: "{{ current_driver_version.stdout | regex_search('nvidia-driver-(\\d+)', '\\1') | int | default(0) }}"

    - name: Extract latest driver version number
      set_fact:
        latest_version_number: "{{ latest_driver_version.stdout | regex_search('nvidia-driver-(\\d+)', '\\1') | int | default(0) }}"

    - name: Debug the current Nvidia driver version
      debug:
        msg: "Current Nvidia driver version: {{ current_driver_version.stdout }} (version: {{ current_version_number }})"

    - name: Debug the latest Nvidia server driver version
      debug:
        msg: "Latest Nvidia server driver version: {{ latest_driver_version.stdout }} (version: {{ latest_version_number }})"

    - name: Compare current and latest Nvidia driver versions
      debug:
        msg: >
          {% if current_version_number >= latest_version_number %}
            "The installed Nvidia driver ({{ current_driver_version.stdout }}) is up to date."
          {% else %}
            "A newer Nvidia driver version is available: {{ latest_driver_version.stdout }} (current: {{ current_driver_version.stdout }})"
          {% endif %}

    - name: Install the latest Nvidia server driver if newer and environment is production
      apt:
        name: "{{ latest_driver_version.stdout }}"
        state: present
      when:
        - var_environment == "production"
        - current_version_number < latest_version_number
      register: driver_installed

    - name: Reboot the server if a new driver was installed and var_reboot is true
      reboot:
        msg: "Rebooting after Nvidia driver installation"
        connect_timeout: 5
        reboot_timeout: 600
      when:
        - var_environment == "production"
        - driver_installed.changed
        - var_reboot | bool