move classification functions to engine; remove required resources

Author: johnnygreco

docs/plugins/example.md

@@ -81,7 +81,6 @@ class IndexMultiplierColumnGenerator(ColumnGenerator[IndexMultiplierColumnConfig
             name="index-multiplier",
             description="Generates values by multiplying the row index by a user-specified multiplier",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=None,
         )
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:
@@ -110,7 +109,6 @@ class IndexMultiplierColumnGenerator(ColumnGenerator[IndexMultiplierColumnConfig
 - `metadata()` describes your generator and its requirements
 - `generation_strategy` can be `FULL_COLUMN`, `CELL_BY_CELL`
 - You have access to the configuration parameters via `self.config`
-- `required_resources` lists any required resources (models, artifact storages, etc.). This parameter will evolve in the near future, so keeping it as `None` is safe for now. That said, if your task will use the model registry, adding `data_designer.engine.resources.ResourceType.MODEL_REGISTRY` will enable automatic model health checking for your column generation task.
 
 !!! info "Understanding generation_strategy"
     The `generation_strategy` specifies how the column generator will generate data.
@@ -179,7 +177,6 @@ class IndexMultiplierColumnGenerator(ColumnGenerator[IndexMultiplierColumnConfig
             name="index-multiplier",
             description="Generates values by multiplying the row index by a user-specified multiplier",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=None,
         )
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:

src/data_designer/config/column_types.py

@@ -62,41 +62,6 @@
 )
 
 
-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,
-    }
-    model_generated_column_types.update(
-        plugin_manager.get_plugin_column_types(
-            DataDesignerColumnType,
-            required_resources=["model_registry"],
-        )
-    )
-    return column_type in model_generated_column_types
-
-
 def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType, **kwargs) -> ColumnConfigT:
     """Create a Data Designer column config object from kwargs.
 

src/data_designer/config/config_builder.py

@@ -18,7 +18,6 @@
 from data_designer.config.column_types import (
     ColumnConfigT,
     DataDesignerColumnType,
-    column_type_is_model_generated,
     get_column_config_from_kwargs,
     get_column_display_order,
 )
@@ -422,23 +421,6 @@ def get_constraints(self, target_column: str) -> list[ColumnConstraintT]:
         """
         return [c for c in self._constraints if c.target_column == target_column]
 
-    def get_llm_gen_columns(self) -> list[ColumnConfigT]:
-        """Get all model-generated column configurations.
-
-        Returns:
-            A list of column configurations that use model generation.
-        """
-        logger.warning("get_llm_gen_columns is deprecated. Use get_model_gen_columns instead.")
-        return self.get_model_gen_columns()
-
-    def get_model_gen_columns(self) -> list[ColumnConfigT]:
-        """Get all model-generated column configurations.
-
-        Returns:
-            A list of column configurations that use model generation.
-        """
-        return [c for c in self._column_configs.values() if column_type_is_model_generated(c.column_type)]
-
     def get_columns_of_type(self, column_type: DataDesignerColumnType) -> list[ColumnConfigT]:
         """Get all column configurations of the specified type.
 

src/data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -31,16 +31,11 @@
 )
 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__)
 
 
 class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
-    @staticmethod
-    def get_required_resources() -> list[ResourceType]:
-        return [ResourceType.MODEL_REGISTRY]
-
     @staticmethod
     def metadata() -> ColumnProfilerMetadata:
         return ColumnProfilerMetadata(

src/data_designer/engine/column_generators/generators/base.py

@@ -12,7 +12,6 @@
 import pandas as pd
 
 from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, DataT, TaskConfigT
-from data_designer.engine.resources.resource_provider import ResourceType
 
 if TYPE_CHECKING:
     from data_designer.config.models import BaseInferenceParams, ModelConfig
@@ -56,10 +55,6 @@ def generate(self, data: pd.DataFrame) -> pd.DataFrame: ...
     @abstractmethod
     def generate(self, data: DataT) -> DataT: ...
 
-    @staticmethod
-    def get_required_resources() -> list[ResourceType]:
-        return []
-
     def log_pre_generation(self) -> None:
         """A shared method to log info before the generator's `generate` method is called.
 
@@ -79,17 +74,10 @@ def generate_from_scratch(self, num_records: int) -> pd.DataFrame: ...
 
 
 class ColumnGeneratorWithModelRegistry(ColumnGenerator[TaskConfigT], ABC):
-    @staticmethod
-    def get_required_resources() -> list[ResourceType]:
-        return [ResourceType.MODEL_REGISTRY]
-
     @property
     def model_registry(self) -> ModelRegistry:
         return self.resource_provider.model_registry
 
-    def get_inference_parameters(self, model_alias: str) -> BaseInferenceParams:
-        return self.get_model_config(model_alias=model_alias).inference_parameters
-
     def get_model(self, model_alias: str) -> ModelFacade:
         return self.model_registry.get_model(model_alias=model_alias)
 
@@ -112,7 +100,7 @@ def model_config(self) -> ModelConfig:
 
     @functools.cached_property
     def inference_parameters(self) -> BaseInferenceParams:
-        return self.get_inference_parameters(model_alias=self.config.model_alias)
+        return self.model_config.inference_parameters
 
     def log_pre_generation(self) -> None:
         logger.info(f"Preparing {self.config.column_type} column generation")

src/data_designer/engine/column_generators/generators/samplers.py

@@ -6,7 +6,7 @@
 import logging
 import random
 from functools import partial
-from typing import TYPE_CHECKING, Callable
+from typing import Callable
 
 import pandas as pd
 
@@ -23,17 +23,10 @@
 from data_designer.engine.sampling_gen.entities.person import load_person_data_sampler
 from data_designer.engine.sampling_gen.generator import DatasetGenerator as SamplingDatasetGenerator
 
-if TYPE_CHECKING:
-    from data_designer.engine.resources.resource_provider import ResourceType
-
 logger = logging.getLogger(__name__)
 
 
 class SamplerColumnGenerator(FromScratchColumnGenerator[SamplerMultiColumnConfig]):
-    @staticmethod
-    def get_required_resources() -> list[ResourceType]:
-        return [ResourceType.BLOB_STORAGE]
-
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(

src/data_designer/engine/column_generators/generators/seed_dataset.py

@@ -19,18 +19,13 @@
 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
 
 logger = logging.getLogger(__name__)
 
 
 class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColumnConfig]):
