Close remaining network paths in local-only loads and validation

The Vulkan-ordinal GGUF preflight downloads before the Phase 2 block, so it
now takes local_files_only and the forced-offline wrap; both download blocks
force offline so the cache-size verification (get_paths_info) stays off the
network. The offline helper now overrides an explicitly falsy
HF_HUB_OFFLINE=0 under force and restores prior values on exit. The validate
route keeps its whole metadata and security preflight offline, not just the
identifier probe. The python load path survives the process boundary: the
orchestrator forwards the flag and the route-resolved snapshot path into the
worker, which re-applies the path rewrite after rebuilding its ModelConfig,
runs the entire load under a scoped offline env, and passes the flag to the
vision processor fallback so a Hub base_model resolves from cache or fails
over instead of downloading.
This commit is contained in:
Unsloth 2026-07-26 22:51:45 -07:00
commit d68330080b
7 changed files with 165 additions and 26 deletions

View file

@ -322,6 +322,7 @@ class InferenceBackend:
hf_token: Optional[str] = None,
trust_remote_code: bool = False,
gpu_ids: Optional[list[int]] = None,
local_files_only: bool = False,
) -> bool:
"""Load any model: base, LoRA adapter, text, or vision."""
# Keep the token so the native-template fallback can fetch a
@ -582,10 +583,14 @@ class InferenceBackend:
)
from transformers import AutoProcessor
# Local-only loads: a LoRA base or export_metadata base_model
# is a Hub repo id; resolve it from cache or fail the load
# (candidate failover) instead of downloading the processor.
processor = AutoProcessor.from_pretrained(
processor_source,
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
local_files_only = local_files_only,
)
logger.info(f"Loaded {type(processor).__name__} from {processor_source}")

View file

@ -510,27 +510,44 @@ def _hf_env_offline() -> bool:
def _hf_offline_if_dns_dead(force: bool = False):
"""Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails;
restores env on exit so a transient hiccup can't quarantine the process.
No-op if the user already set it. ``force`` skips the DNS probe and goes
offline unconditionally (local-only background loads resolve metadata
from the cache without any network)."""
if "HF_HUB_OFFLINE" in os.environ:
yield False
return
if not force and not _probe_dns_dead():
No-op when the user already set it to a truthy value. ``force`` skips the
DNS probe and goes offline unconditionally (local-only background loads
resolve metadata from the cache without any network), overriding even an
explicitly falsy HF_HUB_OFFLINE=0 for the block and restoring it after."""
if _hf_env_offline():
yield False
return
if not force:
if "HF_HUB_OFFLINE" in os.environ:
# A user-pinned falsy value stays authoritative for ordinary loads.
yield False
return
if not _probe_dns_dead():
yield False
return
transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ
hub_prev = os.environ.get("HF_HUB_OFFLINE")
transformers_prev = os.environ.get("TRANSFORMERS_OFFLINE")
wrote_transformers = transformers_prev is None or force
os.environ["HF_HUB_OFFLINE"] = "1"
if not transformers_was_set:
if wrote_transformers:
os.environ["TRANSFORMERS_OFFLINE"] = "1"
logger.warning("huggingface.co unreachable; using local HF cache for this load.")
if force:
logger.info("Local-only load: forcing HF offline for this block.")
else:
logger.warning("huggingface.co unreachable; using local HF cache for this load.")
try:
yield True
finally:
os.environ.pop("HF_HUB_OFFLINE", None)
if not transformers_was_set:
os.environ.pop("TRANSFORMERS_OFFLINE", None)
if hub_prev is None:
os.environ.pop("HF_HUB_OFFLINE", None)
else:
os.environ["HF_HUB_OFFLINE"] = hub_prev
if wrote_transformers:
if transformers_prev is None:
os.environ.pop("TRANSFORMERS_OFFLINE", None)
else:
os.environ["TRANSFORMERS_OFFLINE"] = transformers_prev
try:
@ -6560,11 +6577,15 @@ class LlamaCppBackend:
hf_repo,
)
hf_repo = _resolved_repo
with _hf_offline_if_dns_dead():
# Same local-only policy as the Phase 2 download below: this
# preflight runs FIRST, so without the flag it would be the
# download path a background load slips through.
with _hf_offline_if_dns_dead(force = local_files_only):
_preflight_model_path = self._download_gguf(
hf_repo = hf_repo,
hf_variant = hf_variant,
hf_token = hf_token,
local_files_only = local_files_only,
)
if self._gguf_path_is_diffusion(_preflight_model_path, model_identifier):
raise ValueError(
@ -6599,7 +6620,9 @@ class LlamaCppBackend:
hf_repo,
)
hf_repo = _resolved_repo
with _hf_offline_if_dns_dead():
# Forced offline under local-only so even the cache-size
# verification (get_paths_info) stays off the network.
with _hf_offline_if_dns_dead(force = local_files_only):
model_path = _preflight_model_path or self._download_gguf(
hf_repo = hf_repo,
hf_variant = hf_variant,

View file

@ -524,6 +524,9 @@ class MLXInferenceBackend:
dtype = None,
parallel_mode = None,
distributed_group = None,
# Accepted for worker parity; the load runs under the worker's forced
# offline env when set, which is what enforces local-only here.
local_files_only: bool = False,
) -> bool:
import mlx.core as mx

View file

@ -973,6 +973,7 @@ class InferenceOrchestrator:
subject: Optional[str] = None,
tensor_parallel: bool = False,
mlx_distributed: bool = False,
local_files_only: bool = False,
) -> bool:
"""Load a model for inference.
@ -1000,6 +1001,14 @@ class InferenceOrchestrator:
"gpu_ids": gpu_ids,
"tensor_parallel": bool(tensor_parallel),
"mlx_distributed": bool(mlx_distributed),
"local_files_only": bool(local_files_only),
# Route-resolved local snapshot: the worker rebuilds its
# ModelConfig from model_name, which would lose the rewrite.
"local_snapshot_path": (
getattr(config, "path", None)
if local_files_only and getattr(config, "path", None) != model_name
else None
),
"mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline")
if mlx_distributed
else None,

