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:
RosieOh 2026-10-28 12:03:28 +09:00
parent 8c760fee32
commit d0c6e30838
2 changed files with 39 additions and 4 deletions

View file

@ -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:

View file

@ -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