TPU VM (v4-32, tpu-ubuntu2204-base) hangs when verifying TPU with `xm.get_xla_supported_devices("TPU")`

[renxiangdai]

Environment

• TPU hardware: v4-32 (TPU VM)
• OS image: tpu-ubuntu2204-base

Steps to reproduce

On a fresh TPU VM:

bash
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip install numpy
pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

Then run the verification command:

PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

Expected behavior

The command should return the list of available TPU devices, e.g.:
['TPU:0', 'TPU:1', ...]

Actual behavior

The command hangs indefinitely and never prints any output.
This makes it impossible to confirm TPU accessibility from PyTorch/XLA on the TPU VM.

Additional context

• This was run directly on a TPU VM (not Colab).

• No error messages are shown — just hangs after import.

[ysiraichi]

I tried on a v5litepod-4 with software version v2-alpha-tpuv5-lite , and wasn't able to reproduce the issue.
Instead, I got:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "torch_xla/core/xla_model.py", line 82, in get_xla_supported_devices
    for i, _ in enumerate(devices[:max_devices] if max_devices else devices)
TypeError: slice indices must be integers or None or have an __index__ method

Which is (kind of) expected.
Note that get_xla_supported_devices() devkind parameter was deprecated at #9284.

Removing the "TPU" argument, I get:

['xla:0', 'xla:1', 'xla:2', 'xla:3']

[ysiraichi]

If you stop the execution, what do you get?

[Muinez]

I'm experiencing same issue with v4-32 (4 hosts).

The command PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices())" completes successfully only when run on all hosts simultaneously. If I run it on a single host, the process hangs indefinitely.
When running the official PyTorch TPU Pod examples from the documentation (https://cloud.google.com/tpu/docs/pytorch-pods), the training process appears to be desynchronized. It seems each host trains its own model independently.
The key indicators of this are:
The training steps are out of sync across hosts (e.g., one host is at step 80 while another is at 120).
It produces four separate logs instead of one unified log.
I've observed the exact same behavior with my own custom training script as well.

I have also tried using both .spawn and the .launch script to initiate the training. In both cases, when I launch the script across all hosts, the processes hang indefinitely and never start the actual training

@ysiraichi

[Muinez]

I'm experiencing same issue with v4-32 (4 hosts).

The command PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices())" completes successfully only when run on all hosts simultaneously. If I run it on a single host, the process hangs indefinitely. When running the official PyTorch TPU Pod examples from the documentation (https://cloud.google.com/tpu/docs/pytorch-pods), the training process appears to be desynchronized. It seems each host trains its own model independently. The key indicators of this are: The training steps are out of sync across hosts (e.g., one host is at step 80 while another is at 120). It produces four separate logs instead of one unified log. I've observed the exact same behavior with my own custom training script as well.

I have also tried using both .spawn and the .launch script to initiate the training. In both cases, when I launch the script across all hosts, the processes hang indefinitely and never start the actual training

@ysiraichi

I've recreated the instance, and that has fixed .launch and desync issues. I guess I messed up the network settings during the creation process