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(