Skip to content

Commit d4c7ec7

Browse files
LC044claude
andcommitted
refactor(ai): 统一 ONNX 推理后端选择并新增 OpenVINO 支持
将分散在 face/embedding/image_classification/ticket 四个 service 的硬编码 provider 列表收敛到新增的 onnx_providers.get_onnx_providers(),按 CUDA > OpenVINO > CPU 优先级动态探测可用 EP,避免未安装 EP 报警告。 - 新增 openvino extra(onnxruntime-openvino + openvino 2025.4.1),与 cpu/gpu 互斥;uv.lock 同步更新 - _openvino_device() 按 openvino.Core().available_devices() 探测:NPU 优先、 否则 CPU,避免无 NPU 机器因 device_type=NPU 在会话创建时硬挂;OPENVINO_DEVICE 环境变量保留为强制覆盖口 - image_classification_service 补回 provider_options,与其他 service 一致 - get_onnx_providers() 加 lru_cache,进程内只探测与打日志一次 - ocr_service 注明 OCR 走 RapidOCR 自有引擎口,不接入 OpenVINO EP Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 1b4d2a5 commit d4c7ec7

8 files changed

Lines changed: 302 additions & 93 deletions

File tree

package/ai/app/services/embedding_service.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from app.services.model_downloader import model_downloader
1111
from app.services.model_manager import model_manager
1212
from app.services.ai_config_manager import ai_config_manager
13+
from app.services.onnx_providers import get_onnx_providers
1314

1415
class ONNXCLIPTextWrapper:
1516
def __init__(self, model_dir):
@@ -22,7 +23,7 @@ def __init__(self, model_dir):
2223

2324
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
2425

25-
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
26+
providers, _ = get_onnx_providers()
2627
text_model_path = os.path.join(model_dir, "textual.onnx")
2728
self.text_session = ort.InferenceSession(text_model_path, providers=providers)
2829

@@ -52,8 +53,8 @@ def __init__(self, model_dir):
5253

5354
self.processor = AutoImageProcessor.from_pretrained(model_dir)
5455

55-
# CUDA 优先:onnxruntime-gpu 在 CUDA 运行时库可用时走 GPU,否则回退 CPU。
56-
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
56+
# 推理后端按 CUDA -> OpenVINO -> CPU 自动选择
57+
providers, _ = get_onnx_providers()
5758
vision_model_path = os.path.join(model_dir, "visual.onnx")
5859
self.vision_session = ort.InferenceSession(vision_model_path, providers=providers)
5960

package/ai/app/services/face_service.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,23 +9,23 @@
99
from app.services.model_downloader import model_downloader
1010
from app.services.model_manager import model_manager
1111
from app.services.ai_config_manager import ai_config_manager
12+
from app.services.onnx_providers import get_onnx_providers
1213

1314
def load_insightface_model():
1415
try:
1516
import insightface
1617
from insightface.app import FaceAnalysis
17-
18+
1819
model_name = ai_config_manager.get_model_selection("face")
1920
logging.info(f"Initializing InsightFace with model: {model_name}")
2021

21-
# Initialize InsightFace analysis
22-
# providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] if GPU available
23-
provider_options = [{"device_id": 0}, {}]
22+
# 推理后端按 CUDA -> OpenVINO -> CPU 自动选择,与安装的 extra (gpu/openvino/cpu) 对齐。
23+
providers, provider_options = get_onnx_providers()
2424
model_path = settings.MODEL_PATH.rstrip("/").rstrip("models")
2525
app = FaceAnalysis(
2626
name=model_name, root=model_path,
27-
providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
28-
provider_options = provider_options # 传递 CUDA 配置
27+
providers=providers,
28+
provider_options=provider_options,
2929
)
3030
app.prepare(ctx_id=0, det_size=(640, 640))
3131
logging.info("InsightFace model initialized successfully.")

package/ai/app/services/image_classification_service.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,14 @@
1212
from app.config import settings
1313
from app.services.model_downloader import model_downloader
1414
from app.services.model_manager import model_manager
15+
from app.services.onnx_providers import get_onnx_providers
1516

1617
class ONNXModelWrapper:
1718
def __init__(self, model_path):
1819
import onnxruntime as ort
1920
import ast
20-
self.session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
21+
providers, provider_options = get_onnx_providers()
22+
self.session = ort.InferenceSession(model_path, providers=providers, provider_options=provider_options)
2123
meta = self.session.get_modelmeta()
2224
names_str = meta.custom_metadata_map.get('names', '{}')
2325
try:

