Skip to content

Add optional client-side rate limiting (max_rpm) and preserve partial results on parallel failure - #520

Open
ojassharma7 wants to merge 2 commits into
google:mainfrom
ojassharma7:feat/gemini-throttle-and-partial-results
Open

Add optional client-side rate limiting (max_rpm) and preserve partial results on parallel failure#520
ojassharma7 wants to merge 2 commits into
google:mainfrom
ojassharma7:feat/gemini-throttle-and-partial-results

Conversation

@ojassharma7

@ojassharma7 ojassharma7 commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Description

Addresses the two still-open items of problem #2 in #358 (Gemini 429 RESOURCE_EXHAUSTED handling). The exponential backoff that issue also asked for already landed in 3aab86c, so this PR covers what remains:

  • max_rpm — client-side throttle. A new optional GeminiLanguageModel parameter. When > 0, real-time requests are spaced by a thread-safe leaky bucket so at most max_rpm start per minute (shared across max_workers). This prevents 429 RESOURCE_EXHAUSTED on quota-limited tiers (e.g. the free tier's 15 RPM) rather than only retrying after the fact. Defaults to 0 (disabled) — no behavior change for existing callers. Does not apply to the Batch API.
  • Preserve completed chunks on failure. The parallel infer() path raised on the first failed chunk, which abandoned the still-running futures and discarded every chunk that had already succeeded. It now lets the pool drain, then raises once with the completed work attached to InferenceRuntimeError (partial_results, aligned to the input prompts, plus failed_indices). The all-or-nothing raise is preserved, so annotation's zip(batch, outputs) alignment is unaffected; callers who want to recover already-done work can now read it off the exception.

This is a partial fix for #358 (chunking for very large documents is a separate design question, not touched here), so:

Related to #358

Type: Feature (the max_rpm throttle; the partial-result change also fixes a real loss-of-work bug).

How Has This Been Tested?

Unit tests in tests/gemini_ratelimit_partial_test.py cover the limiter's spacing (mocked clock), max_rpm wiring + validation, and partial-result preservation across the all-success, single-failure, and multi-failure cases.

$ python -m pytest tests/gemini_ratelimit_partial_test.py tests/gemini_retry_test.py

Differential-clean against the full suite: the only failures on my machine are pre-existing and environmental (the optional openai extra isn't installed, plus 3 plugin-packaging tests) — identical with and without this change — and the 10 new tests pass on top. pyink + isort + pylint --rcfile=.pylintrc / --rcfile=tests/.pylintrc (CI's exact configs) all clean.

Checklist:

…n failure

Addresses the two remaining items of problem google#2 in google#358 (Gemini 429 handling);
the exponential backoff it also asked for already landed in 3aab86c.

max_rpm: a new optional GeminiLanguageModel parameter that spaces real-time
requests client-side so at most N start per minute, shared across workers via a
thread-safe leaky bucket. It prevents 429 RESOURCE_EXHAUSTED on quota-limited
tiers (e.g. the free tier's 15 RPM) rather than only retrying after the fact.
Defaults to 0 (disabled), so existing behavior is unchanged. Not applied to the
Batch API.

Partial results: the parallel infer() path raised on the first failed chunk,
abandoning the still-running futures and discarding every chunk that had already
succeeded. It now drains the pool, then raises once with the completed work
attached to InferenceRuntimeError (partial_results, aligned to the input
prompts, plus failed_indices), so a caller can recover instead of losing a whole
batch. The all-or-nothing raise is preserved, so annotation's zip(batch,
outputs) alignment is unaffected.

Adds tests for the limiter's spacing, max_rpm wiring/validation, and partial
preservation across the success, single-failure and multi-failure cases.
fake_time.sleep.side_effect = lambda s: sleeps.append(s) -> sleeps.append.
No behavior change; CI's lint-tests job runs a stricter pylintrc than the
root one I checked before the first push.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size/M Pull request with 150-600 lines changed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant