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

sdg_hub/core/flow/base.py

@@ -4,19 +4,28 @@
 # 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 datasets
 import yaml
 
 # Local
 from ..blocks.base import BaseBlock
 from ..blocks.registry import BlockRegistry
+from ..utils.datautils import safe_concatenate_with_validation, validate_no_duplicates
 from ..utils.error_handling import EmptyDatasetError, FlowValidationError
 from ..utils.logger_config import setup_logger
 from ..utils.path_resolution import resolve_path
-from .metadata import FlowMetadata, FlowParameter
+from ..utils.yaml_utils import save_flow_yaml
+from .checkpointer import FlowCheckpointer
+from .metadata import DatasetRequirements, FlowMetadata, FlowParameter
 from .migration import FlowMigration
 from .validation import FlowValidator
 
@@ -133,7 +142,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 +179,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 +242,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
@@ -275,13 +307,11 @@ class Flow(BaseModel):
 
         # Get block class from registry
         try:
-            block_class = BlockRegistry.get(block_type_name)
+            block_class = BlockRegistry._get(block_type_name)
         except KeyError as exc:
             # Get all available blocks from all categories
-            all_blocks = BlockRegistry.all()
-            available_blocks = ", ".join(
-                [block for blocks in all_blocks.values() for block in blocks]
-            )
+            all_blocks = BlockRegistry.list_blocks()
+            available_blocks = ", ".join(all_blocks)
             raise FlowValidationError(
                 f"Block type '{block_type_name}' not found in registry. "
                 f"Available blocks: {available_blocks}"
@@ -324,6 +354,9 @@ 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,
+        max_concurrency: Optional[int] = None,
     ) -> Dataset:
         """Execute the flow blocks in sequence to generate data.
 
@@ -340,6 +373,14 @@ 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.
+        max_concurrency : Optional[int], optional
+            Maximum number of concurrent requests across all blocks.
+            Controls async request concurrency to prevent overwhelming servers.
 
         Returns
         -------
@@ -353,6 +394,26 @@ 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 max_concurrency parameter
+        if max_concurrency is not None:
+            # Explicitly reject boolean values (bool is a subclass of int in Python)
+            if isinstance(max_concurrency, bool) or not isinstance(
+                max_concurrency, int
+            ):
+                raise FlowValidationError(
+                    f"max_concurrency must be an int, got {type(max_concurrency).__name__}"
+                )
+            if max_concurrency <= 0:
+                raise FlowValidationError(
+                    f"max_concurrency must be greater than 0, got {max_concurrency}"
+                )
+
         # Validate preconditions
         if not self.blocks:
             raise FlowValidationError("Cannot generate with empty flow")
@@ -360,6 +421,8 @@ class Flow(BaseModel):
         if len(dataset) == 0:
             raise EmptyDatasetError("Input dataset is empty")
 
+        validate_no_duplicates(dataset)
+
         # Check if model configuration has been set for flows with LLM blocks
         llm_blocks = self._detect_llm_blocks()
         if llm_blocks and not self._model_config_set:
@@ -376,18 +439,131 @@ class Flow(BaseModel):
                 "Dataset validation failed:\n" + "\n".join(dataset_errors)
             )
 
+        # Log concurrency control if specified
+        if max_concurrency is not None:
+            logger.info(f"Using max_concurrency={max_concurrency} for LLM requests")
+
+        # 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"
+            + (f" (max_concurrency={max_concurrency})" if max_concurrency else "")
         )
 
-        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, max_concurrency
+                )
+                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, max_concurrency
+            )
+
+            # 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]],
+        max_concurrency: Optional[int] = None,
+    ) -> 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.
+        max_concurrency : Optional[int], optional
+            Maximum concurrency for LLM requests across blocks.
+
+        Returns
+        -------
+        Dataset
+            Dataset after processing through all blocks.
+        """
+        current_dataset = dataset
+
         # Execute blocks in sequence
         for i, block in enumerate(self.blocks):
             logger.info(
@@ -398,6 +574,10 @@ class Flow(BaseModel):
             # Prepare block execution parameters
             block_kwargs = self._prepare_block_kwargs(block, runtime_params)
 
+            # Add max_concurrency to block kwargs if provided
+            if max_concurrency is not None:
+                block_kwargs["_flow_max_concurrency"] = max_concurrency
+
             try:
                 # Check if this is a deprecated block and skip validations
                 is_deprecated_block = (
@@ -436,12 +616,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(
@@ -760,6 +934,8 @@ class Flow(BaseModel):
         if len(dataset) == 0:
             raise EmptyDatasetError("Input dataset is empty")
 
+        validate_no_duplicates(dataset)
+
         # Use smaller sample size if dataset is smaller
         actual_sample_size = min(sample_size, len(dataset))
 
@@ -784,9 +960,6 @@ class Flow(BaseModel):
             "execution_time_seconds": 0,
         }
 
-        # Standard
-        import time
-
         start_time = time.time()
 
         try:
@@ -930,6 +1103,186 @@ class Flow(BaseModel):
             "block_names": [block.block_name for block in self.blocks],
         }
 
+    def get_dataset_requirements(self) -> Optional[DatasetRequirements]:
+        """Get the dataset requirements for this flow.
+
+        Returns
+        -------
+        Optional[DatasetRequirements]
+            Dataset requirements object or None if not defined.
+
+        Examples
+        --------
+        >>> flow = Flow.from_yaml("path/to/flow.yaml")
+        >>> requirements = flow.get_dataset_requirements()
+        >>> if requirements:
+        ...     print(f"Required columns: {requirements.required_columns}")
+        """
+        return self.metadata.dataset_requirements
+
+    def get_dataset_schema(self) -> Dataset:
+        """Get an empty dataset with the correct schema for this flow.
+
+        Returns
+        -------
+        Dataset
+            Empty HuggingFace Dataset with the correct schema/features for this flow.
+            Users can add data to this dataset or use it to validate their own dataset schema.
+
+        Examples
+        --------
+        >>> flow = Flow.from_yaml("path/to/flow.yaml")
+        >>> schema_dataset = flow.get_dataset_schema()
+        >>>
+        >>> # Add your data
+        >>> schema_dataset = schema_dataset.add_item({
+        ...     "document": "Your document text",
+        ...     "domain": "Computer Science",
+        ...     "icl_document": "Example document"
+        ... })
+        >>>
+        >>> # Or validate your existing dataset schema
+        >>> my_dataset = Dataset.from_dict(my_data)
+        >>> if my_dataset.features == schema_dataset.features:
+        ...     print("Schema matches!")
+        """
+
+        requirements = self.get_dataset_requirements()
+
+        if requirements is None:
+            # Return empty dataset with no schema requirements
+            return Dataset.from_dict({})
+
+        # Build schema features
+        schema_features = {}
+
+        # Process required columns
+        for col_name in requirements.required_columns:
+            col_type = requirements.column_types.get(col_name, "string")
+            schema_features[col_name] = self._map_column_type_to_feature(col_type)
+
+        # Process optional columns
+        for col_name in requirements.optional_columns:
+            col_type = requirements.column_types.get(col_name, "string")
+            schema_features[col_name] = self._map_column_type_to_feature(col_type)
+
+        # Create empty dataset with the correct features
+        features = datasets.Features(schema_features)
+        empty_data = {col_name: [] for col_name in schema_features.keys()}
+
+        return Dataset.from_dict(empty_data, features=features)
+
+    def _map_column_type_to_feature(self, col_type: str):
+        """Map column type string to HuggingFace feature type."""
+        # Map common type names to HuggingFace types
+        if col_type in ["str", "string", "text"]:
+            return datasets.Value("string")
+        elif col_type in ["int", "integer"]:
+            return datasets.Value("int64")
+        elif col_type in ["float", "number"]:
+            return datasets.Value("float64")
+        elif col_type in ["bool", "boolean"]:
+            return datasets.Value("bool")
+        else:
+            # Default to string for unknown types
+            return datasets.Value("string")
+
+    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 +1305,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."""