diff --git a/src/vidxp/capabilities/actor/definition.py b/src/vidxp/capabilities/actor/definition.py index c9ba378..e05c293 100644 --- a/src/vidxp/capabilities/actor/definition.py +++ b/src/vidxp/capabilities/actor/definition.py @@ -32,21 +32,18 @@ ) from vidxp.capabilities.visual import index_capabilities from vidxp.core.contracts import IndexConfig, VideoSource -from vidxp.core.indexing_common import ProgressCallback +from vidxp.core.indexing_common import ProgressCallback, report_preparation def prepare_models( context: PreparationContext, progress: ProgressCallback | None, ) -> tuple[str, ...]: - if progress is not None: - progress( - { - "state": "preparing", - "stage": "actor_models", - "message": "Preparing OpenCV Zoo YuNet and SFace models.", - } - ) + report_preparation( + progress, + "actor_models", + "Preparing OpenCV Zoo YuNet and SFace models.", + ) get_actor_models(context.runtime, download=True, progress=progress) return (YUNET_MODEL.filename, SFACE_MODEL.filename) diff --git a/src/vidxp/capabilities/actor/models.py b/src/vidxp/capabilities/actor/models.py index 3673fd8..2a487b3 100644 --- a/src/vidxp/capabilities/actor/models.py +++ b/src/vidxp/capabilities/actor/models.py @@ -4,6 +4,7 @@ from typing import Any, Callable from vidxp.ports import ModelRuntimePort +from vidxp.core.indexing_common import report_preparation from vidxp.capabilities.actor.specs import ( SFACE_MODEL, YUNET_MODEL, @@ -38,14 +39,11 @@ def load() -> ActorModels: download=download, progress=progress, ) - if progress is not None: - progress( - { - "state": "preparing", - "stage": "loading_model", - "message": "Loading OpenCV Zoo YuNet and SFace models.", - } - ) + report_preparation( + progress, + "loading_model", + "Loading OpenCV Zoo YuNet and SFace models.", + ) models = ActorModels( detector=cv2.FaceDetectorYN.create( str(detector_path), diff --git a/src/vidxp/capabilities/dialogue/definition.py b/src/vidxp/capabilities/dialogue/definition.py index 950bd8c..0474773 100644 --- a/src/vidxp/capabilities/dialogue/definition.py +++ b/src/vidxp/capabilities/dialogue/definition.py @@ -26,7 +26,7 @@ ) from vidxp.capabilities.schemas import SearchInput, SearchResult from vidxp.core.contracts import IndexConfig, VideoSource -from vidxp.core.indexing_common import ProgressCallback +from vidxp.core.indexing_common import ProgressCallback, report_preparation def filter_requirements_for_source( source: VideoSource, requirements: tuple[Requirement, ...], @@ -49,14 +49,7 @@ def prepare_models( prepared = [] def report(stage: str, message: str) -> None: - if progress is not None: - progress( - { - "state": "preparing", - "stage": stage, - "message": message, - } - ) + report_preparation(progress, stage, message) report( "dialogue_model", diff --git a/src/vidxp/capabilities/dialogue/models.py b/src/vidxp/capabilities/dialogue/models.py index 7a671ea..6979fd3 100644 --- a/src/vidxp/capabilities/dialogue/models.py +++ b/src/vidxp/capabilities/dialogue/models.py @@ -3,6 +3,7 @@ from typing import Any, Callable from vidxp.ports import ModelRuntimePort +from vidxp.core.indexing_common import report_preparation from vidxp.model_contracts import loaded_compute_precision from vidxp.capabilities.dialogue.specs import ( FASTER_WHISPER_MODEL, @@ -28,16 +29,11 @@ def load() -> Any: download=download, progress=progress, ) - if progress is not None: - progress( - { - "state": "preparing", - "stage": "loading_model", - "message": ( - f"Loading {QWEN3_EMBEDDING_MODEL.model_id}." - ), - } - ) + report_preparation( + progress, + "loading_model", + f"Loading {QWEN3_EMBEDDING_MODEL.model_id}.", + ) model = SentenceTransformer( str(snapshot), device=device, @@ -74,16 +70,11 @@ def load() -> Any: download=download, progress=progress, ) - if progress is not None: - progress( - { - "state": "preparing", - "stage": "loading_model", - "message": ( - f"Loading {FASTER_WHISPER_MODEL.model_id}." - ), - } - ) + report_preparation( + progress, + "loading_model", + f"Loading {FASTER_WHISPER_MODEL.model_id}.", + ) model = WhisperModel( str(snapshot), device=device.split(":", 1)[0], diff --git a/src/vidxp/capabilities/scene/definition.py b/src/vidxp/capabilities/scene/definition.py index 366ad4e..9bc3621 100644 --- a/src/vidxp/capabilities/scene/definition.py +++ b/src/vidxp/capabilities/scene/definition.py @@ -19,23 +19,18 @@ from vidxp.capabilities.schemas import SearchInput, SearchResult from vidxp.capabilities.visual import index_capabilities from vidxp.core.contracts import IndexConfig, VideoSource -from vidxp.core.indexing_common import ProgressCallback +from vidxp.core.indexing_common import ProgressCallback, report_preparation def prepare_models( context: PreparationContext, progress: ProgressCallback | None, ) -> tuple[str, ...]: SceneConfig.model_validate(context.settings) - if progress is not None: - progress( - { - "state": "preparing", - "stage": "scene_model", - "message": ( - f"Preparing scene model: SigLIP2 {SIGLIP2_MODEL.model_id}" - ), - } - ) + report_preparation( + progress, + "scene_model", + f"Preparing scene model: SigLIP2 {SIGLIP2_MODEL.model_id}", + ) get_scene_model(context.runtime, download=True, progress=progress) return (SIGLIP2_MODEL.model_id,) diff --git a/src/vidxp/capabilities/scene/models.py b/src/vidxp/capabilities/scene/models.py index 50d8da8..07f0555 100644 --- a/src/vidxp/capabilities/scene/models.py +++ b/src/vidxp/capabilities/scene/models.py @@ -4,6 +4,7 @@ from typing import Any, Callable from vidxp.ports import ModelRuntimePort +from vidxp.core.indexing_common import report_preparation from vidxp.model_contracts import loaded_compute_precision from vidxp.capabilities.scene.specs import SIGLIP2_MODEL @@ -32,14 +33,11 @@ def load() -> SceneModel: download=download, progress=progress, ) - if progress is not None: - progress( - { - "state": "preparing", - "stage": "loading_model", - "message": f"Loading {SIGLIP2_MODEL.model_id}.", - } - ) + report_preparation( + progress, + "loading_model", + f"Loading {SIGLIP2_MODEL.model_id}.", + ) common = { "cache_dir": str(runtime.model_cache), "local_files_only": True, diff --git a/src/vidxp/core/indexing_common.py b/src/vidxp/core/indexing_common.py index 9758c7c..b1b1b75 100644 --- a/src/vidxp/core/indexing_common.py +++ b/src/vidxp/core/indexing_common.py @@ -6,6 +6,28 @@ ProgressCallback = Callable[[dict[str, Any]], None] +def report_preparation( + callback: ProgressCallback | None, + stage: str, + message: str, + *, + current: int | None = None, + total: int | None = None, +) -> None: + if callback is None: + return + event: dict[str, Any] = { + "state": "preparing", + "stage": stage, + "message": message, + } + if current is not None: + event["current"] = current + if total is not None: + event["total"] = total + callback(event) + + def report_progress( callback: ProgressCallback | None, stage: str, diff --git a/src/vidxp/model_contracts.py b/src/vidxp/model_contracts.py index b1ab521..e3852db 100644 --- a/src/vidxp/model_contracts.py +++ b/src/vidxp/model_contracts.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +import hashlib from pathlib import Path import re from typing import Any @@ -174,4 +175,23 @@ def model_artifact_cached( cache: Path, spec: ModelSpec | ArtifactSpec, ) -> bool: - return model_artifact_path(cache, spec).is_file() + return model_artifact_valid(model_artifact_path(cache, spec), spec) + + +def model_artifact_valid( + path: Path, + spec: ModelSpec | ArtifactSpec, +) -> bool: + if not path.is_file(): + return False + expected = ( + spec.weights_sha256 if isinstance(spec, ModelSpec) else spec.sha256 + ) + digest = hashlib.sha256() + try: + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + except OSError: + return False + return digest.hexdigest() == expected diff --git a/src/vidxp/runtime.py b/src/vidxp/runtime.py index abda32c..97086f2 100644 --- a/src/vidxp/runtime.py +++ b/src/vidxp/runtime.py @@ -1,13 +1,12 @@ from __future__ import annotations import platform -import hashlib from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout from contextlib import contextmanager from threading import BoundedSemaphore, Lock, RLock from time import monotonic, sleep -from typing import Any, Callable, Iterator +from typing import Any, Callable, Iterator, TypeVar from pathlib import Path from vidxp.application_models import RuntimeProfile @@ -17,7 +16,9 @@ ModelArtifactUnavailableError, ModelKey, ModelSpec, + model_artifact_valid, ) +from vidxp.core.indexing_common import report_preparation from vidxp.settings import VidXPSettings @@ -32,6 +33,7 @@ class _ModelDownloadVerificationError(RuntimeError): _MODEL_DOWNLOAD_ATTEMPTS = 3 _DOWNLOAD_HEARTBEAT_SECONDS = 5.0 _MINIMUM_PROGRESS_BYTES = 1024 * 1024 +_DownloadResult = TypeVar("_DownloadResult") def _download_failure_reason(exc: Exception) -> str: @@ -83,6 +85,39 @@ def _download_failure_retryable( return None +def _download_with_retries( + spec: ModelSpec | ArtifactSpec, + download: Callable[[int], _DownloadResult], + *, + resumable: bool, + hash_mismatch_is_retryable: bool = False, + on_retry: Callable[[int], None] | None = None, +) -> _DownloadResult: + for attempt in range(1, _MODEL_DOWNLOAD_ATTEMPTS + 1): + try: + return download(attempt) + except Exception as exc: + retryable = _download_failure_retryable( + exc, + hash_mismatch_is_retryable=hash_mismatch_is_retryable, + ) + if retryable is None: + raise + if attempt >= _MODEL_DOWNLOAD_ATTEMPTS or not retryable: + raise ModelArtifactDownloadError( + spec.capability, + spec.model_id, + attempts=attempt, + reason=_download_failure_reason(exc), + resumable=resumable, + retryable=retryable, + ) from exc + if on_retry is not None: + on_retry(attempt + 1) + sleep(2 ** (attempt - 1)) + raise AssertionError("Model download retry loop did not terminate.") + + def _torch_accelerators() -> tuple[bool, bool]: try: import torch @@ -138,14 +173,6 @@ def resolve_backends(requested: str) -> RuntimeProfile: ) -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for block in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(block) - return digest.hexdigest() - - class ResourceScheduler: """Bound concurrent model work without owning workflow state.""" @@ -259,7 +286,7 @@ def update(self, n=1): state["message"] = f"Downloading {spec.model_id}." return result - def download() -> str: + def download_snapshot() -> str: snapshot = Path( snapshot_download( repo_id=spec.model_id, @@ -270,10 +297,7 @@ def download() -> str: ) ) weights = snapshot / spec.weights_file - if ( - not weights.is_file() - or _sha256(weights) != spec.weights_sha256 - ): + if not model_artifact_valid(weights, spec): raise _ModelDownloadVerificationError return str(snapshot) @@ -296,21 +320,17 @@ def report(*, force: bool = False): heartbeat = now - reported_at >= _DOWNLOAD_HEARTBEAT_SECONDS if not force and reported_current >= 0 and not advanced and not heartbeat: return - progress( - { - "state": "preparing", - "stage": "downloading_model", - **event, - } + report_preparation( + progress, + "downloading_model", + event["message"], + current=event["current"], + total=event["total"], ) reported_at = now reported_current = current - last_error: Exception | None = None - last_retryable = False - attempts = 0 - for attempt in range(1, _MODEL_DOWNLOAD_ATTEMPTS + 1): - attempts = attempt + def download(attempt: int) -> str: with state_lock: state["message"] = ( f"Connecting to download {spec.model_id}." @@ -321,49 +341,33 @@ def report(*, force: bool = False): ) ) report(force=True) - try: - with ThreadPoolExecutor(max_workers=1) as pool: - future = pool.submit(download) - while True: - try: - snapshot = future.result(timeout=0.5) - break - except FutureTimeout: - report() - except Exception as exc: - last_error = exc - retryable = _download_failure_retryable(exc) - if retryable is None: - raise - last_retryable = retryable - if ( - attempt >= _MODEL_DOWNLOAD_ATTEMPTS - or not retryable - ): - break - with state_lock: - state["message"] = ( - f"Download interrupted for {spec.model_id}; cached " - "partial files will be resumed." - ) - report(force=True) - sleep(2 ** (attempt - 1)) - continue + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(download_snapshot) + while True: + try: + return future.result(timeout=0.5) + except FutureTimeout: + report() + + def report_retry(_next_attempt: int) -> None: with state_lock: - state["current"] = spec.download_size_bytes - state["message"] = f"Downloaded {spec.model_id}." + state["message"] = ( + f"Download interrupted for {spec.model_id}; cached " + "partial files will be resumed." + ) report(force=True) - return Path(snapshot) - - assert last_error is not None - raise ModelArtifactDownloadError( - spec.capability, - spec.model_id, - attempts=attempts, - reason=_download_failure_reason(last_error), + + snapshot = _download_with_retries( + spec, + download, resumable=True, - retryable=last_retryable, - ) from last_error + on_retry=report_retry, + ) + with state_lock: + state["current"] = spec.download_size_bytes + state["message"] = f"Downloaded {spec.model_id}." + report(force=True) + return Path(snapshot) def resolve_model( self, @@ -390,8 +394,7 @@ def resolve_model( local_weights = local_snapshot / spec.weights_file snapshot = ( local_snapshot - if local_weights.is_file() - and _sha256(local_weights) == spec.weights_sha256 + if model_artifact_valid(local_weights, spec) else None ) except Exception: @@ -424,88 +427,59 @@ def resolve_artifact( try: destination = self.settings.model_cache / spec.provider path = destination / spec.filename - if not path.is_file() or _sha256(path) != spec.sha256: + if not model_artifact_valid(path, spec): if not download or not self.settings.allow_model_downloads: raise ModelArtifactUnavailableError(spec.capability) - if progress is not None: - progress( - { - "state": "preparing", - "stage": "downloading_model", - "message": f"Downloading {spec.model_id}.", - "current": 0, - "total": spec.download_size_bytes, - } - ) + report_preparation( + progress, + "downloading_model", + f"Downloading {spec.model_id}.", + current=0, + total=spec.download_size_bytes, + ) import pooch - last_error = None - for attempt in range(1, _MODEL_DOWNLOAD_ATTEMPTS + 1): - try: - resolved = Path( - pooch.retrieve( - url=spec.url, - known_hash=f"sha256:{spec.sha256}", - fname=spec.filename, - path=destination, - progressbar=False, - ) - ) - break - except Exception as exc: - last_error = exc - retryable = _download_failure_retryable( - exc, - hash_mismatch_is_retryable=True, + def download_artifact(_attempt: int) -> Path: + return Path( + pooch.retrieve( + url=spec.url, + known_hash=f"sha256:{spec.sha256}", + fname=spec.filename, + path=destination, + progressbar=False, ) - if retryable is None: - raise - if ( - attempt >= _MODEL_DOWNLOAD_ATTEMPTS - or not retryable - ): - raise ModelArtifactDownloadError( - spec.capability, - spec.model_id, - attempts=attempt, - reason=_download_failure_reason(exc), - resumable=False, - retryable=retryable, - ) from exc - if progress is not None: - progress( - { - "state": "preparing", - "stage": "downloading_model", - "message": ( - f"Download interrupted for " - f"{spec.model_id}; retrying attempt " - f"{attempt + 1} of " - f"{_MODEL_DOWNLOAD_ATTEMPTS}. This " - "file will restart from zero." - ), - "current": 0, - "total": spec.download_size_bytes, - } - ) - sleep(2 ** (attempt - 1)) - else: - assert last_error is not None - raise last_error + ) + + def report_retry(next_attempt: int) -> None: + report_preparation( + progress, + "downloading_model", + f"Download interrupted for {spec.model_id}; " + f"retrying attempt {next_attempt} of " + f"{_MODEL_DOWNLOAD_ATTEMPTS}. This file will " + "restart from zero.", + current=0, + total=spec.download_size_bytes, + ) + + resolved = _download_with_retries( + spec, + download_artifact, + resumable=False, + hash_mismatch_is_retryable=True, + on_retry=report_retry, + ) else: resolved = path - if not resolved.is_file() or _sha256(resolved) != spec.sha256: + if not model_artifact_valid(resolved, spec): raise ModelArtifactUnavailableError(spec.capability) - if progress is not None: - progress( - { - "state": "preparing", - "stage": "downloading_model", - "message": f"Verified {spec.model_id}.", - "current": spec.download_size_bytes, - "total": spec.download_size_bytes, - } - ) + report_preparation( + progress, + "downloading_model", + f"Verified {spec.model_id}.", + current=spec.download_size_bytes, + total=spec.download_size_bytes, + ) except (ModelArtifactDownloadError, ModelArtifactUnavailableError): raise except Exception as exc: diff --git a/tests/test_models.py b/tests/test_models.py index c56923a..ca3af4a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -40,6 +40,7 @@ ModelArtifactDownloadError, ModelArtifactUnavailableError, ModelKey, + model_artifact_cached, model_artifact_path, ) from vidxp.runtime import ModelRuntime, resolve_backends @@ -204,7 +205,7 @@ def test_normal_model_resolution_never_downloads_implicitly(self): retrieve.assert_not_called() - def test_model_readiness_checks_the_pinned_cache_without_loading(self): + def test_model_readiness_rejects_corrupt_pinned_cache_without_loading(self): with TemporaryDirectory() as directory: cache = Path(directory) / "models" path = model_artifact_path(cache, YUNET_MODEL) @@ -224,7 +225,7 @@ def test_model_readiness_checks_the_pinned_cache_without_loading(self): ( YUNET_MODEL.model_id, YUNET_MODEL.download_size_bytes, - True, + False, ), ( SFACE_MODEL.model_id, @@ -234,6 +235,23 @@ def test_model_readiness_checks_the_pinned_cache_without_loading(self): ], ) + def test_model_cache_requires_the_pinned_checksum(self): + with TemporaryDirectory() as directory: + content = b"verified model" + spec = replace( + YUNET_MODEL, + filename="verified.onnx", + sha256=hashlib.sha256(content).hexdigest(), + ) + cache = Path(directory) / "models" + path = model_artifact_path(cache, spec) + path.parent.mkdir(parents=True) + path.write_bytes(content) + + self.assertTrue(model_artifact_cached(cache, spec)) + path.write_bytes(b"corrupt") + self.assertFalse(model_artifact_cached(cache, spec)) + def test_explicit_snapshot_download_reports_bytes(self): with TemporaryDirectory() as directory: snapshot = Path(directory) / "snapshot" diff --git a/uv.lock b/uv.lock index 59e02df..fcfa563 100644 --- a/uv.lock +++ b/uv.lock @@ -919,14 +919,14 @@ wheels = [ [[package]] name = "gitpython" -version = "3.1.57" +version = "3.1.58" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gitdb" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ba/0d/132ed135c871b6bf91adf16a0e43797cd535b81d4973b5d09291c54fc5ee/gitpython-3.1.57.tar.gz", hash = "sha256:c493ec57c0ef6b19743798b6a5af859c71814b524e7e6f97baa2f8e658961488", size = 225898, upload-time = "2026-07-26T07:33:26.351Z" } +sdist = { url = "https://files.pythonhosted.org/packages/26/d6/5f358ff283325580c2003a6d953aea18cfe10ae87b46f5ebc80fa3a386dc/gitpython-3.1.58.tar.gz", hash = "sha256:621416df10ef3fd0e19fabf9172ddeed0fa704d353d04f194eec56a625a95b22", size = 228498, upload-time = "2026-08-04T15:05:49.47Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/41/6e/2139de986d9c7c3ac86f1f8be43858ce90bdfe2f7175e6c80c650ba15242/gitpython-3.1.57-py3-none-any.whl", hash = "sha256:4ccf7d73c10f5c9e76043fbb2675ac5a1b3ff5b41e648f56bcbed5f63792ecaf", size = 217151, upload-time = "2026-07-26T07:33:24.838Z" }, + { url = "https://files.pythonhosted.org/packages/ec/0c/9d8752098bc442f0726e64aa6135940b3a96809915d1aa4206c1bb97881d/gitpython-3.1.58-py3-none-any.whl", hash = "sha256:d331e722577f0fd7fc1f857419b3ecc07af66282b933d2a4d95f84a042fdd50f", size = 220183, upload-time = "2026-08-04T15:05:48.025Z" }, ] [[package]]