feat: improve podman cli execution

axel7083 · axel7083 · commit 1f0b6f546e6a · 2024-10-01T10:16:51.000+02:00
Signed-off-by: axel7083 &lt;42176370+axel7083@users.noreply.github.com&gt;
diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts
@@ -21,7 +21,7 @@ import os from 'node:os';
 import fs, { type Stats, type PathLike } from 'node:fs';
 import path from 'node:path';
 import { ModelsManager } from './modelsManager';
-import { env, process as coreProcess } from '@podman-desktop/api';
+import { env } from '@podman-desktop/api';
 import type { RunResult, TelemetryLogger, Webview, ContainerProviderConnection } from '@podman-desktop/api';
 import type { CatalogManager } from './catalogManager';
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
@@ -33,7 +33,6 @@ import type { GGUFParseOutput } from '@huggingface/gguf';
 import { gguf } from '@huggingface/gguf';
 import type { PodmanConnection } from './podmanConnection';
 import { VMType } from '@shared/src/models/IPodman';
-import { getPodmanMachineName } from '../utils/podman';
 import type { ConfigurationRegistry } from '../registries/ConfigurationRegistry';
 import { Uploader } from '../utils/uploader';
 
@@ -47,7 +46,6 @@ const mocks = vi.hoisted(() => {
     getTargetMock: vi.fn(),
     getDownloaderCompleter: vi.fn(),
     isCompletionEventMock: vi.fn(),
-    getPodmanCliMock: vi.fn(),
   };
 });
 
@@ -59,11 +57,6 @@ vi.mock('@huggingface/gguf', () => ({
   gguf: vi.fn(),
 }));
 
