From c29a70fe5e6a3878ff7eb4aa114bf1a8316ee114 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Fri, 30 Oct 2026 16:32:40 +0000 Subject: [PATCH] [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):