#!/usr/bin/python

from multiprocessing import Process
import os
import sys, getopt

#------------------
# main
#------------------
def Usage():
	print ""
	print "Usage:"
	print "       start_remote_processes [options] command"
	print ""
	print " This script is used to start a process on several remote"
	print "computers. The same command line will be sent to each of"
	print "the remote machines using \"ssh\". The list of nodes to which"
	print "the command is sent may be specified explicitly on the command"
	print "line or via an input file containing one node name per line."
	print ""
	print "options:"
	print ""
	print "  -h                  Print usage statement"
	print "  -f file             Name of file from which to get node names"
	print "  -n node1,node2,...  Specify remote nodes to execute command"
	print ""
	print "Note that multiple -f and -n arguments can be given. If a node"
	print "appears more than once then it will have the command executed"
	print "more thant once. (So be cautious if specifying more than one"
	print "file which may both contain the same node name.)"
	print ""
	print "Examples:"
	print ""
	print "    start_remote_processes -n L3_nodes.txt start_L3_proc"
	print ""
	print "This will start a program called \"start_L3_proc\" on each"
	print "node listed in the file \"L3nodes.txt\". In this case,"
	print "\"start_L3_proc\" should be a script that launches the actual"
	print "process in the background and then returns immediately."
	print "otherwise, the start_remote_processes will wait until all"
	print "of the remote processes end before finishing."
	print ""
	print "   start_remote_processes -n gluon44,gluon45 file2et run1001_01.evio >& /dev/null &"
	print ""
	print "This will run the program \"file2et\" on both gluon44 and"
	print "gluon45. The stdout and stderr streams of both programs are"
	print "redirected to /dev/null and the process put into the background"
	print "by having \">& /dev/null &\" at the end of the command."
	print "This allows the program to be run in the background and the"
	print "ssh command to return right away without needing to maintain"
	print "a connection to the remote machine."

#------------------
# RunProg
#------------------
def RunProg(node, cmd):
	os.system("ssh %s \"%s\"" % (node, cmd))

#------------------
# main
#------------------
def main(argv):

	if (len(sys.argv) < 2):
		Usage()
		exit(0)

	nodes = []
	try:
		opts, args = getopt.getopt(argv,"hf:n:",["help","file=","nodes="])
	except getopt.GetoptError:
		Usage()
		sys.exit(2)
	for opt, arg in opts:
		if opt in ("-h", "--help"):
			Usage()
			sys.exit()
		elif opt in ("-f", "--file"):
			with open(arg) as f:
				nodes.extend(f.read().splitlines())
		elif opt in ("-n", "--nodes"):
			nodes.extend(arg.split(','))


	# Command is contained in the non-option arguments list "args"
	if(len(args)<1):
		print "\nYOU MUST SPECIFY A COMMAND TO RUN!"
		Usage()
		exit(2)

	command = ""
	for a in args:
		command += a+" "
	command = command.strip()

	# Check for empty node names. This can happen if a file is specified and
	# has extra carriage returns.
	tmp_nodes = []
	for n in nodes:
		tmp = n.strip()
		if( len(tmp)>0 ): tmp_nodes.append(tmp)
	nodes = tmp_nodes
	
	print "   command: '%s'" % (command)
	print "num. nodes: %d" % (len(nodes))


	procs = []
	for n in nodes:
		p = Process(target=RunProg, args=(n,command))
		print "Running command on: %s" % (n)
		p.start()
		procs.append(p)

	for p in procs:
		p.join()


# The following "if" is for Windows
if __name__ == "__main__":
    main(sys.argv[1:])