Add pre-fork hook

This commit is contained in:
Jordan Frazier
2026-06-04 13:55:21 -04:00
parent 35032aac18
commit eb469eb98a
9 changed files with 545 additions and 86 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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