Tighten unsloth start model switching
This commit is contained in:
oobabooga
parent
2735cec36c
commit
b9f10d484f
4 changed files with 163 additions and 161 deletions
|
|
@ -3454,25 +3454,6 @@ async def _wait_for_model_switch_idle(*, current_request_counted: bool) -> None:
|
|||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
# A second waiter map tracks requests during the slow resolve phase, before they
|
||||
# can join the concrete target queue above.
|
||||
_auto_switch_request_waiters: dict[str, int] = {}
|
||||
_auto_switch_request_waiters_guard = threading.Lock()
|
||||
|
||||
|
||||
def _request_waiter_key(requested_model: str) -> str:
|
||||
return requested_model.strip().lower()
|
||||
|
||||
|
||||
def _note_request_waiter(key: str, delta: int) -> None:
|
||||
with _auto_switch_request_waiters_guard:
|
||||
n = _auto_switch_request_waiters.get(key, 0) + delta
|
||||
if n > 0:
|
||||
_auto_switch_request_waiters[key] = n
|
||||
else:
|
||||
_auto_switch_request_waiters.pop(key, None)
|
||||
|
||||
|
||||
def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]:
|
||||
"""The id to report for the loaded GGUF in API responses: the advertised repo
|
||||
id from an auto-switch load, else the cleaned public id, never the on-disk
|
||||
|
|
@ -3618,144 +3599,136 @@ async def _maybe_auto_switch_model(
|
|||
if not auto_switch_on and get_auto_unload_idle_seconds() <= 0:
|
||||
return
|
||||
|
||||
# Register by the raw requested model before resolving, which can be slow.
|
||||
request_key = _request_waiter_key(requested_model)
|
||||
_note_request_waiter(request_key, 1)
|
||||
try:
|
||||
# Off the loop: a cold-cache rebuild walks several model dirs + HF caches.
|
||||
# With auto-switch off (or an omitted-model reload-only request), skip the
|
||||
# resolve so only the reload-stash path runs and no name is ever matched.
|
||||
reload_only = requested_model == _RELOAD_ONLY_MODEL
|
||||
resolved = (
|
||||
await asyncio.to_thread(resolve_local_gguf, requested_model)
|
||||
if auto_switch_on and not reload_only
|
||||
else None
|
||||
)
|
||||
if resolved is None:
|
||||
# Idle-unload may have freed the model; reload exactly what it freed
|
||||
# (path + quant + advertised id) so an alias/unknown name stays servable
|
||||
# and keeps the override keyed by the advertised id, not the load path.
|
||||
last = get_last_unloaded_model()
|
||||
# A non-GGUF (Unsloth/Transformers) model loaded after the idle-unload
|
||||
# leaves the GGUF slot empty but is the live model, so don't resurrect
|
||||
# the stale GGUF over it (that load would tear the active model down).
|
||||
if (
|
||||
not last
|
||||
or get_llama_cpp_backend().is_loaded
|
||||
or getattr(get_inference_backend(), "active_model_name", None)
|
||||
):
|
||||
return
|
||||
if len(last) == 3:
|
||||
target_id, variant, override_id = last
|
||||
else: # pre-3-tuple stash: fall back to the path as the override key
|
||||
target_id, variant = last
|
||||
override_id = target_id
|
||||
else:
|
||||
# load_path is a concrete local path (never the bare repo id), so /load
|
||||
# takes the local branch and cannot trigger a download. override_id is the
|
||||
# advertised repo id, the launch-override key and the public model id.
|
||||
target_id, variant, override_id = resolved
|
||||
backend = get_llama_cpp_backend()
|
||||
# A bare model id (no :VARIANT) is satisfied by any loaded quant of that
|
||||
# repo, so it never reloads a different local quant that already serves it.
|
||||
bare = ":" not in requested_model
|
||||
|
||||
def _already_serving() -> bool:
|
||||
# Match against both the concrete load path and the advertised repo id,
|
||||
# so a model loaded manually by repo id (identifier = repo id) and one
|
||||
# loaded by auto-switch (identifier = path, advertised = repo id) both
|
||||
# count as already serving rather than triggering a needless reswap.
|
||||
if not backend.is_loaded or not backend.model_identifier:
|
||||
return False
|
||||
loaded_keys = {backend.model_identifier.lower()}
|
||||
advertised = getattr(backend, "_openai_advertised_id", None)
|
||||
if advertised:
|
||||
loaded_keys.add(advertised.lower())
|
||||
if loaded_keys.isdisjoint({target_id.lower(), override_id.lower()}):
|
||||
return False
|
||||
if bare:
|
||||
return True
|
||||
if variant:
|
||||
loaded_variant = (getattr(backend, "hf_variant", None) or "").lower()
|
||||
return loaded_variant == variant.lower()
|
||||
return True
|
||||
|
||||
def _record_serving_alias() -> None:
|
||||
# When an advertised alias already resolves to the loaded model (e.g. a
|
||||
# model loaded by local path, requested by its repo/LM Studio id), record
|
||||
# the alias as the public id so /v1/models and responses report it (and
|
||||
# mark it loaded) instead of the path-derived basename. Resolver branch
|
||||
# only: the reload-stash override_id can be the bare path, not a repo id.
|
||||
# Lock-free is safe here: an in-flight request blocks any concurrent swap
|
||||
# (single-slot busy guard), so the loaded model can't change under this.
|
||||
if resolved is None or not override_id:
|
||||
return
|
||||
b = get_llama_cpp_backend()
|
||||
if getattr(b, "_openai_advertised_id", None) != override_id:
|
||||
b._openai_advertised_id = override_id
|
||||
|
||||
if _already_serving():
|
||||
_record_serving_alias()
|
||||
return
|
||||
# An image/audio request naming a different text-only GGUF would load it
|
||||
# here and only 400 below, evicting the working model. Reject before the
|
||||
# swap. Only the resolver branch (an explicit new target); the reload-stash
|
||||
# path just restores the model the request was already using. Both vision and
|
||||
# audio input come from a companion mmproj (a filesystem probe) -- run it off
|
||||
# the loop, like the resolver above.
|
||||
# Off the loop: a cold-cache rebuild walks several model dirs + HF caches.
|
||||
# With auto-switch off (or an omitted-model reload-only request), skip the
|
||||
# resolve so only the reload-stash path runs and no name is ever matched.
|
||||
reload_only = requested_model == _RELOAD_ONLY_MODEL
|
||||
resolved = (
|
||||
await asyncio.to_thread(resolve_local_gguf, requested_model)
|
||||
if auto_switch_on and not reload_only
|
||||
else None
|
||||
)
|
||||
if resolved is None:
|
||||
# Idle-unload may have freed the model; reload exactly what it freed
|
||||
# (path + quant + advertised id) so an alias/unknown name stays servable
|
||||
# and keeps the override keyed by the advertised id, not the load path.
|
||||
last = get_last_unloaded_model()
|
||||
# A non-GGUF (Unsloth/Transformers) model loaded after the idle-unload
|
||||
# leaves the GGUF slot empty but is the live model, so don't resurrect
|
||||
# the stale GGUF over it (that load would tear the active model down).
|
||||
if (
|
||||
require_vision
|
||||
and resolved is not None
|
||||
and not await asyncio.to_thread(_target_is_vision, target_id)
|
||||
not last
|
||||
or get_llama_cpp_backend().is_loaded
|
||||
or getattr(get_inference_backend(), "active_model_name", None)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = openai_error_body(
|
||||
"The requested model does not support the image or audio input in this request.",
|
||||
status = 400,
|
||||
code = "invalid_value",
|
||||
param = "model",
|
||||
),
|
||||
)
|
||||
key = _switch_key(override_id, variant)
|
||||
_note_switch_waiter(key, 1)
|
||||
try:
|
||||
async with _auto_switch_lock():
|
||||
# The asyncio lock is per loop; add a process-wide gate so a swap on
|
||||
# another loop in this process can't race the single slot.
|
||||
await _acquire_swap_gate()
|
||||
try:
|
||||
# Hold the keep-warm gate across the swap so no new inference can
|
||||
# start on the model while it is being torn down and replaced.
|
||||
async with inference_lifecycle_gate():
|
||||
if _already_serving():
|
||||
_record_serving_alias()
|
||||
return
|
||||
# Apply this model's saved launch flags so the swap honors the config.
|
||||
override = get_model_override(override_id)
|
||||
load_kwargs = {"model_path": target_id, "gguf_variant": variant}
|
||||
if override.get("llama_extra_args") is not None:
|
||||
load_kwargs["llama_extra_args"] = override["llama_extra_args"]
|
||||
if override.get("max_seq_length") is not None:
|
||||
load_kwargs["max_seq_length"] = override["max_seq_length"]
|
||||
# Reuse the load impl so its dedup, tensor fallback, and threading
|
||||
# apply. Call the impl directly: we already hold the lifecycle gate
|
||||
# the /load route would otherwise take, so the route would deadlock.
|
||||
await _load_model_impl(
|
||||
LoadRequest(**load_kwargs),
|
||||
fastapi_request,
|
||||
current_subject,
|
||||
current_request_counted = True,
|
||||
)
|
||||
# Advertise the repo id (not the concrete load path) as the loaded
|
||||
# model's public id and override key for /v1/models and idle stash.
|
||||
get_llama_cpp_backend()._openai_advertised_id = override_id
|
||||
finally:
|
||||
_auto_switch_process_lock.release()
|
||||
finally:
|
||||
_note_switch_waiter(key, -1)
|
||||
return
|
||||
if len(last) == 3:
|
||||
target_id, variant, override_id = last
|
||||
else: # pre-3-tuple stash: fall back to the path as the override key
|
||||
target_id, variant = last
|
||||
override_id = target_id
|
||||
else:
|
||||
# load_path is a concrete local path (never the bare repo id), so /load
|
||||
# takes the local branch and cannot trigger a download. override_id is the
|
||||
# advertised repo id, the launch-override key and the public model id.
|
||||
target_id, variant, override_id = resolved
|
||||
backend = get_llama_cpp_backend()
|
||||
# A bare model id (no :VARIANT) is satisfied by any loaded quant of that
|
||||
# repo, so it never reloads a different local quant that already serves it.
|
||||
bare = ":" not in requested_model
|
||||
|
||||
def _already_serving() -> bool:
|
||||
# Match against both the concrete load path and the advertised repo id,
|
||||
# so a model loaded manually by repo id (identifier = repo id) and one
|
||||
# loaded by auto-switch (identifier = path, advertised = repo id) both
|
||||
# count as already serving rather than triggering a needless reswap.
|
||||
if not backend.is_loaded or not backend.model_identifier:
|
||||
return False
|
||||
loaded_keys = {backend.model_identifier.lower()}
|
||||
advertised = getattr(backend, "_openai_advertised_id", None)
|
||||
if advertised:
|
||||
loaded_keys.add(advertised.lower())
|
||||
if loaded_keys.isdisjoint({target_id.lower(), override_id.lower()}):
|
||||
return False
|
||||
if bare:
|
||||
return True
|
||||
if variant:
|
||||
loaded_variant = (getattr(backend, "hf_variant", None) or "").lower()
|
||||
return loaded_variant == variant.lower()
|
||||
return True
|
||||
|
||||
def _record_serving_alias() -> None:
|
||||
# When an advertised alias already resolves to the loaded model (e.g. a
|
||||
# model loaded by local path, requested by its repo/LM Studio id), record
|
||||
# the alias as the public id so /v1/models and responses report it (and
|
||||
# mark it loaded) instead of the path-derived basename. Resolver branch
|
||||
# only: the reload-stash override_id can be the bare path, not a repo id.
|
||||
if resolved is None or not override_id:
|
||||
return
|
||||
b = get_llama_cpp_backend()
|
||||
if getattr(b, "_openai_advertised_id", None) != override_id:
|
||||
b._openai_advertised_id = override_id
|
||||
|
||||
if _already_serving():
|
||||
_record_serving_alias()
|
||||
return
|
||||
# An image/audio request naming a different text-only GGUF would load it
|
||||
# here and only 400 below, evicting the working model. Reject before the
|
||||
# swap. Only the resolver branch (an explicit new target); the reload-stash
|
||||
# path just restores the model the request was already using. Both vision and
|
||||
# audio input come from a companion mmproj (a filesystem probe) -- run it off
|
||||
# the loop, like the resolver above.
|
||||
if (
|
||||
require_vision
|
||||
and resolved is not None
|
||||
and not await asyncio.to_thread(_target_is_vision, target_id)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = openai_error_body(
|
||||
"The requested model does not support the image or audio input in this request.",
|
||||
status = 400,
|
||||
code = "invalid_value",
|
||||
param = "model",
|
||||
),
|
||||
)
|
||||
key = _switch_key(override_id, variant)
|
||||
_note_switch_waiter(key, 1)
|
||||
try:
|
||||
async with _auto_switch_lock():
|
||||
# The asyncio lock is per loop; add a process-wide gate so a swap on
|
||||
# another loop in this process can't race the single slot.
|
||||
await _acquire_swap_gate()
|
||||
try:
|
||||
# Hold the keep-warm gate across the swap so no new inference can
|
||||
# start on the model while it is being torn down and replaced.
|
||||
async with inference_lifecycle_gate():
|
||||
if _already_serving():
|
||||
_record_serving_alias()
|
||||
return
|
||||
# Apply this model's saved launch flags so the swap honors the config.
|
||||
override = get_model_override(override_id)
|
||||
load_kwargs = {"model_path": target_id, "gguf_variant": variant}
|
||||
if override.get("llama_extra_args") is not None:
|
||||
load_kwargs["llama_extra_args"] = override["llama_extra_args"]
|
||||
if override.get("max_seq_length") is not None:
|
||||
load_kwargs["max_seq_length"] = override["max_seq_length"]
|
||||
# Reuse the load impl so its dedup, tensor fallback, and threading
|
||||
# apply. Call the impl directly: we already hold the lifecycle gate
|
||||
# the /load route would otherwise take, so the route would deadlock.
|
||||
await _load_model_impl(
|
||||
LoadRequest(**load_kwargs),
|
||||
fastapi_request,
|
||||
current_subject,
|
||||
current_request_counted = True,
|
||||
)
|
||||
# Advertise the repo id (not the concrete load path) as the loaded
|
||||
# model's public id and override key for /v1/models and idle stash.
|
||||
get_llama_cpp_backend()._openai_advertised_id = override_id
|
||||
finally:
|
||||
_auto_switch_process_lock.release()
|
||||
finally:
|
||||
_note_request_waiter(request_key, -1)
|
||||
_note_switch_waiter(key, -1)
|
||||
|
||||
|
||||
async def _auto_switch_from_request_body(request: Request, current_subject: str):
|
||||
|
|
|
|||
|
|
@ -100,7 +100,6 @@ def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder):
|
|||
# gate that auto-switch already owns, so it calls the impl directly).
|
||||
monkeypatch.setattr(inference_route, "_load_model_impl", recorder)
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
|
||||
|
||||
|
||||
def _run_hook(model = "some/model"):
|
||||
|
|
@ -1584,8 +1583,8 @@ def test_swap_waits_until_concurrent_request_finishes_resolving(monkeypatch):
|
|||
)
|
||||
monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin
|
||||
monkeypatch.setattr(kw, "_pending", 0)
|
||||
# The twin has only registered its raw requested model (not yet a target waiter).
|
||||
inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1)
|
||||
# The twin is still resolving, so it is counted in-flight but has not joined
|
||||
# the concrete target queue yet.
|
||||
|
||||
async def _drive():
|
||||
task = asyncio.create_task(
|
||||
|
|
@ -3166,7 +3165,6 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
|
|||
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
|
||||
monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load)
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
|
||||
|
||||
barrier = threading.Barrier(2)
|
||||
|
||||
|
|
|
|||
|
|
@ -381,7 +381,7 @@ def _http_json(
|
|||
|
||||
|
||||
# A server that WE auto-started (never one we merely found). Kept at module scope so
|
||||
# _run's finally and the atexit backstop can tear it down without threading a handle
|
||||
# failure paths and the atexit backstop can tear it down without threading a handle
|
||||
# through all six agent commands. Only one agent runs per process, so one slot is enough.
|
||||
_auto_served_server: Optional[subprocess.Popen] = None
|
||||
# Model download + load can be slow; give the auto-started server room before giving up.
|
||||
|
|
@ -759,7 +759,8 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess
|
|||
log_path.unlink(missing_ok = True)
|
||||
log = os.fdopen(os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb")
|
||||
# Own session/process group so a mid-session Ctrl+C (cancel a turn) doesn't reach the
|
||||
# server; we tear it down explicitly when the agent exits.
|
||||
# server. It stays available after a successful agent session and is torn down on
|
||||
# startup or launch failure.
|
||||
kwargs: dict = {"stdout": log, "stderr": subprocess.STDOUT, "stdin": subprocess.DEVNULL}
|
||||
if os.name == "nt":
|
||||
kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||||
|
|
@ -1172,7 +1173,8 @@ def _resolve_model(
|
|||
(
|
||||
m
|
||||
for m in models
|
||||
if any(
|
||||
if m.get("loaded") is not False
|
||||
and any(
|
||||
_model_id_matches(m.get("id"), w, allow_casefold = allow_casefold) for w in wanted
|
||||
)
|
||||
),
|
||||
|
|
@ -2228,7 +2230,7 @@ def codex(
|
|||
launch = launch,
|
||||
)
|
||||
# This preflight runs after _connect may have auto-started a server but before _run
|
||||
# installs its teardown finally, so tear the server down here if it rejects the model
|
||||
# takes over its lifecycle, so tear the server down here if it rejects the model
|
||||
# (e.g. a transformers-backend model) rather than leaving it on the atexit backstop.
|
||||
try:
|
||||
_require_gguf_for_codex(base, key, entry["id"])
|
||||
|
|
|
|||
|
|
@ -911,6 +911,35 @@ def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch):
|
|||
assert any(u.endswith("/api/inference/load") for _, u in calls)
|
||||
|
||||
|
||||
def test_resolve_model_does_not_attach_if_catalog_stays_unloaded(monkeypatch):
|
||||
def http_json(
|
||||
method,
|
||||
url,
|
||||
token,
|
||||
payload = None,
|
||||
timeout = 30,
|
||||
error = None,
|
||||
):
|
||||
if url.endswith("/v1/models"):
|
||||
return {
|
||||
"data": [
|
||||
{
|
||||
"id": "unsloth/Gemma-4-GGUF",
|
||||
"loaded": False,
|
||||
"context_length": 131072,
|
||||
}
|
||||
]
|
||||
}
|
||||
if url.endswith("/api/inference/load"):
|
||||
return {"status": "loaded", "model": "unsloth/Gemma-4-GGUF"}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", http_json)
|
||||
|
||||
with pytest.raises(typer.Exit):
|
||||
start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf")
|
||||
|
||||
|
||||
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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue