Package starcluster :: Module threadpool
[hide private]
[frames] | no frames]

Source Code for Module starcluster.threadpool

  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 
13 14 15 -class DaemonWorker(workerpool.workers.Worker):
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 """
21 - def __init__(self, *args, **kwargs):
22 super(DaemonWorker, self).__init__(*args, **kwargs) 23 self.daemon = True
24
25 - def run(self):
26 "Get jobs from the queue and perform them as they arrive." 27 while 1: 28 # Sleep until there is a job to perform. 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
41 42 -def _worker_factory(parent):
43 return DaemonWorker(parent)
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
55 - def run(self):
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
72 73 -class ThreadPool(workerpool.WorkerPool):
74 - def __init__(self, size=1, maxjobs=0, worker_factory=_worker_factory, 75 disable_threads=False):
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
85 - def progress_bar(self):
86 if not self._progress_bar: 87 widgets = ['', progressbar.Fraction(), ' ', 88 progressbar.Bar(marker=progressbar.RotatingMarker()), 89 ' ', progressbar.Percentage(), ' ', ' '] 90 pbar = progressbar.ProgressBar(widgets=widgets, maxval=1, 91 force_update=True) 92 self._progress_bar = pbar 93 return self._progress_bar
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
105 - def get_results(self):
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
119 - def store_exception(self, e):
120 self._exception_queue.put(e)
121
122 - def shutdown(self):
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
146 - def __del__(self):
147 log.debug('del called in threadpool') 148 self.shutdown() 149 self.join()
150
151 152 -def get_thread_pool(size=10, worker_factory=_worker_factory, 153 disable_threads=False):
154 return ThreadPool(size=size, worker_factory=_worker_factory, 155 disable_threads=disable_threads)
156