From 1193c8144abc2dc4d5a0cb3a3ab3da89aca974fe Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 03:19:42 +0000 Subject: [PATCH] Fix/adjust diffusion: round 9 shared release helpers + export-active guard for PR #5754 Round 9 reviewer flagged a pile of handoff asymmetries: every GPU-owning lifecycle change (training, export, chat, images) needed its own bespoke unload sequence and they had drifted out of sync. Some skipped llama-server is_active; some missed safetensors loading_models; export and training did not check is_export_active. Backend handoff (P1) * routes/inference.py: new _release_chat_for / _release_export_for helpers. Both treat llama-server as held when is_loaded OR is_active, safetensors as held when active_model_name OR loading_models is non-empty, and export as held when current_checkpoint OR is_export_active. Both helpers run their unloads in worker threads so async routes do not block the event loop. * routes/training.py: replaces its bespoke inline llama / safe / export unload sequence with await _release_chat_for / _release_ export_for. * routes/export.py: same swap for the chat unload chain (export still does NOT call _release_export_for on itself). * routes/inference.py GGUF + standard chat-load paths: now use _release_export_for to drop a settled export, and the standard path's llama unload now also handles is_active=True (round 9 review #8). Backend reject-on-active export (P1 #5) * routes/inference.py: new _raise_if_export_active. Symmetric with _raise_if_training_active: a long-running export is refused with HTTP 409 instead of being silently killed when /images/load or /load arrives. Diffusion / images load and both chat-load paths call it. * core/inference/diffusion.py _release_other_gpu_owners_for_ diffusion: no longer tears down an in-flight export job. Only drops a SETTLED export checkpoint (current_checkpoint populated, is_export_active False). Round 9 review #5 -- the previous behavior could terminate an in-flight export and leave a partial output artifact. Token leak via logger.exception (P1 #6) * core/inference/diffusion.py: load-failure logging now uses logger.error(..., exc_msg) with the already-scrubbed string and exc_info=False. logger.exception() with the raw Exception would expose any hf_... token that diffusers / huggingface_hub embedded in the message or traceback locals, defeating the earlier in-flight scrub. Dependency pinning (P1 #11) * pyproject.toml: huggingfacenotorch optional extra now pins diffusers>=0.37.0. Previously the floor was only set in studio/backend/requirements/no-torch-runtime.txt, so a normal pip install would resolve diffusers 0.36.0 (no Flux2KleinPipeline) and the default curated FLUX.2 klein Images model would fail at runtime. Cache-delete exact match (P1 #14) * routes/models.py /delete-cached: llama.cpp and safetensors guards now match on exact repo-id (case-insensitive) instead of prefix. Diffusion guard already does this; the chat guards were the remaining surface where loading org/model-v2 blocked deleting org/model. --- pyproject.toml | 7 +- studio/backend/core/inference/diffusion.py | 38 ++--- studio/backend/routes/export.py | 35 +---- studio/backend/routes/inference.py | 174 ++++++++++++++++++++- studio/backend/routes/models.py | 12 +- studio/backend/routes/training.py | 54 ++----- 6 files changed, 219 insertions(+), 101 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index aef88d90f5..97f5caf553 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,7 +83,12 @@ huggingfacenotorch = [ "peft>=0.18.0,!=0.11.0", "huggingface_hub>=0.34.0", "hf_transfer", - "diffusers", + # Studio Images page depends on Flux2KleinPipeline / + # Flux2Pipeline, both shipped in diffusers>=0.37.0. Floor was + # missing here so a `pip install unsloth[huggingfacenotorch]` + # could resolve to 0.36.0 and fail at runtime when the default + # curated FLUX.2 klein model loads. + "diffusers>=0.37.0", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", "trl>=0.18.2,!=0.19.0,<=0.24.0", "sentence-transformers", diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index d0799081ac..8682f3eb9c 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -663,7 +663,15 @@ class DiffusionBackend: exc_msg = re.sub(r"hf_[A-Za-z0-9]{20,}", "", exc_msg) with self._lock: self._last_error = exc_msg - logger.exception("Diffusion load failed for %s", repo_id) + # ``logger.exception`` would emit the raw exception + # (including any unredacted ``hf_...`` token inside + # the message OR traceback locals on rich loggers). + # Use ``logger.error`` with the already-scrubbed + # message and exc_info=False so the bearer token + # cannot leak through structured logging sinks. + logger.error( + "Diffusion load failed for %s: %s", repo_id, exc_msg + ) raise RuntimeError( f"Failed to load diffusion model: {exc_msg}" ) from exc @@ -909,27 +917,21 @@ def _release_other_gpu_owners_for_diffusion() -> None: """Best-effort: shut down export subprocess + active training before a diffusion load. Both can hold multi-GB of VRAM and would OOM the diffusion allocation on consumer GPUs.""" - # Export subprocess. Shut down when EITHER a checkpoint is - # resident OR is_export_active() reports work in flight (a - # checkpoint load that has been kicked off but not yet completed - # the assignment to current_checkpoint). Either case can hold - # GPU memory that would OOM the diffusion allocation. + # Export resident checkpoint. We tear down a SETTLED export + # (current_checkpoint populated) because that means the export + # ran to completion and the user can re-load the result, but we + # do NOT touch is_export_active() here: an in-flight export job + # has unfinished partial output that termination would corrupt. + # The route layer rejects /images/load with HTTP 409 via + # _raise_if_export_active when is_export_active() is True, so + # we only reach this helper when export is either idle or + # holding a previously completed checkpoint. try: from core.export import get_export_backend # type: ignore exp = get_export_backend() - has_checkpoint = bool(getattr(exp, "current_checkpoint", None)) - is_active = False - try: - is_active = bool(exp.is_export_active()) - except Exception: - is_active = False - if has_checkpoint or is_active: - logger.info( - "Shutting down export subprocess (checkpoint=%s active=%s)", - has_checkpoint, - is_active, - ) + if getattr(exp, "current_checkpoint", None): + logger.info("Shutting down idle export subprocess before diffusion load") exp._shutdown_subprocess() exp.current_checkpoint = None exp.is_vision = False diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 4b0b95a809..aabcd1920b 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -107,35 +107,14 @@ async def load_checkpoint( ), ) - # Free GPU memory: shut down any running inference/training subprocesses - # before loading the export checkpoint (they'd compete for VRAM). - try: - from core.inference import get_inference_backend + # Free GPU memory: shut down any chat backend before loading + # the export checkpoint. Routes the unload through the shared + # 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 - inf = get_inference_backend() - if inf.active_model_name: - logger.info( - "Unloading inference model '%s' to free GPU memory for export", - inf.active_model_name, - ) - inf._shutdown_subprocess() - inf.active_model_name = None - inf.models.clear() - except Exception as e: - logger.warning("Could not unload inference model: %s", e) - - # Also unload any active GGUF llama-server (the inference unload - # above only covers the safetensors / Unsloth backend; GGUF - # chat runs as a separate subprocess). - try: - from routes.inference import get_llama_cpp_backend - - llama = get_llama_cpp_backend() - if getattr(llama, "is_loaded", False): - logger.info("Unloading GGUF chat model to free GPU memory for export") - llama.unload_model() - except Exception as e: - logger.debug("llama-server unload skipped for export: %s", e) + await _release_chat_for("export") # Also unload any active diffusion pipeline (Images page); it # competes for the same GPU and would survive the inference diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8b6918c91d..e43e185124 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -288,6 +288,145 @@ def _raise_if_training_active(workload: str) -> None: ) +def _raise_if_export_active(workload: str) -> None: + """Refuse a chat/diffusion load while an export job is active. + + Symmetric with ``_raise_if_training_active``: export is also a + long-running GPU-owning job a user does not want silently killed + by a chat / images load. Treat ``current_checkpoint is not None`` + and ``is_export_active() is True`` as 'export owns the GPU'. + + Same failure-mode split as the training variant: import failure + silently skips, runtime failure fails CLOSED with 503. + """ + try: + from core.export import get_export_backend # type: ignore + except Exception: + return + try: + exp = get_export_backend() + has_checkpoint = bool(getattr(exp, "current_checkpoint", None)) + try: + active = bool(exp.is_export_active()) + except Exception: + active = False + except Exception as exc: + logger.warning( + "Could not verify export status before %s load: %s", + workload, + exc, + ) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not verify export status before loading the " + f"{workload} model. Try again." + ), + ) from exc + if has_checkpoint or active: + raise HTTPException( + status_code = 409, + detail = ( + f"An export job is currently active. Stop the export " + f"job before loading a {workload} model." + ), + ) + + +async def _release_chat_for(workload: str) -> None: + """Shared 'release any GPU-owning chat backend' helper. + + Used by training / export / images / chat handoffs. Treats + llama-server as held when EITHER ``is_loaded`` or ``is_active`` + is true (the latter is mid-download / mid-startup). Treats the + safetensors backend as held when ``active_model_name`` is set + OR ``loading_models`` is non-empty (mid-download / mid-load). + Each unload runs in a worker thread because both backends' + unload paths can block for the full duration of a load. + """ + # GGUF chat (llama-server subprocess). + try: + llama = get_llama_cpp_backend() + is_loaded = bool(getattr(llama, "is_loaded", False)) + is_active = bool(getattr(llama, "is_active", False)) + if is_loaded or is_active: + logger.info( + "Unloading GGUF chat (loaded=%s active=%s) before %s load", + is_loaded, + is_active, + workload, + ) + await asyncio.to_thread(llama.unload_model) + except Exception as e: + logger.debug("llama-server unload skipped for %s: %s", workload, e) + + # Safetensors / Unsloth chat backend. + try: + from core.inference import get_inference_backend as _gib # type: ignore + + inf = _gib() + active_model_name = getattr(inf, "active_model_name", None) + loading_models = set(getattr(inf, "loading_models", set()) or set()) + if active_model_name: + logger.info( + "Unloading safetensors chat '%s' before %s load", + active_model_name, + workload, + ) + await asyncio.to_thread(inf.unload_model, active_model_name) + for loading in loading_models: + if loading == active_model_name: + continue + try: + logger.info( + "Unloading in-flight safetensors chat '%s' before %s load", + loading, + workload, + ) + await asyncio.to_thread(inf.unload_model, loading) + except Exception as inner: + logger.debug( + "loading safetensors unload skipped for %s: %s", + loading, + inner, + ) + except Exception as e: + logger.debug("safetensors unload skipped for %s: %s", workload, e) + + +async def _release_export_for(workload: str) -> None: + """Shared 'shut down export subprocess' helper. + + Treats ``current_checkpoint is not None`` or ``is_export_active()`` + as 'export owns the GPU'. Used by training / chat handoffs. + Diffusion does NOT call this -- it refuses with 409 via + ``_raise_if_export_active`` instead, because killing an in-flight + export would corrupt the user's exported model. + """ + try: + from core.export import get_export_backend # type: ignore + + exp = get_export_backend() + has_checkpoint = bool(getattr(exp, "current_checkpoint", None)) + try: + active = bool(exp.is_export_active()) + except Exception: + active = False + if has_checkpoint or active: + logger.info( + "Shutting down export (checkpoint=%s active=%s) for %s", + has_checkpoint, + active, + 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) + + def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: """Classify reasoning/tool capabilities via the GGUF classifier so flags match across backends. gpt-oss is overridden because Harmony @@ -787,7 +926,14 @@ async def load_model( # training is active. Diffusion and export paths refuse; # without this the GGUF chat load would start llama-server # while training still owned VRAM and double-spend it. + # Also refuse when an export job is in flight: same + # reasoning as diffusion (terminating a live export would + # corrupt the user's exported artifact). _raise_if_training_active("chat") + _raise_if_export_active("chat") + # Drop a settled export checkpoint that is just holding + # GPU memory but is not actively producing output. + await _release_export_for("GGUF chat") llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() @@ -993,17 +1139,29 @@ async def load_model( # ── Standard path: load via Unsloth/transformers ────────── # Symmetric lifecycle guard: refuse a chat load while training - # is active so we do not OOM both the training and inference - # jobs together. + # or an export is active so we do not OOM both jobs together + # and so we do not silently corrupt an in-flight export. _raise_if_training_active("chat") + _raise_if_export_active("chat") + # Drop a settled export checkpoint that is just holding GPU + # memory but is not actively producing output. + await _release_export_for("safetensors chat") backend = get_inference_backend() - # Unload any active GGUF model first + # Unload any active GGUF model first (handles both is_loaded + # and is_active=True so a mid-startup llama-server is also + # killed before we allocate safetensors weights). llama_backend = get_llama_cpp_backend() - if llama_backend.is_loaded: - logger.info("Unloading GGUF model before loading Unsloth model") - llama_backend.unload_model() + llama_loaded = bool(getattr(llama_backend, "is_loaded", False)) + llama_active = bool(getattr(llama_backend, "is_active", False)) + if llama_loaded or llama_active: + logger.info( + "Unloading GGUF model (loaded=%s active=%s) before Unsloth load", + llama_loaded, + llama_active, + ) + await asyncio.to_thread(llama_backend.unload_model) # 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 @@ -1722,7 +1880,11 @@ async def diffusion_load( # Refuse before the long download starts: silently stopping a # running training run to free VRAM was the previous behavior and # left the user with no model loaded plus a dead training job. + # Same logic for export: an export subprocess that is mid-flight + # cannot be safely terminated without corrupting the output, so + # the request is refused with 409 instead of silently killing it. _raise_if_training_active("diffusion") + _raise_if_export_active("diffusion") backend = _get_diffusion_backend() try: status = await asyncio.get_event_loop().run_in_executor( diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1f71379197..7a1b238065 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2681,7 +2681,11 @@ async def delete_cached_model( llama_backend = get_llama_cpp_backend() loaded_id = (llama_backend.model_identifier or "").lower() - wants = loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()) + # Exact match only (case-insensitive). Prefix match would + # block deleting unrelated ``org/model`` while + # ``org/model-v2`` is loaded -- same surface the diffusion + # guard fixed in round 5. + wants = loaded_id == repo_id.lower() if wants and ( llama_backend.is_loaded or getattr(llama_backend, "is_active", False) ): @@ -2707,16 +2711,18 @@ async def delete_cached_model( # Loading set holds model identifiers currently being # downloaded / instantiated; treat them like active loads # so a delete cannot race a partial mmap. + # Exact match only. Prefix matching would block deleting + # ``org/model`` while ``org/model-v2`` is loading. for loading_model in loading_models: ml = (loading_model or "").lower() - if ml == needle or ml.startswith(needle): + if ml == needle: raise HTTPException( status_code = 409, detail = "Cannot delete a model while it is loading", ) if inference_backend.active_model_name: active = inference_backend.active_model_name.lower() - if active == needle or active.startswith(needle): + if active == needle: raise HTTPException( status_code = 400, detail = "Unload the model before deleting", diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c7b33e025e..1bc39423b0 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -265,52 +265,16 @@ async def start_training( ) training_kwargs["trust_remote_code"] = True - # Free GPU memory: shut down any running inference/export subprocesses - # before training starts (they'd compete for VRAM otherwise) - try: - from core.inference import get_inference_backend + # Free GPU memory: shut down any chat backend (llama-server + # subprocess OR safetensors orchestrator) and any settled + # export checkpoint before training starts. The shared + # helpers handle the asymmetric cases (llama is_active, + # safetensors loading_models, export is_export_active) so + # this path stays in sync with /images/load and chat. + from routes.inference import _release_chat_for, _release_export_for - inf_backend = get_inference_backend() - if inf_backend.active_model_name: - logger.info( - "Unloading inference model '%s' to free GPU memory for training", - inf_backend.active_model_name, - ) - inf_backend._shutdown_subprocess() - inf_backend.active_model_name = None - inf_backend.models.clear() - except Exception as e: - logger.warning("Could not unload inference model: %s", e) - - # GGUF chat backend (llama-server subprocess). Without this, - # starting training while a GGUF model is loaded keeps the - # subprocess pinned to VRAM and OOMs the training job. Mirrors - # the symmetric handoffs in routes/inference.py and - # routes/export.py. - try: - from routes.inference import get_llama_cpp_backend - - llama_backend = get_llama_cpp_backend() - if getattr(llama_backend, "is_loaded", False): - logger.info("Unloading GGUF chat model to free GPU memory for training") - llama_backend.unload_model() - except Exception as e: - logger.warning("Could not unload GGUF chat model: %s", e) - - try: - from core.export import get_export_backend - - exp_backend = get_export_backend() - if exp_backend.current_checkpoint: - logger.info( - "Shutting down export subprocess to free GPU memory for training" - ) - exp_backend._shutdown_subprocess() - exp_backend.current_checkpoint = None - exp_backend.is_vision = False - exp_backend.is_peft = False - except Exception as e: - logger.warning("Could not shut down export subprocess: %s", e) + await _release_chat_for("training") + await _release_export_for("training") # Also unload any loaded diffusion pipeline (Images page); it # holds the same GPU and would survive the inference shutdown.