Cascading Cancellation of Dependent Tasks

Hello! Welcome to the next lesson in your preparation for the Anthropic software engineering interview.

In our last session, we built a fully operational DAGScheduler capable of executing a graph of tasks according to their priorities and dependencies. We now have a solid foundation for orchestrating complex, multi-step workflows.

Introduction

Today, we will add a critical feature found in any production-grade workflow system: cascading cancellation. Our goal is to implement a cancel_task method that not only stops a specific task but also propagates that cancellation signal to all of its direct and indirect dependents.

This lesson covers:

• The algorithm for traversing a dependency graph to identify all downstream tasks.
• The distinction between forcible and cooperative cancellation.
• Modifying the scheduler's worker logic to respect cancellation signals, preventing unnecessary work.

This is a problem you've likely navigated when building systems like Rudi at AWS. In Tokio, for instance, calling abort() on a JoinHandle cancels that specific task, but it doesn't automatically propagate to other tasks it may have spawned. This "detach on drop" behavior means that implementing robust, cascading cancellation requires a deliberate, structured approach—exactly what we will build today.

---

1. The Cancellation Strategy: Graph Traversal

To implement cascading cancellation, we must first identify every task that is affected. If a user cancels a task, any other task that depends on it (or depends on one of its dependents, and so on) becomes moot.

This requires a graph traversal, starting from the cancelled task and moving "downstream" to all its dependents.

• Starting Point: The task_id passed to cancel_task.
• Traversal Rule: From any given task, we find its immediate dependents using our _dependents dictionary.
• Algorithm: We can use a simple Breadth-First Search (BFS) or Depth-First Search (DFS) to explore this part of the graph. All tasks visited during this traversal are marked for cancellation.

Imagine our task graph from the previous lesson. If we cancel the clean task, the cancellation must propagate to process_A, process_B, aggregate, and report.

---

2. Semantics of Cancellation in a Concurrent System

Once we've identified which tasks to cancel, how do we actually stop them? This isn't as simple as just removing them from the graph, especially in a concurrent system.

• Task in the ready queue: If a task is waiting in the _ready_queue, we could try to remove it. However, this is inefficient and complex with Python's PriorityQueue. A much simpler and more robust pattern is to let the worker pull the task, see that it's cancelled, and simply discard it.
• Task currently running: In Python's threading model, you cannot safely "kill" a thread from the outside. The running task must be designed to cooperate by periodically checking a cancellation flag and terminating itself gracefully.
• Task already completed: The cancellation request arrives too late. We do nothing.

We will implement a cooperative cancellation model. Our _worker threads will be responsible for checking if a task has been cancelled just before executing it.

This pattern is universal in modern concurrency frameworks. Your experience with Rust's async ecosystem provides a direct parallel.

This post reinforces that what we're building is a form of structured concurrency, ensuring that when one part of a computation is cancelled, the entire dependent subgraph is correctly torn down.

---

3. Your Task: Implement Cascading Cancellation

Now, let's integrate this feature into our DAGScheduler. You will need to:

1. Add a _cancelled_tasks set to track the IDs of all cancelled tasks.
2. Implement the cancel_task(task_id) method, which will perform the graph traversal to populate this set.
3. Modify the _worker method to check the _cancelled_tasks set before executing a task.

Below is the scheduler code from our last lesson. Fill in the new logic where indicated.

import threading
import queue
import time
import itertools
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import Callable, List, Set




# Using the same Task class from the previous lesson
@dataclass(order=True)
class Task:
    priority: int
    task_id: int = field(compare=False)
    func: Callable = field(compare=False)
    args: tuple = field(compare=False, default_factory=tuple)
    kwargs: dict = field(compare=False, default_factory=dict)

