test_mp_dataloader.py 25.8 KB
Newer Older
1
import multiprocessing as mp
Rhett Ying's avatar
Rhett Ying committed
2
3
import os
import tempfile
4
import time
5
import unittest
Rhett Ying's avatar
Rhett Ying committed
6

7
import backend as F
Rhett Ying's avatar
Rhett Ying committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
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

23
24

class NeighborSampler(object):
25
26
27
28
29
30
31
32
    def __init__(
        self,
        g,
        fanouts,
        sample_neighbors,
        use_graphbolt=False,
        return_eids=False,
    ):
33
34
35
        self.g = g
        self.fanouts = fanouts
        self.sample_neighbors = sample_neighbors
36
37
        self.use_graphbolt = use_graphbolt
        self.return_eids = return_eids
38
39
40

    def sample_blocks(self, seeds):
        import torch as th
Rhett Ying's avatar
Rhett Ying committed
41

42
43
44
45
46
        seeds = th.LongTensor(np.asarray(seeds))
        blocks = []
        for fanout in self.fanouts:
            # For each seed node, sample ``fanout`` neighbors.
            frontier = self.sample_neighbors(
47
                self.g, seeds, fanout, use_graphbolt=self.use_graphbolt
Rhett Ying's avatar
Rhett Ying committed
48
49
50
            )
            # Then we compact the frontier into a bipartite graph for
            # message passing.
51
52
53
            block = dgl.to_block(frontier, seeds)
            # Obtain the seed nodes for next layer.
            seeds = block.srcdata[dgl.NID]
54
55
56
            if frontier.num_edges() > 0:
                if not self.use_graphbolt or self.return_eids:
                    block.edata[dgl.EID] = frontier.edata[dgl.EID]
57
58
59
60
61

            blocks.insert(0, block)
        return blocks


Rhett Ying's avatar
Rhett Ying committed
62
63
64
65
66
67
def start_server(
    rank,
    ip_config,
    part_config,
    disable_shared_mem,
    num_clients,
68
    use_graphbolt=False,
Rhett Ying's avatar
Rhett Ying committed
69
70
71
72
73
74
75
76
77
78
):
    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"],
79
        use_graphbolt=use_graphbolt,
Rhett Ying's avatar
Rhett Ying committed
80
    )
81
82
83
    g.start()


Rhett Ying's avatar
Rhett Ying committed
84
85
86
87
88
89
90
91
def start_dist_dataloader(
    rank,
    ip_config,
    part_config,
    num_server,
    drop_last,
    orig_nid,
    orig_eid,
92
93
    use_graphbolt=False,
    return_eids=False,
Rhett Ying's avatar
Rhett Ying committed
94
95
):
    dgl.distributed.initialize(ip_config)
96
    gpb = None
97
    disable_shared_mem = num_server > 1
98
    if disable_shared_mem:
Rhett Ying's avatar
Rhett Ying committed
99
        _, _, _, gpb, _, _, _ = load_partition(part_config, rank)
100
101
102
    num_nodes_to_sample = 202
    batch_size = 32
    train_nid = th.arange(num_nodes_to_sample)
103
104
105
106
107
108
    dist_graph = DistGraph(
        "test_sampling",
        gpb=gpb,
        part_config=part_config,
        use_graphbolt=use_graphbolt,
    )
109

110
    # Create sampler
Rhett Ying's avatar
Rhett Ying committed
111
    sampler = NeighborSampler(
112
113
114
115
116
        dist_graph,
        [5, 10],
        dgl.distributed.sample_neighbors,
        use_graphbolt=use_graphbolt,
        return_eids=return_eids,
Rhett Ying's avatar
Rhett Ying committed
117
    )
118

119
120
121
    # Enable santity check in distributed sampling.
    os.environ["DGL_DIST_DEBUG"] = "1"

122
123
124
125
126
127
128
129
    # 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,
Rhett Ying's avatar
Rhett Ying committed
130
131
            drop_last=drop_last,
        )
132
133
134
135

        groundtruth_g = CitationGraphDataset("cora")[0]
        max_nid = []

136
        for _ in range(2):
Rhett Ying's avatar
Rhett Ying committed
137
138
139
            for idx, blocks in zip(
                range(0, num_nodes_to_sample, batch_size), dataloader
            ):
140
                block = blocks[-1]
