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
+2
View File
@@ -34,6 +34,8 @@ ENHANCE_HOP_S=0.2
# auto (default): beamform if speechbrain loads, else average channels # auto (default): beamform if speechbrain loads, else average channels
# force: require DelaySum_Beamformer, off: average L/R # force: require DelaySum_Beamformer, off: average L/R
BEAMFORM_MODE=auto BEAMFORM_MODE=auto
# One BLAS/Torch thread per worker (20 workers vs 28 cores)
TORCH_NUM_THREADS=1
# Misc # Misc
PUBLISH_TIMEOUT_S=60 PUBLISH_TIMEOUT_S=60
+19 -1
View File
@@ -27,6 +27,7 @@ The preprocessor is created once per worker process and reused.
from __future__ import annotations from __future__ import annotations
import logging import logging
import os
import time import time
import numpy as np import numpy as np
@@ -36,6 +37,20 @@ log = logging.getLogger("audio.cleanup")
SAMPLE_RATE = 16000 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: class AudioCleaner:
"""Stateful audio cleaner. Call .process(...) -> int16 mono. """Stateful audio cleaner. Call .process(...) -> int16 mono.
@@ -48,7 +63,8 @@ class AudioCleaner:
sample_rate: int = SAMPLE_RATE, sample_rate: int = SAMPLE_RATE,
chunk_s: float = 1.0, chunk_s: float = 1.0,
hop_s: float = 0.2, hop_s: float = 0.2,
beamform: str = "auto"): beamform: str = "auto",
torch_num_threads: int = 1):
self.mode = mode self.mode = mode
self.beamform_mode = beamform self.beamform_mode = beamform
self.sample_rate = sample_rate self.sample_rate = sample_rate
@@ -61,6 +77,7 @@ class AudioCleaner:
self._torch = None self._torch = None
self.backend = "off" self.backend = "off"
self.last_dt_ms = -1.0 self.last_dt_ms = -1.0
self.torch_num_threads = max(1, int(torch_num_threads))
self._load_torch() self._load_torch()
if beamform == "force": if beamform == "force":
@@ -92,6 +109,7 @@ class AudioCleaner:
self.backend = "+".join(parts) if parts else "off" self.backend = "+".join(parts) if parts else "off"
def _load_torch(self) -> None: def _load_torch(self) -> None:
apply_torch_thread_limits(self.torch_num_threads)
try: try:
import torch import torch
self._torch = torch self._torch = torch
+4
View File
@@ -40,6 +40,7 @@ class Config:
enhance_hop_s: float = 0.2 enhance_hop_s: float = 0.2
# SpeechBrain Delay-and-Sum (GCC-PHAT) on C920 stereo before enhance # SpeechBrain Delay-and-Sum (GCC-PHAT) on C920 stereo before enhance
beamform_mode: str = "auto" # auto | force | off beamform_mode: str = "auto" # auto | force | off
torch_num_threads: int = 1
# misc # misc
publish_timeout_s: float = 60.0 publish_timeout_s: float = 60.0
log_level: str = "INFO" 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_chunk_s=float(os.environ.get("ENHANCE_CHUNK_S", "1.0")),
enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.2")), enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.2")),
beamform_mode=os.environ.get("BEAMFORM_MODE", "auto").lower(), 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")), publish_timeout_s=float(os.environ.get("PUBLISH_TIMEOUT_S", "60")),
log_level=os.environ.get("LOG_LEVEL", "INFO").upper(), log_level=os.environ.get("LOG_LEVEL", "INFO").upper(),
cameras=_parse_cameras(os.environ.get("CAMERAS", "all")), 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"): if cfg.beamform_mode not in ("auto", "force", "off"):
raise ValueError( raise ValueError(
f"BEAMFORM_MODE must be auto|force|off, got {cfg.beamform_mode}") 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 return cfg
+1
View File
@@ -43,6 +43,7 @@ def worker_args_for(cfg, cam) -> list[str]:
"--enhance-chunk-s", str(cfg.enhance_chunk_s), "--enhance-chunk-s", str(cfg.enhance_chunk_s),
"--enhance-hop-s", str(cfg.enhance_hop_s), "--enhance-hop-s", str(cfg.enhance_hop_s),
"--beamform-mode", cfg.beamform_mode, "--beamform-mode", cfg.beamform_mode,
"--torch-num-threads", str(cfg.torch_num_threads),
] ]
+4 -1
View File
@@ -38,6 +38,7 @@ def main() -> int:
ap.add_argument("--enhance-chunk-s", type=float, default=1.0) 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("--enhance-hop-s", type=float, default=0.2)
ap.add_argument("--beamform-mode", default="auto") ap.add_argument("--beamform-mode", default="auto")
ap.add_argument("--torch-num-threads", type=int, default=1)
args = ap.parse_args() args = ap.parse_args()
cam = json.loads(args.camera_json) cam = json.loads(args.camera_json)
@@ -47,8 +48,9 @@ def main() -> int:
) )
from config import Config # noqa: F401 (ensures .env is loaded) 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 import livekit.rtc as rtc
from audio_cleanup import AudioCleaner
identity = cam["identity"] identity = cam["identity"]
video_dev = cam["video_device"] video_dev = cam["video_device"]
@@ -62,6 +64,7 @@ def main() -> int:
chunk_s=args.enhance_chunk_s, chunk_s=args.enhance_chunk_s,
hop_s=args.enhance_hop_s, hop_s=args.enhance_hop_s,
beamform=args.beamform_mode, beamform=args.beamform_mode,
torch_num_threads=args.torch_num_threads,
) )
log.info("audio backend: %s", cleaner.backend) log.info("audio backend: %s", cleaner.backend)
+4
View File
@@ -9,3 +9,7 @@ def test_parse_cameras_list():
def test_default_beamform_mode_auto(): def test_default_beamform_mode_auto():
assert Config().beamform_mode == "auto" assert Config().beamform_mode == "auto"
def test_torch_num_threads_default_one():
assert Config().torch_num_threads == 1