wehub-resource-sync
1605 lines
49 KiB
Python
1605 lines
49 KiB
Python
"""Define distributed kvstore"""
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
|
|
from .. import backend as F, utils
|
|
from .._ffi.ndarray import empty_shared_mem
|
|
|
|
from . import rpc
|
|
from .graph_partition_book import EdgePartitionPolicy, NodePartitionPolicy
|
|
from .standalone_kvstore import KVClient as SA_KVClient
|
|
|
|
############################ Register KVStore Requsts and Responses ###############################
|
|
|
|
KVSTORE_PULL = 901231
|
|
|
|
|
|
class PullResponse(rpc.Response):
|
|
"""Send the sliced data tensor back to the client.
|
|
|
|
Parameters
|
|
----------
|
|
server_id : int
|
|
ID of current server
|
|
data_tensor : tensor
|
|
sliced data tensor
|
|
"""
|
|
|
|
def __init__(self, server_id, data_tensor):
|
|
self.server_id = server_id
|
|
self.data_tensor = data_tensor
|
|
|
|
def __getstate__(self):
|
|
return self.server_id, self.data_tensor
|
|
|
|
def __setstate__(self, state):
|
|
self.server_id, self.data_tensor = state
|
|
|
|
|
|
class PullRequest(rpc.Request):
|
|
"""Send ID tensor to server and get target data tensor as response.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
a vector storing the data ID
|
|
"""
|
|
|
|
def __init__(self, name, id_tensor):
|
|
self.name = name
|
|
self.id_tensor = id_tensor
|
|
|
|
def __getstate__(self):
|
|
return self.name, self.id_tensor
|
|
|
|
def __setstate__(self, state):
|
|
self.name, self.id_tensor = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
if self.name not in kv_store.part_policy:
|
|
raise RuntimeError(
|
|
"KVServer cannot find partition policy with name: %s"
|
|
% self.name
|
|
)
|
|
if self.name not in kv_store.data_store:
|
|
raise RuntimeError(
|
|
"KVServer Cannot find data tensor with name: %s" % self.name
|
|
)
|
|
local_id = kv_store.part_policy[self.name].to_local(self.id_tensor)
|
|
data = kv_store.pull_handlers[self.name](
|
|
kv_store.data_store, self.name, local_id
|
|
)
|
|
res = PullResponse(kv_store.server_id, data)
|
|
return res
|
|
|
|
|
|
KVSTORE_PUSH = 901232
|
|
|
|
|
|
class PushRequest(rpc.Request):
|
|
"""Send ID tensor and data tensor to server and update kvstore's data.
|
|
|
|
This request has no response.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
a vector storing the data ID
|
|
data_tensor : tensor
|
|
a tensor with the same row size of data ID
|
|
"""
|
|
|
|
def __init__(self, name, id_tensor, data_tensor):
|
|
self.name = name
|
|
self.id_tensor = id_tensor
|
|
self.data_tensor = data_tensor
|
|
|
|
def __getstate__(self):
|
|
return self.name, self.id_tensor, self.data_tensor
|
|
|
|
def __setstate__(self, state):
|
|
self.name, self.id_tensor, self.data_tensor = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
if self.name not in kv_store.part_policy:
|
|
raise RuntimeError(
|
|
"KVServer cannot find partition policy with name: %s"
|
|
% self.name
|
|
)
|
|
if self.name not in kv_store.data_store:
|
|
raise RuntimeError(
|
|
"KVServer Cannot find data tensor with name: %s" % self.name
|
|
)
|
|
local_id = kv_store.part_policy[self.name].to_local(self.id_tensor)
|
|
kv_store.push_handlers[self.name](
|
|
kv_store.data_store, self.name, local_id, self.data_tensor
|
|
)
|
|
|
|
|
|
INIT_DATA = 901233
|
|
INIT_MSG = "Init"
|
|
|
|
|
|
class InitDataResponse(rpc.Response):
|
|
"""Send a confirmation response (just a short string message) of
|
|
InitDataRequest to client.
|
|
|
|
Parameters
|
|
----------
|
|
msg : string
|
|
string message
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class InitDataRequest(rpc.Request):
|
|
"""Send meta data to server and init data tensor
|
|
on server using UDF init function.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
shape : tuple
|
|
data shape
|
|
dtype : str
|
|
data type string, e.g., 'int64', 'float32', etc.
|
|
policy_str : str
|
|
partition-policy string, e.g., 'edge' or 'node'.
|
|
init_func : function
|
|
UDF init function.
|
|
"""
|
|
|
|
def __init__(self, name, shape, dtype, policy_str, init_func):
|
|
self.name = name
|
|
self.shape = shape
|
|
self.dtype = dtype
|
|
self.policy_str = policy_str
|
|
self.init_func = init_func
|
|
|
|
def __getstate__(self):
|
|
return (
|
|
self.name,
|
|
self.shape,
|
|
self.dtype,
|
|
self.policy_str,
|
|
self.init_func,
|
|
)
|
|
|
|
def __setstate__(self, state):
|
|
(
|
|
self.name,
|
|
self.shape,
|
|
self.dtype,
|
|
self.policy_str,
|
|
self.init_func,
|
|
) = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
dtype = F.data_type_dict[self.dtype]
|
|
|
|
# We should see requests from multiple clients. We need to ignore the duplicated
|
|
# reqeusts.
|
|
if self.name in kv_store.data_store:
|
|
assert tuple(F.shape(kv_store.data_store[self.name])) == tuple(
|
|
self.shape
|
|
)
|
|
assert (
|
|
F.reverse_data_type_dict[
|
|
F.dtype(kv_store.data_store[self.name])
|
|
]
|
|
== self.dtype
|
|
)
|
|
assert kv_store.part_policy[self.name].policy_str == self.policy_str
|
|
else:
|
|
if not kv_store.is_backup_server():
|
|
data_tensor = self.init_func(self.shape, dtype)
|
|
kv_store.init_data(
|
|
name=self.name,
|
|
policy_str=self.policy_str,
|
|
data_tensor=data_tensor,
|
|
)
|
|
else:
|
|
kv_store.init_data(name=self.name, policy_str=self.policy_str)
|
|
res = InitDataResponse(INIT_MSG)
|
|
return res
|
|
|
|
|
|
BARRIER = 901234
|
|
BARRIER_MSG = "Barrier"
|
|
|
|
|
|
class BarrierResponse(rpc.Response):
|
|
"""Send an confimation signal (just a short string message) of
|
|
BarrierRequest to client.
|
|
|
|
Parameters
|
|
----------
|
|
msg : string
|
|
string msg
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class BarrierRequest(rpc.Request):
|
|
"""Send a barrier signal (just a short string message) to server.
|
|
|
|
Parameters
|
|
----------
|
|
role : string
|
|
client role
|
|
"""
|
|
|
|
def __init__(self, role):
|
|
self.role = role
|
|
self.group_id = rpc.get_group_id()
|
|
|
|
def __getstate__(self):
|
|
return self.role, self.group_id
|
|
|
|
def __setstate__(self, state):
|
|
self.role, self.group_id = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
roles = server_state.roles
|
|
role = roles[self.group_id]
|
|
barrier_count = kv_store.barrier_count[self.group_id]
|
|
count = barrier_count[self.role]
|
|
barrier_count[self.role] = count + 1
|
|
if barrier_count[self.role] == len(role[self.role]):
|
|
barrier_count[self.role] = 0
|
|
res_list = []
|
|
for client_id, _ in role[self.role]:
|
|
res_list.append((client_id, BarrierResponse(BARRIER_MSG)))
|
|
return res_list
|
|
return None
|
|
|
|
|
|
REGISTER_PULL = 901235
|
|
REGISTER_PULL_MSG = "Register_Pull"
|
|
|
|
|
|
class RegisterPullHandlerResponse(rpc.Response):
|
|
"""Send a confirmation signal (just a short string message) of
|
|
RegisterPullHandler to client.
|
|
|
|
Parameters
|
|
----------
|
|
msg : string
|
|
string message
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class RegisterPullHandlerRequest(rpc.Request):
|
|
"""Send an UDF and register Pull handler on server.
|
|
|
|
Parameters
|
|
----------
|
|
pull_func : func
|
|
UDF pull handler
|
|
"""
|
|
|
|
def __init__(self, name, pull_func):
|
|
self.name = name
|
|
self.pull_func = pull_func
|
|
|
|
def __getstate__(self):
|
|
return self.name, self.pull_func
|
|
|
|
def __setstate__(self, state):
|
|
self.name, self.pull_func = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
kv_store.pull_handlers[self.name] = self.pull_func
|
|
res = RegisterPullHandlerResponse(REGISTER_PULL_MSG)
|
|
return res
|
|
|
|
|
|
REGISTER_PUSH = 901236
|
|
REGISTER_PUSH_MSG = "Register_Push"
|
|
|
|
|
|
class RegisterPushHandlerResponse(rpc.Response):
|
|
"""Send a confirmation signal (just a short string message) of
|
|
RegisterPushHandler to client.
|
|
|
|
Parameters
|
|
----------
|
|
msg : string
|
|
string message
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class RegisterPushHandlerRequest(rpc.Request):
|
|
"""Send an UDF to register Push handler on server.
|
|
|
|
Parameters
|
|
----------
|
|
push_func : func
|
|
UDF push handler
|
|
"""
|
|
|
|
def __init__(self, name, push_func):
|
|
self.name = name
|
|
self.push_func = push_func
|
|
|
|
def __getstate__(self):
|
|
return self.name, self.push_func
|
|
|
|
def __setstate__(self, state):
|
|
self.name, self.push_func = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
kv_store.push_handlers[self.name] = self.push_func
|
|
res = RegisterPushHandlerResponse(REGISTER_PUSH_MSG)
|
|
return res
|
|
|
|
|
|
GET_SHARED = 901237
|
|
GET_SHARED_MSG = "Get_Shared"
|
|
|
|
|
|
class GetSharedDataResponse(rpc.Response):
|
|
"""Send meta data of shared-memory tensor to client.
|
|
|
|
Parameters
|
|
----------
|
|
meta : dict
|
|
a dict of meta, e.g.,
|
|
|
|
{'data_0' : (shape, dtype, policy_str),
|
|
'data_1' : (shape, dtype, policy_str)}
|
|
"""
|
|
|
|
def __init__(self, meta):
|
|
self.meta = meta
|
|
|
|
def __getstate__(self):
|
|
return self.meta
|
|
|
|
def __setstate__(self, state):
|
|
self.meta = state
|
|
|
|
|
|
class GetSharedDataRequest(rpc.Request):
|
|
"""Send a signal (just a short string message) to get the
|
|
meta data of shared-tensor from server.
|
|
|
|
Parameters
|
|
----------
|
|
msg : string
|
|
string message
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
def process_request(self, server_state):
|
|
assert self.msg == GET_SHARED_MSG
|
|
meta = {}
|
|
kv_store = server_state.kv_store
|
|
for name, data in kv_store.data_store.items():
|
|
meta[name] = (
|
|
F.shape(data),
|
|
F.reverse_data_type_dict[F.dtype(data)],
|
|
kv_store.part_policy[name].policy_str,
|
|
)
|
|
res = GetSharedDataResponse(meta)
|
|
return res
|
|
|
|
|
|
GET_PART_SHAPE = 901238
|
|
|
|
|
|
class GetPartShapeResponse(rpc.Response):
|
|
"""Send the partitioned data shape back to client.
|
|
|
|
Parameters
|
|
----------
|
|
shape : tuple
|
|
shape of tensor
|
|
"""
|
|
|
|
def __init__(self, shape):
|
|
self.shape = shape
|
|
|
|
def __getstate__(self):
|
|
return self.shape
|
|
|
|
def __setstate__(self, state):
|
|
# When the shape has only one dimension, state is an integer.
|
|
if isinstance(state, int):
|
|
self.shape = (state,)
|
|
else:
|
|
self.shape = state
|
|
|
|
|
|
class GetPartShapeRequest(rpc.Request):
|
|
"""Send data name to get the partitioned data shape from server.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
"""
|
|
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def __getstate__(self):
|
|
return self.name
|
|
|
|
def __setstate__(self, state):
|
|
self.name = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
if self.name not in kv_store.data_store:
|
|
raise RuntimeError(
|
|
"KVServer Cannot find data tensor with name: %s" % self.name
|
|
)
|
|
data_shape = F.shape(kv_store.data_store[self.name])
|
|
res = GetPartShapeResponse(data_shape)
|
|
return res
|
|
|
|
|
|
SEND_META_TO_BACKUP = 901239
|
|
SEND_META_TO_BACKUP_MSG = "Send_Meta_TO_Backup"
|
|
|
|
|
|
class SendMetaToBackupResponse(rpc.Response):
|
|
"""Send a confirmation signal (just a short string message)
|
|
of SendMetaToBackupRequest to client.
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class SendMetaToBackupRequest(rpc.Request):
|
|
"""Send meta data to backup server and backup server
|
|
will use this meta data to read shared-memory tensor.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
dtype : str
|
|
data type string
|
|
shape : tuple of int
|
|
data shape
|
|
policy_str : str
|
|
partition-policy string, e.g., 'edge' or 'node'.
|
|
pull_handler : callable
|
|
The callback function when data is pulled from kvstore.
|
|
push_handler : callable
|
|
The callback function when data is pushed to kvstore.
|
|
"""
|
|
|
|
def __init__(
|
|
self, name, dtype, shape, policy_str, pull_handler, push_handler
|
|
):
|
|
self.name = name
|
|
self.dtype = dtype
|
|
self.shape = shape
|
|
self.policy_str = policy_str
|
|
self.pull_handler = pull_handler
|
|
self.push_handler = push_handler
|
|
|
|
def __getstate__(self):
|
|
return (
|
|
self.name,
|
|
self.dtype,
|
|
self.shape,
|
|
self.policy_str,
|
|
self.pull_handler,
|
|
self.push_handler,
|
|
)
|
|
|
|
def __setstate__(self, state):
|
|
(
|
|
self.name,
|
|
self.dtype,
|
|
self.shape,
|
|
self.policy_str,
|
|
self.pull_handler,
|
|
self.push_handler,
|
|
) = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
assert kv_store.is_backup_server()
|
|
if self.name not in kv_store.data_store:
|
|
shared_data = empty_shared_mem(
|
|
self.name + "-kvdata-", False, self.shape, self.dtype
|
|
)
|
|
dlpack = shared_data.to_dlpack()
|
|
kv_store.data_store[self.name] = F.zerocopy_from_dlpack(dlpack)
|
|
kv_store.part_policy[self.name] = kv_store.find_policy(
|
|
self.policy_str
|
|
)
|
|
kv_store.pull_handlers[self.name] = self.pull_handler
|
|
kv_store.push_handlers[self.name] = self.push_handler
|
|
else:
|
|
assert tuple(F.shape(kv_store.data_store[self.name])) == tuple(
|
|
self.shape
|
|
)
|
|
assert (
|
|
F.reverse_data_type_dict[
|
|
F.dtype(kv_store.data_store[self.name])
|
|
]
|
|
== self.dtype
|
|
)
|
|
assert kv_store.part_policy[self.name].policy_str == self.policy_str
|
|
assert kv_store.pull_handlers[self.name] == self.pull_handler
|
|
assert kv_store.push_handlers[self.name] == self.push_handler
|
|
res = SendMetaToBackupResponse(SEND_META_TO_BACKUP_MSG)
|
|
return res
|
|
|
|
|
|
DELETE_DATA = 901240
|
|
DELETE_MSG = "Delete_Data"
|
|
|
|
|
|
class DeleteDataResponse(rpc.Response):
|
|
"""Send a confirmation signal (just a short string message)
|
|
of DeleteDataRequest to client.
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class DeleteDataRequest(rpc.Request):
|
|
"""Send message to server to delete data tensor
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
"""
|
|
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def __getstate__(self):
|
|
return self.name
|
|
|
|
def __setstate__(self, state):
|
|
self.name = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
if self.name in kv_store.data_store:
|
|
del kv_store.data_store[self.name]
|
|
del kv_store.part_policy[self.name]
|
|
del kv_store.push_handlers[self.name]
|
|
del kv_store.pull_handlers[self.name]
|
|
res = DeleteDataResponse(DELETE_MSG)
|
|
return res
|
|
|
|
|
|
COUNT_LOCAL_NONZERO = 901241
|
|
|
|
|
|
class CountLocalNonzeroResponse(rpc.Response):
|
|
"""Send the number of nonzero value in local data"""
|
|
|
|
def __init__(self, num_local_nonzero):
|
|
self.num_local_nonzero = num_local_nonzero
|
|
|
|
def __getstate__(self):
|
|
return self.num_local_nonzero
|
|
|
|
def __setstate__(self, state):
|
|
self.num_local_nonzero = state
|
|
|
|
|
|
class CountLocalNonzeroRequest(rpc.Request):
|
|
"""Send data name to server to count local nonzero value
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
"""
|
|
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def __getstate__(self):
|
|
return self.name
|
|
|
|
def __setstate__(self, state):
|
|
self.name = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
num_local_nonzero = kv_store.count_local_nonzero(self.name)
|
|
res = CountLocalNonzeroResponse(num_local_nonzero)
|
|
return res
|
|
|
|
|
|
############################ KVServer ###############################
|
|
|
|
|
|
def default_push_handler(target, name, id_tensor, data_tensor):
|
|
"""Default handler for PUSH message.
|
|
|
|
On default, _push_handler perform scatter_row() operation for the tensor.
|
|
|
|
Parameters
|
|
----------
|
|
target : tensor
|
|
target tensor
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
a vector storing the ID list.
|
|
data_tensor : tensor
|
|
a tensor with the same row size of id
|
|
"""
|
|
# TODO(chao): support Tensorflow backend
|
|
target[name][id_tensor] = data_tensor
|
|
|
|
|
|
def default_pull_handler(target, name, id_tensor):
|
|
"""Default handler for PULL operation.
|
|
|
|
On default, _pull_handler perform gather_row() operation for the tensor.
|
|
|
|
Parameters
|
|
----------
|
|
target : tensor
|
|
target tensor
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
a vector storing the ID list.
|
|
|
|
Return
|
|
------
|
|
tensor
|
|
a tensor with the same row size of ID.
|
|
"""
|
|
# TODO(chao): support Tensorflow backend
|
|
return target[name][id_tensor]
|
|
|
|
|
|
class KVServer(object):
|
|
"""KVServer is a lightweight key-value store service for DGL distributed training.
|
|
|
|
In practice, developers can use KVServer to hold large-scale graph features or
|
|
graph embeddings across machines in a distributed setting. KVServer depends on DGL rpc
|
|
infrastructure thats support backup servers, which means we can lunach many KVServers
|
|
on the same machine for load-balancing.
|
|
|
|
DO NOT use KVServer in mult-threads because this behavior is not defined. For now, KVServer
|
|
can only support CPU-to-CPU communication. We may support GPU-communication in the future.
|
|
|
|
Parameters
|
|
----------
|
|
server_id : int
|
|
ID of current server (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 KVClients that will be connected to the KVServer.
|
|
"""
|
|
|
|
def __init__(self, server_id, ip_config, num_servers, num_clients):
|
|
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 os.path.exists(ip_config), "Cannot open file: %s" % ip_config
|
|
assert num_clients >= 0, (
|
|
"num_clients (%d) cannot be a negative number." % num_clients
|
|
)
|
|
# Register services on server
|
|
rpc.register_service(KVSTORE_PULL, PullRequest, PullResponse)
|
|
rpc.register_service(KVSTORE_PUSH, PushRequest, None)
|
|
rpc.register_service(INIT_DATA, InitDataRequest, InitDataResponse)
|
|
rpc.register_service(BARRIER, BarrierRequest, BarrierResponse)
|
|
rpc.register_service(
|
|
REGISTER_PUSH,
|
|
RegisterPushHandlerRequest,
|
|
RegisterPushHandlerResponse,
|
|
)
|
|
rpc.register_service(
|
|
REGISTER_PULL,
|
|
RegisterPullHandlerRequest,
|
|
RegisterPullHandlerResponse,
|
|
)
|
|
rpc.register_service(
|
|
GET_SHARED, GetSharedDataRequest, GetSharedDataResponse
|
|
)
|
|
rpc.register_service(
|
|
GET_PART_SHAPE, GetPartShapeRequest, GetPartShapeResponse
|
|
)
|
|
rpc.register_service(
|
|
SEND_META_TO_BACKUP,
|
|
SendMetaToBackupRequest,
|
|
SendMetaToBackupResponse,
|
|
)
|
|
rpc.register_service(DELETE_DATA, DeleteDataRequest, DeleteDataResponse)
|
|
rpc.register_service(
|
|
COUNT_LOCAL_NONZERO,
|
|
CountLocalNonzeroRequest,
|
|
CountLocalNonzeroResponse,
|
|
)
|
|
# Store the tensor data with specified data name
|
|
self._data_store = {}
|
|
# Store original tensor data names when instantiating DistGraphServer
|
|
self._orig_data = set()
|
|
# Store the partition information with specified data name
|
|
self._policy_set = set()
|
|
self._part_policy = {}
|
|
# Basic information
|
|
self._server_id = server_id
|
|
self._server_namebook = rpc.read_ip_config(ip_config, num_servers)
|
|
assert (
|
|
server_id in self._server_namebook
|
|
), "Trying to start server {}, but there are {} servers in the config file".format(
|
|
server_id, len(self._server_namebook)
|
|
)
|
|
self._machine_id = self._server_namebook[server_id][0]
|
|
self._group_count = self._server_namebook[server_id][3]
|
|
# We assume partition_id is equal to machine_id
|
|
self._part_id = self._machine_id
|
|
self._num_clients = num_clients
|
|
self._barrier_count = {}
|
|
# push and pull handler
|
|
self._push_handlers = {}
|
|
self._pull_handlers = {}
|
|
|
|
@property
|
|
def server_id(self):
|
|
"""Get server ID"""
|
|
return self._server_id
|
|
|
|
@property
|
|
def barrier_count(self):
|
|
"""Get barrier count"""
|
|
return self._barrier_count
|
|
|
|
@barrier_count.setter
|
|
def barrier_count(self, count):
|
|
"""Set barrier count"""
|
|
self._barrier_count = count
|
|
|
|
@property
|
|
def num_clients(self):
|
|
"""Get number of clients"""
|
|
return self._num_clients
|
|
|
|
@property
|
|
def data_store(self):
|
|
"""Get data store"""
|
|
return self._data_store
|
|
|
|
@property
|
|
def orig_data(self):
|
|
"""Get original data"""
|
|
return self._orig_data
|
|
|
|
@property
|
|
def part_policy(self):
|
|
"""Get part policy"""
|
|
return self._part_policy
|
|
|
|
@property
|
|
def part_id(self):
|
|
"""Get part ID"""
|
|
return self._part_id
|
|
|
|
@property
|
|
def push_handlers(self):
|
|
"""Get push handler"""
|
|
return self._push_handlers
|
|
|
|
@property
|
|
def pull_handlers(self):
|
|
"""Get pull handler"""
|
|
return self._pull_handlers
|
|
|
|
def is_backup_server(self):
|
|
"""Return True if current server is a backup server."""
|
|
if self._server_id % self._group_count == 0:
|
|
return False
|
|
return True
|
|
|
|
def add_part_policy(self, policy):
|
|
"""Add partition policy to kvserver.
|
|
|
|
Parameters
|
|
----------
|
|
policy : PartitionPolicy
|
|
Store the partition information
|
|
"""
|
|
self._policy_set.add(policy)
|
|
|
|
def init_data(self, name, policy_str, data_tensor=None):
|
|
"""Init data tensor on kvserver.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
policy_str : str
|
|
partition-policy string, e.g., 'edge' or 'node'.
|
|
data_tensor : tensor
|
|
If the data_tensor is None, KVServer will
|
|
read shared-memory when client invoking get_shared_data().
|
|
"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
if name in self._data_store:
|
|
raise RuntimeError("Data %s has already exists!" % name)
|
|
self._part_policy[name] = self.find_policy(policy_str)
|
|
if data_tensor is not None: # Create shared-tensor
|
|
data_type = F.reverse_data_type_dict[F.dtype(data_tensor)]
|
|
shared_data = empty_shared_mem(
|
|
name + "-kvdata-", True, data_tensor.shape, data_type
|
|
)
|
|
dlpack = shared_data.to_dlpack()
|
|
self._data_store[name] = F.zerocopy_from_dlpack(dlpack)
|
|
rpc.copy_data_to_shared_memory(self._data_store[name], data_tensor)
|
|
assert (
|
|
self._part_policy[name].get_part_size() == data_tensor.shape[0]
|
|
), "kvserver expect partition {} for {} has {} rows, but gets {} rows".format(
|
|
self._part_policy[name].part_id,
|
|
policy_str,
|
|
self._part_policy[name].get_part_size(),
|
|
data_tensor.shape[0],
|
|
)
|
|
self._pull_handlers[name] = default_pull_handler
|
|
self._push_handlers[name] = default_push_handler
|
|
|
|
def find_policy(self, policy_str):
|
|
"""Find a partition policy from existing policy set
|
|
|
|
Parameters
|
|
----------
|
|
policy_str : str
|
|
partition-policy string, e.g., 'edge' or 'node'.
|
|
"""
|
|
for policy in self._policy_set:
|
|
if policy_str == policy.policy_str:
|
|
return policy
|
|
raise RuntimeError(
|
|
"Cannot find policy_str: %s from kvserver." % policy_str
|
|
)
|
|
|
|
def count_local_nonzero(self, name):
|
|
"""Count nonzero in local data
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
the number of nonzero in local data.
|
|
"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
if name not in self._data_store:
|
|
raise RuntimeError("Data %s has not be created!" % name)
|
|
return F.count_nonzero(self._data_store[name])
|
|
|
|
|
|
############################ KVClient ###############################
|
|
|
|
|
|
class KVClient(object):
|
|
"""KVClient is used to push/pull data to/from KVServer. If the
|
|
target kvclient and kvserver are in the same machine, they can
|
|
communicate with each other using local shared-memory
|
|
automatically, instead of going through the tcp/ip RPC.
|
|
|
|
DO NOT use KVClient in multi-threads because this behavior is
|
|
not defined. For now, KVClient can only support CPU-to-CPU communication.
|
|
We may support GPU-communication in the future.
|
|
|
|
Parameters
|
|
----------
|
|
ip_config : str
|
|
Path of IP configuration file.
|
|
num_servers : int
|
|
Server count on each machine.
|
|
role : str
|
|
We can set different role for kvstore.
|
|
"""
|
|
|
|
def __init__(self, ip_config, num_servers, role="default"):
|
|
assert (
|
|
rpc.get_rank() != -1
|
|
), "Please invoke rpc.connect_to_server() before creating KVClient."
|
|
assert os.path.exists(ip_config), "Cannot open file: %s" % ip_config
|
|
assert num_servers > 0, (
|
|
"num_servers (%d) must be a positive number." % num_servers
|
|
)
|
|
# Register services on client
|
|
rpc.register_service(KVSTORE_PULL, PullRequest, PullResponse)
|
|
rpc.register_service(KVSTORE_PUSH, PushRequest, None)
|
|
rpc.register_service(INIT_DATA, InitDataRequest, InitDataResponse)
|
|
rpc.register_service(BARRIER, BarrierRequest, BarrierResponse)
|
|
rpc.register_service(
|
|
REGISTER_PUSH,
|
|
RegisterPushHandlerRequest,
|
|
RegisterPushHandlerResponse,
|
|
)
|
|
rpc.register_service(
|
|
REGISTER_PULL,
|
|
RegisterPullHandlerRequest,
|
|
RegisterPullHandlerResponse,
|
|
)
|
|
rpc.register_service(
|
|
GET_SHARED, GetSharedDataRequest, GetSharedDataResponse
|
|
)
|
|
rpc.register_service(
|
|
GET_PART_SHAPE, GetPartShapeRequest, GetPartShapeResponse
|
|
)
|
|
rpc.register_service(
|
|
SEND_META_TO_BACKUP,
|
|
SendMetaToBackupRequest,
|
|
SendMetaToBackupResponse,
|
|
)
|
|
rpc.register_service(DELETE_DATA, DeleteDataRequest, DeleteDataResponse)
|
|
rpc.register_service(
|
|
COUNT_LOCAL_NONZERO,
|
|
CountLocalNonzeroRequest,
|
|
CountLocalNonzeroResponse,
|
|
)
|
|
# Store the tensor data with specified data name
|
|
self._data_store = {}
|
|
# Store the partition information with specified data name
|
|
self._part_policy = {}
|
|
# This stores all unique partition policies in the kvstore. The key is the policy name.
|
|
self._all_possible_part_policy = {}
|
|
# Store the full data shape across kvserver
|
|
self._full_data_shape = {}
|
|
# Store all the data name
|
|
self._data_name_list = set()
|
|
# Store all graph data name
|
|
self._gdata_name_list = set()
|
|
# Basic information
|
|
self._server_namebook = rpc.read_ip_config(ip_config, num_servers)
|
|
self._server_count = len(self._server_namebook)
|
|
self._group_count = self._server_namebook[0][3]
|
|
self._machine_count = int(self._server_count / self._group_count)
|
|
self._client_id = rpc.get_rank()
|
|
self._machine_id = rpc.get_machine_id()
|
|
self._part_id = self._machine_id
|
|
self._main_server_id = self._machine_id * self._group_count
|
|
|
|
# push and pull handler
|
|
self._pull_handlers = {}
|
|
self._push_handlers = {}
|
|
# register role on server-0
|
|
self._role = role
|
|
|
|
@property
|
|
def all_possible_part_policy(self):
|
|
"""Get all possible partition policies"""
|
|
return self._all_possible_part_policy
|
|
|
|
@property
|
|
def client_id(self):
|
|
"""Get client ID"""
|
|
return self._client_id
|
|
|
|
@property
|
|
def role(self):
|
|
"""Get client role"""
|
|
return self._role
|
|
|
|
@property
|
|
def machine_id(self):
|
|
"""Get machine ID"""
|
|
return self._machine_id
|
|
|
|
@property
|
|
def num_servers(self):
|
|
"""Get the number of servers"""
|
|
return self._server_count
|
|
|
|
@property
|
|
def group_count(self):
|
|
"""Get the number of groups --num_servers"""
|
|
return self._group_count
|
|
|
|
def barrier(self):
|
|
"""Barrier for all client nodes.
|
|
|
|
This API will be blocked untill all the clients invoke this API.
|
|
"""
|
|
request = BarrierRequest(self._role)
|
|
rpc.send_request(0, request)
|
|
response = rpc.recv_response()
|
|
assert response.msg == BARRIER_MSG
|
|
|
|
def register_push_handler(self, name, func):
|
|
"""Register UDF push function.
|
|
|
|
This UDF is triggered for every push. The signature of the UDF is
|
|
|
|
```
|
|
def push_handler(data_store, name, local_offset, data)
|
|
```
|
|
|
|
``data_store`` is a dict that contains all tensors in the kvstore. ``name`` is the name
|
|
of the tensor where new data is pushed to. ``local_offset`` is the offset where new
|
|
data should be written in the tensor in the local partition. ``data`` is the new data
|
|
to be written.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the tensor
|
|
func : callable
|
|
The function to be called.
|
|
"""
|
|
self.barrier()
|
|
request = RegisterPushHandlerRequest(name, func)
|
|
# send request to all the server nodes
|
|
for server_id in range(self._server_count):
|
|
rpc.send_request(server_id, request)
|
|
# recv response from all the server nodes
|
|
for _ in range(self._server_count):
|
|
response = rpc.recv_response()
|
|
assert response.msg == REGISTER_PUSH_MSG
|
|
self._push_handlers[name] = func
|
|
self.barrier()
|
|
|
|
def register_pull_handler(self, name, func):
|
|
"""Register UDF pull function.
|
|
|
|
This UDF is triggered for every pull. The signature of the UDF is
|
|
|
|
```
|
|
def pull_handler(data_store, name, local_offset)
|
|
```
|
|
|
|
``data_store`` is a dict that contains all tensors in the kvstore. ``name`` is the name
|
|
of the tensor where new data is pushed to. ``local_offset`` is the offset where new
|
|
data should be written in the tensor in the local partition.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the tensor
|
|
func : callable
|
|
The function to be called.
|
|
"""
|
|
self.barrier()
|
|
request = RegisterPullHandlerRequest(name, func)
|
|
# send request to all the server nodes
|
|
for server_id in range(self._server_count):
|
|
rpc.send_request(server_id, request)
|
|
# recv response from all the server nodes
|
|
for _ in range(self._server_count):
|
|
response = rpc.recv_response()
|
|
assert response.msg == REGISTER_PULL_MSG
|
|
self._pull_handlers[name] = func
|
|
self.barrier()
|
|
|
|
def init_data(
|
|
self, name, shape, dtype, part_policy, init_func, is_gdata=True
|
|
):
|
|
"""Send message to kvserver to initialize new data tensor and mapping this
|
|
data from server side to client side.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
shape : list or tuple of int
|
|
data shape
|
|
dtype : dtype
|
|
data type
|
|
part_policy : PartitionPolicy
|
|
partition policy.
|
|
init_func : func
|
|
UDF init function
|
|
is_gdata : bool
|
|
Whether the created tensor is a ndata/edata or not.
|
|
"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
assert len(shape) > 0, "shape cannot be empty"
|
|
assert name not in self._data_name_list, (
|
|
"data name: %s already exists." % name
|
|
)
|
|
self.barrier()
|
|
shape = list(shape)
|
|
|
|
# Send request to the servers to initialize data.
|
|
# The servers may handle the duplicated initializations.
|
|
part_shape = shape.copy()
|
|
part_shape[0] = part_policy.get_part_size()
|
|
request = InitDataRequest(
|
|
name,
|
|
tuple(part_shape),
|
|
F.reverse_data_type_dict[dtype],
|
|
part_policy.policy_str,
|
|
init_func,
|
|
)
|
|
# The request is sent to the servers in one group, which are on the same machine.
|
|
for n in range(self._group_count):
|
|
server_id = part_policy.part_id * self._group_count + n
|
|
rpc.send_request(server_id, request)
|
|
for _ in range(self._group_count):
|
|
response = rpc.recv_response()
|
|
assert response.msg == INIT_MSG
|
|
|
|
self.barrier()
|
|
# Create local shared-data
|
|
local_shape = shape.copy()
|
|
local_shape[0] = part_policy.get_part_size()
|
|
if name in self._part_policy:
|
|
raise RuntimeError("Policy %s has already exists!" % name)
|
|
if name in self._data_store:
|
|
raise RuntimeError("Data %s has already exists!" % name)
|
|
if name in self._full_data_shape:
|
|
raise RuntimeError("Data shape %s has already exists!" % name)
|
|
self._part_policy[name] = part_policy
|
|
self._all_possible_part_policy[part_policy.policy_str] = part_policy
|
|
shared_data = empty_shared_mem(
|
|
name + "-kvdata-",
|
|
False,
|
|
local_shape,
|
|
F.reverse_data_type_dict[dtype],
|
|
)
|
|
dlpack = shared_data.to_dlpack()
|
|
self._data_store[name] = F.zerocopy_from_dlpack(dlpack)
|
|
self._data_name_list.add(name)
|
|
if is_gdata:
|
|
self._gdata_name_list.add(name)
|
|
self._full_data_shape[name] = tuple(shape)
|
|
self._pull_handlers[name] = default_pull_handler
|
|
self._push_handlers[name] = default_push_handler
|
|
|
|
# Now we need to tell the backup server the new tensor.
|
|
request = SendMetaToBackupRequest(
|
|
name,
|
|
F.reverse_data_type_dict[dtype],
|
|
part_shape,
|
|
part_policy.policy_str,
|
|
self._pull_handlers[name],
|
|
self._push_handlers[name],
|
|
)
|
|
# send request to all the backup server nodes
|
|
for i in range(self._group_count - 1):
|
|
server_id = self._machine_id * self._group_count + i + 1
|
|
rpc.send_request(server_id, request)
|
|
# recv response from all the backup server nodes
|
|
for _ in range(self._group_count - 1):
|
|
response = rpc.recv_response()
|
|
assert response.msg == SEND_META_TO_BACKUP_MSG
|
|
self.barrier()
|
|
|
|
def delete_data(self, name):
|
|
"""Send message to kvserver to delete tensor and clear the meta data
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
assert name in self._data_name_list, "data name: %s not exists." % name
|
|
self.barrier()
|
|
part_policy = self._part_policy[name]
|
|
|
|
# send request to every server nodes
|
|
request = DeleteDataRequest(name)
|
|
for n in range(self._group_count):
|
|
server_id = part_policy.part_id * self._group_count + n
|
|
rpc.send_request(server_id, request)
|
|
for _ in range(self._group_count):
|
|
response = rpc.recv_response()
|
|
assert response.msg == DELETE_MSG
|
|
|
|
self.barrier()
|
|
self._data_name_list.remove(name)
|
|
if name in self._gdata_name_list:
|
|
self._gdata_name_list.remove(name)
|
|
# TODO(chao) : remove the delete log print
|
|
del self._data_store[name]
|
|
del self._full_data_shape[name]
|
|
del self._part_policy[name]
|
|
del self._pull_handlers[name]
|
|
del self._push_handlers[name]
|
|
self.barrier()
|
|
|
|
def map_shared_data(self, partition_book):
|
|
"""Mapping shared-memory tensor from server to client.
|
|
|
|
Parameters
|
|
----------
|
|
partition_book : GraphPartitionBook
|
|
Store the partition information
|
|
"""
|
|
# Get all partition policies
|
|
for ntype in partition_book.ntypes:
|
|
policy = NodePartitionPolicy(partition_book, ntype)
|
|
self._all_possible_part_policy[policy.policy_str] = policy
|
|
for etype in partition_book.canonical_etypes:
|
|
policy = EdgePartitionPolicy(partition_book, etype)
|
|
self._all_possible_part_policy[policy.policy_str] = policy
|
|
|
|
# Get shared data from server side
|
|
self.barrier()
|
|
request = GetSharedDataRequest(GET_SHARED_MSG)
|
|
rpc.send_request(self._main_server_id, request)
|
|
response = rpc.recv_response()
|
|
for name, meta in response.meta.items():
|
|
if name not in self._data_name_list:
|
|
shape, dtype, policy_str = meta
|
|
assert policy_str in self._all_possible_part_policy
|
|
shared_data = empty_shared_mem(
|
|
name + "-kvdata-", False, shape, dtype
|
|
)
|
|
dlpack = shared_data.to_dlpack()
|
|
self._data_store[name] = F.zerocopy_from_dlpack(dlpack)
|
|
self._part_policy[name] = self._all_possible_part_policy[
|
|
policy_str
|
|
]
|
|
self._pull_handlers[name] = default_pull_handler
|
|
self._push_handlers[name] = default_push_handler
|
|
# Get full data shape across servers
|
|
for name, meta in response.meta.items():
|
|
if name not in self._data_name_list:
|
|
shape, _, _ = meta
|
|
data_shape = list(shape)
|
|
data_shape[0] = 0
|
|
request = GetPartShapeRequest(name)
|
|
# send request to all main server nodes
|
|
for machine_id in range(self._machine_count):
|
|
server_id = machine_id * self._group_count
|
|
rpc.send_request(server_id, request)
|
|
# recv response from all the main server nodes
|
|
for _ in range(self._machine_count):
|
|
res = rpc.recv_response()
|
|
data_shape[0] += res.shape[0]
|
|
self._full_data_shape[name] = tuple(data_shape)
|
|
# Send meta data to backup servers
|
|
for name, meta in response.meta.items():
|
|
shape, dtype, policy_str = meta
|
|
request = SendMetaToBackupRequest(
|
|
name,
|
|
dtype,
|
|
shape,
|
|
policy_str,
|
|
self._pull_handlers[name],
|
|
self._push_handlers[name],
|
|
)
|
|
# send request to all the backup server nodes
|
|
for i in range(self._group_count - 1):
|
|
server_id = self._machine_id * self._group_count + i + 1
|
|
rpc.send_request(server_id, request)
|
|
# recv response from all the backup server nodes
|
|
for _ in range(self._group_count - 1):
|
|
response = rpc.recv_response()
|
|
assert response.msg == SEND_META_TO_BACKUP_MSG
|
|
self._data_name_list.add(name)
|
|
# map_shared_data happens only at DistGraph initialization
|
|
# TODO(xiangsx): We assume there is no non-graph data initialized at this time
|
|
self._gdata_name_list.add(name)
|
|
self.barrier()
|
|
|
|
def gdata_name_list(self):
|
|
"""Get all the graph data name"""
|
|
return list(self._gdata_name_list)
|
|
|
|
def data_name_list(self):
|
|
"""Get all the data name"""
|
|
return list(self._data_name_list)
|
|
|
|
def get_data_meta(self, name):
|
|
"""Get meta data (data_type, data_shape, partition_policy)"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
data_type = F.dtype(self._data_store[name])
|
|
data_shape = self._full_data_shape[name]
|
|
part_policy = self._part_policy[name]
|
|
return (data_type, data_shape, part_policy)
|
|
|
|
def get_partid(self, name, id_tensor):
|
|
"""
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
a vector storing the global data ID
|
|
"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
id_tensor = utils.toindex(id_tensor)
|
|
id_tensor = id_tensor.tousertensor()
|
|
assert F.ndim(id_tensor) == 1, "ID must be a vector."
|
|
# partition data
|
|
machine_id = self._part_policy[name].to_partid(id_tensor)
|
|
|
|
return machine_id
|
|
|
|
def push(self, name, id_tensor, data_tensor):
|
|
"""Push data to KVServer.
|
|
|
|
Note that, the push() is an non-blocking operation that will return immediately.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
a vector storing the global data ID
|
|
data_tensor : tensor
|
|
a tensor with the same row size of data ID
|
|
"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
id_tensor = utils.toindex(id_tensor)
|
|
id_tensor = id_tensor.tousertensor()
|
|
assert F.ndim(id_tensor) == 1, "ID must be a vector."
|
|
assert (
|
|
F.shape(id_tensor)[0] == F.shape(data_tensor)[0]
|
|
), "The data must has the same row size with ID."
|
|
# partition data
|
|
machine_id = self._part_policy[name].to_partid(id_tensor)
|
|
# sort index by machine id
|
|
sorted_id = F.tensor(np.argsort(F.asnumpy(machine_id)))
|
|
id_tensor = id_tensor[sorted_id]
|
|
data_tensor = data_tensor[sorted_id]
|
|
machine, count = np.unique(F.asnumpy(machine_id), return_counts=True)
|
|
# push data to server by order
|
|
start = 0
|
|
local_id = None
|
|
local_data = None
|
|
for idx, machine_idx in enumerate(machine):
|
|
end = start + count[idx]
|
|
if start == end: # No data for target machine
|
|
continue
|
|
partial_id = id_tensor[start:end]
|
|
partial_data = data_tensor[start:end]
|
|
if machine_idx == self._machine_id: # local push
|
|
# Note that DO NOT push local data right now because we can overlap
|
|
# communication-local_push here
|
|
local_id = self._part_policy[name].to_local(partial_id)
|
|
local_data = partial_data
|
|
else: # push data to remote server
|
|
request = PushRequest(name, partial_id, partial_data)
|
|
rpc.send_request_to_machine(machine_idx, request)
|
|
start += count[idx]
|
|
if local_id is not None: # local push
|
|
self._push_handlers[name](
|
|
self._data_store, name, local_id, local_data
|
|
)
|
|
|
|
def pull(self, name, id_tensor):
|
|
"""Pull message from KVServer.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
id_tensor : tensor
|
|
a vector storing the ID list
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
a data tensor with the same row size of id_tensor.
|
|
"""
|
|
assert len(name) > 0, "name cannot be empty."
|
|
id_tensor = utils.toindex(id_tensor)
|
|
id_tensor = id_tensor.tousertensor()
|
|
assert F.ndim(id_tensor) == 1, "ID must be a vector."
|
|
if self._pull_handlers[name] is default_pull_handler: # Use fast-pull
|
|
part_id = self._part_policy[name].to_partid(id_tensor)
|
|
return rpc.fast_pull(
|
|
name,
|
|
id_tensor,
|
|
part_id,
|
|
KVSTORE_PULL,
|
|
self._machine_count,
|
|
self._group_count,
|
|
self._machine_id,
|
|
self._client_id,
|
|
self._data_store[name],
|
|
self._part_policy[name],
|
|
)
|
|
else:
|
|
# partition data
|
|
machine_id = self._part_policy[name].to_partid(id_tensor)
|
|
# sort index by machine id
|
|
sorted_id = F.tensor(np.argsort(F.asnumpy(machine_id)))
|
|
back_sorted_id = F.tensor(np.argsort(F.asnumpy(sorted_id)))
|
|
id_tensor = id_tensor[sorted_id]
|
|
machine, count = np.unique(
|
|
F.asnumpy(machine_id), return_counts=True
|
|
)
|
|
# pull data from server by order
|
|
start = 0
|
|
pull_count = 0
|
|
local_id = None
|
|
for idx, machine_idx in enumerate(machine):
|
|
end = start + count[idx]
|
|
if start == end: # No data for target machine
|
|
continue
|
|
partial_id = id_tensor[start:end]
|
|
if machine_idx == self._machine_id: # local pull
|
|
# Note that DO NOT pull local data right now because we can overlap
|
|
# communication-local_pull here
|
|
local_id = self._part_policy[name].to_local(partial_id)
|
|
else: # pull data from remote server
|
|
request = PullRequest(name, partial_id)
|
|
rpc.send_request_to_machine(machine_idx, request)
|
|
pull_count += 1
|
|
start += count[idx]
|
|
# recv response
|
|
response_list = []
|
|
if local_id is not None: # local pull
|
|
local_data = self._pull_handlers[name](
|
|
self._data_store, name, local_id
|
|
)
|
|
server_id = self._main_server_id
|
|
local_response = PullResponse(server_id, local_data)
|
|
response_list.append(local_response)
|
|
# wait response from remote server nodes
|
|
for _ in range(pull_count):
|
|
remote_response = rpc.recv_response()
|
|
response_list.append(remote_response)
|
|
# sort response by server_id and concat tensor
|
|
response_list.sort(key=self._take_id)
|
|
data_tensor = F.cat(
|
|
seq=[response.data_tensor for response in response_list], dim=0
|
|
)
|
|
return data_tensor[
|
|
back_sorted_id
|
|
] # return data with original index order
|
|
|
|
def union(self, operand1_name, operand2_name, output_name):
|
|
"""Compute the union of two mask arrays in the KVStore."""
|
|
# Each trainer computes its own result from its local storage.
|
|
self._data_store[output_name][:] = (
|
|
self._data_store[operand1_name] | self._data_store[operand2_name]
|
|
)
|
|
|
|
def _take_id(self, elem):
|
|
"""Used by sort response list"""
|
|
return elem.server_id
|
|
|
|
def count_nonzero(self, name):
|
|
"""Count nonzero value by pull request from KVServers.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
data name
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
the number of nonzero in this data.
|
|
"""
|
|
total = 0
|
|
pull_count = 0
|
|
for machine_id in range(self._machine_count):
|
|
if machine_id == self._machine_id:
|
|
local_id = F.tensor(
|
|
np.arange(
|
|
self._part_policy[name].get_part_size(), dtype=np.int64
|
|
)
|
|
)
|
|
total += F.count_nonzero(self._data_store[name][local_id])
|
|
else:
|
|
request = CountLocalNonzeroRequest(name)
|
|
rpc.send_request_to_machine(machine_id, request)
|
|
pull_count += 1
|
|
for _ in range(pull_count):
|
|
res = rpc.recv_response()
|
|
total += res.num_local_nonzero
|
|
return total
|
|
|
|
@property
|
|
def data_store(self):
|
|
"""Return the local partition of the data storage.
|
|
|
|
Returns
|
|
-------
|
|
dict[str, Tensor]
|
|
The tensor storages of the local partition.
|
|
"""
|
|
return self._data_store
|
|
|
|
|
|
KVCLIENT = None
|
|
|
|
|
|
def init_kvstore(ip_config, num_servers, role):
|
|
"""initialize KVStore"""
|
|
global KVCLIENT
|
|
if KVCLIENT is None:
|
|
if os.environ.get("DGL_DIST_MODE", "standalone") == "standalone":
|
|
KVCLIENT = SA_KVClient()
|
|
else:
|
|
KVCLIENT = KVClient(ip_config, num_servers, role)
|
|
|
|
|
|
def close_kvstore():
|
|
"""Close the current KVClient"""
|
|
global KVCLIENT
|
|
KVCLIENT = None
|
|
|
|
|
|
def get_kvstore():
|
|
"""get the KVClient"""
|
|
return KVCLIENT
|