Skip to content

Drop dead iterator refs from the stats registry - #1429

Open
Shr3yash wants to merge 1 commit into
google:mainfrom
Shr3yash:fix/iter-weakref-registry
Open

Shr3yash wants to merge 1 commit into
google:mainfrom
Shr3yash:fix/iter-weakref-registry

Conversation

@Shr3yash

@Shr3yash Shr3yash commented Oct 12, 2026 •

Copy link
Copy Markdown

stats._iter_weakref_registry keeps a HashableWeakRef for every output iterator, and those refs hash and compare by the referent's id(). Nothing removes a ref when the iterator is collected, so dead entries pile up. CPython reuses those addresses, set.add treats the new iterator as the stale entry, and the live iterator is never registered.

HashableWeakRef now discards its registry entry when the referent is collected, and a dead ref no longer compares equal to a live iterator at the same address.

The new test creates and drops iterators in a loop. It fails on current main and passes with this change:

python3 grain/_src/python/dataset/stats_test.py

Fixes #1427


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

This branch has not been deployed

No deployments
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.

stats._iter_weakref_registry keeps refs of dead iterators, and a new iterator at a reused address is not registered

1 participant