def load_graph(name): cite_graphs = ['cora', 'citeseer', 'pubmed'] if name in cite_graphs: dataset = CitationGraphDataset(name) graph = dataset[0] nodes = graph.nodes() y = graph.ndata['label'] train_mask = graph.ndata['train_mask'] val_mask = graph.ndata['test_mask'] nodes_train, y_train = nodes[train_mask], y[train_mask] nodes_val, y_val = nodes[val_mask], y[val_mask] eval_set = [(nodes_train, y_train), (nodes_val, y_val)] elif name.startswith('ogbn'): dataset = DglNodePropPredDataset(name) graph, y = dataset[0] split_nodes = dataset.get_idx_split() nodes = graph.nodes() train_idx = split_nodes['train'] val_idx = split_nodes['valid'] nodes_train, y_train = nodes[train_idx], y[train_idx] nodes_val, y_val = nodes[val_idx], y[val_idx] eval_set = [(nodes_train, y_train), (nodes_val, y_val)] else: raise ValueError("Dataset name error!") return graph, eval_set
def check_rpc_get_degree_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_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'] 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 check_standalone_etype_sampling(tmpdir, reshuffle): hg = CitationGraphDataset('cora')[0] num_parts = 1 num_hops = 1 partition_graph(hg, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=reshuffle) os.environ['DGL_DIST_MODE'] = 'standalone' dgl.distributed.initialize("rpc_ip_config.txt") dist_graph = DistGraph("test_sampling", part_config=tmpdir / 'test_sampling.json') sampled_graph = sample_etype_neighbors(dist_graph, [0, 10, 99, 66, 1023], dgl.ETYPE, 3) src, dst = sampled_graph.edges() assert sampled_graph.number_of_nodes() == hg.number_of_nodes() assert np.all(F.asnumpy(hg.has_edges_between(src, dst))) eids = hg.edge_ids(src, dst) assert np.array_equal(F.asnumpy(sampled_graph.edata[dgl.EID]), F.asnumpy(eids)) dgl.distributed.exit_client()
def check_standalone_etype_sampling_heterograph(tmpdir, reshuffle): hg = CitationGraphDataset('cora')[0] num_parts = 1 num_hops = 1 src, dst = hg.edges() new_hg = dgl.heterograph( { ('paper', 'cite', 'paper'): (src, dst), ('paper', 'cite-by', 'paper'): (dst, src) }, {'paper': hg.number_of_nodes()}) partition_graph(new_hg, 'test_hetero_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=reshuffle) os.environ['DGL_DIST_MODE'] = 'standalone' dgl.distributed.initialize("rpc_ip_config.txt") dist_graph = DistGraph("test_hetero_sampling", part_config=tmpdir / 'test_hetero_sampling.json') sampled_graph = sample_etype_neighbors( dist_graph, [0, 1, 2, 10, 99, 66, 1023, 1024, 2700, 2701], dgl.ETYPE, 1) src, dst = sampled_graph.edges(etype=('paper', 'cite', 'paper')) assert len(src) == 10 src, dst = sampled_graph.edges(etype=('paper', 'cite-by', 'paper')) assert len(src) == 10 assert sampled_graph.number_of_nodes() == new_hg.number_of_nodes() dgl.distributed.exit_client()
def start_node_dataloader(rank, tmpdir, num_server, num_workers): import dgl import torch as th dgl.distributed.initialize("mp_ip_config.txt", 1, num_workers=num_workers) gpb = None disable_shared_mem = num_server > 1 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') orig_nid = F.zeros((dist_graph.number_of_nodes(), ), dtype=F.int64) orig_eid = F.zeros((dist_graph.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'] # 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.NodeDataLoader(dist_graph, train_nid, sampler, batch_size=batch_size, shuffle=True, drop_last=False, num_workers=num_workers) 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] 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)) 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)) del dataloader dgl.distributed.exit_client( ) # this is needed since there's two test here in one process
def start_dist_dataloader(rank, tmpdir, num_server, drop_last): import dgl import torch as th dgl.distributed.initialize("mp_ip_config.txt") gpb = None disable_shared_mem = num_server > 0 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') orig_nid = F.arange(0, dist_graph.number_of_nodes()) orig_eid = F.arange(0, dist_graph.number_of_edges()) for i in range(num_server): part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json', i) if 'orig_id' in part.ndata: orig_nid[part.ndata[dgl.NID]] = part.ndata['orig_id'] if 'orig_id' in part.edata: orig_eid[part.edata[dgl.EID]] = part.edata['orig_id'] # 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] 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)) # 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_client(rank, tmpdir, disable_shared_mem, num_workers, drop_last): import dgl import torch as th os.environ['DGL_DIST_MODE'] = 'distributed' dgl.distributed.initialize("mp_ip_config.txt", num_workers=4) 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("mp_ip_config.txt", "test_mp", gpb=gpb) # 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)) print(np.unique(np.sort(F.asnumpy(dst_nodes_id)))) 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 dgl.distributed.exit_client( ) # this is needed since there's two test here in one process
def check_rpc_in_subgraph_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_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] time.sleep(3) 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 get_dgl_citation(aggregation_type="sum", dataset="pubmed"): import dgl from dgl.data import CitationGraphDataset tr_set = CitationGraphDataset(dataset) G = tr_set.graph e = len(G.edges) n = len(G.nodes) d_l = tr_set.features.shape[1] is_multilabel = False n_classes = tr_set.num_labels node_labels = torch.tensor(tr_set.features) targets = torch.tensor(tr_set.labels) idx_train = torch.BoolTensor(tr_set.train_mask) idx_valid = torch.BoolTensor(tr_set.val_mask) idx_test = torch.BoolTensor(tr_set.test_mask) edges, agg_matrix = nx_to_format(G, aggregation_type) return Dataset( "cora", n, e, d_l, is_multilabel, n_classes, edges, agg_matrix, node_labels, targets, idx_train, idx_valid, idx_test, )
def test_dist_dataloader(tmpdir, num_server, num_workers, drop_last): ip_config = open("mp_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] print(g.idtype) num_parts = num_server num_hops = 1 partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=False) 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, num_workers+1)) p.start() time.sleep(1) pserver_list.append(p) time.sleep(3) os.environ['DGL_DIST_MODE'] = 'distributed' start_client(0, tmpdir, num_server > 1, num_workers, drop_last) dgl.distributed.exit_client() # this is needed since there's two test here in one process
def load_dataset(name): dataset = name.lower() if dataset == "amazon": from ogb.nodeproppred.dataset_dgl import DglNodePropPredDataset dataset = DglNodePropPredDataset(name="ogbn-products") splitted_idx = dataset.get_idx_split() train_nid = splitted_idx["train"] val_nid = splitted_idx["valid"] test_nid = splitted_idx["test"] g, labels = dataset[0] n_classes = int(labels.max() - labels.min() + 1) g.ndata['label'] = labels.squeeze() g.ndata['feat'] = g.ndata['feat'].float() elif dataset in ["reddit", "cora"]: if dataset == "reddit": from dgl.data import RedditDataset data = RedditDataset(self_loop=True) g = data[0] else: from dgl.data import CitationGraphDataset data = CitationGraphDataset('cora') g = data[0] n_classes = data.num_labels train_mask = g.ndata['train_mask'] val_mask = g.ndata['val_mask'] test_mask = g.ndata['test_mask'] train_nid = torch.LongTensor(train_mask.nonzero().squeeze()) val_nid = torch.LongTensor(val_mask.nonzero().squeeze()) test_nid = torch.LongTensor(test_mask.nonzero().squeeze()) else: print("Dataset {} is not supported".format(name)) assert (0) return g, n_classes, train_nid, val_nid, test_nid
def download(dataset): if dataset == 'cora': return CoraDataset() elif dataset == 'citeseer' or 'pubmed': return CitationGraphDataset(name=dataset) else: return None
def test_standalone(tmpdir): reset_envs() generate_ip_config("mp_ip_config.txt", 1, 1) g = CitationGraphDataset("cora")[0] print(g.idtype) num_parts = 1 num_hops = 1 orig_nid, orig_eid = partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=True, return_mapping=True) os.environ['DGL_DIST_MODE'] = 'standalone' try: start_dist_dataloader(0, tmpdir, 1, True, orig_nid, orig_eid) except Exception as e: print(e) dgl.distributed.exit_client( ) # this is needed since there's two test here in one process
def check_standalone_sampling(tmpdir): g = CitationGraphDataset("cora")[0] num_parts = 1 num_hops = 1 partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=False) dist_graph = DistGraph(None, "test_sampling", conf_file=tmpdir / 'test_sampling.json') sampled_graph = sample_neighbors(dist_graph, [0, 10, 99, 66, 1024, 2008], 3) src, dst = sampled_graph.edges() 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) assert np.array_equal( F.asnumpy(sampled_graph.edata[dgl.EID]), F.asnumpy(eids))
def test_dist_dataloader(tmpdir, num_server, drop_last): ip_config = open("mp_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] print(g.idtype) num_parts = num_server num_hops = 1 partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=False) num_workers = 4 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, num_workers + 1)) p.start() time.sleep(1) pserver_list.append(p) time.sleep(3) sampled_graph = start_client(0, tmpdir, num_server > 1, num_workers, drop_last)
def test_standalone(tmpdir): ip_config = open("mp_ip_config.txt", "w") for _ in range(1): ip_config.write('{}\n'.format(get_local_usable_addr())) ip_config.close() g = CitationGraphDataset("cora")[0] print(g.idtype) num_parts = 1 num_hops = 1 orig_nid, orig_eid = partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=True, return_mapping=True) os.environ['DGL_DIST_MODE'] = 'standalone' try: start_dist_dataloader(0, tmpdir, 1, True, orig_nid, orig_eid) except Exception as e: print(e) dgl.distributed.exit_client( ) # this is needed since there's two test here in one process
def __init__(self, name): t0 = time.time() self.name = name.lower() if self.name == 'cora': dataset = CoraDataset() else: dataset = CitationGraphDataset(self.name) dataset.graph.remove_edges_from(nx.selfloop_edges(dataset.graph)) graph = dgl.DGLGraph(dataset.graph) E = graph.number_of_edges() N = graph.number_of_nodes() D = dataset.features.shape[1] graph.ndata['feat'] = torch.Tensor(dataset.features) graph.edata['feat'] = torch.zeros((E, D)) graph.batch_num_nodes = [N] self.norm_n = torch.FloatTensor(N, 1).fill_(1. / float(N)).sqrt() self.norm_e = torch.FloatTensor(E, 1).fill_(1. / float(E)).sqrt() self.graph = graph self.train_mask = torch.BoolTensor(dataset.train_mask) self.val_mask = torch.BoolTensor(dataset.val_mask) self.test_mask = torch.BoolTensor(dataset.test_mask) self.labels = torch.LongTensor(dataset.labels) self.num_classes = dataset.num_labels self.num_dims = D print("[!] Dataset: ", self.name) print("Time taken: {:.4f}s".format(time.time() - t0))
def __init__(self, name): self.__name = name if name in graph_data: self.__graph_type = 'graph' elif name in node_data: self.__graph_type = 'node' else: print('Wrong dataset name!') self.__name = 'MUTAG' # download dataset root = os.path.abspath( os.path.join(os.path.dirname(__file__), os.path.pardir)) path = os.path.join(root, 'dataset', self.__graph_type) if not os.path.exists(path): os.makedirs(path) if not os.path.exists(os.path.join(path, self.__name)): if self.__name in graph_data: url = 'https://ls11-www.cs.tu-dortmund.de/people/morris/graphkerneldatasets/{0}.zip'.format( self.__name) zip_path = '{0}/{1}.zip'.format(path, self.__name) print('downloading {0}'.format(self.__name)) urllib.request.urlretrieve(url, zip_path) with zipfile.ZipFile(zip_path) as zf: zf.extractall(path) os.remove(zip_path) # process the raw dataset self.graphs, self.aug_view = self.process_raw_data() else: self.cite_graph = CitationGraphDataset(name=self.__name)
def test_dist_dataloader(tmpdir, num_server, num_workers, drop_last, reshuffle): ip_config = open("mp_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] print(g.idtype) num_parts = num_server num_hops = 1 orig_nid, orig_eid = partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=reshuffle, 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, num_workers+1)) p.start() time.sleep(1) pserver_list.append(p) time.sleep(3) os.environ['DGL_DIST_MODE'] = 'distributed' os.environ['DGL_NUM_SAMPLER'] = str(num_workers) ptrainer = ctx.Process(target=start_dist_dataloader, args=( 0, tmpdir, num_server, drop_last, orig_nid, orig_eid)) ptrainer.start() time.sleep(1) for p in pserver_list: p.join() ptrainer.join()
def check_rpc_sampling_shuffle(tmpdir, num_server, num_groups=1): generate_ip_config("rpc_ip_config.txt", num_server, num_server) 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') keep_alive = num_groups > 1 for i in range(num_server): p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling', ['csc', 'coo'], keep_alive)) p.start() time.sleep(1) pserver_list.append(p) pclient_list = [] num_clients = 1 for client_id in range(num_clients): for group_id in range(num_groups): p = ctx.Process(target=start_sample_client_shuffle, args=(client_id, tmpdir, num_server > 1, g, num_server, group_id)) p.start() time.sleep(1) # avoid race condition when instantiating DistGraph pclient_list.append(p) for p in pclient_list: p.join() if keep_alive: for p in pserver_list: assert p.is_alive() # force shutdown server dgl.distributed.shutdown_servers("rpc_ip_config.txt", 1) for p in pserver_list: p.join()
def check_rpc_sampling(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() print(g.idtype) num_parts = num_server num_hops = 1 partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=False) 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() src, dst = sampled_graph.edges() 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) assert np.array_equal( F.asnumpy(sampled_graph.edata[dgl.EID]), F.asnumpy(eids))
def check_rpc_sampling_shuffle(tmpdir): num_server = 2 ip_config = open("rpc_sampling_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)) p.start() time.sleep(1) pserver_list.append(p) time.sleep(3) sampled_graph = start_client(0, tmpdir) 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_dist_dataloader(tmpdir, num_server, num_workers, drop_last, reshuffle, num_groups): reset_envs() # No multiple partitions on single machine for # multiple client groups in case of race condition. if num_groups > 1: num_server = 1 generate_ip_config("mp_ip_config.txt", num_server, num_server) g = CitationGraphDataset("cora")[0] print(g.idtype) num_parts = num_server num_hops = 1 orig_nid, orig_eid = partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=reshuffle, return_mapping=True) pserver_list = [] ctx = mp.get_context('spawn') keep_alive = num_groups > 1 for i in range(num_server): p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, num_workers + 1, keep_alive)) 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 = [] num_trainers = 1 for trainer_id in range(num_trainers): for group_id in range(num_groups): p = ctx.Process(target=start_dist_dataloader, args=(trainer_id, tmpdir, num_server, drop_last, orig_nid, orig_eid, group_id)) p.start() time.sleep(1) # avoid race condition when instantiating DistGraph ptrainer_list.append(p) for p in ptrainer_list: p.join() if keep_alive: for p in pserver_list: assert p.is_alive() # force shutdown server dgl.distributed.shutdown_servers("mp_ip_config.txt", 1) for p in pserver_list: p.join()
def check_rpc_find_edges(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 partition_graph(g, 'test_find_edges', num_parts, tmpdir, num_hops=1, part_method='metis', reshuffle=False) 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')) p.start() time.sleep(1) pserver_list.append(p) time.sleep(3) eids = F.tensor(np.random.randint(g.number_of_edges(), size=100)) u, v = g.find_edges(eids) du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids) assert F.array_equal(u, du) assert F.array_equal(v, dv)
def check_rpc_find_edges_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 orig_nid, orig_eid = partition_graph(g, 'test_find_edges', num_parts, tmpdir, num_hops=1, 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_find_edges', ['csr', 'coo'])) p.start() time.sleep(1) pserver_list.append(p) 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_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 load_dataset(name): dataset = name.lower() if dataset == "amazon": from ogb.nodeproppred.dataset_dgl import DglNodePropPredDataset dataset = DglNodePropPredDataset(name="ogbn-products") splitted_idx = dataset.get_idx_split() train_nid = splitted_idx["train"] val_nid = splitted_idx["valid"] test_nid = splitted_idx["test"] g, labels = dataset[0] labels = labels.squeeze() n_classes = int(labels.max() - labels.min() + 1) features = g.ndata.pop("feat").float() elif dataset in ["reddit", "cora"]: if dataset == "reddit": from dgl.data import RedditDataset data = RedditDataset(self_loop=True) g = data.graph else: from dgl.data import CitationGraphDataset data = CitationGraphDataset('cora') g = dgl.DGLGraph(data.graph) train_mask = data.train_mask val_mask = data.val_mask test_mask = data.test_mask features = torch.Tensor(data.features) labels = torch.LongTensor(data.labels) n_classes = data.num_labels train_nid = torch.LongTensor(np.nonzero(train_mask)[0]) val_nid = torch.LongTensor(np.nonzero(val_mask)[0]) test_nid = torch.LongTensor(np.nonzero(test_mask)[0]) else: print("Dataset {} is not supported".format(name)) assert(0) return g, features, labels, n_classes, train_nid, val_nid, test_nid
def test_dataloader(tmpdir, num_server, num_workers, dataloader_type): ip_config = open("mp_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] print(g.idtype) num_parts = num_server num_hops = 1 partition_graph(g, 'test_sampling', num_parts, tmpdir, num_hops=num_hops, part_method='metis', reshuffle=False) 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, num_workers+1)) p.start() time.sleep(1) pserver_list.append(p) time.sleep(3) os.environ['DGL_DIST_MODE'] = 'distributed' ptrainer_list = [] if dataloader_type == 'node': p = ctx.Process(target=start_node_dataloader, args=( 0, tmpdir, num_server > 1, num_workers)) p.start() time.sleep(1) ptrainer_list.append(p) for p in pserver_list: p.join() for p in ptrainer_list: p.join()
def check_rpc_in_subgraph(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 partition_graph(g, 'test_in_subgraph', num_parts, tmpdir, num_hops=1, part_method='metis', reshuffle=False) 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] time.sleep(3) sampled_graph = start_in_subgraph_client(0, tmpdir, num_server > 1, nodes) for p in pserver_list: p.join() src, dst = sampled_graph.edges() g = dgl.as_heterograph(g) assert sampled_graph.number_of_nodes() == g.number_of_nodes() subg1 = dgl.in_subgraph(g, 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) assert np.array_equal(F.asnumpy(sampled_graph.edata[dgl.EID]), F.asnumpy(eids))
def test_neg_dataloader(tmpdir, num_server, num_workers): g = CitationGraphDataset("cora")[0] check_neg_dataloader(g, tmpdir, num_server, num_workers) g = create_random_hetero() check_neg_dataloader(g, tmpdir, num_server, num_workers)