diff --git a/config/platform-infra/selfmedia-voice.yaml b/config/platform-infra/selfmedia-voice.yaml index 0a1248ce..13cd4bb8 100644 --- a/config/platform-infra/selfmedia-voice.yaml +++ b/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: diff --git a/deploy/selfmedia-voice/api.py b/deploy/selfmedia-voice/api.py index 66a3f4f0..88ef4515 100644 --- a/deploy/selfmedia-voice/api.py +++ b/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, } diff --git a/deploy/selfmedia-voice/compose.d518.yaml b/deploy/selfmedia-voice/compose.d518.yaml index ad902ad0..4d3f8bac 100644 --- a/deploy/selfmedia-voice/compose.d518.yaml +++ b/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 diff --git a/scripts/src/platform-infra-selfmedia-voice.ts b/scripts/src/platform-infra-selfmedia-voice.ts index 3afc53af..e77c1dd5 100644 --- a/scripts/src/platform-infra-selfmedia-voice.ts +++ b/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";