dtoro
bbb8e4850c
Centralize execution provider selection in providers.py with auto-detection and graceful fallback. All ONNX sessions (embedder, detector, face processor, recognizer) now use the configured providers. - New VISION_EXECUTION_PROVIDERS env var: "auto" for GPU auto-detect, or explicit "CUDAExecutionProvider,CPUExecutionProvider" - Provider priority: CUDA > ROCm > OpenVINO > CPU (when set to "auto") - docker-compose.yml includes commented-out NVIDIA GPU deploy section - Supports onnxruntime-gpu as a drop-in replacement for onnxruntime Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
"""
|
|
ONNX Runtime execution provider resolution with GPU auto-detection.
|
|
|
|
Resolves configured execution providers against what's actually available
|
|
in the current ONNX Runtime build. Falls back to CPU if no GPU provider
|
|
is available. Logs the selected provider so users can confirm GPU is active.
|
|
"""
|
|
import logging
|
|
|
|
import onnxruntime as ort
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_resolved: list[str] | None = None
|
|
|
|
|
|
def get_providers(configured: list[str] | None = None) -> list[str]:
|
|
"""Return the best available execution providers.
|
|
|
|
1. If `configured` is provided, filter to only those that are
|
|
actually available in the current ORT build.
|
|
2. If none of the configured providers are available, fall back
|
|
to CPUExecutionProvider.
|
|
3. Auto-detect: if configured is ["auto"], probe for GPU providers.
|
|
|
|
Results are cached after first call.
|
|
"""
|
|
global _resolved
|
|
if _resolved is not None:
|
|
return _resolved
|
|
|
|
available = set(ort.get_available_providers())
|
|
logger.info("ONNX Runtime available providers: %s", sorted(available))
|
|
|
|
if configured is None or configured == ["CPUExecutionProvider"]:
|
|
_resolved = ["CPUExecutionProvider"]
|
|
return _resolved
|
|
|
|
if configured == ["auto"]:
|
|
# Auto-detect: prefer CUDA > ROCm > OpenVINO > CPU
|
|
priority = [
|
|
"CUDAExecutionProvider",
|
|
"ROCMExecutionProvider",
|
|
"OpenVINOExecutionProvider",
|
|
]
|
|
for p in priority:
|
|
if p in available:
|
|
_resolved = [p, "CPUExecutionProvider"]
|
|
logger.info("Auto-detected GPU provider: %s", p)
|
|
return _resolved
|
|
_resolved = ["CPUExecutionProvider"]
|
|
logger.info("No GPU provider detected, using CPU")
|
|
return _resolved
|
|
|
|
# Filter configured list to available providers.
|
|
resolved = [p for p in configured if p in available]
|
|
if not resolved:
|
|
logger.warning(
|
|
"None of the configured providers %s are available. "
|
|
"Falling back to CPU. Available: %s",
|
|
configured,
|
|
sorted(available),
|
|
)
|
|
resolved = ["CPUExecutionProvider"]
|
|
else:
|
|
# Always include CPU as fallback.
|
|
if "CPUExecutionProvider" not in resolved:
|
|
resolved.append("CPUExecutionProvider")
|
|
|
|
_resolved = resolved
|
|
logger.info("Using ONNX Runtime providers: %s", _resolved)
|
|
return _resolved
|
|
|
|
|
|
def create_session(
|
|
model_path: str,
|
|
opts: ort.SessionOptions | None = None,
|
|
configured_providers: list[str] | None = None,
|
|
) -> ort.InferenceSession:
|
|
"""Create an ONNX InferenceSession with the best available providers."""
|
|
providers = get_providers(configured_providers)
|
|
if opts is None:
|
|
opts = ort.SessionOptions()
|
|
opts.inter_op_num_threads = 2
|
|
opts.intra_op_num_threads = 2
|
|
return ort.InferenceSession(model_path, opts, providers=providers)
|