From 60757367ab558447b5511a215ea2c07a758e975b Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Sun, 26 Jul 2026 05:51:58 -0700 Subject: [PATCH] Studio: fix sharded MTP picks and native fallback dedup Addresses the two review points that apply to the current revision. - collapse a split MTP/ drafter to its first shard before ranking, so size ordering cannot select a smaller trailing shard. llama-server takes shard 1 as the model path, matching _local_gguf_load_path - align the shard suffix pattern with _GGUF_SPLIT_FILE_RE - accept the MTP/ fallback during reload deduplication. A native load whose root drafter is out of bounds launches the subdir copy, so root-first detection never equalled the stored path and that layout restarted llama-server on every apply. A deleted drafter still forces a reload The two review points about the companion root during deduplication do not apply: for a quant named directory the load-time and dedup roots both resolve to the repository root. The divergence is limited to subdirectories whose names do not match the quant pattern, which predates this branch and is noted in the description. --- studio/backend/routes/inference.py | 14 ++++- .../tests/test_mtp_drafter_companion.py | 57 +++++++++++++++++++ .../tests/test_native_gguf_companion.py | 43 ++++++++++++++ studio/backend/utils/models/model_config.py | 8 ++- 4 files changed, 118 insertions(+), 4 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 370d76cef3..878a55bc92 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3356,7 +3356,19 @@ def _request_matches_loaded_settings( except OSError: return False if detected_resolved != stored_resolved: - return False + # A native load whose root drafter was out of bounds runs the + # MTP/ fallback instead, so root-first detection never equals + # what launched. Accept the subdir copy as current too, else + # that layout reloads on every apply. + fallback = detect_mtp_file( + llama_backend.gguf_path, search_root = companion_root, skip_root = True + ) + try: + fallback_resolved = Path(fallback).resolve() if fallback else None + except OSError: + return False + if stored_resolved is None or fallback_resolved != stored_resolved: + return False return True diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index a3ac2e5071..2ffeb62d85 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -632,3 +632,60 @@ def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): b._download_companion_gguf = _fake_companion assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None assert reached.get("hit") is True + + +def test_detect_mtp_file_returns_first_shard_of_split_subdir_drafter(tmp_path): + """llama-server takes shard 1 as the model path, so a split MTP/ copy must + not resolve to whichever shard happens to be smallest.""" + weight = tmp_path / "model-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + first = sub / "mtp-model-Q4_0-00001-of-00002.gguf" + first.write_bytes(b"x" * 4096) + (sub / "mtp-model-Q4_0-00002-of-00002.gguf").write_bytes(b"x") + + assert detect_mtp_file(str(weight)) == str(first.resolve()) + + +def test_detect_mtp_file_skip_root_ignores_root_drafter(tmp_path): + """skip_root is how a native load recovers when the root drafter is out + of bounds for its grant.""" + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"x") + (tmp_path / "mtp-model.gguf").write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + subdir_copy = sub / "mtp-model-Q4_0.gguf" + subdir_copy.write_bytes(b"x") + + assert detect_mtp_file(str(weight), str(tmp_path)) == str( + (tmp_path / "mtp-model.gguf").resolve() + ) + assert detect_mtp_file(str(weight), str(tmp_path), skip_root = True) == str(subdir_copy.resolve()) + + +def test_detect_mtp_file_rejects_weight_copy_inside_mtp_dir(tmp_path): + """Everything under MTP/ counts as a drafter for menu exclusion, but only + a published drafter name may be launched as --model-draft.""" + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "gemma-4-E4B-it-qat-Q4_0.gguf").write_bytes(b"x") + + assert detect_mtp_file(str(weight)) is None + + +def test_detect_mtp_file_pairs_k_quant_subdir_drafter(tmp_path): + """Pairing must use the full quant vocabulary, not just Q_/BF16/F16.""" + weight = tmp_path / "gemma-4-12b-it-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + drafter = sub / "mtp-gemma-4-12b-it-UD-Q4_K_XL.gguf" + drafter.write_bytes(b"x") + + assert detect_mtp_file(str(weight)) == str(drafter.resolve()) diff --git a/studio/backend/tests/test_native_gguf_companion.py b/studio/backend/tests/test_native_gguf_companion.py index 33e3c70dbb..105f7b5e1a 100644 --- a/studio/backend/tests/test_native_gguf_companion.py +++ b/studio/backend/tests/test_native_gguf_companion.py @@ -194,3 +194,46 @@ def test_native_companion_rejects_missing_weight(tmp_path): def test_native_companion_none_is_noop(): _validate_native_gguf_companion(None, None, "MTP drafter") + + +def test_reload_dedup_accepts_native_subdir_fallback(tmp_path, monkeypatch): + """A native load whose root drafter was out of bounds launches the MTP/ + copy, so root-first detection never matches it. Dedup must still hold.""" + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"model") + (tmp_path / "mtp-model.gguf").write_bytes(b"root drafter") + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + companion = companion_dir / "mtp-model-Q4_0.gguf" + companion.write_bytes(b"draft") + + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)) + backend = LlamaCppBackend() + backend._gguf_path = str(weight) + backend._mtp_draft_path = str(companion) + + request = LoadRequest(model_path = str(weight)) + assert _request_matches_loaded_settings(request, backend) + + +def test_reload_dedup_still_reloads_when_drafter_disappears(tmp_path, monkeypatch): + """The fallback comparison must not mask a deleted drafter.""" + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"model") + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + companion = companion_dir / "mtp-model-Q4_0.gguf" + companion.write_bytes(b"draft") + + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)) + backend = LlamaCppBackend() + backend._gguf_path = str(weight) + backend._mtp_draft_path = str(companion) + + companion.unlink() + request = LoadRequest(model_path = str(weight)) + assert not _request_matches_loaded_settings(request, backend) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 566abe9d3b..1a6036765a 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1472,7 +1472,7 @@ def detect_mtp_file( stem = stem[len("mtp-") :] # Shard suffix sits outside the quant token, so strip it first or the # anchored strip below cannot match. - stem = re.sub(r"-[0-9]+-of-[0-9]+$", "", stem) + stem = re.sub(r"-[0-9]{5}-of-[0-9]{5}$", "", stem) if stem.endswith("-mtp"): stem = stem[: -len("-mtp")] # Full quant vocabulary, not a subset: K/IQ/UD/MXFP drafters pair too. @@ -1565,11 +1565,13 @@ def detect_mtp_file( continue try: if f.is_file(): - subdir_candidates.append(f) + # llama-server takes shard 1 as the model path, so + # collapse a split copy to it before ranking. + subdir_candidates.append(_local_gguf_load_path(f)) except OSError: continue - for candidate in sorted(subdir_candidates, key = _smallest_first): + for candidate in sorted(dict.fromkeys(subdir_candidates), key = _smallest_first): try: resolved = candidate.resolve() except OSError: