diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 66bbda1ca4e..77dda842302 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -279,6 +279,17 @@ db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value +# The event loop only keeps weak references to tasks, so a background task whose +# only reference was the create_task() call can be collected before it runs. Hold +# each one until it completes. +_background_tasks: Final[set[asyncio.Task]] = set() # mutable-ok: task registry + + +def _hold_background_task(task: asyncio.Task) -> None: + """Keep a strong reference to a fire-and-forget task until it finishes.""" + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) + def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ @@ -2148,10 +2159,12 @@ async def _get_fuzzy_user_object( ) if response is not None and sso_user_id is not None: # update sso_user_id - asyncio.create_task( # background task to update user with sso id - _user_table(UserRepository(prisma_client)).update( - where={"user_id": response.user_id}, - data={"sso_user_id": sso_user_id}, + _hold_background_task( + asyncio.create_task( # background task to update user with sso id + _user_table(UserRepository(prisma_client)).update( + where={"user_id": response.user_id}, + data={"sso_user_id": sso_user_id}, + ) ) ) @@ -4411,10 +4424,12 @@ async def _virtual_key_max_budget_check( key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="token_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="token_budget", + user_info=call_info, + ) ) ) @@ -4515,10 +4530,12 @@ async def _virtual_key_soft_budget_check( event_group=Litellm_EntityType.KEY, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, + ) ) ) @@ -4606,10 +4623,12 @@ async def _virtual_key_max_budget_alert_check( event_group=Litellm_EntityType.KEY, max_budget_alert_emails=alert_email_config, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="max_budget_alert", + user_info=call_info, + ) ) ) else: @@ -4638,10 +4657,12 @@ async def _virtual_key_max_budget_alert_check( event_group=Litellm_EntityType.KEY, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="max_budget_alert", + user_info=call_info, + ) ) ) @@ -4802,10 +4823,12 @@ async def _team_max_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.TEAM, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="team_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) ) ) @@ -4920,10 +4943,12 @@ async def _team_soft_budget_check( alert_emails=alert_emails, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, + ) ) ) @@ -4964,10 +4989,12 @@ async def _project_max_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="project_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="project_budget", + user_info=call_info, + ) ) ) @@ -5016,10 +5043,12 @@ async def _project_soft_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, + ) ) ) @@ -5174,10 +5203,12 @@ async def _organization_max_budget_check( organization_id=org_id, event_group=Litellm_EntityType.ORGANIZATION, ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="organization_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="organization_budget", + user_info=call_info, + ) ) )