lineage helper api

This commit is contained in:
lda
2026-05-25 00:00:15 +07:00 Verified
parent 12c70513ba
commit ea1f746512
2 changed files with 142 additions and 2 deletions
+84 -2
View File
@@ -1,12 +1,13 @@
from __future__ import annotations
from collections.abc import Mapping, Sequence
from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from typing import Any
from wf_core.run_state import ExecutionFrame, RunState, StateWrite
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.ops.state import StatePatch
from wf_core.runtime.ops.state import safe_set_nested_value
@@ -58,3 +59,84 @@ def lineage_writes_for_frame(
if pending is None:
return ()
return pending.patch.writes
def add_lineage(
run: RunState,
*,
scope_id: str,
lineage_id: str,
parent_id: str | None,
) -> None:
"""Create one lineage record inside an existing runtime scope."""
if scope_id not in run.scopes:
raise ValueError(f"unknown scope {scope_id!r}")
if lineage_id in run.lineages:
raise ValueError(f"duplicate lineage {lineage_id!r}")
if parent_id is not None and parent_id not in run.lineages:
raise ValueError(f"unknown parent lineage {parent_id!r}")
run.lineages[lineage_id] = LineageState(
id=lineage_id,
scope_id=scope_id,
parent_id=parent_id,
)
def append_lineage_writes(
run: RunState,
*,
scope_id: str,
lineage_id: str,
writes: Sequence[StateWrite],
) -> None:
"""Append ordered writes to an existing lineage without committing state."""
lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
lineage.writes.extend(writes)
def lineage_patch(
run: RunState,
*,
scope_id: str,
lineage_id: str,
) -> StatePatch:
"""Return a replayable patch for one lineage's pending writes."""
lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
return StatePatch(writes=list(lineage.writes))
def lineage_state_view(
run: RunState,
*,
scope_id: str,
lineage_id: str,
) -> dict[str, Any]:
"""Materialize scope committed state plus ancestor/current lineage writes."""
scope = run.scopes.get(scope_id)
if scope is None:
raise ValueError(f"unknown scope {scope_id!r}")
writes: list[StateWrite] = []
for lineage in _lineage_chain(run, scope_id=scope_id, lineage_id=lineage_id):
writes.extend(lineage.writes)
return LineageStateView(scope.committed_state, writes).to_state_dict()
def _lineage(run: RunState, *, scope_id: str, lineage_id: str) -> LineageState:
lineage = run.lineages.get(lineage_id)
if lineage is None:
raise ValueError(f"unknown lineage {lineage_id!r}")
if lineage.scope_id != scope_id:
raise ValueError(
f"lineage {lineage_id!r} belongs to scope {lineage.scope_id!r}, "
f"not {scope_id!r}"
)
return lineage
def _lineage_chain(
run: RunState, *, scope_id: str, lineage_id: str
) -> Iterator[LineageState]:
lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
if lineage.parent_id is not None:
yield from _lineage_chain(run, scope_id=scope_id, lineage_id=lineage.parent_id)
yield lineage
+58
View File
@@ -10,6 +10,15 @@ from wf_core import (
StateSchema,
Workflow,
)
from wf_core.models.reducers import ReducerRef
from wf_core.paths import StatePath
from wf_core.run_state import StateWrite
from wf_core.runtime.lineage import (
add_lineage,
append_lineage_writes,
lineage_patch,
lineage_state_view,
)
from wf_core.runtime.ops.runs import create_run_state
@@ -31,6 +40,55 @@ def test_create_run_state_initializes_root_scope_and_lineage() -> None:
assert run.frames["root"].parent_lineage_id is None
def test_add_lineage_rejects_duplicate_or_unknown_scope() -> None:
run = create_run_state(_minimal_workflow(), {"value": "seed"})
add_lineage(run, scope_id="root", lineage_id="child", parent_id="root")
assert run.lineages["child"].scope_id == "root"
assert run.lineages["child"].parent_id == "root"
try:
add_lineage(run, scope_id="root", lineage_id="child", parent_id="root")
except ValueError as exc:
assert "duplicate lineage" in str(exc)
else:
raise AssertionError("expected duplicate lineage error")
try:
add_lineage(run, scope_id="missing", lineage_id="other", parent_id="root")
except ValueError as exc:
assert "unknown scope" in str(exc)
else:
raise AssertionError("expected unknown scope error")
def test_lineage_helpers_store_ordered_writes_and_preserve_replay_values() -> None:
run = create_run_state(_minimal_workflow(), {"value": "seed"})
add_lineage(run, scope_id="root", lineage_id="child", parent_id="root")
writes = [
StateWrite(
path=StatePath(("value",)),
incoming_value="incoming",
visible_value="visible",
reducer=ReducerRef(name="wf.std.replace"),
)
]
append_lineage_writes(run, scope_id="root", lineage_id="child", writes=writes)
assert run.lineages["child"].writes[0].incoming_value == "incoming"
assert run.lineages["child"].writes[0].visible_value == "visible"
assert (
lineage_state_view(run, scope_id="root", lineage_id="child")["value"]
== "visible"
)
assert run.state["value"] == "seed"
patch = lineage_patch(run, scope_id="root", lineage_id="child")
assert patch.writes[0].incoming_value == "incoming"
assert patch.writes[0].visible_value == "visible"
def _minimal_workflow() -> Workflow:
return Workflow(
name="lineage_root",