From 48740c2664adba8ec18aa2ae91e457350f2a62c1 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 11:44:55 +0000 Subject: [PATCH] Fix/adjust diffusion: round 24 P1 batch for PR #5754 P1 #1: ``_gpu_workload_busy_for_helper`` in ``utils/datasets/llm_assist.py`` now also gates on the GGUF chat backend (llama-server) AND the safetensors chat backend. Round 23 extended it to training + export but missed Chat, so a helper / advisor GGUF could still race a loaded chat model for VRAM. Both checks fail closed when status is unverifiable. P1 #2 / #3 / #4 / #5: re-ordered the route-level GPU-handoff unloads so the diffusion release runs BEFORE the chat releases. A wedged diffusion unload used to fire AFTER chat was already gone, so the user lost both on a single failure. Drop chat last so an earlier failure preserves it. Applied to ``/training/start`` (training.py), ``/export/load`` (export.py), ``/chat/load`` GGUF branch and ``/chat/load`` safetensors branch (routes/inference.py). P1 #7 + P2 #13: ``/delete-finetuned`` body now hardens ``model_path`` and ``gguf_variant`` via the shared ``_validate_logged_identifier`` helper, so control characters and URL-form HF tokens can no longer log-line-smuggle. P1 #8 + #10: ``/delete-cached`` body hardens ``repo_id`` and ``variant`` the same way. P1 #9: ``/download-progress`` ``repo_id`` query parameter is also hardened; the value flows into log lines deep inside ``_get_repo_size_cached`` on lookup failure. P1 #11: ``CheckFormatRequest.dataset_name`` and ``AiAssistMappingRequest.{dataset_name, model_name}`` in ``models/datasets.py`` now apply the same control-char + embedded-HF-token validators, matching every other public request-body model. All 115 diffusion + training-validation + cached_gguf + export + inference model-validation tests pass locally. (P1 #6 native-path-lease enforcement for diffusion local paths and P1 #12 React Compiler frontend lint deferred -- both need focused design / frontend touchups separate from this batch.) --- studio/backend/models/datasets.py | 27 +++++++++++- studio/backend/routes/export.py | 23 ++++------ studio/backend/routes/inference.py | 43 +++++++++--------- studio/backend/routes/models.py | 18 ++++++++ studio/backend/routes/training.py | 20 ++++----- studio/backend/utils/datasets/llm_assist.py | 49 +++++++++++++++++++++ 6 files changed, 131 insertions(+), 49 deletions(-) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index f20d6f2d15..6f9de26939 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -7,7 +7,12 @@ Dataset-related Pydantic models for API requests and responses. from typing import Any, Dict, List, Optional -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field, field_validator, model_validator + +# Round 24 P1 #11: reuse the chat / diffusion / export identifier +# hardening so dataset routes also reject control characters and +# URL-embedded HF tokens in user-controlled identifiers. +from models.inference import _no_control_chars, _reject_embedded_hf_token class CheckFormatRequest(BaseModel): @@ -27,6 +32,16 @@ class CheckFormatRequest(BaseModel): values.setdefault("train_split", values.pop("split")) return values + @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 CheckFormatResponse(BaseModel): """Response for dataset format check""" @@ -57,6 +72,16 @@ class AiAssistMappingRequest(BaseModel): model_name: Optional[str] = None model_type: Optional[str] = None + @field_validator("dataset_name", "model_name") + @classmethod + def _no_identifier_control_chars(cls, v, info): + return _no_control_chars(v, info.field_name) + + @field_validator("dataset_name", "model_name") + @classmethod + def _no_identifier_embedded_hf_tokens(cls, v, info): + return _reject_embedded_hf_token(v, info.field_name) + class AiAssistMappingResponse(BaseModel): """Response from LLM-assisted column classification and conversion advice.""" diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index faa847f040..b6e3a3c5b8 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -150,22 +150,15 @@ async def load_checkpoint( # reviews #1, #8, #9 flagged. from routes.inference import _release_chat_for, _release_diffusion_for - await _release_chat_for("export") - - # Also unload any active diffusion pipeline (Images page); it - # competes for the same GPU and would survive the inference - # shutdown above. is_loading is treated like is_loaded so an - # in-flight load is also waited out (the diffusion unload - # acquires _load_lock + _generate_lock and blocks until the - # current load completes, then unloads). - # Round 17: previously this was a best-effort try/except that - # swallowed every failure with logger.debug, so a wedged - # diffusion backend let the export checkpoint load anyway and - # OOM at first allocation. ``_release_diffusion_for`` is - # strict: it raises HTTPException 503 if status() or - # unload_model() fails, or if the backend remains loaded or - # loading after the unload call. + # Round 24 P1 #3: release diffusion BEFORE chat so a failing + # diffusion unload does not leave the user with no chat + # model loaded. Same reasoning as the training-start flow + # (round 18 P1 #8 / round 24 P1 #2). Earlier rounds kept the + # chat release first because the helper was best-effort; + # now that ``_release_diffusion_for`` is strict it must run + # while chat is still resident so a failure preserves it. await _release_diffusion_for("export load") + await _release_chat_for("export") # load_checkpoint spawns and waits on a subprocess and can take # minutes. Run it in a worker thread so the event loop stays diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ce7ddd1342..a7e8f74722 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1192,9 +1192,14 @@ async def load_model( # corrupt the user's exported artifact). _raise_if_training_active("chat") _raise_if_export_active("chat") - # Drop a settled export checkpoint that is just holding - # GPU memory but is not actively producing output. + # Round 24 P1 #4: release order is now + # export -> diffusion -> safetensors chat (was + # export -> safetensors chat -> diffusion). A wedged + # diffusion unload used to fire AFTER the safetensors + # chat was already gone, so the user lost both. Drop + # the chat last so an earlier failure preserves it. await _release_export_for("GGUF chat") + await _release_diffusion_for("GGUF chat load") llama_backend = get_llama_cpp_backend() # Round 19 P2 #8: previously also called @@ -1206,19 +1211,13 @@ async def load_model( # ``_release_safetensors_chat_for`` below already # handles missing-backend cases as a no-op. - # Unload any safetensors / Unsloth model first to free - # VRAM. Uses the shared helper so we also drain - # ``loading_models`` (round 10 review #4); the inline - # version only checked ``active_model_name`` and let an - # in-flight safetensors load race the new GGUF allocation. + # Unload any safetensors / Unsloth model. Uses the shared + # helper so we also drain ``loading_models`` (round 10 + # review #4); the inline version only checked + # ``active_model_name`` and let an in-flight safetensors + # load race the new GGUF allocation. await _release_safetensors_chat_for("GGUF chat") - # Round 17 P1 #4: route the diffusion unload through the - # strict ``_release_diffusion_for`` helper so a wedged - # diffusion pipeline blocks the GGUF chat load with 503 - # instead of silently double-owning VRAM. - await _release_diffusion_for("GGUF chat load") - # Inherit llama_extra_args from the previous load when the # request omits the field (the chat-settings Apply path # does not round-trip them; explicit [] still clears). @@ -1392,22 +1391,22 @@ async def load_model( # and so we do not silently corrupt an in-flight export. _raise_if_training_active("chat") _raise_if_export_active("chat") - # Drop a settled export checkpoint that is just holding GPU - # memory but is not actively producing output. + # Round 24 P1 #5: release order is now + # export -> diffusion -> llama-chat (was + # export -> llama-chat -> diffusion). A wedged diffusion + # unload used to fire AFTER the GGUF chat was already gone, + # so the user lost both. Drop llama-chat last so an earlier + # failure preserves it. await _release_export_for("safetensors chat") + await _release_diffusion_for("safetensors chat load") backend = get_inference_backend() - # Unload any active or mid-download llama-server first. - # Shared helper so this stays in sync with the GGUF path's + # Unload any active or mid-download llama-server. Shared + # helper so this stays in sync with the GGUF path's # symmetric ``_release_safetensors_chat_for``. await _release_llama_for("safetensors chat") - # Round 17 P1 #5: strict diffusion unload via the shared - # helper so a wedged pipeline blocks the safetensors chat - # load with 503 instead of silently double-owning VRAM. - await _release_diffusion_for("safetensors chat load") - # Export was already dropped above via the shared # ``await _release_export_for("safetensors chat")`` call # (which checks is_export_active() before the destructive diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 57c6d2ff7f..34d8a55aa9 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1876,6 +1876,14 @@ async def delete_finetuned_model( Only paths under Studio's outputs/exports roots are accepted. Exported GGUF entries can delete one quantization variant at a time. """ + # Round 24 P1 #7 + P2 #13: harden both ``model_path`` and + # ``gguf_variant`` for control characters and embedded HF + # tokens, mirroring the chat / diffusion / training request + # validators. Both fields end up in logger.info(...) lines. + model_path = _validate_logged_identifier(model_path, "model_path") + if gguf_variant is not None: + gguf_variant = _validate_logged_identifier(gguf_variant, "gguf_variant") + if source not in {"training", "exported"}: raise HTTPException( status_code = 400, @@ -2506,6 +2514,10 @@ async def get_download_progress( "progress": 0, "cache_path": None, } + # Round 24 P1 #9: ``repo_id`` flows into log lines deep in + # ``_get_repo_size_cached`` on lookup failure, so the same + # hardening the request-body models use applies here too. + repo_id = _validate_logged_identifier(repo_id, "repo_id") try: if not _is_valid_repo_id(repo_id): return _empty @@ -2769,6 +2781,12 @@ async def delete_cached_model( are removed (e.g. ``UD-Q4_K_XL``). Otherwise the entire repo is deleted. Refuses if the model is currently loaded for inference. """ + # Round 24 P1 #8 + #10: harden both ``repo_id`` and ``variant`` + # against control characters / embedded HF tokens before they + # reach logger.info(...) lines or the HF cache scan. + repo_id = _validate_logged_identifier(repo_id, "repo_id") + if variant is not None: + variant = _validate_logged_identifier(variant, "variant") if not _is_valid_repo_id(repo_id): raise HTTPException(status_code = 400, detail = "Invalid repo_id format") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 355394cc10..4ea755207c 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -284,20 +284,18 @@ async def start_training( # model. The previous order (chat -> export) would drop chat # and then refuse training when a wedged idle export raised, # leaving the user with nothing loaded. + # Round 24 P1 #2: same reasoning extended to diffusion -> + # chat. A wedged diffusion unload used to fire AFTER the chat + # backend was already gone, so the user lost both chat and + # diffusion on a single failure mode. Order is now + # export -> diffusion -> chat, with chat as the last drop so + # earlier failures preserve it. await _release_export_for("training") + await _release_diffusion_for("training") await _release_chat_for("training") - # Also unload any loaded diffusion pipeline (Images page); it - # holds the same GPU and would survive the inference shutdown. - # is_loading=True is also handled (unload_model takes - # _load_lock + _generate_lock and waits the in-flight load out). - # Round 17: previously the diffusion unload was best-effort - # (try/except + logger.warning), so a stuck diffusion backend - # would let training start anyway and immediately OOM the - # subprocess. ``_release_diffusion_for`` is strict: it raises - # HTTPException 503 if status() or unload_model() fails, or if - # the backend remains loaded / loading after the unload call. - await _release_diffusion_for("training") + # (Diffusion release moved above chat in round 24 P1 #2; + # the old trailing call was removed to avoid double-unload.) # start_training now spawns a subprocess (non-blocking) success = backend.start_training(job_id = job_id, **training_kwargs) diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 22387a4f99..51f60ace33 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -139,10 +139,59 @@ def _gpu_workload_busy_for_helper() -> bool: blocks the helper instead of double-owning VRAM. Each step fails closed: an unverifiable status counts as busy so the user's primary workload is preserved over the optional helper. + + Round 24 P1 #1: extended to also catch a Chat-backend GPU owner. + The helper GGUF used to run on top of a loaded GGUF chat model + (llama-server) or safetensors chat model and OOM their shared + GPU; mirror the diffusion check by inspecting llama + ``is_loaded`` / ``is_active`` / ``loading_model_identifier`` and + safetensors ``active_model_name`` / ``loading_models``. """ if _diffusion_image_model_busy(): return True + try: + from routes.inference import get_llama_cpp_backend + except Exception: + pass + else: + try: + llama = get_llama_cpp_backend() + if ( + getattr(llama, "is_loaded", False) + or getattr(llama, "is_active", False) + or getattr(llama, "loading_model_identifier", None) + ): + logger.info( + "Skipping helper GGUF while a GGUF chat model is loaded/loading" + ) + return True + except Exception: + logger.info( + "Skipping helper GGUF because llama-server status is unavailable" + ) + return True + + try: + from core.inference import get_inference_backend + except Exception: + pass + else: + try: + inf = get_inference_backend() + active = getattr(inf, "active_model_name", None) + loading = set(getattr(inf, "loading_models", set()) or set()) + if active or loading: + logger.info( + "Skipping helper GGUF while a safetensors chat model is loaded/loading" + ) + return True + except Exception: + logger.info( + "Skipping helper GGUF because safetensors chat status is unavailable" + ) + return True + try: from core.training import get_training_backend except Exception: