diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index ce014e4b04f9..c4c3f146ed4e 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -767,6 +767,10 @@ def __init__( at the same time. For more information about memory management, see `Use a keyed `ModelHandler _`. # 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 @@ def __call__(self, *args, **kwargs): 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 @@ def __init__( 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 64f5c3b1525d..5890dd087d9a 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 @@ def test_keyed_many_mhs_ignores_inner_no_batching_with_warning(self): 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)]