Implement ZeroMQ-based kernel management

Author: manzt

pyproject.toml

@@ -10,6 +10,7 @@ dependencies = [
   "lsprotocol>=2025.0.0rc1",
   "marimo>=0.15.0",
   "pygls>=2.0.0a4",
+  "pyzmq>=27.0.2",
 ]
 
 [project.scripts]

src/marimo_lsp/kernel_manager.py

@@ -2,51 +2,78 @@
 
 from __future__ import annotations
 
+import json
+import subprocess
 import typing
 
 from marimo._config.settings import GLOBAL_SETTINGS
-from marimo._messaging.types import KernelMessage
 from marimo._runtime.requests import AppMetadata
 from marimo._server.model import SessionMode
 from marimo._server.sessions import KernelManager
-from marimo._utils.typed_connection import TypedConnection
 
 if typing.TYPE_CHECKING:
     from marimo._config.manager import MarimoConfigManager
-    from marimo._server.sessions import QueueManager
 
     from marimo_lsp.app_file_manager import LspAppFileManager
-
-
-def launch_kernel(*args) -> None:  # noqa: ANN002
-    """Launch the marimo kernel with the correct Python environment.
-
-    Runs inside a `multiprocessing.Process` spawned with `ctx.set_executable()`.
-
-    However, multiprocessing reconstructs the parent's `sys.path`, overriding the
-    venv's paths. We fix this by querying sys.executable (which IS correctly set)
-    for its natural `sys.path` and replacing ours before importing marimo.
-    """
-    import json  # noqa: PLC0415
-    import subprocess  # noqa: PLC0415
-    import sys  # noqa: PLC0415
-
-    # Get the natural sys.path from the venv's Python interpreter
-    # sys.executable is correctly set to the venv's Python thanks to set_executable()
-    result = subprocess.run(  # noqa: S603
-        [sys.executable, "-c", "import sys, json; print(json.dumps(sys.path))"],
-        capture_output=True,
+    from marimo_lsp.zeromq.queue_manager import ZeroMQQueueManager
+
+
+def launch_kernel_subprocess(
+    executable: str,
+    queue_manager: "ZeroMQQueueManager",
+    configs: dict,
+    app_metadata: AppMetadata,
+    config_manager: "MarimoConfigManager",
+    virtual_files_supported: bool,
+    redirect_console_to_browser: bool,
+    profile_path: str | None,
+) -> subprocess.Popen:
+    """Launch kernel as a subprocess with ZeroMQ IPC."""
+    from marimo_lsp.zeromq.queue_manager import ZeroMQQueueManager
+
+    # Get connection info for kernel
+    connection_info = queue_manager.get_connection_info()
+    
+    # Prepare kernel arguments
+    kernel_args = {
+        "configs": configs,
+        "app_metadata": {
+            "query_params": app_metadata.query_params,
+            "filename": app_metadata.filename,
+            "cli_args": app_metadata.cli_args,
+            "argv": app_metadata.argv,
+            "app_config": app_metadata.app_config,
+        },
+        "user_config": config_manager.get_config(hide_secrets=False),
+        "virtual_files_supported": virtual_files_supported,
+        "redirect_console_to_browser": redirect_console_to_browser,
+        "profile_path": profile_path,
+        "log_level": GLOBAL_SETTINGS.LOG_LEVEL,
+    }
+    
+    # Launch kernel subprocess
+    kernel_cmd = [
+        executable,
+        "-m",
+        "marimo_lsp.zeromq.kernel_server",
+    ]
+    
+    process = subprocess.Popen(
+        kernel_cmd,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
         text=True,
-        check=True,
     )
-
-    # Replace the inherited (wrong) sys.path with the venv's natural paths
-    sys.path = json.loads(result.stdout)
-
-    # Now we can import marimo from the correct environment
-    from marimo._runtime import runtime  # noqa: PLC0415
-
-    runtime.launch_kernel(*args)
+    
+    # Send connection info and kernel args via stdin
+    if process.stdin:
+        process.stdin.write(json.dumps(connection_info) + "\n")
+        process.stdin.write(json.dumps(kernel_args) + "\n")
+        process.stdin.flush()
+        process.stdin.close()
+    
+    return process
 
 
 class LspKernelManager(KernelManager):
@@ -56,7 +83,7 @@ def __init__(
         self,
         *,
         executable: str,
-        queue_manager: QueueManager,
+        queue_manager: typing.Any,  # Can be QueueManager or ZeroMQQueueManager
         app_file_manager: LspAppFileManager,
         config_manager: MarimoConfigManager,
     ) -> None:
@@ -76,51 +103,24 @@ def __init__(
             virtual_files_supported=False,
         )
         self.executable = executable
+        self.kernel_process = None
 
     def start_kernel(self) -> None:
-        """Start an instance of the marimo kernel."""
-        import multiprocessing as mp  # noqa: PLC0415
-        from multiprocessing import connection  # noqa: PLC0415
-
-        # We use a process in edit mode so that we can interrupt the app
-        # with a SIGINT; we don't mind the additional memory consumption,
-        # since there's only one client sess
-        is_edit_mode = self.mode == SessionMode.EDIT
-
-        # Need to use a socket for windows compatibility
-        listener = connection.Listener(family="AF_INET")
-
-        ctx = mp.get_context("spawn")
-        ctx.set_executable(self.executable)
-
-        kernel_task = ctx.Process(
-            target=launch_kernel,
-            args=(
-                self.queue_manager.control_queue,
-                self.queue_manager.set_ui_element_queue,
-                self.queue_manager.completion_queue,
-                self.queue_manager.input_queue,
-                # stream queue unused
-                None,
-                listener.address,
-                is_edit_mode,
-                self.configs,
-                self.app_metadata,
-                self.config_manager.get_config(hide_secrets=False),
-                self._virtual_files_supported,
-                self.redirect_console_to_browser,
-                self.queue_manager.win32_interrupt_queue,
-                self.profile_path,
-                GLOBAL_SETTINGS.LOG_LEVEL,
-            ),
-            # The process can't be a daemon, because daemonic processes
-            # can't create children
-            # https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Process.daemon  # noqa: E501
-            daemon=False,
+        """Start an instance of the marimo kernel using ZeroMQ IPC."""
+        # Launch kernel subprocess with ZeroMQ
+        self.kernel_process = launch_kernel_subprocess(
+            executable=self.executable,
+            queue_manager=self.queue_manager,
+            configs=self.configs,
+            app_metadata=self.app_metadata,
+            config_manager=self.config_manager,
+            virtual_files_supported=self._virtual_files_supported,
+            redirect_console_to_browser=self.redirect_console_to_browser,
+            profile_path=getattr(self, "profile_path", None),
         )
-
-        kernel_task.start()
-
-        self.kernel_task = kernel_task
-        # First thing kernel does is connect to the socket, so it's safe to call accept
-        self._read_conn = TypedConnection[KernelMessage].of(listener.accept())
+        
+        # Create IOPub connection for receiving kernel messages
+        self._read_conn = self.queue_manager.create_iopub_connection(for_kernel=False)
+        
+        # Store process handle (compatible with mp.Process interface)
+        self.kernel_task = self.kernel_process

src/marimo_lsp/session_manager.py

@@ -9,7 +9,7 @@
 from marimo._config.manager import (
     get_default_config_manager,
 )
-from marimo._server.sessions import KernelManager, QueueManager, Session
+from marimo._server.sessions import Session
 
 from marimo_lsp.app_file_manager import LspAppFileManager
 from marimo_lsp.kernel_manager import LspKernelManager
@@ -66,7 +66,10 @@ def create_session(self, *, server: LanguageServer, notebook_uri: str) -> Sessio
 
         app_file_manager = LspAppFileManager(server=server, notebook_uri=notebook_uri)
         config_manager = get_default_config_manager(current_path=app_file_manager.path)
-        queue_manager = QueueManager(use_multiprocessing=True)
+        
+        # Always use ZeroMQQueueManager for cross-Python-version support
+        from marimo_lsp.zeromq.queue_manager import ZeroMQQueueManager
+        queue_manager = ZeroMQQueueManager(is_kernel_side=False)
 
         kernel_manager = LspKernelManager(
             # TODO: Get executable

src/marimo_lsp/zeromq/__init__.py

@@ -0,0 +1,2 @@
+"""ZeroMQ-based IPC implementation for marimo-lsp."""
+

src/marimo_lsp/zeromq/kernel_server.py

@@ -0,0 +1,86 @@
+"""Standalone kernel server entry point for ZeroMQ-based IPC."""
+
+from __future__ import annotations
+
+import json
+import sys
+import threading
+from multiprocessing import connection
+
+from marimo._config.settings import GLOBAL_SETTINGS
+from marimo._runtime import runtime
+from marimo._runtime.requests import AppMetadata
+
+from marimo_lsp.zeromq.queue_manager import ZeroMQQueueManager
+from marimo_lsp.zeromq.zmq_adapter import create_zmq_to_pipe_adapter
+
+
+def main() -> None:
+    """Launch a marimo kernel using ZeroMQ for IPC.
+
+    This function is the entry point for the kernel subprocess.
+    It reads connection information from stdin and sets up ZeroMQ
+    queues that proxy to marimo's internal kernel.
+    """
+    # Read connection info from stdin
+    connection_info_str = sys.stdin.readline().strip()
+    if not connection_info_str:
+        print("No connection info received", file=sys.stderr)
+        sys.exit(1)
+
+    try:
+        connection_info = json.loads(connection_info_str)
+    except json.JSONDecodeError as e:
+        print(f"Failed to parse connection info: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    # Read additional kernel arguments from stdin
+    kernel_args_str = sys.stdin.readline().strip()
+    if not kernel_args_str:
+        print("No kernel arguments received", file=sys.stderr)
+        sys.exit(1)
+
+    try:
+        kernel_args = json.loads(kernel_args_str)
+    except json.JSONDecodeError as e:
+        print(f"Failed to parse kernel arguments: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    # Create ZeroMQ queue manager for kernel side
+    queue_manager = ZeroMQQueueManager(is_kernel_side=True)
+    queue_manager.connect_kernel(connection_info)
+
+    # Create a pipe adapter that bridges ZeroMQ IOPub to TypedConnection
+    listener = connection.Listener(family="AF_INET")
+
+    adapter_thread = threading.Thread(
+        target=create_zmq_to_pipe_adapter, args=(queue_manager, listener), daemon=True
+    )
+    adapter_thread.start()
+
+    runtime.launch_kernel(
+        control_queue=queue_manager.control_queue,
+        set_ui_element_queue=queue_manager.set_ui_element_queue,
+        completion_queue=queue_manager.completion_queue,
+        input_queue=queue_manager.input_queue,
+        stream_queue=None,  # Not used with ZeroMQ
+        socket_addr=listener.address,
+        is_edit_mode=True,
+        profile_path=kernel_args.get("profile_path"),
+        log_level=kernel_args.get("log_level", GLOBAL_SETTINGS.LOG_LEVEL),
+        interrupt_queue=queue_manager.win32_interrupt_queue,
+        configs=kernel_args["configs"],
+        app_metadata=AppMetadata(**kernel_args["app_metadata"]),
+        user_config=kernel_args["user_config"],
+        virtual_files_supported=kernel_args.get("virtual_files_supported", False),
+        redirect_console_to_browser=kernel_args.get(
+            "redirect_console_to_browser", False
+        ),
+    )
+
+    queue_manager.close_queues()
+
+
+if __name__ == "__main__":
+    main()
+