From d5758ecf58757b4444a905819a0f8a1075d9d1ee Mon Sep 17 00:00:00 2001 From: Unsloth Date: Tue, 28 Jul 2026 23:46:34 -0700 Subject: [PATCH] Address review round 2: proxy-aware detection, shared worker probe, local skip All four reproduced before fixing, and re-measured after. - Proxy-only egress was declared offline. With HTTP(S)_PROXY set, the proxy resolves the hub host, so a failing local lookup says nothing. The DNS shortcut now stands down whenever a proxy applies (and honours NO_PROXY), letting the proxy-aware probe decide. Measured: endpoint probe reachable through the proxy while the guard still forced offline. - The training worker kept its own inline probe hardcoded to huggingface.co, so a reachable HF_ENDPOINT mirror set lifetime offline flags. It now uses the shared endpoint- and proxy-aware helper. - /models/check-vision, /models/config and /picker/chat-template ran the probe even for local paths, which never reach the hub. Measured 0.9s of pure latency per request; now skipped via _hf_offline_if_unreachable_for. DNS/endpoint/proxy helpers now live in utils.utils so llama_cpp and the training worker share one implementation instead of three copies. The static pin in test_offline_inference_parent moved with the probe: the worker block must delegate to the shared helper and must not hardcode a host, and the daemon-thread/no-setdefaulttimeout property is pinned on dns_host_dead where it now lives. Offline path unchanged: load 686s -> 6s, /models/config 378s -> 0s, and a local-path vision check is back to 0s. --- studio/backend/core/inference/llama_cpp.py | 56 +++++++------- studio/backend/core/training/worker.py | 44 ++++------- studio/backend/picker/routes/templates.py | 4 +- studio/backend/routes/models.py | 12 +-- .../tests/test_offline_gguf_cache_fallback.py | 60 +++++++++++++++ .../tests/test_offline_inference_parent.py | 23 +++++- studio/backend/utils/utils.py | 74 +++++++++++++++++++ 7 files changed, 205 insertions(+), 68 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index cdfbfeeba3..c9df5a4599 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -496,36 +496,16 @@ _SWA_CACHE_LOCK = threading.Lock() def _hf_endpoint_host() -> str: - """Host of the configured hub endpoint. Mirror users point HF_ENDPOINT elsewhere, and - probing huggingface.co would then report their working mirror as offline.""" - endpoint = (os.environ.get("HF_ENDPOINT") or "").strip() or "https://huggingface.co" - if "://" not in endpoint: - endpoint = "https://" + endpoint - try: - from urllib.parse import urlparse - return urlparse(endpoint).hostname or "huggingface.co" - except Exception: - return "huggingface.co" + """Host of the configured hub endpoint (HF_ENDPOINT aware).""" + from utils.utils import hf_endpoint_host + return hf_endpoint_host() def _probe_dns_dead(host: Optional[str] = None, timeout: float = 2.0) -> bool: """Quick DNS check on a daemon thread, so concurrent sockets aren't affected by socket.setdefaulttimeout. Defaults to the configured endpoint's host.""" - host = host or _hf_endpoint_host() - result: list[Optional[bool]] = [None] - - def _probe() -> None: - try: - socket.gethostbyname(host) - result[0] = False - except Exception: - result[0] = True - - t = threading.Thread(target = _probe, daemon = True) - t.start() - t.join(timeout) - # Thread still running -> resolver wedged -> dead. - return True if result[0] is None else result[0] + from utils.utils import dns_host_dead, hf_endpoint_host + return dns_host_dead(host or hf_endpoint_host(), timeout) def _hf_env_offline() -> bool: @@ -546,14 +526,16 @@ def _hf_unreachable() -> bool: 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. + The DNS shortcut is skipped when a proxy is configured, since the proxy resolves the + hub host and local DNS then says nothing about reachability. """ - if _probe_dns_dead(): - return True try: - from utils.utils import hf_unreachable - return hf_unreachable() + from utils.utils import hf_dns_dead, hf_unreachable except Exception: return False + if hf_dns_dead(): + return True + return hf_unreachable() @contextlib.contextmanager @@ -599,6 +581,22 @@ def _hf_offline_if_unreachable(): os.environ.pop("TRANSFORMERS_OFFLINE", None) +def _hf_offline_if_unreachable_for(model_name): + """Guard, but only for remote repo ids. + + A local path is served from the filesystem and never reaches the hub, so probing + would add seconds per request and prevent no retry. + """ + try: + from utils.paths import is_local_path + + if isinstance(model_name, str) and is_local_path(model_name): + return contextlib.nullcontext() + except Exception: + pass + return _hf_offline_if_unreachable() + + try: _SLOT_SAVE_MAX_BYTES = int(os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30)) except ValueError: diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index fd6cb21d8b..adca92d0cc 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2346,37 +2346,23 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # 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 + _offline = False + try: + from utils.utils import hf_dns_dead, hf_probe_disabled - # Daemon thread so we don't mutate process-wide setdefaulttimeout. - _result: list = [None] - - def _probe() -> None: - try: - _socket.gethostbyname("huggingface.co") - _result[0] = False - except Exception: - _result[0] = True - - _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. - # These flags last the whole job, so only a connection failure counts: a - # momentary 502/503 must not block every download for the rest of the run. - try: + # hf_dns_dead follows HF_ENDPOINT and stands down when a proxy is configured, + # so a reachable mirror (or proxy-only egress) is never called offline. + _offline = hf_dns_dead() + if not _offline and not hf_probe_disabled(): + # DNS answers even without egress (WAN down, captive portal). These flags + # last the whole job, so only a connection failure counts: a momentary + # 502/503 must not block every download for the rest of the run. from utils.transformers_version import hf_endpoint_unreachable - from utils.utils import hf_probe_disabled - if not hf_probe_disabled() and hf_endpoint_unreachable( - gateway_errors_offline = False - ): - _result[0] = True - except Exception: - pass - if _result[0] is None or _result[0] is True: + + _offline = hf_endpoint_unreachable(gateway_errors_offline = False) + except Exception: + _offline = False + if _offline: os.environ["HF_HUB_OFFLINE"] = "1" os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") os.environ.setdefault("HF_DATASETS_OFFLINE", "1") diff --git a/studio/backend/picker/routes/templates.py b/studio/backend/picker/routes/templates.py index a070c8d054..628cc03841 100644 --- a/studio/backend/picker/routes/templates.py +++ b/studio/backend/picker/routes/templates.py @@ -39,10 +39,10 @@ async def get_default_chat_template_route( ) -> ModelTemplateResponse: # 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 + from core.inference.llama_cpp import _hf_offline_if_unreachable_for def _read(): - with _hf_offline_if_unreachable(): + with _hf_offline_if_unreachable_for(model_name): return read_default_chat_template(model_name, hf_token, gguf_variant) template = await asyncio.to_thread(_read) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index cc2a402b59..58b7433070 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1855,12 +1855,13 @@ 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 + from core.inference.llama_cpp import _hf_offline_if_unreachable_for try: # 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(): + # offline they must all resolve from the HF cache instead of retrying. Local + # paths stay on disk, so they skip the probe entirely. + with _hf_offline_if_unreachable_for(model_name): if not is_local_path(model_name): resolved = resolve_cached_repo_id_case(model_name) if resolved != model_name: @@ -2591,9 +2592,10 @@ async def check_vision_model( logger.info(f"Checking if vision model: {model_name}") # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). # Offline the guard keeps this on the HF cache instead of retrying the hub. - from core.inference.llama_cpp import _hf_offline_if_unreachable + # A local path resolves from disk, so it skips the probe. + from core.inference.llama_cpp import _hf_offline_if_unreachable_for - with _hf_offline_if_unreachable(): + with _hf_offline_if_unreachable_for(model_name): is_vision = is_vision_model(model_name, hf_token = hf_token) logger.info(f"Vision check result for {model_name}: is_vision={is_vision}") diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 3240c2dc5e..80bcf82377 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -1051,6 +1051,66 @@ class TestEndpointAwareOfflineDetection: assert _probe_dns_dead() is True +class TestProxyOnlyEgress: + """With a proxy, the proxy resolves the hub host, so local DNS proves nothing and + must not be used to declare the hub offline.""" + + @pytest.fixture + def dns_all_dead(self, monkeypatch): + def _fail(*a, **k): + raise socket.gaierror(-2, "Name or service not known") + + monkeypatch.setattr(socket, "gethostbyname", _fail) + + def test_dns_shortcut_stands_down_when_proxy_configured(self, monkeypatch, dns_all_dead): + from utils.utils import hf_dns_dead + + monkeypatch.delenv("HF_ENDPOINT", raising = False) + monkeypatch.setenv("HTTPS_PROXY", "http://proxy.internal:3128") + monkeypatch.setenv("HTTP_PROXY", "http://proxy.internal:3128") + assert hf_dns_dead() is False + + def test_dns_shortcut_applies_without_a_proxy(self, monkeypatch, dns_all_dead): + from utils.utils import hf_dns_dead + + for key in ("HTTPS_PROXY", "HTTP_PROXY", "https_proxy", "http_proxy", "ALL_PROXY"): + monkeypatch.delenv(key, raising = False) + monkeypatch.setattr("utils.utils.hf_proxy_configured", lambda: False) + assert hf_dns_dead() is True + + def test_no_proxy_bypass_restores_the_shortcut(self, monkeypatch, dns_all_dead): + """A host listed in NO_PROXY does not go through the proxy, so DNS matters again.""" + from utils.utils import hf_proxy_configured + + monkeypatch.setenv("HF_ENDPOINT", "https://huggingface.co") + monkeypatch.setenv("HTTPS_PROXY", "http://proxy.internal:3128") + monkeypatch.setenv("NO_PROXY", "huggingface.co") + assert hf_proxy_configured() is False + + +class TestGuardSkipsLocalPaths: + """A local path never reaches the hub, so probing costs time and prevents nothing.""" + + def test_local_path_is_a_noop(self, tmp_path, monkeypatch): + from core.inference.llama_cpp import _hf_offline_if_unreachable_for + + called: list = [] + monkeypatch.setattr( + "utils.utils.hf_unreachable", lambda *a, **k: called.append(1) or True + ) + with _hf_offline_if_unreachable_for(str(tmp_path / "model.gguf")) as engaged: + assert engaged is None # nullcontext yields None + assert called == [], "probed the hub for a local path" + + def test_remote_id_still_guarded(self, monkeypatch, clean_offline_env): + import core.inference.llama_cpp as lc + + monkeypatch.setattr(lc, "_hf_unreachable", lambda: True) + with lc._hf_offline_if_unreachable_for("unsloth/Qwen3.5-4B-GGUF") as engaged: + assert engaged is True + assert os.environ.get("HF_HUB_OFFLINE") == "1" + + class TestGatewayErrorsAreNotConnectionFailures: """Lifetime offline flags must not be set by a momentary 502/503/504.""" diff --git a/studio/backend/tests/test_offline_inference_parent.py b/studio/backend/tests/test_offline_inference_parent.py index 3e9f09bb2f..c8c440be46 100644 --- a/studio/backend/tests/test_offline_inference_parent.py +++ b/studio/backend/tests/test_offline_inference_parent.py @@ -218,6 +218,23 @@ class TestTrainingWorkerProbeNoGlobalTimeout: "training worker still calls socket.setdefaulttimeout; " "concurrent sockets would inherit the probe timeout" ) - assert ( - "threading" in block and "Thread" in block - ), "training worker probe must run on a daemon thread" + # The probe now lives in the shared helper (endpoint- and proxy-aware), so the + # worker must delegate to it rather than resolve a hardcoded host itself. + assert "hf_dns_dead" in block, "training worker must use the shared DNS helper" + assert 'gethostbyname("huggingface.co")' not in block, ( + "training worker must not hardcode huggingface.co; a reachable HF_ENDPOINT " + "mirror would be declared offline" + ) + + def test_shared_dns_helper_uses_thread_probe(self): + """The daemon-thread property moved with the probe; pin it where it now lives.""" + import inspect + + from utils.utils import dns_host_dead + + src = inspect.getsource(dns_host_dead) + assert ".setdefaulttimeout(" not in src, ( + "shared DNS probe calls socket.setdefaulttimeout; " + "concurrent sockets would inherit the probe timeout" + ) + assert "Thread" in src and "daemon" in src, "shared DNS probe must run on a daemon thread" diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index b2738f4d56..7e91307a02 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -37,6 +37,80 @@ def hf_env_offline() -> bool: return False +def hf_endpoint_url() -> str: + """Configured hub endpoint, scheme-normalised. Mirror users point this elsewhere.""" + endpoint = (os.environ.get("HF_ENDPOINT") or "").strip() or "https://huggingface.co" + return endpoint if "://" in endpoint else "https://" + endpoint + + +def hf_endpoint_host() -> str: + """Host of the configured endpoint; probing huggingface.co would misjudge a mirror.""" + try: + from urllib.parse import urlparse + + return urlparse(hf_endpoint_url()).hostname or "huggingface.co" + except Exception: + return "huggingface.co" + + +def hf_proxy_configured() -> bool: + """True when egress to the endpoint goes through a proxy. + + The proxy resolves the hub host, so local DNS proves nothing about reachability and + must not be used to declare the hub offline. + """ + try: + import urllib.request + from urllib.parse import urlparse + + url = hf_endpoint_url() + proxies = urllib.request.getproxies() + scheme = urlparse(url).scheme or "https" + if scheme not in proxies and "all" not in proxies: + return False + host = urlparse(url).hostname or "" + try: + if host and urllib.request.proxy_bypass(host): + return False + except Exception: + pass + return True + except Exception: + return False + + +def dns_host_dead(host: str, timeout: float = 2.0) -> bool: + """True when host does not resolve. Runs on a daemon thread so a wedged resolver + cannot block past the deadline and so socket.setdefaulttimeout is left alone.""" + result: list = [None] + + def _probe() -> None: + import socket as _socket + + try: + _socket.gethostbyname(host) + result[0] = False + except Exception: + result[0] = True + + t = threading.Thread(target = _probe, daemon = True) + t.start() + t.join(timeout) + # Still running -> resolver wedged -> treat as dead. + return True if result[0] is None else result[0] + + +def hf_dns_dead(timeout: float = 2.0) -> bool: + """Fast offline shortcut: the endpoint's host does not resolve and no proxy is in play. + + Returns False whenever a proxy is configured, so proxy-only setups fall through to the + real reachability probe instead of being wrongly declared offline. + """ + if hf_proxy_configured(): + return False + return dns_host_dead(hf_endpoint_host(), timeout) + + # One load makes many hub calls, so the verdict is shared briefly to avoid re-probing on # each. Kept short in BOTH directions: a stale "reachable" misses the plug being pulled # (the case this whole path exists for), and a stale "unreachable" sends a load to the