mpiexec.py

fortran并行计算包

#!/usr/bin/env python
#
#   (C) 2001 by Argonne National Laboratory.
#       See COPYRIGHT in top-level directory.
#

"""
usage:
mpiexec [-h or -help or --help]    # get this message
mpiexec -file filename             # (or -f) filename contains XML job description
mpiexec [global args] [local args] executable [args]
   where global args may be
      -l                           # line labels by MPI rank
      -bnr                         # MPICH1 compatibility mode
      -machinefile                 # file mapping procs to machines
      -s <spec>                    # direct stdin to "all" or 1,2 or 2-4,6 
      -1                           # override default of trying 1st proc locally
      -ifhn                        # network interface to use locally
      -tv                          # run procs under totalview (must be installed)
      -tvsu                        # totalview startup only
      -gdb                         # run procs under gdb
      -m                           # merge output lines (default with gdb)
      -a                           # means assign this alias to the job
      -ecfn                        # output_xml_exit_codes_filename
      -g<local arg name>           # global version of local arg (below)
    and local args may be
      -n <n> or -np <n>            # number of processes to start
      -wdir <dirname>              # working directory to start in
      -umask <umask>               # umask for remote process
      -path <dirname>              # place to look for executables
      -host <hostname>             # host to start on
      -soft <spec>                 # modifier of -n value
      -arch <arch>                 # arch type to start on (not implemented)
      -envall                      # pass all env vars in current environment
      -envnone                     # pass no env vars
      -envlist <list of env var names> # pass current values of these vars
      -env <name> <value>          # pass this value of this env var
mpiexec [global args] [local args] executable args : [local args] executable...
mpiexec -gdba jobid                # gdb-attach to existing jobid
mpiexec -configfile filename       # filename contains cmd line segs as lines
  (See User Guide for more details)

Examples:
   mpiexec -l -n 10 cpi 100
   mpiexec -genv QPL_LICENSE 4705 -n 3 a.out

   mpiexec -n 1 -host foo master : -n 4 -host mysmp slave
"""
from time import ctime
__author__ = "Ralph Butler and Rusty Lusk"
__date__ = ctime()
__version__ = "$Revision: 1.90 $"
__credits__ = ""

import signal
if hasattr(signal,'SIGTTIN'):
    signal.signal(signal.SIGTTIN,signal.SIG_IGN)    # asap

import sys, os, socket, re

from  urllib import quote
from  time   import time
from  urllib import unquote
from  mpdlib import mpd_set_my_id, mpd_get_my_username, mpd_version, mpd_print, \
                    mpd_uncaught_except_tb, mpd_handle_signal, mpd_which, \
                    MPDListenSock, MPDStreamHandler, MPDConClientSock, MPDParmDB

try:
    import pwd
    pwd_module_available = 1
except:
    pwd_module_available = 0

global parmdb, nextRange, appnum
global numDoneWithIO, myExitStatus, sigOccurred, outXmlDoc, outECs

recvTimeout = 20  # const