-    @staticmethod
-    def get_required_resources() -> list[ResourceType]:
-        return [ResourceType.DATASTORE]
-
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(

src/data_designer/engine/column_generators/utils/generator_classification.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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

src/data_designer/engine/configurable_task.py

@@ -9,7 +9,7 @@
 
 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)
@@ -65,10 +65,6 @@ def resource_provider(self) -> ResourceProvider:
     @abstractmethod
     def metadata() -> ConfigurableTaskMetadata: ...
 
-    @staticmethod
-    @abstractmethod
-    def get_required_resources() -> list[ResourceType]: ...
-
     def _initialize(self) -> None:
         """An internal method for custom initialization logic, which will be called in the constructor."""
 

src/data_designer/engine/dataset_builders/column_wise_builder.py

@@ -13,7 +13,7 @@
 
 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,
@@ -25,6 +25,7 @@
     ColumnGeneratorWithSingleModel,
     GenerationStrategy,
 )
+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 (
@@ -42,7 +43,7 @@
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
 from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
-from data_designer.engine.resources.resource_provider import ResourceProvider, ResourceType
+from data_designer.engine.resources.resource_provider import ResourceProvider
 
 if TYPE_CHECKING:
     from data_designer.engine.models.usage import ModelUsageStats
@@ -192,7 +193,7 @@ def _run_from_scratch_column_generator(self, generator: ColumnGenerator) -> None
 
     def _run_cell_by_cell_generator(self, generator: ColumnGenerator) -> None:
         max_workers = MAX_CONCURRENCY_PER_NON_LLM_GENERATOR
-        if ResourceType.MODEL_REGISTRY in generator.get_required_resources():
+        if isinstance(generator, ColumnGeneratorWithSingleModel):
             max_workers = generator.inference_parameters.max_parallel_requests
         self._fan_out_with_threads(generator, max_workers=max_workers)