This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 59e32288d54 Correct pickling error check. (#40597)
59e32288d54 is described below
commit 59e32288d541a5340e65a0eaffbda48d0103ef99
Author: tvalentyn <[email protected]>
AuthorDate: Thu Oct 15 07:30:44 2026 -0700
Correct pickling error check. (#40597)
* Correct pickling error check.
* yapf
---
sdks/python/apache_beam/ml/inference/base.py | 13 +++++++-
sdks/python/apache_beam/ml/inference/base_test.py | 40 +++++++++++++++++++++++
2 files changed, 52 insertions(+), 1 deletion(-)
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index 7a19eba3511..69bad902e4b 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/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 @@ class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT,
PredictionT]):
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 a8ffc81e5ac..f381801a141 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 pickle
import random
import sys
import tempfile
+import threading
import time
import unittest
import unittest.mock
@@ -56,6 +57,26 @@ class FakeModel:
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 @@ class RunInferenceBaseTest(unittest.TestCase):
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:
