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

Fix/adjust diffusion: round 25 P1 batch for PR #5754

Browse source 
Five P1 findings from round 25 reviewer aggregate.

1. routes/datasets.py: /download-progress now reuses the same
   identifier hardening that round 24 added to the model route.
   Token-shaped repo_ids like owner/hf_abcdefghij0123456789 used to
   pass the cheap _is_valid_repo_id regex and end up in warning logs.

2. routes/models.py: extend the llama.cpp cache-delete guard to
   path-ownership matching. A GGUF chat model loaded via a local HF
   snapshot path under models--owner--repo/snapshots used to slip
   past the owner/repo string compare and could be rmtree'd while
   llama-server still mmap'd it. Shares one fail-closed HF cache
   scan and an _owned_cache_path_matches helper with the
   safetensors and diffusion guards (round 25 also dedupes the
   diffusion-specific rescan).

3. routes/models.py: extend the safetensors cache-delete guard the
   same way for safetensors models loaded from local snapshot paths.

4. utils/datasets/llm_assist.py: _run_with_helper and
   _run_multi_pass_advisor now acquire the global llama backend via
   routes.inference.get_llama_cpp_backend instead of instantiating
   a private LlamaCppBackend. _gpu_workload_busy_for_helper already
   ensures the global backend is idle on entry, so this is safe, and
   it makes the helper/advisor load visible to the global delete
   guards (loading_model_identifier and friends).

5. requirements/studio.txt: bump huggingface-hub from 0.36.2 pin to
   1.3.0,<2.0 floor and mirror the no-torch-runtime.txt transformers
   and tokenizers constraints. Fresh installs from studio.txt used
   to resolve transformers 5.x with hub 0.36.2, which crashed
   Flux2KleinPipeline import on missing is_offline_mode the first
   time the user hit /api/inference/images/load.

