Add single controller cache for grain. - #4870
Conversation
…rrayRecord datasets Pre-resolves file globbing and extracts ArrayRecord FileInstruction headers once on the coordinator node. Passes pre-extracted FileInstructions directly to worker sidecars via RemoteIteratorWrapper, bypassing dynamic GCS globbing and file header reads on hundreds/thousands of worker hosts. This avoids GCS API contention, prevents watchdog step-timebomb hangs at scale, and speeds up initialization across worker hosts by over 100x.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces FileInstruction and extract_file_instructions to pre-resolve file paths and extract metadata on the coordinator, bypassing remote index discovery for Grain's ArrayRecordDataSource. It also updates the dataset creation logic to handle these pre-resolved instructions for both training and evaluation iterators, and adds a unit test suite. Feedback focuses on: 1) ensuring robust globbing when a list or tuple of patterns is passed to find_data_files; 2) reducing code duplication by unifying the mixture parsing logic in get_datasets; 3) parallelizing the sequential extraction of file instructions for both training and evaluation mixtures to avoid startup latency; and 4) importing FileInstruction and extract_file_instructions directly from the production module in the unit tests to avoid duplication and ensure the production code is properly tested.
sadikneipp
changed the title
- Recursively glob list/tuple of patterns in find_data_files. - Deduplicate mixture branch logic in get_datasets. - Parallelize extract_file_instructions across patterns using ThreadPoolExecutor for both train and eval iterators. - Import FileInstruction and extract_file_instructions from production module in unit tests.
…/without FileInstruction
296fec0
into
abhinavsing/snapshot-scale
1 participant
Cold storage would be accessed for metadata by each and every host at the iterator init before this change.
This change defines a way to init that metadata in the single controller and reuse it at the hosts. The core of the change is a new branch for
create_dataset_from_patternusing the FileInstruction dataclass.Validated for 50 step in v5e-32 capacity.