wehub-resource-sync
1181 lines
38 KiB
Python
1181 lines
38 KiB
Python
import multiprocessing as mp
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
import uuid
|
|
|
|
import backend as F
|
|
import dgl
|
|
import numpy as np
|
|
import pytest
|
|
import torch as th
|
|
from dgl.data import CitationGraphDataset
|
|
from dgl.distributed import (
|
|
DistDataLoader,
|
|
DistGraph,
|
|
DistGraphServer,
|
|
load_partition,
|
|
partition_graph,
|
|
)
|
|
from scipy import sparse as spsp
|
|
from utils import generate_ip_config, reset_envs
|
|
|
|
|
|
def _unique_rand_graph(num_nodes=1000, num_edges=10 * 1000):
|
|
edges_set = set()
|
|
while len(edges_set) < num_edges:
|
|
src = np.random.randint(0, num_nodes - 1)
|
|
dst = np.random.randint(0, num_nodes - 1)
|
|
if (
|
|
src != dst
|
|
and (src, dst) not in edges_set
|
|
and (dst, src) not in edges_set
|
|
):
|
|
edges_set.add((src, dst))
|
|
src_list, dst_list = zip(*edges_set)
|
|
|
|
src = th.tensor(src_list, dtype=th.long)
|
|
dst = th.tensor(dst_list, dtype=th.long)
|
|
g = dgl.graph((th.cat([src, dst]), th.cat([dst, src])))
|
|
E = len(src)
|
|
reverse_eids = th.cat([th.arange(E, 2 * E), th.arange(0, E)])
|
|
return g, reverse_eids
|
|
|
|
|
|
class NeighborSampler(object):
|
|
def __init__(
|
|
self,
|
|
g,
|
|
fanouts,
|
|
sample_neighbors,
|
|
use_graphbolt=False,
|
|
return_eids=False,
|
|
):
|
|
self.g = g
|
|
self.fanouts = fanouts
|
|
self.sample_neighbors = sample_neighbors
|
|
self.use_graphbolt = use_graphbolt
|
|
self.return_eids = return_eids
|
|
|
|
def sample_blocks(self, seeds):
|
|
import torch as th
|
|
|
|
seeds = th.tensor(np.asarray(seeds), dtype=self.g.idtype)
|
|
blocks = []
|
|
for fanout in self.fanouts:
|
|
# For each seed node, sample ``fanout`` neighbors.
|
|
frontier = self.sample_neighbors(
|
|
self.g, seeds, fanout, use_graphbolt=self.use_graphbolt
|
|
)
|
|
# Then we compact the frontier into a bipartite graph for
|
|
# message passing.
|
|
block = dgl.to_block(frontier, seeds)
|
|
# Obtain the seed nodes for next layer.
|
|
seeds = block.srcdata[dgl.NID]
|
|
if frontier.num_edges() > 0:
|
|
if not self.use_graphbolt or self.return_eids:
|
|
block.edata[dgl.EID] = frontier.edata[dgl.EID]
|
|
|
|
blocks.insert(0, block)
|
|
return blocks
|
|
|
|
|
|
def start_server(
|
|
rank,
|
|
ip_config,
|
|
part_config,
|
|
disable_shared_mem,
|
|
num_clients,
|
|
use_graphbolt=False,
|
|
):
|
|
print("server: #clients=" + str(num_clients))
|
|
g = DistGraphServer(
|
|
rank,
|
|
ip_config,
|
|
1,
|
|
num_clients,
|
|
part_config,
|
|
disable_shared_mem=disable_shared_mem,
|
|
graph_format=["csc", "coo"],
|
|
use_graphbolt=use_graphbolt,
|
|
)
|
|
g.start()
|
|
|
|
|
|
def start_dist_dataloader(
|
|
rank,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
drop_last,
|
|
orig_nid,
|
|
orig_eid,
|
|
use_graphbolt=False,
|
|
return_eids=False,
|
|
):
|
|
dgl.distributed.initialize(ip_config)
|
|
gpb = None
|
|
disable_shared_mem = num_server > 1
|
|
if disable_shared_mem:
|
|
_, _, _, gpb, _, _, _ = load_partition(part_config, rank)
|
|
num_nodes_to_sample = 202
|
|
batch_size = 32
|
|
train_nid = th.arange(num_nodes_to_sample)
|
|
graph_name = os.path.splitext(os.path.basename(part_config))[0]
|
|
dist_graph = DistGraph(
|
|
graph_name,
|
|
gpb=gpb,
|
|
part_config=part_config,
|
|
)
|
|
|
|
# Create sampler
|
|
sampler = NeighborSampler(
|
|
dist_graph,
|
|
[5, 10],
|
|
dgl.distributed.sample_neighbors,
|
|
use_graphbolt=use_graphbolt,
|
|
return_eids=return_eids,
|
|
)
|
|
|
|
# Enable santity check in distributed sampling.
|
|
os.environ["DGL_DIST_DEBUG"] = "1"
|
|
|
|
# We need to test creating DistDataLoader multiple times.
|
|
for i in range(2):
|
|
# Create DataLoader for constructing blocks
|
|
dataloader = DistDataLoader(
|
|
dataset=train_nid,
|
|
batch_size=batch_size,
|
|
collate_fn=sampler.sample_blocks,
|
|
shuffle=False,
|
|
drop_last=drop_last,
|
|
)
|
|
|
|
groundtruth_g = CitationGraphDataset("cora")[0]
|
|
max_nid = []
|
|
|
|
for _ 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]
|
|
max_nid.append(np.max(F.asnumpy(dst_nodes_id)))
|
|
|
|
src_nodes_id = orig_nid[src_nodes_id]
|
|
dst_nodes_id = orig_nid[dst_nodes_id]
|
|
has_edges = groundtruth_g.has_edges_between(
|
|
src_nodes_id, dst_nodes_id
|
|
)
|
|
assert np.all(F.asnumpy(has_edges))
|
|
|
|
if use_graphbolt and not return_eids:
|
|
continue
|
|
eids = orig_eid[block.edata[dgl.EID]]
|
|
expected_eids = groundtruth_g.edge_ids(
|
|
src_nodes_id, dst_nodes_id
|
|
)
|
|
assert th.equal(
|
|
eids, expected_eids
|
|
), f"{eids} != {expected_eids}"
|
|
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
|
|
# this is needed since there's two test here in one process
|
|
dgl.distributed.exit_client()
|
|
|
|
|
|
@unittest.skip(reason="Skip due to glitch in CI")
|
|
def test_standalone():
|
|
reset_envs()
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
ip_config = os.path.join(test_dir, "ip_config.txt")
|
|
generate_ip_config(ip_config, 1, 1)
|
|
|
|
g = CitationGraphDataset("cora")[0]
|
|
print(g.idtype)
|
|
num_parts = 1
|
|
num_hops = 1
|
|
graph_name = f"graph_{uuid.uuid4()}"
|
|
orig_nid, orig_eid = partition_graph(
|
|
g,
|
|
graph_name,
|
|
num_parts,
|
|
test_dir,
|
|
num_hops=num_hops,
|
|
part_method="metis",
|
|
return_mapping=True,
|
|
)
|
|
part_config = os.path.join(test_dir, f"{graph_name}.json")
|
|
os.environ["DGL_DIST_MODE"] = "standalone"
|
|
try:
|
|
start_dist_dataloader(
|
|
0, ip_config, part_config, 1, True, orig_nid, orig_eid
|
|
)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
|
|
def start_dist_neg_dataloader(
|
|
rank,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
num_workers,
|
|
orig_nid,
|
|
groundtruth_g,
|
|
):
|
|
import dgl
|
|
import torch as th
|
|
|
|
dgl.distributed.initialize(ip_config)
|
|
gpb = None
|
|
disable_shared_mem = num_server > 1
|
|
if disable_shared_mem:
|
|
_, _, _, gpb, _, _, _ = load_partition(part_config, rank)
|
|
num_edges_to_sample = 202
|
|
batch_size = 32
|
|
graph_name = os.path.splitext(os.path.basename(part_config))[0]
|
|
dist_graph = DistGraph(graph_name, gpb=gpb, part_config=part_config)
|
|
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(part_config, i)
|
|
|
|
num_negs = 5
|
|
sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10])
|
|
negative_sampler = dgl.dataloading.negative_sampler.Uniform(num_negs)
|
|
dataloader = dgl.distributed.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
|
|
# this is needed since there's two test here in one process
|
|
dgl.distributed.exit_client()
|
|
|
|
|
|
def check_neg_dataloader(g, num_server, num_workers):
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
ip_config = "ip_config.txt"
|
|
generate_ip_config(ip_config, num_server, num_server)
|
|
|
|
num_parts = num_server
|
|
num_hops = 1
|
|
graph_name = f"graph_{uuid.uuid4()}"
|
|
orig_nid, orig_eid = partition_graph(
|
|
g,
|
|
graph_name,
|
|
num_parts,
|
|
test_dir,
|
|
num_hops=num_hops,
|
|
part_method="metis",
|
|
return_mapping=True,
|
|
)
|
|
part_config = os.path.join(test_dir, f"{graph_name}.json")
|
|
if not isinstance(orig_nid, dict):
|
|
orig_nid = {g.ntypes[0]: orig_nid}
|
|
if not isinstance(orig_eid, dict):
|
|
orig_eid = {g.etypes[0]: orig_eid}
|
|
|
|
pserver_list = []
|
|
ctx = mp.get_context("spawn")
|
|
for i in range(num_server):
|
|
p = ctx.Process(
|
|
target=start_server,
|
|
args=(
|
|
i,
|
|
ip_config,
|
|
part_config,
|
|
num_server > 1,
|
|
num_workers + 1,
|
|
),
|
|
)
|
|
p.start()
|
|
time.sleep(1)
|
|
pserver_list.append(p)
|
|
os.environ["DGL_DIST_MODE"] = "distributed"
|
|
os.environ["DGL_NUM_SAMPLER"] = str(num_workers)
|
|
ptrainer_list = []
|
|
|
|
p = ctx.Process(
|
|
target=start_dist_neg_dataloader,
|
|
args=(
|
|
0,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
num_workers,
|
|
orig_nid,
|
|
g,
|
|
),
|
|
)
|
|
p.start()
|
|
ptrainer_list.append(p)
|
|
|
|
for p in pserver_list:
|
|
p.join()
|
|
assert p.exitcode == 0
|
|
for p in ptrainer_list:
|
|
p.join()
|
|
assert p.exitcode == 0
|
|
|
|
|
|
@pytest.mark.parametrize("num_server", [1])
|
|
@pytest.mark.parametrize("num_workers", [0, 1])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
@pytest.mark.parametrize("return_eids", [False, True])
|
|
def test_dist_dataloader(num_server, num_workers, use_graphbolt, return_eids):
|
|
if not use_graphbolt and return_eids:
|
|
# return_eids is not supported in non-GraphBolt mode.
|
|
return
|
|
reset_envs()
|
|
os.environ["DGL_DIST_MODE"] = "distributed"
|
|
os.environ["DGL_NUM_SAMPLER"] = str(num_workers)
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
ip_config = "ip_config.txt"
|
|
generate_ip_config(ip_config, num_server, num_server)
|
|
|
|
g = CitationGraphDataset("cora")[0]
|
|
num_parts = num_server
|
|
num_hops = 1
|
|
graph_name = f"graph_{uuid.uuid4()}"
|
|
orig_nid, orig_eid = partition_graph(
|
|
g,
|
|
graph_name,
|
|
num_parts,
|
|
test_dir,
|
|
num_hops=num_hops,
|
|
part_method="metis",
|
|
return_mapping=True,
|
|
use_graphbolt=use_graphbolt,
|
|
store_eids=return_eids,
|
|
)
|
|
|
|
part_config = os.path.join(test_dir, f"{graph_name}.json")
|
|
pserver_list = []
|
|
ctx = mp.get_context("spawn")
|
|
for i in range(num_server):
|
|
p = ctx.Process(
|
|
target=start_server,
|
|
args=(
|
|
i,
|
|
ip_config,
|
|
part_config,
|
|
num_server > 1,
|
|
num_workers + 1,
|
|
use_graphbolt,
|
|
),
|
|
)
|
|
p.start()
|
|
time.sleep(1)
|
|
pserver_list.append(p)
|
|
|
|
ptrainer_list = []
|
|
num_trainers = 1
|
|
for trainer_id in range(num_trainers):
|
|
p = ctx.Process(
|
|
target=start_dist_dataloader,
|
|
args=(
|
|
trainer_id,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
False,
|
|
orig_nid,
|
|
orig_eid,
|
|
use_graphbolt,
|
|
return_eids,
|
|
),
|
|
)
|
|
p.start()
|
|
time.sleep(1) # avoid race condition when instantiating DistGraph
|
|
ptrainer_list.append(p)
|
|
|
|
for p in ptrainer_list:
|
|
p.join()
|
|
assert p.exitcode == 0
|
|
for p in pserver_list:
|
|
p.join()
|
|
assert p.exitcode == 0
|
|
|
|
|
|
def start_node_dataloader(
|
|
rank,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
num_workers,
|
|
orig_nid,
|
|
orig_eid,
|
|
groundtruth_g,
|
|
use_graphbolt=False,
|
|
return_eids=False,
|
|
prob_or_mask=None,
|
|
use_deprecated_dataloader=False,
|
|
):
|
|
dgl.distributed.initialize(ip_config, use_graphbolt=use_graphbolt)
|
|
gpb = None
|
|
disable_shared_mem = num_server > 1
|
|
if disable_shared_mem:
|
|
_, _, _, gpb, _, _, _ = load_partition(part_config, rank)
|
|
num_nodes_to_sample = 202
|
|
batch_size = 32
|
|
graph_name = os.path.splitext(os.path.basename(part_config))[0]
|
|
dist_graph = DistGraph(
|
|
graph_name,
|
|
gpb=gpb,
|
|
part_config=part_config,
|
|
)
|
|
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_nid = th.arange(num_nodes_to_sample, dtype=dist_graph.idtype)
|
|
else:
|
|
train_nid = {
|
|
"n3": th.arange(num_nodes_to_sample, dtype=dist_graph.idtype)
|
|
}
|
|
|
|
for i in range(num_server):
|
|
part, _, _, _, _, _, _ = load_partition(part_config, i)
|
|
|
|
# Create sampler
|
|
_prob = None
|
|
_mask = None
|
|
if prob_or_mask is None:
|
|
pass
|
|
elif prob_or_mask == "prob":
|
|
_prob = "prob"
|
|
elif prob_or_mask == "mask":
|
|
_mask = "mask"
|
|
else:
|
|
raise ValueError(f"Unsupported prob type: {prob_or_mask}")
|
|
sampler = dgl.dataloading.MultiLayerNeighborSampler(
|
|
[
|
|
(
|
|
# test dict for hetero
|
|
{etype: 5 for etype in dist_graph.etypes}
|
|
if len(dist_graph.etypes) > 1
|
|
else 5
|
|
),
|
|
10,
|
|
],
|
|
prob=_prob,
|
|
mask=_mask,
|
|
) # test int for hetero
|
|
|
|
# Enable santity check in distributed sampling.
|
|
os.environ["DGL_DIST_DEBUG"] = "1"
|
|
|
|
# We need to test creating DistDataLoader multiple times.
|
|
for i in range(2):
|
|
# Create DataLoader for constructing blocks
|
|
dataloader_cls = (
|
|
dgl.dataloading.DistNodeDataLoader
|
|
if use_deprecated_dataloader
|
|
else dgl.distributed.DistNodeDataLoader
|
|
)
|
|
dataloader = dataloader_cls(
|
|
dist_graph,
|
|
train_nid,
|
|
sampler,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
drop_last=False,
|
|
num_workers=num_workers,
|
|
)
|
|
|
|
for _ in range(2):
|
|
for idx, (_, _, blocks) in zip(
|
|
range(0, num_nodes_to_sample, batch_size), dataloader
|
|
):
|
|
block = blocks[-1]
|
|
for c_etype in block.canonical_etypes:
|
|
src_type, _, dst_type = c_etype
|
|
o_src, o_dst = block.edges(etype=c_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=c_etype
|
|
)
|
|
assert np.all(F.asnumpy(has_edges))
|
|
|
|
if use_graphbolt and not return_eids:
|
|
assert dgl.EID not in block.edges[c_etype].data
|
|
continue
|
|
eids = orig_eid[c_etype][block.edges[c_etype].data[dgl.EID]]
|
|
expected_eids = groundtruth_g.edge_ids(
|
|
src_nodes_id, dst_nodes_id, etype=c_etype
|
|
)
|
|
assert th.equal(
|
|
eids, expected_eids
|
|
), f"{eids} != {expected_eids}"
|
|
# Verify the prob/mask functionality.
|
|
if prob_or_mask is not None:
|
|
prob_data = groundtruth_g.edges[c_etype].data[
|
|
prob_or_mask
|
|
][eids]
|
|
assert th.all(prob_data > 0)
|
|
del dataloader
|
|
# this is needed since there's two test here in one process
|
|
dgl.distributed.exit_client()
|
|
|
|
|
|
def start_edge_dataloader(
|
|
rank,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
num_workers,
|
|
orig_nid,
|
|
orig_eid,
|
|
groundtruth_g,
|
|
use_graphbolt,
|
|
exclude,
|
|
reverse_eids,
|
|
reverse_etypes,
|
|
negative,
|
|
prob_or_mask,
|
|
use_deprecated_dataloader=False,
|
|
):
|
|
dgl.distributed.initialize(ip_config, use_graphbolt=use_graphbolt)
|
|
gpb = None
|
|
disable_shared_mem = num_server > 1
|
|
if disable_shared_mem:
|
|
_, _, _, gpb, _, _, _ = load_partition(part_config, rank)
|
|
num_edges_to_sample = 202
|
|
batch_size = 32
|
|
graph_name = os.path.splitext(os.path.basename(part_config))[0]
|
|
dist_graph = DistGraph(graph_name, gpb=gpb, part_config=part_config)
|
|
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.canonical_etypes[0]: th.arange(num_edges_to_sample)
|
|
}
|
|
|
|
for i in range(num_server):
|
|
part, _, _, _, _, _, _ = load_partition(part_config, i)
|
|
|
|
# Create sampler
|
|
_prob = None
|
|
_mask = None
|
|
if prob_or_mask is None:
|
|
pass
|
|
elif prob_or_mask == "prob":
|
|
_prob = "prob"
|
|
elif prob_or_mask == "mask":
|
|
_mask = "mask"
|
|
else:
|
|
raise ValueError(f"Unsupported prob type: {prob_or_mask}")
|
|
sampler = dgl.dataloading.MultiLayerNeighborSampler(
|
|
[5, -1], prob=_prob, mask=_mask
|
|
)
|
|
|
|
# Negative sampler.
|
|
negative_sampler = None
|
|
if negative:
|
|
negative_sampler = dgl.dataloading.negative_sampler.Uniform(5)
|
|
|
|
# We need to test creating DistDataLoader multiple times.
|
|
for i in range(2):
|
|
# Create DataLoader for constructing blocks
|
|
dataloader_cls = (
|
|
dgl.dataloading.DistEdgeDataLoader
|
|
if use_deprecated_dataloader
|
|
else dgl.distributed.DistEdgeDataLoader
|
|
)
|
|
dataloader = dataloader_cls(
|
|
dist_graph,
|
|
train_eid,
|
|
sampler,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
drop_last=False,
|
|
num_workers=num_workers,
|
|
exclude=exclude,
|
|
reverse_eids=reverse_eids,
|
|
reverse_etypes=reverse_etypes,
|
|
negative_sampler=negative_sampler,
|
|
)
|
|
|
|
for _ in range(2):
|
|
for _, minibatch in zip(
|
|
range(0, num_edges_to_sample, batch_size), dataloader
|
|
):
|
|
if negative:
|
|
_, pos_pair_graph, neg_pair_graph, blocks = minibatch
|
|
else:
|
|
_, pos_pair_graph, blocks = minibatch
|
|
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]
|
|
)
|
|
)
|
|
if negative:
|
|
assert np.all(
|
|
F.asnumpy(block.dstnodes[dst_type].data[dgl.NID])
|
|
== F.asnumpy(
|
|
neg_pair_graph.nodes[dst_type].data[dgl.NID]
|
|
)
|
|
)
|
|
if (
|
|
dgl.EID
|
|
not in block.edges[(src_type, etype, dst_type)].data
|
|
):
|
|
continue
|
|
sampled_eids = block.edges[
|
|
(src_type, etype, dst_type)
|
|
].data[dgl.EID]
|
|
sampled_orig_eids = orig_eid[(src_type, etype, dst_type)][
|
|
sampled_eids
|
|
]
|
|
raw_src, raw_dst = groundtruth_g.find_edges(
|
|
sampled_orig_eids, etype=(src_type, etype, dst_type)
|
|
)
|
|
sampled_src, sampled_dst = block.edges(
|
|
etype=(src_type, etype, dst_type)
|
|
)
|
|
sampled_orig_src = block.nodes[src_type].data[dgl.NID][
|
|
sampled_src
|
|
]
|
|
sampled_orig_dst = block.nodes[dst_type].data[dgl.NID][
|
|
sampled_dst
|
|
]
|
|
assert th.equal(
|
|
raw_src, orig_nid[src_type][sampled_orig_src]
|
|
)
|
|
assert th.equal(
|
|
raw_dst, orig_nid[dst_type][sampled_orig_dst]
|
|
)
|
|
# Verify the prob/mask functionality.
|
|
if prob_or_mask is not None:
|
|
prob_data = groundtruth_g.edges[etype].data[
|
|
prob_or_mask
|
|
][sampled_orig_eids]
|
|
assert th.all(prob_data > 0)
|
|
# Verify the exclude functionality.
|
|
if dgl.EID not in blocks[-1].edata.keys():
|
|
continue
|
|
for (
|
|
src_type,
|
|
etype,
|
|
dst_type,
|
|
) in pos_pair_graph.canonical_etypes:
|
|
for block in blocks:
|
|
if (
|
|
src_type,
|
|
etype,
|
|
dst_type,
|
|
) not in block.canonical_etypes:
|
|
continue
|
|
current_eids = block.edges[etype].data[dgl.EID]
|
|
seed_eids = pos_pair_graph.edges[etype].data[dgl.EID]
|
|
if exclude is None:
|
|
# seed_eids are not guaranteed to be sampled.
|
|
pass
|
|
elif exclude == "self":
|
|
assert not th.any(th.isin(current_eids, seed_eids))
|
|
elif exclude == "reverse_id":
|
|
src, dst = groundtruth_g.find_edges(seed_eids)
|
|
reverse_seed_eids = groundtruth_g.edge_ids(dst, src)
|
|
assert not th.any(
|
|
th.isin(current_eids, reverse_seed_eids)
|
|
)
|
|
assert not th.any(th.isin(current_eids, seed_eids))
|
|
elif exclude == "reverse_types":
|
|
assert not th.any(th.isin(current_eids, seed_eids))
|
|
reverse_etype = reverse_etypes[
|
|
(src_type, etype, dst_type)
|
|
]
|
|
if reverse_etype in block.canonical_etypes:
|
|
assert not th.any(
|
|
th.isin(
|
|
block.edges[reverse_etype].data[
|
|
dgl.EID
|
|
],
|
|
seed_eids,
|
|
)
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported exclude type: {exclude}"
|
|
)
|
|
del dataloader
|
|
dgl.distributed.exit_client()
|
|
|
|
|
|
def check_dataloader(
|
|
g,
|
|
num_server,
|
|
num_workers,
|
|
dataloader_type,
|
|
use_graphbolt=False,
|
|
return_eids=False,
|
|
exclude=None,
|
|
reverse_eids=None,
|
|
reverse_etypes=None,
|
|
negative=False,
|
|
prob_or_mask=None,
|
|
use_deprecated_dataloader=False,
|
|
):
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
ip_config = "ip_config.txt"
|
|
generate_ip_config(ip_config, num_server, num_server)
|
|
|
|
num_parts = num_server
|
|
num_hops = 1
|
|
graph_name = f"graph_{uuid.uuid4()}"
|
|
orig_nid, orig_eid = partition_graph(
|
|
g,
|
|
graph_name,
|
|
num_parts,
|
|
test_dir,
|
|
num_hops=num_hops,
|
|
part_method="metis",
|
|
return_mapping=True,
|
|
use_graphbolt=use_graphbolt,
|
|
store_eids=return_eids,
|
|
)
|
|
part_config = os.path.join(test_dir, f"{graph_name}.json")
|
|
if not isinstance(orig_nid, dict):
|
|
orig_nid = {g.ntypes[0]: orig_nid}
|
|
if not isinstance(orig_eid, dict):
|
|
orig_eid = {g.canonical_etypes[0]: orig_eid}
|
|
|
|
pserver_list = []
|
|
ctx = mp.get_context("spawn")
|
|
for i in range(num_server):
|
|
p = ctx.Process(
|
|
target=start_server,
|
|
args=(
|
|
i,
|
|
ip_config,
|
|
part_config,
|
|
num_server > 1,
|
|
num_workers + 1,
|
|
use_graphbolt,
|
|
),
|
|
)
|
|
p.start()
|
|
time.sleep(1)
|
|
pserver_list.append(p)
|
|
|
|
os.environ["DGL_DIST_MODE"] = "distributed"
|
|
os.environ["DGL_NUM_SAMPLER"] = str(num_workers)
|
|
ptrainer_list = []
|
|
if dataloader_type == "node":
|
|
p = ctx.Process(
|
|
target=start_node_dataloader,
|
|
args=(
|
|
0,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
num_workers,
|
|
orig_nid,
|
|
orig_eid,
|
|
g,
|
|
use_graphbolt,
|
|
return_eids,
|
|
prob_or_mask,
|
|
use_deprecated_dataloader,
|
|
),
|
|
)
|
|
p.start()
|
|
ptrainer_list.append(p)
|
|
elif dataloader_type == "edge":
|
|
p = ctx.Process(
|
|
target=start_edge_dataloader,
|
|
args=(
|
|
0,
|
|
ip_config,
|
|
part_config,
|
|
num_server,
|
|
num_workers,
|
|
orig_nid,
|
|
orig_eid,
|
|
g,
|
|
use_graphbolt,
|
|
exclude,
|
|
reverse_eids,
|
|
reverse_etypes,
|
|
negative,
|
|
prob_or_mask,
|
|
use_deprecated_dataloader,
|
|
),
|
|
)
|
|
p.start()
|
|
ptrainer_list.append(p)
|
|
for p in pserver_list:
|
|
p.join()
|
|
assert p.exitcode == 0
|
|
for p in ptrainer_list:
|
|
p.join()
|
|
assert p.exitcode == 0
|
|
|
|
|
|
def create_random_hetero():
|
|
num_nodes = {"n1": 10000, "n2": 10010, "n3": 10020}
|
|
etypes = [("n1", "r1", "n2"), ("n1", "r2", "n3"), ("n2", "r3", "n3")]
|
|
edges = {}
|
|
for etype in etypes:
|
|
src_ntype, _, dst_ntype = etype
|
|
arr = spsp.random(
|
|
num_nodes[src_ntype],
|
|
num_nodes[dst_ntype],
|
|
density=0.001,
|
|
format="coo",
|
|
random_state=100,
|
|
)
|
|
edges[etype] = (arr.row, arr.col)
|
|
# Add reverse edges.
|
|
src, dst = edges[("n1", "r1", "n2")]
|
|
edges[("n2", "r21", "n1")] = (dst, src)
|
|
g = dgl.heterograph(edges, num_nodes)
|
|
g.nodes["n1"].data["feat"] = F.unsqueeze(F.arange(0, g.num_nodes("n1")), 1)
|
|
g.edges["r1"].data["feat"] = F.unsqueeze(F.arange(0, g.num_edges("r1")), 1)
|
|
return g
|
|
|
|
|
|
@pytest.mark.parametrize("num_server", [1])
|
|
@pytest.mark.parametrize("num_workers", [0, 1])
|
|
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
@pytest.mark.parametrize("return_eids", [False, True])
|
|
def test_dataloader_homograph(
|
|
num_server, num_workers, dataloader_type, use_graphbolt, return_eids
|
|
):
|
|
if not use_graphbolt and return_eids:
|
|
# return_eids is not supported in non-GraphBolt mode.
|
|
return
|
|
reset_envs()
|
|
g = CitationGraphDataset("cora")[0]
|
|
check_dataloader(
|
|
g,
|
|
num_server,
|
|
num_workers,
|
|
dataloader_type,
|
|
use_graphbolt=use_graphbolt,
|
|
return_eids=return_eids,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("num_workers", [0])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
@pytest.mark.parametrize("exclude", [None, "self", "reverse_id"])
|
|
@pytest.mark.parametrize("negative", [False, True])
|
|
def test_edge_dataloader_homograph(
|
|
num_workers, use_graphbolt, exclude, negative
|
|
):
|
|
num_server = 1
|
|
dataloader_type = "edge"
|
|
reset_envs()
|
|
g, reverse_eids = _unique_rand_graph()
|
|
check_dataloader(
|
|
g,
|
|
num_server,
|
|
num_workers,
|
|
dataloader_type,
|
|
use_graphbolt=use_graphbolt,
|
|
return_eids=True,
|
|
exclude=exclude,
|
|
reverse_eids=reverse_eids,
|
|
negative=negative,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("num_server", [1])
|
|
@pytest.mark.parametrize("num_workers", [1])
|
|
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
@pytest.mark.parametrize("prob_or_mask", ["prob", "mask"])
|
|
def test_dataloader_homograph_prob_or_mask(
|
|
num_server, num_workers, dataloader_type, use_graphbolt, prob_or_mask
|
|
):
|
|
reset_envs()
|
|
g = CitationGraphDataset("cora")[0]
|
|
prob = th.rand(g.num_edges())
|
|
mask = prob > 0.2
|
|
g.edata["prob"] = F.tensor(prob)
|
|
g.edata["mask"] = F.tensor(mask)
|
|
check_dataloader(
|
|
g,
|
|
num_server,
|
|
num_workers,
|
|
dataloader_type,
|
|
use_graphbolt=use_graphbolt,
|
|
return_eids=True,
|
|
prob_or_mask=prob_or_mask,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("num_server", [1])
|
|
@pytest.mark.parametrize("num_workers", [0, 1])
|
|
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
@pytest.mark.parametrize("return_eids", [False, True])
|
|
def test_dataloader_heterograph(
|
|
num_server, num_workers, dataloader_type, use_graphbolt, return_eids
|
|
):
|
|
if not use_graphbolt and return_eids:
|
|
# return_eids is not supported in non-GraphBolt mode.
|
|
return
|
|
reset_envs()
|
|
g = create_random_hetero()
|
|
check_dataloader(
|
|
g,
|
|
num_server,
|
|
num_workers,
|
|
dataloader_type,
|
|
use_graphbolt=use_graphbolt,
|
|
return_eids=return_eids,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("num_workers", [0])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
@pytest.mark.parametrize("exclude", [None, "self", "reverse_types"])
|
|
@pytest.mark.parametrize("negative", [False, True])
|
|
def test_edge_dataloader_heterograph(
|
|
num_workers, use_graphbolt, exclude, negative
|
|
):
|
|
num_server = 1
|
|
dataloader_type = "edge"
|
|
reset_envs()
|
|
g = create_random_hetero()
|
|
reverse_etypes = {("n1", "r1", "n2"): ("n2", "r21", "n1")}
|
|
check_dataloader(
|
|
g,
|
|
num_server,
|
|
num_workers,
|
|
dataloader_type,
|
|
use_graphbolt=use_graphbolt,
|
|
return_eids=True,
|
|
exclude=exclude,
|
|
reverse_etypes=reverse_etypes,
|
|
negative=negative,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("num_server", [1])
|
|
@pytest.mark.parametrize("num_workers", [1])
|
|
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
@pytest.mark.parametrize("prob_or_mask", ["prob", "mask"])
|
|
def test_dataloader_heterograph_prob_or_mask(
|
|
num_server, num_workers, dataloader_type, use_graphbolt, prob_or_mask
|
|
):
|
|
reset_envs()
|
|
g = create_random_hetero()
|
|
for etype in g.canonical_etypes:
|
|
prob = th.rand(g.num_edges(etype))
|
|
mask = prob > prob.median()
|
|
g.edges[etype].data["prob"] = prob
|
|
g.edges[etype].data["mask"] = mask
|
|
check_dataloader(
|
|
g,
|
|
num_server,
|
|
num_workers,
|
|
dataloader_type,
|
|
use_graphbolt=use_graphbolt,
|
|
return_eids=True,
|
|
prob_or_mask=prob_or_mask,
|
|
)
|
|
|
|
|
|
@unittest.skip(reason="Skip due to glitch in CI")
|
|
@pytest.mark.parametrize("num_server", [3])
|
|
@pytest.mark.parametrize("num_workers", [0, 4])
|
|
def test_neg_dataloader(num_server, num_workers):
|
|
reset_envs()
|
|
g = CitationGraphDataset("cora")[0]
|
|
check_neg_dataloader(g, num_server, num_workers)
|
|
g = create_random_hetero()
|
|
check_neg_dataloader(g, num_server, num_workers)
|
|
|
|
|
|
def start_multiple_dataloaders(
|
|
ip_config,
|
|
part_config,
|
|
graph_name,
|
|
orig_g,
|
|
num_dataloaders,
|
|
dataloader_type,
|
|
use_graphbolt,
|
|
):
|
|
dgl.distributed.initialize(ip_config)
|
|
dist_g = dgl.distributed.DistGraph(graph_name, part_config=part_config)
|
|
if dataloader_type == "node":
|
|
train_ids = th.arange(orig_g.num_nodes(), dtype=dist_g.idtype)
|
|
batch_size = orig_g.num_nodes() // 100
|
|
else:
|
|
train_ids = th.arange(orig_g.num_edges())
|
|
batch_size = orig_g.num_edges() // 100
|
|
sampler = dgl.dataloading.NeighborSampler([-1])
|
|
dataloaders = []
|
|
dl_iters = []
|
|
for _ in range(num_dataloaders):
|
|
if dataloader_type == "node":
|
|
dataloader = dgl.distributed.DistNodeDataLoader(
|
|
dist_g, train_ids, sampler, batch_size=batch_size
|
|
)
|
|
else:
|
|
dataloader = dgl.distributed.DistEdgeDataLoader(
|
|
dist_g, train_ids, sampler, batch_size=batch_size
|
|
)
|
|
dataloaders.append(dataloader)
|
|
dl_iters.append(iter(dataloader))
|
|
|
|
# iterate on multiple dataloaders randomly
|
|
while len(dl_iters) > 0:
|
|
next_dl = np.random.choice(len(dl_iters), 1)[0]
|
|
try:
|
|
_ = next(dl_iters[next_dl])
|
|
except StopIteration:
|
|
dl_iters.pop(next_dl)
|
|
del dataloaders[next_dl]
|
|
|
|
dgl.distributed.exit_client()
|
|
|
|
|
|
@pytest.mark.parametrize("num_dataloaders", [4])
|
|
@pytest.mark.parametrize("num_workers", [0])
|
|
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
|
|
@pytest.mark.parametrize("use_graphbolt", [False, True])
|
|
def test_multiple_dist_dataloaders(
|
|
num_dataloaders, num_workers, dataloader_type, use_graphbolt
|
|
):
|
|
reset_envs()
|
|
os.environ["DGL_DIST_MODE"] = "distributed"
|
|
os.environ["DGL_NUM_SAMPLER"] = str(num_workers)
|
|
num_parts = 1
|
|
num_servers = 1
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
ip_config = os.path.join(test_dir, "ip_config.txt")
|
|
generate_ip_config(ip_config, num_parts, num_servers)
|
|
|
|
orig_g = dgl.rand_graph(1000, 10000)
|
|
graph_name = f"graph_{uuid.uuid4()}"
|
|
partition_graph(
|
|
orig_g,
|
|
graph_name,
|
|
num_parts,
|
|
test_dir,
|
|
use_graphbolt=use_graphbolt,
|
|
)
|
|
part_config = os.path.join(test_dir, f"{graph_name}.json")
|
|
|
|
p_servers = []
|
|
ctx = mp.get_context("spawn")
|
|
for i in range(num_servers):
|
|
p = ctx.Process(
|
|
target=start_server,
|
|
args=(
|
|
i,
|
|
ip_config,
|
|
part_config,
|
|
num_servers > 1,
|
|
num_workers + 1,
|
|
use_graphbolt,
|
|
),
|
|
)
|
|
p.start()
|
|
time.sleep(1)
|
|
p_servers.append(p)
|
|
|
|
p_client = ctx.Process(
|
|
target=start_multiple_dataloaders,
|
|
args=(
|
|
ip_config,
|
|
part_config,
|
|
graph_name,
|
|
orig_g,
|
|
num_dataloaders,
|
|
dataloader_type,
|
|
use_graphbolt,
|
|
),
|
|
)
|
|
p_client.start()
|
|
|
|
p_client.join()
|
|
assert p_client.exitcode == 0
|
|
for p in p_servers:
|
|
p.join()
|
|
assert p.exitcode == 0
|
|
reset_envs()
|
|
|
|
|
|
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
|
|
def test_deprecated_dataloader(dataloader_type):
|
|
reset_envs()
|
|
g = CitationGraphDataset("cora")[0]
|
|
check_dataloader(
|
|
g,
|
|
1,
|
|
0,
|
|
dataloader_type,
|
|
use_deprecated_dataloader=True,
|
|
)
|