Includes merge of origin/main (PR #5753 install pin bumps and the
mlx export save_method fix from #5727) so the PR diff stops showing
silent reverts of those landed changes.

Tests: PYTHONPATH=studio/backend pytest
test_diffusion_backend.py test_diffusion_routes.py
test_cached_gguf_routes.py test_llama_cpp_cache_aware_disk_check.py
test_inference_model_validation.py
test_models_get_model_config_case_resolution.py
==> 105 passed locally. The 15 flash-attention test failures and
the test_studio_api SDK suite errors reproduce on HEAD without
these changes (pre-existing, unrelated infrastructure).
This commit is contained in:
Daniel Han-Chen 2026-05-25 12:14:40 +00:00
commit 7b5fe1cf10
4 changed files with 163 additions and 74 deletions
Download patch file Download diff file Expand all files Collapse all files

20
studio/backend/requirements/studio.txt
View file

@ -17,7 +17,25 @@ pyjwt
easydict
addict
# gradio>=4.0.0 # 148 MB - Studio uses React + FastAPI, not Gradio
huggingface-hub==0.36.2
# Round 25 P1 #5: keep the Studio Images dependency set internally
# compatible. ``diffusers>=0.37.0`` ships Flux2KleinPipeline /
# Flux2Pipeline, which transitively import the newer ``transformers``
# (>=4.56) that requires ``huggingface_hub.is_offline_mode`` -- only
# available in ``huggingface_hub>=1.0``. The previous ``==0.36.2``
# pin let fresh installs end up with ``transformers 5.x`` +
# ``huggingface_hub 0.36.2``, which crashed on the first
# ``/api/inference/images/load`` with
# ``Flux2KleinPipeline ... no attribute 'is_offline_mode'``. Bump
# the floor so ``diffusers`` and ``transformers`` resolve into a
# runtime they can actually import.
huggingface-hub>=1.3.0,<2.0
# Mirror the ``transformers`` constraint from
# ``no-torch-runtime.txt``. Without it, the standard install can
# resolve ``transformers 5.4.0+`` which drops Studio-supported
# trainers. ``tokenizers<=0.23.0`` is required because
# ``transformers 4.56..5.3`` declares it explicitly.
tokenizers<=0.23.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.3.0
structlog>=24.1.0
diceware
ddgs

21
studio/backend/routes/datasets.py
View file

@ -68,11 +68,27 @@ if str(backend_path) not in sys.path:
# Import dataset utilities
from utils.datasets import check_dataset_format
from auth.authentication import get_current_subject
from models.inference import _no_control_chars, _reject_embedded_hf_token
router = APIRouter()
logger = get_logger(__name__)
def _validate_logged_identifier(value: str, field_name: str) -> str:
"""Round 25 P1 #1: mirror the helper in routes/models.py so the
dataset ``/download-progress`` route never reaches logger/cache
paths with control characters or embedded HF tokens. Token-shaped
strings like ``owner/hf_abcdefghij0123456789`` would otherwise pass
the cheap ``_is_valid_repo_id`` regex and end up in warning logs.
"""
try:
value = _no_control_chars(value, field_name)
value = _reject_embedded_hf_token(value, field_name)
except ValueError as exc:
raise HTTPException(status_code = 422, detail = str(exc)) from exc
return value
from models.datasets import (
AiAssistMappingRequest,
AiAssistMappingResponse,
@ -370,6 +386,11 @@ async def get_dataset_download_progress(
bytes are observable here. Returns ``cache_path`` so the UI can
show users where the dataset blobs landed on disk.
"""
# Round 25 P1 #1: harden ``repo_id`` before it reaches the
# ``logger.warning`` line at the bottom (or any future log/cache
# path). Matches ``GET /api/models/download-progress`` which
# already validates the same parameter in round 24.
repo_id = _validate_logged_identifier(repo_id, "repo_id")
_empty = {
"downloaded_bytes": 0,
"expected_bytes": 0,

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

@ -2790,6 +2790,73 @@ async def delete_cached_model(
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
# Round 25 P1 #2 / #3: round 15 added a path-ownership check to
# the diffusion guard below, but the llama.cpp and safetensors
# guards still only compared logical ``owner/repo`` strings to
# the loaded/loading identifier. If a chat or safetensors model
# was loaded via a LOCAL HF snapshot path (e.g. through the
# ``/load-local-path`` flow), the loaded identifier is the
# absolute snapshot path -- ``owner/repo`` never appears there,
# the guards passed, and ``DELETE /api/models/delete-cached``
# could rmtree an actively mmap'd snapshot.
#
# Build the HF cache roots for ``repo_id`` ONCE up front and reuse
# them in all three guards (llama, safetensors, diffusion). Failure
# to scan the cache fails CLOSED on the assumption that we cannot
# verify ownership safely; mirrors the diffusion path-scan guard.
needle = repo_id.lower()
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.warning(
"Could not scan HF cache during delete guard preflight: %s",
cache_scan_exc,
)
raise HTTPException(
status_code = 503,
detail = (
"Could not verify cache ownership before deleting. Try again."
),
) from cache_scan_exc
def _owned_cache_path_matches(
value: Optional[str], roots: list[Path]
) -> bool:
"""Return True if ``value`` resolves to (or contains, or is a
child of) any of the HF cache repo roots for the target repo.
Used by the llama / safetensors guards to catch local snapshot
paths the same way the diffusion guard already does.
"""
if not value or not roots:
return False
try:
owned = Path(value).expanduser().resolve()
except Exception:
return False
for root in roots:
try:
if (
owned == root
or _is_path_under(owned, root)
or _is_path_under(root, owned)
):
return True
except Exception:
continue
return False
# Check if model is currently loaded OR loading. is_active and
# not is_loaded means an llama-server download / startup is in
# flight; the cache delete would race the hf_hub_download / mmap.
@ -2800,10 +2867,10 @@ async def delete_cached_model(
from routes.inference import get_llama_cpp_backend
llama_backend = get_llama_cpp_backend()
loaded_id = (llama_backend.model_identifier or "").lower()
loading_id = (
getattr(llama_backend, "loading_model_identifier", None) or ""
).lower()
loaded_id_raw = llama_backend.model_identifier or ""
loaded_id = loaded_id_raw.lower()
loading_id_raw = getattr(llama_backend, "loading_model_identifier", None) or ""
loading_id = loading_id_raw.lower()
loading_variant = (
getattr(llama_backend, "loading_hf_variant", None) or ""
).lower()
@ -2817,9 +2884,14 @@ async def delete_cached_model(
# (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:
# Round 25 P1 #2: also match by HF cache snapshot path so
# local-path GGUF chat loads block the cache delete that
# owns their snapshot.
loading_matches_repo = loading_id == needle or _owned_cache_path_matches(
loading_id_raw, cache_repo_roots
)
if loading_matches_repo:
same_loading_variant = (
not requested_variant
or not loading_variant
@ -2836,7 +2908,10 @@ async def delete_cached_model(
# guard fixed in round 5. Per-variant deletes that target a
# DIFFERENT quant than the loaded one are allowed so the
# llama and diffusion paths stay symmetric (round 14 P1 #7).
if loaded_id == needle and (
loaded_matches_repo = loaded_id == needle or _owned_cache_path_matches(
loaded_id_raw, cache_repo_roots
)
if loaded_matches_repo and (
llama_backend.is_loaded or getattr(llama_backend, "is_active", False)
):
loaded_variant = (getattr(llama_backend, "hf_variant", None) or "").lower()
@ -2864,22 +2939,27 @@ async def delete_cached_model(
try:
inference_backend = get_inference_backend()
loading_models = getattr(inference_backend, "loading_models", set()) or set()
needle = repo_id.lower()
# 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.
# Exact match only on the logical ``owner/repo`` side, but
# also match local snapshot paths (round 25 P1 #3) so a
# safetensors model loaded from a local HF snapshot path
# cannot have its cache rmtree'd out from under it.
for loading_model in loading_models:
ml = (loading_model or "").lower()
if ml == needle:
ml_raw = loading_model or ""
ml = ml_raw.lower()
if ml == needle or _owned_cache_path_matches(ml_raw, cache_repo_roots):
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:
active_model_raw = inference_backend.active_model_name
if active_model_raw:
active = active_model_raw.lower()
if active == needle or _owned_cache_path_matches(
active_model_raw, cache_repo_roots
):
raise HTTPException(
status_code = 400,
detail = "Unload the model before deleting",
@ -2919,46 +2999,11 @@ async def delete_cached_model(
# 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
# 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:
# 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
# ``needle`` and ``cache_repo_roots`` come from the
# preflight scan above; round 25 deduplicated the
# diffusion-specific rescan and now all three guards
# share the same fail-closed cache view.
#
# 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
@ -2969,20 +3014,10 @@ async def delete_cached_model(
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 and _owned_cache_path_matches(
owned_id, cache_repo_roots
):
owned_matches_repo = True
if not owned_matches_repo:
continue
if _variant_delete_is_safe_for_owned_gguf(variant, owned_gguf):

23
studio/backend/utils/datasets/llm_assist.py
View file

@ -245,9 +245,18 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
backend = None
try:
from core.inference.llama_cpp import LlamaCppBackend
# Round 25 P1 #4: use the GLOBAL llama backend instead of a
# private ``LlamaCppBackend()`` instance. The private instance
# was invisible to ``DELETE /api/models/delete-cached`` and the
# other global delete guards because they inspect the singleton
# returned by ``get_llama_cpp_backend()``. A concurrent cache
# delete could rmtree the helper's mid-flight download or
# mmap'd snapshot. ``_gpu_workload_busy_for_helper`` above
# already ensures the global backend is idle before we reach
# here, so taking it over is safe.
from routes.inference import get_llama_cpp_backend
backend = LlamaCppBackend()
backend = get_llama_cpp_backend()
logger.info(f"Loading helper model: {repo} ({variant})")
ok = backend.load_model(
@ -641,9 +650,15 @@ def _run_multi_pass_advisor(
backend = None
try:
from core.inference.llama_cpp import LlamaCppBackend
# Round 25 P1 #4: mirror ``_run_with_helper`` and acquire the
# GLOBAL llama backend so cache-delete and unload guards see
# this advisor load via the singleton's
# ``loading_model_identifier`` / ``model_identifier``. The
# round 23/24 ``_gpu_workload_busy_for_helper`` already
# blocks reach here unless the global llama backend is idle.
from routes.inference import get_llama_cpp_backend
backend = LlamaCppBackend()
backend = get_llama_cpp_backend()
logger.info(f"Loading advisor model: {repo} ({variant})")
t0 = time.monotonic()