perf: preload the app and fork production granian workers

news/+prod-fork-preload.performance.md

@@ -0,0 +1 @@
+Production backend workers on Linux are now forked from a supervisor that has already imported the app, so framework and app code are shared copy-on-write instead of re-imported per worker. A 4-worker blank app drops from about 720 MB to about 220 MB of proportional set size and starts faster. Set `REFLEX_BACKEND_START_METHOD=spawn` for apps that are not fork-safe.

packages/reflex-base/news/+prod-fork-preload.performance.md

@@ -0,0 +1 @@
+Add the `REFLEX_BACKEND_START_METHOD` environment variable to choose how production backend workers are started (`fork`, `spawn`, or `forkserver`).

packages/reflex-base/src/reflex_base/environment.py

@@ -798,6 +798,13 @@ class EnvironmentVariables:
     # Whether to run Granian in a spawn process. This enables Reflex to pick up on environment variable changes between hot reloads.
     REFLEX_STRICT_HOT_RELOAD: EnvVar[bool] = env_var(False)
 
+    # The multiprocessing start method for production backend workers. Unset means "fork" wherever
+    # the interpreter itself defaults to a fork-based method (Linux), so workers share the app
+    # preloaded by the supervisor. Set to "spawn" for apps that are not fork-safe.
+    REFLEX_BACKEND_START_METHOD: EnvVar[
+        Literal["fork", "spawn", "forkserver"] | None
+    ] = env_var(None)
+
     # The path to the reflex log file. If not set, the log file will be stored in the reflex user directory.
     REFLEX_LOG_FILE: EnvVar[Path | None] = env_var(None)
 

reflex/reflex.py

@@ -515,8 +515,11 @@ def _run_prod(running_mode: constants.RunningMode, port: int, host: str):
 
     _skip_compile()
 
-    # Post a telemetry event.
-    telemetry.send("run-prod")
+    # Post the telemetry event once the workers are running: sending it here
+    # would start the telemetry thread before the server forks its workers.
+    def on_started():
+        """Send the run telemetry from the supervisor."""
+        telemetry.send("run-prod")
 
     # Display custom message when there is a keyboard interrupt.
     atexit.register(processes.atexit_handler)
@@ -528,10 +531,14 @@ def _run_prod(running_mode: constants.RunningMode, port: int, host: str):
     )
     if running_mode.has_backend():
         exec.run_backend_prod(
-            host, port, config.loglevel.subprocess_level(), running_mode.has_frontend()
+            host,
+            port,
+            config.loglevel.subprocess_level(),
+            running_mode.has_frontend(),
+            on_started=on_started,
         )
     else:
