File size: 3,823 Bytes
e7abd9e 23c96f8 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 23c96f8 e7abd9e 58582d3 e7abd9e 23c96f8 e7abd9e 23c96f8 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 e7abd9e 58582d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import os
import shutil
from pathlib import Path
from datetime import timedelta
import logging
from app.core.formatting import LogFormatter
from app.config.base import (
CACHE_ROOT,
DATASETS_CACHE,
MODELS_CACHE,
VOTES_CACHE,
EVAL_CACHE,
CACHE_TTL,
)
logger = logging.getLogger(__name__)
class CacheConfig:
def __init__(self):
# Get cache paths from config
self.cache_root = CACHE_ROOT
self.datasets_cache = DATASETS_CACHE
self.models_cache = MODELS_CACHE
self.votes_cache = VOTES_CACHE
self.eval_cache = EVAL_CACHE
# Specific files
self.votes_file = self.votes_cache / "votes_data.jsonl"
self.eval_requests_file = self.eval_cache / "eval_requests.jsonl"
# Cache TTL
self.cache_ttl = timedelta(seconds=CACHE_TTL)
self._initialize_cache_dirs()
self._setup_environment()
def _initialize_cache_dirs(self):
"""Initialize all necessary cache directories"""
try:
logger.info(LogFormatter.section("CACHE INITIALIZATION"))
cache_dirs = {
"Root": self.cache_root,
"Datasets": self.datasets_cache,
"Models": self.models_cache,
"Votes": self.votes_cache,
"Eval": self.eval_cache,
}
for name, cache_dir in cache_dirs.items():
cache_dir.mkdir(parents=True, exist_ok=True)
logger.info(
LogFormatter.success(f"{name} cache directory: {cache_dir}")
)
except Exception as e:
logger.error(LogFormatter.error("Failed to create cache directories", e))
raise
def _setup_environment(self):
"""Configure HuggingFace environment variables"""
logger.info(LogFormatter.subsection("ENVIRONMENT SETUP"))
env_vars = {
"HF_HOME": str(self.cache_root),
"HF_DATASETS_CACHE": str(self.datasets_cache),
}
for var, value in env_vars.items():
os.environ[var] = value
logger.info(LogFormatter.info(f"Set {var}={value}"))
def get_cache_path(self, cache_type: str) -> Path:
"""Returns the path for a specific cache type"""
cache_paths = {
"datasets": self.datasets_cache,
"models": self.models_cache,
"votes": self.votes_cache,
"eval": self.eval_cache,
}
return cache_paths.get(cache_type, self.cache_root)
def flush_cache(self, cache_type: str = None):
"""Flush specified cache or all caches if no type is specified"""
try:
if cache_type:
logger.info(
LogFormatter.section(f"FLUSHING {cache_type.upper()} CACHE")
)
cache_dir = self.get_cache_path(cache_type)
if cache_dir.exists():
stats = {"Cache_Type": cache_type, "Directory": str(cache_dir)}
for line in LogFormatter.tree(stats, "Cache Details"):
logger.info(line)
shutil.rmtree(cache_dir)
cache_dir.mkdir(parents=True, exist_ok=True)
logger.info(LogFormatter.success("Cache cleared successfully"))
else:
logger.info(LogFormatter.section("FLUSHING ALL CACHES"))
for cache_type in ["datasets", "models", "votes", "eval"]:
self.flush_cache(cache_type)
logger.info(LogFormatter.success("All caches cleared successfully"))
except Exception as e:
logger.error(LogFormatter.error("Failed to flush cache", e))
raise
# Singleton instance of cache configuration
cache_config = CacheConfig()
|