AMOOOMA commented on code in PR #40906:
URL: https://github.com/apache/beam/pull/40906#discussion_r4483937716
##########
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_<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 is
+ a no-op.
+ """
+ pass
Review Comment:
Done in b4c101c. `ModelHandler.__init__` now initializes
`self._model_identifier: Optional[str] = None`, and the default
`set_model_identifier` stores the value there (mirroring
`set_environment_vars`), so subclasses can just read `self._model_identifier`.
The wrappers (keyed, maybe-keyed, prebatched, pre/postprocessing) also store it
on themselves before forwarding, so the outermost handler is inspectable too.
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -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)
Review Comment:
Done in b4c101c, in both places. `RunInference.__init__` now calls
`set_model_identifier(self._model_tag)` right after the tag is created, so it
is pickled with the handler and visible at construction time. I added
construction-time tests for the explicit id, the generated uuid (==
`_model_tag`), wrapper chains, and multi-model `KeyedModelHandler`. I kept the
call in `setup()` as well, since the construction-time call mutates the user's
handler: if one handler instance is shared across `RunInference`s with
different identifiers, the worker still re-sets the right value from the DoFn's
own `_model_tag`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
