##// END OF EJS Templates
ENH: ssh now allows sending all current environment args to remote ipengine
Satrajit Ghosh -
Show More
@@ -409,8 +409,9 b' bsub < $1'
409 409 os.chmod(bsub_wrapper.name, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
410 410 return bsub_wrapper
411 411
412 sshx_template="""#!/bin/sh
413 "$@" &> /dev/null &
412 sshx_template_prefix="""#!/bin/sh
413 """
414 sshx_template_suffix=""""$@" &> /dev/null &
414 415 echo $!
415 416 """
416 417
@@ -419,10 +420,11 b" ps -fu `whoami` | grep '[i]pengine' | awk '{print $2}' | xargs kill -TERM"
419 420 """
420 421
421 422 class SSHEngineSet(object):
422 sshx_template=sshx_template
423 sshx_template_prefix=sshx_template_prefix
424 sshx_template_suffix=sshx_template_suffix
423 425 engine_killer_template=engine_killer_template
424 426
425 def __init__(self, engine_hosts, sshx=None, ipengine="ipengine"):
427 def __init__(self, engine_hosts, sshx=None, copyenvs=None, ipengine="ipengine"):
426 428 """Start a controller on localhost and engines using ssh.
427 429
428 430 The engine_hosts argument is a dict with hostnames as keys and
@@ -441,7 +443,11 b' class SSHEngineSet(object):'
441 443 '%s-main-sshx.sh' % os.environ['USER']
442 444 )
443 445 f = open(self.sshx, 'w')
444 f.writelines(self.sshx_template)
446 f.writelines(self.sshx_template_prefix)
447 if copyenvs:
448 for key, val in os.environ.items():
449 f.writelines('export %s=%s\n'%(key,val))
450 f.writelines(self.sshx_template_suffix)
445 451 f.close()
446 452 self.engine_command = ipengine
447 453 self.engine_hosts = engine_hosts
@@ -806,7 +812,8 b' def main_ssh(args):'
806 812 cl = ControllerLauncher(extra_args=cont_args)
807 813 dstart = cl.start()
808 814 def start_engines(cont_pid):
809 ssh_set = SSHEngineSet(clusterfile['engines'], sshx=args.sshx)
815 ssh_set = SSHEngineSet(clusterfile['engines'], sshx=args.sshx,
816 copyenvs=args.copyenvs)
810 817 def shutdown(signum, frame):
811 818 d = ssh_set.kill()
812 819 cl.interrupt_then_kill(1.0)
@@ -991,6 +998,14 b' def get_args():'
991 998 parents=[base_parser]
992 999 )
993 1000 parser_ssh.add_argument(
1001 '-e',
1002 '--copyenvs',
1003 action='store_true',
1004 dest='copyenvs',
1005 help='Copy current shell environment to remote location',
1006 default=False,
1007 )
1008 parser_ssh.add_argument(
994 1009 '--clusterfile',
995 1010 type=str,
996 1011 dest='clusterfile',
General Comments 0
You need to be logged in to leave comments. Login now