From ddf79358e290015dbaf727f96c7520112eb45fa4 Mon Sep 17 00:00:00 2001 From: Grain Team Date: Mon, 23 Mar 2026 02:16:53 -0700 Subject: [PATCH] Add traceback to exceptions raised in prefetch worker. PiperOrigin-RevId: 887969962 --- grain/_src/python/dataset/transformations/process_prefetch.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/grain/_src/python/dataset/transformations/process_prefetch.py b/grain/_src/python/dataset/transformations/process_prefetch.py index 8d5a3fc27..3c590ebca 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch.py +++ b/grain/_src/python/dataset/transformations/process_prefetch.py @@ -23,6 +23,7 @@ from multiprocessing import util import queue import time +import traceback from typing import Any, TypeVar import weakref @@ -238,6 +239,7 @@ def _put_dataset_elements_in_buffer( try: element = it.__next__() except Exception as e: # pylint: disable=broad-except + e.add_note("\n".join(traceback.format_exception(e))) grain_queue.add_element_to_queue( # pytype: disable=wrong-arg-types (None, None, None, e), buffer, should_stop.is_set ) @@ -262,6 +264,8 @@ def _put_dataset_elements_in_buffer( except Exception as e: # pylint: disable=broad-except _clear_queue_and_maybe_unlink_shm(buffer) _clear_queue_and_maybe_unlink_shm(set_state_queue) + e.add_note("\n".join(traceback.format_exception(e))) + grain_queue.add_element_to_queue( # pytype: disable=wrong-arg-types (None, None, None, e), buffer, should_stop.is_set )