Merge pull request #25936 from BerriAI/litellm_health-check-reasoning-tokens

fix(proxy): prioritize reasoning health-check max token precedence
2026-04-18 11:35:04 -07:00 · 2026-04-18 11:35:04 -07:00 · d03c301c79
commit d03c301c79
parent f476447875 76c69bb933
5 changed files with 216 additions and 15 deletions
--- a/docs/my-website/docs/proxy/config_settings.md
+++ b/docs/my-website/docs/proxy/config_settings.md
@ -487,7 +487,8 @@ router_settings:
 | AZURE_STORAGE_CLIENT_ID | The Application Client ID to use for Authentication to Azure Blob Storage logging
 | AZURE_STORAGE_CLIENT_SECRET | The Application Client Secret to use for Authentication to Azure Blob Storage logging
 | AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY | Cost per GB per day for Azure Vector Store service
-| BACKGROUND_HEALTH_CHECK_MAX_TOKENS | Optional global default for `max_tokens` on proxy background health checks when a model has no `health_check_max_tokens`. If unset, non-wildcard models default to 1. Applies to wildcard routes when set. Default is unset
+| BACKGROUND_HEALTH_CHECK_MAX_TOKENS | Optional global default for `max_tokens` on proxy background health checks when a model has no `health_check_max_tokens`. If unset, non-wildcard models default to 5. Applies to wildcard routes when set. Default is unset
+| BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING | For **non-wildcard** reasoning models (`supports_reasoning(model)=true`), this takes precedence over `BACKGROUND_HEALTH_CHECK_MAX_TOKENS` when set. If unset, reasoning models fall back to `BACKGROUND_HEALTH_CHECK_MAX_TOKENS` (if set) or default behavior. Wildcard routes ignore this. Default is unset
 | BATCH_STATUS_POLL_INTERVAL_SECONDS | Interval in seconds for polling batch status. Default is 3600 (1 hour)
 | BATCH_STATUS_POLL_MAX_ATTEMPTS | Maximum number of attempts for polling batch status. Default is 24 (for 24 hours)
 | BEDROCK_MAX_POLICY_SIZE | Maximum size for Bedrock policy. Default is 75
--- a/docs/my-website/docs/proxy/health.md
+++ b/docs/my-website/docs/proxy/health.md
@ -338,7 +338,7 @@ model_list:

 ## Health Check Max Tokens

-By default, health checks use `max_tokens=1` to minimize cost and latency. For wildcard models, the default is `max_tokens=10`.
+By default, health checks use `max_tokens=5` to balance reliability with low cost and latency. For wildcard models, the default is `max_tokens=10`.

 You can override this per-model by setting `health_check_max_tokens` in the `model_info` section of your config.yaml.

