Spawning functions that immediately return hangs

[mruberry]

🐛 Bug

import torch_xla.distributed.xla_multiprocessing as xmp

def _mp_fn(rank, flags):
  return 0

FLAGS={}
xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8,
          start_method='fork')

Hangs. Eventually it errors to:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 116, in _start_fn
    _setup_replication()
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 108, in _setup_replication
    device = xm.xla_device()
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/core/xla_model.py", line 143, in xla_device
    devices = get_xla_supported_devices(devkind=devkind)
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/core/xla_model.py", line 42, in get_xla_supported_devices
    xla_devices = torch_xla._XLAC._xla_get_devices()
RuntimeError: tensorflow/compiler/xla/xla_client/mesh_service.cc:172 : Check failed: impl_->channel->WaitForConnected( std::chrono::system_clock::now() + std::chrono::seconds(connect_wait_seconds)) 
*** Begin stack trace ***
	tensorflow::CurrentStackTrace[abi:cxx11]()
	xla::service::MeshClient::MeshClient(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)
	xla::service::MeshClient::Get()
	xla::ComputationClient::Create()
	xla::ComputationClient::Get()
	...

Reproduces consistently on Colab.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.