ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix/adjust diffusion: round 7 swap-aware guards + race-free generate for PR #5754

Browse source 
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).
This commit is contained in:
Daniel Han-Chen 2026-05-25 02:08:05 +00:00
commit fa8efafcd8
7 changed files with 240 additions and 95 deletions
Download patch file Download diff file Expand all files Collapse all files

57
studio/backend/core/inference/diffusion.py
View file

@ -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

8
studio/backend/models/inference.py
View file

@ -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

66
studio/backend/routes/export.py
View file

@ -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

49
studio/backend/routes/inference.py
View file

@ -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"),

83
studio/backend/routes/models.py
View file

@ -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",

36
studio/frontend/src/features/images/api.ts
View file

@ -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":<int>`` 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":<number>`` 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(

36
studio/frontend/src/features/images/images-page.tsx
View file

@ -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<string>("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"
/>
<Label>Base diffusers repo (optional)</Label>
<Input
value={customBaseRepo}
onChange={(e) => setCustomBaseRepo(e.target.value)}
placeholder="black-forest-labs/FLUX.2-klein-9B"
/>
<p className="text-xs text-muted-foreground">
{"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)."}
</p>
<Label>Pipeline family (override)</Label>
<Select
value={customFamily}
@ -539,7 +555,15 @@ export function ImagesPage() {
/>
<figcaption className="text-xs text-muted-foreground">
{r.width}x{r.height} - {r.num_inference_steps} steps - g={r.guidance_scale.toFixed(1)}
{r.seed !== null && r.seed !== undefined ? ` - seed ${r.seed}` : ""} -
{/* Prefer seed_str (full uint64 precision) since the
numeric seed gets rounded by JSON.parse above
Number.MAX_SAFE_INTEGER and would otherwise
display a value that does not reproduce. */}
{r.seed_str
? ` - seed ${r.seed_str}`
: r.seed !== null && r.seed !== undefined
? ` - seed ${r.seed}`
: ""} -
{` ${(r.duration_ms / 1000).toFixed(1)}s`}
</figcaption>
</figure>