[Studio] Install flash attn at setup time for linux (#4979)
* [Studio] Install flash attn at setup time for linux * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup changes Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Test cases * wheel_utils: narrow url_exists exceptions and log at debug level --------- Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
This commit is contained in:
Datta Nimmaturi
GitHub
parent
dccc0ebada
commit
da78c6be71
6 changed files with 815 additions and 207 deletions
|
|
@ -16,26 +16,31 @@ from __future__ import annotations
|
|||
import structlog
|
||||
from loggers import get_logger
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
import json
|
||||
import subprocess as _sp
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from typing import Any, Callable
|
||||
|
||||
logger = get_logger(__name__)
|
||||
from utils.hardware import apply_gpu_ids
|
||||
from utils.wheel_utils import (
|
||||
direct_wheel_url,
|
||||
flash_attn_wheel_url,
|
||||
install_wheel,
|
||||
probe_torch_wheel_env,
|
||||
url_exists,
|
||||
)
|
||||
|
||||
|
||||
_CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4"
|
||||
_CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1"
|
||||
_MAMBA_SSM_RELEASE_TAG = "v2.3.1"
|
||||
_MAMBA_SSM_PACKAGE_VERSION = "2.3.1"
|
||||
_FLASH_ATTN_RUNTIME_MIN_SEQ_LEN = 32768
|
||||
_FLASH_ATTN_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL"
|
||||
|
||||
|
||||
def _model_wants_causal_conv1d(model_name: str) -> bool:
|
||||
|
|
@ -59,184 +64,60 @@ def _model_wants_causal_conv1d(model_name: str) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _causal_conv1d_platform_tag() -> str | None:
|
||||
machine = platform.machine().lower()
|
||||
if sys.platform.startswith("linux"):
|
||||
if machine in {"x86_64", "amd64"}:
|
||||
return "linux_x86_64"
|
||||
if machine in {"aarch64", "arm64"}:
|
||||
return "linux_aarch64"
|
||||
return None
|
||||
# No prebuilt wheels published for macOS or Windows
|
||||
return None
|
||||
|
||||
|
||||
def _probe_causal_conv1d_env() -> dict[str, str] | None:
|
||||
try:
|
||||
probe = _sp.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
(
|
||||
"import json, sys, re, torch; "
|
||||
"parts = torch.__version__.split('+', 1)[0].split('.')[:2]; "
|
||||
"minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; "
|
||||
"torch_mm = parts[0] + '.' + minor; "
|
||||
"print(json.dumps({"
|
||||
"'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', "
|
||||
"'torch_mm': torch_mm, "
|
||||
"'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', "
|
||||
"'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', "
|
||||
"'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()"
|
||||
"}))"
|
||||
),
|
||||
],
|
||||
stdout = _sp.PIPE,
|
||||
stderr = _sp.PIPE,
|
||||
text = True,
|
||||
timeout = 30,
|
||||
)
|
||||
except _sp.TimeoutExpired:
|
||||
logger.warning("Torch environment probe timed out after 30s")
|
||||
return None
|
||||
if probe.returncode != 0:
|
||||
logger.warning(
|
||||
"Failed to probe torch environment for causal-conv1d wheel:\n%s",
|
||||
probe.stdout,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
return json.loads(probe.stdout.strip())
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(
|
||||
"Failed to parse torch environment probe output: %s", probe.stdout
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _direct_wheel_url(
|
||||
*,
|
||||
filename_prefix: str,
|
||||
package_version: str,
|
||||
release_tag: str,
|
||||
release_base_url: str,
|
||||
env: dict[str, str] | None = None,
|
||||
) -> str | None:
|
||||
env = env or _probe_causal_conv1d_env()
|
||||
platform_tag = _causal_conv1d_platform_tag()
|
||||
if env is None or platform_tag is None or not env.get("cuda_major"):
|
||||
return None
|
||||
|
||||
filename = (
|
||||
f"{filename_prefix}-{package_version}"
|
||||
f"+cu{env['cuda_major']}torch{env['torch_mm']}"
|
||||
f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}-{platform_tag}.whl"
|
||||
)
|
||||
return f"{release_base_url}/{release_tag}/{filename}"
|
||||
|
||||
|
||||
def _url_exists(url: str) -> bool:
|
||||
try:
|
||||
request = urllib.request.Request(url, method = "HEAD")
|
||||
with urllib.request.urlopen(request, timeout = 10):
|
||||
return True
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 404:
|
||||
return False
|
||||
logger.warning("Unexpected HTTP error while probing %s: %s", url, exc)
|
||||
return False
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to probe %s: %s", url, exc)
|
||||
return False
|
||||
|
||||
|
||||
def _install_package_wheel_first(
|
||||
*,
|
||||
event_queue: Any,
|
||||
import_name: str,
|
||||
display_name: str,
|
||||
pypi_name: str,
|
||||
pypi_version: str,
|
||||
filename_prefix: str,
|
||||
release_tag: str,
|
||||
release_base_url: str,
|
||||
) -> None:
|
||||
pypi_version: str | None = None,
|
||||
filename_prefix: str | None = None,
|
||||
release_tag: str | None = None,
|
||||
release_base_url: str | None = None,
|
||||
wheel_url_builder: Callable[[dict[str, str] | None], str | None] | None = None,
|
||||
pypi_spec: str | None = None,
|
||||
pypi_status_message: str | None = None,
|
||||
) -> bool:
|
||||
try:
|
||||
__import__(import_name)
|
||||
logger.info("%s already installed", display_name)
|
||||
return
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
env = _probe_causal_conv1d_env()
|
||||
wheel_url = _direct_wheel_url(
|
||||
filename_prefix = filename_prefix,
|
||||
package_version = pypi_version,
|
||||
release_tag = release_tag,
|
||||
release_base_url = release_base_url,
|
||||
env = env,
|
||||
)
|
||||
env = probe_torch_wheel_env(timeout = 30)
|
||||
if wheel_url_builder is not None:
|
||||
wheel_url = wheel_url_builder(env)
|
||||
else:
|
||||
wheel_url = direct_wheel_url(
|
||||
filename_prefix = filename_prefix,
|
||||
package_version = pypi_version,
|
||||
release_tag = release_tag,
|
||||
release_base_url = release_base_url,
|
||||
env = env,
|
||||
)
|
||||
|
||||
if wheel_url is None:
|
||||
logger.info("No compatible %s wheel candidate", display_name)
|
||||
else:
|
||||
if _url_exists(wheel_url):
|
||||
_send_status(event_queue, f"Installing prebuilt {display_name} wheel...")
|
||||
installed = False
|
||||
# Try uv first if available, then fall back to pip
|
||||
if shutil.which("uv"):
|
||||
uv_cmd = [
|
||||
"uv",
|
||||
"pip",
|
||||
"install",
|
||||
"--python",
|
||||
sys.executable,
|
||||
"--no-deps",
|
||||
wheel_url,
|
||||
]
|
||||
result = _sp.run(
|
||||
uv_cmd,
|
||||
stdout = _sp.PIPE,
|
||||
stderr = _sp.STDOUT,
|
||||
text = True,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
installed = True
|
||||
else:
|
||||
logger.warning(
|
||||
"uv failed to install %s wheel:\n%s",
|
||||
display_name,
|
||||
result.stdout,
|
||||
)
|
||||
if not installed:
|
||||
pip_cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--no-deps",
|
||||
wheel_url,
|
||||
]
|
||||
result = _sp.run(
|
||||
pip_cmd,
|
||||
stdout = _sp.PIPE,
|
||||
stderr = _sp.STDOUT,
|
||||
text = True,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
installed = True
|
||||
else:
|
||||
logger.warning(
|
||||
"pip failed to install %s wheel:\n%s",
|
||||
display_name,
|
||||
result.stdout,
|
||||
)
|
||||
if installed:
|
||||
elif url_exists(wheel_url):
|
||||
_send_status(event_queue, f"Installing prebuilt {display_name} wheel...")
|
||||
for installer, result in install_wheel(
|
||||
wheel_url,
|
||||
python_executable = sys.executable,
|
||||
use_uv = bool(shutil.which("uv")),
|
||||
run = _sp.run,
|
||||
):
|
||||
if result.returncode == 0:
|
||||
logger.info("Installed prebuilt %s wheel successfully", display_name)
|
||||
return
|
||||
else:
|
||||
logger.info("No published %s wheel found: %s", display_name, wheel_url)
|
||||
return True
|
||||
logger.warning(
|
||||
"%s failed to install %s wheel:\n%s",
|
||||
installer,
|
||||
display_name,
|
||||
result.stdout,
|
||||
)
|
||||
else:
|
||||
logger.info("No published %s wheel found: %s", display_name, wheel_url)
|
||||
|
||||
is_hip = env and env.get("hip_version")
|
||||
if is_hip and not shutil.which("hipcc"):
|
||||
|
|
@ -249,43 +130,62 @@ def _install_package_wheel_first(
|
|||
event_queue,
|
||||
f"{display_name}: hipcc not found (ROCm HIP SDK required)",
|
||||
)
|
||||
return
|
||||
return False
|
||||
|
||||
if is_hip:
|
||||
_send_status(
|
||||
event_queue,
|
||||
f"Compiling {display_name} from source for ROCm "
|
||||
"(this may take several minutes)...",
|
||||
)
|
||||
else:
|
||||
_send_status(event_queue, f"Installing {display_name} from PyPI...")
|
||||
if pypi_spec is None:
|
||||
pypi_spec = f"{pypi_name}=={pypi_version}"
|
||||
|
||||
if pypi_status_message is None:
|
||||
if is_hip:
|
||||
pypi_status_message = (
|
||||
f"Compiling {display_name} from source for ROCm "
|
||||
"(this may take several minutes)..."
|
||||
)
|
||||
else:
|
||||
pypi_status_message = f"Installing {display_name} from PyPI..."
|
||||
|
||||
_send_status(event_queue, pypi_status_message)
|
||||
|
||||
# Prefer uv for faster dependency resolution when available
|
||||
if shutil.which("uv"):
|
||||
pypi_cmd = [
|
||||
"uv",
|
||||
"pip",
|
||||
"install",
|
||||
"--python",
|
||||
sys.executable,
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
]
|
||||
# Avoid stale cache artifacts from partial HIP source builds
|
||||
if is_hip:
|
||||
pypi_cmd.append("--no-cache")
|
||||
pypi_cmd.append(f"{pypi_name}=={pypi_version}")
|
||||
plain_pypi_install = pypi_version is None
|
||||
if plain_pypi_install:
|
||||
if shutil.which("uv"):
|
||||
pypi_cmd = [
|
||||
"uv",
|
||||
"pip",
|
||||
"install",
|
||||
"--python",
|
||||
sys.executable,
|
||||
pypi_spec,
|
||||
]
|
||||
else:
|
||||
pypi_cmd = [sys.executable, "-m", "pip", "install", pypi_spec]
|
||||
else:
|
||||
pypi_cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
"--no-cache-dir",
|
||||
f"{pypi_name}=={pypi_version}",
|
||||
]
|
||||
if shutil.which("uv"):
|
||||
pypi_cmd = [
|
||||
"uv",
|
||||
"pip",
|
||||
"install",
|
||||
"--python",
|
||||
sys.executable,
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
]
|
||||
# Avoid stale cache artifacts from partial HIP source builds
|
||||
if is_hip:
|
||||
pypi_cmd.append("--no-cache")
|
||||
pypi_cmd.append(pypi_spec)
|
||||
else:
|
||||
pypi_cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
"--no-cache-dir",
|
||||
pypi_spec,
|
||||
]
|
||||
|
||||
# Source compilation on ROCm can take 10-30 minutes; use a generous
|
||||
# timeout. Non-HIP installs preserve the pre-existing "no timeout"
|
||||
|
|
@ -313,7 +213,7 @@ def _install_package_wheel_first(
|
|||
f"{display_name} installation timed out after "
|
||||
f"{_run_kwargs.get('timeout')}s",
|
||||
)
|
||||
return
|
||||
return False
|
||||
|
||||
if result.returncode != 0:
|
||||
if is_hip:
|
||||
|
|
@ -337,12 +237,13 @@ def _install_package_wheel_first(
|
|||
display_name,
|
||||
result.stdout,
|
||||
)
|
||||
return
|
||||
return False
|
||||
|
||||
if is_hip:
|
||||
logger.info("Compiled and installed %s from source for ROCm", display_name)
|
||||
else:
|
||||
logger.info("Installed %s from PyPI", display_name)
|
||||
return True
|
||||
|
||||
|
||||
def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:
|
||||
|
|
@ -389,6 +290,31 @@ def _ensure_mamba_ssm(event_queue: Any, model_name: str) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool:
|
||||
if os.getenv(_FLASH_ATTN_SKIP_ENV) == "1":
|
||||
return False
|
||||
if max_seq_length < _FLASH_ATTN_RUNTIME_MIN_SEQ_LEN:
|
||||
return False
|
||||
return sys.platform.startswith("linux")
|
||||
|
||||
|
||||
def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None:
|
||||
if not _should_try_runtime_flash_attn_install(max_seq_length):
|
||||
return
|
||||
|
||||
installed = _install_package_wheel_first(
|
||||
event_queue = event_queue,
|
||||
import_name = "flash_attn",
|
||||
display_name = "flash-attn",
|
||||
pypi_name = "flash-attn",
|
||||
wheel_url_builder = flash_attn_wheel_url,
|
||||
pypi_spec = "flash-attn",
|
||||
pypi_status_message = "Installing flash-attn from PyPI for long-context training...",
|
||||
)
|
||||
if not installed:
|
||||
_send_status(event_queue, "Continuing without flash-attn")
|
||||
|
||||
|
||||
def _activate_transformers_version(model_name: str) -> None:
|
||||
"""Activate the correct transformers version BEFORE any ML imports."""
|
||||
# Ensure backend is on path for utils imports
|
||||
|
|
@ -471,6 +397,10 @@ def run_training_process(
|
|||
try:
|
||||
_ensure_causal_conv1d_fast_path(event_queue, model_name)
|
||||
_ensure_mamba_ssm(event_queue, model_name)
|
||||
_ensure_flash_attn_for_long_context(
|
||||
event_queue,
|
||||
int(config.get("max_seq_length", 2048)),
|
||||
)
|
||||
except Exception as exc:
|
||||
event_queue.put(
|
||||
{
|
||||
|
|
|
|||
154
studio/backend/tests/test_training_worker_flash_attn.py
Normal file
154
studio/backend/tests/test_training_worker_flash_attn.py
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import subprocess
|
||||
import sys
|
||||
from unittest import mock
|
||||
|
||||
from core.training import worker
|
||||
|
||||
|
||||
def _missing_flash_attn_import():
|
||||
real_import = builtins.__import__
|
||||
|
||||
def fake_import(name, globals = None, locals = None, fromlist = (), level = 0):
|
||||
if name == "flash_attn":
|
||||
raise ImportError
|
||||
return real_import(name, globals, locals, fromlist, level)
|
||||
|
||||
return fake_import
|
||||
|
||||
|
||||
def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch):
|
||||
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
||||
assert worker._should_try_runtime_flash_attn_install(32767) is False
|
||||
assert worker._should_try_runtime_flash_attn_install(
|
||||
32768
|
||||
) is sys.platform.startswith("linux")
|
||||
|
||||
monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1")
|
||||
assert worker._should_try_runtime_flash_attn_install(32768) is False
|
||||
|
||||
|
||||
def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch):
|
||||
statuses: list[str] = []
|
||||
|
||||
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"flash_attn_wheel_url",
|
||||
lambda env: "https://example.com/fa.whl",
|
||||
)
|
||||
monkeypatch.setattr(worker, "url_exists", lambda url: True)
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"_send_status",
|
||||
lambda queue, message: statuses.append(message),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"install_wheel",
|
||||
lambda *args, **kwargs: [("pip", subprocess.CompletedProcess(["pip"], 0, ""))],
|
||||
)
|
||||
|
||||
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768)
|
||||
|
||||
assert statuses == ["Installing prebuilt flash-attn wheel..."]
|
||||
|
||||
|
||||
def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch):
|
||||
calls: list[list[str]] = []
|
||||
statuses: list[str] = []
|
||||
|
||||
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"probe_torch_wheel_env",
|
||||
lambda timeout = 30: {
|
||||
"python_tag": "cp313",
|
||||
"torch_mm": "2.10",
|
||||
"cuda_major": "13",
|
||||
"cxx11abi": "TRUE",
|
||||
"platform_tag": "linux_x86_64",
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"flash_attn_wheel_url",
|
||||
lambda env: "https://example.com/fa.whl",
|
||||
)
|
||||
monkeypatch.setattr(worker, "url_exists", lambda url: False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: None)
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"_send_status",
|
||||
lambda queue, message: statuses.append(message),
|
||||
)
|
||||
monkeypatch.setattr(worker, "install_wheel", mock.Mock())
|
||||
|
||||
def fake_run(cmd, stdout = None, stderr = None, text = None):
|
||||
calls.append(list(cmd))
|
||||
return subprocess.CompletedProcess(cmd, 0, "")
|
||||
|
||||
monkeypatch.setattr(worker._sp, "run", fake_run)
|
||||
|
||||
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768)
|
||||
|
||||
assert statuses == ["Installing flash-attn from PyPI for long-context training..."]
|
||||
assert calls == [[sys.executable, "-m", "pip", "install", "flash-attn"]]
|
||||
|
||||
|
||||
def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch):
|
||||
monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1")
|
||||
monkeypatch.setattr(worker._sp, "run", mock.Mock())
|
||||
|
||||
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768)
|
||||
|
||||
worker._sp.run.assert_not_called()
|
||||
|
||||
|
||||
def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch):
|
||||
install_mock = mock.Mock(return_value = True)
|
||||
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
||||
|
||||
worker._ensure_causal_conv1d_fast_path(
|
||||
event_queue = [],
|
||||
model_name = "tiiuae/Falcon-H1-0.5B-Instruct",
|
||||
)
|
||||
|
||||
install_mock.assert_called_once_with(
|
||||
event_queue = [],
|
||||
import_name = "causal_conv1d",
|
||||
display_name = "causal-conv1d",
|
||||
pypi_name = "causal-conv1d",
|
||||
pypi_version = worker._CAUSAL_CONV1D_PACKAGE_VERSION,
|
||||
filename_prefix = "causal_conv1d",
|
||||
release_tag = worker._CAUSAL_CONV1D_RELEASE_TAG,
|
||||
release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download",
|
||||
)
|
||||
|
||||
|
||||
def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch):
|
||||
install_mock = mock.Mock(return_value = True)
|
||||
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
||||
|
||||
worker._ensure_mamba_ssm(
|
||||
event_queue = [],
|
||||
model_name = "tiiuae/Falcon-H1-0.5B-Instruct",
|
||||
)
|
||||
|
||||
install_mock.assert_called_once_with(
|
||||
event_queue = [],
|
||||
import_name = "mamba_ssm",
|
||||
display_name = "mamba-ssm",
|
||||
pypi_name = "mamba-ssm",
|
||||
pypi_version = worker._MAMBA_SSM_PACKAGE_VERSION,
|
||||
filename_prefix = "mamba_ssm",
|
||||
release_tag = worker._MAMBA_SSM_RELEASE_TAG,
|
||||
release_base_url = "https://github.com/state-spaces/mamba/releases/download",
|
||||
)
|
||||
170
studio/backend/utils/wheel_utils.py
Normal file
170
studio/backend/utils/wheel_utils.py
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import platform
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from typing import Callable
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
FLASH_ATTN_RELEASE_BASE_URL = (
|
||||
"https://github.com/Dao-AILab/flash-attention/releases/download"
|
||||
)
|
||||
|
||||
|
||||
def linux_wheel_platform_tag() -> str | None:
|
||||
machine = platform.machine().lower()
|
||||
if sys.platform.startswith("linux"):
|
||||
if machine in {"x86_64", "amd64"}:
|
||||
return "linux_x86_64"
|
||||
if machine in {"aarch64", "arm64"}:
|
||||
return "linux_aarch64"
|
||||
# No prebuilt wheels published for macOS or Windows
|
||||
return None
|
||||
|
||||
|
||||
def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | None:
|
||||
platform_tag = linux_wheel_platform_tag()
|
||||
if platform_tag is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
probe = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
(
|
||||
"import json, sys, re, torch; "
|
||||
"parts = torch.__version__.split('+', 1)[0].split('.')[:2]; "
|
||||
"minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; "
|
||||
"torch_mm = parts[0] + '.' + minor; "
|
||||
"print(json.dumps({"
|
||||
"'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', "
|
||||
"'torch_mm': torch_mm, "
|
||||
"'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', "
|
||||
"'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', "
|
||||
"'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()"
|
||||
"}))"
|
||||
),
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.PIPE,
|
||||
text = True,
|
||||
timeout = timeout,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
return None
|
||||
|
||||
if probe.returncode != 0:
|
||||
return None
|
||||
|
||||
try:
|
||||
env = json.loads(probe.stdout.strip())
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
env["platform_tag"] = platform_tag
|
||||
return env
|
||||
|
||||
|
||||
def direct_wheel_url(
|
||||
*,
|
||||
filename_prefix: str,
|
||||
package_version: str,
|
||||
release_tag: str,
|
||||
release_base_url: str,
|
||||
env: dict[str, str] | None,
|
||||
) -> str | None:
|
||||
if env is None or not env.get("cuda_major"):
|
||||
return None
|
||||
|
||||
filename = (
|
||||
f"{filename_prefix}-{package_version}"
|
||||
f"+cu{env['cuda_major']}torch{env['torch_mm']}"
|
||||
f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}"
|
||||
f"-{env['platform_tag']}.whl"
|
||||
)
|
||||
return f"{release_base_url}/{release_tag}/{filename}"
|
||||
|
||||
|
||||
def flash_attn_package_version(torch_mm: str) -> str | None:
|
||||
if torch_mm == "2.10":
|
||||
return "2.8.1"
|
||||
try:
|
||||
major, minor = (int(part) for part in torch_mm.split(".", 1))
|
||||
except ValueError:
|
||||
return None
|
||||
if major == 2 and 4 <= minor <= 9:
|
||||
return "2.8.3"
|
||||
return None
|
||||
|
||||
|
||||
def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None:
|
||||
if env is None:
|
||||
return None
|
||||
package_version = flash_attn_package_version(env["torch_mm"])
|
||||
if package_version is None:
|
||||
return None
|
||||
return direct_wheel_url(
|
||||
filename_prefix = "flash_attn",
|
||||
package_version = package_version,
|
||||
release_tag = f"v{package_version}",
|
||||
release_base_url = FLASH_ATTN_RELEASE_BASE_URL,
|
||||
env = env,
|
||||
)
|
||||
|
||||
|
||||
def install_wheel(
|
||||
wheel_url: str,
|
||||
*,
|
||||
python_executable: str,
|
||||
use_uv: bool,
|
||||
uv_needs_system: bool = False,
|
||||
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
|
||||
) -> list[tuple[str, subprocess.CompletedProcess[str]]]:
|
||||
attempts: list[tuple[str, subprocess.CompletedProcess[str]]] = []
|
||||
|
||||
# Try uv first if available, then fall back to pip
|
||||
if use_uv and shutil.which("uv"):
|
||||
uv_cmd = ["uv", "pip", "install"]
|
||||
if uv_needs_system:
|
||||
uv_cmd.append("--system")
|
||||
uv_cmd.extend(["--python", python_executable, "--no-deps", wheel_url])
|
||||
result = run(
|
||||
uv_cmd,
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
)
|
||||
attempts.append(("uv", result))
|
||||
if result.returncode == 0:
|
||||
return attempts
|
||||
|
||||
pip_cmd = [python_executable, "-m", "pip", "install", "--no-deps", wheel_url]
|
||||
result = run(
|
||||
pip_cmd,
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
)
|
||||
attempts.append(("pip", result))
|
||||
return attempts
|
||||
|
||||
|
||||
def url_exists(url: str) -> bool:
|
||||
try:
|
||||
request = urllib.request.Request(url, method = "HEAD")
|
||||
with urllib.request.urlopen(request, timeout = 10):
|
||||
return True
|
||||
except urllib.error.HTTPError as exc:
|
||||
_logger.debug("url_exists(%s): HTTP %s", url, exc.code)
|
||||
except (urllib.error.URLError, TimeoutError) as exc:
|
||||
_logger.debug("url_exists(%s): %s", url, exc)
|
||||
return False
|
||||
|
|
@ -21,6 +21,14 @@ import tempfile
|
|||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
from backend.utils.wheel_utils import (
|
||||
flash_attn_package_version,
|
||||
flash_attn_wheel_url,
|
||||
install_wheel,
|
||||
probe_torch_wheel_env,
|
||||
url_exists,
|
||||
)
|
||||
|
||||
IS_WINDOWS = sys.platform == "win32"
|
||||
IS_MACOS = sys.platform == "darwin"
|
||||
IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64"
|
||||
|
|
@ -368,7 +376,6 @@ NO_TORCH = _infer_no_torch()
|
|||
VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1"
|
||||
|
||||
# Progress bar state -- updated by _progress() as each install step runs.
|
||||
# _TOTAL counts: pip-upgrade + 7 shared steps + triton (non-Windows) + local-plugin + finalize
|
||||
# Update _TOTAL here if you add or remove install steps in install_python_stack().
|
||||
_STEP: int = 0
|
||||
_TOTAL: int = 0 # set at runtime in install_python_stack() based on platform
|
||||
|
|
@ -535,6 +542,66 @@ NO_TORCH_SKIP_PACKAGES = {
|
|||
"transformers-cfg",
|
||||
}
|
||||
|
||||
|
||||
def _select_flash_attn_version(torch_mm: str) -> str | None:
|
||||
return flash_attn_package_version(torch_mm)
|
||||
|
||||
|
||||
def _build_flash_attn_wheel_url(env: dict[str, str]) -> str | None:
|
||||
return flash_attn_wheel_url(env)
|
||||
|
||||
|
||||
def _print_optional_install_failure(
|
||||
label: str, result: subprocess.CompletedProcess[str]
|
||||
) -> None:
|
||||
_step("warning", f"{label} failed (exit code {result.returncode})", _cyan)
|
||||
if result.stdout:
|
||||
print(result.stdout.strip())
|
||||
|
||||
|
||||
def _flash_attn_install_disabled() -> bool:
|
||||
return os.getenv("UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL") == "1"
|
||||
|
||||
|
||||
def _ensure_flash_attn() -> None:
|
||||
if NO_TORCH or IS_WINDOWS or IS_MACOS:
|
||||
return
|
||||
if _flash_attn_install_disabled():
|
||||
return
|
||||
if (
|
||||
subprocess.run(
|
||||
[sys.executable, "-c", "import flash_attn"],
|
||||
stdout = subprocess.DEVNULL,
|
||||
stderr = subprocess.DEVNULL,
|
||||
).returncode
|
||||
== 0
|
||||
):
|
||||
return
|
||||
|
||||
env = probe_torch_wheel_env()
|
||||
wheel_url = _build_flash_attn_wheel_url(env) if env else None
|
||||
if wheel_url and url_exists(wheel_url):
|
||||
for installer, wheel_result in install_wheel(
|
||||
wheel_url,
|
||||
python_executable = sys.executable,
|
||||
use_uv = USE_UV,
|
||||
uv_needs_system = UV_NEEDS_SYSTEM,
|
||||
):
|
||||
if wheel_result.returncode == 0:
|
||||
return
|
||||
_print_optional_install_failure(
|
||||
f"Installing flash-attn prebuilt wheel with {installer}",
|
||||
wheel_result,
|
||||
)
|
||||
_step("warning", "Continuing without flash-attn", _cyan)
|
||||
return
|
||||
|
||||
if wheel_url is None:
|
||||
_step("warning", "No compatible flash-attn prebuilt wheel found", _cyan)
|
||||
else:
|
||||
_step("warning", "No published flash-attn prebuilt wheel found", _cyan)
|
||||
|
||||
|
||||
# -- uv bootstrap ------------------------------------------------------
|
||||
|
||||
USE_UV = False # Set by _bootstrap_uv() at the start of install_python_stack()
|
||||
|
|
@ -762,10 +829,8 @@ def install_python_stack() -> int:
|
|||
base_total = 10 if IS_WINDOWS else 11
|
||||
if IS_MACOS:
|
||||
base_total -= 1 # triton step is skipped on macOS
|
||||
# ROCm torch check steps (Linux only, non-macOS, non-no-torch):
|
||||
# one early check (step 2b) and one final repair (step 13).
|
||||
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
|
||||
base_total += 2
|
||||
base_total += 3
|
||||
_TOTAL = (base_total - 1) if skip_base else base_total
|
||||
|
||||
# 1. Try to use uv for faster installs (must happen before pip upgrade
|
||||
|
|
@ -979,6 +1044,10 @@ def install_python_stack() -> int:
|
|||
constrain = False,
|
||||
)
|
||||
|
||||
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
|
||||
_progress("flash-attn")
|
||||
_ensure_flash_attn()
|
||||
|
||||
# # 6. Patch: override llama_cpp.py with fix from unsloth-zoo feature/llama-cpp-windows-support branch
|
||||
# patch_package_file(
|
||||
# "unsloth-zoo",
|
||||
|
|
|
|||
282
tests/python/test_flash_attn_install_python_stack.py
Normal file
282
tests/python/test_flash_attn_install_python_stack.py
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
"""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))
|
||||
|
||||
import install_python_stack as ips
|
||||
|
||||
|
||||
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(
|
||||
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_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()
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -419,6 +419,9 @@ class TestInstallPythonStackSubprocessMock:
|
|||
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, "_ensure_flash_attn", return_value = None),
|
||||
mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False),
|
||||
mock.patch.object(ips, "_has_rocm_gpu", return_value = False),
|
||||
mock.patch("subprocess.run", side_effect = mock_run),
|
||||
mock.patch.object(ips, "_bootstrap_uv", return_value = True),
|
||||
mock.patch.object(
|
||||
|
|
@ -637,7 +640,7 @@ class TestInstallShNoTorchFlag:
|
|||
def test_cpu_hint_message_exists(self):
|
||||
"""CPU hint message must exist in install.sh."""
|
||||
assert (
|
||||
"No NVIDIA GPU detected" in self.source
|
||||
"No GPU detected" in self.source
|
||||
), "CPU hint message not found in install.sh"
|
||||
assert (
|
||||
"--no-torch" in self.source
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue