Merge pull request #6864 from Madhu94/timing

Add execution timing to cells

packages/cells/src/widget.ts

@@ -1041,18 +1041,17 @@ export namespace CodeCell {
       model.outputs.clear();
       return;
     }
-
     let cellId = { cellId: model.id };
     metadata = {
       ...model.metadata.toJSON(),
       ...metadata,
       ...cellId
     };
+    const { recordTiming } = metadata;
     model.executionCount = null;
     cell.outputHidden = false;
     cell.setPrompt('*');
     model.trusted = true;
-
     let future: Kernel.IFuture<
       KernelMessage.IExecuteRequestMsg,
       KernelMessage.IExecuteReplyMsg
@@ -1064,10 +1063,49 @@ export namespace CodeCell {
         session,
         metadata
       );
+      // cell.outputArea.future assigned synchronously in `execute`
+      if (recordTiming) {
+        const recordTimingHook = (msg: KernelMessage.IIOPubMessage) => {
+          let label: string;
+          switch (msg.header.msg_type) {
+            case 'status':
+              label = `status.${
+                (msg as KernelMessage.IStatusMsg).content.execution_state
+              }`;
+              break;
+            case 'execute_input':
+              label = 'execute_input';
+              break;
+            default:
+              return;
+          }
+          const value = msg.header.date;
+          if (!value) {
+            return;
+          }
+          const timingInfo: any = model.metadata.get('execution') || {};
+          timingInfo[`iopub.${label}`] = value;
+          model.metadata.set('execution', timingInfo);
+          return true;
+        };
+        cell.outputArea.future.registerMessageHook(recordTimingHook);
+      }
       // Save this execution's future so we can compare in the catch below.
       future = cell.outputArea.future;
       const msg = await msgPromise;
       model.executionCount = msg.content.execution_count;
+      const started = msg.metadata.started as string;
+      if (recordTiming && started) {
+        const timingInfo = (model.metadata.get('execution') as any) || {};
+        if (started) {
+          timingInfo['shell.execute_reply.started'] = started;
+        }
+        const date = msg.header.date as string;
+        if (date) {
+          timingInfo['shell.execute_reply'] = date;
+        }
+        model.metadata.set('execution', timingInfo);
+      }
       return msg;
     } catch (e) {
       // If this is still the current execution, clear the prompt.

packages/notebook-extension/src/index.ts

@@ -173,6 +173,8 @@ namespace CommandIDs {
 
   export const toggleAllLines = 'notebook:toggle-all-cell-line-numbers';
 
+  export const toggleRecordTiming = 'notebook:toggle-record-timing';
+
   export const undoCellAction = 'notebook:undo-cell-action';
 
   export const redoCellAction = 'notebook:redo-cell-action';
@@ -1530,6 +1532,17 @@ function addCommands(
     },
     isEnabled
   });
+  commands.addCommand(CommandIDs.toggleRecordTiming, {
+    label: 'Toggle Recording Cell Timing',
+    execute: args => {
+      const current = getCurrent(args);
+
+      if (current) {
+        return NotebookActions.toggleRecordTiming(current.content);
+      }
+    },
+    isEnabled
+  });
   commands.addCommand(CommandIDs.commandMode, {
     label: 'Enter Command Mode',
     execute: args => {
@@ -1868,6 +1881,7 @@ function populatePalette(
     CommandIDs.deselectAll,
     CommandIDs.clearAllOutputs,
     CommandIDs.toggleAllLines,
+    CommandIDs.toggleRecordTiming,
     CommandIDs.editMode,
     CommandIDs.commandMode,
     CommandIDs.changeKernel,

packages/notebook/src/actions.tsx

@@ -980,6 +980,19 @@ export namespace NotebookActions {
     Private.handleState(notebook, state);
   }
 
+  /**
+   * Toggle whether to record cell timing execution.
+   *
+   * @param notebook - The target notebook widget.
+   */
+  export function toggleRecordTiming(notebook: Notebook): void {
+    if (!notebook.model) {
+      return;
+    }
+    const currentValue = notebook.model.metadata.get('record_timing') || false;
+    notebook.model.metadata.set('record_timing', !currentValue);
+  }
+
   /**
    * Clear the code outputs of the selected cells.
    *
@@ -1472,7 +1485,8 @@ namespace Private {
       case 'code':
         if (session) {
           return CodeCell.execute(cell as CodeCell, session, {
-            deletedCells: notebook.model.deletedCells
+            deletedCells: notebook.model.deletedCells,
+            recordTiming: notebook.model.metadata.get('record_timing') || false
           })
             .then(reply => {
               notebook.model.deletedCells.splice(

tests/test-cells/src/widget.spec.ts

@@ -756,6 +756,29 @@ describe('cells/widget', () => {
         expect(executionCount).not.toEqual(originalCount);
       });
 
+      const TIMING_KEYS = [
+        'iopub.execute_input',
+        'shell.execute_reply.started',
+        'shell.execute_reply',
+        'iopub.status.busy',
+        'iopub.status.idle'
+      ];
+
+      it('should not save timing info by default', async () => {
+        const widget = new CodeCell({ model, rendermime, contentFactory });
+        await CodeCell.execute(widget, session);
+        expect(widget.model.metadata.get('execution')).toBeUndefined();
+      });
+      it('should save timing info if requested', async () => {
+        const widget = new CodeCell({ model, rendermime, contentFactory });
+        await CodeCell.execute(widget, session, { recordTiming: true });
+        expect(widget.model.metadata.get('execution')).toBeDefined();
+        const timingInfo = widget.model.metadata.get('execution') as any;
+        for (const key of TIMING_KEYS) {
+          expect(timingInfo[key]).toBeDefined();
+        }
+      });
+
       it('should set the cell prompt properly while executing', async () => {
         const widget = new CodeCell({ model, rendermime, contentFactory });
         widget.initializeState();