Skip to content

Commit c54e912

Browse files
joyceerhljoyceerhl
authored
A/B test additional entrypoints for launching TensorBoard (#14922)
1 parent 3d53699 commit c54e912

14 files changed

+407
-98
lines changed

‎package.nls.json

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -223,13 +223,15 @@
223223
"Jupyter.extensionRequired": "The Jupyter extension is required to perform that task. Click Yes to open the Jupyter extension installation page.",
224224
"TensorBoard.useCurrentWorkingDirectory": "Use current working directory",
225225
"TensorBoard.currentDirectory": "Current: {0}",
226-
"TensorBoard.logDirectoryPrompt" : "Select a log directory to start TensorBoard with",
227-
"TensorBoard.progressMessage" : "Starting TensorBoard session...",
228-
"TensorBoard.failedToStartSessionError" : "We failed to start a TensorBoard session due to the following error: {0}",
229-
"TensorBoard.nativeTensorBoardPrompt" : "VS Code now has native TensorBoard support. Would you like to launch TensorBoard? (Tip: Launch TensorBoard anytime by opening the command palette and searching for \"Launch TensorBoard\".)",
226+
"TensorBoard.logDirectoryPrompt": "Select a log directory to start TensorBoard with",
227+
"TensorBoard.progressMessage": "Starting TensorBoard session...",
228+
"TensorBoard.failedToStartSessionError": "We failed to start a TensorBoard session due to the following error: {0}",
229+
"TensorBoard.nativeTensorBoardPrompt": "VS Code now has integrated TensorBoard support. Would you like to launch TensorBoard? (Tip: Launch TensorBoard anytime by opening the command palette and searching for \"Launch TensorBoard\".)",
230230
"TensorBoard.selectAFolder": "Select a folder",
231231
"TensorBoard.selectAnotherFolder": "Select another folder",
232232
"TensorBoard.selectAFolderDetail": "Select a log directory containing tfevent files",
233233
"TensorBoard.selectAnotherFolderDetail": "Use the file explorer to select another folder",
234-
"TensorBoard.useCurrentWorkingDirectoryDetail": "TensorBoard will search for tfevent files in all subdirectories of the current working directory"
234+
"TensorBoard.useCurrentWorkingDirectoryDetail": "TensorBoard will search for tfevent files in all subdirectories of the current working directory",
235+
"TensorBoard.launchNativeTensorBoardSessionCodeAction": "Launch TensorBoard session",
236+
"TensorBoard.launchNativeTensorBoardSessionCodeLens": "▶ Launch TensorBoard Session"
235237
}

‎src/client/common/experiments/groups.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,12 @@ export enum LinterInstallationPromptVariants {
9595
noPrompt = 'pythonNotDisplayLinterPrompt'
9696
}
9797

