From d36b6fa60b6190696f89e3f6a6b5fa8f93ed1928 Mon Sep 17 00:00:00 2001
From: mohittalele <mohittalele@gmail.com>
Date: Mon, 19 May 2025 17:00:41 +0200
Subject: [PATCH] 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
---
 litellm/proxy/proxy_cli.py            | 13 ++++++
 tests/litellm/proxy/test_proxy_cli.py | 57 +++++++++++++++++++++------
 2 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
index de78247dafa..9f3123c0b28 100644
--- a/litellm/proxy/proxy_cli.py
+++ b/litellm/proxy/proxy_cli.py
@@ -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,
diff --git a/tests/litellm/proxy/test_proxy_cli.py b/tests/litellm/proxy/test_proxy_cli.py
index 6e1d70553fc..55509a4f78a 100644
--- a/tests/litellm/proxy/test_proxy_cli.py
+++ b/tests/litellm/proxy/test_proxy_cli.py
@@ -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()
+