More fixes to the pbs stuff to make it more robust.

sim/pyconfig/SConscript:
    Embed the jobfile.py script into the binary so that we don't
    need to copy it into the Base directory every time.
test/genini.py:
    Add the util/pbs directory to the path so we can get to
    jobfile.py
    Add a -I argument to set to add to the path.
util/pbs/pbs.py:
    Create a MyPOpen class.  This is a lot like the popen2.Popen3 class
    in the python library except that my version allows redirection of
    standard in and standard out to a file instead of a pipe.
    Use this popen class to execute qsub or ssh qsub.  This was important
    for the ssh version of qsub because we need to pipe the script into
    standard in of ssh so that the script can get to the qsub command.
    (Otherwise we have a problem discovering the path.)
util/pbs/send.py:
    Tweak the script so it figures out paths in NFS correctly.
    Use the new system for running qsub.

--HG--
extra : convert_revision : 1289915ba99cec6fd464b71215c32d2197ff2824

1 files changed, 100 insertions, 42 deletions

diff --git a/util/pbs/pbs.py b/util/pbs/pbs.py
index a71dbbf8e..cd55da722 100755
--- a/util/pbs/pbs.py
+++ b/util/pbs/pbs.py
@@ -26,18 +26,76 @@
 #
 # Authors: Nathan Binkert
 
-import os, re, sys
+import os, popen2, re, sys
 
-def ssh(host, script, tty = False, user = ''):
-    args = [ 'ssh', '-x' ]
-    if user:
-        args.append('-l' + user)
-    if tty:
-        args.append('-t')
-    args.append(host)
-    args.append(script)
+class MyPOpen(object):
+    def __init__(self, cmd, input = None, output = None, bufsize = -1):
+        self.sts = -1
 
-    return os.spawnvp(os.P_WAIT, args[0], args)
+        if input is None:
+            p2c_read, p2c_write = os.pipe()
+            self.tochild = os.fdopen(p2c_write, 'w', bufsize)
+        else:
+            p2c_write = None
+            if isinstance(input, file):
+                p2c_read = input.fileno()
+            elif isinstance(input, str):
+                input = file(input, 'r')
+                p2c_read = input.fileno()
+            elif isinstance(input, int):
+                p2c_read = input
+            else:
+                raise AttributeError
+
+        if output is None:
+            c2p_read, c2p_write = os.pipe()
+            self.fromchild = os.fdopen(c2p_read, 'r', bufsize)
+        else:
+            c2p_read = None
+            if isinstance(output, file):
+                c2p_write = output.fileno()
+            elif isinstance(output, str):
+                output = file(output, 'w')
+                c2p_write = output.fileno()
+            elif isinstance(output, int):
+                c2p_write = output
+            else:
+                raise AttributeError
+
+        self.pid = os.fork()
+        if self.pid == 0:
+            os.dup2(p2c_read, 0)
+            os.dup2(c2p_write, 1)
+            os.dup2(c2p_write, 2)
+            if isinstance(cmd, basestring):
+                cmd = ['/bin/sh', '-c', cmd]
+            if False:
+                for i in range(3, MAXFD):
+                    try:
+                        os.close(i)
+                    except OSError:
+                        pass
+            try:
+                os.execvp(cmd[0], cmd)
+            finally:
+                os._exit(1)
+
+        os.close(p2c_read)
+        os.close(c2p_write)
+
+    def poll(self):
+        if self.sts < 0:
+            pid, sts = os.waitpid(self.pid, os.WNOHANG)
+            if pid == self.pid:
+                self.sts = sts
+        return self.sts
+
+    def wait(self):
+        if self.sts < 0:
+            pid, sts = os.waitpid(self.pid, 0)
+            if pid == self.pid:
+                self.sts = sts
+        return self.sts
 
 class qsub:
     def __init__(self):
@@ -56,37 +114,35 @@ class qsub:
         self.pbshost = ''
         self.qsub = 'qsub'
         self.env = {}
-        self.onlyecho = False
-        self.verbose = False
 
-    def do(self, script, ):
-        args = [self.qsub]
+    def build(self, script, args = []):
+        self.cmd = [ self.qsub ]
 
         if self.env:
             arg = '-v'
             arg += ','.join([ '%s=%s' % i for i in self.env.iteritems() ])
-            args.append(arg)
+            self.cmd.append(arg)
 
         if self.hold:
-            args.append('-h')
+            self.cmd.append('-h')
 
         if len(self.stdout):
-            args.append('-olocalhost:' + self.stdout)
+            self.cmd.append('-olocalhost:' + self.stdout)
 
         if self.keep_stdout and self.keep_stderr:
-            args.append('-koe')
+            self.cmd.append('-koe')
         elif self.keep_stdout:
-            args.append('-ko')
+            self.cmd.append('-ko')
         elif self.keep_stderr:
-            args.append('-ke')
+            self.cmd.append('-ke')
         else:
-            args.append('-kn')
+            self.cmd.append('-kn')
 
         if self.join:
-            args.append('-joe')
+            self.cmd.append('-joe')
 
         if len(self.node_type):
-            args.append('-lnodes=' + self.node_type)
+            self.cmd.append('-lnodes=' + self.node_type)
 
         if self.mail_abort or self.mail_begin or self.mail_end:
             flags = ''
@@ -97,30 +153,32 @@ class qsub:
             if self.mail_end:
                 flags.append('e')
             if len(flags):
-                args.append('-m ' + flags)
+                self.cmd.append('-m ' + flags)
 
         if len(self.name):
-            args.append("-N%s" % self.name)
+            self.cmd.append("-N%s" % self.name)
 
         if self.priority != 0:
-            args.append('-p' + self.priority)
+            self.cmd.append('-p' + self.priority)
 
         if len(self.queue):
-            args.append('-q' + self.queue)
-
-        args.append(script)
-
-        if self.verbose or self.onlyecho:
-            print >>sys.stderr, 'PBS Command:   ', ' '.join(args)
-
-        if self.onlyecho:
-            return 0
-
-        print >>sys.stderr, 'PBS Jobid:      ',
-
-        ec = os.spawnvp(os.P_WAIT, args[0], args)
-
-        if ec != 0 and len(self.pbshost):
-            ec = ssh(self.pbshost, ' '.join(args))
+            self.cmd.append('-q' + self.queue)
+
+        self.cmd.extend(args)
+        self.script = script
+        self.command = ' '.join(self.cmd + [ self.script ])
+
+    def do(self):
+        pbs = MyPOpen(self.cmd + [ self.script ])
+        self.result = pbs.fromchild.read()
+        ec = pbs.wait()
+
+        if ec != 0 and self.pbshost:
+            cmd = ' '.join(self.cmd + [ '-' ])
+            cmd = [ 'ssh', '-x', self.pbshost, cmd ]
+            self.command = ' '.join(cmd)
+            ssh = MyPOpen(cmd, input = self.script)
+            self.result = ssh.fromchild.read()
+            ec = ssh.wait()
 
         return ec