Skip to content

Commit b63f052

Browse files
committed
Merge commit '22f4c31825d52543a0fd19020aad81d94d4bf842' into jupyterlab-kernel-usage
2 parents 0bc37bf + 22f4c31 commit b63f052

File tree

4 files changed

+97
-61
lines changed

4 files changed

+97
-61
lines changed

jupyterlab-kernel-usage/jupyterlab_kernel_usage/__init__.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
21
import json
32
from pathlib import Path
43

jupyterlab-kernel-usage/jupyterlab_kernel_usage/handlers.py

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,19 @@
1-
import asyncio
21
import ipykernel
32
import json
43
import tornado
54
import zmq
65

7-
from functools import partial
8-
96
from jupyter_server.base.handlers import APIHandler
10-
from jupyter_server.utils import url_path_join, ensure_async
7+
from jupyter_server.utils import url_path_join
118
from jupyter_client.jsonutil import date_default
129

1310
from packaging import version
1411

1512

1613
USAGE_IS_SUPPORTED = version.parse("6.9.0") <= version.parse(ipykernel.__version__)
1714

15+
MAX_RETRIES = 3
16+
1817

1918
class RouteHandler(APIHandler):
2019

@@ -38,9 +37,8 @@ async def get(self, matched_part=None, *args, **kwargs):
3837
poller = zmq.Poller()
3938
control_socket = control_channel.socket
4039
poller.register(control_socket, zmq.POLLIN)
41-
while True:
42-
timeout = 100
43-
timeout_ms = int(1000 * timeout)
40+
for i in range(1, MAX_RETRIES + 1):
41+
timeout_ms = 1000 * i
4442
events = dict(poller.poll(timeout_ms))
4543
if not events:
4644
self.write(json.dumps({}))

jupyterlab-kernel-usage/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

@@ -19,24 +18,25 @@ namespace CommandIDs {
1918
const plugin: JupyterFrontEndPlugin<void> = {
2019
id: 'kernelusage:plugin',
2120
requires: [ICommandPalette, INotebookTracker],
22-
optional: [ILauncher],
2321
autoStart: true,
2422
activate: (
2523
app: JupyterFrontEnd,
2624
palette: ICommandPalette,
27-
notebookTracker: INotebookTracker,
28-
launcher: ILauncher | null
25+
notebookTracker: INotebookTracker
2926
) => {
3027
const { commands, shell } = app;
3128
const category = 'Kernel Resource';
3229

33-
async function createPanel(): Promise<KernelUsagePanel> {
34-
const panel = new KernelUsagePanel({
35-
widgetAdded: notebookTracker.widgetAdded,
36-
currentNotebookChanged: notebookTracker.currentChanged
37-
});
38-
shell.add(panel, 'right', { rank: 200 });
39-
return panel;
30+
let panel: KernelUsagePanel | null = null;
31+
32+
function createPanel() {
33+
if (!panel || panel.isDisposed) {
34+
panel = new KernelUsagePanel({
35+
widgetAdded: notebookTracker.widgetAdded,
36+
currentNotebookChanged: notebookTracker.currentChanged
37+
});
38+
shell.add(panel, 'right', { rank: 200 });
39+
}
4040
}
4141

4242
commands.addCommand(CommandIDs.getKernelUsage, {

jupyterlab-kernel-usage/src/widget.tsx

Lines changed: 81 additions & 42 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,
5673
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)