What is the cleanest way to write and run a DAG of tasks?

user_2536838 0
Viewed 1114

I want to write and run a directed acyclic graph (DAG) with several tasks running in serial or parallel. Ideally it would look like:

def task1():
    # ...

def task2():
    # ...

graph = Sequence([
    task1,
    task2,
    Parallel([
        task3,
        task4
    ]),
    task5
]

graph.run()

It would run 1 -> 2 -> (3 and 4 concurrently) -> 5. The tasks need to access the global scope to store results, write logs and access command line parameters.

My use case is writing a deployment script. Parallel tasks are IO-bound: typically waiting on a remote server to complete a step.

I looked into threading, asyncio, Airflow, but did not find any simple library that would allow this without some boilerplate code to traverse and control the graph's execution. Does anything like that exist?

directed-acyclic-graphs concurrency python
2 Answers

Here's a quick proof-of-concept implementation. It can be used like:

graph = sequence(
            lambda: print(1),
            lambda: print(2),
            parallel(
                lambda: print(3),
                lambda: print(4),
                sequence(
                    lambda: print(5),
                    lambda: print(6))),
             lambda: print(7)

graph()

1
2
3
5
6
4
7

sequence produces a function that wraps a for loop, and parallel produces a function that wraps use of a thread pool:

from typing import Callable
from multiprocessing.pool import ThreadPool

Task = Callable[[], None]

_pool: ThreadPool = ThreadPool()

def sequence(*tasks: Task) -> Task:
    def run():
        for task in tasks:
            task()

    return run  # Returning "run" to be used as a task by other "sequence" and "parallel" calls

def parallel(*tasks: Task) -> Task:
    def run():
        _pool.map(lambda f: f(), tasks)  # Delegate to a pool used for IO tasks

    return run

Each call to sequence and parallel returns a new "Task" (a function taking no arguments and returning nothing). That task can then be called by other, outer calls to sequence and parallel.

Things to note about the ThreadPool:

  • While this does use a thread pool for parallel, due to the GIL, this will still only execute one thing at a time. This means parallel is essentially useless for CPU-bound tasks.

  • I haven't specified how many threads the pool should begin with. I think it defaults to the number of cores you have available to you. You could specify how many you want to start with using the first parameter to ThreadPool if you want more.

  • For brevity, I'm not cleaning up the ThreadPool. You should definitely do that though if you use this.

  • Even though ThreadPool is a part of multiprocessing, confusingly it uses threads not processes.

user_30002060

You mentioned that yours tasks are IO bound, that means that asycnio would be a good candidate for this. You can try the aiodag library, which is an extremely light interface on top of asycnio that lets you easily define asynchronous dags:

import asyncio
from aiodag import task

@task
async def task1(x):
    ...

@task
async def task2(x):
    ...

@task
async def task3(x):
    ...

@task
async def task4(x):
    ...

@task
async def task5(x, y):
    ...

# rest of task funcs

async def main():
    t1 = task1()
    t2 = task2(t1)
    t3 = task3(t2)  # t3/t4 take t2, when t2 finishes, will run concurrently
    t4 = task4(t2)
    t5 = task5(t3, t4) # will wait until t3/t4 finish to execute
    await t5

loop = asyncio.new_event_loop()
asyncio.run_until_complete(main())

Check out the readme on the github page for aiodag for a bit of detail on how the dag is constructed/optimally executed. https://github.com/aa1371/aiodag

If you don't want to be tied to async functions, then check out dask's delayed interface. The definition of the dag works the same way as aiodag's, where the dag is constructed by function invocations. Dask will seamlessly handle executing your dag in the optimal parallel scheme, and can distribute over an arbitrarily large cluster to perform the parallel executions as well.

https://docs.dask.org/en/latest/delayed.html

user_54286680
Related