New Kernels API (#22)

Thankyou. See https://github.com/Zsailer/nextgen-kernels-api/blob/main/README.md#a-next-generation-kernels-api-for-jupyter-server

Author: Zsailer

README.md

@@ -8,6 +8,14 @@ This extension is composed of a Python package named `jupyter_rtc_core`
 for the server extension and a NPM package named `@jupyter/rtc-core`
 for the frontend extension.
 
+## Try it out
+
+Run with the proper configuration
+```
+jupyter lab --config jupyter_config.py
+```
+
+
 ## Requirements
 
 - JupyterLab >= 4.0.0

jupyter_rtc_core/__init__.py

@@ -8,6 +8,7 @@
     warnings.warn("Importing 'jupyter_rtc_core' outside a proper installation.")
     __version__ = "dev"
 
+
 from .app import RtcExtensionApp
 
 

jupyter_rtc_core/app.py

@@ -1,4 +1,5 @@
 from jupyter_server.extension.application import ExtensionApp
+from traitlets.config import Config
 
 from .handlers import RouteHandler
 from .websockets import GlobalAwarenessWebsocket, YRoomWebsocket
@@ -13,11 +14,20 @@ class RtcExtensionApp(ExtensionApp):
         # this can be deleted prior to initial release.
         (r"jupyter-rtc-core/get-example/?", RouteHandler),
         # global awareness websocket
-        (r"api/collaboration/room/JupyterLab:globalAwareness/?", GlobalAwarenessWebsocket),
-        # ydoc websocket
-        (r"api/collaboration/room/(.*)", YRoomWebsocket)
+        # (r"api/collaboration/room/JupyterLab:globalAwareness/?", GlobalAwarenessWebsocket),
+        # # ydoc websocket
+        # (r"api/collaboration/room/(.*)", YRoomWebsocket)
     ]
 
     def initialize(self):
         super().initialize()
 
+
+    def _link_jupyter_server_extension(self, server_app):
+        """Setup custom config needed by this extension."""
+        c = Config()
+        c.ServerApp.kernel_websocket_connection_class = "jupyter_rtc_core.kernels.websocket_connection.NextGenKernelWebsocketConnection"
+        c.ServerApp.kernel_manager_class = "jupyter_rtc_core.kernels.multi_kernel_manager.NextGenMappingKernelManager"
+        c.MultiKernelManager.kernel_manager_class = "jupyter_rtc_core.kernels.kernel_manager.NextGenKernelManager"
+        server_app.update_config(c)
+        super()._link_jupyter_server_extension(server_app)

jupyter_rtc_core/handlers.py

@@ -12,3 +12,6 @@ def get(self):
         self.finish(json.dumps({
             "data": "This is /jupyter-rtc-core/get-example endpoint!"
         }))
+
+
+

jupyter_rtc_core/kernels/__init__.py

