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

Twelve P1 findings from round 26 reviewer aggregate, plus the CI
revert of round 25 P1 #5 to a less invasive location.

1. requirements/studio.txt + requirements/single-env/constraints.txt:
   revert the round 25 huggingface-hub bump (broke Studio Update CI,
   Mac Studio Update CI, Mac Studio UI CI, Studio UI CI all with
   ResolutionImpossible against transformers==4.57.6 which requires
   hub<1.0). Standard install path stays on the well-tested 4.57.6 +
   0.36.2 + trl 0.23.1 trio.

2. requirements/no-torch-runtime.txt + pyproject.toml
   [huggingfacenotorch]: bump huggingface_hub floor from >=0.34.0 to
   >=1.3.0,<2.0 -- this is where the actual transformers 5.x +
   hub 0.36.2 broken combo can land because the file installs
   --no-deps. transformers 5.x calls hub.is_offline_mode which only
   exists in hub 1.x.

3. utils/datasets/llm_assist.py: revert round 25 P1 #4 (helper/advisor
   sharing the global llama backend) which introduced three
   regressions: a chat-evict load race after the busy precheck, a
   finally-block that could unload a user chat model, and an
   identifier mismatch the delete guard could not canonicalize. Go
   back to PRIVATE LlamaCppBackend instances and expose the active
   helper/advisor repos through a new thread-safe registry
   (helper_advisor_owns_repo / _register_helper_advisor_repo /
   _unregister_helper_advisor_repo) so DELETE /api/models/delete-cached
   can still block the rmtree.

4. routes/models.py delete_cached_model: check the new helper/advisor
   registry up front and 409 if a helper/advisor still owns the
   target repo. Closes round 26 P1 #13 and #14 (helper/advisor
   identifiers were prefixed and would never equal the raw repo id).

