scope aware state views
This commit is contained in:
@@ -46,6 +46,14 @@ def lineage_writes_for_frame(
|
|||||||
lookup here gives future `RunState.lineages` or subgraph scopes one place to
|
lookup here gives future `RunState.lineages` or subgraph scopes one place to
|
||||||
plug in without making node execution understand foreach internals.
|
plug in without making node execution understand foreach internals.
|
||||||
"""
|
"""
|
||||||
|
lineage = run.lineages.get(frame.lineage_id)
|
||||||
|
if lineage is not None and lineage.scope_id == frame.scope_id and lineage.writes:
|
||||||
|
return tuple(
|
||||||
|
lineage_state_writes(
|
||||||
|
run, scope_id=frame.scope_id, lineage_id=frame.lineage_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
owner = item_frame_owner(frame)
|
owner = item_frame_owner(frame)
|
||||||
if owner is None:
|
if owner is None:
|
||||||
return ()
|
return ()
|
||||||
@@ -61,6 +69,14 @@ def lineage_writes_for_frame(
|
|||||||
return pending.patch.writes
|
return pending.patch.writes
|
||||||
|
|
||||||
|
|
||||||
|
def scope_state_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
|
||||||
|
"""Return the committed state root for the frame's runtime scope."""
|
||||||
|
scope = run.scopes.get(frame.scope_id)
|
||||||
|
if scope is None:
|
||||||
|
raise ValueError(f"unknown scope {frame.scope_id!r}")
|
||||||
|
return scope.committed_state
|
||||||
|
|
||||||
|
|
||||||
def add_lineage(
|
def add_lineage(
|
||||||
run: RunState,
|
run: RunState,
|
||||||
*,
|
*,
|
||||||
@@ -121,6 +137,17 @@ def lineage_state_view(
|
|||||||
return LineageStateView(scope.committed_state, writes).to_state_dict()
|
return LineageStateView(scope.committed_state, writes).to_state_dict()
|
||||||
|
|
||||||
|
|
||||||
|
def lineage_state_writes(
|
||||||
|
run: RunState,
|
||||||
|
*,
|
||||||
|
scope_id: str,
|
||||||
|
lineage_id: str,
|
||||||
|
) -> Iterator[StateWrite]:
|
||||||
|
"""Yield ancestor and current lineage writes in read-visibility order."""
|
||||||
|
for lineage in _lineage_chain(run, scope_id=scope_id, lineage_id=lineage_id):
|
||||||
|
yield from lineage.writes
|
||||||
|
|
||||||
|
|
||||||
def _lineage(run: RunState, *, scope_id: str, lineage_id: str) -> LineageState:
|
def _lineage(run: RunState, *, scope_id: str, lineage_id: str) -> LineageState:
|
||||||
lineage = run.lineages.get(lineage_id)
|
lineage = run.lineages.get(lineage_id)
|
||||||
if lineage is None:
|
if lineage is None:
|
||||||
|
|||||||
@@ -3,7 +3,11 @@ from __future__ import annotations
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from wf_core.run_state import ExecutionFrame, RunState
|
from wf_core.run_state import ExecutionFrame, RunState
|
||||||
from wf_core.runtime.lineage import LineageStateView, lineage_writes_for_frame
|
from wf_core.runtime.lineage import (
|
||||||
|
LineageStateView,
|
||||||
|
lineage_writes_for_frame,
|
||||||
|
scope_state_for_frame,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def state_view_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
|
def state_view_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
|
||||||
@@ -13,8 +17,9 @@ def state_view_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]
|
|||||||
foreach barrier commits. Later nodes in the same item must read those
|
foreach barrier commits. Later nodes in the same item must read those
|
||||||
earlier writes, while sibling item frames must not see them.
|
earlier writes, while sibling item frames must not see them.
|
||||||
"""
|
"""
|
||||||
|
scope_state = scope_state_for_frame(run, frame)
|
||||||
writes = lineage_writes_for_frame(run, frame)
|
writes = lineage_writes_for_frame(run, frame)
|
||||||
if not writes:
|
if not writes:
|
||||||
return run.state
|
return scope_state
|
||||||
|
|
||||||
return LineageStateView(run.state, writes).to_state_dict()
|
return LineageStateView(scope_state, writes).to_state_dict()
|
||||||
|
|||||||
@@ -13,12 +13,15 @@ from wf_core import (
|
|||||||
from wf_core.models.reducers import ReducerRef
|
from wf_core.models.reducers import ReducerRef
|
||||||
from wf_core.paths import StatePath
|
from wf_core.paths import StatePath
|
||||||
from wf_core.run_state import StateWrite
|
from wf_core.run_state import StateWrite
|
||||||
|
from wf_core.run_state import ExecutionFrame, LineageState, RuntimeScope
|
||||||
from wf_core.runtime.lineage import (
|
from wf_core.runtime.lineage import (
|
||||||
add_lineage,
|
add_lineage,
|
||||||
append_lineage_writes,
|
append_lineage_writes,
|
||||||
lineage_patch,
|
lineage_patch,
|
||||||
lineage_state_view,
|
lineage_state_view,
|
||||||
|
scope_state_for_frame,
|
||||||
)
|
)
|
||||||
|
from wf_core.runtime.ops.overlays import state_view_for_frame
|
||||||
from wf_core.runtime.ops.runs import create_run_state
|
from wf_core.runtime.ops.runs import create_run_state
|
||||||
|
|
||||||
|
|
||||||
@@ -89,6 +92,71 @@ def test_lineage_helpers_store_ordered_writes_and_preserve_replay_values() -> No
|
|||||||
assert patch.writes[0].visible_value == "visible"
|
assert patch.writes[0].visible_value == "visible"
|
||||||
|
|
||||||
|
|
||||||
|
def test_state_view_for_frame_reads_from_frame_scope_state() -> None:
|
||||||
|
run = create_run_state(_minimal_workflow(), {"value": "root"})
|
||||||
|
run.scopes["child"] = RuntimeScope(
|
||||||
|
id="child",
|
||||||
|
workflow_name="child_workflow",
|
||||||
|
committed_state={"value": "child"},
|
||||||
|
)
|
||||||
|
run.lineages["child/root"] = LineageState(id="child/root", scope_id="child")
|
||||||
|
frame = ExecutionFrame(
|
||||||
|
id="child-frame",
|
||||||
|
kind="workflow",
|
||||||
|
node_id="finish",
|
||||||
|
scope_id="child",
|
||||||
|
lineage_id="child/root",
|
||||||
|
)
|
||||||
|
|
||||||
|
state_view = state_view_for_frame(run, frame)
|
||||||
|
|
||||||
|
assert scope_state_for_frame(run, frame)["value"] == "child"
|
||||||
|
assert state_view["value"] == "child"
|
||||||
|
assert run.state["value"] == "root"
|
||||||
|
|
||||||
|
|
||||||
|
def test_state_view_for_frame_overlays_writes_onto_frame_scope_state() -> None:
|
||||||
|
run = create_run_state(_minimal_workflow(), {"value": "root"})
|
||||||
|
run.scopes["child"] = RuntimeScope(
|
||||||
|
id="child",
|
||||||
|
workflow_name="child_workflow",
|
||||||
|
committed_state={"value": "child"},
|
||||||
|
)
|
||||||
|
run.lineages["child/root"] = LineageState(id="child/root", scope_id="child")
|
||||||
|
add_lineage(
|
||||||
|
run,
|
||||||
|
scope_id="child",
|
||||||
|
lineage_id="child/branch",
|
||||||
|
parent_id="child/root",
|
||||||
|
)
|
||||||
|
append_lineage_writes(
|
||||||
|
run,
|
||||||
|
scope_id="child",
|
||||||
|
lineage_id="child/branch",
|
||||||
|
writes=[
|
||||||
|
StateWrite(
|
||||||
|
path=StatePath(("value",)),
|
||||||
|
incoming_value="incoming",
|
||||||
|
visible_value="visible",
|
||||||
|
reducer=ReducerRef(name="wf.std.replace"),
|
||||||
|
)
|
||||||
|
],
|
||||||
|
)
|
||||||
|
frame = ExecutionFrame(
|
||||||
|
id="child-frame",
|
||||||
|
kind="workflow",
|
||||||
|
node_id="finish",
|
||||||
|
scope_id="child",
|
||||||
|
lineage_id="child/branch",
|
||||||
|
)
|
||||||
|
|
||||||
|
state_view = state_view_for_frame(run, frame)
|
||||||
|
|
||||||
|
assert state_view["value"] == "visible"
|
||||||
|
assert run.scopes["child"].committed_state["value"] == "child"
|
||||||
|
assert run.state["value"] == "root"
|
||||||
|
|
||||||
|
|
||||||
def _minimal_workflow() -> Workflow:
|
def _minimal_workflow() -> Workflow:
|
||||||
return Workflow(
|
return Workflow(
|
||||||
name="lineage_root",
|
name="lineage_root",
|
||||||
|
|||||||
Reference in New Issue
Block a user