Implement a Weighted Data Batcher with Deterministic Save/Resume (DataRegistry Iterator Interface)

Problem: Implement a Weighted Data Batcher with Deterministic Save/Resume You are given an implemented DataRegistry (provided in the interview Colab) that can return an iterator for a dataset by name. Implement a DataBatcher (or equivalent functions/classes) that mixes samples from multiple datasets according to integer weights and yields fixed-size batches. DataRegistry interface (conceptual) registry.get_iterator(dataset_name) returns an iterator it. Each next(it) yields one sample (treat the sample as an opaque object). Part 1: Weighted mixing into batches (divisibility guaranteed) Implement an iterator/generator that: Inputs: datasets: list of dataset names weights: positive integer weights aligned with datasets batch_size: number of samples per batch Output: an iterator/generator yielding batches (e.g., list/array of samples) of length exactly batch_size Rules Mix datasets according to weights (e.g., A:2, B:1 means A should appear about twice as often as B in the long run). Guarantee: batch_size is divisible by sum(weights) (or an equivalent condition that makes per-batch allocation exact). Define and implement a concrete per-batch allocation strategy (e.g., a fixed repeating pattern). Example tests datasets=[A,B], weights=[1,1], batch_size=2 => each batch: 1 from A, 1 from B. datasets=[A,B], weights=[2,1], batch_size=3 => each batch: 2 from A, 1 from B. datasets=[A,B,C], weights=[1,1,2], batch_size=4 => each batch: 1A, 1B, 2C. Single dataset. Larger weights/batch sizes to verify correctness. Part 2: Add offset + deterministic save/resume (divisibility still guaranteed) Extend Part 1 with: New parameter: offset Behavior: start yielding batches from the offset-th sample position in the global mixed sample stream (i.e., skip the first offset mixed samples). Implement deterministic checkpointing: state = batcher.save_state() returns a serializable state object batcher.load_state(state) restores the batcher Requirements After resume, the output must be identical to uninterrupted execution. offset semantics must match save/resume: starting from offset=k should match “run k samples, then continue”. Example tests Run some batches, save; resume and compare to continuous run. offset=k equals skipping k mixed samples. Verify determinism across weight configurations. Part 3: batch_size not divisible by sum(weights) Extend Part 2 further: No longer assume batch_size % sum(weights) == 0. Still must: yield batches of length batch_size match weights in the long run (as closely as possible) support offset and deterministic save/resume You need a deterministic scheduling strategy to allocate the "remainder" across batches in a way that can be reproduced exactly after restore. Example tests weights=[2,1], batch_size=4 => per-batch ratio may vary; across batches should approach 2:1. weights=[1,1,2], batch_size=3. Mix offset with save/resume and verify identical output. Highly skewed weights like [100,1]. Edge cases like batch_size=1. Follow-up (engineering) In a real training environment, what additional corner cases or engineering concerns would you handle? (e.g., dataset exhaustion vs infinite streams, distributed worker consistency, RNG seeding, checkpoint compatibility, performance/memory, skew, prefetch/backpressure, etc.) Example Input datasets=["A","B"], weights=[1,1], batch_size=2, offset=0 A iterator: a1,a2,a3,... B iterator: b1,b2,b3,... Output batch1 length=2, contains 1 sample from A and 1 sample from B (order deterministic) batch2 length=2, contains 1 from A and 1 from B