update

janhq · Jun 13, 2024 · 5b44c0a · 5b44c0a
1 parent d93378e
commit 5b44c0a
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 8 deletions.
diff --git a/cortex-js/src/download-manager/download-manager.service.ts b/cortex-js/src/download-manager/download-manager.service.ts
@@ -13,6 +13,8 @@ import { firstValueFrom } from 'rxjs';
 @Injectable()
 export class DownloadManagerService {
   private allDownloadStates: DownloadState[] = [];
+  private abortControllers: Record<string, Record<string, AbortController>> =
+    {};
 
   constructor(
     private readonly httpService: HttpService,
@@ -24,6 +26,19 @@ export class DownloadManagerService {
     }, 500);
   }
 
+  async abortDownload(downloadId: string) {
+    if (!this.abortControllers[downloadId]) {
+      return;
+    }
+    Object.keys(this.abortControllers[downloadId]).forEach((destination) => {
+      this.abortControllers[downloadId][destination].abort();
+    });
+    delete this.abortControllers[downloadId];
+    this.allDownloadStates = this.allDownloadStates.filter(
+      (downloadState) => downloadState.id !== downloadId,
+    );
+  }
+
   async submitDownloadRequest(
     downloadId: string,
     title: string,
@@ -67,6 +82,7 @@ export class DownloadManagerService {
     };
 
     this.allDownloadStates.push(downloadState);
+    this.abortControllers[downloadId] = {};
 
     Object.keys(urlToDestination).forEach((url) => {
       const destination = urlToDestination[url];
@@ -79,11 +95,17 @@ export class DownloadManagerService {
     url: string,
     destination: string,
   ) {
+    const controller = new AbortController();
+    // adding to abort controllers
+    this.abortControllers[downloadId][destination] = controller;
+
     const response = await firstValueFrom(
       this.httpService.get(url, {
         responseType: 'stream',
+        signal: controller.signal,
       }),
     );
+
     // check if response is success
     if (!response) {
       throw new Error('Failed to download model');
@@ -93,7 +115,6 @@ export class DownloadManagerService {
     const totalBytes = response.headers['content-length'];
 
     // update download state
-    // TODO: separte this duplicate code to a function
     const currentDownloadState = this.allDownloadStates.find(
       (downloadState) => downloadState.id === downloadId,
     );
@@ -110,6 +131,8 @@ export class DownloadManagerService {
     let transferredBytes = 0;
 
     writer.on('finish', () => {
+      // delete the abort controller
+      delete this.abortControllers[downloadId][destination];
       const currentDownloadState = this.allDownloadStates.find(
         (downloadState) => downloadState.id === downloadId,
       );
@@ -130,8 +153,8 @@ export class DownloadManagerService {
       );
 
       if (allChildrenDownloaded) {
+        delete this.abortControllers[downloadId];
         currentDownloadState.status = DownloadStatus.Downloaded;
-        // TODO: notify if download success so that client can auto refresh
         // remove download state if all children is downloaded
         this.allDownloadStates = this.allDownloadStates.filter(
           (downloadState) => downloadState.id !== downloadId,
@@ -140,6 +163,7 @@ export class DownloadManagerService {
     });
 
     writer.on('error', (error) => {
+      delete this.abortControllers[downloadId][destination];
       const currentDownloadState = this.allDownloadStates.find(
         (downloadState) => downloadState.id === downloadId,
       );
@@ -158,7 +182,6 @@ export class DownloadManagerService {
       currentDownloadState.status = DownloadStatus.Error;
       currentDownloadState.error = error.message;
 
-      // TODO: notify if download error
       // remove download state if all children is downloaded
       this.allDownloadStates = this.allDownloadStates.filter(
         (downloadState) => downloadState.id !== downloadId,

diff --git a/cortex-js/src/infrastructure/controllers/models.controller.ts b/cortex-js/src/infrastructure/controllers/models.controller.ts
@@ -100,17 +100,35 @@ export class ModelsController {
   @ApiOperation({
     summary: 'Download model',
     description: 'Downloads a specific model instance.',
-  })
-  @ApiParam({
-    name: 'modelId',
-    required: true,
-    description: 'The unique identifier of the model.',
+    parameters: [
+      {
+        in: 'path',
+        name: 'modelId',
+        required: true,
+        description: 'The unique identifier of the model.',
+      },
+    ],
   })
   @Get('download/:modelId(*)')
   downloadModel(@Param('modelId') modelId: string) {
     return this.modelsUsecases.downloadModel(modelId);
   }
 
+  @ApiOperation({
+    parameters: [
+      {
+        in: 'path',
+        name: 'download_id',
+        required: true,
+        description: 'The unique identifier of the download.',
+      },
+    ],
+  })
+  @Get('abort-download/:download_id(*)')
+  abortDownloadModel(@Param('download_id') downloadId: string) {
+    return this.modelsUsecases.abortDownloadModel(downloadId);
+  }
+
   @HttpCode(200)
   @ApiResponse({
     status: 200,

diff --git a/cortex-js/src/usecases/models/models.usecases.ts b/cortex-js/src/usecases/models/models.usecases.ts
@@ -292,6 +292,10 @@ export class ModelsUsecases {
     }
   }
 
+  async abortDownloadModel(downloadId: string) {
+    this.downloadManagerService.abortDownload(downloadId);
+  }
+
   /**
    * Populate model metadata from a Model repository (HF, Jan...) and download it
    * @param modelId