Skip to content

Commit a8a8111

Browse files
authored
Increase checkpoint poll interval (#7711)
1 parent 5a845dc commit a8a8111

File tree

2 files changed

+98
-18
lines changed

2 files changed

+98
-18
lines changed

packages/notebook-extension/schema/checkpoints.json

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,15 @@
44
"jupyter.lab.toolbars": {
55
"TopBar": [{ "name": "checkpoint", "rank": 20 }]
66
},
7-
"properties": {},
7+
"properties": {
8+
"checkpointPollingInterval": {
9+
"type": "number",
10+
"title": "Checkpoint Polling Interval (seconds)",
11+
"description": "How often to check for checkpoints (in seconds). Set to 0 to disable polling.",
12+
"default": 30,
13+
"minimum": 0
14+
}
15+
},
816
"additionalProperties": false,
917
"type": "object"
1018
}

packages/notebook-extension/src/index.ts

Lines changed: 89 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ import { PageConfig, Text, Time, URLExt } from '@jupyterlab/coreutils';
1919

2020
import { IDocumentManager } from '@jupyterlab/docmanager';
2121

22+
import { DocumentRegistry } from '@jupyterlab/docregistry';
23+
2224
import { IMainMenu } from '@jupyterlab/mainmenu';
2325

2426
import {
@@ -92,13 +94,14 @@ const checkpoints: JupyterFrontEndPlugin<void> = {
9294
description: 'A plugin for the checkpoint indicator.',
9395
autoStart: true,
9496
requires: [IDocumentManager, ITranslator],
95-
optional: [INotebookShell, IToolbarWidgetRegistry],
97+
optional: [INotebookShell, IToolbarWidgetRegistry, ISettingRegistry],
9698
activate: (
9799
app: JupyterFrontEnd,
98100
docManager: IDocumentManager,
99101
translator: ITranslator,
100102
notebookShell: INotebookShell | null,
101-
toolbarRegistry: IToolbarWidgetRegistry | null
103+
toolbarRegistry: IToolbarWidgetRegistry | null,
104+
settingRegistry: ISettingRegistry | null
102105
) => {
103106
const { shell } = app;
104107
const trans = translator.load('notebook');
@@ -113,18 +116,26 @@ const checkpoints: JupyterFrontEndPlugin<void> = {
113116
});
114117
}
115118

116-
const onChange = async () => {
119+
const getCurrent = () => {
117120
const current = shell.currentWidget;
118121
if (!current) {
119-
return;
122+
return null;
120123
}
121124
const context = docManager.contextForWidget(current);
125+
if (!context) {
126+
return null;
127+
}
128+
return context;
129+
};
122130

123-
context?.fileChanged.disconnect(onChange);
124-
context?.fileChanged.connect(onChange);
125-
126-
const checkpoints = await context?.listCheckpoints();
131+
const updateCheckpointDisplay = async () => {
132+
const current = getCurrent();
133+
if (!current) {
134+
return;
135+
}
136+
const checkpoints = await current.listCheckpoints();
127137
if (!checkpoints || !checkpoints.length) {
138+
node.textContent = '';
128139
return;
129140
}
130141
const checkpoint = checkpoints[checkpoints.length - 1];
@@ -134,19 +145,80 @@ const checkpoints: JupyterFrontEndPlugin<void> = {
134145
);
135146
};
136147

148+
const onSaveState = async (
149+
sender: DocumentRegistry.IContext<DocumentRegistry.IModel>,
150+
state: DocumentRegistry.SaveState
151+
) => {
152+
if (state !== 'completed') {
153+
return;
154+
}
155+
// Add a small artificial delay so that the UI can pick up the newly created checkpoint.
156+
// Since the save state signal is emitted after a file save, but not after a checkpoint has been created.
157+
setTimeout(() => {
158+
void updateCheckpointDisplay();
159+
}, 500);
160+
};
161+
162+
const onChange = async () => {
163+
const context = getCurrent();
164+
if (!context) {
165+
return;
166+
}
167+
168+
context.saveState.disconnect(onSaveState);
169+
context.saveState.connect(onSaveState);
170+
171+
await updateCheckpointDisplay();
172+
};
173+
137174
if (notebookShell) {
138175
notebookShell.currentChanged.connect(onChange);
139176
}
140177

141-
new Poll({
142-
auto: true,
143-
factory: () => onChange(),
144-
frequency: {
145-
interval: 2000,
146-
backoff: false,
147-
},
148-
standby: 'when-hidden',
149-
});
178+
let checkpointPollingInterval = 30; // Default 30 seconds
179+
let poll: Poll | null = null;
180+
181+
const createPoll = () => {
182+
if (poll) {
183+
poll.dispose();
184+
}
185+
if (checkpointPollingInterval > 0) {
186+
poll = new Poll({
187+
auto: true,
188+
factory: () => updateCheckpointDisplay(),
189+
frequency: {
190+
interval: checkpointPollingInterval * 1000,
191+
backoff: false,
192+
},
193+
standby: 'when-hidden',
194+
});
195+
}
196+
};
197+
198+
const updateSettings = (settings: ISettingRegistry.ISettings): void => {
199+
checkpointPollingInterval = settings.get('checkpointPollingInterval')
200+
.composite as number;
201+
createPoll();
202+
};
203+
204+
if (settingRegistry) {
205+
const loadSettings = settingRegistry.load(checkpoints.id);
206+
Promise.all([loadSettings, app.restored])
207+
.then(([settings]) => {
208+
updateSettings(settings);
209+
settings.changed.connect(updateSettings);
210+
})
211+
.catch((reason: Error) => {
212+
console.error(
213+
`Failed to load settings for ${checkpoints.id}: ${reason.message}`
214+
);
215+
// Fall back to creating poll with default settings
216+
createPoll();
217+
});
218+
} else {
219+
// Create poll with default settings
220+
createPoll();
221+
}
150222
},
151223
};
152224

0 commit comments

Comments
 (0)