150 lines
5.1 KiB
Python
150 lines
5.1 KiB
Python
|
"""
|
|||
|
|
AI Box Sophon Adapter 插件
|
|||
|
|
|
|||
|
|
基于 MQTT 的 AI Box 设备心跳适配器，实现：
|
|||
|
|
1. 通过 MQTT 主题 /board_ping 接收设备心跳上报
|
|||
|
|
2. 解析设备系统参数（温度、硬盘、内存、通道、任务、TPU 等）
|
|||
|
|
3. 持久化心跳数据到 SQLite 数据库
|
|||
|
|
4. 提供 HTTP API 查询设备在线状态和详情
|
|||
|
|
5. 后台定时检查离线设备，触发异常告警
|
|||
|
|
|
|||
|
|
心跳超时机制:
|
|||
|
|
- 设备按配置间隔（默认 5 秒）上报心跳
|
|||
|
|
- 服务端每 check_interval 秒检查一次
|
|||
|
|
- 超过 heartbeat_timeout 秒未收到心跳的设备标记为离线
|
|||
|
|
"""
|
|||
|
|
import asyncio
|
|||
|
|
import logging
|
|||
|
|
from datetime import datetime
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
from chongming_worker import auto_bind_models
|
|||
|
|
from chongming_worker.runtime import create_workers as _create_workers
|
|||
|
|
|
|||
|
|
from .router import *
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# 自动绑定 Pydantic 模型（用于 gen_models 生成的模型）
|
|||
|
|
auto_bind_models("aibox_sophon_adapter")
|
|||
|
|
|
|||
|
|
# 离线检查任务引用
|
|||
|
|
_offline_check_task: asyncio.Task | None = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def _check_offline_devices(interval: float):
|
|||
|
|
"""
|
|||
|
|
后台任务：定期检查离线设备
|
|||
|
|
|
|||
|
|
每隔 interval 秒扫描数据库，将超过 heartbeat_timeout 秒
|
|||
|
|
未收到心跳的设备标记为离线，并输出告警日志。
|
|||
|
|
"""
|
|||
|
|
from .global_var import db, heartbeat_timeout
|
|||
|
|
|
|||
|
|
while True:
|
|||
|
|
await asyncio.sleep(interval)
|
|||
|
|
try:
|
|||
|
|
if db is None:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
async with db.get_session() as session:
|
|||
|
|
from .models.orm import BoardHeartbeat
|
|||
|
|
from sqlalchemy import select
|
|||
|
|
|
|||
|
|
stmt = select(BoardHeartbeat).where(BoardHeartbeat.is_online == True)
|
|||
|
|
result = await session.execute(stmt)
|
|||
|
|
records = result.scalars().all()
|
|||
|
|
|
|||
|
|
now = datetime.now()
|
|||
|
|
offline_count = 0
|
|||
|
|
for record in records:
|
|||
|
|
elapsed = (now - record.last_ping_time).total_seconds()
|
|||
|
|
if elapsed > heartbeat_timeout:
|
|||
|
|
record.is_online = False
|
|||
|
|
record.updated_at = now
|
|||
|
|
offline_count += 1
|
|||
|
|
logger.warning(
|
|||
|
|
f"[Offline] 设备 {record.board_id} (IP: {record.board_ip}) "
|
|||
|
|
f"已离线 {elapsed:.0f} 秒（超时阈值: {heartbeat_timeout}s）"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if offline_count > 0:
|
|||
|
|
await session.commit()
|
|||
|
|
logger.info(f"[Offline] 本轮标记 {offline_count} 个设备为离线")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"[OfflineChecker] 检查异常: {e}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def create_workers(session, config, count=1, worker_manager=None):
|
|||
|
|
"""
|
|||
|
|
插件入口 — 初始化数据库和全局资源，创建 Worker 副本
|
|||
|
|
|
|||
|
|
由 FastAPI lifespan 在启动时调用。
|
|||
|
|
"""
|
|||
|
|
from .global_var import set_global_var, heartbeat_timeout as _default_timeout
|
|||
|
|
from chongming_database import Database
|
|||
|
|
|
|||
|
|
# 读取心跳配置
|
|||
|
|
worker_config = config.get("worker", {}).get("heartbeat", {}) if isinstance(config, dict) else {}
|
|||
|
|
hb_interval = worker_config.get("interval", 5)
|
|||
|
|
hb_timeout = worker_config.get("timeout", _default_timeout)
|
|||
|
|
check_interval = worker_config.get("check_interval", 10)
|
|||
|
|
|
|||
|
|
# 设定全局变量
|
|||
|
|
global heartbeat_timeout
|
|||
|
|
heartbeat_timeout = hb_timeout
|
|||
|
|
|
|||
|
|
# 初始化数据库
|
|||
|
|
data_dir = Path(config.get("data_dir", "./data")) if isinstance(config, dict) else Path("./data")
|
|||
|
|
db_path = data_dir / "storage" / "chongming" / "aibox_sophon_adapter.db"
|
|||
|
|
db_path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
db = Database(db_path)
|
|||
|
|
await db.initialize()
|
|||
|
|
|
|||
|
|
# 创建表
|
|||
|
|
from .models.orm import BoardHeartbeat
|
|||
|
|
async with db.get_session() as sess:
|
|||
|
|
from sqlmodel import SQLModel
|
|||
|
|
await sess.run_sync(lambda s: SQLModel.metadata.create_all(s.bind))
|
|||
|
|
|
|||
|
|
set_global_var("db", db)
|
|||
|
|
set_global_var("heartbeat_timeout", hb_timeout)
|
|||
|
|
set_global_var("heartbeat_interval", hb_interval)
|
|||
|
|
|
|||
|
|
logger.info(
|
|||
|
|
f"[aibox_sophon_adapter] 已初始化: "
|
|||
|
|
f"心跳间隔={hb_interval}s, 超时={hb_timeout}s, 检查间隔={check_interval}s"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 启动离线设备检查后台任务
|
|||
|
|
global _offline_check_task
|
|||
|
|
_offline_check_task = asyncio.create_task(_check_offline_devices(check_interval))
|
|||
|
|
|
|||
|
|
# 创建 Worker 副本
|
|||
|
|
return _create_workers(session, config, count, worker_manager)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def cleanup(worker_item: str):
|
|||
|
|
"""
|
|||
|
|
清理资源
|
|||
|
|
|
|||
|
|
在 Worker 关闭时调用，关闭数据库连接和后台任务。
|
|||
|
|
"""
|
|||
|
|
from .global_var import db
|
|||
|
|
|
|||
|
|
global _offline_check_task
|
|||
|
|
if _offline_check_task is not None and not _offline_check_task.done():
|
|||
|
|
_offline_check_task.cancel()
|
|||
|
|
try:
|
|||
|
|
await _offline_check_task
|
|||
|
|
except asyncio.CancelledError:
|
|||
|
|
pass
|
|||
|
|
logger.info("[aibox_sophon_adapter] 离线检查任务已停止")
|
|||
|
|
|
|||
|
|
if db is not None:
|
|||
|
|
await db.dispose()
|
|||
|
|
|
|||
|
|
logger.info(f"[aibox_sophon_adapter] {worker_item} 清理完毕")
|