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 __future__ import annotations
from collections.abc import Mapping, Sequence from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy from copy import deepcopy
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any 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.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 from wf_core.runtime.ops.state import safe_set_nested_value
@@ -58,3 +59,84 @@ def lineage_writes_for_frame(
if pending is None: if pending is None:
return () return ()
return pending.patch.writes 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, StateSchema,
Workflow, 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 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 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: def _minimal_workflow() -> Workflow:
return Workflow( return Workflow(
name="lineage_root", name="lineage_root",