Add opt-in max_response_size cap to lf.llms.REST to bound streamed response memory. - #743
Open
copybara-service[bot] wants to merge 1 commit into
Open
Add opt-in max_response_size cap to lf.llms.REST to bound streamed response memory.#743copybara-service[bot] wants to merge 1 commit into
max_response_size cap to lf.llms.REST to bound streamed response memory.#743copybara-service[bot] wants to merge 1 commit into
Conversation
copybara-service
Bot
force-pushed
the
test_943576714
branch
from
July 7, 2026 02:09
6f65b8b to
58f4c82
Compare
…response memory. Streamed LM responses in `REST._read_response_with_deadline` are buffered fully in memory before parsing. The existing time-based bounds (`inactivity_timeout`, `max_response_timeout`) never fire while a server keeps streaming bytes steadily, so a runaway or pathologically large generation can grow the buffer without bound and OOM-kill the process. This adds an opt-in bound, designed for zero regression on this shared library: - New `REST.max_response_size` field (in BYTES). Default `None` preserves the exact historical (uncapped) behavior, so it is a no-op for existing callers. - The stream reader checks the size BEFORE buffering each chunk, so the buffered body never exceeds the cap; it then closes the socket and raises. - New `ResponseSizeLimitError(LMError)`: a NON-retryable error (not a `RetryableLMError`), so retries (e.g. `max_attempts`) do not re-stream the same oversized response. It subclasses `LMError`, so existing `except LMError` handlers still catch it. - Bind-time validation rejects a non-positive `max_response_size`. Covered by new `MaxResponseSizeTest` cases in rest_test.py: dedicated error type, catchability as `LMError`, non-retryability, exact-boundary allowance, strict buffer bound + socket close, no-regression when unset, and validation. PiperOrigin-RevId: 943576714
copybara-service
Bot
force-pushed
the
test_943576714
branch
from
July 7, 2026 16:01
58f4c82 to
887371d
Compare
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.
Add opt-in
max_response_sizecap to lf.llms.REST to bound streamed response memory.Streamed LM responses in
REST._read_response_with_deadlineare buffered fullyin memory before parsing. The existing time-based bounds (
inactivity_timeout,max_response_timeout) never fire while a server keeps streaming bytes steadily,so a runaway or pathologically large generation can grow the buffer without
bound and OOM-kill the process.
This adds an opt-in bound, designed for zero regression on this shared library:
REST.max_response_sizefield (in BYTES). DefaultNonepreserves theexact historical (uncapped) behavior, so it is a no-op for existing callers.
body never exceeds the cap; it then closes the socket and raises.
ResponseSizeLimitError(LMError): a NON-retryable error (not aRetryableLMError), so retries (e.g.max_attempts) do not re-stream the sameoversized response. It subclasses
LMError, so existingexcept LMErrorhandlers still catch it.
max_response_size.Covered by new
MaxResponseSizeTestcases in rest_test.py: dedicated error type,catchability as
LMError, non-retryability, exact-boundary allowance, strictbuffer bound + socket close, no-regression when unset, and validation.