Files
imagefind/tests/test_models_install.py

536 lines
22 KiB
Python

from __future__ import annotations
import asyncio
import hashlib
import io
import json
import sys
import tarfile
import types
from pathlib import Path
import httpx
import pytest
from imagefind.config import Settings
from imagefind.main import create_app
from imagefind.models import ModelManager
from imagefind.speech import SpeechService
def _bundle(*, unsafe: bool = False) -> bytes:
output = io.BytesIO()
with tarfile.open(fileobj=output, mode="w:gz") as archive:
files = {
"manifest.json": json.dumps({"version": "manual-test"}).encode(),
"visual/image/modules.json": json.dumps([{"idx": 0, "path": "", "type": "Transformer"}]).encode(),
"visual/text/modules.json": json.dumps([{"idx": 0, "path": "", "type": "Transformer"}]).encode(),
}
if unsafe:
files["../outside"] = b"unsafe"
for name, content in files.items():
member = tarfile.TarInfo(name)
member.size = len(content)
archive.addfile(member, io.BytesIO(content))
return output.getvalue()
def _app(tmp_path: Path):
settings = Settings(
data_dir=tmp_path / "data",
embedding_backend="hash",
upload_reserve_gb=0,
model_upload_gb=1,
)
settings.prepare()
app = create_app(settings)
_, token = app.state.services.auth.create_api_token("models")
return app, {"Authorization": f"Bearer {token}"}
def _write_audio_variant(root: Path) -> None:
root.mkdir(parents=True, exist_ok=True)
for name in (
"config.json",
"preprocessor_config.json",
"tokenizer_config.json",
"tokenizer.json",
"openvino_encoder_model.xml",
"openvino_encoder_model.bin",
"openvino_decoder_model.xml",
"openvino_decoder_model.bin",
):
(root / name).write_text("{}", encoding="utf-8")
def test_manual_model_bundle_upload_and_atomic_install(tmp_path: Path):
app, headers = _app(tmp_path)
bundle = _bundle()
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/upload",
headers={
**headers,
"Content-Type": "application/octet-stream",
"X-Model-Filename": "imagefind-models.tar.gz",
},
content=bundle,
)
assert response.status_code == 202
payload = response.json()
assert payload["size_bytes"] == len(bundle)
assert payload["sha256"] == hashlib.sha256(bundle).hexdigest()
with app.state.services.db.read() as conn:
row = conn.execute("SELECT payload_json FROM jobs WHERE id=?", (payload["job_id"],)).fetchone()
job_payload = json.loads(row["payload_json"])
archive_path = Path(job_payload["archive_path"])
result = app.state.services.models.install(
archive_path=archive_path,
expected_sha256=job_payload["sha256"],
)
assert result["manifest"]["version"] == "manual-test"
assert (app.state.services.settings.models_dir / "visual" / "image" / "modules.json").is_file()
assert (app.state.services.settings.models_dir / "visual" / "text" / "modules.json").is_file()
asyncio.run(scenario())
def test_medium_audio_variant_uses_isolated_verification(tmp_path: Path, monkeypatch):
app, _headers = _app(tmp_path)
service = app.state.services
_write_audio_variant(service.settings.models_dir / "audio" / "small")
_write_audio_variant(service.settings.models_dir / "audio" / "medium")
service.settings.audio_model_variant = "small"
service.speech.accelerator.gpu_device = "GPU.0"
called: list[str] = []
def forbidden_main_process_verify():
raise AssertionError("Medium verification must not run in the API process")
def isolated_verify(variant):
called.append(f"{variant}:{service.settings.audio_model_variant}")
service.speech.accelerator.mark_ready("audio", "GPU.0", ["GPU.0"])
return service.speech.accelerator.status()["components"]["audio"]
monkeypatch.setattr(service.speech, "verify_acceleration", forbidden_main_process_verify)
monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)
selected = service.models.set_audio_variant("medium")
assert selected["active"] is True
assert called == ["medium:small"]
assert service.settings.audio_model_variant == "medium"
def test_medium_audio_variant_verification_failure_rolls_back(tmp_path: Path, monkeypatch):
app, _headers = _app(tmp_path)
service = app.state.services
_write_audio_variant(service.settings.models_dir / "audio" / "small")
_write_audio_variant(service.settings.models_dir / "audio" / "medium")
service.settings.audio_model_variant = "small"
service.speech.accelerator.gpu_device = "GPU.0"
def isolated_verify(_variant):
raise RuntimeError("worker crashed")
monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)
with pytest.raises(RuntimeError, match="worker crashed"):
service.models.set_audio_variant("medium")
assert service.settings.audio_model_variant == "small"
def test_speech_config_medium_failure_returns_conflict_and_keeps_small(tmp_path: Path, monkeypatch):
app, headers = _app(tmp_path)
service = app.state.services
_write_audio_variant(service.settings.models_dir / "audio" / "small")
_write_audio_variant(service.settings.models_dir / "audio" / "medium")
service.settings.audio_model_variant = "small"
service.speech.accelerator.gpu_device = "GPU.0"
def isolated_verify(_variant):
raise RuntimeError("worker crashed")
monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.patch(
"/api/v1/speech/config",
headers=headers,
json={"model_variant": "medium"},
)
assert response.status_code == 409
config = await client.get("/api/v1/speech/config", headers=headers)
assert config.json()["model_variant"] == "small"
asyncio.run(scenario())
def test_old_medium_fp16_variant_requires_int8_update(tmp_path: Path):
settings = Settings(data_dir=tmp_path / "data")
settings.prepare()
medium = settings.models_dir / "audio" / "medium"
_write_audio_variant(medium)
manifest = {
"components": {
"audio": {
"variants": {
"medium": {
"version": "OpenVINO/whisper-medium-fp16-ov",
"source": "https://hf-mirror.com/OpenVINO/whisper-medium-fp16-ov",
}
},
"active_variant": "small",
}
}
}
(settings.models_dir / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
manager = ModelManager.__new__(ModelManager)
manager.settings = settings
manager.speech = SpeechService(settings)
manager._manifest_cached_mtime_ns = -1
manager._manifest_cached = {}
variant = manager.audio_variants()["medium"]
assert variant["installed"] is False
assert variant["needs_update"] is True
assert variant["expected_version"] == "OpenVINO/whisper-medium-int8-ov"
assert variant["size_bytes"] > 0
def test_model_status_defaults_to_light_manifest_and_allows_full_manifest(tmp_path: Path):
app, headers = _app(tmp_path)
app.state.services.settings.models_dir.mkdir(parents=True, exist_ok=True)
(app.state.services.settings.models_dir / "manifest.json").write_text(
json.dumps({"version": "test", "components": {}, "files": {"audio/model.bin": "hash"}}),
encoding="utf-8",
)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
light = await client.get("/api/v1/models", headers=headers)
assert light.status_code == 200
assert "files" not in light.json()["manifest"]
full = await client.get("/api/v1/models?summary=false", headers=headers)
assert full.status_code == 200
assert full.json()["manifest"]["files"] == {"audio/model.bin": "hash"}
asyncio.run(scenario())
def test_manual_model_bundle_rejects_unsafe_archive(tmp_path: Path):
app, _ = _app(tmp_path)
archive_path = tmp_path / "unsafe.tar.gz"
archive_path.write_bytes(_bundle(unsafe=True))
try:
app.state.services.models.install(archive_path=archive_path)
except ValueError as exc:
assert "不安全路径" in str(exc)
else:
raise AssertionError("unsafe archive must be rejected")
assert not (tmp_path / "outside").exists()
def test_default_model_install_uses_official_repositories(tmp_path: Path, monkeypatch):
app, _ = _app(tmp_path)
manager = app.state.services.models
rapid_root = tmp_path / "rapidocr"
rapid_models = rapid_root / "models"
rapid_models.mkdir(parents=True)
for name in (
"ch_PP-OCRv4_det_infer.onnx",
"ch_PP-OCRv4_rec_infer.onnx",
"ch_ppocr_mobile_v2.0_cls_infer.onnx",
):
(rapid_models / name).write_bytes(name.encode())
rapid_module = types.ModuleType("rapidocr_onnxruntime")
rapid_module.__file__ = str(rapid_root / "__init__.py")
hub_module = types.ModuleType("huggingface_hub")
requests_module = types.ModuleType("requests")
class Session:
def __init__(self):
self.trust_env = True
self.proxies = {}
def snapshot_download(repo_id, *, local_dir, **_):
target = Path(local_dir)
target.mkdir(parents=True, exist_ok=True)
(target / "modules.json").write_text(
json.dumps([{"idx": 0, "path": "", "type": "Transformer", "repo": repo_id}]),
encoding="utf-8",
)
return str(target)
hub_module.snapshot_download = snapshot_download
hub_module.configure_http_backend = lambda **_: None
requests_module.Session = Session
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", rapid_module)
monkeypatch.setitem(sys.modules, "huggingface_hub", hub_module)
monkeypatch.setitem(sys.modules, "requests", requests_module)
def download_file(url, destination):
destination.parent.mkdir(parents=True, exist_ok=True)
if destination.suffix == ".xml":
destination.write_text("<net><layers/></net>", encoding="utf-8")
else:
destination.write_bytes(b"\0" * (64 * 1024))
monkeypatch.setattr(manager, "_download_file", download_file)
def export_audio_model(_source, destination):
destination.mkdir(parents=True, exist_ok=True)
for name in (
"config.json",
"preprocessor_config.json",
"tokenizer_config.json",
"vocab.json",
"openvino_encoder_model.xml",
"openvino_encoder_model.bin",
"openvino_decoder_model.xml",
"openvino_decoder_model.bin",
"openvino_decoder_with_past_model.xml",
"openvino_decoder_with_past_model.bin",
):
(destination / name).write_text("{}" if name.endswith(".json") else "model")
monkeypatch.setattr(manager, "_export_audio_model", export_audio_model)
monkeypatch.setattr(manager, "repair_component", lambda *_args, **_kwargs: {})
progress = []
result = manager.install(progress=lambda value, message: progress.append((value, message)))
assert result["manifest"]["source"] == "official-repositories"
assert (manager.settings.models_dir / "ocr" / "det.onnx").is_file()
assert (manager.settings.models_dir / "faces" / "detector.xml").is_file()
assert (manager.settings.models_dir / "audio" / "openvino_encoder_model.xml").is_file()
assert any("多语言文本模型" in message for _, message in progress)
def test_component_install_preserves_other_components_and_manifest(tmp_path: Path, monkeypatch):
app, _ = _app(tmp_path)
manager = app.state.services.models
root = manager.settings.models_dir
(root / "faces").mkdir(parents=True)
(root / "faces" / "sentinel.bin").write_bytes(b"existing-face-index-runtime")
(root / "manifest.json").write_text(
json.dumps(
{
"format_version": 2,
"version": "existing",
"components": {"faces": {"version": "faces-existing"}},
"files": {"faces/sentinel.bin": hashlib.sha256(b"existing-face-index-runtime").hexdigest()},
}
),
encoding="utf-8",
)
rapid_root = tmp_path / "component-rapidocr"
rapid_models = rapid_root / "models"
rapid_models.mkdir(parents=True)
for name in (
"ch_PP-OCRv4_det_infer.onnx",
"ch_PP-OCRv4_rec_infer.onnx",
"ch_ppocr_mobile_v2.0_cls_infer.onnx",
):
(rapid_models / name).write_bytes(name.encode())
rapid_module = types.ModuleType("rapidocr_onnxruntime")
rapid_module.__file__ = str(rapid_root / "__init__.py")
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", rapid_module)
progress: list[tuple[float, str]] = []
result = manager.install(component="ocr", progress=lambda value, message: progress.append((value, message)))
assert result["component"] == "ocr"
assert (root / "faces" / "sentinel.bin").read_bytes() == b"existing-face-index-runtime"
assert (root / "ocr" / "det.onnx").is_file()
manifest = json.loads((root / "manifest.json").read_text(encoding="utf-8"))
assert manifest["components"]["faces"]["version"] == "faces-existing"
assert manifest["components"]["ocr"]["version"] == "RapidOCR-PP-OCRv4"
assert "faces/sentinel.bin" in manifest["files"]
assert "ocr/det.onnx" in manifest["files"]
assert any("启用" in message for _, message in progress)
def test_component_activation_rolls_back_directory_and_manifest(tmp_path: Path, monkeypatch):
app, _ = _app(tmp_path)
manager = app.state.services.models
root = manager.settings.models_dir
old = root / "ocr"
old.mkdir(parents=True)
(old / "sentinel.bin").write_bytes(b"old-component")
original_manifest = {"version": "old", "components": {"ocr": {"version": "old-ocr"}}}
(root / "manifest.json").write_text(json.dumps(original_manifest), encoding="utf-8")
staged = tmp_path / "staged-ocr"
staged.mkdir()
for name in ("det.onnx", "rec.onnx", "cls.onnx"):
(staged / name).write_bytes(b"new")
monkeypatch.setattr(manager, "_write_manifest", lambda _manifest: (_ for _ in ()).throw(OSError("disk")))
with pytest.raises(OSError, match="disk"):
manager._activate_component("ocr", staged, {"version": "new-ocr"})
assert (old / "sentinel.bin").read_bytes() == b"old-component"
assert json.loads((root / "manifest.json").read_text(encoding="utf-8")) == original_manifest
def test_component_install_api_queues_missing_models_individually_and_reports_progress(tmp_path: Path):
app, headers = _app(tmp_path)
stale_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "stale"})
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post("/api/v1/models/install", headers=headers, json={"component": "all"})
assert response.status_code == 202
job_ids = response.json()["job_ids"]
assert response.json()["cancelled_audio_jobs"] == 1
# The deterministic hash visual backend is already ready in tests.
assert list(job_ids) == ["ocr", "faces", "audio"]
assert response.json()["job_id"] == job_ids["ocr"]
duplicate = await client.post("/api/v1/models/install", headers=headers, json={"component": "ocr"})
assert duplicate.json()["job_ids"]["ocr"] == job_ids["ocr"]
with app.state.services.db.transaction() as conn:
conn.execute(
"UPDATE jobs SET status='running',progress=.42,message='正在复制 OCR 权重' WHERE id=?",
(job_ids["ocr"],),
)
status = (await client.get("/api/v1/models", headers=headers)).json()
assert status["installations"]["ocr"] == {
"job_id": job_ids["ocr"],
"status": "running",
"progress": 0.42,
"message": "正在复制 OCR 权重",
"error": None,
}
custom = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "faces", "url": "https://example.invalid/models.tar.gz"},
)
assert custom.status_code == 400
with app.state.services.db.read() as conn:
audio = conn.execute("SELECT status FROM jobs WHERE id=?", (stale_audio,)).fetchone()
assert audio["status"] == "cancelled"
asyncio.run(scenario())
def test_audio_medium_install_is_independent_and_does_not_cancel_active_audio_work(tmp_path: Path):
app, headers = _app(tmp_path)
active_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "active"})
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "audio", "variant": "medium"},
)
assert response.status_code == 202
assert response.json()["cancelled_audio_jobs"] == 0
job_id = response.json()["job_ids"]["audio:medium"]
with app.state.services.db.read() as conn:
job = conn.execute("SELECT payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)).fetchone()
active = conn.execute("SELECT status FROM jobs WHERE id=?", (active_audio,)).fetchone()
assert json.loads(job["payload_json"]) == {"component": "audio", "variant": "medium"}
assert job["dedupe_key"] == "install-model:audio:medium"
assert active["status"] == "queued"
asyncio.run(scenario())
def test_pending_audio_install_cancels_stale_audio_jobs_on_startup(tmp_path: Path):
app, _ = _app(tmp_path)
stale_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "stale"})
app.state.services.jobs.enqueue("install_models", {"component": "audio"}, dedupe_key="install-model:audio")
assert app.state.services._cancel_audio_for_pending_model_install() == 1
with app.state.services.db.read() as conn:
row = conn.execute("SELECT status,message FROM jobs WHERE id=?", (stale_audio,)).fetchone()
assert row["status"] == "cancelled"
assert "模型安装优先" in row["message"]
def test_component_install_repairs_ready_runtime_without_redownloading_model(tmp_path: Path, monkeypatch):
app, headers = _app(tmp_path)
manager = app.state.services.models
monkeypatch.setattr(
manager,
"status",
lambda: {
"visual_ready": True,
"ocr_ready": False,
"faces_ready": False,
"audio_ready": False,
"operational_components": {"visual": False},
},
)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "visual"},
)
assert response.status_code == 202
job_id = response.json()["job_ids"]["visual"]
with app.state.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"
asyncio.run(scenario())
def test_outdated_audio_repair_queues_real_model_reinstall(tmp_path: Path, monkeypatch):
app, headers = _app(tmp_path)
manager = app.state.services.models
monkeypatch.setattr(
manager,
"status",
lambda: {
"visual_ready": False,
"ocr_ready": False,
"faces_ready": False,
"audio_ready": True,
"operational_components": {"audio": False},
"component_health": {"audio": {"state": "outdated", "error": "需要转换为 FP16 模型"}},
},
)
async def scenario():
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.post(
"/api/v1/models/install",
headers=headers,
json={"component": "audio"},
)
assert response.status_code == 202
job_id = response.json()["job_ids"]["audio"]
with app.state.services.db.read() as connection:
job = connection.execute("SELECT kind,payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)).fetchone()
assert job["kind"] == "install_models"
assert json.loads(job["payload_json"]) == {"component": "audio"}
assert job["dedupe_key"] == "install-model:audio"
asyncio.run(scenario())