feat(cuda): добавлен optional runtime для Linux и Windows
This commit is contained in:
1 parent
937eafe8db
commit
99ed5ed92f
10 files changed
+282
-36
No files matched your search
@@ -7,6 +7,25 @@ from local_transcriber.backends.faster_whisper import FasterWhisperBackend
|
||||
from local_transcriber.types import Word
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", ["cpu", "cuda"])
|
||||
def test_bootstrap_only_for_explicit_cuda(monkeypatch, device):
|
||||
"""CPU-модель не загружает optional CUDA-библиотеки."""
|
||||
bootstrap = MagicMock()
|
||||
model_factory = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"local_transcriber.backends.faster_whisper.ensure_cublas_loadable", bootstrap
|
||||
)
|
||||
monkeypatch.setattr("faster_whisper.WhisperModel", model_factory)
|
||||
|
||||
model = FasterWhisperBackend().create_model("/model", device, "float32")
|
||||
|
||||
assert model is model_factory.return_value
|
||||
assert bootstrap.call_count == (1 if device == "cuda" else 0)
|
||||
model_factory.assert_called_once_with(
|
||||
"/model", device=device, compute_type="float32", cpu_threads=0
|
||||
)
|
||||
|
||||
|
||||
def test_transcribe_returns_canonical_words(tmp_path):
|
||||
audio = tmp_path / "audio.wav"
|
||||
raw_word = SimpleNamespace(start=0.2, end=0.7, word=" Привет")
|
||||
|
||||
+74
-3
@@ -688,7 +688,7 @@ def test_cli_load_model_called_with_model_name(tmp_path):
|
||||
|
||||
|
||||
def test_cli_windows_cuda_diagnostic(tmp_path):
|
||||
"""CUDA error on Windows prints choco/winget install hint."""
|
||||
"""При проблеме драйвера Windows не предлагает установку Toolkit."""
|
||||
audio = tmp_path / "test.mp3"
|
||||
audio.write_bytes(b"fake")
|
||||
model = _make_model()
|
||||
@@ -712,8 +712,10 @@ def test_cli_windows_cuda_diagnostic(tmp_path):
|
||||
out = runner.invoke(app, [str(audio), "--device", "cuda"])
|
||||
|
||||
assert out.exit_code == 1
|
||||
assert "choco install cuda" in out.output
|
||||
assert "winget install" in out.output
|
||||
assert "драйвер" in out.output
|
||||
assert "CUDA error: no device" in out.output
|
||||
assert "winget install" not in out.output
|
||||
assert "uv sync --extra" not in out.output
|
||||
|
||||
|
||||
def test_cli_linux_cuda_error_no_windows_hint(tmp_path):
|
||||
@@ -744,6 +746,75 @@ def test_cli_linux_cuda_error_no_windows_hint(tmp_path):
|
||||
assert "choco install cuda" not in out.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("from_config", [False, True])
|
||||
@pytest.mark.parametrize("phase", ["load", "single", "batch"])
|
||||
@pytest.mark.parametrize(
|
||||
("message", "hint", "install_hint"),
|
||||
[
|
||||
("Library cublas64_12.dll is not found or cannot be loaded", "Не найдены", True),
|
||||
("libcublas.so.12: cannot open shared object file", "Не найдены", True),
|
||||
("CUDA error: no kernel image is available", "GPU несовместим", False),
|
||||
("CUDA driver version is insufficient for CUDA runtime version", "драйвер", False),
|
||||
("CUDA out of memory", "Недостаточно памяти", False),
|
||||
],
|
||||
)
|
||||
def test_explicit_cuda_errors_preserve_cause_without_fallback(
|
||||
tmp_path, monkeypatch, from_config, phase, message, hint, install_hint
|
||||
):
|
||||
"""CLI/TOML сохраняют CUDA при загрузке, обработке одного файла и батча."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
files = [tmp_path / "one.wav"]
|
||||
if phase == "batch":
|
||||
files.append(tmp_path / "two.wav")
|
||||
for file in files:
|
||||
file.write_bytes(b"audio")
|
||||
if from_config:
|
||||
(tmp_path / ".transcriber.toml").write_text('device = "cuda"\n')
|
||||
args = [str(file) for file in files]
|
||||
if not from_config:
|
||||
args += ["--device", "cuda"]
|
||||
backend = _make_backend()
|
||||
error = RuntimeError(message)
|
||||
if phase == "load":
|
||||
backend.create_model.side_effect = error
|
||||
else:
|
||||
backend.transcribe.side_effect = error
|
||||
|
||||
with patch("local_transcriber.transcriber.get_backend", return_value=backend) as get_backend:
|
||||
out = runner.invoke(app, args)
|
||||
|
||||
assert out.exit_code == 1
|
||||
assert message in out.output
|
||||
assert hint in out.output
|
||||
assert ("uv sync --extra cuda" in out.output) is install_hint
|
||||
assert "winget install" not in out.output
|
||||
assert "Переключение на CPU" not in out.output
|
||||
assert all(call.args[0] == "cuda" for call in get_backend.call_args_list)
|
||||
|
||||
|
||||
def test_default_cli_stays_on_onnx_with_nvidia_driver(tmp_path, monkeypatch):
|
||||
"""Обычный запуск с nvidia-smi использует прежнюю CPU-модель ONNX."""
|
||||
audio = tmp_path / "test.wav"
|
||||
audio.write_bytes(b"audio")
|
||||
backend = _make_backend()
|
||||
backend.actual_ov_device = None
|
||||
backend.actual_compute_type = "int8"
|
||||
backend.transcribe.return_value = _make_result(device_used="onnx")
|
||||
with (
|
||||
patch("local_transcriber.cli.load_config", return_value={}),
|
||||
patch("shutil.which", return_value="/usr/bin/nvidia-smi"),
|
||||
patch("local_transcriber.transcriber.get_backend", return_value=backend) as get_backend,
|
||||
patch("local_transcriber.cli.write_transcript"),
|
||||
patch("local_transcriber._cuda_bootstrap.ensure_cublas_loadable") as bootstrap,
|
||||
):
|
||||
out = runner.invoke(app, [str(audio)])
|
||||
|
||||
assert out.exit_code == 0, out.output
|
||||
get_backend.assert_called_once_with("onnx", compute_type_explicit=False)
|
||||
assert backend.ensure_model_available.call_args.args[:2] == ("gigaam-v3-e2e-rnnt", "int8")
|
||||
bootstrap.assert_not_called()
|
||||
|
||||
|
||||
def test_cli_device_fallback_warning(tmp_path):
|
||||
"""When auto-detected device differs from actual, show fallback warning."""
|
||||
audio = tmp_path / "test.mp3"
|
||||
|
||||
@@ -48,12 +48,73 @@ def test_ensure_cublas_loads_library(monkeypatch, tmp_path):
|
||||
assert calls[0][1] == ctypes.RTLD_GLOBAL
|
||||
|
||||
|
||||
def test_ensure_cublas_skips_non_linux(monkeypatch):
|
||||
"""На не-Linux платформах -- no-op."""
|
||||
monkeypatch.setattr(sys, "platform", "win32")
|
||||
def test_ensure_cublas_skips_macos(monkeypatch):
|
||||
"""На macOS bootstrap не загружает CUDA."""
|
||||
monkeypatch.setattr(sys, "platform", "darwin")
|
||||
ensure_cublas_loadable() # не должно бросать исключений
|
||||
|
||||
|
||||
def test_windows_optional_package_absent(monkeypatch):
|
||||
"""Windows работает без extra и не меняет поиск DLL."""
|
||||
monkeypatch.setattr(sys, "platform", "win32")
|
||||
monkeypatch.setitem(sys.modules, "nvidia.cublas", None)
|
||||
ensure_cublas_loadable()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("broken", [False, True])
|
||||
def test_windows_preloads_dlls_and_keeps_handles(monkeypatch, tmp_path, broken):
|
||||
"""DLL загружаются по порядку один раз, ошибки не скрываются."""
|
||||
from local_transcriber import _cuda_bootstrap as bootstrap
|
||||
|
||||
bin_dir = tmp_path / "bin"
|
||||
bin_dir.mkdir()
|
||||
nvidia = types.ModuleType("nvidia")
|
||||
nvidia.__path__ = []
|
||||
cublas = types.ModuleType("nvidia.cublas")
|
||||
cublas.__path__ = [str(tmp_path)]
|
||||
nvidia.cublas = cublas
|
||||
monkeypatch.setitem(sys.modules, "nvidia", nvidia)
|
||||
monkeypatch.setitem(sys.modules, "nvidia.cublas", cublas)
|
||||
monkeypatch.setattr(sys, "platform", "win32")
|
||||
monkeypatch.setattr(bootstrap, "_dll_directories", {})
|
||||
monkeypatch.setattr(bootstrap, "_dll_libraries", {})
|
||||
calls = []
|
||||
directory_handle = object()
|
||||
library_handle = object()
|
||||
failure = OSError("cublas64_12.dll: dependency not found")
|
||||
|
||||
def add_directory(path):
|
||||
calls.append(("directory", path))
|
||||
return directory_handle
|
||||
|
||||
def load_library(path):
|
||||
calls.append(("load", path))
|
||||
if broken:
|
||||
raise failure
|
||||
return library_handle
|
||||
|
||||
monkeypatch.setattr(os, "add_dll_directory", add_directory, raising=False)
|
||||
monkeypatch.setattr(ctypes, "WinDLL", load_library, raising=False)
|
||||
original_path = os.environ.get("PATH")
|
||||
|
||||
if broken:
|
||||
with pytest.raises(OSError) as caught:
|
||||
ensure_cublas_loadable()
|
||||
assert caught.value is failure
|
||||
assert bootstrap._dll_libraries == {}
|
||||
else:
|
||||
ensure_cublas_loadable()
|
||||
ensure_cublas_loadable()
|
||||
assert calls == [
|
||||
("directory", str(bin_dir)),
|
||||
("load", str(bin_dir / "cublasLt64_12.dll")),
|
||||
("load", str(bin_dir / "cublas64_12.dll")),
|
||||
]
|
||||
assert list(bootstrap._dll_libraries.values()) == [library_handle] * 2
|
||||
assert bootstrap._dll_directories[str(bin_dir)] is directory_handle
|
||||
assert os.environ.get("PATH") == original_path
|
||||
|
||||
|
||||
def _nvidia_cublas_installed() -> bool:
|
||||
"""Проверяет, что pip-пакет nvidia-cublas-cu12 установлен."""
|
||||
try:
|
||||
|
||||
@@ -7,6 +7,7 @@ from local_transcriber.transcriber import (
|
||||
Segment,
|
||||
TranscribeResult,
|
||||
_transcribe_file,
|
||||
cuda_error_hint,
|
||||
ensure_model_available,
|
||||
load_model,
|
||||
transcribe,
|
||||
@@ -16,6 +17,21 @@ from local_transcriber.types import WordTimestampsUnavailableError
|
||||
# === Helpers ===
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[RuntimeError("out of memory"), RuntimeError("CUDA unknown error")],
|
||||
)
|
||||
def test_unknown_error_does_not_suggest_cuda_installation(error):
|
||||
"""CPU OOM и неизвестная причина не получают совет по установке CUDA."""
|
||||
assert cuda_error_hint(error) is None
|
||||
|
||||
|
||||
def test_windows_loader_error_suggests_extra():
|
||||
"""OSError загрузчика DLL распознаётся без обёртки RuntimeError."""
|
||||
error = OSError("[WinError 126] Could not find module 'cublas64_12.dll'")
|
||||
assert "uv sync --extra cuda" in cuda_error_hint(error)
|
||||
|
||||
|
||||
def _make_result(
|
||||
count: int = 2,
|
||||
language: str = "ru",
|
||||
|
||||
Reference in new issue
Block a user