Repository navigation
Conversation
There was a problem hiding this comment.
📋 Review Summary
This PR resolves a critical data-loss bug (#40913) in GeminiModelHandler where passing an entire batch to generate_content caused google-genai to merge the prompts into a single user message, dropping all but the first element's PredictionResult. The fix properly fans out requests per element using a bounded ThreadPoolExecutor (consistent with OpenAIModelHandler), adds strict validation on response counts in request(), and includes comprehensive unit and pipeline tests.
🔍 General Feedback
- Implementation & Concurrency: The
_fan_outhelper cleanly bounds thread concurrency per batch (_MAX_CONCURRENT_REQUESTS_PER_BATCH = 32), preserves ordering viaexecutor.map, and cancels pending futures upon error soRemoteModelHandlerretry logic can handle failures cleanly. - Validation & Backwards Compatibility: Raising a
ValueErrorwhen the response count does not match the batch size prevents silent element dropping viazip(), while still gracefully handling single-element batches returning an unwrapped response object. - Test Coverage: The test suite thoroughly covers per-prompt dispatch, concurrent execution ordering, concurrency capping, error propagation, response count validation, and end-to-end
RunInferencebatching.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
|
Assigning reviewers: R: @shunping 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). |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #40914 +/- ##
==========================================
+ Coverage 79.83% 80.09% +0.25%
==========================================
Files 494 494
Lines 83575 83600 +25
==========================================
+ Hits 66722 66956 +234
+ Misses 16853 16644 -209
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:
|
Fixes #40913
GeminiModelHandlersent a whole batch of prompts in onegenerate_contentcall. google-genai turns a list of strings into a single user message, so Gemini returned one response for the batch, andrequest()wrapped it in a one-item list, so only the first element of each batch got aPredictionResult. The other elements were dropped without an error, and the one result answered the merged message.This PR:
generate_contentrequest per input ingenerate_from_stringandgenerate_image_from_strings_and_images, dispatched on a bounded thread pool (at most 32 in flight per batch), following theOpenAIModelHandlerapproach from [Gemini] Prototype OpenAI RemoteInferenceModelHandler Implementation #40239;request()check that the request function returned one response per input and raise aValueErrorotherwise, so a mismatch can't silently drop elements again. A single response is still accepted for a single-element batch;request_fnmust return one response per input, in order, and adds a CHANGES.md entry.Testing:
gemini_inference_test.py: one request per prompt, concurrency with ordering, the concurrency cap, error propagation, the image request function, the response-count checks, and a pipeline test with batches of 5. 9 of the 10 new tests fail without the fix.gemini_inference_test.pypasses on Linux (Python 3.14) and Windows (Python 3.11): 21 tests. ruff, yapf and pyrefly report no issues.gemini-2.5-flash), batches forced to 3: before the fix, 3 prompts gave 1 result (4 Paris Blue); with the fix, 6 prompts gave 6 correct results. Thegemini_text_classification.pyexample still works.Open questions from the issue: raise vs. warn on a count mismatch, sharing the fan-out helper with
openai_inference.py, and a defaultmax_batch_size.R: @jrmccluskey