Allow overriding the queued commands environment.
[htsworkflow.git] / htsworkflow / util / queuecommands.py
1 """
2 Run up to N simultanous jobs from provided of commands 
3 """
4
5 import logging
6 import os
7 from subprocess import PIPE
8 import subprocess
9 import select
10 import sys
11 import time
12
13 class QueueCommands(object):
14     """
15     Queue up N commands from cmd_list, launching more jobs as the first
16     finish.
17     """
18
19     def __init__(self, cmd_list, N=0, cwd=None, env=None):
20         """
21         cmd_list is a list of elements suitable for subprocess
22         N is the  number of simultanious processes to run. 
23         0 is all of them.
24         
25         WARNING: this will not work on windows
26         (It depends on being able to pass local file descriptors to the 
27         select call with isn't supported by the Win32 API)
28         """
29         self.to_run = cmd_list[:]
30         self.running = {}
31         self.N = N
32         self.cwd = cwd
33         self.env = env
34
35     def under_process_limit(self):
36         """
37         are we still under the total number of allowable jobs?
38         """
39         if self.N == 0:
40             return True
41
42         if len(self.running) < self.N:
43             return True
44
45         return False
46
47     def start_jobs(self):
48         """
49         Launch jobs until we have the maximum allowable running
50         (or have run out of jobs)
51         """
52         queue_log = logging.getLogger('queue')
53
54         while (len(self.to_run) > 0) and self.under_process_limit():
55             queue_log.info('%d left to run', len(self.to_run))
56             cmd = self.to_run.pop(0)
57             p = subprocess.Popen(cmd, 
58                                  stdout=PIPE, 
59                                  shell=True, 
60                                  cwd=self.cwd, 
61                                  env=self.env)
62             self.running[p.stdout] = p
63             queue_log.info("Created process %d from %s" % (p.pid, str(cmd)))
64
65     def run(self):
66         """
67         run up to N jobs until we run out of jobs
68         """
69         queue_log = logging.getLogger('queue')
70         queue_log.debug('using %s as cwd' % (self.cwd,))
71
72         # to_run slowly gets consumed by start_jobs
73         while len(self.to_run) > 0 or len(self.running) > 0:
74             # fill any empty spots in our job queue
75             self.start_jobs()
76
77             # build a list of file descriptors
78             # fds=file desciptors
79             fds = [ x.stdout for x in self.running.values()]
80
81             # wait for something to finish
82             # wl= write list, xl=exception list (not used so get bad names)
83             read_list, wl, xl = select.select(fds, [], fds, 1 )
84
85             # for everything that might have finished...
86             for pending_fd in read_list:
87                 pending = self.running[pending_fd]
88                 # if it really did finish, remove it from running jobs
89                 if pending.poll() is not None:
90                     queue_log.info("Process %d finished [%d]",
91                                    pending.pid, pending.returncode)
92                     del self.running[pending_fd]
93                 else:
94                     # It's still running, but there's some output
95                     buffer = pending_fd.readline()
96                     buffer = buffer.strip()
97                     msg = "%d:(%d) %s" %(pending.pid, len(buffer), buffer)
98                     logging.debug(msg)
99             time.sleep(1)