Fix/adjust diffusion: round 16 P1+P2 batch for PR #5754
Round 16 reviewer aggregate (logs/review_round16_aggregate.md): P1 fixes: - routes/models.py /delete-cached llama guard pairs loading_id with loading_hf_variant so deleting a different cached quant (Q8_0) while another variant (Q4_K_M) is loading is no longer blocked. - core/inference/diffusion.py load_model now calls _release_other_gpu_owners_for_diffusion BEFORE _release_chat_backend_for_diffusion. The other-owners helper RAISES on active training/export, so a route -> worker race or direct backend caller no longer drops the user's chat model before the diffusion load is refused. - routes/models.py /delete-cached diffusion guard fails CLOSED (503) on HF cache scan failure instead of silently falling through to repo-id-only matching, which could miss a loaded local snapshot path. - routes/inference.py _release_llama_for and _release_safetensors_chat_for now raise 503 on actual unload failure (exception or False return), so new GPU workloads do not start while the old chat process still owns VRAM. - core/inference/diffusion.py status() now takes include_internal=False by default and only exposes the guard-facing active_*/pending_* paths when callers opt in. The public /api/inference/images/status route gets the redacted payload; routes/models.py delete guards pass include_internal=True so they still see the raw paths. - core/inference/diffusion.py generate_image_with_metadata routes the response model through _display_repo_id so /images/generate cannot echo back an absolute local path. P2 fixes: - routes/inference.py /images/load now maps backend "Could not verify training/export status" to 503 instead of 409, matching the route-level pre-check. - core/inference/diffusion.py _release_other_gpu_owners_for_diffusion raises "Could not verify export status" when the is_export_active() probe itself raises, instead of silently treating it as active export. - core/inference/diffusion.py detect_family compares compact family spellings (Flux2Klein) against per-token compact strings so unsloth/Flux2Klein-GGUF matches the flux.2-klein family without matching the embedded substring inside flux.20. - main.py installs a RequestValidationError handler that scrubs hf_xxxxx tokens out of the 422 response body so a rejected ``repo_id`` containing a URL-embedded HF token does not echo it back to the browser. Tests: - 3 new regression cases (Flux2Klein compact alias, public status redaction, generate_image_with_metadata redaction). - All 75 diffusion backend + route tests pass.
This commit is contained in:
Daniel Han-Chen
parent
7c8f1eb40d
commit
2ef9b0e09f
5 changed files with 343 additions and 109 deletions
|
|
@ -357,34 +357,37 @@ def detect_family(
|
|||
# P2 #8).
|
||||
needle_norm = re.sub(r"[^a-z0-9]+", "-", needle).strip("-")
|
||||
needle_compact = re.sub(r"[^a-z0-9]+", "", needle)
|
||||
# Per-token compact strings let ``unsloth/Flux2Klein-GGUF`` match
|
||||
# the ``flux2klein`` alias: the whole-needle compact is
|
||||
# ``unslothflux2kleingguf`` and the regex boundary check rejects
|
||||
# the embedded match, but the token ``Flux2Klein`` (between the
|
||||
# ``/`` and the ``-``) compacts to exactly ``flux2klein`` (round
|
||||
# 16 P2 #9).
|
||||
needle_compact_tokens = {
|
||||
re.sub(r"[^a-z0-9]+", "", token)
|
||||
for token in re.split(r"[^a-z0-9]+", needle)
|
||||
if token
|
||||
}
|
||||
|
||||
def _matches_family_token(term: str) -> bool:
|
||||
"""Token-boundary match on the normalised needle. Prevents
|
||||
``owner/flux.20-model`` from matching ``flux.2`` because
|
||||
``flux.20`` does not have a separator after ``flux-2``
|
||||
(round 15 P2 #8). Falls back to compact equality so aliases
|
||||
like ``qwenimage`` still match ``unsloth/QwenImage-GGUF``."""
|
||||
(round 15 P2 #8). Compact spellings (``flux2klein``) match
|
||||
only when they appear as a complete repo-name token, not
|
||||
as a substring of a longer token (round 16 P2 #9)."""
|
||||
term_norm = re.sub(r"[^a-z0-9]+", "-", term.lower()).strip("-")
|
||||
if not term_norm:
|
||||
return False
|
||||
if re.search(rf"(^|-){re.escape(term_norm)}($|-)", needle_norm):
|
||||
return True
|
||||
term_compact = re.sub(r"[^a-z0-9]+", "", term.lower())
|
||||
if term_compact and term_compact in needle_compact:
|
||||
# Compact contiguous match: ``qwenimage`` in
|
||||
# ``qwenimage-gguf`` -> qwenimage-compact in needle_compact.
|
||||
# Use word boundary on the compact form too: the compact
|
||||
# ``flux2`` must not match inside ``flux20``.
|
||||
return (
|
||||
bool(
|
||||
re.search(
|
||||
rf"(^|[^0-9a-z]){re.escape(term_compact)}([^0-9a-z]|$)",
|
||||
needle_compact,
|
||||
)
|
||||
)
|
||||
or term_compact == needle_compact
|
||||
)
|
||||
return False
|
||||
if not term_compact:
|
||||
return False
|
||||
return (
|
||||
term_compact in needle_compact_tokens
|
||||
or term_compact == needle_compact
|
||||
)
|
||||
|
||||
# Scan _FAMILIES first (GGUF-supported), then _FULL_REPO_FAMILIES
|
||||
# so a repo like ``stabilityai/stable-diffusion-xl-base-1.0`` is
|
||||
|
|
@ -496,7 +499,7 @@ class DiffusionBackend:
|
|||
def repo_id(self) -> Optional[str]:
|
||||
return self._repo_id
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
def status(self, *, include_internal: bool = False) -> dict[str, Any]:
|
||||
# Take _lock so the snapshot cannot observe a torn state where
|
||||
# _pipe was already swapped but _family/_repo_id haven't been
|
||||
# updated yet (or vice versa). Frontend polling at 1 Hz would
|
||||
|
|
@ -504,6 +507,14 @@ class DiffusionBackend:
|
|||
# Only echo the GGUF basename; full absolute path leaks the
|
||||
# local HF cache layout (and the system username on default
|
||||
# POSIX layouts) to any authenticated Studio session.
|
||||
#
|
||||
# Round 16 P1 #5: the guard-facing ``active_*`` / ``pending_*``
|
||||
# fields hold the EXACT raw path (so /delete-cached can match
|
||||
# an HF snapshot mmap) but are NOT safe to surface to the
|
||||
# browser. Callers that need the raw path (route-internal
|
||||
# delete guards) pass ``include_internal=True``; the public
|
||||
# ``/api/inference/images/status`` route always uses the
|
||||
# public payload.
|
||||
with self._lock:
|
||||
# UI-facing collapsed basename. Full local path leaks the
|
||||
# HF cache layout + system username; the original caller-
|
||||
|
|
@ -552,7 +563,7 @@ class DiffusionBackend:
|
|||
# guard-facing ``active_*`` / ``pending_*`` fields below
|
||||
# preserve the exact value so delete guards still match
|
||||
# against the snapshot path.
|
||||
return {
|
||||
payload: dict[str, Any] = {
|
||||
"is_loaded": self._pipe is not None,
|
||||
"is_loading": self._loading,
|
||||
"repo_id": _display_repo_id(pending_repo or active_repo),
|
||||
|
|
@ -560,23 +571,30 @@ class DiffusionBackend:
|
|||
"pipeline_class": ui_pipeline_class,
|
||||
"base_repo": _display_repo_id(pending_base or active_base),
|
||||
"gguf_filename": ui_gguf_basename,
|
||||
# Guard-facing fields: every repo / path / GGUF
|
||||
# filename the backend owns RIGHT NOW. Delete routes
|
||||
# iterate both, paired so the variant-filename check
|
||||
# is compared against the SAME repo that owns it
|
||||
# (round 13 P1 #3-5).
|
||||
"active_repo_id": active_repo,
|
||||
"active_base_repo": active_base,
|
||||
"active_gguf_filename": active_gguf,
|
||||
"pending_repo_id": pending_repo,
|
||||
"pending_base_repo": pending_base,
|
||||
"pending_gguf_filename": pending_gguf,
|
||||
"device": self._device,
|
||||
"dtype": self._dtype,
|
||||
"loaded_at": self._loaded_at,
|
||||
"last_error": self._last_error,
|
||||
"supported_families": supported_families(),
|
||||
}
|
||||
if include_internal:
|
||||
# Guard-facing fields: every repo / path / GGUF
|
||||
# filename the backend owns RIGHT NOW. Delete routes
|
||||
# iterate both, paired so the variant-filename check
|
||||
# is compared against the SAME repo that owns it
|
||||
# (round 13 P1 #3-5). Round 16 P1 #5: never returned
|
||||
# by the public /images/status route.
|
||||
payload.update(
|
||||
{
|
||||
"active_repo_id": active_repo,
|
||||
"active_base_repo": active_base,
|
||||
"active_gguf_filename": active_gguf,
|
||||
"pending_repo_id": pending_repo,
|
||||
"pending_base_repo": pending_base,
|
||||
"pending_gguf_filename": pending_gguf,
|
||||
}
|
||||
)
|
||||
return payload
|
||||
|
||||
def _pick_device_and_dtype(self) -> tuple[str, "Any"]:
|
||||
"""Pick (device, dtype) for the current host.
|
||||
|
|
@ -808,20 +826,29 @@ class DiffusionBackend:
|
|||
# transient Hub error on the GGUF download) have now
|
||||
# been validated. Anything past this line allocates
|
||||
# GPU memory, so:
|
||||
# 1. Release competing GPU owners (chat + export).
|
||||
# 2. Release any *previous* diffusion pipeline so the
|
||||
# 1. Verify training is idle and the export job (if
|
||||
# any) is also idle. ``_release_other_gpu_owners
|
||||
# _for_diffusion`` RAISES on conflict, so it must
|
||||
# run BEFORE we unload chat (round 16 P1 #2): a
|
||||
# route precheck -> worker race could otherwise
|
||||
# drop the user's chat model only to bail out
|
||||
# because training started in between, and a
|
||||
# direct ``DiffusionBackend.load_model`` caller
|
||||
# that did not run the route prechecks would also
|
||||
# leave chat unloaded for nothing.
|
||||
# 2. Release the chat backend (llama-server + the
|
||||
# safetensors orchestrator) now that we know the
|
||||
# load can actually proceed.
|
||||
# 3. Release any *previous* diffusion pipeline so the
|
||||
# new transformer / new from_pretrained does not
|
||||
# race the old pipe for VRAM. Switching between
|
||||
# FLUX.2 klein 4B and 9B on a 16-24 GB GPU OOMs
|
||||
# otherwise: from_single_file allocates the new
|
||||
# transformer while the old pipeline still owns
|
||||
# its weights.
|
||||
# 3. THEN call from_single_file / from_pretrained.
|
||||
# Training is *not* unloaded here: the route layer
|
||||
# refuses /images/load with HTTP 409 when training is
|
||||
# active so the user keeps their long run.
|
||||
_release_chat_backend_for_diffusion()
|
||||
# 4. THEN call from_single_file / from_pretrained.
|
||||
_release_other_gpu_owners_for_diffusion()
|
||||
_release_chat_backend_for_diffusion()
|
||||
|
||||
old = self._pipe
|
||||
if old is not None:
|
||||
|
|
@ -1172,8 +1199,12 @@ class DiffusionBackend:
|
|||
with self._generate_lock:
|
||||
image = self._generate_image_unlocked(**kwargs)
|
||||
with self._lock:
|
||||
# Round 16 P1 #6: route ``model`` through
|
||||
# _display_repo_id so a generation response for a
|
||||
# locally-loaded model cannot echo back an absolute
|
||||
# filesystem path to the browser.
|
||||
meta = {
|
||||
"model": self._repo_id,
|
||||
"model": _display_repo_id(self._repo_id),
|
||||
"family": self._family.name if self._family else None,
|
||||
}
|
||||
return image, meta
|
||||
|
|
@ -1343,12 +1374,16 @@ def _release_other_gpu_owners_for_diffusion() -> None:
|
|||
if is_export_active_fn is not None:
|
||||
try:
|
||||
export_is_active = bool(is_export_active_fn())
|
||||
except Exception:
|
||||
# Unverifiable status -> treat as 'might be active' and
|
||||
# refuse so a direct backend caller (test / script /
|
||||
# future route that forgot the higher-level 409 guard)
|
||||
# cannot still terminate an in-flight export.
|
||||
export_is_active = True
|
||||
except Exception as exc:
|
||||
# Round 16 P2 #8: distinguish unverifiable status from
|
||||
# active export. The previous "treat as active" mapping
|
||||
# surfaced as a misleading 409 conflict; raise a
|
||||
# "Could not verify" RuntimeError so the route layer
|
||||
# maps it to 503 (retryable) instead.
|
||||
raise RuntimeError(
|
||||
"Could not verify export status before loading a "
|
||||
"diffusion image model."
|
||||
) from exc
|
||||
if export_is_active:
|
||||
# Round 14 P2 #10: the prior behaviour logged a warning
|
||||
# and continued, so direct ``DiffusionBackend.load_model``
|
||||
|
|
|
|||
|
|
@ -293,6 +293,43 @@ app = FastAPI(
|
|||
lifespan = lifespan,
|
||||
)
|
||||
|
||||
|
||||
# ── Validation error scrubber ────────────────────────────────────
|
||||
# Round 16 P2 #10: FastAPI's default RequestValidationError handler
|
||||
# echoes the rejected ``input`` value back in the 422 body. A
|
||||
# request like
|
||||
# {"repo_id": "https://hf_token@huggingface.co/owner/repo"}
|
||||
# is rejected by ``DiffusionLoadRequest._no_embedded_hf_tokens``,
|
||||
# but the rejected URL would still appear in the response payload,
|
||||
# leaking the token to the browser console / network log. Wrap the
|
||||
# handler so any ``hf_xxxxx`` substring is replaced with
|
||||
# ``<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.responses import JSONResponse as _JSONResponse # noqa: E402
|
||||
import re as _re_validation # noqa: E402
|
||||
|
||||
|
||||
_HF_TOKEN_VALIDATION_RE = _re_validation.compile(r"hf_[A-Za-z0-9]{20,}")
|
||||
|
||||
|
||||
def _scrub_validation_obj(value):
|
||||
if isinstance(value, str):
|
||||
return _HF_TOKEN_VALIDATION_RE.sub("<redacted>", value)
|
||||
if isinstance(value, list):
|
||||
return [_scrub_validation_obj(v) for v in value]
|
||||
if isinstance(value, dict):
|
||||
return {k: _scrub_validation_obj(v) for k, v in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
@app.exception_handler(_RequestValidationError)
|
||||
async def _validation_error_scrubbing_handler(request, exc):
|
||||
return _JSONResponse(
|
||||
status_code = 422,
|
||||
content = {"detail": _scrub_validation_obj(exc.errors())},
|
||||
)
|
||||
|
||||
# Initialize structured logging
|
||||
from loggers.config import LogConfig
|
||||
from loggers.handlers import LoggingMiddleware
|
||||
|
|
|
|||
|
|
@ -366,60 +366,103 @@ async def _release_llama_for(workload: str) -> None:
|
|||
/export/load could start while a long ``_download_gguf`` was in
|
||||
flight; llama-server would then come up afterwards and double-own
|
||||
the GPU.
|
||||
|
||||
Round 16 P1 #4: a missing or unavailable llama backend is a
|
||||
silent no-op (fresh install / no GGUF use), but an unload that
|
||||
actually FAILS raises 503 so the caller does not start a new GPU
|
||||
workload while llama-server is still resident.
|
||||
"""
|
||||
try:
|
||||
llama = get_llama_cpp_backend()
|
||||
is_loaded = bool(getattr(llama, "is_loaded", False))
|
||||
is_active = bool(getattr(llama, "is_active", False))
|
||||
is_loading = bool(getattr(llama, "loading_model_identifier", None))
|
||||
if is_loaded or is_active or is_loading:
|
||||
logger.info(
|
||||
"Unloading GGUF chat (loaded=%s active=%s loading=%s) before %s load",
|
||||
is_loaded,
|
||||
is_active,
|
||||
is_loading,
|
||||
workload,
|
||||
)
|
||||
await asyncio.to_thread(llama.unload_model)
|
||||
except Exception as e:
|
||||
logger.debug("llama-server unload skipped for %s: %s", workload, e)
|
||||
except Exception as exc:
|
||||
logger.debug("llama-server unavailable for %s: %s", workload, exc)
|
||||
return
|
||||
|
||||
is_loaded = bool(getattr(llama, "is_loaded", False))
|
||||
is_active = bool(getattr(llama, "is_active", False))
|
||||
is_loading = bool(getattr(llama, "loading_model_identifier", None))
|
||||
if not (is_loaded or is_active or is_loading):
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"Unloading GGUF chat (loaded=%s active=%s loading=%s) before %s load",
|
||||
is_loaded,
|
||||
is_active,
|
||||
is_loading,
|
||||
workload,
|
||||
)
|
||||
try:
|
||||
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(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Could not unload the existing GGUF chat model before "
|
||||
f"starting {workload}."
|
||||
),
|
||||
) from exc
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Round 16 P1 #4: ``unload_model`` returning ``False`` (subprocess
|
||||
wedged, IPC timeout) used to be silently ignored, leaving the
|
||||
old chat model resident while a new GPU workload started on top.
|
||||
Treat ``False`` as failure and raise 503 so the caller retries
|
||||
instead of double-owning VRAM.
|
||||
"""
|
||||
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,
|
||||
except Exception as exc:
|
||||
logger.debug("safetensors unavailable for %s: %s", workload, exc)
|
||||
return
|
||||
|
||||
async def _unload_required(model_name: str) -> None:
|
||||
try:
|
||||
ok = await asyncio.to_thread(inf.unload_model, model_name)
|
||||
except Exception as exc:
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Could not unload safetensors chat model "
|
||||
f"'{model_name}' before starting {workload}."
|
||||
),
|
||||
) from exc
|
||||
if ok is False:
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
f"Safetensors backend refused to unload "
|
||||
f"'{model_name}' before starting {workload}. "
|
||||
"Try again."
|
||||
),
|
||||
)
|
||||
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)
|
||||
|
||||
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 _unload_required(active_model_name)
|
||||
for loading in loading_models:
|
||||
if loading == active_model_name:
|
||||
continue
|
||||
logger.info(
|
||||
"Unloading in-flight safetensors chat '%s' before %s load",
|
||||
loading,
|
||||
workload,
|
||||
)
|
||||
await _unload_required(loading)
|
||||
|
||||
|
||||
async def _release_chat_for(workload: str) -> None:
|
||||
|
|
@ -1942,18 +1985,21 @@ async def diffusion_load(
|
|||
)
|
||||
return JSONResponse(content = status)
|
||||
except RuntimeError as exc:
|
||||
# Round 15 P2 #7: if a training run / export job starts
|
||||
# between the route-level pre-check and the backend worker,
|
||||
# ``_release_other_gpu_owners_for_diffusion`` raises a
|
||||
# RuntimeError that should surface as a 409 conflict (the
|
||||
# same status the route layer returns), not 400. Match the
|
||||
# known conflict strings the backend raises.
|
||||
# Round 15 P2 #7 / round 16 P2 #7: backend-level conflict
|
||||
# checks raise RuntimeError that surfaces here. Distinguish:
|
||||
# - "Could not verify ..." -> 503 (retryable, status check
|
||||
# itself failed), matching the route-level pre-check.
|
||||
# - explicit "currently active" -> 409 conflict.
|
||||
# - anything else -> 400 (bad request).
|
||||
detail = str(exc)
|
||||
if (
|
||||
"Could not verify training status" in detail
|
||||
or "Could not verify export status" in detail
|
||||
):
|
||||
raise HTTPException(status_code = 503, detail = detail) from exc
|
||||
if (
|
||||
"export job is currently active" in detail
|
||||
or "Training is currently active" in detail
|
||||
or "Could not verify training status" in detail
|
||||
or "Could not verify export status" in detail
|
||||
):
|
||||
raise HTTPException(status_code = 409, detail = detail) from exc
|
||||
raise HTTPException(status_code = 400, detail = detail) from exc
|
||||
|
|
|
|||
|
|
@ -2062,7 +2062,9 @@ async def delete_finetuned_model(
|
|||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
diff_status = diff_backend.status()
|
||||
# include_internal=True so we can iterate active_*/pending_*
|
||||
# raw paths against ``target_path`` (round 16 P1 #5).
|
||||
diff_status = diff_backend.status(include_internal = True)
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
target_str = str(target_path)
|
||||
# Pair each owned repo / path with the GGUF variant it
|
||||
|
|
@ -2757,17 +2759,32 @@ async def delete_cached_model(
|
|||
loading_id = (
|
||||
getattr(llama_backend, "loading_model_identifier", None) or ""
|
||||
).lower()
|
||||
loading_variant = (
|
||||
getattr(llama_backend, "loading_hf_variant", None) or ""
|
||||
).lower()
|
||||
# Also consult the pending-load identifier: a multi-GB HF
|
||||
# download stays in ``loading_model_identifier`` until the
|
||||
# download completes, before ``model_identifier`` is set
|
||||
# (round 13 P1 #6). Without this check the cache directory
|
||||
# the download was writing into could be rmtree'd mid-flight.
|
||||
# Round 16 P1 #1: pair against ``loading_hf_variant`` so a
|
||||
# delete of a DIFFERENT cached quant from the same repo
|
||||
# (loading Q4_K_M, deleting cached Q8_0) is allowed; only
|
||||
# block when the requested variant matches what is being
|
||||
# downloaded. Mirrors the /delete-finetuned pairing.
|
||||
needle = repo_id.lower()
|
||||
requested_variant = (variant or "").lower()
|
||||
if loading_id == needle:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "Cannot delete a model while it is loading",
|
||||
same_loading_variant = (
|
||||
not requested_variant
|
||||
or not loading_variant
|
||||
or requested_variant == loading_variant
|
||||
)
|
||||
if same_loading_variant:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "Cannot delete a model while it is loading",
|
||||
)
|
||||
# Exact match only (case-insensitive). Prefix match would
|
||||
# block deleting unrelated ``org/model`` while
|
||||
# ``org/model-v2`` is loaded -- same surface the diffusion
|
||||
|
|
@ -2778,7 +2795,6 @@ async def delete_cached_model(
|
|||
llama_backend.is_loaded or getattr(llama_backend, "is_active", False)
|
||||
):
|
||||
loaded_variant = (getattr(llama_backend, "hf_variant", None) or "").lower()
|
||||
requested_variant = (variant or "").lower()
|
||||
same_variant = (
|
||||
not requested_variant
|
||||
or not loaded_variant
|
||||
|
|
@ -2854,7 +2870,9 @@ async def delete_cached_model(
|
|||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
diff_status = diff_backend.status()
|
||||
# include_internal=True so we can pair owned raw paths against
|
||||
# the HF cache snapshot root (round 16 P1 #5).
|
||||
diff_status = diff_backend.status(include_internal = True)
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
needle = repo_id.lower()
|
||||
# Round 15 P1 #4: ALSO compare owned paths against the HF
|
||||
|
|
@ -2879,10 +2897,22 @@ async def delete_cached_model(
|
|||
except Exception:
|
||||
pass
|
||||
except Exception as cache_scan_exc:
|
||||
logger.debug(
|
||||
"HF cache scan failed during diffusion delete guard: %s",
|
||||
# Round 16 P1 #3: a transient cache-scan failure here
|
||||
# used to silently fall through to repo-id-only
|
||||
# matching, which misses local snapshot paths and
|
||||
# let /delete-cached unlink an actively mmap'd
|
||||
# snapshot. Fail-closed (503) so the user retries.
|
||||
logger.warning(
|
||||
"Could not scan HF cache during diffusion delete guard: %s",
|
||||
cache_scan_exc,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
"Could not verify diffusion cache ownership before "
|
||||
"deleting. Try again."
|
||||
),
|
||||
) from cache_scan_exc
|
||||
|
||||
# Pair each owned repo with the GGUF variant it actually
|
||||
# owns (active or pending) so a swap in progress does not
|
||||
|
|
|
|||
|
|
@ -174,6 +174,9 @@ def test_get_diffusion_backend_singleton():
|
|||
|
||||
|
||||
def test_status_shape_unloaded():
|
||||
"""Public status() (the browser-facing payload) must NOT contain
|
||||
the guard-only ``active_*`` / ``pending_*`` fields (round 16
|
||||
P1 #5)."""
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
s = get_diffusion_backend().status()
|
||||
|
|
@ -185,12 +188,6 @@ def test_status_shape_unloaded():
|
|||
"pipeline_class",
|
||||
"base_repo",
|
||||
"gguf_filename",
|
||||
"active_repo_id",
|
||||
"active_base_repo",
|
||||
"active_gguf_filename",
|
||||
"pending_repo_id",
|
||||
"pending_base_repo",
|
||||
"pending_gguf_filename",
|
||||
"device",
|
||||
"dtype",
|
||||
"loaded_at",
|
||||
|
|
@ -198,10 +195,23 @@ def test_status_shape_unloaded():
|
|||
"supported_families",
|
||||
}
|
||||
assert expected_keys.issubset(s.keys())
|
||||
# Guard-facing fields are gated behind include_internal=True.
|
||||
for guard_key in (
|
||||
"active_repo_id",
|
||||
"active_base_repo",
|
||||
"active_gguf_filename",
|
||||
"pending_repo_id",
|
||||
"pending_base_repo",
|
||||
"pending_gguf_filename",
|
||||
):
|
||||
assert guard_key not in s, f"public status() must not expose {guard_key}"
|
||||
assert s["is_loaded"] is False
|
||||
assert s["repo_id"] is None
|
||||
assert s["active_gguf_filename"] is None
|
||||
assert s["pending_gguf_filename"] is None
|
||||
|
||||
# Internal status() exposes the guard fields for delete/route use.
|
||||
s_internal = get_diffusion_backend().status(include_internal = True)
|
||||
assert s_internal["active_gguf_filename"] is None
|
||||
assert s_internal["pending_gguf_filename"] is None
|
||||
|
||||
|
||||
# ── encode_png_base64 ───────────────────────────────────────────
|
||||
|
|
@ -1413,7 +1423,7 @@ def test_status_preserves_active_gguf_subdir(monkeypatch):
|
|||
aliases = (),
|
||||
)
|
||||
|
||||
s = backend.status()
|
||||
s = backend.status(include_internal = True)
|
||||
assert s["active_gguf_filename"] == "BF16/model.gguf"
|
||||
# UI-facing field still collapses to the basename.
|
||||
assert s["gguf_filename"] == "model.gguf"
|
||||
|
|
@ -1504,6 +1514,82 @@ def test_detect_family_rejects_substring_collisions():
|
|||
assert fam is not None and fam.name == "flux.2"
|
||||
|
||||
|
||||
def test_detect_family_compact_aliases_with_owner_prefix():
|
||||
"""Round 16 P2 #9: compact aliases must match when the repo has
|
||||
an owner prefix. ``unsloth/Flux2Klein-GGUF`` -> flux.2-klein
|
||||
via the ``flux2-klein`` alias's compact form. Embedded compact
|
||||
matches (e.g. ``flux2`` inside ``flux20``) must NOT match."""
|
||||
from core.inference.diffusion import detect_family
|
||||
|
||||
fam = detect_family("unsloth/Flux2Klein-GGUF")
|
||||
assert fam is not None and fam.name == "flux.2-klein"
|
||||
# 20 is a different number; must not collide with flux.2.
|
||||
assert detect_family("unsloth/Flux20-GGUF") is None
|
||||
|
||||
|
||||
def test_public_status_does_not_leak_local_path_via_active_fields(monkeypatch):
|
||||
"""Round 16 P1 #5: even the guard-facing active_*/pending_* keys
|
||||
must be absent from the public status payload."""
|
||||
import core.inference.diffusion as d
|
||||
|
||||
backend = d.DiffusionBackend()
|
||||
backend._pipe = object()
|
||||
backend._repo_id = "/home/alice/private-flux"
|
||||
backend._base_repo = "/home/alice/base-private"
|
||||
backend._family = d.DiffusionFamily(
|
||||
name = "flux.2-klein",
|
||||
pipeline_class = "Flux2KleinPipeline",
|
||||
transformer_class = "Flux2Transformer2DModel",
|
||||
base_repo = "black-forest-labs/FLUX.2-klein-4B",
|
||||
aliases = (),
|
||||
)
|
||||
|
||||
public = backend.status()
|
||||
# UI-facing fields collapse to leaf and the guard-only fields are absent.
|
||||
assert public["repo_id"] == "private-flux"
|
||||
assert public["base_repo"] == "base-private"
|
||||
for key in (
|
||||
"active_repo_id",
|
||||
"active_base_repo",
|
||||
"active_gguf_filename",
|
||||
"pending_repo_id",
|
||||
"pending_base_repo",
|
||||
"pending_gguf_filename",
|
||||
):
|
||||
assert key not in public
|
||||
|
||||
internal = backend.status(include_internal = True)
|
||||
assert internal["active_repo_id"] == "/home/alice/private-flux"
|
||||
assert internal["active_base_repo"] == "/home/alice/base-private"
|
||||
|
||||
|
||||
def test_generate_image_with_metadata_redacts_local_path(monkeypatch):
|
||||
"""Round 16 P1 #6: the generation response must not echo a raw
|
||||
absolute path back to the browser."""
|
||||
import core.inference.diffusion as d
|
||||
|
||||
backend = d.DiffusionBackend()
|
||||
backend._pipe = object()
|
||||
backend._repo_id = "/home/alice/private/secret-flux"
|
||||
backend._family = d.DiffusionFamily(
|
||||
name = "flux.2-klein",
|
||||
pipeline_class = "Flux2KleinPipeline",
|
||||
transformer_class = "Flux2Transformer2DModel",
|
||||
base_repo = "black-forest-labs/FLUX.2-klein-4B",
|
||||
aliases = (),
|
||||
)
|
||||
|
||||
def _fake_unlocked(**kwargs):
|
||||
from PIL import Image as _Image
|
||||
|
||||
return _Image.new("RGB", (8, 8))
|
||||
|
||||
monkeypatch.setattr(backend, "_generate_image_unlocked", _fake_unlocked)
|
||||
_, meta = backend.generate_image_with_metadata(prompt = "x")
|
||||
assert meta["model"] == "secret-flux"
|
||||
assert "/home/alice" not in meta["model"]
|
||||
|
||||
|
||||
def test_release_other_gpu_owners_raises_on_active_training(monkeypatch):
|
||||
"""Round 15 P1 #3: direct backend callers must not bypass the
|
||||
route layer's training-active 409 guard."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue