From 8ef7c9c345ed194d7f60d2c418be7b5c4eec6bc7 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Fri, 24 Jul 2026 04:37:56 -0700 Subject: [PATCH] Studio: allow native MTP subdirectory companions --- studio/backend/routes/inference.py | 26 +++- .../tests/test_mtp_drafter_companion.py | 56 +++++++++ .../tests/test_native_gguf_companion.py | 119 ++++++++++++++++++ studio/backend/utils/models/model_config.py | 27 ++-- studio/backend/utils/native_path_leases.py | 16 +++ 5 files changed, 232 insertions(+), 12 deletions(-) create mode 100644 studio/backend/tests/test_native_gguf_companion.py diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 445a26f04d..bc48f0b0bf 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1027,6 +1027,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, ) @@ -1065,6 +1066,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, ) @@ -3070,11 +3072,15 @@ 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, ) -> 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 @@ -3096,10 +3102,17 @@ 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 + ): + 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( @@ -4644,7 +4657,10 @@ async def _load_model_impl( # model): drop it rather than fail the load. try: _validate_native_gguf_companion( - config.gguf_mtp_file, config.gguf_file, "MTP drafter" + config.gguf_mtp_file, + config.gguf_file, + "MTP drafter", + allow_mtp_subdir = True, ) except HTTPException as exc: logger.warning("Dropping MTP drafter for native load: %s", exc.detail) diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 999e1bd0ca..ed8a6a5d0e 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -34,6 +34,7 @@ from utils.models.model_config import ( detect_mtp_file, extract_model_size_b, ) +from utils.native_path_leases import native_gguf_companion_parent_allowed # ── Predicate + layering mirrors ───────────────────────────────────── @@ -258,6 +259,61 @@ def test_detect_mtp_file_subdir_skips_foreign_drafter(tmp_path): 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 ──────────────────────────────── 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..d4fe8ba831 --- /dev/null +++ b/studio/backend/tests/test_native_gguf_companion.py @@ -0,0 +1,119 @@ +# 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 + + +def _write_pair(tmp_path: Path, folder: str | None = None) -> tuple[Path, Path]: + 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_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") diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 6fa1b85c08..4c71414559 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1510,20 +1510,33 @@ def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[st subdir_candidates: list[Path] = [] for d in dirs: - mtp_dir = d / "MTP" try: - entries = sorted(mtp_dir.iterdir()) + parent_entries = sorted(d.iterdir()) except OSError: continue - for f in entries: - rel = f"MTP/{f.name}" - if not _is_mtp_drafter(rel) or not _matches_weight(f): + mtp_dirs: list[Path] = [] + for entry in parent_entries: + if entry.name.casefold() != "mtp": continue try: - if f.is_file(): - subdir_candidates.append(f) + 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: + rel = f"MTP/{f.name}" + if not _is_mtp_drafter(rel) or not _matches_weight(f): + continue + try: + if f.is_file(): + subdir_candidates.append(f) + except OSError: + continue for candidate in sorted(subdir_candidates, key = _precision_rank): try: diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py index 3ed7faa7c2..090cedfeb6 100644 --- a/studio/backend/utils/native_path_leases.py +++ b/studio/backend/utils/native_path_leases.py @@ -47,6 +47,22 @@ 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, +) -> 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 + return companion_parent == gguf_parent or bool( + allow_mtp_subdir + and companion_parent.parent == gguf_parent + and companion_parent.name.casefold() == "mtp" + ) + + @dataclass(frozen = True) class NativePathGrant: operation: str