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

Fix/adjust diffusion: round 10 fix export-active asymmetry + GGUF chat helper for PR #5754

Browse source 
Round 10 reviewers found the round 9 export helpers had a
destructive bug: _release_export_for treated is_export_active=True
as a shutdown condition, so any caller (training, chat, diffusion)
could terminate an in-flight export and corrupt the user's output.
Conversely _raise_if_export_active raised 409 on a settled
checkpoint, blocking idle cleanup.

Backend (P1)
  * routes/inference.py: split the export-active surface in two:
      _raise_if_export_active() now ONLY raises when
      is_export_active() is True. A settled current_checkpoint is
      treated as held GPU memory, not an active job.
      _release_export_for() now ONLY shuts down when
      current_checkpoint is set AND is_export_active() is False
      (i.e. a previously completed checkpoint just holding memory).
      An unknown / unverifiable is_export_active is treated as
      'might still be active' so the helper refuses to drop.
  * routes/training.py: now calls _raise_if_export_active before
    _release_chat_for / _release_export_for, mirroring the chat
    and diffusion paths. The previous code went straight to
    _release_export_for and would kill an in-flight export.
  * routes/inference.py: split _release_chat_for into
    _release_llama_for and _release_safetensors_chat_for so the
    GGUF chat-load path can release only the OTHER chat backend
    (round 10 review #4: the previous inline 'if active_model_name'
    check skipped loading_models and let an in-flight safetensors
    load race the new GGUF allocation).
  * routes/inference.py: _raise_if_export_active now fails CLOSED
    (503) when is_export_active() raises, not only when
    get_export_backend() raises. Round 10 review #7.

Dependencies (P1)
  * pyproject.toml huggingfacenotorch extra: pin gguf. The
    Studio Images default curated picker is GGUF-only and
    diffusers.GGUFQuantizationConfig + from_single_file require
    the standalone gguf package at runtime; missing it would 500
    on the first /api/inference/images/load with
    'gguf>=0.10.0 is required'.
This commit is contained in:
Daniel Han-Chen 2026-05-25 03:56:32 +00:00
commit 641cdcc13a
3 changed files with 99 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

6
pyproject.toml
View file

