Skip to content

[RunInference] Send one Gemini request per input so every element gets a result - #40914

Open
tannnmayy wants to merge 1 commit into
apache:masterfrom
tannnmayy:gemini-one-request-per-prompt
Open

tannnmayy wants to merge 1 commit into
apache:masterfrom
tannnmayy:gemini-one-request-per-prompt

Conversation

@tannnmayy

@tannnmayy tannnmayy commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #40913

GeminiModelHandler sent a whole batch of prompts in one generate_content call. google-genai turns a list of strings into a single user message, so Gemini returned one response for the batch, and request() wrapped it in a one-item list, so only the first element of each batch got a PredictionResult. The other elements were dropped without an error, and the one result answered the merged message.

This PR:

  • sends one generate_content request per input in generate_from_string and generate_image_from_strings_and_images, dispatched on a bounded thread pool (at most 32 in flight per batch), following the OpenAIModelHandler approach from [Gemini] Prototype OpenAI RemoteInferenceModelHandler Implementation #40239;
  • makes request() check that the request function returned one response per input and raise a ValueError otherwise, so a mismatch can't silently drop elements again. A single response is still accepted for a single-element batch;
  • documents that a custom request_fn must return one response per input, in order, and adds a CHANGES.md entry.

Testing:

  • New tests in 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.py passes on Linux (Python 3.14) and Windows (Python 3.11): 21 tests. ruff, yapf and pyrefly report no issues.
  • Real Gemini API (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. The gemini_text_classification.py example 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 default max_batch_size.

R: @jrmccluskey

@github-actions github-actions Bot added python size/m M: 50-249 lines changed labels Oct 30, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📋 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_out helper cleanly bounds thread concurrency per batch (_MAX_CONCURRENT_REQUESTS_PER_BATCH = 32), preserves ordering via executor.map, and cancels pending futures upon error so RemoteModelHandler retry logic can handle failures cleanly.
  • Validation & Backwards Compatibility: Raising a ValueError when the response count does not match the batch size prevents silent element dropping via zip(), 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 RunInference batching.

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @shunping for label python.
Backup reviewer: damccorm (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

  • stop reviewer notifications - opt out of the automated review tooling
  • remind me after tests pass - tag the comment author after tests pass
  • waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

@codecov

codecov Bot commented Oct 30, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.09%. Comparing base (12bbb17) to head (b605b4e).
⚠️ Report is 12 commits behind head on master.

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     
Flag Coverage Δ
python 80.09% <100.00%> (+0.25%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: GeminiModelHandler returns one result per batch; the other prompts in the batch get no output

1 participant