- name: install pytorch build requirements (CentOS)
  become: yes
  dnf:
    name:
      - cmake
      - gcc-c++
      - gcc-gfortran
      - lapack-devel
      - libjpeg-devel
      - libpng-devel
      - make
      - openblas-devel
      - patch
      - python3
      - python3-devel
      - python3-nose
      - python3-pillow
      - python3-setuptools
      - python3-virtualenv
      - python3-wheel
      - python3-yaml
    state: present
    enablerepo: powertools
  when: ansible_os_family == 'RedHat'

- name: install pytorch build requirements (Debian)
  become: yes
  apt:
    name:
      - cmake
      - g++
      - gfortran
      - liblapack-dev
      - make
      - libblas-dev
      - libjpeg-dev
      - libpng-dev
      - patch
      - python3
      - python3-dev
      - python3-nose
      - python3-pil
      - python3-setuptools
      - python3-venv
      - python3-wheel
      - python3-yaml
    state: present
  when: ansible_os_family == 'Debian'

- name: create virtualenv
  pip:
    name:
      'pip>=20'
    virtualenv: "{{ virtualenv_dir }}"
    virtualenv_command: "/usr/bin/python3 -m venv"
    virtualenv_site_packages: yes

- name: fetch pytorch source
  git:
    repo: "https://github.com/pytorch/pytorch.git" # noqa 401
    dest: "{{ build_dir.path }}/pytorch/"
    depth: 1
    recursive: yes
    force: yes

- name: install Python packages required by pytorch
  pip:
    requirements: requirements.txt
    chdir: "{{ build_dir.path }}/pytorch/"
    virtualenv: "{{ virtualenv_dir }}"
    virtualenv_command: "/usr/bin/python3 -m venv"
    virtualenv_site_packages: yes
    extra_args: "--extra-index-url {{ pip_extra_index_url }}"
  environment:
    NPY_NUM_BUILD_JOBS: "{{ ansible_processor_vcpus }}"

- name: build pytorch
  shell:
    cmd: |
      set -xe
      source "{{ virtualenv_dir }}/bin/activate"
      USE_CUDA=0 BUILD_CAFFE2_OPS=0 USE_DISTRIBUTED=0 USE_QNNPACK=0 USE_XNNPACK=0 \
        pip wheel -w "{{ wheels_dir }}" \
          --extra-index-url {{ pip_extra_index_url }} \
        .
      find . -name "*.whl" -ls
    chdir: "{{ build_dir.path }}/pytorch/"
    creates: "{{ build_dir.path }}/pytorch/build/"
    executable: /bin/bash

- name: get name of built wheel file
  find:
    path: "{{ wheels_dir }}"
    patterns: "torch-*.whl"
    file_type: file
  register: wheel_file

- name: install pytorch Python package
  pip:
    name: "{{ wheel_file.files[0].path }}"
    virtualenv: "{{ virtualenv_dir }}"
