Blame for scripts/_hf_upload.py

"""HF Upload helper - 计算上传大小 + 使用 upload_large_folder 上传大目录

用法:
    python scripts/_hf_upload.py calc-size <dir> [exclude_patterns...]
    python scripts/_hf_upload.py upload <repo_id> <dir> --token TOKEN [exclude_patterns...]

示例:
    python scripts/_hf_upload.py calc-size . '.git/**' '.venv/**'
    python scripts/_hf_upload.py upload user/space . --token hf_xxx '.git/**' '.venv/**'
"""

from __future__ import annotations

import argparse
import fnmatch
import os
import shutil
import subprocess
import sys
from pathlib import Path

_HF_EXTRA_PATH = "/root/.local/bin"


def _human_size(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if abs(n) < 1024:
            return f"{n:.2f} {unit}"
        n /= 1024
    return f"{n:.2f} PB"


def _should_exclude(rel_path: str, patterns: list[str]) -> bool:
    """判断文件路径是否匹配任意排除模式。"""
    for p in patterns:
        if fnmatch.fnmatch(rel_path, p):
            return True
        parts = rel_path.split(os.sep)
        for i in range(len(parts)):
            prefix = os.sep.join(parts[:i + 1])
            if fnmatch.fnmatch(prefix, p):
                return True
        if p.endswith("/**"):
            dir_p = p[:-3]
            if fnmatch.fnmatch(rel_path, dir_p) or rel_path.startswith(dir_p):
                return True
    return False


def calc_size(dir_path: str, exclude_patterns: list[str]) -> int:
    """计算目录中需要上传的文件总大小（字节），排除匹配的文件。"""
    total = 0
    dir_path = os.path.abspath(dir_path)
    for root, dirs, files in os.walk(dir_path):
        dirs[:] = [d for d in dirs if not _should_exclude(os.path.join(root, d), exclude_patterns)]
        for fname in files:
            full = os.path.join(root, fname)
            rel = os.path.relpath(full, dir_path)
            if _should_exclude(rel, exclude_patterns):
                continue
            try:
                total += os.path.getsize(full)
            except OSError:
                pass
    return total


def _hf_upload_with_progress(
    repo_id: str,
    local_dir: str,
    token: str | None = None,
    ignore_patterns: list[str] | None = None,
    commit_message: str = "Upload",
    repo_type: str = "space",
) -> bool:
    """使用 hf upload CLI 上传（有 tqdm 进度条）。"""
    if ignore_patterns is None:
        ignore_patterns = []

    hf_path = shutil.which("hf") or shutil.which(os.path.join(_HF_EXTRA_PATH, "hf"))
    if not hf_path:
        print("[ERROR] 'hf' CLI not found in PATH", file=sys.stderr)
        return False

    extra_env = None
    if token:
        extra_env = os.environ.copy()
        extra_env["HF_TOKEN"] = token
    if _HF_EXTRA_PATH not in os.environ.get("PATH", ""):
        extra_env = (extra_env or os.environ.copy())
        extra_env["PATH"] = f"{_HF_EXTRA_PATH}:{extra_env['PATH']}"

    cmd = ["hf", "upload", repo_id, local_dir, "--repo-type", repo_type, "--commit-message", commit_message]
    for p in ignore_patterns:
        cmd += ["--exclude", p]

    print("[INFO] Falling back to hf upload CLI (tqdm progress will be shown)...")
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=7200,
            env=extra_env,
        )
        if result.returncode != 0:
            print(f"[ERROR] hf upload failed: {result.stderr.strip()}", file=sys.stderr)
            return False
        return True
    except subprocess.TimeoutExpired:
        print("[ERROR] hf upload timed out (>2h)", file=sys.stderr)
        return False
    except FileNotFoundError:
        print("[ERROR] 'hf' CLI not found", file=sys.stderr)
        return False


def _ensure_repo_exists(
    api,
    repo_id: str,
    repo_type: str,
    private: bool = True,
) -> None:
    """确保目标仓库存在，对于 space 类型需要提供 space_sdk 参数。"""
    import inspect

    # 收集 create_repo 支持的参数（不同版本有差异）
    try:
        sig = inspect.signature(api.create_repo)
        supported_params = set(sig.parameters.keys())
    except (TypeError, ValueError):
        supported_params = {"repo_id", "repo_type", "private", "exist_ok", "space_sdk"}

    create_kwargs = {
        "repo_id": repo_id,
        "repo_type": repo_type,
        "private": private,
        "exist_ok": True,
    }

    # space 类型必须提供 space_sdk（docker / gradio / streamlit / static）
    if repo_type == "space" and "space_sdk" in supported_params:
        create_kwargs["space_sdk"] = "docker"

    # 只传递支持的参数
    create_kwargs = {k: v for k, v in create_kwargs.items() if k in supported_params}

    try:
        api.create_repo(**create_kwargs)
    except Exception as e:
        # 仓库已存在或其他非致命错误：静默继续，让后续 upload 去处理
        err_lower = str(e).lower()
        if "already" not in err_lower and "exists" not in err_lower:
            print(
                f"[DEBUG] create_repo({repo_id}, type={repo_type}) 可能失败但将继续: {e}",
                file=sys.stderr,
            )


