Raise output pickling errors from process prefetch workers

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",
     ],
 )

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,
       ):

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