vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-08-25 08:42:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29
unsloth/studio/backend/tests/test_partial_remaining_bytes.py
Michael Han 77bc10c663
Studio: price a partial GGUF by what is left to fetch (#8989) 
* Studio: price a partial GGUF by what is left to fetch

The size beside a partial variant was the variant total, so continuing a
sharded download that was already 40 GB in still read "56 GB" and looked like
the whole model coming down again. That is the reading behind the reports of
the hub re-downloading models people already have.

The variants endpoint now reports download_remaining_bytes for a partial: the
plan total minus the bytes on disk a transfer can actually reuse. Reuse is
per file, so a finished shard counts and an unresumable partial does not, and
the card shows "16 GB left" instead of "56 GB".

A one-file quant therefore still reads back its full size, because that is
what continuing it costs: huggingface_hub 1.18+ refetches an interrupted file
from zero, and a single-file quant has no other file to keep.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: price local and offline partial rows too

The local and offline listings return before the hub-plan path, so their partial
rows carried no remaining figure and fell back to the full total. The on-device
card asks for exactly those (preferLocalCache), which is where a partial is most
likely to be looked at.

The worker writes its manifest before fetching anything, so those rows can be
priced from the file list that produced them. Capped at the row's own total,
since a manifest counts companions the row's size may not.

The on-device card also merges local rows over remote ones, so it now carries
the remaining figure through instead of dropping it.

* Studio: stop capping a local remainder by the shards on disk

A local listing sizes a variant by summing the shards it can see, so an early
interruption makes that total smaller than the transfer: three 2 GB shards with
one cached advertises 2 GB, and capping the remainder by it reported "2 GB left"
when 4 GB had to be fetched. Under-reporting is the one direction this figure
must never be wrong in.

The cap was there to stop a companion-inclusive remainder reading larger than
the row's own size, but no surface shows the two together: a partial row's size
chip IS the remainder. So the cap is dropped and the manifest total stands.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Price a partial by the bytes it holds, and credit a shared blob once across case-variant repo dirs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Follow the transfer in the remaining label while a download runs

The live overlay carried only the expected size, so a row that says
N left kept whatever the one-time variant fetch had measured, or the
full total for a download started after it. A download 90 percent
through read as though nothing had moved. The running job already
carries its own progress, so derive the remainder from it and keep the
fetched figure for every other row.

* Take both terms of the live remainder from the job

snapshot_progress nets completed_baseline_bytes out of expected_bytes
and downloaded_bytes alike, so the job's two counters are consistent
with each other and not with the catalog totals. Subtracting the job's
transfer from the larger of the two scopes added that baseline straight
back: 1 GB reused and 1 GB fetched of a 5 GB plan read 4 GB left rather
than 3 GB. The catalog total still drives the size the row reports.

* Price a pinned row against the cache root it names

A row can name a snapshot in a previous, legacy or default HF cache
root, and the request already scopes partial detection and the manifest
read to it, but the blob scan always read the active root. Wrong in both
directions: shards already in the pinned root earned no credit, and a
copy of the same blob in the active root earned credit a resume into the
pinned root cannot use, which reported less left than there is. The
docstring called counting the active root only the safe direction to be
wrong in, and it is for the first half, but not for the second.

The fixture stub also had to start honouring an explicit root; ignoring
it would have answered the active root for a pinned row and hidden the
behaviour the new tests are about.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Keep a cancelled download on the measured remainder

A job's progress is not bytes a resume can reuse: from huggingface_hub 1.18 the
partial is process-unique, opened wb and unlinked in a finally, which is how
existing_blob_bytes already prices it. Subtracting a dead job's transfer put
1.0 GB left on a cancelled 18 GB download with all 18 GB still to fetch, so the
overlay now only prices running jobs.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Price a fallback retry off the live transfer counter

An XET run that falls back to HTTP re-claims in the same generation with a completed_baseline_bytes recomputed from disk, so the baseline now covers every blob the XET attempt finalized and the new run reports completed_bytes 0 against a shrunken total. resolveProgressUpdate holds the previous reading through that zero, so the card kept the dead run's finalized bytes while its total moved to the retry's scope, and taking the max of the two counters subtracted 3 GB from a 0.5 GB remainder: the row read 0 B left with the transfer barely started. snapshot_progress builds downloaded_bytes as completed plus in-flight and nets the same baseline out of both, so a single reading never has completed above downloaded and the max could only ever fire on a held figure.

* Hold a retry's remainder off a held transfer reading

resolveProgressUpdate holds the previous downloadedBytes through a poll
that reported zero, and the GGUF row forwarded that held figure as the
current transfer. An XET run that falls back to HTTP re-claims in the
same generation with completed_baseline_bytes recomputed from the
finalized blobs on disk, and snapshot_progress nets that baseline out of
downloaded_bytes, completed_bytes and the total alike, so the retry's
first reading is a legitimate 0 against a shrunken total. Subtracting
the held 3 GB from the 0.5 GB that remained read "0 B left" until the
retry moved its first byte.

Carry the distinction out of resolveProgressUpdate as measuredTransfer
and skip the progress-derived remainder on a held reading, falling back
to the backend's own download_remaining_bytes the way the terminal row
already does.

* Carry the held-transfer marker across a reload

measuredTransfer marks a reading the poll did not measure, so the remainder
is not derived from a stale downloadedBytes. It was not persisted, so a
reload restored the stale bytes beside the new run's shrunken expectedBytes
with the flag reading undefined, which the selector treats as measured. That
is the '0 B left' the guard exists to stop, and on app start it lasts until
the first poll returns rather than a single tick.

Persist and restore the flag. Absent stays absent, so a record written before
this field still means never polled rather than held.

* Carry the held-transfer marker through adoption too

A restored fallback job is adopted by the next running-status probe, which
seeds the persisted byte counters but rebuilt the job without the marker.
That restored it as undefined, which reads as measured, so the remainder was
derived from the dead run's bytes again and the row read 0 B left.

The marker is a statement about those counters, so it travels with them:
seededMeasuredTransfer joins carriesOverSeed in adopt-rules, gated on the
same condition. Not carrying the seed zeroes the counters, so there is no
held figure left and undefined stays the honest answer.

* Do not credit a partial that has not been written yet

blob_bytes_present read a missing st_blocks and a present zero as the same
thing, so a partial set to its final length before its first chunk landed
fell through to st_size and was credited whole -- the row read 0 B left on a
download that had transferred nothing. Confirm the emptiness with SEEK_DATA
rather than inferring it, so a mount that never populates st_blocks keeps the
size fallback.

* Read a pre-marker record's counters as held, not as measured

A record written before measuredTransfer existed cannot say whether its byte
counters were measured, so an absent marker there is not the never-polled it
means in a current record. An upgrade landing mid-reclaim restored the dead
run's held bytes with the guard reading measured, which is the 0 B left the
marker exists to stop. Version the persisted state so the migration can tell
the two absences apart, and read only the pre-marker one conservatively.

* Trim the review commentary

---------

Co-authored-by: shimmyshimmer <michaelhan@Michaels-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: danielhanchen <danielhanchen@users.noreply.github.com>
2026-08-19 07:23:25 -07:00

268 lines
10 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""A partial row is priced by what a resume still has to fetch.
The card used to print the variant total beside a resume button, so continuing a
sharded download that was 40 GB in still read "56 GB". Bytes reused are whole
files: a finished shard is kept, an unresumable partial is refetched, so a
one-file quant really does read back its full size.
"""
from __future__ import annotations
import sys
from pathlib import Path
from typing import Optional
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.services.models.gguf_variants import (
variant_remaining_bytes,
variant_remaining_bytes_from_state,
)
from hub.utils import download_manifest, download_registry, hf_cache_state, state_dir
from hub.utils.download_manifest import ExpectedFile
from hub.utils.gguf_plan import plan_from_expected_files
SHARD_A = "a" * 64
SHARD_B = "b" * 64
GB = 1024**3
@pytest.fixture
def blobs(monkeypatch, tmp_path):
blobs_root = tmp_path / "hub"
blobs_dir = blobs_root / "models--Org--Model" / "blobs"
blobs_dir.mkdir(parents = True)
monkeypatch.setenv("HF_HUB_CACHE", str(blobs_root))
# Honours an explicit root, which is what a row pinned to another cache passes. A stub that
# ignored the argument would hide exactly the bug the pinned tests are about.
def _root(*, root: Optional[Path] = None, **_kw):
return Path(root) if root is not None else blobs_root
monkeypatch.setattr(download_registry, "hf_cache_root", _root)
monkeypatch.setattr(hf_cache_state, "hf_cache_root", _root)
monkeypatch.setattr(hf_cache_state, "hf_cache_roots", lambda *_a, **_k: [blobs_root])
return blobs_dir
def _write(path: Path, size: int) -> Path:
path.write_bytes(b"")
with path.open("wb") as handle:
handle.truncate(size)
return path
def _split_plan():
"""Two shards of one quant, 2 GB each."""
return plan_from_expected_files(
"Q4_K_M",
[
ExpectedFile(path = "model-Q4_K_M-00001-of-00002.gguf", size = 2 * GB, sha256 = SHARD_A),
ExpectedFile(path = "model-Q4_K_M-00002-of-00002.gguf", size = 2 * GB, sha256 = SHARD_B),
],
)
def test_a_finished_shard_is_subtracted(blobs):
_write(blobs / SHARD_A, 2 * GB)
assert variant_remaining_bytes("Org/Model", _split_plan()) == 2 * GB
def test_nothing_on_disk_still_prices_the_whole_variant(blobs):
assert variant_remaining_bytes("Org/Model", _split_plan()) == 4 * GB
def test_an_unresumable_partial_is_priced_as_a_full_refetch(blobs):
# 1.18+ writes <etag>.<nonce>.incomplete and never reopens it, so those bytes are gone.
_write(blobs / f"{SHARD_A}.deadbeef{hf_cache_state.INCOMPLETE_SUFFIX}", 2 * GB)
assert variant_remaining_bytes("Org/Model", _split_plan()) == 4 * GB
def test_a_one_file_quant_reads_back_its_full_size(blobs):
"""Nothing to keep, so a resume costs the whole quant -- the case users report as the
model downloading all over again, and the number has to say so."""
plan = plan_from_expected_files(
"Q4_K_M",
[ExpectedFile(path = "model-Q4_K_M.gguf", size = 4 * GB, sha256 = SHARD_A)],
)
_write(blobs / f"{SHARD_A}.deadbeef{hf_cache_state.INCOMPLETE_SUFFIX}", 3 * GB)
assert variant_remaining_bytes("Org/Model", plan) == 4 * GB
def test_an_unresolvable_plan_reports_nothing_rather_than_guessing(blobs):
assert variant_remaining_bytes("Org/Model", None) is None
empty = plan_from_expected_files(
"Q4_K_M",
[ExpectedFile(path = "model-Q4_K_M.gguf", size = 4 * GB, sha256 = None)],
)
assert variant_remaining_bytes("Org/Model", empty) is None
def test_a_complete_variant_has_nothing_left_to_fetch(blobs):
_write(blobs / SHARD_A, 2 * GB)
_write(blobs / SHARD_B, 2 * GB)
assert variant_remaining_bytes("Org/Model", _split_plan()) == 0
# --------------------------------------------------------------------------------------------
# Offline and local-cache listings, which have no hub plan to price from. The on-device card
# asks for those (preferLocalCache), so leaving them unpriced showed the full total there.
# --------------------------------------------------------------------------------------------
@pytest.fixture
def state(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
return tmp_path
def _write_manifest(files):
assert download_manifest.write_manifest("model", "Org/Model", "Q4_K_M", files, "http")
def test_the_worker_manifest_prices_a_local_partial(blobs, state):
_write_manifest(
[
ExpectedFile(path = "model-Q4_K_M-00001-of-00002.gguf", size = 2 * GB, sha256 = SHARD_A),
ExpectedFile(path = "model-Q4_K_M-00002-of-00002.gguf", size = 2 * GB, sha256 = SHARD_B),
]
)
_write(blobs / SHARD_A, 2 * GB)
assert variant_remaining_bytes_from_state("Org/Model", "Q4_K_M", None) == 2 * GB
def test_a_row_with_no_manifest_stays_unpriced(blobs, state):
assert variant_remaining_bytes_from_state("Org/Model", "Q4_K_M", None) is None
def test_a_companion_the_row_does_not_count_is_still_priced(blobs, state):
# The mmproj comes down with the quant, so it belongs in the transfer even though the
# row's own size does not include it.
_write_manifest(
[
ExpectedFile(path = "model-Q4_K_M.gguf", size = 4 * GB, sha256 = SHARD_A),
ExpectedFile(path = "mmproj-F16.gguf", size = 1 * GB, sha256 = SHARD_B),
]
)
assert variant_remaining_bytes_from_state("Org/Model", "Q4_K_M", None) == 5 * GB
def test_an_unnamed_variant_is_not_priced(blobs, state):
assert variant_remaining_bytes_from_state("Org/Model", "", None) is None
def test_a_local_row_is_not_capped_by_the_shards_it_already_has(blobs, state):
"""A local scan sizes a variant from the shards ON DISK, so an early interruption makes
that total smaller than the transfer still to come."""
shard_c = "c" * 64
_write_manifest(
[
ExpectedFile(path = "m-Q4_K_M-00001-of-00003.gguf", size = 2 * GB, sha256 = SHARD_A),
ExpectedFile(path = "m-Q4_K_M-00002-of-00003.gguf", size = 2 * GB, sha256 = SHARD_B),
ExpectedFile(path = "m-Q4_K_M-00003-of-00003.gguf", size = 2 * GB, sha256 = shard_c),
]
)
_write(blobs / SHARD_A, 2 * GB)
# 2 GB on disk, so the local row advertises 2 GB, but 4 GB is still to fetch.
assert variant_remaining_bytes_from_state("Org/Model", "Q4_K_M", None) == 4 * GB
# --------------------------------------------------------------------------------------------
# A partial is measured by the bytes really on disk, and one shard is credited once however
# many repo directories the cache holds for the same repo. Both were observed against real
# caches.
# --------------------------------------------------------------------------------------------
MB = 1024**2
def _sparse(path: Path, written: int, logical: int) -> Path:
"""A real sparse file: *written* bytes allocated, *logical* bytes reported."""
with path.open("wb") as handle:
handle.write(b"\xa5" * written)
handle.truncate(logical)
return path
def test_a_sparse_partial_is_priced_by_the_bytes_it_actually_holds(blobs):
"""hf_transfer's parallel Range writer leaves a partial whose st_size runs ahead of what has
been written, so crediting the logical size read "0 B left" for a file barely started."""
from filelock import FileLock
plan = plan_from_expected_files(
"Q4_K_M",
[ExpectedFile(path = "model-Q4_K_M.gguf", size = 64 * MB, sha256 = SHARD_A)],
)
partial = _sparse(
blobs / f"{SHARD_A}.deadbeef{hf_cache_state.INCOMPLETE_SUFFIX}", 4 * MB, 64 * MB
)
assert partial.stat().st_size == 64 * MB
assert partial.stat().st_blocks * 512 < 8 * MB
# Held lock: the one state in which a partial no later attempt could reopen still counts,
# because a live writer is finishing it. That is exactly when it is sparsest.
lock_path = blobs.parent.parent / ".locks" / blobs.parent.name / f"{SHARD_A}.lock"
lock_path.parent.mkdir(parents = True, exist_ok = True)
with FileLock(str(lock_path), timeout = 5):
remaining = variant_remaining_bytes("Org/Model", plan)
assert remaining is not None
assert remaining >= 64 * MB - 8 * MB, "credited the sparse file's logical size, not its bytes"
def test_one_shard_in_two_case_variant_repo_dirs_is_credited_once(blobs, monkeypatch):
"""The Hub resolves repo ids case-insensitively while huggingface_hub keeps the caller's
casing in the folder name, so a case-sensitive filesystem holds models--Org--Model beside
models--org--model. Summing the directories counted one shard twice."""
root = blobs.parent.parent
twin = root / "models--org--model" / "blobs"
twin.mkdir(parents = True)
_write(blobs / SHARD_A, 2 * GB)
_write(twin / SHARD_A, 2 * GB)
assert variant_remaining_bytes("Org/Model", _split_plan()) == 2 * GB
# --------------------------------------------------------------------------------------------
# A row pinned to another cache root. A resume writes into the root the row names, so blobs in
# the active root are not bytes it can reuse. Both directions used to be wrong.
# --------------------------------------------------------------------------------------------
def test_a_pinned_root_gets_credit_for_the_shards_it_holds(blobs, tmp_path):
other_root = tmp_path / "previous-hub"
other_blobs = other_root / "models--Org--Model" / "blobs"
other_blobs.mkdir(parents = True)
_write(other_blobs / SHARD_A, 2 * GB)
pinned = other_root / "models--Org--Model"
assert variant_remaining_bytes("Org/Model", _split_plan(), pinned) == 2 * GB
def test_the_active_roots_copy_does_not_pay_for_a_pinned_row(blobs, tmp_path):
# The whole variant sits in the active root and none of it in the pinned one, so counting
# it reported nothing left to fetch for a transfer that has everything still to do.
_write(blobs / SHARD_A, 2 * GB)
_write(blobs / SHARD_B, 2 * GB)
other_root = tmp_path / "previous-hub"
(other_root / "models--Org--Model" / "blobs").mkdir(parents = True)
pinned = other_root / "models--Org--Model"
assert variant_remaining_bytes("Org/Model", _split_plan(), pinned) == 4 * GB
def test_an_unpinned_row_still_reads_the_active_root(blobs):
_write(blobs / SHARD_A, 2 * GB)
assert variant_remaining_bytes("Org/Model", _split_plan(), None) == 2 * GB
Reference in a new issue View git blame Copy permalink