[Python][RunInference] Report the model identifier to model handlers

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,
@@ -360,6 +361,22 @@ 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_<identifier>`` 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 implementation stores the value in ``self._model_identifier``
+    (initialized to ``None`` in ``__init__``), so subclasses can read it
+    directly without overriding this method.
+    """
+    self._model_identifier = model_identifier
+
   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 +1142,13 @@ 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):
+    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():
+      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 +1275,10 @@ 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):
+    super().set_model_identifier(model_identifier)
+    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 +1350,10 @@ 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):
+    super().set_model_identifier(model_identifier)
+    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 +1432,10 @@ 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):
+    super().set_model_identifier(model_identifier)
+    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 +1513,10 @@ 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):
+    super().set_model_identifier(model_identifier)
+    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 +2339,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:

sdks/python/apache_beam/ml/inference/base_test.py

@@ -2829,6 +2829,87 @@ 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_at_load = None
+
+  def load_model(self):
+    # 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):
+    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_default_model_identifier_is_none(self):
+    self.assertIsNone(IdentifierReportingModelHandler()._model_identifier)
+
+  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)
+    handler.set_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),
+    ])
+    keyed.set_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 = (
+          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_in_pipeline(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):