| |
| """Sample process-tree RSS/swap while running a command.""" |
|
|
| import argparse |
| import json |
| import signal |
| import subprocess |
| import sys |
| import time |
|
|
|
|
| def process_tree(root): |
| output = subprocess.check_output(["ps", "-e", "-o", "pid=,ppid="], text=True) |
| children = {} |
| for line in output.splitlines(): |
| pid, ppid = map(int, line.split()) |
| children.setdefault(ppid, []).append(pid) |
| result, todo = [], [root] |
| while todo: |
| pid = todo.pop() |
| result.append(pid) |
| todo.extend(children.get(pid, ())) |
| return result |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--interval", type=float, default=0.25) |
| parser.add_argument("--timeout", type=float, default=7200) |
| parser.add_argument("--rss-limit-gib", type=float) |
| parser.add_argument("command", nargs=argparse.REMAINDER) |
| args = parser.parse_args() |
| command = args.command[1:] if args.command[:1] == ["--"] else args.command |
| if not command: |
| parser.error("missing command") |
|
|
| started = time.monotonic() |
| process = subprocess.Popen(command) |
| peak_rss = peak_swap = 0 |
| exceeded = False |
| while process.poll() is None: |
| rss = swap = 0 |
| for pid in process_tree(process.pid): |
| try: |
| with open(f"/proc/{pid}/status", encoding="ascii") as status: |
| for line in status: |
| if line.startswith("VmRSS:"): |
| rss += int(line.split()[1]) |
| elif line.startswith("VmSwap:"): |
| swap += int(line.split()[1]) |
| except (FileNotFoundError, ProcessLookupError): |
| pass |
| peak_rss = max(peak_rss, rss) |
| peak_swap = max(peak_swap, swap) |
| |
| if args.rss_limit_gib is not None and rss > args.rss_limit_gib * (1024 ** 2): |
| exceeded = True |
| process.send_signal(signal.SIGKILL) |
| break |
| if time.monotonic() - started > args.timeout: |
| process.kill() |
| break |
| time.sleep(args.interval) |
| result = process.wait() |
| print("RUNMEM " + json.dumps({ |
| "peak_kb": peak_rss, |
| "peak_gib": round(peak_rss / 1048576, 3), |
| "limit_gib": args.rss_limit_gib, |
| "limit_exceeded": exceeded, |
| "swap_kb": peak_swap, |
| "wall_s": round(time.monotonic() - started, 2), |
| "exit": result, |
| "cmd": command, |
| "method": "procfs-rss-watchdog", |
| })) |
| return 0 if args.rss_limit_gib is None else (1 if exceeded else result) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
