concurrent foreach work

This commit is contained in:
lda
2026-05-22 13:27:05 +07:00 Verified
parent 4a7f7a5bec
commit 0a6500daf2
15 changed files with 1991 additions and 30 deletions
+46
View File
@@ -6,6 +6,7 @@ from typing import Any, Literal
from wf_core.errors import WorkflowExecutionError
from wf_core.run_state import ExecutionFrame
from wf_core.runtime.ops.state import StatePatch
from wf_core.runtime.scheduler import ForeachIterationMetadata
_BARRIER_METADATA_KEY = "foreach_barriers"
@@ -120,6 +121,7 @@ class ForeachBarrierState:
"""Resumable state owned by one foreach parent frame."""
next_index: int = 0
mode: Literal["serial", "concurrent"] = "serial"
active_frame_ids: tuple[str, ...] = ()
outstanding_frame_ids: tuple[str, ...] = ()
pending_results: dict[int, PendingItemResult] = field(default_factory=dict)
@@ -156,11 +158,14 @@ class ForeachBarrierState:
if not isinstance(raw, dict):
raise WorkflowExecutionError("malformed foreach barrier state")
next_index = raw.get("next_index")
mode = raw.get("mode", "serial")
active_frame_ids = _string_tuple(raw.get("active_frame_ids", ()))
outstanding_frame_ids = _string_tuple(raw.get("outstanding_frame_ids", ()))
pending_results = raw.get("pending_results", {})
if not isinstance(next_index, int):
raise WorkflowExecutionError("malformed foreach barrier next_index")
if mode not in {"serial", "concurrent"}:
raise WorkflowExecutionError("malformed foreach barrier mode")
if not isinstance(pending_results, dict):
raise WorkflowExecutionError("malformed foreach barrier pending results")
parsed_results: dict[int, PendingItemResult] = {}
@@ -174,6 +179,7 @@ class ForeachBarrierState:
parsed_results[index] = PendingItemResult.from_metadata(raw_result)
return cls(
next_index=next_index,
mode=mode,
active_frame_ids=active_frame_ids,
outstanding_frame_ids=outstanding_frame_ids,
pending_results=parsed_results,
@@ -196,6 +202,7 @@ class ForeachBarrierState:
def to_metadata(self) -> dict[str, Any]:
return {
"next_index": self.next_index,
"mode": self.mode,
"active_frame_ids": list(self.active_frame_ids),
"outstanding_frame_ids": list(self.outstanding_frame_ids),
"pending_results": {
@@ -204,6 +211,45 @@ class ForeachBarrierState:
},
}
def start_child(self, frame_id: str) -> None:
"""Record one admitted child frame as active and outstanding."""
if frame_id in self.active_frame_ids or frame_id in self.outstanding_frame_ids:
raise WorkflowExecutionError(
f"foreach child frame {frame_id!r} already active"
)
self.active_frame_ids = (*self.active_frame_ids, frame_id)
self.outstanding_frame_ids = (*self.outstanding_frame_ids, frame_id)
def finish_child(self, frame_id: str) -> None:
"""Record one child frame as no longer active or outstanding."""
self.active_frame_ids = tuple(
item for item in self.active_frame_ids if item != frame_id
)
self.outstanding_frame_ids = tuple(
item for item in self.outstanding_frame_ids if item != frame_id
)
def add_success_patch(
self, *, index: int, frame_id: str, patch: StatePatch
) -> None:
"""Buffer one successful item patch by item index."""
self.pending_results[index] = PendingItemResult(
index=index,
frame_id=frame_id,
status="succeeded",
patch=patch,
)
def item_frame_owner(frame: ExecutionFrame) -> tuple[str, str, int] | None:
"""Return parent frame id, foreach node id, and item index for item frames."""
if frame.kind != "foreach_iteration" or frame.parent_frame_id is None:
return None
metadata = ForeachIterationMetadata.from_frame(frame)
if metadata is None:
return None
return frame.parent_frame_id, metadata.foreach_node_id, metadata.loop_index
def _string_tuple(raw: object) -> tuple[str, ...]:
if isinstance(raw, tuple) and all(isinstance(item, str) for item in raw):
+2 -2
View File
@@ -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)
+240 -23
View File
@@ -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
+23 -3
View File
@@ -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,
+16
View File
@@ -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
+40
View File
@@ -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],
+15
View File
@@ -177,6 +177,21 @@ def wake_parent_if_children_complete(run: RunState, child_frame_id: str) -> None
wake_frame(run, parent_id)
def wake_parent_for_child_progress(run: RunState, child_frame_id: str) -> None:
"""Wake a blocked parent after one child finishes so it can refill slots."""
child = _frame(run, child_frame_id)
parent_id = child.parent_frame_id
if parent_id is None:
return
parent = _frame(run, parent_id)
if parent.status != FrameStatus.BLOCKED:
return
block = BlockedOnChildren.from_frame(parent)
if block is None or child_frame_id not in block.child_frame_ids:
return
wake_frame(run, parent_id)
def resolve_no_ready_frames(run: RunState) -> RunStatus:
"""Classify an empty ready queue into terminal, paused, or deadlocked state."""
if run.status == RunStatus.INTERRUPTED:
+2 -2
View File
@@ -99,7 +99,7 @@ def step_workflow(
elif isinstance(step, InterruptNode):
return handle_interrupt_step(run, step)
elif isinstance(step, ForeachNode):
return step_foreach(workflow, run, step, index)
return step_foreach(workflow, run, step, index, reducers=reducers)
else:
raise WorkflowExecutionError(
f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"
@@ -152,7 +152,7 @@ async def step_workflow_async(
elif isinstance(step, InterruptNode):
return handle_interrupt_step(run, step)
elif isinstance(step, ForeachNode):
return step_foreach(workflow, run, step, index)
return step_foreach(workflow, run, step, index, reducers=reducers)
else:
raise WorkflowExecutionError(
f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"