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

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

Browse source 
P1 #1: ``_release_safetensors_chat_for`` now re-reads
``active_model_name`` and ``loading_models`` after each unload AND
runs a final sweep against the initial owned-name set. The previous
helper trusted ``unload_model() -> True`` even though the
orchestrator can respond ``unloaded`` while still holding weights
or a concurrent ``load`` can repopulate the tracker between calls.
Per-name and global post-state mismatches now raise HTTP 503 so
the caller retries.

P1 #2: same post-state guarantee inside
``_release_chat_backend_for_diffusion`` for direct backend
callers. ``DiffusionBackend.load_model`` now raises RuntimeError
when the safetensors tracker still owns a previously-resident
name after the unload, matching the route-level helper. The route
layer's existing classifier maps the new wording to HTTP 503.

P1 #3: ``DiffusionBackend.load_model`` now preflights the full
diffusers repo (or explicit GGUF ``base_repo``) via
``hf_hub_download(filename="model_index.json")`` BEFORE the
chat / export unload runs. The GGUF path was already covered by
the existing ``hf_hub_download(gguf_filename)`` round-trip; the
full-repo path used to skip validation and let a typo / private /
gated repo only surface inside ``from_pretrained`` AFTER the
user's chat model was already dropped. Local paths are checked
structurally (must be a directory containing ``model_index.json``)
so we do not network-round-trip for an on-disk miss. Error
messages route through ``_display_repo_id`` so an absolute
filesystem path does not leak the operator's layout.

P1 #6: ``/api/inference/unload`` (the direct chat unload endpoint)
now treats ``unload_model() -> False`` AND a leftover state
(``is_loaded`` / ``is_active`` / ``loading_model_identifier`` for
GGUF, ``active_model_name`` / ``loading_models`` for safetensors)
as 503 instead of unconditionally responding
``status="unloaded"``. The UI used to show the model as gone while
the backend still owned VRAM.

P2 #7: extended the /images/load RuntimeError -> HTTPException
marker list with ``still active or loading after unload`` and
``still loading after unload``. Round 18 introduced these exact
phrasings on the backend side; without the extension a retryable
unload failure was returning HTTP 400 to the user instead of 503.

P2 #8: removed the unused ``unsloth_backend = get_inference_backend()``
eager construction in the GGUF chat-load branch. Eager
construction made the GGUF-only path needlessly fail or pay
startup cost when the safetensors backend was unavailable / lazy;
``_release_safetensors_chat_for`` already handles that case as a
no-op.

