test_dist_objects.py 5.46 KB
Newer Older
1
import multiprocessing as mp
2
import os
3
import subprocess
4
import unittest
5
6

import numpy as np
7
8
import pytest
import utils
9
import shutil
10

11
12
13
14
import dgl
import dgl.backend as F
from dgl.distributed import partition_graph

15
16
graph_name = os.environ.get("DIST_DGL_TEST_GRAPH_NAME", "random_test_graph")
target = os.environ.get("DIST_DGL_TEST_OBJECT_TYPE", "")
17
shared_workspace = os.environ.get('DIST_DGL_TEST_WORKSPACE', '/shared_workspace/dgl_dist_tensor_test/')
18

19
20
21
22

def create_graph(num_part, dist_graph_path, hetero):
    if not hetero:
        g = dgl.rand_graph(10000, 42000)
23
24
        g.ndata["feat"] = F.unsqueeze(F.arange(0, g.number_of_nodes()), 1)
        g.edata["feat"] = F.unsqueeze(F.arange(0, g.number_of_edges()), 1)
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
        g.ndata["in_degrees"] = g.in_degrees()
        g.ndata["out_degrees"] = g.out_degrees()

        etype = g.etypes[0]
        ntype = g.ntypes[0]
        edge_u, edge_v = g.find_edges(F.arange(0, g.number_of_edges(etype)))
        g.edges[etype].data["edge_u"] = edge_u
        g.edges[etype].data["edge_v"] = edge_v

        orig_nid, orig_eid = partition_graph(g, graph_name, num_part, dist_graph_path, return_mapping=True)

        orig_nid_f = os.path.join(dist_graph_path, f"orig_nid_array_{ntype}.npy")
        np.save(orig_nid_f, orig_nid.numpy())
        orig_eid_f = os.path.join(dist_graph_path, f"orig_eid_array_{etype}.npy")
        np.save(orig_eid_f, orig_eid.numpy())


42
43
    else:
        from scipy import sparse as spsp
44
45
46

        num_nodes = {"n1": 10000, "n2": 10010, "n3": 10020}
        etypes = [("n1", "r1", "n2"), ("n1", "r2", "n3"), ("n2", "r3", "n3")]
47
48
49
        edges = {}
        for etype in etypes:
            src_ntype, _, dst_ntype = etype
50
51
52
53
54
55
56
            arr = spsp.random(
                num_nodes[src_ntype],
                num_nodes[dst_ntype],
                density=0.001,
                format="coo",
                random_state=100,
            )
57
58
            edges[etype] = (arr.row, arr.col)
        g = dgl.heterograph(edges, num_nodes)
59

60
61
62
63
64
65
        g.nodes["n1"].data["feat"] = F.unsqueeze(
            F.arange(0, g.number_of_nodes("n1")), 1
        )
        g.edges["r1"].data["feat"] = F.unsqueeze(
            F.arange(0, g.number_of_edges("r1")), 1
        )
66
67


68
69
70
71
72
73
74
75
76
77
78
79
80
81
        for _, etype, _ in etypes:
            edge_u, edge_v = g.find_edges(F.arange(0, g.number_of_edges(etype)))
            g.edges[etype].data["edge_u"] = edge_u
            g.edges[etype].data["edge_v"] = edge_v

        orig_nid, orig_eid = partition_graph(g, graph_name, num_part, dist_graph_path, return_mapping=True)

        for n_type, tensor in orig_nid.items():
            orig_nid_f = os.path.join(dist_graph_path, f"orig_nid_array_{n_type}.npy")
            np.save(orig_nid_f, tensor.numpy())
        for e_type, tensor in orig_eid.items():
            orig_eid_f = os.path.join(dist_graph_path, f"orig_eid_array_{e_type}.npy")
            np.save(orig_eid_f, tensor.numpy())

82
83
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@pytest.mark.parametrize("net_type", ["tensorpipe", "socket"])
84
85
86
87
@pytest.mark.parametrize("num_servers", [1, 4])
@pytest.mark.parametrize("num_clients", [1, 4])
@pytest.mark.parametrize("hetero", [False, True])
@pytest.mark.parametrize("shared_mem", [False, True])
88
def test_dist_objects(net_type, num_servers, num_clients, hetero, shared_mem):
89
    if not shared_mem and num_servers > 1:
90
91
92
93
        pytest.skip(
            f"Backup servers are not supported when shared memory is disabled"
        )
    ip_config = os.environ.get("DIST_DGL_TEST_IP_CONFIG", "ip_config.txt")
94
95
96
97

    ips = utils.get_ips(ip_config)
    num_part = len(ips)

98
99
100
    test_bin = os.path.join(
        os.environ.get("DIST_DGL_TEST_PY_BIN_DIR", "."), "run_dist_objects.py"
    )
101

102
    dist_graph_path = os.path.join(
103
        shared_workspace, "hetero_dist_graph" if hetero else "dist_graph"
104
    )
105
106
107
    if not os.path.isdir(dist_graph_path):
        create_graph(num_part, dist_graph_path, hetero)

108
    base_envs = (
109
        f"DIST_DGL_TEST_WORKSPACE={shared_workspace} "
110
111
112
113
114
115
116
        f"DIST_DGL_TEST_NUM_PART={num_part} "
        f"DIST_DGL_TEST_NUM_SERVER={num_servers} "
        f"DIST_DGL_TEST_NUM_CLIENT={num_clients} "
        f"DIST_DGL_TEST_NET_TYPE={net_type} "
        f"DIST_DGL_TEST_GRAPH_PATH={dist_graph_path} "
        f"DIST_DGL_TEST_IP_CONFIG={ip_config} "
    )
117
118
119
120
121
122

    procs = []
    # Start server
    server_id = 0
    for part_id, ip in enumerate(ips):
        for _ in range(num_servers):
123
124
125
126
127
128
129
130
131
            cmd_envs = (
                base_envs + f"DIST_DGL_TEST_SERVER_ID={server_id} "
                f"DIST_DGL_TEST_PART_ID={part_id} "
                f"DIST_DGL_TEST_SHARED_MEM={str(int(shared_mem))} "
                f"DIST_DGL_TEST_MODE=server "
            )
            procs.append(
                utils.execute_remote(f"{cmd_envs} python3 {test_bin}", ip)
            )
132
133
134
135
            server_id += 1
    # Start client processes
    for part_id, ip in enumerate(ips):
        for _ in range(num_clients):
136
137
138
139
140
141
142
143
            cmd_envs = (
                base_envs + f"DIST_DGL_TEST_PART_ID={part_id} "
                f"DIST_DGL_TEST_OBJECT_TYPE={target} "
                f"DIST_DGL_TEST_MODE=client "
            )
            procs.append(
                utils.execute_remote(f"{cmd_envs} python3 {test_bin}", ip)
            )
144
145
146
147

    for p in procs:
        p.join()
        assert p.exitcode == 0
148
149
150
151
152
153
154

def teardown():
    for name in ['dist_graph', 'hetero_dist_graph']:
        path = os.path.join(shared_workspace, name)
        if os.path.exists(path):
            print(f"Removing {path}...")
            shutil.rmtree(path)