1 """
2 ThreadPool module for StarCluster based on WorkerPool
3 """
4 import time
5 import Queue
6 import thread
7 import traceback
8 import workerpool
9
10 from starcluster import exception
11 from starcluster import progressbar
12 from starcluster.logger import log
16 """
17 Improved Worker that sets daemon = True by default and also handles
18 communicating exceptions to the parent pool object by adding them to
19 the parent pool's exception queue
20 """
24
26 "Get jobs from the queue and perform them as they arrive."
27 while 1:
28
29 job = self.jobs.get()
30 try:
31 job.run()
32 except workerpool.exceptions.TerminationNotice:
33 break
34 except Exception, e:
35 tb_msg = traceback.format_exc()
36 jid = job.jobid or str(thread.get_ident())
37 self.jobs.store_exception([e, tb_msg, jid])
38 finally:
39 self.jobs.task_done()
40
44
45
46 -class SimpleJob(workerpool.jobs.SimpleJob):
47 - def __init__(self, method, args=[], kwargs={}, jobid=None,
48 results_queue=None):
49 self.method = method
50 self.args = args
51 self.kwargs = kwargs
52 self.jobid = jobid
53 self.results_queue = results_queue
54
56 if isinstance(self.args, list) or isinstance(self.args, tuple):
57 if isinstance(self.kwargs, dict):
58 r = self.method(*self.args, **self.kwargs)
59 else:
60 r = self.method(*self.args)
61 elif self.args is not None and self.args is not []:
62 if isinstance(self.kwargs, dict):
63 r = self.method(self.args, **self.kwargs)
64 else:
65 r = self.method(self.args)
66 else:
67 r = self.method()
68 if self.results_queue:
69 return self.results_queue.put(r)
70 return r
71
76 self.disable_threads = disable_threads
77 self._exception_queue = Queue.Queue()
78 self._results_queue = Queue.Queue()
79 self._progress_bar = None
80 if self.disable_threads:
81 size = 0
82 workerpool.WorkerPool.__init__(self, size, maxjobs, worker_factory)
83
84 @property
94
95 - def simple_job(self, method, args=[], kwargs={}, jobid=None,
96 results_queue=None):
97 results_queue = results_queue or self._results_queue
98 job = SimpleJob(method, args, kwargs, jobid,
99 results_queue=results_queue)
100 if not self.disable_threads:
101 return self.put(job)
102 else:
103 return job.run()
104
106 results = []
107 for i in range(self._results_queue.qsize()):
108 results.append(self._results_queue.get())
109 return results
110
111 - def map(self, fn, *seq):
112 if self._results_queue.qsize() > 0:
113 self.get_results()
114 args = zip(*seq)
115 for seq in args:
116 self.simple_job(fn, seq)
117 return self.wait(numtasks=len(args))
118
120 self._exception_queue.put(e)
121
123 log.info("Shutting down threads...")
124 workerpool.WorkerPool.shutdown(self)
125 self.wait(numtasks=self.size())
126
127 - def wait(self, numtasks=None, return_results=True):
128 pbar = self.progress_bar.reset()
129 pbar.maxval = self.unfinished_tasks
130 if numtasks is not None:
131 pbar.maxval = max(numtasks, self.unfinished_tasks)
132 while self.unfinished_tasks != 0:
133 finished = pbar.maxval - self.unfinished_tasks
134 pbar.update(finished)
135 log.debug("unfinished_tasks = %d" % self.unfinished_tasks)
136 time.sleep(1)
137 if pbar.maxval != 0:
138 pbar.finish()
139 self.join()
140 if self._exception_queue.qsize() > 0:
141 raise exception.ThreadPoolException(
142 "An error occurred in ThreadPool", self._exception_queue.queue)
143 if return_results:
144 return self.get_results()
145
147 log.debug('del called in threadpool')
148 self.shutdown()
149 self.join()
150
156