From 240767d95411c41a020e08d4cbc1f62a61ef06df Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:12:50 -0700 Subject: [PATCH 01/32] scripts : add Strix host-memory watchdog Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- README.md | 1 + docs/strix-memory-watchdog.md | 40 ++ scripts/strix_memory_watchdog.py | 664 ++++++++++++++++++++++++++++ tests/CMakeLists.txt | 11 + tests/test_strix_memory_watchdog.py | 406 +++++++++++++++++ 5 files changed, 1122 insertions(+) create mode 100644 docs/strix-memory-watchdog.md create mode 100755 scripts/strix_memory_watchdog.py create mode 100644 tests/test_strix_memory_watchdog.py diff --git a/README.md b/README.md index 3067a7a4e71e..dbf547ce5c5b 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,7 @@ Everything else is upstream `llama.cpp`. The additions currently carried here: | Speculative checkpoints on device | | `llama-server` keeps speculative-decoding checkpoints in device memory instead of copying them to the host | | ROCmFPx quant types | `llama-quantize` types `Q4_0_ROCMFP4`, `Q4_0_ROCMFP4_FAST`, `Q2/Q3/Q6/Q8_0_ROCMFPX` and the `_LEAN`/`_COHERENT`/`_STRIX` recipes | Loads the ROCmFP4 GGUFs published for Strix Halo. CPU codecs plus Vulkan dequant, mat-vec, matmul and integer-dot kernels. Weight formats only: not accepted as KV-cache types | | Repeatable output at depth | | Freed KV cells are zeroed so masked-out rows never carry stale K/V, and the Vulkan radix top-k assigns output slots deterministically | +| Host-memory watchdog | [`scripts/strix_memory_watchdog.py`](docs/strix-memory-watchdog.md) | Runs a command in a process group, requires zero active swap, and stops before host-wide memory reaches the 120 GiB validation ceiling | Every ROCm/HIP change above is guarded on architecture, shape and layout, so other devices see upstream behaviour. Run `--help`, or see [tools/server/README.md](tools/server/README.md), for the full options. diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md new file mode 100644 index 000000000000..4adc0e4be7fe --- /dev/null +++ b/docs/strix-memory-watchdog.md @@ -0,0 +1,40 @@ +# Strix host-memory watchdog + +`scripts/strix_memory_watchdog.py` is an external Linux command wrapper for headless Strix Halo validation. It does not change model loading or cache sizing. It measures host-wide memory from procfs and controls the launched command's process group. + +```sh +./scripts/strix_memory_watchdog.py -- ./build/bin/llama-server +``` + +The wrapper performs these checks and actions: + +- It refuses to launch if `/proc/swaps` contains any active entry. +- It calculates used memory as `MemTotal - MemAvailable`. Linux reports these fields in KiB, so the wrapper multiplies each value by 1024 and keeps all accounting as integer bytes. +- It sends `SIGTERM` to the process group at 116 GiB used. +- It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. +- It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. +- It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. + +The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. + +Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. + +The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Child standard input, standard output, and standard error are inherited unchanged. + +Exit classifications are authoritative in the final JSON record. Operational failures use these exit codes: + +| Exit code | Classification | +| ---: | --- | +| 2 | procfs or configuration error | +| 3 | swap active at startup or detected during execution | +| 4 | soft threshold reached | +| 5 | emergency threshold reached | +| 6 | soft-threshold grace period expired | +| 7 | process-group signaling or termination failure | +| 127 | command launch failure | + +No model, backend, or ROCm package is required to run the unit tests: + +```sh +python3 tests/test_strix_memory_watchdog.py +``` diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py new file mode 100755 index 000000000000..0dc532c1d027 --- /dev/null +++ b/scripts/strix_memory_watchdog.py @@ -0,0 +1,664 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import signal +import subprocess +import sys +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import IO, Protocol + + +GIB = 1024**3 +STRICT_CEILING_BYTES = 120 * GIB +DEFAULT_SOFT_BYTES = 116 * GIB +DEFAULT_EMERGENCY_BYTES = 118 * GIB +DEFAULT_GRACE_SECONDS = 30.0 +DEFAULT_SAMPLE_INTERVAL_SECONDS = 1.0 + +EXIT_PROCFS_ERROR = 2 +EXIT_SWAP_ACTIVE = 3 +EXIT_SOFT_LIMIT = 4 +EXIT_EMERGENCY_LIMIT = 5 +EXIT_GRACE_TIMEOUT = 6 +EXIT_SIGNAL_ERROR = 7 +EXIT_LAUNCH_ERROR = 127 + +MEMINFO_VALUE_RE = re.compile(r"([0-9]+) kB") +SWAPS_HEADER = ["Filename", "Type", "Size", "Used", "Priority"] + + +class ProcfsError(RuntimeError): + pass + + +class ProcessGroupError(RuntimeError): + pass + + +class ProcessHandle(Protocol): + pid: int + + def poll(self) -> int | None: + ... + + def wait(self, timeout: float | None = None) -> int: + ... + + +@dataclass(frozen=True) +class HostSnapshot: + total_bytes: int + available_bytes: int + active_swaps: tuple[str, ...] + + @property + def used_bytes(self) -> int: + return self.total_bytes - self.available_bytes + + +@dataclass(frozen=True) +class WatchdogConfig: + command: tuple[str, ...] + procfs_root: Path = Path("/proc") + soft_bytes: int = DEFAULT_SOFT_BYTES + emergency_bytes: int = DEFAULT_EMERGENCY_BYTES + grace_seconds: float = DEFAULT_GRACE_SECONDS + sample_interval_seconds: float = DEFAULT_SAMPLE_INTERVAL_SECONDS + + def validate(self) -> None: + if not self.command: + raise ValueError("a command is required after --") + if self.soft_bytes <= 0: + raise ValueError("soft threshold must be greater than zero") + if self.emergency_bytes <= self.soft_bytes: + raise ValueError("emergency threshold must be greater than soft threshold") + if self.emergency_bytes >= STRICT_CEILING_BYTES: + raise ValueError("emergency threshold must be below 120 GiB") + if not math.isfinite(self.grace_seconds) or self.grace_seconds <= 0: + raise ValueError("grace period must be greater than zero") + if ( + not math.isfinite(self.sample_interval_seconds) + or self.sample_interval_seconds <= 0 + ): + raise ValueError("sample interval must be greater than zero") + + +class ProcfsReader: + def __init__(self, root: Path): + self.root = root + + def _read_text(self, name: str) -> str: + path = self.root / name + try: + return path.read_text(encoding="utf-8") + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcfsError(f"cannot read {path}: {detail}") from exc + + def read_snapshot(self) -> HostSnapshot: + active_swaps = self._parse_swaps(self._read_text("swaps")) + total_bytes, available_bytes = self._parse_meminfo( + self._read_text("meminfo") + ) + return HostSnapshot(total_bytes, available_bytes, active_swaps) + + @staticmethod + def _parse_meminfo(content: str) -> tuple[int, int]: + values: dict[str, int] = {} + required = {"MemTotal", "MemAvailable"} + for line in content.splitlines(): + key, separator, raw_value = line.partition(":") + if not separator or key not in required: + continue + if key in values: + raise ProcfsError(f"duplicate {key} in meminfo") + match = MEMINFO_VALUE_RE.fullmatch(raw_value.strip()) + if match is None: + raise ProcfsError(f"malformed {key} in meminfo") + values[key] = int(match.group(1)) * 1024 + + missing = sorted(required - values.keys()) + if missing: + raise ProcfsError(f"missing {', '.join(missing)} in meminfo") + if values["MemAvailable"] > values["MemTotal"]: + raise ProcfsError("MemAvailable exceeds MemTotal") + return values["MemTotal"], values["MemAvailable"] + + @staticmethod + def _parse_swaps(content: str) -> tuple[str, ...]: + lines = content.splitlines() + if not lines or lines[0].split() != SWAPS_HEADER: + raise ProcfsError("malformed swaps header") + + entries: list[str] = [] + for line in lines[1:]: + if not line.strip(): + continue + fields = line.split() + if len(fields) != len(SWAPS_HEADER): + raise ProcfsError("malformed swaps entry") + try: + int(fields[2]) + int(fields[3]) + int(fields[4]) + except ValueError as exc: + raise ProcfsError("malformed swaps entry") from exc + entries.append(fields[0]) + return tuple(entries) + + +class AuditLogger: + def __init__( + self, + stream: IO[str], + wall_clock: Callable[[], datetime] | None = None, + ): + self.stream = stream + self.wall_clock = wall_clock or ( + lambda: datetime.now(timezone.utc) + ) + + def emit(self, event: str, **fields: object) -> None: + timestamp = self.wall_clock().astimezone(timezone.utc) + record = { + "timestamp": timestamp.isoformat(timespec="milliseconds").replace( + "+00:00", "Z" + ), + "event": event, + **fields, + } + self.stream.write( + json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" + ) + self.stream.flush() + + +def _child_status(returncode: int | None, started: bool = True) -> str: + if not started: + return "not_started" + if returncode is None: + return "running" + return "signaled" if returncode < 0 else "exited" + + +def _state_fields( + snapshot: HostSnapshot | None, + peak_used_bytes: int | None, + child: ProcessHandle | None, + child_returncode: int | None, + process_group_status: str, + threshold_reason: str, +) -> dict[str, object]: + return { + "total_bytes": snapshot.total_bytes if snapshot else None, + "available_bytes": snapshot.available_bytes if snapshot else None, + "used_bytes": snapshot.used_bytes if snapshot else None, + "swap_entries": len(snapshot.active_swaps) if snapshot else None, + "peak_used_bytes": peak_used_bytes, + "child_pid": child.pid if child else None, + "child_status": _child_status( + child_returncode, started=child is not None + ), + "child_returncode": child_returncode, + "process_group_id": child.pid if child else None, + "process_group_status": process_group_status, + "threshold_reason": threshold_reason, + } + + +def _emit_final( + audit: AuditLogger, + classification: str, + exit_code: int, + reason: str, + snapshot: HostSnapshot | None, + peak_used_bytes: int | None, + child: ProcessHandle | None = None, + child_returncode: int | None = None, + process_group_status: str = "not_created", + error: str | None = None, +) -> int: + fields = _state_fields( + snapshot, + peak_used_bytes, + child, + child_returncode, + process_group_status, + reason, + ) + fields.update(classification=classification, exit_code=exit_code) + if error: + fields["error"] = error + audit.emit("final", **fields) + return exit_code + + +def _signal_process_group(process_group_id: int, signal_number: int) -> str: + try: + os.killpg(process_group_id, signal_number) + except ProcessLookupError: + return "missing" + except OSError as exc: + name = signal.Signals(signal_number).name + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot send {name} to process group {process_group_id}: {detail}" + ) from exc + return f"{signal.Signals(signal_number).name.lower()}_sent" + + +def _kill_and_finish( + audit: AuditLogger, + child: ProcessHandle, + snapshot: HostSnapshot, + peak_used_bytes: int, + classification: str, + exit_code: int, + reason: str, + signal_group: Callable[[int, int], str], +) -> int: + try: + group_status = signal_group(child.pid, signal.SIGKILL) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "signal_error", + str(exc), + ) + + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + group_status, + reason, + ), + signal="SIGKILL", + ) + try: + child_returncode = child.wait(timeout=5.0) + except subprocess.TimeoutExpired as exc: + return _emit_final( + audit, + "termination_timeout", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "sigkill_timeout", + str(exc), + ) + return _emit_final( + audit, + classification, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + group_status, + ) + + +def _monitor_child( + config: WatchdogConfig, + reader: ProcfsReader, + audit: AuditLogger, + child: ProcessHandle, + initial_snapshot: HostSnapshot, + signal_group: Callable[[int, int], str], + monotonic: Callable[[], float], + sleeper: Callable[[float], None], +) -> int: + snapshot = initial_snapshot + peak = snapshot.used_bytes + soft_deadline: float | None = None + + while True: + child_returncode = child.poll() + if child_returncode is not None: + soft_stop = soft_deadline is not None + return _emit_final( + audit, + "soft_limit" if soft_stop else "child_exit", + EXIT_SOFT_LIMIT if soft_stop else ( + 128 - child_returncode + if child_returncode < 0 + else child_returncode + ), + ( + "child exited during soft-threshold grace period" + if soft_stop + else "child exited" + ), + snapshot, + peak, + child, + child_returncode, + "leader_exited", + ) + + now = monotonic() + if soft_deadline is not None and now >= soft_deadline: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "grace_timeout", + EXIT_GRACE_TIMEOUT, + "soft-threshold grace period expired", + signal_group, + ) + + try: + snapshot = reader.read_snapshot() + except ProcfsError as exc: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + signal_group, + ) + + peak = max(peak, snapshot.used_bytes) + audit.emit( + "sample", + **_state_fields( + snapshot, peak, child, None, "active", "none" + ), + ) + + if snapshot.active_swaps: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "swap_appeared", + EXIT_SWAP_ACTIVE, + "active swap appeared during execution", + signal_group, + ) + if snapshot.used_bytes >= config.emergency_bytes: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "emergency_limit", + EXIT_EMERGENCY_LIMIT, + "used_bytes >= emergency_bytes", + signal_group, + ) + if soft_deadline is None and snapshot.used_bytes >= config.soft_bytes: + try: + group_status = signal_group(child.pid, signal.SIGTERM) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + "used_bytes >= soft_bytes", + snapshot, + peak, + child, + child.poll(), + "signal_error", + str(exc), + ) + soft_deadline = now + config.grace_seconds + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak, + child, + child.poll(), + group_status, + "used_bytes >= soft_bytes", + ), + signal="SIGTERM", + grace_deadline_monotonic=soft_deadline, + ) + + sleep_seconds = config.sample_interval_seconds + if soft_deadline is not None: + sleep_seconds = min( + sleep_seconds, + max(0.0, soft_deadline - monotonic()), + ) + sleeper(sleep_seconds) + + +def run_watchdog( + config: WatchdogConfig, + *, + reader: ProcfsReader | None = None, + audit: AuditLogger | None = None, + launcher: Callable[..., ProcessHandle] | None = None, + signal_group: Callable[[int, int], str] | None = None, + monotonic: Callable[[], float] | None = None, + sleeper: Callable[[float], None] | None = None, +) -> int: + config.validate() + reader = reader or ProcfsReader(config.procfs_root) + audit = audit or AuditLogger(sys.stderr) + launcher = launcher or subprocess.Popen + signal_group = signal_group or _signal_process_group + monotonic = monotonic or time.monotonic + sleeper = sleeper or time.sleep + + try: + snapshot = reader.read_snapshot() + except ProcfsError as exc: + return _emit_final( + audit, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + None, + None, + error=str(exc), + ) + + audit.emit( + "preflight", + **_state_fields( + snapshot, + snapshot.used_bytes, + None, + None, + "not_created", + "none", + ), + soft_bytes=config.soft_bytes, + emergency_bytes=config.emergency_bytes, + strict_ceiling_bytes=STRICT_CEILING_BYTES, + ) + + if snapshot.active_swaps: + return _emit_final( + audit, + "startup_swap_active", + EXIT_SWAP_ACTIVE, + "active swap present before command launch", + snapshot, + snapshot.used_bytes, + ) + if snapshot.used_bytes >= config.emergency_bytes: + return _emit_final( + audit, + "startup_emergency_limit", + EXIT_EMERGENCY_LIMIT, + "used_bytes >= emergency_bytes before launch", + snapshot, + snapshot.used_bytes, + ) + if snapshot.used_bytes >= config.soft_bytes: + return _emit_final( + audit, + "startup_soft_limit", + EXIT_SOFT_LIMIT, + "used_bytes >= soft_bytes before launch", + snapshot, + snapshot.used_bytes, + ) + + try: + child = launcher(config.command, start_new_session=True) + except (OSError, ValueError) as exc: + detail = getattr(exc, "strerror", None) or str(exc) + return _emit_final( + audit, + "launch_error", + EXIT_LAUNCH_ERROR, + "command launch failed", + snapshot, + snapshot.used_bytes, + error=detail, + ) + + audit.emit( + "child_started", + **_state_fields( + snapshot, + snapshot.used_bytes, + child, + None, + "active", + "none", + ), + command=list(config.command), + ) + return _monitor_child( + config, + reader, + audit, + child, + snapshot, + signal_group, + monotonic, + sleeper, + ) + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be greater than zero") + return parsed + + +def _positive_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError("value must be greater than zero") + return parsed + + +def parse_args(argv: Sequence[str]) -> WatchdogConfig: + parser = argparse.ArgumentParser( + description=( + "Launch a command in a new process group and stop it before " + "host-wide memory use reaches the 120 GiB Strix validation ceiling." + ) + ) + parser.add_argument( + "--procfs-root", + type=Path, + default=Path("/proc"), + help="procfs root containing meminfo and swaps (default: /proc)", + ) + parser.add_argument( + "--soft-gib", + type=_positive_int, + default=116, + help="send SIGTERM at this many GiB used (default: 116)", + ) + parser.add_argument( + "--emergency-gib", + type=_positive_int, + default=118, + help=( + "send SIGKILL at this many GiB used (default: 118, leaving " + "a 2 GiB sampling margin below 120 GiB)" + ), + ) + parser.add_argument( + "--grace-seconds", + type=_positive_float, + default=DEFAULT_GRACE_SECONDS, + help="maximum time after SIGTERM before SIGKILL (default: 30)", + ) + parser.add_argument( + "--sample-interval-seconds", + type=_positive_float, + default=DEFAULT_SAMPLE_INTERVAL_SECONDS, + help="procfs sampling interval (default: 1)", + ) + parser.add_argument( + "command", + nargs=argparse.REMAINDER, + help="command and arguments, preceded by --", + ) + args = parser.parse_args(argv) + command = tuple(args.command) + if command and command[0] == "--": + command = command[1:] + return WatchdogConfig( + command=command, + procfs_root=args.procfs_root, + soft_bytes=args.soft_gib * GIB, + emergency_bytes=args.emergency_gib * GIB, + grace_seconds=args.grace_seconds, + sample_interval_seconds=args.sample_interval_seconds, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + config = parse_args(argv if argv is not None else sys.argv[1:]) + audit = AuditLogger(sys.stderr) + try: + return run_watchdog(config, audit=audit) + except ValueError as exc: + return _emit_final( + audit, + "configuration_error", + EXIT_PROCFS_ERROR, + "invalid configuration", + None, + None, + error=str(exc), + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index ea937784c5a2..e73ea4286197 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -257,6 +257,17 @@ llama_build_and_test(test-chat-template.cpp) # debug tool for chat template differential analysis (not registered as a test, run it manually) llama_build(test-chat-analysis.cpp) llama_build_and_test(test-log.cpp) + +find_package(Python3 3.10 COMPONENTS Interpreter QUIET) +if (CMAKE_SYSTEM_NAME STREQUAL "Linux" AND Python3_Interpreter_FOUND) + llama_test_cmd( + ${Python3_EXECUTABLE} + NAME test-strix-memory-watchdog + LABEL python + ARGS ${CMAKE_CURRENT_SOURCE_DIR}/test_strix_memory_watchdog.py + ) +endif() + llama_build_and_test( test-peg-parser.cpp peg-parser/simple-tokenize.cpp diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py new file mode 100644 index 000000000000..fa4c4aaec969 --- /dev/null +++ b/tests/test_strix_memory_watchdog.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib.util +import io +import json +import signal +import subprocess +import sys +import tempfile +import unittest +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +SCRIPT_PATH = ( + Path(__file__).resolve().parents[1] + / "scripts" + / "strix_memory_watchdog.py" +) +SPEC = importlib.util.spec_from_file_location( + "strix_memory_watchdog", SCRIPT_PATH +) +assert SPEC is not None +assert SPEC.loader is not None +watchdog = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = watchdog +SPEC.loader.exec_module(watchdog) + + +def snapshot( + used_bytes: int, + *, + total_bytes: int = 200, + active_swaps: tuple[str, ...] = (), +) -> Any: + return watchdog.HostSnapshot( + total_bytes=total_bytes, + available_bytes=total_bytes - used_bytes, + active_swaps=active_swaps, + ) + + +class SequenceReader: + def __init__(self, values: list[Any]): + self.values = values + self.index = 0 + + def read_snapshot(self) -> Any: + index = min(self.index, len(self.values) - 1) + self.index += 1 + value = self.values[index] + if isinstance(value, Exception): + raise value + return value + + +class FakeClock: + def __init__(self): + self.value = 0.0 + + def monotonic(self) -> float: + return self.value + + def sleep(self, seconds: float) -> None: + self.value += seconds + + +class FakeProcess: + def __init__(self, returncode: int | None = None): + self.pid = 4321 + self.returncode = returncode + + def poll(self) -> int | None: + return self.returncode + + def wait(self, timeout: float | None = None) -> int: + if self.returncode is None: + raise subprocess.TimeoutExpired("fake", timeout) + return self.returncode + + +class Harness: + def __init__( + self, + values: list[Any], + process: FakeProcess, + signal_handler: Any | None = None, + ): + self.reader = SequenceReader(values) + self.process = process + self.signal_handler = signal_handler + self.clock = FakeClock() + self.stream = io.StringIO() + self.launched = False + self.signals: list[int] = [] + fixed_time = datetime(2026, 1, 1, tzinfo=timezone.utc) + self.audit = watchdog.AuditLogger( + self.stream, wall_clock=lambda: fixed_time + ) + + def launcher(self, command: tuple[str, ...], **kwargs: Any) -> FakeProcess: + self.launched = True + self.command = command + self.launch_kwargs = kwargs + return self.process + + def signal_group(self, process_group_id: int, signal_number: int) -> str: + self.signals.append(signal_number) + if self.signal_handler is not None: + self.signal_handler(self.process, signal_number) + return f"{signal.Signals(signal_number).name.lower()}_sent" + + def run(self, **overrides: Any) -> int: + config = watchdog.WatchdogConfig( + command=("fake-command",), + soft_bytes=100, + emergency_bytes=150, + grace_seconds=2, + sample_interval_seconds=1, + **overrides, + ) + return watchdog.run_watchdog( + config, + reader=self.reader, + audit=self.audit, + launcher=self.launcher, + signal_group=self.signal_group, + monotonic=self.clock.monotonic, + sleeper=self.clock.sleep, + ) + + def records(self) -> list[dict[str, Any]]: + return [ + json.loads(line) + for line in self.stream.getvalue().splitlines() + ] + + +class TestProcfsParsing(unittest.TestCase): + def test_parses_meminfo_as_integer_bytes_and_allows_zero_swap(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + (root / "meminfo").write_text( + "MemTotal: 131072 kB\n" + "MemFree: 4096 kB\n" + "MemAvailable: 32768 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + + result = watchdog.ProcfsReader(root).read_snapshot() + + self.assertEqual(result.total_bytes, 131072 * 1024) + self.assertEqual(result.available_bytes, 32768 * 1024) + self.assertEqual(result.used_bytes, 98304 * 1024) + self.assertEqual(result.active_swaps, ()) + + def test_rejects_active_swap_entry(self) -> None: + content = ( + "Filename Type Size Used Priority\n" + "/swapfile file 1048572 0 -2\n" + ) + self.assertEqual( + watchdog.ProcfsReader._parse_swaps(content), + ("/swapfile",), + ) + + def test_rejects_malformed_or_missing_procfs_data(self) -> None: + with self.assertRaisesRegex( + watchdog.ProcfsError, "malformed MemAvailable" + ): + watchdog.ProcfsReader._parse_meminfo( + "MemTotal: 10 kB\nMemAvailable: unknown\n" + ) + with self.assertRaisesRegex( + watchdog.ProcfsError, "missing MemAvailable" + ): + watchdog.ProcfsReader._parse_meminfo("MemTotal: 10 kB\n") + with self.assertRaisesRegex( + watchdog.ProcfsError, "malformed swaps header" + ): + watchdog.ProcfsReader._parse_swaps("") + with tempfile.TemporaryDirectory() as temp_dir: + with self.assertRaisesRegex( + watchdog.ProcfsError, "cannot read" + ): + watchdog.ProcfsReader( + Path(temp_dir) + ).read_snapshot() + + +class TestWatchdogBehavior(unittest.TestCase): + def test_configuration_rejects_non_finite_timing(self) -> None: + config = watchdog.WatchdogConfig( + command=("fake-command",), + grace_seconds=float("nan"), + ) + with self.assertRaisesRegex(ValueError, "grace period"): + config.validate() + + def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + (root / "meminfo").write_text( + "MemTotal: 131072 kB\nMemAvailable: 65536 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--sample-interval-seconds", + "0.01", + "--", + sys.executable, + "-c", + "raise SystemExit(23)", + ], + capture_output=True, + check=False, + text=True, + ) + + self.assertEqual(result.returncode, 23) + records = [ + json.loads(line) for line in result.stderr.splitlines() + ] + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual(records[-1]["child_returncode"], 23) + + def test_zero_swap_gate_launches_and_propagates_child_exit(self) -> None: + harness = Harness([snapshot(50)], FakeProcess(returncode=37)) + + result = harness.run() + + self.assertEqual(result, 37) + self.assertTrue(harness.launched) + self.assertTrue(harness.launch_kwargs["start_new_session"]) + final = harness.records()[-1] + self.assertEqual(final["classification"], "child_exit") + self.assertEqual(final["total_bytes"], 200) + self.assertEqual(final["available_bytes"], 150) + self.assertEqual(final["used_bytes"], 50) + self.assertEqual(final["peak_used_bytes"], 50) + self.assertEqual(final["child_status"], "exited") + self.assertEqual(final["process_group_status"], "leader_exited") + + def test_signaled_child_exit_uses_shell_exit_convention(self) -> None: + harness = Harness( + [snapshot(50)], + FakeProcess(returncode=-signal.SIGTERM), + ) + + result = harness.run() + + self.assertEqual(result, 128 + signal.SIGTERM) + + def test_active_swap_rejects_startup_without_launch(self) -> None: + harness = Harness( + [snapshot(50, active_swaps=("/swapfile",))], + FakeProcess(), + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SWAP_ACTIVE) + self.assertFalse(harness.launched) + self.assertEqual( + harness.records()[-1]["classification"], + "startup_swap_active", + ) + + def test_soft_limit_sends_sigterm(self) -> None: + def exit_on_term(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGTERM: + process.returncode = -signal.SIGTERM + + harness = Harness( + [snapshot(50), snapshot(110)], + FakeProcess(), + signal_handler=exit_on_term, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SOFT_LIMIT) + self.assertEqual(harness.signals, [signal.SIGTERM]) + self.assertEqual( + harness.records()[-1]["classification"], "soft_limit" + ) + + def test_emergency_limit_sends_sigkill(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(160)], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_EMERGENCY_LIMIT) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "emergency_limit" + ) + + def test_grace_timeout_escalates_to_sigkill(self) -> None: + def ignore_term(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(110)], + FakeProcess(), + signal_handler=ignore_term, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_GRACE_TIMEOUT) + self.assertEqual( + harness.signals, + [signal.SIGTERM, signal.SIGKILL], + ) + self.assertEqual(harness.clock.value, 2.0) + self.assertEqual( + harness.records()[-1]["classification"], "grace_timeout" + ) + + def test_swap_appearing_during_execution_kills_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [ + snapshot(50), + snapshot(60, active_swaps=("/swapfile",)), + ], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SWAP_ACTIVE) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "swap_appeared" + ) + + def test_runtime_procfs_error_kills_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), watchdog.ProcfsError("missing meminfo")], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_PROCFS_ERROR) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "procfs_error" + ) + + def test_launch_failure_is_explicit(self) -> None: + harness = Harness([snapshot(50)], FakeProcess()) + + def fail_launch( + command: tuple[str, ...], **kwargs: Any + ) -> FakeProcess: + raise FileNotFoundError(2, "No such file or directory") + + harness.launcher = fail_launch + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_LAUNCH_ERROR) + self.assertEqual( + harness.records()[-1]["classification"], "launch_error" + ) + + +if __name__ == "__main__": + unittest.main() From 41dbf04fbabc5c94d6cda4930c7f7b92361e7823 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:20:14 -0700 Subject: [PATCH 02/32] scripts : clean up watchdog process group on signals Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 3 + scripts/strix_memory_watchdog.py | 309 +++++++++++++++++++++++----- tests/test_strix_memory_watchdog.py | 150 ++++++++++++++ 3 files changed, 407 insertions(+), 55 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 4adc0e4be7fe..8d6d264b141b 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -13,6 +13,8 @@ The wrapper performs these checks and actions: - It sends `SIGTERM` to the process group at 116 GiB used. - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. +- It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. +- It applies the same bounded process-group cleanup if an unexpected post-launch error occurs. - It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. @@ -31,6 +33,7 @@ Exit classifications are authoritative in the final JSON record. Operational fai | 5 | emergency threshold reached | | 6 | soft-threshold grace period expired | | 7 | process-group signaling or termination failure | +| 70 | unexpected post-launch error | | 127 | command launch failure | No model, backend, or ROCm package is required to run the unit tests: diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 0dc532c1d027..d8738d5041fd 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -31,10 +31,12 @@ EXIT_EMERGENCY_LIMIT = 5 EXIT_GRACE_TIMEOUT = 6 EXIT_SIGNAL_ERROR = 7 +EXIT_INTERNAL_ERROR = 70 EXIT_LAUNCH_ERROR = 127 MEMINFO_VALUE_RE = re.compile(r"([0-9]+) kB") SWAPS_HEADER = ["Filename", "Type", "Size", "Used", "Priority"] +PARENT_SIGNALS = (signal.SIGINT, signal.SIGTERM) class ProcfsError(RuntimeError): @@ -45,6 +47,12 @@ class ProcessGroupError(RuntimeError): pass +class ParentSignal(RuntimeError): + def __init__(self, signal_number: int): + self.signal_number = signal_number + super().__init__(signal.Signals(signal_number).name) + + class ProcessHandle(Protocol): pid: int @@ -66,6 +74,12 @@ def used_bytes(self) -> int: return self.total_bytes - self.available_bytes +@dataclass +class RuntimeState: + snapshot: HostSnapshot + peak_used_bytes: int + + @dataclass(frozen=True) class WatchdogConfig: command: tuple[str, ...] @@ -257,6 +271,39 @@ def _signal_process_group(process_group_id: int, signal_number: int) -> str: return f"{signal.Signals(signal_number).name.lower()}_sent" +def _process_group_alive(process_group_id: int) -> bool: + try: + os.killpg(process_group_id, 0) + except ProcessLookupError: + return False + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot inspect process group {process_group_id}: {detail}" + ) from exc + return True + + +def _raise_parent_signal(signal_number: int, _frame: object) -> None: + raise ParentSignal(signal_number) + + +def _set_parent_signal_handlers( + handler: signal.Handlers, +) -> dict[int, signal.Handlers]: + previous: dict[int, signal.Handlers] = {} + for signal_number in PARENT_SIGNALS: + previous[signal_number] = signal.signal(signal_number, handler) + return previous + + +def _restore_parent_signal_handlers( + previous: dict[int, signal.Handlers], +) -> None: + for signal_number, handler in previous.items(): + signal.signal(signal_number, handler) + + def _kill_and_finish( audit: AuditLogger, child: ProcessHandle, @@ -323,18 +370,105 @@ def _kill_and_finish( ) +def _graceful_cleanup( + audit: AuditLogger, + child: ProcessHandle, + snapshot: HostSnapshot, + peak_used_bytes: int, + classification: str, + exit_code: int, + reason: str, + graceful_signal: int, + grace_seconds: float, + signal_group: Callable[[int, int], str], + group_alive: Callable[[int], bool], + monotonic: Callable[[], float], + sleeper: Callable[[float], None], + error: str | None = None, +) -> int: + signal_events: list[tuple[int, str]] = [] + try: + group_status = signal_group(child.pid, graceful_signal) + signal_events.append((graceful_signal, group_status)) + deadline = monotonic() + grace_seconds + while monotonic() < deadline: + child.poll() + if not group_alive(child.pid): + break + sleeper(min(0.05, deadline - monotonic())) + child.poll() + if group_alive(child.pid): + group_status = signal_group(child.pid, signal.SIGKILL) + signal_events.append((signal.SIGKILL, group_status)) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "signal_error", + str(exc), + ) + + child_returncode = child.poll() + if child_returncode is None: + try: + child_returncode = child.wait(timeout=5.0) + except subprocess.TimeoutExpired as exc: + return _emit_final( + audit, + "termination_timeout", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "termination_timeout", + str(exc), + ) + + for signal_number, status in signal_events: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child_returncode, + status, + reason, + ), + signal=signal.Signals(signal_number).name, + ) + return _emit_final( + audit, + classification, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + signal_events[-1][1], + error, + ) + + def _monitor_child( config: WatchdogConfig, reader: ProcfsReader, audit: AuditLogger, child: ProcessHandle, - initial_snapshot: HostSnapshot, + state: RuntimeState, signal_group: Callable[[int, int], str], monotonic: Callable[[], float], sleeper: Callable[[float], None], ) -> int: - snapshot = initial_snapshot - peak = snapshot.used_bytes soft_deadline: float | None = None while True: @@ -354,8 +488,8 @@ def _monitor_child( if soft_stop else "child exited" ), - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, child, child_returncode, "leader_exited", @@ -366,8 +500,8 @@ def _monitor_child( return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "grace_timeout", EXIT_GRACE_TIMEOUT, "soft-threshold grace period expired", @@ -375,50 +509,60 @@ def _monitor_child( ) try: - snapshot = reader.read_snapshot() + state.snapshot = reader.read_snapshot() except ProcfsError as exc: return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "procfs_error", EXIT_PROCFS_ERROR, str(exc), signal_group, ) - peak = max(peak, snapshot.used_bytes) + state.peak_used_bytes = max( + state.peak_used_bytes, state.snapshot.used_bytes + ) audit.emit( "sample", **_state_fields( - snapshot, peak, child, None, "active", "none" + state.snapshot, + state.peak_used_bytes, + child, + None, + "active", + "none", ), ) - if snapshot.active_swaps: + if state.snapshot.active_swaps: return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "swap_appeared", EXIT_SWAP_ACTIVE, "active swap appeared during execution", signal_group, ) - if snapshot.used_bytes >= config.emergency_bytes: + if state.snapshot.used_bytes >= config.emergency_bytes: return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "emergency_limit", EXIT_EMERGENCY_LIMIT, "used_bytes >= emergency_bytes", signal_group, ) - if soft_deadline is None and snapshot.used_bytes >= config.soft_bytes: + if ( + soft_deadline is None + and state.snapshot.used_bytes >= config.soft_bytes + ): try: group_status = signal_group(child.pid, signal.SIGTERM) except ProcessGroupError as exc: @@ -427,8 +571,8 @@ def _monitor_child( "signal_error", EXIT_SIGNAL_ERROR, "used_bytes >= soft_bytes", - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, child, child.poll(), "signal_error", @@ -438,8 +582,8 @@ def _monitor_child( audit.emit( "process_group_signal", **_state_fields( - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, child, child.poll(), group_status, @@ -465,6 +609,7 @@ def run_watchdog( audit: AuditLogger | None = None, launcher: Callable[..., ProcessHandle] | None = None, signal_group: Callable[[int, int], str] | None = None, + group_alive: Callable[[int], bool] | None = None, monotonic: Callable[[], float] | None = None, sleeper: Callable[[float], None] | None = None, ) -> int: @@ -473,6 +618,7 @@ def run_watchdog( audit = audit or AuditLogger(sys.stderr) launcher = launcher or subprocess.Popen signal_group = signal_group or _signal_process_group + group_alive = group_alive or _process_group_alive monotonic = monotonic or time.monotonic sleeper = sleeper or time.sleep @@ -532,42 +678,95 @@ def run_watchdog( snapshot.used_bytes, ) + previous_mask: set[signal.Signals] | None = signal.pthread_sigmask( + signal.SIG_BLOCK, PARENT_SIGNALS + ) + previous_handlers: dict[int, signal.Handlers] = {} + child: ProcessHandle | None = None + state = RuntimeState(snapshot, snapshot.used_bytes) try: - child = launcher(config.command, start_new_session=True) - except (OSError, ValueError) as exc: - detail = getattr(exc, "strerror", None) or str(exc) - return _emit_final( + try: + child = launcher(config.command, start_new_session=True) + except (OSError, ValueError) as exc: + detail = getattr(exc, "strerror", None) or str(exc) + return _emit_final( + audit, + "launch_error", + EXIT_LAUNCH_ERROR, + "command launch failed", + snapshot, + snapshot.used_bytes, + error=detail, + ) + + previous_handlers = _set_parent_signal_handlers( + _raise_parent_signal + ) + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + previous_mask = None + audit.emit( + "child_started", + **_state_fields( + state.snapshot, + state.peak_used_bytes, + child, + None, + "active", + "none", + ), + command=list(config.command), + ) + return _monitor_child( + config, + reader, audit, - "launch_error", - EXIT_LAUNCH_ERROR, - "command launch failed", - snapshot, - snapshot.used_bytes, - error=detail, + child, + state, + signal_group, + monotonic, + sleeper, ) - - audit.emit( - "child_started", - **_state_fields( - snapshot, - snapshot.used_bytes, + except ParentSignal as exc: + _set_parent_signal_handlers(signal.SIG_IGN) + signal_name = signal.Signals(exc.signal_number).name + return _graceful_cleanup( + audit, child, - None, - "active", - "none", - ), - command=list(config.command), - ) - return _monitor_child( - config, - reader, - audit, - child, - snapshot, - signal_group, - monotonic, - sleeper, - ) + state.snapshot, + state.peak_used_bytes, + "parent_signal", + 128 + exc.signal_number, + f"wrapper received {signal_name}", + exc.signal_number, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + ) + except Exception as exc: + _set_parent_signal_handlers(signal.SIG_IGN) + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "internal_error", + EXIT_INTERNAL_ERROR, + "unexpected post-launch exception", + signal.SIGTERM, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + f"{type(exc).__name__}: {exc}", + ) + finally: + if previous_mask is not None: + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + if previous_handlers: + _restore_parent_signal_handlers(previous_handlers) def _positive_int(value: str) -> int: diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index fa4c4aaec969..c0dbc9c1fcb3 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -5,10 +5,12 @@ import importlib.util import io import json +import os import signal import subprocess import sys import tempfile +import time import unittest from datetime import datetime, timezone from pathlib import Path @@ -113,6 +115,9 @@ def signal_group(self, process_group_id: int, signal_number: int) -> str: self.signal_handler(self.process, signal_number) return f"{signal.Signals(signal_number).name.lower()}_sent" + def group_alive(self, process_group_id: int) -> bool: + return self.process.returncode is None + def run(self, **overrides: Any) -> int: config = watchdog.WatchdogConfig( command=("fake-command",), @@ -128,6 +133,7 @@ def run(self, **overrides: Any) -> int: audit=self.audit, launcher=self.launcher, signal_group=self.signal_group, + group_alive=self.group_alive, monotonic=self.clock.monotonic, sleeper=self.clock.sleep, ) @@ -196,6 +202,128 @@ def test_rejects_malformed_or_missing_procfs_data(self) -> None: class TestWatchdogBehavior(unittest.TestCase): + @staticmethod + def _process_is_running(process_id: int) -> bool: + result = subprocess.run( + ["ps", "-o", "stat=", "-p", str(process_id)], + capture_output=True, + check=False, + text=True, + ) + return result.returncode == 0 and not result.stdout.lstrip().startswith( + "Z" + ) + + def test_parent_signals_leave_no_child_or_grandchild(self) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGINT,signal.SIG_IGN);" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n');" + " time.sleep(30)\n" + ) + for signal_number in (signal.SIGINT, signal.SIGTERM): + with self.subTest(signal=signal.Signals(signal_number).name): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + (root / "meminfo").write_text( + "MemTotal: 131072 kB\n" + "MemAvailable: 65536 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + audit_path = root / "audit.jsonl" + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + self.fail( + "child process group did not start" + ) + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + time.sleep(0.05) + wrapper.send_signal(signal_number) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + if child_pid is not None: + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + self.assertEqual( + wrapper.returncode, 128 + signal_number + ) + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + self.assertEqual( + records[-1]["classification"], "parent_signal" + ) + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual( + forwarded[0], + signal.Signals(signal_number).name, + ) + self.assertEqual(forwarded[-1], "SIGKILL") + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse( + self._process_is_running(process_id) + ) + def test_configuration_rejects_non_finite_timing(self) -> None: config = watchdog.WatchdogConfig( command=("fake-command",), @@ -384,6 +512,28 @@ def exit_on_kill(process: FakeProcess, signal_number: int) -> None: harness.records()[-1]["classification"], "procfs_error" ) + def test_unexpected_monitor_error_cleans_up_process_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), RuntimeError("unexpected")], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_INTERNAL_ERROR) + self.assertEqual( + harness.signals, + [signal.SIGTERM, signal.SIGKILL], + ) + final = harness.records()[-1] + self.assertEqual(final["classification"], "internal_error") + self.assertIn("RuntimeError: unexpected", final["error"]) + def test_launch_failure_is_explicit(self) -> None: harness = Harness([snapshot(50)], FakeProcess()) From da5ce95a8c08a9e01c01cd3519de188c8dca2512 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:37:36 -0700 Subject: [PATCH 03/32] scripts : preserve child signals and monitor descendants Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 3 +- scripts/strix_memory_watchdog.py | 129 ++++++++++++----- tests/test_strix_memory_watchdog.py | 206 +++++++++++++++++++++++++--- 3 files changed, 282 insertions(+), 56 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 8d6d264b141b..5a6558f0a9c0 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -14,6 +14,7 @@ The wrapper performs these checks and actions: - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. - It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. +- It checks the process group after the direct child exits and cleans up remaining descendants before returning the child's classification. - It applies the same bounded process-group cleanup if an unexpected post-launch error occurs. - It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. @@ -21,7 +22,7 @@ The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. -The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Child standard input, standard output, and standard error are inherited unchanged. +The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Signal records are written immediately after each process-group signal. Child standard input, standard output, and standard error are inherited unchanged. Exit classifications are authoritative in the final JSON record. Operational failures use these exit codes: diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index d8738d5041fd..79d7a0013e77 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -378,18 +378,32 @@ def _graceful_cleanup( classification: str, exit_code: int, reason: str, - graceful_signal: int, + graceful_signal: int | None, grace_seconds: float, signal_group: Callable[[int, int], str], group_alive: Callable[[int], bool], monotonic: Callable[[], float], sleeper: Callable[[float], None], + process_group_status: str = "active", error: str | None = None, ) -> int: - signal_events: list[tuple[int, str]] = [] try: - group_status = signal_group(child.pid, graceful_signal) - signal_events.append((graceful_signal, group_status)) + if graceful_signal is not None: + process_group_status = signal_group( + child.pid, graceful_signal + ) + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + reason, + ), + signal=signal.Signals(graceful_signal).name, + ) deadline = monotonic() + grace_seconds while monotonic() < deadline: child.poll() @@ -398,8 +412,21 @@ def _graceful_cleanup( sleeper(min(0.05, deadline - monotonic())) child.poll() if group_alive(child.pid): - group_status = signal_group(child.pid, signal.SIGKILL) - signal_events.append((signal.SIGKILL, group_status)) + process_group_status = signal_group( + child.pid, signal.SIGKILL + ) + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + reason, + ), + signal="SIGKILL", + ) except ProcessGroupError as exc: return _emit_final( audit, @@ -432,19 +459,6 @@ def _graceful_cleanup( str(exc), ) - for signal_number, status in signal_events: - audit.emit( - "process_group_signal", - **_state_fields( - snapshot, - peak_used_bytes, - child, - child_returncode, - status, - reason, - ), - signal=signal.Signals(signal_number).name, - ) return _emit_final( audit, classification, @@ -454,7 +468,7 @@ def _graceful_cleanup( peak_used_bytes, child, child_returncode, - signal_events[-1][1], + process_group_status, error, ) @@ -466,6 +480,7 @@ def _monitor_child( child: ProcessHandle, state: RuntimeState, signal_group: Callable[[int, int], str], + group_alive: Callable[[int], bool], monotonic: Callable[[], float], sleeper: Callable[[float], None], ) -> int: @@ -475,19 +490,50 @@ def _monitor_child( child_returncode = child.poll() if child_returncode is not None: soft_stop = soft_deadline is not None + classification = "soft_limit" if soft_stop else "child_exit" + exit_code = EXIT_SOFT_LIMIT if soft_stop else ( + 128 - child_returncode + if child_returncode < 0 + else child_returncode + ) + reason = ( + "child exited during soft-threshold grace period" + if soft_stop + else "child exited" + ) + if group_alive(child.pid): + grace_seconds = config.grace_seconds + graceful_signal: int | None = signal.SIGTERM + group_status = "active" + if soft_stop: + grace_seconds = max( + 0.0, soft_deadline - monotonic() + ) + graceful_signal = None + group_status = "sigterm_sent" + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + classification, + exit_code, + ( + f"{reason}; process group members still running" + ), + graceful_signal, + grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + group_status, + ) return _emit_final( audit, - "soft_limit" if soft_stop else "child_exit", - EXIT_SOFT_LIMIT if soft_stop else ( - 128 - child_returncode - if child_returncode < 0 - else child_returncode - ), - ( - "child exited during soft-threshold grace period" - if soft_stop - else "child exited" - ), + classification, + exit_code, + reason, state.snapshot, state.peak_used_bytes, child, @@ -678,15 +724,25 @@ def run_watchdog( snapshot.used_bytes, ) - previous_mask: set[signal.Signals] | None = signal.pthread_sigmask( + previous_mask = signal.pthread_sigmask( signal.SIG_BLOCK, PARENT_SIGNALS ) + mask_restored = False previous_handlers: dict[int, signal.Handlers] = {} child: ProcessHandle | None = None state = RuntimeState(snapshot, snapshot.used_bytes) try: + launch_mask = previous_mask + + def restore_child_signal_mask() -> None: + signal.pthread_sigmask(signal.SIG_SETMASK, launch_mask) + try: - child = launcher(config.command, start_new_session=True) + child = launcher( + config.command, + start_new_session=True, + preexec_fn=restore_child_signal_mask, + ) except (OSError, ValueError) as exc: detail = getattr(exc, "strerror", None) or str(exc) return _emit_final( @@ -703,7 +759,7 @@ def run_watchdog( _raise_parent_signal ) signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) - previous_mask = None + mask_restored = True audit.emit( "child_started", **_state_fields( @@ -723,6 +779,7 @@ def run_watchdog( child, state, signal_group, + group_alive, monotonic, sleeper, ) @@ -760,10 +817,10 @@ def run_watchdog( group_alive, monotonic, sleeper, - f"{type(exc).__name__}: {exc}", + error=f"{type(exc).__name__}: {exc}", ) finally: - if previous_mask is not None: + if not mask_restored: signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) if previous_handlers: _restore_parent_signal_handlers(previous_handlers) diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index c0dbc9c1fcb3..2e59eec27512 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -214,6 +214,17 @@ def _process_is_running(process_id: int) -> bool: "Z" ) + @staticmethod + def _write_procfs_fixture(root: Path) -> None: + (root / "meminfo").write_text( + "MemTotal: 131072 kB\nMemAvailable: 65536 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + def test_parent_signals_leave_no_child_or_grandchild(self) -> None: child_code = ( "import os,signal,sys,time;" @@ -232,15 +243,7 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) pid_file = root / "pids" - (root / "meminfo").write_text( - "MemTotal: 131072 kB\n" - "MemAvailable: 65536 kB\n", - encoding="utf-8", - ) - (root / "swaps").write_text( - "Filename Type Size Used Priority\n", - encoding="utf-8", - ) + self._write_procfs_fixture(root) audit_path = root / "audit.jsonl" with audit_path.open("w", encoding="utf-8") as audit: wrapper = subprocess.Popen( @@ -303,16 +306,29 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: self.assertEqual( records[-1]["classification"], "parent_signal" ) - forwarded = [ - record["signal"] + signal_records = [ + record for record in records if record["event"] == "process_group_signal" ] + forwarded = [ + record["signal"] for record in signal_records + ] self.assertEqual( forwarded[0], signal.Signals(signal_number).name, ) self.assertEqual(forwarded[-1], "SIGKILL") + self.assertEqual( + signal_records[0]["child_status"], "running" + ) + self.assertIsNone( + signal_records[0]["child_returncode"] + ) + self.assertLess( + signal_records[0]["timestamp"], + signal_records[-1]["timestamp"], + ) for process_id in (child_pid, grandchild_pid): deadline = time.monotonic() + 2 while ( @@ -324,6 +340,165 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: self._process_is_running(process_id) ) + def test_child_sigterm_handler_exits_without_escalation(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " open(sys.argv[2],'w').write('handled\\n')\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "open(sys.argv[1],'w').write(f'{os.getpid()}\\n')\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + ready_path = root / "ready" + handled_path = root / "handled" + audit_path = root / "audit.jsonl" + self._write_procfs_fixture(root) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.5", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(ready_path), + str(handled_path), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not ready_path.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("SIGTERM child did not become ready") + time.sleep(0.01) + child_pid = int( + ready_path.read_text(encoding="utf-8").strip() + ) + try: + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(wrapper.returncode, 128 + signal.SIGTERM) + self.assertTrue(handled_path.exists()) + self.assertEqual(forwarded, ["SIGTERM"]) + self.assertEqual(records[-1]["child_returncode"], 0) + + def test_leader_exit_cleans_up_surviving_grandchild(self) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n')\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + audit_path = root / "audit.jsonl" + self._write_procfs_fixture(root) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("leader process did not write child PIDs") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + try: + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(wrapper.returncode, 0) + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual(records[-1]["child_returncode"], 0) + self.assertEqual(forwarded, ["SIGTERM", "SIGKILL"]) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + def test_configuration_rejects_non_finite_timing(self) -> None: config = watchdog.WatchdogConfig( command=("fake-command",), @@ -335,14 +510,7 @@ def test_configuration_rejects_non_finite_timing(self) -> None: def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) - (root / "meminfo").write_text( - "MemTotal: 131072 kB\nMemAvailable: 65536 kB\n", - encoding="utf-8", - ) - (root / "swaps").write_text( - "Filename Type Size Used Priority\n", - encoding="utf-8", - ) + self._write_procfs_fixture(root) result = subprocess.run( [ sys.executable, From 93aff41b19864eb89b1cecff358fd50857440d31 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:45:06 -0700 Subject: [PATCH 04/32] scripts : classify soft descendant escalation as timeout Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 1 + scripts/strix_memory_watchdog.py | 22 +++++- tests/test_strix_memory_watchdog.py | 108 ++++++++++++++++++++++++++++ 3 files changed, 130 insertions(+), 1 deletion(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 5a6558f0a9c0..ccd04f7ab667 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -12,6 +12,7 @@ The wrapper performs these checks and actions: - It calculates used memory as `MemTotal - MemAvailable`. Linux reports these fields in KiB, so the wrapper multiplies each value by 1024 and keeps all accounting as integer bytes. - It sends `SIGTERM` to the process group at 116 GiB used. - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. +- It reports `grace_timeout` if any descendant requires `SIGKILL` after the soft-threshold grace period, even when the direct child exited earlier. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. - It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. - It checks the process group after the direct child exits and cleans up remaining descendants before returning the child's classification. diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 79d7a0013e77..00a9fab78d81 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -385,8 +385,10 @@ def _graceful_cleanup( monotonic: Callable[[], float], sleeper: Callable[[float], None], process_group_status: str = "active", + escalation_result: tuple[str, int, str] | None = None, error: str | None = None, ) -> int: + escalated = False try: if graceful_signal is not None: process_group_status = signal_group( @@ -412,9 +414,15 @@ def _graceful_cleanup( sleeper(min(0.05, deadline - monotonic())) child.poll() if group_alive(child.pid): + escalated = True process_group_status = signal_group( child.pid, signal.SIGKILL ) + signal_reason = ( + escalation_result[2] + if escalation_result is not None + else reason + ) audit.emit( "process_group_signal", **_state_fields( @@ -423,7 +431,7 @@ def _graceful_cleanup( child, child.poll(), process_group_status, - reason, + signal_reason, ), signal="SIGKILL", ) @@ -459,6 +467,8 @@ def _graceful_cleanup( str(exc), ) + if escalated and escalation_result is not None: + classification, exit_code, reason = escalation_result return _emit_final( audit, classification, @@ -528,6 +538,16 @@ def _monitor_child( monotonic, sleeper, group_status, + ( + ( + "grace_timeout", + EXIT_GRACE_TIMEOUT, + "soft-threshold grace period expired with " + "process group members still running", + ) + if soft_stop + else None + ), ) return _emit_final( audit, diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 2e59eec27512..898298efb9b9 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -499,6 +499,114 @@ def test_leader_exit_cleans_up_surviving_grandchild(self) -> None: time.sleep(0.01) self.assertFalse(self._process_is_running(process_id)) + def test_soft_limit_descendant_escalation_is_grace_timeout(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "grandchild=os.fork()\n" + "if grandchild == 0:\n" + " signal.signal(signal.SIGTERM,signal.SIG_IGN)\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n')\n" + " time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + audit_path = root / "audit.jsonl" + (root / "meminfo").write_text( + "MemTotal: 3145728 kB\n" + "MemAvailable: 2621440 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--soft-gib", + "1", + "--emergency-gib", + "2", + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("soft-limit process group did not start") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + next_meminfo = root / "meminfo.next" + next_meminfo.write_text( + "MemTotal: 3145728 kB\n" + "MemAvailable: 1572864 kB\n", + encoding="utf-8", + ) + next_meminfo.replace(root / "meminfo") + try: + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + final = records[-1] + self.assertEqual(wrapper.returncode, watchdog.EXIT_GRACE_TIMEOUT) + self.assertEqual(final["classification"], "grace_timeout") + self.assertEqual(final["child_returncode"], 0) + self.assertEqual(forwarded, ["SIGTERM", "SIGKILL"]) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + def test_configuration_rejects_non_finite_timing(self) -> None: config = watchdog.WatchdogConfig( command=("fake-command",), From 1193241b1398d0151aeca607156083e1ce99aada Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:22:51 -0700 Subject: [PATCH 05/32] deepseek41 : admit unified host memory before allocation Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- common/arg.cpp | 58 ++++- common/common.cpp | 8 + common/common.h | 9 +- include/llama.h | 12 +- src/CMakeLists.txt | 1 + src/llama-context.cpp | 5 +- src/llama-dsv41-admission.cpp | 382 ++++++++++++++++++++++++++++ src/llama-dsv41-admission.h | 83 ++++++ src/llama-model.cpp | 8 + src/llama-model.h | 2 + src/models/deepseek41.cpp | 132 +++++++--- src/models/models.h | 4 + tests/CMakeLists.txt | 1 + tests/test-deepseek41-admission.cpp | 263 +++++++++++++++++++ 14 files changed, 929 insertions(+), 39 deletions(-) create mode 100644 src/llama-dsv41-admission.cpp create mode 100644 src/llama-dsv41-admission.h create mode 100644 tests/test-deepseek41-admission.cpp diff --git a/common/arg.cpp b/common/arg.cpp index b9ab6b1a2e5d..47b08f0298e9 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2839,9 +2839,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_NGRAM_DIRECT_IO")); add_opt(common_arg( {"--expert-cache-slots"}, "N", - "DeepSeek V4.1 routed experts resident per layer; requires --expert-cache-mib", + "maximum DeepSeek V4.1 routed experts resident per layer; 0 auto-fits", [](common_params & params, int value) { - if (value <= 0) { + if (value < 0) { throw std::invalid_argument("invalid value"); } params.expert_cache_slots = value; @@ -2849,14 +2849,64 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_EXPERT_CACHE_SLOTS")); add_opt(common_arg( {"--expert-cache-mib"}, "MiB", - "aggregate DeepSeek V4.1 fixed expert slot-tensor capacity; requires --expert-cache-slots", + "exact aggregate DeepSeek V4.1 expert cache capacity; 0 auto-fits", [](common_params & params, int value) { - if (value <= 0) { + if (value < 0) { throw std::invalid_argument("invalid value"); } params.expert_cache_mib = value; } ).set_env("LLAMA_ARG_EXPERT_CACHE_MIB")); + add_opt(common_arg( + {"--dsv41-memory-soft-mib"}, "MiB", + string_format("DeepSeek V4.1 total host-use startup target (default: %d)", params.dsv41_memory_soft_mib), + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.dsv41_memory_soft_mib = value; + } + ).set_env("LLAMA_ARG_DSV41_MEMORY_SOFT_MIB")); + add_opt(common_arg( + {"--dsv41-memory-watchdog-mib"}, "MiB", + string_format("DeepSeek V4.1 external watchdog emergency threshold (default: %d)", params.dsv41_memory_watchdog_mib), + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.dsv41_memory_watchdog_mib = value; + } + ).set_env("LLAMA_ARG_DSV41_MEMORY_WATCHDOG_MIB")); + add_opt(common_arg( + {"--dsv41-memory-hard-mib"}, "MiB", + string_format("DeepSeek V4.1 strict host-use ceiling (default: %d)", params.dsv41_memory_hard_mib), + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.dsv41_memory_hard_mib = value; + } + ).set_env("LLAMA_ARG_DSV41_MEMORY_HARD_MIB")); + add_opt(common_arg( + {"--dsv41-memory-safety-margin-mib"}, "MiB", + string_format("DeepSeek V4.1 explicit startup safety margin (default: %d)", params.dsv41_memory_safety_margin_mib), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.dsv41_memory_safety_margin_mib = value; + } + ).set_env("LLAMA_ARG_DSV41_MEMORY_SAFETY_MARGIN_MIB")); + add_opt(common_arg( + {"--dsv41-procfs-root"}, "PATH", + "procfs root used by DeepSeek V4.1 host-memory admission (default: /proc)", + [](common_params & params, const std::string & value) { + if (value.empty()) { + throw std::invalid_argument("invalid value"); + } + params.dsv41_procfs_root = value; + } + ).set_env("LLAMA_ARG_DSV41_PROCFS_ROOT")); add_opt(common_arg( {"-cmoe", "--cpu-moe"}, "keep all Mixture of Experts (MoE) weights in the CPU", diff --git a/common/common.cpp b/common/common.cpp index b80b7fa5f809..bf3852b0d814 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1699,6 +1699,14 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.ple_cache_mb = params.ple_cache_mb; mparams.expert_cache_slots = params.expert_cache_slots; mparams.expert_cache_bytes = params.expert_cache_mib > 0 ? (size_t) params.expert_cache_mib << 20 : 0; + mparams.dsv41_memory_soft_bytes = (uint64_t) params.dsv41_memory_soft_mib << 20; + mparams.dsv41_memory_watchdog_bytes = (uint64_t) params.dsv41_memory_watchdog_mib << 20; + mparams.dsv41_memory_hard_bytes = (uint64_t) params.dsv41_memory_hard_mib << 20; + mparams.dsv41_memory_safety_margin_bytes = (uint64_t) params.dsv41_memory_safety_margin_mib << 20; + mparams.dsv41_admission_context = params.n_ctx == 0 ? 32768 : params.n_ctx; + mparams.dsv41_admission_sequences = params.n_parallel; + mparams.dsv41_admission_ubatch = params.n_ubatch; + mparams.dsv41_procfs_root = params.dsv41_procfs_root.c_str(); if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; diff --git a/common/common.h b/common/common.h index a838f90b52a8..2efd7e80d707 100644 --- a/common/common.h +++ b/common/common.h @@ -626,8 +626,13 @@ struct common_params { bool ple_direct_io = true; // ... read with O_DIRECT int32_t ple_io_threads = 64; // ... parallel readers (random 4 KiB reads: this NVMe gives 62k IOPS at 16, 130k at 64, ~160k at 128+) int32_t ple_cache_mb = 256; // ... row cache, 0 disables - int32_t expert_cache_slots = 0; // DeepSeek V4.1 routed experts resident per layer - int32_t expert_cache_mib = 0; // aggregate fixed slot-tensor capacity + int32_t expert_cache_slots = 0; // DeepSeek V4.1 routed experts resident per layer, 0 auto-fits + int32_t expert_cache_mib = 0; // exact aggregate cache bytes, 0 auto-fits + int32_t dsv41_memory_soft_mib = 116*1024; + int32_t dsv41_memory_watchdog_mib = 118*1024; + int32_t dsv41_memory_hard_mib = 120*1024; + int32_t dsv41_memory_safety_margin_mib = 2*1024; + std::string dsv41_procfs_root = "/proc"; bool single_turn = false; // single turn chat conversation diff --git a/include/llama.h b/include/llama.h index 43c23eb55d84..dc86b042314b 100644 --- a/include/llama.h +++ b/include/llama.h @@ -348,10 +348,20 @@ extern "C" { int32_t ple_io_threads; // parallel pread workers int32_t ple_cache_mb; // in-memory cache of recently read rows, 0 disables - // DeepSeek V4.1 routed-expert cache. Both values must be non-zero. + // DeepSeek V4.1 routed-expert cache. Zero values auto-fit within admission. size_t expert_cache_bytes; int32_t expert_cache_slots; + // DeepSeek V4.1 unified host-memory admission. Zero values use safe Strix defaults. + uint64_t dsv41_memory_soft_bytes; + uint64_t dsv41_memory_watchdog_bytes; + uint64_t dsv41_memory_hard_bytes; + uint64_t dsv41_memory_safety_margin_bytes; + uint32_t dsv41_admission_context; + uint32_t dsv41_admission_sequences; + uint32_t dsv41_admission_ubatch; + const char * dsv41_procfs_root; + // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5ae5baf8a71a..45643686116f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -17,6 +17,7 @@ add_library(llama llama-chat.cpp llama-context.cpp llama-cparams.cpp + llama-dsv41-admission.cpp llama-dsv41.cpp llama-dsv41-engram.cpp llama-dsv41-expert.cpp diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ff036ee0abd5..1cdf6704d418 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -130,7 +130,8 @@ llama_context::llama_context( cparams.rope_scaling_type = params.rope_scaling_type; cparams.pooling_type = params.pooling_type; - cparams.n_ctx = params.n_ctx == 0 ? hparams.n_ctx_train : params.n_ctx; + const uint32_t n_ctx_default = model.default_context_size(); + cparams.n_ctx = params.n_ctx == 0 ? (n_ctx_default == 0 ? hparams.n_ctx_train : n_ctx_default) : params.n_ctx; cparams.rope_freq_base = params.rope_freq_base == 0.0f ? hparams.rope_freq_base_train : params.rope_freq_base; cparams.rope_freq_scale = params.rope_freq_scale == 0.0f ? hparams.rope_freq_scale_train : params.rope_freq_scale; @@ -255,6 +256,8 @@ llama_context::llama_context( cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ? cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max); + model.validate_context_params(cparams); + // Initialize backend samplers here so they are part of the sampling graph // before the reserve passes run later in this function. This avoids a later // re-reserve when graph nodes change. diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp new file mode 100644 index 000000000000..d442ad52cfc6 --- /dev/null +++ b/src/llama-dsv41-admission.cpp @@ -0,0 +1,382 @@ +#include "llama-dsv41-admission.h" + +#include "llama-dsv41.h" +#include "llama-impl.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +uint64_t checked_add(uint64_t a, uint64_t b, const char * category) { + if (b > std::numeric_limits::max() - a) { + throw std::runtime_error(std::string("DeepSeek V4.1 memory admission overflow: ") + category); + } + return a + b; +} + +uint64_t checked_mul(uint64_t a, uint64_t b, const char * category) { + if (a != 0 && b > std::numeric_limits::max()/a) { + throw std::runtime_error(std::string("DeepSeek V4.1 memory admission overflow: ") + category); + } + return a*b; +} + +uint64_t parse_u64(const std::string & value, const char * field) { + uint64_t result = 0; + const char * begin = value.data(); + const char * end = begin + value.size(); + const auto parsed = std::from_chars(begin, end, result); + if (parsed.ec != std::errc() || parsed.ptr != end) { + throw std::runtime_error(std::string("DeepSeek V4.1 procfs malformed integer: ") + field); + } + return result; +} + +uint64_t read_meminfo_value( + const std::string & line, + const char * expected_key) { + std::istringstream stream(line); + std::string key; + std::string value; + std::string unit; + std::string extra; + if (!(stream >> key >> value >> unit) || stream >> extra || + key != std::string(expected_key) + ":" || unit != "kB") { + throw std::runtime_error(std::string("DeepSeek V4.1 procfs malformed field: ") + expected_key); + } + return checked_mul(parse_u64(value, expected_key), 1024, expected_key); +} + +void validate_context(uint32_t n_ctx) { + switch (n_ctx) { + case 32768: + case 65536: + case 98304: + case 131072: + return; + default: + throw std::runtime_error( + "DeepSeek V4.1 memory admission context must be one of 32768, 65536, 98304, or 131072"); + } +} + +[[noreturn]] void reject( + const char * category, + const llama_dsv41_admission_result & result, + const std::string & detail) { + llama_dsv41_admission_result failure = result; + failure.category = category; + throw std::runtime_error(failure.describe() + ", detail=" + detail); +} + +} + +llama_dsv41_host_memory llama_dsv41_read_host_memory(const std::string & procfs_root) { +#ifndef __linux__ + if (procfs_root == "/proc") { + throw std::runtime_error("DeepSeek V4.1 memory admission requires Linux procfs"); + } +#endif + const std::string root = procfs_root.empty() ? "/proc" : procfs_root; + std::ifstream meminfo(root + "/meminfo"); + if (!meminfo) { + throw std::runtime_error("DeepSeek V4.1 memory admission cannot read " + root + "/meminfo"); + } + + llama_dsv41_host_memory result; + bool have_total = false; + bool have_available = false; + std::string line; + while (std::getline(meminfo, line)) { + if (line.rfind("MemTotal:", 0) == 0) { + if (have_total) { + throw std::runtime_error("DeepSeek V4.1 procfs has duplicate MemTotal"); + } + result.total = read_meminfo_value(line, "MemTotal"); + have_total = true; + } else if (line.rfind("MemAvailable:", 0) == 0) { + if (have_available) { + throw std::runtime_error("DeepSeek V4.1 procfs has duplicate MemAvailable"); + } + result.available = read_meminfo_value(line, "MemAvailable"); + have_available = true; + } + } + if (!meminfo.eof() || !have_total || !have_available || result.available > result.total) { + throw std::runtime_error("DeepSeek V4.1 procfs meminfo is missing or invalid"); + } + result.used = result.total - result.available; + + std::ifstream swaps(root + "/swaps"); + if (!swaps) { + throw std::runtime_error("DeepSeek V4.1 memory admission cannot read " + root + "/swaps"); + } + if (!std::getline(swaps, line)) { + throw std::runtime_error("DeepSeek V4.1 procfs swaps header is missing"); + } + { + std::istringstream header(line); + std::string filename; + std::string type; + std::string size; + std::string used; + std::string priority; + std::string extra; + if (!(header >> filename >> type >> size >> used >> priority) || header >> extra || + filename != "Filename" || type != "Type" || size != "Size" || + used != "Used" || priority != "Priority") { + throw std::runtime_error("DeepSeek V4.1 procfs swaps header is malformed"); + } + } + while (std::getline(swaps, line)) { + if (line.empty()) { + continue; + } + std::istringstream entry(line); + std::string filename; + std::string type; + std::string size; + std::string used; + std::string priority; + std::string extra; + if (!(entry >> filename >> type >> size >> used >> priority) || entry >> extra) { + throw std::runtime_error("DeepSeek V4.1 procfs swaps entry is malformed"); + } + const uint64_t size_bytes = checked_mul(parse_u64(size, "swap size"), 1024, "swap size"); + parse_u64(used, "swap used"); + int64_t priority_value = 0; + const auto parsed_priority = std::from_chars( + priority.data(), priority.data() + priority.size(), priority_value); + if (parsed_priority.ec != std::errc() || + parsed_priority.ptr != priority.data() + priority.size()) { + throw std::runtime_error("DeepSeek V4.1 procfs malformed integer: swap priority"); + } + result.swap_entries = checked_add(result.swap_entries, 1, "swap entries"); + result.swap_bytes = checked_add(result.swap_bytes, size_bytes, "swap bytes"); + } + if (!swaps.eof()) { + throw std::runtime_error("DeepSeek V4.1 procfs swaps read failed"); + } + return result; +} + +uint64_t llama_dsv41_expert_payload_bytes(const std::vector & tensors) { + uint64_t result = 0; + for (const auto & tensor : tensors) { + llama_expert_store_validate_tensor(tensor); + result = checked_add( + result, + checked_mul(tensor.nb[2], tensor.ne[2], "expert payload"), + "expert payload"); + } + return result; +} + +uint64_t llama_dsv41_estimate_graph_workspace(uint32_t n_ctx, uint32_t n_ubatch) { + validate_context(n_ctx); + if (n_ubatch == 0 || n_ubatch > 2048) { + throw std::runtime_error("DeepSeek V4.1 bounded admission requires n_ubatch in 1..2048"); + } + const uint64_t base = 7688ULL << 20; + return checked_add(base, checked_mul(n_ctx, 7424, "graph workspace"), "graph workspace"); +} + +uint64_t llama_dsv41_engram_staging_bytes(uint32_t n_ubatch) { + const uint64_t ids = checked_mul( + checked_mul(n_ubatch, LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, "Engram row IDs"), + sizeof(uint32_t), + "Engram row IDs"); + const uint64_t decoded = checked_mul( + checked_mul(n_ubatch, LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, "Engram decoded rows"), + sizeof(float), + "Engram decoded rows"); + return checked_add(checked_add(ids, decoded, "Engram staging"), n_ubatch, "Engram staging"); +} + +uint64_t llama_dsv41_output_bytes(uint32_t n_vocab, uint32_t n_ubatch) { + if (n_vocab == 0 || n_ubatch == 0) { + throw std::runtime_error("DeepSeek V4.1 output accounting dimensions must be non-zero"); + } + const uint64_t floats = checked_mul( + checked_mul(n_vocab, n_ubatch, "output floats"), + 2*sizeof(float), + "output floats"); + const uint64_t tokens = checked_mul( + checked_mul(n_vocab, n_ubatch, "output tokens"), + sizeof(int32_t), + "output tokens"); + return checked_add(floats, tokens, "outputs"); +} + +llama_dsv41_admission_result llama_dsv41_admit( + const llama_dsv41_host_memory & host, + uint64_t dense_tensor_bytes, + const std::vector & expert_tensors, + const llama_dsv41_admission_params & params) { + llama_dsv41_admission_result result; + result.host_total = host.total; + result.host_available = host.available; + result.host_used = host.used; + result.dense_tensor_bytes = dense_tensor_bytes; + result.soft_bytes = params.soft_bytes; + result.watchdog_bytes = params.watchdog_bytes; + result.hard_bytes = params.hard_bytes; + result.safety_margin_bytes = params.safety_margin_bytes; + result.device_reported_bytes_ignored = params.device_reported_bytes; + result.n_ctx = params.n_ctx; + result.n_seq = params.n_seq; + result.n_ubatch = params.n_ubatch; + + if (host.total == 0 || host.available > host.total || host.used != host.total - host.available) { + reject("host", result, "host memory snapshot is invalid"); + } + if (host.swap_entries != 0 || host.swap_bytes != 0) { + reject("swap", result, format( + "%llu configured swap entries (%llu bytes)", + (unsigned long long) host.swap_entries, + (unsigned long long) host.swap_bytes)); + } + if (!params.direct_io) { + reject("direct_io", result, "buffered expert or Engram I/O is not bounded"); + } + if (!params.unified_memory) { + reject("unified_memory", result, "Strix admission requires one unified host/GPU memory pool"); + } + if (params.soft_bytes == 0 || params.soft_bytes >= params.watchdog_bytes || + params.watchdog_bytes >= params.hard_bytes || + params.hard_bytes > LLAMA_DSV41_ADMISSION_HARD_BYTES) { + reject("thresholds", result, "require soft < watchdog < hard <= 120 GiB"); + } + validate_context(params.n_ctx); + if (params.n_seq != 1) { + reject("context", result, "bounded DeepSeek V4.1 admission currently requires one sequence"); + } + if (params.n_expert_used == 0 || params.n_expert_used > LLAMA_DSV41_N_EXPERT) { + reject("cache", result, "expert top-k is invalid"); + } + + std::vector layer_slot_bytes(LLAMA_DSV41_N_LAYER, 0); + for (const auto & tensor : expert_tensors) { + llama_expert_store_validate_tensor(tensor); + if (tensor.layer < 0 || tensor.layer >= (int32_t) LLAMA_DSV41_N_LAYER || + tensor.ne[2] != LLAMA_DSV41_N_EXPERT) { + reject("cache", result, "expert tensor geometry is invalid"); + } + layer_slot_bytes[tensor.layer] = checked_add( + layer_slot_bytes[tensor.layer], tensor.nb[2], "expert slot"); + result.expert_slot_bytes = checked_add( + result.expert_slot_bytes, tensor.nb[2], "expert slot"); + } + if (expert_tensors.size() != LLAMA_DSV41_N_LAYER*3) { + reject("cache", result, "expected 40 gate/up/down expert tensor sets"); + } + for (uint64_t bytes : layer_slot_bytes) { + if (bytes == 0) { + reject("cache", result, "expert layer has no tensor plane bytes"); + } + result.expert_staging_slot_bytes = std::max(result.expert_staging_slot_bytes, bytes); + } + + const uint64_t bytes_slots = params.configured_cache_bytes == 0 ? + LLAMA_DSV41_N_EXPERT : params.configured_cache_bytes/result.expert_slot_bytes; + if (params.configured_cache_slots != 0 && params.configured_cache_bytes != 0 && + params.configured_cache_slots != bytes_slots) { + reject("cache", result, "configured cache slots and bytes disagree"); + } + uint64_t slot_cap = LLAMA_DSV41_N_EXPERT; + if (params.configured_cache_slots != 0) { + slot_cap = std::min(slot_cap, params.configured_cache_slots); + } + if (params.configured_cache_bytes != 0) { + slot_cap = std::min(slot_cap, bytes_slots); + } + if (slot_cap < params.n_expert_used) { + reject("cache", result, "configured cache is smaller than expert top-k"); + } + + const auto state = llama_dsv41_account_memory( + params.n_ctx, + params.n_seq, + params.n_ubatch, + params.kv_element_size, + params.index_element_size, + 0); + result.state_bytes = state.total(); + result.graph_workspace_bytes = llama_dsv41_estimate_graph_workspace(params.n_ctx, params.n_ubatch); + result.engram_staging_bytes = llama_dsv41_engram_staging_bytes(params.n_ubatch); + result.output_bytes = llama_dsv41_output_bytes(params.n_vocab, params.n_ubatch); + + result.fixed_bytes = result.host_used; + result.fixed_bytes = checked_add(result.fixed_bytes, result.dense_tensor_bytes, "fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, result.state_bytes, "fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, result.graph_workspace_bytes, "fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, result.engram_staging_bytes, "fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, result.output_bytes, "fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, result.safety_margin_bytes, "fixed bytes"); + + if (result.host_used >= result.hard_bytes) { + reject("hard", result, "current host use is at or above the strict hard limit"); + } + if (result.fixed_bytes > result.soft_bytes) { + result.projected_bytes = result.fixed_bytes; + reject("fixed", result, "fixed startup categories exceed the soft limit"); + } + + const uint64_t bytes_per_slot = checked_add( + result.expert_slot_bytes, result.expert_staging_slot_bytes, "expert slot and staging"); + const uint64_t fit_slots = (result.soft_bytes - result.fixed_bytes)/bytes_per_slot; + const uint64_t selected = std::min(slot_cap, fit_slots); + if (selected < params.n_expert_used) { + result.projected_bytes = result.fixed_bytes; + reject("cache", result, "remaining budget cannot hold the minimum expert top-k"); + } + + result.expert_slots = static_cast(selected); + result.expert_cache_bytes = checked_mul(result.expert_slot_bytes, selected, "expert cache"); + result.expert_staging_bytes = checked_mul(result.expert_staging_slot_bytes, selected, "expert staging"); + result.projected_bytes = checked_add( + checked_add(result.fixed_bytes, result.expert_cache_bytes, "projected bytes"), + result.expert_staging_bytes, + "projected bytes"); + if (result.projected_bytes > result.soft_bytes) { + reject("soft", result, "projected startup exceeds the soft limit"); + } + if (result.projected_bytes >= result.hard_bytes) { + reject("hard", result, "projected startup is not strictly below the hard limit"); + } + result.category = "accepted"; + return result; +} + +std::string llama_dsv41_admission_result::describe() const { + return format( + "DeepSeek V4.1 memory admission: category=%s, context=%u, sequences=%u, ubatch=%u, current=%llu, fixed=%llu, " + "dense=%llu, state=%llu, workspace=%llu, engram_staging=%llu, expert_slots=%u, " + "expert_cache=%llu, expert_staging=%llu, outputs=%llu, safety_margin=%llu, projected=%llu, " + "soft=%llu, watchdog=%llu, hard=%llu, device_reported_ignored=%llu", + category.c_str(), + n_ctx, + n_seq, + n_ubatch, + (unsigned long long) host_used, + (unsigned long long) fixed_bytes, + (unsigned long long) dense_tensor_bytes, + (unsigned long long) state_bytes, + (unsigned long long) graph_workspace_bytes, + (unsigned long long) engram_staging_bytes, + expert_slots, + (unsigned long long) expert_cache_bytes, + (unsigned long long) expert_staging_bytes, + (unsigned long long) output_bytes, + (unsigned long long) safety_margin_bytes, + (unsigned long long) projected_bytes, + (unsigned long long) soft_bytes, + (unsigned long long) watchdog_bytes, + (unsigned long long) hard_bytes, + (unsigned long long) device_reported_bytes_ignored); +} diff --git a/src/llama-dsv41-admission.h b/src/llama-dsv41-admission.h new file mode 100644 index 000000000000..e41ef0b33a57 --- /dev/null +++ b/src/llama-dsv41-admission.h @@ -0,0 +1,83 @@ +#pragma once + +#include "llama-expert-store.h" + +#include +#include +#include +#include + +static constexpr uint64_t LLAMA_DSV41_ADMISSION_SOFT_BYTES = 116ULL << 30; +static constexpr uint64_t LLAMA_DSV41_WATCHDOG_EMERGENCY_BYTES = 118ULL << 30; +static constexpr uint64_t LLAMA_DSV41_ADMISSION_HARD_BYTES = 120ULL << 30; +static constexpr uint64_t LLAMA_DSV41_ADMISSION_MARGIN_BYTES = 2ULL << 30; +static constexpr uint32_t LLAMA_DSV41_ADMISSION_CONTEXT = 32768; + +struct llama_dsv41_host_memory { + uint64_t total = 0; + uint64_t available = 0; + uint64_t used = 0; + uint64_t swap_entries = 0; + uint64_t swap_bytes = 0; +}; + +struct llama_dsv41_admission_params { + uint64_t soft_bytes = LLAMA_DSV41_ADMISSION_SOFT_BYTES; + uint64_t watchdog_bytes = LLAMA_DSV41_WATCHDOG_EMERGENCY_BYTES; + uint64_t hard_bytes = LLAMA_DSV41_ADMISSION_HARD_BYTES; + uint64_t safety_margin_bytes = LLAMA_DSV41_ADMISSION_MARGIN_BYTES; + uint64_t configured_cache_bytes = 0; + uint64_t device_reported_bytes = 0; + uint32_t configured_cache_slots = 0; + uint32_t n_ctx = LLAMA_DSV41_ADMISSION_CONTEXT; + uint32_t n_seq = 1; + uint32_t n_ubatch = 2048; + uint32_t n_vocab = 0; + uint32_t n_expert_used = 0; + uint32_t kv_element_size = 2; + uint32_t index_element_size = 2; + bool direct_io = true; + bool unified_memory = true; +}; + +struct llama_dsv41_admission_result { + uint64_t host_total = 0; + uint64_t host_available = 0; + uint64_t host_used = 0; + uint64_t dense_tensor_bytes = 0; + uint64_t state_bytes = 0; + uint64_t graph_workspace_bytes = 0; + uint64_t engram_staging_bytes = 0; + uint64_t expert_staging_bytes = 0; + uint64_t expert_cache_bytes = 0; + uint64_t output_bytes = 0; + uint64_t safety_margin_bytes = 0; + uint64_t device_reported_bytes_ignored = 0; + uint64_t fixed_bytes = 0; + uint64_t projected_bytes = 0; + uint64_t soft_bytes = 0; + uint64_t watchdog_bytes = 0; + uint64_t hard_bytes = 0; + uint64_t expert_slot_bytes = 0; + uint64_t expert_staging_slot_bytes = 0; + uint32_t expert_slots = 0; + uint32_t n_ctx = 0; + uint32_t n_seq = 0; + uint32_t n_ubatch = 0; + std::string category; + + std::string describe() const; +}; + +llama_dsv41_host_memory llama_dsv41_read_host_memory(const std::string & procfs_root); + +uint64_t llama_dsv41_expert_payload_bytes(const std::vector & tensors); +uint64_t llama_dsv41_estimate_graph_workspace(uint32_t n_ctx, uint32_t n_ubatch); +uint64_t llama_dsv41_engram_staging_bytes(uint32_t n_ubatch); +uint64_t llama_dsv41_output_bytes(uint32_t n_vocab, uint32_t n_ubatch); + +llama_dsv41_admission_result llama_dsv41_admit( + const llama_dsv41_host_memory & host, + uint64_t dense_tensor_bytes, + const std::vector & expert_tensors, + const llama_dsv41_admission_params & params); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1424d48580c1..54676c6d8eb7 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2805,6 +2805,14 @@ llama_model_params llama_model_default_params() { /*.ple_cache_mb =*/ 256, /*.expert_cache_bytes =*/ 0, /*.expert_cache_slots =*/ 0, + /*.dsv41_memory_soft_bytes =*/ 116ULL << 30, + /*.dsv41_memory_watchdog_bytes =*/ 118ULL << 30, + /*.dsv41_memory_hard_bytes =*/ 120ULL << 30, + /*.dsv41_memory_safety_margin_bytes =*/ 2ULL << 30, + /*.dsv41_admission_context =*/ 32768, + /*.dsv41_admission_sequences =*/ 1, + /*.dsv41_admission_ubatch =*/ 2048, + /*.dsv41_procfs_root =*/ "/proc", /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, diff --git a/src/llama-model.h b/src/llama-model.h index ab23bad7a0c4..6f5a2ca82d50 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -769,6 +769,8 @@ struct llama_model { virtual void release_runtime_work() const {} virtual void acquire_runtime_context() const {} virtual void release_runtime_context() const {} + virtual uint32_t default_context_size() const { return 0; } + virtual void validate_context_params(const llama_cparams &) const {} // model must define these virtual void load_arch_hparams(llama_model_loader & ml) = 0; diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 7e3ac1644dd3..0b962daaf3eb 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -1,9 +1,13 @@ +#include "llama-dsv41-admission.h" #include "llama-dsv41.h" #include "llama-dsv41-engram.h" #include "llama-dsv41-expert.h" +#include "llama-cparams.h" #include "llama-hparams.h" #include "models.h" +#include "ggml-alloc.h" + #include #include #include @@ -20,6 +24,10 @@ struct llama_model_deepseek41::engram_model { std::array extents; }; +struct llama_model_deepseek41::admission_model { + llama_dsv41_admission_result result; +}; + void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { llama_dsv41_config config = {}; std::string raw_config; @@ -169,20 +177,23 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { const std::initializer_list & ne) { return ml.register_external_tensor(name, layer, projection, ne); }); - if (params.expert_cache_bytes == 0 || params.expert_cache_slots <= 0) { - throw std::runtime_error( - "DeepSeek V4.1 requires non-zero expert_cache_bytes and expert_cache_slots before tensor allocation"); + + for (size_t index = 0; index < LLAMA_ENGRAM_LAYERS; ++index) { + const int32_t il = engram->layout.layer_ids[index]; + const std::string table_name = tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il).str(); + const auto * table = ml.get_weight(table_name.c_str()); + if (table == nullptr) { + throw std::runtime_error("DeepSeek V4.1 is missing required Engram tensor " + table_name); + } + llama_dsv41_engram_extent & extent = engram->extents[index]; + extent.fname = ml.fnames.at(table->idx); + extent.offset = table->offs; + extent.rows = engram->layout.rows[index]; + extent.columns = table->tensor->ne[0]; + extent.row_count = table->tensor->ne[1]; + extent.type = table->tensor->type; + llama_dsv41_validate_engram_extent(extent); } - llama_dsv41_expert_runtime_params expert_params; - expert_params.cache_bytes = params.expert_cache_bytes; - expert_params.cache_slots = params.expert_cache_slots; - expert_params.direct_io = true; - expert_params.allow_buffered_io = false; - expert_params.no_alloc = ml.no_alloc; - experts = std::make_shared( - expert_tensors, - expert_params, - [this](int32_t layer) { return select_buft(layer); }); tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); @@ -226,32 +237,15 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, 0); layer.ffn_exp_probs_b_vl = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B_VL, "bias", il), { n_expert }, TENSOR_NOT_REQUIRED); layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, 0); - layer.ffn_gate_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_GATE); - layer.ffn_down_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_DOWN); - layer.ffn_up_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_UP); layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_exp*n_expert_shared, n_embd }, 0); layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); if (hparams.dsv41_engram_layers.test(il)) { const size_t index = il == (int32_t) engram->layout.layer_ids[0] ? 0 : 1; - const std::string table_name = tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il).str(); - const auto * table = ml.get_weight(table_name.c_str()); - if (table == nullptr) { - throw std::runtime_error("DeepSeek V4.1 is missing required Engram tensor " + table_name); - } - llama_dsv41_engram_extent & extent = engram->extents[index]; - extent.fname = ml.fnames.at(table->idx); - extent.offset = table->offs; - extent.rows = engram->layout.rows[index]; - extent.columns = table->tensor->ne[0]; - extent.row_count = table->tensor->ne[1]; - extent.type = table->tensor->type; - llama_dsv41_validate_engram_extent(extent); - create_tensor( tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il), - { LLAMA_ENGRAM_ROW_BYTES, (int64_t) extent.rows }, + { LLAMA_ENGRAM_ROW_BYTES, (int64_t) engram->extents[index].rows }, TENSOR_SKIP); layer.engram_q_norm = create_tensor( tn(LLM_TENSOR_ENGRAM_Q_NORM, "weight", il), @@ -268,6 +262,62 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { } } + uint64_t dense_tensor_bytes = 0; + for (const auto & item : ml.ctx_map) { + const uint64_t bytes = ggml_backend_alloc_ctx_tensors_from_buft_size( + item.second.get(), item.first.buft); + if (bytes > UINT64_MAX - dense_tensor_bytes) { + throw std::runtime_error("DeepSeek V4.1 dense allocated tensor byte count overflow"); + } + dense_tensor_bytes += bytes; + } + + llama_dsv41_admission_params admission_params; + admission_params.soft_bytes = params.dsv41_memory_soft_bytes == 0 ? + LLAMA_DSV41_ADMISSION_SOFT_BYTES : params.dsv41_memory_soft_bytes; + admission_params.watchdog_bytes = params.dsv41_memory_watchdog_bytes == 0 ? + LLAMA_DSV41_WATCHDOG_EMERGENCY_BYTES : params.dsv41_memory_watchdog_bytes; + admission_params.hard_bytes = params.dsv41_memory_hard_bytes == 0 ? + LLAMA_DSV41_ADMISSION_HARD_BYTES : params.dsv41_memory_hard_bytes; + admission_params.safety_margin_bytes = params.dsv41_memory_safety_margin_bytes == 0 ? + LLAMA_DSV41_ADMISSION_MARGIN_BYTES : params.dsv41_memory_safety_margin_bytes; + admission_params.configured_cache_bytes = params.expert_cache_bytes; + admission_params.configured_cache_slots = std::max(params.expert_cache_slots, 0); + admission_params.n_ctx = params.dsv41_admission_context == 0 ? + LLAMA_DSV41_ADMISSION_CONTEXT : params.dsv41_admission_context; + admission_params.n_seq = params.dsv41_admission_sequences == 0 ? 1 : params.dsv41_admission_sequences; + admission_params.n_ubatch = params.dsv41_admission_ubatch == 0 ? 2048 : params.dsv41_admission_ubatch; + admission_params.n_vocab = n_vocab; + admission_params.n_expert_used = n_expert_used; + admission_params.direct_io = true; + admission_params.unified_memory = true; + + const std::string procfs_root = params.dsv41_procfs_root == nullptr ? "/proc" : params.dsv41_procfs_root; + admission = std::make_shared(); + admission->result = llama_dsv41_admit( + llama_dsv41_read_host_memory(procfs_root), + dense_tensor_bytes, + expert_tensors, + admission_params); + LLAMA_LOG_INFO("%s\n", admission->result.describe().c_str()); + + llama_dsv41_expert_runtime_params expert_params; + expert_params.cache_bytes = admission->result.expert_cache_bytes; + expert_params.cache_slots = admission->result.expert_slots; + expert_params.direct_io = true; + expert_params.allow_buffered_io = false; + expert_params.no_alloc = ml.no_alloc; + experts = std::make_shared( + expert_tensors, + expert_params, + [this](int32_t layer) { return select_buft(layer); }); + + for (int32_t il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + layer.ffn_gate_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_GATE); + layer.ffn_down_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_DOWN); + layer.ffn_up_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_UP); + } } bool llama_model_deepseek41::requires_synchronous_graph() const { @@ -296,6 +346,26 @@ void llama_model_deepseek41::release_runtime_context() const { } } +uint32_t llama_model_deepseek41::default_context_size() const { + return admission ? admission->result.n_ctx : LLAMA_DSV41_ADMISSION_CONTEXT; +} + +void llama_model_deepseek41::validate_context_params(const llama_cparams & cparams) const { + if (!admission) { + throw std::runtime_error("DeepSeek V4.1 context has no host-memory admission result"); + } + if (cparams.n_ctx > admission->result.n_ctx || + cparams.n_seq_max > admission->result.n_seq || + cparams.n_ubatch > admission->result.n_ubatch) { + throw std::runtime_error(format( + "%s, category=context, requested_context=%u, requested_sequences=%u, requested_ubatch=%u", + admission->result.describe().c_str(), + cparams.n_ctx, + cparams.n_seq_max, + cparams.n_ubatch)); + } +} + [[noreturn]] std::unique_ptr llama_model_deepseek41::build_arch_graph(const llm_graph_params &) const { throw std::runtime_error(llama_dsv41_runtime_dependency_error()); } diff --git a/src/models/models.h b/src/models/models.h index 00ae080aa672..18171f2a78c3 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1318,7 +1318,9 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { llama_model_deepseek41(const struct llama_model_params & params) : llama_model_deepseek4(params) {} struct engram_model; + struct admission_model; std::shared_ptr engram; + std::shared_ptr admission; std::shared_ptr experts; void load_arch_hparams(llama_model_loader & ml) override; @@ -1328,6 +1330,8 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { void release_runtime_work() const override; void acquire_runtime_context() const override; void release_runtime_context() const override; + uint32_t default_context_size() const override; + void validate_context_params(const llama_cparams & cparams) const override; [[noreturn]] std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 966e16a7c6c7..e0908454bf00 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -196,6 +196,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW + llama_build_and_test(test-deepseek41-admission.cpp) llama_build_and_test(test-deepseek41-schema.cpp) llama_build_and_test(test-deepseek41-engram.cpp) llama_build_and_test(test-deepseek41-expert.cpp) diff --git a/tests/test-deepseek41-admission.cpp b/tests/test-deepseek41-admission.cpp new file mode 100644 index 000000000000..790ca2112d95 --- /dev/null +++ b/tests/test-deepseek41-admission.cpp @@ -0,0 +1,263 @@ +#include "../src/llama-dsv41-admission.h" +#include "../src/llama-dsv41.h" + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(cond) do { if (!(cond)) { throw std::runtime_error("requirement failed: " #cond); } } while (0) + +namespace { + +struct temp_procfs { + std::filesystem::path path; + + temp_procfs() { + static uint64_t sequence = 0; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path = std::filesystem::temp_directory_path() / + ("llama-dsv41-admission-" + std::to_string(stamp) + "-" + std::to_string(++sequence)); + std::filesystem::create_directories(path); + } + + ~temp_procfs() { + std::error_code ec; + std::filesystem::remove_all(path, ec); + } + + void write(const char * name, const std::string & value) { + std::ofstream file(path / name); + REQUIRE((bool) file); + file << value; + REQUIRE((bool) file); + } +}; + +template +std::string thrown(F && fn) { + try { + fn(); + } catch (const std::exception & e) { + return e.what(); + } + throw std::runtime_error("expected exception"); +} + +std::vector published_tensors() { + std::vector result; + uint64_t offset = 4096; + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + for (llama_expert_projection projection : { + LLAMA_EXPERT_PROJECTION_GATE, + LLAMA_EXPERT_PROJECTION_UP, + LLAMA_EXPERT_PROJECTION_DOWN }) { + llama_expert_store_tensor tensor; + tensor.name = "blk." + std::to_string(il) + ".expert." + std::to_string((int) projection); + tensor.fname = "published.gguf"; + tensor.layer = il; + tensor.projection = projection; + tensor.type = projection == LLAMA_EXPERT_PROJECTION_DOWN ? GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS; + tensor.ne[0] = projection == LLAMA_EXPERT_PROJECTION_DOWN ? 2304 : 5120; + tensor.ne[1] = projection == LLAMA_EXPERT_PROJECTION_DOWN ? 5120 : 2304; + tensor.ne[2] = LLAMA_DSV41_N_EXPERT; + tensor.nb[0] = ggml_type_size(tensor.type); + tensor.nb[1] = ggml_row_size(tensor.type, tensor.ne[0]); + tensor.nb[2] = tensor.nb[1]*tensor.ne[1]; + tensor.file_offset = offset; + tensor.file_size = offset + tensor.nb[2]*tensor.ne[2]; + offset = tensor.file_size; + result.push_back(std::move(tensor)); + } + } + return result; +} + +llama_dsv41_host_memory host_with_used(uint64_t used) { + llama_dsv41_host_memory host; + host.total = 128ULL << 30; + host.available = host.total - used; + host.used = used; + return host; +} + +llama_dsv41_admission_params base_params() { + llama_dsv41_admission_params params; + params.n_vocab = LLAMA_DSV41_N_VOCAB; + params.n_expert_used = LLAMA_DSV41_N_EXPERT_USED; + return params; +} + +void test_procfs() { + temp_procfs procfs; + procfs.write("meminfo", + "MemTotal: 131072000 kB\n" + "MemFree: 100000 kB\n" + "MemAvailable: 120000000 kB\n"); + procfs.write("swaps", "Filename Type Size Used Priority\n"); + const auto memory = llama_dsv41_read_host_memory(procfs.path.string()); + REQUIRE(memory.total == 131072000ULL*1024); + REQUIRE(memory.available == 120000000ULL*1024); + REQUIRE(memory.used == 11072000ULL*1024); + REQUIRE(memory.swap_entries == 0); + + procfs.write("swaps", + "Filename Type Size Used Priority\n" + "/swapfile file 33554428 0 -2\n"); + const auto swapped = llama_dsv41_read_host_memory(procfs.path.string()); + REQUIRE(swapped.swap_entries == 1); + REQUIRE(swapped.swap_bytes == 33554428ULL*1024); + REQUIRE(thrown([&]() { + llama_dsv41_admit(swapped, 0, published_tensors(), base_params()); + }).find("category=swap") != std::string::npos); +} + +void test_procfs_fail_closed() { + temp_procfs procfs; + procfs.write("meminfo", "MemTotal: 10 kB\n"); + procfs.write("swaps", "Filename Type Size Used Priority\n"); + REQUIRE(!thrown([&]() { + llama_dsv41_read_host_memory(procfs.path.string()); + }).empty()); + + procfs.write("meminfo", + "MemTotal: 18446744073709551615 kB\n" + "MemAvailable: 1 kB\n"); + REQUIRE(thrown([&]() { + llama_dsv41_read_host_memory(procfs.path.string()); + }).find("overflow") != std::string::npos); + + procfs.write("meminfo", "MemTotal: 10 bytes\nMemAvailable: 1 kB\n"); + REQUIRE(!thrown([&]() { + llama_dsv41_read_host_memory(procfs.path.string()); + }).empty()); +} + +void test_published_slot_fit() { + const auto tensors = published_tensors(); + auto params = base_params(); + params.device_reported_bytes = 64ULL << 30; + const auto result = llama_dsv41_admit( + host_with_used(8ULL << 30), + 9376ULL << 20, + tensors, + params); + REQUIRE(result.expert_slot_bytes == 398131200); + REQUIRE(result.expert_staging_slot_bytes == 9953280); + REQUIRE(result.expert_slots >= LLAMA_DSV41_N_EXPERT_USED); + REQUIRE(result.expert_slots < LLAMA_DSV41_N_EXPERT); + REQUIRE(result.expert_cache_bytes == result.expert_slot_bytes*result.expert_slots); + REQUIRE(result.projected_bytes <= result.soft_bytes); + REQUIRE(result.device_reported_bytes_ignored == 64ULL << 30); +} + +void test_configured_cache() { + const auto tensors = published_tensors(); + auto params = base_params(); + params.configured_cache_slots = 12; + params.configured_cache_bytes = 12*398131200ULL + 1024; + const auto result = llama_dsv41_admit(host_with_used(0), 0, tensors, params); + REQUIRE(result.expert_slots == 12); + REQUIRE(result.expert_cache_bytes == 12*398131200ULL); + + params.configured_cache_slots = 13; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("disagree") != std::string::npos); + + params.configured_cache_slots = 5; + params.configured_cache_bytes = 0; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("top-k") != std::string::npos); +} + +void test_threshold_boundaries() { + const auto tensors = published_tensors(); + auto params = base_params(); + params.configured_cache_slots = LLAMA_DSV41_N_EXPERT_USED; + params.configured_cache_bytes = params.configured_cache_slots*398131200ULL; + params.safety_margin_bytes = 0; + + const auto zero = llama_dsv41_admit(host_with_used(0), 0, tensors, params); + const uint64_t planned_without_host = zero.projected_bytes; + const auto exact_soft = llama_dsv41_admit( + host_with_used(params.soft_bytes - planned_without_host), 0, tensors, params); + REQUIRE(exact_soft.projected_bytes == params.soft_bytes); + + const std::string watchdog_error = thrown([&]() { + llama_dsv41_admit( + host_with_used(params.watchdog_bytes - planned_without_host), 0, tensors, params); + }); + REQUIRE(watchdog_error.find("category=cache") != std::string::npos); + REQUIRE(watchdog_error.find("watchdog=126701535232") != std::string::npos); + + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(params.hard_bytes), 0, tensors, params); + }).find("category=hard") != std::string::npos); +} + +void test_context_progression() { + const auto tensors = published_tensors(); + uint64_t previous_state = 0; + uint64_t previous_workspace = 0; + for (uint32_t n_ctx : { 32768U, 65536U, 98304U, 131072U }) { + auto params = base_params(); + params.n_ctx = n_ctx; + const auto result = llama_dsv41_admit(host_with_used(0), 0, tensors, params); + REQUIRE(result.n_ctx == n_ctx); + REQUIRE(result.state_bytes > previous_state); + REQUIRE(result.graph_workspace_bytes > previous_workspace); + previous_state = result.state_bytes; + previous_workspace = result.graph_workspace_bytes; + } + + auto params = base_params(); + params.n_ctx = 49152; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("32768, 65536, 98304, or 131072") != std::string::npos); +} + +void test_diagnostics_and_guards() { + const auto tensors = published_tensors(); + auto params = base_params(); + params.direct_io = false; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("category=direct_io") != std::string::npos); + + params.direct_io = true; + params.unified_memory = false; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("category=unified_memory") != std::string::npos); + + params.unified_memory = true; + const auto result = llama_dsv41_admit(host_with_used(0), 0, tensors, params); + const std::string diagnostic = result.describe(); + for (const char * field : { + "category=", "current=", "fixed=", "dense=", "state=", "workspace=", + "expert_slots=", "expert_cache=", "expert_staging=", "soft=", "watchdog=", "hard=" }) { + REQUIRE(diagnostic.find(field) != std::string::npos); + } +} + +} + +int main() { + test_procfs(); + test_procfs_fail_closed(); + test_published_slot_fit(); + test_configured_cache(); + test_threshold_boundaries(); + test_context_progression(); + test_diagnostics_and_guards(); + return 0; +} From 31d5958d7b8deadaf0ec3c85a0015b2d509ad140 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:12:50 -0700 Subject: [PATCH 06/32] scripts : add Strix host-memory watchdog Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- README.md | 1 + docs/strix-memory-watchdog.md | 40 ++ scripts/strix_memory_watchdog.py | 664 ++++++++++++++++++++++++++++ tests/CMakeLists.txt | 11 + tests/test_strix_memory_watchdog.py | 406 +++++++++++++++++ 5 files changed, 1122 insertions(+) create mode 100644 docs/strix-memory-watchdog.md create mode 100755 scripts/strix_memory_watchdog.py create mode 100644 tests/test_strix_memory_watchdog.py diff --git a/README.md b/README.md index 3067a7a4e71e..dbf547ce5c5b 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,7 @@ Everything else is upstream `llama.cpp`. The additions currently carried here: | Speculative checkpoints on device | | `llama-server` keeps speculative-decoding checkpoints in device memory instead of copying them to the host | | ROCmFPx quant types | `llama-quantize` types `Q4_0_ROCMFP4`, `Q4_0_ROCMFP4_FAST`, `Q2/Q3/Q6/Q8_0_ROCMFPX` and the `_LEAN`/`_COHERENT`/`_STRIX` recipes | Loads the ROCmFP4 GGUFs published for Strix Halo. CPU codecs plus Vulkan dequant, mat-vec, matmul and integer-dot kernels. Weight formats only: not accepted as KV-cache types | | Repeatable output at depth | | Freed KV cells are zeroed so masked-out rows never carry stale K/V, and the Vulkan radix top-k assigns output slots deterministically | +| Host-memory watchdog | [`scripts/strix_memory_watchdog.py`](docs/strix-memory-watchdog.md) | Runs a command in a process group, requires zero active swap, and stops before host-wide memory reaches the 120 GiB validation ceiling | Every ROCm/HIP change above is guarded on architecture, shape and layout, so other devices see upstream behaviour. Run `--help`, or see [tools/server/README.md](tools/server/README.md), for the full options. diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md new file mode 100644 index 000000000000..4adc0e4be7fe --- /dev/null +++ b/docs/strix-memory-watchdog.md @@ -0,0 +1,40 @@ +# Strix host-memory watchdog + +`scripts/strix_memory_watchdog.py` is an external Linux command wrapper for headless Strix Halo validation. It does not change model loading or cache sizing. It measures host-wide memory from procfs and controls the launched command's process group. + +```sh +./scripts/strix_memory_watchdog.py -- ./build/bin/llama-server +``` + +The wrapper performs these checks and actions: + +- It refuses to launch if `/proc/swaps` contains any active entry. +- It calculates used memory as `MemTotal - MemAvailable`. Linux reports these fields in KiB, so the wrapper multiplies each value by 1024 and keeps all accounting as integer bytes. +- It sends `SIGTERM` to the process group at 116 GiB used. +- It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. +- It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. +- It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. + +The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. + +Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. + +The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Child standard input, standard output, and standard error are inherited unchanged. + +Exit classifications are authoritative in the final JSON record. Operational failures use these exit codes: + +| Exit code | Classification | +| ---: | --- | +| 2 | procfs or configuration error | +| 3 | swap active at startup or detected during execution | +| 4 | soft threshold reached | +| 5 | emergency threshold reached | +| 6 | soft-threshold grace period expired | +| 7 | process-group signaling or termination failure | +| 127 | command launch failure | + +No model, backend, or ROCm package is required to run the unit tests: + +```sh +python3 tests/test_strix_memory_watchdog.py +``` diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py new file mode 100755 index 000000000000..0dc532c1d027 --- /dev/null +++ b/scripts/strix_memory_watchdog.py @@ -0,0 +1,664 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import signal +import subprocess +import sys +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import IO, Protocol + + +GIB = 1024**3 +STRICT_CEILING_BYTES = 120 * GIB +DEFAULT_SOFT_BYTES = 116 * GIB +DEFAULT_EMERGENCY_BYTES = 118 * GIB +DEFAULT_GRACE_SECONDS = 30.0 +DEFAULT_SAMPLE_INTERVAL_SECONDS = 1.0 + +EXIT_PROCFS_ERROR = 2 +EXIT_SWAP_ACTIVE = 3 +EXIT_SOFT_LIMIT = 4 +EXIT_EMERGENCY_LIMIT = 5 +EXIT_GRACE_TIMEOUT = 6 +EXIT_SIGNAL_ERROR = 7 +EXIT_LAUNCH_ERROR = 127 + +MEMINFO_VALUE_RE = re.compile(r"([0-9]+) kB") +SWAPS_HEADER = ["Filename", "Type", "Size", "Used", "Priority"] + + +class ProcfsError(RuntimeError): + pass + + +class ProcessGroupError(RuntimeError): + pass + + +class ProcessHandle(Protocol): + pid: int + + def poll(self) -> int | None: + ... + + def wait(self, timeout: float | None = None) -> int: + ... + + +@dataclass(frozen=True) +class HostSnapshot: + total_bytes: int + available_bytes: int + active_swaps: tuple[str, ...] + + @property + def used_bytes(self) -> int: + return self.total_bytes - self.available_bytes + + +@dataclass(frozen=True) +class WatchdogConfig: + command: tuple[str, ...] + procfs_root: Path = Path("/proc") + soft_bytes: int = DEFAULT_SOFT_BYTES + emergency_bytes: int = DEFAULT_EMERGENCY_BYTES + grace_seconds: float = DEFAULT_GRACE_SECONDS + sample_interval_seconds: float = DEFAULT_SAMPLE_INTERVAL_SECONDS + + def validate(self) -> None: + if not self.command: + raise ValueError("a command is required after --") + if self.soft_bytes <= 0: + raise ValueError("soft threshold must be greater than zero") + if self.emergency_bytes <= self.soft_bytes: + raise ValueError("emergency threshold must be greater than soft threshold") + if self.emergency_bytes >= STRICT_CEILING_BYTES: + raise ValueError("emergency threshold must be below 120 GiB") + if not math.isfinite(self.grace_seconds) or self.grace_seconds <= 0: + raise ValueError("grace period must be greater than zero") + if ( + not math.isfinite(self.sample_interval_seconds) + or self.sample_interval_seconds <= 0 + ): + raise ValueError("sample interval must be greater than zero") + + +class ProcfsReader: + def __init__(self, root: Path): + self.root = root + + def _read_text(self, name: str) -> str: + path = self.root / name + try: + return path.read_text(encoding="utf-8") + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcfsError(f"cannot read {path}: {detail}") from exc + + def read_snapshot(self) -> HostSnapshot: + active_swaps = self._parse_swaps(self._read_text("swaps")) + total_bytes, available_bytes = self._parse_meminfo( + self._read_text("meminfo") + ) + return HostSnapshot(total_bytes, available_bytes, active_swaps) + + @staticmethod + def _parse_meminfo(content: str) -> tuple[int, int]: + values: dict[str, int] = {} + required = {"MemTotal", "MemAvailable"} + for line in content.splitlines(): + key, separator, raw_value = line.partition(":") + if not separator or key not in required: + continue + if key in values: + raise ProcfsError(f"duplicate {key} in meminfo") + match = MEMINFO_VALUE_RE.fullmatch(raw_value.strip()) + if match is None: + raise ProcfsError(f"malformed {key} in meminfo") + values[key] = int(match.group(1)) * 1024 + + missing = sorted(required - values.keys()) + if missing: + raise ProcfsError(f"missing {', '.join(missing)} in meminfo") + if values["MemAvailable"] > values["MemTotal"]: + raise ProcfsError("MemAvailable exceeds MemTotal") + return values["MemTotal"], values["MemAvailable"] + + @staticmethod + def _parse_swaps(content: str) -> tuple[str, ...]: + lines = content.splitlines() + if not lines or lines[0].split() != SWAPS_HEADER: + raise ProcfsError("malformed swaps header") + + entries: list[str] = [] + for line in lines[1:]: + if not line.strip(): + continue + fields = line.split() + if len(fields) != len(SWAPS_HEADER): + raise ProcfsError("malformed swaps entry") + try: + int(fields[2]) + int(fields[3]) + int(fields[4]) + except ValueError as exc: + raise ProcfsError("malformed swaps entry") from exc + entries.append(fields[0]) + return tuple(entries) + + +class AuditLogger: + def __init__( + self, + stream: IO[str], + wall_clock: Callable[[], datetime] | None = None, + ): + self.stream = stream + self.wall_clock = wall_clock or ( + lambda: datetime.now(timezone.utc) + ) + + def emit(self, event: str, **fields: object) -> None: + timestamp = self.wall_clock().astimezone(timezone.utc) + record = { + "timestamp": timestamp.isoformat(timespec="milliseconds").replace( + "+00:00", "Z" + ), + "event": event, + **fields, + } + self.stream.write( + json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" + ) + self.stream.flush() + + +def _child_status(returncode: int | None, started: bool = True) -> str: + if not started: + return "not_started" + if returncode is None: + return "running" + return "signaled" if returncode < 0 else "exited" + + +def _state_fields( + snapshot: HostSnapshot | None, + peak_used_bytes: int | None, + child: ProcessHandle | None, + child_returncode: int | None, + process_group_status: str, + threshold_reason: str, +) -> dict[str, object]: + return { + "total_bytes": snapshot.total_bytes if snapshot else None, + "available_bytes": snapshot.available_bytes if snapshot else None, + "used_bytes": snapshot.used_bytes if snapshot else None, + "swap_entries": len(snapshot.active_swaps) if snapshot else None, + "peak_used_bytes": peak_used_bytes, + "child_pid": child.pid if child else None, + "child_status": _child_status( + child_returncode, started=child is not None + ), + "child_returncode": child_returncode, + "process_group_id": child.pid if child else None, + "process_group_status": process_group_status, + "threshold_reason": threshold_reason, + } + + +def _emit_final( + audit: AuditLogger, + classification: str, + exit_code: int, + reason: str, + snapshot: HostSnapshot | None, + peak_used_bytes: int | None, + child: ProcessHandle | None = None, + child_returncode: int | None = None, + process_group_status: str = "not_created", + error: str | None = None, +) -> int: + fields = _state_fields( + snapshot, + peak_used_bytes, + child, + child_returncode, + process_group_status, + reason, + ) + fields.update(classification=classification, exit_code=exit_code) + if error: + fields["error"] = error + audit.emit("final", **fields) + return exit_code + + +def _signal_process_group(process_group_id: int, signal_number: int) -> str: + try: + os.killpg(process_group_id, signal_number) + except ProcessLookupError: + return "missing" + except OSError as exc: + name = signal.Signals(signal_number).name + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot send {name} to process group {process_group_id}: {detail}" + ) from exc + return f"{signal.Signals(signal_number).name.lower()}_sent" + + +def _kill_and_finish( + audit: AuditLogger, + child: ProcessHandle, + snapshot: HostSnapshot, + peak_used_bytes: int, + classification: str, + exit_code: int, + reason: str, + signal_group: Callable[[int, int], str], +) -> int: + try: + group_status = signal_group(child.pid, signal.SIGKILL) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "signal_error", + str(exc), + ) + + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + group_status, + reason, + ), + signal="SIGKILL", + ) + try: + child_returncode = child.wait(timeout=5.0) + except subprocess.TimeoutExpired as exc: + return _emit_final( + audit, + "termination_timeout", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "sigkill_timeout", + str(exc), + ) + return _emit_final( + audit, + classification, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + group_status, + ) + + +def _monitor_child( + config: WatchdogConfig, + reader: ProcfsReader, + audit: AuditLogger, + child: ProcessHandle, + initial_snapshot: HostSnapshot, + signal_group: Callable[[int, int], str], + monotonic: Callable[[], float], + sleeper: Callable[[float], None], +) -> int: + snapshot = initial_snapshot + peak = snapshot.used_bytes + soft_deadline: float | None = None + + while True: + child_returncode = child.poll() + if child_returncode is not None: + soft_stop = soft_deadline is not None + return _emit_final( + audit, + "soft_limit" if soft_stop else "child_exit", + EXIT_SOFT_LIMIT if soft_stop else ( + 128 - child_returncode + if child_returncode < 0 + else child_returncode + ), + ( + "child exited during soft-threshold grace period" + if soft_stop + else "child exited" + ), + snapshot, + peak, + child, + child_returncode, + "leader_exited", + ) + + now = monotonic() + if soft_deadline is not None and now >= soft_deadline: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "grace_timeout", + EXIT_GRACE_TIMEOUT, + "soft-threshold grace period expired", + signal_group, + ) + + try: + snapshot = reader.read_snapshot() + except ProcfsError as exc: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + signal_group, + ) + + peak = max(peak, snapshot.used_bytes) + audit.emit( + "sample", + **_state_fields( + snapshot, peak, child, None, "active", "none" + ), + ) + + if snapshot.active_swaps: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "swap_appeared", + EXIT_SWAP_ACTIVE, + "active swap appeared during execution", + signal_group, + ) + if snapshot.used_bytes >= config.emergency_bytes: + return _kill_and_finish( + audit, + child, + snapshot, + peak, + "emergency_limit", + EXIT_EMERGENCY_LIMIT, + "used_bytes >= emergency_bytes", + signal_group, + ) + if soft_deadline is None and snapshot.used_bytes >= config.soft_bytes: + try: + group_status = signal_group(child.pid, signal.SIGTERM) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + "used_bytes >= soft_bytes", + snapshot, + peak, + child, + child.poll(), + "signal_error", + str(exc), + ) + soft_deadline = now + config.grace_seconds + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak, + child, + child.poll(), + group_status, + "used_bytes >= soft_bytes", + ), + signal="SIGTERM", + grace_deadline_monotonic=soft_deadline, + ) + + sleep_seconds = config.sample_interval_seconds + if soft_deadline is not None: + sleep_seconds = min( + sleep_seconds, + max(0.0, soft_deadline - monotonic()), + ) + sleeper(sleep_seconds) + + +def run_watchdog( + config: WatchdogConfig, + *, + reader: ProcfsReader | None = None, + audit: AuditLogger | None = None, + launcher: Callable[..., ProcessHandle] | None = None, + signal_group: Callable[[int, int], str] | None = None, + monotonic: Callable[[], float] | None = None, + sleeper: Callable[[float], None] | None = None, +) -> int: + config.validate() + reader = reader or ProcfsReader(config.procfs_root) + audit = audit or AuditLogger(sys.stderr) + launcher = launcher or subprocess.Popen + signal_group = signal_group or _signal_process_group + monotonic = monotonic or time.monotonic + sleeper = sleeper or time.sleep + + try: + snapshot = reader.read_snapshot() + except ProcfsError as exc: + return _emit_final( + audit, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + None, + None, + error=str(exc), + ) + + audit.emit( + "preflight", + **_state_fields( + snapshot, + snapshot.used_bytes, + None, + None, + "not_created", + "none", + ), + soft_bytes=config.soft_bytes, + emergency_bytes=config.emergency_bytes, + strict_ceiling_bytes=STRICT_CEILING_BYTES, + ) + + if snapshot.active_swaps: + return _emit_final( + audit, + "startup_swap_active", + EXIT_SWAP_ACTIVE, + "active swap present before command launch", + snapshot, + snapshot.used_bytes, + ) + if snapshot.used_bytes >= config.emergency_bytes: + return _emit_final( + audit, + "startup_emergency_limit", + EXIT_EMERGENCY_LIMIT, + "used_bytes >= emergency_bytes before launch", + snapshot, + snapshot.used_bytes, + ) + if snapshot.used_bytes >= config.soft_bytes: + return _emit_final( + audit, + "startup_soft_limit", + EXIT_SOFT_LIMIT, + "used_bytes >= soft_bytes before launch", + snapshot, + snapshot.used_bytes, + ) + + try: + child = launcher(config.command, start_new_session=True) + except (OSError, ValueError) as exc: + detail = getattr(exc, "strerror", None) or str(exc) + return _emit_final( + audit, + "launch_error", + EXIT_LAUNCH_ERROR, + "command launch failed", + snapshot, + snapshot.used_bytes, + error=detail, + ) + + audit.emit( + "child_started", + **_state_fields( + snapshot, + snapshot.used_bytes, + child, + None, + "active", + "none", + ), + command=list(config.command), + ) + return _monitor_child( + config, + reader, + audit, + child, + snapshot, + signal_group, + monotonic, + sleeper, + ) + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be greater than zero") + return parsed + + +def _positive_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError("value must be greater than zero") + return parsed + + +def parse_args(argv: Sequence[str]) -> WatchdogConfig: + parser = argparse.ArgumentParser( + description=( + "Launch a command in a new process group and stop it before " + "host-wide memory use reaches the 120 GiB Strix validation ceiling." + ) + ) + parser.add_argument( + "--procfs-root", + type=Path, + default=Path("/proc"), + help="procfs root containing meminfo and swaps (default: /proc)", + ) + parser.add_argument( + "--soft-gib", + type=_positive_int, + default=116, + help="send SIGTERM at this many GiB used (default: 116)", + ) + parser.add_argument( + "--emergency-gib", + type=_positive_int, + default=118, + help=( + "send SIGKILL at this many GiB used (default: 118, leaving " + "a 2 GiB sampling margin below 120 GiB)" + ), + ) + parser.add_argument( + "--grace-seconds", + type=_positive_float, + default=DEFAULT_GRACE_SECONDS, + help="maximum time after SIGTERM before SIGKILL (default: 30)", + ) + parser.add_argument( + "--sample-interval-seconds", + type=_positive_float, + default=DEFAULT_SAMPLE_INTERVAL_SECONDS, + help="procfs sampling interval (default: 1)", + ) + parser.add_argument( + "command", + nargs=argparse.REMAINDER, + help="command and arguments, preceded by --", + ) + args = parser.parse_args(argv) + command = tuple(args.command) + if command and command[0] == "--": + command = command[1:] + return WatchdogConfig( + command=command, + procfs_root=args.procfs_root, + soft_bytes=args.soft_gib * GIB, + emergency_bytes=args.emergency_gib * GIB, + grace_seconds=args.grace_seconds, + sample_interval_seconds=args.sample_interval_seconds, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + config = parse_args(argv if argv is not None else sys.argv[1:]) + audit = AuditLogger(sys.stderr) + try: + return run_watchdog(config, audit=audit) + except ValueError as exc: + return _emit_final( + audit, + "configuration_error", + EXIT_PROCFS_ERROR, + "invalid configuration", + None, + None, + error=str(exc), + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e0908454bf00..7854663a8173 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -262,6 +262,17 @@ llama_build_and_test(test-chat-template.cpp) # debug tool for chat template differential analysis (not registered as a test, run it manually) llama_build(test-chat-analysis.cpp) llama_build_and_test(test-log.cpp) + +find_package(Python3 3.10 COMPONENTS Interpreter QUIET) +if (CMAKE_SYSTEM_NAME STREQUAL "Linux" AND Python3_Interpreter_FOUND) + llama_test_cmd( + ${Python3_EXECUTABLE} + NAME test-strix-memory-watchdog + LABEL python + ARGS ${CMAKE_CURRENT_SOURCE_DIR}/test_strix_memory_watchdog.py + ) +endif() + llama_build_and_test( test-peg-parser.cpp peg-parser/simple-tokenize.cpp diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py new file mode 100644 index 000000000000..fa4c4aaec969 --- /dev/null +++ b/tests/test_strix_memory_watchdog.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib.util +import io +import json +import signal +import subprocess +import sys +import tempfile +import unittest +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +SCRIPT_PATH = ( + Path(__file__).resolve().parents[1] + / "scripts" + / "strix_memory_watchdog.py" +) +SPEC = importlib.util.spec_from_file_location( + "strix_memory_watchdog", SCRIPT_PATH +) +assert SPEC is not None +assert SPEC.loader is not None +watchdog = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = watchdog +SPEC.loader.exec_module(watchdog) + + +def snapshot( + used_bytes: int, + *, + total_bytes: int = 200, + active_swaps: tuple[str, ...] = (), +) -> Any: + return watchdog.HostSnapshot( + total_bytes=total_bytes, + available_bytes=total_bytes - used_bytes, + active_swaps=active_swaps, + ) + + +class SequenceReader: + def __init__(self, values: list[Any]): + self.values = values + self.index = 0 + + def read_snapshot(self) -> Any: + index = min(self.index, len(self.values) - 1) + self.index += 1 + value = self.values[index] + if isinstance(value, Exception): + raise value + return value + + +class FakeClock: + def __init__(self): + self.value = 0.0 + + def monotonic(self) -> float: + return self.value + + def sleep(self, seconds: float) -> None: + self.value += seconds + + +class FakeProcess: + def __init__(self, returncode: int | None = None): + self.pid = 4321 + self.returncode = returncode + + def poll(self) -> int | None: + return self.returncode + + def wait(self, timeout: float | None = None) -> int: + if self.returncode is None: + raise subprocess.TimeoutExpired("fake", timeout) + return self.returncode + + +class Harness: + def __init__( + self, + values: list[Any], + process: FakeProcess, + signal_handler: Any | None = None, + ): + self.reader = SequenceReader(values) + self.process = process + self.signal_handler = signal_handler + self.clock = FakeClock() + self.stream = io.StringIO() + self.launched = False + self.signals: list[int] = [] + fixed_time = datetime(2026, 1, 1, tzinfo=timezone.utc) + self.audit = watchdog.AuditLogger( + self.stream, wall_clock=lambda: fixed_time + ) + + def launcher(self, command: tuple[str, ...], **kwargs: Any) -> FakeProcess: + self.launched = True + self.command = command + self.launch_kwargs = kwargs + return self.process + + def signal_group(self, process_group_id: int, signal_number: int) -> str: + self.signals.append(signal_number) + if self.signal_handler is not None: + self.signal_handler(self.process, signal_number) + return f"{signal.Signals(signal_number).name.lower()}_sent" + + def run(self, **overrides: Any) -> int: + config = watchdog.WatchdogConfig( + command=("fake-command",), + soft_bytes=100, + emergency_bytes=150, + grace_seconds=2, + sample_interval_seconds=1, + **overrides, + ) + return watchdog.run_watchdog( + config, + reader=self.reader, + audit=self.audit, + launcher=self.launcher, + signal_group=self.signal_group, + monotonic=self.clock.monotonic, + sleeper=self.clock.sleep, + ) + + def records(self) -> list[dict[str, Any]]: + return [ + json.loads(line) + for line in self.stream.getvalue().splitlines() + ] + + +class TestProcfsParsing(unittest.TestCase): + def test_parses_meminfo_as_integer_bytes_and_allows_zero_swap(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + (root / "meminfo").write_text( + "MemTotal: 131072 kB\n" + "MemFree: 4096 kB\n" + "MemAvailable: 32768 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + + result = watchdog.ProcfsReader(root).read_snapshot() + + self.assertEqual(result.total_bytes, 131072 * 1024) + self.assertEqual(result.available_bytes, 32768 * 1024) + self.assertEqual(result.used_bytes, 98304 * 1024) + self.assertEqual(result.active_swaps, ()) + + def test_rejects_active_swap_entry(self) -> None: + content = ( + "Filename Type Size Used Priority\n" + "/swapfile file 1048572 0 -2\n" + ) + self.assertEqual( + watchdog.ProcfsReader._parse_swaps(content), + ("/swapfile",), + ) + + def test_rejects_malformed_or_missing_procfs_data(self) -> None: + with self.assertRaisesRegex( + watchdog.ProcfsError, "malformed MemAvailable" + ): + watchdog.ProcfsReader._parse_meminfo( + "MemTotal: 10 kB\nMemAvailable: unknown\n" + ) + with self.assertRaisesRegex( + watchdog.ProcfsError, "missing MemAvailable" + ): + watchdog.ProcfsReader._parse_meminfo("MemTotal: 10 kB\n") + with self.assertRaisesRegex( + watchdog.ProcfsError, "malformed swaps header" + ): + watchdog.ProcfsReader._parse_swaps("") + with tempfile.TemporaryDirectory() as temp_dir: + with self.assertRaisesRegex( + watchdog.ProcfsError, "cannot read" + ): + watchdog.ProcfsReader( + Path(temp_dir) + ).read_snapshot() + + +class TestWatchdogBehavior(unittest.TestCase): + def test_configuration_rejects_non_finite_timing(self) -> None: + config = watchdog.WatchdogConfig( + command=("fake-command",), + grace_seconds=float("nan"), + ) + with self.assertRaisesRegex(ValueError, "grace period"): + config.validate() + + def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + (root / "meminfo").write_text( + "MemTotal: 131072 kB\nMemAvailable: 65536 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--sample-interval-seconds", + "0.01", + "--", + sys.executable, + "-c", + "raise SystemExit(23)", + ], + capture_output=True, + check=False, + text=True, + ) + + self.assertEqual(result.returncode, 23) + records = [ + json.loads(line) for line in result.stderr.splitlines() + ] + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual(records[-1]["child_returncode"], 23) + + def test_zero_swap_gate_launches_and_propagates_child_exit(self) -> None: + harness = Harness([snapshot(50)], FakeProcess(returncode=37)) + + result = harness.run() + + self.assertEqual(result, 37) + self.assertTrue(harness.launched) + self.assertTrue(harness.launch_kwargs["start_new_session"]) + final = harness.records()[-1] + self.assertEqual(final["classification"], "child_exit") + self.assertEqual(final["total_bytes"], 200) + self.assertEqual(final["available_bytes"], 150) + self.assertEqual(final["used_bytes"], 50) + self.assertEqual(final["peak_used_bytes"], 50) + self.assertEqual(final["child_status"], "exited") + self.assertEqual(final["process_group_status"], "leader_exited") + + def test_signaled_child_exit_uses_shell_exit_convention(self) -> None: + harness = Harness( + [snapshot(50)], + FakeProcess(returncode=-signal.SIGTERM), + ) + + result = harness.run() + + self.assertEqual(result, 128 + signal.SIGTERM) + + def test_active_swap_rejects_startup_without_launch(self) -> None: + harness = Harness( + [snapshot(50, active_swaps=("/swapfile",))], + FakeProcess(), + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SWAP_ACTIVE) + self.assertFalse(harness.launched) + self.assertEqual( + harness.records()[-1]["classification"], + "startup_swap_active", + ) + + def test_soft_limit_sends_sigterm(self) -> None: + def exit_on_term(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGTERM: + process.returncode = -signal.SIGTERM + + harness = Harness( + [snapshot(50), snapshot(110)], + FakeProcess(), + signal_handler=exit_on_term, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SOFT_LIMIT) + self.assertEqual(harness.signals, [signal.SIGTERM]) + self.assertEqual( + harness.records()[-1]["classification"], "soft_limit" + ) + + def test_emergency_limit_sends_sigkill(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(160)], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_EMERGENCY_LIMIT) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "emergency_limit" + ) + + def test_grace_timeout_escalates_to_sigkill(self) -> None: + def ignore_term(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(110)], + FakeProcess(), + signal_handler=ignore_term, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_GRACE_TIMEOUT) + self.assertEqual( + harness.signals, + [signal.SIGTERM, signal.SIGKILL], + ) + self.assertEqual(harness.clock.value, 2.0) + self.assertEqual( + harness.records()[-1]["classification"], "grace_timeout" + ) + + def test_swap_appearing_during_execution_kills_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [ + snapshot(50), + snapshot(60, active_swaps=("/swapfile",)), + ], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SWAP_ACTIVE) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "swap_appeared" + ) + + def test_runtime_procfs_error_kills_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), watchdog.ProcfsError("missing meminfo")], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_PROCFS_ERROR) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "procfs_error" + ) + + def test_launch_failure_is_explicit(self) -> None: + harness = Harness([snapshot(50)], FakeProcess()) + + def fail_launch( + command: tuple[str, ...], **kwargs: Any + ) -> FakeProcess: + raise FileNotFoundError(2, "No such file or directory") + + harness.launcher = fail_launch + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_LAUNCH_ERROR) + self.assertEqual( + harness.records()[-1]["classification"], "launch_error" + ) + + +if __name__ == "__main__": + unittest.main() From 163971607797dd1088b015de61765af2b3068cfa Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:20:14 -0700 Subject: [PATCH 07/32] scripts : clean up watchdog process group on signals Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 3 + scripts/strix_memory_watchdog.py | 309 +++++++++++++++++++++++----- tests/test_strix_memory_watchdog.py | 150 ++++++++++++++ 3 files changed, 407 insertions(+), 55 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 4adc0e4be7fe..8d6d264b141b 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -13,6 +13,8 @@ The wrapper performs these checks and actions: - It sends `SIGTERM` to the process group at 116 GiB used. - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. +- It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. +- It applies the same bounded process-group cleanup if an unexpected post-launch error occurs. - It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. @@ -31,6 +33,7 @@ Exit classifications are authoritative in the final JSON record. Operational fai | 5 | emergency threshold reached | | 6 | soft-threshold grace period expired | | 7 | process-group signaling or termination failure | +| 70 | unexpected post-launch error | | 127 | command launch failure | No model, backend, or ROCm package is required to run the unit tests: diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 0dc532c1d027..d8738d5041fd 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -31,10 +31,12 @@ EXIT_EMERGENCY_LIMIT = 5 EXIT_GRACE_TIMEOUT = 6 EXIT_SIGNAL_ERROR = 7 +EXIT_INTERNAL_ERROR = 70 EXIT_LAUNCH_ERROR = 127 MEMINFO_VALUE_RE = re.compile(r"([0-9]+) kB") SWAPS_HEADER = ["Filename", "Type", "Size", "Used", "Priority"] +PARENT_SIGNALS = (signal.SIGINT, signal.SIGTERM) class ProcfsError(RuntimeError): @@ -45,6 +47,12 @@ class ProcessGroupError(RuntimeError): pass +class ParentSignal(RuntimeError): + def __init__(self, signal_number: int): + self.signal_number = signal_number + super().__init__(signal.Signals(signal_number).name) + + class ProcessHandle(Protocol): pid: int @@ -66,6 +74,12 @@ def used_bytes(self) -> int: return self.total_bytes - self.available_bytes +@dataclass +class RuntimeState: + snapshot: HostSnapshot + peak_used_bytes: int + + @dataclass(frozen=True) class WatchdogConfig: command: tuple[str, ...] @@ -257,6 +271,39 @@ def _signal_process_group(process_group_id: int, signal_number: int) -> str: return f"{signal.Signals(signal_number).name.lower()}_sent" +def _process_group_alive(process_group_id: int) -> bool: + try: + os.killpg(process_group_id, 0) + except ProcessLookupError: + return False + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot inspect process group {process_group_id}: {detail}" + ) from exc + return True + + +def _raise_parent_signal(signal_number: int, _frame: object) -> None: + raise ParentSignal(signal_number) + + +def _set_parent_signal_handlers( + handler: signal.Handlers, +) -> dict[int, signal.Handlers]: + previous: dict[int, signal.Handlers] = {} + for signal_number in PARENT_SIGNALS: + previous[signal_number] = signal.signal(signal_number, handler) + return previous + + +def _restore_parent_signal_handlers( + previous: dict[int, signal.Handlers], +) -> None: + for signal_number, handler in previous.items(): + signal.signal(signal_number, handler) + + def _kill_and_finish( audit: AuditLogger, child: ProcessHandle, @@ -323,18 +370,105 @@ def _kill_and_finish( ) +def _graceful_cleanup( + audit: AuditLogger, + child: ProcessHandle, + snapshot: HostSnapshot, + peak_used_bytes: int, + classification: str, + exit_code: int, + reason: str, + graceful_signal: int, + grace_seconds: float, + signal_group: Callable[[int, int], str], + group_alive: Callable[[int], bool], + monotonic: Callable[[], float], + sleeper: Callable[[float], None], + error: str | None = None, +) -> int: + signal_events: list[tuple[int, str]] = [] + try: + group_status = signal_group(child.pid, graceful_signal) + signal_events.append((graceful_signal, group_status)) + deadline = monotonic() + grace_seconds + while monotonic() < deadline: + child.poll() + if not group_alive(child.pid): + break + sleeper(min(0.05, deadline - monotonic())) + child.poll() + if group_alive(child.pid): + group_status = signal_group(child.pid, signal.SIGKILL) + signal_events.append((signal.SIGKILL, group_status)) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "signal_error", + str(exc), + ) + + child_returncode = child.poll() + if child_returncode is None: + try: + child_returncode = child.wait(timeout=5.0) + except subprocess.TimeoutExpired as exc: + return _emit_final( + audit, + "termination_timeout", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "termination_timeout", + str(exc), + ) + + for signal_number, status in signal_events: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child_returncode, + status, + reason, + ), + signal=signal.Signals(signal_number).name, + ) + return _emit_final( + audit, + classification, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + signal_events[-1][1], + error, + ) + + def _monitor_child( config: WatchdogConfig, reader: ProcfsReader, audit: AuditLogger, child: ProcessHandle, - initial_snapshot: HostSnapshot, + state: RuntimeState, signal_group: Callable[[int, int], str], monotonic: Callable[[], float], sleeper: Callable[[float], None], ) -> int: - snapshot = initial_snapshot - peak = snapshot.used_bytes soft_deadline: float | None = None while True: @@ -354,8 +488,8 @@ def _monitor_child( if soft_stop else "child exited" ), - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, child, child_returncode, "leader_exited", @@ -366,8 +500,8 @@ def _monitor_child( return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "grace_timeout", EXIT_GRACE_TIMEOUT, "soft-threshold grace period expired", @@ -375,50 +509,60 @@ def _monitor_child( ) try: - snapshot = reader.read_snapshot() + state.snapshot = reader.read_snapshot() except ProcfsError as exc: return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "procfs_error", EXIT_PROCFS_ERROR, str(exc), signal_group, ) - peak = max(peak, snapshot.used_bytes) + state.peak_used_bytes = max( + state.peak_used_bytes, state.snapshot.used_bytes + ) audit.emit( "sample", **_state_fields( - snapshot, peak, child, None, "active", "none" + state.snapshot, + state.peak_used_bytes, + child, + None, + "active", + "none", ), ) - if snapshot.active_swaps: + if state.snapshot.active_swaps: return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "swap_appeared", EXIT_SWAP_ACTIVE, "active swap appeared during execution", signal_group, ) - if snapshot.used_bytes >= config.emergency_bytes: + if state.snapshot.used_bytes >= config.emergency_bytes: return _kill_and_finish( audit, child, - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, "emergency_limit", EXIT_EMERGENCY_LIMIT, "used_bytes >= emergency_bytes", signal_group, ) - if soft_deadline is None and snapshot.used_bytes >= config.soft_bytes: + if ( + soft_deadline is None + and state.snapshot.used_bytes >= config.soft_bytes + ): try: group_status = signal_group(child.pid, signal.SIGTERM) except ProcessGroupError as exc: @@ -427,8 +571,8 @@ def _monitor_child( "signal_error", EXIT_SIGNAL_ERROR, "used_bytes >= soft_bytes", - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, child, child.poll(), "signal_error", @@ -438,8 +582,8 @@ def _monitor_child( audit.emit( "process_group_signal", **_state_fields( - snapshot, - peak, + state.snapshot, + state.peak_used_bytes, child, child.poll(), group_status, @@ -465,6 +609,7 @@ def run_watchdog( audit: AuditLogger | None = None, launcher: Callable[..., ProcessHandle] | None = None, signal_group: Callable[[int, int], str] | None = None, + group_alive: Callable[[int], bool] | None = None, monotonic: Callable[[], float] | None = None, sleeper: Callable[[float], None] | None = None, ) -> int: @@ -473,6 +618,7 @@ def run_watchdog( audit = audit or AuditLogger(sys.stderr) launcher = launcher or subprocess.Popen signal_group = signal_group or _signal_process_group + group_alive = group_alive or _process_group_alive monotonic = monotonic or time.monotonic sleeper = sleeper or time.sleep @@ -532,42 +678,95 @@ def run_watchdog( snapshot.used_bytes, ) + previous_mask: set[signal.Signals] | None = signal.pthread_sigmask( + signal.SIG_BLOCK, PARENT_SIGNALS + ) + previous_handlers: dict[int, signal.Handlers] = {} + child: ProcessHandle | None = None + state = RuntimeState(snapshot, snapshot.used_bytes) try: - child = launcher(config.command, start_new_session=True) - except (OSError, ValueError) as exc: - detail = getattr(exc, "strerror", None) or str(exc) - return _emit_final( + try: + child = launcher(config.command, start_new_session=True) + except (OSError, ValueError) as exc: + detail = getattr(exc, "strerror", None) or str(exc) + return _emit_final( + audit, + "launch_error", + EXIT_LAUNCH_ERROR, + "command launch failed", + snapshot, + snapshot.used_bytes, + error=detail, + ) + + previous_handlers = _set_parent_signal_handlers( + _raise_parent_signal + ) + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + previous_mask = None + audit.emit( + "child_started", + **_state_fields( + state.snapshot, + state.peak_used_bytes, + child, + None, + "active", + "none", + ), + command=list(config.command), + ) + return _monitor_child( + config, + reader, audit, - "launch_error", - EXIT_LAUNCH_ERROR, - "command launch failed", - snapshot, - snapshot.used_bytes, - error=detail, + child, + state, + signal_group, + monotonic, + sleeper, ) - - audit.emit( - "child_started", - **_state_fields( - snapshot, - snapshot.used_bytes, + except ParentSignal as exc: + _set_parent_signal_handlers(signal.SIG_IGN) + signal_name = signal.Signals(exc.signal_number).name + return _graceful_cleanup( + audit, child, - None, - "active", - "none", - ), - command=list(config.command), - ) - return _monitor_child( - config, - reader, - audit, - child, - snapshot, - signal_group, - monotonic, - sleeper, - ) + state.snapshot, + state.peak_used_bytes, + "parent_signal", + 128 + exc.signal_number, + f"wrapper received {signal_name}", + exc.signal_number, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + ) + except Exception as exc: + _set_parent_signal_handlers(signal.SIG_IGN) + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "internal_error", + EXIT_INTERNAL_ERROR, + "unexpected post-launch exception", + signal.SIGTERM, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + f"{type(exc).__name__}: {exc}", + ) + finally: + if previous_mask is not None: + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + if previous_handlers: + _restore_parent_signal_handlers(previous_handlers) def _positive_int(value: str) -> int: diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index fa4c4aaec969..c0dbc9c1fcb3 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -5,10 +5,12 @@ import importlib.util import io import json +import os import signal import subprocess import sys import tempfile +import time import unittest from datetime import datetime, timezone from pathlib import Path @@ -113,6 +115,9 @@ def signal_group(self, process_group_id: int, signal_number: int) -> str: self.signal_handler(self.process, signal_number) return f"{signal.Signals(signal_number).name.lower()}_sent" + def group_alive(self, process_group_id: int) -> bool: + return self.process.returncode is None + def run(self, **overrides: Any) -> int: config = watchdog.WatchdogConfig( command=("fake-command",), @@ -128,6 +133,7 @@ def run(self, **overrides: Any) -> int: audit=self.audit, launcher=self.launcher, signal_group=self.signal_group, + group_alive=self.group_alive, monotonic=self.clock.monotonic, sleeper=self.clock.sleep, ) @@ -196,6 +202,128 @@ def test_rejects_malformed_or_missing_procfs_data(self) -> None: class TestWatchdogBehavior(unittest.TestCase): + @staticmethod + def _process_is_running(process_id: int) -> bool: + result = subprocess.run( + ["ps", "-o", "stat=", "-p", str(process_id)], + capture_output=True, + check=False, + text=True, + ) + return result.returncode == 0 and not result.stdout.lstrip().startswith( + "Z" + ) + + def test_parent_signals_leave_no_child_or_grandchild(self) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGINT,signal.SIG_IGN);" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n');" + " time.sleep(30)\n" + ) + for signal_number in (signal.SIGINT, signal.SIGTERM): + with self.subTest(signal=signal.Signals(signal_number).name): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + (root / "meminfo").write_text( + "MemTotal: 131072 kB\n" + "MemAvailable: 65536 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + audit_path = root / "audit.jsonl" + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + self.fail( + "child process group did not start" + ) + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + time.sleep(0.05) + wrapper.send_signal(signal_number) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + if child_pid is not None: + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + self.assertEqual( + wrapper.returncode, 128 + signal_number + ) + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + self.assertEqual( + records[-1]["classification"], "parent_signal" + ) + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual( + forwarded[0], + signal.Signals(signal_number).name, + ) + self.assertEqual(forwarded[-1], "SIGKILL") + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse( + self._process_is_running(process_id) + ) + def test_configuration_rejects_non_finite_timing(self) -> None: config = watchdog.WatchdogConfig( command=("fake-command",), @@ -384,6 +512,28 @@ def exit_on_kill(process: FakeProcess, signal_number: int) -> None: harness.records()[-1]["classification"], "procfs_error" ) + def test_unexpected_monitor_error_cleans_up_process_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), RuntimeError("unexpected")], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_INTERNAL_ERROR) + self.assertEqual( + harness.signals, + [signal.SIGTERM, signal.SIGKILL], + ) + final = harness.records()[-1] + self.assertEqual(final["classification"], "internal_error") + self.assertIn("RuntimeError: unexpected", final["error"]) + def test_launch_failure_is_explicit(self) -> None: harness = Harness([snapshot(50)], FakeProcess()) From a55e6a306e8957b2114f433402e39bddd039139f Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:37:36 -0700 Subject: [PATCH 08/32] scripts : preserve child signals and monitor descendants Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 3 +- scripts/strix_memory_watchdog.py | 129 ++++++++++++----- tests/test_strix_memory_watchdog.py | 206 +++++++++++++++++++++++++--- 3 files changed, 282 insertions(+), 56 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 8d6d264b141b..5a6558f0a9c0 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -14,6 +14,7 @@ The wrapper performs these checks and actions: - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. - It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. +- It checks the process group after the direct child exits and cleans up remaining descendants before returning the child's classification. - It applies the same bounded process-group cleanup if an unexpected post-launch error occurs. - It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. @@ -21,7 +22,7 @@ The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. -The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Child standard input, standard output, and standard error are inherited unchanged. +The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Signal records are written immediately after each process-group signal. Child standard input, standard output, and standard error are inherited unchanged. Exit classifications are authoritative in the final JSON record. Operational failures use these exit codes: diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index d8738d5041fd..79d7a0013e77 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -378,18 +378,32 @@ def _graceful_cleanup( classification: str, exit_code: int, reason: str, - graceful_signal: int, + graceful_signal: int | None, grace_seconds: float, signal_group: Callable[[int, int], str], group_alive: Callable[[int], bool], monotonic: Callable[[], float], sleeper: Callable[[float], None], + process_group_status: str = "active", error: str | None = None, ) -> int: - signal_events: list[tuple[int, str]] = [] try: - group_status = signal_group(child.pid, graceful_signal) - signal_events.append((graceful_signal, group_status)) + if graceful_signal is not None: + process_group_status = signal_group( + child.pid, graceful_signal + ) + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + reason, + ), + signal=signal.Signals(graceful_signal).name, + ) deadline = monotonic() + grace_seconds while monotonic() < deadline: child.poll() @@ -398,8 +412,21 @@ def _graceful_cleanup( sleeper(min(0.05, deadline - monotonic())) child.poll() if group_alive(child.pid): - group_status = signal_group(child.pid, signal.SIGKILL) - signal_events.append((signal.SIGKILL, group_status)) + process_group_status = signal_group( + child.pid, signal.SIGKILL + ) + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + reason, + ), + signal="SIGKILL", + ) except ProcessGroupError as exc: return _emit_final( audit, @@ -432,19 +459,6 @@ def _graceful_cleanup( str(exc), ) - for signal_number, status in signal_events: - audit.emit( - "process_group_signal", - **_state_fields( - snapshot, - peak_used_bytes, - child, - child_returncode, - status, - reason, - ), - signal=signal.Signals(signal_number).name, - ) return _emit_final( audit, classification, @@ -454,7 +468,7 @@ def _graceful_cleanup( peak_used_bytes, child, child_returncode, - signal_events[-1][1], + process_group_status, error, ) @@ -466,6 +480,7 @@ def _monitor_child( child: ProcessHandle, state: RuntimeState, signal_group: Callable[[int, int], str], + group_alive: Callable[[int], bool], monotonic: Callable[[], float], sleeper: Callable[[float], None], ) -> int: @@ -475,19 +490,50 @@ def _monitor_child( child_returncode = child.poll() if child_returncode is not None: soft_stop = soft_deadline is not None + classification = "soft_limit" if soft_stop else "child_exit" + exit_code = EXIT_SOFT_LIMIT if soft_stop else ( + 128 - child_returncode + if child_returncode < 0 + else child_returncode + ) + reason = ( + "child exited during soft-threshold grace period" + if soft_stop + else "child exited" + ) + if group_alive(child.pid): + grace_seconds = config.grace_seconds + graceful_signal: int | None = signal.SIGTERM + group_status = "active" + if soft_stop: + grace_seconds = max( + 0.0, soft_deadline - monotonic() + ) + graceful_signal = None + group_status = "sigterm_sent" + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + classification, + exit_code, + ( + f"{reason}; process group members still running" + ), + graceful_signal, + grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + group_status, + ) return _emit_final( audit, - "soft_limit" if soft_stop else "child_exit", - EXIT_SOFT_LIMIT if soft_stop else ( - 128 - child_returncode - if child_returncode < 0 - else child_returncode - ), - ( - "child exited during soft-threshold grace period" - if soft_stop - else "child exited" - ), + classification, + exit_code, + reason, state.snapshot, state.peak_used_bytes, child, @@ -678,15 +724,25 @@ def run_watchdog( snapshot.used_bytes, ) - previous_mask: set[signal.Signals] | None = signal.pthread_sigmask( + previous_mask = signal.pthread_sigmask( signal.SIG_BLOCK, PARENT_SIGNALS ) + mask_restored = False previous_handlers: dict[int, signal.Handlers] = {} child: ProcessHandle | None = None state = RuntimeState(snapshot, snapshot.used_bytes) try: + launch_mask = previous_mask + + def restore_child_signal_mask() -> None: + signal.pthread_sigmask(signal.SIG_SETMASK, launch_mask) + try: - child = launcher(config.command, start_new_session=True) + child = launcher( + config.command, + start_new_session=True, + preexec_fn=restore_child_signal_mask, + ) except (OSError, ValueError) as exc: detail = getattr(exc, "strerror", None) or str(exc) return _emit_final( @@ -703,7 +759,7 @@ def run_watchdog( _raise_parent_signal ) signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) - previous_mask = None + mask_restored = True audit.emit( "child_started", **_state_fields( @@ -723,6 +779,7 @@ def run_watchdog( child, state, signal_group, + group_alive, monotonic, sleeper, ) @@ -760,10 +817,10 @@ def run_watchdog( group_alive, monotonic, sleeper, - f"{type(exc).__name__}: {exc}", + error=f"{type(exc).__name__}: {exc}", ) finally: - if previous_mask is not None: + if not mask_restored: signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) if previous_handlers: _restore_parent_signal_handlers(previous_handlers) diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index c0dbc9c1fcb3..2e59eec27512 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -214,6 +214,17 @@ def _process_is_running(process_id: int) -> bool: "Z" ) + @staticmethod + def _write_procfs_fixture(root: Path) -> None: + (root / "meminfo").write_text( + "MemTotal: 131072 kB\nMemAvailable: 65536 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + def test_parent_signals_leave_no_child_or_grandchild(self) -> None: child_code = ( "import os,signal,sys,time;" @@ -232,15 +243,7 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) pid_file = root / "pids" - (root / "meminfo").write_text( - "MemTotal: 131072 kB\n" - "MemAvailable: 65536 kB\n", - encoding="utf-8", - ) - (root / "swaps").write_text( - "Filename Type Size Used Priority\n", - encoding="utf-8", - ) + self._write_procfs_fixture(root) audit_path = root / "audit.jsonl" with audit_path.open("w", encoding="utf-8") as audit: wrapper = subprocess.Popen( @@ -303,16 +306,29 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: self.assertEqual( records[-1]["classification"], "parent_signal" ) - forwarded = [ - record["signal"] + signal_records = [ + record for record in records if record["event"] == "process_group_signal" ] + forwarded = [ + record["signal"] for record in signal_records + ] self.assertEqual( forwarded[0], signal.Signals(signal_number).name, ) self.assertEqual(forwarded[-1], "SIGKILL") + self.assertEqual( + signal_records[0]["child_status"], "running" + ) + self.assertIsNone( + signal_records[0]["child_returncode"] + ) + self.assertLess( + signal_records[0]["timestamp"], + signal_records[-1]["timestamp"], + ) for process_id in (child_pid, grandchild_pid): deadline = time.monotonic() + 2 while ( @@ -324,6 +340,165 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: self._process_is_running(process_id) ) + def test_child_sigterm_handler_exits_without_escalation(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " open(sys.argv[2],'w').write('handled\\n')\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "open(sys.argv[1],'w').write(f'{os.getpid()}\\n')\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + ready_path = root / "ready" + handled_path = root / "handled" + audit_path = root / "audit.jsonl" + self._write_procfs_fixture(root) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.5", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(ready_path), + str(handled_path), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not ready_path.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("SIGTERM child did not become ready") + time.sleep(0.01) + child_pid = int( + ready_path.read_text(encoding="utf-8").strip() + ) + try: + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(wrapper.returncode, 128 + signal.SIGTERM) + self.assertTrue(handled_path.exists()) + self.assertEqual(forwarded, ["SIGTERM"]) + self.assertEqual(records[-1]["child_returncode"], 0) + + def test_leader_exit_cleans_up_surviving_grandchild(self) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n')\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + audit_path = root / "audit.jsonl" + self._write_procfs_fixture(root) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("leader process did not write child PIDs") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + try: + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(wrapper.returncode, 0) + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual(records[-1]["child_returncode"], 0) + self.assertEqual(forwarded, ["SIGTERM", "SIGKILL"]) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + def test_configuration_rejects_non_finite_timing(self) -> None: config = watchdog.WatchdogConfig( command=("fake-command",), @@ -335,14 +510,7 @@ def test_configuration_rejects_non_finite_timing(self) -> None: def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) - (root / "meminfo").write_text( - "MemTotal: 131072 kB\nMemAvailable: 65536 kB\n", - encoding="utf-8", - ) - (root / "swaps").write_text( - "Filename Type Size Used Priority\n", - encoding="utf-8", - ) + self._write_procfs_fixture(root) result = subprocess.run( [ sys.executable, From ae910f804ad8beee57d6f9f2240f5334728d00c6 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:45:06 -0700 Subject: [PATCH 09/32] scripts : classify soft descendant escalation as timeout Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 1 + scripts/strix_memory_watchdog.py | 22 +++++- tests/test_strix_memory_watchdog.py | 108 ++++++++++++++++++++++++++++ 3 files changed, 130 insertions(+), 1 deletion(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 5a6558f0a9c0..ccd04f7ab667 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -12,6 +12,7 @@ The wrapper performs these checks and actions: - It calculates used memory as `MemTotal - MemAvailable`. Linux reports these fields in KiB, so the wrapper multiplies each value by 1024 and keeps all accounting as integer bytes. - It sends `SIGTERM` to the process group at 116 GiB used. - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. +- It reports `grace_timeout` if any descendant requires `SIGKILL` after the soft-threshold grace period, even when the direct child exited earlier. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. - It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. - It checks the process group after the direct child exits and cleans up remaining descendants before returning the child's classification. diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 79d7a0013e77..00a9fab78d81 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -385,8 +385,10 @@ def _graceful_cleanup( monotonic: Callable[[], float], sleeper: Callable[[float], None], process_group_status: str = "active", + escalation_result: tuple[str, int, str] | None = None, error: str | None = None, ) -> int: + escalated = False try: if graceful_signal is not None: process_group_status = signal_group( @@ -412,9 +414,15 @@ def _graceful_cleanup( sleeper(min(0.05, deadline - monotonic())) child.poll() if group_alive(child.pid): + escalated = True process_group_status = signal_group( child.pid, signal.SIGKILL ) + signal_reason = ( + escalation_result[2] + if escalation_result is not None + else reason + ) audit.emit( "process_group_signal", **_state_fields( @@ -423,7 +431,7 @@ def _graceful_cleanup( child, child.poll(), process_group_status, - reason, + signal_reason, ), signal="SIGKILL", ) @@ -459,6 +467,8 @@ def _graceful_cleanup( str(exc), ) + if escalated and escalation_result is not None: + classification, exit_code, reason = escalation_result return _emit_final( audit, classification, @@ -528,6 +538,16 @@ def _monitor_child( monotonic, sleeper, group_status, + ( + ( + "grace_timeout", + EXIT_GRACE_TIMEOUT, + "soft-threshold grace period expired with " + "process group members still running", + ) + if soft_stop + else None + ), ) return _emit_final( audit, diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 2e59eec27512..898298efb9b9 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -499,6 +499,114 @@ def test_leader_exit_cleans_up_surviving_grandchild(self) -> None: time.sleep(0.01) self.assertFalse(self._process_is_running(process_id)) + def test_soft_limit_descendant_escalation_is_grace_timeout(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "grandchild=os.fork()\n" + "if grandchild == 0:\n" + " signal.signal(signal.SIGTERM,signal.SIG_IGN)\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n')\n" + " time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + audit_path = root / "audit.jsonl" + (root / "meminfo").write_text( + "MemTotal: 3145728 kB\n" + "MemAvailable: 2621440 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--soft-gib", + "1", + "--emergency-gib", + "2", + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("soft-limit process group did not start") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + next_meminfo = root / "meminfo.next" + next_meminfo.write_text( + "MemTotal: 3145728 kB\n" + "MemAvailable: 1572864 kB\n", + encoding="utf-8", + ) + next_meminfo.replace(root / "meminfo") + try: + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + final = records[-1] + self.assertEqual(wrapper.returncode, watchdog.EXIT_GRACE_TIMEOUT) + self.assertEqual(final["classification"], "grace_timeout") + self.assertEqual(final["child_returncode"], 0) + self.assertEqual(forwarded, ["SIGTERM", "SIGKILL"]) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + def test_configuration_rejects_non_finite_timing(self) -> None: config = watchdog.WatchdogConfig( command=("fake-command",), From af2b27f3ebde8a275f2db04576280c933ca6f4b4 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:28:00 -0700 Subject: [PATCH 10/32] docs : compose admission with Strix watchdog Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- README.md | 2 +- docs/strix-memory-watchdog.md | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index dbf547ce5c5b..e073b9fd1592 100644 --- a/README.md +++ b/README.md @@ -118,7 +118,7 @@ Everything else is upstream `llama.cpp`. The additions currently carried here: | Speculative checkpoints on device | | `llama-server` keeps speculative-decoding checkpoints in device memory instead of copying them to the host | | ROCmFPx quant types | `llama-quantize` types `Q4_0_ROCMFP4`, `Q4_0_ROCMFP4_FAST`, `Q2/Q3/Q6/Q8_0_ROCMFPX` and the `_LEAN`/`_COHERENT`/`_STRIX` recipes | Loads the ROCmFP4 GGUFs published for Strix Halo. CPU codecs plus Vulkan dequant, mat-vec, matmul and integer-dot kernels. Weight formats only: not accepted as KV-cache types | | Repeatable output at depth | | Freed KV cells are zeroed so masked-out rows never carry stale K/V, and the Vulkan radix top-k assigns output slots deterministically | -| Host-memory watchdog | [`scripts/strix_memory_watchdog.py`](docs/strix-memory-watchdog.md) | Runs a command in a process group, requires zero active swap, and stops before host-wide memory reaches the 120 GiB validation ceiling | +| DeepSeek V4.1 memory guard | [`scripts/strix_memory_watchdog.py`](docs/strix-memory-watchdog.md) | In-process admission auto-fits expert slots under 116 GiB before allocation; the external process-group watchdog requires zero swap and stops before the 120 GiB validation ceiling | Every ROCm/HIP change above is guarded on architecture, shape and layout, so other devices see upstream behaviour. Run `--help`, or see [tools/server/README.md](tools/server/README.md), for the full options. diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index ccd04f7ab667..7adbc3281750 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -3,9 +3,18 @@ `scripts/strix_memory_watchdog.py` is an external Linux command wrapper for headless Strix Halo validation. It does not change model loading or cache sizing. It measures host-wide memory from procfs and controls the launched command's process group. ```sh -./scripts/strix_memory_watchdog.py -- ./build/bin/llama-server +./scripts/strix_memory_watchdog.py -- \ + ./build/bin/llama-server \ + -m /mnt/models/deepseek-v41/DeepSeek-V4.1-Flash-Q2.gguf \ + -c 32768 -b 2048 -ub 2048 -ngl 99 ``` +DeepSeek V4.1 also runs an in-process admission check before expert-cache or model backend allocation. The default model parameters read `/proc/meminfo` and `/proc/swaps`, reject any configured swap entry, account unified host/GPU memory once, and auto-fit complete expert slots under 116 GiB total projected host use. The external watchdog is still required for guarded validation because it monitors host-wide use after startup and controls the complete process group. + +Use `--dsv41-procfs-root`, `--dsv41-memory-soft-mib`, `--dsv41-memory-watchdog-mib`, `--dsv41-memory-hard-mib`, and `--dsv41-memory-safety-margin-mib` only when reproducing admission tests or applying a more conservative host policy. `--expert-cache-slots` and `--expert-cache-mib` are optional caps; zero auto-fits. If both cache options are set, their capacity must describe the same number of complete published tensor slots. + +Admission accepts context checkpoints 32768, 65536, 98304, and 131072. It never lowers an explicit context request. A request that does not fit reports current use, fixed tensor bytes, state bytes, graph workspace, Engram and expert staging, output bytes, selected cache slots and bytes, safety margin, all thresholds, and the rejecting category. + The wrapper performs these checks and actions: - It refuses to launch if `/proc/swaps` contains any active entry. From c2a7661f38bca0160b88e812c4474d1e31270fe5 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:33:00 -0700 Subject: [PATCH 11/32] deepseek41 : tighten admission failure guards Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- common/arg.cpp | 2 +- src/llama-dsv41-admission.cpp | 15 ++++++++++++++- src/models/deepseek41.cpp | 6 ++++-- tests/test-deepseek41-admission.cpp | 13 ++++++++++++- 4 files changed, 31 insertions(+), 5 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 47b08f0298e9..df67faf06f1c 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2891,7 +2891,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--dsv41-memory-safety-margin-mib"}, "MiB", string_format("DeepSeek V4.1 explicit startup safety margin (default: %d)", params.dsv41_memory_safety_margin_mib), [](common_params & params, int value) { - if (value < 0) { + if (value <= 0) { throw std::invalid_argument("invalid value"); } params.dsv41_memory_safety_margin_mib = value; diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp index d442ad52cfc6..1b600a41a8af 100644 --- a/src/llama-dsv41-admission.cpp +++ b/src/llama-dsv41-admission.cpp @@ -182,6 +182,8 @@ uint64_t llama_dsv41_estimate_graph_workspace(uint32_t n_ctx, uint32_t n_ubatch) if (n_ubatch == 0 || n_ubatch > 2048) { throw std::runtime_error("DeepSeek V4.1 bounded admission requires n_ubatch in 1..2048"); } + // Conservative ds4 graph bound: 7.884 GiB total state at 32K and 8.951 GiB at 131K. + // Replace this estimate when the full graph can report exact no-alloc reserve bytes before model allocation. const uint64_t base = 7688ULL << 20; return checked_add(base, checked_mul(n_ctx, 7424, "graph workspace"), "graph workspace"); } @@ -206,8 +208,9 @@ uint64_t llama_dsv41_output_bytes(uint32_t n_vocab, uint32_t n_ubatch) { checked_mul(n_vocab, n_ubatch, "output floats"), 2*sizeof(float), "output floats"); + const uint64_t token_rows = checked_add(n_vocab, 1, "output tokens"); const uint64_t tokens = checked_mul( - checked_mul(n_vocab, n_ubatch, "output tokens"), + checked_mul(token_rows, n_ubatch, "output tokens"), sizeof(int32_t), "output tokens"); return checked_add(floats, tokens, "outputs"); @@ -252,6 +255,9 @@ llama_dsv41_admission_result llama_dsv41_admit( params.hard_bytes > LLAMA_DSV41_ADMISSION_HARD_BYTES) { reject("thresholds", result, "require soft < watchdog < hard <= 120 GiB"); } + if (params.safety_margin_bytes == 0) { + reject("thresholds", result, "safety margin must be non-zero"); + } validate_context(params.n_ctx); if (params.n_seq != 1) { reject("context", result, "bounded DeepSeek V4.1 admission currently requires one sequence"); @@ -326,6 +332,10 @@ llama_dsv41_admission_result llama_dsv41_admit( result.projected_bytes = result.fixed_bytes; reject("fixed", result, "fixed startup categories exceed the soft limit"); } + if (result.fixed_bytes > result.host_total) { + result.projected_bytes = result.fixed_bytes; + reject("host", result, "fixed startup categories exceed physical host memory"); + } const uint64_t bytes_per_slot = checked_add( result.expert_slot_bytes, result.expert_staging_slot_bytes, "expert slot and staging"); @@ -346,6 +356,9 @@ llama_dsv41_admission_result llama_dsv41_admit( if (result.projected_bytes > result.soft_bytes) { reject("soft", result, "projected startup exceeds the soft limit"); } + if (result.projected_bytes > result.host_total) { + reject("host", result, "projected startup exceeds physical host memory"); + } if (result.projected_bytes >= result.hard_bytes) { reject("hard", result, "projected startup is not strictly below the hard limit"); } diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 0b962daaf3eb..3caf40c21469 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -357,9 +357,11 @@ void llama_model_deepseek41::validate_context_params(const llama_cparams & cpara if (cparams.n_ctx > admission->result.n_ctx || cparams.n_seq_max > admission->result.n_seq || cparams.n_ubatch > admission->result.n_ubatch) { + llama_dsv41_admission_result failure = admission->result; + failure.category = "context"; throw std::runtime_error(format( - "%s, category=context, requested_context=%u, requested_sequences=%u, requested_ubatch=%u", - admission->result.describe().c_str(), + "%s, requested_context=%u, requested_sequences=%u, requested_ubatch=%u", + failure.describe().c_str(), cparams.n_ctx, cparams.n_seq_max, cparams.n_ubatch)); diff --git a/tests/test-deepseek41-admission.cpp b/tests/test-deepseek41-admission.cpp index 790ca2112d95..3dc80a81d5a7 100644 --- a/tests/test-deepseek41-admission.cpp +++ b/tests/test-deepseek41-admission.cpp @@ -183,7 +183,7 @@ void test_threshold_boundaries() { auto params = base_params(); params.configured_cache_slots = LLAMA_DSV41_N_EXPERT_USED; params.configured_cache_bytes = params.configured_cache_slots*398131200ULL; - params.safety_margin_bytes = 0; + params.safety_margin_bytes = 1; const auto zero = llama_dsv41_admit(host_with_used(0), 0, tensors, params); const uint64_t planned_without_host = zero.projected_bytes; @@ -247,6 +247,17 @@ void test_diagnostics_and_guards() { "expert_slots=", "expert_cache=", "expert_staging=", "soft=", "watchdog=", "hard=" }) { REQUIRE(diagnostic.find(field) != std::string::npos); } + + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), UINT64_MAX, tensors, params); + }).find("overflow") != std::string::npos); + + auto small_host = host_with_used(0); + small_host.total = 8ULL << 30; + small_host.available = small_host.total; + REQUIRE(thrown([&]() { + llama_dsv41_admit(small_host, 0, tensors, params); + }).find("physical host memory") != std::string::npos); } } From 856de69a63c1b2fd0e5c74e3884014f8b977cfad Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:34:03 -0700 Subject: [PATCH 12/32] deepseek41 : remove unused admission helper Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-dsv41-admission.cpp | 12 ------------ src/llama-dsv41-admission.h | 1 - 2 files changed, 13 deletions(-) diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp index 1b600a41a8af..04fb4b4a73da 100644 --- a/src/llama-dsv41-admission.cpp +++ b/src/llama-dsv41-admission.cpp @@ -165,18 +165,6 @@ llama_dsv41_host_memory llama_dsv41_read_host_memory(const std::string & procfs_ return result; } -uint64_t llama_dsv41_expert_payload_bytes(const std::vector & tensors) { - uint64_t result = 0; - for (const auto & tensor : tensors) { - llama_expert_store_validate_tensor(tensor); - result = checked_add( - result, - checked_mul(tensor.nb[2], tensor.ne[2], "expert payload"), - "expert payload"); - } - return result; -} - uint64_t llama_dsv41_estimate_graph_workspace(uint32_t n_ctx, uint32_t n_ubatch) { validate_context(n_ctx); if (n_ubatch == 0 || n_ubatch > 2048) { diff --git a/src/llama-dsv41-admission.h b/src/llama-dsv41-admission.h index e41ef0b33a57..a278079f64e0 100644 --- a/src/llama-dsv41-admission.h +++ b/src/llama-dsv41-admission.h @@ -71,7 +71,6 @@ struct llama_dsv41_admission_result { llama_dsv41_host_memory llama_dsv41_read_host_memory(const std::string & procfs_root); -uint64_t llama_dsv41_expert_payload_bytes(const std::vector & tensors); uint64_t llama_dsv41_estimate_graph_workspace(uint32_t n_ctx, uint32_t n_ubatch); uint64_t llama_dsv41_engram_staging_bytes(uint32_t n_ubatch); uint64_t llama_dsv41_output_bytes(uint32_t n_vocab, uint32_t n_ubatch); From 7a3f93e79e2fa3330fb45e32964a55d65daf7231 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:38:44 -0700 Subject: [PATCH 13/32] deepseek41 : enforce fixed memory ceilings Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-dsv41-admission.cpp | 20 ++++++++++++++++---- tests/test-deepseek41-admission.cpp | 26 ++++++++++++++++++++++++-- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp index 04fb4b4a73da..6d32fb082629 100644 --- a/src/llama-dsv41-admission.cpp +++ b/src/llama-dsv41-admission.cpp @@ -238,10 +238,12 @@ llama_dsv41_admission_result llama_dsv41_admit( if (!params.unified_memory) { reject("unified_memory", result, "Strix admission requires one unified host/GPU memory pool"); } - if (params.soft_bytes == 0 || params.soft_bytes >= params.watchdog_bytes || + if (params.soft_bytes == 0 || params.soft_bytes > LLAMA_DSV41_ADMISSION_SOFT_BYTES || + params.watchdog_bytes > LLAMA_DSV41_WATCHDOG_EMERGENCY_BYTES || + params.soft_bytes >= params.watchdog_bytes || params.watchdog_bytes >= params.hard_bytes || params.hard_bytes > LLAMA_DSV41_ADMISSION_HARD_BYTES) { - reject("thresholds", result, "require soft < watchdog < hard <= 120 GiB"); + reject("thresholds", result, "require soft <= 116 GiB, watchdog <= 118 GiB, and soft < watchdog < hard <= 120 GiB"); } if (params.safety_margin_bytes == 0) { reject("thresholds", result, "safety margin must be non-zero"); @@ -276,10 +278,17 @@ llama_dsv41_admission_result llama_dsv41_admit( result.expert_staging_slot_bytes = std::max(result.expert_staging_slot_bytes, bytes); } + const uint64_t max_cache_bytes = checked_mul( + result.expert_slot_bytes, LLAMA_DSV41_N_EXPERT, "maximum expert cache"); + if (params.configured_cache_slots > LLAMA_DSV41_N_EXPERT || + params.configured_cache_bytes > max_cache_bytes) { + reject("cache", result, "configured cache exceeds the published expert count"); + } const uint64_t bytes_slots = params.configured_cache_bytes == 0 ? LLAMA_DSV41_N_EXPERT : params.configured_cache_bytes/result.expert_slot_bytes; if (params.configured_cache_slots != 0 && params.configured_cache_bytes != 0 && - params.configured_cache_slots != bytes_slots) { + params.configured_cache_bytes != checked_mul( + params.configured_cache_slots, result.expert_slot_bytes, "configured expert cache")) { reject("cache", result, "configured cache slots and bytes disagree"); } uint64_t slot_cap = LLAMA_DSV41_N_EXPERT; @@ -356,7 +365,8 @@ llama_dsv41_admission_result llama_dsv41_admit( std::string llama_dsv41_admission_result::describe() const { return format( - "DeepSeek V4.1 memory admission: category=%s, context=%u, sequences=%u, ubatch=%u, current=%llu, fixed=%llu, " + "DeepSeek V4.1 memory admission: category=%s, context=%u, sequences=%u, ubatch=%u, " + "host_total=%llu, host_available=%llu, current=%llu, fixed=%llu, " "dense=%llu, state=%llu, workspace=%llu, engram_staging=%llu, expert_slots=%u, " "expert_cache=%llu, expert_staging=%llu, outputs=%llu, safety_margin=%llu, projected=%llu, " "soft=%llu, watchdog=%llu, hard=%llu, device_reported_ignored=%llu", @@ -364,6 +374,8 @@ std::string llama_dsv41_admission_result::describe() const { n_ctx, n_seq, n_ubatch, + (unsigned long long) host_total, + (unsigned long long) host_available, (unsigned long long) host_used, (unsigned long long) fixed_bytes, (unsigned long long) dense_tensor_bytes, diff --git a/tests/test-deepseek41-admission.cpp b/tests/test-deepseek41-admission.cpp index 3dc80a81d5a7..9e2e53811bab 100644 --- a/tests/test-deepseek41-admission.cpp +++ b/tests/test-deepseek41-admission.cpp @@ -161,11 +161,16 @@ void test_configured_cache() { const auto tensors = published_tensors(); auto params = base_params(); params.configured_cache_slots = 12; - params.configured_cache_bytes = 12*398131200ULL + 1024; + params.configured_cache_bytes = 12*398131200ULL; const auto result = llama_dsv41_admit(host_with_used(0), 0, tensors, params); REQUIRE(result.expert_slots == 12); REQUIRE(result.expert_cache_bytes == 12*398131200ULL); + params.configured_cache_bytes += 1024; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("disagree") != std::string::npos); + params.configured_cache_slots = 13; REQUIRE(thrown([&]() { llama_dsv41_admit(host_with_used(0), 0, tensors, params); @@ -176,6 +181,11 @@ void test_configured_cache() { REQUIRE(thrown([&]() { llama_dsv41_admit(host_with_used(0), 0, tensors, params); }).find("top-k") != std::string::npos); + + params.configured_cache_slots = LLAMA_DSV41_N_EXPERT + 1; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("published expert count") != std::string::npos); } void test_threshold_boundaries() { @@ -201,6 +211,17 @@ void test_threshold_boundaries() { REQUIRE(thrown([&]() { llama_dsv41_admit(host_with_used(params.hard_bytes), 0, tensors, params); }).find("category=hard") != std::string::npos); + + params.soft_bytes = LLAMA_DSV41_ADMISSION_SOFT_BYTES + 1; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("category=thresholds") != std::string::npos); + + params.soft_bytes = LLAMA_DSV41_ADMISSION_SOFT_BYTES; + params.watchdog_bytes = LLAMA_DSV41_WATCHDOG_EMERGENCY_BYTES + 1; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("category=thresholds") != std::string::npos); } void test_context_progression() { @@ -244,7 +265,8 @@ void test_diagnostics_and_guards() { const std::string diagnostic = result.describe(); for (const char * field : { "category=", "current=", "fixed=", "dense=", "state=", "workspace=", - "expert_slots=", "expert_cache=", "expert_staging=", "soft=", "watchdog=", "hard=" }) { + "host_total=", "host_available=", "expert_slots=", "expert_cache=", "expert_staging=", + "soft=", "watchdog=", "hard=" }) { REQUIRE(diagnostic.find(field) != std::string::npos); } From 1da527ce19ebcd98b2c5537f36a6e7713dedaaad Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:39:33 -0700 Subject: [PATCH 14/32] scripts : cap Strix watchdog thresholds Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 2 +- scripts/strix_memory_watchdog.py | 6 ++++-- tests/test_strix_memory_watchdog.py | 12 ++++++++++++ 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 7adbc3281750..d02a3a7d5426 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -30,7 +30,7 @@ The wrapper performs these checks and actions: The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. -Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. +Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. Threshold overrides may only lower the 116 GiB soft and 118 GiB emergency limits. The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Signal records are written immediately after each process-group signal. Child standard input, standard output, and standard error are inherited unchanged. diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 00a9fab78d81..b91b6f0b44b0 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -94,10 +94,12 @@ def validate(self) -> None: raise ValueError("a command is required after --") if self.soft_bytes <= 0: raise ValueError("soft threshold must be greater than zero") + if self.soft_bytes > DEFAULT_SOFT_BYTES: + raise ValueError("soft threshold must not exceed 116 GiB") if self.emergency_bytes <= self.soft_bytes: raise ValueError("emergency threshold must be greater than soft threshold") - if self.emergency_bytes >= STRICT_CEILING_BYTES: - raise ValueError("emergency threshold must be below 120 GiB") + if self.emergency_bytes > DEFAULT_EMERGENCY_BYTES: + raise ValueError("emergency threshold must not exceed 118 GiB") if not math.isfinite(self.grace_seconds) or self.grace_seconds <= 0: raise ValueError("grace period must be greater than zero") if ( diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 898298efb9b9..d4eea8e00972 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -615,6 +615,18 @@ def test_configuration_rejects_non_finite_timing(self) -> None: with self.assertRaisesRegex(ValueError, "grace period"): config.validate() + def test_configuration_rejects_thresholds_above_policy(self) -> None: + with self.assertRaisesRegex(ValueError, "116 GiB"): + watchdog.WatchdogConfig( + command=("fake-command",), + soft_bytes=watchdog.DEFAULT_SOFT_BYTES + 1, + ).validate() + with self.assertRaisesRegex(ValueError, "118 GiB"): + watchdog.WatchdogConfig( + command=("fake-command",), + emergency_bytes=watchdog.DEFAULT_EMERGENCY_BYTES + 1, + ).validate() + def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) From e18f4e6988d027bccc6d333c804b542e188e4a64 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:40:02 -0700 Subject: [PATCH 15/32] deepseek41 : compare complete cache slot caps Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-dsv41-admission.cpp | 3 +-- tests/test-deepseek41-admission.cpp | 7 +------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp index 6d32fb082629..10fd37df8825 100644 --- a/src/llama-dsv41-admission.cpp +++ b/src/llama-dsv41-admission.cpp @@ -287,8 +287,7 @@ llama_dsv41_admission_result llama_dsv41_admit( const uint64_t bytes_slots = params.configured_cache_bytes == 0 ? LLAMA_DSV41_N_EXPERT : params.configured_cache_bytes/result.expert_slot_bytes; if (params.configured_cache_slots != 0 && params.configured_cache_bytes != 0 && - params.configured_cache_bytes != checked_mul( - params.configured_cache_slots, result.expert_slot_bytes, "configured expert cache")) { + params.configured_cache_slots != bytes_slots) { reject("cache", result, "configured cache slots and bytes disagree"); } uint64_t slot_cap = LLAMA_DSV41_N_EXPERT; diff --git a/tests/test-deepseek41-admission.cpp b/tests/test-deepseek41-admission.cpp index 9e2e53811bab..de769e7f2592 100644 --- a/tests/test-deepseek41-admission.cpp +++ b/tests/test-deepseek41-admission.cpp @@ -161,16 +161,11 @@ void test_configured_cache() { const auto tensors = published_tensors(); auto params = base_params(); params.configured_cache_slots = 12; - params.configured_cache_bytes = 12*398131200ULL; + params.configured_cache_bytes = 12*398131200ULL + 1024; const auto result = llama_dsv41_admit(host_with_used(0), 0, tensors, params); REQUIRE(result.expert_slots == 12); REQUIRE(result.expert_cache_bytes == 12*398131200ULL); - params.configured_cache_bytes += 1024; - REQUIRE(thrown([&]() { - llama_dsv41_admit(host_with_used(0), 0, tensors, params); - }).find("disagree") != std::string::npos); - params.configured_cache_slots = 13; REQUIRE(thrown([&]() { llama_dsv41_admit(host_with_used(0), 0, tensors, params); From 952ff223be836eac55932bdb9fe42364d60bbd60 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:42:17 -0700 Subject: [PATCH 16/32] deepseek41 : account Engram selection staging Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-dsv41-admission.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp index 10fd37df8825..7df73b8ca201 100644 --- a/src/llama-dsv41-admission.cpp +++ b/src/llama-dsv41-admission.cpp @@ -185,7 +185,11 @@ uint64_t llama_dsv41_engram_staging_bytes(uint32_t n_ubatch) { checked_mul(n_ubatch, LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, "Engram decoded rows"), sizeof(float), "Engram decoded rows"); - return checked_add(checked_add(ids, decoded, "Engram staging"), n_ubatch, "Engram staging"); + const uint64_t select = checked_mul(n_ubatch, sizeof(int32_t), "Engram text selection"); + return checked_add( + checked_add(checked_add(ids, decoded, "Engram staging"), n_ubatch, "Engram staging"), + select, + "Engram staging"); } uint64_t llama_dsv41_output_bytes(uint32_t n_vocab, uint32_t n_ubatch) { From ea8f88db93c9fab54bf9495fb60abac6b00d4d73 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:53:08 -0700 Subject: [PATCH 17/32] deepseek41 : close admission envelope gaps Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- common/common.cpp | 13 ++++- docs/strix-memory-watchdog.md | 6 ++- include/llama.h | 3 ++ src/llama-dsv41-admission.cpp | 73 ++++++++++++++++++++++------- src/llama-dsv41-admission.h | 16 ++++++- src/llama-model.cpp | 3 ++ src/models/deepseek41.cpp | 49 +++++++++++++++++-- tests/test-deepseek41-admission.cpp | 62 ++++++++++++++++++++++-- 8 files changed, 195 insertions(+), 30 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index bf3852b0d814..b2c159b5e5bf 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1704,8 +1704,19 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.dsv41_memory_hard_bytes = (uint64_t) params.dsv41_memory_hard_mib << 20; mparams.dsv41_memory_safety_margin_bytes = (uint64_t) params.dsv41_memory_safety_margin_mib << 20; mparams.dsv41_admission_context = params.n_ctx == 0 ? 32768 : params.n_ctx; + mparams.dsv41_admission_batch = std::max(params.n_batch, 1); mparams.dsv41_admission_sequences = params.n_parallel; - mparams.dsv41_admission_ubatch = params.n_ubatch; + mparams.dsv41_admission_ubatch = std::min( + mparams.dsv41_admission_batch, + static_cast(std::max(params.n_ubatch, 1))); + mparams.dsv41_admission_outputs = params.n_outputs_max <= 0 ? + mparams.dsv41_admission_batch : + std::min(params.n_outputs_max, mparams.dsv41_admission_batch); + mparams.dsv41_admission_outputs = std::max( + mparams.dsv41_admission_outputs, std::max(params.n_parallel, 1)); + mparams.dsv41_admission_outputs_per_seq = params.n_outputs_max_per_seq == 0 ? + mparams.dsv41_admission_outputs : + std::min(std::max(params.n_outputs_max_per_seq, 1), mparams.dsv41_admission_outputs); mparams.dsv41_procfs_root = params.dsv41_procfs_root.c_str(); if (params.kv_overrides.empty()) { diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index d02a3a7d5426..7777589c8106 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -6,13 +6,15 @@ ./scripts/strix_memory_watchdog.py -- \ ./build/bin/llama-server \ -m /mnt/models/deepseek-v41/DeepSeek-V4.1-Flash-Q2.gguf \ - -c 32768 -b 2048 -ub 2048 -ngl 99 + -c 32768 -b 2048 -ub 32 -ngl 99 ``` -DeepSeek V4.1 also runs an in-process admission check before expert-cache or model backend allocation. The default model parameters read `/proc/meminfo` and `/proc/swaps`, reject any configured swap entry, account unified host/GPU memory once, and auto-fit complete expert slots under 116 GiB total projected host use. The external watchdog is still required for guarded validation because it monitors host-wide use after startup and controls the complete process group. +DeepSeek V4.1 also runs an in-process admission check before expert-cache or model backend allocation. The default model parameters read `/proc/meminfo` and `/proc/swaps`, reject any configured swap entry, account unified host/GPU memory once, and auto-fit complete expert slots under 116 GiB total projected host use. Admission fails closed unless every selected accelerator reports `GGML_BACKEND_DEVICE_TYPE_IGPU`; CPU-only, discrete GPU, RPC, and tensor-parallel meta-device configurations are not treated as one procfs-accounted pool. The external watchdog is still required for guarded validation because it monitors host-wide use after startup and controls the complete process group. Use `--dsv41-procfs-root`, `--dsv41-memory-soft-mib`, `--dsv41-memory-watchdog-mib`, `--dsv41-memory-hard-mib`, and `--dsv41-memory-safety-margin-mib` only when reproducing admission tests or applying a more conservative host policy. `--expert-cache-slots` and `--expert-cache-mib` are optional caps; zero auto-fits. If both cache options are set, their capacity must describe the same number of complete published tensor slots. +The current expert runtime remaps the unique routed-expert union for one ubatch. Admission therefore requires `min(384, 6 * ubatch)` resident slots instead of only six top-k slots. For example, a 224-slot cache admits at most ubatch 37. Admission reports both the required slot count and the admitted ubatch capacity; it fails rather than lowering an explicit ubatch. DeepSeek V4.1 embedding extraction is rejected because those optional output buffers are not part of the bounded generation profile. + Admission accepts context checkpoints 32768, 65536, 98304, and 131072. It never lowers an explicit context request. A request that does not fit reports current use, fixed tensor bytes, state bytes, graph workspace, Engram and expert staging, output bytes, selected cache slots and bytes, safety margin, all thresholds, and the rejecting category. The wrapper performs these checks and actions: diff --git a/include/llama.h b/include/llama.h index 519a57f66911..fbbfcdbb9938 100644 --- a/include/llama.h +++ b/include/llama.h @@ -358,8 +358,11 @@ extern "C" { uint64_t dsv41_memory_hard_bytes; uint64_t dsv41_memory_safety_margin_bytes; uint32_t dsv41_admission_context; + uint32_t dsv41_admission_batch; uint32_t dsv41_admission_sequences; uint32_t dsv41_admission_ubatch; + uint32_t dsv41_admission_outputs; + uint32_t dsv41_admission_outputs_per_seq; const char * dsv41_procfs_root; // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp index 7df73b8ca201..6e972be87224 100644 --- a/src/llama-dsv41-admission.cpp +++ b/src/llama-dsv41-admission.cpp @@ -192,20 +192,40 @@ uint64_t llama_dsv41_engram_staging_bytes(uint32_t n_ubatch) { "Engram staging"); } -uint64_t llama_dsv41_output_bytes(uint32_t n_vocab, uint32_t n_ubatch) { - if (n_vocab == 0 || n_ubatch == 0) { +uint64_t llama_dsv41_output_bytes( + uint32_t n_vocab, + uint32_t n_batch, + uint32_t n_outputs_max) { + if (n_vocab == 0 || n_batch == 0 || n_outputs_max == 0 || n_outputs_max > n_batch) { throw std::runtime_error("DeepSeek V4.1 output accounting dimensions must be non-zero"); } const uint64_t floats = checked_mul( - checked_mul(n_vocab, n_ubatch, "output floats"), - 2*sizeof(float), + checked_mul(checked_mul(n_vocab, 3, "output floats"), n_outputs_max, "output floats"), + sizeof(float), "output floats"); const uint64_t token_rows = checked_add(n_vocab, 1, "output tokens"); const uint64_t tokens = checked_mul( - checked_mul(token_rows, n_ubatch, "output tokens"), + checked_mul(token_rows, n_outputs_max, "output tokens"), sizeof(int32_t), "output tokens"); - return checked_add(floats, tokens, "outputs"); + const uint64_t output_ids = checked_mul(n_batch, sizeof(int32_t), "output IDs"); + const uint64_t sampling_counts = checked_mul( + checked_mul(n_outputs_max, 3, "sampling counts"), + sizeof(size_t), + "sampling counts"); + return checked_add( + checked_add(floats, tokens, "outputs"), + checked_add(output_ids, sampling_counts, "outputs"), + "outputs"); +} + +bool llama_dsv41_has_unified_topology(const std::vector & device_types) { + return !device_types.empty() && std::all_of( + device_types.begin(), + device_types.end(), + [](enum ggml_backend_dev_type type) { + return type == GGML_BACKEND_DEVICE_TYPE_IGPU; + }); } llama_dsv41_admission_result llama_dsv41_admit( @@ -224,8 +244,11 @@ llama_dsv41_admission_result llama_dsv41_admit( result.safety_margin_bytes = params.safety_margin_bytes; result.device_reported_bytes_ignored = params.device_reported_bytes; result.n_ctx = params.n_ctx; + result.n_batch = params.n_batch; result.n_seq = params.n_seq; result.n_ubatch = params.n_ubatch; + result.n_outputs_max = params.n_outputs_max; + result.n_outputs_max_per_seq = params.n_outputs_max_per_seq; if (host.total == 0 || host.available > host.total || host.used != host.total - host.available) { reject("host", result, "host memory snapshot is invalid"); @@ -256,6 +279,12 @@ llama_dsv41_admission_result llama_dsv41_admit( if (params.n_seq != 1) { reject("context", result, "bounded DeepSeek V4.1 admission currently requires one sequence"); } + if (params.n_batch == 0 || params.n_ubatch == 0 || params.n_ubatch > params.n_batch || + params.n_outputs_max == 0 || params.n_outputs_max > params.n_batch || + params.n_outputs_max_per_seq == 0 || + params.n_outputs_max_per_seq > params.n_outputs_max) { + reject("context", result, "batch, ubatch, and output limits are invalid"); + } if (params.n_expert_used == 0 || params.n_expert_used > LLAMA_DSV41_N_EXPERT) { reject("cache", result, "expert top-k is invalid"); } @@ -288,6 +317,9 @@ llama_dsv41_admission_result llama_dsv41_admit( params.configured_cache_bytes > max_cache_bytes) { reject("cache", result, "configured cache exceeds the published expert count"); } + result.required_expert_slots = static_cast(std::min( + LLAMA_DSV41_N_EXPERT, + checked_mul(params.n_expert_used, params.n_ubatch, "required expert slots"))); const uint64_t bytes_slots = params.configured_cache_bytes == 0 ? LLAMA_DSV41_N_EXPERT : params.configured_cache_bytes/result.expert_slot_bytes; if (params.configured_cache_slots != 0 && params.configured_cache_bytes != 0 && @@ -301,10 +333,6 @@ llama_dsv41_admission_result llama_dsv41_admit( if (params.configured_cache_bytes != 0) { slot_cap = std::min(slot_cap, bytes_slots); } - if (slot_cap < params.n_expert_used) { - reject("cache", result, "configured cache is smaller than expert top-k"); - } - const auto state = llama_dsv41_account_memory( params.n_ctx, params.n_seq, @@ -315,7 +343,10 @@ llama_dsv41_admission_result llama_dsv41_admit( result.state_bytes = state.total(); result.graph_workspace_bytes = llama_dsv41_estimate_graph_workspace(params.n_ctx, params.n_ubatch); result.engram_staging_bytes = llama_dsv41_engram_staging_bytes(params.n_ubatch); - result.output_bytes = llama_dsv41_output_bytes(params.n_vocab, params.n_ubatch); + result.output_bytes = llama_dsv41_output_bytes( + params.n_vocab, + params.n_batch, + params.n_outputs_max); result.fixed_bytes = result.host_used; result.fixed_bytes = checked_add(result.fixed_bytes, result.dense_tensor_bytes, "fixed bytes"); @@ -341,18 +372,17 @@ llama_dsv41_admission_result llama_dsv41_admit( result.expert_slot_bytes, result.expert_staging_slot_bytes, "expert slot and staging"); const uint64_t fit_slots = (result.soft_bytes - result.fixed_bytes)/bytes_per_slot; const uint64_t selected = std::min(slot_cap, fit_slots); - if (selected < params.n_expert_used) { - result.projected_bytes = result.fixed_bytes; - reject("cache", result, "remaining budget cannot hold the minimum expert top-k"); - } - result.expert_slots = static_cast(selected); + result.expert_ubatch_capacity = result.expert_slots/params.n_expert_used; result.expert_cache_bytes = checked_mul(result.expert_slot_bytes, selected, "expert cache"); result.expert_staging_bytes = checked_mul(result.expert_staging_slot_bytes, selected, "expert staging"); result.projected_bytes = checked_add( checked_add(result.fixed_bytes, result.expert_cache_bytes, "projected bytes"), result.expert_staging_bytes, "projected bytes"); + if (selected < result.required_expert_slots) { + reject("cache", result, "selected cache cannot hold the requested ubatch worst-case routed expert union"); + } if (result.projected_bytes > result.soft_bytes) { reject("soft", result, "projected startup exceeds the soft limit"); } @@ -368,15 +398,20 @@ llama_dsv41_admission_result llama_dsv41_admit( std::string llama_dsv41_admission_result::describe() const { return format( - "DeepSeek V4.1 memory admission: category=%s, context=%u, sequences=%u, ubatch=%u, " + "DeepSeek V4.1 memory admission: category=%s, context=%u, batch=%u, sequences=%u, ubatch=%u, " + "outputs=%u, outputs_per_seq=%u, " "host_total=%llu, host_available=%llu, current=%llu, fixed=%llu, " "dense=%llu, state=%llu, workspace=%llu, engram_staging=%llu, expert_slots=%u, " - "expert_cache=%llu, expert_staging=%llu, outputs=%llu, safety_margin=%llu, projected=%llu, " + "required_expert_slots=%u, expert_ubatch_capacity=%u, expert_cache=%llu, expert_staging=%llu, output_bytes=%llu, " + "safety_margin=%llu, projected=%llu, " "soft=%llu, watchdog=%llu, hard=%llu, device_reported_ignored=%llu", category.c_str(), n_ctx, + n_batch, n_seq, n_ubatch, + n_outputs_max, + n_outputs_max_per_seq, (unsigned long long) host_total, (unsigned long long) host_available, (unsigned long long) host_used, @@ -386,6 +421,8 @@ std::string llama_dsv41_admission_result::describe() const { (unsigned long long) graph_workspace_bytes, (unsigned long long) engram_staging_bytes, expert_slots, + required_expert_slots, + expert_ubatch_capacity, (unsigned long long) expert_cache_bytes, (unsigned long long) expert_staging_bytes, (unsigned long long) output_bytes, diff --git a/src/llama-dsv41-admission.h b/src/llama-dsv41-admission.h index a278079f64e0..1742b105285d 100644 --- a/src/llama-dsv41-admission.h +++ b/src/llama-dsv41-admission.h @@ -2,6 +2,8 @@ #include "llama-expert-store.h" +#include "ggml-backend.h" + #include #include #include @@ -30,8 +32,11 @@ struct llama_dsv41_admission_params { uint64_t device_reported_bytes = 0; uint32_t configured_cache_slots = 0; uint32_t n_ctx = LLAMA_DSV41_ADMISSION_CONTEXT; + uint32_t n_batch = 2048; uint32_t n_seq = 1; uint32_t n_ubatch = 2048; + uint32_t n_outputs_max = 2048; + uint32_t n_outputs_max_per_seq = 2048; uint32_t n_vocab = 0; uint32_t n_expert_used = 0; uint32_t kv_element_size = 2; @@ -61,9 +66,14 @@ struct llama_dsv41_admission_result { uint64_t expert_slot_bytes = 0; uint64_t expert_staging_slot_bytes = 0; uint32_t expert_slots = 0; + uint32_t required_expert_slots = 0; + uint32_t expert_ubatch_capacity = 0; uint32_t n_ctx = 0; + uint32_t n_batch = 0; uint32_t n_seq = 0; uint32_t n_ubatch = 0; + uint32_t n_outputs_max = 0; + uint32_t n_outputs_max_per_seq = 0; std::string category; std::string describe() const; @@ -73,7 +83,11 @@ llama_dsv41_host_memory llama_dsv41_read_host_memory(const std::string & procfs_ uint64_t llama_dsv41_estimate_graph_workspace(uint32_t n_ctx, uint32_t n_ubatch); uint64_t llama_dsv41_engram_staging_bytes(uint32_t n_ubatch); -uint64_t llama_dsv41_output_bytes(uint32_t n_vocab, uint32_t n_ubatch); +uint64_t llama_dsv41_output_bytes( + uint32_t n_vocab, + uint32_t n_batch, + uint32_t n_outputs_max); +bool llama_dsv41_has_unified_topology(const std::vector & device_types); llama_dsv41_admission_result llama_dsv41_admit( const llama_dsv41_host_memory & host, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 77d300458133..1cfe70b89ec6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2841,8 +2841,11 @@ llama_model_params llama_model_default_params() { /*.dsv41_memory_hard_bytes =*/ 120ULL << 30, /*.dsv41_memory_safety_margin_bytes =*/ 2ULL << 30, /*.dsv41_admission_context =*/ 32768, + /*.dsv41_admission_batch =*/ 2048, /*.dsv41_admission_sequences =*/ 1, /*.dsv41_admission_ubatch =*/ 2048, + /*.dsv41_admission_outputs =*/ 2048, + /*.dsv41_admission_outputs_per_seq =*/ 2048, /*.dsv41_procfs_root =*/ "/proc", /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 88399009112c..9d2ed52768f3 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -284,12 +284,30 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { admission_params.configured_cache_slots = std::max(params.expert_cache_slots, 0); admission_params.n_ctx = params.dsv41_admission_context == 0 ? LLAMA_DSV41_ADMISSION_CONTEXT : params.dsv41_admission_context; + admission_params.n_batch = params.dsv41_admission_batch == 0 ? 2048 : params.dsv41_admission_batch; admission_params.n_seq = params.dsv41_admission_sequences == 0 ? 1 : params.dsv41_admission_sequences; - admission_params.n_ubatch = params.dsv41_admission_ubatch == 0 ? 2048 : params.dsv41_admission_ubatch; + admission_params.n_ubatch = std::min( + admission_params.n_batch, + params.dsv41_admission_ubatch == 0 ? + admission_params.n_batch : params.dsv41_admission_ubatch); + admission_params.n_outputs_max = std::min( + admission_params.n_batch, + params.dsv41_admission_outputs == 0 ? + admission_params.n_batch : params.dsv41_admission_outputs); + admission_params.n_outputs_max = std::max(admission_params.n_outputs_max, admission_params.n_seq); + admission_params.n_outputs_max_per_seq = std::min( + admission_params.n_outputs_max, + params.dsv41_admission_outputs_per_seq == 0 ? + admission_params.n_outputs_max : params.dsv41_admission_outputs_per_seq); admission_params.n_vocab = n_vocab; admission_params.n_expert_used = n_expert_used; admission_params.direct_io = true; - admission_params.unified_memory = true; + std::vector device_types; + device_types.reserve(devices.size()); + for (const auto & device : devices) { + device_types.push_back(ggml_backend_dev_type(device.dev)); + } + admission_params.unified_memory = llama_dsv41_has_unified_topology(device_types); const std::string procfs_root = params.dsv41_procfs_root == nullptr ? "/proc" : params.dsv41_procfs_root; admission = std::make_shared(); @@ -362,17 +380,38 @@ void llama_model_deepseek41::validate_context_params(const llama_cparams & cpara if (!admission) { throw std::runtime_error("DeepSeek V4.1 context has no host-memory admission result"); } + const uint32_t n_outputs_max = std::min(cparams.n_outputs_max, cparams.n_batch); + const uint32_t output_rows = std::max(n_outputs_max, cparams.n_seq_max); + const uint32_t n_outputs_max_per_seq = std::min(cparams.n_outputs_max_per_seq, output_rows); + const bool has_layer_embeddings = std::any_of( + cparams.embeddings_layer_inp.begin(), + cparams.embeddings_layer_inp.end(), + [](bool enabled) { return enabled; }); if (cparams.n_ctx > admission->result.n_ctx || + cparams.n_batch > admission->result.n_batch || cparams.n_seq_max > admission->result.n_seq || - cparams.n_ubatch > admission->result.n_ubatch) { + cparams.n_ubatch > admission->result.n_ubatch || + output_rows > admission->result.n_outputs_max || + n_outputs_max_per_seq > admission->result.n_outputs_max_per_seq || + cparams.embeddings || + cparams.embeddings_nextn || + has_layer_embeddings) { llama_dsv41_admission_result failure = admission->result; failure.category = "context"; throw std::runtime_error(format( - "%s, requested_context=%u, requested_sequences=%u, requested_ubatch=%u", + "%s, requested_context=%u, requested_batch=%u, requested_sequences=%u, requested_ubatch=%u, " + "requested_outputs=%u, requested_outputs_per_seq=%u, embeddings=%s, embeddings_nextn=%s, " + "layer_embeddings=%s", failure.describe().c_str(), cparams.n_ctx, + cparams.n_batch, cparams.n_seq_max, - cparams.n_ubatch)); + cparams.n_ubatch, + output_rows, + n_outputs_max_per_seq, + cparams.embeddings ? "true" : "false", + cparams.embeddings_nextn ? "true" : "false", + has_layer_embeddings ? "true" : "false")); } } diff --git a/tests/test-deepseek41-admission.cpp b/tests/test-deepseek41-admission.cpp index de769e7f2592..7fb9d743e055 100644 --- a/tests/test-deepseek41-admission.cpp +++ b/tests/test-deepseek41-admission.cpp @@ -89,6 +89,7 @@ llama_dsv41_host_memory host_with_used(uint64_t used) { llama_dsv41_admission_params base_params() { llama_dsv41_admission_params params; + params.n_ubatch = 32; params.n_vocab = LLAMA_DSV41_N_VOCAB; params.n_expert_used = LLAMA_DSV41_N_EXPERT_USED; return params; @@ -160,6 +161,7 @@ void test_published_slot_fit() { void test_configured_cache() { const auto tensors = published_tensors(); auto params = base_params(); + params.n_ubatch = 2; params.configured_cache_slots = 12; params.configured_cache_bytes = 12*398131200ULL + 1024; const auto result = llama_dsv41_admit(host_with_used(0), 0, tensors, params); @@ -175,7 +177,7 @@ void test_configured_cache() { params.configured_cache_bytes = 0; REQUIRE(thrown([&]() { llama_dsv41_admit(host_with_used(0), 0, tensors, params); - }).find("top-k") != std::string::npos); + }).find("routed expert union") != std::string::npos); params.configured_cache_slots = LLAMA_DSV41_N_EXPERT + 1; REQUIRE(thrown([&]() { @@ -186,6 +188,7 @@ void test_configured_cache() { void test_threshold_boundaries() { const auto tensors = published_tensors(); auto params = base_params(); + params.n_ubatch = 1; params.configured_cache_slots = LLAMA_DSV41_N_EXPERT_USED; params.configured_cache_bytes = params.configured_cache_slots*398131200ULL; params.safety_margin_bytes = 1; @@ -260,8 +263,10 @@ void test_diagnostics_and_guards() { const std::string diagnostic = result.describe(); for (const char * field : { "category=", "current=", "fixed=", "dense=", "state=", "workspace=", - "host_total=", "host_available=", "expert_slots=", "expert_cache=", "expert_staging=", - "soft=", "watchdog=", "hard=" }) { + "host_total=", "host_available=", "batch=", "outputs=", "outputs_per_seq=", + "expert_slots=", "required_expert_slots=", "expert_ubatch_capacity=", + "expert_cache=", "expert_staging=", + "output_bytes=", "soft=", "watchdog=", "hard=" }) { REQUIRE(diagnostic.find(field) != std::string::npos); } @@ -277,6 +282,55 @@ void test_diagnostics_and_guards() { }).find("physical host memory") != std::string::npos); } +void test_expert_union_and_outputs() { + const auto tensors = published_tensors(); + auto params = base_params(); + params.n_ubatch = 2; + params.configured_cache_slots = 11; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("worst-case routed expert union") != std::string::npos); + + params.configured_cache_slots = 12; + const auto result = llama_dsv41_admit(host_with_used(0), 0, tensors, params); + REQUIRE(result.required_expert_slots == 12); + REQUIRE(result.expert_slots == 12); + + params = base_params(); + params.n_ubatch = 37; + params.configured_cache_slots = 224; + const auto bounded = llama_dsv41_admit(host_with_used(0), 0, tensors, params); + REQUIRE(bounded.required_expert_slots == 222); + REQUIRE(bounded.expert_slots == 224); + REQUIRE(bounded.expert_ubatch_capacity == 37); + + params.n_ubatch = 38; + REQUIRE(thrown([&]() { + llama_dsv41_admit(host_with_used(0), 0, tensors, params); + }).find("worst-case routed expert union") != std::string::npos); + + const uint64_t expected = + 3*100ULL*10*sizeof(float) + + (100ULL + 1)*10*sizeof(int32_t) + + 16*sizeof(int32_t) + + 3*10*sizeof(size_t); + REQUIRE(llama_dsv41_output_bytes(100, 16, 10) == expected); +} + +void test_unified_topology() { + REQUIRE(!llama_dsv41_has_unified_topology({})); + REQUIRE(!llama_dsv41_has_unified_topology({ GGML_BACKEND_DEVICE_TYPE_CPU })); + REQUIRE(!llama_dsv41_has_unified_topology({ GGML_BACKEND_DEVICE_TYPE_GPU })); + REQUIRE(!llama_dsv41_has_unified_topology({ GGML_BACKEND_DEVICE_TYPE_META })); + REQUIRE(llama_dsv41_has_unified_topology({ GGML_BACKEND_DEVICE_TYPE_IGPU })); + REQUIRE(llama_dsv41_has_unified_topology({ + GGML_BACKEND_DEVICE_TYPE_IGPU, + GGML_BACKEND_DEVICE_TYPE_IGPU })); + REQUIRE(!llama_dsv41_has_unified_topology({ + GGML_BACKEND_DEVICE_TYPE_IGPU, + GGML_BACKEND_DEVICE_TYPE_GPU })); +} + } int main() { @@ -287,5 +341,7 @@ int main() { test_threshold_boundaries(); test_context_progression(); test_diagnostics_and_guards(); + test_expert_union_and_outputs(); + test_unified_topology(); return 0; } From e8bcaf16c98456ddeceee4e3ecdcb791a5c3268f Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:54:11 -0700 Subject: [PATCH 18/32] deepseek41 : keep output admission immutable Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-context.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 67ad5c858f66..0ab55ab6dd0b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1187,6 +1187,9 @@ void llama_context::set_eval_callback(ggml_backend_sched_eval_callback cb_eval, void llama_context::set_embeddings(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); + if (value && model.arch == LLM_ARCH_DEEPSEEK41) { + throw std::runtime_error("DeepSeek V4.1 bounded admission does not support embedding outputs"); + } cparams.embeddings = value; // TODO: not sure yet if we want to reserve here @@ -1196,6 +1199,9 @@ void llama_context::set_embeddings(bool value) { void llama_context::set_embeddings_nextn(bool value, bool masked) { LLAMA_LOG_DEBUG("%s: value = %d, masked = %d\n", __func__, value, masked); + if (value && model.arch == LLM_ARCH_DEEPSEEK41) { + throw std::runtime_error("DeepSeek V4.1 bounded admission does not support next-token embedding outputs"); + } cparams.embeddings_nextn = value; cparams.embeddings_nextn_masked = masked; } @@ -1205,6 +1211,9 @@ void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) { GGML_ASSERT(lid <= model.hparams.n_layer()); + if (enable && model.arch == LLM_ARCH_DEEPSEEK41) { + throw std::runtime_error("DeepSeek V4.1 bounded admission does not support layer embedding outputs"); + } cparams.embeddings_layer_inp[lid] = enable; // note: without this reserve, the draft acceptance drops to zero. not sure why - this is unexpected From 8b1ee951e64b9db3a887f72db24c31a279caf998 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 20:55:07 -0700 Subject: [PATCH 19/32] deepseek41 : report ignored device memory Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/models/deepseek41.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 9d2ed52768f3..16da404237d8 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -306,6 +306,12 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { device_types.reserve(devices.size()); for (const auto & device : devices) { device_types.push_back(ggml_backend_dev_type(device.dev)); + ggml_backend_dev_props properties; + ggml_backend_dev_get_props(device.dev, &properties); + if (properties.memory_total > UINT64_MAX - admission_params.device_reported_bytes) { + throw std::runtime_error("DeepSeek V4.1 device-reported memory byte count overflow"); + } + admission_params.device_reported_bytes += properties.memory_total; } admission_params.unified_memory = llama_dsv41_has_unified_topology(device_types); From 04199520d8d40b3ea996d082db7be19c07621d2c Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 21:01:40 -0700 Subject: [PATCH 20/32] deepseek41 : reject unadmitted embedding outputs Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-context.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0ab55ab6dd0b..1fc643df9ec0 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1188,7 +1188,8 @@ void llama_context::set_embeddings(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); if (value && model.arch == LLM_ARCH_DEEPSEEK41) { - throw std::runtime_error("DeepSeek V4.1 bounded admission does not support embedding outputs"); + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 bounded admission does not support embedding outputs\n", __func__); + return; } cparams.embeddings = value; @@ -1200,7 +1201,8 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) { LLAMA_LOG_DEBUG("%s: value = %d, masked = %d\n", __func__, value, masked); if (value && model.arch == LLM_ARCH_DEEPSEEK41) { - throw std::runtime_error("DeepSeek V4.1 bounded admission does not support next-token embedding outputs"); + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 bounded admission does not support next-token embedding outputs\n", __func__); + return; } cparams.embeddings_nextn = value; cparams.embeddings_nextn_masked = masked; @@ -1212,7 +1214,8 @@ void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) { GGML_ASSERT(lid <= model.hparams.n_layer()); if (enable && model.arch == LLM_ARCH_DEEPSEEK41) { - throw std::runtime_error("DeepSeek V4.1 bounded admission does not support layer embedding outputs"); + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 bounded admission does not support layer embedding outputs\n", __func__); + return; } cparams.embeddings_layer_inp[lid] = enable; From c5d1c4511348bc97975b39f74e6a84dcf4603269 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 21:30:16 -0700 Subject: [PATCH 21/32] deepseek41 : close final admission review gaps Forward SIGHUP through the external watchdog, make unsupported dynamic embedding requests fail the next operation explicitly, and use the bounded ubatch for default model admission. Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 4 ++-- scripts/strix_memory_watchdog.py | 2 +- src/llama-context.cpp | 26 +++++++++++++++++++++++--- src/llama-context.h | 4 ++++ src/llama-model.cpp | 2 +- src/models/deepseek41.cpp | 2 +- tests/test_strix_memory_watchdog.py | 7 ++++++- 7 files changed, 38 insertions(+), 9 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 7777589c8106..84aabcaa9186 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -9,7 +9,7 @@ -c 32768 -b 2048 -ub 32 -ngl 99 ``` -DeepSeek V4.1 also runs an in-process admission check before expert-cache or model backend allocation. The default model parameters read `/proc/meminfo` and `/proc/swaps`, reject any configured swap entry, account unified host/GPU memory once, and auto-fit complete expert slots under 116 GiB total projected host use. Admission fails closed unless every selected accelerator reports `GGML_BACKEND_DEVICE_TYPE_IGPU`; CPU-only, discrete GPU, RPC, and tensor-parallel meta-device configurations are not treated as one procfs-accounted pool. The external watchdog is still required for guarded validation because it monitors host-wide use after startup and controls the complete process group. +DeepSeek V4.1 also runs an in-process admission check before expert-cache or model backend allocation. The default model parameters read `/proc/meminfo` and `/proc/swaps`, reject any configured swap entry, measure the full-graph state through its no-allocation memory implementation, account unified host/GPU memory once, and auto-fit complete expert slots under 116 GiB total projected host use. The context checks the measured scheduler workspace against the admitted conservative workspace envelope before inference. Admission fails closed unless every selected accelerator reports `GGML_BACKEND_DEVICE_TYPE_IGPU`; CPU-only, discrete GPU, RPC, and tensor-parallel meta-device configurations are not treated as one procfs-accounted pool. The external watchdog is still required for guarded validation because it monitors host-wide use after startup and controls the complete process group. Use `--dsv41-procfs-root`, `--dsv41-memory-soft-mib`, `--dsv41-memory-watchdog-mib`, `--dsv41-memory-hard-mib`, and `--dsv41-memory-safety-margin-mib` only when reproducing admission tests or applying a more conservative host policy. `--expert-cache-slots` and `--expert-cache-mib` are optional caps; zero auto-fits. If both cache options are set, their capacity must describe the same number of complete published tensor slots. @@ -25,7 +25,7 @@ The wrapper performs these checks and actions: - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. - It reports `grace_timeout` if any descendant requires `SIGKILL` after the soft-threshold grace period, even when the direct child exited earlier. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. -- It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. +- It forwards wrapper `SIGHUP`, `SIGINT`, or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. - It checks the process group after the direct child exits and cleans up remaining descendants before returning the child's classification. - It applies the same bounded process-group cleanup if an unexpected post-launch error occurs. - It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index b91b6f0b44b0..d66d1536107d 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -36,7 +36,7 @@ MEMINFO_VALUE_RE = re.compile(r"([0-9]+) kB") SWAPS_HEADER = ["Filename", "Type", "Size", "Used", "Priority"] -PARENT_SIGNALS = (signal.SIGINT, signal.SIGTERM) +PARENT_SIGNALS = (signal.SIGHUP, signal.SIGINT, signal.SIGTERM) class ProcfsError(RuntimeError): diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 985138483b1d..0ac3b80f34b8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1201,7 +1201,8 @@ void llama_context::set_embeddings(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); if (value && model.arch == LLM_ARCH_DEEPSEEK41) { - LLAMA_LOG_ERROR("%s: DeepSeek V4.1 bounded admission does not support embedding outputs\n", __func__); + pending_config_error = "DeepSeek V4.1 bounded admission does not support embedding outputs"; + LLAMA_LOG_ERROR("%s: %s\n", __func__, pending_config_error.c_str()); return; } cparams.embeddings = value; @@ -1214,7 +1215,8 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) { LLAMA_LOG_DEBUG("%s: value = %d, masked = %d\n", __func__, value, masked); if (value && model.arch == LLM_ARCH_DEEPSEEK41) { - LLAMA_LOG_ERROR("%s: DeepSeek V4.1 bounded admission does not support next-token embedding outputs\n", __func__); + pending_config_error = "DeepSeek V4.1 bounded admission does not support next-token embedding outputs"; + LLAMA_LOG_ERROR("%s: %s\n", __func__, pending_config_error.c_str()); return; } cparams.embeddings_nextn = value; @@ -1227,7 +1229,8 @@ void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) { GGML_ASSERT(lid <= model.hparams.n_layer()); if (enable && model.arch == LLM_ARCH_DEEPSEEK41) { - LLAMA_LOG_ERROR("%s: DeepSeek V4.1 bounded admission does not support layer embedding outputs\n", __func__); + pending_config_error = "DeepSeek V4.1 bounded admission does not support layer embedding outputs"; + LLAMA_LOG_ERROR("%s: %s\n", __func__, pending_config_error.c_str()); return; } cparams.embeddings_layer_inp[lid] = enable; @@ -1240,6 +1243,15 @@ void llama_context::set_nextn_layer_offset(int32_t offset) { cparams.nextn_layer_offset = offset; } +bool llama_context::consume_pending_config_error() { + if (pending_config_error.empty()) { + return false; + } + LLAMA_LOG_ERROR("%s: %s\n", __func__, pending_config_error.c_str()); + pending_config_error.clear(); + return true; +} + void llama_context::set_causal_attn(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); @@ -1504,6 +1516,10 @@ int llama_context::encode(const llama_batch & batch_inp) { // so accept either present rather than requiring exactly one. GGML_ASSERT(batch_inp.token || batch_inp.embd); + if (consume_pending_config_error()) { + return -1; + } + if (batch_inp.n_tokens == 0) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); return -1; @@ -1742,6 +1758,10 @@ int llama_context::decode(const llama_batch & batch_inp) { // so accept either present rather than requiring exactly one. GGML_ASSERT(batch_inp.token || batch_inp.embd); + if (consume_pending_config_error()) { + return -1; + } + if (!memory) { LLAMA_LOG_DEBUG("%s: cannot decode batches with this context (calling encode() instead)\n", __func__); return encode(batch_inp); diff --git a/src/llama-context.h b/src/llama-context.h index b089e1267aaf..6d6651f4c365 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -12,6 +12,7 @@ #include "ggml-opt.h" #include +#include #include struct llama_model; @@ -255,6 +256,8 @@ struct llama_context { bool set_sampler(llama_seq_id seq_id, llama_sampler * sampler); private: + bool consume_pending_config_error(); + llm_graph_params graph_params( llm_graph_result * res, const llama_ubatch & ubatch, @@ -288,6 +291,7 @@ struct llama_context { llama_cross cross; // TODO: tmp for handling cross-attention - need something better probably llama_memory_ptr memory; + std::string pending_config_error; // decode output (2-dimensional array: [n_outputs][n_vocab]) buffer_view logits = {nullptr, 0}; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 54a3ab9611f5..8aec9d15b7ef 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2864,7 +2864,7 @@ llama_model_params llama_model_default_params() { /*.dsv41_admission_context =*/ 32768, /*.dsv41_admission_batch =*/ 2048, /*.dsv41_admission_sequences =*/ 1, - /*.dsv41_admission_ubatch =*/ 2048, + /*.dsv41_admission_ubatch =*/ 32, /*.dsv41_admission_outputs =*/ 2048, /*.dsv41_admission_outputs_per_seq =*/ 2048, /*.dsv41_admission_type_k =*/ GGML_TYPE_F16, diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 6c6a51e40f6c..57eb8cfd24c3 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -733,7 +733,7 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { admission_params.n_ubatch = std::min( admission_params.n_batch, params.dsv41_admission_ubatch == 0 ? - admission_params.n_batch : params.dsv41_admission_ubatch); + 32U : params.dsv41_admission_ubatch); admission_params.n_outputs_max = std::min( admission_params.n_batch, params.dsv41_admission_outputs == 0 ? diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index d4eea8e00972..ea808f457d8b 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -228,6 +228,7 @@ def _write_procfs_fixture(root: Path) -> None: def test_parent_signals_leave_no_child_or_grandchild(self) -> None: child_code = ( "import os,signal,sys,time;" + "signal.signal(signal.SIGHUP,signal.SIG_IGN);" "signal.signal(signal.SIGINT,signal.SIG_IGN);" "signal.signal(signal.SIGTERM,signal.SIG_IGN);" "grandchild=os.fork();" @@ -238,7 +239,11 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: "f'{os.getpid()} {grandchild}\\n');" " time.sleep(30)\n" ) - for signal_number in (signal.SIGINT, signal.SIGTERM): + for signal_number in ( + signal.SIGHUP, + signal.SIGINT, + signal.SIGTERM, + ): with self.subTest(signal=signal.Signals(signal_number).name): with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) From 8acdb46f83251b531b896658d5c913ba8edecf4a Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 21:36:56 -0700 Subject: [PATCH 22/32] scripts : forward SIGHUP through memory watchdog Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 2 +- scripts/strix_memory_watchdog.py | 2 +- tests/test_strix_memory_watchdog.py | 7 ++++++- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index ccd04f7ab667..0ba9e2e7721e 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -14,7 +14,7 @@ The wrapper performs these checks and actions: - It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. - It reports `grace_timeout` if any descendant requires `SIGKILL` after the soft-threshold grace period, even when the direct child exited earlier. - It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. -- It forwards wrapper `SIGINT` or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. +- It forwards wrapper `SIGHUP`, `SIGINT`, or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. - It checks the process group after the direct child exits and cleans up remaining descendants before returning the child's classification. - It applies the same bounded process-group cleanup if an unexpected post-launch error occurs. - It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 00a9fab78d81..a17c76dc960b 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -36,7 +36,7 @@ MEMINFO_VALUE_RE = re.compile(r"([0-9]+) kB") SWAPS_HEADER = ["Filename", "Type", "Size", "Used", "Priority"] -PARENT_SIGNALS = (signal.SIGINT, signal.SIGTERM) +PARENT_SIGNALS = (signal.SIGHUP, signal.SIGINT, signal.SIGTERM) class ProcfsError(RuntimeError): diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 898298efb9b9..eca818da0f6f 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -228,6 +228,7 @@ def _write_procfs_fixture(root: Path) -> None: def test_parent_signals_leave_no_child_or_grandchild(self) -> None: child_code = ( "import os,signal,sys,time;" + "signal.signal(signal.SIGHUP,signal.SIG_IGN);" "signal.signal(signal.SIGINT,signal.SIG_IGN);" "signal.signal(signal.SIGTERM,signal.SIG_IGN);" "grandchild=os.fork();" @@ -238,7 +239,11 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: "f'{os.getpid()} {grandchild}\\n');" " time.sleep(30)\n" ) - for signal_number in (signal.SIGINT, signal.SIGTERM): + for signal_number in ( + signal.SIGHUP, + signal.SIGINT, + signal.SIGTERM, + ): with self.subTest(signal=signal.Signals(signal_number).name): with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) From a127c1999f45ff0937cda1b808958d37f7a5a6d1 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 21:56:35 -0700 Subject: [PATCH 23/32] deepseek41 : close admission envelope blockers Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- common/arg.cpp | 1 + common/common.cpp | 16 ++++- common/common.h | 5 ++ common/fit.cpp | 12 +++- docs/strix-memory-watchdog.md | 2 +- src/llama-context.cpp | 20 +++++- src/llama-context.h | 1 + src/llama-dsv41-admission.cpp | 30 ++++++++- src/llama-dsv41-admission.h | 2 + src/llama-expert-store.h | 4 +- tests/test-arg-parser.cpp | 23 +++++++ tests/test-deepseek41-admission.cpp | 30 ++++++++- tests/test-deepseek41-runtime.cpp | 94 +++++++++++++++++++++++++++++ 13 files changed, 231 insertions(+), 9 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 6d85c1e1b2d0..b358b04329f6 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1697,6 +1697,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("physical maximum batch size (default: %d)", params.n_ubatch), [](common_params & params, int value) { params.n_ubatch = value; + params.n_ubatch_explicit = true; } ).set_env("LLAMA_ARG_UBATCH")); add_opt(common_arg( diff --git a/common/common.cpp b/common/common.cpp index e58fc9f6fad6..7322045663ea 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1337,6 +1337,11 @@ common_init_result::common_init_result(common_params & params, bool model_only) return; } + char architecture[128] = {}; + if (llama_model_meta_val_str(model, "general.architecture", architecture, sizeof(architecture)) >= 0) { + common_context_params_apply_arch_defaults(architecture, params, cparams); + } + const llama_vocab * vocab = llama_model_get_vocab(model); // load and optionally apply lora adapters @@ -1708,7 +1713,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.dsv41_admission_sequences = params.n_parallel; mparams.dsv41_admission_ubatch = std::min( mparams.dsv41_admission_batch, - static_cast(std::max(params.n_ubatch, 1))); + static_cast(params.n_ubatch_explicit ? std::max(params.n_ubatch, 1) : 32)); mparams.dsv41_admission_outputs = params.n_outputs_max <= 0 ? mparams.dsv41_admission_batch : std::min(params.n_outputs_max, mparams.dsv41_admission_batch); @@ -1743,6 +1748,15 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { return mparams; } +void common_context_params_apply_arch_defaults( + const char * architecture, + const common_params & params, + llama_context_params & cparams) { + if (architecture != nullptr && strcmp(architecture, "deepseek41") == 0 && !params.n_ubatch_explicit) { + cparams.n_ubatch = std::min(cparams.n_batch, 32); + } +} + struct llama_context_params common_context_params_to_llama(const common_params & params) { auto cparams = llama_context_default_params(); diff --git a/common/common.h b/common/common.h index 2efd7e80d707..614799dc4ce5 100644 --- a/common/common.h +++ b/common/common.h @@ -491,6 +491,7 @@ struct common_params { int32_t n_ctx = 0; // context size, 0 == context the model was trained with int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_ubatch = 512; // physical batch size for prompt processing (must be >=32 to use BLAS) + bool n_ubatch_explicit = false; int32_t n_keep = 0; // number of tokens to keep from initial prompt int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) int32_t n_parallel = 1; // number of parallel sequences to decode @@ -1001,6 +1002,10 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode struct llama_model_params common_model_params_to_llama ( common_params & params); struct llama_context_params common_context_params_to_llama(const common_params & params); +void common_context_params_apply_arch_defaults( + const char * architecture, + const common_params & params, + struct llama_context_params & cparams); // clear LoRA adapters from context, then apply new list of adapters void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora); diff --git a/common/fit.cpp b/common/fit.cpp index c601fe405ea5..472ee6a113f0 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -6,9 +6,10 @@ #include #include -#include #include +#include #include +#include #include #include @@ -62,7 +63,14 @@ static std::vector common_get_device_memory_data_impl( throw std::runtime_error("failed to load model"); } - llama_context * ctx = llama_init_from_model(model, *cparams); + llama_context_params cparams_copy = *cparams; + char architecture[128] = {}; + if (llama_model_meta_val_str(model, "general.architecture", architecture, sizeof(architecture)) >= 0 && + strcmp(architecture, "deepseek41") == 0 && + cparams_copy.n_ubatch != mparams_copy.dsv41_admission_ubatch) { + cparams_copy.n_ubatch = mparams_copy.dsv41_admission_ubatch; + } + llama_context * ctx = llama_init_from_model(model, cparams_copy); if (ctx == nullptr) { llama_model_free(model); llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 84aabcaa9186..73933a32db49 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -13,7 +13,7 @@ DeepSeek V4.1 also runs an in-process admission check before expert-cache or mod Use `--dsv41-procfs-root`, `--dsv41-memory-soft-mib`, `--dsv41-memory-watchdog-mib`, `--dsv41-memory-hard-mib`, and `--dsv41-memory-safety-margin-mib` only when reproducing admission tests or applying a more conservative host policy. `--expert-cache-slots` and `--expert-cache-mib` are optional caps; zero auto-fits. If both cache options are set, their capacity must describe the same number of complete published tensor slots. -The current expert runtime remaps the unique routed-expert union for one ubatch. Admission therefore requires `min(384, 6 * ubatch)` resident slots instead of only six top-k slots. For example, a 224-slot cache admits at most ubatch 37. Admission reports both the required slot count and the admitted ubatch capacity; it fails rather than lowering an explicit ubatch. DeepSeek V4.1 embedding extraction is rejected because those optional output buffers are not part of the bounded generation profile. +The current expert runtime remaps the unique routed-expert union for one ubatch. Admission therefore requires `min(384, 6 * ubatch)` resident slots instead of only six top-k slots. For example, a 224-slot cache admits at most ubatch 37. The common CLI and server default to ubatch 32 for DeepSeek V4.1 when `-ub` is not specified; an explicit value is preserved and must fit. Admission includes the resident staging cache, a complete worst-case replacement set, and the largest aligned direct-I/O bounce read. It reports both the required slot count and the admitted ubatch capacity and fails rather than lowering an explicit ubatch. DeepSeek V4.1 embedding extraction is rejected because those optional output buffers are not part of the bounded generation profile. Admission accepts context checkpoints 32768, 65536, 98304, and 131072. It never lowers an explicit context request. A request that does not fit reports current use, fixed tensor bytes, state bytes, graph workspace, Engram and expert staging, output bytes, selected cache slots and bytes, safety margin, all thresholds, and the rejecting category. diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0ac3b80f34b8..3f719721d74b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -33,6 +33,17 @@ static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) { throw std::runtime_error("Unsupported ctx type"); } +struct llama_runtime_context_guard { + const llama_model & model; + bool active = true; + + ~llama_runtime_context_guard() { + if (active) { + model.release_runtime_context(); + } + } +}; + struct llm_fused_op_probe { llm_fused_op op; const char * name; @@ -257,6 +268,8 @@ llama_context::llama_context( cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max); model.validate_context_params(cparams); + model.acquire_runtime_context(); + llama_runtime_context_guard runtime_context_guard { model }; // Initialize backend samplers here so they are part of the sampling graph // before the reserve passes run later in this function. This avoids a later @@ -485,14 +498,17 @@ llama_context::llama_context( } } - model.acquire_runtime_context(); + runtime_context_acquired = true; + runtime_context_guard.active = false; } llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); model.release_runtime_work(); - model.release_runtime_context(); + if (runtime_context_acquired) { + model.release_runtime_context(); + } // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation if (!model.hparams.no_alloc && !opt_ctx) { diff --git a/src/llama-context.h b/src/llama-context.h index 6d6651f4c365..272f375b0613 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -292,6 +292,7 @@ struct llama_context { llama_memory_ptr memory; std::string pending_config_error; + bool runtime_context_acquired = false; // decode output (2-dimensional array: [n_outputs][n_vocab]) buffer_view logits = {nullptr, 0}; diff --git a/src/llama-dsv41-admission.cpp b/src/llama-dsv41-admission.cpp index fd37515d3cb0..6fbf52eb110a 100644 --- a/src/llama-dsv41-admission.cpp +++ b/src/llama-dsv41-admission.cpp @@ -26,6 +26,16 @@ uint64_t checked_mul(uint64_t a, uint64_t b, const char * category) { return a*b; } +uint64_t checked_align_up(uint64_t value, uint64_t alignment, const char * category) { + if (alignment == 0) { + throw std::runtime_error(std::string("DeepSeek V4.1 memory admission invalid alignment: ") + category); + } + return checked_mul( + checked_add(value, alignment - 1, category)/alignment, + alignment, + category); +} + uint64_t parse_u64(const std::string & value, const char * field) { uint64_t result = 0; const char * begin = value.data(); @@ -300,6 +310,15 @@ llama_dsv41_admission_result llama_dsv41_admit( layer_slot_bytes[tensor.layer], tensor.nb[2], "expert slot"); result.expert_slot_bytes = checked_add( result.expert_slot_bytes, tensor.nb[2], "expert slot"); + result.direct_io_bounce_bytes = std::max( + result.direct_io_bounce_bytes, + checked_align_up( + checked_add( + tensor.nb[2], + LLAMA_EXPERT_STORE_DEFAULT_IO_ALIGNMENT - 1, + "direct I/O bounce"), + LLAMA_EXPERT_STORE_DEFAULT_IO_ALIGNMENT, + "direct I/O bounce")); } if (expert_tensors.size() != LLAMA_DSV41_N_LAYER*3) { reject("cache", result, "expected 40 gate/up/down expert tensor sets"); @@ -320,6 +339,10 @@ llama_dsv41_admission_result llama_dsv41_admit( result.required_expert_slots = static_cast(std::min( LLAMA_DSV41_N_EXPERT, checked_mul(params.n_expert_used, params.n_ubatch, "required expert slots"))); + result.expert_replacement_bytes = checked_mul( + result.required_expert_slots, + result.expert_staging_slot_bytes, + "expert replacement staging"); const uint64_t bytes_slots = params.configured_cache_bytes == 0 ? LLAMA_DSV41_N_EXPERT : params.configured_cache_bytes/result.expert_slot_bytes; if (params.configured_cache_slots != 0 && params.configured_cache_bytes != 0 && @@ -349,6 +372,8 @@ llama_dsv41_admission_result llama_dsv41_admit( result.fixed_bytes = checked_add(result.fixed_bytes, result.state_bytes, "fixed bytes"); result.fixed_bytes = checked_add(result.fixed_bytes, result.graph_workspace_bytes, "fixed bytes"); result.fixed_bytes = checked_add(result.fixed_bytes, result.engram_staging_bytes, "fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, result.expert_replacement_bytes, "fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, result.direct_io_bounce_bytes, "fixed bytes"); result.fixed_bytes = checked_add(result.fixed_bytes, result.output_bytes, "fixed bytes"); result.fixed_bytes = checked_add(result.fixed_bytes, result.safety_margin_bytes, "fixed bytes"); @@ -435,7 +460,8 @@ std::string llama_dsv41_admission_result::describe() const { "outputs=%u, outputs_per_seq=%u, " "host_total=%llu, host_available=%llu, current=%llu, fixed=%llu, " "dense=%llu, state=%llu, workspace=%llu, engram_staging=%llu, expert_slots=%u, " - "required_expert_slots=%u, expert_ubatch_capacity=%u, expert_cache=%llu, expert_staging=%llu, output_bytes=%llu, " + "required_expert_slots=%u, expert_ubatch_capacity=%u, expert_cache=%llu, expert_staging=%llu, " + "expert_replacement=%llu, direct_io_bounce=%llu, output_bytes=%llu, " "safety_margin=%llu, projected=%llu, " "soft=%llu, watchdog=%llu, hard=%llu, device_reported_ignored=%llu", category.c_str(), @@ -458,6 +484,8 @@ std::string llama_dsv41_admission_result::describe() const { expert_ubatch_capacity, (unsigned long long) expert_cache_bytes, (unsigned long long) expert_staging_bytes, + (unsigned long long) expert_replacement_bytes, + (unsigned long long) direct_io_bounce_bytes, (unsigned long long) output_bytes, (unsigned long long) safety_margin_bytes, (unsigned long long) projected_bytes, diff --git a/src/llama-dsv41-admission.h b/src/llama-dsv41-admission.h index c7e391c6d830..47f1e34a2093 100644 --- a/src/llama-dsv41-admission.h +++ b/src/llama-dsv41-admission.h @@ -53,6 +53,8 @@ struct llama_dsv41_admission_result { uint64_t graph_workspace_bytes = 0; uint64_t engram_staging_bytes = 0; uint64_t expert_staging_bytes = 0; + uint64_t expert_replacement_bytes = 0; + uint64_t direct_io_bounce_bytes = 0; uint64_t expert_cache_bytes = 0; uint64_t output_bytes = 0; uint64_t safety_margin_bytes = 0; diff --git a/src/llama-expert-store.h b/src/llama-expert-store.h index 10ce3a1a19d7..7ebbe0ce41f9 100644 --- a/src/llama-expert-store.h +++ b/src/llama-expert-store.h @@ -14,6 +14,8 @@ enum llama_expert_projection { LLAMA_EXPERT_PROJECTION_DOWN, }; +static constexpr size_t LLAMA_EXPERT_STORE_DEFAULT_IO_ALIGNMENT = 4096; + struct llama_expert_store_tensor { std::string name; std::string fname; @@ -30,7 +32,7 @@ struct llama_expert_store_tensor { struct llama_expert_store_params { size_t cache_bytes = 0; size_t cache_slots = 0; - size_t io_alignment = 4096; + size_t io_alignment = LLAMA_EXPERT_STORE_DEFAULT_IO_ALIGNMENT; bool direct_io = true; bool allow_buffered_io = false; // opt-in only; page-cache bytes are outside cache_bytes }; diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index e0907631abd8..32dd9fb3ed79 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -170,6 +170,29 @@ static void test(void) { return res; }; + { + common_params default_params; + const auto model_params = common_model_params_to_llama(default_params); + auto context_params = common_context_params_to_llama(default_params); + assert(model_params.dsv41_admission_ubatch == 32); + common_context_params_apply_arch_defaults("deepseek41", default_params, context_params); + assert(context_params.n_ubatch == 32); + + common_params explicit_params; + std::vector explicit_argv = { "binary_name", "-m", "model_file.gguf", "-ub", "37" }; + assert(common_params_parse( + explicit_argv.size(), + list_str_to_char(explicit_argv).data(), + explicit_params, + LLAMA_EXAMPLE_COMMON)); + const auto explicit_model_params = common_model_params_to_llama(explicit_params); + auto explicit_context_params = common_context_params_to_llama(explicit_params); + assert(explicit_params.n_ubatch_explicit); + assert(explicit_model_params.dsv41_admission_ubatch == 37); + common_context_params_apply_arch_defaults("deepseek41", explicit_params, explicit_context_params); + assert(explicit_context_params.n_ubatch == 37); + } + std::vector argv; printf("test-arg-parser: test invalid usage\n\n"); diff --git a/tests/test-deepseek41-admission.cpp b/tests/test-deepseek41-admission.cpp index 44eb54b7c44e..39ba6e564a4c 100644 --- a/tests/test-deepseek41-admission.cpp +++ b/tests/test-deepseek41-admission.cpp @@ -267,7 +267,7 @@ void test_diagnostics_and_guards() { "category=", "current=", "fixed=", "dense=", "state=", "workspace=", "host_total=", "host_available=", "batch=", "outputs=", "outputs_per_seq=", "expert_slots=", "required_expert_slots=", "expert_ubatch_capacity=", - "expert_cache=", "expert_staging=", + "expert_cache=", "expert_staging=", "expert_replacement=", "direct_io_bounce=", "output_bytes=", "soft=", "watchdog=", "hard=" }) { REQUIRE(diagnostic.find(field) != std::string::npos); } @@ -319,6 +319,33 @@ void test_expert_union_and_outputs() { REQUIRE(llama_dsv41_output_bytes(100, 16, 10) == expected); } +void test_expert_replacement_peak() { + auto params = base_params(); + params.n_ubatch = 32; + params.configured_cache_slots = 192; + auto result = llama_dsv41_admit(host_with_used(0), 0, published_tensors(), params); + REQUIRE(result.expert_replacement_bytes == 1911029760); + REQUIRE(result.direct_io_bounce_bytes == 3874816); + REQUIRE(result.expert_replacement_bytes + result.direct_io_bounce_bytes == 1914904576); + + params.n_ubatch = 36; + params.configured_cache_slots = 216; + result = llama_dsv41_admit(host_with_used(0), 0, published_tensors(), params); + REQUIRE(result.expert_replacement_bytes + result.direct_io_bounce_bytes == 2153783296); + + params.n_ubatch = 37; + params.configured_cache_slots = 222; + const auto baseline = llama_dsv41_admit(host_with_used(0), 0, published_tensors(), params); + REQUIRE(baseline.expert_replacement_bytes + baseline.direct_io_bounce_bytes == 2213502976); + + const uint64_t boundary_used = params.soft_bytes - baseline.projected_bytes; + const auto exact = llama_dsv41_admit(host_with_used(boundary_used), 0, published_tensors(), params); + REQUIRE(exact.projected_bytes == params.soft_bytes); + REQUIRE(!thrown([&]() { + llama_dsv41_admit(host_with_used(boundary_used + 1), 0, published_tensors(), params); + }).empty()); +} + void test_runtime_memory_validation() { auto params = base_params(); params.n_ubatch = 1; @@ -365,6 +392,7 @@ int main() { test_context_progression(); test_diagnostics_and_guards(); test_expert_union_and_outputs(); + test_expert_replacement_peak(); test_runtime_memory_validation(); test_unified_topology(); return 0; diff --git a/tests/test-deepseek41-runtime.cpp b/tests/test-deepseek41-runtime.cpp index fed172a333e8..985ff8a72009 100644 --- a/tests/test-deepseek41-runtime.cpp +++ b/tests/test-deepseek41-runtime.cpp @@ -1,5 +1,8 @@ #include "../src/llama-dsv41.h" #include "../src/llama-arch.h" +#include "../src/llama-context.h" +#include "../src/llama-graph.h" +#include "../src/llama-model.h" #include "ggml-backend.h" #include "ggml-cpu.h" @@ -12,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -581,6 +585,95 @@ static void test_graph_construction() { ggml_free(ctx); } +struct reservation_test_model final : llama_model { + mutable bool active = false; + mutable uint32_t acquisitions = 0; + mutable uint32_t releases = 0; + + reservation_test_model() : llama_model(llama_model_default_params()) { + arch = LLM_ARCH_BERT; + hparams.vocab_only = true; + hparams.n_ctx_train = 32; + hparams.causal_attn = true; + } + + void acquire_runtime_context() const override { + ++acquisitions; + if (active) { + throw std::runtime_error("duplicate runtime context"); + } + active = true; + } + + void release_runtime_context() const override { + check(active, "runtime context released without acquisition"); + active = false; + ++releases; + } + + void load_stats(llama_model_loader &) override {} + void load_hparams(llama_model_loader &) override {} + void load_vocab(llama_model_loader &) override {} + bool load_tensors(llama_model_loader &) override { return true; } + void load_arch_hparams(llama_model_loader &) override {} + void load_arch_tensors(llama_model_loader &) override {} + std::unique_ptr build_arch_graph(const llm_graph_params &) const override { + return nullptr; + } +}; + +static llama_context_params reservation_context_params() { + llama_context_params params = llama_context_default_params(); + params.n_ctx = 32; + params.n_batch = 1; + params.n_ubatch = 1; + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + return params; +} + +static void test_runtime_context_reservation() { + reservation_test_model model; + auto first = std::make_unique(model, reservation_context_params()); + check(model.active && model.acquisitions == 1 && model.releases == 0, + "first runtime context did not acquire the model reservation"); + + llama_sampler * invalid_sampler = llama_sampler_init_greedy(); + llama_sampler_seq_config sampler_config = { 0, invalid_sampler }; + llama_context_params duplicate_params = reservation_context_params(); + duplicate_params.samplers = &sampler_config; + duplicate_params.n_samplers = 1; + std::string duplicate_error; + try { + auto duplicate = std::make_unique(model, duplicate_params); + } catch (const std::runtime_error & error) { + duplicate_error = error.what(); + } + check(duplicate_error.find("duplicate runtime context") != std::string::npos, + "duplicate reservation did not reject before later constructor validation"); + check(model.active && model.acquisitions == 2 && model.releases == 0, + "duplicate reservation changed the active context"); + first.reset(); + check(!model.active && model.releases == 1, + "successful context destruction did not release the reservation"); + + reservation_test_model failed_model; + std::string construction_error; + try { + auto failed = std::make_unique(failed_model, duplicate_params); + } catch (const std::runtime_error & error) { + construction_error = error.what(); + } + check(construction_error.find("backend samplers must be of type") != std::string::npos, + "test constructor did not fail after acquiring the reservation"); + check(!failed_model.active && failed_model.acquisitions == 1 && failed_model.releases == 1, + "failed context construction did not release the reservation"); + auto recovered = std::make_unique(failed_model, reservation_context_params()); + check(failed_model.active && failed_model.acquisitions == 2, + "failed construction prevented a later context from acquiring"); + recovered.reset(); + llama_sampler_free(invalid_sampler); +} + int main() { test_hparams(); test_source_maps(); @@ -591,5 +684,6 @@ int main() { test_output_collapse(); test_graph_contract(); test_graph_construction(); + test_runtime_context_reservation(); return 0; } From 59833018814c0883f995848912cd5a52889c5303 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 22:37:06 -0700 Subject: [PATCH 24/32] scripts : publish watchdog-owned validation lease Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 43 ++ scripts/strix_memory_watchdog.py | 914 ++++++++++++++++++++++++++-- tests/test_strix_memory_watchdog.py | 475 ++++++++++++++- 3 files changed, 1385 insertions(+), 47 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 0ba9e2e7721e..5565716a0a3e 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -25,6 +25,48 @@ Use `--procfs-root` to select a different procfs mount or a test fixture. `--sof The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Signal records are written immediately after each process-group signal. Child standard input, standard output, and standard error are inherited unchanged. +## Watchdog-owned validation lease + +Use all three artifact options together when another process must prove that it is inside the active watchdog process group: + +```sh +./scripts/strix_memory_watchdog.py \ + --lease-path /run/deepseek-v41/watchdog-lease.json \ + --heartbeat-path /run/deepseek-v41/watchdog-heartbeat.json \ + --audit-path /run/deepseek-v41/watchdog-audit.jsonl \ + -- \ + python3 tools/deepseek-v41-trace/run_matrix.py +``` + +The watchdog creates the persistent audit before launch, then atomically creates the lease and heartbeat after `Popen` returns. Existing artifact paths are rejected rather than overwritten. The child receives the resolved paths through `STRIX_MEMORY_WATCHDOG_LEASE_PATH`, `STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH`, and `STRIX_MEMORY_WATCHDOG_AUDIT_PATH`. It also receives `STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS`. + +The child can run before the first atomic lease rename. A matching preflight must retry the inherited lease path for a bounded interval and fail closed if a complete valid lease does not appear. It must not accept a lease path supplied separately by the operator. + +Lease format `strix-memory-watchdog-lease`, version 1, contains: + +- `lease_id` and active/final `state` +- `watchdog_pid`, `watchdog_start_time_utc`, Linux `watchdog_start_time_ticks`, `watchdog_command_sha256`, `watchdog_script_path`, and `watchdog_script_sha256` +- exact `soft_bytes`, `emergency_bytes`, and `strict_ceiling_bytes` +- `procfs_root` +- `child_pid`, `child_process_group_id`, `command`, and `child_command_sha256` +- `heartbeat_path`, `max_heartbeat_age_seconds`, and `audit_path` +- the authoritative `final` audit record after termination + +Heartbeat format `strix-memory-watchdog-heartbeat`, version 1, binds `lease_id`, watchdog PID/start ticks, child PID/process group, sequence, state, and update timestamps. Every memory sample atomically replaces the heartbeat and includes the complete sample audit record. A final heartbeat and final lease update remain on disk with the persistent JSONL audit; the watchdog does not delete this evidence. + +A matching Linux preflight must verify all of the following: + +- The inherited lease, heartbeat, and audit paths match the paths inside the lease. +- The expected repository script path hashes to `watchdog_script_sha256`. +- `/proc//stat` start ticks and `/proc//cmdline` SHA-256 match the lease. +- The watchdog command line names the expected script. +- The current process group equals `child_process_group_id`, whose leader is `child_pid` and whose parent is `watchdog_pid`. +- The command identity is expected, the procfs root is `/proc`, and thresholds are exactly 116 GiB soft, 118 GiB emergency, and 120 GiB strict ceiling for the final run. +- The heartbeat identity matches the lease and its monotonic timestamp is not older than `max_heartbeat_age_seconds`. +- The persistent audit exists and contains watchdog JSONL records. + +These checks bind the validation process to the live canonical watchdog. A standalone heartbeat helper has a different PID, start time, command line, script hash, and process group and cannot satisfy the lease. + Exit classifications are authoritative in the final JSON record. Operational failures use these exit codes: | Exit code | Classification | @@ -35,6 +77,7 @@ Exit classifications are authoritative in the final JSON record. Operational fai | 5 | emergency threshold reached | | 6 | soft-threshold grace period expired | | 7 | process-group signaling or termination failure | +| 8 | lease, heartbeat, or persistent audit failure | | 70 | unexpected post-launch error | | 127 | command launch failure | diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index a17c76dc960b..0bb11a0622c3 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -3,10 +3,12 @@ from __future__ import annotations import argparse +import hashlib import json import math import os import re +import secrets import signal import subprocess import sys @@ -15,7 +17,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import IO, Protocol +from typing import IO, Any, Protocol GIB = 1024**3 @@ -24,6 +26,12 @@ DEFAULT_EMERGENCY_BYTES = 118 * GIB DEFAULT_GRACE_SECONDS = 30.0 DEFAULT_SAMPLE_INTERVAL_SECONDS = 1.0 +DEFAULT_HEARTBEAT_MAX_AGE_SECONDS = 5.0 + +LEASE_FORMAT = "strix-memory-watchdog-lease" +LEASE_VERSION = 1 +HEARTBEAT_FORMAT = "strix-memory-watchdog-heartbeat" +HEARTBEAT_VERSION = 1 EXIT_PROCFS_ERROR = 2 EXIT_SWAP_ACTIVE = 3 @@ -31,6 +39,7 @@ EXIT_EMERGENCY_LIMIT = 5 EXIT_GRACE_TIMEOUT = 6 EXIT_SIGNAL_ERROR = 7 +EXIT_LEASE_ERROR = 8 EXIT_INTERNAL_ERROR = 70 EXIT_LAUNCH_ERROR = 127 @@ -47,6 +56,16 @@ class ProcessGroupError(RuntimeError): pass +class ArtifactError(RuntimeError): + def __init__(self, component: str, detail: str): + self.component = component + super().__init__(detail) + + +class LeaseValidationError(RuntimeError): + pass + + class ParentSignal(RuntimeError): def __init__(self, signal_number: int): self.signal_number = signal_number @@ -80,6 +99,13 @@ class RuntimeState: peak_used_bytes: int +@dataclass(frozen=True) +class ArtifactPaths: + lease: Path + heartbeat: Path + audit: Path + + @dataclass(frozen=True) class WatchdogConfig: command: tuple[str, ...] @@ -88,8 +114,16 @@ class WatchdogConfig: emergency_bytes: int = DEFAULT_EMERGENCY_BYTES grace_seconds: float = DEFAULT_GRACE_SECONDS sample_interval_seconds: float = DEFAULT_SAMPLE_INTERVAL_SECONDS + lease_path: Path | None = None + heartbeat_path: Path | None = None + audit_path: Path | None = None + heartbeat_max_age_seconds: float = DEFAULT_HEARTBEAT_MAX_AGE_SECONDS + + @property + def lease_enabled(self) -> bool: + return self.lease_path is not None - def validate(self) -> None: + def validate(self) -> ArtifactPaths | None: if not self.command: raise ValueError("a command is required after --") if self.soft_bytes <= 0: @@ -105,6 +139,45 @@ def validate(self) -> None: or self.sample_interval_seconds <= 0 ): raise ValueError("sample interval must be greater than zero") + lease_paths = ( + self.lease_path, + self.heartbeat_path, + self.audit_path, + ) + if any(path is not None for path in lease_paths) and not all( + path is not None for path in lease_paths + ): + raise ValueError( + "lease, heartbeat, and audit paths must be specified together" + ) + if self.lease_enabled: + assert self.lease_path is not None + assert self.heartbeat_path is not None + assert self.audit_path is not None + try: + paths = ArtifactPaths( + self.lease_path.expanduser().resolve(), + self.heartbeat_path.expanduser().resolve(), + self.audit_path.expanduser().resolve(), + ) + except (OSError, RuntimeError) as exc: + raise ValueError( + f"cannot resolve watchdog artifact path: {exc}" + ) from exc + if len({paths.lease, paths.heartbeat, paths.audit}) != 3: + raise ValueError( + "lease, heartbeat, and audit paths must be distinct" + ) + if ( + not math.isfinite(self.heartbeat_max_age_seconds) + or self.heartbeat_max_age_seconds + <= self.sample_interval_seconds + ): + raise ValueError( + "heartbeat max age must be greater than sample interval" + ) + return paths + return None class ProcfsReader: @@ -171,6 +244,532 @@ def _parse_swaps(content: str) -> tuple[str, ...]: return tuple(entries) +def _timestamp_utc( + wall_clock: Callable[[], datetime] | None = None, +) -> str: + timestamp = (wall_clock or ( + lambda: datetime.now(timezone.utc) + ))().astimezone(timezone.utc) + return timestamp.isoformat(timespec="milliseconds").replace( + "+00:00", "Z" + ) + + +def _sha256_bytes(content: bytes) -> str: + return hashlib.sha256(content).hexdigest() + + +def _sha256_file(path: Path) -> str: + try: + return _sha256_bytes(path.read_bytes()) + except OSError as exc: + detail = exc.strerror or str(exc) + raise ArtifactError( + "lease", f"cannot hash {path}: {detail}" + ) from exc + + +def _command_sha256(command: Sequence[str]) -> str: + encoded = json.dumps( + list(command), + ensure_ascii=True, + separators=(",", ":"), + ).encode("utf-8") + return _sha256_bytes(encoded) + + +def _read_proc_bytes(root: Path, process_id: int, name: str) -> bytes: + path = root / str(process_id) / name + try: + return path.read_bytes() + except OSError as exc: + detail = exc.strerror or str(exc) + raise LeaseValidationError( + f"cannot read {path}: {detail}" + ) from exc + + +def _parse_proc_stat(content: str) -> tuple[int, int, int]: + close_paren = content.rfind(")") + if close_paren < 0: + raise LeaseValidationError("malformed process stat") + fields = content[close_paren + 1:].split() + if len(fields) < 20: + raise LeaseValidationError("malformed process stat") + try: + return int(fields[1]), int(fields[2]), int(fields[19]) + except ValueError as exc: + raise LeaseValidationError("malformed process stat") from exc + + +def _read_proc_stat( + root: Path, process_id: int +) -> tuple[int, int, int]: + content = _read_proc_bytes( + root, process_id, "stat" + ).decode("utf-8") + return _parse_proc_stat(content) + + +def _write_json_atomic( + path: Path, + value: dict[str, object], + *, + create: bool = False, +) -> None: + parent = path.parent + temp_path = parent / ( + f".{path.name}.{os.getpid()}.{secrets.token_hex(8)}.tmp" + ) + payload = ( + json.dumps( + value, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ) + + "\n" + ).encode("utf-8") + try: + descriptor = os.open( + temp_path, + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + 0o600, + ) + with os.fdopen(descriptor, "wb") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + if create: + os.link(temp_path, path) + temp_path.unlink() + else: + os.replace(temp_path, path) + directory_descriptor = os.open(parent, os.O_RDONLY) + try: + os.fsync(directory_descriptor) + finally: + os.close(directory_descriptor) + except OSError as exc: + try: + temp_path.unlink() + except FileNotFoundError: + pass + detail = exc.strerror or str(exc) + action = "create" if create else "write" + raise ArtifactError( + "lease", f"cannot atomically {action} {path}: {detail}" + ) from exc + + +class LeaseManager: + def __init__( + self, + config: WatchdogConfig, + paths: ArtifactPaths, + *, + process_procfs_root: Path = Path("/proc"), + wall_clock: Callable[[], datetime] | None = None, + monotonic_ns: Callable[[], int] | None = None, + ): + self.config = config + self.lease_path = paths.lease + self.heartbeat_path = paths.heartbeat + self.audit_path = paths.audit + self.process_procfs_root = process_procfs_root + self.wall_clock = wall_clock + self.monotonic_ns = monotonic_ns or time.monotonic_ns + self.lease_id = secrets.token_hex(16) + self.sequence = 0 + self.lease: dict[str, object] | None = None + + def _watchdog_identity(self) -> dict[str, object]: + script_path = Path(__file__).resolve() + cmdline_path = ( + self.process_procfs_root / str(os.getpid()) / "cmdline" + ) + proc_start_time_ticks: int | None = None + try: + cmdline = cmdline_path.read_bytes() + _, _, proc_start_time_ticks = _read_proc_stat( + self.process_procfs_root, os.getpid() + ) + except (OSError, LeaseValidationError): + if sys.platform.startswith("linux"): + raise ArtifactError( + "lease", + "cannot read watchdog process identity from procfs", + ) + cmdline = b"\0".join( + os.fsencode(argument) for argument in sys.argv + ) + return { + "pid": os.getpid(), + "start_time_utc": _timestamp_utc(self.wall_clock), + "proc_start_time_ticks": proc_start_time_ticks, + "cmdline_sha256": _sha256_bytes(cmdline), + "script_path": str(script_path), + "script_sha256": _sha256_file(script_path), + } + + def _heartbeat_record( + self, + state: str, + sample: dict[str, object] | None = None, + ) -> dict[str, object]: + assert self.lease is not None + self.sequence += 1 + record: dict[str, object] = { + "format": HEARTBEAT_FORMAT, + "version": HEARTBEAT_VERSION, + "lease_id": self.lease_id, + "sequence": self.sequence, + "state": state, + "updated_at": _timestamp_utc(self.wall_clock), + "updated_monotonic_ns": self.monotonic_ns(), + "watchdog_pid": self.lease["watchdog_pid"], + "watchdog_start_time_ticks": ( + self.lease["watchdog_start_time_ticks"] + ), + "child_pid": self.lease["child_pid"], + "child_process_group_id": ( + self.lease["child_process_group_id"] + ), + } + if sample is not None: + record["sample"] = sample + return record + + def start(self, child: ProcessHandle) -> None: + watchdog_identity = self._watchdog_identity() + self.lease = { + "format": LEASE_FORMAT, + "version": LEASE_VERSION, + "lease_id": self.lease_id, + "state": "active", + "watchdog_pid": watchdog_identity["pid"], + "watchdog_start_time_utc": ( + watchdog_identity["start_time_utc"] + ), + "watchdog_start_time_ticks": ( + watchdog_identity["proc_start_time_ticks"] + ), + "watchdog_command_sha256": ( + watchdog_identity["cmdline_sha256"] + ), + "watchdog_script_path": watchdog_identity["script_path"], + "watchdog_script_sha256": ( + watchdog_identity["script_sha256"] + ), + "soft_bytes": self.config.soft_bytes, + "emergency_bytes": self.config.emergency_bytes, + "strict_ceiling_bytes": STRICT_CEILING_BYTES, + "child_pid": child.pid, + "child_process_group_id": child.pid, + "command": list(self.config.command), + "child_command_sha256": _command_sha256( + self.config.command + ), + "heartbeat_path": str(self.heartbeat_path), + "max_heartbeat_age_seconds": ( + self.config.heartbeat_max_age_seconds + ), + "audit_path": str(self.audit_path), + "procfs_root": str( + self.config.procfs_root.expanduser().resolve() + ), + } + heartbeat = self._heartbeat_record("active") + _write_json_atomic(self.heartbeat_path, heartbeat, create=True) + _write_json_atomic(self.lease_path, self.lease, create=True) + + def update_heartbeat(self, sample: dict[str, object]) -> None: + heartbeat = self._heartbeat_record("active", sample) + _write_json_atomic(self.heartbeat_path, heartbeat) + + def finalize(self, final_record: dict[str, object]) -> None: + if self.lease is None: + return + self.lease["state"] = "final" + self.lease["final"] = final_record + heartbeat = self._heartbeat_record("final") + _write_json_atomic(self.heartbeat_path, heartbeat) + _write_json_atomic(self.lease_path, self.lease) + + +def _read_json_object(path: Path) -> dict[str, object]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise LeaseValidationError( + f"cannot read valid JSON from {path}: {exc}" + ) from exc + if not isinstance(value, dict): + raise LeaseValidationError(f"{path} must contain a JSON object") + return value + + +def _require_int(value: object, field: str) -> int: + if not isinstance(value, int) or isinstance(value, bool): + raise LeaseValidationError(f"lease field {field} is invalid") + return value + + +def _require_string(value: object, field: str) -> str: + if not isinstance(value, str) or not value: + raise LeaseValidationError(f"lease field {field} is invalid") + return value + + +def validate_active_lease( + lease_path: Path, + *, + expected_script_path: Path, + expected_soft_bytes: int = DEFAULT_SOFT_BYTES, + expected_emergency_bytes: int = DEFAULT_EMERGENCY_BYTES, + expected_procfs_root: Path = Path("/proc"), + expected_command_sha256: str | None = None, + expected_heartbeat_path: Path | None = None, + expected_audit_path: Path | None = None, + expected_max_heartbeat_age_seconds: float | None = None, + current_process_id: int | None = None, + process_procfs_root: Path = Path("/proc"), + monotonic_ns: Callable[[], int] | None = None, +) -> dict[str, object]: + lease_path = lease_path.expanduser().resolve() + lease = _read_json_object(lease_path) + if ( + lease.get("format") != LEASE_FORMAT + or lease.get("version") != LEASE_VERSION + or lease.get("state") != "active" + ): + raise LeaseValidationError("lease format, version, or state is invalid") + + script_path = Path( + _require_string( + lease.get("watchdog_script_path"), + "watchdog_script_path", + ) + ).resolve() + expected_script_path = expected_script_path.expanduser().resolve() + if script_path != expected_script_path: + raise LeaseValidationError("watchdog script path does not match") + script_sha256 = _require_string( + lease.get("watchdog_script_sha256"), + "watchdog_script_sha256", + ) + if script_sha256 != _sha256_file(expected_script_path): + raise LeaseValidationError("watchdog script SHA does not match") + + if ( + _require_int( + lease.get("soft_bytes"), "soft_bytes" + ) + != expected_soft_bytes + or _require_int( + lease.get("emergency_bytes"), + "emergency_bytes", + ) + != expected_emergency_bytes + or _require_int( + lease.get("strict_ceiling_bytes"), + "strict_ceiling_bytes", + ) + != STRICT_CEILING_BYTES + ): + raise LeaseValidationError("watchdog thresholds do not match") + lease_procfs_root = Path( + _require_string(lease.get("procfs_root"), "procfs_root") + ).resolve() + if lease_procfs_root != expected_procfs_root.expanduser().resolve(): + raise LeaseValidationError("watchdog procfs root does not match") + + watchdog_pid = _require_int( + lease.get("watchdog_pid"), "watchdog_pid" + ) + watchdog_start_ticks = _require_int( + lease.get("watchdog_start_time_ticks"), + "watchdog_start_time_ticks", + ) + _, _, live_watchdog_start_ticks = _read_proc_stat( + process_procfs_root, watchdog_pid + ) + if live_watchdog_start_ticks != watchdog_start_ticks: + raise LeaseValidationError("watchdog process start time does not match") + live_cmdline = _read_proc_bytes( + process_procfs_root, watchdog_pid, "cmdline" + ) + if _sha256_bytes(live_cmdline) != _require_string( + lease.get("watchdog_command_sha256"), + "watchdog_command_sha256", + ): + raise LeaseValidationError("watchdog command line does not match") + script_named = False + watchdog_cwd: Path | None = None + for raw_argument in live_cmdline.split(b"\0"): + if not raw_argument: + continue + argument_path = Path(os.fsdecode(raw_argument)).expanduser() + if not argument_path.is_absolute(): + if watchdog_cwd is None: + try: + watchdog_cwd = ( + process_procfs_root + / str(watchdog_pid) + / "cwd" + ).resolve() + except OSError as exc: + raise LeaseValidationError( + "cannot resolve watchdog working directory" + ) from exc + argument_path = watchdog_cwd / argument_path + if argument_path.resolve() == expected_script_path: + script_named = True + break + if not script_named: + raise LeaseValidationError( + "watchdog command line does not name the expected script" + ) + + child_pid = _require_int(lease.get("child_pid"), "child_pid") + process_group_id = _require_int( + lease.get("child_process_group_id"), + "child_process_group_id", + ) + child_parent_pid, child_group_id, _ = _read_proc_stat( + process_procfs_root, child_pid + ) + if ( + child_parent_pid != watchdog_pid + or child_group_id != process_group_id + or process_group_id != child_pid + ): + raise LeaseValidationError( + "monitored child parent or process group does not match" + ) + command = lease.get("command") + if ( + not isinstance(command, list) + or not command + or not all(isinstance(argument, str) for argument in command) + ): + raise LeaseValidationError("lease field command is invalid") + command_sha256 = _require_string( + lease.get("child_command_sha256"), + "child_command_sha256", + ) + if command_sha256 != _command_sha256(command): + raise LeaseValidationError("monitored command SHA is invalid") + if ( + expected_command_sha256 is not None + and command_sha256 != expected_command_sha256 + ): + raise LeaseValidationError("monitored command SHA does not match") + + process_id = ( + current_process_id + if current_process_id is not None + else os.getpid() + ) + _, current_group_id, _ = _read_proc_stat( + process_procfs_root, process_id + ) + if current_group_id != process_group_id: + raise LeaseValidationError( + "current process is outside the monitored process group" + ) + + heartbeat_path = Path( + _require_string( + lease.get("heartbeat_path"), "heartbeat_path" + ) + ).resolve() + if ( + expected_heartbeat_path is not None + and heartbeat_path + != expected_heartbeat_path.expanduser().resolve() + ): + raise LeaseValidationError("heartbeat path does not match") + heartbeat_max_age = lease.get("max_heartbeat_age_seconds") + if ( + not isinstance(heartbeat_max_age, (int, float)) + or isinstance(heartbeat_max_age, bool) + or not math.isfinite(heartbeat_max_age) + or heartbeat_max_age <= 0 + ): + raise LeaseValidationError( + "lease field max_heartbeat_age_seconds is invalid" + ) + if ( + expected_max_heartbeat_age_seconds is not None + and heartbeat_max_age != expected_max_heartbeat_age_seconds + ): + raise LeaseValidationError("heartbeat max age does not match") + heartbeat = _read_json_object(heartbeat_path) + lease_id = _require_string(lease.get("lease_id"), "lease_id") + if ( + heartbeat.get("format") != HEARTBEAT_FORMAT + or heartbeat.get("version") != HEARTBEAT_VERSION + or heartbeat.get("state") != "active" + or heartbeat.get("lease_id") != lease_id + or heartbeat.get("watchdog_pid") != watchdog_pid + or heartbeat.get("watchdog_start_time_ticks") + != watchdog_start_ticks + or heartbeat.get("child_pid") != child_pid + or heartbeat.get("child_process_group_id") != process_group_id + ): + raise LeaseValidationError("heartbeat identity does not match lease") + updated_monotonic_ns = _require_int( + heartbeat.get("updated_monotonic_ns"), + "heartbeat.updated_monotonic_ns", + ) + _require_int(heartbeat.get("sequence"), "heartbeat.sequence") + _require_string(heartbeat.get("updated_at"), "heartbeat.updated_at") + now_monotonic_ns = (monotonic_ns or time.monotonic_ns)() + age_ns = now_monotonic_ns - updated_monotonic_ns + if age_ns < 0 or age_ns > int(heartbeat_max_age * 1_000_000_000): + raise LeaseValidationError("watchdog heartbeat is stale") + + audit_path = Path( + _require_string(lease.get("audit_path"), "audit_path") + ).resolve() + if ( + expected_audit_path is not None + and audit_path != expected_audit_path.expanduser().resolve() + ): + raise LeaseValidationError("persistent audit path does not match") + try: + first_line = next( + line + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + if line + ) + first_record = json.loads(first_line) + if ( + not isinstance(first_record, dict) + or not isinstance(first_record.get("event"), str) + or not isinstance(first_record.get("timestamp"), str) + ): + raise LeaseValidationError( + "persistent audit does not contain watchdog records" + ) + except StopIteration as exc: + raise LeaseValidationError("persistent audit is empty") from exc + except (UnicodeError, json.JSONDecodeError) as exc: + raise LeaseValidationError( + "persistent audit does not contain valid JSONL" + ) from exc + except LeaseValidationError: + raise + except OSError as exc: + raise LeaseValidationError( + f"cannot inspect persistent audit {audit_path}: {exc}" + ) from exc + return lease + + class AuditLogger: def __init__( self, @@ -178,23 +777,77 @@ def __init__( wall_clock: Callable[[], datetime] | None = None, ): self.stream = stream - self.wall_clock = wall_clock or ( - lambda: datetime.now(timezone.utc) - ) - - def emit(self, event: str, **fields: object) -> None: - timestamp = self.wall_clock().astimezone(timezone.utc) + self.wall_clock = wall_clock + self.persistent_stream: IO[str] | None = None + self.lease_manager: LeaseManager | None = None + self.finalized = False + self.final_exit_code = EXIT_INTERNAL_ERROR + + def open_persistent(self, path: Path) -> None: + resolved_path = path.expanduser().resolve() + try: + descriptor = os.open( + resolved_path, + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + 0o600, + ) + self.persistent_stream = os.fdopen( + descriptor, "w", encoding="utf-8" + ) + except OSError as exc: + detail = exc.strerror or str(exc) + raise ArtifactError( + "audit", + f"cannot create persistent audit {resolved_path}: {detail}", + ) from exc + + def close(self) -> None: + if self.persistent_stream is not None: + self.persistent_stream.close() + self.persistent_stream = None + + def disable_component(self, component: str) -> None: + if component == "audit": + self.close() + elif component == "lease": + self.lease_manager = None + + def emit(self, event: str, **fields: object) -> dict[str, object]: record = { - "timestamp": timestamp.isoformat(timespec="milliseconds").replace( - "+00:00", "Z" - ), + "timestamp": _timestamp_utc(self.wall_clock), "event": event, **fields, } - self.stream.write( - json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" + line = ( + json.dumps(record, sort_keys=True, separators=(",", ":")) + + "\n" ) + self.stream.write(line) self.stream.flush() + if self.persistent_stream is not None: + try: + self.persistent_stream.write(line) + self.persistent_stream.flush() + os.fsync(self.persistent_stream.fileno()) + except OSError as exc: + detail = exc.strerror or str(exc) + raise ArtifactError( + "audit", + f"cannot write persistent audit: {detail}", + ) from exc + return record + + def heartbeat(self, sample: dict[str, object]) -> None: + if self.lease_manager is not None: + self.lease_manager.update_heartbeat(sample) + + def finalize(self, record: dict[str, object]) -> None: + if self.lease_manager is not None: + self.lease_manager.finalize(record) + + def mark_final(self, exit_code: int) -> None: + self.finalized = True + self.final_exit_code = exit_code def _child_status(returncode: int | None, started: bool = True) -> str: @@ -253,8 +906,35 @@ def _emit_final( fields.update(classification=classification, exit_code=exit_code) if error: fields["error"] = error - audit.emit("final", **fields) - return exit_code + previous_mask = signal.pthread_sigmask( + signal.SIG_BLOCK, PARENT_SIGNALS + ) + try: + try: + record = audit.emit("final", **fields) + audit.finalize(record) + except ArtifactError as exc: + audit.disable_component(exc.component) + fields.update( + classification="lease_error", + exit_code=EXIT_LEASE_ERROR, + threshold_reason="watchdog artifact finalization failed", + error=f"{exc.component}: {exc}", + ) + try: + record = audit.emit("final", **fields) + except ArtifactError as nested_exc: + audit.disable_component(nested_exc.component) + record = audit.emit("final", **fields) + try: + audit.finalize(record) + except ArtifactError as nested_exc: + audit.disable_component(nested_exc.component) + exit_code = EXIT_LEASE_ERROR + audit.mark_final(exit_code) + return exit_code + finally: + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) def _signal_process_group(process_group_id: int, signal_number: int) -> str: @@ -276,6 +956,8 @@ def _process_group_alive(process_group_id: int) -> bool: os.killpg(process_group_id, 0) except ProcessLookupError: return False + except PermissionError: + return True except OSError as exc: detail = exc.strerror or str(exc) raise ProcessGroupError( @@ -289,16 +971,16 @@ def _raise_parent_signal(signal_number: int, _frame: object) -> None: def _set_parent_signal_handlers( - handler: signal.Handlers, -) -> dict[int, signal.Handlers]: - previous: dict[int, signal.Handlers] = {} + handler: Any, +) -> dict[int, Any]: + previous: dict[int, Any] = {} for signal_number in PARENT_SIGNALS: previous[signal_number] = signal.signal(signal_number, handler) return previous def _restore_parent_signal_handlers( - previous: dict[int, signal.Handlers], + previous: dict[int, Any], ) -> None: for signal_number, handler in previous.items(): signal.signal(signal_number, handler) @@ -591,7 +1273,7 @@ def _monitor_child( state.peak_used_bytes = max( state.peak_used_bytes, state.snapshot.used_bytes ) - audit.emit( + sample_record = audit.emit( "sample", **_state_fields( state.snapshot, @@ -602,6 +1284,7 @@ def _monitor_child( "none", ), ) + audit.heartbeat(sample_record) if state.snapshot.active_swaps: return _kill_and_finish( @@ -679,7 +1362,7 @@ def run_watchdog( monotonic: Callable[[], float] | None = None, sleeper: Callable[[float], None] | None = None, ) -> int: - config.validate() + artifact_paths = config.validate() reader = reader or ProcfsReader(config.procfs_root) audit = audit or AuditLogger(sys.stderr) launcher = launcher or subprocess.Popen @@ -688,6 +1371,20 @@ def run_watchdog( monotonic = monotonic or time.monotonic sleeper = sleeper or time.sleep + if artifact_paths is not None: + try: + audit.open_persistent(artifact_paths.audit) + except ArtifactError as exc: + return _emit_final( + audit, + "lease_error", + EXIT_LEASE_ERROR, + "cannot initialize watchdog artifacts", + None, + None, + error=f"{exc.component}: {exc}", + ) + try: snapshot = reader.read_snapshot() except ProcfsError as exc: @@ -701,20 +1398,32 @@ def run_watchdog( error=str(exc), ) - audit.emit( - "preflight", - **_state_fields( + try: + audit.emit( + "preflight", + **_state_fields( + snapshot, + snapshot.used_bytes, + None, + None, + "not_created", + "none", + ), + soft_bytes=config.soft_bytes, + emergency_bytes=config.emergency_bytes, + strict_ceiling_bytes=STRICT_CEILING_BYTES, + ) + except ArtifactError as exc: + audit.disable_component(exc.component) + return _emit_final( + audit, + "lease_error", + EXIT_LEASE_ERROR, + "cannot write watchdog preflight audit", snapshot, snapshot.used_bytes, - None, - None, - "not_created", - "none", - ), - soft_bytes=config.soft_bytes, - emergency_bytes=config.emergency_bytes, - strict_ceiling_bytes=STRICT_CEILING_BYTES, - ) + error=f"{exc.component}: {exc}", + ) if snapshot.active_swaps: return _emit_final( @@ -748,21 +1457,51 @@ def run_watchdog( signal.SIG_BLOCK, PARENT_SIGNALS ) mask_restored = False - previous_handlers: dict[int, signal.Handlers] = {} + previous_handlers: dict[int, Any] = {} child: ProcessHandle | None = None state = RuntimeState(snapshot, snapshot.used_bytes) try: launch_mask = previous_mask + lease_manager = ( + LeaseManager(config, artifact_paths) + if artifact_paths is not None + else None + ) def restore_child_signal_mask() -> None: signal.pthread_sigmask(signal.SIG_SETMASK, launch_mask) try: - child = launcher( - config.command, - start_new_session=True, - preexec_fn=restore_child_signal_mask, - ) + if lease_manager is not None: + child_environment = os.environ.copy() + child_environment.update( + { + "STRIX_MEMORY_WATCHDOG_LEASE_PATH": str( + lease_manager.lease_path + ), + "STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH": str( + lease_manager.heartbeat_path + ), + "STRIX_MEMORY_WATCHDOG_AUDIT_PATH": str( + lease_manager.audit_path + ), + "STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS": ( + str(config.heartbeat_max_age_seconds) + ), + } + ) + child = launcher( + config.command, + start_new_session=True, + preexec_fn=restore_child_signal_mask, + env=child_environment, + ) + else: + child = launcher( + config.command, + start_new_session=True, + preexec_fn=restore_child_signal_mask, + ) except (OSError, ValueError) as exc: detail = getattr(exc, "strerror", None) or str(exc) return _emit_final( @@ -778,6 +1517,9 @@ def restore_child_signal_mask() -> None: previous_handlers = _set_parent_signal_handlers( _raise_parent_signal ) + if lease_manager is not None: + lease_manager.start(child) + audit.lease_manager = lease_manager signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) mask_restored = True audit.emit( @@ -803,8 +1545,41 @@ def restore_child_signal_mask() -> None: monotonic, sleeper, ) + except ArtifactError as exc: + _set_parent_signal_handlers(signal.SIG_IGN) + if exc.component == "audit": + audit.disable_component(exc.component) + if child is None: + return _emit_final( + audit, + "lease_error", + EXIT_LEASE_ERROR, + "watchdog artifact initialization failed", + state.snapshot, + state.peak_used_bytes, + error=f"{exc.component}: {exc}", + ) + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "lease_error", + EXIT_LEASE_ERROR, + "watchdog artifact update failed", + signal.SIGTERM, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + error=f"{exc.component}: {exc}", + ) except ParentSignal as exc: + if audit.finalized: + return audit.final_exit_code _set_parent_signal_handlers(signal.SIG_IGN) + assert child is not None signal_name = signal.Signals(exc.signal_number).name return _graceful_cleanup( audit, @@ -822,7 +1597,19 @@ def restore_child_signal_mask() -> None: sleeper, ) except Exception as exc: + if audit.finalized: + return audit.final_exit_code _set_parent_signal_handlers(signal.SIG_IGN) + if child is None: + return _emit_final( + audit, + "internal_error", + EXIT_INTERNAL_ERROR, + "unexpected pre-launch exception", + state.snapshot, + state.peak_used_bytes, + error=f"{type(exc).__name__}: {exc}", + ) return _graceful_cleanup( audit, child, @@ -900,6 +1687,39 @@ def parse_args(argv: Sequence[str]) -> WatchdogConfig: default=DEFAULT_SAMPLE_INTERVAL_SECONDS, help="procfs sampling interval (default: 1)", ) + parser.add_argument( + "--lease-path", + type=Path, + help=( + "atomically publish the watchdog-owned lease JSON; requires " + "--heartbeat-path and --audit-path" + ), + ) + parser.add_argument( + "--heartbeat-path", + type=Path, + help=( + "atomically update watchdog heartbeat JSON on every sample; " + "requires --lease-path and --audit-path" + ), + ) + parser.add_argument( + "--audit-path", + type=Path, + help=( + "create a persistent JSONL audit in addition to standard error; " + "requires --lease-path and --heartbeat-path" + ), + ) + parser.add_argument( + "--heartbeat-max-age-seconds", + type=_positive_float, + default=DEFAULT_HEARTBEAT_MAX_AGE_SECONDS, + help=( + "maximum heartbeat age accepted by a matching harness " + "(default: 5)" + ), + ) parser.add_argument( "command", nargs=argparse.REMAINDER, @@ -916,6 +1736,10 @@ def parse_args(argv: Sequence[str]) -> WatchdogConfig: emergency_bytes=args.emergency_gib * GIB, grace_seconds=args.grace_seconds, sample_interval_seconds=args.sample_interval_seconds, + lease_path=args.lease_path, + heartbeat_path=args.heartbeat_path, + audit_path=args.audit_path, + heartbeat_max_age_seconds=args.heartbeat_max_age_seconds, ) @@ -934,6 +1758,18 @@ def main(argv: Sequence[str] | None = None) -> int: None, error=str(exc), ) + except Exception as exc: + return _emit_final( + audit, + "internal_error", + EXIT_INTERNAL_ERROR, + "unexpected watchdog error", + None, + None, + error=f"{type(exc).__name__}: {exc}", + ) + finally: + audit.close() if __name__ == "__main__": diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index eca818da0f6f..81c10f84c38f 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -3,6 +3,7 @@ from __future__ import annotations import importlib.util +import hashlib import io import json import os @@ -80,7 +81,7 @@ def poll(self) -> int | None: def wait(self, timeout: float | None = None) -> int: if self.returncode is None: - raise subprocess.TimeoutExpired("fake", timeout) + raise subprocess.TimeoutExpired("fake", timeout or 0.0) return self.returncode @@ -225,6 +226,33 @@ def _write_procfs_fixture(root: Path) -> None: encoding="utf-8", ) + @staticmethod + def _lease_arguments(root: Path) -> list[str]: + return [ + "--lease-path", + str(root / "lease.json"), + "--heartbeat-path", + str(root / "heartbeat.json"), + "--audit-path", + str(root / "persistent-audit.jsonl"), + ] + + @staticmethod + def _proc_stat( + process_id: int, + parent_id: int, + process_group_id: int, + start_time_ticks: int, + ) -> str: + fields = [ + "S", + str(parent_id), + str(process_group_id), + *(["0"] * 16), + str(start_time_ticks), + ] + return f"{process_id} (python) {' '.join(fields)}\n" + def test_parent_signals_leave_no_child_or_grandchild(self) -> None: child_code = ( "import os,signal,sys,time;" @@ -249,14 +277,15 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: root = Path(temp_dir) pid_file = root / "pids" self._write_procfs_fixture(root) - audit_path = root / "audit.jsonl" - with audit_path.open("w", encoding="utf-8") as audit: + stderr_path = root / "stderr.jsonl" + with stderr_path.open("w", encoding="utf-8") as audit: wrapper = subprocess.Popen( [ sys.executable, str(SCRIPT_PATH), "--procfs-root", str(root), + *self._lease_arguments(root), "--grace-seconds", "0.2", "--sample-interval-seconds", @@ -304,7 +333,7 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: ) records = [ json.loads(line) - for line in audit_path.read_text( + for line in stderr_path.read_text( encoding="utf-8" ).splitlines() ] @@ -344,6 +373,32 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: self.assertFalse( self._process_is_running(process_id) ) + lease = json.loads( + (root / "lease.json").read_text( + encoding="utf-8" + ) + ) + heartbeat = json.loads( + (root / "heartbeat.json").read_text( + encoding="utf-8" + ) + ) + persistent_records = [ + json.loads(line) + for line in ( + root / "persistent-audit.jsonl" + ).read_text(encoding="utf-8").splitlines() + ] + self.assertEqual(lease["state"], "final") + self.assertEqual( + lease["final"]["classification"], + "parent_signal", + ) + self.assertEqual(heartbeat["state"], "final") + self.assertEqual( + persistent_records[-1]["classification"], + "parent_signal", + ) def test_child_sigterm_handler_exits_without_escalation(self) -> None: child_code = ( @@ -620,27 +675,124 @@ def test_configuration_rejects_non_finite_timing(self) -> None: with self.assertRaisesRegex(ValueError, "grace period"): config.validate() + def test_final_record_survives_artifact_failures(self) -> None: + class FailingLease: + def finalize(self, record: dict[str, Any]) -> None: + raise watchdog.ArtifactError("lease", "write failed") + + class FailingAudit(io.StringIO): + def write(self, value: str) -> int: + raise OSError("write failed") + + for component in ("lease", "audit"): + with self.subTest(component=component): + stream = io.StringIO() + audit = watchdog.AuditLogger(stream) + if component == "lease": + setattr(audit, "lease_manager", FailingLease()) + else: + audit.persistent_stream = FailingAudit() + result = watchdog._emit_final( + audit, + "child_exit", + 0, + "child exited", + snapshot(50), + 50, + ) + records = [ + json.loads(line) + for line in stream.getvalue().splitlines() + ] + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual( + records[-1]["classification"], "lease_error" + ) + self.assertTrue(audit.finalized) + self.assertEqual( + audit.final_exit_code, watchdog.EXIT_LEASE_ERROR + ) + + def test_invalid_artifact_path_emits_configuration_final(self) -> None: + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--lease-path", + "~strix-watchdog-user-does-not-exist/lease.json", + "--heartbeat-path", + "/tmp/heartbeat.json", + "--audit-path", + "/tmp/audit.jsonl", + "--", + sys.executable, + "-c", + "pass", + ], + capture_output=True, + check=False, + text=True, + ) + + self.assertEqual(result.returncode, watchdog.EXIT_PROCFS_ERROR) + final = json.loads(result.stderr.splitlines()[-1]) + self.assertEqual(final["classification"], "configuration_error") + def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) + child_result_path = root / "child-result.json" self._write_procfs_fixture(root) + child_code = ( + "import json,os,sys,time\n" + "keys=('STRIX_MEMORY_WATCHDOG_LEASE_PATH'," + "'STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH'," + "'STRIX_MEMORY_WATCHDOG_AUDIT_PATH')\n" + "deadline=time.monotonic()+5\n" + "while True:\n" + " try:\n" + " with open(os.environ[keys[0]],encoding='utf-8') as stream:\n" + " lease=json.load(stream)\n" + " break\n" + " except (OSError,json.JSONDecodeError):\n" + " if time.monotonic()>=deadline: raise\n" + " time.sleep(0.01)\n" + "assert os.getpgrp()==lease['child_process_group_id']\n" + "open(sys.argv[1],'w').write(json.dumps({" + "key:os.environ[key] for key in keys}))\n" + "raise SystemExit(23)\n" + ) result = subprocess.run( [ sys.executable, str(SCRIPT_PATH), "--procfs-root", str(root), + *self._lease_arguments(root), "--sample-interval-seconds", "0.01", "--", sys.executable, "-c", - "raise SystemExit(23)", + child_code, + str(child_result_path), ], capture_output=True, check=False, text=True, ) + lease = json.loads( + (root / "lease.json").read_text(encoding="utf-8") + ) + persistent_records = [ + json.loads(line) + for line in ( + root / "persistent-audit.jsonl" + ).read_text(encoding="utf-8").splitlines() + ] + child_result = json.loads( + child_result_path.read_text(encoding="utf-8") + ) self.assertEqual(result.returncode, 23) records = [ @@ -648,6 +800,277 @@ def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: ] self.assertEqual(records[-1]["classification"], "child_exit") self.assertEqual(records[-1]["child_returncode"], 23) + self.assertEqual(lease["format"], watchdog.LEASE_FORMAT) + self.assertEqual(lease["version"], watchdog.LEASE_VERSION) + self.assertEqual(lease["state"], "final") + self.assertEqual(lease["soft_bytes"], 116 * 1024**3) + self.assertEqual( + lease["emergency_bytes"], 118 * 1024**3 + ) + self.assertEqual( + lease["child_command_sha256"], + watchdog._command_sha256( + ( + sys.executable, + "-c", + child_code, + str(child_result_path), + ) + ), + ) + self.assertEqual( + lease["final"]["classification"], "child_exit" + ) + self.assertEqual( + persistent_records[-1]["classification"], "child_exit" + ) + self.assertEqual( + child_result["STRIX_MEMORY_WATCHDOG_LEASE_PATH"], + str((root / "lease.json").resolve()), + ) + self.assertEqual( + child_result["STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH"], + str((root / "heartbeat.json").resolve()), + ) + self.assertEqual( + child_result["STRIX_MEMORY_WATCHDOG_AUDIT_PATH"], + str((root / "persistent-audit.jsonl").resolve()), + ) + + def test_existing_lease_fails_closed_and_stops_child(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + self._write_procfs_fixture(root) + lease_path = root / "lease.json" + lease_path.write_text("untrusted\n", encoding="utf-8") + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--grace-seconds", + "0.1", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + "import time;time.sleep(30)", + ], + capture_output=True, + check=False, + text=True, + timeout=5, + ) + records = [ + json.loads(line) for line in result.stderr.splitlines() + ] + child_pid = records[-1]["child_pid"] + self.assertEqual( + lease_path.read_text(encoding="utf-8"), "untrusted\n" + ) + + self.assertEqual(result.returncode, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(records[-1]["classification"], "lease_error") + self.assertFalse(self._process_is_running(child_pid)) + + def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + process_root = root / "proc" + watchdog_pid = 1200 + child_pid = 1300 + current_pid = 1400 + watchdog_start_ticks = 456789 + script_path = root / "watchdog.py" + script_path.write_text("print('watchdog')\n", encoding="utf-8") + cmdline = ( + b"/usr/bin/python3\0watchdog.py\0--lease-path\0" + ) + for process_id, parent_id, group_id, start_ticks in ( + (watchdog_pid, 1, watchdog_pid, watchdog_start_ticks), + (child_pid, watchdog_pid, child_pid, 456790), + (current_pid, child_pid, child_pid, 456791), + ): + process_dir = process_root / str(process_id) + process_dir.mkdir(parents=True) + (process_dir / "stat").write_text( + self._proc_stat( + process_id, + parent_id, + group_id, + start_ticks, + ), + encoding="utf-8", + ) + (process_root / str(watchdog_pid) / "cwd").symlink_to( + root, target_is_directory=True + ) + (process_root / str(watchdog_pid) / "cmdline").write_bytes( + cmdline + ) + + lease_path = root / "lease.json" + heartbeat_path = root / "heartbeat.json" + audit_path = root / "audit.jsonl" + audit_path.write_text( + '{"event":"child_started","timestamp":"2026-01-01T00:00:00Z"}\n', + encoding="utf-8", + ) + command = ["python3", "run_matrix.py"] + lease = { + "format": watchdog.LEASE_FORMAT, + "version": watchdog.LEASE_VERSION, + "lease_id": "test-lease", + "state": "active", + "watchdog_pid": watchdog_pid, + "watchdog_start_time_utc": "2026-01-01T00:00:00.000Z", + "watchdog_start_time_ticks": watchdog_start_ticks, + "watchdog_command_sha256": hashlib.sha256( + cmdline + ).hexdigest(), + "watchdog_script_path": str(script_path), + "watchdog_script_sha256": hashlib.sha256( + script_path.read_bytes() + ).hexdigest(), + "soft_bytes": watchdog.DEFAULT_SOFT_BYTES, + "emergency_bytes": watchdog.DEFAULT_EMERGENCY_BYTES, + "strict_ceiling_bytes": watchdog.STRICT_CEILING_BYTES, + "child_pid": child_pid, + "child_process_group_id": child_pid, + "command": command, + "child_command_sha256": watchdog._command_sha256( + command + ), + "heartbeat_path": str(heartbeat_path), + "max_heartbeat_age_seconds": 5.0, + "audit_path": str(audit_path), + "procfs_root": "/proc", + } + heartbeat = { + "format": watchdog.HEARTBEAT_FORMAT, + "version": watchdog.HEARTBEAT_VERSION, + "lease_id": "test-lease", + "sequence": 4, + "state": "active", + "updated_at": "2026-01-01T00:00:01.000Z", + "updated_monotonic_ns": 9_000_000_000, + "watchdog_pid": watchdog_pid, + "watchdog_start_time_ticks": ( + watchdog_start_ticks + ), + "child_pid": child_pid, + "child_process_group_id": child_pid, + } + lease_path.write_text(json.dumps(lease), encoding="utf-8") + heartbeat_path.write_text( + json.dumps(heartbeat), encoding="utf-8" + ) + + validated = watchdog.validate_active_lease( + lease_path, + expected_script_path=script_path, + expected_command_sha256=watchdog._command_sha256( + command + ), + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + expected_max_heartbeat_age_seconds=5.0, + current_process_id=current_pid, + process_procfs_root=process_root, + monotonic_ns=lambda: 10_000_000_000, + ) + self.assertEqual(validated["lease_id"], "test-lease") + + with self.subTest("tampered script SHA"): + tampered = dict(lease) + tampered["watchdog_script_sha256"] = "0" * 64 + lease_path.write_text( + json.dumps(tampered), encoding="utf-8" + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, "script SHA" + ): + watchdog.validate_active_lease( + lease_path, + expected_script_path=script_path, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + current_process_id=current_pid, + process_procfs_root=process_root, + monotonic_ns=lambda: 10_000_000_000, + ) + + with self.subTest("stale heartbeat"): + lease_path.write_text( + json.dumps(lease), encoding="utf-8" + ) + heartbeat["updated_monotonic_ns"] = 1 + heartbeat_path.write_text( + json.dumps(heartbeat), encoding="utf-8" + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, "heartbeat is stale" + ): + watchdog.validate_active_lease( + lease_path, + expected_script_path=script_path, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + current_process_id=current_pid, + process_procfs_root=process_root, + monotonic_ns=lambda: 10_000_000_000, + ) + + with self.subTest("arbitrary heartbeat"): + heartbeat["updated_monotonic_ns"] = 9_000_000_000 + heartbeat["lease_id"] = "helper-lease" + heartbeat_path.write_text( + json.dumps(heartbeat), encoding="utf-8" + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "heartbeat identity", + ): + watchdog.validate_active_lease( + lease_path, + expected_script_path=script_path, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + current_process_id=current_pid, + process_procfs_root=process_root, + monotonic_ns=lambda: 10_000_000_000, + ) + + with self.subTest("outside process group"): + heartbeat["lease_id"] = "test-lease" + heartbeat_path.write_text( + json.dumps(heartbeat), encoding="utf-8" + ) + (process_root / str(current_pid) / "stat").write_text( + self._proc_stat( + current_pid, + child_pid, + 9999, + 456791, + ), + encoding="utf-8", + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "outside the monitored process group", + ): + watchdog.validate_active_lease( + lease_path, + expected_script_path=script_path, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + current_process_id=current_pid, + process_procfs_root=process_root, + monotonic_ns=lambda: 10_000_000_000, + ) def test_zero_swap_gate_launches_and_propagates_child_exit(self) -> None: harness = Harness([snapshot(50)], FakeProcess(returncode=37)) @@ -702,13 +1125,29 @@ def exit_on_term(process: FakeProcess, signal_number: int) -> None: signal_handler=exit_on_term, ) - result = harness.run() + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + result = harness.run( + lease_path=root / "lease.json", + heartbeat_path=root / "heartbeat.json", + audit_path=root / "audit.jsonl", + ) + harness.audit.close() + heartbeat = json.loads( + (root / "heartbeat.json").read_text(encoding="utf-8") + ) + lease = json.loads( + (root / "lease.json").read_text(encoding="utf-8") + ) self.assertEqual(result, watchdog.EXIT_SOFT_LIMIT) self.assertEqual(harness.signals, [signal.SIGTERM]) self.assertEqual( harness.records()[-1]["classification"], "soft_limit" ) + self.assertEqual(heartbeat["state"], "final") + self.assertEqual(heartbeat["sequence"], 3) + self.assertEqual(lease["final"]["classification"], "soft_limit") def test_emergency_limit_sends_sigkill(self) -> None: def exit_on_kill(process: FakeProcess, signal_number: int) -> None: @@ -804,7 +1243,23 @@ def exit_on_kill(process: FakeProcess, signal_number: int) -> None: signal_handler=exit_on_kill, ) - result = harness.run() + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + result = harness.run( + lease_path=root / "lease.json", + heartbeat_path=root / "heartbeat.json", + audit_path=root / "audit.jsonl", + ) + harness.audit.close() + lease = json.loads( + (root / "lease.json").read_text(encoding="utf-8") + ) + persistent_records = [ + json.loads(line) + for line in (root / "audit.jsonl").read_text( + encoding="utf-8" + ).splitlines() + ] self.assertEqual(result, watchdog.EXIT_INTERNAL_ERROR) self.assertEqual( @@ -814,6 +1269,10 @@ def exit_on_kill(process: FakeProcess, signal_number: int) -> None: final = harness.records()[-1] self.assertEqual(final["classification"], "internal_error") self.assertIn("RuntimeError: unexpected", final["error"]) + self.assertEqual(lease["final"]["classification"], "internal_error") + self.assertEqual( + persistent_records[-1]["classification"], "internal_error" + ) def test_launch_failure_is_explicit(self) -> None: harness = Harness([snapshot(50)], FakeProcess()) @@ -823,7 +1282,7 @@ def fail_launch( ) -> FakeProcess: raise FileNotFoundError(2, "No such file or directory") - harness.launcher = fail_launch + setattr(harness, "launcher", fail_launch) result = harness.run() From 265f3df3055cbcdbf10a545d8d7c234ddd495fab Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 22:42:24 -0700 Subject: [PATCH 25/32] deepseek41 : bind final validation envelope Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 8 +++++++- tests/test-deepseek41-admission.cpp | 4 ++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 33f83dddaa97..d018f371ce35 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -3,14 +3,20 @@ `scripts/strix_memory_watchdog.py` is an external Linux command wrapper for headless Strix Halo validation. It does not change model loading or cache sizing. It measures host-wide memory from procfs and controls the launched command's process group. ```sh +ROCR_VISIBLE_DEVICES=0 \ +HIP_VISIBLE_DEVICES=0 \ +HIP_LAUNCH_BLOCKING=1 \ ./scripts/strix_memory_watchdog.py -- \ ./build/bin/llama-server \ -m /mnt/models/deepseek-v41/DeepSeek-V4.1-Flash-Q2.gguf \ - -c 32768 -b 2048 -ub 32 -ngl 99 + -c 32768 -b 2048 -ub 32 -ngl 99 -dev ROCm0 \ + --expert-cache-slots 192 --expert-cache-mib 72900 ``` DeepSeek V4.1 also runs an in-process admission check before expert-cache or model backend allocation. The default model parameters read `/proc/meminfo` and `/proc/swaps`, reject any configured swap entry, measure the full-graph state through its no-allocation memory implementation, account unified host/GPU memory once, and auto-fit complete expert slots under 116 GiB total projected host use. The context checks the measured scheduler workspace against the admitted conservative workspace envelope before inference. Admission fails closed unless every selected accelerator reports `GGML_BACKEND_DEVICE_TYPE_IGPU`; CPU-only, discrete GPU, RPC, and tensor-parallel meta-device configurations are not treated as one procfs-accounted pool. The external watchdog is still required for guarded validation because it monitors host-wide use after startup and controls the complete process group. +The final Strix validation preflight must confirm that `ROCm0` reports `gfx1151` before running this command. It must also verify the inherited device and launch-blocking environment, the exact ubatch and cache arguments, and the active watchdog lease. The 72900 MiB budget is exactly 192 published expert slots; admission rejects a disagreement between the byte and slot caps. + Use `--dsv41-procfs-root`, `--dsv41-memory-soft-mib`, `--dsv41-memory-watchdog-mib`, `--dsv41-memory-hard-mib`, and `--dsv41-memory-safety-margin-mib` only when reproducing admission tests or applying a more conservative host policy. `--expert-cache-slots` and `--expert-cache-mib` are optional caps; zero auto-fits. If both cache options are set, their capacity must describe the same number of complete published tensor slots. The current expert runtime remaps the unique routed-expert union for one ubatch. Admission therefore requires `min(384, 6 * ubatch)` resident slots instead of only six top-k slots. For example, a 224-slot cache admits at most ubatch 37. The common CLI and server default to ubatch 32 for DeepSeek V4.1 when `-ub` is not specified; an explicit value is preserved and must fit. Admission includes the resident staging cache, a complete worst-case replacement set, and the largest aligned direct-I/O bounce read. It reports both the required slot count and the admitted ubatch capacity and fails rather than lowering an explicit ubatch. DeepSeek V4.1 embedding extraction is rejected because those optional output buffers are not part of the bounded generation profile. diff --git a/tests/test-deepseek41-admission.cpp b/tests/test-deepseek41-admission.cpp index 39ba6e564a4c..effefdedeeb3 100644 --- a/tests/test-deepseek41-admission.cpp +++ b/tests/test-deepseek41-admission.cpp @@ -323,11 +323,15 @@ void test_expert_replacement_peak() { auto params = base_params(); params.n_ubatch = 32; params.configured_cache_slots = 192; + params.configured_cache_bytes = 72900ULL << 20; auto result = llama_dsv41_admit(host_with_used(0), 0, published_tensors(), params); + REQUIRE(result.expert_slots == 192); + REQUIRE(result.expert_cache_bytes == params.configured_cache_bytes); REQUIRE(result.expert_replacement_bytes == 1911029760); REQUIRE(result.direct_io_bounce_bytes == 3874816); REQUIRE(result.expert_replacement_bytes + result.direct_io_bounce_bytes == 1914904576); + params.configured_cache_bytes = 0; params.n_ubatch = 36; params.configured_cache_slots = 216; result = llama_dsv41_admit(host_with_used(0), 0, published_tensors(), params); From c4598ee747fdb811b474c0b41bdb38eeda3bbc0c Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 23:04:21 -0700 Subject: [PATCH 26/32] scripts : harden watchdog fail-closed lease Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 28 +- scripts/strix_memory_watchdog.py | 862 ++++++++++++++++++++++++---- tests/test_strix_memory_watchdog.py | 682 ++++++++++++++++++---- 3 files changed, 1339 insertions(+), 233 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 5565716a0a3e..49a1c05547fc 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -38,34 +38,40 @@ Use all three artifact options together when another process must prove that it python3 tools/deepseek-v41-trace/run_matrix.py ``` -The watchdog creates the persistent audit before launch, then atomically creates the lease and heartbeat after `Popen` returns. Existing artifact paths are rejected rather than overwritten. The child receives the resolved paths through `STRIX_MEMORY_WATCHDOG_LEASE_PATH`, `STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH`, and `STRIX_MEMORY_WATCHDOG_AUDIT_PATH`. It also receives `STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS`. +The watchdog creates and exclusively locks the persistent audit before launch. It then starts an internal guardian as the new session and process-group leader; the guardian starts the supplied command in that same group without inheriting the private control pipe. After the guardian reports the payload PID, the watchdog atomically creates the lease and heartbeat. Existing artifact paths are rejected rather than overwritten. The payload receives the resolved paths through `STRIX_MEMORY_WATCHDOG_LEASE_PATH`, `STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH`, and `STRIX_MEMORY_WATCHDOG_AUDIT_PATH`. It also receives `STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS`. The child can run before the first atomic lease rename. A matching preflight must retry the inherited lease path for a bounded interval and fail closed if a complete valid lease does not appear. It must not accept a lease path supplied separately by the operator. Lease format `strix-memory-watchdog-lease`, version 1, contains: - `lease_id` and active/final `state` -- `watchdog_pid`, `watchdog_start_time_utc`, Linux `watchdog_start_time_ticks`, `watchdog_command_sha256`, `watchdog_script_path`, and `watchdog_script_sha256` +- `watchdog_pid`, `watchdog_start_time_utc`, Linux `watchdog_start_time_ticks`, `watchdog_executable_path`, `watchdog_command_sha256`, `watchdog_script_path`, and `watchdog_script_sha256` - exact `soft_bytes`, `emergency_bytes`, and `strict_ceiling_bytes` - `procfs_root` -- `child_pid`, `child_process_group_id`, `command`, and `child_command_sha256` +- `guardian_pid`, payload `child_pid`, `child_process_group_id`, `command`, and `child_command_sha256` - `heartbeat_path`, `max_heartbeat_age_seconds`, and `audit_path` +- device, inode, owner, and mode identity for atomic JSON artifacts, plus the watchdog-held audit descriptor identity - the authoritative `final` audit record after termination -Heartbeat format `strix-memory-watchdog-heartbeat`, version 1, binds `lease_id`, watchdog PID/start ticks, child PID/process group, sequence, state, and update timestamps. Every memory sample atomically replaces the heartbeat and includes the complete sample audit record. A final heartbeat and final lease update remain on disk with the persistent JSONL audit; the watchdog does not delete this evidence. +Heartbeat format `strix-memory-watchdog-heartbeat`, version 1, binds `lease_id`, watchdog PID/start ticks, child PID/process group, sequence, state, and update timestamps. Every memory sample first checks swap and memory thresholds, pulses the guardian through the private nonblocking pipe, then atomically replaces the heartbeat with the complete sample audit record and its persistent-audit record hash. A blocked audit or heartbeat write cannot delay the emergency signal. A final heartbeat and final lease update remain on disk with the persistent JSONL audit; the watchdog does not delete this evidence. + +The guardian uses Linux `PR_SET_PDEATHSIG` with a parent-race check. It kills its process group on watchdog death, control-pipe EOF/error, or a missed pulse deadline, including a stopped or wedged watchdog. The payload must call `start_process_group_lease_guard()` before it starts exporter descendants. This validates the lease with bounded startup retries, arms a second parent-death link to the guardian, and starts a thread that kills the process group if the watchdog evidence becomes stale or invalid. A matching Linux preflight must verify all of the following: - The inherited lease, heartbeat, and audit paths match the paths inside the lease. -- The expected repository script path hashes to `watchdog_script_sha256`. -- `/proc//stat` start ticks and `/proc//cmdline` SHA-256 match the lease. -- The watchdog command line names the expected script. -- The current process group equals `child_process_group_id`, whose leader is `child_pid` and whose parent is `watchdog_pid`. +- `/proc//exe` is the exact expected Python executable and argv position 1 is the exact repository watchdog script. `-c`, `-m`, helper-script, inert-argument, and interpreter-option substitutions are rejected. +- The watchdog command line itself supplies the exact 116/118 GiB thresholds, `/proc`, inherited artifact paths, and command after `--`; the lease cannot override those expectations. +- `/proc//stat` start ticks and `/proc//cmdline` SHA-256 match the lease and remain stable across validation. A pidfd is held during validation when Linux provides `pidfd_open`. +- The topology is watchdog parent -> guardian process-group leader -> payload child. The current process must be inside `child_process_group_id`. - The command identity is expected, the procfs root is `/proc`, and thresholds are exactly 116 GiB soft, 118 GiB emergency, and 120 GiB strict ceiling for the final run. -- The heartbeat identity matches the lease and its monotonic timestamp is not older than `max_heartbeat_age_seconds`. -- The persistent audit exists and contains watchdog JSONL records. +- Lease and heartbeat files are regular, mode 0600, owned by the current UID, opened with `O_NOFOLLOW`, and match their recorded device/inode identity. +- The heartbeat identity matches the lease, its monotonic timestamp is not older than `max_heartbeat_age_seconds`, and its audit-record hash exists in the persistent audit. +- The persistent audit matches the watchdog-held descriptor device/inode and remains exclusively locked by the live watchdog. + +These checks reject accidental or helper-process substitution and make regular-file heartbeat forgery unable to keep the process group alive after private pulses stop. They are not a security boundary against intentionally hostile code running as the same UID; use a separately owned systemd user service or cgroup if that threat is in scope. -These checks bind the validation process to the live canonical watchdog. A standalone heartbeat helper has a different PID, start time, command line, script hash, and process group and cannot satisfy the lease. +The guardian controls only the process group. A payload that deliberately calls `setsid()` can escape it. The correctness harness must not do that. If arbitrary payload code is in scope, launch the watchdog in a service/cgroup configured to kill every member when the unit stops. Exit classifications are authoritative in the final JSON record. Operational failures use these exit codes: diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 0bb11a0622c3..dd185a00f5f2 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -3,15 +3,20 @@ from __future__ import annotations import argparse +import ctypes +import fcntl import hashlib import json import math import os import re import secrets +import select import signal +import stat import subprocess import sys +import threading import time from collections.abc import Callable, Sequence from dataclasses import dataclass @@ -32,6 +37,8 @@ LEASE_VERSION = 1 HEARTBEAT_FORMAT = "strix-memory-watchdog-heartbeat" HEARTBEAT_VERSION = 1 +PR_SET_PDEATHSIG = 1 +LEASE_GUARD_SIGNAL = signal.SIGUSR1 EXIT_PROCFS_ERROR = 2 EXIT_SWAP_ACTIVE = 3 @@ -82,6 +89,42 @@ def wait(self, timeout: float | None = None) -> int: ... +@dataclass +class GuardianProcess: + process: subprocess.Popen[bytes] + payload_pid: int + pulse_fd: int + + @property + def pid(self) -> int: + return self.process.pid + + def poll(self) -> int | None: + return self.process.poll() + + def wait(self, timeout: float | None = None) -> int: + return self.process.wait(timeout=timeout) + + def pulse(self) -> None: + try: + os.write(self.pulse_fd, b"\0") + except BlockingIOError as exc: + raise ProcessGroupError( + "guardian pulse pipe is blocked" + ) from exc + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot pulse guardian: {detail}" + ) from exc + + def close(self) -> None: + try: + os.close(self.pulse_fd) + except OSError: + pass + + @dataclass(frozen=True) class HostSnapshot: total_bytes: int @@ -139,6 +182,14 @@ def validate(self) -> ArtifactPaths | None: or self.sample_interval_seconds <= 0 ): raise ValueError("sample interval must be greater than zero") + if ( + not math.isfinite(self.heartbeat_max_age_seconds) + or self.heartbeat_max_age_seconds + <= self.sample_interval_seconds + ): + raise ValueError( + "heartbeat max age must be greater than sample interval" + ) lease_paths = ( self.lease_path, self.heartbeat_path, @@ -168,14 +219,6 @@ def validate(self) -> ArtifactPaths | None: raise ValueError( "lease, heartbeat, and audit paths must be distinct" ) - if ( - not math.isfinite(self.heartbeat_max_age_seconds) - or self.heartbeat_max_age_seconds - <= self.sample_interval_seconds - ): - raise ValueError( - "heartbeat max age must be greater than sample interval" - ) return paths return None @@ -278,6 +321,185 @@ def _command_sha256(command: Sequence[str]) -> str: return _sha256_bytes(encoded) +def _set_parent_death_signal( + signal_number: int, expected_parent_pid: int +) -> None: + if not sys.platform.startswith("linux"): + return + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(PR_SET_PDEATHSIG, signal_number, 0, 0, 0) != 0: + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + if os.getppid() != expected_parent_pid: + os.kill(os.getpid(), signal.SIGKILL) + + +def _kill_own_process_group( + _signal_number: int | None = None, + _frame: object | None = None, +) -> None: + try: + os.killpg(os.getpgrp(), signal.SIGKILL) + except OSError: + os._exit(EXIT_SIGNAL_ERROR) + + +def _guardian_main( + control_fd: int, + status_fd: int, + pulse_timeout_seconds: float, + command: tuple[str, ...], +) -> int: + if not sys.platform.startswith("linux"): + return EXIT_LAUNCH_ERROR + os.set_inheritable(control_fd, False) + os.set_inheritable(status_fd, False) + signal.signal(LEASE_GUARD_SIGNAL, _kill_own_process_group) + for signal_number in PARENT_SIGNALS: + signal.signal(signal_number, signal.SIG_IGN) + _set_parent_death_signal(LEASE_GUARD_SIGNAL, os.getppid()) + + def prepare_payload() -> None: + for signal_number in PARENT_SIGNALS: + signal.signal(signal_number, signal.SIG_DFL) + + try: + payload = subprocess.Popen(command, preexec_fn=prepare_payload) + except (OSError, ValueError) as exc: + os.write( + status_fd, + json.dumps( + {"error": getattr(exc, "strerror", None) or str(exc)} + ).encode("utf-8") + + b"\n", + ) + os.close(status_fd) + return EXIT_LAUNCH_ERROR + + os.write( + status_fd, + json.dumps({"payload_pid": payload.pid}).encode("utf-8") + b"\n", + ) + os.close(status_fd) + poller = select.poll() + poller.register( + control_fd, + select.POLLIN | select.POLLHUP | select.POLLERR, + ) + deadline = time.monotonic() + pulse_timeout_seconds + while True: + remaining = max(0.0, deadline - time.monotonic()) + events = poller.poll(max(1, min(50, int(remaining * 1000)))) + for _, event_mask in events: + if event_mask & (select.POLLHUP | select.POLLERR): + _kill_own_process_group() + try: + pulse = os.read(control_fd, 65536) + except BlockingIOError: + pulse = b"" + if not pulse: + _kill_own_process_group() + deadline = time.monotonic() + pulse_timeout_seconds + if time.monotonic() >= deadline: + _kill_own_process_group() + returncode = payload.poll() + if returncode is not None: + return ( + 128 - returncode + if returncode < 0 + else returncode + ) + + +def _read_guardian_status( + descriptor: int, timeout_seconds: float +) -> int: + poller = select.poll() + poller.register(descriptor, select.POLLIN | select.POLLHUP) + deadline = time.monotonic() + timeout_seconds + content = b"" + while time.monotonic() < deadline: + events = poller.poll( + max(1, int((deadline - time.monotonic()) * 1000)) + ) + if not events: + continue + chunk = os.read(descriptor, 4096) + if not chunk: + break + content += chunk + if b"\n" in content: + break + if not content: + raise OSError("guardian did not report payload startup") + try: + status = json.loads(content.splitlines()[0]) + except (UnicodeError, json.JSONDecodeError) as exc: + raise OSError("guardian returned malformed startup status") from exc + if not isinstance(status, dict): + raise OSError("guardian returned malformed startup status") + if "error" in status: + raise OSError(str(status["error"])) + payload_pid = status.get("payload_pid") + if not isinstance(payload_pid, int): + raise OSError("guardian did not report a payload PID") + return payload_pid + + +def _launch_guardian( + command: tuple[str, ...], + environment: dict[str, str], + pulse_timeout_seconds: float, + launch_mask: set[signal.Signals], +) -> GuardianProcess: + control_read, control_write = os.pipe() + os.set_blocking(control_read, False) + os.set_blocking(control_write, False) + status_read, status_write = os.pipe() + parent_pid = os.getpid() + + def prepare_guardian() -> None: + signal.pthread_sigmask(signal.SIG_SETMASK, launch_mask) + _set_parent_death_signal(signal.SIGKILL, parent_pid) + + guardian_command = ( + sys.executable, + str(Path(__file__).resolve()), + "--internal-guardian", + str(control_read), + str(status_write), + str(pulse_timeout_seconds), + "--", + *command, + ) + try: + process = subprocess.Popen( + guardian_command, + start_new_session=True, + pass_fds=(control_read, status_write), + preexec_fn=prepare_guardian, + env=environment, + ) + finally: + os.close(control_read) + os.close(status_write) + try: + payload_pid = _read_guardian_status(status_read, 5.0) + except OSError: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5.0) + os.close(control_write) + raise + finally: + os.close(status_read) + guardian = GuardianProcess(process, payload_pid, control_write) + guardian.pulse() + return guardian + + def _read_proc_bytes(root: Path, process_id: int, name: str) -> bytes: path = root / str(process_id) / name try: @@ -321,15 +543,6 @@ def _write_json_atomic( temp_path = parent / ( f".{path.name}.{os.getpid()}.{secrets.token_hex(8)}.tmp" ) - payload = ( - json.dumps( - value, - ensure_ascii=True, - sort_keys=True, - separators=(",", ":"), - ) - + "\n" - ).encode("utf-8") try: descriptor = os.open( temp_path, @@ -337,6 +550,23 @@ def _write_json_atomic( 0o600, ) with os.fdopen(descriptor, "wb") as stream: + file_status = os.fstat(stream.fileno()) + record = { + **value, + "file_device": file_status.st_dev, + "file_inode": file_status.st_ino, + "file_uid": file_status.st_uid, + "file_mode": stat.S_IMODE(file_status.st_mode), + } + payload = ( + json.dumps( + record, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ) + + "\n" + ).encode("utf-8") stream.write(payload) stream.flush() os.fsync(stream.fileno()) @@ -394,6 +624,11 @@ def _watchdog_identity(self) -> dict[str, object]: _, _, proc_start_time_ticks = _read_proc_stat( self.process_procfs_root, os.getpid() ) + executable_path = ( + self.process_procfs_root + / str(os.getpid()) + / "exe" + ).resolve() except (OSError, LeaseValidationError): if sys.platform.startswith("linux"): raise ArtifactError( @@ -403,11 +638,13 @@ def _watchdog_identity(self) -> dict[str, object]: cmdline = b"\0".join( os.fsencode(argument) for argument in sys.argv ) + executable_path = Path(sys.executable).resolve() return { "pid": os.getpid(), "start_time_utc": _timestamp_utc(self.wall_clock), "proc_start_time_ticks": proc_start_time_ticks, "cmdline_sha256": _sha256_bytes(cmdline), + "executable_path": str(executable_path), "script_path": str(script_path), "script_sha256": _sha256_file(script_path), } @@ -440,8 +677,16 @@ def _heartbeat_record( record["sample"] = sample return record - def start(self, child: ProcessHandle) -> None: + def start( + self, child: ProcessHandle, audit: AuditLogger + ) -> None: watchdog_identity = self._watchdog_identity() + audit_identity = audit.persistent_identity() + payload_pid = ( + child.payload_pid + if isinstance(child, GuardianProcess) + else child.pid + ) self.lease = { "format": LEASE_FORMAT, "version": LEASE_VERSION, @@ -457,6 +702,9 @@ def start(self, child: ProcessHandle) -> None: "watchdog_command_sha256": ( watchdog_identity["cmdline_sha256"] ), + "watchdog_executable_path": ( + watchdog_identity["executable_path"] + ), "watchdog_script_path": watchdog_identity["script_path"], "watchdog_script_sha256": ( watchdog_identity["script_sha256"] @@ -464,7 +712,8 @@ def start(self, child: ProcessHandle) -> None: "soft_bytes": self.config.soft_bytes, "emergency_bytes": self.config.emergency_bytes, "strict_ceiling_bytes": STRICT_CEILING_BYTES, - "child_pid": child.pid, + "guardian_pid": child.pid, + "child_pid": payload_pid, "child_process_group_id": child.pid, "command": list(self.config.command), "child_command_sha256": _command_sha256( @@ -475,11 +724,19 @@ def start(self, child: ProcessHandle) -> None: self.config.heartbeat_max_age_seconds ), "audit_path": str(self.audit_path), + "audit_device": audit_identity["device"], + "audit_inode": audit_identity["inode"], + "audit_uid": audit_identity["uid"], + "audit_mode": audit_identity["mode"], + "audit_fd": audit_identity["fd"], "procfs_root": str( self.config.procfs_root.expanduser().resolve() ), } - heartbeat = self._heartbeat_record("active") + heartbeat = self._heartbeat_record( + "active", + {"audit_record_sha256": audit.last_record_sha256}, + ) _write_json_atomic(self.heartbeat_path, heartbeat, create=True) _write_json_atomic(self.lease_path, self.lease, create=True) @@ -499,13 +756,34 @@ def finalize(self, final_record: dict[str, object]) -> None: def _read_json_object(path: Path) -> dict[str, object]: try: - value = json.loads(path.read_text(encoding="utf-8")) + descriptor = os.open( + path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + with os.fdopen(descriptor, "r", encoding="utf-8") as stream: + file_status = os.fstat(stream.fileno()) + if ( + not stat.S_ISREG(file_status.st_mode) + or file_status.st_uid != os.getuid() + or stat.S_IMODE(file_status.st_mode) != 0o600 + ): + raise LeaseValidationError( + f"{path} has unsafe type, owner, or mode" + ) + value = json.load(stream) except (OSError, UnicodeError, json.JSONDecodeError) as exc: raise LeaseValidationError( f"cannot read valid JSON from {path}: {exc}" ) from exc if not isinstance(value, dict): raise LeaseValidationError(f"{path} must contain a JSON object") + if ( + value.get("file_device") != file_status.st_dev + or value.get("file_inode") != file_status.st_ino + or value.get("file_uid") != file_status.st_uid + or value.get("file_mode") != stat.S_IMODE(file_status.st_mode) + ): + raise LeaseValidationError(f"{path} identity does not match") return value @@ -525,16 +803,20 @@ def validate_active_lease( lease_path: Path, *, expected_script_path: Path, + expected_executable_path: Path | None = None, expected_soft_bytes: int = DEFAULT_SOFT_BYTES, expected_emergency_bytes: int = DEFAULT_EMERGENCY_BYTES, expected_procfs_root: Path = Path("/proc"), - expected_command_sha256: str | None = None, + expected_command: Sequence[str] | None = None, expected_heartbeat_path: Path | None = None, expected_audit_path: Path | None = None, expected_max_heartbeat_age_seconds: float | None = None, current_process_id: int | None = None, process_procfs_root: Path = Path("/proc"), monotonic_ns: Callable[[], int] | None = None, + pidfd_open: Callable[[int], int] | None = getattr( + os, "pidfd_open", None + ), ) -> dict[str, object]: lease_path = lease_path.expanduser().resolve() lease = _read_json_object(lease_path) @@ -587,6 +869,14 @@ def validate_active_lease( watchdog_pid = _require_int( lease.get("watchdog_pid"), "watchdog_pid" ) + pidfd: int | None = None + if pidfd_open is not None: + try: + pidfd = pidfd_open(watchdog_pid) + except OSError as exc: + raise LeaseValidationError( + "cannot open watchdog pidfd" + ) from exc watchdog_start_ticks = _require_int( lease.get("watchdog_start_time_ticks"), "watchdog_start_time_ticks", @@ -596,6 +886,28 @@ def validate_active_lease( ) if live_watchdog_start_ticks != watchdog_start_ticks: raise LeaseValidationError("watchdog process start time does not match") + expected_executable = ( + expected_executable_path or Path(sys.executable) + ).expanduser().resolve() + try: + live_executable = ( + process_procfs_root / str(watchdog_pid) / "exe" + ).resolve() + except OSError as exc: + raise LeaseValidationError( + "cannot resolve watchdog executable" + ) from exc + if ( + live_executable != expected_executable + or Path( + _require_string( + lease.get("watchdog_executable_path"), + "watchdog_executable_path", + ) + ).resolve() + != expected_executable + ): + raise LeaseValidationError("watchdog executable does not match") live_cmdline = _read_proc_bytes( process_procfs_root, watchdog_pid, "cmdline" ) @@ -604,48 +916,91 @@ def validate_active_lease( "watchdog_command_sha256", ): raise LeaseValidationError("watchdog command line does not match") - script_named = False - watchdog_cwd: Path | None = None - for raw_argument in live_cmdline.split(b"\0"): - if not raw_argument: - continue - argument_path = Path(os.fsdecode(raw_argument)).expanduser() - if not argument_path.is_absolute(): - if watchdog_cwd is None: - try: - watchdog_cwd = ( - process_procfs_root - / str(watchdog_pid) - / "cwd" - ).resolve() - except OSError as exc: - raise LeaseValidationError( - "cannot resolve watchdog working directory" - ) from exc - argument_path = watchdog_cwd / argument_path - if argument_path.resolve() == expected_script_path: - script_named = True - break - if not script_named: + argv = [ + os.fsdecode(argument) + for argument in live_cmdline.split(b"\0") + if argument + ] + if len(argv) < 2 or argv[1] in ("-c", "-m"): raise LeaseValidationError( - "watchdog command line does not name the expected script" + "watchdog script is not in executable argv position" ) + try: + watchdog_cwd = ( + process_procfs_root / str(watchdog_pid) / "cwd" + ).resolve() + except OSError as exc: + raise LeaseValidationError( + "cannot resolve watchdog working directory" + ) from exc + argv_script = Path(argv[1]).expanduser() + if not argv_script.is_absolute(): + argv_script = watchdog_cwd / argv_script + if argv_script.resolve() != expected_script_path: + raise LeaseValidationError( + "watchdog script is not in executable argv position" + ) + try: + live_config = parse_args(argv[2:]) + live_paths = live_config.validate() + except (SystemExit, ValueError) as exc: + raise LeaseValidationError( + "watchdog command line is invalid" + ) from exc + if ( + live_config.soft_bytes != expected_soft_bytes + or live_config.emergency_bytes != expected_emergency_bytes + or live_config.procfs_root.expanduser().resolve() + != expected_procfs_root.expanduser().resolve() + ): + raise LeaseValidationError( + "watchdog command-line policy does not match" + ) + if ( + live_paths is None + or live_paths.lease != lease_path + or ( + expected_heartbeat_path is not None + and live_paths.heartbeat + != expected_heartbeat_path.expanduser().resolve() + ) + or ( + expected_audit_path is not None + and live_paths.audit + != expected_audit_path.expanduser().resolve() + ) + ): + raise LeaseValidationError( + "watchdog command-line artifact paths do not match" + ) + if expected_command is not None and tuple( + expected_command + ) != live_config.command: + raise LeaseValidationError("monitored command does not match") + guardian_pid = _require_int( + lease.get("guardian_pid"), "guardian_pid" + ) child_pid = _require_int(lease.get("child_pid"), "child_pid") process_group_id = _require_int( lease.get("child_process_group_id"), "child_process_group_id", ) + guardian_parent_pid, guardian_group_id, _ = _read_proc_stat( + process_procfs_root, guardian_pid + ) child_parent_pid, child_group_id, _ = _read_proc_stat( process_procfs_root, child_pid ) if ( - child_parent_pid != watchdog_pid + guardian_parent_pid != watchdog_pid + or guardian_group_id != process_group_id + or guardian_pid != process_group_id + or child_parent_pid != guardian_pid or child_group_id != process_group_id - or process_group_id != child_pid ): raise LeaseValidationError( - "monitored child parent or process group does not match" + "watchdog, guardian, child, or process group does not match" ) command = lease.get("command") if ( @@ -660,9 +1015,8 @@ def validate_active_lease( ) if command_sha256 != _command_sha256(command): raise LeaseValidationError("monitored command SHA is invalid") - if ( - expected_command_sha256 is not None - and command_sha256 != expected_command_sha256 + if expected_command is not None and command_sha256 != _command_sha256( + expected_command ): raise LeaseValidationError("monitored command SHA does not match") @@ -725,6 +1079,13 @@ def validate_active_lease( ) _require_int(heartbeat.get("sequence"), "heartbeat.sequence") _require_string(heartbeat.get("updated_at"), "heartbeat.updated_at") + heartbeat_sample = heartbeat.get("sample") + if not isinstance(heartbeat_sample, dict): + raise LeaseValidationError("heartbeat sample is invalid") + audit_record_sha256 = _require_string( + heartbeat_sample.get("audit_record_sha256"), + "heartbeat.sample.audit_record_sha256", + ) now_monotonic_ns = (monotonic_ns or time.monotonic_ns)() age_ns = now_monotonic_ns - updated_monotonic_ns if age_ns < 0 or age_ns > int(heartbeat_max_age * 1_000_000_000): @@ -738,14 +1099,64 @@ def validate_active_lease( and audit_path != expected_audit_path.expanduser().resolve() ): raise LeaseValidationError("persistent audit path does not match") + audit_fd = _require_int(lease.get("audit_fd"), "audit_fd") + audit_device = _require_int( + lease.get("audit_device"), "audit_device" + ) + audit_inode = _require_int( + lease.get("audit_inode"), "audit_inode" + ) + audit_uid = _require_int(lease.get("audit_uid"), "audit_uid") + audit_mode = _require_int(lease.get("audit_mode"), "audit_mode") try: - first_line = next( + audit_status = audit_path.stat(follow_symlinks=False) + live_audit_status = ( + process_procfs_root + / str(watchdog_pid) + / "fd" + / str(audit_fd) + ).stat() + if ( + not stat.S_ISREG(audit_status.st_mode) + or audit_status.st_dev != audit_device + or audit_status.st_ino != audit_inode + or live_audit_status.st_dev != audit_device + or live_audit_status.st_ino != audit_inode + or audit_status.st_uid != audit_uid + or audit_uid != os.getuid() + or stat.S_IMODE(audit_status.st_mode) != audit_mode + or audit_mode != 0o600 + ): + raise LeaseValidationError( + "persistent audit identity does not match" + ) + audit_descriptor = os.open( + audit_path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + try: + try: + fcntl.flock( + audit_descriptor, + fcntl.LOCK_EX | fcntl.LOCK_NB, + ) + except BlockingIOError: + pass + else: + fcntl.flock(audit_descriptor, fcntl.LOCK_UN) + raise LeaseValidationError( + "watchdog does not hold the persistent audit lock" + ) + finally: + os.close(audit_descriptor) + audit_lines = [ line for line in audit_path.read_text( encoding="utf-8" ).splitlines() if line - ) + ] + first_line = next(iter(audit_lines)) first_record = json.loads(first_line) if ( not isinstance(first_record, dict) @@ -755,6 +1166,14 @@ def validate_active_lease( raise LeaseValidationError( "persistent audit does not contain watchdog records" ) + if not any( + _sha256_bytes((line + "\n").encode("utf-8")) + == audit_record_sha256 + for line in audit_lines + ): + raise LeaseValidationError( + "heartbeat audit record does not match persistent audit" + ) except StopIteration as exc: raise LeaseValidationError("persistent audit is empty") from exc except (UnicodeError, json.JSONDecodeError) as exc: @@ -767,9 +1186,104 @@ def validate_active_lease( raise LeaseValidationError( f"cannot inspect persistent audit {audit_path}: {exc}" ) from exc + _, _, final_watchdog_start_ticks = _read_proc_stat( + process_procfs_root, watchdog_pid + ) + if final_watchdog_start_ticks != watchdog_start_ticks: + raise LeaseValidationError( + "watchdog process changed during validation" + ) + if pidfd is not None: + os.close(pidfd) return lease +def start_process_group_lease_guard( + expected_script_path: Path, + *, + startup_timeout_seconds: float = 5.0, + expected_procfs_root: Path = Path("/proc"), + process_procfs_root: Path = Path("/proc"), +) -> threading.Thread: + try: + lease_path = Path( + os.environ["STRIX_MEMORY_WATCHDOG_LEASE_PATH"] + ).resolve() + heartbeat_path = Path( + os.environ["STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH"] + ).resolve() + audit_path = Path( + os.environ["STRIX_MEMORY_WATCHDOG_AUDIT_PATH"] + ).resolve() + max_age_seconds = float( + os.environ[ + "STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS" + ] + ) + except (KeyError, ValueError) as exc: + raise LeaseValidationError( + "watchdog artifact environment is missing or invalid" + ) from exc + current_cmdline = _read_proc_bytes( + process_procfs_root, os.getpid(), "cmdline" + ) + expected_command = tuple( + os.fsdecode(argument) + for argument in current_cmdline.split(b"\0") + if argument + ) + deadline = time.monotonic() + startup_timeout_seconds + while True: + try: + lease = validate_active_lease( + lease_path, + expected_script_path=expected_script_path, + expected_procfs_root=expected_procfs_root, + expected_command=expected_command, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + expected_max_heartbeat_age_seconds=max_age_seconds, + process_procfs_root=process_procfs_root, + ) + break + except LeaseValidationError: + if time.monotonic() >= deadline: + raise + time.sleep(0.01) + + guardian_pid = _require_int( + lease.get("guardian_pid"), "guardian_pid" + ) + signal.signal(LEASE_GUARD_SIGNAL, _kill_own_process_group) + _set_parent_death_signal(LEASE_GUARD_SIGNAL, guardian_pid) + + def monitor() -> None: + interval = min(1.0, max_age_seconds / 3) + while True: + time.sleep(interval) + try: + validate_active_lease( + lease_path, + expected_script_path=expected_script_path, + expected_procfs_root=expected_procfs_root, + expected_command=expected_command, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + expected_max_heartbeat_age_seconds=max_age_seconds, + process_procfs_root=process_procfs_root, + ) + except LeaseValidationError: + _kill_own_process_group() + + guard = threading.Thread( + target=monitor, + name="strix-watchdog-lease-guard", + daemon=True, + ) + guard.start() + return guard + + class AuditLogger: def __init__( self, @@ -782,15 +1296,22 @@ def __init__( self.lease_manager: LeaseManager | None = None self.finalized = False self.final_exit_code = EXIT_INTERNAL_ERROR + self.last_record_sha256: str | None = None def open_persistent(self, path: Path) -> None: resolved_path = path.expanduser().resolve() try: descriptor = os.open( resolved_path, - os.O_CREAT | os.O_EXCL | os.O_WRONLY, + os.O_CREAT + | os.O_EXCL + | os.O_WRONLY + | getattr(os, "O_NOFOLLOW", 0), 0o600, ) + fcntl.flock( + descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB + ) self.persistent_stream = os.fdopen( descriptor, "w", encoding="utf-8" ) @@ -801,6 +1322,20 @@ def open_persistent(self, path: Path) -> None: f"cannot create persistent audit {resolved_path}: {detail}", ) from exc + def persistent_identity(self) -> dict[str, int]: + if self.persistent_stream is None: + raise ArtifactError( + "audit", "persistent audit is not open" + ) + file_status = os.fstat(self.persistent_stream.fileno()) + return { + "device": file_status.st_dev, + "inode": file_status.st_ino, + "uid": file_status.st_uid, + "mode": stat.S_IMODE(file_status.st_mode), + "fd": self.persistent_stream.fileno(), + } + def close(self) -> None: if self.persistent_stream is not None: self.persistent_stream.close() @@ -824,6 +1359,7 @@ def emit(self, event: str, **fields: object) -> dict[str, object]: ) self.stream.write(line) self.stream.flush() + self.last_record_sha256 = _sha256_bytes(line.encode("utf-8")) if self.persistent_stream is not None: try: self.persistent_stream.write(line) @@ -839,7 +1375,12 @@ def emit(self, event: str, **fields: object) -> dict[str, object]: def heartbeat(self, sample: dict[str, object]) -> None: if self.lease_manager is not None: - self.lease_manager.update_heartbeat(sample) + self.lease_manager.update_heartbeat( + { + **sample, + "audit_record_sha256": self.last_record_sha256, + } + ) def finalize(self, record: dict[str, object]) -> None: if self.lease_manager is not None: @@ -952,6 +1493,30 @@ def _signal_process_group(process_group_id: int, signal_number: int) -> str: def _process_group_alive(process_group_id: int) -> bool: + if sys.platform.startswith("linux"): + try: + process_paths = Path("/proc").iterdir() + for process_path in process_paths: + if not process_path.name.isdigit(): + continue + try: + content = ( + process_path / "stat" + ).read_text(encoding="utf-8") + close_paren = content.rfind(")") + fields = content[close_paren + 1:].split() + if ( + close_paren >= 0 + and len(fields) >= 3 + and fields[0] != "Z" + and int(fields[2]) == process_group_id + ): + return True + except (OSError, UnicodeError, ValueError): + continue + return False + except OSError: + pass try: os.killpg(process_group_id, 0) except ProcessLookupError: @@ -1012,18 +1577,23 @@ def _kill_and_finish( str(exc), ) - audit.emit( - "process_group_signal", - **_state_fields( - snapshot, - peak_used_bytes, - child, - child.poll(), - group_status, - reason, - ), - signal="SIGKILL", - ) + artifact_error: ArtifactError | None = None + try: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + group_status, + reason, + ), + signal="SIGKILL", + ) + except ArtifactError as exc: + artifact_error = exc + audit.disable_component(exc.component) try: child_returncode = child.wait(timeout=5.0) except subprocess.TimeoutExpired as exc: @@ -1039,6 +1609,12 @@ def _kill_and_finish( "sigkill_timeout", str(exc), ) + if artifact_error is not None: + classification = "lease_error" + exit_code = EXIT_LEASE_ERROR + error = f"{artifact_error.component}: {artifact_error}" + else: + error = None return _emit_final( audit, classification, @@ -1049,6 +1625,7 @@ def _kill_and_finish( child, child_returncode, group_status, + error, ) @@ -1071,23 +1648,28 @@ def _graceful_cleanup( error: str | None = None, ) -> int: escalated = False + artifact_error: ArtifactError | None = None try: if graceful_signal is not None: process_group_status = signal_group( child.pid, graceful_signal ) - audit.emit( - "process_group_signal", - **_state_fields( - snapshot, - peak_used_bytes, - child, - child.poll(), - process_group_status, - reason, - ), - signal=signal.Signals(graceful_signal).name, - ) + try: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + reason, + ), + signal=signal.Signals(graceful_signal).name, + ) + except ArtifactError as exc: + artifact_error = exc + audit.disable_component(exc.component) deadline = monotonic() + grace_seconds while monotonic() < deadline: child.poll() @@ -1105,18 +1687,23 @@ def _graceful_cleanup( if escalation_result is not None else reason ) - audit.emit( - "process_group_signal", - **_state_fields( - snapshot, - peak_used_bytes, - child, - child.poll(), - process_group_status, - signal_reason, - ), - signal="SIGKILL", - ) + try: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + signal_reason, + ), + signal="SIGKILL", + ) + except ArtifactError as exc: + if artifact_error is None: + artifact_error = exc + audit.disable_component(exc.component) except ProcessGroupError as exc: return _emit_final( audit, @@ -1151,6 +1738,10 @@ def _graceful_cleanup( if escalated and escalation_result is not None: classification, exit_code, reason = escalation_result + if artifact_error is not None: + classification = "lease_error" + exit_code = EXIT_LEASE_ERROR + error = f"{artifact_error.component}: {artifact_error}" return _emit_final( audit, classification, @@ -1173,6 +1764,7 @@ def _monitor_child( state: RuntimeState, signal_group: Callable[[int, int], str], group_alive: Callable[[int], bool], + pulse_guardian: Callable[[], None], monotonic: Callable[[], float], sleeper: Callable[[float], None], ) -> int: @@ -1273,19 +1865,6 @@ def _monitor_child( state.peak_used_bytes = max( state.peak_used_bytes, state.snapshot.used_bytes ) - sample_record = audit.emit( - "sample", - **_state_fields( - state.snapshot, - state.peak_used_bytes, - child, - None, - "active", - "none", - ), - ) - audit.heartbeat(sample_record) - if state.snapshot.active_swaps: return _kill_and_finish( audit, @@ -1308,6 +1887,7 @@ def _monitor_child( "used_bytes >= emergency_bytes", signal_group, ) + soft_signal_fields: dict[str, object] | None = None if ( soft_deadline is None and state.snapshot.used_bytes >= config.soft_bytes @@ -1328,8 +1908,7 @@ def _monitor_child( str(exc), ) soft_deadline = now + config.grace_seconds - audit.emit( - "process_group_signal", + soft_signal_fields = { **_state_fields( state.snapshot, state.peak_used_bytes, @@ -1338,9 +1917,39 @@ def _monitor_child( group_status, "used_bytes >= soft_bytes", ), - signal="SIGTERM", - grace_deadline_monotonic=soft_deadline, + "signal": "SIGTERM", + "grace_deadline_monotonic": soft_deadline, + } + try: + pulse_guardian() + except ProcessGroupError as exc: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "signal_error", + EXIT_SIGNAL_ERROR, + str(exc), + signal_group, + ) + if soft_signal_fields is not None: + audit.emit( + "process_group_signal", + **soft_signal_fields, ) + sample_record = audit.emit( + "sample", + **_state_fields( + state.snapshot, + state.peak_used_bytes, + child, + None, + "active", + "none", + ) + ) + audit.heartbeat(sample_record) sleep_seconds = config.sample_interval_seconds if soft_deadline is not None: @@ -1363,6 +1972,9 @@ def run_watchdog( sleeper: Callable[[float], None] | None = None, ) -> int: artifact_paths = config.validate() + use_guardian = ( + launcher is None and sys.platform.startswith("linux") + ) reader = reader or ProcfsReader(config.procfs_root) audit = audit or AuditLogger(sys.stderr) launcher = launcher or subprocess.Popen @@ -1472,8 +2084,8 @@ def restore_child_signal_mask() -> None: signal.pthread_sigmask(signal.SIG_SETMASK, launch_mask) try: + child_environment = os.environ.copy() if lease_manager is not None: - child_environment = os.environ.copy() child_environment.update( { "STRIX_MEMORY_WATCHDOG_LEASE_PATH": str( @@ -1490,6 +2102,14 @@ def restore_child_signal_mask() -> None: ), } ) + if use_guardian: + child = _launch_guardian( + config.command, + child_environment, + config.heartbeat_max_age_seconds, + launch_mask, + ) + elif lease_manager is not None: child = launcher( config.command, start_new_session=True, @@ -1502,7 +2122,7 @@ def restore_child_signal_mask() -> None: start_new_session=True, preexec_fn=restore_child_signal_mask, ) - except (OSError, ValueError) as exc: + except (OSError, ValueError, subprocess.SubprocessError) as exc: detail = getattr(exc, "strerror", None) or str(exc) return _emit_final( audit, @@ -1518,7 +2138,7 @@ def restore_child_signal_mask() -> None: _raise_parent_signal ) if lease_manager is not None: - lease_manager.start(child) + lease_manager.start(child, audit) audit.lease_manager = lease_manager signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) mask_restored = True @@ -1542,6 +2162,7 @@ def restore_child_signal_mask() -> None: state, signal_group, group_alive, + child.pulse if isinstance(child, GuardianProcess) else lambda: None, monotonic, sleeper, ) @@ -1631,6 +2252,8 @@ def restore_child_signal_mask() -> None: signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) if previous_handlers: _restore_parent_signal_handlers(previous_handlers) + if isinstance(child, GuardianProcess): + child.close() def _positive_int(value: str) -> int: @@ -1744,7 +2367,20 @@ def parse_args(argv: Sequence[str]) -> WatchdogConfig: def main(argv: Sequence[str] | None = None) -> int: - config = parse_args(argv if argv is not None else sys.argv[1:]) + arguments = tuple(argv if argv is not None else sys.argv[1:]) + if arguments and arguments[0] == "--internal-guardian": + if len(arguments) < 6 or arguments[4] != "--": + return EXIT_LAUNCH_ERROR + try: + return _guardian_main( + int(arguments[1]), + int(arguments[2]), + _positive_float(arguments[3]), + tuple(arguments[5:]), + ) + except (OSError, ValueError): + return EXIT_LAUNCH_ERROR + config = parse_args(arguments) audit = AuditLogger(sys.stderr) try: return run_watchdog(config, audit=audit) diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 81c10f84c38f..a61078efb2a3 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -3,6 +3,7 @@ from __future__ import annotations import importlib.util +import fcntl import hashlib import io import json @@ -400,6 +401,199 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: "parent_signal", ) + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_pipe_close_kills_group_without_fd_leak(self) -> None: + child_code = ( + "import json,os,subprocess,sys,time\n" + "targets=[]\n" + "for name in os.listdir('/proc/self/fd'):\n" + " try: targets.append(os.readlink('/proc/self/fd/'+name))\n" + " except OSError: pass\n" + "grandchild=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'])\n" + "open(sys.argv[1],'w').write(json.dumps({" + "'child':os.getpid(),'grandchild':grandchild.pid," + "'fds':targets}))\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + state_path = root / "state.json" + guardian = watchdog._launch_guardian( + ( + sys.executable, + "-c", + child_code, + str(state_path), + ), + os.environ.copy(), + 0.5, + signal.pthread_sigmask(signal.SIG_BLOCK, ()), + ) + control_target = os.readlink( + f"/proc/self/fd/{guardian.pulse_fd}" + ) + deadline = time.monotonic() + 5 + while not state_path.exists(): + if time.monotonic() >= deadline: + self.fail("guardian payload did not become ready") + time.sleep(0.01) + state = json.loads(state_path.read_text(encoding="utf-8")) + os.close(guardian.pulse_fd) + guardian.wait(timeout=5) + + self.assertNotIn(control_target, state["fds"]) + for process_id in (state["child"], state["grandchild"]): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_documents_setsid_escape_limit(self) -> None: + child_code = ( + "import os,subprocess,sys,time\n" + "escaped=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'],start_new_session=True)\n" + "open(sys.argv[1],'w').write(str(escaped.pid))\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + state_path = Path(temp_dir) / "escaped-pid" + guardian = watchdog._launch_guardian( + ( + sys.executable, + "-c", + child_code, + str(state_path), + ), + os.environ.copy(), + 0.5, + signal.pthread_sigmask(signal.SIG_BLOCK, ()), + ) + deadline = time.monotonic() + 5 + while not state_path.exists(): + if time.monotonic() >= deadline: + self.fail("escaped payload did not become ready") + time.sleep(0.01) + escaped_pid = int( + state_path.read_text(encoding="utf-8") + ) + os.close(guardian.pulse_fd) + guardian.wait(timeout=5) + self.assertTrue(self._process_is_running(escaped_pid)) + os.kill(escaped_pid, signal.SIGKILL) + deadline = time.monotonic() + 2 + while ( + self._process_is_running(escaped_pid) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(escaped_pid)) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guard_kills_group_after_watchdog_loss_or_stall(self) -> None: + child_code = ( + "import importlib.util,os,pathlib,subprocess,sys,time\n" + "script=pathlib.Path(sys.argv[1])\n" + "spec=importlib.util.spec_from_file_location('guard_watchdog',script)\n" + "module=importlib.util.module_from_spec(spec)\n" + "sys.modules[spec.name]=module\n" + "spec.loader.exec_module(module)\n" + "module.start_process_group_lease_guard(" + "script,expected_procfs_root=pathlib.Path(sys.argv[2]))\n" + "grandchild=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'])\n" + "open(sys.argv[3],'w').write(" + "f'{os.getpid()} {grandchild.pid}\\n')\n" + "time.sleep(30)\n" + ) + for mode in ("sigkill", "sigstop"): + with self.subTest(mode=mode): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_path = root / "pids" + self._write_procfs_fixture(root) + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--heartbeat-max-age-seconds", + "0.3", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(SCRIPT_PATH), + str(root), + str(pid_path), + ], + stderr=subprocess.PIPE, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_path.exists(): + if wrapper.poll() is not None: + assert wrapper.stderr is not None + self.fail(wrapper.stderr.read()) + if time.monotonic() >= deadline: + self.fail( + "guarded payload did not become ready" + ) + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_path.read_text( + encoding="utf-8" + ).split() + ) + if mode == "sigkill": + wrapper.kill() + else: + os.kill(wrapper.pid, signal.SIGSTOP) + heartbeat_path = root / "heartbeat.json" + heartbeat = json.loads( + heartbeat_path.read_text(encoding="utf-8") + ) + heartbeat["updated_monotonic_ns"] = ( + time.monotonic_ns() + ) + watchdog._write_json_atomic( + heartbeat_path, heartbeat + ) + time.sleep(0.7) + os.kill(wrapper.pid, signal.SIGCONT) + wrapper.wait(timeout=5) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse( + self._process_is_running(process_id) + ) + if wrapper.stderr is not None: + wrapper.stderr.close() + def test_child_sigterm_handler_exits_without_escalation(self) -> None: child_code = ( "import os,signal,sys,time\n" @@ -713,6 +907,89 @@ def write(self, value: str) -> int: audit.final_exit_code, watchdog.EXIT_LEASE_ERROR ) + def test_emergency_signal_precedes_artifact_write(self) -> None: + events: list[str] = [] + + class BlockingAudit(watchdog.AuditLogger): + def __init__(self) -> None: + super().__init__(io.StringIO()) + self.calls = 0 + + def emit( + self, event: str, **fields: object + ) -> dict[str, object]: + self.calls += 1 + events.append(f"audit:{event}") + if self.calls == 3: + raise watchdog.ArtifactError( + "audit", "simulated blocked fsync" + ) + return super().emit(event, **fields) + + def exit_on_kill( + process: FakeProcess, signal_number: int + ) -> None: + events.append(f"signal:{signal_number}") + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(160)], + FakeProcess(), + signal_handler=exit_on_kill, + ) + harness.audit = BlockingAudit() + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(events[2], f"signal:{signal.SIGKILL}") + self.assertEqual(events[3], "audit:process_group_signal") + self.assertEqual(harness.process.returncode, -signal.SIGKILL) + + def test_cleanup_reaps_after_persistent_audit_failure(self) -> None: + class FailingSignalAudit(watchdog.AuditLogger): + def emit( + self, event: str, **fields: object + ) -> dict[str, object]: + if event == "process_group_signal": + raise watchdog.ArtifactError( + "audit", "simulated persistent write failure" + ) + return super().emit(event, **fields) + + process = FakeProcess() + signals: list[int] = [] + clock = FakeClock() + + def signal_group( + process_group_id: int, signal_number: int + ) -> str: + signals.append(signal_number) + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + return f"{signal.Signals(signal_number).name.lower()}_sent" + + result = watchdog._graceful_cleanup( + FailingSignalAudit(io.StringIO()), + process, + snapshot(50), + 50, + "parent_signal", + 128 + signal.SIGTERM, + "wrapper received SIGTERM", + signal.SIGTERM, + 0.1, + signal_group, + lambda _process_group_id: process.returncode is None, + clock.monotonic, + clock.sleep, + ) + + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(signals, [signal.SIGTERM, signal.SIGKILL]) + self.assertEqual(process.returncode, -signal.SIGKILL) + def test_invalid_artifact_path_emits_configuration_final(self) -> None: result = subprocess.run( [ @@ -881,18 +1158,36 @@ def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: root = Path(temp_dir) process_root = root / "proc" watchdog_pid = 1200 + guardian_pid = 1250 child_pid = 1300 current_pid = 1400 watchdog_start_ticks = 456789 script_path = root / "watchdog.py" script_path.write_text("print('watchdog')\n", encoding="utf-8") - cmdline = ( - b"/usr/bin/python3\0watchdog.py\0--lease-path\0" - ) + lease_path = root / "lease.json" + heartbeat_path = root / "heartbeat.json" + audit_path = root / "audit.jsonl" + command = [sys.executable, "run_matrix.py"] + argv = [ + sys.executable, + "watchdog.py", + "--procfs-root", + "/proc", + "--lease-path", + str(lease_path), + "--heartbeat-path", + str(heartbeat_path), + "--audit-path", + str(audit_path), + "--", + *command, + ] + cmdline = b"\0".join(os.fsencode(value) for value in argv) for process_id, parent_id, group_id, start_ticks in ( (watchdog_pid, 1, watchdog_pid, watchdog_start_ticks), - (child_pid, watchdog_pid, child_pid, 456790), - (current_pid, child_pid, child_pid, 456791), + (guardian_pid, watchdog_pid, guardian_pid, 456790), + (child_pid, guardian_pid, guardian_pid, 456791), + (current_pid, child_pid, guardian_pid, 456792), ): process_dir = process_root / str(process_id) process_dir.mkdir(parents=True) @@ -908,18 +1203,31 @@ def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: (process_root / str(watchdog_pid) / "cwd").symlink_to( root, target_is_directory=True ) + (process_root / str(watchdog_pid) / "exe").symlink_to( + Path(sys.executable).resolve() + ) (process_root / str(watchdog_pid) / "cmdline").write_bytes( cmdline ) - lease_path = root / "lease.json" - heartbeat_path = root / "heartbeat.json" - audit_path = root / "audit.jsonl" - audit_path.write_text( - '{"event":"child_started","timestamp":"2026-01-01T00:00:00Z"}\n', - encoding="utf-8", + audit_line = ( + '{"event":"child_started",' + '"timestamp":"2026-01-01T00:00:00Z"}\n' + ) + audit_descriptor = os.open( + audit_path, + os.O_CREAT | os.O_EXCL | os.O_RDWR, + 0o600, ) - command = ["python3", "run_matrix.py"] + os.write(audit_descriptor, audit_line.encode("utf-8")) + os.fsync(audit_descriptor) + fcntl.flock( + audit_descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB + ) + audit_status = os.fstat(audit_descriptor) + fd_root = process_root / str(watchdog_pid) / "fd" + fd_root.mkdir() + (fd_root / "9").symlink_to(audit_path) lease = { "format": watchdog.LEASE_FORMAT, "version": watchdog.LEASE_VERSION, @@ -931,6 +1239,9 @@ def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: "watchdog_command_sha256": hashlib.sha256( cmdline ).hexdigest(), + "watchdog_executable_path": str( + Path(sys.executable).resolve() + ), "watchdog_script_path": str(script_path), "watchdog_script_sha256": hashlib.sha256( script_path.read_bytes() @@ -938,8 +1249,9 @@ def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: "soft_bytes": watchdog.DEFAULT_SOFT_BYTES, "emergency_bytes": watchdog.DEFAULT_EMERGENCY_BYTES, "strict_ceiling_bytes": watchdog.STRICT_CEILING_BYTES, + "guardian_pid": guardian_pid, "child_pid": child_pid, - "child_process_group_id": child_pid, + "child_process_group_id": guardian_pid, "command": command, "child_command_sha256": watchdog._command_sha256( command @@ -947,6 +1259,11 @@ def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: "heartbeat_path": str(heartbeat_path), "max_heartbeat_age_seconds": 5.0, "audit_path": str(audit_path), + "audit_device": audit_status.st_dev, + "audit_inode": audit_status.st_ino, + "audit_uid": audit_status.st_uid, + "audit_mode": 0o600, + "audit_fd": 9, "procfs_root": "/proc", } heartbeat = { @@ -962,115 +1279,262 @@ def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: watchdog_start_ticks ), "child_pid": child_pid, - "child_process_group_id": child_pid, + "child_process_group_id": guardian_pid, + "sample": { + "audit_record_sha256": hashlib.sha256( + audit_line.encode("utf-8") + ).hexdigest() + }, } - lease_path.write_text(json.dumps(lease), encoding="utf-8") - heartbeat_path.write_text( - json.dumps(heartbeat), encoding="utf-8" + watchdog._write_json_atomic( + lease_path, lease, create=True ) - - validated = watchdog.validate_active_lease( - lease_path, - expected_script_path=script_path, - expected_command_sha256=watchdog._command_sha256( - command - ), - expected_heartbeat_path=heartbeat_path, - expected_audit_path=audit_path, - expected_max_heartbeat_age_seconds=5.0, - current_process_id=current_pid, - process_procfs_root=process_root, - monotonic_ns=lambda: 10_000_000_000, + watchdog._write_json_atomic( + heartbeat_path, heartbeat, create=True ) - self.assertEqual(validated["lease_id"], "test-lease") - with self.subTest("tampered script SHA"): - tampered = dict(lease) - tampered["watchdog_script_sha256"] = "0" * 64 - lease_path.write_text( - json.dumps(tampered), encoding="utf-8" + try: + validation_args = { + "expected_script_path": script_path, + "expected_executable_path": Path(sys.executable), + "expected_command": command, + "expected_heartbeat_path": heartbeat_path, + "expected_audit_path": audit_path, + "expected_max_heartbeat_age_seconds": 5.0, + "current_process_id": current_pid, + "process_procfs_root": process_root, + "monotonic_ns": lambda: 10_000_000_000, + "pidfd_open": lambda _pid: os.open( + os.devnull, os.O_RDONLY + ), + } + validated = watchdog.validate_active_lease( + lease_path, **validation_args ) - with self.assertRaisesRegex( - watchdog.LeaseValidationError, "script SHA" + self.assertEqual(validated["lease_id"], "test-lease") + + def publish_lease( + value: dict[str, Any], + process_argv: list[str] = argv, + ) -> None: + process_cmdline = b"\0".join( + os.fsencode(argument) + for argument in process_argv + ) + (process_root / str(watchdog_pid) / "cmdline").write_bytes( + process_cmdline + ) + value["watchdog_command_sha256"] = hashlib.sha256( + process_cmdline + ).hexdigest() + watchdog._write_json_atomic(lease_path, value) + + for name, bad_argv in ( + ( + "helper inert argument", + [ + sys.executable, + "helper.py", + str(script_path), + *argv[2:], + ], + ), + ( + "python command string", + [ + sys.executable, + "-c", + "pass", + str(script_path), + *argv[2:], + ], + ), + ( + "python module", + [ + sys.executable, + "-m", + "helper", + str(script_path), + *argv[2:], + ], + ), + ( + "interpreter option before script", + [ + sys.executable, + "-O", + str(script_path), + *argv[2:], + ], + ), ): - watchdog.validate_active_lease( - lease_path, - expected_script_path=script_path, - expected_heartbeat_path=heartbeat_path, - expected_audit_path=audit_path, - current_process_id=current_pid, - process_procfs_root=process_root, - monotonic_ns=lambda: 10_000_000_000, + with self.subTest(name): + publish_lease(dict(lease), bad_argv) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "executable argv position", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("wrong command-line policy"): + bad_argv = list(argv) + procfs_index = bad_argv.index("--procfs-root") + 1 + bad_argv[procfs_index] = "/tmp/not-proc" + publish_lease(dict(lease), bad_argv) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "command-line policy", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("wrong monitored command"): + bad_argv = [*argv[:-1], "other_matrix.py"] + bad_lease = dict(lease) + bad_lease["command"] = [ + sys.executable, + "other_matrix.py", + ] + bad_lease["child_command_sha256"] = ( + watchdog._command_sha256( + bad_lease["command"] + ) ) + publish_lease(bad_lease, bad_argv) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "monitored command", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) - with self.subTest("stale heartbeat"): - lease_path.write_text( - json.dumps(lease), encoding="utf-8" - ) - heartbeat["updated_monotonic_ns"] = 1 - heartbeat_path.write_text( - json.dumps(heartbeat), encoding="utf-8" - ) - with self.assertRaisesRegex( - watchdog.LeaseValidationError, "heartbeat is stale" - ): - watchdog.validate_active_lease( - lease_path, - expected_script_path=script_path, - expected_heartbeat_path=heartbeat_path, - expected_audit_path=audit_path, - current_process_id=current_pid, - process_procfs_root=process_root, - monotonic_ns=lambda: 10_000_000_000, + with self.subTest("tampered script SHA"): + tampered = dict(lease) + tampered["watchdog_script_sha256"] = "0" * 64 + publish_lease(tampered) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, "script SHA" + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("stale heartbeat"): + publish_lease(dict(lease)) + heartbeat["updated_monotonic_ns"] = 1 + watchdog._write_json_atomic( + heartbeat_path, heartbeat ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "heartbeat is stale", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) - with self.subTest("arbitrary heartbeat"): - heartbeat["updated_monotonic_ns"] = 9_000_000_000 - heartbeat["lease_id"] = "helper-lease" - heartbeat_path.write_text( - json.dumps(heartbeat), encoding="utf-8" - ) - with self.assertRaisesRegex( - watchdog.LeaseValidationError, - "heartbeat identity", - ): - watchdog.validate_active_lease( - lease_path, - expected_script_path=script_path, - expected_heartbeat_path=heartbeat_path, - expected_audit_path=audit_path, - current_process_id=current_pid, - process_procfs_root=process_root, - monotonic_ns=lambda: 10_000_000_000, + with self.subTest("arbitrary heartbeat"): + heartbeat["updated_monotonic_ns"] = 9_000_000_000 + heartbeat["lease_id"] = "helper-lease" + watchdog._write_json_atomic( + heartbeat_path, heartbeat ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "heartbeat identity", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) - with self.subTest("outside process group"): - heartbeat["lease_id"] = "test-lease" - heartbeat_path.write_text( - json.dumps(heartbeat), encoding="utf-8" - ) - (process_root / str(current_pid) / "stat").write_text( - self._proc_stat( - current_pid, - child_pid, - 9999, - 456791, - ), - encoding="utf-8", - ) - with self.assertRaisesRegex( - watchdog.LeaseValidationError, - "outside the monitored process group", - ): - watchdog.validate_active_lease( - lease_path, - expected_script_path=script_path, - expected_heartbeat_path=heartbeat_path, - expected_audit_path=audit_path, - current_process_id=current_pid, - process_procfs_root=process_root, - monotonic_ns=lambda: 10_000_000_000, + with self.subTest("outside process group"): + heartbeat["lease_id"] = "test-lease" + watchdog._write_json_atomic( + heartbeat_path, heartbeat ) + (process_root / str(current_pid) / "stat").write_text( + self._proc_stat( + current_pid, + child_pid, + 9999, + 456792, + ), + encoding="utf-8", + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "outside the monitored process group", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + ( + process_root / str(current_pid) / "stat" + ).write_text( + self._proc_stat( + current_pid, + child_pid, + guardian_pid, + 456792, + ), + encoding="utf-8", + ) + + with self.subTest("environment path mismatch"): + bad_validation_args = { + **validation_args, + "expected_heartbeat_path": root / "other.json", + } + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "artifact paths|heartbeat path", + ): + watchdog.validate_active_lease( + lease_path, **bad_validation_args + ) + + with self.subTest("lease inode mismatch"): + publish_lease(dict(lease)) + lease_record = json.loads( + lease_path.read_text(encoding="utf-8") + ) + lease_record["file_inode"] = 0 + lease_path.write_text( + json.dumps(lease_record), encoding="utf-8" + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "identity does not match", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("watchdog start tick mismatch"): + publish_lease(dict(lease)) + (process_root / str(watchdog_pid) / "stat").write_text( + self._proc_stat( + watchdog_pid, + 1, + watchdog_pid, + watchdog_start_ticks + 1, + ), + encoding="utf-8", + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "start time", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + finally: + os.close(audit_descriptor) def test_zero_swap_gate_launches_and_propagates_child_exit(self) -> None: harness = Harness([snapshot(50)], FakeProcess(returncode=37)) From 0071e8f21701b74b3cc4dfe5c1131950ce338143 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 23:24:49 -0700 Subject: [PATCH 27/32] scripts : harden watchdog cleanup lease Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 14 +- scripts/strix_memory_watchdog.py | 146 +++++++-- tests/test_strix_memory_watchdog.py | 450 ++++++++++++++++++++++++++++ 3 files changed, 579 insertions(+), 31 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 49a1c05547fc..90f80bff7e01 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -21,7 +21,7 @@ The wrapper performs these checks and actions: The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. -Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. +Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. The fail-closed timing bounds are a maximum 30-second grace, maximum one-second sample interval, and maximum five-second heartbeat age. The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Signal records are written immediately after each process-group signal. Child standard input, standard output, and standard error are inherited unchanged. @@ -40,28 +40,28 @@ Use all three artifact options together when another process must prove that it The watchdog creates and exclusively locks the persistent audit before launch. It then starts an internal guardian as the new session and process-group leader; the guardian starts the supplied command in that same group without inheriting the private control pipe. After the guardian reports the payload PID, the watchdog atomically creates the lease and heartbeat. Existing artifact paths are rejected rather than overwritten. The payload receives the resolved paths through `STRIX_MEMORY_WATCHDOG_LEASE_PATH`, `STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH`, and `STRIX_MEMORY_WATCHDOG_AUDIT_PATH`. It also receives `STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS`. -The child can run before the first atomic lease rename. A matching preflight must retry the inherited lease path for a bounded interval and fail closed if a complete valid lease does not appear. It must not accept a lease path supplied separately by the operator. +The child can run before the first atomic lease rename. A matching preflight must retry the inherited lease path for a bounded interval and fail closed if a complete valid lease does not appear. It must not accept a lease path supplied separately by the operator. Consumers must require version 2; version 1 does not describe the guardian topology or timing policy and is rejected. -Lease format `strix-memory-watchdog-lease`, version 1, contains: +Lease format `strix-memory-watchdog-lease`, version 2, contains: - `lease_id` and active/final `state` - `watchdog_pid`, `watchdog_start_time_utc`, Linux `watchdog_start_time_ticks`, `watchdog_executable_path`, `watchdog_command_sha256`, `watchdog_script_path`, and `watchdog_script_sha256` -- exact `soft_bytes`, `emergency_bytes`, and `strict_ceiling_bytes` +- exact `soft_bytes`, `emergency_bytes`, `strict_ceiling_bytes`, `grace_seconds`, and `sample_interval_seconds` - `procfs_root` - `guardian_pid`, payload `child_pid`, `child_process_group_id`, `command`, and `child_command_sha256` - `heartbeat_path`, `max_heartbeat_age_seconds`, and `audit_path` - device, inode, owner, and mode identity for atomic JSON artifacts, plus the watchdog-held audit descriptor identity - the authoritative `final` audit record after termination -Heartbeat format `strix-memory-watchdog-heartbeat`, version 1, binds `lease_id`, watchdog PID/start ticks, child PID/process group, sequence, state, and update timestamps. Every memory sample first checks swap and memory thresholds, pulses the guardian through the private nonblocking pipe, then atomically replaces the heartbeat with the complete sample audit record and its persistent-audit record hash. A blocked audit or heartbeat write cannot delay the emergency signal. A final heartbeat and final lease update remain on disk with the persistent JSONL audit; the watchdog does not delete this evidence. +Heartbeat format `strix-memory-watchdog-heartbeat`, version 2, binds `lease_id`, watchdog PID/start ticks, child PID/process group, sequence, state, and update timestamps. Every memory sample first checks swap and memory thresholds, pulses the guardian through the private nonblocking pipe, then atomically replaces the heartbeat with the complete sample audit record and its persistent-audit record hash. It pulses again after persistence succeeds. A blocked audit or heartbeat write cannot delay the emergency signal; if persistence stalls past the guardian deadline, the guardian fails closed. A final heartbeat and final lease update remain on disk with the persistent JSONL audit; the watchdog does not delete this evidence. -The guardian uses Linux `PR_SET_PDEATHSIG` with a parent-race check. It kills its process group on watchdog death, control-pipe EOF/error, or a missed pulse deadline, including a stopped or wedged watchdog. The payload must call `start_process_group_lease_guard()` before it starts exporter descendants. This validates the lease with bounded startup retries, arms a second parent-death link to the guardian, and starts a thread that kills the process group if the watchdog evidence becomes stale or invalid. +The guardian uses Linux `PR_SET_PDEATHSIG` with a parent-race check. It kills its process group on watchdog death, control-pipe EOF/error, or a missed pulse deadline, including a stopped or wedged watchdog. When the watchdog sends a graceful signal, it also puts the guardian into a bounded grace mode and continues private pulses while it waits. This lets the watchdog own the configured grace deadline and record any `SIGKILL` escalation instead of letting the shorter heartbeat deadline preempt cleanup. The payload must call `start_process_group_lease_guard()` before it starts exporter descendants. This validates the lease with bounded startup retries, arms a second parent-death link to the guardian, and starts a thread that kills the process group if any validation or artifact operation fails or the watchdog evidence becomes stale. A matching Linux preflight must verify all of the following: - The inherited lease, heartbeat, and audit paths match the paths inside the lease. - `/proc//exe` is the exact expected Python executable and argv position 1 is the exact repository watchdog script. `-c`, `-m`, helper-script, inert-argument, and interpreter-option substitutions are rejected. -- The watchdog command line itself supplies the exact 116/118 GiB thresholds, `/proc`, inherited artifact paths, and command after `--`; the lease cannot override those expectations. +- The watchdog command line itself supplies the exact 116/118 GiB thresholds, `/proc`, inherited artifact paths, timing policy, and command after `--`; the lease cannot override those expectations. - `/proc//stat` start ticks and `/proc//cmdline` SHA-256 match the lease and remain stable across validation. A pidfd is held during validation when Linux provides `pidfd_open`. - The topology is watchdog parent -> guardian process-group leader -> payload child. The current process must be inside `child_process_group_id`. - The command identity is expected, the procfs root is `/proc`, and thresholds are exactly 116 GiB soft, 118 GiB emergency, and 120 GiB strict ceiling for the final run. diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index dd185a00f5f2..a16b79743e10 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -32,11 +32,14 @@ DEFAULT_GRACE_SECONDS = 30.0 DEFAULT_SAMPLE_INTERVAL_SECONDS = 1.0 DEFAULT_HEARTBEAT_MAX_AGE_SECONDS = 5.0 +MAX_GRACE_SECONDS = 30.0 +MAX_SAMPLE_INTERVAL_SECONDS = 1.0 +MAX_HEARTBEAT_MAX_AGE_SECONDS = 5.0 LEASE_FORMAT = "strix-memory-watchdog-lease" -LEASE_VERSION = 1 +LEASE_VERSION = 2 HEARTBEAT_FORMAT = "strix-memory-watchdog-heartbeat" -HEARTBEAT_VERSION = 1 +HEARTBEAT_VERSION = 2 PR_SET_PDEATHSIG = 1 LEASE_GUARD_SIGNAL = signal.SIGUSR1 @@ -106,8 +109,14 @@ def wait(self, timeout: float | None = None) -> int: return self.process.wait(timeout=timeout) def pulse(self) -> None: + self._write_control(b"P") + + def begin_grace(self) -> None: + self._write_control(b"G") + + def _write_control(self, value: bytes) -> None: try: - os.write(self.pulse_fd, b"\0") + os.write(self.pulse_fd, value) except BlockingIOError as exc: raise ProcessGroupError( "guardian pulse pipe is blocked" @@ -175,20 +184,32 @@ def validate(self) -> ArtifactPaths | None: raise ValueError("emergency threshold must be greater than soft threshold") if self.emergency_bytes >= STRICT_CEILING_BYTES: raise ValueError("emergency threshold must be below 120 GiB") - if not math.isfinite(self.grace_seconds) or self.grace_seconds <= 0: - raise ValueError("grace period must be greater than zero") + if ( + not math.isfinite(self.grace_seconds) + or self.grace_seconds <= 0 + or self.grace_seconds > MAX_GRACE_SECONDS + ): + raise ValueError( + "grace period must be greater than zero and at most 30 seconds" + ) if ( not math.isfinite(self.sample_interval_seconds) or self.sample_interval_seconds <= 0 + or self.sample_interval_seconds > MAX_SAMPLE_INTERVAL_SECONDS ): - raise ValueError("sample interval must be greater than zero") + raise ValueError( + "sample interval must be greater than zero and at most 1 second" + ) if ( not math.isfinite(self.heartbeat_max_age_seconds) or self.heartbeat_max_age_seconds <= self.sample_interval_seconds + or self.heartbeat_max_age_seconds + > MAX_HEARTBEAT_MAX_AGE_SECONDS ): raise ValueError( - "heartbeat max age must be greater than sample interval" + "heartbeat max age must be greater than sample interval " + "and at most 5 seconds" ) lease_paths = ( self.lease_path, @@ -348,6 +369,7 @@ def _guardian_main( control_fd: int, status_fd: int, pulse_timeout_seconds: float, + grace_timeout_seconds: float, command: tuple[str, ...], ) -> int: if not sys.platform.startswith("linux"): @@ -386,6 +408,7 @@ def prepare_payload() -> None: control_fd, select.POLLIN | select.POLLHUP | select.POLLERR, ) + current_timeout_seconds = pulse_timeout_seconds deadline = time.monotonic() + pulse_timeout_seconds while True: remaining = max(0.0, deadline - time.monotonic()) @@ -399,16 +422,20 @@ def prepare_payload() -> None: pulse = b"" if not pulse: _kill_own_process_group() - deadline = time.monotonic() + pulse_timeout_seconds + if b"G" in pulse: + current_timeout_seconds = grace_timeout_seconds + deadline = time.monotonic() + current_timeout_seconds if time.monotonic() >= deadline: _kill_own_process_group() returncode = payload.poll() if returncode is not None: - return ( - 128 - returncode - if returncode < 0 - else returncode - ) + if returncode >= 0: + return returncode + signal_number = -returncode + if signal_number not in (signal.SIGKILL, signal.SIGSTOP): + signal.signal(signal_number, signal.SIG_DFL) + os.kill(os.getpid(), signal_number) + return 128 + signal_number def _read_guardian_status( @@ -450,6 +477,7 @@ def _launch_guardian( command: tuple[str, ...], environment: dict[str, str], pulse_timeout_seconds: float, + grace_timeout_seconds: float, launch_mask: set[signal.Signals], ) -> GuardianProcess: control_read, control_write = os.pipe() @@ -469,6 +497,7 @@ def prepare_guardian() -> None: str(control_read), str(status_write), str(pulse_timeout_seconds), + str(grace_timeout_seconds), "--", *command, ) @@ -712,6 +741,8 @@ def start( "soft_bytes": self.config.soft_bytes, "emergency_bytes": self.config.emergency_bytes, "strict_ceiling_bytes": STRICT_CEILING_BYTES, + "grace_seconds": self.config.grace_seconds, + "sample_interval_seconds": self.config.sample_interval_seconds, "guardian_pid": child.pid, "child_pid": payload_pid, "child_process_group_id": child.pid, @@ -956,6 +987,16 @@ def validate_active_lease( raise LeaseValidationError( "watchdog command-line policy does not match" ) + if ( + lease.get("grace_seconds") != live_config.grace_seconds + or lease.get("sample_interval_seconds") + != live_config.sample_interval_seconds + or lease.get("max_heartbeat_age_seconds") + != live_config.heartbeat_max_age_seconds + ): + raise LeaseValidationError( + "watchdog lease timing policy does not match" + ) if ( live_paths is None or live_paths.lease != lease_path @@ -1246,7 +1287,7 @@ def start_process_group_lease_guard( process_procfs_root=process_procfs_root, ) break - except LeaseValidationError: + except Exception: if time.monotonic() >= deadline: raise time.sleep(0.01) @@ -1272,7 +1313,7 @@ def monitor() -> None: expected_max_heartbeat_age_seconds=max_age_seconds, process_procfs_root=process_procfs_root, ) - except LeaseValidationError: + except Exception: _kill_own_process_group() guard = threading.Thread( @@ -1291,6 +1332,7 @@ def __init__( wall_clock: Callable[[], datetime] | None = None, ): self.stream = stream + self.stream_enabled = True self.wall_clock = wall_clock self.persistent_stream: IO[str] | None = None self.lease_manager: LeaseManager | None = None @@ -1337,15 +1379,21 @@ def persistent_identity(self) -> dict[str, int]: } def close(self) -> None: - if self.persistent_stream is not None: - self.persistent_stream.close() - self.persistent_stream = None + persistent_stream = self.persistent_stream + self.persistent_stream = None + if persistent_stream is not None: + try: + persistent_stream.close() + except (OSError, ValueError): + pass def disable_component(self, component: str) -> None: if component == "audit": self.close() elif component == "lease": self.lease_manager = None + elif component == "stderr": + self.stream_enabled = False def emit(self, event: str, **fields: object) -> dict[str, object]: record = { @@ -1357,8 +1405,16 @@ def emit(self, event: str, **fields: object) -> dict[str, object]: json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" ) - self.stream.write(line) - self.stream.flush() + if self.stream_enabled: + try: + self.stream.write(line) + self.stream.flush() + except (OSError, ValueError) as exc: + detail = getattr(exc, "strerror", None) or str(exc) + raise ArtifactError( + "stderr", + f"cannot write standard error audit: {detail}", + ) from exc self.last_record_sha256 = _sha256_bytes(line.encode("utf-8")) if self.persistent_stream is not None: try: @@ -1654,6 +1710,11 @@ def _graceful_cleanup( process_group_status = signal_group( child.pid, graceful_signal ) + if ( + isinstance(child, GuardianProcess) + and child.poll() is None + ): + child.begin_grace() try: audit.emit( "process_group_signal", @@ -1675,6 +1736,11 @@ def _graceful_cleanup( child.poll() if not group_alive(child.pid): break + if ( + isinstance(child, GuardianProcess) + and child.poll() is None + ): + child.pulse() sleeper(min(0.05, deadline - monotonic())) child.poll() if group_alive(child.pid): @@ -1908,6 +1974,20 @@ def _monitor_child( str(exc), ) soft_deadline = now + config.grace_seconds + if isinstance(child, GuardianProcess): + try: + child.begin_grace() + except ProcessGroupError as exc: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "signal_error", + EXIT_SIGNAL_ERROR, + str(exc), + signal_group, + ) soft_signal_fields = { **_state_fields( state.snapshot, @@ -1923,6 +2003,8 @@ def _monitor_child( try: pulse_guardian() except ProcessGroupError as exc: + if child.poll() is not None: + continue return _kill_and_finish( audit, child, @@ -1950,6 +2032,21 @@ def _monitor_child( ) ) audit.heartbeat(sample_record) + try: + pulse_guardian() + except ProcessGroupError as exc: + if child.poll() is not None: + continue + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "signal_error", + EXIT_SIGNAL_ERROR, + str(exc), + signal_group, + ) sleep_seconds = config.sample_interval_seconds if soft_deadline is not None: @@ -2107,6 +2204,7 @@ def restore_child_signal_mask() -> None: config.command, child_environment, config.heartbeat_max_age_seconds, + config.grace_seconds + 1.0, launch_mask, ) elif lease_manager is not None: @@ -2168,8 +2266,7 @@ def restore_child_signal_mask() -> None: ) except ArtifactError as exc: _set_parent_signal_handlers(signal.SIG_IGN) - if exc.component == "audit": - audit.disable_component(exc.component) + audit.disable_component(exc.component) if child is None: return _emit_final( audit, @@ -2369,14 +2466,15 @@ def parse_args(argv: Sequence[str]) -> WatchdogConfig: def main(argv: Sequence[str] | None = None) -> int: arguments = tuple(argv if argv is not None else sys.argv[1:]) if arguments and arguments[0] == "--internal-guardian": - if len(arguments) < 6 or arguments[4] != "--": + if len(arguments) < 7 or arguments[5] != "--": return EXIT_LAUNCH_ERROR try: return _guardian_main( int(arguments[1]), int(arguments[2]), _positive_float(arguments[3]), - tuple(arguments[5:]), + _positive_float(arguments[4]), + tuple(arguments[6:]), ) except (OSError, ValueError): return EXIT_LAUNCH_ERROR diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index a61078efb2a3..3a0a534717da 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -401,6 +401,168 @@ def test_parent_signals_leave_no_child_or_grandchild(self) -> None: "parent_signal", ) + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_parent_signal_grace_outlives_guardian_pulse_timeout( + self, + ) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "open(sys.argv[1],'w').write(str(os.getpid()));" + "time.sleep(30)" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pid" + stderr_path = root / "stderr.jsonl" + self._write_procfs_fixture(root) + with stderr_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--grace-seconds", + "0.4", + "--sample-interval-seconds", + "0.05", + "--heartbeat-max-age-seconds", + "0.1", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + child_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + self.fail("child process did not become ready") + time.sleep(0.01) + child_pid = int( + pid_file.read_text(encoding="utf-8") + ) + started = time.monotonic() + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=5) + elapsed = time.monotonic() - started + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + if child_pid is not None: + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in stderr_path.read_text( + encoding="utf-8" + ).splitlines() + ] + signals = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertGreaterEqual(elapsed, 0.35) + self.assertEqual( + wrapper.returncode, + 128 + signal.SIGTERM, + records, + ) + self.assertEqual(signals, ["SIGTERM", "SIGKILL"]) + self.assertEqual( + records[-1]["classification"], "parent_signal" + ) + self.assertFalse(self._process_is_running(child_pid)) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_parent_signal_allows_exit_after_pulse_deadline(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " time.sleep(0.25)\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "open(sys.argv[1],'w').write(str(os.getpid()))\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pid" + stderr_path = root / "stderr.jsonl" + self._write_procfs_fixture(root) + with stderr_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--grace-seconds", + "0.4", + "--sample-interval-seconds", + "0.05", + "--heartbeat-max-age-seconds", + "0.1", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + try: + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + self.fail("child process did not become ready") + time.sleep(0.01) + started = time.monotonic() + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=5) + elapsed = time.monotonic() - started + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + + records = [ + json.loads(line) + for line in stderr_path.read_text( + encoding="utf-8" + ).splitlines() + ] + signals = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertGreaterEqual(elapsed, 0.2) + self.assertLess(elapsed, 0.4) + self.assertEqual(wrapper.returncode, 128 + signal.SIGTERM) + self.assertEqual(signals, ["SIGTERM"]) + self.assertEqual(records[-1]["child_returncode"], 0) + @unittest.skipUnless( sys.platform.startswith("linux"), "Linux guardian lifecycle", @@ -431,6 +593,7 @@ def test_guardian_pipe_close_kills_group_without_fd_leak(self) -> None: ), os.environ.copy(), 0.5, + 1.0, signal.pthread_sigmask(signal.SIG_BLOCK, ()), ) control_target = os.readlink( @@ -478,6 +641,7 @@ def test_guardian_documents_setsid_escape_limit(self) -> None: ), os.environ.copy(), 0.5, + 1.0, signal.pthread_sigmask(signal.SIG_BLOCK, ()), ) deadline = time.monotonic() + 5 @@ -594,6 +758,96 @@ def test_guard_kills_group_after_watchdog_loss_or_stall(self) -> None: if wrapper.stderr is not None: wrapper.stderr.close() + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_payload_guard_fails_closed_on_artifact_error(self) -> None: + child_code = ( + "import importlib.util,os,pathlib,subprocess,sys,time\n" + "script=pathlib.Path(sys.argv[1])\n" + "spec=importlib.util.spec_from_file_location('guard_watchdog',script)\n" + "module=importlib.util.module_from_spec(spec)\n" + "sys.modules[spec.name]=module\n" + "spec.loader.exec_module(module)\n" + "module.start_process_group_lease_guard(" + "script,expected_procfs_root=pathlib.Path(sys.argv[2]))\n" + "def fail(*_args,**_kwargs):\n" + " raise module.ArtifactError('script','unreadable')\n" + "module.validate_active_lease=fail\n" + "grandchild=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'])\n" + "open(sys.argv[3],'w').write(" + "f'{os.getpid()} {grandchild.pid}\\n')\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_path = root / "pids" + self._write_procfs_fixture(root) + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--heartbeat-max-age-seconds", + "0.3", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(SCRIPT_PATH), + str(root), + str(pid_path), + ], + stderr=subprocess.PIPE, + text=True, + ) + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_path.exists(): + if wrapper.poll() is not None: + assert wrapper.stderr is not None + self.fail(wrapper.stderr.read()) + if time.monotonic() >= deadline: + self.fail("guarded payload did not become ready") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_path.read_text( + encoding="utf-8" + ).split() + ) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + if wrapper.stderr is not None: + wrapper.stderr.close() + if child_pid is not None: + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + assert child_pid is not None + assert grandchild_pid is not None + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + def test_child_sigterm_handler_exits_without_escalation(self) -> None: child_code = ( "import os,signal,sys,time\n" @@ -869,6 +1123,143 @@ def test_configuration_rejects_non_finite_timing(self) -> None: with self.assertRaisesRegex(ValueError, "grace period"): config.validate() + def test_configuration_rejects_weakened_liveness_timing(self) -> None: + cases = ( + ( + {"grace_seconds": 31.0}, + "grace period", + ), + ( + {"sample_interval_seconds": 1.1}, + "sample interval", + ), + ( + { + "sample_interval_seconds": 1.0, + "heartbeat_max_age_seconds": 5.1, + }, + "heartbeat max age", + ), + ) + for overrides, message in cases: + with self.subTest(overrides=overrides): + config = watchdog.WatchdogConfig( + command=("fake-command",), + **overrides, + ) + with self.assertRaisesRegex(ValueError, message): + config.validate() + + def test_stderr_failure_does_not_bypass_cleanup(self) -> None: + class FailingStderr(io.StringIO): + def write(self, value: str) -> int: + raise OSError("stderr closed") + + process = FakeProcess() + + def exit_on_kill( + target: FakeProcess, signal_number: int + ) -> None: + if signal_number == signal.SIGKILL: + target.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50)], + process, + signal_handler=exit_on_kill, + ) + harness.audit = watchdog.AuditLogger(FailingStderr()) + with tempfile.TemporaryDirectory() as temp_dir: + persistent_path = Path(temp_dir) / "audit.jsonl" + harness.audit.open_persistent(persistent_path) + result = watchdog._graceful_cleanup( + harness.audit, + process, + snapshot(50), + 50, + "internal_error", + watchdog.EXIT_INTERNAL_ERROR, + "test cleanup", + signal.SIGTERM, + 0.1, + harness.signal_group, + harness.group_alive, + harness.clock.monotonic, + harness.clock.sleep, + ) + harness.audit.close() + records = [ + json.loads(line) + for line in persistent_path.read_text( + encoding="utf-8" + ).splitlines() + ] + + self.assertEqual( + harness.signals, [signal.SIGTERM, signal.SIGKILL] + ) + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(records[-1]["classification"], "lease_error") + + def test_audit_write_and_close_failures_do_not_bypass_cleanup( + self, + ) -> None: + class FailingPersistent(io.StringIO): + def __init__(self) -> None: + super().__init__() + self.close_called = False + + def write(self, value: str) -> int: + raise OSError("persistent write failed") + + def close(self) -> None: + if self.close_called: + super().close() + return + self.close_called = True + raise OSError("persistent close failed") + + process = FakeProcess() + + def exit_on_kill( + target: FakeProcess, signal_number: int + ) -> None: + if signal_number == signal.SIGKILL: + target.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50)], + process, + signal_handler=exit_on_kill, + ) + persistent = FailingPersistent() + harness.audit.persistent_stream = persistent + result = watchdog._graceful_cleanup( + harness.audit, + process, + snapshot(50), + 50, + "internal_error", + watchdog.EXIT_INTERNAL_ERROR, + "test cleanup", + signal.SIGTERM, + 0.1, + harness.signal_group, + harness.group_alive, + harness.clock.monotonic, + harness.clock.sleep, + ) + + self.assertTrue(persistent.close_called) + self.assertEqual( + harness.signals, [signal.SIGTERM, signal.SIGKILL] + ) + self.assertEqual(process.returncode, -signal.SIGKILL) + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual( + harness.records()[-1]["classification"], "lease_error" + ) + def test_final_record_survives_artifact_failures(self) -> None: class FailingLease: def finalize(self, record: dict[str, Any]) -> None: @@ -1114,6 +1505,49 @@ def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: str((root / "persistent-audit.jsonl").resolve()), ) + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_preserves_payload_signal_status(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + self._write_procfs_fixture(root) + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--sample-interval-seconds", + "0.01", + "--", + sys.executable, + "-c", + ( + "import os,signal,time;" + "time.sleep(0.1);" + "os.kill(os.getpid(),signal.SIGTERM)" + ), + ], + capture_output=True, + check=False, + text=True, + timeout=5, + ) + records = [ + json.loads(line) + for line in result.stderr.splitlines() + ] + + self.assertEqual(result.returncode, 128 + signal.SIGTERM) + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual( + records[-1]["child_returncode"], -signal.SIGTERM + ) + self.assertEqual(records[-1]["child_status"], "signaled") + def test_existing_lease_fails_closed_and_stops_child(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) @@ -1249,6 +1683,10 @@ def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: "soft_bytes": watchdog.DEFAULT_SOFT_BYTES, "emergency_bytes": watchdog.DEFAULT_EMERGENCY_BYTES, "strict_ceiling_bytes": watchdog.STRICT_CEILING_BYTES, + "grace_seconds": watchdog.DEFAULT_GRACE_SECONDS, + "sample_interval_seconds": ( + watchdog.DEFAULT_SAMPLE_INTERVAL_SECONDS + ), "guardian_pid": guardian_pid, "child_pid": child_pid, "child_process_group_id": guardian_pid, @@ -1392,6 +1830,18 @@ def publish_lease( lease_path, **validation_args ) + with self.subTest("wrong lease timing policy"): + bad_lease = dict(lease) + bad_lease["grace_seconds"] = 29.0 + publish_lease(bad_lease) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "lease timing policy", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + with self.subTest("wrong monitored command"): bad_argv = [*argv[:-1], "other_matrix.py"] bad_lease = dict(lease) From cd97b50b664bd170819b4225ae52cad97d762ada Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 23:39:59 -0700 Subject: [PATCH 28/32] scripts : enforce cleanup on guardian errors Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 2 +- scripts/strix_memory_watchdog.py | 54 ++++++++---- tests/test_strix_memory_watchdog.py | 123 ++++++++++++++++++++++++++++ 3 files changed, 163 insertions(+), 16 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 90f80bff7e01..da3560caddd9 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -55,7 +55,7 @@ Lease format `strix-memory-watchdog-lease`, version 2, contains: Heartbeat format `strix-memory-watchdog-heartbeat`, version 2, binds `lease_id`, watchdog PID/start ticks, child PID/process group, sequence, state, and update timestamps. Every memory sample first checks swap and memory thresholds, pulses the guardian through the private nonblocking pipe, then atomically replaces the heartbeat with the complete sample audit record and its persistent-audit record hash. It pulses again after persistence succeeds. A blocked audit or heartbeat write cannot delay the emergency signal; if persistence stalls past the guardian deadline, the guardian fails closed. A final heartbeat and final lease update remain on disk with the persistent JSONL audit; the watchdog does not delete this evidence. -The guardian uses Linux `PR_SET_PDEATHSIG` with a parent-race check. It kills its process group on watchdog death, control-pipe EOF/error, or a missed pulse deadline, including a stopped or wedged watchdog. When the watchdog sends a graceful signal, it also puts the guardian into a bounded grace mode and continues private pulses while it waits. This lets the watchdog own the configured grace deadline and record any `SIGKILL` escalation instead of letting the shorter heartbeat deadline preempt cleanup. The payload must call `start_process_group_lease_guard()` before it starts exporter descendants. This validates the lease with bounded startup retries, arms a second parent-death link to the guardian, and starts a thread that kills the process group if any validation or artifact operation fails or the watchdog evidence becomes stale. +The guardian uses Linux `PR_SET_PDEATHSIG` with a parent-race check. It kills its process group on watchdog death, control-pipe EOF/error, or a missed pulse deadline, including a stopped or wedged watchdog. When the watchdog sends a graceful signal, it also puts the guardian into a bounded grace mode and continues private pulses while it waits. This lets the watchdog own the configured grace deadline and record any `SIGKILL` escalation instead of letting the shorter heartbeat deadline preempt cleanup. If the grace control message or a cleanup pulse fails, the watchdog independently sends `SIGKILL` to the process group and reaps the child before it reports `signal_error`. The payload must call `start_process_group_lease_guard()` before it starts exporter descendants. This validates the lease with bounded startup retries, arms a second parent-death link to the guardian, and starts a thread that kills the process group if any validation or artifact operation fails or the watchdog evidence becomes stale. A matching Linux preflight must verify all of the following: diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index a16b79743e10..0eb6df6d78ba 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -1705,16 +1705,12 @@ def _graceful_cleanup( ) -> int: escalated = False artifact_error: ArtifactError | None = None + guardian_control_error: ProcessGroupError | None = None try: if graceful_signal is not None: process_group_status = signal_group( child.pid, graceful_signal ) - if ( - isinstance(child, GuardianProcess) - and child.poll() is None - ): - child.begin_grace() try: audit.emit( "process_group_signal", @@ -1731,8 +1727,19 @@ def _graceful_cleanup( except ArtifactError as exc: artifact_error = exc audit.disable_component(exc.component) + if ( + isinstance(child, GuardianProcess) + and child.poll() is None + ): + try: + child.begin_grace() + except ProcessGroupError as exc: + guardian_control_error = exc deadline = monotonic() + grace_seconds - while monotonic() < deadline: + while ( + guardian_control_error is None + and monotonic() < deadline + ): child.poll() if not group_alive(child.pid): break @@ -1740,19 +1747,31 @@ def _graceful_cleanup( isinstance(child, GuardianProcess) and child.poll() is None ): - child.pulse() + try: + child.pulse() + except ProcessGroupError as exc: + guardian_control_error = exc + break sleeper(min(0.05, deadline - monotonic())) child.poll() - if group_alive(child.pid): + if ( + guardian_control_error is not None + or group_alive(child.pid) + ): escalated = True process_group_status = signal_group( child.pid, signal.SIGKILL ) - signal_reason = ( - escalation_result[2] - if escalation_result is not None - else reason - ) + if guardian_control_error is not None: + signal_reason = ( + "guardian control failed during graceful cleanup" + ) + else: + signal_reason = ( + escalation_result[2] + if escalation_result is not None + else reason + ) try: audit.emit( "process_group_signal", @@ -1802,9 +1821,14 @@ def _graceful_cleanup( str(exc), ) - if escalated and escalation_result is not None: + if guardian_control_error is not None: + classification = "signal_error" + exit_code = EXIT_SIGNAL_ERROR + reason = "guardian control failed during graceful cleanup" + error = str(guardian_control_error) + elif escalated and escalation_result is not None: classification, exit_code, reason = escalation_result - if artifact_error is not None: + if artifact_error is not None and guardian_control_error is None: classification = "lease_error" exit_code = EXIT_LEASE_ERROR error = f"{artifact_error.component}: {artifact_error}" diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 3a0a534717da..d1cfb7f3ef6a 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -563,6 +563,129 @@ def test_parent_signal_allows_exit_after_pulse_deadline(self) -> None: self.assertEqual(signals, ["SIGTERM"]) self.assertEqual(records[-1]["child_returncode"], 0) + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_control_failure_still_kills_and_reaps_group( + self, + ) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n');" + " time.sleep(30)\n" + ) + for mode in ("closed", "blocked"): + with self.subTest(mode=mode): + with tempfile.TemporaryDirectory() as temp_dir: + pid_path = Path(temp_dir) / "pids" + process = subprocess.Popen( + [ + sys.executable, + "-c", + child_code, + str(pid_path), + ], + start_new_session=True, + ) + read_fd, write_fd = os.pipe() + os.set_blocking(write_fd, False) + if mode == "closed": + os.close(write_fd) + write_fd = -1 + else: + try: + while True: + os.write(write_fd, b"x" * 65536) + except BlockingIOError: + pass + guardian = watchdog.GuardianProcess( + process, + process.pid, + write_fd, + ) + stream = io.StringIO() + audit = watchdog.AuditLogger(stream) + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_path.exists(): + if time.monotonic() >= deadline: + self.fail( + "child process group did not start" + ) + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_path.read_text( + encoding="utf-8" + ).split() + ) + result = watchdog._graceful_cleanup( + audit, + guardian, + snapshot(50), + 50, + "parent_signal", + 128 + signal.SIGTERM, + "wrapper received SIGTERM", + signal.SIGTERM, + 0.4, + watchdog._signal_process_group, + watchdog._process_group_alive, + time.monotonic, + time.sleep, + ) + finally: + if process.poll() is None: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5) + guardian.close() + os.close(read_fd) + + records = [ + json.loads(line) + for line in stream.getvalue().splitlines() + ] + signals = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual( + result, watchdog.EXIT_SIGNAL_ERROR + ) + self.assertEqual(signals, ["SIGTERM", "SIGKILL"]) + self.assertEqual( + records[-1]["classification"], "signal_error" + ) + self.assertEqual( + records[-1]["child_returncode"], + -signal.SIGKILL, + ) + assert child_pid is not None + assert grandchild_pid is not None + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse( + self._process_is_running(process_id) + ) + @unittest.skipUnless( sys.platform.startswith("linux"), "Linux guardian lifecycle", From f6b4da49b913f7e6fe739e9ca662387f8c2ef659 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 23:50:22 -0700 Subject: [PATCH 29/32] scripts : preserve watchdog failure cause Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 2 +- scripts/strix_memory_watchdog.py | 59 +++++-- tests/test_strix_memory_watchdog.py | 259 +++++++++++++++++----------- 3 files changed, 202 insertions(+), 118 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index da3560caddd9..f0099b9a55ad 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -73,7 +73,7 @@ These checks reject accidental or helper-process substitution and make regular-f The guardian controls only the process group. A payload that deliberately calls `setsid()` can escape it. The correctness harness must not do that. If arbitrary payload code is in scope, launch the watchdog in a service/cgroup configured to kill every member when the unit stops. -Exit classifications are authoritative in the final JSON record. Operational failures use these exit codes: +Exit classifications are authoritative in the last final JSON record. If final artifact persistence fails after a primary safety failure, the primary classification and exit code remain unchanged and the artifact failure is listed in `secondary_errors`. Operational failures use these exit codes: | Exit code | Classification | | ---: | --- | diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 0eb6df6d78ba..75f3e8852e27 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -1491,6 +1491,7 @@ def _emit_final( child_returncode: int | None = None, process_group_status: str = "not_created", error: str | None = None, + preserve_primary_on_artifact_error: bool = False, ) -> int: fields = _state_fields( snapshot, @@ -1503,31 +1504,52 @@ def _emit_final( fields.update(classification=classification, exit_code=exit_code) if error: fields["error"] = error + + def record_artifact_error(exc: ArtifactError) -> None: + nonlocal exit_code + detail = { + "component": exc.component, + "detail": str(exc), + } + if preserve_primary_on_artifact_error: + secondary_errors = fields.setdefault( + "secondary_errors", [] + ) + assert isinstance(secondary_errors, list) + secondary_errors.append(detail) + else: + fields.update( + classification="lease_error", + exit_code=EXIT_LEASE_ERROR, + threshold_reason="watchdog artifact finalization failed", + error=f"{exc.component}: {exc}", + ) + exit_code = EXIT_LEASE_ERROR + + def emit_final_record() -> dict[str, object]: + try: + return audit.emit("final", **fields) + except ArtifactError as exc: + audit.disable_component(exc.component) + record_artifact_error(exc) + try: + return audit.emit("final", **fields) + except ArtifactError as exc: + audit.disable_component(exc.component) + record_artifact_error(exc) + return audit.emit("final", **fields) + previous_mask = signal.pthread_sigmask( signal.SIG_BLOCK, PARENT_SIGNALS ) try: + record = emit_final_record() try: - record = audit.emit("final", **fields) audit.finalize(record) except ArtifactError as exc: audit.disable_component(exc.component) - fields.update( - classification="lease_error", - exit_code=EXIT_LEASE_ERROR, - threshold_reason="watchdog artifact finalization failed", - error=f"{exc.component}: {exc}", - ) - try: - record = audit.emit("final", **fields) - except ArtifactError as nested_exc: - audit.disable_component(nested_exc.component) - record = audit.emit("final", **fields) - try: - audit.finalize(record) - except ArtifactError as nested_exc: - audit.disable_component(nested_exc.component) - exit_code = EXIT_LEASE_ERROR + record_artifact_error(exc) + emit_final_record() audit.mark_final(exit_code) return exit_code finally: @@ -1843,6 +1865,9 @@ def _graceful_cleanup( child_returncode, process_group_status, error, + preserve_primary_on_artifact_error=( + guardian_control_error is not None + ), ) diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index d1cfb7f3ef6a..8bb209fcbfb4 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -570,6 +570,32 @@ def test_parent_signal_allows_exit_after_pulse_deadline(self) -> None: def test_guardian_control_failure_still_kills_and_reaps_group( self, ) -> None: + class FailingFinalAudit: + def __init__(self, stream: Any): + self.stream = stream + self.write_count = 0 + + def write(self, value: str) -> int: + self.write_count += 1 + if self.write_count == 3: + raise OSError("final audit write failed") + return self.stream.write(value) + + def flush(self) -> None: + self.stream.flush() + + def fileno(self) -> int: + return self.stream.fileno() + + def close(self) -> None: + self.stream.close() + + class FailingFinalLease: + def finalize(self, record: dict[str, Any]) -> None: + raise watchdog.ArtifactError( + "lease", "final lease write failed" + ) + child_code = ( "import os,signal,sys,time;" "signal.signal(signal.SIGTERM,signal.SIG_IGN);" @@ -582,109 +608,142 @@ def test_guardian_control_failure_still_kills_and_reaps_group( " time.sleep(30)\n" ) for mode in ("closed", "blocked"): - with self.subTest(mode=mode): - with tempfile.TemporaryDirectory() as temp_dir: - pid_path = Path(temp_dir) / "pids" - process = subprocess.Popen( - [ - sys.executable, - "-c", - child_code, - str(pid_path), - ], - start_new_session=True, - ) - read_fd, write_fd = os.pipe() - os.set_blocking(write_fd, False) - if mode == "closed": - os.close(write_fd) - write_fd = -1 - else: - try: - while True: - os.write(write_fd, b"x" * 65536) - except BlockingIOError: - pass - guardian = watchdog.GuardianProcess( - process, - process.pid, - write_fd, + for artifact_failure in ("audit", "lease"): + with self.subTest( + mode=mode, + artifact_failure=artifact_failure, + ): + self._assert_guardian_control_failure_cleanup( + mode, + artifact_failure, + child_code, + FailingFinalAudit, + FailingFinalLease, ) - stream = io.StringIO() - audit = watchdog.AuditLogger(stream) - child_pid = None - grandchild_pid = None + + def _assert_guardian_control_failure_cleanup( + self, + mode: str, + artifact_failure: str, + child_code: str, + failing_final_audit: type, + failing_final_lease: type, + ) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + pid_path = Path(temp_dir) / "pids" + process = subprocess.Popen( + [ + sys.executable, + "-c", + child_code, + str(pid_path), + ], + start_new_session=True, + ) + read_fd, write_fd = os.pipe() + os.set_blocking(write_fd, False) + if mode == "closed": + os.close(write_fd) + write_fd = -1 + else: + try: + while True: + os.write(write_fd, b"x" * 65536) + except BlockingIOError: + pass + guardian = watchdog.GuardianProcess( + process, + process.pid, + write_fd, + ) + stream = io.StringIO() + audit = watchdog.AuditLogger(stream) + if artifact_failure == "audit": + persistent_path = Path(temp_dir) / "persistent.jsonl" + persistent_stream = persistent_path.open( + "w", encoding="utf-8" + ) + audit.persistent_stream = failing_final_audit( + persistent_stream + ) + else: + audit.lease_manager = failing_final_lease() + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_path.exists(): + if time.monotonic() >= deadline: + self.fail("child process group did not start") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_path.read_text( + encoding="utf-8" + ).split() + ) + result = watchdog._graceful_cleanup( + audit, + guardian, + snapshot(50), + 50, + "parent_signal", + 128 + signal.SIGTERM, + "wrapper received SIGTERM", + signal.SIGTERM, + 0.4, + watchdog._signal_process_group, + watchdog._process_group_alive, + time.monotonic, + time.sleep, + ) + finally: + if process.poll() is None: try: - deadline = time.monotonic() + 5 - while not pid_path.exists(): - if time.monotonic() >= deadline: - self.fail( - "child process group did not start" - ) - time.sleep(0.01) - child_pid, grandchild_pid = ( - int(value) - for value in pid_path.read_text( - encoding="utf-8" - ).split() - ) - result = watchdog._graceful_cleanup( - audit, - guardian, - snapshot(50), - 50, - "parent_signal", - 128 + signal.SIGTERM, - "wrapper received SIGTERM", - signal.SIGTERM, - 0.4, - watchdog._signal_process_group, - watchdog._process_group_alive, - time.monotonic, - time.sleep, - ) - finally: - if process.poll() is None: - try: - os.killpg(process.pid, signal.SIGKILL) - except ProcessLookupError: - pass - process.wait(timeout=5) - guardian.close() - os.close(read_fd) + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5) + guardian.close() + os.close(read_fd) - records = [ - json.loads(line) - for line in stream.getvalue().splitlines() - ] - signals = [ - record["signal"] - for record in records - if record["event"] == "process_group_signal" - ] - self.assertEqual( - result, watchdog.EXIT_SIGNAL_ERROR - ) - self.assertEqual(signals, ["SIGTERM", "SIGKILL"]) - self.assertEqual( - records[-1]["classification"], "signal_error" - ) - self.assertEqual( - records[-1]["child_returncode"], - -signal.SIGKILL, - ) - assert child_pid is not None - assert grandchild_pid is not None - for process_id in (child_pid, grandchild_pid): - deadline = time.monotonic() + 2 - while ( - self._process_is_running(process_id) - and time.monotonic() < deadline - ): - time.sleep(0.01) - self.assertFalse( - self._process_is_running(process_id) - ) + records = [ + json.loads(line) + for line in stream.getvalue().splitlines() + ] + signals = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(result, watchdog.EXIT_SIGNAL_ERROR) + self.assertEqual(signals, ["SIGTERM", "SIGKILL"]) + self.assertEqual( + records[-1]["classification"], "signal_error" + ) + self.assertEqual(records[-1]["exit_code"], 7) + self.assertEqual( + records[-1]["threshold_reason"], + "guardian control failed during graceful cleanup", + ) + self.assertEqual( + records[-1]["secondary_errors"][0]["component"], + artifact_failure, + ) + self.assertEqual( + records[-1]["child_returncode"], + -signal.SIGKILL, + ) + assert child_pid is not None + assert grandchild_pid is not None + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) @unittest.skipUnless( sys.platform.startswith("linux"), From 778db6f50eae04e6c232c69b9575bdbd0747962b Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 23:56:55 -0700 Subject: [PATCH 30/32] scripts : retain watchdog artifact evidence Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- scripts/strix_memory_watchdog.py | 16 +++++++++++++ tests/test_strix_memory_watchdog.py | 37 +++++++++++++++++++++++------ 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 75f3e8852e27..a06c85de3f96 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -1492,6 +1492,7 @@ def _emit_final( process_group_status: str = "not_created", error: str | None = None, preserve_primary_on_artifact_error: bool = False, + secondary_errors: Sequence[dict[str, str]] | None = None, ) -> int: fields = _state_fields( snapshot, @@ -1504,6 +1505,8 @@ def _emit_final( fields.update(classification=classification, exit_code=exit_code) if error: fields["error"] = error + if secondary_errors: + fields["secondary_errors"] = list(secondary_errors) def record_artifact_error(exc: ArtifactError) -> None: nonlocal exit_code @@ -1868,6 +1871,19 @@ def _graceful_cleanup( preserve_primary_on_artifact_error=( guardian_control_error is not None ), + secondary_errors=( + [ + { + "component": artifact_error.component, + "detail": str(artifact_error), + } + ] + if ( + guardian_control_error is not None + and artifact_error is not None + ) + else None + ), ) diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 8bb209fcbfb4..f9cba175a8d1 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -571,14 +571,15 @@ def test_guardian_control_failure_still_kills_and_reaps_group( self, ) -> None: class FailingFinalAudit: - def __init__(self, stream: Any): + def __init__(self, stream: Any, fail_at: int): self.stream = stream self.write_count = 0 + self.fail_at = fail_at def write(self, value: str) -> int: self.write_count += 1 - if self.write_count == 3: - raise OSError("final audit write failed") + if self.write_count == self.fail_at: + raise OSError("audit write failed") return self.stream.write(value) def flush(self) -> None: @@ -608,7 +609,12 @@ def finalize(self, record: dict[str, Any]) -> None: " time.sleep(30)\n" ) for mode in ("closed", "blocked"): - for artifact_failure in ("audit", "lease"): + for artifact_failure in ( + "term_audit", + "kill_audit", + "final_audit", + "lease", + ): with self.subTest( mode=mode, artifact_failure=artifact_failure, @@ -658,13 +664,18 @@ def _assert_guardian_control_failure_cleanup( ) stream = io.StringIO() audit = watchdog.AuditLogger(stream) - if artifact_failure == "audit": + if artifact_failure.endswith("_audit"): persistent_path = Path(temp_dir) / "persistent.jsonl" persistent_stream = persistent_path.open( "w", encoding="utf-8" ) audit.persistent_stream = failing_final_audit( - persistent_stream + persistent_stream, + { + "term_audit": 1, + "kill_audit": 2, + "final_audit": 3, + }[artifact_failure], ) else: audit.lease_manager = failing_final_lease() @@ -728,7 +739,19 @@ def _assert_guardian_control_failure_cleanup( ) self.assertEqual( records[-1]["secondary_errors"][0]["component"], - artifact_failure, + ( + "audit" + if artifact_failure.endswith("_audit") + else "lease" + ), + ) + self.assertIn( + ( + "audit write failed" + if artifact_failure.endswith("_audit") + else "final lease write failed" + ), + records[-1]["secondary_errors"][0]["detail"], ) self.assertEqual( records[-1]["child_returncode"], From c2c98de7cc0e1876989f90fc49aeb0aafc2d2193 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sun, 13 Sep 2026 00:14:39 -0700 Subject: [PATCH 31/32] scripts : restore approved watchdog artifacts Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/strix-memory-watchdog.md | 2 +- scripts/strix_memory_watchdog.py | 6 ++---- tests/test_strix_memory_watchdog.py | 12 ------------ 3 files changed, 3 insertions(+), 17 deletions(-) diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index e091001c114e..11997ce5de5c 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -38,7 +38,7 @@ The wrapper performs these checks and actions: The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. -Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. Threshold overrides may only lower the 116 GiB soft and 118 GiB emergency limits. The fail-closed timing bounds are a maximum 30-second grace, maximum one-second sample interval, and maximum five-second heartbeat age. +Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The generic watchdog requires the emergency threshold to remain below 120 GiB. Final DeepSeek V4.1 validation must use the exact 116 GiB soft and 118 GiB emergency defaults because the matching preflight rejects any other thresholds. The in-process DeepSeek admission options may only lower these policy limits. The fail-closed timing bounds are a maximum 30-second grace, maximum one-second sample interval, and maximum five-second heartbeat age. The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Signal records are written immediately after each process-group signal. Child standard input, standard output, and standard error are inherited unchanged. diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py index 68f8c890ba35..a06c85de3f96 100755 --- a/scripts/strix_memory_watchdog.py +++ b/scripts/strix_memory_watchdog.py @@ -180,12 +180,10 @@ def validate(self) -> ArtifactPaths | None: raise ValueError("a command is required after --") if self.soft_bytes <= 0: raise ValueError("soft threshold must be greater than zero") - if self.soft_bytes > DEFAULT_SOFT_BYTES: - raise ValueError("soft threshold must not exceed 116 GiB") if self.emergency_bytes <= self.soft_bytes: raise ValueError("emergency threshold must be greater than soft threshold") - if self.emergency_bytes > DEFAULT_EMERGENCY_BYTES: - raise ValueError("emergency threshold must not exceed 118 GiB") + if self.emergency_bytes >= STRICT_CEILING_BYTES: + raise ValueError("emergency threshold must be below 120 GiB") if ( not math.isfinite(self.grace_seconds) or self.grace_seconds <= 0 diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py index 2de953a73d83..f9cba175a8d1 100644 --- a/tests/test_strix_memory_watchdog.py +++ b/tests/test_strix_memory_watchdog.py @@ -1328,18 +1328,6 @@ def test_configuration_rejects_non_finite_timing(self) -> None: with self.assertRaisesRegex(ValueError, "grace period"): config.validate() - def test_configuration_rejects_thresholds_above_policy(self) -> None: - with self.assertRaisesRegex(ValueError, "116 GiB"): - watchdog.WatchdogConfig( - command=("fake-command",), - soft_bytes=watchdog.DEFAULT_SOFT_BYTES + 1, - ).validate() - with self.assertRaisesRegex(ValueError, "118 GiB"): - watchdog.WatchdogConfig( - command=("fake-command",), - emergency_bytes=watchdog.DEFAULT_EMERGENCY_BYTES + 1, - ).validate() - def test_configuration_rejects_weakened_liveness_timing(self) -> None: cases = ( ( From 4f4a47fd703e0d24bf2dcf0bb93dfa6c8635d1b5 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sun, 13 Sep 2026 19:08:17 -0700 Subject: [PATCH 32/32] common : make DeepSeek defaults admission-safe Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- common/arg.cpp | 2 ++ common/common.cpp | 26 +++++++++++++---- common/common.h | 4 ++- common/fit.cpp | 30 +++++++++++++++++--- common/fit.h | 5 ++++ docs/strix-memory-watchdog.md | 2 +- include/llama.h | 2 +- src/llama-context.cpp | 6 ++-- src/llama-dsv41-admission.h | 3 +- src/llama-model.h | 1 + src/models/deepseek41.cpp | 4 +++ src/models/models.h | 1 + tests/test-arg-parser.cpp | 47 ++++++++++++++++++++++++++++++- tests/test-deepseek41-runtime.cpp | 44 ++++++++++++++++++++++++++++- tools/server/server.cpp | 3 +- 15 files changed, 162 insertions(+), 18 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index b358b04329f6..c21ac9dbbfd6 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2569,6 +2569,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex throw std::invalid_argument("error: invalid value for n_parallel\n"); } params.n_parallel = value; + params.n_parallel_explicit = value != -1; } ).set_env("LLAMA_ARG_N_PARALLEL").set_examples({LLAMA_EXAMPLE_SERVER})); } else { @@ -2577,6 +2578,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("number of parallel sequences to decode (default: %d)", params.n_parallel), [](common_params & params, int value) { params.n_parallel = value; + params.n_parallel_explicit = true; } ).set_env("LLAMA_ARG_N_PARALLEL")); } diff --git a/common/common.cpp b/common/common.cpp index 7322045663ea..320436adb38e 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1708,9 +1708,14 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.dsv41_memory_watchdog_bytes = (uint64_t) params.dsv41_memory_watchdog_mib << 20; mparams.dsv41_memory_hard_bytes = (uint64_t) params.dsv41_memory_hard_mib << 20; mparams.dsv41_memory_safety_margin_bytes = (uint64_t) params.dsv41_memory_safety_margin_mib << 20; - mparams.dsv41_admission_context = params.n_ctx == 0 ? 32768 : params.n_ctx; + const uint32_t dsv41_admission_sequences = params.n_parallel_explicit ? + std::max(params.n_parallel, 1) : 1; + mparams.dsv41_admission_context = + params.n_ctx_auto_sized && !params.n_parallel_explicit ? + std::max(params.kv_unified_per_slot, 1) : + (params.n_ctx == 0 ? 32768 : params.n_ctx); mparams.dsv41_admission_batch = std::max(params.n_batch, 1); - mparams.dsv41_admission_sequences = params.n_parallel; + mparams.dsv41_admission_sequences = dsv41_admission_sequences; mparams.dsv41_admission_ubatch = std::min( mparams.dsv41_admission_batch, static_cast(params.n_ubatch_explicit ? std::max(params.n_ubatch, 1) : 32)); @@ -1718,7 +1723,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.dsv41_admission_batch : std::min(params.n_outputs_max, mparams.dsv41_admission_batch); mparams.dsv41_admission_outputs = std::max( - mparams.dsv41_admission_outputs, std::max(params.n_parallel, 1)); + mparams.dsv41_admission_outputs, dsv41_admission_sequences); mparams.dsv41_admission_outputs_per_seq = params.n_outputs_max_per_seq == 0 ? mparams.dsv41_admission_outputs : std::min(std::max(params.n_outputs_max_per_seq, 1), mparams.dsv41_admission_outputs); @@ -1750,9 +1755,20 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { void common_context_params_apply_arch_defaults( const char * architecture, - const common_params & params, + common_params & params, llama_context_params & cparams) { - if (architecture != nullptr && strcmp(architecture, "deepseek41") == 0 && !params.n_ubatch_explicit) { + if (architecture == nullptr || strcmp(architecture, "deepseek41") != 0) { + return; + } + if (!params.n_parallel_explicit) { + params.n_parallel = 1; + cparams.n_seq_max = 1; + if (params.n_ctx_auto_sized) { + params.n_ctx = params.kv_unified_per_slot; + cparams.n_ctx = params.n_ctx; + } + } + if (!params.n_ubatch_explicit) { cparams.n_ubatch = std::min(cparams.n_batch, 32); } } diff --git a/common/common.h b/common/common.h index 614799dc4ce5..4a9007375002 100644 --- a/common/common.h +++ b/common/common.h @@ -489,12 +489,14 @@ struct ggml_opt_optimizer_params common_opt_lr_pars(void * userdata); struct common_params { int32_t n_predict = -1; // max. number of new tokens to predict, -1 == no limit int32_t n_ctx = 0; // context size, 0 == context the model was trained with + bool n_ctx_auto_sized = false; int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_ubatch = 512; // physical batch size for prompt processing (must be >=32 to use BLAS) bool n_ubatch_explicit = false; int32_t n_keep = 0; // number of tokens to keep from initial prompt int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) int32_t n_parallel = 1; // number of parallel sequences to decode + bool n_parallel_explicit = false; int32_t n_sequences = 1; // number of sequences to decode int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch) int32_t n_outputs_max_per_seq = 1; // max outputs per sequence @@ -1004,7 +1006,7 @@ struct llama_model_params common_model_params_to_llama ( common_params & struct llama_context_params common_context_params_to_llama(const common_params & params); void common_context_params_apply_arch_defaults( const char * architecture, - const common_params & params, + common_params & params, struct llama_context_params & cparams); // clear LoRA adapters from context, then apply new list of adapters diff --git a/common/fit.cpp b/common/fit.cpp index 472ee6a113f0..905f930958e8 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -27,6 +27,30 @@ class common_params_fit_exception : public std::runtime_error { using std::runtime_error::runtime_error; }; +void common_fit_context_params_apply_arch_defaults( + const char * architecture, + const llama_model_params & mparams, + llama_context_params & cparams) { + if (architecture == nullptr || strcmp(architecture, "deepseek41") != 0) { + return; + } + + cparams.n_ctx = cparams.n_ctx == 0 ? + mparams.dsv41_admission_context : + std::min(cparams.n_ctx, mparams.dsv41_admission_context); + cparams.n_batch = std::min(cparams.n_batch, mparams.dsv41_admission_batch); + cparams.n_seq_max = std::min(cparams.n_seq_max, mparams.dsv41_admission_sequences); + cparams.n_ubatch = cparams.n_ubatch == 0 || cparams.n_ubatch == UINT32_MAX ? + mparams.dsv41_admission_ubatch : + std::min(cparams.n_ubatch, mparams.dsv41_admission_ubatch); + cparams.n_outputs_max = cparams.n_outputs_max == 0 ? + mparams.dsv41_admission_outputs : + std::min(cparams.n_outputs_max, mparams.dsv41_admission_outputs); + cparams.n_outputs_max_per_seq = cparams.n_outputs_max_per_seq == 0 ? + mparams.dsv41_admission_outputs_per_seq : + std::min(cparams.n_outputs_max_per_seq, mparams.dsv41_admission_outputs_per_seq); +} + static std::vector common_get_device_memory_data_impl( const char * path_model, const llama_model_params * mparams, @@ -65,10 +89,8 @@ static std::vector common_get_device_memory_data_impl( llama_context_params cparams_copy = *cparams; char architecture[128] = {}; - if (llama_model_meta_val_str(model, "general.architecture", architecture, sizeof(architecture)) >= 0 && - strcmp(architecture, "deepseek41") == 0 && - cparams_copy.n_ubatch != mparams_copy.dsv41_admission_ubatch) { - cparams_copy.n_ubatch = mparams_copy.dsv41_admission_ubatch; + if (llama_model_meta_val_str(model, "general.architecture", architecture, sizeof(architecture)) >= 0) { + common_fit_context_params_apply_arch_defaults(architecture, mparams_copy, cparams_copy); } llama_context * ctx = llama_init_from_model(model, cparams_copy); if (ctx == nullptr) { diff --git a/common/fit.h b/common/fit.h index 824d386b07a1..e4fdee4c0ebb 100644 --- a/common/fit.h +++ b/common/fit.h @@ -21,6 +21,11 @@ struct common_fit_extra_model { bool shares_model; }; +void common_fit_context_params_apply_arch_defaults( + const char * architecture, + const llama_model_params & mparams, + llama_context_params & cparams); + // fits mparams and cparams to free device memory (assumes system memory is unlimited) // - returns true if the parameters could be successfully modified to fit device memory // - this function is NOT thread safe because it modifies the global llama logger state diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md index 11997ce5de5c..631f3ab88531 100644 --- a/docs/strix-memory-watchdog.md +++ b/docs/strix-memory-watchdog.md @@ -9,7 +9,7 @@ HIP_LAUNCH_BLOCKING=1 \ ./scripts/strix_memory_watchdog.py -- \ ./build/bin/llama-server \ -m /mnt/models/deepseek-v41/DeepSeek-V4.1-Flash-Q2.gguf \ - -c 32768 -b 2048 -ub 32 -ngl 99 -dev ROCm0 \ + -c 32768 -b 2048 -ub 32 -np 1 -ngl 99 -dev ROCm0 \ --expert-cache-slots 192 --expert-cache-mib 72900 ``` diff --git a/include/llama.h b/include/llama.h index a50b35fec588..ee66c90f2ad6 100644 --- a/include/llama.h +++ b/include/llama.h @@ -403,7 +403,7 @@ extern "C" { struct llama_context_params { uint32_t n_ctx; // text context, 0 = from model uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode - uint32_t n_ubatch; // physical maximum batch size + uint32_t n_ubatch; // physical maximum batch size, UINT32_MAX = model default uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3f719721d74b..77f7fcbbb629 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -261,7 +261,9 @@ llama_context::llama_context( // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; - cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch); + const uint32_t n_ubatch = params.n_ubatch == UINT32_MAX ? + model.default_context_ubatch() : params.n_ubatch; + cparams.n_ubatch = std::min(cparams.n_batch, n_ubatch == 0 ? params.n_batch : n_ubatch); cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max; cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ? @@ -3749,7 +3751,7 @@ llama_context_params llama_context_default_params() { llama_context_params result = { /*.n_ctx =*/ 512, /*.n_batch =*/ 2048, - /*.n_ubatch =*/ 512, + /*.n_ubatch =*/ UINT32_MAX, /*.n_seq_max =*/ 1, /*.n_rs_seq =*/ 0, /*.n_outputs_max =*/ 0, diff --git a/src/llama-dsv41-admission.h b/src/llama-dsv41-admission.h index 47f1e34a2093..a377ef43cc58 100644 --- a/src/llama-dsv41-admission.h +++ b/src/llama-dsv41-admission.h @@ -14,6 +14,7 @@ static constexpr uint64_t LLAMA_DSV41_WATCHDOG_EMERGENCY_BYTES = 118ULL << 30; static constexpr uint64_t LLAMA_DSV41_ADMISSION_HARD_BYTES = 120ULL << 30; static constexpr uint64_t LLAMA_DSV41_ADMISSION_MARGIN_BYTES = 2ULL << 30; static constexpr uint32_t LLAMA_DSV41_ADMISSION_CONTEXT = 32768; +static constexpr uint32_t LLAMA_DSV41_ADMISSION_UBATCH = 32; struct llama_dsv41_host_memory { uint64_t total = 0; @@ -35,7 +36,7 @@ struct llama_dsv41_admission_params { uint32_t n_ctx = LLAMA_DSV41_ADMISSION_CONTEXT; uint32_t n_batch = 2048; uint32_t n_seq = 1; - uint32_t n_ubatch = 2048; + uint32_t n_ubatch = LLAMA_DSV41_ADMISSION_UBATCH; uint32_t n_outputs_max = 2048; uint32_t n_outputs_max_per_seq = 2048; uint32_t n_vocab = 0; diff --git a/src/llama-model.h b/src/llama-model.h index eab504a06374..9c10b7776f27 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -776,6 +776,7 @@ struct llama_model { virtual void acquire_runtime_context() const {} virtual void release_runtime_context() const {} virtual uint32_t default_context_size() const { return 0; } + virtual uint32_t default_context_ubatch() const { return 512; } virtual void validate_context_params(const llama_cparams &) const {} virtual void validate_memory_accounting(uint64_t, uint64_t) const {} diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 8b55df5efaf0..821c79c7b907 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -827,6 +827,10 @@ uint32_t llama_model_deepseek41::default_context_size() const { return admission ? admission->result.n_ctx : LLAMA_DSV41_ADMISSION_CONTEXT; } +uint32_t llama_model_deepseek41::default_context_ubatch() const { + return admission ? admission->result.n_ubatch : LLAMA_DSV41_ADMISSION_UBATCH; +} + void llama_model_deepseek41::validate_context_params(const llama_cparams & cparams) const { if (!admission) { throw std::runtime_error("DeepSeek V4.1 context has no host-memory admission result"); diff --git a/src/models/models.h b/src/models/models.h index 1ecf09b49d98..8235fb26c914 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1339,6 +1339,7 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { void acquire_runtime_context() const override; void release_runtime_context() const override; uint32_t default_context_size() const override; + uint32_t default_context_ubatch() const override; void validate_context_params(const llama_cparams & cparams) const override; void validate_memory_accounting(uint64_t state_bytes, uint64_t graph_workspace_bytes) const override; diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 32dd9fb3ed79..31bea7ec5201 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -1,6 +1,7 @@ #include "arg.h" #include "common.h" #include "download.h" +#include "fit.h" #include "llama.h" #include "speculative.h" @@ -174,12 +175,17 @@ static void test(void) { common_params default_params; const auto model_params = common_model_params_to_llama(default_params); auto context_params = common_context_params_to_llama(default_params); + assert(model_params.dsv41_admission_sequences == 1); assert(model_params.dsv41_admission_ubatch == 32); common_context_params_apply_arch_defaults("deepseek41", default_params, context_params); + assert(default_params.n_parallel == 1); + assert(context_params.n_seq_max == 1); assert(context_params.n_ubatch == 32); common_params explicit_params; - std::vector explicit_argv = { "binary_name", "-m", "model_file.gguf", "-ub", "37" }; + std::vector explicit_argv = { + "binary_name", "-m", "model_file.gguf", "-ub", "37", "-np", "3", + }; assert(common_params_parse( explicit_argv.size(), list_str_to_char(explicit_argv).data(), @@ -188,11 +194,50 @@ static void test(void) { const auto explicit_model_params = common_model_params_to_llama(explicit_params); auto explicit_context_params = common_context_params_to_llama(explicit_params); assert(explicit_params.n_ubatch_explicit); + assert(explicit_params.n_parallel_explicit); + assert(explicit_model_params.dsv41_admission_sequences == 3); assert(explicit_model_params.dsv41_admission_ubatch == 37); common_context_params_apply_arch_defaults("deepseek41", explicit_params, explicit_context_params); + assert(explicit_params.n_parallel == 3); + assert(explicit_context_params.n_seq_max == 3); assert(explicit_context_params.n_ubatch == 37); } + { + common_params server_params; + std::vector server_argv = { "binary_name", "-m", "model_file.gguf" }; + assert(common_params_parse( + server_argv.size(), + list_str_to_char(server_argv).data(), + server_params, + LLAMA_EXAMPLE_SERVER)); + assert(server_params.n_parallel == -1); + assert(!server_params.n_parallel_explicit); + + server_params.n_parallel = 4; + server_params.kv_unified = true; + server_params.kv_unified_per_slot = 32768; + server_params.n_ctx = 4 * server_params.kv_unified_per_slot; + server_params.n_ctx_auto_sized = true; + + const auto model_params = common_model_params_to_llama(server_params); + auto context_params = common_context_params_to_llama(server_params); + assert(model_params.dsv41_admission_sequences == 1); + assert(model_params.dsv41_admission_context == 32768); + + auto fit_context_params = context_params; + common_fit_context_params_apply_arch_defaults("deepseek41", model_params, fit_context_params); + assert(fit_context_params.n_seq_max == 1); + assert(fit_context_params.n_ctx == 32768); + assert(fit_context_params.n_ubatch == 32); + + common_context_params_apply_arch_defaults("deepseek41", server_params, context_params); + assert(server_params.n_parallel == 1); + assert(server_params.n_ctx == 32768); + assert(context_params.n_seq_max == 1); + assert(context_params.n_ctx == 32768); + } + std::vector argv; printf("test-arg-parser: test invalid usage\n\n"); diff --git a/tests/test-deepseek41-runtime.cpp b/tests/test-deepseek41-runtime.cpp index 985ff8a72009..f11f9436f0cb 100644 --- a/tests/test-deepseek41-runtime.cpp +++ b/tests/test-deepseek41-runtime.cpp @@ -585,7 +585,7 @@ static void test_graph_construction() { ggml_free(ctx); } -struct reservation_test_model final : llama_model { +struct reservation_test_model : llama_model { mutable bool active = false; mutable uint32_t acquisitions = 0; mutable uint32_t releases = 0; @@ -622,6 +622,18 @@ struct reservation_test_model final : llama_model { } }; +struct default_ubatch_test_model final : reservation_test_model { + uint32_t default_context_ubatch() const override { + return 32; + } + + void validate_context_params(const llama_cparams & cparams) const override { + if (cparams.n_ubatch != 32) { + throw std::runtime_error("unexpected context ubatch"); + } + } +}; + static llama_context_params reservation_context_params() { llama_context_params params = llama_context_default_params(); params.n_ctx = 32; @@ -631,6 +643,35 @@ static llama_context_params reservation_context_params() { return params; } +static void test_default_context_ubatch() { + default_ubatch_test_model model; + llama_context_params params = llama_context_default_params(); + params.n_ctx = 1024; + params.n_batch = 1024; + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + check(params.n_ubatch == UINT32_MAX, "public context defaults do not preserve model-aware ubatch selection"); + + llama_context * context = llama_init_from_model(&model, params); + check(context != nullptr, "public default context did not use the model ubatch"); + check(llama_n_ubatch(context) == 32, "public default context resolved the wrong model ubatch"); + llama_free(context); + + params.n_ubatch = 512; + context = llama_init_from_model(&model, params); + check(context == nullptr, "explicit context ubatch was silently replaced by the model default"); + llama_free(context); + + reservation_test_model standard_model; + params = llama_context_default_params(); + params.n_ctx = 1024; + params.n_batch = 1024; + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + context = llama_init_from_model(&standard_model, params); + check(context != nullptr, "public default context failed for a standard model"); + check(llama_n_ubatch(context) == 512, "standard model default ubatch changed"); + llama_free(context); +} + static void test_runtime_context_reservation() { reservation_test_model model; auto first = std::make_unique(model, reservation_context_params()); @@ -684,6 +725,7 @@ int main() { test_output_collapse(); test_graph_contract(); test_graph_construction(); + test_default_context_ubatch(); test_runtime_context_reservation(); return 0; } diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 22378b38c5ef..bbe8e173dbb8 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -150,7 +150,7 @@ int llama_server(common_params & params, int argc, char ** argv) { } if (params.n_parallel < 0) { - SRV_TRC("%s", "n_parallel is set to auto, using n_parallel = 4 and kv_unified = true\n"); + SRV_TRC("%s", "n_parallel is set to auto, using n_parallel = 4 unless the model requires a safer default\n"); params.n_parallel = 4; params.kv_unified = true; @@ -165,6 +165,7 @@ int llama_server(common_params & params, int argc, char ** argv) { if (ctx_pool_auto_sized) { params.n_ctx = params.n_parallel * params.kv_unified_per_slot; + params.n_ctx_auto_sized = true; SRV_INF("--kv-unified-per-slot: sizing KV pool to n_parallel * kv_unified_per_slot = %d * %d = %d\n", params.n_parallel, params.kv_unified_per_slot, params.n_ctx); }