def _get_upload_large_folder_kwargs(
    api,
    repo_id: str,
    local_dir: str,
    ignore_patterns: list[str],
    commit_message: str,
    repo_type: str = "space",
) -> dict:
    """根据当前 huggingface_hub 版本的 API 签名，构建 upload_large_folder 的 kwargs。

    不同版本的 HfApi.upload_large_folder 支持的参数不同：
    - 旧版可能只有 repo_id / repo_type / folder_path / ignore_patterns
    - 新版新增了 commit_message / print_report / print_report_every
    - 某些版本可能使用 allow_patterns / delete_patterns 而不是 ignore_patterns
    使用 inspect.signature 动态检测，只传递支持的参数，避免 TypeError。
    """
    import inspect

    try:
        sig = inspect.signature(api.upload_large_folder)
        supported_params = set(sig.parameters.keys())
    except (TypeError, ValueError):
        # 某些旧版本签名无法 inspect，回退到最安全的参数子集
        supported_params = {"repo_id", "repo_type", "folder_path", "ignore_patterns"}

    all_candidates = {
        "repo_id": repo_id,
        "repo_type": repo_type,
        "folder_path": local_dir,
        "ignore_patterns": ignore_patterns,
        "commit_message": commit_message,
        "print_report": True,
        "print_report_every": 20,
    }

    # 如果 upload_large_folder 内部会调用 create_repo，也需要 space_sdk
    if repo_type == "space" and "space_sdk" in supported_params:
        all_candidates["space_sdk"] = "docker"

    kwargs = {k: v for k, v in all_candidates.items() if k in supported_params}

    # 如果 ignore_patterns 不被支持，看看是否有 allow_patterns / delete_patterns
    if "ignore_patterns" not in supported_params and ignore_patterns:
        for alt in ("allow_patterns",):
            if alt in supported_params:
                # allow_patterns 语义不同（白名单 vs 黑名单），这里不做自动转换
                break

    unsupported = [k for k in all_candidates if k not in supported_params]
    if unsupported:
        # 只在首次调用时打印一次，且级别降为 DEBUG（非 stderr 主输出）
        print(
            f"[DEBUG] huggingface_hub upload_large_folder 不支持参数: "
            f"{', '.join(unsupported)} — 将跳过（不影响功能）",
            file=sys.stderr,
        )

    return kwargs


def upload_large_folder(
    repo_id: str,
    local_dir: str,
    token: str | None = None,
    ignore_patterns: list[str] | None = None,
    commit_message: str = "Upload",
    repo_type: str = "space",
) -> bool:
    """使用 HfApi.upload_large_folder 上传目录（有 print_report 进度）。

    参数名使用 ignore_patterns（HfApi 的正式参数名）。
    如果 huggingface_hub 不可用或上传失败，回退到 hf upload CLI（有 tqdm 进度条）。
    为避免不同 huggingface_hub 版本的 API 签名差异（例如某些版本没有 commit_message），
    使用 inspect.signature 动态检测并只传递支持的参数。

    对于 repo_type="space"，会确保先调用 create_repo 并提供 space_sdk="docker"，
    以避免 "No space_sdk provided" 错误。
    """
    if ignore_patterns is None:
        ignore_patterns = []

    try:
        from huggingface_hub import HfApi
    except ModuleNotFoundError:
        print("[WARN] huggingface_hub not found, falling back to hf upload CLI", file=sys.stderr)
        return _hf_upload_with_progress(repo_id, local_dir, token, ignore_patterns, commit_message, repo_type)

    api = HfApi(token=token)

    # 先确保仓库存在（space 类型需要 space_sdk，必须显式提供）
    _ensure_repo_exists(api, repo_id, repo_type=repo_type)

    kwargs = _get_upload_large_folder_kwargs(
        api, repo_id, local_dir, ignore_patterns, commit_message, repo_type=repo_type
    )

    try:
        api.upload_large_folder(**kwargs)
        return True
    except TypeError as e:
        # 如果动态检测仍未覆盖到，例如方法签名有 **kwargs 但内部再次调用时出错，
        # 给出明确提示并回退到 CLI
        print(f"[WARN] HfApi.upload_large_folder 参数签名不兼容: {e}", file=sys.stderr)
        print("[INFO] 回退到 hf upload CLI ...", file=sys.stderr)
        return _hf_upload_with_progress(repo_id, local_dir, token, ignore_patterns, commit_message, repo_type)
    except Exception as e:
        print(f"[WARN] HfApi.upload_large_folder failed: {e}", file=sys.stderr)
        return _hf_upload_with_progress(repo_id, local_dir, token, ignore_patterns, commit_message, repo_type)


# ---- CLI ----
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="HF Upload helper")
    sub = parser.add_subparsers(dest="cmd", required=True)

    calc = sub.add_parser("calc-size", help="Calculate upload size")
    calc.add_argument("dir", help="Directory to scan")
    calc.add_argument("patterns", nargs="*", help="Exclude patterns")

    up = sub.add_parser("upload", help="Upload folder using upload_large_folder")
    up.add_argument("repo_id", help="Target repo ID")
    up.add_argument("dir", help="Directory to upload")
    up.add_argument("--token", help="HF token")
    up.add_argument("--message", default="Upload", help="Commit message")
    up.add_argument("--repo-type", default="space", choices=["space", "dataset", "model"],
                    help="Repository type (default: space)")
    up.add_argument("patterns", nargs="*", help="Exclude patterns (passed as ignore_patterns)")

    args = parser.parse_args()

    if args.cmd == "calc-size":
        size = calc_size(args.dir, args.patterns)
        print(_human_size(size))
        sys.exit(0)

    if args.cmd == "upload":
        ok = upload_large_folder(
            repo_id=args.repo_id,
            local_dir=args.dir,
            token=args.token,
            ignore_patterns=args.patterns,
            commit_message=args.message,
            repo_type=args.repo_type,
        )
        sys.exit(0 if ok else 1)