diff --git a/reflex/config.py b/reflex/config.py index 953a92b64..88230cefe 100644 --- a/reflex/config.py +++ b/reflex/config.py @@ -454,6 +454,14 @@ class PathExistsFlag: ExistingPath = Annotated[Path, PathExistsFlag] +class PerformanceMode(enum.Enum): + """Performance mode for the app.""" + + WARN = "warn" + RAISE = "raise" + OFF = "off" + + class EnvironmentVariables: """Environment variables class to instantiate environment variables.""" @@ -550,6 +558,12 @@ class EnvironmentVariables: # Whether to check for outdated package versions. REFLEX_CHECK_LATEST_VERSION: EnvVar[bool] = env_var(True) + # In which performance mode to run the app. + REFLEX_PERF_MODE: EnvVar[Optional[PerformanceMode]] = env_var(PerformanceMode.WARN) + + # The maximum size of the reflex state in kilobytes. + REFLEX_STATE_SIZE_LIMIT: EnvVar[int] = env_var(1000) + environment = EnvironmentVariables() diff --git a/reflex/state.py b/reflex/state.py index 9ff6f0ea8..442fa57b2 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -43,7 +43,7 @@ from sqlalchemy.orm import DeclarativeBase from typing_extensions import Self from reflex import event -from reflex.config import get_config +from reflex.config import PerformanceMode, get_config from reflex.istate.data import RouterData from reflex.istate.storage import ClientStorageBase from reflex.model import Model @@ -90,6 +90,7 @@ from reflex.utils.exceptions import ( ReflexRuntimeError, SetUndefinedStateVarError, StateSchemaMismatchError, + StateTooLargeError, ) from reflex.utils.exec import is_testing_env from reflex.utils.serializers import serializer @@ -110,10 +111,11 @@ Delta = Dict[str, Any] var = computed_var -# If the state is this large, it's considered a performance issue. -TOO_LARGE_SERIALIZED_STATE = 100 * 1024 # 100kb -# Only warn about each state class size once. -_WARNED_ABOUT_STATE_SIZE: Set[str] = set() +if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF: + # If the state is this large, it's considered a performance issue. + TOO_LARGE_SERIALIZED_STATE = environment.REFLEX_STATE_SIZE_LIMIT.get() * 1024 + # Only warn about each state class size once. + _WARNED_ABOUT_STATE_SIZE: Set[str] = set() # Errors caught during pickling of state HANDLED_PICKLE_ERRORS = ( @@ -2097,7 +2099,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow): state["__dict__"].pop(inherited_var_name, None) return state - def _warn_if_too_large( + def _check_state_size( self, pickle_state_size: int, ): @@ -2105,6 +2107,9 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow): Args: pickle_state_size: The size of the pickled state. + + Raises: + StateTooLargeError: If the state is too large. """ state_full_name = self.get_full_name() if ( @@ -2112,10 +2117,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow): and pickle_state_size > TOO_LARGE_SERIALIZED_STATE and self.substates ): - console.warn( + msg = ( f"State {state_full_name} serializes to {pickle_state_size} bytes " - "which may present performance issues. Consider reducing the size of this state." + + "which may present performance issues. Consider reducing the size of this state." ) + if environment.REFLEX_PERF_MODE.get() == PerformanceMode.WARN: + console.warn(msg) + elif environment.REFLEX_PERF_MODE.get() == PerformanceMode.RAISE: + raise StateTooLargeError(msg) _WARNED_ABOUT_STATE_SIZE.add(state_full_name) @classmethod @@ -2157,7 +2166,8 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow): """ try: pickle_state = pickle.dumps((self._to_schema(), self)) - self._warn_if_too_large(len(pickle_state)) + if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF: + self._check_state_size(len(pickle_state)) return pickle_state except HANDLED_PICKLE_ERRORS as og_pickle_error: error = ( diff --git a/reflex/utils/exceptions.py b/reflex/utils/exceptions.py index 7611e5a4d..714dc912c 100644 --- a/reflex/utils/exceptions.py +++ b/reflex/utils/exceptions.py @@ -151,6 +151,10 @@ class InvalidPropValueError(ReflexError): """Raised when a prop value is invalid.""" +class StateTooLargeError(ReflexError): + """Raised when the state is too large to be serialized.""" + + class SystemPackageMissingError(ReflexError): """Raised when a system package is missing."""