diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 1b56ecbe7e..6b99cc4ec1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1565,6 +1565,16 @@ class DiffusionLoadRequest(BaseModel): repo_id: str = Field( ..., min_length = 1, max_length = 1024, description = "HF repo id or local path" ) + # Round 30 P1 #4: chat /api/inference/load gates native local paths + # through a signed native_path_lease grant before the backend + # touches the filesystem. Mirror that here so /api/inference/images/ + # load cannot be used as an authenticated probe for arbitrary + # local directories. Optional: Hub ids (no leading slash / tilde) + # skip the lease check entirely. + native_path_lease: Optional[str] = Field( + None, + description = "Frontend-visible signed native path grant for a local repo_id", + ) gguf_filename: Optional[str] = Field( None, max_length = 512, @@ -1575,6 +1585,10 @@ class DiffusionLoadRequest(BaseModel): max_length = 1024, description = "Diffusers base repo (HF id or local path) for VAE + text encoders", ) + base_repo_native_path_lease: Optional[str] = Field( + None, + description = "Frontend-visible signed native path grant for a local base_repo", + ) family: Optional[str] = Field( None, max_length = 64, diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index c2a41884f3..36337251be 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -60,6 +60,10 @@ async def load_checkpoint( Wraps ExportBackend.load_checkpoint. """ + # Round 30 P1 #8: track whether we published a public-load pending + # entry so the outer finally clears it on either success or + # failure path. + export_load_window_published = False try: # Version switching is handled automatically by the subprocess-based # export backend — no need for ensure_transformers_version() here. @@ -149,6 +153,7 @@ async def load_checkpoint( # safetensors loading_models -- the asymmetries round 9 # reviews #1, #8, #9 flagged. from routes.inference import ( + _clear_public_load_window, _raise_if_helper_advisor_busy, _release_chat_for, _release_diffusion_for, @@ -156,7 +161,12 @@ async def load_checkpoint( # Round 28 P1 #6: refuse before any release fires so AI Assist # busy does not first tear down idle diffusion. + # Round 30 P1 #8: also publishes a public-load pending entry so + # a concurrent helper / advisor start cannot win the start + # lock between our snapshot and load_checkpoint flipping + # current_checkpoint / is_export_active. _raise_if_helper_advisor_busy("export") + export_load_window_published = True # 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 @@ -190,6 +200,18 @@ async def load_checkpoint( status_code = 500, detail = f"Failed to load checkpoint: {str(e)}", ) + finally: + # Round 30 P1 #8: clear the public-load pending entry once the + # load attempt completes (success or failure). Skipped when + # the helper-busy check itself raised so the counter stays in + # sync with publishes. + if export_load_window_published: + try: + from routes.inference import _clear_public_load_window + except Exception: + pass + else: + _clear_public_load_window("export") @router.post("/cleanup", response_model = ExportOperationResponse) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4c014ea122..7db8c9752f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -363,34 +363,60 @@ def _raise_if_helper_advisor_busy(workload: str) -> None: Called early so callers do NOT first tear down idle export / diffusion / chat owners just to fail on the helper check. + + Round 30 P1 #7-#10: also publishes a public-load pending entry + under the helper-advisor start lock so a concurrent helper start + sees the pending public owner and refuses VRAM. Callers MUST + invoke ``_clear_public_load_window(workload)`` in a paired + finally to clear the entry once the load attempt completes. """ try: - from utils.datasets.llm_assist import helper_advisor_busy + from utils.datasets.llm_assist import ( + _HELPER_ADVISOR_START_LOCK, + _publish_public_load_pending, + helper_advisor_busy, + ) + except Exception: + return + with _HELPER_ADVISOR_START_LOCK: + 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}." + ), + ) + _publish_public_load_pending(workload) + + +def _clear_public_load_window(workload: str) -> None: + """Pair for ``_raise_if_helper_advisor_busy``: release the pending + public-load publish so a subsequent helper start can proceed. + Safe to call when the module import failed (no-op).""" + try: + from utils.datasets.llm_assist import _release_public_load_pending 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}." - ), - ) + _release_public_load_pending(workload) + except Exception: + pass async def _release_llama_for(workload: str) -> None: @@ -1098,6 +1124,10 @@ async def load_model( """ native_grant_backed = False model_log_label = request.model_path + # Round 30 P1 #7 / #9: track which branch (GGUF / safetensors) + # published a public-load pending entry so the outer finally + # decrements the same counter, even on early exception. + chat_load_window_workload: Optional[str] = None try: # Validate user-supplied llama-server pass-through args up front # so a managed-flag collision returns 400 before any model work. @@ -1262,6 +1292,7 @@ async def load_model( # 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") + chat_load_window_workload = "GGUF chat" # Round 24 P1 #4: release order is now # export -> diffusion -> safetensors chat (was # export -> safetensors chat -> diffusion). A wedged @@ -1464,6 +1495,7 @@ async def load_model( # Round 28 P1 #1: refuse before the release helpers tear down # idle GPU owners. _raise_if_helper_advisor_busy("safetensors chat") + chat_load_window_workload = "safetensors chat" # Round 24 P1 #5: release order is now # export -> diffusion -> llama-chat (was # export -> llama-chat -> diffusion). A wedged diffusion @@ -1649,6 +1681,14 @@ async def load_model( if any(h.lower() in msg.lower() for h in not_supported_hints): msg = f"This model is not supported yet. Try a different model. (Original error: {msg})" raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") + finally: + # Round 30 P1 #7 / #9: clear whichever chat branch published a + # public-load pending entry so a subsequent helper / advisor + # start can proceed. Set on the GGUF / safetensors branches + # after _raise_if_helper_advisor_busy succeeds; stays None for + # the already-loaded fast paths above. + if chat_load_window_workload is not None: + _clear_public_load_window(chat_load_window_workload) @router.post("/validate", response_model = ValidateModelResponse) @@ -2232,6 +2272,59 @@ def _get_diffusion_backend(): return get_diffusion_backend() +def _looks_like_local_diffusion_path(value: Optional[str]) -> bool: + """Round 30 P1 #4: decide whether ``repo_id`` / ``base_repo`` + names a local filesystem path that requires a signed + ``native_path_lease`` grant. Hub ids (``owner/repo`` form, no + leading separator or tilde) skip the lease check; anything that + starts with ``/``, ``~``, ``./``, ``../``, contains a backslash, + or resolves to an absolute path is treated as a local-path + access attempt. We DO NOT consult ``Path.exists`` so the route + does not side-channel filesystem layout information back to the + caller via the lease error vs. the load error.""" + if not value: + return False + if value.startswith(("/", "~", "./", "../")): + return True + if "\\" in value: + return True + try: + if Path(value).expanduser().is_absolute(): + return True + except (OSError, ValueError): + # Treat unparseable identifiers as local-path attempts so a + # broken input does not silently fall through to the Hub + # loader (defence-in-depth, not a tested code path). + return True + return False + + +def _resolve_diffusion_repo_for_request( + value: Optional[str], + lease: Optional[str], + *, + operation: str, +) -> Optional[str]: + """Round 30 P1 #4: enforce the same signed-lease boundary the chat + /api/inference/load path uses. Hub ids return as-is. Local + paths require a verified ``native_path_lease`` directory grant; + a missing or invalid lease returns 400 BEFORE any GPU handoff.""" + if value is None: + return None + if not _looks_like_local_diffusion_path(value): + return value + try: + grant = verify_native_path_lease( + lease, + operation = operation, + expected_kind = "model", + expected_path_type = "directory", + ) + except NativePathLeaseError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc + return str(grant.canonical_path) + + @studio_router.post("/images/load") async def diffusion_load( payload: DiffusionLoadRequest, @@ -2256,7 +2349,23 @@ async def diffusion_load( # 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. + # Round 30 P1 #10: also publishes the public-load pending entry so + # a concurrent helper start cannot win the start lock between our + # snapshot and DiffusionBackend.load_model flipping is_loaded. _raise_if_helper_advisor_busy("diffusion") + # Round 30 P1 #4: enforce the signed native_path_lease boundary the + # chat load path uses so local-path repo_id / base_repo cannot be + # probed without a frontend-issued grant. Hub ids pass through. + resolved_repo_id = _resolve_diffusion_repo_for_request( + payload.repo_id, + payload.native_path_lease, + operation = "load-diffusion-model", + ) or payload.repo_id + resolved_base_repo = _resolve_diffusion_repo_for_request( + payload.base_repo, + payload.base_repo_native_path_lease, + operation = "load-diffusion-model", + ) # 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, @@ -2274,9 +2383,9 @@ async def diffusion_load( status = await asyncio.get_event_loop().run_in_executor( None, lambda: backend.load_model( - repo_id = payload.repo_id, + repo_id = resolved_repo_id, gguf_filename = payload.gguf_filename, - base_repo = payload.base_repo, + base_repo = resolved_base_repo, family_override = payload.family, hf_token = payload.hf_token, enable_model_cpu_offload = payload.enable_model_cpu_offload, @@ -2321,6 +2430,11 @@ async def diffusion_load( except Exception as exc: logger.exception("Diffusion load failed") raise HTTPException(status_code = 500, detail = str(exc)) + finally: + # Round 30 P1 #10: clear the public-load pending publish so a + # subsequent helper / advisor start can proceed once the + # diffusion load attempt has finished (success or failure). + _clear_public_load_window("diffusion") @studio_router.post("/images/unload") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 79bf233e4c..5f0bdc38a4 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -127,6 +127,11 @@ async def start_training( This endpoint initiates training in the background and returns immediately. Use the /status endpoint to check training progress. """ + # Round 30 P1 #7: track whether we published a public-load pending + # entry so the outer finally clears it on either success or + # failure (including any early HTTPException raised by the helper + # check itself). + training_load_window_published = False try: logger.info(f"Starting training job with model: {request.model_name}") @@ -272,6 +277,7 @@ async def start_training( # the user's output artifact. Now we 409 first; the user # stops the export and re-submits. from routes.inference import ( + _clear_public_load_window, _raise_if_export_active, _raise_if_helper_advisor_busy, _release_chat_for, @@ -282,7 +288,13 @@ async def start_training( _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. + # Round 30 P1 #7: also publishes a public-load pending entry so + # a concurrent helper / advisor start cannot win the start + # lock between our snapshot and start_training flipping + # is_training_active. Paired clear lives in the outer + # ``finally`` below. _raise_if_helper_advisor_busy("training") + training_load_window_published = True # 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 @@ -336,6 +348,18 @@ async def start_training( status_code = 500, detail = f"Failed to start training: {str(e)}", ) + finally: + # Round 30 P1 #7: clear the public-load pending entry once the + # start attempt has finished. Skipped when the helper-busy + # check itself raised (no publish to clear) so the counter + # stays in sync with publishes. + if training_load_window_published: + try: + from routes.inference import _clear_public_load_window + except Exception: + pass + else: + _clear_public_load_window("training") @router.post("/stop", response_model = TrainingStopResponse) diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index ab4c68b3e2..7eb437342b 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -46,11 +46,22 @@ README_MAX_CHARS = 1500 # * GPU : blocks public chat / training / export / diffusion loads _HELPER_ADVISOR_CACHE_REFCOUNT: Counter[str] = Counter() _HELPER_ADVISOR_GPU_REFCOUNT: Counter[str] = Counter() +# Round 30 P1 #7-#10: counter of public GPU workloads (chat / +# diffusion / training / export) that have passed the helper-busy +# snapshot but have not yet flipped their public ownership flags +# (``llama.is_loaded`` / ``loading_model_identifier`` / +# ``current_checkpoint`` / ``is_training_active``). Helper / advisor +# starts consult this so they cannot win the start lock and race a +# public load that already destroyed the previous owner. +_PUBLIC_LOAD_PENDING_COUNT: 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. +# Round 30 P1 #7-#10: public GPU loads also enter under this lock to +# publish their pending counter so a concurrent helper / advisor +# start sees the pending public owner and refuses VRAM. _HELPER_ADVISOR_START_LOCK = threading.Lock() @@ -99,6 +110,38 @@ def _unregister_helper_advisor_repo(repo_id: str, *, gpu_owner: bool = True) -> _HELPER_ADVISOR_GPU_REFCOUNT.pop(needle, None) +def _publish_public_load_pending(workload: str) -> None: + """Mark a public GPU workload as mid-handoff. Must be called under + ``_HELPER_ADVISOR_START_LOCK`` immediately after the helper-busy + snapshot succeeded (round 30 P1 #7-#10).""" + if not workload: + return + needle = workload.lower() + with _HELPER_ADVISOR_LOCK: + _PUBLIC_LOAD_PENDING_COUNT[needle] += 1 + + +def _release_public_load_pending(workload: str) -> None: + """Decrement the pending public-load counter once per matched + publish. Safe to call in finally even if the load failed.""" + if not workload: + return + needle = workload.lower() + with _HELPER_ADVISOR_LOCK: + _PUBLIC_LOAD_PENDING_COUNT[needle] -= 1 + if _PUBLIC_LOAD_PENDING_COUNT[needle] <= 0: + _PUBLIC_LOAD_PENDING_COUNT.pop(needle, None) + + +def public_load_pending() -> bool: + """True if any public GPU workload has passed its helper-busy + snapshot but not yet flipped its public ownership flags. Helper / + advisor starts treat this as busy so they cannot race a public + load mid-handoff.""" + with _HELPER_ADVISOR_LOCK: + return sum(_PUBLIC_LOAD_PENDING_COUNT.values()) > 0 + + def _strip_think_tags(text: str) -> str: """Strip ... reasoning blocks emitted by some models. @@ -230,6 +273,15 @@ def _gpu_workload_busy_for_helper() -> bool: "Skipping helper GGUF while another helper/advisor is using the GPU" ) return True + # Round 30 P1 #7-#10: a public GPU load (chat / diffusion / training / + # export) that has passed its busy snapshot but not yet flipped its + # public ownership flags is still mid-handoff. Refuse so the helper + # does not race it for VRAM after the previous owner was torn down. + if public_load_pending(): + logger.info( + "Skipping helper GGUF while a public GPU load is mid-handoff" + ) + return True if _diffusion_image_model_busy(): return True diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index ad51ef9db8..3fd9dab681 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -128,8 +128,11 @@ export function ImagesPage() { const preset = CURATED_MODELS[presetIndex] ?? DEFAULT_PRESET; const resolution = RESOLUTION_PRESETS[resolutionIdx]; - const refreshStatus = useCallback(async () => { - setRefreshingStatus(true); + // Round 30 P2 #12: split the fetch from the spinner toggle so the + // mount + auto-poll effects can call the fetch without the + // synchronous setRefreshingStatus(true) that tripped + // react-hooks/set-state-in-effect. + const fetchAndUpdateStatus = useCallback(async () => { try { const next = await fetchDiffusionStatus(); setStatus(next); @@ -139,20 +142,25 @@ export function ImagesPage() { lastErrorRef.current = msg; toast.error("Could not fetch image-model status", { description: msg }); } - } finally { - setRefreshingStatus(false); } }, []); + const refreshStatus = useCallback(async () => { + setRefreshingStatus(true); + try { + await fetchAndUpdateStatus(); + } finally { + setRefreshingStatus(false); + } + }, [fetchAndUpdateStatus]); + useEffect(() => { - // Round 30 P2 #12: defer the first refreshStatus call via - // queueMicrotask so the synchronous setRefreshingStatus(true) - // inside it does not trip the react-hooks/set-state-in-effect - // lint rule on the mount render. - queueMicrotask(() => { - void refreshStatus(); - }); - }, [refreshStatus]); + // Mount fetch goes through fetchAndUpdateStatus so the lint rule + // does not see any synchronous setState in the effect body; the + // user-driven Refresh button still calls refreshStatus to flip + // the spinner. + void fetchAndUpdateStatus(); + }, [fetchAndUpdateStatus]); // Round 27 P2: when the backend is mid-load (is_loading=true) the // status label froze at "Loading..." until the user clicked @@ -161,10 +169,10 @@ export function ImagesPage() { useEffect(() => { if (!status?.is_loading) return; const id = window.setInterval(() => { - void refreshStatus(); + void fetchAndUpdateStatus(); }, 2000); return () => window.clearInterval(id); - }, [status?.is_loading, refreshStatus]); + }, [status?.is_loading, fetchAndUpdateStatus]); const handleLoad = useCallback(async () => { setBusy("loading");