bucket_ops.py

import os
from pathlib import Path
import importlib
from utils import retry_call


BUCKET_CREATE_PRIVATE_DEFAULT = True


def _get_bucket_api_module():
    hh = importlib.import_module("huggingface_hub")
    required = ["create_bucket", "list_bucket_tree", "batch_bucket_files"]
    missing = [name for name in required if not hasattr(hh, name)]
    if missing:
        raise RuntimeError(f"Bucket API is unavailable in current huggingface_hub build. Missing: {', '.join(missing)}")
    return hh


def is_bucket_api_available():
    try:
        _get_bucket_api_module()
        return True
    except Exception:
        return False


def get_bucket_url(bucket_id: str):
    return f"https://huggingface.co/buckets/{bucket_id}"


def get_bucket_handle(bucket_id: str, remote_path: str = ""):
    bucket_id = str(bucket_id).strip().strip("/")
    remote_path = str(remote_path).strip().strip("/")
    return f"hf://buckets/{bucket_id}/{remote_path}" if remote_path else f"hf://buckets/{bucket_id}"


def ensure_bucket(bucket_id: str, hf_token: str, private: bool = BUCKET_CREATE_PRIVATE_DEFAULT):
    hh = _get_bucket_api_module()
    create_bucket = getattr(hh, "create_bucket")
    bucket_info = getattr(hh, "bucket_info")
    retry_call(
        lambda: create_bucket(bucket_id, private=private, exist_ok=True, token=hf_token),
        action=f"create_bucket {bucket_id}",
    )
    return retry_call(
        lambda: bucket_info(bucket_id, token=hf_token),
        action=f"bucket_info {bucket_id}",
    )


def get_bucket_file_info(bucket_id: str, remote_path: str, hf_token: str):
    hh = _get_bucket_api_module()
    list_bucket_tree = getattr(hh, "list_bucket_tree")

    def _list():
        return list(list_bucket_tree(bucket_id, prefix=remote_path, recursive=True, token=hf_token))

    items = retry_call(_list, action=f"list_bucket_tree {bucket_id}:{remote_path}")
    for item in items:
        if getattr(item, "type", "") == "file" and getattr(item, "path", "") == remote_path:
            return item
    return None


def bucket_file_exists(bucket_id: str, remote_path: str, hf_token: str):
    return get_bucket_file_info(bucket_id, remote_path, hf_token) is not None


def get_safe_bucket_filename(filename: str, bucket_id: str, hf_token: str):
    ensure_bucket(bucket_id=bucket_id, hf_token=hf_token, private=BUCKET_CREATE_PRIVATE_DEFAULT)
    path = Path(filename)
    remote_name = path.name
    existing = get_bucket_file_info(bucket_id, remote_name, hf_token)
    if existing is None:
        return filename
    try:
        remote_size = getattr(existing, "size", None)
        local_size = os.path.getsize(filename)
        if remote_size is not None and int(remote_size) == int(local_size):
            print(f"{remote_name} already exists in bucket with same size. keeping original name for skip check.")
            return filename
    except Exception:
        pass
    i = 1
    while True:
        candidate = str(Path(path.parent, f"{path.stem}_{i}{path.suffix}"))
        if get_bucket_file_info(bucket_id, Path(candidate).name, hf_token) is None:
            print(f"{path.name} is already exists in bucket but file size is different. renaming to {Path(candidate).name}.")
            Path(filename).rename(candidate)
            return candidate
        i += 1


def upload_file_to_bucket(local_path: str, bucket_id: str, hf_token: str, remote_path: str = "", private: bool = BUCKET_CREATE_PRIVATE_DEFAULT):
    hh = _get_bucket_api_module()
    batch_bucket_files = getattr(hh, "batch_bucket_files")
    remote_path = remote_path or Path(local_path).name
    ensure_bucket(bucket_id=bucket_id, hf_token=hf_token, private=private)
    if bucket_file_exists(bucket_id, remote_path, hf_token):
        print(f"{remote_path} already exists in bucket. skipping.")
        return get_bucket_handle(bucket_id, remote_path), "skipped"
    retry_call(
        lambda: batch_bucket_files(bucket_id, add=[(str(local_path), remote_path)], token=hf_token),
        action=f"batch_bucket_files {bucket_id}:{remote_path}",
    )
    if not bucket_file_exists(bucket_id, remote_path, hf_token):
        raise RuntimeError(f"Bucket verify failed for {bucket_id}:{remote_path}")
    return get_bucket_handle(bucket_id, remote_path), "uploaded"