From 23c2e9c4c06b7a9aec3a6c584859fe7d8820aff6 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 28 Aug 2026 06:17:33 +0000 Subject: [PATCH] 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. --- .env.example | 2 ++ audio_cleanup.py | 20 +++++++++++++++++++- config.py | 4 ++++ run.py | 1 + run_worker.py | 5 ++++- tests/test_config.py | 4 ++++ 6 files changed, 34 insertions(+), 2 deletions(-) diff --git a/.env.example b/.env.example index f27112b..fe55965 100644 --- a/.env.example +++ b/.env.example @@ -34,6 +34,8 @@ ENHANCE_HOP_S=0.2 # auto (default): beamform if speechbrain loads, else average channels # force: require DelaySum_Beamformer, off: average L/R BEAMFORM_MODE=auto +# One BLAS/Torch thread per worker (20 workers vs 28 cores) +TORCH_NUM_THREADS=1 # Misc PUBLISH_TIMEOUT_S=60 diff --git a/audio_cleanup.py b/audio_cleanup.py index 40bdeaa..9d5ed88 100644 --- a/audio_cleanup.py +++ b/audio_cleanup.py @@ -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 diff --git a/config.py b/config.py index 7047615..4e9c345 100644 --- a/config.py +++ b/config.py @@ -40,6 +40,7 @@ class Config: enhance_hop_s: float = 0.2 # SpeechBrain Delay-and-Sum (GCC-PHAT) on C920 stereo before enhance beamform_mode: str = "auto" # auto | force | off + torch_num_threads: int = 1 # misc publish_timeout_s: float = 60.0 log_level: str = "INFO" @@ -86,6 +87,7 @@ def load_config(env_path: Optional[Path] = None) -> Config: enhance_chunk_s=float(os.environ.get("ENHANCE_CHUNK_S", "1.0")), enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.2")), beamform_mode=os.environ.get("BEAMFORM_MODE", "auto").lower(), + torch_num_threads=int(os.environ.get("TORCH_NUM_THREADS", "1")), publish_timeout_s=float(os.environ.get("PUBLISH_TIMEOUT_S", "60")), log_level=os.environ.get("LOG_LEVEL", "INFO").upper(), cameras=_parse_cameras(os.environ.get("CAMERAS", "all")), @@ -96,4 +98,6 @@ def load_config(env_path: Optional[Path] = None) -> Config: if cfg.beamform_mode not in ("auto", "force", "off"): raise ValueError( f"BEAMFORM_MODE must be auto|force|off, got {cfg.beamform_mode}") + if cfg.torch_num_threads < 1: + raise ValueError("TORCH_NUM_THREADS must be >= 1") return cfg diff --git a/run.py b/run.py index b3c6fed..e1011a8 100644 --- a/run.py +++ b/run.py @@ -43,6 +43,7 @@ def worker_args_for(cfg, cam) -> list[str]: "--enhance-chunk-s", str(cfg.enhance_chunk_s), "--enhance-hop-s", str(cfg.enhance_hop_s), "--beamform-mode", cfg.beamform_mode, + "--torch-num-threads", str(cfg.torch_num_threads), ] diff --git a/run_worker.py b/run_worker.py index a832f0d..173c425 100644 --- a/run_worker.py +++ b/run_worker.py @@ -38,6 +38,7 @@ def main() -> int: ap.add_argument("--enhance-chunk-s", type=float, default=1.0) ap.add_argument("--enhance-hop-s", type=float, default=0.2) ap.add_argument("--beamform-mode", default="auto") + ap.add_argument("--torch-num-threads", type=int, default=1) args = ap.parse_args() cam = json.loads(args.camera_json) @@ -47,8 +48,9 @@ def main() -> int: ) from config import Config # noqa: F401 (ensures .env is loaded) + from audio_cleanup import AudioCleaner, apply_torch_thread_limits + apply_torch_thread_limits(args.torch_num_threads) import livekit.rtc as rtc - from audio_cleanup import AudioCleaner identity = cam["identity"] video_dev = cam["video_device"] @@ -62,6 +64,7 @@ def main() -> int: chunk_s=args.enhance_chunk_s, hop_s=args.enhance_hop_s, beamform=args.beamform_mode, + torch_num_threads=args.torch_num_threads, ) log.info("audio backend: %s", cleaner.backend) diff --git a/tests/test_config.py b/tests/test_config.py index 5c05441..0beda23 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -9,3 +9,7 @@ def test_parse_cameras_list(): def test_default_beamform_mode_auto(): assert Config().beamform_mode == "auto" + + +def test_torch_num_threads_default_one(): + assert Config().torch_num_threads == 1