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

data_designer/config/utils/visualization.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations

data_designer/config/validator_params.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum

data_designer/engine/__init__.py

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

data_designer/engine/analysis/column_profilers/base.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations

data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -31,7 +31,6 @@ from data_designer.engine.analysis.utils.judge_score_processing import (
 )
 from data_designer.engine.models.facade import ModelFacade
 from data_designer.engine.models.recipes.response_recipes import TextResponseRecipe
-from data_designer.engine.resources.resource_provider import ResourceType
 
 logger = logging.getLogger(__name__)
 
@@ -42,7 +41,6 @@ class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
         return ColumnProfilerMetadata(
             name="judge_score_profiler",
             description="Analyzes LLM-as-judge score distributions in a Data Designer dataset.",
-            required_resources=[ResourceType.MODEL_REGISTRY],
             applicable_column_types=[DataDesignerColumnType.LLM_JUDGE],
         )
 

data_designer/engine/analysis/column_profilers/registry.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from data_designer.config.analysis.column_profilers import ColumnProfilerType

data_designer/engine/analysis/column_statistics.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations

data_designer/engine/analysis/dataset_profiler.py

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

data_designer/engine/analysis/errors.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from data_designer.errors import DataDesignerError

data_designer/engine/analysis/utils/column_statistics_calculations.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations

data_designer/engine/analysis/utils/judge_score_processing.py

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

data_designer/engine/column_generators/__init__.py

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

data_designer/engine/column_generators/generators/__init__.py

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

data_designer/engine/column_generators/generators/base.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -16,6 +16,7 @@ from data_designer.engine.configurable_task import ConfigurableTask, Configurabl
 if TYPE_CHECKING:
     from data_designer.config.models import BaseInferenceParams, ModelConfig
     from data_designer.engine.models.facade import ModelFacade
+    from data_designer.engine.models.registry import ModelRegistry
 
 
 logger = logging.getLogger(__name__)
@@ -72,27 +73,38 @@ class FromScratchColumnGenerator(ColumnGenerator[TaskConfigT], ABC):
     def generate_from_scratch(self, num_records: int) -> pd.DataFrame: ...
 
 
-class WithModelGeneration:
+class ColumnGeneratorWithModelRegistry(ColumnGenerator[TaskConfigT], ABC):
+    @property
+    def model_registry(self) -> ModelRegistry:
+        return self.resource_provider.model_registry
+
+    def get_model(self, model_alias: str) -> ModelFacade:
+        return self.model_registry.get_model(model_alias=model_alias)
+
+    def get_model_config(self, model_alias: str) -> ModelConfig:
+        return self.model_registry.get_model_config(model_alias=model_alias)
+
+    def get_model_provider_name(self, model_alias: str) -> str:
+        provider = self.model_registry.get_model_provider(model_alias=model_alias)
+        return provider.name
+
+
+class ColumnGeneratorWithModel(ColumnGeneratorWithModelRegistry[TaskConfigT], ABC):
     @functools.cached_property
     def model(self) -> ModelFacade:
-        return self.resource_provider.model_registry.get_model(model_alias=self.config.model_alias)
+        return self.get_model(model_alias=self.config.model_alias)
 
     @functools.cached_property
     def model_config(self) -> ModelConfig:
-        return self.resource_provider.model_registry.get_model_config(model_alias=self.config.model_alias)
+        return self.get_model_config(model_alias=self.config.model_alias)
 
     @functools.cached_property
     def inference_parameters(self) -> BaseInferenceParams:
         return self.model_config.inference_parameters
 
     def log_pre_generation(self) -> None:
-        logger.info(f"Preparing {self.config.column_type} column generation")
-        logger.info(f"  |-- column name: {self.config.name!r}")
-        logger.info(f"  |-- model config:\n{self.model_config.model_dump_json(indent=4)}")
-        if self.model_config.provider is None:
-            logger.info(f"  |-- default model provider: {self._get_provider_name()!r}")
-
-    def _get_provider_name(self) -> str:
-        model_alias = self.model_config.alias
-        provider = self.resource_provider.model_registry.get_model_provider(model_alias=model_alias)
-        return provider.name
+        logger.info(f"{self.config.column_type} model configuration for generating column '{self.config.name}'")
+        logger.info(f"  |-- model: {self.model_config.model!r}")
+        logger.info(f"  |-- model alias: {self.config.model_alias!r}")
+        logger.info(f"  |-- model provider: {self.get_model_provider_name(model_alias=self.config.model_alias)!r}")
+        logger.info(f"  |-- inference parameters: {self.inference_parameters.format_for_display()}")

