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 d50c0a8235c [RunInference] Reject use_model_manager with a multi-model
KeyedModelHandler (#40574)
d50c0a8235c is described below
commit d50c0a8235c138f9653405e06efebf5fca0cc545
Author: Divyansh Yadav <[email protected]>
AuthorDate: Wed Oct 21 01:42:43 2026 +0530
[RunInference] Reject use_model_manager with a multi-model
KeyedModelHandler (#40574)
---
sdks/python/apache_beam/ml/inference/base.py | 33 +++++++++++++++++++++++
sdks/python/apache_beam/ml/inference/base_test.py | 23 ++++++++++++++++
2 files changed, 56 insertions(+)
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index ce014e4b04f..c4c3f146ed4 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -767,6 +767,10 @@ class KeyedModelHandler(Generic[KeyT, ExampleT,
PredictionT, ModelT],
at the same time. For more information about memory management, see
`Use a keyed `ModelHandler
<https://beam.apache.org/documentation/ml/about-ml/#use-a-keyed-modelhandler-object>_`.
# pylint: disable=line-too-long
+ A KeyedModelHandler with many models defined per key cannot be used with
+ `RunInference(..., use_model_manager=True)`, since it manages its per-key
+ models itself.
+
Args:
unkeyed: Either (a) an implementation of ModelHandler that does not
@@ -1459,6 +1463,19 @@ class OOMProtectedFn:
raise e
+def _is_multi_model_keyed_handler(model_handler: ModelHandler) -> bool:
+ """Returns whether model_handler is a KeyedModelHandler with many models,
+ looking through any batching or pre/postprocessing wrappers."""
+ while isinstance(model_handler,
+ (_PrebatchedModelHandler,
+ _PreProcessingModelHandler,
+ _PostProcessingModelHandler)):
+ model_handler = model_handler._base
+ return (
+ isinstance(model_handler, KeyedModelHandler) and
+ not model_handler._single_model)
+
+
class RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
Iterable[ExampleT]]],
beam.PCollection[PredictionT]]):
@@ -1505,10 +1522,26 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
the same tag for different models will lead to non-deterministic
results, so exercise caution when using this parameter. This only
impacts models which are already being shared across processes.
+ use_model_manager: Whether to load and manage models with a
+ ModelManager, which tracks memory usage and evicts models as needed.
+ This is not supported with a KeyedModelHandler that maps keys to
+ multiple model handlers, since that handler manages its own per-key
+ models (see `max_models_per_worker_hint`); a ValueError is raised
+ in that case.
+ model_manager_args: Keyword arguments used to construct the
+ ModelManager when use_model_manager is True.
monitoring_transform: A PTransform that receives a copy of the
un-postprocessed PCollection of PredictionResult objects produced
directly by inference.
"""
+ if use_model_manager and _is_multi_model_keyed_handler(model_handler):
+ raise ValueError(
+ 'use_model_manager=True is not supported with a KeyedModelHandler '
+ 'that maps keys to multiple model handlers. That handler loads and '
+ 'caches its per-key models itself, so they would not be tracked by '
+ 'the ModelManager. Use max_models_per_worker_hint on the '
+ 'KeyedModelHandler to limit how many models are loaded at once, or '
+ 'set use_model_manager=False.')
self._model_handler = model_handler
self._inference_args = inference_args
self._clock = clock
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py
b/sdks/python/apache_beam/ml/inference/base_test.py
index 64f5c3b1525..5890dd087d9 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -791,6 +791,29 @@ class RunInferenceBaseTest(unittest.TestCase):
self.assertFalse(mh.should_skip_batching())
self.assertTrue(any('skips batching' in line for line in logs.output))
+ def test_keyed_many_model_handlers_with_model_manager_raises(self):
+ mhs = [
+ base.KeyModelMapping([0],
+ FakeModelHandler(
+ state=200, multi_process_shared=True)),
+ base.KeyModelMapping([1, 2, 3],
+ FakeModelHandler(multi_process_shared=True))
+ ]
+ keyed_mh = base.KeyedModelHandler(mhs)
+ with self.assertRaisesRegex(ValueError, r'use_model_manager=True'):
+ base.RunInference(keyed_mh, use_model_manager=True)
+
+ # Wrapped handlers should be rejected as well.
+ with self.assertRaisesRegex(ValueError, r'use_model_manager=True'):
+ base.RunInference(
+ keyed_mh.with_preprocess_fn(lambda x: x).with_no_batching(),
+ use_model_manager=True)
+
+ # Without the model manager, or with a single keyed model, it is allowed.
+ base.RunInference(keyed_mh)
+ base.RunInference(
+ base.KeyedModelHandler(FakeModelHandler()), use_model_manager=True)
+
def test_keyed_model_handler_get_num_bytes(self):
mh = base.KeyedModelHandler(FakeModelHandler(num_bytes_per_element=10))
batch = [('key1', 1), ('key2', 2), ('key1', 3)]
