diff --git a/src/local_transcriber/backends/onnx_asr.py b/src/local_transcriber/backends/onnx_asr.py index 198442b..ab04c9c 100644 --- a/src/local_transcriber/backends/onnx_asr.py +++ b/src/local_transcriber/backends/onnx_asr.py @@ -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)) diff --git a/tests/test_onnx_asr.py b/tests/test_onnx_asr.py index 8f2fa9d..5ab9aaf 100644 --- a/tests/test_onnx_asr.py +++ b/tests/test_onnx_asr.py @@ -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()