Update app.py

app.py

@@ -5,7 +5,7 @@ import asyncio
 import aiohttp
 import zipfile
 import shutil
-from typing import Dict, List, Set, Optional, Tuple
+from typing import Dict, List, Set, Optional, Tuple, Any
 from urllib.parse import quote
 from datetime import datetime
 from pathlib import Path
@@ -23,50 +23,45 @@ HF_TOKEN = os.getenv("HF_TOKEN", "")
 HF_DATASET_ID = os.getenv("HF_DATASET_ID", "Fred808/BG3") # Source dataset for zip files
 HF_OUTPUT_DATASET_ID = os.getenv("HF_OUTPUT_DATASET_ID", "fred808/helium") # Target dataset for captions
 
-# Progress Tracking File
+# Progress and State Tracking
 PROGRESS_FILE = Path("processing_progress.json")
+HF_STATE_FILE = "processing_state_captions.json"  # State file in helium dataset
+LOCAL_STATE_FOLDER = Path(".state")  # Local folder for state file
+LOCAL_STATE_FOLDER.mkdir(exist_ok=True)
+
 # Directory within the HF dataset where the zip files are located
-ZIP_FILE_PREFIX = "frames/"
+ZIP_FILE_PREFIX = "frames_zips/"
 
 # Using the full list from the user's original code for actual deployment
 CAPTION_SERVERS = [
-   "https://fred808-pil-4-1.hf.space/analyze",
-    "https://fred808-pil-4-2.hf.space/analyze",
-    "https://fred808-pil-4-3.hf.space/analyze",
-    "https://fred1012-fred1012-gw0j2h.hf.space/analyze",
-    "https://fred1012-fred1012-wqs6c2.hf.space/analyze",
-    "https://fred1012-fred1012-oncray.hf.space/analyze",
-    "https://fred1012-fred1012-4goge7.hf.space/analyze",
-    "https://fred1012-fred1012-z0eh7m.hf.space/analyze",
-    "https://fred1012-fred1012-u95rte.hf.space/analyze",
-    "https://fred1012-fred1012-igje22.hf.space/analyze",
-    "https://fred1012-fred1012-ibkuf8.hf.space/analyze",
-    "https://fred1012-fred1012-nwqthy.hf.space/analyze",
-    "https://fred1012-fred1012-4ldqj4.hf.space/analyze",
-    "https://fred1012-fred1012-pivlzg.hf.space/analyze",
-    "https://fred1012-fred1012-ptlc5u.hf.space/analyze",
-    "https://fred1012-fred1012-u7lh57.hf.space/analyze",
-    "https://fred1012-fred1012-q8djv1.hf.space/analyze",
-    "https://fredalone-fredalone-ozugrp.hf.space/analyze",
-    "https://fredalone-fredalone-9brxj2.hf.space/analyze",
-    "https://fredalone-fredalone-p8vq9a.hf.space/analyze",
-    "https://fredalone-fredalone-vbli2y.hf.space/analyze",
-    "https://fredalone-fredalone-uggger.hf.space/analyze",
-    "https://fredalone-fredalone-nmi7e8.hf.space/analyze",
-    "https://fredalone-fredalone-d1f26d.hf.space/analyze",
-    "https://fredalone-fredalone-461jp2.hf.space/analyze",
-    "https://fredalone-fredalone-3enfg4.hf.space/analyze",
-    "https://fredalone-fredalone-dqdbpv.hf.space/analyze",
-    "https://fredalone-fredalone-ivtjua.hf.space/analyze",
-    "https://fredalone-fredalone-6bezt2.hf.space/analyze",
-    "https://fredalone-fredalone-e0wfnk.hf.space/analyze",
-    "https://fredalone-fredalone-zu2t7j.hf.space/analyze",
-    "https://fredalone-fredalone-dqtv1o.hf.space/analyze",
-    "https://fredalone-fredalone-wclyog.hf.space/analyze",
-    "https://fredalone-fredalone-t27vig.hf.space/analyze",
-    "https://fredalone-fredalone-gahbxh.hf.space/analyze",
-    "https://fredalone-fredalone-kw2po4.hf.space/analyze",
-    "https://fredalone-fredalone-8h285h.hf.space/analyze"
+"https://favoredone-favoredone-tv88mp.hf.space/analyze",
+"https://favoredone-favoredone-7p1dcf.hf.space/analyze",
+"https://favoredone-favoredone-k7b4mf.hf.space/analyze",
+"https://favoredone-favoredone-mzlxc7.hf.space/analyze",
+"https://favoredone-favoredone-aomfwa.hf.space/analyze",
+"https://favoredone-favoredone-7g6v04.hf.space/analyze",
+"https://favoredone-favoredone-dk1skh.hf.space/analyze",
+"https://favoredone-favoredone-z4yo0y.hf.space/analyze",
+"https://favoredone-favoredone-f6czeq.hf.space/analyze",
+"https://favoredone-favoredone-5fo8ga.hf.space/analyze",
+"https://favoredone-favoredone-zde8x6.hf.space/analyze",
+"https://favoredone-favoredone-r0biih.hf.space/analyze",
+"https://favoredone-favoredone-ljdzkf.hf.space/analyze",
+"https://favoredone-favoredone-irrpe5.hf.space/analyze",
+"https://favoredone-favoredone-bh9rwz.hf.space/analyze",
+"https://favoredone-favoredone-u8c4dt.hf.space/analyze",
+"https://favoredone-favoredone-futwyd.hf.space/analyze",
+"https://favoredone-favoredone-hg2sot.hf.space/analyze",
+"https://favoredone-favoredone-pvweug.hf.space/analyze",
+"https://favoredone-favoredone-z6azk2.hf.space/analyze",
+"https://favoredone-favoredone-4zid9w.hf.space/analyze",
+"https://favoredone-favoredone-be7a1r.hf.space/analyze",
+"https://favoredone-favoredone-ayazxa.hf.space/analyze",
+"https://favoredone-favoredone-6ckj4m.hf.space/analyze",
+"https://favoredone-favoredone-whn0xu.hf.space/analyze",
+"https://favoredone-favoredone-t49exm.hf.space/analyze",
+"https://favoredone-favoredone-cgrh0a.hf.space/analyze",
+"https://favoredone-favoredone-r1kb5g.hf.space/analyze"
 ]
 MODEL_TYPE = "Florence-2-large"
 