data_designer/engine/column_generators/generators/embedding.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 
@@ -6,13 +6,11 @@ from pydantic import BaseModel, computed_field
 
 from data_designer.config.column_configs import EmbeddingColumnConfig
 from data_designer.engine.column_generators.generators.base import (
-    ColumnGenerator,
+    ColumnGeneratorWithModel,
     GenerationStrategy,
     GeneratorMetadata,
-    WithModelGeneration,
 )
 from data_designer.engine.processing.utils import deserialize_json_values, parse_list_string
-from data_designer.engine.resources.resource_provider import ResourceType
 
 
 class EmbeddingGenerationResult(BaseModel):
@@ -27,14 +25,13 @@ class EmbeddingGenerationResult(BaseModel):
         return len(self.embeddings[0]) if len(self.embeddings) > 0 else 0
 
 
-class EmbeddingCellGenerator(WithModelGeneration, ColumnGenerator[EmbeddingColumnConfig]):
+class EmbeddingCellGenerator(ColumnGeneratorWithModel[EmbeddingColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
             name="embedding_cell_generator",
             description="Generate embeddings for a text column.",
             generation_strategy=GenerationStrategy.CELL_BY_CELL,
-            required_resources=[ResourceType.MODEL_REGISTRY],
         )
 
     def generate(self, data: dict) -> dict:

data_designer/engine/column_generators/generators/expression.py

@@ -1,6 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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 pandas as pd
@@ -25,7 +27,6 @@ class ExpressionColumnGenerator(WithJinja2UserTemplateRendering, ColumnGenerator
             name="expression_generator",
             description="Generate a column from a jinja2 expression.",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=None,
         )
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:

data_designer/engine/column_generators/generators/llm_completion.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 import functools
@@ -12,19 +12,18 @@ from data_designer.config.column_configs import (
 )
 from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
 from data_designer.engine.column_generators.generators.base import (
-    ColumnGenerator,
+    ColumnGeneratorWithModel,
     GenerationStrategy,
     GeneratorMetadata,
-    WithModelGeneration,
 )
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
     RecordBasedPromptRenderer,
     create_response_recipe,
 )
+from data_designer.engine.configurable_task import TaskConfigT
 from data_designer.engine.models.recipes.base import ResponseRecipe
 from data_designer.engine.processing.utils import deserialize_json_values
-from data_designer.engine.resources.resource_provider import ResourceType
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +32,7 @@ DEFAULT_MAX_CONVERSATION_RESTARTS = 5
 DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS = 0
 
 
-class WithChatCompletionGeneration(WithModelGeneration):
+class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfigT]):
     @functools.cached_property
     def response_recipe(self) -> ResponseRecipe:
         return create_response_recipe(self.config, self.model_config)
@@ -92,47 +91,43 @@ class WithChatCompletionGeneration(WithModelGeneration):
         return data
 
 
-class LLMTextCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMTextColumnConfig]):
+class LLMTextCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMTextColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
             name="llm_text_generator",
             description="Generate a new dataset cell from a prompt template",
             generation_strategy=GenerationStrategy.CELL_BY_CELL,
-            required_resources=[ResourceType.MODEL_REGISTRY],
         )
 
 
-class LLMCodeCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMCodeColumnConfig]):
+class LLMCodeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMCodeColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
             name="llm_code_generator",
             description="Generate a new dataset cell from a prompt template",
             generation_strategy=GenerationStrategy.CELL_BY_CELL,
-            required_resources=[ResourceType.MODEL_REGISTRY],
         )
 
 
-class LLMStructuredCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMStructuredColumnConfig]):
+class LLMStructuredCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMStructuredColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
             name="llm_structured_generator",
             description="Generate a new dataset cell from a prompt template",
             generation_strategy=GenerationStrategy.CELL_BY_CELL,
-            required_resources=[ResourceType.MODEL_REGISTRY],
         )
 
 
-class LLMJudgeCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMJudgeColumnConfig]):
+class LLMJudgeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMJudgeColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
             name="llm_judge_generator",
             description="Judge a new dataset cell based on a set of rubrics",
             generation_strategy=GenerationStrategy.CELL_BY_CELL,
-            required_resources=[ResourceType.MODEL_REGISTRY],
         )
 
     @property

data_designer/engine/column_generators/generators/samplers.py

@@ -1,6 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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 random
 from functools import partial
