Update app.py

app.py

@@ -394,202 +394,186 @@ async def release_server(server: WhisperServer):
 
 async def process_batch_dynamic(wav_files: List[str], start_batch_index: int, batch_size: int, state: Dict[str, Any], progress: Dict) -> Tuple[int, int]:
     """
-    Dynamically processes a batch of WAV files using available servers.
+    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.
     Returns (next_batch_index, uploaded_count)
     """
     batch_end = min(start_batch_index + batch_size, len(wav_files))
-    current_index = start_batch_index
     uploaded_count = progress.get('uploaded_count', 0)
     
-    # Create tasks for all servers to process files dynamically
-    pending_tasks: Dict[asyncio.Task, Tuple[int, Path, WhisperServer]] = {}
-    
-    print(f"[{FLOW_ID}] Processing batch from index {start_batch_index} to {batch_end}")
+    print(f"[{FLOW_ID}] Processing batch from index {start_batch_index} to {batch_end - 1} ({batch_end - start_batch_index} files)")
 
     # --- Batch-level locking: mark all files in this batch as 'processing' and upload state
     try:
+        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.setdefault("file_states", {})
-            # Only set to processing if it's not already processed/processing
-            if state["file_states"].get(wav_name) not in ("processing", "processed"):
-                state["file_states"][wav_name] = "processing"
+            state["file_states"][wav_name] = "processing"
 
-        # Advance the next_download_index to the end of this batch (1-based index)
+        # Update next_download_index to the end of this batch (0-based)
         state["next_download_index"] = batch_end
 
         # Upload HF state to establish locks for this batch
         if await upload_hf_state(state):
-            print(f"[{FLOW_ID}] ✅ Batch lock uploaded for indices {start_batch_index}..{batch_end - 1}")
+            print(f"[{FLOW_ID}] ✅ Batch locked: files {start_batch_index}-{batch_end - 1} marked 'processing'")
         else:
-            print(f"[{FLOW_ID}] ❌ Failed to upload batch lock for indices {start_batch_index}..{batch_end - 1}")
+            print(f"[{FLOW_ID}] ❌ Failed to upload batch lock")
+            return start_batch_index, uploaded_count
     except Exception as e:
         print(f"[{FLOW_ID}] Error while setting up batch locks: {e}")
+        return start_batch_index, uploaded_count
+
+    # --- Assign files to servers and create tasks
+    pending_tasks: Dict[asyncio.Task, Tuple[int, Path, WhisperServer, str]] = {}
     
     try:
-        while current_index < batch_end or pending_tasks:
-            # Assign new files to available servers
-            while current_index < batch_end:
-                available_servers = await get_available_servers()
-                
-                if not available_servers:
-                    # All servers busy, wait a bit
-                    await asyncio.sleep(0.5)
-                    continue
-                
-                server = available_servers[0]
-                file_index = current_index
-                wav_file = wav_files[file_index]
-                wav_filename = Path(wav_file).name
-                
-                # 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_filename] = "failed"
-                    # Persist failure to HF
-                    await upload_hf_state(state)
-                    current_index += 1
-                    continue
-                
-                # 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)
-                
-                current_index += 1
+        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.)
             
-            # Wait for at least one task to complete
-            if pending_tasks:
-                done, pending_tasks_remaining = await asyncio.wait(
-                    pending_tasks.keys(),
-                    return_when=asyncio.FIRST_COMPLETED
-                )
+            # 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_filename] = "failed"
+                await upload_hf_state(state)
+                print(f"[{FLOW_ID}] ❌ Failed to download {wav_filename}")
+                continue
+            
+            # 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_filename)
+            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:
+            done, pending = await asyncio.wait(
+                pending_tasks.keys(),
+                return_when=asyncio.FIRST_COMPLETED
+            )
+            
+            for task in done:
+                file_index, wav_path, server, wav_filename = pending_tasks.pop(task)
                 
-                # Process completed tasks
-                for task in done:
-                    file_index, wav_path, server = pending_tasks.pop(task)
-                    wav_filename = Path(wav_path).name
+                try:
+                    transcription_result = task.result()
                     
