refactor

Change-Id: Ibe8ce0c54f633092ff8dd63ee6044e4ea0f61fc6

Author: yinsu.zs

src/code/agent/.DS_Store

@@ -112,9 +112,19 @@
 # GPU 函数的 URL，当 COMFYUI_MODE="cpu" 时使用
 GPU_FUNCTION_URL = os.getenv("GPU_FUNCTION_URL", "")
 
+# HTTP Header 常量
+HEADER_SNAPSHOT_NAME = "X-FunArt-Snapshot-Name"
+HEADER_FORWARDED_BY = "X-FunArt-Forwarded-By"
+HEADER_FC_INVOCATION_TYPE = "X-FC-Invocation-Type"
+
 # 日志配置
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
 
+# 队列存储配置
+QUEUE_STORAGE_TYPE = os.getenv('QUEUE_STORAGE_TYPE', 'memory')  # memory | nas
+QUEUE_NAS_DIR = f'{MNT_DIR}/task/queue'  # 队列 NAS 存储目录
+QUEUE_LOCK_TIMEOUT = int(os.getenv('QUEUE_LOCK_TIMEOUT', '5'))  # 队列锁超时（秒）
+
 class ERROR_CODE(Enum):
     UNCLASSIFY = "UNCLASSIFY"
     INVALID_PARAMS = "INVALID_PARAMS"

src/code/agent/constants.py

@@ -6,6 +6,7 @@
 from flask import Blueprint, Flask, jsonify, request
 from flask_sock import Sock
 
