7bd8e64921* Studio: honor custom HF_HOME for model download and load _setup_cache_env always derived HF_HUB_CACHE and HF_XET_CACHE from XDG_CACHE_HOME / ~/.cache, ignoring a user-set HF_HOME. Because it sets HF_HUB_CACHE explicitly and that variable takes precedence over HF_HOME in huggingface_hub, the hub cache was pinned to the standard location: a model already present under a custom HF_HOME was detected but then re-downloaded from scratch on load. Seed HF_HUB_CACHE and HF_XET_CACHE from HF_HOME when the user set it (HF's own default is $HF_HOME/hub and $HF_HOME/xet), and honor the legacy HUGGINGFACE_HUB_CACHE alias. The hub download workers call snapshot_download without a cache_dir for both the Xet and HTTP-fallback paths, so they follow HF_HUB_CACHE; fixing it here unifies detection and both transports on one root. Explicit HF_HUB_CACHE / HF_XET_CACHE stay untouched. Adds tests for the custom-HF_HOME, default, explicit-override, and legacy-alias cases. Fixes #5182. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: do not crash startup when a custom HF_HOME is not writable Seeding HF_HUB_CACHE/HF_XET_CACHE from HF_HOME means _setup_cache_env now mkdir's under a user-controlled path. A non-writable or not-yet-mounted HF_HOME (typo, offline drive) would raise and crash startup, where the old code silently fell back. Make the mkdir best-effort; the env var is still set, so HF reports a clear error at download time. Adds a regression test. * Studio: strip blank HF_HOME and isolate cache-env tests Address review: a whitespace-only HF_HOME no longer derives " /hub"; strip it and fall back to the default (matches studio_root). Tests set UNSLOTH_STUDIO_HOME to a tmp dir so _setup_cache_env's UV/VLLM mkdirs do not touch the real ~/.unsloth/studio. Adds a whitespace regression test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
125 lines
3.9 KiB
Python
125 lines
3.9 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
|
|
|
|
"""_setup_cache_env() must seed HF_HUB_CACHE / HF_XET_CACHE from a user-set
|
|
HF_HOME, so models download to and load from the same custom location (issue
|
|
#5182). Both the Xet and HTTP-fallback download workers call snapshot_download
|
|
without a cache_dir, so they follow HF_HUB_CACHE; getting it right here fixes
|
|
detection and both transports at once.
|
|
"""
|
|
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|
if _BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, _BACKEND_DIR)
|
|
|
|
_STORAGE_ROOTS_PATH = Path(__file__).resolve().parent.parent / "utils/paths/storage_roots.py"
|
|
|
|
|
|
@pytest.fixture(autouse = True)
|
|
def _isolate_studio_home(monkeypatch, tmp_path):
|
|
# Keep _setup_cache_env's UV/VLLM mkdirs out of the real ~/.unsloth/studio.
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio"))
|
|
|
|
|
|
def _load_storage_roots():
|
|
spec = importlib.util.spec_from_file_location("storage_roots_under_test", _STORAGE_ROOTS_PATH)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _clear_hf_env(monkeypatch):
|
|
for key in ("HF_HOME", "HF_HUB_CACHE", "HF_XET_CACHE", "HUGGINGFACE_HUB_CACHE"):
|
|
monkeypatch.delenv(key, raising = False)
|
|
|
|
|
|
def test_custom_hf_home_seeds_hub_and_xet(monkeypatch, tmp_path):
|
|
sr = _load_storage_roots()
|
|
_clear_hf_env(monkeypatch)
|
|
custom = tmp_path / "shared" / "huggingface"
|
|
monkeypatch.setenv("HF_HOME", str(custom))
|
|
|
|
sr._setup_cache_env()
|
|
|
|
import os
|
|
|
|
assert os.environ["HF_HUB_CACHE"] == str(custom / "hub")
|
|
assert os.environ["HF_XET_CACHE"] == str(custom / "xet")
|
|
|
|
|
|
def test_default_when_hf_home_unset(monkeypatch, tmp_path):
|
|
sr = _load_storage_roots()
|
|
_clear_hf_env(monkeypatch)
|
|
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg"))
|
|
|
|
sr._setup_cache_env()
|
|
|
|
import os
|
|
|
|
expected = tmp_path / "xdg" / "huggingface"
|
|
assert os.environ["HF_HUB_CACHE"] == str(expected / "hub")
|
|
|
|
|
|
def test_explicit_hub_cache_is_not_overridden(monkeypatch, tmp_path):
|
|
sr = _load_storage_roots()
|
|
_clear_hf_env(monkeypatch)
|
|
monkeypatch.setenv("HF_HOME", str(tmp_path / "home"))
|
|
explicit = tmp_path / "explicit" / "hub"
|
|
monkeypatch.setenv("HF_HUB_CACHE", str(explicit))
|
|
|
|
sr._setup_cache_env()
|
|
|
|
import os
|
|
|
|
assert os.environ["HF_HUB_CACHE"] == str(explicit)
|
|
|
|
|
|
def test_legacy_huggingface_hub_cache_alias_is_honored(monkeypatch, tmp_path):
|
|
sr = _load_storage_roots()
|
|
_clear_hf_env(monkeypatch)
|
|
monkeypatch.setenv("HF_HOME", str(tmp_path / "home"))
|
|
legacy = tmp_path / "legacy" / "hub"
|
|
monkeypatch.setenv("HUGGINGFACE_HUB_CACHE", str(legacy))
|
|
|
|
sr._setup_cache_env()
|
|
|
|
import os
|
|
|
|
assert os.environ["HF_HUB_CACHE"] == str(legacy)
|
|
|
|
|
|
def test_whitespace_hf_home_falls_back_to_default(monkeypatch, tmp_path):
|
|
# A blank/whitespace HF_HOME must not become " /hub"; fall back to default.
|
|
sr = _load_storage_roots()
|
|
_clear_hf_env(monkeypatch)
|
|
monkeypatch.setenv("HF_HOME", " ")
|
|
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg"))
|
|
|
|
sr._setup_cache_env()
|
|
|
|
import os
|
|
|
|
assert os.environ["HF_HUB_CACHE"] == str(tmp_path / "xdg" / "huggingface" / "hub")
|
|
|
|
|
|
def test_unwritable_hf_home_does_not_crash(monkeypatch, tmp_path):
|
|
# HF_HOME under a regular file -> mkdir fails; startup must not crash and the
|
|
# env var is still set (HF surfaces a clear error later, at download time).
|
|
blocker = tmp_path / "blocker"
|
|
blocker.write_text("not a dir")
|
|
unwritable = blocker / "hf"
|
|
sr = _load_storage_roots()
|
|
_clear_hf_env(monkeypatch)
|
|
monkeypatch.setenv("HF_HOME", str(unwritable))
|
|
|
|
sr._setup_cache_env() # must not raise
|
|
|
|
import os
|
|
|
|
assert os.environ["HF_HUB_CACHE"] == str(unwritable / "hub")
|