@ -89,6 +89,12 @@ huggingfacenotorch = [
# could resolve to 0.36.0 and fail at runtime when the default
# curated FLUX.2 klein model loads.
"diffusers>=0.37.0",
# diffusers.GGUFQuantizationConfig + from_single_file rely on
# the standalone gguf package at runtime. The Studio Images
# default curated picker is GGUF-only so this must install
# with the public huggingfacenotorch extra; missing it makes
# /api/inference/images/load 500 with "gguf>=0.10.0 required".
"gguf",
"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",

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

@ -289,15 +289,21 @@ 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.
"""Refuse a chat/diffusion/training load while an export job is
actively running.
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'.
long-running GPU-owning job a user does not want silently killed.
ONLY raises when ``is_export_active() is True`` (an export
subprocess is currently producing output). A settled
``current_checkpoint`` is NOT an active job -- it is just held
GPU memory and gets dropped by ``_release_export_for``.
Same failure-mode split as the training variant: import failure
silently skips, runtime failure fails CLOSED with 503.
Failure-mode split:
* ``core.export`` cannot be imported -> silently skip.
* ``get_export_backend()`` raises -> 503 fail closed.
* ``is_export_active()`` raises -> 503 fail closed (round 10
review #7).
"""
try:
from core.export import get_export_backend # type: ignore
@ -305,11 +311,21 @@ def _raise_if_export_active(workload: str) -> None:
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 backend 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
try:
active = bool(exp.is_export_active())
except Exception as exc:
logger.warning(
"Could not verify export status before %s load: %s",
@ -323,7 +339,7 @@ def _raise_if_export_active(workload: str) -> None:
f"{workload} model. Try again."
),
) from exc
if has_checkpoint or active:
if active:
raise HTTPException(
status_code = 409,
detail = (
@ -333,18 +349,11 @@ def _raise_if_export_active(workload: str) -> None:
)
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.
async def _release_llama_for(workload: str) -> None:
"""Unload the llama-server (GGUF) chat backend if it owns the
GPU. Treats ``is_loaded`` OR ``is_active`` as held (the latter
is mid-download / mid-startup, before health probes pass).
"""
# GGUF chat (llama-server subprocess).
try:
llama = get_llama_cpp_backend()
is_loaded = bool(getattr(llama, "is_loaded", False))
@ -360,7 +369,11 @@ async def _release_chat_for(workload: str) -> None:
except Exception as e:
logger.debug("llama-server unload skipped for %s: %s", workload, e)
# Safetensors / Unsloth chat backend.
async def _release_safetensors_chat_for(workload: str) -> None:
"""Unload the safetensors / Unsloth chat backend (drains both
``active_model_name`` and ``loading_models``) if it owns the GPU.
"""
try:
from core.inference import get_inference_backend as _gib # type: ignore
@ -394,14 +407,33 @@ async def _release_chat_for(workload: str) -> None:
logger.debug("safetensors unload skipped for %s: %s", workload, e)
async def _release_export_for(workload: str) -> None:
"""Shared 'shut down export subprocess' helper.
async def _release_chat_for(workload: str) -> None:
"""Shared 'release any GPU-owning chat backend' 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.
Used by training / export / diffusion handoffs (which need BOTH
chat backends gone). The GGUF chat-load path uses only
``_release_safetensors_chat_for`` because it is itself starting
llama-server -- we cannot release the backend we are about to
start. Conversely, the standard chat-load path releases only
the llama side.
"""
await _release_llama_for(workload)
await _release_safetensors_chat_for(workload)
async def _release_export_for(workload: str) -> None:
"""Shared 'drop a settled export checkpoint' helper.
ONLY shuts down the export subprocess when ``current_checkpoint``
is set AND ``is_export_active()`` is False -- i.e. a previously
completed load is just holding GPU memory. An in-flight export
job (``is_export_active()`` True) is NEVER touched here; the
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.
"""
try:
from core.export import get_export_backend # type: ignore
@ -411,12 +443,15 @@ async def _release_export_for(workload: str) -> None:
try:
active = bool(exp.is_export_active())
except Exception:
active = False
if has_checkpoint or active:
# Treat unverifiable export state 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:
logger.info(
"Shutting down export (checkpoint=%s active=%s) for %s",
"Shutting down idle export (checkpoint=%s) for %s",
has_checkpoint,
active,
workload,
)
await asyncio.to_thread(exp._shutdown_subprocess)
@ -938,12 +973,12 @@ async def load_model(
llama_backend = get_llama_cpp_backend()
unsloth_backend = get_inference_backend()
# Unload any active Unsloth model first to free VRAM
if unsloth_backend.active_model_name:
logger.info(
f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF"
)
unsloth_backend.unload_model(unsloth_backend.active_model_name)
# Unload any safetensors / Unsloth model first to free
# VRAM. Uses the shared helper so we also drain
# ``loading_models`` (round 10 review #4); the inline
# version only checked ``active_model_name`` and let an
# 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
@ -1149,19 +1184,10 @@ async def load_model(
backend = get_inference_backend()
# 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()
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 or mid-download llama-server first.
# Shared helper so this stays in sync with the GGUF path's
# 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

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

@ -265,14 +265,19 @@ async def start_training(
)
training_kwargs["trust_remote_code"] = True
# 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
# Symmetric lifecycle guard: refuse to start training while
# an export job is in flight. Round 10 review #1 -- the
# previous code went straight to ``_release_export_for``,
# which would terminate the in-flight export and corrupt
# the user's output artifact. Now we 409 first; the user
# stops the export and re-submits.
from routes.inference import (
_raise_if_export_active,
_release_chat_for,
_release_export_for,
)
_raise_if_export_active("training")
await _release_chat_for("training")
await _release_export_for("training")