9725cf2aeb
fb13cab change makefile 1479e37 fixed small bug in mpi submission script 0ca7a63 Update README.md 5ef4830 ok 93a1338 chg note git-subtree-dir: subtree/rabit git-subtree-split: fb13cab216b795f86dc90547b71c0f730766affa
44 lines
1.8 KiB
Python
Executable File
44 lines
1.8 KiB
Python
Executable File
#!/usr/bin/python
|
|
"""
|
|
This is the demo submission script of rabit, it is created to
|
|
submit rabit jobs using hadoop streaming
|
|
"""
|
|
import argparse
|
|
import sys
|
|
import os
|
|
import subprocess
|
|
import rabit_tracker as tracker
|
|
|
|
parser = argparse.ArgumentParser(description='Rabit script to submit rabit job using MPI')
|
|
parser.add_argument('-n', '--nworker', required=True, type=int,
|
|
help = 'number of worker proccess to be launched')
|
|
parser.add_argument('-v', '--verbose', default=0, choices=[0, 1], type=int,
|
|
help = 'print more messages into the console')
|
|
parser.add_argument('-H', '--hostfile', type=str,
|
|
help = 'the hostfile of mpi server')
|
|
parser.add_argument('command', nargs='+',
|
|
help = 'command for rabit program')
|
|
args = parser.parse_args()
|
|
#
|
|
# submission script using MPI
|
|
#
|
|
def mpi_submit(nslave, worker_args):
|
|
"""
|
|
customized submit script, that submit nslave jobs, each must contain args as parameter
|
|
note this can be a lambda function containing additional parameters in input
|
|
Parameters
|
|
nslave number of slave process to start up
|
|
args arguments to launch each job
|
|
this usually includes the parameters of master_uri and parameters passed into submit
|
|
"""
|
|
sargs = ' '.join(args.command + worker_args)
|
|
if args.hostfile is None:
|
|
cmd = ' '.join(['mpirun -n %d' % (nslave)] + args.command + worker_args)
|
|
else:
|
|
cmd = ' '.join(['mpirun -n %d --hostfile %s' % (nslave, args.hostfile)] + args.command + worker_args)
|
|
print cmd
|
|
subprocess.check_call(cmd, shell = True)
|
|
|
|
# call submit, with nslave, the commands to run each job and submit function
|
|
tracker.submit(args.nworker, [], fun_submit = mpi_submit, verbose = args.verbose)
|