fix(hf): load HF models offline-first to avoid Hub 429 rate-limiting

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]

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,

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():

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]
 

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,

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,

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