def __init__(self, config, entity_shape): super(SharedKVClient, self).__init__() self.addr = config.addr self.name = config.name self.servers = config.namebook self.client = KVClient(client_id = self.name, server_namebook = self.servers, client_addr = self.addr) time.sleep(1) self.client.connect() for entity in entity_shape: self.client.init_data(name=entity, shape=entity_shape[entity], init_type='zero')
class SharedKVClient(BaseClient): def __init__(self, config, entity_shape): super(SharedKVClient, self).__init__() self.addr = config.addr self.name = config.name self.servers = config.namebook self.client = KVClient(client_id = self.name, server_namebook = self.servers, client_addr = self.addr) time.sleep(1) self.client.connect() for entity in entity_shape: self.client.init_data(name=entity, shape=entity_shape[entity], init_type='zero') def get_entity_embedding(self, name, emb_id): return self.client.pull(name, emb_id) def get_relation_embedding(self, name, emb_id): return None def put_entity_embedding(self, name, emb_id, data): self.client.push(name, emb_id, data) def put_relation_embedding(self, name, emb_id, data): pass
def start_client():
my_client = KVClient(server_namebook=server_namebook)
my_client.connect()
name_list = my_client.get_data_name_list()
assert len(name_list) == 2
assert 'data_0' in name_list
assert 'data_1' in name_list
meta_0 = my_client.get_data_meta('data_0')
assert meta_0[0] == F.float32
assert_array_equal(meta_0[2], partition_0)
meta_1 = my_client.get_data_meta('data_1')
assert meta_1[0] == F.float32
assert_array_equal(meta_1[2], partition_1)
my_client.push(name='data_0',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor([[1., 1., 1.], [2., 2., 2.],
[3., 3., 3.]]))
res = my_client.pull(name='data_0', id_tensor=F.tensor([0, 1, 2]))
target = F.tensor([[1., 1., 1.], [2., 2., 2.], [3., 3., 3.]])
assert_array_equal(res, target)
my_client.shut_down()
def start_client(args):
"""Start client
"""
server_namebook = dgl.contrib.read_ip_config(filename=args.ip_config)
my_client = KVClient(server_namebook=server_namebook)
my_client.connect()
if my_client.get_id() % args.num_worker == 0:
my_client.set_partition_book(name='entity_embed', partition_book=partition)
else:
time.sleep(3)
my_client.set_partition_book(name='entity_embed')
my_client.print()
my_client.barrier()
print("send request...")
for i in range(4):
my_client.push(name='entity_embed', id_tensor=ID[i], data_tensor=DATA[i])
my_client.barrier()
if my_client.get_id() % args.num_worker == 0:
res = my_client.pull(name='entity_embed', id_tensor=th.tensor([0,1,2,3,4,5,6,7]))
print(res)
my_client.barrier()
my_client.push(name='entity_embed', id_tensor=ID[my_client.get_machine_id()], data_tensor=th.tensor([[0.,0.,0.],[0.,0.,0.]]))
my_client.barrier()
if my_client.get_id() % args.num_worker == 0:
res = my_client.pull(name='entity_embed', id_tensor=th.tensor([0,1,2,3,4,5,6,7]))
print(res)
my_client.shut_down()
def start_client():
my_client = KVClient(server_namebook=server_namebook)
my_client.connect()
my_client.init_data(name='data_2',
shape=(num_entries, dim_size),
dtype=F.float32,
target_name='data_0')
print("Init data from client..")
name_list = my_client.get_data_name_list()
assert len(name_list) == 6
assert 'data_0' in name_list
assert 'data_1' in name_list
assert 'data_2' in name_list
assert 'data_3' in name_list
assert 'data_4' in name_list
assert 'data_5' in name_list
meta_0 = my_client.get_data_meta('data_0')
assert meta_0[0] == F.float32
assert meta_0[1] == tuple(F.shape(data_0))
assert_array_equal(meta_0[2], partition_0)
meta_1 = my_client.get_data_meta('data_1')
assert meta_1[0] == F.float32
assert meta_1[1] == tuple(F.shape(data_1))
assert_array_equal(meta_1[2], partition_1)
meta_2 = my_client.get_data_meta('data_2')
assert meta_2[0] == F.float32
assert meta_2[1] == tuple(F.shape(data_0))
assert_array_equal(meta_2[2], partition_0)
meta_3 = my_client.get_data_meta('data_3')
assert meta_3[0] == F.int64
assert meta_3[1] == tuple(F.shape(data_3))
assert_array_equal(meta_3[2], partition_0)
meta_4 = my_client.get_data_meta('data_4')
assert meta_4[0] == F.float64
assert meta_4[1] == tuple(F.shape(data_4))
assert_array_equal(meta_3[2], partition_0)
meta_5 = my_client.get_data_meta('data_5')
assert meta_5[0] == F.int32
assert meta_5[1] == tuple(F.shape(data_5))
assert_array_equal(meta_3[2], partition_0)
my_client.push(name='data_0',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor([[1., 1., 1.], [2., 2., 2.],
[3., 3., 3.]]))
my_client.push(name='data_2',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor([[1., 1., 1.], [2., 2., 2.],
[3., 3., 3.]]))
my_client.push(name='data_3',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor([[1, 1, 1], [2, 2, 2], [3, 3, 3]]))
my_client.push(name='data_4',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor(
[[1., 1., 1.], [2., 2., 2.], [3., 3., 3.]], F.float64))
my_client.push(name='data_5',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor([[1, 1, 1], [2, 2, 2], [3, 3, 3]],
F.int32))
target = F.tensor([[1., 1., 1.], [2., 2., 2.], [3., 3., 3.]])
res = my_client.pull(name='data_0', id_tensor=F.tensor([0, 1, 2]))
assert_array_equal(res, target)
res = my_client.pull(name='data_2', id_tensor=F.tensor([0, 1, 2]))
assert_array_equal(res, target)
target = F.tensor([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
res = my_client.pull(name='data_3', id_tensor=F.tensor([0, 1, 2]))
assert_array_equal(res, target)
target = F.tensor([[1., 1., 1.], [2., 2., 2.], [3., 3., 3.]], F.float64)
res = my_client.pull(name='data_4', id_tensor=F.tensor([0, 1, 2]))
assert_array_equal(res, target)
target = F.tensor([[1, 1, 1], [2, 2, 2], [3, 3, 3]], F.int32)
res = my_client.pull(name='data_5', id_tensor=F.tensor([0, 1, 2]))
assert_array_equal(res, target)
my_client.shut_down()
def start_client():
my_client = KVClient(server_namebook=server_namebook)
my_client.connect()
my_client.init_data(name='data_2',
shape=[num_entries, dim_size],
dtype=F.float32,
target_name='data_0')
name_list = my_client.get_data_name_list()
assert len(name_list) == 3
assert 'data_0' in name_list
assert 'data_1' in name_list
assert 'data_2' in name_list
meta_0 = my_client.get_data_meta('data_0')
assert meta_0[0] == F.float32
assert_array_equal(meta_0[2], partition_0)
meta_1 = my_client.get_data_meta('data_1')
assert meta_1[0] == F.float32
assert_array_equal(meta_1[2], partition_1)
meta_2 = my_client.get_data_meta('data_2')
assert meta_2[0] == F.float32
assert_array_equal(meta_2[2], partition_0)
my_client.push(name='data_0',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor([[1., 1., 1.], [2., 2., 2.],
[3., 3., 3.]]))
my_client.push(name='data_2',
id_tensor=F.tensor([0, 1, 2]),
data_tensor=F.tensor([[1., 1., 1.], [2., 2., 2.],
[3., 3., 3.]]))
target = F.tensor([[1., 1., 1.], [2., 2., 2.], [3., 3., 3.]])
res = my_client.pull(name='data_0', id_tensor=F.tensor([0, 1, 2]))
assert_array_equal(res, target)
res = my_client.pull(name='data_2', id_tensor=F.tensor([0, 1, 2]))
assert_array_equal(res, target)
my_client.shut_down()
def start_client(args):
"""Start client
"""
server_namebook = dgl.contrib.read_ip_config(filename=args.ip_config)
my_client = KVClient(server_namebook=server_namebook)
my_client.connect()
my_client.print()
my_client.barrier()
local_start = num_entries * my_client.get_machine_id()
local_end = num_entries * (my_client.get_machine_id() + 1)
local_range = np.arange(local_start, local_end)
id_list = []
for i in range(10000):
ids = np.random.choice(local_range, args.batch_size)
id_list.append(F.tensor(ids))
print("Pull from local...")
num_bytes = 0
start = time.time()
for ids in id_list:
tmp = my_client.pull(name='entity_embed', id_tensor=ids)
ndim = tmp.shape[1]
num_bytes += np.prod(tmp.shape) * 4
print("Total time: %.3f, #bytes: %.3f GB" %
(time.time() - start, num_bytes / 1000 / 1000 / 1000))
my_client.barrier()
arr = F.zeros((num_entries, ndim), F.float32, F.cpu())
print('Slice from a tensor...')
num_bytes = 0
start = time.time()
for ids in id_list:
tmp = arr[ids]
num_bytes += np.prod(tmp.shape) * 4
print("Total time: %.3f, #bytes: %.3f GB" %
(time.time() - start, num_bytes / 1000 / 1000 / 1000))
print("Pull from remote...")
if local_start == 0:
remote_range = np.arange(local_end, num_entries * args.num_servers)
elif local_end == num_entries * args.num_servers:
remote_range = np.arange(0, local_start)
else:
range1 = np.arange(0, local_start)
range2 = np.arange(local_end, num_entries * args.num_servers)
remote_range = np.concatenate((range1, range2))
id_list = []
for i in range(1000):
ids = np.random.choice(remote_range, args.batch_size)
id_list.append(F.tensor(ids))
num_bytes = 0
start = time.time()
for ids in id_list:
tmp = my_client.pull(name='entity_embed', id_tensor=ids)
num_bytes += np.prod(tmp.shape) * 4
print("Total time: %.3f, #bytes: %.3f GB" %
(time.time() - start, num_bytes / 1000 / 1000 / 1000))
my_client.barrier()
if my_client.get_id() == 0:
my_client.shut_down()