Only postpone setup script if --num-procs is greater than 1 (#7289)

Co-authored-by: Philipp Rudiger <[email protected]>

panel/command/serve.py

@@ -249,7 +249,7 @@ class Serve(_BkServe):
         ('--setup', dict(
             action  = 'store',
             type    = str,
-            help    = "Path to a setup script to run before server starts.",
+            help    = "Path to a setup script to run before server starts. If --num-procs is enabled it will be run in each process after the server has started.",
             default = None
         )),
         ('--liveness', dict(
@@ -370,6 +370,21 @@ def customize_kwargs(self, args, server_kwargs):
             module.__dict__['__file__'] = fullpath(args.setup)
             state._setup_module = module
 
+            def setup_file():
+                setup_path = state._setup_module.__dict__['__file__']
+                with open(setup_path) as f:
+                    setup_source = f.read()
+                nodes = ast.parse(setup_source, os.fspath(setup_path))
+                code = compile(nodes, filename=setup_path, mode='exec', dont_inherit=True)
+                exec(code, state._setup_module.__dict__)
+
+            if args.num_procs > 1:
+                # We will run the setup_file for each process
+                state._setup_file_callback = setup_file
+            else:
+                state._setup_file_callback = None
+                setup_file()
+
         if args.warm or config.autoreload:
             argvs = {f: args.args for f in files}
             applications = build_single_handler_applications(files, argvs)

panel/io/server.py

@@ -3,7 +3,6 @@
 """
 from __future__ import annotations
 
-import ast
 import asyncio
 import datetime as dt
 import importlib
@@ -295,20 +294,12 @@ def __init__(self, *args, **kwargs):
         if state._admin_context:
             state._admin_context._loop = self._loop
 
-    def setup_file(self):
-        setup_path = state._setup_module.__dict__['__file__']
-        with open(setup_path) as f:
-            setup_source = f.read()
-        nodes = ast.parse(setup_source, os.fspath(setup_path))
-        code = compile(nodes, filename=setup_path, mode='exec', dont_inherit=True)
-        exec(code, state._setup_module.__dict__)
-
     def start(self) -> None:
         super().start()
         if state._admin_context:
             self._loop.add_callback(state._admin_context.run_load_hook)
-        if state._setup_module:
-            self._loop.add_callback(self.setup_file)
+        if state._setup_module and state._setup_file_callback:
+            self._loop.add_callback(state._setup_file_callback)
         if config.autoreload:
             from .reload import setup_autoreload_watcher
             self._autoreload_stop_event = stop_event = asyncio.Event()

panel/tests/command/test_serve.py

@@ -134,3 +134,18 @@ def test_serve_num_procs_setup(tmp_path):
     with run_panel_serve(["--port", "0", py, "--num-procs", 2, "--setup", setup_py], cwd=tmp_path) as p:
         pid1, pid2 = wait_for_regex(p.stdout, regex=regex, count=2)
         assert pid1 != pid2
+
+
+def test_serve_setup(tmp_path):
+    app = "import panel as pn; pn.panel('Hello').servable()"
+    py = tmp_path / "app.py"
+    py.write_text(app)
+
+    setup_app = 'print(f"Setup running before", flush=True)'
+    setup_py = tmp_path / "setup.py"
+    setup_py.write_text(setup_app)
+
+    regex = re.compile('(Setup running before)')
+    with run_panel_serve(["--port", "0", py, "--setup", setup_py], cwd=tmp_path) as p:
+        _, output = wait_for_regex(p.stdout, regex=regex, return_output=True)
+        assert output[0].strip() == "Setup running before"

panel/tests/util.py

@@ -394,7 +394,7 @@ def readline(self, timeout=None):
         except Empty:
             return None
 
-def wait_for_regex(stdout, regex, count=1):
+def wait_for_regex(stdout, regex, count=1, return_output=False):
     nbsr = NBSR(stdout)
     m = None
     output, found = [], []
@@ -415,7 +415,7 @@ def wait_for_regex(stdout, regex, count=1):
             "No matching log line in process output, following output "
             f"was captured:\n\n   {output}"
         )
-    return found
+    return (found, output) if return_output else found
 
 def wait_for_port(stdout):
     return int(wait_for_regex(stdout, APP_PATTERN)[0])