From 5f9308322116372ca0a4334be43463b965dc95bb Mon Sep 17 00:00:00 2001 From: Dominic Bejar Date: Thu, 6 Aug 2026 01:59:13 +0000 Subject: [PATCH] fix(terminal): bound isolated worker memory --- tests/tools/test_process_registry.py | 24 ++++++++++ tools/process_registry.py | 72 +++++++++++++++++++++++++++- 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/tests/tools/test_process_registry.py b/tests/tools/test_process_registry.py index 4d8a02f125225..7d2238f24dd74 100644 --- a/tests/tools/test_process_registry.py +++ b/tests/tools/test_process_registry.py @@ -1651,6 +1651,17 @@ class TestSystemdCgroupIsolation: unit_idx = argv.index("--unit") assert argv[unit_idx + 1].startswith("hermes-worker-"), argv assert argv[unit_idx + 1] == f"hermes-worker-{session.id}", argv # _build_systemd_scope_argv uses bare name + properties = [ + argv[index + 1] + for index, value in enumerate(argv[:-1]) + if value == "--property" + ] + assert "MemoryAccounting=yes" in properties + assert "OOMPolicy=kill" in properties + memory_max = next( + value for value in properties if value.startswith("MemoryMax=") + ) + assert int(memory_max.split("=", 1)[1]) > 0 # The original shell command must still be present at the tail, # after the ``--`` separator that prevents systemd-run from # interpreting command flags as its own. @@ -1782,6 +1793,19 @@ class TestSystemdCgroupIsolation: assert argv[-3:] == ["/bin/bash", "-lic", "set +m; codex"] assert session.systemd_unit == f"hermes-worker-{session.id}.scope" + def test_worker_memory_limit_accepts_valid_byte_override(self, monkeypatch): + import tools.process_registry as pr + + monkeypatch.setenv("HERMES_WORKER_MEMORY_MAX_BYTES", "123456789") + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/systemd-run") + + argv = pr._build_systemd_scope_argv( + ["/bin/bash", "-lc", "true"], + unit_suffix="test", + ) + + assert "MemoryMax=123456789" in argv + def test_kill_recovered_detached_already_exited_stops_persisted_scope( self, registry, monkeypatch ): diff --git a/tools/process_registry.py b/tools/process_registry.py index f29204b999124..8a22956c2e357 100644 --- a/tools/process_registry.py +++ b/tools/process_registry.py @@ -40,6 +40,7 @@ import subprocess import threading import time import uuid +from pathlib import Path _IS_WINDOWS = platform.system() == "Windows" from tools.environments.local import _find_shell, _resolve_safe_cwd, _sanitize_subprocess_env @@ -99,6 +100,64 @@ _SYSTEMD_SCOPE_AVAILABLE: Optional[bool] = None _SYSTEMD_SCOPE_PROBE_LOCK = threading.Lock() _SYSTEMD_SCOPE_PROBED_AT = 0.0 _SYSTEMD_SCOPE_FAILURE_TTL_SECONDS = 60.0 +_MIN_WORKER_MEMORY_MAX_BYTES = 64 * 1024 * 1024 +_DEFAULT_WORKER_MEMORY_MAX_BYTES = 1024 * 1024 * 1024 +_WORKER_MEMORY_MAX_CAP_BYTES = 4 * 1024 * 1024 * 1024 + + +def _worker_memory_max_bytes() -> int: + """Return a finite per-worker cgroup limit without widening host risk. + + An explicit byte override is useful for operators with known workload + requirements. Otherwise retain the tighter of the gateway's current + cgroup-v2 ``memory.max`` and half of physical RAM, capped at 4 GiB. This + keeps the sibling worker outside the gateway cgroup while ensuring the + worker cannot consume memory up to the enclosing user slice or host limit. + """ + override = os.getenv("HERMES_WORKER_MEMORY_MAX_BYTES", "").strip() + if override: + try: + parsed = int(override) + if parsed >= _MIN_WORKER_MEMORY_MAX_BYTES: + return parsed + except ValueError: + pass + logger.warning( + "Ignoring invalid HERMES_WORKER_MEMORY_MAX_BYTES=%r; " + "expected an integer of at least %d bytes", + override, + _MIN_WORKER_MEMORY_MAX_BYTES, + ) + + candidates: List[int] = [] + try: + for line in Path("/proc/self/cgroup").read_text(encoding="utf-8").splitlines(): + if line.startswith("0::"): + relative = line.partition("::")[2].lstrip("/") + raw_limit = ( + Path("/sys/fs/cgroup") / relative / "memory.max" + ).read_text(encoding="utf-8").strip() + if raw_limit.isdigit(): + cgroup_limit = int(raw_limit) + if cgroup_limit >= _MIN_WORKER_MEMORY_MAX_BYTES: + candidates.append(cgroup_limit) + break + except (OSError, ValueError): + pass + + try: + physical_bytes = int(os.sysconf("SC_PHYS_PAGES")) * int( + os.sysconf("SC_PAGE_SIZE") + ) + physical_bound = min( + _WORKER_MEMORY_MAX_CAP_BYTES, + max(_MIN_WORKER_MEMORY_MAX_BYTES, physical_bytes // 2), + ) + candidates.append(physical_bound) + except (OSError, ValueError, TypeError): + pass + + return min(candidates) if candidates else _DEFAULT_WORKER_MEMORY_MAX_BYTES def _systemd_run_user_scope_available() -> bool: @@ -152,7 +211,11 @@ def _systemd_run_user_scope_available() -> bool: [ binary, "--user", "--scope", "--quiet", "--unit", probe_unit, - "--collect", "--", + "--collect", + "--property", "MemoryAccounting=yes", + "--property", f"MemoryMax={_worker_memory_max_bytes()}", + "--property", "OOMPolicy=kill", + "--", "/bin/true", ], capture_output=True, @@ -194,6 +257,7 @@ def _build_systemd_scope_argv( # guard anyway so we never pass None into Popen. return shell_argv unit_name = f"hermes-worker-{unit_suffix}" + memory_max = _worker_memory_max_bytes() return [ binary, "--user", @@ -202,6 +266,12 @@ def _build_systemd_scope_argv( "--unit", unit_name, "--collect", + "--property", + "MemoryAccounting=yes", + "--property", + f"MemoryMax={memory_max}", + "--property", + "OOMPolicy=kill", "--", *shell_argv, ]