refactor(proxy_server.py): add debug logging around license check event (refactor position in startup_event logic)

This commit is contained in:
Krrish Dholakia 2024-11-06 14:29:28 +05:30
parent 0fe8cde7c7
commit cf8fb58bcd
3 changed files with 44 additions and 21 deletions

View file

@ -5,6 +5,9 @@ import json
import os
import traceback
from datetime import datetime
from typing import Optional
import httpx
from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.http_handler import HTTPHandler
@ -44,23 +47,46 @@ class LicenseCheck:
verbose_proxy_logger.error(f"Error reading public key: {str(e)}")
def _verify(self, license_str: str) -> bool:
verbose_proxy_logger.debug(
"litellm.proxy.auth.litellm_license.py::_verify - Checking license against {}/verify_license - {}".format(
self.base_url, license_str
)
)
url = "{}/verify_license/{}".format(self.base_url, license_str)
response: Optional[httpx.Response] = None
try: # don't impact user, if call fails
response = self.http_handler.get(url=url)
num_retries = 3
for i in range(num_retries):
try:
response = self.http_handler.get(url=url)
if response is None:
raise Exception("No response from license server")
response.raise_for_status()
except httpx.HTTPStatusError:
if i == num_retries - 1:
raise
response.raise_for_status()
if response is None:
raise Exception("No response from license server")
response_json = response.json()
premium = response_json["verify"]
assert isinstance(premium, bool)
verbose_proxy_logger.debug(
"litellm.proxy.auth.litellm_license.py::_verify - License={} is premium={}".format(
license_str, premium
)
)
return premium
except Exception as e:
verbose_proxy_logger.error(
"litellm.proxy.auth.litellm_license.py::_verify - Unable to verify License via api. - {}".format(
str(e)
verbose_proxy_logger.exception(
"litellm.proxy.auth.litellm_license.py::_verify - Unable to verify License={} via api. - {}".format(
license_str, str(e)
)
)
return False
@ -72,7 +98,7 @@ class LicenseCheck:
"""
try:
verbose_proxy_logger.debug(
"litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - {}".format(
"litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - LiteLLM License={}".format(
self.license_str
)
)

View file

@ -694,6 +694,9 @@ def run_server( # noqa: PLR0915
import litellm
if detailed_debug is True:
litellm._turn_on_debug()
# DO NOT DELETE - enables global variables to work across files
from litellm.proxy.proxy_server import app # noqa

View file

@ -3060,6 +3060,15 @@ async def startup_event():
user_api_key_cache=user_api_key_cache,
)
## CHECK PREMIUM USER
verbose_proxy_logger.debug(
"litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format(
premium_user
)
)
if premium_user is False:
premium_user = _license_check.is_premium()
### LOAD CONFIG ###
worker_config: Optional[Union[str, dict]] = get_secret("WORKER_CONFIG") # type: ignore
env_config_yaml: Optional[str] = get_secret_str("CONFIG_FILE_PATH")
@ -3107,21 +3116,6 @@ async def startup_event():
if isinstance(worker_config, dict):
await initialize(**worker_config)
## CHECK PREMIUM USER
verbose_proxy_logger.debug(
"litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format(
premium_user
)
)
if premium_user is False:
premium_user = _license_check.is_premium()
verbose_proxy_logger.debug(
"litellm.proxy.proxy_server.py::startup() - PREMIUM USER value - {}".format(
premium_user
)
)
ProxyStartupEvent._initialize_startup_logging(
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,