From 1e730639862033e1fcbee13e739e679e69ebcba7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 06:27:48 +0000 Subject: [PATCH] Studio: address follow-up review on the abort memo and NUMA footprint - Size the NUMA footprint KV estimate for the launched --parallel slots (it defaulted to n_parallel=1), so a multi-slot KV that overflows a node still triggers interleave. - Include the launch settings (context, spec, extra_args) in the scheduler-abort memo key, so an identical replay stays blocked but changing -c, disabling spec, or picking a different quant (the recovery the error message recommends) is allowed to retry. --- studio/backend/core/inference/llama_cpp.py | 21 +++++++++++++------ .../backend/tests/test_cpu_only_defaults.py | 8 ++++--- .../backend/tests/test_sched_reserve_abort.py | 21 ++++++++++++------- 3 files changed, 33 insertions(+), 17 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b37d151dac..7dcd1b3928 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4714,12 +4714,21 @@ class LlamaCppBackend: binary = self._find_llama_server_binary() # Fail fast BEFORE killing the live server (so a known-bad reload is - # non-destructive) if this (binary, model, variant) already aborted in the - # graph scheduler this session: reloading re-reads the weights into the same - # crash (a startup crash 500s and the UI replays /load). Keyed per variant so - # a failed quant does not block trying a different quant in the same repo. + # non-destructive) if this exact launch already aborted in the graph + # scheduler this session: reloading re-reads the weights into the same crash + # (a startup crash 500s and the UI replays /load). The key includes the + # variant AND the launch settings (context, spec) so an identical replay is + # blocked but a user changing quant / lowering -c / disabling spec -- the + # exact recovery the error message recommends -- is allowed to retry. _abort_memo_model = "\x00".join( - [model_identifier or "", hf_variant or "", gguf_path or ""] + [ + model_identifier or "", + hf_variant or "", + gguf_path or "", + str(n_ctx), + str(speculative_type or ""), + " ".join(str(a) for a in (extra_args or [])), + ] ) if LlamaCppBackend._sched_reserve_aborts(binary, _abort_memo_model): logger.warning( @@ -5990,7 +5999,7 @@ class LlamaCppBackend: if model_size and effective_ctx > 0 and self._can_estimate_kv(): try: _numa_footprint = model_size + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv + effective_ctx, cache_type_kv, n_parallel = n_parallel ) except Exception: _numa_footprint = model_size diff --git a/studio/backend/tests/test_cpu_only_defaults.py b/studio/backend/tests/test_cpu_only_defaults.py index 98dfd9930c..5847a97157 100644 --- a/studio/backend/tests/test_cpu_only_defaults.py +++ b/studio/backend/tests/test_cpu_only_defaults.py @@ -134,9 +134,11 @@ def test_cpu_context_floors_to_min_when_weights_exceed_budget(): def test_numa_decision_uses_footprint_not_just_weights(): - """The NUMA interleave decision must use weights + KV, so a model whose weights fit - one node but whose footprint does not still interleaves (PR review fix).""" + """The NUMA interleave decision must use weights + KV (at the launched parallel + slots), so a model whose weights fit one node but whose footprint does not still + interleaves (PR review fixes).""" src = _load_model_src() assert "_numa_footprint" in src - assert "_estimate_kv_cache_bytes(" in src assert "decide_interleave(_numa_footprint" in src + # KV must be sized for the launched --parallel slots, not the n_parallel=1 default. + assert "effective_ctx, cache_type_kv, n_parallel = n_parallel" in src diff --git a/studio/backend/tests/test_sched_reserve_abort.py b/studio/backend/tests/test_sched_reserve_abort.py index 916ff59679..26cb2bbb5f 100644 --- a/studio/backend/tests/test_sched_reserve_abort.py +++ b/studio/backend/tests/test_sched_reserve_abort.py @@ -243,18 +243,23 @@ def test_failfast_guard_runs_before_killing_the_live_server(): assert guard_line < kill_line -def test_abort_memo_key_includes_variant(): - """The memo key must include hf_variant / gguf_path so a failed quant does not - block a different quant of the same repo (PR review fix).""" +def test_abort_memo_key_includes_variant_and_launch_settings(): + """The memo key must include the variant AND the launch settings (context, spec), + so a failed quant does not block a different quant, and changing -c / spec (the + recommended recovery) is allowed to retry while an identical replay stays blocked.""" src = _load_model_src() assert "_abort_memo_model" in src - assert "hf_variant" in src and "gguf_path" in src + for tok in ("hf_variant", "gguf_path", "str(n_ctx)", "speculative_type", "extra_args"): + assert tok in src, tok - # Sanity: distinct variants produce distinct keys for the same model. - def key(model, variant, gguf): - return "\x00".join([model or "", variant or "", gguf or ""]) + def key(model="repo", variant="", gguf="", n_ctx=4096, spec="", extra=""): + return "\x00".join([model, variant, gguf, str(n_ctx), spec, extra]) - assert key("repo", "UD-Q6_K", None) != key("repo", "UD-Q4_K_XL", None) + base = key(variant="UD-Q6_K") + assert base != key(variant="UD-Q4_K_XL") # different quant -> retry allowed + assert base != key(variant="UD-Q6_K", n_ctx=2048) # lower context -> retry allowed + assert base != key(variant="UD-Q6_K", spec="off") # disable spec -> retry allowed + assert base == key(variant="UD-Q6_K") # identical replay -> still blocked def test_load_model_records_abort_on_crash():