data-designer 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

data_designer/interface/data_designer.py

@@ -20,7 +20,7 @@ from data_designer.config.models import (
     ModelProvider,
 )
 from data_designer.config.preview_results import PreviewResults
-from data_designer.config.seed import LocalSeedDatasetReference
+from data_designer.config.run_config import RunConfig
 from data_designer.config.utils.constants import (
     DEFAULT_NUM_RECORDS,
     MANAGED_ASSETS_PATH,
@@ -29,21 +29,23 @@ from data_designer.config.utils.constants import (
     PREDEFINED_PROVIDERS,
 )
 from data_designer.config.utils.info import InfoType, InterfaceInfo
-from data_designer.config.utils.io_helpers import write_seed_dataset
 from data_designer.engine.analysis.dataset_profiler import (
     DataDesignerDatasetProfiler,
     DatasetProfilerConfig,
 )
+from data_designer.engine.compiler import compile_data_designer_config
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.dataset_builders.column_wise_builder import ColumnWiseDatasetBuilder
 from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
 from data_designer.engine.model_provider import resolve_model_provider_registry
-from data_designer.engine.models.registry import create_model_registry
 from data_designer.engine.resources.managed_storage import init_managed_blob_storage
-from data_designer.engine.resources.resource_provider import ResourceProvider
-from data_designer.engine.resources.seed_dataset_data_store import (
-    HfHubSeedDatasetDataStore,
-    LocalSeedDatasetDataStore,
+from data_designer.engine.resources.resource_provider import ResourceProvider, create_resource_provider
+from data_designer.engine.resources.seed_reader import (
+    DataFrameSeedReader,
+    HuggingFaceSeedReader,
+    LocalFileSeedReader,
+    SeedReader,
+    SeedReaderRegistry,
 )
 from data_designer.engine.secret_resolver import (
     CompositeResolver,
@@ -61,6 +63,14 @@ from data_designer.logging import RandomEmoji
 
 DEFAULT_BUFFER_SIZE = 1000
 
+DEFAULT_SECRET_RESOLVER = CompositeResolver([EnvironmentResolver(), PlaintextResolver()])
+
+DEFAULT_SEED_READERS = [
+    HuggingFaceSeedReader(),
+    LocalFileSeedReader(),
+    DataFrameSeedReader(),
+]
+
 logger = logging.getLogger(__name__)
 
 
@@ -79,6 +89,7 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
             uses default providers.
         secret_resolver: Resolver for handling secrets and credentials. Defaults to
             EnvironmentResolver which reads secrets from environment variables.
+        seed_readers: Optional list of seed readers. If None, uses default readers.
         managed_assets_path: Path to the managed assets directory. This is used to point
             to the location of managed datasets and other assets used during dataset generation.
             If not provided, will check for an environment variable called DATA_DESIGNER_MANAGED_ASSETS_PATH.
@@ -92,52 +103,19 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         *,
         model_providers: list[ModelProvider] | None = None,
         secret_resolver: SecretResolver | None = None,
+        seed_readers: list[SeedReader] | None = None,
         managed_assets_path: Path | str | None = None,
     ):
-        self._secret_resolver = secret_resolver or CompositeResolver([EnvironmentResolver(), PlaintextResolver()])
+        self._secret_resolver = secret_resolver or DEFAULT_SECRET_RESOLVER
         self._artifact_path = Path(artifact_path) if artifact_path is not None else Path.cwd() / "artifacts"
         self._buffer_size = DEFAULT_BUFFER_SIZE
+        self._run_config = RunConfig()
         self._managed_assets_path = Path(managed_assets_path or MANAGED_ASSETS_PATH)
         self._model_providers = self._resolve_model_providers(model_providers)
         self._model_provider_registry = resolve_model_provider_registry(
             self._model_providers, get_default_provider_name()
         )
-
-    @staticmethod
-    def make_seed_reference_from_file(file_path: str | Path) -> LocalSeedDatasetReference:
-        """Create a seed dataset reference from an existing file.
-
-        Supported file extensions: .parquet (recommended), .csv, .json, .jsonl
-
-        Args:
-            file_path: Path to an existing dataset file.
-
-        Returns:
-            A LocalSeedDatasetReference pointing to the specified file.
-        """
-        return LocalSeedDatasetReference(dataset=str(file_path))
-
-    @classmethod
-    def make_seed_reference_from_dataframe(
-        cls, dataframe: pd.DataFrame, file_path: str | Path
-    ) -> LocalSeedDatasetReference:
-        """Create a seed dataset reference from a pandas DataFrame.
-
-        This method writes the DataFrame to disk and returns a reference that can
-        be passed to the config builder's `with_seed_dataset` method. If the file
-        already exists, it will be overwritten.
-
-        Supported file extensions: .parquet (recommended), .csv, .json, .jsonl
-
-        Args:
-            dataframe: Pandas DataFrame to use as seed data.
-            file_path: Path where to save dataset.
-
-        Returns:
-            A LocalSeedDatasetReference pointing to the written file.
-        """
-        write_seed_dataset(dataframe, Path(file_path))
-        return cls.make_seed_reference_from_file(file_path)
+        self._seed_reader_registry = SeedReaderRegistry(readers=seed_readers or DEFAULT_SEED_READERS)
 
     @property
     def info(self) -> InterfaceInfo:
@@ -274,6 +252,23 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
             config_builder=config_builder,
         )
 
+    def validate(self, config_builder: DataDesignerConfigBuilder) -> None:
+        """Validate the Data Designer configuration as defined by the DataDesignerConfigBuilder
+        with the configured engine components (SecretResolver, SeedReaders, etc.).
+
+        Args:
+            config_builder: The DataDesignerConfigBuilder containing the dataset
+                configuration (columns, constraints, seed data, etc.).
+
+        Returns:
+            None if the configuration is valid.
+
+        Raises:
+            InvalidConfigError: If the configuration is invalid.
+        """
+        resource_provider = self._create_resource_provider("validate-configuration", config_builder)
+        compile_data_designer_config(config_builder, resource_provider)
+
     def get_default_model_configs(self) -> list[ModelConfig]:
         """Get the default model configurations.
 
@@ -318,6 +313,20 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
             raise InvalidBufferValueError("Buffer size must be greater than 0.")
         self._buffer_size = buffer_size
 
+    def set_run_config(self, run_config: RunConfig) -> None:
+        """Set the runtime configuration for dataset generation.
+
+        Args:
+            run_config: A RunConfig instance containing runtime settings such as
+                early shutdown behavior. Import RunConfig from data_designer.essentials.
+
+        Example:
+            >>> from data_designer.essentials import DataDesigner, RunConfig
+            >>> dd = DataDesigner()
+            >>> dd.set_run_config(RunConfig(disable_early_shutdown=True))
+        """
+        self._run_config = run_config
+
     def _resolve_model_providers(self, model_providers: list[ModelProvider] | None) -> list[ModelProvider]:
         if model_providers is None:
             model_providers = get_default_providers()
@@ -334,11 +343,15 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         return model_providers or []
 
     def _create_dataset_builder(
-        self, config_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider
+        self,
+        config_builder: DataDesignerConfigBuilder,
+        resource_provider: ResourceProvider,
     ) -> ColumnWiseDatasetBuilder:
+        config = compile_data_designer_config(config_builder, resource_provider)
+
         return ColumnWiseDatasetBuilder(
-            column_configs=compile_dataset_builder_column_configs(config_builder.build(raise_exceptions=True)),
-            processor_configs=config_builder.get_processor_configs(),
+            column_configs=compile_dataset_builder_column_configs(config),
+            processor_configs=config.processors or [],
             resource_provider=resource_provider,
         )
 
@@ -356,24 +369,21 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
     def _create_resource_provider(
         self, dataset_name: str, config_builder: DataDesignerConfigBuilder
     ) -> ResourceProvider:
-        model_configs = config_builder.model_configs
         ArtifactStorage.mkdir_if_needed(self._artifact_path)
-        return ResourceProvider(
+
+        seed_dataset_source = None
+        if (seed_config := config_builder.get_seed_config()) is not None:
+            seed_dataset_source = seed_config.source
+
+        return create_resource_provider(
             artifact_storage=ArtifactStorage(artifact_path=self._artifact_path, dataset_name=dataset_name),
-            model_registry=create_model_registry(
-                model_configs=model_configs,
-                model_provider_registry=self._model_provider_registry,
-                secret_resolver=self._secret_resolver,
-            ),
+            model_configs=config_builder.model_configs,
+            secret_resolver=self._secret_resolver,
+            model_provider_registry=self._model_provider_registry,
             blob_storage=init_managed_blob_storage(str(self._managed_assets_path)),
-            datastore=(
-                LocalSeedDatasetDataStore()
-                if (settings := config_builder.get_seed_datastore_settings()) is None
-                else HfHubSeedDatasetDataStore(
-                    endpoint=settings.endpoint,
-                    token=settings.token,
-                )
-            ),
+            seed_dataset_source=seed_dataset_source,
+            seed_reader_registry=self._seed_reader_registry,
+            run_config=self._run_config,
         )
 
     def _get_interface_info(self, model_providers: list[ModelProvider]) -> InterfaceInfo:

data_designer/plugin_manager.py

@@ -50,7 +50,7 @@ class PluginManager:
         type_list = []
         for plugin in self._plugin_registry.get_plugins(PluginType.COLUMN_GENERATOR):
             if required_resources:
-                task_required_resources = plugin.task_cls.metadata().required_resources or []
+                task_required_resources = plugin.impl_cls.metadata().required_resources or []
                 if not all(resource in task_required_resources for resource in required_resources):
                     continue
             type_list.append(enum_type(plugin.name))

data_designer/plugins/errors.py

@@ -4,6 +4,9 @@
 from data_designer.errors import DataDesignerError
 
 
+class PluginLoadError(DataDesignerError): ...
+
+
 class PluginRegistrationError(DataDesignerError): ...
 
 

data_designer/plugins/plugin.py

@@ -1,14 +1,20 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+import ast
+import importlib
+import importlib.util
 from enum import Enum
+from functools import cached_property
 from typing import Literal, get_origin
 
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self
 
 from data_designer.config.base import ConfigBase
-from data_designer.engine.configurable_task import ConfigurableTask
+from data_designer.plugins.errors import PluginLoadError
 
 
 class PluginType(str, Enum):
@@ -26,11 +32,42 @@ class PluginType(str, Enum):
         return self.value.replace("-", " ")
 
 
+def _get_module_and_object_names(fully_qualified_object: str) -> tuple[str, str]:
+    try:
+        module_name, object_name = fully_qualified_object.rsplit(".", 1)
+    except ValueError:
+        # If fully_qualified_object does not have any periods, the rsplit call will return
+        # a list of length 1 and the variable assignment above will raise ValueError
+        raise PluginLoadError("Expected a fully-qualified object name, e.g. 'my_plugin.config.MyConfig'")
+
+    return module_name, object_name
+
+
+def _check_class_exists_in_file(filepath: str, class_name: str) -> None:
+    try:
+        with open(filepath, "r") as file:
+            source = file.read()
+    except FileNotFoundError:
+        raise PluginLoadError(f"Could not read source code at {filepath!r}")
+
+    tree = ast.parse(source)
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ClassDef) and node.name == class_name:
+            return None
+
+    raise PluginLoadError(f"Could not find class named {class_name!r} in {filepath!r}")
+
+
 class Plugin(BaseModel):
-    task_cls: type[ConfigurableTask]
-    config_cls: type[ConfigBase]
-    plugin_type: PluginType
-    emoji: str = "🔌"
+    impl_qualified_name: str = Field(
+        ...,
+        description="The fully-qualified name of the implementation class object, e.g. 'my_plugin.generator.MyColumnGenerator'",
+    )
+    config_qualified_name: str = Field(
+        ..., description="The fully-qualified name o the config class object, e.g. 'my_plugin.config.MyConfig'"
+    )
+    plugin_type: PluginType = Field(..., description="The type of plugin")
+    emoji: str = Field(default="🔌", description="The emoji to use in logs related to the plugin")
 
     @property
     def config_type_as_class_name(self) -> str:
@@ -48,22 +85,55 @@ class Plugin(BaseModel):
     def discriminator_field(self) -> str:
         return self.plugin_type.discriminator_field
 
+    @field_validator("impl_qualified_name", "config_qualified_name", mode="after")
+    @classmethod
+    def validate_class_name(cls, value: str) -> str:
+        module_name, object_name = _get_module_and_object_names(value)
+        try:
+            spec = importlib.util.find_spec(module_name)
+        except:
+            raise PluginLoadError(f"Could not find module {module_name!r}")
+
+        if spec is None or spec.origin is None:
+            raise PluginLoadError(f"Error finding source for module {module_name!r}")
+
+        _check_class_exists_in_file(spec.origin, object_name)
+
+        return value
+
     @model_validator(mode="after")
     def validate_discriminator_field(self) -> Self:
-        cfg = self.config_cls.__name__
+        _, cfg = _get_module_and_object_names(self.config_qualified_name)
         field = self.plugin_type.discriminator_field
         if field not in self.config_cls.model_fields:
-            raise ValueError(f"Discriminator field '{field}' not found in config class {cfg}")
+            raise ValueError(f"Discriminator field {field!r} not found in config class {cfg!r}")
         field_info = self.config_cls.model_fields[field]
         if get_origin(field_info.annotation) is not Literal:
-            raise ValueError(f"Field '{field}' of {cfg} must be a Literal type, not {field_info.annotation}.")
+            raise ValueError(f"Field {field!r} of {cfg!r} must be a Literal type, not {field_info.annotation!r}.")
         if not isinstance(field_info.default, str):
-            raise ValueError(f"The default of '{field}' must be a string, not {type(field_info.default)}.")
+            raise ValueError(f"The default of {field!r} must be a string, not {type(field_info.default)!r}.")
         enum_key = field_info.default.replace("-", "_").upper()
         if not enum_key.isidentifier():
             raise ValueError(
-                f"The default value '{field_info.default}' for discriminator field '{field}' "
-                f"cannot be converted to a valid enum key. The converted key '{enum_key}' "
+                f"The default value {field_info.default!r} for discriminator field {field!r} "
+                f"cannot be converted to a valid enum key. The converted key {enum_key!r} "
                 f"must be a valid Python identifier."
             )
         return self
+
+    @cached_property
+    def config_cls(self) -> type[ConfigBase]:
+        return self._load(self.config_qualified_name)
+
+    @cached_property
+    def impl_cls(self) -> type:
+        return self._load(self.impl_qualified_name)
+
+    @staticmethod
+    def _load(fully_qualified_object: str) -> type:
+        module_name, object_name = _get_module_and_object_names(fully_qualified_object)
+        module = importlib.import_module(module_name)
+        try:
+            return getattr(module, object_name)
+        except AttributeError:
+            raise PluginLoadError(f"Could not find class {object_name!r} in module {module_name!r}")

data_designer/plugins/testing/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.plugins.testing.utils import assert_valid_plugin
+
+__all__ = [
+    assert_valid_plugin.__name__,
+]

data_designer/plugins/testing/stubs.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Literal
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata
+from data_designer.engine.resources.resource_provider import ResourceType
+from data_designer.plugins.plugin import Plugin, PluginType
+
+MODULE_NAME = __name__
+
+
+class ValidTestConfig(SingleColumnConfig):
+    """Valid config for testing plugin creation."""
+
+    column_type: Literal["test-generator"] = "test-generator"
+    name: str
+
+
+class ValidTestTask(ConfigurableTask[ValidTestConfig]):
+    """Valid task for testing plugin creation."""
+
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_generator",
+            description="Test generator",
+            required_resources=None,
+        )
+
+
+class ConfigWithoutDiscriminator(ConfigBase):
+    some_field: str
+
+
+class ConfigWithStringField(ConfigBase):
+    column_type: str = "test-generator"
+
+
+class ConfigWithNonStringDefault(ConfigBase):
+    column_type: Literal["test-generator"] = 123  # type: ignore
+
+
+class ConfigWithInvalidKey(ConfigBase):
+    column_type: Literal["invalid-key-!@#"] = "invalid-key-!@#"
+
+
+class StubPluginConfigA(SingleColumnConfig):
+    column_type: Literal["test-plugin-a"] = "test-plugin-a"
+
+
+class StubPluginConfigB(SingleColumnConfig):
+    column_type: Literal["test-plugin-b"] = "test-plugin-b"
+
+
+class StubPluginTaskA(ConfigurableTask[StubPluginConfigA]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_plugin_a",
+            description="Test plugin A",
+            required_resources=None,
+        )
+
+
+class StubPluginTaskB(ConfigurableTask[StubPluginConfigB]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_plugin_b",
+            description="Test plugin B",
+            required_resources=None,
+        )
+
+
+# Stub plugins requiring different combinations of resources
+
+
+class StubPluginConfigModels(SingleColumnConfig):
+    column_type: Literal["test-plugin-models"] = "test-plugin-models"
+
+
+class StubPluginConfigModelsAndBlobs(SingleColumnConfig):
+    column_type: Literal["test-plugin-models-and-blobs"] = "test-plugin-models-and-blobs"
+
+
+class StubPluginConfigBlobsAndSeeds(SingleColumnConfig):
+    column_type: Literal["test-plugin-blobs-and-seeds"] = "test-plugin-blobs-and-seeds"
+
+
+class StubPluginTaskModels(ConfigurableTask[StubPluginConfigModels]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_plugin_models",
+            description="Test plugin requiring models",
+            required_resources=[ResourceType.MODEL_REGISTRY],
+        )
+
+
+class StubPluginTaskModelsAndBlobs(ConfigurableTask[StubPluginConfigModelsAndBlobs]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_plugin_models_and_blobs",
+            description="Test plugin requiring models and blobs",
+            required_resources=[ResourceType.MODEL_REGISTRY, ResourceType.BLOB_STORAGE],
+        )
+
+
+class StubPluginTaskBlobsAndSeeds(ConfigurableTask[StubPluginConfigBlobsAndSeeds]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_plugin_blobs_and_seeds",
+            description="Test plugin requiring blobs and seeds",
+            required_resources=[ResourceType.BLOB_STORAGE, ResourceType.SEED_READER],
+        )
+
+
+plugin_none = Plugin(
+    config_qualified_name=f"{MODULE_NAME}.StubPluginConfigA",
+    impl_qualified_name=f"{MODULE_NAME}.StubPluginTaskA",
+    plugin_type=PluginType.COLUMN_GENERATOR,
+)
+
+plugin_models = Plugin(
+    config_qualified_name=f"{MODULE_NAME}.StubPluginConfigModels",
+    impl_qualified_name=f"{MODULE_NAME}.StubPluginTaskModels",
+    plugin_type=PluginType.COLUMN_GENERATOR,
+)
+
+plugin_models_and_blobs = Plugin(
+    config_qualified_name=f"{MODULE_NAME}.StubPluginConfigModelsAndBlobs",
+    impl_qualified_name=f"{MODULE_NAME}.StubPluginTaskModelsAndBlobs",
+    plugin_type=PluginType.COLUMN_GENERATOR,
+)
+
+plugin_blobs_and_seeds = Plugin(
+    config_qualified_name=f"{MODULE_NAME}.StubPluginConfigBlobsAndSeeds",
+    impl_qualified_name=f"{MODULE_NAME}.StubPluginTaskBlobsAndSeeds",
+    plugin_type=PluginType.COLUMN_GENERATOR,
+)

data_designer/plugins/testing/utils.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.config.base import ConfigBase
+from data_designer.engine.configurable_task import ConfigurableTask
+from data_designer.plugins.plugin import Plugin
+
+
+def assert_valid_plugin(plugin: Plugin) -> None:
+    assert issubclass(plugin.config_cls, ConfigBase), "Plugin config class is not a subclass of ConfigBase"
+    assert issubclass(plugin.impl_cls, ConfigurableTask), "Plugin impl class is not a subclass of ConfigurableTask"

{data_designer-0.2.2.dist-info → data_designer-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer
-Version: 0.2.2
+Version: 0.3.0
 Summary: General framework for synthetic data generation
 License-Expression: Apache-2.0
 License-File: LICENSE
@@ -24,7 +24,7 @@ Requires-Dist: huggingface-hub<2,>=1.0.1
 Requires-Dist: jinja2<4,>=3.1.6
 Requires-Dist: json-repair<1,>=0.48.0
 Requires-Dist: jsonpath-rust-bindings<2,>=1.0
-Requires-Dist: litellm<2,>=1.73.6
+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: networkx<4,>=3.0
@@ -181,7 +181,7 @@ ModelConfig(
     alias="nv-reasoning",
     model="openai/gpt-oss-20b",
     provider="nvidia",
-    inference_parameters=InferenceParameters(
+    inference_parameters=ChatCompletionInferenceParams(
         temperature=0.3,
         top_p=0.9,
         max_tokens=4096,