test_distributed_sampling.py 39.6 KB
Newer Older
Jinjing Zhou's avatar
Jinjing Zhou committed
1
2
3
import dgl
import unittest
import os
4
import traceback
Jinjing Zhou's avatar
Jinjing Zhou committed
5
from dgl.data import CitationGraphDataset
6
7
from dgl.data import WN18Dataset
from dgl.distributed import sample_neighbors, sample_etype_neighbors
Jinjing Zhou's avatar
Jinjing Zhou committed
8
9
10
11
12
13
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
14
from utils import generate_ip_config, reset_envs
Jinjing Zhou's avatar
Jinjing Zhou committed
15
from pathlib import Path
16
import pytest
17
from scipy import sparse as spsp
18
import random
Jinjing Zhou's avatar
Jinjing Zhou committed
19
20
21
from dgl.distributed import DistGraphServer, DistGraph


22
23
def start_server(rank, tmpdir, disable_shared_mem, graph_name, graph_format=['csc', 'coo'],
                 keep_alive=False):
24
    g = DistGraphServer(rank, "rpc_ip_config.txt", 1, 1,
25
                        tmpdir / (graph_name + '.json'), disable_shared_mem=disable_shared_mem,
26
                        graph_format=graph_format, keep_alive=keep_alive)
Jinjing Zhou's avatar
Jinjing Zhou committed
27
28
29
    g.start()


30
def start_sample_client(rank, tmpdir, disable_shared_mem):
31
32
    gpb = None
    if disable_shared_mem:
33
        _, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json', rank)
34
    dgl.distributed.initialize("rpc_ip_config.txt")
35
    dist_graph = DistGraph("test_sampling", gpb=gpb)
36
37
38
    try:
        sampled_graph = sample_neighbors(dist_graph, [0, 10, 99, 66, 1024, 2008], 3)
    except Exception as e:
39
        print(traceback.format_exc())
40
        sampled_graph = None
41
    dgl.distributed.exit_client()
Jinjing Zhou's avatar
Jinjing Zhou committed
42
43
    return sampled_graph

44

45
46
def start_sample_client_shuffle(rank, tmpdir, disable_shared_mem, g, num_servers, group_id,
        orig_nid, orig_eid):
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
    os.environ['DGL_GROUP_ID'] = str(group_id)
    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)
    sampled_graph = sample_neighbors(dist_graph, [0, 10, 99, 66, 1024, 2008], 3)

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

64
def start_find_edges_client(rank, tmpdir, disable_shared_mem, eids, etype=None):
65
66
    gpb = None
    if disable_shared_mem:
67
        _, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_find_edges.json', rank)
68
    dgl.distributed.initialize("rpc_ip_config.txt")
69
    dist_graph = DistGraph("test_find_edges", gpb=gpb)
70
    try:
71
        u, v = dist_graph.find_edges(eids, etype=etype)
72
    except Exception as e:
73
        print(traceback.format_exc())
74
        u, v = None, None
75
76
    dgl.distributed.exit_client()
    return u, v
Jinjing Zhou's avatar
Jinjing Zhou committed
77

78
79
80
81
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)
82
    dgl.distributed.initialize("rpc_ip_config.txt")
83
84
85
86
87
88
89
    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:
90
        print(traceback.format_exc())
91
92
93
94
        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

95
def check_rpc_sampling(tmpdir, num_server):
96
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
Jinjing Zhou's avatar
Jinjing Zhou committed
97
98
99
100
101
102
103
104
105
106
107
108

    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):
109
        p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling'))
Jinjing Zhou's avatar
Jinjing Zhou committed
110
111
112
113
        p.start()
        time.sleep(1)
        pserver_list.append(p)

114
    sampled_graph = start_sample_client(0, tmpdir, num_server > 1)
Jinjing Zhou's avatar
Jinjing Zhou committed
115
116
117
118
119
120
121
122
123
124
125
    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))

126
def check_rpc_find_edges_shuffle(tmpdir, num_server):
127
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
128
129
130
131

    g = CitationGraphDataset("cora")[0]
    num_parts = num_server

132
133
134
    orig_nid, orig_eid = partition_graph(g, 'test_find_edges', num_parts, tmpdir,
                                         num_hops=1, part_method='metis',
                                         reshuffle=True, return_mapping=True)