@@ -0,0 +1,137 @@
+import asyncio
+import json
+import typing as t
+from traitlets import Set
+from traitlets import Instance
+from traitlets import Any
+from .utils import LRUCache
+from jupyter_client.asynchronous.client import AsyncKernelClient
+import anyio
+
+from jupyter_client.session import Session
+
+class NextGenAsyncKernelClient(AsyncKernelClient): 
+    """
+    A ZMQ-based kernel client class that managers all listeners to a kernel.
+    """
+    # Having this message cache is not ideal. 
+    # Unfortunately, we don't include the parent channel
+    # in the messages that generate IOPub status messages, thus,
+    # we can't differential between the control channel vs.
+    # shell channel status. This message cache gives us 
+    # the ability to map status message back to their source.
+    message_source_cache = Instance(
+        default_value=LRUCache(maxsize=1000), klass=LRUCache
+    )
+
+    # A set of callables that are called when a
+    # ZMQ message comes back from the kernel.
+    _listeners = Set(allow_none=True)
+
+    async def start_listening(self):
+        """Start listening to messages coming from the kernel.
+        
+        Use anyio to setup a task group for listening.
+        """
+        # Wrap a taskgroup so that it can be backgrounded.
+        async def _listening():
+            async with anyio.create_task_group() as tg:
+                for channel_name in ["shell", "control", "stdin", "iopub"]:
+                    tg.start_soon(
+                        self._listen_for_messages, channel_name
+                    )
+    
+        # Background this task.
+        self._listening_task = asyncio.create_task(_listening())
+
+    async def stop_listening(self):
+        # If the listening task isn't defined yet
+        # do nothing.
+        if not self._listening_task:
+            return
+        
+        # Attempt to cancel the task.
+        self._listening_task.cancel()
+        try:
+            # Await cancellation.
+            await self._listening_task
+        except asyncio.CancelledError:
+            self.log.info("Disconnected from client from the kernel.")
+        # Log any exceptions that were raised.
+        except Exception as err:
+            self.log.error(err)
+                        
+    _listening_task: t.Optional[t.Awaitable] = Any(allow_none=True)
+
+    def send_message(self, channel_name, msg):
+        """Use the given session to send the message."""
+        # Cache the message ID and its socket name so that
+        # any response message can be mapped back to the
+        # source channel.
+        header = header = json.loads(msg[0])
+        msg_id = header["msg_id"]
+        self.message_source_cache[msg_id] = channel_name
+        channel = getattr(self, f"{channel_name}_channel")
+        channel.session.send_raw(channel.socket, msg)
+    
+    async def recv_message(self, channel_name, msg):
+        """This is the main method that consumes every
+        message coming back from the kernel. It parses the header
+        (not the content, which might be large) and updates
+        the last_activity, execution_state, and lifecycle_state
+        when appropriate. Then, it routes the message
+        to all listeners.
+        """
+        # Broadcast messages
+        async with anyio.create_task_group() as tg:
+            # Broadcast the message to all listeners.
+            for listener in self._listeners:
+                async def _wrap_listener(listener, channel_name, msg): 
+                    """
+                    Wrap the listener to ensure its async and 
+                    logs (instead of raises) exceptions.
+                    """
+                    try:
+                        listener(channel_name, msg)
+                    except Exception as err:
+                        self.log.error(err)
+                
+                tg.start_soon(_wrap_listener, listener, channel_name, msg)
+
+    def add_listener(self, callback: t.Callable[[dict], None]):
+        """Add a listener to the ZMQ Interface.
+
+        A listener is a callable function/method that takes
+        the deserialized (minus the content) ZMQ message.
+
+        If the listener is already registered, it won't be registered again.
+        """
+        self._listeners.add(callback)
+
+    def remove_listener(self, callback: t.Callable[[dict], None]):
+        """Remove a listener to teh ZMQ interface. If the listener
+        is not found, this method does nothing.
+        """
+        self._listeners.discard(callback)
+
+    async def _listen_for_messages(self, channel_name):
+        """The basic polling loop for listened to kernel messages
+        on a ZMQ socket.
+        """
+        # Wire up the ZMQ sockets
+        # Setup up ZMQSocket broadcasting.
+        channel = getattr(self, f"{channel_name}_channel")
+        while True:
+            # Wait for a message
+            await channel.socket.poll(timeout=float("inf"))
+            raw_msg = await channel.socket.recv_multipart()
+            try:
+                await self.recv_message(channel_name, raw_msg)
+            except Exception as err:
+                self.log.error(err)
+    
+    def kernel_info(self):
+        msg = self.session.msg("kernel_info_request")
+        # Send message, skipping the delimiter and signature
+        msg = self.session.serialize(msg)[2:]
+        self.send_message("shell", msg)

jupyter_rtc_core/kernels/kernel_client.py

@@ -0,0 +1,195 @@
+import typing
+import asyncio
+from traitlets import default
+from traitlets import Instance
+from traitlets import Int
+from traitlets import Dict
+from traitlets import Type
+from traitlets import Unicode
+from traitlets import validate
+from traitlets import observe
+from traitlets import Set
+from traitlets import TraitError
+from traitlets import DottedObjectName
+from traitlets.utils.importstring import import_item
+
+from jupyter_client.manager import AsyncKernelManager
+
+from . import types
+from . import states
+from .kernel_client import AsyncKernelClient
+
+
+class NextGenKernelManager(AsyncKernelManager):
+
+    main_client = Instance(AsyncKernelClient, allow_none=True)
+
+    client_class = DottedObjectName(
+        "jupyter_rtc_core.kernels.kernel_client.NextGenAsyncKernelClient"
+    )
+    
+    client_factory: Type = Type(klass="jupyter_rtc_core.kernels.kernel_client.NextGenAsyncKernelClient")
+
+    # Configurable settings in a kernel manager that I want.
+    time_to_connect: int = Int(
+        default_value=10,
+        help="The timeout for connecting to a kernel."
+    ).tag(config=True)
+    
+    execution_state: types.EXECUTION_STATES = Unicode()
+    
+    @validate("execution_state")
+    def _validate_execution_state(self, proposal: dict):
+        if not proposal["value"] in states.EXECUTION_STATES:
+            raise TraitError(f"execution_state must be one of {states.EXECUTION_STATES}")
+        return proposal["value"]
+
+    lifecycle_state: types.EXECUTION_STATES = Unicode()
+    
+    @validate("lifecycle_state")
+    def _validate_lifecycle_state(self, proposal: dict):
+        if not proposal["value"] in states.LIFECYCLE_STATES:
+            raise TraitError(f"lifecycle_state must be one of {states.LIFECYCLE_STATES}")
+        return proposal["value"]
+    
+    state = Dict()
+    
+    @default('state')
+    def _default_state(self):
+        return {
+            "execution_state": self.execution_state,
+            "lifecycle_state": self.lifecycle_state
+        }
+    
+    @observe('execution_state')
+    def _observer_execution_state(self, change):
+        state = self.state
+        state["execution_state"] = change['new']
+        self.state = state    
+    
+    @observe('lifecycle_state')
+    def _observer_lifecycle_state(self, change):
+        state = self.state
+        state["lifecycle_state"] = change['new']
+        self.state = state    
+    
+    @validate('state')
+    def _validate_state(self, change):
+        value = change['value']
+        if 'execution_state' not in value or 'lifecycle_state' not in value:
+            TraitError("State needs to include execution_state and lifecycle_state")
+        return value
+    
+    @observe('state')
+    def _state_changed(self, change):
+        for observer in self._state_observers:
+            observer(change["new"])
+    
+    _state_observers = Set(allow_none=True)
+    
+    def set_state(
+        self, 
+        lifecycle_state: typing.Optional[types.LIFECYCLE_STATES] = None, 
+        execution_state: typing.Optional[types.EXECUTION_STATES] = None,
+        broadcast=True
+    ):
+        if lifecycle_state:
+            self.lifecycle_state = lifecycle_state
+        if execution_state:
+            self.execution_state = execution_state
+            
+        if broadcast:
+            # Broadcast this state change to all listeners
+            self.broadcast_state()
+
+    async def start_kernel(self, *args, **kwargs):
+        self.set_state("starting", "starting")
+        out = await super().start_kernel(*args, **kwargs)
+        self.set_state("started")
+        await self.connect()
+        return out
+        
+    async def shutdown_kernel(self, *args, **kwargs):
+        self.set_state("terminating")
+        await self.disconnect()
+        out = await super().shutdown_kernel(*args, **kwargs)
+        self.set_state("terminated", "dead")
+     
+    async def restart_kernel(self, *args, **kwargs):
+        self.set_state("restarting")
+        return await super().restart_kernel(*args, **kwargs)
+
+    async def connect(self):
+        """Open a single client interface to the kernel.
+        
+        Ideally this method doesn't care if the kernel
+        is actually started. It will just try a ZMQ 
+        connection anyways and wait. This is helpful for
+        handling 'pending' kernels, which might still 
+        be in a starting phase. We can keep a connection
+        open regardless if the kernel is ready. 
+        """
+        self.set_state("connecting", "busy")
+        # Use the new API for getting a client.
+        self.main_client = self.client()
+        # Track execution state by watching all messages that come through
+        # the kernel client.
+        self.main_client.add_listener(self.execution_state_listener)
+        self.main_client.start_channels()
+        await self.main_client.start_listening()
+        # The Heartbeat channel is paused by default; unpause it here
+        self.main_client.hb_channel.unpause()
+        # Wait for a living heartbeat.
+        attempt = 0
+        while not self.main_client.hb_channel.is_alive():
+            attempt += 1
+            if attempt > self.time_to_connect:
+                # Set the state to unknown.
+                self.set_state("unknown", "unknown")
+                raise Exception("The kernel took too long to connect to the ZMQ sockets.")
+            # Wait a second until the next time we try again.
+            await asyncio.sleep(1)
+        # Send an initial kernel info request on the shell channel.
+        self.main_client.kernel_info()
+        self.set_state("connected")
+        
+    async def disconnect(self):
+        await self.main_client.stop_listening()
+        self.main_client.stop_channels()
+
+    def broadcast_state(self):
+        """Broadcast state to all listeners"""
+        if not self.main_client:
+            return 
+        
+        # Emit this state to all listeners
+        for listener in self.main_client._listeners:
+            # Manufacture a status message
+            session = self.main_client.session
+            msg = session.msg("status", {"execution_state": self.execution_state})
+            msg = session.serialize(msg)
+            listener("iopub", msg)    
+            
+    def execution_state_listener(self, channel_name, msg):
+        """Set the execution state by watching messages returned by the shell channel."""
+        # Only continue if we're on the IOPub where the status is published.
+        if channel_name != "iopub":
+            return
+        session = self.main_client.session        
+        _, smsg = session.feed_identities(msg)
+        # Unpack the message 
+        deserialized_msg = session.deserialize(smsg, content=False)
+        if deserialized_msg["msg_type"] == "status":
+            content = session.unpack(deserialized_msg["content"])
+            status = content["execution_state"]
+            if status == "starting":
+                # Don't broadcast, since this message is already going out.
+                self.set_state("starting", status, broadcast=False)
+            else:
+                parent = deserialized_msg.get("parent_header", {})
+                msg_id = parent.get("msg_id", "")
+                parent_channel = self.main_client.message_source_cache.get(msg_id, None)
+                if parent_channel and parent_channel == "shell":
+                    # Don't broadcast, since this message is already going out.
+                    self.set_state("connected", status, broadcast=False)
+