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:
parent
79da5d910d
commit
c4c9e2aeec
7 changed files with 186 additions and 44 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue