damccorm commented on code in PR #40906:
URL: https://github.com/apache/beam/pull/40906#discussion_r4483153830
##########
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:
Instead of making this a `pass` no-op and requiring every subclass (like the
vLLM handlers in #40760) to initialize `self._model_identifier = None` and
override `set_model_identifier`, can we initialize `self._model_identifier:
Optional[str] = None` in `ModelHandler.__init__` and have
`ModelHandler.set_model_identifier` set `self._model_identifier =
model_identifier` by default (similar to how `set_environment_vars` is
implemented on `ModelHandler`)? That way subclasses can just read
`self._model_identifier` directly without overriding `set_model_identifier`.
##########
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:
Can we call `self._model_handler.set_model_identifier(self._model_tag)` in
`RunInference.__init__` right when `self._model_tag` is created (or in both
places)? At that point the user has already constructed and wrapped their
`ModelHandler`, so setting it in `RunInference.__init__` means
`_model_identifier` is pickled with the handler and is also immediately
inspectable on the handler at pipeline construction time (making unit tests
much simpler).
--
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]
