Hubble metrics to cell latency, execution count, and notebook service calls (jupyter-server#407)

* hubble metrics to measure execution count and latency and all notebook service calls

* Bump to 0.21.1

* add metric for a disconnected kernel

Author: Zsailer

data_studio_jupyter_extensions/__init__.py

@@ -30,6 +30,9 @@ def _jupyter_server_extension_points():  # pragma: no cover
     from data_studio_jupyter_extensions.extensions.kernel_actions.extension import (
         KernelActionsExtension,
     )
+    from data_studio_jupyter_extensions.extensions.kernel_websocket_override.extension import (
+        KernelWebsocketOverrideExtension,
+    )
 
     return [
         {
@@ -56,4 +59,8 @@ def _jupyter_server_extension_points():  # pragma: no cover
             "module": "data_studio_jupyter_extensions.extensions.external_links.extension",
             "app": KernelActionsExtension,
         },
+        {
+            "module": "data_studio_jupyter_extensions.extensions.kernel_websocket_override.extension",
+            "app": KernelWebsocketOverrideExtension,
+        },
     ]

data_studio_jupyter_extensions/_version.py

@@ -1 +1 @@
-__version__ = "0.21.0"  # pragma: no cover
+__version__ = "0.21.1"  # pragma: no cover

data_studio_jupyter_extensions/configurables/kernel_monitor.py

@@ -160,8 +160,8 @@ async def poll(self):
                 # If the kernel is alive, but not communicating,
                 # raise a helpful message in the status menu and
                 # shutdown the kernel.
+                HUBBLE_METRICS.KERNEL_DISCONNECTED.emit()
                 self._disconnected_state()
-                await self.kernel_manager.shutdown_kernel()
                 self._attempt_count = 0
 
     def is_culling_enabled(self) -> bool:

data_studio_jupyter_extensions/configurables/notebook_service.py

@@ -195,23 +195,28 @@ async def get_list_of_kernelspecs_for_notebook(self):
         out = await self.fetch("kernelspecs", method="get")
         return out
 
+    @hubble(HUBBLE_METRICS.NBSERVICE_KERNELSPECS_BY_NAME)
     async def get_kernelspec_by_name(self, name):
         """Get the kernelspec by name."""
         return await self.fetch("kernelspec", name, method="get")
 
+    @hubble(HUBBLE_METRICS.NBSERVICE_START_KERNEL)
     async def start_kernel(self, kernelspec_id):
         """Start a kernel."""
         data = {"kernelSpecId": kernelspec_id}
         return await self.fetch("kernels", method="post", data=data)
 
+    @hubble(HUBBLE_METRICS.NBSERVICE_GET_KERNEL_DETAILS)
     async def get_kernel_details(self, process_id):
         """Get kernel details by id."""
         return await self.fetch(f"/kernels/{process_id}", method="get")
 
+    @hubble(HUBBLE_METRICS.NBSERVICE_STOP_KERNEL)
     async def stop_kernel(self, process_id):
         """Stop kernel by id."""
         return await self.fetch(f"/kernels/{process_id}", method="delete")
 
+    @hubble(HUBBLE_METRICS.NBSERVICE_GET_KERNEL_STATUS)
     async def get_kernel_status(self, process_id, query_params_dict=None):
         """Get kernel status by id."""
         url = f"/kernels/{process_id}/status"
@@ -220,10 +225,12 @@ async def get_kernel_status(self, process_id, query_params_dict=None):
             url = f"{url}?{query_params}"
         return await self.fetch(url, method="get")
 
+    @hubble(HUBBLE_METRICS.NBSERVICE_LIST_KERNELS)
     async def list_kernels(self):
         """List all kernels for the given notebook server."""
         return await self.fetch("kernels", method="get")
 
+    @hubble(HUBBLE_METRICS.NBSERVICE_SHUTDOWN_ALL_KERNELS)
     async def shutdown_all_kernels(self):
         """Shutdown all kernels for the given notebook server."""
         return await self.fetch("kernels", method="delete")

data_studio_jupyter_extensions/extensions/kernel_websocket_override/__init__.py

@@ -0,0 +1,12 @@
+from data_studio_jupyter_extensions.extensions.kernel_websocket_override.extension import (
+    KernelWebsocketOverrideExtension,
+)
+
+
+def _jupyter_server_extension_points():  # pragma: no cover
+    return [
+        {
+            "module": "data_studio_jupyter_extensions.extensions.kernel_websocket_override.extension",
+            "app": KernelWebsocketOverrideExtension,
+        }
+    ]

data_studio_jupyter_extensions/extensions/kernel_websocket_override/extension.py

@@ -0,0 +1,12 @@
+from jupyter_server.extension.application import ExtensionApp
+
+from .handlers import handlers
+
+
+class KernelWebsocketOverrideExtension(ExtensionApp):
+    """Jupyter Server extension that verifies
+    the health of the server.
+    """
+
+    name = "kernel_websocket_override"
+    handlers = handlers

data_studio_jupyter_extensions/extensions/kernel_websocket_override/handlers.py

@@ -0,0 +1,59 @@
+from jupyter_server.base.zmqhandlers import AuthenticatedZMQStreamHandler
+from jupyter_server.services.kernels.handlers import _kernel_id_regex
+from jupyter_server.services.kernels.handlers import ZMQChannelsHandler
+
+from data_studio_jupyter_extensions.hubble.metrics import HUBBLE_METRICS
+
+
+def hubble_execution_time_and_count(msg):
+    """Emit the time it takes to run a cell and the execution cound to hubble."""
+    msg_type = msg.get("msg_type", None)
+    if msg_type == "execute_reply":
+        tdelta = msg["header"]["date"] - msg["parent_header"]["date"]
+        seconds = tdelta.total_seconds()
+        HUBBLE_METRICS.KERNEL_EXECUTION_LATENCY.emit(value=seconds)
+        count = msg["content"]["execution_count"]
+        HUBBLE_METRICS.KERNEL_EXECUTION_COUNT.emit(value=count)
+
+
+class OverrideZMQChannelsHandler(ZMQChannelsHandler):
+    """Adds hubble metrics to Jupyter websocket handler."""
+
+    # This is a complete fork of the Jupyter Websocket handler
+    # method. We don't want to reply on this long time, but should
+    # work on getting this into Jupyter Telemetry.
+    def _on_zmq_reply(self, stream, msg_list):
+        idents, fed_msg_list = self.session.feed_identities(msg_list)
+
+        if self.subprotocol == "v1.kernel.websocket.jupyter.org":
+            msg = {"header": None, "parent_header": None, "content": None}
+        else:
+            msg = self.session.deserialize(fed_msg_list)
+
+        # This chunk is (the only) DataStudio customization for collecting hubble metrics.
+        # if this chunk fails, don't stop messages from going through. just log a warning.
+        try:
+            hubble_execution_time_and_count(msg)
+        except Exception as err:
+            self.log.warning(err)
+
+        channel = getattr(stream, "channel", None)
+        parts = fed_msg_list[1:]
+
+        self._on_error(channel, msg, parts)
+
+        if self._limit_rate(channel, msg, parts):
+            return
+
+        if self.subprotocol == "v1.kernel.websocket.jupyter.org":
+            AuthenticatedZMQStreamHandler._on_zmq_reply(self, stream, parts)
+        else:
+            AuthenticatedZMQStreamHandler._on_zmq_reply(self, stream, msg)
+
+
+handlers = [
+    (
+        rf"/api/kernels/{_kernel_id_regex}/channels",
+        OverrideZMQChannelsHandler,
+    ),
+]

data_studio_jupyter_extensions/hubble/metrics.py

@@ -22,6 +22,9 @@ class _HUBBLE_METRICS:
     KERNEL_POD_LAUNCH_TIME = Metric("kernel_pod_launch_time")
     # Compute how long it takes for a kernel pod + first heartbeat to be established
     KERNEL_READY_LATENCY = Metric("kernel_ready_latency")
+    KERNEL_EXECUTION_COUNT = Metric("kernel_execution_count")
+    KERNEL_EXECUTION_LATENCY = Metric("kernel_execution_latency")
+    KERNEL_DISCONNECTED = Metric("kernel_disconnected")
 
     # Hubble metrics from Jupyter Server handlers.
     API_KERNELS_ROOT = HandlerMetric("/api/kernels")
@@ -43,6 +46,13 @@ class _HUBBLE_METRICS:
     # Hubble metrics from Notebook Service responses.
     NBSERVICE_KERNELSPECS = ResponseMetric("/nbservice/kernelspecs")
     NBSERVICE_LINKS = ResponseMetric("/nbservice/links")
+    NBSERVICE_KERNELSPECS_BY_NAME = ResponseMetric("/nbservice/kernelspecs/*")
+    NBSERVICE_START_KERNEL = ResponseMetric("/nbservice/kernels")
+    NBSERVICE_GET_KERNEL_DETAILS = ResponseMetric("/nbservice/kernels/*")
+    NBSERVICE_STOP_KERNEL = ResponseMetric("/nbservice/kernels/*")
+    NBSERVICE_GET_KERNEL_STATUS = ResponseMetric("/nbservice/kernels/*/status")
+    NBSERVICE_LIST_KERNELS = ResponseMetric("/nbservice/kernels")
+    NBSERVICE_SHUTDOWN_ALL_KERNELS = ResponseMetric("/nbservice/kernels")
 
 
 # Instantiate the Hubble metrics dataclass for reference

pyproject.toml

@@ -23,7 +23,7 @@ testpaths = [
 ]
 
 [tool.tbump.version]
-current = "0.21.0"
+current = "0.21.1"
 regex = '''
   (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)
   ((?P<channel>a|b|rc|.dev)(?P<release>\d+))?

src/status.tsx

@@ -31,7 +31,7 @@ import { ISessionContext, ReactWidget } from '@jupyterlab/apputils';
 import { EventListener } from './eventlistener';
 
 // This is managed by tbump config in pyproject.toml
-const VERSION = '0.21.0';
+const VERSION = '0.21.1';
 
 // Define the error states
 // https://github.pie.apple.com/pie-data-studio/notebook-service/blob/761d63604966db5918d2e491c0f89cce454b7f67/app/com/apple/datastudio/model/ResourceState.scala#L20

version

@@ -1 +1 @@
-0.21.0
+0.21.1