98+
// AB test codeactions vs codelenses as an entrypoint for native TensorBoard sessions
99+
export enum NativeTensorBoardEntrypoints {
100+
codeActions = 'pythonTensorBoardCodeActions',
101+
codeLenses = 'pythonTensorBoardCodeLenses'
102+
}
103+
98104
// Experiment to control which environment discovery mechanism can be used
99105
export enum DiscoveryVariants {
100106
discoverWithFileWatching = 'pythonDiscoveryModule',

‎src/client/common/utils/localize.ts

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ export namespace TensorBoard {
154154
);
155155
export const nativeTensorBoardPrompt = localize(
156156
'TensorBoard.nativeTensorBoardPrompt',
157-
'VS Code now has native TensorBoard support. Would you like to launch TensorBoard? (Tip: Launch TensorBoard anytime by opening the command palette and searching for "Launch TensorBoard".)'
157+
'VS Code now has integrated TensorBoard support. Would you like to launch TensorBoard? (Tip: Launch TensorBoard anytime by opening the command palette and searching for "Launch TensorBoard".)'
158158
);
159159
export const selectAFolder = localize('TensorBoard.selectAFolder', 'Select a folder');
160160
export const selectAFolderDetail = localize(
@@ -166,6 +166,14 @@ export namespace TensorBoard {
166166
'TensorBoard.selectAnotherFolderDetail',
167167
'Use the file explorer to select another folder'
168168
);
169+
export const launchNativeTensorBoardSessionCodeLens = localize(
170+
'TensorBoard.launchNativeTensorBoardSessionCodeLens',
171+
'▶ Launch TensorBoard Session'
172+
);
173+
export const launchNativeTensorBoardSessionCodeAction = localize(
174+
'TensorBoard.launchNativeTensorBoardSessionCodeAction',
175+
'Launch TensorBoard session'
176+
);
169177
}
170178

171179
export namespace LanguageService {

‎src/client/tensorBoard/helpers.ts

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
import { noop } from '../common/utils/misc';
5+
6+
// While it is uncommon for users to `import tensorboard`, TensorBoard is frequently
7+
// included as a submodule of other packages, e.g. torch.utils.tensorboard.
8+
// This is a modified version of the regex from src/client/telemetry/importTracker.ts
9+
// in order to match on imported submodules as well, since the original regex only
10+
// matches the 'main' module.
11+
12+
// eslint-disable-next-line max-len
13+
export const ImportRegEx = /^\s*from (?<fromImport>\w+(?:\.\w+)*) import (?<fromImportTarget>\w+(?:, \w+)*)(?: as \w+)?|import (?<importImport>\w+(?:, \w+)*)(?: as \w+)?$/;
14+
15+
export function containsTensorBoardImport(lines: (string | undefined)[]): boolean {
16+
try {
17+
for (const s of lines) {
18+
const matches = s ? ImportRegEx.exec(s) : null;
19+
if (matches !== null && matches.groups !== undefined) {
20+
let componentsToCheck: string[] = [];
21+
if (matches.groups.fromImport && matches.groups.fromImportTarget) {
22+
// from x.y.z import u, v, w
23+
componentsToCheck = matches.groups.fromImport
24+
.split('.')
25+
.concat(matches.groups.fromImportTarget.split(','));
26+
} else if (matches.groups.importImport) {
27+
// import package1, package2, ...
28+
componentsToCheck = matches.groups.importImport.split(',');
29+
}
30+
for (const component of componentsToCheck) {
31+
if (component && component.trim() === 'tensorboard') {
32+
return true;
33+
}
34+
}
35+
}
36+
}
37+
} catch {
38+
// Don't care about failures.
39+
noop();
40+
}
41+
return false;
42+
}

‎src/client/tensorBoard/serviceRegistry.ts

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,15 @@
33

44
import { IExtensionSingleActivationService } from '../activation/types';
55
import { IServiceManager } from '../ioc/types';
6+
import { TensorBoardCodeActionProvider } from './tensorBoardCodeActionProvider';
7+
import { TensorBoardCodeLensProvider } from './tensorBoardCodeLensProvider';
68
import { TensorBoardFileWatcher } from './tensorBoardFileWatcher';
79
import { TensorBoardImportTracker } from './tensorBoardImportTracker';
810
import { TensorBoardPrompt } from './tensorBoardPrompt';
911
import { TensorBoardSessionProvider } from './tensorBoardSessionProvider';
1012
import { ITensorBoardImportTracker } from './types';
1113

12-
export function registerTypes(serviceManager: IServiceManager) {
14+
export function registerTypes(serviceManager: IServiceManager): void {
1315
serviceManager.addSingleton<IExtensionSingleActivationService>(
1416
IExtensionSingleActivationService,
1517
TensorBoardSessionProvider
@@ -21,4 +23,12 @@ export function registerTypes(serviceManager: IServiceManager) {
2123
serviceManager.addSingleton<TensorBoardPrompt>(TensorBoardPrompt, TensorBoardPrompt);
2224
serviceManager.addSingleton<ITensorBoardImportTracker>(ITensorBoardImportTracker, TensorBoardImportTracker);
2325
serviceManager.addBinding(ITensorBoardImportTracker, IExtensionSingleActivationService);
26+
serviceManager.addSingleton<IExtensionSingleActivationService>(
27+
IExtensionSingleActivationService,
28+
TensorBoardCodeLensProvider
29+
);
30+
serviceManager.addSingleton<IExtensionSingleActivationService>(
31+
IExtensionSingleActivationService,
32+
TensorBoardCodeActionProvider
33+
);
2434
}

‎src/client/tensorBoard/tensorBoardCodeActionProvider.ts

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
import { inject, injectable } from 'inversify';
5+
import {
6+
CancellationToken,
7+
CodeAction,
8+
CodeActionContext,
9+
CodeActionKind,
10+
CodeActionProvider,
11+
languages,
12+
Selection,
13+
TextDocument
14+
} from 'vscode';
15+
import { IExtensionSingleActivationService } from '../activation/types';
16+
import { Commands, PYTHON } from '../common/constants';
17+
import { NativeTensorBoard, NativeTensorBoardEntrypoints } from '../common/experiments/groups';
18+
import { IDisposableRegistry, IExperimentService } from '../common/types';
19+
import { TensorBoard } from '../common/utils/localize';
20+
import { containsTensorBoardImport } from './helpers';
21+
22+
@injectable()
23+
export class TensorBoardCodeActionProvider implements CodeActionProvider, IExtensionSingleActivationService {
24+
constructor(
25+
@inject(IExperimentService) private experimentService: IExperimentService,
26+
@inject(IDisposableRegistry) private disposables: IDisposableRegistry
27+
) {}
28+
29+
public async activate(): Promise<void> {
30+
// Don't hold up activation for this
31+
this.activateInternal().ignoreErrors();
32+
}
33+
34+
// eslint-disable-next-line class-methods-use-this
35+
public provideCodeActions(
36+
document: TextDocument,
37+
range: Selection,
38+
_context: CodeActionContext,
39+
_token: CancellationToken
40+
): CodeAction[] {
41+
const cursorPosition = range.active;
42+
const { text } = document.lineAt(cursorPosition);
43+
if (containsTensorBoardImport([text])) {
44+
const title = TensorBoard.launchNativeTensorBoardSessionCodeAction();
45+
const nativeTensorBoardSession = new CodeAction(title, CodeActionKind.QuickFix);
46+
nativeTensorBoardSession.command = {
47+
title,
48+
command: Commands.LaunchTensorBoard
49+
};
50+
return [nativeTensorBoardSession];
51+
}
52+
return [];
53+
}
54+
55+
private async activateInternal() {
56+
if (
57+
(await this.experimentService.inExperiment(NativeTensorBoard.experiment)) &&
58+
(await this.experimentService.inExperiment(NativeTensorBoardEntrypoints.codeActions))
59+
) {
60+
this.disposables.push(
61+
languages.registerCodeActionsProvider(PYTHON, this, {
62+
providedCodeActionKinds: [CodeActionKind.QuickFix]
63+
})
64+
);
65+
}
66+
}
67+
}

‎src/client/tensorBoard/tensorBoardCodeLensProvider.ts

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
import { inject, injectable } from 'inversify';
5+
import { CancellationToken, CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode';
6+
import { IExtensionSingleActivationService } from '../activation/types';
7+
import { Commands, PYTHON } from '../common/constants';
8+
import { NativeTensorBoard, NativeTensorBoardEntrypoints } from '../common/experiments/groups';
9+
import { IDisposableRegistry, IExperimentService } from '../common/types';
10+
import { TensorBoard } from '../common/utils/localize';
11+
import { containsTensorBoardImport } from './helpers';
12+
13+
@injectable()
14+
export class TensorBoardCodeLensProvider implements IExtensionSingleActivationService {
15+
constructor(
16+
@inject(IExperimentService) private experimentService: IExperimentService,
17+
@inject(IDisposableRegistry) private disposables: IDisposableRegistry
18+
) {}
19+
20+
public async activate(): Promise<void> {
21+
this.activateInternal().ignoreErrors();
22+
}
23+
24+
// eslint-disable-next-line class-methods-use-this
25+
public provideCodeLenses(document: TextDocument, _token: CancellationToken): CodeLens[] {
26+
const command: Command = {
27+
title: TensorBoard.launchNativeTensorBoardSessionCodeLens(),
28+
command: Commands.LaunchTensorBoard
29+
};
30+
const codelenses: CodeLens[] = [];
31+
for (let index = 0; index < document.lineCount; index += 1) {
32+
const line = document.lineAt(index);
33+
if (containsTensorBoardImport([line.text])) {
34+
const range = new Range(new Position(line.lineNumber, 0), new Position(line.lineNumber, 1));
35+
codelenses.push(new CodeLens(range, command));
36+
}
37+
}
38+
return codelenses;
39+
}
40+
41+
private async activateInternal() {
42+
if (
43+
(await this.experimentService.inExperiment(NativeTensorBoard.experiment)) &&
44+
(await this.experimentService.inExperiment(NativeTensorBoardEntrypoints.codeLenses))
45+
) {
46+
this.disposables.push(languages.registerCodeLensProvider(PYTHON, this));
47+
}
48+
}
49+
}

‎src/client/tensorBoard/tensorBoardFileWatcher.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,11 @@ export class TensorBoardFileWatcher implements IExtensionSingleActivationService
2525
@inject(IExperimentService) private experimentService: IExperimentService
2626
) {}
2727

28-
public async activate() {
28+
public async activate(): Promise<void> {
29+
this.activateInternal().ignoreErrors();
30+
}
31+
32+
private async activateInternal() {
2933
if (!(await this.experimentService.inExperiment(NativeTensorBoard.experiment))) {
3034
return;
3135
}

‎src/client/tensorBoard/tensorBoardImportTracker.ts

Lines changed: 27 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,18 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
14
import { inject, injectable } from 'inversify';
2-
import { noop } from 'lodash';
35
import * as path from 'path';
4-
import { Event, EventEmitter, TextEditor, window } from 'vscode';
6+
import { Event, EventEmitter, TextEditor } from 'vscode';
57
import { IExtensionSingleActivationService } from '../activation/types';
68
import { IDocumentManager } from '../common/application/types';
9+
import { isTestExecution } from '../common/constants';
710
import { IDisposableRegistry } from '../common/types';
811
import { getDocumentLines } from '../telemetry/importTracker';
12+
import { containsTensorBoardImport } from './helpers';
913
import { ITensorBoardImportTracker } from './types';
1014

11-
// While it is uncommon for users to `import tensorboard`, TensorBoard is frequently
12-
// included as a submodule of other packages, e.g. torch.utils.tensorboard.
13-
// This is a modified version of the regex from src/client/telemetry/importTracker.ts
14-
// in order to match on imported submodules as well, since the original regex only
15-
// matches the 'main' module.
16-
const ImportRegEx = /^\s*from (?<fromImport>\w+(?:\.\w+)*) import (?<fromImportTarget>\w+(?:, \w+)*)(?: as \w+)?|import (?<importImport>\w+(?:, \w+)*)(?: as \w+)?$/;
17-
15+
const testExecution = isTestExecution();
1816
@injectable()
1917
export class TensorBoardImportTracker implements ITensorBoardImportTracker, IExtensionSingleActivationService {
2018
private pendingChecks = new Map<string, NodeJS.Timer | number>();
@@ -24,26 +22,34 @@ export class TensorBoardImportTracker implements ITensorBoardImportTracker, IExt
2422
constructor(
2523
@inject(IDocumentManager) private documentManager: IDocumentManager,
2624
@inject(IDisposableRegistry) private disposables: IDisposableRegistry
27-
) {
28-
this.documentManager.onDidChangeActiveTextEditor(
29-
(e) => this.onChangedActiveTextEditor(e),
30-
this,
31-
this.disposables
32-
);
33-
}
25+
) {}
3426

3527
// Fires when the active text editor contains a tensorboard import.
3628
public get onDidImportTensorBoard(): Event<void> {
3729
return this._onDidImportTensorBoard.event;
3830
}
3931

40-
public dispose() {
32+
public dispose(): void {
4133
this.pendingChecks.clear();
4234
}
4335

4436
public async activate(): Promise<void> {
45-
// Process active text editor with a timeout delay
46-
this.onChangedActiveTextEditor(window.activeTextEditor);
37+
if (testExecution) {
38+
await this.activateInternal();
39+
} else {
40+
this.activateInternal().ignoreErrors();
41+
}
42+
}
43+
44+
private async activateInternal() {
45+
// Process currently active text editor
46+
this.onChangedActiveTextEditor(this.documentManager.activeTextEditor);
47+
// Process changes to active text editor as well
48+
this.documentManager.onDidChangeActiveTextEditor(
49+
(e) => this.onChangedActiveTextEditor(e),
50+
this,
51+
this.disposables
52+
);
4753
}
4854

4955
private onChangedActiveTextEditor(editor: TextEditor | undefined) {
@@ -56,38 +62,9 @@ export class TensorBoardImportTracker implements ITensorBoardImportTracker, IExt
5662
path.extname(document.fileName) === '.py'
5763
) {
5864
const lines = getDocumentLines(document);
59-
this.lookForImports(lines);
60-
}
61-
}
62-
63-
private lookForImports(lines: (string | undefined)[]) {
64-
try {
65-
for (const s of lines) {
66-
const matches = s ? ImportRegEx.exec(s) : null;
67-
if (matches === null || matches.groups === undefined) {
68-
// eslint-disable-next-line no-continue
69-
continue;
70-
}
71-
let componentsToCheck: string[] = [];
72-
if (matches.groups.fromImport && matches.groups.fromImportTarget) {
73-
// from x.y.z import u, v, w
74-
componentsToCheck = matches.groups.fromImport
75-
.split('.')
76-
.concat(matches.groups.fromImportTarget.split(','));
77-
} else if (matches.groups.importImport) {
78-
// import package1, package2, ...
79-
componentsToCheck = matches.groups.importImport.split(',');
80-
}
81-
for (const component of componentsToCheck) {
82-
if (component && component.trim() === 'tensorboard') {
83-
this._onDidImportTensorBoard.fire();
84-
return;
85-
}
86-
}
65+
if (containsTensorBoardImport(lines)) {
66+
this._onDidImportTensorBoard.fire();
8767
}
88-
} catch {
89-
// Don't care about failures.
90-
noop();
9168
}
9269
}
9370
}

0 commit comments

Comments
 (0)