882 lines
34 KiB
Python
882 lines
34 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import types
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
import imagefind.runtime as runtime_module
|
|
import pytest
|
|
from imagefind.config import Settings
|
|
from imagefind.runtime import AIDependencyManager, RuntimeToolManager
|
|
|
|
|
|
def _settings(tmp_path: Path, requirements: Path | None = None) -> Settings:
|
|
if requirements is not None:
|
|
requirements.mkdir(parents=True, exist_ok=True)
|
|
(requirements / "constraints-cp312.txt").write_text("# test constraints\n")
|
|
for component in ("visual", "ocr", "faces", "audio"):
|
|
(requirements / f"{component}.txt").touch(exist_ok=True)
|
|
settings = Settings(data_dir=tmp_path / "data", runtime_requirements_dir=requirements)
|
|
settings.prepare()
|
|
return settings
|
|
|
|
|
|
def test_ai_dependency_failure_preserves_previous_layer_and_redacts_proxy_password(tmp_path: Path, monkeypatch):
|
|
requirements = tmp_path / "requirements"
|
|
requirements.mkdir()
|
|
(requirements / "audio.txt").write_text("torch==test\n")
|
|
settings = _settings(tmp_path, requirements)
|
|
current = settings.ai_site_path
|
|
current.mkdir()
|
|
(current / "sentinel").write_text("previous")
|
|
(current / ".imagefind-runtime.json").write_text(
|
|
json.dumps({"lock_version": "old", "components": ["visual"]})
|
|
)
|
|
manager = AIDependencyManager(
|
|
settings,
|
|
lambda **_: {
|
|
"enabled": True,
|
|
"url": "http://proxy.invalid:8080",
|
|
"username": "imagefind",
|
|
"password": "top/secret?",
|
|
},
|
|
)
|
|
|
|
def failed_install(*_args, **_kwargs):
|
|
return subprocess.CompletedProcess([], 1, "", "proxy authentication failed: top%2Fsecret%3F")
|
|
|
|
monkeypatch.setattr(runtime_module.subprocess, "run", failed_install)
|
|
with pytest.raises(RuntimeError) as error:
|
|
manager.ensure("audio")
|
|
|
|
assert "top/secret?" not in str(error.value)
|
|
assert "top%2Fsecret%3F" not in str(error.value)
|
|
assert "***" in str(error.value)
|
|
assert (current / "sentinel").read_text() == "previous"
|
|
assert manager.status()["visual"]["state"] == "missing"
|
|
assert manager.status()["audio"]["state"] == "error"
|
|
assert "***" in manager.status()["audio"]["error"]
|
|
|
|
|
|
def test_ai_runtime_uses_configured_pip_sources_and_obeys_proxy_switch(tmp_path: Path, monkeypatch):
|
|
requirements = tmp_path / "requirements"
|
|
requirements.mkdir()
|
|
(requirements / "visual.txt").write_text("torch==test\n")
|
|
settings = _settings(tmp_path, requirements)
|
|
settings.pip_index_url = "https://pypi.tuna.example/simple"
|
|
settings.pytorch_index_url = "https://torch.example/cpu"
|
|
proxy = {
|
|
"enabled": False,
|
|
"url": "http://proxy.example:8080",
|
|
"username": "",
|
|
"password": "",
|
|
}
|
|
manager = AIDependencyManager(settings, lambda **_: proxy)
|
|
environments: list[dict[str, str]] = []
|
|
|
|
def successful_install(command, **kwargs):
|
|
environments.append(kwargs["env"])
|
|
return subprocess.CompletedProcess(command, 0, "", "")
|
|
|
|
monkeypatch.setenv("HTTPS_PROXY", "http://inherited-proxy.invalid:9999")
|
|
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
|
|
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
|
|
manager.ensure("visual")
|
|
|
|
assert environments[0]["PIP_INDEX_URL"] == "https://pypi.tuna.example/simple"
|
|
assert environments[0]["PIP_EXTRA_INDEX_URL"] == "https://torch.example/cpu"
|
|
assert "HTTPS_PROXY" not in environments[0]
|
|
assert environments[0]["PIP_DEFAULT_TIMEOUT"] == "60"
|
|
assert environments[0]["PIP_RETRIES"] == "3"
|
|
|
|
|
|
def test_ai_runtime_injects_enabled_proxy_without_logging_password(tmp_path: Path, monkeypatch):
|
|
requirements = tmp_path / "requirements"
|
|
requirements.mkdir()
|
|
(requirements / "audio.txt").write_text("torch==test\n")
|
|
settings = _settings(tmp_path, requirements)
|
|
manager = AIDependencyManager(
|
|
settings,
|
|
lambda **_: {
|
|
"enabled": True,
|
|
"url": "http://proxy.example:8080",
|
|
"username": "runtime-user",
|
|
"password": "runtime/password",
|
|
},
|
|
)
|
|
environments: list[dict[str, str]] = []
|
|
|
|
def successful_install(command, **kwargs):
|
|
environments.append(kwargs["env"])
|
|
return subprocess.CompletedProcess(command, 0, "", "")
|
|
|
|
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
|
|
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
|
|
manager.ensure("audio")
|
|
|
|
assert environments[0]["HTTPS_PROXY"] == (
|
|
"http://runtime-user:runtime%2Fpassword@proxy.example:8080"
|
|
)
|
|
|
|
|
|
def test_ai_runtime_stamp_binds_python_abi_and_lock_digests(tmp_path: Path, monkeypatch):
|
|
requirements = tmp_path / "requirements"
|
|
settings = _settings(tmp_path, requirements)
|
|
(requirements / "visual.txt").write_text("openvino==test\n")
|
|
manager = AIDependencyManager(settings)
|
|
|
|
monkeypatch.setattr(
|
|
runtime_module.subprocess,
|
|
"run",
|
|
lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "", ""),
|
|
)
|
|
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
|
|
manager.ensure("visual")
|
|
|
|
stamp = json.loads(manager.stamp_path.read_text())
|
|
record = stamp["components"]["visual"]
|
|
assert stamp["schema"] == runtime_module.AI_RUNTIME_SCHEMA
|
|
assert record["python_abi"] == manager._python_abi()
|
|
assert len(record["requirements_sha256"]) == 64
|
|
assert len(record["constraints_sha256"]) == 64
|
|
assert manager.status()["visual"]["state"] == "ready"
|
|
|
|
(requirements / "constraints-cp312.txt").write_text("# comment-only release change\n")
|
|
assert manager.status()["visual"]["state"] == "ready"
|
|
|
|
(requirements / "constraints-cp312.txt").write_text("openvino==changed\n")
|
|
assert manager.status()["visual"]["state"] == "missing"
|
|
|
|
|
|
def test_schema_two_component_stamps_from_0319_all_remain_ready(tmp_path: Path, monkeypatch):
|
|
requirements = tmp_path / "requirements"
|
|
settings = _settings(tmp_path, requirements)
|
|
project_requirements = Path(__file__).resolve().parents[1] / "requirements" / "runtime-ai"
|
|
current_constraints = (project_requirements / "constraints-cp312.txt").read_text()
|
|
_, remainder = current_constraints.split("\n", 1)
|
|
(requirements / "constraints-cp312.txt").write_text(
|
|
"# ImageFind 0.3.19 AI runtime lock for fnOS Python 3.12.\n" + remainder
|
|
)
|
|
for component in runtime_module.AI_IMPORTS:
|
|
(requirements / f"{component}.txt").write_text(
|
|
(project_requirements / f"{component}.txt").read_text()
|
|
)
|
|
manager = AIDependencyManager(settings)
|
|
constraints_digest = runtime_module._sha256(requirements / "constraints-cp312.txt")
|
|
records = {}
|
|
for component in runtime_module.AI_IMPORTS:
|
|
requirement_digest = runtime_module._sha256(requirements / f"{component}.txt")
|
|
combined = hashlib.sha256()
|
|
combined.update(manager._python_abi().encode())
|
|
combined.update(requirement_digest.encode())
|
|
combined.update(constraints_digest.encode())
|
|
for value in runtime_module.AI_NO_DEPENDENCIES.get(component, ()):
|
|
combined.update(value.encode())
|
|
records[component] = {
|
|
"python_abi": manager._python_abi(),
|
|
"requirements_sha256": requirement_digest,
|
|
"constraints_sha256": constraints_digest,
|
|
"digest": combined.hexdigest(),
|
|
"runtime_version": "0.3.19",
|
|
}
|
|
manager.current.mkdir(parents=True, exist_ok=True)
|
|
manager.stamp_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"schema": 2,
|
|
"runtime_version": "0.3.19",
|
|
"components": records,
|
|
}
|
|
)
|
|
)
|
|
monkeypatch.setattr(
|
|
runtime_module.subprocess,
|
|
"run",
|
|
lambda *_args, **_kwargs: pytest.fail("compatible runtime must not invoke pip"),
|
|
)
|
|
|
|
assert {name: value["state"] for name, value in manager.status().items()} == {
|
|
component: "ready" for component in runtime_module.AI_IMPORTS
|
|
}
|
|
for component in runtime_module.AI_IMPORTS:
|
|
manager.ensure(component)
|
|
|
|
|
|
def test_vad_constraint_upgrade_only_invalidates_audio_runtime(tmp_path: Path):
|
|
requirements = tmp_path / "requirements"
|
|
settings = _settings(tmp_path, requirements)
|
|
project_requirements = Path(__file__).resolve().parents[1] / "requirements" / "runtime-ai"
|
|
constraints = runtime_module._requirements_lines(project_requirements / "constraints-cp312.txt")
|
|
old_constraints = tuple(line for line in constraints if line != "webrtcvad-wheels==2.0.14")
|
|
(requirements / "constraints-cp312.txt").write_text("\n".join(constraints) + "\n")
|
|
for component in runtime_module.AI_IMPORTS:
|
|
(requirements / f"{component}.txt").write_text(
|
|
(project_requirements / f"{component}.txt").read_text()
|
|
)
|
|
manager = AIDependencyManager(settings)
|
|
prior_constraints_digest = hashlib.sha256(("\n".join(old_constraints) + "\n").encode()).hexdigest()
|
|
records = {}
|
|
for component in runtime_module.AI_IMPORTS:
|
|
component_lines = runtime_module._requirements_lines(requirements / f"{component}.txt")
|
|
if component == "audio":
|
|
component_lines = tuple(line for line in component_lines if line != "webrtcvad-wheels==2.0.14")
|
|
requirement_digest = hashlib.sha256(("\n".join(component_lines) + "\n").encode()).hexdigest()
|
|
combined = hashlib.sha256()
|
|
combined.update(manager._python_abi().encode())
|
|
combined.update(requirement_digest.encode())
|
|
combined.update(prior_constraints_digest.encode())
|
|
for value in runtime_module.AI_NO_DEPENDENCIES.get(component, ()):
|
|
combined.update(value.encode())
|
|
records[component] = {
|
|
"digest_format": runtime_module.AI_RUNTIME_DIGEST_FORMAT,
|
|
"python_abi": manager._python_abi(),
|
|
"requirements_sha256": requirement_digest,
|
|
"constraints_sha256": prior_constraints_digest,
|
|
"digest": combined.hexdigest(),
|
|
}
|
|
manager.current.mkdir(parents=True, exist_ok=True)
|
|
manager.stamp_path.write_text(json.dumps({"schema": 3, "components": records}))
|
|
|
|
states = {name: value["state"] for name, value in manager.status().items()}
|
|
|
|
assert states == {"visual": "ready", "ocr": "ready", "faces": "ready", "audio": "missing"}
|
|
|
|
|
|
def test_runtime_tools_prefer_system_binary(tmp_path: Path, monkeypatch):
|
|
settings = _settings(tmp_path)
|
|
manager = RuntimeToolManager(settings)
|
|
monkeypatch.setattr(runtime_module.shutil, "which", lambda name: "/system/bin/ffmpeg" if name == "ffmpeg" else None)
|
|
monkeypatch.setattr(manager, "_works", lambda path, name: path == "/system/bin/ffmpeg" and name == "ffmpeg")
|
|
monkeypatch.setattr(manager, "_install_release", lambda _name: pytest.fail("private fallback should not download"))
|
|
|
|
assert manager.ffmpeg() == "/system/bin/ffmpeg"
|
|
|
|
|
|
def test_bundled_ai_runtime_is_ready_and_never_invokes_pip(tmp_path: Path, monkeypatch):
|
|
settings = _settings(tmp_path)
|
|
settings.bundled_ai_runtime = True
|
|
manager = AIDependencyManager(settings)
|
|
imported: list[str] = []
|
|
monkeypatch.setattr(runtime_module.importlib, "import_module", lambda name: imported.append(name))
|
|
monkeypatch.setattr(runtime_module.subprocess, "run", lambda *_args, **_kwargs: pytest.fail("pip must not run"))
|
|
|
|
manager.ensure("audio")
|
|
|
|
assert imported == list(runtime_module.AI_IMPORTS["audio"])
|
|
assert manager.status()["audio"]["source"] == "bundled"
|
|
assert manager.environment_status()["components"] == list(runtime_module.AI_IMPORTS)
|
|
|
|
|
|
def test_runtime_tools_use_bundled_fallback_without_downloading(tmp_path: Path, monkeypatch):
|
|
settings = _settings(tmp_path)
|
|
settings.bundled_tools_dir = tmp_path / "app-bin"
|
|
settings.bundled_tools_dir.mkdir()
|
|
bundled = settings.bundled_tools_dir / "ffmpeg"
|
|
bundled.write_bytes(b"bundled")
|
|
manager = RuntimeToolManager(settings)
|
|
monkeypatch.setattr(runtime_module.shutil, "which", lambda _name: None)
|
|
monkeypatch.setattr(manager, "_works", lambda path, name: path == str(bundled) and name == "ffmpeg")
|
|
monkeypatch.setattr(manager, "_install_release", lambda _name: pytest.fail("fallback should not download"))
|
|
|
|
assert manager.ffmpeg() == str(bundled)
|
|
|
|
|
|
def test_ocr_runtime_uses_headless_opencv_without_desktop_dependency(tmp_path: Path, monkeypatch):
|
|
requirements = tmp_path / "requirements"
|
|
requirements.mkdir()
|
|
(requirements / "ocr.txt").write_text("opencv-python-headless==4.13.0.92\n")
|
|
settings = _settings(tmp_path, requirements)
|
|
manager = AIDependencyManager(settings)
|
|
commands: list[list[str]] = []
|
|
|
|
def successful_install(command, **_kwargs):
|
|
commands.append(command)
|
|
return subprocess.CompletedProcess(command, 0, "", "")
|
|
|
|
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
|
|
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
|
|
manager.ensure("ocr")
|
|
|
|
assert "--requirement" in commands[0]
|
|
assert "opencv-python-headless==4.13.0.92" in (requirements / "ocr.txt").read_text()
|
|
assert "opencv-python==" not in (requirements / "ocr.txt").read_text()
|
|
assert "--no-deps" in commands[1]
|
|
assert commands[1][-1] == "rapidocr-onnxruntime==1.4.4"
|
|
|
|
|
|
def test_installing_one_missing_component_preserves_and_reuses_compatible_runtime(
|
|
tmp_path: Path, monkeypatch
|
|
):
|
|
requirements = tmp_path / "requirements"
|
|
settings = _settings(tmp_path, requirements)
|
|
project_requirements = Path(__file__).resolve().parents[1] / "requirements" / "runtime-ai"
|
|
for component in runtime_module.AI_IMPORTS:
|
|
(requirements / f"{component}.txt").write_text(
|
|
(project_requirements / f"{component}.txt").read_text()
|
|
)
|
|
manager = AIDependencyManager(settings)
|
|
commands: list[list[str]] = []
|
|
|
|
def successful_install(command, **_kwargs):
|
|
commands.append(command)
|
|
return subprocess.CompletedProcess(command, 0, "", "")
|
|
|
|
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
|
|
monkeypatch.setattr(manager, "_validate", lambda *_args: None)
|
|
manager.ensure("visual")
|
|
(manager.current / "preserved-package").write_text("keep")
|
|
|
|
assert manager.status()["visual"]["state"] == "ready"
|
|
assert manager.status()["audio"]["state"] == "missing"
|
|
assert manager.status()["ocr"]["state"] == "missing"
|
|
assert manager.status()["faces"]["state"] == "missing"
|
|
|
|
commands.clear()
|
|
manager.ensure("ocr")
|
|
|
|
assert len(commands) == 2
|
|
assert commands[0][-1] == str(requirements / "ocr.txt")
|
|
assert commands[1][-1] == "rapidocr-onnxruntime==1.4.4"
|
|
assert all(str(requirements / "visual.txt") not in command for command in commands)
|
|
assert (manager.current / "preserved-package").read_text() == "keep"
|
|
assert {name: value["state"] for name, value in manager.status().items()} == {
|
|
"visual": "ready",
|
|
"ocr": "ready",
|
|
"faces": "ready",
|
|
"audio": "missing",
|
|
}
|
|
|
|
monkeypatch.setattr(
|
|
runtime_module.subprocess,
|
|
"run",
|
|
lambda *_args, **_kwargs: pytest.fail("covered components must not invoke pip"),
|
|
)
|
|
manager.ensure("faces")
|
|
|
|
commands.clear()
|
|
monkeypatch.setattr(runtime_module.subprocess, "run", successful_install)
|
|
manager.ensure("audio")
|
|
assert len(commands) == 1
|
|
assert commands[0][-1] == str(requirements / "audio.txt")
|
|
|
|
|
|
def test_preserved_models_trigger_runtime_repair_and_accelerator_refresh(tmp_path: Path, monkeypatch):
|
|
from imagefind.models import ModelManager
|
|
|
|
requirements = tmp_path / "requirements"
|
|
requirements.mkdir()
|
|
settings = _settings(tmp_path, requirements)
|
|
|
|
class Accelerator:
|
|
def __init__(self):
|
|
self.refreshes = 0
|
|
|
|
def refresh(self):
|
|
self.refreshes += 1
|
|
|
|
accelerator = Accelerator()
|
|
embeddings = types.SimpleNamespace(accelerator=accelerator)
|
|
manager = ModelManager(
|
|
settings,
|
|
embeddings,
|
|
types.SimpleNamespace(),
|
|
types.SimpleNamespace(),
|
|
)
|
|
manager.component_versions = lambda: {
|
|
"visual": "visual-v1",
|
|
"ocr": "ocr-v1",
|
|
"faces": None,
|
|
"audio": None,
|
|
}
|
|
manager.runtime_dependencies.status = lambda: {
|
|
"visual": {"state": "missing"},
|
|
"ocr": {"state": "ready"},
|
|
"faces": {"state": "missing"},
|
|
"audio": {"state": "missing"},
|
|
}
|
|
|
|
assert accelerator.refreshes == 1
|
|
assert manager.missing_runtime_components() == ["visual"]
|
|
|
|
ensured = []
|
|
monkeypatch.setattr(
|
|
manager.runtime_dependencies,
|
|
"ensure",
|
|
lambda component, progress=None: ensured.append((component, progress)),
|
|
)
|
|
manager.ensure_runtime("visual")
|
|
|
|
assert ensured == [("visual", None)]
|
|
assert accelerator.refreshes == 2
|
|
|
|
|
|
def test_runtime_repair_resets_stale_component_fallback(tmp_path: Path, monkeypatch):
|
|
from imagefind.models import ModelManager
|
|
|
|
requirements = tmp_path / "requirements"
|
|
requirements.mkdir()
|
|
settings = _settings(tmp_path, requirements)
|
|
|
|
class Accelerator:
|
|
def __init__(self):
|
|
self.refreshes = 0
|
|
|
|
def refresh(self):
|
|
self.refreshes += 1
|
|
|
|
resets: list[str] = []
|
|
accelerator = Accelerator()
|
|
embeddings = types.SimpleNamespace(accelerator=accelerator)
|
|
ocr = types.SimpleNamespace(reset=lambda: resets.append("ocr"))
|
|
manager = ModelManager(settings, embeddings, ocr, types.SimpleNamespace())
|
|
monkeypatch.setattr(manager.runtime_dependencies, "ensure", lambda *_args, **_kwargs: None)
|
|
|
|
manager.ensure_runtime("ocr")
|
|
|
|
assert resets == ["ocr"]
|
|
assert accelerator.refreshes == 2
|
|
|
|
|
|
def test_unavailable_accelerator_makes_installed_component_repairable(tmp_path: Path):
|
|
from imagefind.models import ModelManager
|
|
|
|
settings = _settings(tmp_path)
|
|
accelerator = types.SimpleNamespace(
|
|
refresh=lambda: None,
|
|
status=lambda: {"components": {"visual": {"state": "unavailable"}}},
|
|
)
|
|
manager = ModelManager(
|
|
settings,
|
|
types.SimpleNamespace(accelerator=accelerator),
|
|
types.SimpleNamespace(),
|
|
types.SimpleNamespace(),
|
|
)
|
|
versions = {"visual": "visual-v1", "ocr": None, "faces": None, "audio": None}
|
|
runtime = {name: {"state": "ready"} for name in versions}
|
|
health = {name: {"state": "ready", "error": None} for name in versions}
|
|
|
|
assert manager.operational_components(versions, runtime, health)["visual"] is False
|
|
|
|
|
|
def test_visual_image_export_uses_clip_openvino_task_instead_of_generic_text_backend(
|
|
tmp_path: Path, monkeypatch
|
|
):
|
|
from imagefind.models import ModelManager
|
|
|
|
settings = _settings(tmp_path)
|
|
source = tmp_path / "clip-image"
|
|
source.mkdir()
|
|
(source / "modules.json").write_text(
|
|
'[{"idx":0,"path":"0_CLIPModel","type":"sentence_transformers.models.CLIPModel"}]'
|
|
)
|
|
transformer_source = source / "0_CLIPModel"
|
|
transformer_source.mkdir()
|
|
(transformer_source / "config.json").write_text("{}")
|
|
destination = tmp_path / "exported-image"
|
|
calls = {}
|
|
|
|
class ExportConfig:
|
|
pass
|
|
|
|
def main_export(**kwargs):
|
|
calls["export"] = kwargs
|
|
path = kwargs["output"]
|
|
path.mkdir(parents=True)
|
|
(path / "openvino_model.xml").write_text("<xml/>")
|
|
(path / "openvino_model.bin").write_bytes(b"model")
|
|
|
|
optimum = types.ModuleType("optimum")
|
|
optimum_exporters = types.ModuleType("optimum.exporters")
|
|
optimum_exporters_openvino = types.ModuleType("optimum.exporters.openvino")
|
|
optimum_exporters_openvino.main_export = main_export
|
|
optimum_intel = types.ModuleType("optimum.intel")
|
|
optimum_openvino = types.ModuleType("optimum.intel.openvino")
|
|
optimum_configuration = types.ModuleType("optimum.intel.openvino.configuration")
|
|
optimum_configuration.OVConfig = ExportConfig
|
|
for name, module in {
|
|
"optimum": optimum,
|
|
"optimum.exporters": optimum_exporters,
|
|
"optimum.exporters.openvino": optimum_exporters_openvino,
|
|
"optimum.intel": optimum_intel,
|
|
"optimum.intel.openvino": optimum_openvino,
|
|
"optimum.intel.openvino.configuration": optimum_configuration,
|
|
}.items():
|
|
monkeypatch.setitem(sys.modules, name, module)
|
|
|
|
class Accelerator:
|
|
@staticmethod
|
|
def refresh():
|
|
pass
|
|
|
|
@staticmethod
|
|
def ov_config(device):
|
|
assert device == "CPU"
|
|
return {"INFERENCE_NUM_THREADS": 2}
|
|
|
|
embeddings = types.SimpleNamespace(accelerator=Accelerator())
|
|
manager = ModelManager(settings, embeddings, types.SimpleNamespace(), types.SimpleNamespace())
|
|
manager._export_visual_image_openvino(source, destination)
|
|
|
|
assert calls["export"] == {
|
|
"model_name_or_path": str(transformer_source.resolve()),
|
|
"output": destination / "openvino",
|
|
"task": "zero-shot-image-classification",
|
|
"library_name": "transformers",
|
|
"local_files_only": True,
|
|
"ov_config": calls["export"]["ov_config"],
|
|
}
|
|
assert isinstance(calls["export"]["ov_config"], ExportConfig)
|
|
assert (destination / "modules.json").is_file()
|
|
assert (destination / "0_CLIPModel" / "config.json").is_file()
|
|
assert (destination / "openvino" / "openvino_model.xml").is_file()
|
|
|
|
|
|
def test_visual_text_export_uses_transformer_module_and_explicit_library(
|
|
tmp_path: Path, monkeypatch
|
|
):
|
|
from imagefind.models import ModelManager
|
|
|
|
settings = _settings(tmp_path)
|
|
source = tmp_path / "clip-text"
|
|
source.mkdir()
|
|
(source / "modules.json").write_text(
|
|
'[{"idx":0,"path":"0_Transformer","type":"sentence_transformers.models.Transformer"}]'
|
|
)
|
|
transformer_source = source / "0_Transformer"
|
|
transformer_source.mkdir()
|
|
(transformer_source / "config.json").write_text("{}")
|
|
destination = tmp_path / "exported-text"
|
|
calls = {}
|
|
|
|
class ExportConfig:
|
|
pass
|
|
|
|
def main_export(**kwargs):
|
|
calls["export"] = kwargs
|
|
path = kwargs["output"]
|
|
path.mkdir(parents=True)
|
|
(path / "openvino_model.xml").write_text("<xml/>")
|
|
(path / "openvino_model.bin").write_bytes(b"model")
|
|
|
|
optimum = types.ModuleType("optimum")
|
|
optimum_exporters = types.ModuleType("optimum.exporters")
|
|
optimum_exporters_openvino = types.ModuleType("optimum.exporters.openvino")
|
|
optimum_exporters_openvino.main_export = main_export
|
|
optimum_intel = types.ModuleType("optimum.intel")
|
|
optimum_openvino = types.ModuleType("optimum.intel.openvino")
|
|
optimum_configuration = types.ModuleType("optimum.intel.openvino.configuration")
|
|
optimum_configuration.OVConfig = ExportConfig
|
|
for name, module in {
|
|
"optimum": optimum,
|
|
"optimum.exporters": optimum_exporters,
|
|
"optimum.exporters.openvino": optimum_exporters_openvino,
|
|
"optimum.intel": optimum_intel,
|
|
"optimum.intel.openvino": optimum_openvino,
|
|
"optimum.intel.openvino.configuration": optimum_configuration,
|
|
}.items():
|
|
monkeypatch.setitem(sys.modules, name, module)
|
|
|
|
embeddings = types.SimpleNamespace(
|
|
accelerator=types.SimpleNamespace(refresh=lambda: None)
|
|
)
|
|
manager = ModelManager(settings, embeddings, types.SimpleNamespace(), types.SimpleNamespace())
|
|
manager._export_visual_text_openvino(source, destination)
|
|
|
|
assert calls["export"] == {
|
|
"model_name_or_path": str(transformer_source.resolve()),
|
|
"output": destination / "0_Transformer" / "openvino",
|
|
"task": "feature-extraction",
|
|
"library_name": "transformers",
|
|
"local_files_only": True,
|
|
"ov_config": calls["export"]["ov_config"],
|
|
}
|
|
assert isinstance(calls["export"]["ov_config"], ExportConfig)
|
|
assert (destination / "modules.json").is_file()
|
|
assert (destination / "0_Transformer" / "config.json").is_file()
|
|
assert (destination / "0_Transformer" / "openvino" / "openvino_model.xml").is_file()
|
|
|
|
|
|
def test_ai_and_model_sizes_are_cached_for_status_polling(tmp_path: Path, monkeypatch):
|
|
from imagefind.models import ModelManager
|
|
|
|
requirements = tmp_path / "requirements"
|
|
settings = _settings(tmp_path, requirements)
|
|
runtime = AIDependencyManager(settings)
|
|
runtime_calls = []
|
|
monkeypatch.setattr(
|
|
runtime_module,
|
|
"_directory_size",
|
|
lambda path: runtime_calls.append(path) or 123,
|
|
)
|
|
|
|
assert runtime._current_size() == 123
|
|
assert runtime._current_size() == 123
|
|
assert runtime_calls == [runtime.current]
|
|
|
|
embeddings = types.SimpleNamespace(
|
|
accelerator=types.SimpleNamespace(refresh=lambda: None)
|
|
)
|
|
manager = ModelManager(settings, embeddings, types.SimpleNamespace(), types.SimpleNamespace())
|
|
model_calls = []
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_directory_size",
|
|
lambda path: model_calls.append(path) or 456,
|
|
)
|
|
|
|
assert manager._component_sizes() == {
|
|
"visual": 456,
|
|
"ocr": 456,
|
|
"faces": 456,
|
|
"audio": 456,
|
|
}
|
|
assert manager._component_sizes()["visual"] == 456
|
|
assert len(model_calls) == 4
|
|
|
|
|
|
def test_preserved_model_files_queue_runtime_repair_after_upgrade(tmp_path: Path):
|
|
from imagefind.main import create_app
|
|
|
|
requirements = tmp_path / "requirements"
|
|
requirements.mkdir()
|
|
settings = _settings(tmp_path, requirements)
|
|
(settings.models_dir / "visual" / "image").mkdir(parents=True)
|
|
(settings.models_dir / "visual" / "text").mkdir(parents=True)
|
|
(settings.models_dir / "manifest.json").write_text(
|
|
json.dumps({"components": {"visual": {"version": "visual-v1"}}})
|
|
)
|
|
services = create_app(settings).state.services
|
|
|
|
assert services.models.missing_runtime_components() == ["visual"]
|
|
job_id = services.queue_missing_ai_runtime()
|
|
|
|
with services.db.read() as connection:
|
|
job = connection.execute(
|
|
"SELECT kind,payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)
|
|
).fetchone()
|
|
assert job["kind"] == "prepare_ai_runtime"
|
|
assert json.loads(job["payload_json"]) == {"components": ["visual"]}
|
|
assert job["dedupe_key"] == "prepare-ai-runtime:visual"
|
|
|
|
|
|
def test_private_ffmpeg_and_rclone_releases_coexist_and_failed_update_is_atomic(tmp_path: Path, monkeypatch):
|
|
settings = _settings(tmp_path)
|
|
manager = RuntimeToolManager(settings)
|
|
ffmpeg = b"ffmpeg-test"
|
|
ffprobe = b"ffprobe-test"
|
|
rclone = b"rclone-test"
|
|
ffmpeg_archive = tmp_path / "ffmpeg.tar.xz"
|
|
with tarfile.open(ffmpeg_archive, "w:xz") as archive:
|
|
for name, content in (("release/ffmpeg", ffmpeg), ("release/ffprobe", ffprobe)):
|
|
member = tarfile.TarInfo(name)
|
|
member.size = len(content)
|
|
archive.addfile(member, io.BytesIO(content))
|
|
rclone_archive = tmp_path / "rclone.zip"
|
|
with zipfile.ZipFile(rclone_archive, "w") as archive:
|
|
archive.writestr("release/rclone", rclone)
|
|
releases = {
|
|
"ffmpeg": {
|
|
"version": "test",
|
|
"url": str(ffmpeg_archive),
|
|
"files": {
|
|
"ffmpeg": hashlib.sha256(ffmpeg).hexdigest(),
|
|
"ffprobe": hashlib.sha256(ffprobe).hexdigest(),
|
|
},
|
|
},
|
|
"rclone": {
|
|
"version": "test",
|
|
"url": str(rclone_archive),
|
|
"files": {"rclone": hashlib.sha256(rclone).hexdigest()},
|
|
},
|
|
}
|
|
monkeypatch.setattr(runtime_module, "TOOL_RELEASES", releases)
|
|
|
|
def local_download(url: str, destination: Path):
|
|
destination.write_bytes(Path(url).read_bytes())
|
|
|
|
monkeypatch.setattr(manager, "_download", local_download)
|
|
manager._install_release("ffmpeg")
|
|
manager._install_release("rclone")
|
|
assert (manager.root / "ffmpeg").read_bytes() == ffmpeg
|
|
assert (manager.root / "ffprobe").read_bytes() == ffprobe
|
|
assert (manager.root / "rclone").read_bytes() == rclone
|
|
|
|
releases["rclone"]["files"]["rclone"] = "0" * 64
|
|
with pytest.raises(RuntimeError, match="SHA-256"):
|
|
manager._install_release("rclone")
|
|
assert (manager.root / "ffmpeg").read_bytes() == ffmpeg
|
|
assert (manager.root / "rclone").read_bytes() == rclone
|
|
|
|
|
|
def test_audio_export_declares_transformers_library_and_fp16_configuration(tmp_path: Path, monkeypatch):
|
|
calls: dict[str, object] = {}
|
|
|
|
class ExportConfig:
|
|
def __init__(self, *, dtype):
|
|
self.dtype = dtype
|
|
|
|
class Processor:
|
|
@classmethod
|
|
def from_pretrained(cls, source, *, local_files_only):
|
|
calls["processor_source"] = source
|
|
calls["processor_local"] = local_files_only
|
|
return cls()
|
|
|
|
def save_pretrained(self, destination):
|
|
calls["processor_destination"] = destination
|
|
_write_minimal_audio_processor_files(destination)
|
|
|
|
def main_export(**kwargs):
|
|
calls["export"] = kwargs
|
|
_write_minimal_audio_export(kwargs["output"])
|
|
|
|
modules = {
|
|
"optimum": types.ModuleType("optimum"),
|
|
"optimum.exporters": types.ModuleType("optimum.exporters"),
|
|
"optimum.exporters.openvino": types.ModuleType("optimum.exporters.openvino"),
|
|
"optimum.intel": types.ModuleType("optimum.intel"),
|
|
"optimum.intel.openvino": types.ModuleType("optimum.intel.openvino"),
|
|
"optimum.intel.openvino.configuration": types.ModuleType("optimum.intel.openvino.configuration"),
|
|
"transformers": types.ModuleType("transformers"),
|
|
}
|
|
modules["optimum.exporters.openvino"].main_export = main_export
|
|
modules["optimum.intel.openvino.configuration"].OVConfig = ExportConfig
|
|
modules["transformers"].AutoProcessor = Processor
|
|
for name, module in modules.items():
|
|
monkeypatch.setitem(sys.modules, name, module)
|
|
|
|
from imagefind.models import ModelManager
|
|
|
|
source = tmp_path / "source"
|
|
destination = tmp_path / "output"
|
|
ModelManager._export_audio_model(source, destination)
|
|
|
|
export = calls["export"]
|
|
assert export["library_name"] == "transformers"
|
|
assert export["task"] == "automatic-speech-recognition-with-past"
|
|
assert export["local_files_only"] is True
|
|
assert export["ov_config"].dtype == "fp16"
|
|
assert calls["processor_destination"] == destination
|
|
|
|
|
|
def _write_minimal_audio_export(root: Path, *, include_cache: bool = False, marker: str = "main") -> None:
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
for name in (
|
|
"openvino_encoder_model.xml",
|
|
"openvino_decoder_model.xml",
|
|
):
|
|
(root / name).write_text(f"<{marker}/>")
|
|
for name in (
|
|
"openvino_encoder_model.bin",
|
|
"openvino_decoder_model.bin",
|
|
):
|
|
(root / name).write_bytes(marker.encode())
|
|
if include_cache:
|
|
(root / "openvino_decoder_with_past_model.xml").write_text(f"<{marker}-cache/>")
|
|
(root / "openvino_decoder_with_past_model.bin").write_bytes(marker.encode())
|
|
|
|
|
|
def _write_minimal_audio_processor_files(root: Path) -> None:
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
for name in ("config.json", "preprocessor_config.json", "tokenizer_config.json", "tokenizer.json"):
|
|
(root / name).write_text("{}", encoding="utf-8")
|
|
|
|
|
|
def _install_audio_export_modules(monkeypatch, main_export, *, fallback_model=None):
|
|
class ExportConfig:
|
|
def __init__(self, *, dtype):
|
|
self.dtype = dtype
|
|
|
|
class Processor:
|
|
@classmethod
|
|
def from_pretrained(cls, source, *, local_files_only):
|
|
assert local_files_only is True
|
|
return cls()
|
|
|
|
def save_pretrained(self, destination):
|
|
_write_minimal_audio_processor_files(destination)
|
|
|
|
modules = {
|
|
"optimum": types.ModuleType("optimum"),
|
|
"optimum.exporters": types.ModuleType("optimum.exporters"),
|
|
"optimum.exporters.openvino": types.ModuleType("optimum.exporters.openvino"),
|
|
"optimum.intel": types.ModuleType("optimum.intel"),
|
|
"optimum.intel.openvino": types.ModuleType("optimum.intel.openvino"),
|
|
"optimum.intel.openvino.configuration": types.ModuleType("optimum.intel.openvino.configuration"),
|
|
"transformers": types.ModuleType("transformers"),
|
|
}
|
|
modules["optimum.exporters.openvino"].main_export = main_export
|
|
modules["optimum.intel.openvino.configuration"].OVConfig = ExportConfig
|
|
if fallback_model is not None:
|
|
modules["optimum.intel.openvino"].OVModelForSpeechSeq2Seq = fallback_model
|
|
modules["transformers"].AutoProcessor = Processor
|
|
for name, module in modules.items():
|
|
monkeypatch.setitem(sys.modules, name, module)
|
|
|
|
|
|
def test_audio_export_keeps_main_export_when_optional_cache_export_fails(
|
|
tmp_path: Path, monkeypatch
|
|
):
|
|
def main_export(**kwargs):
|
|
_write_minimal_audio_export(kwargs["output"], marker="main")
|
|
|
|
class FailingFallback:
|
|
@classmethod
|
|
def from_pretrained(cls, *_args, **_kwargs):
|
|
raise ValueError("cache export unavailable")
|
|
|
|
_install_audio_export_modules(monkeypatch, main_export, fallback_model=FailingFallback)
|
|
|
|
from imagefind.models import ModelManager
|
|
|
|
destination = tmp_path / "audio"
|
|
ModelManager._export_audio_model(tmp_path / "source", destination)
|
|
|
|
assert (destination / "openvino_encoder_model.xml").read_text() == "<main/>"
|
|
assert not (destination / "openvino_decoder_with_past_model.xml").exists()
|
|
assert (destination / "tokenizer_config.json").is_file()
|
|
|
|
|
|
def test_audio_export_replaces_main_export_only_when_cache_export_is_complete(
|
|
tmp_path: Path, monkeypatch
|
|
):
|
|
def main_export(**kwargs):
|
|
_write_minimal_audio_export(kwargs["output"], marker="main")
|
|
|
|
class SuccessfulFallback:
|
|
@classmethod
|
|
def from_pretrained(cls, *_args, **_kwargs):
|
|
return cls()
|
|
|
|
def save_pretrained(self, destination):
|
|
_write_minimal_audio_export(destination, include_cache=True, marker="cache")
|
|
|
|
_install_audio_export_modules(monkeypatch, main_export, fallback_model=SuccessfulFallback)
|
|
|
|
from imagefind.models import ModelManager
|
|
|
|
destination = tmp_path / "audio"
|
|
ModelManager._export_audio_model(tmp_path / "source", destination)
|
|
|
|
assert (destination / "openvino_encoder_model.xml").read_text() == "<cache/>"
|
|
assert (destination / "openvino_decoder_with_past_model.xml").read_text() == "<cache-cache/>"
|
|
assert (destination / "tokenizer_config.json").is_file()
|
|
|
|
|
|
def test_audio_export_main_export_failure_reports_root_cause(tmp_path: Path, monkeypatch):
|
|
def main_export(**_kwargs):
|
|
raise ValueError("library name could not be inferred")
|
|
|
|
_install_audio_export_modules(monkeypatch, main_export)
|
|
|
|
from imagefind.models import ModelManager
|
|
|
|
with pytest.raises(RuntimeError, match="音频模型转换阶段失败：ValueError.*library name"):
|
|
ModelManager._export_audio_model(tmp_path / "source", tmp_path / "audio")
|