#!/usr/bin/env python3 # -*- coding: utf-8 -*- """作业 7 · vLLM 故障检测 / 重启 / 钉钉通知。""" from __future__ import annotations import json import re import urllib.request from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Callable, Optional import _bootstrap # noqa: F401 import _config as C TZ8 = timezone(timedelta(hours=8)) Runner = Callable[[str], str] OOM_PATTERNS = ( "CUDA out of memory", "OutOfMemoryError", "out of memory", "OOM", "CUDA error: out of memory", ) def data_dir() -> Path: p = Path(getattr(C, "HOMEWORK7_DATA_DIR", "/opt/workbuddy/data/homework7")) p.mkdir(parents=True, exist_ok=True) return p def simulate_flag_path() -> Path: return data_dir() / "simulate_oom.json" def gpu_exec(cmd: str, runner: Optional[Runner] = None, timeout: int = 60) -> str: if runner is not None: return runner(cmd) import paramiko c = paramiko.SSHClient() c.set_missing_host_key_policy(paramiko.AutoAddPolicy()) c.connect(C.GPU_HOST, username=C.GPU_USER, password=C.GPU_PASS, timeout=20) try: _i, o, e = c.exec_command(cmd, timeout=timeout) out = o.read().decode("utf-8", "replace") err = e.read().decode("utf-8", "replace") if not out.strip() and err.strip(): raise RuntimeError(err.strip()) return out finally: c.close() def read_simulate_oom() -> Optional[dict]: p = simulate_flag_path() if not p.is_file(): return None data = json.loads(p.read_text(encoding="utf-8")) if data.get("active"): return data return None def set_simulate_oom(active: bool = True, reason: str = "CUDA out of memory") -> dict: data = { "active": active, "reason": reason, "container": C.GPU_VLLM_CONTAINER, "ts": datetime.now(TZ8).isoformat(), "mode": "simulate", } simulate_flag_path().write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") return data def clear_simulate_oom() -> None: p = simulate_flag_path() if p.is_file(): p.unlink() def parse_container_status(raw: str) -> dict: line = (raw or "").strip().splitlines()[0] if (raw or "").strip() else "" parts = line.split("|") if len(parts) < 3: return {"running": False, "status": "unknown", "container": C.GPU_VLLM_CONTAINER} return { "container": parts[0].lstrip("/"), "running": parts[1].strip().lower() == "true", "status": parts[2].strip(), } def scan_logs_for_oom(log_text: str) -> list[str]: hits = [] for line in (log_text or "").splitlines(): for pat in OOM_PATTERNS: if pat.lower() in line.lower(): hits.append(line.strip()[:300]) break return hits[-5:] def detect_vllm_anomaly(runner: Optional[Runner] = None) -> dict: sim = read_simulate_oom() if sim: return { "anomaly": True, "kind": "OOM", "source": "simulate_flag", "container": C.GPU_VLLM_CONTAINER, "host": C.GPU_HOST, "evidence": [sim.get("reason", "CUDA out of memory")], "recommendation": "restart_vllm_container", } inspect_cmd = ( "docker inspect -f '{{.Name}}|{{.State.Running}}|{{.State.Status}}|" "{{.State.ExitCode}}|{{.State.Error}}' %s" % C.GPU_VLLM_CONTAINER ) status = parse_container_status(gpu_exec(inspect_cmd, runner=runner)) logs_cmd = "docker logs --since 30m --tail 200 %s 2>&1" % C.GPU_VLLM_CONTAINER logs = gpu_exec(logs_cmd, runner=runner, timeout=90) oom_hits = scan_logs_for_oom(logs) anomaly = False kind = "none" evidence: list[str] = [] if oom_hits: anomaly = True kind = "OOM" evidence = oom_hits elif not status.get("running"): anomaly = True kind = "container_down" evidence = ["container not running: %s" % status.get("status")] return { "anomaly": anomaly, "kind": kind, "source": "docker_logs" if oom_hits else "docker_inspect", "container": status.get("container") or C.GPU_VLLM_CONTAINER, "host": C.GPU_HOST, "running": status.get("running"), "status": status.get("status"), "evidence": evidence, "recommendation": "restart_vllm_container" if anomaly else "none", } def restart_vllm_container(runner: Optional[Runner] = None, force_real: bool = False) -> dict: """重启 vLLM。若存在 simulate_oom 标记且未 force_real，则仅模拟重启（避免 Agent 对话中途断连）。""" if read_simulate_oom() and not force_real: clear_simulate_oom() rec = { "action": "simulate restart", "container": C.GPU_VLLM_CONTAINER, "host": C.GPU_HOST, "status_after": "running", "ok": True, "ts": datetime.now(TZ8).isoformat(), "note": "homework demo: skipped real docker restart while Agent is chatting", } hist = data_dir() / "restart_history.jsonl" with hist.open("a", encoding="utf-8") as f: f.write(json.dumps(rec, ensure_ascii=False) + "\n") return rec cmd = "docker restart %s && sleep 3 && docker inspect -f '{{.State.Status}}' %s" % ( C.GPU_VLLM_CONTAINER, C.GPU_VLLM_CONTAINER, ) out = gpu_exec(cmd, runner=runner, timeout=120) status = (out or "").strip().splitlines()[-1].strip() if out else "unknown" clear_simulate_oom() rec = { "action": "docker restart", "container": C.GPU_VLLM_CONTAINER, "host": C.GPU_HOST, "status_after": status, "ok": status == "running", "ts": datetime.now(TZ8).isoformat(), } hist = data_dir() / "restart_history.jsonl" with hist.open("a", encoding="utf-8") as f: f.write(json.dumps(rec, ensure_ascii=False) + "\n") return rec def send_recovery_notification( message: str, detail: str = "", http_post: Optional[Callable[[str], str]] = None, ) -> dict: kw = getattr(C, "HOMEWORK7_DINGTALK_KEYWORD", "") or getattr(C, "HOMEWORK6_DINGTALK_KEYWORD", "") title = "vLLM 故障恢复通知" text = message if detail: text += "\n\n" + detail if kw and kw not in text: text = "**%s**\n\n%s" % (kw, text) body = {"msgtype": "markdown", "markdown": {"title": title, "text": text}} if http_post is not None: raw = http_post(json.dumps(body, ensure_ascii=False)) return json.loads(raw) req = urllib.request.Request( C.HOMEWORK6_DINGTALK_WEBHOOK, data=json.dumps(body, ensure_ascii=False).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST", ) with urllib.request.urlopen(req, timeout=30) as resp: return json.loads(resp.read().decode("utf-8", "replace"))