diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 36b59ea39..bcfc34489 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1743,7 +1743,10 @@ async def startup_event(): ### START BUDGET SCHEDULER ### scheduler = AsyncIOScheduler() - scheduler.add_job(reset_budget, "interval", seconds=10, args=[prisma_client]) + interval = random.randint( + 7, 14 + ) # 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() diff --git a/tests/test_keys.py b/tests/test_keys.py index 94e96769c..fb665fccf 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -355,9 +355,17 @@ async def test_key_with_budgets(): print(f"hashed_token: {hashed_token}") key_info = await get_key_info(session=session, get_key=key, call_key=key) reset_at_init_value = key_info["info"]["budget_reset_at"] - await asyncio.sleep(30) - key_info = await get_key_info(session=session, get_key=key, call_key=key) - reset_at_new_value = key_info["info"]["budget_reset_at"] + reset_at_new_value = None + i = 0 + while i < 3: + await asyncio.sleep(30) + key_info = await get_key_info(session=session, get_key=key, call_key=key) + reset_at_new_value = key_info["info"]["budget_reset_at"] + try: + assert reset_at_init_value != reset_at_new_value + break + except: + i + 1 assert reset_at_init_value != reset_at_new_value