Scheduling with Task Dependencies (DAG)

Hello! Welcome back.

In our last lesson, we built a robust PriorityScheduler that executes tasks based on importance, not just arrival time. By swapping a FIFO queue for a min-heap, we ensured that critical jobs could preempt less urgent ones.

However, our scheduler still has a blind spot: it doesn't understand dependencies. What happens if task_B requires the output of task_A? Even if task_B has a higher priority, it's invalid to run it before task_A is complete. This is a common requirement in complex workflows, from CI/CD pipelines to data processing jobs like the ones you've managed at AWS.

Introduction

Today, we will address this by extending our scheduler to understand and represent these relationships. Our goal is to model task dependencies using a Directed Acyclic Graph (DAG). We'll focus on modifying our data structures to represent this graph, laying the essential groundwork for a fully dependency-aware execution engine.

This lesson will cover:

• What a Directed Acyclic Graph (DAG) is and why it's the perfect model for this problem.
• How to extend our scheduler's data model to store tasks and their dependencies.
• Implementing the logic to build this graph structure as tasks are added.

By the end of this lesson, you will have a scheduler that can represent a complex workflow, ready for us to implement the execution logic in our next session.

---

1. Modeling Dependencies as a Graph

When one task must precede another, we have a directed relationship. If we represent each task as a node and each dependency as a directed edge, we form a graph. For a system of tasks to be solvable, this graph must not contain any circular dependencies (e.g., A depends on B, and B depends on A), which would cause a deadlock. A directed graph with no cycles is called a Directed Acyclic Graph (DAG).

To understand the core concepts behind DAGs and their processing, let's watch a short video.

The key takeaway is that for any DAG, there exists at least one valid topological ordering—a linear sequence of tasks that respects all dependencies. Our ultimate goal is to make our scheduler execute tasks in such an order.

---

2. Extending the Scheduler's Data Model

To manage a DAG, our scheduler needs to store the graph's structure. We'll move beyond a simple list of tasks and implement a more comprehensive model. The PriorityScheduler was a good start, but a dependency-aware scheduler needs to track the relationships between every task.

Here are the essential data structures we'll add to a new DAGScheduler class:

• _tasks: A dictionary mapping a unique task_id to its corresponding Task object. This gives us fast lookups.
• _dependents: A dictionary representing the graph's adjacency list. It maps a task_id to a set of task_ids that depend on it. Example: {'task_A': {'task_B', 'task_C'}}.
• _dependencies: The reverse of _dependents. It maps a task_id to a set of task_ids it depends on. This is useful for quickly checking a task's prerequisites.
• _in_degree: A dictionary mapping a task_id to an integer representing its number of uncompleted dependencies. A task is ready to run only when its in-degree is 0. This will be the cornerstone of our execution algorithm in the next lesson.

Your experience with leader-follower architectures and coordination using DynamoDB touches upon similar principles of tracking state and dependencies between distributed components. Here, we're applying that same mindset to an in-memory scheduler.

---

3. Your Task: Implement the DAG Data Model

Now it's time to put this into practice. Below is the skeleton for a DAGScheduler. Your task is to implement the add_task method. This method is responsible for taking a new task and its dependencies and correctly updating the graph data structures we just discussed.

For this lesson, we will not implement the worker logic yet. The goal is to correctly build the internal graph representation. The start and stop methods are placeholders for now.

import threading
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Callable, Any, Set, List
import itertools




# We'll use the same Task dataclass from the previous lesson
@dataclass(order=True)
class Task:
    priority: int
    task_id: int
    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):
        self._tasks = {}  # {task_id: Task}



        # Adjacency list: {task_id: {dependent_task_id, ...}}
        self._dependents = defaultdict(set)



        # Reverse lookup: {task_id: {dependency_task_id, ...}}
        self._dependencies = defaultdict(set)



        # In-degree count: {task_id: count}
        self._in_degree = defaultdict(int)

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

    def add_task(self, func: Callable, priority: int, dependencies: List[int] = None, *args, **kwargs) -> int:
        """
        Adds a task to the scheduler and builds the dependency graph.

        Args:
            func: The function to execute.
            priority: The priority of the task.
            dependencies: A list of task_ids that this task depends on.
            *args: Positional arguments for the function.
            **kwargs: Keyword arguments for the function.

        Returns:
            The unique ID of the newly added task.
        """
        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:



            # --- YOUR IMPLEMENTATION START ---
            



            # 1. Add the new task to the `_tasks` dictionary.
            



            # 2. Store the dependencies for the new task.
            #    - Add the given `dependencies` list to `_dependencies[task_id]`.
            



            # 3. For each dependency in the `dependencies` list, update the `_dependents` graph.
            #    - For a dependency `dep_id`, `task_id` is a dependent.
            



            # 4. Set the initial in-degree for the new task.
            #    - The in-degree is the number of tasks it depends on.
            



            # --- YOUR IMPLEMENTATION END ---

        return task_id

    def start(self):
        print("Scheduler start logic will be implemented in the next lesson.")

    def stop(self):
        print("Scheduler stop logic will be implemented in the next lesson.")
        
    def _print_graph_state(self):
        """Helper to visualize the internal state for verification."""
        with self._lock:
            print("\n--- Scheduler Graph State ---")
            print(f"Tasks: {list(self._tasks.keys())}")
            print("In-Degrees:")
            for tid, degree in sorted(self._in_degree.items()):
                print(f"  - Task {tid}: {degree}")
            print("Dependents (Adjacency List):")
            for tid, deps in sorted(self._dependents.items()):
                print(f"  - Task {tid} is a prerequisite for: {sorted(list(deps))}")
            print("---------------------------\n")




