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

Queue model switches behind active inference

Browse source
This commit is contained in:
oobabooga 2026-07-21 16:03:00 -03:00
commit 2735cec36c
2 changed files with 134 additions and 78 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -3406,9 +3406,9 @@ async def _acquire_swap_gate() -> None:
await asyncio.sleep(0.02)
# Counts in-flight auto-switch requests per (target, variant). The busy guard
# subtracts same-target waiters so concurrent requests for one model load once
# instead of each 409-ing the other.
# Counts auto-switch requests waiting to load each (target, variant). These
# requests are queued for a swap but are not generating, so the drain wait below
# excludes them from the active inference count.
_auto_switch_waiters: dict[tuple[str, str], int] = {}
_auto_switch_waiters_guard = threading.Lock()
@ -3426,15 +3426,36 @@ def _note_switch_waiter(key: tuple[str, str], delta: int) -> None:
_auto_switch_waiters.pop(key, None)
def _same_target_waiters(key: tuple[str, str]) -> int:
def _switch_waiter_count() -> int:
with _auto_switch_waiters_guard:
return _auto_switch_waiters.get(key, 0)
return sum(max(0, count) for count in _auto_switch_waiters.values())
# A second waiter map keyed by the raw requested model, registered before the
# (slow) resolve. The middleware counts a concurrent same-model request as
# in-flight before it resolves and joins _auto_switch_waiters, so without this
# the first request would see it as an unrelated request and 409.
async def _wait_for_model_switch_idle(*, current_request_counted: bool) -> None:
"""Wait until a model replacement cannot interrupt active inference.
The caller holds ``inference_lifecycle_gate``, which prevents new inference
from starting while existing requests drain. Auto-switch requests that have
resolved their targets are scheduler waiters, not active generations, so
exclude them to avoid a queue deadlock.
"""
from core.inference.llama_keepwarm import other_inference_request_count
while True:
queued_switches = _switch_waiter_count()
if current_request_counted and queued_switches > 0:
queued_switches -= 1
active_others = other_inference_request_count(
current_request_counted = current_request_counted,
include_pending = False,
)
if active_others <= queued_switches:
return
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()
@ -3452,11 +3473,6 @@ def _note_request_waiter(key: str, delta: int) -> None:
_auto_switch_request_waiters.pop(key, None)
def _same_request_waiters(key: str) -> int:
with _auto_switch_request_waiters_guard:
return _auto_switch_request_waiters.get(key, 0)
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
@ -3582,7 +3598,6 @@ async def _maybe_auto_switch_model(
from core.inference.local_model_resolver import resolve_local_gguf
from core.inference.llama_keepwarm import (
get_last_unloaded_model,
other_inference_request_count,
inference_lifecycle_gate,
)
@ -3603,9 +3618,7 @@ 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):
# the middleware already counts a concurrent same-model request as in-flight,
# so the busy guard must know it shares this target even while it resolves.
# 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:
@ -3718,31 +3731,6 @@ async def _maybe_auto_switch_model(
if _already_serving():
_record_serving_alias()
return
# Single slot: refuse a cross-model swap while another inference
# request is active rather than killing its response. Requests
# heading to this same target (by resolved id or raw name) are
# excluded, so concurrent requests for one model load once. A
# pending request is still in the middleware, not generating, so
# it is not counted here.
same_others = max(
_same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0
)
others = other_inference_request_count(
current_request_counted = True, include_pending = False
)
# Not gated on the GGUF being loaded: _load_model_impl also
# tears down an active Unsloth backend before loading a GGUF,
# so refuse whenever any other inference request is in flight.
if others > same_others:
raise HTTPException(
status_code = 409,
detail = openai_error_body(
"Cannot switch models while another inference request is in progress.",
status = 409,
code = "model_switch_busy",
param = "model",
),
)
# 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}
@ -3757,6 +3745,7 @@ async def _maybe_auto_switch_model(
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.
@ -4223,7 +4212,13 @@ async def load_model(
return await _load_model_impl(request, fastapi_request, current_subject)
async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str):
async def _load_model_impl(
request: LoadRequest,
fastapi_request: Request,
current_subject: str,
*,
current_request_counted: bool = False,
):
from core.inference.llama_cpp import LlamaServerNotFoundError
# A new load starts here; arm the progress throttle so this load's first
@ -4557,6 +4552,12 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
),
)
# Keep the resident model alive until every active generation has
# finished. The lifecycle gate held by the caller blocks new starts.
await _wait_for_model_switch_idle(
current_request_counted = current_request_counted
)
# Unload any active Unsloth model only after every hub conflict check.
if unsloth_backend.active_model_name:
logger.info(
@ -4767,6 +4768,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Unload any active GGUF model first
llama_backend = get_llama_cpp_backend()
await _wait_for_model_switch_idle(
current_request_counted = current_request_counted
)
if llama_backend.is_loaded:
logger.info("Unloading GGUF model before loading Unsloth model")
llama_backend.unload_model()
@ -7096,7 +7100,7 @@ async def openai_chat_completions(
if payload.provider_id or payload.provider_type:
# External provider: this request won't touch the local GGUF, so drop it
# from the keep-warm count or its in-flight stream would falsely block a
# concurrent local auto-switch with model_switch_busy.
# concurrent local model switch from proceeding.
from core.inference.llama_keepwarm import untrack_current_request
untrack_current_request(request.scope)

118
studio/backend/tests/test_openai_auto_switch.py
View file

@ -68,7 +68,13 @@ class _LoadRecorder:
request,
fastapi_request,
current_subject = None,
*,
current_request_counted = False,
):
# Mirror the production load boundary before recording any replacement.
await inference_route._wait_for_model_switch_idle(
current_request_counted = current_request_counted
)
self.calls.append(request)
if self.fail:
from fastapi import HTTPException
@ -1205,10 +1211,9 @@ def test_middleware_ignores_non_post(monkeypatch):
# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ──
def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
# A cross-model swap must 409 (not kill) while another inference request is in
# flight; the requesting call itself is excluded from the count.
from fastapi import HTTPException
def test_auto_switch_waits_for_another_inference_to_finish(monkeypatch):
# A cross-model swap queues while another request is generating, then loads
# after that request drains. The requesting call itself is excluded.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M")
@ -1222,10 +1227,20 @@ def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
)
monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one
monkeypatch.setattr(kw, "_pending", 0)
with pytest.raises(HTTPException) as exc:
_run_hook("org/B-GGUF:Q8_0")
assert exc.value.status_code == 409
assert rec.calls == []
async def _drive():
task = asyncio.create_task(
inference_route._maybe_auto_switch_model(
"org/B-GGUF:Q8_0", object(), "tester"
)
)
await asyncio.sleep(0.05)
assert rec.calls == []
kw._note_end() # the other generation finishes; this request remains counted
await asyncio.wait_for(task, timeout = 1)
asyncio.run(_drive())
assert len(rec.calls) == 1
def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch):
@ -1411,13 +1426,12 @@ def test_concurrent_same_target_requests_load_once(monkeypatch):
monkeypatch.setattr(kw, "_pending", 0)
inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1)
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1 # loads once, no 409
assert len(rec.calls) == 1
def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch):
# A concurrent request heading to a different target still blocks the swap: the
# same-target exclusion must not swallow a genuinely conflicting request.
from fastapi import HTTPException
def test_queued_different_target_does_not_deadlock_current_swap(monkeypatch):
# A concurrent request already queued for another target is not generating,
# so it must not prevent the current serialized swap from proceeding.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@ -1432,10 +1446,8 @@ def test_swap_still_refused_when_other_request_targets_different_model(monkeypat
monkeypatch.setattr(kw, "_inflight", 2)
monkeypatch.setattr(kw, "_pending", 0)
inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1)
with pytest.raises(HTTPException) as exc:
_run_hook("org/B-GGUF:Q8_0")
assert exc.value.status_code == 409
assert rec.calls == []
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1
def test_v1_models_advertises_repo_id_not_load_path(monkeypatch):
@ -1481,6 +1493,22 @@ def test_load_route_holds_lifecycle_gate(monkeypatch):
assert "_load_model_impl" in src
def test_model_replacements_wait_before_either_backend_is_unloaded():
# Both replacement directions share the drain wait. Exact-model reuse exits
# earlier, so an already-loaded model never waits on unrelated inference.
import inspect
src = inspect.getsource(inference_route._load_model_impl)
gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:"))
unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait)
standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1)
unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
already_loaded = src.index('status = "already_loaded"')
assert already_loaded < gguf_wait < unload_unsloth
assert standard_wait < unload_gguf
def _anthropic_payload(max_tokens = None):
from models.inference import AnthropicMessagesRequest, AnthropicMessage
return AnthropicMessagesRequest(
@ -1519,9 +1547,9 @@ def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch):
# ── review round 6: concurrency ordering, external untrack, unload gate, ids ──
def test_pending_same_target_request_does_not_force_409(monkeypatch):
def test_pending_same_target_request_does_not_block_swap(monkeypatch):
# A second same-target request blocked in the middleware (pending, not yet
# generating) must not make the first request 409: pending is excluded.
# generating) must not block the first request: pending is excluded.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@ -1536,13 +1564,13 @@ def test_pending_same_target_request_does_not_force_409(monkeypatch):
monkeypatch.setattr(kw, "_inflight", 1) # just the caller
monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1 # loads once, no 409
assert len(rec.calls) == 1
def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch):
def test_swap_waits_until_concurrent_request_finishes_resolving(monkeypatch):
# The real middleware counts a concurrent same-model request as in-flight
# before it resolves and registers a target waiter. The raw-request waiter,
# registered before resolve, must still exclude it so the first request loads.
# before it resolves and registers a target waiter. Treat it as active until
# its target is known, then recognize it as another queued switch request.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@ -1558,8 +1586,22 @@ def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypat
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)
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1 # loads once, no 409
async def _drive():
task = asyncio.create_task(
inference_route._maybe_auto_switch_model(
"org/B-GGUF:Q8_0", object(), "tester"
)
)
await asyncio.sleep(0.05)
assert rec.calls == []
inference_route._note_switch_waiter(
inference_route._switch_key("org/B-GGUF", "Q8_0"), 1
)
await asyncio.wait_for(task, timeout = 1)
asyncio.run(_drive())
assert len(rec.calls) == 1
def test_external_untrack_decrements_inflight_and_is_idempotent():
@ -1595,11 +1637,9 @@ def test_manual_unload_interrupts_even_while_inference_active(monkeypatch):
assert not backend.is_loaded # torn down despite the active request
def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
def test_auto_switch_waits_when_unsloth_stream_active(monkeypatch):
# The GGUF slot is empty but an Unsloth model is streaming (counted in-flight).
# _load_model_impl would unload it, so auto-switch must 409, not only when a
# GGUF is loaded.
from fastapi import HTTPException
# The replacement waits for it just as it does for a GGUF generation.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend(None) # no GGUF loaded
@ -1613,10 +1653,20 @@ def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
)
monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request
monkeypatch.setattr(kw, "_pending", 0)
with pytest.raises(HTTPException) as exc:
_run_hook("org/B-GGUF:Q8_0")
assert exc.value.status_code == 409
assert rec.calls == [] # the active Unsloth model is not torn down
async def _drive():
task = asyncio.create_task(
inference_route._maybe_auto_switch_model(
"org/B-GGUF:Q8_0", object(), "tester"
)
)
await asyncio.sleep(0.05)
assert rec.calls == []
kw._note_end()
await asyncio.wait_for(task, timeout = 1)
asyncio.run(_drive())
assert len(rec.calls) == 1
def test_public_model_id_prefers_advertised_over_path():
@ -3097,6 +3147,8 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
request,
fastapi_request,
current_subject = None,
*,
current_request_counted = False,
):
with slock:
state["cur"] += 1