Skip to content

Commit

Permalink
Improved KernelWidgetManager creation and kernel restoration.
Browse files Browse the repository at this point in the history
  • Loading branch information
Alan Fleming committed Nov 3, 2024
1 parent 79d6a0c commit e96983c
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 95 deletions.
139 changes: 69 additions & 70 deletions python/jupyterlab_widgets/src/manager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,15 +66,7 @@ export const WIDGET_STATE_MIMETYPE =
/**
* A widget manager that returns Lumino widgets.
*/
export abstract class LabWidgetManager
extends ManagerBase
implements IDisposable
{
constructor(rendermime: IRenderMimeRegistry) {
super();
this._rendermime = rendermime;
}

abstract class LabWidgetManager extends ManagerBase implements IDisposable {
/**
* Default callback handler to emit unhandled kernel messages.
*/
Expand Down Expand Up @@ -181,7 +173,6 @@ export abstract class LabWidgetManager
return;
}
this._isDisposed = true;
this._rendermime = null!;

if (this._commRegistration) {
this._commRegistration.dispose();
Expand Down Expand Up @@ -245,10 +236,6 @@ export abstract class LabWidgetManager

abstract get kernel(): Kernel.IKernelConnection | null;

get rendermime(): IRenderMimeRegistry {
return this._rendermime;
}

/**
* A signal emitted when state is restored to the widget manager.
*
Expand Down Expand Up @@ -331,15 +318,13 @@ export abstract class LabWidgetManager
await this.handle_comm_open(oldComm, msg);
};

static globalRendermime: IRenderMimeRegistry;
static rendermime: IRenderMimeRegistry;

protected _restored = new Signal<this, void>(this);
protected _restoredStatus = false;
protected _kernelRestoreInProgress = false;

private _isDisposed = false;
private _registry: SemVerCache<ExportData> = new SemVerCache<ExportData>();
private _rendermime: IRenderMimeRegistry;

private _commRegistration: IDisposable;

Expand All @@ -352,46 +337,32 @@ export abstract class LabWidgetManager
}

/**
* A singleton widget manager per kernel for the lifecycle of the kernel.
* The factory of the rendermime will be updated to use the widgetManager
* directly if it isn't the globalRendermime.
*
* Note: The rendermime of the instance is always the global rendermime.
* KernelWidgetManager is singleton widget manager per kernel.id.
* This class should not be created directly or subclassed, instead use
* the class method `KernelWidgetManager.getManager(kernel)`. See also
* `KernelWidgetManager.getManagerById`.
*/
export class KernelWidgetManager extends LabWidgetManager {
constructor(
kernel: Kernel.IKernelConnection,
rendermime?: IRenderMimeRegistry,
pendingManagerMessage = 'Loading widget ...'
) {
const instance = Private.managers.get(kernel.id);
if (instance) {
instance._useKernel(kernel);
KernelWidgetManager.configureRendermime(
rendermime,
instance,
pendingManagerMessage
);
return instance;
constructor(kernel: Kernel.IKernelConnection) {
if (Private.managers.has(kernel.id)) {
throw new Error('A manager already exists!');
}
if (!kernel.handleComms) {
throw new Error('Kernel does not have handleComms enabled');
}
super(LabWidgetManager.globalRendermime);
super();
Private.managers.set(kernel.id, this);
this.loadCustomWidgetDefinitions();
LabWidgetManager.WIDGET_REGISTRY.changed.connect(() =>
this.loadCustomWidgetDefinitions()
);
this._useKernel(kernel);
KernelWidgetManager.configureRendermime(
rendermime,
this,
pendingManagerMessage
);
this._updateKernel(kernel);
}

_useKernel(this: KernelWidgetManager, kernel: Kernel.IKernelConnection) {
private _updateKernel(
this: KernelWidgetManager,
kernel: Kernel.IKernelConnection
) {
if (!kernel.handleComms || this._kernel === kernel) {
return;
}
Expand All @@ -409,13 +380,15 @@ export class KernelWidgetManager extends LabWidgetManager {
this._handleKernelConnectionStatusChange,
this
);
this._kernel.disposed.disconnect(this._onKernelDisposed, this);
}
this._kernel = kernel;
this._kernel.statusChanged.connect(this._handleKernelStatusChange, this);
this._kernel.connectionStatusChanged.connect(
kernel.statusChanged.connect(this._handleKernelStatusChange, this);
kernel.connectionStatusChanged.connect(
this._handleKernelConnectionStatusChange,
this
);
kernel.disposed.connect(this._onKernelDisposed, this);
this.restoreWidgets();
}

Expand All @@ -437,7 +410,7 @@ export class KernelWidgetManager extends LabWidgetManager {
manager?: KernelWidgetManager,
pendingManagerMessage = ''
) {
if (!rendermime || rendermime === LabWidgetManager.globalRendermime) {
if (!rendermime || rendermime === LabWidgetManager.rendermime) {
return;
}
rendermime.removeMimeType(WIDGET_VIEW_MIMETYPE);
Expand All @@ -457,27 +430,18 @@ export class KernelWidgetManager extends LabWidgetManager {
): void {
switch (status) {
case 'connected':
// Only restore if we aren't currently trying to restore from the kernel
// (for example, in our initial restore from the constructor).
if (!this._kernelRestoreInProgress) {
this.restoreWidgets();
}
this.restoreWidgets();
break;
case 'disconnected':
this.disconnect();
break;
}
}

static existsWithActiveKenel(id: string) {
const widgetManager = Private.managers.get(id);
return widgetManager?._restoredStatus;
}

/**
* Get the KernelWidgetManager that owns the model.
* Find the KernelWidgetManager that owns the model.
*/
static async getManager(
static async findManager(
model_id: string,
delays = [100, 1000]
): Promise<KernelWidgetManager> {
Expand All @@ -494,6 +458,36 @@ export class KernelWidgetManager extends LabWidgetManager {
);
}

/**
* The correct way to get a KernelWidgetManager
* @param kernel IKernelConnection
* @returns
*/
static async getManager(
kernel: Kernel.IKernelConnection
): Promise<KernelWidgetManager> {
let manager = Private.managers.get(kernel.id);
if (!manager) {
manager = new KernelWidgetManager(kernel);
}
if (kernel.handleComms) {
manager._updateKernel(kernel);
if (!manager.restoredStatus) {
const restored = manager.restored;
await new Promise((resolve) => restored.connect(resolve));
}
}
return manager;
}

/**
* Get a KernelWidgetManager by kernel.id
* @param id The kernel id
*/
static getManagerById(id: string): KernelWidgetManager | undefined {
return Private.managers.get(id);
}

_handleKernelStatusChange(
sender: Kernel.IKernelConnection,
status: Kernel.Status
Expand All @@ -506,23 +500,32 @@ export class KernelWidgetManager extends LabWidgetManager {
}
}

_onKernelDisposed() {
const model = this._kernel.model;
const kernel = KernelWidgetManager.kernels.connectTo({ model });
this._updateKernel(kernel);
}

/**
* Restore widgets from kernel.
*/
async restoreWidgets(): Promise<void> {
if (this._kernelRestoreInProgress) {
return;
}
this._restoredStatus = false;
this._kernelRestoreInProgress = true;
try {
await this.clear_state();
await this._loadFromKernel();
} catch (err) {
// Do nothing
} finally {
this._restoredStatus = true;
this._kernelRestoreInProgress = false;
this.triggerRestored();
}
this.triggerRestored();
}

triggerRestored() {
this._restoredStatus = true;
this._restored.emit();
}
/**
Expand All @@ -533,7 +536,6 @@ export class KernelWidgetManager extends LabWidgetManager {
return;
}
super.dispose();
KernelWidgetManager.configureRendermime(this.rendermime);
Private.managers.delete(this.kernel.id);
this._handleKernelChanged({
name: 'kernel',
Expand All @@ -557,7 +559,7 @@ export class KernelWidgetManager extends LabWidgetManager {
filterModelState(serialized_state: any): any {
return this.filterExistingModelState(serialized_state);
}

static kernels: Kernel.IManager;
private _kernel: Kernel.IKernelConnection;
protected _kernelRestoreInProgress = false;
}
Expand Down Expand Up @@ -630,7 +632,7 @@ export class WidgetManager extends Backbone.Model implements IDisposable {
rendermime: IRenderMimeRegistry,
manager: WidgetManager
) {
if (rendermime === LabWidgetManager.globalRendermime) {
if (rendermime === LabWidgetManager.rendermime) {
return;
}
rendermime.removeMimeType(WIDGET_VIEW_MIMETYPE);
Expand All @@ -648,7 +650,7 @@ export class WidgetManager extends Backbone.Model implements IDisposable {
let wManager: KernelWidgetManager | undefined;
await this.context.sessionContext.ready;
if (this.kernel) {
wManager = new KernelWidgetManager(this.kernel);
wManager = await KernelWidgetManager.getManager(this.kernel);
}
if (wManager === this._widgetManager) {
return;
Expand Down Expand Up @@ -796,7 +798,6 @@ export class WidgetManager extends Backbone.Model implements IDisposable {
this._renderers = null!;
this._context = null!;
this._context = null!;
this._rendermime = null!;
this._settings = null!;
}

Expand Down Expand Up @@ -831,8 +832,6 @@ export class WidgetManager extends Backbone.Model implements IDisposable {
}
}
static loggerRegistry: ILoggerRegistry | null;
// protected _restored = new Signal<this, void>(this);
// protected _restoredStatus = false;
private _isDisposed = false;
private _context: DocumentRegistry.Context;
private _rendermime: IRenderMimeRegistry;
Expand Down
6 changes: 2 additions & 4 deletions python/jupyterlab_widgets/src/output.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import { Panel } from '@lumino/widgets';

import { IRenderMimeRegistry } from '@jupyterlab/rendermime';

import { LabWidgetManager } from './manager';
import { KernelWidgetManager } from './manager';

import { OutputAreaModel, OutputArea } from '@jupyterlab/outputarea';

Expand Down Expand Up @@ -44,7 +44,7 @@ export class OutputModel extends outputBase.OutputModel {
* Reset the message id.
*/
reset_msg_id(): void {
const kernel = this.widget_manager.kernel;
const kernel = (this.widget_manager as KernelWidgetManager).kernel;
const msgId = this.get('msg_id');
const oldMsgId = this.previous('msg_id');

Expand Down Expand Up @@ -98,8 +98,6 @@ export class OutputModel extends outputBase.OutputModel {
}
}

widget_manager: LabWidgetManager;

private _msgHook: (msg: KernelMessage.IIOPubMessage) => boolean;
private _outputs: OutputAreaModel;
static rendermime: IRenderMimeRegistry;
Expand Down
23 changes: 3 additions & 20 deletions python/jupyterlab_widgets/src/plugin.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@ import { WidgetRenderer } from './renderer';

import {
KernelWidgetManager,
LabWidgetManager,
WIDGET_VIEW_MIMETYPE,
WidgetManager,
} from './manager';
Expand Down Expand Up @@ -103,23 +102,7 @@ function activateWidgetExtension(
): base.IJupyterWidgetRegistry {
const { commands } = app;
const trans = (translator ?? nullTranslator).load('jupyterlab_widgets');

app.serviceManager.kernels.runningChanged.connect((models) => {
for (const model of models.running()) {
if (
model &&
model.name === 'python3' &&
model.execution_state !== 'starting' &&
!KernelWidgetManager.existsWithActiveKenel(model.id)
) {
const kernel = app.serviceManager.kernels.connectTo({ model: model });
if (kernel.handleComms) {
new KernelWidgetManager(kernel);
}
}
}
});

KernelWidgetManager.kernels = app.serviceManager.kernels;
if (settingRegistry !== null) {
settingRegistry
.load(managerPlugin.id)
Expand All @@ -132,7 +115,7 @@ function activateWidgetExtension(
});
}
WidgetManager.loggerRegistry = loggerRegistry;
LabWidgetManager.globalRendermime = rendermime;
KernelWidgetManager.rendermime = rendermime;
// Add a default widget renderer.
rendermime.addFactory(
{
Expand Down Expand Up @@ -175,7 +158,7 @@ function activateWidgetExtension(

return {
registerWidget(data: base.IWidgetRegistryData): void {
LabWidgetManager.WIDGET_REGISTRY.push(data);
KernelWidgetManager.WIDGET_REGISTRY.push(data);
},
};
}
Expand Down
2 changes: 1 addition & 1 deletion python/jupyterlab_widgets/src/renderer.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ export class WidgetRenderer
}
if (!this._pendingManagerMessage && !this._managerIsSet) {
try {
this.manager = await KernelWidgetManager.getManager(source.model_id);
this.manager = await KernelWidgetManager.findManager(source.model_id);
} catch {
this.node.textContent = `KernelWidgetManager not found for model: ${model.data['text/plain']}`;
return;
Expand Down

0 comments on commit e96983c

Please sign in to comment.