Merge pull request #24773 from Sameerlite/Sameerlite/healthcheck-max-tokens

feat(health-check): add BACKGROUND_HEALTH_CHECK_MAX_TOKENS env var

litellm/constants.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import List, Literal
+from typing import List, Literal, Optional
 
 from litellm.litellm_core_utils.env_utils import get_env_int
 
@@ -1318,6 +1318,14 @@ BATCH_STATUS_POLL_MAX_ATTEMPTS = int(
 HEALTH_CHECK_TIMEOUT_SECONDS = int(
     os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)
 )  # 60 seconds
+_background_health_check_max_tokens_env = os.getenv(
+    "BACKGROUND_HEALTH_CHECK_MAX_TOKENS"
+)
+BACKGROUND_HEALTH_CHECK_MAX_TOKENS: Optional[int] = (
+    int(_background_health_check_max_tokens_env)
+    if _background_health_check_max_tokens_env
+    else 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"

litellm/proxy/health_check.py

@@ -11,7 +11,11 @@ from typing import List, Optional
 import litellm
 
 logger = logging.getLogger(__name__)
-from litellm.constants import DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS
+from litellm.constants import (
+    BACKGROUND_HEALTH_CHECK_MAX_TOKENS,
+    DEFAULT_HEALTH_CHECK_PROMPT,
+    HEALTH_CHECK_TIMEOUT_SECONDS,
+)
 
 ILLEGAL_DISPLAY_PARAMS = [
     "messages",
@@ -281,6 +285,8 @@ def _update_litellm_params_for_health_check(
     _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 ""
     ):

tests/test_litellm/proxy/test_health_check_max_tokens.py

@@ -1,7 +1,10 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
 import pytest
-from litellm.proxy.health_check import _update_litellm_params_for_health_check
+
 from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers
-from unittest.mock import AsyncMock, patch, MagicMock
+from litellm.proxy import health_check as hc_module
+from litellm.proxy.health_check import _update_litellm_params_for_health_check
 
 
 @pytest.mark.asyncio
@@ -73,3 +76,50 @@ async def test_ahealth_check_wildcard_models_respects_max_tokens():
             litellm_logging_obj=MagicMock(),
         )
         assert model_params["max_tokens"] == 3
+
+
+@pytest.mark.asyncio
+async def test_background_health_check_max_tokens_env_var(monkeypatch):
+    """
+    Test that BACKGROUND_HEALTH_CHECK_MAX_TOKENS env var is used as global default
+    for explicit (non-wildcard) models.
+    """
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10)
+
+    model_info = {}
+    litellm_params = {"model": "azure/gpt-4"}
+
+    updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert updated_params["max_tokens"] == 10
+
+
+@pytest.mark.asyncio
+async def test_per_model_overrides_global_env_var(monkeypatch):
+    """
+    Test that per-model health_check_max_tokens takes priority over
+    BACKGROUND_HEALTH_CHECK_MAX_TOKENS env var.
+    """
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10)
+
+    model_info = {"health_check_max_tokens": 5}
+    litellm_params = {"model": "azure/gpt-4"}
+
+    updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert updated_params["max_tokens"] == 5
+
+
+@pytest.mark.asyncio
+async def test_global_env_var_applies_to_wildcard_models(monkeypatch):
+    """
+    Test that BACKGROUND_HEALTH_CHECK_MAX_TOKENS env var also applies to wildcard models.
+    """
+    monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 15)
+
+    model_info = {}
+    litellm_params = {"model": "openai/*"}
+
+    updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert updated_params["max_tokens"] == 15