diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index abb541779e..733a5d1bbe 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -201,6 +201,60 @@ def _expand_existing_local_path(value: str) -> str: return value +def _preflight_diffusers_subfolder_config( + repo: str, + subfolder: str, + hf_token: Optional[str], +) -> None: + """Round 21 P2 #6: also probe ``{subfolder}/config.json``. + + The full-repo preflight at ``_preflight_full_diffusers_repo`` + only proves ``model_index.json`` exists. For GGUF loads the + follow-up ``from_single_file(..., config=effective_base, + subfolder="transformer")`` still needs a matching + ``transformer/config.json`` on the base companion. Without + this second probe a base that has model_index.json but no + transformer config would still unload chat before the load + failed. + """ + if not repo or not subfolder: + return + try: + local = Path(repo).expanduser() + except (OSError, ValueError): + local = None + if local is not None and local.exists(): + config_path = local / subfolder / "config.json" + if not config_path.is_file(): + raise RuntimeError( + f"Diffusion repo '{_display_repo_id(repo)}' is missing " + f"{subfolder}/config.json." + ) + return + if (local is not None and local.is_absolute()) or repo.startswith("~"): + # Local-only path that does not exist -- _preflight_full_diffusers_repo + # already raised for the absent directory, so reaching here means the + # caller is loading a Hub id that just looks like a path. Fall through + # to the network probe. + pass + try: + from huggingface_hub import hf_hub_download as _hf_hub_download + except Exception: + return + try: + _hf_hub_download( + repo_id = repo, + filename = "config.json", + subfolder = subfolder, + token = hf_token, + ) + except Exception as exc: + raise RuntimeError( + f"Could not access diffusion repo '{_display_repo_id(repo)}' " + f"{subfolder}/config.json before unloading the current model." + ) from exc + + def _preflight_full_diffusers_repo(repo: str, hf_token: Optional[str]) -> None: """Prove a full diffusers repo is accessible before any unloads. @@ -904,6 +958,21 @@ class DiffusionBackend: # ``effective_base`` so a bad companion repo is # caught BEFORE chat / export are released. _preflight_full_diffusers_repo(effective_base, hf_token) + # Round 21 P2 #6: the GGUF transformer path also + # consumes ``effective_base`` via + # ``from_single_file(config=effective_base, + # subfolder="transformer")``. A base that has + # ``model_index.json`` but lacks + # ``transformer/config.json`` would pass the + # round-19 preflight and only fail AFTER the chat + # unload. Run the subfolder probe too so the + # second cheap failure mode is also caught early. + if gguf_filename and fam.transformer_class: + _preflight_diffusers_subfolder_config( + effective_base, + "transformer", + hf_token, + ) # Round 20 P1 #2: ``diffusers.GGUFQuantizationConfig`` # imports the ``gguf`` package lazily at construction @@ -1536,17 +1605,18 @@ def _release_chat_backend_for_diffusion() -> None: ) _require_unload(loading) - # Round 19 P1 #2: final sweep using the initial snapshot of - # owned names. Catches races where a name we did not explicitly - # unload (because it appeared in loading_models between the - # snapshot and the unload calls) is still owned after the loop. - remaining_loading = ( - set(getattr(backend, "loading_models", set()) or set()) & owned_names - ) + # Round 21 P1 #5: final sweep without the owned_names filter. + # A concurrent ``/load`` that appeared AFTER the initial + # snapshot was previously ignored, so a chat model that started + # loading during the diffusion handoff slipped through and + # raced the diffusion allocation for VRAM. Treat ANY surviving + # active / loading entry as a failure so the surrounding + # load_model raises and the caller retries. + remaining_loading = set(getattr(backend, "loading_models", set()) or set()) remaining_active = getattr(backend, "active_model_name", None) - if remaining_loading or (remaining_active in owned_names): + if remaining_loading or remaining_active: raise RuntimeError( - "The existing safetensors chat model is still active or loading " + "A safetensors chat model is still active or loading " "after unload; retry before loading a diffusion image model." ) diff --git a/studio/backend/main.py b/studio/backend/main.py index 4ec5e0fc2e..8d67b2f02a 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -333,7 +333,15 @@ def _scrub_validation_obj(value): if isinstance(value, list): return [_scrub_validation_obj(v) for v in value] if isinstance(value, dict): - return {k: _scrub_validation_obj(v) for k, v in value.items()} + # Round 21 P2 #7: pydantic surfaces ``input`` for ``string_type`` + # validation errors verbatim, including dict KEYS like + # ``{"hf_xxxxx": "owner/repo"}``. Scrub string keys too so the + # token does not leak through the 422 response body. + return { + (_scrub_validation_obj(k) if isinstance(k, str) else k): + _scrub_validation_obj(v) + for k, v in value.items() + } return value diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 14c8414d74..5691c3d0d5 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -73,7 +73,12 @@ class LoadRequest(BaseModel): def _no_identifier_control_chars(cls, v, info): return _no_control_chars(v, info.field_name) - @field_validator("model_path") + # Round 21 P1 #1: also reject embedded HF tokens in + # ``gguf_variant``. A caller can pass a variant string like + # ``Q4_K_M-hf_xxxxxxxx`` that flows into log sinks via the + # GGUF resolver path; without this only ``model_path`` was + # protected. + @field_validator("model_path", "gguf_variant") @classmethod def _no_embedded_hf_tokens(cls, v, info): return _reject_embedded_hf_token(v, info.field_name) @@ -171,7 +176,9 @@ class ValidateModelRequest(BaseModel): def _no_identifier_control_chars(cls, v, info): return _no_control_chars(v, info.field_name) - @field_validator("model_path") + # Round 21 P1 #2: extend embedded-token rejection to + # ``gguf_variant`` here too (mirrors LoadRequest). + @field_validator("model_path", "gguf_variant") @classmethod def _no_embedded_hf_tokens(cls, v, info): return _reject_embedded_hf_token(v, info.field_name) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7d18972a91..a4d8f3d5ce 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -502,19 +502,20 @@ async def _release_safetensors_chat_for(workload: str) -> None: ) await _unload_required(loading) - # Round 19 P1 #1: final sweep using the set of names that were - # initially present. Catches races where a model name we did not - # explicitly unload (because it appeared between the snapshot and - # the unload calls) is still in the owned set after the loop. - remaining_loading = ( - set(getattr(inf, "loading_models", set()) or set()) & owned_names - ) + # Round 21 P1 #4: final sweep without the owned_names filter. + # A concurrent ``/load`` that appeared AFTER the initial + # snapshot was previously ignored here, so a chat model that + # started loading during the unload window let the surrounding + # training / export / GGUF / diffusion start anyway. Treat ANY + # surviving active / loading entry as a failure so the caller + # retries rather than racing the new chat load for VRAM. + remaining_loading = set(getattr(inf, "loading_models", set()) or set()) remaining_active = getattr(inf, "active_model_name", None) - if remaining_loading or (remaining_active in owned_names): + if remaining_loading or remaining_active: raise HTTPException( status_code = 503, detail = ( - "The existing safetensors chat model is still active or loading " + "A safetensors chat model is still active or loading " f"after unload; retry before starting {workload}." ), ) @@ -1667,12 +1668,32 @@ async def unload_model( try: # Check if the GGUF backend has this model loaded or is loading it llama_backend = get_llama_cpp_backend() - if llama_backend.is_active and ( - llama_backend.model_identifier == request.model_path + loaded_identifier = getattr(llama_backend, "model_identifier", None) + loading_identifier = getattr(llama_backend, "loading_model_identifier", None) + # Round 21 P1 #3: a GGUF download that has not yet flipped + # ``is_active`` to True (model_identifier still None, + # ``loading_model_identifier`` populated) used to fall + # through to the safetensors branch, which silently + # responded ``status="unloaded"`` while llama-server kept + # downloading. Match on either the loaded OR loading + # identifier so the explicit unload route can actually + # cancel a pending GGUF load. + llama_matches_request = ( + loaded_identifier == request.model_path + or loading_identifier == request.model_path or is_registered_native_path_label( - llama_backend.model_identifier, request.model_path + loaded_identifier, request.model_path ) - or not llama_backend.is_loaded + 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 19 P1 #6: previously this called # ``llama_backend.unload_model()`` and unconditionally diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 85a45caf15..5a94eef120 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -393,10 +393,15 @@ def _install_fake_diffusers(monkeypatch, *, raise_on_pipeline = False): monkeypatch.setitem(sys.modules, "diffusers", fake) # Pretend HF Hub gave us a local file without actually fetching. + # Round 21: accept arbitrary kwargs (round 20 preflight adds + # ``filename="model_index.json"`` and round 21 preflight adds + # ``subfolder="transformer"``) so existing tests that exercise + # the GGUF path do not hit a TypeError from the fake signature. fake_hub = types.ModuleType("huggingface_hub") - fake_hub.hf_hub_download = ( - lambda repo_id, filename, token = None: f"/fake/{repo_id}/{filename}" - ) + def _fake_download(repo_id, filename, token = None, subfolder = None, **_kwargs): + sub = f"{subfolder}/" if subfolder else "" + return f"/fake/{repo_id}/{sub}{filename}" + fake_hub.hf_hub_download = _fake_download monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub) # Force CPU dtype so the test does not need CUDA. @@ -1308,12 +1313,13 @@ def test_load_model_accepts_relative_local_dir(monkeypatch, tmp_path): def _boom(**kwargs): # Round 20 P1 #1 added a base-repo preflight that downloads # the diffusers ``model_index.json`` of the auto-picked - # companion repo BEFORE the chat unload. Allow that call - # through (it would otherwise hit the network) but still - # reject any attempt to download the GGUF itself, which is - # what this test guards. - if kwargs.get("filename") == "model_index.json": - return "/tmp/model_index.json" + # companion repo BEFORE the chat unload. Round 21 P2 #6 + # added a second preflight for ``transformer/config.json`` + # on that same companion. Allow both preflight kinds through + # but still reject any attempt to download the GGUF itself, + # which is what this test guards. + if kwargs.get("filename") in ("model_index.json", "config.json"): + return "/tmp/preflight" raise AssertionError("hf_hub_download must not run for a local dir") fake_hub = SimpleNamespace(hf_hub_download = _boom)