| |
| """Run a command inside a cgroup-v2 MemoryMax scope and report peak memory.""" |
|
|
| import argparse |
| import json |
| import os |
| import subprocess |
| import sys |
| import time |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--limit-gib", type=int, required=True) |
| parser.add_argument("--timeout", type=float, default=7200) |
| 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") |
| if args.limit_gib not in (2, 4, 8, 16, 32): |
| parser.error("limit must be 2, 4, 8, 16, or 32 GiB") |
|
|
| unit = f"flint-epq-bench-{os.getpid()}" |
| runner = [ |
| "systemd-run", "--quiet", "--wait", "--collect", "--pipe", |
| "--unit", unit, "--property", f"MemoryMax={args.limit_gib}G", |
| "--property", "MemorySwapMax=0", "--", *command, |
| ] |
| started = time.monotonic() |
| process = subprocess.Popen(runner, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) |
| time.sleep(1) |
| if process.poll() is not None: |
| output, _ = process.communicate() |
| if output: |
| sys.stdout.write(output) |
| print("UNTESTED: systemd/cgroup MemoryMax unavailable; RLIMIT_DATA is not a substitute", file=sys.stderr) |
| return 77 |
|
|
| peak = swap_peak = 0 |
| while process.poll() is None: |
| try: |
| cgroup = subprocess.check_output( |
| ["systemctl", "show", f"{unit}.scope", "-p", "ControlGroup", "--value"], |
| text=True, stderr=subprocess.DEVNULL, |
| ).strip() |
| if cgroup: |
| path = os.path.join("/sys/fs/cgroup", cgroup.lstrip("/")) |
| for name, target in (("memory.peak", "peak"), ("memory.swap.peak", "swap")): |
| try: |
| with open(os.path.join(path, name), encoding="ascii") as stream: |
| value = int(stream.read().strip()) |
| if target == "peak": |
| peak = max(peak, value) |
| else: |
| swap_peak = max(swap_peak, value) |
| except OSError: |
| pass |
| except (OSError, subprocess.CalledProcessError): |
| pass |
| if time.monotonic() - started > args.timeout: |
| process.kill() |
| break |
| time.sleep(0.25) |
|
|
| output, _ = process.communicate() |
| if output: |
| sys.stdout.write(output) |
| print("RUNMEM " + json.dumps({ |
| "limit_gib": args.limit_gib, |
| "peak_bytes": peak, |
| "peak_gib": round(peak / (1024 ** 3), 3), |
| "swap_peak_bytes": swap_peak, |
| "wall_s": round(time.monotonic() - started, 2), |
| "exit": process.returncode, |
| "cmd": command, |
| "method": "cgroup-v2-MemoryMax", |
| })) |
| return process.returncode |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
