pbs.py revision 1385
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): 331385Sbinkertn@umich.edu self.sts = -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: 671385Sbinkertn@umich.edu os.dup2(p2c_read, 0) 681385Sbinkertn@umich.edu os.dup2(c2p_write, 1) 691385Sbinkertn@umich.edu os.dup2(c2p_write, 2) 701385Sbinkertn@umich.edu if isinstance(cmd, basestring): 711385Sbinkertn@umich.edu cmd = ['/bin/sh', '-c', cmd] 721385Sbinkertn@umich.edu if False: 731385Sbinkertn@umich.edu for i in range(3, MAXFD): 741385Sbinkertn@umich.edu try: 751385Sbinkertn@umich.edu os.close(i) 761385Sbinkertn@umich.edu except OSError: 771385Sbinkertn@umich.edu pass 781385Sbinkertn@umich.edu try: 791385Sbinkertn@umich.edu os.execvp(cmd[0], cmd) 801385Sbinkertn@umich.edu finally: 811385Sbinkertn@umich.edu os._exit(1) 821385Sbinkertn@umich.edu 831385Sbinkertn@umich.edu os.close(p2c_read) 841385Sbinkertn@umich.edu os.close(c2p_write) 851385Sbinkertn@umich.edu 861385Sbinkertn@umich.edu def poll(self): 871385Sbinkertn@umich.edu if self.sts < 0: 881385Sbinkertn@umich.edu pid, sts = os.waitpid(self.pid, os.WNOHANG) 891385Sbinkertn@umich.edu if pid == self.pid: 901385Sbinkertn@umich.edu self.sts = sts 911385Sbinkertn@umich.edu return self.sts 921385Sbinkertn@umich.edu 931385Sbinkertn@umich.edu def wait(self): 941385Sbinkertn@umich.edu if self.sts < 0: 951385Sbinkertn@umich.edu pid, sts = os.waitpid(self.pid, 0) 961385Sbinkertn@umich.edu if pid == self.pid: 971385Sbinkertn@umich.edu self.sts = sts 981385Sbinkertn@umich.edu return self.sts 991376Sbinkertn@umich.edu 1001376Sbinkertn@umich.educlass qsub: 1011376Sbinkertn@umich.edu def __init__(self): 1021376Sbinkertn@umich.edu self.hold = False 1031376Sbinkertn@umich.edu self.join = False 1041376Sbinkertn@umich.edu self.keep_stdout = False 1051376Sbinkertn@umich.edu self.keep_stderr = False 1061376Sbinkertn@umich.edu self.node_type = '' 1071376Sbinkertn@umich.edu self.mail_abort = False 1081376Sbinkertn@umich.edu self.mail_begin = False 1091376Sbinkertn@umich.edu self.mail_end = False 1101376Sbinkertn@umich.edu self.name = '' 1111376Sbinkertn@umich.edu self.stdout = '' 1121376Sbinkertn@umich.edu self.priority = 0 1131376Sbinkertn@umich.edu self.queue = '' 1141376Sbinkertn@umich.edu self.pbshost = '' 1151376Sbinkertn@umich.edu self.qsub = 'qsub' 1161376Sbinkertn@umich.edu self.env = {} 1171376Sbinkertn@umich.edu 1181385Sbinkertn@umich.edu def build(self, script, args = []): 1191385Sbinkertn@umich.edu self.cmd = [ self.qsub ] 1201376Sbinkertn@umich.edu 1211376Sbinkertn@umich.edu if self.env: 1221376Sbinkertn@umich.edu arg = '-v' 1231376Sbinkertn@umich.edu arg += ','.join([ '%s=%s' % i for i in self.env.iteritems() ]) 1241385Sbinkertn@umich.edu self.cmd.append(arg) 1251376Sbinkertn@umich.edu 1261376Sbinkertn@umich.edu if self.hold: 1271385Sbinkertn@umich.edu self.cmd.append('-h') 1281376Sbinkertn@umich.edu 1291376Sbinkertn@umich.edu if len(self.stdout): 1301385Sbinkertn@umich.edu self.cmd.append('-olocalhost:' + self.stdout) 1311376Sbinkertn@umich.edu 1321376Sbinkertn@umich.edu if self.keep_stdout and self.keep_stderr: 1331385Sbinkertn@umich.edu self.cmd.append('-koe') 1341376Sbinkertn@umich.edu elif self.keep_stdout: 1351385Sbinkertn@umich.edu self.cmd.append('-ko') 1361376Sbinkertn@umich.edu elif self.keep_stderr: 1371385Sbinkertn@umich.edu self.cmd.append('-ke') 1381376Sbinkertn@umich.edu else: 1391385Sbinkertn@umich.edu self.cmd.append('-kn') 1401376Sbinkertn@umich.edu 1411376Sbinkertn@umich.edu if self.join: 1421385Sbinkertn@umich.edu self.cmd.append('-joe') 1431376Sbinkertn@umich.edu 1441376Sbinkertn@umich.edu if len(self.node_type): 1451385Sbinkertn@umich.edu self.cmd.append('-lnodes=' + self.node_type) 1461376Sbinkertn@umich.edu 1471376Sbinkertn@umich.edu if self.mail_abort or self.mail_begin or self.mail_end: 1481376Sbinkertn@umich.edu flags = '' 1491376Sbinkertn@umich.edu if self.mail_abort: 1501376Sbinkertn@umich.edu flags.append('a') 1511376Sbinkertn@umich.edu if self.mail_begin: 1521376Sbinkertn@umich.edu flags.append('b') 1531376Sbinkertn@umich.edu if self.mail_end: 1541376Sbinkertn@umich.edu flags.append('e') 1551376Sbinkertn@umich.edu if len(flags): 1561385Sbinkertn@umich.edu self.cmd.append('-m ' + flags) 1571376Sbinkertn@umich.edu 1581376Sbinkertn@umich.edu if len(self.name): 1591385Sbinkertn@umich.edu self.cmd.append("-N%s" % self.name) 1601376Sbinkertn@umich.edu 1611376Sbinkertn@umich.edu if self.priority != 0: 1621385Sbinkertn@umich.edu self.cmd.append('-p' + self.priority) 1631376Sbinkertn@umich.edu 1641376Sbinkertn@umich.edu if len(self.queue): 1651385Sbinkertn@umich.edu self.cmd.append('-q' + self.queue) 1661376Sbinkertn@umich.edu 1671385Sbinkertn@umich.edu self.cmd.extend(args) 1681385Sbinkertn@umich.edu self.script = script 1691385Sbinkertn@umich.edu self.command = ' '.join(self.cmd + [ self.script ]) 1701376Sbinkertn@umich.edu 1711385Sbinkertn@umich.edu def do(self): 1721385Sbinkertn@umich.edu pbs = MyPOpen(self.cmd + [ self.script ]) 1731385Sbinkertn@umich.edu self.result = pbs.fromchild.read() 1741385Sbinkertn@umich.edu ec = pbs.wait() 1751376Sbinkertn@umich.edu 1761385Sbinkertn@umich.edu if ec != 0 and self.pbshost: 1771385Sbinkertn@umich.edu cmd = ' '.join(self.cmd + [ '-' ]) 1781385Sbinkertn@umich.edu cmd = [ 'ssh', '-x', self.pbshost, cmd ] 1791385Sbinkertn@umich.edu self.command = ' '.join(cmd) 1801385Sbinkertn@umich.edu ssh = MyPOpen(cmd, input = self.script) 1811385Sbinkertn@umich.edu self.result = ssh.fromchild.read() 1821385Sbinkertn@umich.edu ec = ssh.wait() 1831376Sbinkertn@umich.edu 1841376Sbinkertn@umich.edu return ec 185