Fix lineage of maual state updates

2026-09-28 04:25:08 +02:00 · 2024-05-06 11:53:20 -07:00
parent 48df2d1084
commit 251bd9744d
1 changed files with 2 additions and 2 deletions
@@ -482,7 +482,7 @@ class Pregel(
            # apply to checkpoint and save
            _apply_writes(checkpoint, channels, task.writes)
            return self.checkpointer.put(
-                config,
+                saved.config if saved else config,
                create_checkpoint(checkpoint, channels),
                {
                    "source": "update",
@@ -551,7 +551,7 @@ class Pregel(
            # apply to checkpoint and save
            _apply_writes(checkpoint, channels, task.writes)
            return await self.checkpointer.aput(
-                config,
+                saved.config if saved else config,
                create_checkpoint(checkpoint, channels),
                {
                    "source": "update",