Use uuid6 as the id for checkpoints

- this avoids conflicts if multiple processes creating checkpoints in same thread at same time
- uuid6 with a monotically increasing clock_seq is sortable by creation time at ms precision (plus clock_seq for ties, plus 48 bits of randomness for further ties)
This commit is contained in:
Nuno Campos
2024-05-14 16:56:25 -07:00
parent 77db7bf825
commit 52b2c755d7
14 changed files with 71 additions and 85 deletions
+3 -4
View File
@@ -13,6 +13,7 @@ from typing import (
)
from typing_extensions import Self
from uuid6 import uuid6
from langgraph.checkpoint.base import Checkpoint
from langgraph.errors import EmptyChannelError, InvalidUpdateError
@@ -110,13 +111,10 @@ async def AsyncChannelsManager(
def create_checkpoint(
checkpoint: Checkpoint, channels: Mapping[str, BaseChannel]
checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], step: int
) -> Checkpoint:
"""Create a checkpoint for the given channels."""
ts = datetime.now(timezone.utc).isoformat()
assert (
ts > checkpoint["ts"]
), f"Timestamps must be monotonically increasing, got {ts} <= {checkpoint['ts']}"
values: dict[str, Any] = {}
for k, v in channels.items():
try:
@@ -126,6 +124,7 @@ def create_checkpoint(
return Checkpoint(
v=1,
ts=ts,
id=str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
+2 -2
View File
@@ -319,7 +319,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)",
(
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
checkpoint["id"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
@@ -329,6 +329,6 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["ts"],
"thread_ts": checkpoint["id"],
}
}
+6
View File
@@ -12,6 +12,7 @@ from typing import (
)
from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
from uuid6 import uuid6
from langgraph.serde.base import SerializerProtocol
from langgraph.serde.jsonplus import JsonPlusSerializer
@@ -48,6 +49,9 @@ class Checkpoint(TypedDict):
v: int
"""The version of the checkpoint format. Currently 1."""
id: str
"""The ID of the checkpoint. This is both unique and monotonically
increasing, so can be used for sorting checkpoints from first to last."""
ts: str
"""The timestamp of the checkpoint in ISO 8601 format."""
channel_values: dict[str, Any]
@@ -77,6 +81,7 @@ def _seen_dict():
def empty_checkpoint() -> Checkpoint:
return Checkpoint(
v=1,
id=str(uuid6(clock_seq=-2)),
ts=datetime.now(timezone.utc).isoformat(),
channel_values={},
channel_versions=defaultdict(int),
@@ -88,6 +93,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
return Checkpoint(
v=checkpoint["v"],
ts=checkpoint["ts"],
id=checkpoint["id"],
channel_values=checkpoint["channel_values"].copy(),
channel_versions=defaultdict(int, checkpoint["channel_versions"]),
versions_seen=defaultdict(
+2 -2
View File
@@ -192,7 +192,7 @@ class MemorySaver(BaseCheckpointSaver):
"""
self.storage[config["configurable"]["thread_id"]].update(
{
checkpoint["ts"]: (
checkpoint["id"]: (
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
)
@@ -201,7 +201,7 @@ class MemorySaver(BaseCheckpointSaver):
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["ts"],
"thread_ts": checkpoint["id"],
}
}
+2 -2
View File
@@ -426,7 +426,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)",
(
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
checkpoint["id"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
@@ -435,7 +435,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["ts"],
"thread_ts": checkpoint["id"],
}
}
+15 -13
View File
@@ -538,12 +538,13 @@ class Pregel(
)
# apply to checkpoint and save
_apply_writes(checkpoint, channels, task.writes)
step = saved.metadata.get("step", -2) + 1 if saved else -1
return self.checkpointer.put(
saved.config if saved else config,
create_checkpoint(checkpoint, channels),
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": saved.metadata.get("step", 0) + 1 if saved else 0,
"step": step,
"writes": {as_node: values},
},
)
@@ -612,12 +613,13 @@ class Pregel(
)
# apply to checkpoint and save
_apply_writes(checkpoint, channels, task.writes)
step = saved.metadata.get("step", -2) + 1 if saved else -1
return await self.checkpointer.aput(
saved.config if saved else config,
create_checkpoint(checkpoint, channels),
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": saved.metadata.get("step", 0) + 1 if saved else 0,
"step": step,
"writes": {as_node: values},
},
)
@@ -741,7 +743,7 @@ class Pregel(
_apply_writes(checkpoint, channels, input_writes)
# save input checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint = create_checkpoint(checkpoint, channels, start)
bg.append(
executor.submit(
self.checkpointer.put,
@@ -753,7 +755,7 @@ class Pregel(
checkpoint_config = {
"configurable": {
**checkpoint_config["configurable"],
"thread_ts": checkpoint["ts"],
"thread_ts": checkpoint["id"],
}
}
# increment start to 0
@@ -874,7 +876,7 @@ class Pregel(
# save end of step checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint = create_checkpoint(checkpoint, channels, step)
bg.append(
executor.submit(
self.checkpointer.put,
@@ -898,7 +900,7 @@ class Pregel(
checkpoint_config = {
"configurable": {
**checkpoint_config["configurable"],
"thread_ts": checkpoint["ts"],
"thread_ts": checkpoint["id"],
}
}
# yield debug checkpoint
@@ -1031,7 +1033,7 @@ class Pregel(
_apply_writes(checkpoint, channels, input_writes)
# save input checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint = create_checkpoint(checkpoint, channels, start)
bg.append(
asyncio.create_task(
self.checkpointer.aput(
@@ -1044,7 +1046,7 @@ class Pregel(
checkpoint_config = {
"configurable": {
**checkpoint_config["configurable"],
"thread_ts": checkpoint["ts"],
"thread_ts": checkpoint["id"],
}
}
# increment start to 0
@@ -1175,7 +1177,7 @@ class Pregel(
# save end of step checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint = create_checkpoint(checkpoint, channels, step)
bg.append(
asyncio.create_task(
self.checkpointer.aput(
@@ -1200,7 +1202,7 @@ class Pregel(
checkpoint_config = {
"configurable": {
**checkpoint_config["configurable"],
"thread_ts": checkpoint["ts"],
"thread_ts": checkpoint["id"],
}
}
# yield debug checkpoint
@@ -1414,7 +1416,7 @@ def _local_read(
fresh: bool = False,
) -> Union[dict[str, Any], Any]:
if fresh:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint = create_checkpoint(checkpoint, channels, -1)
with ChannelsManager(channels, checkpoint) as channels:
_apply_writes(copy_checkpoint(checkpoint), channels, writes)
return read_channels(channels, select)
Generated
+12 -1
View File
@@ -3868,6 +3868,17 @@ h2 = ["h2 (>=4,<5)"]
socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"]
zstd = ["zstandard (>=0.18.0)"]
[[package]]
name = "uuid6"
version = "2024.1.12"
description = "New time-based UUID formats which are suited for use as a database key"
optional = false
python-versions = ">=3.8"
files = [
{file = "uuid6-2024.1.12-py3-none-any.whl", hash = "sha256:8150093c8d05a331bc0535bc5ef6cf57ac6eceb2404fd319bc10caee2e02c065"},
{file = "uuid6-2024.1.12.tar.gz", hash = "sha256:ed0afb3a973057575f9883201baefe402787ca5e11e1d24e377190f0c43f1993"},
]
[[package]]
name = "watchdog"
version = "4.0.0"
@@ -4094,4 +4105,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "d42a7ee4e0079a588710997ce02f3717820706d93922990537d0e4ab78a78beb"
content-hash = "a2a48cce5d31d0c3e09a9b3b042586c569c3efbe881de7a9c4ce9ae2b4d63499"
+1
View File
@@ -10,6 +10,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
[tool.poetry.dependencies]
python = ">=3.9.0,<4.0"
langchain-core = "^0.1.52"
uuid6 = "^2024.1.12"
[tool.poetry.group.test.dependencies]
+4 -15
View File
@@ -1,8 +1,9 @@
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.channels.base import create_checkpoint
from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, empty_checkpoint
class TestAsyncSqliteSaver:
@@ -18,20 +19,8 @@ class TestAsyncSqliteSaver:
"configurable": {"thread_id": "thread-2", "thread_ts": "2"}
}
self.chkpnt_1: Checkpoint = {
"v": 1,
"ts": "1",
"channel_values": {},
"channel_versions": {},
"versions_seen": {},
}
self.chkpnt_2: Checkpoint = {
"v": 2,
"ts": "2",
"channel_values": {},
"channel_versions": {},
"versions_seen": {},
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.metadata_1: CheckpointMetadata = {
"source": "input",
+4 -15
View File
@@ -1,7 +1,8 @@
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata
from langgraph.channels.base import create_checkpoint
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, empty_checkpoint
from langgraph.checkpoint.memory import MemorySaver
@@ -18,20 +19,8 @@ class TestMemorySaver:
"configurable": {"thread_id": "thread-2", "thread_ts": "2"}
}
self.chkpnt_1: Checkpoint = {
"v": 1,
"ts": "1",
"channel_values": {},
"channel_versions": {},
"versions_seen": {},
}
self.chkpnt_2: Checkpoint = {
"v": 2,
"ts": "2",
"channel_values": {},
"channel_versions": {},
"versions_seen": {},
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.metadata_1: CheckpointMetadata = {
"source": "input",
+4 -15
View File
@@ -1,7 +1,8 @@
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata
from langgraph.channels.base import create_checkpoint
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, empty_checkpoint
from langgraph.checkpoint.sqlite import SqliteSaver, _metadata_predicate, search_where
@@ -18,20 +19,8 @@ class TestSqliteSaver:
"configurable": {"thread_id": "thread-2", "thread_ts": "2"}
}
self.chkpnt_1: Checkpoint = {
"v": 1,
"ts": "1",
"channel_values": {},
"channel_versions": {},
"versions_seen": {},
}
self.chkpnt_2: Checkpoint = {
"v": 2,
"ts": "2",
"channel_values": {},
"channel_versions": {},
"versions_seen": {},
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.metadata_1: CheckpointMetadata = {
"source": "input",
+2 -2
View File
@@ -40,8 +40,8 @@ class MemorySaverAssertImmutable(MemorySaver):
# assert checkpoint hasn't been modified since last written
thread_id = config["configurable"]["thread_id"]
if saved := super().get(config):
assert self.storage_for_copies[thread_id][saved["ts"]] == saved
self.storage_for_copies[thread_id][checkpoint["ts"]] = copy_checkpoint(
assert self.storage_for_copies[thread_id][saved["id"]] == saved
self.storage_for_copies[thread_id][checkpoint["id"]] = copy_checkpoint(
checkpoint
)
# call super to write checkpoint
+7 -7
View File
@@ -818,13 +818,13 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
assert state is not None
assert state.values.get("total") == 2
assert state.next == ()
assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"]
assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["id"]
# total is now 2, so output is 2+3=5
assert app.invoke(3, thread_1) == 5
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"]
assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["id"]
# total is now 2+5=7, so output would be 7+4=11, but raises ValueError
with pytest.raises(ValueError):
app.invoke(4, thread_1)
@@ -879,11 +879,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
assert thread_1_history[-2].values["total"] == 2
# can get each checkpoint using aget with config
assert (
memory.get(thread_1_history[0].config)["ts"]
memory.get(thread_1_history[0].config)["id"]
== thread_1_history[0].config["configurable"]["thread_ts"]
)
assert (
memory.get(thread_1_history[1].config)["ts"]
memory.get(thread_1_history[1].config)["id"]
== thread_1_history[1].config["configurable"]["thread_ts"]
)
@@ -4629,7 +4629,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
config=uconfig,
metadata={
"source": "update",
"step": 0,
"step": -1,
"writes": {START: {"my_key": "key", "market": "DE"}},
},
)
@@ -4645,7 +4645,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
config=tool_two.checkpointer.get_tuple(thread3).config,
metadata={
"source": "loop",
"step": 1,
"step": 0,
"writes": {"prepare": {"my_key": " prepared"}},
},
parent_config=uconfig,
@@ -4661,7 +4661,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
config=tool_two.checkpointer.get_tuple(thread3).config,
metadata={
"source": "loop",
"step": 3,
"step": 2,
"writes": {"finish": {"my_key": " finished"}},
},
parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config,
+7 -7
View File
@@ -801,7 +801,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
assert state.values.get("total") == 2
assert (
state.config["configurable"]["thread_ts"]
== (await memory.aget(thread_1))["ts"]
== (await memory.aget(thread_1))["id"]
)
# total is now 2, so output is 2+3=5
assert await app.ainvoke(3, thread_1) == 5
@@ -810,7 +810,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
assert state.values.get("total") == 7
assert (
state.config["configurable"]["thread_ts"]
== (await memory.aget(thread_1))["ts"]
== (await memory.aget(thread_1))["id"]
)
# total is now 2+5=7, so output would be 7+4=11, but raises ValueError
with pytest.raises(ValueError):
@@ -869,10 +869,10 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
assert thread_1_history[-2].values["total"] == 2
# can get each checkpoint using aget with config
assert (await memory.aget(thread_1_history[0].config))[
"ts"
"id"
] == thread_1_history[0].config["configurable"]["thread_ts"]
assert (await memory.aget(thread_1_history[1].config))[
"ts"
"id"
] == thread_1_history[1].config["configurable"]["thread_ts"]
thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10)
@@ -3907,7 +3907,7 @@ async def test_branch_then() -> None:
config=uconfig,
metadata={
"source": "update",
"step": 0,
"step": -1,
"writes": {START: {"my_key": "key", "market": "DE"}},
},
)
@@ -3923,7 +3923,7 @@ async def test_branch_then() -> None:
config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
metadata={
"source": "loop",
"step": 1,
"step": 0,
"writes": {"prepare": {"my_key": " prepared"}},
},
parent_config=uconfig,
@@ -3939,7 +3939,7 @@ async def test_branch_then() -> None:
config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
metadata={
"source": "loop",
"step": 3,
"step": 2,
"writes": {"finish": {"my_key": " finished"}},
},
parent_config=[