chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,228 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user