135
136
137
138

    pserver_list = []
    ctx = mp.get_context('spawn')
    for i in range(num_server):
139
140
        p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1,
                                                   'test_find_edges', ['csr', 'coo']))
141
142
143
144
145
        p.start()
        time.sleep(1)
        pserver_list.append(p)

    eids = F.tensor(np.random.randint(g.number_of_edges(), size=100))
146
    u, v = g.find_edges(orig_eid[eids])
147
    du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids)
148
149
    du = orig_nid[du]
    dv = orig_nid[dv]
150
151
152
    assert F.array_equal(u, du)
    assert F.array_equal(v, dv)

153
154
155
def create_random_hetero(dense=False, empty=False):
    num_nodes = {'n1': 210, 'n2': 200, 'n3': 220} if dense else \
        {'n1': 1010, 'n2': 1000, 'n3': 1020}
156
157
158
    etypes = [('n1', 'r12', 'n2'),
              ('n1', 'r13', 'n3'),
              ('n2', 'r23', 'n3')]
159
    edges = {}
160
    random.seed(42)
161
162
    for etype in etypes:
        src_ntype, _, dst_ntype = etype
163
164
165
166
        arr = spsp.random(num_nodes[src_ntype] - 10 if empty else num_nodes[src_ntype],
                          num_nodes[dst_ntype] - 10 if empty else num_nodes[dst_ntype],
                          density=0.1 if dense else 0.001,
                          format='coo', random_state=100)
167
        edges[etype] = (arr.row, arr.col)
168
169
170
    g = dgl.heterograph(edges, num_nodes)
    g.nodes['n1'].data['feat'] = F.ones((g.number_of_nodes('n1'), 10), F.float32, F.cpu())
    return g
171
172

def check_rpc_hetero_find_edges_shuffle(tmpdir, num_server):
173
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190

    g = create_random_hetero()
    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)

191
192
    test_etype = g.to_canonical_etype('r12')
    eids = F.tensor(np.random.randint(g.num_edges(test_etype), size=100))
193
194
    expect_except = False
    try:
195
        _, _ = g.find_edges(orig_eid[test_etype][eids], etype=('n1', 'r12'))
196
197
198
    except:
        expect_except = True
    assert expect_except
199
200
    u, v = g.find_edges(orig_eid[test_etype][eids], etype='r12')
    u1, v1 = g.find_edges(orig_eid[test_etype][eids], etype=('n1', 'r12', 'n2'))
201
202
203
    assert F.array_equal(u, u1)
    assert F.array_equal(v, v1)
    du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids, etype='r12')
204
205
206
207
208
    du = orig_nid['n1'][du]
    dv = orig_nid['n2'][dv]
    assert F.array_equal(u, du)
    assert F.array_equal(v, dv)

209
210
211
# 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')
212
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Turn off Mxnet support")
213
214
@pytest.mark.parametrize("num_server", [1, 2])
def test_rpc_find_edges_shuffle(num_server):
215
    reset_envs()
216
217
218
    import tempfile
    os.environ['DGL_DIST_MODE'] = 'distributed'
    with tempfile.TemporaryDirectory() as tmpdirname:
219
        check_rpc_hetero_find_edges_shuffle(Path(tmpdirname), num_server)
220
221
222
        check_rpc_find_edges_shuffle(Path(tmpdirname), num_server)

def check_rpc_get_degree_shuffle(tmpdir, num_server):
223
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
224
225
226
227

    g = CitationGraphDataset("cora")[0]
    num_parts = num_server

228
229
    orig_nid, _ = partition_graph(g, 'test_get_degrees', num_parts, tmpdir,
        num_hops=1, part_method='metis', reshuffle=True, return_mapping=True)
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254

    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)

    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)

# 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')
255
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Turn off Mxnet support")
256
257
@pytest.mark.parametrize("num_server", [1, 2])
def test_rpc_get_degree_shuffle(num_server):
258
    reset_envs()
259
260
261
262
263
    import tempfile
    os.environ['DGL_DIST_MODE'] = 'distributed'
    with tempfile.TemporaryDirectory() as tmpdirname:
        check_rpc_get_degree_shuffle(Path(tmpdirname), num_server)