View file

@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker.
from __future__ import annotations
import base64
import contextlib
import json
from loggers import get_logger
import os
@ -95,6 +96,28 @@ def _clean_token(value: str | None) -> str | None:
return value if value and value.strip() else None
@contextlib.contextmanager
def _local_only_offline_env(enabled: bool):
"""Force HF offline for a local-only background load. The worker is a
fresh per-load process, so scoping env here covers config resolution,
security gates, and every from_pretrained in the load, then restores it so
generation-time fetches (e.g. the chat template fallback) still work."""
if not enabled:
yield
return
prev = {key: os.environ.get(key) for key in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")}
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
try:
yield
finally:
for key, value in prev.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
def _build_model_config(config: dict):
"""Build a ModelConfig from the config dict."""
from utils.models import ModelConfig
@ -107,6 +130,12 @@ def _build_model_config(config: dict):
)
if not mc:
raise ValueError(f"Invalid model identifier: {model_name}")
# The route resolves a cached repo id to its local snapshot for local-only
# loads; from_identifier here rebuilds path as the repo id, losing that
# rewrite across the process boundary, so re-apply it.
snapshot_override = config.get("local_snapshot_path")
if snapshot_override:
mc.path = snapshot_override
return mc
@ -285,7 +314,13 @@ def _run_security_gates(
def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"""Handle a load command: load a model into the backend."""
_offline_guard = contextlib.ExitStack()
try:
# Local-only loads run the WHOLE load offline: config resolution,
# security gates, and from_pretrained can otherwise all reach the Hub.
_offline_guard.enter_context(
_local_only_offline_env(bool(config.get("local_files_only", False)))
)
mc = _build_model_config(config)
hf_token = _clean_token(config.get("hf_token"))
@ -363,6 +398,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"hf_token": hf_token,
"trust_remote_code": trust_remote_code,
"gpu_ids": config.get("resolved_gpu_ids"),
"local_files_only": bool(config.get("local_files_only", False)),
}
if getattr(backend, "device", None) == "mlx":
load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode")
@ -435,6 +471,8 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"stack": traceback.format_exc(limit = 20),
},
)
finally:
_offline_guard.close()
def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool:

View file

