Add tpu-info as a dependency of torch_xla[tpu] (#7938)

pytorch · Sep 3, 2024 · d69f3a5 · d69f3a5
1 parent 9e3b448
commit d69f3a5
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 7 deletions.
diff --git a/setup.py b/setup.py
@@ -312,7 +312,7 @@ def run(self):
     extras_require={
         # On Cloud TPU VM install with:
         # pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
-        'tpu': [f'libtpu-nightly=={_libtpu_version}'],
+        'tpu': [f'libtpu-nightly=={_libtpu_version}', 'tpu-info'],
         # pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
         'pallas': [f'jaxlib=={_jax_version}', f'jax=={_jax_version}'],
     },

diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh
@@ -56,6 +56,4 @@ if [[ -n "$TPU_VERSION" && "$TPU_VERSION" == "4" ]]; then
 fi
 
 # Test `tpu-info` CLI compatibility
-# https://github.com/google/cloud-accelerator-diagnostics/tree/main/tpu_info
-pip install -r test/tpu/tpu_info/requirements.txt
 python3 test/tpu/tpu_info/test_cli.py
diff --git a/test/tpu/tpu_info/requirements.txt b/test/tpu/tpu_info/requirements.txt
diff --git a/test/tpu/tpu_info/test_cli.py b/test/tpu/tpu_info/test_cli.py
@@ -85,8 +85,7 @@ def _torch_xla_spawn(self):
 
       # v2 and v3 may have duplicates due to multithreading
       child_pids = set()
-      for _ in range(self.chip_type.value.accelerators_per_chip *
-                     self.num_chips):
+      for _ in range(self.chip_type.value.devices_per_chip * self.num_chips):
         child_pids.add(q.get(timeout=20.0))
       with contextlib.ExitStack() as e:
         e.callback(done.set)