fix: 限制语音 ONNX CUDA 显存 (#2191)

Co-authored-by: Codex <codex@local>

config/platform-infra/selfmedia-voice.yaml

@@ -128,6 +128,10 @@ runtime:
       max: 2.0
     cudaAllocator: expandable_segments:True,max_split_size_mb:128
     cudaModuleLoading: LAZY
+    onnxCuda:
+      gpuMemLimitMiB: 256
+      arenaExtendStrategy: kSameAsRequested
+      cudnnConvUseMaxWorkspace: false
   operations:
     applyTimeoutSeconds: 3600
   smoke:

deploy/selfmedia-voice/api.py

@@ -6,6 +6,7 @@ from contextlib import asynccontextmanager
 
 import torch
 import torchaudio
+import onnxruntime
 import uvicorn
 from cosyvoice.cli.cosyvoice import AutoModel
 from fastapi import FastAPI, HTTPException
@@ -28,6 +29,7 @@ class Runtime:
         self.model = None
         self.loaded_at = None
         self.onnx_providers = None
+        self.onnx_provider_options = None
         self.error = None
         self.lock = threading.Lock()
 
@@ -39,15 +41,36 @@ class Runtime:
             if not device.startswith("cuda:"):
                 raise RuntimeError("VOICE_DEVICE must select a CUDA device")
             torch.cuda.set_device(int(device.split(":", 1)[1]))
-            self.model = AutoModel(
-                model_dir=os.environ["MODEL_DIR"],
-                load_trt=env_bool("VOICE_LOAD_TRT"),
-                load_vllm=env_bool("VOICE_LOAD_VLLM"),
-                fp16=env_bool("VOICE_FP16"),
-            )
+            original_inference_session = onnxruntime.InferenceSession
+            cuda_options = onnx_cuda_provider_options()
+
+            def inference_session(*args, **kwargs):
+                if kwargs.get("providers") == ["CUDAExecutionProvider"]:
+                    kwargs["providers"] = [("CUDAExecutionProvider", cuda_options)]
+                return original_inference_session(*args, **kwargs)
+
+            onnxruntime.InferenceSession = inference_session
+            try:
+                self.model = AutoModel(
+                    model_dir=os.environ["MODEL_DIR"],
+                    load_trt=env_bool("VOICE_LOAD_TRT"),
+                    load_vllm=env_bool("VOICE_LOAD_VLLM"),
+                    fp16=env_bool("VOICE_FP16"),
+                )
+            finally:
+                onnxruntime.InferenceSession = original_inference_session
             self.onnx_providers = self.model.frontend.speech_tokenizer_session.get_providers()
-            if "CUDAExecutionProvider" not in self.onnx_providers:
+            self.onnx_provider_options = self.model.frontend.speech_tokenizer_session.get_provider_options()
+            if not self.onnx_providers or self.onnx_providers[0] != "CUDAExecutionProvider":
                 raise RuntimeError(f"CUDAExecutionProvider was not instantiated: {self.onnx_providers}")
+            actual_cuda_options = self.onnx_provider_options.get("CUDAExecutionProvider", {})
+            mismatches = {
+                key: {"expected": value, "actual": actual_cuda_options.get(key)}
+                for key, value in cuda_options.items()
+                if actual_cuda_options.get(key) != value
+            }
+            if mismatches:
+                raise RuntimeError(f"CUDAExecutionProvider options were not applied: {mismatches}")
             self.loaded_at = time.time()
         except Exception as error:
             self.error = f"{type(error).__name__}: {error}"
@@ -61,6 +84,17 @@ def env_bool(name: str) -> bool:
     return value == "true"
 
 
+def onnx_cuda_provider_options() -> dict[str, str]:
+    gpu_mem_limit_mib = int(os.environ["VOICE_ONNX_CUDA_GPU_MEM_LIMIT_MIB"])
+    if gpu_mem_limit_mib <= 0:
+        raise RuntimeError("VOICE_ONNX_CUDA_GPU_MEM_LIMIT_MIB must be positive")
+    return {
+        "gpu_mem_limit": str(gpu_mem_limit_mib * 1024 * 1024),
+        "arena_extend_strategy": os.environ["VOICE_ONNX_CUDA_ARENA_EXTEND_STRATEGY"],
+        "cudnn_conv_use_max_workspace": "1" if env_bool("VOICE_ONNX_CUDA_CUDNN_CONV_USE_MAX_WORKSPACE") else "0",
+    }
+
+
 runtime = Runtime()
 
 
@@ -91,6 +125,12 @@ def health() -> dict:
         "workers": int(os.environ["VOICE_WORKERS"]),
         "concurrency": os.environ["VOICE_CONCURRENCY"],
         "onnxProviders": runtime.onnx_providers,
+        "onnxProviderOptions": runtime.onnx_provider_options,
+        "onnxCudaMemory": {
+            "gpuMemLimitMiB": int(os.environ["VOICE_ONNX_CUDA_GPU_MEM_LIMIT_MIB"]),
+            "arenaExtendStrategy": os.environ["VOICE_ONNX_CUDA_ARENA_EXTEND_STRATEGY"],
+            "cudnnConvUseMaxWorkspace": env_bool("VOICE_ONNX_CUDA_CUDNN_CONV_USE_MAX_WORKSPACE"),
+        },
         "loadedAt": runtime.loaded_at,
         "error": runtime.error,
     }

