Skip to content

Python: FileCheckpointStorage save/load symmetry (#8181) - #8193

Closed
lsmlhi_25 (FOWEPJF255) wants to merge 1 commit into
microsoft:mainfrom
FOWEPJF255:fix/file-checkpoint-save-load-symmetry-8181
Closed

Python: FileCheckpointStorage save/load symmetry (#8181)#8193
lsmlhi_25 (FOWEPJF255) wants to merge 1 commit into
microsoft:mainfrom
FOWEPJF255:fix/file-checkpoint-save-load-symmetry-8181

Conversation

@FOWEPJF255

Copy link
Copy Markdown
Contributor

Summary

  • Validate encoded checkpoints against this storage''s allowed_checkpoint_types before writing, so save refuses payloads that load would reject.
  • Wrap invalid JSON on load as WorkflowCheckpointException.
  • Align list_checkpoint_ids with list_checkpoints decode filtering.

Closes #8181

Test plan

  • test_file_storage_rejects_unlisted_user_type_at_save
  • test_file_checkpoint_storage_load_invalid_json_raises
  • test_file_checkpoint_storage_list_ids_matches_list_decode_filter
  • Existing allow-list / block-list file storage tests

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Encoding and invalid-UTF-8 failures can still escape the documented exception contract.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Improves FileCheckpointStorage consistency between saving, loading, and listing.

Changes:

  • Validates checkpoint restorability before saving.
  • Wraps malformed JSON load failures.
  • Aligns checkpoint ID filtering with checkpoint listing.
File summaries
File Description
_checkpoint.py Updates save, load, and listing behavior.
test_checkpoint.py Tests invalid JSON and listing consistency.
test_checkpoint_unrestricted_pickle.py Tests save-time allow-list enforcement.
Review details
  • Files reviewed: 3/3 changed files
  • Comments generated: 3
  • Review effort level: Balanced

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

Comment on lines 331 to +342
checkpoint_dict = checkpoint.to_dict()
encoded_checkpoint = encode_checkpoint_value(checkpoint_dict)
# Fail at save time if this storage could not restore the payload (#8181).
try:
decode_checkpoint_value(encoded_checkpoint, allowed_types=self._allowed_types)
except WorkflowCheckpointException:
raise
except Exception as ex:
raise WorkflowCheckpointException(
f"Checkpoint {checkpoint.checkpoint_id} cannot be restored under this "
"storage's allowed types; refusing to save."
) from ex
encoded_checkpoint = await asyncio.to_thread(_read)
try:
encoded_checkpoint = await asyncio.to_thread(_read)
except json.JSONDecodeError as ex:
Comment on lines +480 to +484
decoded_checkpoint_dict = decode_checkpoint_value(
encoded_checkpoint, allowed_types=self._allowed_types
)
checkpoint = WorkflowCheckpoint.from_dict(decoded_checkpoint_dict)
if checkpoint.workflow_name == workflow_name:
@moonbox3

Copy link
Copy Markdown
Contributor

Closing as PR doesn't use our required pull request template.

@FOWEPJF255

Copy link
Copy Markdown
Contributor Author

Reopened with the required PR template as #8214 (also addresses prior Copilot notes on encode/UTF-8/list_ids).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

python Usage: [Issues, PRs], Target: Python

Projects

None yet

4 participants