diff --git a/packages/backend/package.json b/packages/backend/package.json index 1642a89f3..a1f0538cf 100644 --- a/packages/backend/package.json +++ b/packages/backend/package.json @@ -23,7 +23,7 @@ "views": { "icons/containersList": [ { - "when": "ia-studio-model in containerLabelKeys", + "when": "ai-studio-model-id in containerLabelKeys", "icon": "${brain-icon}" } ] diff --git a/packages/backend/src/managers/playground.spec.ts b/packages/backend/src/managers/playground.spec.ts index 21cc07bc0..1ea8444bc 100644 --- a/packages/backend/src/managers/playground.spec.ts +++ b/packages/backend/src/managers/playground.spec.ts @@ -20,6 +20,7 @@ import { beforeEach, expect, test, vi } from 'vitest'; import { PlayGroundManager } from './playground'; import type { ImageInfo, Webview } from '@podman-desktop/api'; import type { ContainerRegistry } from '../registries/ContainerRegistry'; +import type { PodmanConnection } from './podmanConnection'; const mocks = vi.hoisted(() => ({ postMessage: vi.fn(), @@ -64,6 +65,7 @@ beforeEach(() => { postMessage: mocks.postMessage, } as unknown as Webview, containerRegistryMock, + {} as PodmanConnection, ); }); @@ -123,7 +125,8 @@ test('startPlayground should download image if not present then create container }, Image: 'image1', Labels: { - 'ia-studio-model': 'model1', + 'ai-studio-model-id': 'model1', + 'ai-studio-model-port': '8085', }, }); }); diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts index 6c7de8f8e..e1c5a51b4 100644 --- a/packages/backend/src/managers/playground.ts +++ b/packages/backend/src/managers/playground.ts @@ -17,11 +17,11 @@ ***********************************************************************/ import { - provider, containerEngine, type Webview, - type ProviderContainerConnection, type ImageInfo, + type ProviderContainerConnection, + provider, } from '@podman-desktop/api'; import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo'; import type { ModelResponse } from '@shared/src/models/IModelResponse'; @@ -33,6 +33,10 @@ import type { QueryState } from '@shared/src/models/IPlaygroundQueryState'; import { MSG_NEW_PLAYGROUND_QUERIES_STATE, MSG_PLAYGROUNDS_STATE_UPDATE } from '@shared/Messages'; import type { PlaygroundState, PlaygroundStatus } from '@shared/src/models/IPlaygroundState'; import type { ContainerRegistry } from '../registries/ContainerRegistry'; +import type { PodmanConnection } from './podmanConnection'; + +const LABEL_MODEL_ID = 'ai-studio-model-id'; +const LABEL_MODEL_PORT = 'ai-studio-model-port'; // TODO: this should not be hardcoded const PLAYGROUND_IMAGE = 'quay.io/bootsy/playground:v0'; @@ -55,12 +59,48 @@ export class PlayGroundManager { constructor( private webview: Webview, private containerRegistry: ContainerRegistry, + private podmanConnection: PodmanConnection, ) { this.playgrounds = new Map(); this.queries = new Map(); } - async selectImage(connection: ProviderContainerConnection, image: string): Promise { + adoptRunningPlaygrounds() { + this.podmanConnection.startupSubscribe(() => { + containerEngine + .listContainers() + .then(containers => { + const playgroundContainers = containers.filter( + c => LABEL_MODEL_ID in c.Labels && LABEL_MODEL_PORT in c.Labels && c.State === 'running', + ); + for (const containerToAdopt of playgroundContainers) { + const modelId = containerToAdopt.Labels[LABEL_MODEL_ID]; + if (this.playgrounds.has(modelId)) { + continue; + } + const modelPort = parseInt(containerToAdopt.Labels[LABEL_MODEL_PORT], 10); + if (isNaN(modelPort)) { + continue; + } + const state: PlaygroundState = { + modelId, + status: 'running', + container: { + containerId: containerToAdopt.Id, + engineId: containerToAdopt.engineId, + port: modelPort, + }, + }; + this.updatePlaygroundState(modelId, state); + } + }) + .catch((err: unknown) => { + console.error('error during adoption of existing playground containers', err); + }); + }); + } + + async selectImage(image: string): Promise { const images = (await containerEngine.listImages()).filter(im => im.RepoTags?.some(tag => tag === image)); return images.length > 0 ? images[0] : undefined; } @@ -110,10 +150,10 @@ export class PlayGroundManager { throw new Error('Unable to find an engine to start playground'); } - let image = await this.selectImage(connection, PLAYGROUND_IMAGE); + let image = await this.selectImage(PLAYGROUND_IMAGE); if (!image) { await containerEngine.pullImage(connection.connection, PLAYGROUND_IMAGE, () => {}); - image = await this.selectImage(connection, PLAYGROUND_IMAGE); + image = await this.selectImage(PLAYGROUND_IMAGE); if (!image) { this.setPlaygroundStatus(modelId, 'error'); throw new Error(`Unable to find ${PLAYGROUND_IMAGE} image`); @@ -143,7 +183,8 @@ export class PlayGroundManager { }, }, Labels: { - 'ia-studio-model': modelId, + [LABEL_MODEL_ID]: modelId, + [LABEL_MODEL_PORT]: `${freePort}`, }, Env: [`MODEL_PATH=/models/${path.basename(modelPath)}`], Cmd: ['--models-path', '/models', '--context-size', '700', '--threads', '4'], diff --git a/packages/backend/src/managers/podmanConnection.ts b/packages/backend/src/managers/podmanConnection.ts new file mode 100644 index 000000000..305b433ff --- /dev/null +++ b/packages/backend/src/managers/podmanConnection.ts @@ -0,0 +1,65 @@ +/********************************************************************** + * Copyright (C) 2024 Red Hat, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + ***********************************************************************/ + +import { type RegisterContainerConnectionEvent, provider } from '@podman-desktop/api'; + +type startupHandle = () => void; + +export class PodmanConnection { + #firstFound = false; + #toExecuteAtStartup: startupHandle[] = []; + + init(): void { + // In case the extension has not yet registered, we listen for new registrations + // and retain the first started podman provider + const disposable = provider.onDidRegisterContainerConnection((e: RegisterContainerConnectionEvent) => { + if (e.connection.type !== 'podman' || e.connection.status() !== 'started') { + return; + } + if (this.#firstFound) { + return; + } + this.#firstFound = true; + for (const f of this.#toExecuteAtStartup) { + f(); + } + this.#toExecuteAtStartup = []; + disposable.dispose(); + }); + + // In case at least one extension has already registered, we get one started podman provider + const engines = provider + .getContainerConnections() + .filter(connection => connection.connection.type === 'podman') + .filter(connection => connection.connection.status() === 'started'); + if (engines.length > 0) { + disposable.dispose(); + this.#firstFound = true; + } + } + + // startupSubscribe registers f to be executed when a podman container provider + // registers, or immediately if already registered + startupSubscribe(f: startupHandle): void { + if (this.#firstFound) { + f(); + } else { + this.#toExecuteAtStartup.push(f); + } + } +} diff --git a/packages/backend/src/studio.spec.ts b/packages/backend/src/studio.spec.ts index bf7bb8e2a..75c58ba1d 100644 --- a/packages/backend/src/studio.spec.ts +++ b/packages/backend/src/studio.spec.ts @@ -32,6 +32,11 @@ const mockedExtensionContext = { const studio = new Studio(mockedExtensionContext); +const mocks = vi.hoisted(() => ({ + listContainers: vi.fn(), + getContainerConnections: vi.fn(), +})); + vi.mock('@podman-desktop/api', async () => { return { Uri: class { @@ -48,6 +53,11 @@ vi.mock('@podman-desktop/api', async () => { }, containerEngine: { onEvent: vi.fn(), + listContainers: mocks.listContainers, + }, + provider: { + onDidRegisterContainerConnection: vi.fn(), + getContainerConnections: mocks.getContainerConnections, }, }; }); @@ -66,6 +76,8 @@ afterEach(() => { }); test('check activate ', async () => { + mocks.listContainers.mockReturnValue([]); + mocks.getContainerConnections.mockReturnValue([]); vi.spyOn(fs.promises, 'readFile').mockImplementation(() => { return Promise.resolve(''); }); diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index 9b4563706..5ebe7194b 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -31,6 +31,7 @@ import path from 'node:path'; import os from 'os'; import fs from 'node:fs'; import { ContainerRegistry } from './registries/ContainerRegistry'; +import { PodmanConnection } from './managers/podmanConnection'; // TODO: Need to be configured export const AI_STUDIO_FOLDER = path.join('podman-desktop', 'ai-studio'); @@ -104,9 +105,11 @@ export class Studio { this.rpcExtension = new RpcExtension(this.#panel.webview); const gitManager = new GitManager(); + + const podmanConnection = new PodmanConnection(); const taskRegistry = new TaskRegistry(); const recipeStatusRegistry = new RecipeStatusRegistry(taskRegistry, this.#panel.webview); - this.playgroundManager = new PlayGroundManager(this.#panel.webview, containerRegistry); + this.playgroundManager = new PlayGroundManager(this.#panel.webview, containerRegistry, podmanConnection); // Create catalog manager, responsible for loading the catalog files and watching for changes this.catalogManager = new CatalogManager(appUserDirectory, this.#panel.webview); this.modelsManager = new ModelsManager(appUserDirectory, this.#panel.webview, this.catalogManager); @@ -128,6 +131,8 @@ export class Studio { await this.catalogManager.loadCatalog(); await this.modelsManager.loadLocalModels(); + podmanConnection.init(); + this.playgroundManager.adoptRunningPlaygrounds(); // Register the instance this.rpcExtension.registerInstance(StudioApiImpl, this.studioApi);