Robust Error Handling in Schedulers

Hello! Welcome to the final lesson in our module on building systems incrementally.

In our previous lessons, we've constructed a sophisticated DAGScheduler that handles task priorities, dependencies, and even cascading cancellations. It's a powerful tool, but it currently operates under a dangerous assumption: that the dependency graph provided by the user is always valid.

Introduction

Today, we will address this critical vulnerability by adding robust error handling. Specifically, we will implement a mechanism to detect and reject circular dependencies at the time of their creation.

A circular dependency (e.g., Task A depends on B, B depends on C, and C depends on A) is fatal for our scheduler. The tasks involved in the cycle will wait for each other indefinitely, their in-degrees will never drop to zero, and they will never be placed on the ready queue. This would cause that part of the workflow to hang permanently.

This lesson covers:

• Why circular dependencies break a DAG-based scheduler.
• The classic graph algorithm for cycle detection using Depth-First Search (DFS).
• Integrating this validation check into our scheduler to make it fail-safe.
• A discussion of an alternative, BFS-based approach (Kahn's Algorithm) and its trade-offs.

Your work on production systems at AWS has undoubtedly involved building safeguards to reject invalid configurations. Preventing a user from creating an impossible state is a hallmark of robust system design. Today, we'll apply that same principle to our scheduler.

---

1. The Problem: Cycles in a Directed Graph

A Directed Acyclic Graph (DAG) is, by definition, a graph with directed edges and no cycles. Our entire scheduling model relies on this property.

The core problem is that a cycle creates a logical paradox: to run any task in the cycle, you must first run another task in the same cycle. There is no starting point. Our _in_degree counter for these tasks will never reach zero.

To solve this, we must validate the graph every time a new dependency (an edge) is added. If the new edge would create a cycle, we must reject it and inform the user.

---

2. The Algorithm: Cycle Detection with DFS

The most common and intuitive way to detect a cycle in a directed graph is to perform a Depth-First Search (DFS) and look for a "back edge" — an edge that leads from a node to one of its own ancestors in the current traversal path.

To do this, we track the state of each node using three "colors" or states:

1. UNVISITED: The node has not been explored yet.
2. VISITING: The node is currently in our recursion stack. We have started exploring from it, but we haven't finished exploring all its descendants.
3. VISITED: The node and all its descendants have been fully explored, and no cycles were found in that subgraph.

The algorithm works as follows:

• Start a DFS from a given node. Mark it as VISITING.
• For each of its neighbors (downstream dependents):
  • If the neighbor is VISITING, you have found a back edge. A cycle exists.
  • If the neighbor is UNVISITED, recursively call the DFS on that neighbor.
• After exploring all of a node's neighbors, if no cycles were found, mark the node as VISITED.

This video provides an excellent walkthrough of this exact algorithm, framed as the "Course Schedule" problem, which is a perfect analogy for our task scheduler.

---

3. Your Task: Integrating Cycle Detection

Now, let's implement this logic in our DAGScheduler. A clean way to do this is to refactor our API slightly. Instead of adding dependencies within add_task, we will create a dedicated add_dependency method. This is where our validation will live.

Your task is to:

1. Refactor the scheduler to have separate add_task and add_dependency methods.
2. Implement the cycle detection logic within add_dependency. If adding the dependency creates a cycle, the method should raise a ValueError and not add the dependency.

Here is the refactored code structure. Fill in the missing implementation for the cycle detection.

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




# Using the same Task class from the previous lessons
@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)




# Enum for DFS states
class VisitState:
    UNVISITED = 0
    VISITING = 1
    VISITED = 2

