For input checkpoint use filtered input writes in checkpoint metadata

- if we use the raw input we can end up saving keys that were actually ignored (because not one of the input keys), which is incosistent
This commit is contained in:
Nuno Campos
2024-08-30 13:24:41 -07:00
committed by Nuno Campos
parent c761b80970
commit 8154ab784e
+1 -1
View File
@@ -422,7 +422,7 @@ class PregelLoop:
self.checkpointer_get_next_version,
), "Can't write to SharedValues in graph input"
# save input checkpoint
self._put_checkpoint({"source": "input", "writes": self.input})
self._put_checkpoint({"source": "input", "writes": dict(input_writes)})
else:
raise EmptyInputError(f"Received no input for {input_keys}")
# done with input