pbs.py revision 1928
11376Sbinkertn@umich.edu# Copyright (c) 2005 The Regents of The University of Michigan 21376Sbinkertn@umich.edu# All rights reserved. 31376Sbinkertn@umich.edu# 41376Sbinkertn@umich.edu# Redistribution and use in source and binary forms, with or without 51376Sbinkertn@umich.edu# modification, are permitted provided that the following conditions are 61376Sbinkertn@umich.edu# met: redistributions of source code must retain the above copyright 71376Sbinkertn@umich.edu# notice, this list of conditions and the following disclaimer; 81376Sbinkertn@umich.edu# redistributions in binary form must reproduce the above copyright 91376Sbinkertn@umich.edu# notice, this list of conditions and the following disclaimer in the 101376Sbinkertn@umich.edu# documentation and/or other materials provided with the distribution; 111376Sbinkertn@umich.edu# neither the name of the copyright holders nor the names of its 121376Sbinkertn@umich.edu# contributors may be used to endorse or promote products derived from 131376Sbinkertn@umich.edu# this software without specific prior written permission. 141376Sbinkertn@umich.edu# 151376Sbinkertn@umich.edu# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 161376Sbinkertn@umich.edu# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 171376Sbinkertn@umich.edu# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 181376Sbinkertn@umich.edu# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 191376Sbinkertn@umich.edu# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 201376Sbinkertn@umich.edu# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 211376Sbinkertn@umich.edu# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 221376Sbinkertn@umich.edu# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 231376Sbinkertn@umich.edu# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 241376Sbinkertn@umich.edu# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 251376Sbinkertn@umich.edu# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 261376Sbinkertn@umich.edu# 271376Sbinkertn@umich.edu# Authors: Nathan Binkert 281376Sbinkertn@umich.edu 291385Sbinkertn@umich.eduimport os, popen2, re, sys 301376Sbinkertn@umich.edu 311385Sbinkertn@umich.educlass MyPOpen(object): 321385Sbinkertn@umich.edu def __init__(self, cmd, input = None, output = None, bufsize = -1): 331386Sbinkertn@umich.edu self.status = -1 341376Sbinkertn@umich.edu 351385Sbinkertn@umich.edu if input is None: 361385Sbinkertn@umich.edu p2c_read, p2c_write = os.pipe() 371385Sbinkertn@umich.edu self.tochild = os.fdopen(p2c_write, 'w', bufsize) 381385Sbinkertn@umich.edu else: 391385Sbinkertn@umich.edu p2c_write = None 401385Sbinkertn@umich.edu if isinstance(input, file): 411385Sbinkertn@umich.edu p2c_read = input.fileno() 421385Sbinkertn@umich.edu elif isinstance(input, str): 431385Sbinkertn@umich.edu input = file(input, 'r') 441385Sbinkertn@umich.edu p2c_read = input.fileno() 451385Sbinkertn@umich.edu elif isinstance(input, int): 461385Sbinkertn@umich.edu p2c_read = input 471385Sbinkertn@umich.edu else: 481385Sbinkertn@umich.edu raise AttributeError 491385Sbinkertn@umich.edu 501385Sbinkertn@umich.edu if output is None: 511385Sbinkertn@umich.edu c2p_read, c2p_write = os.pipe() 521385Sbinkertn@umich.edu self.fromchild = os.fdopen(c2p_read, 'r', bufsize) 531385Sbinkertn@umich.edu else: 541385Sbinkertn@umich.edu c2p_read = None 551385Sbinkertn@umich.edu if isinstance(output, file): 561385Sbinkertn@umich.edu c2p_write = output.fileno() 571385Sbinkertn@umich.edu elif isinstance(output, str): 581385Sbinkertn@umich.edu output = file(output, 'w') 591385Sbinkertn@umich.edu c2p_write = output.fileno() 601385Sbinkertn@umich.edu elif isinstance(output, int): 611385Sbinkertn@umich.edu c2p_write = output 621385Sbinkertn@umich.edu else: 631385Sbinkertn@umich.edu raise AttributeError 641385Sbinkertn@umich.edu 651385Sbinkertn@umich.edu self.pid = os.fork() 661385Sbinkertn@umich.edu if self.pid == 0: 671386Sbinkertn@umich.edu os.dup2(p2c_read, sys.stdin.fileno()) 681386Sbinkertn@umich.edu os.dup2(c2p_write, sys.stdout.fileno()) 691386Sbinkertn@umich.edu os.dup2(c2p_write, sys.stderr.fileno()) 701385Sbinkertn@umich.edu try: 711385Sbinkertn@umich.edu os.execvp(cmd[0], cmd) 721385Sbinkertn@umich.edu finally: 731385Sbinkertn@umich.edu os._exit(1) 741385Sbinkertn@umich.edu 751385Sbinkertn@umich.edu os.close(p2c_read) 761385Sbinkertn@umich.edu os.close(c2p_write) 771385Sbinkertn@umich.edu 781385Sbinkertn@umich.edu def poll(self): 791386Sbinkertn@umich.edu if self.status < 0: 801386Sbinkertn@umich.edu pid, status = os.waitpid(self.pid, os.WNOHANG) 811385Sbinkertn@umich.edu if pid == self.pid: 821386Sbinkertn@umich.edu self.status = status 831386Sbinkertn@umich.edu return self.status 841385Sbinkertn@umich.edu 851385Sbinkertn@umich.edu def wait(self): 861386Sbinkertn@umich.edu if self.status < 0: 871386Sbinkertn@umich.edu pid, status = os.waitpid(self.pid, 0) 881385Sbinkertn@umich.edu if pid == self.pid: 891386Sbinkertn@umich.edu self.status = status 901386Sbinkertn@umich.edu return self.status 911376Sbinkertn@umich.edu 921376Sbinkertn@umich.educlass qsub: 931376Sbinkertn@umich.edu def __init__(self): 941916Sbinkertn@umich.edu self.afterok = None 951376Sbinkertn@umich.edu self.hold = False 961376Sbinkertn@umich.edu self.join = False 971376Sbinkertn@umich.edu self.keep_stdout = False 981376Sbinkertn@umich.edu self.keep_stderr = False 991916Sbinkertn@umich.edu self.node_type = None 1001376Sbinkertn@umich.edu self.mail_abort = False 1011376Sbinkertn@umich.edu self.mail_begin = False 1021376Sbinkertn@umich.edu self.mail_end = False 1031916Sbinkertn@umich.edu self.name = None 1041916Sbinkertn@umich.edu self.stdout = None 1051916Sbinkertn@umich.edu self.priority = None 1061916Sbinkertn@umich.edu self.queue = None 1071916Sbinkertn@umich.edu self.pbshost = None 1081376Sbinkertn@umich.edu self.qsub = 'qsub' 1091376Sbinkertn@umich.edu self.env = {} 1101376Sbinkertn@umich.edu 1111385Sbinkertn@umich.edu def build(self, script, args = []): 1121385Sbinkertn@umich.edu self.cmd = [ self.qsub ] 1131376Sbinkertn@umich.edu 1141376Sbinkertn@umich.edu if self.env: 1151376Sbinkertn@umich.edu arg = '-v' 1161376Sbinkertn@umich.edu arg += ','.join([ '%s=%s' % i for i in self.env.iteritems() ]) 1171385Sbinkertn@umich.edu self.cmd.append(arg) 1181376Sbinkertn@umich.edu 1191376Sbinkertn@umich.edu if self.hold: 1201385Sbinkertn@umich.edu self.cmd.append('-h') 1211376Sbinkertn@umich.edu 1221916Sbinkertn@umich.edu if self.stdout: 1231385Sbinkertn@umich.edu self.cmd.append('-olocalhost:' + self.stdout) 1241376Sbinkertn@umich.edu 1251376Sbinkertn@umich.edu if self.keep_stdout and self.keep_stderr: 1261385Sbinkertn@umich.edu self.cmd.append('-koe') 1271376Sbinkertn@umich.edu elif self.keep_stdout: 1281385Sbinkertn@umich.edu self.cmd.append('-ko') 1291376Sbinkertn@umich.edu elif self.keep_stderr: 1301385Sbinkertn@umich.edu self.cmd.append('-ke') 1311376Sbinkertn@umich.edu else: 1321385Sbinkertn@umich.edu self.cmd.append('-kn') 1331376Sbinkertn@umich.edu 1341376Sbinkertn@umich.edu if self.join: 1351385Sbinkertn@umich.edu self.cmd.append('-joe') 1361376Sbinkertn@umich.edu 1371916Sbinkertn@umich.edu if self.node_type: 1381385Sbinkertn@umich.edu self.cmd.append('-lnodes=' + self.node_type) 1391376Sbinkertn@umich.edu 1401376Sbinkertn@umich.edu if self.mail_abort or self.mail_begin or self.mail_end: 1411376Sbinkertn@umich.edu flags = '' 1421376Sbinkertn@umich.edu if self.mail_abort: 1431376Sbinkertn@umich.edu flags.append('a') 1441376Sbinkertn@umich.edu if self.mail_begin: 1451376Sbinkertn@umich.edu flags.append('b') 1461376Sbinkertn@umich.edu if self.mail_end: 1471376Sbinkertn@umich.edu flags.append('e') 1481376Sbinkertn@umich.edu if len(flags): 1491385Sbinkertn@umich.edu self.cmd.append('-m ' + flags) 1501376Sbinkertn@umich.edu 1511916Sbinkertn@umich.edu if self.name: 1521385Sbinkertn@umich.edu self.cmd.append("-N%s" % self.name) 1531376Sbinkertn@umich.edu 1541916Sbinkertn@umich.edu if self.priority: 1551385Sbinkertn@umich.edu self.cmd.append('-p' + self.priority) 1561376Sbinkertn@umich.edu 1571916Sbinkertn@umich.edu if self.queue: 1581385Sbinkertn@umich.edu self.cmd.append('-q' + self.queue) 1591376Sbinkertn@umich.edu 1601916Sbinkertn@umich.edu if self.afterok: 1611928Sbinkertn@umich.edu self.cmd.append('-Wdepend=afterok:%s' % self.afterok) 1621916Sbinkertn@umich.edu 1631385Sbinkertn@umich.edu self.cmd.extend(args) 1641385Sbinkertn@umich.edu self.script = script 1651385Sbinkertn@umich.edu self.command = ' '.join(self.cmd + [ self.script ]) 1661376Sbinkertn@umich.edu 1671385Sbinkertn@umich.edu def do(self): 1681385Sbinkertn@umich.edu pbs = MyPOpen(self.cmd + [ self.script ]) 1691385Sbinkertn@umich.edu self.result = pbs.fromchild.read() 1701385Sbinkertn@umich.edu ec = pbs.wait() 1711376Sbinkertn@umich.edu 1721385Sbinkertn@umich.edu if ec != 0 and self.pbshost: 1731385Sbinkertn@umich.edu cmd = ' '.join(self.cmd + [ '-' ]) 1741385Sbinkertn@umich.edu cmd = [ 'ssh', '-x', self.pbshost, cmd ] 1751385Sbinkertn@umich.edu self.command = ' '.join(cmd) 1761385Sbinkertn@umich.edu ssh = MyPOpen(cmd, input = self.script) 1771385Sbinkertn@umich.edu self.result = ssh.fromchild.read() 1781385Sbinkertn@umich.edu ec = ssh.wait() 1791376Sbinkertn@umich.edu 1801376Sbinkertn@umich.edu return ec 181