Fix Loader restoration before None items

[aswanth-07]

Fixes #1551.

Summary

LoaderIterator used None both as a valid data value and as the empty look-ahead marker. When a checkpoint resumed immediately before a None item, Loader.__iter__() called has_next(), misclassified the item as end-of-input, and reset the pipeline to the beginning. This duplicated all data consumed before the checkpoint.

This change tracks cache occupancy separately from the cached value. None can therefore be cached and returned like any other item, while the existing look-ahead and restart behavior remains unchanged for exhausted iterators.

Reset now also clears the cached prefetch state dictionary together with the cached item. This prevents state_dict() from reporting a position captured before a later load_state_dict() and iterator reset.

The regressions verify that restoring [0, None, 2] after consuming 0 yields exactly [None, 2], and that a state request after has_next(), load_state_dict(), and iter(loader) reports the restored position.

Validation

• PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 py -3.13 -m pytest -q test/nodes/test_loader.py (5 passed)
• py -3.13 -m pre_commit run --files torchdata/nodes/loader.py test/nodes/test_loader.py (passed)
• Direct public-flow probe: the stale-state regression failed before the reset fix with positions 0 and 1, then passed after the fix.
• git diff --check (passed)
• Standalone mypy was not used as a final signal because the installed version rejects the repository's quoted python_version setting and this source checkout lacks the generated torchdata.version module.

AI assistance disclosure: AI-assisted development tools were used during investigation and implementation. The reported tests were run against the final diff.

[aswanth-07]

Hi @divyanshk, when you have time, could you please review this checkpoint restoration fix? It separates look-ahead occupancy from the cached value so None remains valid data. The focused loader tests, pre-commit, mypy, and Meta import checks pass.

[sylvesterkaczmarek]

One adjacent cache survives reset: _cached_state_dict isn't cleared. If has_next() populated it, then load_state_dict() + iter(loader) resets the root and clears _cached_item/_has_cached_item, but a subsequent state_dict() still returns the pre-reset cached state. Should reset() clear _cached_state_dict as well?

[aswanth-07]

@sylvesterkaczmarek Good catch. _cached_state_dict is the checkpoint paired with the prefetched item, so it should be cleared by the same reset. I added that clear and a regression covering has_next(), load_state_dict(), iterator reset, and the next state_dict() in b88467a. The focused loader file passes all 5 tests, and the changed-file pre-commit hooks pass. Thanks.

[sylvesterkaczmarek]

The reset path now clears _cached_state_dict together with the cached item/occupancy flag, and the new regression covers has_next() followed by restore/reset and a subsequent state_dict(). This resolves the stale-checkpoint issue I raised.

[aswanth-07]

Thanks for the careful follow-up and for catching the stale checkpoint case.

[sylvesterkaczmarek]

Thanks, glad it helped. The reset and restore boundary was easy to miss, and the added regression should keep it from coming back.