Rhett Ying's avatar
Rhett Ying committed
141
                o_src, o_dst = block.edges()
142
143
                src_nodes_id = block.srcdata[dgl.NID][o_src]
                dst_nodes_id = block.dstdata[dgl.NID][o_dst]
144
145
146
147
                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]
Rhett Ying's avatar
Rhett Ying committed
148
149
150
                has_edges = groundtruth_g.has_edges_between(
                    src_nodes_id, dst_nodes_id
                )
151
                assert np.all(F.asnumpy(has_edges))
152
153
154
155
156
157
158
159
160
161

                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}"
162
            if drop_last:
Rhett Ying's avatar
Rhett Ying committed
163
164
165
166
167
168
                assert (
                    np.max(max_nid)
                    == num_nodes_to_sample
                    - 1
                    - num_nodes_to_sample % batch_size
                )
169
170
            else:
                assert np.max(max_nid) == num_nodes_to_sample - 1
171
    del dataloader
Rhett Ying's avatar
Rhett Ying committed
172
173
    # this is needed since there's two test here in one process
    dgl.distributed.exit_client()
174
175


Rhett Ying's avatar
Rhett Ying committed
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
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

        orig_nid, orig_eid = partition_graph(
            g,
            "test_sampling",
            num_parts,
            test_dir,
            num_hops=num_hops,
            part_method="metis",
            return_mapping=True,
        )
        part_config = os.path.join(test_dir, "test_sampling.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,
):
215
216
    import dgl
    import torch as th
Rhett Ying's avatar
Rhett Ying committed
217
218

    dgl.distributed.initialize(ip_config)
219
220
221
    gpb = None
    disable_shared_mem = num_server > 1
    if disable_shared_mem:
Rhett Ying's avatar
Rhett Ying committed
222
        _, _, _, gpb, _, _, _ = load_partition(part_config, rank)
223
224
    num_edges_to_sample = 202
    batch_size = 32
Rhett Ying's avatar
Rhett Ying committed
225
    dist_graph = DistGraph("test_mp", gpb=gpb, part_config=part_config)
226
227
228
229
230
231
232
233
    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):
Rhett Ying's avatar
Rhett Ying committed
234
        part, _, _, _, _, _, _ = load_partition(part_config, i)
235
236

    num_negs = 5
Rhett Ying's avatar
Rhett Ying committed
237
238
239
240
241
242
243
244
245
246
247
248
    sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10])
    negative_sampler = dgl.dataloading.negative_sampler.Uniform(num_negs)
    dataloader = dgl.dataloading.DistEdgeDataLoader(
        dist_graph,
        train_eid,
        sampler,
        batch_size=batch_size,
        negative_sampler=negative_sampler,
        shuffle=True,
        drop_last=False,
        num_workers=num_workers,
    )
249
    for _ in range(2):
Rhett Ying's avatar
Rhett Ying committed
250
251
252
        for _, (_, pos_graph, neg_graph, blocks) in zip(
            range(0, num_edges_to_sample, batch_size), dataloader
        ):
253
254
            block = blocks[-1]
            for src_type, etype, dst_type in block.canonical_etypes:
Rhett Ying's avatar
Rhett Ying committed
255
                o_src, o_dst = block.edges(etype=etype)
256
257
258
259
                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]
Rhett Ying's avatar
Rhett Ying committed
260
261
262
                has_edges = groundtruth_g.has_edges_between(
                    src_nodes_id, dst_nodes_id, etype=etype
                )
263
                assert np.all(F.asnumpy(has_edges))
Rhett Ying's avatar
Rhett Ying committed
264
265
266
267
268
269
270
271
                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])
                )
272
273
274
                assert pos_graph.num_edges() * num_negs == neg_graph.num_edges()

    del dataloader
Rhett Ying's avatar
Rhett Ying committed
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
    # 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
        orig_nid, orig_eid = partition_graph(
            g,
            "test_sampling",
            num_parts,
            test_dir,
            num_hops=num_hops,
            part_method="metis",
            return_mapping=True,
        )
        part_config = os.path.join(test_dir, "test_sampling.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,
            ),
        )
333
        p.start()
Rhett Ying's avatar
Rhett Ying committed
334
335
336
337
        ptrainer_list.append(p)

        for p in pserver_list:
            p.join()
338
            assert p.exitcode == 0
