`Dataset.from_generator` errors in debug mode when garbage collection is triggered

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

tf 2.22.0-dev20260823

Custom code

No

OS platform and distribution

macOS 26.5.2, arm64 (Apple silicon)

Mobile device

No response

Python version

3.12

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

If tf.data.experimental.enable_debug_mode() is enabled, any tf.data.Dataset.from_generator dataset stops working after the first garbage collection. Iteration fails with:

ValueError: Could not find callback with key=pyfunc_10 in the registry.

Without debug mode the same dataset iterates fine.

In a real workload this looks nondeterministic, because it depends on when a background collection happens to run. We hit it during a training job with ~1500 examples (~1000 examples passes).

In the reproduction below I use gc.collect() to trigger the issue deterministically.

Standalone code to reproduce the issue

import gc
import tensorflow as tf

tf.data.experimental.enable_debug_mode()

def generator():
    yield 1
    yield 2

dataset = tf.data.Dataset.from_generator(
    generator, output_signature=tf.TensorSpec(shape=(), dtype=tf.int64)
)

iterator = iter(dataset)
print("first:", next(iterator).numpy())
gc.collect()
print("second:", next(iterator).numpy())  # ValueError: Could not find callback ...

Relevant log output

first: 1
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1788872103.355896 1774757 op_kernel.cc:1921] INVALID_ARGUMENT: ValueError: Could not find callback with key=pyfunc_10 in the registry.
Traceback (most recent call last):

  File "/path/.venv/lib/python3.12/site-packages/tensorflow/python/ops/script_ops.py", line 257, in __call__
    raise ValueError(f"Could not find callback with key={token} in the "

ValueError: Could not find callback with key=pyfunc_10 in the registry.


=== Source Location Trace: ===
tensorflow/python/lib/core/py_func.cc:229
tensorflow/python/lib/core/py_func.cc:257

W0000 00:00:1788872103.355944 1770362 local_rendezvous.cc:461] Local rendezvous is aborting with status: INVALID_ARGUMENT: ValueError: Could not find callback with key=pyfunc_10 in the registry.
Traceback (most recent call last):

  File "/path/.venv/lib/python3.12/site-packages/tensorflow/python/ops/script_ops.py", line 257, in __call__
    raise ValueError(f"Could not find callback with key={token} in the "

ValueError: Could not find callback with key=pyfunc_10 in the registry.


         [[{{node EagerPyFunc}}]]
=== Source Location Trace: ===
tensorflow/python/lib/core/py_func.cc:229
tensorflow/python/lib/core/py_func.cc:257
tensorflow/core/common_runtime/executor.cc:1156
tensorflow/core/common_runtime/function.cc:1436
tensorflow/core/data/captured_function.cc:920
tensorflow/core/kernels/data/flat_map_dataset_op.cc:220
tensorflow/core/data/root_dataset.cc:248
tensorflow/core/kernels/data/iterator_ops.cc:950
tensorflow/core/common_runtime/single_threaded_executor.cc:449
tensorflow/core/common_runtime/function.cc:1418

Traceback (most recent call last):
  File "/path/repro.py", line 20, in <module>
    print("second:", next(iterator).numpy())  # ValueError: Could not find callback ...
                     ^^^^^^^^^^^^^^
  File "/path.venv/lib/python3.12/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 826, in __next__
    return self._next_internal()
           ^^^^^^^^^^^^^^^^^^^^^
  File "/path.venv/lib/python3.12/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 776, in _next_internal
    ret = gen_dataset_ops.iterator_get_next(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/.venv/lib/python3.12/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 3086, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/path/.venv/lib/python3.12/site-packages/tensorflow/python/framework/ops.py", line 6083, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__IteratorGetNext_output_types_1_device_/job:localhost/replica:0/task:0/device:CPU:0}} ValueError: Could not find callback with key=pyfunc_10 in the registry.
Traceback (most recent call last):

  File "/path/.venv/lib/python3.12/site-packages/tensorflow/python/ops/script_ops.py", line 257, in __call__
    raise ValueError(f"Could not find callback with key={token} in the "

ValueError: Could not find callback with key=pyfunc_10 in the registry.


         [[{{node EagerPyFunc}}]] [Op:IteratorGetNext] name: 
W0000 00:00:1788872103.466339 1770362 generator_dataset_op.cc:113] Error occurred when finalizing GeneratorDataset iterator: FAILED_PRECONDITION: Python interpreter state is not initialized. The process may be terminated.
         [[{{node EagerPyFunc}}]]
=== Source Location Trace: ===
tensorflow/python/lib/core/py_func.cc:368
tensorflow/core/common_runtime/executor.cc:1156
tensorflow/core/common_runtime/function.cc:1436
tensorflow/core/data/captured_function.cc:953