[mahdi-shafiei]

Fixes #1420.

What

Two fixes, one per cause in the issue.

No reference cycle when a prefetch iterator re-raises. ThreadPrefetchDatasetIterator.__next__ and ProcessPrefetchDatasetIterator.__next__ now delete their err local after raising it (try: raise err / finally: del err, as concurrent.futures.Future.result() does). The producer _put_iterator_elements_in_buffer drops its element and state locals once they are in the buffer, so an exception raised by the next __next__ call does not keep the last element alive through the producer frame. The ProcessPrefetchDatasetIterator change also folds the two raise err sites into one. Behaviour is unchanged: StopIteration is still raised without a reconstructed traceback, and other errors still get the worker traceback.

One closer per SharedMemory. SharedMemoryArray now creates its segments as _SharedMemory, a private SharedMemory subclass. When SharedMemoryArray.__del__ hands a close to the delete thread pool, it sets close_handed_over on the segment, and _SharedMemory.__del__ then leaves the close to the delete thread. When the cyclic GC finalizes the array and its SharedMemory in one pass, only the delete thread closes the segment. If SharedMemory.__del__ runs first, it closes on the collecting thread before the hand-off, and the delete thread's close() finds _fd == -1. Either way os.close runs once.

A plain SharedMemory passed to SharedMemoryArray.from_shared_memory keeps its own __del__, which cannot be coordinated with the delete thread, so __del__ now closes that one synchronously. Every segment Grain creates (SharedMemoryArray(...), from_metadata, unpickling) is a _SharedMemory and still goes through the delete thread pool, so munmap stays off the consumer thread as before. The semaphore flow control is unchanged.

Why not something smaller

Making SharedMemory.close() thread-safe with a lock would still let the collecting thread do the munmap whenever it wins the lock, which is what the delete pool exists to avoid. Disabling async delete, or calling gc.collect() at pool shutdown, hides the race instead of removing the second closer. Fixing only the cycles is not enough either. Any user code that puts a batch in a cycle (an exception stored with a batch in a local, for example) reaches the same double close.

Tests

• SharedMemoryArrayTest.test_cyclic_gc_closes_shm_once_with_async_del: an array that survived one collection is made cyclic and collected with async delete on. os.close is wrapped to record calls and to sleep briefly after closing the segment's fd. On main the fd is closed twice (2 != 1, 3 of 3 runs). With the fix it is closed once.
• ThreadPrefetchIterDatasetTest.test_finished_iterator_is_freed_without_cyclic_gc_{stop_iteration,error,error_without_traceback_filtering}: with the GC disabled, after a ThreadPrefetch iterator ends and is dropped, the last element and the parent iterator are freed. On main stop_iteration and error_without_traceback_filtering fail. error passes on main because traceback filtering happens to break that cycle, and it stays as a guard.
• ThreadPrefetchIterDatasetTest.test_raised_error_does_not_keep_last_element_alive: a caller that keeps the raised error does not keep the last element alive. Fails on main.
• MultiprocessingPrefetchTest.test_finished_iteration_leaves_no_cyclic_garbage: after an mp_prefetch run over 256 KiB elements ends and is closed, a gc.DEBUG_SAVEALL collection finds no SharedMemoryArray, no ProcessPrefetchDatasetIterator and no exception. Fails on main.

Each source change is pinned by its own test. Without the process_prefetch.py change the mp test fails on ProcessPrefetchDatasetIterator. Without the producer del the held-error test fails.

Local runs on Linux, CPython 3.12.6, from a source install (pip install -e .), each module through absltest.main:

module	main	this PR
shared_memory_array_test	11 OK	12 OK
prefetch_test	90 OK (1 skipped)	94 OK (1 skipped)
process_prefetch_test	67 OK	68 OK
batch_test, cache_test, repeat_test, stats_test, stats_utils_test, options_test, traceback_util_test	OK	OK
data_loader_test, dataset_test	6 errors, 1 failure	same 6 errors, same 1 failure

The data_loader_test errors open testdata/ by a relative path and the dataset_test failure is a logging assertion. Both fail identically on main outside Bazel. Bazel itself was not run locally. The new prefetch_test and process_prefetch_test imports have matching BUILD deps.

The repro from the issue prints, with this PR:

Part 1: one mp_prefetch run served to its end, then closed.
  sum of elements: 2016.0
  left in cyclic garbage:
    SharedMemoryArray: 0
    ProcessPrefetchDatasetIterator: 0
    SemLock: 0
    StopIteration: 0
