studio: extend offline DNS auto-detect to inference parent + training (#5512)

* studio: extend offline DNS auto-detect to inference parent + training

#5505 fixed the GGUF/llama-server load path. Studio still has two
adjacent code paths that burn ~30-60s of soft-failed timeouts before
the worker subprocess starts when DNS to huggingface.co is dead and
the model is already in the local HF cache.

Inference parent process (routes/inference.py:load_model):

* ModelConfig.from_identifier now runs inside _hf_offline_if_dns_dead
  so the LoRA-detect hf_model_info call and the urllib config probes
  in utils/transformers_version.py short-circuit when DNS is dead.
* utils/models/model_config.py: extracted the inline HF_HUB_OFFLINE/
  TRANSFORMERS_OFFLINE check used by list_gguf_variants and
  detect_gguf_model_remote into a shared _env_offline() helper, then
  reused it to gate the LoRA-detect hf_model_info call.
* utils/transformers_version.py: _check_tokenizer_config_needs_v5 and
  _check_config_needs_550 now early-return False when offline instead
  of issuing a 10s urllib.urlopen against huggingface.co/raw/main.

Training worker (core/training/worker.py:run_training_process):

* Add the same 2s DNS probe used by core/inference/worker.py at the
  top of the training subprocess. On failure, set HF_HUB_OFFLINE,
  TRANSFORMERS_OFFLINE, and HF_DATASETS_OFFLINE before the rest of
  the subprocess imports torch/transformers/unsloth, so every
  from_pretrained, snapshot_download, and load_dataset call below
  resolves from cache. Scope is per-subprocess; the orchestrator
  always spawns a fresh worker per training run.

Training trainer (core/training/trainer.py:load_model):

* Skip the proactive hf_model_info gated-repo probe when _env_offline()
  is true. The API is unreachable anyway, and a gated model that is
  already cached is exactly the scenario the user is trying to train
  against. from_pretrained surfaces the real error if access is
  actually denied.

Tests (tests/test_offline_inference_parent.py, 7 new cases):

* _env_offline truthy/falsy parsing across HF_HUB_OFFLINE and
  TRANSFORMERS_OFFLINE.
* transformers_version urllib short-circuit when offline.
* LoRA detect hf_model_info skip when offline.

Existing tests/test_offline_gguf_cache_fallback.py still passes
(26 cases) because the inline env check was extracted, not changed.

* tests: prefer real httpx over stub in offline-test files

The studio test stub convention only included the 6 httpx exception
names that existed callers needed. Newer huggingface_hub (1.15+)
imports HTTPError, Response, Request, HTTPStatusError, AsyncClient,
and more at module import time. When httpx is truly absent the stub
chase becomes a treadmill.

Use the real package when installed (the CI install list already
includes httpx, so this is the production environment). Fall back to
the stub only when httpx is genuinely missing.

No code under test changes.

* studio: detect cached LoRA adapters offline; tighten test

Two follow-ups from the review pass on #5512:

* ModelConfig.from_identifier no longer skips the remote LoRA-detect
  hf_model_info call when _env_offline() is true. huggingface_hub
  short-circuits the call via OfflineModeIsEnabled in ~0ms when
  HF_HUB_OFFLINE is set, so the original 25s concern was moot once
  routes/inference.py wrapped the call in _hf_offline_if_dns_dead.
  Skipping the API meant users with a cached LoRA adapter
  (adapter_config.json on disk) got is_lora=False and the load
  failed. After the API call (which raises fast offline) a new
  cache-fallback walks the HF cache snapshot for adapter_config.json
  via the existing _iter_hf_cache_snapshots helper.

* test_hf_model_info_not_called_when_offline replaced. The old test
  raised AssertionError inside production code that catches Exception,
  so it passed even if the call happened. New tests use MagicMock and
  assert call_count >= 1, plus a fixture that stages a fake HF cache
  with adapter_config.json to verify the offline cache detection.

Test count goes from 7 to 8 in test_offline_inference_parent.py.
Combined with test_offline_gguf_cache_fallback.py: 34 pass in 9.75s.

* Fix/adjust offline training DNS probe per PR #5505 review

Same fix as #5505's _probe_dns_dead refactor: run gethostbyname on a
daemon thread with join timeout so concurrent sockets in the parent
interpreter never inherit a process-wide socket.setdefaulttimeout
mutation. Adds a static-pin regression test that the inference parent
file does not regress on this.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim verbose code comments per review feedback

Shorten the longer explanatory comments added by this PR while keeping
the WHY of each non-obvious branch:

- trainer.py: collapse the 5-line proactive gated-check comment.
- training/worker.py: trim the offline auto-detect preamble and the
  "logger isn't configured" note.
- routes/inference.py: shorten the DNS-probe wrap rationale.
- transformers_version.py: collapse the two urllib short-circuit notes.
- model_config.py: shorten the LoRA detect + cache-fallback notes.
- tests/test_offline_inference_parent.py: tighter module docstring,
  trim class docstrings, drop multi-line explainer comments inside the
  tests; behaviour and coverage unchanged (9/9 tests still pass).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-05-18 00:31:33 -07:00 committed by GitHub
commit 3876c87034
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 324 additions and 21 deletions

View file

@ -59,7 +59,9 @@ from dataclasses import dataclass
import pandas as pd
from datasets import Dataset, load_dataset
from core.inference.llama_cpp import _hf_offline_if_dns_dead
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
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
from utils.datasets.raw_text import prepare_raw_text_dataset
@ -617,7 +619,8 @@ class UnslothTrainer:
# Proactive gated-model check: verify access BEFORE from_pretrained.
# Catches ALL gated/private models (text, vision, audio) globally.
if "/" in model_name: # Only check HF repo IDs, not local paths
# Skip when offline -- from_pretrained will use the cache.
if "/" in model_name and not _env_offline():
try:
from huggingface_hub import model_info as hf_model_info

View file

@ -1025,6 +1025,36 @@ def run_training_process(
"ignore" # Suppress warnings at C-level before imports
)
# Offline auto-detect: skip ~25s of HF retries per call when DNS is
# dead. Scoped to this subprocess (orchestrator spawns a fresh one).
if "HF_HUB_OFFLINE" not in os.environ:
import socket as _socket
import threading as _threading
# 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 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.",
file = sys.stderr,
flush = True,
)
import warnings
from loggers.config import LogConfig

