Correct pickling error check.

sdks/python/apache_beam/ml/inference/base.py

@@ -148,6 +148,17 @@ def _to_microseconds(time_ns: int) -> int:
   return int(time_ns / _NANOSECOND_TO_MICROSECOND)
 
 
+def _is_pickling_error(e: BaseException) -> bool:
+  """Returns True if the exception was likely raised while pickling data."""
+  if isinstance(e, pickle.PickleError):
+    return True
+  # Unpicklable objects (e.g. C extension objects, locks, local functions)
+  # typically surface as TypeError or AttributeError, for example
+  # "cannot pickle 'X' object" or "Can't pickle local object ...".
+  return (
+      isinstance(e, (TypeError, AttributeError)) and 'pickle' in str(e).lower())
+
+
 @dataclass(frozen=True)
 class KeyModelPathMapping(Generic[KeyT]):
   """
@@ -2113,7 +2124,7 @@ def _run_inference(self, batch, inference_args):
     except BaseException as e:
       if self._metrics_collector:
         self._metrics_collector.failed_batches_counter.inc()
-      if (e is pickle.PickleError and
+      if (_is_pickling_error(e) and
           self._model_handler.share_model_across_processes()):
         raise TypeError(
             'Pickling error encountered while running inference. '

sdks/python/apache_beam/ml/inference/base_test.py

@@ -24,6 +24,7 @@
 import random
 import sys
 import tempfile
+import threading
 import time
 import unittest
 import unittest.mock
@@ -56,6 +57,26 @@ def predict(self, example: int) -> int:
     return example + 1
 
 
+class _Unpicklable:
+  """Mimics an object that cannot be pickled, e.g. a C extension object."""
+  def __init__(self):
+    self._lock = threading.Lock()
+
+
+class FakeModelHandlerSendsUnpicklableData(base.ModelHandler[int,
+                                                             int,
+                                                             FakeModel]):
+  def load_model(self):
+    return FakeModel()
+
+  def run_inference(self, batch, model, inference_args=None):
+    # Calls on a model shared across processes pickle their arguments.
+    return [model.predict(_Unpicklable()) for _ in batch]
+
+  def share_model_across_processes(self):
+    return True
+
+
 class FakeStatefulModel:
   def __init__(self, state: int):
     if state == 100:
@@ -349,6 +370,25 @@ def test_run_inference_impl_simple_examples_multi_process_shared(self):
           FakeModelHandler(multi_process_shared=True))
       assert_that(actual, equal_to(expected), label='assert:inferences')
 
+  def test_run_inference_unpicklable_data_multi_process_shared(self):
+    with self.assertRaisesRegex(Exception, 'Pickling error encountered'):
+      with TestPipeline() as pipeline:
+        _ = (
+            pipeline
+            | 'start' >> beam.Create([1])
+            | base.RunInference(FakeModelHandlerSendsUnpicklableData()))
+
+  def test_is_pickling_error(self):
+    self.assertTrue(base._is_pickling_error(pickle.PicklingError('x')))
+    self.assertTrue(
+        base._is_pickling_error(
+            TypeError("cannot pickle '_thread.lock' object")))
+    self.assertTrue(
+        base._is_pickling_error(
+            AttributeError("Can't pickle local object 'f.<locals>.g'")))
+    self.assertFalse(base._is_pickling_error(TypeError('unsupported type')))
+    self.assertFalse(base._is_pickling_error(ValueError('pickle')))
+
   def test_run_inference_impl_simple_examples_multi_process_shared_multi_copy(
       self):
     with TestPipeline() as pipeline: