vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-10-31 19:20:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 31
unsloth/tests/python/test_flash_attn_install_python_stack.py
alkinun 6d7c6e962d
Studio installer: drop a flash-attn built for a torch that is gone, and pair torch 2.9 with torchao 0.15.0 (#13247) 
* Studio installer: drop a flash-attn built for a torch that is gone

The Linux torch repair (step 13) runs after the flash-attn step. When it moves
torch, e.g. from 2.12+cu130 to 2.9+cu126 on a CUDA-12 driver, the cu13torch2.10
wheel the earlier step installed can no longer import. The repair re-selected
torchao and evicted a mismatched xFormers, but never looked at flash-attn.

Later passes did not heal it either: _ensure_flash_attn only removed a copy it
had just installed itself. With no wheel for the new torch (flash-attn v2.8.3
ships no cu12/torch2.9/cp313 asset) it warned and kept the broken one. xformers
imports flash_attn whenever find_spec sees it, so every diffusers pipeline then
failed with "libcudart.so.13: cannot open shared object file".

- Re-run _ensure_flash_attn when the step 13 repair moves torch.
- A resident flash-attn that does not import, under a torch that does, is
  removed when no working wheel replaces it, and force-replaced when one does
  (a same-version build for another CUDA otherwise counts as installed).

Fixes #13244 (flash-attn half).

* Studio installer: pair torch 2.9 with torchao 0.15.0

diffusers 0.41.0 imports torchao.quantization.FqnToConfig whenever torchao is
installed, and FqnToConfig first shipped in torchao 0.15.0. torch 2.9 hosts got
0.14.0, so once a broken flash-attn was out of the way, every diffusers pipeline
still failed to import with "cannot import name 'FqnToConfig'". transformers 5's
TorchAoConfig has the same 0.15 floor.

0.15.0 is also the build for this torch: its C++ extensions target torch 2.9.1,
while 0.14.0 targets 2.8.0 and skipped them on every 2.9 release. It is
published on every download.pytorch.org leaf that serves torch 2.9 (cu126,
cu128, cu129, cu130, rocm6.3, rocm6.4, xpu, cpu), which also closes the cu129
gap, where only 0.14.1 existed. torch <= 2.8 keeps 0.14.0.

Fixes #13244 (torchao half).

* Trim comments in PR scope

---------

Co-authored-by: mahiatlinux <mahiatlinux@users.noreply.github.com>
2026-10-11 17:36:38 +13:00

772 lines
30 KiB
Python
Raw Blame History

"""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
from utils import wheel_utils
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 on purpose: a torch minor only joins the reuse table once its wheels have actually 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, so the torch2.9 .so raises "undefined symbol" on 2.10+.
assert set(wheel_utils._PREBUILT_WHEEL_TORCH_MM.values()) == {"2.10"}
def test_direct_wheel_url_reuses_torch210_on_211(self):
# causal-conv1d / mamba go through direct_wheel_url; torch 2.11 reuses the torch2.10 wheel filename just like
# flash-attn does.
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",
"cuda_major": "13",
"cxx11abi": "TRUE",
"platform_tag": "linux_x86_64",
},
)
assert url is not None
assert "causal_conv1d-1.6.1+cu13torch2.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",
"cuda_major": "13",
"cxx11abi": "TRUE",
"platform_tag": "linux_x86_64",
},
)
assert url is not None
assert "mamba_ssm-2.3.1+cu13torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" in url
class TestFlashAttnWheelSelection:
def test_torch_210_maps_to_v281(self):
# v2.8.1 is the newest release still publishing the full torch2.10 asset matrix (cu12 + cu13, cp312 + cp313,
# x86_64 + aarch64).
assert ips._select_flash_attn_version("2.10") == "2.8.1"
def test_selected_version_is_never_a_post_release(self):
# The v2.8.3.post1 respin dropped every torch2.10 asset and stops at torch2.9, whose .so will not load on
# torch 2.10+. A future "just take the newest release" bump must fail here instead of shipping that.
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):
# The raw version table has no torch2.11-tagged wheel;
# the URL builder reuses the torch2.10 wheel instead (see test_torch_211_reuses_torch210_wheel).
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/ABI wheel; the import is what decides.
The #5420 / #6961 Blackwell shape: setup must report that rather than claim
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)
),
),
# Import never succeeds, before or after the install; the removal does.
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):
"""Leaving it installed is not "continuing without flash-attn".
unsloth/models/_utils.py finds it by metadata (_package_available) and then imports
the native module in process, so a wheel that killed the probe kills training too.
"""
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):
# First call is the pre-install check (missing), the second verifies the install.
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_skips_install_without_probing(self):
# flash-attn is Linux-only: on Windows the installer returns before probing the torch env or resolving a
# wheel (no Windows wheels are published upstream).
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") 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()
class TestInstallPythonStackFlashAttnIntegration:
def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int:
flash_attn_calls = 0
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),
):
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_skips_flash_attn_step(self):
assert self._run_install(no_torch = False, is_macos = False, is_windows = True) == 0
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
Reference in a new issue View git blame Copy permalink