From 1122d172edcfbb58fa3a2a45ff6f739495adfbff Mon Sep 17 00:00:00 2001
From: OhYee <oyohyee@oyohyee.com>
Date: Thu, 24 Jul 2025 16:00:08 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=A8=A1=E5=9E=8B=E9=A2=84=E7=83=AD?=
 =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E8=B6=85=E6=97=B6=E6=97=B6=E9=97=B4=EF=BC=8C?=
 =?UTF-8?q?=E8=B6=85=E6=97=B6=E8=87=AA=E5=8A=A8=E8=B7=B3=E8=BF=87=EF=BC=8C?=
 =?UTF-8?q?=E8=B6=85=E6=97=B6=E6=97=B6=E9=97=B4=E5=8F=AF=E9=80=9A=E8=BF=87?=
 =?UTF-8?q?=20PREWARM=5FTIMEOUT=20=E7=8E=AF=E5=A2=83=E5=8F=98=E9=87=8F?=
 =?UTF-8?q?=E9=85=8D=E7=BD=AE=EF=BC=8C=E9=BB=98=E8=AE=A410=E5=88=86?=
 =?UTF-8?q?=E9=92=9F=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/code/agent/constants.py     |  2 ++
 src/code/agent/routes/routes.py | 35 +++++++++++++++++++++++----------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/code/agent/constants.py b/src/code/agent/constants.py
index 60964992..8a53f0b8 100644
--- a/src/code/agent/constants.py
+++ b/src/code/agent/constants.py
@@ -75,7 +75,9 @@
 OSS_OUTPUT_DOMAIN = os.getenv("OSS_OUTPUT_DOMAIN", "")
 OSS_EXPIRES_IN_SECOND = os.getenv("OSS_EXPIRES_IN_SECOND", "")
 
+
 PREWARM_PROMPT = os.getenv("PREWARM_PROMPT", "")
+PREWARM_TIMEOUT = int(os.getenv("PREWARM_TIMEOUT", "600"))  # 单位：秒，默认10分钟
 
 class ERROR_CODE(Enum):
     UNCLASSIFY = "UNCLASSIFY"
diff --git a/src/code/agent/routes/routes.py b/src/code/agent/routes/routes.py
index b590a8dd..d3a8e1e9 100644
--- a/src/code/agent/routes/routes.py
+++ b/src/code/agent/routes/routes.py
@@ -1,6 +1,7 @@
 import json
 import logging
 import threading
+import time
 import traceback
 
 import requests
@@ -54,18 +55,32 @@ def initialize():
                 constants.PREWARM_PROMPT
                 and constants.BACKEND_TYPE == constants.TYPE_COMFYUI
             ):
-                try:
-                    print("prewarm models")
-                    prompt = json.loads(constants.PREWARM_PROMPT)
-                    api = ServerlessApiService()
-                    api.run(prompt)
-                    api.api_clear_history()
-                    print("prewarm models done")
-                except Exception as e:
-                    print(f"prewarm models got exception:\n{e}")
+                prewarm_done = threading.Event()
+                prewarm_result = {"exception": None}
+
+                t = threading.Thread(target=self.prewarm_func, args=(prewarm_done, prewarm_result))
+                t.start()
+                t.join(timeout=constants.PREWARM_TIMEOUT)
+                if not prewarm_done.is_set():
+                    print(f"prewarm timeout after {constants.PREWARM_TIMEOUT} seconds, skip prewarm and continue.")
+                # 如果有异常也继续往下执行
 
             print("FC Initialize End RequestId: " + request_id)
-            return "Function is initialized, request_id: " + request_id + "\n"
+
+    def prewarm_func(self, prewarm_done, prewarm_result):
+        try:
+            print("prewarm models")
+            prompt = json.loads(constants.PREWARM_PROMPT)
+            api = ServerlessApiService()
+            api.run(prompt)
+            api.api_clear_history()
+            print("prewarm models done")
+        except Exception as e:
+            prewarm_result["exception"] = e
+            print(f"prewarm models got exception:\n{e}")
+        finally:
+            prewarm_done.set()
+
 
         @self.app.route("/pre-stop", methods=["GET"])
         def pre_stop():