rahmad7/hermes-openmodel

1#!/usr/bin/env python32"""3Atomic Dataset Restore for HermesFace4Restore state from HF Dataset with integrity validation and local backup.5 6Usage:7    python3 restore_from_dataset_atomic.py <repo_id> <target_dir> [--force]8"""9import hashlib10import json11import logging12import shutil13import sys14import tempfile15import time16from datetime import datetime17from pathlib import Path18from typing import Any, Dict, List, Optional19 20from huggingface_hub import HfApi, hf_hub_download21 22logging.basicConfig(23    level=logging.INFO,24    format='{"timestamp": "%(asctime)s", "level": "%(levelname)s", "module": "atomic-restore", "message": "%(message)s"}',25)26logger = logging.getLogger(__name__)27 28 29class AtomicDatasetRestorer:30    def __init__(self, repo_id: str, dataset_path: str = "state"):31        self.repo_id = repo_id32        self.dataset_path = Path(dataset_path)33        self.api = HfApi()34 35    def calculate_checksum(self, file_path: Path) -> str:36        h = hashlib.sha256()37        with open(file_path, "rb") as f:38            for chunk in iter(lambda: f.read(4096), b""):39                h.update(chunk)40        return h.hexdigest()41 42    def validate_integrity(self, metadata: Dict[str, Any], state_files: List[Path]) -> bool:43        try:44            if "checksum" not in metadata:45                logger.warning("no_checksum_in_metadata")46                return True47            calculated = hashlib.sha256(48                json.dumps(metadata.get("state_data", {}), sort_keys=True).encode()49            ).hexdigest()50            expected = metadata["checksum"]51            valid = calculated == expected52            logger.info(f"integrity_check expected={expected} calculated={calculated} valid={valid}")53            return valid54        except Exception as e:55            logger.error(f"integrity_validation_failed error={e}")56            return False57 58    def create_backup_before_restore(self, target_dir: Path) -> Optional[Path]:59        try:60            if not target_dir.exists():61                return None62            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")63            backup_dir = target_dir.parent / f"state_backup_{timestamp}"64            logger.info(f"creating_local_backup source={target_dir} backup={backup_dir}")65            shutil.copytree(target_dir, backup_dir)66            return backup_dir67        except Exception as e:68            logger.error(f"local_backup_failed error={e}")69            return None70 71    def restore_from_commit(72        self, commit_sha: str, target_dir: Path, force: bool = False73    ) -> Dict[str, Any]:74        operation_id = f"restore_{int(time.time())}"75        logger.info(f"starting_atomic_restore op={operation_id} commit={commit_sha}")76 77        try:78            self.api.repo_info(repo_id=self.repo_id, repo_type="dataset", revision=commit_sha)79        except Exception as e:80            return {"success": False, "operation_id": operation_id,81                    "error": f"Invalid commit: {e}", "timestamp": datetime.now().isoformat()}82 83        backup_dir = self.create_backup_before_restore(target_dir)84 85        try:86            with tempfile.TemporaryDirectory() as _tmpdir:87                files = self.api.list_repo_files(88                    repo_id=self.repo_id, repo_type="dataset", revision=commit_sha89                )90                state_files = [f for f in files if f.startswith(str(self.dataset_path))]91                if not state_files:92                    return {"success": False, "operation_id": operation_id,93                            "error": "No state files found in commit",94                            "timestamp": datetime.now().isoformat()}95 96                downloaded_files: List[Path] = []97                metadata = None98                for file_path in state_files:99                    try:100                        local = hf_hub_download(101                            repo_id=self.repo_id,102                            repo_type="dataset",103                            filename=file_path,104                            revision=commit_sha,105                        )106                        if local:107                            downloaded_files.append(Path(local))108                            if file_path.endswith("metadata.json"):109                                with open(local, "r") as f:110                                    metadata = json.load(f)111                    except Exception as e:112                        logger.error(f"file_download_failed file={file_path} error={e}")113 114                if not metadata:115                    return {"success": False, "operation_id": operation_id,116                            "error": "Metadata not found in state files",117                            "timestamp": datetime.now().isoformat()}118 119                if not self.validate_integrity(metadata, downloaded_files):120                    return {"success": False, "operation_id": operation_id,121                            "error": "Data integrity validation failed",122                            "timestamp": datetime.now().isoformat()}123 124                target_dir.mkdir(parents=True, exist_ok=True)125                restored = []126                for f in downloaded_files:127                    if f.name != "metadata.json":128                        dst = target_dir / f.name129                        shutil.copy2(f, dst)130                        restored.append(str(dst))131 132                result = {133                    "success": True,134                    "operation_id": operation_id,135                    "commit_sha": commit_sha,136                    "backup_dir": str(backup_dir) if backup_dir else None,137                    "timestamp": datetime.now().isoformat(),138                    "restored_files": restored,139                    "metadata": metadata,140                }141                logger.info(f"atomic_restore_completed {result}")142                return result143        except Exception as e:144            logger.error(f"atomic_restore_failed error={e}")145            return {"success": False, "operation_id": operation_id,146                    "error": str(e), "timestamp": datetime.now().isoformat()}147 148    def restore_latest(self, target_dir: Path, force: bool = False) -> Dict[str, Any]:149        try:150            repo_info = self.api.repo_info(repo_id=self.repo_id, repo_type="dataset")151            if not repo_info.sha:152                return {"success": False, "error": "No commit found in repository",153                        "timestamp": datetime.now().isoformat()}154            return self.restore_from_commit(repo_info.sha, target_dir, force)155        except Exception as e:156            return {"success": False, "error": f"Failed to get latest commit: {e}",157                    "timestamp": datetime.now().isoformat()}158 159 160def main() -> None:161    if len(sys.argv) < 3:162        print(json.dumps({163            "error": "Usage: python restore_from_dataset_atomic.py <repo_id> <target_dir> [--force]",164            "status": "error",165        }, indent=2))166        sys.exit(1)167 168    repo_id = sys.argv[1]169    target_dir = sys.argv[2]170    force = "--force" in sys.argv171 172    try:173        restorer = AtomicDatasetRestorer(repo_id)174        result = restorer.restore_latest(Path(target_dir), force)175        print(json.dumps(result, indent=2))176        if not result.get("success", False):177            sys.exit(1)178    except Exception as e:179        print(json.dumps({"error": str(e), "status": "error"}, indent=2))180        sys.exit(1)181 182 183if __name__ == "__main__":184    main()185