From 2735cec36c19cc7859ee579d7ffaefed794b8d6b Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:03:00 -0300 Subject: [PATCH] Queue model switches behind active inference --- studio/backend/routes/inference.py | 94 +++++++------- .../backend/tests/test_openai_auto_switch.py | 118 +++++++++++++----- 2 files changed, 134 insertions(+), 78 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d3e588bb0b..90714d1804 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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) diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 1ee9ef36d3..87ffbf97f4 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -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