---
# ------------------------------------------------------------------------------
# FILE: roles/ollama/tasks/main.yml
# DESCRIPTION: Deploys Ollama with AMD ROCm GPU acceleration on Ubuntu 24.04.
#              Handles:
#                - ROCm repository and driver stack installation
#                - Data disk preparation (LVM, ext4, persistent mount)
#                - Ollama installation and systemd service configuration
#                - HSA override for RX 5000-series (gfx1010) compatibility
#                - Initial model pull
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# ROCm prerequisites
# ------------------------------------------------------------------------------

- name: Install ROCm prerequisite packages
  apt:
    name:
      - wget
      - gnupg
      - ca-certificates
      - lvm2
    state: present
    update_cache: yes

- name: Add ROCm apt repository signing key
  apt_key:
    url: https://repo.radeon.com/rocm/rocm.gpg.key
    state: present

- name: Add ROCm apt repository
  apt_repository:
    repo: "deb [arch=amd64] https://repo.radeon.com/rocm/apt/{{ ollama_rocm_version }} {{ ansible_distribution_release }} main"
    state: present
    filename: rocm
  notify: update apt cache

- name: Flush handlers to update apt cache before ROCm install
  meta: flush_handlers

# ------------------------------------------------------------------------------
# ROCm driver stack
# hip-runtime-amd pulls in the full ROCm stack as a dependency.
# rocminfo is used to verify GPU visibility after install.
# ------------------------------------------------------------------------------

- name: Install ROCm packages
  apt:
    name: "{{ ollama_rocm_packages }}"
    state: present
  notify: update initramfs

- name: Create /opt/rocm symlink to versioned directory
  file:
    src: "/opt/rocm-{{ ollama_rocm_version }}.0"
    dest: /opt/rocm
    state: link
    force: false

- name: Add ROCm binaries to system PATH
  copy:
    dest: /etc/profile.d/rocm.sh
    content: |
      export PATH=$PATH:/opt/rocm/bin
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm/lib
    mode: '0644'

- name: Add wed user to render and video groups for GPU access
  user:
    name: "{{ ansible_user }}"
    groups:
      - render
      - video
    append: yes

- name: Add ollama user to render and video groups (created by Ollama installer)
  user:
    name: ollama
    groups:
      - render
      - video
    append: yes
  ignore_errors: yes
  # ollama user does not exist yet at this point — Ollama installer creates it.
  # This task re-runs after install via the post-install handler.

# ------------------------------------------------------------------------------
# Data disk — LVM setup on /dev/sdb
# Follows the same pattern as the monitoring role (cinderella-castle).
# Wipes any stale GPT/partition signatures before creating the PV to avoid
# the "GPT protective partition" error encountered during monitoring deploy.
# ------------------------------------------------------------------------------

- name: Wipe any existing partition signatures from data disk
  command: wipefs -a {{ ollama_data_disk }}
  args:
    creates: /dev/{{ ollama_data_vg }}
  when: ollama_setup_data_disk | bool

- name: Create LVM physical volume on data disk
  command: pvcreate {{ ollama_data_disk }}
  args:
    creates: /dev/{{ ollama_data_vg }}
  when: ollama_setup_data_disk | bool

- name: Create LVM volume group
  lvg:
    vg: "{{ ollama_data_vg }}"
    pvs: "{{ ollama_data_disk }}"
  when: ollama_setup_data_disk | bool

- name: Create LVM logical volume using 100% of VG
  lvol:
    vg: "{{ ollama_data_vg }}"
    lv: "{{ ollama_data_lv }}"
    size: 100%FREE
    shrink: false
  when: ollama_setup_data_disk | bool

- name: Format logical volume as ext4
  filesystem:
    fstype: ext4
    dev: "/dev/{{ ollama_data_vg }}/{{ ollama_data_lv }}"
    opts: "-L ollama-data"
  when: ollama_setup_data_disk | bool

- name: Create Ollama model directory mount point
  file:
    path: "{{ ollama_data_dir }}"
    state: directory
    mode: '0755'

- name: Mount Ollama data volume and add fstab entry
  mount:
    path: "{{ ollama_data_dir }}"
    src: "/dev/{{ ollama_data_vg }}/{{ ollama_data_lv }}"
    fstype: ext4
    opts: defaults,noatime
    state: mounted
  when: ollama_setup_data_disk | bool
  # noatime: eliminates inode access time writes on every model file read.
  # Meaningful win for large sequential reads like model loading.

# ------------------------------------------------------------------------------
# Ollama installation
# Uses the official install script which creates the ollama user and
# systemd service unit. We override the service with our own config.
# ------------------------------------------------------------------------------

- name: Download Ollama install script
  get_url:
    url: https://ollama.ai/install.sh
    dest: /tmp/ollama_install.sh
    mode: '0755'

- name: Run Ollama install script
  command: /tmp/ollama_install.sh
  environment: "{{ {} if ollama_version == 'latest' else {'OLLAMA_VERSION': ollama_version} }}"
  args:
    creates: /usr/local/bin/ollama

- name: Add ollama user to render and video groups (post-install)
  user:
    name: ollama
    groups:
      - render
      - video
    append: yes

- name: Set ownership of Ollama data directory to ollama user
  file:
    path: "{{ ollama_data_dir }}"
    owner: ollama
    group: ollama
    recurse: yes

# ------------------------------------------------------------------------------
# Systemd service override
# Injects HSA_OVERRIDE_GFX_VERSION for RX 5000-series compatibility and
# sets OLLAMA_MODELS to the data disk mount point.
# ------------------------------------------------------------------------------

- name: Create systemd override directory for ollama service
  file:
    path: /etc/systemd/system/ollama.service.d
    state: directory
    mode: '0755'

- name: Deploy ollama systemd service override
  template:
    src: ollama-override.conf.j2
    dest: /etc/systemd/system/ollama.service.d/override.conf
    mode: '0644'
  notify:
    - reload systemd
    - restart ollama

- name: Flush handlers to apply service config before model pull
  meta: flush_handlers

- name: Enable and start Ollama service
  systemd:
    name: ollama
    state: started
    enabled: yes
    daemon_reload: yes

- name: Wait for Ollama API to become available
  uri:
    url: "http://localhost:{{ ollama_port }}/api/tags"
    status_code: 200
  register: ollama_health
  retries: 12
  delay: 5
  until: ollama_health.status == 200

# ------------------------------------------------------------------------------
# Initial model pull
# ------------------------------------------------------------------------------

- name: Pull initial model
  command: ollama pull {{ ollama_default_model }}
  environment:
    OLLAMA_HOST: "http://localhost:{{ ollama_port }}"
  register: model_pull
  changed_when: "'pulled' in model_pull.stdout or model_pull.rc == 0"
  timeout: 600
  # Qwen3 8B Q4_K_M is ~5GB — allow 10 minutes on first pull.