Merge branch 'dev' into prcix9320

This commit is contained in:
q434343
2026-03-25 14:44:52 +08:00
parent 792504e08c
commit 68029217de
97 changed files with 15504 additions and 8906 deletions

View File

@@ -1,8 +1,10 @@
import argparse
import asyncio
import os
import platform
import shutil
import signal
import subprocess
import sys
import threading
import time
@@ -24,6 +26,84 @@ from unilabos.config.config import load_config, BasicConfig, HTTPConfig
_restart_requested: bool = False
_restart_reason: str = ""
RESTART_EXIT_CODE = 42
def _build_child_argv():
"""Build sys.argv for child process, stripping supervisor-only arguments."""
result = []
skip_next = False
for arg in sys.argv:
if skip_next:
skip_next = False
continue
if arg in ("--restart_mode", "--restart-mode"):
continue
if arg in ("--auto_restart_count", "--auto-restart-count"):
skip_next = True
continue
if arg.startswith("--auto_restart_count=") or arg.startswith("--auto-restart-count="):
continue
result.append(arg)
return result
def _run_as_supervisor(max_restarts: int):
"""
Supervisor process that spawns and monitors child processes.
Similar to Uvicorn's --reload: the supervisor itself does no heavy work,
it only launches the real process as a child and restarts it when the child
exits with RESTART_EXIT_CODE.
"""
child_argv = [sys.executable] + _build_child_argv()
restart_count = 0
print_status(
f"[Supervisor] Restart mode enabled (max restarts: {max_restarts}), "
f"child command: {' '.join(child_argv)}",
"info",
)
while True:
print_status(
f"[Supervisor] Launching process (restart {restart_count}/{max_restarts})...",
"info",
)
try:
process = subprocess.Popen(child_argv)
exit_code = process.wait()
except KeyboardInterrupt:
print_status("[Supervisor] Interrupted, terminating child process...", "info")
process.terminate()
try:
process.wait(timeout=10)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
sys.exit(1)
if exit_code == RESTART_EXIT_CODE:
restart_count += 1
if restart_count > max_restarts:
print_status(
f"[Supervisor] Maximum restart count ({max_restarts}) reached, exiting",
"warning",
)
sys.exit(1)
print_status(
f"[Supervisor] Child requested restart ({restart_count}/{max_restarts}), restarting in 2s...",
"info",
)
time.sleep(2)
else:
if exit_code != 0:
print_status(f"[Supervisor] Child exited with code {exit_code}", "warning")
else:
print_status("[Supervisor] Child exited normally", "info")
sys.exit(exit_code)
def load_config_from_file(config_path):
if config_path is None:
@@ -65,6 +145,13 @@ def parse_args():
action="append",
help="Path to the registry directory",
)
parser.add_argument(
"--devices",
type=str,
default=None,
action="append",
help="Path to Python code directory for AST-based device/resource scanning",
)
parser.add_argument(
"--working_dir",
type=str,
@@ -154,18 +241,18 @@ def parse_args():
action="store_true",
help="Skip environment dependency check on startup",
)
parser.add_argument(
"--complete_registry",
action="store_true",
default=False,
help="Complete registry information",
)
parser.add_argument(
"--check_mode",
action="store_true",
default=False,
help="Run in check mode for CI: validates registry imports and ensures no file changes",
)
parser.add_argument(
"--complete_registry",
action="store_true",
default=False,
help="Complete and rewrite YAML registry files using AST analysis results",
)
parser.add_argument(
"--no_update_feedback",
action="store_true",
@@ -177,6 +264,30 @@ def parse_args():
default=False,
help="Test mode: all actions simulate execution and return mock results without running real hardware",
)
parser.add_argument(
"--external_devices_only",
action="store_true",
default=False,
help="Only load external device packages (--devices), skip built-in unilabos/devices/ scanning and YAML device registry",
)
parser.add_argument(
"--extra_resource",
action="store_true",
default=False,
help="Load extra lab_ prefixed labware resources (529 auto-generated definitions from lab_resources.py)",
)
parser.add_argument(
"--restart_mode",
action="store_true",
default=False,
help="Enable supervisor mode: automatically restart the process when triggered via WebSocket",
)
parser.add_argument(
"--auto_restart_count",
type=int,
default=500,
help="Maximum number of automatic restarts in restart mode (default: 500)",
)
# workflow upload subcommand
workflow_parser = subparsers.add_parser(
"workflow_upload",
@@ -227,16 +338,28 @@ def main():
args = parser.parse_args()
args_dict = vars(args)
# Supervisor mode: spawn child processes and monitor for restart
if args_dict.get("restart_mode", False):
_run_as_supervisor(args_dict.get("auto_restart_count", 5))
return
# 环境检查 - 检查并自动安装必需的包 (可选)
skip_env_check = args_dict.get("skip_env_check", False)
check_mode = args_dict.get("check_mode", False)
if not skip_env_check:
from unilabos.utils.environment_check import check_environment
from unilabos.utils.environment_check import check_environment, check_device_package_requirements
if not check_environment(auto_install=True):
print_status("环境检查失败,程序退出", "error")
os._exit(1)
# 第一次设备包依赖检查build_registry 之前,确保 import map 可用
devices_dirs_for_req = args_dict.get("devices", None)
if devices_dirs_for_req:
if not check_device_package_requirements(devices_dirs_for_req):
print_status("设备包依赖检查失败,程序退出", "error")
os._exit(1)
else:
print_status("跳过环境依赖检查", "warning")
@@ -357,46 +480,63 @@ def main():
BasicConfig.test_mode = args_dict.get("test_mode", False)
if BasicConfig.test_mode:
print_status("启用测试模式:所有动作将模拟执行,不调用真实硬件", "warning")
BasicConfig.extra_resource = args_dict.get("extra_resource", False)
if BasicConfig.extra_resource:
print_status("启用额外资源加载将加载lab_开头的labware资源定义", "info")
BasicConfig.communication_protocol = "websocket"
machine_name = os.popen("hostname").read().strip()
machine_name = platform.node()
machine_name = "".join([c if c.isalnum() or c == "_" else "_" for c in machine_name])
BasicConfig.machine_name = machine_name
BasicConfig.vis_2d_enable = args_dict["2d_vis"]
BasicConfig.check_mode = check_mode
from unilabos.resources.graphio import (
read_node_link_json,
read_graphml,
dict_from_graph,
)
from unilabos.app.communication import get_communication_client
from unilabos.registry.registry import build_registry
from unilabos.app.backend import start_backend
from unilabos.app.web import http_client
from unilabos.app.web import start_server
from unilabos.app.register import register_devices_and_resources
from unilabos.resources.graphio import modify_to_backend_format
from unilabos.resources.resource_tracker import ResourceTreeSet, ResourceDict
# 显示启动横幅
print_unilab_banner(args_dict)
# 注册表 - check_mode 时强制启用 complete_registry
# Step 0: AST 分析优先 + YAML 注册表加载
# check_mode 和 upload_registry 都会执行实际 import 验证
devices_dirs = args_dict.get("devices", None)
complete_registry = args_dict.get("complete_registry", False) or check_mode
lab_registry = build_registry(args_dict["registry_path"], complete_registry, BasicConfig.upload_registry)
external_only = args_dict.get("external_devices_only", False)
lab_registry = build_registry(
registry_paths=args_dict["registry_path"],
devices_dirs=devices_dirs,
upload_registry=BasicConfig.upload_registry,
check_mode=check_mode,
complete_registry=complete_registry,
external_only=external_only,
)
# Check mode: complete_registry 完成后直接退出git diff 检测由 CI workflow 执行
# Check mode: 注册表验证完成后直接退出
if check_mode:
print_status("Check mode: complete_registry 完成,退出", "info")
device_count = len(lab_registry.device_type_registry)
resource_count = len(lab_registry.resource_type_registry)
print_status(f"Check mode: 注册表验证完成 ({device_count} 设备, {resource_count} 资源),退出", "info")
os._exit(0)
# 以下导入依赖 ROS2 环境check_mode 已退出不需要
from unilabos.resources.graphio import (
read_node_link_json,
read_graphml,
dict_from_graph,
modify_to_backend_format,
)
from unilabos.app.communication import get_communication_client
from unilabos.app.backend import start_backend
from unilabos.app.web import http_client
from unilabos.app.web import start_server
from unilabos.app.register import register_devices_and_resources
from unilabos.resources.resource_tracker import ResourceTreeSet, ResourceDict
# Step 1: 上传全部注册表到服务端,同步保存到 unilabos_data
if BasicConfig.upload_registry:
# 设备注册到服务端 - 需要 ak 和 sk
if BasicConfig.ak and BasicConfig.sk:
print_status("开始注册设备到服务端...", "info")
# print_status("开始注册设备到服务端...", "info")
try:
register_devices_and_resources(lab_registry)
print_status("设备注册完成", "info")
# print_status("设备注册完成", "info")
except Exception as e:
print_status(f"设备注册失败: {e}", "error")
else:
@@ -481,12 +621,16 @@ def main():
continue
# 如果从远端获取了物料信息,则与本地物料进行同步
if request_startup_json and "nodes" in request_startup_json:
if file_path is not None and request_startup_json and "nodes" in request_startup_json:
print_status("开始同步远端物料到本地...", "info")
remote_tree_set = ResourceTreeSet.from_raw_dict_list(request_startup_json["nodes"])
resource_tree_set.merge_remote_resources(remote_tree_set)
print_status("远端物料同步完成", "info")
# 第二次设备包依赖检查云端物料同步后community 包可能引入新的 requirements
# TODO: 当 community device package 功能上线后,在这里调用
# install_requirements_txt(community_pkg_path / "requirements.txt", label="community.xxx")
# 使用 ResourceTreeSet 代替 list
args_dict["resources_config"] = resource_tree_set
args_dict["devices_config"] = resource_tree_set
@@ -578,6 +722,10 @@ def main():
open_browser=not args_dict["disable_browser"],
port=BasicConfig.port,
)
if restart_requested:
print_status("[Main] Restart requested, cleaning up...", "info")
cleanup_for_restart()
os._exit(RESTART_EXIT_CODE)
if __name__ == "__main__":

View File

@@ -1,9 +1,8 @@
import json
import time
from typing import Optional, Tuple, Dict, Any
from typing import Any, Dict, Optional, Tuple
from unilabos.utils.log import logger
from unilabos.utils.type_check import TypeEncoder
from unilabos.utils.tools import normalize_json as _normalize_device
def register_devices_and_resources(lab_registry, gather_only=False) -> Optional[Tuple[Dict[str, Any], Dict[str, Any]]]:
@@ -11,50 +10,63 @@ def register_devices_and_resources(lab_registry, gather_only=False) -> Optional[
注册设备和资源到服务器仅支持HTTP
"""
# 注册资源信息 - 使用HTTP方式
from unilabos.app.web.client import http_client
logger.info("[UniLab Register] 开始注册设备和资源...")
# 注册设备信息
devices_to_register = {}
for device_info in lab_registry.obtain_registry_device_info():
devices_to_register[device_info["id"]] = json.loads(
json.dumps(device_info, ensure_ascii=False, cls=TypeEncoder)
)
logger.debug(f"[UniLab Register] 收集设备: {device_info['id']}")
devices_to_register[device_info["id"]] = _normalize_device(device_info)
logger.trace(f"[UniLab Register] 收集设备: {device_info['id']}")
resources_to_register = {}
for resource_info in lab_registry.obtain_registry_resource_info():
resources_to_register[resource_info["id"]] = resource_info
logger.debug(f"[UniLab Register] 收集资源: {resource_info['id']}")
logger.trace(f"[UniLab Register] 收集资源: {resource_info['id']}")
if gather_only:
return devices_to_register, resources_to_register
# 注册设备
if devices_to_register:
try:
start_time = time.time()
response = http_client.resource_registry({"resources": list(devices_to_register.values())})
response = http_client.resource_registry(
{"resources": list(devices_to_register.values())},
tag="device_registry",
)
cost_time = time.time() - start_time
if response.status_code in [200, 201]:
logger.info(f"[UniLab Register] 成功注册 {len(devices_to_register)} 个设备 {cost_time}s")
res_data = response.json() if response.status_code == 200 else {}
skipped = res_data.get("data", {}).get("skipped", False)
if skipped:
logger.info(
f"[UniLab Register] 设备注册跳过(内容未变化)"
f" {len(devices_to_register)}{cost_time:.3f}s"
)
elif response.status_code in [200, 201]:
logger.info(f"[UniLab Register] 成功注册 {len(devices_to_register)} 个设备 {cost_time:.3f}s")
else:
logger.error(f"[UniLab Register] 设备注册失败: {response.status_code}, {response.text} {cost_time}s")
logger.error(f"[UniLab Register] 设备注册失败: {response.status_code}, {response.text} {cost_time:.3f}s")
except Exception as e:
logger.error(f"[UniLab Register] 设备注册异常: {e}")
# 注册资源
if resources_to_register:
try:
start_time = time.time()
response = http_client.resource_registry({"resources": list(resources_to_register.values())})
response = http_client.resource_registry(
{"resources": list(resources_to_register.values())},
tag="resource_registry",
)
cost_time = time.time() - start_time
if response.status_code in [200, 201]:
logger.info(f"[UniLab Register] 成功注册 {len(resources_to_register)} 个资源 {cost_time}s")
res_data = response.json() if response.status_code == 200 else {}
skipped = res_data.get("data", {}).get("skipped", False)
if skipped:
logger.info(
f"[UniLab Register] 资源注册跳过(内容未变化)"
f" {len(resources_to_register)}{cost_time:.3f}s"
)
elif response.status_code in [200, 201]:
logger.info(f"[UniLab Register] 成功注册 {len(resources_to_register)} 个资源 {cost_time:.3f}s")
else:
logger.error(f"[UniLab Register] 资源注册失败: {response.status_code}, {response.text} {cost_time}s")
logger.error(f"[UniLab Register] 资源注册失败: {response.status_code}, {response.text} {cost_time:.3f}s")
except Exception as e:
logger.error(f"[UniLab Register] 资源注册异常: {e}")
logger.info("[UniLab Register] 设备和资源注册完成.")

View File

@@ -1052,7 +1052,7 @@ async def handle_file_import(websocket: WebSocket, request_data: dict):
"result": {},
"schema": lab_registry._generate_unilab_json_command_schema(v["args"], k),
"goal_default": {i["name"]: i["default"] for i in v["args"]},
"handles": [],
"handles": {},
}
# 不生成已配置action的动作
for k, v in enhanced_info["action_methods"].items()
@@ -1340,5 +1340,5 @@ def setup_api_routes(app):
# 启动广播任务
@app.on_event("startup")
async def startup_event():
asyncio.create_task(broadcast_device_status())
asyncio.create_task(broadcast_status_page_data())
asyncio.create_task(broadcast_device_status(), name="web-api-startup-device")
asyncio.create_task(broadcast_status_page_data(), name="web-api-startup-status")

View File

@@ -3,11 +3,13 @@ HTTP客户端模块
提供与远程服务器通信的客户端功能只有host需要用
"""
import gzip
import json
import os
from typing import List, Dict, Any, Optional
from unilabos.utils.tools import fast_dumps as _fast_dumps, fast_dumps_pretty as _fast_dumps_pretty
import requests
from unilabos.resources.resource_tracker import ResourceTreeSet
from unilabos.utils.log import info
@@ -280,22 +282,54 @@ class HTTPClient:
)
return response
def resource_registry(self, registry_data: Dict[str, Any] | List[Dict[str, Any]]) -> requests.Response:
def resource_registry(
self, registry_data: Dict[str, Any] | List[Dict[str, Any]], tag: str = "registry",
) -> requests.Response:
"""
注册资源到服务器
注册资源到服务器,同步保存请求/响应到 unilabos_data
Args:
registry_data: 注册表数据,格式为 {resource_id: resource_info} / [{resource_info}]
tag: 保存文件的标签后缀 (如 "device_registry" / "resource_registry")
Returns:
Response: API响应对象
"""
# 序列化一次,同时用于保存和发送
json_bytes = _fast_dumps(registry_data)
# 保存请求数据到 unilabos_data
req_path = os.path.join(BasicConfig.working_dir, f"req_{tag}_upload.json")
try:
os.makedirs(BasicConfig.working_dir, exist_ok=True)
with open(req_path, "wb") as f:
f.write(_fast_dumps_pretty(registry_data))
logger.trace(f"注册表请求数据已保存: {req_path}")
except Exception as e:
logger.warning(f"保存注册表请求数据失败: {e}")
compressed_body = gzip.compress(json_bytes)
headers = {
"Authorization": f"Lab {self.auth}",
"Content-Type": "application/json",
"Content-Encoding": "gzip",
}
response = requests.post(
f"{self.remote_addr}/lab/resource",
json=registry_data,
headers={"Authorization": f"Lab {self.auth}"},
data=compressed_body,
headers=headers,
timeout=30,
)
# 保存响应数据到 unilabos_data
res_path = os.path.join(BasicConfig.working_dir, f"res_{tag}_upload.json")
try:
with open(res_path, "w", encoding="utf-8") as f:
f.write(f"{response.status_code}\n{response.text}")
logger.trace(f"注册表响应数据已保存: {res_path}")
except Exception as e:
logger.warning(f"保存注册表响应数据失败: {e}")
if response.status_code not in [200, 201]:
logger.error(f"注册资源失败: {response.status_code}, {response.text}")
if response.status_code == 200:

