diff --git a/aTrain_core/load_resources.py b/aTrain_core/load_resources.py index 0c4646de..5895fa0b 100644 --- a/aTrain_core/load_resources.py +++ b/aTrain_core/load_resources.py @@ -1,6 +1,7 @@ import json import os import shutil +import ssl from functools import partial from importlib.resources import files from multiprocessing.managers import DictProxy @@ -12,6 +13,14 @@ from aTrain_core.globals import MODELS_DIR, REQUIRED_MODELS, REQUIRED_MODELS_DIR from aTrain_core.integrity import ModelIntegrityError, find_missing_files, verify_model +OFFLINE_INSTALL_DOCS = ( + "https://github.com/aTrainTranscription/aTrain/blob/develop/docs/offline-model-install.md" +) + + +class ModelDownloadError(Exception): + """A model could not be downloaded, with a hint on what to do about it.""" + class custom_tqdm(tqdm): def __init__(self, progress: DictProxy, total: float, *args, **kwargs): @@ -51,11 +60,42 @@ def download_model(model_path: Path, model_info: dict, progress: DictProxy | Non local_dir=model_path, max_workers=1, ) + except Exception as error: + if _is_certificate_error(error): + raise ModelDownloadError( + f"Could not download model '{model_path.name}': the server certificate " + "could not be verified. A proxy in your network most likely inspects " + "HTTPS traffic with its own certificate. Ask your IT department, or " + f"install the model manually: {OFFLINE_INSTALL_DOCS}" + ) from error + raise finally: file_download.http_get = original_http_get file_download.xet_get = original_xet_get +def _is_certificate_error(error: BaseException) -> bool: + """True if a certificate check failed anywhere in the exception chain. + + The HTTP clients wrap it several layers deep (requests and urllib3 today, + httpx with huggingface_hub 1.x), so the chain is walked instead of + matching one client's exception type. urllib3 keeps the original error in + `args` and `reason` rather than chaining it, so those are followed too. + """ + seen = set() + pending: list[BaseException] = [error] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, ssl.SSLCertVerificationError): + return True + linked = [current.__cause__, current.__context__, getattr(current, "reason", None)] + pending += [e for e in [*linked, *current.args] if isinstance(e, BaseException)] + return False + + def get_model(model: str, progress: DictProxy | None = None) -> Path: """Loads a specific model, downloading and verifying it if necessary.""" models_config = load_model_config_file() diff --git a/tests/unit/test_load_resources.py b/tests/unit/test_load_resources.py index f18bc34a..707b108f 100644 --- a/tests/unit/test_load_resources.py +++ b/tests/unit/test_load_resources.py @@ -75,3 +75,43 @@ def test_hub_download_functions_accept_the_progress_bar(name): # download_model passes `_tqdm_bar` to both; a Hub upgrade that drops it # would only fail at download time. assert "_tqdm_bar" in inspect.signature(getattr(file_download, name)).parameters + + +def _certificate_error(): + """The shape of the chain requests raises when a proxy re-signs TLS. + + requests.SSLError has MaxRetryError as its context, whose `reason` is a + urllib3 SSLError that holds the ssl error in `args` instead of chaining + it. Built from stdlib exceptions, since the unit job installs neither + requests nor urllib3. + """ + import ssl + + urllib3_error = OSError(ssl.SSLCertVerificationError(1, "unable to get local issuer")) + max_retry = OSError("Max retries exceeded") + max_retry.reason = urllib3_error + requests_error = OSError(max_retry) + requests_error.__context__ = max_retry + return requests_error + + +def test_certificate_error_points_to_manual_install(tmp_path, monkeypatch): + original = file_download.http_get + + def fail(**kw): + raise _certificate_error() + + monkeypatch.setattr(load_resources, "snapshot_download", fail) + + with pytest.raises(load_resources.ModelDownloadError) as caught: + load_resources.download_model( + tmp_path / "large-v3-turbo", MODEL_INFO, progress={"current": 0, "total": 0} + ) + + assert "'large-v3-turbo'" in str(caught.value) + assert load_resources.OFFLINE_INSTALL_DOCS in str(caught.value) + assert file_download.http_get is original + + +def test_other_errors_are_not_reported_as_certificate_errors(): + assert not load_resources._is_certificate_error(OSError("connection lost"))