All 85 diffusion-relevant + 98 related backend tests pass locally.
This commit is contained in:
Daniel Han-Chen 2026-05-25 09:45:44 +00:00
commit c20ed25ec6
2 changed files with 191 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -201,6 +201,64 @@ def _expand_existing_local_path(value: str) -> str:
return value
def _preflight_full_diffusers_repo(repo: str, hf_token: Optional[str]) -> None:
"""Prove a full diffusers repo is accessible before any unloads.
Round 19 P1 #3: the GGUF path's ``hf_hub_download(gguf_filename)``
above this function fails fast on a bad / private / gated /
typo'd repo before we touch the chat backend. The full diffusers
path used to skip that round-trip and only discover the issue
inside ``from_pretrained`` AFTER the user's chat model was
already unloaded. Add the same one-file probe (``model_index.json``
is the diffusers manifest; every diffusers repo has one).
Local paths are checked structurally so we do not hit the network
for a missing on-disk directory; both branches raise RuntimeError
so the surrounding load_model bails out before the chat unload.
The display label is collapsed via ``_display_repo_id`` so an
absolute filesystem path in the error message does not leak the
operator's layout (see round 17 P2 #9).
"""
if not repo:
return
try:
local = Path(repo).expanduser()
except (OSError, ValueError):
local = None
if local is not None and local.exists():
if not local.is_dir():
raise RuntimeError(
f"Diffusion repo '{_display_repo_id(repo)}' is not a directory."
)
if not (local / "model_index.json").is_file():
raise RuntimeError(
f"Diffusion repo '{_display_repo_id(repo)}' is missing "
"model_index.json."
)
return
if (local is not None and local.is_absolute()) or repo.startswith("~"):
raise RuntimeError(
f"Local diffusion repo '{_display_repo_id(repo)}' does not exist."
)
try:
from huggingface_hub import hf_hub_download as _hf_hub_download
except Exception:
# diffusers is installed but huggingface_hub is missing -- let
# the downstream loader produce the canonical error.
return
try:
_hf_hub_download(
repo_id = repo,
filename = "model_index.json",
token = hf_token,
)
except Exception as exc:
raise RuntimeError(
f"Could not access diffusion repo '{_display_repo_id(repo)}' "
"before unloading the current model."
) from exc
def _display_repo_id(value: Any) -> Any:
"""Return a public-facing label for a repo_id / base_repo.
@ -835,6 +893,21 @@ class DiffusionBackend:
token = hf_token,
)
# Round 19 P1 #3: the GGUF branch above already
# proved repo + filename are accessible via
# ``hf_hub_download``. The full-diffusers path (no
# ``gguf_filename``) did NOT, so a typo / private /
# gated full repo only surfaced inside
# ``from_pretrained`` AFTER chat was unloaded. Probe
# ``effective_base`` for ``model_index.json`` here so
# the chat model is preserved on a bad full-repo
# request. Also probe when the GGUF caller supplied
# an explicit ``base_repo`` (the base companion is
# ALSO downloaded via from_pretrained further down
# and would OOM-then-fail past the unload).
if not gguf_filename or base_repo:
_preflight_full_diffusers_repo(effective_base, hf_token)
# All cheap failure points (bad gguf_filename, missing
# pipeline / transformer class, gated download token,
# transient Hub error on the GGUF download) have now
@ -1384,6 +1457,7 @@ def _release_chat_backend_for_diffusion() -> None:
backend = get_inference_backend()
active_model_name = getattr(backend, "active_model_name", None)
loading_models = set(getattr(backend, "loading_models", set()) or set())
owned_names = {name for name in ({active_model_name} | loading_models) if name}
def _require_unload(model_name: str) -> None:
try:
@ -1398,6 +1472,20 @@ def _release_chat_backend_for_diffusion() -> None:
f"Safetensors backend refused to unload '{model_name}' "
"before loading a diffusion image model."
)
# Round 19 P1 #2: per-name post-state check. ``unload_model``
# returning ``True`` does not guarantee the orchestrator
# actually dropped the weights; the worker may have responded
# while still holding them, or a concurrent ``load`` may have
# repopulated the tracker. Verify the specific name is gone
# so the surrounding diffusion load bails out instead of
# silently double-owning VRAM.
active_after = getattr(backend, "active_model_name", None)
loading_after = set(getattr(backend, "loading_models", set()) or set())
if active_after == model_name or model_name in loading_after:
raise RuntimeError(
f"Safetensors chat model '{model_name}' is still active "
"or loading after unload; retry before loading a diffusion image model."
)
if active_model_name:
logger.info(
@ -1414,6 +1502,18 @@ def _release_chat_backend_for_diffusion() -> None:
)
_require_unload(loading)
# Round 19 P1 #2: final sweep using the initial snapshot of
# owned names. Catches races where a name we did not explicitly
# unload (because it appeared in loading_models between the
# snapshot and the unload calls) is still owned after the loop.
remaining_loading = set(getattr(backend, "loading_models", set()) or set()) & owned_names
remaining_active = getattr(backend, "active_model_name", None)
if remaining_loading or (remaining_active in owned_names):
raise RuntimeError(
"The existing safetensors chat model is still active or loading "
"after unload; retry before loading a diffusion image model."
)
def _release_other_gpu_owners_for_diffusion() -> None:
"""Best-effort: shut down export subprocess + active training before

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

@ -464,9 +464,27 @@ async def _release_safetensors_chat_for(workload: str) -> None:
"Try again."
),
)
# Round 19 P1 #1: ``unload_model`` returning ``True`` does not
# by itself guarantee the orchestrator dropped the model. The
# worker may have responded ``unloaded`` while still holding
# weights, or a concurrent ``load`` from another tab may have
# repopulated ``loading_models`` between calls. Re-read the
# tracker fields and fail closed if this specific name is
# still active or loading so the caller retries.
remaining_loading = set(getattr(inf, "loading_models", set()) or set())
active_after = getattr(inf, "active_model_name", None)
if active_after == model_name or model_name in remaining_loading:
raise HTTPException(
status_code = 503,
detail = (
f"Safetensors chat model '{model_name}' is still active "
f"or loading after unload; retry before starting {workload}."
),
)
active_model_name = getattr(inf, "active_model_name", None)
loading_models = set(getattr(inf, "loading_models", set()) or set())
owned_names = {name for name in ({active_model_name} | loading_models) if name}
if active_model_name:
logger.info(
"Unloading safetensors chat '%s' before %s load",
@ -484,6 +502,21 @@ async def _release_safetensors_chat_for(workload: str) -> None:
)
await _unload_required(loading)
# Round 19 P1 #1: final sweep using the set of names that were
# initially present. Catches races where a model name we did not
# explicitly unload (because it appeared between the snapshot and
# the unload calls) is still in the owned set after the loop.
remaining_loading = set(getattr(inf, "loading_models", set()) or set()) & owned_names
remaining_active = getattr(inf, "active_model_name", None)
if remaining_loading or (remaining_active in owned_names):
raise HTTPException(
status_code = 503,
detail = (
"The existing safetensors chat model is still active or loading "
f"after unload; retry before starting {workload}."
),
)
async def _release_chat_for(workload: str) -> None:
"""Shared 'release any GPU-owning chat backend' helper.
@ -1161,7 +1194,14 @@ async def load_model(
await _release_export_for("GGUF chat")
llama_backend = get_llama_cpp_backend()
unsloth_backend = get_inference_backend()
# Round 19 P2 #8: previously also called
# ``unsloth_backend = get_inference_backend()`` here, but
# the binding was never used in the GGUF branch. Eager
# construction makes the GGUF-only path needlessly fail
# or pay startup cost when the safetensors backend is
# unavailable / lazy-initialised; the shared
# ``_release_safetensors_chat_for`` below already
# handles missing-backend cases as a no-op.
# Unload any safetensors / Unsloth model first to free
# VRAM. Uses the shared helper so we also drain
@ -1632,16 +1672,57 @@ async def unload_model(
)
or not llama_backend.is_loaded
):
llama_backend.unload_model()
# Round 19 P1 #6: previously this called
# ``llama_backend.unload_model()`` and unconditionally
# returned ``status="unloaded"`` even when the subprocess
# refused to terminate or IPC timed out. The frontend then
# showed the model as unloaded while llama-server was
# still resident. Treat ``False`` / leftover state as a
# 503 so the user retries.
ok = await asyncio.to_thread(llama_backend.unload_model)
if (
ok is False
or getattr(llama_backend, "is_loaded", False)
or getattr(llama_backend, "is_active", False)
or getattr(llama_backend, "loading_model_identifier", None)
):
raise HTTPException(
status_code = 503,
detail = (
"The GGUF model is still active or loading after unload. "
"Try again."
),
)
logger.info(f"Unloaded GGUF model: {request.model_path}")
return UnloadResponse(status = "unloaded", model = request.model_path)
# Otherwise, unload from Unsloth backend
backend = get_inference_backend()
backend.unload_model(request.model_path)
# Round 19 P1 #6: same fail-closed treatment for safetensors.
# ``unload_model`` returning ``False`` or leaving
# ``active_model_name`` / ``loading_models`` populated for the
# requested name must surface to the client so the UI reflects
# the real state.
ok = await asyncio.to_thread(backend.unload_model, request.model_path)
active_after = getattr(backend, "active_model_name", None)
loading_after = set(getattr(backend, "loading_models", set()) or set())
if (
ok is False
or active_after == request.model_path
or request.model_path in loading_after
):
raise HTTPException(
status_code = 503,
detail = (
"The safetensors model is still active or loading after "
"unload. Try again."
),
)
logger.info(f"Unloaded model: {request.model_path}")
return UnloadResponse(status = "unloaded", model = request.model_path)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error unloading model: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = f"Failed to unload model: {str(e)}")
@ -2098,6 +2179,13 @@ async def diffusion_load(
or "Could not unload" in detail
or "refused to unload" in detail
or "still active after unload" in detail
# Round 19 P2 #7: round 18 introduced new RuntimeError
# phrasings (``still active or loading after unload``)
# that the original marker list did not cover, so a
# retryable chat-unload failure was returning HTTP 400
# to the user instead of 503. Match both wordings.
or "still active or loading after unload" in detail
or "still loading after unload" in detail
):
# Round 17 P1 #2: chat unload failures raised by the
# backend helper map to 503 (retryable infra issue),