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:
+19
-1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user