utils/hub_storage.py

from pathlib import Path
from huggingface_hub import HfApi
from .token_handler import TokenHandler

class HubStorage:
    def __init__(self, repo_id):
        self.repo_id = repo_id
        self.api = HfApi()
        
        # Validate repository access
        try:
            self.api.repo_info(repo_id=repo_id, repo_type="space")
        except Exception as e:
            raise RuntimeError(f"Failed to access repository {repo_id}: {e}")

    def save_to_hub(self, file_content, path_in_repo, commit_message):
        """Save file to HuggingFace Hub with retries and better error handling"""
        max_retries = 3
        retry_delay = 1  # seconds
        
        for attempt in range(max_retries):
            try:
                self.api.upload_file(
                    repo_id=self.repo_id,
                    repo_type="space",
                    path_or_fileobj=file_content,
                    path_in_repo=path_in_repo,
                    commit_message=commit_message
                )
                return True
            except Exception as e:
                if attempt == max_retries - 1:
                    raise RuntimeError(f"Failed to save file after {max_retries} attempts: {e}")
                print(f"Attempt {attempt + 1} failed, retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
                retry_delay *= 2  # Exponential backoff