unsloth/studio/backend/utils/paths/storage_roots.py
Roland Tannous 1b3ffd2ceb fix: derive HF_HUB_CACHE from HF_HOME when set
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.
2026-04-16 21:54:47 +00:00

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]