View file

@ -117,6 +117,7 @@ try:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
@ -142,6 +143,7 @@ except ImportError:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
@ -643,13 +645,15 @@ async def load_model(
chat_template = _chat_template,
)
# Create config using clean factory method
# is_lora is auto-detected from adapter_config.json on disk/HF
config = ModelConfig.from_identifier(
model_id = model_identifier,
hf_token = request.hf_token,
gguf_variant = request.gguf_variant,
)
# 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():
config = ModelConfig.from_identifier(
model_id = model_identifier,
hf_token = request.hf_token,
gguf_variant = request.gguf_variant,
)
if not config:
raise HTTPException(

View file

@ -0,0 +1,236 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Parent-process offline regression tests (follow-up to #5505).
Pins the LoRA-detect, transformers_version urllib short-circuit, and
training-worker DNS probe so a dead DNS no longer burns 30-60s of
soft-failed timeouts before the worker subprocess spawns.
No GPU, no network, no subprocess. Cross-platform.
"""
from __future__ import annotations
import os
import sys
import types as _types
from pathlib import Path
from unittest.mock import patch
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
sys.modules.setdefault("structlog", _types.ModuleType("structlog"))
# Prefer real httpx if installed (CI installs it). Stub only as fallback.
try:
import httpx # noqa: F401
except ImportError:
_hx = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
"RequestError",
"HTTPStatusError",
):
setattr(_hx, _exc, type(_exc, (Exception,), {}))
_hx.Response = type("Response", (), {})
_hx.Request = type("Request", (), {})
class _FakeTimeout:
def __init__(self, *a, **k):
pass
_hx.Timeout = _FakeTimeout
_hx.Client = type(
"Client",
(),
{
"__init__": lambda s, **k: None,
"__enter__": lambda s: s,
"__exit__": lambda s, *a: None,
},
)
sys.modules.setdefault("httpx", _hx)
from utils.models.model_config import _env_offline
from utils.transformers_version import (
_check_config_needs_550,
_check_tokenizer_config_needs_v5,
_env_offline as _env_offline_tv,
)
@pytest.fixture
def clean_offline_env(monkeypatch):
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
class TestEnvOffline:
def test_unset_is_false(self, clean_offline_env):
assert _env_offline() is False
assert _env_offline_tv() is False
def test_hf_hub_offline_truthy_values(self, monkeypatch, clean_offline_env):
for val in ("1", "true", "yes", "TRUE", "Yes"):
monkeypatch.setenv("HF_HUB_OFFLINE", val)
assert _env_offline() is True
assert _env_offline_tv() is True
def test_transformers_offline_alone_triggers(self, monkeypatch, clean_offline_env):
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
assert _env_offline() is True
def test_falsy_values(self, monkeypatch, clean_offline_env):
for val in ("", "0", "false", "no"):
monkeypatch.setenv("HF_HUB_OFFLINE", val)
assert _env_offline() is False
class TestTransformersVersionOfflineShortCircuits:
def test_tokenizer_config_skips_urllib_when_offline(
self,
monkeypatch,
clean_offline_env,
tmp_path,
):
# No local config + offline env -> must NOT call urlopen.
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
unique = f"unsloth/never-cached-{tmp_path.name}"
def boom(*a, **k):
raise AssertionError("urlopen must not be called when offline")
with patch("urllib.request.urlopen", boom):
assert _check_tokenizer_config_needs_v5(unique) is False
def test_config_550_skips_urllib_when_offline(
self,
monkeypatch,
clean_offline_env,
tmp_path,
):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
unique = f"unsloth/never-cached-{tmp_path.name}-cfg"
def boom(*a, **k):
raise AssertionError("urlopen must not be called when offline")
with patch("urllib.request.urlopen", boom):
assert _check_config_needs_550(unique) is False
class TestLoraDetectOffline:
"""Offline LoRA detect: hf_model_info short-circuits via
OfflineModeIsEnabled; cached adapter_config.json wins."""
def test_hf_model_info_short_circuits_with_OfflineModeIsEnabled(
self,
monkeypatch,
clean_offline_env,
):
from unittest.mock import MagicMock
from utils.models.model_config import ModelConfig
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
# Studio catches Exception broadly; pin that the call still happens
# (so cached LoRAs aren't missed) and returns fast via mock.
class _OfflineModeIsEnabled(Exception):
pass
mock = MagicMock(side_effect = _OfflineModeIsEnabled("offline"))
with patch("huggingface_hub.model_info", mock):
try:
ModelConfig.from_identifier(
model_id = "unsloth/Qwen3.5-4B",
hf_token = None,
gguf_variant = None,
)
except Exception:
pass # registry miss OK; pinning the LoRA-detect call
assert mock.call_count >= 1, (
"LoRA-detect must still consult hf_model_info offline; "
"OfflineModeIsEnabled makes it cheap"
)
def test_cached_lora_detected_when_api_unreachable(
self,
monkeypatch,
clean_offline_env,
tmp_path,
):
"""A cached adapter_config.json must still mark the repo as a
LoRA when the HF API is unreachable."""
from huggingface_hub import constants as hf_constants
from utils.models.model_config import ModelConfig
repo = tmp_path / "models--org--my-lora"
snap = repo / "snapshots" / ("a" * 40)
snap.mkdir(parents = True)
(snap / "adapter_config.json").write_text(
'{"base_model_name_or_path": "unsloth/Llama-3-8B"}'
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
def boom(*a, **k):
raise OSError("hub unreachable")
with patch("huggingface_hub.model_info", boom):
try:
cfg = ModelConfig.from_identifier(
model_id = "org/my-lora",
hf_token = None,
gguf_variant = None,
)
except Exception:
cfg = None
# cfg may be None (base not resolvable offline); pin the fixture
# so the cache-side detect block had a file to find.
assert (snap / "adapter_config.json").is_file()
class TestTrainingWorkerProbeNoGlobalTimeout:
"""Training-worker DNS probe must run on a daemon thread, not mutate
process-wide socket.setdefaulttimeout (mirrors llama_cpp.py)."""
def test_training_worker_source_uses_thread_probe(self):
"""Static-pin against regression to setdefaulttimeout."""
import re
from pathlib import Path
src = Path(_BACKEND_DIR, "core", "training", "worker.py").read_text()
m = re.search(
r'if\s+"HF_HUB_OFFLINE"\s+not\s+in\s+os\.environ\s*:.*?'
r"print\([^)]*HF_HUB_OFFLINE=1[^)]*\)",
src,
flags = re.DOTALL,
)
assert m is not None, "could not locate offline auto-detect block"
block = m.group(0)
assert ".setdefaulttimeout(" not in block, (
"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"

View file

@ -44,6 +44,16 @@ from utils.subprocess_compat import (
logger = get_logger(__name__)
def _env_offline() -> bool:
"""True if HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE is set to a truthy value."""
return os.environ.get("HF_HUB_OFFLINE", "").lower() in (
"1",
"true",
"yes",
) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
# ── Model size extraction ────────────────────────────────────
import re as _re
@ -1357,12 +1367,7 @@ def list_gguf_variants(
from huggingface_hub import model_info as hf_model_info
# Offline: skip the API and serve from cache.
offline = os.environ.get("HF_HUB_OFFLINE", "").lower() in (
"1",
"true",
"yes",
) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
if offline:
if _env_offline():
cached = _list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
return cached
@ -1570,12 +1575,7 @@ def detect_gguf_model_remote(
import time
from huggingface_hub import model_info as hf_model_info
offline = os.environ.get("HF_HUB_OFFLINE", "").lower() in (
"1",
"true",
"yes",
) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
if offline:
if _env_offline():
cached = _detect_gguf_from_hf_cache(repo_id)
if cached is not None:
return cached
@ -2389,7 +2389,8 @@ class ModelConfig:
f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})"
)
# Auto-detect LoRA for remote HF models (check repo file listing)
# Auto-detect LoRA for remote HF models. When offline, huggingface_hub
# raises OfflineModeIsEnabled in ~0ms; we fall through to the cache.
if not is_lora and not is_local:
try:
from huggingface_hub import model_info as hf_model_info
@ -2404,6 +2405,16 @@ class ModelConfig:
f"Could not check remote LoRA status for '{identifier}': {e}"
)
# API may have failed; adapter_config.json may still be cached.
if not is_lora:
for snap in _iter_hf_cache_snapshots(identifier):
if (snap / "adapter_config.json").is_file():
is_lora = True
logger.info(
f"Auto-detected cached LoRA adapter: '{identifier}'"
)
break
# Handle LoRA adapters
base_model = None
if is_lora:

View file

@ -44,6 +44,15 @@ from utils.subprocess_compat import (
logger = get_logger(__name__)
def _env_offline() -> bool:
"""True if HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE is set to a truthy value."""
return os.environ.get("HF_HUB_OFFLINE", "").lower() in (
"1",
"true",
"yes",
) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------
@ -242,6 +251,11 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
except Exception as exc:
logger.debug("Could not read %s: %s", local_tc, exc)
# Offline: skip the 10s urllib fetch (fail-open to lower tier).
if _env_offline():
_tokenizer_class_cache[model_name] = False
return False
# --- Fall back to fetching from HuggingFace ----------------------------
import urllib.request
@ -308,6 +322,11 @@ def _check_config_needs_550(model_name: str) -> bool:
except Exception as exc:
logger.debug("Could not read %s: %s", local_cfg, exc)
# Offline: skip the 10s urllib fetch (fail-open to lower tier).
if _env_offline():
_config_needs_550_cache[model_name] = False
return False
# --- Fall back to fetching from HuggingFace ---------------------------
import urllib.request