diff --git a/pyproject.toml b/pyproject.toml index a89d2c2a51..e974d8bb3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,7 +92,11 @@ huggingfacenotorch = [ # workflow raises without it), the cache_context child-registry behavior, and the # Flux2 / Z-Image pipelines. An unversioned entry let an upgrade keep an older # diffusers, so advertised models failed to load until the user upgraded by hand. - "diffusers>=0.39.0", + # diffusers dropped 3.9 in 0.38, and we still support it (requires-python >= 3.9), + # so the floor is conditional: pinning it outright leaves pip no candidate at all + # on 3.9 and the whole extra becomes unresolvable. + "diffusers>=0.39.0 ; python_version >= '3.10'", + "diffusers ; python_version < '3.10'", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", "trl>=0.18.2,!=0.19.0,<=0.24.0", "sentence-transformers", diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 4b0c9b9155..f0e6a54707 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -22,7 +22,7 @@ from __future__ import annotations import os import threading -from typing import Any, Optional +from typing import Any, Callable, Optional from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_families import DiffusionFamily, family_sd_cpp_supported @@ -124,6 +124,24 @@ def _activate(name: str, reason: Optional[str]) -> Any: return get_active_diffusion_engine() +def begin_load_on(expected_engine: Any, start: Callable[[], Any]) -> Any: + """Run ``start`` under the transition lock, refusing if the engine changed since selection. + + A load route selects its engine, then yields (device probe, arbiter acquire) before it + registers the load. A second /images/load picking the OTHER engine can transition in that + gap and unload the still-idle engine this request captured, which would then load a model + nothing can reach: generate / status / unload and the arbiter's evictor all resolve through + get_active_diffusion_engine(). Re-checking under the same lock the switch takes makes + selection and registration one operation. + """ + with _transition_lock: + if expected_engine is not get_active_diffusion_engine(): + raise RuntimeError( + "The diffusion engine changed while this load was starting. Retry the load." + ) + return start() + + def select_and_activate_engine( fam: DiffusionFamily, *, diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index 04a7cd08a0..2b96059756 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -32,10 +32,15 @@ def _evict_chat() -> None: from core.inference import get_inference_backend from routes.inference import get_llama_cpp_backend + from core.inference.llama_cpp import chat_load_active + llama = get_llama_cpp_backend() # is_active (process exists), not is_loaded (exists AND healthy): a chat model still starting # up holds VRAM but isn't healthy, so is_loaded would skip it and let the load race diffusion. - if llama.is_active: + # chat_load_active too: an HF load has no process until its GGUF downloaded, so is_active + # alone found nothing to cancel and let that load spawn onto the GPU we just granted away. + # unload_model sets the cancel event the download loop polls, so the pending load aborts. + if llama.is_active or chat_load_active(): llama.unload_model() orchestrator = get_inference_backend() if orchestrator.active_model_name: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 397cba8842..5b2f0fc459 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1382,6 +1382,36 @@ def gguf_load_in_flight(hf_repo: Optional[str]): _LOADS_IN_FLIGHT[key] = remaining +# Chat loads in flight, repo-agnostic (local paths and safetensors included). The repo-keyed +# counter above answers the download manager; this one answers the GPU arbiter, which has to know +# a chat load exists before llama-server is spawned. +_CHAT_LOADS_IN_FLIGHT = 0 + + +@contextlib.contextmanager +def chat_load_in_flight(): + """Track a chat load from before its download until the load call returns.""" + global _CHAT_LOADS_IN_FLIGHT + with _LOADS_IN_FLIGHT_LOCK: + _CHAT_LOADS_IN_FLIGHT += 1 + try: + yield + finally: + with _LOADS_IN_FLIGHT_LOCK: + _CHAT_LOADS_IN_FLIGHT = max(0, _CHAT_LOADS_IN_FLIGHT - 1) + + +def chat_load_active() -> bool: + """Whether a chat load is in flight, spawned or not. + + ``is_active`` only covers a live llama-server process, which an HF load does not have + until its GGUF finished downloading -- minutes during which the arbiter's evictor would + find nothing to cancel and grant the GPU to a second workload. + """ + with _LOADS_IN_FLIGHT_LOCK: + return _CHAT_LOADS_IN_FLIGHT > 0 + + def hf_gguf_load_in_flight(hf_repo: str) -> bool: """Return whether a GGUF load is active for hf_repo.""" key = (hf_repo or "").strip().lower() diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 3280b24346..226bcaa5ff 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -356,15 +356,49 @@ def bf16_unsupported_reason(resolved_family: str) -> Optional[str]: return None +def dit_accelerator_missing_reason(resolved_family: str) -> Optional[str]: + """Reason a DiT family cannot train on this host at all, else None. Never raises. + + nf4 is not a CPU fallback: the 4-bit base load goes through diffusers' bitsandbytes + quantizer, whose validate_environment raises "No GPU found. A GPU is needed for + quantization." unless CUDA, XPU or MPS is present. Without this gate a GPU-less host + accepts the default nf4 start, evicts the resident Images pipeline, downloads the text + encoders, and only then dies in the child. SDXL keeps its own fp32-on-CPU path. + """ + if (resolved_family or "").strip().lower() not in _DIT_TRAIN_FAMILIES: + return None + try: + import torch + xpu = getattr(torch, "xpu", None) + mps = getattr(torch, "mps", None) + if ( + torch.cuda.is_available() + or bool(xpu is not None and xpu.is_available()) + or bool(mps is not None and mps.is_available()) + ): + return None + except Exception: # noqa: BLE001 -- probe failure must not block a start + return None + return ( + "Training the DiT families needs a GPU: even the 4-bit (nf4) base load requires " + "CUDA, XPU or MPS, and this host has none. Train SDXL here, or use a GPU machine." + ) + + def training_precision_preflight_error(resolved_family: str, base_precision: str) -> Optional[str]: """Reason the requested DiT precision cannot run on this host, else None -- checked by the start route BEFORE evicting resident GPU workloads (the trainer's own checks fire only in the child, after eviction). Four gates, all mirroring _resolve_base_precision so a doomed run is - rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); the dense + rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); a host with no + accelerator at all (dit_accelerator_missing_reason, which covers nf4 too); the dense precisions (bf16/int8/fp8/mxfp8) requiring a CUDA GPU; an explicit int8 needing a FUNCTIONAL torchao (its _int8_quantize_base has no fallback); and an explicit mxfp8 needing a Blackwell (sm100+) GPU (its MX GEMM has no kernel below sm100). Never raises.""" reason = bf16_unsupported_reason(resolved_family) + if reason: + return reason + # No accelerator at all: every DiT precision is out, nf4 included. + reason = dit_accelerator_missing_reason(resolved_family) if reason: return reason fam = (resolved_family or "").strip().lower() @@ -430,7 +464,11 @@ def family_train_infos() -> list[dict[str, Any]]: # DiT option that always 400s. Otherwise drop any scheme this family's DiT corrupts (fp8 on # Qwen-Image: activation outliers exceed fp8's range; the inference path denies the same # set), so the UI never offers a mode normalized() would reject. - dit_block = bf16_unsupported_reason(name) if is_dit else None + dit_block = ( + bf16_unsupported_reason(name) or dit_accelerator_missing_reason(name) + if is_dit + else None + ) if not is_dit or dit_block: fam_modes: list[str] = [] else: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1d8cd4a17f..e23eeda7ac 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4395,7 +4395,7 @@ async def _load_model_impl( # validation so a doomed load (bad id, unsupported gpu_ids on GGUF, training 409) can't # evict a working image/video model and then error. Mirrors the image/video loaders, # which validate before acquire_for. - from core.inference.gpu_arbiter import acquire_for, CHAT + from core.inference.gpu_arbiter import acquire_for, current_owner, CHAT # ── Already-loaded check: skip reload if the exact model is active ── backend = get_inference_backend() @@ -4606,44 +4606,56 @@ async def _load_model_impl( gpu_memory_mode = request.gpu_memory_mode, ) + # Mark the load and refuse one the download manager already owns, BEFORE the eviction + # below: this 409 leaves nothing loaded, so checking it afterwards destroyed a working + # Images/Video pipeline for a load that could never start. The marker/check order is the + # handshake with the download manager, so both move together. It also has to run after + # pass-through argument inheritance, since a carried --no-mmproj changes the companion + # requirement exactly as it does for the load. + if config.is_gguf and config.gguf_hf_repo: + from core.inference.llama_cpp import gguf_load_in_flight + gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) + + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + if await asyncio.to_thread( + _hub_download_blocks_gguf_load, + config.gguf_hf_repo, + config.gguf_variant, + require_mmproj = bool( + config.is_vision and not extra_args_disable_mmproj(extra_llama_args) + ), + hf_token = request.hf_token, + ): + raise HTTPException( + status_code = 409, + detail = ( + f"'{model_log_label}' is currently being downloaded " + "by the download manager. Wait for the download to " + "finish (or cancel it), then load the model." + ), + ) + # Load now known viable (valid identifier, gpu_ids ok, fits alongside any active # training): reclaim the GPU for chat, evicting a resident Images/Video pipeline. Doing # this only here -- not before the validation above -- keeps a doomed load from evicting # a working image/video model and then erroring. No-op when chat already owns the GPU. - await asyncio.to_thread(acquire_for, CHAT) + # The in-flight marker is entered UNDER the arbiter lock (the `register` hook), as the + # image and video loads do: a chat load holds no llama-server process until its GGUF has + # downloaded, so a competing Images/Video acquire in that window found nothing to evict + # and both then allocated VRAM at once. With the marker, that evictor cancels this load. + from core.inference.llama_cpp import chat_load_in_flight + + await asyncio.to_thread( + acquire_for, + CHAT, + lambda: gguf_load_stack.enter_context(chat_load_in_flight()), + ) # ── GGUF path: load via llama-server ────────────────────── if config.is_gguf: llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - if config.gguf_hf_repo: - from core.inference.llama_cpp import gguf_load_in_flight - gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) - - # Block cache writes that would race the download manager. This runs - # after pass-through argument inheritance so a carried --no-mmproj - # changes the companion requirement exactly as it does for the load. - if config.gguf_hf_repo: - from core.inference.llama_cpp import _hub_download_blocks_gguf_load - if await asyncio.to_thread( - _hub_download_blocks_gguf_load, - config.gguf_hf_repo, - config.gguf_variant, - require_mmproj = bool( - config.is_vision and not extra_args_disable_mmproj(extra_llama_args) - ), - hf_token = request.hf_token, - ): - raise HTTPException( - status_code = 409, - detail = ( - f"'{model_log_label}' is currently being downloaded " - "by the download manager. Wait for the download to " - "finish (or cancel it), then load the model." - ), - ) - # Keep the resident model alive until every active generation finishes; # the caller's lifecycle gate blocks new starts. await _wait_for_model_switch_idle(current_request_counted = current_request_counted) @@ -4799,6 +4811,20 @@ async def _load_model_impl( detail = f"Failed to load GGUF model: {model_log_label if native_grant_backed else config.display_name}", ) + # An Images/Video acquire can land in the gap between the acquire above and + # load_model clearing the cancel event, so its cancellation is lost and this load + # spawns anyway. Ownership survives that gap: whoever took the GPU keeps it, and this + # load undoes itself rather than leaving two models resident on one device. + if current_owner() != CHAT: + await asyncio.to_thread(llama_backend.unload_model) + raise HTTPException( + status_code = 409, + detail = ( + "An image or video model took the GPU while this model was loading, " + "so the load was cancelled. Unload that model, then try again." + ), + ) + logger.info( f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}" ) @@ -16028,6 +16054,7 @@ async def load_diffusion_model( from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_engine_router import ( annotate_status, + begin_load_on, select_and_activate_engine, ) from core.inference.gpu_arbiter import acquire_for, release, DIFFUSION @@ -16077,7 +16104,7 @@ async def load_diffusion_model( device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) needs_gpu = device != "cpu" - def _begin_load(): + def _start_engine_load(): # Kicks the (slow) load onto a background thread and returns at once (client polls # images/load-progress); begin_load itself validates network-free. return engine.begin_load( @@ -16100,6 +16127,12 @@ async def load_diffusion_model( loras = [(s.id, s.weight) for s in request.loras] if request.loras else None, ) + def _begin_load(): + # Under the router's transition lock, refusing if a competing load switched engines + # since select_and_activate_engine above: begin_load on a deactivated engine leaves a + # resident model that generate / status / unload and the evictor can no longer reach. + return begin_load_on(engine, _start_engine_load) + if needs_gpu: # Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): # otherwise a competing Video/chat acquire in that gap evicts DIFFUSION before the load diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 630f22cf37..5316c79217 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -219,17 +219,28 @@ def test_training_precision_preflight_error(monkeypatch): assert training_precision_preflight_error("sdxl", "int8") is None assert training_precision_preflight_error("", "int8") is None - # On a CUDA-ABSENT host, bf16_unsupported_reason exempts CPU-only, but the DiT trainer's dense - # precisions still require CUDA (mirroring _resolve_base_precision), so bf16/int8/fp8/mxfp8 for - # a DiT family are rejected UP FRONT rather than after eviction. nf4/auto (and SDXL) still pass. + # On a host with NO accelerator every DiT precision is rejected up front rather than after + # eviction -- nf4 included, since its 4-bit load goes through bitsandbytes, which needs a GPU. + # SDXL (its own fp32-on-CPU path) still passes. monkeypatch.setattr(common, "has_functional_torchao", lambda: True) monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(torch.xpu, "is_available", lambda: False) + monkeypatch.setattr(torch.mps, "is_available", lambda: False) for dense in ("bf16", "int8", "fp8", "mxfp8"): reason = training_precision_preflight_error("flux.1", dense) assert reason is not None and "CUDA" in reason + for mode in ("nf4", "auto"): + reason = training_precision_preflight_error("flux.1", mode) + assert reason is not None and "GPU" in reason + assert training_precision_preflight_error("sdxl", "bf16") is None + + # An accelerator that is not CUDA (XPU here) satisfies the 4-bit load, so nf4/auto pass again + # while the dense CUDA-only precisions stay rejected. + monkeypatch.setattr(torch.xpu, "is_available", lambda: True) assert training_precision_preflight_error("flux.1", "nf4") is None assert training_precision_preflight_error("flux.1", "auto") is None - assert training_precision_preflight_error("sdxl", "bf16") is None + assert training_precision_preflight_error("flux.1", "bf16") is not None + monkeypatch.setattr(torch.xpu, "is_available", lambda: False) # mxfp8 needs a Blackwell (sm100+) GPU: its MX GEMM has no kernel below sm100 and would raise at # the first training step, AFTER a full dense-transformer load. On a CUDA GPU that is older than @@ -609,3 +620,36 @@ def test_assert_trusted_base_model_rejects_local_non_pipeline(tmp_path): # An untrusted remote base is still rejected by the trust gate. with pytest.raises(ValueError, match = "untrusted"): common._assert_trusted_base_model("evil/base") + +def test_dit_accelerator_missing_reason_and_info_hide_train_without_a_gpu(monkeypatch): + # Clicking Start on a GPU-less host evicted the resident Images pipeline, downloaded the text + # encoders, and only then died in the child: diffusers' bitsandbytes quantizer refuses 4-bit + # without CUDA/XPU/MPS. Reject it up front, and stop /info advertising a mode that always 400s. + import torch + + from core.training.diffusion_train_common import ( + _DIT_TRAIN_FAMILIES, + dit_accelerator_missing_reason, + family_train_infos, + ) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(torch.xpu, "is_available", lambda: False) + monkeypatch.setattr(torch.mps, "is_available", lambda: False) + assert "GPU" in (dit_accelerator_missing_reason("flux.1") or "") + # SDXL and unknown families keep their own paths. + assert dit_accelerator_missing_reason("sdxl") is None + assert dit_accelerator_missing_reason("") is None + + infos = {info["name"]: info for info in family_train_infos()} + for name, info in infos.items(): + if name in _DIT_TRAIN_FAMILIES: + assert info["precision_modes"] == [] + assert "GPU" in info["vram_note"] + assert info["supports_compile"] is False + else: + assert info["precision_modes"] != [] or name not in _DIT_TRAIN_FAMILIES + + # Any accelerator clears it (MPS here, which bitsandbytes accepts). + monkeypatch.setattr(torch.mps, "is_available", lambda: True) + assert dit_accelerator_missing_reason("flux.1") is None diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 18b9f51b62..dfcfa52be5 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -283,3 +283,48 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch): t.join(2.0) q.join(2.0) assert switch_done.is_set() and query_done.is_set() + + +def test_begin_load_on_refuses_an_engine_that_was_switched_away(monkeypatch): + # A load selects its engine, then yields (device probe, arbiter acquire) before registering. + # A second load choosing the OTHER engine transitions in that gap and unloads the still-idle + # engine this one captured, so registering there would leave a model that generate / status / + # unload and the arbiter's evictor can no longer reach (they all resolve the ACTIVE engine). + diffusers = SimpleNamespace(name = "diffusers") + sd_cpp = SimpleNamespace(name = "sd_cpp") + active = {"engine": diffusers} + monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: active["engine"]) + + started: list[str] = [] + assert r.begin_load_on(diffusers, lambda: started.append("ok") or "status") == "status" + assert started == ["ok"] + + # A competing request switched the active engine after this one captured `diffusers`. + active["engine"] = sd_cpp + with pytest.raises(RuntimeError, match = "engine changed"): + r.begin_load_on(diffusers, lambda: started.append("leaked")) + assert started == ["ok"] + + +def test_begin_load_on_holds_the_transition_lock_while_registering(monkeypatch): + # The check and the registration must be one operation: taken under the same lock a switch + # takes, so no _activate can slip between them. + import threading + + engine = SimpleNamespace(name = "diffusers") + monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: engine) + + inside = threading.Event() + release = threading.Event() + + def _slow_start(): + inside.set() + release.wait(2.0) + return "status" + + t = threading.Thread(target = lambda: r.begin_load_on(engine, _slow_start)) + t.start() + assert inside.wait(2.0) + assert r._transition_lock.locked() + release.set() + t.join(2.0) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index b6dc5bbd7a..28f13cd133 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -1583,6 +1583,34 @@ def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatc assert client._fake.started_with is None +def test_route_start_refuses_a_dit_family_without_a_gpu_before_freeing(client, monkeypatch): + # nf4 is not a CPU fallback: the 4-bit base load needs CUDA/XPU/MPS. Without a pre-teardown + # gate the default Train pick on a GPU-less host unloaded the working Images pipeline, pulled + # the text encoders, and only then failed in the child. + import torch + + import routes.training as tr + + freed = [] + monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(torch.xpu, "is_available", lambda: False) + monkeypatch.setattr(torch.mps, "is_available", lambda: False) + + r = client.post( + "/api/train/diffusion/start", + json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"}, + ) + assert r.status_code == 400 + assert "GPU" in r.json()["detail"] + assert freed == [] + assert client._fake.started_with is None + + # SDXL trains fp32 on CPU (that is its documented fallback), so it is not gated. + r2 = client.post("/api/train/diffusion/start", json = _BODY) + assert r2.status_code == 200, r2.text + + def test_start_ungated_base_preflight_is_noop(client, monkeypatch): # A reachable base (HEAD 200) proceeds to start normally. import urllib.request @@ -1590,6 +1618,12 @@ def test_start_ungated_base_preflight_is_noop(client, monkeypatch): import routes.training as tr monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None) + # This asserts the gated-repo probe is a no-op, not device support: pin the precision + # preflight so a GPU-less test host does not 400 for an unrelated reason. + monkeypatch.setattr( + "core.training.diffusion_train_common.training_precision_preflight_error", + lambda fam, prec: None, + ) monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: object()) r = client.post( "/api/train/diffusion/start", diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 0ab998af39..8fb49e24d7 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -628,6 +628,24 @@ class TestLoadHubDownloadExclusion: with gguf_load_in_flight(None): assert not hf_gguf_load_in_flight("") + def test_chat_load_marker_is_repo_agnostic_and_nests(self): + # The GPU arbiter needs to know a chat load exists before llama-server is spawned, for + # local paths and safetensors too, so this marker carries no repo key. + from core.inference.llama_cpp import chat_load_active, chat_load_in_flight + + assert not chat_load_active() + with chat_load_in_flight(): + assert chat_load_active() + with chat_load_in_flight(): + assert chat_load_active() + assert chat_load_active() + assert not chat_load_active() + + with pytest.raises(RuntimeError): + with chat_load_in_flight(): + raise RuntimeError("boom") + assert not chat_load_active() + def test_marker_cleared_on_exception(self): with pytest.raises(RuntimeError): with gguf_load_in_flight(REPO): @@ -808,27 +826,26 @@ class TestLoadHubDownloadExclusion: asyncio.run(scenario()) - def test_load_marker_precedes_hub_guard_and_unload(self): + def test_load_marker_precedes_hub_guard_which_precedes_the_gpu_handoff(self): source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() - # _load_model_impl has more than one `if config.is_gguf:`, so anchor on - # the branch that actually owns the load marker rather than the first - # one in the file, which belongs to an earlier check. - marker = source.index("enter_context(gguf_load_in_flight") - gguf_branch_start = source.rindex("if config.is_gguf:", 0, marker) - gguf_branch = source[gguf_branch_start:] + impl = source[source.index("async def _load_model_impl") :] - # The gguf_load_in_flight marker must be entered before the hub-download - # guard and the unload so a concurrent load can't race the download - # manager. The llama_extra_args inheritance moved out of the branch into - # _resolve_inherited_extra_args, which must still run BEFORE it: the - # inherited value (e.g. a carried --no-mmproj) shapes the guard's - # require_mmproj. Anchor on the call form so the assertion pins the - # endpoint's call site, not the function definition. - assert source.index("= _resolve_inherited_extra_args(") < gguf_branch_start + # One chain, in this order: + # - _resolve_inherited_extra_args first: the inherited value (e.g. a carried + # --no-mmproj) shapes the guard's require_mmproj. + # - the gguf_load_in_flight marker before the hub-download guard: that pair is the + # handshake that keeps a load and the download manager off the same files. + # - both before the CHAT handoff: the guard's 409 loads nothing, so checking it after + # the handoff destroyed a resident Images/Video pipeline for a load that could never + # start. The handoff registers its own marker under the arbiter lock. + # - the resident unload last. + # Anchored on call forms so each assertion pins a call site, not a definition. assert ( - gguf_branch.index("enter_context(gguf_load_in_flight") - < gguf_branch.index("_hub_download_blocks_gguf_load") - < gguf_branch.index("unsloth_backend.unload_model") + impl.index("= _resolve_inherited_extra_args(") + < impl.index("enter_context(gguf_load_in_flight") + < impl.index("_hub_download_blocks_gguf_load") + < impl.index("enter_context(chat_load_in_flight") + < impl.index("unsloth_backend.unload_model") ) llama_source = ( Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" diff --git a/studio/backend/tests/test_gpu_arbiter.py b/studio/backend/tests/test_gpu_arbiter.py index eabd4747cf..e122585e42 100644 --- a/studio/backend/tests/test_gpu_arbiter.py +++ b/studio/backend/tests/test_gpu_arbiter.py @@ -203,3 +203,49 @@ def test_competing_acquire_blocks_until_register_completes(monkeypatch): assert competitor_done.wait(2.0) assert evicted == ["evict-diffusion"] assert arb.current_owner() == arb.VIDEO + + +def test_evict_chat_cancels_a_chat_load_that_has_not_spawned_yet(monkeypatch): + # An HF chat load has no llama-server process until its GGUF finished downloading, which is + # minutes. Gating only on is_active let the evictor find nothing to cancel, grant the GPU to + # the image/video load, and the chat load then spawned onto the same device: two big models + # allocating at once. The in-flight marker is what makes that load cancellable. + import core.inference as core_inference + import routes.inference as routes_inference + from core.inference.llama_cpp import chat_load_in_flight + + unloaded: list[bool] = [] + + class _FakeLlama: + is_active = False # nothing spawned yet: the download is still running + is_loaded = False + + def unload_model(self): + unloaded.append(True) + + def _wait_for_vram_settle(self, *, since_kill): + pass + + class _FakeOrchestrator: + active_model_name = None + + def unload_model(self, name): + pass + + def _shutdown_subprocess(self, timeout = 5.0): + pass + + monkeypatch.setattr(routes_inference, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(core_inference, "get_inference_backend", lambda: _FakeOrchestrator()) + + # No load in flight: nothing to cancel. + arb._evict_chat() + assert unloaded == [] + + with chat_load_in_flight(): + arb._evict_chat() + assert unloaded == [True] + + # The marker is released with the load, so a later eviction is a no-op again. + arb._evict_chat() + assert unloaded == [True] diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 26e43448c9..d453c8f220 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -1127,6 +1127,7 @@ class TestRouteErrors(unittest.TestCase): gguf_hf_repo = None, gguf_file = "/tmp/test.gguf", gguf_mmproj_file = None, + gguf_mtp_file = None, gguf_variant = None, identifier = "unsloth/test.gguf", display_name = "unsloth/test.gguf", @@ -1136,6 +1137,10 @@ class TestRouteErrors(unittest.TestCase): has_audio_input = False, ) acquired = [] + # Make [0, 1] invalid on any host (a duplicate id is rejected everywhere, GPUs or not): + # the point of the test is the ORDER -- validation before the handoff -- not this + # machine's device count. + request.gpu_ids = [0, 0] with ( patch.object( inference_route, @@ -1145,7 +1150,10 @@ class TestRouteErrors(unittest.TestCase): patch.object(inference_route, "_guard_chat_load_against_training", return_value = None), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), - patch.object(arb, "acquire_for", lambda owner: acquired.append(owner)), + # The chat handoff passes a `register` hook (the in-flight marker), so accept it. + patch.object( + arb, "acquire_for", lambda owner, register = None: acquired.append(owner) + ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -1162,6 +1170,126 @@ class TestRouteErrors(unittest.TestCase): self.assertEqual(exc_info.exception.status_code, 400) self.assertEqual(acquired, []) # no CHAT handoff before the doomed load errored + def test_inference_route_checks_hub_download_conflict_before_the_handoff(self): + # A GGUF the download manager is fetching 409s and loads nothing, so that check has to + # run BEFORE the CHAT handoff: afterwards, it destroyed the resident Images/Video + # pipeline for a load that could never start. + import core.inference.gpu_arbiter as arb + import core.inference.llama_cpp as llama_cpp + + inference_route = _load_route_module( + "inference_route_module_for_hub_conflict_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/Qwen3-4B-GGUF", gguf_variant = "Q4_K_M") + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = "unsloth/Qwen3-4B-GGUF", + gguf_file = None, + gguf_mmproj_file = None, + gguf_variant = "Q4_K_M", + identifier = "unsloth/Qwen3-4B-GGUF", + display_name = "Qwen3-4B", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + acquired = [] + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch.object(inference_route, "_guard_chat_load_against_training", return_value = None), + patch.object(inference_route, "_resolve_inherited_extra_args", lambda *a, **k: None), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(llama_cpp, "_hub_download_blocks_gguf_load", lambda *a, **k: True), + patch.object(arb, "acquire_for", lambda *a, **k: acquired.append(a[0])), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route._load_model_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + self.assertEqual(exc_info.exception.status_code, 409) + self.assertIn("download", exc_info.exception.detail.lower()) + self.assertEqual(acquired, []) # nothing evicted for a load that cannot start + + def test_inference_route_marks_the_chat_load_under_the_arbiter_lock(self): + # A chat load holds no llama-server process until its GGUF downloaded, so the arbiter is + # told about it through acquire_for's `register` hook (which runs under the arbiter lock). + # Passing no register left a competing Images/Video acquire with nothing to cancel. + import core.inference.gpu_arbiter as arb + import core.inference.llama_cpp as llama_cpp + + inference_route = _load_route_module( + "inference_route_module_for_chat_marker_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/Qwen3-4B-GGUF", gguf_variant = "Q4_K_M") + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = "unsloth/Qwen3-4B-GGUF", + gguf_file = None, + gguf_mmproj_file = None, + gguf_variant = "Q4_K_M", + identifier = "unsloth/Qwen3-4B-GGUF", + display_name = "Qwen3-4B", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + marked = [] + + def _acquire(owner, register = None): + # Under the arbiter lock the evictor must already be able to see this load. + if register is not None: + register() + marked.append(llama_cpp.chat_load_active()) + raise RuntimeError("stop the load here") + + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch.object(inference_route, "_guard_chat_load_against_training", return_value = None), + patch.object(inference_route, "_resolve_inherited_extra_args", lambda *a, **k: None), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(llama_cpp, "_hub_download_blocks_gguf_load", lambda *a, **k: False), + patch.object(arb, "acquire_for", _acquire), + ): + with self.assertRaises(HTTPException): + asyncio.run( + inference_route._load_model_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + self.assertEqual(marked, [True]) + # The marker is scoped to the request: it must not outlive the failed load. + self.assertFalse(llama_cpp.chat_load_active()) + def test_training_route_returns_400_for_invalid_gpu_ids(self): training_route = _load_route_module( "training_route_module_for_test", diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 251c2d0fdf..952e07a74b 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1838,12 +1838,16 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // to load the other page's checkpoint as its own kind of model. if (!active) return; const wanted = routeSearch.model; - // Key on the model AND the quant: this page stays mounted across navigation, so a - // model-only marker made every later pick of the same repo a no-op -- including a - // different quant, or the same one after chat evicted it -- without even clearing - // the query string, so the click did nothing at all. - const key = wanted ? `${wanted}${routeSearch.quant ?? ""}` : null; - if (!wanted || !key || handledRouteModel.current === key) return; + // Key on the model AND the quant, and release the marker once the query is gone. This + // page stays mounted, so a marker that outlived the query made re-picking the same + // checkpoint (after chat evicted it) a click that neither loaded nor cleared the URL. + // It only has to survive the re-renders between consuming a pick and the clear landing. + if (!wanted) { + handledRouteModel.current = null; + return; + } + const key = `${wanted}|${routeSearch.quant ?? ""}`; + if (handledRouteModel.current === key) return; handledRouteModel.current = key; void navigateSelf({ to: "/images", search: {}, replace: true }); void loadOrStage( diff --git a/studio/frontend/src/features/video/video-page.tsx b/studio/frontend/src/features/video/video-page.tsx index 372b698fe0..e6d695e655 100644 --- a/studio/frontend/src/features/video/video-page.tsx +++ b/studio/frontend/src/features/video/video-page.tsx @@ -1170,10 +1170,15 @@ export function VideoPage({ active = true }: { active?: boolean }) { // load an image checkpoint as a video model. if (!active) return; const wanted = routeSearch.model; - // Model AND quant, as on the Images page: this page stays mounted, so a model-only - // marker turned every later pick of the same repo into a silent no-op. - const key = wanted ? `${wanted} ${routeSearch.quant ?? ""}` : null; - if (!wanted || !key || handledRouteModel.current === key) return; + // Model AND quant, released once the query is gone, as on the Images page: this page + // stays mounted, so a marker that outlived the query turned re-picking the same + // checkpoint into a click that neither loaded nor cleared the URL. + if (!wanted) { + handledRouteModel.current = null; + return; + } + const key = `${wanted}|${routeSearch.quant ?? ""}`; + if (handledRouteModel.current === key) return; handledRouteModel.current = key; void navigateSelf({ to: "/video", search: {}, replace: true }); void loadOrStage(