diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index add3228a28..99d4f31c5d 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -394,6 +394,9 @@ class LoadResponse(BaseModel): is_vision: bool = Field(False, description = "Whether model is a vision model") is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)") + is_local_model: bool = Field( + False, description = "Whether the loaded model came from a local filesystem path" + ) is_diffusion: bool = Field( False, description = "Whether model is a block-diffusion model (DiffusionGemma)" ) @@ -558,6 +561,9 @@ class InferenceStatusResponse(BaseModel): ) is_vision: bool = Field(False, description = "Whether the active model is a vision model") is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)") + is_local_model: bool = Field( + False, description = "Whether the active model came from a local filesystem path" + ) is_diffusion: bool = Field( False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)" ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1e6959472c..9f193cf02b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1019,8 +1019,11 @@ try: ) from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig + from utils.paths import is_local_path from utils.inference import load_inference_config from utils.models.model_config import ( + _local_gguf_companion_search_root, + colocated_split_shards, detect_mtp_file, load_model_defaults, ) @@ -1028,6 +1031,7 @@ try: NativePathLeaseError, display_label_for_native_path, is_registered_native_path_label, + native_gguf_companion_parent_allowed, redact_native_paths, verify_native_path_lease, ) @@ -1057,8 +1061,11 @@ except ImportError: ) from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig + from utils.paths import is_local_path from utils.inference import load_inference_config from utils.models.model_config import ( + _local_gguf_companion_search_root, + colocated_split_shards, detect_mtp_file, load_model_defaults, ) @@ -1066,6 +1073,7 @@ except ImportError: NativePathLeaseError, display_label_for_native_path, is_registered_native_path_label, + native_gguf_companion_parent_allowed, redact_native_paths, verify_native_path_lease, ) @@ -3140,11 +3148,16 @@ def _monitor_active_model() -> Optional[str]: def _validate_native_gguf_companion( - companion_path: str | None, gguf_path: str | None, label: str + companion_path: str | None, + gguf_path: str | None, + label: str, + *, + allow_mtp_subdir: bool = False, + mtp_search_root: str | Path | None = None, ) -> None: """Reject a companion GGUF (mmproj / MTP drafter) that a native-lease load would otherwise hand to llama-server: must be a regular file (no symlink - escaping the leased directory) living next to the selected GGUF.""" + escaping the leased directory) in a permitted location.""" if not companion_path or not gguf_path: return import stat as _stat_module @@ -3166,10 +3179,20 @@ def _validate_native_gguf_companion( detail = f"Native {label} must be a regular file.", ) try: - if companion.resolve(strict = True).parent != gguf.resolve(strict = True).parent: + if not native_gguf_companion_parent_allowed( + companion, + gguf, + allow_mtp_subdir = allow_mtp_subdir, + mtp_search_root = mtp_search_root, + ): + location = ( + "beside the selected GGUF or in its MTP directory" + if allow_mtp_subdir + else "next to the selected GGUF" + ) raise HTTPException( status_code = 400, - detail = f"Native {label} must live next to the selected GGUF.", + detail = f"Native {label} must live {location}.", ) except OSError as exc: raise HTTPException( @@ -3178,6 +3201,61 @@ def _validate_native_gguf_companion( ) from exc +def _loaded_is_local_model( + llama_backend: LlamaCppBackend, native_grant_backed: bool, model_id: str | None +) -> bool: + """Provenance of the running model, preferring what the load recorded. + + Falls back to the filesystem for a server started before the flag existed. + """ + if native_grant_backed: + return True + stored = getattr(llama_backend, "_is_local_model", None) + if stored is not None: + return bool(stored) + return bool(model_id and is_local_path(model_id)) + + +def _validate_native_mtp_drafter( + companion_path: str | None, + gguf_path: str | None, + *, + mtp_search_root: str | Path | None = None, +) -> None: + """Validate an MTP drafter for a native load, every shard of it. + + llama-server opens the sibling shards of a split drafter implicitly, so + checking only the launch path would let a later shard be a symlink out of + the permitted directory without ever facing the native rules. + """ + if not companion_path or not gguf_path: + return + shards, _ = colocated_split_shards(Path(companion_path)) + for shard in shards or [Path(companion_path)]: + _validate_native_gguf_companion( + str(shard), + gguf_path, + "MTP drafter", + allow_mtp_subdir = True, + mtp_search_root = mtp_search_root, + ) + + +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, as a predicate for + reload dedup. Same rules, so the two cannot disagree.""" + try: + _validate_native_mtp_drafter(companion_path, gguf_path, 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: @@ -3255,6 +3333,7 @@ def _request_matches_loaded_settings( request: LoadRequest, llama_backend: LlamaCppBackend, effective_chat_template_override: Optional[str] = None, + native_grant_backed: bool = False, ) -> bool: """True iff every runtime setting on the request matches the loaded server. Caller has already checked model+variant+is_loaded. See #5401. @@ -3394,7 +3473,29 @@ def _request_matches_loaded_settings( else llama_backend.extra_args ) if not _extra_args_set_spec_type(effective_extras): - detected = detect_mtp_file(llama_backend.gguf_path) + companion_root = _local_gguf_companion_search_root( + 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. + def _usable(candidate: str) -> bool: + return _native_gguf_companion_usable( + candidate, llama_backend.gguf_path, mtp_search_root = companion_root + ) + + if detected and not _usable(detected): + detected = detect_mtp_file( + llama_backend.gguf_path, + search_root = companion_root, + skip_root = True, + accept = _usable, + ) stored = llama_backend.mtp_draft_path try: detected_resolved = Path(detected).resolve() if detected else None @@ -4941,6 +5042,7 @@ async def _load_model_impl( request, llama_backend, effective_chat_template_override, + native_grant_backed = native_grant_backed, ) # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) @@ -4967,6 +5069,9 @@ async def _load_model_impl( is_vision = llama_backend._is_vision, is_lora = False, is_gguf = True, + is_local_model = _loaded_is_local_model( + llama_backend, native_grant_backed, llama_backend.model_identifier + ), is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, @@ -5026,6 +5131,7 @@ async def _load_model_impl( is_vision = _model_info.get("is_vision", False), is_lora = _model_info.get("is_lora", False), is_gguf = False, + is_local_model = native_grant_backed or is_local_path(backend.active_model_name), is_audio = _model_info.get("is_audio", False), audio_type = _model_info.get("audio_type"), has_audio_input = _model_info.get("has_audio_input", False), @@ -5212,13 +5318,41 @@ async def _load_model_impl( if config.gguf_mtp_file: # The drafter is optional (unlike mmproj for a vision # model): drop it rather than fail the load. - try: - _validate_native_gguf_companion( - config.gguf_mtp_file, config.gguf_file, "MTP drafter" + mtp_search_root = _local_gguf_companion_search_root( + config.gguf_file, config.gguf_file + ) + + def _mtp_allowed(candidate: str) -> bool: + try: + _validate_native_mtp_drafter( + candidate, + config.gguf_file, + mtp_search_root = mtp_search_root, + ) + return True + except HTTPException as exc: + logger.warning( + "Dropping MTP drafter for native load: %s", exc.detail + ) + return False + + if not _mtp_allowed(config.gguf_mtp_file): + # The preferred root drafter is out of bounds for a + # grant on a quant subdir, but an MTP/ copy may not + # be. Scan them in preference order rather than + # dropping MTP on the first rejection. + fallback = detect_mtp_file( + config.gguf_file, + search_root = mtp_search_root, + skip_root = True, + accept = _mtp_allowed, ) - except HTTPException as exc: - logger.warning("Dropping MTP drafter for native load: %s", exc.detail) - config.gguf_mtp_file = None + if fallback: + logger.info( + "Using MTP subdirectory drafter for native load: %s", + fallback, + ) + config.gguf_mtp_file = fallback _source_load_kwargs = dict( gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, @@ -5332,6 +5466,10 @@ async def _load_model_impl( _gguf_is_audio = llama_backend._is_audio llama_backend._native_display_label = model_log_label if native_grant_backed else None llama_backend._native_grant_backed = bool(native_grant_backed) + # Provenance is a load-time fact. Re-deriving it per status poll + # would flip a local model to remote if its directory is deleted + # or unmounted underneath a still-running server. + llama_backend._is_local_model = bool(native_grant_backed or config.is_local) if _gguf_is_audio: logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}") @@ -5344,6 +5482,7 @@ async def _load_model_impl( is_vision = llama_backend.is_vision, is_lora = False, is_gguf = True, + is_local_model = config.is_local, is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, @@ -5487,6 +5626,7 @@ async def _load_model_impl( is_vision = config.is_vision, is_lora = config.is_lora, is_gguf = False, + is_local_model = config.is_local, is_audio = config.is_audio, audio_type = config.audio_type, has_audio_input = config.has_audio_input, @@ -6485,6 +6625,9 @@ async def get_status(current_subject: str = Depends(get_current_subject)): model_identifier = None if _native_grant_backed else _model_id, is_vision = llama_backend.is_vision, is_gguf = True, + is_local_model = _loaded_is_local_model( + llama_backend, _native_grant_backed, _model_id + ), is_diffusion = llama_backend.is_diffusion, gguf_variant = llama_backend.hf_variant, is_audio = getattr(llama_backend, "_is_audio", False), @@ -6556,6 +6699,9 @@ async def get_status(current_subject: str = Depends(get_current_subject)): model_identifier = backend.active_model_name, is_vision = is_vision, is_gguf = False, + is_local_model = bool( + backend.active_model_name and is_local_path(backend.active_model_name) + ), is_audio = is_audio, audio_type = audio_type, has_audio_input = has_audio_input, diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 02230632b6..95179b86bd 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -29,11 +29,14 @@ from hub.utils.gguf_plan import ( preferred_mtp_sibling, ) from utils.models.model_config import ( + ModelConfig, _is_mtp_drafter, + _local_gguf_companion_search_root, detect_gguf_model, detect_mtp_file, extract_model_size_b, ) +from utils.native_path_leases import native_gguf_companion_parent_allowed # ── Predicate + layering mirrors ───────────────────────────────────── @@ -211,6 +214,152 @@ def test_detect_mtp_file_search_root(tmp_path): assert found is not None and found.endswith("mtp-gemma-4-12b-it.gguf") +def test_quant_directory_selection_finds_repo_root_mtp(tmp_path): + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + mtp_dir = tmp_path / "MTP" + mtp_dir.mkdir() + drafter = mtp_dir / "mtp-gemma-4-E4B-it-Q4_0.gguf" + drafter.write_bytes(b"x") + + search_root = _local_gguf_companion_search_root(str(quant_dir), str(weight)) + assert Path(search_root).resolve() == tmp_path.resolve() + config = ModelConfig.from_identifier(str(quant_dir)) + assert config.is_local + assert config.gguf_file == str(weight.resolve()) + assert config.gguf_mtp_file == str(drafter.resolve()) + + +def test_bare_relative_gguf_directory_is_local_source(tmp_path, monkeypatch): + model_dir = tmp_path / "outputs" / "gemma" + model_dir.mkdir(parents = True) + weight = model_dir / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + monkeypatch.chdir(tmp_path) + + config = ModelConfig.from_identifier("outputs/gemma") + assert config.is_local + assert config.gguf_file == str(weight.resolve()) + + +def test_detect_mtp_file_falls_back_to_new_scheme_subdir(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "mtp-gemma-4-E4B-it-BF16.gguf").write_bytes(b"x") + q4 = sub / "mtp-gemma-4-E4B-it-Q4_0.gguf" + q4.write_bytes(b"x") + + found = detect_mtp_file(str(weight)) + assert found == str(q4.resolve()) + + +def test_detect_mtp_file_falls_back_to_old_scheme_subdir(tmp_path): + weight = tmp_path / "gemma-4-12b-it-Q4_K_M.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + drafter = sub / "gemma-4-12b-it-Q8_0-MTP.gguf" + drafter.write_bytes(b"x") + + found = detect_mtp_file(str(weight)) + assert found == str(drafter.resolve()) + + +def test_detect_mtp_file_root_still_wins_over_subdir(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + root = tmp_path / "mtp-gemma-4-E4B-it.gguf" + root.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "mtp-gemma-4-E4B-it-Q4_0.gguf").write_bytes(b"x") + + assert detect_mtp_file(str(weight)) == str(root.resolve()) + + +def test_detect_mtp_file_subdir_skips_foreign_drafter(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "mtp-gemma-4-12b-it-Q4_0.gguf").write_bytes(b"x") + + assert detect_mtp_file(str(weight)) is None + + +@pytest.mark.parametrize( + "companion_path", + ["mtp-gemma-4-E4B-it-Q4_0.gguf", "MTP/mtp-gemma-4-E4B-it-Q4_0.gguf"], +) +def test_detect_mtp_file_requires_model_name_boundary(tmp_path, companion_path): + weight = tmp_path / "gemma-4-E4B-item-qat-Q4_0.gguf" + weight.write_bytes(b"x") + companion = tmp_path / companion_path + companion.parent.mkdir(parents = True, exist_ok = True) + companion.write_bytes(b"x") + + assert detect_mtp_file(str(weight)) is None + + +def test_detect_mtp_file_accepts_case_variant_subdir(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "mtp" + sub.mkdir() + drafter = sub / "mtp-gemma-4-E4B-it-Q4_0.gguf" + drafter.write_bytes(b"x") + + assert detect_mtp_file(str(weight)) == str(drafter.resolve()) + + +def test_native_companion_parent_accepts_root_and_mtp_subdir(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + root_drafter = tmp_path / "mtp-gemma-4-E4B-it.gguf" + root_drafter.write_bytes(b"x") + sub = tmp_path / "MtP" + sub.mkdir() + nested_drafter = sub / "mtp-gemma-4-E4B-it-Q4_0.gguf" + nested_drafter.write_bytes(b"x") + + assert native_gguf_companion_parent_allowed(root_drafter, weight) + assert native_gguf_companion_parent_allowed(nested_drafter, weight, allow_mtp_subdir = True) + + +def test_native_companion_parent_rejects_other_nested_directory(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "other" + sub.mkdir() + drafter = sub / "mtp-gemma-4-E4B-it-Q4_0.gguf" + drafter.write_bytes(b"x") + + assert not native_gguf_companion_parent_allowed(drafter, weight) + + +def test_native_companion_parent_rejects_mtp_symlink_escape(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + weight = model_dir / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + outside = tmp_path / "outside" + outside.mkdir() + drafter = outside / "mtp-gemma-4-E4B-it-Q4_0.gguf" + drafter.write_bytes(b"x") + try: + (model_dir / "MTP").symlink_to(outside, target_is_directory = True) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + assert not native_gguf_companion_parent_allowed( + model_dir / "MTP" / drafter.name, weight, allow_mtp_subdir = True + ) + + # ── Reload dedup includes the drafter ──────────────────────────────── @@ -483,3 +632,198 @@ 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()) + + +def test_detect_mtp_file_keeps_snapshot_path_for_sharded_subdir_drafter(tmp_path): + """A split copy stored as HF snapshot symlinks must launch from the + snapshot path: the blob target has no sibling shard names.""" + blobs = tmp_path / "blobs" + snapshot = tmp_path / "snapshots" / "abc" + sub = snapshot / "MTP" + blobs.mkdir(parents = True) + sub.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 = sub / "mtp-model-Q4_0-00001-of-00002.gguf" + second = sub / "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() + + +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() + + +def test_detect_mtp_file_pairs_bpw_qualified_subdir_drafter(tmp_path): + """_extract_quant_label supports bpw-qualified names, so pairing must too.""" + weight = tmp_path / "model-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + drafter = sub / "mtp-model-IQ4_XS-3.53bpw.gguf" + drafter.write_bytes(b"x") + + assert detect_mtp_file(str(weight)) == str(drafter.resolve()) + + +def test_detect_mtp_file_skips_incomplete_split_drafter(tmp_path): + """An incomplete shard set fails llama-server's draft startup, so a + complete copy must win rather than MTP being disabled.""" + weight = tmp_path / "model-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + # Declares two shards but ships only the first. + (sub / "mtp-model-Q4_0-00001-of-00002.gguf").write_bytes(b"x" * 50) + complete = sub / "mtp-model-BF16.gguf" + complete.write_bytes(b"x" * 100) + + assert detect_mtp_file(str(weight)) == str(complete.resolve()) + + +def test_detect_mtp_file_ranks_split_drafter_by_total_size(tmp_path): + """Candidates collapse to shard 1, so a split copy must be summed or it + outranks a smaller single file.""" + weight = tmp_path / "model-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "mtp-model-Q8_0-00001-of-00002.gguf").write_bytes(b"x" * 90) + (sub / "mtp-model-Q8_0-00002-of-00002.gguf").write_bytes(b"x" * 90) + smaller = sub / "mtp-model-BF16.gguf" + smaller.write_bytes(b"x" * 100) + + assert detect_mtp_file(str(weight)) == str(smaller.resolve()) + + +def test_companion_search_root_promotes_bpw_quant_directory(tmp_path): + """A bpw-qualified quant directory must resolve to the repository root, or + the repo-root MTP/ copy is never in scope for it.""" + quant_dir = tmp_path / "IQ4_XS-3.53bpw" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + drafter = sub / "mtp-model.gguf" + drafter.write_bytes(b"x") + + # Directory selection and the file inside it agree on the root. + assert _local_gguf_companion_search_root(str(quant_dir), str(weight)) == str(tmp_path) + assert _local_gguf_companion_search_root(str(weight), str(weight)) == str(tmp_path) + assert detect_mtp_file(str(weight), str(tmp_path)) == str(drafter.resolve()) + + +def test_companion_search_root_keeps_non_quant_directories(tmp_path): + """Sharing the quant vocabulary must not widen what gets promoted.""" + for name in ("DeepSeek-V3-UD-Q2_K_XL", "outputs", "Q4_0-extra", "Q4_0bpw"): + directory = tmp_path / name + directory.mkdir() + weight = directory / "model.gguf" + weight.write_bytes(b"x") + assert _local_gguf_companion_search_root(str(directory), str(weight)) == str(directory) diff --git a/studio/backend/tests/test_native_gguf_companion.py b/studio/backend/tests/test_native_gguf_companion.py new file mode 100644 index 0000000000..3a7b4b0caf --- /dev/null +++ b/studio/backend/tests/test_native_gguf_companion.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Native GGUF companion path validation.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from fastapi import HTTPException + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from routes.inference import _validate_native_gguf_companion +from routes.inference import _request_matches_loaded_settings +from routes.inference import _validate_native_mtp_drafter +from routes.inference import _loaded_is_local_model +from routes.inference import _native_gguf_companion_usable +from utils.models.model_config import detect_mtp_file +from core.inference.llama_cpp import LlamaCppBackend +from models.inference import LoadRequest + + +def _write_pair(tmp_path: Path, folder: str | None = None) -> tuple[Path, Path]: + tmp_path.mkdir(parents = True, exist_ok = True) + weight = tmp_path / "model.gguf" + weight.write_bytes(b"model") + parent = tmp_path if folder is None else tmp_path / folder + parent.mkdir(parents = True, exist_ok = True) + companion = parent / "mtp-model.gguf" + companion.write_bytes(b"draft") + return weight, companion + + +def test_native_companion_allows_model_directory(tmp_path): + weight, companion = _write_pair(tmp_path) + _validate_native_gguf_companion(str(companion), str(weight), "vision companion") + + +@pytest.mark.parametrize("folder", ["MTP", "mtp", "MtP"]) +def test_native_mtp_companion_allows_mtp_directory(tmp_path, folder): + weight, companion = _write_pair(tmp_path, folder) + _validate_native_gguf_companion( + str(companion), str(weight), "MTP drafter", allow_mtp_subdir = True + ) + + +def test_native_mtp_companion_allows_repo_root_mtp_directory(tmp_path): + quant_dir = tmp_path / "Q4_0" + weight, _ = _write_pair(quant_dir) + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + companion = companion_dir / "mtp-model.gguf" + companion.write_bytes(b"draft") + + _validate_native_gguf_companion( + str(companion), + str(weight), + "MTP drafter", + allow_mtp_subdir = True, + mtp_search_root = str(tmp_path), + ) + + +def test_native_mtp_companion_rejects_unrelated_search_root(tmp_path): + quant_dir = tmp_path / "repo" / "Q4_0" + weight, _ = _write_pair(quant_dir) + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + companion = companion_dir / "mtp-model.gguf" + companion.write_bytes(b"draft") + + with pytest.raises(HTTPException, match = "must live beside"): + _validate_native_gguf_companion( + str(companion), + str(weight), + "MTP drafter", + allow_mtp_subdir = True, + mtp_search_root = str(tmp_path), + ) + + +def test_reload_dedup_finds_repo_root_mtp_companion(tmp_path, monkeypatch): + 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.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_matches_quant_directory_selection(tmp_path, monkeypatch): + 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.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(quant_dir), gguf_variant = "Q4_0") + assert _request_matches_loaded_settings(request, backend) + + +def test_native_vision_companion_rejects_mtp_directory(tmp_path): + weight, companion = _write_pair(tmp_path, "MTP") + with pytest.raises(HTTPException, match = "must live next to"): + _validate_native_gguf_companion(str(companion), str(weight), "vision companion") + + +@pytest.mark.parametrize("folder", ["other", "MTP/deeper", "mtp/deeper"]) +def test_native_companion_rejects_arbitrary_nesting(tmp_path, folder): + weight, companion = _write_pair(tmp_path, folder) + with pytest.raises(HTTPException, match = "must live beside") as error: + _validate_native_gguf_companion( + str(companion), str(weight), "MTP drafter", allow_mtp_subdir = True + ) + assert error.value.status_code == 400 + + +def test_native_companion_rejects_file_symlink(tmp_path): + weight, companion = _write_pair(tmp_path) + link = tmp_path / "mtp-link.gguf" + try: + link.symlink_to(companion) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + with pytest.raises(HTTPException, match = "regular file"): + _validate_native_gguf_companion(str(link), str(weight), "MTP drafter") + + +def test_native_companion_rejects_directory_symlink_escape(tmp_path): + model_dir = tmp_path / "model" + outside = tmp_path / "outside" + model_dir.mkdir() + outside.mkdir() + weight = model_dir / "model.gguf" + weight.write_bytes(b"model") + companion = outside / "mtp-model.gguf" + companion.write_bytes(b"draft") + try: + (model_dir / "MTP").symlink_to(outside, target_is_directory = True) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + with pytest.raises(HTTPException, match = "must live beside"): + _validate_native_gguf_companion( + str(model_dir / "MTP" / companion.name), + str(weight), + "MTP drafter", + allow_mtp_subdir = True, + ) + + +def test_native_companion_rejects_missing_file(tmp_path): + weight = tmp_path / "model.gguf" + weight.write_bytes(b"model") + with pytest.raises(HTTPException, match = "no longer accessible"): + _validate_native_gguf_companion(str(tmp_path / "missing.gguf"), str(weight), "MTP drafter") + + +def test_native_companion_rejects_directory(tmp_path): + weight = tmp_path / "model.gguf" + weight.write_bytes(b"model") + companion = tmp_path / "mtp-model.gguf" + companion.mkdir() + with pytest.raises(HTTPException, match = "regular file"): + _validate_native_gguf_companion(str(companion), str(weight), "MTP drafter") + + +def test_native_companion_rejects_missing_weight(tmp_path): + companion = tmp_path / "mtp-model.gguf" + companion.write_bytes(b"draft") + with pytest.raises(HTTPException, match = "no longer accessible"): + _validate_native_gguf_companion( + str(companion), str(tmp_path / "missing.gguf"), "MTP drafter" + ) + + +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, None, native_grant_backed = True) + + +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) + + +def test_reload_dedup_reloads_for_ordinary_load_when_root_drafter_appears(tmp_path, monkeypatch): + """The native fallback exception must not swallow a newly added root + drafter on an ordinary local load, which can reach it.""" + 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) + + request = LoadRequest(model_path = str(weight)) + # No root drafter yet: both routes dedupe. + assert _request_matches_loaded_settings(request, backend, None, native_grant_backed = True) + assert _request_matches_loaded_settings(request, backend, None, native_grant_backed = False) + + (tmp_path / "mtp-model.gguf").write_bytes(b"root drafter") + # Native cannot reach the root drafter, so the subdir copy stays current. + 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) + + +def test_native_mtp_drafter_rejects_symlinked_later_shard(tmp_path): + """llama-server opens sibling shards implicitly, so validating only the + launch path would let a later shard escape the permitted directory.""" + weight = tmp_path / "model-Q4_0.gguf" + weight.write_bytes(b"model") + sub = tmp_path / "MTP" + sub.mkdir() + first = sub / "mtp-model-Q4_0-00001-of-00002.gguf" + first.write_bytes(b"draft") + outside = tmp_path / "outside.bin" + outside.write_bytes(b"secret") + try: + (sub / "mtp-model-Q4_0-00002-of-00002.gguf").symlink_to(outside) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + with pytest.raises(HTTPException, match = "regular file"): + _validate_native_mtp_drafter(str(first), str(weight), mtp_search_root = str(tmp_path)) + + +def test_native_mtp_drafter_accepts_regular_shard_set(tmp_path): + weight = tmp_path / "model-Q4_0.gguf" + weight.write_bytes(b"model") + sub = tmp_path / "MTP" + sub.mkdir() + first = sub / "mtp-model-Q4_0-00001-of-00002.gguf" + first.write_bytes(b"draft") + (sub / "mtp-model-Q4_0-00002-of-00002.gguf").write_bytes(b"draft") + + _validate_native_mtp_drafter(str(first), str(weight), mtp_search_root = str(tmp_path)) + + +def test_status_provenance_survives_deleted_model_directory(tmp_path, monkeypatch): + """Provenance is a load-time fact: a directory removed underneath a running + server must not turn a local model into a remote one.""" + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)) + backend = LlamaCppBackend() + backend._is_local_model = True + # "outputs/gemma" no longer exists, so is_local_path would call it a repo id. + assert _loaded_is_local_model(backend, False, "outputs/gemma") + + stale = LlamaCppBackend() + assert not _loaded_is_local_model(stale, False, "unsloth/gemma-4-12b") + assert _loaded_is_local_model(stale, True, None) + + +def test_native_load_skips_rejected_mtp_candidate_for_next_one(tmp_path): + """MTP/ can hold several compatible copies. If the size-preferred one is + out of the grant, the next must be tried instead of disabling MTP.""" + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"model") + outside = tmp_path.parent / "outside-blob.gguf" + outside.write_bytes(b"d") + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + try: + (companion_dir / "mtp-model-Q4_0.gguf").symlink_to(outside) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + larger = companion_dir / "mtp-model-Q8_0.gguf" + larger.write_bytes(b"d" * 5000) + + def _usable(candidate: str) -> bool: + return _native_gguf_companion_usable(candidate, str(weight), mtp_search_root = str(tmp_path)) + + # Preferred by size, but it resolves out of the permitted directory. + assert not _usable(detect_mtp_file(str(weight), str(tmp_path), skip_root = True)) + assert detect_mtp_file(str(weight), str(tmp_path), skip_root = True, accept = _usable) == str( + larger.resolve() + ) + + +def test_native_load_returns_none_when_no_candidate_passes(tmp_path): + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"model") + outside = tmp_path.parent / "outside-only.gguf" + outside.write_bytes(b"d") + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + try: + (companion_dir / "mtp-model-Q4_0.gguf").symlink_to(outside) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + def _usable(candidate: str) -> bool: + return _native_gguf_companion_usable(candidate, str(weight), mtp_search_root = str(tmp_path)) + + assert detect_mtp_file(str(weight), str(tmp_path), skip_root = True, accept = _usable) is None diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 893b842e11..d3c91bca93 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -29,7 +29,7 @@ import re import subprocess import sys from pathlib import Path -from typing import List, Tuple +from typing import Callable, List, Tuple import hashlib import json import threading @@ -1301,6 +1301,47 @@ def _colocated_first_split_shard(path: Path) -> tuple[Optional[Path], bool]: return first, first is not None and len(indices) == total +def colocated_split_shards(path: Path) -> tuple[list[Path], bool]: + """Every shard beside *path*, and whether the declared set is complete. + + A non-split path is itself a complete one-file set. Callers that hand a + path to llama-server need this: it opens the sibling shards implicitly, so + an incomplete set fails at startup and every shard needs validating. + """ + match = _GGUF_SPLIT_FILE_RE.match(path.name) + if match is None: + return [path], True + + prefix = match.group("prefix").casefold() + total_text = match.group("total") + total = int(total_text) + if total < 1: + return [], False + + found: dict[int, Path] = {} + try: + for sibling in path.parent.iterdir(): + sibling_match = _GGUF_SPLIT_FILE_RE.match(sibling.name) + if ( + sibling_match is None + or sibling_match.group("prefix").casefold() != prefix + or sibling_match.group("total") != total_text + ): + continue + try: + if not sibling.is_file(): + continue + except OSError: + continue + index = int(sibling_match.group("index")) + if 1 <= index <= total: + found[index] = sibling + except OSError: + return [], False + + return [found[i] for i in sorted(found)], len(found) == total + + def _local_gguf_load_path(path: Path) -> Path: """Choose a loadable local path while preserving complete symlink sets.""" if _GGUF_SPLIT_FILE_RE.match(path.name) is None: @@ -1441,7 +1482,12 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional return str(best[1]) -def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[str]: +def detect_mtp_file( + path: str, + search_root: Optional[str] = None, + skip_root: bool = False, + accept: Optional[Callable[[str], bool]] = None, +) -> Optional[str]: """Find the separate MTP drafter (``mtp-*.gguf``) for a local GGUF model. The drafter that pairs with the main weights sits at the repo/snapshot @@ -1455,31 +1501,164 @@ def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[st unsloth names the drafter ``mtp-.gguf`` where ```` prefixes the weight filename across all Gemma 4 repos (e.g. ``mtp-gemma-4-12B-it.gguf`` next to ``gemma-4-12B-it-qat-Q4_0.gguf``). - An unmatched drafter is skipped (fail-safe: no MTP). + If the root drafter is absent, also accept its precision copy under the + repository's ``MTP/`` directory. An unmatched drafter is skipped. + + ``skip_root`` scans only ``MTP/``, for callers that must discard an + out-of-bounds root drafter and still want the subdir copy (native loads). + ``accept`` filters candidates in preference order, so a caller with extra + rules (a native lease) keeps scanning instead of treating the first + rejection as no drafter at all. """ + + def _pairing_stem(name: str) -> str: + stem = Path(name).stem.lower() + if stem.startswith("mtp-"): + 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]{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. + # The optional bpw modifier goes with it, as _extract_quant_label does. + return re.sub( + rf"-(?:{_GGUF_KNOWN_QUANT_RE.pattern})(?:-[0-9]+(?:\.[0-9]+)?bpw)?$", + "", + 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 + stem = _pairing_stem(candidate.name) + return ( + bool(stem) + and weight_name.startswith(stem) + and (len(weight_name) == len(stem) or not weight_name[len(stem)].isalnum()) + ) + + def _launchable(candidate: Path) -> bool: + # An incomplete split set makes llama-server fail its draft startup and + # would disable MTP entirely, so skip it and let a complete copy win. + try: + _, complete = colocated_split_shards(candidate) + except OSError: + return False + return complete + + def _smallest_first(candidate: Path) -> tuple[int, int, str]: + # Cheapest compatible copy wins. Size first: a fixed precision list + # ranked unknown quants behind BF16, so a small K-quant lost to a far + # larger BF16. Precision breaks size ties, name keeps it stable. + # Candidates are collapsed to shard 1, so a split copy must be summed + # across its shards or it would outrank a smaller single file. + name = candidate.name.lower() + try: + shards, _ = colocated_split_shards(candidate) + size = sum(shard.stat().st_size for shard in shards) + except OSError: + size = sys.maxsize + if "-q4_0" in name: + precision = 0 + elif "-q8_0" in name: + precision = 1 + elif "-bf16" in name or "-f16" in name: + precision = 2 + else: + precision = 3 + return size, precision, name + p = Path(path) weight_name = p.name.lower() if p.suffix.lower() == ".gguf" else None start_dir = p.parent if p.is_file() else p dirs = [start_dir] if search_root is not None: dirs.append(Path(search_root)) - for d in dirs: - try: - entries = sorted(d.iterdir()) - except OSError: - continue - for f in entries: - name = f.name.lower() - if not (name.startswith("mtp-") and name.endswith(".gguf")): - continue - stem = name[len("mtp-") : -len(".gguf")] - if not stem or (weight_name is not None and not weight_name.startswith(stem)): - continue + if not skip_root: + for d in dirs: try: - if f.is_file(): - return str(f.resolve()) + entries = sorted(d.iterdir()) except OSError: continue + for f in entries: + name = f.name.lower() + if not (name.startswith("mtp-") and name.endswith(".gguf")): + continue + if not _matches_weight(f): + continue + try: + if not (f.is_file() and _launchable(f)): + continue + launch = _drafter_launch_path(f) + except OSError: + continue + if accept is not None and not accept(launch): + continue + return launch + + subdir_candidates: list[Path] = [] + for d in dirs: + try: + parent_entries = sorted(d.iterdir()) + except OSError: + continue + mtp_dirs: list[Path] = [] + for entry in parent_entries: + if entry.name.casefold() != "mtp": + continue + try: + if entry.is_dir(): + mtp_dirs.append(entry) + except OSError: + continue + for mtp_dir in mtp_dirs: + try: + entries = sorted(mtp_dir.iterdir()) + except OSError: + continue + for f in entries: + # _is_mtp_drafter accepts everything under MTP/ by design (it + # excludes them from variant menus). Too broad to include on: + # a weight copy here would launch as --model-draft. Require a + # published drafter name: mtp- or -MTP. + lower = f.name.lower() + if not lower.endswith(".gguf"): + continue + # 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() and _launchable(f): + # 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: + resolved = _drafter_launch_path(candidate) + except OSError: + continue + if accept is not None and not accept(resolved): + continue + logger.info(f"Detected MTP subdirectory drafter: {resolved}") + return resolved return None @@ -1696,27 +1875,15 @@ def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str selected = Path(selected_path) gguf_path = Path(gguf_file) - if selected.suffix.lower() != ".gguf": - return selected_path - - gguf_dir = gguf_path.parent - if not gguf_dir.name: - return str(gguf_dir) - - quant_dir_re = ( - r"(UD-)?(" - r"MXFP[0-9]+(?:_[A-Z0-9]+)*" - r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" - r"|TQ[0-9]+_[0-9]+" - r"|Q[0-9]+_K_[A-Z]+" - r"|Q[0-9]+_[0-9]+" - r"|Q[0-9]+_K" - r"|BF16|F16|F32" - r")" - ) - if re.fullmatch(quant_dir_re, gguf_dir.name, re.IGNORECASE): - return str(gguf_dir.parent) - return str(gguf_dir) + # One quant vocabulary, shared: a local copy of it silently fell behind on + # the bpw modifier, which left IQ4_XS-3.53bpw unrecognised as a quant dir. + quant_dir_re = rf"{_GGUF_KNOWN_QUANT_RE.pattern}(-[0-9]+(?:\.[0-9]+)?bpw)?" + search_dir = gguf_path.parent if selected.suffix.lower() == ".gguf" else selected + if not search_dir.name: + return str(search_dir) + if re.fullmatch(quant_dir_re, search_dir.name, re.IGNORECASE): + return str(search_dir.parent) + return str(search_dir) def _iter_hf_cache_snapshots(repo_id: str, cache_dir: Optional[str | Path] = None): diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py index 3ed7faa7c2..d17a96eb00 100644 --- a/studio/backend/utils/native_path_leases.py +++ b/studio/backend/utils/native_path_leases.py @@ -47,6 +47,28 @@ class NativePathLeaseError(ValueError): """Raised when a native path grant is missing, invalid, or unsafe.""" +def native_gguf_companion_parent_allowed( + companion_path: str | Path, + gguf_path: str | Path, + *, + allow_mtp_subdir: bool = False, + mtp_search_root: str | Path | None = None, +) -> bool: + """Check whether a GGUF companion is in an allowed directory.""" + companion_parent = Path(companion_path).resolve(strict = True).parent + gguf_parent = Path(gguf_path).resolve(strict = True).parent + if companion_parent == gguf_parent: + return True + if not allow_mtp_subdir or companion_parent.name.casefold() != "mtp": + return False + allowed_roots = {gguf_parent} + if mtp_search_root is not None: + search_root = Path(mtp_search_root).resolve(strict = True) + if search_root in {gguf_parent, gguf_parent.parent}: + allowed_roots.add(search_root) + return companion_parent.parent in allowed_roots + + @dataclass(frozen = True) class NativePathGrant: operation: str diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index fb9331ecc2..1a0a73f095 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1740,6 +1740,7 @@ async function autoLoadSmallestModel(): Promise<{ customContextLength: config.customContextLength, loadedIsMultimodal: isMultimodalResponse(loadResp), loadedIsDiffusion: loadResp.is_diffusion ?? false, + activeModelIsLocal: loadResp.is_local_model ?? false, ...resolveLoadedSpeculativeSettings(loadResp), }); } else { @@ -1765,6 +1766,7 @@ async function autoLoadSmallestModel(): Promise<{ ...resolveLoadedSpeculativeSettings(loadResp), loadedIsMultimodal: isMultimodalResponse(loadResp), loadedIsDiffusion: loadResp.is_diffusion ?? false, + activeModelIsLocal: loadResp.is_local_model ?? false, }); } if (!(loadResp.is_lora ?? false)) { @@ -2043,6 +2045,7 @@ async function autoLoadSmallestModel(): Promise<{ defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), + activeModelIsLocal: loadResp.is_local_model ?? false, ...resolveLoadedSpeculativeSettings(loadResp), }); recordLastLocalModelLoad({ diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 7b310c50d4..ce6fe6d0f5 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -88,7 +88,10 @@ import { providerSupportsBuiltinCodeExecution, providerSupportsFastMode, } from "./provider-capabilities"; -import { useChatRuntimeStore } from "./stores/chat-runtime-store"; +import { + isLocalModelPath, + useChatRuntimeStore, +} from "./stores/chat-runtime-store"; import type { InferenceParams } from "./types/runtime"; export { defaultInferenceParams, type Preset } from "./presets/preset-policy"; @@ -378,6 +381,9 @@ export function ChatSettingsPanel({ const isMobile = useIsMobile(); const isLoadedGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; const currentCheckpoint = params.checkpoint; + const activeModelIsLocal = useChatRuntimeStore( + (s) => s.activeModelIsLocal, + ); const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); // Direct-file / custom-folder GGUFs load without a variant label but still // report a GGUF context, so detect them via the context and the checkpoint @@ -387,6 +393,13 @@ export function ChatSettingsPanel({ isLoadedGguf || ggufContextLength != null || (currentCheckpoint?.toLowerCase().endsWith(".gguf") ?? false); + // activeModelIsLocal is the backend's own classification and covers native + // picks. Two things must not decide this: activeNativePathToken, which + // status reconciliation keeps across a switch to a remote GGUF (no + // replacement token exists), and a bare .gguf suffix, since the backend + // reads a one-slash org/name.gguf as a repository id, not a file. + const isLocalGguf = + isGguf && (activeModelIsLocal || isLocalModelPath(currentCheckpoint ?? "")); const ggufMaxContextLength = useChatRuntimeStore( (s) => s.ggufMaxContextLength, ); @@ -812,7 +825,9 @@ export function ChatSettingsPanel({ : specFallbackReason === "runtime_error" ? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding." : specFallbackReason === "drafter_not_found" - ? "This model supports MTP, but its drafter file could not be downloaded, so MTP is off and it falls back to n-gram speculative decoding where the llama.cpp build supports it. Check your network connection or Hugging Face access, then reload the model to retry the drafter." + ? isLocalGguf + ? "This local model supports MTP, but no matching drafter file was found. Place its mtp-*.gguf beside the model or in its MTP folder, then reload the model." + : "This model supports MTP, but its drafter file could not be downloaded, so MTP is off and it falls back to n-gram speculative decoding where the llama.cpp build supports it. Check your network connection or Hugging Face access, then reload the model to retry the drafter." : `MTP is not available in the installed llama.cpp build, so this model is running without it.${ llamaUpdateStatus?.update_available ? " Update llama.cpp to enable it." diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 48a6168555..2f19ba0864 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -316,6 +316,9 @@ async function syncInferenceStatusToStore(options?: { syncModelCapabilities(checkpointId, statusRes); } } else if (!statusRes.active_model && !isExternalSelectionActive) { + // specFallbackReason survives here, so clearing activeModelIsLocal + // alone would flip a local model's warning to "download failed". Every + // load path and clearCheckpoint set it, so leave it consistent. useChatRuntimeStore.setState({ modelRequiresTrustRemoteCode: false, loadedIsMultimodal: false, @@ -1040,6 +1043,7 @@ export function useChatModelRuntime() { loadedChatTemplateOverride: effectiveChatTemplateOverride, loadedIsMultimodal: isMultimodalResponse(loadResponse), loadedIsDiffusion: loadResponse.is_diffusion ?? false, + activeModelIsLocal: loadResponse.is_local_model ?? false, activeNativePathToken: nativePathToken ?? null, activeNativePathExpiresAtMs: nativePathToken ? nativePathExpiresAtMs @@ -1149,6 +1153,7 @@ export function useChatModelRuntime() { rollbackResponse.speculative_type, ); useChatRuntimeStore.setState({ + activeModelIsLocal: rollbackResponse.is_local_model ?? false, activeNativePathToken: previousActiveNativePathToken ?? null, // Restore the previous token's lease together with the token so a // rollback never pairs restored token A with failed load B's expiry. diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index f85ff3246b..31ccd3721c 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -289,6 +289,7 @@ export function applyActiveModelStatusToStore( defaultChatTemplate: nextDefaultChatTemplate, loadedIsMultimodal: isMultimodalResponse(status), loadedIsDiffusion: status.is_diffusion ?? false, + activeModelIsLocal: status.is_local_model ?? false, specFallbackReason: status.spec_fallback_reason ?? null, // The spec / KV seeds share the GPU-fields reseed mechanism below: a // non-GGUF status leaves their loaded baselines null, so the "unseeded" diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 3f89f70ef1..28d9e6673b 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -1226,6 +1226,7 @@ export function SharedComposer({ // GPU fields on every load path so the gate can't read stale. loadedIsDiffusion: resp.is_diffusion ?? false, loadedIsMultimodal: isMultimodalResponse(resp), + activeModelIsLocal: resp.is_local_model ?? false, // Record the context this pane loaded with (like the single-model path) // so when it becomes the active model, the UI and later reload/save use // its context, not the previous/default one. diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 237cd857f0..aaaee8af07 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -779,6 +779,8 @@ type ChatRuntimeStore = { // lets the attach gates flag a failed load vs "no model picked". lastModelLoadError: string | null; activeGgufVariant: string | null; + /** Whether the backend loaded the active model from a filesystem path. */ + activeModelIsLocal: boolean; ggufContextLength: number | null; ggufMaxContextLength: number | null; ggufNativeContextLength: number | null; @@ -1316,6 +1318,7 @@ export const useChatRuntimeStore = create((set, get) => ({ modelsError: null, lastModelLoadError: null, activeGgufVariant: null, + activeModelIsLocal: false, ggufContextLength: null, ggufMaxContextLength: null, ggufNativeContextLength: null, @@ -1588,7 +1591,17 @@ export const useChatRuntimeStore = create((set, get) => ({ maxTokens: nextMaxTokens, }, activeGgufVariant: ggufVariant ?? null, - ...(checkpointChanged ? { contextUsage: null } : {}), + // Provenance and the spec-fallback reason both describe the model + // being replaced, so they go together on a real change. Dropping only + // one leaves the settings sheet pairing a stale reason with the wrong + // recovery text. The load or status response reseeds both. + ...(checkpointChanged + ? { + contextUsage: null, + activeModelIsLocal: false, + specFallbackReason: null, + } + : {}), // Switching to an external provider disables Deep Research, which only // applies to the local base model. ...(isExternalModelId(modelId) ? { deepResearchEnabled: false } : {}), @@ -1619,6 +1632,7 @@ export const useChatRuntimeStore = create((set, get) => ({ checkpoint: "", }, activeGgufVariant: null, + activeModelIsLocal: false, activeNativePathToken: null, activeNativePathExpiresAtMs: null, ggufContextLength: null, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index d554eb777e..45e7e02708 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -154,6 +154,7 @@ export interface LoadModelResponse { is_vision: boolean; is_lora: boolean; is_gguf?: boolean; + is_local_model?: boolean; is_diffusion?: boolean; is_audio?: boolean; audio_type?: string | null; @@ -208,6 +209,7 @@ export interface InferenceStatusResponse { model_identifier?: string | null; is_vision: boolean; is_gguf?: boolean; + is_local_model?: boolean; is_diffusion?: boolean; gguf_variant?: string | null; is_audio?: boolean; diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index e1aba66b1b..bd24e74a6d 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -26,6 +26,12 @@ def _read(rel: str) -> str: return path.read_text(encoding = "utf-8") +def _read_backend(rel: str) -> str: + path = WORKDIR / "studio" / "backend" / rel + assert path.exists(), f"missing backend source file: {path}" + return path.read_text() + + def test_models_api_sends_token_via_header_not_query(): """getModelConfig / checkVisionModel / checkEmbeddingModel must pass the HF token through hubTokenHeader, never as a ?hf_token= query param (which leaks @@ -336,6 +342,61 @@ def test_local_gguf_diagnostics_gate_on_broad_is_gguf(): assert vram and "isGguf &&" in vram.group(0) and "isLoadedGguf" not in vram.group(0) +def test_local_mtp_warning_covers_path_and_native_gguf_sources(): + """The local MTP recovery text must cover direct files, custom folders, + and native-picker labels instead of classifying only .gguf suffixes.""" + src = _read("features/chat/chat-settings-sheet.tsx") + local = re.search(r"const isLocalGguf =.*?;", src, re.S) + assert local + assert "isGguf &&" in local.group(0) + assert "activeModelIsLocal" in local.group(0) + assert "isLocalModelPath" in local.group(0) + # Two signals must not classify the model here, because both mislabel a + # remote GGUF as local: a native token, which outlives a switch to a remote + # model, and a bare .gguf suffix, since a one-slash org/name.gguf is a + # repository id. activeModelIsLocal is the backend's own answer for both. + assert "activeNativePathToken" not in local.group(0) + assert ".gguf" not in local.group(0) + + # Switching models must drop both together: a kept flag would classify the + # newly selected model by the old one's provenance. + store = _read("features/chat/stores/chat-runtime-store.ts") + reset = re.search(r"setCheckpoint: \(modelId, ggufVariant\) =>.*?\}\),", store, re.S) + assert reset + assert "activeModelIsLocal: false" in reset.group(0) + assert "specFallbackReason: null" in reset.group(0) + assert "isLocalGguf" in src.split('specFallbackReason === "drafter_not_found"', 1)[1] + + +def test_local_mtp_warning_uses_backend_source_metadata(): + types = _read("features/chat/types/api.ts") + assert types.count("is_local_model?: boolean") >= 2 + + status = _read("features/chat/lib/apply-inference-status-to-store.ts") + assert "activeModelIsLocal: status.is_local_model ?? false" in status + + runtime = _read("features/chat/stores/chat-runtime-store.ts") + assert "activeModelIsLocal: boolean" in runtime + assert runtime.count("activeModelIsLocal: false") >= 2 + + load = _read("features/chat/hooks/use-chat-model-runtime.ts") + assert "activeModelIsLocal: loadResponse.is_local_model ?? false" in load + + models = _read_backend("models/inference.py") + assert models.count("is_local_model: bool = Field(") >= 2 + + route = _read_backend("routes/inference.py") + assert route.count("is_local_model = config.is_local") >= 2 + # GGUF status reports the provenance the load recorded. Re-deriving it from + # the filesystem would flip a local model to remote once its directory goes + # away underneath a running server. + assert "llama_backend._is_local_model = bool(native_grant_backed or config.is_local)" in route + # Both GGUF responses report it: the status poll and the already_loaded + # dedup reply. Either one re-deriving it reintroduces the flip. + assert route.count("is_local_model = _loaded_is_local_model(") >= 2 + assert "backend.active_model_name and is_local_path(backend.active_model_name)" in route + + def test_fixed_layer_gguf_pins_displayed_context(): """An already-loaded auto-fit GGUF saved with Manual fixed GPU layers must pin the shown context, so a later fresh load keeps the fitted placement