Spaces:
Paused
Paused
| #!/usr/bin/env python3 | |
| """Kill only the processes opencode spawns, never opencode itself. | |
| Logic: | |
| roots = processes whose comm starts with "opencode" | |
| workers = every descendant of those roots (opencode itself excluded) | |
| total = summed RSS of workers | |
| while total > LIMIT: kill the direct child of opencode holding the most RSS | |
| (whole subtree, so the runaway program A dies with its kids) | |
| """ | |
| import json | |
| import os | |
| import re | |
| import signal | |
| import sys | |
| import time | |
| LIMIT = float(sys.argv[1] if len(sys.argv) > 1 else 80) * 1024 ** 3 | |
| INTERVAL = float(sys.argv[2] if len(sys.argv) > 2 else 2) | |
| PER_PROC = float(sys.argv[3] if len(sys.argv) > 3 else 64) * 1024 ** 3 | |
| MARK = "/tmp/memkill.last.json" | |
| HIST = "/tmp/memkill.jsonl" | |
| PAGE = os.sysconf("SC_PAGE_SIZE") | |
| def read_cmd(p): | |
| try: | |
| with open("/proc/%d/cmdline" % p, "rb") as f: | |
| return f.read().replace(b"\0", b" ").decode(errors="replace").strip() | |
| except Exception: | |
| return "" | |
| def snapshot(): | |
| ppid = {} | |
| rss = {} | |
| comm = {} | |
| for d in os.listdir("/proc"): | |
| if not d.isdigit(): | |
| continue | |
| p = int(d) | |
| try: | |
| with open("/proc/%d/stat" % p) as f: | |
| s = f.read() | |
| ppid[p] = int(s[s.rindex(")") + 2:].split()[1]) | |
| with open("/proc/%d/statm" % p) as f: | |
| rss[p] = int(f.read().split()[1]) * PAGE | |
| with open("/proc/%d/comm" % p) as f: | |
| comm[p] = f.read().strip() | |
| except Exception: | |
| continue | |
| return ppid, rss, comm | |
| def protect(): | |
| """never let the guard (or its own helpers) become a victim""" | |
| out = {os.getpid()} | |
| for d in os.listdir("/proc"): | |
| if d.isdigit(): | |
| try: | |
| with open("/proc/%s/cmdline" % d, "rb") as f: | |
| if b"memkill.py" in f.read(): | |
| out.add(int(d)) | |
| except Exception: | |
| pass | |
| return out | |
| def main(): | |
| try: | |
| with open("/tmp/.memguard.pid", "w") as f: | |
| f.write(str(os.getpid())) | |
| except Exception: | |
| pass | |
| print("memkill: guard armed total_limit=%.0fGiB per_proc_limit=%.0fGiB" | |
| % (LIMIT / 1024 ** 3, PER_PROC / 1024 ** 3), flush=True) | |
| while True: | |
| ppid, rss, comm = snapshot() | |
| keep = protect() | |
| roots = {p for p, c in comm.items() if c.startswith("opencode")} - keep | |
| kids = {} | |
| for child, parent in ppid.items(): | |
| kids.setdefault(parent, []).append(child) | |
| # subtree walk helpers | |
| def sub(p): | |
| out = {p} | |
| stack = [p] | |
| while stack: | |
| for k in kids.get(stack.pop(), ()): | |
| if k not in out: | |
| out.add(k) | |
| stack.append(k) | |
| return out | |
| workers = set() | |
| top = [] # direct children of opencode | |
| for r in roots: | |
| for k in kids.get(r, ()): | |
| if k not in workers: | |
| workers |= sub(k) | |
| top.append(k) | |
| workers -= roots | keep | |
| # threshold counts opencode itself too, so "opencode + everything it | |
| # started" stays under LIMIT -- but victims are only ever children | |
| total = sum(rss.get(p, 0) for p in workers | roots) | |
| print("%s opencode=%.2fGiB workers=%d total=%.2fGiB" | |
| % (time.strftime("%H:%M:%S"), | |
| sum(rss.get(p, 0) for p in roots) / 1024 ** 3, | |
| len(workers), total / 1024 ** 3), flush=True) | |
| victims = set() | |
| top_set = set(top) | |
| if total > LIMIT: | |
| while total > LIMIT and top: | |
| c = max(top, key=lambda p: sum(rss.get(x, 0) for x in sub(p))) | |
| if c in victims: | |
| break | |
| victims |= sub(c) | |
| top.remove(c) | |
| total -= sum(rss.get(x, 0) for x in victims) | |
| if PER_PROC: | |
| for p in list(workers): | |
| if rss.get(p, 0) > PER_PROC: | |
| victims |= sub(p) | |
| doomed = victims - roots - keep | |
| cmds = {v: read_cmd(v) for v in doomed} | |
| for v in sorted(doomed): | |
| try: | |
| os.kill(v, signal.SIGKILL) | |
| print(" KILL pid=%d rss=%.2fGiB (opencode untouched) %s" | |
| % (v, rss.get(v, 0) / 1024 ** 3, cmds[v][:90]), flush=True) | |
| except Exception as e: | |
| print(" skip pid=%d: %s" % (v, e), flush=True) | |
| if doomed: | |
| # the direct child of opencode == the command the assistant ran | |
| heads = [v for v in doomed if v in top_set] or sorted(doomed) | |
| rec = { | |
| "ts": round(time.time(), 3), | |
| "iso": time.strftime("%Y-%m-%d %H:%M:%S"), | |
| "head_pid": heads[0], | |
| "head_cmd": cmds.get(heads[0], ""), | |
| "cmds": [cmds[v] for v in sorted(doomed)][:8], | |
| "rss_gib": round(max(rss.get(v, 0) for v in doomed) / 1024 ** 3, 2), | |
| "total_gib": round(sum(rss.get(v, 0) for v in doomed) / 1024 ** 3, 2), | |
| "limit_gib": round(LIMIT / 1024 ** 3, 2), | |
| "per_proc_gib": round(PER_PROC / 1024 ** 3, 2), | |
| "reason": "over-limit", | |
| } | |
| try: | |
| with open(MARK + ".tmp", "w") as f: | |
| json.dump(rec, f) | |
| os.replace(MARK + ".tmp", MARK) | |
| with open(HIST, "a") as f: | |
| f.write(json.dumps(rec) + "\n") | |
| except Exception as e: | |
| print(" marker failed: %s" % e, flush=True) | |
| time.sleep(3) | |
| time.sleep(INTERVAL) | |
| if __name__ == "__main__": | |
| main() |