This commit is contained in:
Sydney Runkle
2026-04-29 14:05:06 -04:00
parent ee5fd582b8
commit be7101b0ff
3 changed files with 0 additions and 305 deletions
-1
View File
@@ -101,4 +101,3 @@ dmypy.json
.editorconfig
.scratch
.worktrees/
new_pr_desc.md
-167
View File
@@ -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 3100× 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__ == "<lambda>"`.
- `_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).
-137
View File
@@ -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.