app_module_for_backend: wait for _compile in prod mode (#2760)

* Expose reflex.utils.exec.is_prod_mode

Formalize runtime checking of the app's `--env` parameter.

* app_module_for_backend: wait for _compile in prod mode

Prod mode uses separate worker processes that fork from the main process.

If the app is not fully compiled when the fork occurs, any further changes to
the app (like mounting the _upload endpoint) will not be reflected in the
workers. This is not a performance hit because compile is skipped anyway
for backend processes and hot reload is not in the picture for prod mode.

* remove _is_dev_mode and replace it with calls to is_prod_mode()

_is_dev_mode was a private function in the compiler, but now that utils.exec
exposes is_prod_mode, we should use that throughout for consistency

reflex/app_module_for_backend.py

@@ -4,6 +4,7 @@ Only the app attribute is explicitly exposed.
 from concurrent.futures import ThreadPoolExecutor
 
 from reflex import constants
+from reflex.utils.exec import is_prod_mode
 from reflex.utils.prerequisites import get_app, get_compiled_app
 
 if "app" != constants.CompileVars.APP:
@@ -11,14 +12,20 @@ if "app" != constants.CompileVars.APP:
 
 app_module = get_app(reload=False)
 app = getattr(app_module, constants.CompileVars.APP)
-# Force background compile errors to print eagerly
-ThreadPoolExecutor(max_workers=1).submit(app.compile_).add_done_callback(
+compile_future = ThreadPoolExecutor(max_workers=1).submit(app.compile_)
+compile_future.add_done_callback(
+    # Force background compile errors to print eagerly
     lambda f: f.result()
 )
+# Wait for the compile to finish in prod mode to ensure all optional endpoints are mounted.
+if is_prod_mode():
+    compile_future.result()
 
 # ensure only "app" is exposed.
 del app_module
+del compile_future
 del get_app
 del get_compiled_app
+del is_prod_mode
 del constants
 del ThreadPoolExecutor

reflex/compiler/compiler.py

@@ -18,6 +18,7 @@ from reflex.components.component import (
 from reflex.config import get_config
 from reflex.state import BaseState
 from reflex.style import LIGHT_COLOR_MODE
+from reflex.utils.exec import is_prod_mode
 from reflex.utils.imports import ImportVar
 from reflex.vars import Var
 
@@ -66,10 +67,6 @@ def _compile_theme(theme: dict) -> str:
     return templates.THEME.render(theme=theme)
 
 
-def _is_dev_mode() -> bool:
-    return os.environ.get("REFLEX_ENV_MODE", "dev") == "dev"
-
-
 def _compile_contexts(state: Optional[Type[BaseState]], theme: Component) -> str:
     """Compile the initial state and contexts.
 
@@ -88,12 +85,12 @@ def _compile_contexts(state: Optional[Type[BaseState]], theme: Component) -> str
             initial_state=utils.compile_state(state),
             state_name=state.get_name(),
             client_storage=utils.compile_client_storage(state),
-            is_dev_mode=_is_dev_mode(),
+            is_dev_mode=not is_prod_mode(),
             default_color_mode=appearance,
         )
         if state
         else templates.CONTEXT.render(
-            is_dev_mode=_is_dev_mode(),
+            is_dev_mode=not is_prod_mode(),
             default_color_mode=appearance,
         )
     )
@@ -256,7 +253,7 @@ def _compile_stateful_components(
         if (
             isinstance(component, StatefulComponent)
             and component.references > 1
-            and not _is_dev_mode()
+            and is_prod_mode()
         ):
             # Reset this flag to render the actual component.
             component.rendered_as_shared = False
@@ -484,7 +481,7 @@ def remove_tailwind_from_postcss() -> tuple[str, str]:
 
 def purge_web_pages_dir():
     """Empty out .web/pages directory."""
-    if _is_dev_mode() and os.environ.get("REFLEX_PERSIST_WEB_DIR"):
+    if not is_prod_mode() and os.environ.get("REFLEX_PERSIST_WEB_DIR"):
         # Skip purging the web directory in dev mode if REFLEX_PERSIST_WEB_DIR is set.
         return
 

reflex/constants/__init__.py

@@ -2,6 +2,7 @@
 
 from .base import (
     COOKIES,
+    ENV_MODE_ENV_VAR,
     IS_WINDOWS,
     LOCAL_STORAGE,
     POLLING_MAX_HTTP_BUFFER_SIZE,

reflex/constants/base.py

@@ -183,6 +183,9 @@ LOCAL_STORAGE = "local_storage"
 # If this env var is set to "yes", App.compile will be a no-op
 SKIP_COMPILE_ENV_VAR = "__REFLEX_SKIP_COMPILE"
 
+# This env var stores the execution mode of the app
+ENV_MODE_ENV_VAR = "REFLEX_ENV_MODE"
+
 # Testing variables.
 # Testing os env set by pytest when running a test case.
 PYTEST_CURRENT_TEST = "PYTEST_CURRENT_TEST"

reflex/reflex.py

@@ -151,7 +151,7 @@ def _run(
     console.set_log_level(loglevel)
 
     # Set env mode in the environment
-    os.environ["REFLEX_ENV_MODE"] = env.value
+    os.environ[constants.ENV_MODE_ENV_VAR] = env.value
 
     # Show system info
     exec.output_system_info()

reflex/utils/exec.py

@@ -294,3 +294,16 @@ def is_testing_env() -> bool:
         True if the app is running in under pytest.
     """
     return constants.PYTEST_CURRENT_TEST in os.environ
+
+
+def is_prod_mode() -> bool:
+    """Check if the app is running in production mode.
+
+    Returns:
+        True if the app is running in production mode or False if running in dev mode.
+    """
+    current_mode = os.environ.get(
+        constants.ENV_MODE_ENV_VAR,
+        constants.Env.DEV.value,
+    )
+    return current_mode == constants.Env.PROD.value