diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2b98d06aff..1944bc9dd4 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3138,6 +3138,27 @@ def _validate_native_gguf_companion( ) from exc +def _native_gguf_companion_usable( + companion_path: str | None, + gguf_path: str | None, + *, + mtp_search_root: str | Path | None = None, +) -> bool: + """Whether a native load would accept this MTP drafter. Same rules as + _validate_native_gguf_companion, as a predicate for reload dedup.""" + try: + _validate_native_gguf_companion( + companion_path, + gguf_path, + "MTP drafter", + allow_mtp_subdir = True, + mtp_search_root = mtp_search_root, + ) + except HTTPException: + return False + return True + + def _normalise_settings_str(value: Optional[str]) -> Optional[str]: """Lowercase + strip a settings string, mapping blank/None to None.""" if value is None: @@ -3359,6 +3380,25 @@ def _request_matches_loaded_settings( llama_backend.gguf_path, llama_backend.gguf_path ) detected = detect_mtp_file(llama_backend.gguf_path, search_root = companion_root) + if native_grant_backed: + # Mirror the load path's choice, or the comparison is against a + # drafter that never launched. A native grant cannot reach a + # root drafter outside it, so the load falls back to the MTP/ + # copy and, failing that, to no drafter at all. An ordinary + # load reaches the root drafter, so it keeps root-first + # detection and reloads when one appears. + if detected and not _native_gguf_companion_usable( + detected, llama_backend.gguf_path, mtp_search_root = companion_root + ): + detected = detect_mtp_file( + llama_backend.gguf_path, + search_root = companion_root, + skip_root = True, + ) + if detected and not _native_gguf_companion_usable( + detected, llama_backend.gguf_path, mtp_search_root = companion_root + ): + detected = None stored = llama_backend.mtp_draft_path try: detected_resolved = Path(detected).resolve() if detected else None @@ -3366,23 +3406,7 @@ def _request_matches_loaded_settings( except OSError: return False if detected_resolved != stored_resolved: - # 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. Native only: an ordinary - # load can reach the root drafter, so a newly added one must - # still reload. - if not native_grant_backed: - return False - 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 False return True diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 15705753b4..98a233ac2b 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -717,3 +717,43 @@ def test_detect_mtp_file_keeps_snapshot_path_for_sharded_subdir_drafter(tmp_path found = detect_mtp_file(str(weight), str(snapshot)) assert found == str(first) assert (Path(found).parent / second.name).exists() + + +def test_detect_mtp_file_pairs_sharded_old_scheme_subdir_drafter(tmp_path): + """An old-scheme split copy is -Q8_0-MTP-00001-of-00002.gguf, whose + stem does not end in -mtp until the shard suffix comes off.""" + weight = tmp_path / "model-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + first = sub / "model-Q8_0-MTP-00001-of-00002.gguf" + first.write_bytes(b"x" * 4096) + (sub / "model-Q8_0-MTP-00002-of-00002.gguf").write_bytes(b"x") + + assert detect_mtp_file(str(weight)) == str(first) + + +def test_detect_mtp_file_keeps_snapshot_path_for_sharded_root_drafter(tmp_path): + """The root branch needs the same shard handling as the MTP/ branch.""" + blobs = tmp_path / "blobs" + snapshot = tmp_path / "snapshots" / "abc" + blobs.mkdir(parents = True) + snapshot.mkdir(parents = True) + + (blobs / "sha_weight").write_bytes(b"w") + weight = snapshot / "model-Q4_0.gguf" + try: + weight.symlink_to(blobs / "sha_weight") + except OSError: + pytest.skip("symlinks unavailable") + + first = snapshot / "mtp-model-Q4_0-00001-of-00002.gguf" + second = snapshot / "mtp-model-Q4_0-00002-of-00002.gguf" + (blobs / "sha_1").write_bytes(b"d" * 4096) + (blobs / "sha_2").write_bytes(b"d") + first.symlink_to(blobs / "sha_1") + second.symlink_to(blobs / "sha_2") + + found = detect_mtp_file(str(weight), str(snapshot)) + assert found == str(first) + assert (Path(found).parent / second.name).exists() diff --git a/studio/backend/tests/test_native_gguf_companion.py b/studio/backend/tests/test_native_gguf_companion.py index 24f0553580..ab1ba4f7fe 100644 --- a/studio/backend/tests/test_native_gguf_companion.py +++ b/studio/backend/tests/test_native_gguf_companion.py @@ -266,3 +266,23 @@ def test_reload_dedup_reloads_for_ordinary_load_when_root_drafter_appears(tmp_pa assert _request_matches_loaded_settings(request, backend, None, native_grant_backed = True) # An ordinary load would pick the root drafter, so it must reload. assert not _request_matches_loaded_settings(request, backend, None, native_grant_backed = False) + + +def test_reload_dedup_native_load_with_no_admissible_drafter(tmp_path, monkeypatch): + """Root drafter out of the grant and no MTP/ copy: the load stores no + drafter, so dedup must compare against None rather than the root file.""" + 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") + + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)) + backend = LlamaCppBackend() + backend._gguf_path = str(weight) + backend._mtp_draft_path = None + + request = LoadRequest(model_path = str(weight)) + assert _request_matches_loaded_settings(request, backend, None, native_grant_backed = True) + # An ordinary load would launch the root drafter, so it must reload. + assert not _request_matches_loaded_settings(request, backend, None, native_grant_backed = False) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 8f1384df27..6e49294e3f 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1478,6 +1478,15 @@ def detect_mtp_file( # Full quant vocabulary, not a subset: K/IQ/UD/MXFP drafters pair too. return re.sub(rf"-(?:{_GGUF_KNOWN_QUANT_RE.pattern})$", "", stem, flags = re.IGNORECASE) + def _drafter_launch_path(candidate: Path) -> str: + # llama-server takes shard 1 as the model path, and a split copy must + # stay on its snapshot path: the blob target has no sibling shard + # names. Single-file drafters still resolve, as callers expect. + loadable = _local_gguf_load_path(candidate) + if _GGUF_SPLIT_FILE_RE.match(loadable.name): + return str(loadable) + return str(loadable.resolve()) + def _matches_weight(candidate: Path) -> bool: if weight_name is None: return True @@ -1527,7 +1536,7 @@ def detect_mtp_file( continue try: if f.is_file(): - return str(f.resolve()) + return _drafter_launch_path(f) except OSError: continue @@ -1559,29 +1568,28 @@ def detect_mtp_file( lower = f.name.lower() if not lower.endswith(".gguf"): continue - if not (lower.startswith("mtp-") or Path(lower).stem.endswith("-mtp")): + # Drop the shard suffix first: an old-scheme split copy is + # named -Q8_0-MTP-00001-of-00002.gguf, whose stem does + # not end in -mtp. + stem = re.sub(r"-[0-9]{5}-of-[0-9]{5}$", "", Path(lower).stem) + if not (lower.startswith("mtp-") or stem.endswith("-mtp")): continue if not _matches_weight(f): continue try: if f.is_file(): - # llama-server takes shard 1 as the model path, so - # collapse a split copy to it before ranking. + # Collapse a split copy to shard 1 before ranking. subdir_candidates.append(_local_gguf_load_path(f)) except OSError: continue for candidate in sorted(dict.fromkeys(subdir_candidates), key = _smallest_first): try: - # A split copy keeps its snapshot path: resolving to the blob - # drops the sibling shard names llama-server needs to find. - resolved = ( - candidate if _GGUF_SPLIT_FILE_RE.match(candidate.name) else candidate.resolve() - ) + resolved = _drafter_launch_path(candidate) except OSError: continue logger.info(f"Detected MTP subdirectory drafter: {resolved}") - return str(resolved) + return resolved return None diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index c9d59d1b14..d0a525ca18 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -381,9 +381,6 @@ export function ChatSettingsPanel({ const isMobile = useIsMobile(); const isLoadedGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; const currentCheckpoint = params.checkpoint; - const activeNativePathToken = useChatRuntimeStore( - (s) => s.activeNativePathToken, - ); const activeModelIsLocal = useChatRuntimeStore( (s) => s.activeModelIsLocal, ); @@ -396,10 +393,13 @@ export function ChatSettingsPanel({ isLoadedGguf || ggufContextLength != null || (currentCheckpoint?.toLowerCase().endsWith(".gguf") ?? false); + // activeModelIsLocal is the backend's own classification and covers native + // picks. activeNativePathToken must not be used here: status reconciliation + // keeps it across a switch to a remote GGUF (no replacement token exists), + // so a stale token would label that remote model local. const isLocalGguf = isGguf && (activeModelIsLocal || - activeNativePathToken != null || isLocalModelPath(currentCheckpoint ?? "") || (currentCheckpoint?.toLowerCase().endsWith(".gguf") ?? false)); const ggufMaxContextLength = useChatRuntimeStore( diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 50464d31b9..c547f4927b 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -349,9 +349,11 @@ def test_local_mtp_warning_covers_path_and_native_gguf_sources(): local = re.search(r"const isLocalGguf =.*?;", src, re.S) assert local assert "isGguf &&" in local.group(0) - assert "activeNativePathToken" in local.group(0) assert "activeModelIsLocal" in local.group(0) assert "isLocalModelPath" in local.group(0) + # A native token outlives a switch to a remote GGUF, so it must not + # classify the model here; activeModelIsLocal already covers native picks. + assert "activeNativePathToken" not in local.group(0) assert "isLocalGguf" in src.split('specFallbackReason === "drafter_not_found"', 1)[1]