From ff30d64c5cc572c4342fc40135bb87e146c36739 Mon Sep 17 00:00:00 2001
From: Will Cromar <wcromar@google.com>
Date: Tue, 5 Sep 2023 13:20:52 -0700
Subject: [PATCH] Add option to unbundle libtpu (#5534) (#5536)

* Add optiona to unbundle libtpu

* Add clarifying note
---
 infra/ansible/config/env.yaml                    | 3 +--
 infra/ansible/config/vars.yaml                   | 2 ++
 infra/tpu-pytorch-releases/artifacts.auto.tfvars | 3 ++-
 infra/tpu-pytorch-releases/artifacts_builds.tf   | 2 ++
 4 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/infra/ansible/config/env.yaml b/infra/ansible/config/env.yaml
index 783fc69b621c..cecb0ec482d5 100644
--- a/infra/ansible/config/env.yaml
+++ b/infra/ansible/config/env.yaml
@@ -47,5 +47,4 @@ build_env:
   tpu:
     ACCELERATOR: tpu
     TPUVM_MODE: 1
-    BUNDLE_LIBTPU: 1
-
+    BUNDLE_LIBTPU: "{{ bundle_libtpu }}"
diff --git a/infra/ansible/config/vars.yaml b/infra/ansible/config/vars.yaml
index 167ec96720b8..994049d1c506 100644
--- a/infra/ansible/config/vars.yaml
+++ b/infra/ansible/config/vars.yaml
@@ -10,3 +10,5 @@ package_version: 2.1.0
 nightly_release: false
 # Whether to disable XRT during build
 disable_xrt: 0
+# Whether to preinstall libtpu in the PyTorch/XLA wheel. Ignored for GPU build.
+bundle_libtpu: 1
diff --git a/infra/tpu-pytorch-releases/artifacts.auto.tfvars b/infra/tpu-pytorch-releases/artifacts.auto.tfvars
index 524d9ebaaf55..0b1fb0cf187c 100644
--- a/infra/tpu-pytorch-releases/artifacts.auto.tfvars
+++ b/infra/tpu-pytorch-releases/artifacts.auto.tfvars
@@ -46,7 +46,8 @@ versioned_builds = [
     pytorch_git_tag = "v2.1.0-rc2"
     package_version = "2.1"
     accelerator     = "tpu"
-    python_version = "3.10"
+    python_version  = "3.10"
+    bundle_libtpu   = "0"
   },
   {
     git_tag         = "v2.0.0"
diff --git a/infra/tpu-pytorch-releases/artifacts_builds.tf b/infra/tpu-pytorch-releases/artifacts_builds.tf
index 203ccbaa0c84..aa0a0c6b6bd6 100644
--- a/infra/tpu-pytorch-releases/artifacts_builds.tf
+++ b/infra/tpu-pytorch-releases/artifacts_builds.tf
@@ -9,6 +9,7 @@ variable "nightly_builds" {
       cuda_version   = optional(string, "11.8")
       python_version = optional(string, "3.8")
       arch           = optional(string, "amd64")
+      bundle_libtpu  = optional(string, "1")
     })
   )
 
@@ -39,6 +40,7 @@ variable "versioned_builds" {
       python_version  = optional(string, "3.8")
       cuda_version    = optional(string, "11.8")
       arch            = optional(string, "amd64")
+      bundle_libtpu   = optional(string, "1")
     })
   )