diff --git a/README.md b/README.md index 3067a7a4e71e..e073b9fd1592 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 | +| 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/common/arg.cpp b/common/arg.cpp index c55b9a58aef5..c21ac9dbbfd6 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( @@ -2568,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 { @@ -2576,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")); } @@ -2808,9 +2811,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; @@ -2818,14 +2821,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..320436adb38e 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 @@ -1699,6 +1704,32 @@ 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; + 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 = 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)); + 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, 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); + mparams.dsv41_admission_type_k = params.cache_type_k; + mparams.dsv41_procfs_root = params.dsv41_procfs_root.c_str(); + mparams.dsv41_admission_offload_kqv = !params.no_kv_offload; if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; @@ -1722,6 +1753,26 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { return mparams; } +void common_context_params_apply_arch_defaults( + const char * architecture, + common_params & params, + llama_context_params & cparams) { + 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); + } +} + 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 a838f90b52a8..4a9007375002 100644 --- a/common/common.h +++ b/common/common.h @@ -489,11 +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 @@ -626,8 +629,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 @@ -996,6 +1004,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, + 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..905f930958e8 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -6,9 +6,10 @@ #include #include -#include #include +#include #include +#include #include #include @@ -26,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, @@ -62,7 +87,12 @@ 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) { + 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) { llama_model_free(model); llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); 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 new file mode 100644 index 000000000000..631f3ab88531 --- /dev/null +++ b/docs/strix-memory-watchdog.md @@ -0,0 +1,111 @@ +# 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 +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 -np 1 -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. + +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. +- 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 `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`. + +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 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. + +## 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 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. 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 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`, `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 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. 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: + +- 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, 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. +- 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. + +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 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 | +| ---: | --- | +| 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 | +| 8 | lease, heartbeat, or persistent audit failure | +| 70 | unexpected post-launch error | +| 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/include/llama.h b/include/llama.h index 4720bf625606..ee66c90f2ad6 100644 --- a/include/llama.h +++ b/include/llama.h @@ -348,10 +348,24 @@ 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_batch; + uint32_t dsv41_admission_sequences; + uint32_t dsv41_admission_ubatch; + uint32_t dsv41_admission_outputs; + uint32_t dsv41_admission_outputs_per_seq; + enum ggml_type dsv41_admission_type_k; + 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; @@ -376,6 +390,7 @@ extern "C" { bool ple_on_disk; // keep the n-gram hash-embedding table (per_layer_token_embd) on disk: never // mapped or loaded, the rows a batch needs are read from the file (qwen4exp) bool ple_direct_io; // read those rows with O_DIRECT, bypassing the page cache + bool dsv41_admission_offload_kqv; }; struct llama_sampler_seq_config { @@ -388,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/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py new file mode 100755 index 000000000000..a06c85de3f96 --- /dev/null +++ b/scripts/strix_memory_watchdog.py @@ -0,0 +1,2575 @@ +#!/usr/bin/env python3 + +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 +from datetime import datetime, timezone +from pathlib import Path +from typing import IO, Any, 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 +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 = 2 +HEARTBEAT_FORMAT = "strix-memory-watchdog-heartbeat" +HEARTBEAT_VERSION = 2 +PR_SET_PDEATHSIG = 1 +LEASE_GUARD_SIGNAL = signal.SIGUSR1 + +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_LEASE_ERROR = 8 +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.SIGHUP, signal.SIGINT, signal.SIGTERM) + + +class ProcfsError(RuntimeError): + pass + + +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 + super().__init__(signal.Signals(signal_number).name) + + +class ProcessHandle(Protocol): + pid: int + + def poll(self) -> int | None: + ... + + 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: + 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, value) + 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 + available_bytes: int + active_swaps: tuple[str, ...] + + @property + 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 ArtifactPaths: + lease: Path + heartbeat: Path + audit: Path + + +@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 + 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) -> ArtifactPaths | 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 + 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 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 " + "and at most 5 seconds" + ) + 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" + ) + return paths + return None + + +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) + + +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 _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, + grace_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, + ) + current_timeout_seconds = pulse_timeout_seconds + 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() + 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: + 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( + 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, + grace_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), + str(grace_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: + 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" + ) + try: + descriptor = os.open( + temp_path, + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + 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()) + 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() + ) + executable_path = ( + self.process_procfs_root + / str(os.getpid()) + / "exe" + ).resolve() + 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 + ) + 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), + } + + 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, 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, + "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_executable_path": ( + watchdog_identity["executable_path"] + ), + "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, + "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, + "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), + "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", + {"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) + + 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: + 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 + + +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_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: 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) + 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" + ) + 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", + ) + _, _, 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") + 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" + ) + if _sha256_bytes(live_cmdline) != _require_string( + lease.get("watchdog_command_sha256"), + "watchdog_command_sha256", + ): + raise LeaseValidationError("watchdog command line does not match") + 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 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 ( + 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 + 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 ( + 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 + ): + raise LeaseValidationError( + "watchdog, guardian, child, 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 is not None and command_sha256 != _command_sha256( + expected_command + ): + 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") + 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): + 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") + 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: + 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) + 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" + ) + 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: + 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 + _, _, 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 Exception: + 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 Exception: + _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, + stream: IO[str], + 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 + 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 + | 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" + ) + except OSError as exc: + detail = exc.strerror or str(exc) + raise ArtifactError( + "audit", + 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: + 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 = { + "timestamp": _timestamp_utc(self.wall_clock), + "event": event, + **fields, + } + line = ( + json.dumps(record, sort_keys=True, separators=(",", ":")) + + "\n" + ) + 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: + 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, + "audit_record_sha256": self.last_record_sha256, + } + ) + + 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: + 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, + preserve_primary_on_artifact_error: bool = False, + secondary_errors: Sequence[dict[str, 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 + if secondary_errors: + fields["secondary_errors"] = list(secondary_errors) + + 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: + audit.finalize(record) + except ArtifactError as exc: + audit.disable_component(exc.component) + record_artifact_error(exc) + emit_final_record() + 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: + 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 _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: + return False + except PermissionError: + return True + 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: 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, Any], +) -> None: + for signal_number, handler in previous.items(): + signal.signal(signal_number, handler) + + +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), + ) + + 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: + return _emit_final( + audit, + "termination_timeout", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "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, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + group_status, + error, + ) + + +def _graceful_cleanup( + audit: AuditLogger, + child: ProcessHandle, + snapshot: HostSnapshot, + peak_used_bytes: int, + classification: str, + exit_code: int, + reason: str, + 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", + escalation_result: tuple[str, int, str] | None = None, + error: str | None = None, +) -> 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 + ) + 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) + 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 ( + guardian_control_error is None + and monotonic() < deadline + ): + child.poll() + if not group_alive(child.pid): + break + if ( + isinstance(child, GuardianProcess) + and child.poll() is None + ): + try: + child.pulse() + except ProcessGroupError as exc: + guardian_control_error = exc + break + sleeper(min(0.05, deadline - monotonic())) + child.poll() + if ( + guardian_control_error is not None + or group_alive(child.pid) + ): + escalated = True + process_group_status = signal_group( + child.pid, signal.SIGKILL + ) + 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", + **_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, + "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), + ) + + 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 and guardian_control_error is None: + classification = "lease_error" + exit_code = EXIT_LEASE_ERROR + error = f"{artifact_error.component}: {artifact_error}" + return _emit_final( + audit, + classification, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + process_group_status, + error, + 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 + ), + ) + + +def _monitor_child( + config: WatchdogConfig, + reader: ProcfsReader, + audit: AuditLogger, + child: ProcessHandle, + 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: + soft_deadline: float | None = None + + while True: + 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, + ( + ( + "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, + classification, + exit_code, + reason, + state.snapshot, + state.peak_used_bytes, + child, + child_returncode, + "leader_exited", + ) + + now = monotonic() + if soft_deadline is not None and now >= soft_deadline: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "grace_timeout", + EXIT_GRACE_TIMEOUT, + "soft-threshold grace period expired", + signal_group, + ) + + try: + state.snapshot = reader.read_snapshot() + except ProcfsError as exc: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + signal_group, + ) + + state.peak_used_bytes = max( + state.peak_used_bytes, state.snapshot.used_bytes + ) + if state.snapshot.active_swaps: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "swap_appeared", + EXIT_SWAP_ACTIVE, + "active swap appeared during execution", + signal_group, + ) + if state.snapshot.used_bytes >= config.emergency_bytes: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "emergency_limit", + EXIT_EMERGENCY_LIMIT, + "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 + ): + 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", + state.snapshot, + state.peak_used_bytes, + child, + child.poll(), + "signal_error", + 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, + state.peak_used_bytes, + child, + child.poll(), + group_status, + "used_bytes >= soft_bytes", + ), + "signal": "SIGTERM", + "grace_deadline_monotonic": soft_deadline, + } + 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, + ) + 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) + 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: + 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, + group_alive: Callable[[int], bool] | None = None, + monotonic: Callable[[], float] | None = None, + 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 + 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 + + 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: + return _emit_final( + audit, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + None, + None, + error=str(exc), + ) + + 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, + error=f"{exc.component}: {exc}", + ) + + 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, + ) + + previous_mask = signal.pthread_sigmask( + signal.SIG_BLOCK, PARENT_SIGNALS + ) + mask_restored = False + 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_environment = os.environ.copy() + if lease_manager is not None: + 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) + ), + } + ) + if use_guardian: + child = _launch_guardian( + config.command, + child_environment, + config.heartbeat_max_age_seconds, + config.grace_seconds + 1.0, + launch_mask, + ) + elif lease_manager is not None: + 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, subprocess.SubprocessError) 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 + ) + if lease_manager is not None: + lease_manager.start(child, audit) + audit.lease_manager = lease_manager + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + mask_restored = True + 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, + child, + state, + signal_group, + group_alive, + child.pulse if isinstance(child, GuardianProcess) else lambda: None, + monotonic, + sleeper, + ) + except ArtifactError as exc: + _set_parent_signal_handlers(signal.SIG_IGN) + 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, + child, + 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: + 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, + 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, + error=f"{type(exc).__name__}: {exc}", + ) + finally: + if not mask_restored: + 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: + 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( + "--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, + 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, + lease_path=args.lease_path, + heartbeat_path=args.heartbeat_path, + audit_path=args.audit_path, + heartbeat_max_age_seconds=args.heartbeat_max_age_seconds, + ) + + +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) < 7 or arguments[5] != "--": + return EXIT_LAUNCH_ERROR + try: + return _guardian_main( + int(arguments[1]), + int(arguments[2]), + _positive_float(arguments[3]), + _positive_float(arguments[4]), + tuple(arguments[6:]), + ) + except (OSError, ValueError): + return EXIT_LAUNCH_ERROR + config = parse_args(arguments) + 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), + ) + 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__": + sys.exit(main()) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 971d75f5706a..2ea7e9a95947 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,6 +16,7 @@ set(LLAMA_CORE_SOURCES 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 ad8f71bac9be..77f7fcbbb629 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; @@ -130,7 +141,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; @@ -249,12 +261,18 @@ 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 ? 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 // re-reserve when graph nodes change. @@ -482,14 +500,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) { @@ -707,6 +728,14 @@ void llama_context::sched_reserve() { } if (memory) { memory->set_graph_workspace_size(graph_workspace_size); + uint64_t state_bytes = 0; + for (const auto & entry : memory->memory_breakdown()) { + if (entry.second > std::numeric_limits::max() - state_bytes) { + throw std::runtime_error("memory state allocation byte count overflow"); + } + state_bytes += entry.second; + } + model.validate_memory_accounting(state_bytes, graph_workspace_size); } if (n_nodes_pp == n_nodes_tg) { @@ -1189,6 +1218,11 @@ 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) { + 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; // TODO: not sure yet if we want to reserve here @@ -1198,6 +1232,11 @@ 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) { + 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; cparams.embeddings_nextn_masked = masked; } @@ -1207,6 +1246,11 @@ 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) { + 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; // note: without this reserve, the draft acceptance drops to zero. not sure why - this is unexpected @@ -1217,6 +1261,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); @@ -1481,6 +1534,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; @@ -1719,6 +1776,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); @@ -3690,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-context.h b/src/llama-context.h index b089e1267aaf..272f375b0613 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,8 @@ 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; + 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 new file mode 100644 index 000000000000..6fbf52eb110a --- /dev/null +++ b/src/llama-dsv41-admission.cpp @@ -0,0 +1,496 @@ +#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 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(); + 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_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"); + } + // 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"); +} + +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"); + 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_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(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_outputs_max, "output tokens"), + sizeof(int32_t), + "output tokens"); + 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( + 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_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"); + } + 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 > 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 <= 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"); + } + 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_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"); + } + + 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"); + 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"); + } + 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 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"); + } + 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 && + 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 (params.state_bytes == 0) { + reject("state", result, "exact no-allocation state size is missing"); + } + result.state_bytes = params.state_bytes; + 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_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"); + 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"); + + 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"); + } + 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"); + const uint64_t fit_slots = (result.soft_bytes - result.fixed_bytes)/bytes_per_slot; + const uint64_t selected = std::min(slot_cap, fit_slots); + 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"); + } + 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"); + } + result.category = "accepted"; + return result; +} + +llama_dsv41_admission_result llama_dsv41_validate_runtime_memory( + const llama_dsv41_admission_result & admitted, + uint64_t state_bytes, + uint64_t graph_workspace_bytes) { + if (admitted.category != "accepted" || state_bytes == 0 || graph_workspace_bytes == 0 || + admitted.fixed_bytes < admitted.state_bytes || + admitted.fixed_bytes - admitted.state_bytes < admitted.graph_workspace_bytes || + admitted.projected_bytes < admitted.state_bytes || + admitted.projected_bytes - admitted.state_bytes < admitted.graph_workspace_bytes) { + reject("runtime", admitted, "runtime memory accounting inputs are invalid"); + } + + llama_dsv41_admission_result result = admitted; + result.fixed_bytes -= result.state_bytes; + result.fixed_bytes -= result.graph_workspace_bytes; + result.projected_bytes -= result.state_bytes; + result.projected_bytes -= result.graph_workspace_bytes; + result.state_bytes = state_bytes; + result.graph_workspace_bytes = graph_workspace_bytes; + result.fixed_bytes = checked_add(result.fixed_bytes, state_bytes, "runtime fixed bytes"); + result.fixed_bytes = checked_add(result.fixed_bytes, graph_workspace_bytes, "runtime fixed bytes"); + result.projected_bytes = checked_add(result.projected_bytes, state_bytes, "runtime projected bytes"); + result.projected_bytes = checked_add( + result.projected_bytes, graph_workspace_bytes, "runtime projected bytes"); + + if (result.projected_bytes > result.soft_bytes) { + reject("runtime_workspace", result, "measured runtime memory exceeds the admitted soft limit"); + } + if (result.projected_bytes > result.host_total) { + reject("runtime_workspace", result, "measured runtime memory exceeds physical host memory"); + } + if (result.projected_bytes >= result.hard_bytes) { + reject("runtime_workspace", result, "measured runtime memory is not strictly below the hard limit"); + } + return result; +} + +std::string llama_dsv41_admission_result::describe() const { + return format( + "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, " + "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(), + 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, + (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, + required_expert_slots, + 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, + (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..a377ef43cc58 --- /dev/null +++ b/src/llama-dsv41-admission.h @@ -0,0 +1,103 @@ +#pragma once + +#include "llama-expert-store.h" + +#include "ggml-backend.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; +static constexpr uint32_t LLAMA_DSV41_ADMISSION_UBATCH = 32; + +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; + uint64_t state_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 = LLAMA_DSV41_ADMISSION_UBATCH; + 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; + 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_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; + 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 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; +}; + +llama_dsv41_host_memory llama_dsv41_read_host_memory(const std::string & procfs_root); + +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_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, + uint64_t dense_tensor_bytes, + const std::vector & expert_tensors, + const llama_dsv41_admission_params & params); + +llama_dsv41_admission_result llama_dsv41_validate_runtime_memory( + const llama_dsv41_admission_result & admitted, + uint64_t state_bytes, + uint64_t graph_workspace_bytes); 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/src/llama-memory-dsv41.cpp b/src/llama-memory-dsv41.cpp index 2fd088f5a962..573b3c41fb42 100644 --- a/src/llama-memory-dsv41.cpp +++ b/src/llama-memory-dsv41.cpp @@ -470,6 +470,28 @@ llama_memory_dsv41::llama_memory_dsv41( model, type_k, offload, n_ctx, n_seq, n_ubatch, std::move(engram))) { } +uint64_t llama_dsv41_measure_model_state_bytes( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch) { + llama_dsv41_memory_config config = + make_model_config(model, type_k, offload, n_ctx, n_seq, n_ubatch, nullptr); + config.no_alloc = true; + const llama_memory_dsv41 memory(std::move(config)); + + uint64_t total = 0; + for (const auto & entry : memory.memory_breakdown()) { + if (entry.second > std::numeric_limits::max() - total) { + throw std::runtime_error("DeepSeek V4.1 state allocation byte count overflow"); + } + total += entry.second; + } + return total; +} + llama_memory_dsv41::~llama_memory_dsv41() = default; llama_memory_context_ptr llama_memory_dsv41::init_batch( diff --git a/src/llama-memory-dsv41.h b/src/llama-memory-dsv41.h index 0eb074c12fe5..1ca8def9b87a 100644 --- a/src/llama-memory-dsv41.h +++ b/src/llama-memory-dsv41.h @@ -14,6 +14,14 @@ struct ggml_tensor; struct llama_model; +uint64_t llama_dsv41_measure_model_state_bytes( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch); + struct llama_dsv41_memory_config { uint32_t n_ctx = 0; uint32_t n_seq = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 895c29be0489..8aec9d15b7ef 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2545,6 +2545,16 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_DEEPSEEK41: { + if (params.type_k != this->params.dsv41_admission_type_k || + cparams.offload_kqv != this->params.dsv41_admission_offload_kqv) { + throw std::runtime_error(format( + "DeepSeek V4.1 memory parameters differ from admission: " + "type_k=%s, admitted_type_k=%s, offload_kqv=%s, admitted_offload_kqv=%s", + ggml_type_name(params.type_k), + ggml_type_name(this->params.dsv41_admission_type_k), + cparams.offload_kqv ? "true" : "false", + this->params.dsv41_admission_offload_kqv ? "true" : "false")); + } const auto & model_dsv41 = static_cast(*this); res = new llama_memory_dsv41( *this, @@ -2847,6 +2857,18 @@ 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_batch =*/ 2048, + /*.dsv41_admission_sequences =*/ 1, + /*.dsv41_admission_ubatch =*/ 32, + /*.dsv41_admission_outputs =*/ 2048, + /*.dsv41_admission_outputs_per_seq =*/ 2048, + /*.dsv41_admission_type_k =*/ GGML_TYPE_F16, + /*.dsv41_procfs_root =*/ "/proc", /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, @@ -2859,6 +2881,7 @@ llama_model_params llama_model_default_params() { /*.load_mtp =*/ false, /*.ple_on_disk =*/ false, /*.ple_direct_io =*/ true, + /*.dsv41_admission_offload_kqv =*/ true, }; return result; diff --git a/src/llama-model.h b/src/llama-model.h index 2bc5b96e4d92..9c10b7776f27 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -775,6 +775,10 @@ struct llama_model { virtual void release_runtime_work_after_sync(ggml_backend_sched_t) const { release_runtime_work(); } 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 {} // 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 a7e1094f07fc..821c79c7b907 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -1,10 +1,14 @@ +#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 "llama-memory-dsv41.h" #include "models.h" +#include "ggml-alloc.h" + #include #include #include @@ -37,6 +41,10 @@ struct llama_model_deepseek41::engram_model { std::array extents; }; +struct llama_model_deepseek41::admission_model { + llama_dsv41_admission_result result; +}; + std::unique_ptr llama_model_deepseek41::create_memory_engram_runtime( size_t max_tokens) const { return engram ? @@ -607,24 +615,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"); - } - 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](const llama_expert_store_tensor & tensor) { - return select_moe_buft( - tensor.layer, tensor.type, tensor.ne[0], tensor.ne[1], params.expert_cache_slots); - }); + 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); + } 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); output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, 0); @@ -667,32 +674,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), @@ -709,6 +699,96 @@ 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_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 = std::min( + admission_params.n_batch, + params.dsv41_admission_ubatch == 0 ? + 32U : 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.state_bytes = llama_dsv41_measure_model_state_bytes( + *this, + params.dsv41_admission_type_k, + params.dsv41_admission_offload_kqv, + admission_params.n_ctx, + admission_params.n_seq, + admission_params.n_ubatch); + admission_params.direct_io = 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)); + 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); + + 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](const llama_expert_store_tensor & tensor) { + return select_moe_buft( + tensor.layer, tensor.type, tensor.ne[0], tensor.ne[1], admission->result.expert_slots); + }); + + 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 { @@ -743,6 +823,66 @@ 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; +} + +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"); + } + 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 || + 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_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, + output_rows, + n_outputs_max_per_seq, + cparams.embeddings ? "true" : "false", + cparams.embeddings_nextn ? "true" : "false", + has_layer_embeddings ? "true" : "false")); + } +} + +void llama_model_deepseek41::validate_memory_accounting( + uint64_t state_bytes, + uint64_t graph_workspace_bytes) const { + if (!admission) { + throw std::runtime_error("DeepSeek V4.1 context has no host-memory admission result"); + } + admission->result = llama_dsv41_validate_runtime_memory( + admission->result, + state_bytes, + graph_workspace_bytes); + LLAMA_LOG_INFO("%s\n", admission->result.describe().c_str()); +} + namespace { struct dsv41_hc_mix { diff --git a/src/models/models.h b/src/models/models.h index 4acf4d90e1ad..8235fb26c914 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1323,7 +1323,9 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { }; struct engram_model; + struct admission_model; std::shared_ptr engram; + std::shared_ptr admission; std::shared_ptr experts; std::unique_ptr create_memory_engram_runtime(size_t max_tokens) const; @@ -1336,6 +1338,10 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { void release_runtime_work_after_sync(ggml_backend_sched_t sched) const override; 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; std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 98e9cfce4b7a..ae3856065a3e 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) @@ -274,6 +275,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-arg-parser.cpp b/tests/test-arg-parser.cpp index e0907631abd8..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" @@ -170,6 +171,73 @@ 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_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", "-np", "3", + }; + 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_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-admission.cpp b/tests/test-deepseek41-admission.cpp new file mode 100644 index 000000000000..effefdedeeb3 --- /dev/null +++ b/tests/test-deepseek41-admission.cpp @@ -0,0 +1,403 @@ +#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_ubatch = 32; + params.n_vocab = LLAMA_DSV41_N_VOCAB; + params.n_expert_used = LLAMA_DSV41_N_EXPERT_USED; + params.state_bytes = 2ULL << 30; + 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.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); + 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("routed expert union") != 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() { + 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; + + 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); + + 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() { + 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; + params.state_bytes = static_cast(n_ctx)*65536; + 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=", + "host_total=", "host_available=", "batch=", "outputs=", "outputs_per_seq=", + "expert_slots=", "required_expert_slots=", "expert_ubatch_capacity=", + "expert_cache=", "expert_staging=", "expert_replacement=", "direct_io_bounce=", + "output_bytes=", "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); +} + +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_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); + 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; + params.configured_cache_slots = LLAMA_DSV41_N_EXPERT_USED; + const auto admitted = llama_dsv41_admit(host_with_used(0), 0, published_tensors(), params); + + const auto measured = llama_dsv41_validate_runtime_memory( + admitted, admitted.state_bytes, admitted.graph_workspace_bytes - 1); + REQUIRE(measured.projected_bytes == admitted.projected_bytes - 1); + + const uint64_t over_soft = admitted.graph_workspace_bytes + + (admitted.soft_bytes - admitted.projected_bytes) + 1; + REQUIRE(thrown([&]() { + llama_dsv41_validate_runtime_memory(admitted, admitted.state_bytes, over_soft); + }).find("category=runtime_workspace") != std::string::npos); + + REQUIRE(thrown([&]() { + llama_dsv41_validate_runtime_memory(admitted, 0, admitted.graph_workspace_bytes); + }).find("category=runtime") != std::string::npos); +} + +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() { + test_procfs(); + test_procfs_fail_closed(); + test_published_slot_fit(); + test_configured_cache(); + test_threshold_boundaries(); + 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..f11f9436f0cb 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,136 @@ static void test_graph_construction() { ggml_free(ctx); } +struct reservation_test_model : 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; + } +}; + +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; + params.n_batch = 1; + params.n_ubatch = 1; + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + 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()); + 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 +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/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py new file mode 100644 index 000000000000..f9cba175a8d1 --- /dev/null +++ b/tests/test_strix_memory_watchdog.py @@ -0,0 +1,2415 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib.util +import fcntl +import hashlib +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 +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 or 0.0) + 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 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",), + 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, + group_alive=self.group_alive, + 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): + @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" + ) + + @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", + ) + + @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;" + "signal.signal(signal.SIGHUP,signal.SIG_IGN);" + "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.SIGHUP, + 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" + self._write_procfs_fixture(root) + 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", + "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 stderr_path.read_text( + encoding="utf-8" + ).splitlines() + ] + self.assertEqual( + records[-1]["classification"], "parent_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 ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + 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", + ) + + @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", + ) + def test_guardian_control_failure_still_kills_and_reaps_group( + self, + ) -> None: + class FailingFinalAudit: + 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 == self.fail_at: + raise OSError("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);" + "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"): + for artifact_failure in ( + "term_audit", + "kill_audit", + "final_audit", + "lease", + ): + with self.subTest( + mode=mode, + artifact_failure=artifact_failure, + ): + self._assert_guardian_control_failure_cleanup( + mode, + artifact_failure, + child_code, + FailingFinalAudit, + FailingFinalLease, + ) + + 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.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, + { + "term_audit": 1, + "kill_audit": 2, + "final_audit": 3, + }[artifact_failure], + ) + 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: + 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]["exit_code"], 7) + self.assertEqual( + records[-1]["threshold_reason"], + "guardian control failed during graceful cleanup", + ) + self.assertEqual( + records[-1]["secondary_errors"][0]["component"], + ( + "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"], + -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", + ) + 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, + 1.0, + 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, + 1.0, + 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() + + @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" + "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_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",), + grace_seconds=float("nan"), + ) + 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: + 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_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( + [ + 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", + 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 = [ + json.loads(line) for line in result.stderr.splitlines() + ] + 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()), + ) + + @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) + 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 + 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") + 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), + (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) + (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) / "exe").symlink_to( + Path(sys.executable).resolve() + ) + (process_root / str(watchdog_pid) / "cmdline").write_bytes( + cmdline + ) + + 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, + ) + 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, + "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_executable_path": str( + Path(sys.executable).resolve() + ), + "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, + "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, + "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), + "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 = { + "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": guardian_pid, + "sample": { + "audit_record_sha256": hashlib.sha256( + audit_line.encode("utf-8") + ).hexdigest() + }, + } + watchdog._write_json_atomic( + lease_path, lease, create=True + ) + watchdog._write_json_atomic( + heartbeat_path, heartbeat, create=True + ) + + 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 + ) + 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:], + ], + ), + ): + 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 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) + 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("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" + 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" + 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)) + + 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, + ) + + 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: + 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_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, + ) + + 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( + harness.signals, + [signal.SIGTERM, signal.SIGKILL], + ) + 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()) + + def fail_launch( + command: tuple[str, ...], **kwargs: Any + ) -> FakeProcess: + raise FileNotFoundError(2, "No such file or directory") + + setattr(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() 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); }