storage.py (5010B)
1 import os 2 import time 3 import uuid 4 import httpx 5 from huggingface_hub import HfApi 6 7 HF_REPO_ID = os.getenv("HF_REPO_ID", "notamitgamer/cdn") 8 HF_TOKEN = os.getenv("HF_TOKEN") 9 10 api = HfApi(token=HF_TOKEN) 11 12 # LRU Cache implementation 13 _cache = {} 14 CACHE_TTL = 60 15 MAX_CACHE_SIZE = 500 16 17 def _get_cache(key): 18 if key in _cache: 19 expiry, value = _cache[key] 20 if expiry > time.time(): 21 # Move to the end of the dictionary to mark it as most recently used (LRU) 22 _cache[key] = _cache.pop(key) 23 return value 24 else: 25 # Expired 26 del _cache[key] 27 return None 28 29 def _set_cache(key, value): 30 if key in _cache: 31 # Remove it first so we can re-insert it at the end 32 del _cache[key] 33 elif len(_cache) >= MAX_CACHE_SIZE: 34 # Dictionary is full. Remove the oldest item (which is at the front) 35 oldest_key = next(iter(_cache)) 36 del _cache[oldest_key] 37 38 _cache[key] = (time.time() + CACHE_TTL, value) 39 40 async def is_file(path: str) -> bool: 41 if not path: 42 return False 43 44 cache_key = f"is_file_{path}" 45 cached = _get_cache(cache_key) 46 if cached is not None: 47 return cached 48 49 hf_url = f"https://huggingface.co/datasets/{HF_REPO_ID}/resolve/main/{path}" 50 async with httpx.AsyncClient() as client: 51 r = await client.head(hf_url) 52 result = r.status_code == 200 53 _set_cache(cache_key, result) 54 return result 55 56 def _fetch_tree(path: str): 57 return list(api.list_repo_tree( 58 repo_id=HF_REPO_ID, path_in_repo=path, repo_type="dataset", expand=False 59 )) 60 61 def list_directory(path: str): 62 cache_key = f"list_dir_{path}" 63 cached = _get_cache(cache_key) 64 if cached is not None: 65 return cached 66 67 try: 68 items = _fetch_tree(path) 69 except Exception as e: 70 print(f"[storage] list_repo_tree failed for {path!r}: {e}") 71 return [] 72 73 files_and_folders = [] 74 for item in items: 75 name = item.path.split("/")[-1] 76 is_dir = not hasattr(item, "size") 77 size = getattr(item, "size", 0) 78 79 size_str = format_size(size) if not is_dir else "-" 80 81 files_and_folders.append({ 82 "name": name, 83 "path": item.path, 84 "is_dir": is_dir, 85 "size_str": size_str 86 }) 87 88 files_and_folders.sort(key=lambda x: (not x["is_dir"], x["name"].lower())) 89 90 _set_cache(cache_key, files_and_folders) 91 return files_and_folders 92 93 ZIP_MAX_FILES = 300 94 ZIP_MAX_TOTAL_BYTES = 500 * 1024 * 1024 95 96 def list_files_recursive(path: str): 97 try: 98 items = list(api.list_repo_tree( 99 repo_id=HF_REPO_ID, path_in_repo=path, repo_type="dataset", recursive=True 100 )) 101 except Exception: 102 return [] 103 104 files = [] 105 total_size = 0 106 for item in items: 107 if hasattr(item, "size"): 108 size = getattr(item, "size", 0) or 0 109 total_size += size 110 files.append({"path": item.path, "size": size}) 111 if len(files) > ZIP_MAX_FILES or total_size > ZIP_MAX_TOTAL_BYTES: 112 raise ValueError( 113 f"Folder too large to zip (limit: {ZIP_MAX_FILES} files / " 114 f"{ZIP_MAX_TOTAL_BYTES // (1024 * 1024)} MB)." 115 ) 116 return files 117 118 def format_size(size_bytes) -> str: 119 if size_bytes is None: 120 return "-" 121 size = float(size_bytes) 122 units = ["B", "KB", "MB", "GB", "TB"] 123 idx = 0 124 while size >= 1024 and idx < len(units) - 1: 125 size /= 1024 126 idx += 1 127 if idx == 0: 128 return f"{int(size)} {units[idx]}" 129 return f"{size:.2f} {units[idx]}" 130 131 REPO_STATS_CACHE_TTL = 300 132 133 def repo_stats(): 134 cache_key = "repo_stats" 135 now = time.time() 136 if cache_key in _cache: 137 expiry, value = _cache[cache_key] 138 if expiry > now: 139 # Move to end as part of LRU update 140 _cache[cache_key] = _cache.pop(cache_key) 141 return value 142 else: 143 del _cache[cache_key] 144 145 try: 146 items = list(api.list_repo_tree( 147 repo_id=HF_REPO_ID, path_in_repo="", repo_type="dataset", recursive=True 148 )) 149 except Exception as e: 150 print(f"[storage] repo_stats failed: {e}") 151 return None 152 153 file_count = 0 154 total_size = 0 155 for item in items: 156 if hasattr(item, "size"): 157 file_count += 1 158 total_size += getattr(item, "size", 0) or 0 159 160 result = {"file_count": file_count, "size_str": format_size(total_size)} 161 _set_cache(cache_key, result) 162 return result 163 164 def upload_temp_file(temp_path: str, filename: str) -> str: 165 slug = str(uuid.uuid4())[:8] 166 safe_filename = filename.replace(" ", "_") 167 hf_path = f"uploads/{slug}-{safe_filename}" 168 169 api.upload_file( 170 path_or_fileobj=temp_path, 171 path_in_repo=hf_path, 172 repo_id=HF_REPO_ID, 173 repo_type="dataset", 174 token=HF_TOKEN 175 ) 176 177 _cache.clear() 178 return hf_path