UN-4223 [FIX] Cap a streamed LLM completion at a total time limit

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
@@ -95,6 +96,42 @@ 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 finite 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:
+        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
+
+
 # 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 +594,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 +610,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:

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,19 +320,37 @@ 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
             )
             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
 

unstract/sdk1/tests/test_llm_stream_deadline.py

@@ -0,0 +1,128 @@
+"""``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
+
+
+@pytest.mark.parametrize("raw", ["30m", "inf", "-inf", "nan"])
+def test_invalid_limit_falls_back_to_the_default_with_a_warning(
+    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", raw)
+
+    with caplog.at_level(logging.WARNING, logger=llm_module.logger.name):
+        assert llm_module._stream_max_seconds() == 1800.0
+
+    assert f"LLM_STREAM_MAX_SECONDS={raw!r} is not a finite number" in caplog.text

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,97 @@ 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
+
+
+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()

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