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

Start in grain/_src/python/dataset/stats.py: the module-level _iter_weakref_registry set, HashableWeakRef (hash and equality on id()), and the add and discard calls in Stats.__init__. The issue's repro script runs the two sections that show the problem and the control. Done means dead refs leave the registry once their iterator is collected, and a new iterator that reuses an address is registered. Check the existing stats_test.py still passes. The choice between a weakref callback and a WeakSet is open, so confirm with maintainers which one fits.

Versions

• grain 0.2.18 (PyPI), numpy 2.5.3, CPython 3.12, Linux x86_64, CPU only
• The code involved is unchanged on main at 7e36b72 (2026-10-09)

What happens

grain/_src/python/dataset/stats.py keeps a module-level _iter_weakref_registry = set() of HashableWeakRef. HashableWeakRef.__hash__ returns id(obj) and __eq__ compares those ids. Stats.__init__ adds the new iterator's ref and discards each parent's ref. Nothing removes a ref when its iterator is collected.

Two problems follow.

1. Dead refs accumulate. Every output iterator that dies leaves its ref in the set. How fast the set grows depends on address reuse (see 2), but nothing bounds it. In the repro, 100 iterators that are deleted and collected leave 100 dead refs behind.
2. A new iterator can be left out of the registry. CPython reuses freed memory, so a new iterator often gets the address of one that died. set.add(new_ref) then finds the dead ref with the same hash, __eq__ says they are equal, and the set keeps the dead ref. The live iterator is not in the registry. A later discard of an equal ref removes the dead one.

Point 2 is the common case for a loop that re-creates its iterator every epoch. In the repro's 200-pass loop, 175 to 191 iterators per run (8 runs) landed on an earlier pass's address, and none of them was registered. As far as I can tell nothing in the open-source tree reads the registry (it is only written in stats.py, and cleared in stats_test.py). The comment above it says it exists "for collecting execution stats", so whatever reads it (on-demand profiling, I assume) will miss those iterators and see dead entries instead.

Reproduction

Grain and NumPy only. Reads are single-threaded so the run is easy to follow.

"""grain: the iterator stats registry keeps refs of dead iterators and misses new ones.

`grain._src.python.dataset.stats._iter_weakref_registry` is a module-level set of
`HashableWeakRef`, whose hash and equality are the referent's `id()`. `Stats.__init__` adds the
new iterator's ref and discards each parent's ref. Nothing removes a ref when its iterator dies.

1. Dead refs stay in the registry.
2. Re-creating an iterator in a loop (one per epoch): CPython reuses the freed memory, the new
   output iterator often gets a dead iterator's address, `set.add` finds the dead ref with the
   same hash, `__eq__` says equal, the set keeps the dead ref, and the new iterator is not
   registered.
3. Control: the same steps with a registry that drops a ref when its iterator dies.

Reads are single-threaded (`num_threads=0`). The counts in 2 and 3 depend on which addresses the
allocator hands out and change from run to run; that none of the iterators at an earlier pass's
address is registered in 2, and all are in 3, does not.
"""

import gc
import weakref

import grain
import numpy as np
from grain._src.python.dataset import stats


def make_dataset() -> grain.IterDataset:
    dataset = grain.MapDataset.range(8).map(lambda i: np.full(2, i)).batch(4)
    return dataset.to_iter_dataset(grain.ReadOptions(num_threads=0, prefetch_buffer_size=0))


def consume(iterator) -> None:
    for _ in iterator:
        pass


def registry() -> set:
    return stats._iter_weakref_registry


def is_registered(iterator) -> bool:
    return any(ref() is iterator for ref in registry())


def describe() -> str:
    refs = registry()
    return f"registry size {len(refs)}, dead refs {sum(ref() is None for ref in refs)}"


def dead_refs_stay(dataset) -> None:
    iterators = [iter(dataset) for _ in range(100)]
    for iterator in iterators:
        consume(iterator)
    print(f"  100 live iterators consumed: {sum(map(is_registered, iterators))} registered; "
          f"{describe()}")
    del iterators, iterator
    gc.collect()
    print(f"  all 100 deleted and collected: {describe()}")


