mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-27 05:30:39 +00:00
81f4d7707f
* fix(runtime): preserve canonical delta-history order in cache CachedHistorySaver composed a parent checkpoint's channel writes from CheckpointTuple.pending_writes in storage order, but the saver's authoritative get_delta_channel_history orders writes by (task_id, idx). On the in-memory backend pending_writes follows task-arrival order, so a fan-out super-step (several parallel tasks writing one channel in a single checkpoint) was materialized in a different order with the cache enabled than with the bare saver, breaking the cache's transparency invariant. Sort a checkpoint's channel writes by task_id in _channel_writes, restoring the canonical (task_id, idx) order while preserving each task's internal idx sequence. The sort is a no-op on sqlite and postgres, whose pending_writes already arrive ORDER BY task_id, idx. Add async and sync regression tests over a real InMemorySaver seeded with a fan-out checkpoint whose writes arrive out of task_id order; both diverge from the inner saver's native walk without the fix and match with it. Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> * docs(runtime): clarify cached saver ordering contract --------- Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
350 lines
16 KiB
Python
350 lines
16 KiB
Python
"""Read-through delta-history cache wrapper for a canonical BaseCheckpointSaver.
|
|
|
|
The inner saver must return delta history in ascending ``(task_id, idx)``
|
|
order per checkpoint, with each task's ``pending_writes`` in ``idx`` order.
|
|
The in-repo memory, sqlite, and postgres backends satisfy this contract;
|
|
the base saver's default history walk alone does not guarantee it.
|
|
|
|
Correctness argument (spec §3): a checkpoint's delta history is a pure
|
|
function of its sealed ancestor chain — the LangGraph contract excludes the
|
|
target's own pending writes, parent links are fixed at creation, and an
|
|
ancestor's writes are sealed once its child exists. Entries keyed by
|
|
(thread, ns, checkpoint_id, channel) are therefore immutable: no
|
|
invalidation, and shared backends are coherent across processes.
|
|
|
|
The wrapper never caches the "latest checkpoint" resolution; only histories
|
|
keyed by resolved immutable checkpoint_ids.
|
|
|
|
Data lifecycle: thread deletion and prune purge the thread's cached entries
|
|
(source-of-truth removal must not leave residual history payloads in the
|
|
cache); run-scoped deletes cannot be mapped to threads cheaply and rely on
|
|
LRU/TTL bounds.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections.abc import Iterator, Sequence
|
|
from typing import Any
|
|
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple, PendingWrite
|
|
|
|
from deerflow.runtime.checkpoint_cache.base import make_history_key
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Depth budget for recursive compose before falling back to a chain-warming
|
|
# walk. Steady-state runs need ~2 (one intermediate checkpoint per step);
|
|
# deeper cold chains are handled faster by one warming walk than by many
|
|
# recursive single-tuple fetches.
|
|
_COMPOSE_MAX_DEPTH = 8
|
|
|
|
|
|
def _checkpoint_ref(tup: CheckpointTuple) -> tuple[str, str, str]:
|
|
configurable = tup.config["configurable"]
|
|
return (
|
|
str(configurable["thread_id"]),
|
|
str(configurable.get("checkpoint_ns", "")),
|
|
str(configurable["checkpoint_id"]),
|
|
)
|
|
|
|
|
|
def _channel_writes(tup: CheckpointTuple, channel: str) -> list[PendingWrite]:
|
|
"""Writes for one channel in the inner saver's canonical delta order.
|
|
|
|
The wrapper assumes the inner saver's ``get_delta_channel_history`` uses
|
|
ascending ``(task_id, idx)`` order per checkpoint; the base implementation
|
|
instead iterates ``pending_writes`` in storage order. For supported savers,
|
|
``CheckpointTuple.pending_writes`` carries each task's writes in ``idx``
|
|
order, but across tasks it follows
|
|
*storage/completion* order, which the InMemorySaver backend leaves in task
|
|
arrival order rather than ``task_id`` order. Composing from that raw order
|
|
makes the cached history disagree with the inner saver on fan-out
|
|
supersteps (several parallel tasks writing one channel). A stable sort by
|
|
``task_id`` restores the canonical ``(task_id, idx)`` order while preserving
|
|
each task's internal ``idx`` sequence; it is a no-op on backends that
|
|
already return writes ``ORDER BY task_id, idx`` (sqlite, postgres).
|
|
"""
|
|
writes = [w for w in (tup.pending_writes or []) if w[1] == channel]
|
|
writes.sort(key=lambda w: w[0])
|
|
return writes
|
|
|
|
|
|
class CachedHistorySaver(BaseCheckpointSaver):
|
|
def __init__(self, inner: BaseCheckpointSaver, cache: Any, *, key_prefix: str) -> None:
|
|
# Instance attr shadows the base class JsonPlusSerializer default.
|
|
self.serde = inner.serde
|
|
self._inner = inner
|
|
self._cache = cache
|
|
self._key_prefix = key_prefix
|
|
self._compose_hits = 0
|
|
self._full_walks = 0
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
# Safety net for saver-specific extras (e.g. AsyncSqliteSaver.setup).
|
|
# Base-class methods are explicitly delegated below, so this only
|
|
# fires for attributes BaseCheckpointSaver does not define.
|
|
inner = self.__dict__.get("_inner")
|
|
if inner is None:
|
|
raise AttributeError(name)
|
|
return getattr(inner, name)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Key building
|
|
# ------------------------------------------------------------------
|
|
|
|
def _key(self, tup: CheckpointTuple, channel: str) -> str:
|
|
thread_id, ns, checkpoint_id = _checkpoint_ref(tup)
|
|
return make_history_key(self._key_prefix, thread_id, ns, checkpoint_id, channel)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Stats
|
|
# ------------------------------------------------------------------
|
|
|
|
def stats(self) -> dict[str, int]:
|
|
backend = self._cache.stats().as_dict()
|
|
return {**backend, "compose_hits": self._compose_hits, "full_walks": self._full_walks}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Delta history: the only overridden behavior
|
|
# ------------------------------------------------------------------
|
|
|
|
async def aget_delta_channel_history(self, *, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
|
if not channels:
|
|
return {}
|
|
if not getattr(self._cache, "enabled", True):
|
|
# Disabled cache: pass straight through; composing over all-miss
|
|
# entries is strictly more work than the raw saver's walk.
|
|
return await self._walk_inner(config, channels)
|
|
target = await self._inner.aget_tuple(config)
|
|
if target is None:
|
|
return await self._walk_inner(config, channels)
|
|
|
|
keys = {ch: self._key(target, ch) for ch in channels}
|
|
hits = await self._cache.aget_many(list(keys.values()))
|
|
found: dict[str, dict[str, Any]] = {}
|
|
missing: list[str] = []
|
|
for ch in channels:
|
|
entry = hits.get(keys[ch])
|
|
if entry is None:
|
|
missing.append(ch)
|
|
else:
|
|
found[ch] = entry
|
|
|
|
computed: dict[str, dict[str, Any]] = {}
|
|
if missing:
|
|
computed = await self._compose_or_walk(config, target, missing)
|
|
|
|
new_entries = {keys[ch]: computed[ch] for ch in missing if ch in computed}
|
|
if new_entries:
|
|
await self._cache.aset_many(new_entries)
|
|
|
|
return {ch: found.get(ch) or computed.get(ch) or {"writes": []} for ch in channels}
|
|
|
|
async def _compose_or_walk(self, config: RunnableConfig, target: CheckpointTuple, missing: list[str]) -> dict[str, dict[str, Any]]:
|
|
return {ch: await self._aresolve(target, ch, _COMPOSE_MAX_DEPTH) for ch in missing}
|
|
|
|
async def _aresolve(self, tup: CheckpointTuple, channel: str, depth: int) -> dict[str, Any]:
|
|
"""Recursively compose history(tup) from the nearest warm ancestor.
|
|
|
|
Real runs create several checkpoints per super-step and only some are
|
|
ever materialized as targets, so the parent is usually an unwarmed
|
|
intermediate checkpoint (measured: 0 cache hits with single-level
|
|
compose on a 500-step sqlite run). Recursing one level per
|
|
intermediate lands on a warmed ancestor within ~2 levels in steady
|
|
state; each composed level is cached, so the warm frontier follows
|
|
the run. At depth 0 on a cold chain it delegates one inner fast-path
|
|
walk (2 SQL) rather than crawling ancestors tuple-by-tuple.
|
|
"""
|
|
parent_config = tup.parent_config
|
|
if parent_config is None:
|
|
return {"writes": []}
|
|
parent = await self._inner.aget_tuple(parent_config)
|
|
if parent is None:
|
|
return {"writes": []}
|
|
|
|
channel_values = parent.checkpoint.get("channel_values") or {}
|
|
writes = _channel_writes(parent, channel)
|
|
if channel in channel_values:
|
|
self._compose_hits += 1
|
|
return {"writes": writes, "seed": channel_values[channel]}
|
|
|
|
key = self._key(parent, channel)
|
|
hits = await self._cache.aget_many([key])
|
|
parent_history = hits.get(key)
|
|
if parent_history is None:
|
|
if depth > 0:
|
|
parent_history = await self._aresolve(parent, channel, depth - 1)
|
|
else:
|
|
# Depth budget exhausted on a cold chain: delegate ONE inner
|
|
# fast-path walk (2 SQL total) for this level instead of
|
|
# fetching every ancestor tuple individually. Ancestors below
|
|
# stay cold; resolving them later recurses up to the nearest
|
|
# warm level, so the frontier still follows the run.
|
|
self._full_walks += 1
|
|
walked = await self._inner.aget_delta_channel_history(config=parent.config, channels=[channel])
|
|
parent_history = walked.get(channel) or {"writes": []}
|
|
if parent_history is not None:
|
|
await self._cache.aset_many({key: parent_history})
|
|
|
|
self._compose_hits += 1
|
|
entry: dict[str, Any] = {"writes": list(parent_history["writes"]) + writes}
|
|
if "seed" in parent_history:
|
|
entry["seed"] = parent_history["seed"]
|
|
return entry
|
|
|
|
async def _walk_inner(self, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
|
self._full_walks += 1
|
|
return dict(await self._inner.aget_delta_channel_history(config=config, channels=channels))
|
|
|
|
def get_delta_channel_history(self, *, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
|
if not channels:
|
|
return {}
|
|
if not getattr(self._cache, "enabled", True):
|
|
return self._walk_inner_sync(config, channels)
|
|
get_many = getattr(self._cache, "get_many", None)
|
|
set_many = getattr(self._cache, "set_many", None)
|
|
if get_many is None or set_many is None:
|
|
raise TypeError("sync get_delta_channel_history requires a SyncCheckpointHistoryCache (memory backend)")
|
|
target = self._inner.get_tuple(config)
|
|
if target is None:
|
|
return self._walk_inner_sync(config, channels)
|
|
|
|
keys = {ch: self._key(target, ch) for ch in channels}
|
|
hits = get_many(list(keys.values()))
|
|
found: dict[str, dict[str, Any]] = {}
|
|
missing: list[str] = []
|
|
for ch in channels:
|
|
entry = hits.get(keys[ch])
|
|
if entry is None:
|
|
missing.append(ch)
|
|
else:
|
|
found[ch] = entry
|
|
|
|
computed: dict[str, dict[str, Any]] = {}
|
|
if missing:
|
|
computed = {ch: self._resolve_sync(target, ch, _COMPOSE_MAX_DEPTH) for ch in missing}
|
|
|
|
new_entries = {keys[ch]: computed[ch] for ch in missing if ch in computed}
|
|
if new_entries:
|
|
set_many(new_entries)
|
|
return {ch: found.get(ch) or computed.get(ch) or {"writes": []} for ch in channels}
|
|
|
|
def _resolve_sync(self, tup: CheckpointTuple, channel: str, depth: int) -> dict[str, Any]:
|
|
"""Sync twin of _aresolve (recursive compose, see its docstring)."""
|
|
parent_config = tup.parent_config
|
|
if parent_config is None:
|
|
return {"writes": []}
|
|
parent = self._inner.get_tuple(parent_config)
|
|
if parent is None:
|
|
return {"writes": []}
|
|
|
|
channel_values = parent.checkpoint.get("channel_values") or {}
|
|
writes = _channel_writes(parent, channel)
|
|
if channel in channel_values:
|
|
self._compose_hits += 1
|
|
return {"writes": writes, "seed": channel_values[channel]}
|
|
|
|
key = self._key(parent, channel)
|
|
parent_history = self._cache.get_many([key]).get(key)
|
|
if parent_history is None:
|
|
if depth > 0:
|
|
parent_history = self._resolve_sync(parent, channel, depth - 1)
|
|
else:
|
|
# See _aresolve: one inner fast-path walk, no per-tuple crawl.
|
|
self._full_walks += 1
|
|
parent_history = self._inner.get_delta_channel_history(config=parent.config, channels=[channel]).get(channel) or {"writes": []}
|
|
if parent_history is not None:
|
|
self._cache.set_many({key: parent_history})
|
|
|
|
self._compose_hits += 1
|
|
entry: dict[str, Any] = {"writes": list(parent_history["writes"]) + writes}
|
|
if "seed" in parent_history:
|
|
entry["seed"] = parent_history["seed"]
|
|
return entry
|
|
|
|
def _walk_inner_sync(self, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
|
self._full_walks += 1
|
|
return dict(self._inner.get_delta_channel_history(config=config, channels=channels))
|
|
|
|
# ------------------------------------------------------------------
|
|
# Explicit delegation (BaseCheckpointSaver defines these concretely,
|
|
# so __getattr__ never fires for them)
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
|
return self._inner.get_tuple(config)
|
|
|
|
def list(self, config: RunnableConfig | None, *, filter: dict[str, Any] | None = None, before: RunnableConfig | None = None, limit: int | None = None) -> Iterator[CheckpointTuple]:
|
|
return self._inner.list(config, filter=filter, before=before, limit=limit)
|
|
|
|
def put(self, config: RunnableConfig, checkpoint: dict[str, Any], metadata: dict[str, Any], new_versions: dict[str, Any]) -> RunnableConfig:
|
|
return self._inner.put(config, checkpoint, metadata, new_versions)
|
|
|
|
def put_writes(self, config: RunnableConfig, writes: Sequence[tuple[str, str, Any]], task_id: str, task_path: str = "") -> None:
|
|
self._inner.put_writes(config, writes, task_id, task_path)
|
|
|
|
def delete_thread(self, thread_id: str) -> None:
|
|
self._inner.delete_thread(thread_id)
|
|
self._purge_thread_sync(thread_id)
|
|
|
|
def delete_for_runs(self, run_ids: Sequence[str]) -> None:
|
|
# Run-scoped deletes cannot be mapped back to threads without an extra
|
|
# query, so cached entries are left in place: they stay *correct* (the
|
|
# sealed-chain argument is unaffected by other chains) and residual
|
|
# retention is bounded by LRU/TTL. No in-tree callers today.
|
|
self._inner.delete_for_runs(run_ids)
|
|
|
|
def _purge_thread_sync(self, thread_id: str) -> None:
|
|
delete = getattr(self._cache, "delete_thread", None)
|
|
if delete is not None:
|
|
delete(self._key_prefix, thread_id)
|
|
|
|
def copy_thread(self, source_thread_id: str, target_thread_id: str) -> None:
|
|
self._inner.copy_thread(source_thread_id, target_thread_id)
|
|
|
|
def prune(self, thread_ids: Sequence[str], *, strategy: str = "keep_latest") -> None:
|
|
self._inner.prune(thread_ids, strategy=strategy)
|
|
# Pruning rewrites these threads' chains: purge so no cached history
|
|
# references a deleted ancestor (retention) or its pre-prune chain.
|
|
for thread_id in thread_ids:
|
|
self._purge_thread_sync(thread_id)
|
|
|
|
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
|
return await self._inner.aget_tuple(config)
|
|
|
|
def alist(self, config: RunnableConfig | None, *, filter: dict[str, Any] | None = None, before: RunnableConfig | None = None, limit: int | None = None) -> Any:
|
|
return self._inner.alist(config, filter=filter, before=before, limit=limit)
|
|
|
|
async def aput(self, config: RunnableConfig, checkpoint: dict[str, Any], metadata: dict[str, Any], new_versions: dict[str, Any]) -> RunnableConfig:
|
|
return await self._inner.aput(config, checkpoint, metadata, new_versions)
|
|
|
|
async def aput_writes(self, config: RunnableConfig, writes: Sequence[tuple[str, str, Any]], task_id: str, task_path: str = "") -> None:
|
|
await self._inner.aput_writes(config, writes, task_id, task_path)
|
|
|
|
async def adelete_thread(self, thread_id: str) -> None:
|
|
await self._inner.adelete_thread(thread_id)
|
|
await self._apurge_thread(thread_id)
|
|
|
|
async def adelete_for_runs(self, run_ids: Sequence[str]) -> None:
|
|
# See delete_for_runs: unscoped residual retention bounded by LRU/TTL.
|
|
await self._inner.adelete_for_runs(run_ids)
|
|
|
|
async def _apurge_thread(self, thread_id: str) -> None:
|
|
delete = getattr(self._cache, "adelete_thread", None)
|
|
if delete is not None:
|
|
await delete(self._key_prefix, thread_id)
|
|
|
|
async def acopy_thread(self, source_thread_id: str, target_thread_id: str) -> None:
|
|
await self._inner.acopy_thread(source_thread_id, target_thread_id)
|
|
|
|
async def aprune(self, thread_ids: Sequence[str], *, strategy: str = "keep_latest") -> None:
|
|
await self._inner.aprune(thread_ids, strategy=strategy)
|
|
# See prune: rewritten chains must not keep pre-prune cached histories.
|
|
for thread_id in thread_ids:
|
|
await self._apurge_thread(thread_id)
|
|
|
|
def get_next_version(self, current: Any, channel: Any) -> Any:
|
|
return self._inner.get_next_version(current, channel)
|