feat: improve playground error management (#269)

* feat: improving error management in playground system Signed-off-by: axel7083 <[email protected]> * fix: propagating to stop playground Signed-off-by: axel7083 <[email protected]> * fix: prettier Signed-off-by: axel7083 <[email protected]> * fix: keep existing error Signed-off-by: axel7083 <[email protected]> --------- Signed-off-by: axel7083 <[email protected]>
containers · Feb 12, 2024 · d5ea0b1 · d5ea0b1
1 parent f30794d
commit d5ea0b1
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 9 deletions.
diff --git a/packages/backend/src/managers/playground.spec.ts b/packages/backend/src/managers/playground.spec.ts
@@ -217,3 +217,30 @@ test('onMachineStop updates the playground state with no playground running', as
   manager.adoptRunningPlaygrounds();
   expect(sendPlaygroundStateSpy).toHaveBeenCalledOnce();
 });
+
+test('playground error not overwritten', async () => {
+  mocks.postMessage.mockResolvedValue(undefined);
+
+  const states = manager.getPlaygroundsState();
+  expect(states.length).toBe(0);
+
+  manager.setPlaygroundError('random', 'first');
+  expect(manager.getPlaygroundsState().length).toBe(1);
+
+  expect(manager.getPlaygroundsState()[0].error).toBe('first');
+
+  manager.setPlaygroundError('random', 'second');
+  expect(manager.getPlaygroundsState()[0].error).toBe('first');
+});
+
+test('error cleared when status changed', async () => {
+  mocks.postMessage.mockResolvedValue(undefined);
+
+  const states = manager.getPlaygroundsState();
+  expect(states.length).toBe(0);
+
+  manager.setPlaygroundError('random', 'error-msg');
+  manager.setPlaygroundStatus('random', 'running');
+
+  expect(manager.getPlaygroundsState()[0].error).toBeUndefined();
+});
diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts
@@ -122,8 +122,21 @@ export class PlayGroundManager {
     });
   }
 
+  setPlaygroundError(modelId: string, error: string): void {
+    const state: Partial<PlaygroundState> = this.playgrounds.get(modelId) || {};
+    this.updatePlaygroundState(modelId, {
+      modelId: modelId,
+      ...state,
+      status: 'error',
+      error: state.error ?? error, // we never overwrite previous error - we want to keep the first one raised
+    });
+  }
+
   updatePlaygroundState(modelId: string, state: PlaygroundState): void {
-    this.playgrounds.set(modelId, state);
+    this.playgrounds.set(modelId, {
+      ...state,
+      error: state.status === 'error' ? state.error : undefined, // clearing error when status not error
+    });
     this.sendPlaygroundState();
   }
 
@@ -160,25 +173,27 @@ export class PlayGroundManager {
 
     const connection = findFirstProvider();
     if (!connection) {
-      this.setPlaygroundStatus(modelId, 'error');
+      const error = 'Unable to find an engine to start playground';
+      this.setPlaygroundError(modelId, error);
       this.telemetry.logError('playground.start', {
         'model.id': modelId,
-        message: 'unable to find an engine to start playground',
+        message: error,
       });
-      throw new Error('Unable to find an engine to start playground');
+      throw new Error(error);
     }
 
     let image = await this.selectImage(PLAYGROUND_IMAGE);
     if (!image) {
       await containerEngine.pullImage(connection.connection, PLAYGROUND_IMAGE, () => {});
       image = await this.selectImage(PLAYGROUND_IMAGE);
       if (!image) {
-        this.setPlaygroundStatus(modelId, 'error');
+        const error = `Unable to find ${PLAYGROUND_IMAGE} image`;
+        this.setPlaygroundError(modelId, error);
         this.telemetry.logError('playground.start', {
           'model.id': modelId,
           message: 'unable to find playground image',
         });
-        throw new Error(`Unable to find ${PLAYGROUND_IMAGE} image`);
+        throw new Error(error);
       }
     }
 
@@ -277,7 +292,7 @@ export class PlayGroundManager {
       })
       .catch(async (error: unknown) => {
         console.error(error);
-        this.setPlaygroundStatus(modelId, 'error');
+        this.setPlaygroundError(modelId, `Something went wrong while stopping playground: ${String(error)}`);
         this.telemetry.logError('playground.stop', {
           'model.id': modelId,
           message: 'error stopping playground',

diff --git a/packages/backend/src/studio-api-impl.ts b/packages/backend/src/studio-api-impl.ts
@@ -87,11 +87,15 @@ export class StudioApiImpl implements StudioAPI {
 
   async startPlayground(modelId: string): Promise<void> {
     const modelPath = this.modelsManager.getLocalModelPath(modelId);
-    await this.playgroundManager.startPlayground(modelId, modelPath);
+    this.playgroundManager.startPlayground(modelId, modelPath).catch((err: unknown) => {
+      this.playgroundManager.setPlaygroundError(modelId, `Something went wrong while starting the playground: ${err}`);
+    });
   }
 
   async stopPlayground(modelId: string): Promise<void> {
-    await this.playgroundManager.stopPlayground(modelId);
+    this.playgroundManager.stopPlayground(modelId).catch((err: unknown) => {
+      this.playgroundManager.setPlaygroundError(modelId, `Something went wrong while stopping the playground: ${err}`);
+    });
   }
 
   async askPlayground(modelId: string, prompt: string): Promise<number> {

diff --git a/packages/frontend/src/pages/ModelPlayground.svelte b/packages/frontend/src/pages/ModelPlayground.svelte
@@ -44,6 +44,7 @@
       if(playgroundState === undefined) {
         playgroundState = { modelId: model.id,  status: 'none' };
       }
+      error = playgroundState.error ?? error;
     })
 
     return () => {

diff --git a/packages/shared/src/models/IPlaygroundState.ts b/packages/shared/src/models/IPlaygroundState.ts
@@ -8,4 +8,5 @@ export interface PlaygroundState {
   };
   modelId: string;
   status: PlaygroundStatus;
+  error?: string;
 }