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