---

- name: gather GPU facts
  gpu_facts:

# For some reason one machine was running lightdm which was blocking driver install,
# this can probably be removed in the future
- name: stop lightdm if it is running
  service:
    name: lightdm
    state: stopped
    enabled: false
  failed_when: false
  ignore_errors: true

- name: get driver version
  set_fact:
    installed_nvidia_version: "{{hostvars[ansible_host].get('nvidia_smi_log', {}).get('driver_version') | string | trim}}"

- name: check if upgrade is needed
  set_fact:
    needs_nvidia_upgrade: "{{ force_nvidia_install or installed_nvidia_version|string|trim != nvidia_driver_version|string|trim }}"
    upgrade_nvidia: false  # default value

- name: skip install nvidia driver
  debug: msg="{{installed_nvidia_version}} driver is already installed."
  when: not needs_nvidia_upgrade

- name: check for sinfo
  stat:
    path: /usr/local/bin/sinfo
  register: sinfo


- name: Set node_name fact
  set_fact:
    node_name: "{{ inventory_hostname }}"

- name: Update node_name fact if EC2
  set_fact:
    node_name: "{{ ec2_private_dns_name.split('.')[0] }}"
  when: ec2_private_dns_name is defined

- name: check the current state of the host
  shell: "/usr/local/bin/sinfo -h -o '%T' -n {{node_name}} || echo"
  register: slurm_status
  when: sinfo.stat.exists

- debug:
    msg: "inventory_hostname: {{ inventory_hostname }} node_name: {{ node_name }}"

- debug:
    msg: "installed_nvidia_version: {{ installed_nvidia_version }}. needs_nvidia_upgrade: {{ needs_nvidia_upgrade }}"

- debug:
    msg: "Slurm State (empty means node not in Slurm cluster) for node_name {{node_name}} : {{slurm_status.stdout}}"

- name: prompt user to upgrade nvidia drivers
  pause:
    prompt: "Do you want to upgrade nvidia drivers? \nThis will wait for the node to drain. \nSee the previous ansible task for current state. (yes/no)"
  register: upgrade_nvidia_prompt
  when: needs_nvidia_upgrade and not nvidia

- name: save nvidia driver prompt
  set_fact:
    upgrade_nvidia: "{{ nvidia or upgrade_nvidia_prompt.user_input | bool }}"
  when: needs_nvidia_upgrade

