85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
|
# 心跳/任务发现模块
|
||
|
|
# HLDP://zhuyuan-agent/heartbeat
|
||
|
|
|
||
|
|
import os
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
|
||
|
|
class Heartbeat:
|
||
|
|
"""心跳唤醒 + 任务发现"""
|
||
|
|
|
||
|
|
def __init__(self, repo_path: str = "/data/guanghulab", brain_path: str = "/data/guanghulab/brain"):
|
||
|
|
self.repo_path = repo_path
|
||
|
|
self.brain_path = brain_path
|
||
|
|
self.last_check = None
|
||
|
|
self.current_task = None
|
||
|
|
|
||
|
|
def check_brain(self) -> dict:
|
||
|
|
"""读取大脑文件，检查是否有新任务
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
{
|
||
|
|
"has_task": bool,
|
||
|
|
"task_type": "training" | "inference" | "deploy" | null,
|
||
|
|
"task_details": dict,
|
||
|
|
"brain_file": str # 触发任务的大脑文件
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
self.last_check = datetime.now().isoformat()
|
||
|
|
|
||
|
|
# 检查是否有任务标记文件
|
||
|
|
task_file = os.path.join(self.brain_path, "pending-tasks.json")
|
||
|
|
if os.path.exists(task_file):
|
||
|
|
try:
|
||
|
|
with open(task_file, "r") as f:
|
||
|
|
tasks = json.load(f)
|
||
|
|
if isinstance(tasks, list) and len(tasks) > 0:
|
||
|
|
task = tasks[0]
|
||
|
|
return {
|
||
|
|
"has_task": True,
|
||
|
|
"task_type": task.get("type", "unknown"),
|
||
|
|
"task_details": task,
|
||
|
|
"brain_file": "pending-tasks.json"
|
||
|
|
}
|
||
|
|
except (json.JSONDecodeError, FileNotFoundError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
# 检查是否有训练指令文件
|
||
|
|
train_file = os.path.join(self.brain_path, "train-now.json")
|
||
|
|
if os.path.exists(train_file):
|
||
|
|
try:
|
||
|
|
with open(train_file, "r") as f:
|
||
|
|
task = json.load(f)
|
||
|
|
return {
|
||
|
|
"has_task": True,
|
||
|
|
"task_type": "training",
|
||
|
|
"task_details": task,
|
||
|
|
"brain_file": "train-now.json"
|
||
|
|
}
|
||
|
|
except (json.JSONDecodeError, FileNotFoundError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
return {"has_task": False, "task_type": None, "task_details": {}, "brain_file": None}
|
||
|
|
|
||
|
|
def mark_task_done(self, task_file: str):
|
||
|
|
"""标记任务完成（删除或重命名任务文件）"""
|
||
|
|
filepath = os.path.join(self.brain_path, task_file)
|
||
|
|
if os.path.exists(filepath):
|
||
|
|
done_path = filepath + ".done." + datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
|
|
os.rename(filepath, done_path)
|
||
|
|
return done_path
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_wake_summary(self) -> str:
|
||
|
|
"""生成唤醒摘要"""
|
||
|
|
now = datetime.now().strftime("%H:%M:%S")
|
||
|
|
brain_files = []
|
||
|
|
if os.path.exists(self.brain_path):
|
||
|
|
try:
|
||
|
|
brain_files = sorted(os.listdir(self.brain_path))[:10]
|
||
|
|
except:
|
||
|
|
pass
|
||
|
|
return f"[{now}] 心跳唤醒 | 大脑文件: {len(brain_files)}个 | 上次检查: {self.last_check}"
|