fix(proxy_server.py): allow user to disable scheduled reset budget task

This commit is contained in:
Krrish Dholakia 2024-03-19 20:36:22 -07:00
parent 2dfdc8dd69
commit f25b03326b
2 changed files with 11 additions and 4 deletions

View file

@ -591,6 +591,8 @@ general_settings:
"general_settings": { "general_settings": {
"completion_model": "string", "completion_model": "string",
"disable_spend_logs": "boolean", # turn off writing each transaction to the db "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 "key_management_system": "google_kms", # either google_kms or azure_kms
"master_key": "string", "master_key": "string",
"database_url": "string", "database_url": "string",

View file

@ -2658,7 +2658,7 @@ async def startup_event():
if prisma_client is not None: if prisma_client is not None:
create_view_response = await prisma_client.check_view_exists() create_view_response = await prisma_client.check_view_exists()
### START BUDGET SCHEDULER ### ### START BATCH WRITING DB ###
if prisma_client is not None: if prisma_client is not None:
scheduler = AsyncIOScheduler() scheduler = AsyncIOScheduler()
interval = random.randint( interval = random.randint(
@ -2667,9 +2667,14 @@ async def startup_event():
batch_writing_interval = random.randint( batch_writing_interval = random.randint(
proxy_batch_write_at - 3, proxy_batch_write_at + 3 proxy_batch_write_at - 3, proxy_batch_write_at + 3
) # random interval, so multiple workers avoid batch writing at the same time ) # random interval, so multiple workers avoid batch writing at the same time
### RESET BUDGET ###
if general_settings.get("disable_reset_budget", False) == False:
scheduler.add_job( scheduler.add_job(
reset_budget, "interval", seconds=interval, args=[prisma_client] reset_budget, "interval", seconds=interval, args=[prisma_client]
) )
### UPDATE SPEND ###
scheduler.add_job( scheduler.add_job(
update_spend, update_spend,
"interval", "interval",