mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-25 19:15:11 +02:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bc263888b1 | ||
|
|
8c2a30af8a |
@@ -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`
|
||||||
|
|||||||
@@ -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).
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user