bounded_subprocess 2.5.0__tar.gz → 2.7.0__tar.gz

{bounded_subprocess-2.5.0 → bounded_subprocess-2.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bounded_subprocess
-Version: 2.5.0
+Version: 2.7.0
 Summary: A library to facilitate running subprocesses that may misbehave.
 Project-URL: Homepage, https://github.com/arjunguha/bounded_subprocess
 Project-URL: Bug Tracker, https://github.com/arjunguha/bounded_subprocess

{bounded_subprocess-2.5.0 → bounded_subprocess-2.7.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bounded_subprocess"
-version = "2.5.0"
+version = "2.7.0"
 authors = [
   { name="Arjun Guha" },
   { name="Ming-Ho Yee" },

{bounded_subprocess-2.5.0 → bounded_subprocess-2.7.0}/src/bounded_subprocess/bounded_subprocess_async.py

@@ -22,6 +22,133 @@ from .util import (
 logger = logging.getLogger(__name__)
 
 
+def _read_process_group_id(pid: int) -> Optional[int]:
+    """
+    Read a process's PGID from `/proc/<pid>/stat`.
+
+    We use PGID to approximate "all related processes" without requiring
+    cgroups. Parsing `/proc/<pid>/stat` is Linux-specific and somewhat brittle
+    if format assumptions ever change, but it avoids external dependencies and
+    works on typical cluster nodes where cgroups may not be configured for us.
+    """
+    stat_path = f"/proc/{pid}/stat"
+    try:
+        with open(stat_path, "r", encoding="utf-8") as f:
+            stat = f.read().strip()
+    except (FileNotFoundError, ProcessLookupError, PermissionError, OSError):
+        return None
+
+    right_paren = stat.rfind(")")
+    if right_paren == -1:
+        return None
+    rest = stat[right_paren + 1 :].strip().split()
+    # After "comm)", fields are: state, ppid, pgrp, ...
+    if len(rest) < 3:
+        return None
+    try:
+        return int(rest[2])
+    except ValueError:
+        return None
+
+
+def _read_proc_status_kb(pid: int, metric_name: str) -> Optional[int]:
+    """
+    Read a memory metric (in kB) from `/proc/<pid>/status`.
+
+    Returns None when the process is gone or the metric is unavailable.
+    This is intentionally best-effort because `/proc` is racy for short-lived
+    processes and the watchdog tolerates partial visibility.
+    """
+    status_path = f"/proc/{pid}/status"
+    try:
+        with open(status_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if line.startswith(f"{metric_name}:"):
+                    parts = line.split()
+                    if len(parts) >= 2:
+                        return int(parts[1])
+                    return None
+    except (FileNotFoundError, ProcessLookupError, PermissionError, ValueError):
+        return None
+    return None
+
+
+def _read_peak_rss_kb(pid: int) -> Optional[int]:
+    """
+    Read per-process peak resident memory in kB.
+
+    Preferred metric is `VmHWM` (peak RSS); if absent, falls back to `VmRSS`
+    (current RSS). The fallback makes behavior degrade gracefully on kernels
+    that do not expose `VmHWM` consistently, at the cost of weaker "true peak"
+    tracking for those processes.
+    """
+    # VmHWM is the process's peak RSS in kB. Fall back to VmRSS if unavailable.
+    peak = _read_proc_status_kb(pid, "VmHWM")
+    if peak is not None:
+        return peak
+    return _read_proc_status_kb(pid, "VmRSS")
+
+
+def _sum_process_group_peak_rss_kb(process_group_id: int) -> Optional[int]:
+    """
+    Sum per-process peak RSS across all processes in a process group.
+
+    This is a deliberate non-cgroup approximation for aggregate memory. It
+    overcounts shared pages (RSS semantics), can miss processes that start/exit
+    between scans, and is O(number of /proc entries) per poll. Those tradeoffs
+    are accepted here to keep enforcement portable on cluster environments
+    where we cannot assume cgroup control.
+    """
+    total_kb = 0
+    found_any = False
+    for entry in os.scandir("/proc"):
+        if not entry.name.isdigit():
+            continue
+        pid = int(entry.name)
+        pgid = _read_process_group_id(pid)
+        if pgid != process_group_id:
+            continue
+        peak_kb = _read_peak_rss_kb(pid)
+        if peak_kb is not None:
+            total_kb += peak_kb
+            found_any = True
+    return total_kb if found_any else None
+
+
+async def _memory_watchdog(
+    p: subprocess.Popen,
+    process_group_id: int,
+    deadline: float,
+    memory_limit_mb: int,
+    memory_watchdog_interval_seconds: float,
+) -> bool:
+    """
+    Enforce `memory_limit_mb` using periodic process-group peak-RSS checks.
+
+    The watchdog polls aggregate process-group peak RSS and kills the whole
+    process group when the limit is exceeded, returning True in that case.
+    Polling is interval-based (default 1s, configurable) and intentionally
+    approximate: it favors broad compatibility without cgroups over perfect
+    accounting precision.
+    """
+    limit_kb = memory_limit_mb * 1024
+    check_interval = max(0.01, memory_watchdog_interval_seconds)
+    while True:
+        if p.poll() is not None:
+            return False
+        aggregate_peak_rss_kb = _sum_process_group_peak_rss_kb(process_group_id)
+        if aggregate_peak_rss_kb is not None and aggregate_peak_rss_kb > limit_kb:
+            try:
+                os.killpg(process_group_id, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+            return True
+        remaining = deadline - time.time()
+        if remaining <= 0:
+            return False
+        await asyncio.sleep(min(check_interval, remaining))
+
+
 async def run(
     args: List[str],
     timeout_seconds: int = 15,
@@ -29,6 +156,8 @@ async def run(
     env=None,
     stdin_data: Optional[str] = None,
     stdin_write_timeout: Optional[int] = None,
+    memory_limit_mb: Optional[int] = None,
+    memory_watchdog_interval_seconds: float = 1.0,
 ) -> Result:
     """
     Run a subprocess asynchronously with bounded stdout/stderr capture.
@@ -38,7 +167,10 @@ async def run(
     truncated to `max_output_size` bytes each. If the timeout elapses,
     `Result.timeout` is True and `Result.exit_code` is -1. If `stdin_data`
     cannot be fully written before `stdin_write_timeout`, `Result.exit_code`
-    is set to -1 even if the process exits normally.
+    is set to -1 even if the process exits normally. If `memory_limit_mb` is
+    set, a watchdog checks aggregate peak RSS (`VmHWM`, summed across the
+    process group) at a fixed interval and kills the process group when the
+    limit is exceeded.
 
     Example:
 
@@ -91,6 +223,18 @@ async def run(
         except (BrokenPipeError, BlockingIOError):
             pass
 
+    memory_watchdog_task = None
+    if memory_limit_mb is not None:
+        memory_watchdog_task = asyncio.create_task(
+            _memory_watchdog(
+                p=p,
+                process_group_id=process_group_id,
+                deadline=deadline,
+                memory_limit_mb=memory_limit_mb,
+                memory_watchdog_interval_seconds=memory_watchdog_interval_seconds,
+            )
+        )
+
     bufs = await read_to_eof_async(
         [p.stdout, p.stderr],
         timeout_seconds=timeout_seconds,
@@ -110,13 +254,28 @@ async def run(
             break
         await asyncio.sleep(min(0.05, remaining))
 
+    memory_limit_exceeded = False
+    if memory_watchdog_task is not None:
+        if memory_watchdog_task.done():
+            memory_limit_exceeded = memory_watchdog_task.result()
+        else:
+            memory_watchdog_task.cancel()
+            try:
+                await memory_watchdog_task
+            except asyncio.CancelledError:
+                pass
+
     try:
         os.killpg(process_group_id, signal.SIGKILL)
     except ProcessLookupError:
         pass
 
     exit_code = (
-        -1 if is_timeout or (stdin_data is not None and not write_ok) else exit_code
+        -1
+        if is_timeout
+        or (stdin_data is not None and not write_ok)
+        or memory_limit_exceeded
+        else exit_code
     )
 
     return Result(
@@ -165,6 +324,8 @@ async def podman_run(
     stdin_write_timeout: Optional[int] = None,
     volumes: List[str] = [],
     cwd: Optional[str] = None,
+    memory_limit_mb: Optional[int] = None,
+    entrypoint: Optional[str] = None,
 ) -> Result:
     """
     Run a subprocess in a podman container asynchronously with bounded stdout/stderr capture.
@@ -184,6 +345,12 @@ async def podman_run(
         stdin_write_timeout: Optional timeout for writing stdin data.
         volumes: Optional list of volume mount specifications (e.g., ["/host/path:/container/path"]).
         cwd: Optional working directory path inside the container.
+        memory_limit_mb: Optional memory limit in megabytes for the container.
+        entrypoint: Optional override for the container image's ENTRYPOINT.
+            Passed through to podman's `--entrypoint` flag. Pass `""` to clear
+            the image's ENTRYPOINT entirely (so `args` is interpreted as the
+            full command). When `None`, the flag is omitted and the image's
+            ENTRYPOINT is used unchanged.
 
     Example:
 
@@ -228,10 +395,19 @@ async def podman_run(
     for volume in volumes:
         podman_args.extend(["-v", volume])
 
+    # Handle memory limit
+    if memory_limit_mb is not None:
+        podman_args.extend(["--memory", f"{memory_limit_mb}m", "--memory-swap", f"{memory_limit_mb}m"])
+
     # Handle working directory
     if cwd is not None:
         podman_args.extend(["-w", cwd])
 
+    # Handle entrypoint override. Note: `entrypoint=""` is meaningful — it
+    # clears the image's ENTRYPOINT — so we test against None, not falsiness.
+    if entrypoint is not None:
+        podman_args.extend(["--entrypoint", entrypoint])
+
     podman_args.append(image)
     podman_args.extend(args)
 

{bounded_subprocess-2.5.0 → bounded_subprocess-2.7.0}/test/test_async.py

@@ -126,6 +126,59 @@ async def test_read_one_line():
     await assert_no_running_evil()
 
 
+@pytest.mark.asyncio
+async def test_run_memory_limit_ok():
+    result = await run(
+        ["python3", "-c", "import time; x = bytearray(8 * 1024 * 1024); time.sleep(0.5)"],
+        timeout_seconds=5,
+        max_output_size=1024,
+        memory_limit_mb=256,
+        memory_watchdog_interval_seconds=0.05,
+    )
+    assert result.exit_code == 0
+    assert result.timeout is False
+    assert len(result.stderr) == 0
+
+
+@pytest.mark.asyncio
+async def test_run_memory_limit_exceeded():
+    result = await run(
+        ["python3", "-c", "import time; x = bytearray(128 * 1024 * 1024); time.sleep(5)"],
+        timeout_seconds=10,
+        max_output_size=1024,
+        memory_limit_mb=64,
+        memory_watchdog_interval_seconds=0.05,
+    )
+    assert result.exit_code == -1
+    assert result.timeout is False
+
+
+@pytest.mark.asyncio
+async def test_run_memory_limit_exceeded_in_child_process():
+    result = await run(
+        [
+            "python3",
+            "-c",
+            (
+                "import os, time\n"
+                "pid = os.fork()\n"
+                "_ = bytearray(4 * 1024 * 1024)\n"
+                "if pid == 0:\n"
+                "    _ = bytearray(128 * 1024 * 1024)\n"
+                "    time.sleep(5)\n"
+                "else:\n"
+                "    time.sleep(5)\n"
+            ),
+        ],
+        timeout_seconds=10,
+        max_output_size=1024,
+        memory_limit_mb=64,
+        memory_watchdog_interval_seconds=0.05,
+    )
+    assert result.exit_code == -1
+    assert result.timeout is False
+
+
 @pytest.mark.asyncio
 async def test_podman_run_stdin():
     """Test podman_run with stdin_data input."""
@@ -216,3 +269,32 @@ async def test_podman_run_cwd():
     assert result.exit_code == 0
     assert result.timeout is False
     assert result.stdout.strip() == "/tmp"
+
+
+@pytest.mark.asyncio
+async def test_podman_run_memory_limit_ok():
+    """Test podman_run with memory limit that is sufficient."""
+    result = await podman_run(
+        ["echo", "hello"],
+        image="alpine:latest",
+        timeout_seconds=10,
+        max_output_size=1024,
+        memory_limit_mb=64,
+    )
+    assert result.exit_code == 0
+    assert result.timeout is False
+    assert result.stdout.strip() == "hello"
+
+
+@pytest.mark.asyncio
+async def test_podman_run_memory_limit_oom():
+    """Test podman_run where the process exceeds the memory limit."""
+    # Allocate ~128MB in a container limited to 32MB
+    result = await podman_run(
+        ["python3", "-c", "x = bytearray(128 * 1024 * 1024)"],
+        image="python:3.12-alpine",
+        timeout_seconds=30,
+        max_output_size=4096,
+        memory_limit_mb=32,
+    )
+    assert result.exit_code != 0

{bounded_subprocess-2.5.0 → bounded_subprocess-2.7.0}/uv.lock

@@ -1,5 +1,5 @@
 version = 1
-revision = 2
+revision = 3
 requires-python = ">=3.9"
 resolution-markers = [
     "python_full_version >= '3.10'",
@@ -31,7 +31,7 @@ wheels = [
 
 [[package]]
 name = "bounded-subprocess"
-version = "2.5.0"
+version = "2.7.0"
 source = { editable = "." }
 dependencies = [
     { name = "typeguard" },

bounded_subprocess-2.5.0/.git

@@ -1 +0,0 @@
-gitdir: /media/external0/arjun-nosudo/repos/arjunguha/bounded_subprocess/bare/worktrees/main