Skip to content

Commit

Permalink
Optimize to handling of large outputs in notebooks (#11350)
Browse files Browse the repository at this point in the history 
* Optimize to handling of large outputs in notebooks

* Misc

* test

* More logging

* Remove test

* Misc

* Remove test
  • Loading branch information
@DonJayamanne
DonJayamanne authored Oct 13, 2022
1 parent eaf40cb commit 4b4df37
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 136 deletions.
1 change: 1 addition & 0 deletions news/2 Fixes/11031.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Optimize the way stream outputs are handled, to support large output streams without crashing VS Code.
40 changes: 9 additions & 31 deletions src/kernels/execution/cellExecutionMessageHandler.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ export class CellExecutionMessageHandler implements IDisposable {
* If users clear outputs or if we have a new output other than stream, then clear this item.
* Because if after the stream we have an image, then the stream is not the last output item, hence its cleared.
*/
private lastUsedStreamOutput?: { stream: 'stdout' | 'stderr'; text: string; output: NotebookCellOutput };
private lastUsedStreamOutput?: { stream: 'stdout' | 'stderr'; output: NotebookCellOutput };
/**
* When we have nested Output Widgets, we get comm messages one for each output widget.
* The way it works is:
Expand Down Expand Up @@ -818,45 +818,23 @@ export class CellExecutionMessageHandler implements IDisposable {
// Ensure we append to previous output, only if the streams as the same &
// If the last output is the desired stream type.
if (this.lastUsedStreamOutput?.stream === msg.content.name) {
// Get the jupyter output from the vs code output (so we can concatenate the text ourselves).
let existingOutputText = this.lastUsedStreamOutput.text;
let newContent = msg.content.text;
// Look for the ansi code `<char27>[A`. (this means move up)
// Not going to support `[2A` (not for now).
const moveUpCode = `${String.fromCharCode(27)}[A`;
if (msg.content.text.startsWith(moveUpCode)) {
// Split message by lines & strip out the last n lines (where n = number of lines to move cursor up).
const existingOutputLines = existingOutputText.splitLines({
trim: false,
removeEmptyEntries: false
});
if (existingOutputLines.length) {
existingOutputLines.pop();
}
existingOutputText = existingOutputLines.join('\n');
newContent = newContent.substring(moveUpCode.length);
}
// Create a new output item with the concatenated string.
this.lastUsedStreamOutput.text = formatStreamText(
concatMultilineString(`${existingOutputText}${newContent}`)
);
const output = cellOutputToVSCCellOutput({
output_type: 'stream',
name: msg.content.name,
text: this.lastUsedStreamOutput.text
text: msg.content.text
});
traceCellMessage(this.cell, `Replace output items '${this.lastUsedStreamOutput.text.substring(0, 100)}'`);
task?.replaceOutputItems(output.items, this.lastUsedStreamOutput.output).then(noop, noop);
traceCellMessage(this.cell, `Append output items '${msg.content.text.substring(0, 100)}`);
task?.appendOutputItems(output.items, this.lastUsedStreamOutput.output).then(noop, noop);
} else if (previousValueOfClearOutputOnNextUpdateToOutput) {
// Replace the current outputs with a single new output.
const text = formatStreamText(concatMultilineString(msg.content.text));
const text = concatMultilineString(msg.content.text);
const output = cellOutputToVSCCellOutput({
output_type: 'stream',
name: msg.content.name,
text
});
this.lastUsedStreamOutput = { output, stream: msg.content.name, text };
traceCellMessage(this.cell, `Replace output '${this.lastUsedStreamOutput.text.substring(0, 100)}'`);
this.lastUsedStreamOutput = { output, stream: msg.content.name };
traceCellMessage(this.cell, `Replace output with '${text.substring(0, 100)}'`);
task?.replaceOutput([output]).then(noop, noop);
} else {
// Create a new output
Expand All @@ -866,8 +844,8 @@ export class CellExecutionMessageHandler implements IDisposable {
name: msg.content.name,
text
});
this.lastUsedStreamOutput = { output, stream: msg.content.name, text };
traceCellMessage(this.cell, `Append output '${this.lastUsedStreamOutput.text.substring(0, 100)}'`);
this.lastUsedStreamOutput = { output, stream: msg.content.name };
traceCellMessage(this.cell, `Append new output '${text.substring(0, 100)}'`);
task?.appendOutput([output]).then(noop, noop);
}
this.endTemporaryTask();
Expand Down
106 changes: 1 addition & 105 deletions src/test/datascience/notebook/executionService.vscode.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,7 @@ import * as fs from 'fs';
import * as path from '../../../platform/vscode-path/path';
import dedent from 'dedent';
import * as sinon from 'sinon';
import {
commands,
NotebookCell,
NotebookCellExecutionState,
NotebookCellKind,
NotebookCellOutput,
Uri,
window,
workspace
} from 'vscode';
import { commands, NotebookCell, NotebookCellExecutionState, NotebookCellKind, Uri, window, workspace } from 'vscode';
import { Common } from '../../../platform/common/utils/localize';
import { IVSCodeNotebook } from '../../../platform/common/application/types';
import { traceInfo } from '../../../platform/logging';
Expand Down Expand Up @@ -763,74 +754,6 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
)
]);
});
test('Outputs with support for ansic code `\u001b[A`', async function () {
// Ansi Code `<esc>[A` means move cursor up, i.e. replace previous line with the new output (or erase previous line & start there).
const cell1 = await insertCodeCell(
dedent`
import sys
import os
sys.stdout.write(f"Line1{os.linesep}")
sys.stdout.flush()
sys.stdout.write(os.linesep)
sys.stdout.flush()
sys.stdout.write(f"Line3{os.linesep}")
sys.stdout.flush()
sys.stdout.write("Line4")
`,
{ index: 0 }
);
const cell2 = await insertCodeCell(
dedent`
import sys
import os
sys.stdout.write(f"Line1{os.linesep}")
sys.stdout.flush()
sys.stdout.write(os.linesep)
sys.stdout.flush()
sys.stdout.write(u"\u001b[A")
sys.stdout.flush()
sys.stdout.write(f"Line3{os.linesep}")
sys.stdout.flush()
sys.stdout.write("Line4")
`,
{ index: 1 }
);

await Promise.all([
runAllCellsInActiveNotebook(),
waitForTextOutput(cell1, 'Line4', 0, false),
waitForTextOutput(cell2, 'Line4', 0, false)
]);

// In cell 1 we should have the output
// Line1
//
// Line2
// Line3
// & in cell 2 we should have the output
// Line1
// Line2
// Line3

// Work around https://github.com/ipython/ipykernel/issues/729
const ignoreEmptyOutputs = (output: NotebookCellOutput) => {
return output.items.filter((item) => item.mime !== 'text/plain').length > 0;
};
assert.equal(cell1.outputs.filter(ignoreEmptyOutputs).length, 1, 'Incorrect number of output');
assert.equal(cell2.outputs.filter(ignoreEmptyOutputs).length, 1, 'Incorrect number of output');

// Confirm the output
const output1Lines: string[] = getTextOutputValue(cell1.outputs[0]).splitLines({
trim: false,
removeEmptyEntries: false
});
const output2Lines: string[] = getTextOutputValue(cell2.outputs[0]).splitLines({
trim: false,
removeEmptyEntries: false
});
assert.equal(output1Lines.length, 4);
assert.equal(output2Lines.length, 3);
});
test('Stderr & stdout outputs should go into separate outputs', async function () {
await insertCodeCell(
dedent`
Expand Down Expand Up @@ -904,33 +827,6 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
}
});

test('Handling of carriage returns', async () => {
await insertCodeCell('print("one\\r", end="")\nprint("two\\r", end="")\nprint("three\\r", end="")', {
index: 0
});
await insertCodeCell('print("one\\r")\nprint("two\\r")\nprint("three\\r")', { index: 1 });
await insertCodeCell('print("1\\r2\\r3")', { index: 2 });
await insertCodeCell('print("1\\r2")', { index: 3 });
await insertCodeCell(
'import time\nfor i in range(10):\n s = str(i) + "%"\n print("{0}\\r".format(s),end="")\n time.sleep(0.0001)',
{ index: 4 }
);
await insertCodeCell('print("\\rExecute\\rExecute\\nExecute 8\\rExecute 9\\r\\r")', { index: 5 });

const cells = vscodeNotebook.activeNotebookEditor!.notebook.getCells();
await Promise.all([runAllCellsInActiveNotebook(), waitForExecutionCompletedSuccessfully(cells[5])]);

// Wait for the outputs.
await Promise.all([
waitForTextOutput(cells[0], 'three\r', 0, true),
waitForTextOutput(cells[1], 'one\ntwo\nthree\n', 0, true),
waitForTextOutput(cells[2], '3\n', 0, true),
waitForTextOutput(cells[3], '2\n', 0, true),
waitForTextOutput(cells[4], '9%\r', 0, true),
waitForTextOutput(cells[5], 'Execute\nExecute 9\n', 0, true)
]);
});

test('Execute all cells and run after error', async () => {
await insertCodeCell('raise Error("fail")', { index: 0 });
await insertCodeCell('print("after fail")', { index: 1 });
Expand Down

0 comments on commit 4b4df37

Please sign in to comment.