From b7b20664c19e943efca1372505c4f26db5c82500 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 6 Mar 2026 18:09:57 -0800 Subject: [PATCH] Gflags worker parameters (#22931) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add LITELLM_WORKER_STARTUP_HOOKS for per-worker initialization (gflags support) Add support for running user-defined startup hooks in each worker process during proxy_startup_event. This enables re-initialization of in-process state (like gflags.FLAGS) that doesn't survive uvicorn worker spawning. Usage: export LITELLM_WORKER_STARTUP_HOOKS=mymodule:init_fn,other:setup_fn Hooks run early in proxy_startup_event (before config/DB loading). Supports both sync and async callables. Errors propagate to prevent broken workers from serving traffic. No-op when env var is unset. Includes 5 tests covering sync/async hooks, multiple hooks, error propagation, and no-hooks-set scenarios. Co-authored-by: Ishaan Jaff * docs: add Worker Startup Hooks page with gflags usage example - New docs page: docs/proxy/worker_startup_hooks.md - Explains the problem (per-process state lost in multi-worker deployments) - Full gflags example with wrapper module and startup script - Covers multiple hooks, async hooks, error behavior - Architecture diagram showing master→worker flow - Added LITELLM_WORKER_STARTUP_HOOKS to config_settings.md env var table - Added to sidebar under Setup & Deployment Co-authored-by: Ishaan Jaff * Update litellm/proxy/proxy_server.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: Cursor Agent Co-authored-by: Ishaan Jaff Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- docs/my-website/docs/proxy/config_settings.md | 1 + .../docs/proxy/worker_startup_hooks.md | 155 ++++++++++++++++++ docs/my-website/sidebars.js | 1 + litellm/proxy/proxy_server.py | 30 ++++ tests/test_litellm/proxy/test_proxy_cli.py | 141 ++++++++++++++++ 5 files changed, 328 insertions(+) create mode 100644 docs/my-website/docs/proxy/worker_startup_hooks.md diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index af868bc9f9..23c0b1498f 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -815,6 +815,7 @@ router_settings: | LITELLM_TOKEN | Access token for LiteLLM integration | LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES | When set to "true", routes OpenAI /v1/messages requests through chat/completions instead of the Responses API for Anthropic models. Can also be set via `litellm_settings.use_chat_completions_url_for_anthropic_messages` | LITELLM_USER_AGENT | Custom user agent string for LiteLLM API requests. Used for partner telemetry attribution +| LITELLM_WORKER_STARTUP_HOOKS | Comma-separated list of `module.path:function_name` callables to run in each worker process during startup. Runs early in the worker lifecycle (before config/DB loading). Useful for re-initializing per-process state like [gflags](https://github.com/google/python-gflags). See [Worker Startup Hooks](/proxy/worker_startup_hooks) for details | LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD | If true, prints the standard logging payload to the console - useful for debugging | LITELM_ENVIRONMENT | Environment for LiteLLM Instance. This is currently only logged to DeepEval to determine the environment for DeepEval integration. | LITELLM_ASYNCIO_QUEUE_MAXSIZE | Maximum size for asyncio queues (e.g. log queues, spend update queues, and cookbook examples such as realtime audio in `nova_sonic_realtime.py`). Bounds in-memory growth to prevent OOM. Default is 1000. diff --git a/docs/my-website/docs/proxy/worker_startup_hooks.md b/docs/my-website/docs/proxy/worker_startup_hooks.md new file mode 100644 index 0000000000..baf0e51ac9 --- /dev/null +++ b/docs/my-website/docs/proxy/worker_startup_hooks.md @@ -0,0 +1,155 @@ +# Worker Startup Hooks + +Use `LITELLM_WORKER_STARTUP_HOOKS` to run custom initialization functions in **each worker process** during proxy startup. This is essential when using multi-worker deployments (`--num_workers > 1`) with libraries that require per-process initialization, such as [gflags](https://github.com/google/python-gflags). + +## The Problem + +When running the LiteLLM proxy with multiple workers: + +```bash +litellm --config config.yaml --num_workers 4 +``` + +Each worker is a **separate process** spawned by uvicorn or gunicorn. Any in-process state initialized in the master process (before `run_server()`) is **not available** in worker processes. This includes: + +- [python-gflags](https://github.com/google/python-gflags) (`gflags.FLAGS`) +- [absl-py flags](https://abseil.io/docs/python/guides/flags) (`absl.flags.FLAGS`) +- Custom singleton registries or connection pools +- Any module-level state that requires explicit initialization + +## Usage + +Set the `LITELLM_WORKER_STARTUP_HOOKS` environment variable to a comma-separated list of `module.path:function_name` callables: + +```bash +export LITELLM_WORKER_STARTUP_HOOKS="my_module:my_init_function" +``` + +Each hook is called **early** in the worker startup lifecycle — before config loading, database setup, or any request handling. Both sync and async functions are supported. + +## Example: gflags Initialization + +### 1. Define your wrapper module + +```python title="my_litellm_wrapper.py" +import gflags +import json +import os +import sys +from typing import Optional, List, Any + + +def init_gflags( + usage: Optional[Any] = None, + raw_args: Optional[List[str]] = None, + known_only: bool = False, +) -> List[str]: + """Initialize gflags from command-line arguments.""" + try: + gflags.FLAGS.set_gnu_getopt(True) + if raw_args is None: + raw_args = sys.argv + argv = gflags.FLAGS(raw_args, known_only=known_only) + except gflags.Error as e: + if usage is None: + print("%s\nUsage: %s ARGS\n%s" % (e, sys.argv[0], gflags.FLAGS)) + else: + print(usage % dict(cmd=sys.argv[0], flags=gflags.FLAGS)) + sys.exit(1) + return argv + + +def init_gflags_for_worker(): + """Re-initialize gflags in each worker process. + + Reads the original sys.argv from the GFLAGS_ARGV env var + (set by the master process before starting the proxy). + """ + raw_args = json.loads(os.environ.get("GFLAGS_ARGV", "[]")) or sys.argv + init_gflags(raw_args=raw_args, known_only=True) +``` + +### 2. Start the proxy + +```python title="start_proxy.py" +import json +import os +import sys + +from my_litellm_wrapper import init_gflags + +# Store sys.argv so workers can re-parse the same flags +os.environ["GFLAGS_ARGV"] = json.dumps(sys.argv) + +# Tell LiteLLM to call our hook in each worker +os.environ["LITELLM_WORKER_STARTUP_HOOKS"] = "my_litellm_wrapper:init_gflags_for_worker" + +# Initialize gflags in the master process +init_gflags() + +# Start the proxy (programmatic invocation) +from litellm.proxy.proxy_cli import run_server + +run_server( + ["--config", "config.yaml", "--num_workers", "4"], + standalone_mode=False, +) +``` + +Or via shell: + +```bash +export GFLAGS_ARGV='["my_app", "--my_flag=value", "--batch_size=32"]' +export LITELLM_WORKER_STARTUP_HOOKS="my_litellm_wrapper:init_gflags_for_worker" + +litellm --config config.yaml --num_workers 4 +``` + +## How It Works + +``` +Master Process Worker Process (×N) +───────────────── ────────────────────── +1. init_gflags() 3. proxy_startup_event(): +2. run_server() → Read LITELLM_WORKER_STARTUP_HOOKS + → sets env vars → Import & call each hook + → uvicorn.run(workers=N) (gflags.FLAGS re-initialized ✓) + → spawns workers ──────────────────► → Continue with config/DB setup + → Ready to serve requests +``` + +- Hooks run at the **very beginning** of `proxy_startup_event` (the FastAPI lifespan), before config loading, database connections, or any other initialization. +- Environment variables set in the master process are **inherited** by worker processes (standard Unix fork/spawn behavior). +- If a hook **raises an exception**, the worker fails to start — this is intentional, since missing initialization (e.g., uninitialized gflags) would cause downstream errors. + +## Multiple Hooks + +Separate multiple hooks with commas: + +```bash +export LITELLM_WORKER_STARTUP_HOOKS="my_module:init_gflags,my_module:init_metrics,my_module:init_connections" +``` + +Hooks are executed **in order**, left to right. + +## Async Hooks + +Async functions are also supported — they are automatically awaited: + +```python +async def init_async_connections(): + """Example async hook for initializing async resources.""" + await setup_async_connection_pool() +``` + +```bash +export LITELLM_WORKER_STARTUP_HOOKS="my_module:init_async_connections" +``` + +## Reference + +| Environment Variable | Description | +|---|---| +| `LITELLM_WORKER_STARTUP_HOOKS` | Comma-separated `module.path:function_name` callables to run in each worker on startup | + +The hook format follows the standard Python entry point syntax: `module.path:function_name`, where `module.path` is a dotted Python import path and `function_name` is the name of the callable within that module. diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 91edd3118c..a77e905724 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -311,6 +311,7 @@ const sidebars = { "proxy/master_key_rotations", "proxy/model_management", "proxy/prod", + "proxy/worker_startup_hooks", "proxy/release_cycle", ], }, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 120f872b35..b7f11ec270 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1,6 +1,7 @@ import asyncio import copy import enum +import importlib import inspect import io import os @@ -768,6 +769,35 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 import json init_verbose_loggers() + + ## RUN WORKER STARTUP HOOKS (e.g., gflags initialization) ## + _startup_hooks_env = os.environ.get("LITELLM_WORKER_STARTUP_HOOKS", "") + if _startup_hooks_env: + for _hook_spec in _startup_hooks_env.split(","): + _hook_spec = _hook_spec.strip() + if not _hook_spec: + continue + try: + if ":" not in _hook_spec: + raise ValueError( + f"Invalid hook spec '{_hook_spec}': expected format is 'module.path:function_name'" + ) + _module_path, _func_name = _hook_spec.rsplit(":", 1) + _module = importlib.import_module(_module_path) + _hook_fn = getattr(_module, _func_name) + if inspect.iscoroutinefunction(_hook_fn): + await _hook_fn() + else: + _hook_fn() + verbose_proxy_logger.info( + "Worker startup hook '%s' executed successfully", _hook_spec + ) + except Exception as e: + verbose_proxy_logger.error( + "Worker startup hook '%s' failed: %s", _hook_spec, e + ) + raise + ## CHECK PREMIUM USER verbose_proxy_logger.debug( "litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format( diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index b83cbbd3dd..cf6511c18a 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -663,3 +663,144 @@ class TestHealthAppFactory: standalone_mode=False, ) mock_setup_database.assert_called_with(use_migrate=False) + + +# --- Module-level helpers for worker startup hook tests --- + +_dummy_hook_called = False + + +def _dummy_hook(): + """A simple sync hook used by test_should_run_worker_startup_hooks.""" + global _dummy_hook_called + _dummy_hook_called = True + + +_dummy_async_hook_called = False + + +async def _dummy_async_hook(): + """A simple async hook used by test_should_run_async_worker_startup_hook.""" + global _dummy_async_hook_called + _dummy_async_hook_called = True + + +def _failing_hook(): + """A hook that always raises, used by test_should_raise_on_failing_hook.""" + raise RuntimeError("Hook failed on purpose") + + +class TestWorkerStartupHooks: + """Tests for the LITELLM_WORKER_STARTUP_HOOKS mechanism in proxy_startup_event.""" + + @pytest.mark.asyncio + async def test_should_run_worker_startup_hooks(self): + """Sync worker startup hook is called during proxy_startup_event.""" + global _dummy_hook_called + _dummy_hook_called = False + + from litellm.proxy.proxy_server import proxy_startup_event + + env_overrides = { + "LITELLM_WORKER_STARTUP_HOOKS": "tests.test_litellm.proxy.test_proxy_cli:_dummy_hook", + } + # Remove DATABASE_URL to avoid real DB setup + clean_env = { + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env.update(env_overrides) + + with patch.dict(os.environ, clean_env, clear=True): + try: + async with proxy_startup_event(app=None) as _: + pass + except Exception: + pass # We expect errors after the hook (no DB, etc.) + + assert _dummy_hook_called is True, "Sync startup hook was not called" + + @pytest.mark.asyncio + async def test_should_run_async_worker_startup_hook(self): + """Async worker startup hook is awaited during proxy_startup_event.""" + global _dummy_async_hook_called + _dummy_async_hook_called = False + + from litellm.proxy.proxy_server import proxy_startup_event + + env_overrides = { + "LITELLM_WORKER_STARTUP_HOOKS": "tests.test_litellm.proxy.test_proxy_cli:_dummy_async_hook", + } + clean_env = { + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env.update(env_overrides) + + with patch.dict(os.environ, clean_env, clear=True): + try: + async with proxy_startup_event(app=None) as _: + pass + except Exception: + pass + + assert _dummy_async_hook_called is True, "Async startup hook was not called" + + @pytest.mark.asyncio + async def test_should_raise_on_failing_worker_startup_hook(self): + """A failing worker startup hook propagates the error.""" + from litellm.proxy.proxy_server import proxy_startup_event + + env_overrides = { + "LITELLM_WORKER_STARTUP_HOOKS": "tests.test_litellm.proxy.test_proxy_cli:_failing_hook", + } + clean_env = { + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env.update(env_overrides) + + with patch.dict(os.environ, clean_env, clear=True): + with pytest.raises(RuntimeError, match="Hook failed on purpose"): + async with proxy_startup_event(app=None) as _: + pass + + def test_should_skip_when_no_hooks_set(self): + """When LITELLM_WORKER_STARTUP_HOOKS is not set, no hooks are executed.""" + global _dummy_hook_called + _dummy_hook_called = False + + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("LITELLM_WORKER_STARTUP_HOOKS", None) + # The hook block should be skipped entirely when env var is absent + assert "LITELLM_WORKER_STARTUP_HOOKS" not in os.environ + # Verify that an empty env var value also results in no hook execution + assert os.environ.get("LITELLM_WORKER_STARTUP_HOOKS", "") == "" + + @pytest.mark.asyncio + async def test_should_run_multiple_hooks(self): + """Multiple comma-separated hooks are all called.""" + global _dummy_hook_called, _dummy_async_hook_called + _dummy_hook_called = False + _dummy_async_hook_called = False + + from litellm.proxy.proxy_server import proxy_startup_event + + hooks = ( + "tests.test_litellm.proxy.test_proxy_cli:_dummy_hook," + "tests.test_litellm.proxy.test_proxy_cli:_dummy_async_hook" + ) + env_overrides = { + "LITELLM_WORKER_STARTUP_HOOKS": hooks, + } + clean_env = { + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env.update(env_overrides) + + with patch.dict(os.environ, clean_env, clear=True): + try: + async with proxy_startup_event(app=None) as _: + pass + except Exception: + pass + + assert _dummy_hook_called is True, "First hook was not called" + assert _dummy_async_hook_called is True, "Second hook was not called"