Create app.py

app.py

@@ -0,0 +1,327 @@
+import os
+import json
+import asyncio
+import aiohttp
+from typing import Dict, List, Set, Optional
+from pathlib import Path
+from datetime import datetime
+
+from fastapi import FastAPI, BackgroundTasks, HTTPException, status
+from pydantic import BaseModel
+from huggingface_hub import HfApi, HfFileSystem
+import uvicorn
+
+# --- Configuration ---
+# Manager Server will run on port 8000
+MANAGER_PORT = 8000
+
+# Hugging Face Configuration
+HF_TOKEN = "" # User provided token
+HF_DATASET_ID = "Fred808/BG3"
+HF_DATASET_REPO_TYPE = "dataset"
+FRAMES_FOLDER = "frames"
+STATE_FILE_PATH = "flow_processing_state.json"
+
+# Flow Server Configuration (Hardcoded as per user request)
+# NOTE: These URLs must be accessible to the Manager Server.
+# For local testing, you might use localhost with different ports (e.g., 8001 and 8002)
+FLOW_SERVERS = {
+    "flow1": "http://localhost:8001",
+    "flow2": "http://localhost:8002",
+}
+
+# --- State Management Models ---
+class TaskStatus(BaseModel):
+    status: str # UNPROCESSED, IN_PROGRESS, COMPLETED, FAILED
+    assigned_to: Optional[str] = None # flow1 or flow2
+    assigned_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    error_message: Optional[str] = None
+
+class ProcessingState(BaseModel):
+    # Key is the zip file name (which is the course name)
+    tasks: Dict[str, TaskStatus] = {}
+    
+    # Track which flow server is currently processing which task
+    flow_assignments: Dict[str, Optional[str]] = {
+        "flow1": None,
+        "flow2": None,
+    }
+
+class CompleteTaskRequest(BaseModel):
+    flow_id: str
+    course_name: str
+    success: bool
+    error_message: Optional[str] = None
+
+# --- Global State and Initialization ---
+app = FastAPI(
+    title="BG3 Processing Manager",
+    description="Coordinates flow servers for BG3 dataset processing.",
+    version="1.0.0"
+)
+
+api = HfApi(token=HF_TOKEN)
+fs = HfFileSystem(token=HF_TOKEN)
+state = ProcessingState()
+is_coordinating = False
+
+# --- Persistence Functions ---
+
+def get_full_path(filename: str) -> str:
+    return f"{HF_DATASET_ID}/{filename}"
+
+async def load_state_from_hf():
+    global state
+    try:
+        # Check if state file exists
+        if fs.exists(f"{HF_DATASET_ID}/{STATE_FILE_PATH}"):
+            print(f"Loading state from {STATE_FILE_PATH}...")
+            
+            # Use HfApi to download the file content
+            content = api.read_file(
+                path_in_repo=STATE_FILE_PATH,
+                repo_id=HF_DATASET_ID,
+                repo_type=HF_DATASET_REPO_TYPE
+            ).decode('utf-8')
+            
+            data = json.loads(content)
+            state = ProcessingState(**data)
+            print(f"State loaded. Total tasks: {len(state.tasks)}")
+        else:
+            print(f"State file {STATE_FILE_PATH} not found. Initializing.")
+            await initialize_tasks()
+            await save_state_to_hf() # Save initial state
+            
+    except Exception as e:
+        print(f"Error loading state from HF: {e}")
+        # Fallback to initialization if loading fails
+        await initialize_tasks()
+        
+async def save_state_to_hf():
+    try:
+        print(f"Saving state to {STATE_FILE_PATH}...")
+        content = state.model_dump_json(indent=2).encode('utf-8')
+        
+        api.upload_file(
+            path_or_fileobj=content,
+            path_in_repo=STATE_FILE_PATH,
+            repo_id=HF_DATASET_ID,
+            repo_type=HF_DATASET_REPO_TYPE,
+            commit_message="Update processing state"
+        )
+        print("State saved successfully.")
+    except Exception as e:
+        print(f"Error saving state to HF: {e}")
+
+async def initialize_tasks():
+    global state
+    print(f"Discovering zip files in {FRAMES_FOLDER}/...")
+    
+    # List files in the frames folder
+    try:
+        file_list = fs.ls(f"{HF_DATASET_ID}/{FRAMES_FOLDER}", detail=False)
+        
+        zip_files = [
+            Path(f).name
+            for f in file_list
+            if f.endswith(".zip") and not f.endswith(".zip.json")
+        ]
+        
+        new_tasks = {}
+        for zip_file in zip_files:
+            course_name = zip_file.replace(".zip", "")
+            if course_name not in state.tasks:
+                new_tasks[course_name] = TaskStatus(status="UNPROCESSED")
+            else:
+                # Keep existing status if it was already tracked
+                new_tasks[course_name] = state.tasks[course_name]
+        
+        state.tasks = new_tasks
+        print(f"Found {len(zip_files)} zip files. Total tasks: {len(state.tasks)}")
+        
+    except Exception as e:
+        print(f"Error discovering files from HF: {e}")
+        # If discovery fails, we can't proceed.
+        raise RuntimeError(f"Failed to discover files: {e}")
+
+# --- Core Coordination Logic ---
+
+async def assign_next_task(flow_id: str):
+    """
+    Finds the next UNPROCESSED task and assigns it to the given flow server.
+    """
+    global state
+    
+    # 1. Find an UNPROCESSED task
+    next_course = None
+    for course_name, task_status in state.tasks.items():
+        if task_status.status == "UNPROCESSED":
+            next_course = course_name
+            break
+            
+    if next_course is None:
+        print(f"No UNPROCESSED tasks left for {flow_id}.")
+        course_to_assign = None
+        
+    else:
+        # 2. Update state to IN_PROGRESS
+        state.tasks[next_course] = TaskStatus(
+            status="IN_PROGRESS",
+            assigned_to=flow_id,
+            assigned_at=datetime.now()
+        )
+        state.flow_assignments[flow_id] = next_course
+        course_to_assign = next_course
+        
+        # 3. Persist state change
+        await save_state_to_hf()
+        
+    # 4. Notify the Flow Server
+    flow_url = FLOW_SERVERS.get(flow_id)
+    if not flow_url:
+        print(f"Error: Unknown flow_id {flow_id}")
+        return
+
+    try:
+        print(f"Assigning '{course_to_assign}' to {flow_id} at {flow_url}/process_course")
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{flow_url}/process_course",
+                json={"course_name": course_to_assign}
+            ) as response:
+                if response.status != 200:
+                    print(f"Error sending task to {flow_id}: {response.status} - {await response.text()}")
+                    # Revert state if assignment fails
+                    if next_course:
+                        state.tasks[next_course] = TaskStatus(status="UNPROCESSED")
+                        state.flow_assignments[flow_id] = None
+                        await save_state_to_hf()
+                else:
+                    print(f"Successfully assigned {course_to_assign} to {flow_id}.")
+                    
+    except aiohttp.ClientConnectorError as e:
+        print(f"Connection Error: Could not connect to {flow_id} at {flow_url}. Reverting task status. Error: {e}")
+        # Revert state if connection fails
+        if next_course:
+            state.tasks[next_course] = TaskStatus(status="UNPROCESSED")
+            state.flow_assignments[flow_id] = None
+            await save_state_to_hf()
+    except Exception as e:
+        print(f"Unexpected error during assignment to {flow_id}. Error: {e}")
+        # Revert state for safety
+        if next_course:
+            state.tasks[next_course] = TaskStatus(status="UNPROCESSED")
+            state.flow_assignments[flow_id] = None
+            await save_state_to_hf()
+
+
+async def coordinate_loop():
+    """
+    The main coordination loop that runs in the background.
+    """
+    global is_coordinating
+    if is_coordinating:
+        print("Coordinator is already running.")
+        return
+
+    is_coordinating = True
+    print("Starting coordination loop...")
+    
+    try:
+        # Load state and initialize tasks on startup
+        await load_state_from_hf()
+        
+        # Check and assign tasks to any free flow server
+        for flow_id in FLOW_SERVERS.keys():
+            if state.flow_assignments.get(flow_id) is None:
+                asyncio.create_task(assign_next_task(flow_id))
+                
+    except Exception as e:
+        print(f"Coordination loop failed to start: {e}")
+    finally:
+        # The loop is now event-driven based on /task/complete calls
+        pass
+
+# --- API Endpoints ---
+
+@app.on_event("startup")
+async def startup_event():
+    # Start the coordination loop as a background task
+    BackgroundTasks().add_task(coordinate_loop)
+    
+@app.get("/")
+async def root():
+    return {
+        "message": "BG3 Processing Manager API", 
+        "status": "running",
+        "is_coordinating": is_coordinating,
+        "flow_assignments": state.flow_assignments,
+        "total_tasks": len(state.tasks),
+        "unprocessed": sum(1 for t in state.tasks.values() if t.status == "UNPROCESSED"),
+        "in_progress": sum(1 for t in state.tasks.values() if t.status == "IN_PROGRESS"),
+        "completed": sum(1 for t in state.tasks.values() if t.status == "COMPLETED"),
+    }
+
+@app.post("/task/complete")
+async def task_complete(request: CompleteTaskRequest):
+    """
+    Endpoint for flow servers to report task completion.
+    """
+    global state
+    flow_id = request.flow_id
+    course_name = request.course_name
+    
+    if course_name not in state.tasks:
+        raise HTTPException(status_code=404, detail=f"Unknown course: {course_name}")
+        
+    task = state.tasks[course_name]
+    
+    if task.assigned_to != flow_id:
+        # This is a safety check, should not happen in normal operation
+        print(f"Warning: {flow_id} reported completion for a task not assigned to it: {course_name}")
+        
+    if request.success:
+        print(f"Task COMPLETED: {course_name} by {flow_id}")
+        task.status = "COMPLETED"
+        task.completed_at = datetime.now()
+        task.error_message = None
+    else:
+        print(f"Task FAILED: {course_name} by {flow_id}. Error: {request.error_message}")
+        # For now, mark as FAILED. A more robust system might retry.
+        task.status = "FAILED"
+        task.completed_at = datetime.now()
+        task.error_message = request.error_message
+        
+    # Free up the flow server slot
+    state.flow_assignments[flow_id] = None
+    
+    # Persist state change
+    await save_state_to_hf()
+    
+    # Assign the next task to the now-free flow server
+    asyncio.create_task(assign_next_task(flow_id))
+    
+    return {"status": "success", "message": f"Task {course_name} marked as {'COMPLETED' if request.success else 'FAILED'}. Next task assigned."}
+
+@app.post("/start_coordination")
+async def start_coordination(background_tasks: BackgroundTasks):
+    """
+    Manually trigger the coordination loop.
+    """
+    if is_coordinating:
+        return {"status": "info", "message": "Coordination is already running."}
+        
+    background_tasks.add_task(coordinate_loop)
+    return {"status": "success", "message": "Coordination loop started."}
+
+@app.get("/state")
+async def get_state():
+    """
+    Returns the current processing state.
+    """
+    return state
+
+if __name__ == "__main__":
+    # Note: When running in the sandbox, we need to use 0.0.0.0 to expose the port.
+    uvicorn.run(app, host="0.0.0.0", port=MANAGER_PORT)