diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 7a19eba35117..b684718edf81 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -803,10 +803,18 @@ def __init__( 'mh %s defines the following batching kwargs which will be ' 'ignored %s. Batching kwargs are not respected when ' 'more than one model handler is used in a KeyedModelHandler. If ' - 'you would like to specify resource hints, you can do so by ' + 'you would like to specify batching kwargs, you can do so by ' 'overriding the KeyedModelHandler.batch_elements_kwargs() method.', - hints, + mh, batch_kwargs) + if mh.should_skip_batching(): + logging.warning( + 'mh %s skips batching (with_no_batching()), which will be ' + 'ignored. Skipping batching is not respected per model handler ' + 'when more than one model handler is used in a ' + 'KeyedModelHandler. To skip batching, call with_no_batching() on ' + 'the KeyedModelHandler instead.', + mh) env_vars = getattr(mh, '_env_vars', {}) if len(env_vars) > 0: logging.warning( @@ -1054,8 +1062,15 @@ def override_metrics(self, metrics_namespace: str = '') -> bool: return True + def should_skip_batching(self) -> bool: + if self._single_model: + return self._unkeyed.should_skip_batching() + return False + def should_garbage_collect_on_timeout(self) -> bool: - return self._single_model and self.share_model_across_processes() + if self._single_model: + return self._unkeyed.should_garbage_collect_on_timeout() + return False class MaybeKeyedModelHandler(Generic[KeyT, ExampleT, PredictionT, ModelT], @@ -1155,6 +1170,18 @@ def share_model_across_processes(self) -> bool: def model_copies(self) -> int: return self._unkeyed.model_copies() + def update_model_paths( + self, + model: ModelT, + model_paths: Optional[Union[str, list[KeyModelPathMapping]]] = None): + return self._unkeyed.update_model_paths(model, model_paths) + + def override_metrics(self, metrics_namespace: str = '') -> bool: + return self._unkeyed.override_metrics(metrics_namespace) + + def should_garbage_collect_on_timeout(self) -> bool: + return self._unkeyed.should_garbage_collect_on_timeout() + class _PrebatchedModelHandler(Generic[ExampleT, PredictionT, ModelT], ModelHandler[Sequence[ExampleT], @@ -1214,6 +1241,18 @@ def model_copies(self) -> int: def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_postprocess_fns() + def update_model_paths( + self, + model: ModelT, + model_paths: Optional[Union[str, list[KeyModelPathMapping]]] = None): + return self._base.update_model_paths(model, model_paths) + + def override_metrics(self, metrics_namespace: str = '') -> bool: + return self._base.override_metrics(metrics_namespace) + + def should_garbage_collect_on_timeout(self) -> bool: + return self._base.should_garbage_collect_on_timeout() + class _PreProcessingModelHandler(Generic[ExampleT, PredictionT, @@ -1280,6 +1319,18 @@ def model_copies(self) -> int: def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_postprocess_fns() + def update_model_paths( + self, + model: ModelT, + model_paths: Optional[Union[str, list[KeyModelPathMapping]]] = None): + return self._base.update_model_paths(model, model_paths) + + def override_metrics(self, metrics_namespace: str = '') -> bool: + return self._base.override_metrics(metrics_namespace) + + def should_garbage_collect_on_timeout(self) -> bool: + return self._base.should_garbage_collect_on_timeout() + class _PostProcessingModelHandler(Generic[ExampleT, PredictionT, @@ -1345,6 +1396,18 @@ def model_copies(self) -> int: def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_postprocess_fns() + [self._postprocess_fn] + def update_model_paths( + self, + model: ModelT, + model_paths: Optional[Union[str, list[KeyModelPathMapping]]] = None): + return self._base.update_model_paths(model, model_paths) + + def override_metrics(self, metrics_namespace: str = '') -> bool: + return self._base.override_metrics(metrics_namespace) + + def should_garbage_collect_on_timeout(self) -> bool: + return self._base.should_garbage_collect_on_timeout() + class OOMProtectedFn: def __init__(self, func): diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index a8ffc81e5acf..8e386a76b53b 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -17,6 +17,7 @@ """Tests for apache_beam.ml.base.""" import gc +import inspect import math import multiprocessing import os @@ -571,6 +572,185 @@ def mult_two(example: str) -> int: with self.assertRaises(ValueError): base.KeyedModelHandler(mhs) + def test_run_inference_keyed_many_mhs_with_pre_and_post_processing(self): + def mult_two_keyed(element): + return (element[0], element[1] * 2) + + pipeline = TestPipeline() + examples = [1, 5, 3, 10] + keyed_examples = [(i, example) for i, example in enumerate(examples)] + # Key 0 maps to a stateful model which always returns 200, the remaining + # keys map to a model which returns example + 1. + expected = [(0, 400)] + [(i, ((example * 2) + 1) * 2) + for i, example in enumerate(examples)][1:] + mhs = [ + base.KeyModelMapping([0], + FakeModelHandler( + state=200, multi_process_shared=True)), + base.KeyModelMapping([1, 2, 3], + FakeModelHandler(multi_process_shared=True)) + ] + pcoll = pipeline | 'start' >> beam.Create(keyed_examples) + actual = pcoll | base.RunInference( + base.KeyedModelHandler(mhs).with_preprocess_fn( + mult_two_keyed).with_postprocess_fn(mult_two_keyed), + metrics_namespace='test_namespace') + assert_that(actual, equal_to(expected), label='assert:inferences') + result = pipeline.run() + result.wait_until_finish() + + # Per-key metrics are still reported by the wrapped KeyedModelHandler. + metrics_filter = MetricsFilter().with_name('0-_num_inferences') + metrics = result.metrics().query(metrics_filter) + self.assertEqual(metrics['counters'][0].committed, 1) + metrics_filter = MetricsFilter().with_name('1-_num_inferences') + metrics = result.metrics().query(metrics_filter) + self.assertEqual(metrics['counters'][0].committed, 3) + + def test_run_inference_keyed_many_mhs_with_no_batching(self): + with TestPipeline() as pipeline: + examples = [1, 5, 3, 10] + keyed_examples = [(i, example) for i, example in enumerate(examples)] + expected = [(i, example + 1) for i, example in enumerate(examples)] + expected[0] = (0, 200) + mhs = [ + base.KeyModelMapping([0], + FakeModelHandler( + state=200, multi_process_shared=True)), + base.KeyModelMapping([1, 2, 3], + FakeModelHandler(multi_process_shared=True)) + ] + # Each element is a pre-batched list of keyed examples. + pcoll = pipeline | 'start' >> beam.Create([keyed_examples]) + actual = pcoll | base.RunInference( + base.KeyedModelHandler(mhs).with_no_batching()) + assert_that(actual, equal_to(expected), label='assert:inferences') + + def test_wrapped_keyed_many_mhs_delegates_hooks(self): + class FakeModelManager: + def __init__(self): + self.updates = [] + + def update_model_handler(self, key, model_path, previous_key): + self.updates.append((key, model_path, previous_key)) + + wrappers = { + 'preprocess': lambda mh: mh.with_preprocess_fn(lambda x: x), + 'postprocess': lambda mh: mh.with_postprocess_fn(lambda x: x), + 'no_batching': lambda mh: mh.with_no_batching(), + } + for name, wrap in wrappers.items(): + with self.subTest(wrapper=name): + keyed_mh = base.KeyedModelHandler([ + base.KeyModelMapping([0], + FakeModelHandler(multi_process_shared=True)), + base.KeyModelMapping([1, 2], + FakeModelHandler(multi_process_shared=True)) + ]) + wrapped = wrap(keyed_mh) + + self.assertTrue(wrapped.override_metrics('test_namespace')) + self.assertEqual(set(keyed_mh._metrics_collectors), {0, 1}) + self.assertIsNotNone(keyed_mh._default_metrics_collector) + self.assertFalse(wrapped.should_garbage_collect_on_timeout()) + + manager = FakeModelManager() + wrapped.update_model_paths( + manager, + [ + base.KeyModelPathMapping(keys=[0], update_path='path/0'), + base.KeyModelPathMapping(keys=[1, 2], update_path='path/1'), + ]) + self.assertEqual(manager.updates, [(0, 'path/0', 0), (1, 'path/1', 1)]) + + def test_wrapper_model_handlers_delegate_hooks(self): + class HookRecordingModelHandler(FakeModelHandler): + def __init__(self): + super().__init__(multi_process_shared=True) + self.metrics_namespace = None + self.model_paths = None + + def override_metrics(self, metrics_namespace=''): + self.metrics_namespace = metrics_namespace + return True + + def should_garbage_collect_on_timeout(self): + return False + + def update_model_paths(self, model, model_paths=None): + self.model_paths = model_paths + + wrappers = { + 'maybe_keyed': base.MaybeKeyedModelHandler, + 'preprocess': lambda mh: mh.with_preprocess_fn(lambda x: x), + 'postprocess': lambda mh: mh.with_postprocess_fn(lambda x: x), + 'no_batching': lambda mh: mh.with_no_batching(), + } + for name, wrap in wrappers.items(): + with self.subTest(wrapper=name): + inner = HookRecordingModelHandler() + wrapped = wrap(inner) + self.assertTrue(wrapped.override_metrics('test_namespace')) + self.assertEqual(inner.metrics_namespace, 'test_namespace') + self.assertFalse(wrapped.should_garbage_collect_on_timeout()) + wrapped.update_model_paths(None, 'new/path') + self.assertEqual(inner.model_paths, 'new/path') + + def test_wrapper_model_handlers_override_all_model_handler_methods(self): + # Wrappers must forward every ModelHandler hook to the handler they wrap, + # otherwise behavior of the wrapped handler is silently dropped. + not_forwarded = { + 'with_preprocess_fn', + 'with_postprocess_fn', + 'with_no_batching', + 'set_environment_vars', # Wrappers copy _env_vars instead. + } + methods = [ + name for name, _ in inspect.getmembers( + base.ModelHandler, inspect.isfunction) + if not name.startswith('_') and name not in not_forwarded + ] + for cls in (base.MaybeKeyedModelHandler, + base._PrebatchedModelHandler, + base._PreProcessingModelHandler, + base._PostProcessingModelHandler): + for name in methods: + with self.subTest(cls=cls.__name__, method=name): + self.assertIsNot( + getattr(cls, name), + getattr(base.ModelHandler, name), + f'{cls.__name__} must forward {name}() to the wrapped handler') + + def test_keyed_model_handler_respects_inner_no_batching(self): + self.assertTrue( + base.KeyedModelHandler( + FakeModelHandler().with_no_batching()).should_skip_batching()) + self.assertFalse( + base.KeyedModelHandler(FakeModelHandler()).should_skip_batching()) + + with TestPipeline() as pipeline: + examples = [1, 5, 3, 10] + keyed_examples = [(i, example) for i, example in enumerate(examples)] + expected = [(i, example + 1) for i, example in enumerate(examples)] + # Each element is a pre-batched list of keyed examples. + pcoll = pipeline | 'start' >> beam.Create( + [keyed_examples[:2], keyed_examples[2:]]) + actual = pcoll | base.RunInference( + base.KeyedModelHandler(FakeModelHandler().with_no_batching())) + assert_that(actual, equal_to(expected), label='assert:inferences') + + def test_keyed_many_mhs_ignores_inner_no_batching_with_warning(self): + mhs = [ + base.KeyModelMapping( + [0], + FakeModelHandler(multi_process_shared=True).with_no_batching()), + base.KeyModelMapping([1], FakeModelHandler(multi_process_shared=True)) + ] + with self.assertLogs(level='WARNING') as logs: + mh = base.KeyedModelHandler(mhs) + self.assertFalse(mh.should_skip_batching()) + self.assertTrue(any('skips batching' in line for line in logs.output)) + 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)]