test_distributed_sampling.py 9.52 KB
Newer Older
Jinjing Zhou's avatar
Jinjing Zhou committed
1
2
3
4
import dgl
import unittest
import os
from dgl.data import CitationGraphDataset
5
from dgl.distributed import sample_neighbors, find_edges
Jinjing Zhou's avatar
Jinjing Zhou committed
6
7
8
9
10
11
12
13
14
15
16
17
from dgl.distributed import partition_graph, load_partition, load_partition_book
import sys
import multiprocessing as mp
import numpy as np
import backend as F
import time
from utils import get_local_usable_addr
from pathlib import Path

from dgl.distributed import DistGraphServer, DistGraph


18
def start_server(rank, tmpdir, disable_shared_mem, graph_name):
19
    g = DistGraphServer(rank, "rpc_ip_config.txt", 1,
20
                        tmpdir / (graph_name + '.json'), disable_shared_mem=disable_shared_mem)
Jinjing Zhou's avatar
Jinjing Zhou committed
21
22
23
    g.start()


24
def start_sample_client(rank, tmpdir, disable_shared_mem):
25
26
    gpb = None
    if disable_shared_mem:
27
        _, _, _, gpb, _ = load_partition(tmpdir / 'test_sampling.json', rank)
28
    dist_graph = DistGraph("rpc_ip_config.txt", "test_sampling", gpb=gpb)
Jinjing Zhou's avatar
Jinjing Zhou committed
29
    sampled_graph = sample_neighbors(dist_graph, [0, 10, 99, 66, 1024, 2008], 3)
30
    dgl.distributed.exit_client()
Jinjing Zhou's avatar
Jinjing Zhou committed
31
32
    return sampled_graph

33
34
35
36
37
38
39
40
def start_find_edges_client(rank, tmpdir, disable_shared_mem, eids):
    gpb = None
    if disable_shared_mem:
        _, _, _, gpb, _ = load_partition(tmpdir / 'test_find_edges.json', rank)
    dist_graph = DistGraph("rpc_ip_config.txt", "test_find_edges", gpb=gpb)
    u, v = find_edges(dist_graph, eids)
    dgl.distributed.exit_client()
    return u, v
Jinjing Zhou's avatar
Jinjing Zhou committed
41

42
def check_rpc_sampling(tmpdir, num_server):
43
    ip_config = open("rpc_ip_config.txt", "w")
Jinjing Zhou's avatar
Jinjing Zhou committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
    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):
60
        p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling'))
Jinjing Zhou's avatar
Jinjing Zhou committed
61
62
63
64
65
        p.start()
        time.sleep(1)
        pserver_list.append(p)

    time.sleep(3)
66
    sampled_graph = start_sample_client(0, tmpdir, num_server > 1)
Jinjing Zhou's avatar
Jinjing Zhou committed
67
68
69
70
71
72
73
74
75
76
77
    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))

78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
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)

Jinjing Zhou's avatar
Jinjing Zhou committed
106
107
108
109
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(dgl.backend.backend_name == 'tensorflow', reason='Not support tensorflow for now')
def test_rpc_sampling():
    import tempfile
110
    os.environ['DGL_DIST_MODE'] = 'distributed'
Jinjing Zhou's avatar
Jinjing Zhou committed
111
    with tempfile.TemporaryDirectory() as tmpdirname:
112
        check_rpc_sampling(Path(tmpdirname), 2)
Jinjing Zhou's avatar
Jinjing Zhou committed
113

114
def check_rpc_sampling_shuffle(tmpdir, num_server):
115
    ip_config = open("rpc_ip_config.txt", "w")
Jinjing Zhou's avatar
Jinjing Zhou committed
116
117
118
    for _ in range(num_server):
        ip_config.write('{} 1\n'.format(get_local_usable_addr()))
    ip_config.close()
119

Jinjing Zhou's avatar
Jinjing Zhou committed
120
121
122
123
124
125
126
127
128
129
130
    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):
131
        p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling'))
Jinjing Zhou's avatar
Jinjing Zhou committed
132
133
134
135
136
        p.start()
        time.sleep(1)
        pserver_list.append(p)

    time.sleep(3)
137
    sampled_graph = start_sample_client(0, tmpdir, num_server > 1)
Jinjing Zhou's avatar
Jinjing Zhou committed
138
139
140
141
142
143
144
    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):
145
        part, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json', i)
Jinjing Zhou's avatar
Jinjing Zhou committed
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
        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))

# Wait non shared memory graph store
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(dgl.backend.backend_name == 'tensorflow', reason='Not support tensorflow for now')
def test_rpc_sampling_shuffle():
    import tempfile
163
    os.environ['DGL_DIST_MODE'] = 'distributed'
Jinjing Zhou's avatar
Jinjing Zhou committed
164
    with tempfile.TemporaryDirectory() as tmpdirname:
165
166
        check_rpc_sampling_shuffle(Path(tmpdirname), 2)
        check_rpc_sampling_shuffle(Path(tmpdirname), 1)
Jinjing Zhou's avatar
Jinjing Zhou committed
167

168
169
170
171
172
173
174
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)

175
    dist_graph = DistGraph(None, "test_sampling", part_config=tmpdir / 'test_sampling.json')
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
    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))

@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(dgl.backend.backend_name == 'tensorflow', reason='Not support tensorflow for now')
def test_standalone_sampling():
    import tempfile
    os.environ['DGL_DIST_MODE'] = 'standalone'
    with tempfile.TemporaryDirectory() as tmpdirname:
        check_standalone_sampling(Path(tmpdirname))

193
194
195
def start_in_subgraph_client(rank, tmpdir, disable_shared_mem, nodes):
    gpb = None
    if disable_shared_mem:
196
        _, _, _, gpb, _ = load_partition(tmpdir / 'test_in_subgraph.json', rank)
197
198
    dist_graph = DistGraph("rpc_ip_config.txt", "test_in_subgraph", gpb=gpb)
    sampled_graph = dgl.distributed.in_subgraph(dist_graph, nodes)
199
    dgl.distributed.exit_client()
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
    return sampled_graph


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()
    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))

@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(dgl.backend.backend_name == 'tensorflow', reason='Not support tensorflow for now')
def test_rpc_in_subgraph():
    import tempfile
244
    os.environ['DGL_DIST_MODE'] = 'distributed'
245
246
247
    with tempfile.TemporaryDirectory() as tmpdirname:
        check_rpc_in_subgraph(Path(tmpdirname), 2)

Jinjing Zhou's avatar
Jinjing Zhou committed
248
249
250
if __name__ == "__main__":
    import tempfile
    with tempfile.TemporaryDirectory() as tmpdirname:
251
252
253
        os.environ['DGL_DIST_MODE'] = 'standalone'
        check_standalone_sampling(Path(tmpdirname))
        os.environ['DGL_DIST_MODE'] = 'distributed'
254
        check_rpc_in_subgraph(Path(tmpdirname), 2)
255
256
257
258
        check_rpc_sampling_shuffle(Path(tmpdirname), 1)
        check_rpc_sampling_shuffle(Path(tmpdirname), 2)
        check_rpc_sampling(Path(tmpdirname), 2)
        check_rpc_sampling(Path(tmpdirname), 1)
259
260
        check_rpc_find_edges(Path(tmpdirname), 2)
        check_rpc_find_edges(Path(tmpdirname), 1)