AI-Driven cgroup v2 Memory Governance: Let the LLM Read Workload Semantics and Auto-Tune VPS Memory Policy

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

You are a VPS memory governance expert. Based on the following systemd
service information, recommend cgroup v2 memory limits.

Service info:
- unit: nginx.service
- Description: A high performance web server and a reverse proxy server
- Ports: 80, 443
- Dependencies: none
- Historical peak memory (7d): 180M
- Average memory (7d): 95M
- Total available RAM: 3.8G (4G VPS, 200M reserved for OS)
- Other services total peak: 2.1G

Output JSON with these fields:
- role: "gateway" | "database" | "cache" | "worker" | "background" | "api"
- mem_high_mb: soft limit (memory.high), in MB
- mem_max_mb:  hard limit (memory.max), in MB
- mem_low_mb:  retention (memory.low), in MB, can be 0
- reason: one-sentence explanation

Constraints:
- mem_low <= mem_high <= mem_max
- mem_max <= 40% of total available RAM (leave headroom)
- gateway: set mem_high at historical_peak × 1.5
- database: mem_max must be >= configured buffer pool size (if known)
- background: set mem_max at historical_peak × 1.2, do NOT set mem_low

import json, subprocess, os, sys

CGROUP_ROOT = "/sys/fs/cgroup"

def validate_proposal(proposal: dict, total_available_mb: int) -> tuple[bool, str]:
    """Deterministic validation: LLM output must pass this gate."""
    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"order error: low({low}) <= high({high}) <= max({max_})"
    if max_ > total_available_mb * 0.4:
        return False, f"mem_max={max_}MB exceeds 40% of available RAM"
    if max_ == 0:
        return False, "mem_max cannot be 0"
    return True, "OK"

def apply_cgroup(unit: str, high_mb: int, max_mb: int, low_mb: int):
    """systemd-managed services: systemctl set-property.
    Non-systemd services: write cgroupfs directly."""
    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:
        # Fallback: write cgroupfs directly
        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))

    # Persist: write an override drop-in so it survives reboot
    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()  # scan /etc/systemd/system + running units
    total_available = get_available_memory_mb()  # from /proc/meminfo

    for unit in services:
        info = collect_service_context(unit)  # description, ports, deps, history
        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 output not valid JSON, skipping")
            continue

        ok, reason = validate_proposal(proposal, total_available)
        if not ok:
            log.warn(f"{unit}: validation failed: {reason}, keeping current values")
            continue

        current = read_current_cgroup(unit)
        if current == (proposal["mem_high_mb"], proposal["mem_max_mb"],
                       proposal["mem_low_mb"]):
            continue  # no change, don't write

        log.info(f"{unit}: adjusting memory.high={proposal['mem_high_mb']}M "
                 f"memory.max={proposal['mem_max_mb']}M "
                 f"memory.low={proposal['mem_low_mb']}M "
                 f"reason: {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 — LLM-driven cgroup v2 memory governor for VPS"""

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

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:
    """Gather service metadata + 7-day peak memory."""
    ctx = {"unit": unit, "description": "", "ports": [], "dependencies": []}
    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)
            after = re.search(r"After=(.+)$",        content, re.M)
            if desc:  ctx["description"]    = desc.group(1)
            if after: ctx["dependencies"]    = after.group(1).split()
            break
    r = subprocess.run(["ss", "-tlnp"], capture_output=True, text=True)
    if unit.replace(".service", "") in r.stdout:
        ctx["ports"] = list(set(re.findall(r"(\d+):\d+", r.stdout)))
    ctx["hist_peak_mb"] = get_hist_peak(unit)  # 0 if unknown
    return ctx

