llm/scripts/download_hf_range.py

from __future__ import annotations

import argparse
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_url

from common import load_model_card


def selected_files(repo_id: str, allow_files: list[str] | None) -> list[tuple[str, int]]:
    api = HfApi()
    info = api.model_info(repo_id=repo_id, files_metadata=True)
    selected = []
    allowed = set(allow_files or [])
    for sibling in info.siblings:
        name = sibling.rfilename
        if name == ".gitattributes":
            continue
        if allowed and name not in allowed:
            continue
        if sibling.size is None:
            raise ValueError(f"Unknown remote size for {name}")
        selected.append((name, int(sibling.size)))
    return selected


def request_range(url: str, start: int, end: int, token: str | None) -> urllib.request.Request:
    headers = {
        "Range": f"bytes={start}-{end}",
        "User-Agent": "LLM-local-range-downloader/1.0",
    }
    if token:
        headers["Authorization"] = f"Bearer {token}"
    return urllib.request.Request(url, headers=headers, method="GET")


def local_size(path: Path) -> int:
    return path.stat().st_size if path.exists() else 0


def download_file(
    repo_id: str,
    filename: str,
    output: Path,
    expected_size: int,
    *,
    chunk_size: int,
    retries: int,
    retry_delay: float,
    timeout: int,
    token: str | None,
) -> bool:
    output.parent.mkdir(parents=True, exist_ok=True)
    url = hf_hub_url(repo_id, filename)

    while True:
        current_size = local_size(output)
        if current_size == expected_size:
            print(f"OK {filename}: {current_size}/{expected_size}")
            return True
        if current_size > expected_size:
            print(
                f"ERROR {filename}: local file is larger than expected: {current_size}/{expected_size}",
                file=sys.stderr,
            )
            return False

        start = current_size
        end = min(start + chunk_size - 1, expected_size - 1)
        print(f"GET {filename}: bytes {start}-{end} of {expected_size}")

        recompute_range = False
        for attempt in range(1, retries + 1):
            try:
                request = request_range(url, start, end, token)
                with urllib.request.urlopen(request, timeout=timeout) as response:
                    status = response.status
                    if status == 200:
                        raise RuntimeError("server ignored Range request")
                    if status not in (200, 206):
                        raise RuntimeError(f"unexpected HTTP status {status}")

                    mode = "ab" if start else "wb"
                    written = 0
                    with output.open(mode + "") as handle:
                        while True:
                            block = response.read(1024 * 1024)
                            if not block:
                                break
                            handle.write(block)
                            written += len(block)

                if written == 0:
                    raise RuntimeError("downloaded zero bytes")

                new_size = local_size(output)
                if new_size <= current_size:
                    raise RuntimeError(f"file did not grow: {new_size}")
                if new_size > expected_size:
                    raise RuntimeError(f"file grew past expected size: {new_size}/{expected_size}")
                break
            except (OSError, RuntimeError, urllib.error.URLError, urllib.error.HTTPError) as exc:
                new_size = local_size(output)
                if new_size > current_size:
                    print(
                        f"WARN {filename}: partial chunk saved, size is now {new_size}/{expected_size}",
                        file=sys.stderr,
                    )
                    recompute_range = True
                    break
                print(f"WARN {filename}: attempt {attempt}/{retries} failed: {exc}", file=sys.stderr)
                if attempt == retries:
                    return False
                time.sleep(retry_delay)
        if recompute_range:
            continue


def parse_size(value: str) -> int:
    text = value.strip().lower()
    multipliers = {
        "k": 1024,
        "kb": 1024,
        "m": 1024**2,
        "mb": 1024**2,
        "g": 1024**3,
        "gb": 1024**3,
    }
    for suffix, multiplier in multipliers.items():
        if text.endswith(suffix):
            return int(float(text[: -len(suffix)]) * multiplier)
    return int(text)


def main() -> int:
    parser = argparse.ArgumentParser(description="Download Hugging Face files using explicit HTTP ranges.")
    parser.add_argument("card_id")
    parser.add_argument("--local-dir", required=True)
    parser.add_argument("--allow-file", action="append", dest="allow_files")
    parser.add_argument("--chunk-size", default="64mb")
    parser.add_argument("--retries", type=int, default=20)
    parser.add_argument("--retry-delay", type=float, default=2.0)
    parser.add_argument("--timeout", type=int, default=120)
    args = parser.parse_args()

    card = load_model_card(args.card_id)
    repo_id = card.get("upstream_id")
    if not repo_id:
        print(f"{args.card_id}: upstream_id is required", file=sys.stderr)
        return 1

    files = selected_files(repo_id, args.allow_files)
    if not files:
        print("No files selected.", file=sys.stderr)
        return 1

    token = os.environ.get("HF_TOKEN") or None
    local_dir = Path(args.local_dir)
    chunk_size = parse_size(args.chunk_size)
    failed = []

    for filename, expected_size in files:
        ok = download_file(
            repo_id,
            filename,
            local_dir / filename,
            expected_size,
            chunk_size=chunk_size,
            retries=args.retries,
            retry_delay=args.retry_delay,
            timeout=args.timeout,
            token=token,
        )
        if not ok:
            current_size = local_size(local_dir / filename)
            failed.append((filename, current_size, expected_size))

    if failed:
        print("Incomplete downloads:", file=sys.stderr)
        for filename, current_size, expected_size in failed:
            print(f"- {filename}: {current_size}/{expected_size}", file=sys.stderr)
        return 1

    print(f"Downloaded {len(files)} file(s) for {args.card_id}.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())