import os import re import time import json import tempfile import logging from typing import Optional, Dict, List, Tuple, Set from scipy.signal import resample import modal import funcy import numpy as np from botocore.exceptions import ClientError, NoCredentialsError from suno_utils.utils.text import read_jsonl from suno_utils.utils.s3 import list_s3_dir, read_from_s3 from suno_utils.worker.settings import s3_client from suno_utils.audio import Audio from suno_utils.worker.modal_base import get_modal_base_image # Configure logging logging.basicConfig(level=logging.WARN) logger = logging.getLogger(__name__) # Constants S3_BUCKET = "suno-data" BATCH_SIZE = 500 N_MAX_REPLICAS = 200 MAX_RETRIES = 3 MIN_DURATION_SECONDS = 2.0 MIN_REALISTIC_TEMPO = 30 MAX_REALISTIC_TEMPO = 300 BATCH_FILENAME_PATTERN = r'batch_(\d+)_results\.json' BATCH_PROGRESS_LOG_INTERVAL = 50 # Modal setup aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_dict({ "SUNO_ASSETS_PATH": "/suno/models/assets", "XDG_CACHE_HOME": "/suno/models/", }), ] image = ( get_modal_base_image() .pip_install( "git+https://github.com/marc-suno/madmom.git@bf7d502#egg=madmom", "git+https://github.com/CPJKU/beat_this@117ff34", "cvxpy==1.6.5", ) .add_local_python_source("suno_utils", copy=False) ) def make_json_serializable(obj) -> any: """Convert numpy types to JSON serializable types.""" if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, dict): return {key: make_json_serializable(value) for key, value in obj.items()} elif isinstance(obj, list): return [make_json_serializable(item) for item in obj] else: return obj def retry_s3_operation(operation_func, max_retries: int = MAX_RETRIES, base_delay: float = 1.0): """Retry S3 operations with exponential backoff.""" for attempt in range(max_retries): try: return operation_func() except (ClientError, NoCredentialsError, Exception) as e: if attempt == max_retries - 1: logger.error(f"S3 operation failed after {max_retries} attempts: {e}") raise delay = base_delay * (2 ** attempt) logger.warning(f"S3 operation failed (attempt {attempt + 1}/{max_retries}), " f"retrying in {delay}s: {e}") time.sleep(delay) def create_error_result(item_id: str, failure_reason: str, include_beat_times: bool = False) -> Dict: """Create a standardized error result.""" result = { "id": item_id, "inferred_tempo": None, "tempo_std": None, "beats_detected": None, "processing_failed": True, "failure_reason": failure_reason } if include_beat_times: result["beat_times"] = None return result # Modal app setup app = modal.App("audio-analysis-worker", image=image, secrets=SECRETS) @app.cls( cpu=8, gpu="A10G", memory=8000, secrets=SECRETS, timeout=30 * 60, container_idle_timeout=60, concurrency_limit=N_MAX_REPLICAS, ) class AudioAnalysisStub: def __init__(self, output_path: str, include_beat_times: bool = False): from suno_utils.tasks.audio_features.beat_this_downbeat import BeatThisDownbeatExtractor self.downbeat_extractor = BeatThisDownbeatExtractor( device="cuda", model_path="s3://suno-data/m4burns/beat_this_rc_12l.pt", resample_on_device=True ) self.output_path = output_path self.include_beat_times = include_beat_times def extract_beats(self, audio) -> np.ndarray: """Extract beats with robust error handling.""" try: audio_mono = Audio.convert( audio, n_channels=1, sample_rate=audio.sample_rate, byte_width=audio.byte_width ) out = self.downbeat_extractor.extract(audio_mono) # Validate extraction output if not isinstance(out, dict) or "downbeats" not in out: logger.warning("Beat extractor returned invalid format") return np.array([]) downbeats = out["downbeats"] if downbeats is None or len(downbeats) == 0: return np.array([]) # Convert to numpy array and handle different shapes beats_refined = np.array(downbeats) if beats_refined.ndim == 1: return beats_refined elif beats_refined.ndim == 2 and beats_refined.shape[1] > 0: return beats_refined[:, 0] else: logger.warning(f"Unexpected beat array shape: {beats_refined.shape}") return np.array([]) except Exception as e: logger.error(f"Beat extraction failed: {e}") return np.array([]) def calculate_tempo_stats(self, beats: np.ndarray) -> Tuple[Optional[int], Optional[float], int]: """Calculate tempo statistics with robust error handling.""" beats_detected = len(beats) if beats_detected < 2: return None, None, beats_detected try: # Calculate time differences and filter invalid ones beat_diffs = np.diff(beats) valid_diffs = beat_diffs[beat_diffs > 0] if len(valid_diffs) == 0: logger.warning("No valid beat intervals found") return None, None, beats_detected # Calculate and filter tempos tempos = 60.0 / valid_diffs realistic_tempos = tempos[ (tempos >= MIN_REALISTIC_TEMPO) & (tempos <= MAX_REALISTIC_TEMPO) ] if len(realistic_tempos) == 0: logger.warning("No realistic tempo values found") return None, None, beats_detected # Calculate final statistics tempo_std = float(np.std(realistic_tempos)) median_tempo = float(np.median(realistic_tempos)) # Validate results if not (np.isfinite(tempo_std) and np.isfinite(median_tempo)): logger.warning("Non-finite tempo statistics calculated") return None, None, beats_detected return int(round(median_tempo)), tempo_std, beats_detected except Exception as e: logger.error(f"Tempo calculation failed: {e}") return None, None, beats_detected def _analyze_single_audio(self, item_id: str, s3_filepath: str, duration_s: float) -> Dict: """Analyze a single audio file and return statistics.""" try: if duration_s < MIN_DURATION_SECONDS: result = { "id": item_id, "inferred_tempo": None, "tempo_std": None, "beats_detected": None, "processing_failed": False, "failure_reason": "duration_too_short" } if self.include_beat_times: result["beat_times"] = None return result # Load audio from S3 with retry logic audio = retry_s3_operation(lambda: Audio.from_s3(s3_filepath)) # Extract beats and calculate tempo statistics beats = self.extract_beats(audio) est_target_tempo, tempo_std, beats_detected = self.calculate_tempo_stats(beats) result = { "id": item_id, "inferred_tempo": est_target_tempo, "tempo_std": tempo_std, "beats_detected": beats_detected, "processing_failed": False, "failure_reason": None } if self.include_beat_times: # Store beat times as list (will be serialized by make_json_serializable) result["beat_times"] = beats if len(beats) > 0 else None return result except Exception as e: error_msg = str(e)[:200] # Truncate long error messages logger.error(f"Failed processing {item_id} ({s3_filepath}): {error_msg}") return create_error_result(item_id, error_msg, self.include_beat_times) def _save_batch_to_s3(self, results_dict: Dict, batch_index: int) -> None: """Save batch results to S3 with retry logic.""" def upload_operation(): with tempfile.TemporaryDirectory() as td: batch_filename = f"batch_{batch_index:04d}_results.json" batch_path = os.path.join(td, batch_filename) with open(batch_path, 'w') as f: json.dump(results_dict, f, indent=2) s3_batch_path = os.path.join(self.output_path, batch_filename) s3_client.upload_file( batch_path, S3_BUCKET, s3_batch_path, ExtraArgs={"ContentType": "application/json"} ) try: retry_s3_operation(upload_operation) logger.info(f"Batch {batch_index} results saved to S3") except Exception as e: logger.error(f"Failed to save batch {batch_index} to S3: {e}") raise @modal.method() def analyze_batch(self, work_items: List[Dict], batch_index: int) -> Dict: """Analyze a batch of audio files and return results keyed by ID.""" results_dict = {} batch_size = len(work_items) logger.info(f"Processing batch {batch_index} of {batch_size} files") for i, item in enumerate(work_items): try: # Extract and validate required fields item_id = str(item["id"]) s3_filepath = item["s3_filepath"] duration_s = float(item["duration_s"]) # Analyze the audio file result = self._analyze_single_audio(item_id, s3_filepath, duration_s) results_dict[item_id] = make_json_serializable(result) # Log progress periodically if (i + 1) % BATCH_PROGRESS_LOG_INTERVAL == 0 or i == batch_size - 1: logger.info(f"Batch {batch_index}: Processed {i + 1}/{batch_size} files") except KeyError as e: error_result = create_error_result( item.get("id", "unknown"), f"missing_field_{e}", self.include_beat_times ) results_dict[error_result["id"]] = error_result logger.error(f"Missing required field in work item: {e}") except ValueError as e: error_result = create_error_result( item.get("id", "unknown"), f"invalid_data_{e}", self.include_beat_times ) results_dict[error_result["id"]] = error_result logger.error(f"Invalid data type in work item: {e}") except Exception as e: error_result = create_error_result( item.get("id", "unknown"), f"unexpected_error_{str(e)[:100]}", self.include_beat_times ) results_dict[error_result["id"]] = error_result logger.error(f"Unexpected error processing work item: {e}") # Log batch statistics self._log_batch_stats(results_dict, batch_index) # Save batch results to S3 self._save_batch_to_s3(results_dict, batch_index) return results_dict def _log_batch_stats(self, results_dict: Dict, batch_index: int) -> None: """Log statistics for a completed batch.""" results_list = list(results_dict.values()) success_count = sum(1 for r in results_list if not r['processing_failed']) failure_count = len(results_list) - success_count success_rate = success_count / len(results_list) * 100 if results_list else 0 logger.info(f"Batch {batch_index} complete: {success_count}/{len(results_list)} " f"successful ({success_rate:.1f}%)") if failure_count > 0: failed_ids = [r['id'] for r in results_list if r['processing_failed']] display_ids = failed_ids[:5] if len(failed_ids) > 5: display_ids.append('...') logger.warning(f"Batch {batch_index} failed files: {display_ids}") def load_work_items(input_jsonl_path: str, max_files: Optional[int] = None) -> List[Dict]: """Load work items from JSONL file with optional limit.""" print(f"Loading metadata from: {input_jsonl_path}") try: if input_jsonl_path.startswith("s3://"): work_items = retry_s3_operation( lambda: read_from_s3(input_jsonl_path, read_f=read_jsonl) ) else: work_items = read_jsonl(input_jsonl_path) except Exception as e: print(f"Failed to load input file: {e}") raise print(f"Loaded {len(work_items)} total work items") if max_files is not None: work_items = work_items[:max_files] print(f"Limited to {len(work_items)} files") return work_items def get_existing_batch_indices(output_path: str) -> Set[int]: """Get set of existing batch indices from S3.""" try: existing_files = retry_s3_operation( lambda: list_s3_dir(f"s3://{S3_BUCKET}/{output_path}/") ) existing_batch_files = [f[0] for f in existing_files if f[0].endswith("_results.json")] existing_batch_indices = set() for batch_file_path in existing_batch_files: filename = os.path.basename(batch_file_path) match = re.match(BATCH_FILENAME_PATTERN, filename) if match: batch_index = int(match.group(1)) existing_batch_indices.add(batch_index) if existing_batch_indices: max_completed_batch = max(existing_batch_indices) print(f"Found {len(existing_batch_indices)} existing batch files") print(f"Latest completed batch: {max_completed_batch}") else: print("No existing batch files found") return existing_batch_indices except Exception as e: print(f"Warning: Could not check existing batch results: {e}") print("Proceeding without resume...") return set() def filter_batches_to_process(work_batches: List[List[Dict]], existing_batch_indices: Set[int]) -> List[Tuple[int, List[Dict]]]: """Filter out batches that already exist and return (index, batch) tuples.""" batches_to_process = [] skipped_batches = 0 for i, batch in enumerate(work_batches): if i in existing_batch_indices: skipped_batches += 1 else: batches_to_process.append((i, batch)) if skipped_batches > 0: print(f"Skipping {skipped_batches} existing batches") print(f"Will process {len(batches_to_process)} new batches") return batches_to_process def calculate_final_statistics(batch_results: List[Dict], existing_batch_count: int) -> Tuple[int, int, int, Dict[str, int]]: """Calculate final processing statistics.""" processed_files = 0 processed_successful = 0 processed_failed = 0 failure_reasons = {} for batch_result_dict in batch_results: for result in batch_result_dict.values(): processed_files += 1 if result['processing_failed']: processed_failed += 1 reason = result.get('failure_reason', 'unknown') failure_reasons[reason] = failure_reasons.get(reason, 0) + 1 else: processed_successful += 1 return processed_files, processed_successful, processed_failed, failure_reasons def print_final_summary(processed_files: int, processed_successful: int, processed_failed: int, existing_batch_count: int, processing_time: float, output_path: str, failure_reasons: Dict[str, int]) -> None: """Print final processing summary.""" estimated_existing_files = existing_batch_count * BATCH_SIZE total_files = processed_files + estimated_existing_files total_successful = processed_successful + estimated_existing_files total_success_rate = total_successful / total_files * 100 if total_files > 0 else 0 print(f"\n=== ANALYSIS COMPLETE ===") print(f"Processed in this run: {processed_files} files " f"({processed_successful} successful, {processed_failed} failed)") print(f"Existing from previous runs: ~{estimated_existing_files} files (estimated)") print(f"Total across all runs: ~{total_files} files") print(f"Overall success rate: {total_success_rate:.2f}%") print(f"Processing time: {processing_time/60:.1f} minutes") print(f"Results saved to:") print(f" S3: s3://{S3_BUCKET}/{output_path}/") print(f" Files: batch_0000_results.json, batch_0001_results.json, ...") if processed_failed > 0: print(f"New failure breakdown: {dict(failure_reasons)}") print(f"\nTo access results programmatically:") print(f" # Download all batch files") print(f" # Each file contains: {{\"id1\": {{result}}, \"id2\": {{result}}, ...}}") print(f" # aws s3 sync s3://{S3_BUCKET}/{output_path}/ ./results/ " f"--exclude '*' --include '*_results.json'") @app.local_entrypoint() def main( input_jsonl_path: str, output_path: str, max_files: Optional[int] = None, resume: bool = False, include_beat_times: bool = False ) -> None: """ Main entry point for audio analysis. Args: input_jsonl_path: Path to JSONL file containing audio metadata output_path: S3 path for storing results (e.g., "sara/sfx_data_beats") max_files: Optional limit on number of files to process resume: Whether to skip batches that already have results include_beat_times: Whether to store full beat time arrays (increases data size significantly) """ # Load work items work_items = load_work_items(input_jsonl_path, max_files) if not work_items: print("No files to process!") return # Check for existing batches (resume functionality) existing_batch_indices = get_existing_batch_indices(output_path) if resume else set() # Create and filter batches work_batches = list(funcy.chunks(BATCH_SIZE, work_items)) print(f"Created {len(work_batches)} total batches (batch size: {BATCH_SIZE})") batches_to_process = filter_batches_to_process(work_batches, existing_batch_indices) if not batches_to_process: print("All batches already processed!") return # Initialize worker and process batches worker = AudioAnalysisStub(output_path, include_beat_times) print("Starting batch processing...") start_time = time.time() # Use Modal's map for parallel execution across multiple workers batch_data = [(batch_index, batch) for batch_index, batch in batches_to_process] # Extract batches and indices for parallel processing batch_indices = [item[0] for item in batch_data] batches = [item[1] for item in batch_data] # Process batches in parallel using Modal's map functionality batch_results = list(worker.analyze_batch.map(batches, batch_indices)) processing_time = time.time() - start_time # Calculate and display final statistics processed_files, processed_successful, processed_failed, failure_reasons = \ calculate_final_statistics(batch_results, len(existing_batch_indices)) print_final_summary( processed_files, processed_successful, processed_failed, len(existing_batch_indices), processing_time, output_path, failure_reasons )