Update app.py

app.py

@@ -396,10 +396,13 @@ async def process_batch_dynamic(wav_files: List[str], start_batch_index: int, ba
     """
     Processes a batch of WAV files in parallel using available servers.
     Batch size = number of servers. Each server gets one file, processes it, then gets the next.
+    Includes retry mechanism for failed files.
     Returns (next_batch_index, uploaded_count)
     """
     batch_end = min(start_batch_index + batch_size, len(wav_files))
     uploaded_count = progress.get('uploaded_count', 0)
+    max_retries = 3
+    failed_files = []  # Track files that failed for retry
     
     print(f"[{FLOW_ID}] Processing batch from index {start_batch_index} to {batch_end - 1} ({batch_end - start_batch_index} files)")
 
@@ -408,8 +411,7 @@ async def process_batch_dynamic(wav_files: List[str], start_batch_index: int, ba
         state.setdefault("file_states", {})
         for idx in range(start_batch_index, batch_end):
             wav_file = wav_files[idx]
-            wav_name = Path(wav_file).name
-            state["file_states"][wav_name] = "processing"
+            state["file_states"][wav_file] = "processing"
 
         # Update next_download_index to the end of this batch (0-based)
         state["next_download_index"] = batch_end
@@ -424,39 +426,52 @@ async def process_batch_dynamic(wav_files: List[str], start_batch_index: int, ba
         print(f"[{FLOW_ID}] Error while setting up batch locks: {e}")
         return start_batch_index, uploaded_count
 
+    # Create a queue of files to process with retry support
+    files_to_process = [(idx, wav_files[idx], 0) for idx in range(start_batch_index, batch_end)]  # (idx, wav_file, retry_count)
+
     # --- Assign files to servers and create tasks
-    pending_tasks: Dict[asyncio.Task, Tuple[int, Path, WhisperServer, str]] = {}
+    pending_tasks: Dict[asyncio.Task, Tuple[int, Path, WhisperServer, str, int]] = {}
     
     try:
-        for idx in range(start_batch_index, batch_end):
-            file_index = idx
-            wav_file = wav_files[file_index]
-            wav_filename = Path(wav_file).name
-            server = servers[idx - start_batch_index]  # Assign server in order (server 0 -> file 0, server 1 -> file 1, etc.)
-            
-            # Download the WAV file
-            wav_path = await download_wav_file_by_index(file_index + 1, wav_file)
-            if not wav_path:
-                state["file_states"][wav_file] = "failed"
-                await upload_hf_state(state)
-                print(f"[{FLOW_ID}] ❌ Failed to download {wav_filename}")
-                continue
+        while files_to_process or pending_tasks:
+            # Assign new files to available servers
+            while files_to_process:
+                available = await get_available_servers()
+                if not available:
+                    break
+                
+                file_idx, wav_file, retry_count = files_to_process.pop(0)
+                wav_filename = Path(wav_file).name
+                server = available[0]
+                
+                # Download the WAV file
+                wav_path = await download_wav_file_by_index(file_idx + 1, wav_file)
+                if not wav_path:
+                    if retry_count < max_retries:
+                        print(f"[{FLOW_ID}] ⚠️ Download failed for {wav_filename} (retry {retry_count + 1}/{max_retries}), re-queueing...")
+                        files_to_process.append((file_idx, wav_file, retry_count + 1))
+                    else:
+                        state["file_states"][wav_file] = "failed_download"
+                        print(f"[{FLOW_ID}] ❌ Download failed permanently for {wav_filename} after {max_retries} retries")
+                    continue
+                
+                # Assign to server and create task
+                await assign_file_to_server(file_idx, server)
+                task = asyncio.create_task(send_audio_to_whisper(wav_path, server))
+                pending_tasks[task] = (file_idx, wav_path, server, wav_file, retry_count)
+                print(f"[{FLOW_ID}] Assigned {wav_filename} to server {servers.index(server) + 1}")
             
-            # Assign to server and create task
-            await assign_file_to_server(file_index, server)
-            task = asyncio.create_task(send_audio_to_whisper(wav_path, server))
-            pending_tasks[task] = (file_index, wav_path, server, wav_file)
-            print(f"[{FLOW_ID}] Assigned {wav_filename} to server {servers.index(server) + 1}")
-        
-        # --- Wait for all tasks in this batch to complete and process results
-        while pending_tasks:
+            # Wait for at least one task to complete if there are pending tasks
+            if not pending_tasks:
+                break
+                
             done, pending = await asyncio.wait(
                 pending_tasks.keys(),
                 return_when=asyncio.FIRST_COMPLETED
             )
             
             for task in done:
-                file_index, wav_path, server, wav_file = pending_tasks.pop(task)
+                file_idx, wav_path, server, wav_file, retry_count = pending_tasks.pop(task)
                 wav_filename = Path(wav_file).name
                 
                 try:
@@ -473,15 +488,29 @@ async def process_batch_dynamic(wav_files: List[str], start_batch_index: int, ba
                             save_progress(progress)
                             print(f"[{FLOW_ID}] ✅ {wav_filename} uploaded (#{uploaded_count})")
                         else:
-                            state["file_states"][wav_file] = "failed"
-                            print(f"[{FLOW_ID}] ❌ Failed to upload {wav_filename}")
+                            # Retry failed upload
+                            if retry_count < max_retries:
+                                print(f"[{FLOW_ID}] ⚠️ Upload failed for {wav_filename} (retry {retry_count + 1}/{max_retries}), re-queueing...")
+                                files_to_process.append((file_idx, wav_file, retry_count + 1))
+                            else:
+                                state["file_states"][wav_file] = "failed_upload"
+                                print(f"[{FLOW_ID}] ❌ Upload failed permanently for {wav_filename} after {max_retries} retries")
                     else:
-                        state["file_states"][wav_file] = "failed"
-                        print(f"[{FLOW_ID}] ❌ Transcription failed for {wav_filename}")
+                        # Retry failed transcription
+                        if retry_count < max_retries:
+                            print(f"[{FLOW_ID}] ⚠️ Transcription failed for {wav_filename} (retry {retry_count + 1}/{max_retries}), re-queueing...")
+                            files_to_process.append((file_idx, wav_file, retry_count + 1))
+                        else:
+                            state["file_states"][wav_file] = "failed_transcription"
+                            print(f"[{FLOW_ID}] ❌ Transcription failed permanently for {wav_filename} after {max_retries} retries")
                         
                 except Exception as e:
-                    print(f"[{FLOW_ID}] Error processing result for {wav_filename}: {e}")
-                    state["file_states"][wav_file] = "failed"
+                    if retry_count < max_retries:
+                        print(f"[{FLOW_ID}] ⚠️ Error processing {wav_filename}: {e} (retry {retry_count + 1}/{max_retries}), re-queueing...")
+                        files_to_process.append((file_idx, wav_file, retry_count + 1))
+                    else:
+                        print(f"[{FLOW_ID}] ❌ Error processing {wav_filename}: {e} (failed after {max_retries} retries)")
+                        state["file_states"][wav_file] = "failed_error"
                 finally:
                     # Release the server
                     await release_server(server)