Fix/adjust diffusion: round 17 P1+P2 batch for PR #5754
P1: route-layer chat/diffusion/export releases were still asymmetric. Training start and export load called ``diff_backend.unload_model`` inside a best-effort try/except so a wedged diffusion backend let the next workload allocate over the top of the resident pipeline and OOM. Both now use the strict ``_release_diffusion_for`` helper from routes.inference, which raises HTTPException 503 on status/unload failure or post-check mismatch. P2 #9: diffusion load exceptions can include the absolute local repo / base / gguf path verbatim (FileNotFoundError, OSError from diffusers / safetensors). The path flows into ``_last_error``, which ``status()`` returns to every authenticated session. Collapse the known repo_id / effective_base / gguf_filename paths to their leaf name before storing the error, mirroring the ``_display_repo_id`` convention used for the public repo label. P2 #10: when ``repo_id`` is an absolute local path, ``detect_family`` matched _FAMILY_EXCLUDE deny lists against the full path, so models stored under a parent directory containing ``qwen-image-edit`` or ``3.5`` were misclassified as None. Reduce the family-detection needle to the leaf directory when the input looks like a filesystem path; Hub-style ``owner/repo`` ids continue to use the original needle so existing detection rules keep working. P2 #12: ``gguf_filename`` was missing from the ``_reject_embedded_hf_token`` validator. A URL-form quant path like ``https://hf_xxxxx@huggingface.co/.../flux.gguf`` would be stored on ``DiffusionBackend._gguf_filename`` and surface in status() / log lines. Extend the validator to gguf_filename so the token is dropped before it can leak. All 85 diffusion-relevant backend tests pass locally.
This commit is contained in:
Daniel Han-Chen
parent
6ac67571dd
commit
e2f41e4069
7 changed files with 342 additions and 209 deletions
|
|
@ -348,6 +348,23 @@ def detect_family(
|
|||
needle = (repo_id or "").lower()
|
||||
if not needle:
|
||||
return None
|
||||
# Round 17 P2 #10: if repo_id is an absolute local path, the
|
||||
# whole path goes into ``needle`` and the _FAMILY_EXCLUDE deny
|
||||
# lists match against parent-directory names too. That means
|
||||
# ``/home/me/qwen-image-edit-cache/flux-2-klein-4b`` would be
|
||||
# excluded from the Flux family because the parent contains
|
||||
# ``qwen-image-edit``. Reduce to the leaf when the candidate
|
||||
# looks like a filesystem path so excludes only consider the
|
||||
# model directory itself.
|
||||
if "/" in needle or "\\" in needle:
|
||||
try:
|
||||
candidate = Path(repo_id).expanduser()
|
||||
if candidate.is_absolute() or candidate.exists():
|
||||
leaf = candidate.name
|
||||
if leaf:
|
||||
needle = leaf.lower()
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
# Normalise mixed separator spellings (``Qwen_Image-Edit-GGUF``,
|
||||
# ``Qwen-Image_Edit-GGUF``, ``Qwen.Image.Edit-GGUF``) and the
|
||||
# compact concatenation (``QwenImageEdit-GGUF``) so the
|
||||
|
|
@ -989,6 +1006,50 @@ class DiffusionBackend:
|
|||
import re
|
||||
|
||||
exc_msg = re.sub(r"hf_[A-Za-z0-9]{20,}", "<redacted>", exc_msg)
|
||||
# Round 17 P2 #9: diffusers / safetensors raise errors
|
||||
# like ``FileNotFoundError: /home/alice/models/foo.gguf``
|
||||
# or ``OSError: Error while loading state dict from
|
||||
# C:\\Users\\bob\\repos\\flux``. These messages flow
|
||||
# into ``_last_error`` (rendered by status() to every
|
||||
# authenticated browser tab) and the user-facing
|
||||
# RuntimeError, which would leak the operator's
|
||||
# filesystem layout to other sessions. Collapse the
|
||||
# known repo / base / gguf paths to their leaf name
|
||||
# using the same convention as _display_repo_id().
|
||||
def _collapse_local(msg: str, candidate: Optional[str]) -> str:
|
||||
if not candidate or not isinstance(candidate, str):
|
||||
return msg
|
||||
try:
|
||||
p = Path(candidate).expanduser()
|
||||
except (OSError, ValueError):
|
||||
return msg
|
||||
leaf = p.name or candidate
|
||||
abs_str = None
|
||||
if p.is_absolute() or p.exists():
|
||||
try:
|
||||
abs_str = str(p)
|
||||
except (OSError, ValueError):
|
||||
abs_str = None
|
||||
if abs_str and abs_str in msg:
|
||||
msg = msg.replace(abs_str, leaf)
|
||||
if (
|
||||
candidate != leaf
|
||||
and candidate in msg
|
||||
and ("/" in candidate or "\\" in candidate)
|
||||
):
|
||||
msg = msg.replace(candidate, leaf)
|
||||
return msg
|
||||
|
||||
# ``effective_base`` and ``gguf_filename`` are local
|
||||
# to the try block above and may be unbound if the
|
||||
# exception fired before assignment (e.g. the GGUF
|
||||
# repo / filename validation raises before
|
||||
# ``effective_base`` is computed). ``locals().get``
|
||||
# keeps the scrub a no-op in that case.
|
||||
_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"))
|
||||
with self._lock:
|
||||
self._last_error = exc_msg
|
||||
# ``logger.exception`` would emit the raw exception
|
||||
|
|
@ -1247,20 +1308,22 @@ def _release_chat_backend_for_diffusion() -> None:
|
|||
their weights first means a typical 24 GB consumer GPU can host
|
||||
one chat model OR one diffusion model without manual unload steps.
|
||||
|
||||
Best effort: if a chat backend module is not importable (CI,
|
||||
isolated tests, custom builds) or fails on the unload, we log and
|
||||
continue; the diffusion load can still try and surface its own OOM.
|
||||
A missing chat backend module is a silent no-op (fresh install /
|
||||
no GGUF use). An unload that ACTUALLY fails (raises or leaves
|
||||
the backend resident) raises ``RuntimeError`` so the surrounding
|
||||
diffusion ``load_model`` bails out instead of double-owning VRAM
|
||||
(round 17 P1 #2).
|
||||
"""
|
||||
# 1. GGUF chat backend (llama-server subprocess). We unload when
|
||||
# EITHER is_loaded is True (resident model) OR is_active is
|
||||
# True (mid-download / startup) OR loading_model_identifier is
|
||||
# populated (HF GGUF download in progress, before is_active /
|
||||
# is_loaded flip). The last case is what round 13 P1 #8 flagged:
|
||||
# a multi-GB HF download from one workload + a diffusion load
|
||||
# racing on the same GPU would otherwise both end up live.
|
||||
# is_loaded flip). The last case is what round 13 P1 #8 flagged.
|
||||
try:
|
||||
from routes.inference import get_llama_cpp_backend # type: ignore
|
||||
|
||||
except Exception as exc:
|
||||
logger.debug("llama-server unavailable before diffusion load: %s", exc)
|
||||
else:
|
||||
backend = get_llama_cpp_backend()
|
||||
is_loaded = bool(getattr(backend, "is_loaded", False))
|
||||
is_active = bool(getattr(backend, "is_active", False))
|
||||
|
|
@ -1272,45 +1335,67 @@ def _release_chat_backend_for_diffusion() -> None:
|
|||
is_active,
|
||||
is_loading,
|
||||
)
|
||||
backend.unload_model()
|
||||
except Exception as exc:
|
||||
logger.debug("llama-server unload skipped: %s", exc)
|
||||
try:
|
||||
ok = backend.unload_model()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"Could not unload the existing GGUF chat model before "
|
||||
"loading a diffusion image model."
|
||||
) from exc
|
||||
if (
|
||||
ok is False
|
||||
or getattr(backend, "is_loaded", False)
|
||||
or getattr(backend, "is_active", False)
|
||||
):
|
||||
raise RuntimeError(
|
||||
"The existing GGUF chat model is still active after "
|
||||
"unload; retry before loading a diffusion image model."
|
||||
)
|
||||
|
||||
# 2. Safetensors / HF chat backend (the InferenceOrchestrator that
|
||||
# serves FastVisionModel / FastLanguageModel weights). When this
|
||||
# backend has a model resident on the same GPU, a diffusion load
|
||||
# will OOM the same way. The orchestrator's unload_model takes a
|
||||
# model_name; passing it without args raised TypeError and was
|
||||
# swallowed, leaving the chat model resident. We also flush any
|
||||
# loading_models set so a chat load that is mid-download cannot
|
||||
# race the diffusion allocation.
|
||||
# will OOM the same way. We also flush any loading_models set so
|
||||
# a chat load that is mid-download cannot race the diffusion
|
||||
# allocation.
|
||||
try:
|
||||
from core.inference import get_inference_backend # type: ignore
|
||||
|
||||
backend = get_inference_backend()
|
||||
active_model_name = getattr(backend, "active_model_name", None)
|
||||
loading_models = set(getattr(backend, "loading_models", set()) or set())
|
||||
if active_model_name:
|
||||
logger.info(
|
||||
"Unloading safetensors chat backend '%s' before diffusion load",
|
||||
active_model_name,
|
||||
)
|
||||
backend.unload_model(active_model_name)
|
||||
for loading in loading_models:
|
||||
if loading == active_model_name:
|
||||
continue
|
||||
try:
|
||||
logger.info(
|
||||
"Unloading in-flight safetensors chat load '%s' before diffusion",
|
||||
loading,
|
||||
)
|
||||
backend.unload_model(loading)
|
||||
except Exception as inner:
|
||||
logger.debug(
|
||||
"loading safetensors unload skipped for %s: %s", loading, inner
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug("safetensors unload skipped: %s", exc)
|
||||
logger.debug("safetensors unavailable before diffusion load: %s", exc)
|
||||
return
|
||||
|
||||
backend = get_inference_backend()
|
||||
active_model_name = getattr(backend, "active_model_name", None)
|
||||
loading_models = set(getattr(backend, "loading_models", set()) or set())
|
||||
|
||||
def _require_unload(model_name: str) -> None:
|
||||
try:
|
||||
ok = backend.unload_model(model_name)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
f"Could not unload safetensors chat model '{model_name}' "
|
||||
"before loading a diffusion image model."
|
||||
) from exc
|
||||
if ok is False:
|
||||
raise RuntimeError(
|
||||
f"Safetensors backend refused to unload '{model_name}' "
|
||||
"before loading a diffusion image model."
|
||||
)
|
||||
|
||||
if active_model_name:
|
||||
logger.info(
|
||||
"Unloading safetensors chat backend '%s' before diffusion load",
|
||||
active_model_name,
|
||||
)
|
||||
_require_unload(active_model_name)
|
||||
for loading in loading_models:
|
||||
if loading == active_model_name:
|
||||
continue
|
||||
logger.info(
|
||||
"Unloading in-flight safetensors chat load '%s' before diffusion",
|
||||
loading,
|
||||
)
|
||||
_require_unload(loading)
|
||||
|
||||
|
||||
def _release_other_gpu_owners_for_diffusion() -> None:
|
||||
|
|
|
|||
|
|
@ -306,6 +306,7 @@ app = FastAPI(
|
|||
# ``<redacted>`` before serialisation. Scoped to the response body
|
||||
# only; the underlying validator behaviour is unchanged.
|
||||
from fastapi.exceptions import RequestValidationError as _RequestValidationError # noqa: E402
|
||||
from fastapi.encoders import jsonable_encoder as _jsonable_encoder # noqa: E402
|
||||
from fastapi.responses import JSONResponse as _JSONResponse # noqa: E402
|
||||
import re as _re_validation # noqa: E402
|
||||
|
||||
|
|
@ -314,8 +315,21 @@ _HF_TOKEN_VALIDATION_RE = _re_validation.compile(r"hf_[A-Za-z0-9]{20,}")
|
|||
|
||||
|
||||
def _scrub_validation_obj(value):
|
||||
"""Recursively scrub ``hf_xxxxx`` tokens out of a value tree.
|
||||
|
||||
Pydantic v2 nests raw ``ValueError`` (and other ``BaseException``)
|
||||
instances under ``ctx.error``. Convert them to scrubbed strings
|
||||
here; otherwise the default ``JSONResponse`` serializer raises
|
||||
``TypeError: Object of type ValueError is not JSON serializable``
|
||||
and the 422 turns into a 500 (round 17 P1 #1). Tuples become
|
||||
lists so the downstream JSON encoder accepts them.
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
return _HF_TOKEN_VALIDATION_RE.sub("<redacted>", value)
|
||||
if isinstance(value, BaseException):
|
||||
return _scrub_validation_obj(str(value))
|
||||
if isinstance(value, tuple):
|
||||
return [_scrub_validation_obj(v) for v in value]
|
||||
if isinstance(value, list):
|
||||
return [_scrub_validation_obj(v) for v in value]
|
||||
if isinstance(value, dict):
|
||||
|
|
@ -325,9 +339,14 @@ def _scrub_validation_obj(value):
|
|||
|
||||
@app.exception_handler(_RequestValidationError)
|
||||
async def _validation_error_scrubbing_handler(request, exc):
|
||||
# ``jsonable_encoder`` walks the scrubbed payload one more time
|
||||
# to convert anything else Pydantic v2 surfaces (URL objects,
|
||||
# Path objects, Url instances, etc.) into JSON-safe primitives.
|
||||
return _JSONResponse(
|
||||
status_code = 422,
|
||||
content = {"detail": _scrub_validation_obj(exc.errors())},
|
||||
content = _jsonable_encoder(
|
||||
{"detail": _scrub_validation_obj(exc.errors())}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1519,9 +1519,15 @@ class DiffusionLoadRequest(BaseModel):
|
|||
def _no_control_chars(cls, v, info):
|
||||
return _no_control_chars(v, info.field_name)
|
||||
|
||||
@field_validator("repo_id", "base_repo")
|
||||
@field_validator("repo_id", "gguf_filename", "base_repo")
|
||||
@classmethod
|
||||
def _no_embedded_hf_tokens(cls, v, info):
|
||||
# Round 17 P2 #12: ``gguf_filename`` is forwarded to the
|
||||
# backend and stored on ``DiffusionBackend._gguf_filename``,
|
||||
# which is later surfaced via ``status()`` / log lines. If a
|
||||
# user pastes a URL-form quant path like
|
||||
# ``https://hf_xxxxx@huggingface.co/.../flux.gguf`` we drop
|
||||
# the embedded credential before it can leak.
|
||||
return _reject_embedded_hf_token(v, info.field_name)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -148,7 +148,7 @@ async def load_checkpoint(
|
|||
# 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
|
||||
from routes.inference import _release_chat_for, _release_diffusion_for
|
||||
|
||||
await _release_chat_for("export")
|
||||
|
||||
|
|
@ -157,25 +157,15 @@ async def load_checkpoint(
|
|||
# shutdown above. is_loading is treated like is_loaded so an
|
||||
# in-flight load is also waited out (the diffusion unload
|
||||
# acquires _load_lock + _generate_lock and blocks until the
|
||||
# current load completes, then unloads). Best effort; silently
|
||||
# skip if the module is absent.
|
||||
try:
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff = get_diffusion_backend()
|
||||
diff_status = diff.status()
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
logger.info(
|
||||
"Unloading diffusion model (loaded=%s loading=%s) for export",
|
||||
diff_status.get("is_loaded"),
|
||||
diff_status.get("is_loading"),
|
||||
)
|
||||
# Block-move to thread; unload acquires the
|
||||
# diffusion _load_lock + _generate_lock and can take
|
||||
# the full duration of an in-flight load/generation.
|
||||
await asyncio.to_thread(diff.unload_model)
|
||||
except Exception as e:
|
||||
logger.debug("diffusion unload skipped for export: %s", e)
|
||||
# current load completes, then unloads).
|
||||
# Round 17: previously this was a best-effort try/except that
|
||||
# swallowed every failure with logger.debug, so a wedged
|
||||
# diffusion backend let the export checkpoint load anyway and
|
||||
# OOM at first allocation. ``_release_diffusion_for`` is
|
||||
# strict: it raises HTTPException 503 if status() or
|
||||
# unload_model() fails, or if the backend remains loaded or
|
||||
# loading after the unload call.
|
||||
await _release_diffusion_for("export load")
|
||||
|
||||
# load_checkpoint spawns and waits on a subprocess and can take
|
||||
# minutes. Run it in a worker thread so the event loop stays
|
||||
|
|
|
|||
|
|
@ -487,44 +487,139 @@ async def _release_export_for(workload: str) -> None:
|
|||
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.
|
||||
Round 17 P1 #8: idle-export shutdown failures now raise HTTP 503
|
||||
instead of being swallowed, so a wedged export subprocess does
|
||||
not silently leave GPU memory pinned while training / chat /
|
||||
diffusion start on top.
|
||||
"""
|
||||
try:
|
||||
from core.export import get_export_backend # type: ignore
|
||||
except Exception as exc:
|
||||
logger.debug("export backend unavailable for %s: %s", workload, exc)
|
||||
return
|
||||
|
||||
try:
|
||||
exp = get_export_backend()
|
||||
has_checkpoint = bool(getattr(exp, "current_checkpoint", None))
|
||||
# Backends without an async-job tracker (older builds, some
|
||||
# test mocks) cannot report 'active' separately from
|
||||
# 'has_checkpoint'. Treat absence as 'not active' so a
|
||||
# settled checkpoint still gets dropped; on builds that DO
|
||||
# expose it, a True value blocks the drop.
|
||||
is_export_active_fn = getattr(exp, "is_export_active", None)
|
||||
if is_export_active_fn is None:
|
||||
active = False
|
||||
else:
|
||||
try:
|
||||
active = bool(is_export_active_fn())
|
||||
except Exception:
|
||||
# Treat unverifiable 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:
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Could not access export backend before %s: %s", workload, exc
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Could not access export backend before starting {workload}. "
|
||||
"Try again."
|
||||
),
|
||||
) from exc
|
||||
|
||||
has_checkpoint = bool(getattr(exp, "current_checkpoint", None))
|
||||
is_export_active_fn = getattr(exp, "is_export_active", None)
|
||||
if is_export_active_fn is None:
|
||||
active = False
|
||||
else:
|
||||
try:
|
||||
active = bool(is_export_active_fn())
|
||||
except Exception as exc:
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Could not verify export status before starting "
|
||||
f"{workload}. Try again."
|
||||
),
|
||||
) from exc
|
||||
|
||||
if has_checkpoint and not active:
|
||||
try:
|
||||
logger.info(
|
||||
"Shutting down idle export (checkpoint=%s) for %s",
|
||||
has_checkpoint,
|
||||
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)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Could not shut down export for %s: %s", workload, exc
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Could not unload the idle export checkpoint before "
|
||||
f"starting {workload}. Try again."
|
||||
),
|
||||
) from exc
|
||||
exp.current_checkpoint = None
|
||||
exp.is_vision = False
|
||||
exp.is_peft = False
|
||||
|
||||
|
||||
async def _release_diffusion_for(workload: str) -> None:
|
||||
"""Strict diffusion-unload helper for cross-workload handoffs.
|
||||
|
||||
Round 17 P1 #4-7: the GGUF chat load, safetensors chat load,
|
||||
training start, and export load paths each had their own
|
||||
best-effort try/except around ``diff_backend.unload_model()``.
|
||||
A wedged diffusion pipeline therefore stayed resident while a
|
||||
new GPU workload started on top. This helper raises HTTP 503
|
||||
when the unload fails or leaves diffusion resident, so the
|
||||
caller fails closed.
|
||||
"""
|
||||
try:
|
||||
from core.inference.diffusion import get_diffusion_backend # type: ignore
|
||||
except Exception as exc:
|
||||
logger.debug("diffusion backend unavailable for %s: %s", workload, exc)
|
||||
return
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
try:
|
||||
diff_status = diff_backend.status()
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Could not verify diffusion status before %s: %s", workload, exc
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Could not verify diffusion status before starting "
|
||||
f"{workload}. Try again."
|
||||
),
|
||||
) from exc
|
||||
|
||||
if not (diff_status.get("is_loaded") or diff_status.get("is_loading")):
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"Unloading diffusion (loaded=%s loading=%s) before %s",
|
||||
diff_status.get("is_loaded"),
|
||||
diff_status.get("is_loading"),
|
||||
workload,
|
||||
)
|
||||
try:
|
||||
result = await asyncio.to_thread(diff_backend.unload_model)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to unload diffusion before %s: %s", workload, exc)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Could not unload the existing diffusion image model "
|
||||
f"before starting {workload}. Try again."
|
||||
),
|
||||
) from exc
|
||||
|
||||
after = {}
|
||||
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
|
||||
if result is False or after.get("is_loaded") or after.get("is_loading"):
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"The diffusion image model is still active after unload; "
|
||||
f"retry before starting {workload}."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
|
||||
|
|
@ -1045,30 +1140,11 @@ async def load_model(
|
|||
# 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
|
||||
# for VRAM. Also handles is_loading: unload_model takes
|
||||
# _load_lock + _generate_lock and will wait out an
|
||||
# in-flight load before clearing state. Best effort;
|
||||
# silently continue on failure.
|
||||
try:
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
diff_status = diff_backend.status()
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
logger.info(
|
||||
"Unloading diffusion (loaded=%s loading=%s) before GGUF load",
|
||||
diff_status.get("is_loaded"),
|
||||
diff_status.get("is_loading"),
|
||||
)
|
||||
# diff_backend.unload_model takes _load_lock +
|
||||
# _generate_lock and can block for the duration of
|
||||
# an in-flight load / generation. Off-load to a
|
||||
# worker thread to keep the event loop responsive.
|
||||
await asyncio.to_thread(diff_backend.unload_model)
|
||||
except Exception as e:
|
||||
logger.debug("diffusion unload skipped (GGUF path): %s", e)
|
||||
# Round 17 P1 #4: route the diffusion unload through the
|
||||
# strict ``_release_diffusion_for`` helper so a wedged
|
||||
# diffusion pipeline blocks the GGUF chat load with 503
|
||||
# instead of silently double-owning VRAM.
|
||||
await _release_diffusion_for("GGUF chat load")
|
||||
|
||||
# Inherit llama_extra_args from the previous load when the
|
||||
# request omits the field (the chat-settings Apply path
|
||||
|
|
@ -1254,26 +1330,10 @@ async def load_model(
|
|||
# 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
|
||||
# is treated like is_loaded; unload waits behind _load_lock +
|
||||
# _generate_lock so the in-flight load completes first.
|
||||
try:
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
diff_status = diff_backend.status()
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
logger.info(
|
||||
"Unloading diffusion (loaded=%s loading=%s) before chat load",
|
||||
diff_status.get("is_loaded"),
|
||||
diff_status.get("is_loading"),
|
||||
)
|
||||
# Same blocking concern as the GGUF chat path:
|
||||
# _load_lock + _generate_lock serialise the call.
|
||||
await asyncio.to_thread(diff_backend.unload_model)
|
||||
except Exception as e:
|
||||
logger.debug("diffusion unload skipped: %s", e)
|
||||
# Round 17 P1 #5: strict diffusion unload via the shared
|
||||
# helper so a wedged pipeline blocks the safetensors chat
|
||||
# load with 503 instead of silently double-owning VRAM.
|
||||
await _release_diffusion_for("safetensors chat load")
|
||||
|
||||
# Export was already dropped above via the shared
|
||||
# ``await _release_export_for("safetensors chat")`` call
|
||||
|
|
@ -1968,6 +2028,14 @@ 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")
|
||||
backend = _get_diffusion_backend()
|
||||
try:
|
||||
status = await asyncio.get_event_loop().run_in_executor(
|
||||
|
|
@ -1993,7 +2061,13 @@ async def diffusion_load(
|
|||
if (
|
||||
"Could not verify training status" in detail
|
||||
or "Could not verify export status" in detail
|
||||
or "Could not unload" in detail
|
||||
or "refused to unload" in detail
|
||||
or "still active after unload" in detail
|
||||
):
|
||||
# Round 17 P1 #2: chat unload failures raised by the
|
||||
# backend helper map to 503 (retryable infra issue),
|
||||
# matching the route-level _release_*_for helpers.
|
||||
raise HTTPException(status_code = 503, detail = detail) from exc
|
||||
if (
|
||||
"export job is currently active" in detail
|
||||
|
|
|
|||
|
|
@ -274,6 +274,7 @@ async def start_training(
|
|||
from routes.inference import (
|
||||
_raise_if_export_active,
|
||||
_release_chat_for,
|
||||
_release_diffusion_for,
|
||||
_release_export_for,
|
||||
)
|
||||
|
||||
|
|
@ -285,23 +286,13 @@ async def start_training(
|
|||
# holds the same GPU and would survive the inference shutdown.
|
||||
# is_loading=True is also handled (unload_model takes
|
||||
# _load_lock + _generate_lock and waits the in-flight load out).
|
||||
try:
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
diff_status = diff_backend.status()
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
logger.info(
|
||||
"Unloading diffusion (loaded=%s loading=%s) for training",
|
||||
diff_status.get("is_loaded"),
|
||||
diff_status.get("is_loading"),
|
||||
)
|
||||
# Async route: offload the blocking unload to a
|
||||
# worker thread so the event loop stays responsive
|
||||
# during long in-flight load / generate calls.
|
||||
await asyncio.to_thread(diff_backend.unload_model)
|
||||
except Exception as e:
|
||||
logger.warning("Could not unload diffusion model: %s", e)
|
||||
# Round 17: previously the diffusion unload was best-effort
|
||||
# (try/except + logger.warning), so a stuck diffusion backend
|
||||
# would let training start anyway and immediately OOM the
|
||||
# subprocess. ``_release_diffusion_for`` is strict: it raises
|
||||
# HTTPException 503 if status() or unload_model() fails, or if
|
||||
# the backend remains loaded / loading after the unload call.
|
||||
await _release_diffusion_for("training")
|
||||
|
||||
# start_training now spawns a subprocess (non-blocking)
|
||||
success = backend.start_training(job_id = job_id, **training_kwargs)
|
||||
|
|
|
|||
|
|
@ -1056,72 +1056,40 @@ def test_generate_image_does_not_block_status(monkeypatch):
|
|||
t.join(timeout = 5)
|
||||
|
||||
|
||||
def test_load_publishes_pending_target_during_loading(monkeypatch):
|
||||
def test_load_publishes_pending_target_during_loading():
|
||||
"""status() must expose the pending repo_id / base_repo / gguf
|
||||
file while is_loading=True so cache- and finetuned-delete guards
|
||||
can refuse to rmtree the repo being downloaded right now."""
|
||||
import threading
|
||||
can refuse to rmtree the repo being downloaded right now.
|
||||
|
||||
The pending exposure is purely a state-shape contract: load_model
|
||||
sets _loading + _pending_* under _lock at the start, and status()
|
||||
snapshots them under _lock. Test the contract directly instead of
|
||||
racing a fake pipeline through a background thread, which was
|
||||
flaky on the Windows runner (the chat-release helpers' transitive
|
||||
imports of core.training.resume failed there and the load thread
|
||||
exited cleanly before the main thread observed the pending state).
|
||||
"""
|
||||
import core.inference.diffusion as d
|
||||
from PIL import Image
|
||||
|
||||
fake = _install_fake_diffusers(monkeypatch)
|
||||
backend = d.DiffusionBackend()
|
||||
# Simulate the state load_model publishes at the top of its
|
||||
# critical section, before from_pretrained runs.
|
||||
with backend._lock:
|
||||
backend._loading = True
|
||||
backend._pending_repo_id = "unsloth/FLUX.2-klein-4B-GGUF"
|
||||
backend._pending_base_repo = "black-forest-labs/FLUX.2-klein-4B"
|
||||
backend._pending_gguf_filename = "flux-2-klein-4b-Q4_K_S.gguf"
|
||||
|
||||
pending_seen: dict = {}
|
||||
pretrained_blocked = threading.Event()
|
||||
pretrained_release = threading.Event()
|
||||
|
||||
class _SlowPipeline:
|
||||
@classmethod
|
||||
def from_pretrained(cls, base_repo, **kwargs):
|
||||
pretrained_blocked.set()
|
||||
# Capture status() output while the load is blocked.
|
||||
backend = d.get_diffusion_backend()
|
||||
pending_seen.update(backend.status())
|
||||
pretrained_release.wait(timeout = 5)
|
||||
inst = cls()
|
||||
inst.base_repo = base_repo
|
||||
return inst
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
class _Out:
|
||||
pass
|
||||
|
||||
o = _Out()
|
||||
o.images = [Image.new("RGB", (kwargs["width"], kwargs["height"]))]
|
||||
return o
|
||||
|
||||
def enable_model_cpu_offload(self):
|
||||
pass
|
||||
|
||||
def to(self, device):
|
||||
return self
|
||||
|
||||
fake.Flux2KleinPipeline = _SlowPipeline
|
||||
|
||||
backend = d.get_diffusion_backend()
|
||||
backend.unload_model()
|
||||
|
||||
def do_load():
|
||||
try:
|
||||
backend.load_model(
|
||||
"unsloth/FLUX.2-klein-4B-GGUF",
|
||||
gguf_filename = "flux-2-klein-4b-Q4_K_S.gguf",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
t = threading.Thread(target = do_load)
|
||||
t.start()
|
||||
try:
|
||||
assert pretrained_blocked.wait(timeout = 5)
|
||||
# While blocked inside from_pretrained, status reads should
|
||||
# already see the pending repo so deletes can be refused.
|
||||
assert pending_seen.get("is_loading") is True
|
||||
assert pending_seen.get("repo_id") == "unsloth/FLUX.2-klein-4B-GGUF"
|
||||
assert pending_seen.get("base_repo") == "black-forest-labs/FLUX.2-klein-4B"
|
||||
finally:
|
||||
pretrained_release.set()
|
||||
t.join(timeout = 5)
|
||||
public = backend.status()
|
||||
assert public["is_loading"] is True
|
||||
assert public["repo_id"] == "unsloth/FLUX.2-klein-4B-GGUF"
|
||||
assert public["base_repo"] == "black-forest-labs/FLUX.2-klein-4B"
|
||||
# Guard-facing internal payload also reports the pending fields
|
||||
# under their dedicated keys.
|
||||
internal = backend.status(include_internal = True)
|
||||
assert internal["pending_repo_id"] == "unsloth/FLUX.2-klein-4B-GGUF"
|
||||
assert internal["pending_base_repo"] == "black-forest-labs/FLUX.2-klein-4B"
|
||||
assert internal["pending_gguf_filename"] == "flux-2-klein-4b-Q4_K_S.gguf"
|
||||
|
||||
|
||||
def test_unload_waits_for_in_flight_generation(monkeypatch):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue