ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix/adjust diffusion: round 18 P1+P2 batch for PR #5754

Browse source 
P1 #1: ``_release_llama_for()`` now verifies ``llama.unload_model``
did not return False AND that ``is_loaded`` / ``is_active`` /
``loading_model_identifier`` are all cleared after the call. The
previous version only treated raised exceptions as failure, so a
subprocess refusing to terminate or an in-flight GGUF download
let the next workload allocate on top.

P1 #2: ``DiffusionBackend._release_other_gpu_owners_for_diffusion``
now raises RuntimeError when ``exp._shutdown_subprocess`` fails on
a settled checkpoint. Direct backend callers used to log at debug
level and proceed toward diffusion allocation while the export
checkpoint still owned VRAM.

P1 #3 + P1 #7: ``/images/load`` no longer drops chat + idle export
before the cheap backend validation runs. ``DiffusionBackend.load_model``
already calls the strict ``_release_other_gpu_owners_for_diffusion``
and ``_release_chat_backend_for_diffusion`` helpers AFTER family
inference and GGUF filename checks pass, so the GPU is still
freed before allocation and a malformed payload no longer
silently unloads the user's chat / chat-export pair.

P1 #4: ``_release_chat_backend_for_diffusion`` now also rejects a
post-unload state where ``loading_model_identifier`` is still set,
matching the route-level ``_release_llama_for`` strictness. A GGUF
download mid-flight before the diffusion handoff used to slip
through and end up double-owning VRAM after diffusion allocated.

P1 #5: ``_release_diffusion_for`` no longer swallows a post-unload
``status()`` failure as ``after = {}``. Training / chat / export
handoffs need proof that the diffusion pipeline released VRAM;
the helper now raises HTTP 503 when the verification status call
itself raises, so the caller retries.

P1 #6: ``DiffusionBackend._release_other_gpu_owners_for_diffusion``
raises RuntimeError when ``get_export_backend()`` itself raises.
Direct backend callers used to silently ``return`` here and
proceed to GPU allocation without being able to verify export
ownership.

P1 #8: ``/training/start`` releases settled export BEFORE chat,
matching the chat-load helpers. If idle export shutdown fails the
user's chat model is preserved instead of being dropped for a
training run that never starts.

P2 #9: GGUF load-error scrubber also collapses ``local_gguf_path``,
the resolved HF cache path passed to
``transformer_cls.from_single_file()``. Without this an exception
like ``OSError: cannot load /home/alice/.cache/huggingface/.../flux.gguf``
would leak the operator's filesystem layout through ``last_error``
and ``/images/status``.

All 85 diffusion-relevant backend tests pass locally.
This commit is contained in:
Daniel Han-Chen 2026-05-25 08:43:59 +00:00
commit da27143520
3 changed files with 101 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

49
studio/backend/core/inference/diffusion.py
View file

@ -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

68
studio/backend/routes/inference.py
View file

@ -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(

7
studio/backend/routes/training.py
View file

@ -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.