wire up backend to kernel execution state

Author: Zsailer

jupyter_rtc_core/kernels/kernel_client.py

@@ -11,7 +11,7 @@
 from jupyter_client.asynchronous.client import AsyncKernelClient
 import anyio
 from jupyter_rtc_core.rooms.yroom import YRoom
-
+from jupyter_server.utils import ensure_async
 
 class DocumentAwareKernelClient(AsyncKernelClient): 
     """
@@ -140,7 +140,7 @@ async def _wrap_listener(listener_to_wrap, channel_name, msg):
                     logs (instead of raises) exceptions.
                     """
                     try:
-                        listener_to_wrap(channel_name, msg)
+                        await ensure_async(listener_to_wrap(channel_name, msg))
                     except Exception as err:
                         self.log.error(err)
 
@@ -154,19 +154,17 @@ async def handle_outgoing_message(self, channel_name: str, msg: list[bytes]):
         when appropriate. Then, it routes the message
         to all listeners.
         """
-
         # Intercept messages that are IOPub focused.
         if channel_name == "iopub":
             message_returned = await self.handle_iopub_message(msg)
             # TODO: If the message is not returned by the iopub handler, then
             # return here and do not forward to listeners.
             if not message_returned:
-                self.log.warn(f"If message is handled donot forward after adding output manager")
+                self.log.warn(f"If message is handled do not forward after adding output manager")
                 return
 
         # Update the last activity.
         # self.last_activity = self.session.msg_time
-
         await self.send_message_to_listeners(channel_name, msg)
 
     async def handle_iopub_message(self, msg: list[bytes]) -> t.Optional[list[bytes]]:

jupyter_rtc_core/kernels/kernel_manager.py

@@ -59,62 +59,25 @@ def _validate_lifecycle_state(self, proposal: dict):
             raise TraitError(f"lifecycle_state must be one of {LifecycleStates}")
         return 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: LifecycleStates = None, 
         execution_state: ExecutionStates = None,
-        broadcast=True
     ):
         if lifecycle_state:
             self.lifecycle_state = lifecycle_state.value
         if execution_state:
             self.execution_state = execution_state.value
-            
-        if broadcast:
-            # Broadcast this state change to all listeners
-            self._state_observers = None
-            self.broadcast_state()
 
     async def start_kernel(self, *args, **kwargs):
         self.set_state(LifecycleStates.STARTING, ExecutionStates.STARTING)
         out = await super().start_kernel(*args, **kwargs)
         self.set_state(LifecycleStates.STARTED)
-        await self.connect()
+        # Schedule the kernel to connect.
+        # Do not await here, since many clients expect
+        # the server to complete the start flow even
+        # if the kernel is not fully connected yet.
+        task = asyncio.create_task(self.connect())
         return out
 
     async def shutdown_kernel(self, *args, **kwargs):
@@ -137,12 +100,13 @@ async def connect(self):
         be in a starting phase. We can keep a connection
         open regardless if the kernel is ready. 
         """
-        self.set_state(LifecycleStates.CONNECTING, ExecutionStates.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.set_state(LifecycleStates.CONNECTING, ExecutionStates.STARTING)
+        await self.broadcast_state()
         self.main_client.start_channels()
         await self.main_client.start_listening()
         # The Heartbeat channel is paused by default; unpause it here
@@ -157,33 +121,34 @@ async def connect(self):
                 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.send_kernel_info()
-        self.set_state(LifecycleStates.CONNECTED)
+        # Wait for the kernel to reach an idle state.
+        while self.execution_state != ExecutionStates.IDLE.value:
+            self.main_client.send_kernel_info()
+            await asyncio.sleep(1)
 
     async def disconnect(self):
         await self.main_client.stop_listening()
         self.main_client.stop_channels()
 
-    def broadcast_state(self):
+    async 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)
-            _, fed_msg_list = self.session.feed_identities(msg)
-            listener("iopub", fed_msg_list)    
+
+        # Manufacture an IOPub status message from the shell channel.
+        session = self.main_client.session
+        parent_header = session.msg_header("status")
+        parent_msg_id = parent_header["msg_id"]
+        self.main_client.message_source_cache[parent_msg_id] = "shell"
+        msg = session.msg("status", content={"execution_state": self.execution_state}, parent=parent_header)
+        smsg = session.serialize(msg)[1:]
+        await self.main_client.handle_outgoing_message("iopub", smsg)
 
     def execution_state_listener(self, channel_name: str, msg: list[bytes]):
         """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
+            return  
 
         session = self.main_client.session        
         # Unpack the message 
@@ -193,22 +158,18 @@ def execution_state_listener(self, channel_name: str, msg: list[bytes]):
             execution_state = content["execution_state"]
             if execution_state == "starting":
                 # Don't broadcast, since this message is already going out.
-                self.set_state(LifecycleStates.STARTING, ExecutionStates.STARTING, broadcast=False)
+                self.set_state(execution_state=ExecutionStates.STARTING)
             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(LifecycleStates.CONNECTED, ExecutionStates(execution_state), broadcast=False)
-
+                    self.set_state(LifecycleStates.CONNECTED, ExecutionStates(execution_state))
+                    
             kernel_status = {
                 "execution_state": self.execution_state,
                 "lifecycle_state": self.lifecycle_state
-            }
+            }        
             self.log.debug(f"Sending kernel status awareness {kernel_status}")
             self.main_client.send_kernel_awareness(kernel_status)
-            self.log.debug(f"Sent kernel status awareness {kernel_status}")
-
-
-            
+            self.log.debug(f"Sent kernel status awareness {kernel_status}")

jupyter_rtc_core/kernels/websocket_connection.py

@@ -20,7 +20,7 @@ async def connect(self):
         asyncio task happening in parallel.
         """
         self.kernel_manager.main_client.add_listener(self.handle_outgoing_message)
