chore: streamline generation metadata + consolidate sdg json (#226)

Author: nabinchha

src/data_designer/engine/compiler.py

@@ -6,7 +6,6 @@
 import logging
 
 from data_designer.config.column_configs import SeedDatasetColumnConfig
-from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.errors import InvalidConfigError
 from data_designer.engine.resources.resource_provider import ResourceProvider
@@ -16,13 +15,9 @@
 logger = logging.getLogger(__name__)
 
 
-def compile_data_designer_config(
-    config_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider
-) -> DataDesignerConfig:
-    config = config_builder.build()
+def compile_data_designer_config(config: DataDesignerConfig, resource_provider: ResourceProvider) -> DataDesignerConfig:
     _resolve_and_add_seed_columns(config, resource_provider.seed_reader)
     _validate(config)
-
     return config
 
 

src/data_designer/engine/dataset_builders/artifact_storage.py

@@ -24,6 +24,7 @@
 logger = logging.getLogger(__name__)
 
 BATCH_FILE_NAME_FORMAT = "batch_{batch_number:05d}.parquet"
+SDG_CONFIG_FILENAME = "sdg.json"
 
 
 class BatchStage(StrEnum):
@@ -170,12 +171,6 @@ def move_partial_result_to_final_file_path(self, batch_number: int) -> Path:
         shutil.move(partial_result_path, final_file_path)
         return final_file_path
 
-    def write_configs(self, json_file_name: str, configs: list[dict]) -> Path:
-        self.mkdir_if_needed(self.base_dataset_path)
-        with open(self.base_dataset_path / json_file_name, "w") as file:
-            json.dump([c.model_dump(mode="json") for c in configs], file, indent=4)
-        return self.base_dataset_path / json_file_name
-
     def write_batch_to_parquet_file(
         self,
         batch_number: int,
@@ -200,11 +195,89 @@ def write_parquet_file(
         dataframe.to_parquet(file_path, index=False)
         return file_path
 
+    def get_parquet_file_paths(self) -> list[str]:
+        """Get list of parquet file paths relative to base_dataset_path.
+
+        Returns:
+            List of relative paths to parquet files in the final dataset folder.
+        """
+        return [str(f.relative_to(self.base_dataset_path)) for f in sorted(self.final_dataset_path.glob("*.parquet"))]
+
+    def get_processor_file_paths(self) -> dict[str, list[str]]:
+        """Get processor output files organized by processor name.
+
+        Returns:
+            Dictionary mapping processor names to lists of relative file paths.
+        """
+        processor_files: dict[str, list[str]] = {}
+        if self.processors_outputs_path.exists():
+            for processor_dir in sorted(self.processors_outputs_path.iterdir()):
+                if processor_dir.is_dir():
+                    processor_name = processor_dir.name
+                    processor_files[processor_name] = [
+                        str(f.relative_to(self.base_dataset_path))
+                        for f in sorted(processor_dir.rglob("*"))
+                        if f.is_file()
+                    ]
+        return processor_files
+
+    def get_file_paths(self) -> dict[str, list[str] | dict[str, list[str]]]:
+        """Get all file paths organized by type.
+
+        Returns:
+            Dictionary with 'parquet-files' and 'processor-files' keys.
+        """
+        file_paths = {
+            "parquet-files": self.get_parquet_file_paths(),
+        }
+        processor_file_paths = self.get_processor_file_paths()
+        if processor_file_paths:
+            file_paths["processor-files"] = processor_file_paths
+
+        return file_paths
+
+    def read_metadata(self) -> dict:
+        """Read metadata from the metadata.json file.
+
+        Returns:
+            Dictionary containing the metadata.
+
+        Raises:
+            FileNotFoundError: If metadata file doesn't exist.
+        """
+        with open(self.metadata_file_path, "r") as file:
+            return json.load(file)
+
     def write_metadata(self, metadata: dict) -> Path:
+        """Write metadata to the metadata.json file.
+
+        Args:
+            metadata: Dictionary containing metadata to write.
+
+        Returns:
+            Path to the written metadata file.
+        """
         self.mkdir_if_needed(self.base_dataset_path)
         with open(self.metadata_file_path, "w") as file:
-            json.dump(metadata, file)
+            json.dump(metadata, file, indent=4, sort_keys=True)
         return self.metadata_file_path
 
+    def update_metadata(self, updates: dict) -> Path:
+        """Update existing metadata with new fields.
+
+        Args:
+            updates: Dictionary of fields to add/update in metadata.
+
+        Returns:
+            Path to the updated metadata file.
+        """
+        try:
+            existing_metadata = self.read_metadata()
+        except FileNotFoundError:
+            existing_metadata = {}
+
+        existing_metadata.update(updates)
+        return self.write_metadata(existing_metadata)
+
     def _get_stage_path(self, stage: BatchStage) -> Path:
         return getattr(self, resolve_string_enum(stage, BatchStage).value)

src/data_designer/engine/dataset_builders/column_wise_builder.py

@@ -13,6 +13,8 @@
 from typing import TYPE_CHECKING, Callable
 
 from data_designer.config.column_types import ColumnConfigT
+from data_designer.config.config_builder import BuilderConfig
+from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
@@ -25,13 +27,15 @@
     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.compiler import compile_data_designer_config
+from data_designer.engine.dataset_builders.artifact_storage import SDG_CONFIG_FILENAME, ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
-from data_designer.engine.dataset_builders.multi_column_configs import DatasetBuilderColumnConfigT, MultiColumnConfig
+from data_designer.engine.dataset_builders.multi_column_configs import MultiColumnConfig
 from data_designer.engine.dataset_builders.utils.concurrency import (
     MAX_CONCURRENCY_PER_NON_LLM_GENERATOR,
     ConcurrentThreadExecutor,
 )
+from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
 from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
@@ -54,17 +58,20 @@
 class ColumnWiseDatasetBuilder:
     def __init__(
         self,
-        column_configs: list[DatasetBuilderColumnConfigT],
-        processor_configs: list[ProcessorConfig],
+        data_designer_config: DataDesignerConfig,
         resource_provider: ResourceProvider,
         registry: DataDesignerRegistry | None = None,
     ):
         self.batch_manager = DatasetBatchManager(resource_provider.artifact_storage)
         self._resource_provider = resource_provider
         self._records_to_drop: set[int] = set()
         self._registry = registry or DataDesignerRegistry()
-        self._column_configs = column_configs
-        self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(processor_configs)
+
+        self._data_designer_config = compile_data_designer_config(data_designer_config, resource_provider)
+        self._column_configs = compile_dataset_builder_column_configs(self._data_designer_config)
+        self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(
+            self._data_designer_config.processors or []
+        )
         self._validate_column_configs()
 
     @property
@@ -91,9 +98,8 @@ def build(
         num_records: int,
         on_batch_complete: Callable[[Path], None] | None = None,
     ) -> Path:
-        self._write_configs()
         self._run_model_health_check_if_needed()
-
+        self._write_builder_config()
         generators = self._initialize_generators()
         start_time = time.perf_counter()
         group_id = uuid.uuid4().hex
@@ -152,6 +158,12 @@ def _initialize_generators(self) -> list[ColumnGenerator]:
             for config in self._column_configs
         ]
 
+    def _write_builder_config(self) -> None:
+        self.artifact_storage.mkdir_if_needed(self.artifact_storage.base_dataset_path)
+        BuilderConfig(data_designer=self._data_designer_config).to_json(
+            self.artifact_storage.base_dataset_path / SDG_CONFIG_FILENAME
+        )
+
     def _run_batch(
         self, generators: list[ColumnGenerator], *, batch_mode: str, save_partial_results: bool = True, group_id: str
     ) -> None:
@@ -303,16 +315,6 @@ def _worker_error_callback(self, exc: Exception, *, context: dict | None = None)
     def _worker_result_callback(self, result: dict, *, context: dict | None = None) -> None:
         self.batch_manager.update_record(context["index"], result)
 
-    def _write_configs(self) -> None:
-        self.artifact_storage.write_configs(
-            json_file_name="column_configs.json",
-            configs=self._column_configs,
-        )
-        self.artifact_storage.write_configs(
-            json_file_name="model_configs.json",
-            configs=self._resource_provider.model_registry.model_configs.values(),
-        )
-
     def _emit_batch_inference_events(
         self, batch_mode: str, usage_deltas: dict[str, ModelUsageStats], group_id: str
     ) -> None:

src/data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

@@ -91,8 +91,7 @@ def finish_batch(self, on_complete: Callable[[Path], None] | None = None) -> Pat
                     "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": [str(f) for f in sorted(self.artifact_storage.final_dataset_path.glob("*.parquet"))],
-                    "num_records": self.num_records_list[: self._current_batch_number + 1],
+                    "file_paths": self.artifact_storage.get_file_paths(),
                     "num_completed_batches": self._current_batch_number + 1,
                     "dataset_name": self.artifact_storage.dataset_name,
                 }

src/data_designer/interface/data_designer.py

@@ -9,6 +9,7 @@
 
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.default_model_settings import (
     get_default_model_configs,
     get_default_model_providers_missing_api_keys,
@@ -34,7 +35,6 @@
 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.resources.managed_storage import init_managed_blob_storage
 from data_designer.engine.resources.resource_provider import ResourceProvider, create_resource_provider
@@ -165,7 +165,7 @@ def create(
 
         resource_provider = self._create_resource_provider(dataset_name, config_builder)
 
-        builder = self._create_dataset_builder(config_builder, resource_provider)
+        builder = self._create_dataset_builder(config_builder.build(), resource_provider)
 
         try:
             builder.build(num_records=num_records)
@@ -183,6 +183,12 @@ def create(
 
         dataset_metadata = resource_provider.get_dataset_metadata()
 
+        # Update metadata with column statistics from analysis
+        if analysis:
+            builder.artifact_storage.update_metadata(
+                {"column_statistics": [stat.model_dump(mode="json") for stat in analysis.column_statistics]}
+            )
+
         return DatasetCreationResults(
             artifact_storage=builder.artifact_storage,
             analysis=analysis,
@@ -213,7 +219,7 @@ def preview(
         logger.info(f"{RandomEmoji.previewing()} Preview generation in progress")
 
         resource_provider = self._create_resource_provider("preview-dataset", config_builder)
-        builder = self._create_dataset_builder(config_builder, resource_provider)
+        builder = self._create_dataset_builder(config_builder.build(), resource_provider)
 
         try:
             raw_dataset = builder.build_preview(num_records=num_records)
@@ -277,7 +283,7 @@ def validate(self, config_builder: DataDesignerConfigBuilder) -> None:
             InvalidConfigError: If the configuration is invalid.
         """
         resource_provider = self._create_resource_provider("validate-configuration", config_builder)
-        compile_data_designer_config(config_builder, resource_provider)
+        compile_data_designer_config(config_builder.build(), resource_provider)
 
     def get_default_model_configs(self) -> list[ModelConfig]:
         """Get the default model configurations.
@@ -342,14 +348,11 @@ def _resolve_model_providers(self, model_providers: list[ModelProvider] | None)
 
     def _create_dataset_builder(
         self,
-        config_builder: DataDesignerConfigBuilder,
+        data_designer_config: DataDesignerConfig,
         resource_provider: ResourceProvider,
     ) -> ColumnWiseDatasetBuilder:
-        config = compile_data_designer_config(config_builder, resource_provider)
-
         return ColumnWiseDatasetBuilder(
-            column_configs=compile_dataset_builder_column_configs(config),
-            processor_configs=config.processors or [],
+            data_designer_config=data_designer_config,
             resource_provider=resource_provider,
         )
 

tests/engine/conftest.py

@@ -37,6 +37,7 @@ def stub_resource_provider(tmp_path, stub_model_facade):
     mock_provider.artifact_storage = ArtifactStorage(artifact_path=tmp_path)
     mock_provider.blob_storage = Mock(spec=ManagedBlobStorage)
     mock_provider.seed_reader = Mock()
+    mock_provider.seed_reader.get_column_names.return_value = []
     mock_provider.run_config = RunConfig()
     return mock_provider