Repository navigation
[RunInference] Honor model updates in also in the first batch. - #40667
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 80.04% 80.07% +0.02%
==========================================
Files 490 490
Lines 82916 82938 +22
==========================================
+ Hits 66373 66409 +36
+ Misses 16543 16529 -14
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). |
| 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