---
# Linux GPU detection and validation tasks

- name: Check for NVIDIA GPU
  ansible.builtin.command: nvidia-smi
  register: nvidia_smi_output
  changed_when: false
  failed_when: false

- name: Display NVIDIA GPU information
  ansible.builtin.debug:
    msg: "NVIDIA GPU detected: {{ nvidia_smi_output.stdout_lines | join(', ') }}"
  when: nvidia_smi_output.rc == 0

- name: Validate NVIDIA VRAM
  ansible.builtin.set_fact:
    gpu_valid: true
    gpu_vram_gb: "{{ (nvidia_smi_output.stdout | regex_findall('(\\d+)MiB')) | first | default(0) | int / 1024 | int }}"
  when: nvidia_smi_output.rc == 0

- name: Check VRAM requirement (8GB minimum)
  ansible.builtin.debug:
    msg: "GPU VRAM: {{ gpu_vram_gb | default(0) }}GB - Requirement met: {{ gpu_vram_gb | default(0) >= 8 }}"
  when: gpu_valid is defined

- name: Warn about insufficient VRAM
  ansible.builtin.debug:
    msg: "WARNING: NVIDIA GPU has less than 8GB VRAM. Some labs may not function correctly."
  when: gpu_vram_gb is defined and gpu_vram_gb < 8

- name: Check for AMD GPU (ROCm)
  ansible.builtin.command: rocminfo
  register: rocm_output
  changed_when: false
  failed_when: false
  when: nvidia_smi_output.rc != 0

- name: Display AMD GPU information
  ansible.builtin.debug:
    msg: "AMD GPU detected"
  when: rocm_output.rc == 0

- name: Set GPU type fact
  ansible.builtin.set_fact:
    gpu_type: "{{ 'nvidia' if nvidia_smi_output.rc == 0 else ('amd' if rocm_output.rc == 0 else 'none') }}"

- name: Display GPU summary
  ansible.builtin.debug:
    msg: "GPU Type: {{ gpu_type | default('none') }}"
