test

Author: goanpeca

.github/workflows/tpu-tests.yml

@@ -61,7 +61,6 @@ jobs:
       - name: Get pip cache dir
         id: pip-cache
         run: |
-          pip3 install -U "pip<24"
           echo "pip_cache=$(pip cache dir)" >> $GITHUB_OUTPUT
         shell: bash -l {0}
 

scripts/install_torch_xla.py

@@ -1,12 +1,13 @@
 import os
+import sys
 import tempfile
 import zipfile
 from pathlib import Path
 from urllib.request import urlopen
 
 
 def main(python_version: str, xla_version: str) -> None:
-    temp_dir = Path(tempfile.gettempdir()) / "torch_xla"
+    temp_dir = Path(tempfile.mkdtemp()) / "torch_xla"
     temp_dir.mkdir(parents=True, exist_ok=True)
     pyver = python_version.replace(".", "")
     url = f"https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-{xla_version}-cp{pyver}-cp{pyver}-linux_x86_64.whl"
@@ -16,20 +17,19 @@ def main(python_version: str, xla_version: str) -> None:
 
     with open(wheel_path, "b+w") as f:
         f.write(r.read())
-    
+
     with zipfile.ZipFile(wheel_path, "r") as zip_ref:
-       zip_ref.extractall(wheel_path.parent)
+        zip_ref.extractall(wheel_path.parent)
 
     dist = [p for p in os.listdir(temp_dir) if "dist-info" in p]
     if dist:
-        pkg_name =  dist[0].split(".dist-info")[0]
+        pkg_name = dist[0].split(".dist-info")[0]
         wheel_path = wheel_path.rename(f"{pkg_name}{suffix}")
         cmd = ["uv", "pip", "install", str(wheel_path)]
         os.system(" ".join(cmd))
 
-    
+
 if __name__ == "__main__":
-    import sys
     if len(sys.argv) == 3:
         python_version = sys.argv[1]
         xla_version = sys.argv[2]