apodex/frontier-agent-demo

1"""Per-exec cgroup v2 isolation — consumer side of the worker-shell contract."""2 3from __future__ import annotations4 5import contextlib6import logging7import os8import threading9import time10import uuid11from pathlib import Path12 13logger = logging.getLogger(__name__)14 15_ROOT_ENV = "WORKER_EXEC_CGROUP_ROOT"16MEM_MAX_ENV = "WORKER_EXEC_MEM_MAX_BYTES"17_PIDS_MAX_ENV = "WORKER_EXEC_PIDS_MAX"18 19# Probe-once cache, mirroring bwrap_available() / data_cap_effective():20# None = not probed yet; "" = probed and unavailable.21_ROOT_CACHE: str | None = None22_ROOT_PROBED = False23_ROOT_PROBE_LOCK = threading.Lock()24_MKDIR_FAIL_WARNED = False25 26 27def exec_cgroup_root() -> str | None:28    """The delegated ``execs/`` directory, or ``None`` when the feature is dark.29 30    Probed once and cached: env unset is the normal fail-open state on every31    deployment that predates the worker-shell contract and stays silent; env32    set but unusable is a platform-side misconfiguration and warns once.33 34    The probe is locked and ``_ROOT_PROBED`` is published LAST. Setting it35    before the isdir/access syscalls (which release the GIL) let a concurrent36    cold-start exec read PROBED=True with the cache still empty and silently37    run without its cgroup — one uncontained command is exactly the hole this38    module exists to close. The unlocked fast path is safe under the GIL39    because the cache is complete by the time the flag is visible.40    """41    global _ROOT_CACHE, _ROOT_PROBED42    if _ROOT_PROBED:43        return _ROOT_CACHE or None44    with _ROOT_PROBE_LOCK:45        if _ROOT_PROBED:46            return _ROOT_CACHE or None47        cache = ""48        raw = (os.environ.get(_ROOT_ENV) or "").strip()49        if raw:50            if os.path.isdir(raw) and os.access(raw, os.W_OK):51                cache = raw52                logger.info("per-exec cgroup isolation armed at %s", raw)53            else:54                logger.warning(55                    "%s=%r is set but not a writable directory; per-exec cgroup "56                    "isolation stays OFF and commands run with only the "57                    "per-process ulimit cap", _ROOT_ENV, raw,58                )59        _ROOT_CACHE = cache60        _ROOT_PROBED = True61        return cache or None62 63 64def _env_int(name: str) -> int | None:65    raw = (os.environ.get(name) or "").strip()66    try:67        value = int(raw)68    except ValueError:69        return None70    return value if value > 0 else None71 72 73class ExecCgroup:74    """One created ``exec-<id>/`` directory. Kill and remove via :meth:`close`."""75 76    __slots__ = ("path",)77 78    def __init__(self, path: Path) -> None:79        self.path = path80 81    @property82    def procs_path(self) -> str:83        """The ``cgroup.procs`` file the exec's shell writes ``$$`` into."""84        return str(self.path / "cgroup.procs")85 86    @property87    def current_path(self) -> str:88        """``memory.current`` — instantaneous usage of the whole exec tree."""89        return str(self.path / "memory.current")90 91    @property92    def kill_path(self) -> str:93        return str(self.path / "cgroup.kill")94 95    def _read_int(self, name: str) -> int | None:96        try:97            return int((self.path / name).read_text().strip())98        except (OSError, ValueError):99            return None100 101    def oom_kills(self) -> int:102        """Kernel OOM kills charged to THIS exec's cgroup.103 104        Reads the exec cgroup's own ``memory.events`` — never the container's,105        whose ``oom_kill`` counter is RECURSIVE and accumulates every per-exec106        kill below it (a healthy container reads as OOMKilled after two of107        them; measured on a real worker node).108        """109        try:110            body = (self.path / "memory.events").read_text()111        except OSError:112            return 0113        total = 0114        for line in body.splitlines():115            fields = line.split()116            if len(fields) == 2 and fields[0] in ("oom_kill", "oom_group_kill"):117                with contextlib.suppress(ValueError):118                    total = max(total, int(fields[1]))119        return total120 121    def oom_note(self) -> str | None:122        """Human-readable account of a group kill, or ``None`` if none happened.123 124        The whole point of synthesising this: ``memory.oom.group=1`` SIGKILLs125        the shell along with the runaway, so the model sees a bare exit 137126        with no MemoryError and no traceback — nothing it could self-correct127        from. The cgroup still holds the numbers until we rmdir it.128        """129        if self.oom_kills() <= 0:130            return None131        # memory.peak saturates AT the limit (the allocation that would pass it132        # fails), so peak is a floor of what the tree wanted, not an overshoot.133        peak = self._read_int("memory.peak")134        limit = self._read_int("memory.max")135        peak_s = f" (peak usage {peak // (1024 * 1024)}MB)" if peak else ""136        limit_s = f"its {limit // (1024 * 1024)}MB" if limit else "its"137        return (138            f"[memory limit] this command's process tree hit {limit_s} "139            f"per-command memory limit{peak_s} and was killed as a group. "140            "The limit counts ALL its processes together: lower the "141            "parallelism (each concurrent worker holds its own copy), process "142            "the data in chunks, and avoid large anonymous mmap regions."143        )144 145    def kill(self) -> None:146        """Atomically SIGKILL every process in the exec tree (``cgroup.kill``).147 148        Unlike ``killpg`` this also reaches descendants that left the process149        group via ``setsid`` / ``start_new_session`` — cgroup membership is150        inherited across fork/exec and cannot be shed from inside.151        """152        with contextlib.suppress(OSError), open(self.kill_path, "w") as fh:153            fh.write("1")154 155    def close(self) -> None:156        """Kill the tree and remove the directory. Best effort, never raises.157 158        The rmdir is NOT optional hygiene: an exec cgroup pins node-level slab159        memory that no container limit accounts for, and a worker leaking one160        per exec eventually takes the NODE NotReady (kubernetes KEP-5474161        measured ~42k leaked groups exhausting 14GB). The worker shell sweeps162        leftovers at task end as a backstop for SIGKILL'd harnesses, but the163        normal path is this method, on the exec's own ``finally``.164        """165        self.kill()166        # rmdir succeeds only once every member process is gone; SIGKILL'd167        # processes can take a moment to be reaped, hence the short retry.168        for _ in range(10):169            try:170                self.path.rmdir()171                return172            except OSError:173                time.sleep(0.05)174        logger.warning(175            "exec cgroup %s not removable after kill; leaving it to the "176            "worker-shell sweep", self.path,177        )178 179 180def create_exec_cgroup() -> ExecCgroup | None:181    """Create and configure one ``exec-<id>/`` cgroup, or ``None`` (fail-open).182 183    Limits come from the platform-provided env verbatim; a missing variable184    skips its file rather than inventing a value. Each write is suppressed185    individually — one controller not being enabled must not cost the exec the186    others.187 188    ``memory.high`` is deliberately NOT set even though the platform suggests189    a value: it throttles instead of killing, which turns an over-memory190    ``soffice`` into a silent tool timeout — a harder signal for the model191    than a group kill with a synthesized explanation. Revisit once192    ``memory.events`` data from production says otherwise.193    """194    global _MKDIR_FAIL_WARNED195    root = exec_cgroup_root()196    if root is None:197        return None198    path = Path(root) / f"exec-{os.getpid()}-{uuid.uuid4().hex[:8]}"199    try:200        # exist_ok=False: names must never collide across threads — landing two201        # execs in one cgroup would kill both when either overruns.202        path.mkdir(exist_ok=False)203    except OSError as exc:204        # Warn once per failure streak: a persistent cause (cgroupfs turned205        # read-only, execs/ swept away) would otherwise log every exec.206        if not _MKDIR_FAIL_WARNED:207            _MKDIR_FAIL_WARNED = True208            logger.warning("per-exec cgroup mkdir failed (%s); running uncontained", exc)209        return None210    _MKDIR_FAIL_WARNED = False211    cg = ExecCgroup(path)212    mem_max = _env_int(MEM_MAX_ENV)213    settings: list[tuple[str, str]] = []214    if mem_max:215        settings.append(("memory.max", str(mem_max)))216        # Load-bearing, not hygiene: with swap available the kernel swaps217        # anonymous pages out instead of ever reaching memory.max, so nothing218        # is killed AND MAP_SHARED writes keep landing (reproduced under219        # Docker Desktop). EKS nodes have no swap today; do not rely on that.220        settings.append(("memory.swap.max", "0"))221        # Make THIS cgroup the oom_domain: the kernel kills the whole exec222        # tree here and never walks up to the container's oom.group=1.223        settings.append(("memory.oom.group", "1"))224    pids_max = _env_int(_PIDS_MAX_ENV)225    if pids_max:226        settings.append(("pids.max", str(pids_max)))227    for name, value in settings:228        try:229            (path / name).write_text(value)230        except OSError as exc:231            logger.warning("per-exec cgroup: writing %s=%s failed: %s", name, value, exc)232    return cg233