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: