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.
This commit is contained in:
Daniel Han-Chen
parent
0ae20554dc
commit
1193c8144a
6 changed files with 221 additions and 103 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -663,7 +663,15 @@ class DiffusionBackend:
|
|||
exc_msg = re.sub(r"hf_[A-Za-z0-9]{20,}", "<redacted>", 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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue