perf: cap torch threads per camera worker

Default TORCH_NUM_THREADS=1 so 14-20 workers do not each spawn a
28-thread BLAS pool.
This commit is contained in:
root
2026-08-28 06:17:33 +00:00
parent 950128deec
commit 23c2e9c4c0
6 changed files with 34 additions and 2 deletions
+19 -1
View File
@@ -27,6 +27,7 @@ The preprocessor is created once per worker process and reused.
from __future__ import annotations
import logging
import os
import time
import numpy as np
@@ -36,6 +37,20 @@ log = logging.getLogger("audio.cleanup")
SAMPLE_RATE = 16000
def apply_torch_thread_limits(n: int = 1) -> None:
"""Cap BLAS/Torch pools so N camera workers do not oversubscribe."""
n = max(1, int(n))
os.environ.setdefault("OMP_NUM_THREADS", str(n))
os.environ.setdefault("MKL_NUM_THREADS", str(n))
os.environ.setdefault("OPENBLAS_NUM_THREADS", str(n))
try:
import torch
torch.set_num_threads(n)
torch.set_num_interop_threads(1)
except Exception: # noqa: BLE001
pass
class AudioCleaner:
"""Stateful audio cleaner. Call .process(...) -> int16 mono.
@@ -48,7 +63,8 @@ class AudioCleaner:
sample_rate: int = SAMPLE_RATE,
chunk_s: float = 1.0,
hop_s: float = 0.2,
beamform: str = "auto"):
beamform: str = "auto",
torch_num_threads: int = 1):
self.mode = mode
self.beamform_mode = beamform
self.sample_rate = sample_rate
@@ -61,6 +77,7 @@ class AudioCleaner:
self._torch = None
self.backend = "off"
self.last_dt_ms = -1.0
self.torch_num_threads = max(1, int(torch_num_threads))
self._load_torch()
if beamform == "force":
@@ -92,6 +109,7 @@ class AudioCleaner:
self.backend = "+".join(parts) if parts else "off"
def _load_torch(self) -> None:
apply_torch_thread_limits(self.torch_num_threads)
try:
import torch
self._torch = torch