diff --git a/jupyter_resource_usage/api.py b/jupyter_resource_usage/api.py index 57c58eb..08a63cd 100644 --- a/jupyter_resource_usage/api.py +++ b/jupyter_resource_usage/api.py @@ -6,6 +6,7 @@ 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 @@ -19,6 +20,8 @@ USAGE_IS_SUPPORTED = version.parse("6.9.0") <= version.parse(ipykernel.__version__) +MAX_RETRIES = 3 + class ApiHandler(APIHandler): executor = ThreadPoolExecutor(max_workers=5) @@ -104,9 +107,8 @@ async def get(self, matched_part=None, *args, **kwargs): poller = zmq.Poller() control_socket = control_channel.socket poller.register(control_socket, zmq.POLLIN) - while True: - timeout = 100 - timeout_ms = int(1000 * timeout) + 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({})) diff --git a/packages/labextension/src/index.ts b/packages/labextension/src/index.ts index 8945f59..e204eb4 100644 --- a/packages/labextension/src/index.ts +++ b/packages/labextension/src/index.ts @@ -5,7 +5,6 @@ import { 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'; @@ -26,14 +25,12 @@ const extension: JupyterFrontEndPlugin = { id: '@jupyter-server/resource-usage:memory-status-item', autoStart: true, requires: [IStatusBar, ITranslator, ICommandPalette, INotebookTracker], - optional: [ILauncher], activate: ( app: JupyterFrontEnd, statusBar: IStatusBar, translator: ITranslator, palette: ICommandPalette, - notebookTracker: INotebookTracker, - launcher: ILauncher | null + notebookTracker: INotebookTracker ) => { const item = new MemoryUsage(translator); @@ -48,13 +45,16 @@ const extension: JupyterFrontEndPlugin = { const { commands, shell } = app; const category = 'Kernel Resource'; - async function createPanel(): Promise { - const panel = new KernelUsagePanel({ - widgetAdded: notebookTracker.widgetAdded, - currentNotebookChanged: notebookTracker.currentChanged, - }); - shell.add(panel, 'right', { rank: 200 }); - return panel; + 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, { diff --git a/packages/labextension/src/widget.tsx b/packages/labextension/src/widget.tsx index 3355d11..5704e59 100644 --- a/packages/labextension/src/widget.tsx +++ b/packages/labextension/src/widget.tsx @@ -1,4 +1,4 @@ -import React, { useState } from 'react'; +import React, { useState, useEffect } from 'react'; import { ISignal } from '@lumino/signaling'; import { ReactWidget, ISessionContext } from '@jupyterlab/apputils'; import { IChangedArgs } from '@jupyterlab/coreutils'; @@ -32,6 +32,19 @@ type Usage = { const POLL_INTERVAL_SEC = 5; +type KernelChangeCallback = ( + _sender: ISessionContext, + args: IChangedArgs< + Kernel.IKernelConnection | null, + Kernel.IKernelConnection | null, + 'kernel' + > +) => void; +let kernelChangeCallback: { + callback: KernelChangeCallback; + panel: NotebookPanel; +} | null = null; + const KernelUsage = (props: { widgetAdded: ISignal; currentNotebookChanged: ISignal; @@ -44,61 +57,87 @@ const KernelUsage = (props: { useInterval(async () => { if (kernelId && panel.isVisible) { - requestUsage(kernelId).then((usage) => setUsage(usage)); + requestUsage(kernelId) + .then(usage => setUsage(usage)) + .catch(() => { + console.warn(`Request failed for ${kernelId}. Kernel restarting?`); + }); } }, POLL_INTERVAL_SEC * 1000); - const requestUsage = (kid: string) => - requestAPI(`get_usage/${kid}`).then((data) => { + const requestUsage = (kid: string) => { + return requestAPI(`get_usage/${kid}`).then(data => { const usage: Usage = { ...data.content, kernelId: kid, - timestamp: new Date(), + timestamp: new Date() }; return usage; }); + }; - props.currentNotebookChanged.connect( - (sender: INotebookTracker, panel: NotebookPanel | null) => { - panel?.sessionContext.kernelChanged.connect( - ( - _sender: ISessionContext, - args: IChangedArgs< - Kernel.IKernelConnection | null, - Kernel.IKernelConnection | null, - 'kernel' - > - ) => { - /* - const oldKernelId = args.oldValue?.id; - if (oldKernelId) { - const poll = kernelPools.get(oldKernelId); - poll?.poll.dispose(); - kernelPools.delete(oldKernelId); - } - */ - const newKernelId = args.newValue?.id; - if (newKernelId) { - setKernelId(newKernelId); - const path = panel?.sessionContext.session?.model.path; - setPath(path); - requestUsage(newKernelId).then((usage) => setUsage(usage)); - } + useEffect(() => { + const createKernelChangeCallback = (panel: NotebookPanel) => { + return ( + _sender: ISessionContext, + args: IChangedArgs< + Kernel.IKernelConnection | null, + Kernel.IKernelConnection | null, + 'kernel' + > + ) => { + const newKernelId = args.newValue?.id; + if (newKernelId) { + setKernelId(newKernelId); + const path = panel?.sessionContext.session?.model.path; + setPath(path); + requestUsage(newKernelId).then(usage => setUsage(usage)); + } else { + // Kernel was disposed + setKernelId(newKernelId); } - ); - if (panel?.sessionContext.session?.id !== kernelId) { - if (panel?.sessionContext.session?.kernel?.id) { - const kernelId = panel?.sessionContext.session?.kernel?.id; - if (kernelId) { - setKernelId(kernelId); - const path = panel?.sessionContext.session?.model.path; - setPath(path); - requestUsage(kernelId).then((usage) => setUsage(usage)); - } + }; + }; + + const notebookChangeCallback = ( + sender: INotebookTracker, + panel: NotebookPanel | null + ) => { + if (panel === null) { + // Ideally we would switch to a new "select a notebook to get kernel + // usage" screen instead of showing outdated info. + return; + } + if (kernelChangeCallback) { + kernelChangeCallback.panel.sessionContext.kernelChanged.disconnect( + kernelChangeCallback.callback + ); + } + kernelChangeCallback = { + callback: createKernelChangeCallback(panel), + panel + }; + panel.sessionContext.kernelChanged.connect(kernelChangeCallback.callback); + + if (panel.sessionContext.session?.kernel?.id !== kernelId) { + const kernelId = panel.sessionContext.session?.kernel?.id; + if (kernelId) { + setKernelId(kernelId); + const path = panel.sessionContext.session?.model.path; + setPath(path); + requestUsage(kernelId).then(usage => setUsage(usage)); } } - } - ); + }; + props.currentNotebookChanged.connect(notebookChangeCallback); + return () => { + props.currentNotebookChanged.disconnect(notebookChangeCallback); + // In the ideal world we would disconnect kernelChangeCallback from + // last panel here, but this can lead to a race condition. Instead, + // we make sure there is ever only one callback active by holding + // it in a global state. + }; + }, [kernelId]); if (kernelId) { if (usage) {