Install attention backend outside the load locks and surface pip errors
The wheel-only pip install for an optional attention kernel ran inside load_pipeline under _lock and _generate_lock, so a slow or hanging install blocked unload and cancellation for up to the 600s timeout. Resolve and install the kernel before taking the locks (only an explicit backend ever pulls a package, and its resolution ignores the speed tier); the in-lock apply call is then a fast no-op. Also decode and log pip's stderr on a failed install so the fallback to native is diagnosable instead of showing only the exit code.
This commit is contained in:
parent
935eed0cc7
commit
ba3d1f607b
3 changed files with 65 additions and 5 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue