Is it possible to customize TPU initialization timeout?

[ayaka14732]

I always get RuntimeError: Unable to initialize backend 'tpu': UNKNOWN: TPU initialization failed: Worker failed to join a slice within 15m (set JAX_PLATFORMS='' to automatically choose an available backend) when TPU is not configured properly.

Is it possible to manually decrease the timeout value via some environment variables so that I do not need to wait for so long before identifying the error?

@skye

[louwjac]

Have you tried specifying the "initialization_timeout" parameter in the call to jax.distributed.initialize() ?

e.g. jax.distributed.initialize(initialization_timeout=30)

https://jax.readthedocs.io/en/latest/_autosummary/jax.distributed.initialize.html