diff --git a/litellm/constants.py b/litellm/constants.py index c17e776f54b..f278808e31e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -588,6 +588,8 @@ _REALTIME_BODY_CACHE_SIZE = 1000 # Keep realtime helper caches bounded; workloa INITIAL_RETRY_DELAY: Final = get_env_float("INITIAL_RETRY_DELAY", 0.5) MAX_RETRY_DELAY: Final = get_env_float("MAX_RETRY_DELAY", 8.0) JITTER: Final = get_env_float("JITTER", 0.75) +COORDINATION_REDIS_ENV_RETRY_SECONDS: Final = 30.0 +COORDINATION_REDIS_ENV_RETRY_ATTEMPTS: Final = 20 DEFAULT_IN_MEMORY_TTL = get_env_int("DEFAULT_IN_MEMORY_TTL", 5) # default time to live for the in-memory cache DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE: Final = get_env_int( "DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 1000 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a767f5423ed..8ea7d8a75b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -74,6 +74,8 @@ from litellm.constants import ( AWS_ERROR_TYPE_HEADER, BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, BASE_MCP_ROUTE, + COORDINATION_REDIS_ENV_RETRY_ATTEMPTS, + COORDINATION_REDIS_ENV_RETRY_SECONDS, DAILY_TAG_SPEND_BATCH_MULTIPLIER, DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, @@ -5157,6 +5159,97 @@ def build_redis_usage_cache(redis_params: Mapping[str, object]) -> RedisCache: _build_redis_usage_cache: Final = build_redis_usage_cache +# Holds the one background re-probe per process; a set, so scheduling never rebinds a module global. +_coordination_redis_env_retry_tasks: Final[set["asyncio.Task[None]"]] = set() + + +async def _probe_coordination_redis_from_environment() -> RedisCache | None: + """ + Builds the Redis named by the REDIS_* env vars and pings it, off the event loop. + + `RedisCache.__init__` talks to the server, and against a target that drops SYN + rather than refusing it redis-py's own retries take minutes, so it runs in a + worker thread. + """ + try: + candidate: Final = await asyncio.to_thread(_build_redis_usage_cache_from_environment) + except Exception as e: # noqa: BLE001 # a malformed inferred Redis env var must not block startup + verbose_proxy_logger.warning( + "coordination_redis: could not build a Redis client from REDIS_* environment variables " + "(%s); cross-pod coordination stays in-memory. Set general_settings.coordination_redis " + "explicitly to require it.", + e, + ) + return None + if candidate is None: + return None + try: + reachable: Final = await asyncio.wait_for(candidate.ping(), timeout=2.0) + except Exception as e: # noqa: BLE001 # an unreachable inferred Redis must not block startup or writes + verbose_proxy_logger.warning( + "coordination_redis: REDIS_* environment variables named a Redis that is not reachable " + "(%s); retrying in the background. Set general_settings.coordination_redis explicitly to " + "require it.", + e, + ) + return None + return candidate if reachable else None + + +async def _retry_coordination_redis_from_environment(enable_redis_auth_cache: bool) -> None: + """ + Re-probes the REDIS_* Redis after the boot-time probe found it unreachable, and + attaches it to every consumer a healthy boot would have wired, so budgets, rate + limits and auth cache invalidation stop being per-pod. + + Gives up after the window: a REDIS_HOST left over from an unrelated service + must not be probed for the life of the pod. + """ + for _ in range(COORDINATION_REDIS_ENV_RETRY_ATTEMPTS): + await asyncio.sleep(COORDINATION_REDIS_ENV_RETRY_SECONDS) + if redis_usage_cache is not None: + return + recovered = await _probe_coordination_redis_from_environment() # rebind-ok: one probe per attempt + if recovered is None: + continue + if redis_usage_cache is not None: + return + _attach_redis_usage_cache(recovered, enable_redis_auth_cache=enable_redis_auth_cache) + _set_redis_usage_cache(recovered) + proxy_logging_obj.update_values(redis_cache=recovered) + if llm_router is not None and llm_router.cache.redis_cache is None: + llm_router._update_redis_cache(cache=recovered) + proxy_config.start_auth_cache_invalidation_subscriber( + redis_cache=recovered, + user_api_key_cache=user_api_key_cache, + ) + verbose_proxy_logger.info( + "coordination_redis: REDIS_* Redis answered on a retry and is now attached for usage " + "tracking, rate limiting, and cross-pod coordination." + ) + return + verbose_proxy_logger.warning( + "coordination_redis: REDIS_* environment variables named a Redis that stayed unreachable for " + "%s attempts; cross-pod coordination stays in-memory until this pod restarts.", + COORDINATION_REDIS_ENV_RETRY_ATTEMPTS, + ) + + +def _schedule_coordination_redis_env_retry(enable_redis_auth_cache: bool) -> None: + """Start the background re-probe, once per process.""" + if any(not task.done() for task in _coordination_redis_env_retry_tasks): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + task: Final = loop.create_task( + _retry_coordination_redis_from_environment(enable_redis_auth_cache=enable_redis_auth_cache) + ) + _coordination_redis_env_retry_tasks.add(task) + task.add_done_callback(_coordination_redis_env_retry_tasks.discard) + + def environment_has_redis_connection_target() -> bool: """ Whether the REDIS_* environment variables name a Redis to connect to (host, @@ -6252,36 +6345,23 @@ class ProxyConfig: job/service, or a REDIS_CLUSTER_NODES value nothing here ever asked to be parsed. Wrongly guessing "coordination available" must not turn a previously harmless in-memory-only proxy into one that fails to boot or raises on every - cache write, so a malformed value or a failed/slow ping are both treated the - same as no REDIS_* vars at all. + cache write, so a malformed value or a failed/slow ping both leave the proxy + in-memory rather than failing startup. + + A ping can also fail because Redis is simply not up yet, which is the common + case on Kubernetes: the proxy and Redis start together. Giving up for the life + of the pod there multiplies every budget by the replica count, so when the env + does name a Redis, a background re-probe keeps looking for it. """ - try: - env_coordination_redis_cache: Final = _build_redis_usage_cache_from_environment() - except Exception as e: # noqa: BLE001 # a malformed inferred Redis env var must not block startup - verbose_proxy_logger.warning( - "coordination_redis: could not build a Redis client from REDIS_* environment variables " - "(%s); cross-pod coordination stays in-memory. Set general_settings.coordination_redis " - "explicitly to require it.", - e, - ) - return None + enable_redis_auth_cache: Final = litellm_settings.get("enable_redis_auth_cache", False) is True + env_coordination_redis_cache: Final = await _probe_coordination_redis_from_environment() if env_coordination_redis_cache is None: - return None - try: - reachable: Final = await asyncio.wait_for(env_coordination_redis_cache.ping(), timeout=2.0) - except Exception as e: # noqa: BLE001 # an unreachable inferred Redis must not block startup or writes - verbose_proxy_logger.warning( - "coordination_redis: REDIS_* environment variables named a Redis that is not reachable " - "(%s); cross-pod coordination stays in-memory. Set general_settings.coordination_redis " - "explicitly to require it.", - e, - ) - return None - if not reachable: + if environment_has_redis_connection_target(): + _schedule_coordination_redis_env_retry(enable_redis_auth_cache=enable_redis_auth_cache) return None _attach_redis_usage_cache( env_coordination_redis_cache, - enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True, + enable_redis_auth_cache=enable_redis_auth_cache, ) verbose_proxy_logger.info( "coordination_redis: using a standalone Redis built from REDIS_* " diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3647146eb4e..618ddacd0e0 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13487,6 +13487,7 @@ def _patched_coordination_redis_module_state( patch.object(proxy_server_module, "spend_counter_cache", spend_cache), patch.object(proxy_server_module, "user_api_key_cache", DualCache()), patch.object(proxy_server_module, "cli_sso_session_cache", DualCache()), + patch.object(proxy_server_module, "saml_sso_cache", DualCache()), patch.object(proxy_server_module, "llm_router", None), patch.object(proxy_server_module, "litellm_config_cache", config_cache), patch.object(proxy_server_module, "RedisCache", redis_cache_class), @@ -13758,6 +13759,155 @@ async def test_init_coordination_redis_env_fallback_unreachable_stays_in_memory( assert spend_redis is None +@pytest.mark.asyncio +async def test_env_fallback_schedules_a_retry_when_the_named_redis_is_unreachable(): + """An unreachable Redis at boot is usually Redis starting next to the proxy, so + giving up for the life of the pod leaves coordination per-pod.""" + with ( + patch.object(proxy_server_module, "environment_has_redis_connection_target", return_value=True), + patch.object(proxy_server_module, "_schedule_coordination_redis_env_retry") as scheduled, + ): + built, spend_redis = await _run_init_coordination_redis_env_fallback( + litellm_settings={}, + redis_env_kwargs={"host": "unreachable-host", "port": "6390"}, + redis_cache_class=_UnreachableRedisCache, + ) + + assert built is None + assert spend_redis is None + assert scheduled.call_count == 1 + + +@pytest.mark.asyncio +async def test_env_fallback_does_not_schedule_a_retry_without_a_redis_target(): + """With no REDIS_* connection info there is nothing to re-probe.""" + with ( + patch.object(proxy_server_module, "environment_has_redis_connection_target", return_value=False), + patch.object(proxy_server_module, "_schedule_coordination_redis_env_retry") as scheduled, + ): + built, spend_redis = await _run_init_coordination_redis_env_fallback( + litellm_settings={}, + redis_env_kwargs={}, + ) + + assert built is None + assert spend_redis is None + assert scheduled.call_count == 0 + + +@pytest.mark.asyncio +async def test_env_retry_is_scheduled_once_while_one_is_running(): + """A config reload runs the env fallback again, so a second failed boot probe + must not start a second re-probe next to the one still running.""" + release: Final = asyncio.Event() + started: Final[list[bool]] = [] + + async def retry(enable_redis_auth_cache: bool) -> None: + started.append(enable_redis_auth_cache) + await release.wait() + + pending: Final[set[asyncio.Task[None]]] = set() + with ( + patch.object(proxy_server_module, "_retry_coordination_redis_from_environment", retry), + patch.object(proxy_server_module, "_coordination_redis_env_retry_tasks", pending), + ): + proxy_server_module._schedule_coordination_redis_env_retry(enable_redis_auth_cache=False) + proxy_server_module._schedule_coordination_redis_env_retry(enable_redis_auth_cache=False) + assert len(pending) == 1 + + (first,) = tuple(pending) + release.set() + await first + assert not pending, "a finished re-probe is dropped so a later miss can start a new one" + + proxy_server_module._schedule_coordination_redis_env_retry(enable_redis_auth_cache=True) + (second,) = tuple(pending) + await second + + assert started == [False, True] + + +@pytest.mark.asyncio +async def test_retry_attaches_the_env_redis_once_it_answers(): + """The retry has to leave the pod in the state a healthy boot would: spend + counters, the rate limiter and the router on Redis, and the auth cache + invalidation subscriber running, all of which startup wires only once.""" + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + recovered: Final = _EnvBuiltRedisCache(host="env-fallback-host", port="6390") + probe: Final = AsyncMock(side_effect=(None, recovered)) + subscriber: Final = MagicMock() + rate_limiter_cache: Final = DualCache() + router: Final = MagicMock(cache=DualCache()) + + with ( + _patched_coordination_redis_module_state(spend_cache=fresh_spend_cache, config_cache=fresh_config_cache), + patch.object(proxy_server_module, "COORDINATION_REDIS_ENV_RETRY_SECONDS", 0), + patch.object(proxy_server_module, "COORDINATION_REDIS_ENV_RETRY_ATTEMPTS", 5), + patch.object(proxy_server_module, "_probe_coordination_redis_from_environment", probe), + patch.object(proxy_server_module, "llm_router", router), + patch.object( + proxy_server_module.proxy_logging_obj.internal_usage_cache, "dual_cache", rate_limiter_cache + ), + patch.object(proxy_server_module.proxy_config, "start_auth_cache_invalidation_subscriber", subscriber), + ): + await proxy_server_module._retry_coordination_redis_from_environment(enable_redis_auth_cache=False) + attached = proxy_server_module.redis_usage_cache + + assert attached is recovered + assert fresh_spend_cache.redis_cache is recovered + assert rate_limiter_cache.redis_cache is recovered + assert router._update_redis_cache.call_args.kwargs["cache"] is recovered + assert subscriber.call_args.kwargs["redis_cache"] is recovered + + +@pytest.mark.asyncio +async def test_retry_does_not_replace_a_backend_attached_while_it_probed(): + """A probe takes seconds, and the DB-stored coordination settings resolve a + backend in that window. The env fallback is last in startup precedence, so a + late answer must not take over from the backend the proxy already chose.""" + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + chosen_elsewhere: Final = _EnvBuiltRedisCache(host="db-configured-host", port="6390") + late_answer: Final = _EnvBuiltRedisCache(host="env-fallback-host", port="6390") + + async def probe_that_loses_the_race(): + proxy_server_module._set_redis_usage_cache(chosen_elsewhere) + return late_answer + + with ( + _patched_coordination_redis_module_state(spend_cache=fresh_spend_cache, config_cache=fresh_config_cache), + patch.object(proxy_server_module, "COORDINATION_REDIS_ENV_RETRY_SECONDS", 0), + patch.object(proxy_server_module, "COORDINATION_REDIS_ENV_RETRY_ATTEMPTS", 5), + patch.object(proxy_server_module, "_probe_coordination_redis_from_environment", probe_that_loses_the_race), + ): + await proxy_server_module._retry_coordination_redis_from_environment(enable_redis_auth_cache=False) + attached = proxy_server_module.redis_usage_cache + + assert attached is chosen_elsewhere + assert fresh_spend_cache.redis_cache is not late_answer + + +@pytest.mark.asyncio +async def test_retry_gives_up_after_the_configured_attempts(): + """A REDIS_HOST left over from an unrelated service never answers, so the probe + must stop rather than poke a stranger's address for the life of the pod.""" + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + never_answers: Final = AsyncMock(return_value=None) + + with ( + _patched_coordination_redis_module_state(spend_cache=fresh_spend_cache, config_cache=fresh_config_cache), + patch.object(proxy_server_module, "COORDINATION_REDIS_ENV_RETRY_SECONDS", 0), + patch.object(proxy_server_module, "COORDINATION_REDIS_ENV_RETRY_ATTEMPTS", 3), + patch.object(proxy_server_module, "_probe_coordination_redis_from_environment", never_answers), + ): + await proxy_server_module._retry_coordination_redis_from_environment(enable_redis_auth_cache=False) + + assert never_answers.await_count == 3 + assert fresh_spend_cache.redis_cache is None + + @pytest.mark.asyncio async def test_init_coordination_redis_env_fallback_malformed_cluster_nodes_stays_in_memory(): """REDIS_CLUSTER_NODES can be set to a malformed value nothing here ever