optimizing finetuning downloads

wanhakim · wanhakim · commit 44ceaff87c0c · 2025-12-09T13:34:55.000Z
Signed-off-by: wwanarif &lt;wan.abdul.hakim.b.wan.arif@intel.com&gt;
diff --git a/setup-scripts/setup-genai-studio/manifests/studio-manifest.yaml b/setup-scripts/setup-genai-studio/manifests/studio-manifest.yaml
@@ -426,7 +426,7 @@ spec:
   storageClassName: local-path
   resources:
     requests:
-      storage: 1Gi
+      storage: 30Gi
 ---
 apiVersion: v1
 kind: Service
diff --git a/studio-frontend/packages/server/src/controllers/finetuning/index.ts b/studio-frontend/packages/server/src/controllers/finetuning/index.ts
@@ -3,6 +3,9 @@ import { StatusCodes } from 'http-status-codes'
 import { InternalFlowiseError } from '../../errors/internalFlowiseError'
 import finetuningService from '../../services/finetuning'
 
+// Declare timer globals for Node.js
+declare function setTimeout(cb: (...args: any[]) => void, ms?: number): any
+
 /**
  * Upload a training file
  * POST /api/v1/finetuning/files
@@ -154,6 +157,7 @@ const getFineTuningJobLogs = async (req: Request, res: Response, next: NextFunct
 /**
  * Download fine-tuning job output as a zip file
  * GET /api/v1/finetuning/download-ft/:jobId
+ * Creates zip, streams it to client, then deletes the zip file after download completes
  */
 const downloadFineTuningOutput = async (req: Request, res: Response, next: NextFunction) => {
     try {
@@ -166,31 +170,63 @@ const downloadFineTuningOutput = async (req: Request, res: Response, next: NextF
             )
         }
 
-        // Get the zip file path (creates if needed, but returns immediately if already exists)
+        // Get the zip file path from service
         const filePath = await finetuningService.downloadFineTuningOutput(jobId)
+        
         if (!filePath) {
             throw new InternalFlowiseError(
                 StatusCodes.NOT_FOUND,
-                `Error: finetuningController.downloadFineTuningOutput - output not found for job: ${jobId}`
+                `Error: finetuningController.downloadFineTuningOutput - zip file not found for job: ${jobId}. Please request download via WebSocket first.`
             )
         }
 
+        const fs = require('fs')
+
         // Set response headers for file download
         const fileName = `${jobId}-output.zip`
         res.setHeader('Content-Type', 'application/zip')
         res.setHeader('Content-Disposition', `attachment; filename="${fileName}"`)
         
         // Stream the file
-        const fs = require('fs')
         const fileStream = fs.createReadStream(filePath)
+        
+        // Log when stream opens
+        fileStream.on('open', () => {
+            console.debug(`finetuningController.downloadFineTuningOutput - starting to stream: ${filePath}`)
+        })
+        
+        // Delete zip file after response fully finishes (browser completes download)
+        res.on('finish', () => {
+            setTimeout(() => {
+                try {
+                    if (fs.existsSync(filePath)) {
+                        fs.unlinkSync(filePath)
+                        console.debug(`finetuningController.downloadFineTuningOutput - deleted zip after download complete: ${filePath}`)
+                    }
+                } catch (deleteErr: any) {
+                    console.warn(`finetuningController.downloadFineTuningOutput - failed to delete zip: ${deleteErr?.message || deleteErr}`)
+                }
+            }, 100)
+        })
+        
         fileStream.on('error', (err: any) => {
-            console.error('Error streaming fine-tuning output file:', err)
+            console.error('finetuningController.downloadFineTuningOutput - error streaming file:', err)
+            // Try to delete zip on error too
+            try {
+                if (fs.existsSync(filePath)) {
+                    fs.unlinkSync(filePath)
+                    console.debug(`finetuningController.downloadFineTuningOutput - deleted zip after error: ${filePath}`)
+                }
+            } catch (deleteErr: any) {
+                console.warn(`finetuningController.downloadFineTuningOutput - failed to delete zip on error: ${deleteErr?.message || deleteErr}`)
+            }
             if (!res.headersSent) {
                 res.status(StatusCodes.INTERNAL_SERVER_ERROR).json({
                     error: 'Error streaming fine-tuning output file'
                 })
             }
         })
+        
         fileStream.pipe(res)
     } catch (error) {
         next(error)
diff --git a/studio-frontend/packages/server/src/services/finetuning/index.ts b/studio-frontend/packages/server/src/services/finetuning/index.ts
@@ -13,6 +13,10 @@ import { getRunningExpressApp } from '../../utils/getRunningExpressApp'
 import { FineTuningJob } from '../../database/entities/FineTuningJob'
 import logger from '../../utils/logger'
 
+// Declare timer globals for Node.js
+declare function setTimeout(cb: (...args: any[]) => void, ms?: number): any
+declare function clearTimeout(id: any): void
+
 const execAsync = promisify(exec)
 
 const FINETUNING_SERVICE_URL = process.env.FINETUNING_HOST ? `http://${process.env.FINETUNING_HOST}:8015` : 'undefined'
@@ -36,93 +40,6 @@ const axiosClient: AxiosInstance = axios.create({
 // In-memory mapping: filename (raw and decoded) -> { id, rawFilename }
 const uploadedFileIdMap: Map<string, { id: string; rawFilename: string }> = new Map()
 
-/**
- * Helper function to zip a fine-tuning job output directory
- * Checks if zip already exists and is up-to-date before creating a new one
- * @param outputDir - Full path to the output directory for the job
- * @param jobId - ID of the fine-tuning job
- * @returns Path to the zipped file or null if failed
- */
-const ensureFineTuningOutputZip = async (outputDir: string, jobId: string): Promise<string | null> => {
-    try {
-        // eslint-disable-next-line no-console
-        console.debug(`finetuningService.ensureFineTuningOutputZip - processing output for job: ${jobId}`)
-        
-        // Validate output directory exists
-        if (!fs.existsSync(outputDir)) {
-            // eslint-disable-next-line no-console
-            console.warn(`finetuningService.ensureFineTuningOutputZip - output directory not found: ${outputDir}`)
-            return null
-        }
-
-        const zipFilePath = `${outputDir}.zip`
-        const outputStats = fs.statSync(outputDir)
-        
-        // Check if zip exists and is up-to-date
-        if (fs.existsSync(zipFilePath)) {
-            const zipStats = fs.statSync(zipFilePath)
-            // If zip is newer than the output directory, skip re-zipping
-            if (zipStats.mtimeMs > outputStats.mtimeMs) {
-                // eslint-disable-next-line no-console
-                console.debug(`finetuningService.ensureFineTuningOutputZip - zip already up-to-date: ${zipFilePath}`)
-                return zipFilePath
-            }
-            // Remove outdated zip
-            try {
-                fs.unlinkSync(zipFilePath)
-                // eslint-disable-next-line no-console
-                console.debug(`finetuningService.ensureFineTuningOutputZip - removed outdated zip: ${zipFilePath}`)
-            } catch (e) {
-                // eslint-disable-next-line no-console
-                console.warn(`finetuningService.ensureFineTuningOutputZip - failed to remove outdated zip: ${e}`)
-            }
-        }
-
-        // Create zip file using archiver (standard ZIP format compatible with Windows)
-        // eslint-disable-next-line no-console
-        console.debug(`finetuningService.ensureFineTuningOutputZip - starting to zip output for job ${jobId}`)
-        try {
-            return await new Promise((resolve, reject) => {
-                const output = fs.createWriteStream(zipFilePath)
-                const archive = archiver('zip', {
-                    zlib: { level: 6 } // compression level
-                })
-
-                output.on('close', () => {
-                    // eslint-disable-next-line no-console
-                    console.debug(`finetuningService.ensureFineTuningOutputZip - zip created successfully for job ${jobId}: ${zipFilePath} (${archive.pointer()} bytes)`)
-                    resolve(zipFilePath)
-                })
-
-                output.on('error', (err: any) => {
-                    // eslint-disable-next-line no-console
-                    console.error(`finetuningService.ensureFineTuningOutputZip - write stream error: ${err?.message || err}`)
-                    reject(err)
-                })
-
-                archive.on('error', (err: any) => {
-                    // eslint-disable-next-line no-console
-                    console.error(`finetuningService.ensureFineTuningOutputZip - archiver error: ${err?.message || err}`)
-                    reject(err)
-                })
-
-                archive.pipe(output)
-                // Add the entire directory to the archive with basename as root
-                archive.directory(outputDir, path.basename(outputDir))
-                archive.finalize()
-            })
-        } catch (execErr: any) {
-            // eslint-disable-next-line no-console
-            console.error(`finetuningService.ensureFineTuningOutputZip - archiver failed for job ${jobId}: ${execErr?.message || execErr}`)
-            return null
-        }
-    } catch (error: any) {
-        // eslint-disable-next-line no-console
-        console.error(`finetuningService.ensureFineTuningOutputZip - error: ${error?.message || error}`)
-        return null
-    }
-}
-
 /**
  * Upload a training file to the finetuning service
  */
@@ -783,12 +700,12 @@ const deleteFineTuningJob = async (fineTuningJobId: string) => {
 }
 
 /**
- * Download fine-tuning job output as a zip file
- * Creates zip if needed, or returns existing zip immediately
+ * Prepare fine-tuning job output as a zip file for download
+ * Called by WebSocket to create and cache the zip
  * @param jobId - ID of the fine-tuning job
  * @returns Path to the zipped file or null if not found
  */
-const downloadFineTuningOutput = async (jobId: string): Promise<string | null> => {
+const prepareFineTuningOutputZip = async (jobId: string): Promise<string | null> => {
     try {
         if (!jobId) {
             throw new InternalFlowiseError(StatusCodes.BAD_REQUEST, 'Job ID is required')
@@ -816,18 +733,82 @@ const downloadFineTuningOutput = async (jobId: string): Promise<string | null> =
             throw new InternalFlowiseError(StatusCodes.FORBIDDEN, 'Invalid job output path')
         }
 
-        // Ensure the output is zipped (returns immediately if zip is up-to-date)
-        const finalZipPath = await ensureFineTuningOutputZip(jobOutputDir, jobId)
-        if (!finalZipPath) {
-            throw new InternalFlowiseError(
-                StatusCodes.INTERNAL_SERVER_ERROR,
-                `Failed to create zip for job ${jobId}`
-            )
-        }
+        const zipFilePath = `${jobOutputDir}.zip`
 
+        // Create zip file using archiver
         // eslint-disable-next-line no-console
-        console.debug(`finetuningService.downloadFineTuningOutput - file ready for download: ${finalZipPath}`)
-        return finalZipPath
+        console.debug(`finetuningService.downloadFineTuningOutput - creating zip for job ${jobId}`)
+        
+        // Log directory contents for diagnostics
+        try {
+            const dirContents = fs.readdirSync(jobOutputDir)
+            // eslint-disable-next-line no-console
+            console.debug(`finetuningService.downloadFineTuningOutput - output directory contains ${dirContents.length} items: ${dirContents.slice(0, 10).join(', ')}${dirContents.length > 10 ? '...' : ''}`)
+        } catch (e) {
+            // eslint-disable-next-line no-console
+            console.warn(`finetuningService.downloadFineTuningOutput - could not list directory: ${e}`)
+        }
+        
+        try {
+            return await new Promise((resolve, reject) => {
+                const output = fs.createWriteStream(zipFilePath)
+                const archive = archiver('zip', {
+                    zlib: { level: 0 } // no compression for speed
+                })
+                
+                const zipTimeoutMs = 30 * 60 * 1000 // 30 minutes
+                let resolved = false
+                
+                const timeoutHandle = setTimeout(() => {
+                    if (!resolved) {
+                        resolved = true
+                        // eslint-disable-next-line no-console
+                        console.error(`finetuningService.downloadFineTuningOutput - archiver timeout for job ${jobId}`)
+                        try { output.destroy() } catch (e) {}
+                        try { archive.destroy() } catch (e) {}
+                        reject(new Error('Archiver timeout'))
+                    }
+                }, zipTimeoutMs)
+
+                output.on('close', () => {
+                    if (!resolved) {
+                        resolved = true
+                        clearTimeout(timeoutHandle)
+                        // eslint-disable-next-line no-console
+                        console.debug(`finetuningService.downloadFineTuningOutput - zip created: ${zipFilePath}`)
+                        resolve(zipFilePath)
+                    }
+                })
+
+                output.on('error', (err: any) => {
+                    if (!resolved) {
+                        resolved = true
+                        clearTimeout(timeoutHandle)
+                        // eslint-disable-next-line no-console
+                        console.error(`finetuningService.downloadFineTuningOutput - write stream error: ${err?.message || err}`)
+                        reject(err)
+                    }
+                })
+
+                archive.on('error', (err: any) => {
+                    if (!resolved) {
+                        resolved = true
+                        clearTimeout(timeoutHandle)
+                        // eslint-disable-next-line no-console
+                        console.error(`finetuningService.downloadFineTuningOutput - archiver error: ${err?.message || err}`)
+                        reject(err)
+                    }
+                })
+
+                archive.pipe(output)
+                archive.directory(jobOutputDir, path.basename(jobOutputDir))
+                archive.finalize()
+            })
+        } catch (archiverErr: any) {
+            // eslint-disable-next-line no-console
+            console.error(`finetuningService.downloadFineTuningOutput - archiver failed for job ${jobId}: ${archiverErr?.message || archiverErr}`)
+            return null
+        }
     } catch (error: any) {
         if (error instanceof InternalFlowiseError) {
             throw error
@@ -841,6 +822,46 @@ const downloadFineTuningOutput = async (jobId: string): Promise<string | null> =
     }
 }
 
+/**
+ * Download fine-tuning job output - HTTP endpoint
+ * Returns path to cached ZIP file
+ * @param jobId - ID of the fine-tuning job
+ * @returns Path to the zipped file or null if not found
+ */
+const downloadFineTuningOutput = async (jobId: string): Promise<string | null> => {
+    try {
+        if (!jobId) {
+            return null
+        }
+
+        const OUTPUT_BASE_DIR = '/tmp/finetuning/output'
+        const zipFilePath = `${OUTPUT_BASE_DIR}/${jobId}.zip`
+
+        // Check if zip file exists
+        if (fs.existsSync(zipFilePath)) {
+            try {
+                const stat = fs.statSync(zipFilePath)
+                if (stat.size > 0) {
+                    // eslint-disable-next-line no-console
+                    console.debug(`finetuningService.downloadFineTuningOutput - returning cached zip: ${zipFilePath}`)
+                    return zipFilePath
+                }
+            } catch (e) {
+                // eslint-disable-next-line no-console
+                console.warn(`finetuningService.downloadFineTuningOutput - could not stat zip file: ${e}`)
+            }
+        }
+
+        // eslint-disable-next-line no-console
+        console.warn(`finetuningService.downloadFineTuningOutput - zip file not found: ${zipFilePath}`)
+        return null
+    } catch (error: any) {
+        // eslint-disable-next-line no-console
+        console.error(`finetuningService.downloadFineTuningOutput - error: ${error?.message || error}`)
+        return null
+    }
+}
+
 /**
  * Get logs for a fine-tuning job by querying the Ray head node HTTP API.
  * It will call: http://<RAY_HEAD_NODE>/api/jobs/<job_id>/logs
@@ -929,5 +950,6 @@ export default {
     cancelFineTuningJob,
     deleteFineTuningJob,
     getFineTuningJobLogs,
+    prepareFineTuningOutputZip,
     downloadFineTuningOutput
 }
diff --git a/studio-frontend/packages/server/src/ws/finetuningDownload.ts b/studio-frontend/packages/server/src/ws/finetuningDownload.ts
@@ -99,10 +99,18 @@ export const setupFineTuningDownloadHandlers = (io: Server) => {
 
                 // Kick off the async preparation and store the promise so others can join
                 task.status = 'zipping'
+                
+                // Emit progress update to socket immediately
+                socket.emit('download-finetuning-progress', {
+                    jobId,
+                    status: 'zipping',
+                    message: 'Creating zip archive (this may take a few minutes)...'
+                })
+                
                 task.promise = (async () => {
                     try {
-                        // Call the service to prepare the zip file (returns path)
-                        const zipFilePath = await finetuningService.downloadFineTuningOutput(jobId)
+                        // Call the service to prepare the zip file
+                        const zipFilePath = await finetuningService.prepareFineTuningOutputZip(jobId)
 
                         if (!zipFilePath) {
                             task.status = 'error'
