refactor: extract shared kernel lifecycle for subprocess and pyodide#9541
Conversation
`launch_kernel` and `_launch_pyodide_kernel` duplicated hook setup, kernel construction, context init, the UI-merge control loop, and (in the subprocess case) teardown. Pulls those into `marimo/_runtime/kernel_lifecycle.py` shared by both launchers; environment-specific helpers (stream creation, signal handlers, subprocess bootstrap, profiler) stay at the call site. Side effects: - Pyodide now runs `teardown_kernel` on RestartableTask cancellation. - `listen_messages` uniformly catches `handle_message` exceptions across both UI-merge and non-UI branches (was asymmetric in subprocess, absent in pyodide).
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
There was a problem hiding this comment.
No issues found across 7 files
Architecture diagram
sequenceDiagram
participant Subproc as Subprocess Launcher
participant Pyodide as Pyodide Session
participant Lifecycle as kernel_lifecycle module
participant Kernel as Kernel Runtime
participant Stream as Stream/Pipe
participant Queue as Control/UI Queues
participant Context as Kernel Context
Note over Subproc,Pyodide: Entry points call shared lifecycle
Subproc->>Lifecycle: create_kernel(stream, configs, ...)
Pyodide->>Lifecycle: create_kernel(stream, configs, ...)
Lifecycle->>Lifecycle: _build_hooks(is_edit_mode, user_config)
Lifecycle->>Kernel: Kernel(cell_configs, hooks, enqueue_control_request, ...)
Lifecycle->>Context: initialize_kernel_context(kernel, stream, mode, ...)
Lifecycle-->>Subproc: return kernel, ctx
Lifecycle-->>Pyodide: return kernel, ctx
Note over Subproc: Environment-specific setup
Subproc->>Subproc: _bootstrap_subprocess(parent_pid, log_level)
Subproc->>Subproc: _create_streams(socket_addr, stream_queue, ...)
Subproc->>Subproc: _install_subprocess_handlers(kernel, ctx)
Note over Pyodide: Environment-specific setup
Pyodide->>Pyodide: patch_pyodide_networking()
Pyodide->>Pyodide: patch_recursion_limit(1000)
alt is_edit_mode
Pyodide->>Pyodide: signal.signal(SIGINT, interrupt_handler(ctx))
end
Note over Subproc,Pyodide: Unified control loop via listen()
Subproc->>Lifecycle: listen_messages(kernel, control_queue, ui_queue, threaded_queue_reader)
Pyodide->>Lifecycle: listen_messages(kernel, control_queue, ui_queue, asyncio_queue_reader)
loop Until StopKernelCommand
Lifecycle->>Lifecycle: get_request(control_queue)
Lifecycle->>Kernel: kernel.handle_message(request)
alt Exception in handle_message
Lifecycle->>Lifecycle: LOGGER.exception(...)
end
alt UpdateUIElementCommand or ModelCommand
Lifecycle->>Queue: ui_request_mgr.process_request(request)
Lifecycle->>Lifecycle: merged requests
end
end
Note over Subproc,Pyodide: Teardown on stop or cancellation
alt Subprocess: Normal stop
Lifecycle->>Lifecycle: teardown_kernel(kernel, ctx)
else Pyodide: Task cancellation
Pyodide->>Lifecycle: teardown_kernel(kernel, ctx)
end
Lifecycle->>Lifecycle: ctx.virtual_file_registry.shutdown()
Lifecycle->>Lifecycle: ctx.app_kernel_runner_registry.shutdown()
Lifecycle->>Lifecycle: teardown_context()
Lifecycle->>Kernel: kernel.teardown()
There was a problem hiding this comment.
Pull request overview
This PR refactors kernel startup/control-loop logic by extracting shared, environment-agnostic lifecycle helpers into marimo/_runtime/kernel_lifecycle.py, reducing duplication between the subprocess and Pyodide launchers while standardizing UI-merge behavior and exception handling in the message loop.
Changes:
- Extract shared kernel construction, hook setup, context initialization, UI-merge control loop, and teardown into
marimo/_runtime/kernel_lifecycle.py. - Update subprocess launcher (
marimo/_runtime/runtime.py) and Pyodide launcher (marimo/_pyodide/pyodide_session.py) to use the shared lifecycle helpers. - Add targeted tests for
listen_messagesbehavior (stop semantics, UI-merge, and exception swallowing) and for Pyodide teardown-on-cancel behavior.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/_runtime/test_runtime.py | Updates expected import path for initialize_kernel_context after refactor. |
| tests/_runtime/test_kernel_lifecycle.py | Adds unit tests for the new shared listen_messages loop and queue readers. |
| tests/_pyodide/test_pyodide_session.py | Adds regression test ensuring Pyodide triggers teardown on task cancellation. |
| marimo/_runtime/runtime.py | Refactors subprocess launcher to delegate kernel creation/control-loop/teardown to shared lifecycle module. |
| marimo/_runtime/kernel_lifecycle.py | Introduces shared kernel lifecycle primitives (create_kernel, listen_messages, teardown helpers). |
| marimo/_pyodide/pyodide_session.py | Refactors Pyodide kernel launcher to use shared lifecycle and adds teardown in finally. |
| marimo/_messaging/streams.py | Minor clarifications and typing improvement (QueuePipe implements PipeProtocol). |
| start_task = asyncio.create_task(kernel_task.start()) | ||
| # Yield enough times for: outer task → RestartableTask.start → inner | ||
| # task creation → listen() → asyncio.gather → child tasks suspended. | ||
| for _ in range(5): | ||
| await asyncio.sleep(0) | ||
| kernel_task.stop() |
|
🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.23.7-dev20 |
kirangadhave
left a comment
There was a problem hiding this comment.
changes for launch_kernel are in separate PR?
launch_kerneland_launch_pyodide_kernelduplicated hook setup, kernelconstruction, context init, the UI-merge control loop, and (in the
subprocess case) teardown. Pulls those into
marimo/_runtime/kernel_lifecycle.pyshared by both launchers; environment-specific helpers (stream creation,
signal handlers, subprocess bootstrap, profiler) stay at the call site.
Side effects:
teardown_kernelon RestartableTask cancellation.listen_messagesuniformly catcheshandle_messageexceptions acrossboth UI-merge and non-UI branches (was asymmetric in subprocess, absent
in pyodide).