From ebddc334cf6ffddc2f403b9b1ce7d2bf5d49c21e Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Thu, 6 Jun 2024 17:48:15 +0200 Subject: [PATCH 1/3] feat: make extension activation very lazy Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- .../src/managers/applicationManager.spec.ts | 24 +- .../src/managers/applicationManager.ts | 2 +- .../src/managers/modelsManager.spec.ts | 8 +- .../backend/src/managers/modelsManager.ts | 4 + .../src/registries/ContainerRegistry.ts | 12 +- .../LocalRepositoryRegistry.spec.ts | 45 ++- .../src/registries/LocalRepositoryRegistry.ts | 15 +- packages/backend/src/studio.ts | 277 ++++++++++++------ packages/shared/src/messages/MessageProxy.ts | 11 +- 9 files changed, 269 insertions(+), 129 deletions(-) diff --git a/packages/backend/src/managers/applicationManager.spec.ts b/packages/backend/src/managers/applicationManager.spec.ts index 10a4214bd..de130d38c 100644 --- a/packages/backend/src/managers/applicationManager.spec.ts +++ b/packages/backend/src/managers/applicationManager.spec.ts @@ -1001,7 +1001,7 @@ describe('pod detection', async () => { ); }); - test('adoptRunningApplications updates the app state with the found pod', async () => { + test('init updates the app state with the found pod', async () => { vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([ { Labels: { @@ -1016,7 +1016,7 @@ describe('pod detection', async () => { f(); }); const updateApplicationStateSpy = vi.spyOn(manager, 'updateApplicationState'); - manager.adoptRunningApplications(); + manager.init(); await new Promise(resolve => setTimeout(resolve, 0)); expect(updateApplicationStateSpy).toHaveBeenNthCalledWith(1, 'recipe-id-1', 'model-id-1', { pod: { @@ -1037,13 +1037,13 @@ describe('pod detection', async () => { expect(ports).toStrictEqual([5000, 5001]); }); - test('adoptRunningApplications does not update the application state with the found pod without label', async () => { + test('init does not update the application state with the found pod without label', async () => { vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([{} as unknown as PodInfo]); mocks.startupSubscribeMock.mockImplementation((f: startupHandle) => { f(); }); const updateApplicationStateSpy = vi.spyOn(manager, 'updateApplicationState'); - manager.adoptRunningApplications(); + manager.init(); await new Promise(resolve => setTimeout(resolve, 0)); expect(updateApplicationStateSpy).not.toHaveBeenCalled(); }); @@ -1054,7 +1054,7 @@ describe('pod detection', async () => { f(); }); const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.adoptRunningApplications(); + manager.init(); expect(sendApplicationStateSpy).toHaveBeenCalledOnce(); }); @@ -1074,7 +1074,7 @@ describe('pod detection', async () => { return { dispose: vi.fn() }; }); const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.adoptRunningApplications(); + manager.init(); expect(sendApplicationStateSpy).toHaveBeenCalledOnce(); }); @@ -1090,7 +1090,7 @@ describe('pod detection', async () => { return { dispose: vi.fn() }; }); const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.adoptRunningApplications(); + manager.init(); expect(sendApplicationStateSpy).not.toHaveBeenCalledOnce(); }); @@ -1109,7 +1109,7 @@ describe('pod detection', async () => { return { dispose: vi.fn() }; }); const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.adoptRunningApplications(); + manager.init(); expect(sendApplicationStateSpy).not.toHaveBeenCalledOnce(); }); @@ -1142,7 +1142,7 @@ describe('pod detection', async () => { return { dispose: vi.fn() }; }); const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.adoptRunningApplications(); + manager.init(); await new Promise(resolve => setTimeout(resolve, 10)); expect(sendApplicationStateSpy).toHaveBeenCalledTimes(1); }); @@ -1170,7 +1170,7 @@ describe('pod detection', async () => { return { dispose: vi.fn() }; }); const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.adoptRunningApplications(); + manager.init(); await new Promise(resolve => setTimeout(resolve, 10)); expect(sendApplicationStateSpy).toHaveBeenCalledTimes(2); }); @@ -1226,7 +1226,7 @@ describe('pod detection', async () => { expect(podManager.removePod).toHaveBeenCalledWith('engine-1', 'pod-1'); }); - test('adoptRunningApplications should check pods health', async () => { + test('init should check pods health', async () => { vi.mocked(podManager.getHealth).mockResolvedValue('healthy'); vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([ { @@ -1255,7 +1255,7 @@ describe('pod detection', async () => { f(); }); vi.useFakeTimers(); - manager.adoptRunningApplications(); + manager.init(); await vi.advanceTimersByTimeAsync(1100); const state = manager.getApplicationsState(); expect(state).toHaveLength(1); diff --git a/packages/backend/src/managers/applicationManager.ts b/packages/backend/src/managers/applicationManager.ts index 6c073d979..afc7b8919 100644 --- a/packages/backend/src/managers/applicationManager.ts +++ b/packages/backend/src/managers/applicationManager.ts @@ -553,7 +553,7 @@ export class ApplicationManager extends Publisher implements } } - adoptRunningApplications() { + init() { this.podmanConnection.startupSubscribe(() => { this.podManager .getPodsWithLabels([LABEL_RECIPE_ID]) diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts index d7e9b3889..5acbc253d 100644 --- a/packages/backend/src/managers/modelsManager.spec.ts +++ b/packages/backend/src/managers/modelsManager.spec.ts @@ -456,9 +456,9 @@ test('deleteModel deletes the model folder', async () => { maxRetries: 3, }); } - expect(postMessageMock).toHaveBeenCalledTimes(3); + expect(postMessageMock).toHaveBeenCalledTimes(4); // check that a new state is sent with the model removed - expect(postMessageMock).toHaveBeenNthCalledWith(3, { + expect(postMessageMock).toHaveBeenNthCalledWith(4, { id: 'new-models-state', body: [ { @@ -522,9 +522,9 @@ describe('deleting models', () => { maxRetries: 3, }); } - expect(postMessageMock).toHaveBeenCalledTimes(3); + expect(postMessageMock).toHaveBeenCalledTimes(4); // check that a new state is sent with the model non removed - expect(postMessageMock).toHaveBeenNthCalledWith(3, { + expect(postMessageMock).toHaveBeenNthCalledWith(4, { id: 'new-models-state', body: [ { diff --git a/packages/backend/src/managers/modelsManager.ts b/packages/backend/src/managers/modelsManager.ts index 00ee3ac41..ce73868ec 100644 --- a/packages/backend/src/managers/modelsManager.ts +++ b/packages/backend/src/managers/modelsManager.ts @@ -63,6 +63,10 @@ export class ModelsManager implements Disposable { }); }); this.#disposables.push(disposable); + + this.loadLocalModels().catch((err: unknown) => { + console.error('Something went wrong while trying to load local models', err); + }); } dispose(): void { diff --git a/packages/backend/src/registries/ContainerRegistry.ts b/packages/backend/src/registries/ContainerRegistry.ts index 00c1e0288..33e5978bf 100644 --- a/packages/backend/src/registries/ContainerRegistry.ts +++ b/packages/backend/src/registries/ContainerRegistry.ts @@ -26,15 +26,17 @@ export interface ContainerStart { id: string; } -export class ContainerRegistry { +export class ContainerRegistry implements podmanDesktopApi.Disposable { private count: number = 0; private subscribers: Map = new Map(); private readonly _onStartContainerEvent = new podmanDesktopApi.EventEmitter(); readonly onStartContainerEvent: podmanDesktopApi.Event = this._onStartContainerEvent.event; - init(): podmanDesktopApi.Disposable { - return podmanDesktopApi.containerEngine.onEvent(event => { + #eventDisposable: podmanDesktopApi.Disposable | undefined; + + init(): void { + this.#eventDisposable = podmanDesktopApi.containerEngine.onEvent(event => { if (event.status === 'start') { this._onStartContainerEvent.fire({ id: event.id, @@ -52,6 +54,10 @@ export class ContainerRegistry { }); } + dispose(): void { + this.#eventDisposable?.dispose(); + } + subscribe(containerId: string, callback: (status: string) => void): podmanDesktopApi.Disposable { const subscriberId = ++this.count; const nSubs: Subscriber[] = [ diff --git a/packages/backend/src/registries/LocalRepositoryRegistry.spec.ts b/packages/backend/src/registries/LocalRepositoryRegistry.spec.ts index 71602e68d..09e0395d8 100644 --- a/packages/backend/src/registries/LocalRepositoryRegistry.spec.ts +++ b/packages/backend/src/registries/LocalRepositoryRegistry.spec.ts @@ -22,6 +22,7 @@ import type { Webview } from '@podman-desktop/api'; import type { Recipe } from '@shared/src/models/IRecipe'; import fs from 'node:fs'; import path from 'node:path'; +import type { CatalogManager } from '../managers/catalogManager'; const mocks = vi.hoisted(() => ({ DisposableCreateMock: vi.fn(), @@ -45,6 +46,11 @@ vi.mock('node:fs', () => { }; }); +const catalogManagerMock = { + onCatalogUpdate: vi.fn(), + getRecipes: vi.fn(), +} as unknown as CatalogManager; + beforeEach(() => { vi.resetAllMocks(); vi.mock('node:fs'); @@ -57,6 +63,7 @@ test('should not have any repositories by default', () => { postMessage: mocks.postMessageMock, } as unknown as Webview, '/appUserDirectory', + catalogManagerMock, ); expect(localRepositories.getLocalRepositories().length).toBe(0); }); @@ -67,6 +74,7 @@ test('should notify webview when register', () => { postMessage: mocks.postMessageMock, } as unknown as Webview, '/appUserDirectory', + catalogManagerMock, ); localRepositories.register({ path: 'random', sourcePath: 'random', labels: { 'recipe-id': 'random' } }); expect(mocks.postMessageMock).toHaveBeenNthCalledWith(1, { @@ -81,6 +89,7 @@ test('should notify webview when unregister', async () => { postMessage: mocks.postMessageMock, } as unknown as Webview, '/appUserDirectory', + catalogManagerMock, ); vi.spyOn(fs.promises, 'rm').mockResolvedValue(); localRepositories.register({ path: 'random', sourcePath: 'random', labels: { 'recipe-id': 'random' } }); @@ -94,18 +103,26 @@ test('should notify webview when unregister', async () => { test('should register localRepo if it find the folder of the recipe', () => { vi.spyOn(fs, 'existsSync').mockReturnValue(true); + vi.mocked(catalogManagerMock.onCatalogUpdate).mockImplementation((fn: () => void) => { + fn(); + return { dispose: vi.fn() }; + }); + vi.mocked(catalogManagerMock.getRecipes).mockReturnValue([ + { + id: 'recipe', + } as unknown as Recipe, + ]); + const localRepositories = new LocalRepositoryRegistry( { postMessage: mocks.postMessageMock, } as unknown as Webview, '/appUserDirectory', + catalogManagerMock, ); + const registerMock = vi.spyOn(localRepositories, 'register'); - localRepositories.init([ - { - id: 'recipe', - } as unknown as Recipe, - ]); + localRepositories.init(); const folder = path.join('/appUserDirectory', 'recipe'); expect(registerMock).toHaveBeenCalledWith({ @@ -117,18 +134,24 @@ test('should register localRepo if it find the folder of the recipe', () => { test('should NOT register localRepo if it does not find the folder of the recipe', () => { vi.spyOn(fs, 'existsSync').mockReturnValue(false); + vi.mocked(catalogManagerMock.onCatalogUpdate).mockImplementation((fn: () => void) => { + fn(); + return { dispose: vi.fn() }; + }); + vi.mocked(catalogManagerMock.getRecipes).mockReturnValue([ + { + id: 'recipe', + } as unknown as Recipe, + ]); + const localRepositories = new LocalRepositoryRegistry( { postMessage: mocks.postMessageMock, } as unknown as Webview, '/appUserDirectory', + catalogManagerMock, ); const registerMock = vi.spyOn(localRepositories, 'register'); - localRepositories.init([ - { - id: 'recipe', - } as unknown as Recipe, - ]); - + localRepositories.init(); expect(registerMock).not.toHaveBeenCalled(); }); diff --git a/packages/backend/src/registries/LocalRepositoryRegistry.ts b/packages/backend/src/registries/LocalRepositoryRegistry.ts index f9db5a2b7..c8d293779 100644 --- a/packages/backend/src/registries/LocalRepositoryRegistry.ts +++ b/packages/backend/src/registries/LocalRepositoryRegistry.ts @@ -22,23 +22,32 @@ import { Publisher } from '../utils/Publisher'; import type { Recipe } from '@shared/src/models/IRecipe'; import fs from 'node:fs'; import path from 'node:path'; +import type { CatalogManager } from '../managers/catalogManager'; /** * The LocalRepositoryRegistry is responsible for keeping track of the directories where recipe are cloned */ -export class LocalRepositoryRegistry extends Publisher { +export class LocalRepositoryRegistry extends Publisher implements Disposable { // Map path => LocalRepository private repositories: Map = new Map(); + #catalogEventDisposable: Disposable | undefined; constructor( webview: Webview, private appUserDirectory: string, + private catalogManager: CatalogManager, ) { super(webview, Messages.MSG_LOCAL_REPOSITORY_UPDATE, () => this.getLocalRepositories()); } - init(recipes: Recipe[]): void { - this.loadLocalRecipeRepositories(recipes); + dispose(): void { + this.#catalogEventDisposable?.dispose(); + } + + init(): void { + this.#catalogEventDisposable = this.catalogManager.onCatalogUpdate(() => { + this.loadLocalRecipeRepositories(this.catalogManager.getRecipes()); + }); } register(localRepository: LocalRepository): Disposable { diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index 9fb5c4c83..d8328c65e 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -48,16 +48,35 @@ import { InferenceProviderRegistry } from './registries/InferenceProviderRegistr export class Studio { readonly #extensionContext: ExtensionContext; + /** + * Webview panel used by AI Lab + */ #panel: WebviewPanel | undefined; - rpcExtension: RpcExtension | undefined; - studioApi: StudioApiImpl | undefined; - catalogManager: CatalogManager | undefined; - modelsManager: ModelsManager | undefined; - telemetry: TelemetryLogger | undefined; + /** + * API related classes + */ + #rpcExtension: RpcExtension | undefined; + #studioApi: StudioApiImpl | undefined; - #taskRegistry: TaskRegistry | undefined; + #localRepositoryRegistry: LocalRepositoryRegistry | undefined; + #catalogManager: CatalogManager | undefined; + #modelsManager: ModelsManager | undefined; + #telemetry: TelemetryLogger | undefined; #inferenceManager: InferenceManager | undefined; + #podManager: PodManager | undefined; + #builderManager: BuilderManager | undefined; + #containerRegistry: ContainerRegistry | undefined; + #podmanConnection: PodmanConnection | undefined; + #taskRegistry: TaskRegistry | undefined; + #cancellationTokenRegistry: CancellationTokenRegistry | undefined; + #snippetManager: SnippetManager | undefined; + #playgroundManager: PlaygroundV2Manager | undefined; + #applicationManager: ApplicationManager | undefined; + #inferenceProviderRegistry: InferenceProviderRegistry | undefined; + + #initialized: boolean = false; + #lazyInitialization: { init(): void }[] = []; constructor(readonly extensionContext: ExtensionContext) { this.#extensionContext = extensionContext; @@ -79,12 +98,16 @@ export class Studio { public async activate(): Promise { console.log('starting AI Lab extension'); - this.telemetry = env.createTelemetryLogger(); + this.#telemetry = env.createTelemetryLogger(); + /** + * Ensure the running version of podman is compatible with + * our minimum requirement + */ if (!this.checkVersion()) { const min = minVersion(engines['podman-desktop']) ?? { version: 'unknown' }; const current = version ?? 'unknown'; - this.telemetry.logError('start.incompatible', { + this.#telemetry.logError('start.incompatible', { version: current, message: `error activating extension on version below ${min.version}`, }); @@ -93,128 +116,198 @@ export class Studio { ); } - this.telemetry.logUsage('start'); + this.#telemetry.logUsage('start'); - // init webview + /** + * The AI Lab has a webview integrated in Podman Desktop + * We need to initialize and configure it properly + */ this.#panel = await initWebview(this.#extensionContext.extensionUri); this.#extensionContext.subscriptions.push(this.#panel); + this.#panel.onDidChangeViewState((e: WebviewPanelOnDidChangeViewStateEvent) => { + // Lazily init classes + if (!this.#initialized) { + this.#initialized = true; + this.lazyInit(); + } - // Creating cancellation token registry - const cancellationTokenRegistry = new CancellationTokenRegistry(); - this.#extensionContext.subscriptions.push(cancellationTokenRegistry); - - // Creating container registry - const containerRegistry = new ContainerRegistry(); - this.#extensionContext.subscriptions.push(containerRegistry.init()); + this.#telemetry?.logUsage(e.webviewPanel.visible ? 'opened' : 'closed'); + }); - const appUserDirectory = this.extensionContext.storagePath; + /** + * Cancellation token registry store the tokens used to cancel a task + */ + this.#cancellationTokenRegistry = new CancellationTokenRegistry(); + this.#extensionContext.subscriptions.push(this.#cancellationTokenRegistry); - this.rpcExtension = new RpcExtension(this.#panel.webview); - const gitManager = new GitManager(); + /** + * The container registry handle the events linked to containers (start, remove, die...) + */ + this.#containerRegistry = new ContainerRegistry(); + this.#lazyInitialization.push(this.#containerRegistry); + this.#extensionContext.subscriptions.push(this.#containerRegistry); - const podmanConnection = new PodmanConnection(); - this.#taskRegistry = new TaskRegistry(this.#panel.webview); + const appUserDirectory = this.extensionContext.storagePath; - // Init the inference provider registry - const inferenceProviderRegistry = new InferenceProviderRegistry(this.#panel.webview); - this.#extensionContext.subscriptions.push( - inferenceProviderRegistry.register(new LlamaCppPython(this.#taskRegistry)), - ); + /** + * The RpcExtension handle the communication channels between the frontend and the backend + */ + this.#rpcExtension = new RpcExtension(this.#panel.webview); + this.#lazyInitialization.push(this.#rpcExtension); + this.#extensionContext.subscriptions.push(this.#rpcExtension); - // Create catalog manager, responsible for loading the catalog files and watching for changes - this.catalogManager = new CatalogManager(this.#panel.webview, appUserDirectory); - this.catalogManager.init(); + /** + * GitManager is used for cloning, pulling etc. recipes repositories + */ + const gitManager = new GitManager(); - const builderManager = new BuilderManager(this.#taskRegistry); - this.#extensionContext.subscriptions.push(builderManager); + /** + * The podman connection class is responsible for podman machine events (start/stop) + */ + this.#podmanConnection = new PodmanConnection(); + this.#lazyInitialization.push(this.#podmanConnection); + this.#extensionContext.subscriptions.push(this.#podmanConnection); - const podManager = new PodManager(); - podManager.init(); - this.#extensionContext.subscriptions.push(podManager); + /** + * The task registry store the tasks + */ + this.#taskRegistry = new TaskRegistry(this.#panel.webview); - this.modelsManager = new ModelsManager( + /** + * Create catalog manager, responsible for loading the catalog files and watching for changes + */ + this.#catalogManager = new CatalogManager(this.#panel.webview, appUserDirectory); + this.#lazyInitialization.push(this.#catalogManager); + + /** + * The builder manager is handling the building tasks, create corresponding tasks + * through the task registry and cancellation. + */ + this.#builderManager = new BuilderManager(this.#taskRegistry); + this.#extensionContext.subscriptions.push(this.#builderManager); + + /** + * The pod manager is a class responsible for managing the Pods + */ + this.#podManager = new PodManager(); + this.#lazyInitialization.push(this.#podManager); + this.#extensionContext.subscriptions.push(this.#podManager); + + /** + * The ModelManager role is to download and + */ + this.#modelsManager = new ModelsManager( appUserDirectory, this.#panel.webview, - this.catalogManager, - this.telemetry, + this.#catalogManager, + this.#telemetry, this.#taskRegistry, - cancellationTokenRegistry, + this.#cancellationTokenRegistry, ); - this.modelsManager.init(); - const localRepositoryRegistry = new LocalRepositoryRegistry(this.#panel.webview, appUserDirectory); - localRepositoryRegistry.init(this.catalogManager.getRecipes()); - const applicationManager = new ApplicationManager( + this.#lazyInitialization.push(this.#modelsManager); + this.#extensionContext.subscriptions.push(this.#modelsManager); + + /** + * The LocalRepositoryRegistry store and watch for recipes repository locally and expose it. + */ + this.#localRepositoryRegistry = new LocalRepositoryRegistry( + this.#panel.webview, + appUserDirectory, + this.#catalogManager, + ); + this.#lazyInitialization.push(this.#localRepositoryRegistry); + this.#extensionContext.subscriptions.push(this.#localRepositoryRegistry); + + /** + * The application manager is managing the Recipes + */ + this.#applicationManager = new ApplicationManager( appUserDirectory, gitManager, this.#taskRegistry, this.#panel.webview, - podmanConnection, - this.catalogManager, - this.modelsManager, - this.telemetry, - localRepositoryRegistry, - builderManager, - podManager, + this.#podmanConnection, + this.#catalogManager, + this.#modelsManager, + this.#telemetry, + this.#localRepositoryRegistry, + this.#builderManager, + this.#podManager, + ); + this.#lazyInitialization.push(this.#applicationManager); + this.#extensionContext.subscriptions.push(this.#applicationManager); + + /** + * The Inference Provider registry stores all the InferenceProvider (aka backend) which + * can be used to create InferenceServers + */ + this.#inferenceProviderRegistry = new InferenceProviderRegistry(this.#panel.webview); + this.#extensionContext.subscriptions.push( + this.#inferenceProviderRegistry.register(new LlamaCppPython(this.#taskRegistry)), ); + /** + * The inference manager create, stop, manage Inference servers + */ this.#inferenceManager = new InferenceManager( this.#panel.webview, - containerRegistry, - podmanConnection, - this.modelsManager, - this.telemetry, + this.#containerRegistry, + this.#podmanConnection, + this.#modelsManager, + this.#telemetry, this.#taskRegistry, - inferenceProviderRegistry, + this.#inferenceProviderRegistry, ); + this.#lazyInitialization.push(this.#inferenceManager); + this.#extensionContext.subscriptions.push(this.#inferenceManager); - this.#panel.onDidChangeViewState((e: WebviewPanelOnDidChangeViewStateEvent) => { - // Lazily init inference manager - if (this.#inferenceManager && !this.#inferenceManager.isInitialize()) { - this.#inferenceManager.init(); - this.#extensionContext.subscriptions.push(this.#inferenceManager); - } - - this.telemetry?.logUsage(e.webviewPanel.visible ? 'opened' : 'closed'); - }); - - const playgroundV2 = new PlaygroundV2Manager( + /** + * PlaygroundV2Manager handle the conversations of the Playground by using the InferenceServer available + */ + this.#playgroundManager = new PlaygroundV2Manager( this.#panel.webview, this.#inferenceManager, this.#taskRegistry, - this.telemetry, + this.#telemetry, ); - - const snippetManager = new SnippetManager(this.#panel.webview, this.telemetry); - snippetManager.init(); - - // Creating StudioApiImpl - this.studioApi = new StudioApiImpl( - applicationManager, - this.catalogManager, - this.modelsManager, - this.telemetry, - localRepositoryRegistry, + this.#extensionContext.subscriptions.push(this.#playgroundManager); + + /** + * The snippet manager provide code snippet used in the + * InferenceServer details page + */ + this.#snippetManager = new SnippetManager(this.#panel.webview, this.#telemetry); + this.#lazyInitialization.push(this.#snippetManager); + + /** + * The StudioApiImpl is the implementation of our API between backend and frontend + */ + this.#studioApi = new StudioApiImpl( + this.#applicationManager, + this.#catalogManager, + this.#modelsManager, + this.#telemetry, + this.#localRepositoryRegistry, this.#taskRegistry, this.#inferenceManager, - playgroundV2, - snippetManager, - cancellationTokenRegistry, + this.#playgroundManager, + this.#snippetManager, + this.#cancellationTokenRegistry, ); - - await this.modelsManager.loadLocalModels(); - podmanConnection.init(); - applicationManager.adoptRunningApplications(); - this.#extensionContext.subscriptions.push(applicationManager); - // Register the instance - this.rpcExtension.registerInstance(StudioApiImpl, this.studioApi); - this.#extensionContext.subscriptions.push(this.catalogManager); - this.#extensionContext.subscriptions.push(this.modelsManager); - this.#extensionContext.subscriptions.push(podmanConnection); + this.#rpcExtension.registerInstance(StudioApiImpl, this.#studioApi); + } + + /** + * In an effort to limit the IO, lazy init as much classes as possible + * @private + */ + private lazyInit(): void { + this.#lazyInitialization.forEach(item => item.init()); } public async deactivate(): Promise { console.log('stopping AI Lab extension'); - this.telemetry?.logUsage('stop'); + this.#telemetry?.logUsage('stop'); } } diff --git a/packages/shared/src/messages/MessageProxy.ts b/packages/shared/src/messages/MessageProxy.ts index 30c114883..2408d3b73 100644 --- a/packages/shared/src/messages/MessageProxy.ts +++ b/packages/shared/src/messages/MessageProxy.ts @@ -17,7 +17,7 @@ ***********************************************************************/ /* eslint-disable @typescript-eslint/no-explicit-any */ -import type { Webview } from '@podman-desktop/api'; +import type { Webview, Disposable } from '@podman-desktop/api'; import { noTimeoutChannels } from './NoTimeoutChannels'; export interface IMessage { @@ -51,15 +51,20 @@ export function isMessageResponse(content: unknown): content is IMessageResponse return isMessageRequest(content) && 'status' in content; } -export class RpcExtension { +export class RpcExtension implements Disposable { + #webviewDisposable: Disposable | undefined; methods: Map Promise> = new Map(); constructor(private webview: Webview) { this.init(); } + dispose(): void { + this.#webviewDisposable?.dispose(); + } + init() { - this.webview.onDidReceiveMessage(async (message: unknown) => { + this.#webviewDisposable = this.webview.onDidReceiveMessage(async (message: unknown) => { if (!isMessageRequest(message)) { console.error('Received incompatible message.', message); return; From b4b89f5d4a8cd9475856fafcd0270a277bc70b5a Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Fri, 7 Jun 2024 09:17:33 +0200 Subject: [PATCH 2/3] fix: typecheck Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/studio-api-impl.spec.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/backend/src/studio-api-impl.spec.ts b/packages/backend/src/studio-api-impl.spec.ts index e3cd5ec53..fa9064f6f 100644 --- a/packages/backend/src/studio-api-impl.spec.ts +++ b/packages/backend/src/studio-api-impl.spec.ts @@ -117,6 +117,7 @@ beforeEach(async () => { postMessage: vi.fn().mockResolvedValue(undefined), } as unknown as Webview, appUserDirectory, + {} as unknown as CatalogManager, ); const telemetryMock = { From 518931995fcb26e873e0903363638753ab367686 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Mon, 10 Jun 2024 14:08:40 +0200 Subject: [PATCH 3/3] revert: lazy initialization Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/studio.ts | 37 +++++++++------------------------- 1 file changed, 10 insertions(+), 27 deletions(-) diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index d8328c65e..d915668ff 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -75,9 +75,6 @@ export class Studio { #applicationManager: ApplicationManager | undefined; #inferenceProviderRegistry: InferenceProviderRegistry | undefined; - #initialized: boolean = false; - #lazyInitialization: { init(): void }[] = []; - constructor(readonly extensionContext: ExtensionContext) { this.#extensionContext = extensionContext; } @@ -125,12 +122,6 @@ export class Studio { this.#panel = await initWebview(this.#extensionContext.extensionUri); this.#extensionContext.subscriptions.push(this.#panel); this.#panel.onDidChangeViewState((e: WebviewPanelOnDidChangeViewStateEvent) => { - // Lazily init classes - if (!this.#initialized) { - this.#initialized = true; - this.lazyInit(); - } - this.#telemetry?.logUsage(e.webviewPanel.visible ? 'opened' : 'closed'); }); @@ -144,7 +135,7 @@ export class Studio { * The container registry handle the events linked to containers (start, remove, die...) */ this.#containerRegistry = new ContainerRegistry(); - this.#lazyInitialization.push(this.#containerRegistry); + this.#containerRegistry.init(); this.#extensionContext.subscriptions.push(this.#containerRegistry); const appUserDirectory = this.extensionContext.storagePath; @@ -153,7 +144,7 @@ export class Studio { * The RpcExtension handle the communication channels between the frontend and the backend */ this.#rpcExtension = new RpcExtension(this.#panel.webview); - this.#lazyInitialization.push(this.#rpcExtension); + this.#rpcExtension.init(); this.#extensionContext.subscriptions.push(this.#rpcExtension); /** @@ -165,7 +156,7 @@ export class Studio { * The podman connection class is responsible for podman machine events (start/stop) */ this.#podmanConnection = new PodmanConnection(); - this.#lazyInitialization.push(this.#podmanConnection); + this.#podmanConnection.init(); this.#extensionContext.subscriptions.push(this.#podmanConnection); /** @@ -177,7 +168,7 @@ export class Studio { * Create catalog manager, responsible for loading the catalog files and watching for changes */ this.#catalogManager = new CatalogManager(this.#panel.webview, appUserDirectory); - this.#lazyInitialization.push(this.#catalogManager); + this.#catalogManager.init(); /** * The builder manager is handling the building tasks, create corresponding tasks @@ -190,7 +181,7 @@ export class Studio { * The pod manager is a class responsible for managing the Pods */ this.#podManager = new PodManager(); - this.#lazyInitialization.push(this.#podManager); + this.#podManager.init(); this.#extensionContext.subscriptions.push(this.#podManager); /** @@ -204,7 +195,7 @@ export class Studio { this.#taskRegistry, this.#cancellationTokenRegistry, ); - this.#lazyInitialization.push(this.#modelsManager); + this.#modelsManager.init(); this.#extensionContext.subscriptions.push(this.#modelsManager); /** @@ -215,7 +206,7 @@ export class Studio { appUserDirectory, this.#catalogManager, ); - this.#lazyInitialization.push(this.#localRepositoryRegistry); + this.#localRepositoryRegistry.init(); this.#extensionContext.subscriptions.push(this.#localRepositoryRegistry); /** @@ -234,7 +225,7 @@ export class Studio { this.#builderManager, this.#podManager, ); - this.#lazyInitialization.push(this.#applicationManager); + this.#applicationManager.init(); this.#extensionContext.subscriptions.push(this.#applicationManager); /** @@ -258,7 +249,7 @@ export class Studio { this.#taskRegistry, this.#inferenceProviderRegistry, ); - this.#lazyInitialization.push(this.#inferenceManager); + this.#inferenceManager.init(); this.#extensionContext.subscriptions.push(this.#inferenceManager); /** @@ -277,7 +268,7 @@ export class Studio { * InferenceServer details page */ this.#snippetManager = new SnippetManager(this.#panel.webview, this.#telemetry); - this.#lazyInitialization.push(this.#snippetManager); + this.#snippetManager.init(); /** * The StudioApiImpl is the implementation of our API between backend and frontend @@ -298,14 +289,6 @@ export class Studio { this.#rpcExtension.registerInstance(StudioApiImpl, this.#studioApi); } - /** - * In an effort to limit the IO, lazy init as much classes as possible - * @private - */ - private lazyInit(): void { - this.#lazyInitialization.forEach(item => item.init()); - } - public async deactivate(): Promise { console.log('stopping AI Lab extension'); this.#telemetry?.logUsage('stop');