Skip to content

Commit

Permalink
improve ipywidgets performance by synchronizing only output widgets s…
Browse files Browse the repository at this point in the history 
…tate changes
  • Loading branch information
@sdissegna-maystreet
sdissegna-maystreet committed Oct 18, 2021
1 parent 6e5df57 commit 545b139
Showing 1 changed file with 57 additions and 63 deletions.
120 changes: 57 additions & 63 deletions src/client/datascience/ipywidgets/ipyWidgetMessageDispatcher.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,23 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
private totalWaitTime: number = 0;
private totalWaitedMessages: number = 0;
private hookCount: number = 0;
private fullHandleMessage?: { id: string; promise: Deferred<void> };
/**
* This will be true if user has executed something that has resulted in the use of ipywidgets.
* We make this determinination based on whether we see messages coming from backend kernel of a specific shape.
* E.g. if it contains ipywidget mime type, then widgets are in use.
* The Output widget's model can set up or tear down a kernel message hook on state change.
* We need to wait until the kernel message hook has been connected before it's safe to send
* more messages to the UI kernel.
*
* To do this we:
* - Keep track of the id of all the Output widget models in the outputWidgetIds instance variable.
* We add/remove these ids by inspecting messages in onKernelSocketMessage.
* - When a state update message is sent to one of these widgets, we synchronize with the UI and
* stop sending messages until we receive a reply indicating that the state change has been fully handled.
* We keep track of the message we're waiting for in the fullHandleMessage instance variable.
* We start waiting for the state change to finish processing in onKernelSocketMessage,
* and we stop waiting in iopubMessageHandled.
*/
private isUsingIPyWidgets?: boolean;
private outputWidgetIds = new Set<string>();
private fullHandleMessage?: { id: string; promise: Deferred<void> };
private isUsingIPyWidgets = false;
private readonly deserialize: (data: string | ArrayBuffer) => KernelMessage.IMessage<KernelMessage.MessageType>;

constructor(private readonly kernelProvider: IKernelProvider, public readonly document: NotebookDocument) {
Expand Down Expand Up @@ -251,14 +261,12 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
// fully handled on both the UI and extension side before we process the next message incoming
private messageNeedsFullHandle(message: any) {
// We only get a handled callback for iopub messages, so this channel must be iopub
if (message.channel === 'iopub') {
if (message.header?.msg_type === 'comm_msg') {
// IOPub comm messages need to be fully handled
return true;
}
}

return false;
return (
message.channel === 'iopub' &&
message.header?.msg_type === 'comm_msg' &&
message.content?.data?.method === 'update' &&
this.outputWidgetIds.has(message.content?.comm_id)
);
}

// Callback from the UI kernel when an iopubMessage has been fully handled
Expand All @@ -272,46 +280,11 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
}
private async onKernelSocketMessage(data: WebSocketData): Promise<void> {
// Hooks expect serialized data as this normally comes from a WebSocket
let message: undefined | KernelMessage.ICommOpenMsg; // = this.deserialize(data as any) as any;
if (!this.isUsingIPyWidgets) {
// Lets deserialize only if we know we have a potential case
// where this message contains some data we're interested in.
let mustDeserialize = false;
if (typeof data === 'string') {
mustDeserialize = data.includes(WIDGET_MIMETYPE) || data.includes(Identifiers.DefaultCommTarget);
} else {
// Array buffers (non-plain text data) must be deserialized.
mustDeserialize = true;
}
if (!message && mustDeserialize) {
message = this.deserialize(data as any) as any;
}

// Check for hints that would indicate whether ipywidgest are used in outputs.
if (
message &&
message.content &&
message.content.data &&
(message.content.data[WIDGET_MIMETYPE] || message.content.target_name === Identifiers.DefaultCommTarget)
) {
this.isUsingIPyWidgets = true;
}
}

const msgUuid = uuid();
const promise = createDeferred<void>();
this.waitingMessageIds.set(msgUuid, { startTime: Date.now(), resultPromise: promise });

// Check if we need to fully handle this message on UI and Extension side before we move to the next
if (this.isUsingIPyWidgets) {
if (!message) {
message = this.deserialize(data as any) as any;
}
if (this.messageNeedsFullHandle(message)) {
this.fullHandleMessage = { id: message!.header.msg_id, promise: createDeferred<void>() };
}
}

if (typeof data === 'string') {
this.raisePostMessage(IPyWidgetMessages.IPyWidgets_msg, { id: msgUuid, data });
} else {
Expand All @@ -321,21 +294,42 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
});
}

// There are three handling states that we have for messages here
// 1. If we have not detected ipywidget usage at all, we just forward messages to the kernel
// 2. If we have detected ipywidget usage. We wait on our message to be received, but not
// possibly processed yet by the UI kernel. This make sure our ordering is in sync
// 3. For iopub comm messages we wait for them to be fully handled by the UI kernel
// and the Extension kernel as they may be required to do things like
// register message hooks on both sides before we process the nextExtension message

// If there are no ipywidgets thusfar in the notebook, then no need to synchronize messages.
if (this.isUsingIPyWidgets) {
await promise.promise;

// Comm specific iopub messages we need to wait until they are full handled
// by both the UI and extension side before we move forward
if (this.fullHandleMessage) {
// Lets deserialize only if we know we have a potential case
// where this message contains some data we're interested in.
const mustDeserialize =
typeof data !== 'string' ||
data.includes(WIDGET_MIMETYPE) ||
data.includes(Identifiers.DefaultCommTarget) ||
data.includes('comm_open') ||
data.includes('comm_close') ||
data.includes('comm_msg');
if (mustDeserialize) {
const message = this.deserialize(data as any) as any;

// Check for hints that would indicate whether ipywidgest are used in outputs.
if (
message &&
message.content &&
message.content.data &&
(message.content.data[WIDGET_MIMETYPE] || message.content.target_name === Identifiers.DefaultCommTarget)
) {
this.isUsingIPyWidgets = true;
}

const isIPYWidgetOutputModelOpen =
message.header?.msg_type === 'comm_open' &&
message.content?.data?.state?._model_module === '@jupyter-widgets/output' &&
message.content?.data?.state?._model_name === 'OutputModel';
const isIPYWidgetOutputModelClose =
message.header?.msg_type === 'comm_close' && this.outputWidgetIds.has(message.content?.comm_id);

if (isIPYWidgetOutputModelOpen) {
this.outputWidgetIds.add(message.content.comm_id);
} else if (isIPYWidgetOutputModelClose) {
this.outputWidgetIds.delete(message.content.comm_id);
} else if (this.messageNeedsFullHandle(message)) {
this.fullHandleMessage = { id: message.header.msg_id, promise: createDeferred<void>() };
await promise.promise;
await this.fullHandleMessage.promise.promise;
this.fullHandleMessage = undefined;
}
Expand Down

0 comments on commit 545b139

Please sign in to comment.