From d2dceb353718a1c257fba379dc272fab88316a4a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 7 Feb 2024 18:14:26 -0800 Subject: [PATCH] fix(proxy_server.py): check if prisma client is set before scheduling reset budget --- litellm/proxy/proxy_server.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b57f7808956..1f98ef95242 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2004,12 +2004,15 @@ async def startup_event(): ) ### START BUDGET SCHEDULER ### - scheduler = AsyncIOScheduler() - interval = random.randint( - 597, 605 - ) # random interval, so multiple workers avoid resetting budget at the same time - scheduler.add_job(reset_budget, "interval", seconds=interval, args=[prisma_client]) - scheduler.start() + if prisma_client is not None: + scheduler = AsyncIOScheduler() + interval = random.randint( + 597, 605 + ) # random interval, so multiple workers avoid resetting budget at the same time + scheduler.add_job( + reset_budget, "interval", seconds=interval, args=[prisma_client] + ) + scheduler.start() #### API ENDPOINTS ####