From 8849a0eee4882f9f79b9f63964d555b87c2218c8 Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 17:58:53 +0000
Subject: [PATCH 1/9] Change tests to build from built wheels rather than
 release

---
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index 6de990661526..b8c176399355 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -26,7 +26,7 @@
 ARG python_version=3.10
 ARG debian_version=bullseye
 
-FROM python:${python_version}-${debian_version} AS release
+FROM python:${python_version}-${debian_version}
 
 # Install PyTorch wheels. We expect to install three wheels. Example:
 # - torch-2.8.0-cp310-cp310-linux_x86_64.whl
@@ -45,7 +45,7 @@ RUN pip install ansible
 COPY --from=ansible . /ansible
 
 ARG ansible_vars
-RUN ansible-playbook -vvv playbook.yaml -e "stage=release" -e "${ansible_vars}" --tags "install_deps"
+RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}" --tags "install_deps"
 
 WORKDIR /
 

From a54022193306d6a7d54dbe405df2f6b8bb6fced1 Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 18:40:38 +0000
Subject: [PATCH 2/9] Build wheels for torch before trying to install

---
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index b8c176399355..99c66e628e6e 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -26,6 +26,11 @@
 ARG python_version=3.10
 ARG debian_version=bullseye
 
+# Build PyTorch and PyTorch/XLA wheels.
+ARG ansible_vars
+RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}"
+RUN ansible-playbook -vvv playbook.yaml -e "stage=build_plugin" -e "${ansible_vars}" --skip-tags=fetch_srcs,install_deps
+
 FROM python:${python_version}-${debian_version}
 
 # Install PyTorch wheels. We expect to install three wheels. Example:

From bf91e237a06949ede164435de4761280340bb128 Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 20:15:12 +0000
Subject: [PATCH 3/9] set ansible as workdir to download files

---
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index 99c66e628e6e..2320eab16755 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -26,6 +26,10 @@
 ARG python_version=3.10
 ARG debian_version=bullseye
 
+WORKDIR /ansible
+RUN pip install ansible
+COPY . /ansible
+
 # Build PyTorch and PyTorch/XLA wheels.
 ARG ansible_vars
 RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}"
@@ -39,6 +43,7 @@ FROM python:${python_version}-${debian_version}
 # - torchvision-0.22.0a0+966da7e-cp310-cp310-linux_x86_64.whl
 # The precise names will depend on the git commit hash used at build time.
 WORKDIR /tmp/wheels
