def start_get_degrees_client(rank, tmpdir, disable_shared_mem, nids=None):
gpb = None
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_get_degrees.json', rank)
dgl.distributed.initialize("rpc_ip_config.txt", 1)
dist_graph = DistGraph("test_get_degrees", gpb=gpb)
try:
in_deg = dist_graph.in_degrees(nids)
all_in_deg = dist_graph.in_degrees()
out_deg = dist_graph.out_degrees(nids)
all_out_deg = dist_graph.out_degrees()
except Exception as e:
print(e)
in_deg, out_deg, all_in_deg, all_out_deg = None, None, None, None
dgl.distributed.exit_client()
return in_deg, out_deg, all_in_deg, all_out_deg
def check_rpc_get_degree_shuffle(tmpdir, num_server):
ip_config = open("rpc_ip_config.txt", "w")
for _ in range(num_server):
ip_config.write('{}\n'.format(get_local_usable_addr()))
ip_config.close()
g = CitationGraphDataset("cora")[0]
g.readonly()
num_parts = num_server
partition_graph(g,
'test_get_degrees',
num_parts,
tmpdir,
num_hops=1,
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_get_degrees'))
p.start()
time.sleep(1)
pserver_list.append(p)
orig_nid = F.zeros((g.number_of_nodes(), ), dtype=F.int64, ctx=F.cpu())
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(
tmpdir / 'test_get_degrees.json', i)
orig_nid[part.ndata[dgl.NID]] = part.ndata['orig_id']
time.sleep(3)
nids = F.tensor(np.random.randint(g.number_of_nodes(), size=100))
in_degs, out_degs, all_in_degs, all_out_degs = start_get_degrees_client(
0, tmpdir, num_server > 1, nids)
print("Done get_degree")
for p in pserver_list:
p.join()
print('check results')
assert F.array_equal(g.in_degrees(orig_nid[nids]), in_degs)
assert F.array_equal(g.in_degrees(orig_nid), all_in_degs)
assert F.array_equal(g.out_degrees(orig_nid[nids]), out_degs)
assert F.array_equal(g.out_degrees(orig_nid), all_out_degs)
def start_bipartite_sample_client(rank, tmpdir, disable_shared_mem, nodes):
gpb = None
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
rank)
dgl.distributed.initialize("rpc_ip_config.txt")
dist_graph = DistGraph("test_sampling", gpb=gpb)
assert 'feat' in dist_graph.nodes['user'].data
assert 'feat' in dist_graph.nodes['game'].data
if gpb is None:
gpb = dist_graph.get_partition_book()
sampled_graph = sample_neighbors(dist_graph, nodes, 3)
block = dgl.to_block(sampled_graph, nodes)
if sampled_graph.num_edges() > 0:
block.edata[dgl.EID] = sampled_graph.edata[dgl.EID]
dgl.distributed.exit_client()
return block, gpb
def start_dist_dataloader(rank, tmpdir, disable_shared_mem, num_workers, drop_last):
import dgl
import torch as th
dgl.distributed.initialize("mp_ip_config.txt", 1, num_workers=num_workers)
gpb = None
if disable_shared_mem:
_, _, _, gpb, _ = load_partition(tmpdir / 'test_sampling.json', rank)
num_nodes_to_sample = 202
batch_size = 32
train_nid = th.arange(num_nodes_to_sample)
dist_graph = DistGraph("test_mp", gpb=gpb, part_config=tmpdir / 'test_sampling.json')
# Create sampler
sampler = NeighborSampler(dist_graph, [5, 10],
dgl.distributed.sample_neighbors)
# We need to test creating DistDataLoader multiple times.
for i in range(2):
# Create DataLoader for constructing blocks
dataloader = DistDataLoader(
dataset=train_nid.numpy(),
batch_size=batch_size,
collate_fn=sampler.sample_blocks,
shuffle=False,
drop_last=drop_last)
groundtruth_g = CitationGraphDataset("cora")[0]
max_nid = []
for epoch in range(2):
for idx, blocks in zip(range(0, num_nodes_to_sample, batch_size), dataloader):
block = blocks[-1]
o_src, o_dst = block.edges()
src_nodes_id = block.srcdata[dgl.NID][o_src]
dst_nodes_id = block.dstdata[dgl.NID][o_dst]
has_edges = groundtruth_g.has_edges_between(src_nodes_id, dst_nodes_id)
assert np.all(F.asnumpy(has_edges))
max_nid.append(np.max(F.asnumpy(dst_nodes_id)))
# assert np.all(np.unique(np.sort(F.asnumpy(dst_nodes_id))) == np.arange(idx, batch_size))
if drop_last:
assert np.max(max_nid) == num_nodes_to_sample - 1 - num_nodes_to_sample % batch_size
else:
assert np.max(max_nid) == num_nodes_to_sample - 1
del dataloader
dgl.distributed.exit_client() # this is needed since there's two test here in one process
def start_find_edges_client(rank,
tmpdir,
disable_shared_mem,
eids,
etype=None):
gpb = None
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_find_edges.json',
rank)
dgl.distributed.initialize("rpc_ip_config.txt")
dist_graph = DistGraph("test_find_edges", gpb=gpb)
try:
u, v = dist_graph.find_edges(eids, etype=etype)
except Exception as e:
print(e)
u, v = None, None
dgl.distributed.exit_client()
return u, v
def check_rpc_in_subgraph_shuffle(tmpdir, num_server):
generate_ip_config("rpc_ip_config.txt", num_server, num_server)
g = CitationGraphDataset("cora")[0]
g.readonly()
num_parts = num_server
partition_graph(g, 'test_in_subgraph', num_parts, tmpdir,
num_hops=1, 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_in_subgraph'))
p.start()
time.sleep(1)
pserver_list.append(p)
nodes = [0, 10, 99, 66, 1024, 2008]
sampled_graph = start_in_subgraph_client(0, tmpdir, num_server > 1, nodes)
for p in pserver_list:
p.join()
orig_nid = F.zeros((g.number_of_nodes(),), dtype=F.int64, ctx=F.cpu())
orig_eid = F.zeros((g.number_of_edges(),), dtype=F.int64, ctx=F.cpu())
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_in_subgraph.json', i)
orig_nid[part.ndata[dgl.NID]] = part.ndata['orig_id']
orig_eid[part.edata[dgl.EID]] = part.edata['orig_id']
src, dst = sampled_graph.edges()
src = orig_nid[src]
dst = orig_nid[dst]
assert sampled_graph.number_of_nodes() == g.number_of_nodes()
assert np.all(F.asnumpy(g.has_edges_between(src, dst)))
subg1 = dgl.in_subgraph(g, orig_nid[nodes])
src1, dst1 = subg1.edges()
assert np.all(np.sort(F.asnumpy(src)) == np.sort(F.asnumpy(src1)))
assert np.all(np.sort(F.asnumpy(dst)) == np.sort(F.asnumpy(dst1)))
eids = g.edge_ids(src, dst)
eids1 = orig_eid[sampled_graph.edata[dgl.EID]]
assert np.array_equal(F.asnumpy(eids1), F.asnumpy(eids))
def test_split():
prepare_dist()
g = create_random_graph(10000)
num_parts = 4
num_hops = 2
partition_graph(g, 'dist_graph_test', num_parts, '/tmp/dist_graph', num_hops=num_hops, part_method='metis')
node_mask = np.random.randint(0, 100, size=g.number_of_nodes()) > 30
edge_mask = np.random.randint(0, 100, size=g.number_of_edges()) > 30
selected_nodes = np.nonzero(node_mask)[0]
selected_edges = np.nonzero(edge_mask)[0]
for i in range(num_parts):
dgl.distributed.set_num_client(num_parts)
part_g, node_feats, edge_feats, gpb = load_partition('/tmp/dist_graph/dist_graph_test.json', i)
local_nids = F.nonzero_1d(part_g.ndata['inner_node'])
local_nids = F.gather_row(part_g.ndata[dgl.NID], local_nids)
nodes1 = np.intersect1d(selected_nodes, F.asnumpy(local_nids))
nodes2 = node_split(node_mask, gpb, i)
assert np.all(np.sort(nodes1) == np.sort(F.asnumpy(nodes2)))
local_nids = F.asnumpy(local_nids)
for n in nodes1:
assert n in local_nids
dgl.distributed.set_num_client(num_parts * 2)
nodes3 = node_split(node_mask, gpb, i * 2)
nodes4 = node_split(node_mask, gpb, i * 2 + 1)
nodes5 = F.cat([nodes3, nodes4], 0)
assert np.all(np.sort(nodes1) == np.sort(F.asnumpy(nodes5)))
dgl.distributed.set_num_client(num_parts)
local_eids = F.nonzero_1d(part_g.edata['inner_edge'])
local_eids = F.gather_row(part_g.edata[dgl.EID], local_eids)
edges1 = np.intersect1d(selected_edges, F.asnumpy(local_eids))
edges2 = edge_split(edge_mask, gpb, i)
assert np.all(np.sort(edges1) == np.sort(F.asnumpy(edges2)))
local_eids = F.asnumpy(local_eids)
for e in edges1:
assert e in local_eids
dgl.distributed.set_num_client(num_parts * 2)
edges3 = edge_split(edge_mask, gpb, i * 2)
edges4 = edge_split(edge_mask, gpb, i * 2 + 1)
edges5 = F.cat([edges3, edges4], 0)
assert np.all(np.sort(edges1) == np.sort(F.asnumpy(edges5)))
def test_partition():
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',
num_hops=num_hops,
part_method='metis')
for i in range(num_parts):
part_g, node_feats, edge_feats, meta = load_partition(
'/tmp/test.json', i)
num_nodes, num_edges, node_map, edge_map, num_partitions = meta
# Check the metadata
assert num_nodes == g.number_of_nodes()
assert num_edges == g.number_of_edges()
assert num_partitions == num_parts
# Check the node map.
local_nodes = np.nonzero(node_map == i)[0]
part_ids = node_map[F.asnumpy(part_g.ndata[dgl.NID])]
local_nodes1 = F.asnumpy(part_g.ndata[dgl.NID])[part_ids == i]
assert np.all(local_nodes == local_nodes1)
# Check the edge map.
assert np.all(edge_map >= 0)
local_edges = np.nonzero(edge_map == i)[0]
part_ids = edge_map[F.asnumpy(part_g.edata[dgl.EID])]
local_edges1 = F.asnumpy(part_g.edata[dgl.EID])[part_ids == i]
assert np.all(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
def check_rpc_find_edges_shuffle(tmpdir, num_server):
ip_config = open("rpc_ip_config.txt", "w")
for _ in range(num_server):
ip_config.write('{}\n'.format(get_local_usable_addr()))
ip_config.close()
g = CitationGraphDataset("cora")[0]
g.readonly()
num_parts = num_server
partition_graph(g,
'test_find_edges',
num_parts,
tmpdir,
num_hops=1,
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_find_edges',
['csr', 'coo']))
p.start()
time.sleep(1)
pserver_list.append(p)
orig_nid = F.zeros((g.number_of_nodes(), ), dtype=F.int64, ctx=F.cpu())
orig_eid = F.zeros((g.number_of_edges(), ), dtype=F.int64, ctx=F.cpu())
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(
tmpdir / 'test_find_edges.json', i)
orig_nid[part.ndata[dgl.NID]] = part.ndata['orig_id']
orig_eid[part.edata[dgl.EID]] = part.edata['orig_id']
time.sleep(3)
eids = F.tensor(np.random.randint(g.number_of_edges(), size=100))
u, v = g.find_edges(orig_eid[eids])
du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids)
du = orig_nid[du]
dv = orig_nid[dv]
assert F.array_equal(u, du)
assert F.array_equal(v, dv)
def check_rpc_sampling_shuffle(tmpdir, num_server):
ip_config = open("rpc_ip_config.txt", "w")
for _ in range(num_server):
ip_config.write('{} 1\n'.format(get_local_usable_addr()))
ip_config.close()
g = CitationGraphDataset("cora")[0]
g.readonly()
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)
time.sleep(3)
sampled_graph = start_sample_client(0, tmpdir, num_server > 1)
print("Done sampling")
for p in pserver_list:
p.join()
orig_nid = F.zeros((g.number_of_nodes(),), dtype=F.int64)
orig_eid = F.zeros((g.number_of_edges(),), dtype=F.int64)
for i in range(num_server):
part, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json', i)
orig_nid[part.ndata[dgl.NID]] = part.ndata['orig_id']
orig_eid[part.edata[dgl.EID]] = part.edata['orig_id']
src, dst = sampled_graph.edges()
src = orig_nid[src]
dst = orig_nid[dst]
assert sampled_graph.number_of_nodes() == g.number_of_nodes()
assert np.all(F.asnumpy(g.has_edges_between(src, dst)))
eids = g.edge_ids(src, dst)
eids1 = orig_eid[sampled_graph.edata[dgl.EID]]
assert np.array_equal(F.asnumpy(eids1), F.asnumpy(eids))
def test_split():
prepare_dist()
g = create_random_graph(10000)
num_parts = 4
num_hops = 2
partition_graph(g,
'test',
num_parts,
'/tmp',
num_hops=num_hops,
part_method='metis')
node_mask = np.random.randint(0, 100, size=g.number_of_nodes()) > 30
edge_mask = np.random.randint(0, 100, size=g.number_of_edges()) > 30
selected_nodes = np.nonzero(node_mask)[0]
selected_edges = np.nonzero(edge_mask)[0]
for i in range(num_parts):
part_g, node_feats, edge_feats, meta = load_partition(
'/tmp/test.json', i)
num_nodes, num_edges, node_map, edge_map, num_partitions = meta
gpb = GraphPartitionBook(part_id=i,
num_parts=num_partitions,
node_map=node_map,
edge_map=edge_map,
part_graph=part_g)
local_nids = F.nonzero_1d(part_g.ndata['local_node'])
local_nids = F.gather_row(part_g.ndata[dgl.NID], local_nids)
nodes1 = np.intersect1d(selected_nodes, F.asnumpy(local_nids))
nodes2 = node_split(node_mask, gpb, i)
assert np.all(np.sort(nodes1) == np.sort(F.asnumpy(nodes2)))
local_nids = F.asnumpy(local_nids)
for n in nodes1:
assert n in local_nids
local_eids = F.nonzero_1d(part_g.edata['local_edge'])
local_eids = F.gather_row(part_g.edata[dgl.EID], local_eids)
edges1 = np.intersect1d(selected_edges, F.asnumpy(local_eids))
edges2 = edge_split(edge_mask, gpb, i)
assert np.all(np.sort(edges1) == np.sort(F.asnumpy(edges2)))
local_eids = F.asnumpy(local_eids)
for e in edges1:
assert e in local_eids
def start_hetero_sample_client(rank, tmpdir, disable_shared_mem, nodes):
gpb = None
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json', rank)
dgl.distributed.initialize("rpc_ip_config.txt")
dist_graph = DistGraph("test_sampling", gpb=gpb)
assert 'feat' in dist_graph.nodes['n1'].data
assert 'feat' not in dist_graph.nodes['n2'].data
assert 'feat' not in dist_graph.nodes['n3'].data
if gpb is None:
gpb = dist_graph.get_partition_book()
try:
sampled_graph = sample_neighbors(dist_graph, nodes, 3)
block = dgl.to_block(sampled_graph, nodes)
block.edata[dgl.EID] = sampled_graph.edata[dgl.EID]
except Exception as e:
print(e)
block = None
dgl.distributed.exit_client()
return block, gpb
def start_hetero_etype_sample_client(rank,
tmpdir,
disable_shared_mem,
fanout=3,
nodes={'n3': [0, 10, 99, 66, 124, 208]},
etype_sorted=False):
gpb = None
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
rank)
dgl.distributed.initialize("rpc_ip_config.txt")
dist_graph = DistGraph("test_sampling", gpb=gpb)
assert 'feat' in dist_graph.nodes['n1'].data
assert 'feat' not in dist_graph.nodes['n2'].data
assert 'feat' not in dist_graph.nodes['n3'].data
if dist_graph.local_partition is not None:
# Check whether etypes are sorted in dist_graph
local_g = dist_graph.local_partition
local_nids = np.arange(local_g.num_nodes())
for lnid in local_nids:
leids = local_g.in_edges(lnid, form='eid')
letids = F.asnumpy(local_g.edata[dgl.ETYPE][leids])
_, idices = np.unique(letids, return_index=True)
assert np.all(idices[:-1] <= idices[1:])
if gpb is None:
gpb = dist_graph.get_partition_book()
try:
sampled_graph = sample_etype_neighbors(dist_graph,
nodes,
dgl.ETYPE,
fanout,
etype_sorted=etype_sorted)
block = dgl.to_block(sampled_graph, nodes)
block.edata[dgl.EID] = sampled_graph.edata[dgl.EID]
except Exception as e:
print(e)
block = None
dgl.distributed.exit_client()
return block, gpb
def test_graph_partition_book():
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
create_ip_config()
partition_graph(g,
'test',
num_parts,
'/tmp',
num_hops=num_hops,
part_method='metis')
for i in range(num_parts):
part_g, node_feats, edge_feats, meta = load_partition(
'/tmp/test.json', i)
num_nodes, num_edges, node_map, edge_map, num_partitions = meta
gpb = GraphPartitionBook(part_id=i,
num_parts=num_partitions,
node_map=node_map,
edge_map=edge_map,
part_graph=part_g)
assert gpb.num_partitions() == num_parts
gpb_meta = gpb.metadata()
assert len(gpb_meta) == num_parts
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)
assert len(gpb.partid2nids(i)) == gpb_meta[i]['num_nodes']
assert len(gpb.partid2eids(i)) == gpb_meta[i]['num_edges']
local_nid = gpb.nid2localnid(part_g.ndata[dgl.NID], i)
assert np.all(
F.asnumpy(local_nid) == F.asnumpy(F.arange(0, len(local_nid))))
local_eid = gpb.eid2localeid(part_g.edata[dgl.EID], i)
assert np.all(
F.asnumpy(local_eid) == F.asnumpy(F.arange(0, len(local_eid))))
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
partition_graph(hg,
'test',
num_parts,
'/tmp/partition',
num_hops=num_hops,
part_method=part_method,
reshuffle=True)
parts = []
for i in range(num_parts):
part_g, node_feats, edge_feats, gpb, _, ntypes, etypes = load_partition(
'/tmp/partition/test.json', i)
parts.append(part_g)
verify_graph_feats(hg, part_g, node_feats)
verify_hetero_graph(hg, parts)
def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server):
ip_config = open("rpc_ip_config.txt", "w")
for _ in range(num_server):
ip_config.write('{}\n'.format(get_local_usable_addr()))
ip_config.close()
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)
time.sleep(3)
fanout = 3
block, gpb = start_hetero_etype_sample_client(0, tmpdir, num_server > 1,
fanout)
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_rpc_hetero_sampling_shuffle(tmpdir, num_server):
ip_config = open("rpc_ip_config.txt", "w")
for _ in range(num_server):
ip_config.write('{}\n'.format(get_local_usable_addr()))
ip_config.close()
g = create_random_hetero()
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)
time.sleep(3)
block, gpb = start_hetero_sample_client(0, tmpdir, num_server > 1)
print("Done sampling")
for p in pserver_list:
p.join()
orig_nid_map = F.zeros((g.number_of_nodes(), ), dtype=F.int64)
orig_eid_map = F.zeros((g.number_of_edges(), ), dtype=F.int64)
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
i)
F.scatter_row_inplace(orig_nid_map, part.ndata[dgl.NID],
part.ndata['orig_id'])
F.scatter_row_inplace(orig_eid_map, part.edata[dgl.EID],
part.edata['orig_id'])
src, dst = block.edges()
# These are global Ids after shuffling.
shuffled_src = F.gather_row(block.srcdata[dgl.NID], src)
shuffled_dst = F.gather_row(block.dstdata[dgl.NID], dst)
shuffled_eid = block.edata[dgl.EID]
# Get node/edge types.
etype, _ = gpb.map_to_per_etype(shuffled_eid)
src_type, _ = gpb.map_to_per_ntype(shuffled_src)
dst_type, _ = gpb.map_to_per_ntype(shuffled_dst)
etype = F.asnumpy(etype)
src_type = F.asnumpy(src_type)
dst_type = F.asnumpy(dst_type)
# These are global Ids in the original graph.
orig_src = F.asnumpy(F.gather_row(orig_nid_map, shuffled_src))
orig_dst = F.asnumpy(F.gather_row(orig_nid_map, shuffled_dst))
orig_eid = F.asnumpy(F.gather_row(orig_eid_map, shuffled_eid))
etype_map = {g.get_etype_id(etype): etype for etype in g.etypes}
etype_to_eptype = {
g.get_etype_id(etype): (src_ntype, dst_ntype)
for src_ntype, etype, dst_ntype in g.canonical_etypes
}
for e in np.unique(etype):
src_t = src_type[etype == e]
dst_t = dst_type[etype == e]
assert np.all(src_t == src_t[0])
assert np.all(dst_t == dst_t[0])
# Check the node Ids and edge Ids.
orig_src1, orig_dst1 = g.find_edges(orig_eid[etype == e],
etype=etype_map[e])
assert np.all(F.asnumpy(orig_src1) == orig_src[etype == e])
assert np.all(F.asnumpy(orig_dst1) == orig_dst[etype == e])
# Check the node types.
src_ntype, dst_ntype = etype_to_eptype[e]
assert np.all(src_t == g.get_ntype_id(src_ntype))
assert np.all(dst_t == g.get_ntype_id(dst_ntype))
def start_edge_dataloader(rank, tmpdir, num_server, num_workers, orig_nid,
orig_eid, groundtruth_g):
import dgl
import torch as th
dgl.distributed.initialize("mp_ip_config.txt")
gpb = None
disable_shared_mem = num_server > 1
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
rank)
num_edges_to_sample = 202
batch_size = 32
dist_graph = DistGraph("test_mp",
gpb=gpb,
part_config=tmpdir / 'test_sampling.json')
assert len(dist_graph.ntypes) == len(groundtruth_g.ntypes)
assert len(dist_graph.etypes) == len(groundtruth_g.etypes)
if len(dist_graph.etypes) == 1:
train_eid = th.arange(num_edges_to_sample)
else:
train_eid = {dist_graph.etypes[0]: th.arange(num_edges_to_sample)}
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
i)
# Create sampler
sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10])
# We need to test creating DistDataLoader multiple times.
for i in range(2):
# Create DataLoader for constructing blocks
dataloader = dgl.dataloading.EdgeDataLoader(dist_graph,
train_eid,
sampler,
batch_size=batch_size,
shuffle=True,
drop_last=False,
num_workers=num_workers)
for epoch in range(2):
for idx, (input_nodes, pos_pair_graph,
blocks) in zip(range(0, num_edges_to_sample, batch_size),
dataloader):
block = blocks[-1]
for src_type, etype, dst_type in block.canonical_etypes:
o_src, o_dst = block.edges(etype=etype)
src_nodes_id = block.srcnodes[src_type].data[
dgl.NID][o_src]
dst_nodes_id = block.dstnodes[dst_type].data[
dgl.NID][o_dst]
src_nodes_id = orig_nid[src_type][src_nodes_id]
dst_nodes_id = orig_nid[dst_type][dst_nodes_id]
has_edges = groundtruth_g.has_edges_between(src_nodes_id,
dst_nodes_id,
etype=etype)
assert np.all(F.asnumpy(has_edges))
assert np.all(
F.asnumpy(block.dstnodes[dst_type].data[dgl.NID]) == F.
asnumpy(pos_pair_graph.nodes[dst_type].data[dgl.NID]))
# assert np.all(np.unique(np.sort(F.asnumpy(dst_nodes_id))) == np.arange(idx, batch_size))
del dataloader
dgl.distributed.exit_client(
) # this is needed since there's two test here in one process
def test_split_even():
g = create_random_graph(10000)
num_parts = 4
num_hops = 2
partition_graph(g,
'dist_graph_test',
num_parts,
'/tmp/dist_graph',
num_hops=num_hops,
part_method='metis')
node_mask = np.random.randint(0, 100, size=g.number_of_nodes()) > 30
edge_mask = np.random.randint(0, 100, size=g.number_of_edges()) > 30
selected_nodes = np.nonzero(node_mask)[0]
selected_edges = np.nonzero(edge_mask)[0]
all_nodes1 = []
all_nodes2 = []
all_edges1 = []
all_edges2 = []
# The code now collects the roles of all client processes and use the information
# to determine how to split the workloads. Here is to simulate the multi-client
# use case.
def set_roles(num_clients):
dgl.distributed.role.CUR_ROLE = 'default'
dgl.distributed.role.GLOBAL_RANK = {i: i for i in range(num_clients)}
dgl.distributed.role.PER_ROLE_RANK['default'] = {
i: i
for i in range(num_clients)
}
for i in range(num_parts):
set_roles(num_parts)
part_g, node_feats, edge_feats, gpb, _, _, _ = load_partition(
'/tmp/dist_graph/dist_graph_test.json', i)
local_nids = F.nonzero_1d(part_g.ndata['inner_node'])
local_nids = F.gather_row(part_g.ndata[dgl.NID], local_nids)
nodes = node_split(node_mask, gpb, rank=i, force_even=True)
all_nodes1.append(nodes)
subset = np.intersect1d(F.asnumpy(nodes), F.asnumpy(local_nids))
print('part {} get {} nodes and {} are in the partition'.format(
i, len(nodes), len(subset)))
set_roles(num_parts * 2)
nodes1 = node_split(node_mask, gpb, rank=i * 2, force_even=True)
nodes2 = node_split(node_mask, gpb, rank=i * 2 + 1, force_even=True)
nodes3, _ = F.sort_1d(F.cat([nodes1, nodes2], 0))
all_nodes2.append(nodes3)
subset = np.intersect1d(F.asnumpy(nodes), F.asnumpy(nodes3))
print('intersection has', len(subset))
set_roles(num_parts)
local_eids = F.nonzero_1d(part_g.edata['inner_edge'])
local_eids = F.gather_row(part_g.edata[dgl.EID], local_eids)
edges = edge_split(edge_mask, gpb, rank=i, force_even=True)
all_edges1.append(edges)
subset = np.intersect1d(F.asnumpy(edges), F.asnumpy(local_eids))
print('part {} get {} edges and {} are in the partition'.format(
i, len(edges), len(subset)))
set_roles(num_parts * 2)
edges1 = edge_split(edge_mask, gpb, rank=i * 2, force_even=True)
edges2 = edge_split(edge_mask, gpb, rank=i * 2 + 1, force_even=True)
edges3, _ = F.sort_1d(F.cat([edges1, edges2], 0))
all_edges2.append(edges3)
subset = np.intersect1d(F.asnumpy(edges), F.asnumpy(edges3))
print('intersection has', len(subset))
all_nodes1 = F.cat(all_nodes1, 0)
all_edges1 = F.cat(all_edges1, 0)
all_nodes2 = F.cat(all_nodes2, 0)
all_edges2 = F.cat(all_edges2, 0)
all_nodes = np.nonzero(node_mask)[0]
all_edges = np.nonzero(edge_mask)[0]
assert np.all(all_nodes == F.asnumpy(all_nodes1))
assert np.all(all_edges == F.asnumpy(all_edges1))
assert np.all(all_nodes == F.asnumpy(all_nodes2))
assert np.all(all_edges == F.asnumpy(all_edges2))
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 test_split(hetero):
if hetero:
g = create_random_hetero()
ntype = 'n1'
etype = 'r1'
else:
g = create_random_graph(10000)
ntype = '_N'
etype = '_E'
num_parts = 4
num_hops = 2
partition_graph(g,
'dist_graph_test',
num_parts,
'/tmp/dist_graph',
num_hops=num_hops,
part_method='metis')
node_mask = np.random.randint(0, 100, size=g.number_of_nodes(ntype)) > 30
edge_mask = np.random.randint(0, 100, size=g.number_of_edges(etype)) > 30
selected_nodes = np.nonzero(node_mask)[0]
selected_edges = np.nonzero(edge_mask)[0]
# The code now collects the roles of all client processes and use the information
# to determine how to split the workloads. Here is to simulate the multi-client
# use case.
def set_roles(num_clients):
dgl.distributed.role.CUR_ROLE = 'default'
dgl.distributed.role.GLOBAL_RANK = {i: i for i in range(num_clients)}
dgl.distributed.role.PER_ROLE_RANK['default'] = {
i: i
for i in range(num_clients)
}
for i in range(num_parts):
set_roles(num_parts)
part_g, node_feats, edge_feats, gpb, _, _, _ = load_partition(
'/tmp/dist_graph/dist_graph_test.json', i)
local_nids = F.nonzero_1d(part_g.ndata['inner_node'])
local_nids = F.gather_row(part_g.ndata[dgl.NID], local_nids)
if hetero:
ntype_ids, nids = gpb.map_to_per_ntype(local_nids)
local_nids = F.asnumpy(nids)[F.asnumpy(ntype_ids) == 0]
else:
local_nids = F.asnumpy(local_nids)
nodes1 = np.intersect1d(selected_nodes, local_nids)
nodes2 = node_split(node_mask,
gpb,
ntype=ntype,
rank=i,
force_even=False)
assert np.all(np.sort(nodes1) == np.sort(F.asnumpy(nodes2)))
for n in F.asnumpy(nodes2):
assert n in local_nids
set_roles(num_parts * 2)
nodes3 = node_split(node_mask,
gpb,
ntype=ntype,
rank=i * 2,
force_even=False)
nodes4 = node_split(node_mask,
gpb,
ntype=ntype,
rank=i * 2 + 1,
force_even=False)
nodes5 = F.cat([nodes3, nodes4], 0)
assert np.all(np.sort(nodes1) == np.sort(F.asnumpy(nodes5)))
set_roles(num_parts)
local_eids = F.nonzero_1d(part_g.edata['inner_edge'])
local_eids = F.gather_row(part_g.edata[dgl.EID], local_eids)
if hetero:
etype_ids, eids = gpb.map_to_per_etype(local_eids)
local_eids = F.asnumpy(eids)[F.asnumpy(etype_ids) == 0]
else:
local_eids = F.asnumpy(local_eids)
edges1 = np.intersect1d(selected_edges, local_eids)
edges2 = edge_split(edge_mask,
gpb,
etype=etype,
rank=i,
force_even=False)
assert np.all(np.sort(edges1) == np.sort(F.asnumpy(edges2)))
for e in F.asnumpy(edges2):
assert e in local_eids
set_roles(num_parts * 2)
edges3 = edge_split(edge_mask,
gpb,
etype=etype,
rank=i * 2,
force_even=False)
edges4 = edge_split(edge_mask,
gpb,
etype=etype,
rank=i * 2 + 1,
force_even=False)
edges5 = F.cat([edges3, edges4], 0)
assert np.all(np.sort(edges1) == np.sort(F.asnumpy(edges5)))
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):
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_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 = []
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)
# 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 test_split_even():
prepare_dist()
g = create_random_graph(10000)
num_parts = 4
num_hops = 2
partition_graph(g,
'dist_graph_test',
num_parts,
'/tmp/dist_graph',
num_hops=num_hops,
part_method='metis')
node_mask = np.random.randint(0, 100, size=g.number_of_nodes()) > 30
edge_mask = np.random.randint(0, 100, size=g.number_of_edges()) > 30
selected_nodes = np.nonzero(node_mask)[0]
selected_edges = np.nonzero(edge_mask)[0]
all_nodes1 = []
all_nodes2 = []
all_edges1 = []
all_edges2 = []
for i in range(num_parts):
dgl.distributed.set_num_client(num_parts)
part_g, node_feats, edge_feats, gpb, _ = load_partition(
'/tmp/dist_graph/dist_graph_test.json', i)
local_nids = F.nonzero_1d(part_g.ndata['inner_node'])
local_nids = F.gather_row(part_g.ndata[dgl.NID], local_nids)
nodes = node_split(node_mask, gpb, i, force_even=True)
all_nodes1.append(nodes)
subset = np.intersect1d(F.asnumpy(nodes), F.asnumpy(local_nids))
print('part {} get {} nodes and {} are in the partition'.format(
i, len(nodes), len(subset)))
dgl.distributed.set_num_client(num_parts * 2)
nodes1 = node_split(node_mask, gpb, i * 2, force_even=True)
nodes2 = node_split(node_mask, gpb, i * 2 + 1, force_even=True)
nodes3 = F.cat([nodes1, nodes2], 0)
all_nodes2.append(nodes3)
subset = np.intersect1d(F.asnumpy(nodes), F.asnumpy(nodes3))
print('intersection has', len(subset))
dgl.distributed.set_num_client(num_parts)
local_eids = F.nonzero_1d(part_g.edata['inner_edge'])
local_eids = F.gather_row(part_g.edata[dgl.EID], local_eids)
edges = edge_split(edge_mask, gpb, i, force_even=True)
all_edges1.append(edges)
subset = np.intersect1d(F.asnumpy(edges), F.asnumpy(local_eids))
print('part {} get {} edges and {} are in the partition'.format(
i, len(edges), len(subset)))
dgl.distributed.set_num_client(num_parts * 2)
edges1 = edge_split(edge_mask, gpb, i * 2, force_even=True)
edges2 = edge_split(edge_mask, gpb, i * 2 + 1, force_even=True)
edges3 = F.cat([edges1, edges2], 0)
all_edges2.append(edges3)
subset = np.intersect1d(F.asnumpy(edges), F.asnumpy(edges3))
print('intersection has', len(subset))
all_nodes1 = F.cat(all_nodes1, 0)
all_edges1 = F.cat(all_edges1, 0)
all_nodes2 = F.cat(all_nodes2, 0)
all_edges2 = F.cat(all_edges2, 0)
all_nodes = np.nonzero(node_mask)[0]
all_edges = np.nonzero(edge_mask)[0]
assert np.all(all_nodes == F.asnumpy(all_nodes1))
assert np.all(all_edges == F.asnumpy(all_edges1))
assert np.all(all_nodes == F.asnumpy(all_nodes2))
assert np.all(all_edges == F.asnumpy(all_edges2))
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)
def start_dist_neg_dataloader(rank, tmpdir, num_server, num_workers, orig_nid,
groundtruth_g):
import dgl
import torch as th
dgl.distributed.initialize("mp_ip_config.txt")
gpb = None
disable_shared_mem = num_server > 1
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
rank)
num_edges_to_sample = 202
batch_size = 32
dist_graph = DistGraph("test_mp",
gpb=gpb,
part_config=tmpdir / 'test_sampling.json')
assert len(dist_graph.ntypes) == len(groundtruth_g.ntypes)
assert len(dist_graph.etypes) == len(groundtruth_g.etypes)
if len(dist_graph.etypes) == 1:
train_eid = th.arange(num_edges_to_sample)
else:
train_eid = {dist_graph.etypes[0]: th.arange(num_edges_to_sample)}
for i in range(num_server):
part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json',
i)
num_negs = 5
sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10])
negative_sampler = dgl.dataloading.negative_sampler.Uniform(num_negs)
dataloader = dgl.dataloading.DistEdgeDataLoader(
dist_graph,
train_eid,
sampler,
batch_size=batch_size,
negative_sampler=negative_sampler,
shuffle=True,
drop_last=False,
num_workers=num_workers)
for _ in range(2):
for _, (_, pos_graph, neg_graph,
blocks) in zip(range(0, num_edges_to_sample, batch_size),
dataloader):
block = blocks[-1]
for src_type, etype, dst_type in block.canonical_etypes:
o_src, o_dst = block.edges(etype=etype)
src_nodes_id = block.srcnodes[src_type].data[dgl.NID][o_src]
dst_nodes_id = block.dstnodes[dst_type].data[dgl.NID][o_dst]
src_nodes_id = orig_nid[src_type][src_nodes_id]
dst_nodes_id = orig_nid[dst_type][dst_nodes_id]
has_edges = groundtruth_g.has_edges_between(src_nodes_id,
dst_nodes_id,
etype=etype)
assert np.all(F.asnumpy(has_edges))
assert np.all(
F.asnumpy(block.dstnodes[dst_type].data[dgl.NID]) ==
F.asnumpy(pos_graph.nodes[dst_type].data[dgl.NID]))
assert np.all(
F.asnumpy(block.dstnodes[dst_type].data[dgl.NID]) ==
F.asnumpy(neg_graph.nodes[dst_type].data[dgl.NID]))
assert pos_graph.num_edges() * num_negs == neg_graph.num_edges(
)
del dataloader
dgl.distributed.exit_client(
) # this is needed since there's two test here in one process
