From f25b03326bbf6a60ffecc82d5d3d9e75aa24f573 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 19 Mar 2024 20:36:22 -0700 Subject: [PATCH] fix(proxy_server.py): allow user to disable scheduled reset budget task --- docs/my-website/docs/proxy/configs.md | 2 ++ litellm/proxy/proxy_server.py | 13 +++++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index 1ece75529..34271f064 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -591,6 +591,8 @@ general_settings: "general_settings": { "completion_model": "string", "disable_spend_logs": "boolean", # turn off writing each transaction to the db + "disable_reset_budget": "boolean", # turn off reset budget scheduled task + "enable_jwt_auth": "boolean", # allow proxy admin to auth in via jwt tokens with 'litellm_proxy_admin' in claims "key_management_system": "google_kms", # either google_kms or azure_kms "master_key": "string", "database_url": "string", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2f8c9bf18..47b551a07 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2658,7 +2658,7 @@ async def startup_event(): if prisma_client is not None: create_view_response = await prisma_client.check_view_exists() - ### START BUDGET SCHEDULER ### + ### START BATCH WRITING DB ### if prisma_client is not None: scheduler = AsyncIOScheduler() interval = random.randint( @@ -2667,9 +2667,14 @@ async def startup_event(): batch_writing_interval = random.randint( proxy_batch_write_at - 3, proxy_batch_write_at + 3 ) # random interval, so multiple workers avoid batch writing at the same time - scheduler.add_job( - reset_budget, "interval", seconds=interval, args=[prisma_client] - ) + + ### RESET BUDGET ### + if general_settings.get("disable_reset_budget", False) == False: + scheduler.add_job( + reset_budget, "interval", seconds=interval, args=[prisma_client] + ) + + ### UPDATE SPEND ### scheduler.add_job( update_spend, "interval",