264
265
266
#@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
#@unittest.skipIf(dgl.backend.backend_name == 'tensorflow', reason='Not support tensorflow for now')
@unittest.skip('Only support partition with shuffle')
Jinjing Zhou's avatar
Jinjing Zhou committed
267
def test_rpc_sampling():
268
    reset_envs()
Jinjing Zhou's avatar
Jinjing Zhou committed
269
    import tempfile
270
    os.environ['DGL_DIST_MODE'] = 'distributed'
Jinjing Zhou's avatar
Jinjing Zhou committed
271
    with tempfile.TemporaryDirectory() as tmpdirname:
272
        check_rpc_sampling(Path(tmpdirname), 2)
Jinjing Zhou's avatar
Jinjing Zhou committed
273

274
def check_rpc_sampling_shuffle(tmpdir, num_server, num_groups=1):
275
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
276

Jinjing Zhou's avatar
Jinjing Zhou committed
277
278
279
280
    g = CitationGraphDataset("cora")[0]
    num_parts = num_server
    num_hops = 1

281
282
    orig_nids, orig_eids = partition_graph(g, 'test_sampling', num_parts, tmpdir,
        num_hops=num_hops, part_method='metis', reshuffle=True, return_mapping=True)
Jinjing Zhou's avatar
Jinjing Zhou committed
283
284
285

    pserver_list = []
    ctx = mp.get_context('spawn')
286
    keep_alive = num_groups > 1
Jinjing Zhou's avatar
Jinjing Zhou committed
287
    for i in range(num_server):
288
289
        p = ctx.Process(target=start_server, args=(
            i, tmpdir, num_server > 1, 'test_sampling', ['csc', 'coo'], keep_alive))
Jinjing Zhou's avatar
Jinjing Zhou committed
290
291
292
293
        p.start()
        time.sleep(1)
        pserver_list.append(p)

294
295
296
297
    pclient_list = []
    num_clients = 1
    for client_id in range(num_clients):
        for group_id in range(num_groups):
298
299
300
            p = ctx.Process(target=start_sample_client_shuffle,
                args=(client_id, tmpdir, num_server > 1, g, num_server,
                    group_id, orig_nids, orig_eids))
301
            p.start()
302
            time.sleep(1) # avoid race condition when instantiating DistGraph
303
304
305
306
307
308
309
310
            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)
Jinjing Zhou's avatar
Jinjing Zhou committed
311
312
313
    for p in pserver_list:
        p.join()

314
def start_hetero_sample_client(rank, tmpdir, disable_shared_mem, nodes):
315
316
317
    gpb = None
    if disable_shared_mem:
        _, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json', rank)
318
    dgl.distributed.initialize("rpc_ip_config.txt")
319
320
321
322
323
324
325
326
327
328
329
    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:
330
        print(traceback.format_exc())
331
332
333
334
        block = None
    dgl.distributed.exit_client()
    return block, gpb

335
def start_hetero_etype_sample_client(rank, tmpdir, disable_shared_mem, fanout=3,
336
337
                                     nodes={'n3': [0, 10, 99, 66, 124, 208]},
                                     etype_sorted=False):
338
339
340
341
342
343
344
345
    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
346
347
348
349
350
351
352
353
354
355
356

    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:])

357
358
359
    if gpb is None:
        gpb = dist_graph.get_partition_book()
    try:
360
361
        sampled_graph = sample_etype_neighbors(
                dist_graph, nodes, fanout, etype_sorted=etype_sorted)
362
363
364
        block = dgl.to_block(sampled_graph, nodes)
        block.edata[dgl.EID] = sampled_graph.edata[dgl.EID]
    except Exception as e:
365
        print(traceback.format_exc())
366
367
368
369
        block = None
    dgl.distributed.exit_client()
    return block, gpb

370
def check_rpc_hetero_sampling_shuffle(tmpdir, num_server):
371
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
372
373
374
375
376

    g = create_random_hetero()
    num_parts = num_server
    num_hops = 1

377
378
    orig_nid_map, orig_eid_map = partition_graph(g, 'test_sampling', num_parts, tmpdir,
        num_hops=num_hops, part_method='metis', reshuffle=True, return_mapping=True)
379
380
381
382
383
384
385
386
387

    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)

388
389
    block, gpb = start_hetero_sample_client(0, tmpdir, num_server > 1,
                                            nodes = {'n3': [0, 10, 99, 66, 124, 208]})
