mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-10 00:08:58 +00:00
504e253674
The Studio installer skipped Flash Attention on any Blackwell GPU (compute capability >= 10.0: sm_100 B200/B100, sm_120 RTX 50 series), on the assumption that no prebuilt wheels existed for those archs. That assumption is stale: Dao-AILab has shipped wheels with sm_100 / sm_120 kernels since flash-attn v2.7.3, built for CUDA >= 12.8. Replace the blanket has_blackwell_gpu() skip in _ensure_flash_attn and _ensure_flash_attn_for_long_context with a CUDA-version-aware gate. The normal prebuilt-wheel path already builds the exact wheel URL, HEAD-checks it, and falls back gracefully, so it now runs on Blackwell unless the installed torch's CUDA build is older than 12.8 (in which case the matching wheel would lack Blackwell kernels and crash at kernel launch). - wheel_utils: probe now reports full cuda_version; add cuda_supports_blackwell_flash_attn and should_skip_flash_attn_for_blackwell. - install_python_stack / training worker: use the new gate. - Update the corresponding tests.
544 lines
21 KiB
Python
544 lines
21 KiB
Python
"""Tests for the optional FlashAttention installer."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
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 backend.utils import wheel_utils
|
|
|
|
|
|
def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess:
|
|
return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "")
|
|
|
|
|
|
class TestHasBlackwellGpu:
|
|
def setup_method(self):
|
|
wheel_utils.has_blackwell_gpu.cache_clear()
|
|
|
|
def teardown_method(self):
|
|
wheel_utils.has_blackwell_gpu.cache_clear()
|
|
|
|
def test_returns_false_when_nvidia_smi_missing(self):
|
|
with mock.patch.object(wheel_utils.shutil, "which", return_value = None):
|
|
assert wheel_utils.has_blackwell_gpu() is False
|
|
|
|
def test_returns_true_for_sm_100(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is True
|
|
|
|
def test_returns_true_for_sm_120(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is True
|
|
|
|
def test_returns_true_for_sm_121(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is True
|
|
|
|
def test_returns_false_for_sm_90(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is False
|
|
|
|
def test_returns_false_for_sm_89(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is False
|
|
|
|
def test_mixed_gpus_with_one_blackwell_returns_true(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(
|
|
wheel_utils.subprocess,
|
|
"run",
|
|
return_value = _smi_result("8.0\n10.0\n"),
|
|
),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is True
|
|
|
|
def test_returns_false_when_nvidia_smi_fails(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(
|
|
wheel_utils.subprocess,
|
|
"run",
|
|
return_value = _smi_result("", returncode = 1),
|
|
),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is False
|
|
|
|
def test_returns_false_on_subprocess_timeout(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(
|
|
wheel_utils.subprocess,
|
|
"run",
|
|
side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10),
|
|
),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is False
|
|
|
|
def test_returns_false_on_malformed_output(self):
|
|
with (
|
|
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
|
mock.patch.object(
|
|
wheel_utils.subprocess,
|
|
"run",
|
|
return_value = _smi_result("not-a-number\n\n"),
|
|
),
|
|
):
|
|
assert wheel_utils.has_blackwell_gpu() is False
|
|
|
|
|
|
class TestCudaSupportsBlackwellFlashAttn:
|
|
def test_cuda_13_supported(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("13.0") is True
|
|
|
|
def test_cuda_128_is_the_threshold(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("12.8") is True
|
|
|
|
def test_cuda_129_supported(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("12.9") is True
|
|
|
|
def test_major_only_supported(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("13") is True
|
|
|
|
def test_cuda_126_below_threshold(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("12.6") is False
|
|
|
|
def test_cuda_118_below_threshold(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("11.8") is False
|
|
|
|
def test_empty_or_none_is_false(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("") is False
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn(None) is False
|
|
|
|
def test_unparseable_is_false(self):
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("not-a-version") is False
|
|
assert wheel_utils.cuda_supports_blackwell_flash_attn("12.x") is False
|
|
|
|
|
|
class TestShouldSkipFlashAttnForBlackwell:
|
|
def test_non_blackwell_never_skips(self):
|
|
with mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = False):
|
|
assert wheel_utils.should_skip_flash_attn_for_blackwell({"cuda_version": ""}) is False
|
|
|
|
def test_blackwell_with_cuda_13_does_not_skip(self):
|
|
with mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True):
|
|
assert (
|
|
wheel_utils.should_skip_flash_attn_for_blackwell({"cuda_version": "13.0"}) is False
|
|
)
|
|
|
|
def test_blackwell_with_old_cuda_skips(self):
|
|
with mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True):
|
|
assert (
|
|
wheel_utils.should_skip_flash_attn_for_blackwell({"cuda_version": "12.6"}) is True
|
|
)
|
|
|
|
def test_blackwell_probes_when_env_missing(self):
|
|
with (
|
|
mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True),
|
|
mock.patch.object(
|
|
wheel_utils,
|
|
"probe_torch_wheel_env",
|
|
return_value = {"cuda_version": "12.6"},
|
|
),
|
|
):
|
|
assert wheel_utils.should_skip_flash_attn_for_blackwell() is True
|
|
|
|
|
|
class TestFlashAttnWheelSelection:
|
|
def test_torch_210_maps_to_v281(self):
|
|
assert ips._select_flash_attn_version("2.10") == "2.8.1"
|
|
|
|
def test_torch_29_maps_to_v283(self):
|
|
assert ips._select_flash_attn_version("2.9") == "2.8.3"
|
|
|
|
def test_unsupported_torch_has_no_wheel_mapping(self):
|
|
assert ips._select_flash_attn_version("2.11") is None
|
|
|
|
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 TestEnsureFlashAttn:
|
|
def _import_check(self, code: int = 1):
|
|
return subprocess.CompletedProcess(["python", "-c", "import flash_attn"], code)
|
|
|
|
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(wheel_utils, "has_blackwell_gpu", return_value = 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(wheel_utils, "has_blackwell_gpu", return_value = 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
|
|
_, kwargs = install_calls[0]
|
|
assert kwargs["uv_needs_system"] is True
|
|
|
|
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(wheel_utils, "has_blackwell_gpu", return_value = 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(wheel_utils, "has_blackwell_gpu", return_value = 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_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_blackwell_with_supported_cuda_installs_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(wheel_utils, "has_blackwell_gpu", return_value = True),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "13",
|
|
"cuda_version": "13.0",
|
|
"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
|
|
(url,), _ = install_calls[0]
|
|
assert (
|
|
url
|
|
== "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
|
|
)
|
|
|
|
def test_blackwell_with_old_cuda_skips_with_warning(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(wheel_utils, "has_blackwell_gpu", return_value = True),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"cuda_version": "12.6",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
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 any(
|
|
label == "warning" and "Blackwell" in msg and "CUDA" in msg
|
|
for label, msg in step_messages
|
|
)
|
|
|
|
def test_windows_returns_before_blackwell_gate(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", True),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True),
|
|
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
|
|
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_probe.assert_not_called()
|
|
mock_install_wheel.assert_not_called()
|
|
assert not any("Blackwell" in msg for _, msg in step_messages)
|
|
|
|
|
|
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
|