-        exec.run_frontend_prod(host, port)
+        exec.run_frontend_prod(host, port, on_started=on_started)
 
 
 def _run(

reflex/utils/exec.py

@@ -14,7 +14,7 @@
 import subprocess
 import sys
 import threading
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 from typing import Any, NamedTuple, TypedDict
 
@@ -398,24 +398,35 @@ def _frontend_prod_app():
     return Starlette(routes=[get_frontend_mount()])
 
 
-def run_frontend_prod(host: str, port: int):
+def run_frontend_prod(
+    host: str, port: int, on_started: Callable[[], Any] | None = None
+):
     """Run the frontend in production mode by serving compiled static files.
 
     Uses the same granian/uvicorn infrastructure as the backend.
 
     Args:
         host: The host to serve on.
         port: The port to serve on.
+        on_started: Called once the server's workers have been started.
     """
     loglevel = get_config().loglevel.subprocess_level()
 
     if should_use_granian():
         run_granian_backend_prod(
-            host, port, loglevel, app_target=f"{__name__}:_frontend_prod_app"
+            host,
+            port,
+            loglevel,
+            app_target=f"{__name__}:_frontend_prod_app",
+            on_started=on_started,
         )
     else:
         run_uvicorn_backend_prod(
-            host, port, loglevel, app_target=f"{__name__}:_frontend_prod_app"
+            host,
+            port,
+            loglevel,
+            app_target=f"{__name__}:_frontend_prod_app",
+            on_started=on_started,
         )
 
 
@@ -869,6 +880,7 @@ def run_backend_prod(
     port: int,
     loglevel: constants.LogLevel = constants.LogLevel.ERROR,
     mount_frontend_compiled_app: bool = False,
+    on_started: Callable[[], Any] | None = None,
 ):
     """Run the backend.
 
@@ -877,13 +889,14 @@ def run_backend_prod(
         port: The app port
         loglevel: The log level.
         mount_frontend_compiled_app: Whether to mount the compiled frontend app with the backend.
+        on_started: Called once the server's workers have been started.
     """
     environment.REFLEX_MOUNT_FRONTEND_COMPILED_APP.set(mount_frontend_compiled_app)
 
     if should_use_granian():
-        run_granian_backend_prod(host, port, loglevel)
+        run_granian_backend_prod(host, port, loglevel, on_started=on_started)
     else:
-        run_uvicorn_backend_prod(host, port, loglevel)
+        run_uvicorn_backend_prod(host, port, loglevel, on_started=on_started)
 
 
 def _get_backend_workers():
@@ -893,7 +906,11 @@ def _get_backend_workers():
 
 
 def run_uvicorn_backend_prod(
-    host: str, port: int, loglevel: LogLevel, app_target: str | None = None
+    host: str,
+    port: int,
+    loglevel: LogLevel,
+    app_target: str | None = None,
+    on_started: Callable[[], Any] | None = None,
 ):
     """Run the backend in production mode using Uvicorn.
 
@@ -902,6 +919,8 @@ def run_uvicorn_backend_prod(
         port: The app port
         loglevel: The log level.
         app_target: The ASGI app target to run. Defaults to the reflex app instance.
+        on_started: Called right before the server process is launched; it
+            runs in a separate interpreter, so nothing here is forked.
     """
     import os
     import shlex
@@ -944,6 +963,9 @@ def run_uvicorn_backend_prod(
         *("--log-level", loglevel.value),
     ]
 
+    if on_started is not None:
+        on_started()
+
     processes.new_process(
         command,
         run=True,
@@ -954,8 +976,59 @@ def run_uvicorn_backend_prod(
     )
 
 
+def _backend_start_method() -> str | None:
+    """Resolve the multiprocessing start method for production backend workers.
+
+    Returns:
+        The start method to force, or None to keep the interpreter default.
+    """
+    if (method := environment.REFLEX_BACKEND_START_METHOD.get()) is not None:
+        return method
+    import multiprocessing
+
+    # Python defaults to fork (<3.14) or forkserver (3.14+) on Linux and to
+    # spawn elsewhere; only where fork is already the platform norm do we rely
+    # on it so workers can share the supervisor's pages.
+    if multiprocessing.get_start_method() in ("fork", "forkserver"):
+        return "fork"
+    return None
+
+
+def _preload_for_fork(app_target: str | None) -> bool:
+    """Import the app in the supervisor so forked workers share its pages.
+
+    Args:
+        app_target: The ASGI app target; None means the reflex app, which is
+            imported here. Any other target lives in an already-loaded module.
+
+    Returns:
+        Whether forking is safe; False when a telemetry thread is still alive.
+    """
+    import gc
+
+    from reflex_base.utils import serializers
+
+    from reflex.utils import prerequisites, telemetry
+
+    if app_target is None:
+        prerequisites.get_app()
+    serializers._prepare_serializers_for_fork()
+    # Forking while the telemetry thread may hold a lock can deadlock the child.
+    if not telemetry._shutdown_executor():
+        return False
+    # Freezing keeps worker GC passes from writing to the preloaded objects'
+    # headers, which would copy-on-write the shared pages private again.
+    gc.collect()
+    gc.freeze()
+    return True
+
+
 def run_granian_backend_prod(
-    host: str, port: int, loglevel: LogLevel, app_target: str | None = None
+    host: str,
+    port: int,
+    loglevel: LogLevel,
+    app_target: str | None = None,
+    on_started: Callable[[], Any] | None = None,
 ):
     """Run the backend in production mode using Granian.
 
@@ -964,14 +1037,38 @@ def run_granian_backend_prod(
         port: The app port
         loglevel: The log level.
         app_target: The ASGI app target to run. Defaults to the reflex app instance.
+        on_started: Called in the supervisor once the workers have been started,
+            so work it does (e.g. a telemetry thread) is not forked into them.
     """
+    import multiprocessing
+
     from granian.constants import Interfaces
     from granian.log import LogLevels
     from granian.server import Server as Granian
 
     logger.debug("Using Granian for backend")
 
-    granian_app = Granian(
+    if (start_method := _backend_start_method()) is not None:
+        if start_method == "fork" and not _preload_for_fork(app_target):
+            logger.debug("A telemetry send is still running; spawning workers.")
+            start_method = "spawn"
+        multiprocessing.set_start_method(start_method, force=True)
+
+    class NotifyingGranian(Granian):  # pyright: ignore[reportGeneralTypeIssues]
+        """Granian server that reports when its workers have been started."""
+
+        def startup(self, *args, **kwargs):
+            """Start the supervisor and its workers, then notify the caller.
+
+            Args:
+                args: Positional arguments for the Granian startup.
+                kwargs: Keyword arguments for the Granian startup.
+            """
+            super().startup(*args, **kwargs)
+            if on_started is not None:
+                on_started()
+
+    granian_app = NotifyingGranian(
         target=app_target or get_app_instance_from_file(),
         factory=True,
         address=host,

reflex/utils/telemetry.py

@@ -503,6 +503,45 @@ def _get_telemetry_executor() -> ThreadPoolExecutor:
     return _executor
 
 
+def _shutdown_executor(timeout: float = 2) -> bool:
+    """Deliver queued telemetry and stop the worker thread.
+
+    Called before forking so no telemetry thread is alive at the fork; the
+    next send lazily starts a new worker. A send that stalls past the timeout
+    (e.g. on a DNS lookup) is abandoned so it cannot hold up startup.
+
+    Args:
+        timeout: Maximum number of seconds to wait for queued telemetry.
+
+    Returns:
+        Whether the worker thread has stopped, i.e. whether forking is safe.
+    """
+    global _executor
+    if (executor := _executor) is None:
+        return True
+    drained = _flush(timeout)
+    if not drained:
+        logger.debug(f"Telemetry did not drain within {timeout}s before forking.")
+    # Stay published while stopping: a racing send() then fails to submit (and
+    # is suppressed) instead of starting a second pool before the fork.
+    executor.shutdown(wait=drained, cancel_futures=True)
+    with _executor_lock:
+        if _executor is executor:
+            _executor = None
+    return not any(thread.is_alive() for thread in executor._threads)
+
+
+def _reset_executor_after_fork() -> None:
+    """Drop the inherited executor and lock; the child owns neither's thread."""
+    global _executor, _executor_lock
+    _executor_lock = threading.Lock()
+    _executor = None
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_reset_executor_after_fork)
+
+
 def _current_registration_context() -> RegistrationContext | None:
     """Return the caller's RegistrationContext, or None if none is attached.
 

tests/units/test_reflex.py

@@ -10,6 +10,7 @@
 import click
 import click.testing
 import pytest
+from reflex_base import constants
 
 from reflex import reflex
 
@@ -169,6 +170,40 @@ def test_compile_app_worker_flushes_telemetry_on_failure(mocker):
     flush.assert_called_once_with()
 
 
+@pytest.mark.parametrize(
+    ("running_mode", "launcher"),
+    [
+        (constants.RunningMode.BACKEND_ONLY, "run_backend_prod"),
+        (constants.RunningMode.FRONTEND_ONLY, "run_frontend_prod"),
+    ],
+)
+def test_run_prod_sends_telemetry_once_the_server_started(
+    mocker, running_mode: constants.RunningMode, launcher: str
+):
+    """The run-prod event waits for the server so no telemetry thread is forked."""
+    from reflex.utils import build, exec, processes, telemetry
+
+    send = mocker.patch.object(telemetry, "send")
+    mocker.patch.object(reflex, "get_config")
+    mocker.patch.object(reflex, "_compile_app")
+    mocker.patch.object(reflex, "_skip_compile")
+    mocker.patch.object(build, "setup_frontend_prod")
+    mocker.patch.object(processes, "atexit_handler")
+    mocker.patch("atexit.register")
+    mocker.patch.object(exec, "notify_app_running")
+    mocker.patch.object(exec, "notify_frontend")
+
+    def serve(*_args, on_started):
+        send.assert_not_called()
+        on_started()
+
+    mocker.patch.object(exec, launcher, side_effect=serve)
+
+    reflex._run_prod(running_mode, 8000, "127.0.0.1")
+
+    send.assert_called_once_with("run-prod")
+
+
 def test_cloud_commands_registered():
     """The hosting CLI commands import, resolve, and dispatch only on demand."""
     probe = """