AI 驱动的 cgroup v2 内存治理：让 LLM 理解负载语义，实时调优 VPS 内存策略

$ systemctl show postgresql --property=MemoryMax
MemoryMax=∞
$ systemctl show nginx --property=MemoryHigh
MemoryHigh=0

你是一个 VPS 内存治理专家。根据以下 systemd 服务信息，推荐 cgroup v2 内存限额。

服务信息：
- unit: nginx.service
- 描述: A high performance web server and a reverse proxy server
- 端口: 80, 443
- 依赖: 无
- 历史峰值内存（7天）: 180M
- 平均内存（7天）: 95M
- 整机可用内存: 3.8G（4G VPS，系统预留 200M）
- 其他服务总峰值: 2.1G

请输出 JSON，字段：
- role: "gateway" | "database" | "cache" | "worker" | "background" | "api"
- mem_high_mb: 软限（memory.high），单位 MB
- mem_max_mb: 硬限（memory.max），单位 MB
- mem_low_mb: 保留量（memory.low），单位 MB，可以为 0
- reason: 一句话解释

约束：
- mem_low <= mem_high <= mem_max
- mem_max <= 整机可用内存 × 40%（给其他服务留余量）
- 网关类：mem_high 设在历史峰值 × 1.5
- 数据库类：mem_max 必须 >= 配置的 buffer pool 大小（如果已知）
- 后台类：mem_max 设在历史峰值 × 1.2，不要设 mem_low

import json, subprocess, os, sys

CGROUP_ROOT = "/sys/fs/cgroup"

def validate_proposal(proposal: dict, total_available_mb: int) -> tuple[bool, str]:
    """确定性验证：LLM 输出必须过这一关。"""
    high = proposal.get("mem_high_mb", 0)
    max_ = proposal.get("mem_max_mb", 0)
    low  = proposal.get("mem_low_mb", 0)

    if low > high or high > max_:
        return False, f"顺序错误: low({low}) <= high({high}) <= max({max_})"
    if max_ > total_available_mb * 0.4:
        return False, f"mem_max={max_}MB 超过可用内存的 40% 上限"
    if max_ == 0:
        return False, "mem_max 不能为 0"
    return True, "OK"

def apply_cgroup(unit: str, high_mb: int, max_mb: int, low_mb: int):
    """systemd 管理的服务走 systemctl set-property，否则直接写 cgroupfs。"""
    # 方式一：systemctl（推荐，systemd 会同步到 cgroup）
    args = ["systemctl", "set-property", unit,
            f"MemoryHigh={high_mb}M",
            f"MemoryMax={max_mb}M"]
    if low_mb > 0:
        args.append(f"MemoryLow={low_mb}M")
    r = subprocess.run(args, capture_output=True, text=True)
    if r.returncode != 0:
        # 方式二：直接写 cgroupfs（适用于非 systemd 管理的服务）
        path = f"{CGROUP_ROOT}/system.slice/{unit}"
        with open(f"{path}/memory.high", "w") as f:
            f.write(str(high_mb * 1024 * 1024) if high_mb > 0 else "max")
        with open(f"{path}/memory.max", "w") as f:
            f.write(str(max_mb * 1024 * 1024))
        if low_mb > 0:
            with open(f"{path}/memory.low", "w") as f:
                f.write(str(low_mb * 1024 * 1024))

    # 持久化：写进 unit override，防止重启后失效
    override_dir = f"/etc/systemd/system/{unit}.d"
    os.makedirs(override_dir, exist_ok=True)
    with open(f"{override_dir}/memory-governor.conf", "w") as f:
        f.write(f"[Service]\nMemoryHigh={high_mb}M\nMemoryMax={max_mb}M\n")
        if low_mb > 0:
            f.write(f"MemoryLow={low_mb}M\n")
    subprocess.run(["systemctl", "daemon-reload"], capture_output=True)

def run_memory_governor():
    services = get_unit_list()  # 扫描 /etc/systemd/system + 运行中的 unit
    total_available = get_available_memory_mb()  # 从 /proc/meminfo 读

    for unit in services:
        info = collect_service_context(unit)  # 描述、端口、依赖、历史用量
        prompt = build_prompt(info, total_available)
        raw = ollama_chat(prompt)  # qwen2.5-7b, temperature=0
        try:
            proposal = json.loads(raw)
        except json.JSONDecodeError:
            log.warn(f"{unit}: LLM 输出不是合法 JSON，跳过")
            continue

        ok, reason = validate_proposal(proposal, total_available)
        if not ok:
            log.warn(f"{unit}: 验证失败: {reason}，保持原值")
            continue

        current = read_current_cgroup(unit)
        if current == (proposal["mem_high_mb"], proposal["mem_max_mb"], proposal["mem_low_mb"]):
            continue  # 没变化，不写

        log.info(f"{unit}: 调整 memory.high={proposal['mem_high_mb']}M "
                 f"memory.max={proposal['mem_max_mb']}M memory.low={proposal['mem_low_mb']}M "
                 f"原因: {proposal['reason']}")
        apply_cgroup(unit, proposal["mem_high_mb"], proposal["mem_max_mb"],
                     proposal["mem_low_mb"])

