Python DAG: Retries, Rollbacks & Parallel Runs

Python DAG: Retries, Rollbacks & Parallel Runs

Enhance lightweight DAG with robust features for production-ready workflows

The previous article stopped at building a lightweight DAG with the Command pattern. That alone already made long functions more approachable and easier to test. Over the past week I kept iterating on the idea, because real systems rarely end with a single happy-path runner.

Three concerns show up almost immediately in production code: transient failures, partial completion, and throughput. They translate nicely into three extensions on top of the original runner: retry, rollback, and parallel execution.

Before diving into the new capabilities, here is the baseline runner from last week that simply walks the graph in topological order:

from graphlib import TopologicalSorter  
  
class DAGRunner:  
    def __init__(self, graph_definition: dict, commands: dict):  
        self.ts = TopologicalSorter(graph_definition)  
        self.commands = commands  
  
    def run(self):  
        context = {}  
        self.ts.prepare()  
        for task_name in self.ts.static_order():  
            self.commands[task_name].execute(context)  
        return context

This runner is the foundation for every extension in this article.

What matters here is that all three enhancements live inside the runner. The commands stay thin, the workflow definition stays declarative, and the abstraction keeps paying off. Let’s walk through each scenario.

Handling retries without polluting business logic

External integrations fail from time to time. The naive solution is to wrap every command in its own retry loop, but that instantly clutters the code we just spent effort to clean up. Moving retries into the runner keeps the commands focused while letting us tune policies centrally.

import time  
from graphlib import TopologicalSorter  
  
class DAGRunnerWithRetry:  
    def __init__(self, graph_definition: dict, commands: dict, *, max_retries: int = 3, retry_delay: float = 1.0):  
        self.ts = TopologicalSorter(graph_definition)  
        self.commands = commands  
        self.max_retries = max_retries  
        self.retry_delay = retry_delay  
  
    def run(self):  
        context = {}  
        self.ts.prepare()  
        for task_name in self.ts.static_order():  
            command = self.commands[task_name]  
            last_error = None  
            for attempt in range(1, self.max_retries + 1):  
                try:  
                    command.execute(context)  
                    last_error = None  
                    break  
                except Exception as exc:  
                    last_error = exc  
                    if attempt < self.max_retries:  
                        time.sleep(self.retry_delay)  
            if last_error is not None:  
                raise last_error  
        return context

The graph still decides the order, while the runner handles the repeated attempts. Any transient exception is retried in place; once the configured attempts are exhausted, the original error propagates.

Rolling back after a failed step

The next requirement is undoing work if a later step fails. The Command pattern already hinted at this with its symmetric interface, so introducing undo is a natural extension. The runner simply needs to record the successful commands and process them in reverse when an error bubbles up.

from graphlib import TopologicalSorter  
  
class DAGRunnerWithRollback:  
    def __init__(self, graph_definition: dict, commands: dict):  
        self.ts = TopologicalSorter(graph_definition)  
        self.commands = commands  
  
    def run(self):  
        context = {}  
        successful_commands = []  
        try:  
            self.ts.prepare()  
            for task_name in self.ts.static_order():  
                command = self.commands[task_name]  
                command.execute(context)  
                successful_commands.append(command)  
        except Exception as exc:  
            for command in reversed(successful_commands):  
                command.undo(context)  
            raise exc  
        return context

Each command now implements both execute and undo. When a later task fails, the runner cleans up previously completed work in last-in-first-out order, mirroring the graph execution.

Unlocking parallel branches

Many workflows contain independent branches that can safely run side by side once their prerequisites finish. The original runner executes nodes sequentially even when the graph exposes available concurrency. Switching to a thread pool allows the runner to schedule every ready node without rewriting the business logic.

from concurrent.futures import ThreadPoolExecutor, as_completed  
from graphlib import TopologicalSorter  
  
class DAGRunnerWithParallel:  
    def __init__(self, graph_definition: dict, commands: dict):  
        self.ts = TopologicalSorter(graph_definition)  
        self.commands = commands  
  
    def run(self):  
        context = {}  
        self.ts.prepare()  
        with ThreadPoolExecutor() as executor:  
            futures = {}  
            while self.ts.is_active():  
                for task_name in self.ts.get_ready():  
                    command = self.commands[task_name]  
                    futures[executor.submit(command.execute, context)] = task_name  
                for future in as_completed(list(futures)):  
                    task_name = futures.pop(future)  
                    try:  
                        future.result()  
                        self.ts.done(task_name)  
                    except Exception:  
                        for pending in futures:  
                            pending.cancel()  
                        raise  
        return context

The runner keeps asking the topological sorter for ready nodes, submits them to the executor, and marks each task as complete once the future finishes. If any parallel branch fails, the remaining futures are cancelled and the error is surfaced immediately.

Wrap Up

All three extensions reuse the same two abstractions introduced earlier: the DAG encodes dependency, and Command encapsulates each operation behind a stable interface. Because the runner owns orchestration, we can iteratively add capabilities such as retry strategies, rollback semantics, and parallel scheduling without rewriting domain logic.

That’s the real payoff of the abstraction. Once the seams are in place, new requirements stay localized, the codebase remains comprehensible, and future experiments can continue to build on the same foundation.

A message from our Founder

Hey, Sunil here. I wanted to take a moment to thank you for reading until the end and for being a part of this community.

Did you know that our team run these publications as a volunteer effort to over 3.5m monthly readers? We don’t receive any funding, we do this to support the community. ❤️

If you want to show some love, please take a moment to follow me on LinkedIn, TikTok, Instagram. You can also subscribe to our weekly newsletter.

And before you go, don’t forget to clap and follow the writer️!

Originally published on Medium