inkreach-prefect/scripts/metrics_agent.py

"""metrics-agent:宿主机指标采样侧车(独立于 Prefect,常驻循环,compose 服务运行)。

每 5 分钟追加一行 JSONL:data/metrics/metrics-YYYY-MM-DD.jsonl
  {"ts","memPct","memUsed","memTotal","swapPct","swapUsed","swapTotal",
   "gpus":[{"index","name","utilPct","memPct","memUsed","memTotal","tempC"}]}
每日 7 点后做一次 du 快照:data/metrics/disk-YYYYMMDD-HHMM.json
  (根目录 / home 各自的总量 + 顶层目录字节数,device-check 对比最近两份算"增长最快目录")。
保留:采样 14 天、快照 7 份。

compose 中挂载 /:/host:ro(du 扫描)、GPU(nvidia-smi)、./data/metrics(writable)。
调试单次采样:uv run python scripts/metrics_agent.py --once
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import threading
import time
from datetime import datetime
from pathlib import Path

INTERVAL = int(os.environ.get("METRICS_INTERVAL", "300"))
DATA_DIR = Path(os.environ.get("METRICS_DIR", "data/metrics"))
HOST_ROOT = os.environ.get("METRICS_HOST_ROOT", "/host")
SNAPSHOT_AFTER_HOUR = 7
KEEP_SAMPLE_DAYS = 14
KEEP_SNAPSHOTS = 7
# 虚拟文件系统与不关心的挂载点,du 快照跳过
DU_EXCLUDE = {"proc", "sys", "dev", "run", "lost+found"}


def _gb(nbytes: float) -> float:
    return round(nbytes / (1 << 30), 2)


def sample() -> dict:
    import psutil

    vm, sw = psutil.virtual_memory(), psutil.swap_memory()
    return {
        "ts": datetime.now().isoformat(timespec="seconds"),
        "memPct": vm.percent, "memUsed": _gb(vm.used), "memTotal": _gb(vm.total),
        "swapPct": sw.percent, "swapUsed": _gb(sw.used), "swapTotal": _gb(sw.total),
        "gpus": _gpus(),
    }


def _gpus() -> list[dict]:
    try:
        proc = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,name,utilization.gpu,memory.used,memory.total,temperature.gpu",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=20,
        )
        if proc.returncode != 0:
            return []
        out = []
        for line in proc.stdout.strip().splitlines():
            p = [x.strip() for x in line.split(",")]
            if len(p) != 6:
                continue
            idx, name, util, mem_used, mem_total, temp = p
            mt = float(mem_total)
            out.append({
                "index": int(idx), "name": name[:24], "utilPct": float(util),
                "memPct": round(float(mem_used) / mt * 100, 1) if mt > 0 else 0.0,
                "memUsed": _gb(float(mem_used) * (1 << 20)), "memTotal": _gb(mt * (1 << 20)),
                "tempC": float(temp),
            })
        return out
    except Exception:  # noqa: BLE001
        return []


def write_sample(rec: dict) -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    path = DATA_DIR / f"metrics-{rec['ts'][:10]}.jsonl"
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(rec, ensure_ascii=False) + "\n")


def _du_bytes(path: str) -> int:
    # 实际占块(非表面大小,避免 docker 稀疏文件虚增);rc=1 时是个别文件消失,stdout 仍有效
    proc = subprocess.run(["du", "-sB1", path], capture_output=True, text=True, timeout=7200)
    if not proc.stdout:
        return -1
    try:
        return int(proc.stdout.split()[0])
    except (IndexError, ValueError):
        return -1


def _usage_entry(base: str) -> dict:
    import shutil

    st = shutil.disk_usage(base)
    return {"total": st.total, "used": st.used, "free": st.free}


def disk_snapshot() -> Path:
    """du 扫描根目录/home 顶层目录,写 disk-YYYYMMDD-HHMM.json(原子写)。"""
    now = datetime.now()
    out: dict = {"ts": now.isoformat(timespec="seconds")}
    for key, base in (("root", HOST_ROOT), ("home", f"{HOST_ROOT}/home")):
        if not os.path.isdir(base):
            continue
        entry = _usage_entry(base)
        dirs: dict[str, int] = {}
        for name in sorted(os.listdir(base)):
            if name in DU_EXCLUDE:
                continue
            p = os.path.join(base, name)
            if os.path.isdir(p) and not os.path.islink(p) and not os.path.ismount(p):
                size = _du_bytes(p)
                if size >= 0:
                    dirs[name] = size
        entry["dirs"] = dirs
        out[key] = entry
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    path = DATA_DIR / f"disk-{now:%Y%m%d-%H%M}.json"
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(out, ensure_ascii=False), encoding="utf-8")
    tmp.rename(path)
    print(f"[metrics-agent] 磁盘快照完成: {path.name}", flush=True)
    return path


def snapshot_today_exists() -> bool:
    today = datetime.now().strftime("%Y%m%d")
    return any(DATA_DIR.glob(f"disk-{today}-*.json"))


def prune() -> None:
    cutoff = time.time() - KEEP_SAMPLE_DAYS * 86400
    for p in DATA_DIR.glob("metrics-*.jsonl"):
        if p.stat().st_mtime < cutoff:
            p.unlink(missing_ok=True)
    snaps = sorted(DATA_DIR.glob("disk-*.json"))
    for p in snaps[:-KEEP_SNAPSHOTS] if len(snaps) > KEEP_SNAPSHOTS else []:
        p.unlink(missing_ok=True)


def _maybe_snapshot_async() -> None:
    """7 点后补当日快照;du 可能耗时数分钟,放后台线程不阻塞采样。"""
    if datetime.now().hour < SNAPSHOT_AFTER_HOUR or snapshot_today_exists():
        return
    threading.Thread(target=lambda: _safe(disk_snapshot), daemon=True).start()


def _safe(fn, *args):
    try:
        fn(*args)
    except Exception as err:  # noqa: BLE001
        print(f"[metrics-agent] {type(err).__name__}: {err}", flush=True)


def main() -> None:
    ap = argparse.ArgumentParser(description="宿主机指标采样侧车")
    ap.add_argument("--once", action="store_true", help="只采样一次即退出(调试)")
    args = ap.parse_args()
    print(f"[metrics-agent] 启动: interval={INTERVAL}s dir={DATA_DIR} host_root={HOST_ROOT}", flush=True)
    while True:
        _safe(lambda: write_sample(sample()))
        _maybe_snapshot_async()
        _safe(prune)
        if args.once:
            break
        time.sleep(INTERVAL)


if __name__ == "__main__":
    main()