5. routes/models.py get_lora_base_model: validate lora_path with
   _validate_logged_identifier before it is reflected in 404 detail
   and error logs (round 26 P1 #12).

6. routes/inference.py /unload: round 21 P1 #3 added a "or not
   is_loaded" fallback that let an unload of owner/B cancel a pending
   llama load of owner/A. Replace it with a narrow
   llama_is_starting_without_identifier branch that only fires when
   llama-server is mid-startup with neither identifier set (round 26
   P1 #5).

7. routes/inference.py /unload: poll loading_model_identifier for up
   to 5 s after asyncio.to_thread(unload_model) so a legitimate
   pending-load cancel does not 503 because the load thread has not
   yet observed _cancel_event in its finally (round 26 P2 #15).

8. models/training.py TrainingStartRequest: extend identifier
   hardening to hf_dataset, subset, train_split, eval_split. Round 22
   only guarded model_name (round 26 P1 #10).

9. models/data_recipe.py SeedInspectRequest: add _no_control_chars +
   _reject_embedded_hf_token field_validators on dataset_name (round
   26 P1 #11).

Tests: 105 targeted (diffusion + cached_gguf + llama_cpp_cache +
inference_model_validation + models_get_model_config) and 1768
broader backend tests pass locally. Pre-existing
test_desktop_auth.py, test_studio_api.py, and
test_training_worker_flash_attn.py failures reproduce on HEAD
without these changes.
This commit is contained in:
Daniel Han-Chen 2026-05-25 13:13:19 +00:00
commit 65ea3a2c81
9 changed files with 143 additions and 52 deletions

View file

@ -81,7 +81,10 @@ huggingfacenotorch = [
"datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0",
"accelerate>=0.34.1",
"peft>=0.18.0,!=0.11.0",
"huggingface_hub>=0.34.0",
# Round 26 P1 #9: floor at 1.3.0 because the diffusion stack below
# pulls transformers 5.x which calls hub.is_offline_mode (hub 1.x).
# Keep <2.0 to avoid any future hub ABI break.
"huggingface_hub>=1.3.0,<2.0",
"hf_transfer",
# Studio Images page depends on Flux2KleinPipeline /
# Flux2Pipeline, both shipped in diffusers>=0.37.0. Floor was

View file

@ -90,6 +90,18 @@ class SeedInspectRequest(BaseModel):
split: str | None = "train"
preview_size: int = Field(default = 10, ge = 1, le = 50)
# Round 26 P1 #11: dataset_name reaches HF + log/echo paths, so
# mirror the hardening other dataset request models already do.
@field_validator("dataset_name")
@classmethod
def _no_dataset_name_control_chars(cls, v, info):
return _no_control_chars(v, info.field_name)
@field_validator("dataset_name")
@classmethod
def _no_dataset_name_embedded_hf_tokens(cls, v, info):
return _reject_embedded_hf_token(v, info.field_name)
class SeedInspectUploadRequest(BaseModel):
# Legacy single-file flow (mutually exclusive with file_ids)

View file

@ -60,12 +60,14 @@ class TrainingStartRequest(BaseModel):
# Round 22 P1 #1: identifier hardening (round 5 / 15 / 20 / 21
# extended these to chat + diffusion request models; training
# was the last unguarded entry point).
@field_validator("model_name")
# Round 26 P1 #10: hf_dataset / subset / train_split / eval_split
# are reflected in status + error messages, harden them too.
@field_validator("model_name", "hf_dataset", "subset", "train_split", "eval_split")
@classmethod
def _no_model_name_control_chars(cls, v, info):
return _no_control_chars(v, info.field_name)
@field_validator("model_name")
@field_validator("model_name", "hf_dataset")
@classmethod
def _no_model_name_embedded_hf_tokens(cls, v, info):
return _reject_embedded_hf_token(v, info.field_name)

View file

@ -43,7 +43,13 @@ safetensors>=0.4.3
datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0
accelerate>=0.34.1
peft>=0.18.0,!=0.11.0
huggingface_hub>=0.34.0
# Round 26 P1 #8: floor at 1.3.0 because transformers 5.x (allowed by
# the range below) calls huggingface_hub.is_offline_mode, only present
# in hub 1.x. Under --no-deps the resolver does not enforce this
# transitively, so a pre-existing 0.36.2 used to be kept and the next
# `from transformers import AutoConfig` raised ImportError. Upper bound
# <2.0 keeps us off any future ABI break.
huggingface_hub>=1.3.0,<2.0
hf_transfer
# Floor 0.37.0 introduces Flux2KleinPipeline + Flux2Pipeline which the
# Studio Images page imports for the default curated picker.

View file

@ -2,12 +2,7 @@
# Keep compatible with unsloth transformers bounds.
transformers==4.57.6
trl==0.23.1
# Round 25 P1 #5 follow-up: bumped from 0.36.2 because studio.txt now
# requires >=1.3.0,<2.0 (Flux2KleinPipeline transitively imports
# transformers 5.x which needs huggingface_hub.is_offline_mode,
# introduced in hub 1.x). 1.8.0 matches the explicit pin used by
# studio/setup.sh / setup.ps1 for the t5 sub-envs.
huggingface-hub==1.8.0
huggingface-hub==0.36.2
# Studio stack
datasets==4.3.0

View file

@ -17,25 +17,7 @@ pyjwt
easydict
addict
# gradio>=4.0.0 # 148 MB - Studio uses React + FastAPI, not Gradio
# 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
huggingface-hub==0.36.2
structlog>=24.1.0
diceware
ddgs

View file

@ -1683,9 +1683,27 @@ async def unload_model(
or is_registered_native_path_label(loaded_identifier, request.model_path)
or is_registered_native_path_label(loading_identifier, request.model_path)
)
if (getattr(llama_backend, "is_active", False) or loading_identifier) and (
llama_matches_request or not getattr(llama_backend, "is_loaded", False)
):
# Round 26 P1 #5: the previous ``or not is_loaded`` fallback
# let an unload of ``owner/B`` cancel a pending llama download
# of ``owner/A`` and silently leave safetensors ``owner/B``
# alive. Only enter the llama branch when the request actually
# matches the loaded/loading identifier, OR when llama-server
# is starting up without any identifier yet (the original
# narrow case we wanted to catch).
llama_is_starting_without_identifier = (
getattr(llama_backend, "is_active", False)
and not getattr(llama_backend, "is_loaded", False)
and not loaded_identifier
and not loading_identifier
)
should_unload_llama = (
llama_matches_request
and (
getattr(llama_backend, "is_active", False)
or loading_identifier
)
) or llama_is_starting_without_identifier
if should_unload_llama:
# Round 19 P1 #6: previously this called
# ``llama_backend.unload_model()`` and unconditionally
# returned ``status="unloaded"`` even when the subprocess
@ -1694,6 +1712,16 @@ async def unload_model(
# still resident. Treat ``False`` / leftover state as a
# 503 so the user retries.
ok = await asyncio.to_thread(llama_backend.unload_model)
# Round 26 P2 #15: explicit cancel of a pending GGUF load
# leaves loading_model_identifier set briefly until the
# load thread observes _cancel_event in its finally. Wait
# up to 5s so a legitimate cancel does not 503.
deadline = time.monotonic() + 5.0
while (
getattr(llama_backend, "loading_model_identifier", None)
and time.monotonic() < deadline
):
await asyncio.sleep(0.1)
if (
ok is False
or getattr(llama_backend, "is_loaded", False)

View file

@ -2219,6 +2219,9 @@ async def get_lora_base_model(
This endpoint wraps the backend get_base_model_from_lora function.
"""
# Round 26 P1 #12: lora_path is echoed back in 404 detail and logs;
# harden it the same way other reflected identifiers are.
lora_path = _validate_logged_identifier(lora_path, "lora_path")
try:
base_model = get_base_model_from_lora(lora_path)
@ -2853,6 +2856,31 @@ async def delete_cached_model(
continue
return False
# Round 26 P1 #13 / #14: helper/advisor GGUF loads run on a
# PRIVATE LlamaCppBackend, so the global backend below cannot see
# them. utils/datasets/llm_assist.py publishes the active repo
# via helper_advisor_owns_repo() for exactly this guard. Fail
# closed on the variant question (block any variant of the repo)
# because helper/advisor flows do not pass a variant through.
try:
from utils.datasets.llm_assist import helper_advisor_owns_repo
if helper_advisor_owns_repo(repo_id):
raise HTTPException(
status_code = 409,
detail = "Cannot delete a model while AI Assist is using it",
)
except HTTPException:
raise
except Exception as e:
logger.warning(
"Could not check helper/advisor backend status before cache delete: %s", e
)
raise HTTPException(
status_code = 503,
detail = "Could not verify AI Assist load status before deleting cache",
) from e
# 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.

View file

@ -18,6 +18,7 @@ import logging
import os
import re
import textwrap
import threading
import time
from itertools import islice
from typing import Any, Optional
@ -31,6 +32,41 @@ DEFAULT_HELPER_MODEL_VARIANT = "UD-Q4_K_XL"
README_MAX_CHARS = 1500
# Round 26 P1 #13 / #14: helper/advisor run on PRIVATE LlamaCppBackend
# instances (round 25 P1 #4 briefly used the global singleton, which
# caused chat-evict races and finally-eviction bugs and still left
# delete-cache blind because helper/advisor publish prefixed
# identifiers the guard could not match). Expose loading repo ids
# through a thread-safe set so DELETE /api/models/delete-cached can
# block while a helper or advisor still owns the cache.
_HELPER_ADVISOR_ACTIVE_REPOS: set[str] = set()
_HELPER_ADVISOR_LOCK = threading.Lock()
def helper_advisor_owns_repo(repo_id: str) -> bool:
"""Return True if any helper/advisor load currently owns this
HF repo id. Comparison is case-insensitive to match the chat
backend's lowercased needle."""
if not repo_id:
return False
needle = repo_id.lower()
with _HELPER_ADVISOR_LOCK:
return needle in _HELPER_ADVISOR_ACTIVE_REPOS
def _register_helper_advisor_repo(repo_id: str) -> None:
if not repo_id:
return
with _HELPER_ADVISOR_LOCK:
_HELPER_ADVISOR_ACTIVE_REPOS.add(repo_id.lower())
def _unregister_helper_advisor_repo(repo_id: str) -> None:
if not repo_id:
return
with _HELPER_ADVISOR_LOCK:
_HELPER_ADVISOR_ACTIVE_REPOS.discard(repo_id.lower())
def _strip_think_tags(text: str) -> str:
"""Strip <think>...</think> reasoning blocks emitted by some models.
@ -244,19 +280,17 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
)
backend = None
_register_helper_advisor_repo(repo)
try:
# 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
# Round 26 P1 #1 / #3 / #13 / #14: use a PRIVATE backend so the
# helper can never preempt or be preempted by the user's
# chat backend and cannot accidentally unload it in finally.
# The active repo is published via _register_helper_advisor_repo
# above so DELETE /api/models/delete-cached can still block the
# cache rmtree while the helper is downloading or mmap'ing.
from core.inference.llama_cpp import LlamaCppBackend
backend = get_llama_cpp_backend()
backend = LlamaCppBackend()
logger.info(f"Loading helper model: {repo} ({variant})")
ok = backend.load_model(
@ -305,6 +339,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
logger.info("Helper model unloaded")
except Exception:
pass
_unregister_helper_advisor_repo(repo)
# ─── Public API ───────────────────────────────────────────────────────
@ -649,16 +684,15 @@ def _run_multi_pass_advisor(
)
backend = None
_register_helper_advisor_repo(repo)
try:
# 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
# Round 26 P1 #2 / #4 / #13 / #14: mirror ``_run_with_helper``
# and use a PRIVATE backend. Round 25's global-backend swap
# introduced chat-evict races and finally-eviction bugs.
# The registry above keeps delete-cache safe.
from core.inference.llama_cpp import LlamaCppBackend
backend = get_llama_cpp_backend()
backend = LlamaCppBackend()
logger.info(f"Loading advisor model: {repo} ({variant})")
t0 = time.monotonic()
@ -990,6 +1024,7 @@ def _run_multi_pass_advisor(
logger.info("Advisor model unloaded")
except Exception:
pass
_unregister_helper_advisor_repo(repo)
def llm_conversion_advisor(