From 57e53fc6006232911f48bc19667636ea0768c6b3 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 22 Apr 2026 14:12:57 -0400 Subject: [PATCH] =?UTF-8?q?chore(delta-channel):=20remove=20snapshot=5Feve?= =?UTF-8?q?ry=20=E2=80=94=20simpler=20design,=20better=20storage=20savings?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 --- .../langgraph/checkpoint/postgres/aio.py | 8 +- .../langgraph/checkpoint/postgres/base.py | 8 +- .../langgraph/checkpoint/memory/__init__.py | 4 +- libs/langgraph/langgraph/channels/delta.py | 9 +- .../tests/test_delta_channel_benchmark.py | 19 +- libs/langgraph/tests/test_rehydrate_sweep.py | 185 ------------------ 6 files changed, 21 insertions(+), 212 deletions(-) delete mode 100644 libs/langgraph/tests/test_rehydrate_sweep.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 2f461e75f..595e240f3 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -400,10 +400,7 @@ class AsyncPostgresSaver(BasePostgresSaver): channel: str, cur: Any, ) -> list[Any]: - """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest (async). - - Two queries instead of a recursive CTE — see sync version for rationale. - """ + """Async version of _get_channel_writes_cur — see sync version for rationale.""" await cur.execute( "SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints " "WHERE thread_id = %s AND checkpoint_ns = %s", @@ -427,9 +424,8 @@ class AsyncPostgresSaver(BasePostgresSaver): "ORDER BY task_id, idx", (thread_id, checkpoint_ns, channel, ancestor_ids), ) - rows = await cur.fetchall() writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list) - for row in rows: + for row in await cur.fetchall(): writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"])) result = [] for cid in reversed(ancestor_ids): diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 2e679b88e..0a17f052a 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -225,8 +225,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): ) -> list[Any]: """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest. - Two queries instead of a recursive CTE: - 1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread — cheap, just IDs. + Two queries: + 1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread — cheap, IDs only. 2. Walk the ancestor chain in Python, then fetch writes with a plain ANY() filter. """ cur.execute( @@ -237,9 +237,6 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): parent_map: dict[str, str | None] = { row["checkpoint_id"]: row["parent_checkpoint_id"] for row in cur.fetchall() } - # Walk newest→oldest starting from the current checkpoint's parent. - # Writes stored under checkpoint C produced the state *after* C, so we - # want ancestors of the current checkpoint (not the checkpoint itself). ancestor_ids: list[str] = [] cid: str | None = parent_map.get(checkpoint_id) while cid is not None: @@ -257,7 +254,6 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list) for row in cur.fetchall(): writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"])) - # ancestor_ids is newest→oldest; replay oldest→newest result = [] for cid in reversed(ancestor_ids): for type_tag, blob in writes_by_cp.get(cid, []): diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 76b3fe2f1..f265fe860 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -153,7 +153,7 @@ class InMemorySaver( checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = config["configurable"].get("checkpoint_id", "") ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {}) - # Walk the parent chain, collecting checkpoint IDs oldest→newest. + # Walk the parent chain newest→oldest collecting checkpoint IDs. chain: list[str] = [] current: str | None = checkpoint_id while current is not None: @@ -163,7 +163,7 @@ class InMemorySaver( chain.append(current) _, _, parent = entry current = parent - # Collect writes for `channel` from each checkpoint in oldest→newest order. + # Collect writes oldest→newest. result: list[Any] = [] for cp_id in reversed(chain): step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {}) diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index fb8b935c6..3d5e02693 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -41,8 +41,6 @@ class DeltaChannel( def __init__( self, operator: Callable[[list[Value], Any], list[Value]], - *, - snapshot_every: int | None = None, # reserved for future use ) -> None: super().__init__(list) self.operator = operator @@ -83,8 +81,11 @@ class DeltaChannel( except Exception: new.value = [] elif isinstance(checkpoint, list): - # Flat list of individual write values (oldest→newest) from get_channel_writes. - value: Any = new.typ() + # Flat list of write values (oldest→newest) from get_channel_writes. + try: + value: Any = new.typ() + except Exception: + value = [] for write in checkpoint: value = new.operator(value, write) new.value = value diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index 494a83708..3a641dddd 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -219,7 +219,7 @@ def _approx_tokens(n_turns: int) -> str: # Turn counts chosen to demonstrate O(N²) vs O(N) storage growth without running too long. # Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window). -TURN_COUNTS = [10, 25, 50, 100] +TURN_COUNTS = [10, 25, 50, 100, 500] def _checkpointer_factories() -> list[tuple[str, Any]]: @@ -240,6 +240,7 @@ def run_benchmark() -> None: if _POSTGRES_AVAILABLE: try: import psycopg + psycopg.connect(_POSTGRES_URI).close() checkpointers.append(("Postgres (recursive CTE)", "postgres")) except Exception: @@ -292,7 +293,7 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: ) print("-" * W) storage_results = [] - for turns, b_bytes, d_bytes, *_ in rows: + for turns, b_bytes, d_bytes, b_rt, d_rt in rows: if b_bytes < 0: print( f"{turns:>6} {_approx_tokens(turns):>10} {'n/a':>12} {'n/a':>12} {'n/a':>8}" @@ -311,9 +312,7 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: # ── Table 2: Read latency ───────────────────────────────────────────────── print("Read latency (avg of 5 get_state calls)") print("=" * W) - print( - f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}" - ) + print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}") print("-" * W) for turns, b_bytes, d_bytes, b_rt, d_rt in rows: print( @@ -324,9 +323,9 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: print() if storage_results: - best = storage_results[-1] - turns, b_bytes, d_bytes, ratio = best - _, _, _, b_rt, d_rt = rows[-1] + turns, b_bytes, d_bytes, ratio = storage_results[-1] + b_rt = rows[-1][-2] + d_rt = rows[-1][-1] print( f"At {turns} turns: {_fmt_bytes(b_bytes)} → {_fmt_bytes(d_bytes)} ({ratio:.0f}x less storage); " f"read {b_rt * 1000:.1f}ms → {d_rt * 1000:.1f}ms" @@ -335,7 +334,9 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: print("Legend:") print(" add_msgs = Annotated[list, add_messages] — O(N²) storage") - print(" delta = DeltaChannel(add_messages) — O(N) storage, reconstructed from writes") + print( + " delta = DeltaChannel(add_messages) — O(N) storage, full chain replay" + ) print() diff --git a/libs/langgraph/tests/test_rehydrate_sweep.py b/libs/langgraph/tests/test_rehydrate_sweep.py deleted file mode 100644 index 5a72ed9b2..000000000 --- a/libs/langgraph/tests/test_rehydrate_sweep.py +++ /dev/null @@ -1,185 +0,0 @@ -"""Sweep snapshot_every values to find the storage vs. time-travel tradeoff. - -Run directly: python tests/test_rehydrate_sweep.py -Run via pytest: pytest tests/test_rehydrate_sweep.py -s -""" - -from __future__ import annotations - -import sys -import time -from typing import Annotated, Any - -from langchain_core.messages import AIMessage, HumanMessage -from langgraph.checkpoint.memory import MemorySaver -from typing_extensions import TypedDict - -from langgraph.channels.delta import DeltaChannel -from langgraph.graph import END, StateGraph -from langgraph.graph.message import add_messages - -# --------------------------------------------------------------------------- -# Config -# --------------------------------------------------------------------------- - -REHYDRATE_SWEEP = [5, 10, 25, 50, 100, None] # None = no rehydration (pure diff) -TURN_COUNTS = [50, 100, 250, 500] - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_state(snapshot_every: int | None) -> type: - channel = DeltaChannel(add_messages, snapshot_every=snapshot_every) - return TypedDict("S", {"messages": Annotated[list, channel]}) - - -def _make_graph(state_cls: type) -> Any: - def human_node(state: Any) -> dict: - return {} - - def ai_node(state: Any) -> dict: - last = state["messages"][-1] - return {"messages": [AIMessage(content=f"reply-to-{last.id}")]} - - g = StateGraph(state_cls) - g.add_node("human", human_node) - g.add_node("ai", ai_node) - g.add_edge("human", "ai") - g.add_edge("ai", END) - g.set_entry_point("human") - return g.compile(checkpointer=MemorySaver()) - - -def _total_blob_bytes(saver: MemorySaver) -> int: - total = 0 - for (_, _, _, _), (type_tag, blob) in saver.blobs.items(): - if blob is not None: - total += len(blob) - return total - - -def _measure_time_travel_ms(graph: Any, config: dict) -> float: - """Time how long it takes to get state at the very first checkpoint (worst case).""" - history = list(graph.get_state_history(config)) - if not history: - return 0.0 - oldest = history[-1] - t0 = time.perf_counter() - graph.get_state(oldest.config) - return (time.perf_counter() - t0) * 1000 - - -def _run(n_turns: int, snapshot_every: int | None) -> tuple[float, int, float]: - """Returns (write_ms, blob_bytes, time_travel_ms).""" - state_cls = _make_state(snapshot_every) - graph = _make_graph(state_cls) - saver: MemorySaver = graph.checkpointer # type: ignore[assignment] - config = {"configurable": {"thread_id": "sweep"}} - - t0 = time.perf_counter() - for i in range(n_turns): - graph.invoke( - {"messages": [HumanMessage(content=f"msg-{i}", id=f"h{i}")]}, config - ) - write_ms = (time.perf_counter() - t0) * 1000 - - blob_bytes = _total_blob_bytes(saver) - tt_ms = _measure_time_travel_ms(graph, config) - return write_ms, blob_bytes, tt_ms - - -# --------------------------------------------------------------------------- -# ASCII sparkline -# --------------------------------------------------------------------------- - - -def _sparkline(values: list[float], width: int = 20) -> str: - bars = " ▁▂▃▄▅▆▇█" - lo, hi = min(values), max(values) - span = hi - lo or 1 - chars = [bars[round((v - lo) / span * (len(bars) - 1))] for v in values] - return "".join(chars).ljust(width) - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - - -def run_sweep() -> None: - label = {v: (str(v) if v is not None else "None(∞)") for v in REHYDRATE_SWEEP} - - print() - print("snapshot_every sweep — storage vs time-travel cost") - print("=" * 90) - - for turns in TURN_COUNTS: - print(f"\n--- {turns} turns ---") - col_w = 12 - header = ( - f"{'snapshot_every':>18} " - f"{'blob_bytes':>{col_w}} " - f"{'write_ms':>{col_w}} " - f"{'time_travel_ms':>{col_w}}" - ) - print(header) - print("-" * 60) - - tt_vals: list[float] = [] - byte_vals: list[int] = [] - write_vals: list[float] = [] - rows: list[tuple] = [] - - for rv in REHYDRATE_SWEEP: - write_ms, blob_bytes, tt_ms = _run(turns, rv) - rows.append((rv, blob_bytes, write_ms, tt_ms)) - byte_vals.append(blob_bytes) - write_vals.append(write_ms) - tt_vals.append(tt_ms) - - for rv, blob_bytes, write_ms, tt_ms in rows: - print( - f"{label[rv]:>18} " - f"{blob_bytes:>{col_w},} " - f"{write_ms:>{col_w}.1f} " - f"{tt_ms:>{col_w}.2f}" - ) - - print() - print( - f" bytes spark: [{_sparkline(byte_vals)}] " - f"lo={min(byte_vals):,} hi={max(byte_vals):,}" - ) - print( - f" time-travel spark: [{_sparkline(tt_vals)}] " - f"lo={min(tt_vals):.2f}ms hi={max(tt_vals):.2f}ms" - ) - print( - f" write spark: [{_sparkline(write_vals)}] " - f"lo={min(write_vals):.1f}ms hi={max(write_vals):.1f}ms" - ) - - print() - print("=" * 90) - print( - "snapshot_every=None means pure diff (no snapshots) — " - "lowest storage, highest time-travel cost." - ) - print( - "Lower snapshot_every = more frequent full snapshots = " - "faster time-travel, more storage." - ) - print() - - -def test_rehydrate_sweep(capsys: Any) -> None: - with capsys.disabled(): - run_sweep() - - -if __name__ == "__main__": - run_sweep() - sys.exit(0)