From 04ed5dac699a0bbb43171889c8ec1a209c45bb65 Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Thu, 17 Sep 2026 20:28:48 +0300 Subject: [PATCH] fix(runtime): harden process-ownership boundary, adoption lifelines, and power/vram guards --- prime-agent-runtime/src/rlm/bash.py | 6 + prime-agent-runtime/src/rlm/host_safety.py | 499 ++++++++++++++++--- prime-agent-runtime/test/test_host_safety.py | 143 ++++++ 3 files changed, 590 insertions(+), 58 deletions(-) diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 25b6d8ce5..94d2f59d9 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -2714,6 +2714,12 @@ def _signal_group(pid: int, sig: int) -> bool: """ if not _pid_owns_group(pid): return False + from .host_safety import _is_group_fully_owned, get_owned_process + + lease = get_owned_process(pid) + if lease is not None and not lease.created_by_session: + if not _is_group_fully_owned(pid): + return False try: os.killpg(pid, sig) except ProcessLookupError: diff --git a/prime-agent-runtime/src/rlm/host_safety.py b/prime-agent-runtime/src/rlm/host_safety.py index 0dcdc8fe5..e5fbc4c4e 100644 --- a/prime-agent-runtime/src/rlm/host_safety.py +++ b/prime-agent-runtime/src/rlm/host_safety.py @@ -13,6 +13,7 @@ from __future__ import annotations import os import re import signal +import subprocess import sys import threading import time @@ -24,31 +25,108 @@ from typing import Any # Process Ownership Registry & Scoped Process Control # --------------------------------------------------------------------------- -import subprocess + +@dataclass(frozen=True) +class ProcessIdentity: + pid: int + start_time: str + comm: str = "" + + def matches(self, other: ProcessIdentity | None) -> bool: + if other is None or self.pid != other.pid: + return False + if self.start_time != "unknown" and other.start_time != "unknown": + return self.start_time == other.start_time + if self.comm and other.comm: + return ( + self.comm == other.comm + or self.comm.endswith(other.comm) + or other.comm.endswith(self.comm) + ) + return True + + +def get_process_identity(pid: int) -> ProcessIdentity | None: + """Return identity of the live process (start time + comm), or None if dead/inaccessible.""" + if pid <= 1: + return None + + # 1. Linux /proc//stat + proc_stat = Path(f"/proc/{pid}/stat") + if proc_stat.exists(): + try: + content = proc_stat.read_text() + rparen = content.rfind(")") + lparen = content.find("(") + if lparen != -1 and rparen != -1 and rparen > lparen: + comm = content[lparen + 1 : rparen] + rest = content[rparen + 2 :].split() + if len(rest) > 19: + return ProcessIdentity(pid=pid, start_time=rest[19], comm=comm) + except (OSError, ValueError, IndexError): + pass + + # 2. POSIX / macOS fallback: ps -o lstart=,comm= -p + try: + res = subprocess.run( + ["ps", "-o", "lstart=,comm=", "-p", str(pid)], + capture_output=True, + text=True, + timeout=1.5, + ) + if res.returncode == 0 and res.stdout.strip(): + line = res.stdout.strip() + if len(line) >= 24: + return ProcessIdentity( + pid=pid, + start_time=line[:24].strip(), + comm=line[24:].strip(), + ) + return ProcessIdentity(pid=pid, start_time=line, comm="") + except Exception: + pass + + # 3. Fallback when neither /proc nor ps is available: probe live status via signal 0 + try: + os.kill(pid, 0) + return ProcessIdentity(pid=pid, start_time="unknown", comm="") + except (ProcessLookupError, PermissionError): + pass + except OSError: + pass + + return None + @dataclass class ProcessLease: pid: int command: str started_at: float + identity: ProcessIdentity | None = None handle: Any | None = None reaped: bool = False reason: str | None = None + created_by_session: bool = True _OWNED_LEASES: dict[int, ProcessLease] = {} +_ANCESTRY_CACHE: dict[int, tuple[ProcessIdentity, int, ProcessIdentity | None]] = {} _LEASE_LOCK = threading.Lock() def register_owned_process(pid: int, command: str, handle: Any | None = None) -> ProcessLease: """Enroll a process started by this session into the ownership registry.""" + ident = get_process_identity(pid) with _LEASE_LOCK: lease = ProcessLease( pid=pid, command=command, started_at=time.time(), + identity=ident, handle=handle, reaped=False, + created_by_session=True, ) _OWNED_LEASES[pid] = lease return lease @@ -60,6 +138,7 @@ def mark_owned_process_reaped(pid: int) -> None: lease = _OWNED_LEASES.get(pid) if lease is not None: lease.reaped = True + _ANCESTRY_CACHE.pop(pid, None) def _get_parent_pid(pid: int) -> int | None: @@ -79,7 +158,7 @@ def _get_parent_pid(pid: int) -> int | None: ["ps", "-o", "ppid=", "-p", str(pid)], capture_output=True, text=True, - timeout=2.0, + timeout=1.5, ) if res.returncode == 0: val = res.stdout.strip() @@ -90,6 +169,102 @@ def _get_parent_pid(pid: int) -> int | None: return None +def _get_process_group_members(pgid: int) -> list[int]: + """Find all live member PIDs of a process group.""" + members: list[int] = [] + proc_dir = Path("/proc") + if proc_dir.exists(): + try: + for entry in proc_dir.iterdir(): + if entry.name.isdigit(): + pid = int(entry.name) + stat_file = entry / "stat" + try: + content = stat_file.read_text() + rparen = content.rfind(")") + if rparen != -1: + fields = content[rparen + 2 :].split() + if len(fields) > 2 and int(fields[2]) == pgid: + members.append(pid) + except (OSError, ValueError, IndexError): + continue + if members: + return members + except OSError: + pass + + try: + res = subprocess.run( + ["ps", "-o", "pid=", "-g", str(pgid)], + capture_output=True, + text=True, + timeout=1.5, + ) + if res.returncode == 0: + for line in res.stdout.splitlines(): + line = line.strip() + if line.isdigit(): + members.append(int(line)) + return members + except Exception: + pass + + return members + + +def _is_group_fully_owned(pgid: int) -> bool: + """Return True only if every member process of pgid is an owned process of this session.""" + members = _get_process_group_members(pgid) + if not members: + return False + return all(is_pid_owned(m) for m in members) + + +def get_session_lifeline_pids() -> dict[int, str]: + """Identify PIDs critical to session survival (kernel, daemon owner, ssh anchors).""" + lifelines: dict[int, str] = {} + kernel_pid = os.getpid() + lifelines[kernel_pid] = "kernel process" + + ppid = os.getppid() + if ppid > 1: + lifelines[ppid] = "kernel parent process" + + owner_str = os.environ.get("PRIME_AGENT_KERNEL_OWNER_PID") + if owner_str: + try: + owner_pid = int(owner_str) + if owner_pid > 1: + lifelines[owner_pid] = "daemon owner process (PRIME_AGENT_KERNEL_OWNER_PID)" + except ValueError: + pass + + ssh_agent_str = os.environ.get("SSH_AGENT_PID") + if ssh_agent_str: + try: + agent_pid = int(ssh_agent_str) + if agent_pid > 1: + lifelines[agent_pid] = "SSH agent process (SSH_AGENT_PID)" + except ValueError: + pass + + curr = ppid + visited = {kernel_pid} + while curr > 1 and curr not in visited: + visited.add(curr) + ident = get_process_identity(curr) + comm = (ident.comm if ident else "").lower() + if any(anchor in comm for anchor in ("prime-agent", "ssh", "tmux")): + if curr not in lifelines: + lifelines[curr] = f"session anchor process ({comm})" + parent = _get_parent_pid(curr) + if parent is None: + break + curr = parent + + return lifelines + + def is_pid_owned(pid: int, max_hops: int = 8) -> bool: """True if the pid was started by bash(), adopted in this session, or is a descendant.""" if pid <= 1: @@ -97,27 +272,55 @@ def is_pid_owned(pid: int, max_hops: int = 8) -> bool: with _LEASE_LOCK: if pid in _OWNED_LEASES: + lease = _OWNED_LEASES[pid] + if lease.reaped: + return False + curr_ident = get_process_identity(pid) + if curr_ident is None: + lease.reaped = True + return False + if lease.identity is not None and not lease.identity.matches(curr_ident): + lease.reaped = True + return False + if lease.identity is None: + lease.identity = curr_ident return True - # Walk ancestry up to max_hops to discover if any ancestor is an owned lease + if pid in _ANCESTRY_CACHE: + desc_ident, anc_pid, anc_ident = _ANCESTRY_CACHE[pid] + curr_ident = get_process_identity(pid) + if curr_ident is None or not desc_ident.matches(curr_ident): + _ANCESTRY_CACHE.pop(pid, None) + else: + anc_lease = _OWNED_LEASES.get(anc_pid) + if anc_lease is not None and not anc_lease.reaped: + curr_anc_ident = get_process_identity(anc_pid) + if curr_anc_ident is not None and ( + anc_ident is None or anc_ident.matches(curr_anc_ident) + ): + return True + _ANCESTRY_CACHE.pop(pid, None) + current = pid hops = 0 - visited: set[int] = {pid} + visited: list[int] = [pid] while hops < max_hops: ppid = _get_parent_pid(current) if ppid is None or ppid <= 1 or ppid in visited: break - visited.add(ppid) + visited.append(ppid) with _LEASE_LOCK: if ppid in _OWNED_LEASES: - # Cache descendant ownership for O(1) subsequent lookups - _OWNED_LEASES[pid] = ProcessLease( - pid=pid, - command=f"descendant of {ppid}", - started_at=time.time(), - reason=f"ancestry walk matched ancestor {ppid}", - ) - return True + anc_lease = _OWNED_LEASES[ppid] + if not anc_lease.reaped: + curr_anc_ident = get_process_identity(ppid) + if curr_anc_ident is not None and ( + anc_lease.identity is None or anc_lease.identity.matches(curr_anc_ident) + ): + pid_ident = get_process_identity(pid) + if pid_ident is not None: + _ANCESTRY_CACHE[pid] = (pid_ident, ppid, curr_anc_ident) + return True current = ppid hops += 1 @@ -125,18 +328,42 @@ def is_pid_owned(pid: int, max_hops: int = 8) -> bool: def get_owned_process(pid: int) -> ProcessLease | None: - """Return the ownership record for a pid, or None if foreign.""" + """Return the ownership record for a pid, or None if foreign or expired.""" with _LEASE_LOCK: lease = _OWNED_LEASES.get(pid) - if lease is not None: - return lease + if lease is not None: + if lease.reaped: + return None + curr_ident = get_process_identity(pid) + if curr_ident is None or ( + lease.identity is not None and not lease.identity.matches(curr_ident) + ): + lease.reaped = True + return None + return lease + if is_pid_owned(pid): with _LEASE_LOCK: - return _OWNED_LEASES.get(pid) + if pid in _ANCESTRY_CACHE: + desc_ident, anc_pid, _ = _ANCESTRY_CACHE[pid] + anc_lease = _OWNED_LEASES.get(anc_pid) + return ProcessLease( + pid=pid, + command=f"descendant of {anc_pid}", + started_at=time.time(), + identity=desc_ident, + reason=f"ancestry walk matched ancestor {anc_pid}", + created_by_session=anc_lease.created_by_session if anc_lease else False, + ) return None -def _record_journal_adoption(pid: int, reason: str) -> None: +def _record_journal_adoption( + pid: int, + reason: str, + *, + self_lifeline_opt_in: bool = False, +) -> None: """Append an adoption entry to the session process journal.""" from datetime import datetime, timezone import json @@ -158,10 +385,24 @@ def _record_journal_adoption(pid: int, reason: str) -> None: "reason": reason, "ownerPid": owner_pid, "kernelPid": os.getpid(), + "selfLifelineOptIn": self_lifeline_opt_in, "active": True, "recordedAt": datetime.now(timezone.utc).isoformat(), } - data = (json.dumps(record) + "\n").encode() + payload = json.dumps(record) + "\n" + if self_lifeline_opt_in: + audit_record: dict[str, Any] = { + "version": 1, + "pid": pid, + "action": "self_lifeline_override", + "reason": f"Explicit opt-in self_lifeline_ok=True used for session lifeline process {pid}: {reason}", + "ownerPid": owner_pid, + "kernelPid": os.getpid(), + "recordedAt": datetime.now(timezone.utc).isoformat(), + } + payload += json.dumps(audit_record) + "\n" + + data = payload.encode() try: fd = os.open(path, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0o600) try: @@ -178,7 +419,12 @@ def _record_journal_adoption(pid: int, reason: str) -> None: pass -def adopt_process(pid: int, *, reason: str) -> ProcessLease: +def adopt_process( + pid: int, + *, + reason: str, + self_lifeline_ok: bool = False, +) -> ProcessLease: """Adopt an externally started process into the session's owned-process lease.""" if not isinstance(pid, int) or pid <= 1: raise ValueError(f"adopt_process requires a valid non-system PID (> 1), got {pid!r}") @@ -186,23 +432,37 @@ def adopt_process(pid: int, *, reason: str) -> ProcessLease: raise ValueError("adopt_process requires a non-empty reason string") clean_reason = reason.strip() + ident = get_process_identity(pid) + if ident is None: + raise ValueError(f"adopt_process rejects pid {pid}: process does not exist or is already dead") + + lifelines = get_session_lifeline_pids() + is_lifeline = pid in lifelines + if is_lifeline and not self_lifeline_ok: + raise ValueError( + f"adopt_process rejects pid {pid}: process is a critical session lifeline/anchor ({lifelines[pid]}). " + "Adopting session life support requires explicit opt-in with self_lifeline_ok=True." + ) + with _LEASE_LOCK: lease = ProcessLease( pid=pid, command=f"adopted: {clean_reason}", started_at=time.time(), + identity=ident, reason=clean_reason, handle=None, reaped=False, + created_by_session=False, ) _OWNED_LEASES[pid] = lease - _record_journal_adoption(pid, clean_reason) + _record_journal_adoption(pid, clean_reason, self_lifeline_opt_in=bool(is_lifeline and self_lifeline_ok)) return lease def list_owned_processes() -> list[dict[str, Any]]: - """List all processes started or adopted in this session.""" + """List all active processes started or adopted in this session.""" with _LEASE_LOCK: return [ { @@ -213,19 +473,21 @@ def list_owned_processes() -> list[dict[str, Any]]: "reason": lease.reason, } for lease in _OWNED_LEASES.values() + if not lease.reaped ] def list_owned_pids() -> list[int]: - """List PIDs started by bash() in this session.""" + """List active PIDs started or adopted in this session.""" with _LEASE_LOCK: - return sorted(_OWNED_LEASES.keys()) + return sorted(pid for pid, lease in _OWNED_LEASES.items() if not lease.reaped) def clear_owned_processes() -> None: """Internal helper to reset registry during tests.""" with _LEASE_LOCK: _OWNED_LEASES.clear() + _ANCESTRY_CACHE.clear() def kill_process( @@ -239,7 +501,6 @@ def kill_process( `target` must be an integer PID recorded by a BashHandle or the BashHandle itself. Refuses any foreign PID not evidenced as started by this session. """ - # If it's a BashHandle or has .kill() and .pid if hasattr(target, "kill") and hasattr(target, "pid"): target_pid = getattr(target, "pid", None) if isinstance(target_pid, int): @@ -248,6 +509,15 @@ def kill_process( f"kill_process rejects handle for pid {target_pid}: process is not owned by this session. " "Scoped process control only permits signalling processes started by bash() in this session." ) + lease = get_owned_process(target_pid) + curr_ident = get_process_identity(target_pid) + if curr_ident is None or (lease and lease.identity and not lease.identity.matches(curr_ident)): + if lease: + lease.reaped = True + raise ValueError( + f"kill_process rejects handle for pid {target_pid}: process lease has expired " + "(process is gone or PID was reused)." + ) target.kill(sig=sig, grace=grace) return True @@ -264,25 +534,61 @@ def kill_process( ) lease = get_owned_process(target) + if lease is None or lease.reaped: + raise ValueError(f"kill_process rejects pid {target}: process lease has expired or terminated.") + + curr_ident = get_process_identity(target) + if curr_ident is None or (lease.identity is not None and not lease.identity.matches(curr_ident)): + lease.reaped = True + raise ValueError( + f"kill_process rejects pid {target}: process lease has expired " + "(process is gone or PID was reused)." + ) # If the lease has an active BashHandle with .kill: - if lease is not None and lease.handle is not None and hasattr(lease.handle, "kill"): + if lease.handle is not None and hasattr(lease.handle, "kill"): lease.handle.kill(sig=sig, grace=grace) return True - # Fallback to process group signal + # Check whether process group signalling is legitimate: from .bash import _signal_group - delivered = _signal_group(target, sig) - if not delivered: - try: - os.kill(target, sig) - delivered = True - except ProcessLookupError: - delivered = True - except OSError: - delivered = False - return delivered + if lease.created_by_session: + # Group signalling is legitimate for groups the session created + delivered = _signal_group(target, sig) + if not delivered: + try: + os.kill(target, sig) + delivered = True + except (ProcessLookupError, OSError): + delivered = False + return delivered + + # Adopted lease: signal the PID alone (never the group) unless every member is owned + pgid = None + try: + pgid = os.getpgid(target) + except OSError: + pass + + if pgid is not None and _is_group_fully_owned(pgid): + delivered = _signal_group(target, sig) + if not delivered: + try: + os.kill(target, sig) + delivered = True + except (ProcessLookupError, OSError): + delivered = False + return delivered + + # Signal the PID alone + try: + os.kill(target, sig) + return True + except ProcessLookupError: + return True + except OSError: + return False signal_process = kill_process @@ -302,19 +608,23 @@ class HostProfile: vram_budget_gb: float | None = None # Explicit agent budget (if None: max(0, total - headroom)) services_to_survive: tuple[str, ...] = ( "display-manager", "gdm", "sddm", "lightdm", "sway", "hyprland", - "wayland", "X11", "pipewire", "wireplumber", "systemd-homed" + "wayland", "X11", "pipewire", "wireplumber", "systemd-homed", + "logind", "systemd-logind", ) allow_power_control: bool = False _CURRENT_HOST_PROFILE: HostProfile | None = None _PROFILE_LOCK = threading.Lock() +_PROBED_VRAM_GB: float | None = None _POWER_CONTROL_COMMANDS = frozenset({"reboot", "poweroff", "shutdown", "halt", "telinit"}) _SYSTEMCTL_POWER_VERBS = frozenset({ "poweroff", "reboot", "halt", "suspend", "hibernate", "hybrid-sleep", "kexec", }) -_SYSTEMCTL_MUTATION_VERBS = frozenset({"disable", "mask", "stop"}) +_SYSTEMCTL_MUTATION_VERBS = frozenset({ + "disable", "mask", "stop", "restart", "try-restart", "reload-or-restart", "isolate", +}) _POWER_CONTROL_OPT_IN_RE = re.compile( r"(?:^|[\s;&|(])(?:RLM_ALLOW_POWER_CONTROL|PRIME_ALLOW_POWER_CONTROL)=(?:1|true|yes|on)(?:[\s;&|)]|$)", @@ -390,6 +700,67 @@ def _load_profile_from_harness() -> HostProfile | None: return None +def _probe_host_vram_gb() -> float | None: + """Probe total physical GPU memory in GB via nvidia-smi or rocm-smi.""" + global _PROBED_VRAM_GB + if _PROBED_VRAM_GB is not None: + return _PROBED_VRAM_GB + + # 1. Probe nvidia-smi + try: + res = subprocess.run( + ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], + capture_output=True, + text=True, + timeout=1.5, + ) + if res.returncode == 0 and res.stdout.strip(): + total_mib = 0.0 + for line in res.stdout.splitlines(): + line = line.strip() + if line: + total_mib += float(line) + if total_mib > 0: + _PROBED_VRAM_GB = total_mib / 1024.0 + return _PROBED_VRAM_GB + except Exception: + pass + + # 2. Probe rocm-smi + try: + res = subprocess.run( + ["rocm-smi", "--showmeminfo", "vram"], + capture_output=True, + text=True, + timeout=1.5, + ) + if res.returncode == 0 and res.stdout.strip(): + for line in res.stdout.splitlines(): + if "total" in line.lower() and "memory" in line.lower(): + m = re.search(r"(\d+)", line) + if m: + val = float(m.group(1)) + if val > 1_000_000_000: + _PROBED_VRAM_GB = val / (1024.0**3) + return _PROBED_VRAM_GB + except Exception: + pass + + return None + + +def is_headless() -> bool: + """True if the environment has no desktop display session to preserve.""" + if os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY") or os.environ.get("MIR_SOCKET"): + return False + if sys.platform == "darwin": + # macOS has WindowServer and active display sessions; not headless + return False + if os.environ.get("XDG_SESSION_TYPE") in ("wayland", "x11"): + return False + return True + + def get_host_profile() -> HostProfile: """Return the active host profile (configured, env, harness memory, or default).""" with _PROFILE_LOCK: @@ -422,7 +793,8 @@ def set_host_profile( vram_budget_gb=vram_budget_gb, services_to_survive=services_to_survive if services_to_survive is not None else ( "display-manager", "gdm", "sddm", "lightdm", "sway", "hyprland", - "wayland", "X11", "pipewire", "wireplumber", "systemd-homed" + "wayland", "X11", "pipewire", "wireplumber", "systemd-homed", + "logind", "systemd-logind", ), allow_power_control=allow_power_control, ) @@ -433,18 +805,29 @@ def set_host_profile( def reset_host_profile() -> None: """Reset the host profile to default/unconfigured.""" with _PROFILE_LOCK: - global _CURRENT_HOST_PROFILE + global _CURRENT_HOST_PROFILE, _PROBED_VRAM_GB _CURRENT_HOST_PROFILE = None + _PROBED_VRAM_GB = None def available_vram_budget_gb(profile: HostProfile | None = None) -> float | None: - """Return the available VRAM budget for agent operations in GB, or None if unconfigured.""" + """Return the available VRAM budget in GB, deriving from probed GPU or conservative fallback if unconfigured.""" p = profile or get_host_profile() if p.vram_budget_gb is not None: return max(0.0, p.vram_budget_gb) if p.total_vram_gb > 0: return max(0.0, p.total_vram_gb - p.desktop_headroom_gb) - return None + + # Unconfigured: attempt hardware probe + probed = _probe_host_vram_gb() + if probed is not None and probed > 0: + headroom = 0.0 if is_headless() else p.desktop_headroom_gb + return max(0.0, probed - headroom) + + # Conservative unconfigured fallback: fail closed for large models + if is_headless(): + return 80.0 # Server headroom-free ceiling for headless environments + return max(0.0, 16.0 - p.desktop_headroom_gb) # 12.0 GB default desktop budget def preflight_vram( @@ -456,10 +839,13 @@ def preflight_vram( p = profile or get_host_profile() budget = available_vram_budget_gb(p) if budget is not None and requested_vram_gb > budget: + if p.total_vram_gb > 0: + budget_desc = f"{budget:.1f} GB available: {p.total_vram_gb:.1f} GB total minus {p.desktop_headroom_gb:.1f} GB desktop headroom reserved" + else: + budget_desc = f"{budget:.1f} GB available (unconfigured fallback with {p.desktop_headroom_gb:.1f} GB desktop headroom)" raise ValueError( f"VRAM preflight rejected: requested {requested_vram_gb:.1f} GB exceeds host VRAM budget " - f"({budget:.1f} GB available: {p.total_vram_gb:.1f} GB total minus {p.desktop_headroom_gb:.1f} GB " - "desktop headroom reserved). To proceed if you have verified headroom, adjust the host profile via " + f"({budget_desc}). To proceed if you have verified headroom, adjust the host profile via " "set_host_profile(...) or acknowledge with `RLM_ALLOW_OVERBUDGET_VRAM=1`." ) return { @@ -521,21 +907,21 @@ def check_power_control(command: str, words: list[str]) -> None: def estimate_command_vram(command: str, words: list[str], profile: HostProfile) -> float | None: """Estimate requested VRAM in GB from command flags or model names, if discernible.""" + total_ref = profile.total_vram_gb if profile.total_vram_gb > 0 else 24.0 + # 1. Explicit --gpu-memory-utilization flag for i, tok in enumerate(words): if tok.startswith("--gpu-memory-utilization="): val_str = tok.split("=", 1)[1] try: util = float(val_str) - if profile.total_vram_gb > 0: - return util * profile.total_vram_gb + return util * total_ref except ValueError: pass elif tok == "--gpu-memory-utilization" and i + 1 < len(words): try: util = float(words[i + 1]) - if profile.total_vram_gb > 0: - return util * profile.total_vram_gb + return util * total_ref except ValueError: pass @@ -559,12 +945,6 @@ def estimate_command_vram(command: str, words: list[str], profile: HostProfile) if m: try: b_val = float(m.group(1)) - # Reasonable VRAM heuristics: - # 70B: ~70 GB (conservative, 4-bit AWQ is ~38 GB, FP16 is 140 GB) - # 405B: ~200+ GB - # 32B: ~25.6 GB - # 14B: ~11.2 GB - # 8B: ~6.4 GB if b_val >= 70: return b_val elif b_val >= 30: @@ -591,14 +971,17 @@ def check_vram_command(command: str, words: list[str]) -> None: profile = get_host_profile() budget = available_vram_budget_gb(profile) if budget is None: - return # No budget configured; proceed + return estimated = estimate_command_vram(command, words, profile) if estimated is not None and estimated > budget: cmd_display = words[0] + if profile.total_vram_gb > 0: + budget_desc = f"{budget:.1f} GB with {profile.desktop_headroom_gb:.1f} GB desktop headroom reserved on {profile.total_vram_gb:.1f} GB total" + else: + budget_desc = f"{budget:.1f} GB with {profile.desktop_headroom_gb:.1f} GB desktop headroom reserved (unconfigured fallback)" raise ValueError( f"bash() rejects `{cmd_display}`: estimated VRAM requirement ({estimated:.1f} GB) " - f"exceeds host VRAM budget ({budget:.1f} GB with {profile.desktop_headroom_gb:.1f} GB " - f"desktop headroom reserved on {profile.total_vram_gb:.1f} GB total). " + f"exceeds host VRAM budget ({budget_desc}). " "Adjust the host profile via set_host_profile(...) or acknowledge with `RLM_ALLOW_OVERBUDGET_VRAM=1`." ) diff --git a/prime-agent-runtime/test/test_host_safety.py b/prime-agent-runtime/test/test_host_safety.py index 8133b08be..42d615a5b 100644 --- a/prime-agent-runtime/test/test_host_safety.py +++ b/prime-agent-runtime/test/test_host_safety.py @@ -26,6 +26,9 @@ from rlm.bash import ( signal_process, ) from rlm.host_safety import ( + ProcessIdentity, + _ANCESTRY_CACHE, + _OWNED_LEASES, clear_owned_processes, register_owned_process, ) @@ -386,6 +389,146 @@ class HostSafetyTests(unittest.IsolatedAsyncioTestCase): cmd = "vllm serve meta-llama/Llama-3-8B --gpu-memory-utilization 0.7" self.assertIsNone(_check_command_guards(cmd)) + # ----------------------------------------------------------------------- + # 4. Hardened Process Ownership & Identity (Defects 1, 2, 3) + # ----------------------------------------------------------------------- + + def test_dead_process_lease_expires_and_refuses_signal(self): + """A lease for a reaped/dead process expires and cannot authorize signals.""" + proc = subprocess.Popen(["sleep", "0.05"]) + register_owned_process(proc.pid, "sleep 0.05") + proc.wait(timeout=5) + self.assertFalse(is_pid_owned(proc.pid)) + with self.assertRaises(ValueError) as ctx: + kill_process(proc.pid) + self.assertIn("not owned by this session", str(ctx.exception)) + + def test_pid_reuse_identity_mismatch_refuses_signal(self): + """PID reuse with mismatched process identity expires lease and refuses signal.""" + proc = subprocess.Popen(["sleep", "30"]) + try: + lease = register_owned_process(proc.pid, "sleep 30") + lease.identity = ProcessIdentity(pid=proc.pid, start_time="mismatched_start_123", comm="stranger") + self.assertFalse(is_pid_owned(proc.pid)) + with self.assertRaises(ValueError) as ctx: + kill_process(proc.pid) + self.assertIn("not owned by this session", str(ctx.exception)) + finally: + proc.kill() + proc.wait(timeout=5) + + def test_ancestry_walk_does_not_cache_visited_strangers(self): + """Ancestry walk does not permanently enroll visited intermediate PIDs into _OWNED_LEASES.""" + proc = subprocess.Popen(["sh", "-c", "sleep 60 & echo $!; wait"], stdout=subprocess.PIPE, text=True) + try: + register_owned_process(proc.pid, "parent") + assert proc.stdout is not None + grandchild_pid = int(proc.stdout.readline().strip()) + self.assertTrue(is_pid_owned(grandchild_pid)) + self.assertNotIn(grandchild_pid, _OWNED_LEASES) + finally: + proc.kill() + proc.wait(timeout=5) + + def test_adopted_process_signals_pid_alone_not_entire_group(self): + """An adopted process signals the PID alone without os.killpg when group is not fully owned.""" + proc = subprocess.Popen(["sleep", "30"]) + try: + adopt_process(proc.pid, reason="test adoption") + with mock.patch("os.kill") as mock_kill, mock.patch("os.killpg") as mock_killpg: + kill_process(proc.pid, sig=signal.SIGTERM) + mock_kill.assert_called_once_with(proc.pid, signal.SIGTERM) + mock_killpg.assert_not_called() + finally: + proc.kill() + proc.wait(timeout=5) + + def test_adopt_lifeline_refused_without_opt_in(self): + """Adopting kernel self, ppid, or owner PID is refused without self_lifeline_ok=True.""" + with self.assertRaises(ValueError) as ctx: + adopt_process(os.getpid(), reason="attempt self adopt") + self.assertIn("critical session lifeline", str(ctx.exception)) + self.assertIn("self_lifeline_ok=True", str(ctx.exception)) + + with self.assertRaises(ValueError) as ctx: + adopt_process(os.getppid(), reason="attempt ppid adopt") + self.assertIn("critical session lifeline", str(ctx.exception)) + + def test_adopt_lifeline_permitted_with_opt_in_and_audited(self): + """Adopting lifeline with self_lifeline_ok=True succeeds and emits audit entry.""" + import tempfile, json + with tempfile.NamedTemporaryFile(mode="w+", delete=False) as tf: + journal_path = tf.name + try: + with mock.patch.dict(os.environ, {"PRIME_AGENT_INTERNAL_ORPHAN_PROCESS_JOURNAL": journal_path}): + lease = adopt_process(os.getppid(), reason="explicit supervisor restart", self_lifeline_ok=True) + self.assertEqual(lease.pid, os.getppid()) + with open(journal_path) as f: + lines = [json.loads(line) for line in f if line.strip()] + self.assertTrue(any(r.get("selfLifelineOptIn") is True for r in lines)) + self.assertTrue(any(r.get("action") == "self_lifeline_override" for r in lines)) + finally: + if os.path.exists(journal_path): + os.remove(journal_path) + + # ----------------------------------------------------------------------- + # 5. systemctl Restart Verbs Guard (Defect 4) + # ----------------------------------------------------------------------- + + def test_systemctl_restart_verbs_refused_on_desktop_services(self): + """systemctl restart, try-restart, reload-or-restart, isolate are guarded for protected desktop services.""" + set_host_profile(services_to_survive=("display-manager", "gdm", "logind")) + for verb in ["restart", "try-restart", "reload-or-restart", "isolate"]: + for target in ["display-manager", "gdm.service", "logind"]: + with self.subTest(verb=verb, target=target): + with self.assertRaises(ValueError) as ctx: + _check_command_guards(f"systemctl {verb} {target}") + msg = str(ctx.exception) + self.assertIn("protected desktop service", msg) + self.assertIn("RLM_ALLOW_POWER_CONTROL=1", msg) + + def test_systemctl_restart_permitted_for_unrelated_services(self): + """systemctl restart on normal/unrelated services is permitted.""" + set_host_profile(services_to_survive=("display-manager", "gdm")) + for cmd in ["systemctl restart nginx", "systemctl try-restart my-worker.service"]: + with self.subTest(cmd=cmd): + self.assertIsNone(_check_command_guards(cmd)) + + # ----------------------------------------------------------------------- + # 6. Unconfigured VRAM Preflight Fail-Closed (Defect 5) + # ----------------------------------------------------------------------- + + def test_unconfigured_vram_fails_closed_for_large_model(self): + """Unconfigured host profile fails closed for large model requests exceeding default headroom/budget.""" + reset_host_profile() + with self.assertRaises(ValueError) as ctx: + _check_command_guards("vllm serve meta-llama/Llama-3-70B") + msg = str(ctx.exception) + self.assertIn("exceeds host VRAM budget", msg) + self.assertIn("RLM_ALLOW_OVERBUDGET_VRAM=1", msg) + self.assertIn("set_host_profile", msg) + + def test_unconfigured_vram_permits_small_model(self): + """Unconfigured host profile permits small models within conservative headroom.""" + reset_host_profile() + self.assertIsNone(_check_command_guards("vllm serve meta-llama/Llama-3-8B --gpu-memory-utilization 0.5")) + + def test_unconfigured_vram_permits_large_model_with_opt_in(self): + """Unconfigured host profile permits large models with explicit opt-in.""" + reset_host_profile() + self.assertIsNone(_check_command_guards("RLM_ALLOW_OVERBUDGET_VRAM=1 vllm serve meta-llama/Llama-3-70B")) + + def test_unconfigured_vram_headless_allows_normal_work(self): + """On a genuinely headless host, no desktop headroom is reserved and normal server work is permitted.""" + reset_host_profile() + with mock.patch("rlm.host_safety.is_headless", return_value=True): + # Headless unconfigured budget is 80.0 GB, so 70B model passes + self.assertIsNone(_check_command_guards("vllm serve meta-llama/Llama-3-70B")) + # Runaway request (> 80 GB) fails closed even on headless + with self.assertRaises(ValueError) as ctx: + _check_command_guards("vllm serve meta-llama/Llama-3-405B") + self.assertIn("RLM_ALLOW_OVERBUDGET_VRAM=1", str(ctx.exception)) + if __name__ == "__main__": unittest.main() -- 2.54.0