data-designer-engine 0.4.0rc2__tar.gz → 0.5.0rc1__tar.gz

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer-engine
-Version: 0.4.0rc2
+Version: 0.5.0rc1
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Requires-Dist: anyascii<1,>=0.3.3
-Requires-Dist: data-designer-config
+Requires-Dist: data-designer-config==0.5.0rc1
 Requires-Dist: duckdb<2,>=1.1.3
 Requires-Dist: faker<21,>=20.1.0
 Requires-Dist: httpx-retries<1,>=0.4.2
@@ -26,6 +26,7 @@ Requires-Dist: jsonschema<5,>=4.0.0
 Requires-Dist: litellm<1.80.12,>=1.73.6
 Requires-Dist: lxml<7,>=6.0.2
 Requires-Dist: marko<3,>=2.1.2
+Requires-Dist: mcp<2,>=1.26.0
 Requires-Dist: networkx<4,>=3.0
 Requires-Dist: ruff<1,>=0.14.10
 Requires-Dist: scipy<2,>=1.11.0

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "data-designer-engine"
-dynamic = ["version"]
+dynamic = ["version", "dependencies"]
 description = "Generation engine for DataDesigner synthetic data generation"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -18,9 +18,22 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
 ]
 
+[build-system]
+requires = ["hatchling", "uv-dynamic-versioning>=0.7.0"]
+build-backend = "hatchling.build"
+
+[tool.hatch.version]
+source = "uv-dynamic-versioning"
+
+[tool.uv-dynamic-versioning]
+vcs = "git"
+style = "pep440"
+bump = true
+
+[tool.hatch.metadata.hooks.uv-dynamic-versioning]
 dependencies = [
     "anyascii>=0.3.3,<1",
-    "data-designer-config",
+    "data-designer-config=={{ version }}",
     "duckdb>=1.1.3,<2",
     "faker>=20.1.0,<21",
     "httpx>=0.27.2,<1",
@@ -32,6 +45,7 @@ dependencies = [
     "litellm>=1.73.6,<1.80.12",
     "lxml>=6.0.2,<7",
     "marko>=2.1.2,<3",
+    "mcp>=1.26.0,<2",
     "networkx>=3.0,<4",
     "ruff>=0.14.10,<1",
     "scipy>=1.11.0,<2",
@@ -39,24 +53,17 @@ dependencies = [
     "tiktoken>=0.8.0,<1",
 ]
 
-[build-system]
-requires = ["hatchling", "hatch-vcs"]
-build-backend = "hatchling.build"
-
-[tool.hatch.version]
-source = "vcs"
-fallback-version = "0.1.0.dev0"
-raw-options = { root = "../.." }
-
-[tool.hatch.build.hooks.vcs]
-version-file = "src/data_designer/engine/_version.py"
-
 [tool.hatch.build.targets.wheel]
 packages = ["src/data_designer"]
 
 [tool.ruff]
 extend = "../../pyproject.toml"
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_default_fixture_loop_scope = "session"
+env = ["DISABLE_DATA_DESIGNER_PLUGINS=true"]
+
 [tool.uv]
 package = true
 

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/analysis/column_profilers/base.py

@@ -10,8 +10,7 @@ from typing import TYPE_CHECKING
 from pydantic import BaseModel, model_validator
 from typing_extensions import Self
 
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.base import ConfigBase, SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.configurable_task import ConfigurableTask, TaskConfigT
 from data_designer.lazy_heavy_imports import pd

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/analysis/dataset_profiler.py

@@ -12,8 +12,7 @@ from pydantic import Field, field_validator
 
 from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.base import ConfigBase, SingleColumnConfig
 from data_designer.config.column_types import ColumnConfigT
 from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.column_statistics import get_column_statistics_calculator

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/column_generators/generators/base.py

@@ -6,9 +6,9 @@ from __future__ import annotations
 import functools
 import logging
 from abc import ABC, abstractmethod
-from enum import Enum
 from typing import TYPE_CHECKING, overload
 
+from data_designer.config.column_configs import GenerationStrategy
 from data_designer.engine.configurable_task import ConfigurableTask, DataT, TaskConfigT
 from data_designer.lazy_heavy_imports import pd
 
@@ -22,11 +22,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class GenerationStrategy(str, Enum):
-    CELL_BY_CELL = "cell_by_cell"
-    FULL_COLUMN = "full_column"
-
-
 class ColumnGenerator(ConfigurableTask[TaskConfigT], ABC):
     @property
     def can_generate_from_scratch(self) -> bool:

data_designer_engine-0.5.0rc1/src/data_designer/engine/column_generators/generators/custom.py

@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Custom column generator using user-provided callable functions."""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from typing import TYPE_CHECKING, Any
+
+from data_designer.config.column_configs import CustomColumnConfig, GenerationStrategy
+from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.utils.errors import CustomColumnGenerationError
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+class CustomColumnGenerator(ColumnGenerator[CustomColumnConfig]):
+    """Column generator that uses a user-provided callable function.
+
+    Supports two strategies based on config.strategy:
+        - cell_by_cell: Processes rows one at a time (dict -> dict), parallelized by framework.
+        - full_column: Processes entire batch (DataFrame -> DataFrame) for vectorized ops.
+
+    Supported function signatures (validated by parameter name):
+        - fn(row) -> dict                              # cell_by_cell, simple transform
+        - fn(row, generator_params) -> dict            # cell_by_cell, with typed params
+        - fn(row, generator_params, models) -> dict    # cell_by_cell, with LLM access
+        - fn(df) -> DataFrame                          # full_column, simple transform
+        - fn(df, generator_params) -> DataFrame        # full_column, with typed params
+        - fn(df, generator_params, models) -> DataFrame  # full_column, with LLM access
+
+    The models dict provides direct access to ModelFacade instances keyed by alias.
+    """
+
+    def get_generation_strategy(self) -> GenerationStrategy:
+        """Return strategy based on config."""
+        return self.config.generation_strategy
+
+    def generate(self, data: dict | pd.DataFrame) -> dict | pd.DataFrame:
+        """Generate column value(s) for a row (dict) or batch (DataFrame)."""
+        is_full_column = self.config.generation_strategy == GenerationStrategy.FULL_COLUMN
+        is_dataframe = not isinstance(data, dict)
+
+        # Validate data type matches strategy
+        if is_full_column and not is_dataframe:
+            raise CustomColumnGenerationError(
+                f"Custom generator {self.config.name!r} is configured for 'full_column' strategy "
+                "but received a dict. Expected a DataFrame."
+            )
+        if not is_full_column and is_dataframe:
+            raise CustomColumnGenerationError(
+                f"Custom generator {self.config.name!r} is configured for 'cell_by_cell' strategy "
+                "but received a DataFrame. Expected a dict."
+            )
+
+        return self._generate(data, is_dataframe)
+
+    def _generate(self, data: dict | pd.DataFrame, is_dataframe: bool) -> dict | pd.DataFrame:
+        """Unified generation logic for both strategies."""
+        # Get columns/keys using unified accessor
+        get_keys = (lambda d: set(d.columns)) if is_dataframe else (lambda d: set(d.keys()))
+        expected_type = pd.DataFrame if is_dataframe else dict
+        type_name = "DataFrame" if is_dataframe else "dict"
+
+        # Check required columns
+        missing = set(self.config.required_columns) - get_keys(data)
+        if missing:
+            raise CustomColumnGenerationError(
+                f"Missing required columns for custom generator '{self.config.name}': {sorted(missing)}"
+            )
+
+        keys_before = get_keys(data)
+
+        # Invoke generator
+        try:
+            result = self._invoke_generator_function(data)
+        except CustomColumnGenerationError:
+            raise
+        except Exception as e:
+            raise CustomColumnGenerationError(
+                f"Custom generator function failed for column '{self.config.name}': {e}"
+            ) from e
+
+        # Validate return type
+        if not isinstance(result, expected_type):
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' must return a {type_name}, "
+                f"got {type(result).__name__}"
+            )
+
+        return self._validate_output(result, keys_before, is_dataframe)
+
+    def _validate_output(
+        self, result: dict | pd.DataFrame, keys_before: set[str], is_dataframe: bool
+    ) -> dict | pd.DataFrame:
+        """Validate output columns and remove undeclared ones."""
+        # Unified accessors
+        get_keys = (lambda d: set(d.columns)) if is_dataframe else (lambda d: set(d.keys()))
+        container_name = "DataFrame" if is_dataframe else "row"
+
+        expected_new = {self.config.name} | set(self.config.side_effect_columns)
+        result_keys = get_keys(result)
+
+        # Check primary column exists
+        if self.config.name not in result_keys:
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' did not create the expected column. "
+                f"The generator_function must add a key named '{self.config.name}' to the {container_name}."
+            )
+
+        # Check side effect columns exist
+        missing = set(self.config.side_effect_columns) - result_keys
+        if missing:
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' did not create declared side_effect_columns: "
+                f"{sorted(missing)}. Declared side_effect_columns must be added to the {container_name}."
+            )
+
+        # Check no pre-existing columns removed
+        removed = keys_before - result_keys
+        if removed:
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' removed pre-existing columns: "
+                f"{sorted(removed)}. The generator_function must not remove any existing columns."
+            )
+
+        # Remove undeclared columns with warning
+        undeclared = (result_keys - keys_before) - expected_new
+        if undeclared:
+            logger.warning(
+                f"⚠️ Custom generator for column '{self.config.name}' created undeclared columns: "
+                f"{sorted(undeclared)}. These columns will be removed. "
+                f"To keep additional columns, declare them in @custom_column_generator(side_effect_columns=[...])."
+            )
+            if is_dataframe:
+                result = result.drop(columns=list(undeclared))
+            else:
+                for key in undeclared:
+                    del result[key]
+
+        return result
+
+    def _invoke_generator_function(self, data: dict | pd.DataFrame) -> dict | pd.DataFrame:
+        """Invoke the user's generate function with appropriate arguments based on signature."""
+        params = self._get_validated_params()
+
+        if len(params) == 1:
+            return self.config.generator_function(data)
+        elif len(params) == 2:
+            return self.config.generator_function(data, self.config.generator_params)
+        else:
+            models = self._build_models_dict()
+            return self.config.generator_function(data, self.config.generator_params, models)
+
+    def _build_models_dict(self) -> dict[str, Any]:
+        """Build a dict of ModelFacade instances from model_aliases."""
+        return {
+            alias: self.resource_provider.model_registry.get_model(model_alias=alias)
+            for alias in self.config.model_aliases
+        }
+
+    def _get_validated_params(self) -> list[inspect.Parameter]:
+        """Get positional params and validate first param matches generation strategy."""
+        params = [
+            p
+            for p in inspect.signature(self.config.generator_function).parameters.values()
+            if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        ]
+        # Decorator validated param names; here we only check strategy match
+        is_full = self.config.generation_strategy == GenerationStrategy.FULL_COLUMN
+        expected = "df" if is_full else "row"
+        if params[0].name != expected:
+            raise CustomColumnGenerationError(
+                f"Generator '{self.config.name}': strategy is {'full_column' if is_full else 'cell_by_cell'}, "
+                f"first parameter must be '{expected}', got '{params[0].name}'."
+            )
+        return params
+
+    def log_pre_generation(self) -> None:
+        logger.info(f"{self.config.get_column_emoji()} Custom column config for column '{self.config.name}'")
+        logger.info(f"  |-- generator_function: {self.config.generator_function.__name__!r}")
+        logger.info(f"  |-- generation_strategy: {self.config.generation_strategy!r}")
+        logger.info(f"  |-- required_columns: {self.config.required_columns}")
+        if self.config.side_effect_columns:
+            logger.info(f"  |-- side_effect_columns: {self.config.side_effect_columns}")
+        if self.config.model_aliases:
+            logger.info(f"  |-- model_aliases: {self.config.model_aliases}")
+        if self.config.generator_params:
+            logger.info(f"  |-- generator_params: {self.config.generator_params}")

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/column_generators/generators/llm_completion.py

@@ -12,7 +12,8 @@ from data_designer.config.column_configs import (
     LLMStructuredColumnConfig,
     LLMTextColumnConfig,
 )
-from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
+from data_designer.config.utils.constants import REASONING_CONTENT_COLUMN_POSTFIX, TRACE_COLUMN_POSTFIX
+from data_designer.config.utils.trace_type import TraceType
 from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel, GenerationStrategy
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
@@ -66,7 +67,7 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
             for context in self.config.multi_modal_context:
                 multi_modal_context.extend(context.get_contexts(deserialized_record))
 
-        response, reasoning_trace = self.model.generate(
+        response, trace = self.model.generate(
             prompt=self.prompt_renderer.render(
                 record=deserialized_record,
                 prompt_template=self.config.prompt,
@@ -79,6 +80,7 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
             ),
             parser=self.response_recipe.parse,
             multi_modal_context=multi_modal_context,
+            tool_alias=self.config.tool_alias,
             max_correction_steps=self.max_conversation_correction_steps,
             max_conversation_restarts=self.max_conversation_restarts,
             purpose=f"running generation for column '{self.config.name}'",
@@ -87,11 +89,39 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
         serialized_output = self.response_recipe.serialize_output(response)
         data[self.config.name] = self._process_serialized_output(serialized_output)
 
-        if reasoning_trace:
-            data[self.config.name + REASONING_TRACE_COLUMN_POSTFIX] = reasoning_trace
+        effective_trace_type = self.config.with_trace
+
+        if effective_trace_type == TraceType.ALL_MESSAGES:
+            data[self.config.name + TRACE_COLUMN_POSTFIX] = [message.to_dict() for message in trace]
+        elif effective_trace_type == TraceType.LAST_MESSAGE:
+            last_assistant = next((m for m in reversed(trace) if m.role == "assistant"), None)
+            data[self.config.name + TRACE_COLUMN_POSTFIX] = [last_assistant.to_dict()] if last_assistant else []
+
+        if self.config.extract_reasoning_content:
+            data[self.config.name + REASONING_CONTENT_COLUMN_POSTFIX] = self._extract_reasoning_content(trace)
 
         return data
 
+    def _extract_reasoning_content(self, trace: list) -> str | None:
+        """Extract reasoning_content from the final assistant message in the trace.
+
+        Args:
+            trace: List of ChatMessage objects from the generation.
+
+        Returns:
+            The stripped reasoning_content from the final assistant message, or None if not present.
+        """
+        reasoning_value: str | None = None
+        for message in reversed(trace):
+            if message.role == "assistant":
+                reasoning_value = message.reasoning_content
+                break
+
+        if reasoning_value is not None:
+            reasoning_value = reasoning_value.strip() or None
+
+        return reasoning_value
+
     def _process_serialized_output(self, serialized_output: str) -> str | dict | list:
         """Process the serialized output from the model. Subclasses can override to customize deserialization."""
         return serialized_output

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/column_generators/registry.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import (
+    CustomColumnConfig,
     EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
@@ -15,6 +16,7 @@ from data_designer.config.column_configs import (
 )
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.custom import CustomColumnGenerator
 from data_designer.engine.column_generators.generators.embedding import EmbeddingCellGenerator
 from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
 from data_designer.engine.column_generators.generators.llm_completion import (
@@ -40,6 +42,7 @@ class ColumnGeneratorRegistry(TaskRegistry[DataDesignerColumnType, ColumnGenerat
 
 def create_default_column_generator_registry(with_plugins: bool = True) -> ColumnGeneratorRegistry:
     registry = ColumnGeneratorRegistry()
+    registry.register(DataDesignerColumnType.CUSTOM, CustomColumnGenerator, CustomColumnConfig)
     registry.register(DataDesignerColumnType.LLM_TEXT, LLMTextCellGenerator, LLMTextColumnConfig)
     registry.register(DataDesignerColumnType.LLM_CODE, LLMCodeCellGenerator, LLMCodeColumnConfig)
     registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig)

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/column_generators/utils/errors.py

@@ -13,3 +13,6 @@ class ExpressionTemplateRenderError(DataDesignerError): ...
 
 
 class SeedDatasetError(DataDesignerError): ...
+
+
+class CustomColumnGenerationError(DataDesignerError): ...

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 import json
 import logging
 
-from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.base import SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.config.models import ModelConfig
 from data_designer.config.utils.code_lang import CodeLang

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/dataset_builders/column_wise_builder.py

@@ -12,6 +12,7 @@ import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
 
+from data_designer.config.column_configs import CustomColumnConfig
 from data_designer.config.column_types import ColumnConfigT
 from data_designer.config.config_builder import BuilderConfig
 from data_designer.config.data_designer_config import DataDesignerConfig
@@ -34,6 +35,7 @@ from data_designer.engine.dataset_builders.multi_column_configs import MultiColu
 from data_designer.engine.dataset_builders.utils.concurrency import ConcurrentThreadExecutor
 from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
 from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager
+from data_designer.engine.dataset_builders.utils.progress_tracker import ProgressTracker
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
@@ -96,6 +98,7 @@ class ColumnWiseDatasetBuilder:
         on_batch_complete: Callable[[Path], None] | None = None,
     ) -> Path:
         self._run_model_health_check_if_needed()
+        self._run_mcp_tool_check_if_needed()
         self._write_builder_config()
         generators = self._initialize_generators()
         start_time = time.perf_counter()
@@ -124,6 +127,7 @@ class ColumnWiseDatasetBuilder:
 
     def build_preview(self, *, num_records: int) -> pd.DataFrame:
         self._run_model_health_check_if_needed()
+        self._run_mcp_tool_check_if_needed()
 
         generators = self._initialize_generators()
         group_id = uuid.uuid4().hex
@@ -208,11 +212,26 @@ class ColumnWiseDatasetBuilder:
         df = generator.generate(self.batch_manager.get_current_batch(as_dataframe=True))
         self.batch_manager.update_records(df.to_dict(orient="records"))
 
-    def _run_model_health_check_if_needed(self) -> bool:
-        if any(column_type_is_model_generated(config.column_type) for config in self.single_column_configs):
-            self._resource_provider.model_registry.run_health_check(
-                list(set(config.model_alias for config in self.llm_generated_column_configs))
-            )
+    def _run_model_health_check_if_needed(self) -> None:
+        model_aliases: set[str] = set()
+        for config in self.single_column_configs:
+            if column_type_is_model_generated(config.column_type):
+                model_aliases.add(config.model_alias)
+            if isinstance(config, CustomColumnConfig) and config.model_aliases:
+                model_aliases.update(config.model_aliases)
+
+        if model_aliases:
+            self._resource_provider.model_registry.run_health_check(list(model_aliases))
+
+    def _run_mcp_tool_check_if_needed(self) -> None:
+        tool_aliases = sorted(
+            {config.tool_alias for config in self.llm_generated_column_configs if getattr(config, "tool_alias", None)}
+        )
+        if not tool_aliases:
+            return
+        if self._resource_provider.mcp_registry is None:
+            raise DatasetGenerationError(f"Tool alias(es) {tool_aliases!r} specified but no MCPRegistry configured.")
+        self._resource_provider.mcp_registry.run_health_check(tool_aliases)
 
     def _fan_out_with_threads(self, generator: ColumnGeneratorWithModelRegistry, max_workers: int) -> None:
         if generator.get_generation_strategy() != GenerationStrategy.CELL_BY_CELL:
@@ -221,16 +240,18 @@ class ColumnWiseDatasetBuilder:
                 "generator so concurrency through threads is not supported."
             )
 
-        logger.info(
-            f"🐙 Processing {generator.config.column_type} column '{generator.config.name}' "
-            f"with {max_workers} concurrent workers"
+        progress_tracker = ProgressTracker(
+            total_records=self.batch_manager.num_records_batch,
+            label=f"{generator.config.column_type} column '{generator.config.name}'",
         )
+        progress_tracker.log_start(max_workers)
+
         settings = self._resource_provider.run_config
         with ConcurrentThreadExecutor(
             max_workers=max_workers,
             column_name=generator.config.name,
-            result_callback=self._worker_result_callback,
-            error_callback=self._worker_error_callback,
+            result_callback=self._make_result_callback(progress_tracker),
+            error_callback=self._make_error_callback(progress_tracker),
             shutdown_error_rate=settings.shutdown_error_rate,
             shutdown_error_window=settings.shutdown_error_window,
             disable_early_shutdown=settings.disable_early_shutdown,
@@ -238,10 +259,26 @@ class ColumnWiseDatasetBuilder:
             for i, record in self.batch_manager.iter_current_batch():
                 executor.submit(lambda record: generator.generate(record), record, context={"index": i})
 
+        progress_tracker.log_final()
+
         if len(self._records_to_drop) > 0:
             self.batch_manager.drop_records(self._records_to_drop)
             self._records_to_drop.clear()
 
+    def _make_result_callback(self, progress_tracker: ProgressTracker) -> Callable[[dict], None]:
+        def callback(result: dict, *, context: dict | None = None) -> None:
+            self._worker_result_callback(result, context=context)
+            progress_tracker.record_success()
+
+        return callback
+
+    def _make_error_callback(self, progress_tracker: ProgressTracker) -> Callable[[Exception], None]:
+        def callback(exc: Exception, *, context: dict | None = None) -> None:
+            self._worker_error_callback(exc, context=context)
+            progress_tracker.record_failure()
+
+        return callback
+
     def _write_processed_batch(self, dataframe: pd.DataFrame) -> None:
         self.batch_manager.update_records(dataframe.to_dict(orient="records"))
         self.batch_manager.write()

{data_designer_engine-0.4.0rc2 → data_designer_engine-0.5.0rc1}/src/data_designer/engine/dataset_builders/multi_column_configs.py

@@ -8,8 +8,8 @@ from typing import TypeAlias
 
 from pydantic import Field, field_validator
 
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SamplerColumnConfig, SeedDatasetColumnConfig, SingleColumnConfig
+from data_designer.config.base import ConfigBase, SingleColumnConfig
+from data_designer.config.column_configs import SamplerColumnConfig, SeedDatasetColumnConfig
 from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType
 from data_designer.config.sampler_constraints import ColumnConstraintT
 from data_designer.config.seed import SeedConfig