cluster-setup.yaml

---
###
##  Initial Provisioning Host Setup
###
- name: Setup Provisioning host
  hosts: provision
  gather_facts: yes

  tasks:
  - name : Update pip on Provisioning Node
    pip:
      name: pip
      state: latest
    become: yes

  - name: Remove pip urllib3 library
    pip:
      name: urllib3
      state: absent
    become: yes
    notify: Backup urllib3 file

  - name: Install CentOS Prereqs
    yum:
      name:
        - libselinux-python
        - python-urllib3
    become: yes
    when: ansible_distribution == 'CentOS'

  - name: Generate SSH Key
    command:
      cmd: ssh-keygen -b 2048 -t rsa -f "/home/{{ ansible_user }}/.ssh/id_rsa" -q -N ""
      creates: "/home/{{ ansible_user }}/.ssh/id_rsa"

  handlers:
  - name: Backup urllib3 file
    command:
      cmd: mv /usr/lib/python2.7/site-packages/urllib3/packages/ssl_match_hostname /tmp/ssl_match_hostname.bakup
    become: yes
    ignore_errors: yes


###
##  Configure Docker on Provioning Host
###
- name: Setup Docker on Provisioning Host
  hosts: provision
  gather_facts: yes

  tasks:
  - name: Add Docker-ce Repo
    yum_repository:
      name: docker-ce
      baseurl: https://download.docker.com/linux/centos/7/$basearch/stable
      description: docker-ce YUM REPO
      gpgkey: https://download.docker.com/linux/centos/gpg
      gpgcheck: yes
    become: yes
    when: ansible_distribution == 'CentOS'

  - name: Install CentOS Docker
    yum:
      name:
        - docker-ce
        - docker-ce-cli
        - docker-ce-selinux
    become: yes
    when: ansible_distribution == 'CentOS'

  - name: Create Apt Cache Volume
    file:
      path: "{{ apt_cache_path }}"
      owner: root
      group: root
      mode: '0755'
      state: directory
    become: yes

  - name: Update Volume Security context
    sefcontext:
      target: "{{ apt_cache_path }}(/.*)?"
      setype: svirt_sandbox_file_t
      state: present
    become: yes

  - name: Start Docker
    service:
      name: docker
      state: started
      enabled: yes
    become: yes

  - name: Install pip docker
    pip:
      name: docker
    become: yes

###
##  Setup All Cluster Nodes
###
- name: Prep All nodes
  hosts: all
  gather_facts: yes
  vars_prompt:
    - name: ansible_password
      prompt: Please enter the Ansible User SSH Password
      private: yes

  tasks:
  - name: Install CentOS Prereqs
    yum:
      name:
        - libselinux-python
    become: yes
    when: ansible_distribution == 'CentOS'

  - name: Copy SSH Keys to all hosts
    authorized_key:
      user: "{{ ansible_user }}"
      key: "{{ lookup('file', lookup('env','HOME') + '/.ssh/id_rsa.pub') }}"
      state: present
    when: "'provision' not in group_names"

  - name: Copy hosts file from Provisioning Node
    copy:
      src: /etc/hosts
      dest: /etc/hosts
    become: yes
    when: "'provision' not in group_names"

  - name: Set hostname
    hostname:
      name: "{{ inventory_hostname }}"
    become: yes

