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
|
||||
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)
|
||||
if owner is None:
|
||||
return ()
|
||||
@@ -61,6 +69,14 @@ def lineage_writes_for_frame(
|
||||
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(
|
||||
run: RunState,
|
||||
*,
|
||||
@@ -121,6 +137,17 @@ def lineage_state_view(
|
||||
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:
|
||||
lineage = run.lineages.get(lineage_id)
|
||||
if lineage is None:
|
||||
|
||||
@@ -3,7 +3,11 @@ from __future__ import annotations
|
||||
from typing import Any
|
||||
|
||||
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]:
|
||||
@@ -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
|
||||
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)
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user