wehub-resource-sync
1435 lines
37 KiB
Python
1435 lines
37 KiB
Python
"""RPC components. They are typically functions or utilities used by both
|
|
server and clients."""
|
|
import abc
|
|
import os
|
|
import pickle
|
|
import random
|
|
|
|
import numpy as np
|
|
|
|
from .. import backend as F
|
|
from .._ffi.function import _init_api
|
|
from .._ffi.object import ObjectBase, register_object
|
|
from ..base import DGLError
|
|
from .constants import SERVER_EXIT
|
|
|
|
__all__ = [
|
|
"set_rank",
|
|
"get_rank",
|
|
"Request",
|
|
"Response",
|
|
"register_service",
|
|
"create_sender",
|
|
"create_receiver",
|
|
"finalize_sender",
|
|
"finalize_receiver",
|
|
"wait_for_senders",
|
|
"connect_receiver",
|
|
"read_ip_config",
|
|
"get_group_id",
|
|
"get_num_machines",
|
|
"set_num_machines",
|
|
"get_machine_id",
|
|
"set_machine_id",
|
|
"send_request",
|
|
"recv_request",
|
|
"send_response",
|
|
"recv_response",
|
|
"remote_call",
|
|
"send_request_to_machine",
|
|
"remote_call_to_machine",
|
|
"fast_pull",
|
|
"DistConnectError",
|
|
"get_num_client",
|
|
"set_num_client",
|
|
"client_barrier",
|
|
"copy_data_to_shared_memory",
|
|
]
|
|
|
|
REQUEST_CLASS_TO_SERVICE_ID = {}
|
|
RESPONSE_CLASS_TO_SERVICE_ID = {}
|
|
SERVICE_ID_TO_PROPERTY = {}
|
|
|
|
DEFUALT_PORT = 30050
|
|
|
|
|
|
def read_ip_config(filename, num_servers):
|
|
"""Read network configuration information of server from file.
|
|
|
|
For exampple, the following TXT shows a 4-machine configuration:
|
|
|
|
172.31.40.143
|
|
172.31.36.140
|
|
172.31.47.147
|
|
172.31.30.180
|
|
|
|
Users can also set user-specified port for this network configuration. For example:
|
|
|
|
172.31.40.143 20090
|
|
172.31.36.140 20090
|
|
172.31.47.147 20090
|
|
172.31.30.180 20090
|
|
|
|
Note that, DGL supports multiple backup servers that shares data with each others
|
|
on the same machine via shared-memory tensor. The num_servers should be >= 1. For example,
|
|
if we set num_servers to 5, it means that we have 1 main server and 4 backup servers on
|
|
current machine.
|
|
|
|
Parameters
|
|
----------
|
|
filename : str
|
|
Path of IP configuration file.
|
|
|
|
num_servers : int
|
|
Server count on each machine.
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
server namebook.
|
|
The key is server_id (int)
|
|
The value is [machine_id, ip, port, num_servers] ([int, str, int, int])
|
|
|
|
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]}
|
|
"""
|
|
assert len(filename) > 0, "filename cannot be empty."
|
|
assert num_servers > 0, (
|
|
"num_servers (%d) must be a positive number." % num_servers
|
|
)
|
|
server_namebook = {}
|
|
try:
|
|
server_id = 0
|
|
machine_id = 0
|
|
lines = [line.rstrip("\n") for line in open(filename)]
|
|
for line in lines:
|
|
result = line.split()
|
|
if len(result) == 2:
|
|
port = int(result[1])
|
|
elif len(result) == 1:
|
|
port = DEFUALT_PORT
|
|
else:
|
|
raise RuntimeError("length of result can only be 1 or 2.")
|
|
ip_addr = result[0]
|
|
for s_count in range(num_servers):
|
|
server_namebook[server_id] = [
|
|
machine_id,
|
|
ip_addr,
|
|
port + s_count,
|
|
num_servers,
|
|
]
|
|
server_id += 1
|
|
machine_id += 1
|
|
except RuntimeError:
|
|
print("Error: data format on each line should be: [ip] [port]")
|
|
return server_namebook
|
|
|
|
|
|
def reset():
|
|
"""Reset the rpc context"""
|
|
_CAPI_DGLRPCReset()
|
|
|
|
|
|
def create_sender(max_queue_size):
|
|
"""Create rpc sender of this process.
|
|
|
|
Parameters
|
|
----------
|
|
max_queue_size : int
|
|
Maximal size (bytes) of network queue buffer.
|
|
"""
|
|
max_thread_count = int(os.getenv("DGL_SOCKET_MAX_THREAD_COUNT", "0"))
|
|
_CAPI_DGLRPCCreateSender(int(max_queue_size), max_thread_count)
|
|
|
|
|
|
def create_receiver(max_queue_size):
|
|
"""Create rpc receiver of this process.
|
|
|
|
Parameters
|
|
----------
|
|
max_queue_size : int
|
|
Maximal size (bytes) of network queue buffer.
|
|
"""
|
|
max_thread_count = int(os.getenv("DGL_SOCKET_MAX_THREAD_COUNT", "0"))
|
|
_CAPI_DGLRPCCreateReceiver(int(max_queue_size), max_thread_count)
|
|
|
|
|
|
def finalize_sender():
|
|
"""Finalize rpc sender of this process."""
|
|
_CAPI_DGLRPCFinalizeSender()
|
|
|
|
|
|
def finalize_receiver():
|
|
"""Finalize rpc receiver of this process."""
|
|
_CAPI_DGLRPCFinalizeReceiver()
|
|
|
|
|
|
def wait_for_senders(ip_addr, port, num_senders):
|
|
"""Wait all of the senders' connections.
|
|
|
|
This api will be blocked until all the senders connect to the receiver.
|
|
|
|
Parameters
|
|
----------
|
|
ip_addr : str
|
|
receiver's IP address, e,g, '192.168.8.12'
|
|
port : int
|
|
receiver's port
|
|
num_senders : int
|
|
total number of senders
|
|
"""
|
|
_CAPI_DGLRPCWaitForSenders(ip_addr, int(port), int(num_senders))
|
|
|
|
|
|
def connect_receiver(ip_addr, port, recv_id, group_id=-1):
|
|
"""Connect to target receiver
|
|
|
|
Parameters
|
|
----------
|
|
ip_addr : str
|
|
receiver's IP address, e,g, '192.168.8.12'
|
|
port : int
|
|
receiver's listening port
|
|
recv_id : int
|
|
receiver's ID
|
|
"""
|
|
target_id = (
|
|
recv_id if group_id == -1 else register_client(recv_id, group_id)
|
|
)
|
|
if target_id < 0:
|
|
raise DGLError("Invalid target id: {}".format(target_id))
|
|
return _CAPI_DGLRPCConnectReceiver(ip_addr, int(port), int(target_id))
|
|
|
|
|
|
def connect_receiver_finalize(max_try_times):
|
|
"""Finalize the action to connect to receivers. Make sure that either all connections are
|
|
successfully established or connection fails.
|
|
|
|
When "socket" network backend is in use, the function issues actual requests to receiver
|
|
sockets to establish connections.
|
|
|
|
Parameters
|
|
----------
|
|
max_try_times : int
|
|
maximum try times
|
|
"""
|
|
return _CAPI_DGLRPCConnectReceiverFinalize(max_try_times)
|
|
|
|
|
|
def set_rank(rank):
|
|
"""Set the rank of this process.
|
|
|
|
If the process is a client, this is equal to client ID. Otherwise, the process
|
|
is a server and this is equal to server ID.
|
|
|
|
Parameters
|
|
----------
|
|
rank : int
|
|
Rank value
|
|
"""
|
|
_CAPI_DGLRPCSetRank(int(rank))
|
|
|
|
|
|
def get_rank():
|
|
"""Get the rank of this process.
|
|
|
|
If the process is a client, this is equal to client ID. Otherwise, the process
|
|
is a server and this is equal to server ID.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
Rank value
|
|
"""
|
|
return _CAPI_DGLRPCGetRank()
|
|
|
|
|
|
def set_machine_id(machine_id):
|
|
"""Set current machine ID
|
|
|
|
Parameters
|
|
----------
|
|
machine_id : int
|
|
Current machine ID
|
|
"""
|
|
_CAPI_DGLRPCSetMachineID(int(machine_id))
|
|
|
|
|
|
def get_machine_id():
|
|
"""Get current machine ID
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
machine ID
|
|
"""
|
|
return _CAPI_DGLRPCGetMachineID()
|
|
|
|
|
|
def set_num_machines(num_machines):
|
|
"""Set number of machine
|
|
|
|
Parameters
|
|
----------
|
|
num_machines : int
|
|
Number of machine
|
|
"""
|
|
_CAPI_DGLRPCSetNumMachines(int(num_machines))
|
|
|
|
|
|
def get_num_machines():
|
|
"""Get number of machines
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
number of machines
|
|
"""
|
|
return _CAPI_DGLRPCGetNumMachines()
|
|
|
|
|
|
def set_num_server(num_server):
|
|
"""Set the total number of server."""
|
|
_CAPI_DGLRPCSetNumServer(int(num_server))
|
|
|
|
|
|
def get_num_server():
|
|
"""Get the total number of server."""
|
|
return _CAPI_DGLRPCGetNumServer()
|
|
|
|
|
|
def set_num_client(num_client):
|
|
"""Set the total number of client."""
|
|
_CAPI_DGLRPCSetNumClient(int(num_client))
|
|
|
|
|
|
def get_num_client():
|
|
"""Get the total number of client."""
|
|
return _CAPI_DGLRPCGetNumClient()
|
|
|
|
|
|
def set_num_server_per_machine(num_server):
|
|
"""Set the total number of server per machine"""
|
|
_CAPI_DGLRPCSetNumServerPerMachine(num_server)
|
|
|
|
|
|
def get_num_server_per_machine():
|
|
"""Get the total number of server per machine"""
|
|
return _CAPI_DGLRPCGetNumServerPerMachine()
|
|
|
|
|
|
def incr_msg_seq():
|
|
"""Increment the message sequence number and return the old one.
|
|
|
|
Returns
|
|
-------
|
|
long
|
|
Message sequence number
|
|
"""
|
|
return _CAPI_DGLRPCIncrMsgSeq()
|
|
|
|
|
|
def get_msg_seq():
|
|
"""Get the current message sequence number.
|
|
|
|
Returns
|
|
-------
|
|
long
|
|
Message sequence number
|
|
"""
|
|
return _CAPI_DGLRPCGetMsgSeq()
|
|
|
|
|
|
def set_msg_seq(msg_seq):
|
|
"""Set the current message sequence number.
|
|
|
|
Parameters
|
|
----------
|
|
msg_seq : int
|
|
sequence number of current rpc message.
|
|
"""
|
|
_CAPI_DGLRPCSetMsgSeq(int(msg_seq))
|
|
|
|
|
|
def register_service(service_id, req_cls, res_cls=None):
|
|
"""Register a service to RPC.
|
|
|
|
Parameter
|
|
---------
|
|
service_id : int
|
|
Service ID.
|
|
req_cls : class
|
|
Request class.
|
|
res_cls : class, optional
|
|
Response class. If none, the service has no response.
|
|
"""
|
|
REQUEST_CLASS_TO_SERVICE_ID[req_cls] = service_id
|
|
if res_cls is not None:
|
|
RESPONSE_CLASS_TO_SERVICE_ID[res_cls] = service_id
|
|
SERVICE_ID_TO_PROPERTY[service_id] = (req_cls, res_cls)
|
|
|
|
|
|
def get_service_property(service_id):
|
|
"""Get service property.
|
|
|
|
Parameters
|
|
----------
|
|
service_id : int
|
|
Service ID.
|
|
|
|
Returns
|
|
-------
|
|
(class, class)
|
|
(Request class, Response class)
|
|
"""
|
|
return SERVICE_ID_TO_PROPERTY[service_id]
|
|
|
|
|
|
class Request:
|
|
"""Base request class"""
|
|
|
|
@abc.abstractmethod
|
|
def __getstate__(self):
|
|
"""Get serializable states.
|
|
|
|
Must be inherited by subclasses. For array members, return them as
|
|
individual return values (i.e., do not put them in containers like
|
|
dictionary or list).
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def __setstate__(self, state):
|
|
"""Construct the request object from serialized states.
|
|
|
|
Must be inherited by subclasses.
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def process_request(self, server_state):
|
|
"""Server-side function to process the request.
|
|
|
|
Must be inherited by subclasses.
|
|
|
|
Parameters
|
|
----------
|
|
server_state : ServerState
|
|
Server state data.
|
|
|
|
Returns
|
|
-------
|
|
Response
|
|
Response of this request or None if no response.
|
|
"""
|
|
|
|
@property
|
|
def service_id(self):
|
|
"""Get service ID."""
|
|
cls = self.__class__
|
|
sid = REQUEST_CLASS_TO_SERVICE_ID.get(cls, None)
|
|
if sid is None:
|
|
raise DGLError(
|
|
"Request class {} has not been registered as a service.".format(
|
|
cls
|
|
)
|
|
)
|
|
return sid
|
|
|
|
|
|
class Response:
|
|
"""Base response class"""
|
|
|
|
@abc.abstractmethod
|
|
def __getstate__(self):
|
|
"""Get serializable states.
|
|
|
|
Must be inherited by subclasses. For array members, return them as
|
|
individual return values (i.e., do not put them in containers like
|
|
dictionary or list).
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def __setstate__(self, state):
|
|
"""Construct the response object from serialized states.
|
|
|
|
Must be inherited by subclasses.
|
|
"""
|
|
|
|
@property
|
|
def service_id(self):
|
|
"""Get service ID."""
|
|
cls = self.__class__
|
|
sid = RESPONSE_CLASS_TO_SERVICE_ID.get(cls, None)
|
|
if sid is None:
|
|
raise DGLError(
|
|
"Response class {} has not been registered as a service.".format(
|
|
cls
|
|
)
|
|
)
|
|
return sid
|
|
|
|
|
|
def serialize_to_payload(serializable):
|
|
"""Serialize an object to payloads.
|
|
|
|
The object must have implemented the __getstate__ function.
|
|
|
|
Parameters
|
|
----------
|
|
serializable : object
|
|
Any serializable object.
|
|
|
|
Returns
|
|
-------
|
|
bytearray
|
|
Serialized payload buffer.
|
|
list[Tensor]
|
|
A list of tensor payloads.
|
|
"""
|
|
state = serializable.__getstate__()
|
|
if not isinstance(state, tuple):
|
|
state = (state,)
|
|
nonarray_pos = []
|
|
nonarray_state = []
|
|
array_state = []
|
|
for i, arr_state in enumerate(state):
|
|
if F.is_tensor(arr_state):
|
|
array_state.append(arr_state)
|
|
else:
|
|
nonarray_state.append(arr_state)
|
|
nonarray_pos.append(i)
|
|
data = bytearray(pickle.dumps((nonarray_pos, nonarray_state)))
|
|
return data, array_state
|
|
|
|
|
|
class PlaceHolder:
|
|
"""PlaceHolder object for deserialization"""
|
|
|
|
|
|
_PLACEHOLDER = PlaceHolder()
|
|
|
|
|
|
def deserialize_from_payload(cls, data, tensors):
|
|
"""Deserialize and reconstruct the object from payload.
|
|
|
|
The object must have implemented the __setstate__ function.
|
|
|
|
Parameters
|
|
----------
|
|
cls : class
|
|
The object class.
|
|
data : bytearray
|
|
Serialized data buffer.
|
|
tensors : list[Tensor]
|
|
A list of tensor payloads.
|
|
|
|
Returns
|
|
-------
|
|
object
|
|
De-serialized object of class cls.
|
|
"""
|
|
pos, nonarray_state = pickle.loads(data)
|
|
# Use _PLACEHOLDER to distinguish with other deserizliaed elements
|
|
state = [_PLACEHOLDER] * (len(nonarray_state) + len(tensors))
|
|
for i, no_state in zip(pos, nonarray_state):
|
|
state[i] = no_state
|
|
if len(tensors) != 0:
|
|
j = 0
|
|
state_len = len(state)
|
|
for i in range(state_len):
|
|
if state[i] is _PLACEHOLDER:
|
|
state[i] = tensors[j]
|
|
j += 1
|
|
if len(state) == 1:
|
|
state = state[0]
|
|
else:
|
|
state = tuple(state)
|
|
obj = cls.__new__(cls)
|
|
obj.__setstate__(state)
|
|
return obj
|
|
|
|
|
|
@register_object("rpc.RPCMessage")
|
|
class RPCMessage(ObjectBase):
|
|
"""Serialized RPC message that can be sent to remote processes.
|
|
|
|
This class can be used as argument or return value for C API.
|
|
|
|
Attributes
|
|
----------
|
|
service_id : int
|
|
The remote service ID the message wishes to invoke.
|
|
msg_seq : int
|
|
Sequence number of this message.
|
|
client_id : int
|
|
The client ID.
|
|
server_id : int
|
|
The server ID.
|
|
data : bytearray
|
|
Payload buffer carried by this request.
|
|
tensors : list[tensor]
|
|
Extra payloads in the form of tensors.
|
|
group_id : int
|
|
The group ID
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
service_id,
|
|
msg_seq,
|
|
client_id,
|
|
server_id,
|
|
data,
|
|
tensors,
|
|
group_id=0,
|
|
):
|
|
self.__init_handle_by_constructor__(
|
|
_CAPI_DGLRPCCreateRPCMessage,
|
|
int(service_id),
|
|
int(msg_seq),
|
|
int(client_id),
|
|
int(server_id),
|
|
data,
|
|
[F.zerocopy_to_dgl_ndarray(tsor) for tsor in tensors],
|
|
int(group_id),
|
|
)
|
|
|
|
@property
|
|
def service_id(self):
|
|
"""Get service ID."""
|
|
return _CAPI_DGLRPCMessageGetServiceId(self)
|
|
|
|
@property
|
|
def msg_seq(self):
|
|
"""Get message sequence number."""
|
|
return _CAPI_DGLRPCMessageGetMsgSeq(self)
|
|
|
|
@property
|
|
def client_id(self):
|
|
"""Get client ID."""
|
|
return _CAPI_DGLRPCMessageGetClientId(self)
|
|
|
|
@property
|
|
def server_id(self):
|
|
"""Get server ID."""
|
|
return _CAPI_DGLRPCMessageGetServerId(self)
|
|
|
|
@property
|
|
def data(self):
|
|
"""Get payload buffer."""
|
|
return _CAPI_DGLRPCMessageGetData(self)
|
|
|
|
@property
|
|
def tensors(self):
|
|
"""Get tensor payloads."""
|
|
rst = _CAPI_DGLRPCMessageGetTensors(self)
|
|
return [F.zerocopy_from_dgl_ndarray(tsor) for tsor in rst]
|
|
|
|
@property
|
|
def group_id(self):
|
|
"""Get group ID."""
|
|
return _CAPI_DGLRPCMessageGetGroupId(self)
|
|
|
|
|
|
def send_request(target, request):
|
|
"""Send one request to the target server.
|
|
|
|
Serialize the given request object to an :class:`RPCMessage` and send it
|
|
out.
|
|
|
|
The operation is non-blocking -- it does not guarantee the payloads have
|
|
reached the target or even have left the sender process. However,
|
|
all the payloads (i.e., data and arrays) can be safely freed after this
|
|
function returns.
|
|
|
|
Parameters
|
|
----------
|
|
target : int
|
|
ID of target server.
|
|
request : Request
|
|
The request to send.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
service_id = request.service_id
|
|
msg_seq = incr_msg_seq()
|
|
client_id = get_rank()
|
|
server_id = target
|
|
data, tensors = serialize_to_payload(request)
|
|
msg = RPCMessage(
|
|
service_id,
|
|
msg_seq,
|
|
client_id,
|
|
server_id,
|
|
data,
|
|
tensors,
|
|
group_id=get_group_id(),
|
|
)
|
|
send_rpc_message(msg, server_id)
|
|
|
|
|
|
def send_request_to_machine(target, request):
|
|
"""Send one request to the target machine, which will randomly
|
|
select a server node to process this request.
|
|
|
|
The operation is non-blocking -- it does not guarantee the payloads have
|
|
reached the target or even have left the sender process. However,
|
|
all the payloads (i.e., data and arrays) can be safely freed after this
|
|
function returns.
|
|
|
|
Parameters
|
|
----------
|
|
target : int
|
|
ID of target machine.
|
|
request : Request
|
|
The request to send.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
service_id = request.service_id
|
|
msg_seq = incr_msg_seq()
|
|
client_id = get_rank()
|
|
server_id = random.randint(
|
|
target * get_num_server_per_machine(),
|
|
(target + 1) * get_num_server_per_machine() - 1,
|
|
)
|
|
data, tensors = serialize_to_payload(request)
|
|
msg = RPCMessage(
|
|
service_id, msg_seq, client_id, server_id, data, tensors, get_group_id()
|
|
)
|
|
send_rpc_message(msg, server_id)
|
|
|
|
|
|
def send_response(target, response, group_id):
|
|
"""Send one response to the target client.
|
|
|
|
Serialize the given response object to an :class:`RPCMessage` and send it
|
|
out.
|
|
|
|
The operation is non-blocking -- it does not guarantee the payloads have
|
|
reached the target or even have left the sender process. However,
|
|
all the payloads (i.e., data and arrays) can be safely freed after this
|
|
function returns.
|
|
|
|
Parameters
|
|
----------
|
|
target : int
|
|
ID of target client.
|
|
response : Response
|
|
The response to send.
|
|
group_id : int
|
|
Group ID of target client.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
service_id = response.service_id
|
|
msg_seq = get_msg_seq()
|
|
client_id = target
|
|
server_id = get_rank()
|
|
data, tensors = serialize_to_payload(response)
|
|
msg = RPCMessage(
|
|
service_id, msg_seq, client_id, server_id, data, tensors, group_id
|
|
)
|
|
send_rpc_message(msg, get_client(client_id, group_id))
|
|
|
|
|
|
def recv_request(timeout=0):
|
|
"""Receive one request.
|
|
|
|
Receive one :class:`RPCMessage` and de-serialize it into a proper Request object.
|
|
|
|
The operation is blocking -- it returns when it receives any message
|
|
or it times out.
|
|
|
|
Parameters
|
|
----------
|
|
timeout : int, optional
|
|
The timeout value in milliseconds. If zero, wait indefinitely.
|
|
|
|
Returns
|
|
-------
|
|
req : request
|
|
One request received from the target, or None if it times out.
|
|
client_id : int
|
|
Client' ID received from the target, or -1 if it times out.
|
|
group_id : int
|
|
Group' ID received from the target, or -1 if it times out.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
msg = recv_rpc_message(timeout)
|
|
if msg is None:
|
|
return None, -1, -1
|
|
set_msg_seq(msg.msg_seq)
|
|
req_cls, _ = SERVICE_ID_TO_PROPERTY[msg.service_id]
|
|
if req_cls is None:
|
|
raise DGLError(
|
|
"Got request message from service ID {}, "
|
|
"but no request class is registered.".format(msg.service_id)
|
|
)
|
|
req = deserialize_from_payload(req_cls, msg.data, msg.tensors)
|
|
if msg.server_id != get_rank():
|
|
raise DGLError(
|
|
"Got request sent to server {}, "
|
|
"different from my rank {}!".format(msg.server_id, get_rank())
|
|
)
|
|
return req, msg.client_id, msg.group_id
|
|
|
|
|
|
def recv_response(timeout=0):
|
|
"""Receive one response.
|
|
|
|
Receive one :class:`RPCMessage` and de-serialize it into a proper Response object.
|
|
|
|
The operation is blocking -- it returns when it receives any message
|
|
or it times out.
|
|
|
|
Parameters
|
|
----------
|
|
timeout : int, optional
|
|
The timeout value in milliseconds. If zero, wait indefinitely.
|
|
|
|
Returns
|
|
-------
|
|
res : Response
|
|
One response received from the target, or None if it times out.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
msg = recv_rpc_message(timeout)
|
|
if msg is None:
|
|
return None
|
|
_, res_cls = SERVICE_ID_TO_PROPERTY[msg.service_id]
|
|
if res_cls is None:
|
|
raise DGLError(
|
|
"Got response message from service ID {}, "
|
|
"but no response class is registered.".format(msg.service_id)
|
|
)
|
|
res = deserialize_from_payload(res_cls, msg.data, msg.tensors)
|
|
if msg.client_id != get_rank() and get_rank() != -1:
|
|
raise DGLError(
|
|
"Got response of request sent by client {}, "
|
|
"different from my rank {}!".format(msg.client_id, get_rank())
|
|
)
|
|
if msg.group_id != get_group_id():
|
|
raise DGLError(
|
|
"Got response of request sent by group {}, "
|
|
"different from my group {}!".format(msg.group_id, get_group_id())
|
|
)
|
|
return res
|
|
|
|
|
|
def remote_call(target_and_requests, timeout=0):
|
|
"""Invoke registered services on remote servers and collect responses.
|
|
|
|
The operation is blocking -- it returns when it receives all responses
|
|
or it times out.
|
|
|
|
If the target server state is available locally, it invokes local computation
|
|
to calculate the response.
|
|
|
|
Parameters
|
|
----------
|
|
target_and_requests : list[(int, Request)]
|
|
A list of requests and the server they should be sent to.
|
|
timeout : int, optional
|
|
The timeout value in milliseconds. If zero, wait indefinitely.
|
|
|
|
Returns
|
|
-------
|
|
list[Response]
|
|
Responses for each target-request pair. If the request does not have
|
|
response, None is placed.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
all_res = [None] * len(target_and_requests)
|
|
msgseq2pos = {}
|
|
num_res = 0
|
|
myrank = get_rank()
|
|
for pos, (target, request) in enumerate(target_and_requests):
|
|
# send request
|
|
service_id = request.service_id
|
|
msg_seq = incr_msg_seq()
|
|
client_id = get_rank()
|
|
server_id = random.randint(
|
|
target * get_num_server_per_machine(),
|
|
(target + 1) * get_num_server_per_machine() - 1,
|
|
)
|
|
data, tensors = serialize_to_payload(request)
|
|
msg = RPCMessage(
|
|
service_id,
|
|
msg_seq,
|
|
client_id,
|
|
server_id,
|
|
data,
|
|
tensors,
|
|
get_group_id(),
|
|
)
|
|
send_rpc_message(msg, server_id)
|
|
# check if has response
|
|
res_cls = get_service_property(service_id)[1]
|
|
if res_cls is not None:
|
|
num_res += 1
|
|
msgseq2pos[msg_seq] = pos
|
|
while num_res != 0:
|
|
# recv response
|
|
msg = recv_rpc_message(timeout)
|
|
if msg is None:
|
|
raise DGLError(
|
|
f"Timed out for receiving message within {timeout} milliseconds"
|
|
)
|
|
num_res -= 1
|
|
_, res_cls = SERVICE_ID_TO_PROPERTY[msg.service_id]
|
|
if res_cls is None:
|
|
raise DGLError(
|
|
"Got response message from service ID {}, "
|
|
"but no response class is registered.".format(msg.service_id)
|
|
)
|
|
res = deserialize_from_payload(res_cls, msg.data, msg.tensors)
|
|
if msg.client_id != myrank:
|
|
raise DGLError(
|
|
"Got reponse of request sent by client {}, "
|
|
"different from my rank {}!".format(msg.client_id, myrank)
|
|
)
|
|
# set response
|
|
all_res[msgseq2pos[msg.msg_seq]] = res
|
|
return all_res
|
|
|
|
|
|
def send_requests_to_machine(target_and_requests):
|
|
"""Send requests to the remote machines.
|
|
|
|
This operation isn't block. It returns immediately once it sends all requests.
|
|
|
|
Parameters
|
|
----------
|
|
target_and_requests : list[(int, Request)]
|
|
A list of requests and the machine they should be sent to.
|
|
timeout : int, optional
|
|
The timeout value in milliseconds. If zero, wait indefinitely.
|
|
|
|
Returns
|
|
-------
|
|
msgseq2pos : dict
|
|
map the message sequence number to its position in the input list.
|
|
"""
|
|
msgseq2pos = {}
|
|
for pos, (target, request) in enumerate(target_and_requests):
|
|
# send request
|
|
service_id = request.service_id
|
|
msg_seq = incr_msg_seq()
|
|
client_id = get_rank()
|
|
|
|
server_id = random.randint(
|
|
target * get_num_server_per_machine(),
|
|
(target + 1) * get_num_server_per_machine() - 1,
|
|
)
|
|
data, tensors = serialize_to_payload(request)
|
|
msg = RPCMessage(
|
|
service_id,
|
|
msg_seq,
|
|
client_id,
|
|
server_id,
|
|
data,
|
|
tensors,
|
|
get_group_id(),
|
|
)
|
|
send_rpc_message(msg, server_id)
|
|
# check if has response
|
|
res_cls = get_service_property(service_id)[1]
|
|
if res_cls is not None:
|
|
msgseq2pos[msg_seq] = pos
|
|
return msgseq2pos
|
|
|
|
|
|
def recv_responses(msgseq2pos, timeout=0):
|
|
"""Receive responses
|
|
|
|
It returns the responses in the same order as the requests. The order of requests
|
|
are stored in msgseq2pos.
|
|
|
|
The operation is blocking -- it returns when it receives all responses
|
|
or it times out.
|
|
|
|
Parameters
|
|
----------
|
|
msgseq2pos : dict
|
|
map the message sequence number to its position in the input list.
|
|
timeout : int, optional
|
|
The timeout value in milliseconds. If zero, wait indefinitely.
|
|
|
|
Returns
|
|
-------
|
|
list[Response]
|
|
Responses for each target-request pair. If the request does not have
|
|
response, None is placed.
|
|
"""
|
|
myrank = get_rank()
|
|
size = np.max(list(msgseq2pos.values())) + 1
|
|
all_res = [None] * size
|
|
num_res = len(msgseq2pos)
|
|
while num_res != 0:
|
|
# recv response
|
|
msg = recv_rpc_message(timeout)
|
|
if msg is None:
|
|
raise DGLError(
|
|
f"Timed out for receiving message within {timeout} milliseconds"
|
|
)
|
|
num_res -= 1
|
|
_, res_cls = SERVICE_ID_TO_PROPERTY[msg.service_id]
|
|
if res_cls is None:
|
|
raise DGLError(
|
|
"Got response message from service ID {}, "
|
|
"but no response class is registered.".format(msg.service_id)
|
|
)
|
|
res = deserialize_from_payload(res_cls, msg.data, msg.tensors)
|
|
if msg.client_id != myrank:
|
|
raise DGLError(
|
|
"Got reponse of request sent by client {}, "
|
|
"different from my rank {}!".format(msg.client_id, myrank)
|
|
)
|
|
# set response
|
|
all_res[msgseq2pos[msg.msg_seq]] = res
|
|
return all_res
|
|
|
|
|
|
def remote_call_to_machine(target_and_requests, timeout=0):
|
|
"""Invoke registered services on remote machine
|
|
(which will ramdom select a server to process the request) and collect responses.
|
|
|
|
The operation is blocking -- it returns when it receives all responses
|
|
or it times out.
|
|
|
|
If the target server state is available locally, it invokes local computation
|
|
to calculate the response.
|
|
|
|
Parameters
|
|
----------
|
|
target_and_requests : list[(int, Request)]
|
|
A list of requests and the machine they should be sent to.
|
|
timeout : int, optional
|
|
The timeout value in milliseconds. If zero, wait indefinitely.
|
|
|
|
Returns
|
|
-------
|
|
list[Response]
|
|
Responses for each target-request pair. If the request does not have
|
|
response, None is placed.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
msgseq2pos = send_requests_to_machine(target_and_requests)
|
|
return recv_responses(msgseq2pos, timeout)
|
|
|
|
|
|
def send_rpc_message(msg, target):
|
|
"""Send one message to the target server.
|
|
|
|
The operation is non-blocking -- it does not guarantee the payloads have
|
|
reached the target or even have left the sender process. However,
|
|
all the payloads (i.e., data and arrays) can be safely freed after this
|
|
function returns.
|
|
|
|
The data buffer in the requst will be copied to internal buffer for actual
|
|
transmission, while no memory copy for tensor payloads (a.k.a. zero-copy).
|
|
The underlying sending threads will hold references to the tensors until
|
|
the contents have been transmitted.
|
|
|
|
Parameters
|
|
----------
|
|
msg : RPCMessage
|
|
The message to send.
|
|
target : int
|
|
target ID
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
_CAPI_DGLRPCSendRPCMessage(msg, int(target))
|
|
|
|
|
|
def recv_rpc_message(timeout=0):
|
|
"""Receive one message.
|
|
|
|
The operation is blocking -- it returns when it receives any message
|
|
or it times out.
|
|
|
|
Parameters
|
|
----------
|
|
timeout : int, optional
|
|
The timeout value in milliseconds. If zero, wait indefinitely.
|
|
|
|
Returns
|
|
-------
|
|
msg : RPCMessage
|
|
One rpc message received from the target, or None if it times out.
|
|
|
|
Raises
|
|
------
|
|
ConnectionError if there is any problem with the connection.
|
|
"""
|
|
msg = _CAPI_DGLRPCCreateEmptyRPCMessage()
|
|
status = _CAPI_DGLRPCRecvRPCMessage(timeout, msg)
|
|
return msg if status == 0 else None
|
|
|
|
|
|
def client_barrier():
|
|
"""Barrier all client processes"""
|
|
req = ClientBarrierRequest()
|
|
send_request(0, req)
|
|
res = recv_response()
|
|
assert res.msg == "barrier"
|
|
|
|
|
|
def finalize_server():
|
|
"""Finalize resources of current server"""
|
|
finalize_sender()
|
|
finalize_receiver()
|
|
print("Server (%d) shutdown." % get_rank())
|
|
|
|
|
|
def fast_pull(
|
|
name,
|
|
id_tensor,
|
|
part_id,
|
|
service_id,
|
|
machine_count,
|
|
group_count,
|
|
machine_id,
|
|
client_id,
|
|
local_data,
|
|
policy,
|
|
):
|
|
"""Fast-pull api used by kvstore.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
data ID
|
|
part_id : tensor
|
|
partition ID of id_tensor
|
|
service_id : int
|
|
service_id of pull request
|
|
machine_count : int
|
|
total number of machine
|
|
group_count : int
|
|
total number of server inside machine
|
|
machine_id : int
|
|
current machine ID
|
|
client_id : int
|
|
current client ID
|
|
local_data : tensor
|
|
local data tensor
|
|
policy : PartitionPolicy
|
|
store the partition information
|
|
"""
|
|
msg_seq = incr_msg_seq()
|
|
pickle_data = bytearray(pickle.dumps(([0], [name])))
|
|
global_id = _CAPI_DGLRPCGetGlobalIDFromLocalPartition(
|
|
F.zerocopy_to_dgl_ndarray(id_tensor),
|
|
F.zerocopy_to_dgl_ndarray(part_id),
|
|
machine_id,
|
|
)
|
|
global_id = F.zerocopy_from_dgl_ndarray(global_id)
|
|
g2l_id = policy.to_local(global_id)
|
|
res_tensor = _CAPI_DGLRPCFastPull(
|
|
name,
|
|
int(machine_id),
|
|
int(machine_count),
|
|
int(group_count),
|
|
int(client_id),
|
|
int(service_id),
|
|
int(msg_seq),
|
|
pickle_data,
|
|
F.zerocopy_to_dgl_ndarray(id_tensor),
|
|
F.zerocopy_to_dgl_ndarray(part_id),
|
|
F.zerocopy_to_dgl_ndarray(g2l_id),
|
|
F.zerocopy_to_dgl_ndarray(local_data),
|
|
)
|
|
return F.zerocopy_from_dgl_ndarray(res_tensor)
|
|
|
|
|
|
def register_sig_handler():
|
|
"""Register for handling signal event."""
|
|
_CAPI_DGLRPCHandleSignal()
|
|
|
|
|
|
def copy_data_to_shared_memory(dst, source):
|
|
"""Copy tensor data to shared-memory tensor"""
|
|
F.zerocopy_to_dgl_ndarray(dst).copyfrom(F.zerocopy_to_dgl_ndarray(source))
|
|
|
|
|
|
############### Some basic services will be defined here #############
|
|
|
|
CLIENT_REGISTER = 22451
|
|
|
|
|
|
class ClientRegisterRequest(Request):
|
|
"""This request will send client's ip to server.
|
|
|
|
Parameters
|
|
----------
|
|
ip_addr : str
|
|
client's IP address
|
|
"""
|
|
|
|
def __init__(self, ip_addr):
|
|
self.ip_addr = ip_addr
|
|
|
|
def __getstate__(self):
|
|
return self.ip_addr
|
|
|
|
def __setstate__(self, state):
|
|
self.ip_addr = state
|
|
|
|
def process_request(self, server_state):
|
|
return None # do nothing
|
|
|
|
|
|
class ClientRegisterResponse(Response):
|
|
"""This response will send assigned ID to client.
|
|
|
|
Parameters
|
|
----------
|
|
ID : int
|
|
client's ID
|
|
"""
|
|
|
|
def __init__(self, client_id):
|
|
self.client_id = client_id
|
|
|
|
def __getstate__(self):
|
|
return self.client_id
|
|
|
|
def __setstate__(self, state):
|
|
self.client_id = state
|
|
|
|
|
|
SHUT_DOWN_SERVER = 22452
|
|
|
|
|
|
class ShutDownRequest(Request):
|
|
"""Client send this request to shut-down a server.
|
|
|
|
This request has no response.
|
|
|
|
Parameters
|
|
----------
|
|
client_id : int
|
|
client's ID
|
|
"""
|
|
|
|
def __init__(self, client_id, force_shutdown_server=False):
|
|
self.client_id = client_id
|
|
self.force_shutdown_server = force_shutdown_server
|
|
|
|
def __getstate__(self):
|
|
return self.client_id, self.force_shutdown_server
|
|
|
|
def __setstate__(self, state):
|
|
self.client_id, self.force_shutdown_server = state
|
|
|
|
def process_request(self, server_state):
|
|
assert self.client_id == 0
|
|
finalize_server()
|
|
return SERVER_EXIT
|
|
|
|
|
|
GET_NUM_CLIENT = 22453
|
|
|
|
|
|
class GetNumberClientsResponse(Response):
|
|
"""This reponse will send total number of clients.
|
|
|
|
Parameters
|
|
----------
|
|
num_client : int
|
|
total number of clients
|
|
"""
|
|
|
|
def __init__(self, num_client):
|
|
self.num_client = num_client
|
|
|
|
def __getstate__(self):
|
|
return self.num_client
|
|
|
|
def __setstate__(self, state):
|
|
self.num_client = state
|
|
|
|
|
|
class GetNumberClientsRequest(Request):
|
|
"""Client send this request to get the total number of client.
|
|
|
|
Parameters
|
|
----------
|
|
client_id : int
|
|
client's ID
|
|
"""
|
|
|
|
def __init__(self, client_id):
|
|
self.client_id = client_id
|
|
|
|
def __getstate__(self):
|
|
return self.client_id
|
|
|
|
def __setstate__(self, state):
|
|
self.client_id = state
|
|
|
|
def process_request(self, server_state):
|
|
res = GetNumberClientsResponse(get_num_client())
|
|
return res
|
|
|
|
|
|
CLIENT_BARRIER = 22454
|
|
|
|
|
|
class ClientBarrierResponse(Response):
|
|
"""Send the barrier confirmation to client
|
|
|
|
Parameters
|
|
----------
|
|
msg : str
|
|
string msg
|
|
"""
|
|
|
|
def __init__(self, msg="barrier"):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class ClientBarrierRequest(Request):
|
|
"""Send the barrier information to server
|
|
|
|
Parameters
|
|
----------
|
|
msg : str
|
|
string msg
|
|
"""
|
|
|
|
def __init__(self, msg="barrier"):
|
|
self.msg = msg
|
|
self.group_id = get_group_id()
|
|
|
|
def __getstate__(self):
|
|
return self.msg, self.group_id
|
|
|
|
def __setstate__(self, state):
|
|
self.msg, self.group_id = state
|
|
|
|
def process_request(self, server_state):
|
|
_CAPI_DGLRPCSetBarrierCount(
|
|
_CAPI_DGLRPCGetBarrierCount(self.group_id) + 1, self.group_id
|
|
)
|
|
if _CAPI_DGLRPCGetBarrierCount(self.group_id) == get_num_client():
|
|
_CAPI_DGLRPCSetBarrierCount(0, self.group_id)
|
|
res_list = []
|
|
for target_id in range(get_num_client()):
|
|
res_list.append((target_id, ClientBarrierResponse()))
|
|
return res_list
|
|
return None
|
|
|
|
|
|
def set_group_id(group_id):
|
|
"""Set current group ID
|
|
|
|
Parameters
|
|
----------
|
|
group_id : int
|
|
Current group ID
|
|
"""
|
|
_CAPI_DGLRPCSetGroupID(int(group_id))
|
|
|
|
|
|
def get_group_id():
|
|
"""Get current group ID
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
group ID
|
|
"""
|
|
return _CAPI_DGLRPCGetGroupID()
|
|
|
|
|
|
def register_client(client_id, group_id):
|
|
"""Register client
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
unique client ID
|
|
"""
|
|
return _CAPI_DGLRPCRegisterClient(int(client_id), int(group_id))
|
|
|
|
|
|
def get_client(client_id, group_id):
|
|
"""Get global client ID
|
|
|
|
Parameters
|
|
----------
|
|
client_id : int
|
|
client ID
|
|
group_id : int
|
|
group ID
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
global client ID
|
|
"""
|
|
return _CAPI_DGLRPCGetClient(int(client_id), int(group_id))
|
|
|
|
|
|
class DistConnectError(DGLError):
|
|
"""Exception raised for errors if fail to connect peer.
|
|
|
|
Attributes
|
|
----------
|
|
kv_store : KVServer
|
|
reference for KVServer
|
|
"""
|
|
|
|
def __init__(self, max_try_times, ip="", port=""):
|
|
peer_str = "peer[{}:{}]".format(ip, port) if ip != "" else "peer"
|
|
self.message = (
|
|
"Failed to build conncetion with {} after {} retries. "
|
|
"Please check network availability or increase max try "
|
|
"times via 'DGL_DIST_MAX_TRY_TIMES'.".format(
|
|
peer_str, max_try_times
|
|
)
|
|
)
|
|
super().__init__(self.message)
|
|
|
|
|
|
_init_api("dgl.distributed.rpc")
|