main.yaml

- name: Read PyTorch pin
  ansible.builtin.command: cat {{ (src_root, 'pytorch/xla/.torch_pin') | path_join }}
  register: torch_pin
  # Pin may not exist
  ignore_errors: true

- name: Checkout PyTorch pin
  # ansible.builtin.git wants to fetch the entire history, so check out the pin manually
  ansible.builtin.shell:
    cmd: |
      set -xe
      PIN="{{ torch_pin.stdout }}"
      if [[ $PIN = \#* ]]; then
        PRNUM="${PIN//[!0-9]/}"
        git fetch --no-recurse-submodules origin "pull/$PRNUM/head"
      else
        git fetch --no-recurse-submodules origin {{ torch_pin.stdout }}
      fi
      git checkout --no-recurse-submodules FETCH_HEAD
      git submodule update --init --recursive
    chdir: "{{ (src_root, 'pytorch') | path_join }}"
  args:
    executable: /bin/bash
  when: torch_pin is succeeded

- name: Build PyTorch
  ansible.builtin.command:
    cmd: python setup.py bdist_wheel
    chdir: "{{ (src_root, 'pytorch') | path_join }}"
    creates: "{{ (src_root, 'pytorch/dist/*.whl') | path_join }}"
  # Set `USE_CUDA=0` as PyTorch cannot be used with GPU in eager and XLA mode.
  environment: "{{ env_vars | combine({'USE_CUDA': 0}) }}"

- name: Find PyTorch *.whl files in pytorch/dist
  ansible.builtin.find:
    path: "{{ (src_root, 'pytorch/dist') | path_join }}"
    pattern: "*.whl"
  register: pytorch_wheels

- name: Install PyTorch wheels
  ansible.builtin.pip:
    name: "{{ pytorch_wheels.files | map(attribute='path') }}"
    state: "forcereinstall"

- name: Build PyTorch/XLA
  ansible.builtin.command:
    cmd: python setup.py bdist_wheel
    chdir: "{{ (src_root, 'pytorch/xla') | path_join }}"
  environment: "{{ env_vars }}"

- name: Find XLA *.whl files in pytorch/xla/dist
  ansible.builtin.find:
    path: "{{ (src_root, 'pytorch/xla/dist') | path_join }}"
    pattern: "*.whl"
  register: xla_wheels

- name: Install XLA wheels
  ansible.builtin.pip:
    name: "{{ xla_wheels.files | map(attribute='path') }}"
    state: "forcereinstall"

- name: Create a temp directory for collecting wheels
  ansible.builtin.file:
    path: /tmp/staging-wheels
    state: directory
    mode: '0755'

- name: Create /dist directory for exported wheels
  ansible.builtin.file:
    path: /dist
    state: directory
    mode: '0755'

- name: Rename and append +YYYYMMDD suffix to nightly wheels
  ansible.builtin.shell: |
    pushd /tmp/staging-wheels
    cp {{ item.dir }}/*.whl .
    rename -v "s/^{{ item.prefix }}-(.*?)-cp/{{ item.prefix }}-nightly-cp/" *.whl
    mv /tmp/staging-wheels/* /dist/
    popd

    rename -v "s/^{{ item.prefix }}-(.*?)-cp/{{ item.prefix }}-nightly+$(date -u +%Y%m%d)-cp/" *.whl
  args:
    executable: /bin/bash
    chdir: "{{ item.dir }}"
  loop:
    - { dir: "{{ (src_root, 'pytorch/dist') | path_join }}", prefix: "torch" }
    - { dir: "{{ (src_root, 'pytorch/xla/dist') | path_join }}", prefix: "torch_xla" }
  when: nightly_release

- name: Copy wheels to /dist
  ansible.builtin.shell: "cp {{ item }}/*.whl /dist"
  args:
    executable: /bin/bash
  loop:
    - "{{ (src_root, 'pytorch/dist') | path_join }}"
    - "{{ (src_root, 'pytorch/xla/dist') | path_join }}"

- name: Delete temp directory
  ansible.builtin.file:
    path: /tmp/staging-wheels
    state: absent
    mode: '0755'

- name: Create temp directory for C++ tests
  ansible.builtin.file:
    path: /tmp/test/bin
    state: directory
    mode: '0755'
  when: build_cpp_tests

- name: Collect C++ test files
  ansible.builtin.shell: |
    cd pytorch/xla/build/temp*
    bazel query 'kind(".*_test", tests(//:cpp_tests))' --output=label | xargs -n 1 bazel cquery --output=files | xargs cp -t /tmp/test/bin
  args:
    executable: bash
    chdir: "{{ src_root }}"
  when: build_cpp_tests

- name: Read Torchvision pin
  ansible.builtin.command: cat {{ (src_root, 'pytorch') | path_join }}/.github/ci_commit_pins/vision.txt
  register: torchvision_pin

- name: Build Torchvision wheel
  ansible.builtin.command:
    cmd: pip wheel --no-build-isolation --no-deps "git+https://github.com/pytorch/vision.git@{{ torchvision_pin.stdout }}"
    chdir: /dist
  environment: "{{ env_vars }}"

- name: Find Torchvision *.whl files in /dist
  ansible.builtin.find:
    path: /dist
    pattern: "torchvision*.whl"
  register: torchvision_wheels

- name: Install Torchvision wheel
  ansible.builtin.pip:
    name: "{{ torchvision_wheels.files | map(attribute='path') }}"

- name: "Tests"
  include_tasks: tests.yaml
  tags:
    - tests