rahmad7/hermes-openmodel

1#!/usr/bin/env python32"""3Atomic Dataset Persistence for HermesFace4Save state to Hugging Face Dataset with atomic commit, checksum, and backup.5 6Usage:7    python3 save_to_dataset_atomic.py <repo_id> <source_path1> [source_path2...]8 9Env vars:10    HF_TOKEN              HF access token (read from env by huggingface_hub)11"""12import hashlib13import json14import logging15import os16import shutil17import sys18import tempfile19import time20from datetime import datetime21from pathlib import Path22from typing import Any, Dict, List, Optional23 24from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download25from huggingface_hub.utils import RepositoryNotFoundError26 27logging.basicConfig(28    level=logging.INFO,29    format='{"timestamp": "%(asctime)s", "level": "%(levelname)s", "module": "atomic-save", "message": "%(message)s"}',30)31logger = logging.getLogger(__name__)32 33 34class AtomicDatasetSaver:35    def __init__(self, repo_id: str, dataset_path: str = "state"):36        self.repo_id = repo_id37        self.dataset_path = Path(dataset_path)38        self.api = HfApi()39        self.max_retries = 340        self.base_delay = 1.041        self.max_backups = 342        logger.info(f"init repo_id={repo_id} dataset_path={dataset_path}")43 44    def calculate_checksum(self, file_path: Path) -> str:45        h = hashlib.sha256()46        with open(file_path, "rb") as f:47            for chunk in iter(lambda: f.read(4096), b""):48                h.update(chunk)49        return h.hexdigest()50 51    def create_backup(self, current_commit: Optional[str]) -> Optional[str]:52        if not current_commit:53            return None54        try:55            files = self.api.list_repo_files(56                repo_id=self.repo_id, repo_type="dataset", revision=current_commit57            )58            state_files = [f for f in files if f.startswith(str(self.dataset_path))]59            if not state_files:60                return None61 62            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")63            backup_path = f"backups/state_{timestamp}"64            logger.info(f"creating_backup path={backup_path} files={len(state_files)}")65 66            with tempfile.TemporaryDirectory() as tmpdir:67                tmpdir_path = Path(tmpdir)68                ops = []69                for file_path in state_files:70                    local = hf_hub_download(71                        repo_id=self.repo_id,72                        repo_type="dataset",73                        filename=file_path,74                        revision=current_commit,75                    )76                    if local:77                        dst = tmpdir_path / Path(file_path).name78                        shutil.copy2(local, dst)79                        ops.append(CommitOperationAdd(80                            path_in_repo=f"{backup_path}/{Path(file_path).name}",81                            path_or_fileobj=str(dst),82                        ))83 84                if ops:85                    info = self.api.create_commit(86                        repo_id=self.repo_id,87                        repo_type="dataset",88                        operations=ops,89                        commit_message=f"Backup state before update - {timestamp}",90                        parent_commit=current_commit,91                    )92                    logger.info(f"backup_created commit={info.oid}")93                    return info.oid94        except Exception as e:95            logger.error(f"backup_failed error={e}")96        return None97 98    def save_state_atomic(99        self, state_data: Dict[str, Any], source_paths: List[str]100    ) -> Dict[str, Any]:101        operation_id = f"save_{int(time.time())}"102        logger.info(f"starting_atomic_save op={operation_id} sources={source_paths}")103 104        try:105            try:106                repo_info = self.api.repo_info(repo_id=self.repo_id, repo_type="dataset")107                current_commit = repo_info.sha108            except RepositoryNotFoundError:109                current_commit = None110 111            backup_commit = self.create_backup(current_commit)112 113            with tempfile.TemporaryDirectory() as tmpdir:114                state_dir = Path(tmpdir) / self.dataset_path115                state_dir.mkdir(parents=True, exist_ok=True)116 117                metadata = {118                    "timestamp": datetime.now().isoformat(),119                    "operation_id": operation_id,120                    "checksum": None,121                    "backup_commit": backup_commit,122                    "state_data": state_data,123                }124                metadata_path = state_dir / "metadata.json"125                with open(metadata_path, "w") as f:126                    json.dump(metadata, f, indent=2)127 128                ops = [CommitOperationAdd(129                    path_in_repo="state/metadata.json",130                    path_or_fileobj=str(metadata_path),131                )]132 133                for source_path in source_paths:134                    src = Path(source_path)135                    if src.exists():136                        dst = state_dir / src.name137                        shutil.copy2(src, dst)138                        checksum = self.calculate_checksum(dst)139                        ops.append(CommitOperationAdd(140                            path_in_repo=f"state/{src.name}",141                            path_or_fileobj=str(dst),142                        ))143                        logger.info(f"file_added source={source_path} sha256={checksum[:12]}")144 145                metadata["checksum"] = hashlib.sha256(146                    json.dumps(state_data, sort_keys=True).encode()147                ).hexdigest()148                with open(metadata_path, "w") as f:149                    json.dump(metadata, f, indent=2)150 151                info = self.api.create_commit(152                    repo_id=self.repo_id,153                    repo_type="dataset",154                    operations=ops,155                    commit_message=f"Atomic state update - {operation_id}",156                    parent_commit=current_commit,157                )158 159                result = {160                    "success": True,161                    "operation_id": operation_id,162                    "commit_id": info.oid,163                    "backup_commit": backup_commit,164                    "timestamp": datetime.now().isoformat(),165                    "files_count": len(source_paths),166                }167                logger.info(f"atomic_save_completed {result}")168                return result169 170        except Exception as e:171            logger.error(f"atomic_save_failed error={e}")172            raise173 174 175def main() -> None:176    if len(sys.argv) < 3:177        print(json.dumps({178            "error": "Usage: python save_to_dataset_atomic.py <repo_id> <source_path1> [source_path2...]",179            "status": "error",180        }, indent=2))181        sys.exit(1)182 183    repo_id = sys.argv[1]184    source_paths = sys.argv[2:]185    for p in source_paths:186        if not os.path.exists(p):187            print(json.dumps({"error": f"Source path does not exist: {p}", "status": "error"}, indent=2))188            sys.exit(1)189 190    state_data = {191        "environment": "production",192        "version": "1.0.0",193        "platform": "huggingface-spaces",194        "app": "hermesface",195        "timestamp": datetime.now().isoformat(),196    }197 198    try:199        saver = AtomicDatasetSaver(repo_id)200        result = saver.save_state_atomic(state_data, source_paths)201        print(json.dumps(result, indent=2))202    except Exception as e:203        print(json.dumps({"error": str(e), "status": "error"}, indent=2))204        sys.exit(1)205 206 207if __name__ == "__main__":208    main()209