ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add a file-scoped flavour to the Hub download job

Browse source 
Lets a consumer that reads a deliberate subset of a repo stage it through the
normal download manager. Keyed as "@scope" so it never collides with a quant or
with the repo's full snapshot, and the file list rides the registry so an
XET to HTTP retry respawns the same scoped job.
This commit is contained in:
Unsloth 2026-07-26 04:07:05 -07:00
commit 032561ae21
6 changed files with 337 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

10
studio/backend/hub/schemas/downloads.py
View file

@ -28,6 +28,16 @@ class DownloadModelRequest(BaseModel):
True,
description = "Use Xet parallel chunked transport. Default True; set False for HTTP Range-resume.",
)
scope_id: Optional[str] = Field(
None,
description = "Marks a partial-by-design download of `files` only (e.g. 'diffusion', "
"whose loader reads a scoped subset of a repo). Keyed separately from the full "
"snapshot of the same repo, so neither one's manifest describes the other.",
)
files: List[str] = Field(
default_factory = list,
description = "Exact files to fetch. Required with scope_id, ignored without it.",
)
class CancelDownloadRequest(BaseModel):

22
studio/backend/hub/services/download_lifecycle.py
View file

@ -11,7 +11,7 @@ import sys
import time
import threading
from pathlib import Path
from typing import Callable, Mapping, Optional
from typing import Callable, Mapping, Optional, Sequence
from fastapi import HTTPException
@ -51,6 +51,22 @@ def resolve_transport(use_xet: bool) -> str:
return transport
def write_files_manifest(files: Sequence[str]) -> str:
"""Stage a scoped job's file list in a temp JSON file and return its path.
The worker deletes it after reading. A pipeline repo lists hundreds of files, well past
what is comfortable on a command line."""
import json
import tempfile
handle = tempfile.NamedTemporaryFile(
mode = "w", suffix = ".json", prefix = "unsloth-dl-files-", delete = False, encoding = "utf-8"
)
with handle:
json.dump(list(files), handle)
return handle.name
def spawn_worker(
args: list[str],
hf_token: Optional[str],
@ -447,6 +463,10 @@ def _try_http_retry(
args.append("--dataset")
elif variant:
args.extend(["--variant", variant])
# A scoped job must retry as the SAME scoped download; without its file list the
# HTTP worker would fall through to a full snapshot of the repo.
if original_metadata.scoped_files:
args.extend(["--files-json", write_files_manifest(original_metadata.scoped_files)])
# Re-query at spawn time: sibling state may have changed since XET failed.
peer_hashes = registry.peer_blob_hashes(key) if variant else frozenset()

85
studio/backend/hub/services/models/downloads.py
View file

@ -6,7 +6,7 @@
from __future__ import annotations
import asyncio
from typing import Optional, TYPE_CHECKING
from typing import Optional, Sequence, TYPE_CHECKING
from fastapi import HTTPException
from loggers import get_logger
@ -21,6 +21,7 @@ from hub.utils import download_registry
from hub.utils import download_manifest
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import has_active_incomplete_blobs
from hub.utils.snapshot_filters import blob_hashes_for_siblings
from hub.utils.paths import (
is_valid_gguf_variant as _is_valid_gguf_variant,
is_valid_repo_id as _is_valid_repo_id,
@ -44,6 +45,31 @@ def _download_job_key(repo_id: str, variant: Optional[str]) -> str:
)
# A scope rides the variant slot as "@name". No GGUF quant label starts with "@", so a
# scoped job never collides with a real variant or with the repo's full snapshot, and its
# manifest/cancel marker stay in their own space.
_SCOPE_PREFIX = "@"
def _scope_variant(scope_id: Optional[str]) -> Optional[str]:
scope = (scope_id or "").strip()
return f"{_SCOPE_PREFIX}{scope}" if scope else None
def scoped_file_blob_hashes(
repo_id: str, files: Sequence[str], hf_token: Optional[str]
) -> frozenset[str]:
"""Blob hashes for exactly ``files``, so a scoped job's progress, purge and peer
protection cover its own files and nothing else in the repo."""
from huggingface_hub import HfApi
wanted = set(files)
info = HfApi().model_info(repo_id, files_metadata = True, token = hf_token)
return blob_hashes_for_siblings(
[s for s in info.siblings if getattr(s, "rfilename", None) in wanted]
)
def _job_status(
key: str,
*,
@ -91,10 +117,14 @@ def _spawn_download_worker(
use_xet: bool = True,
protected_blob_hashes: Optional[frozenset[str]] = None,
cache_env: Optional[dict[str, str]] = None,
files: Optional[Sequence[str]] = None,
) -> subprocess.Popen:
args = ["--repo-id", repo_id]
if variant:
args.extend(["--variant", variant])
if files:
# Via a temp file, not argv: a pipeline repo's list runs to hundreds of names.
args.extend(["--files-json", download_lifecycle.write_files_manifest(files)])
return download_lifecycle.spawn_worker(
args,
hf_token,
@ -124,6 +154,22 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
status_code = 400,
detail = f"Invalid gguf_variant: {variant!r}",
)
# A scoped job fetches only `files` and keys itself apart from the repo's full snapshot.
scope_variant = _scope_variant(body.scope_id)
scoped_files = [f for f in (body.files or []) if f and f.strip()]
if scope_variant is not None:
if variant is not None:
raise HTTPException(
status_code = 400,
detail = "scope_id and gguf_variant are mutually exclusive.",
)
if not scoped_files:
raise HTTPException(
status_code = 400, detail = "scope_id requires a non-empty files list."
)
if not _is_valid_gguf_variant(scope_variant):
raise HTTPException(status_code = 400, detail = f"Invalid scope_id: {body.scope_id!r}")
variant = scope_variant
key = _download_job_key(repo_id, variant)
use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
transport = download_lifecycle.resolve_transport(use_xet)
@ -136,20 +182,27 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
completed_baseline_bytes = 0
if variant is not None:
try:
variant_blob_hashes = await asyncio.to_thread(
gguf_variants.gguf_variant_blob_hashes,
repo_id,
variant,
hf_token,
include_companions = False,
)
variant_progress_blob_hashes = await asyncio.to_thread(
gguf_variants.gguf_variant_blob_hashes,
repo_id,
variant,
hf_token,
include_companions = True,
)
if scope_variant is not None:
# A scope owns exactly its own files: same set for purge and for progress.
variant_blob_hashes = await asyncio.to_thread(
scoped_file_blob_hashes, repo_id, scoped_files, hf_token
)
variant_progress_blob_hashes = variant_blob_hashes
else:
variant_blob_hashes = await asyncio.to_thread(
gguf_variants.gguf_variant_blob_hashes,
repo_id,
variant,
hf_token,
include_companions = False,
)
variant_progress_blob_hashes = await asyncio.to_thread(
gguf_variants.gguf_variant_blob_hashes,
repo_id,
variant,
hf_token,
include_companions = True,
)
except Exception as e:
logger.warning(
"GGUF hash pre-resolution failed for %s [%s]; continuing without "
@ -183,6 +236,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
admission_check = lambda: not _load_in_flight(repo_id),
hub_cache = str(cache_paths.hub_cache),
xet_cache = str(cache_paths.xet_cache),
scoped_files = scoped_files if scope_variant is not None else None,
)
generation = _registry.current_generation(key)
if not claimed:
@ -218,6 +272,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
use_xet = use_xet,
protected_blob_hashes = protected_blob_hashes,
cache_env = cache_env,
files = scoped_files if scope_variant is not None else None,
),
hf_token = hf_token,
label = label,

7
studio/backend/hub/utils/download_registry.py
View file

@ -47,7 +47,7 @@ import time
import weakref
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Callable, Iterator, Literal, Optional
from typing import Callable, Iterator, Literal, Optional, Sequence
from loggers import get_logger
@ -773,6 +773,9 @@ class DownloadMetadata:
completed_baseline_bytes: int = 0
hub_cache: Optional[str] = None
xet_cache: Optional[str] = None
# Scoped jobs only: the exact files to fetch. Kept here so the XET -> HTTP retry
# respawns the same scoped download instead of a full snapshot.
scoped_files: tuple[str, ...] = ()
@dataclass(frozen = True)
@ -1099,6 +1102,7 @@ class DownloadRegistry:
cancel_marker_transport: Optional[str] = None,
hub_cache: Optional[str] = None,
xet_cache: Optional[str] = None,
scoped_files: Optional[Sequence[str]] = None,
) -> tuple[bool, str]:
key = normalize_job_key(key)
repo = _repo_of_key(key)
@ -1174,6 +1178,7 @@ class DownloadRegistry:
),
hub_cache = hub_cache,
xet_cache = xet_cache,
scoped_files = tuple(scoped_files or ()),
)
if cancel_marker_transport is not None:
self._cancel_marker_transports[key] = cancel_marker_transport

