Skip to content

Commit

Permalink
Fix #171385. Ignore controller affinity in MRU. (#171570)
Browse files Browse the repository at this point in the history 
* Move kernel resolving into strategy.

* Single controller for notebook will be always be used.

* Fix unit test.
  • Loading branch information
@rebornix
rebornix authored Jan 18, 2023
1 parent a3b2f85 commit 78b91ac
Show file tree
Hide file tree
Showing 5 changed files with 97 additions and 70 deletions.
67 changes: 12 additions & 55 deletions src/vs/workbench/contrib/notebook/browser/services/notebookExecutionServiceImpl.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,16 @@ import { CancellationTokenSource } from 'vs/base/common/cancellation';
import { IDisposable } from 'vs/base/common/lifecycle';
import * as nls from 'vs/nls';
import { ICommandService } from 'vs/platform/commands/common/commands';
import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
import { ILogService } from 'vs/platform/log/common/log';
import { IWorkspaceTrustRequestService } from 'vs/platform/workspace/common/workspaceTrust';
import { SELECT_KERNEL_ID } from 'vs/workbench/contrib/notebook/browser/controller/coreActions';
import { KernelPickerFlatStrategy, KernelPickerMRUStrategy } from 'vs/workbench/contrib/notebook/browser/viewParts/notebookKernelQuickPickStrategy';
import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
import { CellKind, INotebookTextModel, NotebookCellExecutionState } from 'vs/workbench/contrib/notebook/common/notebookCommon';
import { INotebookExecutionService } from 'vs/workbench/contrib/notebook/common/notebookExecutionService';
import { INotebookCellExecution, INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
import { INotebookKernel, INotebookKernelService } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
import { INotebookKernel, INotebookKernelHistoryService, INotebookKernelService } from 'vs/workbench/contrib/notebook/common/notebookKernelService';

export class NotebookExecutionService implements INotebookExecutionService, IDisposable {
declare _serviceBrand: undefined;
Expand All @@ -24,9 +25,11 @@ export class NotebookExecutionService implements INotebookExecutionService, IDis
constructor(
@ICommandService private readonly _commandService: ICommandService,
@INotebookKernelService private readonly _notebookKernelService: INotebookKernelService,
@INotebookKernelHistoryService private readonly _notebookKernelHistoryService: INotebookKernelHistoryService,
@IWorkspaceTrustRequestService private readonly _workspaceTrustRequestService: IWorkspaceTrustRequestService,
@ILogService private readonly _logService: ILogService,
@INotebookExecutionStateService private readonly _notebookExecutionStateService: INotebookExecutionStateService
@INotebookExecutionStateService private readonly _notebookExecutionStateService: INotebookExecutionStateService,
@IConfigurationService private readonly _configurationService: IConfigurationService,
) {
}

Expand Down Expand Up @@ -54,15 +57,13 @@ export class NotebookExecutionService implements INotebookExecutionService, IDis
cellExecutions.push([cell, this._notebookExecutionStateService.createCellExecution(notebook.uri, cell.handle)]);
}

let kernel = this._notebookKernelService.getSelectedOrSuggestedKernel(notebook);
const noPendingKernelDetections = this._notebookKernelService.getKernelDetectionTasks(notebook).length === 0;
// do not auto run source action when there is pending kernel detections
if (!kernel && noPendingKernelDetections) {
kernel = await this.resolveSourceActions(notebook, contextKeyService);
}
let kernel: INotebookKernel | undefined;
const kernelPickerType = this._configurationService.getValue<'all' | 'mru'>('notebook.kernelPicker.type');

if (!kernel) {
kernel = await this.resolveKernelFromKernelPicker(notebook);
if (kernelPickerType === 'mru') {
kernel = await KernelPickerMRUStrategy.resolveKernel(notebook, this._notebookKernelService, this._notebookKernelHistoryService, this._commandService);
} else {
kernel = await KernelPickerFlatStrategy.resolveKernel(notebook, this._notebookKernelService, this._commandService, contextKeyService);
}

if (!kernel) {
Expand Down Expand Up @@ -94,50 +95,6 @@ export class NotebookExecutionService implements INotebookExecutionService, IDis
}
}

private async resolveKernelFromKernelPicker(notebook: INotebookTextModel, attempt: number = 1): Promise<INotebookKernel | undefined> {
if (attempt > 3) {
// we couldnt resolve kernels through kernel picker multiple times, skip
return;
}

await this._commandService.executeCommand(SELECT_KERNEL_ID);
const runningSourceActions = this._notebookKernelService.getRunningSourceActions(notebook);

if (runningSourceActions.length) {
await Promise.all(runningSourceActions.map(action => action.runAction()));

const kernel = this._notebookKernelService.getSelectedOrSuggestedKernel(notebook);
if (kernel) {
return kernel;
}

attempt += 1;
return this.resolveKernelFromKernelPicker(notebook, attempt);
} else {
return this._notebookKernelService.getSelectedOrSuggestedKernel(notebook);
}
}

private async resolveSourceActions(notebook: INotebookTextModel, contextKeyService: IContextKeyService) {
let kernel: INotebookKernel | undefined;
const info = this._notebookKernelService.getMatchingKernel(notebook);
if (info.all.length === 0) {
// no kernel at all
const sourceActions = this._notebookKernelService.getSourceActions(notebook, contextKeyService);
const primaryActions = sourceActions.filter(action => action.isPrimary);
const action = sourceActions.length === 1
? sourceActions[0]
: (primaryActions.length === 1 ? primaryActions[0] : undefined);

if (action) {
await action.runAction();
kernel = this._notebookKernelService.getSelectedOrSuggestedKernel(notebook);
}
}

return kernel;
}

async cancelNotebookCellHandles(notebook: INotebookTextModel, cells: Iterable<number>): Promise<void> {
const cellsArr = Array.from(cells);
this._logService.debug(`NotebookExecutionService#cancelNotebookCellHandles ${JSON.stringify(cellsArr)}`);
Expand Down
4 changes: 2 additions & 2 deletions src/vs/workbench/contrib/notebook/browser/services/notebookKernelHistoryServiceImpl.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,8 @@ export class NotebookKernelHistoryService extends Disposable implements INoteboo
const allAvailableKernels = this._notebookKernelService.getMatchingKernel(notebook);
const allKernels = allAvailableKernels.all;
const selectedKernel = allAvailableKernels.selected;
const suggested = (allAvailableKernels.suggestions.length === 1 ? allAvailableKernels.suggestions[0] : undefined)
?? (allAvailableKernels.all.length === 1) ? allAvailableKernels.all[0] : undefined;
// We will suggest the only kernel
const suggested = allAvailableKernels.all.length === 1 ? allAvailableKernels.all[0] : undefined;

const mostRecentKernelIds = this._mostRecentKernelsMap[notebook.viewType] ? [...this._mostRecentKernelsMap[notebook.viewType].values()] : [];

Expand Down
87 changes: 78 additions & 9 deletions src/vs/workbench/contrib/notebook/browser/viewParts/notebookKernelQuickPickStrategy.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ import { IExtensionService } from 'vs/workbench/services/extensions/common/exten
import { IPaneCompositePartService } from 'vs/workbench/services/panecomposite/browser/panecomposite';
import { URI } from 'vs/base/common/uri';
import { IOpenerService } from 'vs/platform/opener/common/opener';
import { INotebookTextModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
import { SELECT_KERNEL_ID } from 'vs/workbench/contrib/notebook/browser/controller/coreActions';

type KernelPick = IQuickPickItem & { kernel: INotebookKernel };
function isKernelPick(item: QuickPickInput<IQuickPickItem>): item is KernelPick {
Expand Down Expand Up @@ -568,6 +570,65 @@ export class KernelPickerFlatStrategy extends KernelPickerStrategyBase {
action.tooltip = '';
}
}

static async resolveKernel(notebook: INotebookTextModel, notebookKernelService: INotebookKernelService, commandService: ICommandService, contextKeyService: IContextKeyService): Promise<INotebookKernel | undefined> {
let kernel = notebookKernelService.getSelectedOrSuggestedKernel(notebook);
const noPendingKernelDetections = notebookKernelService.getKernelDetectionTasks(notebook).length === 0;
// do not auto run source action when there is pending kernel detections
if (!kernel && noPendingKernelDetections) {
kernel = await KernelPickerFlatStrategy.resolveSourceActions(notebook, notebookKernelService, contextKeyService);
}

if (!kernel) {
kernel = await this.resolveKernelFromKernelPicker(notebook, commandService, notebookKernelService);
}

return kernel;
}

private static async resolveSourceActions(notebook: INotebookTextModel, notebookKernelService: INotebookKernelService, contextKeyService: IContextKeyService): Promise<INotebookKernel | undefined> {
let kernel: INotebookKernel | undefined;
const info = notebookKernelService.getMatchingKernel(notebook);
if (info.all.length === 0) {
// no kernel at all
const sourceActions = notebookKernelService.getSourceActions(notebook, contextKeyService);
const primaryActions = sourceActions.filter(action => action.isPrimary);
const action = sourceActions.length === 1
? sourceActions[0]
: (primaryActions.length === 1 ? primaryActions[0] : undefined);

if (action) {
await action.runAction();
kernel = notebookKernelService.getSelectedOrSuggestedKernel(notebook);
}
}

return kernel;
}

private static async resolveKernelFromKernelPicker(notebook: INotebookTextModel, commandService: ICommandService, notebookKernelService: INotebookKernelService, attempt: number = 1): Promise<INotebookKernel | undefined> {
if (attempt > 3) {
// we couldnt resolve kernels through kernel picker multiple times, skip
return;
}

await commandService.executeCommand(SELECT_KERNEL_ID);
const runningSourceActions = notebookKernelService.getRunningSourceActions(notebook);

if (runningSourceActions.length) {
await Promise.all(runningSourceActions.map(action => action.runAction()));

const kernel = notebookKernelService.getSelectedOrSuggestedKernel(notebook);
if (kernel) {
return kernel;
}

attempt += 1;
return KernelPickerFlatStrategy.resolveKernelFromKernelPicker(notebook, commandService, notebookKernelService, attempt);
} else {
return notebookKernelService.getSelectedOrSuggestedKernel(notebook);
}
}
}

export class KernelPickerMRUStrategy extends KernelPickerStrategyBase {
Expand Down Expand Up @@ -877,7 +938,7 @@ export class KernelPickerMRUStrategy extends KernelPickerStrategyBase {
}
}

static updateKernelStatusAction(notebook: NotebookTextModel, action: IAction, notebookKernelService: INotebookKernelService) {
static updateKernelStatusAction(notebook: NotebookTextModel, action: IAction, notebookKernelService: INotebookKernelService, notebookKernelHistoryService: INotebookKernelHistoryService) {
const detectionTasks = notebookKernelService.getKernelDetectionTasks(notebook);
if (detectionTasks.length) {
const info = notebookKernelService.getMatchingKernel(notebook);
Expand Down Expand Up @@ -909,20 +970,28 @@ export class KernelPickerMRUStrategy extends KernelPickerStrategyBase {
return updateActionFromSourceAction(runningActions[0] /** TODO handle multiple actions state */, true);
}

const info = notebookKernelService.getMatchingKernel(notebook);
const suggested = (info.suggestions.length === 1 ? info.suggestions[0] : undefined)
?? (info.all.length === 1) ? info.all[0] : undefined;

const selectedOrSuggested = info.selected ?? suggested;
const { selected } = notebookKernelHistoryService.getKernels(notebook);

if (selectedOrSuggested) {
action.label = selectedOrSuggested.label;
if (selected) {
action.label = selected.label;
action.class = ThemeIcon.asClassName(selectKernelIcon);
action.tooltip = selectedOrSuggested.description ?? selectedOrSuggested.detail ?? '';
action.tooltip = selected.description ?? selected.detail ?? '';
} else {
action.label = localize('select', "Select Kernel");
action.class = ThemeIcon.asClassName(selectKernelIcon);
action.tooltip = '';
}
}

static async resolveKernel(notebook: INotebookTextModel, notebookKernelService: INotebookKernelService, notebookKernelHistoryService: INotebookKernelHistoryService, commandService: ICommandService): Promise<INotebookKernel | undefined> {
const { selected } = notebookKernelHistoryService.getKernels(notebook);

if (selected) {
return selected;
}

await commandService.executeCommand(SELECT_KERNEL_ID);
const kernel = notebookKernelService.getSelectedOrSuggestedKernel(notebook);
return kernel;
}
}
5 changes: 3 additions & 2 deletions src/vs/workbench/contrib/notebook/browser/viewParts/notebookKernelView.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ import { selectKernelIcon } from 'vs/workbench/contrib/notebook/browser/notebook
import { KernelPickerFlatStrategy, KernelPickerMRUStrategy, KernelQuickPickContext } from 'vs/workbench/contrib/notebook/browser/viewParts/notebookKernelQuickPickStrategy';
import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
import { NOTEBOOK_IS_ACTIVE_EDITOR, NOTEBOOK_KERNEL_COUNT } from 'vs/workbench/contrib/notebook/common/notebookContextKeys';
import { INotebookKernelService } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
import { INotebookKernelHistoryService, INotebookKernelService } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
import { IEditorService } from 'vs/workbench/services/editor/common/editorService';

function getEditorFromContext(editorService: IEditorService, context?: KernelQuickPickContext): INotebookEditor | undefined {
Expand Down Expand Up @@ -146,6 +146,7 @@ export class NotebooKernelActionViewItem extends ActionViewItem {
actualAction: IAction,
private readonly _editor: { onDidChangeModel: Event<void>; textModel: NotebookTextModel | undefined; scopedContextKeyService?: IContextKeyService } | INotebookEditor,
@INotebookKernelService private readonly _notebookKernelService: INotebookKernelService,
@INotebookKernelHistoryService private readonly _notebookKernelHistoryService: INotebookKernelHistoryService,
@IConfigurationService private readonly _configurationService: IConfigurationService,
) {
super(
Expand Down Expand Up @@ -187,7 +188,7 @@ export class NotebooKernelActionViewItem extends ActionViewItem {

const kernelPickerType = this._configurationService.getValue<'all' | 'mru'>('notebook.kernelPicker.type');
if (kernelPickerType === 'mru') {
KernelPickerMRUStrategy.updateKernelStatusAction(notebook, this._action, this._notebookKernelService);
KernelPickerMRUStrategy.updateKernelStatusAction(notebook, this._action, this._notebookKernelService, this._notebookKernelHistoryService);
} else {
KernelPickerFlatStrategy.updateKernelStatusAction(notebook, this._action, this._notebookKernelService, this._editor.scopedContextKeyService);
}
Expand Down
4 changes: 2 additions & 2 deletions src/vs/workbench/contrib/notebook/test/browser/notebookKernelHistory.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,8 @@ suite('NotebookKernelHistoryService', () => {

info = kernelHistoryService.getKernels({ uri: u1, viewType: 'foo' });
assert.equal(info.all.length, 0);
// suggested kernel should be visible
assert.deepStrictEqual(info.selected, k2);
// MRU only auto selects kernel if there is only one
assert.deepStrictEqual(info.selected, undefined);
});

test('notebook kernel history restore', function () {
Expand Down

0 comments on commit 78b91ac

Please sign in to comment.