dataset-uploader/scripts/process_datasets_parallel.py

#!/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()