For my application, I needed a lightweight way to launch multiple tasks and monitor their completion in a main loop periodically or use callbacks. The main loop is going to be busy most of the time performing analysis after receiving inputs from tasks. I was writing a lot of boilerplate code for asyncio each time and decided to write this multi-threaded queue/dispatcher outlined below. New to Python and still learning the ropes, so looking for two things:
- Are there known packages for my use case?
- Could I run into thread locks with the approach below for the intended scale (2-3 threads, 20-30 tasks each)?
Dispatcher package
import asyncio
import threading
import time
from datetime import datetime
from threading import Thread, Lock
from SeTools.Printer import mprint
"""
Task Ion is a holder for the task and attributes such as coro, callback, and timeout associated with it.
"""
class TaskIon:
def __init__(self, name, taskCoro, *args, **kwargs):
"""
Initialize a taskIon which is then used to invoke an asyncio Task
:param name: name of task; this is the primary key you will need to cancel or fetch task results
:param taskCoro: the coroutine to be called
:param args: args for the coroutine
:param kwargs: supported values are:
"callback" - the routine to call after tasks completed or is cancelled. Passes the task object to callback
"timeout" - timeout in seconds for the Task
"wrappedCall" - when set to True, the coro is called using executor. Useful when you don't want a synchronous task to block
"""
self.name = name
self.coro = taskCoro
self.coroArgs = args
self.callback = kwargs["callback"] if "callback" in kwargs else None #callbackRoutine
self.timeout = kwargs["timeout"] if "timeout" in kwargs else None
self.wrappedCall = kwargs["wrappedCall"] if "wrappedCall" in kwargs else False
"""
A workerIon is an independent asyncio loop running any number of tasks and storing results
"""
class workerIon:
def _init_caller(self, future, ion):
pass
def sleepThread(self, sleepTime):
"""
Ask the worker to make thread sleep for time
:param sleepTime:
:return:
"""
self.loopStatus = "paused"
time.sleep(sleepTime)
self.loopStatus = "started"
def __init__(self, name, taskRequests: list[TaskIon]):
self.loopStartTime: datetime = None
self.loopEndTime: datetime = None
self.name = name
self.results = {}
self.tasks: dict[str, asyncio.Task] = {}
self.loopHandle = asyncio.new_event_loop()
self.taskStatus = {}
self.loopStatus = "initializing"
self._taskIons = taskRequests
for t in taskRequests: self.taskStatus[t] = "not started"
def startLoop(self):
self.loopHandle.run_until_complete(self._loop(self._taskIons))
# self._loop(self._taskIons)
async def _wrappedTask(self, coro, *args):
return await self.loopHandle.run_in_executor(None, coro, *args)
async def _loop(self, rTasks: list[TaskIon]):
self.loopStartTime = datetime.now()
self.loopStatus = "started"
mprint(f"Dispatcher-worker : {self.name} loop started at time: {self.loopStartTime}")
rtask: TaskIon
for rtask in rTasks:
if Dispatcher.taskExists(rtask.name):
raise KeyError(f"Global task tracker: {rtask.name} already started as task and hasn't been fetched")
if rtask.wrappedCall:
t = self.loopHandle.create_task(self._wrappedTask(rtask.coro, *rtask.coroArgs), name=rtask.name)
else:
t = self.loopHandle.create_task(rtask.coro(*rtask.coroArgs), name=rtask.name)
self.taskStatus[rtask.name] = "started"
if rtask.callback is not None:
t.add_done_callback(rtask.callback)
self.tasks[rtask.name] = t
await Dispatcher.addTaskAsync(rtask.name, t)
while True:
try:
done, pending = await asyncio.wait(self.tasks.values(), return_when=asyncio.FIRST_COMPLETED)
for t in done:
tName = t.get_name()
self.results[tName] = t.result()
self.taskStatus[tName] = "completed"
except asyncio.CancelledError:
self.taskStatus[tName] = "cancelled"
finally:
self.tasks.pop(t.get_name())
if len(self.tasks) <= 0:
break
self.loopStatus = "completed"
self.loopEndTime = datetime.now()
self.loopHandle.stop()
mprint(f"Dispatcher-worker : {self.name} loop ending at time: {self.loopEndTime}")
Dispatcher.unregisterThread(self)
"""
Dispatcher allows you to launch many concurrent tasks synchronously or asynchronously
"""
class Dispatcher:
_tasks: dict[str, asyncio.Task] = {}
_workers: dict[str, workerIon] = {}
_threads: dict[str, Thread] = {}
_asyncioLock = asyncio.Lock()
_threadLock = Lock()
@staticmethod
def unregisterThread(worker: workerIon):
"""
Immidiately
:param worker:
:return:
"""
if worker.name not in Dispatcher._threads:
raise KeyError(f"Thread {worker.name} not registered")
else:
with Dispatcher._threadLock:
Dispatcher._threads.pop(worker.name)
mprint(f"Thread {worker.name} unregistered after completion")
@staticmethod
def isTaskDone(name, silent=False):
if name not in Dispatcher._tasks:
if not silent:
raise KeyError(f"Task with name: {name} not registered in tasks")
else:
return False
return Dispatcher._tasks[name].done()
@staticmethod
def taskExists(name) -> bool:
"""
Returns whether a task was registered. A task doesn't leave a queue until the results are fetched.
:param name: name of task to search
:return: bool
"""
return name in Dispatcher._tasks
@staticmethod
async def addTaskAsync(name, task) -> None:
"""
Adds a task to global tracker. Thread safe
:param name: name of task to add. A task doesn't leave a queue until the results are fetched.
:param task: the task object created inside the worker
:return: None
"""
with Dispatcher._threadLock:
async with Dispatcher._asyncioLock:
Dispatcher._tasks[name] = task
# mprint(f"task added: {task.get_name()}")
@staticmethod
def addWorker(name, worker) -> None:
"""
Add a worked to queue. A worker does not leave queue until all tasks in worker are finished and results fetched
:param name: Name of worker to add
:param worker: the workerIon object
:return: None
"""
if name in Dispatcher._workers:
raise KeyError(
f"Worker with {name} already added. Try fetching all results to remove from queue or change name")
Dispatcher._workers[name] = worker
@staticmethod
def launchTasks(tasks: list[TaskIon], name=None) -> None:
"""
Create a new workerIon and start one or more tasks in parallel. Launch tasks is synchronous and will block
execution until all tasks complete
:param tasks:
:param name: name of workerIon
:return: None
"""
name = datetime.now().timestamp() if name is None else name
w = workerIon(name, tasks)
for rtask in tasks:
Dispatcher.addWorker(rtask.name, w)
w.startLoop()
@staticmethod
async def addWorkerAsync(name, worker):
async with Dispatcher._asyncioLock:
Dispatcher._workers[name] = worker
@staticmethod
def launchTasksAsync(tasks: list[TaskIon], name=None) -> str:
"""
:param tasks:
:param name:
:return:
"""
name = datetime.now().timestamp() if name is None else name
with Dispatcher._threadLock:
if name in Dispatcher._threads:
raise KeyError(
f"Thread with {name} is already in queue, and alive status: {Dispatcher._threads[name].is_alive()}")
else:
Dispatcher._threads[name] = Thread(target=Dispatcher.launchTasks, name=name, args=(tasks, name))
Dispatcher._threads[name].start()
time.sleep(0.1)
return name
@staticmethod
def wait(name=None) -> None:
"""
:param name:
:return:
"""
if name is not None:
if name not in Dispatcher._threads:
return None
else:
return Dispatcher._threads[name].join()
else:
for t in threading.enumerate():
if t.name in Dispatcher._threads: # "MainThread":
mprint(f"Thread {t.name} will be waited")
t.join()
# Dispatcher._threads[t].join()
return None
@staticmethod
def popTaskResult(taskName, silent=False):
res = Dispatcher.getTaskResult(taskName, silent)
if Dispatcher.taskStatus(taskName, silent) in ("completed", "cancelled"):
Dispatcher.cleanupTask(taskName, silent=silent)
return res
@staticmethod
def getTaskResult(taskName, silent=False):
"""
:param taskName:
:param silent:
:return:
"""
if taskName not in Dispatcher._tasks:
if not silent:
raise KeyError(f"Task with name: {taskName} not registered in tasks")
else:
return None
if not Dispatcher.isTaskDone(taskName):
if not silent:
raise Exception(f"Task name: {taskName} has not finished")
else:
return None
if taskName not in Dispatcher._workers:
if not silent:
raise Exception(f"Worker name: {taskName} has not been registered yet")
else:
return None
if taskName not in Dispatcher._workers[taskName].results:
if not silent:
raise Exception(f"Task name: {taskName} has not registered results yet")
else:
return None
result = Dispatcher._workers[taskName].results[taskName]
# Dispatcher.cleanup(taskName)
# Dispatcher._tasks.pop(taskName)
return result
@staticmethod
def taskNamesForWorker(workerName) -> [str]:
"""
:param workerName:
:return:
"""
names = []
for key, _w in Dispatcher._workers.items():
if _w.name == workerName:
names.append(key)
return names
@staticmethod
def taskStatus(taskName, silent=False) -> str:
"""
:param taskName:
:param silent:
:return:
"""
if taskName not in Dispatcher._workers or taskName not in Dispatcher._workers[taskName].taskStatus:
if not silent:
raise KeyError(f"Task {taskName} not registered yet")
else:
return "not registered"
else:
return Dispatcher._workers[taskName].taskStatus[taskName]
@staticmethod
def cleanupTask(taskName=None, force=False, silent=False):
if taskName is not None:
with Dispatcher._threadLock:
task = Dispatcher._tasks[taskName]
if not task.done():
if not force:
if silent:
return None
else:
raise Exception(f"Task: {taskName} is in progress")
else:
# mprint(f"Cancelling due to invoke of cleanup with force for task: {taskName}")
Dispatcher._workers[taskName].loopHandle.call_soon_threadsafe(task.cancel)
time.sleep(0.01)
# Now that task is done we can cleanup the task store and result store
worker = Dispatcher._workers[taskName]
if taskName in worker.results: _ = worker.results.pop(taskName) # cleanup results
# if taskName in worker.tasks: _ = worker.tasks.pop(taskName) # The worker will cleanup the task
if taskName in Dispatcher._tasks: _ = Dispatcher._tasks.pop(taskName)
# mprint(f"task removed: {taskName}")
if taskName in Dispatcher._workers: _ = Dispatcher._workers.pop(
taskName) # cleanup the task to worker mapping table
# Check if last task in thread and release thread if so.
taskNamesLeft = Dispatcher.taskNamesForWorker(worker.name)
if len(taskNamesLeft) <= 0:
mprint(f"No more tasks left in worker {taskNamesLeft}")
else: # cleanup all tasks
tNames = list(Dispatcher._tasks.keys())
for tName in tNames:
mprint(f"Cleaning up task: {tName}, force={force}, silent={silent}")
Dispatcher.cleanupTask(tName, force, silent)
@staticmethod
def cleanupThread(threadName=None, force=False, silent=False):
if threadName is not None:
with Dispatcher._threadLock:
thread = Dispatcher._threads[threadName]
if not thread.is_alive():
if not force:
if silent:
return None
else:
raise Exception(f"Thread: {thread.name} is in progress")
else:
for wTask in Dispatcher.taskNamesForWorker(threadName):
Dispatcher.cleanupTask(wTask, force, silent)
time.sleep(0.01)
if threadName in Dispatcher._threads: _ = Dispatcher._threads.pop(threadName)
else: # cleanup all tasks
tNames = list(Dispatcher._threads.keys())
for tName in tNames:
Dispatcher.cleanupThread(tName, force, silent)
@staticmethod
def printThreads():
mprint(f"Threads: {Dispatcher._threads.keys()}")
Testcase
from Dispatcher import *
from SeTools.Printer import mprint
# Testing the file directly
if __name__ == '__main__':
async def printHelloAndSleep(name, j):
try:
await asyncio.sleep(j)
return f"Result -- TaskName: {name}: -- sleep-time: {j} -- end-time: {datetime.now()}"
except asyncio.exceptions.CancelledError:
mprint(f"Task {name} cancelled, printed from coro")
raise
def callbackTest(future: asyncio.Task):
if not future.cancelled():
mprint(f"Another way to get results: {future.result()}")
else:
mprint(f"Task {future.get_name()} cancelled, printed from callback")
mprint(f"Main loop start time: {datetime.now()}")
mprint(f"Task1 status: {Dispatcher.taskStatus('task1', silent=True)}")
nameT1 = Dispatcher.launchTasksAsync([
TaskIon("task1", printHelloAndSleep("Task1", 4)),
TaskIon("task2", printHelloAndSleep("Task2", 200), callbackTest)])
nameT2 = Dispatcher.launchTasksAsync([
TaskIon("task3", printHelloAndSleep("Task3", 10)),
TaskIon("task4", printHelloAndSleep("Task4", 11))])
mprint(f"Task1 status: {Dispatcher.taskStatus('task1', silent=True)}")
mprint(f"Main loop tasks launched: {datetime.now()}")
time.sleep(1)
Dispatcher.cleanupTask("task2", force=True)
mprint(f"Task1 status: {Dispatcher.taskStatus('task1', silent=True)}")
Dispatcher.printThreads()
while True:
x = Dispatcher.popTaskResult("task1", silent=True)
if x is not None:
mprint(x)
break
mprint(f"Task1 status: {Dispatcher.taskStatus('task1', silent=True)}")
Dispatcher.printThreads()
time.sleep(2)
Dispatcher.printThreads()
Dispatcher.wait()
Dispatcher.printThreads()
Output
[Info] -- [Dispatcher_copy.py, 21] -- Main loop start time: 2022-02-25 13:33:42.676244
[Info] -- [Dispatcher_copy.py, 22] -- Task1 status: not registered
[Info] -- [Dispatcher.py, 58] -- Dispatcher-worker : 1645824822.760012 loop started at time: 2022-02-25 13:33:42.768744
[Info] -- [Dispatcher.py, 58] -- Dispatcher-worker : 1645824822.868929 loop started at time: 2022-02-25 13:33:42.870863
[Info] -- [Dispatcher_copy.py, 30] -- Task1 status: started
[Info] -- [Dispatcher_copy.py, 31] -- Main loop tasks launched: 2022-02-25 13:33:42.984696
[Info] -- [Dispatcher_copy.py, 10] -- Task Task2 cancelled, printed from coro
[Info] -- [Dispatcher_copy.py, 35] -- Task1 status: started
[Info] -- [Dispatcher.py, 348] -- Threads: dict_keys([1645824822.760012, 1645824822.868929])
[Info] -- [Dispatcher_copy.py, 18] -- Task task2 cancelled, printed from callback
[Info] -- [Dispatcher.py, 317] -- No more tasks left in worker []
[Info] -- [Dispatcher_copy.py, 40] -- Result -- TaskName: Task1: -- sleep-time: 4 -- end-time: 2022-02-25 13:33:46.928457
[Info] -- [Dispatcher_copy.py, 42] -- Task1 status: not registered
[Info] -- [Dispatcher.py, 88] -- Dispatcher-worker : 1645824822.760012 loop ending at time: 2022-02-25 13:33:46.928457
[Info] -- [Dispatcher.py, 348] -- Threads: dict_keys([1645824822.760012, 1645824822.868929])
[Info] -- [Dispatcher.py, 116] -- Thread 1645824822.760012 unregistered after completion
[Info] -- [Dispatcher.py, 348] -- Threads: dict_keys([1645824822.868929])
[Info] -- [Dispatcher.py, 88] -- Dispatcher-worker : 1645824822.868929 loop ending at time: 2022-02-25 13:33:53.919234
[Info] -- [Dispatcher.py, 116] -- Thread 1645824822.868929 unregistered after completion
[Info] -- [Dispatcher.py, 348] -- Threads: dict_keys([])
Process finished with exit code 0
Lock
. If you want concurrency, you don't want locks. If you think you need locks, you should redesign your code. I may come back to you with a detailed review if I find the time, since writing concurrent code is running with scissors. \$\endgroup\$