Lineage state view
This commit is contained in:
@@ -1,13 +1,41 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from collections.abc import Mapping, Sequence
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from wf_core.run_state import ExecutionFrame, RunState
|
from wf_core.run_state import ExecutionFrame, RunState, StateWrite
|
||||||
from wf_core.runtime.foreach_state import ForeachBarrierState, item_frame_owner
|
from wf_core.runtime.foreach_state import ForeachBarrierState, item_frame_owner
|
||||||
from wf_core.runtime.ops.state import safe_set_nested_value
|
from wf_core.runtime.ops.state import safe_set_nested_value
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(slots=True)
|
||||||
|
class LineageStateView:
|
||||||
|
"""Committed state plus writes visible inside one child lineage.
|
||||||
|
|
||||||
|
Today concurrent foreach supplies the writes from barrier metadata. Future
|
||||||
|
native subgraphs and fork/gather should use the same primitive instead of
|
||||||
|
rebuilding foreach-specific overlay logic.
|
||||||
|
"""
|
||||||
|
|
||||||
|
base_state: Mapping[str, Any]
|
||||||
|
writes: Sequence[StateWrite]
|
||||||
|
|
||||||
|
def to_state_dict(self) -> dict[str, Any]:
|
||||||
|
"""Materialize the lineage-visible state as an isolated mutable dict."""
|
||||||
|
# Correctness first: this full copy isolates sibling reads. If state grows
|
||||||
|
# large, replace this with a lazy/copy-on-write overlay.
|
||||||
|
state_view = deepcopy(dict(self.base_state))
|
||||||
|
for write in self.writes:
|
||||||
|
safe_set_nested_value(
|
||||||
|
state_view,
|
||||||
|
list(write.path.parts),
|
||||||
|
write.visible_value,
|
||||||
|
)
|
||||||
|
return state_view
|
||||||
|
|
||||||
|
|
||||||
def state_view_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
|
def state_view_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
|
||||||
"""Return committed parent state plus this frame's item-local overlay.
|
"""Return committed parent state plus this frame's item-local overlay.
|
||||||
|
|
||||||
@@ -29,9 +57,4 @@ def state_view_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]
|
|||||||
if pending is None:
|
if pending is None:
|
||||||
return run.state
|
return run.state
|
||||||
|
|
||||||
# Correctness first: this full copy isolates sibling reads. If state grows
|
return LineageStateView(run.state, pending.patch.writes).to_state_dict()
|
||||||
# large, replace this with a lazy/copy-on-write overlay.
|
|
||||||
state_view = deepcopy(run.state)
|
|
||||||
for write in pending.patch.writes:
|
|
||||||
safe_set_nested_value(state_view, list(write.path.parts), write.visible_value)
|
|
||||||
return state_view
|
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from wf_core.runtime.foreach_state import (
|
|||||||
ItemErrorRecord,
|
ItemErrorRecord,
|
||||||
PendingItemResult,
|
PendingItemResult,
|
||||||
)
|
)
|
||||||
|
from wf_core.runtime.ops.overlays import LineageStateView
|
||||||
from wf_core.runtime.ops.state import StatePatch
|
from wf_core.runtime.ops.state import StatePatch
|
||||||
|
|
||||||
|
|
||||||
@@ -85,6 +86,34 @@ def test_foreach_barrier_state_round_trips_reducer_write_records() -> None:
|
|||||||
assert write.reducer.name == "wf.std.add"
|
assert write.reducer.name == "wf.std.add"
|
||||||
|
|
||||||
|
|
||||||
|
def test_lineage_state_view_materializes_visible_values_without_mutating_base() -> None:
|
||||||
|
base_state = {"count": 2, "nested": {"value": "old"}}
|
||||||
|
view = LineageStateView(
|
||||||
|
base_state,
|
||||||
|
[
|
||||||
|
StateWrite(
|
||||||
|
path=StatePath(("count",)),
|
||||||
|
incoming_value=3,
|
||||||
|
visible_value=5,
|
||||||
|
reducer=ReducerRef(name="wf.std.add"),
|
||||||
|
),
|
||||||
|
StateWrite(
|
||||||
|
path=StatePath(("nested", "value")),
|
||||||
|
incoming_value="new",
|
||||||
|
visible_value="new",
|
||||||
|
reducer=ReducerRef(name="wf.std.replace"),
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
state_view = view.to_state_dict()
|
||||||
|
|
||||||
|
assert state_view["count"] == 5
|
||||||
|
assert state_view["nested"]["value"] == "new"
|
||||||
|
assert base_state["count"] == 2
|
||||||
|
assert base_state["nested"]["value"] == "old"
|
||||||
|
|
||||||
|
|
||||||
def test_foreach_barrier_state_returns_none_when_missing() -> None:
|
def test_foreach_barrier_state_returns_none_when_missing() -> None:
|
||||||
frame = ExecutionFrame(id="root", kind="root", node_id="each")
|
frame = ExecutionFrame(id="root", kind="root", node_id="each")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user