data-designer-engine 0.4.0__py3-none-any.whl

data_designer/engine/dataset_builders/utils/config_compiler.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.data_designer_config import DataDesignerConfig
+from data_designer.config.processors import ProcessorConfig
+from data_designer.engine.dataset_builders.multi_column_configs import (
+    DatasetBuilderColumnConfigT,
+    SamplerMultiColumnConfig,
+    SeedDatasetMultiColumnConfig,
+)
+from data_designer.engine.dataset_builders.utils.dag import topologically_sort_column_configs
+from data_designer.engine.dataset_builders.utils.errors import ConfigCompilationError
+
+
+def compile_dataset_builder_column_configs(config: DataDesignerConfig) -> list[DatasetBuilderColumnConfigT]:
+    seed_column_configs = []
+    sampler_column_configs = []
+    generated_column_configs = []
+
+    for column_config in topologically_sort_column_configs(config.columns):
+        if column_config.column_type == DataDesignerColumnType.SEED_DATASET:
+            seed_column_configs.append(column_config)
+        elif column_config.column_type == DataDesignerColumnType.SAMPLER:
+            sampler_column_configs.append(column_config)
+        else:
+            generated_column_configs.append(column_config)
+
+    compiled_column_configs = []
+
+    if len(seed_column_configs) > 0:
+        if config.seed_config is None:
+            raise ConfigCompilationError("🛑 Seed column configs require a seed configuration.")
+        compiled_column_configs.append(
+            SeedDatasetMultiColumnConfig(
+                columns=seed_column_configs,
+                source=config.seed_config.source,
+                sampling_strategy=config.seed_config.sampling_strategy,
+                selection_strategy=config.seed_config.selection_strategy,
+            )
+        )
+
+    if len(sampler_column_configs) > 0:
+        compiled_column_configs.append(
+            SamplerMultiColumnConfig(
+                columns=sampler_column_configs,
+                constraints=config.constraints or [],
+            )
+        )
+
+    if len(generated_column_configs) > 0:
+        compiled_column_configs.extend(generated_column_configs)
+
+    return compiled_column_configs
+
+
+def compile_dataset_builder_processor_configs(
+    config: DataDesignerConfig,
+) -> list[ProcessorConfig]:
+    return config.processors or []

data_designer/engine/dataset_builders/utils/dag.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from data_designer.config.column_types import ColumnConfigT
+from data_designer.engine.column_generators.utils.generator_classification import column_type_used_in_execution_dag
+from data_designer.engine.dataset_builders.utils.errors import DAGCircularDependencyError
+from data_designer.lazy_heavy_imports import nx
+
+if TYPE_CHECKING:
+    import networkx as nx
+
+logger = logging.getLogger(__name__)
+
+
+def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> list[ColumnConfigT]:
+    dag = nx.DiGraph()
+
+    non_dag_column_config_list = [
+        col for col in column_configs if not column_type_used_in_execution_dag(col.column_type)
+    ]
+    dag_column_config_dict = {
+        col.name: col for col in column_configs if column_type_used_in_execution_dag(col.column_type)
+    }
+
+    if len(dag_column_config_dict) == 0:
+        return non_dag_column_config_list
+
+    side_effect_dict = {n: list(c.side_effect_columns) for n, c in dag_column_config_dict.items()}
+
+    logger.info("⛓️ Sorting column configs into a Directed Acyclic Graph")
+    for name, col in dag_column_config_dict.items():
+        dag.add_node(name)
+        for req_col_name in col.required_columns:
+            if req_col_name in list(dag_column_config_dict.keys()):
+                logger.debug(f"  |-- 🔗 `{name}` depends on `{req_col_name}`")
+                dag.add_edge(req_col_name, name)
+
+            # If the required column is a side effect of another column,
+            # add an edge from the parent column to the current column.
+            elif req_col_name in sum(side_effect_dict.values(), []):
+                for parent, cols in side_effect_dict.items():
+                    if req_col_name in cols:
+                        logger.debug(f"  |-- 🔗 `{name}` depends on `{parent}` via `{req_col_name}`")
+                        dag.add_edge(parent, name)
+                        break
+
+    if not nx.is_directed_acyclic_graph(dag):
+        raise DAGCircularDependencyError(
+            "🛑 The Data Designer column configurations contain cyclic dependencies. Please "
+            "inspect the column configurations and ensure they can be sorted without "
+            "circular references."
+        )
+
+    sorted_columns = non_dag_column_config_list
+    sorted_columns.extend([dag_column_config_dict[n] for n in list(nx.topological_sort(dag))])
+
+    return sorted_columns

