chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,813 @@
+import json
+import os
+from itertools import product
+
+import dgl
+import dgl.backend as F
+
+import numpy as np
+from dgl.distributed import edge_split, load_partition_book, node_split
+
+mode = os.environ.get("DIST_DGL_TEST_MODE", "")
+graph_name = os.environ.get("DIST_DGL_TEST_GRAPH_NAME", "random_test_graph")
+num_part = int(os.environ.get("DIST_DGL_TEST_NUM_PART"))
+num_servers_per_machine = int(os.environ.get("DIST_DGL_TEST_NUM_SERVER"))
+num_client_per_machine = int(os.environ.get("DIST_DGL_TEST_NUM_CLIENT"))
+shared_workspace = os.environ.get("DIST_DGL_TEST_WORKSPACE")
+graph_path = os.environ.get("DIST_DGL_TEST_GRAPH_PATH")
+part_id = int(os.environ.get("DIST_DGL_TEST_PART_ID"))
+ip_config = os.environ.get("DIST_DGL_TEST_IP_CONFIG", "ip_config.txt")
+
+os.environ["DGL_DIST_MODE"] = "distributed"
+
+
+def batched_assert_zero(tensor, size):
+    BATCH_SIZE = 2**16
+    curr_pos = 0
+    while curr_pos < size:
+        end = min(curr_pos + BATCH_SIZE, size)
+        assert F.sum(tensor[F.arange(curr_pos, end)], 0) == 0
+        curr_pos = end
+
+
+def zeros_init(shape, dtype):
+    return F.zeros(shape, dtype=dtype, ctx=F.cpu())
+
+
+def rand_init(shape, dtype):
+    return F.tensor((np.random.randint(0, 100, size=shape) > 30), dtype=dtype)
+
+
+def run_server(
+    graph_name,
+    server_id,
+    server_count,
+    num_clients,
+    shared_mem,
+):
+    # server_count = num_servers_per_machine
+    g = dgl.distributed.DistGraphServer(
+        server_id,
+        ip_config,
+        server_count,
+        num_clients,
+        graph_path + "/{}.json".format(graph_name),
+        disable_shared_mem=not shared_mem,
+        graph_format=["csc", "coo"],
+    )
+    print("start server", server_id)
+    g.start()
+
+
+##########################################
+############### DistGraph ###############
+##########################################
+
+
+def node_split_test(g, force_even, ntype="_N"):
+    gpb = g.get_partition_book()
+
+    selected_nodes_dist_tensor = dgl.distributed.DistTensor(
+        [g.num_nodes(ntype)], F.uint8, init_func=rand_init
+    )
+
+    nodes = node_split(
+        selected_nodes_dist_tensor, gpb, ntype=ntype, force_even=force_even
+    )
+    g.barrier()
+
+    selected_nodes_dist_tensor[nodes] = F.astype(
+        F.zeros_like(nodes), selected_nodes_dist_tensor.dtype
+    )
+    g.barrier()
+
+    if g.rank() == 0:
+        batched_assert_zero(selected_nodes_dist_tensor, g.num_nodes(ntype))
+
+    g.barrier()
+
+
+def edge_split_test(g, force_even, etype="_E"):
+    gpb = g.get_partition_book()
+
+    selected_edges_dist_tensor = dgl.distributed.DistTensor(
+        [g.num_edges(etype)], F.uint8, init_func=rand_init
+    )
+
+    edges = edge_split(
+        selected_edges_dist_tensor, gpb, etype=etype, force_even=force_even
+    )
+    g.barrier()
+
+    selected_edges_dist_tensor[edges] = F.astype(
+        F.zeros_like(edges), selected_edges_dist_tensor.dtype
+    )
+    g.barrier()
+
+    if g.rank() == 0:
+        batched_assert_zero(selected_edges_dist_tensor, g.num_edges(etype))
+
+    g.barrier()
+
+
+def test_dist_graph(g):
+    gpb_path = graph_path + "/{}.json".format(graph_name)
+    with open(gpb_path) as conf_f:
+        part_metadata = json.load(conf_f)
+    assert "num_nodes" in part_metadata
+    assert "num_edges" in part_metadata
+    num_nodes = part_metadata["num_nodes"]
+    num_edges = part_metadata["num_edges"]
+
+    assert g.num_nodes() == num_nodes
+    assert g.num_edges() == num_edges
+
+    num_nodes = {ntype: g.num_nodes(ntype) for ntype in g.ntypes}
+    num_edges = {etype: g.num_edges(etype) for etype in g.etypes}
+
+    for key, n_nodes in num_nodes.items():
+        assert g.num_nodes(key) == n_nodes
+        node_split_test(g, force_even=False, ntype=key)
+        node_split_test(g, force_even=True, ntype=key)
+
+    for key, n_edges in num_edges.items():
+        assert g.num_edges(key) == n_edges
+        edge_split_test(g, force_even=False, etype=key)
+        edge_split_test(g, force_even=True, etype=key)
+
+
+##########################################
+########### DistGraphServices ###########
+##########################################
+
+
+def find_edges_test(g, orig_nid_map):
+    etypes = g.canonical_etypes
+
+    etype_eids_uv_map = dict()
+    for u_type, etype, v_type in etypes:
+        orig_u = g.edges[etype].data["edge_u"]
+        orig_v = g.edges[etype].data["edge_v"]
+        eids = F.tensor(np.random.randint(g.num_edges(etype), size=100))
+        u, v = g.find_edges(eids, etype=etype)
+        assert F.allclose(orig_nid_map[u_type][u], orig_u[eids])
+        assert F.allclose(orig_nid_map[v_type][v], orig_v[eids])
+        etype_eids_uv_map[etype] = (eids, F.cat([u, v], dim=0))
+    return etype_eids_uv_map
+
+
+def edge_subgraph_test(g, etype_eids_uv_map):
+    etypes = g.canonical_etypes
+    all_eids = dict()
+    for t in etypes:
+        all_eids[t] = etype_eids_uv_map[t[1]][0]
+
+    sg = g.edge_subgraph(all_eids)
+    for t in etypes:
+        assert sg.num_edges(t[1]) == len(all_eids[t])
+        assert F.allclose(sg.edges[t].data[dgl.EID], all_eids[t])
+
+    for u_type, etype, v_type in etypes:
+        uv = etype_eids_uv_map[etype][1]
+        sg_u_nids = sg.nodes[u_type].data[dgl.NID]
+        sg_v_nids = sg.nodes[v_type].data[dgl.NID]
+        sg_uv = F.cat([sg_u_nids, sg_v_nids], dim=0)
+        for node_id in uv:
+            assert node_id in sg_uv
+
+
+def sample_neighbors_with_args(g, size, fanout):
+    num_nodes = {ntype: g.num_nodes(ntype) for ntype in g.ntypes}
+    etypes = g.canonical_etypes
+
+    sampled_graph = g.sample_neighbors(
+        {
+            ntype: np.random.randint(0, n, size=size)
+            for ntype, n in num_nodes.items()
+        },
+        fanout,
+    )
+
+    for ntype, n in num_nodes.items():
+        assert sampled_graph.num_nodes(ntype) == n
+    for t in etypes:
+        src, dst = sampled_graph.edges(etype=t)
+        eids = sampled_graph.edges[t].data[dgl.EID]
+        dist_u, dist_v = g.find_edges(eids, etype=t[1])
+        assert F.allclose(dist_u, src)
+        assert F.allclose(dist_v, dst)
+
+
+def sample_neighbors_test(g):
+    sample_neighbors_with_args(g, size=1024, fanout=3)
+    sample_neighbors_with_args(g, size=1, fanout=10)
+    sample_neighbors_with_args(g, size=1024, fanout=2)
+    sample_neighbors_with_args(g, size=10, fanout=-1)
+    sample_neighbors_with_args(g, size=2**10, fanout=1)
+    sample_neighbors_with_args(g, size=2**12, fanout=1)
+
+
+def test_dist_graph_services(g):
+    # in_degrees and out_degrees does not support heterograph
+    if len(g.etypes) == 1:
+        nids = F.arange(0, 128)
+
+        # Test in_degrees
+        orig_in_degrees = g.ndata["in_degrees"]
+        local_in_degrees = g.in_degrees(nids)
+        F.allclose(local_in_degrees, orig_in_degrees[nids])
+
+        # Test out_degrees
+        orig_out_degrees = g.ndata["out_degrees"]
+        local_out_degrees = g.out_degrees(nids)
+        F.allclose(local_out_degrees, orig_out_degrees[nids])
+
+    num_nodes = {ntype: g.num_nodes(ntype) for ntype in g.ntypes}
+
+    orig_nid_map = dict()
+    dtype = g.edges[g.etypes[0]].data["edge_u"].dtype
+    for ntype, _ in num_nodes.items():
+        orig_nid = F.tensor(
+            np.load(graph_path + f"/orig_nid_array_{ntype}.npy"), dtype
+        )
+        orig_nid_map[ntype] = orig_nid
+
+    etype_eids_uv_map = find_edges_test(g, orig_nid_map)
+    edge_subgraph_test(g, etype_eids_uv_map)
+    sample_neighbors_test(g)
+
+
+##########################################
+############### DistTensor ###############
+##########################################
+
+
+def dist_tensor_test_sanity(data_shape, name=None):
+    local_rank = dgl.distributed.get_rank() % num_client_per_machine
+    dist_ten = dgl.distributed.DistTensor(
+        data_shape, F.int32, init_func=zeros_init, name=name
+    )
+    # arbitrary value
+    stride = 3
+    pos = (part_id // 2) * num_client_per_machine + local_rank
+    if part_id % 2 == 0:
+        dist_ten[pos * stride : (pos + 1) * stride] = F.ones(
+            (stride, 2), dtype=F.int32, ctx=F.cpu()
+        ) * (pos + 1)
+
+    dgl.distributed.client_barrier()
+    assert F.allclose(
+        dist_ten[pos * stride : (pos + 1) * stride],
+        F.ones((stride, 2), dtype=F.int32, ctx=F.cpu()) * (pos + 1),
+    )
+
+
+def dist_tensor_test_destroy_recreate(data_shape, name):
+    dist_ten = dgl.distributed.DistTensor(
+        data_shape, F.float32, name, init_func=zeros_init
+    )
+    del dist_ten
+
+    dgl.distributed.client_barrier()
+
+    new_shape = (data_shape[0], 4)
+    dist_ten = dgl.distributed.DistTensor(
+        new_shape, F.float32, name, init_func=zeros_init
+    )
+
+
+def dist_tensor_test_persistent(data_shape):
+    dist_ten_name = "persistent_dist_tensor"
+    dist_ten = dgl.distributed.DistTensor(
+        data_shape,
+        F.float32,
+        dist_ten_name,
+        init_func=zeros_init,
+        persistent=True,
+    )
+    del dist_ten
+    try:
+        dist_ten = dgl.distributed.DistTensor(
+            data_shape, F.float32, dist_ten_name
+        )
+        raise Exception("")
+    except BaseException:
+        pass
+
+
+def test_dist_tensor(g):
+    first_type = g.ntypes[0]
+    data_shape = (g.num_nodes(first_type), 2)
+    dist_tensor_test_sanity(data_shape)
+    dist_tensor_test_sanity(data_shape, name="DistTensorSanity")
+    dist_tensor_test_destroy_recreate(data_shape, name="DistTensorRecreate")
+    dist_tensor_test_persistent(data_shape)
+
+
+##########################################
+############# DistEmbedding ##############
+##########################################
+
+
+def dist_embedding_check_sanity(num_nodes, optimizer, name=None):
+    local_rank = dgl.distributed.get_rank() % num_client_per_machine
+
+    emb = dgl.distributed.DistEmbedding(
+        num_nodes, 1, name=name, init_func=zeros_init
+    )
+    lr = 0.001
+    optim = optimizer(params=[emb], lr=lr)
+
+    stride = 3
+
+    pos = (part_id // 2) * num_client_per_machine + local_rank
+    idx = F.arange(pos * stride, (pos + 1) * stride)
+
+    if part_id % 2 == 0:
+        with F.record_grad():
+            value = emb(idx)
+            optim.zero_grad()
+            loss = F.sum(value + 1, 0)
+        loss.backward()
+        optim.step()
+
+    dgl.distributed.client_barrier()
+    value = emb(idx)
+    F.allclose(value, F.ones((len(idx), 1), dtype=F.int32, ctx=F.cpu()) * -lr)
+
+    not_update_idx = F.arange(
+        ((num_part + 1) / 2) * num_client_per_machine * stride, num_nodes
+    )
+    value = emb(not_update_idx)
+    assert np.all(F.asnumpy(value) == np.zeros((len(not_update_idx), 1)))
+
+
+def dist_embedding_check_existing(num_nodes):
+    dist_emb_name = "UniqueEmb"
+    emb = dgl.distributed.DistEmbedding(
+        num_nodes, 1, name=dist_emb_name, init_func=zeros_init
+    )
+    try:
+        emb1 = dgl.distributed.DistEmbedding(
+            num_nodes, 2, name=dist_emb_name, init_func=zeros_init
+        )
+        raise Exception("")
+    except BaseException:
+        pass
+
+
+def test_dist_embedding(g):
+    num_nodes = g.num_nodes(g.ntypes[0])
+    dist_embedding_check_sanity(num_nodes, dgl.distributed.optim.SparseAdagrad)
+    dist_embedding_check_sanity(
+        num_nodes, dgl.distributed.optim.SparseAdagrad, name="SomeEmbedding"
+    )
+    dist_embedding_check_sanity(
+        num_nodes, dgl.distributed.optim.SparseAdam, name="SomeEmbedding"
+    )
+
+    dist_embedding_check_existing(num_nodes)
+
+
+##########################################
+############# DistOptimizer ##############
+##########################################
+
+
+def dist_optimizer_check_store(g):
+    num_nodes = g.num_nodes(g.ntypes[0])
+    rank = g.rank()
+    try:
+        emb = dgl.distributed.DistEmbedding(
+            num_nodes, 1, name="optimizer_test", init_func=zeros_init
+        )
+        emb2 = dgl.distributed.DistEmbedding(
+            num_nodes, 5, name="optimizer_test2", init_func=zeros_init
+        )
+        emb_optimizer = dgl.distributed.optim.SparseAdam([emb, emb2], lr=0.1)
+        if rank == 0:
+            name_to_state = {}
+            for _, emb_states in emb_optimizer._state.items():
+                for state in emb_states:
+                    name_to_state[state.name] = F.uniform(
+                        state.shape, F.float32, F.cpu(), 0, 1
+                    )
+                    state[
+                        F.arange(0, num_nodes, F.int64, F.cpu())
+                    ] = name_to_state[state.name]
+        emb_optimizer.save("emb.pt")
+        new_emb_optimizer = dgl.distributed.optim.SparseAdam(
+            [emb, emb2], lr=000.1, eps=2e-08, betas=(0.1, 0.222)
+        )
+        new_emb_optimizer.load("emb.pt")
+        if rank == 0:
+            for _, emb_states in new_emb_optimizer._state.items():
+                for new_state in emb_states:
+                    state = name_to_state[new_state.name]
+                    new_state = new_state[
+                        F.arange(0, num_nodes, F.int64, F.cpu())
+                    ]
+                    assert F.allclose(state, new_state, 0.0, 0.0)
+            assert new_emb_optimizer._lr == emb_optimizer._lr
+            assert new_emb_optimizer._eps == emb_optimizer._eps
+            assert new_emb_optimizer._beta1 == emb_optimizer._beta1
+            assert new_emb_optimizer._beta2 == emb_optimizer._beta2
+        g.barrier()
+    finally:
+        file = f"emb.pt_{rank}"
+        if os.path.exists(file):
+            os.remove(file)
+
+
+def test_dist_optimizer(g):
+    dist_optimizer_check_store(g)
+
+
+##########################################
+############# DistDataLoader #############
+##########################################
+
+
+class NeighborSampler(object):
+    def __init__(self, g, fanouts, sample_neighbors):
+        self.g = g
+        self.fanouts = fanouts
+        self.sample_neighbors = sample_neighbors
+
+    def sample_blocks(self, seeds):
+        import torch as th
+
+        seeds = th.LongTensor(np.asarray(seeds))
+        blocks = []
+        for fanout in self.fanouts:
+            # For each seed node, sample ``fanout`` neighbors.
+            frontier = self.sample_neighbors(
+                self.g, seeds, fanout, replace=True
+            )
+            # Then we compact the frontier into a bipartite graph for
+            # message passing.
+            block = dgl.to_block(frontier, seeds)
+            # Obtain the seed nodes for next layer.
+            seeds = block.srcdata[dgl.NID]
+            block.edata["original_eids"] = frontier.edata[dgl.EID]
+
+            blocks.insert(0, block)
+        return blocks
+
+
+def distdataloader_test(g, batch_size, drop_last, shuffle):
+    # We sample only a subset to minimize the test runtime
+    num_nodes_to_sample = int(g.num_nodes() * 0.05)
+    # To make sure that drop_last is tested
+    if num_nodes_to_sample % batch_size == 0:
+        num_nodes_to_sample -= 1
+
+    orig_nid_map = dict()
+    dtype = g.edges[g.etypes[0]].data["edge_u"].dtype
+    for ntype in g.ntypes:
+        orig_nid = F.tensor(
+            np.load(graph_path + f"/orig_nid_array_{ntype}.npy"), dtype
+        )
+        orig_nid_map[ntype] = orig_nid
+
+    orig_uv_map = dict()
+    for etype in g.etypes:
+        orig_uv_map[etype] = (
+            g.edges[etype].data["edge_u"],
+            g.edges[etype].data["edge_v"],
+        )
+
+    if len(g.ntypes) == 1:
+        train_nid = F.arange(0, num_nodes_to_sample)
+    else:
+        train_nid = {g.ntypes[0]: F.arange(0, num_nodes_to_sample)}
+
+    sampler = NeighborSampler(g, [5, 10], dgl.distributed.sample_neighbors)
+
+    dataloader = dgl.dataloading.DistDataLoader(
+        dataset=train_nid.numpy(),
+        batch_size=batch_size,
+        collate_fn=sampler.sample_blocks,
+        shuffle=shuffle,
+        drop_last=drop_last,
+    )
+
+    for _ in range(2):
+        max_nid = []
+        for idx, blocks in zip(
+            range(0, num_nodes_to_sample, batch_size), dataloader
+        ):
+            block = blocks[-1]
+            for src_type, etype, dst_type in block.canonical_etypes:
+                orig_u, orig_v = orig_uv_map[etype]
+                o_src, o_dst = block.edges(etype=etype)
+                src_nodes_id = block.srcnodes[src_type].data[dgl.NID][o_src]
+                dst_nodes_id = block.dstnodes[dst_type].data[dgl.NID][o_dst]
+                max_nid.append(np.max(F.asnumpy(dst_nodes_id)))
+
+                src_nodes_id = orig_nid_map[src_type][src_nodes_id]
+                dst_nodes_id = orig_nid_map[dst_type][dst_nodes_id]
+                eids = block.edata["original_eids"]
+                F.allclose(src_nodes_id, orig_u[eids])
+                F.allclose(dst_nodes_id, orig_v[eids])
+        if not shuffle and len(max_nid) > 0:
+            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
+
+
+def distnodedataloader_test(
+    g, batch_size, drop_last, shuffle, num_workers, orig_nid_map, orig_uv_map
+):
+    # We sample only a subset to minimize the test runtime
+    num_nodes_to_sample = int(g.num_nodes(g.ntypes[-1]) * 0.05)
+    # To make sure that drop_last is tested
+    if num_nodes_to_sample % batch_size == 0:
+        num_nodes_to_sample -= 1
+
+    if len(g.ntypes) == 1:
+        train_nid = F.arange(0, num_nodes_to_sample)
+    else:
+        train_nid = {g.ntypes[-1]: F.arange(0, num_nodes_to_sample)}
+
+    if len(g.etypes) > 1:
+        sampler = dgl.dataloading.MultiLayerNeighborSampler(
+            [
+                {etype: 5 for etype in g.etypes},
+                10,
+            ]
+        )
+    else:
+        sampler = dgl.dataloading.MultiLayerNeighborSampler(
+            [
+                5,
+                10,
+            ]
+        )
+
+    dataloader = dgl.dataloading.DistNodeDataLoader(
+        g,
+        train_nid,
+        sampler,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        drop_last=drop_last,
+        num_workers=num_workers,
+    )
+
+    for _ in range(2):
+        for _, (_, _, blocks) in zip(
+            range(0, num_nodes_to_sample, batch_size), dataloader
+        ):
+            block = blocks[-1]
+            for src_type, etype, dst_type in block.canonical_etypes:
+                orig_u, orig_v = orig_uv_map[etype]
+                o_src, o_dst = block.edges(etype=etype)
+                src_nodes_id = block.srcnodes[src_type].data[dgl.NID][o_src]
+                dst_nodes_id = block.dstnodes[dst_type].data[dgl.NID][o_dst]
+                src_nodes_id = orig_nid_map[src_type][src_nodes_id]
+                dst_nodes_id = orig_nid_map[dst_type][dst_nodes_id]
+                eids = block.edges[etype].data[dgl.EID]
+                F.allclose(src_nodes_id, orig_u[eids])
+                F.allclose(dst_nodes_id, orig_v[eids])
+    del dataloader
+
+
+def distedgedataloader_test(
+    g,
+    batch_size,
+    drop_last,
+    shuffle,
+    num_workers,
+    orig_nid_map,
+    orig_uv_map,
+    num_negs,
+):
+    # We sample only a subset to minimize the test runtime
+    num_edges_to_sample = int(g.num_edges(g.etypes[-1]) * 0.05)
+    # To make sure that drop_last is tested
+    if num_edges_to_sample % batch_size == 0:
+        num_edges_to_sample -= 1
+
+    if len(g.etypes) == 1:
+        train_eid = F.arange(0, num_edges_to_sample)
+    else:
+        train_eid = {g.etypes[-1]: F.arange(0, num_edges_to_sample)}
+
+    sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10])
+
+    dataloader = dgl.dataloading.DistEdgeDataLoader(
+        g,
+        train_eid,
+        sampler,
+        batch_size=batch_size,
+        negative_sampler=dgl.dataloading.negative_sampler.Uniform(num_negs)
+        if num_negs > 0
+        else None,
+        shuffle=shuffle,
+        drop_last=drop_last,
+        num_workers=num_workers,
+    )
+    for _ in range(2):
+        for _, sampled_data in zip(
+            range(0, num_edges_to_sample, batch_size), dataloader
+        ):
+            blocks = sampled_data[3 if num_negs > 0 else 2]
+            block = blocks[-1]
+            for src_type, etype, dst_type in block.canonical_etypes:
+                orig_u, orig_v = orig_uv_map[etype]
+                o_src, o_dst = block.edges(etype=etype)
+                src_nodes_id = block.srcnodes[src_type].data[dgl.NID][o_src]
+                dst_nodes_id = block.dstnodes[dst_type].data[dgl.NID][o_dst]
+                src_nodes_id = orig_nid_map[src_type][src_nodes_id]
+                dst_nodes_id = orig_nid_map[dst_type][dst_nodes_id]
+                eids = block.edges[etype].data[dgl.EID]
+                F.allclose(src_nodes_id, orig_u[eids])
+                F.allclose(dst_nodes_id, orig_v[eids])
+                if num_negs == 0:
+                    pos_pair_graph = sampled_data[1]
+                    assert np.all(
+                        F.asnumpy(block.dstnodes[dst_type].data[dgl.NID])
+                        == F.asnumpy(
+                            pos_pair_graph.nodes[dst_type].data[dgl.NID]
+                        )
+                    )
+                else:
+                    pos_graph, neg_graph = sampled_data[1:3]
+                    assert np.all(
+                        F.asnumpy(block.dstnodes[dst_type].data[dgl.NID])
+                        == F.asnumpy(pos_graph.nodes[dst_type].data[dgl.NID])
+                    )
+                    assert np.all(
+                        F.asnumpy(block.dstnodes[dst_type].data[dgl.NID])
+                        == F.asnumpy(neg_graph.nodes[dst_type].data[dgl.NID])
+                    )
+                    assert (
+                        pos_graph.num_edges() * num_negs
+                        == neg_graph.num_edges()
+                    )
+    del dataloader
+
+
+def multi_distdataloader_test(g, dataloader_class):
+    total_num_items = (
+        g.num_nodes(g.ntypes[-1])
+        if "Node" in dataloader_class.__name__
+        else g.num_edges(g.etypes[-1])
+    )
+
+    num_dataloaders = 4
+    batch_size = 32
+    sampler = dgl.dataloading.NeighborSampler([-1])
+    dataloaders = []
+    dl_iters = []
+
+    # We sample only a subset to minimize the test runtime
+    num_items_to_sample = int(total_num_items * 0.05)
+    # To make sure that drop_last is tested
+    if num_items_to_sample % batch_size == 0:
+        num_items_to_sample -= 1
+
+    if len(g.ntypes) == 1:
+        train_ids = F.arange(0, num_items_to_sample)
+    else:
+        train_ids = {
+            g.ntypes[-1]
+            if "Node" in dataloader_class.__name__
+            else g.etypes[-1]: F.arange(0, num_items_to_sample)
+        }
+
+    for _ in range(num_dataloaders):
+        dataloader = dataloader_class(
+            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:
+        current_dl = np.random.choice(len(dl_iters), 1)[0]
+        try:
+            _ = next(dl_iters[current_dl])
+        except StopIteration:
+            dl_iters.pop(current_dl)
+            del dataloaders[current_dl]
+
+
+def test_dist_dataloader(g):
+    orig_nid_map = dict()
+    dtype = g.edges[g.etypes[0]].data["edge_u"].dtype
+    for ntype in g.ntypes:
+        orig_nid = F.tensor(
+            np.load(graph_path + f"/orig_nid_array_{ntype}.npy"), dtype
+        )
+        orig_nid_map[ntype] = orig_nid
+
+    orig_uv_map = dict()
+    for etype in g.etypes:
+        orig_uv_map[etype] = (
+            g.edges[etype].data["edge_u"],
+            g.edges[etype].data["edge_v"],
+        )
+
+    batch_size_l = [64]
+    drop_last_l = [False, True]
+    num_workers_l = [0, 4]
+    shuffle_l = [False, True]
+
+    for batch_size, drop_last, shuffle, num_workers in product(
+        batch_size_l, drop_last_l, shuffle_l, num_workers_l
+    ):
+        if len(g.ntypes) == 1 and num_workers == 0:
+            distdataloader_test(g, batch_size, drop_last, shuffle)
+        distnodedataloader_test(
+            g,
+            batch_size,
+            drop_last,
+            shuffle,
+            num_workers,
+            orig_nid_map,
+            orig_uv_map,
+        )
+        # No negssampling
+        distedgedataloader_test(
+            g,
+            batch_size,
+            drop_last,
+            shuffle,
+            num_workers,
+            orig_nid_map,
+            orig_uv_map,
+            num_negs=0,
+        )
+        # negsampling 15
+        distedgedataloader_test(
+            g,
+            batch_size,
+            drop_last,
+            shuffle,
+            num_workers,
+            orig_nid_map,
+            orig_uv_map,
+            num_negs=15,
+        )
+
+    multi_distdataloader_test(g, dgl.dataloading.DistNodeDataLoader)
+    multi_distdataloader_test(g, dgl.dataloading.DistEdgeDataLoader)
+
+
+if mode == "server":
+    shared_mem = bool(int(os.environ.get("DIST_DGL_TEST_SHARED_MEM")))
+    server_id = int(os.environ.get("DIST_DGL_TEST_SERVER_ID"))
+    run_server(
+        graph_name,
+        server_id,
+        server_count=num_servers_per_machine,
+        num_clients=num_part * num_client_per_machine,
+        shared_mem=shared_mem,
+    )
+elif mode == "client":
+    os.environ["DGL_NUM_SERVER"] = str(num_servers_per_machine)
+    dgl.distributed.initialize(ip_config)
+
+    gpb, graph_name, _, _ = load_partition_book(
+        graph_path + "/{}.json".format(graph_name), part_id
+    )
+    g = dgl.distributed.DistGraph(graph_name, gpb=gpb)
+
+    target_func_map = {
+        "DistGraph": test_dist_graph,
+        "DistGraphServices": test_dist_graph_services,
+        "DistTensor": test_dist_tensor,
+        "DistEmbedding": test_dist_embedding,
+        "DistOptimizer": test_dist_optimizer,
+        "DistDataLoader": test_dist_dataloader,
+    }
+
+    targets = os.environ.get("DIST_DGL_TEST_OBJECT_TYPE", "")
+    targets = targets.replace(" ", "").split(",") if targets else []
+    blacklist = os.environ.get("DIST_DGL_TEST_OBJECT_TYPE_BLACKLIST", "")
+    blacklist = blacklist.replace(" ", "").split(",") if blacklist else []
+
+    for to_bl in blacklist:
+        target_func_map.pop(to_bl, None)
+
+    if not targets:
+        for test_func in target_func_map.values():
+            test_func(g)
+    else:
+        for target in targets:
+            if target in target_func_map:
+                target_func_map[target](g)
+            else:
+                print(f"Tests not implemented for target '{target}'")
+
+else:
+    exit(1)