Spaces:
Runtime error
Runtime error
File size: 4,377 Bytes
5a38d4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | 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"
|