- block:
  - name: will upgrade/install nvidia driver
    debug:
      msg: "Nvidia driver needs upgrading. Current '{{installed_nvidia_version}}' Want '{{nvidia_driver_version}}'"

  - name: install prereqs
    apt:
      name:
        - gcc
        - make
      state: present

  - name: install linux image extras
    apt:
      name: "linux-image-{{ansible_kernel}}"
      install_recommends: yes

  #http://us.download.nvidia.com/XFree86/Linux-x86_64/367.35/NVIDIA-Linux-x86_64-367.35.run
  - name: download nvidia {{nvidia_driver_version}} driver runfile
    get_url:
      url: "http://us.download.nvidia.com/XFree86/Linux-x86_64/{{nvidia_driver_version}}/NVIDIA-Linux-x86_64-{{nvidia_driver_version}}.run"
      dest: "/tmp/NVIDIA-Linux-x86_64-{{nvidia_driver_version}}.run"
      mode: 0700


  - name: "Drain the node"
    shell: "/usr/local/bin/scontrol update NodeName={{node_name}} state=DRAIN reason=\"{{drain_reason}}\""
    become: yes
    when: sinfo.stat.exists and 'drain' not in slurm_status.stdout and slurm_status.stdout != ""

  - name: "Wait until the node is drained"
    shell: "/usr/local/bin/sinfo -hN -n {{node_name}} | grep -v drain || true"
    register: result
    until: result.stdout_lines | length == 0
    retries: 1440
    delay: 60 # wait max 1 day before failing
    when: sinfo.stat.exists and 'drain' not in slurm_status.stdout and slurm_status.stdout != ""

  # nvidia's installer can't upgrade the driver if slurm is currently using it
  # don't fail if slurm isn't on this machine
  - name: stop slurm to upgrade gpu
    service:
      name: slurm
      state: stopped
    register: slurm_stopped
    failed_when: false

  - name: stop collectors to upgrade gpu
    service:
      name: zoox-pipedream-collectors
      state: stopped
    register: collectors_stopped
    failed_when: false

  - name: stop xserver to upgrade gpu
    service:
      name: zoox-xserver-idler
      state: stopped
    register: xserver_stopped
    failed_when: false

  - name: stop traffic control to upgrade gpu
    docker_service:
      project_src: /etc/docker-compose/traffic-control
      state: absent
    register: traffic_control_stopped
    failed_when: false

  - name: ensure no other containers are running other than collectors
    shell: "docker ps | tail -n +2"
    register: docker_ps
    failed_when: docker_ps.stdout_lines | length > 0

  - name: stop docker to upgrade gpu
    service:
      name: docker
      state: stopped
    register: docker_stopped
    failed_when: false

  - name: stop cron to upgrade gpu
    service:
      name: cron
      state: stopped
    register: cron_stopped
    failed_when: false

  - name: stop gpu watchdog to upgrade gpu
    service:
      name: gpu_watchdog
      state: stopped
    register: gpu_watchdog_stopped
    failed_when: false

  - name: unload necessary kernel modules
    modprobe:
      name: "{{ item }}"
      state: absent
    with_items:
      - nvidia_uvm
      - nvidia_drm
      - nvidia_modeset
      - nvidia
      - nouveau

  # accept license, answer questions with default, no curses ui, don't overwrite opengl files
  - name: install nvidia driver from runfile
    shell: "/tmp/NVIDIA-Linux-x86_64-{{nvidia_driver_version}}.run -a -q --ui=none --install-libglvnd"

  # Reload kernel modules
  - name: load nvidia drm/uvm kernel modules
    modprobe:
      name: nvidia_drm
      state: present
    with_items:
      - nvidia_drm
      - nvidia_uvm

  # start up slurm and gpu_watchdog if we stopped it previously
  - name: start up slurm again
    service:
      name: slurm
      state: started
    when: slurm_stopped.changed

  - name: start up docker again
    service:
      name: docker
      state: started
    when: docker_stopped.changed

  - name: start up cron again
    service:
      name: cron
      state: started
    when: cron_stopped.changed

  - name: start up gpu watchdog again
    service:
      name: gpu_watchdog
      state: started
    when: gpu_watchdog_stopped.changed

  - name: start xserver again
    service:
      name: zoox-xserver-idler
      state: started
    when: xserver_stopped.changed

  - name: start collectors again
    service:
      name: zoox-pipedream-collectors
      state: started
    when: collectors_stopped.changed

  - name: start traffic control again
    docker_service:
      project_src: /etc/docker-compose/traffic-control
      state: present
      stopped: no
    when: traffic_control_stopped.changed

  - name: run nvidia-modprobe and nvidia-smi to make sure all gpus devices are setup
    command: "{{ item }}"
    changed_when: false
    with_items:
      - "{{ nvidia_modprobe_command }}"
      - "{{ nvidia_smi_command }}"

  - name: verify driver version
    shell: "{{ nvidia_driver_version_command }}"
    register: updated_nvidia_version
    failed_when: "updated_nvidia_version.stdout != \"{{ nvidia_driver_version}}\""

  - name: verify gpus
    shell: "{{ nvidia_number_gpus_command }}"
    register: num_gpus
    failed_when: "{{ num_gpus.stdout != gpus|string }}"

  - name: "Undrain the node"
    shell: "/usr/local/bin/scontrol update NodeName={{node_name}} state=UNDRAIN reason=\"{{drain_reason}}\""
    become: yes
    when: sinfo.stat.exists and 'drain' not in slurm_status.stdout and slurm_status.stdout != ""

  when: upgrade_nvidia

