From 899465ed8079ac164a7961ddb6974438687bb62c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 06:31:00 +0000 Subject: [PATCH] Studio: close arbiter load-registration race and surface native progress + local pipeline folders Publish native sd.cpp generate progress (_gen) before LoRA resolution so a reload probe reads active during setup, matching the diffusers path. Register the diffusion/video GPU load under the arbiter lock (acquire_for now takes a register callback) so a competing acquire cannot evict an owner before its load is marked in-flight and let two loaders allocate VRAM at once. Admit local diffusers pipeline folders (root model_index.json, weights in component subdirs) in the local model scan so they reach task tagging and the On Device picker. --- studio/backend/core/inference/gpu_arbiter.py | 16 ++++- .../backend/core/inference/sd_cpp_backend.py | 7 +- studio/backend/routes/inference.py | 52 ++++++++------ studio/backend/routes/models.py | 10 ++- studio/backend/routes/video.py | 44 +++++++----- studio/backend/tests/test_diffusion_routes.py | 8 ++- studio/backend/tests/test_gpu_arbiter.py | 72 +++++++++++++++++++ .../backend/tests/test_local_model_format.py | 19 +++++ studio/backend/tests/test_sd_cpp_backend.py | 26 +++++++ studio/backend/tests/test_video_routes.py | 8 ++- 10 files changed, 216 insertions(+), 46 deletions(-) diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index 0757eb83e8..53da970ce6 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -12,7 +12,7 @@ under the lock, so a transfer is atomic vs other acquires. from __future__ import annotations import threading -from typing import Optional +from typing import Any, Callable, Optional from loggers import get_logger @@ -64,8 +64,17 @@ def _evict_video() -> None: _EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion, VIDEO: _evict_video} -def acquire_for(owner: str) -> None: - """Make ``owner`` the sole GPU owner, evicting the other if it holds it.""" +def acquire_for(owner: str, register: Optional[Callable[[], Any]] = None) -> Any: + """Make ``owner`` the sole GPU owner, evicting the other if it holds it. + + ``register``, if given, runs under the arbiter lock right after ownership transfers, + and its return value is returned. Registering the in-flight load HERE -- not after + ``acquire_for`` returns -- closes the window where a competing acquire could evict this + owner before its load is marked in-flight: eviction would then find nothing to cancel + and both loaders would allocate VRAM at once. ``register`` must be quick (it holds the + lock) and must not re-enter the arbiter. If it raises, ownership stays with ``owner`` -- + matching the pre-register behaviour where a failed load left the handoff in place. + """ global _owner if owner not in _EVICTORS: raise ValueError(f"unknown GPU owner: {owner!r}") @@ -74,6 +83,7 @@ def acquire_for(owner: str) -> None: logger.info("gpu_arbiter: evicting %s for %s", _owner, owner) _EVICTORS[_owner]() _owner = owner + return register() if register is not None else None def release(owner: str) -> None: diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index a6847ae020..3fb40b5c7f 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -764,6 +764,12 @@ class SdCppDiffusionBackend: self._state = None raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) self._active_generate_cancel = cancel + # Publish an active (step 0) state now, before the slow pre-generate setup + # (LoRA listing/download), so a reload's progress probe doesn't read idle + # while this generation already holds _generate_lock and let a second generate + # queue behind it. The parsed sd-cli progress lines advance this step count. + # Mirrors DiffusionBackend.generate, which publishes _gen before its setup. + self._gen = _SdGen(total_steps = int(steps)) try: if seed is None: seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1) @@ -789,7 +795,6 @@ class SdCppDiffusionBackend: lora_resolved = diffusion_lora.resolve_specs( active_loras, hf_token = state.hf_token, cancel_event = cancel ) - self._gen = _SdGen(total_steps = int(steps)) if state.mode == "server" and state.server is not None: images, seeds = self._generate_server( state, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index fab3652b7f..20aca01fe6 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14338,10 +14338,36 @@ async def load_diffusion_model( # engine name -- else we'd evict a resident chat model for a load that can't use the GPU. device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) needs_gpu = device != "cpu" + + def _begin_load(): + # Kicks the (slow) load onto a background thread and returns at once (the client + # polls images/load-progress); begin_load itself validates network-free. + return engine.begin_load( + request.model_path, + gguf_filename = request.gguf_filename, + base_repo = request.base_repo, + family_override = request.family_override, + hf_token = request.hf_token, + cpu_offload = request.cpu_offload, + memory_mode = request.memory_mode, + speed_mode = request.speed_mode, + text_encoder_quant = request.text_encoder_quant, + transformer_quant = request.transformer_quant, + transformer_quant_fast_accum = request.transformer_quant_fast_accum, + transformer_prequant_path = request.transformer_prequant_path, + attention_backend = request.attention_backend, + transformer_cache = request.transformer_cache, + transformer_cache_threshold = request.transformer_cache_threshold, + model_kind = kind, + ) + if needs_gpu: - # 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) + # Register the in-flight load UNDER the arbiter lock (not after acquire_for + # returns): a competing Video/chat acquire in that gap would otherwise evict + # DIFFUSION before begin_load marks a load in-flight, so eviction finds nothing + # to cancel and both loaders allocate VRAM at once. begin_load returns at once, + # so the lock is held only briefly. + status_dict = await asyncio.to_thread(acquire_for, DIFFUSION, _begin_load) else: # A CPU-only native load never touches the GPU, so it neither acquires nor is # tracked by the arbiter. But switching here FROM a previous diffusers/GPU load @@ -14349,25 +14375,7 @@ async def load_diffusion_model( # "evict" this CPU model for no reason. Release that stale ownership -- release() # is owner-guarded, so it's a no-op when diffusion never owned the GPU. await asyncio.to_thread(release, DIFFUSION) - status_dict = await asyncio.to_thread( - engine.begin_load, - request.model_path, - gguf_filename = request.gguf_filename, - base_repo = request.base_repo, - family_override = request.family_override, - hf_token = request.hf_token, - cpu_offload = request.cpu_offload, - memory_mode = request.memory_mode, - speed_mode = request.speed_mode, - text_encoder_quant = request.text_encoder_quant, - transformer_quant = request.transformer_quant, - transformer_quant_fast_accum = request.transformer_quant_fast_accum, - transformer_prequant_path = request.transformer_prequant_path, - attention_backend = request.attention_backend, - transformer_cache = request.transformer_cache, - transformer_cache_threshold = request.transformer_cache_threshold, - model_kind = kind, - ) + status_dict = await asyncio.to_thread(_begin_load) return DiffusionStatusResponse(**annotate_status(status_dict)) except (ValueError, FileNotFoundError) as exc: raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index fe11446561..d1337e4210 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -338,7 +338,15 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca has_config = (child / "config.json").exists() or ( child / "adapter_config.json" ).exists() - has_model_files = has_gguf or has_non_gguf_weights or has_config + # A standard diffusers PIPELINE folder keeps its weights/configs in component + # subdirs (transformer/, vae/, ...) and carries only model_index.json at the + # root, so the checks above miss it. The Images/Video load path accepts such a + # local pipeline dir, so admit it here too (task tagging then classifies it via + # _local_is_diffusers); otherwise it is hidden from the On Device picker. + has_pipeline_index = (child / "model_index.json").is_file() + has_model_files = ( + has_gguf or has_non_gguf_weights or has_config or has_pipeline_index + ) except OSError: # Skip unreadable children rather than failing the scan. continue diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 7e239151c5..4abc513935 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -112,26 +112,36 @@ async def load_video_model( # Take the GPU from chat only for a non-CPU load; a CPU load never touches GPU memory, # so key off the device. Release stale VIDEO ownership on a CPU load (owner-guarded no-op). device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) + + def _begin_load(): + # Kicks the (slow) load onto a background thread and returns at once; + # begin_load itself validates network-free. + return backend.begin_load( + request.model_path, + gguf_filename = request.gguf_filename, + base_repo = request.base_repo, + family_override = request.family_override, + hf_token = request.hf_token, + memory_mode = request.memory_mode, + speed_mode = request.speed_mode, + attention_backend = request.attention_backend, + transformer_cache = request.transformer_cache, + transformer_cache_threshold = request.transformer_cache_threshold, + transformer_quant = request.transformer_quant, + text_encoder_quant = request.text_encoder_quant, + model_kind = kind, + ) + if device != "cpu": - await asyncio.to_thread(acquire_for, VIDEO) + # Register the in-flight load UNDER the arbiter lock (not after acquire_for + # returns): a competing Images/chat acquire in that gap would otherwise evict + # VIDEO before begin_load marks a load in-flight, so eviction finds nothing to + # cancel and both loaders allocate VRAM at once. begin_load returns at once, so + # the lock is held only briefly. Mirrors the images/load handoff. + status_dict = await asyncio.to_thread(acquire_for, VIDEO, _begin_load) else: await asyncio.to_thread(release, VIDEO) - status_dict = await asyncio.to_thread( - backend.begin_load, - request.model_path, - gguf_filename = request.gguf_filename, - base_repo = request.base_repo, - family_override = request.family_override, - hf_token = request.hf_token, - memory_mode = request.memory_mode, - speed_mode = request.speed_mode, - attention_backend = request.attention_backend, - transformer_cache = request.transformer_cache, - transformer_cache_threshold = request.transformer_cache_threshold, - transformer_quant = request.transformer_quant, - text_encoder_quant = request.text_encoder_quant, - model_kind = kind, - ) + status_dict = await asyncio.to_thread(_begin_load) return VideoStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index c00d9e71d5..322c188534 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -789,7 +789,13 @@ def _force_engine(monkeypatch, backend, *, engine_name, device): devmod, "resolve_diffusion_device_target", lambda: _types.SimpleNamespace(device = device) ) acquired: list = [] - monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role)) + + def _fake_acquire(role, register = None): + # Mirror the real arbiter: record the handoff and run the (registered) load under it. + acquired.append(role) + return register() if register is not None else None + + monkeypatch.setattr(gpu_arbiter, "acquire_for", _fake_acquire) return acquired diff --git a/studio/backend/tests/test_gpu_arbiter.py b/studio/backend/tests/test_gpu_arbiter.py index 0487830def..a73e786f90 100644 --- a/studio/backend/tests/test_gpu_arbiter.py +++ b/studio/backend/tests/test_gpu_arbiter.py @@ -104,3 +104,75 @@ def test_evict_chat_unloads_a_still_loading_chat_backend(monkeypatch): arb._evict_chat() assert unloaded == [True] # still-loading chat backend was unloaded, not skipped + + +def test_register_runs_under_ownership_and_returns_result(calls): + # A register callback runs after ownership transfers (owner already set) and its + # return value is forwarded -- the route uses this to register the in-flight load. + seen_owner: list = [] + + def register(): + seen_owner.append(arb.current_owner()) + return "status-dict" + + result = arb.acquire_for(arb.DIFFUSION, register) + assert result == "status-dict" + assert seen_owner == [arb.DIFFUSION] + assert arb.current_owner() == arb.DIFFUSION + + +def test_register_failure_leaves_ownership_in_place(calls): + # A failing register (e.g. begin_load reporting a load already in progress) propagates + # but must not drop ownership -- the prior handoff (chat already evicted) stands. + arb.acquire_for(arb.CHAT) + + def register(): + raise RuntimeError("A diffusion load is already in progress.") + + with pytest.raises(RuntimeError): + arb.acquire_for(arb.DIFFUSION, register) + assert calls == ["evict-chat"] + assert arb.current_owner() == arb.DIFFUSION + + +def test_competing_acquire_blocks_until_register_completes(monkeypatch): + # The window this closes: while DIFFUSION registers its load, a competing VIDEO acquire + # must not evict DIFFUSION until the load is marked in-flight. Holding the lock across + # register makes the competitor wait, so eviction never races an unregistered load. + import threading + import time + + monkeypatch.setattr(arb, "_owner", None) + evicted: list = [] + monkeypatch.setitem(arb._EVICTORS, arb.DIFFUSION, lambda: evicted.append("evict-diffusion")) + monkeypatch.setitem(arb._EVICTORS, arb.VIDEO, lambda: evicted.append("evict-video")) + + in_register = threading.Event() + release_register = threading.Event() + + def register(): + in_register.set() + # Hold the arbiter lock here; a competing acquire_for(VIDEO) must block until we return. + assert release_register.wait(2.0) + return "loading" + + loader = threading.Thread(target = lambda: arb.acquire_for(arb.DIFFUSION, register)) + loader.start() + assert in_register.wait(2.0) + + competitor_done = threading.Event() + threading.Thread( + target = lambda: (arb.acquire_for(arb.VIDEO), competitor_done.set()), + ).start() + + # The competitor cannot evict DIFFUSION while register still holds the lock. + time.sleep(0.1) + assert evicted == [] + assert not competitor_done.is_set() + + # Let register finish; ownership is now safely registered, so the competitor proceeds. + release_register.set() + loader.join(2.0) + assert competitor_done.wait(2.0) + assert evicted == ["evict-diffusion"] + assert arb.current_owner() == arb.VIDEO diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 202b9a1b72..1fbd7efcd6 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -117,6 +117,25 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): assert row.model_format == "gguf" +def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): + # A standard diffusers PIPELINE folder keeps its weights/configs in component subdirs + # (transformer/, vae/, ...) and carries only model_index.json at the root. The Images/Video + # load path accepts such a local pipeline dir, so the scan must surface it -- otherwise the + # weights-in-subdirs layout is missed and it never reaches task tagging / the On Device + # picker. It is not a GGUF, so model_format stays None (task tagging classifies it later). + root = tmp_path / "models" + pipe = root / "my-pipeline" + _touch(pipe / "model_index.json") + _touch(pipe / "transformer" / "config.json") + _touch(pipe / "transformer" / "diffusion_pytorch_model.safetensors") + _touch(pipe / "vae" / "diffusion_pytorch_model.safetensors") + + rows = {Path(m.path).name: m for m in models_route._scan_models_dir(root)} + + assert "my-pipeline" in rows + assert rows["my-pipeline"].model_format is None + + # ── Images picker task tag for local (non-GGUF) diffusers models ────────────── from models.models import LocalModelInfo # noqa: E402 diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 1d0268effa..1cb01df417 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -279,6 +279,32 @@ def test_generate_progress_tracks_parsed_steps(): assert b.generate_progress()["step"] == 4 +def test_generate_publishes_progress_before_lora_resolution(monkeypatch): + # Native LoRA resolution (listing/downloading a not-yet-cached adapter) happens during the + # pre-generate setup while _generate_lock is already held. A reload/progress probe in that + # window must read ACTIVE, not idle, or the UI queues a second generate behind the first. + # So _gen is published before LoRA resolution, mirroring the diffusers path. + from core.inference import diffusion_lora + + eng = _FakeEngine() + b = _loaded_backend(engine = eng) + monkeypatch.setattr(diffusion_lora, "supports_lora", lambda **_k: True) + + seen: dict = {} + + def _resolve(active, *, hf_token = None, cancel_event = None): + # Mid-setup: the in-flight generation must already be reported as active. + seen["progress"] = b.generate_progress() + return [] + + monkeypatch.setattr(diffusion_lora, "resolve_specs", _resolve) + + out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, loras = [("some/lora", 1.0)]) + assert out["images"] + assert seen["progress"]["active"] is True + assert seen["progress"]["total_steps"] == 8 + + # ── load validation + binary install ────────────────────────────────────────── diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 0793d13788..f7bc076225 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -263,7 +263,13 @@ def test_load_happy_path_and_arbiter_acquired(client, monkeypatch): devmod, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cuda") ) acquired: list = [] - monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role)) + + def _fake_acquire(role, register = None): + # Mirror the real arbiter: record the handoff and run the (registered) load under it. + acquired.append(role) + return register() if register is not None else None + + monkeypatch.setattr(gpu_arbiter, "acquire_for", _fake_acquire) resp = client.post( "/api/inference/video/load",