From e9784fd2123fc890562eaee02b1affbd3780724b Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Sun, 26 Jul 2026 18:34:57 -0700 Subject: [PATCH] Studio: fix shard handling on the root branch and native dedup comparison All four review points on the previous revision were correct. - share one launch-path helper between the root and MTP/ branches. The root branch still resolved to the blob, so a sharded snapshot drafter lost its sibling shard names. It began matching once _pairing_stem learned shard suffixes, so the two branches had drifted - strip the shard suffix before the -mtp name check. An old-scheme split copy is -Q8_0-MTP-00001-of-00002.gguf, whose stem does not end in -mtp, so every shard was discarded - mirror the load path's admissibility choice during reload dedup for native loads instead of special casing a stored None. With a root drafter outside the grant and no MTP/ copy the load stores no drafter, and dedup compared that None against the root file and reloaded every time - drop activeNativePathToken from the local GGUF classification. Status reconciliation keeps the token across a switch to a remote GGUF because no replacement token exists, so a stale token labelled that remote model local and showed placement guidance instead of the download recovery text. activeModelIsLocal is the backend's own classification and already covers native picks The model picker contract now asserts the token is absent rather than present, since it pinned the stale-token behaviour. --- studio/backend/routes/inference.py | 58 +++++++++++++------ .../tests/test_mtp_drafter_companion.py | 40 +++++++++++++ .../tests/test_native_gguf_companion.py | 20 +++++++ studio/backend/utils/models/model_config.py | 28 +++++---- .../src/features/chat/chat-settings-sheet.tsx | 8 +-- tests/studio/test_model_picker_contracts.py | 4 +- 6 files changed, 126 insertions(+), 32 deletions(-) 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]