threadpool.py

   1 """Easy to use object-oriented thread pool framework.
   2
   3 A thread pool is an object that maintains a pool of worker threads to perform
   4 time consuming operations in parallel. It assigns jobs to the threads
   5 by putting them in a work request queue, where they are picked up by the
   6 next available thread. This then performs the requested operation in the
   7 background and puts the results in a another queue.
   8
   9 The thread pool object can then collect the results from all threads from
  10 this queue as soon as they become available or after all threads have
  11 finished their work. It's also possible, to define callbacks to handle
  12 each result as it comes in.
  13
  14 The basic concept and some code was taken from the book "Python in a Nutshell"
  15 by Alex Martelli, copyright 2003, ISBN 0-596-00188-6, from section 14.5
  16 "Threaded Program Architecture". I wrapped the main program logic in the
  17 ThreadPool class, added the WorkRequest class and the callback system and
  18 tweaked the code here and there. Kudos also to Florent Aide for the exception
  19 handling mechanism.
  20
  21 Basic usage:
  22
  23 >>> pool = TreadPool(poolsize)
  24 >>> requests = makeRequests(some_callable, list_of_args, callback)
  25 >>> [pool.putRequest(req) for req in requests]
  26 >>> pool.wait()
  27
  28 See the end of the module code for a brief, annotated usage example.
  29
  30 Website : http://chrisarndt.de/en/software/python/threadpool/
  31 """
  32
  33 __all__ = [
  34   'makeRequests',
  35   'NoResultsPending',
  36   'NoWorkersAvailable',
  37   'ThreadPool',
  38   'WorkRequest',
  39   'WorkerThread'
  40 ]
  41
  42 __author__ = "Christopher Arndt"
  43 __version__ = "1.2.3"
  44 __revision__ = "$Revision: 1.5 $"
  45 __date__ = "$Date: 2006/06/23 12:32:25 $"
  46 __license__ = 'Python license'
  47
  48 # standard library modules
  49 import sys
  50 import threading
  51 import Queue
  52
  53 # exceptions
  54 class NoResultsPending(Exception):
  55     """All work requests have been processed."""
  56     pass
  57
  58 class NoWorkersAvailable(Exception):
  59     """No worker threads available to process remaining requests."""
  60     pass
  61
  62 # classes
  63 class WorkerThread(threading.Thread):
  64     """Background thread connected to the requests/results queues.
  65
  66     A worker thread sits in the background and picks up work requests from
  67     one queue and puts the results in another until it is dismissed.
  68     """
  69
  70     def __init__(self, requestsQueue, resultsQueue, **kwds):
  71         """Set up thread in daemonic mode and start it immediatedly.
  72
  73         requestsQueue and resultQueue are instances of Queue.Queue passed
  74         by the ThreadPool class when it creates a new worker thread.
  75         """
  76
  77         threading.Thread.__init__(self, **kwds)
  78         self.setDaemon(1)
  79         self.workRequestQueue = requestsQueue
  80         self.resultQueue = resultsQueue
  81         self._dismissed = threading.Event()
  82         self.start()
  83
  84     def run(self):
  85         """Repeatedly process the job queue until told to exit."""
  86
  87         while not self._dismissed.isSet():
  88             # thread blocks here, if queue empty
  89             request = self.workRequestQueue.get()
  90             if self._dismissed.isSet():
  91                 # if told to exit, return the work request we just picked up
  92                 self.workRequestQueue.put(request)
  93                 break # and exit
  94             try:
  95                 self.resultQueue.put(
  96                     (request, request.callable(*request.args, **request.kwds))
  97                 )
  98             except:
  99                 request.exception = True
 100                 self.resultQueue.put((request, sys.exc_info()))
 101
 102     def dismiss(self):
 103         """Sets a flag to tell the thread to exit when done with current job.
 104         """
 105
 106         self._dismissed.set()
 107
 108
 109 class WorkRequest:
 110     """A request to execute a callable for putting in the request queue later.
 111
 112     See the module function makeRequests() for the common case
 113     where you want to build several WorkRequests for the same callable
 114     but with different arguments for each call.
 115     """
 116
 117     def __init__(self, callable, args=None, kwds=None, requestID=None,
 118       callback=None, exc_callback=None):
 119         """Create a work request for a callable and attach callbacks.
 120
 121         A work request consists of the a callable to be executed by a
 122         worker thread, a list of positional arguments, a dictionary
 123         of keyword arguments.
 124
 125         A callback function can be specified, that is called when the results
 126         of the request are picked up from the result queue. It must accept
 127         two arguments, the request object and the results of the callable,
 128         in that order. If you want to pass additional information to the
 129         callback, just stick it on the request object.
 130
 131         You can also give a callback for when an exception occurs. It should
 132         also accept two arguments, the work request and a tuple with the
 133         exception details as returned by sys.exc_info().
 134
 135         requestID, if given, must be hashable since it is used by the
 136         ThreadPool object to store the results of that work request in a
 137         dictionary. It defaults to the return value of id(self).
 138         """
 139
 140         if requestID is None:
 141             self.requestID = id(self)
 142         else:
 143             try:
 144                 hash(requestID)
 145             except TypeError:
 146                 raise TypeError("requestID must be hashable.")
 147             self.requestID = requestID
 148         self.exception = False
 149         self.callback = callback
 150         self.exc_callback = exc_callback
 151         self.callable = callable
 152         self.args = args or []
 153         self.kwds = kwds or {}
 154
 155
 156 class ThreadPool:
 157     """A thread pool, distributing work requests and collecting results.
 158
 159     See the module doctring for more information.
 160     """
 161
 162     def __init__(self, num_workers, q_size=0):
 163         """Set up the thread pool and start num_workers worker threads.
 164
 165         num_workers is the number of worker threads to start initialy.
 166         If q_size > 0 the size of the work request queue is limited and
 167         the thread pool blocks when the queue is full and it tries to put
 168         more work requests in it (see putRequest method).
 169         """
 170
 171         self.requestsQueue = Queue.Queue(q_size)
 172         self.resultsQueue = Queue.Queue()
 173         self.workers = []
 174         self.workRequests = {}
 175         self.createWorkers(num_workers)
 176
 177     def createWorkers(self, num_workers):
 178         """Add num_workers worker threads to the pool."""
 179
 180         for i in range(num_workers):
 181             self.workers.append(WorkerThread(self.requestsQueue,
 182               self.resultsQueue))
 183
 184     def dismissWorkers(self, num_workers):
 185         """Tell num_workers worker threads to quit after their current task.
 186         """
 187
 188         for i in range(min(num_workers, len(self.workers))):
 189             worker = self.workers.pop()
 190             worker.dismiss()
 191
 192     def putRequest(self, request, block=True, timeout=0):
 193         """Put work request into work queue and save its id for later."""
 194
 195         assert isinstance(request, WorkRequest)
 196         self.requestsQueue.put(request, block, timeout)
 197         self.workRequests[request.requestID] = request
 198
 199     def poll(self, block=False):
 200         """Process any new results in the queue."""
 201
 202         while True:
 203             # still results pending?
 204             if not self.workRequests:
 205                 raise NoResultsPending
 206             # are there still workers to process remaining requests?
 207             elif block and not self.workers:
 208                 raise NoWorkersAvailable
 209             try:
 210                 # get back next results
 211                 request, result = self.resultsQueue.get(block=block)
 212                 # has an exception occured?
 213                 if request.exception and request.exc_callback:
 214                     request.exc_callback(request, result)
 215                 # hand results to callback, if any
 216                 if request.callback and not \
 217                   (request.exception and request.exc_callback):
 218                     request.callback(request, result)
 219                 del self.workRequests[request.requestID]
 220             except Queue.Empty:
 221                 break
 222
 223     def wait(self):
 224         """Wait for results, blocking until all have arrived."""
 225
 226         while 1:
 227             try:
 228                 self.poll(True)
 229             except NoResultsPending:
 230                 break
 231
 232 # helper functions
 233 def makeRequests(callable, args_list, callback=None, exc_callback=None):
 234     """Create several work requests for same callable with different arguments.
 235
 236     Convenience function for creating several work requests for the same
 237     callable where each invocation of the callable receives different values
 238     for its arguments.
 239
 240     args_list contains the parameters for each invocation of callable.
 241     Each item in 'args_list' should be either a 2-item tuple of the list of
 242     positional arguments and a dictionary of keyword arguments or a single,
 243     non-tuple argument.
 244
 245     See docstring for WorkRequest for info on callback and exc_callback.
 246     """
 247
 248     requests = []
 249     for item in args_list:
 250         if isinstance(item, tuple):
 251             requests.append(
 252               WorkRequest(callable, item[0], item[1], callback=callback,
 253                 exc_callback=exc_callback)
 254             )
 255         else:
 256             requests.append(
 257               WorkRequest(callable, [item], None, callback=callback,
 258                 exc_callback=exc_callback)
 259             )
 260     return requests
 261
 262 ################
 263 # USAGE EXAMPLE
 264 ################
 265
 266 if __name__ == '__main__':
 267     import random
 268     import time
 269
 270     # the work the threads will have to do (rather trivial in our example)
 271     def do_something(data):
 272         time.sleep(random.randint(1,5))
 273         result = round(random.random() * data, 5)
 274         # just to show off, we throw an exception once in a while
 275         if result > 3:
 276             raise RuntimeError("Something extraordinary happened!")
 277         return result
 278
 279     # this will be called each time a result is available
 280     def print_result(request, result):
 281         print "**Result: %s from request #%s" % (result, request.requestID)
 282
 283     # this will be called when an exception occurs within a thread
 284     def handle_exception(request, exc_info):
 285         print "Exception occured in request #%s: %s" % \
 286           (request.requestID, exc_info[1])
 287
 288     # assemble the arguments for each job to a list...
 289     data = [random.randint(1,10) for i in range(20)]
 290     # ... and build a WorkRequest object for each item in data
 291     requests = makeRequests(do_something, data, print_result, handle_exception)
 292
 293     # or the other form of args_lists accepted by makeRequests: ((,), {})
 294     data = [((random.randint(1,10),), {}) for i in range(20)]
 295     requests.extend(
 296       makeRequests(do_something, data, print_result, handle_exception)
 297     )
 298
 299     # we create a pool of 3 worker threads
 300     main = ThreadPool(3)
 301
 302     # then we put the work requests in the queue...
 303     for req in requests:
 304         main.putRequest(req)
 305         print "Work request #%s added." % req.requestID
 306     # or shorter:
 307     # [main.putRequest(req) for req in requests]
 308
 309     # ...and wait for the results to arrive in the result queue
 310     # by using ThreadPool.wait(). This would block until results for
 311     # all work requests have arrived:
 312     # main.wait()
 313
 314     # instead we can poll for results while doing something else:
 315     i = 0
 316     while 1:
 317         try:
 318             main.poll()
 319             print "Main thread working..."
 320             time.sleep(0.5)
 321             if i == 10:
 322                 print "Adding 3 more worker threads..."
 323                 main.createWorkers(3)
 324             i += 1
 325         except KeyboardInterrupt:
 326             print "Interrupted!"
 327             break
 328         except NoResultsPending:
 329             print "All results collected."
 330             break