6492d36719
- Overlapping requests lost offline mid-flight. A later guard saw the HF_HUB_OFFLINE that an earlier one had set and took the no-op branch, so when the earlier guard exited it restored the constants and sessions while the later request was still resolving hub files, dropping it back onto the retry path. Each guard now holds its own reference on the refcounted force_hf_offline window. A user-supplied offline variable is still left untouched, told apart via force_hf_offline_active(). - The socket-timeout fallback trusted a TCP handshake to the proxy, which only proves the proxy is up, not that it can reach the hub. A live proxy with a blackholed upstream therefore read as reachable. With a proxy configured the timeout now stays unreachable; the TCP check is only evidence when connecting to the endpoint directly. Verified: second guard engages and offline survives the first guard's exit, state fully restored after both; dead-upstream proxy reads unreachable while a slow direct endpoint still reads reachable; 9 concurrent metadata requests against an unreachable hub all return 200 in 5.1s total.
606 lines
22 KiB
Python
606 lines
22 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Shared backend utilities."""
|
|
|
|
import os
|
|
import structlog
|
|
import threading
|
|
import time
|
|
from loggers import get_logger
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import shutil
|
|
import tempfile
|
|
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
# ── Offline / HF-cache helpers ──────────────────────────────────
|
|
# An offline load must never touch the network (a DNS-dead session hangs on hub retries);
|
|
# these read the local HF cache the load itself uses.
|
|
|
|
_HF_OFFLINE_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
|
|
|
|
|
|
def hf_env_offline() -> bool:
|
|
"""True when HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE requests offline mode.
|
|
|
|
Also honors TRANSFORMERS_OFFLINE (hub honors only HF_HUB_OFFLINE) since users set it
|
|
to keep transformers loads local.
|
|
"""
|
|
for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"):
|
|
if os.environ.get(var, "").strip().lower() in _HF_OFFLINE_TRUE_VALUES:
|
|
return True
|
|
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.
|
|
|
|
Uses getaddrinfo, not gethostbyname: the latter is IPv4-only and would call an
|
|
AAAA-only mirror or an IPv6 literal dead.
|
|
"""
|
|
result: list = [None]
|
|
|
|
def _probe() -> None:
|
|
import socket as _socket
|
|
try:
|
|
_socket.getaddrinfo(host, None)
|
|
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_connect_target(endpoint: Optional[str] = None):
|
|
"""(host, port) egress actually has to reach: the proxy when one applies, else the endpoint."""
|
|
from urllib.parse import urlparse
|
|
|
|
url = endpoint or hf_endpoint_url()
|
|
parsed = urlparse(url)
|
|
default_port = 443 if parsed.scheme == "https" else 80
|
|
try:
|
|
import urllib.request
|
|
|
|
proxies = urllib.request.getproxies()
|
|
proxy = proxies.get(parsed.scheme) or proxies.get("all")
|
|
host = parsed.hostname or ""
|
|
if proxy and not (host and urllib.request.proxy_bypass(host)):
|
|
p = urlparse(proxy if "://" in proxy else "http://" + proxy)
|
|
return p.hostname, p.port or 80
|
|
except Exception:
|
|
pass
|
|
return parsed.hostname, parsed.port or default_port
|
|
|
|
|
|
def hf_tcp_reachable(timeout: float = 3.0, endpoint: Optional[str] = None) -> bool:
|
|
"""True when a TCP connection to the hub (or its proxy) can be established.
|
|
|
|
Separates "no egress" from "slow to answer": a loaded server still completes the
|
|
handshake promptly, whereas a blackholed route times out. A refused connection also
|
|
counts as reachable, since something answered.
|
|
"""
|
|
import socket as _socket
|
|
|
|
host, port = hf_connect_target(endpoint)
|
|
if not host:
|
|
return False
|
|
try:
|
|
with _socket.create_connection((host, port), timeout = timeout):
|
|
return True
|
|
except ConnectionRefusedError:
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
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
|
|
# cache after the user reconnected, failing it if the model is not cached.
|
|
_HF_REACHABILITY_TTL_S = 5.0
|
|
_hf_reachability: Optional[tuple] = None
|
|
_hf_reachability_lock = threading.Lock()
|
|
|
|
|
|
def _reachability_fresh(entry) -> bool:
|
|
"""True while a cached (timestamp, unreachable) verdict may still be reused."""
|
|
return entry is not None and (time.monotonic() - entry[0]) < _HF_REACHABILITY_TTL_S
|
|
|
|
|
|
def hf_probe_disabled() -> bool:
|
|
"""True when UNSLOTH_OFFLINE_PROBE opts out of the reachability probe."""
|
|
return os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {
|
|
"0",
|
|
"false",
|
|
"no",
|
|
"off",
|
|
}
|
|
|
|
|
|
def reset_hf_reachability_cache() -> None:
|
|
"""Drop the memoised verdict so the next call re-probes (tests, network changes)."""
|
|
global _hf_reachability
|
|
with _hf_reachability_lock:
|
|
_hf_reachability = None
|
|
|
|
|
|
def hf_unreachable(timeout: int = 3) -> bool:
|
|
"""True when the HF endpoint is unreachable, memoised for _HF_REACHABILITY_TTL_S.
|
|
|
|
Resolving DNS does not mean the Hub is reachable: a router that is up with the WAN
|
|
down, a captive portal or a stale DNS cache all answer lookups while every request
|
|
then burns huggingface_hub's retry backoff. This is the bounded, proxy-aware probe
|
|
the export path already uses; disable it with UNSLOTH_OFFLINE_PROBE=0.
|
|
|
|
Fails open: an unavailable probe reports reachable, so the load decides as it does today.
|
|
"""
|
|
if hf_probe_disabled():
|
|
return False
|
|
|
|
global _hf_reachability
|
|
cached = _hf_reachability
|
|
if _reachability_fresh(cached):
|
|
return cached[1]
|
|
|
|
with _hf_reachability_lock:
|
|
cached = _hf_reachability
|
|
if _reachability_fresh(cached):
|
|
return cached[1]
|
|
try:
|
|
from utils.transformers_version import hf_endpoint_unreachable
|
|
unreachable = hf_endpoint_unreachable(timeout)
|
|
except Exception:
|
|
unreachable = False
|
|
_hf_reachability = (time.monotonic(), unreachable)
|
|
return unreachable
|
|
|
|
|
|
def _reset_hf_sessions() -> None:
|
|
"""Drop cached hub sessions so they remount with the current offline adapter."""
|
|
try:
|
|
from huggingface_hub.utils import _http
|
|
|
|
for name in ("_get_session_from_cache", "get_session"):
|
|
cache_clear = getattr(getattr(_http, name, None), "cache_clear", None)
|
|
if cache_clear is not None:
|
|
cache_clear()
|
|
reset = getattr(_http, "reset_sessions", None)
|
|
if reset is not None:
|
|
reset()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# Process-global, so nested/concurrent loads refcount rather than restore out from under
|
|
# each other.
|
|
_force_offline_depth = 0
|
|
_force_offline_saved: list = []
|
|
_force_offline_saved_env: dict = {}
|
|
_force_offline_lock = threading.Lock()
|
|
|
|
_OFFLINE_ENV_KEYS = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")
|
|
_OFFLINE_CONSTANTS = (
|
|
("huggingface_hub.constants", ("HF_HUB_OFFLINE",)),
|
|
("transformers.utils.hub", ("_is_offline_mode", "OFFLINE")),
|
|
)
|
|
|
|
|
|
def force_hf_offline_active() -> bool:
|
|
"""True while a force_hf_offline window is open anywhere in this process.
|
|
|
|
Lets a concurrent caller tell our own forced offline apart from one the user set, so
|
|
it can hold its own reference instead of no-opping and losing offline when the first
|
|
window exits.
|
|
"""
|
|
with _force_offline_lock:
|
|
return _force_offline_depth > 0
|
|
|
|
|
|
@contextmanager
|
|
def force_hf_offline():
|
|
"""Force HF offline for this block, in-process.
|
|
|
|
Setting the env vars is not enough once the process is running: huggingface_hub and
|
|
transformers read their offline constants at import, and hub sessions cache a
|
|
non-offline adapter. Flip the constants too and rebuild the sessions, so hub calls
|
|
fail fast instead of retrying. Everything is restored on exit.
|
|
"""
|
|
global _force_offline_depth, _force_offline_saved, _force_offline_saved_env
|
|
import importlib
|
|
|
|
with _force_offline_lock:
|
|
if _force_offline_depth == 0:
|
|
saved: list = []
|
|
saved_env: dict = {}
|
|
# Snapshot constants BEFORE forcing the env, else a module first imported
|
|
# inside the window reads the "1" and we would restore it as offline.
|
|
for mod_name, attrs in _OFFLINE_CONSTANTS:
|
|
try:
|
|
mod = importlib.import_module(mod_name)
|
|
except Exception:
|
|
continue
|
|
for attr in attrs:
|
|
if hasattr(mod, attr):
|
|
saved.append((mod, attr, getattr(mod, attr)))
|
|
for key in _OFFLINE_ENV_KEYS:
|
|
saved_env[key] = os.environ.get(key)
|
|
os.environ[key] = "1"
|
|
for mod, attr, _ in saved:
|
|
try:
|
|
setattr(mod, attr, True)
|
|
except Exception:
|
|
pass
|
|
_force_offline_saved = saved
|
|
_force_offline_saved_env = saved_env
|
|
_reset_hf_sessions()
|
|
_force_offline_depth += 1
|
|
try:
|
|
yield
|
|
finally:
|
|
with _force_offline_lock:
|
|
_force_offline_depth -= 1
|
|
if _force_offline_depth == 0:
|
|
for mod, attr, val in _force_offline_saved:
|
|
try:
|
|
setattr(mod, attr, val)
|
|
except Exception:
|
|
pass
|
|
_force_offline_saved = []
|
|
for key, val in _force_offline_saved_env.items():
|
|
if val is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = val
|
|
_force_offline_saved_env = {}
|
|
_reset_hf_sessions()
|
|
|
|
|
|
def st_repo_id_candidates(model_name: str) -> list:
|
|
"""Repo ids a Sentence-Transformers load may resolve model_name to; a slashless name
|
|
also resolves under the sentence-transformers/ namespace, so both are candidates."""
|
|
name = (model_name or "").strip().strip("/")
|
|
if not name:
|
|
return []
|
|
candidates = [name]
|
|
if "/" not in name:
|
|
candidates.append(f"sentence-transformers/{name}")
|
|
return candidates
|
|
|
|
|
|
def _expand_path(raw: str) -> Path:
|
|
"""Expand ~ and $VARS as huggingface_hub does, so the gate resolves the loader's dir."""
|
|
return Path(os.path.expandvars(os.path.expanduser(raw)))
|
|
|
|
|
|
def _hf_cache_roots() -> list:
|
|
"""Cache roots to search for a model's local snapshot, most-authoritative first.
|
|
|
|
The app's selected hub cache (set via /settings) is searched first: after a
|
|
no-restart cache switch the process env is stale, yet the loader reads the
|
|
selected cache via ``cache_folder=active_hf_hub_cache()``, so the snapshot
|
|
and offline security lookups must match where it actually loads. The env
|
|
precedence (SENTENCE_TRANSFORMERS_HOME, HF_HUB_CACHE, HF_HOME/hub,
|
|
~/.cache/huggingface/hub) follows so a copy still in a previous cache resolves."""
|
|
roots: list = []
|
|
seen: set = set()
|
|
|
|
def _add(path) -> None:
|
|
if path is None:
|
|
return
|
|
expanded = _expand_path(str(path))
|
|
key = str(expanded)
|
|
if key not in seen:
|
|
seen.add(key)
|
|
roots.append(expanded)
|
|
|
|
try:
|
|
from utils.hf_cache_settings import get_hf_cache_paths
|
|
_add(get_hf_cache_paths().hub_cache)
|
|
except Exception:
|
|
pass
|
|
|
|
if st_home := os.environ.get("SENTENCE_TRANSFORMERS_HOME"):
|
|
_add(st_home)
|
|
if hub := (os.environ.get("HF_HUB_CACHE") or os.environ.get("HUGGINGFACE_HUB_CACHE")):
|
|
_add(hub)
|
|
if hf_home := os.environ.get("HF_HOME"):
|
|
_add(_expand_path(hf_home) / "hub")
|
|
if not roots:
|
|
_add(Path.home() / ".cache" / "huggingface" / "hub")
|
|
return roots
|
|
|
|
|
|
def hf_cache_snapshot_dir(model_name: str) -> Optional[Path]:
|
|
"""Active local snapshot dir for model_name's main revision, or None if not cached.
|
|
Reads refs/main then snapshots/<commit>; no network. Tries the ST alias for slashless names."""
|
|
try:
|
|
from huggingface_hub.file_download import repo_folder_name
|
|
except Exception:
|
|
repo_folder_name = None
|
|
for cache_root in _hf_cache_roots():
|
|
for repo_id in st_repo_id_candidates(model_name):
|
|
try:
|
|
if repo_folder_name is not None:
|
|
folder = repo_folder_name(repo_id = repo_id, repo_type = "model")
|
|
else:
|
|
folder = "models--" + repo_id.replace("/", "--")
|
|
repo_dir = cache_root / folder
|
|
ref = repo_dir / "refs" / "main"
|
|
if not ref.is_file():
|
|
continue
|
|
commit = ref.read_text(encoding = "utf-8").strip()
|
|
if not commit:
|
|
continue
|
|
snapshot = repo_dir / "snapshots" / commit
|
|
if snapshot.is_dir():
|
|
return snapshot
|
|
except (OSError, UnicodeDecodeError):
|
|
continue
|
|
return None
|
|
|
|
|
|
# A weight file plus a config distinguishes a real cached model from a metadata-only
|
|
# partial cache that resolves refs/main but would fail at load time.
|
|
_LOADABLE_WEIGHT_SUFFIXES = frozenset({".safetensors", ".bin", ".gguf", ".pt", ".pth", ".ckpt"})
|
|
|
|
|
|
def hf_cache_snapshot_is_loadable(model_name: str) -> bool:
|
|
"""True when model_name's snapshot is cached and loadable: a config (config.json or
|
|
modules.json) plus at least one weight file, not a metadata-only partial cache. No network."""
|
|
snapshot = hf_cache_snapshot_dir(model_name)
|
|
if snapshot is None:
|
|
return False
|
|
try:
|
|
has_config = (snapshot / "config.json").is_file() or (snapshot / "modules.json").is_file()
|
|
if not has_config:
|
|
return False
|
|
for path in snapshot.rglob("*"):
|
|
if path.suffix.lower() in _LOADABLE_WEIGHT_SUFFIXES and path.is_file():
|
|
return True
|
|
except OSError:
|
|
return False
|
|
return False
|
|
|
|
|
|
# ── Client-safe error helpers ───────────────────────────────────
|
|
# Never return raw exception text to clients; log server-side, return generic.
|
|
|
|
|
|
def safe_error_detail(error: Exception, fallback: str = "An internal error occurred") -> str:
|
|
"""Map an exception to a generic, client-safe message (never raw
|
|
``str(error)``, which can leak paths). Log the real exception server-side.
|
|
"""
|
|
text = str(error).lower()
|
|
if (
|
|
isinstance(error, (ConnectionError, TimeoutError))
|
|
or "connection" in text
|
|
or "timed out" in text
|
|
or "timeout" in text
|
|
):
|
|
return "Could not reach an upstream service. Please try again."
|
|
if "out of memory" in text or "cuda error" in text:
|
|
return "Ran out of memory. Try a smaller model or shorter input."
|
|
return fallback
|
|
|
|
|
|
def safe_curated_detail(error: Exception, fallback: str = "An internal error occurred") -> str:
|
|
"""Client-safe text for curated domain/validation exceptions.
|
|
|
|
Keeps the message (paths stripped) instead of a generic fallback; for known
|
|
exception types only (use ``safe_error_detail`` for generic ``Exception``).
|
|
"""
|
|
from utils.native_path_leases import redact_native_paths
|
|
|
|
msg = redact_native_paths(str(error)).strip()
|
|
return msg or fallback
|
|
|
|
|
|
def log_and_http_error(
|
|
error: Exception,
|
|
status_code: int,
|
|
public_message: str,
|
|
*,
|
|
event: str = "request_failed",
|
|
log = None,
|
|
):
|
|
"""Log ``error`` in full server-side and return an ``HTTPException`` whose
|
|
``detail`` is only ``public_message`` -- never the raw exception text.
|
|
|
|
Usage: raise log_and_http_error(e, 500, "Failed to start training")
|
|
"""
|
|
from fastapi import HTTPException
|
|
|
|
# exc_info=error works for both structlog and stdlib loggers.
|
|
(log or logger).error(f"{event}: {error}", exc_info = error)
|
|
return HTTPException(status_code = status_code, detail = public_message)
|
|
|
|
|
|
@contextmanager
|
|
def without_hf_auth():
|
|
"""
|
|
Temporarily disable HuggingFace authentication.
|
|
|
|
Usage:
|
|
with without_hf_auth():
|
|
# Code that should run without cached tokens
|
|
model_info(model_name, token=None)
|
|
"""
|
|
saved_env = {}
|
|
env_vars = ["HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_HOME"]
|
|
for var in env_vars:
|
|
if var in os.environ:
|
|
saved_env[var] = os.environ[var]
|
|
del os.environ[var]
|
|
|
|
saved_disable = os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN")
|
|
os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1"
|
|
|
|
# Move token files aside temporarily
|
|
token_files = []
|
|
token_locations = [
|
|
Path.home() / ".cache" / "huggingface" / "token",
|
|
Path.home() / ".huggingface" / "token",
|
|
]
|
|
|
|
for token_loc in token_locations:
|
|
if token_loc.exists():
|
|
temp = tempfile.NamedTemporaryFile(delete = False)
|
|
temp.close()
|
|
shutil.move(str(token_loc), temp.name)
|
|
token_files.append((token_loc, temp.name))
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
# Restore tokens
|
|
for original, temp in token_files:
|
|
try:
|
|
original.parent.mkdir(parents = True, exist_ok = True)
|
|
shutil.move(temp, str(original))
|
|
except Exception as e:
|
|
logger.error(f"Failed to restore token {original}: {e}")
|
|
|
|
# Restore env
|
|
for var, value in saved_env.items():
|
|
os.environ[var] = value
|
|
|
|
if saved_disable is not None:
|
|
os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = saved_disable
|
|
else:
|
|
os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None)
|
|
|
|
|
|
def is_hf_authentication_error(error: Exception) -> bool:
|
|
"""Return whether an exception chain contains a definitive HF auth failure."""
|
|
seen: set[int] = set()
|
|
current: BaseException | None = error
|
|
while current is not None and id(current) not in seen:
|
|
seen.add(id(current))
|
|
response = getattr(current, "response", None)
|
|
status = getattr(response, "status_code", None)
|
|
try:
|
|
if status is not None and int(status) == 401:
|
|
return True
|
|
except (TypeError, ValueError):
|
|
pass
|
|
message = str(current).lower()
|
|
if "invalid user token" in message or "invalid hf token" in message:
|
|
return True
|
|
current = current.__cause__ or current.__context__
|
|
return False
|
|
|
|
|
|
def format_error_message(error: Exception, model_name: str) -> str:
|
|
"""
|
|
Format a user-friendly error message for common load issues.
|
|
|
|
Args:
|
|
error: The exception that occurred
|
|
model_name: Name of the model being loaded
|
|
"""
|
|
error_str = str(error).lower()
|
|
model_short = model_name.split("/")[-1] if "/" in model_name else model_name
|
|
|
|
if "repository not found" in error_str or "404" in error_str:
|
|
return f"Model '{model_short}' not found. Check the model name."
|
|
|
|
if "401" in error_str or "unauthorized" in error_str:
|
|
return f"Authentication failed for '{model_short}'. Please provide a valid HF token."
|
|
|
|
if "gated" in error_str or "access to model" in error_str:
|
|
return f"Model '{model_short}' requires authentication. Please provide a valid HF token."
|
|
|
|
if "invalid user token" in error_str:
|
|
return "Invalid HF token. Please check your token and try again."
|
|
|
|
if (
|
|
"out of memory" in error_str
|
|
or "out of device memory" in error_str
|
|
or "out_of_device_memory" in error_str # ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY
|
|
or "out_of_host_memory" in error_str # ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY
|
|
or "not enough memory" in error_str
|
|
or "cannot allocate memory" in error_str
|
|
or "memory allocation failed" in error_str
|
|
or "cublas_status_alloc_failed" in error_str # cuBLAS workspace OOM
|
|
or ("cuda error" in error_str and "alloc" in error_str)
|
|
or ("xpu" in error_str and ("alloc" in error_str or "memory" in error_str))
|
|
or isinstance(error, MemoryError)
|
|
or ("mlx" in error_str and ("memory" in error_str or "allocate" in error_str))
|
|
):
|
|
# Resolve get_device() at call time (not import time) so tests that
|
|
# monkey-patch utils.hardware.get_device after this module is loaded
|
|
# still see the patched backend.
|
|
from utils.hardware import get_device
|
|
|
|
device = get_device()
|
|
device_label = {
|
|
"cuda": "GPU",
|
|
"xpu": "Intel GPU",
|
|
"mlx": "Apple Silicon GPU",
|
|
"cpu": "system",
|
|
}.get(device.value, "GPU")
|
|
return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory."
|
|
|
|
return str(error)
|