diff --git a/src/vidxp/application.py b/src/vidxp/application.py index fcbd1cf..dd25eaf 100644 --- a/src/vidxp/application.py +++ b/src/vidxp/application.py @@ -27,6 +27,7 @@ ImportMediaCommand, IndexSnapshotReference, MediaAsset, + ModelDownloadError, ModelUnavailableError, PrepareModelsCommand, PrepareModelsResult, @@ -59,7 +60,10 @@ from vidxp.ports import IndexBackend, ModelRuntimePort, QueryModelPort from vidxp.query_service import GroundedQueryService from vidxp.search_fusion import fuse_search_results -from vidxp.model_contracts import ModelArtifactUnavailableError +from vidxp.model_contracts import ( + ModelArtifactDownloadError, + ModelArtifactUnavailableError, +) from vidxp.repository_layout import RepositoryLayout from vidxp.settings import VidXPSettings from vidxp.control_plane import ControlPlaneApplication @@ -119,6 +123,15 @@ def _capability_dependencies( capabilities, self.registry.install_hint(capabilities), ) from exc + except ModelArtifactDownloadError as exc: + raise ModelDownloadError( + exc.capability, + exc.model_id, + attempts=exc.attempts, + reason=exc.reason, + resumable=exc.resumable, + retryable=exc.retryable, + ) from exc except ModelArtifactUnavailableError as exc: raise ModelUnavailableError(exc.capability) from exc except (ModuleNotFoundError, CapabilityDependencyError) as exc: diff --git a/src/vidxp/application_models.py b/src/vidxp/application_models.py index 9975e3a..807a0a0 100644 --- a/src/vidxp/application_models.py +++ b/src/vidxp/application_models.py @@ -199,6 +199,43 @@ def __init__(self, capability: str) -> None: ) +class ModelDownloadError(ApplicationError): + def __init__( + self, + capability: str, + model_id: str, + *, + attempts: int, + reason: str, + resumable: bool, + retryable: bool, + ) -> None: + modality = capability.split(".", 1)[0] + remediation = f"vidxp prepare --modalities {modality}" + retry_message = ( + "Partial files were kept and the next preparation attempt will " + "resume them." + if resumable + else "The next preparation attempt will restart this file." + ) + super().__init__( + "model_download_failed", + ErrorCategory.unavailable, + f"Downloading {model_id} failed after {attempts} attempt(s) " + f"({reason}). {retry_message} Run " + f"`{remediation}` again when the connection is available.", + details={ + "capability": capability, + "model": model_id, + "attempts": attempts, + "reason": reason, + "partial_files_preserved": resumable, + "remediation": remediation, + }, + retryable=retryable, + ) + + class InvalidRequestError(ApplicationError): def __init__( self, diff --git a/src/vidxp/cli_commands/runtime.py b/src/vidxp/cli_commands/runtime.py index f5992bc..560827b 100644 --- a/src/vidxp/cli_commands/runtime.py +++ b/src/vidxp/cli_commands/runtime.py @@ -471,9 +471,8 @@ def prepare( typer.confirm("Download these models?", abort=True) show_progress = not state.quiet and output_format == OutputFormat.rich if show_progress: - emit_progress( - "Starting model preparation for " + ", ".join(selected) + "." - ) + action = "Downloading and validating" if missing else "Validating cached" + emit_progress(f"{action} models for " + ", ".join(selected) + ".") job = state.jobs.submit_prepare_models( PrepareModelsCommand( modalities=selected, diff --git a/src/vidxp/model_contracts.py b/src/vidxp/model_contracts.py index 6b6788e..b1ab521 100644 --- a/src/vidxp/model_contracts.py +++ b/src/vidxp/model_contracts.py @@ -24,6 +24,29 @@ def __init__(self, capability: str) -> None: super().__init__(f"Model artifacts for {capability} are unavailable.") +class ModelArtifactDownloadError(RuntimeError): + def __init__( + self, + capability: str, + model_id: str, + *, + attempts: int, + reason: str, + resumable: bool, + retryable: bool, + ) -> None: + self.capability = capability + self.model_id = model_id + self.attempts = attempts + self.reason = reason + self.resumable = resumable + self.retryable = retryable + super().__init__( + f"Downloading {model_id} failed after {attempts} attempt(s): " + f"{reason}." + ) + + @dataclass(frozen=True) class ModelKey: capability: str diff --git a/src/vidxp/runtime.py b/src/vidxp/runtime.py index 4fe70d7..abda32c 100644 --- a/src/vidxp/runtime.py +++ b/src/vidxp/runtime.py @@ -6,13 +6,14 @@ from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout from contextlib import contextmanager from threading import BoundedSemaphore, Lock, RLock -from time import monotonic +from time import monotonic, sleep from typing import Any, Callable, Iterator from pathlib import Path from vidxp.application_models import RuntimeProfile from vidxp.model_contracts import ( ArtifactSpec, + ModelArtifactDownloadError, ModelArtifactUnavailableError, ModelKey, ModelSpec, @@ -24,6 +25,64 @@ class RuntimeBackendUnavailableError(RuntimeError): """Raised when an explicitly requested compute backend cannot be used.""" +class _ModelDownloadVerificationError(RuntimeError): + """Raised internally when a completed transfer is missing pinned weights.""" + + +_MODEL_DOWNLOAD_ATTEMPTS = 3 +_DOWNLOAD_HEARTBEAT_SECONDS = 5.0 +_MINIMUM_PROGRESS_BYTES = 1024 * 1024 + + +def _download_failure_reason(exc: Exception) -> str: + if isinstance(exc, _ModelDownloadVerificationError): + return "artifact verification failed" + response = getattr(exc, "response", None) + status = getattr(response, "status_code", None) + if isinstance(status, int): + return f"HTTP {status} {type(exc).__name__}" + return type(exc).__name__ + + +def _download_failure_retryable( + exc: Exception, + *, + hash_mismatch_is_retryable: bool = False, +) -> bool | None: + if isinstance(exc, _ModelDownloadVerificationError): + return True + response = getattr(exc, "response", None) + status = getattr(response, "status_code", None) + if isinstance(status, int): + return status in {408, 409, 425, 429} or status >= 500 + if isinstance(exc, (ConnectionError, TimeoutError)): + return True + try: + import httpx + except ModuleNotFoundError: + pass + else: + if isinstance(exc, httpx.TransportError): + return True + try: + from requests import exceptions as requests_exceptions + except ModuleNotFoundError: + pass + else: + if isinstance( + exc, + (requests_exceptions.ConnectionError, requests_exceptions.Timeout), + ): + return True + if hash_mismatch_is_retryable and isinstance(exc, ValueError): + return True + if isinstance(exc, OSError): + return False + if type(exc).__module__.startswith(("huggingface_hub", "pooch")): + return False + return None + + def _torch_accelerators() -> tuple[bool, bool]: try: import torch @@ -190,60 +249,121 @@ def update(self, n=1): result = super().update(n) if self.unit == "B": with state_lock: - state.update( - { - "current": int(self.n), - "total": ( - int(self.total) - if self.total - else None - ), - "message": f"Downloading {spec.model_id}.", - } + previous = getattr(self, "_vidxp_reported_bytes", 0) + current = int(self.n) + self._vidxp_reported_bytes = current + state["current"] = min( + spec.download_size_bytes, + int(state["current"]) + max(0, current - previous), ) + state["message"] = f"Downloading {spec.model_id}." return result def download() -> str: - return snapshot_download( - repo_id=spec.model_id, - revision=spec.revision, - cache_dir=str(cache), - local_files_only=False, - tqdm_class=ReportingTqdm, + snapshot = Path( + snapshot_download( + repo_id=spec.model_id, + revision=spec.revision, + cache_dir=str(cache), + local_files_only=False, + tqdm_class=ReportingTqdm, + ) ) + weights = snapshot / spec.weights_file + if ( + not weights.is_file() + or _sha256(weights) != spec.weights_sha256 + ): + raise _ModelDownloadVerificationError + return str(snapshot) reported_at = 0.0 - with ThreadPoolExecutor(max_workers=1) as pool: - future = pool.submit(download) - while True: - try: - snapshot = future.result(timeout=0.5) - if progress is not None: - with state_lock: - event = dict(state) - progress( - { - "state": "preparing", - "stage": "downloading_model", - **event, - } - ) + reported_current = -1 + progress_step = max( + _MINIMUM_PROGRESS_BYTES, + spec.download_size_bytes // 100, + ) + + def report(*, force: bool = False): + nonlocal reported_at, reported_current + if progress is None: + return + now = monotonic() + with state_lock: + event = dict(state) + current = int(event["current"]) + advanced = current - reported_current >= progress_step + 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, + } + ) + 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 + with state_lock: + state["message"] = ( + f"Connecting to download {spec.model_id}." + if attempt == 1 + else ( + f"Retrying download of {spec.model_id} " + f"(attempt {attempt} of {_MODEL_DOWNLOAD_ATTEMPTS})." + ) + ) + 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 - except FutureTimeout: - now = monotonic() - if progress is None or now - reported_at < 1: - continue - with state_lock: - event = dict(state) - progress( - { - "state": "preparing", - "stage": "downloading_model", - **event, - } + with state_lock: + state["message"] = ( + f"Download interrupted for {spec.model_id}; cached " + "partial files will be resumed." ) - reported_at = now - return Path(snapshot) + report(force=True) + sleep(2 ** (attempt - 1)) + continue + with state_lock: + state["current"] = spec.download_size_bytes + state["message"] = f"Downloaded {spec.model_id}." + 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), + resumable=True, + retryable=last_retryable, + ) from last_error def resolve_model( self, @@ -257,8 +377,9 @@ def resolve_model( from huggingface_hub import snapshot_download try: + snapshot: Path | None try: - snapshot = Path( + local_snapshot = Path( snapshot_download( repo_id=spec.model_id, revision=spec.revision, @@ -266,7 +387,16 @@ def resolve_model( local_files_only=True, ) ) + local_weights = local_snapshot / spec.weights_file + snapshot = ( + local_snapshot + if local_weights.is_file() + and _sha256(local_weights) == spec.weights_sha256 + else None + ) except Exception: + snapshot = None + if snapshot is None: if not download or not self.settings.allow_model_downloads: raise ModelArtifactUnavailableError(spec.capability) snapshot = self._download_snapshot( @@ -274,10 +404,7 @@ def resolve_model( cache=self.settings.model_cache, progress=progress, ) - weights = snapshot / spec.weights_file - if not weights.is_file() or _sha256(weights) != spec.weights_sha256: - raise ModelArtifactUnavailableError(spec.capability) - except ModelArtifactUnavailableError: + except (ModelArtifactDownloadError, ModelArtifactUnavailableError): raise except Exception as exc: raise ModelArtifactUnavailableError(spec.capability) from exc @@ -306,24 +433,80 @@ def resolve_artifact( "state": "preparing", "stage": "downloading_model", "message": f"Downloading {spec.model_id}.", + "current": 0, + "total": spec.download_size_bytes, } ) import pooch - resolved = Path( - pooch.retrieve( - url=spec.url, - known_hash=f"sha256:{spec.sha256}", - fname=spec.filename, - path=destination, - progressbar=False, - ) - ) + 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, + ) + 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 else: resolved = path if not resolved.is_file() or _sha256(resolved) != spec.sha256: raise ModelArtifactUnavailableError(spec.capability) - except ModelArtifactUnavailableError: + 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, + } + ) + except (ModelArtifactDownloadError, ModelArtifactUnavailableError): raise except Exception as exc: raise ModelArtifactUnavailableError(spec.capability) from exc diff --git a/tests/test_api.py b/tests/test_api.py index b862c1e..658c0e1 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -19,6 +19,7 @@ from vidxp.application_models import ( ApplicationError, ErrorCategory, + ErrorDetail, ComponentReadiness, Job, JobKind, @@ -405,6 +406,35 @@ def test_job_submission_is_thin_idempotent_delegation(self): ), ) + def test_failed_model_preparation_job_is_structured_over_http(self): + with TemporaryDirectory() as directory: + context = self.context(Path(directory)) + context.jobs.get.return_value = Job( + job_id=JOB_ID, + kind=JobKind.prepare_models, + state=JobState.failed, + queue=JobQueue.cpu, + error=ErrorDetail( + code="model_download_failed", + category=ErrorCategory.unavailable, + message="The model download failed after three attempts.", + details={ + "model": "publisher/model", + "partial_files_preserved": True, + "remediation": "vidxp prepare --modalities dialogue", + }, + retryable=True, + ), + ) + with TestClient(create_app(context=context)) as client: + response = client.get(f"/api/v1/jobs/{JOB_ID}") + + self.assertEqual(response.status_code, 200) + error = response.json()["error"] + self.assertEqual(error["code"], "model_download_failed") + self.assertTrue(error["retryable"]) + self.assertTrue(error["details"]["partial_files_preserved"]) + def test_missing_models_fail_before_job_submission(self): with TemporaryDirectory() as directory: context = self.context(Path(directory)) diff --git a/tests/test_application.py b/tests/test_application.py index f3bb2d4..ee55144 100644 --- a/tests/test_application.py +++ b/tests/test_application.py @@ -21,6 +21,7 @@ IndexResult, IndexSnapshotReference, ModelUnavailableError, + ModelDownloadError, PrepareModelsCommand, QueryAnswerMode, QueryVideoCommand, @@ -30,7 +31,10 @@ from vidxp.core.media import MediaUnavailableError from vidxp.core.contracts import IndexConfig, IndexSchemaError from vidxp.infrastructure.local_index import LocalIndexBackend -from vidxp.model_contracts import ModelArtifactUnavailableError +from vidxp.model_contracts import ( + ModelArtifactDownloadError, + ModelArtifactUnavailableError, +) from vidxp.capabilities.contracts import ( CapabilityDefinition, CapabilityExecutor, @@ -818,6 +822,54 @@ def test_missing_model_is_not_misclassified_as_a_package_dependency(self): json.dumps(raised.exception.to_dict()), ) + def test_model_download_failure_preserves_retry_details(self): + definition = CapabilityDefinition( + name="prepare-only", + description="Prepare a provider.", + extra="prepare-only", + operations={ + "noop": OperationDefinition( + input_model=SearchInput, + output_model=SearchResult, + requires_index=False, + ) + }, + prepares_models=True, + ) + failure = ModelArtifactDownloadError( + "prepare-only.embedding", + "publisher/model", + attempts=3, + reason="ConnectionError", + resumable=True, + retryable=True, + ) + plugin = CapabilityPlugin( + definition=definition, + executor_factory=lambda: CapabilityExecutor( + operations={"noop": Mock()}, + prepare=Mock(side_effect=failure), + ), + ) + registry = CapabilityRegistry((plugin,)) + registry.dependency_checks = Mock(return_value=()) + application, _ = self.application("unused", registry=registry) + + with self.assertRaises(ModelDownloadError) as raised: + application.prepare_models( + PrepareModelsCommand(modalities=("prepare-only",)) + ) + + payload = raised.exception.to_dict() + self.assertEqual(payload["code"], "model_download_failed") + self.assertTrue(payload["retryable"]) + self.assertEqual(payload["details"]["attempts"], 3) + self.assertTrue(payload["details"]["partial_files_preserved"]) + self.assertEqual( + payload["details"]["remediation"], + "vidxp prepare --modalities prepare-only", + ) + def test_unexpected_handler_error_is_not_misclassified(self): failure = RuntimeError("implementation bug") definition = CapabilityDefinition( diff --git a/tests/test_cli.py b/tests/test_cli.py index b3b4466..b819fe9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -25,6 +25,8 @@ CreateIndexCommand, DependencyCheckResult, DependencyKind, + ErrorCategory, + ErrorDetail, FusedSearchResult, FusionProvenance, IndexJobResult, @@ -145,6 +147,33 @@ def test_grouped_commands_are_exposed(self): ): self.assertIn(command, result.output) + def test_failed_model_preparation_job_is_structured_in_cli_json(self): + self.jobs.get.return_value = Job( + job_id=JOB_ID, + kind=JobKind.prepare_models, + state=JobState.failed, + queue=JobQueue.cpu, + error=ErrorDetail( + code="model_download_failed", + category=ErrorCategory.unavailable, + message="The model download failed after three attempts.", + details={ + "model": "publisher/model", + "partial_files_preserved": True, + "remediation": "vidxp prepare --modalities dialogue", + }, + retryable=True, + ), + ) + + result = self.invoke(["jobs", "show", JOB_ID]) + + self.assertEqual(result.exit_code, 0, result.output) + error = json.loads(result.output)["error"] + self.assertEqual(error["code"], "model_download_failed") + self.assertTrue(error["retryable"]) + self.assertTrue(error["details"]["partial_files_preserved"]) + def test_mcp_config_is_copy_paste_json_without_opening_repository(self): result = self.invoke(["mcp-config"]) @@ -813,10 +842,48 @@ def test_prepare_announces_start_and_subscribes_to_job_progress(self): self.assertIn("1.43 GiB", result.output) self.assertRegex( result.output, - r"\[\d{2}:\d{2}:\d{2}\] Starting model preparation for scene\.", + r"\[\d{2}:\d{2}:\d{2}\] Downloading and validating models for " + r"scene\.", ) self.assertTrue(callable(self.jobs.wait.call_args.kwargs["progress"])) + def test_prepare_distinguishes_cached_model_verification(self): + self.service.model_readiness.return_value = DependencyCheckResult( + ok=True, + modalities=("scene",), + checks=(), + ) + prepared = PrepareModelsResult( + prepared=("scene-model",), + modalities=("scene",), + runtime={ + "requested": "cpu", + "torch_device": "cpu", + "transcription_device": "cpu", + }, + ) + self.jobs.submit_prepare_models.return_value = Job( + job_id=JOB_ID, + kind=JobKind.prepare_models, + state=JobState.queued, + queue=JobQueue.cpu, + ) + self.jobs.wait.return_value = Job( + job_id=JOB_ID, + kind=JobKind.prepare_models, + state=JobState.succeeded, + queue=JobQueue.cpu, + result=PrepareModelsJobResult(result=prepared), + ) + + result = self.invoke(["prepare", "--modalities", "scene"]) + + self.assertEqual(result.exit_code, 0, result.output) + self.assertRegex( + result.output, + r"\[\d{2}:\d{2}:\d{2}\] Validating cached models for scene\.", + ) + def test_prepare_discloses_size_and_requires_confirmation(self): self.service.model_readiness.return_value = DependencyCheckResult( ok=False, diff --git a/tests/test_job_contracts.py b/tests/test_job_contracts.py index 817f05f..34f67aa 100644 --- a/tests/test_job_contracts.py +++ b/tests/test_job_contracts.py @@ -9,6 +9,8 @@ from vidxp.application_models import ( ActorOverlayJobRequest, ApplicationError, + ErrorCategory, + ErrorDetail, CreateActorOverlayCommand, CreateIndexCommand, Job, @@ -303,6 +305,45 @@ def test_job_backend_errors_are_normalized_for_every_adapter(self): service.get(JOB_ID) self.assertEqual(raised.exception.code, "job_backend_unavailable") + def test_failed_model_preparation_error_round_trips_through_job_service(self): + error = ErrorDetail( + code="model_download_failed", + category=ErrorCategory.unavailable, + message="The model download failed after three attempts.", + details={ + "capability": "dialogue.transcription", + "model": "publisher/model", + "attempts": 3, + "reason": "ConnectionError", + "partial_files_preserved": True, + "remediation": "vidxp prepare --modalities dialogue", + }, + retryable=True, + ) + backend = Mock() + backend.get.return_value = Job( + job_id=JOB_ID, + kind=JobKind.prepare_models, + state=JobState.failed, + queue=JobQueue.cpu, + error=error, + ) + service = JobService( + settings=VidXPSettings( + repository_root=Path("repository"), + runtime_backend="cpu", + ), + backend=backend, + ) + + with self.assertRaises(ApplicationError) as raised: + service.result(JOB_ID) + + self.assertEqual( + raised.exception.to_dict(), + error.model_dump(mode="json"), + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_mcp.py b/tests/test_mcp.py index eda1d88..39de1e7 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -24,6 +24,7 @@ ApplicationError, Artifact, ErrorCategory, + ErrorDetail, IndexStatus, Job, JobKind, @@ -401,6 +402,42 @@ async def test_discovery_tools_use_shared_media_and_job_pages(self): ) self.assertEqual(context.jobs.list.call_args.args[0].page_size, 9) + async def test_failed_model_preparation_job_is_structured_over_mcp(self): + with TemporaryDirectory() as directory: + context = self.context(Path(directory)) + context.jobs.get.return_value = Job( + job_id=JOB_ID, + kind=JobKind.prepare_models, + state=JobState.failed, + queue=JobQueue.cpu, + error=ErrorDetail( + code="model_download_failed", + category=ErrorCategory.unavailable, + message="The model download failed after three attempts.", + details={ + "model": "publisher/model", + "partial_files_preserved": True, + "remediation": "vidxp prepare --modalities dialogue", + }, + retryable=True, + ), + ) + server = create_mcp_server( + context, + default_principal=Principal( + subject="agent", + scopes=frozenset({"vidxp.read"}), + ), + ) + async with Client(server) as client: + result = await client.call_tool("get_job", {"job_id": JOB_ID}) + + self.assertFalse(result.is_error) + error = result.structured_content["error"] + self.assertEqual(error["code"], "model_download_failed") + self.assertTrue(error["retryable"]) + self.assertTrue(error["details"]["partial_files_preserved"]) + async def test_retry_job_uses_shared_stable_idempotency(self): with TemporaryDirectory() as directory: context = self.context(Path(directory)) diff --git a/tests/test_models.py b/tests/test_models.py index 74af819..c56923a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -37,6 +37,7 @@ SERVER_INDEX_RUNTIME_CHECKS, ) from vidxp.model_contracts import ( + ModelArtifactDownloadError, ModelArtifactUnavailableError, ModelKey, model_artifact_path, @@ -150,6 +151,39 @@ def test_actor_download_uses_pinned_media_object_not_lfs_pointer(self): "yunet", ) + def test_actor_download_retries_without_claiming_partial_resume(self): + with TemporaryDirectory() as directory: + spec = replace(YUNET_MODEL, filename="missing.onnx") + runtime = self.runtime( + directory, + allowed_specs=(spec,), + allow_model_downloads=True, + ) + retrieve = Mock(side_effect=ConnectionError("interrupted")) + events = [] + with patch.dict( + sys.modules, + {"pooch": types.SimpleNamespace(retrieve=retrieve)}, + ), patch("vidxp.runtime.sleep"), self.assertRaises( + ModelArtifactDownloadError + ) as raised: + runtime.resolve_artifact( + spec, + download=True, + progress=events.append, + ) + + self.assertEqual(retrieve.call_count, 3) + self.assertFalse(raised.exception.resumable) + self.assertTrue(raised.exception.retryable) + self.assertTrue( + any( + event["stage"] == "downloading_model" + and "restart from zero" in event["message"] + for event in events + ) + ) + def test_normal_model_resolution_never_downloads_implicitly(self): with TemporaryDirectory() as directory: spec = replace(YUNET_MODEL, filename="missing.onnx") @@ -203,6 +237,14 @@ def test_model_readiness_checks_the_pinned_cache_without_loading(self): def test_explicit_snapshot_download_reports_bytes(self): with TemporaryDirectory() as directory: snapshot = Path(directory) / "snapshot" + weights = snapshot / FASTER_WHISPER_MODEL.weights_file + weights.parent.mkdir(parents=True) + content = b"verified weights" + weights.write_bytes(content) + spec = replace( + FASTER_WHISPER_MODEL, + weights_sha256=hashlib.sha256(content).hexdigest(), + ) events = [] def download(**options): @@ -226,7 +268,7 @@ def download(**options): False, ): resolved = ModelRuntime._download_snapshot( - FASTER_WHISPER_MODEL, + spec, cache=Path(directory), progress=events.append, ) @@ -235,15 +277,197 @@ def download(**options): ) self.assertEqual(resolved, snapshot) + download_events = [ + event for event in events if event["stage"] == "downloading_model" + ] + self.assertTrue(download_events) + self.assertTrue( + all( + event["total"] == spec.download_size_bytes + for event in download_events + ) + ) self.assertTrue( any( event["stage"] == "downloading_model" - and event["current"] == 1024 - and event["total"] == 1024 + and event["current"] + == spec.download_size_bytes + and event["total"] + == spec.download_size_bytes for event in events ) ) + def test_snapshot_download_retries_and_resumes_partial_cache(self): + with TemporaryDirectory() as directory: + snapshot = Path(directory) / "snapshot" + weights = snapshot / FASTER_WHISPER_MODEL.weights_file + weights.parent.mkdir(parents=True) + content = b"verified weights" + weights.write_bytes(content) + spec = replace( + FASTER_WHISPER_MODEL, + weights_sha256=hashlib.sha256(content).hexdigest(), + ) + download = Mock( + side_effect=(ConnectionError("interrupted"), str(snapshot)) + ) + events = [] + + with patch( + "huggingface_hub.snapshot_download", + download, + ), patch("vidxp.runtime.sleep") as retry_sleep: + resolved = ModelRuntime._download_snapshot( + spec, + cache=Path(directory), + progress=events.append, + ) + + self.assertEqual(resolved, snapshot) + self.assertEqual(download.call_count, 2) + retry_sleep.assert_called_once_with(1) + self.assertTrue( + any( + event["stage"] == "downloading_model" + and "partial files will be resumed" in event["message"] + for event in events + ) + ) + + def test_snapshot_download_retries_an_incomplete_returned_snapshot(self): + with TemporaryDirectory() as directory: + snapshot = Path(directory) / "snapshot" + weights = snapshot / FASTER_WHISPER_MODEL.weights_file + content = b"verified weights" + spec = replace( + FASTER_WHISPER_MODEL, + weights_sha256=hashlib.sha256(content).hexdigest(), + ) + calls = 0 + + def download(**_options): + nonlocal calls + calls += 1 + if calls == 2: + weights.parent.mkdir(parents=True) + weights.write_bytes(content) + return str(snapshot) + + events = [] + with patch( + "huggingface_hub.snapshot_download", + side_effect=download, + ), patch("vidxp.runtime.sleep") as retry_sleep: + resolved = ModelRuntime._download_snapshot( + spec, + cache=Path(directory), + progress=events.append, + ) + + self.assertEqual(resolved, snapshot) + self.assertEqual(calls, 2) + retry_sleep.assert_called_once_with(1) + self.assertTrue( + any( + event["stage"] == "downloading_model" + and "partial files will be resumed" in event["message"] + for event in events + ) + ) + + def test_snapshot_download_reports_actionable_terminal_failure(self): + with TemporaryDirectory() as directory: + download = Mock(side_effect=ConnectionError("private detail")) + + with patch( + "huggingface_hub.snapshot_download", + download, + ), patch("vidxp.runtime.sleep"), self.assertRaises( + ModelArtifactDownloadError + ) as raised: + ModelRuntime._download_snapshot( + FASTER_WHISPER_MODEL, + cache=Path(directory), + progress=None, + ) + + self.assertEqual(download.call_count, 3) + self.assertEqual(raised.exception.attempts, 3) + self.assertEqual(raised.exception.reason, "ConnectionError") + self.assertTrue(raised.exception.resumable) + self.assertTrue(raised.exception.retryable) + self.assertNotIn("private detail", str(raised.exception)) + + def test_snapshot_download_does_not_mask_programming_errors(self): + with TemporaryDirectory() as directory: + download = Mock(side_effect=TypeError("implementation bug")) + + with patch( + "huggingface_hub.snapshot_download", + download, + ), self.assertRaisesRegex(TypeError, "implementation bug"): + ModelRuntime._download_snapshot( + FASTER_WHISPER_MODEL, + cache=Path(directory), + progress=None, + ) + + download.assert_called_once() + + def test_snapshot_download_does_not_retry_terminal_http_errors(self): + with TemporaryDirectory() as directory: + response = Mock(status_code=404) + failure = RuntimeError("not found") + failure.response = response + download = Mock(side_effect=failure) + + with patch( + "huggingface_hub.snapshot_download", + download, + ), self.assertRaises(ModelArtifactDownloadError) as raised: + ModelRuntime._download_snapshot( + FASTER_WHISPER_MODEL, + cache=Path(directory), + progress=None, + ) + + download.assert_called_once() + self.assertEqual(raised.exception.reason, "HTTP 404 RuntimeError") + self.assertFalse(raised.exception.retryable) + + def test_incomplete_cached_snapshot_is_resumed_during_prepare(self): + with TemporaryDirectory() as directory: + cache = Path(directory) / "models" + incomplete = Path(directory) / "incomplete" + complete = Path(directory) / "complete" + content = b"verified weights" + weights = complete / FASTER_WHISPER_MODEL.weights_file + weights.parent.mkdir(parents=True) + weights.write_bytes(content) + spec = replace( + FASTER_WHISPER_MODEL, + weights_sha256=hashlib.sha256(content).hexdigest(), + ) + runtime = self.runtime( + directory, + allowed_specs=(spec,), + allow_model_downloads=True, + ) + + with patch( + "huggingface_hub.snapshot_download", + return_value=str(incomplete), + ), patch.object( + ModelRuntime, + "_download_snapshot", + return_value=complete, + ) as resume: + resolved = runtime.resolve_model(spec, download=True) + + self.assertEqual(resolved, complete) + resume.assert_called_once_with(spec, cache=cache, progress=None) + def test_runtime_rejects_specs_not_declared_by_enabled_capabilities(self): with TemporaryDirectory() as directory: runtime = self.runtime(directory)