###
##  Setup Base Compute/GPU Nodes (Ubuntu ONLY)
###
- name: Setup Compute/GPU Nodes (Ubuntu)
  hosts: node
  gather_facts: yes

  tasks:

  # Configture APT Proxy
  - name: Configure Apt Proxy
    template:
      src: templates/01proxy.j2
      dest: /etc/apt/apt.conf.d/01proxy
      owner: root
      group: root
      mode: '0644'
    become: yes

  # Update All APT Packages
  - name: Run upgrade
    apt:
      update_cache: yes
      cache_valid_time: 3600 #One hour
    become: yes
    retries: 3
    delay: 10
    register: apt_result
    until: "apt_result is not failed"

  # Configure Time and TimeZone
  - name: Configure Time
    block:
      - name: Set Timezone
        timezone:
          name: "{{ cluster_timezone }}"
        become: yes

      - name: Unconfigure SystemD TimeSync/NTP
        shell: |
          timedatectl set-ntp false
        become: yes

      - name: Install Chronyd
        apt:
          pkg:
          - chrony
          update_cache: yes
          cache_valid_time: 86400 #One day
        become: yes
        notify:
          - Restart Chrony
          - Force Timesync
        retries: 3
        delay: 10
        register: apt_result
        until: "apt_result is not failed"

      - name: Configure chrony
        template:
          src: templates/chrony.conf.j2
          dest: /etc/chrony/chrony.conf
          owner: root
          mode: '0644'
        become: yes
        notify:
          - Restart Chrony
          - Force Timesync

      - name: Enable Chrony
        service:
          name: chrony
          enabled: true
          state: started
        become: yes

  # Install PAN FS Driver
  - name: Install PanFS Driver
    block:
      - name: Fail if Kernel driver is available
        fail:
          msg: "No PanFS Driver found for detected Kernel '{{ ansible_kernel }}'. Put the deb file in ./files and update the map in ./group_vars/all.yml"
        when: panfs_kernel_driver_map[ansible_kernel] is not defined

      - name: Check if Kernel driver is available
        debug:
          msg: "Detected Kernel '{{ ansible_kernel }}'. Using installer '{{ panfs_kernel_driver_map[ansible_kernel] }}'"
        when: panfs_kernel_driver_map[ansible_kernel] is defined

      # - name: Get version of PanFS Packing Installed
      #   shell: |
      #     dpkg -L panfs |grep -P '^/lib/modules/\d.+-generic$' |awk -F/ '{print $4}'
      #   register: dpkg_panfs_kernel
      #   become: yes

      # - name: Check PanFS Installed version
      #   debug:
      #     msg: "Found installed PanFs Driver for Kernel {{ dpkg_panfs_kernel.stdout }}"
      #   when: allow_panfs_driver_uninstall|bool is True and
      #         dpkg_panfs_kernel.stdout is not ''

      # - name: Remove PanFS Driver
      #   apt:
      #     package: panfs
      #     state: absent
      #   become: yes
      #   when: allow_panfs_driver_uninstall|bool is True and
      #         dpkg_panfs_kernel.stdout is not ansible_kernel and
      #         dpkg_panfs_kernel.stdout is not ''

      - name: Copy PanFS DEB
        copy:
          src: "files/{{ panfs_kernel_driver_map[ansible_kernel] }}"
          dest: "/tmp/{{ panfs_kernel_driver_map[ansible_kernel] }}"
      
      - name: Install PanFS Driver
        apt:
          deb: "/tmp/{{ panfs_kernel_driver_map[ansible_kernel] }}"
        become: yes

      - name: Probe for PanFS Module
        shell: |
          modprobe panfs
        become: yes

      - name: Check for PanFS Module
        shell: |
          lsmod |grep -ic panfs
        register: lsmod_out
        become: yes

      - name: Is PanFS Installed
        debug:
          msg: "Panfs successfully loaded"
        when: lsmod_out.stdout == '1'
        
    when: install_panfs_driver|bool == true

  # Reboot if required
  - name: Check if Reboot is required
    stat:
      path: /var/run/reboot-required
    register: reboot_required

  - name: Reboot if required
    reboot:
      msg: Reboot is required, rebooting now
      reboot_timeout: 1800
    when: reboot_required.stat.exists == true
          and allow_node_reboot == true
    become: yes

  handlers:
  - name: Restart Chrony
    service:
      name: chrony
      enabled: true
      state: restarted
    become: yes

  - name: Force Timesync
    shell: |
      chronyc -a makestep
    become: yes


###
##  Setup DGX Nodes (Ubuntu)
###
- name: Prep DGX Nodes
  hosts: dgx
  gather_facts: no

  tasks:
  - name: Download the Nvidia GPG Package
    get_url:
      url: https://international.download.nvidia.com/dgx/repos/bionic/pool/multiverse/d/dgx-repo-keys/dgx-repo-keys_2.0_amd64.deb
      dest: /tmp/dgx-repo-keys_2.0_amd64.deb
      checksum: sha256:c4d69304707f94bd255e9cf7b26c1f117a82b835b1363306ceebca73e39a3aff
    delegate_to: localhost
    retries: 3
    delay: 5

  - name: Copy Nvidia GPG Package to DGX nodes
    copy:
      src: /tmp/dgx-repo-keys_2.0_amd64.deb
      dest: /tmp/dgx-repo-keys_2.0_amd64.deb

  - name: Install Nvidia GPG Package on DGX nodes
    apt:
      deb: /tmp/dgx-repo-keys_2.0_amd64.deb
    become: yes

  - name: Download the Docker GPG Key
    get_url:
      url: https://download.docker.com/linux/ubuntu/gpg
      dest: /tmp/docker-ubuntu-key.gpg
      checksum: sha256:1500c1f56fa9e26b9b8f42452a553675796ade0807cdce11975eb98170b3a570
    delegate_to: localhost
    retries: 3
    delay: 5

  - name: Copy Docker GPG Key to DGX nodes
    copy:
      src: /tmp/docker-ubuntu-key.gpg
      dest: /tmp/docker-ubuntu-key.gpg

  - name: Install the Docker GPG Key on DGX Nodes
    apt_key:
      id: 9DC858229FC7DD38854AE2D88D81803C0EBFCD88
      file: /tmp/docker-ubuntu-key.gpg
      state: present
    become: yes

  - name: Configure Nvidia DGX OS Repos
    copy:
      src: files/dgx.list
      dest: /etc/apt/sources.list.d/dgx.list
      owner: root
      group: root
      mode: '0644'
    become: yes

  - name: Configure Nvidia DGX Cuda Repos
    copy:
      src: files/dgx-bionic-r418-cuda10-1-repo.list
      dest: /etc/apt/sources.list.d/dgx-bionic-r418-cuda10-1-repo.list
      owner: root
      group: root
      mode: '0644'
    become: yes

  - name: Configure Nvidia APT Preferences
    copy:
      src: files/nvidia
      dest: /etc/apt/preferences.d/nvidia
      owner: root
      group: root
      mode: '0644'
    become: yes