# --- Example Usage ---
if __name__ == '__main__':
    scheduler = DAGScheduler()




    # Let's model a simple data processing pipeline:
    #                          ---> process_data_B --\
    # fetch_data -> clean_data                         --> aggregate_results -> generate_report
    #                          ---> process_data_A --/
    
    print("Building the task graph...")




    # Wave 1: No dependencies
    fetch_task = scheduler.add_task(print, priority=1, func_name="fetch_data")
    



    # Wave 2: Depends on fetch_task
    clean_task = scheduler.add_task(print, priority=2, dependencies=[fetch_task], func_name="clean_data")
    



    # Wave 3: Depends on clean_task
    process_a_task = scheduler.add_task(print, priority=3, dependencies=[clean_task], func_name="process_data_A")
    process_b_task = scheduler.add_task(print, priority=3, dependencies=[clean_task], func_name="process_data_B")
    



    # Wave 4: Depends on both processing tasks
    aggregate_task = scheduler.add_task(print, priority=4, dependencies=[process_a_task, process_b_task], func_name="aggregate_results")
    



    # Wave 5: Depends on aggregation
    report_task = scheduler.add_task(print, priority=5, dependencies=[aggregate_task], func_name="generate_report")




    # Verify the internal state
    scheduler._print_graph_state()

After implementing add_task, running the script should produce output that correctly reflects the dependencies we've defined, with accurate in-degree counts for each task. For example, aggregate_task (ID 4) should have an in-degree of 2.

---

4. Preview: How to Process the Graph

Now that our scheduler can model a DAG, the next logical question is: how do we execute it? As previewed in the video, we need a topological sorting algorithm. There are two main approaches.

Kahn's Algorithm (Queue-based)

This is an iterative approach that aligns perfectly with a scheduler's design.

The core idea of Kahn's algorithm is to maintain a queue of "ready" tasks (those with an in-degree of 0). Our PriorityScheduler's heap is perfect for this. When a worker finishes a task, it effectively removes it from the graph, decrements the in-degree of its dependents, and adds any newly-ready tasks to the priority queue.

DFS-based Algorithm (Recursive)

Another common method involves a Depth-First Search (DFS). This is a classic graph traversal pattern often seen in interview questions like "Course Schedule".

Trade-offs:

• Kahn's Algorithm is generally more suitable for a scheduler. It naturally identifies all runnable tasks at any given moment, allowing for parallel execution. The "wave-based" scheduling discussed in the dag-task-scheduler resource (LINK) is a direct application of this principle.
• DFS-based Algorithm is often simpler to write for simply finding one valid ordering, but it's less natural for a dynamic scheduling system that needs to manage a pool of ready tasks.

In our next lesson, we will focus on implementing Kahn's algorithm to bring our DAGScheduler to life.

---

Conclusion

Great work today! We've made a significant leap from a simple task processor to a system capable of understanding complex workflows.

Key Takeaways:

• Task dependencies are best modeled as a Directed Acyclic Graph (DAG).
• A robust scheduler needs internal data structures to represent the graph: _tasks for lookup, _dependents for the adjacency list, and _in_degree to track readiness.
• By populating these structures upon task creation, we've laid the foundation for a dependency-aware execution strategy.

Next Lesson Preview

In the next lesson, we will fully integrate a topological sort algorithm into the scheduler to process dependent tasks in a valid order. We will implement the worker loop, turning our static graph model into a dynamic, running system that respects every dependency while still honoring task priorities.