Support running user code in a thread

Author: davidbrochart

plugins/akernel_task/fps_akernel_task/akernel_task.py

@@ -8,8 +8,9 @@
 
 
 class AKernelTask(_Kernel):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, execute_in_thread: bool = False, **kwargs):
         super().__init__()
+        self.execute_in_thread = execute_in_thread
 
     async def start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) -> None:
         async with (
@@ -37,6 +38,7 @@ async def start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) ->
                 self._to_stdin_receive_stream,
                 self._from_stdin_send_stream,
                 self._from_iopub_send_stream,
+                execute_in_thread=self.execute_in_thread,
             )
             self.task_group.start_soon(self.kernel.start)
             task_status.started()

plugins/akernel_task/fps_akernel_task/main.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from functools import partial
+
 from fps import Module
 
 from jupyverse_api.kernel import KernelFactory
@@ -9,6 +11,16 @@
 
 
 class AKernelTaskModule(Module):
+    def __init__(self, *args, execute_in_thread: bool = False, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.execute_in_thread = execute_in_thread
+
+    async def prepare(self) -> None:
+        kernels = await self.get(Kernels)
+        kernels.register_kernel_factory("akernel", KernelFactory(partial(AKernelTask, execute_in_thread=self.execute_in_thread)))
+
+
+class AKernelThreadTaskModule(Module):
     async def prepare(self) -> None:
         kernels = await self.get(Kernels)
-        kernels.register_kernel_factory("akernel", KernelFactory(AKernelTask))
+        kernels.register_kernel_factory("akernel-thread", KernelFactory(partial(AKernelTask, execute_in_thread=True)))

plugins/akernel_task/pyproject.toml

@@ -29,5 +29,5 @@ text = "MIT"
 Homepage = "https://github.com/davidbrochart/akernel"
 
 [project.entry-points]
-"fps.modules"       = {akernel_task = "fps_akernel_task.main:AKernelTaskModule"}
-"jupyverse.modules" = {akernel_task = "fps_akernel_task.main:AKernelTaskModule"}
+"fps.modules"       = {akernel_task = "fps_akernel_task.main:AKernelTaskModule", akernelthread_task = "fps_akernel_task.main:AKernelThreadTaskModule"}
+"jupyverse.modules" = {akernel_task = "fps_akernel_task.main:AKernelTaskModule", akernelthread_task = "fps_akernel_task.main:AKernelThreadTaskModule"}

pyproject.toml

@@ -44,8 +44,8 @@ test = [
 
 [project.optional-dependencies]
 subprocess = [
-    "zmq-anyio >=0.3.9,<0.4.0",
-    "typer >=0.4.0",
+    "zmq-anyio >=0.3.13,<0.4.0",
+    "cyclopts >=4.3.0,<5.0.0",
 ]
 
 react = [
@@ -57,14 +57,15 @@ cache = [
 ]
 
 [project.scripts]
-akernel = "akernel.akernel:cli"
+akernel = "akernel.akernel:app"
 
 [tool.hatch.build.targets.wheel]
 ignore-vcs = true
 packages = ["src/akernel"]
 
 [tool.hatch.build.targets.wheel.shared-data]
 "share/jupyter/kernels/akernel/kernel.json" = "share/jupyter/kernels/akernel/kernel.json"
+"share/jupyter/kernels/akernel-thread/kernel.json" = "share/jupyter/kernels/akernel-thread/kernel.json"
 
 [project.urls]
 Homepage = "https://github.com/davidbrochart/akernel"

share/jupyter/kernels/akernel-thread/kernel.json

@@ -0,0 +1,11 @@
+{
+  "argv": [
+    "akernel",
+    "launch",
+    "--execute-in-thread",
+    "-f",
+    "{connection_file}"
+  ],
+  "display_name": "Python 3 (akernel-thread)",
+  "language": "python"
+}

src/akernel/akernel.py

@@ -1,26 +1,30 @@
 from __future__ import annotations
 
 import json
-from typing import Optional, cast
+from typing import Annotated, cast
 
-import typer
 from anyio import create_memory_object_stream, create_task_group, run, sleep_forever
+from cyclopts import App, Parameter
 
 from .connect import connect_channel
 from .kernel import Kernel
 from .kernelspec import write_kernelspec
 
 
-cli = typer.Typer()
+app = App()
 
 
-@cli.command()
+@app.command()
 def install(
-    mode: str = typer.Argument("", help="Mode of the kernel to install."),
-    cache_dir: Optional[str] = typer.Option(
-        None, "-c", help="Path to the cache directory, if mode is 'cache'."
-    ),
-):
+    mode: str = "",
+    cache_dir: str | None = None,
+) -> None:
+    """Install the kernel.
+
+    Args:
+        mode: Mode of the kernel to install.
+        cache_dir: Path to the cache directory, if mode is 'cache'.
+    """
     kernel_name = "akernel"
     if mode:
         modes = mode.split("-")
@@ -31,20 +35,27 @@ def install(
     write_kernelspec(kernel_name, mode, display_name, cache_dir)
 
 
-@cli.command()
+@app.command()
 def launch(
-    mode: str = typer.Argument("", help="Mode of the kernel to launch."),
-    cache_dir: Optional[str] = typer.Option(
-        None, "-c", help="Path to the cache directory, if mode is 'cache'."
-    ),
-    connection_file: str = typer.Option(..., "-f", help="Path to the connection file."),
+    connection_file: Annotated[str, Parameter(alias=["-f"])],
+    mode: str = "",
+    cache_dir: str | None = None,
+    execute_in_thread: bool = False,
 ):
-    akernel = AKernel(mode, cache_dir, connection_file)
+    """Launch the kernel.
+
+    Args:
+        mode: Mode of the kernel to launch.
+        cache_dir: Path to the cache directory, if mode is 'cache'.
+        connection_file: Path to the connection file.
+        execute_in_thread: Whether to run user code in a thread.
+    """
+    akernel = AKernel(mode, cache_dir, connection_file, execute_in_thread)
     run(akernel.start)
 
 
 class AKernel:
-    def __init__(self, mode, cache_dir, connection_file):
+    def __init__(self, mode, cache_dir, connection_file, execute_in_thread):
         self._to_shell_send_stream, self._to_shell_receive_stream = create_memory_object_stream[list[bytes]]()
         self._from_shell_send_stream, self._from_shell_receive_stream = create_memory_object_stream[list[bytes]]()
         self._to_control_send_stream, self._to_control_receive_stream = create_memory_object_stream[list[bytes]]()
@@ -62,6 +73,7 @@ def __init__(self, mode, cache_dir, connection_file):
             self._from_iopub_send_stream,
             mode,
             cache_dir,
+            execute_in_thread,
         )
         with open(connection_file) as f:
             connection_cfg = json.load(f)
@@ -136,4 +148,4 @@ async def from_iopub(self) -> None:
 
 
 if __name__ == "__main__":
-    cli()
+    app()