Skip to content

Start interleave children concurrently after snapshot or restore - #1416

Closed
vshmyhlo wants to merge 1 commit into
google:mainfrom
vshmyhlo:fix-interleave-prefetch-startup
Closed

vshmyhlo wants to merge 1 commit into
google:mainfrom
vshmyhlo:fix-interleave-prefetch-startup

Conversation

@vshmyhlo

@vshmyhlo vshmyhlo commented Oct 7, 2026 •

Copy link
Copy Markdown

An initial get_state() call, including the one made by a ThreadPrefetchDatasetIterator wrapper, materializes interleave children with prefetch disabled. set_state() also leaves restored children stopped. The next interleave read then waits for its selected child before starting the other children, serializing multiprocessing worker startup across the first round of batches.

Call the existing, idempotent start_prefetch() method before waiting for the selected child. This starts all active lanes while keeping round-robin consumption and checkpoint state unchanged. Repeating the call also covers children stopped or replaced by restore.

Add eight regression cases covering initial snapshots, zero-buffer wrappers, restore into an existing iterator, and restore into a fresh iterator, with debug mode on and off. Each lane waits until all lanes have started, then the test checks the exact output sequence. The snapshot case fails against unpatched 0.2.18 and all eight pass with the fix. Fourteen downstream placement/replay cases and a real three-process snapshot/restore check also pass.


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

@vshmyhlo vshmyhlo closed this by deleting the head repository Oct 7, 2026
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.

1 participant