diff --git a/src/lfx/src/lfx/base/models/unified_models/credentials.py b/src/lfx/src/lfx/base/models/unified_models/credentials.py index 9b28aa366d..e3b9589947 100644 --- a/src/lfx/src/lfx/base/models/unified_models/credentials.py +++ b/src/lfx/src/lfx/base/models/unified_models/credentials.py @@ -217,7 +217,7 @@ def get_all_variables_for_provider(user_id: UUID | str | None, provider: str) -> continue # Honor the request's no-env-fallback contract: a served flow under # no_env_fallback must stay isolated from process-wide credentials even on - # this post-DB-miss rotation fallback (concerns A3). + # this post-DB-miss rotation fallback. if is_env_fallback_disabled(): continue env_value = _env_value_for(var_key) diff --git a/src/lfx/src/lfx/cli/_running_commands.py b/src/lfx/src/lfx/cli/_running_commands.py index 6c64186913..d65995715d 100644 --- a/src/lfx/src/lfx/cli/_running_commands.py +++ b/src/lfx/src/lfx/cli/_running_commands.py @@ -136,6 +136,27 @@ def register(app: typer.Typer) -> None: "Defaults to in-memory only when omitted." ), ), + max_requests: int | None = typer.Option( + None, + "--max-requests", + help=( + "Recycle each worker after this many requests (gunicorn, Unix-only, --workers > 1). " + "Set to 1 for per-request worker recycling. Default (unset) means workers are never " + "recycled. Not supported on Windows, where multi-worker serving uses uvicorn. " + "For full per-request isolation, combine with --limit-concurrency 1." + ), + ), + limit_concurrency: int | None = typer.Option( + None, + "--limit-concurrency", + help=( + "Max in-flight requests per worker (--workers > 1); excess get HTTP 503. " + "Recycling alone does NOT stop a worker from accepting a 2nd concurrent request, so " + "without this two requests may share one process/os.environ. Set to 1 (with " + "--max-requests 1) so each worker handles exactly one request in its own process — " + "strict cross-request isolation. Default (unset) means unlimited concurrency." + ), + ), *, stdin: bool = typer.Option( False, @@ -178,4 +199,6 @@ def register(app: typer.Typer) -> None: stdin=stdin, check_variables=check_variables, no_env_fallback=no_env_fallback, + max_requests=max_requests, + limit_concurrency=limit_concurrency, ) diff --git a/src/lfx/src/lfx/cli/commands.py b/src/lfx/src/lfx/cli/commands.py index 3f79c7ea63..d0d53f3381 100644 --- a/src/lfx/src/lfx/cli/commands.py +++ b/src/lfx/src/lfx/cli/commands.py @@ -180,6 +180,27 @@ def serve_command( "Defaults to in-memory only when omitted." ), ), + max_requests: int | None = typer.Option( + None, + "--max-requests", + help=( + "Recycle each worker after this many requests (gunicorn, Unix-only, --workers > 1). " + "Set to 1 for per-request worker recycling. Default (unset) means workers are never " + "recycled. Not supported on Windows, where multi-worker serving uses uvicorn. " + "For full per-request isolation, combine with --limit-concurrency 1." + ), + ), + limit_concurrency: int | None = typer.Option( + None, + "--limit-concurrency", + help=( + "Max in-flight requests per worker (--workers > 1); excess get HTTP 503. " + "Recycling alone does NOT stop a worker from accepting a 2nd concurrent request, so " + "without this two requests may share one process/os.environ. Set to 1 (with " + "--max-requests 1) so each worker handles exactly one request in its own process — " + "strict cross-request isolation. Default (unset) means unlimited concurrency." + ), + ), *, stdin: bool = typer.Option( False, # noqa: FBT003 @@ -255,10 +276,9 @@ def serve_command( if workers > 1 and flow_dir is None: typer.echo( "Warning: --workers > 1 without --flow-dir means each worker has an isolated " - "in-memory registry. Flows uploaded to one worker will not be visible to others, " - "and because workers are recycled after each request for per-request isolation, " - "uploaded flows do not survive at all without a shared store. " - "Pass --flow-dir to enable shared flow storage across workers.", + "in-memory registry. Flows uploaded to one worker will not be visible to others " + "(and with --max-requests recycling, uploaded flows do not survive worker recycling " + "at all). Pass --flow-dir to enable shared flow storage across workers.", err=True, ) @@ -361,6 +381,8 @@ def serve_command( script_paths=script_paths, temp_file_to_cleanup=temp_file_to_cleanup, verbose_print=verbose_print, + max_requests=max_requests, + limit_concurrency=limit_concurrency, ) else: serve_app = create_multi_serve_app(registry=registry) @@ -389,34 +411,40 @@ def _launch_workers( script_paths: list[str] | None, temp_file_to_cleanup: str | None, verbose_print: Callable[[str], None], + max_requests: int | None, + limit_concurrency: int | None, ) -> None: - """Launch ``workers`` isolated worker processes for ``lfx serve --workers N``. + """Launch ``workers`` worker processes for ``lfx serve --workers N``. On Unix this runs gunicorn with ``preload_app=True`` (the master builds the - warm app once and forks workers via copy-on-write) and ``max_requests=1`` so - each request executes in a freshly-forked worker that is recycled after one - request — making cross-request ``os.environ`` credential leakage structurally + warm app once and forks workers via copy-on-write). ``max_requests`` controls + per-request recycling: ``None`` (the default) maps to gunicorn's ``0`` (workers + are never recycled — warm and shared, but no per-request isolation), while + ``--max-requests 1`` recycles each worker after one request. + + ``limit_concurrency`` caps the in-flight requests a single worker accepts + (excess get HTTP 503). Recycling alone does NOT prevent a worker from accepting + a second concurrent request — two requests can then share one process / + ``os.environ``. ``--limit-concurrency 1`` together with ``--max-requests 1`` + closes that window: each worker handles exactly one request, in its own process, + then recycles — making cross-request ``os.environ`` leakage structurally impossible. - gunicorn is Unix-only, so on Windows multi-worker serving is refused with a - clear message; run with ``--workers 1`` or deploy on Linux/macOS instead. + gunicorn is Unix-only. On Windows it cannot run at all, so multi-worker serving + falls back to uvicorn's own multi-worker supervisor (no preload, no per-request + recycling). ``--limit-concurrency`` is still honored there (uvicorn-native), but + ``--max-requests`` (recycling) is refused, since it cannot be supported. """ from lfx.cli.serve_app import ( _SERVE_FLOW_DIR_ENV, + _SERVE_LIMIT_CONCURRENCY_ENV, _SERVE_NO_ENV_FALLBACK_ENV, _SERVE_STARTUP_PATHS_ENV, ) - if sys.platform == "win32": - verbose_print( - "Multi-worker isolated serving uses gunicorn, which is not available on " - "Windows. Run with --workers 1 (single worker, no per-request process " - "recycling) or deploy on Linux/macOS for full isolation." - ) - raise typer.Exit(1) - - # Set env vars so the gunicorn preload master's build_registry_from_env() can - # reconstruct config. When flow_dir is set, startup flows are already in the + # Set env vars so each worker can reconstruct config: the gunicorn preload + # master via build_registry_from_env(), or each uvicorn factory worker via + # create_serve_app(). When flow_dir is set, startup flows are already in the # store (written by _build_serve_registry) so workers load them via # warm_from_store(). When flow_dir is NOT set, workers re-read original files. os.environ[_SERVE_FLOW_DIR_ENV] = str(flow_dir) if flow_dir else "" @@ -428,29 +456,68 @@ def _launch_workers( elif temp_file_to_cleanup: startup_paths_for_workers = [temp_file_to_cleanup] os.environ[_SERVE_STARTUP_PATHS_ENV] = json.dumps(startup_paths_for_workers) + # Read per worker by LFXUvicornWorker (Unix); passed to uvicorn.run on Windows. + if limit_concurrency is not None: + os.environ[_SERVE_LIMIT_CONCURRENCY_ENV] = str(limit_concurrency) try: - from lfx.cli.serve_gunicorn import LFXGunicornApp + if sys.platform == "win32": + if max_requests is not None: + verbose_print( + "Error: --max-requests enables per-request worker recycling via gunicorn, " + "which is not available on Windows. Omit --max-requests to run multi-worker " + "without isolation, or deploy on Linux/macOS for per-request isolation." + ) + raise typer.Exit(1) + # gunicorn cannot run on Windows; fall back to uvicorn's multi-worker + # supervisor. No preload/COW and no per-request recycling (no isolation), + # though --limit-concurrency is still honored (uvicorn-native). + verbose_print( + "Note: multi-worker serving on Windows uses uvicorn (no per-request recycling); " + "deploy on Linux/macOS and pass --max-requests 1 for full isolation." + ) + # +1: uvicorn counts the active connection, so its limit must exceed the + # desired in-flight count (limit_concurrency=1 would reject everything). + uvicorn_limit = (limit_concurrency + 1) if limit_concurrency is not None else None + uvicorn.run( + "lfx.cli.serve_app:create_serve_app", + host=host, + port=port, + workers=workers, + log_level=log_level, + factory=True, + limit_concurrency=uvicorn_limit, + ) + else: + from lfx.cli.serve_gunicorn import LFXGunicornApp - LFXGunicornApp( - "lfx.cli.serve_preloaded_app:app", - { - "bind": f"{host}:{port}", - "workers": workers, - "worker_class": "uvicorn.workers.UvicornWorker", - "preload_app": True, - "max_requests": 1, - "max_requests_jitter": 0, - "loglevel": log_level, - # Sized for long synchronous flows; gunicorn's default 30s timeout - # would kill long LLM flows. Expose --timeout later if needed (D4). - "timeout": 120, - }, - ).run() + LFXGunicornApp( + "lfx.cli.serve_preloaded_app:app", + { + "bind": f"{host}:{port}", + "workers": workers, + # Custom worker applies LFX_SERVE_LIMIT_CONCURRENCY (gunicorn's + # UvicornWorker cannot forward uvicorn's limit_concurrency). + "worker_class": "lfx.cli.serve_gunicorn.LFXUvicornWorker", + "preload_app": True, + # None -> 0 (gunicorn's default: never recycle). 1 -> recycle per request. + "max_requests": max_requests if max_requests is not None else 0, + "max_requests_jitter": 0, + "loglevel": log_level, + # Sized for long synchronous flows; gunicorn's default 30s timeout + # would kill long LLM flows. Expose --timeout later if needed. + "timeout": 120, + }, + ).run() finally: # Only remove the keys we set above — a prefix sweep would also delete any # LFX_SERVE_* var the operator intentionally exported before launch. - for k in (_SERVE_FLOW_DIR_ENV, _SERVE_NO_ENV_FALLBACK_ENV, _SERVE_STARTUP_PATHS_ENV): + for k in ( + _SERVE_FLOW_DIR_ENV, + _SERVE_NO_ENV_FALLBACK_ENV, + _SERVE_STARTUP_PATHS_ENV, + _SERVE_LIMIT_CONCURRENCY_ENV, + ): os.environ.pop(k, None) diff --git a/src/lfx/src/lfx/cli/serve_app.py b/src/lfx/src/lfx/cli/serve_app.py index 08f6db2707..ec421631a6 100644 --- a/src/lfx/src/lfx/cli/serve_app.py +++ b/src/lfx/src/lfx/cli/serve_app.py @@ -54,14 +54,17 @@ _SERVE_ENV_PREFIX = "LFX_SERVE_" _SERVE_FLOW_DIR_ENV = f"{_SERVE_ENV_PREFIX}FLOW_DIR" _SERVE_NO_ENV_FALLBACK_ENV = f"{_SERVE_ENV_PREFIX}NO_ENV_FALLBACK" _SERVE_STARTUP_PATHS_ENV = f"{_SERVE_ENV_PREFIX}STARTUP_PATHS" +# uvicorn limit_concurrency applied per worker by LFXUvicornWorker (gunicorn's +# UvicornWorker doesn't expose it). Set by serve_command; read in each worker. +_SERVE_LIMIT_CONCURRENCY_ENV = f"{_SERVE_ENV_PREFIX}LIMIT_CONCURRENCY" api_key_query = APIKeyQuery(name=API_KEY_NAME, scheme_name="API key query", auto_error=False) api_key_header = APIKeyHeader(name=API_KEY_NAME, scheme_name="API key header", auto_error=False) -# One in-flight flow execution per worker process. With gunicorn max_requests=1 a -# worker is recycled after a single request, but an async UvicornWorker can accept -# a second connection before recycling; this guard ensures the env-sensitive -# execution section (where request-scoped vars are active and flow code may touch -# os.environ) is never entered by two requests in the same process at once. +# One in-flight flow execution per worker process. An async UvicornWorker can +# accept a second connection while one request is mid-flight; this guard ensures +# the env-sensitive execution section (where request-scoped vars are active and +# flow code may touch os.environ) is never entered by two requests in the same +# process at once. Matters most with per-request recycling (--max-requests 1). _EXECUTE_GUARD = asyncio.Semaphore(1) @@ -304,10 +307,9 @@ class FlowRegistry: raw_json = self._store.read(flow_id) if raw_json is None: return None - # Cache-miss reconstruction from the store. Under gunicorn max_requests=1 - # (per-request worker recycling), a flow not folded into the preload image - # pays this graph-rebuild cost on every request — log it so that per-request - # overhead is observable (concerns B2). + # Cache-miss reconstruction from the store. With per-request worker recycling + # (gunicorn --max-requests 1), a flow not folded into the preload image pays + # this graph-rebuild cost on every request — log it so the overhead is observable. logger.info(f"Reconstructing flow '{flow_id}' from store on cache miss") graph, meta = self._reconstruct(flow_id, raw_json) # Cache under the authoritative JSON id so requests by UUID find it. diff --git a/src/lfx/src/lfx/cli/serve_gunicorn.py b/src/lfx/src/lfx/cli/serve_gunicorn.py index 31e36f25a8..2381433408 100644 --- a/src/lfx/src/lfx/cli/serve_gunicorn.py +++ b/src/lfx/src/lfx/cli/serve_gunicorn.py @@ -1,21 +1,128 @@ """gunicorn process manager for ``lfx serve --workers N`` on Unix. Uses ``preload_app=True`` so the master builds the warm app once and forks -workers (copy-on-write). ``max_requests=1`` recycles each worker after a single +workers (copy-on-write). ``--max-requests 1`` recycles each worker after a single request so no request inherits another's process environment. + +The ``pre_fork`` hook (mirrors ``langflow.server.LangflowApplication``) runs in +the master immediately before each worker is forked: it ``gc.freeze()``s the +fully-built, warmed preload heap so workers inherit it via copy-on-write without +the GC dirtying shared pages (preserving the preload memory savings), and warns +about any fork-unsafe state — background threads or open TCP connections (e.g. a +DB connection pool) left alive in the master, which would be silently dead or +corrupt in the forked workers. """ from __future__ import annotations from gunicorn.app.base import BaseApplication +from uvicorn.workers import UvicornWorker + + +class LFXUvicornWorker(UvicornWorker): + """UvicornWorker that applies ``LFX_SERVE_LIMIT_CONCURRENCY`` per worker. + + gunicorn's stock ``UvicornWorker`` maps ``max_requests`` (recycling) but never + exposes uvicorn's ``limit_concurrency``, and gunicorn has no setting that + forwards to it. So we read the limit from the ``LFX_SERVE_*`` environment (set + by ``serve_command`` and inherited by each forked worker) and apply it to the + worker's uvicorn ``Config``. + + The env var holds the user-facing cap N (max in-flight requests per worker). + uvicorn's ``limit_concurrency`` counts the active connection itself, so it + rejects once ``len(connections) >= limit`` — i.e. ``limit_concurrency=1`` + rejects EVERY request. To allow exactly N in-flight we set uvicorn's value to + ``N + 1`` (empirically verified). With ``--limit-concurrency 1`` a worker then + accepts one in-flight request and returns HTTP 503 for a concurrent second — + combined with ``--max-requests 1`` this guarantees no two requests ever share a + worker process / ``os.environ``. + """ + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._apply_limit_concurrency(self.config) + + @staticmethod + def _apply_limit_concurrency(config) -> None: + import os + + from lfx.cli.serve_app import _SERVE_LIMIT_CONCURRENCY_ENV + + raw = os.environ.get(_SERVE_LIMIT_CONCURRENCY_ENV) + if raw: + # +1: uvicorn counts the active connection, so its limit must exceed the + # desired in-flight count (limit_concurrency=1 would reject everything). + config.limit_concurrency = int(raw) + 1 class LFXGunicornApp(BaseApplication): def __init__(self, app_import_string: str, options: dict) -> None: self._app_import_string = app_import_string self._options = options or {} + # Freeze + fork-safety diagnostics happen in the master right before each + # fork. Doing the freeze here (not at app-import time) means the fully + # warmed preload heap is frozen immediately before fork, so workers still + # inherit it via copy-on-write — the preload memory savings are preserved. + self._options.setdefault("pre_fork", self.pre_fork) super().__init__() + # Thread-name prefixes known to be benign before fork: they never survive into + # workers and produce no side effects when their fd is inherited. + _BENIGN_THREAD_PREFIXES = ( + "OTel", # OpenTelemetry SDK (BatchSpanProcessor, etc.) + "opentelemetry", # alternate OTel naming + "prometheus", # Prometheus client background threads + "loguru", # loguru enqueue=True worker + "asyncio", # event-loop helper threads (Python internals) + "ThreadPoolExecutor", # stdlib executor - harmless in parent + "concurrent.futures", # same pool, different prefix + ) + + @classmethod + def _is_benign_thread(cls, thread) -> bool: + return any(thread.name.startswith(prefix) for prefix in cls._BENIGN_THREAD_PREFIXES) + + @classmethod + def pre_fork(cls, server, _worker) -> None: + """Run in the master before each fork: warn on fork-unsafe state, then freeze. + + Any non-benign background thread or non-listening TCP connection still alive + here will be dead/corrupt in the forked workers (e.g. a DB engine pool opened + during preload). lfx serve is DB-less by default and opens nothing fork-unsafe + at preload, so this is normally silent — but it loudly flags the day something + leaves a live connection or thread in the master. + """ + import gc + import threading + + non_main = [t for t in threading.enumerate() if t.is_alive() and t is not threading.main_thread()] + suspicious = [t for t in non_main if not cls._is_benign_thread(t)] + if suspicious: + server.log.warning( + "Ghost threads found before fork (these will be dead in workers): %s", + [t.name for t in suspicious], + ) + + try: + import psutil + + ghost_conns = [c for c in psutil.Process().net_connections(kind="tcp") if c.status != "LISTEN"] + if ghost_conns: + server.log.warning( + "Ghost TCP connections found before fork (these will be dead/corrupt in workers): %s", + [(c.laddr, c.raddr, c.status) for c in ghost_conns], + ) + except ImportError: + server.log.debug("psutil not installed; skipping ghost TCP connection check") + except Exception as e: # noqa: BLE001 + server.log.warning("Failed to inspect TCP connections before fork: %s", e) + + try: + gc.collect() + except Exception as e: # noqa: BLE001 + server.log.warning("gc.collect() raised during pre-fork hook: %s", e) + gc.freeze() + def load_config(self) -> None: for key, value in self._options.items(): if value is not None and key in self.cfg.settings: diff --git a/src/lfx/src/lfx/cli/serve_preloaded_app.py b/src/lfx/src/lfx/cli/serve_preloaded_app.py index 01b75f614b..8d00fc4495 100644 --- a/src/lfx/src/lfx/cli/serve_preloaded_app.py +++ b/src/lfx/src/lfx/cli/serve_preloaded_app.py @@ -3,18 +3,17 @@ gunicorn's master imports this module ONCE before forking. Building + warming the registry here means the library, imported component modules, and warm flow graphs all live in the master heap and are inherited by every forked worker via -copy-on-write. ``gc.freeze()`` keeps that heap out of the GC so refcount/GC -traversal doesn't dirty the shared pages. Workers are recycled per request -(``max_requests=1``); each new fork re-inherits this frozen, warm image. -""" +copy-on-write. Workers are recycled per request when ``--max-requests 1`` is set; +each new fork re-inherits this warm image. -import gc +The ``gc.freeze()`` that keeps this heap out of the GC (so refcount/GC traversal +doesn't dirty the shared pages) is NOT done here — it runs in the gunicorn +``pre_fork`` hook (:meth:`lfx.cli.serve_gunicorn.LFXGunicornApp.pre_fork`), i.e. +in the master immediately before each fork, which is the correct point to freeze +the fully-built heap while preserving copy-on-write sharing. +""" from lfx.cli.serve_app import build_registry_from_env, create_multi_serve_app registry = build_registry_from_env() app = create_multi_serve_app(registry=registry) - -# Must run in the master, after the full preload, before any fork. Placing it at -# module-import bottom guarantees that under --preload. -gc.freeze() diff --git a/src/lfx/tests/unit/cli/test_serve.py b/src/lfx/tests/unit/cli/test_serve.py index 731d7e21c1..d2f1cfeed3 100644 --- a/src/lfx/tests/unit/cli/test_serve.py +++ b/src/lfx/tests/unit/cli/test_serve.py @@ -678,13 +678,8 @@ def test_startup_scan_store_flows_accessible_lazily(tmp_path): assert result[1].title == "Pre-existing" -def test_serve_command_passes_workers_to_gunicorn(): - """--workers N must be forwarded to the gunicorn launcher as workers. - - Multi-worker serving runs gunicorn with preload + max_requests=1 so each - request executes in a freshly-forked, recycled worker. The app is loaded by - the preload master via the ``serve_preloaded_app`` import string. - """ +def _run_serve_capturing_gunicorn(*, workers, max_requests, captured, limit_concurrency=None): + """Run serve_command under a stubbed gunicorn launcher, capturing options + launch env.""" import json import os import tempfile @@ -697,15 +692,13 @@ def test_serve_command_passes_workers_to_gunicorn(): mock_graph = MagicMock() mock_graph.context = {} - captured: dict = {} - class FakeGunicornApp: def __init__(self, app_import_string, options): captured["app_import_string"] = app_import_string captured["options"] = options def run(self): - pass + captured["env"] = dict(os.environ) # capture env set before launch with tempfile.TemporaryDirectory() as tmp: p = Path(tmp) / "flow.json" @@ -715,11 +708,13 @@ def test_serve_command_passes_workers_to_gunicorn(): patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), patch("lfx.cli.serve_gunicorn.LFXGunicornApp", FakeGunicornApp), ): + # Direct call (not via CLI): typer defaults are not applied, so pass + # max_requests / limit_concurrency explicitly even when None. serve_command( script_paths=[str(p)], host="127.0.0.1", port=9999, - workers=4, + workers=workers, verbose=False, env_file=None, log_level="warning", @@ -728,14 +723,63 @@ def test_serve_command_passes_workers_to_gunicorn(): stdin=False, check_variables=False, no_env_fallback=False, + max_requests=max_requests, + limit_concurrency=limit_concurrency, ) + +def test_serve_command_passes_workers_to_gunicorn(): + """--workers N is forwarded to gunicorn; default --max-requests means no recycling. + + Multi-worker serving runs gunicorn with preload; the app is loaded by the + preload master via the ``serve_preloaded_app`` import string. With no + --max-requests, max_requests defaults to 0 (gunicorn: never recycle, no + per-request isolation). + """ + captured: dict = {} + _run_serve_capturing_gunicorn(workers=4, max_requests=None, captured=captured) + assert captured["options"]["workers"] == 4 # The preload master loads the warm app from this import string, not a factory. assert captured["app_import_string"] == "lfx.cli.serve_preloaded_app:app" - # Per-request isolation: preload + recycle-after-one-request. assert captured["options"]["preload_app"] is True + # Default: no recycling (0), so workers persist (no per-request isolation). + assert captured["options"]["max_requests"] == 0 + # The custom worker (applies limit_concurrency) is used, not the stock one. + assert captured["options"]["worker_class"] == "lfx.cli.serve_gunicorn.LFXUvicornWorker" + + +def test_serve_command_max_requests_propagates_to_gunicorn(): + """--max-requests N is forwarded to gunicorn as max_requests (per-request recycling).""" + captured: dict = {} + _run_serve_capturing_gunicorn(workers=4, max_requests=1, captured=captured) + assert captured["options"]["max_requests"] == 1 + assert captured["options"]["preload_app"] is True + + +def test_serve_command_limit_concurrency_sets_env_for_workers(): + """--limit-concurrency N is exported so each LFXUvicornWorker applies it.""" + from lfx.cli.serve_app import _SERVE_LIMIT_CONCURRENCY_ENV + + captured: dict = {} + _run_serve_capturing_gunicorn(workers=4, max_requests=1, limit_concurrency=1, captured=captured) + + # Set in the environment before launch (inherited by forked workers), then cleaned up. + assert captured["env"].get(_SERVE_LIMIT_CONCURRENCY_ENV) == "1" + import os as _os + + assert _SERVE_LIMIT_CONCURRENCY_ENV not in _os.environ # cleaned up after launch + + +def test_serve_command_no_limit_concurrency_leaves_env_unset(): + """Without --limit-concurrency, the env var is never set.""" + from lfx.cli.serve_app import _SERVE_LIMIT_CONCURRENCY_ENV + + captured: dict = {} + _run_serve_capturing_gunicorn(workers=4, max_requests=None, limit_concurrency=None, captured=captured) + + assert _SERVE_LIMIT_CONCURRENCY_ENV not in captured["env"] def test_serve_command_sets_startup_paths_env_for_multi_worker(tmp_path): diff --git a/src/lfx/tests/unit/cli/test_serve_process_model.py b/src/lfx/tests/unit/cli/test_serve_process_model.py index d8a507dbfb..5e9d14055f 100644 --- a/src/lfx/tests/unit/cli/test_serve_process_model.py +++ b/src/lfx/tests/unit/cli/test_serve_process_model.py @@ -6,10 +6,13 @@ Covers the preload/warm/guard/fork-safety pieces of the gunicorn - ``build_registry_from_env`` builds and warms a registry from the LFX_SERVE_* environment (the single source reused by the uvicorn factory and the gunicorn preload master). -- ``serve_preloaded_app`` builds the warm app at import time and ``gc.freeze()``s. -- ``LFXGunicornApp`` translates options into a gunicorn config. +- ``serve_preloaded_app`` builds the warm app at import time (freezing happens in + the gunicorn ``pre_fork`` hook, not at import). +- ``LFXGunicornApp`` translates options into a gunicorn config and runs the + ``pre_fork`` fork-safety hook (gc.freeze + ghost-thread/connection diagnostics). - The single-in-flight execute guard serializes the env-sensitive section. -- Fork-safety regressions (DB-less serve, uuid4-based ids). +- Fork-safety regressions: DB-less serve, uuid4-based ids, preload must not open + the DB or start background threads, and the pre_fork guardrail must fire. """ import json @@ -53,8 +56,7 @@ def test_build_registry_from_env_warms_store(tmp_path, monkeypatch): assert registry._no_env_fallback is True -def test_preloaded_app_module_builds_warm_app_and_freezes(tmp_path, monkeypatch): - import gc +def test_preloaded_app_module_builds_warm_app(tmp_path, monkeypatch): import importlib from fastapi import FastAPI @@ -73,7 +75,52 @@ def test_preloaded_app_module_builds_warm_app_and_freezes(tmp_path, monkeypatch) assert isinstance(mod.app, FastAPI) assert mod.registry.get("flow-1") is not None - assert gc.get_freeze_count() > 0 # gc.freeze() ran at import + # The module does NOT freeze at import; freezing happens in the gunicorn + # pre_fork hook (see test_gunicorn_pre_fork_freezes_heap). + + +@pytest.mark.skipif(sys.platform == "win32", reason="gunicorn is Unix-only") +def test_gunicorn_pre_fork_freezes_heap_and_runs_clean(monkeypatch): + """The pre_fork hook calls gc.freeze() (preserving COW) and runs without raising. + + Spies on gc.freeze directly rather than gc.get_freeze_count(): the hook also + calls gc.collect(), which by itself bumps the freeze count, so a count-based + assertion would pass even if gc.freeze() were removed. + """ + import gc + from unittest.mock import MagicMock + + from lfx.cli.serve_gunicorn import LFXGunicornApp + + class _FakeLog: + def warning(self, *_a, **_k): + pass + + def debug(self, *_a, **_k): + pass + + class _FakeServer: + log = _FakeLog() + + freeze_spy = MagicMock() + monkeypatch.setattr(gc, "freeze", freeze_spy) + + LFXGunicornApp.pre_fork(_FakeServer(), None) + + freeze_spy.assert_called_once() # gc.freeze() ran in the hook + + +@pytest.mark.skipif(sys.platform == "win32", reason="gunicorn is Unix-only") +def test_gunicorn_app_registers_pre_fork_hook(): + from lfx.cli.serve_gunicorn import LFXGunicornApp + + gapp = LFXGunicornApp("lfx.cli.serve_preloaded_app:app", {"workers": 2}) + gapp.load_config() + # The pre_fork hook is wired into gunicorn's config so it runs before each fork. + # (Bound classmethods aren't identity-equal across accesses; compare the function. + # getattr default keeps this a clean assertion failure if the hook isn't wired, + # since gunicorn's default pre_fork is a plain function with no __func__.) + assert getattr(gapp.cfg.pre_fork, "__func__", None) is LFXGunicornApp.pre_fork.__func__ @pytest.mark.skipif(sys.platform == "win32", reason="gunicorn is Unix-only") @@ -99,7 +146,59 @@ def test_gunicorn_app_sets_options(): assert gapp.cfg.worker_class_str == "uvicorn.workers.UvicornWorker" -def test_serve_command_refuses_multiworker_on_windows(monkeypatch, tmp_path): +def _assert_worker_env_cleaned(): + import os as _os + + from lfx.cli.serve_app import ( + _SERVE_FLOW_DIR_ENV, + _SERVE_NO_ENV_FALLBACK_ENV, + _SERVE_STARTUP_PATHS_ENV, + ) + + for key in (_SERVE_FLOW_DIR_ENV, _SERVE_NO_ENV_FALLBACK_ENV, _SERVE_STARTUP_PATHS_ENV): + assert key not in _os.environ + + +def test_windows_multiworker_default_falls_back_to_uvicorn(monkeypatch, tmp_path): + """On Windows with no --max-requests, multi-worker falls back to uvicorn (no isolation).""" + from lfx.cli import commands + + monkeypatch.setattr(commands.sys, "platform", "win32") + + called = {} + + def fake_uvicorn_run(app_str, **kwargs): + called["app"] = app_str + called["kwargs"] = kwargs + + monkeypatch.setattr(commands.uvicorn, "run", fake_uvicorn_run) + + commands._launch_workers( + host="127.0.0.1", + port=8000, + workers=2, + log_level="warning", + flow_dir=tmp_path, + no_env_fallback=False, + script_paths=None, + temp_file_to_cleanup=None, + verbose_print=lambda *_a, **_k: None, + max_requests=None, + limit_concurrency=1, + ) + + # Falls back to the uvicorn factory multi-worker path, not gunicorn. + assert called["app"] == "lfx.cli.serve_app:create_serve_app" + assert called["kwargs"].get("workers") == 2 + assert called["kwargs"].get("factory") is True + # --limit-concurrency is honored on Windows (uvicorn-native); +1 because uvicorn + # counts the active connection, so "1 in-flight" maps to uvicorn's 2. + assert called["kwargs"].get("limit_concurrency") == 2 + _assert_worker_env_cleaned() + + +def test_windows_multiworker_with_max_requests_errors(monkeypatch, tmp_path): + """On Windows, requesting isolation via --max-requests is refused (gunicorn unavailable).""" import typer from lfx.cli import commands @@ -116,19 +215,35 @@ def test_serve_command_refuses_multiworker_on_windows(monkeypatch, tmp_path): script_paths=None, temp_file_to_cleanup=None, verbose_print=lambda *_a, **_k: None, + max_requests=1, + limit_concurrency=None, ) - # The Windows refusal must not leave the worker env vars set. - import os as _os + _assert_worker_env_cleaned() - from lfx.cli.serve_app import ( - _SERVE_FLOW_DIR_ENV, - _SERVE_NO_ENV_FALLBACK_ENV, - _SERVE_STARTUP_PATHS_ENV, - ) - for key in (_SERVE_FLOW_DIR_ENV, _SERVE_NO_ENV_FALLBACK_ENV, _SERVE_STARTUP_PATHS_ENV): - assert key not in _os.environ +def test_lfx_uvicorn_worker_applies_limit_concurrency(monkeypatch): + """LFXUvicornWorker reads LFX_SERVE_LIMIT_CONCURRENCY and sets it on the uvicorn config. + + Tested via the static helper so we don't have to construct a full gunicorn worker. + """ + from types import SimpleNamespace + + from lfx.cli.serve_app import _SERVE_LIMIT_CONCURRENCY_ENV + from lfx.cli.serve_gunicorn import LFXUvicornWorker + + cfg = SimpleNamespace(limit_concurrency=None) + monkeypatch.setenv(_SERVE_LIMIT_CONCURRENCY_ENV, "1") + LFXUvicornWorker._apply_limit_concurrency(cfg) + # +1: uvicorn counts the active connection, so "1 in-flight" maps to uvicorn's 2 + # (uvicorn's limit_concurrency=1 would reject every request). + assert cfg.limit_concurrency == 2 + + # Unset -> left at the uvicorn default (None), i.e. unlimited. + cfg2 = SimpleNamespace(limit_concurrency=None) + monkeypatch.delenv(_SERVE_LIMIT_CONCURRENCY_ENV, raising=False) + LFXUvicornWorker._apply_limit_concurrency(cfg2) + assert cfg2.limit_concurrency is None async def test_guarded_execute_serializes(monkeypatch): @@ -213,3 +328,105 @@ def test_uuid4_is_unique_and_fork_safe(): ids = {str(uuid.uuid4()) for _ in range(10_000)} assert len(ids) == 10_000 + + +def _build_preloaded(monkeypatch, tmp_path): + """Run the exact work the gunicorn preload master does: build + warm the app.""" + from lfx.cli import serve_app + + flow_dir = tmp_path / "flows" + flow_dir.mkdir() + _write_flow(flow_dir, "flow-1") + monkeypatch.setenv(serve_app._SERVE_FLOW_DIR_ENV, str(flow_dir)) + monkeypatch.setenv(serve_app._SERVE_NO_ENV_FALLBACK_ENV, "0") + monkeypatch.setenv(serve_app._SERVE_STARTUP_PATHS_ENV, "") + registry = serve_app.build_registry_from_env() + return serve_app.create_multi_serve_app(registry=registry) + + +def test_preload_does_not_open_the_database(tmp_path, monkeypatch): + """Fork-safety invariant: building/warming the preload app must NOT open a DB session. + + This is the property that keeps lfx serve fork-safe with a pluggable DB: if the + preload master opened a session, an eager-engine DB service would create its + connection pool pre-fork and every forked worker would inherit/corrupt it. A + spy on get_db_service() catches all DB access (session_scope resolves it at + call-time), so this fails loudly if any future preload step touches the DB. + """ + from contextlib import asynccontextmanager + + from lfx.services import deps + + opened = {"count": 0} + + class _SpyDB: + @asynccontextmanager + async def _with_session(self): + opened["count"] += 1 + from lfx.services.session import NoopSession + + async with NoopSession() as session: + yield session + + monkeypatch.setattr(deps, "get_db_service", lambda: _SpyDB()) + + _build_preloaded(monkeypatch, tmp_path) + + assert opened["count"] == 0, "preload opened a DB session — fork-unsafe with a real DB service" + + +def test_preload_starts_no_fork_unsafe_background_threads(tmp_path, monkeypatch): + """Fork-safety invariant: building the preload app must not leave non-benign threads. + + Background threads (telemetry, threaded logging, schedulers) are NOT inherited + by forks, so any started during preload would silently die in workers — and one + holding a lock at fork could deadlock the child. + """ + import threading + + from lfx.cli.serve_gunicorn import LFXGunicornApp + + before = {t.ident for t in threading.enumerate()} + _build_preloaded(monkeypatch, tmp_path) + new_alive = [t for t in threading.enumerate() if t.ident not in before and t.is_alive()] + suspicious = [t for t in new_alive if not LFXGunicornApp._is_benign_thread(t)] + + assert not suspicious, f"preload started fork-unsafe background threads: {[t.name for t in suspicious]}" + + +@pytest.mark.skipif(sys.platform == "win32", reason="gunicorn is Unix-only") +def test_pre_fork_flags_ghost_thread_but_not_benign(): + """The pre_fork guardrail must warn about a non-benign thread alive before fork + and must NOT flag a benign-named one. + """ + import threading + + from lfx.cli.serve_gunicorn import LFXGunicornApp + + warnings: list[str] = [] + + class _Log: + def warning(self, msg, *args): + warnings.append(msg % args if args else msg) + + def debug(self, *_a, **_k): + pass + + class _Server: + log = _Log() + + stop = threading.Event() + ghost = threading.Thread(target=stop.wait, name="EvilGhostThread", daemon=True) + benign = threading.Thread(target=stop.wait, name="OTel-benign", daemon=True) + ghost.start() + benign.start() + try: + LFXGunicornApp.pre_fork(_Server(), None) + finally: + stop.set() + ghost.join(timeout=2) + benign.join(timeout=2) + + blob = "\n".join(warnings) + assert "Ghost threads" in blob and "EvilGhostThread" in blob, warnings + assert "OTel-benign" not in blob # benign-named threads are filtered out diff --git a/src/lfx/tests/unit/services/test_no_env_fallback_credentials.py b/src/lfx/tests/unit/services/test_no_env_fallback_credentials.py index a928387d80..10865f3c46 100644 --- a/src/lfx/tests/unit/services/test_no_env_fallback_credentials.py +++ b/src/lfx/tests/unit/services/test_no_env_fallback_credentials.py @@ -5,7 +5,7 @@ database lookup, fills any still-missing provider keys from ``os.environ`` (so a ``SECRET_KEY`` rotation that silently drops decrypted values doesn't leave the assistant rejecting requests). That loop must honor the request's ``no_env_fallback`` contract — a served flow under ``no_env_fallback`` must stay -isolated from process-wide credentials (concerns A3). +isolated from process-wide credentials. These tests drive the ``user_id``-not-None path (the only one reaching the post-async fallback) with a forced DB miss by monkeypatching ``run_until_complete``