native subgraph execution

src/wf_core/runtime/engine.py

@@ -3,17 +3,19 @@ from __future__ import annotations
 from collections.abc import Mapping
 from typing import Any
 
+from wf_core.errors import WorkflowExecutionError
 from wf_core.models.workflow import Workflow
 from wf_core.runtime.ops.flow import finalize_run
 from wf_core.runtime.ops.merges import ReducerDefinition
 from wf_core.runtime.ops.nodes import AsyncNodeHandler, NodeHandler
 from wf_core.runtime.ops.runs import create_run_state
 from wf_core.runtime.scheduler import resolve_no_ready_frames, select_next_frame
-from wf_core.run_state import RunState, RunStatus
+from wf_core.run_state import ROOT_SCOPE_ID, RunState, RunStatus
 from wf_core.tokens import END
 
 from .preparation import prepare_new_run, prepare_resume
 from .step import step_workflow, step_workflow_async
+from .subgraphs import PreparedSubgraph, resolve_prepared_subgraph
 
 
 def execute_workflow(
@@ -22,13 +24,20 @@ def execute_workflow(
     registry: Mapping[str, NodeHandler],
     *,
     reducers: Mapping[str, ReducerDefinition] | None = None,
+    subgraphs: Mapping[str, PreparedSubgraph[NodeHandler]] | None = None,
 ) -> RunState:
     """Create a run and execute a workflow synchronously until it stops."""
     run = create_run_state(workflow, workflow_input)
 
     try:
         prepare_new_run(workflow, workflow_input, run)
-        return resume_workflow(workflow, run, registry, reducers=reducers)
+        return resume_workflow(
+            workflow,
+            run,
+            registry,
+            reducers=reducers,
+            subgraphs=subgraphs,
+        )
     except Exception as exc:
         run.status = RunStatus.FAILED
         run.error = str(exc)
@@ -41,13 +50,20 @@ async def execute_workflow_async(
     registry: Mapping[str, AsyncNodeHandler],
     *,
     reducers: Mapping[str, ReducerDefinition] | None = None,
+    subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None,
 ) -> RunState:
     """Create a run and execute a workflow asynchronously until it stops."""
     run = create_run_state(workflow, workflow_input)
 
     try:
         prepare_new_run(workflow, workflow_input, run)
-        return await resume_workflow_async(workflow, run, registry, reducers=reducers)
+        return await resume_workflow_async(
+            workflow,
+            run,
+            registry,
+            reducers=reducers,
+            subgraphs=subgraphs,
+        )
     except Exception as exc:
         run.status = RunStatus.FAILED
         run.error = str(exc)
@@ -62,6 +78,7 @@ def resume_workflow(
     resume_payload: dict[str, Any] | None = None,
     resume_outcome: str = "submitted",
     reducers: Mapping[str, ReducerDefinition] | None = None,
+    subgraphs: Mapping[str, PreparedSubgraph[NodeHandler]] | None = None,
 ) -> RunState:
     """Resume a synchronous run from its current state."""
     index = prepare_resume(
@@ -77,17 +94,22 @@ def resume_workflow(
         return run
 
     while True:
-        if select_next_frame(run) is None:
+        frame = select_next_frame(run)
+        if frame is None:
             status = resolve_no_ready_frames(run)
             if status == RunStatus.COMPLETED:
                 break
             return run
+        active_workflow, active_registry, active_reducers = _sync_execution_target(
+            workflow, registry, reducers, run, subgraphs
+        )
         step_workflow(
-            workflow,
+            active_workflow,
             run,
-            registry,
-            index=index,
-            reducers=reducers,
+            active_registry,
+            index=index if frame.scope_id == ROOT_SCOPE_ID else None,
+            reducers=active_reducers,
+            subgraphs=subgraphs,
         )
         if run.status == RunStatus.INTERRUPTED:
             return run
@@ -103,6 +125,7 @@ async def resume_workflow_async(
     resume_payload: dict[str, Any] | None = None,
     resume_outcome: str = "submitted",
     reducers: Mapping[str, ReducerDefinition] | None = None,
+    subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None,
 ) -> RunState:
     """Resume an async run from its current state."""
     index = prepare_resume(
@@ -118,19 +141,68 @@ async def resume_workflow_async(
         return run
 
     while True:
-        if select_next_frame(run) is None:
+        frame = select_next_frame(run)
+        if frame is None:
             status = resolve_no_ready_frames(run)
             if status == RunStatus.COMPLETED:
                 break
             return run
+        active_workflow, active_registry, active_reducers = _async_execution_target(
+            workflow, registry, reducers, run, subgraphs
+        )
         await step_workflow_async(
-            workflow,
+            active_workflow,
             run,
-            registry,
-            index=index,
-            reducers=reducers,
+            active_registry,
+            index=index if frame.scope_id == ROOT_SCOPE_ID else None,
+            reducers=active_reducers,
+            subgraphs=subgraphs,
         )
         if run.status == RunStatus.INTERRUPTED:
             return run
 
     return finalize_run(workflow, run)
+
+
+def _sync_execution_target(
+    root_workflow: Workflow,
+    root_registry: Mapping[str, NodeHandler],
+    root_reducers: Mapping[str, ReducerDefinition] | None,
+    run: RunState,
+    subgraphs: Mapping[str, PreparedSubgraph[NodeHandler]] | None,
+) -> tuple[Workflow, Mapping[str, NodeHandler], Mapping[str, ReducerDefinition] | None]:
+    """Return the workflow dependencies owned by the selected frame scope."""
+    frame = run.current_frame()
+    if frame.scope_id == ROOT_SCOPE_ID:
+        return root_workflow, root_registry, root_reducers
+    scope = run.scopes.get(frame.scope_id)
+    if scope is None or scope.workflow_ref is None:
+        raise WorkflowExecutionError(
+            f"child frame {frame.id!r} has no prepared workflow scope"
+        )
+    child = resolve_prepared_subgraph(scope.workflow_ref, subgraphs)
+    return child.workflow, child.registry, child.reducers
+
+
+def _async_execution_target(
+    root_workflow: Workflow,
+    root_registry: Mapping[str, AsyncNodeHandler],
+    root_reducers: Mapping[str, ReducerDefinition] | None,
+    run: RunState,
+    subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None,
+) -> tuple[
+    Workflow,
+    Mapping[str, AsyncNodeHandler],
+    Mapping[str, ReducerDefinition] | None,
+]:
+    """Return async workflow dependencies owned by the selected frame scope."""
+    frame = run.current_frame()
+    if frame.scope_id == ROOT_SCOPE_ID:
+        return root_workflow, root_registry, root_reducers
+    scope = run.scopes.get(frame.scope_id)
+    if scope is None or scope.workflow_ref is None:
+        raise WorkflowExecutionError(
+            f"child frame {frame.id!r} has no prepared workflow scope"
+        )
+    child = resolve_prepared_subgraph(scope.workflow_ref, subgraphs)
+    return child.workflow, child.registry, child.reducers