From 65ea3a2c81d21bd55db311f009b16203317a06ba Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 13:13:19 +0000 Subject: [PATCH] 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. --- pyproject.toml | 5 +- studio/backend/models/data_recipe.py | 12 +++ studio/backend/models/training.py | 6 +- .../backend/requirements/no-torch-runtime.txt | 8 +- .../requirements/single-env/constraints.txt | 7 +- studio/backend/requirements/studio.txt | 20 +---- studio/backend/routes/inference.py | 34 ++++++++- studio/backend/routes/models.py | 28 +++++++ studio/backend/utils/datasets/llm_assist.py | 73 ++++++++++++++----- 9 files changed, 142 insertions(+), 51 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 106dcef3dd..4ccf8583b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index fe607a3f92..06fdcb0963 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -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) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index e0eec81197..234474fd23 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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) diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 117de55c51..76da097a71 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -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. diff --git a/studio/backend/requirements/single-env/constraints.txt b/studio/backend/requirements/single-env/constraints.txt index 0e25124a9f..156f78567e 100644 --- a/studio/backend/requirements/single-env/constraints.txt +++ b/studio/backend/requirements/single-env/constraints.txt @@ -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 diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 3268f41431..b360261e44 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a7e8f74722..cf93153fbe 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index d7ed9aed4b..8b3025e056 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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. diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 0f244a5d56..5299a7ac75 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -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 ... 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(