File size: 3,528 Bytes
f8b48da | 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 110 111 112 113 114 115 | """Back up AskMyDocs state to a Hugging Face Storage Bucket.
Creates a compressed archive of PostgreSQL data, Qdrant storage, and uploaded
documents, then uploads it to the specified HF Storage Bucket. The Space's
``start.sh`` restores this archive on boot when ``RESTORE_ON_BOOT=true``.
Usage:
python scripts/hf_state_backup.py --bucket username/askmydocs-data
python scripts/hf_state_backup.py --bucket username/askmydocs-data --data-root /home/user
"""
from __future__ import annotations
import argparse
import subprocess
import tempfile
from pathlib import Path
def create_archive(data_root: Path, output: Path) -> bool:
"""Create a zstd-compressed tar of the state directories."""
dirs_to_include = ["pgdata", "qdrant_storage"]
uploads_dir = data_root / "app" / "data" / "uploads"
if uploads_dir.exists():
dirs_to_include.append("app/data/uploads")
existing = [d for d in dirs_to_include if (data_root / d).exists()]
if not existing:
print("No state directories found to back up.")
return False
cmd = [
"tar",
"--zstd",
"-cf",
str(output),
"-C",
str(data_root),
*existing,
]
print(f"Creating archive: {' '.join(cmd)}")
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
print(f"tar failed: {result.stderr}")
return False
size_mb = output.stat().st_size / (1024 * 1024)
print(f"Archive created: {output} ({size_mb:.1f} MB)")
return True
def upload_to_bucket(archive: Path, bucket: str, token: str | None = None) -> bool:
"""Upload the archive to an HF Storage Bucket.
Storage Buckets are not repos: ``upload_file(repo_type="space_storage")``
does not exist (supported repo types are model/dataset/space). The bucket
API is ``batch_bucket_files(bucket_id=..., add=[(local, remote)])``.
"""
try:
from huggingface_hub import batch_bucket_files
except ImportError:
print("Error: huggingface_hub is required. Install with: pip install huggingface_hub")
return False
print(f"Uploading {archive.name} to bucket {bucket}...")
batch_bucket_files(
bucket_id=bucket,
add=[(archive, "state.tar.zst")],
token=token,
)
print(f"Uploaded to hf://buckets/{bucket}/state.tar.zst")
return True
def main() -> int:
parser = argparse.ArgumentParser(description="Back up AskMyDocs state to HF Storage Bucket.")
parser.add_argument(
"--bucket",
required=True,
help="HF Storage Bucket name (e.g. 'username/askmydocs-data')",
)
parser.add_argument(
"--data-root",
type=Path,
default=Path("/home/user"),
help="Data root directory (default: /home/user for HF Space)",
)
parser.add_argument(
"--token",
default=None,
help="HF token for private bucket access (default: HF_TOKEN env var)",
)
args = parser.parse_args()
import os
token = args.token or os.environ.get("HF_TOKEN")
with tempfile.NamedTemporaryFile(suffix=".tar.zst", delete=False) as tmp:
tmp_path = Path(tmp.name)
try:
if not create_archive(args.data_root, tmp_path):
return 1
if not upload_to_bucket(tmp_path, args.bucket, token):
return 1
finally:
tmp_path.unlink(missing_ok=True)
print("State backup complete.")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|