mpdman.py

mpi并行计算的c++代码 可用vc或gcc编译通过 可以用来搭建并行计算试验环境

            ## NOTE: if you spawn a non-MPI job, it may not send these msgs
            ## in which case adding 2 to numWithIO will cause the pgm to hang.
            totspawns = int(parsedMsg['totspawns'])
            spawnssofar = int(parsedMsg['spawnssofar'])
            if spawnssofar == 1: # this is the first of possibly several spawn msgs
                self.numWithIO += 2
                self.tpsf = 0             # total processes spawned so far
                self.spawnExecs = {}      # part of MPI_Spawn_multiple args
                self.spawnHosts = {}      # comes from info
                self.spawnUsers = {}      # always the current user
                self.spawnCwds  = {}      # could come from info, but doesn't yet
                self.spawnUmasks = {}     # could come from info, but doesn't yet
                self.spawnPaths = {}      # could come from info, but doesn't yet
                self.spawnEnvvars = {}    # whole environment from mpiexec, plus appnum
                self.spawnLimits = {}
                self.spawnArgs = {}
            self.spawnNprocs  = int(parsedMsg['nprocs']) # num procs in this spawn
            pmiInfo = {}
            for i in range(0,int(parsedMsg['info_num'])):
                info_key = parsedMsg['info_key_%d' % i]
                info_val = parsedMsg['info_val_%d' % i]
                pmiInfo[info_key] = info_val

            if pmiInfo.has_key('host'):
                try:
                    toIfhn = socket.gethostbyname_ex(pmiInfo['host'])[2][0]
                    self.spawnHosts[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = toIfhn
                except:
                    mpd_print(1, "unable to obtain host info for :%s:" % (pmiInfo['host']))
                    pmiMsgToSend = 'cmd=spawn_result rc=-2 status=unknown_host\n'
                    self.pmiSock.send_char_msg(pmiMsgToSend)
                    return
            else:
                self.spawnHosts[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = '_any_'
            if pmiInfo.has_key('path'):
                self.spawnPaths[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = pmiInfo['path']
            else:
                self.spawnPaths[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = os.environ['MPDMAN_CLI_PATH']
            if pmiInfo.has_key('wdir'):
                self.spawnCwds[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = pmiInfo['wdir']
            else:
                self.spawnCwds[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = os.environ['MPDMAN_CWD']
            if pmiInfo.has_key('umask'):
                self.spawnUmasks[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = pmiInfo['umask'] 
            else:
                self.spawnUmasks[(self.tpsf,self.tpsf+self.spawnNprocs-1)]  = os.environ['MPDMAN_UMASK']
            self.spawnExecs[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = parsedMsg['execname']
            self.spawnUsers[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = mpd_get_my_username()
            self.spawnEnv = {}
            self.spawnEnv.update(os.environ)
            self.spawnEnv['MPI_APPNUM'] = str(spawnssofar-1)
            self.spawnEnvvars[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = self.spawnEnv
            self.spawnLimits[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = {} # not implemented yet
            ##### args[(tpsf,tpsf+spawnNprocs-1) = [ parsedMsg['args'] ]
            ##### args[(tpsf,tpsf+spawnNprocs-1) = [ 'AA', 'BB', 'CC' ]
            cliArgs = []
            cliArgcnt = int(parsedMsg['argcnt'])
            for i in range(1,cliArgcnt+1):    # start at 1
                cliArgs.append(parsedMsg['arg%d' % i])
            self.spawnArgs[(self.tpsf,self.tpsf+self.spawnNprocs-1)] = cliArgs
            self.tpsf += self.spawnNprocs

            if totspawns == spawnssofar:    # This is the last in the spawn sequence
                self.spawnedCnt += 1    # non-zero to use for creating kvsname in msg below
                msgToSend = { 'cmd'          : 'spawn',
                              'conhost'      : self.myHost,
                              'conifhn'      : self.myIfhn,
                              'conport'      : self.listenNonRingPort,
                              'spawned'      : self.spawnedCnt,
                              'jobid'        : self.jobid,
                              'nstarted'     : 0,
                              'nprocs'       : self.tpsf,
                              'hosts'        : self.spawnHosts,
                              'execs'        : self.spawnExecs,
                              'users'        : self.spawnUsers,
                              'cwds'         : self.spawnCwds,
                              'umasks'       : self.spawnUmasks,
                              'paths'        : self.spawnPaths,
                              'args'         : self.spawnArgs,
                              'envvars'      : self.spawnEnvvars,
                              'limits'       : self.spawnLimits,
                              'singinitpid'  : 0,
                              'singinitport' : 0,
                            }
                msgToSend['line_labels'] = self.lineLabelFmt
                self.mpdSock.send_dict_msg(msgToSend)
                # I could send the preput_info along but will keep it here
                # and let the spawnee call me up and ask for it; he will
                # call me anyway since I am his parent in the tree.  So, I
                # will create a KVS to hold the info until he calls
                self.spawnedKVSname = 'mpdman_kvs_for_spawned_' + str(self.spawnedCnt)
                self.KVSs[self.spawnedKVSname] = {}
                preput_num = int(parsedMsg['preput_num'])
                for i in range(0,preput_num):
                    preput_key = parsedMsg['preput_key_%d' % i]
                    preput_val = parsedMsg['preput_val_%d' % i]
                    self.KVSs[self.spawnedKVSname][preput_key] = preput_val
        elif parsedMsg['cmd'] == 'finalize':
            # the following lines are present to support a process that runs
            # 2 MPI pgms in tandem (e.g. mpish at ANL)
            self.KVSs = {}
            self.KVSs[self.default_kvsname] = {}
            self.kvs_next_id = 1
            self.jobEndingEarly = 0
            self.pmiCollectiveJob = 0
            self.spawnedCnt = 0
            pmiMsgToSend = 'cmd=finalize_ack\n' 
            self.pmiSock.send_char_msg(pmiMsgToSend)
        elif parsedMsg['cmd'] == 'client_bnr_fence_in':    ## BNR
            self.pmiBarrierInRecvd = 1
            if self.myRank == 0  or  self.holdingPMIBarrierLoop1:
                msgToSend = { 'cmd' : 'pmi_barrier_loop_1' }
                self.ring.rhsSock.send_dict_msg(msgToSend)
        elif parsedMsg['cmd'] == 'client_bnr_put':         ## BNR
            key = parsedMsg['attr']
            value = parsedMsg['val']
            try:
                self.KVSs[self.default_kvsname][key] = value
                pmiMsgToSend = 'cmd=put_result rc=0\n'
                self.pmiSock.send_char_msg(pmiMsgToSend)
            except Exception, errmsg:
                pmiMsgToSend = 'cmd=put_result rc=-1 msg="%s"\n' % errmsg
                self.pmiSock.send_char_msg(pmiMsgToSend)
        elif parsedMsg['cmd'] == 'client_bnr_get':          ## BNR
            key = parsedMsg['attr']
            if self.KVSs[self.default_kvsname].has_key(key):
                value = self.KVSs[self.default_kvsname][key]
                pmiMsgToSend = 'cmd=client_bnr_get_output rc=0 val=%s\n' % (value)
                self.pmiSock.send_char_msg(pmiMsgToSend)
            else:
                msgToSend = { 'cmd' : 'bnr_get', 'key' : key,
                              'kvsname' : kvsname, 'from_rank' : self.myRank }
                self.ring.rhsSock.send_dict_msg(msgToSend)
        elif parsedMsg['cmd'] == 'client_ready':               ## BNR
            ## continue to wait for accepting_signals
            pass
        elif parsedMsg['cmd'] == 'accepting_signals':          ## BNR
            ## handle it like a barrier_in ??
            self.pmiBarrierInRecvd = 1
            self.doingBNR = 1    ## BNR  # set again is OK
        elif parsedMsg['cmd'] == 'interrupt_peer_with_msg':    ## BNR
            self.ring.rhsSock.send_dict_msg(parsedMsg)
        else:
            mpd_print(1, "unrecognized pmi msg :%s:" % line )
    def handle_console_input(self,sock):
        msg = self.conSock.recv_dict_msg()
        if not msg:
            if self.conSock:
                self.streamHandler.del_handler(self.conSock)
                self.conSock.close()
                self.conSock = 0
            if self.parentStdoutSock:
                self.streamHandler.del_handler(self.parentStdoutSock)
                self.parentStdoutSock.close()
                self.parentStdoutSock = 0
            if self.parentStderrSock:
                self.streamHandler.del_handler(self.parentStderrSock)
                self.parentStderrSock.close()
                self.parentStderrSock = 0
            if self.ring.rhsSock:
                msgToSend = { 'cmd' : 'signal', 'signo' : 'SIGKILL' }
                self.ring.rhsSock.send_dict_msg(msgToSend)
            try:
                pgrp = clientPid * (-1)   # neg Pid -> group
                os.kill(pgrp,signal.SIGKILL)   # may be reaped by sighandler
            except:
                pass
        elif msg['cmd'] == 'signal':
            if msg['signo'] == 'SIGINT':
                mpd_print(0000,"RECVD ON CON: SIGINT")
                self.ring.rhsSock.send_dict_msg(msg)
                for s in self.spawnedChildSocks:
                    s.send_dict_msg(msg)
                if self.gdb:
                    os.kill(clientPid,signal.SIGINT)
                else:
                    try:
                        pgrp = clientPid * (-1)   # neg Pid -> group
                        os.kill(pgrp,signal.SIGKILL)   # may be reaped by sighandler
                    except:
                        pass
            elif msg['signo'] == 'SIGKILL':
                self.ring.rhsSock.send_dict_msg(msg)
                for s in self.spawnedChildSocks:
                    s.send_dict_msg(msg)
                if self.gdb:
                    os.kill(clientPid,signal.SIGUSR1)    # tell gdb driver to kill all
                else:
                    try:
                        pgrp = clientPid * (-1)   # neg Pid -> group
                        os.kill(pgrp,signal.SIGKILL)   # may be reaped by sighandler
                    except:
                        pass
            elif msg['signo'] == 'SIGTSTP':
                msg['dest'] = self.myId
                self.ring.rhsSock.send_dict_msg(msg)
                try:
                    pgrp = clientPid * (-1)   # neg Pid -> group
                    os.kill(pgrp,signal.SIGTSTP)   # may be reaped by sighandler
                except:
                    pass
            elif msg['signo'] == 'SIGCONT':
                msg['dest'] = self.myId
                self.ring.rhsSock.send_dict_msg(msg)
                try:
                    pgrp = clientPid * (-1)   # neg Pid -> group
                    os.kill(pgrp,signal.SIGCONT)   # may be reaped by sighandler
                except:
                    pass
        elif msg['cmd'] == 'stdin_from_user':
            msg['src'] = self.myId
            self.ring.rhsSock.send_dict_msg(msg)
            if in_stdinRcvrs(self.myRank,self.stdinDest):
                try:
                    if msg.has_key('eof'):
                        if self.subproc:    # must close subproc's file (not just the fd)
                            self.subproc.stdin.close()
                        else:
                            os.close(self.fd_write_cli_stdin)
                    else:
                        os.write(self.fd_write_cli_stdin,msg['line'])
                except:
                    mpd_print(1, 'cannot send stdin to client')
        elif msg['cmd'] == 'stdin_dest':
            self.stdinDest = msg['stdin_procs']
            msg['src'] = self.myId
            self.ring.rhsSock.send_dict_msg(msg)
        elif msg['cmd'] == 'tv_ready':
            self.tvReady = 1
            msg['src'] = self.myId
            self.ring.rhsSock.send_dict_msg(msg)
            if self.pmiSock:    # should be valid sock if running tv
                pmiMsgToSend = 'cmd=tv_ready\n'
                self.pmiSock.send_char_msg(pmiMsgToSend)
        else:
            mpd_print(1, 'unexpected msg recvd on conSock :%s:' % msg )
    def handle_mpd_input(self,sock):
        msg = self.mpdSock.recv_dict_msg()
        mpd_print(0000, 'msg recvd on mpdSock :%s:' % msg )
        if not msg:
            if self.conSock:
                msgToSend = { 'cmd' : 'job_aborted', 'reason' : 'mpd disappeared',
                              'jobid' : self.jobid }
                self.conSock.send_dict_msg(msgToSend,errprint=0)
                self.streamHandler.del_handler(self.conSock)
                self.conSock.close()
                self.conSock = 0
            try:
                os.kill(0,signal.SIGKILL)  # pid 0 -> all in my process group
            except:
                pass
            sys.exit(0)
        if msg['cmd'] == 'abortjob':
            mpd_print(1, "job aborted by mpd; reason=%s" % (msg['reason']))
        elif msg['cmd'] == 'signal_to_handle'  and  msg.has_key('sigtype'):
            if msg['sigtype'].isdigit():
                signum = int(msg['sigtype'])
            else:
                exec