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

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

Browse source 
Twelve actionable P1/P2 findings from round 28 reviewer aggregate.
Skipped #3 (studio.txt huggingface-hub bump) because the empirical
CI evidence in round 26 contradicts that suggestion: bumping the
pin there breaks installs that apply constraints.txt
(transformers==4.57.6 requires hub<1.0). The actual broken combo
only happens via the --no-deps no-torch path which is already
bumped in no-torch-runtime.txt and pyproject.toml huggingfacenotorch.

1. utils/datasets/llm_assist.py: split _HELPER_ADVISOR_REFCOUNT
   into CACHE vs GPU counters. helper_advisor_owns_repo (used by
   delete-cache) reads CACHE; helper_advisor_busy (used by public
   handoffs) reads GPU. precache_helper_gguf now registers with
   gpu_owner=False so a background pre-cache download does not
   503 every chat / training / export / diffusion load.

2. utils/datasets/llm_assist.py: introduce _HELPER_ADVISOR_START_LOCK
   and wrap the busy precheck + register pair in _run_with_helper
   and _run_multi_pass_advisor. Two concurrent helper / advisor
   invocations could both pass _gpu_workload_busy_for_helper before
   either registered, then OOM each other.

3. utils/datasets/llm_assist.py: _gpu_workload_busy_for_helper now
   also returns True when another helper/advisor already holds the
   private LlamaCppBackend.

4. routes/inference.py: add _raise_if_helper_advisor_busy(workload)
   that 503s when AI Assist owns the GPU. Wire it into both chat
   load branches (GGUF + safetensors) BEFORE the existing
   _release_export_for / _release_diffusion_for calls so we do not
   first tear down an idle export / diffusion just to fail on the
   helper check.

5. routes/training.py + routes/export.py + diffusion.load_model:
   call the helper-busy check FIRST before any release helper
   fires. Mirrors the chat-load ordering.

6. routes/inference.py _release_llama_for: poll
   loading_model_identifier for up to 5 s after unload_model() so a
   cancelled pending GGUF download has time to clear its
   identifier. Mirrors the same wait round 26 added to the explicit
   /api/inference/unload route.

7. core/inference/diffusion.py _release_chat_backend_for_diffusion:
   same 5 s settling wait for cancelled pending GGUF downloads.

8. models/inference.py LoadRequest: validate every llama_extra_args
   entry through _no_control_chars + _reject_embedded_hf_token.
   The list was forwarded verbatim to a logged llama-server command
   line, so a smuggled control char or hf_... token would land in
   logs and subprocess args.

9. routes/models.py /gguf-download-progress: apply
   _validate_logged_identifier to repo_id and variant, matching the
   round 24 hardening on the adjacent generic /download-progress.

10. routes/inference.py diffusion-load RuntimeError classifier:
    treat "AI Assist ..." messages as retryable 503 instead of 400
    (round 28 P2 #15). Mirrors the round 18/19 markers for chat
    unload failures.

Tests: 105 targeted + 1768 broader backend tests pass locally.
This commit is contained in:
Daniel Han-Chen 2026-05-25 14:21:37 +00:00
commit c4c9e2aeec
7 changed files with 186 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1035,8 +1035,13 @@ class DiffusionBackend:
# transformer while the old pipeline still owns
# its weights.
# 4. THEN call from_single_file / from_pretrained.
_release_other_gpu_owners_for_diffusion()
# Round 28 P1 #4: helper/advisor check must fire BEFORE
# _release_other_gpu_owners_for_diffusion. Otherwise a
# blocked Images load could first tear down an idle
# export checkpoint just to then RuntimeError on the
# helper check inside _release_chat_backend_for_diffusion.
_release_chat_backend_for_diffusion()
_release_other_gpu_owners_for_diffusion()
old = self._pipe
if old is not None:
@ -1558,6 +1563,16 @@ def _release_chat_backend_for_diffusion() -> None:
"Could not unload the existing GGUF chat model before "
"loading a diffusion image model."
) from exc
# Round 28 P1 #12: a cancelled pending GGUF download takes
# up to a few seconds to clear loading_model_identifier in
# its finally block. Wait briefly so the same retryable
# cancel path used by the unload route does not 503 us.
deadline = time.monotonic() + 5.0
while (
getattr(backend, "loading_model_identifier", None)
and time.monotonic() < deadline
):
time.sleep(0.1)
# Round 18 P1 #4: also reject when ``loading_model_identifier``
# is still set after the unload call. Without this, a GGUF
# download / startup that was already in flight before the

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

@ -127,6 +127,27 @@ class LoadRequest(BaseModel):
),
)
# Round 28 P1 #13: each entry is forwarded verbatim to a logged
# subprocess command line and reflected in errors. Reject control
# chars and embedded HF tokens for every list entry; allow None.
@field_validator("llama_extra_args")
@classmethod
def _no_extra_args_control_chars(cls, v):
if v is None:
return v
for i, entry in enumerate(v):
_no_control_chars(entry, f"llama_extra_args[{i}]")
return v
@field_validator("llama_extra_args")
@classmethod
def _no_extra_args_embedded_hf_tokens(cls, v):
if v is None:
return v
for i, entry in enumerate(v):
_reject_embedded_hf_token(entry, f"llama_extra_args[{i}]")
return v
class UnloadRequest(BaseModel):
"""Request to unload a model"""

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

@ -148,8 +148,15 @@ async def load_checkpoint(
# helper so we cover llama-server is_active=True and
# safetensors loading_models -- the asymmetries round 9
# reviews #1, #8, #9 flagged.
from routes.inference import _release_chat_for, _release_diffusion_for
from routes.inference import (
_raise_if_helper_advisor_busy,
_release_chat_for,
_release_diffusion_for,
)
# Round 28 P1 #6: refuse before any release fires so AI Assist
# busy does not first tear down idle diffusion.
_raise_if_helper_advisor_busy("export")
# Round 24 P1 #3: release diffusion BEFORE chat so a failing
# diffusion unload does not leave the user with no chat
# model loaded. Same reasoning as the training-start flow

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

@ -357,6 +357,42 @@ def _raise_if_export_active(workload: str) -> None:
)
def _raise_if_helper_advisor_busy(workload: str) -> None:
"""Round 28 P1 #1 / #4 / #5 / #6: refuse a new GPU workload while
AI Assist helper / advisor still owns its PRIVATE LlamaCppBackend.
Called early so callers do NOT first tear down idle export /
diffusion / chat owners just to fail on the helper check.
"""
try:
from utils.datasets.llm_assist import helper_advisor_busy
except Exception:
return
try:
busy = helper_advisor_busy()
except Exception as exc:
logger.warning(
"Could not verify helper/advisor status before %s load: %s",
workload,
exc,
)
raise HTTPException(
status_code = 503,
detail = (
f"Could not verify AI Assist status before starting {workload}. "
f"Try again."
),
) from exc
if busy:
raise HTTPException(
status_code = 503,
detail = (
f"AI Assist (helper / advisor GGUF) is still using the GPU. "
f"Wait for it to finish before starting {workload}."
),
)
async def _release_llama_for(workload: str) -> None:
"""Unload the llama-server (GGUF) chat backend if it owns the
GPU. Treats ``is_loaded`` OR ``is_active`` OR
@ -403,6 +439,18 @@ async def _release_llama_for(workload: str) -> None:
),
) from exc
# Round 28 P1 #11: a pending HF GGUF download cancelled by
# unload_model() takes up to a few seconds to settle (the load
# thread observes _cancel_event in its finally and clears
# loading_model_identifier). Wait briefly so a legitimate cancel
# does not 503. Mirrors the /api/inference/unload settling wait.
deadline = time.monotonic() + 5.0
while (
bool(getattr(llama, "loading_model_identifier", None))
and time.monotonic() < deadline
):
await asyncio.sleep(0.1)
# Round 18 P1 #1: previously only the raised-exception path was
# treated as failure. ``llama.unload_model()`` returning ``False``
# (subprocess refused to terminate, IPC timeout) or leaving
@ -1210,6 +1258,10 @@ async def load_model(
# corrupt the user's exported artifact).
_raise_if_training_active("chat")
_raise_if_export_active("chat")
# Round 28 P1 #1: refuse before the release helpers fire
# so we do not tear down an idle export / diffusion just to
# then 503 on the helper check.
_raise_if_helper_advisor_busy("GGUF chat")
# Round 24 P1 #4: release order is now
# export -> diffusion -> safetensors chat (was
# export -> safetensors chat -> diffusion). A wedged
@ -1409,6 +1461,9 @@ async def load_model(
# and so we do not silently corrupt an in-flight export.
_raise_if_training_active("chat")
_raise_if_export_active("chat")
# Round 28 P1 #1: refuse before the release helpers tear down
# idle GPU owners.
_raise_if_helper_advisor_busy("safetensors chat")
# Round 24 P1 #5: release order is now
# export -> diffusion -> llama-chat (was
# export -> llama-chat -> diffusion). A wedged diffusion
@ -2196,6 +2251,12 @@ async def diffusion_load(
# the request is refused with 409 instead of silently killing it.
_raise_if_training_active("diffusion")
_raise_if_export_active("diffusion")
# Round 28 P1 #4: AI Assist helper/advisor owns a private llama
# backend invisible to _release_chat_backend_for_diffusion's
# global checks. Refuse early so we do not first tear down an
# idle export checkpoint just to fail on the helper check inside
# load_model.
_raise_if_helper_advisor_busy("diffusion")
# Round 18 P1 #3 + P1 #7: the route used to drop chat and idle
# export BEFORE ``backend.load_model`` ran its cheap validation
# (family inference, GGUF filename checks, gated-token failures,
@ -2243,6 +2304,9 @@ async def diffusion_load(
# to the user instead of 503. Match both wordings.
or "still active or loading after unload" in detail
or "still loading after unload" in detail
# Round 28 P2 #15: AI Assist running (raised by
# _release_chat_backend_for_diffusion) is retryable.
or "AI Assist" in detail
):
# Round 17 P1 #2: chat unload failures raised by the
# backend helper map to 503 (retryable infra issue),

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

@ -2440,6 +2440,13 @@ async def get_gguf_download_progress(
Tracks completed shard downloads in snapshots and in-progress downloads
in the blobs directory (incomplete files).
"""
# Round 28 P1 #14: mirror the hardening on the generic
# /download-progress route. Both repo_id and variant are echoed
# into the cache-scan path and can reach logs on the failure
# branch via the surrounding try/except.
repo_id = _validate_logged_identifier(repo_id, "repo_id")
if variant:
variant = _validate_logged_identifier(variant, "variant")
try:
if not _is_valid_repo_id(repo_id):
return {

4
studio/backend/routes/training.py
View file

@ -273,12 +273,16 @@ async def start_training(
# stops the export and re-submits.
from routes.inference import (
_raise_if_export_active,
_raise_if_helper_advisor_busy,
_release_chat_for,
_release_diffusion_for,
_release_export_for,
)
_raise_if_export_active("training")
# Round 28 P1 #5: refuse before any release fires so AI Assist
# busy does not first tear down idle diffusion/export.
_raise_if_helper_advisor_busy("training")
# Round 18 P1 #8: release settled export FIRST so an export
# cleanup failure preserves the user's currently loaded chat
# model. The previous order (chat -> export) would drop chat

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

@ -34,57 +34,69 @@ DEFAULT_HELPER_MODEL_VARIANT = "UD-Q4_K_XL"
README_MAX_CHARS = 1500
# Round 26 P1 #13 / #14: helper/advisor run on PRIVATE LlamaCppBackend
# instances (round 25 P1 #4 briefly used the global singleton, which
# 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 Counter so DELETE /api/models/delete-cached
# can block while a helper or advisor still owns the cache.
# instances. Expose loading repo ids through thread-safe Counters 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()
# Round 28 P1 #2: split into CACHE vs GPU refcounts. precache_helper_gguf
# downloads files (cache ownership) without occupying VRAM (GPU
# ownership), so collapsing them caused the public GPU handoffs to
# 503 during a background precache that did not need the GPU.
# * CACHE: blocks delete-cache for any active downloader / loader
# * GPU : blocks public chat / training / export / diffusion loads
_HELPER_ADVISOR_CACHE_REFCOUNT: Counter[str] = Counter()
_HELPER_ADVISOR_GPU_REFCOUNT: Counter[str] = Counter()
_HELPER_ADVISOR_LOCK = threading.Lock()
# Round 28 P1 #7 / #8 / #10: serialize helper / advisor STARTS so two
# concurrent invocations cannot both pass the busy precheck before
# either registers. Held only across the precheck + register window,
# not across the full helper run.
_HELPER_ADVISOR_START_LOCK = threading.Lock()
def helper_advisor_owns_repo(repo_id: str) -> bool:
"""Return True if any helper/advisor load currently owns this
HF repo id. Comparison is case-insensitive to match the chat
backend's lowercased needle."""
"""Return True if any helper/advisor activity (precache OR live
helper / advisor load) currently owns this HF repo id."""
if not repo_id:
return False
needle = repo_id.lower()
with _HELPER_ADVISOR_LOCK:
return _HELPER_ADVISOR_REFCOUNT.get(needle, 0) > 0
return _HELPER_ADVISOR_CACHE_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."""
"""True if any helper/advisor load is currently OCCUPYING THE GPU.
Round 28 P1 #2: must not return True for a precache-only download
(it owns disk cache, not VRAM)."""
with _HELPER_ADVISOR_LOCK:
return sum(_HELPER_ADVISOR_REFCOUNT.values()) > 0
return sum(_HELPER_ADVISOR_GPU_REFCOUNT.values()) > 0
def _register_helper_advisor_repo(repo_id: str) -> None:
if not repo_id:
return
with _HELPER_ADVISOR_LOCK:
_HELPER_ADVISOR_REFCOUNT[repo_id.lower()] += 1
def _unregister_helper_advisor_repo(repo_id: str) -> None:
def _register_helper_advisor_repo(repo_id: str, *, gpu_owner: bool = True) -> None:
"""Register a helper/advisor activity. Set ``gpu_owner=False`` for
precache-only downloads that need cache-delete protection but do
not load weights into VRAM."""
if not repo_id:
return
needle = repo_id.lower()
with _HELPER_ADVISOR_LOCK:
_HELPER_ADVISOR_REFCOUNT[needle] -= 1
if _HELPER_ADVISOR_REFCOUNT[needle] <= 0:
_HELPER_ADVISOR_REFCOUNT.pop(needle, None)
_HELPER_ADVISOR_CACHE_REFCOUNT[needle] += 1
if gpu_owner:
_HELPER_ADVISOR_GPU_REFCOUNT[needle] += 1
def _unregister_helper_advisor_repo(repo_id: str, *, gpu_owner: bool = True) -> None:
if not repo_id:
return
needle = repo_id.lower()
with _HELPER_ADVISOR_LOCK:
_HELPER_ADVISOR_CACHE_REFCOUNT[needle] -= 1
if _HELPER_ADVISOR_CACHE_REFCOUNT[needle] <= 0:
_HELPER_ADVISOR_CACHE_REFCOUNT.pop(needle, None)
if gpu_owner:
_HELPER_ADVISOR_GPU_REFCOUNT[needle] -= 1
if _HELPER_ADVISOR_GPU_REFCOUNT[needle] <= 0:
_HELPER_ADVISOR_GPU_REFCOUNT.pop(needle, None)
def _strip_think_tags(text: str) -> str:
@ -128,10 +140,11 @@ def precache_helper_gguf():
)
# Round 27 P1 #4: register the repo so DELETE /api/models/delete-cached
# cannot rmtree the cache directory while we are mid-download. Helper
# / advisor runtime calls already register, but the startup precache
# was the asymmetric gap that let cache delete race the first download.
_register_helper_advisor_repo(repo)
# cannot rmtree the cache directory while we are mid-download.
# Round 28 P1 #2: precache only downloads files; it does NOT occupy
# VRAM. Use gpu_owner=False so helper_advisor_busy() does not block
# public GPU workloads during a background pre-cache.
_register_helper_advisor_repo(repo, gpu_owner = False)
try:
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
@ -163,7 +176,7 @@ def precache_helper_gguf():
except Exception as e:
logger.warning(f"Failed to pre-cache helper GGUF: {e}")
finally:
_unregister_helper_advisor_repo(repo)
_unregister_helper_advisor_repo(repo, gpu_owner = False)
try:
enable_progress_bars()
except Exception as e:
@ -207,7 +220,14 @@ def _gpu_workload_busy_for_helper() -> bool:
GPU; mirror the diffusion check by inspecting llama
``is_loaded`` / ``is_active`` / ``loading_model_identifier`` and
safetensors ``active_model_name`` / ``loading_models``.
Round 28 P1 #9: also catch another helper / advisor that already
owns a private LlamaCppBackend. Without this two concurrent
helpers could both pass the precheck and OOM each other.
"""
if helper_advisor_busy():
logger.info("Skipping helper GGUF while another helper/advisor is using the GPU")
return True
if _diffusion_image_model_busy():
return True
@ -296,16 +316,18 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
# 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
# Round 28 P1 #7 / #10: serialize the busy check + register pair
# so two concurrent helper invocations cannot both pass the
# precheck before either registers and then OOM each other.
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
with _HELPER_ADVISOR_START_LOCK:
if _gpu_workload_busy_for_helper():
return None
_register_helper_advisor_repo(repo)
backend = None
_register_helper_advisor_repo(repo)
try:
# Round 26 P1 #1 / #3 / #13 / #14: use a PRIVATE backend so the
# helper can never preempt or be preempted by the user's
@ -700,16 +722,18 @@ def _run_multi_pass_advisor(
# 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
# Round 28 P1 #8 / #10: serialize the precheck + register pair so
# two concurrent advisor invocations cannot both pass before
# either registers and then OOM each other.
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
with _HELPER_ADVISOR_START_LOCK:
if _gpu_workload_busy_for_helper():
return None
_register_helper_advisor_repo(repo)
backend = None
_register_helper_advisor_repo(repo)
try:
# Round 26 P1 #2 / #4 / #13 / #14: mirror ``_run_with_helper``
# and use a PRIVATE backend. Round 25's global-backend swap