diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 305ad6ea94..3bdce54819 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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 diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 5691c3d0d5..1b56ecbe7e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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""" diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index b6e3a3c5b8..c2a41884f3 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7649c453e1..636bb5cc9c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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), diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 3ec3a9372e..377d2174be 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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 { diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4ea755207c..79bf233e4c 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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 diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index d6dea74c47..a994489653 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -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