-        self.kernel_manager.broadcast_state()
+        await self.kernel_manager.broadcast_state()
         self.log.info("Kernel websocket is now listening to kernel.")
 
     def disconnect(self):

jupyter_rtc_core/rooms/yroom.py

@@ -384,7 +384,7 @@ def send_server_awareness(self, type: str, changes: tuple[dict[str, Any], Any])
         Arguments:
             type: The change type.
             changes: The awareness changes.
-        """
+        """        
         if type != "update" or changes[1] != "local":
             return
 

src/executionindicator.tsx

@@ -13,10 +13,6 @@ import {
   ExecutionIndicatorComponent, 
   ExecutionIndicator as E,
 } from '@jupyterlab/notebook';
-import { Poll } from '@lumino/polling';
-
-
-const EXECUTION_STATE_KEY = "executionState";
 
 
 /**
@@ -77,22 +73,7 @@ export namespace AwarenessExecutionIndicator {
       if (!nb) {
         return;
       }
-      (this as any)._currentNotebook = nb;
-      // Artificially toggle the execution state.
-      new Poll({
-        auto: true,
-        factory: () => {
-          let fullState = nb?.model?.sharedModel.awareness.getLocalState();
-          let state = 'idle'
-          if (fullState && fullState[EXECUTION_STATE_KEY] === "idle") {
-            state = 'busy'
-          }
-          nb?.model?.sharedModel.awareness.setLocalState({ executionState: state}) 
-          return Promise.resolve()
-        },
-        frequency: { interval: 2000, backoff: false }
-      });
-  
+      (this as any)._currentNotebook = nb;  
       (this as any)._notebookExecutionProgress.set(nb, {
         executionStatus: 'idle',
         kernelStatus: 'idle',
@@ -107,13 +88,17 @@ export namespace AwarenessExecutionIndicator {
 
       const contextStatusChanged = (ctx: ISessionContext) => {
         if (state) {
-          let fullState = nb?.model?.sharedModel.awareness.getLocalState();
-          if (fullState) {
-            let currentState = fullState[EXECUTION_STATE_KEY];
-            state.kernelStatus = currentState;
+          let awarenessStates = nb?.model?.sharedModel.awareness.getStates();
+          if (awarenessStates) {
+            for (let [_, clientState] of awarenessStates) {
+              if ('kernel' in clientState) {
+                state.kernelStatus = clientState['kernel']['execution_state'];
+                this.stateChanged.emit(void 0);
+                return;
+              }
+            }
           }
         }
-        this.stateChanged.emit(void 0);
       };
 
       nb?.model?.sharedModel.awareness.on('change', contextStatusChanged);

src/kernelstatus.tsx

@@ -16,8 +16,6 @@ import {
 import { NotebookPanel } from '@jupyterlab/notebook';
 import { Session } from '@jupyterlab/services';
 
-const EXECUTION_STATE_KEY = "executionState";
-
 /**
  * A pure functional component for rendering kernel status.
  */
@@ -133,11 +131,15 @@ export namespace AwarenessKernelStatus {
       let panel = (widget as NotebookPanel);
       const stateChanged = () => {
         if (this) {
-          let fullState = panel?.model?.sharedModel.awareness.getLocalState();
-          if (fullState) {
-            let currentState = fullState[EXECUTION_STATE_KEY];
-            (this as any)._kernelStatus = currentState;
-            (this as any).stateChanged.emit(void 0);
+          let awarenessStates = panel?.model?.sharedModel.awareness.getStates();
+          if (awarenessStates) {
+            for (let [_, clientState] of awarenessStates) {
+              if ('kernel' in clientState) {
+                (this as any)._kernelStatus = clientState['kernel']['execution_state'];
+                (this as any).stateChanged.emit(void 0);
+                return;
+              }
+            }
           }
         }
       };