blob: 1d4de3ffebf20c2e88193e786a2c3ab07ec8c2ee [file] [log] [blame]
mblighb0fab822007-07-25 16:40:19 +00001__author__ = """Copyright Andy Whitcroft, Martin J. Bligh - 2006, 2007"""
2
3import sys, os, subprocess
4
5
6def parallel(tasklist):
7 """Run an set of predefined subcommands in parallel"""
8 pids = []
9 error = False
10 for task in tasklist:
11 task.fork_start()
12 for task in tasklist:
13 status = task.fork_waitfor()
14 if status != 0:
15 error = True
16 if error:
17 raise "One or more subcommands failed"
18
19
20def parallel_simple(function, arglist):
21 """Each element in the arglist used to create a subcommand object,
22 where that arg is used both as a subdir name, and a single argument
23 to pass to "function".
24 We create a subcommand object for each element in the list,
25 then execute those subcommand objects in parallel."""
26 subcommands = []
27 for arg in arglist:
mbligh0c9e7822007-07-25 22:47:51 +000028 args = [arg]
29 subdir = str(arg)
30 subcommands.append(subcommand(function, args, subdir))
mblighb0fab822007-07-25 16:40:19 +000031 parallel(subcommands)
32
33
mblighcee25b12007-08-31 08:53:05 +000034def _where_art_thy_filehandles():
mblighb0fab822007-07-25 16:40:19 +000035 os.system("ls -l /proc/%d/fd >> /dev/tty" % os.getpid())
36
mblighdc735a22007-08-02 16:54:37 +000037
mblighcee25b12007-08-31 08:53:05 +000038def _print_to_tty(string):
mblighb0fab822007-07-25 16:40:19 +000039 open('/dev/tty', 'w').write(string + '\n')
40
41
mblighcee25b12007-08-31 08:53:05 +000042def _redirect_stream(fd, output):
mblighb0fab822007-07-25 16:40:19 +000043 newfd = os.open(output, os.O_WRONLY | os.O_CREAT)
44 os.dup2(newfd, fd)
45 os.close(newfd)
46 if fd == 1:
47 sys.stdout = os.fdopen(fd, 'w')
48 if fd == 2:
49 sys.stderr = os.fdopen(fd, 'w')
50
51
mblighcee25b12007-08-31 08:53:05 +000052def _redirect_stream_tee(fd, output, tag):
mblighb0fab822007-07-25 16:40:19 +000053 """Use the low-level fork & pipe operations here to get a fd,
54 not a filehandle. This ensures that we get both the
55 filehandle and fd for stdout/stderr redirected correctly."""
56 r, w = os.pipe()
57 pid = os.fork()
58 if pid: # Parent
59 os.dup2(w, fd)
60 os.close(r)
61 os.close(w)
62 if fd == 1:
63 sys.stdout = os.fdopen(fd, 'w', 1)
64 if fd == 2:
65 sys.stderr = os.fdopen(fd, 'w', 1)
66 return
67 else: # Child
68 os.close(w)
69 log = open(output, 'w')
70 f = os.fdopen(r, 'r')
71 for line in iter(f.readline, ''):
72 # Tee straight to file
73 log.write(line)
74 log.flush()
75 # Prepend stdout with the tag
76 print tag + ' : ' + line,
77 sys.stdout.flush()
78 log.close()
79 os._exit(0)
80
81
82class subcommand:
mblighd7685d32007-08-10 22:08:42 +000083 def __init__(self, func, args, subdir = None, stdprint = True):
mblighb0fab822007-07-25 16:40:19 +000084 # func(args) - the subcommand to run
85 # subdir - the subdirectory to log results in
mblighd7685d32007-08-10 22:08:42 +000086 # stdprint - whether to print results to stdout/stderr
87 if subdir:
88 self.subdir = os.path.abspath(subdir)
89 if os.path.exists(self.subdir):
90 os.system("rm -rf %s" % self.subdir)
91 os.mkdir(self.subdir)
92 self.debug = os.path.join(self.subdir, 'debug')
mblighc526da32007-08-28 10:07:27 +000093 os.mkdir(self.debug)
mblighd7685d32007-08-10 22:08:42 +000094 self.stdout = os.path.join(self.debug, 'stdout')
95 self.stderr = os.path.join(self.debug, 'stderr')
96 else:
97 self.subdir = None
98 self.debug = '/dev/null'
99 self.stdout = '/dev/null'
100 self.stderr = '/dev/null'
101
mblighb0fab822007-07-25 16:40:19 +0000102 self.func = func
103 self.args = args
104 self.lambda_function = lambda: func(*args)
105 self.pid = None
mblighd7685d32007-08-10 22:08:42 +0000106 self.stdprint = stdprint
mblighb0fab822007-07-25 16:40:19 +0000107
108
109 def redirect_output(self):
mblighd7685d32007-08-10 22:08:42 +0000110 if self.stdprint:
111 if self.subdir:
112 tag = os.path.basename(self.subdir)
mblighcee25b12007-08-31 08:53:05 +0000113 _redirect_stream_tee(1, self.stdout, tag)
114 _redirect_stream_tee(2, self.stderr, tag)
mblighb491d022007-08-09 23:04:56 +0000115 else:
mblighcee25b12007-08-31 08:53:05 +0000116 _redirect_stream(1, self.stdout)
117 _redirect_stream(2, self.stderr)
mblighb0fab822007-07-25 16:40:19 +0000118
119
120 def fork_start(self):
121 sys.stdout.flush()
122 sys.stderr.flush()
123 self.pid = os.fork()
124
125 if self.pid: # I am the parent
126 return
127
128 # We are the child from this point on. Never return.
mblighd7685d32007-08-10 22:08:42 +0000129 if self.subdir:
130 os.chdir(self.subdir)
mblighb0fab822007-07-25 16:40:19 +0000131 self.redirect_output()
132
133 try:
134 self.lambda_function()
135
136 except:
137 raise
138 sys.stdout.flush()
139 sys.stderr.flush()
140 os._exit(1)
141
142 sys.stdout.flush()
143 sys.stderr.flush()
144 os._exit(0)
145
146
147 def fork_waitfor(self):
148 (pid, status) = os.waitpid(self.pid, 0)
149
150 if status != 0:
151 print "subcommand failed pid %d" % pid
152 print "%s(%s)" % (self.func, self.args)
153 print "rc=%d" % status
154 print
155 if os.path.exists(self.stderr):
156 for line in open(self.stderr).readlines():
157 print line,
158 print "\n--------------------------------------------\n"
159 return status