Repository navigation
[Python][RunInference] Report the model identifier to model handlers - #40906
Conversation
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_<identifier> 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 apache#40714. The hook is a no-op by default and is forwarded by the keyed, maybe-keyed, prebatched, preprocessing and postprocessing wrappers.
|
R: @damccorm |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #40906 +/- ##
==========================================
+ Coverage 79.83% 80.11% +0.27%
==========================================
Files 494 494
Lines 83575 83673 +98
==========================================
+ Hits 66722 67034 +312
+ Misses 16853 16639 -214
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| 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) |
There was a problem hiding this 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).
There was a problem hiding this 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 RunInferences with different identifiers, the worker still re-sets the right value from the DoFn's own _model_tag.
There was a problem hiding this comment.
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 RunInferences with different identifiers, the worker still re-sets the right value from the DoFn's own _model_tag.
Oh interesting, I hadn't thought about this modifying the original underlying model handler in that case. Given that, lets just keep this set call instead of having it in both places. Sorry for the churn
There was a problem hiding this comment.
No worries! Updated now.
…ference.__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.
github-actions
Bot
added
size/m
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.
damccorm
left a comment
damccorm
left a comment
There was a problem hiding this comment.
Thanks! Letting tests run and then I will merge
Split out of #40760, following review feedback to separate the generic pieces from the vLLM-specific ones. This PR is a small generic RunInference hook.
Why
Runners that autoscale on
num_loaded_models_<model_identifier>(Dataflow's GPU autoscaling) look the metric up by the identifier RunInference annotates the inference step with: themodel_identifierpassed toRunInference, or a uuid it generates.ModelManagerpublishes under that identifier (and since #40714 keeps it stable across side-input model updates). But a model handler that publishes the metric itself, e.g. one that fronts an out-of-process model server, has no way to learn that value.What
ModelHandler.set_model_identifier(model_identifier): a no-op by default._RunInferenceDoFn.setup()calls it on each worker, before the firstload_model(), with the same tag it annotates the step with and givesModelManager.KeyedModelHandlerforwards to every per-key handler).No behavior change for existing handlers.
Testing
base_test.py:model_identifierreaches the handler beforeload_model().model_identifierannotation onRunInference/BeamML_RunInference.KeyedModelHandlerand pre/post-processing wrappers.Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123orfixes #456).CHANGES.mdwith noteworthy changes.