From aedd5c20ff02d5776c9380945503f4cbe6eca371 Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Mon, 5 Oct 2026 12:52:32 +0530 Subject: [PATCH 1/2] UN-4223 [FIX] Cap a streamed LLM completion at a total time limit A read timeout resets on every chunk, so it never stops a model that keeps generating. In production single Gemini calls streamed for 2+ hours with no timeout, error or retry logged, and memory grew to ~10 GB until pod liveness took whole executor pods offline. collect_with_retry takes an optional max_seconds that caps the whole call, retries included, and raises DeadlineExceededError when it passes; the stream is closed and not replayed. LLM.complete() applies it to every streamed completion: LLM_STREAM_MAX_SECONDS, default 1800 s, 0 disables it. Co-Authored-By: Claude Opus 5.5 --- unstract/sdk1/src/unstract/sdk1/llm.py | 36 ++++- .../src/unstract/sdk1/utils/retry_utils.py | 22 ++- .../sdk1/tests/test_llm_stream_deadline.py | 126 ++++++++++++++++++ .../tests/utils/test_collect_with_retry.py | 78 ++++++++++- workers/sample.env | 4 + 5 files changed, 263 insertions(+), 3 deletions(-) create mode 100644 unstract/sdk1/tests/test_llm_stream_deadline.py diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 2082c24d9a..90bf21e98c 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -95,6 +95,38 @@ def _inject_mock_response(completion_kwargs: dict[str, object]) -> None: completion_kwargs["mock_response"] = mock +# Total wall-clock cap on a streamed completion, retries included. A read +# timeout resets on every chunk, so it never stops a model that keeps +# generating; this does. 30 minutes clears the longest legitimate generations +# seen (~16 min) and stays under the executor's 1 h result timeout. Set to 0 +# to disable. +_STREAM_MAX_SECONDS_ENV = "LLM_STREAM_MAX_SECONDS" +_DEFAULT_STREAM_MAX_SECONDS = 1800.0 + + +@cache +def _warn_invalid_stream_max_seconds(value: str) -> None: + logger.warning( + "%s=%r is not a number; using the default of %.0fs.", + _STREAM_MAX_SECONDS_ENV, + value, + _DEFAULT_STREAM_MAX_SECONDS, + ) + + +def _stream_max_seconds() -> float | None: + """Total time limit for a streamed completion, or None when disabled.""" + raw = os.environ.get(_STREAM_MAX_SECONDS_ENV, "").strip() + if not raw: + return _DEFAULT_STREAM_MAX_SECONDS + try: + value = float(raw) + except ValueError: + _warn_invalid_stream_max_seconds(raw) + return _DEFAULT_STREAM_MAX_SECONDS + return value if value > 0 else None + + # Drop unsupported params rather than raising errors. # Set once at module level instead of per-call to avoid repeated # global mutation in concurrent environments. @@ -557,7 +589,8 @@ def _complete_via_stream( Retries follow ``collect_with_retry``: a failure before the first content chunk is retried like a non-streaming call; a failure after content started is raised as-is so a long generation is never - replayed. ``litellm.stream_chunk_builder`` reassembles content, + replayed. The whole call is capped at ``LLM_STREAM_MAX_SECONDS`` + (default 1800 s). ``litellm.stream_chunk_builder`` reassembles content, thinking blocks, ``finish_reason``, usage (including cache tokens) and the provider response headers. """ @@ -572,6 +605,7 @@ def _complete_via_stream( retry_predicate=is_retryable_litellm_error, is_content=self._chunk_has_content, description=self._get_adapter_info(), + max_seconds=_stream_max_seconds(), ) response = litellm.stream_chunk_builder(chunks, messages=messages) if response is None: diff --git a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py index 794082e218..bf58fbf74e 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py @@ -274,6 +274,10 @@ def iter_with_retry[T]( time.sleep(delay) +class DeadlineExceededError(Exception): + """A drained stream ran past its total wall-clock limit.""" + + def collect_with_retry[T]( fn: Callable[[], Iterable[T]], *, @@ -282,6 +286,7 @@ def collect_with_retry[T]( is_content: Callable[[T], bool], description: str = "", logger_instance: logging.Logger | None = None, + max_seconds: float | None = None, ) -> list[T]: """Drain fn() into a list with retry. Only retries before content. @@ -292,9 +297,16 @@ def collect_with_retry[T]( and is raised immediately: replaying it would re-run a generation that may already have consumed minutes. Items from a failed attempt are discarded so a retry never duplicates them. + + ``max_seconds`` caps the whole call, retries included, and raises + ``DeadlineExceededError`` once it passes. A read timeout cannot do this: + it resets on every item, so a stream that keeps producing (a model stuck + generating) would otherwise run forever. The check runs as each item + arrives; a stream that goes silent is still bounded by the read timeout. """ _validate_max_retries(max_retries) log = logger_instance or logger + deadline = None if max_seconds is None else time.monotonic() + max_seconds for attempt in range(max_retries + 1): items: list[T] = [] has_content = False @@ -308,13 +320,21 @@ def collect_with_retry[T]( for item in gen: items.append(item) has_content = has_content or is_content(item) + if deadline is not None and time.monotonic() > deadline: + raise DeadlineExceededError( + f"stream exceeded its {max_seconds:g}s total time " + f"limit after {len(items)} chunk(s)" + ) return items except Exception as e: # Release the in-flight HTTP/socket resources before retrying. close = getattr(gen, "close", None) if callable(close): close() - if has_content: + if has_content or isinstance(e, DeadlineExceededError): + raise + # Out of time: surface the failure rather than start another attempt. + if deadline is not None and time.monotonic() >= deadline: raise delay = _get_retry_delay( e, attempt, max_retries, retry_predicate, description, log diff --git a/unstract/sdk1/tests/test_llm_stream_deadline.py b/unstract/sdk1/tests/test_llm_stream_deadline.py new file mode 100644 index 0000000000..f81eb3b0ff --- /dev/null +++ b/unstract/sdk1/tests/test_llm_stream_deadline.py @@ -0,0 +1,126 @@ +"""``LLM.complete()`` caps a streamed completion at ``LLM_STREAM_MAX_SECONDS``. + +In production (UN-4223) single ``vertex_ai/gemini-3.1-flash-lite`` calls kept +streaming for 2+ hours: no timeout fired, memory climbed to ~10 GB, and pod +liveness eventually took whole executor pods offline. A read timeout resets on +every chunk, so only a total wall-clock limit stops a model stuck generating. +""" + +from __future__ import annotations + +import logging +import time +from functools import lru_cache +from importlib import import_module +from itertools import cycle +from typing import TYPE_CHECKING +from unittest.mock import patch + +import litellm +import pytest + +if TYPE_CHECKING: + from collections.abc import Iterator + +ANTHROPIC_ADAPTER_ID = "anthropic|90ebd4cd-2f19-4cef-a884-9eeb6ac0f203" +_REAL_COMPLETION = litellm.completion + + +@lru_cache(maxsize=1) +def _load_llm_module() -> object: + import sys + from types import ModuleType + + sys.modules.setdefault("magic", ModuleType("magic")) + return import_module("unstract.sdk1.llm") + + +@pytest.fixture +def no_cost() -> Iterator[None]: + llm_module = _load_llm_module() + with patch.object(llm_module.litellm, "cost_per_token", return_value=(0.0, 0.0)): + yield + + +def _content_chunks() -> list[object]: + """Real litellm stream chunks that carry text, produced without network.""" + chunks = list( + _REAL_COMPLETION( + model="anthropic/claude-sonnet-4-6", + messages=[{"role": "user", "content": "hi"}], + stream=True, + mock_response="loop ", + api_key="test-key", + ) + ) + return [c for c in chunks if c["choices"] and c["choices"][0].delta.content] + + +def test_endless_stream_fails_at_the_limit( + no_cost: None, monkeypatch: pytest.MonkeyPatch +) -> None: + llm_module = _load_llm_module() + llm = llm_module.LLM( + adapter_id=ANTHROPIC_ADAPTER_ID, + adapter_metadata={ + "model": "claude-sonnet-4-6", + "api_key": "test-key", + "max_retries": 2, + }, + ) + chunks = _content_chunks() + calls: list[int] = [] + + def endless(**_: object) -> Iterator[object]: + calls.append(1) + for chunk in cycle(chunks): + time.sleep(0.005) + yield chunk + + monkeypatch.setenv("LLM_STREAM_MAX_SECONDS", "0.05") + with ( + patch.object(llm_module.litellm, "completion", endless), + pytest.raises( + llm_module.LLMError, match="exceeded its 0.05s total time limit" + ) as exc, + ): + llm.complete("hi") + + assert len(calls) == 1 # a runaway generation is never replayed + # The adapter is named once, by LLM's wrapper. + assert str(exc.value).startswith("Error from LLM adapter 'anthropic': stream") + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, 1800.0), + ("", 1800.0), + ("600", 600.0), + (" 900.5 ", 900.5), + ("0", None), + ("-5", None), + ], +) +def test_limit_is_read_from_the_environment( + raw: str | None, expected: float | None, monkeypatch: pytest.MonkeyPatch +) -> None: + llm_module = _load_llm_module() + if raw is None: + monkeypatch.delenv("LLM_STREAM_MAX_SECONDS", raising=False) + else: + monkeypatch.setenv("LLM_STREAM_MAX_SECONDS", raw) + + assert llm_module._stream_max_seconds() == expected + + +def test_invalid_limit_falls_back_to_the_default_with_a_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + llm_module = _load_llm_module() + monkeypatch.setenv("LLM_STREAM_MAX_SECONDS", "30m") + + with caplog.at_level(logging.WARNING, logger=llm_module.logger.name): + assert llm_module._stream_max_seconds() == 1800.0 + + assert "LLM_STREAM_MAX_SECONDS='30m' is not a number" in caplog.text diff --git a/unstract/sdk1/tests/utils/test_collect_with_retry.py b/unstract/sdk1/tests/utils/test_collect_with_retry.py index 974a6e6e00..c1ecf9231f 100644 --- a/unstract/sdk1/tests/utils/test_collect_with_retry.py +++ b/unstract/sdk1/tests/utils/test_collect_with_retry.py @@ -9,12 +9,13 @@ from __future__ import annotations +import time from typing import TYPE_CHECKING from unittest.mock import patch import pytest from unstract.sdk1.utils import retry_utils -from unstract.sdk1.utils.retry_utils import collect_with_retry +from unstract.sdk1.utils.retry_utils import DeadlineExceededError, collect_with_retry if TYPE_CHECKING: from collections.abc import Callable, Iterator @@ -150,3 +151,78 @@ def close(self) -> None: ) assert result == ["c1"] assert closed == [True] + + +# ── max_seconds: a total wall-clock cap ────────────────────────────────────── + + +def _endless(item: str, pause: float, closed: list[bool]) -> Iterator[str]: + """A stream that never ends: a model stuck generating.""" + try: + while True: + time.sleep(pause) + yield item + finally: + closed.append(True) + + +@pytest.mark.parametrize("item", ["c1", "meta"]) +def test_endless_stream_is_stopped_at_the_deadline(item: str) -> None: + """A read timeout never fires while items keep arriving; the deadline does. + + It also fires before any content, and the stream is not retried: a second + attempt would hit the same wall. + """ + closed: list[bool] = [] + calls: list[int] = [] + + def fn() -> Iterator[str]: + calls.append(1) + return _endless(item, pause=0.01, closed=closed) + + with pytest.raises( + DeadlineExceededError, match="exceeded its 0.05s total time limit" + ): + collect_with_retry( + fn, + max_retries=2, + retry_predicate=lambda _: True, + is_content=_is_content, + max_seconds=0.05, + ) + assert len(calls) == 1 + assert closed == [True] + + +def test_no_retry_once_the_deadline_has_passed() -> None: + calls: list[int] = [] + + def fn() -> Iterator[str]: + calls.append(1) + yield "meta" + time.sleep(0.06) + raise TimeoutError("read timed out") + + with pytest.raises(TimeoutError): + collect_with_retry( + fn, + max_retries=2, + retry_predicate=lambda _: True, + is_content=_is_content, + max_seconds=0.05, + ) + assert len(calls) == 1 + + +def test_retries_normally_within_the_deadline() -> None: + fn, calls = _stream_factory([["meta", TimeoutError()], ["meta", "c1"]]) + with patch.object(retry_utils.time, "sleep"): + result = collect_with_retry( + fn, + max_retries=2, + retry_predicate=lambda _: True, + is_content=_is_content, + max_seconds=60, + ) + assert result == ["meta", "c1"] + assert len(calls) == 2 diff --git a/workers/sample.env b/workers/sample.env index b6c1078571..8de8e6956c 100644 --- a/workers/sample.env +++ b/workers/sample.env @@ -325,6 +325,10 @@ EXECUTOR_AUTOSCALE=2,1 EXECUTOR_RESULT_TIMEOUT=3600 EXECUTOR_TASK_TIME_LIMIT=3600 EXECUTOR_TASK_SOFT_TIME_LIMIT=3300 +# Total time limit for one streamed LLM completion, retries included (default +# 1800). Stops a model stuck generating, which a read timeout never does. +# Keep it below EXECUTOR_RESULT_TIMEOUT. 0 disables it. +# LLM_STREAM_MAX_SECONDS=1800 # Notification Worker NOTIFICATION_WORKER_NAME=notification-worker From 2dcc27fa1546a22b152bbc6df45b4d02aeaedde2 Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Mon, 5 Oct 2026 13:06:00 +0530 Subject: [PATCH 2/2] UN-4223 [FIX] Don't retry past the stream deadline; reject non-finite limits A retryable failure just before the limit could sleep through the backoff and then send a new request after the deadline. Raise instead when the backoff would end past it. LLM_STREAM_MAX_SECONDS values "inf" and "nan" parse as floats and would quietly switch the limit off; treat them like any other invalid value and fall back to the default with a warning. Co-Authored-By: Claude Opus 5.5 --- unstract/sdk1/src/unstract/sdk1/llm.py | 7 ++++++- .../src/unstract/sdk1/utils/retry_utils.py | 10 ++++++++++ .../sdk1/tests/test_llm_stream_deadline.py | 8 +++++--- .../tests/utils/test_collect_with_retry.py | 19 +++++++++++++++++++ 4 files changed, 40 insertions(+), 4 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 90bf21e98c..cc34d09a02 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -1,4 +1,5 @@ import logging +import math import os import re from collections.abc import Callable, Generator, Mapping, Sequence @@ -107,7 +108,7 @@ def _inject_mock_response(completion_kwargs: dict[str, object]) -> None: @cache def _warn_invalid_stream_max_seconds(value: str) -> None: logger.warning( - "%s=%r is not a number; using the default of %.0fs.", + "%s=%r is not a finite number; using the default of %.0fs.", _STREAM_MAX_SECONDS_ENV, value, _DEFAULT_STREAM_MAX_SECONDS, @@ -122,6 +123,10 @@ def _stream_max_seconds() -> float | None: try: value = float(raw) except ValueError: + value = math.nan + # ``float()`` also accepts "inf" and "nan", either of which would quietly + # switch the limit off. + if not math.isfinite(value): _warn_invalid_stream_max_seconds(raw) return _DEFAULT_STREAM_MAX_SECONDS return value if value > 0 else None diff --git a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py index bf58fbf74e..0f3bb9d958 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py @@ -341,6 +341,16 @@ def collect_with_retry[T]( ) if delay is None: raise + # A retry that would start after the deadline only spends another + # request on a call that has already run out of time. + if deadline is not None and time.monotonic() + delay >= deadline: + log.warning( + "Not retrying %s: the %gs total time limit would pass " + "during the backoff", + description, + max_seconds, + ) + raise time.sleep(delay) raise RuntimeError("unreachable") # for type-checker: loop always returns or raises diff --git a/unstract/sdk1/tests/test_llm_stream_deadline.py b/unstract/sdk1/tests/test_llm_stream_deadline.py index f81eb3b0ff..96b8782af9 100644 --- a/unstract/sdk1/tests/test_llm_stream_deadline.py +++ b/unstract/sdk1/tests/test_llm_stream_deadline.py @@ -114,13 +114,15 @@ def test_limit_is_read_from_the_environment( assert llm_module._stream_max_seconds() == expected +@pytest.mark.parametrize("raw", ["30m", "inf", "-inf", "nan"]) def test_invalid_limit_falls_back_to_the_default_with_a_warning( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + raw: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: + """A typo, or a non-finite value ``float()`` accepts, keeps the limit on.""" llm_module = _load_llm_module() - monkeypatch.setenv("LLM_STREAM_MAX_SECONDS", "30m") + monkeypatch.setenv("LLM_STREAM_MAX_SECONDS", raw) with caplog.at_level(logging.WARNING, logger=llm_module.logger.name): assert llm_module._stream_max_seconds() == 1800.0 - assert "LLM_STREAM_MAX_SECONDS='30m' is not a number" in caplog.text + assert f"LLM_STREAM_MAX_SECONDS={raw!r} is not a finite number" in caplog.text diff --git a/unstract/sdk1/tests/utils/test_collect_with_retry.py b/unstract/sdk1/tests/utils/test_collect_with_retry.py index c1ecf9231f..08ef71a530 100644 --- a/unstract/sdk1/tests/utils/test_collect_with_retry.py +++ b/unstract/sdk1/tests/utils/test_collect_with_retry.py @@ -226,3 +226,22 @@ def test_retries_normally_within_the_deadline() -> None: ) assert result == ["meta", "c1"] assert len(calls) == 2 + + +def test_no_retry_when_the_backoff_would_pass_the_deadline() -> None: + """A retry that would start after the limit is never sent.""" + fn, calls = _stream_factory([["meta", TimeoutError()], ["meta", "c1"]]) + with ( + patch.object(retry_utils, "calculate_delay", return_value=10.0), + patch.object(retry_utils.time, "sleep") as fake_sleep, + pytest.raises(TimeoutError), + ): + collect_with_retry( + fn, + max_retries=2, + retry_predicate=lambda _: True, + is_content=_is_content, + max_seconds=5, + ) + assert len(calls) == 1 + fake_sleep.assert_not_called()