From 0d8cc009ce3bac9323b414ac55bb16e614bff526 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Thu, 18 Jul 2024 15:50:46 +0200 Subject: [PATCH] feat(backend): split application manager (#1396) Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- .../application/applicationManager.spec.ts | 367 +++++ .../{ => application}/applicationManager.ts | 323 ++--- .../src/managers/applicationManager.spec.ts | 1277 ----------------- .../src/managers/recipes/BuilderManager.ts | 30 +- .../managers/recipes/RecipeManager.spec.ts | 207 +++ .../src/managers/recipes/RecipeManager.ts | 194 +++ packages/backend/src/studio-api-impl.spec.ts | 4 +- packages/backend/src/studio-api-impl.ts | 6 +- packages/backend/src/studio.ts | 23 +- packages/backend/src/utils/RecipeConstants.ts | 31 + packages/shared/src/models/IRecipe.ts | 11 + 11 files changed, 943 insertions(+), 1530 deletions(-) create mode 100644 packages/backend/src/managers/application/applicationManager.spec.ts rename packages/backend/src/managers/{ => application}/applicationManager.ts (68%) delete mode 100644 packages/backend/src/managers/applicationManager.spec.ts create mode 100644 packages/backend/src/managers/recipes/RecipeManager.spec.ts create mode 100644 packages/backend/src/managers/recipes/RecipeManager.ts create mode 100644 packages/backend/src/utils/RecipeConstants.ts diff --git a/packages/backend/src/managers/application/applicationManager.spec.ts b/packages/backend/src/managers/application/applicationManager.spec.ts new file mode 100644 index 000000000..6ec069e53 --- /dev/null +++ b/packages/backend/src/managers/application/applicationManager.spec.ts @@ -0,0 +1,367 @@ +/********************************************************************** + * 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 { beforeEach, describe, expect, test, vi } from 'vitest'; +import type { TaskRegistry } from '../../registries/TaskRegistry'; +import type { PodInfo, TelemetryLogger, Webview } from '@podman-desktop/api'; +import { containerEngine, window } from '@podman-desktop/api'; +import type { PodmanConnection } from '../podmanConnection'; +import type { CatalogManager } from '../catalogManager'; +import type { ModelsManager } from '../modelsManager'; +import type { PodManager } from '../recipes/PodManager'; +import type { RecipeManager } from '../recipes/RecipeManager'; +import { ApplicationManager } from './applicationManager'; +import type { Recipe, RecipeImage } from '@shared/src/models/IRecipe'; +import type { ModelInfo } from '@shared/src/models/IModelInfo'; +import { VMType } from '@shared/src/models/IPodman'; +import { POD_LABEL_MODEL_ID, POD_LABEL_RECIPE_ID } from '../../utils/RecipeConstants'; + +const taskRegistryMock = { + createTask: vi.fn(), + updateTask: vi.fn(), + deleteByLabels: vi.fn(), +} as unknown as TaskRegistry; + +const webviewMock = { + postMessage: vi.fn(), +} as unknown as Webview; + +const podmanConnectionMock = { + startupSubscribe: vi.fn(), + onMachineStop: vi.fn(), + getVMType: vi.fn(), +} as unknown as PodmanConnection; + +const catalogManagerMock = {} as unknown as CatalogManager; + +const modelsManagerMock = { + requestDownloadModel: vi.fn(), + uploadModelToPodmanMachine: vi.fn(), +} as unknown as ModelsManager; + +const telemetryMock = { + logError: vi.fn(), + logUsage: vi.fn(), +} as unknown as TelemetryLogger; + +const podManager = { + onStartPodEvent: vi.fn(), + onRemovePodEvent: vi.fn(), + getPodsWithLabels: vi.fn(), + createPod: vi.fn(), + getPod: vi.fn(), + findPodByLabelsValues: vi.fn(), + startPod: vi.fn(), + stopPod: vi.fn(), + removePod: vi.fn(), +} as unknown as PodManager; + +const recipeManager = { + cloneRecipe: vi.fn(), + buildRecipe: vi.fn(), +} as unknown as RecipeManager; + +vi.mock('@podman-desktop/api', () => ({ + window: { + withProgress: vi.fn(), + }, + ProgressLocation: { + TASK_WIDGET: 'task-widget', + }, + provider: { + getContainerConnections: vi.fn(), + }, + containerEngine: { + createContainer: vi.fn(), + }, + Disposable: { + create: vi.fn(), + }, +})); + +const recipeMock: Recipe = { + id: 'recipe-test', + name: 'Test Recipe', + categories: [], + description: 'test recipe description', + repository: 'http://test-repository.test', + readme: 'test recipe readme', +}; + +const remoteModelMock: ModelInfo = { + id: 'model-test', + name: 'Test Model', + description: 'test model description', + hw: 'cpu', + url: 'http://test-repository.test', +}; + +const recipeImageInfoMock: RecipeImage = { + name: 'test recipe image info', + id: 'test-recipe-image-info', + appName: 'test-app-name', + engineId: 'test-engine-id', + ports: [], + modelService: false, + recipeId: recipeMock.id, +}; + +beforeEach(() => { + vi.resetAllMocks(); + + vi.mocked(webviewMock.postMessage).mockResolvedValue(true); + vi.mocked(podmanConnectionMock.getVMType).mockResolvedValue(VMType.WSL); + vi.mocked(recipeManager.buildRecipe).mockResolvedValue([recipeImageInfoMock]); + vi.mocked(podManager.createPod).mockResolvedValue({ engineId: 'test-engine-id', Id: 'test-pod-id' }); + vi.mocked(podManager.getPod).mockResolvedValue({ engineId: 'test-engine-id', Id: 'test-pod-id' } as PodInfo); + vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([]); + vi.mocked(taskRegistryMock.createTask).mockImplementation((name, state, labels) => ({ + name, + state, + labels, + id: 'fake-task', + })); + vi.mocked(modelsManagerMock.uploadModelToPodmanMachine).mockResolvedValue('downloaded-model-path'); +}); + +function getInitializedApplicationManager(): ApplicationManager { + const manager = new ApplicationManager( + taskRegistryMock, + webviewMock, + podmanConnectionMock, + catalogManagerMock, + modelsManagerMock, + telemetryMock, + podManager, + recipeManager, + ); + + manager.init(); + return manager; +} + +describe('requestPullApplication', () => { + test('task should be set to error if pull application raise an error', async () => { + vi.mocked(window.withProgress).mockRejectedValue(new Error('pull application error')); + const trackingId = await getInitializedApplicationManager().requestPullApplication(recipeMock, remoteModelMock); + + // ensure the task is created + await vi.waitFor(() => { + expect(taskRegistryMock.createTask).toHaveBeenCalledWith(`Pulling ${recipeMock.name} recipe`, 'loading', { + trackingId: trackingId, + 'recipe-pulling': recipeMock.id, + }); + }); + + // ensure the task is updated + await vi.waitFor(() => { + expect(taskRegistryMock.updateTask).toHaveBeenCalledWith( + expect.objectContaining({ + state: 'error', + }), + ); + }); + }); +}); + +describe('stopApplication', () => { + test('calling stop with exited pod should not create task', async () => { + vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ + engineId: 'test-engine-id', + Id: 'test-pod-id-existing', + Labels: { + [POD_LABEL_MODEL_ID]: remoteModelMock.id, + [POD_LABEL_RECIPE_ID]: recipeMock.id, + }, + Status: 'Exited', + } as unknown as PodInfo); + + await getInitializedApplicationManager().stopApplication(recipeMock.id, remoteModelMock.id); + + expect(taskRegistryMock.createTask).not.toHaveBeenCalled(); + expect(podManager.stopPod).not.toHaveBeenCalled(); + }); + + test('calling stop application with running pod should create stop task ', async () => { + vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ + engineId: 'test-engine-id', + Id: 'test-pod-id-existing', + Labels: { + [POD_LABEL_MODEL_ID]: remoteModelMock.id, + [POD_LABEL_RECIPE_ID]: recipeMock.id, + }, + Status: 'Running', + } as unknown as PodInfo); + + await getInitializedApplicationManager().stopApplication(recipeMock.id, remoteModelMock.id); + + expect(taskRegistryMock.createTask).toHaveBeenCalledWith('Stopping AI App', 'loading', { + 'recipe-id': recipeMock.id, + 'model-id': remoteModelMock.id, + }); + expect(podManager.stopPod).toHaveBeenCalledWith('test-engine-id', 'test-pod-id-existing'); + }); + + test('error raised should make the task as failed', async () => { + vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ + engineId: 'test-engine-id', + Id: 'test-pod-id-existing', + Labels: { + [POD_LABEL_MODEL_ID]: remoteModelMock.id, + [POD_LABEL_RECIPE_ID]: recipeMock.id, + }, + Status: 'Running', + } as unknown as PodInfo); + + vi.mocked(podManager.stopPod).mockRejectedValue(new Error('stop pod error')); + + await expect(() => { + return getInitializedApplicationManager().stopApplication(recipeMock.id, remoteModelMock.id); + }).rejects.toThrowError('stop pod error'); + + expect(taskRegistryMock.updateTask).toHaveBeenCalledWith( + expect.objectContaining({ + state: 'error', + }), + ); + }); +}); + +describe('startApplication', () => { + test('expect startPod in podManager to be properly called', async () => { + vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ + engineId: 'test-engine-id', + Id: 'test-pod-id-existing', + Labels: { + [POD_LABEL_MODEL_ID]: remoteModelMock.id, + [POD_LABEL_RECIPE_ID]: recipeMock.id, + }, + Status: 'Exited', + } as unknown as PodInfo); + + await getInitializedApplicationManager().startApplication(recipeMock.id, remoteModelMock.id); + + expect(podManager.startPod).toHaveBeenCalledWith('test-engine-id', 'test-pod-id-existing'); + }); + + test('error raised should make the task as failed', async () => { + vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ + engineId: 'test-engine-id', + Id: 'test-pod-id-existing', + Labels: { + [POD_LABEL_MODEL_ID]: remoteModelMock.id, + [POD_LABEL_RECIPE_ID]: recipeMock.id, + }, + Status: 'Exited', + } as unknown as PodInfo); + + vi.mocked(podManager.startPod).mockRejectedValue(new Error('start pod error')); + + await expect(() => { + return getInitializedApplicationManager().startApplication(recipeMock.id, remoteModelMock.id); + }).rejects.toThrowError('start pod error'); + + expect(taskRegistryMock.updateTask).toHaveBeenCalledWith( + expect.objectContaining({ + state: 'error', + }), + ); + }); +}); + +describe('pullApplication', () => { + test('labels should be propagated', async () => { + await getInitializedApplicationManager().pullApplication(recipeMock, remoteModelMock, { + 'test-label': 'test-value', + }); + + // clone the recipe + expect(recipeManager.cloneRecipe).toHaveBeenCalledWith(recipeMock, { + 'test-label': 'test-value', + 'model-id': remoteModelMock.id, + }); + // download model + expect(modelsManagerMock.requestDownloadModel).toHaveBeenCalledWith(remoteModelMock, { + 'test-label': 'test-value', + 'recipe-id': recipeMock.id, + 'model-id': remoteModelMock.id, + }); + // upload model to podman machine + expect(modelsManagerMock.uploadModelToPodmanMachine).toHaveBeenCalledWith(remoteModelMock, { + 'test-label': 'test-value', + 'recipe-id': recipeMock.id, + 'model-id': remoteModelMock.id, + }); + // build the recipe + expect(recipeManager.buildRecipe).toHaveBeenCalledWith(recipeMock, { + 'test-label': 'test-value', + 'recipe-id': recipeMock.id, + 'model-id': remoteModelMock.id, + }); + // create AI App task must be created + expect(taskRegistryMock.createTask).toHaveBeenCalledWith('Creating AI App', 'loading', { + 'test-label': 'test-value', + 'recipe-id': recipeMock.id, + 'model-id': remoteModelMock.id, + }); + + // a pod must have been created + expect(podManager.createPod).toHaveBeenCalledWith({ + name: expect.any(String), + portmappings: [], + labels: { + [POD_LABEL_MODEL_ID]: remoteModelMock.id, + [POD_LABEL_RECIPE_ID]: recipeMock.id, + }, + }); + + expect(containerEngine.createContainer).toHaveBeenCalledWith('test-engine-id', { + Image: recipeImageInfoMock.id, + name: expect.any(String), + Env: [], + HealthCheck: undefined, + HostConfig: undefined, + Detach: true, + pod: 'test-pod-id', + start: false, + }); + + // finally the pod must be started + expect(podManager.startPod).toHaveBeenCalledWith('test-engine-id', 'test-pod-id'); + }); + + test('existing application should be removed', async () => { + vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ + engineId: 'test-engine-id', + Id: 'test-pod-id-existing', + Labels: { + [POD_LABEL_MODEL_ID]: remoteModelMock.id, + [POD_LABEL_RECIPE_ID]: recipeMock.id, + }, + } as unknown as PodInfo); + + await getInitializedApplicationManager().pullApplication(recipeMock, remoteModelMock); + + // removing existing application should create a task to notify the user + expect(taskRegistryMock.createTask).toHaveBeenCalledWith('Removing AI App', 'loading', { + 'recipe-id': recipeMock.id, + 'model-id': remoteModelMock.id, + }); + // the remove pod should have been called + expect(podManager.removePod).toHaveBeenCalledWith('test-engine-id', 'test-pod-id-existing'); + }); +}); diff --git a/packages/backend/src/managers/applicationManager.ts b/packages/backend/src/managers/application/applicationManager.ts similarity index 68% rename from packages/backend/src/managers/applicationManager.ts rename to packages/backend/src/managers/application/applicationManager.ts index b746eb470..b078f4609 100644 --- a/packages/backend/src/managers/applicationManager.ts +++ b/packages/backend/src/managers/application/applicationManager.ts @@ -16,80 +16,56 @@ * SPDX-License-Identifier: Apache-2.0 ***********************************************************************/ -import type { Recipe } from '@shared/src/models/IRecipe'; -import type { GitCloneInfo, GitManager } from './gitManager'; -import fs from 'fs'; +import type { Recipe, RecipeImage } from '@shared/src/models/IRecipe'; import * as path from 'node:path'; +import { containerEngine, Disposable, window, ProgressLocation } from '@podman-desktop/api'; import type { - HealthConfig, - HostConfig, - PodContainerInfo, PodCreatePortOptions, - PodInfo, TelemetryLogger, + PodInfo, Webview, + HostConfig, + HealthConfig, + PodContainerInfo, } from '@podman-desktop/api'; -import { containerEngine, Disposable, window, ProgressLocation } from '@podman-desktop/api'; -import type { AIConfig, AIConfigFile, ContainerConfig } from '../models/AIConfig'; -import { parseYamlFile } from '../models/AIConfig'; -import type { Task } from '@shared/src/models/ITask'; import type { ModelInfo } from '@shared/src/models/IModelInfo'; -import type { ModelsManager } from './modelsManager'; -import { getPortsFromLabel, getPortsInfo } from '../utils/ports'; -import { goarch } from '../utils/arch'; -import { getDurationSecondsSince, timeout } from '../utils/utils'; -import type { LocalRepositoryRegistry } from '../registries/LocalRepositoryRegistry'; +import type { ModelsManager } from '../modelsManager'; +import { getPortsFromLabel, getPortsInfo } from '../../utils/ports'; +import { getDurationSecondsSince, timeout } from '../../utils/utils'; import type { ApplicationState } from '@shared/src/models/IApplicationState'; -import type { PodmanConnection } from './podmanConnection'; +import type { PodmanConnection } from '../podmanConnection'; import { Messages } from '@shared/Messages'; -import type { CatalogManager } from './catalogManager'; -import { ApplicationRegistry } from '../registries/ApplicationRegistry'; -import type { TaskRegistry } from '../registries/TaskRegistry'; -import { Publisher } from '../utils/Publisher'; -import { getModelPropertiesForEnvironment } from '../utils/modelsUtils'; -import { getRandomName, getRandomString } from '../utils/randomUtils'; -import type { BuilderManager } from './recipes/BuilderManager'; -import type { PodManager } from './recipes/PodManager'; -import { SECOND } from '../workers/provider/LlamaCppPython'; +import type { CatalogManager } from '../catalogManager'; +import { ApplicationRegistry } from '../../registries/ApplicationRegistry'; +import type { TaskRegistry } from '../../registries/TaskRegistry'; +import { Publisher } from '../../utils/Publisher'; +import { getModelPropertiesForEnvironment } from '../../utils/modelsUtils'; +import { getRandomName, getRandomString } from '../../utils/randomUtils'; +import type { PodManager } from '../recipes/PodManager'; +import { SECOND } from '../../workers/provider/LlamaCppPython'; +import type { RecipeManager } from '../recipes/RecipeManager'; +import { + POD_LABEL_APP_PORTS, + POD_LABEL_MODEL_ID, + POD_LABEL_MODEL_PORTS, + POD_LABEL_RECIPE_ID, +} from '../../utils/RecipeConstants'; import { VMType } from '@shared/src/models/IPodman'; -export const LABEL_MODEL_ID = 'ai-lab-model-id'; -export const LABEL_MODEL_PORTS = 'ai-lab-model-ports'; - -export const LABEL_RECIPE_ID = 'ai-lab-recipe-id'; -export const LABEL_APP_PORTS = 'ai-lab-app-ports'; - -export const CONFIG_FILENAME = 'ai-lab.yaml'; - -interface AIContainers { - aiConfigFile: AIConfigFile; - containers: ContainerConfig[]; -} - -export interface ImageInfo { - id: string; - modelService: boolean; - ports: string[]; - appName: string; -} - export class ApplicationManager extends Publisher implements Disposable { #applications: ApplicationRegistry; protectTasks: Set = new Set(); #disposables: Disposable[]; constructor( - private appUserDirectory: string, - private git: GitManager, private taskRegistry: TaskRegistry, webview: Webview, private podmanConnection: PodmanConnection, private catalogManager: CatalogManager, private modelsManager: ModelsManager, private telemetry: TelemetryLogger, - private localRepositories: LocalRepositoryRegistry, - private builderManager: BuilderManager, private podManager: PodManager, + private recipeManager: RecipeManager, ) { super(webview, Messages.MSG_APPLICATIONS_STATE_UPDATE, () => this.getApplicationsState()); this.#applications = new ApplicationRegistry(); @@ -161,29 +137,6 @@ export class ApplicationManager extends Publisher implements } } - public async cloneApplication(recipe: Recipe, labels?: { [key: string]: string }): Promise { - const localFolder = path.join(this.appUserDirectory, recipe.id); - - // clone the recipe repository on the local folder - const gitCloneInfo: GitCloneInfo = { - repository: recipe.repository, - ref: recipe.ref, - targetDirectory: localFolder, - }; - await this.doCheckout(gitCloneInfo, { - ...labels, - 'recipe-id': recipe.id, - }); - - this.localRepositories.register({ - path: gitCloneInfo.targetDirectory, - sourcePath: path.join(gitCloneInfo.targetDirectory, recipe.basedir ?? ''), - labels: { - 'recipe-id': recipe.id, - }, - }); - } - /** * This method will execute the following tasks * - git clone @@ -204,14 +157,8 @@ export class ApplicationManager extends Publisher implements model: ModelInfo, labels: Record = {}, ): Promise { - const localFolder = path.join(this.appUserDirectory, recipe.id); - - // clone the application - await this.cloneApplication(recipe, { ...labels, 'model-id': model.id }); - - // load and parse the recipe configuration file and filter containers based on architecture, gpu accelerator - // and backend (that define which model supports) - const configAndFilteredContainers = this.getConfigAndFilterContainers(recipe.basedir, localFolder); + // clone the recipe + await this.recipeManager.cloneRecipe(recipe, { ...labels, 'model-id': model.id }); // get model by downloading it or retrieving locally await this.modelsManager.requestDownloadModel(model, { @@ -228,16 +175,11 @@ export class ApplicationManager extends Publisher implements }); // build all images, one per container (for a basic sample we should have 2 containers = sample app + model service) - const images = await this.builderManager.build( - recipe, - configAndFilteredContainers.containers, - configAndFilteredContainers.aiConfigFile.path, - { - ...labels, - 'recipe-id': recipe.id, - 'model-id': model.id, - }, - ); + const images = await this.recipeManager.buildRecipe(recipe, { + ...labels, + 'recipe-id': recipe.id, + 'model-id': model.id, + }); // first delete any existing pod with matching labels if (await this.hasApplicationPod(recipe.id, model.id)) { @@ -257,28 +199,32 @@ export class ApplicationManager extends Publisher implements * @param podInfo * @param labels */ - async runApplication(podInfo: PodInfo, labels?: { [key: string]: string }): Promise { + protected async runApplication(podInfo: PodInfo, labels?: { [key: string]: string }): Promise { const task = this.taskRegistry.createTask('Starting AI App', 'loading', labels); // it starts the pod - await this.podManager.startPod(podInfo.engineId, podInfo.Id); + try { + await this.podManager.startPod(podInfo.engineId, podInfo.Id); - // check if all containers have started successfully - for (const container of podInfo.Containers ?? []) { - await this.waitContainerIsRunning(podInfo.engineId, container); - } + // check if all containers have started successfully + for (const container of podInfo.Containers ?? []) { + await this.waitContainerIsRunning(podInfo.engineId, container); + } - // Update task registry - this.taskRegistry.updateTask({ - ...task, - state: 'success', - name: 'AI App is running', - }); + task.state = 'success'; + task.name = 'AI App is running'; + } catch (err: unknown) { + task.state = 'error'; + task.error = String(err); + throw err; + } finally { + this.taskRegistry.updateTask(task); + } return this.checkPodsHealth(); } - async waitContainerIsRunning(engineId: string, container: PodContainerInfo): Promise { + protected async waitContainerIsRunning(engineId: string, container: PodContainerInfo): Promise { const TIME_FRAME_MS = 5000; const MAX_ATTEMPTS = 60 * (60000 / TIME_FRAME_MS); // try for 1 hour for (let i = 0; i < MAX_ATTEMPTS; i++) { @@ -291,10 +237,10 @@ export class ApplicationManager extends Publisher implements throw new Error(`Container ${container.Id} not started in time`); } - async createApplicationPod( + protected async createApplicationPod( recipe: Recipe, model: ModelInfo, - images: ImageInfo[], + images: RecipeImage[], modelPath: string, labels?: { [key: string]: string }, ): Promise { @@ -332,9 +278,9 @@ export class ApplicationManager extends Publisher implements return podInfo; } - async createContainerAndAttachToPod( + protected async createContainerAndAttachToPod( podInfo: PodInfo, - images: ImageInfo[], + images: RecipeImage[], modelInfo: ModelInfo, modelPath: string, ): Promise { @@ -393,7 +339,7 @@ export class ApplicationManager extends Publisher implements ); } - async createPod(recipe: Recipe, model: ModelInfo, images: ImageInfo[]): Promise { + protected async createPod(recipe: Recipe, model: ModelInfo, images: RecipeImage[]): Promise { // find the exposed port of the sample app so we can open its ports on the new pod const sampleAppImageInfo = images.find(image => !image.modelService); if (!sampleAppImageInfo) { @@ -420,22 +366,24 @@ export class ApplicationManager extends Publisher implements // create new pod const labels: Record = { - [LABEL_RECIPE_ID]: recipe.id, - [LABEL_MODEL_ID]: model.id, + [POD_LABEL_RECIPE_ID]: recipe.id, + [POD_LABEL_MODEL_ID]: model.id, }; + // collecting all modelService ports const modelPorts = images .filter(img => img.modelService) .flatMap(img => img.ports) .map(port => portmappings.find(pm => `${pm.container_port}` === port)?.host_port); if (modelPorts.length) { - labels[LABEL_MODEL_PORTS] = modelPorts.join(','); + labels[POD_LABEL_MODEL_PORTS] = modelPorts.join(','); } + // collecting all application ports (excluding service ports) const appPorts = images .filter(img => !img.modelService) .flatMap(img => img.ports) .map(port => portmappings.find(pm => `${pm.container_port}` === port)?.host_port); if (appPorts.length) { - labels[LABEL_APP_PORTS] = appPorts.join(','); + labels[POD_LABEL_APP_PORTS] = appPorts.join(','); } const { engineId, Id } = await this.podManager.createPod({ name: getRandomName(`pod-${sampleAppImageInfo.appName}`), @@ -475,8 +423,10 @@ export class ApplicationManager extends Publisher implements stoppingTask.state = 'success'; stoppingTask.name = `AI App Stopped`; } catch (err: unknown) { + stoppingTask.state = 'error'; stoppingTask.error = `Error removing the pod.: ${String(err)}`; stoppingTask.name = 'Error stopping AI App'; + throw err; } finally { this.taskRegistry.updateTask(stoppingTask); await this.checkPodsHealth(); @@ -499,112 +449,10 @@ export class ApplicationManager extends Publisher implements }); } - private getConfigAndFilterContainers( - recipeBaseDir: string | undefined, - localFolder: string, - labels?: { [key: string]: string }, - ): AIContainers { - // Adding loading configuration task - const task = this.taskRegistry.createTask('Loading configuration', 'loading', labels); - - let aiConfigFile: AIConfigFile; - try { - // load and parse the recipe configuration file - aiConfigFile = this.getConfiguration(recipeBaseDir, localFolder); - } catch (e) { - task.error = `Something went wrong while loading configuration: ${String(e)}.`; - this.taskRegistry.updateTask(task); - throw e; - } - - // filter the containers based on architecture, gpu accelerator and backend (that define which model supports) - const filteredContainers: ContainerConfig[] = this.filterContainers(aiConfigFile.aiConfig); - if (filteredContainers.length > 0) { - // Mark as success. - task.state = 'success'; - this.taskRegistry.updateTask(task); - } else { - // Mark as failure. - task.error = 'No containers available.'; - this.taskRegistry.updateTask(task); - throw new Error('No containers available.'); - } - - return { - aiConfigFile: aiConfigFile, - containers: filteredContainers, - }; - } - - filterContainers(aiConfig: AIConfig): ContainerConfig[] { - return aiConfig.application.containers.filter( - container => container.gpu_env.length === 0 && container.arch.some(arc => arc === goarch()), - ); - } - - getConfiguration(recipeBaseDir: string | undefined, localFolder: string): AIConfigFile { - let configFile: string; - if (recipeBaseDir !== undefined) { - configFile = path.join(localFolder, recipeBaseDir, CONFIG_FILENAME); - } else { - configFile = path.join(localFolder, CONFIG_FILENAME); - } - - if (!fs.existsSync(configFile)) { - throw new Error(`The file located at ${configFile} does not exist.`); - } - - // If the user configured the config as a directory we check for "ai-lab.yaml" inside. - if (fs.statSync(configFile).isDirectory()) { - const tmpPath = path.join(configFile, CONFIG_FILENAME); - // If it has the ai-lab.yaml we use it. - if (fs.existsSync(tmpPath)) { - configFile = tmpPath; - } - } - - // Parsing the configuration - console.log(`Reading configuration from ${configFile}.`); - let aiConfig: AIConfig; - try { - aiConfig = parseYamlFile(configFile, goarch()); - } catch (err) { - console.error('Cannot load configure file.', err); - throw new Error(`Cannot load configuration file.`); - } - - // Mark as success. - return { - aiConfig, - path: configFile, - }; - } - - private async doCheckout(gitCloneInfo: GitCloneInfo, labels?: { [id: string]: string }): Promise { - // Creating checkout task - const checkoutTask: Task = this.taskRegistry.createTask('Checking out repository', 'loading', { - ...labels, - git: 'checkout', - }); - - try { - await this.git.processCheckout(gitCloneInfo); - checkoutTask.state = 'success'; - } catch (err: unknown) { - checkoutTask.state = 'error'; - checkoutTask.error = String(err); - // propagate error - throw err; - } finally { - // Update task registry - this.taskRegistry.updateTask(checkoutTask); - } - } - init() { this.podmanConnection.startupSubscribe(() => { this.podManager - .getPodsWithLabels([LABEL_RECIPE_ID]) + .getPodsWithLabels([POD_LABEL_RECIPE_ID]) .then(pods => { pods.forEach(pod => this.adoptPod(pod)); }) @@ -651,17 +499,17 @@ export class ApplicationManager extends Publisher implements ); } - private adoptPod(pod: PodInfo) { + protected adoptPod(pod: PodInfo) { if (!pod.Labels) { return; } - const recipeId = pod.Labels[LABEL_RECIPE_ID]; - const modelId = pod.Labels[LABEL_MODEL_ID]; + const recipeId = pod.Labels[POD_LABEL_RECIPE_ID]; + const modelId = pod.Labels[POD_LABEL_MODEL_ID]; if (!recipeId || !modelId) { return; } - const appPorts = getPortsFromLabel(pod.Labels, LABEL_APP_PORTS); - const modelPorts = getPortsFromLabel(pod.Labels, LABEL_MODEL_PORTS); + const appPorts = getPortsFromLabel(pod.Labels, POD_LABEL_APP_PORTS); + const modelPorts = getPortsFromLabel(pod.Labels, POD_LABEL_MODEL_PORTS); if (this.#applications.has({ recipeId, modelId })) { return; } @@ -676,7 +524,7 @@ export class ApplicationManager extends Publisher implements this.updateApplicationState(recipeId, modelId, state); } - private forgetPodById(podId: string) { + protected forgetPodById(podId: string) { const app = Array.from(this.#applications.values()).find(p => p.pod.Id === podId); if (!app) { return; @@ -684,8 +532,8 @@ export class ApplicationManager extends Publisher implements if (!app.pod.Labels) { return; } - const recipeId = app.pod.Labels[LABEL_RECIPE_ID]; - const modelId = app.pod.Labels[LABEL_MODEL_ID]; + const recipeId = app.pod.Labels[POD_LABEL_RECIPE_ID]; + const modelId = app.pod.Labels[POD_LABEL_MODEL_ID]; if (!recipeId || !modelId) { return; } @@ -706,13 +554,13 @@ export class ApplicationManager extends Publisher implements } } - private async checkPodsHealth(): Promise { - const pods = await this.podManager.getPodsWithLabels([LABEL_RECIPE_ID, LABEL_MODEL_ID]); + protected async checkPodsHealth(): Promise { + const pods = await this.podManager.getPodsWithLabels([POD_LABEL_RECIPE_ID, POD_LABEL_MODEL_ID]); let changes = false; for (const pod of pods) { - const recipeId = pod.Labels[LABEL_RECIPE_ID]; - const modelId = pod.Labels[LABEL_MODEL_ID]; + const recipeId = pod.Labels[POD_LABEL_RECIPE_ID]; + const modelId = pod.Labels[POD_LABEL_MODEL_ID]; if (!this.#applications.has({ recipeId, modelId })) { // a fresh pod could not have been added yet, we will handle it at next iteration continue; @@ -736,7 +584,7 @@ export class ApplicationManager extends Publisher implements } } - updateApplicationState(recipeId: string, modelId: string, state: ApplicationState): void { + protected updateApplicationState(recipeId: string, modelId: string, state: ApplicationState): void { this.#applications.set({ recipeId, modelId }, state); this.notify(); } @@ -745,7 +593,7 @@ export class ApplicationManager extends Publisher implements return Array.from(this.#applications.values()); } - private clearTasks(recipeId: string, modelId: string): void { + protected clearTasks(recipeId: string, modelId: string): void { // clear any existing status / tasks related to the pair recipeId-modelId. this.taskRegistry.deleteByLabels({ 'recipe-id': recipeId, @@ -785,7 +633,7 @@ export class ApplicationManager extends Publisher implements const appPod = await this.getApplicationPod(recipeId, modelId); await this.removeApplication(recipeId, modelId); const recipe = this.catalogManager.getRecipeById(recipeId); - const model = this.catalogManager.getModelById(appPod.Labels[LABEL_MODEL_ID]); + const model = this.catalogManager.getModelById(appPod.Labels[POD_LABEL_MODEL_ID]); // init the recipe const podInfo = await this.initApplication(recipe, model); @@ -798,15 +646,14 @@ export class ApplicationManager extends Publisher implements } async getApplicationPorts(recipeId: string, modelId: string): Promise { - const recipe = this.catalogManager.getRecipeById(recipeId); const state = this.#applications.get({ recipeId, modelId }); if (state) { return state.appPorts; } - throw new Error(`Recipe ${recipe.name} has no ports available`); + throw new Error(`Recipe ${recipeId} has no ports available`); } - async getApplicationPod(recipeId: string, modelId: string): Promise { + protected async getApplicationPod(recipeId: string, modelId: string): Promise { const appPod = await this.findPod(recipeId, modelId); if (!appPod) { throw new Error(`no pod found with recipe Id ${recipeId} and model Id ${modelId}`); @@ -814,7 +661,7 @@ export class ApplicationManager extends Publisher implements return appPod; } - private async hasApplicationPod(recipeId: string, modelId: string): Promise { + protected async hasApplicationPod(recipeId: string, modelId: string): Promise { const pod = await this.podManager.findPodByLabelsValues({ LABEL_RECIPE_ID: recipeId, LABEL_MODEL_ID: modelId, @@ -822,10 +669,10 @@ export class ApplicationManager extends Publisher implements return !!pod; } - private async findPod(recipeId: string, modelId: string): Promise { + protected async findPod(recipeId: string, modelId: string): Promise { return this.podManager.findPodByLabelsValues({ - [LABEL_RECIPE_ID]: recipeId, - [LABEL_MODEL_ID]: modelId, + [POD_LABEL_RECIPE_ID]: recipeId, + [POD_LABEL_MODEL_ID]: modelId, }); } diff --git a/packages/backend/src/managers/applicationManager.spec.ts b/packages/backend/src/managers/applicationManager.spec.ts deleted file mode 100644 index f8484f041..000000000 --- a/packages/backend/src/managers/applicationManager.spec.ts +++ /dev/null @@ -1,1277 +0,0 @@ -/********************************************************************** - * 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 { beforeEach, describe, expect, type MockInstance, test, vi } from 'vitest'; -import { ApplicationManager, CONFIG_FILENAME, type ImageInfo } from './applicationManager'; -import type { GitManager } from './gitManager'; -import os from 'os'; -import fs, { type PathLike } from 'node:fs'; -import type { Recipe } from '@shared/src/models/IRecipe'; -import type { ModelInfo } from '@shared/src/models/IModelInfo'; -import { ModelsManager } from './modelsManager'; -import path from 'node:path'; -import type { AIConfig, ContainerConfig } from '../models/AIConfig'; -import * as portsUtils from '../utils/ports'; -import { goarch } from '../utils/arch'; -import * as utils from '../utils/utils'; -import type { Disposable, PodInfo, TelemetryLogger, Webview } from '@podman-desktop/api'; -import type { CatalogManager } from './catalogManager'; -import type { LocalRepositoryRegistry } from '../registries/LocalRepositoryRegistry'; -import type { machineStopHandle, PodmanConnection, startupHandle } from './podmanConnection'; -import { TaskRegistry } from '../registries/TaskRegistry'; -import type { CancellationTokenRegistry } from '../registries/CancellationTokenRegistry'; -import type { BuilderManager } from './recipes/BuilderManager'; -import type { PodEvent, PodManager } from './recipes/PodManager'; -import { VMType } from '@shared/src/models/IPodman'; - -const mocks = vi.hoisted(() => { - return { - parseYamlFileMock: vi.fn(), - listImagesMock: vi.fn(), - getImageInspectMock: vi.fn(), - createContainerMock: vi.fn(), - startContainerMock: vi.fn(), - inspectContainerMock: vi.fn(), - logUsageMock: vi.fn(), - logErrorMock: vi.fn(), - registerLocalRepositoryMock: vi.fn(), - postMessageMock: vi.fn(), - getContainerConnectionsMock: vi.fn(), - pullImageMock: vi.fn(), - stopContainerMock: vi.fn(), - containerRegistrySubscribeMock: vi.fn(), - startupSubscribeMock: vi.fn(), - onMachineStopMock: vi.fn(), - listContainersMock: vi.fn(), - performDownloadMock: vi.fn(), - getTargetMock: vi.fn(), - onEventDownloadMock: vi.fn(), - // TaskRegistry - getTaskMock: vi.fn(), - createTaskMock: vi.fn(), - updateTaskMock: vi.fn(), - deleteMock: vi.fn(), - deleteAllMock: vi.fn(), - getTasksMock: vi.fn(), - getTasksByLabelsMock: vi.fn(), - deleteByLabelsMock: vi.fn(), - }; -}); -vi.mock('../models/AIConfig', () => ({ - parseYamlFile: mocks.parseYamlFileMock, -})); - -vi.mock('../utils/downloader', () => ({ - Downloader: class { - onEvent = mocks.onEventDownloadMock; - perform = mocks.performDownloadMock; - getTarget = mocks.getTargetMock; - }, -})); - -vi.mock('@podman-desktop/api', () => ({ - provider: { - getContainerConnections: mocks.getContainerConnectionsMock, - }, - containerEngine: { - listImages: mocks.listImagesMock, - getImageInspect: mocks.getImageInspectMock, - createContainer: mocks.createContainerMock, - startContainer: mocks.startContainerMock, - inspectContainer: mocks.inspectContainerMock, - pullImage: mocks.pullImageMock, - stopContainer: mocks.stopContainerMock, - listContainers: mocks.listContainersMock, - }, - Disposable: { - create: vi.fn(), - }, -})); - -const telemetryLogger = { - logUsage: mocks.logUsageMock, - logError: mocks.logErrorMock, -} as unknown as TelemetryLogger; - -const taskRegistry = { - getTask: mocks.getTaskMock, - createTask: mocks.createTaskMock, - updateTask: mocks.updateTaskMock, - delete: mocks.deleteMock, - deleteAll: mocks.deleteAllMock, - getTasks: mocks.getTasksMock, - getTasksByLabels: mocks.getTasksByLabelsMock, - deleteByLabels: mocks.deleteByLabelsMock, -} as unknown as TaskRegistry; - -const builderManager = { - build: vi.fn(), -} as unknown as BuilderManager; - -const podManager = { - getAllPods: vi.fn(), - findPodByLabelsValues: vi.fn(), - getPodsWithLabels: vi.fn(), - getHealth: vi.fn(), - getPod: vi.fn(), - createPod: vi.fn(), - stopPod: vi.fn(), - removePod: vi.fn(), - startPod: vi.fn(), - onStartPodEvent: vi.fn(), - onStopPodEvent: vi.fn(), - onRemovePodEvent: vi.fn(), -} as unknown as PodManager; - -const localRepositoryRegistry = { - register: mocks.registerLocalRepositoryMock, -} as unknown as LocalRepositoryRegistry; - -const podmanConnection = { - getVMType: vi.fn(), - startupSubscribe: mocks.startupSubscribeMock, - onMachineStop: mocks.onMachineStopMock, -} as unknown as PodmanConnection; - -beforeEach(() => { - vi.resetAllMocks(); - - mocks.createTaskMock.mockImplementation((name, state, labels) => ({ - id: 'random', - name: name, - state: state, - labels: labels ?? {}, - error: undefined, - })); -}); - -describe('pullApplication', () => { - interface mockForPullApplicationOptions { - recipeFolderExists: boolean; - } - const processCheckoutMock = vi.fn(); - let manager: ApplicationManager; - let modelsManager: ModelsManager; - vi.spyOn(utils, 'timeout').mockResolvedValue(); - - function mockForPullApplication(options: mockForPullApplicationOptions) { - vi.spyOn(os, 'homedir').mockReturnValue('/home/user'); - vi.spyOn(fs, 'mkdirSync').mockReturnValue(undefined); - vi.spyOn(fs, 'existsSync').mockImplementation((path: PathLike) => { - path = path.toString(); - if (path.endsWith('recipe1')) { - return options.recipeFolderExists; - } else if (path.endsWith('ai-lab.yaml')) { - return true; - } else if (path.endsWith('contextdir1')) { - return true; - } - return false; - }); - (vi.spyOn(fs, 'statSync') as unknown as MockInstance<(path: PathLike) => fs.Stats>).mockImplementation( - (path: PathLike): fs.Stats => { - path = path.toString(); - if (path.endsWith('recipe1')) { - const stat = new fs.Stats(); - stat.isDirectory = () => true; - return stat; - } else if (path.endsWith('ai-lab.yaml')) { - const stat = new fs.Stats(); - stat.isDirectory = () => false; - return stat; - } - throw new Error('should never be reached'); - }, - ); - (vi.spyOn(fs, 'readFileSync') as unknown as MockInstance<() => string>).mockImplementation(() => { - return ''; - }); - mocks.parseYamlFileMock.mockReturnValue({ - application: { - containers: [ - { - name: 'container1', - contextdir: 'contextdir1', - containerfile: 'Containerfile', - arch: [goarch()], - gpu_env: [], - }, - ], - }, - }); - mocks.inspectContainerMock.mockResolvedValue({ - State: { - Running: true, - }, - }); - vi.mocked(builderManager.build).mockResolvedValue([ - { - modelService: false, - appName: 'dummy-app-name', - ports: [], - id: 'dummy-id', - }, - ]); - mocks.listImagesMock.mockResolvedValue([ - { - RepoTags: ['recipe1-container1:latest'], - engineId: 'engine', - Id: 'id1', - }, - ]); - vi.mocked(podManager.createPod).mockResolvedValue({ - engineId: 'engine', - Id: 'id', - }); - mocks.createContainerMock.mockResolvedValue({ - id: 'id', - }); - modelsManager = new ModelsManager( - 'appdir', - {} as Webview, - { - getModels(): ModelInfo[] { - return []; - }, - } as CatalogManager, - telemetryLogger, - new TaskRegistry({ postMessage: vi.fn().mockResolvedValue(undefined) } as unknown as Webview), - { - createCancellationTokenSource: vi.fn(), - } as unknown as CancellationTokenRegistry, - ); - manager = new ApplicationManager( - '/home/user/aistudio', - { - processCheckout: processCheckoutMock, - isGitInstalled: () => true, - } as unknown as GitManager, - taskRegistry, - {} as Webview, - podmanConnection, - {} as CatalogManager, - modelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - } - test('pullApplication should clone repository and call downloadModelMain and buildImage', async () => { - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([]); - mockForPullApplication({ - recipeFolderExists: false, - }); - vi.mocked(podManager.getAllPods).mockResolvedValue([]); - vi.mocked(podManager.getPod).mockResolvedValue({ - engineId: 'dummyEngineId', - Id: 'dummyPodId', - } as unknown as PodInfo); - vi.mocked(podmanConnection.getVMType).mockResolvedValue(VMType.WSL); - vi.spyOn(modelsManager, 'isModelOnDisk').mockReturnValue(false); - vi.spyOn(modelsManager, 'uploadModelToPodmanMachine').mockResolvedValue('path'); - mocks.performDownloadMock.mockResolvedValue('path'); - const recipe: Recipe = { - id: 'recipe1', - name: 'Recipe 1', - categories: [], - description: '', - ref: '000000', - readme: '', - repository: 'repo', - }; - const model: ModelInfo = { - id: 'model1', - description: '', - hw: '', - license: '', - name: 'Model 1', - registry: '', - url: 'dummy-url', - memory: 1000, - }; - mocks.inspectContainerMock.mockResolvedValue({ - State: { - Running: true, - }, - }); - vi.spyOn(utils, 'getDurationSecondsSince').mockReturnValue(99); - await manager.pullApplication(recipe, model); - const gitCloneOptions = { - repository: 'repo', - ref: '000000', - targetDirectory: '\\home\\user\\aistudio\\recipe1', - }; - if (process.platform === 'win32') { - expect(processCheckoutMock).toHaveBeenNthCalledWith(1, gitCloneOptions); - } else { - gitCloneOptions.targetDirectory = '/home/user/aistudio/recipe1'; - expect(processCheckoutMock).toHaveBeenNthCalledWith(1, gitCloneOptions); - } - expect(mocks.performDownloadMock).toHaveBeenCalledOnce(); - expect(builderManager.build).toHaveBeenCalledOnce(); - expect(builderManager.build).toHaveBeenCalledWith( - { - categories: [], - description: '', - id: 'recipe1', - name: 'Recipe 1', - readme: '', - ref: '000000', - repository: 'repo', - }, - [ - { - arch: ['amd64'], - containerfile: 'Containerfile', - contextdir: 'contextdir1', - gpu_env: [], - name: 'container1', - }, - ], - expect.anything(), - { - 'model-id': 'model1', - 'recipe-id': 'recipe1', - }, - ); - expect(mocks.logUsageMock).toHaveBeenNthCalledWith(1, 'recipe.pull', { - 'recipe.id': 'recipe1', - 'recipe.name': 'Recipe 1', - durationSeconds: 99, - }); - }); - test('pullApplication should clone repository and call downloadModelMain and fail on buildImage', async () => { - mockForPullApplication({ - recipeFolderExists: false, - }); - vi.mocked(builderManager.build).mockRejectedValue(new Error('Build failed')); - vi.mocked(podManager.getAllPods).mockResolvedValue([]); - vi.mocked(podmanConnection.getVMType).mockResolvedValue(VMType.WSL); - vi.spyOn(modelsManager, 'isModelOnDisk').mockReturnValue(false); - vi.spyOn(modelsManager, 'uploadModelToPodmanMachine').mockResolvedValue('path'); - mocks.performDownloadMock.mockResolvedValue('path'); - const recipe: Recipe = { - id: 'recipe1', - name: 'Recipe 1', - categories: [], - description: '', - ref: '000000', - readme: '', - repository: 'repo', - }; - const model: ModelInfo = { - id: 'model1', - description: '', - hw: '', - license: '', - name: 'Model 1', - registry: '', - url: 'dummy-url', - memory: 1000, - }; - mocks.inspectContainerMock.mockResolvedValue({ - State: { - Running: true, - }, - }); - vi.spyOn(utils, 'getDurationSecondsSince').mockReturnValue(99); - let error: unknown = undefined; - try { - await manager.pullApplication(recipe, model); - } catch (err: unknown) { - error = err; - } - expect(error).toBeDefined(); - const gitCloneOptions = { - repository: 'repo', - ref: '000000', - targetDirectory: '\\home\\user\\aistudio\\recipe1', - }; - if (process.platform === 'win32') { - expect(processCheckoutMock).toHaveBeenNthCalledWith(1, gitCloneOptions); - } else { - gitCloneOptions.targetDirectory = '/home/user/aistudio/recipe1'; - expect(processCheckoutMock).toHaveBeenNthCalledWith(1, gitCloneOptions); - } - expect(mocks.performDownloadMock).toHaveBeenCalledOnce(); - expect(builderManager.build).toHaveBeenCalledOnce(); - expect(mocks.logErrorMock).toHaveBeenNthCalledWith( - 1, - 'recipe.pull', - expect.objectContaining({ - 'recipe.id': 'recipe1', - 'recipe.name': 'Recipe 1', - durationSeconds: 99, - message: 'error pulling application', - }), - ); - }); - test('pullApplication should not download model if already on disk', async () => { - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([]); - mockForPullApplication({ - recipeFolderExists: true, - }); - vi.mocked(podManager.getAllPods).mockResolvedValue([]); - vi.mocked(podManager.getPod).mockResolvedValue({ - engineId: 'dummyEngineId', - Id: 'dummyPodId', - } as unknown as PodInfo); - vi.spyOn(modelsManager, 'isModelOnDisk').mockReturnValue(true); - vi.spyOn(modelsManager, 'uploadModelToPodmanMachine').mockResolvedValue('path'); - vi.spyOn(modelsManager, 'getLocalModelPath').mockReturnValue('path'); - const recipe: Recipe = { - id: 'recipe1', - name: 'Recipe 1', - categories: [], - ref: '000000', - description: '', - readme: '', - repository: 'repo', - }; - const model: ModelInfo = { - id: 'model1', - description: '', - hw: '', - license: '', - name: 'Model 1', - registry: '', - url: '', - memory: 1000, - }; - await manager.pullApplication(recipe, model); - expect(mocks.performDownloadMock).not.toHaveBeenCalled(); - }); - - test('pullApplication should mark the loading config as error if not container are found', async () => { - mockForPullApplication({ - recipeFolderExists: true, - }); - - const recipe: Recipe = { - id: 'recipe1', - name: 'Recipe 1', - categories: [], - description: '', - ref: '000000', - readme: '', - repository: 'repo', - }; - const model: ModelInfo = { - id: 'model1', - description: '', - hw: '', - license: '', - name: 'Model 1', - registry: '', - url: '', - memory: 1000, - }; - - mocks.parseYamlFileMock.mockReturnValue({ - application: { - containers: [], - }, - }); - - await expect(manager.pullApplication(recipe, model)).rejects.toThrowError('No containers available.'); - expect(mocks.performDownloadMock).not.toHaveBeenCalled(); - }); -}); - -describe('getConfiguration', () => { - test('throws error if config file do not exists', async () => { - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - vi.spyOn(fs, 'existsSync').mockReturnValue(false); - expect(() => manager.getConfiguration('config', 'local')).toThrowError( - `The file located at ${path.join('local', 'config', CONFIG_FILENAME)} does not exist.`, - ); - }); - - test('return AIConfigFile', async () => { - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - vi.spyOn(fs, 'existsSync').mockReturnValue(true); - const stats = { - isDirectory: vi.fn().mockReturnValue(false), - } as unknown as fs.Stats; - vi.spyOn(fs, 'statSync').mockReturnValue(stats); - vi.spyOn(fs, 'readFileSync').mockReturnValue(''); - const aiConfig = { - application: { - containers: [ - { - name: 'container1', - contextdir: 'contextdir1', - containerfile: 'Containerfile', - }, - ], - }, - }; - mocks.parseYamlFileMock.mockReturnValue(aiConfig); - - const result = manager.getConfiguration('config', 'local'); - expect(result.path).toEqual(path.join('local', 'config', CONFIG_FILENAME)); - expect(result.aiConfig).toEqual(aiConfig); - }); -}); - -describe('filterContainers', () => { - test('return empty array when no container fit the system', () => { - const aiConfig: AIConfig = { - application: { - containers: [ - { - name: 'container2', - contextdir: 'contextdir2', - containerfile: 'Containerfile', - arch: ['arm64'], - modelService: false, - gpu_env: [], - }, - ], - }, - }; - Object.defineProperty(process, 'arch', { - value: 'amd64', - }); - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - const containers = manager.filterContainers(aiConfig); - expect(containers.length).toBe(0); - }); - test('return one container when only one fit the system', () => { - const aiConfig: AIConfig = { - application: { - containers: [ - { - name: 'container1', - contextdir: 'contextdir1', - containerfile: 'Containerfile', - arch: ['amd64'], - modelService: false, - gpu_env: [], - }, - { - name: 'container2', - contextdir: 'contextdir2', - containerfile: 'Containerfile', - arch: ['arm64'], - modelService: false, - gpu_env: [], - }, - ], - }, - }; - Object.defineProperty(process, 'arch', { - value: 'amd64', - }); - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - const containers = manager.filterContainers(aiConfig); - expect(containers.length).toBe(1); - expect(containers[0].name).equal('container1'); - }); - test('return 2 containers when two fit the system', () => { - const containerConfig: ContainerConfig[] = [ - { - name: 'container1', - contextdir: 'contextdir1', - containerfile: 'Containerfile', - arch: ['amd64'], - modelService: false, - gpu_env: [], - }, - { - name: 'container2', - contextdir: 'contextdir2', - containerfile: 'Containerfile', - arch: ['arm64'], - modelService: false, - gpu_env: [], - }, - { - name: 'container3', - contextdir: 'contextdir3', - containerfile: 'Containerfile', - arch: ['amd64'], - modelService: false, - gpu_env: [], - }, - ]; - const aiConfig: AIConfig = { - application: { - containers: containerConfig, - }, - }; - Object.defineProperty(process, 'arch', { - value: 'amd64', - }); - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - const containers = manager.filterContainers(aiConfig); - expect(containers.length).toBe(2); - expect(containers[0].name).equal('container1'); - expect(containers[1].name).equal('container3'); - }); -}); - -describe('createPod', async () => { - const imageInfo1: ImageInfo = { - id: 'id', - appName: 'appName', - modelService: false, - ports: ['8080', '8081'], - }; - const imageInfo2: ImageInfo = { - id: 'id2', - appName: 'appName2', - modelService: true, - ports: ['8082'], - }; - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - test('throw an error if there is no sample image', async () => { - const images = [imageInfo2]; - await expect( - manager.createPod({ id: 'recipe-id' } as Recipe, { id: 'model-id' } as ModelInfo, images), - ).rejects.toThrowError('no sample app found'); - }); - test('call createPod with sample app exposed port', async () => { - const images = [imageInfo1, imageInfo2]; - vi.spyOn(portsUtils, 'getPortsInfo').mockResolvedValueOnce('9000'); - vi.spyOn(portsUtils, 'getPortsInfo').mockResolvedValueOnce('9001'); - vi.spyOn(portsUtils, 'getPortsInfo').mockResolvedValueOnce('9002'); - vi.mocked(podManager.createPod).mockResolvedValue({ - Id: 'podId', - engineId: 'engineId', - }); - await manager.createPod({ id: 'recipe-id' } as Recipe, { id: 'model-id' } as ModelInfo, images); - expect(podManager.createPod).toBeCalledWith({ - name: expect.anything(), - portmappings: [ - { - container_port: 8080, - host_port: 9002, - host_ip: '', - protocol: '', - range: 1, - }, - { - container_port: 8081, - host_port: 9001, - host_ip: '', - protocol: '', - range: 1, - }, - { - container_port: 8082, - host_port: 9000, - host_ip: '', - protocol: '', - range: 1, - }, - ], - labels: { - 'ai-lab-recipe-id': 'recipe-id', - 'ai-lab-app-ports': '9002,9001', - 'ai-lab-model-id': 'model-id', - 'ai-lab-model-ports': '9000', - }, - }); - }); -}); - -describe('createApplicationPod', () => { - const imageInfo1: ImageInfo = { - id: 'id', - appName: 'appName', - modelService: false, - ports: ['8080', '8081'], - }; - const imageInfo2: ImageInfo = { - id: 'id2', - appName: 'appName2', - modelService: true, - ports: ['8082'], - }; - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - const images = [imageInfo1, imageInfo2]; - test('throw if createPod fails', async () => { - vi.spyOn(manager, 'createPod').mockRejectedValue('error createPod'); - await expect( - manager.createApplicationPod({ id: 'recipe-id' } as Recipe, { id: 'model-id' } as ModelInfo, images, 'path'), - ).rejects.toThrowError('error createPod'); - expect(mocks.updateTaskMock).toBeCalledWith({ - error: 'Something went wrong while creating pod: error createPod', - id: expect.any(String), - state: 'error', - name: 'Creating AI App', - labels: {}, - }); - }); - test('call createAndAddContainersToPod after pod is created', async () => { - const pod: PodInfo = { - engineId: 'engine', - Id: 'id', - } as unknown as PodInfo; - vi.spyOn(manager, 'createPod').mockResolvedValue(pod); - // mock createContainerAndAttachToPod - const createAndAddContainersToPodMock = vi - .spyOn(manager, 'createContainerAndAttachToPod') - .mockResolvedValue(undefined); - await manager.createApplicationPod({ id: 'recipe-id' } as Recipe, { id: 'model-id' } as ModelInfo, images, 'path'); - expect(createAndAddContainersToPodMock).toBeCalledWith(pod, images, { id: 'model-id' }, 'path'); - expect(mocks.updateTaskMock).toBeCalledWith({ - id: expect.any(String), - state: 'success', - name: 'Creating AI App', - labels: { - 'pod-id': pod.Id, - }, - }); - }); - test('throw if createAndAddContainersToPod fails', async () => { - const pod: PodInfo = { - engineId: 'engine', - Id: 'id', - } as unknown as PodInfo; - vi.spyOn(manager, 'createPod').mockResolvedValue(pod); - vi.spyOn(manager, 'createContainerAndAttachToPod').mockRejectedValue('error'); - await expect(() => - manager.createApplicationPod({ id: 'recipe-id' } as Recipe, { id: 'model-id' } as ModelInfo, images, 'path'), - ).rejects.toThrowError('error'); - expect(mocks.updateTaskMock).toHaveBeenLastCalledWith({ - id: expect.any(String), - state: 'error', - error: 'Something went wrong while creating pod: error', - name: 'Creating AI App', - labels: { - 'pod-id': pod.Id, - }, - }); - }); -}); - -describe('runApplication', () => { - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - {} as PodmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - const pod: PodInfo = { - engineId: 'engine', - Id: 'id', - Containers: [ - { - Id: 'dummyContainerId', - }, - ], - } as unknown as PodInfo; - test('check startPod is called and also waitContainerIsRunning for sample app', async () => { - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([]); - const waitContainerIsRunningMock = vi.spyOn(manager, 'waitContainerIsRunning').mockResolvedValue(undefined); - vi.spyOn(utils, 'timeout').mockResolvedValue(); - await manager.runApplication(pod); - expect(podManager.startPod).toBeCalledWith(pod.engineId, pod.Id); - expect(waitContainerIsRunningMock).toBeCalledWith(pod.engineId, { - Id: 'dummyContainerId', - }); - }); -}); - -describe('createAndAddContainersToPod', () => { - const manager = new ApplicationManager( - '/home/user/aistudio', - {} as unknown as GitManager, - taskRegistry, - {} as Webview, - podmanConnection, - {} as CatalogManager, - {} as unknown as ModelsManager, - telemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - const pod: PodInfo = { - engineId: 'engine', - Id: 'id', - portmappings: [], - } as unknown as PodInfo; - const imageInfo1: ImageInfo = { - id: 'id', - appName: 'appName', - modelService: false, - ports: ['8080', '8081'], - }; - const imageInfo2: ImageInfo = { - id: 'id2', - appName: 'appName2', - modelService: true, - ports: ['8085'], - }; - async function checkContainers(modelInfo: ModelInfo, extraEnvs: string[]) { - mocks.createContainerMock.mockResolvedValue({ - id: 'container-1', - }); - vi.mocked(podmanConnection.getVMType).mockResolvedValue(VMType.WSL); - await manager.createContainerAndAttachToPod(pod, [imageInfo1, imageInfo2], modelInfo, 'path'); - expect(mocks.createContainerMock).toHaveBeenNthCalledWith(1, 'engine', { - Image: 'id', - Detach: true, - Env: ['MODEL_ENDPOINT=http://localhost:8085'], - start: false, - name: expect.anything(), - pod: 'id', - HealthCheck: { - Interval: 5000000000, - Retries: 20, - Test: ['CMD-SHELL', 'curl -s localhost:8080 > /dev/null'], - Timeout: 2000000000, - }, - }); - expect(mocks.createContainerMock).toHaveBeenNthCalledWith(2, 'engine', { - Image: 'id2', - Detach: true, - Env: ['MODEL_PATH=/path', ...extraEnvs], - start: false, - name: expect.anything(), - pod: 'id', - HostConfig: { - Mounts: [ - { - Mode: 'Z', - Source: 'path', - Target: '/path', - Type: 'bind', - }, - ], - }, - HealthCheck: { - Interval: 5000000000, - Retries: 20, - Test: ['CMD-SHELL', 'curl -s localhost:8085 > /dev/null'], - Timeout: 2000000000, - }, - }); - } - - test('check that containers are correctly created with no model properties', async () => { - await checkContainers({} as ModelInfo, []); - }); - - test('check that containers are correctly created with model properties', async () => { - await checkContainers( - { - properties: { - modelName: 'myModel', - }, - } as unknown as ModelInfo, - ['MODEL_MODEL_NAME=myModel'], - ); - }); -}); - -describe('pod detection', async () => { - let manager: ApplicationManager; - - beforeEach(() => { - vi.resetAllMocks(); - - mocks.createTaskMock.mockImplementation((name, state, labels) => ({ - id: 'random', - name: name, - state: state, - labels: labels ?? {}, - error: undefined, - })); - - manager = new ApplicationManager( - '/path/to/user/dir', - {} as GitManager, - taskRegistry, - { - postMessage: mocks.postMessageMock, - } as unknown as Webview, - podmanConnection, - { - getRecipeById: vi.fn().mockReturnValue({ name: 'MyRecipe' } as Recipe), - } as unknown as CatalogManager, - {} as ModelsManager, - {} as TelemetryLogger, - localRepositoryRegistry, - builderManager, - podManager, - ); - }); - - test('init updates the app state with the found pod', async () => { - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([ - { - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - 'ai-lab-app-ports': '5000,5001', - 'ai-lab-model-ports': '8000,8001', - }, - } as unknown as PodInfo, - ]); - mocks.startupSubscribeMock.mockImplementation((f: startupHandle) => { - f(); - }); - const updateApplicationStateSpy = vi.spyOn(manager, 'updateApplicationState'); - manager.init(); - await new Promise(resolve => setTimeout(resolve, 0)); - expect(updateApplicationStateSpy).toHaveBeenNthCalledWith(1, 'recipe-id-1', 'model-id-1', { - pod: { - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - 'ai-lab-app-ports': '5000,5001', - 'ai-lab-model-ports': '8000,8001', - }, - }, - recipeId: 'recipe-id-1', - modelId: 'model-id-1', - appPorts: [5000, 5001], - modelPorts: [8000, 8001], - health: 'starting', - }); - const ports = await manager.getApplicationPorts('recipe-id-1', 'model-id-1'); - expect(ports).toStrictEqual([5000, 5001]); - }); - - 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.init(); - await new Promise(resolve => setTimeout(resolve, 0)); - expect(updateApplicationStateSpy).not.toHaveBeenCalled(); - }); - - test('onMachineStop updates the applications state with no application running', async () => { - vi.mocked(podManager.getAllPods).mockResolvedValue([]); - mocks.onMachineStopMock.mockImplementation((f: machineStopHandle) => { - f(); - }); - const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.init(); - expect(sendApplicationStateSpy).toHaveBeenCalledOnce(); - }); - - test('onPodStart updates the applications state with the started pod', async () => { - vi.mocked(podManager.getAllPods).mockResolvedValue([]); - mocks.onMachineStopMock.mockImplementation((_f: machineStopHandle) => {}); - vi.mocked(podManager.onStartPodEvent).mockImplementation((f: (e: PodInfo) => void): Disposable => { - f({ - engineId: 'engine-1', - engineName: 'Engine 1', - kind: 'podman', - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - } as unknown as PodInfo); - return { dispose: vi.fn() }; - }); - const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.init(); - expect(sendApplicationStateSpy).toHaveBeenCalledOnce(); - }); - - test('onPodStart does no update the applications state with the started pod without labels', async () => { - vi.mocked(podManager.getAllPods).mockResolvedValue([]); - mocks.onMachineStopMock.mockImplementation((_f: machineStopHandle) => {}); - vi.mocked(podManager.onStartPodEvent).mockImplementation((f: (e: PodInfo) => void): Disposable => { - f({ - engineId: 'engine-1', - engineName: 'Engine 1', - kind: 'podman', - } as unknown as PodInfo); - return { dispose: vi.fn() }; - }); - const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.init(); - expect(sendApplicationStateSpy).not.toHaveBeenCalledOnce(); - }); - - test('onPodStart does no update the applications state with the started pod without specific labels', async () => { - vi.mocked(podManager.getAllPods).mockResolvedValue([]); - mocks.onMachineStopMock.mockImplementation((_f: machineStopHandle) => {}); - vi.mocked(podManager.onStartPodEvent).mockImplementation((f: (e: PodInfo) => void): Disposable => { - f({ - engineId: 'engine-1', - engineName: 'Engine 1', - kind: 'podman', - Labels: { - label1: 'value1', - }, - } as unknown as PodInfo); - return { dispose: vi.fn() }; - }); - const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.init(); - expect(sendApplicationStateSpy).not.toHaveBeenCalledOnce(); - }); - - test('onPodStop updates the applications state by removing the stopped pod', async () => { - mocks.startupSubscribeMock.mockImplementation((f: startupHandle) => { - f(); - }); - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([ - { - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - } as unknown as PodInfo, - ]); - mocks.onMachineStopMock.mockImplementation((_f: machineStopHandle) => {}); - vi.mocked(podManager.onStopPodEvent).mockImplementation((f: (e: PodInfo) => void): Disposable => { - setTimeout(() => { - f({ - engineId: 'engine-1', - engineName: 'Engine 1', - kind: 'podman', - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - } as unknown as PodInfo); - return { dispose: vi.fn() }; - }, 1); - return { dispose: vi.fn() }; - }); - const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.init(); - await new Promise(resolve => setTimeout(resolve, 10)); - expect(sendApplicationStateSpy).toHaveBeenCalledTimes(1); - }); - - test('onPodRemove updates the applications state by removing the removed pod', async () => { - mocks.startupSubscribeMock.mockImplementation((f: startupHandle) => { - f(); - }); - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([ - { - Id: 'pod-id-1', - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - } as unknown as PodInfo, - ]); - mocks.onMachineStopMock.mockImplementation((_f: machineStopHandle) => {}); - vi.mocked(podManager.onRemovePodEvent).mockImplementation((f: (e: PodEvent) => void): Disposable => { - setTimeout(() => { - f({ - podId: 'pod-id-1', - }); - }, 1); - return { dispose: vi.fn() }; - }); - const sendApplicationStateSpy = vi.spyOn(manager, 'notify').mockResolvedValue(); - manager.init(); - await new Promise(resolve => setTimeout(resolve, 10)); - expect(sendApplicationStateSpy).toHaveBeenCalledTimes(2); - }); - - test('getApplicationPod', async () => { - vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - } as unknown as PodInfo); - const result = await manager.getApplicationPod('recipe-id-1', 'model-id-1'); - expect(result).toEqual({ - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - }); - expect(podManager.findPodByLabelsValues).toHaveBeenCalledWith({ - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }); - }); - - test('removeApplication calls stopPod and removePod', async () => { - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([]); - vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ - engineId: 'engine-1', - Id: 'pod-1', - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - } as unknown as PodInfo); - await manager.removeApplication('recipe-id-1', 'model-id-1'); - expect(podManager.stopPod).toHaveBeenCalledWith('engine-1', 'pod-1'); - expect(podManager.removePod).toHaveBeenCalledWith('engine-1', 'pod-1'); - }); - - test('removeApplication calls stopPod and removePod even if stopPod fails because pod already stopped', async () => { - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([]); - vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({ - engineId: 'engine-1', - Id: 'pod-1', - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - }, - } as unknown as PodInfo); - vi.mocked(podManager.stopPod).mockRejectedValue('something went wrong, pod already stopped...'); - await manager.removeApplication('recipe-id-1', 'model-id-1'); - expect(podManager.stopPod).toHaveBeenCalledWith('engine-1', 'pod-1'); - expect(podManager.removePod).toHaveBeenCalledWith('engine-1', 'pod-1'); - }); - - test('init should check pods health', async () => { - vi.mocked(podManager.getHealth).mockResolvedValue('healthy'); - vi.mocked(podManager.getPodsWithLabels).mockResolvedValue([ - { - Id: 'pod1', - engineId: 'engine1', - Labels: { - 'ai-lab-recipe-id': 'recipe-id-1', - 'ai-lab-model-id': 'model-id-1', - 'ai-lab-app-ports': '5000,5001', - 'ai-lab-model-ports': '8000,8001', - }, - Containers: [ - { - Id: 'container1', - }, - { - Id: 'container2', - }, - { - Id: 'container3', - }, - ], - } as unknown as PodInfo, - ]); - mocks.startupSubscribeMock.mockImplementation((f: startupHandle) => { - f(); - }); - vi.useFakeTimers(); - manager.init(); - await vi.advanceTimersByTimeAsync(1100); - const state = manager.getApplicationsState(); - expect(state).toHaveLength(1); - expect(state[0].health).toEqual('healthy'); - - expect(podManager.getHealth).toHaveBeenCalledWith({ - Id: 'pod1', - engineId: 'engine1', - Labels: expect.anything(), - Containers: expect.anything(), - }); - }); -}); diff --git a/packages/backend/src/managers/recipes/BuilderManager.ts b/packages/backend/src/managers/recipes/BuilderManager.ts index 451293073..c39ed3ad1 100644 --- a/packages/backend/src/managers/recipes/BuilderManager.ts +++ b/packages/backend/src/managers/recipes/BuilderManager.ts @@ -17,15 +17,19 @@ ***********************************************************************/ import { type BuildImageOptions, type Disposable, containerEngine } from '@podman-desktop/api'; import type { TaskRegistry } from '../../registries/TaskRegistry'; -import type { Recipe } from '@shared/src/models/IRecipe'; +import type { RecipeImage, Recipe } from '@shared/src/models/IRecipe'; import type { ContainerConfig } from '../../models/AIConfig'; -import type { ImageInfo } from '../applicationManager'; -import { LABEL_RECIPE_ID } from '../applicationManager'; import type { Task } from '@shared/src/models/ITask'; import path from 'node:path'; import { getParentDirectory } from '../../utils/pathUtils'; import fs from 'fs'; import { getImageTag } from '../../utils/imagesUtils'; +import { + IMAGE_LABEL_APP_PORTS, + IMAGE_LABEL_APPLICATION_NAME, + IMAGE_LABEL_MODEL_SERVICE, + IMAGE_LABEL_RECIPE_ID, +} from '../../utils/RecipeConstants'; export class BuilderManager implements Disposable { private controller: Map = new Map(); @@ -43,8 +47,8 @@ export class BuilderManager implements Disposable { recipe: Recipe, containers: ContainerConfig[], configPath: string, - labels?: { [key: string]: string }, - ): Promise { + labels: { [key: string]: string } = {}, + ): Promise { const containerTasks: { [key: string]: Task } = Object.fromEntries( containers.map(container => [ container.name, @@ -52,7 +56,7 @@ export class BuilderManager implements Disposable { ]), ); - const imageInfoList: ImageInfo[] = []; + const imageInfoList: RecipeImage[] = []; // Promise all the build images const abortController = new AbortController(); @@ -85,7 +89,11 @@ export class BuilderManager implements Disposable { containerFile: container.containerfile, tag: imageTag, labels: { - [LABEL_RECIPE_ID]: labels !== undefined && 'recipe-id' in labels ? labels['recipe-id'] : '', + ...labels, + [IMAGE_LABEL_RECIPE_ID]: recipe.id, + [IMAGE_LABEL_MODEL_SERVICE]: container.modelService ? 'true' : 'false', + [IMAGE_LABEL_APPLICATION_NAME]: container.name, + [IMAGE_LABEL_APP_PORTS]: (container.ports ?? []).join(','), }, abortController: abortController, }; @@ -142,11 +150,19 @@ export class BuilderManager implements Disposable { throw new Error(`no image found for ${container.name}:latest`); } + let imageName: string | undefined = undefined; + if (image.RepoTags && image.RepoTags.length > 0) { + imageName = image.RepoTags[0]; + } + imageInfoList.push({ id: image.Id, + engineId: image.engineId, + name: imageName, modelService: container.modelService, ports: container.ports?.map(p => `${p}`) ?? [], appName: container.name, + recipeId: recipe.id, }); task.state = 'success'; diff --git a/packages/backend/src/managers/recipes/RecipeManager.spec.ts b/packages/backend/src/managers/recipes/RecipeManager.spec.ts new file mode 100644 index 000000000..451605009 --- /dev/null +++ b/packages/backend/src/managers/recipes/RecipeManager.spec.ts @@ -0,0 +1,207 @@ +/********************************************************************** + * 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 { beforeEach, describe, expect, test, vi } from 'vitest'; +import type { TaskRegistry } from '../../registries/TaskRegistry'; +import type { BuilderManager } from './BuilderManager'; +import type { GitManager } from '../gitManager'; +import type { LocalRepositoryRegistry } from '../../registries/LocalRepositoryRegistry'; +import { RecipeManager } from './RecipeManager'; +import { containerEngine } from '@podman-desktop/api'; +import type { Recipe } from '@shared/src/models/IRecipe'; +import type { Stats } from 'node:fs'; +import { existsSync, statSync } from 'node:fs'; +import { parseYamlFile } from '../../models/AIConfig'; +import { goarch } from '../../utils/arch'; + +const taskRegistryMock = { + createTask: vi.fn(), + updateTask: vi.fn(), +} as unknown as TaskRegistry; + +const builderManagerMock = { + build: vi.fn(), +} as unknown as BuilderManager; + +const gitManagerMock = { + processCheckout: vi.fn(), +} as unknown as GitManager; + +const localRepositoriesMock = { + register: vi.fn(), +} as unknown as LocalRepositoryRegistry; + +const recipeMock: Recipe = { + id: 'recipe-test', + name: 'Test Recipe', + categories: [], + description: 'test recipe description', + repository: 'http://test-repository.test', + readme: 'test recipe readme', +}; + +vi.mock('../../models/AIConfig', () => ({ + parseYamlFile: vi.fn(), +})); + +vi.mock('node:fs', () => ({ + existsSync: vi.fn(), + statSync: vi.fn(), +})); + +vi.mock('@podman-desktop/api', () => ({ + containerEngine: { + listImages: vi.fn(), + }, +})); + +vi.mock('../../utils/arch', () => ({ + goarch: vi.fn(), +})); + +beforeEach(() => { + vi.resetAllMocks(); + + vi.mocked(containerEngine.listImages).mockResolvedValue([]); + vi.mocked(taskRegistryMock.createTask).mockImplementation((name, state, labels) => ({ + name, + state, + labels, + id: 'fake-task', + })); + + vi.mocked(existsSync).mockReturnValue(true); + vi.mocked(statSync).mockReturnValue({ + isDirectory: () => true, + } as unknown as Stats); + + vi.mocked(parseYamlFile).mockReturnValue({ + application: { + containers: [ + { + arch: ['dummy-arch'], + modelService: false, + name: 'test-container', + gpu_env: [], + contextdir: '.', + }, + ], + }, + }); + + vi.mocked(goarch).mockReturnValue('dummy-arch'); +}); + +async function getInitializedRecipeManager(): Promise { + const manager = new RecipeManager( + 'test-app-user-directory', + gitManagerMock, + taskRegistryMock, + builderManagerMock, + localRepositoriesMock, + ); + manager.init(); + return manager; +} + +describe('cloneRecipe', () => { + test('error in checkout should set the task to error and propagate it', async () => { + vi.mocked(gitManagerMock.processCheckout).mockRejectedValue(new Error('clone error')); + + const manager = await getInitializedRecipeManager(); + + await expect(() => { + return manager.cloneRecipe(recipeMock); + }).rejects.toThrowError('clone error'); + + expect(taskRegistryMock.updateTask).toHaveBeenCalledWith( + expect.objectContaining({ + state: 'error', + }), + ); + }); + + test('labels should be propagated', async () => { + const manager = await getInitializedRecipeManager(); + await manager.cloneRecipe(recipeMock, { + 'test-label': 'test-value', + }); + + expect(gitManagerMock.processCheckout).toHaveBeenCalledWith({ + repository: recipeMock.repository, + ref: recipeMock.ref, + targetDirectory: expect.any(String), + }); + + expect(taskRegistryMock.createTask).toHaveBeenCalledWith('Checking out repository', 'loading', { + 'test-label': 'test-value', + 'recipe-id': recipeMock.id, + git: 'checkout', + }); + + expect(localRepositoriesMock.register).toHaveBeenCalledWith({ + path: expect.any(String), + sourcePath: expect.any(String), + labels: { + 'recipe-id': recipeMock.id, + }, + }); + }); +}); + +describe('buildRecipe', () => { + test('error in build propagate it', async () => { + vi.mocked(builderManagerMock.build).mockRejectedValue(new Error('build error')); + + const manager = await getInitializedRecipeManager(); + + await expect(() => { + return manager.buildRecipe(recipeMock); + }).rejects.toThrowError('build error'); + }); + + test('labels should be propagated', async () => { + const manager = await getInitializedRecipeManager(); + + await manager.buildRecipe(recipeMock, { + 'test-label': 'test-value', + }); + + expect(taskRegistryMock.createTask).toHaveBeenCalledWith('Loading configuration', 'loading', { + 'test-label': 'test-value', + 'recipe-id': recipeMock.id, + }); + + expect(builderManagerMock.build).toHaveBeenCalledWith( + recipeMock, + [ + { + arch: ['dummy-arch'], + modelService: false, + name: 'test-container', + gpu_env: [], + contextdir: '.', + }, + ], + expect.any(String), + { + 'test-label': 'test-value', + 'recipe-id': recipeMock.id, + }, + ); + }); +}); diff --git a/packages/backend/src/managers/recipes/RecipeManager.ts b/packages/backend/src/managers/recipes/RecipeManager.ts new file mode 100644 index 000000000..b23eacc27 --- /dev/null +++ b/packages/backend/src/managers/recipes/RecipeManager.ts @@ -0,0 +1,194 @@ +/********************************************************************** + * 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 { GitCloneInfo, GitManager } from '../gitManager'; +import type { TaskRegistry } from '../../registries/TaskRegistry'; +import type { Recipe, RecipeImage } from '@shared/src/models/IRecipe'; +import path from 'node:path'; +import type { Task } from '@shared/src/models/ITask'; +import type { LocalRepositoryRegistry } from '../../registries/LocalRepositoryRegistry'; +import type { AIConfig, AIConfigFile, ContainerConfig } from '../../models/AIConfig'; +import { parseYamlFile } from '../../models/AIConfig'; +import { existsSync, statSync } from 'node:fs'; +import { goarch } from '../../utils/arch'; +import type { BuilderManager } from './BuilderManager'; +import { type Disposable } from '@podman-desktop/api'; +import { CONFIG_FILENAME } from '../../utils/RecipeConstants'; + +export interface AIContainers { + aiConfigFile: AIConfigFile; + containers: ContainerConfig[]; +} + +export class RecipeManager implements Disposable { + constructor( + private appUserDirectory: string, + private git: GitManager, + private taskRegistry: TaskRegistry, + private builderManager: BuilderManager, + private localRepositories: LocalRepositoryRegistry, + ) {} + + dispose(): void {} + + init(): void {} + + private async doCheckout(gitCloneInfo: GitCloneInfo, labels?: { [id: string]: string }): Promise { + // Creating checkout task + const checkoutTask: Task = this.taskRegistry.createTask('Checking out repository', 'loading', { + ...labels, + git: 'checkout', + }); + + try { + await this.git.processCheckout(gitCloneInfo); + checkoutTask.state = 'success'; + } catch (err: unknown) { + checkoutTask.state = 'error'; + checkoutTask.error = String(err); + // propagate error + throw err; + } finally { + // Update task registry + this.taskRegistry.updateTask(checkoutTask); + } + } + + public async cloneRecipe(recipe: Recipe, labels?: { [key: string]: string }): Promise { + const localFolder = path.join(this.appUserDirectory, recipe.id); + + // clone the recipe repository on the local folder + const gitCloneInfo: GitCloneInfo = { + repository: recipe.repository, + ref: recipe.ref, + targetDirectory: localFolder, + }; + await this.doCheckout(gitCloneInfo, { + ...labels, + 'recipe-id': recipe.id, + }); + + this.localRepositories.register({ + path: gitCloneInfo.targetDirectory, + sourcePath: path.join(gitCloneInfo.targetDirectory, recipe.basedir ?? ''), + labels: { + 'recipe-id': recipe.id, + }, + }); + } + + public async buildRecipe(recipe: Recipe, labels?: { [key: string]: string }): Promise { + const localFolder = path.join(this.appUserDirectory, recipe.id); + + // load and parse the recipe configuration file and filter containers based on architecture + const configAndFilteredContainers = this.getConfigAndFilterContainers(recipe.basedir, localFolder, { + ...labels, + 'recipe-id': recipe.id, + }); + + return await this.builderManager.build( + recipe, + configAndFilteredContainers.containers, + configAndFilteredContainers.aiConfigFile.path, + { + ...labels, + 'recipe-id': recipe.id, + }, + ); + } + + private getConfigAndFilterContainers( + recipeBaseDir: string | undefined, + localFolder: string, + labels?: { [key: string]: string }, + ): AIContainers { + // Adding loading configuration task + const task = this.taskRegistry.createTask('Loading configuration', 'loading', labels); + + let aiConfigFile: AIConfigFile; + try { + // load and parse the recipe configuration file + aiConfigFile = this.getConfiguration(recipeBaseDir, localFolder); + } catch (e) { + task.error = `Something went wrong while loading configuration: ${String(e)}.`; + this.taskRegistry.updateTask(task); + throw e; + } + + // filter the containers based on architecture, gpu accelerator and backend (that define which model supports) + const filteredContainers: ContainerConfig[] = this.filterContainers(aiConfigFile.aiConfig); + if (filteredContainers.length > 0) { + // Mark as success. + task.state = 'success'; + this.taskRegistry.updateTask(task); + } else { + // Mark as failure. + task.error = 'No containers available.'; + this.taskRegistry.updateTask(task); + throw new Error('No containers available.'); + } + + return { + aiConfigFile: aiConfigFile, + containers: filteredContainers, + }; + } + + private filterContainers(aiConfig: AIConfig): ContainerConfig[] { + return aiConfig.application.containers.filter( + container => container.gpu_env.length === 0 && container.arch.some(arc => arc === goarch()), + ); + } + + private getConfiguration(recipeBaseDir: string | undefined, localFolder: string): AIConfigFile { + let configFile: string; + if (recipeBaseDir !== undefined) { + configFile = path.join(localFolder, recipeBaseDir, CONFIG_FILENAME); + } else { + configFile = path.join(localFolder, CONFIG_FILENAME); + } + + if (!existsSync(configFile)) { + throw new Error(`The file located at ${configFile} does not exist.`); + } + + // If the user configured the config as a directory we check for "ai-lab.yaml" inside. + if (statSync(configFile).isDirectory()) { + const tmpPath = path.join(configFile, CONFIG_FILENAME); + // If it has the ai-lab.yaml we use it. + if (existsSync(tmpPath)) { + configFile = tmpPath; + } + } + + // Parsing the configuration + console.log(`Reading configuration from ${configFile}.`); + let aiConfig: AIConfig; + try { + aiConfig = parseYamlFile(configFile, goarch()); + } catch (err) { + console.error('Cannot load configure file.', err); + throw new Error(`Cannot load configuration file.`); + } + + // Mark as success. + return { + aiConfig, + path: configFile, + }; + } +} diff --git a/packages/backend/src/studio-api-impl.spec.ts b/packages/backend/src/studio-api-impl.spec.ts index 572f73678..61440f1c8 100644 --- a/packages/backend/src/studio-api-impl.spec.ts +++ b/packages/backend/src/studio-api-impl.spec.ts @@ -20,7 +20,7 @@ import { beforeEach, expect, test, vi, describe } from 'vitest'; import content from './tests/ai-test.json'; -import type { ApplicationManager } from './managers/applicationManager'; +import type { ApplicationManager } from './managers/application/applicationManager'; import { StudioApiImpl } from './studio-api-impl'; import type { InferenceManager } from './managers/inference/inferenceManager'; import type { ProviderContainerConnection, TelemetryLogger, Webview } from '@podman-desktop/api'; @@ -39,6 +39,7 @@ import path from 'node:path'; import type { LocalModelImportInfo } from '@shared/src/models/ILocalModelInfo'; import * as podman from './utils/podman'; import type { ConfigurationRegistry } from './registries/ConfigurationRegistry'; +import type { RecipeManager } from './managers/recipes/RecipeManager'; vi.mock('./ai.json', () => { return { @@ -145,6 +146,7 @@ beforeEach(async () => { {} as unknown as SnippetManager, {} as unknown as CancellationTokenRegistry, {} as unknown as ConfigurationRegistry, + {} as unknown as RecipeManager, ); vi.mock('node:fs'); diff --git a/packages/backend/src/studio-api-impl.ts b/packages/backend/src/studio-api-impl.ts index 8f73f34a3..4b41d6c2c 100644 --- a/packages/backend/src/studio-api-impl.ts +++ b/packages/backend/src/studio-api-impl.ts @@ -17,7 +17,7 @@ ***********************************************************************/ import type { StudioAPI } from '@shared/src/StudioAPI'; -import type { ApplicationManager } from './managers/applicationManager'; +import type { ApplicationManager } from './managers/application/applicationManager'; import type { ModelCheckerInfo, ModelInfo } from '@shared/src/models/IModelInfo'; import * as podmanDesktopApi from '@podman-desktop/api'; @@ -47,6 +47,7 @@ import { checkContainerConnectionStatusAndResources, getPodmanConnection } from import type { ContainerConnectionInfo } from '@shared/src/models/IContainerConnectionInfo'; import type { ExtensionConfiguration } from '@shared/src/models/IExtensionConfiguration'; import type { ConfigurationRegistry } from './registries/ConfigurationRegistry'; +import type { RecipeManager } from './managers/recipes/RecipeManager'; interface PortQuickPickItem extends podmanDesktopApi.QuickPickItem { port: number; @@ -65,6 +66,7 @@ export class StudioApiImpl implements StudioAPI { private snippetManager: SnippetManager, private cancellationTokenRegistry: CancellationTokenRegistry, private configurationRegistry: ConfigurationRegistry, + private recipeManager: RecipeManager, ) {} async requestDeleteConversation(conversationId: string): Promise { @@ -193,7 +195,7 @@ export class StudioApiImpl implements StudioAPI { const recipe = this.catalogManager.getRecipes().find(recipe => recipe.id === recipeId); if (!recipe) throw new Error(`recipe with if ${recipeId} not found`); - return this.applicationManager.cloneApplication(recipe); + return this.recipeManager.cloneRecipe(recipe); } async requestPullApplication(recipeId: string, modelId: string): Promise { diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index 51719948e..f96f24fa7 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -26,7 +26,7 @@ import type { } from '@podman-desktop/api'; import { RpcExtension } from '@shared/src/messages/MessageProxy'; import { StudioApiImpl } from './studio-api-impl'; -import { ApplicationManager } from './managers/applicationManager'; +import { ApplicationManager } from './managers/application/applicationManager'; import { GitManager } from './managers/gitManager'; import { TaskRegistry } from './registries/TaskRegistry'; import { CatalogManager } from './managers/catalogManager'; @@ -45,6 +45,7 @@ import { initWebview } from './webviewUtils'; import { LlamaCppPython } from './workers/provider/LlamaCppPython'; import { InferenceProviderRegistry } from './registries/InferenceProviderRegistry'; import { ConfigurationRegistry } from './registries/ConfigurationRegistry'; +import { RecipeManager } from './managers/recipes/RecipeManager'; export class Studio { readonly #extensionContext: ExtensionContext; @@ -74,6 +75,7 @@ export class Studio { #snippetManager: SnippetManager | undefined; #playgroundManager: PlaygroundV2Manager | undefined; #applicationManager: ApplicationManager | undefined; + #recipeManager: RecipeManager | undefined; #inferenceProviderRegistry: InferenceProviderRegistry | undefined; #configurationRegistry: ConfigurationRegistry | undefined; @@ -217,20 +219,30 @@ export class Studio { this.#extensionContext.subscriptions.push(this.#localRepositoryRegistry); /** - * The application manager is managing the Recipes + * The recipe manage offer some andy methods to manage recipes, build get images etc. */ - this.#applicationManager = new ApplicationManager( + this.#recipeManager = new RecipeManager( appUserDirectory, gitManager, + this.#taskRegistry, + this.#builderManager, + this.#localRepositoryRegistry, + ); + this.#recipeManager.init(); + this.#extensionContext.subscriptions.push(this.#recipeManager); + + /** + * The application manager is managing the Recipes + */ + this.#applicationManager = new ApplicationManager( this.#taskRegistry, this.#panel.webview, this.#podmanConnection, this.#catalogManager, this.#modelsManager, this.#telemetry, - this.#localRepositoryRegistry, - this.#builderManager, this.#podManager, + this.#recipeManager, ); this.#applicationManager.init(); this.#extensionContext.subscriptions.push(this.#applicationManager); @@ -293,6 +305,7 @@ export class Studio { this.#snippetManager, this.#cancellationTokenRegistry, this.#configurationRegistry, + this.#recipeManager, ); // Register the instance this.#rpcExtension.registerInstance(StudioApiImpl, this.#studioApi); diff --git a/packages/backend/src/utils/RecipeConstants.ts b/packages/backend/src/utils/RecipeConstants.ts new file mode 100644 index 000000000..43e78887a --- /dev/null +++ b/packages/backend/src/utils/RecipeConstants.ts @@ -0,0 +1,31 @@ +/********************************************************************** + * 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 + ***********************************************************************/ + +export const CONFIG_FILENAME = 'ai-lab.yaml'; + +// pod labels +export const POD_LABEL_RECIPE_ID = 'ai-lab-recipe-id'; +export const POD_LABEL_MODEL_ID = 'ai-lab-model-id'; +export const POD_LABEL_MODEL_PORTS = 'ai-lab-model-ports'; +export const POD_LABEL_APP_PORTS = 'ai-lab-application-ports'; + +// image labels +export const IMAGE_LABEL_RECIPE_ID = 'ai-lab-recipe-id'; +export const IMAGE_LABEL_APP_PORTS = 'ai-lab-application-ports'; +export const IMAGE_LABEL_MODEL_SERVICE = 'ai-lab-model-service'; +export const IMAGE_LABEL_APPLICATION_NAME = 'ai-lab-application-name'; diff --git a/packages/shared/src/models/IRecipe.ts b/packages/shared/src/models/IRecipe.ts index cda534550..ab2384cde 100644 --- a/packages/shared/src/models/IRecipe.ts +++ b/packages/shared/src/models/IRecipe.ts @@ -16,6 +16,17 @@ * SPDX-License-Identifier: Apache-2.0 ***********************************************************************/ +export interface RecipeImage { + id: string; + engineId: string; + name?: string; + // recipe related + recipeId: string; + modelService: boolean; + ports: string[]; + appName: string; +} + export interface Recipe { id: string; name: string;