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

Source Code for Module starcluster.threadpool

  1  #!/usr/bin/env python 
  2  """ 
  3  ThreadPool module for StarCluster based on WorkerPool 
  4  """ 
  5  import time 
  6  import Queue 
  7  import thread 
  8  import traceback 
  9  import workerpool 
 10   
 11  from starcluster import exception 
 12  from starcluster import progressbar 
 13  from starcluster.logger import log 
14 15 16 -class DaemonWorker(workerpool.workers.Worker):
17 """ 18 Improved Worker that sets daemon = True by default and also handles 19 communicating exceptions to the parent pool object by adding them to 20 the parent pool's exception queue 21 """
22 - def __init__(self, *args, **kwargs):
23 super(DaemonWorker, self).__init__(*args, **kwargs) 24 self.daemon = True
25
26 - def run(self):
27 "Get jobs from the queue and perform them as they arrive." 28 while 1: 29 # Sleep until there is a job to perform. 30 job = self.jobs.get() 31 try: 32 job.run() 33 except workerpool.exceptions.TerminationNotice: 34 break 35 except Exception, e: 36 tb_msg = traceback.format_exc() 37 jid = job.jobid or str(thread.get_ident()) 38 self.jobs.store_exception([e, tb_msg, jid]) 39 finally: 40 self.jobs.task_done()
41
42 43 -def _worker_factory(parent):
44 return DaemonWorker(parent)
45
46 47 -class SimpleJob(workerpool.jobs.SimpleJob):
48 - def __init__(self, method, args=[], jobid=None):
49 self.method = method 50 self.args = args 51 self.jobid = jobid
52
53 - def run(self):
54 if isinstance(self.args, list) or isinstance(self.args, tuple): 55 r = self.method(*self.args) 56 elif isinstance(self.args, dict): 57 r = self.method(**self.args) 58 else: 59 r = self.method(self.args) 60 return r
61
62 63 -class ThreadPool(workerpool.WorkerPool):
64 - def __init__(self, size=1, maxjobs=0, worker_factory=_worker_factory, 65 disable_threads=False):
66 self.disable_threads = disable_threads 67 self._exception_queue = Queue.Queue() 68 self._progress_bar = None 69 if self.disable_threads: 70 size = 0 71 workerpool.WorkerPool.__init__(self, size, maxjobs, worker_factory)
72 73 @property
74 - def progress_bar(self):
75 if not self._progress_bar: 76 widgets = ['', progressbar.Fraction(), ' ', 77 progressbar.Bar(marker=progressbar.RotatingMarker()), 78 ' ', progressbar.Percentage(), ' ', ' '] 79 pbar = progressbar.ProgressBar(widgets=widgets, 80 maxval=1, 81 force_update=True) 82 self._progress_bar = pbar 83 return self._progress_bar
84
85 - def simple_job(self, method, args=[], jobid=None):
86 job = SimpleJob(method, args, jobid) 87 if not self.disable_threads: 88 return self.put(job) 89 else: 90 return job.run()
91
92 - def store_exception(self, e):
93 self._exception_queue.put(e)
94
95 - def shutdown(self):
96 log.info("Shutting down threads...") 97 workerpool.WorkerPool.shutdown(self) 98 self.wait(numtasks=self.size())
99
100 - def wait(self, numtasks=None):
101 pbar = self.progress_bar.reset() 102 pbar.maxval = self.unfinished_tasks 103 if numtasks is not None: 104 pbar.maxval = max(numtasks, self.unfinished_tasks) 105 while self.unfinished_tasks != 0: 106 finished = pbar.maxval - self.unfinished_tasks 107 pbar.update(finished) 108 log.debug("unfinished_tasks = %d" % self.unfinished_tasks) 109 time.sleep(1) 110 if pbar.maxval != 0: 111 pbar.finish() 112 self.join() 113 if self._exception_queue.qsize() > 0: 114 raise exception.ThreadPoolException( 115 "An error occured in ThreadPool", self._exception_queue.queue)
116
117 - def __del__(self):
118 log.debug('del called in threadpool') 119 self.shutdown() 120 self.join()
121
122 123 -def get_thread_pool(size=10, worker_factory=_worker_factory, 124 disable_threads=False):
125 return ThreadPool(size=size, worker_factory=_worker_factory, 126 disable_threads=disable_threads)
127