fix: loadLocalModels issue when models being used (#304)

* fix: deletion models issue Signed-off-by: axel7083 <[email protected]> * fix: prettier&linter Signed-off-by: axel7083 <[email protected]> * fix: prettier&linter Signed-off-by: axel7083 <[email protected]> --------- Signed-off-by: axel7083 <[email protected]>
containers · Feb 14, 2024 · defd5d5 · defd5d5
1 parent cf6bba8
commit defd5d5
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 9 deletions.
diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts
@@ -18,7 +18,7 @@
 
 import { type MockInstance, beforeEach, describe, expect, test, vi } from 'vitest';
 import os from 'os';
-import fs from 'node:fs';
+import fs, { type Stats, type PathLike } from 'node:fs';
 import path from 'node:path';
 import type { DownloadModelResult } from './modelsManager';
 import { ModelsManager } from './modelsManager';
@@ -198,6 +198,48 @@ test('getModelsInfo should return an empty array if the models folder does not e
   }
 });
 
+test('getLocalModelsFromDisk should return undefined Date and size when stat fail', async () => {
+  const now = new Date();
+  mockFiles(now);
+  const statSyncSpy = vi.spyOn(fs, 'statSync');
+  statSyncSpy.mockImplementation((path: PathLike) => {
+    if (`${path}`.endsWith('model-id-1')) throw new Error('random-error');
+    return { isDirectory: () => true } as Stats;
+  });
+
+  let appdir: string;
+  if (process.platform === 'win32') {
+    appdir = 'C:\\home\\user\\aistudio';
+  } else {
+    appdir = '/home/user/aistudio';
+  }
+  const manager = new ModelsManager(
+    appdir,
+    {
+      postMessage: vi.fn(),
+    } as unknown as Webview,
+    {
+      getModels(): ModelInfo[] {
+        return [{ id: 'model-id-1', name: 'model-id-1-model' } as ModelInfo];
+      },
+    } as CatalogManager,
+    telemetryLogger,
+  );
+  await manager.loadLocalModels();
+  expect(manager.getModelsInfo()).toEqual([
+    {
+      id: 'model-id-1',
+      name: 'model-id-1-model',
+      file: {
+        size: undefined,
+        creation: undefined,
+        path: path.resolve(dirent[0].path, dirent[0].name),
+        file: 'model-id-1-model',
+      },
+    },
+  ]);
+});
+
 test('loadLocalModels should post a message with the message on disk and on catalog', async () => {
   const now = new Date();
   mockFiles(now);
@@ -334,6 +376,12 @@ test('deleteLocalModel fails to delete the model folder', async () => {
     body: [
       {
         id: 'model-id-1',
+        file: {
+          creation: now,
+          file: 'model-id-1-model',
+          size: 32000,
+          path: path.resolve(dirent[0].path, dirent[0].name),
+        },
       },
     ],
   });

diff --git a/packages/backend/src/managers/modelsManager.ts b/packages/backend/src/managers/modelsManager.ts
@@ -43,6 +43,7 @@ interface DownloadModelFailureResult {
 export class ModelsManager {
   #modelsDir: string;
   #models: Map<string, ModelInfo>;
+  #watcher?: podmanDesktopApi.FileSystemWatcher;
 
   constructor(
     private appUserDirectory: string,
@@ -60,10 +61,13 @@ export class ModelsManager {
       this.getLocalModelsFromDisk();
       await this.sendModelsInfo();
     };
-    const watcher = apiFs.createFileSystemWatcher(this.#modelsDir);
-    watcher.onDidCreate(reloadLocalModels);
-    watcher.onDidDelete(reloadLocalModels);
-    watcher.onDidChange(reloadLocalModels);
+    if (this.#watcher === undefined) {
+      this.#watcher = apiFs.createFileSystemWatcher(this.#modelsDir);
+      this.#watcher.onDidCreate(reloadLocalModels);
+      this.#watcher.onDidDelete(reloadLocalModels);
+      this.#watcher.onDidChange(reloadLocalModels);
+    }
+
     // Initialize the local models manually
     await reloadLocalModels();
   }
@@ -98,7 +102,14 @@ export class ModelsManager {
       }
       const modelFile = modelEntries[0];
       const fullPath = path.resolve(d.path, d.name, modelFile);
-      const info = fs.statSync(fullPath);
+
+      let info: { size?: number; mtime?: Date } = { size: undefined, mtime: undefined };
+      try {
+        info = fs.statSync(fullPath);
+      } catch (err: unknown) {
+        console.error('Something went wrong while getting file stats (probably in use).', err);
+      }
+
       const model = this.#models.get(d.name);
       if (model) {
         model.file = {
@@ -148,15 +159,19 @@ export class ModelsManager {
       try {
         await fs.promises.rm(modelDir, { recursive: true });
         this.telemetry.logUsage('model.delete', { 'model.id': modelId });
+        model.file = model.state = undefined;
       } catch (err: unknown) {
         this.telemetry.logError('model.delete', {
           'model.id': modelId,
           message: 'error deleting model from disk',
           error: err,
         });
         await podmanDesktopApi.window.showErrorMessage(`Error deleting model ${modelId}. ${String(err)}`);
+
+        // Let's reload the models manually to avoid any issue
+        model.state = undefined;
+        this.getLocalModelsFromDisk();
       } finally {
-        model.file = model.state = undefined;
         await this.sendModelsInfo();
       }
     }

diff --git a/packages/shared/src/models/ILocalModelInfo.ts b/packages/shared/src/models/ILocalModelInfo.ts
@@ -1,6 +1,6 @@
 export interface LocalModelInfo {
   file: string;
   path: string;
-  size: number;
-  creation: Date;
+  size?: number;
+  creation?: Date;
 }