fix(device): закреплён ONNX CPU для автоматического выбора

This commit is contained in:
Dmitry Dementiev committed 2026-09-21 22:20:27 +03:00
1 parent 2b54f5bc9e
commit 937eafe8db
4 files changed
+35 -15

No files matched your search

+3 -2
View File
@@ -200,8 +200,9 @@ class OnnxAsrBackend:
model = onnx_asr.load_model( model = onnx_asr.load_model(
model=model_path, model=model_path,
quantization=quantization, quantization=quantization,
providers=["CPUExecutionProvider"],
) )
vad = onnx_asr.load_vad("silero") vad = onnx_asr.load_vad("silero", providers=["CPUExecutionProvider"])
self._vad = vad self._vad = vad
return model.with_vad(vad).with_timestamps() return model.with_vad(vad).with_timestamps()
@@ -278,7 +279,7 @@ class OnnxAsrBackend:
fallback = _whisper_fallback_model(model_name) fallback = _whisper_fallback_model(model_name)
raise ValueError( raise ValueError(
f"Модель '{model_name}' относится к Whisper и не поддерживается " f"Модель '{model_name}' относится к Whisper и не поддерживается "
"ONNX-бэкендом. Без CUDA --device auto выбирает ONNX; " "ONNX-бэкендом. --device auto всегда выбирает ONNX CPU; "
f"укажите --device openvino-cpu --model {model_name} на x86, " f"укажите --device openvino-cpu --model {model_name} на x86, "
f"--device cpu --model {fallback} на любой платформе " f"--device cpu --model {fallback} на любой платформе "
f"или --device cuda --model {fallback} при NVIDIA GPU." f"или --device cuda --model {fallback} при NVIDIA GPU."
+2 -5
View File
@@ -2,7 +2,6 @@
import glob import glob
import platform import platform
import shutil
import subprocess import subprocess
import warnings import warnings
from pathlib import Path from pathlib import Path
@@ -16,15 +15,13 @@ SUPPORTED_EXTENSIONS = {
def detect_device(requested: str = "auto") -> str: def detect_device(requested: str = "auto") -> str:
"""Определяет устройство для вычислений. """Определяет устройство для вычислений.
При ``requested="auto"`` выбирает CUDA при наличии ``nvidia-smi``, При ``requested="auto"`` всегда выбирает ONNX на CPU.
иначе ONNX на CPU. CUDA включается только явным выбором пользователя.
``"openvino"`` — алиас для авто-детекта внутри OpenVINO ``"openvino"`` — алиас для авто-детекта внутри OpenVINO
(GPU если доступен, иначе CPU). (GPU если доступен, иначе CPU).
Возвращает всегда конкретное значение (не абстрактный ``"openvino"``). Возвращает всегда конкретное значение (не абстрактный ``"openvino"``).
""" """
if requested == "auto": if requested == "auto":
if shutil.which("nvidia-smi") is not None:
return "cuda"
return "onnx" return "onnx"
if requested == "openvino": if requested == "openvino":
if _is_openvino_gpu_available(): if _is_openvino_gpu_available():
+25 -3
View File
@@ -1,6 +1,7 @@
"""Tests for onnx-asr backend.""" """Tests for onnx-asr backend."""
import warnings import warnings
from unittest.mock import MagicMock
import pytest import pytest
@@ -64,12 +65,12 @@ class TestEnsureModelAvailable:
def with_timestamps(self): def with_timestamps(self):
return self return self
def fake_load_model(*, model, quantization): def fake_load_model(*, model, quantization, providers):
quantizations.append(quantization) quantizations.append(quantization)
return FakeAsrAdapter() return FakeAsrAdapter()
monkeypatch.setattr("onnx_asr.load_model", fake_load_model) monkeypatch.setattr("onnx_asr.load_model", fake_load_model)
monkeypatch.setattr("onnx_asr.load_vad", lambda model: None) monkeypatch.setattr("onnx_asr.load_vad", lambda model, **kwargs: None)
backend = OnnxAsrBackend(compute_type_explicit=False) backend = OnnxAsrBackend(compute_type_explicit=False)
model_id = backend.ensure_model_available( model_id = backend.ensure_model_available(
@@ -87,6 +88,26 @@ class TestEnsureModelAvailable:
class TestCreateModel: class TestCreateModel:
def test_cpu_provider_is_explicit_for_asr_and_vad(self, monkeypatch):
"""Доступность CoreML/CUDA не меняет исполнение ASR и VAD."""
monkeypatch.setattr(
"onnxruntime.get_available_providers",
lambda: ["CoreMLExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"],
)
load_asr = MagicMock()
load_vad = MagicMock()
monkeypatch.setattr("onnx_asr.load_model", load_asr)
monkeypatch.setattr("onnx_asr.load_vad", load_vad)
OnnxAsrBackend().create_model("gigaam-v3-e2e-rnnt", "onnx", "int8")
load_asr.assert_called_once_with(
model="gigaam-v3-e2e-rnnt",
quantization="int8",
providers=["CPUExecutionProvider"],
)
load_vad.assert_called_once_with("silero", providers=["CPUExecutionProvider"])
def test_wraps_vad_model_with_timestamps(self, monkeypatch): def test_wraps_vad_model_with_timestamps(self, monkeypatch):
timestamped_model = object() timestamped_model = object()
@@ -99,7 +120,7 @@ class TestCreateModel:
return FakeVadAdapter() return FakeVadAdapter()
monkeypatch.setattr("onnx_asr.load_model", lambda **kwargs: FakeAsrAdapter()) monkeypatch.setattr("onnx_asr.load_model", lambda **kwargs: FakeAsrAdapter())
monkeypatch.setattr("onnx_asr.load_vad", lambda model: object()) monkeypatch.setattr("onnx_asr.load_vad", lambda model, **kwargs: object())
model = OnnxAsrBackend().create_model("gigaam-v3-e2e-rnnt", "onnx", "int8") model = OnnxAsrBackend().create_model("gigaam-v3-e2e-rnnt", "onnx", "int8")
@@ -127,6 +148,7 @@ class TestCreateModel:
return self return self
monkeypatch.setattr("onnx_asr.load_model", fake_load_model) monkeypatch.setattr("onnx_asr.load_model", fake_load_model)
monkeypatch.setattr("onnx_asr.load_vad", lambda model, **kwargs: None)
backend = OnnxAsrBackend() backend = OnnxAsrBackend()
backend.actual_compute_type = "int8" backend.actual_compute_type = "int8"
+5 -5
View File
@@ -66,17 +66,17 @@ def test_detect_device_explicit_passthrough():
assert detect_device("openvino-cpu") == "openvino-cpu" assert detect_device("openvino-cpu") == "openvino-cpu"
def test_detect_device_auto_cuda_when_nvidia_smi_available(): def test_detect_device_auto_onnx_when_nvidia_smi_available():
"""При доступном nvidia-smi auto выбирает CUDA.""" """При доступном nvidia-smi auto остаётся на ONNX CPU."""
with patch( with patch(
"local_transcriber.utils.shutil.which", return_value="/usr/bin/nvidia-smi" "shutil.which", return_value="/usr/bin/nvidia-smi"
): ):
assert detect_device("auto") == "cuda" assert detect_device("auto") == "onnx"
def test_detect_device_auto_onnx_without_cuda(): def test_detect_device_auto_onnx_without_cuda():
"""Без CUDA auto выбирает ONNX CPU.""" """Без CUDA auto выбирает ONNX CPU."""
with patch("local_transcriber.utils.shutil.which", return_value=None): with patch("shutil.which", return_value=None):
assert detect_device("auto") == "onnx" assert detect_device("auto") == "onnx"