starcluster.threadpool

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