data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

@@ -0,0 +1,200 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Container, Iterator
+
+from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage, BatchStage
+from data_designer.engine.dataset_builders.utils.errors import DatasetBatchManagementError
+from data_designer.lazy_heavy_imports import pd, pq
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow.parquet as pq
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetBatchManager:
+    def __init__(self, artifact_storage: ArtifactStorage):
+        self._buffer: list[dict] = []
+        self._current_batch_number = 0
+        self._num_records_list: list[int] | None = None
+        self._buffer_size: int | None = None
+        self.artifact_storage = artifact_storage
+
+    @property
+    def num_batches(self) -> int:
+        if self._num_records_list is None:
+            return 0
+        return len(self._num_records_list)
+
+    @property
+    def num_records_batch(self) -> int:
+        if self._num_records_list is None or self._current_batch_number >= len(self._num_records_list):
+            raise DatasetBatchManagementError("🛑 Invalid batch number or num_records_list not set.")
+        return self._num_records_list[self._current_batch_number]
+
+    @property
+    def num_records_list(self) -> list[int]:
+        if self._num_records_list is None:
+            raise DatasetBatchManagementError("🛑 `num_records_list` is not set. Call start() first.")
+        return self._num_records_list
+
+    @property
+    def num_records_in_buffer(self) -> int:
+        return len(self._buffer)
+
+    @property
+    def buffer_is_empty(self) -> bool:
+        return len(self._buffer) == 0
+
+    @property
+    def buffer_size(self) -> int:
+        if self._buffer_size is None:
+            raise DatasetBatchManagementError("🛑 `buffer_size` is not set. Call start() first.")
+        return self._buffer_size
+
+    def add_record(self, record: dict) -> None:
+        self.add_records([record])
+
+    def add_records(self, records: list[dict]) -> None:
+        self._buffer.extend(records)
+        if len(self._buffer) > self.buffer_size:
+            raise DatasetBatchManagementError(
+                f"🛑 Buffer size exceeded. Current: {len(self._buffer)}, Max: {self.buffer_size}. "
+                "Flush the batch before adding more records."
+            )
+
+    def drop_records(self, index: Container[int]) -> None:
+        self._buffer = [record for i, record in enumerate(self._buffer) if i not in index]
+
+    def finish_batch(self, on_complete: Callable[[Path], None] | None = None) -> Path | None:
+        """Finish the batch by moving the results from the partial results path to the final parquet folder.
+
+        Returns:
+            The path to the written parquet file.
+        """
+        if self._current_batch_number >= self.num_batches:
+            raise DatasetBatchManagementError("🛑 All batches have been processed.")
+
+        if self.write() is not None:
+            final_file_path = self.artifact_storage.move_partial_result_to_final_file_path(self._current_batch_number)
+
+            self.artifact_storage.write_metadata(
+                {
+                    "target_num_records": sum(self.num_records_list),
+                    "total_num_batches": self.num_batches,
+                    "buffer_size": self._buffer_size,
+                    "schema": {field.name: str(field.type) for field in pq.read_schema(final_file_path)},
+                    "file_paths": self.artifact_storage.get_file_paths(),
+                    "num_completed_batches": self._current_batch_number + 1,
+                    "dataset_name": self.artifact_storage.dataset_name,
+                }
+            )
+
+            if on_complete:
+                on_complete(final_file_path)
+        else:
+            final_file_path = None
+
+            logger.warning(
+                f"⚠️ Batch {self._current_batch_number + 1} finished without any results to write. "
+                "A partial dataset containing the currently available columns has been written to the partial results "
+                f"directory: {self.artifact_storage.partial_results_path}"
+            )
+
+        self._current_batch_number += 1
+        self._buffer: list[dict] = []
+
+        return final_file_path
+
+    def finish(self) -> None:
+        """Finish the dataset writing process by deleting the partial results path if it exists and is empty."""
+
+        # If the partial results path is empty, delete it.
+        if not any(self.artifact_storage.partial_results_path.iterdir()):
+            self.artifact_storage.partial_results_path.rmdir()
+
+        # Otherwise, log a warning, since existing partial results means the dataset is not complete.
+        else:
+            logger.warning("⚠️ Dataset writing finished with partial results.")
+
+        self.reset()
+
+    def get_current_batch_number(self) -> int:
+        return self._current_batch_number
+
+    def get_current_batch(self, *, as_dataframe: bool = False) -> pd.DataFrame | list[dict]:
+        if as_dataframe:
+            return pd.DataFrame(self._buffer)
+        return self._buffer
+
+    def iter_current_batch(self) -> Iterator[tuple[int, dict]]:
+        for i, record in enumerate(self._buffer):
+            yield i, record
+
+    def reset(self, delete_files: bool = False) -> None:
+        self._current_batch_number = 0
+        self._buffer: list[dict] = []
+        if delete_files:
+            for dir_path in [
+                self.artifact_storage.final_dataset_path,
+                self.artifact_storage.partial_results_path,
+                self.artifact_storage.dropped_columns_dataset_path,
+                self.artifact_storage.base_dataset_path,
+            ]:
+                if dir_path.exists():
+                    try:
+                        shutil.rmtree(dir_path)
+                    except OSError as e:
+                        raise DatasetBatchManagementError(f"🛑 Failed to delete directory {dir_path}: {e}")
+
+    def start(self, *, num_records: int, buffer_size: int) -> None:
+        if num_records <= 0:
+            raise DatasetBatchManagementError("🛑 num_records must be positive.")
+        if buffer_size <= 0:
+            raise DatasetBatchManagementError("🛑 buffer_size must be positive.")
+
+        self._buffer_size = buffer_size
+        self._num_records_list = [buffer_size] * (num_records // buffer_size)
+        if remaining_records := num_records % buffer_size:
+            self._num_records_list.append(remaining_records)
+        self.reset()
+
+    def write(self) -> Path | None:
+        """Write the current batch to a parquet file.
+
+        This method always writes results to the partial results path.
+
+        Returns:
+            The path to the written parquet file. If the buffer is empty, returns None.
+        """
+        if len(self._buffer) == 0:
+            return None
+        try:
+            file_path = self.artifact_storage.write_batch_to_parquet_file(
+                batch_number=self._current_batch_number,
+                dataframe=pd.DataFrame(self._buffer),
+                batch_stage=BatchStage.PARTIAL_RESULT,
+            )
+            return file_path
+        except Exception as e:
+            raise DatasetBatchManagementError(f"🛑 Failed to write batch {self._current_batch_number}: {e}")
+
+    def update_record(self, index: int, record: dict) -> None:
+        if index < 0 or index >= len(self._buffer):
+            raise IndexError(f"🛑 Index {index} is out of bounds for buffer of size {len(self._buffer)}.")
+        self._buffer[index] = record
+
+    def update_records(self, records: list[dict]) -> None:
+        if len(records) != len(self._buffer):
+            raise DatasetBatchManagementError(
+                f"🛑 Number of records to update ({len(records)}) must match "
+                f"the number of records in the buffer ({len(self._buffer)})."
+            )
+        self._buffer = records

data_designer/engine/dataset_builders/utils/errors.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.engine.errors import DataDesignerError
+
+
+class DatasetBatchManagementError(DataDesignerError): ...
+
+
+class ConfigCompilationError(DataDesignerError): ...
+
+
+class DAGCircularDependencyError(DataDesignerError): ...

data_designer/engine/dataset_builders/utils/progress_tracker.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import time
+from threading import Lock
+
+from data_designer.logging import RandomEmoji
+
+logger = logging.getLogger(__name__)
+
+
+class ProgressTracker:
+    """
+    Thread-safe progress tracker for monitoring concurrent task completion.
+
+    Tracks completed, successful, and failed task counts and logs progress
+    at configurable intervals. Designed for use with ConcurrentThreadExecutor
+    to provide visibility into long-running batch operations.
+
+    Example usage:
+        tracker = ProgressTracker(total_records=100, label="LLM_TEXT column 'response'")
+        tracker.log_start(max_workers=8)
+
+        # In callbacks from ConcurrentThreadExecutor:
+        tracker.record_success()  # or tracker.record_failure()
+
+        # After executor completes:
+        tracker.log_final()
+    """
+
+    def __init__(self, total_records: int, label: str, log_interval_percent: int = 10):
+        """
+        Initialize the progress tracker.
+
+        Args:
+            total_records: Total number of records to process.
+            label: Human-readable label for log messages (e.g., "LLM_TEXT column 'response'").
+            log_interval_percent: How often to log progress as a percentage (default 10%).
+        """
+        self.total_records = total_records
+        self.label = label
+
+        self.completed = 0
+        self.success = 0
+        self.failed = 0
+
+        interval_fraction = max(1, log_interval_percent) / 100.0
+        self.log_interval = max(1, int(total_records * interval_fraction)) if total_records > 0 else 1
+        self.next_log_at = self.log_interval
+
+        self.start_time = time.perf_counter()
+        self.lock = Lock()
+        self._random_emoji = RandomEmoji()
+
+    def log_start(self, max_workers: int) -> None:
+        """Log the start of processing with worker count and interval information."""
+        logger.info(
+            "🐙 Processing %s with %d concurrent workers",
+            self.label,
+            max_workers,
+        )
+        logger.info(
+            "🧭 %s will report progress every %d record(s).",
+            self.label,
+            self.log_interval,
+        )
+
+    def record_success(self) -> None:
+        """Record a successful task completion and log progress if at interval."""
+        self._record_completion(success=True)
+
+    def record_failure(self) -> None:
+        """Record a failed task completion and log progress if at interval."""
+        self._record_completion(success=False)
+
+    def log_final(self) -> None:
+        """Log final progress summary."""
+        with self.lock:
+            if self.completed > 0:
+                self._log_progress_unlocked()
+
+    def _record_completion(self, *, success: bool) -> None:
+        should_log = False
+        with self.lock:
+            self.completed += 1
+            if success:
+                self.success += 1
+            else:
+                self.failed += 1
+
+            if self.completed >= self.next_log_at and self.completed < self.total_records:
+                should_log = True
+                while self.next_log_at <= self.completed:
+                    self.next_log_at += self.log_interval
+
+        if should_log:
+            with self.lock:
+                self._log_progress_unlocked()
+
+    def _log_progress_unlocked(self) -> None:
+        """Log current progress. Must be called while holding the lock."""
+        elapsed = time.perf_counter() - self.start_time
+        rate = self.completed / elapsed if elapsed > 0 else 0.0
+        remaining = max(0, self.total_records - self.completed)
+        eta = f"{(remaining / rate):.1f}s" if rate > 0 else "unknown"
+        percent = (self.completed / self.total_records) * 100 if self.total_records else 100.0
+
+        logger.info(
+            "  |-- %s %s progress: %d/%d (%.0f%%) complete, %d ok, %d failed, %.2f rec/s, eta %s",
+            self._random_emoji.progress(percent),
+            self.label,
+            self.completed,
+            self.total_records,
+            percent,
+            self.success,
+            self.failed,
+            rate,
+            eta,
+        )

data_designer/engine/errors.py

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+from data_designer.errors import DataDesignerError
+
+
+class DataDesignerRuntimeError(DataDesignerError): ...
+
+
+class UnknownModelAliasError(DataDesignerError): ...
+
+
+class UnknownProviderError(DataDesignerError): ...
+
+
+class NoModelProvidersError(DataDesignerError): ...
+
+
+class SecretResolutionError(DataDesignerError): ...
+
+
+class RemoteValidationSchemaError(DataDesignerError): ...
+
+
+class LocalCallableValidationError(DataDesignerError): ...
+
+
+class ErrorTrap(BaseModel):
+    error_count: int = 0
+    task_errors: dict[str, int] = Field(default_factory=dict)
+
+    def _track_error(self, error: DataDesignerError) -> None:
+        """
+        Track a specific error type.
+        """
+        error_type = type(error).__name__
+        if error_type not in self.task_errors:
+            self.task_errors[error_type] = 0
+        self.task_errors[error_type] += 1
+
+    def handle_error(self, error: Exception) -> None:
+        self.error_count += 1
+
+        if not isinstance(error, DataDesignerError):
+            error = DataDesignerError(str(error))
+
+        self._track_error(error)

data_designer/engine/model_provider.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from functools import cached_property
+
+from pydantic import BaseModel, field_validator, model_validator
+from typing_extensions import Self
+
+from data_designer.config.models import ModelProvider
+from data_designer.engine.errors import NoModelProvidersError, UnknownProviderError
+
+
+class ModelProviderRegistry(BaseModel):
+    providers: list[ModelProvider]
+    default: str | None = None
+
+    @field_validator("providers", mode="after")
+    @classmethod
+    def validate_providers_not_empty(cls, v: list[ModelProvider]) -> list[ModelProvider]:
+        if len(v) == 0:
+            raise ValueError("At least one model provider must be defined")
+        return v
+
+    @field_validator("providers", mode="after")
+    @classmethod
+    def validate_providers_have_unique_names(cls, v: list[ModelProvider]) -> list[ModelProvider]:
+        names = set()
+        dupes = set()
+        for provider in v:
+            if provider.name in names:
+                dupes.add(provider.name)
+            names.add(provider.name)
+
+        if len(dupes) > 0:
+            raise ValueError(f"Model providers must have unique names, found duplicates: {dupes}")
+        return v
+
+    @model_validator(mode="after")
+    def check_implicit_default(self) -> Self:
+        if self.default is None and len(self.providers) != 1:
+            raise ValueError("A default provider must be specified if multiple model providers are defined")
+        return self
+
+    @model_validator(mode="after")
+    def check_default_exists(self) -> Self:
+        if self.default and self.default not in self._providers_dict:
+            raise ValueError(f"Specified default {self.default!r} not found in providers list")
+        return self
+
+    def get_default_provider_name(self) -> str:
+        return self.default or self.providers[0].name
+
+    @cached_property
+    def _providers_dict(self) -> dict[str, ModelProvider]:
+        return {p.name: p for p in self.providers}
+
+    def get_provider(self, name: str | None) -> ModelProvider:
+        if name is None:
+            name = self.get_default_provider_name()
+
+        try:
+            return self._providers_dict[name]
+        except KeyError:
+            raise UnknownProviderError(f"No provider named {name!r} registered")
+
+
+def resolve_model_provider_registry(
+    model_providers: list[ModelProvider], default_provider_name: str | None = None
+) -> ModelProviderRegistry:
+    if len(model_providers) == 0:
+        raise NoModelProvidersError("At least one model provider must be defined")
+    return ModelProviderRegistry(
+        providers=model_providers,
+        default=default_provider_name or model_providers[0].name,
+    )

data_designer/engine/models/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0