provision.yml

---
- name: "Configure GPU VM"
  hosts: 'all'
  become: yes
  become_method: sudo

  vars:
    # The following environment variables should be set in your environment before calling this playbook based on the CI system:
    #
    # Azure Pipelines:
    #  AZURE_DEVOPS_ORGANIZATION : Azure DevOps organization under which the pipeline is created
    #  AZURE_DEVOPS_GPU_POOL: Azure DevOps build agent pool name this agent will register into
    #  AZURE_DEVOPS_PAT_TOKEN: Azure DevOps Personal Access Token to grant API access for the organization. This should be a secret environment variable.
    # GitHub Actions:
    #  GITHUB_ACTIONS: set to true by GitHub Actions when running in that context. We use that to determine which CI system to set up for.
    #  GITHUB_TOKEN: GitHub token granting API access to the repository, set automatically but GitHub Actions when running in that context.
    #  GITHUB_REPOSITORY: the GitHub repo hosting this project, set automatically by GitHub Actions when running in that context.
    #  GITHUB_PAT: Unfortunately the GITHUB_TOKEN generated automatically by GitHub Actions doesn't have administration permission so we need our own

    # CI System: automatically set based on presence or absence of GITHUB_ACTIONS environment variable, override with --extra-vars if needed
    ci_system: "{{ 'github_actions' if lookup('env','GITHUB_ACTIONS') == 'true' else 'azure_pipelines' }}"
    # Cloud Provider: defaults to "azure", can also specify "aws" or "gcp". Typically set as an --extra-vars on the ansible-playbook command line
    cloud_provider: "azure"

    # Set ci_agent_version to "latest" to use latest version of the Azure Pipelines / GitHub Actions agent, otherwise lock it to a specific version.
    # Don't prepend "v" to the semantic version number if you want to lock the agent version, using for instance "2.165.2"
    ci_agent_version: "latest"
    # What to ask for from the GitHub versions API
    ci_agent_version_api: "{{ 'latest' if ci_agent_version == 'latest' else 'tags/v{{ ci_agent_version }}' }}"
    ci_agent_dir: "ci_agent"
    ci_agent_local_filename: "/tmp/ci_agent.tar.gz"
    ci_agent_repo:
      azure_pipelines: "Microsoft/azure-pipelines-agent"
      github_actions:  "actions/runner"
    ci_agent_download_prefix:
      azure_pipelines: "https://vstsagentpackage.azureedge.net/agent/"
      github_actions:  "https://github.com/actions/runner/releases/download/v"
    ci_agent_base_filename:
      azure_pipelines: "vsts-agent-linux-x64"
      github_actions:  "actions-runner-linux-x64"
    # For GitHub Actions this will get replaced by a specific registration token for the runner API
    ci_agent_registration_token:
      azure_pipelines: "{{ lookup('env','AZURE_DEVOPS_PAT_TOKEN') }}"
      github_actions:  "{{ lookup('env','GITHUB_PAT') }}"
      
    github_repository: "{{ lookup('env','GITHUB_REPOSITORY') }}"

    ci_agent_registration_url:
      azure_pipelines: "https://dev.azure.com/{{ lookup('env','AZURE_DEVOPS_ORGANIZATION') }}"
      github_actions: "https://github.com/{{ github_repository }}"
    ci_agent_registration_options:
      azure_pipelines: "--auth pat 
        --pool '{{ lookup('env','AZURE_DEVOPS_GPU_POOL') }}' 
        --agent {{ ansible_hostname }} 
        --acceptTeeEula"
      github_actions: "--work _work
        --name {{ ansible_hostname }}"

    # Pre-licensed GRID driver to install based on cloud_provider variable
    nvidia_grid_driver_source:
      # See https://docs.microsoft.com/en-us/azure/virtual-machines/linux/n-series-driver-setup
      # and https://github.com/Azure/azhpc-extensions/blob/master/NvidiaGPU/resources.json for direct
      # locations of pre-licensed GRID driver. The URL https://go.microsoft.com/fwlink/?linkid=874272
      # provides the "latest" driver available but we need more control.
      azure: "https://download.microsoft.com/download/1/a/5/1a537cae-5b52-4348-acd2-2f210fc412b0/NVIDIA-Linux-x86_64-430.46-grid.run"
      # See https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/install-nvidia-driver.html
      aws: "http://ec2-linux-nvidia-drivers.s3.amazonaws.com/grid-9.1/NVIDIA-Linux-x86_64-430.46-grid.run"
      # See https://cloud.google.com/compute/docs/gpus/add-gpus#installing_grid_drivers_for_virtual_workstations
      gcp: "https://storage.googleapis.com/nvidia-drivers-us-public/GRID/GRID9.1/NVIDIA-Linux-x86_64-430.46-grid.run"

 # roles:
    
  tasks:
  - name: Update and upgrade apt packages
    apt:
      update_cache: yes
      upgrade: yes

  - name: Install support for repository access via https
    apt:
      name: "{{ packages }}"
    vars:
      packages:
        - apt-transport-https
        - ca-certificates
        - curl
        - gnupg-agent
        - software-properties-common

  - name: Add an Apt signing key to a specific keyring file
    apt_key:
      url: https://download.docker.com/linux/ubuntu/gpg
      state: present

  - name: Add Stable Docker repository
    apt_repository:
      repo: deb [arch=amd64] https://download.docker.com/linux/ubuntu {{ ansible_distribution_release }} stable
      state: present

  - name: Install Docker CE
    apt:
      name: "{{ packages }}"
    vars:
      packages:
        - docker-ce
        - docker-ce-cli
        - containerd.io

  - name: Allow Docker access to non-root user
    user:
      name: '{{ ansible_user }}'
      groups: docker
      append: yes

  # Need GCC and Make to install NVIDIA driver, xorg/gdm3 for OpenGL apps
  - name: Install Xorg, GCC and make
    apt:
      name: "{{ packages }}"
    vars:
      packages:
        - gcc
        - make
        - mesa-utils
        - xorg
        - daemon

  - name: Download NVIDIA GRID driver for the current cloud provider
    get_url:
      url: "{{ nvidia_grid_driver_source[cloud_provider] }}"
      dest: "/tmp/NVIDIA_Driver_GRID.run"
      mode: '0755'

  - name: Install the NVIDIA driver in silent mode, will not upgrade a previous version (you need fancier check)
    command: "/tmp/NVIDIA_Driver_GRID.run -s"
    args:
      creates: /proc/driver/nvidia/version

  # BoardName doesn't seem to matter much
  - name: Xorg config magic to get NVIDIA hardware OpenGL
    copy:
      dest: /usr/share/X11/xorg.conf.d/20-nvidia.conf
      content: |
        Section "Device"
          Identifier "Nvidia Card"
          Driver "nvidia"
          VendorName "NVIDIA Corporation"
          BoardName "Tesla M60"
        EndSection

  # rc.local still supported as a systemd service in Ubuntu 18.04
  - name: Xorg startup script
    copy:
      dest: /etc/rc.local
      mode: '0755'
      owner: root
      group: root
      content: |
        #!/bin/bash
        /usr/bin/X11/X &
        exit 0

  # start Xorg
  - name: Start Xorg
    systemd:
      enabled: True
      state: started
      name: rc-local

  - name: Add an Apt signing key for NVIDIA Docker container repository
    apt_key:
      url: https://nvidia.github.io/nvidia-docker/gpgkey
      state: present

  - name: Add NVIDIA container repositories
    apt_repository:
      repo: "{{ item }}"
      state: present
    loop: "{{ lookup('url', 'https://nvidia.github.io/nvidia-docker/{{ ansible_distribution|lower }}{{ ansible_distribution_version }}/nvidia-docker.list', wantlist=True) }}"

  - name: Install nvidia-container-toolkit
    apt:
      name: "{{ packages }}"
    vars:
      packages:
          - nvidia-container-toolkit

  # Should test first if we actually installed nvidia-container-toolkit
  - name: Restart Docker after installing nvidia-container-toolkit
    systemd:
      state: restarted
      name: docker

  # Figure out version number corresponding to "latest" agent release
  # We're bypassing the more specific API to get the latest GitHub Actions self-hosted runner download
  # documented at https://developer.github.com/v3/actions/self_hosted_runners/ to use instead the generic GitHub releases API
  - name: Query latest version of CI agent
    uri:
      url: "https://api.github.com/repos/{{ ci_agent_repo[ci_system] }}/releases/{{ ci_agent_version_api }}"
    register: ci_agent_version_json
    when:
      ci_agent_version == "latest"

  # Update to the actual "latest" version
  - name: Update to the actual latest CI agent version
    set_fact:
      ci_agent_version: "{{ ci_agent_version_json.json.tag_name | regex_replace('^v(.+)$', '\\1') }}"
    when:
      ci_agent_version == "latest"

  - name: Download CI agent for Linux
    get_url:
      url: "{{ ci_agent_download_prefix[ci_system] }}{{ ci_agent_version }}/{{ ci_agent_base_filename[ci_system] }}-{{ ci_agent_version }}.tar.gz"
      dest: "{{ ci_agent_local_filename }}"
      mode: '0644'

  - name: Create install dir for CI agent
    file:
      path: "~/{{ ci_agent_dir }}"
      state: directory
    become: no
  
  - name: Unpack the CI agent archive into installation dir
    unarchive:
      src: "{{ ci_agent_local_filename }}"
      dest: "~/{{ ci_agent_dir }}"
      remote_src: yes
    become: no
    args:
      creates: "~/{{ ci_agent_dir }}/config.sh"

  # For GitHub Actions we need a specific registration token
  - name: Get GitHub Actions registration token
    uri:
      url:     "https://api.github.com/repos/{{ github_repository }}/actions/runners/registration-token"
      method: POST
      status_code: 201
      headers:
         Authorization: "token {{ ci_agent_registration_token[ci_system] }}"
    register: github_registration_token
    no_log: False
    when:
      ci_system == "github_actions"

  - name: Set CI agent registration token
    set_fact:
      ci_agent_token: "{{ github_registration_token.json.token if ci_system == 'github_actions' else ci_agent_registration_token[ci_system] }}"

  - name: Configure the CI agent
    command: "./config.sh
      --url {{ ci_agent_registration_url[ci_system] }}
      --token {{ ci_agent_token }}
      --unattended
      --replace
      {{ ci_agent_registration_options[ci_system] }}"
    become: no
    no_log: False
    args:
      chdir: "~/{{ ci_agent_dir}}"
      creates: "~/{{ ci_agent_dir}}/.agent"

  - name: Test for service already created
    shell: "./svc.sh status | grep -q -v -e '^$' -e  'not installed'"
    register: ci_agent_service_installed
    ignore_errors: true
    args:
      chdir: "~{{ ansible_user }}/{{ ci_agent_dir }}"

  - name: Install CI agent as a service
    command: "./svc.sh install"
    args:
      chdir: "~{{ ansible_user }}/{{ ci_agent_dir }}"
    when:
      ci_agent_service_installed.rc == 1

  # No harm done if already started
  - name: Start the CI agent service
    command: "./svc.sh start"
    args:
      chdir: "~{{ ansible_user }}/{{ ci_agent_dir }}"