diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index e3ee73d7b7..c2112dce23 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1047,10 +1047,19 @@ class DiffusionBackend: # repo / filename validation raises before # ``effective_base`` is computed). ``locals().get`` # keeps the scrub a no-op in that case. + # Round 18 P2 #9: also scrub ``local_gguf_path``. The + # GGUF quant is loaded via + # ``transformer_cls.from_single_file(local_gguf_path)``, + # and diffusers / safetensors errors include the + # resolved absolute HF cache path + # (``/home/alice/.cache/huggingface/hub/.../flux.gguf``). + # Without this the cache path would leak into + # ``_last_error`` (and therefore status() / log lines). _locals = locals() exc_msg = _collapse_local(exc_msg, repo_id) exc_msg = _collapse_local(exc_msg, _locals.get("effective_base")) exc_msg = _collapse_local(exc_msg, _locals.get("gguf_filename")) + exc_msg = _collapse_local(exc_msg, _locals.get("local_gguf_path")) with self._lock: self._last_error = exc_msg # ``logger.exception`` would emit the raw exception @@ -1343,14 +1352,21 @@ def _release_chat_backend_for_diffusion() -> None: "Could not unload the existing GGUF chat model before " "loading a diffusion image model." ) from exc + # Round 18 P1 #4: also reject when ``loading_model_identifier`` + # is still set after the unload call. Without this, a GGUF + # download / startup that was already in flight before the + # diffusion handoff (and which never flipped is_active to + # True before the unload landed) keeps allocating into VRAM + # while diffusion proceeds, double-owning the GPU. if ( ok is False or getattr(backend, "is_loaded", False) or getattr(backend, "is_active", False) + or getattr(backend, "loading_model_identifier", None) ): raise RuntimeError( - "The existing GGUF chat model is still active after " - "unload; retry before loading a diffusion image model." + "The existing GGUF chat model is still active or loading " + "after unload; retry before loading a diffusion image model." ) # 2. Safetensors / HF chat backend (the InferenceOrchestrator that @@ -1447,11 +1463,19 @@ def _release_other_gpu_owners_for_diffusion() -> None: logger.debug("export module not importable: %s", exc) return + # Round 18 P1 #6: ``get_export_backend()`` raising used to be a + # silent ``return`` so direct ``DiffusionBackend.load_model`` + # callers could proceed toward GPU allocation without being able + # to verify export ownership. Fail closed instead, matching the + # route-level helper which already maps "Could not verify" / + # "Could not access" failures to HTTP 503. try: exp = get_export_backend() except Exception as exc: - logger.debug("export backend not available: %s", exc) - return + raise RuntimeError( + "Could not verify export status before loading a " + "diffusion image model." + ) from exc is_export_active_fn = getattr(exp, "is_export_active", None) if is_export_active_fn is not None: @@ -1480,14 +1504,23 @@ def _release_other_gpu_owners_for_diffusion() -> None: ) if getattr(exp, "current_checkpoint", None): + # Round 18 P1 #2: a wedged ``_shutdown_subprocess`` used to log + # at debug level and continue, so direct backend callers could + # allocate diffusion VRAM on top of an export checkpoint that + # still owned the GPU. Mirror the route-level helper and raise + # so the surrounding ``load_model`` bails out with a clean + # RuntimeError that the route layer maps to HTTP 503. try: logger.info("Shutting down idle export subprocess before diffusion load") exp._shutdown_subprocess() - exp.current_checkpoint = None - exp.is_vision = False - exp.is_peft = False except Exception as exc: - logger.debug("idle export shutdown failed: %s", exc) + raise RuntimeError( + "Could not unload the idle export checkpoint before " + "loading a diffusion image model." + ) from exc + exp.current_checkpoint = None + exp.is_vision = False + exp.is_peft = False # Note: active training is *not* stopped here. The route layer # (`_raise_if_training_active` in routes/inference.py) refuses diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b88175e270..079a90e231 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -392,7 +392,7 @@ async def _release_llama_for(workload: str) -> None: workload, ) try: - await asyncio.to_thread(llama.unload_model) + ok = await asyncio.to_thread(llama.unload_model) except Exception as exc: logger.warning("Failed to unload GGUF chat before %s load: %s", workload, exc) raise HTTPException( @@ -403,6 +403,28 @@ async def _release_llama_for(workload: str) -> None: ), ) from exc + # Round 18 P1 #1: previously only the raised-exception path was + # treated as failure. ``llama.unload_model()`` returning ``False`` + # (subprocess refused to terminate, IPC timeout) or leaving + # ``is_loaded`` / ``is_active`` / ``loading_model_identifier`` + # populated after the call meant the next workload could allocate + # while llama-server was still resident. Re-read the same three + # fields and fail closed if anything is still set so the caller + # retries instead of double-owning VRAM. + if ( + ok is False + or bool(getattr(llama, "is_loaded", False)) + or bool(getattr(llama, "is_active", False)) + or bool(getattr(llama, "loading_model_identifier", None)) + ): + raise HTTPException( + status_code = 503, + detail = ( + "The existing GGUF chat model is still active or loading " + f"after unload; retry before starting {workload}." + ), + ) + async def _release_safetensors_chat_for(workload: str) -> None: """Unload the safetensors / Unsloth chat backend (drains both @@ -599,13 +621,27 @@ async def _release_diffusion_for(workload: str) -> None: ), ) from exc - after = {} + # Round 18 P1 #5: a successful pre-check status() and a + # success-shaped unload result used to mask a post-unload + # status() failure (after = {}) and let the caller proceed + # without proof that diffusion released VRAM. Fail closed + # instead so training / chat / export retry rather than + # double-owning the GPU. try: after = diff_backend.status() - except Exception: - # status() failure here is unusual but should not mask the - # primary outcome. Fall back to assuming the unload finished. - pass + except Exception as exc: + logger.warning( + "Could not verify diffusion status after unload before %s: %s", + workload, + exc, + ) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not verify diffusion unload before starting " + f"{workload}. Try again." + ), + ) from exc if result is False or after.get("is_loaded") or after.get("is_loading"): raise HTTPException( status_code = 503, @@ -2022,14 +2058,18 @@ async def diffusion_load( # the request is refused with 409 instead of silently killing it. _raise_if_training_active("diffusion") _raise_if_export_active("diffusion") - # Round 17 P1 #3: drop the chat backends through the strict - # route-level helpers BEFORE the diffusion load. The backend's - # own ``_release_chat_backend_for_diffusion`` is now strict - # too (round 17 P1 #2), but doing it here keeps the public API - # path symmetric with training / export / chat handoffs that - # already use ``_release_chat_for``. - await _release_chat_for("diffusion") - await _release_export_for("diffusion") + # Round 18 P1 #3 + P1 #7: the route used to drop chat and idle + # export BEFORE ``backend.load_model`` ran its cheap validation + # (family inference, GGUF filename checks, gated-token failures, + # missing diffusers). A malformed image request would therefore + # unload the user's chat model and then return a 400 with nothing + # loaded; if export cleanup raised, chat had already been dropped. + # ``DiffusionBackend.load_model`` itself calls + # ``_release_other_gpu_owners_for_diffusion`` (strict idle-export + # shutdown after round 18 P1 #2) and + # ``_release_chat_backend_for_diffusion`` (strict GGUF + safetensors + # unload after round 17 P1 #2 + round 18 P1 #4), so the GPU is + # still freed before any allocation, just AFTER validation. backend = _get_diffusion_backend() try: status = await asyncio.get_event_loop().run_in_executor( diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index eac5837279..355394cc10 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -279,8 +279,13 @@ async def start_training( ) _raise_if_export_active("training") - await _release_chat_for("training") + # Round 18 P1 #8: release settled export FIRST so an export + # cleanup failure preserves the user's currently loaded chat + # model. The previous order (chat -> export) would drop chat + # and then refuse training when a wedged idle export raised, + # leaving the user with nothing loaded. await _release_export_for("training") + await _release_chat_for("training") # Also unload any loaded diffusion pipeline (Images page); it # holds the same GPU and would survive the inference shutdown.