From 545a62174d8fe6418e0f088d08c75a3839cc02e8 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:46:05 -0300 Subject: [PATCH] Recheck sidecar reservation after inference drain --- studio/backend/routes/inference.py | 27 ++++++++++++------- .../backend/tests/test_openai_auto_switch.py | 13 +++++---- 2 files changed, 25 insertions(+), 15 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f87a2cc67f..7a5852b2ce 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4152,6 +4152,16 @@ def _maybe_unsupported_message(msg: str) -> str: return msg +def _raise_if_sidecar_swap_in_progress() -> None: + from utils.transformers_version import sidecar_swap_in_progress + + if sidecar_swap_in_progress(): + raise HTTPException( + status_code = 409, + detail = "A transformers installation is in progress. Retry when it completes.", + ) + + @router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, @@ -4172,20 +4182,12 @@ async def load_model( # install can reserve while this request queues on the gate, so the pre-gate # check alone is only a fast path. from core.inference.llama_keepwarm import inference_lifecycle_gate - from utils.transformers_version import sidecar_swap_in_progress - - _swap_409 = HTTPException( - status_code = 409, - detail = "A transformers installation is in progress. Retry when it completes.", - ) - if sidecar_swap_in_progress(): - raise _swap_409 + _raise_if_sidecar_swap_in_progress() # Hold the lifecycle gate across the load so idle auto-unload can't unload the # model mid-load. Auto-switch calls _load_model_impl directly since it already # holds this gate. async with inference_lifecycle_gate(): - if sidecar_swap_in_progress(): - raise _swap_409 + _raise_if_sidecar_swap_in_progress() return await _load_model_impl(request, fastapi_request, current_subject) @@ -4532,6 +4534,10 @@ async def _load_model_impl( # 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) + # The installer reserves its sidecar swap before waiting on this gate. + # It can do so while active inference drains, after the route-level + # checks above, so honor that reservation before replacing either backend. + _raise_if_sidecar_swap_in_progress() # Unload any active Unsloth model only after every hub conflict check. if unsloth_backend.active_model_name: @@ -4744,6 +4750,7 @@ async def _load_model_impl( # Unload any active GGUF model first llama_backend = get_llama_cpp_backend() await _wait_for_model_switch_idle(current_request_counted = current_request_counted) + _raise_if_sidecar_swap_in_progress() if llama_backend.is_loaded: logger.info("Unloading GGUF model before loading Unsloth model") llama_backend.unload_model() diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index d8db447c27..2dbc3d36a2 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -1490,20 +1490,23 @@ 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. +def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded(): + # Both replacement directions drain active inference, then recheck whether a + # sidecar install reserved the lifecycle gate during that 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:")) + gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait) unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait) standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1) + standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait) 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 + assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth + assert standard_wait < standard_sidecar_check < unload_gguf def _anthropic_payload(max_tokens = None):