ckriti/HuggingClaw

1#!/usr/bin/env python32 3import os4import sys5import json6import hashlib7import time8import tarfile9import tempfile10import shutil11from datetime import datetime12from pathlib import Path13from typing import Dict, Any, Optional, List14import requests15import logging16 17from huggingface_hub import HfApi18from huggingface_hub.utils import RepositoryNotFoundError19from huggingface_hub import hf_hub_download20 21logging.basicConfig(22    level=logging.INFO,23    format='{"timestamp": "%(asctime)s", "level": "%(levelname)s", "module": "atomic-restore", "message": "%(message)s"}'24)25logger = logging.getLogger(__name__)26 27class AtomicDatasetRestorer:28    29    def __init__(self, repo_id: str, dataset_path: str = "state"):30        self.repo_id = repo_id31        self.dataset_path = Path(dataset_path)32        self.api = HfApi()33        self.max_retries = 334        self.base_delay = 1.035        36        logger.info("init", {37            "repo_id": repo_id,38            "dataset_path": dataset_path,39            "max_retries": self.max_retries40        })41    42    def calculate_checksum(self, file_path: Path) -> str:43        sha256_hash = hashlib.sha256()44        with open(file_path, "rb") as f:45            for chunk in iter(lambda: f.read(4096), b""):46                sha256_hash.update(chunk)47        return sha256_hash.hexdigest()48    49    def validate_integrity(self, metadata: Dict[str, Any], state_files: List[Path]) -> bool:50        """Validate data integrity using checksums"""51        try:52            if "checksum" not in metadata:53                logger.warning("no_checksum_in_metadata", {"action": "skipping_validation"})54                return True55            56            state_data = metadata.get("state_data", {})57            calculated_checksum = hashlib.sha256(58                json.dumps(state_data, sort_keys=True).encode()59            ).hexdigest()60            61            expected_checksum = metadata["checksum"]62            63            is_valid = calculated_checksum == expected_checksum64            65            logger.info("integrity_check", {66                "expected": expected_checksum,67                "calculated": calculated_checksum,68                "valid": is_valid69            })70            71            return is_valid72            73        except Exception as e:74            logger.error("integrity_validation_failed", {"error": str(e)})75            return False76    77    def create_backup_before_restore(self, target_dir: Path) -> Optional[Path]:78        try:79            if not target_dir.exists():80                return None81            82            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")83            backup_dir = target_dir.parent / f"state_backup_{timestamp}"84            85            logger.info("creating_local_backup", {86                "source": str(target_dir),87                "backup": str(backup_dir)88            })89            90            shutil.copytree(target_dir, backup_dir)91            return backup_dir92            93        except Exception as e:94            logger.error("local_backup_failed", {"error": str(e)})95            return None96    97    def restore_from_commit(self, commit_sha: str, target_dir: Path, force: bool = False) -> Dict[str, Any]:98        """99        Restore state from specific commit100        101        Args:102            commit_sha: Git commit hash to restore from103            target_dir: Directory to restore state to104            force: Force restore without confirmation105            106        Returns:107            Dictionary with operation result108        """109        operation_id = f"restore_{int(time.time())}"110        111        logger.info("starting_atomic_restore", {112            "operation_id": operation_id,113            "commit_sha": commit_sha,114            "target_dir": str(target_dir),115            "force": force116        })117        118        try:119            # Validate commit exists120            try:121                repo_info = self.api.repo_info(122                    repo_id=self.repo_id,123                    repo_type="dataset",124                    revision=commit_sha125                )126                logger.info("commit_validated", {"commit": commit_sha})127            except Exception as e:128                error_result = {129                    "success": False,130                    "operation_id": operation_id,131                    "error": f"Invalid commit: {str(e)}",132                    "timestamp": datetime.now().isoformat()133                }134                logger.error("commit_validation_failed", error_result)135                return error_result136            137            # Create backup before restore138            backup_dir = self.create_backup_before_restore(target_dir)139            140            # Create temporary directory for download141            with tempfile.TemporaryDirectory() as tmpdir:142                tmpdir_path = Path(tmpdir)143                144                # List files in the commit145                files = self.api.list_repo_files(146                    repo_id=self.repo_id,147                    repo_type="dataset",148                    revision=commit_sha149                )150                151                # Find state files152                state_files = [f for f in files if f.startswith(str(self.dataset_path))]153                if not state_files:154                    error_result = {155                        "success": False,156                        "operation_id": operation_id,157                        "error": "No state files found in commit",158                        "timestamp": datetime.now().isoformat()159                    }160                    logger.error("no_state_files", error_result)161                    return error_result162                163                # Download state files164                downloaded_files = []165                metadata = None166                167                for file_path in state_files:168                    try:169                        local_path = hf_hub_download(170                            repo_id=self.repo_id,171                            repo_type="dataset",172                            filename=file_path,173                            revision=commit_sha,174                            local_files_only=False175                        )176                        177                        if local_path:178                            downloaded_files.append(Path(local_path))179                            180                            # Load metadata if this is metadata.json181                            if file_path.endswith("metadata.json"):182                                with open(local_path, "r") as f:183                                    metadata = json.load(f)184                                    185                    except Exception as e:186                        logger.error("file_download_failed", {"file": file_path, "error": str(e)})187                        continue188                189                if not metadata:190                    error_result = {191                        "success": False,192                        "operation_id": operation_id,193                        "error": "Metadata not found in state files",194                        "timestamp": datetime.now().isoformat()195                    }196                    logger.error("metadata_not_found", error_result)197                    return error_result198                199                # Validate data integrity200                if not self.validate_integrity(metadata, downloaded_files):201                    error_result = {202                        "success": False,203                        "operation_id": operation_id,204                        "error": "Data integrity validation failed",205                        "timestamp": datetime.now().isoformat()206                    }207                    logger.error("integrity_validation_failed", error_result)208                    return error_result209                210                # Create target directory211                target_dir.mkdir(parents=True, exist_ok=True)212                213                # Restore files (except metadata.json which is for reference)214                restored_files = []215                for file_path in downloaded_files:216                    if file_path.name != "metadata.json":217                        dest_path = target_dir / file_path.name218                        shutil.copy2(file_path, dest_path)219                        restored_files.append(str(dest_path))220                        221                        logger.info("file_restored", {222                            "source": str(file_path),223                            "destination": str(dest_path)224                        })225                226                result = {227                    "success": True,228                    "operation_id": operation_id,229                    "commit_sha": commit_sha,230                    "backup_dir": str(backup_dir) if backup_dir else None,231                    "timestamp": datetime.now().isoformat(),232                    "restored_files": restored_files,233                    "metadata": metadata234                }235                236                logger.info("atomic_restore_completed", result)237                return result238                239        except Exception as e:240            error_result = {241                "success": False,242                "operation_id": operation_id,243                "error": str(e),244                "timestamp": datetime.now().isoformat()245            }246            247            logger.error("atomic_restore_failed", error_result)248            return error_result249    250    def restore_latest(self, target_dir: Path, force: bool = False) -> Dict[str, Any]:251        """Restore from the latest commit"""252        try:253            repo_info = self.api.repo_info(254                repo_id=self.repo_id,255                repo_type="dataset"256            )257            258            if not repo_info.sha:259                error_result = {260                    "success": False,261                    "error": "No commit found in repository",262                    "timestamp": datetime.now().isoformat()263                }264                logger.error("no_commit_found", error_result)265                return error_result266            267            return self.restore_from_commit(repo_info.sha, target_dir, force)268            269        except Exception as e:270            error_result = {271                "success": False,272                "error": f"Failed to get latest commit: {str(e)}",273                "timestamp": datetime.now().isoformat()274            }275            logger.error("latest_commit_failed", error_result)276            return error_result277 278def main():279    """Main function for command line usage"""280    if len(sys.argv) < 3:281        print(json.dumps({282            "error": "Usage: python restore_from_dataset_atomic.py <repo_id> <target_dir> [--force]",283            "status": "error"284        }, indent=2))285        sys.exit(1)286    287    repo_id = sys.argv[1]288    target_dir = sys.argv[2]289    force = "--force" in sys.argv290    291    try:292        target_path = Path(target_dir)293        restorer = AtomicDatasetRestorer(repo_id)294        result = restorer.restore_latest(target_path, force)295        296        print(json.dumps(result, indent=2))297        298        if not result.get("success", False):299            sys.exit(1)300            301    except Exception as e:302        print(json.dumps({303            "error": str(e),304            "status": "error"305        }, indent=2))306        sys.exit(1)307 308if __name__ == "__main__":309    main()