From cf0f5d55042812da174f652c9718248408f68d01 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 20:57:44 +0000 Subject: [PATCH] Serialize video-load GPU placement with eviction; skip LoRA adapters as base picks Two evict/OOM fixes on the diffusion load paths: - The video load moved a pipeline onto the GPU (apply_memory_plan) and committed it while holding no lock, so an unload / GPU-arbiter eviction -- which bumps the load token and then barriers on _generate_lock before freeing -- could hand VIDEO to chat/images and let the new owner allocate concurrently with the in-flight placement, OOMing. Hold _generate_lock across placement + the locked commit, mirroring the image backend, so an evicting owner waits until this worker's placement is torn down or committed. Lock order stays _generate_lock -> _lock (unload takes _lock then releases it before the barrier), so there is no deadlock. - resolve_local_single_file reinterpreted an On-Device folder as a base single_file load whenever it held exactly one .safetensors, so a PEFT LoRA adapter folder (adapter_config.json + adapter_model.safetensors) with a family-token name was picked as a base checkpoint, evicting the resident model before from_single_file failed on the adapter weights. Skip adapter folders (adapter_config.json) and the adapter_model basename so the pick stays a pipeline load and 400s in validation, before the GPU handoff. Adds regression tests for both. --- studio/backend/core/inference/diffusion.py | 15 +- studio/backend/core/inference/video.py | 201 +++++++++--------- .../backend/tests/test_diffusion_backend.py | 15 ++ studio/backend/tests/test_video_backend.py | 56 +++++ 4 files changed, 185 insertions(+), 102 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index c04907f9c5..c09f12231a 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -159,13 +159,24 @@ def resolve_local_single_file(model_path: str) -> Optional[str]: advertised model is unusable. The images load route uses this to reinterpret such a pick as a ``single_file`` load of the sole checkpoint. A real pipeline dir (has ``model_index.json``) or an ambiguous one (0 or more than 1 ``.safetensors``, e.g. a sharded pipeline) returns None and - loads unchanged. Never raises.""" + loads unchanged. A PEFT LoRA adapter folder is also skipped (see below). Never raises.""" try: root = Path(model_path).expanduser() if not root.is_dir() or (root / "model_index.json").is_file(): return None + # A PEFT LoRA adapter folder (adapter_config.json + adapter_model.safetensors) is not a + # base checkpoint: from_single_file would fail on the adapter weights AFTER the route + # evicted the resident GPU model. Skip it so the pick stays a pipeline load and 400s in + # validation, before the GPU handoff. Also drop a bare adapter_model.safetensors so a + # config-less adapter export is never reinterpreted as the sole checkpoint. + if (root / "adapter_config.json").is_file(): + return None checkpoints = [ - p.name for p in root.iterdir() if p.is_file() and p.suffix.lower() == ".safetensors" + p.name + for p in root.iterdir() + if p.is_file() + and p.suffix.lower() == ".safetensors" + and p.stem.lower() != "adapter_model" ] except OSError: return None diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 4e72a6fb1f..8839a8ed36 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -1291,111 +1291,112 @@ class VideoBackend: if view is pipe: attention_engaged = engaged speed_optims = tuple(k for k, v in applied.items() if v) - # A cancelled/superseded load must not place weights on the GPU the arbiter - # may already have handed to another backend; recheck right before placement - # (the commit below still does the final locked check). - if _load_token is not None and _load_token != self._load_token: - del pipe - clear_gpu_cache() - raise RuntimeError("Video load was cancelled or superseded.") - offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger) - # A dual-DiT MoE pipe (Wan2.2-A14B) needs no extra per-expert offload pass here: - # apply_memory_plan's group tier (_apply_group_offload) already block-streams every - # DiT it finds on the pipe -- transformer AND transformer_2 -- and model/sequential - # offload hook every top-level module, so the second expert is covered under all tiers. - # A second _apply_group_offload on transformer_2 would re-register the group-offload - # hooks it already carries, which diffusers rejects with a duplicate-hook ValueError. - if not vae_tiling: - # Decode of a whole clip is the video memory peak; tiling is near-free - # in quality and keeps the decode bounded, so it is always on. - try: - pipe.vae.enable_tiling() - vae_tiling = True - except Exception as exc: # noqa: BLE001 -- tiling is an optimisation only - logger.warning("video.vae_tiling_failed: %s", exc) - - resolved = build_resolved_record( - { - "memory_mode": ( - memory_mode, - plan.requested_mode, - f"planned '{plan.offload_policy}' offload from the family size table", - ), - "speed_mode": ( - speed_mode, - effective_speed, - "quantized transformer requires compile" - if transformer_quant_engaged is not None - else "clip denoises amortise the one-time compile within a single run" - if speed_mode is None - else "requested", - ), - "attention_backend": ( - attention_backend, - attention_engaged or "native", - "cuDNN fused attention on NVIDIA when a speed profile is active", - ), - "transformer_cache": ( - None if cache_auto else transformer_cache, - cache_engaged or "off", - cache_reason, - ), - "transformer_quant": ( - transformer_quant, - transformer_quant_engaged or "off", - "dense DiT(s) torchao-quantised onto the low-precision tensor cores" - if transformer_quant_engaged is not None - else ( - "skipped: offload moves the DiT, unsupported for torchao " - "tensors; pin a resident memory mode to combine them" - if quant_skipped_for_offload - else "not engaged (dense bf16 DiT loaded)" - ), - ), - "text_encoder_quant": ( - text_encoder_quant, - text_encoder_quant_engaged or "off", - "dense text encoder quantised in place" - if text_encoder_quant_engaged is not None - else "not engaged (dense bf16 text encoder loaded)", - ), - } - ) - - with self._lock: + with self._generate_lock: + # A cancelled/superseded load must not place weights on the GPU the arbiter + # may already have handed to another backend; recheck right before placement + # (the commit below still does the final locked check). if _load_token is not None and _load_token != self._load_token: del pipe clear_gpu_cache() raise RuntimeError("Video load was cancelled or superseded.") - self._state = _VideoLoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - kind = kind, - gguf_filename = gguf_filename, - offload_policy = offload_policy, - vae_tiling = vae_tiling, - memory_mode = plan.requested_mode, - speed_mode = effective_speed, - # Already filtered above to only the optimisations that engaged; - # apply_speed_optims returns every flag True/False and the view - # loop keeps just the True names. - speed_optims = speed_optims, - backend_flags = backend_flags, - attention_backend = attention_engaged, - transformer_cache = cache_engaged, - cache_auto = cache_may_toggle, - cache_quant_active = cache_quant_active, - cache_threshold = transformer_cache_threshold, - transformer_quant = transformer_quant_engaged, - text_encoder_quant = text_encoder_quant_engaged, - resolved = resolved, + offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger) + # A dual-DiT MoE pipe (Wan2.2-A14B) needs no extra per-expert offload pass here: + # apply_memory_plan's group tier (_apply_group_offload) already block-streams every + # DiT it finds on the pipe -- transformer AND transformer_2 -- and model/sequential + # offload hook every top-level module, so the second expert is covered under all tiers. + # A second _apply_group_offload on transformer_2 would re-register the group-offload + # hooks it already carries, which diffusers rejects with a duplicate-hook ValueError. + if not vae_tiling: + # Decode of a whole clip is the video memory peak; tiling is near-free + # in quality and keeps the decode bounded, so it is always on. + try: + pipe.vae.enable_tiling() + vae_tiling = True + except Exception as exc: # noqa: BLE001 -- tiling is an optimisation only + logger.warning("video.vae_tiling_failed: %s", exc) + + resolved = build_resolved_record( + { + "memory_mode": ( + memory_mode, + plan.requested_mode, + f"planned '{plan.offload_policy}' offload from the family size table", + ), + "speed_mode": ( + speed_mode, + effective_speed, + "quantized transformer requires compile" + if transformer_quant_engaged is not None + else "clip denoises amortise the one-time compile within a single run" + if speed_mode is None + else "requested", + ), + "attention_backend": ( + attention_backend, + attention_engaged or "native", + "cuDNN fused attention on NVIDIA when a speed profile is active", + ), + "transformer_cache": ( + None if cache_auto else transformer_cache, + cache_engaged or "off", + cache_reason, + ), + "transformer_quant": ( + transformer_quant, + transformer_quant_engaged or "off", + "dense DiT(s) torchao-quantised onto the low-precision tensor cores" + if transformer_quant_engaged is not None + else ( + "skipped: offload moves the DiT, unsupported for torchao " + "tensors; pin a resident memory mode to combine them" + if quant_skipped_for_offload + else "not engaged (dense bf16 DiT loaded)" + ), + ), + "text_encoder_quant": ( + text_encoder_quant, + text_encoder_quant_engaged or "off", + "dense text encoder quantised in place" + if text_encoder_quant_engaged is not None + else "not engaged (dense bf16 text encoder loaded)", + ), + } ) - # Ownership of the globals transferred to _state / _teardown_state. - self._precommit_globals = None + + with self._lock: + if _load_token is not None and _load_token != self._load_token: + del pipe + clear_gpu_cache() + raise RuntimeError("Video load was cancelled or superseded.") + self._state = _VideoLoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + kind = kind, + gguf_filename = gguf_filename, + offload_policy = offload_policy, + vae_tiling = vae_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + # Already filtered above to only the optimisations that engaged; + # apply_speed_optims returns every flag True/False and the view + # loop keeps just the True names. + speed_optims = speed_optims, + backend_flags = backend_flags, + attention_backend = attention_engaged, + transformer_cache = cache_engaged, + cache_auto = cache_may_toggle, + cache_quant_active = cache_quant_active, + cache_threshold = transformer_cache_threshold, + transformer_quant = transformer_quant_engaged, + text_encoder_quant = text_encoder_quant_engaged, + resolved = resolved, + ) + # Ownership of the globals transferred to _state / _teardown_state. + self._precommit_globals = None logger.info( "video.loaded: %s (%s, %s, offload=%s, speed=%s, quant=%s)", repo_id, diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index b6e8ee6570..4973f7bd94 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1272,6 +1272,21 @@ def test_resolve_local_single_file(tmp_path): # A remote repo id (not a local dir) -> None. assert resolve_local_single_file("unsloth/Qwen-Image-2512-GGUF") is None + # A PEFT LoRA adapter folder (adapter_config.json + adapter_model.safetensors), even with a + # family-token name, is NOT a base checkpoint: from_single_file would fail on the adapter + # weights AFTER the route evicted the resident GPU model, so it must not be reinterpreted as a + # single_file pick -> None (the pipeline pick then 400s in validation, before the handoff). + adapter = tmp_path / "flux-style-lora" + adapter.mkdir() + (adapter / "adapter_config.json").write_text("{}") + (adapter / "adapter_model.safetensors").write_bytes(b"w") + assert resolve_local_single_file(str(adapter)) is None + # A bare adapter_model.safetensors (no config) is likewise not treated as the sole checkpoint. + adapter2 = tmp_path / "z-image-lora" + adapter2.mkdir() + (adapter2 / "adapter_model.safetensors").write_bytes(b"w") + assert resolve_local_single_file(str(adapter2)) is None + def test_resolve_base_repo_drops_untrusted_card_tag(monkeypatch): # When no base_repo is passed, the base is resolved from the GGUF repo's base_model card diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 14d074d45c..1a25813c5d 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -699,6 +699,62 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): assert status["loaded"] is False +def test_load_holds_generate_lock_across_placement(fake_runtime, tmp_path, monkeypatch): + # The video load must hold _generate_lock across GPU placement (apply_memory_plan) so an + # unload / arbiter eviction -- which barriers on _generate_lock before freeing -- cannot hand + # the GPU to another backend while a multi-GB pipeline is still being moved onto it (mirrors + # the image backend, which places + commits under this lock). Verify unload() blocks until + # placement releases the lock, and the superseded load then aborts without committing. + import threading + + from core.inference import video as video_mod + + backend = VideoBackend() + placement_started = threading.Event() + release_placement = threading.Event() + real_apply = video_mod.apply_memory_plan + + def blocking_apply(pipe, plan, **kw): + placement_started.set() + assert release_placement.wait(timeout = 5), "test placement barrier never released" + return real_apply(pipe, plan, **kw) + + monkeypatch.setattr(video_mod, "apply_memory_plan", blocking_apply) + + load_exc = [] + + def do_load(): + try: + _load_gguf(backend, tmp_path) + except Exception as e: # noqa: BLE001 -- the concurrent unload supersedes this load + load_exc.append(e) + + load_thread = threading.Thread(target = do_load) + load_thread.start() + assert placement_started.wait(timeout = 5), "load never reached placement" + + # Placement is in flight, holding _generate_lock. unload() must block on its barrier. + unload_done = [] + + def do_unload(): + backend.unload() + unload_done.append(True) + + unload_thread = threading.Thread(target = do_unload) + unload_thread.start() + unload_thread.join(timeout = 0.5) + assert not unload_done, "unload() returned while placement still held _generate_lock (the race)" + + # Release placement; unload()'s barrier then passes and its teardown runs strictly AFTER + # the load's placement+commit -- never concurrently -- so no two pipelines are ever resident. + release_placement.set() + unload_thread.join(timeout = 5) + load_thread.join(timeout = 5) + assert unload_done, "unload() did not complete after placement released _generate_lock" + assert not load_thread.is_alive() and not load_exc + assert backend._state is None # unload's teardown ran after the load, leaving nothing resident + + def test_load_records_engaged_speed_optims(fake_runtime, tmp_path, monkeypatch): # Regression: the load tail once re-ran the already-filtered speed_optims # tuple through ``.items()`` as if it were still the raw applied dict, so