Add hierarchical cgroups and pattern-specific cgroup kill

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

src/branching/agent/patterns.py

@@ -37,6 +37,7 @@ def __init__(
         isolate_processes: bool = False,
         timeout: float | None = None,
         resource_limits: ResourceLimits | None = None,
+        group_limits: ResourceLimits | None = None,
     ):
         """
         Args:
@@ -49,20 +50,56 @@ def __init__(
             timeout: Overall timeout in seconds for all candidates.
             resource_limits: Optional per-branch resource limits (implies
                 process isolation).
+            group_limits: Optional resource limits applied to the root
+                cgroup that contains all branches.
         """
         self._candidates = list(candidates)
         self._first_wins = first_wins
         self._max_parallel = max_parallel or len(self._candidates)
         self._isolate_processes = isolate_processes
         self._timeout = timeout
         self._resource_limits = resource_limits
+        self._group_limits = group_limits
 
     def __call__(self, workspace: Workspace) -> SpeculationOutcome:
+        import os as _os
+
+        root_cgroup: Optional[Path] = None
+        if self._resource_limits is not None and self._group_limits is not None:
+            try:
+                from ..process._cgroup import create_group
+                root_cgroup = create_group(
+                    f"speculate-{_os.getpid()}",
+                    limits=self._group_limits,
+                )
+            except OSError:
+                root_cgroup = None
+
+        try:
+            return self._run(workspace, root_cgroup)
+        finally:
+            if root_cgroup is not None:
+                from ..process._cgroup import kill_scope
+                kill_scope(root_cgroup)
+
+    def _run(self, workspace: Workspace, root_cgroup: Optional[Path]) -> SpeculationOutcome:
         results: list[SpeculationResult] = [None] * len(self._candidates)  # type: ignore
         winner: Optional[SpeculationResult] = None
         committed = False
         cancel_event = threading.Event()
 
+        # Track live cgroup scope paths so we can kill losers immediately.
+        # Populated by scope_callback from run_in_process; read by
+        # _kill_scopes.  dict ops are GIL-protected in CPython.
+        branch_scopes: dict[int, Path] = {}
+
+        def _kill_scopes(exclude: int = -1) -> None:
+            """Kill all tracked cgroup scopes except *exclude*."""
+            from ..process._cgroup import kill_scope
+            for idx, scope in list(branch_scopes.items()):
+                if idx != exclude:
+                    kill_scope(scope)
+
         def _run_candidate(index: int) -> SpeculationResult:
             branch_name = f"speculate_{index}"
             result = SpeculationResult(branch_index=index, success=False)
@@ -81,7 +118,13 @@ def _run_candidate(index: int) -> SpeculationResult:
                         return result
 
                     if self._resource_limits is not None:
-                        success = self._run_with_limits(b.path, index)
+                        def _on_scope(scope_path: Path, _i: int = index) -> None:
+                            branch_scopes[_i] = scope_path
+
+                        success = self._run_with_limits(
+                            b.path, index, root_cgroup,
+                            scope_callback=_on_scope,
+                        )
                     elif self._isolate_processes:
                         success = self._run_in_process(b.path, index)
                     else:
@@ -93,6 +136,7 @@ def _run_candidate(index: int) -> SpeculationResult:
                     if result.success and not cancel_event.is_set():
                         if self._first_wins:
                             cancel_event.set()
+                            _kill_scopes(index)
                         try:
                             b.commit()
                         except ConflictError:
@@ -131,6 +175,7 @@ def _run_candidate(index: int) -> SpeculationResult:
             except TimeoutError:
                 # Signal remaining candidates to abort
                 cancel_event.set()
+                _kill_scopes()
 
             # Wait briefly for in-flight branches to finish cleanup
             for f in futures:
@@ -173,7 +218,10 @@ def target(workspace: Path) -> None:
             except ProcessBranchError:
                 return False
 
-    def _run_with_limits(self, path: Path, index: int) -> bool:
+    def _run_with_limits(
+        self, path: Path, index: int, parent_cgroup: Path | None = None,
+        scope_callback=None,
+    ) -> bool:
         """Run a candidate in a forked child with resource limits."""
         from ..process.runner import run_in_process
         from ..exceptions import ProcessBranchError
@@ -191,6 +239,8 @@ def _run_with_limits(self, path: Path, index: int) -> bool:
                 workspace=path,
                 limits=self._resource_limits,
                 timeout=per_candidate,
+                parent_cgroup=parent_cgroup,
+                scope_callback=scope_callback,
             )
             return bool(result)
         except ProcessBranchError:

src/branching/agent/speculate.py

@@ -33,28 +33,89 @@ def _own_cgroup() -> Path:
     raise OSError("No cgroup v2 entry found in /proc/self/cgroup")
 
 
-def create_scope(name: str) -> Path:
-    """Create a child cgroup under our own cgroup for process grouping.
+def _enable_subtree_controllers(cgroup_dir: Path) -> None:
+    """Enable memory and cpu controllers for children of *cgroup_dir*.
 
