Merge pull request #116 from anubrag/improve-compile-process

Improve compile process

Author: anubrag

README.md

@@ -35,10 +35,11 @@
 <br>
 
 <a href="https://discord.gg/asTZktHrcH" target="blank">
-<img alt="Discord" src="https://img.shields.io/discord/1111984809821089883?style=social&logo=discord&label=Join%20Discord%20to%20access%20the%20nextpy%20bot%F0%9F%A4%96" width="500px">
+<img alt="Discord" src="https://img.shields.io/discord/1111984809821089883?style=for-the-badge&logo=discord&logoColor=white&label=Live%20Support%20%26%20Coding%20Bots%20on%20Discord&labelColor=%23684DFF&link=https%3A%2F%2Fdiscord.gg%2FasTZktHrcH">
+
 </a>
 
-<hr>
+![-----------------------------------------------------](https://res.cloudinary.com/dzznkbdrb/image/upload/v1694798498/divider_1_rej288.gif)
 
 <h3><i>Streamlit's simplicity (but 4-10x faster) + FastAPI's full power + (Pydantic & SQL Alchemy)'s robustness</i></h3>
 

nextpy/app.py

@@ -9,7 +9,9 @@
 import contextlib
 import copy
 import functools
+import multiprocessing
 import os
+import platform
 from typing import (
     Any,
     AsyncIterator,
@@ -49,11 +51,13 @@
     State,
     StateManager,
     StateUpdate,
+    code_uses_state_contexts,
 )
 from nextpy.base import Base
 from nextpy.build import prerequisites
 from nextpy.build.compiler import compiler
 from nextpy.build.compiler import utils as compiler_utils
+from nextpy.build.compiler.compiler import ExecutorSafeFunctions
 from nextpy.build.config import get_config
 from nextpy.data.model import Model
 from nextpy.frontend.components import connection_modal
@@ -554,6 +558,8 @@ def get_frontend_packages(self, imports: Dict[str, set[ReactImportVar]]):
         Example:
             >>> get_frontend_packages({"react": "16.14.0", "react-dom": "16.14.0"})
         """
+        if getattr(self, '_has_installed_frontend_packages', False):
+            return
         page_imports = {
             i
             for i, tags in imports.items()
@@ -582,6 +588,7 @@ def get_frontend_packages(self, imports: Dict[str, set[ReactImportVar]]):
             _frontend_packages.append(package)
         page_imports.update(_frontend_packages)
         prerequisites.install_frontend_packages(page_imports)
+        self._has_installed_frontend_packages = True
 
     def _app_root(self, app_wrappers: dict[tuple[int, str], Component]) -> Component:
         for component in tuple(app_wrappers.values()):
@@ -633,6 +640,17 @@ def compile_(self):
             TimeElapsedColumn(),
         )
 
+        # try to be somewhat accurate - but still not 100%
+        adhoc_steps_without_executor = 6
+        fixed_pages_within_executor = 7
+        progress.start()
+        task = progress.add_task(
+            "Compiling:",
+            total=len(self.pages)
+            + fixed_pages_within_executor
+            + adhoc_steps_without_executor,
+        )
+
         # Get the env mode.
         config = get_config()
 
@@ -641,7 +659,6 @@ def compile_(self):
 
         # Compile the pages in parallel.
         custom_components = set()
-        # TODO Anecdotally, processes=2 works 10% faster (cpu_count=12)
         all_imports = {}
         app_wrappers: Dict[tuple[int, str], Component] = {
             # Default app wrap component renders {children}
@@ -651,116 +668,141 @@ def compile_(self):
             # If a theme component was provided, wrap the app with it
             app_wrappers[(20, "Theme")] = self.theme
 
-        with progress, concurrent.futures.ThreadPoolExecutor() as thread_pool:
-            fixed_pages = 7
-            task = progress.add_task("Compiling:", total=len(self.pages) + fixed_pages)
+        progress.advance(task)
 
-            def mark_complete(_=None):
-                progress.advance(task)
+        for _route, component in self.pages.items():
+            # Merge the component style with the app style.
+            component.add_style(self.style)
 
-            for _route, component in self.pages.items():
-                # Merge the component style with the app style.
-                component.add_style(self.style)
+            component.apply_theme(self.theme)
 
-                component.apply_theme(self.theme)
+            # Add component.get_imports() to all_imports.
+            all_imports.update(component.get_imports())
 
-                # Add component.get_imports() to all_imports.
-                all_imports.update(component.get_imports())
+            # Add the app wrappers from this component.
+            app_wrappers.update(component.get_app_wrap_components())
 
-                # Add the app wrappers from this component.
-                app_wrappers.update(component.get_app_wrap_components())
+            # Add the custom components from the page to the set.
+            custom_components |= component.get_custom_components()
 
-                # Add the custom components from the page to the set.
-                custom_components |= component.get_custom_components()
+        progress.advance(task)
 
-            # Perform auto-memoization of stateful components.
-            (
-                stateful_components_path,
-                stateful_components_code,
-                page_components,
-            ) = compiler.compile_stateful_components(self.pages.values())
-            compile_results.append((stateful_components_path, stateful_components_code))
 
-            result_futures = []
+        # Perform auto-memoization of stateful components.
+        (
+            stateful_components_path,
+            stateful_components_code,
+            page_components,
+        ) = compiler.compile_stateful_components(self.pages.values())
 
-            def submit_work(fn, *args, **kwargs):
-                """Submit work to the thread pool and add a callback to mark the task as complete.
 
-                The Future will be added to the `result_futures` list.
+        progress.advance(task)
+
+        # Catch "static" apps (that do not define a xt.State subclass) which are trying to access xt.State.
+        if code_uses_state_contexts(stateful_components_code) and self.state is None:
+            raise RuntimeError(
+                "To access xt.State in frontend components, at least one "
+                "subclass of xt.State must be defined in the app."
+            )
+        compile_results.append((stateful_components_path, stateful_components_code))
 
-                Args:
-                    fn: The function to submit.
-                    *args: The args to submit.
-                    **kwargs: The kwargs to submit.
-                """
-                f = thread_pool.submit(fn, *args, **kwargs)
-                f.add_done_callback(mark_complete)
+        app_root = self._app_root(app_wrappers=app_wrappers)
+
+        progress.advance(task)
+
+        # Prepopulate the global ExecutorSafeFunctions class with input data required by the compile functions.
+        # This is required for multiprocessing to work, in presence of non-picklable inputs.
+        for route, component in zip(self.pages, page_components):
+            ExecutorSafeFunctions.COMPILE_PAGE_ARGS_BY_ROUTE[route] = (
+                route,
+                component,
+                self.state,
+            )
+
+        ExecutorSafeFunctions.COMPILE_APP_APP_ROOT = app_root
+        ExecutorSafeFunctions.CUSTOM_COMPONENTS = custom_components
+        ExecutorSafeFunctions.HEAD_COMPONENTS = self.head_components
+        ExecutorSafeFunctions.STYLE = self.style
+        ExecutorSafeFunctions.STATE = self.state
+
+        # Use a forking process pool, if possible.  Much faster, especially for large sites.
+        # Fallback to ThreadPoolExecutor as something that will always work.
+        executor = None
+        if platform.system() in ("Linux", "Darwin"):
+            executor = concurrent.futures.ProcessPoolExecutor(
+                mp_context=multiprocessing.get_context("fork")
+            )
+        else:
+            executor = concurrent.futures.ThreadPoolExecutor()
+
+        with executor:
+            result_futures = []
+
+            def _mark_complete(_=None):
+                progress.advance(task)
+
+            def _submit_work(fn, *args, **kwargs):
+                f = executor.submit(fn, *args, **kwargs)
+                f.add_done_callback(_mark_complete)
                 result_futures.append(f)
 
             # Compile all page components.
-            for route, component in zip(self.pages, page_components):
-                submit_work(
-                    compiler.compile_page,
-                    route,
-                    component,
-                    self.state,
-                )
+            for route in self.pages:
+                _submit_work(ExecutorSafeFunctions.compile_page, route)
+
 
             # Compile the app wrapper.
-            app_root = self._app_root(app_wrappers=app_wrappers)
-            submit_work(compiler.compile_app, app_root)
+            _submit_work(ExecutorSafeFunctions.compile_app)
+
 
             # Compile the custom components.
-            submit_work(compiler.compile_components, custom_components)
+            _submit_work(ExecutorSafeFunctions.compile_custom_components)
 
             # Compile the root stylesheet with base styles.
-            submit_work(compiler.compile_root_stylesheet, self.stylesheets)
+            _submit_work(compiler.compile_root_stylesheet, self.stylesheets)
 
             # Compile the root document.
-            submit_work(compiler.compile_document_root, self.head_components)
+            _submit_work(ExecutorSafeFunctions.compile_document_root)
 
             # Compile the theme.
-            submit_work(compiler.compile_theme, style=self.style)
+            _submit_work(ExecutorSafeFunctions.compile_theme)
 
             # Compile the contexts.
-            submit_work(compiler.compile_contexts, self.state)
+            _submit_work(ExecutorSafeFunctions.compile_contexts)
 
             # Compile the Tailwind config.
             if config.tailwind is not None:
                 config.tailwind["content"] = config.tailwind.get(
                     "content", constants.Tailwind.CONTENT
                 )
-            submit_work(compiler.compile_tailwind, config.tailwind)
-
-            # Get imports from AppWrap components.
-            all_imports.update(app_root.get_imports())
-
-            # Iterate through all the custom components and add their imports to the all_imports.
-            for component in custom_components:
-                all_imports.update(component.get_imports())
+                _submit_work(compiler.compile_tailwind, config.tailwind)
+            else:
+                _submit_work(compiler.remove_tailwind_from_postcss)
 
             # Wait for all compilation tasks to complete.
             for future in concurrent.futures.as_completed(result_futures):
                 compile_results.append(future.result())
 
-            # Empty the .web pages directory.
-            compiler.purge_web_pages_dir()
+        # Get imports from AppWrap components.
+        all_imports.update(app_root.get_imports())
 
-            # Avoid flickering when installing frontend packages
-            progress.stop()
+        # Iterate through all the custom components and add their imports to the all_imports.
+        for component in custom_components:
+            all_imports.update(component.get_imports())
 
-            # Install frontend packages.
-            self.get_frontend_packages(all_imports)
+        progress.advance(task)
 
-            # Write the pages at the end to trigger the NextJS hot reload only once.
-            write_page_futures = []
-            for output_path, code in compile_results:
-                write_page_futures.append(
-                    thread_pool.submit(compiler_utils.write_page, output_path, code)
-                )
-            for future in concurrent.futures.as_completed(write_page_futures):
-                future.result()
+        # Empty the .web pages directory.
+        compiler.purge_web_pages_dir()
+
+        progress.advance(task)
+        progress.stop()
+
+        # Install frontend packages.
+        self.get_frontend_packages(all_imports)
 
+        for output_path, code in compile_results:
+            compiler_utils.write_page(output_path, code)
     @contextlib.asynccontextmanager
     async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
         """Modify the state out of band.

nextpy/backend/state.py

@@ -1327,6 +1327,9 @@ def __init__(self, state_instance):
         """
         super().__init__(state_instance)
         # compile is not relevant to backend logic
+        #TODO: We're currently using this weirdo mechanism for installing initial packages
+        # gets to load_module -> compile -> get_frontend_packages -> install_frontend_packages
+        # We can improve this
         self._self_app = getattr(prerequisites.get_app(), constants.CompileVars.APP)
         self._self_substate_path = state_instance.get_full_name().split(".")
         self._self_actx = None
@@ -2142,3 +2145,14 @@ def _mark_dirty(
         return super()._mark_dirty(
             wrapped=wrapped, instance=instance, args=args, kwargs=kwargs
         )
+    
+def code_uses_state_contexts(javascript_code: str) -> bool:
+    """Check if the rendered Javascript uses state contexts.
+
+    Args:
+        javascript_code: The Javascript code to check.
+
+    Returns:
+        True if the code attempts to access a member of StateContexts.
+    """
+    return bool("useContext(StateContexts" in javascript_code)