Rhett Ying's avatar
Rhett Ying committed
339
340
        for p in ptrainer_list:
            p.join()
341
            assert p.exitcode == 0
Rhett Ying's avatar
Rhett Ying committed
342
343


344
@pytest.mark.parametrize("num_server", [1])
345
@pytest.mark.parametrize("num_workers", [0, 1])
346
347
348
349
350
351
@pytest.mark.parametrize("drop_last", [False, True])
@pytest.mark.parametrize("use_graphbolt", [False, True])
@pytest.mark.parametrize("return_eids", [False, True])
def test_dist_dataloader(
    num_server, num_workers, drop_last, use_graphbolt, return_eids
):
352
    reset_envs()
353
354
    os.environ["DGL_DIST_MODE"] = "distributed"
    os.environ["DGL_NUM_SAMPLER"] = str(num_workers)
Rhett Ying's avatar
Rhett Ying committed
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
    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

        orig_nid, orig_eid = partition_graph(
            g,
            "test_sampling",
            num_parts,
            test_dir,
            num_hops=num_hops,
            part_method="metis",
            return_mapping=True,
371
372
            use_graphbolt=use_graphbolt,
            store_eids=return_eids,
Rhett Ying's avatar
Rhett Ying committed
373
374
375
376
377
378
379
380
381
382
383
384
385
386
        )

        part_config = os.path.join(test_dir, "test_sampling.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,
387
                    use_graphbolt,
Rhett Ying's avatar
Rhett Ying committed
388
389
                ),
            )
390
            p.start()
Rhett Ying's avatar
Rhett Ying committed
391
392
393
394
395
396
            time.sleep(1)
            pserver_list.append(p)

        ptrainer_list = []
        num_trainers = 1
        for trainer_id in range(num_trainers):
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
            p = ctx.Process(
                target=start_dist_dataloader,
                args=(
                    trainer_id,
                    ip_config,
                    part_config,
                    num_server,
                    drop_last,
                    orig_nid,
                    orig_eid,
                    use_graphbolt,
                    return_eids,
                ),
            )
            p.start()
            time.sleep(1)  # avoid race condition when instantiating DistGraph
            ptrainer_list.append(p)
Rhett Ying's avatar
Rhett Ying committed
414
415
416

        for p in ptrainer_list:
            p.join()
417
            assert p.exitcode == 0
418
        for p in pserver_list:
Rhett Ying's avatar
Rhett Ying committed
419
            p.join()
420
            assert p.exitcode == 0
Rhett Ying's avatar
Rhett Ying committed
421
422
423
424
425
426
427
428
429
430
431


def start_node_dataloader(
    rank,
    ip_config,
    part_config,
    num_server,
    num_workers,
    orig_nid,
    orig_eid,
    groundtruth_g,
432
433
    use_graphbolt=False,
    return_eids=False,
Rhett Ying's avatar
Rhett Ying committed
434
435
):
    dgl.distributed.initialize(ip_config)
436
    gpb = None
437
    disable_shared_mem = num_server > 1
438
    if disable_shared_mem:
Rhett Ying's avatar
Rhett Ying committed
439
        _, _, _, gpb, _, _, _ = load_partition(part_config, rank)
440
441
    num_nodes_to_sample = 202
    batch_size = 32
442
443
444
445
446
447
    dist_graph = DistGraph(
        "test_sampling",
        gpb=gpb,
        part_config=part_config,
        use_graphbolt=use_graphbolt,
    )
448
449
450
451
452
    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)
    else:
Rhett Ying's avatar
Rhett Ying committed
453
        train_nid = {"n3": th.arange(num_nodes_to_sample)}
454

455
    for i in range(num_server):
Rhett Ying's avatar
Rhett Ying committed
456
        part, _, _, _, _, _, _ = load_partition(part_config, i)
457

458
    # Create sampler
Rhett Ying's avatar
Rhett Ying committed
459
460
461
462
463
464
465
466
467
    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,
        ]
    )  # test int for hetero
468

469
470
471
    # Enable santity check in distributed sampling.
    os.environ["DGL_DIST_DEBUG"] = "1"

472
473
474
    # We need to test creating DistDataLoader multiple times.
    for i in range(2):
        # Create DataLoader for constructing blocks
