"""Tests for the optional FlashAttention installer.""" from __future__ import annotations import os import subprocess import sys from pathlib import Path from unittest import mock def _shared_setup_1(fake_run, step_messages): with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "13", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ), mock.patch.object(ips, "url_exists", return_value = True), mock.patch.object( ips, "install_wheel", return_value = [("uv", subprocess.CompletedProcess(["uv"], 0, ""))], ), mock.patch.object( ips, "_step", side_effect = lambda label, value, color_fn = None: step_messages.append((label, value)), ), mock.patch("subprocess.run", side_effect = fake_run), ): ips._ensure_flash_attn() STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio" sys.path.insert(0, str(STUDIO_DIR)) sys.path.insert(0, str(STUDIO_DIR / "backend")) import install_python_stack as ips import pytest from utils import wheel_utils @pytest.fixture(autouse = True) def _an_sm80_gpu(monkeypatch): # The host's real GPU must not decide these tests; the capability cases below override it. monkeypatch.setattr(ips, "_flash_attn_gpu_below_floor", lambda: False) class TestPrebuiltWheelTorchMapping: def test_torch_211_maps_to_torch210(self): assert wheel_utils.prebuilt_wheel_torch_mm("2.11") == "2.10" def test_torch_212_maps_to_torch210(self): assert wheel_utils.prebuilt_wheel_torch_mm("2.12") == "2.10" def test_other_versions_pass_through(self): # 2.13 stays unmapped until its wheels have been measured. for torch_mm in ("2.9", "2.10", "2.13"): assert wheel_utils.prebuilt_wheel_torch_mm(torch_mm) == torch_mm def test_reuse_never_targets_a_pre_210_wheel(self): # torch broke extension ABI between 2.9 and 2.10. assert set(wheel_utils._PREBUILT_WHEEL_TORCH_MM.values()) == {"2.10"} def test_direct_wheel_url_reuses_torch210_on_211(self): url = wheel_utils.direct_wheel_url( filename_prefix = "causal_conv1d", package_version = "1.6.1", release_tag = "v1.6.1.post4", release_base_url = "https://example.test/download", env = { "python_tag": "cp313", "torch_mm": "2.11", # cu12: cu13 cp311-cp313 resolves to our own build instead. "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ) assert url is not None assert "causal_conv1d-1.6.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url def test_direct_wheel_url_reuses_torch210_on_212(self): url = wheel_utils.direct_wheel_url( filename_prefix = "mamba_ssm", package_version = "2.3.1", release_tag = "v2.3.1", release_base_url = "https://example.test/download", env = { "python_tag": "cp312", "torch_mm": "2.12", # cu12: cu13 cp311-cp313 resolves to our own build instead. "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ) assert url is not None assert "mamba_ssm-2.3.1+cu12torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" in url class TestFlashAttnWheelSelection: def test_torch_210_maps_to_v281(self): assert ips._select_flash_attn_version("2.10") == "2.8.1" def test_selected_version_is_never_a_post_release(self): # v2.8.3.post1 dropped torch2.10 assets; a naive newest-release bump must fail here. for torch_mm in ("2.4", "2.7", "2.9", "2.10"): version = ips._select_flash_attn_version(torch_mm) assert version is not None assert ".post" not in version def test_torch_29_maps_to_v283(self): assert ips._select_flash_attn_version("2.9") == "2.8.3" def test_torch_211_has_no_native_version_entry(self): assert ips._select_flash_attn_version("2.11") is None def test_torch_211_reuses_torch210_wheel(self): url = ips._build_flash_attn_wheel_url( { "python_tag": "cp313", "torch_mm": "2.11", "cuda_major": "13", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", } ) assert url is not None assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url def test_torch_212_reuses_torch210_wheel(self): url = ips._build_flash_attn_wheel_url( { "python_tag": "cp313", "torch_mm": "2.12", "cuda_major": "13", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", } ) assert url is not None assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url def test_exact_wheel_url_uses_full_env_tuple(self): url = ips._build_flash_attn_wheel_url( { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", } ) assert url is not None assert "v2.8.1" in url assert "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url def test_missing_cuda_major_disables_wheel_lookup(self): assert ( ips._build_flash_attn_wheel_url( { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", } ) is None ) class TestFlashAttnImportProbe: """The probe is bounded: a native extension can hang in its initialiser, not just fail.""" def test_clean_exit_is_importable(self): with mock.patch( "subprocess.run", return_value = subprocess.CompletedProcess(["python"], 0), ): assert ips._flash_attn_importable() is True def test_non_zero_exit_is_not_importable(self): with mock.patch( "subprocess.run", return_value = subprocess.CompletedProcess(["python"], 1), ): assert ips._flash_attn_importable() is False def test_a_hung_import_is_not_importable(self): with mock.patch( "subprocess.run", side_effect = subprocess.TimeoutExpired(cmd = "python", timeout = 300), ): assert ips._flash_attn_importable() is False def test_the_probe_is_bounded(self): with mock.patch( "subprocess.run", return_value = subprocess.CompletedProcess(["python"], 0), ) as run: ips._flash_attn_importable() assert run.call_args.kwargs["timeout"] == ips._FLASH_ATTN_IMPORT_PROBE_TIMEOUT class TestEnsureFlashAttn: def _import_check(self, code: int = 1): return subprocess.CompletedProcess(["python", "-c", "import flash_attn"], code) def _import_fails_removal_works(self, cmd, **_kwargs): """flash_attn never imports, but uninstalling it succeeds.""" if "uninstall" in cmd: return subprocess.CompletedProcess(cmd, 0) return self._import_check() def test_prefers_exact_match_wheel(self): install_calls = [] def fake_install_wheel(*args, **kwargs): install_calls.append((args, kwargs)) return [("uv", subprocess.CompletedProcess(["uv"], 0, ""))] with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object(ips, "USE_UV", True), mock.patch.object(ips, "UV_NEEDS_SYSTEM", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ), mock.patch.object(ips, "url_exists", return_value = True), mock.patch.object(ips, "install_wheel", side_effect = fake_install_wheel), mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() assert len(install_calls) == 1 args, kwargs = install_calls[0] assert args == ( "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl", ) assert kwargs["python_executable"] == sys.executable assert kwargs["use_uv"] is True assert kwargs["uv_needs_system"] is False def test_uv_install_respects_system_flag(self): install_calls = [] def fake_install_wheel(*args, **kwargs): install_calls.append((args, kwargs)) return [("uv", subprocess.CompletedProcess(["uv"], 0, ""))] with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object(ips, "USE_UV", True), mock.patch.object(ips, "UV_NEEDS_SYSTEM", True), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ), mock.patch.object(ips, "url_exists", return_value = True), mock.patch.object(ips, "install_wheel", side_effect = fake_install_wheel), mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() assert len(install_calls) == 1 _, kwargs = install_calls[0] assert kwargs["uv_needs_system"] is True def test_wheel_that_does_not_import_is_not_trusted(self): """pip exits 0 on a wrong-arch or ABI wheel, so only a successful import shows flash-attn is ready.""" step_messages: list[tuple[str, str]] = [] with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "13", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ), mock.patch.object(ips, "url_exists", return_value = True), mock.patch.object( ips, "install_wheel", return_value = [("uv", subprocess.CompletedProcess(["uv"], 0, ""))], ), mock.patch.object( ips, "_step", side_effect = lambda label, value, color_fn = None: step_messages.append( (label, value) ), ), mock.patch("subprocess.run", side_effect = self._import_fails_removal_works), ): ips._ensure_flash_attn() assert ( "warning", "flash-attn wheel installed but is not importable on this GPU; removed it", ) in step_messages assert ("warning", "Continuing without flash-attn") in step_messages def test_rejected_wheel_is_uninstalled(self): """Rejected wheels must be uninstalled: _package_available trusts metadata and imports natively.""" step_messages: list[tuple[str, str]] = [] removals: list[list[str]] = [] def fake_run(cmd, **kwargs): if "uninstall" in cmd: removals.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0) return self._import_check() _shared_setup_1(fake_run, step_messages) assert removals, "the rejected wheel must be uninstalled, not left in site-packages" assert any("flash-attn" in cmd for cmd in removals), removals assert ("warning", "Continuing without flash-attn") in step_messages def test_uninstall_targets_the_interpreter_install_wheel_used(self): """install_wheel passes --python as well as --system, and its pip fallback runs sys.executable directly, so --system alone would uninstall from the wrong Python.""" commands: list[list[str]] = [] def fake_run(cmd, **kwargs): commands.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0) with ( mock.patch.object(ips, "USE_UV", True), mock.patch.object(ips, "UV_NEEDS_SYSTEM", True), mock.patch.object(ips.shutil, "which", return_value = "/usr/bin/uv"), mock.patch("subprocess.run", side_effect = fake_run), ): assert ips._remove_rejected_flash_attn() is True assert commands == [ ["uv", "pip", "uninstall", "--system", "--python", sys.executable, "flash-attn"] ] def test_uninstall_targets_the_interpreter_without_system_mode(self): commands: list[list[str]] = [] def fake_run(cmd, **kwargs): commands.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0) with ( mock.patch.object(ips, "USE_UV", True), mock.patch.object(ips, "UV_NEEDS_SYSTEM", False), mock.patch.object(ips.shutil, "which", return_value = "/usr/bin/uv"), mock.patch("subprocess.run", side_effect = fake_run), ): assert ips._remove_rejected_flash_attn() is True assert commands == [["uv", "pip", "uninstall", "--python", sys.executable, "flash-attn"]] def test_a_failed_removal_is_not_reported_as_removed(self): """Still importable in process, so it must not read like a clean skip.""" step_messages: list[tuple[str, str]] = [] def fake_run(cmd, **kwargs): if "uninstall" in cmd: return subprocess.CompletedProcess(cmd, 1) return self._import_check() _shared_setup_1(fake_run, step_messages) warnings = [value for _, value in step_messages] assert any("could not be removed" in value for value in warnings), warnings assert not any("removed it" in value for value in warnings), warnings def test_working_wheel_reports_no_warning(self): """The happy path stays silent: install exits 0 and the module imports.""" step_messages: list[tuple[str, str]] = [] import_calls: list[int] = [] def fake_run(cmd, **kwargs): import_calls.append(1) return self._import_check(1 if len(import_calls) == 1 else 0) _shared_setup_1(fake_run, step_messages) assert step_messages == [] assert len(import_calls) == 2, "expected a verification import after the install" def test_wheel_failure_warns_and_continues(self): step_messages: list[tuple[str, str]] = [] printed_failures: list[str] = [] def fake_step( label: str, value: str, color_fn = None, ): step_messages.append((label, value)) with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object(ips, "USE_UV", True), mock.patch.object(ips, "UV_NEEDS_SYSTEM", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ), mock.patch.object(ips, "url_exists", return_value = True), mock.patch.object( ips, "install_wheel", return_value = [ ("uv", subprocess.CompletedProcess(["uv"], 1, "uv wheel failed")), ( "pip", subprocess.CompletedProcess(["pip"], 1, "pip wheel failed"), ), ], ), mock.patch.object( ips, "_print_optional_install_failure", side_effect = lambda label, result: printed_failures.append(label), ), mock.patch.object(ips, "_step", side_effect = fake_step), mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() assert printed_failures == [ "Installing flash-attn prebuilt wheel with uv", "Installing flash-attn prebuilt wheel with pip", ] assert ("warning", "Continuing without flash-attn") in step_messages def test_wheel_missing_skips_install_at_setup_time(self): step_messages: list[tuple[str, str]] = [] def fake_step( label: str, value: str, color_fn = None, ): step_messages.append((label, value)) with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "13", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ), mock.patch.object(ips, "url_exists", return_value = False), mock.patch.object(ips, "install_wheel") as mock_install_wheel, mock.patch.object(ips, "_step", side_effect = fake_step), mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() mock_install_wheel.assert_not_called() assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages def test_an_unchecked_wheel_is_not_reported_as_unpublished(self): step_messages = [] def fake_step(kind, message, *args, **kwargs): step_messages.append((kind, message)) with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "torch_mm": "2.10", "cuda_major": "13", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ), mock.patch.object(ips, "url_exists", return_value = None), mock.patch.object(ips, "install_wheel") as mock_install_wheel, mock.patch.object(ips, "_step", side_effect = fake_step), mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() mock_install_wheel.assert_not_called() assert ("warning", "No published flash-attn prebuilt wheel found") not in step_messages assert any("could not" in message.lower() for _kind, message in step_messages) def _run_with_resident_copy( self, *, resident: str | None = "2.8.1", env: dict[str, str] | None = None, torch_imports: bool = True, url_available: bool | None = False, install_returncode: int = 0, fake_run = None, ) -> tuple[list[tuple[str, str]], list[list[str]], mock.Mock]: """#13244 leaves cu13torch2.10 under torch 2.9+cu126 on cp313, unsupported by v2.8.3.""" step_messages: list[tuple[str, str]] = [] removals: list[list[str]] = [] def default_run(cmd, **kwargs): if "uninstall" in cmd: removals.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0) return self._import_check() if env is None: env = { "python_tag": "cp313", "torch_mm": "2.9", "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", } with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = env if torch_imports else None ), mock.patch.object(ips, "_installed_distribution_version", return_value = resident), mock.patch.object(ips, "url_exists", return_value = url_available), mock.patch.object( ips, "install_wheel", return_value = [("uv", subprocess.CompletedProcess(["uv"], install_returncode, ""))], ) as mock_install_wheel, mock.patch.object( ips, "_step", side_effect = lambda label, value, color_fn = None: step_messages.append( (label, value) ), ), mock.patch("subprocess.run", side_effect = fake_run or default_run), ): ips._ensure_flash_attn() return step_messages, removals, mock_install_wheel def test_stale_copy_is_removed_when_no_wheel_is_published(self): """if kept, xformers imports it and every diffusers pipeline fails (#13244).""" step_messages, removals, mock_install_wheel = self._run_with_resident_copy() mock_install_wheel.assert_not_called() assert any("flash-attn" in cmd for cmd in removals), removals assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages assert ( "warning", "flash-attn was built for another torch or CUDA and does not import; removed it", ) in step_messages def test_stale_copy_is_removed_when_no_wheel_matches(self): step_messages, removals, _ = self._run_with_resident_copy( env = { "python_tag": "cp313", "torch_mm": "2.3", "cuda_major": "12", "cxx11abi": "TRUE", "platform_tag": "linux_x86_64", }, ) assert any("flash-attn" in cmd for cmd in removals), removals assert ("warning", "No compatible flash-attn prebuilt wheel found") in step_messages def test_stale_copy_is_removed_when_the_replacement_fails_to_install(self): step_messages, removals, _ = self._run_with_resident_copy( url_available = True, install_returncode = 1 ) assert any("flash-attn" in cmd for cmd in removals), removals assert ("warning", "Continuing without flash-attn") in step_messages def test_stale_copy_is_force_replaced(self): """a same-version build for another CUDA counts as installed unless forced.""" imports: list[int] = [] def fake_run(cmd, **kwargs): imports.append(1) return self._import_check(1 if len(imports) == 1 else 0) step_messages, _, mock_install_wheel = self._run_with_resident_copy( url_available = True, fake_run = fake_run ) assert mock_install_wheel.call_args.kwargs["reinstall"] is True assert step_messages == [] def test_a_failed_stale_removal_is_reported(self): def fake_run(cmd, **kwargs): if "uninstall" in cmd: return subprocess.CompletedProcess(cmd, 1) return self._import_check() step_messages, _, _ = self._run_with_resident_copy(fake_run = fake_run) warnings = [value for _, value in step_messages] assert any("could not be removed" in value for value in warnings), warnings assert not any("removed it" in value for value in warnings), warnings def test_resident_copy_is_kept_when_torch_does_not_import(self): """the import probe loads torch first, so a broken torch says nothing about flash-attn.""" _, removals, mock_install_wheel = self._run_with_resident_copy(torch_imports = False) assert removals == [] mock_install_wheel.assert_not_called() def test_nothing_is_uninstalled_when_flash_attn_is_absent(self): step_messages, removals, _ = self._run_with_resident_copy(resident = None) assert removals == [] assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages def test_skip_env_disables_setup_install(self): with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", False), mock.patch.object(ips, "IS_MACOS", False), mock.patch.dict(os.environ, {"UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL": "1"}), mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, mock.patch.object(ips, "install_wheel") as mock_install_wheel, mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() mock_probe.assert_not_called() mock_install_wheel.assert_not_called() def test_windows_probes_with_windows_opt_in_and_installs_the_wheel(self): url = "https://example.invalid/flash_attn-2.8.4+cu13torch2.13cxx11abiTRUE-cp313-cp313-win_amd64.whl" with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", True), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object(ips, "USE_UV", True), mock.patch.object(ips, "UV_NEEDS_SYSTEM", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = {"platform_tag": "win_amd64"} ) as mock_probe, mock.patch.object(ips, "_build_flash_attn_wheel_url", return_value = url), mock.patch.object(ips, "url_exists", return_value = True), mock.patch.object( ips, "install_wheel", return_value = [("uv", subprocess.CompletedProcess(["uv"], 0, ""))], ) as mock_install_wheel, mock.patch.object(ips, "_flash_attn_importable", side_effect = [False, True]), ): ips._ensure_flash_attn() mock_probe.assert_called_once_with( timeout = ips._FLASH_ATTN_IMPORT_PROBE_TIMEOUT, include_windows = True ) assert mock_install_wheel.call_args.args[0] == url def test_windows_without_a_wheel_notes_it_quietly(self): notes: list[str] = [] steps: list[tuple[str, str]] = [] with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", True), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object( ips, "probe_torch_wheel_env", return_value = {"platform_tag": "win_amd64"} ), mock.patch.object(ips, "_build_flash_attn_wheel_url", return_value = None), mock.patch.object(ips, "install_wheel") as mock_install_wheel, mock.patch.object(ips, "_flash_attn_importable", return_value = False), mock.patch.object(ips, "_note", side_effect = lambda message, *a: notes.append(message)), mock.patch.object( ips, "_step", side_effect = lambda kind, message, *a: steps.append((kind, message)) ), ): ips._ensure_flash_attn() mock_install_wheel.assert_not_called() assert steps == [] assert any("flash-attn" in note for note in notes) class TestInstallPythonStackFlashAttnIntegration: def _run_install( self, *, no_torch: bool, is_macos: bool, is_windows: bool, flavor_moves_torch: bool = False, ) -> int: flash_attn_calls = 0 torch_version = ["2.13.0+cpu"] def fix_flavor(tag): if flavor_moves_torch: torch_version[0] = "2.13.0+cu130" return True def fake_run(cmd, **kw): return subprocess.CompletedProcess(cmd, 0, b"", b"") def count_flash_attn(): nonlocal flash_attn_calls flash_attn_calls += 1 with ( mock.patch.object(ips, "NO_TORCH", no_torch), mock.patch.object(ips, "IS_MACOS", is_macos), mock.patch.object(ips, "IS_WINDOWS", is_windows), mock.patch.object(ips, "USE_UV", True), mock.patch.object(ips, "UV_NEEDS_SYSTEM", False), mock.patch.object(ips, "VERBOSE", False), mock.patch.object(ips, "_bootstrap_uv", return_value = True), mock.patch.object(ips, "_ensure_flash_attn", side_effect = count_flash_attn), mock.patch("subprocess.run", side_effect = fake_run), mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False), mock.patch.object(ips, "_has_rocm_gpu", return_value = False), mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")), mock.patch("pathlib.Path.is_dir", return_value = True), mock.patch("pathlib.Path.is_file", return_value = True), mock.patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}, clear = False), mock.patch.object(ips, "_ensure_expected_torch_flavor", side_effect = fix_flavor), mock.patch.object( ips, "_probe_installed_torch_version", side_effect = lambda: torch_version[0] ), ): ips.install_python_stack() return flash_attn_calls def test_linux_torch_install_calls_flash_attn_step(self): assert self._run_install(no_torch = False, is_macos = False, is_windows = False) == 1 def test_no_torch_install_skips_flash_attn_step(self): assert self._run_install(no_torch = True, is_macos = False, is_windows = False) == 0 def test_macos_install_skips_flash_attn_step(self): assert self._run_install(no_torch = False, is_macos = True, is_windows = False) == 0 def test_windows_install_calls_flash_attn_step(self): assert self._run_install(no_torch = False, is_macos = False, is_windows = True) == 1 def test_flash_attn_is_rechecked_after_the_linux_torch_repair(self): """step 13 can move torch from 2.12+cu130 to 2.9+cu126 under cu13torch2.10 (#13244).""" source = (STUDIO_DIR / "install_python_stack.py").read_text(encoding = "utf-8") step = source.split('_progress(_torch_step_label("final"))', 1)[1] step = step.split("# 13w.", 1)[0] moved = step.split( "if _torch_after_repair and _torch_after_repair != _torch_before_repair:", 1 )[1] moved = moved.split("_evict_xformers_built_for_another_torch(", 1)[0] assert "_ensure_flash_attn()" in moved def test_flash_attn_is_rechecked_when_the_windows_flavor_invariant_moves_torch(self): assert ( self._run_install( no_torch = False, is_macos = False, is_windows = True, flavor_moves_torch = True ) == 2 ) class TestWarningColor: def test_optional_install_warning_uses_shell_installer_warn_color(self, capsys): import re install_sh = (STUDIO_DIR.parent / "install.sh").read_text(encoding = "utf-8") code = re.search(r'C_WARN="\$\{_ESC\}\[([0-9;]+)m"', install_sh).group(1) result = subprocess.CompletedProcess(["uv"], 1, "", "") with mock.patch.object(ips, "_HAS_COLOR", True): ips._print_optional_install_failure("flash-attn install", result) out = capsys.readouterr().out assert f"\033[{code}mflash-attn install failed (exit code 1)\033[0m" in out assert "\033[96m" not in out class TestFlashAttnCapabilityFloor: """Unsloth only enables FlashAttention 2 from sm80: below it the wheel is never downloaded.""" def _run(self, monkeypatch, capability): from utils import kernel_install monkeypatch.setattr( ips, "_flash_attn_gpu_below_floor", lambda: kernel_install.gpu_below_flash_attn_floor(lambda *a, **k: None), ) monkeypatch.setattr(kernel_install, "_gpu_capability", lambda run: capability) calls = {"head": 0, "install": 0} notes = [] with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", True), mock.patch.object(ips, "IS_MACOS", False), mock.patch.object(ips, "_flash_attn_importable", return_value = False), mock.patch.object(ips, "_installed_distribution_version", return_value = None), mock.patch.object( ips, "probe_torch_wheel_env", return_value = { "python_tag": "cp313", "abi_tag": "cp313", "torch_mm": "2.13", "cuda_major": "13", "cxx11abi": "TRUE", "platform_tag": "win_amd64", "hip_version": "", }, ), mock.patch.object( ips, "url_exists", side_effect = lambda url: calls.__setitem__("head", calls["head"] + 1) or True, ), mock.patch.object( ips, "install_prebuilt", side_effect = lambda *a, **k: calls.__setitem__("install", calls["install"] + 1) or "installed", ), mock.patch.object(ips, "_note", side_effect = notes.append), mock.patch.object(ips, "_count_install_action"), ): ips._ensure_flash_attn() return calls, notes def test_a_t4_does_not_download_flash_attn(self, monkeypatch): calls, notes = self._run(monkeypatch, (7, 5)) assert calls == {"head": 0, "install": 0} assert any("sm80" in note for note in notes) @pytest.mark.parametrize("capability", [(8, 0), (12, 0), None]) def test_sm80_and_newer_or_no_visible_gpu_still_install(self, monkeypatch, capability): calls, _ = self._run(monkeypatch, capability) assert calls == {"head": 1, "install": 1}