Skip to content

Commit 6703e02

Browse files
davidbrochartdavidbrochart
committed
1 parent 657b727 commit 6703e02

File tree

3 files changed

+98
-57
lines changed

3 files changed

+98
-57
lines changed

‎jupyter_resource_usage/api.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import zmq
77
from jupyter_client.jsonutil import date_default
88
from jupyter_server.base.handlers import APIHandler
9+
from jupyter_server.utils import url_path_join
910
from packaging import version
1011
from tornado import web
1112
from tornado.concurrent import run_on_executor
@@ -19,6 +20,8 @@
1920

2021
USAGE_IS_SUPPORTED = version.parse("6.9.0") <= version.parse(ipykernel.__version__)
2122

23+
MAX_RETRIES = 3
24+
2225

2326
class ApiHandler(APIHandler):
2427
executor = ThreadPoolExecutor(max_workers=5)
@@ -104,9 +107,8 @@ async def get(self, matched_part=None, *args, **kwargs):
104107
poller = zmq.Poller()
105108
control_socket = control_channel.socket
106109
poller.register(control_socket, zmq.POLLIN)
107-
while True:
108-
timeout = 100
109-
timeout_ms = int(1000 * timeout)
110+
for i in range(1, MAX_RETRIES + 1):
111+
timeout_ms = 1000 * i
110112
events = dict(poller.poll(timeout_ms))
111113
if not events:
112114
self.write(json.dumps({}))

‎packages/labextension/src/index.ts

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import {
55
import { INotebookTracker } from '@jupyterlab/notebook';
66
import { LabIcon } from '@jupyterlab/ui-components';
77
import { ICommandPalette } from '@jupyterlab/apputils';
8-
import { ILauncher } from '@jupyterlab/launcher';
98
import { KernelUsagePanel } from './panel';
109
import tachometer from '../style/tachometer.svg';
1110

@@ -26,14 +25,12 @@ const extension: JupyterFrontEndPlugin<void> = {
2625
id: '@jupyter-server/resource-usage:memory-status-item',
2726
autoStart: true,
2827
requires: [IStatusBar, ITranslator, ICommandPalette, INotebookTracker],
29-
optional: [ILauncher],
3028
activate: (
3129
app: JupyterFrontEnd,
3230
statusBar: IStatusBar,
3331
translator: ITranslator,
3432
palette: ICommandPalette,
35-
notebookTracker: INotebookTracker,
36-
launcher: ILauncher | null
33+
notebookTracker: INotebookTracker
3734
) => {
3835
const item = new MemoryUsage(translator);
3936

@@ -48,13 +45,16 @@ const extension: JupyterFrontEndPlugin<void> = {
4845
const { commands, shell } = app;
4946
const category = 'Kernel Resource';
5047

51-
async function createPanel(): Promise<KernelUsagePanel> {
52-
const panel = new KernelUsagePanel({
53-
widgetAdded: notebookTracker.widgetAdded,
54-
currentNotebookChanged: notebookTracker.currentChanged,
55-
});
56-
shell.add(panel, 'right', { rank: 200 });
57-
return panel;
48+
let panel: KernelUsagePanel | null = null;
49+
50+
function createPanel() {
51+
if (!panel || panel.isDisposed) {
52+
panel = new KernelUsagePanel({
53+
widgetAdded: notebookTracker.widgetAdded,
54+
currentNotebookChanged: notebookTracker.currentChanged
55+
});
56+
shell.add(panel, 'right', { rank: 200 });
57+
}
5858
}
5959

6060
commands.addCommand(CommandIDs.getKernelUsage, {

‎packages/labextension/src/widget.tsx

Lines changed: 82 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import React, { useState } from 'react';
1+
import React, { useState, useEffect } from 'react';
22
import { ISignal } from '@lumino/signaling';
33
import { ReactWidget, ISessionContext } from '@jupyterlab/apputils';
44
import { IChangedArgs } from '@jupyterlab/coreutils';
@@ -32,6 +32,19 @@ type Usage = {
3232

3333
const POLL_INTERVAL_SEC = 5;
3434

35+
type KernelChangeCallback = (
36+
_sender: ISessionContext,
37+
args: IChangedArgs<
38+
Kernel.IKernelConnection | null,
39+
Kernel.IKernelConnection | null,
40+
'kernel'
41+
>
42+
) => void;
43+
let kernelChangeCallback: {
44+
callback: KernelChangeCallback;
45+
panel: NotebookPanel;
46+
} | null = null;
47+
3548
const KernelUsage = (props: {
3649
widgetAdded: ISignal<INotebookTracker, NotebookPanel | null>;
3750
currentNotebookChanged: ISignal<INotebookTracker, NotebookPanel | null>;
@@ -44,61 +57,87 @@ const KernelUsage = (props: {
4457

4558
useInterval(async () => {
4659
if (kernelId && panel.isVisible) {
47-
requestUsage(kernelId).then((usage) => setUsage(usage));
60+
requestUsage(kernelId)
61+
.then(usage => setUsage(usage))
62+
.catch(() => {
63+
console.warn(`Request failed for ${kernelId}. Kernel restarting?`);
64+
});
4865
}
4966
}, POLL_INTERVAL_SEC * 1000);
5067

51-
const requestUsage = (kid: string) =>
52-
requestAPI<any>(`get_usage/${kid}`).then((data) => {
68+
const requestUsage = (kid: string) => {
69+
return requestAPI<any>(`get_usage/${kid}`).then(data => {
5370
const usage: Usage = {
5471
...data.content,
5572
kernelId: kid,
56-
timestamp: new Date(),
73+
timestamp: new Date()
5774
};
5875
return usage;
5976
});
77+
};
6078

61-
props.currentNotebookChanged.connect(
62-
(sender: INotebookTracker, panel: NotebookPanel | null) => {
63-
panel?.sessionContext.kernelChanged.connect(
64-
(
65-
_sender: ISessionContext,
66-
args: IChangedArgs<
67-
Kernel.IKernelConnection | null,
68-
Kernel.IKernelConnection | null,
69-
'kernel'
70-
>
71-
) => {
72-
/*
73-
const oldKernelId = args.oldValue?.id;
74-
if (oldKernelId) {
75-
const poll = kernelPools.get(oldKernelId);
76-
poll?.poll.dispose();
77-
kernelPools.delete(oldKernelId);
78-
}
79-
*/
80-
const newKernelId = args.newValue?.id;
81-
if (newKernelId) {
82-
setKernelId(newKernelId);
83-
const path = panel?.sessionContext.session?.model.path;
84-
setPath(path);
85-
requestUsage(newKernelId).then((usage) => setUsage(usage));
86-
}
79+
useEffect(() => {
80+
const createKernelChangeCallback = (panel: NotebookPanel) => {
81+
return (
82+
_sender: ISessionContext,
83+
args: IChangedArgs<
84+
Kernel.IKernelConnection | null,
85+
Kernel.IKernelConnection | null,
86+
'kernel'
87+
>
88+
) => {
89+
const newKernelId = args.newValue?.id;
90+
if (newKernelId) {
91+
setKernelId(newKernelId);
92+
const path = panel?.sessionContext.session?.model.path;
93+
setPath(path);
94+
requestUsage(newKernelId).then(usage => setUsage(usage));
95+
} else {
96+
// Kernel was disposed
97+
setKernelId(newKernelId);
8798
}
88-
);
89-
if (panel?.sessionContext.session?.id !== kernelId) {
90-
if (panel?.sessionContext.session?.kernel?.id) {
91-
const kernelId = panel?.sessionContext.session?.kernel?.id;
92-
if (kernelId) {
93-
setKernelId(kernelId);
94-
const path = panel?.sessionContext.session?.model.path;
95-
setPath(path);
96-
requestUsage(kernelId).then((usage) => setUsage(usage));
97-
}
99+
};
100+
};
101+
102+
const notebookChangeCallback = (
103+
sender: INotebookTracker,
104+
panel: NotebookPanel | null
105+
) => {
106+
if (panel === null) {
107+
// Ideally we would switch to a new "select a notebook to get kernel
108+
// usage" screen instead of showing outdated info.
109+
return;
110+
}
111+
if (kernelChangeCallback) {
112+
kernelChangeCallback.panel.sessionContext.kernelChanged.disconnect(
113+
kernelChangeCallback.callback
114+
);
115+
}
116+
kernelChangeCallback = {
117+
callback: createKernelChangeCallback(panel),
118+
panel
119+
};
120+
panel.sessionContext.kernelChanged.connect(kernelChangeCallback.callback);
121+
122+
if (panel.sessionContext.session?.kernel?.id !== kernelId) {
123+
const kernelId = panel.sessionContext.session?.kernel?.id;
124+
if (kernelId) {
125+
setKernelId(kernelId);
126+
const path = panel.sessionContext.session?.model.path;
127+
setPath(path);
128+
requestUsage(kernelId).then(usage => setUsage(usage));
98129
}
99130
}
100-
}
101-
);
131+
};
132+
props.currentNotebookChanged.connect(notebookChangeCallback);
133+
return () => {
134+
props.currentNotebookChanged.disconnect(notebookChangeCallback);
135+
// In the ideal world we would disconnect kernelChangeCallback from
136+
// last panel here, but this can lead to a race condition. Instead,
137+
// we make sure there is ever only one callback active by holding
138+
// it in a global state.
139+
};
140+
}, [kernelId]);
102141

103142
if (kernelId) {
104143
if (usage) {

0 commit comments

Comments
 (0)