Repository navigation
Preserve ModelManager metric_tag across RunInference side-input model updates - #40714
Merged
tvalentyn merged 4 commits intoOct 23, 2026
Merged
Conversation
AMOOOMA
marked this pull request as ready for review
October 21, 2026 21:25
Contributor
Author
|
R: @tvalentyn |
Contributor
|
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❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #40714 +/- ##
========================================
Coverage 80.04% 80.05%
========================================
Files 490 490
Lines 82916 83031 +115
========================================
+ Hits 66373 66467 +94
- Misses 16543 16564 +21
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:
|
tvalentyn
reviewed
Oct 21, 2026
| self.models = models | ||
| self.use_model_manager = not isinstance(models, list) | ||
| self.model_tag = model_tag | ||
| self.metric_tag = metric_tag or model_tag |
Contributor
There was a problem hiding this comment.
could we clarify in the docstring the difference between these two concepts?
…te; document model_tag vs metric_tag
…note that model updates should only change weights
Contributor
|
is this ready to merge? thanks! |
Contributor
Author
Yep thanks! |
tvalentyn
approved these changes
Oct 23, 2026
2 of 3 tasks
AMOOOMA
added a commit
to AMOOOMA/beam
that referenced
this pull request
Oct 30, 2026
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.
AMOOOMA
added a commit
to AMOOOMA/beam
that referenced
this pull request
Nov 3, 2026
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.
AMOOOMA
added a commit
to AMOOOMA/beam
that referenced
this pull request
Nov 5, 2026
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.
damccorm
pushed a commit
that referenced
this pull request
Nov 5, 2026
…40906) * [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_<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 #40714. The hook is a no-op by default and is forwarded by the keyed, maybe-keyed, prebatched, preprocessing and postprocessing wrappers. * Store model identifier on ModelHandler by default and set it in RunInference.__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. * Only set the model identifier in _RunInferenceDoFn.setup() 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
When
RunInferenceupdates its model via side input (ModelMetadata),_RunInferenceDoFn._load_modelacquires the new model under the updatedmodel_idtag (self._cur_tag) soModelManagercaches and evicts the new model version separately from the previous version, while theBeamML_RunInferencetransform annotation remains fixed at construction time (self._model_tag/model_identifier).This PR passes
metric_tag=self._model_tagthrough_SharedModelWrappertoModelManager.acquire_model(..., metric_tag=...), and mapsmetric_tagto the active version tag inModelManagerandResourceEstimatorsonum_loaded_models_<metric_tag>andmemory_estimate_mb_<metric_tag>continue reporting under the originalmetric_tagacross runtime side-input model updates.Testing
test_run_inference_dofn_model_update_acquires_and_releases_matching_taginbase_test.pyto verifymetric_tagpreservation across model updates.test_model_manager_metrics_with_metric_taginmodel_manager_test.py.