From 7d17e2600cd51f27245fb90a36b14d85588fc18e Mon Sep 17 00:00:00 2001
From: Valentyn Tymofieiev <valentyn@google.com>
Date: Wed, 14 Oct 2026 16:37:04 -0700
Subject: [PATCH 1/2] Correct pickling error check.

---
 sdks/python/apache_beam/ml/inference/base.py  | 14 ++++++-
 .../apache_beam/ml/inference/base_test.py     | 40 +++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 7a19eba35117..04bbf9438740 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -148,6 +148,18 @@ 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 +2125,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. '
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py
index a8ffc81e5acf..f381801a1410 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/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:

From 94964252fc96c65f47377854b362359a0a5c2784 Mon Sep 17 00:00:00 2001
From: Valentyn Tymofieiev <valentyn@google.com>
Date: Wed, 14 Oct 2026 17:32:29 -0700
Subject: [PATCH 2/2] yapf

---
 sdks/python/apache_beam/ml/inference/base.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 04bbf9438740..69bad902e4b8 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -156,8 +156,7 @@ def _is_pickling_error(e: BaseException) -> bool:
   # 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())
+      isinstance(e, (TypeError, AttributeError)) and 'pickle' in str(e).lower())
 
 
 @dataclass(frozen=True)