feat(onnx): бюджет потоков применяется к сессиям ASR и VAD
--threads больше не игнорируется ONNX adapter'ом: cpu_threads > 0 задаёт intra_op_num_threads через SessionOptions обеим сессиям, 0 оставляет настройки onnxruntime. Фактический бюджет виден в runtime_info(). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
4dea7f4bc6
commit
e22b2f1bae
2 files changed
+49
-2
No files matched your search
@@ -141,6 +141,7 @@ class OnnxAsrBackend:
|
||||
self._vad: Any = None
|
||||
self._providers: list[str] = []
|
||||
self._quantization: str | None = None
|
||||
self._cpu_threads = 0
|
||||
|
||||
@property
|
||||
def word_timestamps_available(self) -> bool:
|
||||
@@ -194,6 +195,8 @@ class OnnxAsrBackend:
|
||||
|
||||
compute_type маппится в onnx-asr ``quantization`` — это суффикс файла
|
||||
модели; для unquantized (float32/fp32) нужно None, не строку.
|
||||
cpu_threads > 0 задаёт intra_op_num_threads сессиям ASR и VAD;
|
||||
0 оставляет настройки потоков onnxruntime.
|
||||
"""
|
||||
import onnx_asr
|
||||
|
||||
@@ -201,15 +204,19 @@ class OnnxAsrBackend:
|
||||
quantization = _normalize_quantization(actual_compute_type)
|
||||
|
||||
providers = ["CPUExecutionProvider"]
|
||||
session_kwargs: dict[str, Any] = {"providers": providers}
|
||||
if cpu_threads > 0:
|
||||
session_kwargs["sess_options"] = _session_options(cpu_threads)
|
||||
model = onnx_asr.load_model(
|
||||
model=model_path,
|
||||
quantization=quantization,
|
||||
providers=providers,
|
||||
**session_kwargs,
|
||||
)
|
||||
vad = onnx_asr.load_vad("silero", providers=providers)
|
||||
vad = onnx_asr.load_vad("silero", **session_kwargs)
|
||||
self._vad = vad
|
||||
self._providers = providers
|
||||
self._quantization = quantization
|
||||
self._cpu_threads = cpu_threads
|
||||
return model.with_vad(vad).with_timestamps()
|
||||
|
||||
def runtime_info(self) -> dict[str, str]:
|
||||
@@ -229,6 +236,7 @@ class OnnxAsrBackend:
|
||||
"asr_providers": configured,
|
||||
"vad_providers": configured,
|
||||
"quantization": self._quantization or "float32",
|
||||
"intra_op_threads": str(self._cpu_threads) if self._cpu_threads else "по умолчанию",
|
||||
}
|
||||
|
||||
def transcribe(
|
||||
@@ -339,6 +347,15 @@ class OnnxAsrBackend:
|
||||
)
|
||||
|
||||
|
||||
def _session_options(cpu_threads: int) -> Any:
|
||||
"""SessionOptions с бюджетом потоков; общий объект для ASR и VAD."""
|
||||
import onnxruntime
|
||||
|
||||
options = onnxruntime.SessionOptions()
|
||||
options.intra_op_num_threads = cpu_threads
|
||||
return options
|
||||
|
||||
|
||||
def _format_compute_types(quantizations: frozenset[str | None]) -> str:
|
||||
values = [_compute_type_for_quantization(value) for value in quantizations]
|
||||
return ", ".join(sorted(values))
|
||||
|
||||
@@ -108,6 +108,36 @@ class TestCreateModel:
|
||||
)
|
||||
load_vad.assert_called_once_with("silero", providers=["CPUExecutionProvider"])
|
||||
|
||||
def test_thread_budget_reaches_asr_and_vad_sessions(self, monkeypatch):
|
||||
"""--threads задаёт intra_op_num_threads обеим сессиям через SessionOptions."""
|
||||
load_asr = MagicMock()
|
||||
load_vad = MagicMock()
|
||||
monkeypatch.setattr("onnx_asr.load_model", load_asr)
|
||||
monkeypatch.setattr("onnx_asr.load_vad", load_vad)
|
||||
|
||||
backend = OnnxAsrBackend()
|
||||
backend.create_model("gigaam-v3-e2e-rnnt", "onnx", "int8", cpu_threads=4)
|
||||
|
||||
asr_options = load_asr.call_args.kwargs["sess_options"]
|
||||
vad_options = load_vad.call_args.kwargs["sess_options"]
|
||||
assert asr_options.intra_op_num_threads == 4
|
||||
assert vad_options.intra_op_num_threads == 4
|
||||
assert backend.runtime_info()["intra_op_threads"] == "4"
|
||||
|
||||
def test_zero_thread_budget_keeps_library_session_defaults(self, monkeypatch):
|
||||
"""0 не передаёт SessionOptions: потоки остаются на усмотрение onnxruntime."""
|
||||
load_asr = MagicMock()
|
||||
load_vad = MagicMock()
|
||||
monkeypatch.setattr("onnx_asr.load_model", load_asr)
|
||||
monkeypatch.setattr("onnx_asr.load_vad", load_vad)
|
||||
|
||||
backend = OnnxAsrBackend()
|
||||
backend.create_model("gigaam-v3-e2e-rnnt", "onnx", "int8", cpu_threads=0)
|
||||
|
||||
assert "sess_options" not in load_asr.call_args.kwargs
|
||||
assert "sess_options" not in load_vad.call_args.kwargs
|
||||
assert backend.runtime_info()["intra_op_threads"] == "по умолчанию"
|
||||
|
||||
def test_wraps_vad_model_with_timestamps(self, monkeypatch):
|
||||
timestamped_model = object()
|
||||
|
||||
|
||||
Reference in new issue
Block a user