Add option to unbundle libtpu (#5534) (#5536)

* Add optiona to unbundle libtpu * Add clarifying note
pytorch · Sep 14, 2023 · ff30d64 · ff30d64
1 parent 2d7e92f
commit ff30d64
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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")
     })
   )