D100 训练Watchdog: GPU服务器状态采集脚本

scripts/training_watchdog.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""铸渊训练状态Watchdog - GPU服务器端"""
+import os, json, time, re, sys, subprocess
+sys.stdout.reconfigure(line_buffering=True)
+
+LOG_FILE = "/root/autodl-tmp/train_mother.log"
+CODE_LOG = "/root/autodl-tmp/train_coder.log"
+STATUS_FILE = "/root/autodl-tmp/training_status.json"
+POLL_INTERVAL = 300
+
+def run(cmd):
+    try:
+        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
+        return r.stdout.strip(), r.returncode
+    except:
+        return "", -1
+
+def get_gpu_info():
+    out, _ = run("nvidia-smi --query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu --format=csv,noheader,nounits")
+    parts = out.split(", ")
+    if len(parts) >= 4:
+        return {"util_pct": float(parts[0]), "mem_used_mb": int(parts[1]), 
+                "mem_total_mb": int(parts[2]), "temp_c": float(parts[3])}
+    return None
+
+def parse_mother_status():
+    status = {"status": "unknown", "status_label": "⚪ 未知",
+              "progress": "N/A", "current_step": 0, "current_loss": None,
+              "current_epoch": 0, "elapsed_hours": 0, "has_error": False}
+    if not os.path.exists(LOG_FILE):
+        status["status"] = "pending"; status["status_label"] = "⚪ 未启动"
+        return status
+    with open(LOG_FILE, 'r') as f:
+        content = f.read()
+    if "Traceback" in content:
+        status["status"] = "error"; status["status_label"] = "🔴 报错"; status["has_error"] = True
+    if "DONE!" in content:
+        status["status"] = "completed"; status["status_label"] = "🟢 已完成"
+        return status
+    if "[5/5] Starting training!" in content:
+        status["status"] = "training"; status["status_label"] = "🟢 训练中"
+        steps = re.findall(r"global_step[\s:=]+(\d+)|Step:\s*(\d+)", content)
+        losses = re.findall(r"loss['\s:=]+([\d.]+)", content)
+        if steps: status["current_step"] = int([s for pair in steps for s in pair if s][-1])
+        if losses: status["current_loss"] = float(losses[-1])
+    elif "Tokenize:" in content:
+        matches = re.findall(r"Tokenize:\s*\d+%\|.*?\|\s*(\d+)/(\d+)", content)
+        if matches:
+            done, total = matches[-1]
+            status["progress"] = f"{done} / {total} 条"
+            pct = int(done) / int(total) * 100
+            status["status_label"] = f"🟡 分词中 ({pct:.0f}%)"
+            status["status"] = "tokenizing"
+    return status
+
+def parse_coder_status():
+    if not os.path.exists(CODE_LOG):
+        return {"status": "pending", "status_label": "⚪ 未启动"}
+    with open(CODE_LOG, 'r') as f:
+        content = f.read()
+    if "Traceback" in content:
+        return {"status": "error", "status_label": "🔴 报错", "has_error": True}
+    if "DONE!" in content:
+        return {"status": "completed", "status_label": "🟢 已完成"}
+    if "[5/5] Starting training!" in content:
+        return {"status": "training", "status_label": "🟢 训练中"}
+    return {"status": "starting", "status_label": "🟡 启动中"}
+
+def check_output():
+    mother_out = "/root/autodl-tmp/output/qwen25-7b-sft/final"
+    coder_out = "/root/autodl-tmp/output/qwen25-coder-7b-sft/final"
+    return {
+        "mother_exists": os.path.isdir(mother_out) and bool(os.listdir(mother_out)),
+        "coder_exists": os.path.isdir(coder_out) and bool(os.listdir(coder_out))
+    }
+
+def main():
+    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] Watchdog started"); sys.stdout.flush()
+    while True:
+        try:
+            gpu = get_gpu_info()
+            mother = parse_mother_status()
+            coder = parse_coder_status()
+            output = check_output()
+            status = {
+                "last_updated": time.strftime("%Y-%m-%dT%H:%M:%S+08:00"),
+                "mother_model": {**mother, "gpu": gpu},
+                "code_model": coder,
+                "output": output,
+                "alerts": []
+            }
+            if mother.get("has_error"):
+                status["alerts"].append({"level": "error", "message": "❌ 母模型报错！找铸渊"})
+            if coder.get("has_error"):
+                status["alerts"].append({"level": "error", "message": "❌ 代码模型报错！找铸渊"})
+            with open(STATUS_FILE, 'w') as f:
+                json.dump(status, f, indent=2, ensure_ascii=False)
+            gpu_str = f"GPU:{gpu['mem_used_mb']}/{gpu['mem_total_mb']}MB {gpu['temp_c']}C" if gpu else "N/A"
+            print(f"[{time.strftime('%H:%M:%S')}] {mother['status_label']} | {gpu_str}")
+            sys.stdout.flush()
+        except Exception as e:
+            print(f"Error: {e}"); sys.stdout.flush()
+        time.sleep(POLL_INTERVAL)
+
+if __name__ == "__main__":
+    main()