def test_sddmm(g, shp, lhs_target, rhs_target, msg, index_dtype):
if dgl.backend.backend_name == 'mxnet' and g.number_of_edges() == 0:
pytest.skip() # mxnet do not support zero shape tensor
if dgl.backend.backend_name == 'tensorflow' and index_dtype == 'int32':
pytest.skip() # tensorflow dlpack has problem with int32 ndarray.
if index_dtype == 'int32':
g = g.int()
else:
g = g.long()
print(g)
print(g.idtype)
len_lhs = select(lhs_target, g.number_of_src_nodes(), g.number_of_edges(),
g.number_of_dst_nodes())
lhs_shp = (len_lhs, ) + shp[0]
len_rhs = select(rhs_target, g.number_of_src_nodes(), g.number_of_edges(),
g.number_of_dst_nodes())
rhs_shp = (len_rhs, ) + shp[1]
feat_lhs = F.tensor(np.random.rand(*lhs_shp) + 1)
feat_rhs = F.tensor(np.random.rand(*rhs_shp) + 1)
print('lhs shape: {}, rhs shape: {}'.format(F.shape(feat_lhs),
F.shape(feat_rhs)))
lhs_frame = select(lhs_target, g.srcdata, g.edata, g.dstdata)
rhs_frame = select(rhs_target, g.srcdata, g.edata, g.dstdata)
lhs_frame['x'] = F.attach_grad(F.clone(feat_lhs))
rhs_frame['y'] = F.attach_grad(F.clone(feat_rhs))
msg_func = lhs_target + '_' + msg + '_' + rhs_target
print('SDDMM(message func: {})'.format(msg_func))
lhs = F.attach_grad(F.clone(feat_lhs))
rhs = F.attach_grad(F.clone(feat_rhs))
with F.record_grad():
e = gsddmm(g,
msg,
lhs,
rhs,
lhs_target=lhs_target,
rhs_target=rhs_target)
F.backward(F.reduce_sum(e))
grad_lhs = F.grad(lhs)
grad_rhs = F.grad(rhs)
with F.record_grad():
g.apply_edges(udf_apply_edges[msg_func])
if g.number_of_edges() > 0:
e1 = g.edata['m']
assert F.allclose(e, e1)
print('forward passed')
F.backward(F.reduce_sum(e1))
if msg != 'copy_rhs':
assert F.allclose(F.grad(lhs_frame['x']), grad_lhs)
if msg != 'copy_lhs':
assert F.allclose(F.grad(rhs_frame['y']), grad_rhs)
print('backward passed')
lhs_frame.pop('x')
rhs_frame.pop('y')
if 'm' in g.edata: g.edata.pop('m')
def test_row1():
# test row getter/setter
data = create_test_data()
f = FrameRef(Frame(data))
# getter
# test non-duplicate keys
rowid = Index(F.tensor([0, 2]))
rows = f[rowid]
for k, v in rows.items():
assert tuple(F.shape(v)) == (len(rowid), D)
assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
# test duplicate keys
rowid = Index(F.tensor([8, 2, 2, 1]))
rows = f[rowid]
for k, v in rows.items():
assert tuple(F.shape(v)) == (len(rowid), D)
assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
# setter
rowid = Index(F.tensor([0, 2, 4]))
vals = {'a1' : F.zeros((len(rowid), D)),
'a2' : F.zeros((len(rowid), D)),
'a3' : F.zeros((len(rowid), D)),
}
f[rowid] = vals
for k, v in f[rowid].items():
assert F.allclose(v, F.zeros((len(rowid), D)))
# setting rows with new column should raise error with error initializer
f.set_initializer(lambda shape, dtype : assert_(False))
def failed_update_rows():
vals['a4'] = F.ones((len(rowid), D))
f[rowid] = vals
assert check_fail(failed_update_rows)
def _test_layer_sampler(prefetch=False):
g = generate_rand_graph(100)
nid = g.nodes()
src, dst, eid = g.all_edges(form='all', order='eid')
n_batches = 5
batch_size = 50
seed_batches = [
np.sort(np.random.choice(F.asnumpy(nid), batch_size, replace=False))
for i in range(n_batches)
]
seed_nodes = np.hstack(seed_batches)
layer_sizes = [50] * 3
LayerSampler = getattr(dgl.contrib.sampling, 'LayerSampler')
sampler = LayerSampler(g,
batch_size,
layer_sizes,
'in',
seed_nodes=seed_nodes,
num_workers=4,
prefetch=prefetch)
for sub_g in sampler:
assert all(
sub_g.layer_size(i) < size for i, size in enumerate(layer_sizes))
sub_nid = F.arange(0, sub_g.number_of_nodes())
assert all(
np.all(np.isin(F.asnumpy(sub_g.layer_nid(i)), F.asnumpy(sub_nid)))
for i in range(sub_g.num_layers))
assert np.all(
np.isin(F.asnumpy(sub_g.map_to_parent_nid(sub_nid)),
F.asnumpy(nid)))
sub_eid = F.arange(0, sub_g.number_of_edges())
assert np.all(
np.isin(F.asnumpy(sub_g.map_to_parent_eid(sub_eid)),
F.asnumpy(eid)))
assert any(
np.all(
np.sort(F.asnumpy(sub_g.layer_parent_nid(-1))) == seed_batch)
for seed_batch in seed_batches)
sub_src, sub_dst = sub_g.all_edges(order='eid')
for i in range(sub_g.num_blocks):
block_eid = sub_g.block_eid(i)
block_src = sub_g.map_to_parent_nid(sub_src[block_eid])
block_dst = sub_g.map_to_parent_nid(sub_dst[block_eid])
block_parent_eid = sub_g.block_parent_eid(i)
block_parent_src = src[block_parent_eid]
block_parent_dst = dst[block_parent_eid]
assert np.all(F.asnumpy(block_src == block_parent_src))
n_layers = sub_g.num_layers
sub_n = sub_g.number_of_nodes()
assert sum(F.shape(sub_g.layer_nid(i))[0]
for i in range(n_layers)) == sub_n
n_blocks = sub_g.num_blocks
sub_m = sub_g.number_of_edges()
assert sum(F.shape(sub_g.block_eid(i))[0]
for i in range(n_blocks)) == sub_m
def test_readonly():
g = dgl.DGLGraph()
g.add_nodes(5)
g.add_edges([0, 1, 2, 3], [1, 2, 3, 4])
g.ndata['x'] = F.zeros((5, 3))
g.edata['x'] = F.zeros((4, 4))
g.readonly(False)
assert g._graph.is_readonly() == False
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
g.readonly()
assert g._graph.is_readonly() == True
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
try:
g.add_nodes(5)
fail = False
except DGLError:
fail = True
finally:
assert fail
g.readonly()
assert g._graph.is_readonly() == True
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
try:
g.add_nodes(5)
fail = False
except DGLError:
fail = True
finally:
assert fail
g.readonly(False)
assert g._graph.is_readonly() == False
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
try:
g.add_nodes(10)
g.add_edges([4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
[5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
fail = False
except DGLError:
fail = True
finally:
assert not fail
assert g.number_of_nodes() == 15
assert F.shape(g.ndata['x']) == (15, 3)
assert g.number_of_edges() == 14
assert F.shape(g.edata['x']) == (14, 4)
def test_spmm(idtype, g, shp, msg, reducer):
g = g.astype(idtype).to(F.ctx())
if dgl.backend.backend_name == 'tensorflow' and (reducer in ['min', 'max']):
pytest.skip() # tensorflow dlpack has problem writing into int32 arrays on GPU.
print(g)
print(g.idtype)
hu = F.tensor(np.random.rand(*((g.number_of_src_nodes(),) + shp[0])) + 1)
he = F.tensor(np.random.rand(*((g.number_of_edges(),) + shp[1])) + 1)
print('u shape: {}, e shape: {}'.format(F.shape(hu), F.shape(he)))
g.srcdata['x'] = F.attach_grad(F.clone(hu))
g.edata['w'] = F.attach_grad(F.clone(he))
print('SpMM(message func: {}, reduce func: {})'.format(msg, reducer))
u = F.attach_grad(F.clone(hu))
e = F.attach_grad(F.clone(he))
with F.record_grad():
v = gspmm(g, msg, reducer, u, e)
non_degree_indices = F.tensor(
np.nonzero(F.asnumpy(g.in_degrees()) != 0)[0])
v = F.gather_row(v, non_degree_indices)
if g.number_of_edges() > 0:
F.backward(F.reduce_sum(v))
if msg != 'copy_rhs':
grad_u = F.grad(u)
if msg != 'copy_lhs':
grad_e = F.grad(e)
with F.record_grad():
g.update_all(udf_msg[msg], udf_reduce[reducer])
if g.number_of_edges() > 0:
v1 = F.gather_row(g.dstdata['v'], non_degree_indices)
assert F.allclose(v, v1)
print('forward passed')
F.backward(F.reduce_sum(v1))
if msg != 'copy_rhs':
if reducer in ['min', 'max']: # there might be some numerical errors
rate = F.reduce_sum(F.abs(F.grad(g.srcdata['x']) - grad_u)) /\
F.reduce_sum(F.abs(grad_u))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.srcdata['x']), grad_u)
if msg != 'copy_lhs':
if reducer in ['min', 'max']:
rate = F.reduce_sum(F.abs(F.grad(g.edata['w']) - grad_e)) /\
F.reduce_sum(F.abs(grad_e))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.edata['w']), grad_e)
print('backward passed')
g.srcdata.pop('x')
g.edata.pop('w')
if 'v' in g.dstdata: g.dstdata.pop('v')
def test_spmm(idtype, g, shp, msg, reducer):
g = g.astype(idtype).to(F.ctx())
print(g)
print(g.idtype)
hu = F.tensor(np.random.rand(*((g.number_of_src_nodes(), ) + shp[0])) + 1)
he = F.tensor(np.random.rand(*((g.number_of_edges(), ) + shp[1])) + 1)
print('u shape: {}, e shape: {}'.format(F.shape(hu), F.shape(he)))
g.srcdata['x'] = F.attach_grad(F.clone(hu))
g.edata['w'] = F.attach_grad(F.clone(he))
print('SpMM(message func: {}, reduce func: {})'.format(msg, reducer))
u = F.attach_grad(F.clone(hu))
e = F.attach_grad(F.clone(he))
with F.record_grad():
v = gspmm(g, msg, reducer, u, e)
if reducer in ['max', 'min']:
v = F.replace_inf_with_zero(v)
if g.number_of_edges() > 0:
F.backward(F.reduce_sum(v))
if msg != 'copy_rhs':
grad_u = F.grad(u)
if msg != 'copy_lhs':
grad_e = F.grad(e)
with F.record_grad():
g.update_all(udf_msg[msg], udf_reduce[reducer])
if g.number_of_edges() > 0:
v1 = g.dstdata['v']
assert F.allclose(v, v1)
print('forward passed')
F.backward(F.reduce_sum(v1))
if msg != 'copy_rhs':
if reducer in ['min',
'max']: # there might be some numerical errors
rate = F.reduce_sum(F.abs(F.grad(g.srcdata['x']) - grad_u)) /\
F.reduce_sum(F.abs(grad_u))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.srcdata['x']), grad_u)
if msg != 'copy_lhs':
if reducer in ['min', 'max']:
rate = F.reduce_sum(F.abs(F.grad(g.edata['w']) - grad_e)) /\
F.reduce_sum(F.abs(grad_e))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.edata['w']), grad_e)
print('backward passed')
g.srcdata.pop('x')
g.edata.pop('w')
if 'v' in g.dstdata: g.dstdata.pop('v')
def __call__(self, edges):
sdata = edges.src[self.src_field]
edata = edges.data[self.edge_field]
# Due to the different broadcasting semantics of different backends,
# we need to broadcast the sdata and edata to be of the same rank.
rank = max(F.ndim(sdata), F.ndim(edata))
sshape = F.shape(sdata)
eshape = F.shape(edata)
sdata = F.reshape(sdata, sshape + (1, ) * (rank - F.ndim(sdata)))
edata = F.reshape(edata, eshape + (1, ) * (rank - F.ndim(edata)))
ret = self.mul_op(sdata, edata)
return {self.out_field: ret}
def test_segment_reduce(reducer):
ctx = F.ctx()
value = F.tensor(np.random.rand(10, 5))
v1 = F.attach_grad(F.clone(value))
v2 = F.attach_grad(F.clone(value))
seglen = F.tensor([2, 3, 0, 4, 1, 0, 0])
u = F.copy_to(F.arange(0, F.shape(value)[0], F.int32), ctx)
v = F.repeat(F.copy_to(F.arange(0, len(seglen), F.int32), ctx),
seglen,
dim=0)
num_nodes = {'_U': len(u), '_V': len(seglen)}
g = dgl.convert.heterograph({('_U', '_E', '_V'): (u, v)},
num_nodes_dict=num_nodes)
with F.record_grad():
rst1 = gspmm(g, 'copy_lhs', reducer, v1, None)
if reducer in ['max', 'min']:
rst1 = F.replace_inf_with_zero(rst1)
F.backward(F.reduce_sum(rst1))
grad1 = F.grad(v1)
with F.record_grad():
rst2 = segment_reduce(seglen, v2, reducer=reducer)
F.backward(F.reduce_sum(rst2))
assert F.allclose(rst1, rst2)
print('forward passed')
grad2 = F.grad(v2)
assert F.allclose(grad1, grad2)
print('backward passed')
def test_append1():
# test append API on Frame
data = create_test_data()
f1 = Frame()
f2 = Frame(data)
f1.append(data)
assert f1.num_rows == N
f1.append(f2)
assert f1.num_rows == 2 * N
c1 = f1['a1']
assert tuple(F.shape(c1.data)) == (2 * N, D)
truth = F.cat([data['a1'], data['a1']], 0)
assert F.allclose(truth, c1.data)
# append dict of different length columns should fail
f3 = {'a1' : F.zeros((3, D)), 'a2' : F.zeros((3, D)), 'a3' : F.zeros((2, D))}
def failed_append():
f1.append(f3)
assert check_fail(failed_append)
def check_partition(g, part_method, reshuffle):
g.ndata['labels'] = F.arange(0, g.number_of_nodes())
g.ndata['feats'] = F.tensor(np.random.randn(g.number_of_nodes(), 10),
F.float32)
g.edata['feats'] = F.tensor(np.random.randn(g.number_of_edges(), 10),
F.float32)
g.update_all(fn.copy_src('feats', 'msg'), fn.sum('msg', 'h'))
g.update_all(fn.copy_edge('feats', 'msg'), fn.sum('msg', 'eh'))
num_parts = 4
num_hops = 2
orig_nids, orig_eids = partition_graph(g,
'test',
num_parts,
'/tmp/partition',
num_hops=num_hops,
part_method=part_method,
reshuffle=reshuffle,
return_mapping=True)
part_sizes = []
shuffled_labels = []
shuffled_edata = []
for i in range(num_parts):
part_g, node_feats, edge_feats, gpb, _, ntypes, etypes = load_partition(
'/tmp/partition/test.json', i)
# Check the metadata
assert gpb._num_nodes() == g.number_of_nodes()
assert gpb._num_edges() == g.number_of_edges()
assert gpb.num_partitions() == num_parts
gpb_meta = gpb.metadata()
assert len(gpb_meta) == num_parts
assert len(gpb.partid2nids(i)) == gpb_meta[i]['num_nodes']
assert len(gpb.partid2eids(i)) == gpb_meta[i]['num_edges']
part_sizes.append((gpb_meta[i]['num_nodes'], gpb_meta[i]['num_edges']))
nid = F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node'])
local_nid = gpb.nid2localnid(nid, i)
assert F.dtype(local_nid) in (F.int64, F.int32)
assert np.all(F.asnumpy(local_nid) == np.arange(0, len(local_nid)))
eid = F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge'])
local_eid = gpb.eid2localeid(eid, i)
assert F.dtype(local_eid) in (F.int64, F.int32)
assert np.all(F.asnumpy(local_eid) == np.arange(0, len(local_eid)))
# Check the node map.
local_nodes = F.boolean_mask(part_g.ndata[dgl.NID],
part_g.ndata['inner_node'])
llocal_nodes = F.nonzero_1d(part_g.ndata['inner_node'])
local_nodes1 = gpb.partid2nids(i)
assert F.dtype(local_nodes1) in (F.int32, F.int64)
assert np.all(
np.sort(F.asnumpy(local_nodes)) == np.sort(F.asnumpy(
local_nodes1)))
assert np.all(F.asnumpy(llocal_nodes) == np.arange(len(llocal_nodes)))
# Check the edge map.
local_edges = F.boolean_mask(part_g.edata[dgl.EID],
part_g.edata['inner_edge'])
llocal_edges = F.nonzero_1d(part_g.edata['inner_edge'])
local_edges1 = gpb.partid2eids(i)
assert F.dtype(local_edges1) in (F.int32, F.int64)
assert np.all(
np.sort(F.asnumpy(local_edges)) == np.sort(F.asnumpy(
local_edges1)))
assert np.all(F.asnumpy(llocal_edges) == np.arange(len(llocal_edges)))
# Verify the mapping between the reshuffled IDs and the original IDs.
part_src_ids, part_dst_ids = part_g.edges()
part_src_ids = F.gather_row(part_g.ndata[dgl.NID], part_src_ids)
part_dst_ids = F.gather_row(part_g.ndata[dgl.NID], part_dst_ids)
part_eids = part_g.edata[dgl.EID]
orig_src_ids = F.gather_row(orig_nids, part_src_ids)
orig_dst_ids = F.gather_row(orig_nids, part_dst_ids)
orig_eids1 = F.gather_row(orig_eids, part_eids)
orig_eids2 = g.edge_ids(orig_src_ids, orig_dst_ids)
assert F.shape(orig_eids1)[0] == F.shape(orig_eids2)[0]
assert np.all(F.asnumpy(orig_eids1) == F.asnumpy(orig_eids2))
if reshuffle:
part_g.ndata['feats'] = F.gather_row(g.ndata['feats'],
part_g.ndata['orig_id'])
part_g.edata['feats'] = F.gather_row(g.edata['feats'],
part_g.edata['orig_id'])
# when we read node data from the original global graph, we should use orig_id.
local_nodes = F.boolean_mask(part_g.ndata['orig_id'],
part_g.ndata['inner_node'])
local_edges = F.boolean_mask(part_g.edata['orig_id'],
part_g.edata['inner_edge'])
else:
part_g.ndata['feats'] = F.gather_row(g.ndata['feats'],
part_g.ndata[dgl.NID])
part_g.edata['feats'] = F.gather_row(g.edata['feats'],
part_g.edata[dgl.NID])
part_g.update_all(fn.copy_src('feats', 'msg'), fn.sum('msg', 'h'))
part_g.update_all(fn.copy_edge('feats', 'msg'), fn.sum('msg', 'eh'))
assert F.allclose(F.gather_row(g.ndata['h'], local_nodes),
F.gather_row(part_g.ndata['h'], llocal_nodes))
assert F.allclose(F.gather_row(g.ndata['eh'], local_nodes),
F.gather_row(part_g.ndata['eh'], llocal_nodes))
for name in ['labels', 'feats']:
assert '_N/' + name in node_feats
assert node_feats['_N/' + name].shape[0] == len(local_nodes)
true_feats = F.gather_row(g.ndata[name], local_nodes)
ndata = F.gather_row(node_feats['_N/' + name], local_nid)
assert np.all(F.asnumpy(true_feats) == F.asnumpy(ndata))
for name in ['feats']:
assert '_E/' + name in edge_feats
assert edge_feats['_E/' + name].shape[0] == len(local_edges)
true_feats = F.gather_row(g.edata[name], local_edges)
edata = F.gather_row(edge_feats['_E/' + name], local_eid)
assert np.all(F.asnumpy(true_feats) == F.asnumpy(edata))
# This only works if node/edge IDs are shuffled.
if reshuffle:
shuffled_labels.append(node_feats['_N/labels'])
shuffled_edata.append(edge_feats['_E/feats'])
# Verify that we can reconstruct node/edge data for original IDs.
if reshuffle:
shuffled_labels = F.asnumpy(F.cat(shuffled_labels, 0))
shuffled_edata = F.asnumpy(F.cat(shuffled_edata, 0))
orig_labels = np.zeros(shuffled_labels.shape,
dtype=shuffled_labels.dtype)
orig_edata = np.zeros(shuffled_edata.shape, dtype=shuffled_edata.dtype)
orig_labels[F.asnumpy(orig_nids)] = shuffled_labels
orig_edata[F.asnumpy(orig_eids)] = shuffled_edata
assert np.all(orig_labels == F.asnumpy(g.ndata['labels']))
assert np.all(orig_edata == F.asnumpy(g.edata['feats']))
if reshuffle:
node_map = []
edge_map = []
for i, (num_nodes, num_edges) in enumerate(part_sizes):
node_map.append(np.ones(num_nodes) * i)
edge_map.append(np.ones(num_edges) * i)
node_map = np.concatenate(node_map)
edge_map = np.concatenate(edge_map)
nid2pid = gpb.nid2partid(F.arange(0, len(node_map)))
assert F.dtype(nid2pid) in (F.int32, F.int64)
assert np.all(F.asnumpy(nid2pid) == node_map)
eid2pid = gpb.eid2partid(F.arange(0, len(edge_map)))
assert F.dtype(eid2pid) in (F.int32, F.int64)
assert np.all(F.asnumpy(eid2pid) == edge_map)
def _fmsg(edges):
assert tuple(F.shape(edges.src['h'])) == (5, D)
return {'m': edges.src['h']}
def reduce_func(nodes):
msgs = nodes.mailbox['m']
reduce_msg_shapes.add(tuple(msgs.shape))
assert F.ndim(msgs) == 3
assert F.shape(msgs)[2] == D
return {'accum': F.sum(msgs, 1)}
def message_func(edges):
assert F.ndim(edges.src['h']) == 2
assert F.shape(edges.src['h'])[1] == D
return {'m': edges.src['h']}
def start_client(num_clients, num_servers):
os.environ['DGL_DIST_MODE'] = 'distributed'
# Note: connect to server first !
dgl.distributed.initialize(ip_config='kv_ip_config.txt')
# Init kvclient
kvclient = dgl.distributed.KVClient(ip_config='kv_ip_config.txt', num_servers=num_servers)
kvclient.map_shared_data(partition_book=gpb)
assert dgl.distributed.get_num_client() == num_clients
kvclient.init_data(name='data_1',
shape=F.shape(data_1),
dtype=F.dtype(data_1),
part_policy=edge_policy,
init_func=init_zero_func)
kvclient.init_data(name='data_2',
shape=F.shape(data_2),
dtype=F.dtype(data_2),
part_policy=node_policy,
init_func=init_zero_func)
# Test data_name_list
name_list = kvclient.data_name_list()
print(name_list)
assert 'data_0' in name_list
assert 'data_0_1' in name_list
assert 'data_0_2' in name_list
assert 'data_0_3' in name_list
assert 'data_1' in name_list
assert 'data_2' in name_list
# Test get_meta_data
meta = kvclient.get_data_meta('data_0')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0)
assert shape == F.shape(data_0)
assert policy.policy_str == 'node:_N'
meta = kvclient.get_data_meta('data_0_1')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0_1)
assert shape == F.shape(data_0_1)
assert policy.policy_str == 'node:_N'
meta = kvclient.get_data_meta('data_0_2')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0_2)
assert shape == F.shape(data_0_2)
assert policy.policy_str == 'node:_N'
meta = kvclient.get_data_meta('data_0_3')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0_3)
assert shape == F.shape(data_0_3)
assert policy.policy_str == 'node:_N'
meta = kvclient.get_data_meta('data_1')
dtype, shape, policy = meta
assert dtype == F.dtype(data_1)
assert shape == F.shape(data_1)
assert policy.policy_str == 'edge:_E'
meta = kvclient.get_data_meta('data_2')
dtype, shape, policy = meta
assert dtype == F.dtype(data_2)
assert shape == F.shape(data_2)
assert policy.policy_str == 'node:_N'
# Test push and pull
id_tensor = F.tensor([0,2,4], F.int64)
data_tensor = F.tensor([[6.,6.],[6.,6.],[6.,6.]], F.float32)
kvclient.push(name='data_0',
id_tensor=id_tensor,
data_tensor=data_tensor)
kvclient.push(name='data_1',
id_tensor=id_tensor,
data_tensor=data_tensor)
kvclient.push(name='data_2',
id_tensor=id_tensor,
data_tensor=data_tensor)
res = kvclient.pull(name='data_0', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_1', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_2', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
# Register new push handler
kvclient.register_push_handler('data_0', udf_push)
kvclient.register_push_handler('data_1', udf_push)
kvclient.register_push_handler('data_2', udf_push)
# Test push and pull
kvclient.push(name='data_0',
id_tensor=id_tensor,
data_tensor=data_tensor)
kvclient.push(name='data_1',
id_tensor=id_tensor,
data_tensor=data_tensor)
kvclient.push(name='data_2',
id_tensor=id_tensor,
data_tensor=data_tensor)
kvclient.barrier()
data_tensor = data_tensor * data_tensor
res = kvclient.pull(name='data_0', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_1', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_2', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
# Test delete data
kvclient.delete_data('data_0')
kvclient.delete_data('data_1')
kvclient.delete_data('data_2')
# Register new push handler
kvclient.init_data(name='data_3',
shape=F.shape(data_2),
dtype=F.dtype(data_2),
part_policy=node_policy,
init_func=init_zero_func)
kvclient.register_push_handler('data_3', add_push)
data_tensor = F.tensor([[6.,6.],[6.,6.],[6.,6.]], F.float32)
kvclient.barrier()
time.sleep(kvclient.client_id + 1)
print("add...")
kvclient.push(name='data_3',
id_tensor=id_tensor,
data_tensor=data_tensor)
kvclient.barrier()
res = kvclient.pull(name='data_3', id_tensor=id_tensor)
data_tensor = data_tensor * num_clients
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
def test_edge_batch():
d = 10
g = dgl.DGLGraph(nx.path_graph(20))
nfeat = F.randn((g.number_of_nodes(), d))
efeat = F.randn((g.number_of_edges(), d))
g.ndata['x'] = nfeat
g.edata['x'] = efeat
# test all
eid = ALL
u, v, _ = g._graph.edges('eid')
src_data = g.get_n_repr(u)
edge_data = g.get_e_repr(eid)
dst_data = g.get_n_repr(v)
ebatch = EdgeBatch(g, (u, v, eid), src_data, edge_data, dst_data)
assert F.shape(ebatch.src['x'])[0] == g.number_of_edges() and\
F.shape(ebatch.src['x'])[1] == d
assert F.shape(ebatch.dst['x'])[0] == g.number_of_edges() and\
F.shape(ebatch.dst['x'])[1] == d
assert F.shape(ebatch.data['x'])[0] == g.number_of_edges() and\
F.shape(ebatch.data['x'])[1] == d
assert F.allclose(ebatch.edges()[0], u.tousertensor())
assert F.allclose(ebatch.edges()[1], v.tousertensor())
assert F.allclose(ebatch.edges()[2], F.arange(0, g.number_of_edges()))
assert ebatch.batch_size() == g.number_of_edges()
assert len(ebatch) == g.number_of_edges()
# test partial
eid = utils.toindex(F.tensor([0, 3, 5, 7, 11, 13, 15, 27]))
u, v, _ = g._graph.find_edges(eid)
src_data = g.get_n_repr(u)
edge_data = g.get_e_repr(eid)
dst_data = g.get_n_repr(v)
ebatch = EdgeBatch(g, (u, v, eid), src_data, edge_data, dst_data)
assert F.shape(ebatch.src['x'])[0] == 8 and\
F.shape(ebatch.src['x'])[1] == d
assert F.shape(ebatch.dst['x'])[0] == 8 and\
F.shape(ebatch.dst['x'])[1] == d
assert F.shape(ebatch.data['x'])[0] == 8 and\
F.shape(ebatch.data['x'])[1] == d
assert F.allclose(ebatch.edges()[0], u.tousertensor())
assert F.allclose(ebatch.edges()[1], v.tousertensor())
assert F.allclose(ebatch.edges()[2], eid.tousertensor())
assert ebatch.batch_size() == 8
assert len(ebatch) == 8
def start_client():
# Note: connect to server first !
dgl.distributed.connect_to_server(ip_config='kv_ip_config.txt')
# Init kvclient
kvclient = dgl.distributed.KVClient(ip_config='kv_ip_config.txt')
kvclient.init_data(name='data_1',
shape=F.shape(data_1),
dtype=F.dtype(data_1),
policy_str='edge',
partition_book=gpb,
init_func=init_zero_func)
kvclient.init_data(name='data_2',
shape=F.shape(data_2),
dtype=F.dtype(data_2),
policy_str='node',
partition_book=gpb,
init_func=init_zero_func)
kvclient.map_shared_data(partition_book=gpb)
# Test data_name_list
name_list = kvclient.data_name_list()
print(name_list)
assert 'data_0' in name_list
assert 'data_0_1' in name_list
assert 'data_0_2' in name_list
assert 'data_0_3' in name_list
assert 'data_1' in name_list
assert 'data_2' in name_list
# Test get_meta_data
meta = kvclient.get_data_meta('data_0')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0)
assert shape == F.shape(data_0)
assert policy.policy_str == 'node'
meta = kvclient.get_data_meta('data_0_1')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0_1)
assert shape == F.shape(data_0_1)
assert policy.policy_str == 'node'
meta = kvclient.get_data_meta('data_0_2')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0_2)
assert shape == F.shape(data_0_2)
assert policy.policy_str == 'node'
meta = kvclient.get_data_meta('data_0_3')
dtype, shape, policy = meta
assert dtype == F.dtype(data_0_3)
assert shape == F.shape(data_0_3)
assert policy.policy_str == 'node'
meta = kvclient.get_data_meta('data_1')
dtype, shape, policy = meta
assert dtype == F.dtype(data_1)
assert shape == F.shape(data_1)
assert policy.policy_str == 'edge'
meta = kvclient.get_data_meta('data_2')
dtype, shape, policy = meta
assert dtype == F.dtype(data_2)
assert shape == F.shape(data_2)
assert policy.policy_str == 'node'
# Test push and pull
id_tensor = F.tensor([0, 2, 4], F.int64)
data_tensor = F.tensor([[6., 6.], [6., 6.], [6., 6.]], F.float32)
kvclient.push(name='data_0', id_tensor=id_tensor, data_tensor=data_tensor)
kvclient.push(name='data_1', id_tensor=id_tensor, data_tensor=data_tensor)
kvclient.push(name='data_2', id_tensor=id_tensor, data_tensor=data_tensor)
res = kvclient.pull(name='data_0', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_1', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_2', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
# Register new push handler
kvclient.register_push_handler('data_0', udf_push)
kvclient.register_push_handler('data_1', udf_push)
kvclient.register_push_handler('data_2', udf_push)
# Test push and pull
kvclient.push(name='data_0', id_tensor=id_tensor, data_tensor=data_tensor)
kvclient.push(name='data_1', id_tensor=id_tensor, data_tensor=data_tensor)
kvclient.push(name='data_2', id_tensor=id_tensor, data_tensor=data_tensor)
data_tensor = data_tensor * data_tensor
res = kvclient.pull(name='data_0', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_1', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
res = kvclient.pull(name='data_2', id_tensor=id_tensor)
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
# clean up
dgl.distributed.shutdown_servers()
dgl.distributed.finalize_client()
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()
