mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-21 14:53:58 +00:00
947b4bbd7c
* Studio: verify the flash-attn import after installing it A prebuilt wheel can install with exit code 0 and then fail to load, so the exit code on its own is not proof the install is usable. Both Blackwell incidents were that shape: in #5420 the older-arch wheels installed and raised on import, and the arch gate added to work around it became the bug in #6961 once Dao-AILab started shipping sm_100 wheels. Verify the import after a zero exit code in the two paths that did not, and treat a failure as not installed. ssm_runtime._install_kernel already did exactly this, for the same reason, so this is that check applied to the setup installer and the long-context training worker. worker._is_importable() catches any exception rather than only ImportError: an arch/ABI mismatch surfaces as OSError or RuntimeError ("undefined symbol"), which the old pre-check would have let escape mid-training. Also drop has_blackwell_gpu() and its two call sites. It has returned False unconditionally since03cbe21, so both blocks were already unreachable and removing them changes no behaviour. An arch gate encodes a snapshot of what upstream publishes and goes stale silently in both directions, whereas the import check catches a wheel that will not load whatever the cause. Checked on an 8x B200 host (compute_cap 10.0): the wheel this resolver builds for cu13/torch2.10/cp313 carries sm_100 and sm_120 cubins, installs through install_wheel(), and runs, matching SDPA to 0.0078 max abs difference on a 2x4096x16x128 bf16 causal forward. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten the comments added by this PR (no behaviour change) * Stub utils.native_tls in the MLX worker config test worker.py imports six utils.* submodules and calls activate_native_tls() at import time; _load_worker_module stubbed five of them. The file therefore loaded only when another test had already imported the real utils.native_tls, so it passed in a full run and failed on its own with "No module named 'utils.native_tls'". Reproduced on upstream main ataaf99488, so this predates this branch: running the file alone, or in any narrow selection, errors during collection. Cross-platform CI runs the PR's own test files as a narrow selection, which is how it surfaced. * Replace a rejected wheel on the fallback, and verify that install too Rejecting a wheel that installed but would not import left the broken distribution in site-packages, and the fallback installed over it rather than replacing it. pip reports it as already satisfied and uv audits it as no change, both exiting 0, and that exit code reached an unconditional return True. Long-context training then continued with an unusable extension, which is the failure the wheel check was meant to stop. Measured against a real flash_attn 2.8.1 install: pip install --no-build-isolation --no-deps flash-attn -> "Requirement already satisfied: flash-attn ... (2.8.1)", exit 0 uv pip install --no-build-isolation --no-deps flash-attn -> "Audited 1 package", "Would make no changes", exit 0 So pass --force-reinstall (pip) / --reinstall (uv) on the fallback when a wheel was rejected, and verify the import afterwards instead of trusting rc=0. The flag is gated on the rejection so the ordinary path keeps installing over nothing and does not rebuild. test_runtime_flash_attn_falls_back_to_pypi mocked an install that exits 0 without ever making the module importable, which under the post-install check is a failed install, not a successful one. It now flips the import stub when the install runs. * Uninstall a rejected wheel instead of reinstalling over it, and probe out of process Two problems with the previous commit. --force-reinstall was the wrong tool. pip documents it as "Reinstall all packages even if they are already up-to-date" and uv's --reinstall as "Reinstall all packages": both scope to the whole resolved transaction, not the named one. flash-attn depends on torch, so on the plain fallback path (which carries no --no-deps) that could reinstall or downgrade the torch the worker is currently running on. Measured with a stand-in whose metadata resolves, since flash-attn's sdist cannot build metadata here: pip install --dry-run requests -> 0 packages pip install --force-reinstall --dry-run requests -> 5 packages (requests + its whole closure) It also fails outright on that path: --force-reinstall makes pip rebuild flash-attn from sdist, and the plain branch passes no --no-build-isolation, so the build dies on "No module named 'torch'". This file already pairs --force-reinstall with --no-deps elsewhere for the same reason. So remove the rejected distribution and leave the fallback command untouched. Nothing is reported as already satisfied afterwards, the transaction stays exactly as wide as it was before, and no rebuild flags change. Second, the post-install probe now runs in a child. A wheel built for the wrong arch can abort or segfault inside the extension's initialiser rather than raising, and except Exception cannot catch that: it would kill the worker and take the fallback with it. The child turns it into a return code, negative for a fatal signal. install_python_stack already probes this way; the worker now matches it. The cheap in-process check stays for "is it already installed", which costs no spawn on every training start and imports the module the worker wants imported anyway. The module name is passed as argv rather than formatted into the -c body. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Bound the installer's flash-attn import probe The worker bounds this same untrusted import at 300s and handles TimeoutExpired; the installer probe had no timeout, so a native extension that hangs in its initialiser rather than failing would leave setup waiting forever and never reach the warning. The asymmetry arrived with the post-install probe added earlier in this branch, which is exactly the call that hands an unvalidated wheel to the import machinery. Bound it the same way and treat a hang, or a probe that cannot be spawned, as an unusable wheel. The timeout is a named constant on both sides so the two paths cannot drift apart again silently. * Tighten the comments added in review (no behaviour change) * Remove a rejected flash-attn wheel in the installer too The setup path warned and continued while leaving the unusable distribution in site-packages, so "Continuing without flash-attn" was not true. unsloth/models/_utils.py gates on _package_available("flash_attn"), which reads metadata rather than importing, and then does an in-process "from flash_attn.flash_attn_interface import ...". A wheel that aborts during native initialisation therefore takes the training process down, and that is exactly the wheel this branch has just rejected. Confirmed with a package that calls os.abort() at import: importing it in process kills the interpreter with SIGABRT, while the out-of-process probe returns rc -6 and survives. Uninstalling a properly installed distribution removes it cleanly. The worker path already uninstalls before its fallback; the installer now matches. * Uninstall with the mode uv was installed with, and stop claiming a failed removal Two problems in the cleanup added in the previous commit. It hard-coded "uv pip uninstall --python". _bootstrap_uv sets UV_NEEDS_SYSTEM exactly when the --python probe FAILED and --system succeeded, so on those hosts the cleanup used the one mode already known not to work there: the uninstall fails, the unusable wheel stays in site-packages, and unsloth/models/_utils.py still finds it by metadata and imports it in process. Mirror the install mode instead. uv documents --system as "Use the system Python to uninstall packages". It also printed "removed it" unconditionally, so a failed removal produced two contradicting warnings in the same run: "Could not remove the unusable flash-attn install" followed by "...; removed it". The helper now returns whether the package is actually gone, and the caller says which happened. A wheel still installed is not the same state as never having installed one, so the failure case says so and points at the manual uninstall. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route every rejected install through one discard path The PyPI fallback returned False without removing the distribution it had just rejected, so the unusable extension stayed where unsloth/models/_utils.py finds it: that gate reads package METADATA and then imports the native module in process, which is exactly the import the isolated probe could not survive. Two ways in: a failed uninstall leaves the wheel "already satisfied" so the fallback no-ops, or the source build itself produces an incompatible extension. This is the third round on the same defect, in a third place, so fix the shape rather than the instance. _reject_install is now the single discard path: it uninstalls, and it reports which state we actually ended in. _uninstall_package returns whether the distribution is gone rather than only logging. The pre-fallback uninstall stays a plain call: a failure there is not fatal, because the fallback then no-ops on "already satisfied", the probe rejects it, and _reject_install reports the real state. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enforce the discard invariant in one place instead of at each return The timeout and install-failure exits returned without discarding the rejected distribution, so a wheel that had already failed its import check stayed where unsloth/models/_utils.py finds it: that gate reads METADATA and only then imports the native module in process. The comment added last round claimed the fallback path would clean up on its own, and that is only true when the fallback exits 0. This is the fourth defect of the same shape, each one an exit somebody did not think to clean up, so enforce the invariant structurally rather than adding a fourth call. _install_package_wheel_first now keeps the two "touch nothing" guards (already importable, offline) and delegates the rest to _attempt_package_install, discarding whatever is left in a finally. Any unsuccessful exit, including ones added later, is covered. The discard is state-based so it is safe to run everywhere: _distribution_present reads metadata via importlib.metadata without importing, so it never loads an extension that would abort, and _reject_install no-ops when there is nothing installed. Verified on the B200 that a working flash-attn 2.8.1 is still untouched: three consecutive runs return True with zero install or uninstall subprocesses, and it stays importable. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * State the installed-but-broken state in the rejection tests Two rejection tests asserted on the discard status message without pinning _distribution_present, so after the discard became state-based they only passed where flash-attn happened to be installed. That was true in the venv I had been running them in and false on a clean runner, where the discard correctly no-ops and no status is sent. Cross-platform CI caught it on ubuntu-latest; the same two tests fail in a fresh venv here and pass with the state pinned. No production change: the invariant holds either way. The tests were reading the machine rather than the code. * Uninstall from the interpreter install_wheel actually installed into install_wheel always targets sys.executable: its uv command passes --python in addition to --system, and its pip fallback runs that interpreter directly. The cleanup passed --system INSTEAD of --python, so on a UV_NEEDS_SYSTEM host it uninstalled from the system Python while the wheel sat in the venv, and setup then reported the wheel removed when it was still there for the metadata gate to import. Pass both, so the removal mirrors the install rather than half of it. The earlier test asserting the --system-only command is replaced, since it pinned the broken command. The worker path already targets --python sys.executable on both sides and needs no change. * Tighten the review comments (no behaviour change) --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
661 lines
26 KiB
Python
661 lines
26 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 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+. Reuse may only point at torch2.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()
|
|
|
|
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()
|
|
|
|
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()
|
|
|
|
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()
|
|
|
|
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)
|
|
|
|
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()
|
|
|
|
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_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
|