From 6c528fb0138616dce7a732f3e0d09c124d121ab8 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 13:39:40 +0000 Subject: [PATCH] Fix/adjust diffusion: round 27 P1 + P2 batch for PR #5754 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. --- studio/backend/core/inference/diffusion.py | 14 +++++++++ studio/backend/routes/inference.py | 18 +++++++++++ studio/backend/utils/datasets/llm_assist.py | 31 +++++++++++++++---- .../src/features/images/images-page.tsx | 22 ++++++++++++- 4 files changed, 78 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index f670fb0f05..305ad6ea94 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 673c31d508..7649c453e1 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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) diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 5299a7ac75..37c17329c8 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -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: diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index b5aa124fd7..265786e6bd 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -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() {
- {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