#!/usr/bin/env python3 """Process multiple RDF datasets in parallel with rich progress display. This script provides parallel processing capabilities for downloading, converting, and uploading RDF datasets with clear progress tracking for each dataset. USAGE: # Process multiple datasets in parallel python scripts/process_datasets_parallel.py --dataset wordnet --dataset yago-4.5 # Process all small datasets with 4 parallel workers python scripts/process_datasets_parallel.py --category small --workers 4 # Process with specific steps only python scripts/process_datasets_parallel.py --dataset wordnet \ --dataset schema-org --skip-upload # Dry run to see what would happen python scripts/process_datasets_parallel.py --category small --dry-run FEATURES: - Parallel processing with configurable worker count - Rich progress display showing all datasets simultaneously - Automatic resource management based on dataset size - Graceful error handling and recovery - Memory-efficient streaming for large datasets """ from __future__ import annotations import argparse import sys import time from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed from pathlib import Path # Import our modules from dataset_registry import DATASET_REGISTRY from rich.console import Console from rich.live import Live from rich.progress import ( BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TaskID, TextColumn, TimeElapsedColumn, ) from rich.table import Table console = Console() class ParallelDatasetProcessor: """Manages parallel processing of multiple datasets with progress tracking.""" def __init__( self, base_dir: Path, max_workers: int = 4, skip_download: bool = False, skip_convert: bool = False, skip_upload: bool = False, dry_run: bool = False, remove_after: bool = False, ): """Initialize the parallel processor. Args: base_dir: Base directory for dataset processing max_workers: Maximum number of parallel workers skip_download: Skip download step skip_convert: Skip conversion step skip_upload: Skip upload step dry_run: Show what would be done without executing remove_after: Remove downloaded files after processing """ self.base_dir = base_dir self.max_workers = max_workers self.skip_download = skip_download self.skip_convert = skip_convert self.skip_upload = skip_upload self.dry_run = dry_run self.remove_after = remove_after # Create progress tracking self.overall_progress = Progress( SpinnerColumn(), TextColumn("[bold blue]Overall Progress"), BarColumn(), MofNCompleteColumn(), TimeElapsedColumn(), ) # Individual dataset progress trackers self.dataset_progress = Progress( TextColumn("[bold cyan]{task.fields[dataset_id]:20}"), SpinnerColumn(), TextColumn("{task.description}"), BarColumn(), TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), TimeElapsedColumn(), expand=True, ) # Results tracking self.results: dict[str, tuple[bool, str]] = {} def process_dataset_task( self, dataset_id: str, task_id: TaskID ) -> tuple[str, bool, str]: """Process a single dataset (runs in worker process/thread). Args: dataset_id: Dataset identifier task_id: Progress task ID for this dataset Returns: Tuple of (dataset_id, success, message) """ import shutil import subprocess try: dataset_info = DATASET_REGISTRY.get(dataset_id) if not dataset_info: self.dataset_progress.update( task_id, description="[red]✗ Unknown dataset[/red]", completed=100 ) return (dataset_id, False, f"Unknown dataset: {dataset_id}") if not dataset_info.available: self.dataset_progress.update( task_id, description="[yellow]⊘ Unavailable[/yellow]", completed=100 ) return (dataset_id, False, "Dataset marked as unavailable") # Setup directories download_dir = self.base_dir / "downloads" / dataset_id hf_dataset_dir = self.base_dir / "hf_datasets" / dataset_id # Calculate progress steps total_steps = 3 - ( self.skip_download + self.skip_convert + self.skip_upload ) if total_steps == 0: total_steps = 1 step_progress = 100 / total_steps current_progress = 0 # Step 1: Download if not self.skip_download: self.dataset_progress.update( task_id, description="[yellow]📥 Downloading...[/yellow]", completed=current_progress, total=100, ) if not self.dry_run: result = subprocess.run( [ sys.executable, str(Path(__file__).parent / "rdf_dataset_downloader.py"), dataset_id, "-o", str(download_dir), ], capture_output=True, text=True, ) if result.returncode != 0: self.dataset_progress.update( task_id, description="[red]✗ Download failed[/red]", completed=100, ) return ( dataset_id, False, f"Download failed: {result.stderr[:100]}", ) current_progress += step_progress self.dataset_progress.update(task_id, completed=current_progress) # Step 2: Convert if not self.skip_convert: self.dataset_progress.update( task_id, description="[cyan]🔄 Converting to HF...[/cyan]", completed=current_progress, total=100, ) if not self.dry_run: # Find the downloaded file rdf_files = ( list(download_dir.rglob("*.ttl")) + list(download_dir.rglob("*.nt")) + list(download_dir.rglob("*.rdf")) + list(download_dir.rglob("*.owl")) ) if dataset_info.format == "tsv": rdf_files += list(download_dir.rglob("*.txt")) + list( download_dir.rglob("*.tsv") ) if not rdf_files: self.dataset_progress.update( task_id, description="[red]✗ No RDF files found[/red]", completed=100, ) return (dataset_id, False, "No RDF files found after download") rdf_file = rdf_files[0] # Determine converter based on dataset size use_streaming = dataset_info.category in [ "medium", "large", "xlarge", ] converter_script = ( "convert_rdf_to_hf_dataset_streaming.py" if use_streaming else "convert_rdf_to_hf_dataset.py" ) result = subprocess.run( [ sys.executable, str(Path(__file__).parent / converter_script), str(rdf_file), str(hf_dataset_dir), "--format", dataset_info.format, ], capture_output=True, text=True, ) if result.returncode != 0: self.dataset_progress.update( task_id, description="[red]✗ Conversion failed[/red]", completed=100, ) return ( dataset_id, False, f"Conversion failed: {result.stderr[:100]}", ) current_progress += step_progress self.dataset_progress.update(task_id, completed=current_progress) # Step 3: Upload if not self.skip_upload: self.dataset_progress.update( task_id, description="[green]📤 Uploading to HF...[/green]", completed=current_progress, total=100, ) if not self.dry_run: # Here you would add the actual upload logic # For now, just simulate with a delay time.sleep(1) current_progress += step_progress self.dataset_progress.update(task_id, completed=current_progress) # Cleanup if requested if self.remove_after and not self.dry_run and download_dir.exists(): shutil.rmtree(download_dir) self.dataset_progress.update( task_id, description="[green]✓ Complete[/green]", completed=100, total=100, ) return (dataset_id, True, "Successfully processed") except Exception as e: self.dataset_progress.update( task_id, description=f"[red]✗ Failed: {str(e)[:30]}[/red]", completed=100, ) return (dataset_id, False, str(e)) def process_datasets(self, dataset_ids: list[str]) -> dict[str, tuple[bool, str]]: """Process multiple datasets in parallel. Args: dataset_ids: List of dataset IDs to process Returns: Dictionary mapping dataset_id to (success, message) """ # Filter to available datasets only available_datasets = [ d_id for d_id in dataset_ids if DATASET_REGISTRY.get(d_id) and DATASET_REGISTRY[d_id].available ] if not available_datasets: console.print("[red]No available datasets to process[/red]") return {} console.print( f"\n[bold cyan]Processing {len(available_datasets)} datasets " f"with {self.max_workers} workers[/bold cyan]\n" ) # Create task IDs for each dataset task_ids = {} for dataset_id in available_datasets: task_id = self.dataset_progress.add_task( "[dim]Waiting...[/dim]", dataset_id=dataset_id, total=100 ) task_ids[dataset_id] = task_id # Create overall progress task overall_task = self.overall_progress.add_task( "Processing datasets", total=len(available_datasets) ) # Create progress table def generate_progress_table() -> Table: """Generate a table showing all progress bars.""" table = Table(title="Dataset Processing Progress", expand=True) table.add_column("Progress", ratio=1) # Add overall progress table.add_row(self.overall_progress) table.add_row("") # Spacer # Add individual dataset progress table.add_row(self.dataset_progress) return table # Process datasets in parallel with live progress display with Live( generate_progress_table(), console=console, refresh_per_second=4 ) as live: # Determine executor type based on operations # Use ProcessPoolExecutor for CPU-intensive conversion # Use ThreadPoolExecutor for I/O-bound download/upload if self.skip_convert: executor_class = ThreadPoolExecutor else: executor_class = ProcessPoolExecutor with executor_class(max_workers=self.max_workers) as executor: # Submit all tasks futures = { executor.submit( self.process_dataset_task, dataset_id, task_ids[dataset_id] ): dataset_id for dataset_id in available_datasets } # Process completed tasks for completed, future in enumerate(as_completed(futures), start=1): dataset_id = futures[future] try: dataset_id, success, message = future.result() self.results[dataset_id] = (success, message) except Exception as e: self.results[dataset_id] = (False, str(e)) self.overall_progress.update(overall_task, completed=completed) # Update live display live.update(generate_progress_table()) return self.results def print_summary(self): """Print a summary of processing results.""" console.print("\n[bold cyan]Processing Summary:[/bold cyan]\n") table = Table(show_header=True) table.add_column("Dataset", style="cyan") table.add_column("Status", style="green") table.add_column("Message", style="yellow") successful = 0 failed = 0 for dataset_id, (success, message) in self.results.items(): if success: status = "[green]✓ Success[/green]" successful += 1 else: status = "[red]✗ Failed[/red]" failed += 1 # Truncate message for display display_message = message[:50] + "..." if len(message) > 50 else message table.add_row(dataset_id, status, display_message) console.print(table) console.print(f"\n[bold]Total: {successful} successful, {failed} failed[/bold]") return successful, failed def main(): """Main entry point for parallel dataset processing.""" parser = argparse.ArgumentParser( description="Process RDF datasets in parallel with progress tracking" ) parser.add_argument( "--dataset", "-d", type=str, action="append", help="Specific dataset(s) to process (can be specified multiple times)", ) parser.add_argument( "--category", "-c", choices=["small", "medium", "large", "xlarge", "all"], help="Process datasets by category", ) parser.add_argument( "--workers", "-w", type=int, default=4, help="Number of parallel workers (default: 4)", ) parser.add_argument( "--base-dir", type=Path, default=Path("./dataset_processing"), help="Base directory for work", ) parser.add_argument( "--skip-download", action="store_true", help="Skip download step", ) parser.add_argument( "--skip-convert", action="store_true", help="Skip conversion step", ) parser.add_argument( "--skip-upload", action="store_true", help="Skip upload step", ) parser.add_argument( "--dry-run", action="store_true", help="Show what would be done without doing it", ) parser.add_argument( "--rm", action="store_true", help="Remove downloaded files after processing", ) parser.add_argument( "--list", action="store_true", help="List datasets that would be processed", ) args = parser.parse_args() # Determine which datasets to process if args.dataset: datasets_to_process = args.dataset elif args.category: if args.category == "all": datasets_to_process = list(DATASET_REGISTRY.keys()) else: datasets_to_process = [ d_id for d_id, d_info in DATASET_REGISTRY.items() if d_info.category == args.category ] else: parser.error("Please specify either --dataset or --category") # List mode if args.list: console.print("\n[bold cyan]Datasets to Process:[/bold cyan]\n") table = Table(show_header=True) table.add_column("ID", style="cyan") table.add_column("Name", style="green") table.add_column("Category", style="yellow") table.add_column("Size", style="blue") table.add_column("Available", style="magenta") for dataset_id in datasets_to_process: if dataset_id in DATASET_REGISTRY: info = DATASET_REGISTRY[dataset_id] available = "[green]✓[/green]" if info.available else "[red]✗[/red]" table.add_row( dataset_id, info.name, info.category, f"{info.size_gb} GB", available, ) console.print(table) available_count = sum( 1 for d_id in datasets_to_process if d_id in DATASET_REGISTRY and DATASET_REGISTRY[d_id].available ) console.print( f"\n[bold]Total: {len(datasets_to_process)} datasets " f"({available_count} available)[/bold]" ) return # Create processor and run processor = ParallelDatasetProcessor( base_dir=args.base_dir, max_workers=args.workers, skip_download=args.skip_download, skip_convert=args.skip_convert, skip_upload=args.skip_upload, dry_run=args.dry_run, remove_after=args.rm, ) # Process datasets start_time = time.time() processor.process_datasets(datasets_to_process) elapsed = time.time() - start_time # Print summary _successful, failed = processor.print_summary() console.print(f"\n[bold]Processing completed in {elapsed:.1f} seconds[/bold]") # Exit with appropriate code sys.exit(0 if failed == 0 else 1) if __name__ == "__main__": main()