Avoid storing Sends twice in memory and postgres checkpointers

- Sends are stored through put_writes, so we don't need to also store them inside checkpoint object
- On reading checkpoint, reconstruct pending_sends from the stored writes
This commit is contained in:
Nuno Campos
2024-08-27 11:10:18 -07:00
parent 238b562e78
commit 9b6100bc4f
5 changed files with 102 additions and 44 deletions
@@ -64,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
pipeline (bool): whether to use AsyncPipeline
Returns:
PostgresSaver: A new PostgresSaver instance.
AsyncPostgresSaver: A new AsyncPostgresSaver instance.
"""
async with await AsyncConnection.connect(
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
@@ -137,12 +137,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
"checkpoint_id": value["checkpoint_id"],
}
},
{
**self._load_checkpoint(value["checkpoint"]),
"channel_values": await asyncio.to_thread(
self._load_blobs, value["channel_values"]
),
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
{
"configurable": {
@@ -196,12 +196,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
"checkpoint_id": value["checkpoint_id"],
}
},
{
**self._load_checkpoint(value["checkpoint"]),
"channel_values": await asyncio.to_thread(
self._load_blobs, value["channel_values"]
),
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
{
"configurable": {