utils/cache_manager.py

"""
Hugging Face キャッシュ管理ユーティリティ
"""
import os
import shutil
import logging
from typing import Optional, Dict, Any
from huggingface_hub import scan_cache_dir, HfApi
from pathlib import Path


class CacheManager:
    """Hugging Face キャッシュを管理するクラス"""
    
    def __init__(self, cache_size_limit_gb: float = 40.0):
        """
        キャッシュマネージャーを初期化
        
        Args:
            cache_size_limit_gb: キャッシュサイズの制限（GB）
        """
        self.cache_size_limit_gb = cache_size_limit_gb
        self.cache_size_limit_bytes = int(cache_size_limit_gb * 1024 * 1024 * 1024)
        self.logger = logging.getLogger(__name__)
        
    def get_cache_info(self) -> Dict[str, Any]:
        """
        キャッシュの現在の状態を取得
        
        Returns:
            キャッシュ情報の辞書
        """
        try:
            cache_info = scan_cache_dir()
            
            total_size = sum(repo.size_on_disk for repo in cache_info.repos)
            
            # キャッシュパスを取得
            cache_path = "N/A"
            if cache_info.repos:
                # 最初のリポジトリのパスから推測
                first_repo = next(iter(cache_info.repos))
                cache_path = str(first_repo.repo_path.parent.parent)
            else:
                # デフォルトのキャッシュパス
                import os
                cache_path = os.getenv("HF_HUB_CACHE", 
                                     os.path.expanduser("~/.cache/huggingface/hub"))
            
            return {
                "total_size_bytes": total_size,
                "total_size_gb": total_size / (1024 * 1024 * 1024),
                "num_repos": len(cache_info.repos),
                "cache_path": cache_path,
                "repos": [
                    {
                        "repo_id": repo.repo_id,
                        "size_on_disk": repo.size_on_disk,
                        "size_on_disk_gb": repo.size_on_disk / (1024 * 1024 * 1024),
                        "last_accessed": repo.last_accessed,
                        "last_modified": repo.last_modified
                    }
                    for repo in cache_info.repos
                ]
            }
        except Exception as e:
            self.logger.error(f"キャッシュ情報の取得に失敗: {e}")
            return {"error": str(e)}
    
    def cleanup_cache(self, force: bool = False) -> Dict[str, Any]:
        """
        キャッシュをクリーンアップ
        
        Args:
            force: 強制的にクリーンアップする場合はTrue
            
        Returns:
            クリーンアップ結果の辞書
        """
        try:
            cache_info = scan_cache_dir()
            total_size = sum(repo.size_on_disk for repo in cache_info.repos)
            
            # キャッシュサイズが制限を超えているか、強制実行の場合
            if total_size > self.cache_size_limit_bytes or force:
                # 最後のアクセス時刻が古い順にソート
                sorted_repos = sorted(
                    cache_info.repos,
                    key=lambda x: x.last_accessed if x.last_accessed else 0
                )
                
                # 削除するリポジトリを選択
                to_delete = []
                freed_size = 0
                
                for repo in sorted_repos:
                    if total_size - freed_size <= self.cache_size_limit_bytes * 0.8:  # 80%まで削減
                        break
                    to_delete.append(repo)
                    freed_size += repo.size_on_disk
                
                if to_delete:
                    # 削除実行
                    revision_hashes = []
                    for repo in to_delete:
                        revision_hashes.extend([rev.revision_hash for rev in repo.revisions])
                    
                    delete_strategy = cache_info.delete_revisions(*revision_hashes)
                    expected_freed_size = delete_strategy.expected_freed_size
                    delete_strategy.execute()
                    
                    self.logger.info(f"キャッシュクリーンアップ完了: {expected_freed_size / (1024 * 1024 * 1024):.2f}GB削除")
                    
                    return {
                        "success": True,
                        "deleted_repos": len(to_delete),
                        "freed_size_bytes": expected_freed_size,
                        "freed_size_gb": expected_freed_size / (1024 * 1024 * 1024)
                    }
                else:
                    return {
                        "success": True,
                        "deleted_repos": 0,
                        "freed_size_bytes": 0,
                        "freed_size_gb": 0,
                        "message": "削除するリポジトリがありません"
                    }
            else:
                return {
                    "success": True,
                    "deleted_repos": 0,
                    "freed_size_bytes": 0,
                    "freed_size_gb": 0,
                    "message": "キャッシュサイズが制限内です"
                }
                
        except Exception as e:
            self.logger.error(f"キャッシュクリーンアップに失敗: {e}")
            return {"success": False, "error": str(e)}
    
    def get_disk_usage(self) -> Dict[str, Any]:
        """
        ディスク使用状況を取得
        
        Returns:
            ディスク使用状況の辞書
        """
        try:
            # ワーキングディレクトリの使用状況
            working_dir = Path.cwd()
            working_dir_usage = shutil.disk_usage(working_dir)
            
            # キャッシュディレクトリの使用状況
            cache_info = self.get_cache_info()
            
            return {
                "working_directory": {
                    "path": str(working_dir),
                    "total_gb": working_dir_usage.total / (1024 * 1024 * 1024),
                    "used_gb": working_dir_usage.used / (1024 * 1024 * 1024),
                    "free_gb": working_dir_usage.free / (1024 * 1024 * 1024),
                    "usage_percent": (working_dir_usage.used / working_dir_usage.total) * 100
                },
                "cache": cache_info
            }
        except Exception as e:
            self.logger.error(f"ディスク使用状況の取得に失敗: {e}")
            return {"error": str(e)}
    
    def should_cleanup(self) -> bool:
        """
        クリーンアップが必要かどうかを判定
        
        Returns:
            クリーンアップが必要な場合True
        """
        try:
            disk_usage = self.get_disk_usage()
            
            if "error" in disk_usage:
                return False
                
            # ディスク使用量が85%を超えている場合、またはキャッシュサイズが制限を超えている場合
            usage_percent = disk_usage["working_directory"]["usage_percent"]
            cache_size_gb = disk_usage["cache"].get("total_size_gb", 0)
            
            return usage_percent > 85 or cache_size_gb > self.cache_size_limit_gb
            
        except Exception as e:
            self.logger.error(f"クリーンアップ判定に失敗: {e}")
            return False
    
    def emergency_cleanup(self) -> Dict[str, Any]:
        """
        緊急時のキャッシュクリーンアップ（より積極的な削除）
        
        Returns:
            クリーンアップ結果の辞書
        """
        try:
            cache_info = scan_cache_dir()
            
            # 全てのリポジトリを削除対象とする
            revision_hashes = []
            for repo in cache_info.repos:
                revision_hashes.extend([rev.revision_hash for rev in repo.revisions])
            
            if revision_hashes:
                delete_strategy = cache_info.delete_revisions(*revision_hashes)
                expected_freed_size = delete_strategy.expected_freed_size
                delete_strategy.execute()
                
                self.logger.warning(f"緊急キャッシュクリーンアップ完了: {expected_freed_size / (1024 * 1024 * 1024):.2f}GB削除")
                
                return {
                    "success": True,
                    "deleted_repos": len(cache_info.repos),
                    "freed_size_bytes": expected_freed_size,
                    "freed_size_gb": expected_freed_size / (1024 * 1024 * 1024)
                }
            else:
                return {
                    "success": True,
                    "deleted_repos": 0,
                    "freed_size_bytes": 0,
                    "freed_size_gb": 0,
                    "message": "削除するキャッシュがありません"
                }
                
        except Exception as e:
            self.logger.error(f"緊急キャッシュクリーンアップに失敗: {e}")
            return {"success": False, "error": str(e)}