from __future__ import annotations import os import sys import types from pathlib import Path import imagefind.accelerator as accelerator_module import imagefind.ai as ai_module import imagefind.main as main_module import pytest from imagefind.accelerator import AcceleratorService from imagefind.ai import EmbeddingService, ModelUnavailable, _OpenVINOClipImageEncoder, _OpenVINOOrtSession from imagefind.config import Settings from imagefind.media import MediaService def _openvino(monkeypatch, devices: list[str]): class Core: available_devices = devices @staticmethod def get_property(device, _name): return {"CPU": "Intel CPU", "GPU.0": "Intel UHD Graphics"}.get(device, device) monkeypatch.setitem(sys.modules, "openvino", types.SimpleNamespace(Core=Core)) def _render_nodes(monkeypatch, nodes: list[Path], accessible: bool = True): original_glob = Path.glob original_access = os.access def fake_glob(path: Path, pattern: str): if path == Path("/dev/dri") and pattern == "renderD*": return iter(nodes) return original_glob(path, pattern) monkeypatch.setattr(Path, "glob", fake_glob) monkeypatch.setattr( accelerator_module.os, "access", lambda path, mode: accessible if str(path).startswith("/dev/dri/renderD") else original_access(path, mode), ) def test_accelerator_prefers_gpu_and_falls_back_per_component(monkeypatch, tmp_path: Path): _openvino(monkeypatch, ["CPU", "GPU.0"]) _render_nodes(monkeypatch, [Path("/dev/dri/renderD128")]) service = AcceleratorService(Settings(data_dir=tmp_path, ai_cpu_threads=3)) assert service.gpu_device == "GPU.0" assert service.device_for("visual") == "GPU.0" assert service.status()["device_names"]["GPU.0"] == "Intel UHD Graphics" service.mark_ready("visual", "GPU.0") assert service.fall_back("visual", RuntimeError("secret driver detail")) is True assert service.device_for("visual") == "CPU" assert service.device_for("ocr") == "GPU.0" assert service.fall_back("visual", RuntimeError("do not retry")) is False visual = service.status()["components"]["visual"] assert visual["state"] == "fallback" assert "secret driver detail" not in visual["fallback_reason"] assert visual["requested_device"] == "GPU.0" assert visual["actual_device"] == "CPU" assert visual["failure_stage"] == "inference" assert service.ov_config("CPU") == {"INFERENCE_NUM_THREADS": 3, "PERFORMANCE_HINT": "LATENCY"} gpu_config = service.ov_config("GPU.0") assert gpu_config["PERFORMANCE_HINT"] == "LATENCY" assert gpu_config["NUM_STREAMS"] == "1" assert gpu_config["CACHE_DIR"].endswith("openvino-cache") service.reset("visual") visual = service.status()["components"]["visual"] assert visual["state"] == "not_loaded" assert visual["device"] is None assert visual["requested_device"] == "GPU.0" assert visual["actual_device"] == "CPU" assert visual["execution_devices"] == ["CPU"] assert visual["last_verified_at"] is not None assert visual["failure_stage"] is None assert visual["fallback_reason"] is None assert service.device_for("visual") == "GPU.0" @pytest.mark.parametrize( ("mode", "hint", "streams", "batch_size"), [("low", "LATENCY", "1", 1), ("balanced", "THROUGHPUT", "2", 2), ("throughput", "THROUGHPUT", "3", 4)], ) def test_gpu_profiles_bound_streams_and_batches(monkeypatch, tmp_path: Path, mode, hint, streams, batch_size): _openvino(monkeypatch, ["CPU", "GPU.0"]) _render_nodes(monkeypatch, [Path("/dev/dri/renderD128")]) service = AcceleratorService(Settings(data_dir=tmp_path, ai_gpu_mode=mode)) config = service.ov_config("GPU.0") assert config["PERFORMANCE_HINT"] == hint assert config["NUM_STREAMS"] == streams assert service.gpu_profile()["batch_size"] == batch_size assert service.settings.ai_cpu_threads == 1 def test_accelerator_reset_without_prior_inference_remains_never_loaded(monkeypatch, tmp_path: Path): _openvino(monkeypatch, ["CPU", "GPU.0"]) _render_nodes(monkeypatch, [Path("/dev/dri/renderD128")]) service = AcceleratorService(Settings(data_dir=tmp_path)) service.reset("faces") faces = service.status()["components"]["faces"] assert faces["state"] == "not_loaded" assert faces["requested_device"] is None assert faces["actual_device"] is None assert faces["execution_devices"] == [] assert faces["last_verified_at"] is None def test_accelerator_cpu_verification_is_transient_but_can_become_fallback( monkeypatch, tmp_path: Path, ): _openvino(monkeypatch, ["CPU", "GPU.0"]) _render_nodes(monkeypatch, [Path("/dev/dri/renderD128")]) service = AcceleratorService(Settings(data_dir=tmp_path)) service.mark_ready("audio", "GPU.0", ["GPU.0"]) service.mark_verifying_cpu("audio") verifying = service.status()["components"]["audio"] assert verifying["state"] == "verifying_cpu" assert verifying["actual_device"] == "CPU" assert service.device_for("audio") == "GPU.0" assert service.fall_back("audio", "CPU sample contained speech", stage="empty_result") is True fallback = service.status()["components"]["audio"] assert fallback["state"] == "fallback" assert fallback["actual_device"] == "CPU" assert service.device_for("audio") == "CPU" def test_audio_transient_device_failures_only_open_circuit_after_three_tasks( monkeypatch, tmp_path: Path, ): _openvino(monkeypatch, ["CPU", "GPU.0"]) _render_nodes(monkeypatch, [Path("/dev/dri/renderD128")]) service = AcceleratorService(Settings(data_dir=tmp_path)) assert service.record_transient_failure("audio", "stall one", stage="inference_stall") is False assert service.device_for("audio") == "GPU.0" assert service.record_transient_failure("audio", "stall two", stage="inference_stall") is False assert service.device_for("audio") == "GPU.0" assert service.record_transient_failure("audio", "stall three", stage="inference_stall") is True status = service.status()["components"]["audio"] assert status["state"] == "fallback" assert status["fallback_scope"] == "component" assert status["circuit_state"] == "open" assert status["failure_count"] == 3 assert status["retry_at"] assert service.device_for("audio") == "CPU" service.mark_inference_success("audio", "GPU.0", ["GPU.0"]) status = service.status()["components"]["audio"] assert status["circuit_state"] == "closed" assert status["failure_count"] == 0 def test_accelerator_reports_render_permission_problem(monkeypatch, tmp_path: Path): _openvino(monkeypatch, ["CPU"]) _render_nodes(monkeypatch, [Path("/dev/dri/renderD128")], accessible=False) service = AcceleratorService(Settings(data_dir=tmp_path)) assert service.gpu_device is None assert service.status()["render_device"] == {"available": True, "accessible": False, "count": 1} assert service.status()["unavailable_reason"] == "Intel render 设备权限不足" def test_gpu_sysfs_metrics_need_two_samples_and_do_not_invent_shared_vram( monkeypatch, tmp_path: Path, ): _openvino(monkeypatch, ["CPU", "GPU.0"]) _render_nodes(monkeypatch, [Path("/dev/dri/renderD128")]) service = AcceleratorService(Settings(data_dir=tmp_path)) busy_path = Path("/sys/class/drm/card0/engine/rcs0/busy") original_glob = Path.glob values = iter((1_000_000_000, 1_500_000_000)) def fake_glob(path: Path, pattern: str): if path == Path("/sys/class/drm") and pattern == "card*/engine/*/busy": return iter((busy_path,)) if path == Path("/sys/class/drm") and pattern == "card*/device": return iter(()) return original_glob(path, pattern) original_read_text = Path.read_text monkeypatch.setattr(Path, "glob", fake_glob) monkeypatch.setattr( Path, "read_text", lambda path, **kwargs: str(next(values)) if path == busy_path else original_read_text(path, **kwargs), ) clock = iter((10.0, 11.0)) monkeypatch.setattr(accelerator_module.time, "monotonic", lambda: next(clock)) first = service.hardware_metrics() second = service.hardware_metrics() assert first["supported"] is True assert first["utilization_percent"] is None assert second["utilization_percent"] == 50.0 assert second["engines"] == {"card0/rcs0": 50.0} assert second["memory_supported"] is False assert second["memory_total_bytes"] is None def test_visual_models_refuse_runtime_export_and_require_persisted_ir(tmp_path: Path): settings = Settings(data_dir=tmp_path) for name in ("image", "text"): root = settings.models_dir / "visual" / name root.mkdir(parents=True) (root / "modules.json").write_text('[{"idx":0,"path":"","type":"Transformer"}]') class Accelerator: @staticmethod def ov_config(_device): return {"PERFORMANCE_HINT": "THROUGHPUT"} @staticmethod def mark_ready(_component, _device): pass calls = [] class Model: def __init__(self, path, **kwargs): calls.append((path, kwargs)) service = EmbeddingService(settings, Accelerator()) with pytest.raises(ModelUnavailable, match="OpenVINO 模型不完整"): service._build_models(Model, "openvino", "GPU.0") assert calls == [] @pytest.mark.parametrize("text_module", ["", "0_Transformer"]) def test_visual_models_reuse_existing_openvino_exports( tmp_path: Path, monkeypatch, text_module: str ): settings = Settings(data_dir=tmp_path) for name in ("image", "text"): root = settings.models_dir / "visual" / name module = root / text_module if name == "text" and text_module else root export = module / "openvino" export.mkdir(parents=True) module_type = "Transformer" if name == "text" else "CLIPModel" module_path = text_module if name == "text" else "" (root / "modules.json").write_text( f'[' f'{{"idx":0,"path":"{module_path}","type":"{module_type}"}}' f']' ) (module / "config.json").write_text("{}") (export / "openvino_model.xml").write_text("") (export / "openvino_model.bin").write_bytes(b"model") (export / "config.json").write_text("{}") class Accelerator: @staticmethod def ov_config(_device): return {"PERFORMANCE_HINT": "THROUGHPUT"} @staticmethod def execution_devices_from(*_models): return ["GPU.0"] @staticmethod def mark_ready(_component, _device, _execution_devices): pass text_calls = [] image_calls = [] class ImageModel: def __init__(self, path, device, ov_config): image_calls.append((path, device, ov_config)) class Model: def __init__(self, path, **kwargs): text_calls.append((path, kwargs)) monkeypatch.setattr(ai_module, "_OpenVINOClipImageEncoder", ImageModel) EmbeddingService(settings, Accelerator())._build_models(Model, "openvino", "GPU.0") assert len(image_calls) == 1 assert Path(image_calls[0][0]).is_absolute() assert image_calls[0][1:] == ("GPU", {"PERFORMANCE_HINT": "THROUGHPUT"}) assert len(text_calls) == 1 assert Path(text_calls[0][0]).is_absolute() assert text_calls[0][1]["local_files_only"] is True assert text_calls[0][1]["model_kwargs"] == { "device": "GPU", "ov_config": {"PERFORMANCE_HINT": "THROUGHPUT"}, "export": False, "file_name": "openvino_model.xml", "subfolder": "openvino", } def test_clip_image_encoder_uses_image_aware_openvino_model(monkeypatch, tmp_path: Path): root = tmp_path / "image" export = root / "openvino" export.mkdir(parents=True) processor_root = root / "0_CLIPModel" processor_root.mkdir() (root / "modules.json").write_text( '[{"idx":0,"path":"0_CLIPModel","type":"sentence_transformers.models.CLIPModel"}]' ) (processor_root / "config.json").write_text("{}") (export / "openvino_model.xml").write_text("") (export / "openvino_model.bin").write_bytes(b"model") (export / "config.json").write_text("{}") calls = {} class Matrix: def __init__(self, rows): self.rows = rows def __truediv__(self, denominators): return Matrix( [[value / denominators[index][0] for value in row] for index, row in enumerate(self.rows)] ) def tolist(self): return self.rows numpy = types.ModuleType("numpy") numpy.float32 = "float32" numpy.asarray = lambda value, dtype=None: Matrix(value) numpy.maximum = lambda values, _minimum: values numpy.finfo = lambda _dtype: types.SimpleNamespace(eps=1e-7) numpy.linalg = types.SimpleNamespace( norm=lambda values, axis, keepdims: [ [sum(item * item for item in row) ** 0.5] for row in values.rows ] ) class OpenVINOModel: @classmethod def from_pretrained(cls, _path, **_kwargs): raise AssertionError("public loader must not perform Hugging Face library inference") @classmethod def _from_pretrained(cls, path, **kwargs): calls["model"] = (path, kwargs) return cls() def __call__(self, **inputs): calls["inputs"] = inputs return types.SimpleNamespace(image_embeds=[[3.0, 4.0]]) class Processor: @classmethod def from_pretrained(cls, path, **kwargs): calls["processor"] = (path, kwargs) return cls() def __call__(self, **kwargs): calls["processor_inputs"] = kwargs return {"input_ids": [[1]], "pixel_values": [[0.0]]} class Config: @classmethod def from_pretrained(cls, path, **kwargs): calls["config"] = (path, kwargs) return cls() optimum = types.ModuleType("optimum") optimum_intel = types.ModuleType("optimum.intel") optimum_openvino = types.ModuleType("optimum.intel.openvino") optimum_openvino.OVModelForZeroShotImageClassification = OpenVINOModel transformers = types.ModuleType("transformers") transformers.AutoConfig = Config transformers.AutoProcessor = Processor for name, module in { "optimum": optimum, "optimum.intel": optimum_intel, "optimum.intel.openvino": optimum_openvino, "transformers": transformers, "numpy": numpy, }.items(): monkeypatch.setitem(sys.modules, name, module) encoder = _OpenVINOClipImageEncoder(root, "GPU.0", {"NUM_STREAMS": "1"}) vectors = encoder.encode([object()]) assert calls["model"] == ( str(export.resolve()), { "config": calls["model"][1]["config"], "file_name": "openvino_model.xml", "local_files_only": True, "device": "GPU.0", "ov_config": {"NUM_STREAMS": "1"}, }, ) assert isinstance(calls["model"][1]["config"], Config) assert calls["config"] == (str(processor_root.resolve()), {"local_files_only": True}) assert calls["processor"] == (str(processor_root.resolve()), {"local_files_only": True}) assert calls["processor_inputs"]["text"] == [""] assert set(calls["inputs"]) == {"input_ids", "pixel_values"} assert [round(value, 6) for value in vectors.tolist()[0]] == [0.6, 0.8] def test_visual_text_encoder_bypasses_optimum_hub_inference( tmp_path: Path, monkeypatch ): settings = Settings(data_dir=tmp_path) for name in ("image", "text"): root = settings.models_dir / "visual" / name export = root / "openvino" export.mkdir(parents=True) module_type = "CLIPModel" if name == "image" else "Transformer" (root / "modules.json").write_text( f'[' f'{{"idx":0,"path":"","type":"{module_type}"}}' f']' ) (root / "config.json").write_text("{}") (export / "openvino_model.xml").write_text("") (export / "openvino_model.bin").write_bytes(b"model") (export / "config.json").write_text("{}") calls = {} class OpenVINOTextModel: @classmethod def from_pretrained(cls, _path, **_kwargs): raise AssertionError("public Optimum loader must not infer a Hub library") @classmethod def _from_pretrained(cls, **kwargs): calls["local"] = kwargs return object() class SentenceTransformerModel: __module__ = "sentence_transformers.sentence_transformer.model" def __init__(self, path, **kwargs): calls["wrapper"] = (path, kwargs) options = kwargs["model_kwargs"] self.model = OpenVINOTextModel.from_pretrained( path, config="local-config", export=options["export"], file_name=options["file_name"], subfolder=options["subfolder"], local_files_only=kwargs["local_files_only"], device=options["device"], ov_config=options["ov_config"], ) optimum = types.ModuleType("optimum") optimum_intel = types.ModuleType("optimum.intel") optimum_openvino = types.ModuleType("optimum.intel.openvino") optimum_openvino.OVModelForFeatureExtraction = OpenVINOTextModel for name, module in { "optimum": optimum, "optimum.intel": optimum_intel, "optimum.intel.openvino": optimum_openvino, }.items(): monkeypatch.setitem(sys.modules, name, module) class Accelerator: @staticmethod def ov_config(_device): return {"NUM_STREAMS": "1"} @staticmethod def execution_devices_from(*_models): return ["GPU"] @staticmethod def mark_ready(*_args): pass monkeypatch.setattr(ai_module, "_OpenVINOClipImageEncoder", lambda *_args: object()) EmbeddingService(settings, Accelerator())._build_models( SentenceTransformerModel, "openvino", "GPU" ) assert calls["local"] == { "model_id": str((settings.models_dir / "visual" / "text").resolve()), "config": "local-config", "file_name": "openvino_model.xml", "subfolder": "openvino", "local_files_only": True, "device": "GPU", "ov_config": {"NUM_STREAMS": "1"}, } with pytest.raises(AssertionError, match="public Optimum loader"): OpenVINOTextModel.from_pretrained("/vol1/imagefind/models/visual/text") def test_hugging_face_validation_errors_are_not_gpu_device_errors(): error_type = type("HFValidationError", (ValueError,), {"__module__": "huggingface_hub.errors"}) assert AcceleratorService.is_device_error(error_type("invalid repo id")) is False assert AcceleratorService.is_device_error(ModelUnavailable("OpenVINO IR 尚未准备完成")) is False assert AcceleratorService.is_device_error(RuntimeError("OpenVINO GPU compile_model failed")) is True def test_execution_devices_are_read_from_compiled_wrappers(): class Compiled: @staticmethod def get_property(name): assert name == "EXECUTION_DEVICES" return ["GPU.0"] class Request: @staticmethod def get_compiled_model(): return Compiled() wrapper = types.SimpleNamespace(encoder=types.SimpleNamespace(request=Request())) assert AcceleratorService.execution_devices_from(wrapper) == ["GPU.0"] def test_execution_device_discovery_ignores_unsupported_gpu_properties(): class UnsupportedProperties: @staticmethod def get_property(_name): raise ValueError("property is not supported by this GPU wrapper") @staticmethod def get_compiled_model(): raise ValueError("compiled model is exposed by a sibling request") class Compiled: @staticmethod def get_property(name): assert name == "EXECUTION_DEVICES" return ["GPU.0"] wrapper = types.SimpleNamespace( model=UnsupportedProperties(), request=Compiled(), ) assert AcceleratorService.execution_devices_from(wrapper) == ["GPU.0"] def test_openvino_value_errors_retry_on_cpu_but_model_layout_errors_do_not(): assert AcceleratorService.is_device_error(ValueError("GPU plugin property failed")) is True assert AcceleratorService.is_device_error(ValueError("modules.json is invalid")) is False def test_visual_gpu_value_error_retries_complete_model_on_cpu(tmp_path: Path, monkeypatch): settings = Settings(data_dir=tmp_path) for name in ("image", "text"): root = settings.models_dir / "visual" / name export = root / "openvino" export.mkdir(parents=True) (root / "modules.json").write_text('[{"idx":0,"path":"","type":"Transformer"}]') (root / "config.json").write_text("{}") (export / "openvino_model.xml").write_text("") (export / "openvino_model.bin").write_bytes(b"model") (export / "config.json").write_text("{}") calls = [] unavailable = [] class Accelerator: def __init__(self): self.fallback = False def device_for(self, _component): return "CPU" if self.fallback else "GPU.0" @staticmethod def is_device_error(error): return isinstance(error, ValueError) def fall_back(self, _component, _error, stage): assert stage == "image_encoder_compile" self.fallback = True return True @staticmethod def ov_config(device): return {"device_config": device} @staticmethod def execution_devices_from(*_models): return ["CPU"] @staticmethod def mark_ready(component, device, execution_devices): calls.append((component, device, execution_devices)) @staticmethod def mark_unavailable(component, error, stage): unavailable.append((component, error, stage)) class ImageModel: def __init__(self, _path, device, _ov_config): if device == "GPU": raise ValueError("Intel GPU property rejected by wrapper") class TextModel: def __init__(self, _path, **kwargs): assert kwargs["model_kwargs"]["device"] == "CPU" sentence_transformers = types.ModuleType("sentence_transformers") sentence_transformers.SentenceTransformer = TextModel monkeypatch.setitem(sys.modules, "sentence_transformers", sentence_transformers) monkeypatch.setattr(ai_module, "_OpenVINOClipImageEncoder", ImageModel) service = EmbeddingService(settings, Accelerator()) service._load() assert service._image_model is not None assert service._text_model is not None assert calls == [("visual", "CPU", ["CPU"])] assert unavailable == [] def test_openvino_ocr_adapter_compiles_requested_device(monkeypatch): captured = {} class Metadata: value = "a\nb" class Model: @staticmethod def get_rt_info(): return {"framework": {"character": Metadata()}} class Compiled: @staticmethod def output(_index): return "output" def __call__(self, values): captured["values"] = values return {"output": "result"} class Core: @staticmethod def read_model(path): captured["model"] = path return Model() @staticmethod def compile_model(model, device, config): captured.update(compiled_model=model, device=device, config=config) return Compiled() monkeypatch.setitem(sys.modules, "openvino", types.SimpleNamespace(Core=Core)) session = _OpenVINOOrtSession("ocr.onnx", "GPU.0", {"PERFORMANCE_HINT": "THROUGHPUT"}) assert session("pixels") == ["result"] assert session.have_key() is True assert session.get_character_list() == ["a", "b"] assert captured["model"] == "ocr.onnx" assert isinstance(captured["compiled_model"], Model) assert captured["device"] == "GPU.0" assert captured["config"] == {"PERFORMANCE_HINT": "THROUGHPUT"} assert captured["values"] == ["pixels"] def test_remote_media_uses_runtime_internal_port(tmp_path: Path): settings = Settings(data_dir=tmp_path, port=8765, internal_media_port=43123) media = MediaService(settings).input_for({"id": "remote-video", "source_kind": "webdav"}) assert media.value == "http://127.0.0.1:43123/api/internal/remote/remote-video" assert media.headers and media.headers.startswith("X-ImageFind-Internal: ") assert media.headers.endswith("\r\n") def test_gateway_only_serve_creates_private_random_tcp_listener(monkeypatch, tmp_path: Path): gateway_socket = tmp_path / "imagefind.sock" settings = Settings( data_dir=tmp_path / "data", gateway_socket=gateway_socket, direct_access=False, host="0.0.0.0", port=8765, internal_media_port=0, ) captured = {} class Listener: def __init__(self, name): self.name = name def getsockname(self): return self.name def close(self): pass class Config: def __init__(self, app, **kwargs): captured.update(app=app, config=kwargs) class Server: def __init__(self, _config): pass @staticmethod def run(sockets): captured["listeners"] = [listener.getsockname() for listener in sockets] monkeypatch.setattr(main_module, "create_app", lambda current: ("app", current.internal_media_port)) monkeypatch.setattr( main_module, "_tcp_socket", lambda host, port: captured.setdefault("tcp_request", (host, port)) and Listener((host, 43123)), ) monkeypatch.setattr(main_module, "_unix_socket", lambda path, _mode: Listener(str(path))) monkeypatch.setitem(sys.modules, "uvicorn", types.SimpleNamespace(Config=Config, Server=Server)) main_module.serve(settings) assert captured["tcp_request"] == ("127.0.0.1", 0) assert settings.internal_media_port == 43123 assert captured["app"] == ("app", settings.internal_media_port) assert captured["listeners"][0] == ("127.0.0.1", settings.internal_media_port) assert captured["listeners"][1] == str(gateway_socket) assert all(listener != ("0.0.0.0", 8765) for listener in captured["listeners"]) assert not gateway_socket.exists()