266 lines
11 KiB
Python
266 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# 铸渊Agent · 自主守护进程
|
||
# HLDP://zhuyuan-agent/agent
|
||
#
|
||
# 运行在3090 GPU服务器上，心跳唤醒，推送到主服务器仪表盘。
|
||
# 冰朔离开WorkBuddy后，通过 guanghulab.com/console/ 看实时进度。
|
||
#
|
||
# 使用: python3 agent.py [--config config.json]
|
||
# PM2: pm2 start agent.py --name zhuyuan-agent --interpreter python3
|
||
|
||
import os
|
||
import sys
|
||
import json
|
||
import time
|
||
import signal
|
||
import traceback
|
||
from datetime import datetime
|
||
|
||
from gpu_monitor import collect_gpu_metrics, gpu_summary
|
||
from log_pusher import LogPusher
|
||
from heartbeat import Heartbeat
|
||
from training_runner import TrainingRunner
|
||
|
||
# 配置路径
|
||
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")
|
||
|
||
# 全局状态
|
||
running = True
|
||
current_training = None # 当前训练进程信息
|
||
|
||
|
||
def load_config() -> dict:
|
||
"""加载配置"""
|
||
config_path = CONFIG_PATH
|
||
for arg in sys.argv[1:]:
|
||
if arg.startswith("--config="):
|
||
config_path = arg.split("=", 1)[1]
|
||
|
||
if not os.path.exists(config_path):
|
||
print(f"[铸渊Agent] 配置文件不存在: {config_path}")
|
||
print("[铸渊Agent] 请先设置 config.json 中的 api_key")
|
||
sys.exit(1)
|
||
|
||
with open(config_path, "r") as f:
|
||
config = json.load(f)
|
||
|
||
# 检查API key
|
||
if config.get("api_key") == "__FROM_KEY_DELIVERY__" or not config.get("api_key"):
|
||
# 尝试从环境变量读取
|
||
env_key = os.environ.get("ZHUYUAN_API_KEY", "")
|
||
if env_key:
|
||
config["api_key"] = env_key
|
||
else:
|
||
print("[铸渊Agent] ⚠️ 未配置API Key!")
|
||
print("[铸渊Agent] 请在 guanghulab.com/console/ 密钥投递面板设置")
|
||
print("[铸渊Agent] 然后将API Key写入 config.json 的 api_key 字段")
|
||
print("[铸渊Agent] 或者设置环境变量 ZHUYUAN_API_KEY")
|
||
|
||
return config
|
||
|
||
|
||
def handle_signal(signum, frame):
|
||
"""处理退出信号"""
|
||
global running
|
||
print(f"\n[铸渊Agent] 收到信号 {signum}，优雅退出...")
|
||
running = False
|
||
|
||
|
||
def main():
|
||
global running, current_training
|
||
|
||
print("=" * 60)
|
||
print(" 铸渊Agent · ICE-GL-ZY001 · 自主守护进程")
|
||
print(" 曜冥纪元 · HoloLake Era · AGE v1.0")
|
||
print("=" * 60)
|
||
|
||
# 注册信号处理
|
||
signal.signal(signal.SIGINT, handle_signal)
|
||
signal.signal(signal.SIGTERM, handle_signal)
|
||
|
||
# 加载配置
|
||
config = load_config()
|
||
hostname = config.get("hostname", "3090-server")
|
||
poll_interval = config.get("poll_interval_seconds", 30)
|
||
|
||
# 初始化模块
|
||
pusher = LogPusher(
|
||
base_url=config["main_server"],
|
||
api_key=config.get("api_key", ""),
|
||
hostname=hostname
|
||
)
|
||
heartbeat = Heartbeat(
|
||
repo_path=config.get("brain_repo_path", "/data/guanghulab"),
|
||
brain_path=config.get("brain_path", "/data/guanghulab/brain")
|
||
)
|
||
|
||
# 检查是否有API key
|
||
if not config.get("api_key"):
|
||
print("[铸渊Agent] 无API Key，仅本地监控模式（不上报到仪表盘）")
|
||
print("[铸渊Agent] GPU指标将仅输出到终端")
|
||
|
||
has_key = bool(config.get("api_key"))
|
||
|
||
# 启动日记
|
||
if has_key:
|
||
pusher.push_diary("checkpoint", "铸渊Agent启动", f"主机: {hostname}, 轮询间隔: {poll_interval}s")
|
||
pusher.push_log("info", f"铸渊Agent v1.0 启动 · 主机: {hostname}")
|
||
|
||
print(f"[铸渊Agent] 主机: {hostname}")
|
||
print(f"[铸渊Agent] 主服务器: {config['main_server']}")
|
||
print(f"[铸渊Agent] 轮询间隔: {poll_interval}s")
|
||
print(f"[铸渊Agent] 上报仪表盘: {'是' if has_key else '否（仅本地）'}")
|
||
print(f"[铸渊Agent] 开始守护循环...")
|
||
print()
|
||
|
||
cycle = 0
|
||
|
||
while running:
|
||
cycle += 1
|
||
cycle_start = time.time()
|
||
|
||
try:
|
||
# ── 1. 心跳唤醒 ──
|
||
brain_status = heartbeat.check_brain()
|
||
|
||
if brain_status["has_task"]:
|
||
task = brain_status["task_details"]
|
||
task_type = brain_status["task_type"]
|
||
print(f"[心跳 #{cycle}] 发现任务: {task_type} — {task.get('name', task.get('title', '未命名'))}")
|
||
if has_key:
|
||
pusher.push_diary("decision", f"发现新任务: {task_type}",
|
||
json.dumps(task, ensure_ascii=False)[:200])
|
||
else:
|
||
print(f"[心跳 #{cycle}] {heartbeat.get_wake_summary()}")
|
||
|
||
# ── 2. GPU监控 ──
|
||
gpu_data = collect_gpu_metrics()
|
||
|
||
if gpu_data["gpus"]:
|
||
summary = gpu_summary(gpu_data["gpus"])
|
||
print(f"[GPU #{cycle}] {summary}")
|
||
|
||
if has_key:
|
||
ok = pusher.push_gpu(gpu_data)
|
||
if not ok:
|
||
print(f"[GPU #{cycle}] ⚠️ 推送上仪表盘失败")
|
||
elif gpu_data.get("error"):
|
||
print(f"[GPU #{cycle}] ⚠️ {gpu_data['error']}")
|
||
if has_key:
|
||
pusher.push_log("warn", f"GPU监控异常: {gpu_data['error']}")
|
||
else:
|
||
print(f"[GPU #{cycle}] 未检测到GPU")
|
||
|
||
# ── 3. 训练状态检查/执行 ──
|
||
if current_training is not None:
|
||
# 检查训练进程状态
|
||
if current_training.get("status") == "running":
|
||
# 训练正在运行中（由 training_runner 自主上报进度）
|
||
pass
|
||
elif current_training.get("status") == "done":
|
||
if has_key:
|
||
pusher.push_diary("checkpoint", "训练任务完成",
|
||
f"结果: {json.dumps(current_training, ensure_ascii=False)[:200]}")
|
||
pusher.push_log("success", "训练任务完成")
|
||
heartbeat.mark_task_done(brain_status.get("brain_file", ""))
|
||
current_training = None
|
||
elif current_training.get("status") == "error":
|
||
if has_key:
|
||
pusher.push_diary("error", "训练任务失败",
|
||
current_training.get("message", "未知错误"))
|
||
pusher.push_log("error", f"训练失败: {current_training.get('message', '')}")
|
||
current_training = None
|
||
|
||
elif brain_status["has_task"] and brain_status["task_type"] == "training":
|
||
# 启动新训练
|
||
task = brain_status["task_details"]
|
||
print(f"[Agent #{cycle}] 启动训练任务: {task.get('name', 'HLDP训练')}")
|
||
|
||
if has_key:
|
||
pusher.push_diary("decision", f"开始HLDP训练",
|
||
f"模型: {config['training'].get('model_name')}, "
|
||
f"语料: {config['training'].get('corpus_dir')}")
|
||
pusher.push_log("info", f"启动训练: {task.get('name', 'HLDP-3B')}")
|
||
|
||
current_training = {
|
||
"status": "starting",
|
||
"task": task,
|
||
"started_at": datetime.now().isoformat()
|
||
}
|
||
|
||
# 异步启动训练（简化版：同步执行）
|
||
try:
|
||
runner = TrainingRunner(
|
||
config=config.get("training", {}),
|
||
)
|
||
|
||
# 定义进度回调
|
||
def on_progress(step, loss, total_steps, extra_info):
|
||
global current_training
|
||
progress_pct = step / total_steps * 100 if total_steps > 0 else 0
|
||
print(f"[训练 #{cycle}] Step {step}/{total_steps} ({progress_pct:.0f}%) | Loss: {loss:.4f}")
|
||
|
||
if has_key:
|
||
pusher.push_training({
|
||
"job_id": "hldp-3b-test",
|
||
"status": "running",
|
||
"step": step,
|
||
"total_steps": total_steps,
|
||
"loss": loss,
|
||
"loss_history": extra_info.get("loss_history", []),
|
||
"eta_seconds": extra_info.get("eta_seconds"),
|
||
"elapsed_seconds": extra_info.get("elapsed_seconds", 0),
|
||
"learning_rate": extra_info.get("learning_rate"),
|
||
"model_name": config["training"].get("model_name", ""),
|
||
"message": f"HLDP原生格式训练 · {step}/{total_steps}步"
|
||
})
|
||
|
||
result = runner.train(
|
||
corpus_dir=config["training"].get("corpus_dir", "/data/corpus/notion-hldp"),
|
||
progress_callback=on_progress
|
||
)
|
||
|
||
current_training = result
|
||
|
||
except Exception as e:
|
||
current_training = {
|
||
"status": "error",
|
||
"message": f"训练异常: {str(e)}"
|
||
}
|
||
print(f"[Agent #{cycle}] 训练异常: {e}")
|
||
traceback.print_exc()
|
||
|
||
# ── 4. 空闲日志 ──
|
||
if has_key and cycle % 10 == 0:
|
||
# 每10个周期发送一次心跳日志
|
||
pusher.push_log("info", f"铸渊Agent守护中 · 周期#{cycle} · " +
|
||
(gpu_summary(gpu_data["gpus"]) if gpu_data["gpus"] else "无GPU"))
|
||
|
||
except Exception as e:
|
||
print(f"[Agent #{cycle}] 循环异常: {e}")
|
||
traceback.print_exc()
|
||
if has_key:
|
||
pusher.push_log("error", f"Agent循环异常 #{cycle}: {str(e)[:200]}")
|
||
|
||
# ── 5. 等待下一个周期 ──
|
||
elapsed = time.time() - cycle_start
|
||
sleep_time = max(0, poll_interval - elapsed)
|
||
|
||
if sleep_time > 0:
|
||
# 分段sleep以响应退出信号
|
||
for _ in range(int(sleep_time)):
|
||
if not running:
|
||
break
|
||
time.sleep(1)
|
||
|
||
# 退出清理
|
||
print("\n[铸渊Agent] 守护循环结束")
|
||
if has_key:
|
||
pusher.push_diary("checkpoint", "铸渊Agent停止", f"共运行 {cycle} 个周期")
|
||
pusher.push_log("warn", f"铸渊Agent停止 · 运行了{cycle}个周期")
|
||
print("[铸渊Agent] 再见。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|