update worker shutdown gracefully

FlowiseAI · Jan 22, 2025 · 2262965 · 2262965
1 parent 67aa942
commit 2262965
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 38 deletions.
diff --git a/packages/server/src/commands/worker.ts b/packages/server/src/commands/worker.ts
@@ -59,7 +59,7 @@ export default class Worker extends BaseCommand {
         await appDataSource.initialize()
         await appDataSource.runMigrations({ transaction: 'each' })
 
-        // Initialize chatflow pool
+        // Initialize abortcontroller pool
         const abortControllerPool = new AbortControllerPool()
 
         // Init telemetry
@@ -85,10 +85,14 @@ export default class Worker extends BaseCommand {
 
     async stopProcess() {
         try {
+            const queueManager = QueueManager.getInstance()
+            const predictionWorker = queueManager.getQueue('prediction').getWorker()
             logger.info(`Shutting down Flowise Prediction Worker ${this.predictionWorkerId}...`)
+            await predictionWorker.close()
+
+            const upsertWorker = queueManager.getQueue('upsert').getWorker()
             logger.info(`Shutting down Flowise Upsertion Worker ${this.upsertionWorkerId}...`)
-            //const serverApp = Server.getInstance()
-            //if (serverApp) await serverApp.stopApp()
+            await upsertWorker.close()
         } catch (error) {
             logger.error('There was an error shutting down Flowise Worker...', error)
             await this.failExit()

diff --git a/packages/server/src/controllers/internal-predictions/index.ts b/packages/server/src/controllers/internal-predictions/index.ts
@@ -23,7 +23,7 @@ const createInternalPrediction = async (req: Request, res: Response, next: NextF
 const createAndStreamInternalPrediction = async (req: Request, res: Response, next: NextFunction) => {
     const chatId = req.body.chatId
     const sseStreamer = getRunningExpressApp().sseStreamer
-    const redisSubscriber = getRunningExpressApp().redisSubscriber
+
     try {
         sseStreamer.addClient(chatId, res)
         res.setHeader('Content-Type', 'text/event-stream')
@@ -33,7 +33,7 @@ const createAndStreamInternalPrediction = async (req: Request, res: Response, ne
         res.flushHeaders()
 
         if (process.env.MODE === MODE.QUEUE) {
-            redisSubscriber.subscribe(chatId)
+            getRunningExpressApp().redisSubscriber.subscribe(chatId)
         }
 
         const apiResponse = await utilBuildChatflow(req, true)

diff --git a/packages/server/src/controllers/predictions/index.ts b/packages/server/src/controllers/predictions/index.ts
@@ -56,7 +56,6 @@ const createPrediction = async (req: Request, res: Response, next: NextFunction)
             const isStreamingRequested = req.body.streaming === 'true' || req.body.streaming === true
             if (streamable?.isStreaming && isStreamingRequested) {
                 const sseStreamer = getRunningExpressApp().sseStreamer
-                const redisSubscriber = getRunningExpressApp().redisSubscriber
 
                 let chatId = req.body.chatId
                 if (!req.body.chatId) {
@@ -72,11 +71,11 @@ const createPrediction = async (req: Request, res: Response, next: NextFunction)
                     res.flushHeaders()
 
                     if (process.env.MODE === MODE.QUEUE) {
-                        redisSubscriber.subscribe(chatId)
+                        getRunningExpressApp().redisSubscriber.subscribe(chatId)
                     }
 
                     const apiResponse = await predictionsServices.buildChatflow(req)
-                    if (apiResponse) sseStreamer.streamMetadataEvent(apiResponse.chatId, apiResponse)
+                    sseStreamer.streamMetadataEvent(apiResponse.chatId, apiResponse)
                 } catch (error) {
                     if (chatId) {
                         sseStreamer.streamErrorEvent(chatId, getErrorMessage(error))
@@ -87,7 +86,7 @@ const createPrediction = async (req: Request, res: Response, next: NextFunction)
                 }
             } else {
                 const apiResponse = await predictionsServices.buildChatflow(req)
-                if (apiResponse) return res.json(apiResponse)
+                return res.json(apiResponse)
             }
         } else {
             const isStreamingRequested = req.body.streaming === 'true' || req.body.streaming === true

diff --git a/packages/server/src/queue/BaseQueue.ts b/packages/server/src/queue/BaseQueue.ts
@@ -11,6 +11,7 @@ export abstract class BaseQueue {
     protected queue: Queue
     protected queueEvents: QueueEvents
     protected connection: RedisOptions
+    private worker: Worker
 
     constructor(queueName: string, connection: RedisOptions) {
         this.connection = connection
@@ -27,13 +28,17 @@ export abstract class BaseQueue {
 
     abstract getQueue(): Queue
 
+    public getWorker(): Worker {
+        return this.worker
+    }
+
     public async addJob(jobData: any): Promise<Job> {
         const jobId = jobData.id || uuidv4()
         return await this.queue.add(jobId, jobData, { removeOnFail: true })
     }
 
     public createWorker(concurrency: number = WORKER_CONCURRENCY): Worker {
-        return new Worker(
+        this.worker = new Worker(
             this.queue.name,
             async (job: Job) => {
                 const start = new Date().getTime()
@@ -48,6 +53,7 @@ export abstract class BaseQueue {
                 concurrency
             }
         )
+        return this.worker
     }
 
     public async getJobs(): Promise<Job[]> {

diff --git a/packages/server/src/utils/index.ts b/packages/server/src/utils/index.ts
@@ -37,7 +37,6 @@ import {
     IMessage,
     FlowiseMemory,
     IFileUpload,
-    mapExtToInputField,
     getS3Config
 } from 'flowise-components'
 import { randomBytes } from 'crypto'
@@ -1759,33 +1758,6 @@ export const getAPIOverrideConfig = (chatflow: IChatFlow) => {
         const apiOverrideStatus: boolean =
             apiConfig.overrideConfig && apiConfig.overrideConfig.status ? apiConfig.overrideConfig.status : false
 
-        /* For "files" input, add a new node override with the actual input name such as pdfFile, txtFile, etc, to allow overriding the input
-         * https://github.com/FlowiseAI/Flowise/pull/3569
-         */
-        for (const nodeLabel in nodeOverrides) {
-            const params = nodeOverrides[nodeLabel]
-            const enabledFileParam = params.find((param) => param.enabled && param.name === 'files')
-            if (enabledFileParam) {
-                if (enabledFileParam.type.includes(',')) {
-                    const fileInputFieldsFromExt = enabledFileParam.type.split(',').map((fileType) => mapExtToInputField(fileType.trim()))
-                    for (const fileInputFieldFromExt of fileInputFieldsFromExt) {
-                        if (nodeOverrides[nodeLabel].some((param) => param.name === fileInputFieldFromExt)) {
-                            continue
-                        }
-                        nodeOverrides[nodeLabel].push({
-                            ...enabledFileParam,
-                            name: fileInputFieldFromExt
-                        })
-                    }
-                } else {
-                    const fileInputFieldFromExt = mapExtToInputField(enabledFileParam.type)
-                    nodeOverrides[nodeLabel].push({
-                        ...enabledFileParam,
-                        name: fileInputFieldFromExt
-                    })
-                }
-            }
-        }
         return { nodeOverrides, variableOverrides, apiOverrideStatus }
     } catch (error) {
         return { nodeOverrides: {}, variableOverrides: [], apiOverrideStatus: false }