diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 9555b3c38a..bcda3ba829 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -68,6 +68,7 @@ from .diffusion_speed import ( from .diffusion_attention import ( apply_attention_backend, select_attention_backend, + _ensure_attention_backend_installed, ) from . import diffusion_compile_cache as compile_cache from . import diffusion_gguf_compile as gguf_compile @@ -940,6 +941,23 @@ class DiffusionBackend: import diffusers + # Pre-install the optional attention kernel BEFORE taking the load locks. The + # wheel-only pip install can run up to 600s, and doing it under _lock / + # _generate_lock (as the in-lock apply_attention_backend otherwise would) blocks + # unload() and cancellation for that whole window. Only an explicit backend pulls + # a package -- auto resolves to cuDNN / native, which ship with torch -- and an + # explicit backend's resolution ignores the speed tier, so it can run here without + # effective_speed. Best-effort: the authoritative resolve + set still happens under + # the lock, where the now-satisfied install call is a fast no-op. + try: + preinstall_backend = select_attention_backend( + target, attention_backend, speed_active = True + ) + if preinstall_backend is not None: + _ensure_attention_backend_installed(preinstall_backend, logger) + except Exception: # noqa: BLE001 — the locked path re-resolves and validates + pass + # Signal an in-flight denoise to abort, then take _generate_lock to WAIT for # it to actually exit before allocating the replacement: a load is about to # claim VRAM, so unlike unload() it must not overlap a still-live pipeline. diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 912b20e73e..c04c569fa0 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -213,11 +213,24 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non ) except Exception as exc: # noqa: BLE001 — no wheel / no network -> native fallback if logger is not None: - logger.warning( - "diffusion.attention: could not install %s (%s); falling back to default", - package, - exc, - ) + # A failed pip install raises CalledProcessError whose str() shows only the + # exit code and command; the real reason (no matching wheel, resolver error) + # is in exc.stderr. Surface it so a fallback to native is diagnosable. + stderr = getattr(exc, "stderr", None) + if stderr: + if isinstance(stderr, bytes): + stderr = stderr.decode("utf-8", errors = "replace") + logger.warning( + "diffusion.attention: could not install %s; pip failed with: %s", + package, + stderr.strip() or str(exc), + ) + else: + logger.warning( + "diffusion.attention: could not install %s (%s); falling back to default", + package, + exc, + ) def apply_attention_backend( diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index dce7dd2bfa..dc16cc634f 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -290,6 +290,35 @@ def test_install_never_attempted_for_builtin_backends(monkeypatch): assert run.calls == [] +def test_install_failure_logs_pip_stderr(monkeypatch): + # A CalledProcessError's str() hides the pip reason; the warning must surface the + # captured stderr (decoding bytes) so a fallback to native is diagnosable. + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + import importlib.util + import subprocess as sp + + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + + def _boom(cmd, **kwargs): + raise sp.CalledProcessError( + returncode = 1, cmd = cmd, stderr = b"ERROR: No matching distribution found" + ) + + _stub_subprocess(monkeypatch, _boom) + + warnings: list[str] = [] + + class _Logger: + def info(self, *a, **k): + pass + + def warning(self, msg, *args): + warnings.append(msg % args if args else msg) + + att._ensure_attention_backend_installed("sage", _Logger()) + assert warnings and "No matching distribution found" in warnings[-1] + + def test_install_failure_falls_back_to_native(monkeypatch): # pip failing (no wheel for this platform) must not break the load: the apply # path proceeds, set_attention_backend raises on the missing package, and the