From 6e2adffd92f4ab1ec45d3b4a451bea6685727063 Mon Sep 17 00:00:00 2001
From: Philippe Martin <phmartin@redhat.com>
Date: Tue, 30 Jan 2024 17:21:08 +0100
Subject: [PATCH] adopt playground containers

---
 packages/backend/package.json               |  2 +-
 packages/backend/src/managers/playground.ts | 33 ++++++++++++++++++++-
 packages/backend/src/studio.ts              |  1 +
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/packages/backend/package.json b/packages/backend/package.json
index 1642a89f3..a1f0538cf 100644
--- a/packages/backend/package.json
+++ b/packages/backend/package.json
@@ -23,7 +23,7 @@
     "views": {
       "icons/containersList": [
         {
-          "when": "ia-studio-model in containerLabelKeys",
+          "when": "ai-studio-model-id in containerLabelKeys",
           "icon": "${brain-icon}"
         }
       ]
diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts
index 6c7de8f8e..9cda4089b 100644
--- 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
index 9b4563706..cfe6c7560 100644
--- 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);