+import constants
 from services.management_service import ManagementService, BackendStatus
 from utils.logger import log
 from services.gateway import (
@@ -77,36 +78,33 @@ def comfyui_compatible_ws(ws):
             """
             try:
                 # 从查询参数获取 clientId（ComfyUI 前端重连时会传递）
-                from flask import request as flask_request
-                client_id = flask_request.args.get('clientId', '')
+                client_id = request.args.get('clientId', '')
 
                 if client_id:
                     # 复用已有的 client_id（重连场景）
                     log("INFO", f"WebSocket reconnecting with existing client_id: {client_id}")
                 else:
                     # 生成新的 client_id（首次连接）
-                    client_id = f"cpu_client_{int(time.time() * 1000)}"
+                    client_id = f"funart_client_{int(time.time() * 1000)}"
                     log("INFO", f"New ComfyUI WebSocket connection with client_id: {client_id}")
 
                 # 添加连接到管理器
                 ws_manager.add_connection(ws)
 
-                # 发送初始状态消息（模拟ComfyUI原生行为）
-                try:
-                    ws.send(json.dumps({
-                        "type": "status",
-                        "data": {
-                            "sid": client_id,
-                            "status": {
-                                "exec_info": {
-                                    "queue_remaining": get_task_queue()._get_pending_task_count()
-                                }
+                # 通过消息队列发送初始状态消息，保证线程安全
+                initial_status = {
+                    "type": "status",
+                    "data": {
+                        "sid": client_id,
+                        "status": {
+                            "exec_info": {
+                                "queue_remaining": get_task_queue().get_running_task_count()
                             }
                         }
-                    }))
-                except Exception as e:
-                    log("ERROR", f"Failed to send initial status: {e}")
-                    return
+                    }
+                }
+                # 使用 _send_sync 发送初始状态
+                ws_manager._send_sync(ws, initial_status)
 
                 # 设置客户端ID，用于后续关联任务
                 setattr(ws, '_comfyui_client_id', client_id)
@@ -116,21 +114,19 @@ def comfyui_compatible_ws(ws):
 
                 # 如果是重连，重新订阅该客户端的所有进行中的任务
                 ws_manager.resubscribe_client_tasks(ws, client_id)
-                
-                # TODO 可能是多余的
+
                 while True:
                     try:
                         message = ws.receive()
                         log("DEBUG", f"Received message from ComfyUI frontend: {message[:100]}...")
-                        
+
                     except Exception as e:
                         error_str = str(e)
                         if "Connection closed" in error_str or "closed" in error_str.lower():
                             log("INFO", f"Connection closed by client")
                             break
                         log("ERROR", f"Error receiving message: {e}\n{traceback.format_exc()}")
                         break
-                    
             except Exception as e:
                 log("ERROR", f"Connection error: {e}\n{traceback.format_exc()}")
             finally:
@@ -202,7 +198,7 @@ def handle_serverless_run():
             
             调用方式:
             - 默认: 异步调用（与 /api/prompt 处理一致）
-            - Header X-Art-Invocation-Type: Sync 时: 同步调用，等待GPU返回结果
+            - Header X-FC-Invocation-Type: Sync 时: 同步调用，等待GPU返回结果
             
             异步模式:
             - 将请求转发到GPU函数（异步调用）
@@ -221,12 +217,12 @@ def handle_serverless_run():
             try:
                 gateway_service = CpuGatewayService()
 
-                # 检查调用类型：Header X-Art-Invocation-Type: Sync 表示同步调用
-                invocation_type = request.headers.get("X-Art-Invocation-Type", "").strip()
+                # TODO:  支持X-FC-Invocation-Type透传到Runtime,当前默认使用异步
+                invocation_type = request.headers.get(constants.HEADER_FC_INVOCATION_TYPE, "Async").strip()
                 is_sync = invocation_type.lower() == "sync"
 
                 if is_sync:
-                    log("DEBUG", f"Processing /serverless/run in SYNC mode (X-Art-Invocation-Type: Sync)")
+                    log("DEBUG", f"Processing /serverless/run in SYNC mode (X-FC-Invocation-Type: Sync)")
                     return gateway_service.handle_serverless_run_sync()
                 else:
                     log("DEBUG", f"Processing /serverless/run in ASYNC mode (default)")
@@ -267,10 +263,9 @@ def handle_history(subpath=""):
                 }), 500
 
     def _register_userdata_handler(self):
-        """在 prod 模式下，阻止保存 userdata 文件"""
         @self.bp.route("/userdata/<path:file>", methods=["POST"])
         def block_userdata_save(file):
-            log("WARN", f"Attempt to save userdata blocked in prod mode: {file}")
+            log("INFO", f"Disable Saving Userdata in prod mode: {file}")
             return jsonify({
                 "error": {
                     "type": "forbidden",

src/code/agent/routes/cpu_routes.py

@@ -1,5 +1,7 @@
 import json
+import logging
 import os
+import threading
 import traceback
 
 from flask import Flask, jsonify, request, Response
@@ -9,6 +11,7 @@
 import constants
 from exceptions.exceptions import CustomError
 from services.management_service import ManagementService, Action, BackendStatus
+from utils.logger import log
 from .management_routes import ManagementRoutes
 from .serverless_api_routes import ServerlessApiRoutes
 from .cpu_routes import CpuRoutes
@@ -19,16 +22,17 @@ class Routes:
     def __init__(self):
         self.app = Flask(__name__)
         self._sock = Sock(self.app)
+        # 重启锁，防止并发重启
+        self._reboot_lock = threading.Lock()
         self.setup_routes()
-        import logging
-        log = logging.getLogger('werkzeug')
-        log.setLevel(logging.ERROR)
+        # 设置 Werkzeug 日志级别为 ERROR，只显示错误日志，不输出每个请求
+        logging.getLogger('werkzeug').setLevel(logging.ERROR)
 
 
     def setup_routes(self):
         # 管控API
-        management = ManagementRoutes()
-        management.register(self.app)
+        self.management = ManagementRoutes()
+        self.management.register(self.app)
 
         # ServerlessAPI
         if constants.BACKEND_TYPE == constants.TYPE_COMFYUI:
@@ -39,11 +43,118 @@ def setup_routes(self):
             cpu_router = CpuRoutes()
             cpu_router.register(self.app)
 
+        @self.app.route("/api/manager/reboot", methods=["GET", "POST"])
+        def manager_reboot():
+            """
+            拦截 ComfyUI-Manager 的 reboot 请求，使用管控接口实现重启
+            支持从 CPU 函数触发：
+            - X-FunArt-Snapshot-Name header: 指定要使用的 snapshot 名称
+            - X-Forwarded-By header: 标识请求来源
+            
+            重启是异步的，立即返回响应，通过/management/status接口查询重启状态
+            """
+            # 检查是否已有重启在进行中
+            if not self._reboot_lock.acquire(blocking=False):
+                log("WARNING", "Reboot request rejected: reboot already in progress")
+                return jsonify({
+                    "status": "failed",
+                    "message": "Reboot already in progress, please wait"
+                }), 409
+            
+            # 优先从 header 中获取 snapshot 名称（CPU 函数传递）
+            snapshot_from_header = request.headers.get(constants.HEADER_SNAPSHOT_NAME)
+            forwarded_by = request.headers.get(constants.HEADER_FORWARDED_BY, 'Direct')
+            
+            service = ManagementService()
+            current_snapshot = service.cur_snapshot_name or 'latest-dev'
+            
+            log("INFO", f"Intercepted /api/manager/reboot request from {forwarded_by}, current snapshot: {current_snapshot}")
+            
+            # 异步执行重启逻辑
+            def do_reboot():
+                try:
+                    # 如果是 CPU 函数触发且指定了 snapshot，直接使用，不再保存
+                    if snapshot_from_header:
+                        log("INFO", f"Using snapshot from CPU function: {snapshot_from_header}")
+                        snapshot_to_load = snapshot_from_header
+                        skip_save = True
+                    else:
+                        snapshot_to_load = current_snapshot
+                        skip_save = False
+                    
+                    # 若最近一次管控操作为Start 且未指定 snapshot，则在重启前保存工作空间
+                    # TODO： 是否会影响pre-stop逻辑？
+                    if not skip_save and service.latest_action and service.latest_action == Action.START:
+                        try:
+                            from services.workspace.snapshot_manager import SnapshotManager
+                            log("INFO", "Saving workspace before reboot...")
+                            result_map = service.save(SnapshotManager.TYPE_DEV)
+                            log("INFO", f"Save result before reboot: {json.dumps(result_map, indent=2)}")
+                            
+                            # 使用新保存的 snapshot 名称
+                            if 'snapshot' in result_map:
+                                snapshot_to_load = result_map['snapshot']
+                                log("INFO", f"Will restart with newly saved snapshot: {snapshot_to_load}")
+                        except Exception as e:
+                            log("WARNING", f"Failed to save workspace before reboot: {str(e)}")
+                    else:
+                        log("INFO", "Skip saving workspace (latest_action is not START)")
+                    
+                    # 停止服务
+                    try:
+                        service.stop()
+                    except Exception as e:
+                        log("WARNING", f"Error during stop: {e}")
+                    
+                    # 如果是 CPU 模式，异步触发 GPU 函数的重启
+                    if constants.COMFYUI_MODE == 'cpu' and constants.GPU_FUNCTION_URL:
+                        def trigger_gpu_reboot():
+                            try:
+                                gpu_reboot_url = f"{constants.GPU_FUNCTION_URL.rstrip('/')}/api/manager/reboot"
+                                log("INFO", f"Triggering GPU function reboot: {gpu_reboot_url}")
+                                
+                                gpu_headers = {
+                                    constants.HEADER_SNAPSHOT_NAME: snapshot_to_load,
+                                    constants.HEADER_FORWARDED_BY: 'CPU-Reboot-Trigger',
+                                    constants.HEADER_FC_INVOCATION_TYPE: 'Async'
+                                }
+                                
+                                gpu_resp = requests.post(gpu_reboot_url, headers=gpu_headers, timeout=10)
+                                log("INFO", f"GPU function reboot triggered: status={gpu_resp.status_code}")
+                            except requests.exceptions.Timeout:
+                                log("WARNING", "GPU function reboot request timed out (expected for async call)")
+                            except Exception as e:
+                                log("WARNING", f"Failed to trigger GPU function reboot: {str(e)}")
+                        
+                        # 触发 GPU 重启
+                        threading.Thread(target=trigger_gpu_reboot, daemon=True).start()
+                    
+                    # 重新启动本地服务（不安装依赖），使用新保存的 snapshot
+                    log("INFO", f"Restarting local service with snapshot: {snapshot_to_load}")
+                    service.start(snapshot_to_load, nodes_map=service.SKIP_INSTALL_SENTINEL)
+                    log("INFO", "Reboot completed successfully")
+                    
+                except Exception as e:
+                    error_msg = f"Failed to restart ComfyUI: {str(e)}"
+                    log("ERROR", f"{error_msg}\nStacktrace:\n{traceback.format_exc()}")
+                finally:
+                    # 释放重启锁
+                    self._reboot_lock.release()
+            
+            # 在后台线程中执行重启
+            threading.Thread(target=do_reboot, daemon=True).start()
+            
+            # 立即返回响应
+            return jsonify({
+                "status": "success",
+                "message": "Reboot request accepted, restarting in background"
+            }), 202
+        
         @self.app.route("/initialize", methods=["POST"])
         def initialize():
             # See FC docs for all the HTTP headers: https://www.alibabacloud.com/help/doc-detail/132044.htm#common-headers
             request_id = request.headers.get("x-fc-request-id", "")
-            print("FC Initialize Start RequestId: " + request_id)
+            log("INFO", f"FC Initialize Start RequestId: {request_id}")
 
             # Use the following code to get temporary credentials
             # access_key_id = request.headers['x-fc-access-key-id']
@@ -56,7 +167,7 @@ def initialize():
 
             # 使用环境变量指定的snapshot，默认为latest-dev
             snapshot_name = os.environ.get('AUTO_LAUNCH_SNAPSHOT_NAME', 'latest-dev')
-            print(f"Initializing function with ComfyUI mode: {constants.COMFYUI_MODE}, snapshot: {snapshot_name}")
+            log("INFO", f"Initializing function with ComfyUI mode: {constants.COMFYUI_MODE}, snapshot: {snapshot_name}")
             service.start(snapshot_name, nodes_map={})
 
             if (
@@ -65,22 +176,22 @@ def initialize():
                 and constants.COMFYUI_MODE == "gpu"
             ):
                 try:
-                    print("prewarm models")
+                    log("INFO", "prewarm models")
                     prompt = json.loads(constants.PREWARM_PROMPT)
                     api = ServerlessApiService()
                     api.run(prompt)
                     api.api_clear_history()
-                    print("prewarm models done")
+                    log("INFO", "prewarm models done")
                 except Exception as e:
-                    print(f"prewarm models got exception:\n{e}")
+                    log("ERROR", f"prewarm models got exception:\n{e}")
 
-            print("FC Initialize End RequestId: " + request_id)
+            log("INFO", f"FC Initialize End RequestId: {request_id}")
             return "Function is initialized, request_id: " + request_id + "\n"
 
         @self.app.route("/pre-stop", methods=["GET"])
         def pre_stop():
             request_id = request.headers.get("x-fc-request-id", "")
-            print("FC PreStop Start RequestId: " + request_id)
+            log("INFO", f"FC PreStop Start RequestId: {request_id}")
 
             service = ManagementService()  # singleton
             # 若最近一次管控操作为Start，且实例非预期销毁时，需要在pre-stop中保存工作空间从而兜底;
@@ -89,18 +200,18 @@ def pre_stop():
                 try:
                     from services.workspace.snapshot_manager import SnapshotManager
                     result_map = service.save(SnapshotManager.TYPE_DEV)
-                    print(f"save resp when preStop: {json.dumps(result_map, indent=2)}")
+                    log("INFO", f"save resp when preStop: {json.dumps(result_map, indent=2)}")
                 except Exception as e:
-                    print(f"error occur when preStop: {str(e)}")
+                    log("ERROR", f"error occur when preStop: {str(e)}")
             else:
-                print("Do nothing in pre-stop")
-            print("FC PreStop End RequestId: " + request_id)
+                log("INFO", "Do nothing in pre-stop")
+            log("INFO", f"FC PreStop End RequestId: {request_id}")
             return "OK"
 
         @self.app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
         @self.app.route("/", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
         def proxy(path=""):
-            backend_status = management.service.status
+            backend_status = self.management.service.status
             if backend_status not in (BackendStatus.RUNNING, BackendStatus.SAVING):
                 return jsonify({
                     "status": "failed",
@@ -146,7 +257,7 @@ def handle_base_error(error):
 
         def _handle_exception(e):
             err_msg = traceback.format_exc()
-            print(f"{str(e)}\nStacktrace:\n{err_msg}")
+            log("ERROR", f"{str(e)}\nStacktrace:\n{err_msg}")
 
             if isinstance(e, CustomError):
                 # 处理自定义异常