nvidia.yml

---

- hosts: homelab

  vars:
    application: nvidia

    # https://www.nvidia.com/en-us/drivers/unix/
    # check compatibility table https://github.com/keylase/nvidia-patch
    nvidia_driver_version: "550.120"

    docker_network: "{{ networks.pub }}"

  tasks:
    - name: Create config folder
      ansible.builtin.file:
        path: "{{ config_directory }}"
        state: directory
        owner: "{{ common_user }}"
        group: "{{ common_group }}"
        mode: "0771"

    - name: Copy Docker driver check and restart script
      ansible.builtin.copy:
        content: |
          #!/usr/bin/env bash

          FUNCTION=$1

          if [ "$FUNCTION" == "STOP" ]; then
            for i in $(docker ps -q); do
              if docker inspect "$i" | jq -re '.[].HostConfig.DeviceRequests?' > /dev/null; then
                docker stop "$i" > /dev/null;
              fi
            done
          fi

          if [ "$FUNCTION" == "START" ]; then
            for i in $(docker ps -a -q); do
              if docker inspect "$i" | jq -re '.[].HostConfig.DeviceRequests?' > /dev/null; then
                docker start "$i" > /dev/null;
              fi
            done
          fi

          if [ "$FUNCTION" == "MONITOR" ]; then
            for i in $(docker ps -q); do
              if docker inspect "$i" | jq -e '.[].State.Status == "running" and .[].HostConfig.DeviceRequests?' > /dev/null; then
                if ! docker exec "$i" nvidia-smi > /dev/null; then
                  docker restart "$i" > /dev/null;
                fi
              fi
            done
          fi
        dest: "{{ config_directory }}/docker-container-check-restart.sh"
        mode: '0755'

    - name: Update apt cache
      ansible.builtin.apt:
        update_cache: true
        cache_valid_time: 600
      changed_when: false

    - name: Disable nouveau
      ansible.builtin.copy:
        content: |
          blacklist nouveau
          options nouveau modeset=0
        dest: /etc/modprobe.d/blacklist-nouveau.conf
        mode: '0644'
      register: __disable_nouveau

    - name: Regenerate the kernel initramfs
      ansible.builtin.command: update-initramfs -u
      changed_when: true
      when: __disable_nouveau.changed

    - name: Reboot to disable nouveau
      ansible.builtin.reboot:
        reboot_timeout: 360
      when: __disable_nouveau.changed

    - name: Install prerequisites
      ansible.builtin.apt:
        name:
          - pkg-config
          - xorg-dev
          - libglvnd-dev
          - dkms
          - libvulkan1
          - nvtop
        state: present
      register: __install_nvidia_prerequisites
      until: __install_nvidia_prerequisites is success
      retries: 5
      delay: 5

    - name: Create xorg modules folder
      ansible.builtin.file:
        path: /usr/lib/xorg/modules
        state: directory
        owner: "{{ common_root_id }}"
        group: "{{ common_root_group }}"
        mode: "0755"

    - name: Get current driver version
      ansible.builtin.command: nvidia-smi --query-gpu=driver_version --format=csv,noheader --id=0
      changed_when: false
      failed_when: false
      register: _nvidia_installed_version

    # https://www.nvidia.com/en-us/drivers/unix/
    - name: Download NVIDIA driver
      ansible.builtin.get_url:
        url: "https://international.download.nvidia.com/XFree86/Linux-x86_64/{{ nvidia_driver_version }}/NVIDIA-Linux-x86_64-{{ nvidia_driver_version }}.run"
        dest: "{{ config_directory }}/{{ nvidia_driver_version }}.run"
        mode: '0755'
      register: __download_nvidia_drivers
      when: _nvidia_installed_version.stdout != nvidia_driver_version

    - name: Stop containers using NVIDIA GPU
      ansible.builtin.command: "{{ config_directory }}/docker-container-check-restart.sh STOP"
      when: _nvidia_installed_version.stdout != nvidia_driver_version
      changed_when: true

    - name: Install NVIDIA driver
      ansible.builtin.command: "{{ config_directory }}/{{ nvidia_driver_version }}.run --silent"
      changed_when: true
      register: __install_nvidia_drivers
      when: _nvidia_installed_version.stdout != nvidia_driver_version

    - name: Create persistenced override directory
      ansible.builtin.file:
        path: /etc/systemd/system/nvidia-persistenced.service.d/
        state: directory
        recurse: true

    # https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker
    - name: Add NVIDIA's Docker GPG key
      ansible.builtin.get_url:
        url: https://nvidia.github.io/nvidia-docker/gpgkey
        dest: "/etc/apt/trusted.gpg.d/{{ application }}"
        owner: "{{ common_root_id }}"
        group: "{{ common_root_group }}"
        mode: "0644"
      register: result
      until: result is success
      retries: 5
      delay: 5

    # modified from https://nvidia.github.io/nvidia-docker/ubuntu22.04/nvidia-docker.list
    - name: Add NVIDIA's Docker Repository to APT
      ansible.builtin.copy:
        content: |
          deb [arch=amd64 signed-by=/etc/apt/trusted.gpg.d/{{ application }}] https://nvidia.github.io/libnvidia-container/stable/ubuntu18.04/$(ARCH) /
          #deb [arch=amd64 signed-by=/etc/apt/trusted.gpg.d/{{ application }}] https://nvidia.github.io/libnvidia-container/experimental/ubuntu18.04/$(ARCH) /
          deb [arch=amd64 signed-by=/etc/apt/trusted.gpg.d/{{ application }}] https://nvidia.github.io/nvidia-container-runtime/stable/ubuntu18.04/$(ARCH) /
          #deb [arch=amd64 signed-by=/etc/apt/trusted.gpg.d/{{ application }}] https://nvidia.github.io/nvidia-container-runtime/experimental/ubuntu18.04/$(ARCH) /
          deb [arch=amd64 signed-by=/etc/apt/trusted.gpg.d/{{ application }}] https://nvidia.github.io/nvidia-docker/ubuntu18.04/$(ARCH) /
        dest: /etc/apt/sources.list.d/{{ application }}.list
        owner: "{{ common_root_id }}"
        group: "{{ common_root_group }}"
        mode: "0644"
      register: __apt_repository

    - name: Update apt cache
      ansible.builtin.apt:
        update_cache: true
      when: __apt_repository.changed
      register: result
      until: result is success
      retries: 5
      delay: 5

    - name: Install NVIDIA Container Toolkit
      ansible.builtin.apt:
        name: nvidia-container-toolkit
        state: present
      register: __install_nvidia_toolkit
      until: __install_nvidia_toolkit is success
      retries: 5
      delay: 5

    - name: Restart Docker to complete installation
      ansible.builtin.service:
        name: docker
        state: restarted
      when:
        - not __install_nvidia_drivers.changed
        - __install_nvidia_toolkit.changed

    - name: Reboot to load NVIDIA drivers
      ansible.builtin.reboot:
        reboot_timeout: 360
      when: __install_nvidia_drivers.changed

    - name: Clone NVIDIA patch repo
      ansible.builtin.git:
        repo: "https://github.com/keylase/nvidia-patch"
        dest: "{{ config_directory }}/nvidia-patch"
        update: true

    - name: Run NVENV video encoding patch
      ansible.builtin.command: "bash {{ config_directory }}/nvidia-patch/patch.sh"
      register: _nvidia_patch_video_encoding
      changed_when: "'Patched!' in _nvidia_patch_video_encoding.stdout"

    - name: Run NvFBC unlock patch
      ansible.builtin.command: "bash {{ config_directory }}/nvidia-patch/patch-fbc.sh"
      register: _nvidia_patch_nvfbc_unlock
      changed_when: "'Patched!' in _nvidia_patch_nvfbc_unlock.stdout"

    - name: Start containers using NVIDIA GPU
      ansible.builtin.command: "{{ config_directory }}/docker-container-check-restart.sh START"
      when: _nvidia_installed_version.stdout != nvidia_driver_version
      changed_when: true

    - name: Create dcgm-exporter container
      ansible.builtin.include_role:
        name: docker_container
      vars:
        name: nvidia-dcgm-exporter
        image: nvcr.io/nvidia/k8s/dcgm-exporter:3.3.8-3.6.0-ubuntu22.04
        capabilities:
          - sys_admin
        device_requests:
          - driver: nvidia
            count: -1
            device_ids: []
            options: {}
            capabilities:
              - - gpu
                - video
                - compute
                - utility
        metrics:
          - port: 9400

    - name: Automatically restart containers when driver is not found
      ansible.builtin.cron:
        name: "nvidia gpu driver docker check"
        minute: "*/5"
        job: "{{ config_directory }}/docker-container-check-restart.sh MONITOR"

# The NVIDIA driver install can mess with the Docker daemon settings.
# Rerun docker again to be safe.
- ansible.builtin.import_playbook: docker.yml
  when: __install_nvidia_drivers.changed