diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 2c4b2c359b..7377f8d919 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -216,6 +216,12 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non timeout = 600, check = True, ) + # The wheel just landed in site-packages, but the import system caches each + # directory's listing; the very next find_spec / import in this same process can + # still miss the freshly installed package when the install lands within the + # directory mtime's resolution -- silently falling back to native on the first + # use. Invalidate the finder caches so set_attention_backend picks it up now. + importlib.invalidate_caches() except Exception as exc: # noqa: BLE001 — no wheel / no network -> native fallback if logger is not None: # A failed pip install raises CalledProcessError whose str() shows only the diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 89e496eca7..5abaa86fa7 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -290,6 +290,41 @@ def test_install_runs_wheel_only_for_missing_kernel(monkeypatch): assert "--only-binary" in cmd and ":all:" in cmd and "sageattention" in cmd +def test_install_invalidates_import_caches_on_success(monkeypatch): + # A wheel written to site-packages after the finder cached that directory can be + # missed by the very next import, so a successful install must invalidate the caches + # (otherwise set_attention_backend imports the missing package and falls back). + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + import importlib + import importlib.util + + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + _stub_subprocess(monkeypatch, _Recorder()) + invalidated = [] + monkeypatch.setattr(importlib, "invalidate_caches", lambda: invalidated.append(True)) + att._ensure_attention_backend_installed("sage") + assert invalidated == [True] + + +def test_install_failure_skips_cache_invalidation(monkeypatch): + # A failed install left nothing to import, so the finder caches must be left alone. + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + import importlib + 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) + + _stub_subprocess(monkeypatch, _boom) + invalidated = [] + monkeypatch.setattr(importlib, "invalidate_caches", lambda: invalidated.append(True)) + att._ensure_attention_backend_installed("sage") + assert invalidated == [] + + def test_install_never_attempted_for_builtin_backends(monkeypatch): monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") run = _Recorder()