Implement Dependency-Aware Task Scheduler

Problem Statement

In distributed annotation and model training pipelines, tasks often have inter-dependencies and variable priorities. You need to design an in-memory task scheduler that manages tasks with dependencies, resolves execution readiness dynamically, and dispatches tasks based on priority.

Implement the TaskScheduler class with the following methods:

• __init__(): Initializes the scheduler.
• add_task(task_id: str, priority: int, dependencies: list[str]) -> bool: Registers a new task with an integer priority (higher value indicates higher priority) and a list of dependencies (task IDs that must be completed before task_id can execute). Dependencies may refer to tasks that have not yet been registered or tasks already completed.
  • If task_id already exists, return False without making changes.
  • If adding this task introduces a circular dependency among currently known tasks, return False without modifying the scheduler state.
  • Otherwise, register the task in state PENDING and return True.
• get_next_task() -> str | None: Selects and returns the task_id of the ready task with the highest priority.
  • A task is considered ready if it is in the PENDING state and all its dependencies are in the COMPLETED state (dependencies that have never been registered are considered incomplete).
  • If multiple ready tasks have the same highest priority, break ties by selecting the one with the lexicographically smallest task_id.
  • Upon being returned, the task transitions from PENDING to IN_PROGRESS.
  • If no task is ready, return None.
• complete_task(task_id: str) -> bool: Transitions a task from IN_PROGRESS to COMPLETED.
  • Return True if task_id was in IN_PROGRESS and successfully transitions to COMPLETED.
  • Return False if task_id does not exist or was not currently IN_PROGRESS.
  • Completing a task should unblock any dependent tasks whose remaining prerequisite tasks are all completed.

Method Signatures

class TaskScheduler:
    def __init__(self):
        pass

    def add_task(self, task_id: str, priority: int, dependencies: list[str]) -> bool:
        pass

    def get_next_task(self) -> str | None:
        pass

    def complete_task(self, task_id: str) -> bool:
        pass

Examples

Example 1

scheduler = TaskScheduler()
scheduler.add_task("B", 10, ["A"]) # returns True; A is not yet added
scheduler.get_next_task()          # returns None (B is blocked by uncompleted A)
scheduler.add_task("A", 5, [])    # returns True; A has no dependencies
scheduler.get_next_task()          # returns "A" (A is ready; state -> IN_PROGRESS)
scheduler.complete_task("A")       # returns True (A -> COMPLETED; B is now ready)
scheduler.get_next_task()          # returns "B" (B is now ready; state -> IN_PROGRESS)
scheduler.complete_task("B")       # returns True

Example 2 (Cycle Detection)

scheduler = TaskScheduler()
scheduler.add_task("X", 10, ["Y"]) # returns True
scheduler.add_task("Y", 10, ["Z"]) # returns True
scheduler.add_task("Z", 10, ["X"]) # returns False (creates cycle: X -> Y -> Z -> X)

Constraints

• Number of calls to add_task <= 2,000
• Total calls to get_next_task and complete_task <= 10,000
• task_id consists of 1 to 30 alphanumeric characters ([a-zA-Z0-9_-])
• priority is an integer in the range [-10^9, 10^9]
• Each task specifies at most 50 dependencies in dependencies

Edge Cases

• Dependencies listed in add_task that have already been marked COMPLETED should not block the task.
• Calling complete_task on a task that is PENDING, already COMPLETED, or never registered must return False.
• Calling get_next_task multiple times when no new tasks have become ready should continue returning None.