Xla: xmp.spawn hangs

[magicknight]

❓ Questions and Help

Hey guys,

I am trying to run the multiple processes tutorial on a n1-standard-8 VM with a v3-8 TPU, using torch-nightly VM image. The codes below are from tutorial, but it hangs for a few minutes before printing the error output.

import torch
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp

# "Map function": acquires a corresponding Cloud TPU core, creates a tensor on it,
# and prints its core
def simple_map_fn(index, flags):
  torch.manual_seed(11)
  device = xm.xla_device()  
  t = torch.randn((2, 2), device=device)
  print("Process", index ,"is using", xm.xla_real_devices([str(device)])[0])

# the Cloud TPU
flags = {}
# Note: Colab only supports start_method='fork'
xmp.spawn(simple_map_fn, args=(flags,), nprocs=8, start_method='fork')

output:

Process 0 is using TPU:0
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-1-43e75d6c3925> in <module>
     15 flags = {}
     16 # Note: Colab only supports start_method='fork'
---> 17 xmp.spawn(simple_map_fn, args=(flags,), nprocs=8, start_method='fork')

/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
    180         join=join,
    181         daemon=daemon,
--> 182         start_method=start_method)

/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
    156 
    157     # Loop on join until it returns True or raises an exception.
--> 158     while not context.join():
    159         pass
    160 

/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    117         msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
    118         msg += original_trace
--> 119         raise Exception(msg)
    120 
    121 

Exception: 

-- Process 2 terminated with the following error:
Traceback (most recent call last):
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
    fn(i, *args)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 116, in _start_fn
    _setup_replication()
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 108, in _setup_replication
    device = xm.xla_device()
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 137, in xla_device
    devkind=[devkind] if devkind is not None else None)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 41, in get_xla_supported_devices
    xla_devices = torch_xla._XLAC._xla_get_devices()
RuntimeError: tensorflow/compiler/xla/xla_client/mesh_service.cc:212 : 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()


    _PyCFunction_FastCallDict

    _PyEval_EvalFrameDefault



    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault
    PyEval_EvalCodeEx

    PyObject_Call
    _PyEval_EvalFrameDefault
    PyEval_EvalCodeEx

    PyObject_Call
    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault
    _PyFunction_FastCallDict
    _PyObject_FastCallDict
    _PyObject_Call_Prepend
    PyObject_Call


    _PyObject_FastCallDict

    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault



    _PyEval_EvalFrameDefault



    _PyEval_EvalFrameDefault
    PyEval_EvalCodeEx
    PyEval_EvalCode

    _PyCFunction_FastCallDict

    _PyEval_EvalFrameDefault
    _PyGen_Send
    _PyEval_EvalFrameDefault
    _PyGen_Send
    _PyEval_EvalFrameDefault
    _PyGen_Send
    _PyCFunction_FastCallDict

    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault

    _PyFunction_FastCallDict
    _PyObject_FastCallDict
    _PyObject_Call_Prepend
    PyObject_Call
    _PyEval_EvalFrameDefault



    _PyEval_EvalFrameDefault

    _PyCFunction_FastCallDict

    _PyEval_EvalFrameDefault



    _PyEval_EvalFrameDefault

    _PyCFunction_FastCallDict

    _PyEval_EvalFrameDefault



    _PyEval_EvalFrameDefault

    _PyCFunction_FastCallDict

    _PyEval_EvalFrameDefault

    _PyFunction_FastCallDict
    _PyObject_FastCallDict
    _PyObject_Call_Prepend
    PyObject_Call
    _PyEval_EvalFrameDefault
    _PyGen_Send
    _PyCFunction_FastCallDict

    _PyEval_EvalFrameDefault


    _PyEval_EvalFrameDefault

    _PyFunction_FastCallDict
    _PyObject_FastCallDict

    _PyObject_FastCallDict
*** End stack trace ***
Failed to connect to client mesh master: localhost:40947

Any suggestions would be really appreciated!

[JackCaoG]

Failure pattern is similar to https://github.com/pytorch/xla/issues/1795

Failed to connect to client mesh master

Might be the similar issue?

[dlibenzi]

The master exits before the clients can connect.
Add this at the bottom (inside the map_fn):

xm.rendezvous('checking_out')

[magicknight]

The master exits before the clients can connect.
Add this at the bottom (inside the map_fn):

xm.rendezvous('checking_out')

Wow it works like a magic. Thank you!