phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(proxy_server.py): speed up proxy startup time

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-30 16:17:23 -08:00
parent 37de964da4
commit 2019347f0d
1 changed files with 39 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/proxy_server.py
View file

@ -1267,7 +1267,7 @@ async def generate_key_helper_fn(
update_key_values: Optional[dict] = None,
key_alias: Optional[str] = None,
):
global prisma_client, custom_db_client
global prisma_client, custom_db_client, user_api_key_cache
if prisma_client is None and custom_db_client is None:
raise Exception(
@ -1357,6 +1357,18 @@ async def generate_key_helper_fn(
}
if general_settings.get("allow_user_auth", False) == True:
key_data["key_name"] = f"sk-...{token[-4:]}"
saved_token = copy.deepcopy(key_data)
if isinstance(saved_token["aliases"], str):
saved_token["aliases"] = json.loads(saved_token["aliases"])
if isinstance(saved_token["config"], str):
saved_token["config"] = json.loads(saved_token["config"])
if isinstance(saved_token["metadata"], str):
saved_token["metadata"] = json.loads(saved_token["metadata"])
user_api_key_cache.set_cache(
key=key_data["token"],
value=LiteLLM_VerificationToken(**saved_token), # type: ignore
ttl=60,
)
if prisma_client is not None:
## CREATE USER (If necessary)
verbose_proxy_logger.debug(f"prisma_client: Creating User={user_data}")
@ -1671,7 +1683,8 @@ async def startup_event():
if prisma_client is not None and master_key is not None:
# add master key to db
await generate_key_helper_fn(
asyncio.create_task(
generate_key_helper_fn(
duration=None,
models=[],
aliases={},
@ -1680,6 +1693,7 @@ async def startup_event():
token=master_key,
user_id="default_user_id",
)
)
if prisma_client is not None and litellm.max_budget > 0:
if litellm.budget_duration is None:
@ -1688,7 +1702,8 @@ async def startup_event():
)
# add proxy budget to db in the user table
await generate_key_helper_fn(
asyncio.create_task(
generate_key_helper_fn(
user_id=litellm_proxy_budget_name,
duration=None,
models=[],
@ -1703,6 +1718,7 @@ async def startup_event():
"budget_duration": litellm.budget_duration,
},
)
)
verbose_proxy_logger.debug(
f"custom_db_client client {custom_db_client}. Master_key: {master_key}"