ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: harden MTP companion pairing

Browse source
This commit is contained in:
Michael Han 2026-07-24 14:41:34 -07:00
commit af4464a0e0
5 changed files with 103 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

17
studio/backend/routes/inference.py
View file

@ -1020,6 +1020,7 @@ try:
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import (
_local_gguf_companion_search_root,
detect_mtp_file,
load_model_defaults,
)
@ -1059,6 +1060,7 @@ except ImportError:
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import (
_local_gguf_companion_search_root,
detect_mtp_file,
load_model_defaults,
)
@ -3077,6 +3079,7 @@ def _validate_native_gguf_companion(
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
@ -3103,7 +3106,10 @@ def _validate_native_gguf_companion(
)
try:
if not native_gguf_companion_parent_allowed(
companion, gguf, allow_mtp_subdir = allow_mtp_subdir
companion,
gguf,
allow_mtp_subdir = allow_mtp_subdir,
mtp_search_root = mtp_search_root,
):
location = (
"beside the selected GGUF or in its MTP directory"
@ -3337,7 +3343,10 @@ 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)
stored = llama_backend.mtp_draft_path
try:
detected_resolved = Path(detected).resolve() if detected else None
@ -4656,11 +4665,15 @@ async def _load_model_impl(
# The drafter is optional (unlike mmproj for a vision
# model): drop it rather than fail the load.
try:
mtp_search_root = _local_gguf_companion_search_root(
config.gguf_file, config.gguf_file
)
_validate_native_gguf_companion(
config.gguf_mtp_file,
config.gguf_file,
"MTP drafter",
allow_mtp_subdir = True,
mtp_search_root = mtp_search_root,
)
except HTTPException as exc:
logger.warning("Dropping MTP drafter for native load: %s", exc.detail)

14
studio/backend/tests/test_mtp_drafter_companion.py
View file

@ -259,6 +259,20 @@ def test_detect_mtp_file_subdir_skips_foreign_drafter(tmp_path):
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")

58
studio/backend/tests/test_native_gguf_companion.py
View file

@ -16,9 +16,13 @@ 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 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
@ -41,6 +45,60 @@ def test_native_mtp_companion_allows_mtp_directory(tmp_path, folder):
)
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_native_vision_companion_rejects_mtp_directory(tmp_path):
weight, companion = _write_pair(tmp_path, "MTP")
with pytest.raises(HTTPException, match = "must live next to"):

6
studio/backend/utils/models/model_config.py
View file

@ -1471,7 +1471,11 @@ def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[st
if weight_name is None:
return True
stem = _pairing_stem(candidate.name)
return bool(stem) and weight_name.startswith(stem)
return (
bool(stem)
and weight_name.startswith(stem)
and (len(weight_name) == len(stem) or not weight_name[len(stem)].isalnum())
)
def _precision_rank(candidate: Path) -> tuple[int, str]:
name = candidate.name.lower()

16
studio/backend/utils/native_path_leases.py
View file

@ -52,15 +52,21 @@ def native_gguf_companion_parent_allowed(
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
return companion_parent == gguf_parent or bool(
allow_mtp_subdir
and companion_parent.parent == gguf_parent
and companion_parent.name.casefold() == "mtp"
)
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)