From b94c014e346e0cc4ac432561d92a27382c9d464b Mon Sep 17 00:00:00 2001
From: Siyuan Liu <lsiyuan@google.com>
Date: Wed, 24 Apr 2024 13:14:28 -0700
Subject: [PATCH] Build torchvision from the pytorch pin in ansible (#6756)
 (#6960)

Co-authored-by: Will Cromar <wcromar@google.com>
---
 infra/ansible/config/pip.yaml                 |  2 --
 infra/ansible/e2e_tests.Dockerfile            |  8 ++------
 .../ansible/roles/build_srcs/tasks/main.yaml  | 20 +++++++++++++++++++
 .../ansible/roles/fetch_srcs/tasks/main.yaml  | 19 ------------------
 test/tpu/xla_test_job.yaml                    |  2 --
 5 files changed, 22 insertions(+), 29 deletions(-)

diff --git a/infra/ansible/config/pip.yaml b/infra/ansible/config/pip.yaml
index 7082c33f6e8..1c82075abef 100644
--- a/infra/ansible/config/pip.yaml
+++ b/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
diff --git a/infra/ansible/e2e_tests.Dockerfile b/infra/ansible/e2e_tests.Dockerfile
index 7bedc096c84..ab29d2dc76c 100644
--- a/infra/ansible/e2e_tests.Dockerfile
+++ b/infra/ansible/e2e_tests.Dockerfile
@@ -23,14 +23,10 @@ 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
@@ -38,4 +34,4 @@ RUN pip install *.whl
 WORKDIR /
 
 # Clean-up unused directories.
-RUN rm -rf /ansible /tmp/wheels
\ No newline at end of file
+RUN rm -rf /ansible /tmp/wheels
diff --git a/infra/ansible/roles/build_srcs/tasks/main.yaml b/infra/ansible/roles/build_srcs/tasks/main.yaml
index 3a785a0f8e7..b2b5859fb8f 100644
--- a/infra/ansible/roles/build_srcs/tasks/main.yaml
+++ b/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:
diff --git a/infra/ansible/roles/fetch_srcs/tasks/main.yaml b/infra/ansible/roles/fetch_srcs/tasks/main.yaml
index 5a0ed7c36f5..438e015b25c 100644
--- a/infra/ansible/roles/fetch_srcs/tasks/main.yaml
+++ b/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:
diff --git a/test/tpu/xla_test_job.yaml b/test/tpu/xla_test_job.yaml
index 4b179fe66a5..5091555a508 100644
--- a/test/tpu/xla_test_job.yaml
+++ b/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