This repository was archived by the owner on Apr 14, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 9
Expand file tree
/
Copy pathpool.py
More file actions
514 lines (429 loc) · 20.8 KB
/
pool.py
File metadata and controls
514 lines (429 loc) · 20.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors
#
# This module is part of async and is released under
# the New BSD License: http://www.opensource.org/licenses/bsd-license.php
"""Implementation of a thread-pool working with channels"""
from .thread import WorkerThread
from threading import Lock
from .util import AsyncQueue
from .graph import Graph
from .channel import (
ChannelWriter,
Channel,
SerialChannel,
CallbackChannelReader
)
import sys
import weakref
import logging
from functools import reduce
log = logging.root
__all__ = ('PoolReader', 'Pool', 'ThreadPool')
class PoolReader(CallbackChannelReader):
"""A reader designed to read from channels which take part in pools
It acts like a handle to the underlying task in the pool."""
__slots__ = ('_task_ref', '_pool_ref')
def __init__(self, channel, task, pool):
CallbackChannelReader.__init__(self, channel)
self._task_ref = weakref.ref(task)
self._pool_ref = weakref.ref(pool)
def __del__(self):
"""Assures that our task will be deleted if we were the last reader"""
task = self._task_ref()
if task is None:
return
pool = self._pool_ref()
if pool is None:
return
# if this is the last reader to the wc we just handled, there
# is no way anyone will ever read from the task again. If so,
# delete the task in question, it will take care of itself and orphans
# it might leave
# 1 is ourselves, + 1 for the call + 1, and 3 magical ones which
# I can't explain, but appears to be normal in the destructor
# On the caller side, getrefcount returns 2, as expected
# When just calling remove_task,
# it has no way of knowing that the write channel is about to diminsh.
# which is why we pass the info as a private kwarg - not nice, but
# okay for now
if sys.getrefcount(self) < 6:
pool.remove_task(task, _from_destructor_ = True)
# END handle refcount based removal of task
#{ Internal
def _read(self, count=0, block=True, timeout=None):
return CallbackChannelReader.read(self, count, block, timeout)
def pool_ref(self):
""":return: reference to the pool we belong to"""
return self._pool_ref
def task_ref(self):
""":return: reference to the task producing our items"""
return self._task_ref
#} END internal
#{ Interface
def task(self):
""":return: task we read from
:raise ValueError: If the instance is not attached to at task"""
task = self._task_ref()
if task is None:
raise ValueError("PoolReader is not associated with at task anymore")
return task
def pool(self):
""":return: pool our task belongs to
:raise ValueError: if the instance does not belong to a pool"""
pool = self._pool_ref()
if pool is None:
raise ValueError("PoolReader is not associated with a pool anymore")
return pool
#} END interface
def read(self, count=0, block=True, timeout=None):
"""Read an item that was processed by one of our threads
:note: Triggers task dependency handling needed to provide the necessary input"""
# NOTE: we always queue the operation that would give us count items
# as tracking the scheduled items or testing the channels size
# is in herently unsafe depending on the design of the task network
# If we put on tasks onto the queue for every request, we are sure
# to always produce enough items, even if the task.min_count actually
# provided enough - its better to have some possibly empty task runs
# than having and empty queue that blocks.
# if the user tries to use us to read from a done task, we will never
# compute as all produced items are already in the channel
task = self._task_ref()
if task is None:
return list()
# END abort if task was deleted
skip_compute = task.is_done() or task.error()
########## prepare ##############################
if not skip_compute:
self._pool_ref()._prepare_channel_read(task, count)
# END prepare pool scheduling
####### read data ########
##########################
# read actual items, tasks were setup to put their output into our channel ( as well )
items = CallbackChannelReader.read(self, count, block, timeout)
##########################
return items
class Pool(object):
"""A thread pool maintains a set of one or more worker threads, but supports
a fully serial mode in which case the amount of threads is zero.
Work is distributed via Channels, which form a dependency graph. The evaluation
is lazy, as work will only be done once an output is requested.
The thread pools inherent issue is the global interpreter lock that it will hit,
which gets worse considering a few c extensions specifically lock their part
globally as well. The only way this will improve is if custom c extensions
are written which do some bulk work, but release the GIL once they have acquired
their resources.
Due to the nature of having multiple objects in git, its easy to distribute
that work cleanly among threads.
:note: the current implementation returns channels which are meant to be
used only from the main thread, hence you cannot consume their results
from multiple threads unless you use a task for it."""
__slots__ = ( '_tasks', # a graph of tasks
'_num_workers', # list of workers
'_queue', # master queue for tasks
'_taskorder_cache', # map task id -> ordered dependent tasks
'_taskgraph_lock', # lock for accessing the task graph
)
# CONFIGURATION
# The type of worker to create - its expected to provide the Thread interface,
# taking the taskqueue as only init argument
# as well as a method called stop_and_join() to terminate it
WorkerCls = None
# The type of lock to use to protect critical sections, providing the
# threading.Lock interface
LockCls = None
# the type of the task queue to use - it must provide the Queue interface
TaskQueueCls = None
def __init__(self, size=0):
self._tasks = Graph()
self._num_workers = 0
self._queue = self.TaskQueueCls()
self._taskgraph_lock = self.LockCls()
self._taskorder_cache = dict()
self.set_size(size)
def __del__(self):
self.set_size(0)
#{ Internal
def _prepare_channel_read(self, task, count):
"""Process the tasks which depend on the given one to be sure the input
channels are filled with data once we process the actual task
Tasks have two important states: either they are done, or they are done
and have an error, so they are likely not to have finished all their work.
Either way, we will put them onto a list of tasks to delete them, providng
information about the failed ones.
Tasks which are not done will be put onto the queue for processing, which
is fine as we walked them depth-first."""
# for the walk, we must make sure the ordering does not change. Even
# when accessing the cache, as it is related to graph changes
self._taskgraph_lock.acquire()
try:
try:
dfirst_tasks = self._taskorder_cache[id(task)]
except KeyError:
# have to retrieve the list from the graph
dfirst_tasks = self._tasks.input_inclusive_dfirst_reversed(task)
self._taskorder_cache[id(task)] = dfirst_tasks
# END handle cached order retrieval
finally:
self._taskgraph_lock.release()
# END handle locking
# check the min count on all involved tasks, and be sure that we don't
# have any task which produces less than the maximum min-count of all tasks
# The actual_count is used when chunking tasks up for the queue, whereas
# the count is usued to determine whether we still have enough output
# on the queue, checking qsize ( ->revise )
# ABTRACT: If T depends on T-1, and the client wants 1 item, T produces
# at least 10, T-1 goes with 1, then T will block after 1 item, which
# is read by the client. On the next read of 1 item, we would find T's
# queue empty and put in another 10, which could put another thread into
# blocking state. T-1 produces one more item, which is consumed right away
# by the two threads running T. Although this works in the end, it leaves
# many threads blocking and waiting for input, which is not desired.
# Setting the min-count to the max of the mincount of all tasks assures
# we have enough items for all.
# Addition: in serial mode, we would enter a deadlock if one task would
# ever wait for items !
actual_count = count
min_counts = (((t.min_count is not None and t.min_count) or count) for t in dfirst_tasks)
min_count = reduce(lambda m1, m2: max(m1, m2), min_counts)
if 0 < count < min_count:
actual_count = min_count
# END set actual count
# the list includes our tasks - the first one to evaluate first, the
# requested one last
for task in dfirst_tasks:
# if task.error() or task.is_done():
# in theory, the should never be consumed task in the pool, right ?
# They delete themselves once they are done. But as we run asynchronously,
# It can be that someone reads, while a task realizes its done, and
# we get here to prepare the read although it already is done.
# Its not a problem though, the task wiill not do anything.
# Hence we don't waste our time with checking for it
# raise AssertionError("Shouldn't have consumed tasks on the pool, they delete themeselves, what happend ?")
# END skip processing
# but use the actual count to produce the output, we may produce
# more than requested
numchunks = 1
chunksize = actual_count
remainder = 0
# we need the count set for this - can't chunk up unlimited items
# In serial mode we could do this by checking for empty input channels,
# but in dispatch mode its impossible ( == not easily possible )
# Only try it if we have enough demand
if task.max_chunksize and actual_count > task.max_chunksize:
numchunks = actual_count / task.max_chunksize
chunksize = task.max_chunksize
remainder = actual_count - (numchunks * chunksize)
# END handle chunking
# These both need to be integers, not floats
chunksize = int(chunksize)
numchunks = int(numchunks)
# the following loops are kind of unrolled - code duplication
# should make things execute faster. Putting the if statements
# into the loop would be less code, but ... slower
if self._num_workers:
# respect the chunk size, and split the task up if we want
# to process too much. This can be defined per task
qput = self._queue.put
if numchunks > 1:
for i in range(numchunks):
qput((task.process, chunksize))
# END for each chunk to put
else:
qput((task.process, chunksize))
# END try efficient looping
if remainder:
qput((task.process, remainder))
# END handle chunksize
else:
# no workers, so we have to do the work ourselves
if numchunks > 1:
for i in range(numchunks):
task.process(chunksize)
# END for each chunk to put
else:
task.process(chunksize)
# END try efficient looping
if remainder:
task.process(remainder)
# END handle chunksize
# END handle serial mode
# END for each task to process
def _remove_task_if_orphaned(self, task, from_destructor):
"""Check the task, and delete it if it is orphaned"""
# 1 for writer on task, 1 for the getrefcount call + 1 for each other writer/reader
# If we are getting here from the destructor of an RPool channel,
# its totally valid to virtually decrement the refcount by 1 as
# we can expect it to drop once the destructor completes, which is when
# we finish all recursive calls
max_ref_count = 3 + from_destructor
if sys.getrefcount(task.writer().channel) < max_ref_count:
self.remove_task(task, from_destructor)
#} END internal
#{ Interface
def size(self):
""":return: amount of workers in the pool
:note: method is not threadsafe !"""
return self._num_workers
def set_size(self, size=0):
"""Set the amount of workers to use in this pool. When reducing the size,
threads will continue with their work until they are done before effectively
being removed.
:return: self
:param size: if 0, the pool will do all work itself in the calling thread,
otherwise the work will be distributed among the given amount of threads.
If the size is 0, newly added tasks will use channels which are NOT
threadsafe to optimize item throughput.
:note: currently NOT threadsafe !"""
assert size > -1, "Size cannot be negative"
# Enforce sync operation in py3 - it doesn't work. More information in-code at async.test.lib.py:9
if True:
log.debug("async: Actual asynchronous operation was disabled as it is slower that way")
size = 0
# end
# either start new threads, or kill existing ones.
# If we end up with no threads, we process the remaining chunks on the queue
# ourselves
cur_count = self._num_workers
if cur_count < size:
# we can safely increase the size, even from serial mode, as we would
# only be able to do this if the serial ( sync ) mode finished processing.
# Just adding more workers is not a problem at all.
add_count = size - cur_count
for i in range(add_count):
self.WorkerCls(self._queue).start()
# END for each new worker to create
self._num_workers += add_count
elif cur_count > size:
# We don't care which thread exactly gets hit by our stop request
# On their way, they will consume remaining tasks, but new ones
# could be added as we speak.
del_count = cur_count - size
for i in range(del_count):
self._queue.put((self.WorkerCls.stop, True)) # arg doesnt matter
# END for each thread to stop
self._num_workers -= del_count
# END handle count
if size == 0:
# NOTE: we do not preocess any tasks still on the queue, as we ill
# naturally do that once we read the next time, only on the tasks
# that are actually required. The queue will keep the tasks,
# and once we are deleted, they will vanish without additional
# time spend on them. If there shouldn't be any consumers anyway.
# If we should reenable some workers again, they will continue on the
# remaining tasks, probably with nothing to do.
# We can't clear the task queue if we have removed workers
# as they will receive the termination signal through it, and if
# we had added workers, we wouldn't be here ;).
pass
# END process queue
return self
def num_tasks(self):
""":return: amount of tasks"""
self._taskgraph_lock.acquire()
try:
return len(self._tasks.nodes)
finally:
self._taskgraph_lock.release()
def remove_task(self, task, _from_destructor_ = False):
"""
Delete the task.
Additionally we will remove orphaned tasks, which can be identified if their
output channel is only held by themselves, so no one will ever consume
its items.
This method blocks until all tasks to be removed have been processed, if
they are currently being processed.
:return: self"""
self._taskgraph_lock.acquire()
try:
# it can be that the task is already deleted, but its chunk was on the
# queue until now, so its marked consumed again
if not task in self._tasks.nodes:
return self
# END early abort
# the task we are currently deleting could also be processed by
# a thread right now. We don't care about it as its taking care about
# its write channel itself, and sends everything it can to it.
# For it it doesn't matter that its not part of our task graph anymore.
# now delete our actual node - be sure its done to prevent further
# processing in case there are still client reads on their way.
task.set_done()
# keep its input nodes as we check whether they were orphaned
in_tasks = task.in_nodes
self._tasks.remove_node(task)
self._taskorder_cache.clear()
finally:
self._taskgraph_lock.release()
# END locked deletion
for t in in_tasks:
self._remove_task_if_orphaned(t, _from_destructor_)
# END handle orphans recursively
return self
def add_task(self, task):
"""Add a new task to be processed.
:return: a read channel to retrieve processed items. If that handle is lost,
the task will be considered orphaned and will be deleted on the next
occasion."""
# create a write channel for it
ctype = Channel
# adjust the task with our pool ref, if it has the slot and is empty
# For now, we don't allow tasks to be used in multiple pools, except
# for by their channels
if hasattr(task, 'pool'):
their_pool = task.pool()
if their_pool is None:
task.set_pool(self)
elif their_pool is not self:
raise ValueError("Task %r is already registered to another pool" % task.id)
# END handle pool exclusivity
# END handle pool aware tasks
self._taskgraph_lock.acquire()
try:
self._taskorder_cache.clear()
self._tasks.add_node(task)
# Use a non-threadsafe queue
# This brings about 15% more performance, but sacrifices thread-safety
if self.size() == 0:
ctype = SerialChannel
# END improve locks
# setup the tasks channel - respect the task creators choice though
# if it is set.
wc = task.writer()
ch = None
if wc is None:
ch = ctype()
wc = ChannelWriter(ch)
task.set_writer(wc)
else:
ch = wc.channel
# END create write channel ifunset
rc = PoolReader(ch, task, self)
finally:
self._taskgraph_lock.release()
# END sync task addition
# If the input channel is one of our read channels, we add the relation
if hasattr(task, 'reader'):
ic = task.reader()
if hasattr(ic, 'pool_ref') and ic.pool_ref()() is self:
self._taskgraph_lock.acquire()
try:
self._tasks.add_edge(ic._task_ref(), task)
# additionally, bypass ourselves when reading from the
# task, if possible
if hasattr(ic, '_read'):
task.set_read(ic._read)
# END handle read bypass
finally:
self._taskgraph_lock.release()
# END handle edge-adding
# END add task relation
# END handle input channels for connections
return rc
#} END interface
class ThreadPool(Pool):
"""A pool using threads as worker"""
WorkerCls = WorkerThread
LockCls = Lock
# NOTE: Since threading.Lock is a method not a class, we need to prevent
# conversion to an unbound method.
LockCls = staticmethod(Lock)
TaskQueueCls = AsyncQueue