From be7101b0ff9f5c4c1550d957abf5b3e7d12e2abc Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 29 Apr 2026 14:05:06 -0400 Subject: [PATCH] nits --- .gitignore | 1 - new_pr_desc.md | 167 ----------------------------- notes/delta_channel_query_bench.md | 137 ----------------------- 3 files changed, 305 deletions(-) delete mode 100644 new_pr_desc.md delete mode 100644 notes/delta_channel_query_bench.md diff --git a/.gitignore b/.gitignore index d34e545b9..b154d3b56 100644 --- a/.gitignore +++ b/.gitignore @@ -101,4 +101,3 @@ dmypy.json .editorconfig .scratch .worktrees/ -new_pr_desc.md diff --git a/new_pr_desc.md b/new_pr_desc.md deleted file mode 100644 index 1be12ac31..000000000 --- a/new_pr_desc.md +++ /dev/null @@ -1,167 +0,0 @@ -# DeltaChannel: sentinel-based checkpoint blobs + write-replay reconstruction - -## TL;DR - -Introduces `DeltaChannel`, a new fold-reducer channel that stores only a -zero-byte sentinel in checkpoint blobs instead of the full accumulated value. -On load, the runtime replays the channel's ancestor writes through the reducer -to reconstruct state. For long-running threads with large accumulating state -(e.g. message histories), this trades read-time work for dramatically smaller -checkpoint blobs and avoids redundant duplication of the full value at every -step. - -An optional `snapshot_frequency=N` parameter bounds replay depth by writing a -full `_DeltaSnapshot` blob every N steps — a configurable storage vs. latency -tradeoff. - ---- - -## DeltaChannel - -### New channel type (`libs/langgraph/langgraph/channels/delta.py`) - -`DeltaChannel(typ, operator, *, snapshot_frequency=None)` is a fold-reducer -channel (same semantics as `BinaryOperatorAggregate`) with a different -checkpoint strategy: - -- **`checkpoint()`** always returns `DELTA_SENTINEL` (a zero-byte sentinel), - never the accumulated value. -- **`from_checkpoint()`** reconstructs value from a seed blob + replayed - writes returned by `_get_channel_writes_history`. -- **`snapshot_frequency=N`**: `create_checkpoint` writes a `_DeltaSnapshot` - blob every N steps, bounding the ancestor walk to at most N checkpoints. - Snapshots are written eagerly — even if the channel had no write that step, - a version bump forces `put()` to store the blob. - -The constructor signature mirrors `BinaryOperatorAggregate`: `typ` is required -as the first positional argument and is normalized to its concrete counterpart -(e.g. `Sequence → list`, `Mapping → dict`) in `__init__`. `_is_field_channel` -in `graph/state.py` was updated to reconstruct the channel with the correct -`typ` from the `Annotated` outer type, instead of patching `item.typ` after -construction. - -Both `__init__` and the clone path in `copy()` / `from_checkpoint()` start -with `value = MISSING` — no inconsistency between fresh construction and -checkpoint-loaded clones. - -### New sentinel and snapshot types (`libs/checkpoint/langgraph/checkpoint/serde/types.py`) - -- **`_DeltaSentinel` / `DELTA_SENTINEL`**: singleton marker. Identity-compared - (`is DELTA_SENTINEL`) throughout; `loads_typed` always returns the same - instance. -- **`_DeltaSnapshot(NamedTuple)`**: wraps the full accumulated value at a - snapshot step. Serialized via a dedicated msgpack ext code - (`EXT_DELTA_SNAPSHOT = 7`), so it round-trips through the standard serde - path without a separate type tag. - -### Serializer support (`libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py`) - -Both delta types serialize through msgpack, keeping them in the same codec path with no special string type tags: - -- `DELTA_SENTINEL` → msgpack ext code 8 (`EXT_DELTA_SENTINEL`, zero data bytes). -- `_DeltaSnapshot` → msgpack ext code 7 (`EXT_DELTA_SNAPSHOT`, value packed as nested msgpack). -- Ext hooks decode both back to their singleton / NamedTuple counterparts. - -### Ancestor-walk API on `BaseCheckpointSaver` (`libs/checkpoint/langgraph/checkpoint/base/__init__.py`) - -Three new methods on `BaseCheckpointSaver` (all experimental / underscore-prefixed): - -- **`_ChannelWritesHistory(NamedTuple)`** — return type carrying `seed: Any` - (nearest non-sentinel ancestor blob, or `DELTA_SENTINEL` if none found) and - `writes: list[PendingWrite]` (oldest→newest on-path deltas). -- **`_get_tuple_raw(config)`** / **`_aget_tuple_raw(config)`** — pure storage - reads used by the base walk implementation. Default delegates to `get_tuple`; - savers whose `get_tuple` performs channel hydration (which calls - `channels_from_checkpoint` which calls `_get_channel_writes_history`) - override this to break the re-entrancy cycle. -- **`_get_channel_writes_history(config, channel)`** / - **`_aget_channel_writes_history(config, channel)`** — reference - implementation walks the parent chain via `_get_tuple_raw`, collecting - `pending_writes` for the target channel and stopping at the first - non-sentinel blob. - -### `InMemorySaver` optimized override (`libs/checkpoint/langgraph/checkpoint/memory/__init__.py`) - -`InMemorySaver._get_channel_writes_history` reads directly from -`self.storage`, `self.blobs`, and `self.writes` in one pass without repeated -`get_tuple` calls. Handles two distinct blob-termination cases: - -- **Pre-delta blob** (plain value from `BinaryOperatorAggregate` era): the - blob IS the full state at that ancestor — do not replay its pending_writes - (they are already baked in). Return immediately. -- **`_DeltaSnapshot` blob**: the snapshot IS the state, but pending_writes at - that ancestor encode the NEXT step's transition and must be collected before - returning the snapshot as `seed`. - -Also adds **`InMemorySaver.prune(thread_ids, *, strategy)`** — delta-aware -pruning that retains the minimal ancestor chain needed to reconstruct sentinel -channels for the latest checkpoint. - -### PostgresSaver / AsyncPostgresSaver optimized override (`libs/checkpoint-postgres/`) - -Both sync and async postgres savers override -`_get_channel_writes_history` / `_aget_channel_writes_history` with a -**single-roundtrip UNION ALL query** (`SELECT_DELTA_COMBINED_SQL`) that -fetches rows from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` -in one query. Rows are assembled by the shared pure helper -`_build_delta_channel_writes_history` on `BasePostgresSaver`. - -- **`_DeltaCombinedRow(TypedDict, total=False)`** — typed view of the nine - columns emitted by the UNION ALL (kind-tagged `"p"` / `"w"` / `"b"`). -- **`_load_blobs`** parameter typed from `Any` to - `Sequence[tuple[bytes, bytes, bytes]]`. - -A benchmark (`notes/delta_channel_query_bench.md`) shows the prior recursive -CTE carried a hidden O(ancestors × blobs_in_thread) join; the plain UNION ALL -is 3–100× faster in the realistic depth range. - -### Pregel integration (`libs/langgraph/langgraph/pregel/`) - -**`_checkpoint.py`**: -- `channels_from_checkpoint` gains `saver` and `config` keyword args; for - each channel where `_needs_replay` is true (DeltaChannel with a sentinel - blob), it calls `saver._get_channel_writes_history` and replays writes. -- `achannels_from_checkpoint` added — async counterpart, used by - `AsyncPregelLoop`. -- `create_checkpoint` gains `get_next_version` kwarg; on snapshot steps it - writes `_DeltaSnapshot` blobs and bumps channel versions eagerly. -- `_needs_replay(spec, stored)` helper — True iff spec is a `DeltaChannel` - and stored blob is `MISSING` or `DELTA_SENTINEL`. - -**`_loop.py`**: -- `SyncPregelLoop.__enter__` passes `saver` + `config` to - `channels_from_checkpoint`. -- `AsyncPregelLoop.__aenter__` calls `achannels_from_checkpoint`. -- **Async write-ordering safety**: `AsyncPregelLoop` maintains - `_delta_write_futs`, a list of in-flight `aput_writes` futures for - DeltaChannel channels. In `accept_writes`, any write to a `DeltaChannel` - appends its future to this list. In `_checkpointer_put_after_previous`, - the list is drained via `await asyncio.gather()` before `aput()` is - called. This ensures checkpoint_writes are durable before the sentinel - blob is stored — a sentinel backed by missing writes would cause silent - data loss on replay. The sync loop does not need this guard: all - background tasks complete before `invoke()` returns via - `BackgroundExecutor.__exit__`. - -### `binop.py` refactoring - -- `_strip_extras`: fixed ordering — `Required`/`NotRequired` check must come - before the generic `__origin__` recurse to avoid swallowing the inner arg. -- `_operators_equal` extracted as a shared helper (used by both - `BinaryOperatorAggregate.__eq__` and `DeltaChannel.__eq__`): lambdas are - treated as always-equal since they all share `__name__ == ""`. -- `_get_overwrite`: `set(value.keys()) == {OVERWRITE}` → `len(value) == 1 and - OVERWRITE in value` (avoids throwaway set allocation on every call). - ---- - -## Follow-ups - -- **Batch reconstruction across multiple DeltaChannels**: each channel - currently issues its own `_get_channel_writes_history` call. A single - ancestor walk collecting writes for all sentinel channels at once would - reduce roundtrips proportionally to the number of DeltaChannel fields in a - state schema. -- **Store delta epoch ids in writes table**: this would allow for more targeted -reads of the writes table when reconstructing delta channels (especialy valuable -for postgres). diff --git a/notes/delta_channel_query_bench.md b/notes/delta_channel_query_bench.md deleted file mode 100644 index 8dd69728a..000000000 --- a/notes/delta_channel_query_bench.md +++ /dev/null @@ -1,137 +0,0 @@ -# Delta-channel reconstruction: query strategy benchmark - -**Branch:** `delta-channel-writes-based` -**Question (Nuno):** Is the recursive CTE the right query shape for reconstructing a delta channel inside `get_tuple`, or would a plain `SELECT WHERE` be cheaper even though it returns more rows? -**Answer:** Plain `SELECT WHERE` wins at every realistic depth. The recursion isn't the problem — the JSON-expression join inside the CTE is. - -## Setup - -- Postgres 16 on `localhost:5441` (the `compose-postgres.yml` instance, run directly without docker for this round) -- Single delta channel `messages`, one write per checkpoint, `DELTA_SENTINEL` blob per checkpoint -- Linear chain (`branch=1`) and 5-way branching at every step (`branch=5`) — branching is the case where plain over-fetches sibling rows -- Median of 20 timed runs after 3 warmups, fresh psycopg cursor per strategy -- Bench script: `bench_get_tuple_strategies.py` at repo root - -Three strategies compared: - -| name | roundtrips | shape | -|------|-----------|-------| -| `cte` | 1 | Current prod: recursive CTE walks ancestors, LEFT JOINs writes + blobs | -| `plain` | 3 | Nuno's suggestion: thread-wide `SELECT WHERE` per table, Python walks parent chain and filters | -| `cte+narrow` | 2 | CTE returns ancestor IDs only, then one `UNION ALL` of writes + blobs filtered by `ANY(ids)` | - -## Results (ms per get_tuple, median of 20) - -``` - depth branch cte plain cte+narrow rows_cte rows_plain plain/cte - 10 1 0.14ms 0.26ms 0.24ms 9 30 1.89x - 10 5 0.21ms 0.23ms 0.17ms 9 110 1.11x - 50 1 0.89ms 0.27ms 0.33ms 49 150 0.31x - 50 5 2.35ms 0.66ms 0.51ms 49 550 0.28x - 200 1 11.61ms 0.78ms 1.30ms 199 600 0.07x - 200 5 34.79ms 2.33ms 3.07ms 199 2200 0.07x - 1000 1 274.60ms 2.59ms 13.29ms 999 3000 0.01x - 1000 5 856.01ms 10.14ms 15.31ms 999 11000 0.01x -``` - -Lower is better. `plain/cte < 1` means plain is faster. - -### Headline numbers - -- depth 50: plain is **3x** faster -- depth 200: plain is **15x** faster -- depth 1000: plain is **~100x** faster -- Branching makes plain over-fetch (3000 rows → 11000 rows at d=1000), but it remains ~85x faster than the CTE - -## Why the CTE collapses - -`EXPLAIN (ANALYZE, BUFFERS)` of the CTE at depth 1000 (linear). Excerpt with the load-bearing nodes: - -``` -Sort ... actual time=137.798..137.827 rows=999 - CTE ancestors - -> Recursive Union ... actual time=0.005..2.443 rows=999 - ^^^^^^ - recursion is 2.4 ms — fine - -> Nested Loop Left Join ... actual time=2.676..137.529 rows=999 - Join Filter: (cw.checkpoint_id = a.cid) - Rows Removed by Join Filter: 998001 - ^^^^^^^ - 999 ancestors x ~1000 writes - -> Nested Loop Left Join ... actual time=2.669..85.061 rows=999 - Join Filter: (bl.version = ((c.checkpoint -> 'channel_versions'::text) ->> bl.channel)) - Rows Removed by Join Filter: 998001 - ^^^^^^^ - same quadratic blow-up on the blob join -``` - -Two pathological things are happening: - -1. **The blob join filter is on a JSON expression**: `bl.version = (c.checkpoint -> 'channel_versions' ->> bl.channel)`. The planner cannot push this into an index lookup, so it materializes `checkpoint_blobs` for the thread and does a nested-loop comparison against every ancestor — a Cartesian product that grows as `O(ancestors × blobs_in_thread)`. -2. **The writes join is similar**: writes for the thread are materialized once, then nested-loop joined against ancestors with a `Join Filter` rather than a hash/merge join over the indexed `checkpoint_id`. - -At depth 1000 that's **~2 million rows evaluated, 99.9% of them discarded**. The recursion itself is a rounding error. - -For comparison, the plain Q1 (`SELECT … FROM checkpoints WHERE thread_id=? AND checkpoint_ns=?`) at depth 1000: - -``` -Seq Scan on checkpoints ... actual time=0.012..0.121 rows=1000 -Execution Time: 0.140 ms -``` - -A simple seq scan over 57 buffers. Q2 and Q3 follow the same shape and complete in well under 1 ms each. - -## Crossover and remote-DB reasoning - -- Pure local Postgres: plain wins from depth ~30 onward; CTE wins by fractions of a ms below that -- Remote Postgres at ~5 ms RTT adds ~10 ms to plain (3 roundtrips vs 1). Crossover shifts to ~depth 30. Above that, the CTE's quadratic SQL cost still dominates the RTT savings. - -There is no realistic conversation depth where the CTE wins on a remote DB. At depth 200+ (anything resembling a real multi-turn agent run) plain is faster regardless of network. - -## Recommendation - -**Switch to plain SELECT WHERE, one delta channel at a time.** - -Three indexed queries per delta channel: - -```sql --- Q1: parent chain + per-checkpoint version of this channel -SELECT checkpoint_id, - parent_checkpoint_id, - checkpoint -> 'channel_versions' ->> 'channel_name' AS ver -FROM checkpoints -WHERE thread_id = ? AND checkpoint_ns = ?; - --- Q2: writes for this channel, anywhere in the thread -SELECT checkpoint_id, type, blob, task_id, idx -FROM checkpoint_writes -WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?; - --- Q3: blobs for this channel, anywhere in the thread -SELECT version, type, blob -FROM checkpoint_blobs -WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?; -``` - -Python then: -- Builds `parent_of: dict[cid, parent_cid]` from Q1 -- Walks from target's parent newest → oldest -- Filters Q2 rows by `ancestor_set`, processes oldest → newest, applies overwrite-terminator -- Picks seed blob via the per-ancestor `ver` map, terminates at first non-sentinel blob - -All O(n) on n = thread checkpoints, with tight constants (dict lookups). No recursion, no JSON-expression joins, no quadratic plans. - -If the 3-roundtrip cost ever shows up on remote-DB benchmarks, fold Q2 + Q3 into one `UNION ALL` to get back to 2 roundtrips. Bench says it isn't worth the SQL complexity right now. - -## Bonus: code simplification from single-channel scope - -Multi-channel reconstruction in the current `_reconstruct_delta_channels_cur` carries: - -- `rows_by_cid` nested dicts, keyed by cid then channel -- `seen_blob: set[(cid, channel)]` and `seen_write: set[(cid, channel, task_id, idx)]` dedup -- `collected: dict[channel, list]`, `done: set[channel]`, `seeds: dict[channel, value]` -- Inner `for ch in channels_list` loops and an early-exit `if len(done) == len(channels_list)` - -Single-channel collapses these to a single list, a single bool, and one `Optional[Any]`. Roughly half the Python in that function, plus an obvious shape for splitting pure post-processing into `base.py` so sync and async stop duplicating it. - -If multi-channel coalescing turns out to matter later, it can come back as a SQL-level optimization without re-introducing the bookkeeping in Python.