475
        dataloader = dgl.dataloading.DistNodeDataLoader(
476
477
478
479
480
481
            dist_graph,
            train_nid,
            sampler,
            batch_size=batch_size,
            shuffle=True,
            drop_last=False,
Rhett Ying's avatar
Rhett Ying committed
482
483
            num_workers=num_workers,
        )
484

485
        for _ in range(2):
Rhett Ying's avatar
Rhett Ying committed
486
487
488
            for idx, (_, _, blocks) in zip(
                range(0, num_nodes_to_sample, batch_size), dataloader
            ):
489
                block = blocks[-1]
490
                for src_type, etype, dst_type in block.canonical_etypes:
Rhett Ying's avatar
Rhett Ying committed
491
                    o_src, o_dst = block.edges(etype=etype)
492
493
494
495
                    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]
Rhett Ying's avatar
Rhett Ying committed
496
497
498
                    has_edges = groundtruth_g.has_edges_between(
                        src_nodes_id, dst_nodes_id, etype=etype
                    )
499
                    assert np.all(F.asnumpy(has_edges))
500
501
502
503
504
505
506
507
508
509

                    if use_graphbolt and not return_eids:
                        continue
                    eids = orig_eid[etype][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}"
510
    del dataloader
Rhett Ying's avatar
Rhett Ying committed
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
    # 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,
):
    dgl.distributed.initialize(ip_config)
526
527
528
    gpb = None
    disable_shared_mem = num_server > 1
    if disable_shared_mem:
Rhett Ying's avatar
Rhett Ying committed
529
        _, _, _, gpb, _, _, _ = load_partition(part_config, rank)
530
531
    num_edges_to_sample = 202
    batch_size = 32
532
    dist_graph = DistGraph("test_sampling", gpb=gpb, part_config=part_config)
533
534
535
536
537
538
539
540
    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):
Rhett Ying's avatar
Rhett Ying committed
541
        part, _, _, _, _, _, _ = load_partition(part_config, i)
542
543
544

    # Create sampler
    sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10])
545

546
547
548
    # We need to test creating DistDataLoader multiple times.
    for i in range(2):
        # Create DataLoader for constructing blocks
549
        dataloader = dgl.dataloading.DistEdgeDataLoader(
550
551
552
553
554
555
            dist_graph,
            train_eid,
            sampler,
            batch_size=batch_size,
            shuffle=True,
            drop_last=False,
Rhett Ying's avatar
Rhett Ying committed
556
557
            num_workers=num_workers,
        )
558
559

        for epoch in range(2):
Rhett Ying's avatar
Rhett Ying committed
560
561
562
            for idx, (input_nodes, pos_pair_graph, blocks) in zip(
                range(0, num_edges_to_sample, batch_size), dataloader
            ):
563
564
                block = blocks[-1]
                for src_type, etype, dst_type in block.canonical_etypes:
Rhett Ying's avatar
Rhett Ying committed
565
                    o_src, o_dst = block.edges(etype=etype)
566
567
568
569
                    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]
Rhett Ying's avatar
Rhett Ying committed
570
571
572
                    has_edges = groundtruth_g.has_edges_between(
                        src_nodes_id, dst_nodes_id, etype=etype
                    )
573
                    assert np.all(F.asnumpy(has_edges))
Rhett Ying's avatar
Rhett Ying committed
574
575
576
577
578
579
                    assert np.all(
                        F.asnumpy(block.dstnodes[dst_type].data[dgl.NID])
                        == F.asnumpy(
                            pos_pair_graph.nodes[dst_type].data[dgl.NID]
                        )
                    )
580
    del dataloader
Rhett Ying's avatar
Rhett Ying committed
581
582
583
    dgl.distributed.exit_client()


584
585
586
587
588
589
590
591
def check_dataloader(
    g,
    num_server,
    num_workers,
    dataloader_type,
    use_graphbolt=False,
    return_eids=False,
):
Rhett Ying's avatar
Rhett Ying committed
592
593
594
595
596
597
598
599
600
601
602
603
604
605
    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
        orig_nid, orig_eid = partition_graph(
            g,
            "test_sampling",
            num_parts,
            test_dir,
            num_hops=num_hops,
            part_method="metis",
            return_mapping=True,
606
607
            use_graphbolt=use_graphbolt,
            store_eids=return_eids,
Rhett Ying's avatar
Rhett Ying committed
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
        )
        part_config = os.path.join(test_dir, "test_sampling.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,
626
                    use_graphbolt,
Rhett Ying's avatar
Rhett Ying committed
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
                ),
            )
            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,
