Merge jupyterlab-kernel-usage

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.

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,13 @@
 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 jupyter_server.utils import url_path_join
+from packaging import version
 from tornado import web
 from tornado.concurrent import run_on_executor
 
@@ -13,6 +18,11 @@
     from .utils import Callable
 
 
+USAGE_IS_SUPPORTED = version.parse("6.9.0") <= version.parse(ipykernel.__version__)
+
+MAX_RETRIES = 3
+
+
 class ApiHandler(APIHandler):
     executor = ThreadPoolExecutor(max_workers=5)
 
@@ -74,3 +84,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,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,57 @@
     "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"
+  },
+  "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,35 @@ import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin,
 } from '@jupyterlab/application';
+import { INotebookTracker } from '@jupyterlab/notebook';
+import { LabIcon } from '@jupyterlab/ui-components';
+import { ICommandPalette } from '@jupyterlab/apputils';
+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],
   activate: (
     app: JupyterFrontEnd,
     statusBar: IStatusBar,
-    translator: ITranslator
+    translator: ITranslator,
+    palette: ICommandPalette,
+    notebookTracker: INotebookTracker
   ) => {
     const item = new MemoryUsage(translator);
 
@@ -30,6 +41,35 @@ const extension: JupyterFrontEndPlugin<void> = {
       isActive: () => item.model.metricsAvailable,
       activeStateChanged: item.model.stateChanged,
     });
+
+    const { commands, shell } = app;
+    const category = 'Kernel Resource';
+
+    let panel: KernelUsagePanel | null = null;
+
+    function createPanel() {
+      if (!panel || panel.isDisposed) {
+        panel = new KernelUsagePanel({
+          widgetAdded: notebookTracker.widgetAdded,
+          currentNotebookChanged: notebookTracker.currentChanged,
+        });
+        shell.add(panel, 'right', { rank: 200 });
+      }
+    }
+
+    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();
   },
 };