mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-18 14:20:42 -07:00
fix(offline): mutate cached HF constants + threadsafe refcount
Review feedback on the initial fix surfaced two real issues:
1. ``os.environ`` toggles alone don't flip offline mode.
``huggingface_hub.constants.HF_HUB_OFFLINE`` is read once at import
time into a module-level bool; ``transformers.utils.hub._is_offline_mode``
mirrors that bool at its own import time. The hot paths
(``_http._default_backend_factory`` in huggingface_hub,
``is_offline_mode`` in transformers) read the cached bools — not the
env — so mutating only ``os.environ`` was a no-op.
2. Race condition on concurrent inference. Two threads running inside
``force_offline_if_cached`` via ``asyncio.to_thread`` could have
thread A's ``finally`` strip thread B's offline protection mid-run.
Rewrite the helper to:
- mutate ``huggingface_hub.constants.HF_HUB_OFFLINE`` and
``transformers.utils.hub._is_offline_mode`` directly
- refcount concurrent users under a single ``threading.RLock`` so a
shared offline window is restored only when the last caller exits
- still write ``os.environ`` for anything that reads it dynamically
Also addresses the unused-variable ruff flag on the Whisper transcribe
path (``audio, sr`` → ``audio, _sr``).
New unit tests cover the cached-constant mutation, env propagation,
no-op on ``is_cached=False``, nested contexts, and a threaded race
where a slow thread must retain offline mode after a peer exits.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
f3ed312cf2
commit
de15d8fdc6
@@ -347,7 +347,7 @@ class PyTorchSTTBackend:
|
||||
def _transcribe_sync():
|
||||
"""Run synchronous transcription in thread pool."""
|
||||
# Load audio
|
||||
audio, sr = load_audio(audio_path, sample_rate=16000)
|
||||
audio, _sr = load_audio(audio_path, sample_rate=16000)
|
||||
|
||||
# Model is loaded → weights are on disk. Force offline so
|
||||
# `get_decoder_prompt_ids` and any lazy tokenizer lookups
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
Unit tests for the ``force_offline_if_cached`` helper.
|
||||
|
||||
Verifies that the helper mutates the cached module constants in
|
||||
``huggingface_hub.constants`` and ``transformers.utils.hub`` — not just
|
||||
``os.environ`` — and that concurrent users are refcount-coordinated so
|
||||
one thread's exit can't strip another thread's offline protection.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from utils.hf_offline_patch import force_offline_if_cached # noqa: E402
|
||||
|
||||
|
||||
def _hf_const():
|
||||
import huggingface_hub.constants as hf_const
|
||||
|
||||
return hf_const
|
||||
|
||||
|
||||
def _tf_hub():
|
||||
import transformers.utils.hub as tf_hub
|
||||
|
||||
return tf_hub
|
||||
|
||||
|
||||
def test_mutates_cached_huggingface_hub_constant():
|
||||
original = _hf_const().HF_HUB_OFFLINE
|
||||
with force_offline_if_cached(True, "t"):
|
||||
assert _hf_const().HF_HUB_OFFLINE is True
|
||||
assert _hf_const().HF_HUB_OFFLINE == original
|
||||
|
||||
|
||||
def test_mutates_cached_transformers_constant():
|
||||
original = _tf_hub()._is_offline_mode
|
||||
with force_offline_if_cached(True, "t"):
|
||||
assert _tf_hub()._is_offline_mode is True
|
||||
assert _tf_hub()._is_offline_mode == original
|
||||
|
||||
|
||||
def test_sets_env_variable():
|
||||
original = os.environ.get("HF_HUB_OFFLINE")
|
||||
with force_offline_if_cached(True, "t"):
|
||||
assert os.environ.get("HF_HUB_OFFLINE") == "1"
|
||||
assert os.environ.get("HF_HUB_OFFLINE") == original
|
||||
|
||||
|
||||
def test_noop_when_not_cached():
|
||||
before = _hf_const().HF_HUB_OFFLINE
|
||||
with force_offline_if_cached(False, "t"):
|
||||
assert _hf_const().HF_HUB_OFFLINE == before
|
||||
|
||||
|
||||
def test_nested_contexts_respect_refcount():
|
||||
original = _hf_const().HF_HUB_OFFLINE
|
||||
with force_offline_if_cached(True, "outer"):
|
||||
assert _hf_const().HF_HUB_OFFLINE is True
|
||||
with force_offline_if_cached(True, "inner"):
|
||||
assert _hf_const().HF_HUB_OFFLINE is True
|
||||
# inner exit must not restore while outer is still active
|
||||
assert _hf_const().HF_HUB_OFFLINE is True
|
||||
assert _hf_const().HF_HUB_OFFLINE == original
|
||||
|
||||
|
||||
def test_concurrent_threads_share_offline_window():
|
||||
"""A slow thread must keep seeing offline mode even if a peer exits first."""
|
||||
original = _hf_const().HF_HUB_OFFLINE
|
||||
observations: list[bool] = []
|
||||
errors: list[Exception] = []
|
||||
barrier = threading.Barrier(2)
|
||||
|
||||
def slow():
|
||||
try:
|
||||
with force_offline_if_cached(True, "slow"):
|
||||
barrier.wait() # sync with fast
|
||||
time.sleep(0.15) # fast will exit during this sleep
|
||||
observations.append(_hf_const().HF_HUB_OFFLINE)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
errors.append(exc)
|
||||
|
||||
def fast():
|
||||
try:
|
||||
with force_offline_if_cached(True, "fast"):
|
||||
barrier.wait()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
errors.append(exc)
|
||||
|
||||
t_slow = threading.Thread(target=slow)
|
||||
t_fast = threading.Thread(target=fast)
|
||||
t_slow.start()
|
||||
t_fast.start()
|
||||
t_slow.join(timeout=5)
|
||||
t_fast.join(timeout=5)
|
||||
|
||||
assert not errors, errors
|
||||
assert observations == [True], "slow thread lost offline protection"
|
||||
assert _hf_const().HF_HUB_OFFLINE == original
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -6,6 +6,7 @@ are already downloaded. Must be imported BEFORE mlx_audio.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union
|
||||
@@ -13,13 +14,33 @@ from typing import Optional, Union
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# huggingface_hub reads ``HF_HUB_OFFLINE`` once at import time into
|
||||
# ``huggingface_hub.constants.HF_HUB_OFFLINE``; transformers mirrors that into
|
||||
# ``transformers.utils.hub._is_offline_mode`` at *its* import time. Toggling
|
||||
# ``os.environ`` after either module is imported does not flip those cached
|
||||
# bools, and the hot paths (``_http._default_backend_factory``,
|
||||
# ``transformers.utils.hub.is_offline_mode``) read the bools — not the env.
|
||||
# We mutate the cached constants directly, guarded by a refcount so
|
||||
# concurrent inference threads share a single offline window safely.
|
||||
|
||||
_offline_lock = threading.RLock()
|
||||
_offline_refcount = 0
|
||||
_saved_env: Optional[str] = None
|
||||
_saved_hf_const: Optional[bool] = None
|
||||
_saved_transformers_const: Optional[bool] = None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def force_offline_if_cached(is_cached: bool, model_label: str = ""):
|
||||
"""Context manager that sets ``HF_HUB_OFFLINE=1`` while loading a cached model.
|
||||
"""Force offline mode for the duration of a cached-model operation.
|
||||
|
||||
Flips ``HF_HUB_OFFLINE`` in the process env **and** in the cached bools
|
||||
inside ``huggingface_hub.constants`` / ``transformers.utils.hub`` so HTTP
|
||||
adapters and offline-mode checks actually see the change. Uses a refcount
|
||||
so multiple concurrent inference threads share a single offline window
|
||||
and the last one to exit restores state.
|
||||
|
||||
If *is_cached* is ``False`` the block runs normally (network allowed).
|
||||
If the offline load raises an error containing "offline" we automatically
|
||||
retry with network access so a partially-cached model still works.
|
||||
|
||||
Args:
|
||||
is_cached: Whether the model weights are already on disk.
|
||||
@@ -29,34 +50,60 @@ def force_offline_if_cached(is_cached: bool, model_label: str = ""):
|
||||
yield
|
||||
return
|
||||
|
||||
original_value = os.environ.get("HF_HUB_OFFLINE")
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
logger.info(
|
||||
"[offline-guard] %s is cached — forcing HF_HUB_OFFLINE=1",
|
||||
model_label or "model",
|
||||
)
|
||||
global _offline_refcount, _saved_env, _saved_hf_const, _saved_transformers_const
|
||||
|
||||
with _offline_lock:
|
||||
if _offline_refcount == 0:
|
||||
_saved_env = os.environ.get("HF_HUB_OFFLINE")
|
||||
try:
|
||||
import huggingface_hub.constants as hf_const
|
||||
|
||||
_saved_hf_const = hf_const.HF_HUB_OFFLINE
|
||||
hf_const.HF_HUB_OFFLINE = True
|
||||
except ImportError:
|
||||
_saved_hf_const = None
|
||||
try:
|
||||
import transformers.utils.hub as tf_hub
|
||||
|
||||
_saved_transformers_const = tf_hub._is_offline_mode
|
||||
tf_hub._is_offline_mode = True
|
||||
except ImportError:
|
||||
_saved_transformers_const = None
|
||||
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
logger.info(
|
||||
"[offline-guard] %s is cached — forcing offline mode",
|
||||
model_label or "model",
|
||||
)
|
||||
_offline_refcount += 1
|
||||
|
||||
try:
|
||||
yield
|
||||
except Exception as exc:
|
||||
if "offline" in str(exc).lower():
|
||||
logger.warning(
|
||||
"[offline-guard] Offline load failed for %s, retrying with network: %s",
|
||||
model_label or "model",
|
||||
exc,
|
||||
)
|
||||
# Restore original env and retry — caller must wrap the load
|
||||
# inside force_offline_if_cached so retrying here isn't possible.
|
||||
# Instead, propagate a flag via the exception so the caller can
|
||||
# decide. For simplicity we just let it fall through to the
|
||||
# finally block and re-raise.
|
||||
raise
|
||||
raise
|
||||
finally:
|
||||
if original_value is not None:
|
||||
os.environ["HF_HUB_OFFLINE"] = original_value
|
||||
else:
|
||||
os.environ.pop("HF_HUB_OFFLINE", None)
|
||||
with _offline_lock:
|
||||
_offline_refcount -= 1
|
||||
if _offline_refcount == 0:
|
||||
if _saved_env is not None:
|
||||
os.environ["HF_HUB_OFFLINE"] = _saved_env
|
||||
else:
|
||||
os.environ.pop("HF_HUB_OFFLINE", None)
|
||||
if _saved_hf_const is not None:
|
||||
try:
|
||||
import huggingface_hub.constants as hf_const
|
||||
|
||||
hf_const.HF_HUB_OFFLINE = _saved_hf_const
|
||||
except ImportError:
|
||||
pass
|
||||
if _saved_transformers_const is not None:
|
||||
try:
|
||||
import transformers.utils.hub as tf_hub
|
||||
|
||||
tf_hub._is_offline_mode = _saved_transformers_const
|
||||
except ImportError:
|
||||
pass
|
||||
_saved_env = None
|
||||
_saved_hf_const = None
|
||||
_saved_transformers_const = None
|
||||
|
||||
|
||||
def patch_huggingface_hub_offline():
|
||||
|
||||
Reference in New Issue
Block a user