soorena62
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
from typing import Set
|
|
from radiuma_api.task import Task, Status
|
|
from radiuma_api.execution_context import ExecutionContext
|
|
# Write Codes Here:
|
|
|
|
|
|
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):
|
|
"""
|
|
Entry point: execute a Workflow or Task graph.
|
|
"""
|
|
self._execute_task(root, context)
|
|
|
|
def _execute_task(self, task: Task, context: ExecutionContext):
|
|
"""
|
|
Execute a task respecting DAG dependencies.
|
|
If 'task' is a Workflow (has children), delegate to _run_workflow.
|
|
Otherwise, run the leaf task when inputs are ready.
|
|
"""
|
|
# Skip finished tasks
|
|
if task.status in (Status.COMPLETED, Status.FAILED, Status.STOPPED):
|
|
return
|
|
|
|
# Workflow (composite)
|
|
if hasattr(task, "children"):
|
|
self._run_workflow(task, context)
|
|
return
|
|
|
|
# Leaf Task: must have inputs ready
|
|
if not task.check_inputs_ready(context):
|
|
# Optional: diagnostic log
|
|
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)
|
|
# Optional: diagnostic log
|
|
print(f"[Scheduler] Task '{task.name}' failed: {e}")
|
|
# Re-raise to surface error if caller needs to handle
|
|
raise
|
|
|
|
def _run_workflow(self, workflow: Task, context: ExecutionContext):
|
|
"""
|
|
Execute children respecting dependencies.
|
|
Progress until all children are executed or a deadlock is detected.
|
|
"""
|
|
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:
|
|
# Deadlock: dependencies not resolvable with current data
|
|
raise RuntimeError(
|
|
f"Workflow '{workflow.name}' deadlock: unresolved dependencies"
|
|
)
|
|
|
|
# Final workflow status based on children results
|
|
if any(c.status == Status.FAILED for c in workflow.children):
|
|
workflow._set_status(Status.FAILED)
|
|
else:
|
|
workflow._set_status(Status.COMPLETED)
|