#!/usr/bin/env python3 """Test pipeline for data monitoring on 100k rows.""" import sys import time import json import logging from pathlib import Path from typing import Dict, Any, List import argparse # Add parent directory to path sys.path.append(str(Path(__file__).parent.parent.parent)) from sunodata.data_monitor.analyzers.id_analyzer import IDAnalyzer from sunodata.data_monitor.analyzers.tags_analyzer import TagsAnalyzer from sunodata.data_monitor.analyzers.text_analyzer import TextAnalyzer from sunodata.data_monitor.analyzers.lang_analyzer import LanguageAnalyzer from sunodata.data_monitor.analyzers.weight_analyzer import WeightAnalyzer from sunodata.data_monitor.analyzers.stems_analyzer import StemsAnalyzer from sunodata.data_monitor.analyzers.paths_analyzer import PathsAnalyzer from sunodata.data_monitor.analyzers.lists_analyzer import ListsAnalyzer from sunodata.data_monitor.analyzers.sample_collector import SampleCollector from sunodata.data_monitor.utils.file_reader import FileReader from sunodata.data_monitor.utils.cache_manager import CacheManager from sunodata.data_monitor.utils.chunk_processor import ChunkProcessor # Setup logging logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) def load_config(config_path: str = "config.yaml") -> Dict: """Load configuration from YAML file.""" import yaml config_path = Path(__file__).parent / config_path with open(config_path, "r") as f: return yaml.safe_load(f) def test_sequential(file_path: Path, sample_size: int, output_dir: Path) -> Dict[str, Any]: """Test all analyzers sequentially on sample data. Args: file_path: Path to dataset file sample_size: Number of records to process output_dir: Directory for outputs Returns: Results from all analyzers """ logger.info(f"Testing sequential processing on {sample_size:,} records") # Initialize analyzers analyzers = [ IDAnalyzer(output_dir), TagsAnalyzer(output_dir), TextAnalyzer(output_dir), LanguageAnalyzer(output_dir), WeightAnalyzer(output_dir), StemsAnalyzer(output_dir), PathsAnalyzer(output_dir), ListsAnalyzer(output_dir), SampleCollector(output_dir), ] # Read sample records reader = FileReader(file_path) logger.info(f"Reading {sample_size:,} records from {file_path.name}...") records = list(reader.read_records(end=sample_size)) logger.info(f"Read {len(records):,} records") # Process with each analyzer results = {} total_time = 0 for analyzer in analyzers: logger.info(f"Running {analyzer.name} analyzer...") start_time = time.time() try: # Process as single chunk chunk_result = analyzer.process_chunk(records) # Aggregate (with single chunk) final_result = analyzer.aggregate([chunk_result]) results[analyzer.name] = final_result elapsed = time.time() - start_time total_time += elapsed logger.info(f" Completed in {elapsed:.2f}s") # Save individual results analyzer.save_results(final_result, suffix="_test_sequential") except Exception as e: logger.error(f" Failed: {e}") results[analyzer.name] = {"error": str(e)} # Add metadata results["_metadata"] = { "mode": "sequential", "file_path": str(file_path), "sample_size": sample_size, "records_processed": len(records), "total_time": total_time, "analyzers": [a.name for a in analyzers], } return results def test_parallel( file_path: Path, sample_size: int, output_dir: Path, chunk_size: int = 10000, num_workers: int = 4 ) -> Dict[str, Any]: """Test parallel processing on sample data. Args: file_path: Path to dataset file sample_size: Number of records to process output_dir: Directory for outputs chunk_size: Records per chunk num_workers: Number of parallel workers Returns: Results from all analyzers """ logger.info(f"Testing parallel processing on {sample_size:,} records") logger.info(f" Chunk size: {chunk_size:,}, Workers: {num_workers}") # Create temporary file with sample data temp_file = Path(output_dir) / "temp_sample.jsonl" reader = FileReader(file_path) logger.info(f"Creating temporary sample file...") with open(temp_file, "w") as f: for record_dict in reader.read_records(end=sample_size): f.write(json.dumps(record_dict, ensure_ascii=False) + "\n") # Initialize components analyzers = [ IDAnalyzer(output_dir), TagsAnalyzer(output_dir), TextAnalyzer(output_dir), LanguageAnalyzer(output_dir), WeightAnalyzer(output_dir), StemsAnalyzer(output_dir), PathsAnalyzer(output_dir), ListsAnalyzer(output_dir), SampleCollector(output_dir), ] # Process with ChunkProcessor processor = ChunkProcessor(temp_file, chunk_size=chunk_size, num_workers=num_workers) start_time = time.time() results = processor.process_file(analyzers, output_dir) elapsed = time.time() - start_time logger.info(f"Parallel processing completed in {elapsed:.2f}s") # Save results for analyzer_name, analyzer_results in results.items(): if analyzer_name.startswith("_"): continue for analyzer in analyzers: if analyzer.name == analyzer_name: analyzer.save_results(analyzer_results, suffix="_test_parallel") break # Update metadata results["_metadata"]["mode"] = "parallel" results["_metadata"]["chunk_size"] = chunk_size results["_metadata"]["num_workers"] = num_workers results["_metadata"]["processing_time"] = elapsed # Clean up temp file temp_file.unlink() return results def compare_results(seq_results: Dict, par_results: Dict) -> Dict[str, Any]: """Compare sequential and parallel results for consistency. Args: seq_results: Results from sequential processing par_results: Results from parallel processing Returns: Comparison report """ logger.info("Comparing sequential and parallel results...") comparison = { "identical": [], "differences": {}, "summary": {"analyzers_compared": 0, "identical_count": 0, "different_count": 0}, } # Compare each analyzer's results for analyzer_name in seq_results: if analyzer_name.startswith("_"): continue comparison["summary"]["analyzers_compared"] += 1 if analyzer_name not in par_results: comparison["differences"][analyzer_name] = "Missing in parallel results" comparison["summary"]["different_count"] += 1 continue seq_data = seq_results[analyzer_name] par_data = par_results[analyzer_name] # Compare summaries if "summary" in seq_data and "summary" in par_data: seq_summary = seq_data["summary"] par_summary = par_data["summary"] is_identical = True diffs = [] for key in seq_summary: if key in par_summary: # Allow small floating point differences if isinstance(seq_summary[key], (int, float)) and isinstance( par_summary[key], (int, float) ): if abs(seq_summary[key] - par_summary[key]) > 0.0001: diffs.append(f"{key}: seq={seq_summary[key]}, par={par_summary[key]}") is_identical = False elif seq_summary[key] != par_summary[key]: diffs.append(f"{key}: seq={seq_summary[key]}, par={par_summary[key]}") is_identical = False if is_identical: comparison["identical"].append(analyzer_name) comparison["summary"]["identical_count"] += 1 else: comparison["differences"][analyzer_name] = diffs comparison["summary"]["different_count"] += 1 else: comparison["identical"].append(analyzer_name) comparison["summary"]["identical_count"] += 1 return comparison def main(): """Main test pipeline.""" parser = argparse.ArgumentParser(description="Test data monitoring pipeline") parser.add_argument("--dataset", default="v9", help="Dataset version to test") parser.add_argument("--sample-size", type=int, default=100000, help="Number of records to test") parser.add_argument( "--chunk-size", type=int, default=10000, help="Chunk size for parallel processing" ) parser.add_argument("--workers", type=int, default=4, help="Number of parallel workers") parser.add_argument("--skip-sequential", action="store_true", help="Skip sequential test") parser.add_argument("--skip-parallel", action="store_true", help="Skip parallel test") parser.add_argument("--skip-comparison", action="store_true", help="Skip comparison") args = parser.parse_args() # Load configuration config = load_config() dataset_config = config["datasets"][args.dataset] train_file = Path(dataset_config["train"]) # Setup output directory output_base = Path.home() / "data" / "suno_data_monitor" / "outputs" output_dir = output_base / f"test_{args.dataset}_{time.strftime('%Y%m%d_%H%M%S')}" output_dir.mkdir(parents=True, exist_ok=True) logger.info(f"Testing dataset: {args.dataset}") logger.info(f"File: {train_file}") logger.info(f"Output directory: {output_dir}") # Run sequential test seq_results = None if not args.skip_sequential: logger.info("\n" + "=" * 50) logger.info("SEQUENTIAL TEST") logger.info("=" * 50) seq_results = test_sequential(train_file, args.sample_size, output_dir) # Save sequential results (clean non-serializable objects) def clean_for_json(obj): if isinstance(obj, dict): return {k: clean_for_json(v) for k, v in obj.items() if k != "id_set"} elif isinstance(obj, list): return [clean_for_json(item) for item in obj] elif isinstance(obj, set): return list(obj)[:1000] else: return obj with open(output_dir / "results_sequential.json", "w") as f: json.dump(clean_for_json(seq_results), f, indent=2, ensure_ascii=False) # Run parallel test par_results = None if not args.skip_parallel: logger.info("\n" + "=" * 50) logger.info("PARALLEL TEST") logger.info("=" * 50) par_results = test_parallel( train_file, args.sample_size, output_dir, chunk_size=args.chunk_size, num_workers=args.workers, ) # Save parallel results with open(output_dir / "results_parallel.json", "w") as f: json.dump(clean_for_json(par_results), f, indent=2, ensure_ascii=False) # Compare results if not args.skip_comparison and seq_results and par_results: logger.info("\n" + "=" * 50) logger.info("COMPARISON") logger.info("=" * 50) comparison = compare_results(seq_results, par_results) # Save comparison with open(output_dir / "comparison.json", "w") as f: json.dump(comparison, f, indent=2) # Print summary logger.info(f"Comparison Summary:") logger.info(f" Analyzers compared: {comparison['summary']['analyzers_compared']}") logger.info(f" Identical: {comparison['summary']['identical_count']}") logger.info(f" Different: {comparison['summary']['different_count']}") if comparison["differences"]: logger.warning("Differences found:") for name, diffs in comparison["differences"].items(): logger.warning(f" {name}: {diffs}") else: logger.info("✓ All results are consistent!") logger.info(f"\nTest complete. Results saved to: {output_dir}") if __name__ == "__main__": main()