From fa8efafcd8ab7913c5c79ee3063e10a0356f26f3 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 02:08:05 +0000 Subject: [PATCH] Fix/adjust diffusion: round 7 swap-aware guards + race-free generate for PR #5754 Round 7 reviewer surfaced a handful of swap-window races, fail-open guards, and seed precision mismatches. This commit closes them. Lifecycle / state (P1) * core/inference/diffusion.py: status() now emits active_repo_id, active_base_repo, pending_repo_id, pending_base_repo, and pending_gguf_filename alongside the existing UI-facing fields. During a swap (model A loaded, model B loading) the previous coalesced 'repo_id or pending_repo_id' hid the loading target from delete guards. Splitting the fields lets guards block deletion of either repo currently owned by the backend. * core/inference/diffusion.py: generate_image() now takes _generate_lock BEFORE snapshotting _pipe / _device. Snapshotting outside the lock let a concurrent unload/load clear or replace the backend between the snapshot and the forward, so the freed or swapped pipeline would still run. Symmetric handoffs (P1) * routes/export.py: training-active check now runs BEFORE the chat / inference / diffusion unload helpers, so a 409 does not leave the user's chat session torn down for nothing. Also explicitly fails CLOSED with 503 when is_training_active() raises. * routes/inference.py: _raise_if_training_active now fails closed with 503 when the training backend is importable but its status check raises. The previous best-effort log-and-continue could let chat / diffusion loads collide with unverifiable training. Delete guards (P1) * routes/models.py /delete-cached: chat guard now also blocks when llama-server is_active (i.e. mid-download) and when the inference backend's loading_models set contains the target. Round 7 review #7 flagged that the PR's diffusion-side loading guard had no chat-side parallel, so deleting a chat repo while it was downloading could still race the cache. * routes/models.py /delete-cached: diffusion guard iterates the new active_* + pending_* status fields so a delete during a swap is refused on either repo. * routes/models.py /delete-finetuned: same active_+ pending handling, plus the guard now also refuses deletes of a parent directory that contains the loaded pipeline (round 7 review #6: rm -rf /exports/flux-model/ could unlink model_index.json that the live pipeline is reading via mmap). Seed precision (P2) * models/inference.py + routes/inference.py: DiffusionGenerate- Response now carries seed_str alongside the existing numeric seed. Seeds above Number.MAX_SAFE_INTEGER are rounded by JSON.parse in the browser; seed_str ships full decimal precision for display and reproduction. * frontend/api.ts: DiffusionGenerateResponse types seed_str; images-page.tsx prefers seed_str over seed in the figure caption so the displayed value reproduces the image. * frontend/api.ts: stringifyWithBigInt no longer regex-replaces sentinel strings over the full JSON output. It pulls the seed BigInt out, JSON-serialises the remaining payload, and splices the seed's decimal digits into the resulting object literal at the known position. Avoids the round 7 #10 case where a user-supplied prompt equal to '__bigint__:123' was rewritten into a JSON integer and rejected as a non-string prompt. Custom HF repo (P2) * frontend/images-page.tsx: custom panel now exposes a 'Base diffusers repo' input that maps to DiffusionLoadRequest. base_repo. Required when a private / mirrored GGUF needs a non-default base (e.g. a 9B Klein transformer would otherwise fall back to the 4B base default). --- studio/backend/core/inference/diffusion.py | 57 ++++++++----- studio/backend/models/inference.py | 8 ++ studio/backend/routes/export.py | 66 +++++++++------ studio/backend/routes/inference.py | 49 ++++++++--- studio/backend/routes/models.py | 83 ++++++++++++++----- studio/frontend/src/features/images/api.ts | 36 ++++++-- .../src/features/images/images-page.tsx | 36 ++++++-- 7 files changed, 240 insertions(+), 95 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8ae3a13f0f..6b9e1c1af0 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -320,24 +320,36 @@ class DiffusionBackend: # POSIX layouts) to any authenticated Studio session. with self._lock: gguf_basename = Path(self._gguf_path).name if self._gguf_path else None - # During an in-flight load, expose _pending_* so cache / - # finetuned delete guards can refuse to wipe the repo - # that is mid-download. After the load completes (success - # or failure), the pending fields are cleared so status() - # reverts to publishing only the resident pipeline's id. - effective_repo = self._repo_id or self._pending_repo_id - effective_base = self._base_repo or self._pending_base_repo - effective_gguf = gguf_basename or self._pending_gguf_filename + # Expose BOTH the resident pipeline's id AND the pending + # load target. Delete guards must check both: when model A + # is already loaded and a swap to model B is in flight, + # only checking one would let the user rmtree whichever + # repo the guard ignored. UI-facing ``repo_id`` / + # ``base_repo`` / ``gguf_filename`` still prefer pending + # during a swap so the panel shows the load target the + # user just clicked. + active_repo = self._repo_id + active_base = self._base_repo + pending_repo = self._pending_repo_id if self._loading else None + pending_base = self._pending_base_repo if self._loading else None + pending_gguf = self._pending_gguf_filename if self._loading else None return { "is_loaded": self._pipe is not None, "is_loading": self._loading, - "repo_id": effective_repo, + "repo_id": pending_repo or active_repo, "family": self._family.name if self._family else None, "pipeline_class": ( self._family.pipeline_class if self._family else None ), - "base_repo": effective_base, - "gguf_filename": effective_gguf, + "base_repo": pending_base or active_base, + "gguf_filename": pending_gguf or gguf_basename, + # Guard-facing fields: every repo / path the backend + # owns RIGHT NOW. Delete routes iterate both. + "active_repo_id": active_repo, + "active_base_repo": active_base, + "pending_repo_id": pending_repo, + "pending_base_repo": pending_base, + "pending_gguf_filename": pending_gguf, "device": self._device, "dtype": self._dtype, "loaded_at": self._loaded_at, @@ -703,19 +715,18 @@ class DiffusionBackend: import torch - with self._lock: - if self._pipe is None: - raise RuntimeError("No diffusion model is loaded.") - pipe = self._pipe - device = self._device or "cpu" - - # _generate_lock outside _lock: only one forward at a time, but - # status() / unload() callers do not block on a running forward - # pass. unload_model takes _load_lock + _lock; the pipe object - # itself is kept alive by the local ``pipe`` reference until - # this function returns, so a concurrent unload during forward - # cannot free the weights from under us. + # Take _generate_lock FIRST so a concurrent unload/load that + # observes us holding it will queue behind this generation + # (and `unload_model` then waits its turn before clearing + # state). Snapshotting `self._pipe` outside the lock and then + # taking the lock let a load/unload race in between, so the + # forward could run against a freed or swapped pipeline. with self._generate_lock: + with self._lock: + if self._pipe is None: + raise RuntimeError("No diffusion model is loaded.") + pipe = self._pipe + device = self._device or "cpu" generator = None if seed is not None: # Match the device of the pipeline so determinism holds diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 2767474e5f..08b16d89e7 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1526,7 +1526,15 @@ class DiffusionGenerateResponse(BaseModel): height: int num_inference_steps: int guidance_scale: float + # ``seed`` ships as a JSON number for backwards compatibility with + # the gallery and existing API consumers, but JavaScript rounds + # integers above Number.MAX_SAFE_INTEGER on JSON.parse so seeds + # bigger than 2**53 would render different from the value the + # backend actually used. ``seed_str`` is the exact decimal + # representation; the frontend reads it for reproducibility and + # falls back to ``seed`` when not supplied. seed: Optional[int] = None + seed_str: Optional[str] = None duration_ms: int model: Optional[str] = None family: Optional[str] = None diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index d45659a386..5834767032 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -64,6 +64,48 @@ async def load_checkpoint( # Version switching is handled automatically by the subprocess-based # export backend — no need for ensure_transformers_version() here. + # Symmetric lifecycle guard: refuse to load an export + # checkpoint while training is active so we do not silently + # terminate someone's long-running training job and possibly + # fail the export load on top of that. Mirrors the + # _raise_if_training_active checks in routes/inference.py for + # chat and /images/load. + # Run BEFORE the chat / inference / diffusion unload helpers + # below: otherwise a 409 from this guard would still leave + # the user's chat / inference / diffusion GPU owners freed + # for nothing, which is the asymmetry round 7 review #5 + # flagged. Fail-CLOSED (503) when the training backend is + # importable but its status check raises. + try: + from core.training import get_training_backend # type: ignore + + trn = get_training_backend() + try: + active = trn.is_training_active() + except Exception as e: + logger.warning( + "Could not verify training status before export load: %s", e + ) + raise HTTPException( + status_code = 503, + detail = ( + "Could not verify training status before loading " + "an export checkpoint. Try again." + ), + ) from e + if active: + raise HTTPException( + status_code = 409, + detail = ( + "Training is currently active. Stop the training " + "run before loading an export checkpoint." + ), + ) + except HTTPException: + raise + except Exception as e: + logger.debug("training activity check skipped for export: %s", e) + # Free GPU memory: shut down any running inference/training subprocesses # before loading the export checkpoint (they'd compete for VRAM). try: @@ -94,30 +136,6 @@ async def load_checkpoint( except Exception as e: logger.debug("llama-server unload skipped for export: %s", e) - # Symmetric lifecycle guard: refuse to load an export - # checkpoint while training is active so we do not silently - # terminate someone's long-running training job and possibly - # fail the export load on top of that. Mirrors the - # _raise_if_training_active checks in routes/inference.py for - # chat and /images/load. Fail-closed (503) when the training - # backend can be imported but its status check raises. - try: - from core.training import get_training_backend # type: ignore - - trn = get_training_backend() - if trn.is_training_active(): - raise HTTPException( - status_code = 409, - detail = ( - "Training is currently active. Stop the training " - "run before loading an export checkpoint." - ), - ) - except HTTPException: - raise - except Exception as e: - logger.debug("training activity check skipped for export: %s", e) - # 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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f49d28e095..4af5e9d3e4 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -248,8 +248,15 @@ def _raise_if_training_active(workload: str) -> None: Without this guard the load path would either (a) silently stop a running training run via _release_other_gpu_owners_for_diffusion or (b) double-spend VRAM and OOM both jobs. Both are worse for the - user than a 409 explaining why the request was refused. Best-effort - import so unit-test backends without core.training do not 500. + user than a 409 explaining why the request was refused. + + Failure modes are split: + * ``core.training`` cannot be imported (CI, isolated tests, + custom builds) -> silently return; nothing to protect. + * ``core.training`` is importable but ``get_training_backend()`` + or ``is_training_active()`` raises -> 503 fail-closed. We + cannot verify the GPU is free, so taking the safer route + avoids OOMing an unverifiable training run. """ try: from core.training import get_training_backend # type: ignore @@ -257,18 +264,28 @@ def _raise_if_training_active(workload: str) -> None: return try: trn = get_training_backend() - if trn.is_training_active(): - raise HTTPException( - status_code = 409, - detail = ( - f"Training is currently active. Stop the training run " - f"before loading a {workload} model." - ), - ) - except HTTPException: - raise + active = trn.is_training_active() except Exception as exc: - logger.debug("training activity check skipped: %s", exc) + logger.warning( + "Could not verify training status before %s load: %s", + workload, + exc, + ) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not verify training status before loading the " + f"{workload} model. Try again." + ), + ) from exc + if active: + raise HTTPException( + status_code = 409, + detail = ( + f"Training is currently active. Stop the training run " + f"before loading a {workload} model." + ), + ) def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: @@ -1789,6 +1806,12 @@ async def diffusion_generate( num_inference_steps = payload.num_inference_steps, guidance_scale = payload.guidance_scale, seed = payload.seed, + # str() of a Python int has full precision; JavaScript can + # display it via BigInt without rounding. The numeric ``seed`` + # field above is kept for backwards compatibility with older + # clients but is unsafe to use for seeds above 2**53 on the + # browser side. + seed_str = str(payload.seed) if payload.seed is not None else None, duration_ms = duration_ms, model = status.get("repo_id"), family = status.get("family"), diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index a8a18c081b..c061f42c23 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1973,9 +1973,15 @@ async def delete_finetuned_model( # the merged repo locally, then loaded it via /images/load with a # local path as repo_id). Without this guard /delete-finetuned # could rmtree the directory the diffusion backend is reading from. - # is_loading is also blocked: status() exposes _pending_repo_id / - # _pending_base_repo during the load window so deletes during a - # mid-flight from_pretrained are refused. + # is_loading is also blocked: status() exposes pending_repo_id / + # pending_base_repo during the load window so deletes during a + # mid-flight from_pretrained are refused. During a swap we still + # see the previous load's active_repo_id, so every owned path is + # checked rather than just the UI-facing one. + # Block both DIRECTIONS: + # * loaded path is the same as target (or a parent), and + # * loaded path is a child of target (so the user cannot rmtree + # a parent directory that contains the pipeline's mmap'd file). # Fail-CLOSED on exception (503) like the llama.cpp / safetensors # guards above: an unverifiable diffusion state means we cannot # confirm the target is safe to rmtree. @@ -1985,12 +1991,18 @@ async def delete_finetuned_model( diff_backend = get_diffusion_backend() diff_status = diff_backend.status() if diff_status.get("is_loaded") or diff_status.get("is_loading"): - diff_repo = diff_status.get("repo_id") or "" - diff_base = diff_status.get("base_repo") or "" + candidates: list[str] = [] + for key in ( + "active_repo_id", + "active_base_repo", + "pending_repo_id", + "pending_base_repo", + ): + v = diff_status.get(key) or "" + if v: + candidates.append(v) target_str = str(target_path) - for candidate in (diff_repo, diff_base): - if not candidate: - continue + for candidate in candidates: try: candidate_path = Path(candidate).expanduser() except Exception: @@ -2005,6 +2017,7 @@ async def delete_finetuned_model( candidate_resolved == target_path or str(candidate_resolved) == target_str or _is_path_under(candidate_resolved, target_path) + or _is_path_under(target_path, candidate_resolved) ): raise HTTPException( status_code = 400, @@ -2654,18 +2667,26 @@ async def delete_cached_model( if not _is_valid_repo_id(repo_id): raise HTTPException(status_code = 400, detail = "Invalid repo_id format") - # Check if model is currently loaded + # 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. try: from routes.inference import get_llama_cpp_backend llama_backend = get_llama_cpp_backend() - if llama_backend.is_loaded and llama_backend.model_identifier: - loaded_id = llama_backend.model_identifier.lower() - if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()): - raise HTTPException( - status_code = 400, - detail = "Unload the model before deleting", - ) + loaded_id = (llama_backend.model_identifier or "").lower() + wants = ( + loaded_id == repo_id.lower() + or loaded_id.startswith(repo_id.lower()) + ) + if wants and ( + llama_backend.is_loaded + or getattr(llama_backend, "is_active", False) + ): + raise HTTPException( + status_code = 400, + detail = "Unload the model before deleting", + ) except HTTPException: raise except Exception: @@ -2673,9 +2694,21 @@ async def delete_cached_model( try: inference_backend = get_inference_backend() + loading_models = getattr(inference_backend, "loading_models", set()) or set() + needle = repo_id.lower() + # Loading set holds model identifiers currently being + # downloaded / instantiated; treat them like active loads + # so a delete cannot race a partial mmap. + for loading_model in loading_models: + ml = (loading_model or "").lower() + if ml == needle or ml.startswith(needle): + raise HTTPException( + status_code = 409, + detail = "Cannot delete a model while it is loading", + ) if inference_backend.active_model_name: active = inference_backend.active_model_name.lower() - if active == repo_id.lower() or active.startswith(repo_id.lower()): + if active == needle or active.startswith(needle): raise HTTPException( status_code = 400, detail = "Unload the model before deleting", @@ -2685,7 +2718,7 @@ async def delete_cached_model( except Exception: pass - # Also refuse to delete the cache underlying a loaded or *loading* + # Also refuse to delete the cache underlying a loaded OR loading # diffusion pipeline. The diffusion backend mmap's the GGUF + base # repo weights and continues to read from the cache long after # load; deleting them out from under it would corrupt generation. @@ -2694,6 +2727,9 @@ async def delete_cached_model( # Match exactly on repo_id (case-insensitive) instead of prefix to # avoid blocking unrelated deletes like "org/model" while # "org/model-v2" is loaded. + # During a swap (model A loaded, model B loading), status() + # exposes both via ``active_*`` and ``pending_*`` so we check + # every repo the backend currently owns. # Fail-CLOSED on exception (return 503) like the neighboring # llama.cpp / safetensors guards: we cannot verify whether the # delete is safe, so refuse rather than risk corrupting the @@ -2704,10 +2740,15 @@ async def delete_cached_model( diff_backend = get_diffusion_backend() diff_status = diff_backend.status() if diff_status.get("is_loaded") or diff_status.get("is_loading"): - diff_repo = (diff_status.get("repo_id") or "").lower() - diff_base = (diff_status.get("base_repo") or "").lower() needle = repo_id.lower() - if diff_repo == needle or diff_base == needle: + owned = { + (diff_status.get("active_repo_id") or "").lower(), + (diff_status.get("active_base_repo") or "").lower(), + (diff_status.get("pending_repo_id") or "").lower(), + (diff_status.get("pending_base_repo") or "").lower(), + } + owned.discard("") + if needle in owned: raise HTTPException( status_code = 400, detail = "Unload the diffusion image model before deleting", diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 0a09d971de..b84cfab0dd 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -60,7 +60,14 @@ export interface DiffusionGenerateResponse { height: number; num_inference_steps: number; guidance_scale: number; + /** + * Numeric seed. Safe ONLY for values <= Number.MAX_SAFE_INTEGER. + * For larger seeds, prefer ``seed_str`` (full-precision decimal). + */ seed: number | null; + /** Decimal string with full uint64 precision. Use this for display + * and reproduction when the user pastes the seed back in. */ + seed_str: string | null; duration_ms: number; model: string | null; family: string | null; @@ -95,14 +102,27 @@ export async function unloadDiffusionModel(): Promise<{ is_loaded: boolean }> { ); } -/** JSON.stringify cannot serialise BigInt directly. We only ever - * have BigInts in the seed field, which is an integer; emit the - * literal digits so the server receives a JSON integer rather than - * a string. Pydantic v2 accepts arbitrarily large ints. */ -function stringifyWithBigInt(value: unknown): string { - return JSON.stringify(value, (_, v) => - typeof v === "bigint" ? `__bigint__:${v.toString()}` : v, - ).replace(/"__bigint__:(-?\d+)"/g, "$1"); +/** JSON.stringify cannot serialise BigInt directly. Pull the seed + * BigInt out, stringify the rest of the payload normally, then + * splice the seed's decimal digits back into the JSON literal at the + * exact ``"seed":`` slot. + * + * Avoids the previous regex-over-JSON approach, which could be + * tripped by a user-supplied prompt that exactly matched the + * sentinel string. With this approach the only thing we touch is + * the literal ``"seed":`` substring we wrote ourselves. + */ +function stringifyWithBigInt(value: DiffusionGenerateRequest): string { + const { seed, ...rest } = value; + if (typeof seed !== "bigint") { + return JSON.stringify(value); + } + // Serialise the rest without seed, then inject the seed at the end + // of the object literal as a JSON integer. Strip the trailing "}" + // and re-append once the field is added. + const base = JSON.stringify(rest); + const inner = base.length === 2 /* '{}' */ ? "" : base.slice(1, -1) + ","; + return `{${inner}"seed":${seed.toString()}}`; } export async function generateDiffusionImage( diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 64f06bfbb0..b5aa124fd7 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -110,6 +110,7 @@ export function ImagesPage() { const [presetIndex, setPresetIndex] = useState(0); const [customRepoId, setCustomRepoId] = useState(""); const [customGguf, setCustomGguf] = useState(""); + const [customBaseRepo, setCustomBaseRepo] = useState(""); const [customFamily, setCustomFamily] = useState("auto"); const [useCustom, setUseCustom] = useState(false); const [hfToken, setHfToken] = useState(""); @@ -162,10 +163,14 @@ export function ImagesPage() { : customFamily : preset.family; // Always pass base_repo for curated entries; custom-repo mode - // lets the backend either infer it from the family default or - // (when no GGUF is given) treat the repo as a full diffusers - // checkpoint and call from_pretrained on it directly. - const baseRepo = useCustom ? undefined : preset.base_repo; + // now also lets the user pin one because private / mirrored + // GGUFs (e.g. a 9B klein transformer) would otherwise fall + // back to the family-default 4B base and 500 on load. Empty + // string still falls back to the backend's smart-base / + // repo-id defaults. + const baseRepo = useCustom + ? customBaseRepo.trim() || undefined + : preset.base_repo; if (!repo) { toast.error("Pick a model first"); return; @@ -191,7 +196,7 @@ export function ImagesPage() { } finally { setBusy("idle"); } - }, [useCustom, customRepoId, customGguf, customFamily, preset, hfToken, refreshStatus]); + }, [useCustom, customRepoId, customGguf, customBaseRepo, customFamily, preset, hfToken, refreshStatus]); const handleUnload = useCallback(async () => { setBusy("unloading"); @@ -350,6 +355,17 @@ export function ImagesPage() { onChange={(e) => setCustomGguf(e.target.value)} placeholder="FLUX.2-klein-4B-Q4_K_S.gguf" /> + + setCustomBaseRepo(e.target.value)} + placeholder="black-forest-labs/FLUX.2-klein-9B" + /> +

+ {"Optional. Defaults to the family base. Set this when "} + {"your GGUF expects a non-default base (for example a 9B "} + {"transformer that would otherwise fall back to a 4B base)."} +