Compare commits

...
Author SHA1 Message Date
Quanzheng LongandCursor bc263888b1 chore: ruff format fixes
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-05-01 10:41:56 -07:00
Quanzheng Long 8c2a30af8a combine 2026-05-01 10:30:54 -07:00
9 changed files with 466 additions and 297 deletions
@@ -2,7 +2,7 @@ from __future__ import annotations
import threading import threading
from collections import defaultdict from collections import defaultdict
from collections.abc import Iterator, Sequence from collections.abc import Iterator, Mapping, Sequence
from contextlib import contextmanager from contextlib import contextmanager
from typing import Any, cast from typing import Any, cast
@@ -27,10 +27,9 @@ from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import ( from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_STAGE1_SQL,
SELECT_DELTA_STAGE2_SQL, SELECT_DELTA_STAGE2_SQL,
BasePostgresSaver, BasePostgresSaver,
_DeltaStage1Row, _build_delta_stage1_sql,
_DeltaStage2Row, _DeltaStage2Row,
) )
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -444,53 +443,79 @@ class PostgresSaver(BasePostgresSaver):
with conn.cursor(binary=True, row_factory=dict_row) as cur: with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur yield cur
def _get_channel_writes_history( def _get_all_delta_channels_writes_history(
self, config: RunnableConfig, channel: str self, config: RunnableConfig, channels: Sequence[str]
) -> _ChannelWritesHistory: ) -> Mapping[str, _ChannelWritesHistory]:
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`. """Fast-path override of `BaseCheckpointSaver._get_all_delta_channels_writes_history`.
Two-stage query: stage 1 scans checkpoint metadata to walk the parent Two-stage query, both stages cover ALL requested channels in a single
chain and locate the nearest snapshot; stage 2 fetches only the Postgres roundtrip each:
chain-limited writes and single seed blob.
* Stage 1: dynamic SELECT over `checkpoints` with K parallel JSONB
key lookups (one column pair per channel) — no subquery, no
aggregation. Returns one row per checkpoint with versions and
snapshot flags for every requested channel.
* Stage 2: one UNION ALL over `checkpoint_writes` and
`checkpoint_blobs` filtered by `channel = ANY(?)` and per-channel
chain_cids / seed_versions (collapsed across channels).
""" """
if not channels:
return {}
channels = list(channels)
thread_id = config["configurable"]["thread_id"] thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config) checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None: if checkpoint_id is None:
target = self.get_tuple(config) target = self.get_tuple(config)
if target is None: if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) return {
ch: _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
for ch in channels
}
checkpoint_id = target.config["configurable"]["checkpoint_id"] checkpoint_id = target.config["configurable"]["checkpoint_id"]
# Stage 1: K parallel JSONB lookups per row, one query for all channels.
stage1_sql = _build_delta_stage1_sql(channels)
stage1_params: list[Any] = []
for ch in channels:
stage1_params.extend([ch, ch])
stage1_params.extend([thread_id, checkpoint_ns])
with self._cursor() as cur: with self._cursor() as cur:
cur.execute( cur.execute(stage1_sql, stage1_params)
SELECT_DELTA_STAGE1_SQL,
(channel, channel, thread_id, checkpoint_ns),
)
stage1_rows = cur.fetchall() stage1_rows = cur.fetchall()
chain_cids, seed_version = self._walk_stage1(
cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id chain_by_ch, seed_ver_by_ch = self._walk_stage1_multi(
cast("list[Mapping[str, Any]]", stage1_rows), checkpoint_id, channels
) )
seed_versions = [seed_version] if seed_version else [] # Union of chain cids and seed versions across all channels.
union_chain_cids: list[str] = sorted(
{cid for chain in chain_by_ch.values() for cid in chain}
)
union_seed_versions: list[str] = sorted(
{ver for ver in seed_ver_by_ch.values() if ver is not None}
)
# Stage 2: chain-limited writes + chain-limited seed blobs for all channels.
with self._cursor() as cur: with self._cursor() as cur:
cur.execute( cur.execute(
SELECT_DELTA_STAGE2_SQL, SELECT_DELTA_STAGE2_SQL,
( (
thread_id, thread_id,
checkpoint_ns, checkpoint_ns,
channel, channels,
chain_cids, union_chain_cids,
thread_id, thread_id,
checkpoint_ns, checkpoint_ns,
channel, channels,
seed_versions, union_seed_versions,
), ),
) )
stage2_rows = cur.fetchall() stage2_rows = cur.fetchall()
return self._build_delta_channel_writes_history( return self._build_delta_channels_writes_history(
channel=channel, channels=channels,
chain_cids=chain_cids, chain_by_ch=chain_by_ch,
seed_version=seed_version, seed_ver_by_ch=seed_ver_by_ch,
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows), stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
) )
@@ -2,7 +2,7 @@ from __future__ import annotations
import asyncio import asyncio
from collections import defaultdict from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Sequence from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from typing import Any, cast from typing import Any, cast
@@ -27,10 +27,9 @@ from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import ( from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_STAGE1_SQL,
SELECT_DELTA_STAGE2_SQL, SELECT_DELTA_STAGE2_SQL,
BasePostgresSaver, BasePostgresSaver,
_DeltaStage1Row, _build_delta_stage1_sql,
_DeltaStage2Row, _DeltaStage2Row,
) )
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
@@ -405,53 +404,68 @@ class AsyncPostgresSaver(BasePostgresSaver):
async with conn.cursor(binary=True, row_factory=dict_row) as cur: async with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur yield cur
async def _aget_channel_writes_history( async def _aget_all_delta_channels_writes_history(
self, config: RunnableConfig, channel: str self, config: RunnableConfig, channels: Sequence[str]
) -> _ChannelWritesHistory: ) -> Mapping[str, _ChannelWritesHistory]:
"""Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`. """Fast-path override of `BaseCheckpointSaver._aget_all_delta_channels_writes_history`.
Two-stage query: stage 1 scans checkpoint metadata to walk the parent Two-stage query, both stages cover ALL requested channels in a single
chain and locate the nearest snapshot; stage 2 fetches only the Postgres roundtrip each. See `PostgresSaver._get_all_delta_channels_writes_history`
chain-limited writes and single seed blob. for design notes.
""" """
if not channels:
return {}
channels = list(channels)
thread_id = config["configurable"]["thread_id"] thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config) checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None: if checkpoint_id is None:
target = await self.aget_tuple(config) target = await self.aget_tuple(config)
if target is None: if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) return {
ch: _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
for ch in channels
}
checkpoint_id = target.config["configurable"]["checkpoint_id"] checkpoint_id = target.config["configurable"]["checkpoint_id"]
stage1_sql = _build_delta_stage1_sql(channels)
stage1_params: list[Any] = []
for ch in channels:
stage1_params.extend([ch, ch])
stage1_params.extend([thread_id, checkpoint_ns])
async with self._cursor() as cur: async with self._cursor() as cur:
await cur.execute( await cur.execute(stage1_sql, stage1_params)
SELECT_DELTA_STAGE1_SQL,
(channel, channel, thread_id, checkpoint_ns),
)
stage1_rows = await cur.fetchall() stage1_rows = await cur.fetchall()
chain_cids, seed_version = self._walk_stage1(
cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id chain_by_ch, seed_ver_by_ch = self._walk_stage1_multi(
cast("list[Mapping[str, Any]]", stage1_rows), checkpoint_id, channels
) )
seed_versions = [seed_version] if seed_version else [] union_chain_cids: list[str] = sorted(
{cid for chain in chain_by_ch.values() for cid in chain}
)
union_seed_versions: list[str] = sorted(
{ver for ver in seed_ver_by_ch.values() if ver is not None}
)
async with self._cursor() as cur: async with self._cursor() as cur:
await cur.execute( await cur.execute(
SELECT_DELTA_STAGE2_SQL, SELECT_DELTA_STAGE2_SQL,
( (
thread_id, thread_id,
checkpoint_ns, checkpoint_ns,
channel, channels,
chain_cids, union_chain_cids,
thread_id, thread_id,
checkpoint_ns, checkpoint_ns,
channel, channels,
seed_versions, union_seed_versions,
), ),
) )
stage2_rows = await cur.fetchall() stage2_rows = await cur.fetchall()
return self._build_delta_channel_writes_history( return self._build_delta_channels_writes_history(
channel=channel, channels=channels,
chain_cids=chain_cids, chain_by_ch=chain_by_ch,
seed_version=seed_version, seed_ver_by_ch=seed_ver_by_ch,
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows), stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
) )
@@ -2,7 +2,7 @@ from __future__ import annotations
import random import random
import warnings import warnings
from collections.abc import Sequence from collections.abc import Mapping, Sequence
from importlib.metadata import version as get_version from importlib.metadata import version as get_version
from typing import Any, TypedDict, cast from typing import Any, TypedDict, cast
@@ -161,6 +161,7 @@ class _DeltaStage2Row(TypedDict, total=False):
_kind: str # "w" or "b" _kind: str # "w" or "b"
checkpoint_id: str | None # "w" rows only checkpoint_id: str | None # "w" rows only
channel: str | None # set on both "w" and "b" rows
type: str | None type: str | None
blob: bytes | None blob: bytes | None
task_id: str | None # "w" rows only task_id: str | None # "w" rows only
@@ -168,48 +169,85 @@ class _DeltaStage2Row(TypedDict, total=False):
version: str | None # "b" rows only version: str | None # "b" rows only
# Two-stage DeltaChannel reconstruction. Stage 1 scans checkpoint # Multi-channel two-stage DeltaChannel reconstruction.
# metadata (no blob bytes) to walk the parent chain and locate the
# nearest snapshot marker. Stage 2 fetches only the chain-limited
# writes and the single seed snapshot blob.
# #
# Parameter order: # Stage 1 scans checkpoint metadata (no blob bytes) and emits one row per
# stage1: (channel, channel, thread_id, checkpoint_ns) # checkpoint with K parallel JSONB key lookups (one column pair per
# stage2: (thread_id, checkpoint_ns, channel, chain_cids[], # requested delta channel: ver_i / hs_i). No subqueries, no aggregation.
# thread_id, checkpoint_ns, channel, seed_versions[]) # Python walks the parent chain once across all channels.
#
# Stage 2 fetches all writes and the seed blobs for ALL channels in a
# single roundtrip via `channel = ANY(%s)` and chain/seed-version
# filtering.
#
# Empirical comparison vs an alternative "ship full channel_versions /
# channel_values JSONB and let Python pick" form (1000 checkpoints,
# 8 total channels in graph, 3 delta channels requested):
#
# Postgres execution: A=0.24ms vs B=0.38ms (both negligible)
# End-to-end latency: A=6.83ms vs B=2.28ms (B is 3.0x faster)
# Wire payload: A=836KB vs B=330KB (61% smaller)
# Buffer hits: identical (167 blocks)
#
# B (this dynamic-columns design) wins because it avoids JSONB
# serialization on the wire and JSONB-to-dict deserialization in
# psycopg. Even at K=8 (8 delta channels = 16 dynamic columns), B
# still beats A end-to-end (4.2ms vs 6.8ms).
def _build_delta_stage1_sql(channels: Sequence[str]) -> str:
"""Build stage 1 SQL with 2K parallel JSONB key lookups.
For channels=["messages", "files"] the result is::
SELECT checkpoint_id, parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver_0,
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_0,
checkpoint -> 'channel_versions' ->> %s AS ver_1,
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_1
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
Channel names are passed as `%s` parameters (safe from SQL injection).
Only the column aliases `ver_i` / `hs_i` are interpolated into the
SQL string (i is bounded by len(channels) and uses safe identifiers).
Caller must extend params with `[ch_0, ch_0, ch_1, ch_1, ...,
thread_id, ns]`.
"""
cols = []
for i in range(len(channels)):
cols.append(
f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}, "
f"(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_{i}"
)
return (
"SELECT checkpoint_id, parent_checkpoint_id, "
+ ", ".join(cols)
+ " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s"
)
SELECT_DELTA_STAGE1_SQL = """
SELECT checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver,
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS has_snapshot
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
"""
SELECT_DELTA_STAGE2_SQL = """ SELECT_DELTA_STAGE2_SQL = """
SELECT 'w'::text AS _kind, SELECT 'w'::text AS _kind,
checkpoint_id, checkpoint_id, channel,
type, blob, task_id, idx, NULL::text AS version type, blob, task_id, idx, NULL::text AS version
FROM checkpoint_writes FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s WHERE thread_id = %s AND checkpoint_ns = %s AND channel = ANY(%s)
AND checkpoint_id = ANY(%s) AND checkpoint_id = ANY(%s)
UNION ALL UNION ALL
SELECT 'b', NULL, SELECT 'b', NULL, channel,
type, blob, NULL, NULL, version type, blob, NULL, NULL, version
FROM checkpoint_blobs FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s WHERE thread_id = %s AND checkpoint_ns = %s AND channel = ANY(%s)
AND version = ANY(%s) AND version = ANY(%s)
""" """
class _DeltaStage1Row(TypedDict): # Stage 1 rows are dynamic-shape dicts: {checkpoint_id, parent_checkpoint_id,
"""One row from `SELECT_DELTA_STAGE1_SQL`.""" # ver_0, hs_0, ver_1, hs_1, ...}. Walking is parameterized by the channel
# list to map indices back to channel names — no static TypedDict here.
checkpoint_id: str # `dict[str, Any]` is the practical signature.
parent_checkpoint_id: str | None
ver: str | None
has_snapshot: bool
class BasePostgresSaver(BaseCheckpointSaver[str]): class BasePostgresSaver(BaseCheckpointSaver[str]):
@@ -255,86 +293,119 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
} }
@staticmethod @staticmethod
def _walk_stage1( def _walk_stage1_multi(
stage1_rows: Sequence[_DeltaStage1Row], stage1_rows: Sequence[Mapping[str, Any]],
target_id: str, target_id: str,
) -> tuple[list[str], str | None]: channels: Sequence[str],
"""Walk the parent chain from stage 1 metadata rows. ) -> tuple[dict[str, list[str]], dict[str, str | None]]:
"""Walk the parent chain once for all requested channels.
Returns (chain_cids, seed_version): Each row carries `ver_i` / `hs_i` per channel index. We walk the
chain_cids: ancestor checkpoint IDs from target's parent down to parent chain from target's parent toward the root; for each
the seed (or root), in newest-first order. channel we stop at the nearest ancestor where `hs_i` is true and
seed_version: the channel blob version at the nearest ancestor record that ancestor's `ver_i` as the seed version. All
with has_snapshot=True, or None if pure delta. ancestors visited up to (and including) a channel's seed are in
that channel's `chain_cids`.
Returns:
chain_cids_by_channel: per-channel list of ancestor cids in
newest-first order.
seed_version_by_channel: per-channel seed version (None if
walk reached root with no snapshot).
""" """
parent_of: dict[str, str | None] = {} parent_of: dict[str, str | None] = {}
ver_of: dict[str, str | None] = {} # For each channel index, store ver and has_snapshot per cid.
snapshot_of: dict[str, bool] = {} ver_by_i_by_cid: list[dict[str, str | None]] = [
{} for _ in range(len(channels))
]
hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in range(len(channels))]
for r in stage1_rows: for r in stage1_rows:
cid = r["checkpoint_id"] cid = cast(str, r["checkpoint_id"])
parent_of[cid] = r["parent_checkpoint_id"] parent_of[cid] = cast("str | None", r["parent_checkpoint_id"])
ver_of[cid] = r["ver"] for i in range(len(channels)):
snapshot_of[cid] = r["has_snapshot"] ver_by_i_by_cid[i][cid] = cast("str | None", r.get(f"ver_{i}"))
hs_by_i_by_cid[i][cid] = bool(r.get(f"hs_{i}"))
chain_cids: list[str] = [] chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
seed_version: str | None = None seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
cur_cid: str | None = parent_of.get(target_id) # For each channel, walk from target's parent until we hit a
while cur_cid is not None: # snapshot or the root. Walks share the parent_of mapping but
chain_cids.append(cur_cid) # are otherwise independent.
if snapshot_of.get(cur_cid, False): for i, ch in enumerate(channels):
seed_version = ver_of.get(cur_cid) cur_cid: str | None = parent_of.get(target_id)
break while cur_cid is not None:
cur_cid = parent_of.get(cur_cid) chain_by_ch[ch].append(cur_cid)
return chain_cids, seed_version if hs_by_i_by_cid[i].get(cur_cid, False):
seed_ver_by_ch[ch] = ver_by_i_by_cid[i].get(cur_cid)
break
cur_cid = parent_of.get(cur_cid)
return chain_by_ch, seed_ver_by_ch
def _build_delta_channel_writes_history( def _build_delta_channels_writes_history(
self, self,
*, *,
channel: str, channels: Sequence[str],
chain_cids: list[str], chain_by_ch: dict[str, list[str]],
seed_version: str | None, seed_ver_by_ch: dict[str, str | None],
stage2_rows: Sequence[_DeltaStage2Row], stage2_rows: Sequence[_DeltaStage2Row],
) -> _ChannelWritesHistory: ) -> dict[str, _ChannelWritesHistory]:
"""Reconstruct delta channel history from two-stage query results. """Demux stage 2 rows per channel; produce per-channel histories.
chain_cids are in newest-first order (target's parent first). stage2_rows carry `channel` on every row. We build per-channel
stage2_rows contain only writes for chain_cids and the single `writes_by_cid` and per-channel `seed_blob` dicts, then assemble
seed blob at seed_version. a `_ChannelWritesHistory` per requested channel.
""" """
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {} # writes_by_ch_by_cid[channel][cid] = list of (type, blob, task_id, idx)
seed_blob: tuple[str, bytes] | None = None writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
ch: {} for ch in channels
}
# seed_blob_by_ver[(channel, version)] = (type, blob)
seed_blob_by_ver: dict[tuple[str, str], tuple[str, bytes]] = {}
for r in stage2_rows: for r in stage2_rows:
ch = cast(str, r["channel"])
kind = r["_kind"] kind = r["_kind"]
if kind == "w": if kind == "w":
cid = cast(str, r["checkpoint_id"]) cid = cast(str, r["checkpoint_id"])
writes_by_cid.setdefault(cid, []).append( writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append(
cast( cast(
"tuple[str, bytes, str, int]", "tuple[str, bytes, str, int]",
(r["type"], r["blob"], r["task_id"], r["idx"]), (r["type"], r["blob"], r["task_id"], r["idx"]),
) )
) )
else: # kind == "b" else: # kind == "b"
seed_blob = cast("tuple[str, bytes]", (r["type"], r["blob"])) ver = cast(str, r["version"])
seed_blob_by_ver[(ch, ver)] = cast(
"tuple[str, bytes]", (r["type"], r["blob"])
)
for ws in writes_by_cid.values(): # Sort writes per (channel, cid) newest-first by (task_id, idx)
ws.sort(key=lambda w: (w[2], w[3]), reverse=True) for cid_map in writes_by_ch_by_cid.values():
for ws in cid_map.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
if not chain_cids: result: dict[str, _ChannelWritesHistory] = {}
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) for ch in channels:
chain_cids = chain_by_ch.get(ch, [])
seed_version = seed_ver_by_ch.get(ch)
collected: list[PendingWrite] = [] collected: list[PendingWrite] = []
for cid in chain_cids: cid_writes = writes_by_ch_by_cid.get(ch, {})
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []): for cid in chain_cids:
val = self.serde.loads_typed((type_tag, write_blob)) for type_tag, write_blob, task_id, _idx in cid_writes.get(cid, []):
collected.append((task_id, channel, val)) val = self.serde.loads_typed((type_tag, write_blob))
collected.append((task_id, ch, val))
seed: Any = DELTA_SENTINEL seed: Any = DELTA_SENTINEL
if seed_blob is not None and seed_blob[0] != "empty": if seed_version is not None:
seed = self.serde.loads_typed(seed_blob) blob = seed_blob_by_ver.get((ch, seed_version))
if blob is not None and blob[0] != "empty":
seed = self.serde.loads_typed(blob)
collected.reverse() collected.reverse()
return _ChannelWritesHistory(seed=seed, writes=collected) result[ch] = _ChannelWritesHistory(seed=seed, writes=collected)
return result
def _dump_blobs( def _dump_blobs(
self, self,
@@ -125,7 +125,8 @@ class CheckpointTuple(NamedTuple):
class _ChannelWritesHistory(NamedTuple): class _ChannelWritesHistory(NamedTuple):
"""Result of `BaseCheckpointSaver._get_channel_writes_history`. """Result of `BaseCheckpointSaver._get_all_delta_channels_writes_history`
(a per-channel entry from the returned mapping).
Storage-level view of what one channel wrote across the ancestor chain Storage-level view of what one channel wrote across the ancestor chain
of a target checkpoint: of a target checkpoint:
@@ -487,12 +488,12 @@ class BaseCheckpointSaver(Generic[V]):
raise NotImplementedError raise NotImplementedError
def _get_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None: def _get_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None:
"""Pure storage read used by `_get_channel_writes_history`. """Pure storage read used by `_get_all_delta_channels_writes_history`.
Must return the same value as `get_tuple` but must NOT trigger channel Must return the same value as `get_tuple` but must NOT trigger channel
reconstruction; otherwise the channel-hydration path would re-enter reconstruction; otherwise the channel-hydration path would re-enter
`_get_channel_writes_history`. Override only if `get_tuple` itself `_get_all_delta_channels_writes_history`. Override only if `get_tuple`
performs channel hydration. itself performs channel hydration.
""" """
return self.get_tuple(config) return self.get_tuple(config)
@@ -500,14 +501,17 @@ class BaseCheckpointSaver(Generic[V]):
"""Async version of `_get_tuple_raw`. See docstring there.""" """Async version of `_get_tuple_raw`. See docstring there."""
return await self.aget_tuple(config) return await self.aget_tuple(config)
def _get_channel_writes_history( def _get_all_delta_channels_writes_history(
self, config: RunnableConfig, channel: str self, config: RunnableConfig, channels: Sequence[str]
) -> _ChannelWritesHistory: ) -> Mapping[str, _ChannelWritesHistory]:
"""**Experimental.** Query one channel's writes along the parent chain. """**Experimental.** Query multiple delta channels' writes along the parent chain.
Storage-level query, not channel semantics: returns `(seed, writes)` Storage-level query, not channel semantics: returns a per-channel
reflecting what storage knows about a single channel across the `(seed, writes)` reflecting what storage knows about each channel
ancestor chain of the target checkpoint identified by `config`. across the ancestor chain of the target checkpoint identified by
`config`.
For every channel in `channels`:
* `writes` — on-path deltas oldest→newest as `PendingWrite` tuples. * `writes` — on-path deltas oldest→newest as `PendingWrite` tuples.
Writes stored at the target `checkpoint_id` itself are pending Writes stored at the target `checkpoint_id` itself are pending
@@ -520,69 +524,86 @@ class BaseCheckpointSaver(Generic[V]):
Walks the **parent chain** (not `list(before=...)`): for forked Walks the **parent chain** (not `list(before=...)`): for forked
threads, only on-path ancestors contribute. threads, only on-path ancestors contribute.
Reference implementation walks `get_tuple` + `parent_config`, Reference implementation walks `get_tuple` + `parent_config` ONCE
inspecting each ancestor's `channel_values[channel]` for the seed for all channels (each ancestor visited once, not once per channel),
terminator. Savers with direct storage access (`InMemorySaver`, inspecting each ancestor's `channel_values[channel]` for that
`PostgresSaver`) override for performance; the return contract is channel's seed terminator. Savers with direct storage access
fixed here. (`InMemorySaver`, `PostgresSaver`) override for performance; the
return contract is fixed here.
Underscore-prefixed because the method surface is experimental. Empty `channels` returns `{}`. Underscore-prefixed because the
method surface is experimental.
""" """
collected: list[PendingWrite] = [] # newest first; reversed at the end if not channels:
return {}
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels}
seed_by_ch: dict[str, Any] = {c: DELTA_SENTINEL for c in channels}
remaining: set[str] = set(channels)
target_tuple = self._get_tuple_raw(config) target_tuple = self._get_tuple_raw(config)
cursor_config: RunnableConfig | None = ( cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None target_tuple.parent_config if target_tuple else None
) )
while cursor_config is not None: while cursor_config is not None and remaining:
tup = self._get_tuple_raw(cursor_config) tup = self._get_tuple_raw(cursor_config)
if tup is None: if tup is None:
break break
# Collect this ancestor's writes FIRST — they encode the # Collect each ancestor's writes for any channel still searching.
# transition from this ancestor's state to its child's, so
# they must be included whether or not this ancestor is the
# seed terminator.
if tup.pending_writes: if tup.pending_writes:
# Within a superstep, pending_writes are oldest→newest;
# reverse to scan newest-first.
for write in reversed(tup.pending_writes): for write in reversed(tup.pending_writes):
if write[1] != channel: ch = write[1]
continue if ch in remaining:
collected.append(write) collected_by_ch[ch].append(write)
# Seed terminator: any non-sentinel blob on an ancestor # Per-channel seed terminator: a non-sentinel blob value at this
# establishes the reconstruction base. Stop here. # ancestor establishes that channel's reconstruction base.
ancestor_value = tup.checkpoint["channel_values"].get(channel) for ch in list(remaining):
if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: ancestor_value = tup.checkpoint["channel_values"].get(ch)
collected.reverse() if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL:
return _ChannelWritesHistory(seed=ancestor_value, writes=collected) seed_by_ch[ch] = ancestor_value
remaining.discard(ch)
cursor_config = tup.parent_config cursor_config = tup.parent_config
collected.reverse() return {
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) ch: _ChannelWritesHistory(
seed=seed_by_ch[ch],
writes=list(reversed(collected_by_ch[ch])),
)
for ch in channels
}
async def _aget_channel_writes_history( async def _aget_all_delta_channels_writes_history(
self, config: RunnableConfig, channel: str self, config: RunnableConfig, channels: Sequence[str]
) -> _ChannelWritesHistory: ) -> Mapping[str, _ChannelWritesHistory]:
"""Async version of `_get_channel_writes_history`. See docstring there.""" """Async version of `_get_all_delta_channels_writes_history`."""
collected: list[PendingWrite] = [] if not channels:
return {}
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels}
seed_by_ch: dict[str, Any] = {c: DELTA_SENTINEL for c in channels}
remaining: set[str] = set(channels)
target_tuple = await self._aget_tuple_raw(config) target_tuple = await self._aget_tuple_raw(config)
cursor_config: RunnableConfig | None = ( cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None target_tuple.parent_config if target_tuple else None
) )
while cursor_config is not None: while cursor_config is not None and remaining:
tup = await self._aget_tuple_raw(cursor_config) tup = await self._aget_tuple_raw(cursor_config)
if tup is None: if tup is None:
break break
if tup.pending_writes: if tup.pending_writes:
for write in reversed(tup.pending_writes): for write in reversed(tup.pending_writes):
if write[1] != channel: ch = write[1]
continue if ch in remaining:
collected.append(write) collected_by_ch[ch].append(write)
ancestor_value = tup.checkpoint["channel_values"].get(channel) for ch in list(remaining):
if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: ancestor_value = tup.checkpoint["channel_values"].get(ch)
collected.reverse() if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL:
return _ChannelWritesHistory(seed=ancestor_value, writes=collected) seed_by_ch[ch] = ancestor_value
remaining.discard(ch)
cursor_config = tup.parent_config cursor_config = tup.parent_config
collected.reverse() return {
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) ch: _ChannelWritesHistory(
seed=seed_by_ch[ch],
writes=list(reversed(collected_by_ch[ch])),
)
for ch in channels
}
def get_next_version(self, current: V | None, channel: None) -> V: def get_next_version(self, current: V | None, channel: None) -> V:
"""Generate the next version ID for a channel. """Generate the next version ID for a channel.
@@ -6,7 +6,7 @@ import pickle
import random import random
import shutil import shutil
from collections import defaultdict from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Sequence from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
from types import TracebackType from types import TracebackType
from typing import Any from typing import Any
@@ -141,17 +141,24 @@ class InMemorySaver(
result[k] = self.serde.loads_typed(vv) result[k] = self.serde.loads_typed(vv)
return result return result
def _get_channel_writes_history( def _get_all_delta_channels_writes_history(
self, config: RunnableConfig, channel: str self, config: RunnableConfig, channels: Sequence[str]
) -> _ChannelWritesHistory: ) -> Mapping[str, _ChannelWritesHistory]:
"""Override: walk the parent chain ONCE for all requested channels.
For each channel we track its own seed terminator independently.
On a snapshot or pre-delta ancestor for a given channel, that
channel stops collecting further writes; other channels keep
walking until they find their own terminator or hit the root.
"""
if not channels:
return {}
thread_id = config["configurable"]["thread_id"] thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id", "") checkpoint_id = config["configurable"].get("checkpoint_id", "")
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {}) ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
# Walk the parent chain newest→oldest. Skip the target itself —
# writes stored AT `checkpoint_id` are pending for the next step # Build the parent chain (newest→oldest), skipping the target.
# (pregel applies them via `apply_writes`; they aren't part of the
# snapshot value AT `checkpoint_id`).
chain: list[str] = [] chain: list[str] = []
target_entry = ns_storage.get(checkpoint_id) target_entry = ns_storage.get(checkpoint_id)
current: str | None = target_entry[2] if target_entry is not None else None current: str | None = target_entry[2] if target_entry is not None else None
@@ -162,77 +169,73 @@ class InMemorySaver(
chain.append(current) chain.append(current)
_, _, parent = entry _, _, parent = entry
current = parent current = parent
# Scan newest→oldest. A pre-delta blob on an ancestor terminates the
# walk and is bound as `seed`; without this, a thread migrated from
# pre-delta storage would replay ancestor writes all the way to the
# root AND miss any value that lived only in the old blob (e.g. from
# `update_state`).
#
# At each ancestor, check the blob BEFORE processing its pending
# writes: a pre-delta blob represents the state AT that ancestor,
# which already subsumes any writes stored under it. Processing
# those writes first would fold them into the reconstructed value
# twice (once via the blob, once via replay).
collected: list[PendingWrite] = [] # newest first
for cp_id in chain: # newest → oldest
entry = ns_storage.get(cp_id)
if entry is not None:
ckpt = self.serde.loads_typed(entry[0])
ver = ckpt.get("channel_versions", {}).get(channel)
if ver is not None:
blob_entry = self.blobs.get(
(thread_id, checkpoint_ns, channel, ver)
)
if blob_entry is not None and blob_entry[0] != "empty":
blob_value = self.serde.loads_typed(blob_entry)
if blob_value is not DELTA_SENTINEL:
if isinstance(blob_value, _DeltaSnapshot):
# Step-based snapshot: the blob is state AT this
# ancestor, but the ancestor's pending_writes
# encode the NEXT step's transition and are NOT
# subsumed by the snapshot — collect them first.
step_writes = self.writes.get(
(thread_id, checkpoint_ns, cp_id), {}
)
for (_task_id, _idx), (
tid,
ch,
serialized,
_,
) in sorted(step_writes.items(), reverse=True):
if ch != channel:
continue
collected.append(
(tid, ch, self.serde.loads_typed(serialized))
)
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
)
# Pre-delta blob: state AT this ancestor already
# subsumes its pending_writes — skip them.
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
)
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels}
seed_by_ch: dict[str, Any] = {c: DELTA_SENTINEL for c in channels}
remaining: set[str] = set(channels)
for cp_id in chain: # newest → oldest
if not remaining:
break
entry = ns_storage.get(cp_id)
ckpt = self.serde.loads_typed(entry[0]) if entry is not None else None
# Per-channel: check seed terminator at this ancestor first.
terminated_here: set[str] = set()
blob_value_by_ch: dict[str, Any] = {}
if ckpt is not None:
versions = ckpt.get("channel_versions", {})
for ch in remaining:
ver = versions.get(ch)
if ver is None:
continue
blob_entry = self.blobs.get((thread_id, checkpoint_ns, ch, ver))
if blob_entry is None or blob_entry[0] == "empty":
continue
blob_value = self.serde.loads_typed(blob_entry)
if blob_value is DELTA_SENTINEL:
continue
blob_value_by_ch[ch] = blob_value
terminated_here.add(ch)
# Process step writes: filter by channel, collect newest-first.
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {}) step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
# Within a superstep, sorted by (task_id, idx) = oldest → newest;
# reverse for newest-first scan.
for (_task_id, _idx), (tid, ch, serialized, _) in sorted( for (_task_id, _idx), (tid, ch, serialized, _) in sorted(
step_writes.items(), reverse=True step_writes.items(), reverse=True
): ):
if ch != channel: if ch not in remaining:
continue continue
val = self.serde.loads_typed(serialized) blob_value = blob_value_by_ch.get(ch)
collected.append((tid, ch, val)) if blob_value is not None and not isinstance(
collected.reverse() blob_value, _DeltaSnapshot
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) ):
# Pre-delta blob terminator: state at this ancestor
# already subsumes these writes — skip them.
continue
# Either no terminator at this ancestor for this channel,
# OR a `_DeltaSnapshot` terminator (writes here encode the
# transition to the child and are NOT subsumed).
collected_by_ch[ch].append(
(tid, ch, self.serde.loads_typed(serialized))
)
async def _aget_channel_writes_history( # Now apply terminators: channels that found a seed are done.
self, config: RunnableConfig, channel: str for ch in terminated_here:
) -> _ChannelWritesHistory: seed_by_ch[ch] = blob_value_by_ch[ch]
return self._get_channel_writes_history(config, channel) remaining.discard(ch)
return {
ch: _ChannelWritesHistory(
seed=seed_by_ch[ch],
writes=list(reversed(collected_by_ch[ch])),
)
for ch in channels
}
async def _aget_all_delta_channels_writes_history(
self, config: RunnableConfig, channels: Sequence[str]
) -> Mapping[str, _ChannelWritesHistory]:
return self._get_all_delta_channels_writes_history(config, channels)
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
"""Get a checkpoint tuple from the in-memory storage. """Get a checkpoint tuple from the in-memory storage.
@@ -36,7 +36,7 @@ class _DeltaSnapshot(NamedTuple):
"""Snapshot blob for a DeltaChannel with finite snapshot_frequency. """Snapshot blob for a DeltaChannel with finite snapshot_frequency.
Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code. Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code.
The ancestor walk in `_get_channel_writes_history` terminates when it The ancestor walk in `_get_all_delta_channels_writes_history` terminates when it
encounters this type (any non-sentinel blob stops the walk). encounters this type (any non-sentinel blob stops the walk).
`from_checkpoint` reconstructs the channel value directly from `.value` `from_checkpoint` reconstructs the channel value directly from `.value`
+36 -22
View File
@@ -335,9 +335,10 @@ class TestInMemorySaverDeltaChannel:
assert channel not in result assert channel not in result
def test_get_channel_writes_collects_ancestor_writes_only(self) -> None: def test_get_channel_writes_collects_ancestor_writes_only(self) -> None:
"""_get_channel_writes_history collects ancestor writes oldest→newest, """_get_all_delta_channels_writes_history collects ancestor writes
and excludes writes stored at the target checkpoint itself (those are oldest→newest, and excludes writes stored at the target checkpoint
pending writes for the next step, applied separately by pregel).""" itself (those are pending writes for the next step, applied separately
by pregel)."""
saver = InMemorySaver() saver = InMemorySaver()
serde = JsonPlusSerializer() serde = JsonPlusSerializer()
@@ -375,7 +376,9 @@ class TestInMemorySaverDeltaChannel:
"checkpoint_id": "cp2", "checkpoint_id": "cp2",
} }
} }
result = saver._get_channel_writes_history(config, channel) result = saver._get_all_delta_channels_writes_history(config, [channel])[
channel
]
assert result.seed is DELTA_SENTINEL assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes] values = [v for _, _, v in result.writes]
assert values == [{"content": "hi"}] assert values == [{"content": "hi"}]
@@ -405,15 +408,17 @@ class TestInMemorySaverDeltaChannel:
"checkpoint_id": "cp1", "checkpoint_id": "cp1",
} }
} }
result = saver._get_channel_writes_history(config, channel) result = saver._get_all_delta_channels_writes_history(config, [channel])[
channel
]
assert result.seed is DELTA_SENTINEL assert result.seed is DELTA_SENTINEL
assert result.writes == [] assert result.writes == []
class TestBaseFallbackGetChannelWrites: class TestBaseFallbackGetChannelWrites:
"""Exercises the `BaseCheckpointSaver._get_channel_writes_history` default """Exercises the `BaseCheckpointSaver._get_all_delta_channels_writes_history`
implementation — the path third-party savers inherit when they don't default implementation — the path third-party savers inherit when they
override `_get_channel_writes_history` themselves. don't override `_get_all_delta_channels_writes_history` themselves.
Regression guard for a bug where the fallback passed the caller's config Regression guard for a bug where the fallback passed the caller's config
(with `checkpoint_id`) straight to `self.list()`, which most savers (with `checkpoint_id`) straight to `self.list()`, which most savers
@@ -429,11 +434,11 @@ class TestBaseFallbackGetChannelWrites:
""" """
class _ThirdPartyStyleSaver(InMemorySaver): class _ThirdPartyStyleSaver(InMemorySaver):
_get_channel_writes_history = ( _get_all_delta_channels_writes_history = (
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined] InMemorySaver.__mro__[1]._get_all_delta_channels_writes_history # type: ignore[attr-defined]
) )
_aget_channel_writes_history = ( _aget_all_delta_channels_writes_history = (
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] InMemorySaver.__mro__[1]._aget_all_delta_channels_writes_history # type: ignore[attr-defined]
) )
saver = _ThirdPartyStyleSaver() saver = _ThirdPartyStyleSaver()
@@ -477,7 +482,9 @@ class TestBaseFallbackGetChannelWrites:
} }
} }
result = saver._get_channel_writes_history(config, "messages") result = saver._get_all_delta_channels_writes_history(config, ["messages"])[
"messages"
]
assert result.seed is DELTA_SENTINEL assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes] values = [v for _, _, v in result.writes]
@@ -494,7 +501,9 @@ class TestBaseFallbackGetChannelWrites:
} }
} }
result = await saver._aget_channel_writes_history(config, "messages") result = (
await saver._aget_all_delta_channels_writes_history(config, ["messages"])
)["messages"]
assert result.seed is DELTA_SENTINEL assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes] values = [v for _, _, v in result.writes]
@@ -503,9 +512,9 @@ class TestBaseFallbackGetChannelWrites:
async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None: async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None:
"""Regression: the re-entrancy guard must be task-local, not thread-local. """Regression: the re-entrancy guard must be task-local, not thread-local.
Two concurrent `_aget_channel_writes_history` calls on the same Two concurrent `_aget_all_delta_channels_writes_history` calls on the
event-loop thread must each see their full reconstructed writes. A same event-loop thread must each see their full reconstructed writes.
`threading.local()` guard would let whichever task set it first A `threading.local()` guard would let whichever task set it first
short-circuit the other to `writes=[]`. short-circuit the other to `writes=[]`.
""" """
import asyncio import asyncio
@@ -533,12 +542,13 @@ class TestBaseFallbackGetChannelWrites:
} }
results = await asyncio.gather( results = await asyncio.gather(
saver._aget_channel_writes_history(config, "messages"), saver._aget_all_delta_channels_writes_history(config, ["messages"]),
saver._aget_channel_writes_history(config, "messages"), saver._aget_all_delta_channels_writes_history(config, ["messages"]),
) )
expected_values = [{"content": "first"}, {"content": "second"}] expected_values = [{"content": "first"}, {"content": "second"}]
for result in results: for result_map in results:
result = result_map["messages"]
assert result.seed is DELTA_SENTINEL assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes] values = [v for _, _, v in result.writes]
assert values == expected_values assert values == expected_values
@@ -625,7 +635,9 @@ class TestPreDeltaBlobTerminator:
} }
} }
result = saver._get_channel_writes_history(config, channel) result = saver._get_all_delta_channels_writes_history(config, [channel])[
channel
]
# Seed came from the pre-delta blob at cp1. # Seed came from the pre-delta blob at cp1.
assert result.seed == ["A"] assert result.seed == ["A"]
@@ -647,7 +659,9 @@ class TestPreDeltaBlobTerminator:
} }
} }
result = saver._get_channel_writes_history(config, channel) result = saver._get_all_delta_channels_writes_history(config, [channel])[
channel
]
values = [v for _, _, v in result.writes] values = [v for _, _, v in result.writes]
# The pre-delta write under cp1 must not appear (the blob subsumes it). # The pre-delta write under cp1 must not appear (the blob subsumes it).
+32 -11
View File
@@ -114,9 +114,12 @@ def channels_from_checkpoint(
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])` For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
is sufficient. `DeltaChannel` is the exception: sentinel blobs require an is sufficient. `DeltaChannel` is the exception: sentinel blobs require an
ancestor walk via `saver._get_channel_writes_history`. The walk terminates ancestor walk via `saver._get_all_delta_channels_writes_history`. All
at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain delta channels needing replay are batched into a single saver call to
value, so read depth is bounded by `snapshot_frequency`. save K-1 redundant scans of `checkpoint_writes` (which has no channel
index). The walk terminates per-channel at the nearest `_DeltaSnapshot`
blob or pre-migration plain value, so read depth is bounded by
`snapshot_frequency`.
""" """
channel_specs: dict[str, BaseChannel] = {} channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {} managed_specs: dict[str, ManagedValueSpec] = {}
@@ -126,18 +129,26 @@ def channels_from_checkpoint(
else: else:
managed_specs[k] = v managed_specs[k] = v
delta_channels: list[str] = [
k
for k, spec in channel_specs.items()
if _needs_replay(spec, checkpoint["channel_values"].get(k, MISSING))
]
histories: Mapping[str, Any] = {}
if delta_channels and saver is not None and config is not None:
histories = saver._get_all_delta_channels_writes_history(config, delta_channels)
channels: dict[str, BaseChannel] = {} channels: dict[str, BaseChannel] = {}
for k, spec in channel_specs.items(): for k, spec in channel_specs.items():
ch: BaseChannel ch: BaseChannel
stored = checkpoint["channel_values"].get(k, MISSING) if k in histories:
if _needs_replay(spec, stored) and saver is not None and config is not None:
delta_spec = cast(DeltaChannel, spec) delta_spec = cast(DeltaChannel, spec)
history = saver._get_channel_writes_history(config, k) history = histories[k]
replay_ch = delta_spec.from_checkpoint(history.seed) replay_ch = delta_spec.from_checkpoint(history.seed)
replay_ch.replay_writes(history.writes) replay_ch.replay_writes(history.writes)
ch = replay_ch ch = replay_ch
else: else:
ch = spec.from_checkpoint(stored) ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
channels[k] = ch channels[k] = ch
return channels, managed_specs return channels, managed_specs
@@ -158,18 +169,28 @@ async def achannels_from_checkpoint(
else: else:
managed_specs[k] = v managed_specs[k] = v
delta_channels: list[str] = [
k
for k, spec in channel_specs.items()
if _needs_replay(spec, checkpoint["channel_values"].get(k, MISSING))
]
histories: Mapping[str, Any] = {}
if delta_channels and saver is not None and config is not None:
histories = await saver._aget_all_delta_channels_writes_history(
config, delta_channels
)
channels: dict[str, BaseChannel] = {} channels: dict[str, BaseChannel] = {}
for k, spec in channel_specs.items(): for k, spec in channel_specs.items():
ch: BaseChannel ch: BaseChannel
stored = checkpoint["channel_values"].get(k, MISSING) if k in histories:
if _needs_replay(spec, stored) and saver is not None and config is not None:
delta_spec = cast(DeltaChannel, spec) delta_spec = cast(DeltaChannel, spec)
history = await saver._aget_channel_writes_history(config, k) history = histories[k]
replay_ch = delta_spec.from_checkpoint(history.seed) replay_ch = delta_spec.from_checkpoint(history.seed)
replay_ch.replay_writes(history.writes) replay_ch.replay_writes(history.writes)
ch = replay_ch ch = replay_ch
else: else:
ch = spec.from_checkpoint(stored) ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
channels[k] = ch channels[k] = ch
return channels, managed_specs return channels, managed_specs
@@ -6,7 +6,7 @@ checkpointer — pre-migration state visible at each *settled* ancestor
checkpoint is preserved, and post-migration writes fold on top through checkpoint is preserved, and post-migration writes fold on top through
the reducer. the reducer.
Mechanism under test: the saver's `_get_channel_writes_history(config, Mechanism under test: the saver's `_get_all_delta_channels_writes_history(config,
channel)` walks the parent chain; when it encounters an ancestor whose channel)` walks the parent chain; when it encounters an ancestor whose
`channel_values[channel]` is a real value (not `DELTA_SENTINEL`), it `channel_values[channel]` is a real value (not `DELTA_SENTINEL`), it
returns that as the `seed`. `DeltaChannel.from_checkpoint(seed)` uses returns that as the `seed`. `DeltaChannel.from_checkpoint(seed)` uses
@@ -29,7 +29,7 @@ Scenarios covered:
pre-migration seed. pre-migration seed.
4. **Base-saver fallback path**: a third-party-style subclass that 4. **Base-saver fallback path**: a third-party-style subclass that
removes the optimized `InMemorySaver` override and falls back to removes the optimized `InMemorySaver` override and falls back to
`BaseCheckpointSaver._get_channel_writes_history` must produce the `BaseCheckpointSaver._get_all_delta_channels_writes_history` must produce the
same result as the optimized path. same result as the optimized path.
5. **Channel-type isolation across threads**: two threads on the same 5. **Channel-type isolation across threads**: two threads on the same
checkpointer under the delta-channel graph — one freshly-started, checkpointer under the delta-channel graph — one freshly-started,
@@ -266,7 +266,7 @@ def test_continuing_migrated_thread_folds_deltas_on_seed() -> None:
class _ThirdPartyStyleSaver(InMemorySaver): class _ThirdPartyStyleSaver(InMemorySaver):
"""Simulates a third-party saver that inherits the reference """Simulates a third-party saver that inherits the reference
`_get_channel_writes_history` implementation from `_get_all_delta_channels_writes_history` implementation from
`BaseCheckpointSaver` rather than overriding it. `BaseCheckpointSaver` rather than overriding it.
We rebind the two methods to the base-class versions (via MRO) so We rebind the two methods to the base-class versions (via MRO) so
@@ -275,11 +275,11 @@ class _ThirdPartyStyleSaver(InMemorySaver):
""" """
# MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...] # MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...]
_get_channel_writes_history = ( # type: ignore[assignment] _get_all_delta_channels_writes_history = ( # type: ignore[assignment]
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined] InMemorySaver.__mro__[1]._get_all_delta_channels_writes_history # type: ignore[attr-defined]
) )
_aget_channel_writes_history = ( # type: ignore[assignment] _aget_all_delta_channels_writes_history = ( # type: ignore[assignment]
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] InMemorySaver.__mro__[1]._aget_all_delta_channels_writes_history # type: ignore[attr-defined]
) )
@@ -326,7 +326,7 @@ def test_delta_and_migrated_threads_do_not_cross_contaminate() -> None:
"""Two threads sharing a checkpointer — one migrated from """Two threads sharing a checkpointer — one migrated from
pre-migration state, one freshly-started under DeltaChannel — must pre-migration state, one freshly-started under DeltaChannel — must
maintain independent state. The parent-chain walk in maintain independent state. The parent-chain walk in
`_get_channel_writes_history` must be scoped to the target thread. `_get_all_delta_channels_writes_history` must be scoped to the target thread.
""" """
checkpointer = InMemorySaver() checkpointer = InMemorySaver()