Add optional client-side rate limiting (max_rpm) and preserve partial results on parallel failure - #520
Open
ojassharma7 wants to merge 2 commits into
Open
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Addresses the two still-open items of problem #2 in #358 (Gemini
429 RESOURCE_EXHAUSTEDhandling). The exponential backoff that issue also asked for already landed in3aab86c, so this PR covers what remains:max_rpm— client-side throttle. A new optionalGeminiLanguageModelparameter. When> 0, real-time requests are spaced by a thread-safe leaky bucket so at mostmax_rpmstart per minute (shared acrossmax_workers). This prevents429 RESOURCE_EXHAUSTEDon quota-limited tiers (e.g. the free tier's 15 RPM) rather than only retrying after the fact. Defaults to0(disabled) — no behavior change for existing callers. Does not apply to the Batch API.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 toInferenceRuntimeError(partial_results, aligned to the input prompts, plusfailed_indices). The all-or-nothing raise is preserved, soannotation'szip(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_rpmthrottle; 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.pycover the limiter's spacing (mocked clock),max_rpmwiring + validation, and partial-result preservation across the all-success, single-failure, and multi-failure cases.Differential-clean against the full suite: the only failures on my machine are pre-existing and environmental (the optional
openaiextra 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:
max_rpmconstructor docstring; the retry-family params aren't documented outside the docstring either).pylintover the affected code.