Spaces:
Paused
Paused
File size: 3,726 Bytes
8523e75 | 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 | #!/usr/bin/env python3
"""Restore data from HF Dataset snapshot before server starts.
Handles three cases:
1. PAPERCLIP_DATA_REPO_ID or HF_TOKEN not set -> skip (exit 0)
2. Dataset repo does not exist (first deploy) -> skip gracefully (exit 0)
3. Restore succeeds -> files copied to PAPERCLIP_HOME (exit 0)
4. Unexpected error -> FATAL, exit 1 (refuse to start)
"""
import os
import shutil
import sys
import tempfile
from huggingface_hub import snapshot_download
from huggingface_hub.errors import RepositoryNotFoundError
def _get_env(name, default=""):
"""Read env var at call time (not import time) so tests can monkeypatch."""
return os.environ.get(name, default)
def _get_instance_name():
"""Derive instance name from SPACE_ID or fallback env var.
SPACE_ID on HF Spaces is like 'user/space-name'.
We sanitize '/' to '-' for use as a subdirectory name.
"""
name = (
_get_env("SPACE_ID")
or _get_env("PAPERCLIP_INSTANCE_NAME")
or _get_env("PAPERCLIP_INSTANCE_ID", "default")
)
return name.replace("/", "-")
def _sync_restore(src_dir: str, dst_dir: str) -> None:
"""Restore snapshot using atomic file replacement.
For files: write to a temp path next to the target, then os.replace()
(atomic on same filesystem). If the process crashes mid-restore,
old files remain intact and temp files are harmless.
For directories: recurse depth-first so leaf files are replaced atomically.
Local-only entries (not in snapshot) are left untouched.
"""
for entry in os.listdir(src_dir):
src_path = os.path.join(src_dir, entry)
dst_path = os.path.join(dst_dir, entry)
if os.path.isdir(src_path):
os.makedirs(dst_path, exist_ok=True)
_sync_restore(src_path, dst_path)
else:
os.makedirs(dst_dir, exist_ok=True)
tmp_path = dst_path + ".restoring"
shutil.copy2(src_path, tmp_path)
os.replace(tmp_path, dst_path)
def main():
paperclip_home = _get_env("PAPERCLIP_HOME", "/paperclip")
data_repo_id = _get_env("PAPERCLIP_DATA_REPO_ID")
hf_token = _get_env("HF_TOKEN") or _get_env("HUGGING_FACE_HUB_TOKEN")
if not data_repo_id:
print("restore_snapshot: PAPERCLIP_DATA_REPO_ID not set, skipping restore")
return
if not hf_token:
print("restore_snapshot: HF_TOKEN not set, skipping restore")
return
instance_name = _get_instance_name()
print(f"restore_snapshot: restoring {instance_name} from {data_repo_id}")
try:
with tempfile.TemporaryDirectory() as tmp_dir:
snapshot_download(
repo_id=data_repo_id,
repo_type="dataset",
local_dir=tmp_dir,
token=hf_token,
allow_patterns=[f"{instance_name}/**"],
)
# snapshot_download preserves subdirectory structure, so
# files land at tmp_dir/{instance_name}/...
instance_dir = os.path.join(tmp_dir, instance_name)
if os.path.isdir(instance_dir):
_sync_restore(instance_dir, paperclip_home)
print(f"restore_snapshot: restored data to {paperclip_home}")
else:
print(
f"restore_snapshot: no data for instance {instance_name}, skipping"
)
except RepositoryNotFoundError:
print(
f"restore_snapshot: Dataset {data_repo_id} not found "
"(first deploy?), skipping restore"
)
except Exception as e:
print(f"FATAL: restore_snapshot failed: {e}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()
|