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"