`test/tpu/tpu_info/test_cli.py` is SegFaulting on CI.

[ysiraichi]

🐛 Bug

This issue is blocking the following PRs:

• #9564
• #9566

+ python3 /home/runner/_work/xla/xla/pytorch/xla/test/tpu/tpu_info/test_cli.py
test/tpu/run_training_tests.sh: line 39: 36020 Segmentation fault      python3 "$_TPU_DIR/tpu_info/test_cli.py"
Error: Process completed with exit code 139.

I can see it first breaking on master after #9556 was merged.
That said, it doesn't look like it was the culprit.

[ysiraichi]

@tengyifei @bhavya01 @zhanyong-wan @qihqi
Any idea what might be causing it?

[ysiraichi]

I tried reproducing this error on a v5litepod-4, but didn't get to the same error.
Got a TPU initialization error, instead.

[qihqi]

Let's remove this test (permanently).

The background of tpu-info is:
It is originally source-collocated in Pytorch/XLA, becuase it is written by one of our Engs. Later on it is moved to https://github.com/AI-Hypercomputer/cloud-accelerator-diagnostics because it is not a Pytorch specific tool. It's test remained.

Now, speculating on what could have triggered the segfault:
Recently I have seen this error:

F0813 02:28:04.280416  245819 pjrt_c_api_helpers.cc:258] Unexpected error status Unexpected PJRT_Plugin_Attributes_Args size: expected 32, got 24. The plugin is likely built with a later version than the framework. This plugin is built with PJRT API version 0.73.
*** Check failure stack trac

in attempting to load torch_xla with a more recent libtpu.

Looks like the version bump happened around 2 weeks ago (it's now 75). (openxla/xla@877a57d)

Then, given that tpu-info tool is installed from pip, it might pull it's own libtpu.

So if torch_xla (uses old libtpu) + new libtpu = warning & exception, then it's reasonable to believe that tpu-info(uses new libtpu) + old libtpu = other failure modes, that would include segfault.

[ysiraichi]

@qihqi Thank you for chiming in, and for landing the PR with the fix.
I did investigate it a bit more, and was actually able to reproduce the error only when using Python 3.12.
I think it could be due to the PyPI version conflict you mentioned, since the segmentation fault seems to be coming from the import of tpu_info.

~/pytorch/xla$ python -X faulthandler test/tpu/tpu_info/test_cli.py
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
Fatal Python error: Segmentation fault

Current thread 0x00007f6f9241a800 (most recent call first):
  File "<frozen importlib._bootstrap>", line 488 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 1293 in create_module
  File "<frozen importlib._bootstrap>", line 813 in module_from_spec
  File "<frozen importlib._bootstrap>", line 921 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 1331 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 1360 in _find_and_load
  File "<frozen importlib._bootstrap>", line 488 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1415 in _handle_fromlist
  File "/home/ysiraichi/.local/lib/python3.12/site-packages/tpu_info/cli_helper.py", line 29 in <module>
  File "<frozen importlib._bootstrap>", line 488 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 999 in exec_module
  File "<frozen importlib._bootstrap>", line 935 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 1331 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 1360 in _find_and_load
  File "<frozen importlib._bootstrap>", line 488 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1415 in _handle_fromlist
  File "/home/ysiraichi/.local/lib/python3.12/site-packages/tpu_info/cli.py", line 26 in <module>
  File "<frozen importlib._bootstrap>", line 488 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 999 in exec_module
  File "<frozen importlib._bootstrap>", line 935 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 1331 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 1360 in _find_and_load
  File "<frozen importlib._bootstrap>", line 488 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1415 in _handle_fromlist
  File "/home/ysiraichi/.local/lib/python3.12/site-packages/tpu_info/__init__.py", line 16 in <module>
  File "<frozen importlib._bootstrap>", line 488 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 999 in exec_module
  File "<frozen importlib._bootstrap>", line 935 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 1331 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 1360 in _find_and_load
  File "/home/ysiraichi/pytorch/xla/test/tpu/tpu_info/test_cli.py", line 10 in <module>

Extension modules: numpy._core._multiarray_umath, numpy.linalg._umath_linalg, torch._C, torch._C._dynamo.autograd_compiler, torch._C._dynamo.eval_frame, torch._C._dynamo.guards, torch._C._dynamo.utils, torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._nn, torch._C._sparse, torch._C._special, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._pcg64, numpy.random._mt19937, numpy.random._generator, numpy.random._philox, numpy.random._sfc64, numpy.random.mtrand, grpc._cython.cygrpc, google._upb._message (total: 24)
Segmentation fault (core dumped)

~/pytorch/xla$ python
Python 3.12.11 (main, Jun  4 2025, 08:56:18) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from tpu_info import cli
Segmentation fault (core dumped)