@@ -94,10 +89,10 @@ class CaptionServer:
 servers = [CaptionServer(url) for url in CAPTION_SERVERS]
 server_index = 0
 
-# --- Progress Tracking Functions ---
+# --- Progress and State Management Functions ---
 
 def load_progress() -> Dict:
-    """Loads the processing progress from the JSON file."""
+    """Loads the local processing progress from the JSON file."""
     if PROGRESS_FILE.exists():
         try:
             with PROGRESS_FILE.open('r') as f:
@@ -114,13 +109,127 @@ def load_progress() -> Dict:
     }
 
 def save_progress(progress_data: Dict):
-    """Saves the processing progress to the JSON file."""
+    """Saves the local processing progress to the JSON file."""
     try:
         with PROGRESS_FILE.open('w') as f:
             json.dump(progress_data, f, indent=4)
     except Exception as e:
         print(f"[{FLOW_ID}] CRITICAL ERROR: Could not save progress to {PROGRESS_FILE}: {e}")
 
+def load_json_state(file_path: str, default_value: Dict[str, Any]) -> Dict[str, Any]:
+    """Load state from JSON file with migration logic for new structure."""
+    if os.path.exists(file_path):
+        try:
+            with open(file_path, "r") as f:
+                data = json.load(f)
+                
+                # Migration Logic
+                if "file_states" not in data or not isinstance(data["file_states"], dict):
+                    print(f"[{FLOW_ID}] Initializing 'file_states' dictionary.")
+                    data["file_states"] = {}
+                    
+                if "next_download_index" not in data:
+                    data["next_download_index"] = 0
+                    
+                return data
+        except json.JSONDecodeError:
+            print(f"[{FLOW_ID}] WARNING: Corrupted state file: {file_path}")
+    return default_value
+
+def save_json_state(file_path: str, data: Dict[str, Any]):
+    """Save state to JSON file"""
+    with open(file_path, "w") as f:
+        json.dump(data, f, indent=2)
+
+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
+    default_state = {"next_download_index": 0, "file_states": {}}
+    
+    try:
+        # Check if the file exists in the helium repo
+        files = HfApi(token=HF_TOKEN).list_repo_files(
+            repo_id=HF_OUTPUT_DATASET_ID,
+            repo_type="dataset"
+        )
+        
+        if HF_STATE_FILE not in files:
+            print(f"[{FLOW_ID}] State file not found in {HF_OUTPUT_DATASET_ID}. Starting fresh.")
+            return default_state
+
+        # Download the file
+        hf_hub_download(
+            repo_id=HF_OUTPUT_DATASET_ID,
+            filename=HF_STATE_FILE,
+            repo_type="dataset",
+            local_dir=LOCAL_STATE_FOLDER,
+            local_dir_use_symlinks=False,
+            token=HF_TOKEN
+        )
+        
+        print(f"[{FLOW_ID}] Successfully downloaded state file.")
+        return load_json_state(str(local_path), default_state)
+        
+    except Exception as e:
+        print(f"[{FLOW_ID}] Failed to download state file: {str(e)}. Starting fresh.")
+        return default_state
+
+async def upload_hf_state(state: Dict[str, Any]) -> bool:
+    """Uploads the state file to Hugging Face."""
+    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
+
+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."""
+    print(f"[{FLOW_ID}] 🔒 Attempting to lock file: {zip_filename}")
+    
+    # Update state locally
+    state["file_states"][zip_filename] = "processing"
+    
+    # Upload the updated state file immediately to establish the lock
+    if await upload_hf_state(state):
+        print(f"[{FLOW_ID}] ✅ Successfully locked file: {zip_filename}")
+        return True
+    else:
+        print(f"[{FLOW_ID}] ❌ Failed to lock file: {zip_filename}")
+        # Revert local state
+        if zip_filename in state["file_states"]:
+            del state["file_states"][zip_filename]
+        return False
+
+async def unlock_file_as_processed(zip_filename: str, state: Dict[str, Any], next_index: int) -> bool:
+    """Marks a file as 'processed', updates the index, and uploads the state."""
+    print(f"[{FLOW_ID}] 🔓 Marking file as processed: {zip_filename}")
+    
+    # Update state locally
+    state["file_states"][zip_filename] = "processed"
+    state["next_download_index"] = next_index
+    
+    # Upload the updated state
+    if await upload_hf_state(state):
+        print(f"[{FLOW_ID}] ✅ Successfully marked as processed: {zip_filename}")
+        return True
+    else:
+        print(f"[{FLOW_ID}] ❌ Failed to update state for: {zip_filename}")
+        return False
+
 # --- Hugging Face Utility Functions ---
 
 async def get_zip_file_list(progress_data: Dict) -> List[str]:
@@ -325,7 +434,9 @@ async def send_image_for_captioning(image_path: Path, course_name: str, progress
 async def process_dataset_task(start_index: int):
     """Main task to process the dataset sequentially starting from a given index."""
     
+    # Load both local progress and HF state
     progress = load_progress()
+    current_state = await download_hf_state()
     file_list = await get_zip_file_list(progress)
     
     if not file_list:
@@ -350,11 +461,18 @@ async def process_dataset_task(start_index: int):
         zip_full_name = Path(repo_file_full_path).name 
         course_name = zip_full_name.replace('.zip', '') # Use the file name as the course/job name
         
-        # Check if the file has already been successfully processed
-        if str(file_index) in progress['processed_files']:
-            print(f"[{FLOW_ID}] Skipping file #{file_index} ({zip_full_name}): Already processed according to progress file.")
-            progress['last_processed_index'] = file_index
-            save_progress(progress)
+        # Check file state in both local and HF state
+        file_state = current_state["file_states"].get(zip_full_name)
+        if file_state == "processed":
+            print(f"[{FLOW_ID}] Skipping {zip_full_name}: Already processed in global state.")
+            continue
+        elif file_state == "processing":
+            print(f"[{FLOW_ID}] Skipping {zip_full_name}: Currently being processed by another worker.")
+            continue
+            
+        # Try to lock the file
+        if not await lock_file_for_processing(zip_full_name, current_state):
+            print(f"[{FLOW_ID}] Failed to lock {zip_full_name}. Skipping.")
             continue
             
         extract_dir = None
@@ -433,16 +551,22 @@ async def process_dataset_task(start_index: int):
                 shutil.rmtree(extract_dir, ignore_errors=True)
                 
             if current_file_success:
-                # Update progress only on successful completion of the file
+                # Update both local progress and HF state
                 progress['last_processed_index'] = file_index
                 progress['processed_files'][str(file_index)] = repo_file_full_path
                 save_progress(progress)
-                print(f"[{FLOW_ID}] Progress saved: File #{file_index} marked as processed.")
+                
+                # Update HF state and unlock the file
+                if await unlock_file_as_processed(zip_full_name, current_state, file_index + 1):
+                    print(f"[{FLOW_ID}] Progress saved and file unlocked: {zip_full_name}")
+                else:
+                    print(f"[{FLOW_ID}] Warning: File processed but state update failed: {zip_full_name}")
             else:
-                # If a file fails, we stop the continuous loop to allow for manual intervention or a fresh start
-                print(f"[{FLOW_ID}] File #{file_index} failed. Stopping continuous processing.")
+                # Mark as failed in the state and continue with next file
+                current_state["file_states"][zip_full_name] = "failed"
+                await upload_hf_state(current_state)
+                print(f"[{FLOW_ID}] File {zip_full_name} marked as failed. Continuing with next file.")
                 global_success = False
-                break
                 
     print(f"[{FLOW_ID}] All processing loops complete. Overall success: {global_success}")
     return global_success