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,13 @@ 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 model
+            # id carries no explicit revision here, so the default "main" is used.
+            revision = model.revision if model is not None else "main"
+            with hf_offline_loading(model_name, revision or "main"):
+                processor = AutoProcessor.from_pretrained(model_name)
+                mdl = AutoModelForSpeechSeq2Seq.from_pretrained(model_name, dtype=dtype)
             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

@@ -20,6 +20,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:
@@ -84,9 +85,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(
+                    pipeline,
+                    repo_id=model.path_or_uri,
+                    revision=model.revision or "main",
+                    **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,122 @@ 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. ``hf_offline_loading`` serializes its
+# temporary toggle through this lock so concurrent loaders in the same process
+# don't observe each other's env mutation. Model *loads* (not inference) thus
+# run sequentially while the env is flipped — acceptable since loads are
+# one-time and cached. (Mirrors the global-state-lock precedent of
+# ``speechbrain_loading_cwd``.)
+_hf_offline_lock = threading.RLock()
+
+
+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
+    with _hf_offline_lock:
+        saved = {var: os.environ.get(var) for var in _HF_OFFLINE_ENV_VARS}
+        for var in _HF_OFFLINE_ENV_VARS:
+            os.environ[var] = "1"
+        try:
+            yield True
+        finally:
+            for var, prev in saved.items():
+                if prev is None:
+                    os.environ.pop(var, None)
+                else:
+                    os.environ[var] = prev
+
+
+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