Merge pull request #164 from davidbrochart/jupyterlab-kernel-usage

Jupyterlab kernel usage

Author: jtpio

.github/workflows/build.yml

@@ -25,6 +25,10 @@ jobs:
     - name: Install the extension
       run: |
         python -m pip install .
+        jupyter server extension enable --py jupyter_resource_usage  --sys-prefix
+        jupyter serverextension enable --py jupyter_resource_usage --sys-prefix
+        jupyter nbextension install --py jupyter_resource_usage --sys-prefix
+        jupyter nbextension enable --py jupyter_resource_usage --sys-prefix
 
     - name: Check the server, classic and lab extensions are installed
       run: |

CONTRIBUTING.md

@@ -96,7 +96,7 @@ python -m pip install -e ".[dev]"
 jupyter labextension develop . --overwrite
 
 # go to the labextension directory
-cd labextension/
+cd packages/labextension/
 
 # Rebuild extension Typescript source after making changes
 jlpm run build

README.md

@@ -19,6 +19,9 @@ displays an indication of how much resources your current notebook server and
 its children (kernels, terminals, etc) are using. This is displayed in the
 main toolbar in the notebook itself, refreshing every 5s.
 
+Kernel resource usage can be displayed in a sidebar for IPython kernels with
+[ipykernel](https://github.com/ipython/ipykernel) >= 6.11.0.
+
 ## Installation
 
 You can currently install this package from PyPI.
@@ -36,9 +39,9 @@ conda install -c conda-forge jupyter-resource-usage
 **If your notebook version is < 5.3**, you need to enable the extension manually.
 
 ```
-jupyter serverextension enable --py jupyter-resource-usage --sys-prefix
-jupyter nbextension install --py jupyter-resource-usage --sys-prefix
-jupyter nbextension enable --py jupyter-resource-usage --sys-prefix
+jupyter serverextension enable --py jupyter_resource_usage --sys-prefix
+jupyter nbextension install --py jupyter_resource_usage --sys-prefix
+jupyter nbextension enable --py jupyter_resource_usage --sys-prefix
 ```
 
 ## Configuration

jupyter_resource_usage/__init__.py

@@ -1,13 +1,12 @@
 import json
-import os.path as osp
+from pathlib import Path
 
 from ._version import __version__
 from .server_extension import load_jupyter_server_extension
 
-HERE = osp.abspath(osp.dirname(__file__))
+HERE = Path(__file__).parent.resolve()
 
-with open(osp.join(HERE, "labextension", "package.json")) as fid:
-    data = json.load(fid)
+data = json.loads((HERE / "labextension" / "package.json").read_text())
 
 
 def _jupyter_labextension_paths():

jupyter_resource_usage/api.py

@@ -2,7 +2,11 @@
 from concurrent.futures import ThreadPoolExecutor
 
 import psutil
+import zmq
+from jupyter_client.jsonutil import date_default
 from jupyter_server.base.handlers import APIHandler
+from jupyter_server.utils import url_path_join
+from packaging import version
 from tornado import web
 from tornado.concurrent import run_on_executor
 
@@ -13,6 +17,16 @@
     from .utils import Callable
 
 
+try:
+    import ipykernel
+
+    USAGE_IS_SUPPORTED = version.parse("6.9.0") <= version.parse(ipykernel.__version__)
+except ImportError:
+    USAGE_IS_SUPPORTED = False
+
+MAX_RETRIES = 3
+
+
 class ApiHandler(APIHandler):
     executor = ThreadPoolExecutor(max_workers=5)
 
@@ -74,3 +88,37 @@ def get_cpu_percent(p):
                 return 0
 
         return sum([get_cpu_percent(p) for p in all_processes])
+
+
+class KernelUsageHandler(APIHandler):
+    @web.authenticated
+    async def get(self, matched_part=None, *args, **kwargs):
+
+        if not USAGE_IS_SUPPORTED:
+            self.write(json.dumps({}))
+            return
+
+        kernel_id = matched_part
+        km = self.kernel_manager
+        lkm = km.pinned_superclass.get_kernel(km, kernel_id)
+        session = lkm.session
+        client = lkm.client()
+
+        control_channel = client.control_channel
+        usage_request = session.msg("usage_request", {})
+
+        control_channel.send(usage_request)
+        poller = zmq.Poller()
+        control_socket = control_channel.socket
+        poller.register(control_socket, zmq.POLLIN)
+        for i in range(1, MAX_RETRIES + 1):
+            timeout_ms = 1000 * i
+            events = dict(poller.poll(timeout_ms))
+            if not events:
+                self.write(json.dumps({}))
+                break
+            if control_socket not in events:
+                continue
+            res = await client.control_channel.get_msg(timeout=0)
+            self.write(json.dumps(res, default=date_default))
+            break

jupyter_resource_usage/server_extension.py

@@ -2,6 +2,7 @@
 from tornado import ioloop
 
 from jupyter_resource_usage.api import ApiHandler
+from jupyter_resource_usage.api import KernelUsageHandler
 from jupyter_resource_usage.config import ResourceUseDisplay
 from jupyter_resource_usage.metrics import PSUtilMetricsLoader
 from jupyter_resource_usage.prometheus import PrometheusHandler
@@ -18,6 +19,17 @@ def load_jupyter_server_extension(server_app):
     server_app.web_app.add_handlers(
         ".*", [(url_path_join(base_url, "/api/metrics/v1"), ApiHandler)]
     )
+    server_app.web_app.add_handlers(
+        ".*$",
+        [
+            (
+                url_path_join(
+                    base_url, "/api/metrics/v1/kernel_usage", r"get_usage/(.+)$"
+                ),
+                KernelUsageHandler,
+            )
+        ],
+    )
 
     if resuseconfig.enable_prometheus_metrics:
         callback = ioloop.PeriodicCallback(

packages/labextension/package.json

@@ -19,6 +19,7 @@
   ],
   "main": "lib/index.js",
   "types": "lib/index.d.ts",
+  "style": "style/index.css",
   "repository": {
     "type": "git",
     "url": "https://github.com/jupyter-server/jupyter-resource-usage.git"
@@ -41,22 +42,48 @@
     "watch:labextension": "jupyter labextension watch ."
   },
   "dependencies": {
-    "@jupyterlab/application": "^3.0.0",
-    "@jupyterlab/apputils": "^3.0.0",
-    "@jupyterlab/coreutils": "^5.0.0",
-    "@jupyterlab/services": "^6.0.0",
-    "@jupyterlab/statusbar": "^3.0.0",
-    "@jupyterlab/translation": "^3.0.0",
-    "@lumino/polling": "^1.3.3",
-    "typestyle": "^2.0.4"
+    "@jupyterlab/application": "^3.5.1",
+    "@jupyterlab/apputils": "^3.5.1",
+    "@jupyterlab/coreutils": "^5.5.1",
+    "@jupyterlab/notebook": "^3.5.1",
+    "@jupyterlab/services": "^6.5.1",
+    "@jupyterlab/statusbar": "^3.5.1",
+    "@jupyterlab/translation": "^3.5.1",
+    "@lumino/polling": "^1.11.3",
+    "typestyle": "^2.4.0"
   },
   "devDependencies": {
-    "@jupyterlab/builder": "^3.0.0",
+    "@jupyterlab/builder": "^3.5.1",
+    "@typescript-eslint/eslint-plugin": "^4.8.1",
+    "@typescript-eslint/parser": "^4.8.1",
+    "eslint": "^7.14.0",
+    "eslint-config-prettier": "^6.15.0",
+    "eslint-plugin-prettier": "^3.1.4",
+    "mkdirp": "^1.0.3",
     "npm-run-all": "^4.1.5",
+    "prettier": "^2.1.1",
     "rimraf": "^3.0.2",
-    "typescript": "~4.0.3"
+    "typescript": "~4.1.3"
+  },
+  "sideEffects": [
+    "style/*.css",
+    "style/index.js"
+  ],
+  "styleModule": "style/index.js",
+  "publishConfig": {
+    "access": "public"
   },
   "jupyterlab": {
+    "discovery": {
+      "server": {
+        "managers": [
+          "pip"
+        ],
+        "base": {
+          "name": "jupyterlab_kernel_usage"
+        }
+      }
+    },
     "extension": true,
     "outputDir": "../../jupyter_resource_usage/labextension"
   }

packages/labextension/src/format.ts

@@ -0,0 +1,60 @@
+/**
+ * The type of unit used for reporting memory usage.
+ */
+export type MemoryUnit = 'B' | 'KB' | 'MB' | 'GB' | 'TB' | 'PB';
+
+/**
+ * The number of bytes in each memory unit.
+ */
+export const MEMORY_UNIT_LIMITS: {
+  readonly [U in MemoryUnit]: number;
+} = {
+  B: 1,
+  KB: 1024,
+  MB: 1048576,
+  GB: 1073741824,
+  TB: 1099511627776,
+  PB: 1125899906842624,
+};
+
+export function formatForDisplay(numBytes: number | undefined): string {
+  const lu = convertToLargestUnit(numBytes);
+  return lu[0].toFixed(2) + ' ' + lu[1];
+}
+
+/**
+ * Given a number of bytes, convert to the most human-readable
+ * format, (GB, TB, etc).
+ */
+export function convertToLargestUnit(
+  numBytes: number | undefined
+): [number, MemoryUnit] {
+  if (!numBytes) {
+    return [0, 'B'];
+  }
+  if (numBytes < MEMORY_UNIT_LIMITS.KB) {
+    return [numBytes, 'B'];
+  } else if (
+    MEMORY_UNIT_LIMITS.KB === numBytes ||
+    numBytes < MEMORY_UNIT_LIMITS.MB
+  ) {
+    return [numBytes / MEMORY_UNIT_LIMITS.KB, 'KB'];
+  } else if (
+    MEMORY_UNIT_LIMITS.MB === numBytes ||
+    numBytes < MEMORY_UNIT_LIMITS.GB
+  ) {
+    return [numBytes / MEMORY_UNIT_LIMITS.MB, 'MB'];
+  } else if (
+    MEMORY_UNIT_LIMITS.GB === numBytes ||
+    numBytes < MEMORY_UNIT_LIMITS.TB
+  ) {
+    return [numBytes / MEMORY_UNIT_LIMITS.GB, 'GB'];
+  } else if (
+    MEMORY_UNIT_LIMITS.TB === numBytes ||
+    numBytes < MEMORY_UNIT_LIMITS.PB
+  ) {
+    return [numBytes / MEMORY_UNIT_LIMITS.TB, 'TB'];
+  } else {
+    return [numBytes / MEMORY_UNIT_LIMITS.PB, 'PB'];
+  }
+}

packages/labextension/src/handler.ts

@@ -0,0 +1,44 @@
+import { URLExt } from '@jupyterlab/coreutils';
+import { ServerConnection } from '@jupyterlab/services';
+
+/**
+ * Call the API extension
+ *
+ * @param endPoint API REST end point for the extension
+ * @param init Initial values for the request
+ * @returns The response body interpreted as JSON
+ */
+export async function requestAPI<T>(
+  endPoint = '',
+  init: RequestInit = {}
+): Promise<T> {
+  const settings = ServerConnection.makeSettings();
+  const requestUrl = URLExt.join(
+    settings.baseUrl,
+    'api/metrics/v1/kernel_usage', // API Namespace
+    endPoint
+  );
+
+  let response: Response;
+  try {
+    response = await ServerConnection.makeRequest(requestUrl, init, settings);
+  } catch (error) {
+    throw new ServerConnection.NetworkError(error as any);
+  }
+
+  let data: any = await response.text();
+
+  if (data.length > 0) {
+    try {
+      data = JSON.parse(data);
+    } catch (error) {
+      console.log('Not a JSON response body.', response);
+    }
+  }
+
+  if (!response.ok) {
+    throw new ServerConnection.ResponseError(response, data.message || data);
+  }
+
+  return data;
+}