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]
@@ -55,3 +56,6 @@ ignore = [
   "PLC",
   "PLR",
 ]
+
+[tool.uv.sources]
+marimo = { path = "../marimo", editable = true }

src/marimo_lsp/kernel_manager.py

@@ -2,51 +2,77 @@
 
 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
+
+from marimo_lsp.zeromq.queue_manager import encode_connection_info
 
 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 ConnectionInfo, ZeroMqQueueManager
+
+
+def launch_kernel_subprocess(
+    executable: str,
+    connection_info: ConnectionInfo,
+    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."""
+
+    # 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)
+    assert process.stdin, "Expected stdin"
 
-    # Now we can import marimo from the correct environment
-    from marimo._runtime import runtime  # noqa: PLC0415
+    process.stdin.write(encode_connection_info(connection_info) + "\n")
+    process.stdin.write(json.dumps(kernel_args) + "\n")
+    process.stdin.flush()
+    process.stdin.close()
 
-    runtime.launch_kernel(*args)
+    return process
 
 
 class LspKernelManager(KernelManager):
@@ -56,7 +82,8 @@ def __init__(
         self,
         *,
         executable: str,
-        queue_manager: QueueManager,
+        connection_info: ConnectionInfo,
+        queue_manager: ZeroMqQueueManager,
         app_file_manager: LspAppFileManager,
         config_manager: MarimoConfigManager,
     ) -> None:
@@ -75,52 +102,26 @@ def __init__(
             redirect_console_to_browser=False,
             virtual_files_supported=False,
         )
+        self.kernel_process = None
         self.executable = executable
+        self.connection_info = connection_info
 
     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,
+            connection_info=self.connection_info,
+            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()
+        # Create IOPub connection for receiving kernel messages
+        self._read_conn = self.queue_manager.create_iopub_connection(for_kernel=False)
 
-        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())
+        # Store process handle (compatible with mp.Process interface)
+        self.kernel_task = self.kernel_process

src/marimo_lsp/session_consumer.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
+import json
 from typing import TYPE_CHECKING
 
-from marimo._messaging.ops import MessageOperation, serialize
+from marimo._messaging.msgspec_encoder import asdict
 from marimo._server.model import ConnectionState, SessionConsumer
 from marimo._types.ids import ConsumerId
 
@@ -13,6 +14,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+    from marimo._messaging.ops import MessageOperation
     from marimo._messaging.types import KernelMessage
     from pygls.lsp.server import LanguageServer
 
@@ -41,11 +43,15 @@ def on_start(self) -> Callable[[KernelMessage], None]:
         def handle_message(msg: KernelMessage) -> None:
             """Forward kernel message over LSP."""
             try:
-                op_name, data = msg
-
+                op_name, msg_bytes = msg
                 self.server.protocol.notify(
                     "marimo/operation",
-                    {"notebookUri": self.notebook_uri, "op": op_name, "data": data},
+                    {
+                        "notebookUri": self.notebook_uri,
+                        "op": op_name,
+                        # TODO: Just send bytes directly
+                        "data": json.loads(msg_bytes),
+                    },
                 )
 
                 logger.debug(f"Forwarded {op_name} to {self.notebook_uri}")
@@ -66,7 +72,7 @@ def write_operation(self, op: MessageOperation) -> None:
             {
                 "notebookUri": self.notebook_uri,
                 "op": op.name,
-                "data": serialize(op),
+                "data": asdict(op),
             },
         )
         logger.debug(f"Sent {op.name} operation to {self.notebook_uri}")

src/marimo_lsp/session_manager.py

@@ -9,16 +9,17 @@
 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
 from marimo_lsp.loggers import get_logger
 from marimo_lsp.session_consumer import LspSessionConsumer
+from marimo_lsp.zeromq.queue_manager import ZeroMqQueueManager
 
 if typing.TYPE_CHECKING:
     from marimo._server.file_manager import AppFileManager
-    from marimo._server.sessions import KernelManager
+    from marimo._server.sessions import KernelManager, QueueManager
     from pygls.lsp.server import LanguageServer
 
 
@@ -66,22 +67,24 @@ 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)
+
+        queue_manager, connection_info = ZeroMqQueueManager.create_host()
 
         kernel_manager = LspKernelManager(
             # TODO: Get executable
             executable=sys.executable,
             queue_manager=queue_manager,
             app_file_manager=app_file_manager,
             config_manager=config_manager,
+            connection_info=connection_info,
         )
 
         logger.info(f"Creating new session for {notebook_uri}")
 
         session = Session(
             initialization_id=str(uuid4()),
             session_consumer=LspSessionConsumer(server, notebook_uri),
-            queue_manager=queue_manager,
+            queue_manager=typing.cast("QueueManager", queue_manager),
             kernel_manager=typing.cast("KernelManager", kernel_manager),
             app_file_manager=typing.cast("AppFileManager", app_file_manager),
             config_manager=config_manager,

src/marimo_lsp/zeromq/__init__.py

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

src/marimo_lsp/zeromq/kernel_server.py

@@ -0,0 +1,82 @@
+"""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, decode_connection_info
+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 = decode_connection_info(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)
+
+    queue_manager = ZeroMqQueueManager.from_connection_info(connection_info)
+
+    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()