adopt playground containers (#200)

* adopt playground containers * Update packages/backend/src/managers/playground.ts Co-authored-by: axel7083 <[email protected]> * fix unit tests * use provider.onDidRegisterContainerConnection to know when providers are up * start adopt when podman only provider starts * create a dedicated class for listening the registration of podman provider * remove the getConnection method from PodmanConnection --------- Co-authored-by: axel7083 <[email protected]>
containers · Feb 2, 2024 · 10d4b0a · 10d4b0a
1 parent 9daccfe
commit 10d4b0a
Show file tree

Hide file tree

Showing 6 changed files with 135 additions and 9 deletions.
diff --git a/packages/backend/package.json b/packages/backend/package.json
@@ -23,7 +23,7 @@
     "views": {
       "icons/containersList": [
         {
-          "when": "ia-studio-model in containerLabelKeys",
+          "when": "ai-studio-model-id in containerLabelKeys",
           "icon": "${brain-icon}"
         }
       ]

diff --git a/packages/backend/src/managers/playground.spec.ts b/packages/backend/src/managers/playground.spec.ts
@@ -20,6 +20,7 @@ import { beforeEach, expect, test, vi } from 'vitest';
 import { PlayGroundManager } from './playground';
 import type { ImageInfo, Webview } from '@podman-desktop/api';
 import type { ContainerRegistry } from '../registries/ContainerRegistry';
+import type { PodmanConnection } from './podmanConnection';
 
 const mocks = vi.hoisted(() => ({
   postMessage: vi.fn(),
@@ -64,6 +65,7 @@ beforeEach(() => {
       postMessage: mocks.postMessage,
     } as unknown as Webview,
     containerRegistryMock,
+    {} as PodmanConnection,
   );
 });
 
@@ -123,7 +125,8 @@ test('startPlayground should download image if not present then create container
     },
     Image: 'image1',
     Labels: {
-      'ia-studio-model': 'model1',
+      'ai-studio-model-id': 'model1',
+      'ai-studio-model-port': '8085',
     },
   });
 });

diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts
@@ -17,11 +17,11 @@
  ***********************************************************************/
 
 import {
-  provider,
   containerEngine,
   type Webview,
-  type ProviderContainerConnection,
   type ImageInfo,
+  type ProviderContainerConnection,
+  provider,
 } from '@podman-desktop/api';
 import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo';
 import type { ModelResponse } from '@shared/src/models/IModelResponse';
@@ -33,6 +33,10 @@ import type { QueryState } from '@shared/src/models/IPlaygroundQueryState';
 import { MSG_NEW_PLAYGROUND_QUERIES_STATE, MSG_PLAYGROUNDS_STATE_UPDATE } from '@shared/Messages';
 import type { PlaygroundState, PlaygroundStatus } from '@shared/src/models/IPlaygroundState';
 import type { ContainerRegistry } from '../registries/ContainerRegistry';
+import type { PodmanConnection } from './podmanConnection';
+
+const LABEL_MODEL_ID = 'ai-studio-model-id';
+const LABEL_MODEL_PORT = 'ai-studio-model-port';
 
 // TODO: this should not be hardcoded
 const PLAYGROUND_IMAGE = 'quay.io/bootsy/playground:v0';
@@ -55,12 +59,48 @@ export class PlayGroundManager {
   constructor(
     private webview: Webview,
     private containerRegistry: ContainerRegistry,
+    private podmanConnection: PodmanConnection,
   ) {
     this.playgrounds = new Map<string, PlaygroundState>();
     this.queries = new Map<number, QueryState>();
   }
 
-  async selectImage(connection: ProviderContainerConnection, image: string): Promise<ImageInfo | undefined> {
+  adoptRunningPlaygrounds() {
+    this.podmanConnection.startupSubscribe(() => {
+      containerEngine
+        .listContainers()
+        .then(containers => {
+          const playgroundContainers = containers.filter(
+            c => LABEL_MODEL_ID in c.Labels && LABEL_MODEL_PORT in c.Labels && c.State === 'running',
+          );
+          for (const containerToAdopt of playgroundContainers) {
+            const modelId = containerToAdopt.Labels[LABEL_MODEL_ID];
+            if (this.playgrounds.has(modelId)) {
+              continue;
+            }
+            const modelPort = parseInt(containerToAdopt.Labels[LABEL_MODEL_PORT], 10);
+            if (isNaN(modelPort)) {
+              continue;
+            }
+            const state: PlaygroundState = {
+              modelId,
+              status: 'running',
+              container: {
+                containerId: containerToAdopt.Id,
+                engineId: containerToAdopt.engineId,
+                port: modelPort,
+              },
+            };
+            this.updatePlaygroundState(modelId, state);
+          }
+        })
+        .catch((err: unknown) => {
+          console.error('error during adoption of existing playground containers', err);
+        });
+    });
+  }
+
+  async selectImage(image: string): Promise<ImageInfo | undefined> {
     const images = (await containerEngine.listImages()).filter(im => im.RepoTags?.some(tag => tag === image));
     return images.length > 0 ? images[0] : undefined;
   }
@@ -110,10 +150,10 @@ export class PlayGroundManager {
       throw new Error('Unable to find an engine to start playground');
     }
 
-    let image = await this.selectImage(connection, PLAYGROUND_IMAGE);
+    let image = await this.selectImage(PLAYGROUND_IMAGE);
     if (!image) {
       await containerEngine.pullImage(connection.connection, PLAYGROUND_IMAGE, () => {});
-      image = await this.selectImage(connection, PLAYGROUND_IMAGE);
+      image = await this.selectImage(PLAYGROUND_IMAGE);
       if (!image) {
         this.setPlaygroundStatus(modelId, 'error');
         throw new Error(`Unable to find ${PLAYGROUND_IMAGE} image`);
@@ -143,7 +183,8 @@ export class PlayGroundManager {
         },
       },
       Labels: {
-        'ia-studio-model': modelId,
+        [LABEL_MODEL_ID]: modelId,
+        [LABEL_MODEL_PORT]: `${freePort}`,
       },
       Env: [`MODEL_PATH=/models/${path.basename(modelPath)}`],
       Cmd: ['--models-path', '/models', '--context-size', '700', '--threads', '4'],

diff --git a/packages/backend/src/managers/podmanConnection.ts b/packages/backend/src/managers/podmanConnection.ts
@@ -0,0 +1,65 @@
+/**********************************************************************
+ * Copyright (C) 2024 Red Hat, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ***********************************************************************/
+
+import { type RegisterContainerConnectionEvent, provider } from '@podman-desktop/api';
+
+type startupHandle = () => void;
+
+export class PodmanConnection {
+  #firstFound = false;
+  #toExecuteAtStartup: startupHandle[] = [];
+
+  init(): void {
+    // In case the extension has not yet registered, we listen for new registrations
+    // and retain the first started podman provider
+    const disposable = provider.onDidRegisterContainerConnection((e: RegisterContainerConnectionEvent) => {
+      if (e.connection.type !== 'podman' || e.connection.status() !== 'started') {
+        return;
+      }
+      if (this.#firstFound) {
+        return;
+      }
+      this.#firstFound = true;
+      for (const f of this.#toExecuteAtStartup) {
+        f();
+      }
+      this.#toExecuteAtStartup = [];
+      disposable.dispose();
+    });
+
+    // In case at least one extension has already registered, we get one started podman provider
+    const engines = provider
+      .getContainerConnections()
+      .filter(connection => connection.connection.type === 'podman')
+      .filter(connection => connection.connection.status() === 'started');
+    if (engines.length > 0) {
+      disposable.dispose();
+      this.#firstFound = true;
+    }
+  }
+
+  // startupSubscribe registers f to be executed when a podman container provider
+  // registers, or immediately if already registered
+  startupSubscribe(f: startupHandle): void {
+    if (this.#firstFound) {
+      f();
+    } else {
+      this.#toExecuteAtStartup.push(f);
+    }
+  }
+}
diff --git a/packages/backend/src/studio.spec.ts b/packages/backend/src/studio.spec.ts
@@ -32,6 +32,11 @@ const mockedExtensionContext = {
 
 const studio = new Studio(mockedExtensionContext);
 
+const mocks = vi.hoisted(() => ({
+  listContainers: vi.fn(),
+  getContainerConnections: vi.fn(),
+}));
+
 vi.mock('@podman-desktop/api', async () => {
   return {
     Uri: class {
@@ -48,6 +53,11 @@ vi.mock('@podman-desktop/api', async () => {
     },
     containerEngine: {
       onEvent: vi.fn(),
+      listContainers: mocks.listContainers,
+    },
+    provider: {
+      onDidRegisterContainerConnection: vi.fn(),
+      getContainerConnections: mocks.getContainerConnections,
     },
   };
 });
@@ -66,6 +76,8 @@ afterEach(() => {
 });
 
 test('check activate ', async () => {
+  mocks.listContainers.mockReturnValue([]);
+  mocks.getContainerConnections.mockReturnValue([]);
   vi.spyOn(fs.promises, 'readFile').mockImplementation(() => {
     return Promise.resolve('<html></html>');
   });

diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts
@@ -31,6 +31,7 @@ import path from 'node:path';
 import os from 'os';
 import fs from 'node:fs';
 import { ContainerRegistry } from './registries/ContainerRegistry';
+import { PodmanConnection } from './managers/podmanConnection';
 
 // TODO: Need to be configured
 export const AI_STUDIO_FOLDER = path.join('podman-desktop', 'ai-studio');
@@ -104,9 +105,11 @@ export class Studio {
 
     this.rpcExtension = new RpcExtension(this.#panel.webview);
     const gitManager = new GitManager();
+
+    const podmanConnection = new PodmanConnection();
     const taskRegistry = new TaskRegistry();
     const recipeStatusRegistry = new RecipeStatusRegistry(taskRegistry, this.#panel.webview);
-    this.playgroundManager = new PlayGroundManager(this.#panel.webview, containerRegistry);
+    this.playgroundManager = new PlayGroundManager(this.#panel.webview, containerRegistry, podmanConnection);
     // Create catalog manager, responsible for loading the catalog files and watching for changes
     this.catalogManager = new CatalogManager(appUserDirectory, this.#panel.webview);
     this.modelsManager = new ModelsManager(appUserDirectory, this.#panel.webview, this.catalogManager);
@@ -128,6 +131,8 @@ export class Studio {
 
     await this.catalogManager.loadCatalog();
     await this.modelsManager.loadLocalModels();
+    podmanConnection.init();
+    this.playgroundManager.adoptRunningPlaygrounds();
 
     // Register the instance
     this.rpcExtension.registerInstance<StudioApiImpl>(StudioApiImpl, this.studioApi);