perf: memory leak (#5370)

* perf: memory leak

* perf: workflow share buffer;Circle checker;Get file from stream

* doc

* remove report.md

projects/app/src/service/core/dataset/queues/generateQA.ts

@@ -44,157 +44,161 @@ export async function generateQA(): Promise<any> {
   if (global.qaQueueLen >= max) return;
   global.qaQueueLen++;
 
-  while (true) {
-    const startTime = Date.now();
-    // get training data
-    const {
-      data,
-      text,
-      done = false,
-      error = false
-    } = await (async () => {
-      try {
-        const data = await MongoDatasetTraining.findOneAndUpdate(
-          {
-            mode: TrainingModeEnum.qa,
-            retryCount: { $gt: 0 },
-            lockTime: { $lte: addMinutes(new Date(), -10) }
-          },
-          {
-            lockTime: new Date(),
-            $inc: { retryCount: -1 }
-          }
-        )
-          .populate<PopulateType>([
+  try {
+    while (true) {
+      const startTime = Date.now();
+      // get training data
+      const {
+        data,
+        text,
+        done = false,
+        error = false
+      } = await (async () => {
+        try {
+          const data = await MongoDatasetTraining.findOneAndUpdate(
             {
-              path: 'dataset',
-              select: 'agentModel vectorModel vlmModel'
+              mode: TrainingModeEnum.qa,
+              retryCount: { $gt: 0 },
+              lockTime: { $lte: addMinutes(new Date(), -10) }
             },
             {
-              path: 'collection',
-              select: 'qaPrompt'
+              lockTime: new Date(),
+              $inc: { retryCount: -1 }
             }
-          ])
-          .lean();
+          )
+            .populate<PopulateType>([
+              {
+                path: 'dataset',
+                select: 'agentModel vectorModel vlmModel'
+              },
+              {
+                path: 'collection',
+                select: 'qaPrompt'
+              }
+            ])
+            .lean();
 
-        // task preemption
-        if (!data) {
+          // task preemption
+          if (!data) {
+            return {
+              done: true
+            };
+          }
           return {
-            done: true
+            data,
+            text: data.q
+          };
+        } catch (error) {
+          return {
+            error: true
           };
         }
-        return {
-          data,
-          text: data.q
-        };
-      } catch (error) {
-        return {
-          error: true
-        };
+      })();
+
+      if (done || !data) {
+        break;
+      }
+      if (error) {
+        addLog.error(`[QA Queue] Error`, error);
+        await delay(500);
+        continue;
       }
-    })();
 
-    if (done || !data) {
-      break;
-    }
-    if (error) {
-      addLog.error(`[QA Queue] Error`, error);
-      await delay(500);
-      continue;
-    }
+      if (!data.dataset || !data.collection) {
+        addLog.info(`[QA Queue] Dataset or collection not found`, data);
+        // Delete data
+        await MongoDatasetTraining.deleteOne({ _id: data._id });
+        continue;
+      }
+      // auth balance
+      if (!(await checkTeamAiPointsAndLock(data.teamId))) {
+        continue;
+      }
 
-    if (!data.dataset || !data.collection) {
-      addLog.info(`[QA Queue] Dataset or collection not found`, data);
-      // Delete data
-      await MongoDatasetTraining.deleteOne({ _id: data._id });
-      continue;
-    }
-    // auth balance
-    if (!(await checkTeamAiPointsAndLock(data.teamId))) {
-      continue;
-    }
+      addLog.info(`[QA Queue] Start`);
 
-    addLog.info(`[QA Queue] Start`);
+      try {
+        const modelData = getLLMModel(data.dataset.agentModel);
+        const prompt = `${data.collection.qaPrompt || Prompt_AgentQA.description}
+  ${replaceVariable(Prompt_AgentQA.fixedText, { text })}`;
 
-    try {
-      const modelData = getLLMModel(data.dataset.agentModel);
-      const prompt = `${data.collection.qaPrompt || Prompt_AgentQA.description}
-${replaceVariable(Prompt_AgentQA.fixedText, { text })}`;
-
-      // request LLM to get QA
-      const messages: ChatCompletionMessageParam[] = [
-        {
-          role: 'user',
-          content: prompt
-        }
-      ];
-
-      const { response: chatResponse } = await createChatCompletion({
-        body: llmCompletionsBodyFormat(
+        // request LLM to get QA
+        const messages: ChatCompletionMessageParam[] = [
           {
-            model: modelData.model,
-            temperature: 0.3,
-            messages: await loadRequestMessages({ messages, useVision: false }),
-            stream: true
+            role: 'user',
+            content: prompt
+          }
+        ];
+
+        const { response: chatResponse } = await createChatCompletion({
+          body: llmCompletionsBodyFormat(
+            {
+              model: modelData.model,
+              temperature: 0.3,
+              messages: await loadRequestMessages({ messages, useVision: false }),
+              stream: true
+            },
+            modelData
+          )
+        });
+        const { text: answer, usage } = await formatLLMResponse(chatResponse);
+        const inputTokens = usage?.prompt_tokens || (await countGptMessagesTokens(messages));
+        const outputTokens = usage?.completion_tokens || (await countPromptTokens(answer));
+
+        const qaArr = await formatSplitText({ answer, rawText: text, llmModel: modelData }); // 格式化后的QA对
+
+        // get vector and insert
+        await pushDataListToTrainingQueue({
+          teamId: data.teamId,
+          tmbId: data.tmbId,
+          datasetId: data.datasetId,
+          collectionId: data.collectionId,
+          mode: TrainingModeEnum.chunk,
+          data: qaArr.map((item) => ({
+            ...item,
+            chunkIndex: data.chunkIndex
+          })),
+          billId: data.billId,
+          vectorModel: data.dataset.vectorModel,
+          agentModel: data.dataset.agentModel,
+          vlmModel: data.dataset.vlmModel
+        });
+
+        // delete data from training
+        await MongoDatasetTraining.findByIdAndDelete(data._id);
+
+        // Push usage
+        pushLLMTrainingUsage({
+          teamId: data.teamId,
+          tmbId: data.tmbId,
+          inputTokens,
+          outputTokens,
+          billId: data.billId,
+          model: modelData.model,
+          mode: 'qa'
+        });
+
+        addLog.info(`[QA Queue] Finish`, {
+          time: Date.now() - startTime,
+          splitLength: qaArr.length,
+          usage
+        });
+      } catch (err: any) {
+        addLog.error(`[QA Queue] Error`, err);
+        await MongoDatasetTraining.updateOne(
+          {
+            _id: data._id
           },
-          modelData
-        )
-      });
-      const { text: answer, usage } = await formatLLMResponse(chatResponse);
-      const inputTokens = usage?.prompt_tokens || (await countGptMessagesTokens(messages));
-      const outputTokens = usage?.completion_tokens || (await countPromptTokens(answer));
+          {
+            errorMsg: getErrText(err, 'unknown error')
+          }
+        );
 
-      const qaArr = await formatSplitText({ answer, rawText: text, llmModel: modelData }); // 格式化后的QA对
-
-      // get vector and insert
-      await pushDataListToTrainingQueue({
-        teamId: data.teamId,
-        tmbId: data.tmbId,
-        datasetId: data.datasetId,
-        collectionId: data.collectionId,
-        mode: TrainingModeEnum.chunk,
-        data: qaArr.map((item) => ({
-          ...item,
-          chunkIndex: data.chunkIndex
-        })),
-        billId: data.billId,
-        vectorModel: data.dataset.vectorModel,
-        agentModel: data.dataset.agentModel,
-        vlmModel: data.dataset.vlmModel
-      });
-
-      // delete data from training
-      await MongoDatasetTraining.findByIdAndDelete(data._id);
-
-      // Push usage
-      pushLLMTrainingUsage({
-        teamId: data.teamId,
-        tmbId: data.tmbId,
-        inputTokens,
-        outputTokens,
-        billId: data.billId,
-        model: modelData.model,
-        mode: 'qa'
-      });
-
-      addLog.info(`[QA Queue] Finish`, {
-        time: Date.now() - startTime,
-        splitLength: qaArr.length,
-        usage
-      });
-    } catch (err: any) {
-      addLog.error(`[QA Queue] Error`, err);
-      await MongoDatasetTraining.updateOne(
-        {
-          _id: data._id
-        },
-        {
-          errorMsg: getErrText(err, 'unknown error')
-        }
-      );
-
-      await delay(100);
+        await delay(100);
+      }
     }
+  } catch (error) {
+    addLog.error(`[QA Queue] Error`, error);
   }
 
   if (reduceQueue()) {

projects/app/src/service/core/dataset/queues/generateVector.ts

@@ -42,116 +42,120 @@ export async function generateVector(): Promise<any> {
   if (global.vectorQueueLen >= max) return;
   global.vectorQueueLen++;
 
-  while (true) {
-    const start = Date.now();
+  try {
+    while (true) {
+      const start = Date.now();
 
-    // get training data
-    const {
-      data,
-      done = false,
-      error = false
-    } = await (async () => {
-      try {
-        const data = await MongoDatasetTraining.findOneAndUpdate(
-          {
-            mode: TrainingModeEnum.chunk,
-            retryCount: { $gt: 0 },
-            lockTime: { $lte: addMinutes(new Date(), -3) }
-          },
-          {
-            lockTime: new Date(),
-            $inc: { retryCount: -1 }
-          }
-        )
-          .populate<PopulateType>([
+      // get training data
+      const {
+        data,
+        done = false,
+        error = false
+      } = await (async () => {
+        try {
+          const data = await MongoDatasetTraining.findOneAndUpdate(
             {
-              path: 'dataset',
-              select: 'vectorModel'
+              mode: TrainingModeEnum.chunk,
+              retryCount: { $gt: 0 },
+              lockTime: { $lte: addMinutes(new Date(), -3) }
             },
             {
-              path: 'collection',
-              select: 'name indexPrefixTitle'
-            },
-            {
-              path: 'data',
-              select: '_id indexes'
+              lockTime: new Date(),
+              $inc: { retryCount: -1 }
             }
-          ])
-          .lean();
+          )
+            .populate<PopulateType>([
+              {
+                path: 'dataset',
+                select: 'vectorModel'
+              },
+              {
+                path: 'collection',
+                select: 'name indexPrefixTitle'
+              },
+              {
+                path: 'data',
+                select: '_id indexes'
+              }
+            ])
+            .lean();
 
-        // task preemption
-        if (!data) {
+          // task preemption
+          if (!data) {
+            return {
+              done: true
+            };
+          }
           return {
-            done: true
+            data
+          };
+        } catch (error) {
+          return {
+            error: true
           };
-        }
-        return {
-          data
-        };
-      } catch (error) {
-        return {
-          error: true
-        };
-      }
-    })();
-
-    // Break loop
-    if (done || !data) {
-      break;
-    }
-    if (error) {
-      addLog.error(`[Vector Queue] Error`, error);
-      await delay(500);
-      continue;
-    }
-
-    if (!data.dataset || !data.collection) {
-      addLog.info(`[Vector Queue] Dataset or collection not found`, data);
-      // Delete data
-      await MongoDatasetTraining.deleteOne({ _id: data._id });
-      continue;
-    }
-
-    // auth balance
-    if (!(await checkTeamAiPointsAndLock(data.teamId))) {
-      continue;
-    }
-
-    addLog.info(`[Vector Queue] Start`);
-
-    try {
-      const { tokens } = await (async () => {
-        if (data.dataId) {
-          return rebuildData({ trainingData: data });
-        } else {
-          return insertData({ trainingData: data });
         }
       })();
 
-      // push usage
-      pushGenerateVectorUsage({
-        teamId: data.teamId,
-        tmbId: data.tmbId,
-        inputTokens: tokens,
-        model: data.dataset.vectorModel,
-        billId: data.billId
-      });
+      // Break loop
+      if (done || !data) {
+        break;
+      }
+      if (error) {
+        addLog.error(`[Vector Queue] Error`, error);
+        await delay(500);
+        continue;
+      }
 
-      addLog.info(`[Vector Queue] Finish`, {
-        time: Date.now() - start
-      });
-    } catch (err: any) {
-      addLog.error(`[Vector Queue] Error`, err);
-      await MongoDatasetTraining.updateOne(
-        {
-          _id: data._id
-        },
-        {
-          errorMsg: getErrText(err, 'unknown error')
-        }
-      );
-      await delay(100);
+      if (!data.dataset || !data.collection) {
+        addLog.info(`[Vector Queue] Dataset or collection not found`, data);
+        // Delete data
+        await MongoDatasetTraining.deleteOne({ _id: data._id });
+        continue;
+      }
+
+      // auth balance
+      if (!(await checkTeamAiPointsAndLock(data.teamId))) {
+        continue;
+      }
+
+      addLog.info(`[Vector Queue] Start`);
+
+      try {
+        const { tokens } = await (async () => {
+          if (data.dataId) {
+            return rebuildData({ trainingData: data });
+          } else {
+            return insertData({ trainingData: data });
+          }
+        })();
+
+        // push usage
+        pushGenerateVectorUsage({
+          teamId: data.teamId,
+          tmbId: data.tmbId,
+          inputTokens: tokens,
+          model: data.dataset.vectorModel,
+          billId: data.billId
+        });
+
+        addLog.info(`[Vector Queue] Finish`, {
+          time: Date.now() - start
+        });
+      } catch (err: any) {
+        addLog.error(`[Vector Queue] Error`, err);
+        await MongoDatasetTraining.updateOne(
+          {
+            _id: data._id
+          },
+          {
+            errorMsg: getErrText(err, 'unknown error')
+          }
+        );
+        await delay(100);
+      }
     }
+  } catch (error) {
+    addLog.error(`[Vector Queue] Error`, error);
   }
 
   if (reduceQueue()) {