Skip to content

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
mainfrom
test_943576714
Open

Add opt-in max_response_size cap to lf.llms.REST to bound streamed response memory.#743
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_943576714

Conversation

@copybara-service

Copy link
Copy Markdown

Add opt-in max_response_size cap to lf.llms.REST to bound streamed 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.

…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant