chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,150 @@
|
||||
"""Functions used by server."""
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
from ..base import DGLError
|
||||
from . import rpc
|
||||
from .constants import MAX_QUEUE_SIZE, SERVER_EXIT
|
||||
|
||||
|
||||
def start_server(
|
||||
server_id,
|
||||
ip_config,
|
||||
num_servers,
|
||||
num_clients,
|
||||
server_state,
|
||||
max_queue_size=MAX_QUEUE_SIZE,
|
||||
):
|
||||
"""Start DGL server, which will be shared with all the rpc services.
|
||||
|
||||
This is a blocking function -- it returns only when the server shutdown.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
server_id : int
|
||||
Current server ID (starts from 0).
|
||||
ip_config : str
|
||||
Path of IP configuration file.
|
||||
num_servers : int
|
||||
Server count on each machine.
|
||||
num_clients : int
|
||||
Total number of clients that will be connected to the server.
|
||||
Note that, we do not support dynamic connection for now. It means
|
||||
that when all the clients connect to server, no client will can be added
|
||||
to the cluster.
|
||||
server_state : ServerSate object
|
||||
Store in main data used by server.
|
||||
max_queue_size : int
|
||||
Maximal size (bytes) of server queue buffer (~20 GB on default).
|
||||
Note that the 20 GB is just an upper-bound because DGL uses zero-copy and
|
||||
it will not allocate 20GB memory at once.
|
||||
"""
|
||||
assert server_id >= 0, (
|
||||
"server_id (%d) cannot be a negative number." % server_id
|
||||
)
|
||||
assert num_servers > 0, (
|
||||
"num_servers (%d) must be a positive number." % num_servers
|
||||
)
|
||||
assert num_clients >= 0, (
|
||||
"num_client (%d) cannot be a negative number." % num_clients
|
||||
)
|
||||
assert max_queue_size > 0, (
|
||||
"queue_size (%d) cannot be a negative number." % max_queue_size
|
||||
)
|
||||
# Register signal handler.
|
||||
rpc.register_sig_handler()
|
||||
# Register some basic services
|
||||
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.set_rank(server_id)
|
||||
server_namebook = rpc.read_ip_config(ip_config, num_servers)
|
||||
machine_id = server_namebook[server_id][0]
|
||||
rpc.set_machine_id(machine_id)
|
||||
ip_addr = server_namebook[server_id][1]
|
||||
port = server_namebook[server_id][2]
|
||||
rpc.create_sender(max_queue_size)
|
||||
rpc.create_receiver(max_queue_size)
|
||||
# wait all the senders connect to server.
|
||||
# Once all the senders connect to server, server will not
|
||||
# accept new sender's connection
|
||||
print(
|
||||
"Server is waiting for connections on [{}:{}]...".format(ip_addr, port)
|
||||
)
|
||||
rpc.wait_for_senders(ip_addr, port, num_clients)
|
||||
rpc.set_num_client(num_clients)
|
||||
recv_clients = {}
|
||||
while True:
|
||||
# go through if any client group is ready for connection
|
||||
for group_id in list(recv_clients.keys()):
|
||||
ips = recv_clients[group_id]
|
||||
if len(ips) < rpc.get_num_client():
|
||||
continue
|
||||
|
||||
del recv_clients[group_id]
|
||||
# a new client group is ready
|
||||
ips.sort()
|
||||
client_namebook = dict(enumerate(ips))
|
||||
time.sleep(3) # wait for clients' receivers ready
|
||||
max_try_times = int(os.environ.get("DGL_DIST_MAX_TRY_TIMES", 120))
|
||||
for client_id, addr in client_namebook.items():
|
||||
client_ip, client_port = addr.split(":")
|
||||
try_times = 0
|
||||
while not rpc.connect_receiver(
|
||||
client_ip, client_port, client_id, group_id
|
||||
):
|
||||
try_times += 1
|
||||
if try_times % 200 == 0:
|
||||
print(
|
||||
"Server~{} is trying to connect client receiver: {}:{}".format(
|
||||
server_id, client_ip, client_port
|
||||
)
|
||||
)
|
||||
if try_times >= max_try_times:
|
||||
raise rpc.DistConnectError(
|
||||
max_try_times, client_ip, client_port
|
||||
)
|
||||
time.sleep(1)
|
||||
if not rpc.connect_receiver_finalize(max_try_times):
|
||||
raise rpc.DistConnectError(max_try_times)
|
||||
if rpc.get_rank() == 0: # server_0 send all the IDs
|
||||
for client_id, _ in client_namebook.items():
|
||||
register_res = rpc.ClientRegisterResponse(client_id)
|
||||
rpc.send_response(client_id, register_res, group_id)
|
||||
# receive incomming client requests
|
||||
timeout = 60 * 1000 # in milliseconds
|
||||
req, client_id, group_id = rpc.recv_request(timeout)
|
||||
if req is None:
|
||||
continue
|
||||
if isinstance(req, rpc.ClientRegisterRequest):
|
||||
if group_id not in recv_clients:
|
||||
recv_clients[group_id] = []
|
||||
recv_clients[group_id].append(req.ip_addr)
|
||||
continue
|
||||
|
||||
res = req.process_request(server_state)
|
||||
if res is not None:
|
||||
if isinstance(res, list):
|
||||
for response in res:
|
||||
target_id, res_data = response
|
||||
rpc.send_response(target_id, res_data, group_id)
|
||||
elif isinstance(res, str):
|
||||
if res == SERVER_EXIT:
|
||||
print("Server is exiting...")
|
||||
return
|
||||
else:
|
||||
raise DGLError("Unexpected response: {}".format(res))
|
||||
else:
|
||||
rpc.send_response(client_id, res, group_id)
|
||||
Reference in New Issue
Block a user