249 lines
8.9 KiB
Python
249 lines
8.9 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Any, Generic, TypeVar
|
|
|
|
from wf_core.errors import WorkflowExecutionError
|
|
from wf_core.models.steps import SubgraphNode
|
|
from wf_core.models.workflow import Workflow
|
|
from wf_core.models.workflow_refs import WorkflowRef
|
|
from wf_core.run_state import (
|
|
ExecutionFrame,
|
|
FrameStatus,
|
|
LineageState,
|
|
RunState,
|
|
RuntimeScope,
|
|
StepExecutionResult,
|
|
)
|
|
from wf_core.runtime.input_bindings import resolve_step_input_bindings
|
|
from wf_core.runtime.lineage import commit_foreach_aware_patch
|
|
from wf_core.runtime.ops.frames import frame_context_values
|
|
from wf_core.runtime.ops.merges import ReducerDefinition
|
|
from wf_core.runtime.ops.overlays import state_view_for_frame
|
|
from wf_core.runtime.ops.runs import initial_state
|
|
from wf_core.runtime.ops.schemas import validate_payload_against_schema
|
|
from wf_core.runtime.ops.state import build_output_patch, project_output
|
|
from wf_core.runtime.scheduler import add_frame, block_frame_on_children
|
|
|
|
HandlerT = TypeVar("HandlerT", bound=Callable[..., object])
|
|
_ACTIVATION_KEY = "subgraph_activation"
|
|
|
|
|
|
@dataclass(slots=True, frozen=True)
|
|
class PreparedSubgraph(Generic[HandlerT]):
|
|
"""Executable local child dependency supplied by the caller.
|
|
|
|
Core owns child execution semantics but does not load artifacts or resolve
|
|
deployment/source bindings. Higher layers must resolve those concerns into
|
|
this prepared dependency before a run starts.
|
|
"""
|
|
|
|
workflow: Workflow
|
|
registry: Mapping[str, HandlerT]
|
|
reducers: Mapping[str, ReducerDefinition] | None = None
|
|
|
|
|
|
@dataclass(slots=True, frozen=True)
|
|
class SubgraphActivation:
|
|
"""Runtime ownership record for one in-flight subgraph boundary."""
|
|
|
|
workflow_ref: WorkflowRef
|
|
scope_id: str
|
|
lineage_id: str
|
|
child_frame_id: str
|
|
child_input: dict[str, Any]
|
|
|
|
@classmethod
|
|
def from_frame(cls, frame: ExecutionFrame) -> SubgraphActivation | None:
|
|
raw = frame.metadata.get(_ACTIVATION_KEY)
|
|
if raw is None:
|
|
return None
|
|
if not isinstance(raw, Mapping):
|
|
raise WorkflowExecutionError(
|
|
f"malformed subgraph activation for frame {frame.id!r}"
|
|
)
|
|
try:
|
|
return cls(
|
|
workflow_ref=WorkflowRef.model_validate(raw["workflow_ref"]),
|
|
scope_id=str(raw["scope_id"]),
|
|
lineage_id=str(raw["lineage_id"]),
|
|
child_frame_id=str(raw["child_frame_id"]),
|
|
child_input=dict(raw["child_input"]),
|
|
)
|
|
except (KeyError, TypeError, ValueError) as exc:
|
|
raise WorkflowExecutionError(
|
|
f"malformed subgraph activation for frame {frame.id!r}"
|
|
) from exc
|
|
|
|
def save_to_frame(self, frame: ExecutionFrame) -> None:
|
|
frame.metadata[_ACTIVATION_KEY] = {
|
|
"workflow_ref": self.workflow_ref.model_dump(mode="json"),
|
|
"scope_id": self.scope_id,
|
|
"lineage_id": self.lineage_id,
|
|
"child_frame_id": self.child_frame_id,
|
|
"child_input": dict(self.child_input),
|
|
}
|
|
|
|
|
|
def resolve_prepared_subgraph(
|
|
ref: WorkflowRef,
|
|
subgraphs: Mapping[str, PreparedSubgraph[HandlerT]] | None,
|
|
) -> PreparedSubgraph[HandlerT]:
|
|
"""Resolve a caller-prepared child; artifact loading is not a core concern.
|
|
|
|
Local refs use their registry name. Saved refs use their structural display
|
|
key only as an already-prepared dependency lookup key; loading immutable
|
|
artifacts and resolving deployment bindings remains platform work.
|
|
"""
|
|
key = ref.name if ref.name is not None else ref.display
|
|
prepared = None if subgraphs is None else subgraphs.get(key)
|
|
if prepared is None:
|
|
raise WorkflowExecutionError(
|
|
f"no prepared child workflow registered for {ref.display!r}"
|
|
)
|
|
return prepared
|
|
|
|
|
|
def step_subgraph(
|
|
workflow: Workflow,
|
|
run: RunState,
|
|
step: SubgraphNode,
|
|
*,
|
|
subgraphs: Mapping[str, PreparedSubgraph[HandlerT]] | None,
|
|
reducers: Mapping[str, ReducerDefinition] | None,
|
|
) -> StepExecutionResult | None:
|
|
"""Start or finish one native child activation.
|
|
|
|
Returning ``None`` means the parent frame is blocked while child frames run.
|
|
Returning a result means child execution completed and the parent boundary
|
|
can advance normally through the child's terminal workflow outcome.
|
|
"""
|
|
frame = run.current_frame()
|
|
activation = SubgraphActivation.from_frame(frame)
|
|
prepared = resolve_prepared_subgraph(step.workflow, subgraphs)
|
|
if activation is None:
|
|
_start_subgraph(run, frame, step, prepared)
|
|
return None
|
|
return _finish_subgraph(workflow, run, frame, step, activation, prepared, reducers)
|
|
|
|
|
|
def _start_subgraph(
|
|
run: RunState,
|
|
frame: ExecutionFrame,
|
|
step: SubgraphNode,
|
|
prepared: PreparedSubgraph[HandlerT],
|
|
) -> None:
|
|
prepared.workflow.validate_structure().raise_for_errors()
|
|
parent_scope = run.scopes[frame.scope_id]
|
|
child_input = resolve_step_input_bindings(
|
|
step.input,
|
|
state=state_view_for_frame(run, frame),
|
|
workflow_input=parent_scope.workflow_input,
|
|
context=frame_context_values(frame),
|
|
label=f"subgraph {step.id!r}",
|
|
)
|
|
validate_payload_against_schema(
|
|
step.input_schema, child_input, f"subgraph input for {step.id}"
|
|
)
|
|
validate_payload_against_schema(
|
|
prepared.workflow.input_schema,
|
|
child_input,
|
|
f"child workflow input for {step.id}",
|
|
)
|
|
|
|
scope_id = f"{frame.id}:subgraph:{step.id}"
|
|
lineage_id = f"{scope_id}:root"
|
|
child_frame_id = f"{scope_id}:frame"
|
|
if scope_id in run.scopes or lineage_id in run.lineages:
|
|
raise WorkflowExecutionError(
|
|
f"duplicate subgraph activation identifiers for step {step.id!r}"
|
|
)
|
|
run.scopes[scope_id] = RuntimeScope(
|
|
id=scope_id,
|
|
workflow_name=prepared.workflow.name,
|
|
workflow_input=dict(child_input),
|
|
committed_state=initial_state(prepared.workflow, child_input),
|
|
workflow_ref=step.workflow,
|
|
)
|
|
run.lineages[lineage_id] = LineageState(id=lineage_id, scope_id=scope_id)
|
|
add_frame(
|
|
run,
|
|
ExecutionFrame(
|
|
id=child_frame_id,
|
|
kind="subgraph_root",
|
|
node_id=prepared.workflow.start,
|
|
status=FrameStatus.PENDING,
|
|
parent_frame_id=frame.id,
|
|
scope_id=scope_id,
|
|
lineage_id=lineage_id,
|
|
),
|
|
ready=True,
|
|
)
|
|
SubgraphActivation(
|
|
workflow_ref=step.workflow,
|
|
scope_id=scope_id,
|
|
lineage_id=lineage_id,
|
|
child_frame_id=child_frame_id,
|
|
child_input=child_input,
|
|
).save_to_frame(frame)
|
|
block_frame_on_children(run, frame.id, (child_frame_id,))
|
|
|
|
|
|
def _finish_subgraph(
|
|
workflow: Workflow,
|
|
run: RunState,
|
|
frame: ExecutionFrame,
|
|
step: SubgraphNode,
|
|
activation: SubgraphActivation,
|
|
prepared: PreparedSubgraph[HandlerT],
|
|
reducers: Mapping[str, ReducerDefinition] | None,
|
|
) -> StepExecutionResult:
|
|
child_frame = run.frames.get(activation.child_frame_id)
|
|
if child_frame is None or child_frame.status != FrameStatus.COMPLETED:
|
|
raise WorkflowExecutionError(
|
|
f"subgraph step {step.id!r} resumed before its child completed"
|
|
)
|
|
child_scope = run.scopes.get(activation.scope_id)
|
|
if child_scope is None:
|
|
raise WorkflowExecutionError(
|
|
f"subgraph step {step.id!r} is missing child scope {activation.scope_id!r}"
|
|
)
|
|
child_outcome = child_frame.metadata.get("workflow_outcome")
|
|
if not isinstance(child_outcome, str):
|
|
raise WorkflowExecutionError(
|
|
f"subgraph step {step.id!r} child completed without a workflow outcome"
|
|
)
|
|
child_output = project_output(
|
|
prepared.workflow,
|
|
child_scope.committed_state,
|
|
workflow_input=child_scope.workflow_input,
|
|
context=frame_context_values(child_frame),
|
|
)
|
|
validate_payload_against_schema(
|
|
prepared.workflow.output_schema,
|
|
child_output,
|
|
f"child workflow output for {step.id}",
|
|
)
|
|
validate_payload_against_schema(
|
|
step.output_schema, child_output, f"subgraph output for {step.id}"
|
|
)
|
|
patch = build_output_patch(
|
|
workflow,
|
|
step.output,
|
|
child_output,
|
|
state_view_for_frame(run, frame),
|
|
reducers=reducers,
|
|
missing_field_message="subgraph output did not include required field {field}",
|
|
)
|
|
# Foreach-aware routing (root, serial parent, concurrent lineage) is
|
|
# owned by the shared helper so subgraph output commits exactly like
|
|
# node output. Closed or superseded activations fail closed inside.
|
|
state_changes = commit_foreach_aware_patch(run, frame, patch)
|
|
return StepExecutionResult(
|
|
outcome=child_outcome,
|
|
resolved_input=activation.child_input,
|
|
output=child_output,
|
|
state_changes=state_changes,
|
|
)
|