-                    try:
-                        transcription_result = task.result()
-                        
-                        if transcription_result:
-                            # Upload transcription (immediately)
-                            uploaded_ok = await upload_transcription_to_hf(wav_filename, transcription_result)
-                            if uploaded_ok:
-                                state["file_states"][wav_filename] = "processed"
-                                uploaded_count += 1
-                                progress['uploaded_count'] = uploaded_count
-                                save_progress(progress)
-                            else:
-                                state["file_states"][wav_filename] = "failed"
-                            # Persist state change for this file immediately
-                            await upload_hf_state(state)
+                    if transcription_result:
+                        # Upload transcription immediately
+                        uploaded_ok = await upload_transcription_to_hf(wav_filename, transcription_result)
+                        if uploaded_ok:
+                            state["file_states"][wav_filename] = "processed"
+                            uploaded_count += 1
+                            progress['uploaded_count'] = uploaded_count
+                            save_progress(progress)
+                            print(f"[{FLOW_ID}] ✅ {wav_filename} uploaded (#{uploaded_count})")
                         else:
                             state["file_states"][wav_filename] = "failed"
-                            await upload_hf_state(state)
-                            
-                    except Exception as e:
-                        print(f"[{FLOW_ID}] Error processing result for {wav_filename}: {e}")
+                            print(f"[{FLOW_ID}] ❌ Failed to upload {wav_filename}")
+                    else:
                         state["file_states"][wav_filename] = "failed"
-                    finally:
-                        # Release the server
-                        await release_server(server)
-                        # Clean up the WAV file
-                        if wav_path.exists():
-                            wav_path.unlink()
-                
-                # Update pending_tasks with remaining
-                pending_tasks = {task: pending_tasks[task] for task in pending_tasks_remaining}
-            
-            # Update HF state periodically
-            await upload_hf_state(state)
-    
+                        print(f"[{FLOW_ID}] ❌ Transcription failed for {wav_filename}")
+                    
+                    # Persist state change for this file to HF
+                    await upload_hf_state(state)
+                        
+                except Exception as e:
+                    print(f"[{FLOW_ID}] Error processing result for {wav_filename}: {e}")
+                    state["file_states"][wav_filename] = "failed"
+                    await upload_hf_state(state)
+                finally:
+                    # Release the server
+                    await release_server(server)
+                    # Clean up the WAV file
+                    if wav_path.exists():
+                        wav_path.unlink()
     except Exception as e:
         print(f"[{FLOW_ID}] Error in process_batch_dynamic: {e}")
+        
+        return batch_end, uploaded_count
     
-    return current_index, uploaded_count
-
-async def process_dataset_task(start_index: int):
-    """Main task to process the dataset using dynamic server assignment."""
+    async def process_dataset_task(start_index: int):
+        """Main task to process the dataset using dynamic server assignment."""
     
-    # Load both local progress and HF state
-    progress = load_progress()
-    current_state = await download_hf_state()
-    file_list = await get_audio_file_list(progress)
+        # Load both local progress and HF state
+        progress = load_progress()
+        current_state = await download_hf_state()
+        file_list = await get_audio_file_list(progress)
     
-    if not file_list:
-        print(f"[{FLOW_ID}] ERROR: Cannot proceed. File list is empty.")
-        return False
+        if not file_list:
+            print(f"[{FLOW_ID}] ERROR: Cannot proceed. File list is empty.")
+            return False
     
-    # Ensure start_index is within bounds
-    if start_index > len(file_list):
-        print(f"[{FLOW_ID}] WARNING: Start index {start_index} is greater than the total number of files ({len(file_list)}). Exiting.")
-        return True
+        # Ensure start_index is within bounds
+        if start_index > len(file_list):
+            print(f"[{FLOW_ID}] WARNING: Start index {start_index} is greater than the total number of files ({len(file_list)}). Exiting.")
+            return True
     
-    # Determine the actual starting index in the 0-indexed list
-    start_list_index = start_index - 1
+        # Determine the actual starting index in the 0-indexed list
+        start_list_index = start_index - 1
     
-    print(f"[{FLOW_ID}] Starting audio transcription from file index: {start_index} out of {len(file_list)}.")
-    print(f"[{FLOW_ID}] Using {len(servers)} Whisper servers for dynamic processing.")
-    print(f"[{FLOW_ID}] Upload pause enabled: {UPLOAD_PAUSE_ENABLED}, Max uploads before pause: {MAX_UPLOADS_BEFORE_PAUSE}")
+        print(f"[{FLOW_ID}] Starting audio transcription from file index: {start_index} out of {len(file_list)}.")
+        print(f"[{FLOW_ID}] Using {len(servers)} Whisper servers for dynamic processing.")
+        print(f"[{FLOW_ID}] Upload pause enabled: {UPLOAD_PAUSE_ENABLED}, Max uploads before pause: {MAX_UPLOADS_BEFORE_PAUSE}")
     
