zhuyuan-agent: 自主Agent守护进程+GPU监控+HLDP训练管道
This commit is contained in:
parent
070e2d8183
commit
fb3a6e0ba4
265
zhuyuan-agent/agent.py
Normal file
265
zhuyuan-agent/agent.py
Normal file
@ -0,0 +1,265 @@
|
||||
#!/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()
|
||||
26
zhuyuan-agent/config.json
Normal file
26
zhuyuan-agent/config.json
Normal file
@ -0,0 +1,26 @@
|
||||
{
|
||||
"agent_name": "铸渊Agent · ICE-GL-ZY001",
|
||||
"hostname": "3090-GPU-SERVER",
|
||||
"main_server": "https://guanghulab.com",
|
||||
"api_key": "__FROM_KEY_DELIVERY__",
|
||||
"poll_interval_seconds": 30,
|
||||
"brain_repo": "https://guanghulab.com/bingshuo/guanghulab.git",
|
||||
"brain_repo_path": "/data/guanghulab",
|
||||
"brain_path": "/data/guanghulab/brain",
|
||||
"training": {
|
||||
"model_name": "Qwen/Qwen2.5-3B",
|
||||
"output_dir": "/data/models/shuangyan-3b-hldp",
|
||||
"corpus_dir": "/data/corpus/notion-hldp",
|
||||
"lora_r": 16,
|
||||
"lora_alpha": 32,
|
||||
"learning_rate": 2e-4,
|
||||
"batch_size": 2,
|
||||
"gradient_accumulation": 4,
|
||||
"max_seq_length": 2048,
|
||||
"warmup_steps": 100,
|
||||
"save_steps": 50,
|
||||
"max_steps": 500,
|
||||
"use_4bit": true,
|
||||
"bnb_4bit_compute_dtype": "float16"
|
||||
}
|
||||
}
|
||||
99
zhuyuan-agent/gpu_monitor.py
Normal file
99
zhuyuan-agent/gpu_monitor.py
Normal file
@ -0,0 +1,99 @@
|
||||
# GPU监控模块 · 采集nvidia-smi数据
|
||||
# HLDP://zhuyuan-agent/gpu-monitor
|
||||
|
||||
import subprocess
|
||||
import json
|
||||
import re
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
|
||||
def collect_gpu_metrics() -> Dict:
|
||||
"""采集所有GPU的实时指标
|
||||
|
||||
Returns:
|
||||
{
|
||||
"gpus": [
|
||||
{
|
||||
"index": 0,
|
||||
"name": "NVIDIA GeForce RTX 3090",
|
||||
"uuid": "GPU-xxxx",
|
||||
"utilization_gpu": 85, # %
|
||||
"memory_used_mb": 18432, # MB
|
||||
"memory_total_mb": 24576, # MB
|
||||
"temperature_gpu": 72, # °C
|
||||
"power_draw_w": 285.5, # W
|
||||
"fan_speed": 65 # %
|
||||
}
|
||||
],
|
||||
"error": null
|
||||
}
|
||||
"""
|
||||
try:
|
||||
# 查询GPU关键指标
|
||||
result = subprocess.run(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=index,name,uuid,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw,fan.speed",
|
||||
"--format=csv,noheader,nounits"
|
||||
],
|
||||
capture_output=True, text=True, timeout=10
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
return {"gpus": [], "error": f"nvidia-smi failed: {result.stderr.strip()}"}
|
||||
|
||||
gpus = []
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
if not line.strip():
|
||||
continue
|
||||
parts = [p.strip() for p in line.split(",")]
|
||||
if len(parts) < 9:
|
||||
continue
|
||||
|
||||
try:
|
||||
gpu = {
|
||||
"index": int(parts[0]),
|
||||
"name": parts[1],
|
||||
"uuid": parts[2],
|
||||
"utilization_gpu": int(parts[3]) if parts[3] != "[Not Supported]" else 0,
|
||||
"memory_used_mb": int(parts[4]) if parts[4] != "[Not Supported]" else 0,
|
||||
"memory_total_mb": int(parts[5]) if parts[5] != "[Not Supported]" else 0,
|
||||
"temperature_gpu": int(parts[6]) if parts[6] != "[Not Supported]" else 0,
|
||||
"power_draw_w": float(parts[7]) if parts[7] not in ("[Not Supported]", "") else 0.0,
|
||||
"fan_speed": int(parts[8]) if parts[8] not in ("[Not Supported]", "") else 0,
|
||||
}
|
||||
gpus.append(gpu)
|
||||
except (ValueError, IndexError):
|
||||
continue
|
||||
|
||||
return {"gpus": gpus, "error": None}
|
||||
|
||||
except FileNotFoundError:
|
||||
return {"gpus": [], "error": "nvidia-smi not found - not a GPU machine?"}
|
||||
except subprocess.TimeoutExpired:
|
||||
return {"gpus": [], "error": "nvidia-smi timed out"}
|
||||
except Exception as e:
|
||||
return {"gpus": [], "error": str(e)}
|
||||
|
||||
|
||||
def gpu_summary(gpus: List[Dict]) -> str:
|
||||
"""生成GPU状态的一行摘要"""
|
||||
if not gpus:
|
||||
return "无GPU"
|
||||
parts = []
|
||||
for g in gpus:
|
||||
util = g.get("utilization_gpu", 0)
|
||||
temp = g.get("temperature_gpu", 0)
|
||||
mem = g.get("memory_used_mb", 0)
|
||||
mem_total = g.get("memory_total_mb", 0)
|
||||
mem_pct = int(mem / mem_total * 100) if mem_total > 0 else 0
|
||||
parts.append(f"GPU{g['index']}: {util}%/{temp}°C/{mem_pct}%VRAM")
|
||||
return " | ".join(parts)
|
||||
|
||||
|
||||
# 快速测试
|
||||
if __name__ == "__main__":
|
||||
data = collect_gpu_metrics()
|
||||
print(json.dumps(data, indent=2, ensure_ascii=False))
|
||||
if data["gpus"]:
|
||||
print(f"\n摘要: {gpu_summary(data['gpus'])}")
|
||||
84
zhuyuan-agent/heartbeat.py
Normal file
84
zhuyuan-agent/heartbeat.py
Normal file
@ -0,0 +1,84 @@
|
||||
# 心跳/任务发现模块
|
||||
# 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}"
|
||||
80
zhuyuan-agent/log_pusher.py
Normal file
80
zhuyuan-agent/log_pusher.py
Normal file
@ -0,0 +1,80 @@
|
||||
# 日志推送模块 · HTTP POST到主服务器
|
||||
# HLDP://zhuyuan-agent/log-pusher
|
||||
|
||||
import json
|
||||
import urllib.request
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class LogPusher:
|
||||
"""向主服务器推送操作日志、日记、GPU指标、训练进度"""
|
||||
|
||||
def __init__(self, base_url: str, api_key: str, hostname: str = "3090-server"):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.hostname = hostname
|
||||
self._headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def _post(self, path: str, data: dict, timeout: int = 15) -> bool:
|
||||
"""POST请求到主服务器"""
|
||||
url = f"{self.base_url}{path}"
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
url,
|
||||
data=json.dumps(data).encode("utf-8"),
|
||||
headers=self._headers,
|
||||
method="POST"
|
||||
)
|
||||
resp = urllib.request.urlopen(req, timeout=timeout)
|
||||
result = json.loads(resp.read())
|
||||
return result.get("ok", False)
|
||||
except Exception as e:
|
||||
print(f"[PUSH ERROR] {path}: {e}")
|
||||
return False
|
||||
|
||||
def push_gpu(self, gpu_data: dict) -> bool:
|
||||
"""推送GPU指标"""
|
||||
data = {
|
||||
"hostname": self.hostname,
|
||||
"gpus": gpu_data.get("gpus", [])
|
||||
}
|
||||
return self._post("/api/gpu/status", data)
|
||||
|
||||
def push_training(self, training_data: dict) -> bool:
|
||||
"""推送训练进度"""
|
||||
return self._post("/api/training/status", training_data)
|
||||
|
||||
def push_log(self, level: str, message: str, category: str = "agent") -> bool:
|
||||
"""推送操作日志"""
|
||||
return self._post("/api/agent/log", {
|
||||
"level": level,
|
||||
"message": message,
|
||||
"category": category
|
||||
})
|
||||
|
||||
def push_diary(self, entry_type: str, title: str, description: str = "") -> bool:
|
||||
"""推送日记条目"""
|
||||
return self._post("/api/agent/diary", {
|
||||
"type": entry_type,
|
||||
"title": title,
|
||||
"description": description
|
||||
})
|
||||
|
||||
def log_info(self, msg: str):
|
||||
"""快捷:info级别日志"""
|
||||
self.push_log("info", msg)
|
||||
|
||||
def log_success(self, msg: str):
|
||||
"""快捷:success级别日志"""
|
||||
self.push_log("success", msg)
|
||||
|
||||
def log_warn(self, msg: str):
|
||||
"""快捷:warn级别日志"""
|
||||
self.push_log("warn", msg)
|
||||
|
||||
def log_error(self, msg: str):
|
||||
"""快捷:error级别日志"""
|
||||
self.push_log("error", msg)
|
||||
6
zhuyuan-agent/requirements.txt
Normal file
6
zhuyuan-agent/requirements.txt
Normal file
@ -0,0 +1,6 @@
|
||||
requests>=2.28.0
|
||||
torch>=2.0.0
|
||||
transformers>=4.38.0
|
||||
peft>=0.8.0
|
||||
accelerate>=0.27.0
|
||||
bitsandbytes>=0.41.0
|
||||
286
zhuyuan-agent/training_runner.py
Normal file
286
zhuyuan-agent/training_runner.py
Normal file
@ -0,0 +1,286 @@
|
||||
# HLDP原生格式训练执行器
|
||||
# HLDP://zhuyuan-agent/training-runner
|
||||
#
|
||||
# 核心验证:Notion原生页面格式(HLDP标记)直接训练,不转JSONL。
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import sys
|
||||
from typing import Optional, Callable
|
||||
|
||||
|
||||
# 特殊Token定义(HLDP结构标记)
|
||||
SPECIAL_TOKENS = [
|
||||
"[HLDP_PATH]", "[/HLDP_PATH]",
|
||||
"[PERSONA]", "[/PERSONA]",
|
||||
"[COGNITIVE_JUMP]", "[/COGNITIVE_JUMP]",
|
||||
"[CAUSAL_CHAIN]", "[/CAUSAL_CHAIN]",
|
||||
"[TITLE]", "[/TITLE]",
|
||||
"[CONTENT]", "[/CONTENT]",
|
||||
"[QUALITY_HIGH]", "[QUALITY_MEDIUM]",
|
||||
"[THINKING]", "[/THINKING]",
|
||||
"[HEADING_1]", "[/HEADING_1]",
|
||||
"[HEADING_2]", "[/HEADING_2]",
|
||||
"[HEADING_3]", "[/HEADING_3]",
|
||||
"[CODE_BLOCK]", "[/CODE_BLOCK]",
|
||||
"[QUOTE]", "[/QUOTE]",
|
||||
"[CALLOUT]", "[/CALLOUT]",
|
||||
"[LIST_ITEM]", "[/LIST_ITEM]",
|
||||
]
|
||||
|
||||
|
||||
class TrainingRunner:
|
||||
"""HLDP原生格式训练管道"""
|
||||
|
||||
def __init__(self, config: dict, progress_callback: Optional[Callable] = None):
|
||||
"""
|
||||
Args:
|
||||
config: 训练配置(来自config.json的training部分)
|
||||
progress_callback: 每步回调,接收 (step, loss, total_steps)
|
||||
"""
|
||||
self.config = config
|
||||
self.progress_callback = progress_callback
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.start_time = None
|
||||
self.loss_history = []
|
||||
|
||||
def prepare(self):
|
||||
"""准备训练环境:注册特殊token,加载模型"""
|
||||
print("[铸渊Agent] 准备HLDP训练环境...")
|
||||
|
||||
try:
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
||||
import torch
|
||||
|
||||
model_name = self.config.get("model_name", "Qwen/Qwen2.5-3B")
|
||||
use_4bit = self.config.get("use_4bit", True)
|
||||
|
||||
print(f"[铸渊Agent] 加载模型: {model_name}")
|
||||
|
||||
# 加载tokenizer并添加HLDP特殊token
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_name,
|
||||
trust_remote_code=True,
|
||||
padding_side="right"
|
||||
)
|
||||
|
||||
# 设置pad_token
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
# 添加HLDP特殊token到tokenizer
|
||||
num_added = self.tokenizer.add_tokens(SPECIAL_TOKENS)
|
||||
print(f"[铸渊Agent] 添加了 {num_added} 个HLDP特殊token到tokenizer")
|
||||
|
||||
# 加载模型(4bit量化以适配3090 24GB)
|
||||
load_kwargs = {
|
||||
"trust_remote_code": True,
|
||||
"torch_dtype": torch.float16,
|
||||
"device_map": "auto",
|
||||
}
|
||||
|
||||
if use_4bit:
|
||||
try:
|
||||
from transformers import BitsAndBytesConfig
|
||||
load_kwargs["quantization_config"] = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
print("[铸渊Agent] 使用4bit量化加载")
|
||||
except ImportError:
|
||||
print("[铸渊Agent] bitsandbytes不可用,使用float16")
|
||||
load_kwargs.pop("quantization_config", None)
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(model_name, **load_kwargs)
|
||||
|
||||
# 扩展embedding层以支持新token
|
||||
if num_added > 0:
|
||||
self.model.resize_token_embeddings(len(self.tokenizer))
|
||||
|
||||
# 准备LoRA
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
lora_config = LoraConfig(
|
||||
r=self.config.get("lora_r", 16),
|
||||
lora_alpha=self.config.get("lora_alpha", 32),
|
||||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
self.model.print_trainable_parameters()
|
||||
|
||||
print("[铸渊Agent] HLDP训练环境准备完成")
|
||||
return True
|
||||
|
||||
except ImportError as e:
|
||||
print(f"[铸渊Agent] 缺少依赖: {e}")
|
||||
print("请安装: pip install transformers peft accelerate bitsandbytes torch")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"[铸渊Agent] 准备失败: {e}")
|
||||
return False
|
||||
|
||||
def load_hldp_corpus(self, corpus_dir: str) -> list:
|
||||
"""加载HLDP格式语料(不转JSONL,保留原生HLDP标记)
|
||||
|
||||
Returns:
|
||||
list of str: 每条是一个HLDP标记的完整训练文本
|
||||
"""
|
||||
texts = []
|
||||
|
||||
if not os.path.exists(corpus_dir):
|
||||
print(f"[铸渊Agent] 语料目录不存在: {corpus_dir}")
|
||||
return texts
|
||||
|
||||
for root, dirs, files in os.walk(corpus_dir):
|
||||
for filename in files:
|
||||
filepath = os.path.join(root, filename)
|
||||
|
||||
if filename.endswith(".hldp"):
|
||||
# HLDP原生格式文件
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
texts.append(f.read())
|
||||
|
||||
elif filename.endswith(".md"):
|
||||
# Markdown文件 → 按HLDP结构包装
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
# 包裹HLDP标记
|
||||
wrapped = f"[HLDP_PATH]{filepath}[/HLDP_PATH]\n[CONTENT]\n{content}\n[/CONTENT]"
|
||||
texts.append(wrapped)
|
||||
|
||||
print(f"[铸渊Agent] 加载了 {len(texts)} 条HLDP语料")
|
||||
return texts
|
||||
|
||||
def train(self, corpus_dir: str, progress_callback: Optional[Callable] = None):
|
||||
"""执行HLDP原生格式训练"""
|
||||
if self.model is None:
|
||||
if not self.prepare():
|
||||
return {"status": "error", "message": "模型准备失败"}
|
||||
|
||||
texts = self.load_hldp_corpus(corpus_dir)
|
||||
if not texts:
|
||||
return {"status": "error", "message": "无语料数据"}
|
||||
|
||||
# 如果没有transformers Trainer,用simulated training输出结构
|
||||
# 实际训练在3090上运行时才加载完整transformers
|
||||
try:
|
||||
from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
max_steps = self.config.get("max_steps", 500)
|
||||
batch_size = self.config.get("batch_size", 2)
|
||||
grad_accum = self.config.get("gradient_accumulation", 4)
|
||||
lr = self.config.get("learning_rate", 2e-4)
|
||||
output_dir = self.config.get("output_dir", "/data/models/shuangyan-3b-hldp")
|
||||
max_seq_length = self.config.get("max_seq_length", 2048)
|
||||
|
||||
# 准备数据集
|
||||
class HLDPDataset(Dataset):
|
||||
def __init__(self, texts, tokenizer, max_length):
|
||||
self.encodings = tokenizer(
|
||||
texts, truncation=True, padding="max_length",
|
||||
max_length=max_length, return_tensors="pt"
|
||||
)
|
||||
|
||||
def __len__(self): return len(self.encodings["input_ids"])
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return {k: v[idx] for k, v in self.encodings.items()}
|
||||
|
||||
dataset = HLDPDataset(texts, self.tokenizer, max_seq_length)
|
||||
data_collator = DataCollatorForLanguageModeling(
|
||||
tokenizer=self.tokenizer, mlm=False
|
||||
)
|
||||
|
||||
# 训练参数
|
||||
training_args = TrainingArguments(
|
||||
output_dir=output_dir,
|
||||
per_device_train_batch_size=batch_size,
|
||||
gradient_accumulation_steps=grad_accum,
|
||||
learning_rate=lr,
|
||||
warmup_steps=self.config.get("warmup_steps", 100),
|
||||
max_steps=max_steps,
|
||||
logging_steps=5,
|
||||
save_steps=self.config.get("save_steps", 50),
|
||||
save_total_limit=3,
|
||||
fp16=True,
|
||||
report_to=[],
|
||||
dataloader_pin_memory=False,
|
||||
)
|
||||
|
||||
# 自定义callback用于进度上报
|
||||
class ProgressCallback:
|
||||
def __init__(self, runner, total_steps):
|
||||
self.runner = runner
|
||||
self.total_steps = total_steps
|
||||
self.current_step = 0
|
||||
self.start = time.time()
|
||||
|
||||
def on_log(self, args, state, control, logs=None, **kwargs):
|
||||
if logs and "loss" in logs:
|
||||
self.current_step = state.global_step
|
||||
loss = logs["loss"]
|
||||
self.runner.loss_history.append(loss)
|
||||
elapsed = time.time() - self.start
|
||||
|
||||
# 估算ETA
|
||||
if self.current_step > 0:
|
||||
eta = (elapsed / self.current_step) * (self.total_steps - self.current_step)
|
||||
else:
|
||||
eta = 0
|
||||
|
||||
# 回调上报进度
|
||||
cb = progress_callback or self.runner.progress_callback
|
||||
if cb:
|
||||
cb(self.current_step, loss, self.total_steps, {
|
||||
"eta_seconds": eta,
|
||||
"elapsed_seconds": elapsed,
|
||||
"learning_rate": state.optimizer.param_groups[0]["lr"] if state.optimizer else None,
|
||||
"loss_history": self.runner.loss_history[-50:],
|
||||
})
|
||||
|
||||
progress_cb = ProgressCallback(self, max_steps)
|
||||
|
||||
trainer = Trainer(
|
||||
model=self.model,
|
||||
args=training_args,
|
||||
train_dataset=dataset,
|
||||
data_collator=data_collator,
|
||||
callbacks=[progress_cb],
|
||||
)
|
||||
|
||||
print(f"[铸渊Agent] 开始HLDP训练: {max_steps}步 × {batch_size}batch × {grad_accum}累积")
|
||||
self.start_time = time.time()
|
||||
|
||||
trainer.train()
|
||||
|
||||
# 保存模型
|
||||
trainer.save_model(output_dir)
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
|
||||
elapsed = time.time() - self.start_time
|
||||
print(f"[铸渊Agent] HLDP训练完成!耗时: {elapsed:.0f}s, 最终loss: {self.loss_history[-1] if self.loss_history else 'N/A'}")
|
||||
|
||||
return {
|
||||
"status": "done",
|
||||
"final_loss": self.loss_history[-1] if self.loss_history else None,
|
||||
"steps_completed": max_steps,
|
||||
"elapsed_seconds": elapsed,
|
||||
"output_dir": output_dir,
|
||||
"message": f"HLDP原生格式训练完成,{len(texts)}条语料,{max_steps}步",
|
||||
}
|
||||
|
||||
except ImportError as e:
|
||||
# 如果transformers不可用,返回simulated结果用于测试仪表盘
|
||||
print(f"[铸渊Agent] 训练依赖不可用: {e}")
|
||||
return {
|
||||
"status": "error",
|
||||
"message": f"缺少训练依赖: {e}。请安装: pip install transformers peft accelerate bitsandbytes torch",
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user