94
studio/backend/hub/workers/hf_download.py
View file

@ -671,6 +671,78 @@ def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mod
)
def _download_scoped_snapshot(
repo_id: str, scope: str, files: list[str], hf_token: str | None, mode: str
) -> None:
"""Fetch exactly ``files`` from ``repo_id``, keyed under ``scope``.
For consumers that read a deliberate subset of a repo (the diffusion loader skips the
packaged root single, transformer/ shards and fp16 twins). Keyed apart from the repo's
full snapshot so neither manifest describes the other, and the repo is not later judged
partial against expectations it was never meant to meet."""
from huggingface_hub import HfApi, snapshot_download
from hub.utils.download_registry import prepare_cache_for_transport
from hub.utils import download_manifest
from hub.utils.download_manifest import ExpectedFile
wanted = set(files)
try:
info = _model_info_with_retry(repo_id, hf_token)
except Exception as e:
print(
f"metadata unavailable for scoped download of {repo_id} "
f"({type(e).__name__}: {e})",
file = sys.stderr,
)
info = None
expected_files: list[ExpectedFile] = []
blob_hashes: frozenset[str] = frozenset()
if info is not None:
siblings = [s for s in info.siblings if getattr(s, "rfilename", None) in wanted]
expected_files = [
ExpectedFile(
path = s.rfilename,
size = int(getattr(s, "size", 0) or 0),
sha256 = sibling_sha256(s),
)
for s in siblings
]
from hub.utils.snapshot_filters import blob_hashes_for_siblings
blob_hashes = blob_hashes_for_siblings(siblings)
download_manifest.write_manifest("model", repo_id, scope, expected_files, mode)
download_manifest.clear_cancel_marker("model", repo_id, scope)
purged = prepare_cache_for_transport(
"model",
repo_id,
mode,
scope,
only_blob_hashes = blob_hashes or None,
protected_blob_hashes = _protected_blob_hashes(),
)
if purged:
print(
f"Purged {purged} untrusted partial blob(s) for {repo_id} [{scope}] "
f"before starting {mode} download.",
file = sys.stderr,
)
_preflight_disk_space("model", repo_id, expected_files)
snapshot_path = snapshot_download(
repo_id = repo_id,
token = _hf_token_arg(hf_token),
allow_patterns = files,
max_workers = 1,
)
_verify_completed_download(
"model",
repo_id,
scope,
snapshot_path,
metadata_unavailable = info is None,
)
def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None:
from huggingface_hub import snapshot_download
from hub.utils.download_registry import prepare_cache_for_transport
@ -739,8 +811,26 @@ def main() -> None:
parser.add_argument("--dataset", action = "store_true")
parser.add_argument("--transport", choices = ("http", "xet"), default = "http")
parser.add_argument("--parent-pid", type = int, default = None)
parser.add_argument(
"--files-json",
default = None,
help = "Temp JSON file holding a scoped job's exact file list (deleted after reading).",
)
args = parser.parse_args()
scoped_files: list[str] = []
if args.files_json:
import json
try:
with open(args.files_json, encoding = "utf-8") as handle:
scoped_files = [str(f) for f in json.load(handle)]
finally:
try:
os.unlink(args.files_json)
except OSError:
pass
_install_signal_handlers()
_install_parent_death_watchdog(args.parent_pid)
@ -749,6 +839,10 @@ def main() -> None:
try:
if args.dataset:
_download_dataset(args.repo_id, hf_token, args.transport)
elif scoped_files:
_download_scoped_snapshot(
args.repo_id, args.variant, scoped_files, hf_token, args.transport
)
elif args.variant:
_download_gguf_variant(args.repo_id, args.variant, hf_token, args.transport)
else:

