diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1944bc9dd4..6f204b7e38 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1022,6 +1022,7 @@ try: 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, ) @@ -1063,6 +1064,7 @@ except ImportError: 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, ) @@ -3138,22 +3140,56 @@ 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. Same rules as - _validate_native_gguf_companion, as a predicate for reload dedup.""" + """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_gguf_companion( - companion_path, - gguf_path, - "MTP drafter", - allow_mtp_subdir = True, - mtp_search_root = mtp_search_root, - ) + _validate_native_mtp_drafter(companion_path, gguf_path, mtp_search_root = mtp_search_root) except HTTPException: return False return True @@ -4726,11 +4762,9 @@ async def _load_model_impl( def _mtp_allowed(candidate: str) -> bool: try: - _validate_native_gguf_companion( + _validate_native_mtp_drafter( candidate, config.gguf_file, - "MTP drafter", - allow_mtp_subdir = True, mtp_search_root = mtp_search_root, ) return True @@ -4865,6 +4899,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}") @@ -5999,7 +6037,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 = _native_grant_backed or bool(_model_id and is_local_path(_model_id)), + 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), diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 98a233ac2b..bf1ef1a061 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -757,3 +757,45 @@ def test_detect_mtp_file_keeps_snapshot_path_for_sharded_root_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_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()) diff --git a/studio/backend/tests/test_native_gguf_companion.py b/studio/backend/tests/test_native_gguf_companion.py index ab1ba4f7fe..3b0df808dd 100644 --- a/studio/backend/tests/test_native_gguf_companion.py +++ b/studio/backend/tests/test_native_gguf_companion.py @@ -17,6 +17,8 @@ if _BACKEND_DIR not in sys.path: 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 core.inference.llama_cpp import LlamaCppBackend from models.inference import LoadRequest @@ -286,3 +288,49 @@ def test_reload_dedup_native_load_with_no_admissible_drafter(tmp_path, monkeypat 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) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 6e49294e3f..523b2a7b30 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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: @@ -1476,7 +1517,13 @@ def detect_mtp_file( if stem.endswith("-mtp"): stem = stem[: -len("-mtp")] # 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) + # 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 @@ -1497,13 +1544,25 @@ def detect_mtp_file( 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: - size = candidate.stat().st_size + 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: @@ -1535,7 +1594,7 @@ def detect_mtp_file( if not _matches_weight(f): continue try: - if f.is_file(): + if f.is_file() and _launchable(f): return _drafter_launch_path(f) except OSError: continue @@ -1577,7 +1636,7 @@ def detect_mtp_file( if not _matches_weight(f): continue try: - if f.is_file(): + 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: diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index c547f4927b..dedda3f297 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -376,7 +376,11 @@ def test_local_mtp_warning_uses_backend_source_metadata(): route = _read_backend("routes/inference.py") assert route.count("is_local_model = config.is_local") >= 2 - assert "is_local_model = _native_grant_backed" in route + # 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 + assert "is_local_model = _loaded_is_local_model(" in route assert "backend.active_model_name and is_local_path(backend.active_model_name)" in route