wehub-resource-sync
229 lines
7.1 KiB
Python
229 lines
7.1 KiB
Python
"""Functions used by client."""
|
|
|
|
import atexit
|
|
import logging
|
|
import os
|
|
import socket
|
|
import time
|
|
|
|
from . import rpc
|
|
from .constants import MAX_QUEUE_SIZE
|
|
|
|
if os.name != "nt":
|
|
import fcntl
|
|
import struct
|
|
|
|
|
|
def local_ip4_addr_list():
|
|
"""Return a set of IPv4 address
|
|
|
|
You can use
|
|
`logging.getLogger("dgl-distributed-socket").setLevel(logging.WARNING+1)`
|
|
to disable the warning here
|
|
"""
|
|
assert os.name != "nt", "Do not support Windows rpc yet."
|
|
nic = set()
|
|
logger = logging.getLogger("dgl-distributed-socket")
|
|
for if_nidx in socket.if_nameindex():
|
|
name = if_nidx[1]
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
try:
|
|
ip_of_ni = fcntl.ioctl(
|
|
sock.fileno(),
|
|
0x8915, # SIOCGIFADDR
|
|
struct.pack("256s", name[:15].encode("UTF-8")),
|
|
)
|
|
except OSError as e:
|
|
if e.errno == 99: # EADDRNOTAVAIL
|
|
logger.warning(
|
|
"Warning! Interface: %s \n"
|
|
"IP address not available for interface.",
|
|
name,
|
|
)
|
|
continue
|
|
raise e
|
|
|
|
ip_addr = socket.inet_ntoa(ip_of_ni[20:24])
|
|
nic.add(ip_addr)
|
|
return nic
|
|
|
|
|
|
def get_local_machine_id(server_namebook):
|
|
"""Given server_namebook, find local machine ID
|
|
|
|
Parameters
|
|
----------
|
|
server_namebook: dict
|
|
IP address namebook of server nodes, where key is the server's ID
|
|
(start from 0) and value is the server's machine_id, IP address,
|
|
port, and group_count, e.g.,
|
|
|
|
{0:'[0, '172.31.40.143', 30050, 2],
|
|
1:'[0, '172.31.40.143', 30051, 2],
|
|
2:'[1, '172.31.36.140', 30050, 2],
|
|
3:'[1, '172.31.36.140', 30051, 2],
|
|
4:'[2, '172.31.47.147', 30050, 2],
|
|
5:'[2, '172.31.47.147', 30051, 2],
|
|
6:'[3, '172.31.30.180', 30050, 2],
|
|
7:'[3, '172.31.30.180', 30051, 2]}
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
local machine ID
|
|
"""
|
|
res = 0
|
|
ip_list = local_ip4_addr_list()
|
|
for _, data in server_namebook.items():
|
|
machine_id = data[0]
|
|
ip_addr = data[1]
|
|
if ip_addr in ip_list:
|
|
res = machine_id
|
|
break
|
|
return res
|
|
|
|
|
|
def get_local_usable_addr(probe_addr):
|
|
"""Get local usable IP and port
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
IP address, e.g., '192.168.8.12:50051'
|
|
"""
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
try:
|
|
# should get the address on the same subnet as probe_addr's
|
|
sock.connect((probe_addr, 1))
|
|
ip_addr = sock.getsockname()[0]
|
|
except ValueError:
|
|
ip_addr = "127.0.0.1"
|
|
finally:
|
|
sock.close()
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.bind(("", 0))
|
|
sock.listen(1)
|
|
port = sock.getsockname()[1]
|
|
sock.close()
|
|
|
|
return ip_addr + ":" + str(port)
|
|
|
|
|
|
def connect_to_server(
|
|
ip_config,
|
|
num_servers,
|
|
max_queue_size=MAX_QUEUE_SIZE,
|
|
group_id=0,
|
|
):
|
|
"""Connect this client to server.
|
|
|
|
Parameters
|
|
----------
|
|
ip_config : str
|
|
Path of server IP configuration file.
|
|
num_servers : int
|
|
server count on each machine.
|
|
max_queue_size : int
|
|
Maximal size (bytes) of client queue buffer (~20 GB on default).
|
|
Note that the 20 GB is just an upper-bound and DGL uses zero-copy and
|
|
it will not allocate 20GB memory at once.
|
|
group_id : int
|
|
Indicates which group this client belongs to. Clients that are
|
|
booted together in each launch are gathered as a group and should
|
|
have same unique group_id.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError : If anything wrong with the connection.
|
|
"""
|
|
assert num_servers > 0, (
|
|
"num_servers (%d) must be a positive number." % num_servers
|
|
)
|
|
assert max_queue_size > 0, (
|
|
"queue_size (%d) cannot be a negative number." % max_queue_size
|
|
)
|
|
# Register some basic service
|
|
rpc.register_service(
|
|
rpc.CLIENT_REGISTER,
|
|
rpc.ClientRegisterRequest,
|
|
rpc.ClientRegisterResponse,
|
|
)
|
|
rpc.register_service(rpc.SHUT_DOWN_SERVER, rpc.ShutDownRequest, None)
|
|
rpc.register_service(
|
|
rpc.GET_NUM_CLIENT,
|
|
rpc.GetNumberClientsRequest,
|
|
rpc.GetNumberClientsResponse,
|
|
)
|
|
rpc.register_service(
|
|
rpc.CLIENT_BARRIER, rpc.ClientBarrierRequest, rpc.ClientBarrierResponse
|
|
)
|
|
rpc.register_sig_handler()
|
|
server_namebook = rpc.read_ip_config(ip_config, num_servers)
|
|
num_servers = len(server_namebook)
|
|
rpc.set_num_server(num_servers)
|
|
# group_count means how many servers
|
|
# (main_server + bakcup_server) in total inside a machine.
|
|
group_count = []
|
|
max_machine_id = 0
|
|
for server_info in server_namebook.values():
|
|
group_count.append(server_info[3])
|
|
if server_info[0] > max_machine_id:
|
|
max_machine_id = server_info[0]
|
|
rpc.set_num_server_per_machine(group_count[0])
|
|
num_machines = max_machine_id + 1
|
|
rpc.set_num_machines(num_machines)
|
|
machine_id = get_local_machine_id(server_namebook)
|
|
rpc.set_machine_id(machine_id)
|
|
rpc.set_group_id(group_id)
|
|
rpc.create_sender(max_queue_size)
|
|
rpc.create_receiver(max_queue_size)
|
|
# Get connected with all server nodes
|
|
max_try_times = int(os.environ.get("DGL_DIST_MAX_TRY_TIMES", 1024))
|
|
for server_id, addr in server_namebook.items():
|
|
server_ip = addr[1]
|
|
server_port = addr[2]
|
|
try_times = 0
|
|
while not rpc.connect_receiver(server_ip, server_port, server_id):
|
|
try_times += 1
|
|
if try_times % 200 == 0:
|
|
print(
|
|
"Client is trying to connect server receiver: {}:{}".format(
|
|
server_ip, server_port
|
|
)
|
|
)
|
|
if try_times >= max_try_times:
|
|
raise rpc.DistConnectError(
|
|
max_try_times, server_ip, server_port
|
|
)
|
|
time.sleep(3)
|
|
if not rpc.connect_receiver_finalize(max_try_times):
|
|
raise rpc.DistConnectError(max_try_times)
|
|
# Get local usable IP address and port
|
|
ip_addr = get_local_usable_addr(server_ip)
|
|
client_ip, client_port = ip_addr.split(":")
|
|
# Register client on server
|
|
register_req = rpc.ClientRegisterRequest(ip_addr)
|
|
for server_id in range(num_servers):
|
|
rpc.send_request(server_id, register_req)
|
|
# wait server connect back
|
|
rpc.wait_for_senders(client_ip, client_port, num_servers)
|
|
print(
|
|
"Client [{}] waits on {}:{}".format(os.getpid(), client_ip, client_port)
|
|
)
|
|
# recv client ID from server
|
|
res = rpc.recv_response()
|
|
rpc.set_rank(res.client_id)
|
|
print(
|
|
"Machine (%d) group (%d) client (%d) connect to server successfuly!"
|
|
% (machine_id, group_id, rpc.get_rank())
|
|
)
|
|
# get total number of client
|
|
get_client_num_req = rpc.GetNumberClientsRequest(rpc.get_rank())
|
|
rpc.send_request(0, get_client_num_req)
|
|
res = rpc.recv_response()
|
|
rpc.set_num_client(res.num_client)
|
|
from .dist_context import exit_client, set_initialized
|
|
|
|
atexit.register(exit_client)
|
|
set_initialized(True)
|