#!/usr/bin/env python3
"""ai_cgroup_governor.py — VPS 内存治理 LLM 调度器"""

import json, os, re, subprocess, time, logging, urllib.request

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("cgroup-governor")

OLLAMA_URL = "http://127.0.0.1:11434/api/chat"
OLLAMA_MODEL = "qwen2.5:7b"
CGROUP_ROOT = "/sys/fs/cgroup"

def get_available_memory_mb() -> int:
    with open("/proc/meminfo") as f:
        for line in f:
            if line.startswith("MemAvailable:"):
                return int(line.split()[1]) // 1024
    return 2048

def collect_service_context(unit: str) -> dict:
    """收集服务的元信息 + 最近 7 天的峰值内存。"""
    ctx = {"unit": unit, "description": "", "ports": [], "dependencies": []}
    # 读 unit 文件
    for path in [f"/etc/systemd/system/{unit}", f"/usr/lib/systemd/system/{unit}"]:
        if os.path.exists(path):
            with open(path) as f:
                content = f.read()
            desc = re.search(r"^Description=(.+)$", content, re.M)
            if desc: ctx["description"] = desc.group(1)
            after = re.search(r"After=(.+)$", content, re.M)
            if after: ctx["dependencies"] = after.group(1).split()
            break
    # 读端口（ss 或 解析配置）
    r = subprocess.run(["ss", "-tlnp"], capture_output=True, text=True)
    if unit.replace(".service", "") in r.stdout:
        ports = re.findall(r"(\d+):\d+", r.stdout)
        ctx["ports"] = list(set(ports))
    # 历史峰值：从 Prometheus 拉（如果有）
    ctx["hist_peak_mb"] = get_hist_peak(unit)  # 简化：返回 0 表示未知
    ctx["hist_avg_mb"]  = get_hist_avg(unit)
    return ctx

def get_hist_peak(unit: str) -> int:
    """从 Prometheus 查询最近 7 天的容器内存峰值。"""
    try:
        query = f'max by (name) (container_memory_usage_bytes{{name=~".*{unit}.*"}}[7d])'
        url = f"http://127.0.0.1:9090/api/v1/query?query={urllib.parse.quote(query)}"
        r = urllib.request.urlopen(url, timeout=5)
        data = json.loads(r.read())
        if data.get("data", {}).get("result"):
            return int(data["data"]["result"][0]["value"][1]) // 1024 // 1024
    except Exception:
        pass
    return 0

def ollama_chat(prompt: str) -> str:
    payload = json.dumps({
        "model": OLLAMA_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False,
        "options": {"temperature": 0}
    }).encode()
    req = urllib.request.Request(OLLAMA_URL, data=payload,
                                 headers={"Content-Type": "application/json"})
    r = urllib.request.urlopen(req, timeout=60)
    data = json.loads(r.read())
    return data["message"]["content"]

def build_prompt(ctx: dict, total_available_mb: int) -> str:
    return f"""你是一个 VPS 内存治理专家。根据以下 systemd 服务信息，推荐 cgroup v2 内存限额。

服务信息：
- unit: {ctx['unit']}
- 描述: {ctx['description'] or '未知'}
- 端口: {', '.join(map(str, ctx['ports'])) or '未检测到'}
- 依赖: {', '.join(ctx['dependencies']) or '无'}
- 历史峰值内存（7天）: {ctx['hist_peak_mb']}MB（0=未知）
- 整机可用内存: {total_available_mb}MB

请输出 JSON（不要 markdown 代码块），字段：
- role: "gateway"|"database"|"cache"|"worker"|"background"|"api"
- mem_high_mb: 软限，整数
- mem_max_mb: 硬限，整数
- mem_low_mb: 保留量，整数，可为 0
- reason: 一句话中文解释

约束：
- 0 <= mem_low <= mem_high <= mem_max <= {total_available_mb // 2}
- 历史峰值已知时，mem_high 设在峰值×1.5，mem_max 设在峰值×2
- 未知时，按 role 经验值：gateway 256/512, database 512/1024, cache 256/512, worker 128/256, background 128/256, api 256/512
- background 类不要设 mem_low
"""

def read_current_cgroup(unit: str) -> tuple:
    path = f"{CGROUP_ROOT}/system.slice/{unit}"
    def read(f, default=0):
        try:
            v = open(f"{path}/{f}").read().strip()
            return int(v) // 1024 // 1024 if v != "max" else default
        except Exception:
            return default
    return (read("memory.high"), read("memory.max"), read("memory.low"))

def apply_cgroup(unit: str, high_mb: int, max_mb: int, low_mb: int):
    log.info(f"  → {unit}: high={high_mb}M max={max_mb}M low={low_mb}M")
    r = subprocess.run(
        ["systemctl", "set-property", unit,
         f"MemoryHigh={high_mb}M", f"MemoryMax={max_mb}M"]
        + ([f"MemoryLow={low_mb}M"] if low_mb > 0 else []),
        capture_output=True, text=True)
    if r.returncode != 0:
        log.warning(f"  systemctl set-property 失败: {r.stderr}，尝试直接写 cgroupfs")
        try:
            path = f"{CGROUP_ROOT}/system.slice/{unit}"
            open(f"{path}/memory.high","w").write(str(high_mb*1024*1024))
            open(f"{path}/memory.max","w").write(str(max_mb*1024*1024))
            if low_mb: open(f"{path}/memory.low","w").write(str(low_mb*1024*1024))
        except Exception as e:
            log.error(f"  cgroupfs 写入失败: {e}")

    # 持久化
    odir = f"/etc/systemd/system/{unit}.d"
    os.makedirs(odir, exist_ok=True)
    with open(f"{odir}/memory-governor.conf", "w") as f:
        f.write(f"[Service]\nMemoryHigh={high_mb}M\nMemoryMax={max_mb}M\n")
        if low_mb: f.write(f"MemoryLow={low_mb}M\n")
    subprocess.run(["systemctl","daemon-reload"], capture_output=True)

def main():
    total_mb = get_available_memory_mb()
    log.info(f"整机可用内存: {total_mb}MB，开始扫描服务")

    # 扫描运行中的服务
    r = subprocess.run(["systemctl", "list-units", "--type=service",
                        "--state=running", "--no-legend", "--no-pager"],
                       capture_output=True, text=True)
    units = [line.split()[0] for line in r.stdout.strip().splitlines()]
    # 过滤掉 systemd 内部服务
    units = [u for u in units if not u.startswith("systemd-") and u != "dbus.service"]

    log.info(f"共 {len(units)} 个服务待评估")
    for unit in units:
        ctx = collect_service_context(unit)
        prompt = build_prompt(ctx, total_mb)
        try:
            raw = ollama_chat(prompt)
            # 提取 JSON（LLM 有时会带 markdown 围栏）
            raw = re.sub(r"^```json\s*|```$", "", raw, flags=re.M).strip()
            proposal = json.loads(raw)
        except Exception as e:
            log.warning(f"{unit}: LLM 解析失败 ({e})，跳过")
            continue

        high  = proposal.get("mem_high_mb", 0)
        max_  = proposal.get("mem_max_mb", 0)
        low   = proposal.get("mem_low_mb", 0)

        if low > high or high > max_ or max_ == 0:
            log.warning(f"{unit}: 数值顺序错误 (low={low}, high={high}, max={max_})，跳过")
            continue
        if max_ > total_mb // 2:
            log.warning(f"{unit}: mem_max={max_} 超过安全上限 {total_mb//2}，截断")
            max_ = total_mb // 2
            high = min(high, max_)

        current = read_current_cgroup(unit)
        if current == (high, max_, low):
            log.debug(f"{unit}: 无变化，跳过")
            continue

        log.info(f"{unit}: 调整 ({current[0]},{current[1]},{current[2]}) "
                 f"→ ({high},{max_},{low})  原因: {proposal.get('reason','')}")
        apply_cgroup(unit, high, max_, low)
        time.sleep(1)  # 避免连续写 cgroupfs 触发内核回收风暴

    log.info("治理完成")

if __name__ == "__main__":
    main()

# 1. 安装 Ollama 并拉模型
curl -fsSL https://ollama.com/install.sh | sh
ollama pull qwen2.5:7b

# 2. 确认 cgroup v2 已启用（2024 年后多数发行版默认是）
stat -fc %T /sys/fs/cgroup/   # 应输出 "tmpfs"（v2）或 "cgroup2fs"
# 如果是 "cgroupfs"（v1），需要:
# echo CGROUP2=1 > /etc/default/grub && grub-mkconfig -s && reboot

# 3. 放置脚本
sudo cp ai_cgroup_governor.py /usr/local/bin/
sudo chmod +x /usr/local/bin/ai_cgroup_governor.py

# 4. 建 timer（每天凌晨 3 点）
sudo tee /etc/systemd/system/ai-cgroup-governor.service <<'EOF'
[Unit]
Description=AI cgroup memory governor
After=network.target ollama.service

[Service]
User=root
ExecStart=/usr/local/bin/ai_cgroup_governor.py
Restart=on-failure

[Install]
WantedBy=multi-user.target
EOF

sudo tee /etc/systemd/system/ai-cgroup-governor.timer <<'EOF'
[Unit]
Description=Run AI cgroup governor daily

[Timer]
OnCalendar=*-*-* 03:00:00
Persistent=true

[Install]
WantedBy=timers.target
EOF

sudo systemctl enable --now ai-cgroup-governor.timer
sudo systemctl start ollama