Skip to content

Restore ConcatThenSplitIterDataset from a JSON-serialized state - #1434

Open
AshishKumar4 wants to merge 1 commit into
google:mainfrom
AshishKumar4:fix/concat-then-split-json-state
Open

AshishKumar4 wants to merge 1 commit into
google:mainfrom
AshishKumar4:fix/concat-then-split-json-state

Conversation

@AshishKumar4

@AshishKumar4 AshishKumar4 commented Oct 19, 2026 •

Copy link
Copy Markdown

Fixes #1433.

_CtsState marks a missing (or whole, unsliced) remainder with the tuple _EMPTY_SLICE = (-1, -1). A state read back from JSON holds the slices as lists, which never equal that tuple, so set_state took the next parent element as a remainder sliced [-1:-1] and dropped it. _CtsState.__post_init__ now holds the slices as tuples.

test_checkpointing gains through_json, which restores each checkpoint from json.loads(json.dumps(state)). With the fix, the module's tests pass (100 passed). Without it, the four through_json=True cases whose checkpoints include step 0 fail.


📚 Documentation preview 📚: https://google-grain--1434.org.readthedocs.build/

A state read back from JSON holds the remainder slices as lists, which never
equal the tuple sentinel _EMPTY_SLICE, so set_state took the next parent
element as a sliced remainder and dropped it. _CtsState now holds the slices
as tuples. test_checkpointing also restores through JSON.

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ConcatThenSplitIterDataset drops a document when restored from a JSON-serialized state

1 participant