@@ -17,7 +19,6 @@ from data_designer.engine.column_generators.generators.base import (
 from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
 from data_designer.engine.processing.utils import concat_datasets
 from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
-from data_designer.engine.resources.resource_provider import ResourceType
 from data_designer.engine.sampling_gen.data_sources.sources import SamplerType
 from data_designer.engine.sampling_gen.entities.person import load_person_data_sampler
 from data_designer.engine.sampling_gen.generator import DatasetGenerator as SamplingDatasetGenerator
@@ -32,7 +33,6 @@ class SamplerColumnGenerator(FromScratchColumnGenerator[SamplerMultiColumnConfig
             name="sampler_column_generator",
             description="Generate columns using sampling-based method.",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=[ResourceType.BLOB_STORAGE],
         )
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:

data_designer/engine/column_generators/generators/seed_dataset.py

@@ -1,6 +1,9 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+
+from __future__ import annotations
+
 import functools
 import logging
 
@@ -16,7 +19,6 @@ from data_designer.engine.column_generators.generators.base import (
 from data_designer.engine.column_generators.utils.errors import SeedDatasetError
 from data_designer.engine.dataset_builders.multi_column_configs import SeedDatasetMultiColumnConfig
 from data_designer.engine.processing.utils import concat_datasets
-from data_designer.engine.resources.resource_provider import ResourceType
 
 MAX_ZERO_RECORD_RESPONSE_FACTOR = 2
 
@@ -30,7 +32,6 @@ class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColu
             name="seed_dataset_column_generator",
             description="Sample columns from a seed dataset.",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=[ResourceType.SEED_READER],
         )
 
     @property

data_designer/engine/column_generators/generators/validation.py

@@ -1,6 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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 pandas as pd
@@ -50,7 +52,6 @@ class ValidationColumnGenerator(ColumnGenerator[ValidationColumnConfig]):
             name="validate",
             description="Validate data.",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=None,
         )
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:

data_designer/engine/column_generators/registry.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from data_designer.config.base import ConfigBase

data_designer/engine/column_generators/utils/errors.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from data_designer.engine.errors import DataDesignerError

data_designer/engine/column_generators/utils/generator_classification.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.utils.type_helpers import resolve_string_enum
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModelRegistry
+from data_designer.plugin_manager import PluginManager
+
+plugin_manager = PluginManager()
+
+
+def column_type_used_in_execution_dag(column_type: str | DataDesignerColumnType) -> bool:
+    """Return True if the column type is used in the workflow execution DAG."""
+    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
+    dag_column_types = {
+        DataDesignerColumnType.EXPRESSION,
+        DataDesignerColumnType.LLM_CODE,
+        DataDesignerColumnType.LLM_JUDGE,
+        DataDesignerColumnType.LLM_STRUCTURED,
+        DataDesignerColumnType.LLM_TEXT,
+        DataDesignerColumnType.VALIDATION,
+        DataDesignerColumnType.EMBEDDING,
+    }
+    dag_column_types.update(plugin_manager.get_plugin_column_types(DataDesignerColumnType))
+    return column_type in dag_column_types
+
+
+def column_type_is_model_generated(column_type: str | DataDesignerColumnType) -> bool:
+    """Return True if the column type is a model-generated column."""
+    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
+    model_generated_column_types = {
+        DataDesignerColumnType.LLM_TEXT,
+        DataDesignerColumnType.LLM_CODE,
+        DataDesignerColumnType.LLM_STRUCTURED,
+        DataDesignerColumnType.LLM_JUDGE,
+        DataDesignerColumnType.EMBEDDING,
+    }
+    for plugin in plugin_manager.get_column_generator_plugins():
+        if issubclass(plugin.impl_cls, ColumnGeneratorWithModelRegistry):
+            model_generated_column_types.add(plugin.name)
+    return column_type in model_generated_column_types

data_designer/engine/column_generators/utils/judge_score_factory.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum

data_designer/engine/column_generators/utils/prompt_renderer.py

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

data_designer/engine/compiler.py

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

data_designer/engine/configurable_task.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
@@ -9,7 +9,7 @@ import pandas as pd
 
 from data_designer.config.base import ConfigBase
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
-from data_designer.engine.resources.resource_provider import ResourceProvider, ResourceType
+from data_designer.engine.resources.resource_provider import ResourceProvider
 
 DataT = TypeVar("DataT", dict, pd.DataFrame)
 TaskConfigT = TypeVar("ConfigT", bound=ConfigBase)
@@ -18,14 +18,12 @@ TaskConfigT = TypeVar("ConfigT", bound=ConfigBase)
 class ConfigurableTaskMetadata(ConfigBase):
     name: str
     description: str
-    required_resources: list[ResourceType] | None
 
 
 class ConfigurableTask(ABC, Generic[TaskConfigT]):
-    def __init__(self, config: TaskConfigT, *, resource_provider: ResourceProvider | None):
+    def __init__(self, config: TaskConfigT, resource_provider: ResourceProvider):
         self._config = self.get_config_type().model_validate(config)
         self._resource_provider = resource_provider
-        self._validate_resources()
         self._validate()
         self._initialize()
 
@@ -61,8 +59,6 @@ class ConfigurableTask(ABC, Generic[TaskConfigT]):
 
     @property
     def resource_provider(self) -> ResourceProvider:
-        if self._resource_provider is None:
-            raise ValueError(f"No resource provider provided for the `{self.metadata().name}` task.")
         return self._resource_provider
 
     @staticmethod
@@ -74,9 +70,3 @@ class ConfigurableTask(ABC, Generic[TaskConfigT]):
 
     def _validate(self) -> None:
         """An internal method for custom validation logic, which will be called in the constructor."""
-
-    def _validate_resources(self) -> None:
-        for resource in self.metadata().required_resources or []:
-            if resource is not None:
-                if getattr(self.resource_provider, ResourceType(resource).value) is None:
-                    raise ValueError(f"Resource {resource} is required for the `{self.metadata().name}`")

data_designer/engine/dataset_builders/artifact_storage.py

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

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -1,5 +1,6 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+
 from __future__ import annotations
 
 import functools
@@ -13,7 +14,7 @@ from typing import TYPE_CHECKING, Callable
 
 import pandas as pd
 
-from data_designer.config.column_types import ColumnConfigT, column_type_is_model_generated
+from data_designer.config.column_types import ColumnConfigT
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
@@ -22,9 +23,10 @@ from data_designer.config.processors import (
 )
 from data_designer.engine.column_generators.generators.base import (
     ColumnGenerator,
+    ColumnGeneratorWithModel,
     GenerationStrategy,
-    WithModelGeneration,
 )
+from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
 from data_designer.engine.dataset_builders.multi_column_configs import (
@@ -45,6 +47,7 @@ from data_designer.engine.registry.data_designer_registry import DataDesignerReg
 from data_designer.engine.resources.resource_provider import ResourceProvider
 
 if TYPE_CHECKING:
+    from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModelRegistry
     from data_designer.engine.models.usage import ModelUsageStats
 
 logger = logging.getLogger(__name__)
@@ -192,7 +195,7 @@ class ColumnWiseDatasetBuilder:
 
     def _run_cell_by_cell_generator(self, generator: ColumnGenerator) -> None:
         max_workers = MAX_CONCURRENCY_PER_NON_LLM_GENERATOR
-        if isinstance(generator, WithModelGeneration):
+        if isinstance(generator, ColumnGeneratorWithModel):
             max_workers = generator.inference_parameters.max_parallel_requests
         self._fan_out_with_threads(generator, max_workers=max_workers)
 
@@ -206,7 +209,7 @@ class ColumnWiseDatasetBuilder:
                 list(set(config.model_alias for config in self.llm_generated_column_configs))
             )
 
-    def _fan_out_with_threads(self, generator: WithModelGeneration, max_workers: int) -> None:
+    def _fan_out_with_threads(self, generator: ColumnGeneratorWithModelRegistry, max_workers: int) -> None:
         if generator.generation_strategy != GenerationStrategy.CELL_BY_CELL:
             raise DatasetGenerationError(
                 f"Generator {generator.metadata().name} is not a {GenerationStrategy.CELL_BY_CELL} "

data_designer/engine/dataset_builders/errors.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from data_designer.engine.errors import DataDesignerError

data_designer/engine/dataset_builders/multi_column_configs.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC

data_designer/engine/dataset_builders/utils/__init__.py

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

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations

data_designer/engine/dataset_builders/utils/config_compiler.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from data_designer.config.column_types import DataDesignerColumnType

data_designer/engine/dataset_builders/utils/dag.py

@@ -1,11 +1,12 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
 
 import networkx as nx
 
-from data_designer.config.column_types import ColumnConfigT, column_type_used_in_execution_dag
+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
 
 logger = logging.getLogger(__name__)

data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

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

data_designer/engine/dataset_builders/utils/errors.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from data_designer.engine.errors import DataDesignerError