Tensorflow exceptions in dataset from generator logged before except block

Question

I am trying to encapsulate my Tensorflow dataset generator in a try-except block. However, the whole traceback is logged before the code enters the catch block.

In this minimal case:

try:
    def data_generator():
        for i in range(5):
            # Also tried tf.debugging.Assert(i != 4, ["Let's raise to test"])
            if i ==4:
                raise ValueError("Lets raise to test")
            yield i
    dataset = tf.data.Dataset.from_generator(
        generator=data_generator,
        output_signature=tf.TensorSpec(shape=(), dtype=tf.int32)
    )
    for element in dataset:
        print(element.numpy())
except Exception as exc:
    print("Something went wrong-->", exc)

This is logged:

2023-11-29 11:13:20.003212: W tensorflow/core/framework/op_kernel.cc:1733] INVALID_ARGUMENT: ValueError: Lets raise to test
Traceback (most recent call last):

  File "/my-pacakge-dir/lib/python3.8/site-packages/tensorflow/python/ops/script_ops.py", line 270, in __call__
    ret = func(*args)

  File "/my-pacakge-dir/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 642, in wrapper
    return func(*args, **kwargs)

  File "/my-pacakge-dir/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 1030, in generator_py_func
    values = next(generator_state.get_iterator(iterator_id))

  File "/my_script_dir/test_tf_logger.py", line 16, in data_generator
    raise ValueError("Lets raise to test")

ValueError: Lets raise to test


Something went wrong--> ValueError: Lets raise to test
Traceback (most recent call last):

  File "/my-pacakge-dir/lib/python3.8/site-packages/tensorflow/python/ops/script_ops.py", line 270, in __call__
    ret = func(*args)

  File "/my-pacakge-dir/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 642, in wrapper
    return func(*args, **kwargs)

  File "/my-pacakge-dir/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 1030, in generator_py_func
    values = next(generator_state.get_iterator(iterator_id))

  File "/my_script_dir/test_tf_logger.py", line 16, in data_generator
    raise ValueError("Lets raise to test")

ValueError: Lets raise to test

To test it with another exception raised from an invalid tf op.

try:
    numerator = tf.constant(1.0, dtype=tf.float32)
    denominator = tf.constant(0.0, dtype=tf.float32)
    result = tf.divide(numerator, denominator)
    result = tf.debugging.check_numerics(result, message="Operation result is not valid.")
except Exception as exc:
    print("Invalid TF operation-->",exc)

This correctly only logs the message in the except block.