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:
parent
370195c702
commit
d68330080b
7 changed files with 165 additions and 26 deletions
|
|
@ -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}")
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue