Build torchvision from the pytorch pin in ansible (#6756) (#6960)

Co-authored-by: Will Cromar <wcromar@google.com>

infra/ansible/config/pip.yaml

@@ -49,5 +49,3 @@ pip:
   # Packages that will be installed with the `--nodeps` flag.
   pkgs_nodeps:
     release_common:
-      - torchvision
-      - pillow

infra/ansible/e2e_tests.Dockerfile

@@ -23,19 +23,15 @@ RUN ansible-playbook -vvv playbook.yaml -e "stage=release" -e "${ansible_vars}"
 # Copy test sources.
 RUN mkdir -p /src/pytorch/xla
 COPY --from=build /src/pytorch/xla/test /src/pytorch/xla/test
-# Copy ci_commit_pins from upstream
-RUN mkdir -p /src/pytorch/.github
-COPY --from=build /src/pytorch/.github/ci_commit_pins /src/pytorch/.github/ci_commit_pins
 
 # Copy and install wheels.
 WORKDIR /tmp/wheels
-COPY --from=build /src/pytorch/dist/*.whl ./
-COPY --from=build /src/pytorch/xla/dist/*.whl ./
+COPY --from=build /dist/*.whl ./
 
 RUN echo "Installing the following wheels" && ls *.whl
 RUN pip install *.whl
 
 WORKDIR /
 
 # Clean-up unused directories.
-RUN rm -rf /ansible /tmp/wheels
+RUN rm -rf /ansible /tmp/wheels

infra/ansible/roles/build_srcs/tasks/main.yaml

@@ -84,6 +84,26 @@
     state: absent
     mode: '0755'
 
+- 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:

infra/ansible/roles/fetch_srcs/tasks/main.yaml

@@ -20,25 +20,6 @@
       dest: "{{ (src_root, 'pytorch/xla') | path_join }}"
       version: "{{ xla_git_rev }}"
 
-- name: Find *.diff files in pytorch/xla/openxla_patches
-  ansible.builtin.find:
-    path: "{{ (src_root, 'pytorch/xla/openxla_patches') | path_join }}"
-    pattern: "*.diff"
-  register: openxla_patches
-
-- name: Apply patches to OpenXLA
-  ansible.posix.patch:
-    src: "{{ item }}"
-    # Use source file on the target machine instead of the one where
-    # the playbook is located. Has no effect when the target machine is
-    # localhost.
-    remote_src: true
-    strip: 1
-    ignore_whitespace: true
-    basedir: "{{ (src_root, 'pytorch/xla/third_party/xla') | path_join }}"
-  loop: "{{ openxla_patches.files | map(attribute='path') }}"
-  ignore_errors: true
-
 - name: "Tests"
   include_tasks: tests.yaml
   tags:

test/tpu/xla_test_job.yaml

@@ -41,8 +41,6 @@ spec:
     - bash
     - -cxe
     - |
-      # Torchvision installed from a commit requires torch, need to install after torch is compiled.
-      pip install --no-use-pep517 "git+https://github.com/pytorch/vision.git@$(cat /src/pytorch/.github/ci_commit_pins/vision.txt)"
       pip install expecttest==0.1.6
       pip install rich
       pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html