Spaces:

SZFSPC
/

apple

Paused

apple / scripts /_hf_upload.py

feat: deploy to hf space | model=mistral-21461b

2bfa191 verified 5 months ago

11.2 kB

"""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)

	"""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)