+RUN ls -la
 COPY ./*.whl ./
 
 RUN echo "Installing the following wheels" && ls *.whl

From e09a6011f8d05e23f4f3252c9f2ca3e10f98204c Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 20:29:54 +0000
Subject: [PATCH 4/9] correct env stage

---
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index 2320eab16755..c02c76ab4427 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -26,6 +26,8 @@
 ARG python_version=3.10
 ARG debian_version=bullseye
 
+FROM python:${python_version}-${debian_version} AS build
+
 WORKDIR /ansible
 RUN pip install ansible
 COPY . /ansible
@@ -35,8 +37,6 @@ ARG ansible_vars
 RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}"
 RUN ansible-playbook -vvv playbook.yaml -e "stage=build_plugin" -e "${ansible_vars}" --skip-tags=fetch_srcs,install_deps
 
-FROM python:${python_version}-${debian_version}
-
 # Install PyTorch wheels. We expect to install three wheels. Example:
 # - torch-2.8.0-cp310-cp310-linux_x86_64.whl
 # - torch_xla-2.8.0+gitd4b0a48-cp310-cp310-linux_x86_64.whl

From dee39d1ea29e5fdac45845423c5cd94824df81a3 Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 20:31:27 +0000
Subject: [PATCH 5/9] add python reference post pytorch build

---
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index c02c76ab4427..72a18a044d24 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -37,6 +37,8 @@ ARG ansible_vars
 RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}"
 RUN ansible-playbook -vvv playbook.yaml -e "stage=build_plugin" -e "${ansible_vars}" --skip-tags=fetch_srcs,install_deps
 
+FROM python:${python_version}-${debian_version}
+
 # Install PyTorch wheels. We expect to install three wheels. Example:
 # - torch-2.8.0-cp310-cp310-linux_x86_64.whl
 # - torch_xla-2.8.0+gitd4b0a48-cp310-cp310-linux_x86_64.whl

From dc74d2f45b407bbe1536ec6ff96876541e4dc7d1 Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 21:31:04 +0000
Subject: [PATCH 6/9] Use newer version of python

---
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index 72a18a044d24..d3c51d00d346 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -23,7 +23,7 @@
 #
 # (Googlers only) Refer to go/ptxla-torchprime-trigger for information on retention policy of
 # the docker images.
-ARG python_version=3.10
+ARG python_version=3.12
 ARG debian_version=bullseye
 
 FROM python:${python_version}-${debian_version} AS build
@@ -45,7 +45,6 @@ FROM python:${python_version}-${debian_version}
 # - torchvision-0.22.0a0+966da7e-cp310-cp310-linux_x86_64.whl
 # The precise names will depend on the git commit hash used at build time.
 WORKDIR /tmp/wheels
-RUN ls -la
 COPY ./*.whl ./
 
 RUN echo "Installing the following wheels" && ls *.whl

From 82451fc8c4b1a3f1de8a3ec0257a9f71f20d903c Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 22:11:35 +0000
Subject: [PATCH 7/9] test vars change

---
 infra/ansible/config/vars.yaml                       | 2 +-
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/infra/ansible/config/vars.yaml b/infra/ansible/config/vars.yaml
index f34e2c3cb632..4c053a5e3ad5 100644
--- a/infra/ansible/config/vars.yaml
+++ b/infra/ansible/config/vars.yaml
@@ -7,7 +7,7 @@ cuda_compute_capabilities: 5.2,7.0,7.5,8.0,9.0
 llvm_debian_repo: bullseye
 clang_version: 17
 # PyTorch and PyTorch/XLA wheel versions.
-package_version: 2.9.0
+package_version: 2.8.0
 # If set to true, wheels will be renamed to $WHEEL_NAME-nightly-cp38-cp38-linux_x86_64.whl.
 nightly_release: false
 # Whether to preinstall libtpu in the PyTorch/XLA wheel. Ignored for GPU build.
diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index d3c51d00d346..f80f667c3bb1 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -26,7 +26,7 @@
 ARG python_version=3.12
 ARG debian_version=bullseye
 
-FROM python:${python_version}-${debian_version} AS build
+FROM python:${python_version}-${debian_version}
 
 WORKDIR /ansible
 RUN pip install ansible
@@ -37,8 +37,6 @@ ARG ansible_vars
 RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}"
 RUN ansible-playbook -vvv playbook.yaml -e "stage=build_plugin" -e "${ansible_vars}" --skip-tags=fetch_srcs,install_deps
 
-FROM python:${python_version}-${debian_version}
-
 # Install PyTorch wheels. We expect to install three wheels. Example:
 # - torch-2.8.0-cp310-cp310-linux_x86_64.whl
 # - torch_xla-2.8.0+gitd4b0a48-cp310-cp310-linux_x86_64.whl

From 04f48f1446d61a14a2d7f4c52d77356a5ed25e7e Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 23:22:42 +0000
Subject: [PATCH 8/9] Update python version everywhere for test

---
 infra/ansible/config/vars.yaml                      |  2 +-
 .../ansible/ptxla_docker_for_torchprime.Dockerfile  | 13 +------------
 infra/ansible/publish_torchprime_e2e_test_docker.sh |  2 +-
 3 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/infra/ansible/config/vars.yaml b/infra/ansible/config/vars.yaml
index 4c053a5e3ad5..f34e2c3cb632 100644
--- a/infra/ansible/config/vars.yaml
+++ b/infra/ansible/config/vars.yaml
@@ -7,7 +7,7 @@ cuda_compute_capabilities: 5.2,7.0,7.5,8.0,9.0
 llvm_debian_repo: bullseye
 clang_version: 17
 # PyTorch and PyTorch/XLA wheel versions.
-package_version: 2.8.0
+package_version: 2.9.0
 # If set to true, wheels will be renamed to $WHEEL_NAME-nightly-cp38-cp38-linux_x86_64.whl.
 nightly_release: false
 # Whether to preinstall libtpu in the PyTorch/XLA wheel. Ignored for GPU build.
diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index f80f667c3bb1..90ac695411f3 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -26,17 +26,6 @@
 ARG python_version=3.12
 ARG debian_version=bullseye
 
-FROM python:${python_version}-${debian_version}
-
-WORKDIR /ansible
-RUN pip install ansible
-COPY . /ansible
-
-# Build PyTorch and PyTorch/XLA wheels.
-ARG ansible_vars
-RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}"
-RUN ansible-playbook -vvv playbook.yaml -e "stage=build_plugin" -e "${ansible_vars}" --skip-tags=fetch_srcs,install_deps
-
 # Install PyTorch wheels. We expect to install three wheels. Example:
 # - torch-2.8.0-cp310-cp310-linux_x86_64.whl
 # - torch_xla-2.8.0+gitd4b0a48-cp310-cp310-linux_x86_64.whl
@@ -54,7 +43,7 @@ RUN pip install ansible
 COPY --from=ansible . /ansible
 
 ARG ansible_vars
-RUN ansible-playbook -vvv playbook.yaml -e "stage=build" -e "${ansible_vars}" --tags "install_deps"
+RUN ansible-playbook -vvv playbook.yaml -e "stage=release" -e "${ansible_vars}" --tags "install_deps"
 
 WORKDIR /
 
diff --git a/infra/ansible/publish_torchprime_e2e_test_docker.sh b/infra/ansible/publish_torchprime_e2e_test_docker.sh
index 4e8c3d23009d..cee169b95d79 100755
--- a/infra/ansible/publish_torchprime_e2e_test_docker.sh
+++ b/infra/ansible/publish_torchprime_e2e_test_docker.sh
@@ -66,7 +66,7 @@ docker build -t "${DOCKER_URL}" \
     "${DEFAULT_CONTEXT_PATH}" \
     -f "${DOCKERFILE_PATH}" \
     --build-arg ansible_vars="${ANSIBLE_VARS_COMPACT}" \
-    --build-arg python_version=3.10 \
+    --build-arg python_version=3.12 \
     --build-arg debian_version=bullseye
 docker push "${DOCKER_URL}"
 

From 9593df21a27b1a003b12ac83c87143c5a958e51a Mon Sep 17 00:00:00 2001
From: Pedro Goncalves Mokarzel <pgmoka@google.com>
Date: Tue, 15 Jul 2025 23:50:44 +0000
Subject: [PATCH 9/9] Minor dockerfile mistake

---
 infra/ansible/ptxla_docker_for_torchprime.Dockerfile | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
index 90ac695411f3..e9885bd49284 100644
--- a/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
+++ b/infra/ansible/ptxla_docker_for_torchprime.Dockerfile
@@ -26,6 +26,8 @@
 ARG python_version=3.12
 ARG debian_version=bullseye
 
+FROM python:${python_version}-${debian_version} AS release
+
 # Install PyTorch wheels. We expect to install three wheels. Example:
 # - torch-2.8.0-cp310-cp310-linux_x86_64.whl
 # - torch_xla-2.8.0+gitd4b0a48-cp310-cp310-linux_x86_64.whl