mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(proxy_cli): import proxy_server once on script-style boot (#42584)
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
b277be0867
commit
ecce7cdd9c
2 changed files with 52 additions and 46 deletions
|
|
@ -33,21 +33,20 @@ else:
|
|||
FastAPI = Any
|
||||
|
||||
|
||||
def _deprioritize_script_dir_in_sys_path() -> None:
|
||||
def _drop_script_dir_from_sys_path() -> None:
|
||||
"""Stop ``litellm/proxy`` modules from shadowing installed packages.
|
||||
|
||||
Running this file as a script puts its own directory at ``sys.path[0]``, so
|
||||
``import a2a`` resolves to ``litellm/proxy/a2a`` instead of the ``a2a`` SDK
|
||||
and A2A agent calls fail. The entry is moved to the end rather than dropped,
|
||||
because the sibling-import fallbacks in this module (``from proxy_server
|
||||
import ...``) still need it. No-op under the ``litellm`` console script.
|
||||
and ``proxy_server`` resolves to a second copy of
|
||||
``litellm.proxy.proxy_server``. No-op under the ``litellm`` console script.
|
||||
"""
|
||||
script_dir: Final = os.path.dirname(os.path.abspath(__file__))
|
||||
if sys.path and os.path.abspath(sys.path[0]) == script_dir:
|
||||
sys.path.append(sys.path.pop(0))
|
||||
sys.path.pop(0)
|
||||
|
||||
|
||||
_deprioritize_script_dir_in_sys_path()
|
||||
_drop_script_dir_from_sys_path()
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
config_filename: Final = "litellm.secrets"
|
||||
|
|
@ -881,7 +880,7 @@ class ProxyInitializationHelpers:
|
|||
default=False,
|
||||
help="Use prisma db push instead of prisma migrate for database schema updates",
|
||||
)
|
||||
@click.option("--local", is_flag=True, default=False, help="for local debugging")
|
||||
@click.option("--local", is_flag=True, default=False, help="no-op, kept for backwards compatibility")
|
||||
@click.option(
|
||||
"--skip_server_startup",
|
||||
is_flag=True,
|
||||
|
|
@ -1058,35 +1057,15 @@ def run_server(
|
|||
return
|
||||
|
||||
args: Final = locals()
|
||||
if local:
|
||||
from proxy_server import (
|
||||
try:
|
||||
from litellm.proxy.proxy_server import (
|
||||
KeyManagementSettings,
|
||||
ProxyConfig,
|
||||
app,
|
||||
save_worker_config,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
from .proxy_server import (
|
||||
KeyManagementSettings,
|
||||
ProxyConfig,
|
||||
app,
|
||||
save_worker_config,
|
||||
)
|
||||
except ModuleNotFoundError as e:
|
||||
raise ModuleNotFoundError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`")
|
||||
except ImportError as e:
|
||||
if "litellm[proxy]" in str(e):
|
||||
# user is missing a proxy dependency, ask them to pip install litellm[proxy]
|
||||
raise e
|
||||
else:
|
||||
# this is just a local/relative import error, user git cloned litellm
|
||||
from proxy_server import (
|
||||
KeyManagementSettings,
|
||||
ProxyConfig,
|
||||
app,
|
||||
save_worker_config,
|
||||
)
|
||||
except ModuleNotFoundError as e:
|
||||
raise ModuleNotFoundError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`") from e
|
||||
if version is True:
|
||||
ProxyInitializationHelpers._echo_litellm_version()
|
||||
return
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ import pytest
|
|||
|
||||
|
||||
import builtins
|
||||
import runpy
|
||||
import sys
|
||||
import types
|
||||
import urllib.parse as urlparse
|
||||
|
||||
|
|
@ -18,6 +20,7 @@ import yaml
|
|||
from uvicorn.config import LOOP_FACTORIES
|
||||
from uvicorn.importer import import_from_string
|
||||
|
||||
from litellm.proxy import proxy_cli
|
||||
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
|
||||
|
||||
|
||||
|
|
@ -636,6 +639,36 @@ class TestProxyInitializationHelpers:
|
|||
), f"exit_code={result.exit_code}, output={result.output}"
|
||||
mock_uvicorn_run.assert_called_once()
|
||||
|
||||
@patch("uvicorn.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False)
|
||||
def test_script_boot_imports_the_package_proxy_server(
|
||||
self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run
|
||||
):
|
||||
package_proxy_server = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
sibling_proxy_server = types.ModuleType("proxy_server")
|
||||
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{"proxy_server": sibling_proxy_server, "litellm.proxy.proxy_server": package_proxy_server},
|
||||
),
|
||||
patch.object(sys, "argv", ["proxy_cli.py", "--skip_server_startup"]),
|
||||
patch.object(sys, "path", list(sys.path)),
|
||||
pytest.raises(SystemExit) as exit_info,
|
||||
):
|
||||
runpy.run_path(proxy_cli.__file__, run_name="__main__")
|
||||
|
||||
assert exit_info.value.code == 0
|
||||
package_proxy_server.save_worker_config.assert_called_once()
|
||||
|
||||
@patch("uvicorn.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
|
|
@ -1945,13 +1978,18 @@ class TestProxyInitializationHelpers:
|
|||
mock_proxy_config_instance.get_config = mock_get_config
|
||||
mock_proxy_config.return_value = mock_proxy_config_instance
|
||||
|
||||
mock_proxy_server_module = MagicMock(app=mock_app)
|
||||
mock_proxy_server_module = MagicMock(
|
||||
app=mock_app,
|
||||
ProxyConfig=mock_proxy_config,
|
||||
KeyManagementSettings=mock_key_mgmt,
|
||||
save_worker_config=mock_save_worker_config,
|
||||
)
|
||||
|
||||
# Only remove DATABASE_URL and DIRECT_URL to prevent the database setup
|
||||
# code path from running. Do NOT use clear=True as it removes PATH, HOME,
|
||||
# etc., which causes imports inside run_server to break in CI (the real
|
||||
# litellm.proxy.proxy_server import at line 820 of proxy_cli.py has heavy
|
||||
# side effects that fail without a proper environment).
|
||||
# litellm.proxy.proxy_server import has heavy side effects that fail
|
||||
# without a proper environment).
|
||||
env_overrides = {
|
||||
"DATABASE_URL": "",
|
||||
"DIRECT_URL": "",
|
||||
|
|
@ -1967,18 +2005,7 @@ class TestProxyInitializationHelpers:
|
|||
with (
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"proxy_server": MagicMock(
|
||||
app=mock_app,
|
||||
ProxyConfig=mock_proxy_config,
|
||||
KeyManagementSettings=mock_key_mgmt,
|
||||
save_worker_config=mock_save_worker_config,
|
||||
),
|
||||
# Also mock litellm.proxy.proxy_server to prevent the real
|
||||
# import at line 820 of proxy_cli.py which has heavy side
|
||||
# effects (FastAPI app init, logging setup, etc.)
|
||||
"litellm.proxy.proxy_server": mock_proxy_server_module,
|
||||
},
|
||||
{"litellm.proxy.proxy_server": mock_proxy_server_module},
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue