Support separate kernel environment (#534)

* Support separate kernel environment

* Fix types

* Again

Author: davidbrochart

jupyverse_api/jupyverse_api/kernels/__init__.py

@@ -251,3 +251,4 @@ class KernelsConfig(Config):
         default=None,
     )
     require_yjs: bool = False
+    kernelenv_path: str = ""

plugins/kernel_subprocess/fps_kernel_subprocess/kernel_subprocess.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import os
 import signal
 import subprocess
 import sys
@@ -9,8 +10,9 @@
 from typing import cast
 
 import anyio
-from anyio import TASK_STATUS_IGNORED, create_task_group, open_process
+from anyio import TASK_STATUS_IGNORED, create_task_group, open_file, open_process, run_process
 from anyio.abc import TaskStatus
+from anyio.streams.text import TextReceiveStream
 
 from jupyverse_api.kernel import Kernel
 
@@ -32,6 +34,7 @@ class KernelSubprocess(Kernel):
     kernel_cwd: str | None
     capture_output: bool
     connection_cfg: cfg_t | None = None
+    kernelenv_path: str = ""
 
     def __post_init__(self):
         super().__init__()
@@ -45,6 +48,8 @@ def __post_init__(self):
                 raise RuntimeError("No connection_cfg")
         self.key = cast(str, self.connection_cfg["key"])
         self.wait_for_ready = True
+        self._process = None
+        self._pid = None
 
     async def start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) -> None:
         async with (
@@ -64,23 +69,48 @@ async def start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) ->
             self._from_iopub_receive_stream,
             create_task_group() as self.task_group,
         ):
-            with open(self.kernelspec_path) as f:
-                kernelspec = json.load(f)
-            cmd = [s.format(connection_file=self.connection_file) for s in kernelspec["argv"]]
-            if cmd and cmd[0] in {
-                "python",
-                f"python{sys.version_info[0]}",
-                "python" + ".".join(map(str, sys.version_info[:2])),
-            }:
-                cmd[0] = sys.executable
+            async with await open_file(self.kernelspec_path) as f:
+                contents = await f.read()
+            kernelspec = json.loads(contents)
+            launch_kernel_cmd = [
+                s.format(connection_file=self.connection_file) for s in kernelspec["argv"]
+            ]
             if self.capture_output:
                 stdout = subprocess.DEVNULL
                 stderr = subprocess.STDOUT
             else:
                 stdout = None
                 stderr = None
             kernel_cwd = self.kernel_cwd if self.kernel_cwd else None
-            self._process = await open_process(cmd, stdout=stdout, stderr=stderr, cwd=kernel_cwd)
+            kernelenv = ""
+            if self.kernelenv_path:
+                path = anyio.Path(self.kernelenv_path)
+                if await path.is_file():
+                    kernelenv = await path.read_text()
+            if kernelenv:
+                import yaml  # type: ignore[import-untyped]
+                env_name = yaml.load(kernelenv, Loader=yaml.CLoader)["name"]
+                cmd = f"micromamba create -f {self.kernelenv_path} --yes"
+                result = await run_process(cmd)
+                if result.returncode == 0:
+                    cmd = """bash -c 'eval "$(micromamba shell hook --shell bash)";""" + \
+                        f"micromamba activate {env_name};" + \
+                        " ".join(launch_kernel_cmd) + "' & echo $!"
+                    process = await open_process(cmd)
+                    assert process.stdout is not None
+                    async for text in TextReceiveStream(process.stdout):
+                        self._pid = int(text)
+                        break
+            else:
+                if launch_kernel_cmd and launch_kernel_cmd[0] in {
+                    "python",
+                    f"python{sys.version_info[0]}",
+                    "python" + ".".join(map(str, sys.version_info[:2])),
+                }:
+                    launch_kernel_cmd[0] = sys.executable
+                self._process = await open_process(
+                    launch_kernel_cmd, stdout=stdout, stderr=stderr, cwd=kernel_cwd
+                )
 
             assert self.connection_cfg is not None
             identity = uuid.uuid4().hex.encode("ascii")
@@ -108,14 +138,17 @@ async def start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) ->
             self.started.set()
 
     async def stop(self) -> None:
-        try:
-            self._process.terminate()
-        except ProcessLookupError:
-            pass
-        await self._process.wait()
-        if self.write_connection_file:
-            path = anyio.Path(self.connection_file)
-            await path.unlink(missing_ok=True)
+        if self._process:
+            try:
+                self._process.terminate()
+            except ProcessLookupError:
+                pass
+            await self._process.wait()
+            if self.write_connection_file:
+                path = anyio.Path(self.connection_file)
+                await path.unlink(missing_ok=True)
+        else:
+            os.kill(self._pid, signal.SIGTERM)
 
         await self.shell_channel.stop()
         await self.stdin_channel.stop()
@@ -124,7 +157,10 @@ async def stop(self) -> None:
         self.task_group.cancel_scope.cancel()
 
     async def interrupt(self) -> None:
-        self._process.send_signal(signal.SIGINT)
+        if self._process:
+            self._process.send_signal(signal.SIGINT)
+        else:
+            os.kill(self._pid, signal.SIGINT)
 
     async def forward_messages_to_shell(self) -> None:
         async for msg in self._to_shell_receive_stream:

plugins/kernel_subprocess/fps_kernel_subprocess/main.py

@@ -10,4 +10,4 @@
 class KernelSubprocessModule(Module):
     async def prepare(self) -> None:
         default_kernel_factory = DefaultKernelFactory(KernelSubprocess)
-        self.put(default_kernel_factory, DefaultKernelFactory)
+        self.put(default_kernel_factory)

plugins/kernels/fps_kernels/kernel_server/server.py

@@ -52,6 +52,7 @@ def __init__(
         self,
         default_kernel_factory: DefaultKernelFactory,
         kernelspec_path: str = "",
+        kernelenv_path: str = "",
         kernel_cwd: str = "",
         connection_file: str = "",
         write_connection_file: bool = True,
@@ -60,6 +61,7 @@ def __init__(
         self.default_kernel_factory = default_kernel_factory
         self.capture_kernel_output = capture_kernel_output
         self.kernelspec_path = kernelspec_path
+        self.kernelenv_path = kernelenv_path
         self.kernel_cwd = kernel_cwd
         self.connection_file = connection_file
         self.write_connection_file = write_connection_file
@@ -109,13 +111,15 @@ async def start(
                     self.kernel = self.default_kernel_factory(
                         write_connection_file=self.write_connection_file,
                         kernelspec_path=self.kernelspec_path,
+                        kernelenv_path=self.kernelenv_path,
                         connection_file=self.connection_file,
                         kernel_cwd=self.kernel_cwd,
                         capture_output=self.capture_kernel_output,
                     )
                 else:
                     self.kernel = kernel_factory(
                         kernelspec_path=self.kernelspec_path,
+                        kernelenv_path=self.kernelenv_path,
                         connection_file=self.connection_file,
                         kernel_cwd=self.kernel_cwd,
                         capture_output=self.capture_kernel_output,

plugins/kernels/fps_kernels/routes.py

@@ -223,6 +223,7 @@ async def create_session(
                 kernel_cwd = kernel_cwd.parent
             kernel_server = KernelServer(
                 kernelspec_path=Path(find_kernelspec(kernel_name)).as_posix(),
+                kernelenv_path=self.kernels_config.kernelenv_path,
                 kernel_cwd=str(kernel_cwd),
                 default_kernel_factory=self.default_kernel_factory,
             )

pyproject.toml

@@ -89,6 +89,7 @@ pre-install-commands = [
   "pip install -e ./plugins/auth_jupyterhub",
   "pip install -e ./plugins/jupyterlab",
   "pip install -e ./plugins/notebook",
+  "pip install pyyaml",
 ]
 features = ["test"]