Support multiplier memory limits (#395)

Author: JasonGross

coq_tools/diagnose_error.py

@@ -271,6 +271,7 @@ def make_reg_string(output, strict_whitespace=False):
 
 
 TIMEOUT = {}
+MEMORY_USAGE = {}
 
 
 def get_timeout(coqc=None):
@@ -287,6 +288,21 @@ def reset_timeout():
     TIMEOUT = {}
 
 
+def get_memory_usage(key):
+    return MEMORY_USAGE.get(key)
+
+
+def set_memory_usage(key, value):
+    if value is None:
+        return
+    MEMORY_USAGE[key] = value
+
+
+def reset_memory_usage():
+    global MEMORY_USAGE
+    MEMORY_USAGE = {}
+
+
 def timeout_Popen_communicate(log, *args, **kwargs):
     ret = {"value": ("", ""), "returncode": None, "rusage": None}
     timeout = kwargs.pop("timeout", None)
@@ -298,9 +314,34 @@ def timeout_Popen_communicate(log, *args, **kwargs):
     # Extract memory limit parameters
     max_mem_rss = kwargs.pop("max_mem_rss", None)
     max_mem_as = kwargs.pop("max_mem_as", None)
+    max_mem_rss_multiplier = kwargs.pop("max_mem_rss_multiplier", None)
+    max_mem_as_multiplier = kwargs.pop("max_mem_as_multiplier", None)
+    memory_usage_key = kwargs.pop("memory_usage_key", None)
     cgroup = kwargs.pop("cgroup", None)
     mem_limit_method = kwargs.pop("mem_limit_method", "prlimit")
 
+    def resolve_dynamic_limit(limit_value, multiplier_value, limit_kind):
+        if multiplier_value is None:
+            return limit_value
+        if multiplier_value <= 0:
+            return limit_value
+        if memory_usage_key is None:
+            return None
+        usage = get_memory_usage((memory_usage_key, limit_kind))
+        if usage is None:
+            return None
+        resolved = int(math.ceil(usage * multiplier_value))
+        return resolved if resolved > 0 else None
+
+    max_mem_rss = resolve_dynamic_limit(max_mem_rss, max_mem_rss_multiplier, "rss")
+    max_mem_as = resolve_dynamic_limit(max_mem_as, max_mem_as_multiplier, "as")
+
+    def record_memory_usage(usage_bytes):
+        if memory_usage_key is None or usage_bytes is None:
+            return
+        set_memory_usage((memory_usage_key, "rss"), usage_bytes)
+        set_memory_usage((memory_usage_key, "as"), usage_bytes)
+
     # Get command from args
     cmd = list(args[0]) if args else []
     cg_path = None
@@ -333,15 +374,19 @@ def get_peak_rss():
     thread.join(timeout)
     if not thread.is_alive():
         cleanup_cgroup(cg_path)
-        return (ret["value"], ret["returncode"], get_peak_rss())
+        peak_rss_bytes = get_peak_rss()
+        record_memory_usage(peak_rss_bytes)
+        return (ret["value"], ret["returncode"], peak_rss_bytes)
 
     p.terminate()
     thread.join()
     cleanup_cgroup(cg_path)
+    peak_rss_bytes = get_peak_rss()
+    record_memory_usage(peak_rss_bytes)
     return (
         tuple(map((lambda s: (s if s else "") + TIMEOUT_POSTFIX), ret["value"])),
         ret["returncode"],
-        get_peak_rss(),
+        peak_rss_bytes,
     )
 
 
@@ -577,11 +622,17 @@ def get_coq_output(
         return COQ_OUTPUT[key][1]
 
     start = time.time()
-    extra_kwargs = {
-        k: kwargs[k]
-        for k in ["max_mem_rss", "max_mem_as", "cgroup", "mem_limit_method"]
-        if k in kwargs
-    }
+    extra_keys = [
+        "max_mem_rss",
+        "max_mem_as",
+        "cgroup",
+        "mem_limit_method",
+        "max_mem_rss_multiplier",
+        "max_mem_as_multiplier",
+        "memory_usage_key",
+    ]
+    extra_kwargs = {k: kwargs[k] for k in extra_keys if k in kwargs}
+    extra_kwargs.setdefault("memory_usage_key", coqc_prog)
     ((stdout, stderr), returncode, peak_rss_bytes) = (
         memory_robust_timeout_Popen_communicate(
             kwargs["log"],

coq_tools/find_bug.py

@@ -78,12 +78,12 @@
 from .strip_comments import strip_comments
 from .strip_newlines import strip_newlines
 from .util import (
+    MEMORY_LIMIT_METHODS,
     PY3,
     BooleanOptionalAction,
     list_diff,
+    parse_memory_limit_with_multiplier,
     yes_no_prompt,
-    parse_memory_bytes,
-    MEMORY_LIMIT_METHODS,
 )
 
 if PY3:
@@ -4132,6 +4132,17 @@ def prepend_coqbin(prog):
     args = adjust_no_error_defaults(args)
     bug_file_name = args.bug_file.name
     output_file_name = args.output_file
+
+    try:
+        max_mem_rss_value, max_mem_rss_multiplier = parse_memory_limit_with_multiplier(
+            args.max_mem_rss, "--max-mem-rss"
+        )
+        max_mem_as_value, max_mem_as_multiplier = parse_memory_limit_with_multiplier(
+            args.max_mem_as, "--max-mem-as"
+        )
+    except ValueError as exc:
+        parser.error(str(exc))
+
     env = {
         "only_inline": args.only_inline,
         "fast_merge_imports": args.fast_merge_imports,
@@ -4276,14 +4287,10 @@ def prepend_coqbin(prog):
         "add_proof_using_before_admit": args.add_proof_using_before_admit,
         "prefer_final_proof_using": args.prefer_final_proof_using,
         "remove_non_definitions": args.remove_non_definitions,
-        "max_mem_rss": (
-            parse_memory_bytes(args.max_mem_rss)
-            if args.max_mem_rss is not None
-            else None
-        ),
-        "max_mem_as": (
-            parse_memory_bytes(args.max_mem_as) if args.max_mem_as is not None else None
-        ),
+        "max_mem_rss": max_mem_rss_value,
+        "max_mem_as": max_mem_as_value,
+        "max_mem_rss_multiplier": max_mem_rss_multiplier,
+        "max_mem_as_multiplier": max_mem_as_multiplier,
         "cgroup": args.cgroup,
         "mem_limit_method": args.mem_limit_method,
     }
@@ -4332,7 +4339,11 @@ def prepend_coqbin(prog):
                 level=LOG_ALWAYS,
             )
             env["mem_limit_method"] = "prlimit"
-        if env["mem_limit_method"] == "cgexec" and env["cgroup"] is None and env["max_mem_rss"] is None:
+        if (
+            env["mem_limit_method"] == "cgexec"
+            and env["cgroup"] is None
+            and env["max_mem_rss"] is None
+        ):
             env["log"](
                 "\nError: --mem-limit-method=cgexec requires --cgroup or --max-mem-rss.",
                 force_stdout=True,

coq_tools/util.py

@@ -474,7 +474,9 @@ def parse_memory_bytes(mem_str: str) -> int:
     return int(number * multiplier)
 
 
-def wrap_with_prlimit(cmd: List[str], as_bytes: int = None, rss_bytes: int = None) -> List[str]:
+def wrap_with_prlimit(
+    cmd: List[str], as_bytes: int = None, rss_bytes: int = None
+) -> List[str]:
     """
     Wrap command with prlimit to set resource limits.
 
@@ -518,7 +520,9 @@ def wrap_with_ulimit(cmd: List[str], as_bytes: int = None) -> List[str]:
     return ["sh", "-c", f"ulimit -v {kb}; exec {escaped_cmd}"]
 
 
-def wrap_with_cgexec(cmd: List[str], cgroup: str, controllers: str = "memory") -> List[str]:
+def wrap_with_cgexec(
+    cmd: List[str], cgroup: str, controllers: str = "memory"
+) -> List[str]:
     """
     Wrap command with cgexec to run in an existing cgroup.
 
@@ -541,6 +545,35 @@ def wrap_with_cgexec(cmd: List[str], cgroup: str, controllers: str = "memory") -
     return ["cgexec", "-g", f"{controllers}:{cgroup}"] + cmd
 
 
+def parse_memory_limit_with_multiplier(value, arg_name):
+    """returns (absolute_bytes_or_minus_one_or_none, multiplier_or_none)"""
+    if value is None:
+        return None, None
+    value_str = str(value).strip()
+    if not value_str:
+        return None, None
+    lower_val = value_str.lower()
+    if lower_val.endswith("x"):
+        multiplier_str = lower_val[:-1].strip()
+        if not multiplier_str:
+            raise ValueError(
+                f"{arg_name} multiplier must include a numeric value before 'x'"
+            )
+        try:
+            multiplier = float(multiplier_str)
+        except ValueError as exc:
+            raise ValueError(
+                f"{arg_name} multiplier '{value_str}' is not a valid number"
+            ) from exc
+        if multiplier <= 0:
+            raise ValueError(f"{arg_name} multiplier '{value_str}' must be positive")
+        return None, multiplier
+    try:
+        return parse_memory_bytes(value_str), None
+    except ValueError as exc:
+        raise ValueError(f"{arg_name} has invalid value '{value_str}': {exc}") from exc
+
+
 def wrap_with_cgexec_and_create(
     cmd: List[str],
     mem_bytes: int,
@@ -607,10 +640,7 @@ def wrap_with_systemd_run(
     Note:
         Requires systemd and user session (--user) or root privileges.
     """
-    has_limits = any(
-        v is not None and v > 0
-        for v in [mem_bytes, as_bytes, swap_bytes]
-    )
+    has_limits = any(v is not None and v > 0 for v in [mem_bytes, as_bytes, swap_bytes])
     if not has_limits:
         return cmd
 
@@ -687,7 +717,9 @@ def apply_memory_limit(
 
     elif method == "ulimit":
         if rss_bytes is not None and rss_bytes > 0:
-            raise ValueError("ulimit method only supports --max-mem-as, not --max-mem-rss")
+            raise ValueError(
+                "ulimit method only supports --max-mem-as, not --max-mem-rss"
+            )
         return wrap_with_ulimit(cmd, as_bytes=as_bytes), None
 
     elif method == "cgexec":
@@ -702,10 +734,7 @@ def apply_memory_limit(
 
     elif method == "systemd-run":
         return wrap_with_systemd_run(
-            cmd,
-            mem_bytes=rss_bytes,
-            as_bytes=as_bytes,
-            swap_bytes=swap_bytes
+            cmd, mem_bytes=rss_bytes, as_bytes=as_bytes, swap_bytes=swap_bytes
         ), None
 
     else: