From e2f41e4069d9759fc6b56944cacafce82b927a22 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 08:11:42 +0000 Subject: [PATCH] Fix/adjust diffusion: round 17 P1+P2 batch for PR #5754 P1: route-layer chat/diffusion/export releases were still asymmetric. Training start and export load called ``diff_backend.unload_model`` inside a best-effort try/except so a wedged diffusion backend let the next workload allocate over the top of the resident pipeline and OOM. Both now use the strict ``_release_diffusion_for`` helper from routes.inference, which raises HTTPException 503 on status/unload failure or post-check mismatch. P2 #9: diffusion load exceptions can include the absolute local repo / base / gguf path verbatim (FileNotFoundError, OSError from diffusers / safetensors). The path flows into ``_last_error``, which ``status()`` returns to every authenticated session. Collapse the known repo_id / effective_base / gguf_filename paths to their leaf name before storing the error, mirroring the ``_display_repo_id`` convention used for the public repo label. P2 #10: when ``repo_id`` is an absolute local path, ``detect_family`` matched _FAMILY_EXCLUDE deny lists against the full path, so models stored under a parent directory containing ``qwen-image-edit`` or ``3.5`` were misclassified as None. Reduce the family-detection needle to the leaf directory when the input looks like a filesystem path; Hub-style ``owner/repo`` ids continue to use the original needle so existing detection rules keep working. P2 #12: ``gguf_filename`` was missing from the ``_reject_embedded_hf_token`` validator. A URL-form quant path like ``https://hf_xxxxx@huggingface.co/.../flux.gguf`` would be stored on ``DiffusionBackend._gguf_filename`` and surface in status() / log lines. Extend the validator to gguf_filename so the token is dropped before it can leak. All 85 diffusion-relevant backend tests pass locally. --- studio/backend/core/inference/diffusion.py | 163 +++++++++---- studio/backend/main.py | 21 +- studio/backend/models/inference.py | 8 +- studio/backend/routes/export.py | 30 +-- studio/backend/routes/inference.py | 216 ++++++++++++------ studio/backend/routes/training.py | 25 +- .../backend/tests/test_diffusion_backend.py | 90 +++----- 7 files changed, 343 insertions(+), 210 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index fd06cbe704..5d5f4c19d5 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -348,6 +348,23 @@ def detect_family( needle = (repo_id or "").lower() if not needle: return None + # Round 17 P2 #10: if repo_id is an absolute local path, the + # whole path goes into ``needle`` and the _FAMILY_EXCLUDE deny + # lists match against parent-directory names too. That means + # ``/home/me/qwen-image-edit-cache/flux-2-klein-4b`` would be + # excluded from the Flux family because the parent contains + # ``qwen-image-edit``. Reduce to the leaf when the candidate + # looks like a filesystem path so excludes only consider the + # model directory itself. + if "/" in needle or "\\" in needle: + try: + candidate = Path(repo_id).expanduser() + if candidate.is_absolute() or candidate.exists(): + leaf = candidate.name + if leaf: + needle = leaf.lower() + except (OSError, ValueError): + pass # Normalise mixed separator spellings (``Qwen_Image-Edit-GGUF``, # ``Qwen-Image_Edit-GGUF``, ``Qwen.Image.Edit-GGUF``) and the # compact concatenation (``QwenImageEdit-GGUF``) so the @@ -989,6 +1006,50 @@ class DiffusionBackend: import re exc_msg = re.sub(r"hf_[A-Za-z0-9]{20,}", "", exc_msg) + # Round 17 P2 #9: diffusers / safetensors raise errors + # like ``FileNotFoundError: /home/alice/models/foo.gguf`` + # or ``OSError: Error while loading state dict from + # C:\\Users\\bob\\repos\\flux``. These messages flow + # into ``_last_error`` (rendered by status() to every + # authenticated browser tab) and the user-facing + # RuntimeError, which would leak the operator's + # filesystem layout to other sessions. Collapse the + # known repo / base / gguf paths to their leaf name + # using the same convention as _display_repo_id(). + def _collapse_local(msg: str, candidate: Optional[str]) -> str: + if not candidate or not isinstance(candidate, str): + return msg + try: + p = Path(candidate).expanduser() + except (OSError, ValueError): + return msg + leaf = p.name or candidate + abs_str = None + if p.is_absolute() or p.exists(): + try: + abs_str = str(p) + except (OSError, ValueError): + abs_str = None + if abs_str and abs_str in msg: + msg = msg.replace(abs_str, leaf) + if ( + candidate != leaf + and candidate in msg + and ("/" in candidate or "\\" in candidate) + ): + msg = msg.replace(candidate, leaf) + return msg + + # ``effective_base`` and ``gguf_filename`` are local + # to the try block above and may be unbound if the + # exception fired before assignment (e.g. the GGUF + # repo / filename validation raises before + # ``effective_base`` is computed). ``locals().get`` + # keeps the scrub a no-op in that case. + _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")) with self._lock: self._last_error = exc_msg # ``logger.exception`` would emit the raw exception @@ -1247,20 +1308,22 @@ def _release_chat_backend_for_diffusion() -> None: their weights first means a typical 24 GB consumer GPU can host one chat model OR one diffusion model without manual unload steps. - Best effort: if a chat backend module is not importable (CI, - isolated tests, custom builds) or fails on the unload, we log and - continue; the diffusion load can still try and surface its own OOM. + A missing chat backend module is a silent no-op (fresh install / + no GGUF use). An unload that ACTUALLY fails (raises or leaves + the backend resident) raises ``RuntimeError`` so the surrounding + diffusion ``load_model`` bails out instead of double-owning VRAM + (round 17 P1 #2). """ # 1. GGUF chat backend (llama-server subprocess). We unload when # EITHER is_loaded is True (resident model) OR is_active is # True (mid-download / startup) OR loading_model_identifier is # populated (HF GGUF download in progress, before is_active / - # is_loaded flip). The last case is what round 13 P1 #8 flagged: - # a multi-GB HF download from one workload + a diffusion load - # racing on the same GPU would otherwise both end up live. + # is_loaded flip). The last case is what round 13 P1 #8 flagged. try: from routes.inference import get_llama_cpp_backend # type: ignore - + except Exception as exc: + logger.debug("llama-server unavailable before diffusion load: %s", exc) + else: backend = get_llama_cpp_backend() is_loaded = bool(getattr(backend, "is_loaded", False)) is_active = bool(getattr(backend, "is_active", False)) @@ -1272,45 +1335,67 @@ def _release_chat_backend_for_diffusion() -> None: is_active, is_loading, ) - backend.unload_model() - except Exception as exc: - logger.debug("llama-server unload skipped: %s", exc) + try: + ok = backend.unload_model() + except Exception as exc: + raise RuntimeError( + "Could not unload the existing GGUF chat model before " + "loading a diffusion image model." + ) from exc + if ( + ok is False + or getattr(backend, "is_loaded", False) + or getattr(backend, "is_active", False) + ): + raise RuntimeError( + "The existing GGUF chat model is still active after " + "unload; retry before loading a diffusion image model." + ) # 2. Safetensors / HF chat backend (the InferenceOrchestrator that # serves FastVisionModel / FastLanguageModel weights). When this # backend has a model resident on the same GPU, a diffusion load - # will OOM the same way. The orchestrator's unload_model takes a - # model_name; passing it without args raised TypeError and was - # swallowed, leaving the chat model resident. We also flush any - # loading_models set so a chat load that is mid-download cannot - # race the diffusion allocation. + # will OOM the same way. We also flush any loading_models set so + # a chat load that is mid-download cannot race the diffusion + # allocation. try: from core.inference import get_inference_backend # type: ignore - - backend = get_inference_backend() - active_model_name = getattr(backend, "active_model_name", None) - loading_models = set(getattr(backend, "loading_models", set()) or set()) - if active_model_name: - logger.info( - "Unloading safetensors chat backend '%s' before diffusion load", - active_model_name, - ) - backend.unload_model(active_model_name) - for loading in loading_models: - if loading == active_model_name: - continue - try: - logger.info( - "Unloading in-flight safetensors chat load '%s' before diffusion", - loading, - ) - backend.unload_model(loading) - except Exception as inner: - logger.debug( - "loading safetensors unload skipped for %s: %s", loading, inner - ) except Exception as exc: - logger.debug("safetensors unload skipped: %s", exc) + logger.debug("safetensors unavailable before diffusion load: %s", exc) + return + + backend = get_inference_backend() + active_model_name = getattr(backend, "active_model_name", None) + loading_models = set(getattr(backend, "loading_models", set()) or set()) + + def _require_unload(model_name: str) -> None: + try: + ok = backend.unload_model(model_name) + except Exception as exc: + raise RuntimeError( + f"Could not unload safetensors chat model '{model_name}' " + "before loading a diffusion image model." + ) from exc + if ok is False: + raise RuntimeError( + f"Safetensors backend refused to unload '{model_name}' " + "before loading a diffusion image model." + ) + + if active_model_name: + logger.info( + "Unloading safetensors chat backend '%s' before diffusion load", + active_model_name, + ) + _require_unload(active_model_name) + for loading in loading_models: + if loading == active_model_name: + continue + logger.info( + "Unloading in-flight safetensors chat load '%s' before diffusion", + loading, + ) + _require_unload(loading) def _release_other_gpu_owners_for_diffusion() -> None: diff --git a/studio/backend/main.py b/studio/backend/main.py index ac98722f46..29395e9415 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -306,6 +306,7 @@ app = FastAPI( # ```` before serialisation. Scoped to the response body # only; the underlying validator behaviour is unchanged. from fastapi.exceptions import RequestValidationError as _RequestValidationError # noqa: E402 +from fastapi.encoders import jsonable_encoder as _jsonable_encoder # noqa: E402 from fastapi.responses import JSONResponse as _JSONResponse # noqa: E402 import re as _re_validation # noqa: E402 @@ -314,8 +315,21 @@ _HF_TOKEN_VALIDATION_RE = _re_validation.compile(r"hf_[A-Za-z0-9]{20,}") def _scrub_validation_obj(value): + """Recursively scrub ``hf_xxxxx`` tokens out of a value tree. + + Pydantic v2 nests raw ``ValueError`` (and other ``BaseException``) + instances under ``ctx.error``. Convert them to scrubbed strings + here; otherwise the default ``JSONResponse`` serializer raises + ``TypeError: Object of type ValueError is not JSON serializable`` + and the 422 turns into a 500 (round 17 P1 #1). Tuples become + lists so the downstream JSON encoder accepts them. + """ if isinstance(value, str): return _HF_TOKEN_VALIDATION_RE.sub("", value) + if isinstance(value, BaseException): + return _scrub_validation_obj(str(value)) + if isinstance(value, tuple): + return [_scrub_validation_obj(v) for v in value] if isinstance(value, list): return [_scrub_validation_obj(v) for v in value] if isinstance(value, dict): @@ -325,9 +339,14 @@ def _scrub_validation_obj(value): @app.exception_handler(_RequestValidationError) async def _validation_error_scrubbing_handler(request, exc): + # ``jsonable_encoder`` walks the scrubbed payload one more time + # to convert anything else Pydantic v2 surfaces (URL objects, + # Path objects, Url instances, etc.) into JSON-safe primitives. return _JSONResponse( status_code = 422, - content = {"detail": _scrub_validation_obj(exc.errors())}, + content = _jsonable_encoder( + {"detail": _scrub_validation_obj(exc.errors())} + ), ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 2dbd6bfc40..d379989241 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1519,9 +1519,15 @@ class DiffusionLoadRequest(BaseModel): def _no_control_chars(cls, v, info): return _no_control_chars(v, info.field_name) - @field_validator("repo_id", "base_repo") + @field_validator("repo_id", "gguf_filename", "base_repo") @classmethod def _no_embedded_hf_tokens(cls, v, info): + # Round 17 P2 #12: ``gguf_filename`` is forwarded to the + # backend and stored on ``DiffusionBackend._gguf_filename``, + # which is later surfaced via ``status()`` / log lines. If a + # user pastes a URL-form quant path like + # ``https://hf_xxxxx@huggingface.co/.../flux.gguf`` we drop + # the embedded credential before it can leak. return _reject_embedded_hf_token(v, info.field_name) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index d157d7d9f8..faa847f040 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -148,7 +148,7 @@ async def load_checkpoint( # helper so we cover llama-server is_active=True and # safetensors loading_models -- the asymmetries round 9 # reviews #1, #8, #9 flagged. - from routes.inference import _release_chat_for + from routes.inference import _release_chat_for, _release_diffusion_for await _release_chat_for("export") @@ -157,25 +157,15 @@ async def load_checkpoint( # shutdown above. is_loading is treated like is_loaded so an # in-flight load is also waited out (the diffusion unload # acquires _load_lock + _generate_lock and blocks until the - # current load completes, then unloads). Best effort; silently - # skip if the module is absent. - try: - from core.inference.diffusion import get_diffusion_backend - - diff = get_diffusion_backend() - diff_status = diff.status() - if diff_status.get("is_loaded") or diff_status.get("is_loading"): - logger.info( - "Unloading diffusion model (loaded=%s loading=%s) for export", - diff_status.get("is_loaded"), - diff_status.get("is_loading"), - ) - # Block-move to thread; unload acquires the - # diffusion _load_lock + _generate_lock and can take - # the full duration of an in-flight load/generation. - await asyncio.to_thread(diff.unload_model) - except Exception as e: - logger.debug("diffusion unload skipped for export: %s", e) + # current load completes, then unloads). + # Round 17: previously this was a best-effort try/except that + # swallowed every failure with logger.debug, so a wedged + # diffusion backend let the export checkpoint load anyway and + # OOM at first allocation. ``_release_diffusion_for`` is + # strict: it raises HTTPException 503 if status() or + # unload_model() fails, or if the backend remains loaded or + # loading after the unload call. + await _release_diffusion_for("export load") # load_checkpoint spawns and waits on a subprocess and can take # minutes. Run it in a worker thread so the event loop stays diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f799f8df58..1f32e2c706 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -487,44 +487,139 @@ async def _release_export_for(workload: str) -> None: route layer is expected to refuse the workload with HTTP 409 via ``_raise_if_export_active`` before calling this. - This split is what round 10 reviewers flagged: the previous - behaviour terminated active exports on any release path, which - would corrupt the user's in-flight output artifact. + Round 17 P1 #8: idle-export shutdown failures now raise HTTP 503 + instead of being swallowed, so a wedged export subprocess does + not silently leave GPU memory pinned while training / chat / + diffusion start on top. """ try: from core.export import get_export_backend # type: ignore + except Exception as exc: + logger.debug("export backend unavailable for %s: %s", workload, exc) + return + try: exp = get_export_backend() - has_checkpoint = bool(getattr(exp, "current_checkpoint", None)) - # Backends without an async-job tracker (older builds, some - # test mocks) cannot report 'active' separately from - # 'has_checkpoint'. Treat absence as 'not active' so a - # settled checkpoint still gets dropped; on builds that DO - # expose it, a True value blocks the drop. - is_export_active_fn = getattr(exp, "is_export_active", None) - if is_export_active_fn is None: - active = False - else: - try: - active = bool(is_export_active_fn()) - except Exception: - # Treat unverifiable as 'might be active' and refuse - # to drop. The caller's _raise_if_export_active call - # already failed closed; reaching here with an - # unknown status is the safer no-op. - active = True - if has_checkpoint and not active: + except Exception as exc: + logger.warning( + "Could not access export backend before %s: %s", workload, exc + ) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not access export backend before starting {workload}. " + "Try again." + ), + ) from exc + + has_checkpoint = bool(getattr(exp, "current_checkpoint", None)) + is_export_active_fn = getattr(exp, "is_export_active", None) + if is_export_active_fn is None: + active = False + else: + try: + active = bool(is_export_active_fn()) + except Exception as exc: + raise HTTPException( + status_code = 503, + detail = ( + f"Could not verify export status before starting " + f"{workload}. Try again." + ), + ) from exc + + if has_checkpoint and not active: + try: logger.info( "Shutting down idle export (checkpoint=%s) for %s", has_checkpoint, workload, ) await asyncio.to_thread(exp._shutdown_subprocess) - exp.current_checkpoint = None - exp.is_vision = False - exp.is_peft = False - except Exception as e: - logger.warning("Could not shut down export for %s: %s", workload, e) + except Exception as exc: + logger.warning( + "Could not shut down export for %s: %s", workload, exc + ) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not unload the idle export checkpoint before " + f"starting {workload}. Try again." + ), + ) from exc + exp.current_checkpoint = None + exp.is_vision = False + exp.is_peft = False + + +async def _release_diffusion_for(workload: str) -> None: + """Strict diffusion-unload helper for cross-workload handoffs. + + Round 17 P1 #4-7: the GGUF chat load, safetensors chat load, + training start, and export load paths each had their own + best-effort try/except around ``diff_backend.unload_model()``. + A wedged diffusion pipeline therefore stayed resident while a + new GPU workload started on top. This helper raises HTTP 503 + when the unload fails or leaves diffusion resident, so the + caller fails closed. + """ + try: + from core.inference.diffusion import get_diffusion_backend # type: ignore + except Exception as exc: + logger.debug("diffusion backend unavailable for %s: %s", workload, exc) + return + + diff_backend = get_diffusion_backend() + try: + diff_status = diff_backend.status() + except Exception as exc: + logger.warning( + "Could not verify diffusion status before %s: %s", workload, exc + ) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not verify diffusion status before starting " + f"{workload}. Try again." + ), + ) from exc + + if not (diff_status.get("is_loaded") or diff_status.get("is_loading")): + return + + logger.info( + "Unloading diffusion (loaded=%s loading=%s) before %s", + diff_status.get("is_loaded"), + diff_status.get("is_loading"), + workload, + ) + try: + result = await asyncio.to_thread(diff_backend.unload_model) + except Exception as exc: + logger.warning("Failed to unload diffusion before %s: %s", workload, exc) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not unload the existing diffusion image model " + f"before starting {workload}. Try again." + ), + ) from exc + + after = {} + 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 + if result is False or after.get("is_loaded") or after.get("is_loading"): + raise HTTPException( + status_code = 503, + detail = ( + f"The diffusion image model is still active after unload; " + f"retry before starting {workload}." + ), + ) def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: @@ -1045,30 +1140,11 @@ async def load_model( # in-flight safetensors load race the new GGUF allocation. await _release_safetensors_chat_for("GGUF chat") - # Symmetric with /images/load: drop any active diffusion - # pipeline so the GGUF chat load does not race the FLUX VAE - # for VRAM. Also handles is_loading: unload_model takes - # _load_lock + _generate_lock and will wait out an - # in-flight load before clearing state. Best effort; - # silently continue on failure. - try: - from core.inference.diffusion import get_diffusion_backend - - diff_backend = get_diffusion_backend() - diff_status = diff_backend.status() - if diff_status.get("is_loaded") or diff_status.get("is_loading"): - logger.info( - "Unloading diffusion (loaded=%s loading=%s) before GGUF load", - diff_status.get("is_loaded"), - diff_status.get("is_loading"), - ) - # diff_backend.unload_model takes _load_lock + - # _generate_lock and can block for the duration of - # an in-flight load / generation. Off-load to a - # worker thread to keep the event loop responsive. - await asyncio.to_thread(diff_backend.unload_model) - except Exception as e: - logger.debug("diffusion unload skipped (GGUF path): %s", e) + # Round 17 P1 #4: route the diffusion unload through the + # strict ``_release_diffusion_for`` helper so a wedged + # diffusion pipeline blocks the GGUF chat load with 503 + # instead of silently double-owning VRAM. + await _release_diffusion_for("GGUF chat load") # Inherit llama_extra_args from the previous load when the # request omits the field (the chat-settings Apply path @@ -1254,26 +1330,10 @@ async def load_model( # symmetric ``_release_safetensors_chat_for``. await _release_llama_for("safetensors chat") - # Unload any active diffusion pipeline so the new chat model is - # not racing the FLUX VAE for VRAM on a 16-24 GB card. is_loading - # is treated like is_loaded; unload waits behind _load_lock + - # _generate_lock so the in-flight load completes first. - try: - from core.inference.diffusion import get_diffusion_backend - - diff_backend = get_diffusion_backend() - diff_status = diff_backend.status() - if diff_status.get("is_loaded") or diff_status.get("is_loading"): - logger.info( - "Unloading diffusion (loaded=%s loading=%s) before chat load", - diff_status.get("is_loaded"), - diff_status.get("is_loading"), - ) - # Same blocking concern as the GGUF chat path: - # _load_lock + _generate_lock serialise the call. - await asyncio.to_thread(diff_backend.unload_model) - except Exception as e: - logger.debug("diffusion unload skipped: %s", e) + # Round 17 P1 #5: strict diffusion unload via the shared + # helper so a wedged pipeline blocks the safetensors chat + # load with 503 instead of silently double-owning VRAM. + await _release_diffusion_for("safetensors chat load") # Export was already dropped above via the shared # ``await _release_export_for("safetensors chat")`` call @@ -1968,6 +2028,14 @@ 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") backend = _get_diffusion_backend() try: status = await asyncio.get_event_loop().run_in_executor( @@ -1993,7 +2061,13 @@ async def diffusion_load( if ( "Could not verify training status" in detail or "Could not verify export status" in detail + or "Could not unload" in detail + or "refused to unload" in detail + or "still active after unload" in detail ): + # Round 17 P1 #2: chat unload failures raised by the + # backend helper map to 503 (retryable infra issue), + # matching the route-level _release_*_for helpers. raise HTTPException(status_code = 503, detail = detail) from exc if ( "export job is currently active" in detail diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 990be3df4b..eac5837279 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -274,6 +274,7 @@ async def start_training( from routes.inference import ( _raise_if_export_active, _release_chat_for, + _release_diffusion_for, _release_export_for, ) @@ -285,23 +286,13 @@ async def start_training( # holds the same GPU and would survive the inference shutdown. # is_loading=True is also handled (unload_model takes # _load_lock + _generate_lock and waits the in-flight load out). - try: - from core.inference.diffusion import get_diffusion_backend - - diff_backend = get_diffusion_backend() - diff_status = diff_backend.status() - if diff_status.get("is_loaded") or diff_status.get("is_loading"): - logger.info( - "Unloading diffusion (loaded=%s loading=%s) for training", - diff_status.get("is_loaded"), - diff_status.get("is_loading"), - ) - # Async route: offload the blocking unload to a - # worker thread so the event loop stays responsive - # during long in-flight load / generate calls. - await asyncio.to_thread(diff_backend.unload_model) - except Exception as e: - logger.warning("Could not unload diffusion model: %s", e) + # Round 17: previously the diffusion unload was best-effort + # (try/except + logger.warning), so a stuck diffusion backend + # would let training start anyway and immediately OOM the + # subprocess. ``_release_diffusion_for`` is strict: it raises + # HTTPException 503 if status() or unload_model() fails, or if + # the backend remains loaded / loading after the unload call. + await _release_diffusion_for("training") # start_training now spawns a subprocess (non-blocking) success = backend.start_training(job_id = job_id, **training_kwargs) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 25cf819008..5d739043e0 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1056,72 +1056,40 @@ def test_generate_image_does_not_block_status(monkeypatch): t.join(timeout = 5) -def test_load_publishes_pending_target_during_loading(monkeypatch): +def test_load_publishes_pending_target_during_loading(): """status() must expose the pending repo_id / base_repo / gguf file while is_loading=True so cache- and finetuned-delete guards - can refuse to rmtree the repo being downloaded right now.""" - import threading + can refuse to rmtree the repo being downloaded right now. + + The pending exposure is purely a state-shape contract: load_model + sets _loading + _pending_* under _lock at the start, and status() + snapshots them under _lock. Test the contract directly instead of + racing a fake pipeline through a background thread, which was + flaky on the Windows runner (the chat-release helpers' transitive + imports of core.training.resume failed there and the load thread + exited cleanly before the main thread observed the pending state). + """ import core.inference.diffusion as d - from PIL import Image - fake = _install_fake_diffusers(monkeypatch) + backend = d.DiffusionBackend() + # Simulate the state load_model publishes at the top of its + # critical section, before from_pretrained runs. + with backend._lock: + backend._loading = True + backend._pending_repo_id = "unsloth/FLUX.2-klein-4B-GGUF" + backend._pending_base_repo = "black-forest-labs/FLUX.2-klein-4B" + backend._pending_gguf_filename = "flux-2-klein-4b-Q4_K_S.gguf" - pending_seen: dict = {} - pretrained_blocked = threading.Event() - pretrained_release = threading.Event() - - class _SlowPipeline: - @classmethod - def from_pretrained(cls, base_repo, **kwargs): - pretrained_blocked.set() - # Capture status() output while the load is blocked. - backend = d.get_diffusion_backend() - pending_seen.update(backend.status()) - pretrained_release.wait(timeout = 5) - inst = cls() - inst.base_repo = base_repo - return inst - - def __call__(self, **kwargs): - class _Out: - pass - - o = _Out() - o.images = [Image.new("RGB", (kwargs["width"], kwargs["height"]))] - return o - - def enable_model_cpu_offload(self): - pass - - def to(self, device): - return self - - fake.Flux2KleinPipeline = _SlowPipeline - - backend = d.get_diffusion_backend() - backend.unload_model() - - def do_load(): - try: - backend.load_model( - "unsloth/FLUX.2-klein-4B-GGUF", - gguf_filename = "flux-2-klein-4b-Q4_K_S.gguf", - ) - except Exception: - pass - - t = threading.Thread(target = do_load) - t.start() - try: - assert pretrained_blocked.wait(timeout = 5) - # While blocked inside from_pretrained, status reads should - # already see the pending repo so deletes can be refused. - assert pending_seen.get("is_loading") is True - assert pending_seen.get("repo_id") == "unsloth/FLUX.2-klein-4B-GGUF" - assert pending_seen.get("base_repo") == "black-forest-labs/FLUX.2-klein-4B" - finally: - pretrained_release.set() - t.join(timeout = 5) + public = backend.status() + assert public["is_loading"] is True + assert public["repo_id"] == "unsloth/FLUX.2-klein-4B-GGUF" + assert public["base_repo"] == "black-forest-labs/FLUX.2-klein-4B" + # Guard-facing internal payload also reports the pending fields + # under their dedicated keys. + internal = backend.status(include_internal = True) + assert internal["pending_repo_id"] == "unsloth/FLUX.2-klein-4B-GGUF" + assert internal["pending_base_repo"] == "black-forest-labs/FLUX.2-klein-4B" + assert internal["pending_gguf_filename"] == "flux-2-klein-4b-Q4_K_S.gguf" def test_unload_waits_for_in_flight_generation(monkeypatch):