a00fe86c13* Studio: read model text as utf-8 so umlauts survive on Windows Chat rejects or mangles non-ASCII on Windows: "ä ö ü" in a prompt, a chat template, or a model path comes back as mojibake, or the load dies with UnicodeDecodeError. open() and Path.read_text() fall back to locale.getencoding() when no encoding is passed. On Windows that is the ANSI codepage (cp1252, cp932, cp1251, ... by system locale), never UTF-8. Hugging Face writes these files as raw UTF-8, so every read of one decodes with the wrong codec: - tokenizer_config.json, which holds the chat template. Templates routinely carry -> arrows, smart quotes and CJK, so this is the common path into chat - config.json and adapter_config.json - modules.json, Ollama manifests, and the .py sources the remote-code scanner reads before a model is allowed to load The llama-server and embedding-server stdout readers have the same problem via subprocess(text = True); they now decode utf-8 with errors = "replace" so a stray byte cannot kill a log reader. Encoding arguments only, no logic changes. tests/test_chat_text_encoding.py covers a config.json and a chat template holding umlauts, arrows and CJK, plus the remote-code scanner reading a source file with umlauts. Those pass anywhere the locale is already UTF-8, so a fourth test re-runs the readers under -X warn_default_encoding and fails on any platform if an encoding argument goes missing again. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: name utf-8 explicitly on the remaining text I/O, with an AST guard (#7465) * Studio: name utf-8 explicitly on the remaining text I/O Follow-up to the model-text reads in #7467, covering the rest of the backend: system probes (nvidia-smi, amd-smi, powershell, git, node), package installers, /proc and /sys readers, and internal marker files (pid, install id, bootstrap password, Colab credentials). Same reason as #7467. open(), Path.read_text()/write_text() and subprocess(text = True) fall back to locale.getencoding(), which on Windows is the ANSI codepage rather than UTF-8. These paths are mostly ASCII today, so this is hardening, not a live bug. Encoding arguments only, no logic changes. Adds tests/test_text_io_encoding.py: an AST guard walking every backend source and asserting text I/O names its encoding, so the class of bug cannot creep back in one call at a time. 275 files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Catch aliased subprocess and positional Path.open, migrate legacy JSONL The guard only matched a receiver literally named subprocess, so worker.py's `import subprocess as _sp` hid three text = True installs that decode pip output with the ANSI codepage. It also skipped any .open() with more than one positional argument, though Path.open takes buffering/encoding/errors/newline positionally. Resuming a scrape written by an older release is the other half: those JSONL lines are in the locale codepage, so the UTF-8 preload raised, the dedup keys were silently forgotten and duplicates were appended to a now mixed-encoding file. Decode with the locale codepage as fallback and rewrite as UTF-8 before the append handle opens, since Windows cannot replace a file it holds open. * Stream the JSONL preload and keep a torn line from relabelling the shard Reading the whole shard to migrate it was wrong twice over. These files reach gigabytes on a large scrape, so the preload now streams line by line and the rewrite streams through a temp file. Worse, one interrupted append used to condemn the file: the whole-file UTF-8 decode failed, every byte was retried as cp1252, and the rewrite persisted mojibake over records that were fine. A line now counts as legacy only if the locale codepage both decodes it and yields valid JSON, which a torn UTF-8 line does not. Damaged lines are skipped and copied through byte for byte. When the rewrite cannot be written at all, the append handle opens with the legacy encoding rather than mixing UTF-8 into the file. install_wheel takes run = subprocess.run as a parameter, so the guard cannot see it. Both wheel installs there now name their encoding. * Decide the shard's encoding from the file, not one line at a time Some byte strings parse both ways. cp1251 `Р°` is D0 B0, which is also valid UTF-8 for `а`, so a UTF-8-first parse quietly showed the wrong text instead of migrating it. A line now yields both readings, and the file decides. Any line that parses under the codepage but not as UTF-8 is unambiguous evidence, and ambiguous lines then follow that verdict, which is enough for any real shard: ordinary Cyrillic or Japanese prose is invalid UTF-8 several times per line. Keys for ambiguous lines are re-derived from the legacy reading during the rewrite. A shard is undecidable only if every line is ambiguous, and nothing can tell those apart. latin-1 is also tried after the locale codepage, so a scrape carried from Windows to a UTF-8 machine still has a reading rather than none. Requiring valid JSON, not just a decode, keeps that from claiming torn lines. * Weigh the whole shard, and never lose a record on the fallback path One structurally valid JSON line carrying a stray 0x96 parses as cp1252, so a single-line verdict let it relabel a healthy shard and mojibake every good record in it. Each line with non-ASCII bytes now votes: parsing only under the codepage is evidence for legacy, parsing as UTF-8 is evidence against, since codepage text rarely forms valid multibyte UTF-8. Ties leave the file alone. When the migration cannot be written the append handle uses the legacy codepage, and errors = "replace" quietly turned characters it cannot hold into question marks while write() still reported success. That path now escapes to \uXXXX instead, which is ASCII, so every codepage holds it and json.loads returns the exact characters. Nothing needs replacing, so errors = "strict" is safe. stream_installer runs sys.executable, so its output is now decoded as UTF-8 by utf8_child_env rather than read as the ANSI codepage. * Only rewrite a shard we can attribute, and append ASCII when we cannot latin-1 was doing too much work. It reads any byte, so it gave a moved shard a reading, but it is the right text only for cp1252: cp1251 Привет came back as Ïðèâåò and the rewrite made that permanent. The codepage is now trusted only when it is the locale's, and an untrusted reading is never written back. That leaves three cases where the file holds bytes UTF-8 cannot read and we are not converting it: no codepage to attribute it to, ambiguous lines outvoting the unambiguous ones, and a preload that could not read the file at all. All three used to append UTF-8 into it. They now append pure ASCII, which every ASCII-compatible codepage stores identically, so the file keeps decoding exactly as it did and no record is lost. Keys from the two readings are also kept apart. A damaged line in a healthy shard was marked seen through its codepage reading, so the retry that would have replaced the unreadable record was refused as a duplicate. * Let the flash-attn install stub take the kwargs the installer now passes _run_kwargs gained encoding and errors, so the one stub in this file that spelled its signature out rejected the call. The other four here already take **kwargs; this one now matches. * Do not let a stuck temp file mask the migration failure unlink() on the failure path could raise in its own right, on a stale .utf8.tmp directory or a temp another process holds. That escaped the constructor instead of returning False, so the caller never reached the ASCII append fallback that keeps the shard single-encoding. The pip fallback in install_wheel also spawns a Python child, so it gets utf8_child_env like the probe above it already had. The uv and nvidia-smi children are native binaries, where PYTHONIOENCODING would do nothing. * Stop converting legacy shards; the encoding that wrote them is unknowable trusted only ever meant that the bytes parse under this machine's codepage, which for a single-byte codepage is nearly always true. A cp1251 shard opened on a cp1252 Windows box decodes cleanly and would have been rewritten with Привет as Ïðèâåò. That is the fourth way this rewrite could corrupt a shard, and the common cause is that a file's encoding cannot be recovered from its bytes. So the rewrite is gone. The shard is left exactly as found, and appends are pure ASCII whenever it holds bytes UTF-8 cannot read, which is what actually delivered the no-mixed-encoding guarantee the rewrite was added for. Dedup keys still come from whichever reading parses, since ids are ASCII either way. This also removes the temp file, so there is no longer any file mode or ACL to carry across. * Scan the sandbox shim; it is shipped code, not a build artifact sandbox_site is on the sandboxed child's PYTHONPATH for every Python run (tools.py:332, 2660), so excluding it let two unannotated text calls through in code we ship. Both read and write the remap sidecar, which holds file paths. The exclusion list is meant for build output only, so the directory comes off it and the two calls name their encoding. * Force the worker's pip children to UTF-8, and read DBCS keys with a DBCS codec The three installer calls run sys.executable -m pip with an inherited environment, so the parent decoded UTF-8 while the child emitted the ANSI codepage. They now go through utf8_child_env like the other Python children. Two tests asserted no env kwarg was passed as a stand-in for no HIP flag being injected. They now assert the flag itself, which is the guarantee they were written for and does not depend on how the env is delivered. Separately, latin-1 cannot stand in for a double-byte codepage while recovering dedup keys: cp932 表 is 95 5C, and the trail byte reads as a JSON backslash, so the record failed to parse and its id was forgotten, appending a duplicate on resume. cp932, cp936, cp949 and cp950 are tried too. The reading is still only ever used for keys, which are ASCII and identical whichever codec parses. * Require more than one legacy line before trusting its dedup keys A shard whose valid records are all ASCII casts no UTF-8 votes, so a single damaged line won the vote by itself, its key was remembered, and the retry that would have replaced the unreadable record was refused. One such line is genuinely undecidable: a legacy record with one accented character and an ASCII record with one stray byte are the same shape. Reading it as damage costs a duplicate; reading it as legacy loses the record for good. Only one of those is recoverable, so it is now read as damage. A real legacy shard has a legacy line for every record carrying an umlaut, so its dedup is unaffected. * Append ASCII whenever the shard already holds non-ASCII bytes The gate asked whether any line was undecodable as UTF-8, which misses a shard where every legacy line happens to be valid UTF-8 too. A cp1251 shard of Р° records is bytes D0 B0 throughout, so appending 世界 as UTF-8 left a file where cp1251 reads the old records correctly and the new one as mojibake, and UTF-8 does the reverse. No single decoding recovered the whole scrape. The gate is now simply whether the shard holds any non-ASCII byte at all, which covers both cases and is easier to reason about: if what is already there reads differently under different encodings, do not add more bytes that do. Appending ASCII costs only \uXXXX escapes, which json.loads turns back into the exact characters, and it leaves the new record correct under either reading. * Skip the two Linux-gated flash-attn tests off Linux _should_try_runtime_flash_attn_install ends in sys.platform.startswith( "linux"), and the threshold test one line above already asserts exactly that, so the two tests that drive _ensure_flash_attn_for_long_context past the gate cannot pass anywhere else: the call returns before it reports a status. They were written on Linux and only surface once the suite actually runs on Windows or macOS, where both fail on an empty status list. This PR is about making the backend behave on Windows, so its own suite should be runnable there. * Fail closed when a KFD topology node does not decode This PR pins that read to utf-8, which turns an undecodable byte into UnicodeDecodeError. That is a ValueError, not an OSError, so it slips past the handler one line below and escapes a helper whose docstring promises to fail closed on any unreadable node. The caller would then lose the whole HIP-order map on a machine that has AMD GPUs, and the reason the helper fails closed is that dropping a node shifts every later ordinal and lets a similar-capacity GPU pass the total-size guard while showing another card's usage. Widening the handler is the same one-line change main already made in #7487, so the two agree and the eventual merge is clean. * Tighten the comments added in this branch * Treat an undecodable marker and undecodable metadata as malformed, not fatal Two more places where pinning the decode changed the failure mode. A UnicodeDecodeError is a ValueError, so neither `except OSError` nor `except (JSONDecodeError, OSError)` catches it, and both sites had a documented fallback that stopped being reached. An undecodable .transport marker used to read as an unknown value, and the caller then safely purged and restarted the partial download. It now aborts prepare_cache_for_transport instead, so the transfer fails rather than retrying. Undecodable .meta.json used to fall back to the file's own name, the same way invalid JSON does. It now aborts URI construction for the entire unstructured seed, so one corrupt byte in original_filename takes out the whole dataset. Both handlers are widened, matching the KFD fix earlier on this branch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Widen two more decode guards, and pin the kernel installer's pipe Same shape as the ones already fixed here: the read was pinned to UTF-8 while the handler around it still only catches OSError, and UnicodeDecodeError is a ValueError. hf_cache_snapshot_dir answers whether a model is already on disk, and the offline embedding checks turn a raise into a 500. A torn refs/main used to decode into a nonsense commit and miss the snapshot dir; it now skips that cache root and keeps looking. _remove_pid_file runs first in _graceful_shutdown, so a corrupt studio.pid raising there abandoned the inference, export, training and tunnel children the rest of that function exists to kill. ssm_runtime's source-build path builds its subprocess kwargs in a dict and splats them through _run_with_heartbeat, so neither the encoding guard nor the earlier sweep saw the text = True in it: pip's output was still decoded with the Windows ANSI codepage, where a non-ASCII path or a compiler diagnostic mojibakes or raises over an install that was going fine. It now pins the same utf-8/replace pair install_wheel uses, and the HIP branch extends that env rather than replacing it. The guard learned the dict-literal shape and reddens on the old code (ssm_runtime.py:253). * Tighten the comments around the UTF-8 text I/O pins Collapse the multi-line rationales added with the encoding pins down to a line or two each, drop what the code already says, and use one wording for the repeated child-env note. * Do not let an unreadable bootstrap password stop startup, and narrow the kwargs guard ensure_default_admin calls _load_bootstrap_password for every existing admin and the lifespan calls that with no handler, so pinning the decode turned a damaged or pre-pin .bootstrap_password file into a backend that will not start. We write that file ourselves in UTF-8, so a byte that will not decode belongs to a file whose plaintext is worthless anyway; it now reads as no bootstrap password, the same answer as an absent file. A readable one still loads. The new kwargs check also judged every dict literal in the tree, so an unrelated payload carrying "text": True would have been reported as subprocess configuration with a misleading message, and a dict that fills in its encoding on a later line would have been reported too. It now only judges a dict that actually reaches a call, either splatted through a name or written at the call site, and treats a later kw["encoding"] assignment as satisfying it. The ssm_runtime shape it was written for is still caught, and a test pins both directions. * Stop reading a UTF-8 record a second time _read_line always parsed the line under the codepage as well, even when it had already read as UTF-8. Both callers take the UTF-8 reading when there is one and never look at the other, so on a healthy shard the second parse is pure waste, and this file reads all of one on every resume of a scrape it expects to reach gigabytes. Measured on 200,000 records, 76 MB: 1.96s before, 0.81s after, so the double reading was costing 2.8x. The early return is limited to a record, since the key lookup deliberately falls through to the codepage reading when UTF-8 yields something that is not one. A line UTF-8 cannot read still tries the codepage, latin-1 and the double-byte encodings as before, which is what the second reading is for. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pin the scanned source fixture's line endings test_remote_code_scan_reads_non_ascii_sources compared a file's contents against the string it wrote, but wrote it in text mode, so Windows translated the line ends on the way out and the read back differed by a carriage return. That is the writer's doing, not the encoding the test is about, and it was the one failure on the Windows runner that belonged to this branch. The fixture now writes with newline = "" so the bytes on disk are the string on every platform. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim the newer comments to their point Shorten the widened-guard and state store notes added since the last pass, and collapse the line-ending note on the scanned source fixture. * Read the scraper checkpoint as UTF-8 only, never as a codepage A checkpoint holds nothing but base64 cursors and booleans, so one written by an older locale-encoded release is byte-identical to a UTF-8 one and already reads back. The codepage fallback can therefore only ever contribute non-ASCII: if a single-byte reading of the file were all ASCII, the UTF-8 read would have succeeded first. So the only file it changes the answer for is a damaged one, and there it turns a safe reset into a resume on a mojibaked cursor. GitHub answers that with INVALID_CURSOR_ARGUMENTS at HTTP 200, gh_client returns the partial document, and the scraper reads zero nodes and an empty pageInfo, which marks the stream done. Every later resume then skips it entirely. Reading UTF-8 only restores the earlier behaviour of dropping a checkpoint that will not decode, which re-scrapes from the first page while the writers dedup the replay. The shard scan below keeps its codepage reading; those records do carry non-ASCII. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate the remaining tilelang install tests to Linux _tilelang_platform_supported() returns False off Linux, so _ensure_tilelang_backend returns before the install and the subprocess mock these six assert on is never called. They fail on macOS runners for that reason alone. The rest of the file already carries this marker; these were missed. * Gate the Windows-incompatible worker and ROCm tests Two different gates, because the production code has two. The causal-conv1d and flash-linear-attention installers bail out on sys.platform == 'win32' alone and run everywhere else including macOS, so those cases get not_on_windows; marking them linux_only would skip tests that legitimately pass off Linux. The DRM and KFD readers return early unless platform.system() is Linux, and their fixtures build a fake sysfs tree needing PCI addresses like 0000:00:02.0 as directory names, which Windows cannot represent, so those get linux_only. The two visible-utilization cases failed for a different reason: on Windows get_visible_gpu_utilization takes the AMD adapter branch ahead of the torch fallback under test, and probing it imports torch, which the runner lacks. Stubbing that branch empty leaves every other platform unchanged. * Treat unparseable JSON nesting as a parse failure, and guard os.fdopen json.loads answers nesting it cannot descend with RecursionError, a RuntimeError, so _parse let it escape where the catch-all it replaced discarded the record. Both callers run _parse outside any further handler, so one damaged checkpoint or shard line aborted the scraper at startup. The encoding guard also missed os.fdopen, which is open() on a descriptor and takes the same locale default in text mode. It flags exactly the two text-mode calls that were left unencoded; the swap lock file's reader was already pinned to UTF-8 while its writer still used the codepage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Write the non-ASCII source fixture without a 3.10-only argument Path.write_text() only grew newline in 3.10, and pyproject declares requires-python >=3.9, so this raised TypeError there. open() takes the same argument on every supported version and pins the bytes on disk the same way. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten encoding comments * Follow subprocess calls through callable aliases in the encoding guard --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com> --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com>
1514 lines
59 KiB
Python
1514 lines
59 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
|
|
|
|
"""HF cache inspection, download registry state, and orphan-worker reaping.
|
|
|
|
Worker spawning and exit handling live in
|
|
:mod:`hub.services.download_lifecycle`; this module owns the registry state
|
|
machine plus the cache/marker inspection those workers depend on.
|
|
|
|
Resume model
|
|
------------
|
|
Only the HTTP transport supports true partial-file resume:
|
|
huggingface_hub's HTTP resumer opens ``<etag>.incomplete`` in append mode
|
|
and sends ``Range: bytes={resume_size}-`` to continue from disk.
|
|
|
|
The XET transport CANNOT resume from a ``.incomplete`` partial:
|
|
``hf_xet.download_files`` rewrites the destination from scratch.
|
|
Network-level dedup still happens, but through the separate chunk cache at
|
|
``~/.cache/huggingface/xet/chunk-cache``, which these helpers never touch.
|
|
|
|
Cross-transport corruption: a partial written by XET (or ``hf_transfer``'s
|
|
parallel-Range writer) can be sparse — high reported size, zero-filled
|
|
gaps below. Feeding it to the HTTP resumer would produce a correct-sized
|
|
blob whose internal bytes are silently wrong. To prevent that, we keep
|
|
transport markers at the download's scope (repo for snapshots/datasets,
|
|
variant for GGUF) and refuse to inherit an HTTP partial unless the marker
|
|
proves the previous writer was the same single-stream sequential writer.
|
|
|
|
Marker writes go through tmp+rename in :func:`prepare_cache_for_transport`
|
|
before the worker hands off to ``snapshot_download``, so the next process
|
|
always sees a consistent provenance signal.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import re
|
|
import shlex
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import weakref
|
|
from dataclasses import dataclass, field, replace
|
|
from pathlib import Path
|
|
from typing import Callable, Iterator, Literal, Optional
|
|
|
|
from loggers import get_logger
|
|
|
|
from hub.utils import state_dir
|
|
from hub.utils.state_dir import RepoType
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
from hub.utils.hf_cache_state import (
|
|
INCOMPLETE_SUFFIX,
|
|
TRANSPORT_HTTP,
|
|
TRANSPORT_XET,
|
|
TRANSPORT_MARKER_NAME,
|
|
VALID_TRANSPORTS,
|
|
has_active_incomplete_blobs,
|
|
iter_repo_cache_dirs,
|
|
iter_active_repo_cache_dirs,
|
|
repo_cache_dir_name,
|
|
target_dir_name,
|
|
hf_cache_root,
|
|
)
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class DownloadTransportCapability:
|
|
available: bool
|
|
reason: Optional[str] = None
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class DownloadTransportCapabilities:
|
|
http: DownloadTransportCapability
|
|
xet: DownloadTransportCapability
|
|
|
|
|
|
def get_download_transport_capabilities() -> DownloadTransportCapabilities:
|
|
xet_available = importlib.util.find_spec("hf_xet") is not None
|
|
return DownloadTransportCapabilities(
|
|
http = DownloadTransportCapability(available = True),
|
|
xet = DownloadTransportCapability(
|
|
available = xet_available,
|
|
reason = None
|
|
if xet_available
|
|
else "Xet transport is unavailable because hf_xet is not installed.",
|
|
),
|
|
)
|
|
|
|
|
|
def download_transport_unavailable_reason(transport: str) -> Optional[str]:
|
|
if transport == TRANSPORT_HTTP:
|
|
return None
|
|
if transport == TRANSPORT_XET:
|
|
caps = get_download_transport_capabilities().xet
|
|
return None if caps.available else caps.reason
|
|
return f"Unsupported download transport: {transport}"
|
|
|
|
|
|
def _worker_breadcrumb_path(key: str) -> Optional[Path]:
|
|
parent = state_dir.workers_dir()
|
|
if parent is None:
|
|
return None
|
|
safe = hashlib.sha256(key.encode("utf-8")).hexdigest()[:32]
|
|
return parent / f"{safe}.json"
|
|
|
|
|
|
def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMetadata"]) -> None:
|
|
"""Record a live worker's PID so a restarted backend can reap it. Best
|
|
effort: a write failure only forfeits boot-time reaping for this worker,
|
|
still covered by the worker's own parent-death watchdog."""
|
|
path = _worker_breadcrumb_path(key)
|
|
if path is None:
|
|
return
|
|
payload = {
|
|
"pid": int(pid),
|
|
"repo_type": metadata.repo_type if metadata is not None else None,
|
|
"repo_id": metadata.repo_id if metadata is not None else None,
|
|
"variant": metadata.variant if metadata is not None else None,
|
|
"transport": metadata.transport if metadata is not None else None,
|
|
"cancel_marker_transport": metadata.cancel_marker_transport
|
|
if metadata is not None
|
|
else None,
|
|
"hub_cache": metadata.hub_cache if metadata is not None else None,
|
|
"xet_cache": metadata.xet_cache if metadata is not None else None,
|
|
}
|
|
tmp = path.with_name(f".{path.name}.tmp-{pid}")
|
|
try:
|
|
tmp.write_text(json.dumps(payload), encoding = "utf-8")
|
|
os.replace(tmp, path)
|
|
except OSError as exc:
|
|
logger.debug("Could not write worker breadcrumb %s: %s", path, exc)
|
|
try:
|
|
tmp.unlink(missing_ok = True)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def remove_worker_breadcrumb(key: str) -> None:
|
|
path = _worker_breadcrumb_path(key)
|
|
if path is None:
|
|
return
|
|
_safe_unlink(path)
|
|
|
|
|
|
def _safe_unlink(path: Path) -> None:
|
|
try:
|
|
path.unlink(missing_ok = True)
|
|
except OSError as exc:
|
|
logger.debug("Could not remove %s: %s", path, exc)
|
|
|
|
|
|
def _process_alive(pid: int) -> bool:
|
|
if sys.platform == "win32":
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
|
|
SYNCHRONIZE = 0x00100000
|
|
ERROR_INVALID_PARAMETER = 87
|
|
kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
|
|
kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
|
|
kernel32.OpenProcess.restype = wintypes.HANDLE
|
|
kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
|
|
kernel32.CloseHandle.restype = wintypes.BOOL
|
|
ctypes.set_last_error(0)
|
|
handle = kernel32.OpenProcess(SYNCHRONIZE, False, pid)
|
|
if not handle:
|
|
return ctypes.get_last_error() != ERROR_INVALID_PARAMETER
|
|
kernel32.CloseHandle(handle)
|
|
return True
|
|
try:
|
|
os.kill(pid, 0)
|
|
return True
|
|
except ProcessLookupError:
|
|
return False
|
|
except OSError:
|
|
return True
|
|
|
|
|
|
def _read_process_cmdline(pid: int) -> Optional[str]:
|
|
proc_cmdline = Path(f"/proc/{pid}/cmdline")
|
|
try:
|
|
if proc_cmdline.exists():
|
|
raw = proc_cmdline.read_bytes()
|
|
return raw.replace(b"\x00", b" ").decode("utf-8", "replace")
|
|
except OSError:
|
|
pass
|
|
try:
|
|
import psutil
|
|
return " ".join(psutil.Process(pid).cmdline())
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _cmdline_repo_id(cmdline: str) -> Optional[str]:
|
|
try:
|
|
args = shlex.split(cmdline)
|
|
except ValueError:
|
|
args = cmdline.split()
|
|
for i, arg in enumerate(args):
|
|
if arg == "--repo-id" and i + 1 < len(args):
|
|
return args[i + 1]
|
|
if arg.startswith("--repo-id="):
|
|
return arg.split("=", 1)[1]
|
|
return None
|
|
|
|
|
|
def _is_our_worker(pid: int, repo_id: Optional[str]) -> bool:
|
|
cmdline = _read_process_cmdline(pid)
|
|
if cmdline is None:
|
|
return False
|
|
if "hub.workers.hf_download" not in cmdline:
|
|
return False
|
|
# Exact --repo-id match: a substring match would let a stale breadcrumb for
|
|
# Org/Model reap a live worker for Org/Model-v2.
|
|
if isinstance(repo_id, str) and repo_id:
|
|
return _cmdline_repo_id(cmdline) == repo_id
|
|
return True
|
|
|
|
|
|
def _kill_orphan(pid: int) -> None:
|
|
try:
|
|
os.kill(pid, signal.SIGTERM if sys.platform == "win32" else signal.SIGKILL)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _settle_orphaned_download(
|
|
repo_type: Optional[str],
|
|
repo_id: Optional[str],
|
|
variant: Optional[str],
|
|
transport: Optional[str],
|
|
hub_cache: Optional[str] = None,
|
|
) -> None:
|
|
"""Persist a cancel marker for a reaped orphan still mid-download so the next
|
|
launch settles it to a resumable "cancelled" state instead of a phantom-running
|
|
row.
|
|
|
|
Gated on surviving partial state and on the recorded manifest not already
|
|
verifying against an active snapshot, so a download that finished before its
|
|
breadcrumb was cleaned up is never mislabeled cancelled. For a GGUF variant
|
|
manifest with blob hashes, the partial-state check is scoped to those hashes so
|
|
a sibling variant cannot contaminate this orphan's state. The recorded
|
|
transport is preserved so the resume affordance stays accurate."""
|
|
if repo_type not in ("model", "dataset") or not repo_id:
|
|
return
|
|
from hub.utils import download_manifest
|
|
|
|
cache_root = Path(hub_cache) if isinstance(hub_cache, str) and hub_cache else None
|
|
|
|
manifest = download_manifest.read_manifest(
|
|
repo_type,
|
|
repo_id,
|
|
variant,
|
|
hub_cache = cache_root,
|
|
)
|
|
if repo_type == "model" and variant and manifest is None:
|
|
return
|
|
if manifest is None:
|
|
if not has_active_incomplete_blobs(repo_type, repo_id, root = cache_root):
|
|
return
|
|
else:
|
|
if _manifest_verifies_against_active_cache(
|
|
repo_type,
|
|
repo_id,
|
|
manifest,
|
|
root = cache_root,
|
|
):
|
|
return
|
|
if not _manifest_has_active_incomplete_blobs(
|
|
repo_type,
|
|
repo_id,
|
|
manifest,
|
|
root = cache_root,
|
|
):
|
|
return
|
|
persist_cancel_marker(
|
|
repo_type,
|
|
repo_id,
|
|
variant,
|
|
transport,
|
|
hub_cache = hub_cache,
|
|
logger = logger,
|
|
)
|
|
|
|
|
|
def reap_orphan_workers() -> None:
|
|
"""Kill download workers left running by a previous backend instance.
|
|
|
|
Verifies each breadcrumb's PID is alive AND its command line is one of our
|
|
workers before terminating, so a recycled PID can't take down an unrelated
|
|
process. Partial blobs are never touched, so a reaped download stays
|
|
resumable; an interrupted one with bytes on disk is settled to a cancelled
|
|
marker (see :func:`_settle_orphaned_download`) so its resume affordance
|
|
survives a hard crash like a graceful shutdown's does. Runs once at startup
|
|
and never raises."""
|
|
parent = state_dir.workers_dir()
|
|
if parent is None:
|
|
return
|
|
try:
|
|
entries = list(parent.iterdir())
|
|
except OSError:
|
|
return
|
|
for entry in entries:
|
|
if not entry.is_file() or not entry.name.endswith(".json"):
|
|
continue
|
|
try:
|
|
data = json.loads(entry.read_text(encoding = "utf-8"))
|
|
except (OSError, ValueError):
|
|
_safe_unlink(entry)
|
|
continue
|
|
pid = data.get("pid") if isinstance(data, dict) else None
|
|
repo_id = data.get("repo_id") if isinstance(data, dict) else None
|
|
if not isinstance(pid, int) or pid <= 0:
|
|
_safe_unlink(entry)
|
|
continue
|
|
try:
|
|
if _process_alive(pid) and _is_our_worker(pid, repo_id):
|
|
_kill_orphan(pid)
|
|
logger.warning(
|
|
"Reaped orphaned download worker pid=%s repo=%s from a "
|
|
"previous backend instance.",
|
|
pid,
|
|
repo_id,
|
|
)
|
|
_settle_orphaned_download(
|
|
data.get("repo_type"),
|
|
repo_id,
|
|
data.get("variant"),
|
|
data.get("cancel_marker_transport") or data.get("transport"),
|
|
data.get("hub_cache"),
|
|
)
|
|
except Exception as exc:
|
|
logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc)
|
|
_safe_unlink(entry)
|
|
|
|
|
|
def _purge_incomplete_blobs(
|
|
entry: Path,
|
|
only_hashes: Optional[frozenset[str]] = None,
|
|
protected_hashes: Optional[frozenset[str]] = None,
|
|
) -> int:
|
|
"""Delete matching ``*.incomplete`` blobs beneath *entry*; return the count
|
|
removed. Per-file failures are swallowed.
|
|
|
|
``only_hashes`` whitelists which partials may be purged; ``None`` means
|
|
every partial (full-repo snapshot/dataset). ``protected_hashes`` is honoured
|
|
unconditionally, even when ``only_hashes`` is ``None``, so a blob a
|
|
concurrent same-repo peer is writing is never purged from under it."""
|
|
blobs_dir = entry / "blobs"
|
|
if not blobs_dir.is_dir():
|
|
return 0
|
|
removed = 0
|
|
try:
|
|
candidates = list(blobs_dir.iterdir())
|
|
except OSError:
|
|
return 0
|
|
for blob in candidates:
|
|
try:
|
|
if not blob.is_file():
|
|
continue
|
|
if not blob.name.endswith(INCOMPLETE_SUFFIX):
|
|
continue
|
|
blob_hash = blob.name[: -len(INCOMPLETE_SUFFIX)]
|
|
if protected_hashes and blob_hash in protected_hashes:
|
|
continue
|
|
if only_hashes is not None and blob_hash not in only_hashes:
|
|
continue
|
|
blob.unlink()
|
|
removed += 1
|
|
except OSError:
|
|
# Swallow; downstream snapshot_download surfaces a precise error if
|
|
# it actually can't proceed.
|
|
continue
|
|
return removed
|
|
|
|
|
|
def _iter_active_snapshot_dirs(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
*,
|
|
root: Optional[Path] = None,
|
|
) -> Iterator[Path]:
|
|
for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root):
|
|
snapshots_dir = entry / "snapshots"
|
|
if not snapshots_dir.is_dir():
|
|
continue
|
|
try:
|
|
snapshots = list(snapshots_dir.iterdir())
|
|
except OSError:
|
|
continue
|
|
for snapshot in snapshots:
|
|
if snapshot.is_dir():
|
|
yield snapshot
|
|
|
|
|
|
def _manifest_verifies_against_active_cache(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
manifest,
|
|
*,
|
|
root: Optional[Path] = None,
|
|
) -> bool:
|
|
from hub.utils import download_manifest
|
|
for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id, root = root):
|
|
if download_manifest.verify_against_disk(manifest, snapshot_dir).ok:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _manifest_has_active_incomplete_blobs(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
manifest,
|
|
*,
|
|
root: Optional[Path] = None,
|
|
) -> bool:
|
|
if not getattr(manifest, "variant", None):
|
|
return has_active_incomplete_blobs(repo_type, repo_id, root = root)
|
|
expected_hashes = frozenset(
|
|
expected.sha256 for expected in manifest.expected_files if expected.sha256
|
|
)
|
|
if not expected_hashes:
|
|
return has_active_incomplete_blobs(repo_type, repo_id, root = root)
|
|
return bool(
|
|
incomplete_blob_hashes(
|
|
repo_type,
|
|
repo_id,
|
|
active_only = True,
|
|
root = root,
|
|
).intersection(expected_hashes)
|
|
)
|
|
|
|
|
|
def _marker_path(entry: Path, variant: Optional[str] = None) -> Path:
|
|
if not variant:
|
|
return entry / TRANSPORT_MARKER_NAME
|
|
digest = hashlib.sha256(variant.strip().lower().encode("utf-8")).hexdigest()[:24]
|
|
return entry / f"{TRANSPORT_MARKER_NAME}.gguf-{digest}"
|
|
|
|
|
|
def _is_transport_marker_file(path: Path) -> bool:
|
|
# Matches ".transport", its tmps, and variant-scoped ".transport.gguf-*".
|
|
# Real HF cache entries (blobs/refs/snapshots/.no_exist) never start with
|
|
# ".transport.".
|
|
return path.name == TRANSPORT_MARKER_NAME or path.name.startswith(f"{TRANSPORT_MARKER_NAME}.")
|
|
|
|
|
|
def _companion_marker_path(entry: Path) -> Path:
|
|
return entry / f"{TRANSPORT_MARKER_NAME}.companion"
|
|
|
|
|
|
def _read_marker_value(marker: Path) -> Optional[str]:
|
|
try:
|
|
if not marker.exists():
|
|
return None
|
|
value = marker.read_text(encoding = "utf-8").strip()
|
|
except (OSError, UnicodeDecodeError):
|
|
# UnicodeDecodeError is a ValueError, so it would escape and abort
|
|
# prepare_cache_for_transport. An unknown value just purges and restarts.
|
|
return None
|
|
return value if value in VALID_TRANSPORTS else None
|
|
|
|
|
|
def _write_marker_value(marker: Path, mode: str) -> None:
|
|
try:
|
|
# tmp + rename so a SIGKILL mid-write can't leave a half-written marker.
|
|
# The tmp name is per-process so concurrent writers don't clobber tmps.
|
|
tmp = marker.with_name(f"{marker.name}.tmp-{os.getpid()}")
|
|
tmp.write_text(mode, encoding = "utf-8")
|
|
os.replace(tmp, marker)
|
|
except OSError:
|
|
# Best-effort: a missing marker next run purges the partial defensively,
|
|
# the safe failure mode.
|
|
pass
|
|
|
|
|
|
def _read_marker(entry: Path, variant: Optional[str] = None) -> Optional[str]:
|
|
return _read_marker_value(_marker_path(entry, variant))
|
|
|
|
|
|
def _write_marker(
|
|
entry: Path,
|
|
mode: str,
|
|
variant: Optional[str] = None,
|
|
) -> None:
|
|
_write_marker_value(_marker_path(entry, variant), mode)
|
|
|
|
|
|
def _read_companion_marker(entry: Path) -> Optional[str]:
|
|
return _read_marker_value(_companion_marker_path(entry))
|
|
|
|
|
|
def _write_companion_marker(entry: Path, mode: str) -> None:
|
|
_write_marker_value(_companion_marker_path(entry), mode)
|
|
|
|
|
|
def prepare_cache_for_transport(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
mode: str,
|
|
variant: Optional[str] = None,
|
|
only_blob_hashes: Optional[frozenset[str]] = None,
|
|
companion_blob_hashes: Optional[frozenset[str]] = None,
|
|
protected_blob_hashes: Optional[frozenset[str]] = None,
|
|
root: Optional[Path] = None,
|
|
) -> int:
|
|
"""Guarantee any pre-existing ``.incomplete`` blobs are SAFE to resume under
|
|
*mode*. Returns the number of partial blobs purged for untrusted provenance.
|
|
|
|
Two marker scopes govern GGUF downloads. ``only_blob_hashes`` are the
|
|
variant's own (main quant) blobs, judged by the ``variant``-scoped marker;
|
|
``None`` widens the scope to every partial for full-repo snapshots/datasets.
|
|
``companion_blob_hashes`` are blobs shared across sibling variants (a vision
|
|
mmproj), judged by a separate repo-scoped companion marker — so a companion
|
|
partial is trusted against the transport that wrote it, not against
|
|
whichever sibling variant resumes next.
|
|
|
|
The contract:
|
|
- HTTP mode: a partial is trusted ONLY when its governing marker equals
|
|
``"http"``. Any other case (missing/unreadable/mismatched marker) purges,
|
|
since the HTTP resumer would otherwise append to a sparse
|
|
XET/parallel-Range partial and silently produce a corrupt blob.
|
|
- XET mode: incomplete blobs are purged (``hf_xet.download_files`` rewrites
|
|
from scratch, so this only fixes UI accounting — bytes already in CAS are
|
|
reused via the chunk-cache). Scoped to ``only_blob_hashes``: companion
|
|
blobs fall outside that set and survive (shared, and XET overwrites them).
|
|
|
|
``protected_blob_hashes`` are blobs a concurrent same-repo peer is writing;
|
|
they are excluded from every purge so a shared companion is never deleted
|
|
mid-write.
|
|
|
|
Scope: ``root`` selects the cache captured by the caller. It defaults to the
|
|
active ``HF_HUB_CACHE`` root for workers that inherit their cache through
|
|
the environment. Markers are written for the new mode before returning.
|
|
"""
|
|
if mode not in VALID_TRANSPORTS:
|
|
raise ValueError(f"Invalid transport mode: {mode!r}")
|
|
root = hf_cache_root(create = True) if root is None else hf_cache_root(create = True, root = root)
|
|
if root is None:
|
|
return 0
|
|
target = target_dir_name(repo_type, repo_id)
|
|
try:
|
|
entries = [e for e in root.iterdir() if e.name.lower() == target]
|
|
except OSError:
|
|
return 0
|
|
if not entries:
|
|
# First download: pre-create the repo dir so the marker lands before the
|
|
# worker writes any bytes. Otherwise a SIGKILL mid-download leaves a
|
|
# partial with no marker that the resume then purges.
|
|
canonical = repo_cache_dir_name(repo_type, repo_id)
|
|
new_entry = root / canonical
|
|
try:
|
|
new_entry.mkdir(exist_ok = True)
|
|
except OSError:
|
|
return 0
|
|
entries = [new_entry]
|
|
protected = protected_blob_hashes or frozenset()
|
|
has_companion = bool(companion_blob_hashes)
|
|
total_purged = 0
|
|
for entry in entries:
|
|
if mode == TRANSPORT_XET:
|
|
total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected)
|
|
else:
|
|
if _read_marker(entry, variant) != mode:
|
|
total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected)
|
|
if companion_blob_hashes and _read_companion_marker(entry) != mode:
|
|
total_purged += _purge_incomplete_blobs(entry, companion_blob_hashes, protected)
|
|
_write_marker(entry, mode, variant)
|
|
if has_companion:
|
|
_write_companion_marker(entry, mode)
|
|
return total_purged
|
|
|
|
|
|
_HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}")
|
|
_BEARER_RE = re.compile(r"(?i)bearer\s+[A-Za-z0-9._\-]+")
|
|
|
|
|
|
def scrub_secrets(text: str, *, hf_token: Optional[str] = None) -> str:
|
|
if not text:
|
|
return text
|
|
cleaned = text
|
|
if hf_token:
|
|
cleaned = cleaned.replace(hf_token, "***")
|
|
cleaned = _BEARER_RE.sub("Bearer ***", cleaned)
|
|
cleaned = _HF_TOKEN_RE.sub("hf_***", cleaned)
|
|
return cleaned
|
|
|
|
|
|
def purge_empty_marker_dir(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
variant: Optional[str] = None,
|
|
) -> bool:
|
|
"""Remove the failed download's own transport marker from a marker-only dir.
|
|
|
|
``prepare_cache_for_transport`` pre-creates the dir + marker before any
|
|
download; a failure during validation/auth/network setup leaves the dir as
|
|
marker-only litter. Only the failed download's OWN marker is removed (the
|
|
repo-scope ``.transport`` or the variant-scoped ``.transport.gguf-*`` plus
|
|
its ``.tmp-*`` siblings); a sibling variant's marker and the shared
|
|
``.transport.companion`` are left intact, so cancelling one quant never
|
|
strips a peer's provenance. A dir holding ``blobs/``/``snapshots/``/``refs/``
|
|
won't match and is left untouched, so a resumable partial isn't blown away.
|
|
"""
|
|
cleaned = False
|
|
for entry in iter_repo_cache_dirs(repo_type, repo_id):
|
|
try:
|
|
contents = list(entry.iterdir())
|
|
except OSError:
|
|
continue
|
|
if not contents or not all(_is_transport_marker_file(item) for item in contents):
|
|
continue
|
|
own_name = _marker_path(entry, variant).name
|
|
own_markers = [
|
|
item
|
|
for item in contents
|
|
if item.name == own_name or item.name.startswith(f"{own_name}.tmp")
|
|
]
|
|
if not own_markers:
|
|
continue
|
|
try:
|
|
for marker in own_markers:
|
|
marker.unlink()
|
|
except OSError:
|
|
continue
|
|
cleaned = True
|
|
try:
|
|
entry.rmdir()
|
|
except OSError:
|
|
continue
|
|
return cleaned
|
|
|
|
|
|
def read_active_transport_marker(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
variant: Optional[str] = None,
|
|
) -> Optional[str]:
|
|
for entry in iter_active_repo_cache_dirs(repo_type, repo_id):
|
|
value = _read_marker(entry, variant)
|
|
if value is not None:
|
|
return value
|
|
return None
|
|
|
|
|
|
def is_resumable_partial(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
variant: Optional[str] = None,
|
|
) -> bool:
|
|
"""True only when a partial exists AND was produced by a byte-resumable
|
|
writer (the HTTP transport). XET partials exist on disk but are discarded on
|
|
the next download attempt."""
|
|
if not has_active_incomplete_blobs(repo_type, repo_id):
|
|
return False
|
|
return read_active_transport_marker(repo_type, repo_id, variant) == TRANSPORT_HTTP
|
|
|
|
|
|
def incomplete_blob_hashes(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
*,
|
|
active_only: bool = False,
|
|
root: Optional[Path] = None,
|
|
) -> set[str]:
|
|
out: set[str] = set()
|
|
entries = (
|
|
iter_active_repo_cache_dirs(repo_type, repo_id, root = root)
|
|
if active_only
|
|
else iter_repo_cache_dirs(repo_type, repo_id)
|
|
)
|
|
for entry in entries:
|
|
blobs_dir = entry / "blobs"
|
|
if not blobs_dir.is_dir():
|
|
continue
|
|
try:
|
|
for blob in blobs_dir.iterdir():
|
|
if blob.is_file() and blob.name.endswith(INCOMPLETE_SUFFIX):
|
|
out.add(blob.name[: -len(INCOMPLETE_SUFFIX)])
|
|
except OSError:
|
|
continue
|
|
return out
|
|
|
|
|
|
def completed_blob_bytes(
|
|
repo_type: str,
|
|
repo_id: str,
|
|
blob_hashes: frozenset[str],
|
|
*,
|
|
root: Optional[Path] = None,
|
|
) -> int:
|
|
"""Sum finalized blob bytes for *blob_hashes* in a single HF cache root.
|
|
|
|
A worker only writes to its captured ``HF_HUB_CACHE`` root, so a baseline
|
|
must be scoped to that root (``root``), not re-resolved to whatever cache is
|
|
active now; otherwise a runtime cache switch makes the retry baseline count
|
|
bytes from the wrong disk.
|
|
"""
|
|
if not blob_hashes:
|
|
return 0
|
|
total = 0
|
|
for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root):
|
|
blobs_dir = entry / "blobs"
|
|
if not blobs_dir.is_dir():
|
|
continue
|
|
for blob_hash in blob_hashes:
|
|
blob = blobs_dir / blob_hash
|
|
try:
|
|
if blob.is_file():
|
|
total += max(0, int(blob.stat().st_size))
|
|
except OSError:
|
|
continue
|
|
return total
|
|
|
|
|
|
def existing_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int:
|
|
"""Bytes already on disk (finalized + ``.incomplete``) for *blob_hashes* in
|
|
the active HF cache root. A blob is in exactly one state, so summing both
|
|
candidate names never double-counts. Used to size what a (possibly resumed)
|
|
download still needs to write before the run starts."""
|
|
if not blob_hashes:
|
|
return 0
|
|
total = 0
|
|
for entry in iter_active_repo_cache_dirs(repo_type, repo_id):
|
|
blobs_dir = entry / "blobs"
|
|
if not blobs_dir.is_dir():
|
|
continue
|
|
for blob_hash in blob_hashes:
|
|
for name in (blob_hash, f"{blob_hash}{INCOMPLETE_SUFFIX}"):
|
|
blob = blobs_dir / name
|
|
try:
|
|
if blob.is_file():
|
|
total += max(0, int(blob.stat().st_size))
|
|
except OSError:
|
|
continue
|
|
return total
|
|
|
|
|
|
JobState = Literal["idle", "running", "cancelling", "cancelled", "complete", "error"]
|
|
|
|
TERMINAL_STATES = frozenset({"complete", "cancelled", "error"})
|
|
_ACTIVE_STATES = frozenset({"running", "cancelling"})
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class DownloadState:
|
|
state: JobState
|
|
error: Optional[str] = None
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class DownloadMetadata:
|
|
repo_type: RepoType
|
|
repo_id: str
|
|
variant: Optional[str]
|
|
transport: Optional[str]
|
|
cancel_marker_transport: Optional[str] = None
|
|
# GGUF variant main/writable hashes, identifying the variant-specific shards
|
|
# for concurrency decisions.
|
|
blob_hashes: frozenset[str] = field(default_factory = frozenset)
|
|
# Full required hash set for progress/completion (includes the shared mmproj
|
|
# companion for vision GGUF repos).
|
|
progress_blob_hashes: frozenset[str] = field(default_factory = frozenset)
|
|
# Bytes already complete before this job started; not counted as this run's
|
|
# progress.
|
|
completed_baseline_bytes: int = 0
|
|
hub_cache: Optional[str] = None
|
|
xet_cache: Optional[str] = None
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class ActiveDownloadRef:
|
|
key: str
|
|
state: str
|
|
metadata: Optional[DownloadMetadata]
|
|
generation: int
|
|
|
|
|
|
def normalize_repo_key(repo_id: str) -> str:
|
|
return repo_id.strip().lower()
|
|
|
|
|
|
def normalize_job_key(key: str) -> str:
|
|
repo, sep, variant = key.partition("::")
|
|
repo_key = normalize_repo_key(repo)
|
|
return f"{repo_key}{sep}{variant.strip().lower()}" if sep else repo_key
|
|
|
|
|
|
def _repo_of_key(key: str) -> str:
|
|
return normalize_repo_key(key.split("::", 1)[0])
|
|
|
|
|
|
def variant_from_key(key: str) -> Optional[str]:
|
|
"""Parse the variant suffix from a 'repo_id::variant' key. Empty
|
|
variant returns None — matches the manifest/marker calling
|
|
convention for full-snapshot models and datasets."""
|
|
if "::" not in key:
|
|
return None
|
|
_, _, variant = key.partition("::")
|
|
return variant or None
|
|
|
|
|
|
def persist_cancel_marker(
|
|
repo_type: Optional[RepoType],
|
|
repo_id: Optional[str],
|
|
variant: Optional[str],
|
|
transport: Optional[str],
|
|
*,
|
|
hub_cache: Optional[str] = None,
|
|
logger = logger,
|
|
) -> None:
|
|
if not repo_type or not repo_id:
|
|
return
|
|
try:
|
|
from hub.utils.download_manifest import write_cancel_marker
|
|
if not write_cancel_marker(
|
|
repo_type,
|
|
repo_id,
|
|
variant,
|
|
transport = transport,
|
|
hub_cache = hub_cache,
|
|
):
|
|
logger.debug("write_cancel_marker returned False for %s", repo_id)
|
|
except Exception as exc:
|
|
logger.debug("write_cancel_marker failed for %s: %s", repo_id, exc)
|
|
|
|
|
|
_REGISTRIES: "weakref.WeakSet[DownloadRegistry]" = weakref.WeakSet()
|
|
_NAMED_REGISTRIES: dict[str, "DownloadRegistry"] = {}
|
|
_NAMED_REGISTRIES_LOCK = threading.Lock()
|
|
|
|
|
|
def terminate_active_downloads() -> None:
|
|
"""Best-effort shutdown hook called from the FastAPI lifespan.
|
|
|
|
Walks every live DownloadRegistry instance and SIGKILLs any in-flight
|
|
workers so the parent exit path doesn't leak zombies. The WeakSet drops
|
|
ad-hoc registries (e.g. test fixtures) automatically once their last
|
|
strong reference is gone; the long-lived named registries stay reachable
|
|
via ``_NAMED_REGISTRIES``. Quiet on its own failures: shutdown must not
|
|
raise.
|
|
"""
|
|
for registry in list(_REGISTRIES):
|
|
try:
|
|
registry.terminate_all("download")
|
|
except Exception as exc:
|
|
logger.warning("terminate_active_downloads: %s", exc)
|
|
|
|
|
|
class DownloadRegistry:
|
|
"""Thread-safe state machine for background HF download jobs.
|
|
|
|
One instance backs model downloads (keys ``repo_id::variant``) and another
|
|
backs dataset downloads (keys ``repo_id``). Repo-scoped tracking serializes
|
|
full snapshots, datasets, cross-transport work, and deletes; same-transport
|
|
GGUF variants may run concurrently.
|
|
"""
|
|
|
|
def __init__(self, max_terminal: int = 64) -> None:
|
|
self._jobs: dict[str, DownloadState] = {}
|
|
self._processes: dict[str, subprocess.Popen] = {}
|
|
self._repo_active: dict[str, set[str]] = {}
|
|
self._metadata: dict[str, DownloadMetadata] = {}
|
|
self._cancel_marker_transports: dict[str, str] = {}
|
|
self._pending_cancel: dict[str, Optional[int]] = {}
|
|
self._generations: dict[str, int] = {}
|
|
# Monotonic across keys so an evicted then re-claimed key never reuses a
|
|
# prior generation (which would let a stale cancel match a new run).
|
|
self._generation_seq = 0
|
|
self._deleting: dict[str, set[Optional[str]]] = {}
|
|
self._lock = threading.Lock()
|
|
_REGISTRIES.add(self)
|
|
self._max_terminal = max_terminal
|
|
|
|
def _put_terminal_job_locked(
|
|
self,
|
|
key: str,
|
|
state: JobState,
|
|
error: Optional[str] = None,
|
|
) -> None:
|
|
self._jobs.pop(key, None)
|
|
self._jobs[key] = DownloadState(state, error)
|
|
if len(self._jobs) > self._max_terminal:
|
|
for stale_key, stale in list(self._jobs.items()):
|
|
if stale.state in TERMINAL_STATES and stale_key != key:
|
|
self._jobs.pop(stale_key, None)
|
|
self._metadata.pop(stale_key, None)
|
|
self._generations.pop(stale_key, None)
|
|
if len(self._jobs) <= self._max_terminal:
|
|
break
|
|
|
|
def set_job(
|
|
self,
|
|
key: str,
|
|
state: JobState,
|
|
error: Optional[str] = None,
|
|
) -> None:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
if state in TERMINAL_STATES:
|
|
self._put_terminal_job_locked(key, state, error)
|
|
self._pending_cancel.pop(key, None)
|
|
self._cancel_marker_transports.pop(key, None)
|
|
repo = _repo_of_key(key)
|
|
active = self._repo_active.get(repo)
|
|
if active is not None:
|
|
active.discard(key)
|
|
if not active:
|
|
self._repo_active.pop(repo, None)
|
|
else:
|
|
self._jobs[key] = DownloadState(state, error)
|
|
|
|
def set_error_unless_cancelled(
|
|
self, key: str, error: str
|
|
) -> tuple[JobState, Optional[DownloadMetadata]]:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
current = self._jobs.get(key, DownloadState("idle")).state
|
|
has_pending_cancel = key in self._pending_cancel
|
|
pending_generation = self._pending_cancel.get(key)
|
|
metadata = self._metadata.get(key)
|
|
should_cancel = current == "cancelling" or (
|
|
has_pending_cancel and self._generation_matches_locked(key, pending_generation)
|
|
)
|
|
terminal_state: JobState = "cancelled" if should_cancel else "error"
|
|
marker_transport = self._cancel_marker_transports.pop(key, None)
|
|
if marker_transport is None and metadata is not None:
|
|
marker_transport = metadata.cancel_marker_transport
|
|
self._put_terminal_job_locked(
|
|
key,
|
|
terminal_state,
|
|
None if should_cancel else error,
|
|
)
|
|
self._pending_cancel.pop(key, None)
|
|
repo = _repo_of_key(key)
|
|
active = self._repo_active.get(repo)
|
|
if active is not None:
|
|
active.discard(key)
|
|
if not active:
|
|
self._repo_active.pop(repo, None)
|
|
if should_cancel and metadata is not None and marker_transport is not None:
|
|
metadata = replace(metadata, transport = marker_transport)
|
|
return terminal_state, metadata
|
|
|
|
def update_job_transport(self, key: str, transport: str) -> None:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
metadata = self._metadata.get(key)
|
|
if metadata is None or metadata.transport == transport:
|
|
return
|
|
self._metadata[key] = replace(metadata, transport = transport)
|
|
|
|
def release_active_slot(self, key: str) -> None:
|
|
key = normalize_job_key(key)
|
|
repo = _repo_of_key(key)
|
|
with self._lock:
|
|
active = self._repo_active.get(repo)
|
|
if active is None:
|
|
return
|
|
active.discard(key)
|
|
if not active:
|
|
self._repo_active.pop(repo, None)
|
|
|
|
def get_job(self, key: str) -> DownloadState:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
return self._jobs.get(key, DownloadState("idle"))
|
|
|
|
def current_generation(self, key: str) -> int:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
return self._generations.get(key, 0)
|
|
|
|
def get_job_metadata(self, key: str) -> Optional[DownloadMetadata]:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
return self._metadata.get(key)
|
|
|
|
def _generation_matches_locked(self, key: str, generation: Optional[int]) -> bool:
|
|
key = normalize_job_key(key)
|
|
return generation is None or self._generations.get(key, 0) == generation
|
|
|
|
def register_process(self, key: str, proc: subprocess.Popen) -> bool:
|
|
"""Register *proc* for *key*. Returns ``False`` when a cancel was
|
|
requested during the claim→register window (the caller must kill
|
|
*proc* immediately); ``True`` otherwise."""
|
|
key = normalize_job_key(key)
|
|
metadata_to_persist: Optional[DownloadMetadata] = None
|
|
registered = False
|
|
breadcrumb_metadata: Optional[DownloadMetadata] = None
|
|
with self._lock:
|
|
has_pending_cancel = key in self._pending_cancel
|
|
pending_generation = self._pending_cancel.pop(key, None)
|
|
if has_pending_cancel and self._generation_matches_locked(
|
|
key,
|
|
pending_generation,
|
|
):
|
|
self._put_terminal_job_locked(key, "cancelled")
|
|
metadata_to_persist = self._metadata.pop(key, None)
|
|
marker_transport = self._cancel_marker_transports.pop(key, None)
|
|
if marker_transport is None and metadata_to_persist is not None:
|
|
marker_transport = metadata_to_persist.cancel_marker_transport
|
|
if metadata_to_persist is not None and marker_transport is not None:
|
|
metadata_to_persist = replace(
|
|
metadata_to_persist,
|
|
transport = marker_transport,
|
|
)
|
|
repo = _repo_of_key(key)
|
|
active = self._repo_active.get(repo)
|
|
if active is not None:
|
|
active.discard(key)
|
|
if not active:
|
|
self._repo_active.pop(repo, None)
|
|
else:
|
|
self._processes[key] = proc
|
|
breadcrumb_metadata = self._metadata.get(key)
|
|
registered = True
|
|
if registered:
|
|
try:
|
|
write_worker_breadcrumb(key, proc.pid, breadcrumb_metadata)
|
|
except Exception as exc:
|
|
logger.debug("Could not record worker breadcrumb: %s", exc)
|
|
return True
|
|
if metadata_to_persist is not None:
|
|
persist_cancel_marker(
|
|
metadata_to_persist.repo_type,
|
|
metadata_to_persist.repo_id,
|
|
metadata_to_persist.variant,
|
|
metadata_to_persist.transport,
|
|
hub_cache = metadata_to_persist.hub_cache,
|
|
)
|
|
return False
|
|
|
|
def mark_pending_cancel(
|
|
self,
|
|
key: str,
|
|
generation: Optional[int] = None,
|
|
) -> bool:
|
|
"""Record a cancel for an active job whose worker process hasn't
|
|
registered yet. Returns ``True`` when the pending cancel was armed,
|
|
so :func:`register_process` will kill the process on arrival."""
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
if self._jobs.get(key, DownloadState("idle")).state not in _ACTIVE_STATES:
|
|
return False
|
|
if not self._generation_matches_locked(key, generation):
|
|
return False
|
|
self._pending_cancel[key] = generation
|
|
self._jobs[key] = DownloadState("cancelling")
|
|
return True
|
|
|
|
def cancel_requested(self, key: str) -> bool:
|
|
"""True when *we* initiated a stop for *key* (a pending cancel armed
|
|
before the worker registered, or the job already moved to
|
|
``cancelling``). Lets exit classification tell an intentional kill
|
|
apart from an OOM/external SIGKILL."""
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
if key in self._pending_cancel:
|
|
return True
|
|
return self._jobs.get(key, DownloadState("idle")).state == "cancelling"
|
|
|
|
def get_process(self, key: str) -> Optional[subprocess.Popen]:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
return self._processes.get(key)
|
|
|
|
def drop_process(self, key: str, proc: subprocess.Popen) -> bool:
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
if self._processes.get(key) is not proc:
|
|
return False
|
|
self._processes.pop(key, None)
|
|
remove_worker_breadcrumb(key)
|
|
return True
|
|
|
|
def claim(
|
|
self,
|
|
key: str,
|
|
transport: str,
|
|
*,
|
|
repo_type: Optional[RepoType] = None,
|
|
repo_id: Optional[str] = None,
|
|
variant: Optional[str] = None,
|
|
blob_hashes: Optional[frozenset[str]] = None,
|
|
progress_blob_hashes: Optional[frozenset[str]] = None,
|
|
completed_baseline_bytes: int = 0,
|
|
admission_check: Optional[Callable[[], bool]] = None,
|
|
generation: Optional[int] = None,
|
|
replace_active: bool = False,
|
|
metadata_transport: Optional[str] = None,
|
|
cancel_marker_transport: Optional[str] = None,
|
|
hub_cache: Optional[str] = None,
|
|
xet_cache: Optional[str] = None,
|
|
) -> tuple[bool, str]:
|
|
key = normalize_job_key(key)
|
|
repo = _repo_of_key(key)
|
|
requested_hashes = blob_hashes or frozenset()
|
|
requested_progress_hashes = progress_blob_hashes or frozenset()
|
|
with self._lock:
|
|
# Run the final external admission check while the registry lock is
|
|
# held, immediately before inspecting and publishing active state.
|
|
# The GGUF load path establishes its marker before calling
|
|
# its active-job probe, so either this claim observes that marker
|
|
# or the load's later probe observes this claim.
|
|
if admission_check is not None and not admission_check():
|
|
return False, "admission_blocked"
|
|
deleting_scopes = self._deleting.get(repo)
|
|
if deleting_scopes is not None and (
|
|
None in deleting_scopes or variant_from_key(key) in deleting_scopes
|
|
):
|
|
return False, "deleting"
|
|
active = self._repo_active.get(repo, set())
|
|
stale_keys: list[str] = []
|
|
conflict_state: Optional[str] = None
|
|
for other_key in active:
|
|
if other_key == key:
|
|
continue
|
|
other_status = self._jobs.get(other_key)
|
|
if other_status is None or other_status.state not in _ACTIVE_STATES:
|
|
stale_keys.append(other_key)
|
|
continue
|
|
other_metadata = self._metadata.get(other_key)
|
|
# Same-transport variants of one model run concurrently: each
|
|
# worker purges only its own re-resolved main blobs and the
|
|
# shared companion is guarded by its marker. Cross-transport
|
|
# stays serialized so an HTTP resume and an XET rewrite never
|
|
# write one shared blob at once.
|
|
concurrent_gguf_variants = (
|
|
repo_type == "model"
|
|
and bool(variant)
|
|
and other_metadata is not None
|
|
and other_metadata.repo_type == "model"
|
|
and bool(other_metadata.variant)
|
|
and other_metadata.transport == transport
|
|
)
|
|
if concurrent_gguf_variants:
|
|
continue
|
|
conflict_state = other_status.state
|
|
break
|
|
for stale_key in stale_keys:
|
|
active.discard(stale_key)
|
|
if conflict_state is not None:
|
|
return False, conflict_state
|
|
current = self._jobs.get(key, DownloadState("idle")).state
|
|
if current in _ACTIVE_STATES and not replace_active:
|
|
return False, current
|
|
if generation is None:
|
|
self._generation_seq += 1
|
|
self._generations[key] = self._generation_seq
|
|
else:
|
|
self._generations[key] = generation
|
|
self._jobs[key] = DownloadState("running")
|
|
self._repo_active.setdefault(repo, active).add(key)
|
|
if repo_type and repo_id:
|
|
self._metadata[key] = DownloadMetadata(
|
|
repo_type = repo_type,
|
|
repo_id = repo_id,
|
|
variant = variant,
|
|
transport = metadata_transport if metadata_transport is not None else transport,
|
|
cancel_marker_transport = cancel_marker_transport,
|
|
blob_hashes = requested_hashes,
|
|
progress_blob_hashes = requested_progress_hashes,
|
|
completed_baseline_bytes = max(
|
|
0,
|
|
int(completed_baseline_bytes or 0),
|
|
),
|
|
hub_cache = hub_cache,
|
|
xet_cache = xet_cache,
|
|
)
|
|
if cancel_marker_transport is not None:
|
|
self._cancel_marker_transports[key] = cancel_marker_transport
|
|
else:
|
|
self._cancel_marker_transports.pop(key, None)
|
|
else:
|
|
self._metadata.pop(key, None)
|
|
self._cancel_marker_transports.pop(key, None)
|
|
return True, "running"
|
|
|
|
def adoptable(self, key: str) -> bool:
|
|
"""True when *key* itself has a live job a client can attach to.
|
|
|
|
Lets a rejected claim distinguish a collision with this key's own
|
|
in-flight job (pollable) from one blocked by a different repo job
|
|
or an in-progress delete, where no job exists for this key."""
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
return self._jobs.get(key, DownloadState("idle")).state in _ACTIVE_STATES
|
|
|
|
def _active_job_variant_locked(self, key: str) -> Optional[str]:
|
|
metadata = self._metadata.get(key)
|
|
if metadata is not None:
|
|
return (metadata.variant or "").strip().lower() or None
|
|
return variant_from_key(key)
|
|
|
|
def _delete_blocked_by_active_locked(self, repo_id: str, variant: Optional[str]) -> bool:
|
|
"""Whether an active download conflicts with deleting *repo_id*/*variant*.
|
|
|
|
A whole-repo delete (``variant is None``) conflicts with any active
|
|
download. A variant delete conflicts only with that same variant or a
|
|
whole-repo download writing the shared snapshot; other quantizations
|
|
download concurrently and never block it."""
|
|
active_keys = self._repo_active.get(repo_id, set())
|
|
for key in active_keys:
|
|
job = self._jobs.get(key)
|
|
if job is None or job.state not in _ACTIVE_STATES:
|
|
continue
|
|
if variant is None:
|
|
return True
|
|
other_variant = self._active_job_variant_locked(key)
|
|
if other_variant is None or other_variant == variant:
|
|
return True
|
|
for key, job in self._jobs.items():
|
|
if key in active_keys or _repo_of_key(key) != repo_id:
|
|
continue
|
|
if job.state not in _ACTIVE_STATES:
|
|
continue
|
|
if variant is None:
|
|
return True
|
|
other_variant = self._active_job_variant_locked(key)
|
|
if other_variant is None or other_variant == variant:
|
|
return True
|
|
return False
|
|
|
|
def peer_blob_hashes(self, key: str) -> frozenset[str]:
|
|
"""Union of the writable blob hashes of every OTHER active download for
|
|
this key's repo. A worker excludes these from its purge so it never
|
|
deletes an ``.incomplete`` a concurrent same-repo variant is writing
|
|
(e.g. a shared mmproj bundled with two GGUF quants)."""
|
|
key = normalize_job_key(key)
|
|
repo = _repo_of_key(key)
|
|
out: set[str] = set()
|
|
with self._lock:
|
|
for other_key in self._repo_active.get(repo, set()):
|
|
if other_key == key:
|
|
continue
|
|
job = self._jobs.get(other_key)
|
|
if job is None or job.state not in _ACTIVE_STATES:
|
|
continue
|
|
metadata = self._metadata.get(other_key)
|
|
if metadata is not None:
|
|
out |= set(metadata.progress_blob_hashes or metadata.blob_hashes)
|
|
return frozenset(out)
|
|
|
|
def active_jobs(self, repo_id: str) -> dict[str, str]:
|
|
"""Map of every active job key for *repo_id* to its state."""
|
|
repo_id = normalize_repo_key(repo_id)
|
|
with self._lock:
|
|
result: dict[str, str] = {}
|
|
for key in self._repo_active.get(repo_id, set()):
|
|
job = self._jobs.get(key)
|
|
if job is not None and job.state in _ACTIVE_STATES:
|
|
metadata = self._metadata.get(key)
|
|
display_key = (
|
|
f"{_repo_of_key(key)}::{metadata.variant}"
|
|
if metadata is not None and metadata.variant
|
|
else key
|
|
)
|
|
result[display_key] = job.state
|
|
return result
|
|
|
|
def active_job_refs(self, repo_id: Optional[str] = None) -> list[ActiveDownloadRef]:
|
|
repo_key = normalize_repo_key(repo_id) if repo_id else None
|
|
with self._lock:
|
|
if repo_key:
|
|
candidate_keys = list(self._repo_active.get(repo_key, set()))
|
|
else:
|
|
candidate_keys = [key for active in self._repo_active.values() for key in active]
|
|
# An XET->HTTP retry handoff briefly drops its key from _repo_active
|
|
# while its job stays active; include those released-but-active jobs
|
|
# so the waiting retry still lists and can be adopted or cancelled.
|
|
seen = set(candidate_keys)
|
|
for key, job in self._jobs.items():
|
|
if key in seen or job.state not in _ACTIVE_STATES:
|
|
continue
|
|
if repo_key is not None and _repo_of_key(key) != repo_key:
|
|
continue
|
|
candidate_keys.append(key)
|
|
refs: list[ActiveDownloadRef] = []
|
|
for key in candidate_keys:
|
|
job = self._jobs.get(key)
|
|
if job is None or job.state not in _ACTIVE_STATES:
|
|
continue
|
|
refs.append(
|
|
ActiveDownloadRef(
|
|
key = key,
|
|
state = job.state,
|
|
metadata = self._metadata.get(key),
|
|
generation = self._generations.get(key, 0),
|
|
)
|
|
)
|
|
return refs
|
|
|
|
def has_active_variant(self, repo_id: str, variant: Optional[str]) -> bool:
|
|
"""Whether an active model job targets this exact GGUF variant.
|
|
|
|
Scans the job table rather than only ``_repo_active`` so an XET-to-HTTP
|
|
retry handoff remains visible while it has temporarily released its
|
|
active slot.
|
|
"""
|
|
repo_key = normalize_repo_key(repo_id)
|
|
target = (variant or "").strip().lower() or None
|
|
with self._lock:
|
|
for key, job in self._jobs.items():
|
|
if _repo_of_key(key) != repo_key or job.state not in _ACTIVE_STATES:
|
|
continue
|
|
if self._active_job_variant_locked(key) == target:
|
|
return True
|
|
return False
|
|
|
|
def begin_delete(
|
|
self,
|
|
repo_id: str,
|
|
variant: Optional[str] = None,
|
|
) -> bool:
|
|
"""Reserve *repo_id* (or one GGUF *variant* of it) for deletion. Returns
|
|
``False`` when a conflicting download is active (a whole-repo delete vs
|
|
any download, a variant delete vs that same variant or a whole-repo
|
|
download), so sibling quantizations keep downloading. On success the
|
|
scope is marked so :func:`claim` rejects overlapping downloads until
|
|
:func:`end_delete` runs, closing the check-then-delete race against a
|
|
concurrently spawned worker."""
|
|
repo_id = normalize_repo_key(repo_id)
|
|
variant_key = (variant or "").strip().lower() or None
|
|
with self._lock:
|
|
if self._delete_blocked_by_active_locked(repo_id, variant_key):
|
|
return False
|
|
self._deleting.setdefault(repo_id, set()).add(variant_key)
|
|
return True
|
|
|
|
def end_delete(
|
|
self,
|
|
repo_id: str,
|
|
variant: Optional[str] = None,
|
|
) -> None:
|
|
repo_id = normalize_repo_key(repo_id)
|
|
variant_key = (variant or "").strip().lower() or None
|
|
with self._lock:
|
|
scopes = self._deleting.get(repo_id)
|
|
if scopes is None:
|
|
return
|
|
scopes.discard(variant_key)
|
|
if not scopes:
|
|
self._deleting.pop(repo_id, None)
|
|
|
|
def has_active_peer_variant(self, repo_id: str, variant: Optional[str]) -> bool:
|
|
"""Whether a DIFFERENT quantization of *repo_id* is downloading while
|
|
*variant* is being deleted. When one is, the delete reclaims only this
|
|
variant's files and leaves the shared companion (mmproj) for the live
|
|
sibling. Point-in-time (a sibling may claim just after it returns), but
|
|
safe: the finalized companion is held by deletion's reference-count
|
|
walk and a sibling starting mid-delete re-fetches it, so protection
|
|
never depends on the sibling having resolved its blob hashes."""
|
|
repo_id = normalize_repo_key(repo_id)
|
|
target = (variant or "").strip().lower() or None
|
|
with self._lock:
|
|
active_keys = self._repo_active.get(repo_id, set())
|
|
for key in active_keys:
|
|
job = self._jobs.get(key)
|
|
if job is None or job.state not in _ACTIVE_STATES:
|
|
continue
|
|
if self._active_job_variant_locked(key) != target:
|
|
return True
|
|
# An XET->HTTP retry peer between release_active_slot() and its reclaim
|
|
# is briefly absent from _repo_active while its job stays active and
|
|
# still owns the shared companion; mirror the released-but-active scan
|
|
# used by _delete_blocked_by_active_locked so it still blocks companion
|
|
# deletion of a different variant.
|
|
for key, job in self._jobs.items():
|
|
if key in active_keys or _repo_of_key(key) != repo_id:
|
|
continue
|
|
if job.state not in _ACTIVE_STATES:
|
|
continue
|
|
if self._active_job_variant_locked(key) != target:
|
|
return True
|
|
return False
|
|
|
|
def request_cancel(
|
|
self,
|
|
key: str,
|
|
proc: subprocess.Popen,
|
|
generation: Optional[int] = None,
|
|
) -> bool:
|
|
"""Authorize a SIGKILL for the registered *proc*. Idempotent across an
|
|
active job's lifetime: a repeated cancel while already ``cancelling``
|
|
still returns ``True`` so a kill that raced and lost can be re-sent."""
|
|
key = normalize_job_key(key)
|
|
with self._lock:
|
|
if self._processes.get(key) is not proc:
|
|
return False
|
|
if not self._generation_matches_locked(key, generation):
|
|
return False
|
|
if self._jobs.get(key, DownloadState("idle")).state not in _ACTIVE_STATES:
|
|
return False
|
|
self._jobs[key] = DownloadState("cancelling")
|
|
return True
|
|
|
|
def terminate_all(self, kind: str = "download") -> None:
|
|
settled_no_proc: list[Optional[DownloadMetadata]] = []
|
|
with self._lock:
|
|
live = [
|
|
(key, proc, self._metadata.get(key))
|
|
for key, proc in self._processes.items()
|
|
if proc.poll() is None
|
|
]
|
|
live_keys = {key for key, _proc, _metadata in live}
|
|
# Flag as an intentional stop so the watcher's exit classification
|
|
# reports them cancelled rather than an OOM/crash once SIGKILL lands.
|
|
for key, _proc, _metadata in live:
|
|
if self._jobs.get(key, DownloadState("idle")).state == "running":
|
|
self._jobs[key] = DownloadState("cancelling")
|
|
# Settle active jobs without a live worker too. Two cases: an
|
|
# XET->HTTP retry parked in the reclaim wait loop has dropped its
|
|
# worker and slot guard, so it is absent from `live`; and a
|
|
# registered worker that already exited with an error but whose
|
|
# watcher has not yet run would otherwise stay `running` and spawn an
|
|
# HTTP retry after this shutdown snapshot. Skip a registered worker
|
|
# that exited cleanly (rc == 0): it completed and the watcher will
|
|
# mark it done, so marking it cancelling would strand a stale marker.
|
|
for key, job in list(self._jobs.items()):
|
|
if job.state not in _ACTIVE_STATES or key in live_keys:
|
|
continue
|
|
proc = self._processes.get(key)
|
|
if proc is not None:
|
|
if proc.poll() == 0:
|
|
continue
|
|
# A registered worker that exited nonzero on its own over HTTP
|
|
# is a genuine terminal download failure, not a shutdown cancel
|
|
# and not retry-capable: leave its error status intact rather
|
|
# than persisting a cancel marker that would read as
|
|
# cancelled/resumable after restart. Only an exited XET worker
|
|
# could still spawn a post-shutdown HTTP retry, so only that
|
|
# needs settling here.
|
|
metadata = self._metadata.get(key)
|
|
if metadata is not None and metadata.transport == TRANSPORT_HTTP:
|
|
continue
|
|
self._pending_cancel[key] = self._generations.get(key)
|
|
self._jobs[key] = DownloadState("cancelling")
|
|
settled_no_proc.append(self._metadata.get(key))
|
|
# Persist a cancel marker for each settled no-live-worker job outside the
|
|
# lock (mirroring the reaped path) so shutdown records resumable/cancelled
|
|
# state even if it returns before the daemon watcher wakes to do so.
|
|
for metadata in settled_no_proc:
|
|
if metadata is not None:
|
|
persist_cancel_marker(
|
|
metadata.repo_type,
|
|
metadata.repo_id,
|
|
metadata.variant,
|
|
metadata.cancel_marker_transport or metadata.transport,
|
|
hub_cache = metadata.hub_cache,
|
|
)
|
|
reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = []
|
|
for key, proc, metadata in live:
|
|
try:
|
|
proc.kill()
|
|
except ProcessLookupError:
|
|
pass
|
|
except Exception as e:
|
|
logger.warning(f"shutdown: failed to kill {kind} worker for {key}: {e}")
|
|
if metadata is not None:
|
|
persist_cancel_marker(
|
|
metadata.repo_type,
|
|
metadata.repo_id,
|
|
metadata.variant,
|
|
metadata.cancel_marker_transport or metadata.transport,
|
|
hub_cache = metadata.hub_cache,
|
|
)
|
|
continue
|
|
reaped.append((key, proc, metadata))
|
|
deadline = time.monotonic() + 10.0
|
|
for key, proc, metadata in reaped:
|
|
try:
|
|
proc.wait(timeout = max(0.0, deadline - time.monotonic()))
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning(f"shutdown: {kind} worker for {key} did not exit after kill")
|
|
except Exception:
|
|
pass
|
|
# Mark only genuinely interrupted workers (rc != 0, or None on wait
|
|
# timeout); persisting before the exit is known would strand a stale
|
|
# marker on a worker that completed cleanly during shutdown.
|
|
if metadata is not None and proc.poll() != 0:
|
|
persist_cancel_marker(
|
|
metadata.repo_type,
|
|
metadata.repo_id,
|
|
metadata.variant,
|
|
metadata.cancel_marker_transport or metadata.transport,
|
|
hub_cache = metadata.hub_cache,
|
|
)
|
|
|
|
|
|
def _named_registry(name: str) -> DownloadRegistry:
|
|
with _NAMED_REGISTRIES_LOCK:
|
|
registry = _NAMED_REGISTRIES.get(name)
|
|
if registry is None:
|
|
registry = DownloadRegistry()
|
|
_NAMED_REGISTRIES[name] = registry
|
|
return registry
|
|
|
|
|
|
def get_models_registry() -> DownloadRegistry:
|
|
return _named_registry("models")
|
|
|
|
|
|
def get_datasets_registry() -> DownloadRegistry:
|
|
return _named_registry("datasets")
|