From c29a70fe5e6a3878ff7eb4aa114bf1a8316ee114 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Fri, 30 Oct 2026 16:32:40 +0000 Subject: [PATCH 1/3] [RunInference] Report the model identifier to model handlers Add ModelHandler.set_model_identifier(), which RunInference calls on each worker before the first load_model() with the step's model_identifier (the value passed to RunInference, or the uuid it generates). Handlers that publish num_loaded_models_ themselves, such as handlers that front an out-of-process model server, can then use the exact tag the step is annotated with. The value is stable across side-input model updates, matching the metric tag ModelManager uses since #40714. The hook is a no-op by default and is forwarded by the keyed, maybe-keyed, prebatched, preprocessing and postprocessing wrappers. --- sdks/python/apache_beam/ml/inference/base.py | 32 ++++++++++ .../apache_beam/ml/inference/base_test.py | 61 +++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4d5eb4f74c0a..e44d940ed881 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -360,6 +360,19 @@ def update_model_paths( """ pass + def set_model_identifier(self, model_identifier: str): + """Receives the identifier RunInference annotates this model with. + + Called by RunInference in the worker before the first ``load_model`` with + the ``model_identifier`` passed to ``RunInference`` (or the random id it + generated when none was passed). This is the value Dataflow's autoscaler + uses to find the ``num_loaded_models_`` metric for this step, + and it is stable across side-input model updates, so handlers that publish + that metric themselves should tag it with this identifier. The default is + a no-op. + """ + pass + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: """Gets all preprocessing functions to be run before batching/inference. Functions are in order that they should be applied.""" @@ -1125,6 +1138,12 @@ def update_model_path(self, model_path: Optional[str] = None): 'KeyedModelHandlers with multiple different per-key ' + 'ModelHandlers.') + def set_model_identifier(self, model_identifier: str): + if self._single_model: + return self._unkeyed.set_model_identifier(model_identifier) + for mh in self._id_to_mh_map.values(): + mh.set_model_identifier(model_identifier) + def share_model_across_processes(self) -> bool: if self._single_model: return self._unkeyed.share_model_across_processes() @@ -1251,6 +1270,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._unkeyed.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._unkeyed.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._unkeyed.get_preprocess_fns() @@ -1322,6 +1344,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._base.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_preprocess_fns() @@ -1400,6 +1425,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._base.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return [self._preprocess_fn] + self._base.get_preprocess_fns() @@ -1477,6 +1505,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._base.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_preprocess_fns() @@ -2299,6 +2330,7 @@ def get_metrics_collector(self, prefix: str = ''): def setup(self): self._metrics_collector = self.get_metrics_collector() self._model_handler.set_environment_vars() + self._model_handler.set_model_identifier(self._model_tag) self._model_metadata = load_model_status( self._model_tag, self._model_handler.share_model_across_processes()) if not self._load_model_at_runtime: diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index d4cfce8ec4f4..92b3d2b20d8a 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2829,6 +2829,67 @@ def run_inference(self, batch, model, inference_args=None): return [f'{s}:{max_len}' for s in batch] +class IdentifierReportingModelHandler(base.ModelHandler[int, str, FakeModel]): + """Emits the identifier RunInference reported, as seen at load_model time. + + The DoFn works on a pickled copy of the handler, so the only way to observe + the hook is through the pipeline output. + """ + def __init__(self): + super().__init__() + self._identifier = None + self._identifier_at_load = None + + def set_model_identifier(self, model_identifier): + self._identifier = model_identifier + + def load_model(self): + # Captures whether set_model_identifier ran before load_model. + self._identifier_at_load = self._identifier + return FakeModel() + + def run_inference(self, batch, model, inference_args=None): + return [f'{x}:{self._identifier_at_load}' for x in batch] + + +class RunInferenceModelIdentifierTest(unittest.TestCase): + """RunInference hands handlers the identifier the autoscaler looks up.""" + def test_explicit_model_identifier_is_reported_before_load(self): + with TestPipeline('FnApiRunner') as p: + out = ( + p | beam.Create([1, 2], reshuffle=False) + | base.RunInference( + IdentifierReportingModelHandler(), + model_identifier='my-step-model')) + assert_that(out, equal_to(['1:my-step-model', '2:my-step-model'])) + + def test_generated_uuid_identifier_matches_annotation(self): + # No model_identifier given: RunInference generates a uuid hex, annotates + # the inference step with it, and reports that same value to the handler. + with TestPipeline('FnApiRunner') as p: + out = ( + p | beam.Create([1], reshuffle=False) + | base.RunInference(IdentifierReportingModelHandler())) + annotated = [ + t.annotations['model_identifier'].decode() + for t in p.to_runner_api().components.transforms.values() + if t.unique_name == 'RunInference/BeamML_RunInference' + ] + self.assertEqual(len(annotated), 1) + self.assertRegex(annotated[0], r'^[0-9a-f]{32}$') + assert_that(out, equal_to([f'1:{annotated[0]}'])) + + def test_identifier_forwarded_through_wrappers(self): + handler = base.KeyedModelHandler( + IdentifierReportingModelHandler()).with_preprocess_fn( + lambda kv: kv).with_postprocess_fn(lambda kv: kv) + with TestPipeline('FnApiRunner') as p: + out = ( + p | beam.Create([('k', 1)], reshuffle=False) + | base.RunInference(handler, model_identifier='wrapped-id')) + assert_that(out, equal_to([('k', '1:wrapped-id')])) + + class RunInferenceLengthAwareBatchingTest(unittest.TestCase): """End-to-end tests for PR2 length-aware batching in RunInference.""" def test_run_inference_with_length_aware_batch_elements(self): From b4c101c4842a5a1c1bf7eab2d3ef3ebc98095ab6 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 5 Nov 2026 17:46:18 +0000 Subject: [PATCH 2/3] Store model identifier on ModelHandler by default and set it in RunInference.__init__ Address review feedback: ModelHandler.set_model_identifier now stores the value in self._model_identifier (initialized to None in __init__), so subclasses can read it without overriding. RunInference.__init__ reports the identifier at construction time so it is pickled with the handler and inspectable before the pipeline runs; _RunInferenceDoFn.setup() still re-reports it on the worker. --- sdks/python/apache_beam/ml/inference/base.py | 21 +++++++-- .../apache_beam/ml/inference/base_test.py | 43 ++++++++++++++++--- 2 files changed, 55 insertions(+), 9 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index e44d940ed881..049167ec122f 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -251,6 +251,7 @@ def __init__( before loading the model. """ self._env_vars = kwargs.get('env_vars', {}) + self._model_identifier: Optional[str] = None _validate_batching_args( max_batch_weight=max_batch_weight, @@ -368,10 +369,13 @@ def set_model_identifier(self, model_identifier: str): generated when none was passed). This is the value Dataflow's autoscaler uses to find the ``num_loaded_models_`` metric for this step, and it is stable across side-input model updates, so handlers that publish - that metric themselves should tag it with this identifier. The default is - a no-op. + that metric themselves should tag it with this identifier. + + The default implementation stores the value in ``self._model_identifier`` + (initialized to ``None`` in ``__init__``), so subclasses can read it + directly without overriding this method. """ - pass + self._model_identifier = model_identifier def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: """Gets all preprocessing functions to be run before batching/inference. @@ -1139,6 +1143,7 @@ def update_model_path(self, model_path: Optional[str] = None): 'ModelHandlers.') def set_model_identifier(self, model_identifier: str): + super().set_model_identifier(model_identifier) if self._single_model: return self._unkeyed.set_model_identifier(model_identifier) for mh in self._id_to_mh_map.values(): @@ -1271,6 +1276,7 @@ def update_model_path(self, model_path: Optional[str] = None): return self._unkeyed.update_model_path(model_path=model_path) def set_model_identifier(self, model_identifier: str): + super().set_model_identifier(model_identifier) return self._unkeyed.set_model_identifier(model_identifier) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: @@ -1345,6 +1351,7 @@ def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) def set_model_identifier(self, model_identifier: str): + super().set_model_identifier(model_identifier) return self._base.set_model_identifier(model_identifier) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: @@ -1426,6 +1433,7 @@ def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) def set_model_identifier(self, model_identifier: str): + super().set_model_identifier(model_identifier) return self._base.set_model_identifier(model_identifier) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: @@ -1506,6 +1514,7 @@ def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) def set_model_identifier(self, model_identifier: str): + super().set_model_identifier(model_identifier) return self._base.set_model_identifier(model_identifier) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: @@ -1659,6 +1668,12 @@ def __init__( self._model_tag = model_identifier if model_identifier is None: self._model_tag = uuid.uuid4().hex + # Report the identifier at construction time so it is pickled with the + # handler and inspectable on it before the pipeline runs. + # _RunInferenceDoFn.setup() reports it again on the worker in case the + # same handler instance was shared by RunInference transforms with + # different identifiers. + self._model_handler.set_model_identifier(self._model_tag) def annotations(self): extra = {} diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 92b3d2b20d8a..7b4708174b38 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2837,15 +2837,12 @@ class IdentifierReportingModelHandler(base.ModelHandler[int, str, FakeModel]): """ def __init__(self): super().__init__() - self._identifier = None self._identifier_at_load = None - def set_model_identifier(self, model_identifier): - self._identifier = model_identifier - def load_model(self): - # Captures whether set_model_identifier ran before load_model. - self._identifier_at_load = self._identifier + # Relies on the default ModelHandler.set_model_identifier, which stores + # the identifier in self._model_identifier before load_model runs. + self._identifier_at_load = self._model_identifier return FakeModel() def run_inference(self, batch, model, inference_args=None): @@ -2854,6 +2851,40 @@ def run_inference(self, batch, model, inference_args=None): class RunInferenceModelIdentifierTest(unittest.TestCase): """RunInference hands handlers the identifier the autoscaler looks up.""" + def test_default_model_identifier_is_none(self): + self.assertIsNone(IdentifierReportingModelHandler()._model_identifier) + + def test_explicit_model_identifier_set_at_construction(self): + handler = IdentifierReportingModelHandler() + base.RunInference(handler, model_identifier='my-step-model') + self.assertEqual(handler._model_identifier, 'my-step-model') + + def test_generated_model_identifier_set_at_construction(self): + handler = IdentifierReportingModelHandler() + run_inference = base.RunInference(handler) + self.assertRegex(handler._model_identifier, r'^[0-9a-f]{32}$') + self.assertEqual(handler._model_identifier, run_inference._model_tag) + + def test_identifier_set_on_wrappers_at_construction(self): + inner = IdentifierReportingModelHandler() + keyed = base.KeyedModelHandler(inner) + handler = keyed.with_preprocess_fn(lambda kv: kv).with_postprocess_fn( + lambda kv: kv) + base.RunInference(handler, model_identifier='wrapped-id') + for mh in (handler, keyed, inner): + self.assertEqual(mh._model_identifier, 'wrapped-id') + + def test_identifier_forwarded_to_every_per_key_handler(self): + mh1 = IdentifierReportingModelHandler() + mh2 = IdentifierReportingModelHandler() + keyed = base.KeyedModelHandler([ + base.KeyModelMapping(['k1'], mh1), + base.KeyModelMapping(['k2'], mh2), + ]) + base.RunInference(keyed, model_identifier='multi-id') + self.assertEqual(mh1._model_identifier, 'multi-id') + self.assertEqual(mh2._model_identifier, 'multi-id') + def test_explicit_model_identifier_is_reported_before_load(self): with TestPipeline('FnApiRunner') as p: out = ( From 9e3e9e2c5398a31080d01c64845d5e0719c280b5 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 5 Nov 2026 18:41:01 +0000 Subject: [PATCH 3/3] Only set the model identifier in _RunInferenceDoFn.setup() Drop the construction-time call in RunInference.__init__ so the user's handler is not mutated (a handler shared across RunInference transforms would otherwise end up with the last identifier). Wrapper forwarding is now tested by calling set_model_identifier directly. --- sdks/python/apache_beam/ml/inference/base.py | 6 ------ .../apache_beam/ml/inference/base_test.py | 19 ++++--------------- 2 files changed, 4 insertions(+), 21 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 049167ec122f..53bd5540cac9 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1668,12 +1668,6 @@ def __init__( self._model_tag = model_identifier if model_identifier is None: self._model_tag = uuid.uuid4().hex - # Report the identifier at construction time so it is pickled with the - # handler and inspectable on it before the pipeline runs. - # _RunInferenceDoFn.setup() reports it again on the worker in case the - # same handler instance was shared by RunInference transforms with - # different identifiers. - self._model_handler.set_model_identifier(self._model_tag) def annotations(self): extra = {} diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 7b4708174b38..6b8f9998aada 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2854,23 +2854,12 @@ class RunInferenceModelIdentifierTest(unittest.TestCase): def test_default_model_identifier_is_none(self): self.assertIsNone(IdentifierReportingModelHandler()._model_identifier) - def test_explicit_model_identifier_set_at_construction(self): - handler = IdentifierReportingModelHandler() - base.RunInference(handler, model_identifier='my-step-model') - self.assertEqual(handler._model_identifier, 'my-step-model') - - def test_generated_model_identifier_set_at_construction(self): - handler = IdentifierReportingModelHandler() - run_inference = base.RunInference(handler) - self.assertRegex(handler._model_identifier, r'^[0-9a-f]{32}$') - self.assertEqual(handler._model_identifier, run_inference._model_tag) - - def test_identifier_set_on_wrappers_at_construction(self): + def test_identifier_forwarded_through_wrappers(self): inner = IdentifierReportingModelHandler() keyed = base.KeyedModelHandler(inner) handler = keyed.with_preprocess_fn(lambda kv: kv).with_postprocess_fn( lambda kv: kv) - base.RunInference(handler, model_identifier='wrapped-id') + handler.set_model_identifier('wrapped-id') for mh in (handler, keyed, inner): self.assertEqual(mh._model_identifier, 'wrapped-id') @@ -2881,7 +2870,7 @@ def test_identifier_forwarded_to_every_per_key_handler(self): base.KeyModelMapping(['k1'], mh1), base.KeyModelMapping(['k2'], mh2), ]) - base.RunInference(keyed, model_identifier='multi-id') + keyed.set_model_identifier('multi-id') self.assertEqual(mh1._model_identifier, 'multi-id') self.assertEqual(mh2._model_identifier, 'multi-id') @@ -2910,7 +2899,7 @@ def test_generated_uuid_identifier_matches_annotation(self): self.assertRegex(annotated[0], r'^[0-9a-f]{32}$') assert_that(out, equal_to([f'1:{annotated[0]}'])) - def test_identifier_forwarded_through_wrappers(self): + def test_identifier_forwarded_through_wrappers_in_pipeline(self): handler = base.KeyedModelHandler( IdentifierReportingModelHandler()).with_preprocess_fn( lambda kv: kv).with_postprocess_fn(lambda kv: kv)