pbs.py revision 1376:190de61fed5a
12SN/A# Copyright (c) 2005 The Regents of The University of Michigan
21762SN/A# All rights reserved.
32SN/A#
42SN/A# Redistribution and use in source and binary forms, with or without
52SN/A# modification, are permitted provided that the following conditions are
62SN/A# met: redistributions of source code must retain the above copyright
72SN/A# notice, this list of conditions and the following disclaimer;
82SN/A# redistributions in binary form must reproduce the above copyright
92SN/A# notice, this list of conditions and the following disclaimer in the
102SN/A# documentation and/or other materials provided with the distribution;
112SN/A# neither the name of the copyright holders nor the names of its
122SN/A# contributors may be used to endorse or promote products derived from
132SN/A# this software without specific prior written permission.
142SN/A#
152SN/A# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
162SN/A# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
172SN/A# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
182SN/A# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
192SN/A# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
202SN/A# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
212SN/A# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
222SN/A# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
232SN/A# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
242SN/A# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
252SN/A# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
262SN/A#
272665Ssaidi@eecs.umich.edu# Authors: Nathan Binkert
282665Ssaidi@eecs.umich.edu
292665Ssaidi@eecs.umich.eduimport os, re, sys
302SN/A
312SN/Adef ssh(host, script, tty = False, user = ''):
321717SN/A    args = [ 'ssh', '-x' ]
331717SN/A    if user:
342SN/A        args.append('-l' + user)
352SN/A    if tty:
362SN/A        args.append('-t')
374182Sgblack@eecs.umich.edu    args.append(host)
38707SN/A    args.append(script)
391858SN/A
4056SN/A    return os.spawnvp(os.P_WAIT, args[0], args)
414776Sgblack@eecs.umich.edu
422856Srdreslin@umich.educlass qsub:
432SN/A    def __init__(self):
443520Sgblack@eecs.umich.edu        self.hold = False
453520Sgblack@eecs.umich.edu        self.join = False
463520Sgblack@eecs.umich.edu        self.keep_stdout = False
473520Sgblack@eecs.umich.edu        self.keep_stderr = False
482190SN/A        self.node_type = ''
492315SN/A        self.mail_abort = False
502680Sktlim@umich.edu        self.mail_begin = False
512SN/A        self.mail_end = False
522856Srdreslin@umich.edu        self.name = ''
532SN/A        self.stdout = ''
544182Sgblack@eecs.umich.edu        self.priority = 0
554182Sgblack@eecs.umich.edu        self.queue = ''
564182Sgblack@eecs.umich.edu        self.pbshost = ''
574182Sgblack@eecs.umich.edu        self.qsub = 'qsub'
584182Sgblack@eecs.umich.edu        self.env = {}
592356SN/A        self.onlyecho = False
602356SN/A        self.verbose = False
612356SN/A
622356SN/A    def do(self, script, ):
632356SN/A        args = [self.qsub]
642356SN/A
652356SN/A        if self.env:
662356SN/A            arg = '-v'
673126Sktlim@umich.edu            arg += ','.join([ '%s=%s' % i for i in self.env.iteritems() ])
682356SN/A            args.append(arg)
692356SN/A
702356SN/A        if self.hold:
712356SN/A            args.append('-h')
722356SN/A
732356SN/A        if len(self.stdout):
742856Srdreslin@umich.edu            args.append('-olocalhost:' + self.stdout)
752SN/A
761634SN/A        if self.keep_stdout and self.keep_stderr:
771634SN/A            args.append('-koe')
781695SN/A        elif self.keep_stdout:
793814Ssaidi@eecs.umich.edu            args.append('-ko')
803814Ssaidi@eecs.umich.edu        elif self.keep_stderr:
811634SN/A            args.append('-ke')
821634SN/A        else:
832359SN/A            args.append('-kn')
841695SN/A
851695SN/A        if self.join:
861695SN/A            args.append('-joe')
873814Ssaidi@eecs.umich.edu
883814Ssaidi@eecs.umich.edu        if len(self.node_type):
891634SN/A            args.append('-lnodes=' + self.node_type)
903495Sktlim@umich.edu
913495Sktlim@umich.edu        if self.mail_abort or self.mail_begin or self.mail_end:
923495Sktlim@umich.edu            flags = ''
933495Sktlim@umich.edu            if self.mail_abort:
943495Sktlim@umich.edu                flags.append('a')
953495Sktlim@umich.edu            if self.mail_begin:
963495Sktlim@umich.edu                flags.append('b')
973495Sktlim@umich.edu            if self.mail_end:
983495Sktlim@umich.edu                flags.append('e')
993495Sktlim@umich.edu            if len(flags):
1003495Sktlim@umich.edu                args.append('-m ' + flags)
1013495Sktlim@umich.edu
1023495Sktlim@umich.edu        if len(self.name):
1033495Sktlim@umich.edu            args.append("-N%s" % self.name)
1041858SN/A
1052SN/A        if self.priority != 0:
1063520Sgblack@eecs.umich.edu            args.append('-p' + self.priority)
1073520Sgblack@eecs.umich.edu
1083520Sgblack@eecs.umich.edu        if len(self.queue):
1092SN/A            args.append('-q' + self.queue)
1102SN/A
1112SN/A        args.append(script)
1122SN/A
1132SN/A        if self.verbose or self.onlyecho:
1144103Ssaidi@eecs.umich.edu            print >>sys.stderr, 'PBS Command:   ', ' '.join(args)
1152SN/A
1163521Sgblack@eecs.umich.edu        if self.onlyecho:
1173521Sgblack@eecs.umich.edu            return 0
1181917SN/A
1191917SN/A        print >>sys.stderr, 'PBS Jobid:      ',
1201917SN/A
1211917SN/A        ec = os.spawnvp(os.P_WAIT, args[0], args)
1221917SN/A
1231917SN/A        if ec != 0 and len(self.pbshost):
1241917SN/A            ec = ssh(self.pbshost, ' '.join(args))
1251917SN/A
1261917SN/A        return ec
1271917SN/A