def mpiexec():
    global parmdb, nextRange, appnum
    global numDoneWithIO, myExitStatus, sigOccurred, outXmlDoc, outECs

    import sys  # for sys.excepthook on next line
    sys.excepthook = mpd_uncaught_except_tb

    myExitStatus = 0
    if len(sys.argv) < 2  or  sys.argv[1] == '-h'  \
    or  sys.argv[1] == '-help'  or  sys.argv[1] == '--help':
	usage()
    myHost = socket.gethostname()
    mpd_set_my_id(myid='mpiexec_%s' % (myHost) )
    try:
        hostinfo = socket.gethostbyname_ex(myHost)
    except:
        print 'mpiexec failed: gethostbyname_ex failed for %s' % (myHost)
        sys.exit(-1)
    myIP = hostinfo[2][0]

    parmdb = MPDParmDB(orderedSources=['cmdline','xml','env','rcfile','thispgm'])
    parmsToOverride = {
                        'MPD_USE_ROOT_MPD'            :  0,
                        'MPD_SECRETWORD'              :  '',
                        'MPIEXEC_SHOW_LINE_LABELS'    :  0,
                        'MPIEXEC_LINE_LABEL_FMT'      :  '%r',
                        'MPIEXEC_JOB_ALIAS'           :  '',
                        'MPIEXEC_USIZE'               :  0,
                        'MPIEXEC_GDB'                 :  0,
                        'MPIEXEC_IFHN'                :  '',  # use one from mpd as default
                        'MPIEXEC_MERGE_OUTPUT'        :  0,
                        'MPIEXEC_STDIN_DEST'          :  '0',
                        'MPIEXEC_MACHINEFILE'         :  '',
                        'MPIEXEC_BNR'                 :  0,
                        'MPIEXEC_TOTALVIEW'           :  0,
                        'MPIEXEC_TVSU'                :  0,
                        'MPIEXEC_EXITCODES_FILENAME'  :  '',
                        'MPIEXEC_TRY_1ST_LOCALLY'     :  1,
                        'MPIEXEC_TIMEOUT'             :  0,
                        'MPIEXEC_HOST_LIST'           :  [],
                        'MPIEXEC_HOST_CHECK'          :  0,
                      }
    for (k,v) in parmsToOverride.items():
        parmdb[('thispgm',k)] = v
    parmdb[('thispgm','mship')] = ''
    parmdb[('thispgm','rship')] = ''
    parmdb[('thispgm','userpgm')] = ''
    parmdb[('thispgm','nprocs')] = 0
    parmdb[('thispgm','ecfn_format')] = ''
    parmdb[('thispgm','-gdba')] = ''
    parmdb[('thispgm','singinitpid')] = 0
    parmdb[('thispgm','singinitport')] = 0
    parmdb[('thispgm','ignore_rcfile')] = 0
    parmdb[('thispgm','ignore_environ')] = 0
    parmdb[('thispgm','inXmlFilename')] = ''
    parmdb[('thispgm','print_parmdb_all')] = 0
    parmdb[('thispgm','print_parmdb_def')] = 0

    appnum = 0
    nextRange = 0
    localArgSets = { 0 : [] }

    if sys.argv[1] == '-gdba':
	if len(sys.argv) != 3:
            print '-gdba arg must appear only with jobid'
	    usage()
        parmdb[('cmdline','-gdba')] = sys.argv[2]
        parmdb[('cmdline','MPIEXEC_GDB')] = 1
        parmdb[('cmdline','MPIEXEC_MERGE_OUTPUT')] = 1       # implied
        parmdb[('cmdline','MPIEXEC_SHOW_LINE_LABELS')] = 1   # implied
        parmdb[('cmdline','MPIEXEC_STDIN_DEST')]   = 'all'   # implied
    elif sys.argv[1] == '-file'  or  sys.argv[1] == '-f':
	if len(sys.argv) != 3:
            print '-file (-f) arg must appear alone'
	    usage()
        parmdb[('cmdline','inXmlFilename')] = sys.argv[2]
    elif sys.argv[1] == '-pmi_args':
        parmdb[('cmdline','singinitport')] = sys.argv[2]
        # ignoring interface name (where app is listening) and authentication key, for now
        parmdb[('cmdline','singinitpid')]  = sys.argv[5]
        parmdb[('cmdline','userpgm')] = 'unknown_pgmname'
        parmdb[('cmdline','nprocs')] = 1
        parmdb[('cmdline','MPIEXEC_TRY_1ST_LOCALLY')] = 1
        machineFileInfo = {}
        tempargv = [sys.argv[0],'unknown_pgmname']
        collect_args(tempargv,localArgSets)
    else:
        if sys.argv[1] == '-configfile':
	    if len(sys.argv) != 3:
	        usage()
            configFile = open(sys.argv[2],'r',0)
            configLines = configFile.readlines()
            configLines = [ x.strip() + ' : '  for x in configLines if x[0] != '#' and x.strip() != '' ]
            tempargv = []
            for line in configLines:
                line = 'mpddummyarg ' + line  # gets pitched in shells that can't handle --
                (shellIn,shellOut) = \
                    os.popen4("/bin/sh -c 'for a in $*; do echo _$a; done' -- %s" % (line))
                for shellLine in shellOut:
                    if shellLine.startswith('_mpddummyarg'):
                        continue
                    tempargv.append(shellLine[1:].strip())    # 1: strips off the leading _
	    tempargv = [sys.argv[0]] + tempargv[0:-1]   # strip off the last : I added
            collect_args(tempargv,localArgSets)
        else:
            collect_args(sys.argv,localArgSets)
        machineFileInfo = read_machinefile(parmdb['MPIEXEC_MACHINEFILE'])

    # set some default values for mpd; others added as discovered below
    msgToMPD = { 'cmd'            : 'mpdrun',
                 'conhost'        : myHost,
                 'spawned'        : 0,
                 'nstarted'       : 0,
                 'hosts'          : {},
                 'execs'          : {},
                 'users'          : {},
                 'cwds'           : {},
                 'umasks'         : {},
                 'paths'          : {},
                 'args'           : {},
                 'limits'         : {},
                 'envvars'        : {},
                 'ifhns'          : {},
               }

    if parmdb['inXmlFilename']:
        get_parms_from_xml_file(msgToMPD)  # fills in some more values of msgToMPD
    else:
        parmdb.get_parms_from_env(parmsToOverride)
        parmdb.get_parms_from_rcfile(parmsToOverride)

    # mostly old mpdrun below here
    numDoneWithIO = 0
    outXmlDoc = ''
    outECs = ''
    outECFile = None
    sigOccurred = 0

    listenSock = MPDListenSock('',0,name='socket_to_listen_for_man')
    listenPort = listenSock.getsockname()[1]
    if (hasattr(os,'getuid')  and  os.getuid() == 0)  or  parmdb['MPD_USE_ROOT_MPD']:
        fullDirName = os.path.abspath(os.path.split(sys.argv[0])[0])  # normalize
        mpdroot = os.path.join(fullDirName,'mpdroot')
        conSock = MPDConClientSock(mpdroot=mpdroot,secretword=parmdb['MPD_SECRETWORD'])
    else:
        conSock = MPDConClientSock(secretword=parmdb['MPD_SECRETWORD'])

    if parmdb['MPIEXEC_HOST_CHECK']:    # if this was requested in the xml file
        msgToSend = { 'cmd' : 'verify_hosts_in_ring',
                      'host_list' : parmdb['MPIEXEC_HOST_LIST'] }
        conSock.send_dict_msg(msgToSend)
        msg = conSock.recv_dict_msg(timeout=recvTimeout)
        if not msg:
            mpd_print(1,'no msg recvd from mpd for verify_hosts_in_ring')
            sys.exit(-1)
        elif msg['cmd'] != 'verify_hosts_in_ring_response':
            mpd_print(1,'unexpected msg from mpd :%s:' % (msg) )
            sys.exit(-1)
        if msg['host_list']:
            print 'These hosts are not in the mpd ring:'
            for host in  msg['host_list']:
                if host[0].isdigit():
                    print '    %s' % (host),
                    try:
                        print ' (%s)' % (socket.gethostbyaddr(host)[0])
                    except:
                        print ''
                else:
                    print '    %s' % (host)
            sys.exit(-1)

    msgToSend = { 'cmd' : 'get_mpdrun_values' }
    conSock.send_dict_msg(msgToSend)
    msg = conSock.recv_dict_msg(timeout=recvTimeout)
    if not msg:
        mpd_print(1, 'no msg recvd from mpd during version check')
        sys.exit(-1)
    elif msg['cmd'] != 'response_get_mpdrun_values':
        mpd_print(1,'unexpected msg from mpd :%s:' % (msg) )
        sys.exit(-1)
    if msg['mpd_version'] != mpd_version():
        mpd_print(1,'mpd version %s does not match mpiexec version %s' % \
                  (msg['mpd_version'],mpd_version()) )
        sys.exit(-1)

    # if using/testing the INET CONSOLE
    if os.environ.has_key('MPD_CON_INET_HOST_PORT'):
        try:
            myIfhn = socket.gethostbyname_ex(myHost)[2][0]
        except:
            print 'mpiexec failed: gethostbyname_ex failed for %s' % (myHost)
            sys.exit(-1)
        parmdb[('thispgm','MPIEXEC_IFHN')] = myIfhn
    elif not parmdb['MPIEXEC_IFHN']:    # if user did not specify one, use mpd's
        parmdb[('thispgm','MPIEXEC_IFHN')] = msg['mpd_ifhn']    # not really thispgm here

    if parmdb['-gdba']:
        get_vals_for_attach(parmdb,conSock,msgToMPD)
    elif not parmdb['inXmlFilename']:
        parmdb[('cmdline','nprocs')] = 0  # for incr later
        for k in localArgSets.keys():
	    handle_local_argset(localArgSets[k],machineFileInfo,msgToMPD)

    if parmdb['MPIEXEC_MERGE_OUTPUT']  and  not parmdb['MPIEXEC_SHOW_LINE_LABELS']:
        parmdb[('thispgm','MPIEXEC_SHOW_LINE_LABELS')] = 1   # causes line labels also

    if parmdb['print_parmdb_all']:
        parmdb.printall()
    if parmdb['print_parmdb_def']:
        parmdb.printdef()

    if parmdb['mship']:
        mshipSock = MPDListenSock('',0,name='socket_for_mship')
        mshipPort = mshipSock.getsockname()[1]
        mshipPid = os.fork()
        if mshipPid == 0:
            conSock.close()
            os.environ['MPDCP_AM_MSHIP'] = '1'
            os.environ['MPDCP_MSHIP_PORT'] = str(mshipPort)
            os.environ['MPDCP_MSHIP_FD'] = str(mshipSock.fileno())
            os.environ['MPDCP_MSHIP_NPROCS'] = str(parmdb['nprocs'])
            try:
                os.execvpe(parmdb['mship'],[parmdb['MPIEXEC_MSHIP']],os.environ)
            except Exception, errmsg:
                mpd_print(1,'execvpe failed for copgm %s; errmsg=:%s:' % \
                          (parmdb['MPIEXEC_MSHIP'],errmsg))
                sys.exit(-1)
            os._exit(0);  # do NOT do cleanup
        mshipSock.close()
    else:
        mshipPid = 0

    # make sure to do this after nprocs has its value
    linesPerRank = {}  # keep this a dict instead of a list
    for i in range(parmdb['nprocs']):
        linesPerRank[i] = []

    if parmdb['MPIEXEC_EXITCODES_FILENAME']:
        if parmdb['ecfn_format'] == 'xml':
            try:
                import xml.dom.minidom
            except:
                print 'you requested to save the exit codes in an xml file, but'
                print '  I was unable to import the xml.dom.minidom module'
                sys.exit(-1)
            outXmlDoc = xml.dom.minidom.Document()
            outECs = outXmlDoc.createElement('exit-codes')
            outXmlDoc.appendChild(outECs)
        else:
            outECs = 'exit-codes\n'

    msgToMPD['nprocs'] = parmdb['nprocs']
    msgToMPD['limits'][(0,parmdb['nprocs']-1)]  = {}
    msgToMPD['conport'] = listenPort
    msgToMPD['conip'] = myIP
    msgToMPD['conifhn'] = parmdb['MPIEXEC_IFHN']
    if parmdb['MPIEXEC_JOB_ALIAS']:
        msgToMPD['jobalias'] = parmdb['MPIEXEC_JOB_ALIAS']
    else:
        msgToMPD['jobalias'] = ''
    if parmdb['MPIEXEC_TRY_1ST_LOCALLY']:
        msgToMPD['try_1st_locally'] = 1
    if parmdb['rship']:
        msgToMPD['rship'] = parmdb['rship']
        msgToMPD['mship_host'] = socket.gethostname()
        msgToMPD['mship_port'] = mshipPort
    if parmdb['MPIEXEC_BNR']:
        msgToMPD['doing_bnr'] = 1
    if parmdb['MPIEXEC_STDIN_DEST'] == 'all':
        stdinDest = '0-%d' % (parmdb['nprocs']-1)
    else:
        stdinDest = parmdb['MPIEXEC_STDIN_DEST']
    if parmdb['MPIEXEC_SHOW_LINE_LABELS']:
        msgToMPD['line_labels'] = parmdb['MPIEXEC_LINE_LABEL_FMT']
    else:
        msgToMPD['line_labels'] = ''
    msgToMPD['stdin_dest'] = stdinDest
    msgToMPD['gdb'] = parmdb['MPIEXEC_GDB']