136
studio/backend/tests/test_scoped_download_job.py Normal file
View file

@ -0,0 +1,136 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The scoped download flavour: fetch an explicit file list through the normal download
manager, so the Images/Video pages stage models the same way Chat and the Hub do.
A diffusion load reads a deliberate subset of a repo (no packaged root single, no
transformer/ shards, no fp16 twins), so a plain snapshot would pull tens of GB it never
opens. These cover the scoping, the separate job key, and the XET -> HTTP retry.
"""
from __future__ import annotations
import asyncio
import json
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)
from hub.schemas.downloads import DownloadModelRequest
from hub.services import download_lifecycle
from hub.services.models import downloads as dl
from hub.utils.paths import is_valid_gguf_variant
FILES = ["model_index.json", "vae/diffusion_pytorch_model.safetensors"]
def _request(**over) -> DownloadModelRequest:
body = {
"repo_id": "black-forest-labs/FLUX.1-dev",
"scope_id": "diffusion",
"files": list(FILES),
"use_xet": False,
}
body.update(over)
return DownloadModelRequest(**body)
def test_scope_keys_apart_from_the_full_snapshot():
# Same repo, two jobs: the scoped one must not adopt or overwrite the full snapshot's
# manifest, or the repo would read as partial against expectations it never had.
full = dl._download_job_key("black-forest-labs/FLUX.1-dev", None)
scoped = dl._download_job_key("black-forest-labs/FLUX.1-dev", dl._scope_variant("diffusion"))
assert full != scoped
assert scoped.endswith("@diffusion")
# It rides the variant slot, so it must satisfy the same validator.
assert is_valid_gguf_variant("@diffusion")
# The "@" prefix is what keeps a scope out of the quant namespace: a job scoped
# "diffusion" and a (hypothetical) quant named "diffusion" stay distinct.
assert dl._download_job_key("org/m", "diffusion") != dl._download_job_key(
"org/m", dl._scope_variant("diffusion")
)
def test_scope_requires_files_and_rejects_a_variant(monkeypatch):
monkeypatch.setattr(dl, "_reject_if_load_in_flight", lambda repo_id: None)
monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo)
with pytest.raises(Exception) as no_files:
asyncio.run(dl.download_model_response(_request(files = [])))
assert "files" in str(no_files.value)
with pytest.raises(Exception) as both:
asyncio.run(dl.download_model_response(_request(gguf_variant = "Q4_K_M")))
assert "mutually exclusive" in str(both.value)
def test_scoped_start_spawns_a_file_scoped_worker(monkeypatch):
spawned: dict = {}
monkeypatch.setattr(dl, "_reject_if_load_in_flight", lambda repo_id: None)
monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo)
monkeypatch.setattr(dl, "scoped_file_blob_hashes", lambda *a, **k: frozenset({"h1"}))
def _fake_launch(registry, key, *, spawn, **kwargs):
spawn()
return "running"
def _fake_spawn(args, hf_token, **kwargs):
spawned["args"] = args
return object()
monkeypatch.setattr(download_lifecycle, "launch_worker", _fake_launch)
monkeypatch.setattr(download_lifecycle, "spawn_worker", _fake_spawn)
result = asyncio.run(dl.download_model_response(_request()))
assert result["accepted"] is True
assert result["job_key"].endswith("@diffusion")
args = spawned["args"]
assert "--variant" in args and args[args.index("--variant") + 1] == "@diffusion"
# The file list travels in a temp JSON file, not argv: a pipeline repo lists hundreds.
manifest_path = args[args.index("--files-json") + 1]
assert json.loads(Path(manifest_path).read_text(encoding = "utf-8")) == FILES
Path(manifest_path).unlink(missing_ok = True)
def test_scoped_files_survive_into_the_registry(monkeypatch):
# The XET -> HTTP retry rebuilds worker args from registry metadata alone. Without the
# file list there, a retried scoped job would silently become a full snapshot.
captured: dict = {}
real_claim = dl._registry.claim
def _spy_claim(key, transport, **kwargs):
captured.update(kwargs)
return real_claim(key, transport, **kwargs)
monkeypatch.setattr(dl, "_reject_if_load_in_flight", lambda repo_id: None)
monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo)
monkeypatch.setattr(dl, "scoped_file_blob_hashes", lambda *a, **k: frozenset())
monkeypatch.setattr(dl._registry, "claim", _spy_claim)
monkeypatch.setattr(
download_lifecycle, "launch_worker", lambda *a, **k: "running"
)
asyncio.run(dl.download_model_response(_request()))
assert captured["scoped_files"] == FILES
metadata = dl._registry.get_job_metadata(
dl._download_job_key("black-forest-labs/FLUX.1-dev", "@diffusion")
)
assert metadata is not None and list(metadata.scoped_files) == FILES
def test_files_manifest_round_trips():
path = download_lifecycle.write_files_manifest(FILES)
try:
assert json.loads(Path(path).read_text(encoding = "utf-8")) == FILES
finally:
Path(path).unlink(missing_ok = True)