adopt playground containers

containers · Jan 31, 2024 · 6e2adff · 6e2adff
1 parent a1cb858
commit 6e2adff
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 2 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.ts b/packages/backend/src/managers/playground.ts
@@ -34,6 +34,9 @@ import { MSG_NEW_PLAYGROUND_QUERIES_STATE, MSG_PLAYGROUNDS_STATE_UPDATE } from '
 import type { PlaygroundState, PlaygroundStatus } from '@shared/src/models/IPlaygroundState';
 import type { ContainerRegistry } from '../registries/ContainerRegistry';
 
+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';
 
@@ -60,6 +63,33 @@ export class PlayGroundManager {
     this.queries = new Map<number, QueryState>();
   }
 
+  async adoptRunningPlaygrounds() {
+    const containers = await containerEngine.listContainers();
+    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);
+    }
+  }
+
   async selectImage(connection: ProviderContainerConnection, 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;
@@ -143,7 +173,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/studio.ts b/packages/backend/src/studio.ts
@@ -128,6 +128,7 @@ export class Studio {
 
     await this.catalogManager.loadCatalog();
     await this.modelsManager.loadLocalModels();
+    await this.playgroundManager.adoptRunningPlaygrounds();
 
     // Register the instance
     this.rpcExtension.registerInstance<StudioApiImpl>(StudioApiImpl, this.studioApi);