8
\$\begingroup\$

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:

  1. Are there known packages for my use case?
  2. 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
\$\endgroup\$
4
  • 3
    \$\begingroup\$ The first thing that came to my eye is your import of 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\$ Commented Feb 25, 2022 at 21:58
  • 2
    \$\begingroup\$ IO-bound or CPU-bound? \$\endgroup\$
    – Reinderien
    Commented Feb 25, 2022 at 22:14
  • 1
    \$\begingroup\$ tasks are io-bound, main thread will be cpu-bound \$\endgroup\$
    – ROMSCore
    Commented Feb 25, 2022 at 22:22
  • \$\begingroup\$ I might be wrong as I'm not an expert, but Apache Airflow seems to be a potential tool for this. I've been playing with it on my local machine. You can create a graph of tasks (DAGs) with a Python code, and you can see all sort of summaries in a browser. You can even watch on the graph view all the task states. There are several executors including celery, kubernetes. And there are managed versions on cloud. airflow.apache.org/docs/apache-airflow/stable/index.html \$\endgroup\$ Commented Jul 7, 2023 at 7:54

1 Answer 1

1
\$\begingroup\$

are there known packages for my use case?

Yes.

Use 0mq, RabbitMQ, kafka, or similar pub/sub platform, to address horizontal scaling issues.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.