def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server):
generate_ip_config("rpc_ip_config.txt", num_server, num_server)
g = create_random_hetero(dense=True)
num_parts = num_server
num_hops = 1
partition_graph(g, 'test_sampling', num_parts, tmpdir,
num_hops=num_hops, part_method='metis', reshuffle=True)
pserver_list = []
ctx = mp.get_context('spawn')
for i in range(num_server):
p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling'))
p.start()
time.sleep(1)
pserver_list.append(p)
fanout = 3
block, gpb = start_hetero_etype_sample_client(0, tmpdir, num_server > 1, fanout,
nodes={'n3': [0, 10, 99, 66, 124, 208]})
print("Done sampling")
for p in pserver_list:
p.join()
src, dst = block.edges(etype=('n1', 'r2', 'n3'))
assert len(src) == 18
src, dst = block.edges(etype=('n2', 'r3', 'n3'))
assert len(src) == 18
orig_nid_map = {ntype: F.zeros((g.number_of_nodes(ntype),), dtype=F.int64) for ntype in g.ntypes}
orig_eid_map = {etype: F.zeros((g.number_of_edges(etype),), dtype=F.int64) for etype in g.etypes}
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json', i)
ntype_ids, type_nids = gpb.map_to_per_ntype(part.ndata[dgl.NID])
for ntype_id, ntype in enumerate(g.ntypes):
idx = ntype_ids == ntype_id
F.scatter_row_inplace(orig_nid_map[ntype], F.boolean_mask(type_nids, idx),
F.boolean_mask(part.ndata['orig_id'], idx))
etype_ids, type_eids = gpb.map_to_per_etype(part.edata[dgl.EID])
for etype_id, etype in enumerate(g.etypes):
idx = etype_ids == etype_id
F.scatter_row_inplace(orig_eid_map[etype], F.boolean_mask(type_eids, idx),
F.boolean_mask(part.edata['orig_id'], idx))
for src_type, etype, dst_type in block.canonical_etypes:
src, dst = block.edges(etype=etype)
# These are global Ids after shuffling.
shuffled_src = F.gather_row(block.srcnodes[src_type].data[dgl.NID], src)
shuffled_dst = F.gather_row(block.dstnodes[dst_type].data[dgl.NID], dst)
shuffled_eid = block.edges[etype].data[dgl.EID]
orig_src = F.asnumpy(F.gather_row(orig_nid_map[src_type], shuffled_src))
orig_dst = F.asnumpy(F.gather_row(orig_nid_map[dst_type], shuffled_dst))
orig_eid = F.asnumpy(F.gather_row(orig_eid_map[etype], shuffled_eid))
# Check the node Ids and edge Ids.
orig_src1, orig_dst1 = g.find_edges(orig_eid, etype=etype)
assert np.all(F.asnumpy(orig_src1) == orig_src)
assert np.all(F.asnumpy(orig_dst1) == orig_dst)
def check_hetero_partition(hg, part_method):
hg.nodes['n1'].data['labels'] = F.arange(0, hg.number_of_nodes('n1'))
hg.nodes['n1'].data['feats'] = F.tensor(
np.random.randn(hg.number_of_nodes('n1'), 10), F.float32)
hg.edges['r1'].data['feats'] = F.tensor(
np.random.randn(hg.number_of_edges('r1'), 10), F.float32)
num_parts = 4
num_hops = 1
orig_nids, orig_eids = partition_graph(hg,
'test',
num_parts,
'/tmp/partition',
num_hops=num_hops,
part_method=part_method,
reshuffle=True,
return_mapping=True)
assert len(orig_nids) == len(hg.ntypes)
assert len(orig_eids) == len(hg.etypes)
for ntype in hg.ntypes:
assert len(orig_nids[ntype]) == hg.number_of_nodes(ntype)
for etype in hg.etypes:
assert len(orig_eids[etype]) == hg.number_of_edges(etype)
parts = []
for i in range(num_parts):
part_g, node_feats, edge_feats, gpb, _, ntypes, etypes = load_partition(
'/tmp/partition/test.json', i)
# Verify the mapping between the reshuffled IDs and the original IDs.
# These are partition-local IDs.
part_src_ids, part_dst_ids = part_g.edges()
# These are reshuffled global homogeneous IDs.
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]
# These are reshuffled per-type IDs.
src_ntype_ids, part_src_ids = gpb.map_to_per_ntype(part_src_ids)
dst_ntype_ids, part_dst_ids = gpb.map_to_per_ntype(part_dst_ids)
etype_ids, part_eids = gpb.map_to_per_etype(part_eids)
# These are original per-type IDs.
for etype_id, etype in enumerate(hg.etypes):
part_src_ids1 = F.boolean_mask(part_src_ids, etype_ids == etype_id)
src_ntype_ids1 = F.boolean_mask(src_ntype_ids,
etype_ids == etype_id)
part_dst_ids1 = F.boolean_mask(part_dst_ids, etype_ids == etype_id)
dst_ntype_ids1 = F.boolean_mask(dst_ntype_ids,
etype_ids == etype_id)
part_eids1 = F.boolean_mask(part_eids, etype_ids == etype_id)
assert np.all(F.asnumpy(src_ntype_ids1 == src_ntype_ids1[0]))
assert np.all(F.asnumpy(dst_ntype_ids1 == dst_ntype_ids1[0]))
src_ntype = hg.ntypes[F.as_scalar(src_ntype_ids1[0])]
dst_ntype = hg.ntypes[F.as_scalar(dst_ntype_ids1[0])]
orig_src_ids1 = F.gather_row(orig_nids[src_ntype], part_src_ids1)
orig_dst_ids1 = F.gather_row(orig_nids[dst_ntype], part_dst_ids1)
orig_eids1 = F.gather_row(orig_eids[etype], part_eids1)
orig_eids2 = hg.edge_ids(orig_src_ids1, orig_dst_ids1, etype=etype)
assert len(orig_eids1) == len(orig_eids2)
assert np.all(F.asnumpy(orig_eids1) == F.asnumpy(orig_eids2))
parts.append(part_g)
verify_graph_feats(hg, part_g, node_feats)
verify_hetero_graph(hg, parts)
def check_partition(reshuffle):
g = create_random_graph(10000)
g.ndata['labels'] = F.arange(0, g.number_of_nodes())
g.ndata['feats'] = F.tensor(np.random.randn(g.number_of_nodes(), 10))
num_parts = 4
num_hops = 2
partition_graph(g, 'test', num_parts, '/tmp/partition', num_hops=num_hops,
part_method='metis', reshuffle=reshuffle)
part_sizes = []
for i in range(num_parts):
part_g, node_feats, edge_feats, gpb = 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']))
local_nid = gpb.nid2localnid(F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node']), i)
assert np.all(F.asnumpy(local_nid) == np.arange(0, len(local_nid)))
local_eid = gpb.eid2localeid(F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge']), i)
assert np.all(F.asnumpy(local_eid) == np.arange(0, len(local_eid)))
# Check the node map.
local_nodes = F.asnumpy(F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node']))
local_nodes1 = F.asnumpy(gpb.partid2nids(i))
assert np.all(np.sort(local_nodes) == np.sort(local_nodes1))
# Check the edge map.
local_edges = F.asnumpy(F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge']))
local_edges1 = F.asnumpy(gpb.partid2eids(i))
assert np.all(np.sort(local_edges) == np.sort(local_edges1))
for name in ['labels', 'feats']:
assert name in node_feats
assert node_feats[name].shape[0] == len(local_nodes)
assert len(local_nodes) == len(node_feats[name])
assert np.all(F.asnumpy(g.ndata[name])[local_nodes] == F.asnumpy(node_feats[name]))
assert len(edge_feats) == 0
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)
assert np.all(F.asnumpy(gpb.nid2partid(F.arange(0, len(node_map)))) == node_map)
assert np.all(F.asnumpy(gpb.eid2partid(F.arange(0, len(edge_map)))) == edge_map)
def verify_graph_feats(g, part, node_feats):
for ntype in g.ntypes:
ntype_id = g.get_ntype_id(ntype)
for name in g.nodes[ntype].data:
if name in [dgl.NID, 'inner_node']:
continue
inner_node_mask = _get_inner_node_mask(part, ntype_id)
inner_nids = F.boolean_mask(part.ndata[dgl.NID], inner_node_mask)
min_nids = F.min(inner_nids, 0)
orig_id = F.boolean_mask(part.ndata['orig_id'], inner_node_mask)
true_feats = F.gather_row(g.nodes[ntype].data[name], orig_id)
ndata = F.gather_row(node_feats[ntype + '/' + name],
inner_nids - min_nids)
assert np.all(F.asnumpy(ndata == true_feats))
def verify_graph_feats(g, gpb, part, node_feats, edge_feats):
for ntype in g.ntypes:
ntype_id = g.get_ntype_id(ntype)
inner_node_mask = _get_inner_node_mask(part, ntype_id)
inner_nids = F.boolean_mask(part.ndata[dgl.NID], inner_node_mask)
ntype_ids, inner_type_nids = gpb.map_to_per_ntype(inner_nids)
partid = gpb.nid2partid(inner_type_nids, ntype)
assert np.all(F.asnumpy(ntype_ids) == ntype_id)
assert np.all(F.asnumpy(partid) == gpb.partid)
orig_id = F.boolean_mask(part.ndata['orig_id'], inner_node_mask)
local_nids = gpb.nid2localnid(inner_type_nids, gpb.partid, ntype)
for name in g.nodes[ntype].data:
if name in [dgl.NID, 'inner_node']:
continue
true_feats = F.gather_row(g.nodes[ntype].data[name], orig_id)
ndata = F.gather_row(node_feats[ntype + '/' + name], local_nids)
assert np.all(F.asnumpy(ndata == true_feats))
for etype in g.etypes:
etype_id = g.get_etype_id(etype)
inner_edge_mask = _get_inner_edge_mask(part, etype_id)
inner_eids = F.boolean_mask(part.edata[dgl.EID], inner_edge_mask)
etype_ids, inner_type_eids = gpb.map_to_per_etype(inner_eids)
partid = gpb.eid2partid(inner_type_eids, etype)
assert np.all(F.asnumpy(etype_ids) == etype_id)
assert np.all(F.asnumpy(partid) == gpb.partid)
orig_id = F.boolean_mask(part.edata['orig_id'], inner_edge_mask)
local_eids = gpb.eid2localeid(inner_type_eids, gpb.partid, etype)
for name in g.edges[etype].data:
if name in [dgl.EID, 'inner_edge']:
continue
true_feats = F.gather_row(g.edges[etype].data[name], orig_id)
edata = F.gather_row(edge_feats[etype + '/' + name], local_eids)
assert np.all(F.asnumpy(edata == true_feats))
def check_partition(part_method, reshuffle):
g = create_random_graph(10000)
g.ndata['labels'] = F.arange(0, g.number_of_nodes())
g.ndata['feats'] = F.tensor(np.random.randn(g.number_of_nodes(), 10))
g.edata['feats'] = F.tensor(np.random.randn(g.number_of_edges(), 10))
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
partition_graph(g, 'test', num_parts, '/tmp/partition', num_hops=num_hops,
part_method=part_method, reshuffle=reshuffle)
part_sizes = []
for i in range(num_parts):
part_g, node_feats, edge_feats, gpb, _ = 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']))
local_nid = gpb.nid2localnid(F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node']), i)
assert F.dtype(local_nid) in (F.int64, F.int32)
assert np.all(F.asnumpy(local_nid) == np.arange(0, len(local_nid)))
local_eid = gpb.eid2localeid(F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge']), 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)))
# Check the edge map.
local_edges = F.boolean_mask(part_g.edata[dgl.EID], 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)))
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 name in node_feats
assert node_feats[name].shape[0] == len(local_nodes)
assert np.all(F.asnumpy(g.ndata[name])[F.asnumpy(local_nodes)] == F.asnumpy(node_feats[name]))
for name in ['feats']:
assert name in edge_feats
assert edge_feats[name].shape[0] == len(local_edges)
assert np.all(F.asnumpy(g.edata[name])[F.asnumpy(local_edges)] == F.asnumpy(edge_feats[name]))
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 verify_hetero_graph(g, parts):
num_nodes = {ntype: 0 for ntype in g.ntypes}
num_edges = {etype: 0 for etype in g.etypes}
for part in parts:
assert len(g.ntypes) == len(F.unique(part.ndata[dgl.NTYPE]))
assert len(g.etypes) == len(F.unique(part.edata[dgl.ETYPE]))
for ntype in g.ntypes:
ntype_id = g.get_ntype_id(ntype)
inner_node_mask = _get_inner_node_mask(part, ntype_id)
num_inner_nodes = F.sum(F.astype(inner_node_mask, F.int64), 0)
num_nodes[ntype] += num_inner_nodes
for etype in g.etypes:
etype_id = g.get_etype_id(etype)
inner_edge_mask = _get_inner_edge_mask(part, etype_id)
num_inner_edges = F.sum(F.astype(inner_edge_mask, F.int64), 0)
num_edges[etype] += num_inner_edges
# Verify the number of nodes are correct.
for ntype in g.ntypes:
print('node {}: {}, {}'.format(ntype, g.number_of_nodes(ntype),
num_nodes[ntype]))
assert g.number_of_nodes(ntype) == num_nodes[ntype]
# Verify the number of edges are correct.
for etype in g.etypes:
print('edge {}: {}, {}'.format(etype, g.number_of_edges(etype),
num_edges[etype]))
assert g.number_of_edges(etype) == num_edges[etype]
nids = {ntype: [] for ntype in g.ntypes}
eids = {etype: [] for etype in g.etypes}
for part in parts:
src, dst, eid = part.edges(form='all')
orig_src = F.gather_row(part.ndata['orig_id'], src)
orig_dst = F.gather_row(part.ndata['orig_id'], dst)
orig_eid = F.gather_row(part.edata['orig_id'], eid)
etype_arr = F.gather_row(part.edata[dgl.ETYPE], eid)
eid_type = F.gather_row(part.edata[dgl.EID], eid)
for etype in g.etypes:
etype_id = g.get_etype_id(etype)
src1 = F.boolean_mask(orig_src, etype_arr == etype_id)
dst1 = F.boolean_mask(orig_dst, etype_arr == etype_id)
eid1 = F.boolean_mask(orig_eid, etype_arr == etype_id)
exist = g.has_edges_between(src1, dst1, etype=etype)
assert np.all(F.asnumpy(exist))
eid2 = g.edge_ids(src1, dst1, etype=etype)
assert np.all(F.asnumpy(eid1 == eid2))
eids[etype].append(F.boolean_mask(eid_type, etype_arr == etype_id))
# Make sure edge Ids fall into a range.
inner_edge_mask = _get_inner_edge_mask(part, etype_id)
inner_eids = np.sort(
F.asnumpy(F.boolean_mask(part.edata[dgl.EID],
inner_edge_mask)))
assert np.all(
inner_eids == np.arange(inner_eids[0], inner_eids[-1] + 1))
for ntype in g.ntypes:
ntype_id = g.get_ntype_id(ntype)
# Make sure inner nodes have Ids fall into a range.
inner_node_mask = _get_inner_node_mask(part, ntype_id)
inner_nids = F.boolean_mask(part.ndata[dgl.NID], inner_node_mask)
assert np.all(
F.asnumpy(
inner_nids == F.arange(F.as_scalar(inner_nids[0]),
F.as_scalar(inner_nids[-1]) + 1)))
nids[ntype].append(inner_nids)
for ntype in nids:
nids_type = F.cat(nids[ntype], 0)
uniq_ids = F.unique(nids_type)
# We should get all nodes.
assert len(uniq_ids) == g.number_of_nodes(ntype)
for etype in eids:
eids_type = F.cat(eids[etype], 0)
uniq_ids = F.unique(eids_type)
assert len(uniq_ids) == g.number_of_edges(etype)
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 check_hetero_partition(hg,
part_method,
num_parts=4,
num_trainers_per_machine=1,
load_feats=True):
hg.nodes['n1'].data['labels'] = F.arange(0, hg.number_of_nodes('n1'))
hg.nodes['n1'].data['feats'] = F.tensor(
np.random.randn(hg.number_of_nodes('n1'), 10), F.float32)
hg.edges['r1'].data['feats'] = F.tensor(
np.random.randn(hg.number_of_edges('r1'), 10), F.float32)
hg.edges['r1'].data['labels'] = F.arange(0, hg.number_of_edges('r1'))
num_hops = 1
orig_nids, orig_eids = partition_graph(
hg,
'test',
num_parts,
'/tmp/partition',
num_hops=num_hops,
part_method=part_method,
reshuffle=True,
return_mapping=True,
num_trainers_per_machine=num_trainers_per_machine)
assert len(orig_nids) == len(hg.ntypes)
assert len(orig_eids) == len(hg.etypes)
for ntype in hg.ntypes:
assert len(orig_nids[ntype]) == hg.number_of_nodes(ntype)
for etype in hg.etypes:
assert len(orig_eids[etype]) == hg.number_of_edges(etype)
parts = []
shuffled_labels = []
shuffled_elabels = []
for i in range(num_parts):
part_g, node_feats, edge_feats, gpb, _, ntypes, etypes = load_partition(
'/tmp/partition/test.json', i, load_feats=load_feats)
if not load_feats:
assert not node_feats
assert not edge_feats
node_feats, edge_feats = load_partition_feats(
'/tmp/partition/test.json', i)
if num_trainers_per_machine > 1:
for ntype in hg.ntypes:
name = ntype + '/trainer_id'
assert name in node_feats
part_ids = F.floor_div(node_feats[name],
num_trainers_per_machine)
assert np.all(F.asnumpy(part_ids) == i)
for etype in hg.etypes:
name = etype + '/trainer_id'
assert name in edge_feats
part_ids = F.floor_div(edge_feats[name],
num_trainers_per_machine)
assert np.all(F.asnumpy(part_ids) == i)
# Verify the mapping between the reshuffled IDs and the original IDs.
# These are partition-local IDs.
part_src_ids, part_dst_ids = part_g.edges()
# These are reshuffled global homogeneous IDs.
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]
# These are reshuffled per-type IDs.
src_ntype_ids, part_src_ids = gpb.map_to_per_ntype(part_src_ids)
dst_ntype_ids, part_dst_ids = gpb.map_to_per_ntype(part_dst_ids)
etype_ids, part_eids = gpb.map_to_per_etype(part_eids)
# These are original per-type IDs.
for etype_id, etype in enumerate(hg.etypes):
part_src_ids1 = F.boolean_mask(part_src_ids, etype_ids == etype_id)
src_ntype_ids1 = F.boolean_mask(src_ntype_ids,
etype_ids == etype_id)
part_dst_ids1 = F.boolean_mask(part_dst_ids, etype_ids == etype_id)
dst_ntype_ids1 = F.boolean_mask(dst_ntype_ids,
etype_ids == etype_id)
part_eids1 = F.boolean_mask(part_eids, etype_ids == etype_id)
assert np.all(F.asnumpy(src_ntype_ids1 == src_ntype_ids1[0]))
assert np.all(F.asnumpy(dst_ntype_ids1 == dst_ntype_ids1[0]))
src_ntype = hg.ntypes[F.as_scalar(src_ntype_ids1[0])]
dst_ntype = hg.ntypes[F.as_scalar(dst_ntype_ids1[0])]
orig_src_ids1 = F.gather_row(orig_nids[src_ntype], part_src_ids1)
orig_dst_ids1 = F.gather_row(orig_nids[dst_ntype], part_dst_ids1)
orig_eids1 = F.gather_row(orig_eids[etype], part_eids1)
orig_eids2 = hg.edge_ids(orig_src_ids1, orig_dst_ids1, etype=etype)
assert len(orig_eids1) == len(orig_eids2)
assert np.all(F.asnumpy(orig_eids1) == F.asnumpy(orig_eids2))
parts.append(part_g)
verify_graph_feats(hg, gpb, part_g, node_feats, edge_feats)
shuffled_labels.append(node_feats['n1/labels'])
shuffled_elabels.append(edge_feats['r1/labels'])
verify_hetero_graph(hg, parts)
shuffled_labels = F.asnumpy(F.cat(shuffled_labels, 0))
shuffled_elabels = F.asnumpy(F.cat(shuffled_elabels, 0))
orig_labels = np.zeros(shuffled_labels.shape, dtype=shuffled_labels.dtype)
orig_elabels = np.zeros(shuffled_elabels.shape,
dtype=shuffled_elabels.dtype)
orig_labels[F.asnumpy(orig_nids['n1'])] = shuffled_labels
orig_elabels[F.asnumpy(orig_eids['r1'])] = shuffled_elabels
assert np.all(orig_labels == F.asnumpy(hg.nodes['n1'].data['labels']))
assert np.all(orig_elabels == F.asnumpy(hg.edges['r1'].data['labels']))
def check_rpc_bipartite_etype_sampling_shuffle(tmpdir, num_server):
"""sample on bipartite via sample_etype_neighbors() which yields non-empty sample results"""
generate_ip_config("rpc_ip_config.txt", num_server, num_server)
g = create_random_bipartite()
num_parts = num_server
num_hops = 1
orig_nids, _ = partition_graph(g,
'test_sampling',
num_parts,
tmpdir,
num_hops=num_hops,
part_method='metis',
reshuffle=True,
return_mapping=True)
pserver_list = []
ctx = mp.get_context('spawn')
for i in range(num_server):
p = ctx.Process(target=start_server,
args=(i, tmpdir, num_server > 1, 'test_sampling'))
p.start()
time.sleep(1)
pserver_list.append(p)
fanout = 3
deg = get_degrees(g, orig_nids['game'], 'game')
nids = F.nonzero_1d(deg > 0)
block, gpb = start_bipartite_etype_sample_client(0,
tmpdir,
num_server > 1,
fanout,
nodes={
'game': nids,
'user': [0]
})
print("Done sampling")
for p in pserver_list:
p.join()
orig_nid_map = {
ntype: F.zeros((g.number_of_nodes(ntype), ), dtype=F.int64)
for ntype in g.ntypes
}
orig_eid_map = {
etype: F.zeros((g.number_of_edges(etype), ), dtype=F.int64)
for etype in g.etypes
}
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
i)
ntype_ids, type_nids = gpb.map_to_per_ntype(part.ndata[dgl.NID])
for ntype_id, ntype in enumerate(g.ntypes):
idx = ntype_ids == ntype_id
F.scatter_row_inplace(orig_nid_map[ntype],
F.boolean_mask(type_nids, idx),
F.boolean_mask(part.ndata['orig_id'], idx))
etype_ids, type_eids = gpb.map_to_per_etype(part.edata[dgl.EID])
for etype_id, etype in enumerate(g.etypes):
idx = etype_ids == etype_id
F.scatter_row_inplace(orig_eid_map[etype],
F.boolean_mask(type_eids, idx),
F.boolean_mask(part.edata['orig_id'], idx))
for src_type, etype, dst_type in block.canonical_etypes:
src, dst = block.edges(etype=etype)
# These are global Ids after shuffling.
shuffled_src = F.gather_row(block.srcnodes[src_type].data[dgl.NID],
src)
shuffled_dst = F.gather_row(block.dstnodes[dst_type].data[dgl.NID],
dst)
shuffled_eid = block.edges[etype].data[dgl.EID]
orig_src = F.asnumpy(F.gather_row(orig_nid_map[src_type],
shuffled_src))
orig_dst = F.asnumpy(F.gather_row(orig_nid_map[dst_type],
shuffled_dst))
orig_eid = F.asnumpy(F.gather_row(orig_eid_map[etype], shuffled_eid))
# Check the node Ids and edge Ids.
orig_src1, orig_dst1 = g.find_edges(orig_eid, etype=etype)
assert np.all(F.asnumpy(orig_src1) == orig_src)
assert np.all(F.asnumpy(orig_dst1) == orig_dst)