###
##  Setup Lambda Nodes (Ubuntu ONLY)
###
- name: Prep Lambda Nodes
  hosts: lambda
  gather_facts: no
  vars:
    restart_docker: true

  tasks:
  - name: Install Lambda Repos
    block:
      - name: Download Lambda DEB
        get_url:
          url: https://lambdalabs.com/static/misc/lambda-stack-repo.deb
          dest: /tmp/lambda-stack-repo.deb
        delegate_to: localhost
        retries: 3
        delay: 5

      - name: Copy to Lambda Nodes
        copy:
          src: /tmp/lambda-stack-repo.deb
          dest: /tmp/lambda-stack-repo.deb

      - name: Install DEB
        apt:
          deb: /tmp/lambda-stack-repo.deb
        become: yes
        register: lambda_deb
    when: install_lambda_repos|bool == true

  - name: Install Lambda Stack and Reboot
    block:
      - name: Accept CUDA EULA
        shell: |
          echo "cudnn cudnn/license_preseed select ACCEPT" |debconf-set-selections
        become: yes

      - name: Install Lambda Stack
        apt:
          pkg:
          - lambda-stack-cuda
          update_cache: yes
          install_recommends: false
        become: yes
        register: lambda_stack
        retries: 3
        delay: 10
        register: lambda_stack
        until: "lambda_stack is not failed"
    
      - name: Reboot Node
        reboot:
          msg: Reboot requested by Ansible Nofity
          reboot_timeout: 1800
        when: (lambda_deb.changed|bool == true or lambda_stack.changed|bool == true)
              and allow_node_reboot
        become: yes
    when: install_lambda_stack == true

  - name: Install Docker and Nvidia Toolkit
    apt:
      pkg:
      - docker.io
      - nvidia-container-toolkit
      update_cache: yes
      cache_valid_time: 0
    become: yes
    notify: Restart Docker
    retries: 3
    delay: 10
    register: apt_result
    until: "apt_result is not failed"

  handlers:
  - name: Restart Docker
    service:
      name: docker
      enabled: true
      state: restarted
    become: yes

###
##  Setup Docker and Proxy on Compute/GPU Node
###
- name: Setup Docker on Compute/GPU Nodes (Ubuntu)
  hosts: node
  gather_facts: no
  vars:
    update_certs: True
    deamon_reload: True

  tasks:
  # Review /etc/systemd/system/docker files
  - name: Create Docker System.d Directory
    file:
      path: /etc/systemd/system/docker.service.d
      owner: root
      group: root
      mode: '0755'
      state: directory
    become: yes

  - name: Configure Docker Proxy
    template:
      src: templates/http-proxy.conf.j2
      dest: /etc/systemd/system/docker.service.d/http-proxy.conf
      owner: root
      mode: '0644'
    become: yes
    notify: Restart Docker

  - name: Install docker-registry-proxy CA cert
    get_url:
      url: "http://{{ docker_cache_host }}:{{ docker_cache_port }}/ca.crt"
      dest: /usr/share/ca-certificates/docker_registry_proxy.crt
      owner: root
      group: root
      mode: '0644'
    become: yes
    notify:
      - Update CA Certs
      - Reload Systemd
      - Restart Docker

  - name: Add the Cert to the root certs
    lineinfile:
      path: /etc/ca-certificates.conf
      line: docker_registry_proxy.crt
    become: yes
    notify:
      - Update CA Certs
      - Reload Systemd
      - Restart Docker

  # Configure Docker to use Nvidia Runtime
  - name: Configure Nvidia Container Runtime (DGX)
    copy:
      src: files/daemon.json
      dest: /etc/docker/daemon.json
    become: yes
    notify: Restart Docker
    when: "'lambda' not in group_names"

  - name: Configure Nvidia Container Runtime (Lambda)
    copy:
      src: files/lambda.daemon.json
      dest: /etc/docker/daemon.json
    become: yes
    notify: Restart Docker
    when: "'lambda' in group_names"

  handlers:
  - name: Update CA Certs
    command:
      cmd: /usr/sbin/update-ca-certificates --fresh
    become: yes
    when: update_certs is defined

  - name: Reload Systemd
    systemd: daemon_reload=yes
    become: yes
    when: deamon_reload is defined

  - name: Restart Docker
    service:
      name: docker
      enabled: true
      state: restarted
    become: yes