mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-26 17:42:24 +02:00
nits
This commit is contained in:
@@ -101,4 +101,3 @@ dmypy.json
|
||||
.editorconfig
|
||||
.scratch
|
||||
.worktrees/
|
||||
new_pr_desc.md
|
||||
|
||||
-167
@@ -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__ == "<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).
|
||||
@@ -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.
|
||||
Reference in New Issue
Block a user