mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-01 05:55:14 +02:00
chore(langgraph): use two phase read to avoid unnecessary data transport (#7660)
## Summary Replaces the single-roundtrip `UNION ALL` DeltaChannel read with a two-stage query that avoids fetching unused snapshot blobs, then removes the old combined path entirely. ### Problem `_get_channel_writes_history` used a single `UNION ALL` query that fetched **all** checkpoint metadata, writes, and blobs for a `(thread_id, channel)` in one shot. With `snapshot_frequency=N`, this pulled back O(N/freq) full-size snapshot blobs even though only the nearest one is needed to seed reconstruction. At 500 turns with `snapshot_frequency=10`, this meant fetching ~100 complete message-history snapshots per read. ### Solution Two-stage read: - **Stage 1** — lightweight scan of `checkpoints` only (no blob bytes): walks the parent chain from the target checkpoint and stops at the first ancestor with a snapshot, returning `chain_cids` and `seed_version` - **Stage 2** — targeted fetch: only the writes for `chain_cids` and the single seed blob at `seed_version` The two-stage path is now unconditional — the old combined query and `LG_DELTA_TWO_STAGE_QUERY` env-var gate have been removed. ### Sentinel cleanup `DELTA_SENTINEL` is now a pure in-memory signal and is never written to storage: - Postgres `put()` already stripped it from `channel_values` before writing blobs - Memory saver `put()` now stores `"empty"` instead of serializing the sentinel - `EXT_DELTA_SENTINEL` (msgpack ext code 8) removed from `JsonPlusSerializer` - `DELTA_SENTINEL` is kept as an in-memory marker: `DeltaChannel.checkpoint()` returns it so savers know to skip it, and `_ChannelWritesHistory.seed` uses it to mean "no snapshot found, start from empty" ## Performance Benchmarked at `snapshot_frequency=10` on Postgres (`~100 tok/msg`): | turns | old combined query | two-stage | |------:|-------------------:|----------:| | 50 | 6.0ms | 2.8ms (2.1x faster) | | 100 | 10.1ms | 5.6ms (1.8x faster) | | 500 | **216.1ms** | 15.3ms (**14x faster**) | The old query's read time grew super-linearly with turn count because each read fetched O(N/freq) full snapshot blobs. Two-stage keeps read depth bounded by `snapshot_frequency` regardless of thread length. ## Test plan - `make test` in `libs/checkpoint`, `libs/checkpoint-postgres`, `libs/langgraph` - Removed `test_delta_sentinel_serde_round_trip` (sentinel no longer serializable) - Updated `test_memory.py` — delta channel blobs stored as `"empty"`, not serialized sentinel - Updated `test_channels.py` — `channel_values` no longer contains sentinel key for DeltaChannels - Deleted `test_delta_channel_two_stage_benchmark.py` (one-stage vs two-stage comparison; path no longer exists) --------- Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Sydney Runkle
Claude Sonnet 4.6
parent
c8c58a0768
commit
d8b7800183
@@ -19,6 +19,7 @@ from langgraph.checkpoint.base import (
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
@@ -26,9 +27,11 @@ from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
SELECT_DELTA_COMBINED_SQL,
|
||||
SELECT_DELTA_STAGE1_SQL,
|
||||
SELECT_DELTA_STAGE2_SQL,
|
||||
BasePostgresSaver,
|
||||
_DeltaCombinedRow,
|
||||
_DeltaStage1Row,
|
||||
_DeltaStage2Row,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
|
||||
@@ -308,7 +311,12 @@ class PostgresSaver(BasePostgresSaver):
|
||||
# others are stored in blobs table
|
||||
blob_values = {}
|
||||
for k, v in checkpoint["channel_values"].items():
|
||||
if v is None or isinstance(v, (str, int, float, bool)):
|
||||
if v is DELTA_SENTINEL:
|
||||
copy["channel_values"].pop(k)
|
||||
elif isinstance(v, _DeltaSnapshot):
|
||||
blob_values[k] = copy["channel_values"].pop(k)
|
||||
copy["channel_values"][k] = True
|
||||
elif v is None or isinstance(v, (str, int, float, bool)):
|
||||
pass
|
||||
else:
|
||||
blob_values[k] = copy["channel_values"].pop(k)
|
||||
@@ -441,41 +449,49 @@ class PostgresSaver(BasePostgresSaver):
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
|
||||
|
||||
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
|
||||
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
|
||||
single roundtrip; the ancestor walk runs in Python.
|
||||
Two-stage query: stage 1 scans checkpoint metadata to walk the parent
|
||||
chain and locate the nearest snapshot; stage 2 fetches only the
|
||||
chain-limited writes and single seed blob.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
# Caller didn't specify a target — resolve to the latest
|
||||
# checkpoint on the thread. `get_tuple` without `checkpoint_id`
|
||||
# returns the newest; its config carries the resolved id.
|
||||
target = self.get_tuple(config)
|
||||
if target is None:
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
||||
|
||||
with self._cursor() as cur:
|
||||
cur.execute(
|
||||
SELECT_DELTA_COMBINED_SQL,
|
||||
SELECT_DELTA_STAGE1_SQL,
|
||||
(channel, channel, thread_id, checkpoint_ns),
|
||||
)
|
||||
stage1_rows = cur.fetchall()
|
||||
chain_cids, seed_version = self._walk_stage1(
|
||||
cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id
|
||||
)
|
||||
seed_versions = [seed_version] if seed_version else []
|
||||
with self._cursor() as cur:
|
||||
cur.execute(
|
||||
SELECT_DELTA_STAGE2_SQL,
|
||||
(
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
chain_cids,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
seed_versions,
|
||||
),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
stage2_rows = cur.fetchall()
|
||||
return self._build_delta_channel_writes_history(
|
||||
channel=channel,
|
||||
target_id=checkpoint_id,
|
||||
rows=cast("list[_DeltaCombinedRow]", rows),
|
||||
chain_cids=chain_cids,
|
||||
seed_version=seed_version,
|
||||
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
|
||||
)
|
||||
|
||||
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
|
||||
Reference in New Issue
Block a user