-    Creates under the current process's cgroup, which is compatible
-    with systemd's delegated hierarchy — we never write into a level
-    that systemd manages directly.
+    Reads ``cgroup.controllers`` to discover available controllers, then
+    writes ``+memory +cpu`` (intersection with available) to
+    ``cgroup.subtree_control``.  Best-effort — silently ignores errors.
+    """
+    try:
+        available = (cgroup_dir / "cgroup.controllers").read_text().split()
+    except OSError:
+        return
+    wanted = [c for c in ("memory", "cpu") if c in available]
+    if not wanted:
+        return
+    payload = " ".join(f"+{c}" for c in wanted)
+    try:
+        (cgroup_dir / "cgroup.subtree_control").write_text(payload)
+    except OSError:
+        pass
+
+
+def create_scope(name: str, *, parent: Path | None = None) -> Path:
+    """Create a child cgroup under *parent* (or our own cgroup) for process grouping.
+
+    When *parent* is given the scope is created as a child of that
+    directory and ``_enable_subtree_controllers`` is called on the parent
+    first so that memory/cpu controllers are available in the child.
 
     Args:
         name: Scope name suffix (e.g. PID).
+        parent: Optional parent cgroup directory.  Defaults to
+            ``_own_cgroup()`` when ``None``.
 
     Returns:
         Path to the cgroup directory.
 
     Raises:
         OSError: If cgroup creation fails.
     """
-    parent = _own_cgroup()
-    scope_dir = parent / f"branching-{name}.scope"
+    if parent is not None:
+        _enable_subtree_controllers(parent)
+        parent_dir = parent
+    else:
+        parent_dir = _own_cgroup()
+    scope_dir = parent_dir / f"branching-{name}.scope"
     scope_dir.mkdir(exist_ok=True)
     return scope_dir
 
 
+def create_group(
+    name: str,
+    *,
+    parent: Path | None = None,
+    limits: "ResourceLimits | None" = None,
+) -> Path:
+    """Create an intermediate cgroup for nesting (never holds PIDs directly).
+
+    Calls ``_enable_subtree_controllers`` on the new directory so that
+    child scopes can use memory/cpu controllers.  Applies optional
+    group-level limits (total budget for all children).
+
+    Args:
+        name: Group name suffix.
+        parent: Optional parent cgroup directory.  Defaults to
+            ``_own_cgroup()`` when ``None``.
+        limits: Optional resource limits applied to the group itself.
+
+    Returns:
+        Path to the new group cgroup directory.
+    """
+    if parent is not None:
+        _enable_subtree_controllers(parent)
+        parent_dir = parent
+    else:
+        parent_dir = _own_cgroup()
+    group_dir = parent_dir / f"branching-{name}.scope"
+    group_dir.mkdir(exist_ok=True)
+    _enable_subtree_controllers(group_dir)
+    if limits is not None:
+        set_limits(group_dir, limits)
+    return group_dir
+
+
 def add_pid(scope: Path, pid: int) -> None:
     """Add a PID to a cgroup scope.
 
@@ -68,7 +129,7 @@ def add_pid(scope: Path, pid: int) -> None:
     (scope / "cgroup.procs").write_text(str(pid))
 
 
-def set_limits(scope: Path, limits: ResourceLimits) -> None:
+def set_limits(scope: Path, limits: "ResourceLimits") -> None:
     """Apply resource limits to a cgroup scope.
 
     Best-effort — skips ``None`` fields and ignores write errors so that
@@ -79,6 +140,16 @@ def set_limits(scope: Path, limits: ResourceLimits) -> None:
             (scope / "memory.max").write_text(str(limits.memory))
         except OSError:
             pass
+    if limits.memory_high is not None:
+        try:
+            (scope / "memory.high").write_text(str(limits.memory_high))
+        except OSError:
+            pass
+    if limits.oom_group:
+        try:
+            (scope / "memory.oom.group").write_text("1")
+        except OSError:
+            pass
     if limits.cpu is not None:
         period = 100_000  # 100 ms
         quota = int(limits.cpu * period)
@@ -89,11 +160,21 @@ def set_limits(scope: Path, limits: ResourceLimits) -> None:
 
 
 def kill_scope(scope: Path) -> None:
-    """Kill all processes in a cgroup scope and remove it.
+    """Kill all processes in a cgroup scope and remove it recursively.
+
+    Iterates child directories bottom-up, kills each, then kills and
+    removes self.  Handles the cgroup v2 requirement that ``rmdir``
+    needs no child cgroups.
 
     Best-effort cleanup — ignores errors.
     """
     try:
+        # Recurse into children first (bottom-up)
+        if scope.is_dir():
+            for child in sorted(scope.iterdir()):
+                if child.is_dir():
+                    kill_scope(child)
+        # Kill processes in this scope
         kill_file = scope / "cgroup.kill"
         if kill_file.exists():
             kill_file.write_text("1")

src/branching/process/_cgroup.py

@@ -42,6 +42,7 @@ def __init__(
         isolate: bool = False,
         close_fds: bool = False,
         limits: ResourceLimits | None = None,
+        parent_cgroup: Path | None = None,
     ):
         """
         Args:
@@ -52,12 +53,16 @@ def __init__(
                      since each child needs its own user ns for bind-mount).
             close_fds: BR_CLOSE_FDS — close inherited fds (3+) in child.
             limits: Optional resource limits applied via cgroup v2.
+            parent_cgroup: Optional parent cgroup directory.  When given,
+                the child's scope is created under this directory instead
+                of the process's own cgroup, enabling hierarchical nesting.
         """
         self._target = target
         self._workspace = workspace
         self._isolate = isolate
         self._close_fds = close_fds
         self._limits = limits
+        self._parent_cgroup = parent_cgroup
         self._pid: Optional[int] = None
         self._exited = False
         self._cgroup_scope: Optional[Path] = None
@@ -68,6 +73,11 @@ def pid(self) -> int:
             raise ProcessBranchError("Process not started")
         return self._pid
 
+    @property
+    def cgroup_scope(self) -> Optional[Path]:
+        """The cgroup v2 scope directory, or ``None`` if unavailable."""
+        return self._cgroup_scope
+
     @property
     def alive(self) -> bool:
         if self._pid is None or self._exited:
@@ -175,7 +185,9 @@ def __enter__(self) -> "BranchContext":
         # (required for per-branch limits).
         scope_name = f"{os.getpid()}-{uuid.uuid4().hex[:8]}"
         try:
-            self._cgroup_scope = _cgroup.create_scope(scope_name)
+            self._cgroup_scope = _cgroup.create_scope(
+                scope_name, parent=self._parent_cgroup,
+            )
         except OSError:
             self._cgroup_scope = None
 
@@ -253,6 +265,7 @@ def create(
         isolate: bool = False,
         close_fds: bool = False,
         limits: ResourceLimits | None = None,
+        parent_cgroup: Path | None = None,
     ) -> Iterator[list["BranchContext"]]:
         """Create N branch contexts, mirroring branch(BR_CREATE, n_branches=N).
 
@@ -264,6 +277,7 @@ def create(
             isolate: BR_ISOLATE — separate user ns per child.
             close_fds: BR_CLOSE_FDS — close inherited fds in children.
             limits: Optional resource limits applied via cgroup v2.
+            parent_cgroup: Optional parent cgroup for hierarchical nesting.
 
         Yields:
             List of entered BranchContext instances (already forked).
@@ -276,7 +290,7 @@ def create(
             for target, workspace in zip(targets, workspaces):
                 ctx = BranchContext(
                     target, workspace, isolate=isolate, close_fds=close_fds,
-                    limits=limits,
+                    limits=limits, parent_cgroup=parent_cgroup,
                 )
                 ctx.__enter__()
                 contexts.append(ctx)

src/branching/process/context.py

@@ -52,3 +52,17 @@ class ResourceLimits:
 
     cpu: float | None = None
     """Fraction of one CPU (0.5 = 50%).  Written to ``cpu.max``."""
+
+    memory_high: int | None = None
+    """Soft memory throttle in bytes (written to ``memory.high``).
+
+    When usage exceeds this value the kernel reclaims aggressively but
+    does *not* OOM-kill the process.
+    """
+
+    oom_group: bool = False
+    """Atomic OOM termination (written to ``memory.oom.group``).
+
+    When ``True``, all processes in the cgroup are killed together on
+    OOM rather than picking a single victim.
+    """

src/branching/process/limits.py

@@ -21,6 +21,8 @@ def run_in_process(
     *,
     limits: ResourceLimits | None = None,
     timeout: float | None = None,
+    parent_cgroup: Path | None = None,
+    scope_callback: Callable[[Path], None] | None = None,
 ) -> Any:
     """Run *fn(*args)* in a forked child process, optionally with cgroup limits.
 
@@ -34,6 +36,10 @@ def run_in_process(
             BranchContext workspace).
         limits: Optional resource limits applied to the child's cgroup.
         timeout: Maximum seconds to wait for the child.
+        parent_cgroup: Optional parent cgroup for hierarchical nesting.
+        scope_callback: Optional callback invoked with the cgroup scope path
+            after the child's scope is created.  Allows callers to track live
+            cgroup paths for external kill/throttle.
 
     Returns:
         Whatever *fn* returned.
@@ -61,7 +67,12 @@ def _target(ws_path: Path) -> None:
                 pass
             raise
 
-    with BranchContext(_target, workspace=workspace, limits=limits) as ctx:
+    with BranchContext(
+        _target, workspace=workspace, limits=limits,
+        parent_cgroup=parent_cgroup,
+    ) as ctx:
+        if scope_callback is not None and ctx.cgroup_scope is not None:
+            scope_callback(ctx.cgroup_scope)
         try:
             ctx.wait(timeout=timeout)
         except ProcessBranchError: