Spaces:
Paused
Paused
File size: 5,787 Bytes
1949064 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | #!/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() |