mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
add skip server startup flag to cli (#10665)
* add flag * fix ruff linting * add unit test * fix ruff errors * fix lintings issues * fix linting errors * add global noqa for print
This commit is contained in:
parent
8a880ccd2f
commit
d36b6fa60b
2 changed files with 59 additions and 11 deletions
|
|
@ -1,3 +1,4 @@
|
|||
# ruff: noqa: T201
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
|
|
@ -458,6 +459,12 @@ class ProxyInitializationHelpers:
|
|||
help="Use prisma migrate instead of prisma db push for database schema updates",
|
||||
)
|
||||
@click.option("--local", is_flag=True, default=False, help="for local debugging")
|
||||
@click.option(
|
||||
"--skip_server_startup",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Skip starting the server after setup (useful for migrations only)",
|
||||
)
|
||||
def run_server( # noqa: PLR0915
|
||||
host,
|
||||
port,
|
||||
|
|
@ -493,6 +500,7 @@ def run_server( # noqa: PLR0915
|
|||
ssl_certfile_path,
|
||||
log_config,
|
||||
use_prisma_migrate,
|
||||
skip_server_startup,
|
||||
):
|
||||
args = locals()
|
||||
if local:
|
||||
|
|
@ -751,6 +759,11 @@ def run_server( # noqa: PLR0915
|
|||
# DO NOT DELETE - enables global variables to work across files
|
||||
from litellm.proxy.proxy_server import app # noqa
|
||||
|
||||
# Skip server startup if requested (after all setup is done)
|
||||
if skip_server_startup:
|
||||
print("LiteLLM: Setup complete. Skipping server startup as requested.") # noqa
|
||||
return
|
||||
|
||||
uvicorn_args = ProxyInitializationHelpers._get_default_unvicorn_init_args(
|
||||
host=host,
|
||||
port=port,
|
||||
|
|
|
|||
|
|
@ -1,22 +1,12 @@
|
|||
import importlib
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import click
|
||||
import httpx
|
||||
from unittest.mock import MagicMock, patch
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
) # Adds the parent directory to the system-path
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.proxy_cli import ProxyInitializationHelpers
|
||||
|
||||
|
||||
|
|
@ -165,3 +155,48 @@ class TestProxyInitializationHelpers:
|
|||
# Test on Linux
|
||||
with patch("sys.platform", "linux"):
|
||||
assert ProxyInitializationHelpers._get_loop_type() == "uvloop"
|
||||
|
||||
@patch("uvicorn.run")
|
||||
@patch("builtins.print")
|
||||
def test_skip_server_startup(self, mock_print, mock_uvicorn_run):
|
||||
"""Test that the skip_server_startup flag prevents server startup when True"""
|
||||
from click.testing import CliRunner
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
mock_app = MagicMock()
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_key_mgmt = MagicMock()
|
||||
mock_save_worker_config = MagicMock()
|
||||
|
||||
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
|
||||
)
|
||||
}), \
|
||||
patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args") as mock_get_args:
|
||||
|
||||
mock_get_args.return_value = {
|
||||
"app": "litellm.proxy.proxy_server:app",
|
||||
"host": "localhost",
|
||||
"port": 8000
|
||||
}
|
||||
|
||||
result = runner.invoke(run_server, ["--local", "--skip_server_startup"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_uvicorn_run.assert_not_called()
|
||||
mock_print.assert_any_call("LiteLLM: Setup complete. Skipping server startup as requested.")
|
||||
|
||||
mock_uvicorn_run.reset_mock()
|
||||
mock_print.reset_mock()
|
||||
|
||||
result = runner.invoke(run_server, ["--local"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_uvicorn_run.assert_called_once()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue