scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #1696 from BerriAI/litellm_proxy_cold_start_Time

Browse source 
fix(proxy_server.py): speed up proxy startup time
This commit is contained in:
Krish Dholakia 2024-01-30 21:34:45 -08:00 • committed by GitHub
commit 67bb11836a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 46 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/_types.py
View file

@ -325,17 +325,18 @@ class LiteLLM_VerificationToken(LiteLLMBase):
key_alias: Optional[str] = None
spend: float = 0.0
max_budget: Optional[float] = None
expires: Union[str, None]
expires: Union[datetime, str, None]
models: List[str]
aliases: Dict[str, str] = {}
config: Dict[str, str] = {}
user_id: Union[str, None]
max_parallel_requests: Union[int, None]
metadata: Dict[str, str] = {}
aliases: Dict = {}
config: Dict = {}
user_id: Optional[str] = None
max_parallel_requests: Optional[int] = None
metadata: Dict = {}
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
budget_duration: Optional[str] = None
budget_reset_at: Optional[datetime] = None
team_id: Optional[str] = None
class LiteLLM_Config(LiteLLMBase):

62
litellm/proxy/proxy_server.py
View file

@ -1269,7 +1269,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(
@ -1359,6 +1359,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}")
@ -1673,14 +1685,16 @@ 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(
duration=None,
models=[],
aliases={},
config={},
spend=0,
token=master_key,
user_id="default_user_id",
asyncio.create_task(
generate_key_helper_fn(
duration=None,
models=[],
aliases={},
config={},
spend=0,
token=master_key,
user_id="default_user_id",
)
)
if prisma_client is not None and litellm.max_budget > 0:
@ -1690,20 +1704,22 @@ async def startup_event():
)
# add proxy budget to db in the user table
await generate_key_helper_fn(
user_id=litellm_proxy_budget_name,
duration=None,
models=[],
aliases={},
config={},
spend=0,
max_budget=litellm.max_budget,
budget_duration=litellm.budget_duration,
query_type="update_data",
update_key_values={
"max_budget": litellm.max_budget,
"budget_duration": litellm.budget_duration,
},
asyncio.create_task(
generate_key_helper_fn(
user_id=litellm_proxy_budget_name,
duration=None,
models=[],
aliases={},
config={},
spend=0,
max_budget=litellm.max_budget,
budget_duration=litellm.budget_duration,
query_type="update_data",
update_key_values={
"max_budget": litellm.max_budget,
"budget_duration": litellm.budget_duration,
},
)
)
verbose_proxy_logger.debug(