soorena62
65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
from typing import Set
|
|
from api.task import Task, Status
|
|
from api.execution_context import ExecutionContext
|
|
|
|
|
|
class Scheduler:
|
|
"""
|
|
DAG-based execution scheduler.
|
|
Decides WHEN a Task can run based on port readiness and orchestrates execution.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._completed: Set[Task] = set()
|
|
self._failed: Set[Task] = set()
|
|
|
|
def run(self, root: Task, context: ExecutionContext):
|
|
self._execute_task(root, context)
|
|
|
|
def _execute_task(self, task: Task, context: ExecutionContext):
|
|
if task.status in (Status.COMPLETED, Status.FAILED, Status.STOPPED):
|
|
return
|
|
|
|
if hasattr(task, "children"):
|
|
self._run_workflow(task, context)
|
|
return
|
|
|
|
if not task.check_inputs_ready(context):
|
|
print(f"[Scheduler] Task '{task.name}' not ready, skipping for now.")
|
|
return
|
|
|
|
try:
|
|
task._set_status(Status.RUNNING)
|
|
task.run(context)
|
|
task._set_status(Status.COMPLETED)
|
|
self._completed.add(task)
|
|
except Exception as e:
|
|
task._set_status(Status.FAILED)
|
|
self._failed.add(task)
|
|
print(f"[Scheduler] Task '{task.name}' failed: {e}")
|
|
raise
|
|
|
|
def _run_workflow(self, workflow: Task, context: ExecutionContext):
|
|
workflow._set_status(Status.RUNNING)
|
|
|
|
remaining = set(workflow.children)
|
|
|
|
while remaining:
|
|
progress_made = False
|
|
|
|
for child in list(remaining):
|
|
if child.check_inputs_ready(context):
|
|
self._execute_task(child, context)
|
|
remaining.remove(child)
|
|
progress_made = True
|
|
|
|
if not progress_made:
|
|
raise RuntimeError(
|
|
f"Workflow '{workflow.name}' deadlock: unresolved dependencies"
|
|
)
|
|
|
|
if any(c.status == Status.FAILED for c in workflow.children):
|
|
workflow._set_status(Status.FAILED)
|
|
else:
|
|
workflow._set_status(Status.COMPLETED)
|