data-designer 0.3.8rc1__py3-none-any.whl → 0.4.0__py3-none-any.whl

data_designer/engine/dataset_builders/utils/dag.py

@@ -1,62 +0,0 @@
-# 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

@@ -1,200 +0,0 @@
-# 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

@@ -1,15 +0,0 @@
-# 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/errors.py

@@ -1,51 +0,0 @@
-# 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

@@ -1,77 +0,0 @@
-# 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

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