From c6c4378f38cc39be26b738d731fd93cbc17ae52f Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 11:20:05 +0000 Subject: [PATCH] Fix/adjust diffusion: round 23 P1+P2 batch for PR #5754 P1 #1 + #2 + #6: extended the chat / diffusion / training identifier hardening to every export-side request model. ExportCommonOptions (parent of ExportMergedModelRequest / ExportBaseModelRequest / ExportLoRAAdapterRequest) now applies _no_control_chars and _reject_embedded_hf_token to repo_id and base_model_id; ExportGGUFRequest gets the same on its repo_id plus a control-char check on quantization_method; and LoadCheckpointRequest validates checkpoint_path. Previously "/api/export/*" accepted newline-smuggled identifiers and URL-form ``hf_xxxxx`` tokens that flowed into log lines. P1 #3 + #4: ``_run_with_helper`` and ``_run_multi_pass_advisor`` now use a shared ``_gpu_workload_busy_for_helper`` that gates on diffusion (round 22 already), training, AND export. The round 22 guard only checked diffusion, so the dataset helper / advisor could still load llama-server on top of an active training run or a resident export checkpoint. Each step fails closed (unverifiable status counts as busy) so the user's primary workload is preserved. P1 #5: PublishDatasetRequest in models/data_recipe.py also applies the identifier hardening to repo_id; the publish path previously accepted control characters and URL-form tokens. P1 #7-10: added _validate_logged_identifier helper to routes/models.py and applied it to the path / query parameter endpoints that flow into logger.info(...) calls -- ``/config/{model_name}``, ``/check-vision/{model_name}``, ``/check-embedding/{model_name}``, ``/gguf-variants``. Mapped the validator's ValueError to HTTP 422 so the client sees the same shape as a Pydantic validation failure. P2 #11 + #12: ``Loading diffusion model %s`` and ``Diffusion load failed for %s`` log lines route ``repo_id`` / ``effective_base`` through ``_display_repo_id`` (collapses absolute local paths to the leaf, still scrubs HF tokens) instead of plain ``_redact_hf_tokens``. The error path was already collapsed in the user-facing 400 / RuntimeError, but the structured-log lines kept the full path. All 97 diffusion + training-validation + related tests pass locally. --- studio/backend/core/inference/diffusion.py | 17 +++++- studio/backend/models/data_recipe.py | 18 +++++- studio/backend/models/export.py | 55 +++++++++++++++++ studio/backend/routes/models.py | 29 ++++++++- studio/backend/utils/datasets/llm_assist.py | 68 ++++++++++++++++++--- 5 files changed, 173 insertions(+), 14 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index be86e47282..f670fb0f05 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -907,13 +907,21 @@ class DiffusionBackend: # Scrub them BEFORE the logger formats the line so the # token never reaches structured-log sinks (round 14 # P2 #9). + # Round 23 P2 #11: ``_redact_hf_tokens`` only scrubs + # ``hf_xxxxx`` substrings, so an absolute local + # path like ``/home/alice/private/FLUX.2-klein-GGUF`` + # used to land in this log line verbatim. Route + # through ``_display_repo_id`` so the leaf is + # logged when the value is a filesystem path, with + # the token-redaction step inside that helper as a + # belt-and-braces defence. logger.info( "Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)", - _redact_hf_tokens(repo_id), + _display_repo_id(repo_id), fam.name, device, dtype, - _redact_hf_tokens(effective_base), + _display_repo_id(effective_base), ) transformer = None @@ -1249,9 +1257,12 @@ class DiffusionBackend: # Use ``logger.error`` with the already-scrubbed # message and exc_info=False so the bearer token # cannot leak through structured logging sinks. + # Round 23 P2 #12: same fix as the start-of-load + # log above. ``_redact_hf_tokens`` alone left + # absolute local repo paths in this failure line. logger.error( "Diffusion load failed for %s: %s", - _redact_hf_tokens(repo_id), + _display_repo_id(repo_id), exc_msg, ) raise RuntimeError( diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index b382ddb3d0..fe607a3f92 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -9,7 +9,13 @@ from __future__ import annotations from typing import Any -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field, field_validator, model_validator + +# Round 23 P1 #5: identifier hardening reused from the chat models +# so /api/data_recipe/publish rejects control characters and +# URL-form ``hf_xxxxx`` tokens in ``repo_id`` before they reach +# log lines or the HF API. +from models.inference import _no_control_chars, _reject_embedded_hf_token class RecipePayload(BaseModel): @@ -60,6 +66,16 @@ class PublishDatasetRequest(BaseModel): description = "Execution artifact path captured by the UI for completed runs", ) + @field_validator("repo_id") + @classmethod + def _no_repo_id_control_chars(cls, v, info): + return _no_control_chars(v, info.field_name) + + @field_validator("repo_id") + @classmethod + def _no_repo_id_embedded_hf_tokens(cls, v, info): + return _reject_embedded_hf_token(v, info.field_name) + class PublishDatasetResponse(BaseModel): success: bool = True diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 86ce2b05bf..a881714a0c 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -10,6 +10,13 @@ from pathlib import Path from pydantic import BaseModel, Field, field_validator from typing import List, Optional, Literal, Dict, Any +# Round 23 P1 #1 / #2 / #6: reuse the chat identifier validators +# so export requests reject newline / tab / control characters and +# URL-form ``hf_xxxxx`` tokens in any user-supplied identifier +# (Hub ``repo_id``, ``base_model_id``, the local +# ``checkpoint_path``) that flows into log lines or HF API calls. +from models.inference import _no_control_chars, _reject_embedded_hf_token + def _validate_save_directory(value: str) -> str: """Reject save_directory values that escape the export root.""" @@ -54,6 +61,19 @@ class LoadCheckpointRequest(BaseModel): description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.", ) + # Round 23 P1 #6: ``checkpoint_path`` is logged verbatim by the + # export route. Apply the same control-char + embedded-token + # rejection the chat / diffusion / training request models use. + @field_validator("checkpoint_path") + @classmethod + def _no_checkpoint_control_chars(cls, v, info): + return _no_control_chars(v, info.field_name) + + @field_validator("checkpoint_path") + @classmethod + def _no_checkpoint_embedded_hf_tokens(cls, v, info): + return _reject_embedded_hf_token(v, info.field_name) + class ExportStatusResponse(BaseModel): """Current export backend status.""" @@ -117,6 +137,20 @@ class ExportCommonOptions(BaseModel): description = "HuggingFace model ID of the base model (for model card metadata)", ) + # Round 23 P1 #1: ``repo_id`` (Hub destination) and + # ``base_model_id`` (model card metadata) both feed log lines + # and the HF API. Reject control characters and URL-form + # ``hf_xxxxx`` tokens before they reach those sinks. + @field_validator("repo_id", "base_model_id") + @classmethod + def _no_identifier_control_chars(cls, v, info): + return _no_control_chars(v, info.field_name) + + @field_validator("repo_id", "base_model_id") + @classmethod + def _no_identifier_embedded_hf_tokens(cls, v, info): + return _reject_embedded_hf_token(v, info.field_name) + class ExportMergedModelRequest(ExportCommonOptions): """Request for exporting a merged PEFT model.""" @@ -163,6 +197,27 @@ class ExportGGUFRequest(BaseModel): description = "Hugging Face token for GGUF upload", ) + # Round 23 P1 #2: GGUF export endpoint defines its own + # ``repo_id`` (does not inherit from ExportCommonOptions), so + # the chat-style hardening needs to be applied here separately. + # ``quantization_method`` is forwarded to the export worker + # command line, so it gets the control-char check too even + # though it does not normally carry tokens. + @field_validator("repo_id") + @classmethod + def _no_repo_id_control_chars(cls, v, info): + return _no_control_chars(v, info.field_name) + + @field_validator("repo_id") + @classmethod + def _no_repo_id_embedded_hf_tokens(cls, v, info): + return _reject_embedded_hf_token(v, info.field_name) + + @field_validator("quantization_method") + @classmethod + def _no_quantization_control_chars(cls, v, info): + return _no_control_chars(v, info.field_name) + class ExportLoRAAdapterRequest(ExportCommonOptions): """Request for exporting only the LoRA adapter (not merged).""" diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1ebf525f56..57c6d2ff7f 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -134,11 +134,30 @@ from models.responses import ( VisionCheckResponse, EmbeddingCheckResponse, ) +from models.inference import _no_control_chars, _reject_embedded_hf_token router = APIRouter() logger = get_logger(__name__) +def _validate_logged_identifier(value: str, field_name: str) -> str: + """Round 23 P1 #7 / #8 / #9 / #10: path / query parameters that + flow into ``logger.info("... %s", value)`` lines were the last + unguarded entry points. Newline / tab / control characters let + a caller smuggle forged log entries; URL-form ``hf_xxxxx`` + tokens would leak into structured-log sinks. Mirror the + request-body validators by running both checks here and + mapping the validator's ``ValueError`` to HTTP 422 so the + client sees the same shape as a Pydantic validation failure. + """ + try: + value = _no_control_chars(value, field_name) + value = _reject_embedded_hf_token(value, field_name) + except ValueError as exc: + raise HTTPException(status_code = 422, detail = str(exc)) from exc + return value + + def derive_model_type( is_vision: bool, audio_type: Optional[str], is_embedding: bool = False ) -> ModelType: @@ -1571,6 +1590,7 @@ async def get_model_config( This endpoint wraps the backend load_model_defaults function. """ + model_name = _validate_logged_identifier(model_name, "model_name") try: if not is_local_path(model_name): resolved = resolve_cached_repo_id_case(model_name) @@ -1580,7 +1600,11 @@ async def get_model_config( resolved, model_name, ) - model_name = resolved + # Round 23 P1 #7: re-validate the cache-resolved value + # (case-only resolver should be a no-op for these + # checks, but defend in depth in case the resolver + # ever broadens its match heuristic). + model_name = _validate_logged_identifier(resolved, "model_name") logger.info(f"Getting model config for: {model_name}") from utils.models.model_config import detect_audio_type @@ -2220,6 +2244,7 @@ async def check_vision_model( This endpoint wraps the backend is_vision_model function. """ + model_name = _validate_logged_identifier(model_name, "model_name") try: logger.info(f"Checking if vision model: {model_name}") is_vision = is_vision_model(model_name) @@ -2248,6 +2273,7 @@ async def check_embedding_model( This endpoint wraps the backend is_embedding_model function. """ + model_name = _validate_logged_identifier(model_name, "model_name") try: logger.info(f"Checking if embedding model: {model_name}") is_embedding = is_embedding_model(model_name, hf_token = hf_token) @@ -2285,6 +2311,7 @@ async def get_gguf_variants( with file sizes, whether the model supports vision, and the recommended default variant. """ + repo_id = _validate_logged_identifier(repo_id, "repo_id") try: from utils.models.model_config import is_local_path, list_local_gguf_variants diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 8aad3230e0..c0bedc049d 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -131,6 +131,57 @@ def _diffusion_image_model_busy() -> bool: return bool(status.get("is_loaded") or status.get("is_loading")) +def _gpu_workload_busy_for_helper() -> bool: + """Round 23 P1 #3 / #4: the diffusion-only guard from round 22 + let the helper / advisor GGUF run on top of a live training run + or a resident export checkpoint. Extend the busy check to those + workloads too so any GPU owner (Images, Training, Export) + 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. + """ + if _diffusion_image_model_busy(): + return True + + try: + from core.training import get_training_backend + except Exception: + pass + else: + try: + if get_training_backend().is_training_active(): + logger.info( + "Skipping helper GGUF while training is active" + ) + return True + except Exception: + logger.info( + "Skipping helper GGUF because training status is unavailable" + ) + return True + + try: + from core.export import get_export_backend + except Exception: + return False + + try: + exp = get_export_backend() + is_active = getattr(exp, "is_export_active", None) + if (is_active and is_active()) or getattr( + exp, "current_checkpoint", None + ): + logger.info("Skipping helper GGUF while export owns the GPU") + return True + except Exception: + logger.info( + "Skipping helper GGUF because export status is unavailable" + ) + return True + + return False + + def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: """ Load helper model, run one chat completion, unload. @@ -140,10 +191,10 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"): return None - if _diffusion_image_model_busy(): - logger.info( - "Skipping helper GGUF while a diffusion image model is loaded/loading" - ) + # Round 23 P1 #3: round 22 only guarded against a busy + # diffusion pipeline. Training / export own the same GPU too, + # so use the broader helper that gates on all three workloads. + if _gpu_workload_busy_for_helper(): return None repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) @@ -536,11 +587,10 @@ def _run_multi_pass_advisor( if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"): return None - # Round 22 P1 #3: same diffusion-busy guard as ``_run_with_helper``. - if _diffusion_image_model_busy(): - logger.info( - "Skipping advisor GGUF while a diffusion image model is loaded/loading" - ) + # Round 23 P1 #4: extend the round 22 diffusion-only check to + # training + export so the advisor cannot race the user's + # active workload for GPU memory. + if _gpu_workload_busy_for_helper(): return None repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)