def get_hist_peak(unit: str) -> int:
    """Query Prometheus for 7-day container memory peak."""
    try:
        q   = (f'max by (name) '
               f'(container_memory_usage_bytes{{name=~".*{unit}.*"}}[7d])')
        url = f"http://127.0.0.1:9090/api/v1/query?query={urllib.parse.quote(q)}"
        data = json.loads(urllib.request.urlopen(url, timeout=5).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"})
    data = json.loads(urllib.request.urlopen(req, timeout=60).read())
    return data["message"]["content"]

def build_prompt(ctx: dict, total_mb: int) -> str:
    peak = ctx.get("hist_peak_mb", 0)
    return f"""You are a VPS memory governance expert. Recommend cgroup v2
memory limits for this systemd service.

Service info:
- unit: {ctx['unit']}
- description: {ctx.get('description','unknown')}
- ports: {', '.join(map(str, ctx.get('ports',[]))) or 'none detected'}
- dependencies: {', '.join(ctx.get('dependencies',[])) or 'none'}
- 7-day peak memory: {peak}MB (0 = unknown)
- total available RAM: {total_mb}MB

Output JSON only (no markdown fences):
{{
  "role": "gateway"|"database"|"cache"|"worker"|"background"|"api",
  "mem_high_mb": <int>,
  "mem_max_mb":  <int>,
  "mem_low_mb":  <int, can be 0>,
  "reason": "<one sentence>"
}}

Constraints:
- 0 <= mem_low <= mem_high <= mem_max <= {total_mb // 2}
- if peak known: mem_high = peak*1.5, mem_max = peak*2
- if peak unknown, use role defaults:
    gateway: 256/512, database: 512/1024, cache: 256/512,
    worker: 128/256, background: 128/256, api: 256/512
- background services: do NOT set mem_low
"""

def read_current_cgroup(unit: str) -> tuple:
    path = f"{CGROUP_ROOT}/system.slice/{unit}"
    def read(fname, default=0):
        try:
            v = open(f"{path}/{fname}").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 failed: {r.stderr}; "
                    f"writing cgroupfs directly")
        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 write failed: {e}")

    # Persist via override drop-in
    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"Available RAM: {total_mb}MB — starting service scan")

    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()]
    units = [u for u in units
             if not u.startswith("systemd-") and u != "dbus.service"]

    log.info(f"Evaluating {len(units)} services")
    for unit in units:
        ctx    = collect_service_context(unit)
        prompt = build_prompt(ctx, total_mb)
        try:
            raw = ollama_chat(prompt)
            raw = re.sub(r"^```json\s*|```$", "", raw, flags=re.M).strip()
            proposal = json.loads(raw)
        except Exception as e:
            log.warning(f"{unit}: LLM parse failed ({e}), skipping")
            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}: invalid ranges "
                        f"(low={low}, high={high}, max={max_}), skipping")
            continue
        cap = total_mb // 2
        if max_ > cap:
            log.warning(f"{unit}: mem_max={max_} exceeds cap {cap}, clamping")
            max_ = cap
            high = min(high, max_)

        current = read_current_cgroup(unit)
        if current == (high, max_, low):
            log.debug(f"{unit}: no change, skipping")
            continue

        log.info(f"{unit}: ({current[0]},{current[1]},{current[2]}) "
                 f"-> ({high},{max_},{low})  reason: {proposal.get('reason','')}")
        apply_cgroup(unit, high, max_, low)
        time.sleep(1)  # avoid triggering kernel reclaim storm from rapid writes

    log.info("Governance complete")

if __name__ == "__main__":
    main()

# 1. Install Ollama and pull the model
curl -fsSL https://ollama.com/install.sh | sh
ollama pull qwen2.5:7b

# 2. Verify cgroup v2 is active (default on most post-2024 distros)
stat -fc %T /sys/fs/cgroup/
# Should print "cgroup2fs" (or "tmpfs" on some kernels).
# If it prints "cgroupfs" (v1), run:
#   echo "GRUB_CMDLINE_LINUX=\"... cgroup_v2=1\"" > /etc/default/grub
#   update-grub && reboot

# 3. Install the script
sudo cp ai_cgroup_governor.py /usr/local/bin/
sudo chmod +x /usr/local/bin/ai_cgroup_governor.py

# 4. Create a systemd timer (daily at 3 AM)
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