-vi.mock('../utils/podman', () => ({
-  getPodmanCli: mocks.getPodmanCliMock,
-  getPodmanMachineName: vi.fn(),
-}));
-
 vi.mock('@podman-desktop/api', () => {
   return {
     Disposable: {
@@ -72,9 +65,6 @@ vi.mock('@podman-desktop/api', () => {
     env: {
       isWindows: false,
     },
-    process: {
-      exec: vi.fn(),
-    },
     fs: {
       createFileSystemWatcher: (): unknown => ({
         onDidCreate: vi.fn(),
@@ -102,6 +92,7 @@ vi.mock('../utils/downloader', () => ({
 
 const podmanConnectionMock = {
   getContainerProviderConnections: vi.fn(),
+  executeSSH: vi.fn(),
 } as unknown as PodmanConnection;
 
 const cancellationTokenRegistryMock = {
@@ -598,8 +589,7 @@ describe('deleting models', () => {
   });
 
   test('deleting on windows should check for all connections', async () => {
-    vi.mocked(coreProcess.exec).mockResolvedValue({} as RunResult);
-    mocks.getPodmanCliMock.mockReturnValue('dummyCli');
+    vi.mocked(podmanConnectionMock.executeSSH).mockResolvedValue({} as RunResult);
     vi.mocked(env).isWindows = true;
     const connections: ContainerProviderConnection[] = [
       {
@@ -622,7 +612,6 @@ describe('deleting models', () => {
       },
     ];
     vi.mocked(podmanConnectionMock.getContainerProviderConnections).mockReturnValue(connections);
-    vi.mocked(getPodmanMachineName).mockReturnValue('machine-2');
 
     const rmSpy = vi.spyOn(fs.promises, 'rm');
     rmSpy.mockResolvedValue(undefined);
@@ -659,10 +648,7 @@ describe('deleting models', () => {
 
     expect(podmanConnectionMock.getContainerProviderConnections).toHaveBeenCalledOnce();
 
-    expect(coreProcess.exec).toHaveBeenCalledWith('dummyCli', [
-      'machine',
-      'ssh',
-      'machine-2',
+    expect(podmanConnectionMock.executeSSH).toHaveBeenCalledWith(connections[1], [
       'rm',
       '-f',
       '/home/user/ai-lab/models/dummyFile',
diff --git a/packages/backend/src/managers/modelsManager.ts b/packages/backend/src/managers/modelsManager.ts
@@ -30,8 +30,7 @@ import type { Task } from '@shared/src/models/ITask';
 import type { BaseEvent } from '../models/baseEvent';
 import { isCompletionEvent, isProgressEvent } from '../models/baseEvent';
 import { Uploader } from '../utils/uploader';
-import { deleteRemoteModel, getLocalModelFile, isModelUploaded } from '../utils/modelsUtils';
-import { getPodmanMachineName } from '../utils/podman';
+import { getLocalModelFile, getRemoteModelFile } from '../utils/modelsUtils';
 import type { CancellationTokenRegistry } from '../registries/CancellationTokenRegistry';
 import { getHash, hasValidSha } from '../utils/sha';
 import type { GGUFParseOutput } from '@huggingface/gguf';
@@ -231,17 +230,32 @@ export class ModelsManager implements Disposable {
     for (const connection of connections) {
       // ignore non-wsl machines
       if (connection.vmType !== VMType.WSL) continue;
-      // Get the corresponding machine name
-      const machineName = getPodmanMachineName(connection);
 
-      // check if model already loaded on the podman machine
-      const existsRemote = await isModelUploaded(machineName, modelInfo);
-      if (!existsRemote) return;
+      // check if remote model is
+      try {
+        await this.podmanConnection.executeSSH(connection, ['stat', getRemoteModelFile(modelInfo)]);
+      } catch (err: unknown) {
+        console.warn(err);
+        continue;
+      }
 
-      await deleteRemoteModel(machineName, modelInfo);
+      await this.deleteRemoteModelByConnection(connection, modelInfo);
     }
   }
 
+  /**
+   * Delete a model given a {@link ContainerProviderConnection}
+   * @param connection
+   * @param modelInfo
+   * @protected
+   */
+  protected async deleteRemoteModelByConnection(
+    connection: ContainerProviderConnection,
+    modelInfo: ModelInfo,
+  ): Promise<void> {
+    await this.podmanConnection.executeSSH(connection, ['rm', '-f', getRemoteModelFile(modelInfo)]);
+  }
+
   /**
    * This method will resolve when the provided model will be downloaded.
    *
@@ -439,7 +453,7 @@ export class ModelsManager implements Disposable {
       connection: connection.name,
     });
 
-    const uploader = new Uploader(connection, model);
+    const uploader = new Uploader(this.podmanConnection, connection, model);
     uploader.onEvent(event => this.onDownloadUploadEvent(event, 'upload'), this);
 
     // perform download
diff --git a/packages/backend/src/utils/modelsUtils.spec.ts b/packages/backend/src/utils/modelsUtils.spec.ts
@@ -16,16 +16,8 @@
  * SPDX-License-Identifier: Apache-2.0
  ***********************************************************************/
 import { beforeEach, describe, expect, test, vi } from 'vitest';
-import { process as apiProcess } from '@podman-desktop/api';
-import {
-  deleteRemoteModel,
-  getLocalModelFile,
-  getRemoteModelFile,
-  isModelUploaded,
-  MACHINE_BASE_FOLDER,
-} from './modelsUtils';
+import { getLocalModelFile, getRemoteModelFile, MACHINE_BASE_FOLDER } from './modelsUtils';
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
-import { getPodmanCli } from './podman';
 
 vi.mock('@podman-desktop/api', () => {
   return {
@@ -35,14 +27,8 @@ vi.mock('@podman-desktop/api', () => {
   };
 });
 
-vi.mock('./podman', () => ({
-  getPodmanCli: vi.fn(),
-}));
-
 beforeEach(() => {
   vi.resetAllMocks();
-
-  vi.mocked(getPodmanCli).mockReturnValue('dummyPodmanCli');
 });
 
 describe('getLocalModelFile', () => {
@@ -94,48 +80,3 @@ describe('getRemoteModelFile', () => {
     expect(path).toBe(`${MACHINE_BASE_FOLDER}dummy.guff`);
   });
 });
-
-describe('isModelUploaded', () => {
-  test('execute stat on targeted machine', async () => {
-    expect(
-      await isModelUploaded('dummyMachine', {
-        id: 'dummyModelId',
-        file: {
-          path: 'dummyPath',
-          file: 'dummy.guff',
-        },
-      } as unknown as ModelInfo),
-    ).toBeTruthy();
-
-    expect(getPodmanCli).toHaveBeenCalled();
-    expect(apiProcess.exec).toHaveBeenCalledWith('dummyPodmanCli', [
-      'machine',
-      'ssh',
-      'dummyMachine',
-      'stat',
-      expect.anything(),
-    ]);
-  });
-});
-
-describe('deleteRemoteModel', () => {
-  test('execute stat on targeted machine', async () => {
-    await deleteRemoteModel('dummyMachine', {
-      id: 'dummyModelId',
-      file: {
-        path: 'dummyPath',
-        file: 'dummy.guff',
-      },
-    } as unknown as ModelInfo);
-
-    expect(getPodmanCli).toHaveBeenCalled();
-    expect(apiProcess.exec).toHaveBeenCalledWith('dummyPodmanCli', [
-      'machine',
-      'ssh',
-      'dummyMachine',
-      'rm',
-      '-f',
-      expect.anything(),
-    ]);
-  });
-});
diff --git a/packages/backend/src/utils/modelsUtils.ts b/packages/backend/src/utils/modelsUtils.ts
@@ -17,8 +17,6 @@
  ***********************************************************************/
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
 import { join, posix } from 'node:path';
-import { getPodmanCli } from './podman';
-import { process } from '@podman-desktop/api';
 
 export const MACHINE_BASE_FOLDER = '/home/user/ai-lab/models/';
 
@@ -42,36 +40,6 @@ export function getRemoteModelFile(modelInfo: ModelInfo): string {
   return posix.join(MACHINE_BASE_FOLDER, modelInfo.file.file);
 }
 
-/**
- * utility method to determine if a model is already uploaded to the podman machine
- * @param machine
- * @param modelInfo
- */
-export async function isModelUploaded(machine: string, modelInfo: ModelInfo): Promise<boolean> {
-  try {
-    const remotePath = getRemoteModelFile(modelInfo);
-    await process.exec(getPodmanCli(), ['machine', 'ssh', machine, 'stat', remotePath]);
-    return true;
-  } catch (err: unknown) {
-    console.error('Something went wrong while trying to stat remote model path', err);
-    return false;
-  }
-}
-
-/**
- * Given a machine and a modelInfo, delete the corresponding file on the podman machine
- * @param machine the machine to target
- * @param modelInfo the model info
- */
-export async function deleteRemoteModel(machine: string, modelInfo: ModelInfo): Promise<void> {
-  try {
-    const remotePath = getRemoteModelFile(modelInfo);
-    await process.exec(getPodmanCli(), ['machine', 'ssh', machine, 'rm', '-f', remotePath]);
-  } catch (err: unknown) {
-    console.error('Something went wrong while trying to stat remote model path', err);
-  }
-}
-
 export function getModelPropertiesForEnvironment(modelInfo: ModelInfo): string[] {
   const envs: string[] = [];
   if (modelInfo.properties) {
diff --git a/packages/backend/src/utils/podman.ts b/packages/backend/src/utils/podman.ts
@@ -32,6 +32,11 @@ export type MachineJSON = {
   VMType?: string;
 };
 
+/**
+ * We should be using the {@link @podman-desktop/api#extensions.getExtensions} function to get podman
+ * exec method
+ * @deprecated
+ */
 export function getPodmanCli(): string {
   // If we have a custom binary path regardless if we are running Windows or not
   const customBinaryPath = getCustomBinaryPath();
@@ -52,8 +57,9 @@ export function getCustomBinaryPath(): string | undefined {
 }
 
 /**
- * In the ${link ContainerProviderConnection.name} property the name is not usage, and we need to transform it
+ * In the {@link ContainerProviderConnection.name} property the name is not usage, and we need to transform it
  * @param connection
+ * @deprecated
  */
 export function getPodmanMachineName(connection: ContainerProviderConnection): string {
   const runningConnectionName = connection.name;
diff --git a/packages/backend/src/utils/uploader.spec.ts b/packages/backend/src/utils/uploader.spec.ts
@@ -24,15 +24,13 @@ import { Uploader } from './uploader';
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
 import type { ContainerProviderConnection } from '@podman-desktop/api';
 import { VMType } from '@shared/src/models/IPodman';
+import type { PodmanConnection } from '../managers/podmanConnection';
 
 vi.mock('@podman-desktop/api', async () => {
   return {
     env: {
       isWindows: false,
     },
-    process: {
-      exec: vi.fn(),
-    },
     EventEmitter: vi.fn().mockImplementation(() => {
       return {
         fire: vi.fn(),
@@ -41,6 +39,10 @@ vi.mock('@podman-desktop/api', async () => {
   };
 });
 
+const podmanConnectionMock: PodmanConnection = {
+  executeSSH: vi.fn(),
+} as unknown as PodmanConnection;
+
 const connectionMock: ContainerProviderConnection = {
   name: 'machine2',
   type: 'podman',
@@ -51,7 +53,7 @@ const connectionMock: ContainerProviderConnection = {
   },
 };
 
-const uploader = new Uploader(connectionMock, {
+const uploader = new Uploader(podmanConnectionMock, connectionMock, {
   id: 'dummyModelId',
   file: {
     file: 'dummyFile.guff',
diff --git a/packages/backend/src/utils/uploader.ts b/packages/backend/src/utils/uploader.ts
@@ -24,18 +24,20 @@ import type { ModelInfo } from '@shared/src/models/IModelInfo';
 import { getLocalModelFile } from './modelsUtils';
 import type { IWorker } from '../workers/IWorker';
 import type { UploaderOptions } from '../workers/uploader/UploaderOptions';
+import type { PodmanConnection } from '../managers/podmanConnection';
 
 export class Uploader {
   readonly #_onEvent = new EventEmitter<BaseEvent>();
   readonly onEvent: Event<BaseEvent> = this.#_onEvent.event;
   readonly #workers: IWorker<UploaderOptions, string>[] = [];
 
   constructor(
+    podman: PodmanConnection,
     private connection: ContainerProviderConnection,
     private modelInfo: ModelInfo,
     private abortSignal?: AbortSignal,
   ) {
-    this.#workers = [new WSLUploader()];
+    this.#workers = [new WSLUploader(podman)];
   }
 
   /**
diff --git a/packages/backend/src/workers/uploader/WSLUploader.spec.ts b/packages/backend/src/workers/uploader/WSLUploader.spec.ts
diff --git a/packages/backend/src/workers/uploader/WSLUploader.ts b/packages/backend/src/workers/uploader/WSLUploader.ts