648
649
                    use_graphbolt,
                    return_eids,
Rhett Ying's avatar
Rhett Ying committed
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
                ),
            )
            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,
                ),
            )
            p.start()
            ptrainer_list.append(p)
        for p in pserver_list:
            p.join()
672
            assert p.exitcode == 0
Rhett Ying's avatar
Rhett Ying committed
673
674
        for p in ptrainer_list:
            p.join()
675
            assert p.exitcode == 0
Rhett Ying's avatar
Rhett Ying committed
676

677

678
def create_random_hetero():
Rhett Ying's avatar
Rhett Ying committed
679
680
    num_nodes = {"n1": 10000, "n2": 10010, "n3": 10020}
    etypes = [("n1", "r1", "n2"), ("n1", "r2", "n3"), ("n2", "r3", "n3")]
681
682
683
    edges = {}
    for etype in etypes:
        src_ntype, _, dst_ntype = etype
Rhett Ying's avatar
Rhett Ying committed
684
685
686
687
688
689
690
        arr = spsp.random(
            num_nodes[src_ntype],
            num_nodes[dst_ntype],
            density=0.001,
            format="coo",
            random_state=100,
        )
691
692
        edges[etype] = (arr.row, arr.col)
    g = dgl.heterograph(edges, num_nodes)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
693
694
    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)
695
696
    return g

Rhett Ying's avatar
Rhett Ying committed
697

698
699
@pytest.mark.parametrize("num_server", [1])
@pytest.mark.parametrize("num_workers", [0, 1])
700
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
701
702
703
704
705
706
707
708
@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 dataloader_type == "edge" and use_graphbolt:
        # GraphBolt does not support edge dataloader.
        return
709
    reset_envs()
710
    g = CitationGraphDataset("cora")[0]
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
    check_dataloader(
        g,
        num_server,
        num_workers,
        dataloader_type,
        use_graphbolt=use_graphbolt,
        return_eids=return_eids,
    )


@unittest.skip(reason="Skip due to glitch in CI")
@pytest.mark.parametrize("num_server", [1])
@pytest.mark.parametrize("num_workers", [0, 1])
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
def test_dataloader_heterograph(num_server, num_workers, dataloader_type):
    reset_envs()
727
    g = create_random_hetero()
Rhett Ying's avatar
Rhett Ying committed
728
729
    check_dataloader(g, num_server, num_workers, dataloader_type)

730

731
@unittest.skip(reason="Skip due to glitch in CI")
732
733
@pytest.mark.parametrize("num_server", [3])
@pytest.mark.parametrize("num_workers", [0, 4])
Rhett Ying's avatar
Rhett Ying committed
734
def test_neg_dataloader(num_server, num_workers):
735
    reset_envs()
736
    g = CitationGraphDataset("cora")[0]
Rhett Ying's avatar
Rhett Ying committed
737
    check_neg_dataloader(g, num_server, num_workers)
738
    g = create_random_hetero()
Rhett Ying's avatar
Rhett Ying committed
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
    check_neg_dataloader(g, num_server, num_workers)


def start_multiple_dataloaders(
    ip_config, part_config, graph_name, orig_g, num_dataloaders, dataloader_type
):
    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())
        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.dataloading.DistNodeDataLoader(
                dist_g, train_ids, sampler, batch_size=batch_size
            )
        else:
            dataloader = dgl.dataloading.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()


780
@unittest.skip(reason="Skip due to glitch in CI")
Rhett Ying's avatar
Rhett Ying committed
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
@pytest.mark.parametrize("num_dataloaders", [1, 4])
@pytest.mark.parametrize("num_workers", [0, 1, 4])
@pytest.mark.parametrize("dataloader_type", ["node", "edge"])
def test_multiple_dist_dataloaders(
    num_dataloaders, num_workers, dataloader_type
):
    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 = "test"
        partition_graph(orig_g, graph_name, num_parts, test_dir)
        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,
                ),
            )
            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,
            ),
        )
        p_client.start()

        p_client.join()
832
        assert p_client.exitcode == 0
Rhett Ying's avatar
Rhett Ying committed
833
834
        for p in p_servers:
            p.join()
835
            assert p.exitcode == 0
Rhett Ying's avatar
Rhett Ying committed
836
    reset_envs()