Part 2: ... open_a_file_in_between=False.
  os.close(7) calls: [('Thread-1 (worker)', 7)]
Part 2: ... open_a_file_in_between=True.
  os.close(7) calls: [('Thread-1 (worker)', 7)]
  file opened in between got fd 7: still open

Not in this PR

On a worker error, traceback_util.reconstruct_traceback still leaves a cycle (its _ local from sys.exc_info(), and its frame is the f_back of the reconstructed frames). It holds the ProcessPrefetchDatasetIterator and its SemLocks but no SharedMemoryArray, so it can still cause the ResourceTracker reentrancy warning on error paths. PrefetchDatasetIterator.__next__, batch.py parallel stacking and parallel.run_in_parallel have the same re-raise pattern on error paths. These are separate changes.

---

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

[mahdi-shafiei]

Two more effects of the cycle in #1420, both seen downstream and both gone with this PR's head (65bac7d).

The consumer's frame keeps its last element alive. The traceback in the cycle reaches the frame that iterated the dataset, so after a function that loops over an mp_prefetch dataset returns, its last element stays alive until the cyclic GC runs. Dropping the dataset doesn't free it either. The frame keeps all its locals, so in a training loop that is the last batch along with anything made from it. In our loader we saw it hold a jax.Array (on the CPU backend; on an accelerator that would be device memory).

"""Repro: after an mp_prefetch run ends, the consumer's frame keeps its last element alive.

Needs only grain and numpy. Run on Linux:
  python repro.py
"""

import gc
import sys
import time
import weakref

from absl import flags
import grain
import numpy as np


def _element(index):
  # 256 KiB, so it is sent from the worker through shared memory.
  return np.full((64, 1024), index, np.float32)


def consume(ds):
  last = None
  for x in ds:
    last = x
  return weakref.ref(last)


if __name__ == "__main__":
  flags.FLAGS(sys.argv[:1])
  print(f"grain {grain.__version__}, Python {sys.version.split()[0]}")
  ds = grain.MapDataset.range(64).map(_element).to_iter_dataset()
  ds = ds.mp_prefetch(grain.MultiprocessingOptions(num_workers=2))
  gc.collect()
  gc.disable()
  last = consume(ds)
  print(f"  consume() returned:          last element alive = {last() is not None}")
  del ds
  print(f"  dataset dropped:             last element alive = {last() is not None}")
  gc.collect()
  print(f"  after gc.collect():          last element alive = {last() is not None}")
  # Let the delete thread pool finish before exit. Without the wait, the
  # resource tracker can report a segment still queued for deletion at exit,
  # with or without the fix.
  time.sleep(1)

Command, run in an empty directory, once per build:

env -i HOME=$HOME PATH=/usr/bin:/bin:$HOME/.local/bin uv run -q --isolated --no-project --python 3.12 --with grain==0.2.18 --with numpy python repro.py
env -i HOME=$HOME PATH=/usr/bin:/bin:$HOME/.local/bin uv run -q --isolated --no-project --python 3.12 --with "grain @ git+https://github.com/mahdi-shafiei/grain@65bac7d0a44d17f5ec57e6aeef95ac65c81c7204" --with numpy python repro.py

Output, the same in 3 of 3 runs of each (absl's jax profiler warnings removed):

grain 0.2.18, Python 3.12.6
  consume() returned:          last element alive = True
  dataset dropped:             last element alive = True
  after gc.collect():          last element alive = False

grain 0.2.19, Python 3.12.6
  consume() returned:          last element alive = False
  dataset dropped:             last element alive = False
  after gc.collect():          last element alive = False

(0.2.19 is the version string on main, which the PR branch is based on.)

A collection on the consumer thread does Grain's cleanup. Whichever thread triggers the collection that frees the cycle also runs the SharedMemory finalizers. We saw this in a test that counts the Python calls the consumer thread makes while iterating. The test collects before it starts counting, and that collection finalizes earlier runs' arrays, which hand their segments to the delete pool. One segment was still held in a cycle afterwards and was freed by the next collection, inside the counted window, on the consumer thread (SharedMemory.__del__ and close). The test failed in CI on some runs and not others. We reran a failing CI shard locally with the same test order and seed, swapping only Grain (0.2.18 from PyPI, CPython 3.13):

Grain	result	SharedMemory calls on the consumer thread inside the counted window
0.2.18	fails	2
0.2.18 with only del err after the two StopIteration re-raises	passes	0
0.2.18 with this PR's three source files patched in	passes	0

That test is in our own suite and needs our package, so it isn't posted here. The repro above needs only Grain and NumPy.