dbb06ff60eAdds a configurable Hugging Face model download cache location to Unsloth Studio, selectable from Settings, with per-cache download manifests, scoped deletion, and read-only inventory of previously selected caches.
414 lines
14 KiB
Python
414 lines
14 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
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from hub.schemas.datasets import CheckFormatRequest, LocalDatasetItem
|
|
from hub.services.datasets import cache_inventory, downloads, formatting, local
|
|
from hub.utils import download_manifest, download_registry, state_dir
|
|
|
|
|
|
class _Upload:
|
|
def __init__(self, filename: str, payload: bytes):
|
|
self.filename = filename
|
|
self._payload = payload
|
|
self._offset = 0
|
|
|
|
async def read(self, size: int) -> bytes:
|
|
if self._offset >= len(self._payload):
|
|
return b""
|
|
chunk = self._payload[self._offset : self._offset + size]
|
|
self._offset += len(chunk)
|
|
return chunk
|
|
|
|
|
|
def test_dataset_cache_scan_merges_raw_and_processed_rows(monkeypatch):
|
|
raw_repo = SimpleNamespace(
|
|
repo_id = "Org/Data",
|
|
repo_type = "dataset",
|
|
repo_path = "/cache/datasets--Org--Data",
|
|
size_on_disk = 100,
|
|
revisions = [SimpleNamespace(files = [], commit_hash = "abc")],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_collect_hf_cache_scans",
|
|
lambda: ([SimpleNamespace(repos = [raw_repo])], {"/cache"}),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda _repo_type, _repo_id, _cache_dir: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_scan_hub_dataset_cache_dirs",
|
|
lambda: [],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_scan_processed_dataset_caches",
|
|
lambda: [
|
|
{
|
|
"repo_id": "org/data",
|
|
"size_bytes": 250,
|
|
"cache_path": "/processed/org___data",
|
|
"processed_cache": True,
|
|
"partial": False,
|
|
}
|
|
],
|
|
)
|
|
|
|
rows = cache_inventory._scan_hf_dataset_caches()
|
|
|
|
assert len(rows) == 1
|
|
assert rows[0]["repo_id"] == "Org/Data"
|
|
assert rows[0]["size_bytes"] == 250
|
|
assert rows[0]["partial"] is False
|
|
|
|
|
|
def test_delete_cached_dataset_scopes_delete_to_selected_root(monkeypatch, tmp_path):
|
|
"""A dataset present in the active cache and a previously selected cache is
|
|
deleted only from the selected root, so the other cache's copy survives."""
|
|
calls = []
|
|
target_hub = tmp_path / "active" / "hub"
|
|
other_hub = tmp_path / "previous" / "hub"
|
|
for hub in (target_hub, other_hub):
|
|
(hub / "datasets--Org--Data").mkdir(parents = True)
|
|
|
|
class _DeleteStrategy:
|
|
def __init__(self, label: str):
|
|
self.label = label
|
|
|
|
def execute(self):
|
|
calls.append(self.label)
|
|
|
|
def _cache(label: str, hub):
|
|
return SimpleNamespace(
|
|
cache_dir = label,
|
|
repos = [
|
|
SimpleNamespace(
|
|
repo_type = "dataset",
|
|
repo_id = "Org/Data",
|
|
repo_path = str(hub / "datasets--Org--Data"),
|
|
revisions = [SimpleNamespace(commit_hash = f"{label}-rev")],
|
|
)
|
|
],
|
|
delete_revisions = lambda *_revs, _label = label: _DeleteStrategy(_label),
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_collect_hf_cache_scans",
|
|
lambda: ([_cache("active", target_hub), _cache("previous", other_hub)], set()),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_delete_processed_dataset_cache",
|
|
lambda _repo_id, **_kwargs: (False, []),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.download_manifest,
|
|
"purge_all_state_for_repo",
|
|
lambda *_args, **_kwargs: 0,
|
|
)
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = target_hub),
|
|
)
|
|
monkeypatch.setattr(
|
|
"hub.utils.hf_cache_state.hf_cache_roots",
|
|
lambda: [target_hub, other_hub],
|
|
)
|
|
|
|
result = cache_inventory._delete_cached_dataset_blocking("Org/Data")
|
|
|
|
assert result == {"status": "deleted", "repo_id": "Org/Data"}
|
|
# Only the selected (active) cache's revision is deleted; the previous
|
|
# cache's copy is never touched.
|
|
assert calls == ["active"]
|
|
assert not (target_hub / "datasets--Org--Data").exists()
|
|
assert (other_hub / "datasets--Org--Data").exists()
|
|
|
|
|
|
def test_delete_processed_only_dataset_accepts_processed_cache_path(monkeypatch, tmp_path):
|
|
"""A processed-only dataset row sends its Arrow cache path (<owner>___<repo>
|
|
under HF_DATASETS_CACHE), which is not a Hub datasets-- dir. The delete must
|
|
accept it and run the processed-cache delete instead of raising 400."""
|
|
datasets_root = tmp_path / "datasets"
|
|
processed_dir = datasets_root / "Org___Data"
|
|
processed_dir.mkdir(parents = True)
|
|
|
|
# No Hub-cache copy exists; only the processed Arrow cache holds this repo.
|
|
monkeypatch.setattr(cache_inventory, "_collect_hf_cache_scans", lambda: ([], set()))
|
|
monkeypatch.setattr(cache_inventory, "_hf_datasets_cache_roots", lambda: [datasets_root])
|
|
processed_calls: list[str] = []
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_delete_processed_dataset_cache",
|
|
lambda repo_id, **_kwargs: (processed_calls.append(repo_id) or True, []),
|
|
)
|
|
|
|
result = cache_inventory._delete_cached_dataset_blocking("Org/Data", str(processed_dir))
|
|
|
|
assert result == {"status": "deleted", "repo_id": "Org/Data"}
|
|
assert processed_calls == ["Org/Data"]
|
|
|
|
|
|
def test_delete_processed_dataset_scopes_to_selected_root(monkeypatch, tmp_path):
|
|
"""A dataset processed under two HF_DATASETS_CACHE roots is deleted only from
|
|
the selected root; the copy under the other cache home survives (real delete,
|
|
not stubbed)."""
|
|
selected_root = tmp_path / "selected" / "datasets"
|
|
other_root = tmp_path / "other" / "datasets"
|
|
for root in (selected_root, other_root):
|
|
(root / "Org___Data").mkdir(parents = True)
|
|
|
|
monkeypatch.setattr(cache_inventory, "_collect_hf_cache_scans", lambda: ([], set()))
|
|
monkeypatch.setattr(
|
|
cache_inventory, "_hf_datasets_cache_roots", lambda: [selected_root, other_root]
|
|
)
|
|
|
|
result = cache_inventory._delete_cached_dataset_blocking(
|
|
"Org/Data", str(selected_root / "Org___Data")
|
|
)
|
|
|
|
assert result == {"status": "deleted", "repo_id": "Org/Data"}
|
|
assert not (selected_root / "Org___Data").exists() # the selected copy is deleted
|
|
assert (other_root / "Org___Data").exists() # the other cache home is untouched
|
|
|
|
|
|
def test_delete_cached_dataset_purges_blob_only_repo_dir(monkeypatch):
|
|
"""A blob-only ``datasets--owner--repo`` dir (no usable snapshot/refs) is
|
|
fully removable: purge_partial_repo alone clears only ``.incomplete`` files
|
|
and would leave the complete blobs and the row."""
|
|
purged_dirs: list[str] = []
|
|
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_collect_hf_cache_scans",
|
|
lambda: ([], set()),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_delete_processed_dataset_cache",
|
|
lambda _repo_id, **_kwargs: (False, []),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"purge_repo_cache_dirs",
|
|
lambda _repo_type, repo_id, **_kwargs: purged_dirs.append(repo_id) or True,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"purge_partial_repo",
|
|
lambda *_args, **_kwargs: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.download_manifest,
|
|
"purge_all_state_for_repo",
|
|
lambda *_args, **_kwargs: 0,
|
|
)
|
|
|
|
result = cache_inventory._delete_cached_dataset_blocking("Org/Data")
|
|
|
|
assert result == {"status": "deleted", "repo_id": "Org/Data"}
|
|
assert purged_dirs == ["Org/Data"]
|
|
|
|
|
|
def test_delete_cached_dataset_absent_everywhere_raises_404(monkeypatch):
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_collect_hf_cache_scans",
|
|
lambda: ([], set()),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"_delete_processed_dataset_cache",
|
|
lambda _repo_id, **_kwargs: (False, []),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"purge_repo_cache_dirs",
|
|
lambda *_args, **_kwargs: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"purge_partial_repo",
|
|
lambda *_args, **_kwargs: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.download_manifest,
|
|
"purge_all_state_for_repo",
|
|
lambda *_args, **_kwargs: 0,
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
cache_inventory._delete_cached_dataset_blocking("Org/Missing")
|
|
|
|
assert exc_info.value.status_code == 404
|
|
|
|
|
|
def test_check_format_rejects_invalid_path_as_400():
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
formatting.check_format_response(CheckFormatRequest(dataset_name = "../../etc/passwd"))
|
|
|
|
assert exc_info.value.status_code == 400
|
|
|
|
|
|
def test_dataset_download_status_preserves_idle_shape():
|
|
status = downloads._dataset_status("Org/Data")
|
|
|
|
assert status.state == "idle"
|
|
assert status.error is None
|
|
|
|
|
|
def test_dataset_download_registry_key_is_case_insensitive():
|
|
registry = download_registry.DownloadRegistry()
|
|
|
|
claimed, state = registry.claim(
|
|
"Org/Data",
|
|
download_registry.TRANSPORT_HTTP,
|
|
repo_type = "dataset",
|
|
repo_id = "Org/Data",
|
|
)
|
|
duplicate_claimed, duplicate_state = registry.claim(
|
|
"org/data",
|
|
download_registry.TRANSPORT_HTTP,
|
|
repo_type = "dataset",
|
|
repo_id = "org/data",
|
|
)
|
|
|
|
assert claimed is True
|
|
assert state == "running"
|
|
assert duplicate_claimed is False
|
|
assert duplicate_state == "running"
|
|
assert registry.active_jobs("ORG/DATA") == {"org/data": "running"}
|
|
|
|
|
|
def test_dataset_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
|
|
monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry())
|
|
assert download_manifest.write_cancel_marker("dataset", "Owner/Data", None, "http")
|
|
|
|
status = asyncio.run(downloads.get_dataset_download_status_response("owner/data"))
|
|
|
|
assert status.state == "cancelled"
|
|
assert status.error is None
|
|
|
|
|
|
def test_dataset_claim_register_cancel_uses_registry_marker_owner(monkeypatch):
|
|
killed = []
|
|
|
|
class _Registry:
|
|
def claim(self, *_args, **_kwargs):
|
|
return True, "running"
|
|
|
|
def current_generation(self, _key):
|
|
return 1
|
|
|
|
def register_process(self, _key, _proc):
|
|
return False
|
|
|
|
def persist_cancel_for_key(self, *_args, **_kwargs):
|
|
raise AssertionError("register_process owns pending-cancel markers")
|
|
|
|
def get_job(self, _key):
|
|
return SimpleNamespace(state = "cancelled", error = None)
|
|
|
|
monkeypatch.setattr(downloads, "_registry", _Registry())
|
|
monkeypatch.setattr(
|
|
downloads,
|
|
"resolve_cached_repo_id_case",
|
|
lambda repo_id, **_kwargs: repo_id,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloads.download_registry,
|
|
"download_transport_unavailable_reason",
|
|
lambda _transport: None,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloads.download_lifecycle,
|
|
"spawn_worker",
|
|
lambda *_args, **_kwargs: object(),
|
|
)
|
|
monkeypatch.setattr(
|
|
downloads.download_lifecycle,
|
|
"kill_and_reap_process",
|
|
lambda proc, **_kwargs: killed.append(proc),
|
|
)
|
|
|
|
result = asyncio.run(
|
|
downloads.download_dataset_response(SimpleNamespace(repo_id = "Org/Data", use_xet = False))
|
|
)
|
|
|
|
assert result["state"] == "cancelled"
|
|
assert killed
|
|
|
|
|
|
def test_dataset_cancel_pending_spawn_arms_pending_cancel(monkeypatch):
|
|
events = []
|
|
|
|
class _Registry:
|
|
def get_process(self, _key):
|
|
return None
|
|
|
|
def mark_pending_cancel(self, key, generation):
|
|
events.append(("pending", key, generation))
|
|
return True
|
|
|
|
def get_job(self, _key):
|
|
return SimpleNamespace(state = "running")
|
|
|
|
monkeypatch.setattr(downloads, "_registry", _Registry())
|
|
monkeypatch.setattr(
|
|
downloads,
|
|
"resolve_cached_repo_id_case",
|
|
lambda repo_id, **_kwargs: repo_id,
|
|
)
|
|
|
|
result = asyncio.run(
|
|
downloads.cancel_dataset_download_response(
|
|
SimpleNamespace(repo_id = "Org/Data", generation = 5)
|
|
)
|
|
)
|
|
|
|
assert result == {"repo_id": "Org/Data", "state": "cancelling"}
|
|
assert events == [("pending", "org/data", 5)]
|
|
|
|
|
|
def test_upload_dataset_response_writes_non_empty_file(monkeypatch, tmp_path):
|
|
payload = b'{"text":"hello"}\n'
|
|
monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", tmp_path)
|
|
|
|
response = asyncio.run(local.upload_dataset_response(_Upload("../train.jsonl", payload)))
|
|
|
|
stored_path = Path(response.stored_path)
|
|
assert response.filename == "train.jsonl"
|
|
assert stored_path.parent == tmp_path
|
|
assert stored_path.name.endswith("_train.jsonl")
|
|
assert stored_path.read_bytes() == payload
|
|
|
|
|
|
def test_local_dataset_items_expose_recipe_and_upload_source(monkeypatch, tmp_path):
|
|
recipe_root = tmp_path / "recipes"
|
|
upload_root = tmp_path / "uploads"
|
|
parquet_dir = recipe_root / "recipe_alpha" / "parquet-files"
|
|
parquet_dir.mkdir(parents = True)
|
|
(parquet_dir / "part.parquet").write_bytes(b"parquet")
|
|
upload_root.mkdir()
|
|
(upload_root / "manual.jsonl").write_text('{"text":"hello"}\n', encoding = "utf-8")
|
|
monkeypatch.setattr(local, "LOCAL_DATASETS_ROOT", recipe_root)
|
|
monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", upload_root)
|
|
|
|
response = local.list_local_datasets_response()
|
|
|
|
assert "source" in LocalDatasetItem.__annotations__
|
|
by_id = {item.id: item for item in response.datasets}
|
|
assert by_id["recipe_alpha"].source == "recipe"
|
|
assert by_id["manual.jsonl"].source == "upload"
|