package/ai/app/services/ocr_service.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ def load_paddleocr_model():
1515
# GPU 优先:ONNXRUNTIME 引擎开启 use_cuda,RapidOCR 会把 CUDAExecutionProvider
1616
# 插到 provider 列表首位;CUDA 运行时库不可用时自动回退 CPU(onnxruntime.get_device()
1717
# 非 GPU 即跳过,不会报错)。安装了 torch 时仍走下面的 TORCH 引擎 GPU 路径。
18+
# 注意:OCR 走 RapidOCR 自己的引擎配置口,不经过 onnx_providers.get_onnx_providers(),
19+
# 因此 openvino extra 不会让 OCR 走 OpenVINO EP(仅人脸 / 分类 / 票据 / embedding 受影响)。
1820
params = {
1921
"Det.engine_type": EngineType.ONNXRUNTIME,
2022
"Det.lang_type": LangDet.CH,
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import functools
2+
import logging
3+
import os
4+
5+
logger = logging.getLogger("app.services.onnx_providers")
6+
7+
# 推理后端优先级:CUDA (gpu extra) > OpenVINO (openvino extra) > CPU (cpu extra)。
8+
# 通过 ort.get_available_providers() 动态探测,避免硬编码导致未安装的 EP 报警告。
9+
_PRIORITY = ["CUDAExecutionProvider", "OpenVINOExecutionProvider", "CPUExecutionProvider"]
10+
11+
12+
def _openvino_device() -> str:
13+
"""选择 OpenVINO EP 的 device_type:NPU 优先,否则 CPU。
14+
15+
OPENVINO_DEVICE 环境变量可强制覆盖(设为 NPU / GPU / CPU);留空时按
16+
openvino.Core().available_devices() 探测——有 NPU 走 NPU,否则回退 CPU,
17+
避免在无 NPU 的机器上指定 NPU 导致会话创建失败。
18+
"""
19+
override = os.getenv("OPENVINO_DEVICE", "").strip()
20+
if override:
21+
return override
22+
try:
23+
from openvino import Core
24+
devices = set(Core().available_devices)
25+
if "NPU" in devices:
26+
return "NPU"
27+
except Exception as e:
28+
logger.warning(f"Failed to probe OpenVINO devices, fallback to CPU: {e}")
29+
return "CPU"
30+
31+
32+
@functools.lru_cache(maxsize=1)
33+
def get_onnx_providers():
34+
"""返回 (providers, provider_options),按可用性筛选并保持优先级。
35+
36+
结果在进程生命周期内缓存(可用 EP 不会变),避免每次加载模型都重复探测与打日志。
37+
38+
provider_options 与 providers 一一对齐:
39+
- CUDAExecutionProvider -> {"device_id": 0}
40+
- OpenVINOExecutionProvider -> {"device_type": "NPU" | "CPU"}(NPU 优先)
41+
- CPUExecutionProvider -> {}
42+
"""
43+
try:
44+
import onnxruntime as ort
45+
available = set(ort.get_available_providers())
46+
except Exception as e:
47+
logger.warning(f"Failed to probe onnxruntime providers, fallback to CPU: {e}")
48+
return ["CPUExecutionProvider"], [{}]
49+
50+
providers = [p for p in _PRIORITY if p in available]
51+
if not providers:
52+
providers = ["CPUExecutionProvider"]
53+
54+
options = []
55+
for p in providers:
56+
if p == "CUDAExecutionProvider":
57+
options.append({"device_id": 0})
58+
elif p == "OpenVINOExecutionProvider":
59+
options.append({"device_type": _openvino_device()})
60+
else:
61+
options.append({})
62+
63+
logger.info(f"ONNX Runtime providers selected: {providers} with options {options}")
64+
return providers, options

package/ai/app/services/ticket_service.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from app.services.ai_config_manager import ai_config_manager
1111
from app.services.ticket_parser import parse_ticket_info, extract_text
1212
from app.services.fly_ticket_parser import extract_flight_info
13+
from app.services.onnx_providers import get_onnx_providers
1314

1415
def load_modelscope_model():
1516
"""
@@ -36,7 +37,7 @@ def load_onnx_model():
3637
model_path = os.path.join(model_dir, "ticket-recognition.onnx")
3738

3839
# Initialize ONNX model
39-
providers = ['CPUExecutionProvider', 'CUDAExecutionProvider']
40+
providers, _ = get_onnx_providers()
4041
session = ort.InferenceSession(model_path, providers=providers)
4142

4243
# Set names attribute manually since ONNX model might not have it in the same format

package/ai/pyproject.toml

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,21 @@ dependencies = [
1515
"rapidocr>=3.0.0",
1616
"uvicorn>=0.27.0",
1717
"httpx>=0.27.0",
18+
"transformers>=4.37.0"
1819
]
1920

2021
[project.optional-dependencies]
2122
cpu = [
22-
"onnxruntime>=1.23.2",
23-
"transformers>=4.37.0"
23+
"onnxruntime>=1.23.2"
24+
]
25+
26+
openvino = [
27+
"onnxruntime-openvino>=1.24.1",
28+
"openvino==2025.4.1"
2429
]
30+
2531
gpu = [
26-
"onnxruntime-gpu>=1.23.2",
27-
"transformers>=4.37.0"
32+
"onnxruntime-gpu>=1.23.2"
2833
]
2934

3035
[dependency-groups]
@@ -37,6 +42,7 @@ dev = [
3742
conflicts = [
3843
[
3944
{ extra = "cpu" },
45+
{ extra = "openvino" },
4046
{ extra = "gpu" },
4147
],
4248
]

0 commit comments

Comments
 (0)