deploy/selfmedia-voice/compose.d518.yaml

@@ -67,6 +67,9 @@ services:
       VOICE_PORT: ${VOICE_PORT}
       PYTORCH_CUDA_ALLOC_CONF: ${PYTORCH_CUDA_ALLOC_CONF}
       CUDA_MODULE_LOADING: ${CUDA_MODULE_LOADING}
+      VOICE_ONNX_CUDA_GPU_MEM_LIMIT_MIB: ${VOICE_ONNX_CUDA_GPU_MEM_LIMIT_MIB}
+      VOICE_ONNX_CUDA_ARENA_EXTEND_STRATEGY: ${VOICE_ONNX_CUDA_ARENA_EXTEND_STRATEGY}
+      VOICE_ONNX_CUDA_CUDNN_CONV_USE_MAX_WORKSPACE: ${VOICE_ONNX_CUDA_CUDNN_CONV_USE_MAX_WORKSPACE}
       LD_LIBRARY_PATH: ${TORCH_LIBRARY_PATH}
     volumes:
       - ${VOICE_CACHE_DIR}/model:/model:ro

scripts/src/platform-infra-selfmedia-voice.ts

@@ -68,6 +68,7 @@ interface VoiceConfig {
       speed: { min: number; max: number };
       cudaAllocator: string;
       cudaModuleLoading: string;
+      onnxCuda: { gpuMemLimitMiB: number; arenaExtendStrategy: "kSameAsRequested" | "kNextPowerOfTwo"; cudnnConvUseMaxWorkspace: boolean };
     };
     operations: { applyTimeoutSeconds: number };
     smoke: { healthPath: string; speechPath: string; text: string; voice: string; timeoutSeconds: number };
@@ -207,6 +208,10 @@ function readConfig(): VoiceConfig {
   if (config.runtime.inference.concurrency !== "serial") throw new Error(`${configLabel}.runtime.inference.concurrency must be serial`);
   if (config.runtime.inference.referenceVoice !== config.image.referenceVoice.id) throw new Error(`${configLabel}.runtime.inference.referenceVoice must select image.referenceVoice.id`);
   if (typeof config.runtime.inference.speed.min !== "number" || typeof config.runtime.inference.speed.max !== "number" || config.runtime.inference.speed.min >= config.runtime.inference.speed.max) throw new Error(`${configLabel}.runtime.inference.speed must declare an increasing numeric range`);
+  requireInteger(config.runtime.inference.onnxCuda.gpuMemLimitMiB, "runtime.inference.onnxCuda.gpuMemLimitMiB");
+  if (config.runtime.inference.onnxCuda.gpuMemLimitMiB <= 0) throw new Error(`${configLabel}.runtime.inference.onnxCuda.gpuMemLimitMiB must be positive`);
+  if (!["kSameAsRequested", "kNextPowerOfTwo"].includes(config.runtime.inference.onnxCuda.arenaExtendStrategy)) throw new Error(`${configLabel}.runtime.inference.onnxCuda.arenaExtendStrategy is unsupported`);
+  if (typeof config.runtime.inference.onnxCuda.cudnnConvUseMaxWorkspace !== "boolean") throw new Error(`${configLabel}.runtime.inference.onnxCuda.cudnnConvUseMaxWorkspace must be a boolean`);
   return config;
 }
 
@@ -383,6 +388,9 @@ function runtimeEnv(config: VoiceConfig, target: VoiceTarget): string {
     VOICE_SPEED_MAX: config.runtime.inference.speed.max,
     PYTORCH_CUDA_ALLOC_CONF: config.runtime.inference.cudaAllocator,
     CUDA_MODULE_LOADING: config.runtime.inference.cudaModuleLoading,
+    VOICE_ONNX_CUDA_GPU_MEM_LIMIT_MIB: config.runtime.inference.onnxCuda.gpuMemLimitMiB,
+    VOICE_ONNX_CUDA_ARENA_EXTEND_STRATEGY: config.runtime.inference.onnxCuda.arenaExtendStrategy,
+    VOICE_ONNX_CUDA_CUDNN_CONV_USE_MAX_WORKSPACE: String(config.runtime.inference.onnxCuda.cudnnConvUseMaxWorkspace),
     FRPC_IMAGE: config.frp.images.frpc,
   };
   return Object.entries(values).map(([key, value]) => `${key}=${value}`).join("\n") + "\n";