Fix case-variant model matching and GGUF cache reuse in unsloth start (#6900)
* fix: handle case-variant GGUF cache hits for unsloth start * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gguf cache: keep split shards co-located and isolate cache tests properly When a cached main shard was reused from an older snapshot, the extra shards were resolved independently and could come from a different snapshot dir (or a fresh download into the current ref), leaving llama.cpp unable to load a multi-shard GGUF whose pieces are split across directories. Only reuse a cached main shard when every sibling shard sits in the same snapshot; otherwise fetch the whole set together so they stay co-located. Also patch huggingface_hub.constants.HF_HUB_CACHE (not just the HF_HUB_CACHE env var) in the two cache tests that seeded a temp cache: the snapshot lookup reads the module constant, so the env-only override let the real cache leak in and skip an asserted download. * Do not let a companion-only cache snapshot shadow real GGUF variants When listing GGUF variants from the local HF cache, a newer snapshot may contain only a companion file (for example a vision projector fetched on demand) while the actual quant files live in an older snapshot. The prior scan returned the first snapshot whose vision flag was set, yielding an empty variant list and hiding the real quants. Keep scanning older snapshots for actual variants and carry the vision flag across snapshots. Also record the disk-space fallback variant's size in expected_sizes so the later cache-reuse probe can size-verify the fallback main shard instead of only checking for its existence. * Propagate cached repo casing to companions and preflight split co-location Two fixes to the case-variant GGUF cache reuse: - Resolve the requested repo id to its cached canonical casing once in load_model, up front, and pass it to the main GGUF and its companions (mmproj / MTP drafter). Previously only _download_gguf resolved the casing internally, so a case-variant request loaded the main file from the canonical cache dir while the companions kept the requested casing and missed the cached vision projector / drafter offline. Extracted the resolution into a shared _resolve_repo_id_casing helper. - Apply the split-shard co-location check in the disk-space preflight. When a split GGUF's shards are cached across different snapshots the whole set is refetched later, so counting them as cached made the preflight read 0 bytes to download, skip the smaller-variant fallback, and then fail the full download on a low-disk machine. * Reuse a co-located split GGUF snapshot and fix split fallback size probe - When reusing a cached split GGUF, scan snapshots for one that holds the whole set co-located instead of taking the newest snapshot's first shard. A newer snapshot with only the first shard no longer shadows an older complete snapshot, so an already-cached split model is reused rather than refetched (which would fail offline). - The disk-space fallback records its size in expected_sizes only for a single-file fallback. _find_smallest_fitting_variant returns the whole variant size, so using it as the first shard's expected size rejected a valid cached first shard of a split fallback and forced a re-download. * Scan for a complete split snapshot in the preflight; require a loaded catalog hit - The disk-space preflight now uses the same co-located snapshot scan as the download path (_cached_colocated_split_main) instead of the newest-snapshot probe, so a newer snapshot holding only the first shard no longer masks an older complete one and trips the smaller-variant fallback for a fully cached split model. - _resolve_model only attaches to a /v1/models entry that is actually loaded (loaded != False). /v1/models also lists cached-but-unloaded catalog entries, and matching one by case skipped /api/inference/load and left the agent pointed at a model that is not resident. * Restrict cross-snapshot GGUF cache reuse to offline Reusing a same-name blob from an older or case-variant snapshot bypasses the Hub revision/etag check, so a repo that updates a GGUF in place could serve stale weights online. Gate the cross-snapshot and case-variant reuse (both the disk-space preflight accounting and the download path) on HF_HUB_OFFLINE. Online, hf_hub_download fetches the current revision and resumes a partial download, so the reuse is unnecessary there; offline it remains the resilience fallback. Marked the two reuse regression tests as the offline scenarios they represent and added an online test asserting a fresh fetch. * Harden offline cache reuse and hub-id detection Three follow-ups on the case-variant GGUF cache path: - Honor every truthy HF_HUB_OFFLINE spelling (1/true/yes/on), not just "1", when gating the cross-snapshot and case-variant cache reuse. With HF_HUB_OFFLINE=true the Hub calls are already offline, so the reuse must trigger or the cached GGUF fails to load; route both the preflight accounting and the download path through the same offline parse the rest of the backend uses. - Resolve mmproj/MTP companions from the actual cached snapshot when offline. resolve_cached_repo_id_case can keep a partial lower-case spelling when any dir exists under the requested casing, so an hf_hub_download on that casing misses the canonical companion; scan every case-variant snapshot and return the cached path. - Restrict the case-insensitive model-id match to syntactically valid hub ids (a single namespace/name over the HF charset). A server-side relative path such as models/Llama/Foo.gguf is no longer treated as a hub id, so it cannot casefold-match a differently cased path on a case-sensitive filesystem. This is host independent, unlike the local-existence probe which cannot see a server path. * Only casefold-match model ids against a loopback Studio A two-segment string like Models/Foo is indistinguishable from a hub id, and the local Path.exists() probe in _is_hub_model_id cannot see a path that exists only on a remote Studio host. So against a remote server, casefolding could attach to a distinct server-side path (Models/Foo vs models/foo) on a case-sensitive filesystem. Gate the case-insensitive match on is_loopback_url(base): only a local Studio, where the existence probe is authoritative, casefolds. For a remote Studio the match is exact and a case-mismatched request falls through to /api/inference/load, whose already-loaded dedup resolves it correctly. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
This commit is contained in:
parent
a113f893ea
commit
df6b5a57d9
5 changed files with 874 additions and 34 deletions
|
|
@ -597,6 +597,59 @@ def _loaded_models(base: str, key: str) -> list:
|
|||
return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", [])
|
||||
|
||||
|
||||
_HF_REPO_ID_SEGMENT_RE = re.compile(r"^[A-Za-z0-9._-]+$")
|
||||
|
||||
|
||||
def _is_hub_model_id(value: object) -> bool:
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
text = value.strip()
|
||||
if "\\" in text:
|
||||
return False
|
||||
if text.startswith(("/", "./", "../", "~")):
|
||||
return False
|
||||
if len(text) >= 2 and text[1] == ":" and text[0].isalpha():
|
||||
return False
|
||||
# A hub id is exactly "namespace/name" over a restricted charset. Anything with
|
||||
# extra path segments (e.g. a server-side relative path such as
|
||||
# models/Llama/Foo.gguf on a remote Studio) is not a hub id and must not be
|
||||
# casefold-matched against a differently cased path on a case-sensitive
|
||||
# filesystem. This is host independent, unlike the existence probe below which
|
||||
# cannot see a path that only exists on the server.
|
||||
parts = text.split("/")
|
||||
if len(parts) != 2:
|
||||
return False
|
||||
if any(part in ("", ".", "..") or not _HF_REPO_ID_SEGMENT_RE.match(part) for part in parts):
|
||||
return False
|
||||
try:
|
||||
if Path(os.path.expanduser(text)).exists():
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _model_id_matches(
|
||||
actual: object,
|
||||
requested: object,
|
||||
*,
|
||||
allow_casefold: bool = True,
|
||||
) -> bool:
|
||||
if actual == requested:
|
||||
return True
|
||||
# Case-insensitive matching is only safe when the local existence probe in
|
||||
# _is_hub_model_id is authoritative, i.e. against a loopback Studio on this host.
|
||||
# Against a remote Studio a two-segment string is indistinguishable from a
|
||||
# server-side relative path (e.g. Models/Foo vs models/foo), so casefolding it
|
||||
# could attach to the wrong model on a case-sensitive server; defer to an exact
|
||||
# match there and let the load endpoint resolve the requested path.
|
||||
if not allow_casefold:
|
||||
return False
|
||||
if not (_is_hub_model_id(actual) and _is_hub_model_id(requested)):
|
||||
return False
|
||||
return str(actual).casefold() == str(requested).casefold()
|
||||
|
||||
|
||||
def _resolve_model(
|
||||
base: str,
|
||||
key: str,
|
||||
|
|
@ -604,6 +657,9 @@ def _resolve_model(
|
|||
load: LoadOptions = LoadOptions(),
|
||||
) -> dict:
|
||||
models = _loaded_models(base, key)
|
||||
# Only casefold-match ids against a loopback Studio, where _is_hub_model_id's
|
||||
# local existence probe can actually reject a server-side path; see the note there.
|
||||
allow_casefold = is_loopback_url(base)
|
||||
# /v1/models reports the model id but not the active GGUF variant or runtime load
|
||||
# settings, so an id match alone can hide the wrong quant (Q8_0 serving while the
|
||||
# user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to
|
||||
|
|
@ -613,10 +669,21 @@ def _resolve_model(
|
|||
load_has_overrides = bool(
|
||||
load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel
|
||||
)
|
||||
# /v1/models also lists cached-but-unloaded catalog entries (loaded == False);
|
||||
# matching one would skip /api/inference/load and leave the agent pointed at a
|
||||
# model that is not resident, so only attach to an entry that is actually loaded.
|
||||
match = (
|
||||
None
|
||||
if requested and load_has_overrides
|
||||
else next((m for m in models if m["id"] == requested), None)
|
||||
else next(
|
||||
(
|
||||
m
|
||||
for m in models
|
||||
if _model_id_matches(m.get("id"), requested, allow_casefold = allow_casefold)
|
||||
and m.get("loaded") is not False
|
||||
),
|
||||
None,
|
||||
)
|
||||
)
|
||||
if requested and match is None:
|
||||
typer.echo(
|
||||
|
|
@ -651,7 +718,16 @@ def _resolve_model(
|
|||
if isinstance(loaded, dict):
|
||||
wanted |= {loaded.get("model"), loaded.get("display_name")} - {None}
|
||||
models = _loaded_models(base, key)
|
||||
match = next((m for m in models if m["id"] in wanted), None)
|
||||
match = next(
|
||||
(
|
||||
m
|
||||
for m in models
|
||||
if any(
|
||||
_model_id_matches(m.get("id"), w, allow_casefold = allow_casefold) for w in wanted
|
||||
)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if match is not None:
|
||||
return match
|
||||
if requested:
|
||||
|
|
|
|||
|
|
@ -457,6 +457,189 @@ def test_connect_codex_no_launch(fake_studio, tmp_path):
|
|||
assert (home / "unsloth_api.config.toml").exists()
|
||||
|
||||
|
||||
def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, tmp_path):
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
[
|
||||
"codex",
|
||||
"--no-launch",
|
||||
"--model",
|
||||
"unsloth/gemma-4-26b-a4b-it-gguf",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
home = tmp_path / "agents" / "codex"
|
||||
profile = _parse_toml((home / "unsloth_api.config.toml").read_text())
|
||||
assert profile["model"] == MODEL["id"]
|
||||
|
||||
|
||||
def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch):
|
||||
calls = []
|
||||
state = {"loaded": False}
|
||||
|
||||
def http_json(
|
||||
method,
|
||||
url,
|
||||
token,
|
||||
payload = None,
|
||||
timeout = 30,
|
||||
error = None,
|
||||
):
|
||||
calls.append((method, url, payload))
|
||||
if url.endswith("/v1/models"):
|
||||
return {
|
||||
"data": [
|
||||
{
|
||||
"id": "unsloth/gemma-4-E2B-it-GGUF" if state["loaded"] else "other/model",
|
||||
"context_length": 131072,
|
||||
}
|
||||
]
|
||||
}
|
||||
if url.endswith("/api/inference/load"):
|
||||
state["loaded"] = True
|
||||
return {"model": "unsloth/gemma-4-E2B-it-GGUF"}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", http_json)
|
||||
|
||||
entry = start._resolve_model(
|
||||
BASE,
|
||||
"sk-test",
|
||||
"unsloth/gemma-4-e2b-it-gguf",
|
||||
start.LoadOptions(gguf_variant = "UD-Q4_K_XL"),
|
||||
)
|
||||
|
||||
assert entry["id"] == "unsloth/gemma-4-E2B-it-GGUF"
|
||||
assert any(c[1].endswith("/api/inference/load") for c in calls)
|
||||
|
||||
|
||||
def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch):
|
||||
# A cached-but-unloaded catalog entry (loaded == False) that only case-differs must
|
||||
# not be treated as ready; the load endpoint must still be called so the requested
|
||||
# model becomes resident instead of the agent preflighting a different backend.
|
||||
calls = []
|
||||
state = {"loaded": False}
|
||||
|
||||
def http_json(
|
||||
method,
|
||||
url,
|
||||
token,
|
||||
payload = None,
|
||||
timeout = 30,
|
||||
error = None,
|
||||
):
|
||||
calls.append((method, url))
|
||||
if url.endswith("/v1/models"):
|
||||
return {
|
||||
"data": [
|
||||
{
|
||||
"id": "unsloth/Gemma-4-GGUF",
|
||||
"loaded": state["loaded"],
|
||||
"context_length": 131072,
|
||||
}
|
||||
]
|
||||
}
|
||||
if url.endswith("/api/inference/load"):
|
||||
state["loaded"] = True
|
||||
return {"model": "unsloth/Gemma-4-GGUF"}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", http_json)
|
||||
|
||||
entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf")
|
||||
|
||||
assert entry["id"] == "unsloth/Gemma-4-GGUF"
|
||||
assert any(u.endswith("/api/inference/load") for _, u in calls)
|
||||
|
||||
|
||||
def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch):
|
||||
# The mirror case: a loaded entry (loaded == True) that case-matches attaches with
|
||||
# no /api/inference/load call.
|
||||
calls = []
|
||||
|
||||
def http_json(
|
||||
method,
|
||||
url,
|
||||
token,
|
||||
payload = None,
|
||||
timeout = 30,
|
||||
error = None,
|
||||
):
|
||||
calls.append((method, url))
|
||||
if url.endswith("/v1/models"):
|
||||
return {
|
||||
"data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}]
|
||||
}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", http_json)
|
||||
|
||||
entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf")
|
||||
|
||||
assert entry["id"] == "unsloth/Gemma-4-GGUF"
|
||||
assert not any(u.endswith("/api/inference/load") for _, u in calls)
|
||||
|
||||
|
||||
def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch):
|
||||
# Against a remote Studio the local existence probe cannot see server-side paths,
|
||||
# so a case-variant loaded id must NOT attach without a load: it could be a distinct
|
||||
# server-side path on a case-sensitive host. The load endpoint resolves the request.
|
||||
calls = []
|
||||
state = {"loaded": False}
|
||||
|
||||
def http_json(
|
||||
method,
|
||||
url,
|
||||
token,
|
||||
payload = None,
|
||||
timeout = 30,
|
||||
error = None,
|
||||
):
|
||||
calls.append((method, url))
|
||||
if url.endswith("/v1/models"):
|
||||
return {
|
||||
"data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}]
|
||||
}
|
||||
if url.endswith("/api/inference/load"):
|
||||
state["loaded"] = True
|
||||
return {"model": "unsloth/Gemma-4-GGUF"}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", http_json)
|
||||
|
||||
entry = start._resolve_model("http://10.0.0.5:8888", "sk-test", "unsloth/gemma-4-gguf")
|
||||
|
||||
# The load endpoint was consulted (no casefold shortcut), and we still attach to the
|
||||
# server's canonical id it reports back.
|
||||
assert entry["id"] == "unsloth/Gemma-4-GGUF"
|
||||
assert any(u.endswith("/api/inference/load") for _, u in calls)
|
||||
|
||||
|
||||
def test_model_id_matching_does_not_casefold_local_paths(tmp_path):
|
||||
existing_local = tmp_path / "Org" / "Foo"
|
||||
existing_local.mkdir(parents = True)
|
||||
|
||||
assert start._model_id_matches("Org/Foo", "org/foo")
|
||||
assert not start._model_id_matches(str(existing_local), str(existing_local).lower())
|
||||
assert not start._model_id_matches("./Models/Foo", "./models/foo")
|
||||
assert not start._model_id_matches(r".\Models\Foo", r".\models\foo")
|
||||
# A server-side relative path (extra path segments) is not a hub id even when it
|
||||
# does not exist on the CLI host, so it must not casefold-match a differently
|
||||
# cased path on a case-sensitive server filesystem.
|
||||
assert not start._is_hub_model_id("models/Llama/Foo.gguf")
|
||||
assert not start._model_id_matches("models/Llama/Foo.gguf", "models/llama/foo.gguf")
|
||||
# A genuine two-segment hub id still matches case-insensitively.
|
||||
assert start._is_hub_model_id("unsloth/Gemma-3-4b-it-GGUF")
|
||||
assert start._model_id_matches("unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf")
|
||||
# Casefolding is gated to loopback studios (allow_casefold). With it disabled (a
|
||||
# remote studio, where a two-segment string could be a server-side path), even a
|
||||
# genuine hub-id case variant must not match, so the load endpoint resolves it.
|
||||
assert not start._model_id_matches(
|
||||
"unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf", allow_casefold = False
|
||||
)
|
||||
assert start._model_id_matches("unsloth/Foo", "unsloth/Foo", allow_casefold = False)
|
||||
|
||||
|
||||
def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch):
|
||||
# Launch mode writes config to a throwaway temp CODEX_HOME and removes it after
|
||||
# the agent exits; the user's real ~/.codex is never the target.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue