Repository navigation
Conversation
There was a problem hiding this comment.
📋 Review Summary
This PR resolves a bug (#40665) where RunInference with per-key models (KeyedModelHandler) failed to apply KeyModelPathMapping side-input updates on the first batch when load_model_at_runtime=True. It also avoids redundant update_model calls on every batch when the side input is empty or the per-key model path mapping has not changed.
🔍 General Feedback
- The fix in
_RunInferenceDoFn._load_modelandprocessis clean, well-tested, and properly addresses both the first-batch stale model bug and redundantupdate_modelinvocations. - One minor edge case was noted in
processwhen usingModelMetadataside inputs together with timeout-based model invalidation, along with a minorCHANGES.mdtemplate placeholder removal.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
|
/gemini review |
There was a problem hiding this comment.
📋 Review Summary
This PR fixes a bug (#40665) where RunInference with per-key models (KeyedModelHandler) failed to apply KeyModelPathMapping side-input updates on the first batch when load_model_at_runtime=True. It also prevents redundant update_model calls on every batch when side inputs are empty or unchanged, while properly handling model reloading across all branches when a model is marked invalid.
🔍 General Feedback
- The updates in
_RunInferenceDoFn._load_modeland_RunInferenceDoFn.processare clean, consistent across all side-input branches, and well-covered by the new unit tests.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #40667 +/- ##
=============================================
+ Coverage 56.16% 80.02% +23.85%
=============================================
Files 1124 490 -634
Lines 178336 82810 -95526
Branches 1489 0 -1489
=============================================
- Hits 100170 66267 -33903
+ Misses 75645 16543 -59102
+ Partials 2521 0 -2521
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:
|
|
Assigning reviewers: R: @jrmccluskey for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
| # The side input value that the currently loaded model reflects: a model | ||
| # id from ModelMetadata, a list[KeyModelPathMapping], or None if no side | ||
| # input value has been applied. Used to skip redundant model updates. |
There was a problem hiding this comment.
This a pretty verbose comment that could be summarized by applying a type annotation to the _side_input_path field instead. This should also make the check later in the file a little easier to understand.
| needs_update = ( | ||
| self._model is None or # No model loaded yet. | ||
| self._side_input_path != si_model_metadata.model_id or # New model. | ||
| not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. |
There was a problem hiding this comment.
The needs_update boolean evaluation is repeated multiple times with only slight variations based on si_model_metadata, consider a refactor into a helper (which also streamlines some of the repeated comments)
Honor model updates in first batch.
Don't initiate load model calls when there are no updates. In another branch, do an update if the model was invalidated.
Drive-by entries to changes.md from earlier PRs.
Fixes: #40665