-    # Initialize progress tracking
-    if 'uploaded_count' not in progress:
-        progress['uploaded_count'] = 0
+        # Initialize progress tracking
+        if 'uploaded_count' not in progress:
+            progress['uploaded_count'] = 0
     
-    # If there was no HF state in the repo, upload a fresh initial state file
-    try:
-        if not current_state.get("file_states") and current_state.get("next_download_index", 0) == 0:
-            print(f"[{FLOW_ID}] No HF state detected; uploading initial state file to {HF_OUTPUT_DATASET_ID}...")
-            # Ensure structure
-            current_state.setdefault("file_states", {})
-            current_state.setdefault("next_download_index", 0)
-            if await upload_hf_state(current_state):
-                print(f"[{FLOW_ID}] ✅ Initial HF state uploaded.")
-            else:
-                print(f"[{FLOW_ID}] ❌ Failed to upload initial HF state.")
-    except Exception as e:
-        print(f"[{FLOW_ID}] Error while uploading initial HF state: {e}")
-    global_success = True
-    current_batch_index = start_list_index
-    batch_size = len(servers) * 2  # Process 2 batches per server at a time
+        # If there was no HF state in the repo, upload a fresh initial state file
+        try:
+            if not current_state.get("file_states") and current_state.get("next_download_index", 0) == 0:
+                print(f"[{FLOW_ID}] No HF state detected; uploading initial state file to {HF_OUTPUT_DATASET_ID}...")
+                # Ensure structure
+                current_state.setdefault("file_states", {})
+                current_state.setdefault("next_download_index", 0)
+                if await upload_hf_state(current_state):
+                    print(f"[{FLOW_ID}] ✅ Initial HF state uploaded.")
+                else:
+                    print(f"[{FLOW_ID}] ❌ Failed to upload initial HF state.")
+        except Exception as e:
+            print(f"[{FLOW_ID}] Error while uploading initial HF state: {e}")
+        global_success = True
+        current_batch_index = start_list_index
+        batch_size = len(servers)  # Batch size = number of servers (20 files per batch)
     
-    try:
-        while current_batch_index < len(file_list):
-            # Process a batch dynamically
-            next_index, uploaded_count = await process_batch_dynamic(
-                file_list, 
-                current_batch_index, 
-                batch_size, 
-                current_state, 
-                progress
-            )
-            
-            # Update progress
-            progress['last_processed_index'] = next_index
-            progress['uploaded_count'] = uploaded_count
-            save_progress(progress)
-            
-            # Update current batch index
-            current_batch_index = next_index
+        try:
+            while current_batch_index < len(file_list):
+                # Process a batch dynamically
+                next_index, uploaded_count = await process_batch_dynamic(
+                    file_list, 
+                    current_batch_index, 
+                    batch_size, 
+                    current_state, 
+                    progress
+                )
+                
+                # Update progress
+                progress['last_processed_index'] = next_index
+                progress['uploaded_count'] = uploaded_count
+                save_progress(progress)
+                
+                # Update current batch index
+                current_batch_index = next_index
+                
+                # Log statistics
+                print(f"[{FLOW_ID}] Batch complete. Progress: {current_batch_index}/{len(file_list)}, Uploaded: {uploaded_count}")
+                
+                # Print server statistics
+                print(f"[{FLOW_ID}] Server Statistics:")
+                for i, server in enumerate(servers):
+                    print(f"  Server {i+1}: {server.total_processed} files, {server.total_time:.2f}s total, {server.fps:.2f} files/sec")
             
-            # Log statistics
-            print(f"[{FLOW_ID}] Batch complete. Progress: {current_batch_index}/{len(file_list)}, Uploaded: {uploaded_count}")
+            print(f"[{FLOW_ID}] All files processed successfully!")
+            return True
             
-            # Print server statistics
-            print(f"[{FLOW_ID}] Server Statistics:")
-            for i, server in enumerate(servers):
-                print(f"  Server {i+1}: {server.total_processed} files, {server.total_time:.2f}s total, {server.fps:.2f} files/sec")
-        
-        print(f"[{FLOW_ID}] All files processed successfully!")
-        return True
-        
-    except Exception as e:
-        print(f"[{FLOW_ID}] Critical error in process_dataset_task: {e}")
-        global_success = False
-        return global_success
+        except Exception as e:
+            print(f"[{FLOW_ID}] Critical error in process_dataset_task: {e}")
+            global_success = False
+            return global_success
 
 # --- FastAPI App and Endpoints ---