Studio: complete shard handling and persist local provenance
All five review points on the previous revision were correct. - pair bpw-qualified drafters such as mtp-model-IQ4_XS-3.53bpw.gguf. The anchored quant strip left the modifier behind, so the name never matched even though _extract_quant_label supports these filenames - skip an incomplete split candidate instead of launching shard 1 of a set that llama-server cannot start, which also let it outrank a complete copy and disable MTP outright - rank a split copy by its total shard size. Candidates collapse to shard 1, so a 90+90 split was beating a 100 byte single file - validate every shard of a split drafter under a native lease. Sibling shards are opened implicitly by llama-server, so checking only the launch path let a later shard be a symlink out of the permitted directory - record provenance at load time and report that from status. Deleting or unmounting a bare relative model directory underneath a running server made is_local_path read the identifier as an org/model repo id and flipped a local model to remote colocated_split_shards is the one place that enumerates a shard set, so detection, ranking and native validation cannot disagree about it. The model picker contract now pins the provenance helper rather than the inline status expression it replaced.
This commit is contained in:
Michael Han
parent
e9784fd212
commit
2d1e2cc180
5 changed files with 208 additions and 15 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue