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:
mohittalele 2025-05-19 17:00:41 +02:00 • committed by GitHub
parent 8a880ccd2f
commit d36b6fa60b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 59 additions and 11 deletions

View file

@ -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,

View file

@ -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()