refactor: DRY up code between wasm and native kernel (#9591)

Author: mscolnick

marimo/_pyodide/pyodide_session.py

@@ -116,13 +116,19 @@ def __init__(
         user_config: MarimoConfig,
     ) -> None:
         """Initialize kernel and client connection to it."""
+        from marimo._runtime.kernel_lifecycle import make_control_enqueuer
+
         self.app_manager = app
         self.mode = mode
         self.app_metadata = app_metadata
         self._queue_manager = AsyncQueueManager()
         self.session_consumer = on_write
         self.session_view = SessionView()
         self._initial_user_config = user_config
+        self._enqueue_control_request = make_control_enqueuer(
+            self._queue_manager.control_queue,
+            self._queue_manager.set_ui_element_queue,
+        )
 
         self.consumers: list[Callable[[KernelMessage], None]] = [
             lambda msg: self.session_consumer(msg),
@@ -148,12 +154,7 @@ async def start(self) -> None:
         await self.kernel_task.start()
 
     def put_control_request(self, request: commands.CommandMessage) -> None:
-        self._queue_manager.control_queue.put_nowait(request)
-        if isinstance(
-            request,
-            (commands.UpdateUIElementCommand, commands.ModelCommand),
-        ):
-            self._queue_manager.set_ui_element_queue.put_nowait(request)
+        self._enqueue_control_request(request)
 
     def put_completion_request(
         self, request: commands.CodeCompletionCommand
@@ -442,6 +443,7 @@ def _launch_pyodide_kernel(
         KernelArgs,
         asyncio_queue_reader,
         create_kernel,
+        drain_stale,
         listen_messages,
         teardown_kernel,
     )
@@ -485,14 +487,9 @@ def _launch_pyodide_kernel(
 
     async def listen_completion() -> None:
         while True:
-            request = await completion_queue.get()
-            while not completion_queue.empty():
-                # discard stale requests to avoid choking the runtime
-                request = await completion_queue.get()
-            LOGGER.debug("received completion request %s", request)
-            # 5 is arbitrary, but is a good limit:
-            # too high will cause long load times
-            # too low can be not as useful
+            request = drain_stale(
+                completion_queue, latest=await completion_queue.get()
+            )
             kernel.code_completion(request, docstrings_limit=5)
 
     async def listen() -> None:

marimo/_runtime/complete.py

@@ -23,7 +23,6 @@
 from marimo._output.md import _md
 from marimo._runtime import dataflow
 from marimo._runtime.commands import CodeCompletionCommand
-from marimo._session.queue import QueueType
 from marimo._types.ids import RequestId
 from marimo._utils.docs import MarimoConverter
 from marimo._utils.format_signature import format_signature
@@ -360,17 +359,6 @@ def _write_no_completions(stream: Stream, completion_id: RequestId) -> None:
     _write_completion_result(stream, completion_id, 0, [])
 
 
-def _drain_queue(
-    completion_queue: QueueType[CodeCompletionCommand],
-) -> CodeCompletionCommand:
-    """Drain the queue of completion requests, returning the most recent one"""
-
-    request = completion_queue.get()
-    while not completion_queue.empty():
-        request = completion_queue.get()
-    return request
-
-
 def _get_completions_with_script(
     codes: list[str], document: str
 ) -> tuple[jedi.Script, list[jedi.api.classes.Completion]]:
@@ -739,32 +727,3 @@ def complete(
             pass
         else:
             LOGGER.debug("Completion worker released globals lock.")
-
-
-def completion_worker(
-    completion_queue: QueueType[CodeCompletionCommand],
-    graph: dataflow.DirectedGraph,
-    glbls: dict[str, Any],
-    glbls_lock: threading.RLock,
-    stream: Stream,
-) -> None:
-    """Code completion worker.
-
-
-    Args:
-        completion_queue: queue from which requests are pulled.
-        graph: dataflow graph backing the marimo program
-        glbls: dictionary of global variables in interpreter memory
-        glbls_lock: lock protecting globals
-        stream: stream used to communicate completion results
-    """
-
-    while True:
-        request = _drain_queue(completion_queue)
-        complete(
-            request=request,
-            graph=graph,
-            glbls=glbls,
-            glbls_lock=glbls_lock,
-            stream=stream,
-        )

marimo/_runtime/kernel_lifecycle.py

@@ -10,6 +10,7 @@
 import asyncio
 import contextlib
 import dataclasses
+import queue as _queue
 from typing import TYPE_CHECKING, Any, TypeVar
 
 from marimo import _loggers
@@ -83,6 +84,20 @@ def is_edit_mode(self) -> bool:
 # Lets each caller pin listen_messages and its reader to the same queue type
 # (threading vs asyncio).
 _Q = TypeVar("_Q")
+_T = TypeVar("_T")
+
+
+def drain_stale(queue: Any, *, latest: _T) -> _T:
+    """Discard stale items queued behind ``latest`` and return the newest.
+
+    Drains via ``get_nowait()`` until exhausted; ``empty()`` is intentionally
+    avoided because ``multiprocessing.Queue.empty()`` can lie.
+    """
+    while True:
+        try:
+            latest = queue.get_nowait()
+        except (asyncio.QueueEmpty, _queue.Empty):
+            return latest
 
 
 def _build_hooks(
@@ -103,6 +118,21 @@ def _build_hooks(
     return hooks
 
 
+def make_control_enqueuer(
+    control_queue: ControlQueue,
+    set_ui_element_queue: UIElementQueue,
+) -> Callable[[CommandMessage], None]:
+    """Build a callable that routes control requests, mirroring UI-element
+    commands onto the batching queue."""
+
+    def enqueue(req: CommandMessage) -> None:
+        control_queue.put_nowait(req)
+        if isinstance(req, (UpdateUIElementCommand, ModelCommand)):
+            set_ui_element_queue.put_nowait(req)
+
+    return enqueue
+
+
 def create_kernel(
     args: KernelArgs,
 ) -> tuple[Kernel, KernelRuntimeContext]:
@@ -113,11 +143,6 @@ def create_kernel(
         user_config["runtime"]["on_cell_change"] = "autorun"
         user_config["runtime"]["auto_reload"] = "off"
 
-    def _enqueue_control_request(req: CommandMessage) -> None:
-        args.control_queue.put_nowait(req)
-        if isinstance(req, (UpdateUIElementCommand, ModelCommand)):
-            args.set_ui_element_queue.put_nowait(req)
-
     # Deferred to break the runtime.py <-> kernel_lifecycle.py import cycle.
     from marimo._runtime.runtime import Kernel
 
@@ -133,7 +158,10 @@ def _enqueue_control_request(req: CommandMessage) -> None:
         ),
         debugger_override=args.debugger,
         user_config=user_config,
-        enqueue_control_request=_enqueue_control_request,
+        enqueue_control_request=make_control_enqueuer(
+            args.control_queue,
+            args.set_ui_element_queue,
+        ),
         hooks=_build_hooks(args.is_edit_mode, user_config),
     )
     ctx = initialize_kernel_context(

marimo/_runtime/runtime.py

@@ -702,19 +702,16 @@ def start_completion_worker(
         self, completion_queue: QueueType[CodeCompletionCommand]
     ) -> None:
         """Must be called after context is initialized"""
-        from marimo._runtime.complete import completion_worker
+        from marimo._runtime.kernel_lifecycle import drain_stale
 
-        threading.Thread(
-            target=completion_worker,
-            args=(
-                completion_queue,
-                self.graph,
-                self.globals,
-                self._globals_lock,
-                get_context().stream,
-            ),
-            daemon=True,
-        ).start()
+        def _worker() -> None:
+            while True:
+                request = drain_stale(
+                    completion_queue, latest=completion_queue.get()
+                )
+                self.code_completion(request, docstrings_limit=80)
+
+        threading.Thread(target=_worker, daemon=True).start()
         self._completion_worker_started = True
 
     @kernel_tracer.start_as_current_span("code_completion")
@@ -728,7 +725,7 @@ def code_completion(
             self.graph,
             self.globals,
             self._globals_lock,
-            get_context().stream,
+            self.stream,
             docstrings_limit,
         )
 

tests/_runtime/test_kernel_lifecycle.py

@@ -2,21 +2,25 @@
 from __future__ import annotations
 
 import asyncio
+import queue as _queue
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
 from marimo._runtime.commands import (
     ExecuteCellsCommand,
+    ModelCommand,
     StopKernelCommand,
     UpdateUIElementCommand,
 )
 from marimo._runtime.kernel_lifecycle import (
     asyncio_queue_reader,
+    drain_stale,
     listen_messages,
+    make_control_enqueuer,
 )
-from marimo._types.ids import CellId_t, UIElementId
+from marimo._types.ids import CellId_t, UIElementId, WidgetModelId
 
 
 @pytest.fixture
@@ -151,3 +155,76 @@ async def test_listen_messages_merges_ui_updates(
     dispatched = kernel.handle_message.await_args.args[0]
     assert isinstance(dispatched, UpdateUIElementCommand)
     assert dispatched.values == [2]
+
+
+@pytest.mark.parametrize(
+    "queue_factory",
+    [asyncio.Queue, _queue.Queue],
+    ids=["asyncio", "threading"],
+)
+def test_drain_stale_returns_latest_when_queue_empty(
+    queue_factory: Any,
+) -> None:
+    q = queue_factory()
+    latest = _execute("only")
+    assert drain_stale(q, latest=latest) is latest
+
+
+@pytest.mark.parametrize(
+    "queue_factory",
+    [asyncio.Queue, _queue.Queue],
+    ids=["asyncio", "threading"],
+)
+def test_drain_stale_returns_newest_pending(queue_factory: Any) -> None:
+    q = queue_factory()
+    initial = _execute("initial")
+    newer = _execute("newer")
+    newest = _execute("newest")
+    q.put_nowait(newer)
+    q.put_nowait(newest)
+
+    assert drain_stale(q, latest=initial) is newest
+    # Drained: nothing else remains.
+    assert q.empty()
+
+
+def test_make_control_enqueuer_routes_plain_command_to_control_only() -> None:
+    control: asyncio.Queue[Any] = asyncio.Queue()
+    ui: asyncio.Queue[Any] = asyncio.Queue()
+    enqueue = make_control_enqueuer(control, ui)
+
+    cmd = _execute()
+    enqueue(cmd)
+
+    assert control.get_nowait() is cmd
+    assert ui.empty()
+
+
+def test_make_control_enqueuer_mirrors_ui_element_command() -> None:
+    control: asyncio.Queue[Any] = asyncio.Queue()
+    ui: asyncio.Queue[Any] = asyncio.Queue()
+    enqueue = make_control_enqueuer(control, ui)
+
+    cmd = _ui_update("u", 1)
+    enqueue(cmd)
+
+    assert control.get_nowait() is cmd
+    assert ui.get_nowait() is cmd
+
+
+def test_make_control_enqueuer_mirrors_model_command() -> None:
+    from marimo._runtime.commands import ModelUpdateMessage
+
+    control: asyncio.Queue[Any] = asyncio.Queue()
+    ui: asyncio.Queue[Any] = asyncio.Queue()
+    enqueue = make_control_enqueuer(control, ui)
+
+    cmd = ModelCommand(
+        model_id=WidgetModelId("m1"),
+        message=ModelUpdateMessage(state={"x": 1}, buffer_paths=[]),
+        buffers=[],
+    )
+    enqueue(cmd)
+
+    assert control.get_nowait() is cmd
+    assert ui.get_nowait() is cmd