diff --git a/litellm/router_rust_core.py b/litellm/router_rust_core.py index ab33fd2f2d2..b815c313465 100644 --- a/litellm/router_rust_core.py +++ b/litellm/router_rust_core.py @@ -87,7 +87,6 @@ from litellm.router_rust_fallbacks import ( ) from litellm.router_rust_health import ( RustHealthCallbacks, - cooled_in_running_loop, failure_status, ) from litellm.router_rust_health_routing import ( @@ -896,7 +895,7 @@ class RustHealth(RouterHealth): handle: Final = self._core._handle # pyright: ignore[reportPrivateUsage] # see above call: Final = handle.report_failure(deployment_id, failure, count_only=exception_status == "") done: Final = self._core._drive_sync_steps(call) # pyright: ignore[reportPrivateUsage] # see above - return cooled_in_running_loop(bool(done.get("cooled"))) + return bool(done.get("cooled")) def team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": return self._core._hosted._core_team_ptu_ceiling(team_id, model_group) # pyright: ignore[reportPrivateUsage] # the host registry's ported ceiling diff --git a/litellm/router_rust_health.py b/litellm/router_rust_health.py index 0794577689d..006b88f4964 100644 --- a/litellm/router_rust_health.py +++ b/litellm/router_rust_health.py @@ -35,16 +35,6 @@ def failure_status(exception_status: object) -> int | None: return status if status is not None and 0 <= status <= 65535 else None -def cooled_in_running_loop(cooled: bool) -> bool: - """``set_cooldown_deployments`` schedules the cooldown event with ``asyncio.create_task`` - after writing the cooldown, so outside a running loop it raises ``RuntimeError`` then.""" - import asyncio - - if cooled: - asyncio.get_running_loop() - return cooled - - def _configured_cooldown(model_info: object, litellm_params: Mapping[str, object]) -> float | None: """``_first_present(model_info, litellm_params, key="cooldown_time")``.""" for source in (as_str_mapping(model_info), litellm_params): @@ -140,7 +130,7 @@ class RustHealthCallbacks: requested_group=group if isinstance(group, str) else None, cooldown_time=_configured_cooldown(model_info, litellm_params), ) - return cooled_in_running_loop(bool(self._drive_sync_steps(call).get("cooled"))) + return bool(self._drive_sync_steps(call).get("cooled")) def sync_deployment_callback_on_success( self, @@ -194,6 +184,5 @@ class RustHealthCallbacks: __all__ = [ "RustHealthCallbacks", - "cooled_in_running_loop", "failure_status", ] diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 2cea7389808..962ca6e3f2f 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -24,6 +24,8 @@ class BaseRoutingStrategy(ABC): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] self._sync_task: asyncio.Task[None] | None = None + self._sync_interval: float | None = None + self._wants_sync_task = False self._retired = False self._flush_after_retire: asyncio.Task[None] | None = None if should_batch_redis_writes: @@ -32,18 +34,30 @@ class BaseRoutingStrategy(ABC): self.in_memory_keys_to_update: set[str] = set() # Set with max size of 1000 keys def setup_sync_task(self, default_sync_interval: float | None): - """Setup the sync task in a way that's compatible with FastAPI""" - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) + """Start the periodic Redis sync on the running loop. + Built from synchronous code there is no loop to run it on, so it starts on the first + increment made inside one (``_increment_value_in_current_window``) instead.""" + self._sync_interval = default_sync_interval + self._wants_sync_task = True + self._start_sync_task() + + def _start_sync_task(self) -> None: + if not self._wants_sync_task or self._retired: + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + running: Final = self._sync_task + if running is not None and not running.done() and running.get_loop() is loop: + return self._sync_task = loop.create_task( - self.periodic_sync_in_memory_spend_with_redis(default_sync_interval=default_sync_interval) + self.periodic_sync_in_memory_spend_with_redis(default_sync_interval=self._sync_interval) ) def cancel_sync_task(self) -> None: + self._wants_sync_task = False if self._sync_task is not None: self._sync_task.cancel() @@ -99,6 +113,7 @@ class BaseRoutingStrategy(ABC): - Increments the spend in memory cache (so spend instantly updated in memory) - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) """ + self._start_sync_task() result: Final = await self.dual_cache.in_memory_cache.async_increment( key=key, value=value, diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 1d502dd8d95..7da1657fa21 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -128,7 +128,8 @@ class RouterBudgetLimiting(CustomLogger): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None - self._sync_task: asyncio.Task[None] = asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._sync_task: asyncio.Task[None] | None = None + self._pending_provider_budget_inits: tuple[tuple[str, GenericBudgetInfo], ...] = () self._retired = False self._flush_after_retire: asyncio.Task[None] | None = None self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config @@ -137,13 +138,37 @@ class RouterBudgetLimiting(CustomLogger): self._init_provider_budgets() self._init_deployment_budgets(model_list=model_list) self._init_tag_budgets() + self._start_background_work() # Add self to litellm callbacks if it's a list if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(self) + def _start_background_work(self) -> None: + """Start the periodic Redis sync and the provider budget cache seeding once a loop runs. + + Built from synchronous code there is no loop to run them on, so they wait for the first + async hook (``async_filter_deployments``, ``async_log_success_event``) instead; a sync task + left on a loop other than the running one (a short-lived loop that has since closed) is + replaced by one on the running loop.""" + if self._retired: + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + running: Final = self._sync_task + if running is not None and not running.done() and running.get_loop() is loop: + return + self._sync_task = loop.create_task(self.periodic_sync_in_memory_spend_with_redis()) + pending: Final = self._pending_provider_budget_inits + self._pending_provider_budget_inits = () + for provider, budget_config in pending: + _ = loop.create_task(self._init_provider_budget_in_cache(provider=provider, budget_config=budget_config)) + def retire(self) -> None: - self._sync_task.cancel() + if self._sync_task is not None: + self._sync_task.cancel() self._retired = True self._schedule_flush_after_retire() @@ -184,6 +209,7 @@ class RouterBudgetLimiting(CustomLogger): and openai spend > openai budget limit then skip this deployment """ + self._start_background_work() # If a single deployment is passed, convert it to a list if isinstance(healthy_deployments, dict): @@ -541,6 +567,7 @@ class RouterBudgetLimiting(CustomLogger): callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs) if callback_fields.get("call_type") == "_arealtime": return + self._start_background_work() standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: raise ValueError("standard_logging_payload is required") @@ -917,12 +944,7 @@ class RouterBudgetLimiting(CustomLogger): budget_limit=config.get("budget_limit"), time_period=config.get("time_period"), ) - asyncio.create_task( - self._init_provider_budget_in_cache( - provider=provider, - budget_config=self.provider_budget_config[provider], - ) - ) + self._pending_provider_budget_inits += ((provider, self.provider_budget_config[provider]),) verbose_router_logger.debug("Initalized Provider budget config: %s", self.provider_budget_config) diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 25821703029..e80ef95e2f8 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -29,6 +29,19 @@ async def router_cooldown_event_callback( - Updates deployment state on Prometheus - Increments cooldown metric for deployment on Prometheus """ + emit_router_cooldown_event( + ctx=ctx, deployment_id=deployment_id, exception_status=exception_status, cooldown_time=cooldown_time + ) + + +def emit_router_cooldown_event( + ctx: "RoutingContext", + deployment_id: str, + exception_status: str | int, + cooldown_time: float | None, +) -> None: + """``router_cooldown_event_callback``'s work, run inline: nothing in it awaits, so a cooldown + recorded outside a running event loop still reaches Prometheus.""" verbose_logger.debug("In router_cooldown_event_callback - updating prometheus") _deployment: Final = ctx.registry().deployment(deployment_id) if _deployment is None: @@ -74,8 +87,6 @@ async def router_cooldown_event_callback( exception_status=str(exception_status), ) - return - def get_prometheus_logger_from_callbacks() -> PrometheusLogger | None: """ diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index e50d5bb47a6..c40c04f6405 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -31,7 +31,7 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials from litellm.router_deprecated import renamed_router_param from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCacheValue -from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback +from litellm.router_utils.cooldown_callbacks import emit_router_cooldown_event, router_cooldown_event_callback from litellm.types.router import Deployment from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -615,15 +615,28 @@ def set_cooldown_deployments( cooldown_time=time_to_cooldown, ) - # Trigger cooldown callback handler - asyncio.create_task( - router_cooldown_event_callback( - ctx=ctx, - deployment_id=deployment, - exception_status=exception_status, - cooldown_time=time_to_cooldown, + # Trigger cooldown callback handler: a background task inside a running loop, inline outside one. + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + try: + emit_router_cooldown_event( + ctx=ctx, + deployment_id=deployment, + exception_status=exception_status, + cooldown_time=time_to_cooldown, + ) + except Exception as e: # noqa: BLE001 # as the task would, a failed event never undoes the cooldown + verbose_router_logger.exception("router cooldown event failed for %s: %s", deployment, e) + else: + _ = loop.create_task( + router_cooldown_event_callback( + ctx=ctx, + deployment_id=deployment, + exception_status=exception_status, + cooldown_time=time_to_cooldown, + ) ) - ) return True return False diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index c0583566531..edce24dc0b3 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -537,3 +537,43 @@ async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirma await second._push_in_memory_increments_to_redis() assert float(await client.get("provider_spend:openai:1d")) == 0.5 await client.aclose() + + +def test_limiter_built_outside_a_loop_starts_its_background_work_on_the_first_async_hook( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + started: Final[list[str]] = [] + + async def sync_loop(self: RouterBudgetLimiting) -> None: + started.append("sync") + + monkeypatch.setattr(RouterBudgetLimiting, "periodic_sync_in_memory_spend_with_redis", sync_loop) + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}} + ) + assert limiter._sync_task is None # pyright: ignore[reportPrivateUsage] # the background task under test + + async def first_hook() -> None: + await limiter.async_filter_deployments(model="g", healthy_deployments=[], messages=None) # pyright: ignore[reportUnknownMemberType] # an untyped deployment list + for _ in range(5): + await asyncio.sleep(0) + + asyncio.run(first_hook()) + assert started == ["sync"] + assert cache.get_cache("provider_spend:openai:1d") == 0.0 # pyright: ignore[reportUnknownMemberType] # DualCache reads are untyped + assert cache.get_cache("provider_budget_start_time:openai") is not None # pyright: ignore[reportUnknownMemberType] # see above + + +@pytest.mark.asyncio +async def test_limiter_built_inside_a_loop_starts_one_sync_task(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + limiter: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None) + task: Final = limiter._sync_task # pyright: ignore[reportPrivateUsage] # the background task under test + assert task is not None and not task.done() + + await limiter.async_filter_deployments(model="g", healthy_deployments=[], messages=None) # pyright: ignore[reportUnknownMemberType] # an untyped deployment list + + assert limiter._sync_task is task # pyright: ignore[reportPrivateUsage] # the background task under test + limiter.retire() diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index 129bb10eccc..c32683567fe 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -717,3 +717,29 @@ def setup_and_teardown(): if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() yield + + +def test_handler_built_outside_a_loop_starts_its_redis_sync_on_the_first_increment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started: Final[list[float | None]] = [] + + async def sync_loop(self: LowestTPMLoggingHandler, default_sync_interval: float | None) -> None: + started.append(default_sync_interval) + + monkeypatch.setattr(LowestTPMLoggingHandler, "periodic_sync_in_memory_spend_with_redis", sync_loop) + + def refuse_to_install(loop: object) -> None: + pytest.fail("must not install an event loop") + + monkeypatch.setattr(asyncio, "set_event_loop", refuse_to_install) + handler: Final = LowestTPMLoggingHandler(router_cache=DualCache()) + assert handler._sync_task is None # pyright: ignore[reportPrivateUsage] # the background task under test + + async def first_increment() -> None: + await handler._increment_value_in_current_window(key="k", value=1, ttl=60) # pyright: ignore[reportPrivateUsage] # the increment that starts it + await handler._increment_value_in_current_window(key="k", value=1, ttl=60) # pyright: ignore[reportPrivateUsage] # the increment that starts it + await asyncio.sleep(0) + + asyncio.run(first_increment()) + assert started == [0.1] diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 71c37b5b1d5..a3ff1c30942 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -3672,3 +3672,64 @@ def test_the_refusal_names_each_cooled_deployment_some_cooldown_state_routes_the with pytest.raises(RouterRateLimitError) as refusal: _select(_tenant_router(EVERY_TENANT, **BOTH_ON), "async", "tenants", copy.deepcopy(request_kwargs)) assert {deployment["model_id"] for deployment in refusal.value.cooldown_deployments or ()} == reachable + + +def _sync_router() -> Router: + return Router( + model_list=[ + {"model_name": "chat", "litellm_params": {"model": "openai/a", "api_key": "k"}, "model_info": {"id": "a"}}, + {"model_name": "chat", "litellm_params": {"model": "openai/b", "api_key": "k"}, "model_info": {"id": "b"}}, + ], + allowed_fails=0, + ) + + +def _cool_down_a(router: Router) -> bool: + exception: Final = litellm.InternalServerError(message="boom", model="a", llm_provider="openai") + return set_cooldown_deployments( + ctx=RoutingContext(router), + original_exception=exception, + exception_status=500, + deployment="a", + time_to_cooldown=60, + ) + + +def test_a_cooldown_set_outside_a_running_loop_reports_the_event_inline() -> None: + router: Final = _sync_router() + + with patch("litellm.router_utils.cooldown_handlers.emit_router_cooldown_event") as emit: + cooled: Final = _cool_down_a(router) + + assert cooled is True + assert emit.call_count == 1 + assert emit.call_args.kwargs["deployment_id"] == "a" + assert emit.call_args.kwargs["exception_status"] == 500 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == ["a"] + + +def test_a_failing_cooldown_event_outside_a_running_loop_keeps_the_cooldown() -> None: + router: Final = _sync_router() + + with patch( + "litellm.router_utils.cooldown_handlers.emit_router_cooldown_event", side_effect=ValueError("metrics down") + ): + cooled: Final = _cool_down_a(router) + + assert cooled is True + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == ["a"] + + +@pytest.mark.asyncio +async def test_a_cooldown_set_inside_a_running_loop_reports_the_event_as_a_task() -> None: + router: Final = _sync_router() + with ( + patch("litellm.router_utils.cooldown_handlers.router_cooldown_event_callback", new=MagicMock()) as callback, + patch("litellm.router_utils.cooldown_handlers.emit_router_cooldown_event") as emit, + ): + callback.return_value = asyncio.sleep(0) + assert _cool_down_a(router) is True + await asyncio.sleep(0) + + assert callback.call_count == 1 + assert emit.call_count == 0 diff --git a/tests/unit/test_router/core/test_construction_precall_characterization.py b/tests/unit/test_router/core/test_construction_precall_characterization.py index a86e7ae5022..6256f1021f5 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization.py @@ -767,11 +767,25 @@ def test_an_empty_optional_pre_call_check_list_registers_nothing(make_core: Make assert host.lists["callbacks"] == [] -def test_router_budget_limiting_cannot_be_built_outside_an_event_loop(make_core: MakeCore) -> None: - """Oddity: the budget limiter starts a background task in its constructor, so a router that needs - one cannot be built from synchronous code with no running loop.""" - with pytest.raises(RuntimeError, match="no running event loop"): - _one(make_core, optional_pre_call_checks=["router_budget_limiting"]) +def test_router_budget_limiting_can_be_built_outside_an_event_loop(make_core: MakeCore, host: ScriptedHost) -> None: + """Built from synchronous code, the budget limiter defers its background Redis sync until a loop runs.""" + core: Final = _one(make_core, optional_pre_call_checks=["router_budget_limiting"]) + limiter: Final = core.registry().plugins("budget_limiter") + assert isinstance(limiter, RouterBudgetLimiting) + assert host.lists["callbacks"] == [limiter] + core.retire() + + +def test_a_provider_budget_config_can_be_built_outside_an_event_loop(make_core: MakeCore) -> None: + core: Final = _one(make_core, provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}}) + assert isinstance(core.registry().plugins("budget_limiter"), RouterBudgetLimiting) + core.retire() + + +def test_a_deployment_budget_can_be_built_outside_an_event_loop(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", max_budget=1, budget_duration="1d")]) + assert isinstance(core.registry().plugins("budget_limiter"), RouterBudgetLimiting) + core.retire() @pytest.mark.asyncio diff --git a/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py b/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py index 4aeb11c58ec..2a331bf4da8 100644 --- a/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py +++ b/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py @@ -822,12 +822,32 @@ async def test_an_unhealthy_deployment_stays_excluded_when_the_healthy_one_is_co assert host.deployments_called("acompletion") == ["b"] -def test_a_cooldown_recorded_outside_an_event_loop_raises_after_writing_it(make_core: MakeCore) -> None: +def test_a_cooldown_recorded_outside_an_event_loop_is_written_and_reported(make_core: MakeCore) -> None: core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) - with pytest.raises(RuntimeError, match="no running event loop"): - core.health().record_failure("a", server_error(), 500) + cooled: Final = core.health().record_failure("a", server_error(), 500) + assert cooled is True + assert set(cooled_down(core, ["a", "b"])) == {"a"} + + +def test_a_failure_below_the_threshold_outside_an_event_loop_is_only_counted(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=3) + + cooled: Final = core.health().record_failure("a", server_error(), 500) + + assert cooled is False + assert core.health().failures_this_minute("a") == 1 + assert cooled_down(core, ["a", "b"]) == {} + + +def test_failure_callback_outside_an_event_loop_cools_down(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + cooled: Final = core.deployment_callback_on_failure(failure_kwargs({"id": "a"}, server_error()), None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # PythonCore's public member + + assert cooled is True + assert core.health().failures_this_minute("a") == 1 assert set(cooled_down(core, ["a", "b"])) == {"a"}