Previously HF_HUB_CACHE always defaulted to ~/.cache/huggingface/hub even when HF_HOME was explicitly set (e.g. in Docker). This caused models to download to the wrong location instead of the configured HF_HOME path.
305 lines
8.8 KiB
Python
305 lines
8.8 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
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import tempfile
|
|
|
|
|
|
def studio_root() -> Path:
|
|
return Path.home() / ".unsloth" / "studio"
|
|
|
|
|
|
def cache_root() -> Path:
|
|
"""Central cache directory for all studio downloads (models, datasets, etc.)."""
|
|
return Path.home() / ".unsloth" / "studio" / "cache"
|
|
|
|
|
|
def assets_root() -> Path:
|
|
return studio_root() / "assets"
|
|
|
|
|
|
def datasets_root() -> Path:
|
|
return assets_root() / "datasets"
|
|
|
|
|
|
def dataset_uploads_root() -> Path:
|
|
return datasets_root() / "uploads"
|
|
|
|
|
|
def recipe_datasets_root() -> Path:
|
|
return datasets_root() / "recipes"
|
|
|
|
|
|
def outputs_root() -> Path:
|
|
return studio_root() / "outputs"
|
|
|
|
|
|
def exports_root() -> Path:
|
|
return studio_root() / "exports"
|
|
|
|
|
|
def auth_root() -> Path:
|
|
return studio_root() / "auth"
|
|
|
|
|
|
def auth_db_path() -> Path:
|
|
return auth_root() / "auth.db"
|
|
|
|
|
|
def studio_db_path() -> Path:
|
|
return studio_root() / "studio.db"
|
|
|
|
|
|
def tmp_root() -> Path:
|
|
return Path(tempfile.gettempdir()) / "unsloth-studio"
|
|
|
|
|
|
def seed_uploads_root() -> Path:
|
|
return datasets_root() / "seed-uploads"
|
|
|
|
|
|
def unstructured_seed_cache_root() -> Path:
|
|
return tmp_root() / "unstructured-seed-cache"
|
|
|
|
|
|
def unstructured_uploads_root() -> Path:
|
|
return datasets_root() / "unstructured-uploads"
|
|
|
|
|
|
def oxc_validator_tmp_root() -> Path:
|
|
return tmp_root() / "oxc-validator"
|
|
|
|
|
|
def tensorboard_root() -> Path:
|
|
return studio_root() / "runs"
|
|
|
|
|
|
def ensure_dir(path: Path) -> Path:
|
|
path.mkdir(parents = True, exist_ok = True)
|
|
return path
|
|
|
|
|
|
def legacy_hf_cache_dir() -> Path:
|
|
"""Old Unsloth-specific HF hub cache, kept for backward-compat scanning."""
|
|
return cache_root() / "huggingface" / "hub"
|
|
|
|
|
|
def hf_default_cache_dir() -> Path:
|
|
"""Return the platform default HuggingFace hub cache (ignoring env overrides).
|
|
|
|
This is the location HF uses when no ``HF_HUB_CACHE`` / ``HF_HOME``
|
|
env var is set. We scan it so that models a user downloaded *before*
|
|
installing Unsloth Studio are still discovered.
|
|
"""
|
|
return Path.home() / ".cache" / "huggingface" / "hub"
|
|
|
|
|
|
def lmstudio_model_dirs() -> list[Path]:
|
|
"""Return LM Studio model directories that exist on disk."""
|
|
dirs: list[Path] = []
|
|
seen: set[Path] = set()
|
|
|
|
def _add(p: Path) -> None:
|
|
resolved = p.resolve()
|
|
if resolved not in seen and p.is_dir():
|
|
seen.add(resolved)
|
|
dirs.append(p)
|
|
|
|
# 1. Check LM Studio settings.json for custom downloads folder
|
|
settings_path = Path.home() / ".lmstudio" / "settings.json"
|
|
if settings_path.is_file():
|
|
try:
|
|
with open(settings_path) as f:
|
|
settings = json.load(f)
|
|
downloads = settings.get("downloadsFolder", "")
|
|
if downloads:
|
|
_add(Path(downloads).expanduser())
|
|
except Exception:
|
|
pass
|
|
|
|
# 2. LM Studio current default models directory (all platforms)
|
|
_add(Path.home() / ".lmstudio" / "models")
|
|
|
|
# 3. Legacy LM Studio cache location
|
|
_add(Path.home() / ".cache" / "lm-studio" / "models")
|
|
|
|
return dirs
|
|
|
|
|
|
def well_known_model_dirs() -> list[Path]:
|
|
"""Return directories commonly used by other local LLM tools.
|
|
|
|
Used by the folder browser to offer quick-pick chips. Returns only
|
|
paths that exist on disk, so the UI never shows dead chips. Order
|
|
reflects a rough "likelihood the user has models here" -- LM Studio
|
|
and Ollama first, then the generic fallbacks.
|
|
"""
|
|
candidates: list[Path] = []
|
|
|
|
# LM Studio (reuses the logic above, including settings.json override)
|
|
candidates.extend(lmstudio_model_dirs())
|
|
|
|
# Ollama -- both the user-level and common system-wide install paths
|
|
# (https://github.com/ollama/ollama/issues/733).
|
|
ollama_env = os.environ.get("OLLAMA_MODELS")
|
|
if ollama_env:
|
|
candidates.append(Path(ollama_env).expanduser())
|
|
candidates.append(Path.home() / ".ollama" / "models")
|
|
candidates.append(Path("/usr/share/ollama/.ollama/models"))
|
|
candidates.append(Path("/var/lib/ollama/.ollama/models"))
|
|
|
|
# HF hub cache root (separate from the explicit HF cache chip)
|
|
candidates.append(Path.home() / ".cache" / "huggingface" / "hub")
|
|
|
|
# Generic "my models" spots users tend to drop things into
|
|
for name in ("models", "Models"):
|
|
candidates.append(Path.home() / name)
|
|
|
|
# Deduplicate while preserving order; keep only extant dirs
|
|
out: list[Path] = []
|
|
seen: set[str] = set()
|
|
for p in candidates:
|
|
try:
|
|
resolved = str(p.resolve())
|
|
except OSError:
|
|
continue
|
|
if resolved in seen:
|
|
continue
|
|
if Path(resolved).is_dir():
|
|
seen.add(resolved)
|
|
out.append(Path(resolved))
|
|
return out
|
|
|
|
|
|
def _setup_cache_env() -> None:
|
|
"""Set cache environment variables for HuggingFace, uv, and vLLM.
|
|
|
|
Respects the standard HF cache resolution chain: explicit ``HF_HOME``
|
|
/ ``HF_HUB_CACHE`` env vars take priority, then ``XDG_CACHE_HOME``,
|
|
then the platform default (``~/.cache/huggingface``). The legacy
|
|
Unsloth cache is still *scanned* for models but is never set as the
|
|
active download target.
|
|
|
|
Only sets variables that are not already set by the user, so
|
|
explicit overrides (e.g. HF_HOME=/data/hf) are respected.
|
|
Works on Linux, macOS, and Windows.
|
|
"""
|
|
root = cache_root()
|
|
xdg_cache = Path(
|
|
os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")
|
|
).expanduser()
|
|
hf_default = xdg_cache / "huggingface"
|
|
hf_home = Path(os.environ.get("HF_HOME", str(hf_default)))
|
|
defaults: dict[str, str] = {
|
|
"HF_HOME": str(hf_default),
|
|
"HF_HUB_CACHE": str(hf_home / "hub"),
|
|
"HF_XET_CACHE": str(hf_home / "xet"),
|
|
"UV_CACHE_DIR": str(root / "uv"),
|
|
"VLLM_CACHE_ROOT": str(root / "vllm"),
|
|
}
|
|
for key, value in defaults.items():
|
|
if key not in os.environ:
|
|
os.environ[key] = value
|
|
Path(value).mkdir(parents = True, exist_ok = True)
|
|
|
|
|
|
def ensure_studio_directories() -> None:
|
|
"""Create all standard studio directories on startup."""
|
|
for dir_fn in (
|
|
studio_root,
|
|
assets_root,
|
|
datasets_root,
|
|
dataset_uploads_root,
|
|
recipe_datasets_root,
|
|
unstructured_uploads_root,
|
|
outputs_root,
|
|
exports_root,
|
|
auth_root,
|
|
tensorboard_root,
|
|
):
|
|
ensure_dir(dir_fn())
|
|
_setup_cache_env()
|
|
|
|
|
|
def _clean_relative_path(
|
|
path_value: str, *, strip_prefixes: tuple[str, ...] = ()
|
|
) -> Path:
|
|
path = Path(path_value).expanduser()
|
|
parts = [part for part in path.parts if part not in ("", ".")]
|
|
while parts and parts[0] in strip_prefixes:
|
|
parts = parts[1:]
|
|
return Path(*parts) if parts else Path()
|
|
|
|
|
|
def resolve_under_root(
|
|
path_value: str | None,
|
|
*,
|
|
root: Path,
|
|
strip_prefixes: tuple[str, ...] = (),
|
|
) -> Path:
|
|
if not path_value or not str(path_value).strip():
|
|
return root
|
|
|
|
path = Path(str(path_value).strip()).expanduser()
|
|
if path.is_absolute():
|
|
return path
|
|
|
|
cleaned = _clean_relative_path(str(path), strip_prefixes = strip_prefixes)
|
|
return root / cleaned
|
|
|
|
|
|
def resolve_output_dir(path_value: str | None = None) -> Path:
|
|
return resolve_under_root(
|
|
path_value,
|
|
root = outputs_root(),
|
|
strip_prefixes = ("outputs",),
|
|
)
|
|
|
|
|
|
def resolve_export_dir(path_value: str | None = None) -> Path:
|
|
return resolve_under_root(
|
|
path_value,
|
|
root = exports_root(),
|
|
strip_prefixes = ("exports",),
|
|
)
|
|
|
|
|
|
def resolve_tensorboard_dir(path_value: str | None = None) -> Path:
|
|
return resolve_under_root(
|
|
path_value,
|
|
root = tensorboard_root(),
|
|
strip_prefixes = ("runs", "tensorboard"),
|
|
)
|
|
|
|
|
|
def resolve_dataset_path(path_value: str) -> Path:
|
|
path = Path(path_value).expanduser()
|
|
if path.is_absolute():
|
|
return path
|
|
|
|
parts = [part for part in Path(path_value).parts if part not in ("", ".")]
|
|
if parts[:2] == ["assets", "datasets"]:
|
|
parts = parts[2:]
|
|
if parts and parts[0] == "uploads":
|
|
cleaned = Path(*parts[1:]) if len(parts) > 1 else Path()
|
|
return dataset_uploads_root() / cleaned
|
|
if parts and parts[0] == "recipes":
|
|
cleaned = Path(*parts[1:]) if len(parts) > 1 else Path()
|
|
return recipe_datasets_root() / cleaned
|
|
|
|
cleaned = Path(*parts) if parts else Path()
|
|
candidates = [
|
|
dataset_uploads_root() / cleaned,
|
|
recipe_datasets_root() / cleaned,
|
|
datasets_root() / cleaned,
|
|
dataset_uploads_root() / cleaned.name,
|
|
recipe_datasets_root() / cleaned.name,
|
|
]
|
|
for candidate in candidates:
|
|
if candidate.exists():
|
|
return candidate
|
|
return candidates[0]
|