fix: prevent double upload of model into Podman machine (#2067)

Fixes #2065 Signed-off-by: Jeff MAURY <[email protected]>
containers · Nov 19, 2024 · 8d43a1b · 8d43a1b
1 parent 7e63d90
commit 8d43a1b
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 8 deletions.
diff --git a/packages/backend/src/managers/application/applicationManager.spec.ts b/packages/backend/src/managers/application/applicationManager.spec.ts
@@ -29,6 +29,7 @@ import type { Recipe, RecipeImage } from '@shared/src/models/IRecipe';
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
 import { VMType } from '@shared/src/models/IPodman';
 import { POD_LABEL_MODEL_ID, POD_LABEL_RECIPE_ID } from '../../utils/RecipeConstants';
+import type { InferenceServer } from '@shared/src/models/IInference';
 
 const taskRegistryMock = {
   createTask: vi.fn(),
@@ -350,6 +351,71 @@ describe('pullApplication', () => {
     expect(podManager.startPod).toHaveBeenCalledWith('test-engine-id', 'test-pod-id');
   });
 
+  test('requestDownloadModel skipped with inference server', async () => {
+    vi.mocked(recipeManager.buildRecipe).mockResolvedValue({
+      images: [recipeImageInfoMock],
+      inferenceServer: {
+        connection: {
+          port: 56001,
+        },
+      } as InferenceServer,
+    });
+    await getInitializedApplicationManager().pullApplication(connectionMock, recipeMock, remoteModelMock, {
+      'test-label': 'test-value',
+    });
+
+    // clone the recipe
+    expect(recipeManager.cloneRecipe).toHaveBeenCalledWith(recipeMock, {
+      'test-label': 'test-value',
+      'model-id': remoteModelMock.id,
+    });
+    // download model
+    expect(modelsManagerMock.requestDownloadModel).toHaveBeenCalledWith(remoteModelMock, {
+      'test-label': 'test-value',
+      'recipe-id': recipeMock.id,
+      'model-id': remoteModelMock.id,
+    });
+    // upload model to podman machine
+    expect(modelsManagerMock.uploadModelToPodmanMachine).not.toHaveBeenCalled();
+    // build the recipe
+    expect(recipeManager.buildRecipe).toHaveBeenCalledWith(connectionMock, recipeMock, remoteModelMock, {
+      'test-label': 'test-value',
+      'recipe-id': recipeMock.id,
+      'model-id': remoteModelMock.id,
+    });
+    // create AI App task must be created
+    expect(taskRegistryMock.createTask).toHaveBeenCalledWith('Creating AI App', 'loading', {
+      'test-label': 'test-value',
+      'recipe-id': recipeMock.id,
+      'model-id': remoteModelMock.id,
+    });
+
+    // a pod must have been created
+    expect(podManager.createPod).toHaveBeenCalledWith({
+      provider: connectionMock,
+      name: expect.any(String),
+      portmappings: [],
+      labels: {
+        [POD_LABEL_MODEL_ID]: remoteModelMock.id,
+        [POD_LABEL_RECIPE_ID]: recipeMock.id,
+      },
+    });
+
+    expect(containerEngine.createContainer).toHaveBeenCalledWith('test-engine-id', {
+      Image: recipeImageInfoMock.id,
+      name: expect.any(String),
+      Env: ['MODEL_ENDPOINT=http://host.containers.internal:56001'],
+      HealthCheck: undefined,
+      HostConfig: undefined,
+      Detach: true,
+      pod: 'test-pod-id',
+      start: false,
+    });
+
+    // finally the pod must be started
+    expect(podManager.startPod).toHaveBeenCalledWith('test-engine-id', 'test-pod-id');
+  });
+
   test('existing application should be removed', async () => {
     vi.mocked(podManager.findPodByLabelsValues).mockResolvedValue({
       engineId: 'test-engine-id',

diff --git a/packages/backend/src/managers/application/applicationManager.ts b/packages/backend/src/managers/application/applicationManager.ts
@@ -182,14 +182,7 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
     await this.recipeManager.cloneRecipe(recipe, { ...labels, 'model-id': model.id });
 
     // get model by downloading it or retrieving locally
-    await this.modelsManager.requestDownloadModel(model, {
-      ...labels,
-      'recipe-id': recipe.id,
-      'model-id': model.id,
-    });
-
-    // upload model to podman machine if user system is supported
-    const modelPath = await this.modelsManager.uploadModelToPodmanMachine(connection, model, {
+    let modelPath = await this.modelsManager.requestDownloadModel(model, {
       ...labels,
       'recipe-id': recipe.id,
       'model-id': model.id,
@@ -202,6 +195,15 @@ export class ApplicationManager extends Publisher<ApplicationState[]> implements
       'model-id': model.id,
     });
 
+    // upload model to podman machine if user system is supported
+    if (!recipeComponents.inferenceServer) {
+      modelPath = await this.modelsManager.uploadModelToPodmanMachine(connection, model, {
+        ...labels,
+        'recipe-id': recipe.id,
+        'model-id': model.id,
+      });
+    }
+
     // first delete any existing pod with matching labels
     if (await this.hasApplicationPod(recipe.id, model.id)) {
       await this.removeApplication(recipe.id, model.id);