Topological Sorting for Task Scheduling

Hello! Welcome back to your training.

In our last lesson, we successfully designed the data structures for a DAGScheduler. We can now represent a complex workflow of tasks and their dependencies as a Directed Acyclic Graph (DAG), tracking _dependencies and _in_degree for each task. However, our scheduler is still just a static model—it can't execute anything yet.

Introduction

Today, we're going to bring that model to life. Our goal is to integrate a topological sort algorithm into the scheduler to process dependent tasks in a valid order. We will implement the core worker logic that pulls ready tasks, executes them, and unlocks subsequent tasks in the dependency chain.

This lesson will cover:

• Implementing the scheduler's start and _worker methods based on Kahn's algorithm.
• Using a priority queue to manage tasks that are ready to run, ensuring we respect both dependencies and priorities.
• Handling the necessary concurrency controls to safely update the graph state as tasks complete.

This problem of orchestrating dependent units of work is central to many distributed systems. Your work on the Rudi inference scheduler at AWS, managing complex request flows, involves a similar kind of coordination, albeit in a distributed, high-throughput context. Here, we're building the fundamental in-memory engine that powers such systems.

1. The Execution Strategy: Kahn's Algorithm

As we previewed, we'll use Kahn's algorithm because it naturally fits a scheduler's design. It works by identifying all tasks that have no pending dependencies and making them available for execution.

Here's a quick refresher on the process:

1. Initialization: Compute the "in-degree" (number of dependencies) for every task. We've already done this in our add_task method.
2. Find Starting Points: Identify all tasks with an in-degree of 0. These are the "root" tasks that can run immediately. Add them to a queue of ready tasks.
3. Process and Propagate:
   • Take a task from the ready queue and execute it.
   • Once it completes, consider it "removed" from the graph.
   • For each of this task's dependents, decrement their in-degree by one.
   • If a dependent's in-degree drops to 0, it is now ready to run. Add it to the queue.
4. Repeat: Continue until the ready queue is empty.

This iterative process discovers "waves" of tasks that can be executed, which is ideal for parallel workers.

The blog post you'll read next connects this classic algorithm to modern systems development, and it's written from a Rust perspective, which should feel familiar.

2. Your Task: Implement the DAGScheduler Worker Logic

Now it's time to implement the execution engine. We will complete the DAGScheduler by filling in the start and _worker methods.

Our implementation will differ slightly from the classic Kahn's algorithm:

• We'll use a queue.PriorityQueue to hold ready tasks. This ensures that among all tasks with an in-degree of 0, we still execute the one with the highest priority first.
• We'll use multiple worker threads to process tasks in parallel.
• We must use a threading.Lock to protect shared access to the _in_degree and _dependents dictionaries when a task completes.

Below is the skeleton code. Your job is to fill in the logic in start and _worker.

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




# Using the same Task class from the previous lesson
@dataclass(order=True)
class Task:
    priority: int



    # task_id is not used for comparison in the priority queue
    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()

    def add_task(self, func: Callable, priority: int, dependencies: List[int] = None, *args, **kwargs) -> int:
        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

    def start(self):
        """
        Initializes the ready queue and starts the worker threads.
        """
        with self._lock:



            # --- YOUR IMPLEMENTATION START ---
            



            # 1. Find all tasks with an in-degree of 0. These are the initial tasks.
            #    Iterate through `self._in_degree` or `self._tasks`.
            #
            # 2. For each initial task, add it to the `_ready_queue`.
            #    Remember that `_ready_queue` is a PriorityQueue, so it will handle the priority.
            
            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])




            # --- YOUR IMPLEMENTATION END ---




        # Start worker threads
        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):
        """
        The main loop for a worker thread.
        Pulls tasks from the ready queue and executes them.
        """
        while not self._stop_event.is_set():
            try:



                # --- YOUR IMPLEMENTATION START ---
                



                # 1. Get a task from the `_ready_queue`. Use a timeout to prevent blocking forever
                #    when the stop event is set.
                task = self._ready_queue.get(timeout=1)




                # 2. 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}")




                # 3. Task is complete. Now, update the dependencies.
                #    This is a critical section and must be protected by the lock.
                with self._lock:



                    # a. Get the set of tasks that depend on the completed task.
                    dependents = self._dependents.get(task.task_id, set())
                    



                    # b. For each dependent task...
                    for dep_id in dependents:



                        # i. ...decrement its in-degree.
                        self._in_degree[dep_id] -= 1
                        



                        # ii. ...if its in-degree is now 0, it's ready! Add it to the `_ready_queue`.
                        if self._in_degree[dep_id] == 0:
                            print(f"Task {dep_id} is now ready. Adding to queue.")
                            self._ready_queue.put(self._tasks[dep_id])




                # --- YOUR IMPLEMENTATION END ---

            except queue.Empty:



                # This is expected when the queue is empty. Continue the loop.
                continue
        
        print(f"Worker {worker_id} stopping.")

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