390
391
392
393
    print("Done sampling")
    for p in pserver_list:
        p.join()

394
395
    for c_etype in block.canonical_etypes:
        src_type, etype, dst_type = c_etype
396
397
398
399
400
401
402
403
        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))
404
        orig_eid = F.asnumpy(F.gather_row(orig_eid_map[c_etype], shuffled_eid))
405
406

        # Check the node Ids and edge Ids.
407
408
409
        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)
410

411
412
413
414
415
416
417
418
419
420
def get_degrees(g, nids, ntype):
    deg = F.zeros((len(nids),), dtype=F.int64)
    for srctype, etype, dsttype in g.canonical_etypes:
        if srctype == ntype:
            deg += g.out_degrees(u=nids, etype=etype)
        elif dsttype == ntype:
            deg += g.in_degrees(v=nids, etype=etype)
    return deg

def check_rpc_hetero_sampling_empty_shuffle(tmpdir, num_server):
421
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449

    g = create_random_hetero(empty=True)
    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)

    deg = get_degrees(g, orig_nids['n3'], 'n3')
    empty_nids = F.nonzero_1d(deg == 0)
    block, gpb = start_hetero_sample_client(0, tmpdir, num_server > 1,
                                            nodes = {'n3': empty_nids})
    print("Done sampling")
    for p in pserver_list:
        p.join()

    assert block.number_of_edges() == 0
    assert len(block.etypes) == len(g.etypes)

450
def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server, graph_formats=None):
451
452
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)

453
454
455
456
    g = create_random_hetero(dense=True)
    num_parts = num_server
    num_hops = 1

457
    orig_nid_map, orig_eid_map = partition_graph(g, 'test_sampling', num_parts, tmpdir,
458
459
        num_hops=num_hops, part_method='metis', reshuffle=True, return_mapping=True,
        graph_formats=graph_formats)
460
461
462
463

    pserver_list = []
    ctx = mp.get_context('spawn')
    for i in range(num_server):
464
        p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling', ['csc', 'coo']))
465
466
467
468
        p.start()
        time.sleep(1)
        pserver_list.append(p)

469
    fanout = {etype: 3 for etype in g.canonical_etypes}
470
471
472
    etype_sorted = False
    if graph_formats is not None:
        etype_sorted = 'csc' in graph_formats or 'csr' in graph_formats
473
    block, gpb = start_hetero_etype_sample_client(0, tmpdir, num_server > 1, fanout,
474
475
                                                  nodes={'n3': [0, 10, 99, 66, 124, 208]},
                                                  etype_sorted=etype_sorted)
476
477
478
479
    print("Done sampling")
    for p in pserver_list:
        p.join()

480
    src, dst = block.edges(etype=('n1', 'r13', 'n3'))
481
    assert len(src) == 18
482
    src, dst = block.edges(etype=('n2', 'r23', 'n3'))
483
484
    assert len(src) == 18

485
486
    for c_etype in block.canonical_etypes:
        src_type, etype, dst_type = c_etype
487
488
489
490
491
492
493
494
        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))
495
        orig_eid = F.asnumpy(F.gather_row(orig_eid_map[c_etype], shuffled_eid))
496
497
498
499
500
501

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

502
def check_rpc_hetero_etype_sampling_empty_shuffle(tmpdir, num_server):
503
504
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)

505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
    g = create_random_hetero(dense=True, empty=True)
    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['n3'], 'n3')
    empty_nids = F.nonzero_1d(deg == 0)
    block, gpb = start_hetero_etype_sample_client(0, tmpdir, num_server > 1, fanout,
                                                  nodes={'n3': empty_nids})
    print("Done sampling")
    for p in pserver_list:
        p.join()

    assert block.number_of_edges() == 0
    assert len(block.etypes) == len(g.etypes)

533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584

