mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
refactor(proxy): keep the env re-probe task in a set instead of a module global
The strict ruff gate now rejects new global statements (PLW0603), so the once-per-process re-probe is tracked in a set that drops a task when it finishes. A test covers that a second miss doesn't start a second re-probe while one is running, and that a finished one is replaced. The coordination Redis test helper now also isolates saml_sso_cache, which _attach_redis_usage_cache writes to, so these tests stop leaking a stand-in Redis into it.
This commit is contained in:
parent
8c760fee32
commit
d0c6e30838
2 changed files with 39 additions and 4 deletions
|
|
@ -5110,7 +5110,8 @@ def build_redis_usage_cache(redis_params: Mapping[str, object]) -> RedisCache:
|
|||
_build_redis_usage_cache: Final = build_redis_usage_cache
|
||||
|
||||
|
||||
_coordination_redis_env_retry_task: "asyncio.Task | None" = None
|
||||
# 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:
|
||||
|
|
@ -5187,16 +5188,17 @@ async def _retry_coordination_redis_from_environment(enable_redis_auth_cache: bo
|
|||
|
||||
def _schedule_coordination_redis_env_retry(enable_redis_auth_cache: bool) -> None:
|
||||
"""Start the background re-probe, once per process."""
|
||||
global _coordination_redis_env_retry_task # rebind-ok: one retry task per process
|
||||
if _coordination_redis_env_retry_task is not None and not _coordination_redis_env_retry_task.done():
|
||||
if any(not task.done() for task in _coordination_redis_env_retry_tasks):
|
||||
return
|
||||
try:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return
|
||||
_coordination_redis_env_retry_task = loop.create_task(
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -13169,6 +13169,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),
|
||||
|
|
@ -13476,6 +13477,38 @@ async def test_env_fallback_does_not_schedule_a_retry_without_a_redis_target():
|
|||
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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue