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

sdg_hub/core/flow/base.py

@@ -2,13 +2,22 @@
 """Pydantic-based Flow class for managing data generation pipelines."""
 
 # Standard
+from datetime import datetime
 from pathlib import Path
 from typing import Any, Optional, Union
 import time
+import uuid
 
 # Third Party
 from datasets import Dataset
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    field_validator,
+    model_validator,
+)
 from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
@@ -21,6 +30,7 @@ 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.flow_metrics import display_metrics_summary, save_metrics_to_json
 from ..utils.logger_config import setup_logger
 from ..utils.path_resolution import resolve_path
 from ..utils.yaml_utils import save_flow_yaml
@@ -67,6 +77,9 @@ class Flow(BaseModel):
     _migrated_runtime_params: dict[str, dict[str, Any]] = {}
     _llm_client: Any = None  # Only used for backward compatibility with old YAMLs
     _model_config_set: bool = False  # Track if model configuration has been set
+    _block_metrics: list[dict[str, Any]] = PrivateAttr(
+        default_factory=list
+    )  # Track block execution metrics
 
     @field_validator("blocks")
     @classmethod
@@ -356,6 +369,7 @@ class Flow(BaseModel):
         runtime_params: Optional[dict[str, dict[str, Any]]] = None,
         checkpoint_dir: Optional[str] = None,
         save_freq: Optional[int] = None,
+        log_dir: Optional[str] = None,
         max_concurrency: Optional[int] = None,
     ) -> Dataset:
         """Execute the flow blocks in sequence to generate data.
@@ -378,6 +392,10 @@ class Flow(BaseModel):
         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.
+        log_dir : Optional[str], optional
+            Directory to save execution logs. If provided, logs will be written to both
+            console and a log file in this directory. Maintains backward compatibility
+            when None.
         max_concurrency : Optional[int], optional
             Maximum number of concurrent requests across all blocks.
             Controls async request concurrency to prevent overwhelming servers.
@@ -400,6 +418,23 @@ class Flow(BaseModel):
                 f"save_freq must be greater than 0, got {save_freq}"
             )
 
+        # Set up file logging if log_dir is provided
+        flow_logger = logger  # Use global logger by default
+        if log_dir is not None:
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            flow_name = self.metadata.name.replace(" ", "_").lower()
+            log_filename = f"{flow_name}_{timestamp}.log"
+
+            # Create a flow-specific logger for this execution
+            unique_id = str(uuid.uuid4())[:8]  # Short unique ID
+            flow_logger_name = f"{__name__}.flow_{flow_name}_{timestamp}_{unique_id}"
+            flow_logger = setup_logger(
+                flow_logger_name, log_dir=log_dir, log_filename=log_filename
+            )
+            flow_logger.propagate = False
+            flow_logger.info(
+                f"Flow logging enabled - logs will be saved to: {log_dir}/{log_filename}"
+            )
         # Validate max_concurrency parameter
         if max_concurrency is not None:
             # Explicitly reject boolean values (bool is a subclass of int in Python)
@@ -459,84 +494,145 @@ class Flow(BaseModel):
             )
 
             if len(remaining_dataset) == 0:
-                logger.info("All samples already completed, returning existing results")
+                flow_logger.info(
+                    "All samples already completed, returning existing results"
+                )
+                if log_dir is not None and flow_logger is not logger:
+                    for h in list(getattr(flow_logger, "handlers", [])):
+                        try:
+                            h.flush()
+                            h.close()
+                        except Exception:
+                            pass
+                        finally:
+                            flow_logger.removeHandler(h)
+
                 return completed_dataset
 
             dataset = remaining_dataset
-            logger.info(f"Resuming with {len(dataset)} remaining samples")
+            flow_logger.info(f"Resuming with {len(dataset)} remaining samples")
 
-        logger.info(
+        flow_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 "")
         )
 
+        # Reset metrics for this execution
+        self._block_metrics = []
+        run_start = time.perf_counter()
+
         # 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 = []
+        # Execute flow with metrics capture, ensuring metrics are always displayed/saved
+        final_dataset = None
+        execution_successful = False
 
-            # 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))
+        try:
+            # Process dataset in chunks if checkpointing with save_freq
+            if checkpointer and save_freq:
+                all_processed = []
 
-                logger.info(
-                    f"Processing chunk {i // save_freq + 1}: samples {i} to {chunk_end - 1}"
-                )
+                # 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))
 
-                # Execute all blocks on this chunk
-                processed_chunk = self._execute_blocks_on_dataset(
-                    chunk_dataset, runtime_params, max_concurrency
-                )
-                all_processed.append(processed_chunk)
+                    flow_logger.info(
+                        f"Processing chunk {i // save_freq + 1}: samples {i} to {chunk_end - 1}"
+                    )
 
-                # Save checkpoint after chunk completion
-                checkpointer.add_completed_samples(processed_chunk)
+                    # Execute all blocks on this chunk
+                    processed_chunk = self._execute_blocks_on_dataset(
+                        chunk_dataset, runtime_params, flow_logger, max_concurrency
+                    )
+                    all_processed.append(processed_chunk)
 
-            # Save final checkpoint for any remaining samples
-            checkpointer.save_final_checkpoint()
+                    # Save checkpoint after chunk completion
+                    checkpointer.add_completed_samples(processed_chunk)
 
-            # Combine all processed chunks
-            final_dataset = safe_concatenate_with_validation(
-                all_processed, "processed chunks from flow execution"
-            )
+                # Save final checkpoint for any remaining samples
+                checkpointer.save_final_checkpoint()
 
-            # Combine with previously completed samples if any
-            if checkpointer and completed_dataset:
+                # Combine all processed chunks
                 final_dataset = safe_concatenate_with_validation(
-                    [completed_dataset, final_dataset],
-                    "completed checkpoint data with newly processed data",
+                    all_processed, "processed chunks from flow execution"
                 )
 
-        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:
+                if checkpointer and 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"
-        )
+            else:
+                # Process entire dataset at once
+                final_dataset = self._execute_blocks_on_dataset(
+                    dataset, runtime_params, flow_logger, 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",
+                        )
+
+            execution_successful = True
+
+        finally:
+            # Always display metrics and save JSON, even if execution failed
+            display_metrics_summary(
+                self._block_metrics, self.metadata.name, final_dataset
+            )
+
+            # Save metrics to JSON if log_dir is provided
+            if log_dir is not None:
+                # Ensure necessary variables exist
+                if "timestamp" not in locals() or "flow_name" not in locals():
+                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                    flow_name = self.metadata.name.replace(" ", "_").lower()
+
+                save_metrics_to_json(
+                    self._block_metrics,
+                    self.metadata.name,
+                    self.metadata.version,
+                    execution_successful,
+                    run_start,
+                    log_dir,
+                    timestamp,
+                    flow_name,
+                    flow_logger,
+                )
+
+        # Keep a basic log entry for file logs (only if execution was successful)
+        if execution_successful and final_dataset is not None:
+            flow_logger.info(
+                f"Flow '{self.metadata.name}' completed successfully: "
+                f"{len(final_dataset)} final samples, "
+                f"{len(final_dataset.column_names)} final columns"
+            )
+
+        # Close file handlers if we opened a flow-specific logger
+        if log_dir is not None and flow_logger is not logger:
+            for h in list(getattr(flow_logger, "handlers", [])):
+                try:
+                    h.flush()
+                    h.close()
+                except Exception:
+                    pass
+                finally:
+                    flow_logger.removeHandler(h)
 
         return final_dataset
 
@@ -544,6 +640,7 @@ class Flow(BaseModel):
         self,
         dataset: Dataset,
         runtime_params: dict[str, dict[str, Any]],
+        flow_logger=None,
         max_concurrency: Optional[int] = None,
     ) -> Dataset:
         """Execute all blocks in sequence on the given dataset.
@@ -554,6 +651,8 @@ class Flow(BaseModel):
             Dataset to process through all blocks.
         runtime_params : Dict[str, Dict[str, Any]]
             Runtime parameters for block execution.
+        flow_logger : logging.Logger, optional
+            Logger to use for this execution. Falls back to global logger if None.
         max_concurrency : Optional[int], optional
             Maximum concurrency for LLM requests across blocks.
 
@@ -562,11 +661,13 @@ class Flow(BaseModel):
         Dataset
             Dataset after processing through all blocks.
         """
+        # Use provided logger or fall back to global logger
+        exec_logger = flow_logger if flow_logger is not None else logger
         current_dataset = dataset
 
         # Execute blocks in sequence
         for i, block in enumerate(self.blocks):
-            logger.info(
+            exec_logger.info(
                 f"Executing block {i + 1}/{len(self.blocks)}: "
                 f"{block.block_name} ({block.__class__.__name__})"
             )
@@ -578,6 +679,11 @@ class Flow(BaseModel):
             if max_concurrency is not None:
                 block_kwargs["_flow_max_concurrency"] = max_concurrency
 
+            # Capture metrics before execution
+            start_time = time.perf_counter()
+            input_rows = len(current_dataset)
+            input_cols = set(current_dataset.column_names)
+
             try:
                 # Check if this is a deprecated block and skip validations
                 is_deprecated_block = (
@@ -587,7 +693,7 @@ class Flow(BaseModel):
                 )
 
                 if is_deprecated_block:
-                    logger.debug(
+                    exec_logger.debug(
                         f"Skipping validations for deprecated block: {block.block_name}"
                     )
                     # Call generate() directly to skip validations, but keep the runtime params
@@ -602,14 +708,51 @@ class Flow(BaseModel):
                         f"Block '{block.block_name}' produced empty dataset"
                     )
 
-                logger.info(
+                # Capture metrics after successful execution
+                execution_time = time.perf_counter() - start_time
+                output_rows = len(current_dataset)
+                output_cols = set(current_dataset.column_names)
+                added_cols = output_cols - input_cols
+                removed_cols = input_cols - output_cols
+
+                # Store block metrics
+                self._block_metrics.append(
+                    {
+                        "block_name": block.block_name,
+                        "block_type": block.__class__.__name__,
+                        "execution_time": execution_time,
+                        "input_rows": input_rows,
+                        "output_rows": output_rows,
+                        "added_cols": list(added_cols),
+                        "removed_cols": list(removed_cols),
+                        "status": "success",
+                    }
+                )
+
+                exec_logger.info(
                     f"Block '{block.block_name}' completed successfully: "
                     f"{len(current_dataset)} samples, "
                     f"{len(current_dataset.column_names)} columns"
                 )
 
             except Exception as exc:
-                logger.error(
+                # Capture metrics for failed execution
+                execution_time = time.perf_counter() - start_time
+                self._block_metrics.append(
+                    {
+                        "block_name": block.block_name,
+                        "block_type": block.__class__.__name__,
+                        "execution_time": execution_time,
+                        "input_rows": input_rows,
+                        "output_rows": 0,
+                        "added_cols": [],
+                        "removed_cols": [],
+                        "status": "failed",
+                        "error": str(exc),
+                    }
+                )
+
+                exec_logger.error(
                     f"Block '{block.block_name}' failed during execution: {exc}"
                 )
                 raise FlowValidationError(
@@ -960,7 +1103,7 @@ class Flow(BaseModel):
             "execution_time_seconds": 0,
         }
 
-        start_time = time.time()
+        start_time = time.perf_counter()
 
         try:
             # Execute the flow with sample data
@@ -968,7 +1111,7 @@ class Flow(BaseModel):
             runtime_params = runtime_params or {}
 
             for i, block in enumerate(self.blocks):
-                block_start_time = time.time()
+                block_start_time = time.perf_counter()
                 input_rows = len(current_dataset)
 
                 logger.info(
@@ -1027,7 +1170,7 @@ class Flow(BaseModel):
                 else {},
             }
 
-            execution_time = time.time() - start_time
+            execution_time = time.perf_counter() - start_time
             dry_run_results["execution_time_seconds"] = execution_time
 
             logger.info(
@@ -1038,7 +1181,7 @@ class Flow(BaseModel):
             return dry_run_results
 
         except Exception as exc:
-            execution_time = time.time() - start_time
+            execution_time = time.perf_counter() - start_time
             dry_run_results["execution_successful"] = False
             dry_run_results["execution_time_seconds"] = execution_time
             dry_run_results["error"] = str(exc)

sdg_hub/core/utils/datautils.py

@@ -1,5 +1,6 @@
 # Third Party
 from datasets import Dataset, concatenate_datasets
+import numpy as np
 
 # Local
 from .error_handling import FlowValidationError
@@ -19,7 +20,8 @@ def validate_no_duplicates(dataset: Dataset) -> None:
     """
     Validate that the input dataset contains only unique rows.
 
