""" 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)