Problem 88 · Concurrency & Python

Parallel DAG Task Executor

HardClassic DSA & systems

01Problem

Given tasks and their dependencies, run every task once its dependencies finish, executing independent tasks in parallel and passing upstream results as arguments; detect cycles.

02Key insight

Kahn's algorithm driven by completion events: submit all in-degree-0 tasks, and on each completion decrement children and submit newly-ready ones. A cycle shows up as tasks that never become ready. Threads suit IO-bound tasks; use processes for CPU-bound ones because of the GIL.

03Solution

solution.pypython43 lines

from collections import defaultdict
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
from typing import Any, Callable


def run_dag(tasks: dict[str, Callable[..., Any]], deps: dict[str, list[str]],
            max_workers: int = 4) -> dict[str, Any]:
    indegree = {name: 0 for name in tasks}
    children: dict[str, list[str]] = defaultdict(list)
    for name, parents in deps.items():
        for parent in parents:
            if parent not in tasks:
                raise KeyError(f"unknown dependency {parent!r}")
            indegree[name] += 1
            children[parent].append(name)

    results: dict[str, Any] = {}
    pool = ThreadPoolExecutor(max_workers=max_workers)

    def submit(name: str) -> Future:
        args = [results[p] for p in deps.get(name, [])]
        return pool.submit(tasks[name], *args)

    try:
        running = {submit(n): n for n, d in indegree.items() if d == 0}
        while running:
            done, _ = wait(running, return_when=FIRST_COMPLETED)
            for fut in done:
                name = running.pop(fut)
                results[name] = fut.result()  # re-raises the task's error
                for child in children[name]:
                    indegree[child] -= 1
                    if indegree[child] == 0:
                        running[submit(child)] = child
    except BaseException:
        # Fail fast: cancel queued tasks and return immediately; threads can't be
        # interrupted, so already-running tasks finish in the background.
        pool.shutdown(wait=False, cancel_futures=True)
        raise
    pool.shutdown()
    if len(results) != len(tasks):
        raise ValueError(f"cycle among {sorted(set(tasks) - set(results))}")
    return results

Try it first.

TimeO(V+E) scheduling overhead

SpaceO(V+E)

tests.py · 4 passing

tests.pypython20 lines

import threading
barrier = threading.Barrier(2, timeout=2)  # only passes if a and b run concurrently
def a(): barrier.wait(); return 2
def b(): barrier.wait(); return 3
tasks = {"a": a, "b": b, "c": lambda x, y: x * y, "d": lambda z: z + 1}
res = run_dag(tasks, {"c": ["a", "b"], "d": ["c"]})
assert res == {"a": 2, "b": 3, "c": 6, "d": 7}
try:
    run_dag({"x": lambda v: v, "y": lambda v: v, "z": lambda: 0}, {"x": ["y"], "y": ["x"]})
    assert False
except ValueError as e:
    assert "x" in str(e) and "y" in str(e)
import time
def boom(): raise RuntimeError("boom")
def slow(): time.sleep(0.5); return 1
t0 = time.perf_counter()
try:
    run_dag({"a": boom, "b": slow, "c": lambda v: v}, {"c": ["b"]}); assert False
except RuntimeError:
    assert time.perf_counter() - t0 < 0.3  # did not wait for the slow sibling

04Staff-level follow-up

How would you make this resumable like Airflow? Persist task state (pending/running/succeeded/failed) and outputs keyed by run_id, make tasks idempotent, and on restart rebuild in-degrees counting only unfinished parents; add per-task retries and timeouts.