Merge jupyterlab-kernel-usage

Author: davidbrochart

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

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

@@ -1,8 +1,12 @@
 import json
 from concurrent.futures import ThreadPoolExecutor
 
+import ipykernel
 import psutil
+import zmq
+from jupyter_client.jsonutil import date_default
 from jupyter_server.base.handlers import APIHandler
+from packaging import version
 from tornado import web
 from tornado.concurrent import run_on_executor
 
@@ -13,6 +17,9 @@
     from .utils import Callable
 
 
+USAGE_IS_SUPPORTED = version.parse("6.9.0") <= version.parse(ipykernel.__version__)
+
+
 class ApiHandler(APIHandler):
     executor = ThreadPoolExecutor(max_workers=5)
 
@@ -74,3 +81,38 @@ 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)
+        while True:
+            timeout = 100
+            timeout_ms = int(1000 * timeout)
+            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,15 @@ 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, "jupyterlab_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,23 +42,58 @@
     "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/launcher": "^3.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"
+  },
+  "jupyter-releaser": {
+    "hooks": {
+      "before-build-npm": [
+        "python -m pip install jupyterlab~=3.1",
+        "jlpm"
+      ]
+    }
   }
 }

packages/labextension/src/format.ts

@@ -0,0 +1,57 @@
+// Taken from https://github.com/jupyter-server/jupyter-resource-usage/blob/e6ec53fa69fdb6de8e878974bcff006310658408/packages/labextension/src/memoryUsage.tsx#L272
+
+type MemoryUnit = 'B' | 'KB' | 'MB' | 'GB' | 'TB' | 'PB';
+
+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).
+ * Taken from https://github.com/jupyter-server/jupyter-resource-usage/blob/e6ec53fa69fdb6de8e878974bcff006310658408/packages/labextension/src/memoryUsage.tsx#L272
+ */
+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,
+    'jupyterlab_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;
+}

packages/labextension/src/index.ts

@@ -2,24 +2,38 @@ import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin,
 } from '@jupyterlab/application';
+import { INotebookTracker } from '@jupyterlab/notebook';
+import { LabIcon } from '@jupyterlab/ui-components';
+import { ICommandPalette } from '@jupyterlab/apputils';
+import { ILauncher } from '@jupyterlab/launcher';
+import { KernelUsagePanel } from './panel';
+import tachometer from '../style/tachometer.svg';
 
 import { IStatusBar } from '@jupyterlab/statusbar';
 
 import { ITranslator } from '@jupyterlab/translation';
 
 import { MemoryUsage } from './memoryUsage';
 
+namespace CommandIDs {
+  export const getKernelUsage = 'kernel-usage:get';
+}
+
 /**
  * Initialization data for the jupyter-resource-usage extension.
  */
 const extension: JupyterFrontEndPlugin<void> = {
   id: '@jupyter-server/resource-usage:memory-status-item',
   autoStart: true,
-  requires: [IStatusBar, ITranslator],
+  requires: [IStatusBar, ITranslator, ICommandPalette, INotebookTracker],
+  optional: [ILauncher],
   activate: (
     app: JupyterFrontEnd,
     statusBar: IStatusBar,
-    translator: ITranslator
+    translator: ITranslator,
+    palette: ICommandPalette,
+    notebookTracker: INotebookTracker,
+    launcher: ILauncher | null
   ) => {
     const item = new MemoryUsage(translator);
 
@@ -30,6 +44,32 @@ const extension: JupyterFrontEndPlugin<void> = {
       isActive: () => item.model.metricsAvailable,
       activeStateChanged: item.model.stateChanged,
     });
+
+    const { commands, shell } = app;
+    const category = 'Kernel Resource';
+
+    async function createPanel(): Promise<KernelUsagePanel> {
+      const panel = new KernelUsagePanel({
+        widgetAdded: notebookTracker.widgetAdded,
+        currentNotebookChanged: notebookTracker.currentChanged,
+      });
+      shell.add(panel, 'right', { rank: 200 });
+      return panel;
+    }
+
+    commands.addCommand(CommandIDs.getKernelUsage, {
+      label: 'Kernel Usage',
+      caption: 'Kernel Usage',
+      icon: new LabIcon({
+        name: 'jupyterlab-kernel-usage:icon',
+        svgstr: tachometer,
+      }),
+      execute: createPanel,
+    });
+
+    palette.addItem({ command: CommandIDs.getKernelUsage, category });
+
+    createPanel();
   },
 };