Update app.py

app.py

@@ -7,13 +7,14 @@ import zipfile
 import shutil
 from typing import Dict, List, Set, Optional, Tuple, Any
 from urllib.parse import quote
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
 import io
 
 from fastapi import FastAPI, BackgroundTasks, HTTPException, status
 from pydantic import BaseModel, Field
 from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub.utils import HfHubHTTPError
 
 # --- Configuration ---
 AUTO_START_INDEX = 1# Hardcoded default start index if no progress is found
@@ -191,6 +192,16 @@ def save_json_state(file_path: str, data: Dict[str, Any]):
     with open(file_path, "w") as f:
         json.dump(data, f, indent=2)
 
+async def handle_rate_limit_error(e: Exception, context: str) -> bool:
+    """Handles 429 rate limit error by pausing for 1 hour."""
+    if isinstance(e, HfHubHTTPError) and e.response.status_code == 429:
+        pause_duration = 3600  # 1 hour in seconds
+        print(f"[{FLOW_ID}] 🚨 RATE LIMIT (429) DETECTED during {context}. Pausing for {pause_duration // 3600} hour(s)...")
+        print(f"[{FLOW_ID}] Resuming at: {datetime.now() + timedelta(seconds=pause_duration)}")
+        await asyncio.sleep(pause_duration)
+        return True
+    return False
+
 async def download_hf_state() -> Dict[str, Any]:
     """Downloads the state file from Hugging Face or returns a default state."""
     local_path = LOCAL_STATE_FOLDER / HF_STATE_FILE
@@ -225,26 +236,38 @@ async def download_hf_state() -> Dict[str, Any]:
         return default_state
 
 async def upload_hf_state(state: Dict[str, Any]) -> bool:
-    """Uploads the state file to Hugging Face."""
+    """Uploads the state file to Hugging Face with 429 rate limit handling."""
     local_path = LOCAL_STATE_FOLDER / HF_STATE_FILE
     
-    try:
-        # Save state locally first
-        save_json_state(str(local_path), state)
-        
-        # Upload to helium dataset
-        HfApi(token=HF_TOKEN).upload_file(
-            path_or_fileobj=str(local_path),
-            path_in_repo=HF_STATE_FILE,
-            repo_id=HF_OUTPUT_DATASET_ID,
-            repo_type="dataset",
-            commit_message=f"Update caption processing state: next_index={state['next_download_index']}"
-        )
-        print(f"[{FLOW_ID}] Successfully uploaded state file.")
-        return True
-    except Exception as e:
-        print(f"[{FLOW_ID}] Failed to upload state file: {str(e)}")
-        return False
+    MAX_RETRIES = 5
+    for attempt in range(MAX_RETRIES):
+        try:
+            # Save state locally first
+            save_json_state(str(local_path), state)
+            
+            # Upload to helium dataset
+            HfApi(token=HF_TOKEN).upload_file(
+                path_or_fileobj=str(local_path),
+                path_in_repo=HF_STATE_FILE,
+                repo_id=HF_OUTPUT_DATASET_ID,
+                repo_type="dataset",
+                commit_message=f"Update caption processing state: next_index={state['next_download_index']}"
+            )
+            print(f"[{FLOW_ID}] Successfully uploaded state file.")
+            return True
+        except Exception as e:
+            if await handle_rate_limit_error(e, "state file upload"):
+                continue # Retry after pause
+            
+            print(f"[{FLOW_ID}] Failed to upload state file on attempt {attempt + 1}: {str(e)}")
+            # If it's not a rate limit error, or we've exhausted retries, break
+            if attempt == MAX_RETRIES - 1:
+                break
+            # Add a small delay before non-429 retries
+            await asyncio.sleep(5)
+    
+    print(f"[{FLOW_ID}] CRITICAL: Failed to upload state file after {MAX_RETRIES} attempts.")
+    return False
 
 async def lock_file_for_processing(zip_filename: str, state: Dict[str, Any]) -> bool:
     """Marks a file as 'processing' in the state file and uploads the lock."""
@@ -362,31 +385,43 @@ async def download_and_extract_zip_by_index(file_index: int, repo_file_full_path
         return None
 
 async def upload_captions_to_hf(zip_full_name: str, captions: List[Dict]) -> bool:
-    """Uploads the final captions JSON file to the output dataset."""
+    """Uploads the final captions JSON file to the output dataset with 429 rate limit handling."""
     # Use the full zip name, replacing the extension with .json
     caption_filename = Path(zip_full_name).with_suffix('.json').name
     
-    try:
-        print(f"[{FLOW_ID}] Uploading {len(captions)} captions for {zip_full_name} as {caption_filename} to {HF_OUTPUT_DATASET_ID}...")
-        
-        # Create JSON content in memory
-        json_content = json.dumps(captions, indent=2, ensure_ascii=False).encode('utf-8')
-        
-        api = HfApi(token=HF_TOKEN)
-        api.upload_file(
-            path_or_fileobj=io.BytesIO(json_content),
-            path_in_repo=caption_filename,
-            repo_id=HF_OUTPUT_DATASET_ID,
-            repo_type="dataset",
-            commit_message=f"[{FLOW_ID}] Captions for {zip_full_name}"
-        )
-        
-        print(f"[{FLOW_ID}] Successfully uploaded captions for {zip_full_name}.")
-        return True
-        
-    except Exception as e:
-        print(f"[{FLOW_ID}] Error uploading captions for {zip_full_name}: {e}")
-        return False
+    MAX_RETRIES = 5
+    for attempt in range(MAX_RETRIES):
+        try:
+            print(f"[{FLOW_ID}] Uploading {len(captions)} captions for {zip_full_name} as {caption_filename} to {HF_OUTPUT_DATASET_ID}...")
+            
+            # Create JSON content in memory
+            json_content = json.dumps(captions, indent=2, ensure_ascii=False).encode('utf-8')
+            
+            api = HfApi(token=HF_TOKEN)
+            api.upload_file(
+                path_or_fileobj=io.BytesIO(json_content),
+                path_in_repo=caption_filename,
+                repo_id=HF_OUTPUT_DATASET_ID,
+                repo_type="dataset",
+                commit_message=f"[{FLOW_ID}] Captions for {zip_full_name}"
+            )
+            
+            print(f"[{FLOW_ID}] Successfully uploaded captions for {zip_full_name}.")
+            return True
+            
+        except Exception as e:
+            if await handle_rate_limit_error(e, f"captions upload for {zip_full_name}"):
+                continue # Retry after pause
+            
+            print(f"[{FLOW_ID}] Error uploading captions for {zip_full_name} on attempt {attempt + 1}: {e}")
+            # If it's not a rate limit error, or we've exhausted retries, break
+            if attempt == MAX_RETRIES - 1:
+                break
+            # Add a small delay before non-429 retries
+            await asyncio.sleep(5)
+            
+    print(f"[{FLOW_ID}] CRITICAL: Failed to upload captions for {zip_full_name} after {MAX_RETRIES} attempts.")
+    return False
 
 # --- Core Processing Functions (Modified) ---