diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index ed66e554f9..d48f8269f5 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -482,55 +482,67 @@ class DiffusionBackend: # first so unload can restore them: TF32 / cudnn.benchmark are global, # and a later `off` load must not inherit this load's settings. backend_flags_before = snapshot_backend_flags() - speed_applied = apply_speed_optims( - pipe, - target, - is_gguf = bool(gguf_filename), - family = fam, - speed_mode = effective_speed, - logger = logger, - ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. - te_quant = quantize_text_encoders( - pipe, - target, - mode = text_encoder_quant, - logger = logger, - ) + # apply_speed_optims mutates PROCESS-WIDE flags (TF32 / cudnn.benchmark); + # they are only restored via _LoadState.backend_flags_before on unload. If + # the build fails after this but before _state commits (e.g. an OOM in + # apply_memory_plan / pipe.to), nothing would restore them and a later `off` + # generation would be contaminated, so restore on any non-committed exit. + committed = False + try: + speed_applied = apply_speed_optims( + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = effective_speed, + logger = logger, + ) + # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), + # also before placement so the offload hooks move the smaller weights. + te_quant = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + logger = logger, + ) - # Decide placement from MEASURED free device memory vs the model's - # estimated resident size (transformer GGUF dequantised + the - # companion text-encoder / VAE already cached for `base`), then - # apply it. Computed here, after the build but before placement, - # because the weights are still on CPU so free VRAM is the real - # budget. `cpu_offload=True` stays an explicit override. - plan = self._plan_memory( - target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload - ) - # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it - # may fall back to whole-module offload, and tiling is a no-op on a - # pipeline with no tiling control), so status stays honest. - effective_policy, effective_tiling = apply_memory_plan( - pipe, plan, device = device, logger = logger - ) + # Decide placement from MEASURED free device memory vs the model's + # estimated resident size (transformer GGUF dequantised + the + # companion text-encoder / VAE already cached for `base`), then + # apply it. Computed here, after the build but before placement, + # because the weights are still on CPU so free VRAM is the real + # budget. `cpu_offload=True` stays an explicit override. + plan = self._plan_memory( + target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload + ) + # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it + # may fall back to whole-module offload, and tiling is a no-op on a + # pipeline with no tiling control), so status stays honest. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) - self._state = _LoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - cpu_offload = effective_policy != OFFLOAD_NONE, - offload_policy = effective_policy, - vae_tiling = effective_tiling, - memory_mode = plan.requested_mode, - speed_mode = effective_speed, - speed_optims = tuple(k for k, v in speed_applied.items() if v), - backend_flags_before = backend_flags_before, - text_encoder_quant = te_quant, - ) + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = effective_policy != OFFLOAD_NONE, + offload_policy = effective_policy, + vae_tiling = effective_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(k for k, v in speed_applied.items() if v), + backend_flags_before = backend_flags_before, + text_encoder_quant = te_quant, + ) + committed = True + finally: + if not committed: + restore_backend_flags(backend_flags_before) + clear_gpu_cache() logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 2548c0f0dc..9a88b5f317 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -42,31 +42,49 @@ SPEED_MODES = (SPEED_OFF, SPEED_DEFAULT, SPEED_MAX) def snapshot_backend_flags() -> Optional[dict]: """Capture the process-wide torch backend flags this layer may mutate, so the - caller can restore them on unload. None if torch is unavailable.""" + caller can restore them on unload. None if torch is unavailable. Each flag is read + defensively so a build/platform missing one (e.g. no cuda.matmul on CPU/MPS) still + captures the rest -- otherwise a single missing attribute would skip the whole + snapshot and a real mutated flag would leak.""" try: import torch - return { - "matmul_tf32": bool(torch.backends.cuda.matmul.allow_tf32), - "cudnn_tf32": bool(torch.backends.cudnn.allow_tf32), - "cudnn_benchmark": bool(torch.backends.cudnn.benchmark), - } - except Exception: # noqa: BLE001 — best-effort; no snapshot -> no restore + except Exception: # noqa: BLE001 — no torch -> nothing to snapshot/restore return None + state: dict[str, bool] = {} + matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None) + if matmul is not None and hasattr(matmul, "allow_tf32"): + state["matmul_tf32"] = bool(matmul.allow_tf32) + cudnn = getattr(torch.backends, "cudnn", None) + if cudnn is not None: + if hasattr(cudnn, "allow_tf32"): + state["cudnn_tf32"] = bool(cudnn.allow_tf32) + if hasattr(cudnn, "benchmark"): + state["cudnn_benchmark"] = bool(cudnn.benchmark) + return state def restore_backend_flags(state: Optional[dict]) -> None: - """Restore the flags captured by ``snapshot_backend_flags``. No-op on None.""" + """Restore the flags captured by ``snapshot_backend_flags``. No-op on None. Each + flag is restored independently so one failure can't leave the others leaked.""" if not state: return try: import torch - - torch.backends.cuda.matmul.allow_tf32 = state["matmul_tf32"] - torch.backends.cudnn.allow_tf32 = state["cudnn_tf32"] - torch.backends.cudnn.benchmark = state["cudnn_benchmark"] - except Exception: # noqa: BLE001 — best-effort restore + except Exception: # noqa: BLE001 — no torch -> nothing to restore return + def _set(obj: Any, attr: str, key: str) -> None: + if obj is not None and key in state and hasattr(obj, attr): + try: + setattr(obj, attr, state[key]) + except Exception: # noqa: BLE001 — best-effort per-flag restore + pass + + _set(getattr(getattr(torch.backends, "cuda", None), "matmul", None), "allow_tf32", "matmul_tf32") + cudnn = getattr(torch.backends, "cudnn", None) + _set(cudnn, "allow_tf32", "cudnn_tf32") + _set(cudnn, "benchmark", "cudnn_benchmark") + def normalize_speed_mode(value: Optional[str]) -> str: """Lower/strip a requested speed mode (dashes ok); None / "" -> off.""" diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index fb7c0fdb79..b0ca10af4c 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -116,6 +116,42 @@ def test_restore_backend_flags_tolerates_none(): restore_backend_flags(None) # no torch needed, no-op +def test_snapshot_partial_when_some_backends_missing(monkeypatch): + # A build/platform without cuda.matmul (e.g. CPU/MPS) must still snapshot + restore the + # flags it does have, rather than skipping the whole snapshot on one missing attribute. + torch = types.ModuleType("torch") + torch.backends = types.SimpleNamespace( + cuda = types.SimpleNamespace(), # no .matmul + cudnn = types.SimpleNamespace(benchmark = True), # no .allow_tf32 + ) + monkeypatch.setitem(sys.modules, "torch", torch) + snap = snapshot_backend_flags() + assert snap == {"cudnn_benchmark": True} + torch.backends.cudnn.benchmark = False + restore_backend_flags(snap) + assert torch.backends.cudnn.benchmark is True + + +def test_restore_is_independent_per_flag(monkeypatch): + # A read-only / failing attribute must not abort restoring the remaining flags. + torch = _stub_torch(monkeypatch) + + class _NoMatmulSet: + @property + def allow_tf32(self): + return False + + @allow_tf32.setter + def allow_tf32(self, value): + raise RuntimeError("read-only on this build") + + torch.backends.cuda.matmul = _NoMatmulSet() + snap = {"matmul_tf32": False, "cudnn_tf32": False, "cudnn_benchmark": False} + torch.backends.cudnn.benchmark = True + restore_backend_flags(snap) # matmul setter raises, cudnn still restored + assert torch.backends.cudnn.benchmark is False + + # ── applier ───────────────────────────────────────────────────────────────────