diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index 47b5247d..218a34bf 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -425,6 +425,14 @@ def litellm_proxy_config( # storms or deployment cooldown fast-fails (#830). "num_retries": 0, "disable_cooldowns": True, + # ...with ONE narrow exception honoring #830's rationale: transient + # upstream 5xx (InternalServerError) are not deterministic rejects, + # and some agents silently swallow a passed-through 500 into a + # zero-event turn (observed: MiMo Code ending turns with no error on + # a deepseek 500, repro x3). Absorb up to 2 transient 500s at the + # proxy; every deterministic class (4xx, auth, content) still fails + # fast with zero retries. + "retry_policy": {"InternalServerErrorRetries": 2}, }, "litellm_settings": { "callbacks": [f"{callback_module}.proxy_handler_instance"], diff --git a/src/benchflow/providers/litellm_retry_patch.py b/src/benchflow/providers/litellm_retry_patch.py new file mode 100644 index 00000000..a12fc2b6 --- /dev/null +++ b/src/benchflow/providers/litellm_retry_patch.py @@ -0,0 +1,102 @@ +"""LiteLLM startup patch for transient upstream 5xx retry policy. + +LiteLLM 1.89.0's ``RetryPolicy`` model accepts +``InternalServerErrorRetries``, but the Router helper that reads retry-policy +fields never checks that attribute. BenchFlow still keeps the proxy-wide +``num_retries=0`` fail-fast default; this patch makes the existing +``InternalServerErrorRetries`` policy work for the one transient 5xx class. + +This module is copied into the per-run LiteLLM runtime and imported by +``sitecustomize``. Keep it standalone: sandbox proxy processes may not be able +to import the BenchFlow package tree. +""" + +from __future__ import annotations + +from typing import Any, cast + + +def _effective_retry_policy( + *, + retry_policy: Any, + model_group: str | None, + model_group_retry_policy: dict[str, Any] | None, +) -> Any: + if ( + model_group_retry_policy is not None + and model_group is not None + and model_group in model_group_retry_policy + ): + return model_group_retry_policy.get(model_group) + return retry_policy + + +def _internal_server_error_retries(policy: Any) -> int | None: + if policy is None: + return None + if isinstance(policy, dict): + value = policy.get("InternalServerErrorRetries") + else: + value = getattr(policy, "InternalServerErrorRetries", None) + if value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _patch_internal_server_error_retry_policy() -> None: + try: + import litellm + import litellm.router as router_mod + from litellm.router_utils import get_retry_from_policy as policy_mod + except Exception: + return + + original = getattr(router_mod, "_get_num_retries_from_retry_policy", None) + if original is None: + original = getattr(policy_mod, "get_num_retries_from_retry_policy", None) + if original is None or getattr(original, "__benchflow_retry_patch__", False): + return + + internal_server_error = getattr(litellm, "InternalServerError", None) + if internal_server_error is None: + return + + def patched_get_num_retries_from_retry_policy( + exception: Exception, + retry_policy: Any = None, + model_group: str | None = None, + model_group_retry_policy: dict[str, Any] | None = None, + ) -> int | None: + retries = original( + exception=exception, + retry_policy=retry_policy, + model_group=model_group, + model_group_retry_policy=model_group_retry_policy, + ) + if retries is not None: + return retries + if not isinstance(exception, internal_server_error): + return None + policy = _effective_retry_policy( + retry_policy=retry_policy, + model_group=model_group, + model_group_retry_policy=model_group_retry_policy, + ) + return _internal_server_error_retries(policy) + + patched_any = cast(Any, patched_get_num_retries_from_retry_policy) + patched_any.__benchflow_retry_patch__ = True + policy_module_any = cast(Any, policy_mod) + router_module_any = cast(Any, router_mod) + policy_module_any.get_num_retries_from_retry_policy = ( + patched_get_num_retries_from_retry_policy + ) + router_module_any._get_num_retries_from_retry_policy = ( + patched_get_num_retries_from_retry_policy + ) + + +_patch_internal_server_error_retry_policy() diff --git a/src/benchflow/providers/litellm_retry_preflight.py b/src/benchflow/providers/litellm_retry_preflight.py new file mode 100644 index 00000000..3c7273ac --- /dev/null +++ b/src/benchflow/providers/litellm_retry_preflight.py @@ -0,0 +1,123 @@ +"""Fail-closed checks for BenchFlow's LiteLLM 5xx retry patch.""" + +from __future__ import annotations + +import shlex +import subprocess +from typing import Any + +from benchflow.providers.litellm_bedrock_preflight import ( + _exec_details, + _host_python_for_litellm, +) + + +class RetryPatchPreflightError(RuntimeError): + """Raised when the transient-5xx retry patch cannot be proven active.""" + + +RETRY_PATCH_PREFLIGHT_SOURCE = """\ +import sys + +failures = [] +try: + import litellm + import litellm.router as router_mod + + helper = getattr(router_mod, "_get_num_retries_from_retry_policy", None) + if not getattr(helper, "__benchflow_retry_patch__", False): + failures.append("transient-5xx retry helper is unpatched") + else: + policy = {"InternalServerErrorRetries": 2} + transient = litellm.InternalServerError( + "upstream 500", + llm_provider="openai", + model="openai/test", + ) + retries = helper(exception=transient, retry_policy=policy) + if retries != 2: + failures.append( + f"InternalServerError retry probe returned {retries!r}, expected 2" + ) + + permanent = litellm.BadRequestError( + "bad request", + llm_provider="openai", + model="openai/test", + ) + permanent_retries = helper(exception=permanent, retry_policy=policy) + if permanent_retries is not None: + failures.append( + "BadRequestError retry probe returned " + f"{permanent_retries!r}, expected fail-fast None" + ) +except Exception as exc: # noqa: BLE001 - report any LiteLLM shape drift + failures.append(f"retry preflight probe failed: {exc}") + +if failures: + print("; ".join(failures)) + sys.exit(1) +""" + + +def _failure_message(runtime: str, detail: str) -> str: + return ( + "BenchFlow LiteLLM transient-5xx retry patch is NOT active in the " + f"{runtime} LiteLLM runtime: {detail[:2000]}. Failing closed before " + "agent launch - a silent fallback would expose transient upstream " + "500s to agents instead of retrying them at the proxy." + ) + + +def preflight_host_retry_patch( + *, + env: dict[str, str], + litellm_executable: str, +) -> None: + """Fail closed if the retry patch is not active for the host proxy.""" + try: + result = subprocess.run( + [ + _host_python_for_litellm(litellm_executable, env=env), + "-c", + RETRY_PATCH_PREFLIGHT_SOURCE, + ], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + except Exception as exc: + raise RetryPatchPreflightError( + _failure_message("host", f"preflight execution failed: {exc}") + ) from exc + if result.returncode != 0: + detail = (result.stdout or "").strip() or (result.stderr or "").strip() + raise RetryPatchPreflightError(_failure_message("host", detail)) + + +async def preflight_sandbox_retry_patch( + sandbox: Any, + *, + python: str, + runtime_dir: str, + preflight_path: str, +) -> None: + """Fail closed if the retry patch is not active in a sandbox proxy.""" + command = ( + f"PYTHONPATH={shlex.quote(runtime_dir)} " + f"{shlex.quote(python)} {shlex.quote(preflight_path)}" + ) + try: + result = await sandbox.exec(command, timeout_sec=120) + except Exception as exc: + raise RetryPatchPreflightError( + _failure_message("sandbox", f"preflight execution failed: {exc}") + ) from exc + if result.return_code != 0: + raise RetryPatchPreflightError( + _failure_message( + "sandbox", + _exec_details("retry patch preflight", result), + ) + ) diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 1f53f442..8aee6cfc 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -48,6 +48,11 @@ extract_usage_from_trajectory, trajectory_from_litellm_callback_log, ) +from benchflow.providers.litellm_retry_preflight import ( + RETRY_PATCH_PREFLIGHT_SOURCE, + preflight_host_retry_patch, + preflight_sandbox_retry_patch, +) from benchflow.sandbox.providers import OFF_BOX_MODEL_PROVIDERS from benchflow.trajectories.types import Trajectory from benchflow.usage_tracking import UsageTrackingConfig, usage_unavailable @@ -57,7 +62,8 @@ LITELLM_VERSION_SPEC = "litellm[proxy]==1.89.0" LITELLM_SANDBOX_ROOT = "/tmp/benchflow-litellm" _CALLBACK_MODULE = "benchflow_litellm_callback" -_PATCH_MODULE = "benchflow_litellm_bedrock_patch" +_BEDROCK_PATCH_MODULE = "benchflow_litellm_bedrock_patch" +_RETRY_PATCH_MODULE = "benchflow_litellm_retry_patch" # The proxy is an internal single-route gateway — it must never register the # FastAPI Swagger docs route. litellm's `_get_docs_url()` honours an inherited @@ -409,13 +415,18 @@ def _write_runtime_files( ) -> tuple[Path, Path, Path]: runtime_dir.mkdir(parents=True, exist_ok=True) callback_path = runtime_dir / f"{_CALLBACK_MODULE}.py" - patch_path = runtime_dir / f"{_PATCH_MODULE}.py" + patch_path = runtime_dir / f"{_BEDROCK_PATCH_MODULE}.py" + retry_patch_path = runtime_dir / f"{_RETRY_PATCH_MODULE}.py" sitecustomize_path = runtime_dir / "sitecustomize.py" config_path = runtime_dir / "config.yaml" callback_path.write_text(callback_module_source()) patch_source = Path(__file__).with_name("litellm_bedrock_patch.py").read_text() patch_path.write_text(patch_source) - sitecustomize_path.write_text(f"import {_PATCH_MODULE}\n") + retry_patch_source = Path(__file__).with_name("litellm_retry_patch.py").read_text() + retry_patch_path.write_text(retry_patch_source) + sitecustomize_path.write_text( + f"import {_BEDROCK_PATCH_MODULE}\nimport {_RETRY_PATCH_MODULE}\n" + ) config_path.write_text(yaml.safe_dump(config, sort_keys=False)) return config_path, callback_path, patch_path @@ -536,6 +547,11 @@ async def _start_host_litellm( ) try: await _poll_host_health(runner) + await asyncio.to_thread( + preflight_host_retry_patch, + env=env, + litellm_executable=litellm_executable, + ) if route_requires_bedrock_patch(route): # Fail closed before the agent launches when the Bedrock 4.8+ # thinking patch did not activate (#602). @@ -626,7 +642,8 @@ async def _upload_runtime_files_to_sandbox( paths = { "config": f"{runtime_dir}/config.yaml", "callback": f"{runtime_dir}/{_CALLBACK_MODULE}.py", - "patch": f"{runtime_dir}/{_PATCH_MODULE}.py", + "patch": f"{runtime_dir}/{_BEDROCK_PATCH_MODULE}.py", + "retry_patch": f"{runtime_dir}/{_RETRY_PATCH_MODULE}.py", "sitecustomize": f"{runtime_dir}/sitecustomize.py", "launcher": f"{runtime_dir}/launcher.py", "stdout": f"{runtime_dir}/stdout.log", @@ -637,6 +654,7 @@ async def _upload_runtime_files_to_sandbox( "venv": f"{runtime_dir}/venv", "launch_config": f"{runtime_dir}/launch_config.json", "preflight": f"{runtime_dir}/bedrock_patch_preflight.py", + "retry_preflight": f"{runtime_dir}/retry_patch_preflight.py", } result = await sandbox.exec(f"mkdir -p {shlex.quote(runtime_dir)}", timeout_sec=20) if result.return_code != 0: @@ -652,12 +670,24 @@ async def _upload_runtime_files_to_sandbox( ".py", ) await _upload_text( - sandbox, f"import {_PATCH_MODULE}\n", paths["sitecustomize"], ".py" + sandbox, + Path(__file__).with_name("litellm_retry_patch.py").read_text(), + paths["retry_patch"], + ".py", + ) + await _upload_text( + sandbox, + f"import {_BEDROCK_PATCH_MODULE}\nimport {_RETRY_PATCH_MODULE}\n", + paths["sitecustomize"], + ".py", ) await _upload_text(sandbox, _sandbox_launcher_source(), paths["launcher"], ".py") await _upload_text( sandbox, BEDROCK_PATCH_PREFLIGHT_SOURCE, paths["preflight"], ".py" ) + await _upload_text( + sandbox, RETRY_PATCH_PREFLIGHT_SOURCE, paths["retry_preflight"], ".py" + ) return paths @@ -838,6 +868,12 @@ async def _start_sandbox_litellm( port=port, stderr_path=paths["stderr"], ) + await preflight_sandbox_retry_patch( + sandbox, + python=python, + runtime_dir=runtime_dir, + preflight_path=paths["retry_preflight"], + ) if route_requires_bedrock_patch(route): # Fail closed before the agent launches when the Bedrock 4.8+ # thinking patch did not activate (#602). On Daytona this proxy is diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index 6679c097..1c19ea1b 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -181,6 +181,10 @@ def test_proxy_config_registers_plain_and_openai_aliases(): assert config["router_settings"] == { "num_retries": 0, "disable_cooldowns": True, + # transient-5xx-only exception to the no-retry rule (#830 honored: + # deterministic classes keep zero retries) — absorbs upstream 500s that + # some agents (MiMo) silently swallow into zero-event turns. + "retry_policy": {"InternalServerErrorRetries": 2}, } diff --git a/tests/test_litellm_retry_patch.py b/tests/test_litellm_retry_patch.py new file mode 100644 index 00000000..16f94dab --- /dev/null +++ b/tests/test_litellm_retry_patch.py @@ -0,0 +1,273 @@ +"""Regression coverage for the BenchFlow LiteLLM retry patch.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + +from benchflow.providers import litellm_retry_preflight as preflight_mod +from benchflow.providers import litellm_runtime as runtime_mod +from benchflow.providers.litellm_config import resolve_litellm_route + + +def test_runtime_retry_patch_retries_only_transient_5xx(tmp_path): + """Guards PR #882: LiteLLM Router retries upstream 500s, but not 400s.""" + runtime_mod._write_runtime_files(tmp_path, config={"model_list": []}) + env = dict(os.environ) + env["PYTHONPATH"] = f"{tmp_path}{os.pathsep}{env.get('PYTHONPATH', '')}" + probe = textwrap.dedent( + r""" + import asyncio + + import litellm + from litellm import Router + from litellm.utils import ModelResponse + + + async def healthy_deployments(*args, **kwargs): + return ([], [{"litellm_params": {}}]) + + + def router(): + routed = Router( + model_list=[ + { + "model_name": "benchflow-test", + "litellm_params": { + "model": "openai/test", + "api_key": "sk-test", + }, + } + ], + num_retries=0, + retry_policy={"InternalServerErrorRetries": 2}, + disable_cooldowns=True, + retry_after=0, + ) + routed._async_get_healthy_deployments = healthy_deployments + routed._time_to_sleep_before_retry = lambda **kwargs: 0 + return routed + + + async def main(): + routed = router() + calls = 0 + + async def transient_then_success(original_function, *args, **kwargs): + nonlocal calls + calls += 1 + if calls < 3: + raise litellm.InternalServerError( + "upstream 500", + llm_provider="openai", + model="openai/test", + ) + return ModelResponse(model="openai/test", choices=[]) + + routed.make_call = transient_then_success + response = await routed.async_function_with_retries( + original_function=lambda: None, + model="benchflow-test", + messages=[], + metadata={}, + num_retries=0, + ) + assert calls == 3, calls + headers = response._hidden_params["additional_headers"] + assert headers["x-litellm-attempted-retries"] == 2, headers + assert headers["x-litellm-max-retries"] == 2, headers + + async def assert_fails_fast(error_factory, expected_type): + routed = router() + calls = 0 + + async def fail(original_function, *args, **kwargs): + nonlocal calls + calls += 1 + raise error_factory() + + routed.make_call = fail + try: + await routed.async_function_with_retries( + original_function=lambda: None, + model="benchflow-test", + messages=[], + metadata={}, + num_retries=0, + ) + except expected_type: + pass + else: + raise AssertionError(f"{expected_type.__name__} should fail fast") + assert calls == 1, calls + + await assert_fails_fast( + lambda: litellm.BadRequestError( + "bad request", + llm_provider="openai", + model="openai/test", + ), + litellm.BadRequestError, + ) + await assert_fails_fast( + lambda: litellm.ContextWindowExceededError( + "context window exceeded", + llm_provider="openai", + model="openai/test", + ), + litellm.ContextWindowExceededError, + ) + + + asyncio.run(main()) + """ + ) + result = subprocess.run( + [sys.executable, "-c", probe], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_retry_patch_preflight_passes_when_runtime_files_on_pythonpath(tmp_path): + """Guards PR #882: proxy startup must prove the retry monkeypatch loaded.""" + runtime_mod._write_runtime_files(tmp_path, config={"model_list": []}) + env = dict(os.environ) + env["PYTHONPATH"] = str(tmp_path) + result = subprocess.run( + [sys.executable, "-c", preflight_mod.RETRY_PATCH_PREFLIGHT_SOURCE], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_retry_patch_preflight_fails_closed_when_patch_not_loaded(tmp_path): + """Guards PR #882 against silently booting a proxy without 5xx retry behavior.""" + env = dict(os.environ) + env.pop("PYTHONPATH", None) + result = subprocess.run( + [sys.executable, "-c", preflight_mod.RETRY_PATCH_PREFLIGHT_SOURCE], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode != 0 + assert "unpatched" in result.stdout + + +def test_host_retry_preflight_uses_litellm_runtime_python(tmp_path): + """Guards PR #882 host proxy startup: preflight runs in the LiteLLM runtime.""" + runtime_mod._write_runtime_files(tmp_path, config={"model_list": []}) + litellm = tmp_path / "litellm" + litellm.write_text(f"#!{sys.executable}\n") + env = dict(os.environ) + env["PYTHONPATH"] = str(tmp_path) + + preflight_mod.preflight_host_retry_patch( + env=env, + litellm_executable=str(litellm), + ) + + +class _ExecResult: + def __init__(self, return_code: int = 0, stdout: str = "", stderr: str = ""): + self.return_code = return_code + self.stdout = stdout + self.stderr = stderr + + +class _FakeSandbox: + def __init__(self, *, fail_retry_preflight: bool = False): + self.uploaded: dict[str, str] = {} + self.exec_calls: list[str] = [] + self.exec_timeouts: list[int | None] = [] + self.fail_retry_preflight = fail_retry_preflight + self._started = False + + async def upload_file(self, local_path, remote_path) -> None: + self.uploaded[str(remote_path)] = Path(local_path).read_text() + + async def exec(self, command: str, timeout_sec: int | None = None) -> _ExecResult: + self.exec_calls.append(command) + self.exec_timeouts.append(timeout_sec) + if "retry_patch_preflight.py" in command: + if self.fail_retry_preflight: + return _ExecResult(1, stdout="transient-5xx retry helper is unpatched") + return _ExecResult(0) + if "urllib.request" in command: + return _ExecResult(0) + if "launcher.py" in command: + self._started = True + return _ExecResult(0) + if command.strip().startswith("cat") and "state.json" in command: + if self._started: + return _ExecResult(0, stdout='{"pid": 4242, "port": 45999}') + return _ExecResult(0, stdout="") + if command.strip().startswith("rm -rf"): + return _ExecResult(0) + return _ExecResult(0) + + +@pytest.mark.asyncio +async def test_sandbox_litellm_runs_retry_preflight_before_returning_proxy(): + """Guards PR #882: sandbox proxy startup probes the retry patch first.""" + route = resolve_litellm_route( + "minimax/MiniMax-M3", + {"MINIMAX_API_KEY": "k", "MINIMAX_BASE_URL": "https://api.minimax.io/v1"}, + ) + sandbox = _FakeSandbox() + + proc = await runtime_mod._start_sandbox_litellm( + sandbox=sandbox, + route=route, + master_key="sk-master", + agent_env={ + "MINIMAX_API_KEY": "k", + "MINIMAX_BASE_URL": "https://api.minimax.io/v1", + }, + session_id="s", + agent_name="openhands", + ) + + assert proc.base_url == "http://127.0.0.1:45999" + retry_preflights = [ + c for c in sandbox.exec_calls if "retry_patch_preflight.py" in c + ] + assert retry_preflights, "sandbox LiteLLM must fail-close probe retry patch" + + +@pytest.mark.asyncio +async def test_sandbox_litellm_retry_preflight_failure_fails_closed(): + """Guards PR #882: inactive sandbox retry patch tears down the proxy.""" + route = resolve_litellm_route( + "minimax/MiniMax-M3", + {"MINIMAX_API_KEY": "k", "MINIMAX_BASE_URL": "https://api.minimax.io/v1"}, + ) + sandbox = _FakeSandbox(fail_retry_preflight=True) + + with pytest.raises(RuntimeError, match="transient-5xx retry patch is NOT active"): + await runtime_mod._start_sandbox_litellm( + sandbox=sandbox, + route=route, + master_key="sk-master", + agent_env={ + "MINIMAX_API_KEY": "k", + "MINIMAX_BASE_URL": "https://api.minimax.io/v1", + }, + session_id="s", + agent_name="openhands", + ) + + assert any(call.strip().startswith("rm -rf") for call in sandbox.exec_calls)