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

Fix/adjust diffusion: round 27 P1 + P2 batch for PR #5754

Browse source 
Round 27 findings (Opus parallel concurrency + frontend reviews).

Backend P1 fixes:

1. utils/datasets/llm_assist.py: the round 26 helper/advisor active
   registry used a plain set, so two concurrent helper / advisor
   loads of the same DEFAULT_HELPER_MODEL_REPO would both
   set.add() (no-op the second time) and then the first finally
   set.discard() would underflow the registration while the second
   call was still mmap'ing the GGUF. Switch to a Counter with
   proper refcount increment/decrement so the repo stays registered
   until the last user releases it.

2. routes/inference.py _release_chat_for and
   core/inference/diffusion.py _release_chat_backend_for_diffusion:
   helper/advisor GGUF runs on a PRIVATE LlamaCppBackend (round 26
   P1 #1), so the global llama checks below could not see them.
   A user-driven /training/start, /export/load-checkpoint, or
   /images/load would skip the unload and allocate FLUX VRAM on top
   of the helper's resident weights, OOMing on 16-24 GB consumer
   GPUs. Both release paths now consult helper_advisor_busy() and
   fail 503 (or RuntimeError for the in-backend path) so the user
   retries instead of double-owning VRAM.

Frontend P2 fixes:

3. studio/frontend/src/features/images/images-page.tsx: handleUnload
   now calls refreshStatus() in the catch path so a partial unload
   (503 from the backend) does not leave the UI showing a stale
   "Loaded:" label. Matches the handleLoad pattern.

4. images-page.tsx: when status.is_loading is true, auto-poll
   refreshStatus every 2 s so the user sees real progress instead
   of a frozen "Loading..." label until they manually click Refresh.

5. images-page.tsx: aria-label="Inference steps" / "Guidance scale"
   on the two sliders so screen readers can announce them.

6. images-page.tsx: defensive (r.guidance_scale ?? 0).toFixed(1)
   in the results caption so a future backend that serialises
   NaN/None for guidance does not throw at render.

Tests: 105 targeted (diffusion + cached_gguf + inference_validation)
and 1768 broader backend tests pass locally. Frontend
`npm run typecheck` passes.
This commit is contained in:
Daniel Han-Chen 2026-05-25 13:39:40 +00:00
commit 6c528fb013
4 changed files with 78 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1516,6 +1516,20 @@ def _release_chat_backend_for_diffusion() -> None:
diffusion ``load_model`` bails out instead of double-owning VRAM
(round 17 P1 #2).
"""
# Round 27 P1 #2: helper / advisor GGUF loads run on a PRIVATE
# LlamaCppBackend so the global llama check below cannot see them.
# Refuse the diffusion handoff while a helper / advisor still owns
# its private backend so we do not allocate FLUX VRAM on top.
try:
from utils.datasets.llm_assist import helper_advisor_busy
except Exception:
pass
else:
if helper_advisor_busy():
raise RuntimeError(
"AI Assist (helper / advisor GGUF) is still using the GPU. "
"Wait for it to finish before loading a diffusion image model."
)
# 1. GGUF chat backend (llama-server subprocess). We unload when
# EITHER is_loaded is True (resident model) OR is_active is
# True (mid-download / startup) OR loading_model_identifier is

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

@ -531,6 +531,24 @@ async def _release_chat_for(workload: str) -> None:
start. Conversely, the standard chat-load path releases only
the llama side.
"""
# Round 27 P1 #2: helper / advisor GGUF loads run on a PRIVATE
# LlamaCppBackend (round 26 P1 #1) so the global llama checks
# below do not see them. Refuse the handoff while a helper /
# advisor still owns its private backend so a new GPU workload
# does not allocate on top of helper VRAM and OOM.
try:
from utils.datasets.llm_assist import helper_advisor_busy
except Exception:
pass
else:
if helper_advisor_busy():
raise HTTPException(
status_code = 503,
detail = (
f"AI Assist (helper / advisor GGUF) is still using the "
f"GPU. Wait for it to finish before starting {workload}."
),
)
await _release_llama_for(workload)
await _release_safetensors_chat_for(workload)

31
studio/backend/utils/datasets/llm_assist.py
View file

@ -20,6 +20,7 @@ import re
import textwrap
import threading
import time
from collections import Counter
from itertools import islice
from typing import Any, Optional
@ -37,9 +38,15 @@ README_MAX_CHARS = 1500
# caused chat-evict races and finally-eviction bugs and still left
# delete-cache blind because helper/advisor publish prefixed
# identifiers the guard could not match). Expose loading repo ids
# through a thread-safe set so DELETE /api/models/delete-cached can
# block while a helper or advisor still owns the cache.
_HELPER_ADVISOR_ACTIVE_REPOS: set[str] = set()
# through a thread-safe Counter so DELETE /api/models/delete-cached
# can block while a helper or advisor still owns the cache.
#
# Round 27 P1 #1: must refcount, not a plain set. A helper and an
# advisor (or two concurrent helpers) often share the default repo
# unsloth/gemma-4-E2B-it-GGUF. With a set, the first finally call
# discarded the repo while the second invocation was still loading,
# and the delete-cache guard then let rmtree race the live mmap.
_HELPER_ADVISOR_REFCOUNT: Counter[str] = Counter()
_HELPER_ADVISOR_LOCK = threading.Lock()
@ -51,21 +58,33 @@ def helper_advisor_owns_repo(repo_id: str) -> bool:
return False
needle = repo_id.lower()
with _HELPER_ADVISOR_LOCK:
return needle in _HELPER_ADVISOR_ACTIVE_REPOS
return _HELPER_ADVISOR_REFCOUNT.get(needle, 0) > 0
def helper_advisor_busy() -> bool:
"""Round 27 P1 #2: True if ANY helper/advisor load is in flight.
Used by diffusion / training / export release paths so they do
not allocate on top of the helper's VRAM while it owns its
private LlamaCppBackend instance."""
with _HELPER_ADVISOR_LOCK:
return sum(_HELPER_ADVISOR_REFCOUNT.values()) > 0
def _register_helper_advisor_repo(repo_id: str) -> None:
if not repo_id:
return
with _HELPER_ADVISOR_LOCK:
_HELPER_ADVISOR_ACTIVE_REPOS.add(repo_id.lower())
_HELPER_ADVISOR_REFCOUNT[repo_id.lower()] += 1
def _unregister_helper_advisor_repo(repo_id: str) -> None:
if not repo_id:
return
needle = repo_id.lower()
with _HELPER_ADVISOR_LOCK:
_HELPER_ADVISOR_ACTIVE_REPOS.discard(repo_id.lower())
_HELPER_ADVISOR_REFCOUNT[needle] -= 1
if _HELPER_ADVISOR_REFCOUNT[needle] <= 0:
_HELPER_ADVISOR_REFCOUNT.pop(needle, None)
def _strip_think_tags(text: str) -> str:

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

@ -148,6 +148,18 @@ export function ImagesPage() {
void refreshStatus();
}, [refreshStatus]);
// Round 27 P2: when the backend is mid-load (is_loading=true) the
// status label froze at "Loading..." until the user clicked
// Refresh. Auto-poll every 2 s while a load is in flight so the
// UI tracks real backend progress.
useEffect(() => {
if (!status?.is_loading) return;
const id = window.setInterval(() => {
void refreshStatus();
}, 2000);
return () => window.clearInterval(id);
}, [status?.is_loading, refreshStatus]);
const handleLoad = useCallback(async () => {
setBusy("loading");
try {
@ -207,6 +219,12 @@ export function ImagesPage() {
toast.error("Failed to unload image model", {
description: err instanceof Error ? err.message : String(err),
});
// Round 27 P2: a partial unload (subprocess refused to terminate,
// 503 from the backend) used to leave the UI showing the old
// "Loaded:" label even though the backend state was half torn
// down. Refresh so the button states match reality (mirrors
// handleLoad above which always re-fetches on catch).
await refreshStatus();
} finally {
setBusy("idle");
}
@ -491,6 +509,7 @@ export function ImagesPage() {
<div className="flex flex-col gap-1">
<Label>Steps: {steps}</Label>
<Slider
aria-label="Inference steps"
min={1}
max={60}
step={1}
@ -501,6 +520,7 @@ export function ImagesPage() {
<div className="flex flex-col gap-1">
<Label>Guidance: {guidance.toFixed(1)}</Label>
<Slider
aria-label="Guidance scale"
min={0}
max={15}
step={0.1}
@ -554,7 +574,7 @@ export function ImagesPage() {
data-testid="diffusion-result-image"
/>
<figcaption className="text-xs text-muted-foreground">
{r.width}x{r.height} - {r.num_inference_steps} steps - g={r.guidance_scale.toFixed(1)}
{r.width}x{r.height} - {r.num_inference_steps} steps - g={(r.guidance_scale ?? 0).toFixed(1)}
{/* Prefer seed_str (full uint64 precision) since the
numeric seed gets rounded by JSON.parse above
Number.MAX_SAFE_INTEGER and would otherwise