diff --git a/src/senselab/audio/tasks/forced_alignment/forced_alignment.py b/src/senselab/audio/tasks/forced_alignment/forced_alignment.py index feaaa02a8..1deff744e 100644 --- a/src/senselab/audio/tasks/forced_alignment/forced_alignment.py +++ b/src/senselab/audio/tasks/forced_alignment/forced_alignment.py @@ -25,6 +25,7 @@ ) from senselab.audio.tasks.preprocessing import extract_segments, pad_audios from senselab.utils.data_structures import DeviceType, HFModel, Language, ScriptLine, _select_device_and_dtype +from senselab.utils.dependencies import hf_offline_loading try: from nltk.tokenize.punkt import PunktParameters, PunktSentenceTokenizer @@ -654,14 +655,18 @@ def _load_mms_aligner( """ key = (MMS_MODEL_ID, iso3) if key not in cache: - processor = Wav2Vec2Processor.from_pretrained(MMS_MODEL_ID) - processor.tokenizer.set_target_lang(iso3) # type: ignore[attr-defined] - model = Wav2Vec2ForCTC.from_pretrained( - MMS_MODEL_ID, - target_lang=iso3, - ignore_mismatched_sizes=True, - ).to(device.value) # type: ignore[arg-type] - model.load_adapter(iso3) + # Cache-once (cross-node locked) then load the base weights + per-language + # adapter from the local cache only, so the many parallel alignment jobs + # don't 429 on the HF Hub revision check for facebook/mms-1b-all. + with hf_offline_loading(MMS_MODEL_ID): + processor = Wav2Vec2Processor.from_pretrained(MMS_MODEL_ID) + processor.tokenizer.set_target_lang(iso3) # type: ignore[attr-defined] + model = Wav2Vec2ForCTC.from_pretrained( + MMS_MODEL_ID, + target_lang=iso3, + ignore_mismatched_sizes=True, + ).to(device.value) # type: ignore[arg-type] + model.load_adapter(iso3) cache[key] = (processor, model) return cache[key] @@ -736,10 +741,11 @@ def align_transcriptions( # otherwise silently serve the first revision's weights. cache_key = f"{model_variant.path_or_uri}@{model_variant.revision or 'main'}" if cache_key not in loaded_processors_and_models: - processor = Wav2Vec2Processor.from_pretrained( - model_variant.path_or_uri, revision=model_variant.revision - ) - _w2v = Wav2Vec2ForCTC.from_pretrained(model_variant.path_or_uri, revision=model_variant.revision) + with hf_offline_loading(model_variant.path_or_uri, model_variant.revision or "main"): + processor = Wav2Vec2Processor.from_pretrained( + model_variant.path_or_uri, revision=model_variant.revision + ) + _w2v = Wav2Vec2ForCTC.from_pretrained(model_variant.path_or_uri, revision=model_variant.revision) model = _w2v.to(device.value) # type: ignore[arg-type] loaded_processors_and_models[cache_key] = (processor, model) processor, model = loaded_processors_and_models[cache_key] diff --git a/src/senselab/audio/tasks/speech_to_text/canary_qwen.py b/src/senselab/audio/tasks/speech_to_text/canary_qwen.py index 0c0840554..3adee997b 100644 --- a/src/senselab/audio/tasks/speech_to_text/canary_qwen.py +++ b/src/senselab/audio/tasks/speech_to_text/canary_qwen.py @@ -25,6 +25,7 @@ from senselab.audio.data_structures import Audio from senselab.utils.data_structures import DeviceType, HFModel, ScriptLine, _select_device_and_dtype +from senselab.utils.dependencies import hf_subprocess_env from senselab.utils.subprocess_venv import _clean_subprocess_env, ensure_venv, parse_subprocess_result, venv_python # Dedicated venv — kept separate from the existing nemo-diarization venv. @@ -174,7 +175,11 @@ def transcribe_with_canary_qwen( } ) - env = _clean_subprocess_env() + # Ensure the model is cached once (cross-node locked) and tell the + # child to load from the local cache only, so its SALM.from_pretrained + # skips the HF Hub revision check that 429s under many parallel jobs. + revision = model.revision if model is not None else "main" + env = hf_subprocess_env(model_name, revision or "main", base_env=_clean_subprocess_env()) result = subprocess.run( [python, "-c", _CANARY_WORKER_SCRIPT], input=input_json, diff --git a/src/senselab/audio/tasks/speech_to_text/granite.py b/src/senselab/audio/tasks/speech_to_text/granite.py index 6d9241bb2..23cd7fbc0 100644 --- a/src/senselab/audio/tasks/speech_to_text/granite.py +++ b/src/senselab/audio/tasks/speech_to_text/granite.py @@ -23,6 +23,7 @@ from senselab.audio.data_structures import Audio from senselab.utils.data_structures import DeviceType, HFModel, ScriptLine, _select_device_and_dtype +from senselab.utils.dependencies import hf_offline_loading class GraniteSpeechASR: @@ -81,8 +82,14 @@ def transcribe_with_granite( cache_key = f"{model_name}@{device_type.value}" if cache_key not in cls._cache: - processor = AutoProcessor.from_pretrained(model_name) - mdl = AutoModelForSpeechSeq2Seq.from_pretrained(model_name, dtype=dtype) + # Cache-once (cross-node locked) then load from the local cache only, + # so parallel jobs don't 429 on the HF Hub revision check. The same + # revision must be forwarded to both loaders, else offline mode looks + # for the wrong (default "main") snapshot and cache-misses. + revision = (model.revision if model is not None else None) or "main" + with hf_offline_loading(model_name, revision): + processor = AutoProcessor.from_pretrained(model_name, revision=revision) + mdl = AutoModelForSpeechSeq2Seq.from_pretrained(model_name, dtype=dtype, revision=revision) if device_type == DeviceType.CUDA and torch.cuda.is_available(): mdl = mdl.cuda() elif device_type == DeviceType.MPS and torch.backends.mps.is_available(): diff --git a/src/senselab/audio/tasks/speech_to_text/huggingface.py b/src/senselab/audio/tasks/speech_to_text/huggingface.py index e18d413b6..da4312364 100644 --- a/src/senselab/audio/tasks/speech_to_text/huggingface.py +++ b/src/senselab/audio/tasks/speech_to_text/huggingface.py @@ -6,6 +6,7 @@ """ import time +from functools import partial from typing import Any, Dict, List, Optional, Union, cast from transformers import Pipeline, pipeline @@ -20,6 +21,7 @@ ) from senselab.utils.data_structures.logging import logger from senselab.utils.data_structures.model import get_huggingface_token +from senselab.utils.dependencies import load_hf_resilient class HuggingFaceASR: @@ -72,10 +74,14 @@ def _get_hf_asr_pipeline( # When the caller asked for return_timestamps=False, the safest # path across all three is to OMIT the kwarg and let each # pipeline class default to its no-timestamps mode. + # ``revision`` is NOT put in pipeline_kwargs: ``load_hf_resilient`` + # reserves the ``revision`` kwarg for its cache step and does not + # forward it to the loader, so passing it both ways collides. Bind it + # onto the loader via ``partial`` and pass it to the cache step too. + revision = model.revision or "main" pipeline_kwargs: Dict[str, Any] = { "task": "automatic-speech-recognition", "model": model.path_or_uri, - "revision": model.revision, "max_new_tokens": max_new_tokens, "chunk_length_s": chunk_length_s, "batch_size": batch_size, @@ -84,9 +90,17 @@ def _get_hf_asr_pipeline( } if return_timestamps is not False: pipeline_kwargs["return_timestamps"] = return_timestamps + # Load via the resilient path: download-once (cross-node locked) then + # load from the local cache only, so parallel jobs don't 429 on the + # HF Hub revision check. No external env vars required. cls._pipelines[key] = cast( Pipeline, - pipeline(**pipeline_kwargs), # type: ignore[call-overload] + load_hf_resilient( + partial(pipeline, revision=revision), + repo_id=model.path_or_uri, + revision=revision, + **pipeline_kwargs, # type: ignore[arg-type] + ), ) return cls._pipelines[key] diff --git a/src/senselab/audio/tasks/speech_to_text/nemo.py b/src/senselab/audio/tasks/speech_to_text/nemo.py index d0e235030..23fabe610 100644 --- a/src/senselab/audio/tasks/speech_to_text/nemo.py +++ b/src/senselab/audio/tasks/speech_to_text/nemo.py @@ -15,6 +15,7 @@ from senselab.audio.data_structures import Audio from senselab.utils.data_structures import DeviceType, HFModel, ScriptLine, _select_device_and_dtype +from senselab.utils.dependencies import hf_subprocess_env from senselab.utils.subprocess_venv import _clean_subprocess_env, ensure_venv, parse_subprocess_result, venv_python # Reuse the same NeMo venv as diarization — it already has nemo_toolkit[asr] @@ -134,7 +135,11 @@ def transcribe_with_nemo( } ) - env = _clean_subprocess_env() + # Ensure the model is cached once (cross-node locked) and load the + # child offline so its from_pretrained skips the HF Hub revision + # check that 429s under many parallel jobs. + revision = model.revision if model is not None else "main" + env = hf_subprocess_env(model_name, revision or "main", base_env=_clean_subprocess_env()) result = subprocess.run( [python, "-c", _ASR_WORKER_SCRIPT], input=input_json, diff --git a/src/senselab/audio/tasks/speech_to_text/qwen.py b/src/senselab/audio/tasks/speech_to_text/qwen.py index 3018ae85d..042f30238 100644 --- a/src/senselab/audio/tasks/speech_to_text/qwen.py +++ b/src/senselab/audio/tasks/speech_to_text/qwen.py @@ -23,6 +23,7 @@ from senselab.audio.data_structures import Audio from senselab.utils.data_structures import DeviceType, HFModel, ScriptLine, _select_device_and_dtype +from senselab.utils.dependencies import hf_subprocess_env from senselab.utils.subprocess_venv import _clean_subprocess_env, ensure_venv, parse_subprocess_result, venv_python _QWEN_VENV = "qwen-asr" @@ -185,7 +186,14 @@ def transcribe_with_qwen( } ) - env = _clean_subprocess_env() + # Ensure the ASR model (and, when used, the companion forced aligner) + # are cached once (cross-node locked) and load the child offline so + # its from_pretrained skips the HF Hub revision check that 429s under + # many parallel jobs. Both Qwen3-ASR and Qwen3-ForcedAligner were + # observed rate-limiting, so both must be cached before going offline. + revision = model.revision if model is not None else "main" + also = [(aligner_name, "main")] if return_timestamps else None + env = hf_subprocess_env(model_name, revision or "main", also=also, base_env=_clean_subprocess_env()) result = subprocess.run( [python, "-c", _QWEN_WORKER_SCRIPT], input=input_json, diff --git a/src/senselab/utils/dependencies.py b/src/senselab/utils/dependencies.py index 6e07d7f34..30462d2f9 100644 --- a/src/senselab/utils/dependencies.py +++ b/src/senselab/utils/dependencies.py @@ -8,7 +8,7 @@ import time from functools import lru_cache from pathlib import Path -from typing import Callable, Iterator, Optional, TypeVar +from typing import Callable, Iterable, Iterator, Optional, TypeVar logger = logging.getLogger("senselab") @@ -343,6 +343,137 @@ def hf_local_files_only(repo_id: str, revision: str = "main") -> bool: return False +# Env vars that switch HuggingFace libraries to local-cache-only mode. Both are +# honored at call time (read fresh on each download/HEAD), so toggling them +# around a load reliably suppresses the network revision-check that 429s under +# many parallel jobs. +_HF_OFFLINE_ENV_VARS = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + +# ``HF_HUB_OFFLINE`` is process-global. Rather than hold a lock across the whole +# (potentially slow) model load — which would serialize even unrelated loads — +# ``hf_offline_loading`` uses a reference count: the first engaged loader captures +# the prior env and sets the offline vars; the last to exit restores them. This +# lets concurrent in-process loads run in parallel while keeping the global flag +# stable for all of them, and makes nested calls in one thread correct without +# reentrancy tricks. The lock guards only the (fast) counter/env transitions. +# (Global-state precedent: ``speechbrain_loading_cwd``.) +_hf_offline_lock = threading.Lock() +_hf_offline_depth = 0 +_hf_offline_saved_env: dict = {} + + +def hf_subprocess_env( + repo_id: "str | os.PathLike[str]", + revision: str = "main", + *, + also: Optional[Iterable[tuple[str, str]]] = None, + base_env: Optional[dict] = None, +) -> dict: + """Build an environment dict for a subprocess that loads HF model(s) offline. + + Ensures every referenced snapshot is present locally first (download-once + across processes/nodes via :func:`ensure_hf_model`), then returns a copy of + ``base_env`` (default ``os.environ``) with ``HF_HUB_OFFLINE`` / + ``TRANSFORMERS_OFFLINE`` set to ``"1"`` so the child's ``from_pretrained`` + skips the network revision check that rate-limits (429) under many parallel + jobs. + + The offline flag is only set when **all** referenced models are cached — if + any is missing, the env is returned unchanged so the child can still + download it online. Use ``also`` for workers that load more than one model + (e.g. Qwen3-ASR + its companion forced aligner). + + This is the subprocess analogue of :func:`hf_offline_loading`: in-process + loaders use the context manager, subprocess-venv workers (NeMo / Qwen / + Granite) inherit the offline flag through this env. No caller-set env vars + are required. + """ + env = dict(os.environ if base_env is None else base_env) + repos = [(os.fspath(repo_id), revision), *(also or [])] + if all(hf_local_files_only(rid, rev) for rid, rev in repos): + for var in _HF_OFFLINE_ENV_VARS: + env[var] = "1" + return env + + +@contextlib.contextmanager +def hf_offline_loading(repo_id: "str | os.PathLike[str]", revision: str = "main") -> Iterator[bool]: + """Force local-cache-only HuggingFace loading for the duration of the block. + + Ensures the snapshot is present first (download-once across processes/nodes + via :func:`ensure_hf_model`: cross-process heartbeat lock + retry-on-429), + then sets ``HF_HUB_OFFLINE`` / ``TRANSFORMERS_OFFLINE`` so ``from_pretrained`` + / ``pipeline`` make **no** network revision-check calls during the load — + the source of 429 storms when many jobs load the same model at once. + Requires **no** caller configuration. + + If the model cannot be cached (e.g. genuinely offline and never downloaded), + the env is left untouched and the block runs normally (online), so behavior + degrades gracefully rather than failing. + + Yields: + ``True`` if offline mode was engaged (model is cached), else ``False``. + """ + repo_id = os.fspath(repo_id) + if not hf_local_files_only(repo_id, revision): + yield False + return + global _hf_offline_depth + # First engaged loader captures the prior env and flips the offline vars; + # concurrent/nested loaders just bump the count and share the flag. + with _hf_offline_lock: + if _hf_offline_depth == 0: + _hf_offline_saved_env.clear() + _hf_offline_saved_env.update({var: os.environ.get(var) for var in _HF_OFFLINE_ENV_VARS}) + for var in _HF_OFFLINE_ENV_VARS: + os.environ[var] = "1" + _hf_offline_depth += 1 + try: + yield True + finally: + # Last loader out restores the env captured by the first one in. + with _hf_offline_lock: + _hf_offline_depth -= 1 + if _hf_offline_depth == 0: + for var, prev in _hf_offline_saved_env.items(): + if prev is None: + os.environ.pop(var, None) + else: + os.environ[var] = prev + _hf_offline_saved_env.clear() + + +def load_hf_resilient( + loader: Callable[..., _T], + *args: object, + repo_id: "str | os.PathLike[str]", + revision: str = "main", + **kwargs: object, +) -> _T: + """Load an HF model resiliently: cache-once, load local-only, retry on 429. + + Loader-agnostic wrapper for any in-process model constructor + (``transformers.pipeline``, ``*.from_pretrained``, SpeechBrain + ``from_hparams``, ...). It: + + 1. Ensures the snapshot is present exactly once across processes/nodes + (:func:`ensure_hf_model` — cross-process heartbeat lock + retry). + 2. Runs ``loader(*args, **kwargs)`` with HF libs in local-cache-only mode + (:func:`hf_offline_loading`) so no network revision-check fires. + 3. Retries the load on any residual transient error (5xx / 429 / timeout). + + ``repo_id`` / ``revision`` identify the model for the cache step and are + **not** forwarded to ``loader`` — pass the loader's own model/revision + arguments via ``*args`` / ``**kwargs`` as usual. + """ + + def _call() -> _T: + with hf_offline_loading(repo_id, revision): + return loader(*args, **kwargs) + + return retry_on_transient_error(_call) + + def _cached_error(cached: dict) -> Exception: """Reconstruct an exception from a cached error result.""" from huggingface_hub.errors import RepositoryNotFoundError, RevisionNotFoundError diff --git a/src/tests/audio/tasks/speech_to_text/granite_test.py b/src/tests/audio/tasks/speech_to_text/granite_test.py index 95e521fcc..865bbd7f8 100644 --- a/src/tests/audio/tasks/speech_to_text/granite_test.py +++ b/src/tests/audio/tasks/speech_to_text/granite_test.py @@ -18,7 +18,7 @@ from senselab.audio.data_structures import Audio from senselab.audio.tasks.preprocessing import downmix_audios_to_mono, resample_audios from senselab.audio.tasks.speech_to_text.granite import GraniteSpeechASR -from senselab.utils.data_structures import HFModel +from senselab.utils.data_structures import DeviceType, HFModel REPO_ROOT = Path(__file__).resolve().parents[5] FIXTURE_WAV = REPO_ROOT / "src" / "tests" / "data_for_testing" / "audio_48khz_mono_16bits.wav" @@ -36,6 +36,55 @@ def _load_16k_mono_fixture() -> Audio: return audio +def test_granite_forwards_revision_to_from_pretrained(monkeypatch: pytest.MonkeyPatch) -> None: + """The requested model revision reaches both ``from_pretrained`` calls. + + Regression: ``hf_offline_loading`` caches/locks the *requested* revision and + flips HF to offline, but the loaders were called without ``revision`` and so + defaulted to ``"main"`` — an offline cache-miss (or wrong version) for any + non-``main`` revision. + """ + import contextlib + from collections.abc import Iterator + + import transformers + + from senselab.audio.tasks.speech_to_text import granite as granite_mod + + recorded: dict = {} + + def _fake_processor(name: str, **kwargs: object) -> object: + recorded["processor_revision"] = kwargs.get("revision") + return object() + + def _fake_model(name: str, **kwargs: object) -> object: + recorded["model_revision"] = kwargs.get("revision") + return object() + + @contextlib.contextmanager + def _noop_offline(repo_id: object, revision: str = "main") -> Iterator[bool]: + yield True + + monkeypatch.setattr(transformers.AutoProcessor, "from_pretrained", _fake_processor) + monkeypatch.setattr(transformers.AutoModelForSpeechSeq2Seq, "from_pretrained", _fake_model) + monkeypatch.setattr(granite_mod, "hf_offline_loading", _noop_offline) + + class _FakeModel: + path_or_uri = "fake/granite-speech" + revision = "rev-xyz" + + granite_mod.GraniteSpeechASR._cache.clear() + result = granite_mod.GraniteSpeechASR.transcribe_with_granite( + audios=[], + model=_FakeModel(), # type: ignore[arg-type] + device=DeviceType.CPU, + ) + + assert result == [] # no audios -> no transcription, model just loaded + assert recorded["processor_revision"] == "rev-xyz" + assert recorded["model_revision"] == "rev-xyz" + + @pytest.mark.skipif( not granite_available, reason=f"ibm-granite/granite-speech-3.3-8b not in HF cache at {GRANITE_CACHE_DIR}", diff --git a/src/tests/audio/tasks/speech_to_text/huggingface_no_timestamps_test.py b/src/tests/audio/tasks/speech_to_text/huggingface_no_timestamps_test.py index f54676a7d..686288b41 100644 --- a/src/tests/audio/tasks/speech_to_text/huggingface_no_timestamps_test.py +++ b/src/tests/audio/tasks/speech_to_text/huggingface_no_timestamps_test.py @@ -44,6 +44,49 @@ def _load_16k_mono_fixture() -> Audio: return audio +def test_hf_pipeline_forwards_revision_without_collision(monkeypatch: pytest.MonkeyPatch) -> None: + """_get_hf_asr_pipeline forwards the revision to the loader without collision. + + Regression: ``revision`` was placed in ``pipeline_kwargs`` AND passed as the + reserved ``revision=`` argument of ``load_hf_resilient``, raising + ``TypeError: got multiple values for keyword argument 'revision'`` for every + HF-pipeline transcription. This test builds the pipeline with a fake loader + (no model download) and asserts the revision reaches it exactly once. + """ + from senselab.audio.tasks.speech_to_text import huggingface as hf_mod + from senselab.utils import dependencies + from senselab.utils.data_structures import DeviceType + + # Not cached -> load_hf_resilient just calls the loader once (no offline toggle). + monkeypatch.setattr(dependencies, "hf_local_files_only", lambda *a, **k: False) + + recorded: dict = {} + + def _fake_pipeline(**kwargs: object) -> object: + recorded.update(kwargs) + return object() + + monkeypatch.setattr(hf_mod, "pipeline", _fake_pipeline) + + class _FakeModel: + path_or_uri = "openai/whisper-tiny" + revision = "rev-xyz" + + hf_mod.HuggingFaceASR._pipelines.clear() + out = hf_mod.HuggingFaceASR._get_hf_asr_pipeline( + model=_FakeModel(), # type: ignore[arg-type] + return_timestamps=False, + max_new_tokens=128, + chunk_length_s=30, + batch_size=1, + device=DeviceType.CPU, + ) + + assert out is not None + assert recorded["revision"] == "rev-xyz" # forwarded to the loader exactly once + assert recorded["model"] == "openai/whisper-tiny" + + @pytest.mark.skipif( not ctc_available, reason=f"facebook/wav2vec2-base-960h not in HF cache at {CTC_CACHE_DIR}", diff --git a/src/tests/utils/dependencies_test.py b/src/tests/utils/dependencies_test.py new file mode 100644 index 000000000..8deea0828 --- /dev/null +++ b/src/tests/utils/dependencies_test.py @@ -0,0 +1,197 @@ +"""Unit tests for the HuggingFace-resilient loading helpers in ``dependencies``. + +Covers the offline-loading context, the subprocess env builder, and the +``load_hf_resilient`` wrapper. No network calls and no real model downloads: +``hf_local_files_only`` (the cache-ensuring step) is monkeypatched so the tests +exercise only the env-toggling / retry logic this module owns. +""" + +import os + +import pytest + +from senselab.utils import dependencies +from senselab.utils.dependencies import ( + hf_offline_loading, + hf_subprocess_env, + load_hf_resilient, + retry_on_transient_error, +) + +_OFFLINE_VARS = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + + +@pytest.fixture +def cached(monkeypatch: pytest.MonkeyPatch) -> None: + """Pretend every model is already cached (no download, no network).""" + monkeypatch.setattr(dependencies, "hf_local_files_only", lambda *a, **k: True) + + +@pytest.fixture +def not_cached(monkeypatch: pytest.MonkeyPatch) -> None: + """Pretend no model is cached / cannot be cached.""" + monkeypatch.setattr(dependencies, "hf_local_files_only", lambda *a, **k: False) + + +# ── hf_offline_loading ───────────────────────────────────────────── + + +def test_offline_loading_sets_and_restores_when_cached(cached: None, monkeypatch: pytest.MonkeyPatch) -> None: + """Inside the block the offline vars are "1"; afterwards the prior state is restored.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "preexisting") + + with hf_offline_loading("some/model") as engaged: + assert engaged is True + assert os.environ["HF_HUB_OFFLINE"] == "1" + assert os.environ["TRANSFORMERS_OFFLINE"] == "1" + + # Unset var goes back to unset; preexisting value is restored verbatim. + assert "HF_HUB_OFFLINE" not in os.environ + assert os.environ["TRANSFORMERS_OFFLINE"] == "preexisting" + + +def test_offline_loading_noop_when_not_cached(not_cached: None, monkeypatch: pytest.MonkeyPatch) -> None: + """If the model can't be cached, env is untouched and the block runs online.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + with hf_offline_loading("some/model") as engaged: + assert engaged is False + assert "HF_HUB_OFFLINE" not in os.environ + assert "HF_HUB_OFFLINE" not in os.environ + + +def test_offline_loading_restores_on_exception(cached: None, monkeypatch: pytest.MonkeyPatch) -> None: + """Env is restored even when the wrapped block raises.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + with pytest.raises(ValueError): + with hf_offline_loading("some/model"): + assert os.environ["HF_HUB_OFFLINE"] == "1" + raise ValueError("boom") + assert "HF_HUB_OFFLINE" not in os.environ + + +def test_offline_loading_nested_keeps_env_until_outermost_exit(cached: None, monkeypatch: pytest.MonkeyPatch) -> None: + """Nested blocks keep the offline vars set until the OUTERMOST block exits.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + + with hf_offline_loading("model_a") as engaged_a: + assert engaged_a is True + assert os.environ["HF_HUB_OFFLINE"] == "1" + with hf_offline_loading("model_b") as engaged_b: + assert engaged_b is True + assert os.environ["HF_HUB_OFFLINE"] == "1" + # Inner exit must NOT clear the flag while the outer block is still active. + assert os.environ["HF_HUB_OFFLINE"] == "1" + + assert "HF_HUB_OFFLINE" not in os.environ + + +def test_offline_loading_concurrent_does_not_serialize(cached: None, monkeypatch: pytest.MonkeyPatch) -> None: + """Two cached loads can be inside the block simultaneously (no whole-load serialization). + + A rendezvous barrier only releases if both threads are inside their + ``hf_offline_loading`` block at once. A design that holds a lock across the + whole block would keep the second thread out until the first exits, so the + barrier would time out (BrokenBarrierError). + """ + import threading + + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + + barrier = threading.Barrier(2, timeout=5) + errors: list = [] + + def worker() -> None: + try: + with hf_offline_loading("some/model") as engaged: + assert engaged is True + assert os.environ["HF_HUB_OFFLINE"] == "1" + barrier.wait() # both threads must reach here together + except Exception as exc: # BrokenBarrierError if the loads serialized + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"concurrent cached loads were serialized: {errors!r}" + assert "HF_HUB_OFFLINE" not in os.environ # env cleared once the last holder exits + + +# ── hf_subprocess_env ────────────────────────────────────────────── + + +def test_subprocess_env_offline_when_all_cached(cached: None) -> None: + """All referenced models cached → offline vars set in the returned env.""" + env = hf_subprocess_env("a/model", "main", also=[("b/aligner", "main")], base_env={"PATH": "/x"}) + assert env["PATH"] == "/x" # base_env preserved + for var in _OFFLINE_VARS: + assert env[var] == "1" + + +def test_subprocess_env_online_when_any_missing(monkeypatch: pytest.MonkeyPatch) -> None: + """If any referenced model is uncached, offline vars are NOT set (child may download).""" + # First repo cached, companion not. + calls = {"a/model": True, "b/aligner": False} + monkeypatch.setattr(dependencies, "hf_local_files_only", lambda rid, rev="main": calls[rid]) + env = hf_subprocess_env("a/model", "main", also=[("b/aligner", "main")], base_env={}) + for var in _OFFLINE_VARS: + assert var not in env + + +def test_subprocess_env_does_not_mutate_base(cached: None) -> None: + """The passed base_env dict is copied, not mutated in place.""" + base = {"PATH": "/x"} + hf_subprocess_env("a/model", base_env=base) + assert base == {"PATH": "/x"} + + +# ── load_hf_resilient ────────────────────────────────────────────── + + +def test_load_hf_resilient_returns_loader_result(not_cached: None) -> None: + """The wrapper forwards args/kwargs to the loader but consumes repo_id/revision.""" + sentinel = object() + + def loader(model_id: str, *, device: str) -> object: + assert model_id == "openai/whisper-tiny" + assert device == "cpu" + return sentinel + + # repo_id/revision steer the cache step and are NOT passed to the loader; + # the loader's own args go through *args / **kwargs. + out = load_hf_resilient(loader, "openai/whisper-tiny", repo_id="openai/whisper-tiny", revision="main", device="cpu") + assert out is sentinel + + +def test_load_hf_resilient_retries_transient(not_cached: None, monkeypatch: pytest.MonkeyPatch) -> None: + """A transient failure on first attempt is retried and then succeeds.""" + monkeypatch.setattr(dependencies.time, "sleep", lambda _s: None) # don't actually back off + attempts = {"n": 0} + + def flaky() -> str: + attempts["n"] += 1 + if attempts["n"] == 1: + raise TimeoutError("transient") + return "ok" + + out = load_hf_resilient(flaky, repo_id="x/y") + assert out == "ok" + assert attempts["n"] == 2 + + +def test_retry_reraises_non_transient(monkeypatch: pytest.MonkeyPatch) -> None: + """Non-transient errors (e.g. ValueError) are raised immediately, not retried.""" + attempts = {"n": 0} + + def boom() -> None: + attempts["n"] += 1 + raise ValueError("permanent") + + with pytest.raises(ValueError): + retry_on_transient_error(boom) + assert attempts["n"] == 1