import asyncio import aiohttp import json import os from pathlib import Path from typing import Dict, List, Set import logging from datetime import datetime from suno_utils.utils.text import read_jsonl # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class AsyncFileProcessor: def __init__( self, output_dir: str = "processed_chunks", chunk_size: int = 1000, max_concurrent_requests: int = 50, progress_file: str = "processing_progress.json" ): self.output_dir = Path(output_dir) self.output_dir.mkdir(exist_ok=True) self.chunk_size = chunk_size self.max_concurrent_requests = max_concurrent_requests self.progress_file = progress_file # Load existing progress self.processed_files = self._load_progress() def _load_progress(self) -> Set[str]: """Load previously processed files from progress file.""" if os.path.exists(self.progress_file): try: with open(self.progress_file, 'r') as f: data = json.load(f) return set(data.get('processed_files', [])) except Exception as e: logger.warning(f"Could not load progress file: {e}") return set() def _save_progress(self, processed_files: Set[str]): """Save current progress to file.""" progress_data = { 'processed_files': list(processed_files), 'last_updated': datetime.now().isoformat(), 'total_processed': len(processed_files) } with open(self.progress_file, 'w') as f: json.dump(progress_data, f, indent=2) def _save_chunk_results(self, chunk_index: int, results: Dict[str, dict]): """Save a chunk of results to a file.""" chunk_file = self.output_dir / f"chunk_{chunk_index:06d}.json" with open(chunk_file, 'w') as f: json.dump(results, f, indent=2) logger.info(f"Saved chunk {chunk_index} with {len(results)} results to {chunk_file}") async def _process_single_file( self, session: aiohttp.ClientSession, s3_path: str, semaphore: asyncio.Semaphore ) -> tuple[str, dict]: """Process a single file with the API endpoint.""" async with semaphore: # Limit concurrent requests params = { "gen_id": s3_path, "min_loop_length_bars": 2, "max_loop_length_bars": 4, "output_stems": "false" } max_retries = 3 for attempt in range(max_retries): try: async with session.get( "https://suno-ai--loop-extraction-data-extract-loop-points.modal.run", params=params, timeout=aiohttp.ClientTimeout(total=600) # 5 minute timeout ) as response: if response.status == 200: response_json = await response.json() return s3_path, response_json else: logger.warning(f"HTTP {response.status} for {s3_path}") if attempt == max_retries - 1: return s3_path, {"error": f"HTTP {response.status}"} except asyncio.TimeoutError: logger.warning(f"Timeout for {s3_path} (attempt {attempt + 1})") if attempt == max_retries - 1: return s3_path, {"error": "timeout"} except Exception as e: logger.warning(f"Error processing {s3_path}: {e} (attempt {attempt + 1})") if attempt == max_retries - 1: return s3_path, {"error": str(e)} # Wait before retry if attempt < max_retries - 1: await asyncio.sleep(2 ** attempt) # Exponential backoff async def _process_chunk( self, session: aiohttp.ClientSession, chunk_files: List[str], chunk_index: int, semaphore: asyncio.Semaphore ) -> Dict[str, dict]: """Process a chunk of files concurrently.""" logger.info(f"Processing chunk {chunk_index} with {len(chunk_files)} files") # Create tasks for all files in the chunk tasks = [ self._process_single_file(session, s3_path, semaphore) for s3_path in chunk_files ] # Wait for all tasks in the chunk to complete results = await asyncio.gather(*tasks, return_exceptions=True) # Process results and handle exceptions chunk_results = {} for result in results: if isinstance(result, Exception): logger.error(f"Task failed with exception: {result}") continue s3_path, response_data = result chunk_results[s3_path] = response_data return chunk_results async def process_all_files(self, stems_to_process: List[str]): """Process all files with chunking and resume capability.""" # Filter out already processed files remaining_files = [f for f in stems_to_process if f not in self.processed_files] if not remaining_files: logger.info("All files have already been processed!") return logger.info(f"Processing {len(remaining_files)} remaining files out of {len(stems_to_process)} total") logger.info(f"Already processed: {len(self.processed_files)} files") # Create chunks chunks = [ remaining_files[i:i + self.chunk_size] for i in range(0, len(remaining_files), self.chunk_size) ] # Calculate starting chunk index based on existing files existing_chunks = len([f for f in self.output_dir.glob("chunk_*.json")]) # Create semaphore to limit concurrent requests semaphore = asyncio.Semaphore(self.max_concurrent_requests) # Create aiohttp session with connection limits connector = aiohttp.TCPConnector( limit=self.max_concurrent_requests * 2, limit_per_host=self.max_concurrent_requests ) async with aiohttp.ClientSession(connector=connector) as session: for i, chunk_files in enumerate(chunks): chunk_index = existing_chunks + i try: # Process the chunk chunk_results = await self._process_chunk( session, chunk_files, chunk_index, semaphore ) # Save chunk results self._save_chunk_results(chunk_index, chunk_results) # Update progress self.processed_files.update(chunk_files) self._save_progress(self.processed_files) logger.info(f"Completed chunk {chunk_index}. Total processed: {len(self.processed_files)}") except Exception as e: logger.error(f"Failed to process chunk {chunk_index}: {e}") # Continue with next chunk even if one fails logger.info(f"Processing complete! Processed {len(self.processed_files)} files total") def combine_all_results(self, output_file: str = "combined_results.json") -> Dict[str, dict]: """Combine all chunk files into a single result dictionary.""" combined_results = {} chunk_files = sorted(self.output_dir.glob("chunk_*.json")) for chunk_file in chunk_files: try: with open(chunk_file, 'r') as f: chunk_data = json.load(f) combined_results.update(chunk_data) logger.info(f"Loaded {len(chunk_data)} results from {chunk_file}") except Exception as e: logger.error(f"Failed to load {chunk_file}: {e}") # Save combined results with open(output_file, 'w') as f: json.dump(combined_results, f, indent=2) logger.info(f"Combined {len(combined_results)} results into {output_file}") return combined_results # Usage example async def main(): extreme_meta = read_jsonl("/home/sara/sfx/extreme_stems_consolidated_analysis.jsonl") stems_to_process = [] for meta in extreme_meta: if not meta["is_mostly_silent"] and meta["duration_s"] < 300: stems_to_process.append(str(meta["s3_path"])) # Create processor with custom settings processor = AsyncFileProcessor( output_dir="processed_chunks", chunk_size=1000, # Adjust based on your needs max_concurrent_requests=1000, # Adjust based on API limits progress_file="processing_progress.json" ) # Process all files await processor.process_all_files(stems_to_process) # Optionally combine all results into one file all_results = processor.combine_all_results("final_results.json") print(f"Processing complete! Total results: {len(all_results)}") # Run the async function if __name__ == "__main__": asyncio.run(main())