concurrent foreach work
This commit is contained in:
@@ -15,7 +15,7 @@ from wf_core.runtime.ops.schemas import validate_payload_against_schema
|
||||
from wf_core.runtime.ops.state import project_output
|
||||
from wf_core.runtime.scheduler import (
|
||||
mark_frame_pending,
|
||||
wake_parent_if_children_complete,
|
||||
wake_parent_for_child_progress,
|
||||
)
|
||||
from wf_core.tokens import END
|
||||
|
||||
@@ -81,7 +81,7 @@ def advance_frame(
|
||||
if next_node_id == END:
|
||||
frame.status = FrameStatus.COMPLETED
|
||||
frame.finished_at_node_id = END
|
||||
wake_parent_if_children_complete(run, frame.id)
|
||||
wake_parent_for_child_progress(run, frame.id)
|
||||
else:
|
||||
frame.finished_at_node_id = None
|
||||
mark_frame_pending(run, frame.id)
|
||||
|
||||
@@ -1,19 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
|
||||
from wf_core.conditions import safe_resolve_path
|
||||
from wf_core.errors import WorkflowExecutionError
|
||||
from wf_core.models.steps import ForeachNode
|
||||
from wf_core.models.steps import ForeachNode, NodeUse
|
||||
from wf_core.models.workflow import Workflow
|
||||
from wf_core.run_state import ExecutionFrame, FrameStatus, RunState, StepExecutionResult
|
||||
from wf_core.runtime.foreach_state import ForeachBarrierState
|
||||
from wf_core.runtime.ops.flow import advance_frame, append_step_result_trace
|
||||
from wf_core.runtime.ops.frames import frame_context_values
|
||||
from wf_core.runtime.ops.index import WorkflowIndex
|
||||
from wf_core.runtime.ops.merges import ReducerDefinition
|
||||
from wf_core.runtime.ops.state import build_barrier_patch, commit_state_patch
|
||||
from wf_core.runtime.scheduler import (
|
||||
ForeachIterationMetadata,
|
||||
add_frame,
|
||||
block_frame_on_children,
|
||||
)
|
||||
from wf_core.tokens import END
|
||||
|
||||
|
||||
def step_foreach(
|
||||
@@ -21,25 +26,32 @@ def step_foreach(
|
||||
run: RunState,
|
||||
step: ForeachNode,
|
||||
index: WorkflowIndex,
|
||||
*,
|
||||
reducers: Mapping[str, ReducerDefinition] | None = None,
|
||||
) -> RunState:
|
||||
if step.mode == "serial":
|
||||
return _step_foreach_serial(workflow, run, step, index)
|
||||
return _step_foreach_concurrent(
|
||||
workflow,
|
||||
run,
|
||||
step,
|
||||
index,
|
||||
reducers=reducers,
|
||||
)
|
||||
|
||||
|
||||
def _step_foreach_serial(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
step: ForeachNode,
|
||||
index: WorkflowIndex,
|
||||
) -> RunState:
|
||||
if step.mode != "serial":
|
||||
raise WorkflowExecutionError(
|
||||
"concurrent foreach execution is not implemented yet"
|
||||
)
|
||||
raise WorkflowExecutionError("serial foreach helper received non-serial mode")
|
||||
|
||||
frame = run.current_frame()
|
||||
barrier = ForeachBarrierState.from_frame(frame, step.id) or ForeachBarrierState()
|
||||
|
||||
iterable = safe_resolve_path(
|
||||
str(step.over),
|
||||
state=run.state,
|
||||
workflow_input=run.workflow_input,
|
||||
context=frame_context_values(frame),
|
||||
)
|
||||
if not isinstance(iterable, list):
|
||||
raise WorkflowExecutionError(
|
||||
f"foreach source {str(step.over)!r} must resolve to a list"
|
||||
)
|
||||
iterable = _resolve_foreach_iterable(run, frame, step)
|
||||
|
||||
loop_index = barrier.next_index
|
||||
if loop_index >= len(iterable):
|
||||
@@ -62,17 +74,10 @@ def step_foreach(
|
||||
return run
|
||||
|
||||
loop_start = index.next_node_id(frame.node_id, "loop")
|
||||
|
||||
item = iterable[loop_index]
|
||||
barrier.next_index = loop_index + 1
|
||||
barrier.save_to_frame(frame, step.id)
|
||||
child_id = f"{frame.id}:{step.id}:{loop_index}"
|
||||
child_metadata = ForeachIterationMetadata(
|
||||
foreach_node_id=step.id,
|
||||
loop_index=loop_index,
|
||||
loop_item=item,
|
||||
loop_alias=step.as_,
|
||||
)
|
||||
add_frame(
|
||||
run,
|
||||
ExecutionFrame(
|
||||
@@ -81,7 +86,12 @@ def step_foreach(
|
||||
node_id=loop_start,
|
||||
status=FrameStatus.PENDING,
|
||||
parent_frame_id=frame.id,
|
||||
metadata=child_metadata.to_metadata(),
|
||||
metadata=ForeachIterationMetadata(
|
||||
foreach_node_id=step.id,
|
||||
loop_index=loop_index,
|
||||
loop_item=item,
|
||||
loop_alias=step.as_,
|
||||
).to_metadata(),
|
||||
),
|
||||
ready=True,
|
||||
)
|
||||
@@ -101,3 +111,210 @@ def step_foreach(
|
||||
)
|
||||
run.sync_from_current_frame()
|
||||
return run
|
||||
|
||||
|
||||
def _step_foreach_concurrent(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
step: ForeachNode,
|
||||
index: WorkflowIndex,
|
||||
*,
|
||||
reducers: Mapping[str, ReducerDefinition] | None = None,
|
||||
) -> RunState:
|
||||
if step.item_error.action != "fail":
|
||||
raise WorkflowExecutionError(
|
||||
"concurrent foreach v1 only supports item_error.action='fail'"
|
||||
)
|
||||
if step.concurrent is None:
|
||||
raise WorkflowExecutionError("concurrent foreach requires concurrent policy")
|
||||
_validate_single_node_loop_body(index, step)
|
||||
|
||||
frame = run.current_frame()
|
||||
barrier = ForeachBarrierState.from_frame(frame, step.id)
|
||||
if barrier is None:
|
||||
barrier = ForeachBarrierState(mode="concurrent")
|
||||
elif barrier.mode != "concurrent":
|
||||
raise WorkflowExecutionError("malformed concurrent foreach barrier mode")
|
||||
|
||||
_finish_completed_children(run, barrier)
|
||||
iterable = _resolve_foreach_iterable(run, frame, step)
|
||||
_admit_concurrent_children(
|
||||
run=run,
|
||||
frame=frame,
|
||||
step=step,
|
||||
index=index,
|
||||
barrier=barrier,
|
||||
iterable=iterable,
|
||||
)
|
||||
|
||||
if barrier.next_index >= len(iterable) and not barrier.outstanding_frame_ids:
|
||||
return _finish_concurrent_foreach(
|
||||
workflow=workflow,
|
||||
run=run,
|
||||
frame=frame,
|
||||
step=step,
|
||||
index=index,
|
||||
barrier=barrier,
|
||||
reducers=reducers,
|
||||
)
|
||||
|
||||
barrier.save_to_frame(frame, step.id)
|
||||
block_frame_on_children(run, frame.id, barrier.outstanding_frame_ids)
|
||||
run.sync_from_current_frame()
|
||||
return run
|
||||
|
||||
|
||||
def _resolve_foreach_iterable(
|
||||
run: RunState,
|
||||
frame: ExecutionFrame,
|
||||
step: ForeachNode,
|
||||
) -> list[object]:
|
||||
iterable = safe_resolve_path(
|
||||
str(step.over),
|
||||
state=run.state,
|
||||
workflow_input=run.workflow_input,
|
||||
context=frame_context_values(frame),
|
||||
)
|
||||
if not isinstance(iterable, list):
|
||||
raise WorkflowExecutionError(
|
||||
f"foreach source {str(step.over)!r} must resolve to a list"
|
||||
)
|
||||
return iterable
|
||||
|
||||
|
||||
def _validate_single_node_loop_body(index: WorkflowIndex, step: ForeachNode) -> None:
|
||||
"""Reject multi-step concurrent item bodies until item overlays are real.
|
||||
|
||||
The current slice has a no-op item-state overlay seam. Without a real overlay,
|
||||
multi-node item bodies would read stale parent state after earlier item-local
|
||||
writes, so V1 only allows loop -> one node -> END.
|
||||
"""
|
||||
loop_start = index.next_node_id(step.id, "loop")
|
||||
loop_step = index.nodes_by_id.get(loop_start)
|
||||
if not isinstance(loop_step, NodeUse):
|
||||
raise WorkflowExecutionError(
|
||||
"concurrent foreach v1 only supports loop bodies with one node"
|
||||
)
|
||||
node_def = index.node_defs[loop_step.node]
|
||||
for outcome in node_def.outcomes:
|
||||
if index.next_node_id(loop_step.id, outcome) != END:
|
||||
raise WorkflowExecutionError(
|
||||
"concurrent foreach v1 only supports loop bodies with one node"
|
||||
)
|
||||
|
||||
|
||||
def _finish_completed_children(run: RunState, barrier: ForeachBarrierState) -> None:
|
||||
for child_id in tuple(barrier.outstanding_frame_ids):
|
||||
child = run.frames[child_id]
|
||||
if child.status == FrameStatus.COMPLETED:
|
||||
barrier.finish_child(child_id)
|
||||
elif child.status == FrameStatus.FAILED:
|
||||
message = child.metadata.get("error", "unknown item failure")
|
||||
raise WorkflowExecutionError(
|
||||
f"concurrent foreach item frame {child_id!r} failed: {message}"
|
||||
)
|
||||
|
||||
|
||||
def _admit_concurrent_children(
|
||||
*,
|
||||
run: RunState,
|
||||
frame: ExecutionFrame,
|
||||
step: ForeachNode,
|
||||
index: WorkflowIndex,
|
||||
barrier: ForeachBarrierState,
|
||||
iterable: list[object],
|
||||
) -> int:
|
||||
if step.concurrent is None:
|
||||
raise WorkflowExecutionError("concurrent foreach requires concurrent policy")
|
||||
|
||||
admitted = 0
|
||||
loop_start = index.next_node_id(frame.node_id, "loop")
|
||||
while (
|
||||
barrier.next_index < len(iterable)
|
||||
and len(barrier.active_frame_ids) < step.concurrent.max_active
|
||||
and len(barrier.outstanding_frame_ids) < step.concurrent.max_outstanding
|
||||
):
|
||||
loop_index = barrier.next_index
|
||||
item = iterable[loop_index]
|
||||
child_id = f"{frame.id}:{step.id}:{loop_index}"
|
||||
active_count = len(barrier.active_frame_ids)
|
||||
barrier.next_index = loop_index + 1
|
||||
barrier.start_child(child_id)
|
||||
add_frame(
|
||||
run,
|
||||
ExecutionFrame(
|
||||
id=child_id,
|
||||
kind="foreach_iteration",
|
||||
node_id=loop_start,
|
||||
status=FrameStatus.PENDING,
|
||||
parent_frame_id=frame.id,
|
||||
metadata=ForeachIterationMetadata(
|
||||
foreach_node_id=step.id,
|
||||
loop_index=loop_index,
|
||||
loop_item=item,
|
||||
loop_alias=step.as_,
|
||||
).to_metadata(),
|
||||
),
|
||||
ready=True,
|
||||
)
|
||||
append_step_result_trace(
|
||||
run,
|
||||
frame_id=frame.id,
|
||||
node_id=frame.node_id,
|
||||
step_type=step.type,
|
||||
next_node_id=loop_start,
|
||||
result=StepExecutionResult(
|
||||
outcome="loop",
|
||||
resolved_input={
|
||||
"item": item,
|
||||
"index": loop_index,
|
||||
"active_count": active_count,
|
||||
},
|
||||
output={},
|
||||
state_changes={},
|
||||
),
|
||||
)
|
||||
admitted += 1
|
||||
return admitted
|
||||
|
||||
|
||||
def _finish_concurrent_foreach(
|
||||
*,
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
frame: ExecutionFrame,
|
||||
step: ForeachNode,
|
||||
index: WorkflowIndex,
|
||||
barrier: ForeachBarrierState,
|
||||
reducers: Mapping[str, ReducerDefinition] | None = None,
|
||||
) -> RunState:
|
||||
next_node_id = index.next_node_id(frame.node_id, "done")
|
||||
combined = build_barrier_patch(
|
||||
workflow,
|
||||
[
|
||||
barrier.pending_results[item_index].patch
|
||||
for item_index in sorted(barrier.pending_results)
|
||||
],
|
||||
run.state,
|
||||
reducers=reducers,
|
||||
)
|
||||
state_changes = commit_state_patch(run.state, combined)
|
||||
append_step_result_trace(
|
||||
run,
|
||||
frame_id=frame.id,
|
||||
node_id=frame.node_id,
|
||||
step_type=step.type,
|
||||
next_node_id=next_node_id,
|
||||
result=StepExecutionResult(
|
||||
outcome="done",
|
||||
resolved_input={
|
||||
"count": barrier.next_index,
|
||||
"index": barrier.next_index,
|
||||
"committed_items": len(barrier.pending_results),
|
||||
},
|
||||
output={},
|
||||
state_changes=state_changes,
|
||||
),
|
||||
)
|
||||
advance_frame(run, frame, outcome="done", next_node_id=next_node_id)
|
||||
return run
|
||||
|
||||
@@ -11,10 +11,12 @@ from wf_core.models.schemas import NodeDef
|
||||
from wf_core.models.steps import InputPathBinding, InputValueBinding, NodeUse
|
||||
from wf_core.models.workflow import Workflow
|
||||
from wf_core.run_state import RunState, RuntimeContext, StepExecutionResult
|
||||
from wf_core.runtime.foreach_state import ForeachBarrierState, item_frame_owner
|
||||
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.schemas import validate_payload_against_schema
|
||||
from wf_core.runtime.ops.state import apply_output_bindings
|
||||
from wf_core.runtime.ops.state import build_output_patch, commit_state_patch
|
||||
|
||||
NodeHandler = Callable[[dict[str, Any], RuntimeContext], NodeResult | dict[str, Any]]
|
||||
AsyncNodeHandler = Callable[
|
||||
@@ -32,6 +34,7 @@ def _resolve_node_execution(
|
||||
) -> tuple[dict[str, Any], RuntimeContext]:
|
||||
frame = run.current_frame()
|
||||
context_values = frame_context_values(frame)
|
||||
state_view = state_view_for_frame(run, frame)
|
||||
resolved_input: dict[str, Any] = {}
|
||||
for binding in node.input:
|
||||
if isinstance(binding, InputValueBinding):
|
||||
@@ -39,7 +42,7 @@ def _resolve_node_execution(
|
||||
elif isinstance(binding, InputPathBinding):
|
||||
value = safe_resolve_path(
|
||||
str(binding.path),
|
||||
state=run.state,
|
||||
state=state_view,
|
||||
workflow_input=run.workflow_input,
|
||||
context=context_values,
|
||||
)
|
||||
@@ -85,13 +88,30 @@ def _finalize_node_execution(
|
||||
validate_payload_against_schema(
|
||||
node_def.output_schema, result.output, f"node output for {node.id}"
|
||||
)
|
||||
state_changes = apply_output_bindings(
|
||||
patch = build_output_patch(
|
||||
workflow,
|
||||
node.output,
|
||||
result.output,
|
||||
run.state,
|
||||
reducers=reducers,
|
||||
)
|
||||
owner = item_frame_owner(run.current_frame())
|
||||
if owner is None:
|
||||
state_changes = commit_state_patch(run.state, patch)
|
||||
else:
|
||||
parent_frame_id, foreach_node_id, item_index = owner
|
||||
parent_frame = run.frames[parent_frame_id]
|
||||
barrier = ForeachBarrierState.from_frame(parent_frame, foreach_node_id)
|
||||
if barrier is not None and barrier.mode == "concurrent":
|
||||
barrier.add_success_patch(
|
||||
index=item_index,
|
||||
frame_id=run.current_frame().id,
|
||||
patch=patch,
|
||||
)
|
||||
barrier.save_to_frame(parent_frame, foreach_node_id)
|
||||
state_changes = {}
|
||||
else:
|
||||
state_changes = commit_state_patch(run.state, patch)
|
||||
return StepExecutionResult(
|
||||
outcome=result.outcome,
|
||||
resolved_input=resolved_input,
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from wf_core.run_state import ExecutionFrame, RunState
|
||||
|
||||
|
||||
def state_view_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
|
||||
"""Return the state view visible to one execution frame.
|
||||
|
||||
This is intentionally a no-op seam for concurrent foreach V1. The first
|
||||
sync-concurrent slice only supports single-node item bodies, so item frames
|
||||
do not need to read their own prior buffered writes yet. The overlay slice
|
||||
should replace this with parent-state plus item-local staged writes.
|
||||
"""
|
||||
return run.state
|
||||
@@ -145,6 +145,46 @@ def commit_state_patch(state: dict[str, Any], patch: StatePatch) -> dict[str, An
|
||||
return dict(patch.changes)
|
||||
|
||||
|
||||
def build_barrier_patch(
|
||||
workflow: Workflow,
|
||||
item_patches: Sequence[StatePatch],
|
||||
state: dict[str, Any],
|
||||
*,
|
||||
reducers: Mapping[str, ReducerDefinition] | None = None,
|
||||
) -> StatePatch:
|
||||
"""Build one committed barrier patch by replaying item writes in order.
|
||||
|
||||
Item patches are built against the parent-visible state. Their prepared
|
||||
writes cannot be blindly merged because reducers must see the value produced
|
||||
by earlier item patches. The barrier therefore replays trace-facing incoming
|
||||
changes against a single staged state in deterministic item order.
|
||||
"""
|
||||
state_fields = workflow.state_schema.field_index()
|
||||
staged_state = deepcopy(state)
|
||||
prepared_patch: dict[StatePath, tuple[list[str], Any]] = {}
|
||||
committed_changes: dict[str, Any] = {}
|
||||
for item_patch in item_patches:
|
||||
for destination, incoming_value in item_patch.changes.items():
|
||||
destination_path = StatePath.parse(destination)
|
||||
key_path, merged_value = prepare_state_value(
|
||||
workflow,
|
||||
staged_state,
|
||||
destination_path,
|
||||
incoming_value,
|
||||
reducers=reducers,
|
||||
state_fields=state_fields,
|
||||
)
|
||||
safe_set_nested_value(staged_state, key_path, merged_value)
|
||||
prepared_patch[destination_path] = (key_path, merged_value)
|
||||
committed_changes[destination] = merged_value
|
||||
validate_staged_state_patch(staged_state, prepared_patch, state_fields)
|
||||
return StatePatch(
|
||||
changes=committed_changes,
|
||||
_prepared_writes=prepared_patch,
|
||||
_staged_state=staged_state,
|
||||
)
|
||||
|
||||
|
||||
def apply_mapped_state(
|
||||
workflow: Workflow,
|
||||
source_data: dict[str, Any],
|
||||
|
||||
Reference in New Issue
Block a user