chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,97 @@
|
||||
import operator
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
import pytest
|
||||
from utils import create_random_graph, generate_ip_config, reset_envs
|
||||
|
||||
dist_g = None
|
||||
|
||||
|
||||
def rand_mask(shape, dtype):
|
||||
return F.randn(shape) > 0
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
dgl.backend.backend_name == "tensorflow",
|
||||
reason="TF doesn't support some of operations in DistGraph",
|
||||
)
|
||||
@unittest.skipIf(
|
||||
dgl.backend.backend_name == "mxnet", reason="Turn off Mxnet support"
|
||||
)
|
||||
def setup_module():
|
||||
global dist_g
|
||||
|
||||
reset_envs()
|
||||
os.environ["DGL_DIST_MODE"] = "standalone"
|
||||
|
||||
dist_g = create_random_graph(10000)
|
||||
# Partition the graph.
|
||||
num_parts = 1
|
||||
graph_name = "dist_graph_test_3"
|
||||
dist_g.ndata["features"] = F.unsqueeze(F.arange(0, dist_g.num_nodes()), 1)
|
||||
dist_g.edata["features"] = F.unsqueeze(F.arange(0, dist_g.num_edges()), 1)
|
||||
dgl.distributed.partition_graph(
|
||||
dist_g, graph_name, num_parts, "/tmp/dist_graph"
|
||||
)
|
||||
|
||||
dgl.distributed.initialize("kv_ip_config.txt")
|
||||
dist_g = dgl.distributed.DistGraph(
|
||||
graph_name, part_config="/tmp/dist_graph/{}.json".format(graph_name)
|
||||
)
|
||||
dist_g.edata["mask1"] = dgl.distributed.DistTensor(
|
||||
(dist_g.num_edges(),), F.bool, init_func=rand_mask
|
||||
)
|
||||
dist_g.edata["mask2"] = dgl.distributed.DistTensor(
|
||||
(dist_g.num_edges(),), F.bool, init_func=rand_mask
|
||||
)
|
||||
|
||||
|
||||
def check_binary_op(key1, key2, key3, op):
|
||||
for i in range(0, dist_g.num_edges(), 1000):
|
||||
i_end = min(i + 1000, dist_g.num_edges())
|
||||
assert F.array_equal(
|
||||
dist_g.edata[key3][i:i_end],
|
||||
op(dist_g.edata[key1][i:i_end], dist_g.edata[key2][i:i_end]),
|
||||
)
|
||||
# Test with different index dtypes. int32 is not supported.
|
||||
with pytest.raises(
|
||||
dgl.utils.internal.InconsistentDtypeException,
|
||||
match="DGL now requires the input tensor to have",
|
||||
):
|
||||
_ = dist_g.edata[key3][F.tensor([100, 20, 10], F.int32)]
|
||||
_ = dist_g.edata[key3][F.tensor([100, 20, 10], F.int64)]
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
dgl.backend.backend_name == "tensorflow",
|
||||
reason="TF doesn't support some of operations in DistGraph",
|
||||
)
|
||||
@unittest.skipIf(
|
||||
dgl.backend.backend_name == "mxnet", reason="Turn off Mxnet support"
|
||||
)
|
||||
def test_op():
|
||||
dist_g.edata["mask3"] = dist_g.edata["mask1"] | dist_g.edata["mask2"]
|
||||
check_binary_op("mask1", "mask2", "mask3", operator.or_)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
dgl.backend.backend_name == "tensorflow",
|
||||
reason="TF doesn't support some of operations in DistGraph",
|
||||
)
|
||||
@unittest.skipIf(
|
||||
dgl.backend.backend_name == "mxnet", reason="Turn off Mxnet support"
|
||||
)
|
||||
def teardown_module():
|
||||
# Since there are two tests in one process, this is needed to make sure
|
||||
# the client exits properly.
|
||||
dgl.distributed.exit_client()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
setup_module()
|
||||
test_op()
|
||||
teardown_module()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,390 @@
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
from numpy.testing import assert_array_equal
|
||||
from utils import generate_ip_config, reset_envs
|
||||
|
||||
|
||||
# Create an one-part Graph
|
||||
node_map = {"_N": F.tensor([[0, 6]], F.int64)}
|
||||
edge_map = {("_N", "_E", "_N"): F.tensor([[0, 7]], F.int64)}
|
||||
global_nid = F.tensor([0, 1, 2, 3, 4, 5], F.int64)
|
||||
global_eid = F.tensor([0, 1, 2, 3, 4, 5, 6], F.int64)
|
||||
|
||||
g = dgl.graph([])
|
||||
g.add_nodes(6)
|
||||
g.add_edges(0, 1) # 0
|
||||
g.add_edges(0, 2) # 1
|
||||
g.add_edges(0, 3) # 2
|
||||
g.add_edges(2, 3) # 3
|
||||
g.add_edges(1, 1) # 4
|
||||
g.add_edges(0, 4) # 5
|
||||
g.add_edges(2, 5) # 6
|
||||
|
||||
g.ndata[dgl.NID] = global_nid
|
||||
g.edata[dgl.EID] = global_eid
|
||||
|
||||
gpb = dgl.distributed.graph_partition_book.RangePartitionBook(
|
||||
part_id=0,
|
||||
num_parts=1,
|
||||
node_map=node_map,
|
||||
edge_map=edge_map,
|
||||
ntypes={ntype: i for i, ntype in enumerate(g.ntypes)},
|
||||
etypes={etype: i for i, etype in enumerate(g.canonical_etypes)},
|
||||
)
|
||||
|
||||
node_policy = dgl.distributed.PartitionPolicy(
|
||||
policy_str="node~_N", partition_book=gpb
|
||||
)
|
||||
|
||||
edge_policy = dgl.distributed.PartitionPolicy(
|
||||
policy_str="edge~_N:_E:_N", partition_book=gpb
|
||||
)
|
||||
|
||||
data_0 = F.tensor(
|
||||
[[1.0, 1.0], [1.0, 1.0], [1.0, 1.0], [1.0, 1.0], [1.0, 1.0], [1.0, 1.0]],
|
||||
F.float32,
|
||||
)
|
||||
data_0_1 = F.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], F.float32)
|
||||
data_0_2 = F.tensor([1, 2, 3, 4, 5, 6], F.int32)
|
||||
data_0_3 = F.tensor([1, 2, 3, 4, 5, 6], F.int64)
|
||||
data_1 = F.tensor(
|
||||
[
|
||||
[2.0, 2.0],
|
||||
[2.0, 2.0],
|
||||
[2.0, 2.0],
|
||||
[2.0, 2.0],
|
||||
[2.0, 2.0],
|
||||
[2.0, 2.0],
|
||||
[2.0, 2.0],
|
||||
],
|
||||
F.float32,
|
||||
)
|
||||
data_2 = F.tensor(
|
||||
[[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
|
||||
F.float32,
|
||||
)
|
||||
|
||||
|
||||
def init_zero_func(shape, dtype):
|
||||
return F.zeros(shape, dtype, F.cpu())
|
||||
|
||||
|
||||
def udf_push(target, name, id_tensor, data_tensor):
|
||||
target[name][id_tensor] = data_tensor * data_tensor
|
||||
|
||||
|
||||
def add_push(target, name, id_tensor, data_tensor):
|
||||
target[name][id_tensor] += data_tensor
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
os.name == "nt" or os.getenv("DGLBACKEND") == "tensorflow",
|
||||
reason="Do not support windows and TF yet",
|
||||
)
|
||||
def test_partition_policy():
|
||||
assert node_policy.part_id == 0
|
||||
assert edge_policy.part_id == 0
|
||||
local_nid = node_policy.to_local(F.tensor([0, 1, 2, 3, 4, 5]))
|
||||
local_eid = edge_policy.to_local(F.tensor([0, 1, 2, 3, 4, 5, 6]))
|
||||
assert_array_equal(
|
||||
F.asnumpy(local_nid), F.asnumpy(F.tensor([0, 1, 2, 3, 4, 5], F.int64))
|
||||
)
|
||||
assert_array_equal(
|
||||
F.asnumpy(local_eid),
|
||||
F.asnumpy(F.tensor([0, 1, 2, 3, 4, 5, 6], F.int64)),
|
||||
)
|
||||
nid_partid = node_policy.to_partid(F.tensor([0, 1, 2, 3, 4, 5], F.int64))
|
||||
eid_partid = edge_policy.to_partid(F.tensor([0, 1, 2, 3, 4, 5, 6], F.int64))
|
||||
assert_array_equal(
|
||||
F.asnumpy(nid_partid), F.asnumpy(F.tensor([0, 0, 0, 0, 0, 0], F.int64))
|
||||
)
|
||||
assert_array_equal(
|
||||
F.asnumpy(eid_partid),
|
||||
F.asnumpy(F.tensor([0, 0, 0, 0, 0, 0, 0], F.int64)),
|
||||
)
|
||||
assert node_policy.get_part_size() == len(local_nid)
|
||||
assert edge_policy.get_part_size() == len(local_eid)
|
||||
|
||||
|
||||
def start_server(server_id, num_clients, num_servers):
|
||||
# Init kvserver
|
||||
print("Sleep 5 seconds to test client re-connect.")
|
||||
time.sleep(5)
|
||||
kvserver = dgl.distributed.KVServer(
|
||||
server_id=server_id,
|
||||
ip_config="kv_ip_config.txt",
|
||||
num_servers=num_servers,
|
||||
num_clients=num_clients,
|
||||
)
|
||||
kvserver.add_part_policy(node_policy)
|
||||
kvserver.add_part_policy(edge_policy)
|
||||
if kvserver.is_backup_server():
|
||||
kvserver.init_data("data_0", "node~_N")
|
||||
kvserver.init_data("data_0_1", "node~_N")
|
||||
kvserver.init_data("data_0_2", "node~_N")
|
||||
kvserver.init_data("data_0_3", "node~_N")
|
||||
else:
|
||||
kvserver.init_data("data_0", "node~_N", data_0)
|
||||
kvserver.init_data("data_0_1", "node~_N", data_0_1)
|
||||
kvserver.init_data("data_0_2", "node~_N", data_0_2)
|
||||
kvserver.init_data("data_0_3", "node~_N", data_0_3)
|
||||
# start server
|
||||
server_state = dgl.distributed.ServerState(
|
||||
kv_store=kvserver, local_g=None, partition_book=None
|
||||
)
|
||||
dgl.distributed.start_server(
|
||||
server_id=server_id,
|
||||
ip_config="kv_ip_config.txt",
|
||||
num_servers=num_servers,
|
||||
num_clients=num_clients,
|
||||
server_state=server_state,
|
||||
)
|
||||
|
||||
|
||||
def start_server_mul_role(server_id, num_clients, num_servers):
|
||||
# Init kvserver
|
||||
kvserver = dgl.distributed.KVServer(
|
||||
server_id=server_id,
|
||||
ip_config="kv_ip_mul_config.txt",
|
||||
num_servers=num_servers,
|
||||
num_clients=num_clients,
|
||||
)
|
||||
kvserver.add_part_policy(node_policy)
|
||||
if kvserver.is_backup_server():
|
||||
kvserver.init_data("data_0", "node~_N")
|
||||
else:
|
||||
kvserver.init_data("data_0", "node~_N", data_0)
|
||||
# start server
|
||||
server_state = dgl.distributed.ServerState(
|
||||
kv_store=kvserver, local_g=None, partition_book=None
|
||||
)
|
||||
dgl.distributed.start_server(
|
||||
server_id=server_id,
|
||||
ip_config="kv_ip_mul_config.txt",
|
||||
num_servers=num_servers,
|
||||
num_clients=num_clients,
|
||||
server_state=server_state,
|
||||
)
|
||||
|
||||
|
||||
def start_client(num_clients, num_servers):
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
# Note: connect to server first !
|
||||
dgl.distributed.initialize(ip_config="kv_ip_config.txt")
|
||||
# Init kvclient
|
||||
kvclient = dgl.distributed.KVClient(
|
||||
ip_config="kv_ip_config.txt", num_servers=num_servers
|
||||
)
|
||||
kvclient.map_shared_data(partition_book=gpb)
|
||||
assert dgl.distributed.get_num_client() == num_clients
|
||||
kvclient.init_data(
|
||||
name="data_1",
|
||||
shape=F.shape(data_1),
|
||||
dtype=F.dtype(data_1),
|
||||
part_policy=edge_policy,
|
||||
init_func=init_zero_func,
|
||||
)
|
||||
kvclient.init_data(
|
||||
name="data_2",
|
||||
shape=F.shape(data_2),
|
||||
dtype=F.dtype(data_2),
|
||||
part_policy=node_policy,
|
||||
init_func=init_zero_func,
|
||||
)
|
||||
|
||||
# Test data_name_list
|
||||
name_list = kvclient.data_name_list()
|
||||
print(name_list)
|
||||
assert "data_0" in name_list
|
||||
assert "data_0_1" in name_list
|
||||
assert "data_0_2" in name_list
|
||||
assert "data_0_3" in name_list
|
||||
assert "data_1" in name_list
|
||||
assert "data_2" in name_list
|
||||
# Test get_meta_data
|
||||
meta = kvclient.get_data_meta("data_0")
|
||||
dtype, shape, policy = meta
|
||||
assert dtype == F.dtype(data_0)
|
||||
assert shape == F.shape(data_0)
|
||||
assert policy.policy_str == "node~_N"
|
||||
|
||||
meta = kvclient.get_data_meta("data_0_1")
|
||||
dtype, shape, policy = meta
|
||||
assert dtype == F.dtype(data_0_1)
|
||||
assert shape == F.shape(data_0_1)
|
||||
assert policy.policy_str == "node~_N"
|
||||
|
||||
meta = kvclient.get_data_meta("data_0_2")
|
||||
dtype, shape, policy = meta
|
||||
assert dtype == F.dtype(data_0_2)
|
||||
assert shape == F.shape(data_0_2)
|
||||
assert policy.policy_str == "node~_N"
|
||||
|
||||
meta = kvclient.get_data_meta("data_0_3")
|
||||
dtype, shape, policy = meta
|
||||
assert dtype == F.dtype(data_0_3)
|
||||
assert shape == F.shape(data_0_3)
|
||||
assert policy.policy_str == "node~_N"
|
||||
|
||||
meta = kvclient.get_data_meta("data_1")
|
||||
dtype, shape, policy = meta
|
||||
assert dtype == F.dtype(data_1)
|
||||
assert shape == F.shape(data_1)
|
||||
assert policy.policy_str == "edge~_N:_E:_N"
|
||||
|
||||
meta = kvclient.get_data_meta("data_2")
|
||||
dtype, shape, policy = meta
|
||||
assert dtype == F.dtype(data_2)
|
||||
assert shape == F.shape(data_2)
|
||||
assert policy.policy_str == "node~_N"
|
||||
|
||||
# Test push and pull
|
||||
id_tensor = F.tensor([0, 2, 4], F.int64)
|
||||
data_tensor = F.tensor([[6.0, 6.0], [6.0, 6.0], [6.0, 6.0]], F.float32)
|
||||
kvclient.push(name="data_0", id_tensor=id_tensor, data_tensor=data_tensor)
|
||||
kvclient.push(name="data_1", id_tensor=id_tensor, data_tensor=data_tensor)
|
||||
kvclient.push(name="data_2", id_tensor=id_tensor, data_tensor=data_tensor)
|
||||
res = kvclient.pull(name="data_0", id_tensor=id_tensor)
|
||||
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
|
||||
res = kvclient.pull(name="data_1", id_tensor=id_tensor)
|
||||
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
|
||||
res = kvclient.pull(name="data_2", id_tensor=id_tensor)
|
||||
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
|
||||
# Register new push handler
|
||||
kvclient.register_push_handler("data_0", udf_push)
|
||||
kvclient.register_push_handler("data_1", udf_push)
|
||||
kvclient.register_push_handler("data_2", udf_push)
|
||||
# Test push and pull
|
||||
kvclient.push(name="data_0", id_tensor=id_tensor, data_tensor=data_tensor)
|
||||
kvclient.push(name="data_1", id_tensor=id_tensor, data_tensor=data_tensor)
|
||||
kvclient.push(name="data_2", id_tensor=id_tensor, data_tensor=data_tensor)
|
||||
kvclient.barrier()
|
||||
data_tensor = data_tensor * data_tensor
|
||||
res = kvclient.pull(name="data_0", id_tensor=id_tensor)
|
||||
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
|
||||
res = kvclient.pull(name="data_1", id_tensor=id_tensor)
|
||||
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
|
||||
res = kvclient.pull(name="data_2", id_tensor=id_tensor)
|
||||
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
|
||||
|
||||
# Test delete data
|
||||
kvclient.delete_data("data_0")
|
||||
kvclient.delete_data("data_1")
|
||||
kvclient.delete_data("data_2")
|
||||
|
||||
# Register new push handler
|
||||
kvclient.init_data(
|
||||
name="data_3",
|
||||
shape=F.shape(data_2),
|
||||
dtype=F.dtype(data_2),
|
||||
part_policy=node_policy,
|
||||
init_func=init_zero_func,
|
||||
)
|
||||
kvclient.register_push_handler("data_3", add_push)
|
||||
data_tensor = F.tensor([[6.0, 6.0], [6.0, 6.0], [6.0, 6.0]], F.float32)
|
||||
kvclient.barrier()
|
||||
time.sleep(kvclient.client_id + 1)
|
||||
print("add...")
|
||||
kvclient.push(name="data_3", id_tensor=id_tensor, data_tensor=data_tensor)
|
||||
kvclient.barrier()
|
||||
res = kvclient.pull(name="data_3", id_tensor=id_tensor)
|
||||
data_tensor = data_tensor * num_clients
|
||||
assert_array_equal(F.asnumpy(res), F.asnumpy(data_tensor))
|
||||
|
||||
|
||||
def start_client_mul_role(i):
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
# Initialize creates kvstore !
|
||||
dgl.distributed.initialize(ip_config="kv_ip_mul_config.txt")
|
||||
if i == 0: # block one trainer
|
||||
time.sleep(5)
|
||||
kvclient = dgl.distributed.kvstore.get_kvstore()
|
||||
kvclient.barrier()
|
||||
print("i: %d role: %s" % (i, kvclient.role))
|
||||
|
||||
assert dgl.distributed.role.get_num_trainers() == 2
|
||||
assert dgl.distributed.role.get_trainer_rank() < 2
|
||||
print(
|
||||
"trainer rank: %d, global rank: %d"
|
||||
% (
|
||||
dgl.distributed.role.get_trainer_rank(),
|
||||
dgl.distributed.role.get_global_rank(),
|
||||
)
|
||||
)
|
||||
dgl.distributed.exit_client()
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
os.name == "nt" or os.getenv("DGLBACKEND") == "tensorflow",
|
||||
reason="Do not support windows and TF yet",
|
||||
)
|
||||
def test_kv_store():
|
||||
reset_envs()
|
||||
num_servers = 2
|
||||
num_clients = 2
|
||||
generate_ip_config("kv_ip_config.txt", 1, num_servers)
|
||||
ctx = mp.get_context("spawn")
|
||||
pserver_list = []
|
||||
pclient_list = []
|
||||
os.environ["DGL_NUM_SERVER"] = str(num_servers)
|
||||
for i in range(num_servers):
|
||||
pserver = ctx.Process(
|
||||
target=start_server, args=(i, num_clients, num_servers)
|
||||
)
|
||||
pserver.start()
|
||||
pserver_list.append(pserver)
|
||||
for i in range(num_clients):
|
||||
pclient = ctx.Process(
|
||||
target=start_client, args=(num_clients, num_servers)
|
||||
)
|
||||
pclient.start()
|
||||
pclient_list.append(pclient)
|
||||
for i in range(num_clients):
|
||||
pclient_list[i].join()
|
||||
for i in range(num_servers):
|
||||
pserver_list[i].join()
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
os.name == "nt" or os.getenv("DGLBACKEND") == "tensorflow",
|
||||
reason="Do not support windows and TF yet",
|
||||
)
|
||||
def test_kv_multi_role():
|
||||
reset_envs()
|
||||
num_servers = 2
|
||||
num_trainers = 2
|
||||
num_samplers = 2
|
||||
generate_ip_config("kv_ip_mul_config.txt", 1, num_servers)
|
||||
# There are two trainer processes and each trainer process has two sampler processes.
|
||||
num_clients = num_trainers * (1 + num_samplers)
|
||||
ctx = mp.get_context("spawn")
|
||||
pserver_list = []
|
||||
pclient_list = []
|
||||
os.environ["DGL_NUM_SAMPLER"] = str(num_samplers)
|
||||
os.environ["DGL_NUM_SERVER"] = str(num_servers)
|
||||
for i in range(num_servers):
|
||||
pserver = ctx.Process(
|
||||
target=start_server_mul_role, args=(i, num_clients, num_servers)
|
||||
)
|
||||
pserver.start()
|
||||
pserver_list.append(pserver)
|
||||
for i in range(num_trainers):
|
||||
pclient = ctx.Process(target=start_client_mul_role, args=(i,))
|
||||
pclient.start()
|
||||
pclient_list.append(pclient)
|
||||
for i in range(num_trainers):
|
||||
pclient_list[i].join()
|
||||
for i in range(num_servers):
|
||||
pserver_list[i].join()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_partition_policy()
|
||||
test_kv_store()
|
||||
test_kv_multi_role()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,425 @@
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
import pytest
|
||||
from numpy.testing import assert_array_equal
|
||||
from utils import generate_ip_config, reset_envs
|
||||
|
||||
if os.name != "nt":
|
||||
import fcntl
|
||||
import struct
|
||||
|
||||
INTEGER = 2
|
||||
STR = "hello world!"
|
||||
HELLO_SERVICE_ID = 901231
|
||||
TENSOR = F.zeros((1000, 1000), F.int64, F.cpu())
|
||||
|
||||
|
||||
def foo(x, y):
|
||||
assert x == 123
|
||||
assert y == "abc"
|
||||
|
||||
|
||||
class MyRequest(dgl.distributed.Request):
|
||||
def __init__(self):
|
||||
self.x = 123
|
||||
self.y = "abc"
|
||||
self.z = F.randn((3, 4))
|
||||
self.foo = foo
|
||||
|
||||
def __getstate__(self):
|
||||
return self.x, self.y, self.z, self.foo
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.x, self.y, self.z, self.foo = state
|
||||
|
||||
def process_request(self, server_state):
|
||||
pass
|
||||
|
||||
|
||||
class MyResponse(dgl.distributed.Response):
|
||||
def __init__(self):
|
||||
self.x = 432
|
||||
|
||||
def __getstate__(self):
|
||||
return self.x
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.x = state
|
||||
|
||||
|
||||
def simple_func(tensor):
|
||||
return tensor
|
||||
|
||||
|
||||
class HelloResponse(dgl.distributed.Response):
|
||||
def __init__(self, hello_str, integer, tensor):
|
||||
self.hello_str = hello_str
|
||||
self.integer = integer
|
||||
self.tensor = tensor
|
||||
|
||||
def __getstate__(self):
|
||||
return self.hello_str, self.integer, self.tensor
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.hello_str, self.integer, self.tensor = state
|
||||
|
||||
|
||||
class HelloRequest(dgl.distributed.Request):
|
||||
def __init__(self, hello_str, integer, tensor, func):
|
||||
self.hello_str = hello_str
|
||||
self.integer = integer
|
||||
self.tensor = tensor
|
||||
self.func = func
|
||||
|
||||
def __getstate__(self):
|
||||
return self.hello_str, self.integer, self.tensor, self.func
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.hello_str, self.integer, self.tensor, self.func = state
|
||||
|
||||
def process_request(self, server_state):
|
||||
assert self.hello_str == STR
|
||||
assert self.integer == INTEGER
|
||||
new_tensor = self.func(self.tensor)
|
||||
res = HelloResponse(self.hello_str, self.integer, new_tensor)
|
||||
return res
|
||||
|
||||
|
||||
TIMEOUT_SERVICE_ID = 123456789
|
||||
TIMEOUT_META = "timeout_test"
|
||||
|
||||
|
||||
class TimeoutResponse(dgl.distributed.Response):
|
||||
def __init__(self, meta):
|
||||
self.meta = meta
|
||||
|
||||
def __getstate__(self):
|
||||
return self.meta
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.meta = state
|
||||
|
||||
|
||||
class TimeoutRequest(dgl.distributed.Request):
|
||||
def __init__(self, meta, timeout, response=True):
|
||||
self.meta = meta
|
||||
self.timeout = timeout
|
||||
self.response = response
|
||||
|
||||
def __getstate__(self):
|
||||
return self.meta, self.timeout, self.response
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.meta, self.timeout, self.response = state
|
||||
|
||||
def process_request(self, server_state):
|
||||
assert self.meta == TIMEOUT_META
|
||||
# convert from milliseconds to seconds
|
||||
time.sleep(self.timeout / 1000)
|
||||
if not self.response:
|
||||
return None
|
||||
res = TimeoutResponse(self.meta)
|
||||
return res
|
||||
|
||||
|
||||
def start_server(
|
||||
num_clients,
|
||||
ip_config,
|
||||
server_id=0,
|
||||
num_servers=1,
|
||||
):
|
||||
print("Sleep 1 seconds to test client re-connect.")
|
||||
time.sleep(1)
|
||||
server_state = dgl.distributed.ServerState(
|
||||
None, local_g=None, partition_book=None
|
||||
)
|
||||
dgl.distributed.register_service(
|
||||
HELLO_SERVICE_ID, HelloRequest, HelloResponse
|
||||
)
|
||||
dgl.distributed.register_service(
|
||||
TIMEOUT_SERVICE_ID, TimeoutRequest, TimeoutResponse
|
||||
)
|
||||
print("Start server {}".format(server_id))
|
||||
dgl.distributed.start_server(
|
||||
server_id=server_id,
|
||||
ip_config=ip_config,
|
||||
num_servers=num_servers,
|
||||
num_clients=num_clients,
|
||||
server_state=server_state,
|
||||
)
|
||||
|
||||
|
||||
def start_client(ip_config, group_id=0, num_servers=1):
|
||||
dgl.distributed.register_service(
|
||||
HELLO_SERVICE_ID, HelloRequest, HelloResponse
|
||||
)
|
||||
dgl.distributed.connect_to_server(
|
||||
ip_config=ip_config,
|
||||
num_servers=num_servers,
|
||||
group_id=group_id,
|
||||
)
|
||||
req = HelloRequest(STR, INTEGER, TENSOR, simple_func)
|
||||
# test send and recv
|
||||
dgl.distributed.send_request(0, req)
|
||||
res = dgl.distributed.recv_response()
|
||||
assert res.hello_str == STR
|
||||
assert res.integer == INTEGER
|
||||
assert_array_equal(F.asnumpy(res.tensor), F.asnumpy(TENSOR))
|
||||
# test remote_call
|
||||
target_and_requests = []
|
||||
for i in range(10):
|
||||
target_and_requests.append((0, req))
|
||||
res_list = dgl.distributed.remote_call(target_and_requests)
|
||||
for res in res_list:
|
||||
assert res.hello_str == STR
|
||||
assert res.integer == INTEGER
|
||||
assert_array_equal(F.asnumpy(res.tensor), F.asnumpy(TENSOR))
|
||||
# test send_request_to_machine
|
||||
dgl.distributed.send_request_to_machine(0, req)
|
||||
res = dgl.distributed.recv_response()
|
||||
assert res.hello_str == STR
|
||||
assert res.integer == INTEGER
|
||||
assert_array_equal(F.asnumpy(res.tensor), F.asnumpy(TENSOR))
|
||||
# test remote_call_to_machine
|
||||
target_and_requests = []
|
||||
for i in range(10):
|
||||
target_and_requests.append((0, req))
|
||||
res_list = dgl.distributed.remote_call_to_machine(target_and_requests)
|
||||
for res in res_list:
|
||||
assert res.hello_str == STR
|
||||
assert res.integer == INTEGER
|
||||
assert_array_equal(F.asnumpy(res.tensor), F.asnumpy(TENSOR))
|
||||
|
||||
|
||||
def start_client_timeout(ip_config, group_id=0, num_servers=1):
|
||||
dgl.distributed.register_service(
|
||||
TIMEOUT_SERVICE_ID, TimeoutRequest, TimeoutResponse
|
||||
)
|
||||
dgl.distributed.connect_to_server(
|
||||
ip_config=ip_config,
|
||||
num_servers=num_servers,
|
||||
group_id=group_id,
|
||||
)
|
||||
timeout = 1 * 1000 # milliseconds
|
||||
req = TimeoutRequest(TIMEOUT_META, timeout)
|
||||
# test send and recv
|
||||
dgl.distributed.send_request(0, req)
|
||||
res = dgl.distributed.recv_response(timeout=int(timeout / 2))
|
||||
assert res is None
|
||||
res = dgl.distributed.recv_response()
|
||||
assert res.meta == TIMEOUT_META
|
||||
# test remote_call
|
||||
req = TimeoutRequest(TIMEOUT_META, timeout, response=False)
|
||||
target_and_requests = []
|
||||
for i in range(3):
|
||||
target_and_requests.append((0, req))
|
||||
expect_except = False
|
||||
try:
|
||||
res_list = dgl.distributed.remote_call(
|
||||
target_and_requests, timeout=int(timeout / 2)
|
||||
)
|
||||
except dgl.DGLError:
|
||||
expect_except = True
|
||||
assert expect_except
|
||||
# test send_request_to_machine
|
||||
req = TimeoutRequest(TIMEOUT_META, timeout)
|
||||
dgl.distributed.send_request_to_machine(0, req)
|
||||
res = dgl.distributed.recv_response(timeout=int(timeout / 2))
|
||||
assert res is None
|
||||
res = dgl.distributed.recv_response()
|
||||
assert res.meta == TIMEOUT_META
|
||||
# test remote_call_to_machine
|
||||
req = TimeoutRequest(TIMEOUT_META, timeout, response=False)
|
||||
target_and_requests = []
|
||||
for i in range(3):
|
||||
target_and_requests.append((0, req))
|
||||
expect_except = False
|
||||
try:
|
||||
res_list = dgl.distributed.remote_call_to_machine(
|
||||
target_and_requests, timeout=int(timeout / 2)
|
||||
)
|
||||
except dgl.DGLError:
|
||||
expect_except = True
|
||||
assert expect_except
|
||||
|
||||
|
||||
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
|
||||
def test_rpc_timeout():
|
||||
reset_envs()
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
ip_config = "rpc_ip_config.txt"
|
||||
generate_ip_config(ip_config, 1, 1)
|
||||
ctx = mp.get_context("spawn")
|
||||
pserver = ctx.Process(target=start_server, args=(1, ip_config, 0, 1))
|
||||
pclient = ctx.Process(target=start_client_timeout, args=(ip_config, 0, 1))
|
||||
pserver.start()
|
||||
pclient.start()
|
||||
pserver.join()
|
||||
pclient.join()
|
||||
|
||||
|
||||
def test_serialize():
|
||||
reset_envs()
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
from dgl.distributed.rpc import (
|
||||
deserialize_from_payload,
|
||||
serialize_to_payload,
|
||||
)
|
||||
|
||||
SERVICE_ID = 12345
|
||||
dgl.distributed.register_service(SERVICE_ID, MyRequest, MyResponse)
|
||||
req = MyRequest()
|
||||
data, tensors = serialize_to_payload(req)
|
||||
req1 = deserialize_from_payload(MyRequest, data, tensors)
|
||||
req1.foo(req1.x, req1.y)
|
||||
assert req.x == req1.x
|
||||
assert req.y == req1.y
|
||||
assert F.array_equal(req.z, req1.z)
|
||||
|
||||
res = MyResponse()
|
||||
data, tensors = serialize_to_payload(res)
|
||||
res1 = deserialize_from_payload(MyResponse, data, tensors)
|
||||
assert res.x == res1.x
|
||||
|
||||
|
||||
def test_rpc_msg():
|
||||
reset_envs()
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
from dgl.distributed.rpc import (
|
||||
deserialize_from_payload,
|
||||
RPCMessage,
|
||||
serialize_to_payload,
|
||||
)
|
||||
|
||||
SERVICE_ID = 32452
|
||||
dgl.distributed.register_service(SERVICE_ID, MyRequest, MyResponse)
|
||||
req = MyRequest()
|
||||
data, tensors = serialize_to_payload(req)
|
||||
rpcmsg = RPCMessage(SERVICE_ID, 23, 0, 1, data, tensors)
|
||||
assert rpcmsg.service_id == SERVICE_ID
|
||||
assert rpcmsg.msg_seq == 23
|
||||
assert rpcmsg.client_id == 0
|
||||
assert rpcmsg.server_id == 1
|
||||
assert len(rpcmsg.data) == len(data)
|
||||
assert len(rpcmsg.tensors) == 1
|
||||
assert F.array_equal(rpcmsg.tensors[0], req.z)
|
||||
|
||||
|
||||
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
|
||||
def test_multi_client():
|
||||
reset_envs()
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
ip_config = "rpc_ip_config_mul_client.txt"
|
||||
generate_ip_config(ip_config, 1, 1)
|
||||
ctx = mp.get_context("spawn")
|
||||
num_clients = 20
|
||||
pserver = ctx.Process(
|
||||
target=start_server,
|
||||
args=(num_clients, ip_config, 0, 1),
|
||||
)
|
||||
pclient_list = []
|
||||
for i in range(num_clients):
|
||||
pclient = ctx.Process(target=start_client, args=(ip_config, 0, 1))
|
||||
pclient_list.append(pclient)
|
||||
pserver.start()
|
||||
for i in range(num_clients):
|
||||
pclient_list[i].start()
|
||||
for i in range(num_clients):
|
||||
pclient_list[i].join()
|
||||
pserver.join()
|
||||
|
||||
|
||||
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
|
||||
def test_multi_thread_rpc():
|
||||
reset_envs()
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
num_servers = 2
|
||||
ip_config = "rpc_ip_config_multithread.txt"
|
||||
generate_ip_config(ip_config, num_servers, num_servers)
|
||||
ctx = mp.get_context("spawn")
|
||||
pserver_list = []
|
||||
for i in range(num_servers):
|
||||
pserver = ctx.Process(target=start_server, args=(1, ip_config, i, 1))
|
||||
pserver.start()
|
||||
pserver_list.append(pserver)
|
||||
|
||||
def start_client_multithread(ip_config):
|
||||
import threading
|
||||
|
||||
dgl.distributed.connect_to_server(
|
||||
ip_config=ip_config,
|
||||
num_servers=1,
|
||||
)
|
||||
dgl.distributed.register_service(
|
||||
HELLO_SERVICE_ID, HelloRequest, HelloResponse
|
||||
)
|
||||
|
||||
req = HelloRequest(STR, INTEGER, TENSOR, simple_func)
|
||||
dgl.distributed.send_request(0, req)
|
||||
|
||||
def subthread_call(server_id):
|
||||
req = HelloRequest(STR, INTEGER, TENSOR, simple_func)
|
||||
dgl.distributed.send_request(server_id, req)
|
||||
|
||||
subthread = threading.Thread(target=subthread_call, args=(1,))
|
||||
subthread.start()
|
||||
subthread.join()
|
||||
|
||||
res0 = dgl.distributed.recv_response()
|
||||
res1 = dgl.distributed.recv_response()
|
||||
# Order is not guaranteed
|
||||
assert_array_equal(F.asnumpy(res0.tensor), F.asnumpy(TENSOR))
|
||||
assert_array_equal(F.asnumpy(res1.tensor), F.asnumpy(TENSOR))
|
||||
dgl.distributed.exit_client()
|
||||
|
||||
start_client_multithread(ip_config)
|
||||
pserver.join()
|
||||
|
||||
|
||||
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
|
||||
def test_multi_client_connect():
|
||||
reset_envs()
|
||||
os.environ["DGL_DIST_MODE"] = "distributed"
|
||||
ip_config = "rpc_ip_config_mul_client.txt"
|
||||
generate_ip_config(ip_config, 1, 1)
|
||||
ctx = mp.get_context("spawn")
|
||||
num_clients = 1
|
||||
pserver = ctx.Process(
|
||||
target=start_server,
|
||||
args=(num_clients, ip_config, 0, 1),
|
||||
)
|
||||
|
||||
# small max try times
|
||||
os.environ["DGL_DIST_MAX_TRY_TIMES"] = "1"
|
||||
expect_except = False
|
||||
try:
|
||||
start_client(ip_config, 0, 1)
|
||||
except dgl.distributed.DistConnectError as err:
|
||||
print("Expected error: {}".format(err))
|
||||
expect_except = True
|
||||
assert expect_except
|
||||
|
||||
# large max try times
|
||||
os.environ["DGL_DIST_MAX_TRY_TIMES"] = "1024"
|
||||
pclient = ctx.Process(target=start_client, args=(ip_config, 0, 1))
|
||||
pclient.start()
|
||||
pserver.start()
|
||||
pclient.join()
|
||||
pserver.join()
|
||||
reset_envs()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_serialize()
|
||||
test_rpc_msg()
|
||||
test_multi_client("socket")
|
||||
test_multi_client("tesnsorpipe")
|
||||
test_multi_thread_rpc()
|
||||
test_multi_client_connect("socket")
|
||||
@@ -0,0 +1,62 @@
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
|
||||
import dgl
|
||||
|
||||
import numpy as np
|
||||
import scipy.sparse as spsp
|
||||
|
||||
|
||||
def generate_ip_config(file_name, num_machines, num_servers):
|
||||
"""Get local IP and available ports, writes to file."""
|
||||
# get available IP in localhost
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
try:
|
||||
# doesn't even have to be reachable
|
||||
sock.connect(("10.255.255.255", 1))
|
||||
ip = sock.getsockname()[0]
|
||||
except ValueError:
|
||||
ip = "127.0.0.1"
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
# scan available PORT
|
||||
ports = []
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
start = random.randint(10000, 30000)
|
||||
for port in range(start, 65535):
|
||||
try:
|
||||
sock.connect((ip, port))
|
||||
ports = []
|
||||
except:
|
||||
ports.append(port)
|
||||
if len(ports) == num_machines * num_servers:
|
||||
break
|
||||
sock.close()
|
||||
if len(ports) < num_machines * num_servers:
|
||||
raise RuntimeError(
|
||||
"Failed to get available IP/PORT with required numbers."
|
||||
)
|
||||
with open(file_name, "w") as f:
|
||||
for i in range(num_machines):
|
||||
f.write("{} {}\n".format(ip, ports[i * num_servers]))
|
||||
|
||||
|
||||
def reset_envs():
|
||||
"""Reset common environment variable which are set in tests."""
|
||||
for key in [
|
||||
"DGL_ROLE",
|
||||
"DGL_NUM_SAMPLER",
|
||||
"DGL_NUM_SERVER",
|
||||
"DGL_DIST_MODE",
|
||||
"DGL_NUM_CLIENT",
|
||||
"DGL_DIST_MAX_TRY_TIMES",
|
||||
"DGL_DIST_DEBUG",
|
||||
]:
|
||||
if key in os.environ:
|
||||
os.environ.pop(key)
|
||||
|
||||
|
||||
def create_random_graph(n):
|
||||
return dgl.rand_graph(n, int(n * n * 0.001))
|
||||
Reference in New Issue
Block a user