class DAGScheduler:
    def __init__(self, num_workers: int = 4):
        self._tasks: Dict[int, Task] = {}



        # _dependents maps a task to the tasks that depend on it (downstream)
        self._dependents: Dict[int, Set[int]] = defaultdict(set)



        # _dependencies maps a task to the tasks it depends on (upstream)
        self._dependencies: Dict[int, Set[int]] = defaultdict(set)
        self._in_degree: Dict[int, int] = 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()
        self._cancelled_tasks: Set[int] = set()

    def add_task(self, func: Callable, priority: int, *args, **kwargs) -> int:
        """Adds a new task to the graph without any 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:
            self._tasks[task_id] = task
            self._in_degree[task_id] = 0
        
        return task_id




    # --- YOUR IMPLEMENTATION START ---
    def add_dependency(self, source_id: int, target_id: int):
        """
        Adds a dependency from source_id to target_id.
        The target_id task will run after the source_id task.
        Raises ValueError if the dependency creates a cycle.
        """
        with self._lock:
            if source_id not in self._tasks or target_id not in self._tasks:
                raise ValueError("Source or target task not found.")
            
            if target_id in self._dependents[source_id]:
                return # Dependency already exists




            # 1. Tentatively add the edge
            self._dependents[source_id].add(target_id)
            self._dependencies[target_id].add(source_id)




            # 2. Run cycle detection
            # We only need to check for cycles starting from the target node,
            # as a new cycle must involve the new edge.
            if self._detect_cycle(target_id):



                # 3. If a cycle is found, revert the change and raise an error
                self._dependents[source_id].remove(target_id)
                self._dependencies[target_id].remove(source_id)
                raise ValueError(f"Adding dependency from {source_id} to {target_id} creates a cycle.")
            



            # 4. If no cycle, commit the in-degree change
            self._in_degree[target_id] += 1

    def _detect_cycle(self, start_node: int) -> bool:
        """
        Detects if a cycle exists in the graph reachable from start_node.
        Returns True if a cycle is detected, False otherwise.
        """
        states = defaultdict(lambda: VisitState.UNVISITED)
        



        # The recursion stack for the DFS traversal
        recursion_stack = {start_node}
        
        return self._detect_cycle_util(start_node, states, recursion_stack)

    def _detect_cycle_util(self, node_id: int, states: Dict[int, int], recursion_stack: Set[int]) -> bool:
        """
        Recursive helper for DFS-based cycle detection.
        `recursion_stack` tracks nodes in the current traversal path (VISITING state).
        `states` tracks nodes that are fully explored (VISITED state).
        """
        states[node_id] = VisitState.VISITING

        for neighbor_id in self._dependents[node_id]:



            # If the neighbor is in the current recursion stack, we found a cycle.
            if neighbor_id in recursion_stack:
                return True
            



            # If the neighbor is unvisited, recurse.
            if states[neighbor_id] == VisitState.UNVISITED:
                recursion_stack.add(neighbor_id)
                if self._detect_cycle_util(neighbor_id, states, recursion_stack):
                    return True
                recursion_stack.remove(neighbor_id)




        # Mark as fully explored for this traversal path
        states[node_id] = VisitState.VISITED
        return False



    # --- YOUR IMPLEMENTATION END ---




    # The rest of the methods (start, _worker, stop, cancel_task) remain the same.
    # ...
    def start(self):
        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:
                if task_id in self._tasks: # Check if task exists
                    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)
                
                with self._lock:
                    if task.task_id in self._cancelled_tasks:
                        print(f"Worker {worker_id}: Skipping cancelled task {task.task_id}")
                    else:
                        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}")
                
                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 and 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])

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

    def is_done(self) -> bool:
        with self._lock:



            # Check if all tasks that were added are accounted for (completed, cancelled, or still have dependencies)
            processed_count = len([tid for tid, deg in self._in_degree.items() if deg <= 0])
            return self._ready_queue.empty() and processed_count >= len(self._tasks)

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



        # Unchanged from previous lesson
        pass




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

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




    # A -> B -> C
    task_a = scheduler.add_task(sample_task, priority=1, name="A")
    task_b = scheduler.add_task(sample_task, priority=1, name="B")
    task_c = scheduler.add_task(sample_task, priority=1, name="C")
    
    scheduler.add_dependency(task_a, task_b)
    scheduler.add_dependency(task_b, task_c)
    
    print("Graph: A -> B -> C")
    print("Attempting to add dependency C -> A (which would create a cycle)...")
    
    try:
        scheduler.add_dependency(task_c, task_a)
    except ValueError as e:
        print(f"\nSUCCESS: Correctly caught error: {e}")
        
    print("\nStarting scheduler with the valid graph (A -> B -> C).")
    scheduler.start()
    scheduler.stop()

Note: A more complete implementation of _detect_cycle_util would need to handle disconnected components in the graph by iterating through all nodes. For our use case—checking if a single new edge creates a cycle—starting the check from the target_id is a sound optimization.

---

4. Alternative Approach: Kahn's Algorithm

The DFS-based method is not the only way to detect cycles. Another popular algorithm, which also produces a topological sort, is Kahn's Algorithm. It is BFS-based and works by tracking the in-degrees of nodes.

The core idea is:

1. Initialize a queue with all nodes that have an in-degree of 0.
2. While the queue is not empty, dequeue a node, add it to your sorted list, and "remove" it from the graph.
3. To "remove" it, iterate through its neighbors and decrement their in-degrees.
4. If a neighbor's in-degree becomes 0, add it to the queue.
5. After the loop, if the number of nodes in your sorted list is less than the total number of nodes in the graph, it means some nodes were never processed. These are the nodes involved in a cycle.

This video gives a great overview of Kahn's algorithm.

Trade-offs: DFS vs. Kahn's Algorithm

• Cycle Reporting: The DFS approach excels at finding and reporting the exact path of a cycle as soon as it's found (a "back edge"). Kahn's algorithm tells you that a cycle exists and which nodes are part of some cycle (those with in_degree > 0 at the end), but it requires extra work to reconstruct a specific cycle path.
• Implementation: DFS is naturally recursive, which can be elegant but risks stack overflow on very deep graphs. Kahn's algorithm is iterative (using a queue), avoiding recursion limits.
• Parallelism: The structure of Kahn's algorithm can expose opportunities for parallel execution more naturally. The set of nodes with an in-degree of 0 at any given time can all be processed concurrently. This is very similar to how our scheduler's main loop already functions.

Given your background building schedulers like Rudi, the in-degree counting in Kahn's algorithm might feel very familiar, as it's a common pattern for managing resource dependencies and readiness.

---

Conclusion

Congratulations! You have now completed the "Building Systems Incrementally" module. By adding cycle detection, you've transformed our scheduler from a functional prototype into a robust, production-aware system. It now correctly validates its inputs and guarantees that it won't enter an unrecoverable state due to a malformed dependency graph.

Key Takeaways:

• Fail-fast validation is crucial. Detecting configuration errors like cycles before execution is a fundamental principle of robust system design.
• Three-state DFS is a classic and effective algorithm for finding cycles in a directed graph by identifying "back edges."
• Kahn's algorithm provides an alternative, iterative approach based on in-degrees, which has its own set of trade-offs regarding implementation and cycle reporting.

Next Lesson Preview

In our next module, "Real-World Coding: Concurrent Web Crawler," we will shift our focus to a new problem domain. We'll leverage Python's asyncio library to build a high-performance web crawler, applying many of the concurrency and data structure concepts we've discussed to a practical, network-bound application. Our first step will be to handle the starting point of all well-behaved web crawlers: parsing a robots.txt file.