Repository: awslabs/dgl-ke (language: python)

Issue:
update url
*Issue #, if available:*

*Description of changes:*


By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.



Discussion / hints:

Fix patch:
diff --git a/examples/distributed/Distributed Demo.md b/examples/distributed/Distributed Demo.md
new file mode 100644
index 0000000..a0993e3
--- /dev/null
+++ b/examples/distributed/Distributed Demo.md
@@ -0,0 +1,33 @@
+##Distributed Demo
+
+First setup distributed enviroment as [https://github.com/dmlc/dgl/tree/master/examples/pytorch/graphsage/experimental#distributed-training](https://github.com/dmlc/dgl/tree/master/examples/pytorch/graphsage/experimental#distributed-training)
+
+To part the graph into parts needed:
+
+```bash
+python3 ~/workspace/dgl-ke/python/partion_graph.py --nump_parts NUM_PARTS
+```
+
+Run the demo:
+
+```bash
+python3 ~/workspace/dgl-ke/python/dglke/launch.py \
+		--workspace ~/workspace/dgl-ke/examples/distributed \
+		--num_trainers 1 \
+		--num_samplers 0 \
+		--num_servers 1 \
+		--part_config data/FB15k.json \
+		--ip_config ip_config.txt "python3 dist_test.py"
+```
+
+To run in standalone mode, first construct a graph with one partition:
+
+```bash
+python3 ~/workspace/dgl-ke/python/partion_graph.py --nump_parts 1
+```
+
+To run the demo:
+
+```bash
+python3 ~/workspace/dgl-ke/python/dist_test.py --standalone
+```
\ No newline at end of file
diff --git a/python/dglke/dist_main.py b/python/dglke/dist_main.py
new file mode 100644
index 0000000..9fe6222
--- /dev/null
+++ b/python/dglke/dist_main.py
@@ -0,0 +1,176 @@
+import os
+os.environ['DGLBACKEND']='pytorch'
+from multiprocessing import Process
+import argparse, time, math
+import numpy as np
+from functools import wraps
+import tqdm
+import random
+
+import dgl
+from dgl import DGLGraph
+from dgl.data import register_data_args, load_data
+from dgl.data.utils import load_graphs
+import dgl.function as fn
+import dgl.nn.pytorch as dglnn
+from dgl.distributed import DistDataLoader
+
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.multiprocessing as mp
+from torch.utils.data import DataLoader
+
+from .utils.argparser import TrainArgParser
+from .utils.misc import prepare_args
+from .utils.logging import Logger
+from .data import get_dataset, TrainDataset, TestDataset, ValidDataset
+from .data.dataloader import KGETrainDataLoaderGenerator, KGEEvalDataLoaderGenerator
+from .utils import EMB_INIT_EPS
+from .nn.modules import KGEEncoder, TransREncoder
+from .nn.modules import KGEDecoder, AttHDecoder, TransRDecoder
+from .nn.loss import sLCWAKGELossGenerator
+from .nn.loss import BCELoss, HingeLoss, LogisticLoss, LogsigmoidLoss
+from .regularizer import Regularizer
+from .nn.modules import TransEScore, TransRScore, DistMultScore, ComplExScore, RESCALScore, RotatEScore, SimplEScore
+from .nn.metrics import RankingMetricsEvaluator
+from functools import partial
+import torch as th
+from torch import nn
+import time
+from .nn.modules import KEModel
+
+def dist_main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--local_rank', type=int, help='get rank of the process')
+    parser.add_argument('--standalone', action='store_true', help='run in the standalone mode')
+    args = parser.parse_args()
+
+    dgl.distributed.initialize('ip_config.txt')
+    g = dgl.distributed.DistGraph('FB15k', part_config='data/FB15k.json')
+
+    if not args.standalone:
+        th.distributed.init_process_group(backend='gloo')
+
+    pb = g.get_partition_book()
+
+    train_mask = th.cat((th.ones((483142,), dtype=th.bool), th.zeros((59071,), dtype=th.bool)))
+    test_mask = th.cat((th.zeros((483142,), dtype=th.bool), th.ones((59071,), dtype=th.bool)))
+
+    train_eids = dgl.distributed.edge_split(train_mask, pb, force_even=True)
+    test_eids = dgl.distributed.edge_split(test_mask, pb, force_even=True)
+    nids = pb.partid2nids(pb.partid)
+
+    emb_init = (143.0 + 1.0) / 200
+    init_func = [partial(th.nn.init.uniform_, a=-emb_init, b=emb_init), partial(th.nn.init.uniform_, a=-emb_init, b=emb_init)]
+    encoder = KGEEncoder(hidden_dim=200,
+                             n_entity=g.num_nodes(),
+                             n_relation=1345, #need to get numbers of relations here
+                             init_func=init_func,
+                             score_func='TransE')
+
+    
+    emb_init = (143.0 + 2.0) / 200
+    score_func = TransEScore(143.0, dist_func='l2')
+    loss_gen = sLCWAKGELossGenerator(neg_adversarial_sampling=True,
+                                         adversarial_temperature=1.0,
+                                         pairwise=False,
+                                         label_smooth=.0)
+    criterion = LogsigmoidLoss()
+
+    loss_gen.set_criterion(criterion)
+    metrics_evaluator = RankingMetricsEvaluator(False)
+    decoder = KGEDecoder('TransE',
+                         score_func,
+                         loss_gen,
+                         metrics_evaluator)
+    
+    if not args.standalone:
+        encoder = th.nn.parallel.DistributedDataParallel(encoder)
+    
+
+    def sample(eids):
+        s, d = g.find_edges(eids)
+        s = th.tensor(s)
+        d = th.tensor(d)
+        rel = g.edata['tid'][eids]
+        rel = th.tensor(rel)
+        neg = random.sample(list(nids.numpy()), 128)
+        neg = th.tensor(neg)
+        data = {'head' : s, 
+                'tail' : d, 
+                'rel' : rel, 
+                'neg' : neg}
+        return data
+
+    dataloader = DistDataLoader(
+        dataset=train_eids.numpy(),
+        batch_size=512,
+        collate_fn=sample,
+        shuffle=True,
+        drop_last=False)
+
+    emb_optimizer = th.optim.SparseAdam(list(encoder.parameters()), lr=1e-1)
+    print('optimize Pytorch sparse embedding:', encoder)
+
+    regularizer = Regularizer(coef=2e-9, norm=3)
+
+
+    for epoch in range(10):
+        for step, data in enumerate(dataloader):
+            encode_results = encoder.forward(data, -1)
+            data['neg_type'] = 'head'
+            data['chunk_size'] = data['head'].shape[0]
+            data['neg_sample_size'] = 128
+            decode_results = decoder.forward(encode_results, data, -1)
+            loss = {}
+            loss['decode'] = decoder.get_loss(decode_results)
+            #loss['reg'] = regularizer.compute_regularization(encode_results)
+
+            total_loss = 0
+            for k, v in loss.items():
+                total_loss += v
+
+            emb_optimizer.zero_grad()
+            total_loss.backward()
+            emb_optimizer.step()
+            if step % 100 == 0:
+                print(loss.items())
+    
+    def test_sample(eids):
+        s, d = g.find_edges(eids)
+        s = th.tensor(s)
+        d = th.tensor(d)
+        rel = g.edata['tid'][eids]
+        rel = th.tensor(rel)
+        neg = random.sample(list(nids.numpy()), 128)
+        neg = th.tensor(neg)
+        data = {'head' : s, 
+                'tail' : d, 
+                'rel' : rel, 
+                'neg' : neg}
+        return data
+
+    test_dataloader = DistDataLoader(
+        dataset=test_eids.numpy(),
+        batch_size=16,
+        collate_fn=test_sample,
+        shuffle=False,
+        drop_last=False)
+    
+
+    with th.no_grad():
+        logs = []
+        for step, data in enumerate(test_dataloader):
+            data['chunk_size'] = data['head'].shape[0]
+            data['neg_sample_size'] = 128
+            data['neg_type'] = 'head_tail'
+            encode_results = encoder.forward(data, -1)
+            results = decoder.forward(encode_results, data, -1)
+            logs += decoder.evaluate(results, data, g)
+        metrics = {}
+        if len(logs) > 0:
+            for metric in logs[0].keys():
+                metrics[metric] = sum([log[metric] for log in logs]) / len(logs)
+        for k, v in metrics.items():
+            print('average {}: {}'.format(k, v))
+    
\ No newline at end of file
diff --git a/python/dglke/partition_graph.py b/python/dglke/partition_graph.py
new file mode 100644
index 0000000..8dfbe2d
--- /dev/null
+++ b/python/dglke/partition_graph.py
@@ -0,0 +1,220 @@
+from .data import get_dataset, TestDataset
+from torch import nn
+import time
+import dgl
+import numpy as np
+import torch as th
+import argparse
+import time
+import scipy as sp
+import dgl.backend as F
+
+class BaseDataset(object):
+    def __init__(self):
+        pass
+
+    def create_label(self, triples, n_entity, args, mode='train', type='tail'):
+        path = os.path.join(args.data_path, args.dataset, f'{mode}_{type}_label.pkl')
+        if not os.path.exists(path):
+            print(f"{mode}_{type}_label.pkl does not exist, create label from scratch...")
+            if type == 'tail':
+                hr_dict = {}
+                for idx in range(len(triples[0])):
+                    if (triples[0][idx], triples[1][idx]) not in hr_dict:
+                        hr_dict[(triples[0][idx], triples[1][idx])] = []
+                    hr_dict[(triples[0][idx], triples[1][idx])].append(triples[2][idx])
+                label_dict = {}
+                for k, v in tqdm.tqdm(hr_dict.items()):
+                    hr_dict[k] = np.hstack(v)
+                    label_dict[k] = np.zeros(n_entity)
+                    label_dict[k][hr_dict[k]] = 1
+            else:
+                rt_dict = {}
+                for idx in range(len(triples[1])):
+                    if(triples[1][idx], triples[2][idx]) not in rt_dict:
+                        rt_dict[(triples[1][idx], triples[2][idx])] = []
+                    rt_dict[(triples[1][idx], triples[2][idx])].append(triples[0][idx])
+                label_dict = {}
+                for k, v in tqdm.tqdm(rt_dict.items()):
+                    rt_dict[k] = np.stack(v)
+                    label_dict[k] = np.zeros(n_entity)
+                    label_dict[k][rt_dict[k]] = 1
+
+            print(f"save label dict...")
+            with open(path, 'wb') as f:
+                pickle.dump(label_dict, f)
+        else:
+            print(f"load label from cache...")
+            with open(path, 'rb') as f:
+                label_dict = pickle.load(f)
+        return label_dict
+
+    def partition(self, rank, world_size) -> th.utils.data.Dataset:
+        raise NotImplementedError
+
+
+class TestDataset(BaseDataset):
+    def __init__(self, dataset, args):
+        src = [dataset.train[0]]
+        etype_id = [dataset.train[1]]
+        dst = [dataset.train[2]]
+        self.num_train = len(dataset.train[0])
+        self.batch_size = args.batch_size_eval
+        self.num_valid = 0
+        if dataset.test is not None:
+            src.append(dataset.test[0])
+            etype_id.append(dataset.test[1])
+            dst.append(dataset.test[2])
+            self.num_test = len(dataset.test[0])
+        else:
+            self.num_test = 0
+        assert len(src) > 1, "we need to have at least validation set or test set."
+        src = np.concatenate(src)
+        etype_id = np.concatenate(etype_id)
+        dst = np.concatenate(dst)
+
+        coo = sp.sparse.coo_matrix((np.ones(len(src)), (src, dst)),
+                                   shape=[dataset.n_entities, dataset.n_entities])
+        g = dgl.DGLGraph(coo, readonly=True, multigraph=True, sort_csr=True)
+        g.edata['tid'] = F.tensor(etype_id, F.int64)
+        self.g = g
+        if args.eval_percent < 1:
+            self.test = th.randint(low=0, high=self.num_test, size=(int(self.num_test * args.eval_percent),))
+            self.test += self.num_train + self.num_valid
+        else:
+            self.test = th.arange(self.num_train + self.num_valid, self.g.number_of_edges())
+        print('|test|:', len(self.test))
+
+    def partition(self, rank, world_size):
+        edges = self.test
+        step_size = (edges.shape[0] + world_size - 1) // world_size
+        beg = step_size * rank
+        end = min(step_size * (rank + 1), edges.shape[0])
+        edges = edges[beg: end]
+        heads, tails = self.g.all_edges(order='eid')
+        heads = heads[edges]
+        tails = tails[edges]
+        rels = self.g.edata['tid'][edges]
+        return PartitionEvalDataset(heads, rels, tails)
+
+def create_dataset_graph(args):
+    g = None
+    dataset = get_dataset(data_path=args.data_path,
+                          data_name=args.dataset,
+                          format_str=args.format,
+                          delimiter=args.delimiter,
+                          files=args.data_files,
+                          has_edge_importance=args.has_edge_importance,
+                          inverse_rel=args.inverse_rel
+                          )
+    train_dataset, eval_dataset, test_dataset = None, None, None
+    # create training dataset needed parameters
+    train_dataset = TestDataset(dataset, args)
+    g = train_dataset.g
+    args.strict_rel_part = args.mix_cpu_gpu and (train_dataset.cross_part is False)
+    args.rel_parts = train_dataset.rel_parts if args.strict_rel_part else None
+    args.n_entities = dataset.n_entities
+    args.n_relations = dataset.n_relations
+
+    return [train_dataset, eval_dataset, test_dataset], g
+
+def partition_graph():
+    argparser = argparse.ArgumentParser("Partition builtin graphs")
+    argparser.add_argument('--num_parts', type=int, default=2,
+                           help='number of partitions')
+    argparser.add_argument('--part_method', type=str, default='metis',
+                           help='the partition method')
+    argparser.add_argument('--balance_train', action='store_true',
+                           help='balance the training size in each partition.')
+    argparser.add_argument('--undirected', action='store_true',
+                           help='turn the graph into an undirected graph.')
+    argparser.add_argument('--balance_edges', action='store_true',
+                           help='balance the number of edges in each partition.')
+    argparser.add_argument('--output', type=str, default='data',
+                           help='Output path of partitioned graph.')
+    argparser.add_argument('--data_path', type=str, default='data',
+                      help='The path of the directory where DGL-KE loads knowledge graph data.')
+    argparser.add_argument('--dataset', type=str, default='FB15k',
+                      help='The name of the builtin knowledge graph. Currently, the builtin knowledge ' \
+                           'graphs include FB15k, FB15k-237, wn18, wn18rr and Freebase. ' \
+                           'DGL-KE automatically downloads the knowledge graph and keep it under data_path.')
+    argparser.add_argument('--format', type=str, default='built_in',
+                      help='The format of the dataset. For builtin knowledge graphs,' \
+                           'the foramt should be built_in. For users own knowledge graphs,' \
+                           'it needs to be raw_udd_{htr} or udd_{htr}.')
+    argparser.add_argument('--data_files', type=str, default=None, nargs='+',
+                      help='A list of data file names. This is used if users want to train KGE' \
+                           'on their own datasets. If the format is raw_udd_{htr},' \
+                           'users need to provide train_file [valid_file] [test_file].' \
+                           'If the format is udd_{htr}, users need to provide' \
+                           'entity_file relation_file train_file [valid_file] [test_file].' \
+                           'In both cases, valid_file and test_file are optional.')
+    argparser.add_argument('--delimiter', type=str, default='\t',
+                      help='Delimiter used in data files. Note all files should use the same delimiter.')
+
+    argparser.add_argument('--mix_cpu_gpu', action='store_true',
+                      help='Training a knowledge graph embedding model with both CPUs and GPUs.' \
+                           'The embeddings are stored in CPU memory and the training is performed in GPUs.' \
+                           'This is usually used for training a large knowledge graph embeddings.')
+    argparser.add_argument('--valid', action='store_true',
+                      help='Evaluate the model on the validation set in the training.')
+    argparser.add_argument('--eval_interval', type=int, default=1,
+                      help='Print evaluation results on the validation dataset every x steps' \
+                           'if validation is turned on')
+    argparser.add_argument('--rel_part', action='store_true',
+                      help='Enable relation partitioning for multi-GPU training.')
+    argparser.add_argument('--has_edge_importance', action='store_true',
+                      help='Allow providing edge importance score for each edge during training.' \
+                           'The positive score will be adjusted ' \
+                           'as pos_score = pos_score * edge_importance')
+    argparser.add_argument('--inverse_rel', action='store_true', dest='inverse_rel',
+                      help='If specified, create a->inv_rel->b  based on b->rel->a.')
+    argparser.add_argument('--num_proc', type=int, default=1,
+                      help='The number of processes to train the model in parallel.' \
+                           'In multi-GPU training, the number of processes by default is set to match the number of GPUs.' \
+                           'If set explicitly, the number of processes needs to be divisible by the number of GPUs.')
+    argparser.add_argument('--num_thread', type=int, default=1,
+                      help='The number of CPU threads to train the model in each process.' \
+                           'This argument is used for multiprocessing training.')
+    argparser.add_argument('--batch_size', type=int, default=1024,
+                      help='The batch size for training.')
+    argparser.add_argument('--num_workers', type=int, default=0,
+                      help='Number of process to fetch data for training/validation dataset.')
+    argparser.add_argument('--max_step', type=int, default=80000,
+                      help='The maximal number of steps to train the model.' \
+                           'A step trains the model with a batch of data.')
+    argparser.add_argument('--batch_size_eval', type=int, default=8,
+                      help='The batch size used for validation and test.')
+    argparser.add_argument('--neg_sample_size', type=int, default=256,
+                      help='The number of negative samples we use for each positive sample in the training.')
+    argparser.add_argument('--neg_deg_sample', action='store_true',
+                      help='Construct negative samples proportional to vertex degree in the training.' \
+                           'When this option is turned on, the number of negative samples per positive edge' \
+                           'will be doubled. Half of the negative samples are generated uniformly while' \
+                           'the other half are generated proportional to vertex degree.')
+    argparser.add_argument('--neg_deg_sample_eval', action='store_true',
+                      help='Construct negative samples proportional to vertex degree in the evaluation.')
+    argparser.add_argument('--neg_sample_size_eval', type=int, default=-1,
+                      help='The number of negative samples we use to evaluate a positive sample.')
+    argparser.add_argument('--eval_percent', type=float, default=1,
+                      help='Randomly sample some percentage of edges for evaluation.')
+    argparser.add_argument('--no_eval_filter', action='store_false', dest='eval_filter',
+                      help='Disable filter positive edges from randomly constructed negative edges for evaluation')
+    argparser.add_argument('--self_loop_filter', action='store_true', dest='self_loop_filter',
+                      help='Disable filter triple like (head - relation - head) score for evaluation')
+    argparser.add_argument('--test', action='store_true',
+                      help='Evaluate the model on the test set after the model is trained.')
+    
+    
+    args = argparser.parse_args()
+
+    start = time.time()
+
+    dataset, g = create_dataset_graph(args)
+
+    print('load {} takes {:.3f} seconds'.format(args.dataset, time.time() - start))
+    print('|V|={}, |E|={}'.format(g.number_of_nodes(), g.number_of_edges()))
+    dgl.distributed.partition_graph(g, args.dataset, args.num_parts, args.output,
+                                    part_method=args.part_method,
+                                    balance_ntypes=None,
+                                    balance_edges=args.balance_edges)
\ No newline at end of file
diff --git a/python/launch.py b/python/launch.py
new file mode 100644
index 0000000..b682dd5
--- /dev/null
+++ b/python/launch.py
@@ -0,0 +1,286 @@
+"""Launching tool for DGL distributed training"""
+import os
+import stat
+import sys
+import subprocess
+import argparse
+import signal
+import logging
+import time
+import json
+import multiprocessing
+import re
+from functools import partial
+from threading import Thread
+
+DEFAULT_PORT = 30050
+
+def cleanup_proc(get_all_remote_pids, conn):
+    '''This process tries to clean up the remote training tasks.
+    '''
+    print('cleanupu process runs')
+    # This process should not handle SIGINT.
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    data = conn.recv()
+    # If the launch process exits normally, this process doesn't need to do anything.
+    if data == 'exit':
+        sys.exit(0)
+    else:
+        remote_pids = get_all_remote_pids()
+        # Otherwise, we need to ssh to each machine and kill the training jobs.
+        for (ip, port), pids in remote_pids.items():
+            kill_process(ip, port, pids)
+    print('cleanup process exits')
+
+def kill_process(ip, port, pids):
+    '''ssh to a remote machine and kill the specified processes.
+    '''
+    curr_pid = os.getpid()
+    killed_pids = []
+    # If we kill child processes first, the parent process may create more again. This happens
+    # to Python's process pool. After sorting, we always kill parent processes first.
+    pids.sort()
+    for pid in pids:
+        assert curr_pid != pid
+        print('kill process {} on {}:{}'.format(pid, ip, port), flush=True)
+        kill_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'kill {}\''.format(pid)
+        subprocess.run(kill_cmd, shell=True)
+        killed_pids.append(pid)
+    # It's possible that some of the processes are not killed. Let's try again.
+    for i in range(3):
+        killed_pids = get_killed_pids(ip, port, killed_pids)
+        if len(killed_pids) == 0:
+            break
+        else:
+            killed_pids.sort()
+            for pid in killed_pids:
+                print('kill process {} on {}:{}'.format(pid, ip, port), flush=True)
+                kill_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'kill -9 {}\''.format(pid)
+                subprocess.run(kill_cmd, shell=True)
+
+def get_killed_pids(ip, port, killed_pids):
+    '''Get the process IDs that we want to kill but are still alive.
+    '''
+    killed_pids = [str(pid) for pid in killed_pids]
+    killed_pids = ','.join(killed_pids)
+    ps_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'ps -p {} -h\''.format(killed_pids)
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    pids = []
+    for p in res.stdout.decode('utf-8').split('\n'):
+        l = p.split()
+        if len(l) > 0:
+            pids.append(int(l[0]))
+    return pids
+
+def execute_remote(cmd, ip, port, thread_list):
+    """execute command line on remote machine via ssh"""
+    cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'' + cmd + '\''
+    # thread func to run the job
+    def run(cmd):
+        subprocess.check_call(cmd, shell = True)
+
+    thread = Thread(target = run, args=(cmd,))
+    thread.setDaemon(True)
+    thread.start()
+    thread_list.append(thread)
+
+def get_remote_pids(ip, port, cmd_regex):
+    """Get the process IDs that run the command in the remote machine.
+    """
+    pids = []
+    curr_pid = os.getpid()
+    # Here we want to get the python processes. We may get some ssh processes, so we should filter them out.
+    ps_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'ps -aux | grep python | grep -v StrictHostKeyChecking\''
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    for p in res.stdout.decode('utf-8').split('\n'):
+        l = p.split()
+        if len(l) < 2:
+            continue
+        # We only get the processes that run the specified command.
+        res = re.search(cmd_regex, p)
+        if res is not None and int(l[1]) != curr_pid:
+            pids.append(l[1])
+
+    pid_str = ','.join([str(pid) for pid in pids])
+    ps_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'pgrep -P {}\''.format(pid_str)
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    pids1 = res.stdout.decode('utf-8').split('\n')
+    all_pids = []
+    for pid in set(pids + pids1):
+        if pid == '' or int(pid) == curr_pid:
+            continue
+        all_pids.append(int(pid))
+    all_pids.sort()
+    return all_pids
+
+def get_all_remote_pids(hosts, ssh_port, udf_command):
+    '''Get all remote processes.
+    '''
+    remote_pids = {}
+    for node_id, host in enumerate(hosts):
+        ip, _ = host
+        # When creating training processes in remote machines, we may insert some arguments
+        # in the commands. We need to use regular expressions to match the modified command.
+        cmds = udf_command.split()
+        new_udf_command = ' .*'.join(cmds)
+        pids = get_remote_pids(ip, ssh_port, new_udf_command)
+        remote_pids[(ip, ssh_port)] = pids
+    return remote_pids
+
+def submit_jobs(args, udf_command):
+    """Submit distributed jobs (server and client processes) via ssh"""
+    hosts = []
+    thread_list = []
+    server_count_per_machine = 0
+
+    # Get the IP addresses of the cluster.
+    ip_config = args.workspace + '/' + args.ip_config
+    with open(ip_config) as f:
+        for line in f:
+            result = line.strip().split()
+            if len(result) == 2:
+                ip = result[0]
+                port = int(result[1])
+                hosts.append((ip, port))
+            elif len(result) == 1:
+                ip = result[0]
+                port = DEFAULT_PORT
+                hosts.append((ip, port))
+            else:
+                raise RuntimeError("Format error of ip_config.")
+            server_count_per_machine = args.num_servers
+    # Get partition info of the graph data
+    part_config = args.workspace + '/' + args.part_config
+    with open(part_config) as conf_f:
+        part_metadata = json.load(conf_f)
+    assert 'num_parts' in part_metadata, 'num_parts does not exist.'
+    # The number of partitions must match the number of machines in the cluster.
+    assert part_metadata['num_parts'] == len(hosts), \
+            'The number of graph partitions has to match the number of machines in the cluster.'
+
+    tot_num_clients = args.num_trainers * (1 + args.num_samplers) * len(hosts)
+    # launch server tasks
+    server_cmd = 'DGL_ROLE=server DGL_NUM_SAMPLER=' + str(args.num_samplers)
+    server_cmd = server_cmd + ' ' + 'OMP_NUM_THREADS=' + str(args.num_server_threads)
+    server_cmd = server_cmd + ' ' + 'DGL_NUM_CLIENT=' + str(tot_num_clients)
+    server_cmd = server_cmd + ' ' + 'DGL_CONF_PATH=' + str(args.part_config)
+    server_cmd = server_cmd + ' ' + 'DGL_IP_CONFIG=' + str(args.ip_config)
+    server_cmd = server_cmd + ' ' + 'DGL_NUM_SERVER=' + str(args.num_servers)
+    server_cmd = server_cmd + ' ' + 'DGL_GRAPH_FORMAT=' + str(args.graph_format)
+    for i in range(len(hosts)*server_count_per_machine):
+        ip, _ = hosts[int(i / server_count_per_machine)]
+        cmd = server_cmd + ' ' + 'DGL_SERVER_ID=' + str(i)
+        cmd = cmd + ' ' + udf_command
+        cmd = 'cd ' + str(args.workspace) + '; ' + cmd
+        execute_remote(cmd, ip, args.ssh_port, thread_list)
+
+    # launch client tasks
+    client_cmd = 'DGL_DIST_MODE="distributed" DGL_ROLE=client DGL_NUM_SAMPLER=' + str(args.num_samplers)
+    client_cmd = client_cmd + ' ' + 'DGL_NUM_CLIENT=' + str(tot_num_clients)
+    client_cmd = client_cmd + ' ' + 'DGL_CONF_PATH=' + str(args.part_config)
+    client_cmd = client_cmd + ' ' + 'DGL_IP_CONFIG=' + str(args.ip_config)
+    client_cmd = client_cmd + ' ' + 'DGL_NUM_SERVER=' + str(args.num_servers)
+    if os.environ.get('OMP_NUM_THREADS') is not None:
+        client_cmd = client_cmd + ' ' + 'OMP_NUM_THREADS=' + os.environ.get('OMP_NUM_THREADS')
+    else:
+        client_cmd = client_cmd + ' ' + 'OMP_NUM_THREADS=' + str(args.num_omp_threads)
+    if os.environ.get('PYTHONPATH') is not None:
+        client_cmd = client_cmd + ' ' + 'PYTHONPATH=' + os.environ.get('PYTHONPATH')
+    client_cmd = client_cmd + ' ' + 'DGL_GRAPH_FORMAT=' + str(args.graph_format)
+
+    torch_cmd = '-m torch.distributed.launch'
+    torch_cmd = torch_cmd + ' ' + '--nproc_per_node=' + str(args.num_trainers)
+    torch_cmd = torch_cmd + ' ' + '--nnodes=' + str(len(hosts))
+    torch_cmd = torch_cmd + ' ' + '--node_rank=' + str(0)
+    torch_cmd = torch_cmd + ' ' + '--master_addr=' + str(hosts[0][0])
+    torch_cmd = torch_cmd + ' ' + '--master_port=' + str(1234)
+    for node_id, host in enumerate(hosts):
+        ip, _ = host
+        new_torch_cmd = torch_cmd.replace('node_rank=0', 'node_rank='+str(node_id))
+        if 'python3' in udf_command:
+            new_udf_command = udf_command.replace('python3', 'python3 ' + new_torch_cmd)
+        elif 'python2' in udf_command:
+            new_udf_command = udf_command.replace('python2', 'python2 ' + new_torch_cmd)
+        else:
+            new_udf_command = udf_command.replace('python', 'python ' + new_torch_cmd)
+        cmd = client_cmd + ' ' + new_udf_command
+        cmd = 'cd ' + str(args.workspace) + '; ' + cmd
+        execute_remote(cmd, ip, args.ssh_port, thread_list)
+
+    # Start a cleanup process dedicated for cleaning up remote training jobs.
+    conn1,conn2 = multiprocessing.Pipe()
+    func = partial(get_all_remote_pids, hosts, args.ssh_port, udf_command)
+    process = multiprocessing.Process(target=cleanup_proc, args=(func, conn1))
+    process.start()
+
+    def signal_handler(signal, frame):
+        logging.info('Stop launcher')
+        # We need to tell the cleanup process to kill remote training jobs.
+        conn2.send('cleanup')
+        sys.exit(0)
+    signal.signal(signal.SIGINT, signal_handler)
+
+    for thread in thread_list:
+        thread.join()
+    # The training processes complete. We should tell the cleanup process to exit.
+    conn2.send('exit')
+    process.join()
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Launch a distributed job')
+    parser.add_argument('--ssh_port', type=int, default=22, help='SSH Port.')
+    parser.add_argument('--workspace', type=str,
+                        help='Path of user directory of distributed tasks. \
+                        This is used to specify a destination location where \
+                        the contents of current directory will be rsyncd')
+    parser.add_argument('--num_trainers', type=int,
+                        help='The number of trainer processes per machine')
+    parser.add_argument('--num_omp_threads', type=int,
+                        help='The number of OMP threads per trainer')
+    parser.add_argument('--num_samplers', type=int, default=0,
+                        help='The number of sampler processes per trainer process')
+    parser.add_argument('--num_servers', type=int,
+                        help='The number of server processes per machine')
+    parser.add_argument('--part_config', type=str,
+                        help='The file (in workspace) of the partition config')
+    parser.add_argument('--ip_config', type=str,
+                        help='The file (in workspace) of IP configuration for server processes')
+    parser.add_argument('--num_server_threads', type=int, default=1,
+                        help='The number of OMP threads in the server process. \
+                        It should be small if server processes and trainer processes run on \
+                        the same machine. By default, it is 1.')
+    parser.add_argument('--graph_format', type=str, default='csc',
+                        help='The format of the graph structure of each partition. \
+                        The allowed formats are csr, csc and coo. A user can specify multiple \
+                        formats, separated by ",". For example, the graph format is "csr,csc".')
+    args, udf_command = parser.parse_known_args()
+    assert len(udf_command) == 1, 'Please provide user command line.'
+    assert args.num_trainers is not None and args.num_trainers > 0, \
+            '--num_trainers must be a positive number.'
+    assert args.num_samplers is not None and args.num_samplers >= 0, \
+            '--num_samplers must be a non-negative number.'
+    assert args.num_servers is not None and args.num_servers > 0, \
+            '--num_servers must be a positive number.'
+    assert args.num_server_threads > 0, '--num_server_threads must be a positive number.'
+    assert args.workspace is not None, 'A user has to specify a workspace with --workspace.'
+    assert args.part_config is not None, \
+            'A user has to specify a partition configuration file with --part_config.'
+    assert args.ip_config is not None, \
+            'A user has to specify an IP configuration file with --ip_config.'
+    if args.num_omp_threads is None:
+        # Here we assume all machines have the same number of CPU cores as the machine
+        # where the launch script runs.
+        args.num_omp_threads = max(multiprocessing.cpu_count() // 2 // args.num_trainers, 1)
+        print('The number of OMP threads per trainer is set to', args.num_omp_threads)
+
+    udf_command = str(udf_command[0])
+    if 'python' not in udf_command:
+        raise RuntimeError("DGL launching script can only support Python executable file.")
+    submit_jobs(args, udf_command)
+
+if __name__ == '__main__':
+    fmt = '%(asctime)s %(levelname)s %(message)s'
+    logging.basicConfig(format=fmt, level=logging.INFO)
+    main()
diff --git a/python/partition_graph.py b/python/partition_graph.py
new file mode 100644
index 0000000..1c2d10b
--- /dev/null
+++ b/python/partition_graph.py
@@ -0,0 +1,6 @@
+from dglke.partition_graph import partition_graph
+import sys
+import re
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(partition_graph())


Test patch:
diff --git a/python/dist_test.py b/python/dist_test.py
new file mode 100644
index 0000000..504853e
--- /dev/null
+++ b/python/dist_test.py
@@ -0,0 +1,6 @@
+from dglke.dist_main import dist_main
+import sys
+import re
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(dist_main())