def epoch_loop(dataset, passes: int) -> None:
    previous_addresses = set()
    registered = reused = reused_registered = 0
    first_reuse = None
    for _ in range(passes):
        iterator = iter(dataset)
        consume(iterator)
        is_reused = id(iterator) in previous_addresses
        is_in_registry = is_registered(iterator)
        registered += is_in_registry
        reused += is_reused
        reused_registered += is_reused and is_in_registry
        if is_reused and first_reuse is None:
            stale = [ref for ref in registry() if hash(ref) == id(iterator)]
            first_reuse = (is_in_registry, len(stale), all(ref() is None for ref in stale))
        previous_addresses.add(id(iterator))
        del iterator
        gc.collect()
    print(f"  {passes} passes, a new iterator each: {registered} registered, "
          f"{passes - registered} not registered")
    print(f"  passes whose iterator got an earlier pass's address: {reused}, "
          f"of which registered: {reused_registered}")
    if first_reuse is not None:
        is_in_registry, matching, all_dead = first_reuse
        print(f"  first such iterator: registered {is_in_registry}; refs in the registry at its "
              f"address: {matching}, all dead: {all_dead}")
    print(f"  after the loop: {describe()}")


class DropOnDeathRegistry(set):
    """The suggested direction: drop a ref as soon as its iterator is collected."""

    def add(self, ref) -> None:
        iterator = ref()
        if iterator is not None:
            weakref.finalize(iterator, self.discard, ref)
        super().add(ref)


print(f"grain {grain.__version__}, numpy {np.__version__}")
dataset = make_dataset()

print("1. Dead refs stay in the registry")
dead_refs_stay(dataset)

print("2. An iterator re-created each epoch is often not registered")
epoch_loop(dataset, passes=200)

print("3. Control: a registry that drops refs when their iterator dies")
stats._iter_weakref_registry = DropOnDeathRegistry()
dead_refs_stay(dataset)
epoch_loop(dataset, passes=200)

Run in an empty directory:

uv run --isolated --no-project --python 3.12 --with grain==0.2.18 --with numpy python repro.py

Output:

WARNING:absl:Failed to load jax profiler: No module named 'jax'
grain 0.2.18, numpy 2.5.3
1. Dead refs stay in the registry
  100 live iterators consumed: 100 registered; registry size 100, dead refs 0
  all 100 deleted and collected: registry size 100, dead refs 100
2. An iterator re-created each epoch is often not registered
  200 passes, a new iterator each: 11 registered, 189 not registered
  passes whose iterator got an earlier pass's address: 184, of which registered: 0
  first such iterator: registered False; refs in the registry at its address: 1, all dead: True
  after the loop: registry size 111, dead refs 111
3. Control: a registry that drops refs when their iterator dies
  100 live iterators consumed: 100 registered; registry size 100, dead refs 0
  all 100 deleted and collected: registry size 0, dead refs 0
  200 passes, a new iterator each: 200 registered, 0 not registered
  passes whose iterator got an earlier pass's address: 184, of which registered: 184
  first such iterator: registered True; refs in the registry at its address: 1, all dead: False
  after the loop: registry size 0, dead refs 0

The first line is grain's import-time warning when jax is not installed; the repro installs only grain and NumPy, and the behaviour does not depend on jax.

The counts in sections 2 and 3 change between runs because they depend on the allocator. A second run gave 189 reused addresses in section 2, with 0 registered, and 180 in section 3, all registered. The 189 iterators in section 2 that were not registered include a few that landed on the addresses of section 1's dead iterators rather than an earlier pass's.

Where

On main (7e36b7231574d0f9ab106f903a8d832f58157594), grain/_src/python/dataset/stats.py:

• L47: _iter_weakref_registry = set()
• L439-455: HashableWeakRef, with __hash__ returning self._obj_id and __eq__ comparing _obj_id
• L604: _iter_weakref_registry.add(config.iter_weakref) in Stats.__init__
• L607: _iter_weakref_registry.discard(p._config.iter_weakref) for each parent

v0.2.18 has the same code at L48, L440-456, L605 and L608.

The refs are created in MapDataset._initialize_stats and DatasetIterator._initialize_stats (dataset.py L946 and L1689 on main), transformations/prefetch.py L66, transformations/interleave.py L406 and transformations/process_prefetch.py L413.

The open #1426 (py_permanently_disable_stats) keeps the registry empty when stats are turned off. With stats on, which is the default, the behaviour above is unchanged.

Suggested direction

Drop a ref when its referent dies, so a reused address never meets a stale entry. A weakref.ref callback or weakref.finalize that discards the key would do it, as the control above shows. A weakref.WeakSet of the iterators would do the same, provided the iterator types stay hashable by identity. There may be a better fit depending on what reads the registry.