lib: When copying checkpoint, make it a child of the parent

2026-08-30 19:59:40 +02:00 · 2024-11-15 16:56:17 -08:00
parent 5494855ffa
commit 66b9a7dee7
1 changed files with 2 additions and 2 deletions
@@ -961,7 +961,7 @@ class Pregel(PregelProtocol):
                next_checkpoint = create_checkpoint(checkpoint, None, step)
                # copy checkpoint
                next_config = checkpointer.put(
-                    checkpoint_config,
+                    saved.parent_config or saved.config if saved else checkpoint_config,
                    next_checkpoint,
                    {
                        **checkpoint_metadata,
@@ -1222,7 +1222,7 @@ class Pregel(PregelProtocol):
                next_checkpoint = create_checkpoint(checkpoint, None, step)
                # copy checkpoint
                next_config = await checkpointer.aput(
-                    checkpoint_config,
+                    saved.parent_config or saved.config if saved else checkpoint_config,
                    next_checkpoint,
                    {
                        **checkpoint_metadata,