diff --git a/machine-learning/app/models/base.py b/machine-learning/app/models/base.py index 6909a935c..ad48624b4 100644 --- a/machine-learning/app/models/base.py +++ b/machine-learning/app/models/base.py @@ -165,6 +165,14 @@ class InferenceModel(ABC): def providers_default(self) -> list[str]: available_providers = set(ort.get_available_providers()) log.debug(f"Available ORT providers: {available_providers}") + if (openvino := "OpenVINOExecutionProvider") in available_providers: + device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids() + log.debug(f"Available OpenVINO devices: {device_ids}") + + gpu_devices = [device_id for device_id in device_ids if device_id.startswith("GPU")] + if not gpu_devices: + log.warning("No GPU device found in OpenVINO. Falling back to CPU.") + available_providers.remove(openvino) return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers] @property @@ -184,15 +192,7 @@ class InferenceModel(ABC): case "CPUExecutionProvider" | "CUDAExecutionProvider": option = {"arena_extend_strategy": "kSameAsRequested"} case "OpenVINOExecutionProvider": - try: - device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids() - log.debug(f"Available OpenVINO devices: {device_ids}") - gpu_devices = [device_id for device_id in device_ids if device_id.startswith("GPU")] - option = {"device_id": gpu_devices[0]} if gpu_devices else {} - except AttributeError as e: - log.warning("Failed to get OpenVINO device IDs. Using default options.") - log.error(e) - option = {} + option = {"device_type": "GPU_FP32"} case _: option = {} options.append(option) diff --git a/machine-learning/app/test_main.py b/machine-learning/app/test_main.py index e25099e67..0f802997f 100644 --- a/machine-learning/app/test_main.py +++ b/machine-learning/app/test_main.py @@ -45,11 +45,23 @@ class TestBase: assert encoder.providers == self.CUDA_EP @pytest.mark.providers(OV_EP) - def test_sets_openvino_provider_if_available(self, providers: list[str]) -> None: + def test_sets_openvino_provider_if_available(self, providers: list[str], mocker: MockerFixture) -> None: + mocked = mocker.patch("app.models.base.ort.capi._pybind_state") + mocked.get_available_openvino_device_ids.return_value = ["GPU.0", "CPU"] + encoder = OpenCLIPEncoder("ViT-B-32__openai") assert encoder.providers == self.OV_EP + @pytest.mark.providers(OV_EP) + def test_avoids_openvino_if_gpu_not_available(self, providers: list[str], mocker: MockerFixture) -> None: + mocked = mocker.patch("app.models.base.ort.capi._pybind_state") + mocked.get_available_openvino_device_ids.return_value = ["CPU"] + + encoder = OpenCLIPEncoder("ViT-B-32__openai") + + assert encoder.providers == self.CPU_EP + @pytest.mark.providers(CUDA_EP_OUT_OF_ORDER) def test_sets_providers_in_correct_order(self, providers: list[str]) -> None: encoder = OpenCLIPEncoder("ViT-B-32__openai") @@ -68,22 +80,14 @@ class TestBase: assert encoder.providers == providers - def test_sets_default_provider_options(self) -> None: - encoder = OpenCLIPEncoder("ViT-B-32__openai", providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"]) - - assert encoder.provider_options == [ - {}, - {"arena_extend_strategy": "kSameAsRequested"}, - ] - - def test_sets_openvino_device_id_if_possible(self, mocker: MockerFixture) -> None: + def test_sets_default_provider_options(self, mocker: MockerFixture) -> None: mocked = mocker.patch("app.models.base.ort.capi._pybind_state") mocked.get_available_openvino_device_ids.return_value = ["GPU.0", "CPU"] encoder = OpenCLIPEncoder("ViT-B-32__openai", providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"]) assert encoder.provider_options == [ - {"device_id": "GPU.0"}, + {"device_type": "GPU_FP32"}, {"arena_extend_strategy": "kSameAsRequested"}, ]