@ -352,6 +352,30 @@ model_list:
      health_check_max_tokens: 5 # 👈 OVERRIDE HEALTH CHECK MAX TOKENS
 ```

+### Reasoning vs non-reasoning defaults
+
+Reasoning models (per `supports_reasoning` in the model map) often need a higher health-check `max_tokens` because providers count reasoning tokens toward the completion budget. You can set **separate** limits without listing every model:
+
+**Per deployment (`model_info`)** — used when `health_check_max_tokens` is not set. Ignored for wildcard routes (`*` in `litellm_params.model`, i.e. the deployment model string; not `health_check_model`).
+
+```yaml
+model_list:
+  - model_name: openai-stack
+    litellm_params:
+      model: openai/gpt-5-nano
+      api_key: os.environ/OPENAI_API_KEY
+    model_info:
+      health_check_max_tokens_reasoning: 128
+      health_check_max_tokens_non_reasoning: 1
+```
+
+**Global (environment)**:
+
+- `BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING` — for non-wildcard reasoning models, this value takes precedence when set
+- `BACKGROUND_HEALTH_CHECK_MAX_TOKENS` — global fallback for all models (including wildcard routes)
+
+If neither is set, non-wildcard models default to `5` and wildcard routes omit `max_tokens`.
+
 ## `/health/readiness`

 Unprotected endpoint for checking if proxy is ready to accept requests
--- a/litellm/constants.py
+++ b/litellm/constants.py
@ -1360,6 +1360,25 @@ try:
    )
 except (ValueError, TypeError):
    BACKGROUND_HEALTH_CHECK_MAX_TOKENS = None
+
+
+_background_health_check_max_tokens_reasoning_env = os.getenv(
+    "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING"
+)
+try:
+    _raw_background_health_check_max_tokens_reasoning = (
+        _background_health_check_max_tokens_reasoning_env.strip()
+        if _background_health_check_max_tokens_reasoning_env is not None
+        else ""
+    )
+    BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: Optional[int] = (
+        int(_raw_background_health_check_max_tokens_reasoning)
+        if _raw_background_health_check_max_tokens_reasoning
+        else None
+    )
+except (ValueError, TypeError):
+    BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING = None
+
 LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check"
 LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli"
 LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs"
--- a/litellm/proxy/health_check.py
+++ b/litellm/proxy/health_check.py
@ -13,6 +13,7 @@ import litellm
 logger = logging.getLogger(__name__)
 from litellm.constants import (
    BACKGROUND_HEALTH_CHECK_MAX_TOKENS,
+    BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING,
    DEFAULT_HEALTH_CHECK_PROMPT,
    HEALTH_CHECK_TIMEOUT_SECONDS,
 )
@ -292,6 +293,69 @@ def build_deployment_health_states(
    return states


+def _deployment_model_string_for_health_check(litellm_params: dict) -> str:
+    """Deployment model from litellm_params (before Bedrock rewrite).
+
+    Used for reasoning vs non-reasoning max_tokens and wildcard detection only.
+    Does not use ``health_check_model``; that override applies later to the request.
+    """
+    return litellm_params.get("model") or ""
+
+
+def _health_check_deployment_is_wildcard(litellm_params: dict) -> bool:
+    return "*" in _deployment_model_string_for_health_check(litellm_params)
+
+
+def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> Optional[int]:
+    """
+    Pick max_tokens for the health check request.
+
+    Priority:
+    1. model_info.health_check_max_tokens (explicit override)
+    2. For non-wildcard routes: health_check_max_tokens_reasoning / _non_reasoning
+       from model_info based on litellm.supports_reasoning(litellm_params["model"])
+    3. For non-wildcard reasoning routes: BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING
+       from env (if set)
+    4. BACKGROUND_HEALTH_CHECK_MAX_TOKENS (global, any route including wildcards)
+    5. Non-wildcard default: 5
+    6. Wildcard and nothing from (1)(4): leave unset (caller omits max_tokens)
+    """
+    explicit = model_info.get("health_check_max_tokens", None)
+    if explicit is not None:
+        return int(explicit)
+
+    is_wildcard = _health_check_deployment_is_wildcard(litellm_params)
+    deployment_model = _deployment_model_string_for_health_check(litellm_params)
+
+    if not is_wildcard:
+        try:
+            is_reasoning = litellm.supports_reasoning(deployment_model)
+        except Exception:
+            is_reasoning = False
+        tokens_reasoning = model_info.get("health_check_max_tokens_reasoning", None)
+        tokens_non_reasoning = model_info.get(
+            "health_check_max_tokens_non_reasoning", None
+        )
+        if tokens_reasoning is not None or tokens_non_reasoning is not None:
+            if is_reasoning and tokens_reasoning is not None:
+                return int(tokens_reasoning)
+            if not is_reasoning and tokens_non_reasoning is not None:
+                return int(tokens_non_reasoning)
+        if (
+            is_reasoning
+            and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None
+        ):
+            return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING)
+
+    if BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None:
+        return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS)
+
+    if not is_wildcard:
+        return 5
+
+    return None
+
+
 def _update_litellm_params_for_health_check(
    model_info: dict, litellm_params: dict
 ) -> dict:
@ -304,15 +368,9 @@ def _update_litellm_params_for_health_check(
    - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID
    """
    litellm_params["messages"] = _get_random_llm_message()
-    _health_check_max_tokens = model_info.get("health_check_max_tokens", None)
-    if _health_check_max_tokens is not None:
-        litellm_params["max_tokens"] = _health_check_max_tokens
-    elif BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None:
-        litellm_params["max_tokens"] = BACKGROUND_HEALTH_CHECK_MAX_TOKENS
-    elif "*" not in (
-        model_info.get("health_check_model") or litellm_params.get("model") or ""
-    ):
-        litellm_params["max_tokens"] = 1
+    _resolved_max_tokens = _resolve_health_check_max_tokens(model_info, litellm_params)
+    if _resolved_max_tokens is not None:
+        litellm_params["max_tokens"] = _resolved_max_tokens

    _health_check_model = model_info.get("health_check_model", None)
    if _health_check_model is not None:
--- a/tests/test_litellm/proxy/test_health_check_max_tokens.py
+++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py
@ -4,20 +4,25 @@ import pytest

 from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers
 from litellm.proxy import health_check as hc_module
-from litellm.proxy.health_check import _update_litellm_params_for_health_check
+from litellm.proxy.health_check import (
+    _resolve_health_check_max_tokens,
+    _update_litellm_params_for_health_check,
+)


@pytest.mark.asyncio
-async def test_update_litellm_params_max_tokens_default():
+async def test_update_litellm_params_max_tokens_default(monkeypatch):
    """
-    Test that max_tokens defaults to 1 for non-wildcard models.
+    Test that max_tokens defaults to 5 for non-wildcard models.
    """
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None)
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", None)
    model_info = {}
    litellm_params = {"model": "gpt-4"}

    updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)

-    assert updated_params["max_tokens"] == 1
+    assert updated_params["max_tokens"] == 5


@pytest.mark.asyncio
@ -126,3 +131,97 @@ async def test_global_env_var_applies_to_wildcard_models(monkeypatch):
    updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)

    assert updated_params["max_tokens"] == 15
+
+
+def test_resolve_health_check_max_tokens_reasoning_specific_model_info():
+    model_info = {
+        "health_check_max_tokens_reasoning": 64,
+        "health_check_max_tokens_non_reasoning": 2,
+    }
+    litellm_params = {"model": "openai/gpt-4o"}
+
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=False):
+        assert _resolve_health_check_max_tokens(model_info, litellm_params) == 2
+
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=True):
+        assert _resolve_health_check_max_tokens(model_info, litellm_params) == 64
+
+
+def test_explicit_health_check_max_tokens_beats_reasoning_specific():
+    model_info = {
+        "health_check_max_tokens": 9,
+        "health_check_max_tokens_reasoning": 64,
+        "health_check_max_tokens_non_reasoning": 2,
+    }
+    litellm_params = {"model": "openai/gpt-4o"}
+
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=True):
+        assert _resolve_health_check_max_tokens(model_info, litellm_params) == 9
+
+
+def test_reasoning_specific_falls_through_when_wrong_branch_only(monkeypatch):
+    """Only non-reasoning key set but model is reasoning → fall back to default 5."""
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None)
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", None)
+    model_info = {"health_check_max_tokens_non_reasoning": 3}
+    litellm_params = {"model": "openai/o1"}
+
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=True):
+        assert _resolve_health_check_max_tokens(model_info, litellm_params) == 5
+
+
+@pytest.mark.asyncio
+async def test_background_split_env_reasoning_vs_non_reasoning(monkeypatch):
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None)
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", 50)
+
+    model_info = {}
+    litellm_params = {"model": "azure/gpt-4"}
+
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=False):
+        updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+        assert updated["max_tokens"] == 5
+
+    litellm_params2 = {"model": "openai/o1"}
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=True):
+        updated2 = _update_litellm_params_for_health_check(model_info, litellm_params2)
+        assert updated2["max_tokens"] == 50
+
+
+@pytest.mark.asyncio
+async def test_reasoning_env_precedence_over_global(monkeypatch):
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10)
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", 20)
+
+    model_info = {}
+    litellm_params = {"model": "openai/gpt-5.4"}
+
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=True):
+        updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+        assert updated["max_tokens"] == 20
+
+
+@pytest.mark.asyncio
+async def test_non_reasoning_uses_global_when_reasoning_env_set(monkeypatch):
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10)
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", 20)
+
+    model_info = {}
+    litellm_params = {"model": "azure/gpt-4"}
+
+    with patch.object(hc_module.litellm, "supports_reasoning", return_value=False):
+        updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+        assert updated["max_tokens"] == 10
+
+
+def test_wildcard_ignores_reasoning_split_model_info(monkeypatch):
+    """Wildcard routes do not use reasoning/non-reasoning model_info split."""
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None)
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", None)
+    model_info = {
+        "health_check_max_tokens_reasoning": 99,
+        "health_check_max_tokens_non_reasoning": 7,
+    }
+    litellm_params = {"model": "openai/*"}
+
+    assert _resolve_health_check_max_tokens(model_info, litellm_params) is None