class DAGScheduler:
    def __init__(self, num_workers: int = 4):
        self._tasks = {}
        self._dependents = defaultdict(set)
        self._dependencies = defaultdict(set)
        self._in_degree = defaultdict(int)
        
        self._ready_queue = queue.PriorityQueue()
        self._workers = []
        self._num_workers = num_workers
        self._stop_event = threading.Event()

        self._task_id_counter = itertools.count()
        self._lock = threading.Lock()




        # --- YOUR IMPLEMENTATION START (1) ---
        # Add a set to store the IDs of cancelled tasks.
        self._cancelled_tasks: Set[int] = set()



        # --- YOUR IMPLEMENTATION END (1) ---

    def add_task(self, func: Callable, priority: int, dependencies: List[int] = None, *args, **kwargs) -> int:



        # This method remains unchanged from the previous lesson.
        if dependencies is None:
            dependencies = []

        task_id = next(self._task_id_counter)
        task = Task(priority=priority, task_id=task_id, func=func, args=args, kwargs=kwargs)

        with self._lock:
            if any(dep_id not in self._tasks for dep_id in dependencies):
                raise ValueError("One or more dependencies do not exist.")
            
            self._tasks[task_id] = task
            self._dependencies[task_id] = set(dependencies)
            
            for dep_id in dependencies:
                self._dependents[dep_id].add(task_id)
            
            self._in_degree[task_id] = len(dependencies)
        
        return task_id




    # --- YOUR IMPLEMENTATION START (2) ---
    def cancel_task(self, task_id: int):
        """
        Cancels a task and all of its downstream dependents.
        """
        if task_id not in self._tasks:
            print(f"Warning: Task {task_id} not found for cancellation.")
            return

        print(f"Initiating cancellation for task {task_id} and its dependents...")
        



        # Use a queue for BFS traversal to find all dependents.
        tasks_to_cancel = deque([task_id])
        all_cancelled = set()

        with self._lock:
            while tasks_to_cancel:
                current_task_id = tasks_to_cancel.popleft()
                if current_task_id in all_cancelled:
                    continue
                
                all_cancelled.add(current_task_id)
                



                # Add all direct dependents to the queue for traversal.
                for dependent_id in self._dependents.get(current_task_id, set()):
                    tasks_to_cancel.append(dependent_id)
            



            # Atomically update the main cancelled set.
            self._cancelled_tasks.update(all_cancelled)
            print(f"Tasks marked for cancellation: {all_cancelled}")



    # --- YOUR IMPLEMENTATION END (2) ---

    def start(self):



        # This method remains mostly unchanged.
        with self._lock:
            initial_ready_tasks = [
                task_id for task_id, degree in self._in_degree.items() if degree == 0
            ]
            for task_id in initial_ready_tasks:
                self._ready_queue.put(self._tasks[task_id])

        for i in range(self._num_workers):
            worker = threading.Thread(target=self._worker, args=(i,))
            self._workers.append(worker)
            worker.start()
        print(f"{self._num_workers} worker threads started.")

    def _worker(self, worker_id: int):
        while not self._stop_event.is_set():
            try:
                task = self._ready_queue.get(timeout=1)




                # --- YOUR IMPLEMENTATION START (3) ---
                # Check if the task has been cancelled.
                with self._lock:
                    if task.task_id in self._cancelled_tasks:
                        print(f"Worker {worker_id}: Skipping cancelled task {task.task_id}")



                        # Even though the task is cancelled, we must treat it as "complete"
                        # to unblock the rest of the graph. The logic below will handle this.
                        # We just skip the execution of task.func().
                        pass
                    else:



                        # If not cancelled, execute the task.
                        print(f"Worker {worker_id}: Starting task {task.task_id} (Priority: {task.priority})")
                        task.func(*task.args, **task.kwargs)
                        print(f"Worker {worker_id}: Finished task {task.task_id}")



                # --- YOUR IMPLEMENTATION END (3) ---




                # This logic now runs for both completed and cancelled tasks,
                # ensuring the graph state is always updated correctly.
                with self._lock:
                    dependents = self._dependents.get(task.task_id, set())
                    
                    for dep_id in dependents:
                        self._in_degree[dep_id] -= 1
                        
                        if self._in_degree[dep_id] == 0:



                            # Add a final check: don't enqueue a task that's already cancelled.
                            if dep_id not in self._cancelled_tasks:
                                print(f"Task {dep_id} is now ready. Adding to queue.")
                                self._ready_queue.put(self._tasks[dep_id])
                            else:
                                print(f"Task {dep_id} is ready but was already cancelled. Not adding to queue.")


            except queue.Empty:
                continue
        
        print(f"Worker {worker_id} stopping.")

    def stop(self):
        print("Scheduler stopping...")
        self._stop_event.set()
        for worker in self._workers:
            worker.join()
        print("All workers have stopped.")




