1 

     2 # -*- coding: UTF-8 -*-

     3 """Easy to use object-oriented thread pool framework.

     4 

     5 A thread pool is an object that maintains a pool of worker threads to perform

     6 time consuming operations in parallel. It assigns jobs to the threads

     7 by putting them in a work request queue, where they are picked up by the

     8 next available thread. This then performs the requested operation in the

     9 background and puts the results in a another queue.

    10 

    11 The thread pool object can then collect the results from all threads from

    12 this queue as soon as they become available or after all threads have

    13 finished their work. It's also possible, to define callbacks to handle

    14 each result as it comes in.

    15 

    16 The basic concept and some code was taken from the book "Python in a Nutshell"

    17 by Alex Martelli, copyright 2003, ISBN 0-596-00188-6, from section 14.5

    18 "Threaded Program Architecture". I wrapped the main program logic in the

    19 ThreadPool class, added the WorkRequest class and the callback system and

    20 tweaked the code here and there. Kudos also to Florent Aide for the exception

    21 handling mechanism.

    22 

    23 Basic usage:

    24 

    25 >>> pool = TreadPool(poolsize)

    26 >>> requests = makeRequests(some_callable, list_of_args, callback)

    27 >>> [pool.putRequest(req) for req in requests]

    28 >>> pool.wait()

    29 

    30 See the end of the module code for a brief, annotated usage example.

    31 

    32 Website : http://chrisarndt.de/en/software/python/threadpool/

    33 """

    34 

    35 __all__ = [

    36   'makeRequests',

    37   'NoResultsPending',

    38   'NoWorkersAvailable',

    39   'ThreadPool',

    40   'WorkRequest',

    41   'WorkerThread'

    42 ]

    43 

    44 __author__ = "Christopher Arndt"

    45 __version__ = "1.2.3"

    46 __revision__ = "$Revision: 1.5 $"

    47 __date__ = "$Date: 2006/06/23 12:32:25 $"

    48 __license__ = 'Python license'

    49 

    50 # standard library modules

    51 import sys

    52 import threading

    53 import Queue

    54 

    55 # exceptions

    56 class NoResultsPending(Exception):

    57     """All work requests have been processed."""

    58     pass

    59 

    60 class NoWorkersAvailable(Exception):

    61     """No worker threads available to process remaining requests."""

    62     pass

    63 

    64 # classes

    65 class WorkerThread(threading.Thread):

    66     """Background thread connected to the requests/results queues.

    67 

    68     A worker thread sits in the background and picks up work requests from

    69     one queue and puts the results in another until it is dismissed.

    70     """

    71 

    72     def __init__(self, requestsQueue, resultsQueue, **kwds):

    73         """Set up thread in daemonic mode and start it immediatedly.

    74 

    75         requestsQueue and resultQueue are instances of Queue.Queue passed

    76         by the ThreadPool class when it creates a new worker thread.

    77         """

    78 

    79         threading.Thread.__init__(self, **kwds)

    80         self.setDaemon(1)

    81         self.workRequestQueue = requestsQueue

    82         self.resultQueue = resultsQueue

    83         self._dismissed = threading.Event()

    84         self.start()

    85 

    86     def run(self):

    87         """Repeatedly process the job queue until told to exit."""

    88 

    89         while not self._dismissed.isSet():

    90             # thread blocks here, if queue empty

    91             request = self.workRequestQueue.get()

    92             if self._dismissed.isSet():

    93                 # if told to exit, return the work request we just picked up

    94                 self.workRequestQueue.put(request)

    95                 break # and exit

    96             try:

    97                 self.resultQueue.put(

    98                     (request, request.callable(*request.args, **request.kwds))

    99                 )

   100             except:

   101                 request.exception = True

   102                 self.resultQueue.put((request, sys.exc_info()))

   103 

   104     def dismiss(self):

   105         """Sets a flag to tell the thread to exit when done with current job.

   106         """

   107 

   108         self._dismissed.set()

   109 

   110 

   111 class WorkRequest:

   112     """A request to execute a callable for putting in the request queue later.

   113 

   114     See the module function makeRequests() for the common case

   115     where you want to build several WorkRequests for the same callable

   116     but with different arguments for each call.

   117     """

   118 

   119     def __init__(self, callable, args=None, kwds=None, requestID=None,

   120       callback=None, exc_callback=None):

   121         """Create a work request for a callable and attach callbacks.

   122 

   123         A work request consists of the a callable to be executed by a

   124         worker thread, a list of positional arguments, a dictionary

   125         of keyword arguments.

   126 

   127         A callback function can be specified, that is called when the results

   128         of the request are picked up from the result queue. It must accept

   129         two arguments, the request object and the results of the callable,

   130         in that order. If you want to pass additional information to the

   131         callback, just stick it on the request object.

   132 

   133         You can also give a callback for when an exception occurs. It should

   134         also accept two arguments, the work request and a tuple with the

   135         exception details as returned by sys.exc_info().

   136 

   137         requestID, if given, must be hashable since it is used by the

   138         ThreadPool object to store the results of that work request in a

   139         dictionary. It defaults to the return value of id(self).

   140         """

   141 

   142         if requestID is None:

   143             self.requestID = id(self)

   144         else:

   145             try:

   146                 hash(requestID)

   147             except TypeError:

   148                 raise TypeError("requestID must be hashable.")

   149             self.requestID = requestID

   150         self.exception = False

   151         self.callback = callback

   152         self.exc_callback = exc_callback

   153         self.callable = callable

   154         self.args = args or []

   155         self.kwds = kwds or {}

   156 

   157 

   158 class ThreadPool:

   159     """A thread pool, distributing work requests and collecting results.

   160 

   161     See the module doctring for more information.

   162     """

   163 

   164     def __init__(self, num_workers, q_size=0):

   165         """Set up the thread pool and start num_workers worker threads.

   166 

   167         num_workers is the number of worker threads to start initialy.

   168         If q_size > 0 the size of the work request queue is limited and

   169         the thread pool blocks when the queue is full and it tries to put

   170         more work requests in it (see putRequest method).

   171         """

   172 

   173         self.requestsQueue = Queue.Queue(q_size)

   174         self.resultsQueue = Queue.Queue()

   175         self.workers = []

   176         self.workRequests = {}

   177         self.createWorkers(num_workers)

   178 

   179     def createWorkers(self, num_workers):

   180         """Add num_workers worker threads to the pool."""

   181 

   182         for i in range(num_workers):

   183             self.workers.append(WorkerThread(self.requestsQueue,

   184               self.resultsQueue))

   185 

   186     def dismissWorkers(self, num_workers):

   187         """Tell num_workers worker threads to quit after their current task.

   188         """

   189 

   190         for i in range(min(num_workers, len(self.workers))):

   191             worker = self.workers.pop()

   192             worker.dismiss()

   193 

   194     def addWork(self, callable, args=None, kwds=None, requestID=None, callback=None, exc_callback=None, block=True, timeout=0):

   195         request = WorkRequest(callable, args, kwds, requestID, callback, exc_callback)

   196         self.putRequest(request, block, timeout)

   197         

   198     def putRequest(self, request, block=True, timeout=0):

   199         """Put work request into work queue and save its id for later."""

   200 

   201         assert isinstance(request, WorkRequest)

   202         self.requestsQueue.put(request, block, timeout)

   203         self.workRequests[request.requestID] = request

   204 

   205     def poll(self, block=False):

   206         """Process any new results in the queue."""

   207 

   208         while True:

   209             # still results pending?

   210             if not self.workRequests:

   211                 raise NoResultsPending

   212             # are there still workers to process remaining requests?

   213             elif block and not self.workers:

   214                 raise NoWorkersAvailable

   215             try:

   216                 # get back next results

   217                 request, result = self.resultsQueue.get(block=block)

   218                 # has an exception occured?

   219                 if request.exception and request.exc_callback:

   220                     request.exc_callback(request, result)

   221                 # hand results to callback, if any

   222                 if request.callback and not \

   223                   (request.exception and request.exc_callback):

   224                     request.callback(request, result)

   225                 del self.workRequests[request.requestID]

   226             except Queue.Empty:

   227                 break

   228 

   229     def wait(self):

   230         """Wait for results, blocking until all have arrived."""

   231 

   232         while 1:

   233             try:

   234                 self.poll(True)

   235             except NoResultsPending:

   236                 break

   237 

   238 # helper functions

   239 def makeRequests(callable, args_list, callback=None, exc_callback=None):

   240     """Create several work requests for same callable with different arguments.

   241 

   242     Convenience function for creating several work requests for the same

   243     callable where each invocation of the callable receives different values

   244     for its arguments.

   245 

   246     args_list contains the parameters for each invocation of callable.

   247     Each item in 'args_list' should be either a 2-item tuple of the list of

   248     positional arguments and a dictionary of keyword arguments or a single,

   249     non-tuple argument.

   250 

   251     See docstring for WorkRequest for info on callback and exc_callback.

   252     """

   253 

   254     requests = []

   255     for item in args_list:

   256         if isinstance(item, tuple):

   257             requests.append(

   258               WorkRequest(callable, item[0], item[1], callback=callback,

   259                 exc_callback=exc_callback)

   260             )

   261         else:

   262             requests.append(

   263               WorkRequest(callable, [item], None, callback=callback,

   264                 exc_callback=exc_callback)

   265             )

   266     return requests

   267 

   268 ################

   269 # USAGE EXAMPLE

   270 ################

   271 

   272 if __name__ == '__main__':

   273     import random

   274     import time

   275 

   276     # the work the threads will have to do (rather trivial in our example)

   277     def do_something(data):

   278         time.sleep(random.randint(1, 5))

   279         result = round(random.random() * data, 5)

   280         # just to show off, we throw an exception once in a while

   281         if result > 3:

   282             raise RuntimeError("Something extraordinary happened!")

   283         return result

   284 

   285     # this will be called each time a result is available

   286     def print_result(request, result):

   287         print "**Result: %s from request #%s" % (result, request.requestID)

   288 

   289     # this will be called when an exception occurs within a thread

   290     def handle_exception(request, exc_info):

   291         print "Exception occured in request #%s: %s" % \

   292           (request.requestID, exc_info[1])

   293 

   294     # assemble the arguments for each job to a list...

   295     data = [random.randint(1, 10) for i in range(20)]

   296     # ... and build a WorkRequest object for each item in data

   297     requests = makeRequests(do_something, data, print_result, handle_exception)

   298 

   299     # or the other form of args_lists accepted by makeRequests: ((,), {})

   300     data = [((random.randint(1, 10), ), {}) for i in range(20)]

   301     requests.extend(

   302       makeRequests(do_something, data, print_result, handle_exception)

   303     )

   304 

   305     # we create a pool of 3 worker threads

   306     main = ThreadPool(3)

   307 

   308     # then we put the work requests in the queue...

   309     for req in requests:

   310         main.putRequest(req)

   311         print "Work request #%s added." % req.requestID

   312     # or shorter:

   313     # [main.putRequest(req) for req in requests]

   314 

   315     # ...and wait for the results to arrive in the result queue

   316     # by using ThreadPool.wait(). This would block until results for

   317     # all work requests have arrived:

   318     # main.wait()

   319 

   320     # instead we can poll for results while doing something else:

   321     i = 0

   322     while 1:

   323         try:

   324             main.poll()

   325             print "Main thread working..."

   326             time.sleep(0.5)

   327             if i == 10:

   328                 print "Adding 3 more worker threads..."

   329                 main.createWorkers(3)

   330             i += 1

   331         except KeyboardInterrupt:

   332             print "Interrupted!"

   333             break

   334         except NoResultsPending:

   335             print "All results collected."

   336             break

   337