Add All Folders

Dagster/API_R&D/api/scheduler.py

@@ -0,0 +1,64 @@
+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)