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

Studio: allow native MTP subdirectory companions

Browse source
This commit is contained in:
Michael Han 2026-07-24 04:37:56 -07:00
commit 8ef7c9c345
5 changed files with 232 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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)

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

@ -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 ────────────────────────────────

119
studio/backend/tests/test_native_gguf_companion.py Normal file
View file

@ -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")

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

@ -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:

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

@ -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