diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4a6c45944e..718814ced0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10178,6 +10178,29 @@ async def _openai_passthrough_non_streaming( # ────────────────────────────────────────────────────────────────────────── +def _guard_diffusion_load_against_training() -> None: + """Refuse loading an image model while a training run is active. Unlike chat, + a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the + load is refused outright rather than fit-checked. No-op when training is + inactive or its state can't be read. Raises HTTP 409.""" + from core.training import get_training_backend + + try: + if not get_training_backend().is_training_active(): + return + except Exception as e: + logger.warning("Could not check training state for image-load guard: %s", e) + return + raise HTTPException( + status_code = 409, + detail = ( + "Can't load an image model while training is running: the diffusion " + "pipeline would compete with the training run for GPU memory. Training " + "was left untouched. Try again after training finishes." + ), + ) + + @studio_router.post("/images/load", response_model = DiffusionStatusResponse) async def load_diffusion_model( request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject) @@ -10196,6 +10219,10 @@ async def load_diffusion_model( gguf_filename = request.gguf_filename, family_override = request.family_override, ) + # Refuse while training is running: a multi-GB diffusion pipeline would + # compete with the training subprocess for VRAM. The chat path does the + # same via _guard_chat_load_against_training; this is its image sibling. + _guard_diffusion_load_against_training() # Take the GPU from the chat backend, then kick the (slow) load onto a # background thread and return at once — the client polls images/load-progress. await asyncio.to_thread(acquire_for, DIFFUSION) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index e223e8411c..d5fa9f2664 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3331,6 +3331,24 @@ async def delete_cached_model( except Exception: pass + # Also refuse if the diffusion (Images) backend has this repo loaded; its + # delete guard is otherwise chat-only, so its GGUF could be removed from + # under a live pipeline. Repo-level match, like the chat guards above. + try: + from core.inference.diffusion import get_diffusion_backend + diffusion_status = get_diffusion_backend().status() + if diffusion_status.get("loaded") and diffusion_status.get("repo_id"): + loaded_id = str(diffusion_status["repo_id"]).lower() + if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()): + raise HTTPException( + status_code = 400, + detail = "Unload the model before deleting", + ) + except HTTPException: + raise + except Exception: + pass + try: cache_scans = _all_hf_cache_scans() diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3818fe9f73..4056493af9 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -362,6 +362,25 @@ async def start_training( except Exception as e: logger.warning("Could not shut down export subprocess: %s", e) + try: + # A resident or in-flight diffusion (Images) pipeline also holds + # GPU memory the training run needs, and it can't be cheaply sized, + # so tear it down unconditionally like the export subprocess above + # (the chat block below fit-checks; diffusion can't). unload() is a + # no-op when nothing is loaded and also preempts an in-flight load; + # release the arbiter so it doesn't think the gone pipeline owns + # the GPU. Must precede the chat block, which early-returns. + from core.inference import gpu_arbiter + from core.inference.diffusion import get_diffusion_backend + + diffusion = get_diffusion_backend() + if diffusion.is_loaded: + logger.info("Unloading diffusion (Images) model to free GPU memory for training") + diffusion.unload() + gpu_arbiter.release(gpu_arbiter.DIFFUSION) + except Exception as e: + logger.warning("Could not unload diffusion model for training: %s", e) + try: from routes.training_vram import ( can_keep_chat_during_training, diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 046eb2543d..b8ffb5a3e6 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -706,3 +706,39 @@ def test_gguf_download_progress_counts_quant_subdir(monkeypatch, tmp_path): assert result["downloaded_bytes"] == 20_000 assert result["progress"] == 1.0 + + +def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch): + # The cached-delete guard refuses deleting a repo the diffusion (Images) + # backend has loaded, mirroring the chat guard, so its GGUF can't be removed + # from under a live pipeline. + from fastapi import HTTPException + import core.inference.diffusion as diffusion_mod + import routes.inference as routes_inference + + # Chat and orchestrator report nothing loaded; only diffusion holds the repo. + # delete_cached_model resolves get_inference_backend from the models module + # namespace, so patch it there (not on core.inference) to isolate that guard. + monkeypatch.setattr( + routes_inference, "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False, model_identifier = None), + ) + monkeypatch.setattr( + models_route, "get_inference_backend", + lambda: SimpleNamespace(active_model_name = None), + ) + monkeypatch.setattr( + diffusion_mod, "get_diffusion_backend", + lambda: SimpleNamespace(status = lambda: {"loaded": True, "repo_id": "org/Z-Image-GGUF"}), + ) + + try: + asyncio.run( + models_route.delete_cached_model( + repo_id = "org/Z-Image-GGUF", variant = None, current_subject = "u", + ) + ) + assert False, "expected HTTPException refusing the delete" + except HTTPException as e: + assert e.status_code == 400 + assert "Unload the model before deleting" in e.detail diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index f2b61e852f..878de0d477 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -282,6 +282,31 @@ def test_load_validation_failure_does_not_evict_chat(client, monkeypatch): assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT +def test_load_refused_during_training_does_not_evict_chat(client, monkeypatch): + # An image load while training is active is refused (409) before the GPU is + # taken, so the training run and the loaded chat model are both untouched. + import core.training as core_training + + monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT) + evicted = [] + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: evicted.append(True)) + + class _Training: + def is_training_active(self): + return True + + monkeypatch.setattr(core_training, "get_training_backend", lambda: _Training()) + + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}, + ) + assert resp.status_code == 409 + assert "training" in resp.json()["detail"].lower() + assert evicted == [] # chat backend was never evicted + assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT + + def test_load_progress_route(client): # Before load: idle. idle = client.get("/api/inference/images/load-progress")