generated from daniil-berg/boilerplate-py
Compare commits
2 Commits
Author | SHA1 | Date | |
---|---|---|---|
410e73e68b | |||
727f0b7c8b |
@ -1,6 +1,6 @@
|
||||
[metadata]
|
||||
name = asyncio-taskpool
|
||||
version = 0.1.5
|
||||
version = 0.1.7
|
||||
author = Daniil Fajnberg
|
||||
author_email = mail@daniil.fajnberg.de
|
||||
description = Dynamically manage pools of asyncio tasks
|
||||
|
@ -3,6 +3,6 @@ MSG_BYTES = 1024
|
||||
CMD_START = 'start'
|
||||
CMD_STOP = 'stop'
|
||||
CMD_STOP_ALL = 'stop_all'
|
||||
CMD_SIZE = 'size'
|
||||
CMD_NUM_RUNNING = 'num_running'
|
||||
CMD_FUNC = 'func'
|
||||
CLIENT_EXIT = 'exit'
|
||||
|
@ -438,7 +438,7 @@ class TaskPool(BaseTaskPool):
|
||||
break
|
||||
await q.put(arg) # This blocks as long as the queue is full.
|
||||
|
||||
async def _queue_consumer(self, q: Queue, func: CoroutineFunc, arg_stars: int = 0,
|
||||
async def _queue_consumer(self, q: Queue, first_batch_started: Event, func: CoroutineFunc, arg_stars: int = 0,
|
||||
end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
|
||||
"""
|
||||
Wrapper around the `_start_task()` taking the next element from the arguments queue set up in `_map()`.
|
||||
@ -447,6 +447,9 @@ class TaskPool(BaseTaskPool):
|
||||
Args:
|
||||
q:
|
||||
The queue of function arguments to consume for starting the next task.
|
||||
first_batch_started:
|
||||
The event flag to wait for, before launching the next consumer.
|
||||
It can only set by the `_map()` method, which happens after the first batch of task has been started.
|
||||
func:
|
||||
The coroutine function to use for spawning the tasks within the task pool.
|
||||
arg_stars (optional):
|
||||
@ -466,15 +469,17 @@ class TaskPool(BaseTaskPool):
|
||||
await self._start_task(
|
||||
star_function(func, arg, arg_stars=arg_stars),
|
||||
ignore_closed=True,
|
||||
end_callback=partial(TaskPool._queue_callback, self, q=q, func=func, arg_stars=arg_stars,
|
||||
end_callback=end_callback, cancel_callback=cancel_callback),
|
||||
end_callback=partial(TaskPool._queue_callback, self, q=q, first_batch_started=first_batch_started,
|
||||
func=func, arg_stars=arg_stars, end_callback=end_callback,
|
||||
cancel_callback=cancel_callback),
|
||||
cancel_callback=cancel_callback
|
||||
)
|
||||
finally:
|
||||
q.task_done()
|
||||
|
||||
async def _queue_callback(self, task_id: int, q: Queue, func: CoroutineFunc, arg_stars: int = 0,
|
||||
end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
|
||||
async def _queue_callback(self, task_id: int, q: Queue, first_batch_started: Event, func: CoroutineFunc,
|
||||
arg_stars: int = 0, end_callback: EndCallbackT = None,
|
||||
cancel_callback: CancelCallbackT = None) -> None:
|
||||
"""
|
||||
Wrapper around an end callback function passed into the `_map()` method.
|
||||
Triggers the next `_queue_consumer` with the same arguments.
|
||||
@ -484,6 +489,9 @@ class TaskPool(BaseTaskPool):
|
||||
The ID of the ending task.
|
||||
q:
|
||||
The queue of function arguments to consume for starting the next task.
|
||||
first_batch_started:
|
||||
The event flag to wait for, before launching the next consumer.
|
||||
It can only set by the `_map()` method, which happens after the first batch of task has been started.
|
||||
func:
|
||||
The coroutine function to use for spawning the tasks within the task pool.
|
||||
arg_stars (optional):
|
||||
@ -495,7 +503,9 @@ class TaskPool(BaseTaskPool):
|
||||
The callback that was specified to execute after cancellation of the task (and the next one).
|
||||
It is run with the `task_id` as its only positional argument.
|
||||
"""
|
||||
await self._queue_consumer(q, func, arg_stars, end_callback=end_callback, cancel_callback=cancel_callback)
|
||||
await first_batch_started.wait()
|
||||
await self._queue_consumer(q, first_batch_started, func, arg_stars,
|
||||
end_callback=end_callback, cancel_callback=cancel_callback)
|
||||
await execute_optional(end_callback, args=(task_id,))
|
||||
|
||||
def _set_up_args_queue(self, args_iter: ArgsT, num_tasks: int) -> Queue:
|
||||
@ -574,10 +584,18 @@ class TaskPool(BaseTaskPool):
|
||||
if not self.is_open:
|
||||
raise exceptions.PoolIsClosed("Cannot start new tasks")
|
||||
args_queue = self._set_up_args_queue(args_iter, num_tasks)
|
||||
# We need a flag to ensure that starting all tasks from the first batch here will not be blocked by the
|
||||
# `_queue_callback` triggered by one or more of them.
|
||||
# This could happen, e.g. if the pool has just enough room for one more task, but the queue here contains more
|
||||
# than one element, and the pool remains full until after the first task of the first batch ends. Then the
|
||||
# callback might trigger the next `_queue_consumer` before this method can, which will keep it blocked.
|
||||
first_batch_started = Event()
|
||||
for _ in range(args_queue.qsize()):
|
||||
# This is where blocking can occur, if the pool is full.
|
||||
await self._queue_consumer(args_queue, func,
|
||||
await self._queue_consumer(args_queue, first_batch_started, func,
|
||||
arg_stars=arg_stars, end_callback=end_callback, cancel_callback=cancel_callback)
|
||||
# Now the callbacks can immediately trigger more tasks.
|
||||
first_batch_started.set()
|
||||
|
||||
async def map(self, func: CoroutineFunc, arg_iter: ArgsT, num_tasks: int = 1,
|
||||
end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
|
||||
@ -638,9 +656,45 @@ class TaskPool(BaseTaskPool):
|
||||
|
||||
|
||||
class SimpleTaskPool(BaseTaskPool):
|
||||
"""
|
||||
Simplified task pool class.
|
||||
|
||||
A `SimpleTaskPool` instance can manage an arbitrary number of concurrent tasks,
|
||||
but they **must** come from a single coroutine function, called with the same arguments.
|
||||
|
||||
The coroutine function and its arguments are defined upon initialization.
|
||||
|
||||
As long as there is room in the pool, more tasks can be added. (By default, there is no pool size limit.)
|
||||
Each task started in the pool receives a unique ID, which can be used to cancel specific tasks at any moment.
|
||||
However, since all tasks come from the same function-arguments-combination, the specificity of the `cancel()` method
|
||||
is probably unnecessary. Instead, a simpler `stop()` method is introduced.
|
||||
|
||||
Adding tasks blocks **only if** the pool is full at that moment.
|
||||
"""
|
||||
|
||||
def __init__(self, func: CoroutineFunc, args: ArgsT = (), kwargs: KwArgsT = None,
|
||||
end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None,
|
||||
name: str = None) -> None:
|
||||
pool_size: int = inf, name: str = None) -> None:
|
||||
"""
|
||||
|
||||
Args:
|
||||
func:
|
||||
The function to use for spawning new tasks within the pool.
|
||||
args (optional):
|
||||
The positional arguments to pass into each function call.
|
||||
kwargs (optional):
|
||||
The keyword-arguments to pass into each function call.
|
||||
end_callback (optional):
|
||||
A callback to execute after a task has ended.
|
||||
It is run with the task's ID as its only positional argument.
|
||||
cancel_callback (optional):
|
||||
A callback to execute after cancellation of a task.
|
||||
It is run with the task's ID as its only positional argument.
|
||||
pool_size (optional):
|
||||
The maximum number of tasks allowed to run concurrently in the pool
|
||||
name (optional):
|
||||
An optional name for the pool.
|
||||
"""
|
||||
if not iscoroutinefunction(func):
|
||||
raise exceptions.NotCoroutine(f"Not a coroutine function: {func}")
|
||||
self._func: CoroutineFunc = func
|
||||
@ -648,32 +702,39 @@ class SimpleTaskPool(BaseTaskPool):
|
||||
self._kwargs: KwArgsT = kwargs if kwargs is not None else {}
|
||||
self._end_callback: EndCallbackT = end_callback
|
||||
self._cancel_callback: CancelCallbackT = cancel_callback
|
||||
super().__init__(name=name)
|
||||
super().__init__(pool_size=pool_size, name=name)
|
||||
|
||||
@property
|
||||
def func_name(self) -> str:
|
||||
"""Returns the name of the coroutine function used in the pool."""
|
||||
return self._func.__name__
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
return self.num_running
|
||||
|
||||
async def _start_one(self) -> int:
|
||||
"""Starts a single new task within the pool and returns its ID."""
|
||||
return await self._start_task(self._func(*self._args, **self._kwargs),
|
||||
end_callback=self._end_callback, cancel_callback=self._cancel_callback)
|
||||
|
||||
async def start(self, num: int = 1) -> List[int]:
|
||||
return [await self._start_one() for _ in range(num)]
|
||||
"""Starts `num` new tasks within the pool and returns their IDs as a list."""
|
||||
ids = await gather(*(self._start_one() for _ in range(num)))
|
||||
assert isinstance(ids, list) # for PyCharm (see above to-do-item)
|
||||
return ids
|
||||
|
||||
def stop(self, num: int = 1) -> List[int]:
|
||||
num = min(num, self.size)
|
||||
"""
|
||||
Cancels `num` running tasks within the pool and returns their IDs as a list.
|
||||
|
||||
The tasks are canceled in LIFO order, meaning tasks started later will be stopped before those started earlier.
|
||||
If `num` is greater than or equal to the number of currently running tasks, naturally all tasks are cancelled.
|
||||
"""
|
||||
ids = []
|
||||
for i, task_id in enumerate(reversed(self._running)):
|
||||
if i >= num:
|
||||
break
|
||||
break # We got the desired number of task IDs, there may well be more tasks left to keep running
|
||||
ids.append(task_id)
|
||||
self.cancel(*ids)
|
||||
return ids
|
||||
|
||||
def stop_all(self) -> List[int]:
|
||||
return self.stop(self.size)
|
||||
"""Cancels all running tasks and returns their IDs as a list."""
|
||||
return self.stop(self.num_running)
|
||||
|
@ -63,8 +63,8 @@ class ControlServer(ABC): # TODO: Implement interface for normal TaskPool insta
|
||||
writer.write(str(self._pool.stop_all()).encode())
|
||||
|
||||
def _pool_size(self, writer: StreamWriter) -> None:
|
||||
log.debug("%s requests pool size", self.client_class.__name__)
|
||||
writer.write(str(self._pool.size).encode())
|
||||
log.debug("%s requests number of running tasks", self.client_class.__name__)
|
||||
writer.write(str(self._pool.num_running).encode())
|
||||
|
||||
def _pool_func(self, writer: StreamWriter) -> None:
|
||||
log.debug("%s requests pool function", self.client_class.__name__)
|
||||
@ -83,7 +83,7 @@ class ControlServer(ABC): # TODO: Implement interface for normal TaskPool insta
|
||||
self._stop_tasks(writer, arg)
|
||||
elif cmd == constants.CMD_STOP_ALL:
|
||||
self._stop_all_tasks(writer)
|
||||
elif cmd == constants.CMD_SIZE:
|
||||
elif cmd == constants.CMD_NUM_RUNNING:
|
||||
self._pool_size(writer)
|
||||
elif cmd == constants.CMD_FUNC:
|
||||
self._pool_func(writer)
|
||||
|
@ -446,14 +446,14 @@ class TaskPoolTestCase(CommonTestCase):
|
||||
q, arg = Queue(), 420.69
|
||||
q.put_nowait(arg)
|
||||
mock_func, stars = MagicMock(), 3
|
||||
end_cb, cancel_cb = MagicMock(), MagicMock()
|
||||
self.assertIsNone(await self.task_pool._queue_consumer(q, mock_func, stars, end_cb, cancel_cb))
|
||||
mock_flag, end_cb, cancel_cb = MagicMock(), MagicMock(), MagicMock()
|
||||
self.assertIsNone(await self.task_pool._queue_consumer(q, mock_flag, mock_func, stars, end_cb, cancel_cb))
|
||||
self.assertTrue(q.empty())
|
||||
mock__start_task.assert_awaited_once_with(awaitable, ignore_closed=True,
|
||||
end_callback=queue_callback, cancel_callback=cancel_cb)
|
||||
mock_star_function.assert_called_once_with(mock_func, arg, arg_stars=stars)
|
||||
mock_partial.assert_called_once_with(pool.TaskPool._queue_callback, self.task_pool,
|
||||
q=q, func=mock_func, arg_stars=stars,
|
||||
q=q, first_batch_started=mock_flag, func=mock_func, arg_stars=stars,
|
||||
end_callback=end_cb, cancel_callback=cancel_cb)
|
||||
mock__start_task.reset_mock()
|
||||
mock_star_function.reset_mock()
|
||||
@ -470,9 +470,13 @@ class TaskPoolTestCase(CommonTestCase):
|
||||
async def test__queue_callback(self, mock__queue_consumer: AsyncMock, mock_execute_optional: AsyncMock):
|
||||
task_id, mock_q = 420, MagicMock()
|
||||
mock_func, stars = MagicMock(), 3
|
||||
mock_wait = AsyncMock()
|
||||
mock_flag = MagicMock(wait=mock_wait)
|
||||
end_cb, cancel_cb = MagicMock(), MagicMock()
|
||||
self.assertIsNone(await self.task_pool._queue_callback(task_id, mock_q, mock_func, stars, end_cb, cancel_cb))
|
||||
mock__queue_consumer.assert_awaited_once_with(mock_q, mock_func, stars,
|
||||
self.assertIsNone(await self.task_pool._queue_callback(task_id, mock_q, mock_flag, mock_func, stars,
|
||||
end_callback=end_cb, cancel_callback=cancel_cb))
|
||||
mock_wait.assert_awaited_once_with()
|
||||
mock__queue_consumer.assert_awaited_once_with(mock_q, mock_flag, mock_func, stars,
|
||||
end_callback=end_cb, cancel_callback=cancel_cb)
|
||||
mock_execute_optional.assert_awaited_once_with(end_cb, args=(task_id,))
|
||||
|
||||
@ -521,13 +525,16 @@ class TaskPoolTestCase(CommonTestCase):
|
||||
mock__queue_producer.assert_not_called()
|
||||
mock_create_task.assert_not_called()
|
||||
|
||||
@patch.object(pool, 'Event')
|
||||
@patch.object(pool.TaskPool, '_queue_consumer')
|
||||
@patch.object(pool.TaskPool, '_set_up_args_queue')
|
||||
@patch.object(pool.TaskPool, 'is_open', new_callable=PropertyMock)
|
||||
async def test__map(self, mock_is_open: MagicMock, mock__set_up_args_queue: MagicMock,
|
||||
mock__queue_consumer: AsyncMock):
|
||||
mock__queue_consumer: AsyncMock, mock_event_cls: MagicMock):
|
||||
qsize = 4
|
||||
mock__set_up_args_queue.return_value = mock_q = MagicMock(qsize=MagicMock(return_value=qsize))
|
||||
mock_flag_set = MagicMock()
|
||||
mock_event_cls.return_value = mock_flag = MagicMock(set=mock_flag_set)
|
||||
|
||||
mock_func, stars = MagicMock(), 3
|
||||
args_iter, num_tasks = (FOO, BAR, 1, 2, 3), 2
|
||||
@ -539,14 +546,16 @@ class TaskPoolTestCase(CommonTestCase):
|
||||
mock_is_open.assert_called_once_with()
|
||||
mock__set_up_args_queue.assert_not_called()
|
||||
mock__queue_consumer.assert_not_awaited()
|
||||
mock_flag_set.assert_not_called()
|
||||
|
||||
mock_is_open.reset_mock()
|
||||
|
||||
mock_is_open.return_value = True
|
||||
self.assertIsNone(await self.task_pool._map(mock_func, args_iter, stars, num_tasks, end_cb, cancel_cb))
|
||||
mock__set_up_args_queue.assert_called_once_with(args_iter, num_tasks)
|
||||
mock__queue_consumer.assert_has_awaits(qsize * [call(mock_q, mock_func, arg_stars=stars,
|
||||
mock__queue_consumer.assert_has_awaits(qsize * [call(mock_q, mock_flag, mock_func, arg_stars=stars,
|
||||
end_callback=end_cb, cancel_callback=cancel_cb)])
|
||||
mock_flag_set.assert_called_once_with()
|
||||
|
||||
@patch.object(pool.TaskPool, '_map')
|
||||
async def test_map(self, mock__map: AsyncMock):
|
||||
@ -574,3 +583,91 @@ class TaskPoolTestCase(CommonTestCase):
|
||||
self.assertIsNone(await self.task_pool.doublestarmap(mock_func, kwargs_iter, num_tasks, end_cb, cancel_cb))
|
||||
mock__map.assert_awaited_once_with(mock_func, kwargs_iter, arg_stars=2, num_tasks=num_tasks,
|
||||
end_callback=end_cb, cancel_callback=cancel_cb)
|
||||
|
||||
|
||||
class SimpleTaskPoolTestCase(CommonTestCase):
|
||||
TEST_CLASS = pool.SimpleTaskPool
|
||||
task_pool: pool.SimpleTaskPool
|
||||
|
||||
TEST_POOL_FUNC = AsyncMock(__name__=FOO)
|
||||
TEST_POOL_ARGS = (FOO, BAR)
|
||||
TEST_POOL_KWARGS = {'a': 1, 'b': 2}
|
||||
TEST_POOL_END_CB = MagicMock()
|
||||
TEST_POOL_CANCEL_CB = MagicMock()
|
||||
|
||||
def get_task_pool_init_params(self) -> dict:
|
||||
return super().get_task_pool_init_params() | {
|
||||
'func': self.TEST_POOL_FUNC,
|
||||
'args': self.TEST_POOL_ARGS,
|
||||
'kwargs': self.TEST_POOL_KWARGS,
|
||||
'end_callback': self.TEST_POOL_END_CB,
|
||||
'cancel_callback': self.TEST_POOL_CANCEL_CB,
|
||||
}
|
||||
|
||||
def setUp(self) -> None:
|
||||
self.base_class_init_patcher = patch.object(pool.BaseTaskPool, '__init__')
|
||||
self.base_class_init = self.base_class_init_patcher.start()
|
||||
super().setUp()
|
||||
|
||||
def tearDown(self) -> None:
|
||||
self.base_class_init_patcher.stop()
|
||||
|
||||
def test_init(self):
|
||||
self.assertEqual(self.TEST_POOL_FUNC, self.task_pool._func)
|
||||
self.assertEqual(self.TEST_POOL_ARGS, self.task_pool._args)
|
||||
self.assertEqual(self.TEST_POOL_KWARGS, self.task_pool._kwargs)
|
||||
self.assertEqual(self.TEST_POOL_END_CB, self.task_pool._end_callback)
|
||||
self.assertEqual(self.TEST_POOL_CANCEL_CB, self.task_pool._cancel_callback)
|
||||
self.base_class_init.assert_called_once_with(pool_size=self.TEST_POOL_SIZE, name=self.TEST_POOL_NAME)
|
||||
|
||||
with self.assertRaises(exceptions.NotCoroutine):
|
||||
pool.SimpleTaskPool(MagicMock())
|
||||
|
||||
def test_func_name(self):
|
||||
self.assertEqual(self.TEST_POOL_FUNC.__name__, self.task_pool.func_name)
|
||||
|
||||
@patch.object(pool.SimpleTaskPool, '_start_task')
|
||||
async def test__start_one(self, mock__start_task: AsyncMock):
|
||||
mock__start_task.return_value = expected_output = 99
|
||||
self.task_pool._func = MagicMock(return_value=BAR)
|
||||
output = await self.task_pool._start_one()
|
||||
self.assertEqual(expected_output, output)
|
||||
self.task_pool._func.assert_called_once_with(*self.task_pool._args, **self.task_pool._kwargs)
|
||||
mock__start_task.assert_awaited_once_with(BAR, end_callback=self.task_pool._end_callback,
|
||||
cancel_callback=self.task_pool._cancel_callback)
|
||||
|
||||
@patch.object(pool.SimpleTaskPool, '_start_one')
|
||||
async def test_start(self, mock__start_one: AsyncMock):
|
||||
mock__start_one.return_value = FOO
|
||||
num = 5
|
||||
output = await self.task_pool.start(num)
|
||||
expected_output = num * [FOO]
|
||||
self.assertListEqual(expected_output, output)
|
||||
mock__start_one.assert_has_awaits(num * [call()])
|
||||
|
||||
@patch.object(pool.SimpleTaskPool, 'cancel')
|
||||
def test_stop(self, mock_cancel: MagicMock):
|
||||
num = 2
|
||||
id1, id2, id3 = 5, 6, 7
|
||||
self.task_pool._running = {id1: FOO, id2: BAR, id3: FOO + BAR}
|
||||
output = self.task_pool.stop(num)
|
||||
expected_output = [id3, id2]
|
||||
self.assertEqual(expected_output, output)
|
||||
mock_cancel.assert_called_once_with(*expected_output)
|
||||
mock_cancel.reset_mock()
|
||||
|
||||
num = 50
|
||||
output = self.task_pool.stop(num)
|
||||
expected_output = [id3, id2, id1]
|
||||
self.assertEqual(expected_output, output)
|
||||
mock_cancel.assert_called_once_with(*expected_output)
|
||||
|
||||
@patch.object(pool.SimpleTaskPool, 'num_running', new_callable=PropertyMock)
|
||||
@patch.object(pool.SimpleTaskPool, 'stop')
|
||||
def test_stop_all(self, mock_stop: MagicMock, mock_num_running: MagicMock):
|
||||
mock_num_running.return_value = num = 9876
|
||||
mock_stop.return_value = expected_output = 'something'
|
||||
output = self.task_pool.stop_all()
|
||||
self.assertEqual(expected_output, output)
|
||||
mock_num_running.assert_called_once_with()
|
||||
mock_stop.assert_called_once_with(num)
|
||||
|
156
usage/USAGE.md
156
usage/USAGE.md
@ -2,7 +2,7 @@
|
||||
|
||||
## Minimal example for `SimpleTaskPool`
|
||||
|
||||
The minimum required setup is a "worker" coroutine function that can do something asynchronously, a main coroutine function that sets up the `SimpleTaskPool` and starts/stops the tasks as desired, eventually awaiting them all.
|
||||
The minimum required setup is a "worker" coroutine function that can do something asynchronously, and a main coroutine function that sets up the `SimpleTaskPool`, starts/stops the tasks as desired, and eventually awaits them all.
|
||||
|
||||
The following demo code enables full log output first for additional clarity. It is complete and should work as is.
|
||||
|
||||
@ -77,6 +77,156 @@ did 4
|
||||
did 4
|
||||
```
|
||||
|
||||
## Advanced example
|
||||
## Advanced example for `TaskPool`
|
||||
|
||||
...
|
||||
This time, we want to start tasks from _different_ coroutine functions **and** with _different_ arguments. For this we need an instance of the more generalized `TaskPool` class.
|
||||
|
||||
As with the simple example, we need "worker" coroutine functions that can do something asynchronously, as well as a main coroutine function that sets up the pool, starts the tasks, and eventually awaits them.
|
||||
|
||||
The following demo code enables full log output first for additional clarity. It is complete and should work as is.
|
||||
|
||||
### Code
|
||||
```python
|
||||
import logging
|
||||
import asyncio
|
||||
|
||||
from asyncio_taskpool.pool import TaskPool
|
||||
|
||||
|
||||
logging.getLogger().setLevel(logging.NOTSET)
|
||||
logging.getLogger('asyncio_taskpool').addHandler(logging.StreamHandler())
|
||||
|
||||
|
||||
async def work(start: int, stop: int, step: int = 1) -> None:
|
||||
"""Pseudo-worker function counting through a range with a second of sleep in between each iteration."""
|
||||
for i in range(start, stop, step):
|
||||
await asyncio.sleep(1)
|
||||
print("work with", i)
|
||||
|
||||
|
||||
async def other_work(a: int, b: int) -> None:
|
||||
"""Different pseudo-worker counting through a range with half a second of sleep in between each iteration."""
|
||||
for i in range(a, b):
|
||||
await asyncio.sleep(0.5)
|
||||
print("other_work with", i)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
# Initialize a new task pool instance and limit its size to 3 tasks.
|
||||
pool = TaskPool(3)
|
||||
# Queue up two tasks (IDs 0 and 1) to run concurrently (with the same positional arguments).
|
||||
print("Called `apply`")
|
||||
await pool.apply(work, kwargs={'start': 100, 'stop': 200, 'step': 10}, num=2)
|
||||
# Let the tasks work for a bit.
|
||||
await asyncio.sleep(1.5)
|
||||
# Now, let us enqueue four more tasks (which will receive IDs 2, 3, 4, and 5), each created with different
|
||||
# positional arguments by using `starmap`, but have **no more than two of those** run concurrently.
|
||||
# Since we set our pool size to 3, and already have two tasks working within the pool,
|
||||
# only the first one of these will start immediately (and receive ID 2).
|
||||
# The second one will start (with ID 3), only once there is room in the pool,
|
||||
# which -- in this example -- will be the case after ID 2 ends;
|
||||
# until then the `starmap` method call **will block**!
|
||||
# Once there is room in the pool again, the third one will start immediately (and receive ID 4).
|
||||
# The last one will start (with ID 5) **only** once there is room in the pool **and** no more than one of these
|
||||
# last four tasks is running.
|
||||
args_list = [(0, 10), (10, 20), (20, 30), (30, 40)]
|
||||
print("Calling `starmap`...")
|
||||
await pool.starmap(other_work, args_list, num_tasks=2)
|
||||
print("`starmap` returned")
|
||||
# Now we close the pool, so that we can safely await all our tasks.
|
||||
pool.close()
|
||||
# Finally, we block, until all tasks have ended.
|
||||
print("Called `gather`")
|
||||
await pool.gather()
|
||||
print("Done.")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### Output
|
||||
Additional comments for the output are provided with `<---` next to the output lines.
|
||||
|
||||
(Keep in mind that the logger and `print` asynchronously write to `stdout`.)
|
||||
```
|
||||
TaskPool-0 initialized
|
||||
Started TaskPool-0_Task-0
|
||||
Started TaskPool-0_Task-1
|
||||
Called `apply`
|
||||
work with 100
|
||||
work with 100
|
||||
Calling `starmap`... <--- notice that this blocks as expected
|
||||
Started TaskPool-0_Task-2
|
||||
work with 110
|
||||
work with 110
|
||||
other_work with 0
|
||||
other_work with 1
|
||||
work with 120
|
||||
work with 120
|
||||
other_work with 2
|
||||
other_work with 3
|
||||
work with 130
|
||||
work with 130
|
||||
other_work with 4
|
||||
other_work with 5
|
||||
work with 140
|
||||
work with 140
|
||||
other_work with 6
|
||||
other_work with 7
|
||||
work with 150
|
||||
work with 150
|
||||
other_work with 8
|
||||
Ended TaskPool-0_Task-2 <--- here Task-2 makes room in the pool and unblocks `main()`
|
||||
TaskPool-0 is closed!
|
||||
Started TaskPool-0_Task-3
|
||||
other_work with 9
|
||||
`starmap` returned
|
||||
Called `gather`
|
||||
work with 160
|
||||
work with 160
|
||||
other_work with 10
|
||||
other_work with 11
|
||||
work with 170
|
||||
work with 170
|
||||
other_work with 12
|
||||
other_work with 13
|
||||
work with 180
|
||||
work with 180
|
||||
other_work with 14
|
||||
other_work with 15
|
||||
Ended TaskPool-0_Task-0
|
||||
Ended TaskPool-0_Task-1 <--- even though there is room in the pool now, Task-5 will not start
|
||||
Started TaskPool-0_Task-4
|
||||
work with 190
|
||||
work with 190
|
||||
other_work with 16
|
||||
other_work with 20
|
||||
other_work with 17
|
||||
other_work with 21
|
||||
other_work with 18
|
||||
other_work with 22
|
||||
other_work with 19
|
||||
Ended TaskPool-0_Task-3 <--- now that only Task-4 is left, Task-5 will start
|
||||
Started TaskPool-0_Task-5
|
||||
other_work with 23
|
||||
other_work with 30
|
||||
other_work with 24
|
||||
other_work with 31
|
||||
other_work with 25
|
||||
other_work with 32
|
||||
other_work with 26
|
||||
other_work with 33
|
||||
other_work with 27
|
||||
other_work with 34
|
||||
other_work with 28
|
||||
other_work with 35
|
||||
Ended TaskPool-0_Task-4
|
||||
other_work with 29
|
||||
other_work with 36
|
||||
other_work with 37
|
||||
other_work with 38
|
||||
other_work with 39
|
||||
Done.
|
||||
Ended TaskPool-0_Task-5
|
||||
```
|
||||
|
Reference in New Issue
Block a user