Feature/specfied CPU Threads (VOICEVOX#286)

* Add: argument bridge to make_synthesis_engine

* Add: cpu_num_threads to initializer args

* Fix: cpu-num-threads value assertion

* Fix: Passing args regardless of specify

* Update README.md

* Fix: wrong japanese

* Apply: pysen lint

* Fix: too long line

* Fix: error message with stderr

Co-authored-by: takana-v <[email protected]>

* Fix: bridge cpu_num_threads run.py -> make_synthesis_engine

* Enhanced: correspondence Cancellable_engine

Co-authored-by: takana-v <[email protected]>

README.md

@@ -240,6 +240,23 @@ python run.py --voicevox_dir=$VOICEVOX_DIR --voicelib_dir=$VOICELIB_DIR
 python run.py
 ```
 
+
+### CPUスレッド数を指定する
+CPUスレッド数が未指定の場合は、論理コア数の半分か物理コア数が使われます。（殆どのCPUで、これは全体の処理能力の半分です）  
+もしIaaS上で実行していたり、専用サーバーで実行している場合など、  
+VOICEVOX ENGINEが使う処理能力を調節したい場合は、CPUスレッド数を指定することで実現できます。
+
+- 実行時引数で指定する
+    ```bash
+    python run.py --voicevox_dir=$VOICEVOX_DIR --cpu_num_threads=4
+    ```
+
+- 環境変数で指定する
+    ```bash
+    CPU_NUM_THREADS=4
+    python run.py --voicevox_dir=$VOICEVOX_DIR
+    ```
+
 ## コードフォーマット
 
 このソフトウェアでは、リモートにプッシュする前にコードフォーマットを確認する仕組み(静的解析ツール)を利用できます。

run.py

@@ -3,6 +3,7 @@
 import base64
 import json
 import multiprocessing
+import os
 import zipfile
 from functools import lru_cache
 from pathlib import Path
@@ -494,8 +495,16 @@ def speaker_info(speaker_uuid: str):
     parser.add_argument("--voicelib_dir", type=Path, default=None)
     parser.add_argument("--enable_cancellable_synthesis", action="store_true")
     parser.add_argument("--init_processes", type=int, default=2)
+    parser.add_argument("--cpu_num_threads", type=int, default=None)
     args = parser.parse_args()
 
+    cpu_num_threads: Optional[int] = args.cpu_num_threads
+
+    # 引数へcpu_num_threadsの指定がなければ、環境変数をロールします。
+    # 環境変数にもない場合は、Noneのままとします。
+    if cpu_num_threads is None:
+        cpu_num_threads = os.getenv("CPU_NUM_THREADS", None)
+
     # voicelib_dir が Noneのとき、音声ライブラリの Python モジュールと同じディレクトリにあるとする
     voicelib_dir: Optional[Path] = args.voicelib_dir
     if voicelib_dir is None:
@@ -506,14 +515,15 @@ def speaker_info(speaker_uuid: str):
 
     cancellable_engine = None
     if args.enable_cancellable_synthesis:
-        cancellable_engine = CancellableEngine(args, voicelib_dir)
+        cancellable_engine = CancellableEngine(args, voicelib_dir, cpu_num_threads)
 
     uvicorn.run(
         generate_app(
             make_synthesis_engine(
                 use_gpu=args.use_gpu,
                 voicelib_dir=voicelib_dir,
                 voicevox_dir=args.voicevox_dir,
+                cpu_num_threads=cpu_num_threads,
             )
         ),
         host=args.host,

voicevox_engine/cancellable_engine.py

@@ -33,13 +33,14 @@ class CancellableEngine:
         （音声合成中のプロセスは入っていない）
     """
 
-    def __init__(self, args, voicelib_dir) -> None:
+    def __init__(self, args, voicelib_dir, cpu_num_threads: Optional[int]) -> None:
         """
         変数の初期化を行う
         また、args.init_processesの数だけプロセスを起動し、procs_and_consに格納する
         """
         self.args = args
         self.voicelib_dir = voicelib_dir
+        self.cpu_num_threads = cpu_num_threads
         if not self.args.enable_cancellable_synthesis:
             raise HTTPException(
                 status_code=404,
@@ -71,6 +72,7 @@ def start_new_proc(
                 "args": self.args,
                 "voicelib_dir": self.voicelib_dir,
                 "sub_proc_con": sub_proc_con2,
+                "cpu_num_threads": self.cpu_num_threads,
             },
             daemon=True,
         )
@@ -170,7 +172,10 @@ async def catch_disconnection(self):
 
 
 def start_synthesis_subprocess(
-    args: argparse.Namespace, voicelib_dir: Path, sub_proc_con: Connection
+    args: argparse.Namespace,
+    voicelib_dir: Path,
+    sub_proc_con: Connection,
+    cpu_num_threads: Optional[int],
 ):
     """
     音声合成を行うサブプロセスで行うための関数
@@ -188,6 +193,7 @@ def start_synthesis_subprocess(
         use_gpu=args.use_gpu,
         voicevox_dir=args.voicevox_dir,
         voicelib_dir=voicelib_dir,
+        cpu_num_threads=cpu_num_threads,
     )
     while True:
         try:

voicevox_engine/synthesis_engine/make_synthesis_engine.py

@@ -9,6 +9,7 @@ def make_synthesis_engine(
     use_gpu: bool,
     voicelib_dir: Path,
     voicevox_dir: Optional[Path] = None,
+    cpu_num_threads: Optional[int] = None,
 ) -> SynthesisEngineBase:
     """
     音声ライブラリをロードして、音声合成エンジンを生成
@@ -22,6 +23,9 @@ def make_synthesis_engine(
     voicevox_dir: Path, optional, default=None
         音声ライブラリの Python モジュールがあるディレクトリ
         None のとき、Python 標準のモジュール検索パスのどれかにあるとする
+    cpu_num_threads: int, optional, default=None
+        音声ライブラリが、推論に用いるCPUスレッド数を設定する
+        Noneのとき、ライブラリ側の挙動により論理コア数の半分か、物理コア数が指定される
     """
 
     # Python モジュール検索パスへ追加
@@ -47,7 +51,14 @@ def make_synthesis_engine(
             file=sys.stderr,
         )
 
-    core.initialize(voicelib_dir.as_posix() + "/", use_gpu)
+    if cpu_num_threads == 0:
+        print(
+            "Warning: cpu_num_threads is set to 0. "
+            + "( The library leaves the decision to the synthesis runtime )",
+            file=sys.stderr,
+        )
+
+    core.initialize(voicelib_dir.as_posix() + "/", use_gpu, cpu_num_threads or 0)
 
     if has_voicevox_core:
         return SynthesisEngine(