JAX doesn't see local devices in multi-host setting

[kimihailv]

I try to run simple multi-host program on two nodes (each node has 4 gpus) which are connected via ethernet:

import jax
import sys

print(sys.argv)
jax.distributed.initialize(
    sys.argv[1],
    num_processes=2,
    process_id=int(sys.argv[2]),
    local_device_ids=[0, 1, 2, 3]
)

print('hello', sys.argv[2])
print(jax.local_device_count())

On the first node I use this command

NCCL_SOCKET_IFNAME=enp36s0f1 python3 multinode.py <ip>:9000 0

On the second node I use this command:

NCCL_SOCKET_IFNAME=enp36s0f1 python3 multinode.py <ip>:9000 1

Program prints the following output (on the first and second nodes respectively):

['multinode.py', '192.168.6.2:9000', '0']
hello 0

['multinode.py', '192.168.6.2:9000', '1']
hello 1
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
1

So, the first node didn't print the number of local devices, the second one couldn't even find them. However, when I run jax scripts which use all gpus on the first/second node, everything works as expected.