server.py 1.14 KB
Newer Older
1
# This is a simple MXNet server demo shows how to use DGL distributed kvstore.
Chao Ma's avatar
Chao Ma committed
2
3
import dgl
import argparse
Chao Ma's avatar
Chao Ma committed
4
import torch as th
Chao Ma's avatar
Chao Ma committed
5

6
7
ndata_g2l = []
edata_g2l = []
Chao Ma's avatar
Chao Ma committed
8

9
10
11
12
13
14
15
16
17
ndata_g2l.append({'ndata':th.tensor([0,1,0,0,0,0,0,0])})
ndata_g2l.append({'ndata':th.tensor([0,0,0,1,0,0,0,0])})
ndata_g2l.append({'ndata':th.tensor([0,0,0,0,0,1,0,0])})
ndata_g2l.append({'ndata':th.tensor([0,0,0,0,0,0,0,1])})

edata_g2l.append({'edata':th.tensor([0,1,0,0,0,0,0,0])})
edata_g2l.append({'edata':th.tensor([0,0,0,1,0,0,0,0])})
edata_g2l.append({'edata':th.tensor([0,0,0,0,0,1,0,0])})
edata_g2l.append({'edata':th.tensor([0,0,0,0,0,0,0,1])})
Chao Ma's avatar
Chao Ma committed
18

19
20
21
22
23
24
25
26
27
28
def start_server(args):
    
    dgl.contrib.start_server(
        server_id=args.id,
        ip_config='ip_config.txt',
        num_client=4,
        ndata={'ndata':th.tensor([[0.,0.,0.],[0.,0.,0.]])},
        edata={'edata':th.tensor([[0.,0.,0.],[0.,0.,0.]])},
        ndata_g2l=ndata_g2l[args.id],
        edata_g2l=edata_g2l[args.id])
Chao Ma's avatar
Chao Ma committed
29

Chao Ma's avatar
Chao Ma committed
30
31
32
33
34
35

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='kvstore')
    parser.add_argument("--id", type=int, default=0, help="node ID")
    args = parser.parse_args()

36
    start_server(args)