From 391fa75e8e12092e824310461a17756addd4361b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20H=C3=B8xbro=20Hansen?= <simon.hansen@me.com>
Date: Mon, 26 Aug 2024 20:24:38 +0200
Subject: [PATCH] Improving support for warm and autoreload with num-procs

---
 panel/command/serve.py            | 10 +++++----
 panel/tests/command/test_serve.py | 35 +++++++++++++++++++++++++++++++
 panel/tests/util.py               |  4 ++--
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/panel/command/serve.py b/panel/command/serve.py
index 6dc5f33bf8..67d44659f2 100644
--- a/panel/command/serve.py
+++ b/panel/command/serve.py
@@ -9,6 +9,7 @@
 import logging
 import os
 import pathlib
+import sys
 
 from glob import glob
 from types import ModuleType
@@ -280,11 +281,11 @@ def customize_applications(self, args, applications):
                 applications['/'] = applications[f'/{index}']
         return super().customize_applications(args, applications)
 
-    def warm_applications(self, applications, reuse_sessions, error=True):
+    def warm_applications(self, applications, reuse_sessions, error=True, initialize_session=True):
         from ..io.session import generate_session
         for path, app in applications.items():
             try:
-                session = generate_session(app)
+                session = generate_session(app, initialize=initialize_session)
             except Exception as e:
                 if error:
                     raise e
@@ -366,13 +367,14 @@ def customize_kwargs(self, args, server_kwargs):
         if args.warm or args.autoreload:
             argvs = {f: args.args for f in files}
             applications = build_single_handler_applications(files, argvs)
+            initialize_session = not (args.num_procs and sys.version_info < (3, 12))
             if args.autoreload:
                 with record_modules(list(applications.values())):
                     self.warm_applications(
-                        applications, args.reuse_sessions, error=False
+                        applications, args.reuse_sessions, error=False, initialize_session=initialize_session
                     )
             else:
-                self.warm_applications(applications, args.reuse_sessions)
+                self.warm_applications(applications, args.reuse_sessions, initialize_session=initialize_session)
 
         if args.liveness:
             argvs = {f: args.args for f in files}
diff --git a/panel/tests/command/test_serve.py b/panel/tests/command/test_serve.py
index fbeedfaf36..42d0dc0f20 100644
--- a/panel/tests/command/test_serve.py
+++ b/panel/tests/command/test_serve.py
@@ -1,4 +1,5 @@
 import os
+import re
 import tempfile
 
 import pytest
@@ -103,3 +104,37 @@ def test_serve_markdown():
         r = requests.get(f"http://localhost:{port}/")
         assert r.status_code == 200
         assert '<title>My app</title>' in r.content.decode('utf-8')
+
+
+@linux_only
+@pytest.mark.parametrize("arg", ["--warm", "--autoreload"])
+def test_serve_num_procs(arg):
+    app = "import panel as pn; pn.panel('Hello').servable()"
+    py = tempfile.NamedTemporaryFile(mode='w', suffix='.py')
+    write_file(app, py.file)
+
+    regex = re.compile(r'Starting Bokeh server with process id: (\d+)')
+    with run_panel_serve(["--port", "0", py.name, "--num-procs", "2", arg]) as p:
+        pid1 = wait_for_port(p.stdout, regex=regex)
+        pid2 = wait_for_port(p.stdout, regex=regex)
+        assert pid1 != pid2
+#
+@linux_only
+def test_serve_num_procs_setup():
+    app = "import panel as pn; pn.panel('Hello').servable()"
+    py = tempfile.NamedTemporaryFile(mode='w', suffix='.py')
+    write_file(app, py.file)
+
+    setup_app = """\
+import os
+import time
+time.sleep(1)
+print(f"Setup PID {os.getpid()}", flush=True)"""
+    setup_py = tempfile.NamedTemporaryFile(mode='w', suffix='.py')
+    write_file(setup_app, setup_py.file)
+
+    regex = re.compile(r'Setup PID (\d+)')
+    with run_panel_serve(["--port", "0", py.name, "--num-procs", "2", "--setup", setup_py.name]) as p:
+        pid1 = wait_for_port(p.stdout, regex=regex)
+        pid2 = wait_for_port(p.stdout, regex=regex)
+        assert pid1 != pid2
diff --git a/panel/tests/util.py b/panel/tests/util.py
index a6e2e45a74..27c3116aaf 100644
--- a/panel/tests/util.py
+++ b/panel/tests/util.py
@@ -371,7 +371,7 @@ def readline(self, timeout=None):
         except Empty:
             return None
 
-def wait_for_port(stdout):
+def wait_for_port(stdout, regex=APP_PATTERN):
     nbsr = NBSR(stdout)
     m = None
     output = []
@@ -381,7 +381,7 @@ def wait_for_port(stdout):
             continue
         out = o.decode('utf-8')
         output.append(out)
-        m = APP_PATTERN.search(out)
+        m = regex.search(out)
         if m is not None:
             break
     if m is None: