-
Notifications
You must be signed in to change notification settings - Fork 297
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Revision to displaying cached remote kernel specs and live kernels (#…
…9985) * Misc * Misc * Misc * Misc * oops * Misc * Add tests * Misc * Misc * Address code review comments * Misc * oops * Fix tests
- Loading branch information
1 parent
3675351
commit 1143125
Showing
19 changed files
with
877 additions
and
86 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
100 changes: 100 additions & 0 deletions
100
src/kernels/jupyter/liveRemoteKernelConnectionTracker.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,100 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| import { inject, injectable, named } from 'inversify'; | ||
| import { Memento, Uri } from 'vscode'; | ||
| import { IExtensionSyncActivationService } from '../../platform/activation/types'; | ||
| import { GLOBAL_MEMENTO, IDisposableRegistry, IMemento } from '../../platform/common/types'; | ||
| import { noop } from '../../platform/common/utils/misc'; | ||
| import { LiveRemoteKernelConnectionMetadata } from '../types'; | ||
| import { computeServerId } from './jupyterUtils'; | ||
| import { IJupyterServerUriStorage, ILiveRemoteKernelConnectionUsageTracker } from './types'; | ||
|
|
||
| export const mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources = 'removeKernelUrisAndSessionsUsedByResources'; | ||
|
|
||
| type ServerId = string; | ||
| type KernelId = string; | ||
| type UriString = string; | ||
| type UriSessionUsedByResources = Record<ServerId, Record<KernelId, UriString[]>>; | ||
|
|
||
| /** | ||
| * Class to track the remote kernels that have been used by notebooks. | ||
| */ | ||
| @injectable() | ||
| export class LiveRemoteKernelConnectionUsageTracker | ||
| implements IExtensionSyncActivationService, ILiveRemoteKernelConnectionUsageTracker | ||
| { | ||
| private usedRemoteKernelServerIdsAndSessions: UriSessionUsedByResources = {}; | ||
| constructor( | ||
| @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, | ||
| @inject(IJupyterServerUriStorage) private readonly uriStorage: IJupyterServerUriStorage, | ||
| @inject(IMemento) @named(GLOBAL_MEMENTO) private readonly memento: Memento | ||
| ) {} | ||
| public activate(): void { | ||
| this.usedRemoteKernelServerIdsAndSessions = this.memento.get<UriSessionUsedByResources>( | ||
| mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, | ||
| {} | ||
| ); | ||
| this.uriStorage.onDidRemoveUri(this.onDidRemoveUri, this, this.disposables); | ||
| } | ||
|
|
||
| public wasKernelUsed(connection: LiveRemoteKernelConnectionMetadata) { | ||
| return ( | ||
| connection.serverId in this.usedRemoteKernelServerIdsAndSessions && | ||
| typeof connection.kernelModel.id === 'string' && | ||
| connection.kernelModel.id in this.usedRemoteKernelServerIdsAndSessions[connection.serverId] | ||
| ); | ||
| } | ||
| public trackKernelIdAsUsed(resource: Uri, serverId: string, kernelId: string) { | ||
| this.usedRemoteKernelServerIdsAndSessions[serverId] = this.usedRemoteKernelServerIdsAndSessions[serverId] || {}; | ||
| this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId] = | ||
| this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId] || []; | ||
| const uris = this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId]; | ||
| if (uris.includes(resource.toString())) { | ||
| return; | ||
| } | ||
| uris.push(resource.toString()); | ||
| this.memento | ||
| .update( | ||
| mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, | ||
| this.usedRemoteKernelServerIdsAndSessions | ||
| ) | ||
| .then(noop, noop); | ||
| } | ||
| public trackKernelIdAsNotUsed(resource: Uri, serverId: string, kernelId: string) { | ||
| if (!(serverId in this.usedRemoteKernelServerIdsAndSessions)) { | ||
| return; | ||
| } | ||
| if (!(kernelId in this.usedRemoteKernelServerIdsAndSessions[serverId])) { | ||
| return; | ||
| } | ||
| const uris = this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId]; | ||
| if (!Array.isArray(uris) || !uris.includes(resource.toString())) { | ||
| return; | ||
| } | ||
| uris.splice(uris.indexOf(resource.toString()), 1); | ||
| if (uris.length === 0) { | ||
| delete this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId]; | ||
| } | ||
| if (Object.keys(this.usedRemoteKernelServerIdsAndSessions[serverId]).length === 0) { | ||
| delete this.usedRemoteKernelServerIdsAndSessions[serverId]; | ||
| } | ||
|
|
||
| this.memento | ||
| .update( | ||
| mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, | ||
| this.usedRemoteKernelServerIdsAndSessions | ||
| ) | ||
| .then(noop, noop); | ||
| } | ||
| private onDidRemoveUri(uri: string) { | ||
| const serverId = computeServerId(uri); | ||
| delete this.usedRemoteKernelServerIdsAndSessions[serverId]; | ||
| this.memento | ||
| .update( | ||
| mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, | ||
| this.usedRemoteKernelServerIdsAndSessions | ||
| ) | ||
| .then(noop, noop); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| import { inject, injectable } from 'inversify'; | ||
| import { Disposable, NotebookDocument } from 'vscode'; | ||
| import { IVSCodeNotebookController } from '../../notebooks/controllers/types'; | ||
| import { INotebookControllerManager } from '../../notebooks/types'; | ||
| import { IExtensionSyncActivationService } from '../../platform/activation/types'; | ||
| import { IDisposableRegistry } from '../../platform/common/types'; | ||
| import { noop } from '../../platform/common/utils/misc'; | ||
| import { traceInfo } from '../../platform/logging'; | ||
| import { IKernel, IKernelProvider, isLocalConnection } from '../types'; | ||
| import { PreferredRemoteKernelIdProvider } from './preferredRemoteKernelIdProvider'; | ||
| import { ILiveRemoteKernelConnectionUsageTracker } from './types'; | ||
|
|
||
| @injectable() | ||
| export class RemoteKernelConnectionHandler implements IExtensionSyncActivationService { | ||
| constructor( | ||
| @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, | ||
| @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, | ||
| @inject(INotebookControllerManager) private readonly controllers: INotebookControllerManager, | ||
| @inject(ILiveRemoteKernelConnectionUsageTracker) | ||
| private readonly liveKernelTracker: ILiveRemoteKernelConnectionUsageTracker, | ||
| @inject(PreferredRemoteKernelIdProvider) | ||
| private readonly preferredRemoteKernelIdProvider: PreferredRemoteKernelIdProvider | ||
| ) {} | ||
| activate(): void { | ||
| this.kernelProvider.onDidStartKernel(this.onDidStartKernel, this, this.disposables); | ||
| this.controllers.onNotebookControllerSelectionChanged( | ||
| this.onNotebookControllerSelectionChanged, | ||
| this, | ||
| this.disposables | ||
| ); | ||
| } | ||
| private onNotebookControllerSelectionChanged({ | ||
| selected, | ||
| notebook, | ||
| controller | ||
| }: { | ||
| selected: boolean; | ||
| notebook: NotebookDocument; | ||
| controller: IVSCodeNotebookController; | ||
| }) { | ||
| if (controller.connection.kind === 'connectToLiveRemoteKernel' && controller.connection.kernelModel.id) { | ||
| if (selected) { | ||
| this.liveKernelTracker.trackKernelIdAsUsed( | ||
| notebook.uri, | ||
| controller.connection.serverId, | ||
| controller.connection.kernelModel.id | ||
| ); | ||
| } else { | ||
| this.liveKernelTracker.trackKernelIdAsNotUsed( | ||
| notebook.uri, | ||
| controller.connection.serverId, | ||
| controller.connection.kernelModel.id | ||
| ); | ||
| } | ||
| } | ||
| if (isLocalConnection(controller.connection)) { | ||
| this.preferredRemoteKernelIdProvider.clearPreferredRemoteKernelId(notebook.uri).catch(noop); | ||
| } | ||
| } | ||
| private onDidStartKernel(kernel: IKernel) { | ||
| if (kernel.creator !== 'jupyterExtension' || !kernel.resourceUri) { | ||
| return; | ||
| } | ||
| const resource = kernel.resourceUri; | ||
| if (kernel.kernelConnectionMetadata.kind === 'startUsingRemoteKernelSpec') { | ||
| const serverId = kernel.kernelConnectionMetadata.serverId; | ||
| const subscription = kernel.kernelSocket.subscribe((info) => { | ||
| const kernelId = info?.options.id; | ||
| if (!kernel.disposed && !kernel.disposing && kernelId) { | ||
| traceInfo(`Updating preferred kernel for remote notebook ${kernelId}`); | ||
| this.preferredRemoteKernelIdProvider.storePreferredRemoteKernelId(resource, kernelId).catch(noop); | ||
| this.liveKernelTracker.trackKernelIdAsUsed(resource, serverId, kernelId); | ||
| } | ||
| }); | ||
| this.disposables.push(new Disposable(() => subscription.unsubscribe())); | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.