setup-jetson.yml

---
- name: Initial setup of device
  hosts: jetsons
  become: yes
  become_user: root
  become_method: sudo


  tasks:
  - name: Update system package cache
    ansible.builtin.apt:
      update_cache: true
  - name: Install jetpack
    ansible.builtin.apt:
      name: nvidia-jetpack
  - name: Install useful developer tools
    apt:
      name:
        - python3-pip
        - python3-virtualenv
        - htop
        - curl
  - name: Install jtop
    pip:
      name: jetson-stats
      state: latest
  - name: Add user to docker group to be able to execute commands
    ansible.builtin.user:
      append: true
      groups: ["docker"]
      name: intrinsic
  - name: Set nvidia-container-runtime as default Docker runtime
    copy:
      src: daemon.json
      dest: /etc/docker/daemon.json
      owner: root
      group: root
      mode: '0644'
      backup: no

- name: User setup
  hosts: jetsons
  become: no
  vars:
    vimrc: '{{ansible_env.HOME}}/.vimrc'

  tasks:
  - name: Initial vim settings
    copy:
      src: vimrc
      dest: '{{vimrc}}'
      backup: no
      mode: '0640'

## PyTorch play. Based on documented steps from https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html
- name: ML Tools
  hosts: jetsons
  vars:
    pytorch:
      '5.1.1':
        url: 'https://developer.download.nvidia.com/compute/redist/jp/v511/pytorch/torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl'
      '5.1.2':
        url: 'https://developer.download.nvidia.com/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl'

  tasks:
  - name: Collect current JetPack version
    become: yes
    become_user: root
    become_method: sudo
    shell: >
      apt info nvidia-jetpack 2>/dev/null | grep 'Version:' | cut -f 2 -d' ' | cut -f1 -d'-'
    register: jetpack_version
  - name: log jetpack_version
    debug: msg='{{jetpack_version.stdout}}'
  - name: log url
    debug: msg='{{pytorch[jetpack_version.stdout].url}}'
  - name: Install PyTorch deps
    become: yes
    become_user: root
    become_method: sudo
    apt:
      name:
      - autoconf
      - bc
      - build-essential
      - g++-8
      - gcc-8
      - clang-8
      - lld-8
      - gettext-base
      - gfortran-8
      - iputils-ping
      - libbz2-dev
      - libc++-dev
      - libcgal-dev
      - libffi-dev
      - libfreetype6-dev
      - libhdf5-dev
      - libjpeg-dev
      - liblzma-dev
      - libncurses5-dev
      - libncursesw5-dev
      - libpng-dev
      - libreadline-dev
      - libssl-dev
      - libsqlite3-dev
      - libxml2-dev
      - libxslt-dev
      - locales
      - moreutils
      - openssl
      - python-openssl
      - rsync
      - scons
      - python3-pip
      - libopenblas-dev
  - name: Create working directory for install
    file:
      path: '{{ansible_env.HOME}}/ml-setup'
      state: directory
      mode: '0755'
  - name: Install pytorch from wheel
    environment:
      LD_LIBRARY_PATH: '/usr/lib/llvm-8/lib'
    pip:
      name:
        - '{{pytorch[jetpack_version.stdout].url}}'
      extra_args: "--no-cache"
  - name: Check PyTorch installation
    register: cuda_enabled
    shell: >
      python3 -c 'import torch;print(torch.cuda.is_available())'
  - debug: msg="Torch installed with CUDA support? {{cuda_enabled.stdout}}"