-    Uses pandas `.duplicated()` for efficient duplicate detection.
+    Uses pandas `.duplicated()` for efficient duplicate detection, with preprocessing
+    to handle numpy arrays that cause TypeError in pandas duplicate detection.
     Raises FlowValidationError if duplicates are found, including a count
     of the duplicate rows detected.
 
@@ -33,9 +35,50 @@ def validate_no_duplicates(dataset: Dataset) -> None:
     FlowValidationError
         If duplicate rows are detected in the dataset.
     """
+    if len(dataset) == 0:
+        return
+
     df = dataset.to_pandas()
-    duplicate_count = int(df.duplicated(keep="first").sum())
 
+    def is_hashable(x):
+        try:
+            hash(x)
+            return True
+        except TypeError:
+            return False
+
+    def make_hashable(x):
+        if is_hashable(x):
+            # int, float, str, bytes, None etc. are already hashable
+            return x
+        if isinstance(x, np.ndarray):
+            if x.ndim == 0:
+                return make_hashable(x.item())
+            return tuple(make_hashable(i) for i in x)
+        if isinstance(x, dict):
+            # sort robustly even with heterogeneous key types
+            return tuple(
+                sorted(
+                    ((k, make_hashable(v)) for k, v in x.items()),
+                    key=lambda kv: repr(kv[0]),
+                )
+            )
+        if isinstance(x, (set, frozenset)):
+            # order‑insensitive
+            return frozenset(make_hashable(i) for i in x)
+        if hasattr(x, "__iter__"):
+            # lists, tuples, custom iterables
+            return tuple(make_hashable(i) for i in x)
+        # last‑resort fallback to a stable representation
+        return repr(x)
+
+    # Apply to the whole dataframe to ensure every cell is hashable
+    if hasattr(df, "map"):
+        df = df.map(make_hashable)
+    else:
+        df = df.applymap(make_hashable)
+
+    duplicate_count = int(df.duplicated(keep="first").sum())
     if duplicate_count > 0:
         raise FlowValidationError(
             f"Input dataset contains {duplicate_count} duplicate rows. "