diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index edae02e50f..6f5efd1642 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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}") diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b74fb8a35b..30d251727b 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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, diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index d19c67a01a..03a98c18dc 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -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 diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 616384386d..5aa0a181a6 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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, diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 367de196f7..e7e9e7336b 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -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: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a2c218269e..7f2119d118 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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. diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 91aa0cbf43..d402741d10 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -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