sdg-hub 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

sdg_hub/core/flow/base.py

@@ -4,18 +4,26 @@
 # Standard
 from pathlib import Path
 from typing import Any, Optional, Union
+import time
 
 # Third Party
 from datasets import Dataset
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+from rich.tree import Tree
 import yaml
 
 # Local
 from ..blocks.base import BaseBlock
 from ..blocks.registry import BlockRegistry
+from ..utils.datautils import safe_concatenate_with_validation
 from ..utils.error_handling import EmptyDatasetError, FlowValidationError
 from ..utils.logger_config import setup_logger
 from ..utils.path_resolution import resolve_path
+from ..utils.yaml_utils import save_flow_yaml
+from .checkpointer import FlowCheckpointer
 from .metadata import FlowMetadata, FlowParameter
 from .migration import FlowMigration
 from .validation import FlowValidator
@@ -133,7 +141,17 @@ class Flow(BaseModel):
         -------
         Flow
             Validated Flow instance.
+
+        Raises
+        ------
+        FlowValidationError
+            If yaml_path is None or the file doesn't exist.
         """
+        if yaml_path is None:
+            raise FlowValidationError(
+                "Flow path cannot be None. Please provide a valid YAML file path or check that the flow exists in the registry."
+            )
+
         yaml_path = resolve_path(yaml_path, [])
         yaml_dir = Path(yaml_path).parent
 
@@ -160,6 +178,8 @@ class Flow(BaseModel):
             flow_config, migrated_runtime_params = FlowMigration.migrate_to_new_format(
                 flow_config, yaml_path
             )
+            # Save migrated config back to YAML to persist id
+            save_flow_yaml(yaml_path, flow_config, "migrated to new format")
 
         # Validate YAML structure
         validator = FlowValidator()
@@ -221,6 +241,17 @@ class Flow(BaseModel):
         # Create and validate the flow
         try:
             flow = cls(blocks=blocks, metadata=metadata, parameters=parameters)
+            # Persist generated id back to the YAML file (only on initial load)
+            # If the file had no metadata.id originally, update and rewrite
+            if not flow_config.get("metadata", {}).get("id"):
+                flow_config.setdefault("metadata", {})["id"] = flow.metadata.id
+                save_flow_yaml(
+                    yaml_path,
+                    flow_config,
+                    f"added generated id: {flow.metadata.id}",
+                )
+            else:
+                logger.debug(f"Flow already had id: {flow.metadata.id}")
             # Store migrated runtime params and client for backward compatibility
             if migrated_runtime_params:
                 flow._migrated_runtime_params = migrated_runtime_params
@@ -324,6 +355,8 @@ class Flow(BaseModel):
         self,
         dataset: Dataset,
         runtime_params: Optional[dict[str, dict[str, Any]]] = None,
+        checkpoint_dir: Optional[str] = None,
+        save_freq: Optional[int] = None,
     ) -> Dataset:
         """Execute the flow blocks in sequence to generate data.
 
@@ -340,6 +373,11 @@ class Flow(BaseModel):
                 "block_name": {"param1": value1, "param2": value2},
                 "other_block": {"param3": value3}
             }
+        checkpoint_dir : Optional[str], optional
+            Directory to save/load checkpoints. If provided, enables checkpointing.
+        save_freq : Optional[int], optional
+            Number of completed samples after which to save a checkpoint.
+            If None, only saves final results when checkpointing is enabled.
 
         Returns
         -------
@@ -353,6 +391,12 @@ class Flow(BaseModel):
         FlowValidationError
             If flow validation fails or if model configuration is required but not set.
         """
+        # Validate save_freq parameter early to prevent range() errors
+        if save_freq is not None and save_freq <= 0:
+            raise FlowValidationError(
+                f"save_freq must be greater than 0, got {save_freq}"
+            )
+
         # Validate preconditions
         if not self.blocks:
             raise FlowValidationError("Cannot generate with empty flow")
@@ -376,18 +420,119 @@ class Flow(BaseModel):
                 "Dataset validation failed:\n" + "\n".join(dataset_errors)
             )
 
+        # Initialize checkpointer if enabled
+        checkpointer = None
+        completed_dataset = None
+        if checkpoint_dir:
+            checkpointer = FlowCheckpointer(
+                checkpoint_dir=checkpoint_dir,
+                save_freq=save_freq,
+                flow_id=self.metadata.id,
+            )
+
+            # Load existing progress
+            remaining_dataset, completed_dataset = checkpointer.load_existing_progress(
+                dataset
+            )
+
+            if len(remaining_dataset) == 0:
+                logger.info("All samples already completed, returning existing results")
+                return completed_dataset
+
+            dataset = remaining_dataset
+            logger.info(f"Resuming with {len(dataset)} remaining samples")
+
         logger.info(
             f"Starting flow '{self.metadata.name}' v{self.metadata.version} "
             f"with {len(dataset)} samples across {len(self.blocks)} blocks"
         )
 
-        current_dataset = dataset
         # Merge migrated runtime params with provided ones (provided ones take precedence)
         merged_runtime_params = self._migrated_runtime_params.copy()
         if runtime_params:
             merged_runtime_params.update(runtime_params)
         runtime_params = merged_runtime_params
 
+        # Process dataset in chunks if checkpointing with save_freq
+        if checkpointer and save_freq:
+            all_processed = []
+
+            # Process in chunks of save_freq
+            for i in range(0, len(dataset), save_freq):
+                chunk_end = min(i + save_freq, len(dataset))
+                chunk_dataset = dataset.select(range(i, chunk_end))
+
+                logger.info(
+                    f"Processing chunk {i // save_freq + 1}: samples {i} to {chunk_end - 1}"
+                )
+
+                # Execute all blocks on this chunk
+                processed_chunk = self._execute_blocks_on_dataset(
+                    chunk_dataset, runtime_params
+                )
+                all_processed.append(processed_chunk)
+
+                # Save checkpoint after chunk completion
+                checkpointer.add_completed_samples(processed_chunk)
+
+            # Save final checkpoint for any remaining samples
+            checkpointer.save_final_checkpoint()
+
+            # Combine all processed chunks
+            final_dataset = safe_concatenate_with_validation(
+                all_processed, "processed chunks from flow execution"
+            )
+
+            # Combine with previously completed samples if any
+            if checkpointer and completed_dataset:
+                final_dataset = safe_concatenate_with_validation(
+                    [completed_dataset, final_dataset],
+                    "completed checkpoint data with newly processed data",
+                )
+
+        else:
+            # Process entire dataset at once
+            final_dataset = self._execute_blocks_on_dataset(dataset, runtime_params)
+
+            # Save final checkpoint if checkpointing enabled
+            if checkpointer:
+                checkpointer.add_completed_samples(final_dataset)
+                checkpointer.save_final_checkpoint()
+
+                # Combine with previously completed samples if any
+                if completed_dataset:
+                    final_dataset = safe_concatenate_with_validation(
+                        [completed_dataset, final_dataset],
+                        "completed checkpoint data with newly processed data",
+                    )
+
+        logger.info(
+            f"Flow '{self.metadata.name}' completed successfully: "
+            f"{len(final_dataset)} final samples, "
+            f"{len(final_dataset.column_names)} final columns"
+        )
+
+        return final_dataset
+
+    def _execute_blocks_on_dataset(
+        self, dataset: Dataset, runtime_params: dict[str, dict[str, Any]]
+    ) -> Dataset:
+        """Execute all blocks in sequence on the given dataset.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            Dataset to process through all blocks.
+        runtime_params : Dict[str, Dict[str, Any]]
+            Runtime parameters for block execution.
+
+        Returns
+        -------
+        Dataset
+            Dataset after processing through all blocks.
+        """
+        current_dataset = dataset
+
         # Execute blocks in sequence
         for i, block in enumerate(self.blocks):
             logger.info(
@@ -436,12 +581,6 @@ class Flow(BaseModel):
                     f"Block '{block.block_name}' execution failed: {exc}"
                 ) from exc
 
-        logger.info(
-            f"Flow '{self.metadata.name}' completed successfully: "
-            f"{len(current_dataset)} final samples, "
-            f"{len(current_dataset.column_names)} final columns"
-        )
-
         return current_dataset
 
     def _prepare_block_kwargs(
@@ -784,9 +923,6 @@ class Flow(BaseModel):
             "execution_time_seconds": 0,
         }
 
-        # Standard
-        import time
-
         start_time = time.time()
 
         try:
@@ -930,6 +1066,102 @@ class Flow(BaseModel):
             "block_names": [block.block_name for block in self.blocks],
         }
 
+    def print_info(self) -> None:
+        """
+        Print an interactive summary of the Flow in the console.
+
+        The summary contains:
+        1. Flow metadata (name, version, author, description)
+        2. Defined runtime parameters with type hints and defaults
+        3. A table of all blocks with their input and output columns
+
+        Notes
+        -----
+        Uses the `rich` library for colourised output; install with
+        `pip install rich` if not already present.
+
+        Returns
+        -------
+        None
+        """
+
+        console = Console()
+
+        # Create main tree structure
+        flow_tree = Tree(
+            f"[bold bright_blue]{self.metadata.name}[/bold bright_blue] Flow"
+        )
+
+        # Metadata section
+        metadata_branch = flow_tree.add(
+            "[bold bright_green]Metadata[/bold bright_green]"
+        )
+        metadata_branch.add(
+            f"Version: [bright_cyan]{self.metadata.version}[/bright_cyan]"
+        )
+        metadata_branch.add(
+            f"Author: [bright_cyan]{self.metadata.author}[/bright_cyan]"
+        )
+        if self.metadata.description:
+            metadata_branch.add(
+                f"Description: [white]{self.metadata.description}[/white]"
+            )
+
+        # Parameters section
+        if self.parameters:
+            params_branch = flow_tree.add(
+                "[bold bright_yellow]Parameters[/bold bright_yellow]"
+            )
+            for name, param in self.parameters.items():
+                param_info = f"[bright_cyan]{name}[/bright_cyan]: [white]{param.type_hint}[/white]"
+                if param.default is not None:
+                    param_info += f" = [bright_white]{param.default}[/bright_white]"
+                params_branch.add(param_info)
+
+        # Blocks overview
+        flow_tree.add(
+            f"[bold bright_magenta]Blocks[/bold bright_magenta] ({len(self.blocks)} total)"
+        )
+
+        # Create blocks table
+        blocks_table = Table(show_header=True, header_style="bold bright_white")
+        blocks_table.add_column("Block Name", style="bright_cyan")
+        blocks_table.add_column("Type", style="bright_green")
+        blocks_table.add_column("Input Cols", style="bright_yellow")
+        blocks_table.add_column("Output Cols", style="bright_red")
+
+        for block in self.blocks:
+            input_cols = getattr(block, "input_cols", None)
+            output_cols = getattr(block, "output_cols", None)
+
+            blocks_table.add_row(
+                block.block_name,
+                block.__class__.__name__,
+                str(input_cols) if input_cols else "[bright_black]None[/bright_black]",
+                str(output_cols)
+                if output_cols
+                else "[bright_black]None[/bright_black]",
+            )
+
+        # Print everything
+        console.print()
+        console.print(
+            Panel(
+                flow_tree,
+                title="[bold bright_white]Flow Information[/bold bright_white]",
+                border_style="bright_blue",
+            )
+        )
+        console.print()
+        console.print(
+            Panel(
+                blocks_table,
+                title="[bold bright_white]Block Details[/bold bright_white]",
+                border_style="bright_magenta",
+            )
+        )
+        console.print()
+
     def to_yaml(self, output_path: str) -> None:
         """Save flow configuration to YAML file.
 
@@ -952,10 +1184,7 @@ class Flow(BaseModel):
                 name: param.model_dump() for name, param in self.parameters.items()
             }
 
-        with open(output_path, "w", encoding="utf-8") as f:
-            yaml.dump(config, f, default_flow_style=False, sort_keys=False)
-
-        logger.info(f"Flow configuration saved to: {output_path}")
+        save_flow_yaml(output_path, config)
 
     def __len__(self) -> int:
         """Number of blocks in the flow."""

sdg_hub/core/flow/checkpointer.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flow-level checkpointing with sample-level tracking for data generation pipelines."""
+
+# Standard
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+import json
+import os
+import uuid
+
+# Third Party
+from datasets import Dataset
+
+# Local
+from ..utils.datautils import safe_concatenate_with_validation
+from ..utils.logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+class FlowCheckpointer:
+    """Enhanced checkpointer for Flow execution with sample-level tracking.
+
+    Provides data-level checkpointing where progress is saved after processing
+    a specified number of samples through the entire flow pipeline.
+    """
+
+    def __init__(
+        self,
+        checkpoint_dir: Optional[str] = None,
+        save_freq: Optional[int] = None,
+        flow_id: Optional[str] = None,
+    ):
+        """Initialize the FlowCheckpointer.
+
+        Parameters
+        ----------
+        checkpoint_dir : Optional[str]
+            Directory to save/load checkpoints. If None, checkpointing is disabled.
+        save_freq : Optional[int]
+            Number of completed samples after which to save a checkpoint.
+            If None, only final results are saved.
+        flow_id : Optional[str]
+            Unique ID of the flow for checkpoint identification.
+        """
+        self.checkpoint_dir = checkpoint_dir
+        self.save_freq = save_freq
+        self.flow_id = flow_id or "unknown_flow"
+
+        # Internal state
+        self._samples_processed = 0
+        self._checkpoint_counter = 0
+        self._pending_samples: List[Dict[str, Any]] = []
+
+        # Ensure checkpoint directory exists
+        if self.checkpoint_dir:
+            Path(self.checkpoint_dir).mkdir(parents=True, exist_ok=True)
+
+    @property
+    def is_enabled(self) -> bool:
+        """Check if checkpointing is enabled."""
+        return self.checkpoint_dir is not None
+
+    @property
+    def metadata_path(self) -> str:
+        """Path to the flow metadata file."""
+        return os.path.join(self.checkpoint_dir, "flow_metadata.json")
+
+    def load_existing_progress(
+        self, input_dataset: Dataset
+    ) -> Tuple[Dataset, Optional[Dataset]]:
+        """Load existing checkpoint data and determine remaining work.
+
+        Parameters
+        ----------
+        input_dataset : Dataset
+            Original input dataset for the flow.
+
+        Returns
+        -------
+        Tuple[Dataset, Optional[Dataset]]
+            (remaining_samples_to_process, completed_samples_dataset)
+            If no checkpoints exist, returns (input_dataset, None)
+        """
+        if not self.is_enabled:
+            return input_dataset, None
+
+        try:
+            # Load flow metadata
+            metadata = self._load_metadata()
+            if not metadata:
+                logger.info(f"No existing checkpoints found in {self.checkpoint_dir}")
+                return input_dataset, None
+
+            # Validate flow identity to prevent mixing checkpoints from different flows
+            saved_flow_id = metadata.get("flow_id")
+            if saved_flow_id and saved_flow_id != self.flow_id:
+                logger.warning(
+                    f"Flow ID mismatch: saved checkpoints are for flow ID '{saved_flow_id}' "
+                    f"but current flow ID is '{self.flow_id}'. Starting fresh to avoid "
+                    f"mixing incompatible checkpoint data."
+                )
+                return input_dataset, None
+
+            # Load all completed samples from checkpoints
+            completed_dataset = self._load_completed_samples()
+            if completed_dataset is None or len(completed_dataset) == 0:
+                logger.info("No completed samples found in checkpoints")
+                return input_dataset, None
+
+            # Find samples that still need processing
+            remaining_dataset = self._find_remaining_samples(
+                input_dataset, completed_dataset
+            )
+
+            self._samples_processed = len(completed_dataset)
+            self._checkpoint_counter = metadata.get("checkpoint_counter", 0)
+
+            logger.info(
+                f"Loaded {len(completed_dataset)} completed samples, "
+                f"{len(remaining_dataset)} samples remaining"
+            )
+
+            return remaining_dataset, completed_dataset
+
+        except Exception as exc:
+            logger.warning(f"Failed to load checkpoints: {exc}. Starting from scratch.")
+            return input_dataset, None
+
+    def add_completed_samples(self, samples: Dataset) -> None:
+        """Add samples that have completed the entire flow.
+
+        Parameters
+        ----------
+        samples : Dataset
+            Samples that have completed processing through all blocks.
+        """
+        if not self.is_enabled:
+            return
+
+        # Add to pending samples
+        for sample in samples:
+            self._pending_samples.append(sample)
+            self._samples_processed += 1
+
+            # Check if we should save a checkpoint
+            if self.save_freq and len(self._pending_samples) >= self.save_freq:
+                self._save_checkpoint()
+
+    def save_final_checkpoint(self) -> None:
+        """Save any remaining pending samples as final checkpoint."""
+        if not self.is_enabled:
+            return
+
+        if self._pending_samples:
+            sample_count = len(self._pending_samples)
+            self._save_checkpoint()
+            logger.info(f"Saved final checkpoint with {sample_count} samples")
+
+    def _save_checkpoint(self) -> None:
+        """Save current pending samples to a checkpoint file."""
+        if not self._pending_samples:
+            return
+
+        self._checkpoint_counter += 1
+        checkpoint_file = os.path.join(
+            self.checkpoint_dir, f"checkpoint_{self._checkpoint_counter:04d}.jsonl"
+        )
+
+        # Convert pending samples to dataset and save
+        checkpoint_dataset = Dataset.from_list(self._pending_samples)
+        checkpoint_dataset.to_json(checkpoint_file, orient="records", lines=True)
+
+        # Update metadata
+        self._save_metadata()
+
+        logger.info(
+            f"Saved checkpoint {self._checkpoint_counter} with "
+            f"{len(self._pending_samples)} samples to {checkpoint_file}"
+        )
+
+        # Clear pending samples
+        self._pending_samples.clear()
+
+    def _save_metadata(self) -> None:
+        """Save flow execution metadata."""
+        metadata = {
+            "flow_id": self.flow_id,
+            "save_freq": self.save_freq,
+            "samples_processed": self._samples_processed,
+            "checkpoint_counter": self._checkpoint_counter,
+            "last_updated": str(uuid.uuid4()),  # Simple versioning
+        }
+
+        with open(self.metadata_path, "w", encoding="utf-8") as f:
+            json.dump(metadata, f, indent=2)
+
+    def _load_metadata(self) -> Optional[Dict[str, Any]]:
+        """Load flow execution metadata."""
+        if not os.path.exists(self.metadata_path):
+            return None
+
+        try:
+            with open(self.metadata_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception as exc:
+            logger.warning(f"Failed to load metadata: {exc}")
+            return None
+
+    def _load_completed_samples(self) -> Optional[Dataset]:
+        """Load all completed samples from checkpoint files."""
+        checkpoint_files = []
+        checkpoint_dir = Path(self.checkpoint_dir)
+
+        # Find all checkpoint files
+        for file_path in checkpoint_dir.glob("checkpoint_*.jsonl"):
+            checkpoint_files.append(str(file_path))
+
+        if not checkpoint_files:
+            return None
+
+        # Sort checkpoint files by number
+        checkpoint_files.sort()
+
+        # Load and concatenate all checkpoint datasets
+        datasets = []
+        for file_path in checkpoint_files:
+            try:
+                dataset = Dataset.from_json(file_path)
+                if len(dataset) > 0:
+                    datasets.append(dataset)
+                    logger.debug(
+                        f"Loaded checkpoint: {file_path} ({len(dataset)} samples)"
+                    )
+            except Exception as exc:
+                logger.warning(f"Failed to load checkpoint {file_path}: {exc}")
+
+        if not datasets:
+            return None
+
+        return safe_concatenate_with_validation(datasets, "checkpoint files")
+
+    def _find_remaining_samples(
+        self, input_dataset: Dataset, completed_dataset: Dataset
+    ) -> Dataset:
+        """Find samples from input_dataset that are not in completed_dataset.
+
+        Note: Assumes input_dataset contains unique samples. For datasets with
+        duplicates, multiset semantics with collections.Counter would be needed.
+
+        Parameters
+        ----------
+        input_dataset : Dataset
+            Original input dataset (assumed to contain unique samples).
+        completed_dataset : Dataset
+            Dataset of completed samples.
+
+        Returns
+        -------
+        Dataset
+            Samples that still need processing.
+        """
+        # Get common columns for comparison
+        input_columns = set(input_dataset.column_names)
+        completed_columns = set(completed_dataset.column_names)
+        common_columns = list(input_columns & completed_columns)
+
+        if not common_columns:
+            logger.warning(
+                "No common columns found between input and completed datasets. "
+                "Processing all input samples."
+            )
+            return input_dataset
+
+        # Convert to pandas for easier comparison
+        input_df = input_dataset.select_columns(common_columns).to_pandas()
+        completed_df = completed_dataset.select_columns(common_columns).to_pandas()
+
+        # Find rows that haven't been completed
+        # Use tuple representation for comparison
+        input_tuples = set(input_df.apply(tuple, axis=1))
+        completed_tuples = set(completed_df.apply(tuple, axis=1))
+        remaining_tuples = input_tuples - completed_tuples
+
+        # Filter input dataset to only remaining samples
+        remaining_mask = input_df.apply(tuple, axis=1).isin(remaining_tuples)
+        remaining_indices = input_df[remaining_mask].index.tolist()
+
+        if not remaining_indices:
+            # Return empty dataset with same structure
+            return input_dataset.select([])
+
+        return input_dataset.select(remaining_indices)
+
+    def get_progress_info(self) -> Dict[str, Any]:
+        """Get information about current progress.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Progress information including samples processed, checkpoints saved, etc.
+        """
+        return {
+            "checkpoint_dir": self.checkpoint_dir,
+            "save_freq": self.save_freq,
+            "flow_id": self.flow_id,
+            "samples_processed": self._samples_processed,
+            "checkpoint_counter": self._checkpoint_counter,
+            "pending_samples": len(self._pending_samples),
+            "is_enabled": self.is_enabled,
+        }
+
+    def cleanup_checkpoints(self) -> None:
+        """Remove all checkpoint files and metadata."""
+        if not self.is_enabled:
+            return
+
+        checkpoint_dir = Path(self.checkpoint_dir)
+        if not checkpoint_dir.exists():
+            return
+
+        # Remove all checkpoint files
+        for file_path in checkpoint_dir.glob("checkpoint_*.jsonl"):
+            file_path.unlink()
+            logger.debug(f"Removed checkpoint file: {file_path}")
+
+        # Remove metadata file
+        metadata_path = Path(self.metadata_path)
+        if metadata_path.exists():
+            metadata_path.unlink()
+            logger.debug(f"Removed metadata file: {metadata_path}")
+
+        logger.info(f"Cleaned up all checkpoints in {self.checkpoint_dir}")