Pin torch_xla base image to 20250217 (#113)

* Pin torch_xla base image to 20250217

* Update Dockerfile

* Update Dockerfile

* Add the workaround flag in thunk

Author: tengyifei

.github/workflows/cpu_test.yml

@@ -15,7 +15,7 @@ jobs:
       matrix:
         python-version: ["3.10", "3.11"]
     container:
-      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_${{ matrix.python-version }}_tpuvm
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_${{ matrix.python-version }}_tpuvm_20250217
     steps:
     - uses: actions/checkout@v4
     - name: Install torchax

torchprime/launcher/Dockerfile

@@ -1,7 +1,6 @@
 # syntax=docker/dockerfile:experimental
 # Use torch_xla Python 3.10 as the base image
-# TODO(https://github.com/pytorch/xla/issues/8683): Go back to nightly once the linked segfault is fixed.
-FROM us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11
+FROM us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_20250217
 
 ARG USE_TRANSFORMERS=false
 # Install system dependencies

torchprime/launcher/thunk.py

@@ -4,6 +4,14 @@
 from datetime import datetime
 from pathlib import Path
 
+# Workaround for MegaScale crash
+#
+# TODO(https://github.com/pytorch/xla/issues/8683): Remove the
+# `--megascale_grpc_enable_xor_tracer=false` flag when libtpu is updated
+xla_flags = os.environ.get("LIBTPU_INIT_ARGS", "")
+xla_flags = f"{xla_flags} --megascale_grpc_enable_xor_tracer=false"
+os.environ["LIBTPU_INIT_ARGS"] = xla_flags
+
 # Get the artifact dir from env var.
 gcs_artifact_dir = os.environ["TORCHPRIME_ARTIFACT_DIR"]
 assert gcs_artifact_dir.startswith(