View File

@@ -86,7 +86,7 @@ def setup_server() -> FastAPI:
# 设置页面路由
try:
setup_web_pages(pages)
info("[Web] 已加载Web UI模块")
# info("[Web] 已加载Web UI模块")
except ImportError as e:
info(f"[Web] 未找到Web页面模块: {str(e)}")
except Exception as e:
@@ -138,7 +138,7 @@ def start_server(host: str = "0.0.0.0", port: int = 8002, open_browser: bool = T
server_thread = threading.Thread(target=server.run, daemon=True, name="uvicorn_server")
server_thread.start()
info("[Web] Server started, monitoring for restart requests...")
# info("[Web] Server started, monitoring for restart requests...")
# 监控重启标志
import unilabos.app.main as main_module

View File

@@ -23,9 +23,10 @@ from typing import Optional, Dict, Any, List
from urllib.parse import urlparse
from enum import Enum
from jedi.inference.gradual.typing import TypedDict
from typing_extensions import TypedDict
from unilabos.app.model import JobAddReq
from unilabos.resources.resource_tracker import ResourceDictType
from unilabos.ros.nodes.presets.host_node import HostNode
from unilabos.utils.type_check import serialize_result_info
from unilabos.app.communication import BaseCommunicationClient
@@ -164,7 +165,7 @@ class DeviceActionManager:
job_info.set_ready_timeout(10) # 设置10秒超时
self.active_jobs[device_key] = job_info
job_log = format_job_log(job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name)
logger.info(f"[DeviceActionManager] Job {job_log} can start immediately for {device_key}")
logger.trace(f"[DeviceActionManager] Job {job_log} can start immediately for {device_key}")
return True
def start_job(self, job_id: str) -> bool:
@@ -231,8 +232,9 @@ class DeviceActionManager:
job_info.update_timestamp()
# 从all_jobs中移除已结束的job
del self.all_jobs[job_id]
job_log = format_job_log(job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name)
logger.info(f"[DeviceActionManager] Job {job_log} ended for {device_key}")
# job_log = format_job_log(job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name)
# logger.debug(f"[DeviceActionManager] Job {job_log} ended for {device_key}")
pass
else:
job_log = format_job_log(job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name)
logger.warning(f"[DeviceActionManager] Job {job_log} was not active for {device_key}")
@@ -248,7 +250,7 @@ class DeviceActionManager:
next_job_log = format_job_log(
next_job.job_id, next_job.task_id, next_job.device_id, next_job.action_name
)
logger.info(f"[DeviceActionManager] Next job {next_job_log} can start for {device_key}")
logger.trace(f"[DeviceActionManager] Next job {next_job_log} can start for {device_key}")
return next_job
return None
@@ -302,7 +304,7 @@ class DeviceActionManager:
# 从all_jobs中移除
del self.all_jobs[job_id]
job_log = format_job_log(job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name)
logger.info(f"[DeviceActionManager] Active job {job_log} cancelled for {device_key}")
logger.trace(f"[DeviceActionManager] Active job {job_log} cancelled for {device_key}")
# 启动下一个任务
if device_key in self.device_queues and self.device_queues[device_key]:
@@ -315,7 +317,7 @@ class DeviceActionManager:
next_job_log = format_job_log(
next_job.job_id, next_job.task_id, next_job.device_id, next_job.action_name
)
logger.info(f"[DeviceActionManager] Next job {next_job_log} can start after cancel")
logger.trace(f"[DeviceActionManager] Next job {next_job_log} can start after cancel")
return True
# 如果是排队中的任务
@@ -329,7 +331,7 @@ class DeviceActionManager:
job_log = format_job_log(
job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name
)
logger.info(f"[DeviceActionManager] Queued job {job_log} cancelled for {device_key}")
logger.trace(f"[DeviceActionManager] Queued job {job_log} cancelled for {device_key}")
return True
job_log = format_job_log(job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name)
@@ -407,6 +409,7 @@ class MessageProcessor:
# 线程控制
self.is_running = False
self.thread = None
self._loop = None # asyncio event loop引用用于外部关闭websocket
self.reconnect_count = 0
logger.info(f"[MessageProcessor] Initialized for URL: {websocket_url}")
@@ -433,22 +436,31 @@ class MessageProcessor:
def stop(self) -> None:
"""停止消息处理线程"""
self.is_running = False
# 主动关闭websocket以快速中断消息接收循环
ws = self.websocket
loop = self._loop
if ws and loop and loop.is_running():
try:
asyncio.run_coroutine_threadsafe(ws.close(), loop)
except Exception:
pass
if self.thread and self.thread.is_alive():
self.thread.join(timeout=2)
logger.info("[MessageProcessor] Stopped")
def _run(self):
"""运行消息处理主循环"""
loop = asyncio.new_event_loop()
self._loop = asyncio.new_event_loop()
try:
asyncio.set_event_loop(loop)
loop.run_until_complete(self._connection_handler())
asyncio.set_event_loop(self._loop)
self._loop.run_until_complete(self._connection_handler())
except Exception as e:
logger.error(f"[MessageProcessor] Thread error: {str(e)}")
logger.error(traceback.format_exc())
finally:
if loop:
loop.close()
if self._loop:
self._loop.close()
self._loop = None
async def _connection_handler(self):
"""处理WebSocket连接和重连逻辑"""
@@ -465,8 +477,10 @@ class MessageProcessor:
async with websockets.connect(
self.websocket_url,
ssl=ssl_context,
open_timeout=20,
ping_interval=WSConfig.ping_interval,
ping_timeout=10,
close_timeout=5,
additional_headers={
"Authorization": f"Lab {BasicConfig.auth_secret()}",
"EdgeSession": f"{self.session_id}",
@@ -477,77 +491,94 @@ class MessageProcessor:
self.connected = True
self.reconnect_count = 0
logger.trace(f"[MessageProcessor] Connected to {self.websocket_url}")
logger.info(f"[MessageProcessor] 已连接到 {self.websocket_url}")
# 启动发送协程
send_task = asyncio.create_task(self._send_handler())
send_task = asyncio.create_task(self._send_handler(), name="websocket-send_task")
# 每次连接(含重连)后重新向服务端注册,
# 否则服务端不知道客户端已上线,不会推送消息。
if self.websocket_client:
self.websocket_client.publish_host_ready()
try:
# 接收消息循环
await self._message_handler()
finally:
# 必须在 async with __aexit__ 之前停止 send_task
# 否则 send_task 会在关闭握手期间继续发送数据,
# 干扰 websockets 库的内部清理,导致 task 泄漏。
self.connected = False
send_task.cancel()
try:
await send_task
except asyncio.CancelledError:
pass
self.connected = False
except websockets.exceptions.ConnectionClosed:
logger.warning("[MessageProcessor] Connection closed")
self.connected = False
logger.warning("[MessageProcessor] 与服务端连接中断")
except TimeoutError:
logger.warning(
f"[MessageProcessor] 与服务端连接通信超时 (已尝试 {self.reconnect_count + 1} 次),请检查您的网络状况"
)
except websockets.exceptions.InvalidStatus as e:
logger.warning(
f"[MessageProcessor] 收到服务端注册码 {e.response.status_code}, 上一进程可能还未退出"
)
except Exception as e:
logger.error(f"[MessageProcessor] Connection error: {str(e)}")
logger.error(traceback.format_exc())
self.connected = False
logger.error(f"[MessageProcessor] 尝试重连时出错 {str(e)}")
finally:
self.connected = False
self.websocket = None
# 重连逻辑
if self.is_running and self.reconnect_count < WSConfig.max_reconnect_attempts:
if not self.is_running:
break
if self.reconnect_count < WSConfig.max_reconnect_attempts:
self.reconnect_count += 1
backoff = WSConfig.reconnect_interval
logger.info(
f"[MessageProcessor] Reconnecting in {WSConfig.reconnect_interval}s "
f"(attempt {self.reconnect_count}/{WSConfig.max_reconnect_attempts})"
f"[MessageProcessor] 即将在 {backoff} 秒后重连 (已尝试 {self.reconnect_count}/{WSConfig.max_reconnect_attempts})"
)
await asyncio.sleep(WSConfig.reconnect_interval)
elif self.reconnect_count >= WSConfig.max_reconnect_attempts:
await asyncio.sleep(backoff)
else:
logger.error("[MessageProcessor] Max reconnection attempts reached")
break
else:
self.reconnect_count -= 1
async def _message_handler(self):
"""处理接收到的消息"""
"""处理接收到的消息
ConnectionClosed 不在此处捕获,让其向上传播到 _connection_handler
以便 async with websockets.connect() 的 __aexit__ 能感知连接已断,
正确清理内部 task避免 task 泄漏。
"""
if not self.websocket:
logger.error("[MessageProcessor] WebSocket connection is None")
return
try:
async for message in self.websocket:
try:
data = json.loads(message)
message_type = data.get("action", "")
message_data = data.get("data")
if self.session_id and self.session_id == data.get("edge_session"):
await self._process_message(message_type, message_data)
async for message in self.websocket:
try:
data = json.loads(message)
message_type = data.get("action", "")
message_data = data.get("data")
if self.session_id and self.session_id == data.get("edge_session"):
await self._process_message(message_type, message_data)
else:
if message_type.endswith("_material"):
logger.trace(
f"[MessageProcessor] 收到一条归属 {data.get('edge_session')} 的旧消息:{data}"
)
logger.debug(
f"[MessageProcessor] 跳过了一条归属 {data.get('edge_session')} 的旧消息: {data.get('action')}"
)
else:
if message_type.endswith("_material"):
logger.trace(f"[MessageProcessor] 收到一条归属 {data.get('edge_session')} 的旧消息:{data}")
logger.debug(f"[MessageProcessor] 跳过了一条归属 {data.get('edge_session')} 的旧消息: {data.get('action')}")
else:
await self._process_message(message_type, message_data)
except json.JSONDecodeError:
logger.error(f"[MessageProcessor] Invalid JSON received: {message}")
except Exception as e:
logger.error(f"[MessageProcessor] Error processing message: {str(e)}")
logger.error(traceback.format_exc())
except websockets.exceptions.ConnectionClosed:
logger.info("[MessageProcessor] Message handler stopped - connection closed")
except Exception as e:
logger.error(f"[MessageProcessor] Message handler error: {str(e)}")
logger.error(traceback.format_exc())
await self._process_message(message_type, message_data)
except json.JSONDecodeError:
logger.error(f"[MessageProcessor] Invalid JSON received: {message}")
except Exception as e:
logger.error(f"[MessageProcessor] Error processing message: {str(e)}")
logger.error(traceback.format_exc())
async def _send_handler(self):
"""处理发送队列中的消息"""
@@ -596,6 +627,7 @@ class MessageProcessor:
except asyncio.CancelledError:
logger.debug("[MessageProcessor] Send handler cancelled")
raise
except Exception as e:
logger.error(f"[MessageProcessor] Fatal error in send handler: {str(e)}")
logger.error(traceback.format_exc())
@@ -604,7 +636,7 @@ class MessageProcessor:
async def _process_message(self, message_type: str, message_data: Dict[str, Any]):
"""处理收到的消息"""
logger.debug(f"[MessageProcessor] Processing message: {message_type}")
logger.trace(f"[MessageProcessor] Processing message: {message_type}")
try:
if message_type == "pong":
@@ -627,6 +659,10 @@ class MessageProcessor:
# elif message_type == "session_id":
# self.session_id = message_data.get("session_id")
# logger.info(f"[MessageProcessor] Session ID: {self.session_id}")
elif message_type == "add_device":
await self._handle_device_manage(message_data, "add")
elif message_type == "remove_device":
await self._handle_device_manage(message_data, "remove")
elif message_type == "request_restart":
await self._handle_request_restart(message_data)
else:
@@ -698,13 +734,13 @@ class MessageProcessor:
await self._send_action_state_response(
device_id, action_name, task_id, job_id, "query_action_status", True, 0
)
logger.info(f"[MessageProcessor] Job {job_log} can start immediately")
logger.trace(f"[MessageProcessor] Job {job_log} can start immediately")
else:
# 需要排队
await self._send_action_state_response(
device_id, action_name, task_id, job_id, "query_action_status", False, 10
)
logger.info(f"[MessageProcessor] Job {job_log} queued")
logger.trace(f"[MessageProcessor] Job {job_log} queued")
# 通知QueueProcessor有新的队列更新
if self.queue_processor:
@@ -718,6 +754,32 @@ class MessageProcessor:
req = JobAddReq(**data)
job_log = format_job_log(req.job_id, req.task_id, req.device_id, req.action)
# 服务端对always_free动作可能跳过query_action_state直接发job_start
# 此时job尚未注册需要自动补注册
existing_job = self.device_manager.get_job_info(req.job_id)
if not existing_job:
action_name = req.action
device_action_key = f"/devices/{req.device_id}/{action_name}"
action_always_free = self._check_action_always_free(req.device_id, action_name)
if action_always_free:
job_info = JobInfo(
job_id=req.job_id,
task_id=req.task_id,
device_id=req.device_id,
action_name=action_name,
device_action_key=device_action_key,
status=JobStatus.QUEUE,
start_time=time.time(),
always_free=True,
)
self.device_manager.add_queue_request(job_info)
logger.info(f"[MessageProcessor] Job {job_log} always_free, auto-registered from direct job_start")
else:
logger.error(f"[MessageProcessor] Job {job_log} not registered (missing query_action_state)")
return
success = self.device_manager.start_job(req.job_id)
if not success:
logger.error(f"[MessageProcessor] Failed to start job {job_log}")
@@ -911,9 +973,7 @@ class MessageProcessor:
device_action_groups[key_add] = []
device_action_groups[key_add].append(item["uuid"])
logger.info(
f"[资源同步] 跨站Transfer: {item['uuid'][:8]} from {device_old_id} to {device_id}"
)
logger.info(f"[资源同步] 跨站Transfer: {item['uuid'][:8]} from {device_old_id} to {device_id}")
else:
# 正常update
key = (device_id, "update")
@@ -927,7 +987,9 @@ class MessageProcessor:
device_action_groups[key] = []
device_action_groups[key].append(item["uuid"])
logger.trace(f"[资源同步] 动作 {action} 分组数量: {len(device_action_groups)}, 总数量: {len(resource_uuid_list)}")
logger.trace(
f"[资源同步] 动作 {action} 分组数量: {len(device_action_groups)}, 总数量: {len(resource_uuid_list)}"
)
# 为每个(device_id, action)创建独立的更新线程
for (device_id, actual_action), items in device_action_groups.items():
@@ -963,45 +1025,77 @@ class MessageProcessor:
)
thread.start()
async def _handle_device_manage(self, device_list: list[ResourceDictType], action: str):
"""Handle add_device / remove_device from LabGo server."""
if not device_list:
return
for item in device_list:
target_node_id = item.get("target_node_id", "host_node")
def _notify(target_id: str, act: str, cfg: ResourceDictType):
try:
host_node = HostNode.get_instance(timeout=5)
if not host_node:
logger.error(f"[DeviceManage] HostNode not available for {act}_device")
return
success = host_node.notify_device_manage(target_id, act, cfg)
if success:
logger.info(f"[DeviceManage] {act}_device completed on {target_id}")
else:
logger.warning(f"[DeviceManage] {act}_device failed on {target_id}")
except Exception as e:
logger.error(f"[DeviceManage] Error in {act}_device: {e}")
logger.error(traceback.format_exc())
thread = threading.Thread(
target=_notify,
args=(target_node_id, action, item),
daemon=True,
name=f"DeviceManage-{action}-{item.get('id', '')}",
)
thread.start()
async def _handle_request_restart(self, data: Dict[str, Any]):
"""
处理重启请求
当LabGo发送request_restart时执行清理并触发重启
"""
reason = data.get("reason", "unknown")
delay = data.get("delay", 2) # 默认延迟2秒
logger.info(f"[MessageProcessor] Received restart request, reason: {reason}, delay: {delay}s")
# 发送确认消息
if self.websocket_client:
await self.websocket_client.send_message({
"action": "restart_acknowledged",
"data": {"reason": reason, "delay": delay}
})
self.send_message(
{"action": "restart_acknowledged", "data": {"reason": reason, "delay": delay}}
)
# 设置全局重启标志
import unilabos.app.main as main_module
main_module._restart_requested = True
main_module._restart_reason = reason
# 延迟后执行清理
await asyncio.sleep(delay)
# 在新线程中执行清理,避免阻塞当前事件循环
def do_cleanup():
import time
time.sleep(0.5) # 给当前消息处理完成的时间
logger.info(f"[MessageProcessor] Starting cleanup for restart, reason: {reason}")
try:
from unilabos.app.utils import cleanup_for_restart
if cleanup_for_restart():
logger.info("[MessageProcessor] Cleanup successful, main() will restart")
else:
logger.error("[MessageProcessor] Cleanup failed")
except Exception as e:
logger.error(f"[MessageProcessor] Error during cleanup: {e}")
cleanup_thread = threading.Thread(target=do_cleanup, name="RestartCleanupThread", daemon=True)
cleanup_thread.start()
logger.info(f"[MessageProcessor] Restart cleanup scheduled")
@@ -1077,6 +1171,7 @@ class QueueProcessor:
def stop(self) -> None:
"""停止队列处理线程"""
self.is_running = False
self.queue_update_event.set() # 立即唤醒等待中的线程
if self.thread and self.thread.is_alive():
self.thread.join(timeout=2)
logger.info("[QueueProcessor] Stopped")
@@ -1197,7 +1292,7 @@ class QueueProcessor:
success = self.message_processor.send_message(message)
job_log = format_job_log(job_info.job_id, job_info.task_id, job_info.device_id, job_info.action_name)
if success:
logger.debug(f"[QueueProcessor] Sent busy/need_more for queued job {job_log}")
logger.trace(f"[QueueProcessor] Sent busy/need_more for queued job {job_log}")
else:
logger.warning(f"[QueueProcessor] Failed to send busy status for job {job_log}")
@@ -1220,7 +1315,7 @@ class QueueProcessor:
job_info.action_name,
)
logger.info(f"[QueueProcessor] Job {job_log} completed with status: {status}")
logger.trace(f"[QueueProcessor] Job {job_log} completed with status: {status}")
# 结束任务,获取下一个可执行的任务
next_job = self.device_manager.end_job(job_id)
@@ -1240,8 +1335,8 @@ class QueueProcessor:
},
}
self.message_processor.send_message(message)
next_job_log = format_job_log(next_job.job_id, next_job.task_id, next_job.device_id, next_job.action_name)
logger.info(f"[QueueProcessor] Notified next job {next_job_log} can start")
# next_job_log = format_job_log(next_job.job_id, next_job.task_id, next_job.device_id, next_job.action_name)
# logger.debug(f"[QueueProcessor] Notified next job {next_job_log} can start")
# 立即触发下一轮状态检查
self.notify_queue_update()
@@ -1330,8 +1425,8 @@ class WebSocketClient(BaseCommunicationClient):
message = {"action": "normal_exit", "data": {"session_id": session_id}}
self.message_processor.send_message(message)
logger.info(f"[WebSocketClient] Sent normal_exit message with session_id: {session_id}")
# 给一点时间让消息发送出去
time.sleep(1)
# send_handler 每100ms检查一次队列等300ms足以让消息发
time.sleep(0.3)
except Exception as e:
logger.warning(f"[WebSocketClient] Failed to send normal_exit message: {str(e)}")
@@ -1383,7 +1478,7 @@ class WebSocketClient(BaseCommunicationClient):
except (KeyError, AttributeError):
logger.warning(f"[WebSocketClient] Failed to remove job {item.job_id} from HostNode status")
logger.info(f"[WebSocketClient] Intercepting final status for job_id: {item.job_id} - {status}")
# logger.debug(f"[WebSocketClient] Intercepting final status for job_id: {item.job_id} - {status}")
# 通知队列处理器job完成包括timeout的job
self.queue_processor.handle_job_completed(item.job_id, status)
@@ -1444,15 +1539,17 @@ class WebSocketClient(BaseCommunicationClient):
# 收集设备信息
devices = []
machine_name = BasicConfig.machine_name
try:
host_node = HostNode.get_instance(0)
if host_node:
# 获取设备信息
for device_id, namespace in host_node.devices_names.items():
device_key = f"{namespace}/{device_id}" if namespace.startswith("/") else f"/{namespace}/{device_id}"
device_key = (
f"{namespace}/{device_id}" if namespace.startswith("/") else f"/{namespace}/{device_id}"
)
is_online = device_key in host_node._online_devices
# 获取设备的动作信息
actions = {}
for action_id, client in host_node._action_clients.items():
@@ -1463,16 +1560,18 @@ class WebSocketClient(BaseCommunicationClient):
"action_path": action_id,
"action_type": str(type(client).__name__),
}
devices.append({
"device_id": device_id,
"namespace": namespace,
"device_key": device_key,
"is_online": is_online,
"machine_name": host_node.device_machine_names.get(device_id, machine_name),
"actions": actions,
})
devices.append(
{
"device_id": device_id,
"namespace": namespace,
"device_key": device_key,
"is_online": is_online,
"machine_name": host_node.device_machine_names.get(device_id, machine_name),
"actions": actions,
}
)
logger.info(f"[WebSocketClient] Collected {len(devices)} devices for host_ready")
except Exception as e:
logger.warning(f"[WebSocketClient] Error collecting device info: {e}")