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:
my_client.set_partition_book(name='entity_embed')
my_client.print()
my_client.barrier()
print("send request...")
for i in range(100):
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=mx.nd.array([0, 1, 2, 3, 4, 5, 6, 7],
dtype='int64'))
print(res)
my_client.barrier()
for i in range(100):
my_client.push(name='entity_embed',
id_tensor=ID[my_client.get_machine_id()],
data_tensor=mx.nd.array([[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=mx.nd.array([0, 1, 2, 3, 4, 5, 6, 7],
dtype='int64'))
print(res)
if my_client.get_id() == 0:
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()
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()
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()