@ -4897,6 +4897,9 @@ async def _load_model_impl(
approved_remote_code_fingerprint = request.approved_remote_code_fingerprint,
gpu_ids = effective_gpu_ids,
subject = current_subject,
# Threads the rewritten snapshot path and forced-offline load into
# the worker subprocess, which rebuilds its own ModelConfig.
local_files_only = request.local_files_only,
)
if not success:
@ -5130,18 +5133,24 @@ async def validate_model(
native_grant_backed = False
model_log_label = request.model_path
# Local-only validation (background auto-loads) covers the WHOLE metadata
# and security preflight with a forced-offline env, not just the identifier
# probe: the upgrade, remote-code, and file-security helpers below all read
# Hub configs and would otherwise refetch them for a cache-only candidate.
import contextlib
_local_only_offline = contextlib.ExitStack()
try:
model_identifier, model_log_label, native_grant_backed = (
_resolve_model_identifier_for_request(request, operation = "validate-model")
)
# Local-only validation (background auto-loads) forces offline so the
# metadata probe resolves from cache instead of the Hub.
with _hf_offline_if_dns_dead(force = request.local_files_only):
config = ModelConfig.from_identifier(
model_id = model_identifier,
hf_token = request.hf_token,
gguf_variant = request.gguf_variant,
)
if request.local_files_only:
_local_only_offline.enter_context(_hf_offline_if_dns_dead(force = True))
config = ModelConfig.from_identifier(
model_id = model_identifier,
hf_token = request.hf_token,
gguf_variant = request.gguf_variant,
)
if not config:
raise HTTPException(
@ -5400,6 +5409,8 @@ async def validate_model(
status_code = 400,
detail = "Invalid model",
)
finally:
_local_only_offline.close()
# studio_router only: admin action, kept off the OpenAI-compatible /v1 mount.

View file

@ -1100,8 +1100,9 @@ def test_background_candidate_filters_have_no_side_effects():
assert "local_files_only: bool = Field(" in validate_schema
route = _read_backend("routes/inference.py")
# Both /load and /validate force offline resolution under local-only.
assert route.count("with _hf_offline_if_dns_dead(force = request.local_files_only):") == 2
# /load forces offline resolution under local-only; /validate covers its
# whole preflight via the ExitStack wrap (asserted separately below).
assert route.count("with _hf_offline_if_dns_dead(force = request.local_files_only):") == 1
# Audio gate present in both routes, GGUF exempt (llama.cpp path already
# resolves companions cached-or-skipped under the flag).
assert route.count("needs audio codec downloads") == 2
@ -1111,13 +1112,62 @@ def test_background_candidate_filters_have_no_side_effects():
llama = _read_backend("core/inference/llama_cpp.py")
assert "def _hf_offline_if_dns_dead(force: bool = False):" in llama
assert "if not force and not _probe_dns_dead():" in llama
# force must also override an explicitly falsy HF_HUB_OFFLINE=0: only a
# TRUTHY env value short-circuits, and prior values are restored on exit.
assert "if _hf_env_offline():" in llama
assert 'hub_prev = os.environ.get("HF_HUB_OFFLINE")' in llama
assert 'os.environ["HF_HUB_OFFLINE"] = hub_prev' in llama
helper = _read_backend("hub/utils/local_snapshot.py")
assert "def _snapshot_dir_fallback(" in helper
assert "config.json" in helper
def test_local_only_covers_every_load_and_validate_network_path():
"""Round-14 gates. The Vulkan-ordinal GGUF preflight downloads FIRST, so it
takes the same local-only flag and forced-offline wrap as the Phase 2
download (whose cache-size check would otherwise call get_paths_info).
The validate route keeps its whole metadata/security preflight offline, not
just the identifier probe. The python load path survives the process
boundary: the orchestrator forwards the flag and the route-resolved
snapshot path into the worker, which runs the entire load under a scoped
offline env and passes the flag to the vision processor fallback."""
llama = _read_backend("core/inference/llama_cpp.py")
# Both GGUF download blocks in load_model force offline under local-only.
assert llama.count("with _hf_offline_if_dns_dead(force = local_files_only):") == 2
preflight = llama.split("_preflight_model_path = self._download_gguf(", 1)[1]
preflight = preflight.split(")", 1)[0]
assert "local_files_only = local_files_only," in preflight
route = _read_backend("routes/inference.py")
validate_src = route.split('operation = "validate-model"', 1)[1]
assert "_local_only_offline.enter_context(_hf_offline_if_dns_dead(force = True))" in validate_src
assert "_local_only_offline.close()" in validate_src
# The non-GGUF load threads the flag into the subprocess backend.
load_call = route.split("backend.load_model,\n config = config,", 1)[1]
load_call = load_call.split(")", 1)[0]
assert "local_files_only = request.local_files_only," in load_call
orchestrator = _read_backend("core/inference/orchestrator.py")
assert '"local_files_only": bool(local_files_only),' in orchestrator
assert '"local_snapshot_path"' in orchestrator
worker = _read_backend("core/inference/worker.py")
assert "def _local_only_offline_env(" in worker
assert "_offline_guard.enter_context(" in worker
assert "_offline_guard.close()" in worker
# The route's snapshot rewrite is re-applied after the worker rebuilds
# its ModelConfig from the identifier.
assert 'snapshot_override = config.get("local_snapshot_path")' in worker
assert 'mc.path = snapshot_override' in worker
assert '"local_files_only": bool(config.get("local_files_only", False)),' in worker
inference = _read_backend("core/inference/inference.py")
processor_call = inference.split("processor = AutoProcessor.from_pretrained(", 1)[1]
processor_call = processor_call.split("logger.info", 1)[0]
assert "local_files_only = local_files_only," in processor_call
def test_gguf_background_loads_never_download_companions():
"""A cached GGUF load can still fetch from the Hub through its optional
companions (mmproj, MTP drafter) or a cache-miss main quant. Background