Buckets:
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import re | |
| import sys | |
| from datetime import datetime, timezone | |
| from pathlib import Path, PurePosixPath | |
| from huggingface_hub import HfApi | |
| EXCLUDED_DIRECTORIES = { | |
| ".git", | |
| ".hg", | |
| ".svn", | |
| "__pycache__", | |
| "node_modules", | |
| "checkpoints", | |
| "weights", | |
| "secrets", | |
| ".secrets", | |
| ".ssh", | |
| ".aws", | |
| ".cache", | |
| ".pytest_cache", | |
| ".ruff_cache", | |
| ".mypy_cache", | |
| ".tox", | |
| ".nox", | |
| } | |
| EXCLUDED_SUFFIXES = { | |
| ".pt", | |
| ".pth", | |
| ".ckpt", | |
| ".safetensors", | |
| ".pkl", | |
| ".pickle", | |
| ".onnx", | |
| ".bin", | |
| ".pem", | |
| ".key", | |
| ".p12", | |
| ".pfx", | |
| ".pyc", | |
| ".pyo", | |
| } | |
| RESERVED_NAMES = {"upload_status.json", "upload_manifest.json"} | |
| SECRET_PATTERN = re.compile(rb"hf_[A-Za-z0-9]{20,}|-----BEGIN [A-Z ]*PRIVATE KEY-----") | |
| def excluded(path, code): | |
| parts = path.parts | |
| if any( | |
| part in EXCLUDED_DIRECTORIES | |
| or part.startswith((".venv", "venv", ".env")) | |
| or (code and part in {"data", "outputs", "logs", "external", "build", "dist"}) | |
| for part in parts | |
| ): | |
| return True | |
| name = path.name.lower() | |
| return ( | |
| any( | |
| name == reserved or name.startswith(reserved + ".") | |
| for reserved in RESERVED_NAMES | |
| ) | |
| or path.suffix.lower() in EXCLUDED_SUFFIXES | |
| or any(word in name for word in ("token", "credential", "secret", "checkpoint")) | |
| or name in {"id_rsa", "id_ed25519", ".netrc", ".npmrc", ".pypirc"} | |
| or bool(re.search(r"hf_[A-Za-z0-9]{20,}", str(path))) | |
| ) | |
| def inspect_file(path, sha256): | |
| digest = hashlib.sha256() if sha256 else None | |
| before = path.stat() | |
| previous = b"" | |
| with path.open("rb") as stream: | |
| while chunk := stream.read(1024 * 1024): | |
| if SECRET_PATTERN.search(previous + chunk): | |
| return None | |
| previous = chunk[-256:] | |
| if digest is not None: | |
| digest.update(chunk) | |
| after = path.stat() | |
| if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns): | |
| raise RuntimeError("An input file changed while preparing the upload") | |
| record = {"size": after.st_size, "mtime_ns": after.st_mtime_ns} | |
| if digest is not None: | |
| record["sha256"] = digest.hexdigest() | |
| return record | |
| def collect(folder, destination, code, sha256, token_file): | |
| records = [] | |
| skipped = 0 | |
| for root, directories, names in os.walk(folder, followlinks=False): | |
| parent = Path(root) | |
| directories[:] = sorted( | |
| name | |
| for name in directories | |
| if not (parent / name).is_symlink() | |
| and not excluded((parent / name).relative_to(folder), code) | |
| ) | |
| for name in sorted(names): | |
| path = parent / name | |
| relative = path.relative_to(folder) | |
| if ( | |
| path.is_symlink() | |
| or not path.is_file() | |
| or excluded(relative, code) | |
| or (token_file is not None and path.resolve() == token_file) | |
| ): | |
| skipped += 1 | |
| continue | |
| record = inspect_file(path, sha256) | |
| if record is None: | |
| skipped += 1 | |
| continue | |
| record.update( | |
| path=str(destination / PurePosixPath(relative.as_posix())), | |
| local=path, | |
| ) | |
| records.append(record) | |
| return records, skipped | |
| def write_status(path, status): | |
| temporary = path.with_name(path.name + ".tmp") | |
| temporary.write_text(json.dumps(status, indent=2) + "\n") | |
| temporary.replace(path) | |
| def upload(args, token): | |
| api = HfApi(token=token) | |
| if not api.bucket_info(args.bucket).private: | |
| raise RuntimeError("The destination bucket must be private") | |
| destination = PurePosixPath(args.prefix) | |
| token_file = args.token_file.resolve() if args.token_file is not None else None | |
| records, skipped = collect(args.folder, destination, False, args.sha256, token_file) | |
| if args.code is not None: | |
| code, code_skipped = collect( | |
| args.code, destination / "code", True, args.sha256, token_file | |
| ) | |
| records.extend(code) | |
| skipped += code_skipped | |
| paths = [record["path"] for record in records] | |
| if len(paths) != len(set(paths)): | |
| raise ValueError("Result files and the code snapshot have conflicting paths") | |
| if not records: | |
| raise ValueError("No eligible files to upload") | |
| status_path = args.folder / "upload_status.json" | |
| status = { | |
| "bucket": args.bucket, | |
| "prefix": args.prefix, | |
| "started_at": datetime.now(timezone.utc).isoformat(), | |
| "status": "uploading", | |
| "files": len(records), | |
| "bytes": sum(record["size"] for record in records), | |
| "skipped_files": skipped, | |
| "sha256": args.sha256, | |
| } | |
| write_status(status_path, status) | |
| try: | |
| for start in range(0, len(records), 128): | |
| batch = records[start : start + 128] | |
| api.batch_bucket_files( | |
| args.bucket, add=[(record["local"], record["path"]) for record in batch] | |
| ) | |
| print( | |
| json.dumps( | |
| {"uploaded_files": start + len(batch), "files": len(records)} | |
| ), | |
| flush=True, | |
| ) | |
| remote = { | |
| item.path: item.size | |
| for item in api.list_bucket_tree( | |
| args.bucket, prefix=args.prefix or None, recursive=True | |
| ) | |
| if hasattr(item, "size") | |
| } | |
| for record in records: | |
| current = record["local"].stat() | |
| if (current.st_size, current.st_mtime_ns) != ( | |
| record["size"], | |
| record["mtime_ns"], | |
| ): | |
| raise RuntimeError( | |
| "An input file changed during the upload; run the uploader again" | |
| ) | |
| if remote.get(record["path"]) != record["size"]: | |
| raise RuntimeError("Remote file size verification failed") | |
| manifest = { | |
| **status, | |
| "status": "verified_sizes", | |
| "completed_at": datetime.now(timezone.utc).isoformat(), | |
| "entries": [ | |
| {key: value for key, value in record.items() if key != "local"} | |
| for record in records | |
| ], | |
| } | |
| payload = (json.dumps(manifest, indent=2) + "\n").encode() | |
| manifest_path = str(destination / "upload_manifest.json") | |
| api.batch_bucket_files(args.bucket, add=[(payload, manifest_path)]) | |
| if not any( | |
| item.path == manifest_path and getattr(item, "size", None) == len(payload) | |
| for item in api.list_bucket_tree( | |
| args.bucket, prefix=manifest_path, recursive=True | |
| ) | |
| ): | |
| raise RuntimeError("Remote manifest size verification failed") | |
| status.update( | |
| status="verified_sizes", | |
| completed_at=manifest["completed_at"], | |
| manifest=manifest_path, | |
| ) | |
| write_status(status_path, status) | |
| print(json.dumps(status), flush=True) | |
| except Exception: | |
| status.update(status="failed", failed_at=datetime.now(timezone.utc).isoformat()) | |
| write_status(status_path, status) | |
| raise | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--folder", type=Path, required=True) | |
| parser.add_argument("--prefix", required=True) | |
| parser.add_argument("--bucket", default="twanghcmut/ego-real2sim") | |
| parser.add_argument("--token-file", type=Path) | |
| parser.add_argument( | |
| "--code", nargs="?", type=Path, const=Path(__file__).resolve().parents[1] | |
| ) | |
| parser.add_argument("--sha256", action="store_true") | |
| args = parser.parse_args() | |
| prefix = args.prefix.replace("\\", "/") | |
| if ( | |
| prefix.startswith("/") | |
| or ".." in prefix.split("/") | |
| or re.search(r"hf_[A-Za-z0-9]{20,}", prefix) | |
| ): | |
| parser.error( | |
| "prefix must be a relative bucket path without credentials or parent traversal" | |
| ) | |
| args.prefix = str(PurePosixPath(prefix)) if prefix else "" | |
| args.folder = args.folder.resolve() | |
| if not args.folder.is_dir(): | |
| parser.error("folder must be an existing directory") | |
| if args.code is not None: | |
| args.code = args.code.resolve() | |
| if not args.code.is_dir(): | |
| parser.error("code must be an existing directory") | |
| token = "" | |
| try: | |
| token = ( | |
| args.token_file.read_text().strip() | |
| if args.token_file | |
| else os.environ.get("HF_TOKEN", "").strip() | |
| ) | |
| if not token: | |
| raise ValueError("Provide --token-file or HF_TOKEN") | |
| upload(args, token) | |
| except (OSError, ValueError, RuntimeError) as error: | |
| message = str(error) | |
| if token: | |
| message = message.replace(token, "[redacted]") | |
| message = re.sub(r"hf_[A-Za-z0-9]+", "[redacted]", message) | |
| print(f"Upload failed: {message}", file=sys.stderr) | |
| return 1 | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |
Xet Storage Details
- Size:
- 9.19 kB
- Xet hash:
- 36650d83890be41cb51746af6232a194c9168ad43b7d99356b65de37be05323a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.