From 28603990a79b46cd64413b77764ab713349e87ef Mon Sep 17 00:00:00 2001 From: Rio Yu <52408936+rioyu123@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:56:51 +0800 Subject: [PATCH] Raise output pickling errors from process prefetch workers Signed-off-by: Rio Yu <52408936+rioyu123@users.noreply.github.com> --- .../_src/python/dataset/transformations/BUILD | 2 + .../transformations/process_prefetch.py | 36 +++++++- .../transformations/process_prefetch_test.py | 92 +++++++++++++++++++ 3 files changed, 129 insertions(+), 1 deletion(-) diff --git a/grain/_src/python/dataset/transformations/BUILD b/grain/_src/python/dataset/transformations/BUILD index 01ff40058..a8a2e6a8b 100644 --- a/grain/_src/python/dataset/transformations/BUILD +++ b/grain/_src/python/dataset/transformations/BUILD @@ -403,9 +403,11 @@ py_test( "//grain/_src/python:options", "//grain/_src/python/dataset", "//grain/_src/python/dataset:base", + "//grain/_src/python/ipc:shared_memory_array", "@abseil-py//absl/logging", "@abseil-py//absl/testing:absltest", "@abseil-py//absl/testing:parameterized", + "@pypi//cloudpickle:pkg", "@pypi//numpy:pkg", ], ) diff --git a/grain/_src/python/dataset/transformations/process_prefetch.py b/grain/_src/python/dataset/transformations/process_prefetch.py index 93e0f44e9..0eb483ec4 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch.py +++ b/grain/_src/python/dataset/transformations/process_prefetch.py @@ -18,6 +18,7 @@ from collections.abc import Callable, Sequence import functools from multiprocessing import queues +from multiprocessing import reduction from multiprocessing import sharedctypes from multiprocessing import synchronize from multiprocessing import util @@ -198,6 +199,18 @@ def _element_spec(self) -> Any: return dataset.get_element_spec(self._parent) +class _SerializedElement: + """Pickles eagerly so failures are raised in the worker loop.""" + + def __init__(self, element): + self._data = bytes(reduction.ForkingPickler.dumps(element)) + + def __reduce__(self): + # The queue's feeder only serializes these bytes, not the element again. + # Queue readers still receive the original tuple, including during cleanup. + return reduction.ForkingPickler.loads, (self._data,) + + def _put_dataset_elements_in_buffer( pickled_parse_debug_flags_fn: bytes, pickled_worker_init_fn: bytes, @@ -289,8 +302,29 @@ def _put_dataset_elements_in_buffer( it._stats.record_bytes_produced(element) # pylint: disable=protected-access if next_index is not None: next_index += 1 + try: + buffered_element = _SerializedElement( + (element, it.get_state(), next_index, None, None) + ) + except Exception as e: # pylint: disable=broad-except + shared_memory_array.unlink_shm(element) + grain_queue.add_element_to_queue( + ( + None, + None, + None, + e, + traceback_util.PicklableTraceback.from_traceback( + e.__traceback__ + ), + ), + buffer, # pyrefly: ignore[bad-argument-type] + should_stop.is_set, + ) + parent_exhausted = True + continue if not grain_queue.add_element_to_queue( - (element, it.get_state(), next_index, None, None), + buffered_element, buffer, # pyrefly: ignore[bad-argument-type] should_stop.is_set, ): diff --git a/grain/_src/python/dataset/transformations/process_prefetch_test.py b/grain/_src/python/dataset/transformations/process_prefetch_test.py index f3ff30d65..486e03ec7 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch_test.py +++ b/grain/_src/python/dataset/transformations/process_prefetch_test.py @@ -15,7 +15,11 @@ from concurrent import futures import dataclasses import logging as std_logging +from multiprocessing import reduction +from multiprocessing import shared_memory import os +import pickle +import queue import sys import time import traceback @@ -26,6 +30,7 @@ from absl import logging from absl.testing import absltest from absl.testing import parameterized +import cloudpickle from grain._src.core import config from grain._src.core import transforms import multiprocessing as mp @@ -34,6 +39,7 @@ from grain._src.python.dataset import dataset from grain._src.python.dataset.transformations import process_prefetch from grain._src.python.dataset.transformations import testing_util +from grain._src.python.ipc import shared_memory_array import numpy as np @@ -62,6 +68,70 @@ def setUp(self): .filter(FilterKeepingOddElementsOnly()) ) + def test_serializes_output_once(self): + reducer = mock.Mock(return_value=(int, (17,))) + + class Output: + + def __reduce__(self): + return reducer() + + element = (Output(), {'index': 1}, 1, None, None) + serialized = process_prefetch._SerializedElement(element) + actual = reduction.ForkingPickler.loads( + reduction.ForkingPickler.dumps(serialized) + ) + self.assertEqual(actual, (17, {'index': 1}, 1, None, None)) + reducer.assert_called_once_with() + + def test_unlinks_output_when_serialization_fails(self): + ds = ( + dataset.MapDataset.range(1) + .map(lambda _: (np.ones(4), lambda: None)) + .to_iter_dataset() + ) + buffer = queue.Queue() + should_stop = mp.Event() + emitted = [] + + def record_error_and_stop(element, **unused_kwargs): + emitted.append(element) + should_stop.set() + + with ( + mock.patch.object(process_prefetch, '_is_in_worker_process', False), + mock.patch.object(buffer, 'put', side_effect=record_error_and_stop), + mock.patch.object( + shared_memory_array, + 'unlink_shm', + wraps=shared_memory_array.unlink_shm, + ) as unlink, + ): + process_prefetch._put_dataset_elements_in_buffer( + cloudpickle.dumps(lambda _: None), + cloudpickle.dumps(None), + cloudpickle.dumps(ds), + buffer, + should_stop, + mp.Value('i', 0), + queue.Queue(), + None, + None, + None, + {}, + ) + unlink.assert_called_once() + metadata = unlink.call_args.args[0][0] + self.assertIsInstance( + metadata, shared_memory_array.SharedMemoryArrayMetadata + ) + with self.assertRaises(FileNotFoundError): + shared_memory.SharedMemory(metadata.name) + self.assertLen(emitted, 1) + self.assertIsInstance( + emitted[0][3], (AttributeError, TypeError, pickle.PicklingError) + ) + @parameterized.parameters( ( 'off', @@ -361,6 +431,28 @@ def __getstate__(self): ): list(ds) + @parameterized.parameters(1, 4) + def test_reports_unpicklable_output(self, buffer_size): + def transform(i): + extra = (lambda: i) if i == 2 else i + return np.full(4, i), {'extra': extra} + + ds = dataset.MapDataset.range(5).map(transform).to_iter_dataset() + ds = process_prefetch.ProcessPrefetchIterDataset( + ds, buffer_size=buffer_size + ) + iterator = iter(ds) + try: + self.assertEqual(int(next(iterator)[0][0]), 0) + self.assertEqual(int(next(iterator)[0][0]), 1) + with self.assertRaisesRegex( + (AttributeError, TypeError, pickle.PicklingError), 'lambda' + ): + next(iterator) + self.assertIsNone(iterator._prefetch_process) + finally: + iterator.close() + def test_reports_worker_crash(self): def failing_transform(element): del element