feat: Control which GPU the model selects by the CUDA_VISIBLE_DEVICES environment variable

Author: yxf0314

vox_box/backends/stt/faster_whisper.py

@@ -28,6 +28,7 @@ def __init__(
         preprocessor_config_path = os.path.join(
             self._cfg.model, "preprocessor_config.json"
         )
+
         if os.path.exists(preprocessor_config_path):
             with open(preprocessor_config_path, "r", encoding="utf-8") as f:
                 self._preprocessor_config_json = json.load(f)
@@ -40,22 +41,17 @@ def load(self):
         if self._cfg.device == "cpu":
             cpu_threads = 8
 
+        device = self._cfg.device
+        if device.startswith("cuda:"):
+            device = device.split(":")[0]
+
         compute_type = "default"
         if platform.system() == "Darwin":
             compute_type = "int8"
 
-        device = self._cfg.device
-        device_index = 0
-        if self._cfg.device != "cpu":
-            arr = device.split(":")
-            device = arr[0]
-            if len(arr) > 1:
-                device_index = int(arr[1])
-
         self._model = WhisperModel(
             self._cfg.model,
             device=device,
-            device_index=device_index,
             cpu_threads=cpu_threads,
             compute_type=compute_type,
         )

vox_box/cmd/start.py

@@ -7,6 +7,7 @@
 from vox_box.config import Config
 from vox_box.server.model import ModelInstance
 from vox_box.server.server import Server
+from vox_box.utils.model import parse_and_set_cuda_visible_devices
 
 
 logger = logging.getLogger(__name__)
@@ -91,6 +92,7 @@ def run(args: argparse.Namespace):
     try:
         cfg = parse_args(args)
         setup_logging(cfg.debug)
+        parse_and_set_cuda_visible_devices(cfg)
 
         logger.info("Starting with arguments: %s", args._get_kwargs())
 

vox_box/utils/model.py

@@ -1,6 +1,12 @@
+import logging
+import os
 import time
 from typing import Dict
 
+from vox_box.config import Config
+
+logger = logging.getLogger(__name__)
+
 
 def create_model_dict(id: str, **kwargs) -> Dict:
     d = {
@@ -16,3 +22,16 @@ def create_model_dict(id: str, **kwargs) -> Dict:
             d[k] = v
 
     return d
+
+
+def parse_and_set_cuda_visible_devices(cfg: Config):
+    """
+    Parse CUDA device in format cuda:1 and set CUDA_VISIBLE_DEVICES accordingly.
+    """
+    if cfg.device.startswith("cuda:"):
+        device_index = cfg.device.split(":")[1]
+        if device_index.isdigit():
+            os.environ["CUDA_VISIBLE_DEVICES"] = device_index
+            logger.info(f"Set CUDA_VISIBLE_DEVICES = {device_index}")
+        else:
+            raise ValueError(f"Invalid CUDA device index: {device_index}")