def create_random_bipartite():
    g = dgl.rand_bipartite('user', 'buys', 'game', 500, 1000, 1000)
    g.nodes['user'].data['feat'] = F.ones(
        (g.num_nodes('user'), 10), F.float32, F.cpu())
    g.nodes['game'].data['feat'] = F.ones(
        (g.num_nodes('game'), 10), F.float32, F.cpu())
    return g


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_bipartite_etype_sample_client(rank, tmpdir, disable_shared_mem, fanout=3,
                                        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 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()
585
    sampled_graph = sample_etype_neighbors(dist_graph, nodes, fanout)
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
    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 check_rpc_bipartite_sampling_empty(tmpdir, num_server):
    """sample on bipartite via sample_neighbors() which yields 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)

    deg = get_degrees(g, orig_nids['game'], 'game')
    empty_nids = F.nonzero_1d(deg == 0)
    block, _ = start_bipartite_sample_client(0, tmpdir, num_server > 1,
                                             nodes={'game': empty_nids, 'user': [1]})

    print("Done sampling")
    for p in pserver_list:
        p.join()

    assert block.number_of_edges() == 0
    assert len(block.etypes) == len(g.etypes)


def check_rpc_bipartite_sampling_shuffle(tmpdir, num_server):
    """sample on bipartite via sample_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

634
635
    orig_nid_map, orig_eid_map = partition_graph(g, 'test_sampling', num_parts, tmpdir,
        num_hops=num_hops, part_method='metis', reshuffle=True, return_mapping=True)
636
637
638
639
640
641
642
643
644
645

    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)

646
    deg = get_degrees(g, orig_nid_map['game'], 'game')
647
648
649
650
651
652
653
    nids = F.nonzero_1d(deg > 0)
    block, gpb = start_bipartite_sample_client(0, tmpdir, num_server > 1,
                                               nodes={'game': nids, 'user': [0]})
    print("Done sampling")
    for p in pserver_list:
        p.join()

654
655
    for c_etype in block.canonical_etypes:
        src_type, etype, dst_type = c_etype
656
657
658
659
660
661
662
663
664
665
666
667
        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))
668
        orig_eid = F.asnumpy(F.gather_row(orig_eid_map[c_etype], shuffled_eid))
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717

        # 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_bipartite_etype_sampling_empty(tmpdir, num_server):
    """sample on bipartite via sample_etype_neighbors() which yields 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)

    deg = get_degrees(g, orig_nids['game'], 'game')
    empty_nids = F.nonzero_1d(deg == 0)
    block, gpb = start_bipartite_etype_sample_client(0, tmpdir, num_server > 1,
                                                     nodes={'game': empty_nids, 'user': [1]})

    print("Done sampling")
    for p in pserver_list:
        p.join()

    assert block is not None
    assert block.number_of_edges() == 0
    assert len(block.etypes) == len(g.etypes)


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

718
719
    orig_nid_map, orig_eid_map = partition_graph(g, 'test_sampling', num_parts, tmpdir,
        num_hops=num_hops, part_method='metis', reshuffle=True, return_mapping=True)
720
721
722
723
724
725
726
727
728
729
730

    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
731
    deg = get_degrees(g, orig_nid_map['game'], 'game')
732
733
734
735
736
737
738
    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()

739
740
    for c_etype in block.canonical_etypes:
        src_type, etype, dst_type = c_etype
741
742
743
744
745
746
747
748
749
750
751
752
        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))
753
        orig_eid = F.asnumpy(F.gather_row(orig_eid_map[c_etype], shuffled_eid))
754
755
756
757
758
759

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

Jinjing Zhou's avatar
Jinjing Zhou committed
760
761
762
# 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')
763
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Turn off Mxnet support")
764
765
@pytest.mark.parametrize("num_server", [1, 2])
def test_rpc_sampling_shuffle(num_server):
766
    reset_envs()
Jinjing Zhou's avatar
Jinjing Zhou committed
767
    import tempfile
768
    os.environ['DGL_DIST_MODE'] = 'distributed'
Jinjing Zhou's avatar
Jinjing Zhou committed
769
    with tempfile.TemporaryDirectory() as tmpdirname:
770
        check_rpc_sampling_shuffle(Path(tmpdirname), num_server)
771
772
773
        # [TODO][Rhett] Tests for multiple groups may fail sometimes and
        # root cause is unknown. Let's disable them for now.
        #check_rpc_sampling_shuffle(Path(tmpdirname), num_server, num_groups=2)
774
        check_rpc_hetero_sampling_shuffle(Path(tmpdirname), num_server)
775
        check_rpc_hetero_sampling_empty_shuffle(Path(tmpdirname), num_server)
776
        check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server)
777
778
779
        check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server, ['csc'])
        check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server, ['csr'])
        check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server, ['csc', 'coo'])
780
        check_rpc_hetero_etype_sampling_empty_shuffle(Path(tmpdirname), num_server)
781
782
783
784
        check_rpc_bipartite_sampling_empty(Path(tmpdirname), num_server)
        check_rpc_bipartite_sampling_shuffle(Path(tmpdirname), num_server)
        check_rpc_bipartite_etype_sampling_empty(Path(tmpdirname), num_server)
        check_rpc_bipartite_etype_sampling_shuffle(Path(tmpdirname), num_server)
Jinjing Zhou's avatar
Jinjing Zhou committed
785

786
def check_standalone_sampling(tmpdir, reshuffle):
787
    g = CitationGraphDataset("cora")[0]
788
789
790
791
    prob = np.maximum(np.random.randn(g.num_edges()), 0)
    mask = (prob > 0)
    g.edata['prob'] = F.tensor(prob)
    g.edata['mask'] = F.tensor(mask)
792
793
794
    num_parts = 1
    num_hops = 1
    partition_graph(g, 'test_sampling', num_parts, tmpdir,
795
                    num_hops=num_hops, part_method='metis', reshuffle=reshuffle)
796

797
    os.environ['DGL_DIST_MODE'] = 'standalone'
798
    dgl.distributed.initialize("rpc_ip_config.txt")
799
    dist_graph = DistGraph("test_sampling", part_config=tmpdir / 'test_sampling.json')
800
801
802
803
804
805
806
807
    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))
808
809
810
811
812
813
814
815
816
817

    sampled_graph = sample_neighbors(
            dist_graph, [0, 10, 99, 66, 1024, 2008], 3, prob='mask')
    eid = F.asnumpy(sampled_graph.edata[dgl.EID])
    assert mask[eid].all()

    sampled_graph = sample_neighbors(
            dist_graph, [0, 10, 99, 66, 1024, 2008], 3, prob='prob')
    eid = F.asnumpy(sampled_graph.edata[dgl.EID])
    assert (prob[eid] > 0).all()
818
    dgl.distributed.exit_client()
819

820
821
def check_standalone_etype_sampling(tmpdir, reshuffle):
    hg = CitationGraphDataset('cora')[0]
822
823
824
825
    prob = np.maximum(np.random.randn(hg.num_edges()), 0)
    mask = (prob > 0)
    hg.edata['prob'] = F.tensor(prob)
    hg.edata['mask'] = F.tensor(mask)
826
827
828
829
830
831
832
833
    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')
834
    sampled_graph = sample_etype_neighbors(dist_graph, [0, 10, 99, 66, 1023], 3)
835
836
837
838
839
840
841

    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))
842
843
844
845
846
847
848
849
850
851

    sampled_graph = sample_etype_neighbors(
            dist_graph, [0, 10, 99, 66, 1023], 3, prob='mask')
    eid = F.asnumpy(sampled_graph.edata[dgl.EID])
    assert mask[eid].all()

    sampled_graph = sample_etype_neighbors(
            dist_graph, [0, 10, 99, 66, 1023], 3, prob='prob')
    eid = F.asnumpy(sampled_graph.edata[dgl.EID])
    assert (prob[eid] > 0).all()
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
    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')
867
868
    sampled_graph = sample_etype_neighbors(
            dist_graph, [0, 1, 2, 10, 99, 66, 1023, 1024, 2700, 2701], 1)
869
870
871
872
873
874
875
    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()

876
877
878
@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():
879
    reset_envs()
880
881
882
    import tempfile
    os.environ['DGL_DIST_MODE'] = 'standalone'
    with tempfile.TemporaryDirectory() as tmpdirname:
883
884
        check_standalone_sampling(Path(tmpdirname), False)
        check_standalone_sampling(Path(tmpdirname), True)
885

886
887
def start_in_subgraph_client(rank, tmpdir, disable_shared_mem, nodes):
    gpb = None
888
    dgl.distributed.initialize("rpc_ip_config.txt")
889
    if disable_shared_mem:
890
        _, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_in_subgraph.json', rank)
891
    dist_graph = DistGraph("test_in_subgraph", gpb=gpb)
892
893
894
    try:
        sampled_graph = dgl.distributed.in_subgraph(dist_graph, nodes)
    except Exception as e:
895
        print(traceback.format_exc())
896
        sampled_graph = None
897
    dgl.distributed.exit_client()
898
899
900
    return sampled_graph


901
def check_rpc_in_subgraph_shuffle(tmpdir, num_server):
902
    generate_ip_config("rpc_ip_config.txt", num_server, num_server)
903
904
905
906

    g = CitationGraphDataset("cora")[0]
    num_parts = num_server

907
908
    orig_nid, orig_eid = partition_graph(g, 'test_in_subgraph', num_parts, tmpdir,
        num_hops=1, part_method='metis', reshuffle=True, return_mapping=True)
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923

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

    src, dst = sampled_graph.edges()
924
925
    src = orig_nid[src]
    dst = orig_nid[dst]
926
    assert sampled_graph.number_of_nodes() == g.number_of_nodes()
927
928
929
    assert np.all(F.asnumpy(g.has_edges_between(src, dst)))

    subg1 = dgl.in_subgraph(g, orig_nid[nodes])
930
931
932
933
    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)
934
935
    eids1 = orig_eid[sampled_graph.edata[dgl.EID]]
    assert np.array_equal(F.asnumpy(eids1), F.asnumpy(eids))
936
937
938
939

@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():
940
    reset_envs()
941
    import tempfile
942
    os.environ['DGL_DIST_MODE'] = 'distributed'
943
    with tempfile.TemporaryDirectory() as tmpdirname:
944
        check_rpc_in_subgraph_shuffle(Path(tmpdirname), 2)
945

946
947
948
949
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(dgl.backend.backend_name == 'tensorflow', reason='Not support tensorflow for now')
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Turn off Mxnet support")
def test_standalone_etype_sampling():
950
    reset_envs()
951
952
953
954
955
956
957
958
    import tempfile
    with tempfile.TemporaryDirectory() as tmpdirname:
        os.environ['DGL_DIST_MODE'] = 'standalone'
        check_standalone_etype_sampling_heterograph(Path(tmpdirname), True)
    with tempfile.TemporaryDirectory() as tmpdirname:
        os.environ['DGL_DIST_MODE'] = 'standalone'
        check_standalone_etype_sampling(Path(tmpdirname), True)

Jinjing Zhou's avatar
Jinjing Zhou committed
959
960
961
if __name__ == "__main__":
    import tempfile
    with tempfile.TemporaryDirectory() as tmpdirname:
962
        os.environ['DGL_DIST_MODE'] = 'standalone'
963
964
965
966
967
968
        check_standalone_etype_sampling_heterograph(Path(tmpdirname), True)

    with tempfile.TemporaryDirectory() as tmpdirname:
        os.environ['DGL_DIST_MODE'] = 'standalone'
        check_standalone_etype_sampling(Path(tmpdirname), True)
        check_standalone_etype_sampling(Path(tmpdirname), False)
969
970
        check_standalone_sampling(Path(tmpdirname), True)
        check_standalone_sampling(Path(tmpdirname), False)
971
        os.environ['DGL_DIST_MODE'] = 'distributed'
972
973
        check_rpc_sampling(Path(tmpdirname), 2)
        check_rpc_sampling(Path(tmpdirname), 1)
974
975
        check_rpc_get_degree_shuffle(Path(tmpdirname), 1)
        check_rpc_get_degree_shuffle(Path(tmpdirname), 2)
976
977
        check_rpc_find_edges_shuffle(Path(tmpdirname), 2)
        check_rpc_find_edges_shuffle(Path(tmpdirname), 1)
978
979
        check_rpc_hetero_find_edges_shuffle(Path(tmpdirname), 1)
        check_rpc_hetero_find_edges_shuffle(Path(tmpdirname), 2)
980
981
982
983
        check_rpc_in_subgraph_shuffle(Path(tmpdirname), 2)
        check_rpc_sampling_shuffle(Path(tmpdirname), 1)
        check_rpc_hetero_sampling_shuffle(Path(tmpdirname), 1)
        check_rpc_hetero_sampling_shuffle(Path(tmpdirname), 2)
984
985
986
987
        check_rpc_hetero_sampling_empty_shuffle(Path(tmpdirname), 1)
        check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), 1)
        check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), 2)
        check_rpc_hetero_etype_sampling_empty_shuffle(Path(tmpdirname), 1)