From f69c4502cf981766b7f044721230803a8e33cf35 Mon Sep 17 00:00:00 2001 From: lda Date: Fri, 4 Sep 2026 07:25:55 +0700 Subject: [PATCH] feat: identify dynamic foreach activations --- src/wf_core/runtime/foreach_state.py | 205 +++++++++++++++++- src/wf_core/runtime/lineage.py | 17 +- src/wf_core/runtime/ops/foreach.py | 44 ++-- src/wf_core/runtime/ops/nodes.py | 22 +- src/wf_core/runtime/scheduler.py | 41 +++- src/wf_core/runtime/step.py | 24 +- tests/core/test_concurrent_foreach.py | 50 +++-- tests/core/test_concurrent_foreach_errors.py | 4 +- .../test_concurrent_foreach_interrupts.py | 6 +- tests/core/test_foreach_activations.py | 102 +++++++++ tests/core/test_foreach_barrier_state.py | 38 ++-- tests/core/test_scheduler.py | 22 ++ 12 files changed, 499 insertions(+), 76 deletions(-) create mode 100644 tests/core/test_foreach_activations.py diff --git a/src/wf_core/runtime/foreach_state.py b/src/wf_core/runtime/foreach_state.py index 1da9065f..6aa8236a 100644 --- a/src/wf_core/runtime/foreach_state.py +++ b/src/wf_core/runtime/foreach_state.py @@ -11,6 +11,31 @@ from wf_core.runtime.ops.state import StatePatch from wf_core.runtime.scheduler import ForeachIterationMetadata _BARRIER_METADATA_KEY = "foreach_barriers" +_ACTIVATION_METADATA_KEY = "foreach_activations" + + +@dataclass(slots=True) +class ForeachActivationState: + """Persisted state for one dynamic visit to a foreach controller. + + A parent frame creates a fresh activation on first entry, reuses it while + admitting items, and closes it before emitting ``done``. The id is opaque: + callers compare it by name and never parse it. + """ + + id: str + foreach_node_id: str + barrier: ForeachBarrierState + + +@dataclass(frozen=True, slots=True) +class ForeachItemOwner: + """Named ownership record for one foreach item frame.""" + + parent_frame_id: str + foreach_node_id: str + activation_id: str + item_index: int @dataclass(slots=True) @@ -320,14 +345,188 @@ class ForeachBarrierState: ) -def item_frame_owner(frame: ExecutionFrame) -> tuple[str, str, int] | None: - """Return parent frame id, foreach node id, and item index for item frames.""" +def load_or_begin_foreach_activation( + frame: ExecutionFrame, + foreach_node_id: str, + *, + mode: Literal["serial", "concurrent"], +) -> ForeachActivationState: + """Load the active activation or begin a fresh visit. + + The first entry for one visit allocates an opaque id from the parent frame + id, foreach node id, and a persisted per-frame sequence. Later calls reuse + the active activation; closing it makes the next visit allocate a new id + with fresh barrier state. Mode mismatches and malformed tables fail fast. + """ + table = _activation_table(frame) + entry = table.get(foreach_node_id) + if entry is None: + entry = {"next_sequence": 0, "active": None} + table[foreach_node_id] = entry + if not isinstance(entry, dict): + raise WorkflowExecutionError( + f"malformed foreach activation entry for frame {frame.id!r}" + ) + next_sequence = entry.get("next_sequence", 0) + if not isinstance(next_sequence, int) or next_sequence < 0: + raise WorkflowExecutionError( + f"malformed foreach activation sequence for frame {frame.id!r}" + ) + active = entry.get("active") + if active is not None: + activation = _activation_from_metadata( + active, frame_id=frame.id, foreach_node_id=foreach_node_id + ) + if activation.barrier.mode != mode: + raise WorkflowExecutionError( + f"foreach {foreach_node_id!r} activation {activation.id!r} " + f"has mode {activation.barrier.mode!r}, got {mode!r}" + ) + return activation + activation_id = f"{frame.id}:{foreach_node_id}#{next_sequence}" + activation = ForeachActivationState( + id=activation_id, + foreach_node_id=foreach_node_id, + barrier=ForeachBarrierState(mode=mode), + ) + entry["next_sequence"] = next_sequence + 1 + entry["active"] = { + "id": activation.id, + "barrier": activation.barrier.to_metadata(), + } + return activation + + +def save_foreach_activation( + frame: ExecutionFrame, activation: ForeachActivationState +) -> None: + """Persist barrier progress for the named active activation.""" + table = _activation_table(frame) + entry = table.get(activation.foreach_node_id) + if not isinstance(entry, dict): + raise WorkflowExecutionError( + f"malformed foreach activation entry for frame {frame.id!r}" + ) + active = entry.get("active") + if not isinstance(active, dict) or active.get("id") != activation.id: + raise WorkflowExecutionError( + f"cannot save stale foreach activation {activation.id!r} " + f"for frame {frame.id!r}" + ) + active["barrier"] = activation.barrier.to_metadata() + + +def close_foreach_activation( + frame: ExecutionFrame, activation: ForeachActivationState +) -> None: + """Close the named active activation, preserving the visit sequence. + + The barrier is removed so a later visit starts fresh; the sequence keeps + increasing so child and lineage ids cannot collide across visits. + """ + table = _activation_table(frame) + entry = table.get(activation.foreach_node_id) + if not isinstance(entry, dict): + raise WorkflowExecutionError( + f"malformed foreach activation entry for frame {frame.id!r}" + ) + active = entry.get("active") + if not isinstance(active, dict) or active.get("id") != activation.id: + raise WorkflowExecutionError( + f"cannot close stale foreach activation {activation.id!r} " + f"for frame {frame.id!r}" + ) + entry["active"] = None + + +def load_foreach_activation( + frame: ExecutionFrame, foreach_node_id: str, activation_id: str +) -> ForeachActivationState | None: + """Return the active activation only when its id matches the child. + + A child result naming a closed or different activation must fail closed in + the caller rather than buffering into the wrong barrier. + """ + table = _activation_table(frame) + entry = table.get(foreach_node_id) + if not isinstance(entry, dict): + raise WorkflowExecutionError( + f"malformed foreach activation entry for frame {frame.id!r}" + ) + active = entry.get("active") + if active is None: + return None + activation = _activation_from_metadata( + active, frame_id=frame.id, foreach_node_id=foreach_node_id + ) + if activation.id != activation_id: + return None + return activation + + +def require_foreach_activation( + frame: ExecutionFrame, foreach_node_id: str, activation_id: str +) -> ForeachActivationState: + """Load the named activation or raise when it is closed or superseded.""" + activation = load_foreach_activation(frame, foreach_node_id, activation_id) + if activation is None: + raise WorkflowExecutionError( + f"foreach item activation {activation_id!r} for node " + f"{foreach_node_id!r} is closed or superseded" + ) + return activation + + +def item_frame_owner(frame: ExecutionFrame) -> ForeachItemOwner | None: + """Return the named foreach ownership record for item frames. + + Malformed item metadata fails closed via ``ForeachIterationMetadata``; + only non-item frames return ``None``. + """ 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 + return ForeachItemOwner( + parent_frame_id=frame.parent_frame_id, + foreach_node_id=metadata.foreach_node_id, + activation_id=metadata.activation_id, + item_index=metadata.loop_index, + ) + + +def _activation_table(frame: ExecutionFrame) -> dict[str, Any]: + raw = frame.metadata.get(_ACTIVATION_METADATA_KEY) + if raw is None: + table: dict[str, Any] = {} + frame.metadata[_ACTIVATION_METADATA_KEY] = table + return table + if not isinstance(raw, dict): + raise WorkflowExecutionError( + f"malformed foreach activation table for frame {frame.id!r}" + ) + return raw + + +def _activation_from_metadata( + raw: object, *, frame_id: str, foreach_node_id: str +) -> ForeachActivationState: + if not isinstance(raw, dict): + raise WorkflowExecutionError( + f"malformed foreach activation for frame {frame_id!r}" + ) + activation_id = raw.get("id") + barrier_raw = raw.get("barrier") + if not isinstance(activation_id, str) or not activation_id: + raise WorkflowExecutionError( + f"malformed foreach activation id for frame {frame_id!r}" + ) + return ForeachActivationState( + id=activation_id, + foreach_node_id=foreach_node_id, + barrier=ForeachBarrierState.from_metadata(barrier_raw), + ) def _string_tuple(raw: object) -> tuple[str, ...]: diff --git a/src/wf_core/runtime/lineage.py b/src/wf_core/runtime/lineage.py index d2d3d97e..800d82a9 100644 --- a/src/wf_core/runtime/lineage.py +++ b/src/wf_core/runtime/lineage.py @@ -7,7 +7,7 @@ from typing import Any from wf_core.errors import WorkflowExecutionError from wf_core.run_state import ExecutionFrame, LineageState, RunState, StateWrite -from wf_core.runtime.foreach_state import ForeachBarrierState, item_frame_owner +from wf_core.runtime.foreach_state import item_frame_owner, load_foreach_activation from wf_core.runtime.ops.state import ( StatePatch, commit_state_patch, @@ -61,21 +61,24 @@ def lineage_writes_for_frame( # Compatibility fallback: concurrent foreach used barrier-local patches # before `RunState.lineages` became the primary write store. Keep reading # those patches so old serialized runs and direct barrier tests still work. + # Barrier lookup includes the activation so a stale visit cannot read a + # later activation's buffered writes. owner = item_frame_owner(frame) if owner is None: return () - parent_frame_id, foreach_node_id, item_index = owner - parent_frame = run.frames.get(parent_frame_id) + parent_frame = run.frames.get(owner.parent_frame_id) if parent_frame is None: raise WorkflowExecutionError( "foreach lineage compatibility state references missing parent frame " - f"{parent_frame_id!r} for child frame {frame.id!r}" + f"{owner.parent_frame_id!r} for child frame {frame.id!r}" ) - barrier = ForeachBarrierState.from_frame(parent_frame, foreach_node_id) - if barrier is None or barrier.mode != "concurrent": + activation = load_foreach_activation( + parent_frame, owner.foreach_node_id, owner.activation_id + ) + if activation is None or activation.barrier.mode != "concurrent": return () - pending = barrier.pending_results.get(item_index) + pending = activation.barrier.pending_results.get(owner.item_index) if pending is None: return () return pending.patch.writes diff --git a/src/wf_core/runtime/ops/foreach.py b/src/wf_core/runtime/ops/foreach.py index db38d71c..b45b15ef 100644 --- a/src/wf_core/runtime/ops/foreach.py +++ b/src/wf_core/runtime/ops/foreach.py @@ -8,9 +8,12 @@ from wf_core.models.steps import ForeachNode from wf_core.models.workflow import Workflow from wf_core.run_state import ExecutionFrame, FrameStatus, RunState, StepExecutionResult from wf_core.runtime.foreach_state import ( + ForeachActivationState, ForeachBarrierState, ItemErrorRecord, PendingItemResult, + load_or_begin_foreach_activation, + save_foreach_activation, ) from wf_core.runtime.lineage import ( add_lineage, @@ -63,7 +66,8 @@ def _step_foreach_serial( raise WorkflowExecutionError("serial foreach helper received non-serial mode") frame = run.current_frame() - barrier = ForeachBarrierState.from_frame(frame, step.id) or ForeachBarrierState() + activation = load_or_begin_foreach_activation(frame, step.id, mode="serial") + barrier = activation.barrier iterable = _resolve_foreach_iterable(run, frame, step) loop_index = barrier.next_index @@ -89,9 +93,9 @@ def _step_foreach_serial( 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_lineage_id = _child_lineage_id(frame, step, loop_index) + save_foreach_activation(frame, activation) + child_id = _child_frame_id(activation, loop_index) + child_lineage_id = _child_lineage_id(activation, loop_index) add_frame( run, ExecutionFrame( @@ -105,6 +109,7 @@ def _step_foreach_serial( parent_lineage_id=frame.lineage_id, metadata=ForeachIterationMetadata( foreach_node_id=step.id, + activation_id=activation.id, loop_index=loop_index, loop_item=item, loop_alias=step.as_, @@ -141,11 +146,8 @@ def _step_foreach_concurrent( if step.concurrent is None: raise WorkflowExecutionError("concurrent foreach requires concurrent policy") 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") + activation = load_or_begin_foreach_activation(frame, step.id, mode="concurrent") + barrier = activation.barrier _finish_completed_children(run, step, barrier) iterable = _resolve_foreach_iterable(run, frame, step) @@ -154,6 +156,7 @@ def _step_foreach_concurrent( frame=frame, step=step, index=index, + activation=activation, barrier=barrier, iterable=iterable, ) @@ -169,7 +172,7 @@ def _step_foreach_concurrent( reducers=reducers, ) - barrier.save_to_frame(frame, step.id) + save_foreach_activation(frame, activation) block_frame_on_children(run, frame.id, barrier.outstanding_frame_ids) run.sync_from_current_frame() return run @@ -242,6 +245,7 @@ def _admit_concurrent_children( frame: ExecutionFrame, step: ForeachNode, index: WorkflowIndex, + activation: ForeachActivationState, barrier: ForeachBarrierState, iterable: list[object], ) -> None: @@ -256,8 +260,8 @@ def _admit_concurrent_children( ): loop_index = barrier.next_index item = iterable[loop_index] - child_id = f"{frame.id}:{step.id}:{loop_index}" - child_lineage_id = _child_lineage_id(frame, step, loop_index) + child_id = _child_frame_id(activation, loop_index) + child_lineage_id = _child_lineage_id(activation, loop_index) add_lineage( run, scope_id=frame.scope_id, @@ -280,6 +284,7 @@ def _admit_concurrent_children( parent_lineage_id=frame.lineage_id, metadata=ForeachIterationMetadata( foreach_node_id=step.id, + activation_id=activation.id, loop_index=loop_index, loop_item=item, loop_alias=step.as_, @@ -372,13 +377,22 @@ def _finish_concurrent_foreach( return run -def _child_lineage_id(frame: ExecutionFrame, step: ForeachNode, loop_index: int) -> str: - """Return a deterministic opaque lineage id for one foreach child frame. +def _child_frame_id(activation: ForeachActivationState, loop_index: int) -> str: + """Return a deterministic opaque child frame id for one activation item. + + The id embeds the activation so a later visit at item zero cannot collide + with the first visit. Compare full ids; never parse them. + """ + return f"{activation.id}:{loop_index}" + + +def _child_lineage_id(activation: ForeachActivationState, loop_index: int) -> str: + """Return a deterministic opaque lineage id for one activation item. The readable shape is only for diagnostics. Runtime code should compare the full id, not parse it; future structured lineage refs can replace this. """ - return f"{frame.lineage_id}/{step.id}[{loop_index}]" + return f"{activation.id}[{loop_index}]" def _patch_for_successful_item( diff --git a/src/wf_core/runtime/ops/nodes.py b/src/wf_core/runtime/ops/nodes.py index 4df383bf..81d6165e 100644 --- a/src/wf_core/runtime/ops/nodes.py +++ b/src/wf_core/runtime/ops/nodes.py @@ -15,7 +15,11 @@ from wf_core.run_state import ( RuntimeContext, StepExecutionResult, ) -from wf_core.runtime.foreach_state import ForeachBarrierState, item_frame_owner +from wf_core.runtime.foreach_state import ( + item_frame_owner, + require_foreach_activation, + save_foreach_activation, +) from wf_core.runtime.input_bindings import resolve_step_input_bindings from wf_core.runtime.lineage import ( append_lineage_writes, @@ -116,10 +120,14 @@ def _finalize_node_execution( if owner is None: state_changes = commit_patch_for_frame(run, frame, 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": + parent_frame = run.frames[owner.parent_frame_id] + # Fail closed when the child names a closed or superseded activation: + # its writes must not land in a later visit's barrier. + activation = require_foreach_activation( + parent_frame, owner.foreach_node_id, owner.activation_id + ) + barrier = activation.barrier + if barrier.mode == "concurrent": # New concurrent foreach stores writes in the child lineage; the # barrier keeps only result metadata plus old patch fallback. append_lineage_writes( @@ -129,12 +137,12 @@ def _finalize_node_execution( writes=patch.writes, ) barrier.add_success_patch( - index=item_index, + index=owner.item_index, frame_id=frame.id, patch=StatePatch(), lineage_id=frame.lineage_id, ) - barrier.save_to_frame(parent_frame, foreach_node_id) + save_foreach_activation(parent_frame, activation) state_changes = {} else: state_changes = commit_patch_for_frame(run, parent_frame, patch) diff --git a/src/wf_core/runtime/scheduler.py b/src/wf_core/runtime/scheduler.py index 002c15d7..a8d2b896 100644 --- a/src/wf_core/runtime/scheduler.py +++ b/src/wf_core/runtime/scheduler.py @@ -38,9 +38,15 @@ class BlockedOnChildren: @dataclass(slots=True, frozen=True) class ForeachIterationMetadata: - """Typed metadata for a foreach iteration frame.""" + """Typed metadata for a foreach iteration frame. + + ``activation_id`` names the dynamic foreach visit that owns this item. + It separates fresh barrier state from earlier visits to the same node use + and must survive checkpoint serialization. + """ foreach_node_id: str + activation_id: str loop_index: int loop_item: Any loop_alias: str @@ -51,12 +57,17 @@ class ForeachIterationMetadata: return None metadata = frame.metadata foreach_node_id = metadata.get("foreach_node_id") + activation_id = metadata.get("activation_id") loop_index = metadata.get("loop_index") loop_alias = metadata.get("loop_alias") if not isinstance(foreach_node_id, str) or not foreach_node_id: raise WorkflowExecutionError( f"malformed foreach node id for frame {frame.id!r}" ) + if not isinstance(activation_id, str) or not activation_id: + raise WorkflowExecutionError( + f"malformed foreach activation id for frame {frame.id!r}" + ) if not isinstance(loop_index, int): raise WorkflowExecutionError( f"malformed foreach loop index for frame {frame.id!r}" @@ -71,6 +82,7 @@ class ForeachIterationMetadata: ) return cls( foreach_node_id=foreach_node_id, + activation_id=activation_id, loop_index=loop_index, loop_item=metadata["loop_item"], loop_alias=loop_alias, @@ -79,6 +91,7 @@ class ForeachIterationMetadata: def to_metadata(self) -> dict[str, object]: return { "foreach_node_id": self.foreach_node_id, + "activation_id": self.activation_id, "loop_index": self.loop_index, "loop_item": self.loop_item, "loop_alias": self.loop_alias, @@ -178,7 +191,11 @@ def wake_parent_if_children_complete(run: RunState, child_frame_id: str) -> None 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.""" + """Wake a blocked parent after one child finishes so it can refill slots. + + The wake-up includes the foreach activation: a child naming a closed or + superseded activation cannot wake a parent waiting on a later visit. + """ child = _frame(run, child_frame_id) parent_id = child.parent_frame_id if parent_id is None: @@ -189,6 +206,26 @@ def wake_parent_for_child_progress(run: RunState, child_frame_id: str) -> None: block = BlockedOnChildren.from_frame(parent) if block is None or child_frame_id not in block.child_frame_ids: return + # Lazy import avoids a cycle: foreach_state owns activation persistence on + # top of this scheduler's frame metadata types. + from wf_core.runtime.foreach_state import ( + item_frame_owner, + load_foreach_activation, + ) + + try: + owner = item_frame_owner(child) + except WorkflowExecutionError: + raise + if owner is not None: + activation = load_foreach_activation( + parent, owner.foreach_node_id, owner.activation_id + ) + if activation is None: + raise WorkflowExecutionError( + f"foreach item frame {child_frame_id!r} names closed activation " + f"{owner.activation_id!r} and cannot wake parent {parent_id!r}" + ) wake_frame(run, parent_id) diff --git a/src/wf_core/runtime/step.py b/src/wf_core/runtime/step.py index 730ebb1d..ddea7f87 100644 --- a/src/wf_core/runtime/step.py +++ b/src/wf_core/runtime/step.py @@ -16,7 +16,7 @@ from wf_core.models.steps import ( ) 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, item_frame_owner +from wf_core.runtime.foreach_state import item_frame_owner, load_foreach_activation from wf_core.runtime.ops.flow import advance_frame, append_step_result_trace from wf_core.runtime.ops.foreach import step_foreach from wf_core.runtime.ops.handlers import ( @@ -361,10 +361,15 @@ def _claim_matching_async_item_frames( index: WorkflowIndex, first_frame: ExecutionFrame, ) -> list[ExecutionFrame]: + """Claim sibling item frames from the same activation for async batching. + + Batching never mixes activations: only frames naming the same parent, + foreach, and activation id run together, preserving deterministic barrier + commits across revisits. + """ owner = item_frame_owner(first_frame) if owner is None: return [] - parent_frame_id, foreach_node_id, _item_index = owner claimed: list[ExecutionFrame] = [] remaining_ready: list[str] = [] for frame_id in run.ready_frame_ids: @@ -373,7 +378,9 @@ def _claim_matching_async_item_frames( if ( frame.status == FrameStatus.PENDING and frame_owner is not None - and frame_owner[:2] == (parent_frame_id, foreach_node_id) + and frame_owner.parent_frame_id == owner.parent_frame_id + and frame_owner.foreach_node_id == owner.foreach_node_id + and frame_owner.activation_id == owner.activation_id and isinstance(index.nodes_by_id.get(frame.node_id), NodeUse) ): frame.status = FrameStatus.RUNNING @@ -392,14 +399,15 @@ def _can_batch_async_foreach_item( owner = item_frame_owner(frame) if owner is None: return False - parent_frame_id, foreach_node_id, _item_index = owner - parent_frame = run.frames.get(parent_frame_id) + parent_frame = run.frames.get(owner.parent_frame_id) if parent_frame is None: return False - barrier = ForeachBarrierState.from_frame(parent_frame, foreach_node_id) + activation = load_foreach_activation( + parent_frame, owner.foreach_node_id, owner.activation_id + ) return ( - barrier is not None - and barrier.mode == "concurrent" + activation is not None + and activation.barrier.mode == "concurrent" and isinstance(index.nodes_by_id.get(frame.node_id), NodeUse) ) diff --git a/tests/core/test_concurrent_foreach.py b/tests/core/test_concurrent_foreach.py index 2ba18c59..c73afdb6 100644 --- a/tests/core/test_concurrent_foreach.py +++ b/tests/core/test_concurrent_foreach.py @@ -19,7 +19,11 @@ from wf_core import ( execute_workflow, ) from wf_core.run_state import ExecutionFrame, RunState, RuntimeContext -from wf_core.runtime.foreach_state import ForeachBarrierState +from wf_core.runtime.foreach_state import ( + ForeachItemOwner, + item_frame_owner, + load_or_begin_foreach_activation, +) from wf_core.runtime.scheduler import ForeachIterationMetadata @@ -143,7 +147,11 @@ def test_concurrent_foreach_item_frames_use_distinct_lineages() -> None: assert run.frames["root"].lineage_id == "root" assert run.frames["root"].parent_lineage_id is None assert len(item_frames) == 2 - assert item_lineage_ids == {"root/each[0]", "root/each[1]"} + assert item_lineage_ids == {"root:each#0[0]", "root:each#0[1]"} + for frame in item_frames: + owner = item_frame_owner(frame) + assert isinstance(owner, ForeachItemOwner) + assert owner.activation_id == "root:each#0" assert set(context_lineage_ids) == item_lineage_ids assert all(frame.scope_id == "root" for frame in item_frames) assert all(frame.parent_lineage_id == "root" for frame in item_frames) @@ -162,15 +170,27 @@ def test_nested_concurrent_foreach_records_parent_child_lineages() -> None: inner_frames = _foreach_frames(run, "inner_each") assert {frame.lineage_id for frame in outer_frames} == { - "root/outer_each[0]", - "root/outer_each[1]", + "root:outer_each#0[0]", + "root:outer_each#0[1]", } assert all(frame.parent_lineage_id == "root" for frame in outer_frames) assert {(frame.parent_lineage_id, frame.lineage_id) for frame in inner_frames} == { - ("root/outer_each[0]", "root/outer_each[0]/inner_each[0]"), - ("root/outer_each[0]", "root/outer_each[0]/inner_each[1]"), - ("root/outer_each[1]", "root/outer_each[1]/inner_each[0]"), - ("root/outer_each[1]", "root/outer_each[1]/inner_each[1]"), + ( + "root:outer_each#0[0]", + "root:outer_each#0:0:inner_each#0[0]", + ), + ( + "root:outer_each#0[0]", + "root:outer_each#0:0:inner_each#0[1]", + ), + ( + "root:outer_each#0[1]", + "root:outer_each#0:1:inner_each#0[0]", + ), + ( + "root:outer_each#0[1]", + "root:outer_each#0:1:inner_each#0[1]", + ), } @@ -264,12 +284,14 @@ def test_sync_concurrent_foreach_barrier_replays_add_reducer_inputs() -> None: assert run.state["number"] == 6 assert run.output["number"] == 6 - assert run.lineages["root/each[0]"].writes[0].incoming_value == 3 - assert run.lineages["root/each[1]"].writes[0].incoming_value == 1 - barrier = ForeachBarrierState.from_frame(run.frames["root"], "each") - assert barrier is not None - assert barrier.pending_results[0].lineage_id == "root/each[0]" - assert barrier.pending_results[0].patch.writes == [] + assert run.lineages["root:each#0[0]"].writes[0].incoming_value == 3 + assert run.lineages["root:each#0[1]"].writes[0].incoming_value == 1 + active = load_or_begin_foreach_activation( + run.frames["root"], "each", mode="concurrent" + ) + assert active.id == "root:each#0" + assert active.barrier.pending_results[0].lineage_id == "root:each#0[0]" + assert active.barrier.pending_results[0].patch.writes == [] foreach_entries = [entry for entry in run.trace if entry.step_type == "foreach"] assert foreach_entries[-1].state_changes["state.number"] == 6 diff --git a/tests/core/test_concurrent_foreach_errors.py b/tests/core/test_concurrent_foreach_errors.py index 609f9508..18c1e3a4 100644 --- a/tests/core/test_concurrent_foreach_errors.py +++ b/tests/core/test_concurrent_foreach_errors.py @@ -27,7 +27,7 @@ def test_concurrent_foreach_skip_emits_completed_with_errors() -> None: ) assert run.state["seen"] == ["a", "c"] - assert run.frames["root:each:1"].status == "failed" + assert run.frames["root:each#0:1"].status == "failed" foreach_entries = [entry for entry in run.trace if entry.step_type == "foreach"] assert foreach_entries[-1].outcome == "completed_with_errors" assert foreach_entries[-1].resolved_input["failed_items"] == 1 @@ -47,7 +47,7 @@ def test_concurrent_foreach_collect_writes_ordered_error_records() -> None: assert len(run.state["errors"]) == 1 error = run.state["errors"][0] assert error["index"] == 1 - assert error["frame_id"] == "root:each:1" + assert error["frame_id"] == "root:each#0:1" assert error["node_id"] == "record" assert error["error_type"] == "ValueError" assert error["message"] == "bad item" diff --git a/tests/core/test_concurrent_foreach_interrupts.py b/tests/core/test_concurrent_foreach_interrupts.py index f171e700..3463e473 100644 --- a/tests/core/test_concurrent_foreach_interrupts.py +++ b/tests/core/test_concurrent_foreach_interrupts.py @@ -31,8 +31,8 @@ async def test_concurrent_foreach_interrupt_returns_before_refill() -> None: assert run.status is RunStatus.INTERRUPTED assert run.interrupt is not None assert run.interrupt.payload["item"] == "b" - assert run.frames["root:each:1"].status == "interrupted" - assert "root:each:2" not in run.frames + assert run.frames["root:each#0:1"].status == "interrupted" + assert "root:each#0:2" not in run.frames assert "seen" not in run.state @@ -54,7 +54,7 @@ async def test_resume_prioritizes_interrupted_item_before_siblings() -> None: assert resumed.status is RunStatus.COMPLETED assert resumed.state["seen"] == ["a", "b", "c"] - assert resumed.trace[interrupted_trace_len].frame_id == "root:each:1" + assert resumed.trace[interrupted_trace_len].frame_id == "root:each#0:1" assert resumed.trace[interrupted_trace_len].step_type == "interrupt" assert resumed.trace[interrupted_trace_len].outcome == "submitted" foreach_entries = [entry for entry in resumed.trace if entry.step_type == "foreach"] diff --git a/tests/core/test_foreach_activations.py b/tests/core/test_foreach_activations.py new file mode 100644 index 00000000..cb94d0f6 --- /dev/null +++ b/tests/core/test_foreach_activations.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import pytest + +from wf_core.errors import WorkflowExecutionError +from wf_core.run_state import ExecutionFrame +from wf_core.runtime.foreach_state import ( + close_foreach_activation, + item_frame_owner, + load_or_begin_foreach_activation, + save_foreach_activation, +) +from wf_core.runtime.scheduler import ForeachIterationMetadata + + +def _frame() -> ExecutionFrame: + return ExecutionFrame(id="root", kind="workflow", node_id="each") + + +def test_activation_lifecycle_reuses_active_then_fresh_after_close() -> None: + frame = _frame() + + first = load_or_begin_foreach_activation(frame, "each", mode="serial") + save_foreach_activation(frame, first) + restored = load_or_begin_foreach_activation(frame, "each", mode="serial") + + assert restored.id == first.id + + close_foreach_activation(frame, restored) + second = load_or_begin_foreach_activation(frame, "each", mode="serial") + + assert second.id != first.id + assert second.barrier.next_index == 0 + + +def test_activation_rejects_malformed_metadata() -> None: + frame = ExecutionFrame( + id="root", + kind="workflow", + node_id="each", + metadata={"foreach_activations": "corrupt"}, + ) + + with pytest.raises(WorkflowExecutionError, match="activation"): + load_or_begin_foreach_activation(frame, "each", mode="serial") + + +def test_activation_rejects_mode_mismatch() -> None: + frame = _frame() + activation = load_or_begin_foreach_activation(frame, "each", mode="serial") + save_foreach_activation(frame, activation) + + with pytest.raises(WorkflowExecutionError, match="mode"): + load_or_begin_foreach_activation(frame, "each", mode="concurrent") + + +def test_closing_stale_activation_fails_closed() -> None: + frame = _frame() + first = load_or_begin_foreach_activation(frame, "each", mode="serial") + save_foreach_activation(frame, first) + close_foreach_activation(frame, first) + second = load_or_begin_foreach_activation(frame, "each", mode="serial") + save_foreach_activation(frame, second) + + with pytest.raises(WorkflowExecutionError, match="stale|closed|active"): + close_foreach_activation(frame, first) + + +def test_activation_json_round_trip_through_frame_metadata() -> None: + frame = _frame() + activation = load_or_begin_foreach_activation(frame, "each", mode="serial") + activation.barrier.next_index = 2 + save_foreach_activation(frame, activation) + + dumped = dict(frame.metadata) + restored_frame = ExecutionFrame( + id="root", kind="workflow", node_id="each", metadata=dumped + ) + restored = load_or_begin_foreach_activation(restored_frame, "each", mode="serial") + + assert restored.id == activation.id + assert restored.barrier.next_index == 2 + + +def test_item_metadata_requires_activation_identity() -> None: + frame = ExecutionFrame( + id="root:each#0:0", + kind="foreach_iteration", + node_id="work", + parent_frame_id="root", + metadata={ + "foreach_node_id": "each", + "loop_index": 0, + "loop_item": "a", + "loop_alias": "item", + }, + ) + + with pytest.raises(WorkflowExecutionError, match="activation"): + ForeachIterationMetadata.from_frame(frame) + with pytest.raises(WorkflowExecutionError, match="activation"): + item_frame_owner(frame) diff --git a/tests/core/test_foreach_barrier_state.py b/tests/core/test_foreach_barrier_state.py index a428648d..9bf44abf 100644 --- a/tests/core/test_foreach_barrier_state.py +++ b/tests/core/test_foreach_barrier_state.py @@ -8,9 +8,13 @@ from wf_core.paths import StatePath from wf_core.run_state import ExecutionFrame, RunState, RunStatus, StateWrite from wf_core.runtime.foreach_state import ( ForeachBarrierState, + ForeachItemOwner, ItemErrorRecord, PendingItemResult, _state_write_from_metadata, + item_frame_owner, + load_or_begin_foreach_activation, + save_foreach_activation, ) from wf_core.runtime.lineage import LineageStateView, lineage_writes_for_frame from wf_core.runtime.ops.state import StatePatch @@ -134,20 +138,28 @@ def test_lineage_state_view_materializes_visible_values_without_mutating_base() def test_lineage_writes_for_frame_reads_current_foreach_pending_result() -> None: parent = ExecutionFrame(id="root", kind="workflow", node_id="each") + activation = load_or_begin_foreach_activation(parent, "each", mode="concurrent") + child_lineage_id = f"{activation.id}[0]" child = ExecutionFrame( - id="root:each:0", + id=f"{activation.id}:0", kind="foreach_iteration", node_id="work", parent_frame_id="root", - lineage_id="root/each[0]", + lineage_id=child_lineage_id, parent_lineage_id="root", metadata={ "foreach_node_id": "each", + "activation_id": activation.id, "loop_index": 0, "loop_item": "a", "loop_alias": "item", }, ) + # Ownership is named, not positional. + owner = item_frame_owner(child) + assert isinstance(owner, ForeachItemOwner) + assert owner.activation_id == activation.id + assert owner.item_index == 0 patch = StatePatch( writes=[ StateWrite( @@ -158,19 +170,14 @@ def test_lineage_writes_for_frame_reads_current_foreach_pending_result() -> None ) ] ) - barrier = ForeachBarrierState( - mode="concurrent", - pending_results={ - 0: PendingItemResult( - index=0, - frame_id=child.id, - status="succeeded", - lineage_id=child.lineage_id, - patch=patch, - ) - }, + activation.barrier.pending_results[0] = PendingItemResult( + index=0, + frame_id=child.id, + status="succeeded", + lineage_id=child.lineage_id, + patch=patch, ) - barrier.save_to_frame(parent, "each") + save_foreach_activation(parent, activation) run = RunState( workflow_name="lineage", status=RunStatus.PENDING, @@ -188,12 +195,13 @@ def test_lineage_writes_for_frame_reads_current_foreach_pending_result() -> None def test_lineage_writes_for_frame_rejects_missing_compatibility_parent_frame() -> None: child = ExecutionFrame( - id="missing:each:0", + id="missing:each#0:0", kind="foreach_iteration", node_id="work", parent_frame_id="missing", metadata={ "foreach_node_id": "each", + "activation_id": "missing:each#0", "loop_index": 0, "loop_item": "a", "loop_alias": "item", diff --git a/tests/core/test_scheduler.py b/tests/core/test_scheduler.py index 741d24e7..b88689aa 100644 --- a/tests/core/test_scheduler.py +++ b/tests/core/test_scheduler.py @@ -163,8 +163,19 @@ def test_child_completion_wakes_blocked_parent() -> None: def test_wake_parent_when_child_finishes_for_refill() -> None: + from wf_core.runtime.foreach_state import ( + ForeachItemOwner, + item_frame_owner, + load_or_begin_foreach_activation, + save_foreach_activation, + ) + run = _run() add_frame(run, ExecutionFrame(id="parent", kind="root", node_id="foreach")) + activation = load_or_begin_foreach_activation( + run.frames["parent"], "foreach", mode="serial" + ) + save_foreach_activation(run.frames["parent"], activation) add_frame( run, ExecutionFrame( @@ -172,11 +183,22 @@ def test_wake_parent_when_child_finishes_for_refill() -> None: kind="foreach_iteration", node_id="__end__", parent_frame_id="parent", + metadata={ + "foreach_node_id": "foreach", + "activation_id": activation.id, + "loop_index": 0, + "loop_item": "a", + "loop_alias": "item", + }, ), ) block_frame_on_children(run, "parent", ("child", "other")) run.frames["child"].status = FrameStatus.COMPLETED + owner = item_frame_owner(run.frames["child"]) + assert isinstance(owner, ForeachItemOwner) + assert owner.activation_id == activation.id + wake_parent_for_child_progress(run, "child") assert run.frames["parent"].status == FrameStatus.PENDING