diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 144aa1fd37..47917004f9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -527,29 +527,63 @@ def _hf_env_offline() -> bool: return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} +def _hf_unreachable() -> bool: + """True when the Hub can't be reached: dead DNS, or DNS that resolves with no egress. + + DNS alone misses the common offline shapes (WAN down behind a live router, captive + portal, stale DNS cache), leaving every hub call to burn its full retry backoff. + """ + if _probe_dns_dead(): + return True + try: + from utils.utils import hf_unreachable + return hf_unreachable() + except Exception: + return False + + @contextlib.contextmanager -def _hf_offline_if_dns_dead(): - """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; +def _hf_offline_if_unreachable(): + """Set HF_HUB_OFFLINE for this block only when the Hub is unreachable; restores env on exit so a transient hiccup can't quarantine the process. No-op if the user already set it.""" if "HF_HUB_OFFLINE" in os.environ: yield False return - if not _probe_dns_dead(): + if not _hf_unreachable(): yield False return + logger.warning("Hugging Face endpoint unreachable; using local HF cache for this load.") + + # Env vars alone do not take effect mid-process: huggingface_hub and transformers read + # their offline constants at import, so the calls below would still retry. + force_ctx = None + try: + from utils.utils import force_hf_offline + + force_ctx = force_hf_offline() + force_ctx.__enter__() + except Exception: + force_ctx = None + transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ - os.environ["HF_HUB_OFFLINE"] = "1" - if not transformers_was_set: - os.environ["TRANSFORMERS_OFFLINE"] = "1" - logger.warning("huggingface.co unreachable; using local HF cache for this load.") + if force_ctx is None: + os.environ["HF_HUB_OFFLINE"] = "1" + if not transformers_was_set: + os.environ["TRANSFORMERS_OFFLINE"] = "1" try: yield True finally: - os.environ.pop("HF_HUB_OFFLINE", None) - if not transformers_was_set: - os.environ.pop("TRANSFORMERS_OFFLINE", None) + if force_ctx is not None: + try: + force_ctx.__exit__(None, None, None) + except Exception: + pass + else: + os.environ.pop("HF_HUB_OFFLINE", None) + if not transformers_was_set: + os.environ.pop("TRANSFORMERS_OFFLINE", None) try: @@ -6889,7 +6923,7 @@ class LlamaCppBackend: hf_repo, ) hf_repo = _resolved_repo - with _hf_offline_if_dns_dead(): + with _hf_offline_if_unreachable(): _preflight_model_path = self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, @@ -6931,7 +6965,7 @@ class LlamaCppBackend: hf_repo, ) hf_repo = _resolved_repo - with _hf_offline_if_dns_dead(): + with _hf_offline_if_unreachable(): model_path = _preflight_model_path or self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index b858fe6f17..daee45950f 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -68,7 +68,7 @@ import pandas as pd from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset -from core.inference.llama_cpp import _hf_offline_if_dns_dead +from core.inference.llama_cpp import _hf_offline_if_unreachable from utils.models import is_vision_model, detect_audio_type from utils.models.model_config import _env_offline from utils.datasets import format_and_template_dataset diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index baf6329dae..23b5c4a371 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2344,7 +2344,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> flush = True, ) - # Offline auto-detect: skip ~25s of HF retries per call when DNS is dead. + # Offline auto-detect: skip ~25s of HF retries per call when the hub is unreachable. if "HF_HUB_OFFLINE" not in os.environ: import socket as _socket import threading as _threading @@ -2362,13 +2362,24 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> _t = _threading.Thread(target = _probe, daemon = True) _t.start() _t.join(2.0) + if _result[0] is False: + # DNS answers even when there is no egress (WAN down, captive portal), so + # confirm with the bounded, proxy-aware reachability probe. HF_ENDPOINT aware. + try: + from utils.transformers_version import hf_endpoint_unreachable + from utils.utils import hf_probe_disabled + + if not hf_probe_disabled() and hf_endpoint_unreachable(): + _result[0] = True + except Exception: + pass if _result[0] is None or _result[0] is True: os.environ["HF_HUB_OFFLINE"] = "1" os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") os.environ.setdefault("HF_DATASETS_OFFLINE", "1") # logger isn't configured yet; print to stderr instead. print( - "huggingface.co unreachable; HF_HUB_OFFLINE=1 set for this worker.", + "Hugging Face endpoint unreachable; HF_HUB_OFFLINE=1 set for this worker.", file = sys.stderr, flush = True, ) diff --git a/studio/backend/picker/routes/templates.py b/studio/backend/picker/routes/templates.py index 02b8bf7184..a070c8d054 100644 --- a/studio/backend/picker/routes/templates.py +++ b/studio/backend/picker/routes/templates.py @@ -37,9 +37,15 @@ async def get_default_chat_template_route( hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ) -> ModelTemplateResponse: - template = await asyncio.to_thread( - read_default_chat_template, model_name, hf_token, gguf_variant - ) + # Cached repos resolve from disk, but a cache miss falls through to the hub and + # offline that costs one retry backoff per candidate template file. + from core.inference.llama_cpp import _hf_offline_if_unreachable + + def _read(): + with _hf_offline_if_unreachable(): + return read_default_chat_template(model_name, hf_token, gguf_variant) + + template = await asyncio.to_thread(_read) if template is not None and len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: template = None return ModelTemplateResponse(model_name = model_name, chat_template = template) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d0a2d97f74..b317133d5c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1007,7 +1007,7 @@ try: _canonicalize_spec_mode, _extra_args_n_ubatch, _extra_args_set_spec_type, - _hf_offline_if_dns_dead, + _hf_offline_if_unreachable, _kv_bytes_per_elem, _kv_unified_from_args, _planned_main_cache_types, @@ -1051,7 +1051,7 @@ except ImportError: _canonicalize_spec_mode, _extra_args_n_ubatch, _extra_args_set_spec_type, - _hf_offline_if_dns_dead, + _hf_offline_if_unreachable, _kv_bytes_per_elem, _kv_unified_from_args, _planned_main_cache_types, @@ -3672,7 +3672,10 @@ def _target_is_vision(load_path: str) -> bool: # paths, where the token is unused, but the rule requires it regardless). from utils.models.model_config import is_vision_model try: - return bool(is_vision_model(load_path, hf_token = os.environ.get("HF_TOKEN"))) + # Guarded: this runs per request, so an unreachable hub would re-pay its retry + # backoff on every image/audio call. + with _hf_offline_if_unreachable(): + return bool(is_vision_model(load_path, hf_token = os.environ.get("HF_TOKEN"))) except Exception as exc: # Detection failure: don't block the swap, let the load decide. logger.debug("auto-switch: vision probe failed for %s: %s", load_path, exc) @@ -5467,7 +5470,7 @@ async def _load_model_impl( # is_lora auto-detected from adapter_config.json on disk/HF. # DNS-probe wrap so offline loads skip 30-60s of soft-failed network # checks before the worker starts. - with _hf_offline_if_dns_dead(): + with _hf_offline_if_unreachable(): config = ModelConfig.from_identifier( model_id = model_identifier, hf_token = request.hf_token, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 96c5b96d73..cc2a402b59 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1855,75 +1855,80 @@ async def get_model_config( ): """Get configuration for a specific model (wraps load_model_defaults).""" hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) + from core.inference.llama_cpp import _hf_offline_if_unreachable + try: - if not is_local_path(model_name): - resolved = resolve_cached_repo_id_case(model_name) - if resolved != model_name: - logger.info( - "Using cached repo_id casing '%s' for requested '%s'", - resolved, - model_name, - ) - model_name = resolved + # Each probe below can reach the hub, so the guard wraps the whole handler: + # offline they must all resolve from the HF cache instead of retrying. + with _hf_offline_if_unreachable(): + if not is_local_path(model_name): + resolved = resolve_cached_repo_id_case(model_name) + if resolved != model_name: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved, + model_name, + ) + model_name = resolved - logger.info(f"Getting model config for: {model_name}") - from utils.models.model_config import detect_audio_type + logger.info(f"Getting model config for: {model_name}") + from utils.models.model_config import detect_audio_type - config_dict = load_model_defaults(model_name) + config_dict = load_model_defaults(model_name) - # Detect capabilities (HF token for gated models). - is_vision = is_vision_model(model_name, hf_token = hf_token) - is_embedding = is_embedding_model(model_name, hf_token = hf_token) - audio_type = detect_audio_type(model_name, hf_token = hf_token) + # Detect capabilities (HF token for gated models). + is_vision = is_vision_model(model_name, hf_token = hf_token) + is_embedding = is_embedding_model(model_name, hf_token = hf_token) + audio_type = detect_audio_type(model_name, hf_token = hf_token) - is_lora = False - base_model = None - max_position_embeddings = None - try: - model_config = ModelConfig.from_identifier(model_name) - is_lora = model_config.is_lora - base_model = model_config.base_model if is_lora else None - max_position_embeddings = _get_max_position_embeddings(model_config) - except Exception: - pass - - # Fallback: read raw config.json (declarative fields only) -- a selection-time - # metadata probe that must never execute a repo's auto_map Python. - if max_position_embeddings is None: + is_lora = False + base_model = None + max_position_embeddings = None try: - from utils.transformers_version import _load_config_json - from types import SimpleNamespace - - _cfg = _load_config_json(model_name, hf_token = hf_token) - if _cfg is not None: - - def _to_ns(d): - if isinstance(d, dict): - return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()}) - return d - - max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg)) + model_config = ModelConfig.from_identifier(model_name) + is_lora = model_config.is_lora + base_model = model_config.base_model if is_lora else None + max_position_embeddings = _get_max_position_embeddings(model_config) except Exception: pass - logger.info( - f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}, max_position_embeddings={max_position_embeddings}" - ) - return ModelDetails( - id = model_name, - model_name = model_name, - config = config_dict, - is_vision = is_vision, - is_embedding = is_embedding, - is_lora = is_lora, - is_audio = audio_type is not None, - audio_type = audio_type, - has_audio_input = is_audio_input_type(audio_type), - model_type = derive_model_type(is_vision, audio_type, is_embedding), - base_model = base_model, - max_position_embeddings = max_position_embeddings, - model_size_bytes = _get_model_size_bytes(model_name, hf_token), - ) + # Fallback: read raw config.json (declarative fields only) -- a selection-time + # metadata probe that must never execute a repo's auto_map Python. + if max_position_embeddings is None: + try: + from utils.transformers_version import _load_config_json + from types import SimpleNamespace + + _cfg = _load_config_json(model_name, hf_token = hf_token) + if _cfg is not None: + + def _to_ns(d): + if isinstance(d, dict): + return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()}) + return d + + max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg)) + except Exception: + pass + + logger.info( + f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}, max_position_embeddings={max_position_embeddings}" + ) + return ModelDetails( + id = model_name, + model_name = model_name, + config = config_dict, + is_vision = is_vision, + is_embedding = is_embedding, + is_lora = is_lora, + is_audio = audio_type is not None, + audio_type = audio_type, + has_audio_input = is_audio_input_type(audio_type), + model_type = derive_model_type(is_vision, audio_type, is_embedding), + base_model = base_model, + max_position_embeddings = max_position_embeddings, + model_size_bytes = _get_model_size_bytes(model_name, hf_token), + ) except Exception as e: raise log_and_http_error( @@ -2585,7 +2590,11 @@ async def check_vision_model( try: logger.info(f"Checking if vision model: {model_name}") # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). - is_vision = is_vision_model(model_name, hf_token = hf_token) + # Offline the guard keeps this on the HF cache instead of retrying the hub. + from core.inference.llama_cpp import _hf_offline_if_unreachable + + with _hf_offline_if_unreachable(): + is_vision = is_vision_model(model_name, hf_token = hf_token) logger.info(f"Vision check result for {model_name}: is_vision={is_vision}") return VisionCheckResponse( diff --git a/studio/backend/tests/test_active_generations.py b/studio/backend/tests/test_active_generations.py index aa087fe4ea..2f86f8efdc 100644 --- a/studio/backend/tests/test_active_generations.py +++ b/studio/backend/tests/test_active_generations.py @@ -344,7 +344,7 @@ def test_a_forced_load_that_fails_preflight_leaves_the_chats_alone(monkeypatch): from models.inference import LoadRequest inf_mod = _stub_load_route(monkeypatch, active_model_name = "org/OTHER") - monkeypatch.setattr(inf_mod, "_hf_offline_if_dns_dead", contextlib.nullcontext) + monkeypatch.setattr(inf_mod, "_hf_offline_if_unreachable", contextlib.nullcontext) # Stands in for any preflight refusal; a None here is the route's own 400. monkeypatch.setattr(inf_mod.ModelConfig, "from_identifier", staticmethod(lambda **kwargs: None)) @@ -375,7 +375,7 @@ def _stub_standard_load_route(monkeypatch): _stub_load_route(monkeypatch, active_model_name = "org/OTHER") # _stub_load_route neutralises the sidecar guard; this test is about it. monkeypatch.setattr(inf_mod, "_raise_if_sidecar_swap_in_progress", real_sidecar_check) - monkeypatch.setattr(inf_mod, "_hf_offline_if_dns_dead", contextlib.nullcontext) + monkeypatch.setattr(inf_mod, "_hf_offline_if_unreachable", contextlib.nullcontext) monkeypatch.setattr(inf_mod, "_mlx_distributed_launch_detected", lambda: False) monkeypatch.setattr( inf_mod.ModelConfig, @@ -1582,7 +1582,7 @@ def test_a_forced_load_that_loses_to_a_sidecar_install_leaves_the_chats_alone(mo from models.inference import LoadRequest inf_mod = _stub_load_route(monkeypatch, active_model_name = "org/OTHER") - monkeypatch.setattr(inf_mod, "_hf_offline_if_dns_dead", contextlib.nullcontext) + monkeypatch.setattr(inf_mod, "_hf_offline_if_unreachable", contextlib.nullcontext) monkeypatch.setattr( inf_mod.ModelConfig, "from_identifier", diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 6ec9c44e88..1ead0afe4b 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1163,7 +1163,7 @@ class TestLoadModelGuardIntegration(unittest.TestCase): patch.object(self.route, "resolve_effective_chat_template_override", return_value = None), patch.object(self.route, "get_inference_backend", return_value = inf), patch.object(self.route, "get_llama_cpp_backend", return_value = llama), - patch.object(self.route, "_hf_offline_if_dns_dead", lambda: contextlib.nullcontext()), + patch.object(self.route, "_hf_offline_if_unreachable", lambda: contextlib.nullcontext()), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), _stub_guard_deps(training_active = True, decision = (False, info)), ): diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index ccbe50bcb9..8b68af885a 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -914,7 +914,7 @@ class TestLoadHubDownloadExclusion: patch.object(route, "_resolve_gguf_gpu_ids_for_request", _no_gguf_gpu_ids), patch.object(route, "_guard_chat_load_against_training", return_value = None), patch.object(route, "_effective_load_in_4bit", return_value = False), - patch.object(route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(route, "_hf_offline_if_unreachable", nullcontext), patch.object(route.asyncio, "to_thread", new = _inline_to_thread), patch.object(llama_cpp_module, "_hub_download_blocks_gguf_load", _fake_blocks), ): diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index 43365bd3ca..354ce1dccc 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -737,7 +737,7 @@ def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch): ) monkeypatch.setattr( llama_cpp_module, - "_hf_offline_if_dns_dead", + "_hf_offline_if_unreachable", lambda: __import__("contextlib").nullcontext(), ) @@ -786,7 +786,7 @@ def test_remote_vulkan_preflight_download_failure_keeps_active_server(monkeypatc monkeypatch.setattr(llama_cpp_module, "_resolve_repo_id_casing", lambda repo: repo) monkeypatch.setattr( llama_cpp_module, - "_hf_offline_if_dns_dead", + "_hf_offline_if_unreachable", lambda: __import__("contextlib").nullcontext(), ) diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 7999eb4f73..93dff5cbd5 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -1025,7 +1025,7 @@ class TestRouteErrors(unittest.TestCase): return_value = None, ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), - patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(inference_route, "_hf_offline_if_unreachable", nullcontext), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -1095,7 +1095,7 @@ class TestRouteErrors(unittest.TestCase): return_value = None, ) as training_guard, patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), - patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(inference_route, "_hf_offline_if_unreachable", nullcontext), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -1194,7 +1194,7 @@ class TestRouteErrors(unittest.TestCase): return_value = None, ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), - patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(inference_route, "_hf_offline_if_unreachable", nullcontext), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -1340,7 +1340,7 @@ class TestRouteErrors(unittest.TestCase): return_value = None, ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), - patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(inference_route, "_hf_offline_if_unreachable", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -1411,7 +1411,7 @@ class TestRouteErrors(unittest.TestCase): return_value = None, ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), - patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object(inference_route, "_hf_offline_if_unreachable", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index d1e61d0546..be9c92eb3a 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -81,7 +81,7 @@ from core.inference.llama_cpp import ( LlamaCppBackend, _cached_colocated_split_main, _gguf_files_for_variant, - _hf_offline_if_dns_dead, + _hf_offline_if_unreachable, _probe_dns_dead, _resolve_repo_id_casing, ) @@ -871,7 +871,7 @@ class TestDetectGgufModelRemoteOffline: # --------------------------------------------------------------------------- -# _probe_dns_dead / _hf_offline_if_dns_dead +# _probe_dns_dead / _hf_offline_if_unreachable # --------------------------------------------------------------------------- @@ -900,6 +900,18 @@ def dns(monkeypatch): return _DnsState(monkeypatch) +@pytest.fixture +def reachable(monkeypatch): + """Control the endpoint reachability probe; defaults to reachable so no test hits the network.""" + import utils.utils as _utils + + def _set(unreachable: bool): + monkeypatch.setattr(_utils, "hf_unreachable", lambda *a, **k: unreachable) + + _set(False) + return _set + + class TestProbeDnsDead: def test_returns_false_on_success(self, dns): dns.ok() @@ -919,11 +931,11 @@ class TestProbeDnsDead: socket.setdefaulttimeout(None) -class TestHfOfflineIfDnsDead: - def test_dns_fail_sets_env_inside_block_only(self, dns, clean_offline_env): +class TestHfOfflineIfUnreachable: + def test_dns_fail_sets_env_inside_block_only(self, dns, reachable, clean_offline_env): dns.fail() assert "HF_HUB_OFFLINE" not in os.environ - with _hf_offline_if_dns_dead() as did_set: + with _hf_offline_if_unreachable() as did_set: assert did_set is True assert os.environ.get("HF_HUB_OFFLINE") == "1" assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" @@ -931,38 +943,54 @@ class TestHfOfflineIfDnsDead: assert "HF_HUB_OFFLINE" not in os.environ assert "TRANSFORMERS_OFFLINE" not in os.environ - def test_dns_ok_is_noop(self, dns, clean_offline_env): + def test_dns_ok_and_reachable_is_noop(self, dns, reachable, clean_offline_env): dns.ok() - with _hf_offline_if_dns_dead() as did_set: + with _hf_offline_if_unreachable() as did_set: assert did_set is False assert "HF_HUB_OFFLINE" not in os.environ - def test_dns_recovers_between_calls(self, dns, clean_offline_env): + def test_dns_ok_but_endpoint_unreachable_engages(self, dns, reachable, clean_offline_env): + # WAN down behind a live router: DNS answers, egress does not. + dns.ok() + reachable(True) + with _hf_offline_if_unreachable() as did_set: + assert did_set is True + assert os.environ.get("HF_HUB_OFFLINE") == "1" + assert "HF_HUB_OFFLINE" not in os.environ + + def test_probe_opt_out_keeps_dns_only_behaviour(self, dns, clean_offline_env, monkeypatch): + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "0") + dns.ok() + with _hf_offline_if_unreachable() as did_set: + assert did_set is False + assert "HF_HUB_OFFLINE" not in os.environ + + def test_dns_recovers_between_calls(self, dns, reachable, clean_offline_env): # First call: DNS dead -> env set inside, cleared on exit. dns.fail() - with _hf_offline_if_dns_dead(): + with _hf_offline_if_unreachable(): pass assert "HF_HUB_OFFLINE" not in os.environ # Second call: DNS healthy -> no env mutation. dns.ok() - with _hf_offline_if_dns_dead() as did_set: + with _hf_offline_if_unreachable() as did_set: assert did_set is False assert "HF_HUB_OFFLINE" not in os.environ - def test_user_set_hf_hub_offline_is_preserved(self, dns, clean_offline_env, monkeypatch): + def test_user_set_hf_hub_offline_is_preserved(self, dns, reachable, clean_offline_env, monkeypatch): # User explicitly set offline before launching Unsloth. monkeypatch.setenv("HF_HUB_OFFLINE", "1") dns.fail() - with _hf_offline_if_dns_dead() as did_set: + with _hf_offline_if_unreachable() as did_set: assert did_set is False assert os.environ.get("HF_HUB_OFFLINE") == "1" # Helper must not pop a variable it did not set. assert os.environ.get("HF_HUB_OFFLINE") == "1" - def test_user_set_transformers_offline_is_preserved(self, dns, clean_offline_env, monkeypatch): + def test_user_set_transformers_offline_is_preserved(self, dns, reachable, clean_offline_env, monkeypatch): monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") dns.fail() - with _hf_offline_if_dns_dead(): + with _hf_offline_if_unreachable(): assert os.environ.get("HF_HUB_OFFLINE") == "1" assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" # HF_HUB_OFFLINE was set by helper -> removed. @@ -970,16 +998,75 @@ class TestHfOfflineIfDnsDead: # TRANSFORMERS_OFFLINE pre-existed -> preserved. assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" - def test_exception_inside_block_still_restores_env(self, dns, clean_offline_env): + def test_exception_inside_block_still_restores_env(self, dns, reachable, clean_offline_env): dns.fail() with pytest.raises(RuntimeError, match = "boom"): - with _hf_offline_if_dns_dead(): + with _hf_offline_if_unreachable(): raise RuntimeError("boom") # Cleanup must happen on exception as well. assert "HF_HUB_OFFLINE" not in os.environ assert "TRANSFORMERS_OFFLINE" not in os.environ +class TestHfUnreachableProbe: + """``utils.utils.hf_unreachable``: memoised, opt-outable, fails open.""" + + @pytest.fixture(autouse = True) + def _reset(self): + from utils.utils import reset_hf_reachability_cache + + reset_hf_reachability_cache() + yield + reset_hf_reachability_cache() + + def _patch_probe(self, monkeypatch, result, calls): + import utils.transformers_version as tv + + def _probe(*_a, **_k): + calls.append(1) + if isinstance(result, Exception): + raise result + return result + + monkeypatch.setattr(tv, "hf_endpoint_unreachable", _probe) + + def test_probes_once_then_memoises(self, monkeypatch, clean_offline_env): + from utils.utils import hf_unreachable + + calls: list = [] + self._patch_probe(monkeypatch, True, calls) + assert hf_unreachable() is True + assert hf_unreachable() is True + assert len(calls) == 1 + + def test_opt_out_skips_probe(self, monkeypatch, clean_offline_env): + from utils.utils import hf_unreachable + + calls: list = [] + self._patch_probe(monkeypatch, True, calls) + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "0") + assert hf_unreachable() is False + assert calls == [] + + def test_probe_failure_reports_reachable(self, monkeypatch, clean_offline_env): + from utils.utils import hf_unreachable + + calls: list = [] + self._patch_probe(monkeypatch, RuntimeError("boom"), calls) + # Fail open: a broken probe must not strand a working install offline. + assert hf_unreachable() is False + + def test_reset_forces_reprobe(self, monkeypatch, clean_offline_env): + from utils.utils import hf_unreachable, reset_hf_reachability_cache + + calls: list = [] + self._patch_probe(monkeypatch, True, calls) + assert hf_unreachable() is True + reset_hf_reachability_cache() + assert hf_unreachable() is True + assert len(calls) == 2 + + class TestExtractQuantLabelSubdir: """``_extract_quant_label`` must consider parent dirs when the basename has no quant token (subdir layouts like ``BF16/foo.gguf``).""" diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index e4964b8d04..0871681fb6 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -5,6 +5,8 @@ import os import structlog +import threading +import time from loggers import get_logger from contextlib import contextmanager from pathlib import Path @@ -35,6 +37,150 @@ def hf_env_offline() -> bool: return False +# One load makes many hub calls, so the reachability verdict is shared for a short window. +_HF_REACHABILITY_TTL_S = 60.0 +_hf_reachability: Optional[tuple] = None +_hf_reachability_lock = threading.Lock() + + +def hf_probe_disabled() -> bool: + """True when UNSLOTH_OFFLINE_PROBE opts out of the reachability probe.""" + return os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in { + "0", + "false", + "no", + "off", + } + + +def reset_hf_reachability_cache() -> None: + """Drop the memoised verdict so the next call re-probes (tests, network changes).""" + global _hf_reachability + with _hf_reachability_lock: + _hf_reachability = None + + +def hf_unreachable(timeout: int = 3) -> bool: + """True when the HF endpoint is unreachable, memoised for _HF_REACHABILITY_TTL_S. + + Resolving DNS does not mean the Hub is reachable: a router that is up with the WAN + down, a captive portal or a stale DNS cache all answer lookups while every request + then burns huggingface_hub's retry backoff. This is the bounded, proxy-aware probe + the export path already uses; disable it with UNSLOTH_OFFLINE_PROBE=0. + + Fails open: an unavailable probe reports reachable, so the load decides as it does today. + """ + if hf_probe_disabled(): + return False + + global _hf_reachability + cached = _hf_reachability + if cached is not None and time.monotonic() - cached[0] < _HF_REACHABILITY_TTL_S: + return cached[1] + + with _hf_reachability_lock: + cached = _hf_reachability + if cached is not None and time.monotonic() - cached[0] < _HF_REACHABILITY_TTL_S: + return cached[1] + try: + from utils.transformers_version import hf_endpoint_unreachable + + unreachable = hf_endpoint_unreachable(timeout) + except Exception: + unreachable = False + _hf_reachability = (time.monotonic(), unreachable) + return unreachable + + +def _reset_hf_sessions() -> None: + """Drop cached hub sessions so they remount with the current offline adapter.""" + try: + from huggingface_hub.utils import _http + + for name in ("_get_session_from_cache", "get_session"): + cache_clear = getattr(getattr(_http, name, None), "cache_clear", None) + if cache_clear is not None: + cache_clear() + reset = getattr(_http, "reset_sessions", None) + if reset is not None: + reset() + except Exception: + pass + + +# Process-global, so nested/concurrent loads refcount rather than restore out from under +# each other. +_force_offline_depth = 0 +_force_offline_saved: list = [] +_force_offline_saved_env: dict = {} +_force_offline_lock = threading.Lock() + +_OFFLINE_ENV_KEYS = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") +_OFFLINE_CONSTANTS = ( + ("huggingface_hub.constants", ("HF_HUB_OFFLINE",)), + ("transformers.utils.hub", ("_is_offline_mode", "OFFLINE")), +) + + +@contextmanager +def force_hf_offline(): + """Force HF offline for this block, in-process. + + Setting the env vars is not enough once the process is running: huggingface_hub and + transformers read their offline constants at import, and hub sessions cache a + non-offline adapter. Flip the constants too and rebuild the sessions, so hub calls + fail fast instead of retrying. Everything is restored on exit. + """ + global _force_offline_depth, _force_offline_saved, _force_offline_saved_env + import importlib + + with _force_offline_lock: + if _force_offline_depth == 0: + saved: list = [] + saved_env: dict = {} + # Snapshot constants BEFORE forcing the env, else a module first imported + # inside the window reads the "1" and we would restore it as offline. + for mod_name, attrs in _OFFLINE_CONSTANTS: + try: + mod = importlib.import_module(mod_name) + except Exception: + continue + for attr in attrs: + if hasattr(mod, attr): + saved.append((mod, attr, getattr(mod, attr))) + for key in _OFFLINE_ENV_KEYS: + saved_env[key] = os.environ.get(key) + os.environ[key] = "1" + for mod, attr, _ in saved: + try: + setattr(mod, attr, True) + except Exception: + pass + _force_offline_saved = saved + _force_offline_saved_env = saved_env + _reset_hf_sessions() + _force_offline_depth += 1 + try: + yield + finally: + with _force_offline_lock: + _force_offline_depth -= 1 + if _force_offline_depth == 0: + for mod, attr, val in _force_offline_saved: + try: + setattr(mod, attr, val) + except Exception: + pass + _force_offline_saved = [] + for key, val in _force_offline_saved_env.items(): + if val is None: + os.environ.pop(key, None) + else: + os.environ[key] = val + _force_offline_saved_env = {} + _reset_hf_sessions() + + def st_repo_id_candidates(model_name: str) -> list: """Repo ids a Sentence-Transformers load may resolve model_name to; a slashless name also resolves under the sentence-transformers/ namespace, so both are candidates."""