mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-28 14:10:37 +00:00
5cef1a11ef
* fix(checkpoints): wire conservative retention after runs * fix(checkpoints): fence retention before checkpoint reads * fix(checkpoints): address retention review feedback * fix(checkpoints): retry delta reads after retention race * fix(checkpoints): bound admission retention to the root tail --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
669 lines
32 KiB
Python
669 lines
32 KiB
Python
"""Thread-level checkpoint retention enforcing the #4189 item 3 contract.
|
|
|
|
Implements exactly the two deletion shapes proven safe by
|
|
``docs/checkpoint-retention-contract.md`` and its executable suite
|
|
(``tests/test_checkpoint_retention_contract.py``) — nothing else:
|
|
|
|
- **Trailing duration-only leaves** — ``persist_run_durations`` appends
|
|
metadata-only checkpoints after a run finishes; while no later run has
|
|
forked from one, it is nobody's ancestor and can be dropped (contract
|
|
scenario E1).
|
|
- **Leaf sibling branches** *(opt-in)* — a checkpoint forked off an older
|
|
turn that has no children (contract scenario E2). Opt-in because a
|
|
superseded line's checkpoints may still be explicit resume targets a
|
|
client holds (protected set item 1); ``RetentionPolicy.protect_checkpoint_ids``
|
|
is the escape hatch until a TTL semantic is agreed for that item.
|
|
|
|
Everything on the resume head's ancestor chain — including duration-only
|
|
chain links, which would need grafting before deletion — every explicitly
|
|
protected id, and any checkpoint that still owns ``writes`` rows (protected
|
|
set item 3) is never deleted. Deletion runs at the storage layer, mirroring
|
|
the contract's per-backend data model, and removes the writes rows orphaned
|
|
by a deleted checkpoint in the same step (contract "deletion mechanics").
|
|
Postgres ``checkpoint_blobs`` rows are keyed by ``version`` = the id of the
|
|
checkpoint that wrote the blob, so the same join keys clean them.
|
|
|
|
The Gateway invokes the conservative production policy after the next run has
|
|
won durable same-thread admission but before that worker reads a checkpoint.
|
|
It removes at most the trailing duration/history-cache leaf; the sibling-branch
|
|
shape remains opt-in. The admission path reads only the root head and parent,
|
|
checks that head for writes, and leaves shared blobs in place. The exhaustive
|
|
maintenance policy alone collects before/after per-thread stats in the same
|
|
normalized shape as
|
|
``scripts/benchmark/checkpoint/bench_channels.py``.
|
|
|
|
History fast-path interaction: the trailing duration-only leaf is also the
|
|
carrier of the run-history metadata cache (``run_durations`` /
|
|
``run_message_ids``). Retaining it until the next admission lets repeated
|
|
history reads keep their fast path; pruning immediately before checkpoint
|
|
selection prevents the next run from making a regenerated cache leaf a
|
|
permanent ancestor.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections import defaultdict
|
|
from contextlib import AbstractAsyncContextManager, nullcontext
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
|
|
from app.gateway.checkpoint_lineage import (
|
|
checkpoint_configurable,
|
|
is_duration_only_checkpoint,
|
|
)
|
|
|
|
__all__ = [
|
|
"RetentionPolicy",
|
|
"RetentionReport",
|
|
"enforce_completed_run_retention",
|
|
"enforce_thread_retention",
|
|
"enforce_trailing_duration_retention",
|
|
]
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetentionPolicy:
|
|
"""What this run may delete. Defaults prune only what the contract proves safe unconditionally."""
|
|
|
|
prune_trailing_duration_leaves: bool = True
|
|
prune_leaf_sibling_branches: bool = False
|
|
protect_checkpoint_ids: frozenset[str] = frozenset()
|
|
strict_pending_write_guard: bool = True
|
|
max_delete_per_run: int | None = None
|
|
|
|
|
|
@dataclass
|
|
class RetentionReport:
|
|
"""Outcome of one retention pass over one thread."""
|
|
|
|
thread_id: str
|
|
# Resume heads that survived this pass, per namespace (checkpoint_ns -> id).
|
|
# The root namespace (the "" key) is what an unsaved ``aget_tuple`` resolves
|
|
# as the thread's latest state; a persistent subgraph contributes its own
|
|
# child namespace whose head is protected too and would be invisible in a
|
|
# singular field. Consumers that only care about the thread's latest state
|
|
# read the root key.
|
|
protected_head_ids: dict[str, str] = field(default_factory=dict)
|
|
deleted_checkpoint_ids: list[str] = field(default_factory=list)
|
|
stats_before: dict[str, int] = field(default_factory=dict)
|
|
stats_after: dict[str, int] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class _Node:
|
|
ns: str
|
|
cp_id: str
|
|
parent_ns: str | None
|
|
parent_id: str | None
|
|
duration_only: bool
|
|
metadata_source: str | None = None
|
|
carries_run_durations: bool = False
|
|
versions: frozenset = frozenset()
|
|
|
|
|
|
def _thread_config(thread_id: str) -> dict[str, Any]:
|
|
return {"configurable": {"thread_id": thread_id}}
|
|
|
|
|
|
def _row_field(row: Any, name: str, index: int) -> Any:
|
|
"""Read one column from a DB-API row whatever row factory produced it.
|
|
|
|
Postgres cursors are opened with ``row_factory=dict_row`` — both by this
|
|
codebase (``deerflow/runtime/checkpointer/async_provider.py``) and inside
|
|
``langgraph-checkpoint-postgres`` itself (``aio.py`` opens every cursor as
|
|
``conn.cursor(binary=True, row_factory=dict_row)``) — so their rows are
|
|
name-addressable and positional access raises ``KeyError: 0``. aiosqlite
|
|
rows are plain tuples. Reading through this helper keeps the per-backend
|
|
helpers independent of which saver opened the cursor: the Postgres leg of
|
|
``tests/test_checkpoint_retention_service.py`` failed on exactly that
|
|
assumption, on six scenarios, the first time CI ran it.
|
|
"""
|
|
try:
|
|
return row[name]
|
|
except (TypeError, KeyError, IndexError):
|
|
return row[index]
|
|
|
|
|
|
def _mark_duration_leaves_without_the_marker(nodes: dict[tuple[str, str], _Node]) -> None:
|
|
"""Classify metadata-only duration leaves on backends that drop the marker.
|
|
|
|
``persist_run_history_metadata`` is the only writer of these checkpoints and
|
|
stamps ``metadata["writes"]["runtime_run_duration"]``, which is what
|
|
:func:`is_duration_only_checkpoint` reads. The Postgres saver funnels
|
|
metadata through langgraph's ``get_serializable_checkpoint_metadata``, which
|
|
pops ``writes`` before the row is written (``checkpoint/base/__init__.py``;
|
|
only the Postgres savers call it), so on Postgres that marker never comes
|
|
back — every duration leaf would look resumable and the default E1 shape
|
|
would silently never fire on the production backend.
|
|
|
|
The writer's other stamps survive a Postgres round trip and together are
|
|
exact: ``source == "update"``, a non-empty accumulated ``run_durations``
|
|
map, and the leaf being a verbatim copy of its parent (``channel_versions``
|
|
copied unchanged — the writer only replaces ``id``/``ts``). A client
|
|
``update_state`` writes a newly versioned channel, so it fails the last
|
|
condition and is never swept into this class.
|
|
"""
|
|
for node in nodes.values():
|
|
if node.duration_only or node.parent_id is None:
|
|
continue
|
|
parent = nodes.get((node.parent_ns, node.parent_id))
|
|
if parent is None or node.metadata_source != "update" or not node.carries_run_durations:
|
|
continue
|
|
if node.versions and node.versions == parent.versions:
|
|
node.duration_only = True
|
|
|
|
|
|
def _storage_saver(saver: BaseCheckpointSaver) -> BaseCheckpointSaver:
|
|
"""Resolve a wrapper through its public source-saver contract."""
|
|
storage = getattr(saver, "source_saver", saver)
|
|
return storage if isinstance(storage, BaseCheckpointSaver) else saver
|
|
|
|
|
|
def _ensure_supported_saver(saver: BaseCheckpointSaver) -> BaseCheckpointSaver:
|
|
"""Fail fast before any row is read or written on an untested saver.
|
|
|
|
The per-backend helpers below model exactly three storage layouts
|
|
(memory, SQLite, Postgres) and would otherwise fall through to "assume
|
|
Postgres" SQL for any other ``BaseCheckpointSaver``. A shallow Postgres
|
|
saver (no ``checkpoint_blobs``/``checkpoint_writes`` tables) or a
|
|
third-party saver would then issue DELETEs and die partway — with
|
|
autocommit on the Postgres connection, after the ``checkpoints`` row is
|
|
already gone. For a destructive tool, an explicit allowlist that raises
|
|
``NotImplementedError`` beats a partial deletion and a confusing
|
|
traceback.
|
|
"""
|
|
storage = _storage_saver(saver)
|
|
if isinstance(storage, (InMemorySaver, AsyncSqliteSaver)):
|
|
return storage
|
|
if callable(getattr(storage, "_cursor", None)):
|
|
return storage
|
|
raise NotImplementedError(f"checkpoint retention requires a supported storage saver, got {type(storage).__module__}.{type(storage).__name__}")
|
|
|
|
|
|
async def _thread_storage_stats(saver: Any, thread_id: str) -> dict[str, int]:
|
|
"""Per-thread rows/bytes, same normalized shape as ``bench_channels``."""
|
|
if isinstance(saver, InMemorySaver):
|
|
checkpoint_rows = checkpoint_bytes = blob_rows = blob_bytes = write_rows = write_bytes = 0
|
|
for namespace in saver.storage.get(thread_id, {}).values():
|
|
for checkpoint, metadata, _parent in namespace.values():
|
|
checkpoint_rows += 1
|
|
checkpoint_bytes += len(checkpoint[1]) + len(metadata[1])
|
|
for (stored_thread, _ns, _channel, _version), (_type_tag, blob) in saver.blobs.items():
|
|
if stored_thread != thread_id:
|
|
continue
|
|
blob_rows += 1
|
|
blob_bytes += len(blob)
|
|
for (stored_thread, _ns, _cp_id), writes in saver.writes.items():
|
|
if stored_thread != thread_id:
|
|
continue
|
|
for _task_id, _channel, (_type_tag, blob), _path in writes.values():
|
|
write_rows += 1
|
|
write_bytes += len(blob)
|
|
return {
|
|
"logical_checkpoint_bytes": checkpoint_bytes + blob_bytes,
|
|
"logical_write_bytes": write_bytes,
|
|
"checkpoint_rows": checkpoint_rows,
|
|
"checkpoint_bytes": checkpoint_bytes,
|
|
"blob_rows": blob_rows,
|
|
"blob_bytes": blob_bytes,
|
|
"write_rows": write_rows,
|
|
"write_bytes": write_bytes,
|
|
}
|
|
if isinstance(saver, AsyncSqliteSaver):
|
|
sqls = (
|
|
("checkpoint_rows", "checkpoint_bytes", "SELECT COUNT(*), COALESCE(SUM(LENGTH(checkpoint) + LENGTH(metadata)), 0) FROM checkpoints WHERE thread_id = ?"),
|
|
("write_rows", "write_bytes", "SELECT COUNT(*), COALESCE(SUM(LENGTH(value)), 0) FROM writes WHERE thread_id = ?"),
|
|
)
|
|
stats: dict[str, int] = {}
|
|
for row_key, bytes_key, sql in sqls:
|
|
async with saver.conn.execute(sql, (thread_id,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
stats[row_key] = int(row[0])
|
|
stats[bytes_key] = int(row[1] or 0)
|
|
stats["blob_rows"] = 0
|
|
stats["blob_bytes"] = 0
|
|
stats["logical_checkpoint_bytes"] = stats["checkpoint_bytes"]
|
|
stats["logical_write_bytes"] = stats["write_bytes"]
|
|
return stats
|
|
# Reachable only for the allowlisted Postgres saver (see
|
|
# ``_ensure_supported_saver``): its checkpoint rows are split into
|
|
# ``checkpoints`` + ``checkpoint_blobs``.
|
|
sqls = (
|
|
("checkpoint_rows", "checkpoint_bytes", "SELECT COUNT(*) AS rows, COALESCE(SUM(pg_column_size(checkpoint) + pg_column_size(metadata)), 0) AS bytes FROM checkpoints WHERE thread_id = %s"),
|
|
("blob_rows", "blob_bytes", "SELECT COUNT(*) AS rows, COALESCE(SUM(octet_length(blob)), 0) AS bytes FROM checkpoint_blobs WHERE thread_id = %s"),
|
|
("write_rows", "write_bytes", "SELECT COUNT(*) AS rows, COALESCE(SUM(octet_length(blob)), 0) AS bytes FROM checkpoint_writes WHERE thread_id = %s"),
|
|
)
|
|
stats = {}
|
|
for row_key, bytes_key, sql in sqls:
|
|
async with saver._cursor() as cursor:
|
|
await cursor.execute(sql, (thread_id,))
|
|
row = await cursor.fetchone()
|
|
stats[row_key] = int(_row_field(row, "rows", 0))
|
|
stats[bytes_key] = int(_row_field(row, "bytes", 1) or 0)
|
|
stats["logical_checkpoint_bytes"] = stats["checkpoint_bytes"] + stats["blob_bytes"]
|
|
stats["logical_write_bytes"] = stats["write_bytes"]
|
|
return stats
|
|
|
|
|
|
async def _checkpoint_has_writes(saver: Any, key: tuple[str, str], thread_id: str) -> bool:
|
|
"""Check one checkpoint's pending writes without scanning its thread."""
|
|
ns, cp_id = key
|
|
if isinstance(saver, InMemorySaver):
|
|
return bool(saver.writes.get((thread_id, ns, cp_id)))
|
|
if isinstance(saver, AsyncSqliteSaver):
|
|
async with saver.conn.execute(
|
|
"SELECT 1 FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? LIMIT 1",
|
|
(thread_id, ns, cp_id),
|
|
) as cursor:
|
|
return await cursor.fetchone() is not None
|
|
async with saver._cursor() as cursor:
|
|
await cursor.execute(
|
|
"SELECT 1 FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s LIMIT 1",
|
|
(thread_id, ns, cp_id),
|
|
)
|
|
return await cursor.fetchone() is not None
|
|
|
|
|
|
async def _checkpoint_ids_with_writes(saver: Any, thread_id: str) -> set[tuple[str, str]]:
|
|
"""(checkpoint_ns, checkpoint_id) pairs that still own writes rows.
|
|
|
|
Protected set item 3: pending/uncommitted writes are retained state, not
|
|
garbage, so v1 refuses to delete any checkpoint that still owns writes
|
|
rows. On the memory backend ``InMemorySaver.writes`` also holds an *empty*
|
|
dict for checkpoints whose task produced no writes — counting key presence
|
|
would spare those phantom entries and silently disable pruning on the
|
|
memory saver, so a pair qualifies only when its writes dict is non-empty,
|
|
matching how :func:`_thread_storage_stats` counts rows rather than keys.
|
|
Production checkpoints normally accumulate their own committed writes rows
|
|
too, so the conservative default still makes pruning a no-op on hot
|
|
threads; a policy that distinguishes in-flight from orphaned writes belongs
|
|
to the contract's next revision, not to a fast path here.
|
|
"""
|
|
if isinstance(saver, InMemorySaver):
|
|
return {(ns, cp_id) for (stored_thread, ns, cp_id), writes in saver.writes.items() if stored_thread == thread_id and writes}
|
|
if isinstance(saver, AsyncSqliteSaver):
|
|
async with saver.conn.execute(
|
|
"SELECT DISTINCT checkpoint_ns, checkpoint_id FROM writes WHERE thread_id = ?",
|
|
(thread_id,),
|
|
) as cursor:
|
|
rows = await cursor.fetchall()
|
|
# aiosqlite rows are plain tuples; the Postgres rows below are not.
|
|
return {(row[0] or "", row[1]) for row in rows}
|
|
async with saver._cursor() as cursor:
|
|
await cursor.execute(
|
|
"SELECT DISTINCT checkpoint_ns, checkpoint_id FROM checkpoint_writes WHERE thread_id = %s",
|
|
(thread_id,),
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return {(_row_field(row, "checkpoint_ns", 0) or "", _row_field(row, "checkpoint_id", 1)) for row in rows}
|
|
|
|
|
|
async def _delete_checkpoint_rows(saver: Any, thread_id: str, key: tuple[str, str]) -> None:
|
|
"""Remove one checkpoint row and the writes rows it owns, jointly.
|
|
|
|
Blob rows are handled by the survivor-reachability pass
|
|
(:func:`_delete_unreachable_blobs`), never per-checkpoint: versions are
|
|
shared between a checkpoint and its clones (see the contract doc).
|
|
"""
|
|
ns, cp_id = key
|
|
if isinstance(saver, InMemorySaver):
|
|
saver.storage.get(thread_id, {}).get(ns, {}).pop(cp_id, None)
|
|
saver.writes.pop((thread_id, ns, cp_id), None)
|
|
return
|
|
if isinstance(saver, AsyncSqliteSaver):
|
|
await saver.conn.execute(
|
|
"DELETE FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
|
|
(thread_id, ns, cp_id),
|
|
)
|
|
await saver.conn.execute(
|
|
"DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
|
|
(thread_id, ns, cp_id),
|
|
)
|
|
await saver.conn.commit()
|
|
return
|
|
# Reachable only for the allowlisted Postgres saver (see
|
|
# ``_ensure_supported_saver``).
|
|
async with saver._cursor() as cursor:
|
|
await cursor.execute(
|
|
"DELETE FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s",
|
|
(thread_id, ns, cp_id),
|
|
)
|
|
await cursor.execute(
|
|
"DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s",
|
|
(thread_id, ns, cp_id),
|
|
)
|
|
|
|
|
|
async def _delete_unreachable_blobs(saver: Any, thread_id: str, survivor_versions: set) -> None:
|
|
"""Whole-thread blob GC: drop exactly the versions no surviving checkpoint
|
|
references (memory ``saver.blobs`` / Postgres ``checkpoint_blobs``)."""
|
|
if isinstance(saver, InMemorySaver):
|
|
for key in [key for key in saver.blobs if key[0] == thread_id and key[3] not in survivor_versions]:
|
|
del saver.blobs[key]
|
|
return
|
|
if isinstance(saver, AsyncSqliteSaver):
|
|
return
|
|
async with saver._cursor() as cursor:
|
|
await cursor.execute("SELECT DISTINCT version FROM checkpoint_blobs WHERE thread_id = %s", (thread_id,))
|
|
rows = await cursor.fetchall()
|
|
orphans = [version for version in (_row_field(row, "version", 0) for row in rows) if version not in survivor_versions]
|
|
if orphans:
|
|
async with saver._cursor() as cursor:
|
|
await cursor.execute(
|
|
"DELETE FROM checkpoint_blobs WHERE thread_id = %s AND version = ANY(%s)",
|
|
(thread_id, orphans),
|
|
)
|
|
|
|
|
|
async def enforce_thread_retention(
|
|
saver: BaseCheckpointSaver,
|
|
thread_id: str,
|
|
policy: RetentionPolicy | None = None,
|
|
*,
|
|
thread_lock: AbstractAsyncContextManager[None] | None = None,
|
|
collect_stats: bool = True,
|
|
) -> RetentionReport:
|
|
"""Apply *policy* to one thread's checkpoints and return what happened.
|
|
|
|
Classification walks the parent chain the same way
|
|
``app/gateway/checkpoint_lineage.py`` does; each namespace's resume head is
|
|
that namespace's newest non-duration-only checkpoint by checkpoint id
|
|
(LangGraph ids are time-ordered), and every head's whole ancestor chain is
|
|
protected. Anything off those chains is only deletable when it is a leaf,
|
|
not explicitly protected, free of writes rows under the strict guard, and
|
|
matches one of the two contract-proven shapes. A node whose parent is
|
|
already missing is left alone: partial damage must not be silently
|
|
compounded. Only savers the deletion mechanics have been validated on are
|
|
accepted (see :func:`_ensure_supported_saver`).
|
|
|
|
``max_delete_per_run`` is validated before any row is read: a negative cap
|
|
raises ``ValueError`` rather than reaching the slice that applies it, where
|
|
Python's negative indexing would widen the batch instead of disabling it.
|
|
An empty thread returns a no-op report whose ``stats_after`` mirrors
|
|
``stats_before``, so the measurement shape does not depend on whether the
|
|
thread had anything to classify.
|
|
|
|
Concurrency requirement: classification and deletion are two separate
|
|
passes over the store, so a run that forks from a node classified as a
|
|
leaf in between leaves a dangling ``parent_config`` — the contract's own
|
|
"converts a cleanup into a thread-level outage" failure class. Callers
|
|
must therefore serialize per-thread mutation against the runtime writer
|
|
by passing the thread's checkpoint lock
|
|
(``deerflow.runtime.runs.worker._checkpoint_thread_lock(thread_id)``) as
|
|
*thread_lock*; without one, retention must only run while the thread is
|
|
guaranteed quiescent.
|
|
|
|
The lock returned by ``_checkpoint_thread_lock`` is a plain, *non-reentrant*
|
|
``asyncio.Lock`` (``AsyncKeyedLockTable.hold``). A caller that already
|
|
holds it and then passes it here self-deadlocks — which matters because
|
|
the runtime's own ``persist_run_history_metadata`` enters that same lock
|
|
before writing, so a post-run-hook call site must invoke retention *after*
|
|
releasing it, not from inside the held section.
|
|
"""
|
|
storage_saver = _ensure_supported_saver(saver)
|
|
effective = policy or RetentionPolicy()
|
|
if effective.max_delete_per_run is not None and effective.max_delete_per_run < 0:
|
|
# Fail closed *before* any store read. ``deletable[:cap]`` with a
|
|
# negative cap selects every candidate except the last few, so an
|
|
# invalid value meant to disable the run would instead maximize it —
|
|
# a batch bound must never widen a destructive pass.
|
|
raise ValueError(f"max_delete_per_run must be >= 0, got {effective.max_delete_per_run}")
|
|
report = RetentionReport(thread_id=thread_id)
|
|
lock: AbstractAsyncContextManager[None] = thread_lock if thread_lock is not None else nullcontext()
|
|
async with lock:
|
|
if collect_stats:
|
|
report.stats_before = await _thread_storage_stats(storage_saver, thread_id)
|
|
|
|
nodes: dict[tuple[str, str], _Node] = {}
|
|
async for tuple_ in storage_saver.alist(_thread_config(thread_id), limit=None):
|
|
configurable = checkpoint_configurable(tuple_)
|
|
cp_id = configurable.get("checkpoint_id")
|
|
if not cp_id:
|
|
continue
|
|
ns = configurable.get("checkpoint_ns") or ""
|
|
parent_ns: str | None = None
|
|
parent_id: str | None = None
|
|
parent_config = getattr(tuple_, "parent_config", None)
|
|
if isinstance(parent_config, dict):
|
|
parent = parent_config.get("configurable") or {}
|
|
parent_ns = parent.get("checkpoint_ns") or ""
|
|
parent_id = parent.get("checkpoint_id")
|
|
checkpoint = getattr(tuple_, "checkpoint", None) or {}
|
|
channel_versions = checkpoint.get("channel_versions")
|
|
metadata = getattr(tuple_, "metadata", None) or {}
|
|
metadata = metadata if isinstance(metadata, dict) else {}
|
|
run_durations = metadata.get("run_durations")
|
|
nodes[(ns, cp_id)] = _Node(
|
|
ns=ns,
|
|
cp_id=cp_id,
|
|
parent_ns=parent_ns,
|
|
parent_id=parent_id,
|
|
# Pruning is destructive, so the stamp fallback is disabled:
|
|
# title and goal leaves inherit the writer's stamps, and only
|
|
# the shape confirmation in
|
|
# _mark_duration_leaves_without_the_marker may widen the class.
|
|
duration_only=is_duration_only_checkpoint(tuple_, stamp_fallback=False),
|
|
metadata_source=metadata.get("source"),
|
|
carries_run_durations=isinstance(run_durations, dict) and bool(run_durations),
|
|
versions=frozenset(channel_versions.values()) if isinstance(channel_versions, dict) else frozenset(),
|
|
)
|
|
if not nodes:
|
|
# Nothing to classify: an empty (or unknown) thread is a no-op, but
|
|
# the before/after pair is part of the report contract, so it is
|
|
# mirrored rather than truncated — a caller aggregating measurements
|
|
# must not have to special-case "thread had no checkpoints".
|
|
if collect_stats:
|
|
report.stats_after = dict(report.stats_before)
|
|
return report
|
|
_mark_duration_leaves_without_the_marker(nodes)
|
|
|
|
children: dict[tuple[str, str], list[tuple[str, str]]] = defaultdict(list)
|
|
for key, node in nodes.items():
|
|
if node.parent_id is not None:
|
|
children[(node.parent_ns, node.parent_id)].append(key)
|
|
|
|
resumable = [key for key, node in nodes.items() if not node.duration_only]
|
|
# Head = newest by checkpoint id. LangGraph ids are time-ordered (uuid7):
|
|
# metadata step restarts from the fork point after a branch-resume, so it
|
|
# is not a thread-global sequence, while max-id matches what an unsaved
|
|
# ``aget_tuple`` resolves as the thread's latest state.
|
|
#
|
|
# Heads are selected *per namespace*. ``alist`` is called with the thread
|
|
# id alone, so ``nodes`` holds every namespace in the thread — a
|
|
# persistent subgraph (compiled with ``checkpointer=True``) keeps its own
|
|
# checkpoints under e.g. ``tools:<task>``. Protecting only one global
|
|
# head leaves the child namespace's latest checkpoint looking like an
|
|
# off-chain sibling leaf, which the opt-in E2 shape would then delete,
|
|
# rolling that graph's saved state back one step.
|
|
heads: dict[str, tuple[str, str]] = {}
|
|
for key in resumable:
|
|
current = heads.get(key[0])
|
|
if current is None or key[1] > current[1]:
|
|
heads[key[0]] = key
|
|
report.protected_head_ids = {ns: key[1] for ns, key in heads.items()}
|
|
|
|
chain: set[tuple[str, str]] = set()
|
|
for head in heads.values():
|
|
cursor: tuple[str, str] | None = head
|
|
while cursor is not None:
|
|
if cursor not in nodes:
|
|
# A head whose ancestor row is missing (partial damage from
|
|
# an earlier policy revision or manual cleanup) ends this
|
|
# walk instead of crashing the whole pass; the deletable loop
|
|
# below already leaves nodes with missing parents alone.
|
|
break
|
|
chain.add(cursor)
|
|
node = nodes[cursor]
|
|
cursor = (node.parent_ns, node.parent_id) if node.parent_id is not None else None
|
|
|
|
guarded: set[tuple[str, str]] = set()
|
|
if effective.strict_pending_write_guard:
|
|
guarded = await _checkpoint_ids_with_writes(storage_saver, thread_id)
|
|
|
|
deletable: list[tuple[str, str]] = []
|
|
for key, node in nodes.items():
|
|
if key in chain:
|
|
continue
|
|
if children.get(key):
|
|
continue
|
|
if key[1] in effective.protect_checkpoint_ids:
|
|
continue
|
|
if key in guarded:
|
|
continue
|
|
if node.parent_id is not None and (node.parent_ns, node.parent_id) not in nodes:
|
|
continue
|
|
if node.duration_only:
|
|
if not effective.prune_trailing_duration_leaves:
|
|
continue
|
|
elif not effective.prune_leaf_sibling_branches:
|
|
continue
|
|
deletable.append(key)
|
|
|
|
if effective.max_delete_per_run is not None:
|
|
deletable = deletable[: effective.max_delete_per_run]
|
|
|
|
# Blob GC, contract deletion mechanics: a blob row is an orphan only if no
|
|
# SURVIVING checkpoint references its version. A real duration-only leaf
|
|
# copies its parent's channel_versions verbatim, so its blobs are the
|
|
# parent's rows — deleting "blobs keyed by the removed checkpoint's own
|
|
# versions" would corrupt the surviving state.
|
|
deleted_keys = set(deletable)
|
|
survivor_versions: set[Any] = set()
|
|
for key, node in nodes.items():
|
|
if key not in deleted_keys:
|
|
survivor_versions.update(node.versions)
|
|
|
|
# Invalidate before touching source-of-truth rows. If cache deletion
|
|
# fails, fail closed with every checkpoint intact; doing this after the
|
|
# first DELETE could leave a stale wrapper entry for a missing node.
|
|
invalidate_cache = getattr(saver, "ainvalidate_history_cache", None)
|
|
if deletable and callable(invalidate_cache):
|
|
await invalidate_cache(thread_id)
|
|
|
|
for key in deletable:
|
|
await _delete_checkpoint_rows(storage_saver, thread_id, key)
|
|
report.deleted_checkpoint_ids.append(key[1])
|
|
|
|
if deletable:
|
|
await _delete_unreachable_blobs(storage_saver, thread_id, survivor_versions)
|
|
|
|
if collect_stats:
|
|
report.stats_after = await _thread_storage_stats(storage_saver, thread_id)
|
|
return report
|
|
|
|
|
|
async def enforce_trailing_duration_retention(
|
|
saver: BaseCheckpointSaver,
|
|
thread_id: str,
|
|
*,
|
|
protect_checkpoint_ids: frozenset[str] = frozenset(),
|
|
thread_lock: AbstractAsyncContextManager[None] | None = None,
|
|
) -> RetentionReport:
|
|
"""Prune at most the current root duration leaf with bounded storage work.
|
|
|
|
Read the head and its parent directly, then probe only the head's pending
|
|
writes. Identical channel-version maps prove that every blob referenced by
|
|
the leaf remains reachable through its parent, so no history scan, thread
|
|
statistics or blob GC is needed. Explicit admission selectors are protected.
|
|
Invalidate wrapper caches before deleting the source row; invalidation
|
|
failure must leave the row intact for delta reads.
|
|
|
|
As with the exhaustive service, callers must exclude same-thread mutation
|
|
for this entire operation. Gateway supplies durable admission and the local
|
|
checkpoint lock before the worker's first checkpoint read.
|
|
"""
|
|
storage = _ensure_supported_saver(saver)
|
|
report = RetentionReport(thread_id=thread_id)
|
|
lock: AbstractAsyncContextManager[None] = thread_lock if thread_lock is not None else nullcontext()
|
|
async with lock:
|
|
latest = await storage.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
|
|
if latest is None:
|
|
return report
|
|
latest_config = checkpoint_configurable(latest)
|
|
latest_id = latest_config.get("checkpoint_id")
|
|
if not latest_id or latest_config.get("thread_id") != thread_id or latest_config.get("checkpoint_ns") or latest_id in protect_checkpoint_ids:
|
|
return report
|
|
parent_config = getattr(latest, "parent_config", None)
|
|
parent_values = parent_config.get("configurable") if isinstance(parent_config, dict) else None
|
|
if not isinstance(parent_values, dict):
|
|
return report
|
|
parent_id = parent_values.get("checkpoint_id")
|
|
if not parent_id or parent_id == latest_id or parent_values.get("thread_id") != thread_id or parent_values.get("checkpoint_ns"):
|
|
return report
|
|
parent = await storage.aget_tuple(parent_config)
|
|
if parent is None or checkpoint_configurable(parent).get("checkpoint_id") != parent_id:
|
|
return report
|
|
|
|
latest_versions = latest.checkpoint.get("channel_versions")
|
|
parent_versions = parent.checkpoint.get("channel_versions")
|
|
if not isinstance(latest_versions, dict) or latest_versions != parent_versions:
|
|
return report
|
|
metadata = latest.metadata or {}
|
|
run_durations = metadata.get("run_durations")
|
|
latest_node = _Node(
|
|
ns="",
|
|
cp_id=latest_id,
|
|
parent_ns="",
|
|
parent_id=parent_id,
|
|
duration_only=is_duration_only_checkpoint(latest),
|
|
metadata_source=metadata.get("source"),
|
|
carries_run_durations=isinstance(run_durations, dict) and bool(run_durations),
|
|
versions=frozenset(latest_versions.values()),
|
|
)
|
|
parent_node = _Node(ns="", cp_id=parent_id, parent_ns=None, parent_id=None, duration_only=False, versions=latest_node.versions)
|
|
_mark_duration_leaves_without_the_marker({("", latest_id): latest_node, ("", parent_id): parent_node})
|
|
if not latest_node.duration_only or await _checkpoint_has_writes(storage, ("", latest_id), thread_id):
|
|
return report
|
|
|
|
invalidate_cache = getattr(saver, "ainvalidate_history_cache", None)
|
|
if callable(invalidate_cache):
|
|
await invalidate_cache(thread_id)
|
|
await _delete_checkpoint_rows(storage, thread_id, ("", latest_id))
|
|
report.protected_head_ids = {"": parent_id}
|
|
report.deleted_checkpoint_ids.append(latest_id)
|
|
return report
|
|
|
|
|
|
async def enforce_completed_run_retention(
|
|
*,
|
|
checkpointer: BaseCheckpointSaver,
|
|
thread_id: str,
|
|
protect_checkpoint_ids: frozenset[str] = frozenset(),
|
|
) -> RetentionReport:
|
|
"""Run the conservative policy inside the next run's durable admission.
|
|
|
|
Successful runs append one metadata-only duration checkpoint. If it is
|
|
left in place, the next run forks from it and turns that otherwise
|
|
disposable leaf into a protected ancestor forever. The run worker calls
|
|
this hook only after its active durable row excludes all same-thread run,
|
|
branch and checkpoint-write admissions, and before it reads a checkpoint.
|
|
The process-local checkpoint lock additionally serializes the local saver.
|
|
|
|
Only the contract-proven trailing duration leaf is eligible. Sibling
|
|
branches and resumable checkpoints remain protected. Storage work is bounded
|
|
to the root head, its parent and one pending-write probe, independently of
|
|
thread history length.
|
|
"""
|
|
from deerflow.runtime.runs.worker import _checkpoint_thread_lock
|
|
|
|
report = await enforce_trailing_duration_retention(
|
|
checkpointer,
|
|
thread_id,
|
|
protect_checkpoint_ids=protect_checkpoint_ids,
|
|
thread_lock=_checkpoint_thread_lock(thread_id),
|
|
)
|
|
if report.deleted_checkpoint_ids:
|
|
logger.info(
|
|
"Checkpoint retention removed trailing metadata checkpoints for thread %s: %s",
|
|
thread_id,
|
|
report.deleted_checkpoint_ids,
|
|
)
|
|
return report
|