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

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

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