lineage helper api
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
Reference in New Issue
Block a user