Closes#3704
## Problem
Transforms can fail on bad data (e.g. nulls/NaNs from incomplete user
surveys). Today any transform exception aborts iteration, and there is
no way to skip invalid rows during loading.
## Solution
New `on_transform_error` parameter on `StreamingDataset`:
- `"raise"` (default, matches current behavior and the convention in
tf.data / WebDataset / Ray Data)
- `"skip"` — drop the failing rows and continue
- `"warn"` — like skip, plus a logged warning per failing batch
- a WebDataset-style callable `handler(exc) -> bool`, so users can skip
only expected error types
Key design points:
- **Row-granular skipping**: when a batch fails, the transform is re-run
on single-row slices so only the rows that actually fail are dropped
(avoids Ray-style whole-block loss). Skips are counted in a new
`rows_skipped` property.
- **No crash on uneven skips**: the round-robin loop now ends the epoch
at the last cycle where every split still has a row, instead of hitting
`IndexError` when a split runs dry early.
- **Exact resumability under skips**: checkpoints are now
position-based. `state_dict` gains `positions_consumed_per_split` (exact
for owned splits), and a new `merge_state_dicts` static method combines
per-rank states via elementwise max for elastic resume across topology
changes. Old checkpoints without the new key still load. Positions equal
sample counts when nothing is skipped, so existing behavior is
unchanged.
- **Guardrail**: transforms returning the wrong number of rows now raise
a clear `ValueError` instead of silently corrupting split accounting.
### Answers to the issue's open questions
- *Can we do this?* Yes — all transforms funnel through one guarded call
in the Stage 2 pipeline.
- *What do other libraries do?* tf.data `ignore_errors()`, WebDataset
`handler=`, Ray `max_errored_blocks`; MosaicML StreamingDataset offers
nothing (skipping conflicts with its determinism model). This design
follows the common conventions: raise by default, opt-in skipping,
count/log drops.
- *Error handling or pre-filtering?* Both: the existing `filter=`
remains the recommended tool for predictable bad data (splits are built
post-filter, so all guarantees hold — now documented);
`on_transform_error` covers failures not expressible as a predicate.
- *Impact on splits / elastic determinism?* Per-split sample sequences
stay deterministic (skips are data-dependent, not topology-dependent).
With unequal bad-row counts across splits the last few global steps of
an epoch can differ across topologies (bounded by the skew), which is
documented on the parameter. With equal counts per split, full
determinism is preserved — covered by a test.
## Testing
15 new tests in `test_elastic_dataloader.py` covering: default raise,
invalid values, uniform and uneven skips (including epoch-end
truncation), warn logging, selective callable handlers, wrong-row-count
guardrail, determinism across runs and across world sizes (1/2/3/4) with
skips, exact mid-epoch resume with skips on the same topology, elastic
resume via `merge_state_dicts` (ws=2 → ws=1), merge validation, and
backward-compat loading of old checkpoints.
Note: relying on CI for the test run — my local machine OOMs during the
final link of the native extension. The change itself is pure Python.
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Xuanwo
Sravan Avvaru