```grasp
{
  "type": "exercise",
  "id": "820abe6e-6d6c-4ab0-b874-75ccc5bfcd53"
}

--- Example Usage ---

def sample_task(duration, name):
"""A sample task that sleeps for a given duration."""
print(f" -> Running task '{name}' for {duration}s...")
time.sleep(duration)
print(f" -> Task '{name}' complete.")

if name == 'main':
scheduler = DAGScheduler(num_workers=4)

# Model a simple parallel pipeline:
#                          ---> process_B (P=3) --\
# fetch (P=1) -> clean (P=2)                         --> aggregate (P=4) -> report (P=5)
#                          ---> process_A (P=3) --/

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



# Add tasks and store their IDs
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=3, name="process_A")
t_proc_b = scheduler.add_task(sample_task, priority=3, dependencies=[t_clean], 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")




# Start the scheduler
scheduler.start()




# Let the scheduler run until all tasks are done.
# A more robust implementation would use a different condition.
while not scheduler._ready_queue.empty() or any(thread.is_alive() for thread in scheduler._workers):
    try:
        time.sleep(1)
    except KeyboardInterrupt:
        break
        



# Check if there are uncompleted tasks (cycle detection)
with scheduler._lock:
    remaining_tasks = [tid for tid, deg in scheduler._in_degree.items() if deg > 0]
    if remaining_tasks:
        print(f"\nWarning: Cycle detected or tasks remain. Uncompleted tasks: {remaining_tasks}")

scheduler.stop()

Run the script after filling in the code. You should see the workers picking up tasks in a valid order. `fetch` will run first. After it finishes, `clean` will become ready. After `clean`, `process_A` and `process_B` will run in parallel. Only after both are done will `aggregate` run, followed finally by `report`.

---




### 3. From In-Memory to Distributed Schedulers

Our `DAGScheduler` is a powerful in-memory coordinator. How does this concept scale to a distributed system like the ones you build at AWS?

The core principles remain the same, but the implementation details change. Instead of in-memory dictionaries and locks, a distributed scheduler would use external systems.

```grasp
{
  "type": "video",
  "title": "20: Distributed Job Scheduler | Systems Design Interview Questions With Ex-Google SWE",
  "id": "[LINK](https://www.youtube.com/watch?v=pzDwYHRzEnk)",
  "video_id": "pzDwYHRzEnk",
  "relevant_section_indices": [
    2
  ],
  "par_intro": "This video segment from 'Distributed Job Scheduler' discusses exactly this transition. It outlines how a database table can be used to manage DAGs, tracking dependencies and triggering child jobs.",
  "par_directions": "Watch from 05:43 to 08:08. Pay attention to how the database table models the same concepts we used:\n\n*   **Children/Dependencies**: Our `_dependents` and `_dependencies` dictionaries.\n*   **Epochs/State**: A way to track completion, analogous to our `_in_degree` count.\n*   **Atomic Transactions**: The equivalent of our `threading.Lock` to ensure consistency when updating state."
}

This shows that the logic you've just implemented is the fundamental building block for much larger, more complex systems.

Conclusion

Excellent work! You have now built a fully functional, priority-driven, dependency-aware task scheduler in Python. This is a significant piece of engineering that showcases core concepts in concurrency, data structures, and algorithms.

Key Takeaways:

• Kahn's algorithm provides a natural, iterative model for scheduling tasks in a DAG, especially when parallel execution is desired.
• A PriorityQueue is the perfect data structure for the "ready set," as it allows the scheduler to respect both dependencies (in-degree is 0) and priority.
• In a concurrent environment, any modification to the shared graph state (like decrementing in-degrees) must be protected by a lock to prevent race conditions.

Next Lesson Preview

Our scheduler is now quite robust, but it's missing a key feature for production systems: graceful failure management. What happens if a task needs to be cancelled? Should all the tasks that depend on it also be cancelled?

In our next lesson, we will implement a cascading cancellation feature where cancelling a task also cancels all its dependents. This will require us to traverse the dependency graph again, but this time in the opposite direction—from a task to its dependents.