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

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

Browse source 
Round 15 reviewer aggregate (logs/review_round15_aggregate.md):

P1 fixes:
- core/inference/llama_cpp.py publishes loading_model_identifier +
  loading_hf_variant AFTER acquiring _serial_load_lock; previously
  a queued second load could overwrite or clear the identifier
  currently in flight, breaking delete-safety and GPU handoff guards.
- routes/models.py /delete-finetuned compares the pending llama
  load against loading_hf_variant (new), not the stale hf_variant
  from the previous loaded model. Without this, a Q4-loaded
  directory loading Q8 would still accept a Q8 delete.
- core/inference/diffusion.py _release_other_gpu_owners_for_diffusion
  now also raises when training is active so direct backend callers
  cannot bypass the route layer's 409 guard. Mirrors the
  export-active check the same helper already enforces.
- routes/models.py /delete-cached diffusion guard compares owned
  diffusion paths against the HF cache root for the target repo
  via _all_hf_cache_scans + _is_path_under. Without this, loading
  from a local models--owner--model/snapshots/<sha> path let the
  cache delete proceed while the snapshot was still mmap'd.
- models/inference.py DiffusionLoadRequest refuses URL-embedded
  hf_xxxxx tokens in repo_id / base_repo at the API boundary, so
  the value never reaches self._repo_id and status() can never
  echo it back to other authenticated sessions.

P2 fixes:
- core/inference/diffusion.py status() routes UI-facing repo_id /
  base_repo through _display_repo_id, which collapses absolute
  local paths to the leaf name (delete guards still see the full
  path via active_*/pending_*).
- routes/inference.py /images/load maps backend RuntimeError that
  reports an export/training conflict to HTTP 409 instead of 400.
- core/inference/diffusion.py detect_family now uses token-boundary
  matching so owner/flux.20-model does not collide with flux.2.

P3 fixes:
- tests/test_diffusion_routes.py drops the partial routes.inference
  module from sys.modules if exec_module() raises, so the real
  ImportError surfaces instead of a misleading AttributeError on
  follow-up tests.

Tests:
- 5 new regression cases (display_repo_id, token-boundary family
  detection, training-active raise from backend helper, embedded HF
  token rejection).
