Spaces:
Paused
Paused
| #!/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() | |