# --- Example Usage ---
def sample_task(duration, name):
    print(f"  -> Running task '{name}' for {duration}s...")
    time.sleep(duration)
    print(f"  -> Task '{name}' complete.")

if __name__ == '__main__':
    scheduler = DAGScheduler(num_workers=4)

    print("Building the task graph...")



    #      t_clean(P2) -> t_proc_a(P3)
    #    /                            \
    # t_fetch(P1)                       -> t_agg(P4) -> t_report(P5)
    #    \                            /
    #      ------------------------> t_proc_b(P3)
    t_fetch = scheduler.add_task(sample_task, priority=1, duration=1, name="fetch")
    t_clean = scheduler.add_task(sample_task, priority=2, dependencies=[t_fetch], duration=2, name="clean")
    t_proc_a = scheduler.add_task(sample_task, priority=3, dependencies=[t_clean], duration=4, name="process_A")
    t_proc_b = scheduler.add_task(sample_task, priority=3, dependencies=[t_fetch], duration=1, name="process_B")
    t_agg = scheduler.add_task(sample_task, priority=4, dependencies=[t_proc_a, t_proc_b], duration=1, name="aggregate")
    t_report = scheduler.add_task(sample_task, priority=5, dependencies=[t_agg], duration=1, name="report")

    scheduler.start()




    # Let the scheduler run for a bit, then cancel a task.
    time.sleep(2) # Let fetch and clean start/finish
    scheduler.cancel_task(t_proc_a) # Cancel process_A

    while not scheduler._ready_queue.empty() or any(thread.is_alive() for thread in scheduler._workers):
        try:
            time.sleep(1)



            # A simple condition to break the loop once the queue is empty and workers might be idle.
            if scheduler._ready_queue.empty():



                # Give workers a moment to finish their current task.
                time.sleep(2)
                if scheduler._ready_queue.empty():
                    break
        except KeyboardInterrupt:
            break

    scheduler.stop()

When you run the example, observe the output. The fetch task runs. After it completes, both clean and process_B become ready. While they are running, we call cancel_task on process_A. Because process_A is cancelled, its dependents (aggregate and report) will also be marked for cancellation and will be skipped by the workers. However, process_B should complete successfully.

---

Conclusion

Fantastic work! You have now equipped the DAGScheduler with a robust, concurrent-safe, cascading cancellation mechanism. This is a non-trivial feature that demonstrates a deep understanding of how to manage complex state in a multi-threaded application.

Key Takeaways:

• Cascading cancellation requires a "downstream" graph traversal from the point of cancellation to identify all affected dependent tasks.
• In multi-threaded Python, cooperative cancellation is the standard, safe pattern. Workers must check for a cancellation signal before starting a unit of work.
• By treating a cancelled task as "completed" (for the purpose of dependency resolution), we allow the scheduler to correctly prune the cancelled branch while letting unrelated branches of the graph execute.

This in-memory implementation directly maps to distributed patterns. In a system like those at AWS, the _cancelled_tasks set would be stored externally in a service like Redis or DynamoDB, and workers would query it before processing a job.

Next Lesson Preview

Our scheduler is now quite powerful, but it still assumes a "happy path" for task creation. What happens if a user accidentally defines a set of tasks with a circular dependency (e.g., A -> B -> C -> A)? Our current scheduler would hang, as the tasks in the cycle would never reach an in-degree of zero.

In our next and final lesson on this scheduler, we will implement robust error handling, specifically to detect and reject circular dependencies upon task creation.