- All 72 diffusion backend + route tests pass.
This commit is contained in:
Daniel Han-Chen 2026-05-25 07:00:20 +00:00
commit 59aa75b8ff
7 changed files with 341 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -201,6 +201,29 @@ def _expand_existing_local_path(value: str) -> str:
return value
def _display_repo_id(value: Any) -> Any:
"""Return a public-facing label for a repo_id / base_repo.
For Hub-style identifiers (``owner/repo``) the value passes
through unchanged so the Images panel and result figcaption
stay informative. Absolute local paths (``/home/me/exports/...``
or ``C:\\Users\\...``) collapse to the leaf name so
``/images/status`` does not leak the user's filesystem layout
to other authenticated browser sessions (round 15 P2 #6). HF
tokens are scrubbed defensively in case they slipped past the
request-side validator.
"""
if not isinstance(value, str) or not value:
return value
try:
candidate = Path(value).expanduser()
if candidate.is_absolute() or candidate.exists():
return candidate.name or value
except (OSError, ValueError):
pass
return _redact_hf_tokens(value)
_HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}")
@ -334,6 +357,32 @@ def detect_family(
# P2 #8).
needle_norm = re.sub(r"[^a-z0-9]+", "-", needle).strip("-")
needle_compact = re.sub(r"[^a-z0-9]+", "", needle)
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``."""
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
# Scan _FAMILIES first (GGUF-supported), then _FULL_REPO_FAMILIES
# so a repo like ``stabilityai/stable-diffusion-xl-base-1.0`` is
# auto-detected as SDXL instead of returning None.
@ -346,10 +395,10 @@ def detect_family(
for e in excludes
):
continue
if fam.name in needle:
if _matches_family_token(fam.name):
return fam
for alias in fam.aliases:
if alias and alias in needle:
if alias and _matches_family_token(alias):
return fam
return None
@ -493,13 +542,20 @@ class DiffusionBackend:
# variants like ``BF16/model.gguf`` (round 14 P1 #4-5).
ui_gguf = pending_gguf or active_gguf
ui_gguf_basename = Path(ui_gguf).name if ui_gguf else None
# UI-facing ``repo_id`` / ``base_repo`` collapse absolute
# local paths to their leaf name so ``/images/status``
# does not leak the user's filesystem layout to other
# authenticated browser sessions (round 15 P2 #6). The
# guard-facing ``active_*`` / ``pending_*`` fields below
# preserve the exact value so delete guards still match
# against the snapshot path.
return {
"is_loaded": self._pipe is not None,
"is_loading": self._loading,
"repo_id": pending_repo or active_repo,
"repo_id": _display_repo_id(pending_repo or active_repo),
"family": ui_family,
"pipeline_class": ui_pipeline_class,
"base_repo": pending_base or active_base,
"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
@ -1242,6 +1298,32 @@ def _release_other_gpu_owners_for_diffusion() -> None:
# helper repeats the local check anyway so that direct backend
# callers (tests, scripts, future routes that forget the
# higher-level guard) cannot still kill an active export.
# Training-active check runs FIRST so direct backend callers
# (tests, scripts, future routes) cannot bypass the route layer's
# 409 by calling ``load_model`` directly while a training run is
# active (round 15 P1 #3). The route layer's
# ``_raise_if_training_active`` still runs ahead of the load to
# surface the conflict as 409; this helper re-raises so direct
# callers see the same RuntimeError the export-active path raises.
try:
from core.training import get_training_backend # type: ignore
except Exception as exc:
logger.debug("training module not importable: %s", exc)
else:
try:
training_active = bool(get_training_backend().is_training_active())
except Exception as exc:
# Unverifiable status -> fail closed (might be active).
raise RuntimeError(
"Could not verify training status before loading a "
"diffusion image model."
) from exc
if training_active:
raise RuntimeError(
"Training is currently active. Stop the training run "
"before loading a diffusion image model."
)
try:
from core.export import get_export_backend # type: ignore
except Exception as exc:

105
studio/backend/core/inference/llama_cpp.py
View file

@ -618,7 +618,12 @@ class LlamaCppBackend:
# ``loading_model_identifier`` so a multi-GB HF download cannot
# have its cache rmtree'd or be ignored by /images/load,
# /training/start, /export/load while it is still resolving.
# ``_loading_hf_variant`` mirrors the same lifetime so the
# per-variant delete guard at routes/models.py:/delete-finetuned
# compares against the NEW variant rather than the previous
# loaded ``hf_variant`` (round 15 P1 #2).
self._loading_model_identifier: Optional[str] = None
self._loading_hf_variant: Optional[str] = None
self._gguf_path: Optional[str] = None
self._hf_repo: Optional[str] = None
self._hf_variant: Optional[str] = None
@ -733,6 +738,20 @@ class LlamaCppBackend:
concurrent /images/load that thinks llama-server is idle."""
return self._loading_model_identifier
@property
def loading_hf_variant(self) -> Optional[str]:
"""``hf_variant`` of the load currently in progress, or None.
Mirrors ``loading_model_identifier``'s lifetime so the
per-variant delete guards (routes/models.py /delete-cached and
/delete-finetuned) can compare against the NEW variant rather
than the previously-loaded one (round 15 P1 #2). Without this,
a directory with Q4 loaded and Q8 loading would still see the
stale Q4 ``hf_variant``, and a Q8 delete would be wrongly
allowed even though Q8 is being downloaded into the same
directory."""
return self._loading_hf_variant
@property
def is_vision(self) -> bool:
return self._is_vision
@ -2616,43 +2635,47 @@ class LlamaCppBackend:
Returns True if server started and health check passed.
"""
# Publish ``_loading_model_identifier`` BEFORE any phase of
# the load can begin and clear it AFTER the load fully settles
# (success or failure, including the duplicate-state fast path
# and every internal early ``return False``). Round 14 P1 #2:
# the prior inline try/finally only wrapped the download, so
# /delete-cached and the cross-workload handoff helpers saw
# the backend as idle once the GGUF bytes had landed but the
# subprocess had not yet spawned. Mark the load as pending
# for the entire duration -- download, metadata read,
# VRAM settle, process spawn, health check, audio probe.
self._loading_model_identifier = model_identifier
try:
# Serialise the whole load so concurrent /load calls never
# leave two llama-server processes alive (#5401 / #5161).
# Does not block /unload, /status, /load-progress.
return self._load_model_impl(
gguf_path = gguf_path,
mmproj_path = mmproj_path,
hf_repo = hf_repo,
hf_variant = hf_variant,
hf_token = hf_token,
model_identifier = model_identifier,
is_vision = is_vision,
n_ctx = n_ctx,
chat_template_override = chat_template_override,
cache_type_kv = cache_type_kv,
speculative_type = speculative_type,
spec_draft_n_max = spec_draft_n_max,
n_threads = n_threads,
n_gpu_layers = n_gpu_layers,
n_parallel = n_parallel,
extra_args = extra_args,
)
finally:
self._loading_model_identifier = None
# Serialise the whole load so concurrent /load calls never
# leave two llama-server processes alive (#5401 / #5161). Does
# not block /unload, /status, /load-progress.
#
# Publish ``_loading_model_identifier`` + ``_loading_hf_variant``
# AFTER acquiring ``_serial_load_lock``. Round 15 P1 #1: the
# previous round 14 version set them outside the lock so a
# second queued ``load_model`` would overwrite or clear the
# identifier of the load currently holding the lock, breaking
# the delete-safety and GPU handoff guards. Cleared in
# ``finally`` so failure / cancellation leaves the pending
# state empty. Round 15 P1 #2 added ``_loading_hf_variant``
# so per-variant delete guards can compare against the
# NEW variant rather than the previous loaded one.
with self._serial_load_lock:
self._loading_model_identifier = model_identifier
self._loading_hf_variant = hf_variant
try:
return self._load_model_impl_locked(
gguf_path = gguf_path,
mmproj_path = mmproj_path,
hf_repo = hf_repo,
hf_variant = hf_variant,
hf_token = hf_token,
model_identifier = model_identifier,
is_vision = is_vision,
n_ctx = n_ctx,
chat_template_override = chat_template_override,
cache_type_kv = cache_type_kv,
speculative_type = speculative_type,
spec_draft_n_max = spec_draft_n_max,
n_threads = n_threads,
n_gpu_layers = n_gpu_layers,
n_parallel = n_parallel,
extra_args = extra_args,
)
finally:
self._loading_model_identifier = None
self._loading_hf_variant = None
def _load_model_impl(
def _load_model_impl_locked(
self,
*,
gguf_path: Optional[str] = None,
@ -2672,11 +2695,11 @@ class LlamaCppBackend:
n_parallel: int = 1,
extra_args: Optional[List[str]] = None,
) -> bool:
"""Internal body of ``load_model``. Kept as a separate method
so ``load_model`` can wrap it in a single try/finally that
publishes ``_loading_model_identifier`` for the WHOLE load
instead of only the download window."""
with self._serial_load_lock:
"""Internal body of ``load_model``. The caller is responsible
for holding ``_serial_load_lock`` and for publishing /
clearing ``_loading_model_identifier`` + ``_loading_hf_variant``
in the surrounding try/finally."""
if True:
# Duplicate /load that raced past the route-level check
# (the first one hadn't published _healthy=True yet). If the
# live server already satisfies this request, do nothing.

31
studio/backend/models/inference.py
View file

@ -1450,6 +1450,32 @@ def _no_control_chars(value: Optional[str], field_name: str) -> Optional[str]:
return value
import re as _re
_EMBEDDED_HF_TOKEN_RE = _re.compile(r"hf_[A-Za-z0-9]{20,}")
def _reject_embedded_hf_token(
value: Optional[str], field_name: str
) -> Optional[str]:
"""Refuse identifiers that contain an embedded ``hf_xxx`` token.
Round 15 P1 #5: ``repo_id`` and ``base_repo`` accept URL-style
strings (``https://hf_token@huggingface.co/owner/repo``). The
token would otherwise be stored in ``self._repo_id`` and echoed
back through ``status()`` to every authenticated browser session.
Log redaction (``_redact_hf_tokens``) covers the logger sink, but
the public status payload also needed to refuse the input. Use
the dedicated ``hf_token`` field for authentication.
"""
if value is not None and _EMBEDDED_HF_TOKEN_RE.search(value):
raise ValueError(
f"{field_name} must not embed a Hugging Face token; "
"pass it via the dedicated hf_token field instead."
)
return value
class DiffusionLoadRequest(BaseModel):
"""Load a diffusion image-generation model.
@ -1495,6 +1521,11 @@ class DiffusionLoadRequest(BaseModel):
def _no_control_chars(cls, v, info):
return _no_control_chars(v, info.field_name)
@field_validator("repo_id", "base_repo")
@classmethod
def _no_embedded_hf_tokens(cls, v, info):
return _reject_embedded_hf_token(v, info.field_name)
# torch.Generator.manual_seed packs into signed int64; values outside
# [-2**63, 2**63 - 1] raise ``Overflow when unpacking long long`` deep

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

@ -1942,7 +1942,21 @@ async def diffusion_load(
)
return JSONResponse(content = status)
except RuntimeError as exc:
raise HTTPException(status_code = 400, detail = str(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.
detail = str(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
except Exception as exc:
logger.exception("Diffusion load failed")
raise HTTPException(status_code = 500, detail = str(exc))

55
studio/backend/routes/models.py
View file

@ -1944,7 +1944,12 @@ async def delete_finetuned_model(
# ``loading_model_identifier`` is set before the download starts
# and cleared after the subprocess settles, so the user cannot
# rmtree the directory llama.cpp is writing into mid-flight.
# Round 15 P1 #2: compare against ``loading_hf_variant`` (the
# variant being downloaded) rather than ``hf_variant`` (the
# PREVIOUS loaded variant, which is stale until the new load
# completes its late-metadata update).
loading_identifier = getattr(llama_backend, "loading_model_identifier", None)
loading_variant = getattr(llama_backend, "loading_hf_variant", None)
if (
loading_identifier
and _loaded_model_matches_deleted_path(
@ -1953,8 +1958,8 @@ async def delete_finetuned_model(
)
and (
not gguf_variant
or not getattr(llama_backend, "hf_variant", None)
or llama_backend.hf_variant.lower() == gguf_variant.lower()
or not loading_variant
or loading_variant.lower() == gguf_variant.lower()
)
):
raise HTTPException(
@ -2852,6 +2857,33 @@ async def delete_cached_model(
diff_status = diff_backend.status()
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
# cache root for this repo. The user may have loaded the
# diffusion model from a local snapshot path under
# ``models--owner--model/snapshots/<sha>``; the string
# ``owner/model`` then never appears in ``owned_id`` and
# the previous string-only check would let the cache
# delete proceed while the snapshot was still mmap'd.
cache_repo_roots: list[Path] = []
try:
for hf_cache in _all_hf_cache_scans():
for repo_info in hf_cache.repos:
if (
repo_info.repo_type == "model"
and repo_info.repo_id.lower() == needle
):
try:
cache_repo_roots.append(
Path(repo_info.repo_path).expanduser().resolve()
)
except Exception:
pass
except Exception as cache_scan_exc:
logger.debug(
"HF cache scan failed during diffusion delete guard: %s",
cache_scan_exc,
)
# Pair each owned repo with the GGUF variant it actually
# owns (active or pending) so a swap in progress does not
# collapse both quants into the pending one (round 13
@ -2859,7 +2891,24 @@ async def delete_cached_model(
# requested variant differs from the variant that owns
# the matched repo.
for owned_id, owned_gguf in _diffusion_owned_targets(diff_status):
if not owned_id or owned_id.lower() != needle:
if not owned_id:
continue
owned_matches_repo = owned_id.lower() == needle
if not owned_matches_repo and cache_repo_roots:
try:
owned_path = Path(owned_id).expanduser().resolve()
except Exception:
owned_path = None
if owned_path is not None:
for repo_root in cache_repo_roots:
if (
owned_path == repo_root
or _is_path_under(owned_path, repo_root)
or _is_path_under(repo_root, owned_path)
):
owned_matches_repo = True
break
if not owned_matches_repo:
continue
if _variant_delete_is_safe_for_owned_gguf(variant, owned_gguf):
continue

56
studio/backend/tests/test_diffusion_backend.py
View file

@ -1473,6 +1473,62 @@ def test_smart_base_repo_uses_windows_leaf_only_already_set_separator_round14():
assert _smart_base_repo(fam, repo) == "black-forest-labs/FLUX.2-klein-9B"
def test_display_repo_id_collapses_absolute_path():
"""Round 15 P2 #6: absolute local paths must NOT leak through
status(). Hub-style repo ids pass through unchanged."""
from core.inference.diffusion import _display_repo_id
# Hub id passes through.
assert (
_display_repo_id("black-forest-labs/FLUX.2-klein-4B")
== "black-forest-labs/FLUX.2-klein-4B"
)
# Absolute local path collapses to leaf.
assert _display_repo_id("/home/alice/exports/private-flux") == "private-flux"
# HF tokens are scrubbed defensively.
leaky = "https://hf_abcdefghij0123456789@huggingface.co/owner/repo"
out = _display_repo_id(leaky)
assert "hf_" not in out
def test_detect_family_rejects_substring_collisions():
"""Round 15 P2 #8: ``flux.20-model`` must NOT match ``flux.2``."""
from core.inference.diffusion import detect_family
# ``flux.20`` is a different number and must not collide with ``flux.2``.
assert detect_family("owner/flux.20-model") is None
# ``stable-diffusion-30`` must not match ``stable-diffusion-3``.
assert detect_family("foo/stable-diffusion-30") is None
# Legitimate ``flux.2`` still matches.
fam = detect_family("black-forest-labs/FLUX.2-dev")
assert fam is not None and fam.name == "flux.2"
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."""
import core.inference.diffusion as d
fake_training_mod = types.ModuleType("core.training")
fake_training_mod.get_training_backend = lambda: SimpleNamespace(
is_training_active = lambda: True
)
monkeypatch.setitem(sys.modules, "core.training", fake_training_mod)
# Ensure export module import does not fail the test before the
# training raise lands.
fake_export_mod = types.ModuleType("core.export")
fake_export_mod.get_export_backend = lambda: SimpleNamespace(
is_export_active = lambda: False,
current_checkpoint = None,
)
monkeypatch.setitem(sys.modules, "core.export", fake_export_mod)
with pytest.raises(RuntimeError) as exc_info:
d._release_other_gpu_owners_for_diffusion()
assert "Training is currently active" in str(exc_info.value)
def test_generate_image_with_metadata_blocks_concurrent_unload(monkeypatch):
"""Round 13 P2 #9: _generate_lock serialises the forward AND the
meta snapshot, so a queued unload cannot wipe state in between."""

38
studio/backend/tests/test_diffusion_routes.py
View file

@ -57,7 +57,15 @@ def _import_inference_module():
assert spec and spec.loader, "could not build spec for routes/inference.py"
module = importlib.util.module_from_spec(spec)
sys.modules["routes.inference"] = module
spec.loader.exec_module(module)
# Round 15 P3 #9: drop the half-initialised module from
# sys.modules if exec_module() raises, otherwise later tests pick
# up the poisoned entry and report a misleading AttributeError
# instead of the original ImportError.
try:
spec.loader.exec_module(module)
except Exception:
sys.modules.pop("routes.inference", None)
raise
return module
@ -249,6 +257,34 @@ def test_unload_clears_state(app_with_stub):
assert r.json()["is_loaded"] is False
def test_load_rejects_embedded_hf_token(app_with_stub):
"""Round 15 P1 #5: URL-embedded ``hf_xxxxx`` tokens in repo_id /
base_repo must be rejected with 422 so they never reach
``self._repo_id`` and get echoed back by ``status()``."""
app, _ = app_with_stub
c = TestClient(app)
r = c.post(
"/api/inference/images/load",
json = {
"repo_id": "https://hf_abcdefghij0123456789@huggingface.co/owner/repo",
},
)
assert r.status_code == 422, r.text
body = r.json()
text = repr(body).lower()
assert "hf_token" in text or "embed" in text
# base_repo is also rejected.
r = c.post(
"/api/inference/images/load",
json = {
"repo_id": "owner/repo",
"gguf_filename": "x.gguf",
"base_repo": "https://hf_abcdefghij0123456789@huggingface.co/base/repo",
},
)
assert r.status_code == 422, r.text
def test_load_rejects_control_chars_in_repo_id(app_with_stub):
"""Newline-laden repo ids must be rejected by Pydantic BEFORE the
log line that echoes them. Catches log-injection from authenticated