dmlc--dgl
Chao Ma
5fc334fc30
* fix get-shape * update * update * update * update * update * fix typo * update * update
136 行
5.0 KiB
Python
136 行
5.0 KiB
Python
import backend as F
|
|
import numpy as np
|
|
import scipy as sp
|
|
import dgl
|
|
from dgl import utils
|
|
from dgl.contrib import KVServer
|
|
from dgl.contrib import KVClient
|
|
from numpy.testing import assert_array_equal
|
|
|
|
import os
|
|
import time
|
|
|
|
num_entries = 10
|
|
dim_size = 3
|
|
|
|
server_namebook = {0:[0, '127.0.0.1', 30070, 1]}
|
|
|
|
data_0 = F.zeros((num_entries, dim_size), F.float32, F.cpu())
|
|
g2l_0 = F.arange(0, num_entries)
|
|
partition_0 = F.zeros(num_entries, F.int64, F.cpu())
|
|
|
|
data_1 = F.zeros((num_entries*2, dim_size), F.float32, F.cpu())
|
|
g2l_1 = F.arange(0, num_entries*2)
|
|
partition_1 = F.zeros(num_entries*2, F.int64, F.cpu())
|
|
|
|
data_3 = F.zeros((num_entries, dim_size), F.int64, F.cpu())
|
|
data_4 = F.zeros((num_entries, dim_size), F.float64, F.cpu())
|
|
data_5 = F.zeros((num_entries, dim_size), F.int32, F.cpu())
|
|
|
|
def start_server():
|
|
my_server = KVServer(server_id=0, server_namebook=server_namebook, num_client=1)
|
|
my_server.set_global2local(name='data_0', global2local=g2l_0)
|
|
my_server.set_global2local(name='data_1', global2local=g2l_1)
|
|
my_server.set_global2local(name='data_3', global2local=g2l_0)
|
|
my_server.set_global2local(name='data_4', global2local=g2l_0)
|
|
my_server.set_global2local(name='data_5', global2local=g2l_0)
|
|
my_server.set_partition_book(name='data_0', partition_book=partition_0)
|
|
my_server.set_partition_book(name='data_1', partition_book=partition_1)
|
|
my_server.set_partition_book(name='data_3', partition_book=partition_0)
|
|
my_server.set_partition_book(name='data_4', partition_book=partition_0)
|
|
my_server.set_partition_book(name='data_5', partition_book=partition_0)
|
|
my_server.init_data(name='data_0', data_tensor=data_0)
|
|
my_server.init_data(name='data_1', data_tensor=data_1)
|
|
my_server.init_data(name='data_3', data_tensor=data_3)
|
|
my_server.init_data(name='data_4', data_tensor=data_4)
|
|
my_server.init_data(name='data_5', data_tensor=data_5)
|
|
|
|
my_server.start()
|
|
|
|
|
|
def start_client():
|
|
my_client = KVClient(server_namebook=server_namebook)
|
|
my_client.connect()
|
|
|
|
my_client.init_data(name='data_2', shape=(num_entries, dim_size), dtype=F.float32, target_name='data_0')
|
|
print("Init data from client..")
|
|
|
|
name_list = my_client.get_data_name_list()
|
|
assert len(name_list) == 6
|
|
assert 'data_0' in name_list
|
|
assert 'data_1' in name_list
|
|
assert 'data_2' in name_list
|
|
assert 'data_3' in name_list
|
|
assert 'data_4' in name_list
|
|
assert 'data_5' in name_list
|
|
|
|
meta_0 = my_client.get_data_meta('data_0')
|
|
assert meta_0[0] == F.float32
|
|
assert meta_0[1] == tuple(F.shape(data_0))
|
|
assert_array_equal(meta_0[2], partition_0)
|
|
|
|
meta_1 = my_client.get_data_meta('data_1')
|
|
assert meta_1[0] == F.float32
|
|
assert meta_1[1] == tuple(F.shape(data_1))
|
|
assert_array_equal(meta_1[2], partition_1)
|
|
|
|
meta_2 = my_client.get_data_meta('data_2')
|
|
assert meta_2[0] == F.float32
|
|
assert meta_2[1] == tuple(F.shape(data_0))
|
|
assert_array_equal(meta_2[2], partition_0)
|
|
|
|
meta_3 = my_client.get_data_meta('data_3')
|
|
assert meta_3[0] == F.int64
|
|
assert meta_3[1] == tuple(F.shape(data_3))
|
|
assert_array_equal(meta_3[2], partition_0)
|
|
|
|
meta_4 = my_client.get_data_meta('data_4')
|
|
assert meta_4[0] == F.float64
|
|
assert meta_4[1] == tuple(F.shape(data_4))
|
|
assert_array_equal(meta_3[2], partition_0)
|
|
|
|
meta_5 = my_client.get_data_meta('data_5')
|
|
assert meta_5[0] == F.int32
|
|
assert meta_5[1] == tuple(F.shape(data_5))
|
|
assert_array_equal(meta_3[2], partition_0)
|
|
|
|
my_client.push(name='data_0', id_tensor=F.tensor([0, 1, 2]), data_tensor=F.tensor([[1.,1.,1.],[2.,2.,2.],[3.,3.,3.]]))
|
|
my_client.push(name='data_2', id_tensor=F.tensor([0, 1, 2]), data_tensor=F.tensor([[1.,1.,1.],[2.,2.,2.],[3.,3.,3.]]))
|
|
my_client.push(name='data_3', id_tensor=F.tensor([0, 1, 2]), data_tensor=F.tensor([[1,1,1],[2,2,2],[3,3,3]]))
|
|
my_client.push(name='data_4', id_tensor=F.tensor([0, 1, 2]), data_tensor=F.tensor([[1.,1.,1.],[2.,2.,2.],[3.,3.,3.]], F.float64))
|
|
my_client.push(name='data_5', id_tensor=F.tensor([0, 1, 2]), data_tensor=F.tensor([[1,1,1],[2,2,2],[3,3,3]], F.int32))
|
|
|
|
target = F.tensor([[1.,1.,1.],[2.,2.,2.],[3.,3.,3.]])
|
|
|
|
res = my_client.pull(name='data_0', id_tensor=F.tensor([0, 1, 2]))
|
|
assert_array_equal(res, target)
|
|
|
|
res = my_client.pull(name='data_2', id_tensor=F.tensor([0, 1, 2]))
|
|
assert_array_equal(res, target)
|
|
|
|
target = F.tensor([[1,1,1],[2,2,2],[3,3,3]])
|
|
|
|
res = my_client.pull(name='data_3', id_tensor=F.tensor([0, 1, 2]))
|
|
assert_array_equal(res, target)
|
|
|
|
target = F.tensor([[1.,1.,1.],[2.,2.,2.],[3.,3.,3.]], F.float64)
|
|
|
|
res = my_client.pull(name='data_4', id_tensor=F.tensor([0, 1, 2]))
|
|
assert_array_equal(res, target)
|
|
|
|
target = F.tensor([[1,1,1],[2,2,2],[3,3,3]], F.int32)
|
|
|
|
res = my_client.pull(name='data_5', id_tensor=F.tensor([0, 1, 2]))
|
|
assert_array_equal(res, target)
|
|
|
|
my_client.shut_down()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pid = os.fork()
|
|
if pid == 0:
|
|
start_server()
|
|
else:
|
|
time.sleep(2) # wait trainer start
|
|
start_client()
|