comments

Author: andreatgretel

src/data_designer/config/processors.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import json
 from abc import ABC
 from enum import Enum
 from typing import Any, Literal
@@ -9,6 +10,7 @@
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.dataset_builders import BuildStage
+from data_designer.config.errors import InvalidConfigError
 
 SUPPORTED_STAGES = [BuildStage.POST_BATCH]
 
@@ -72,3 +74,12 @@ class SchemaTransformProcessorConfig(ProcessorConfig):
         """,
     )
     processor_type: Literal[ProcessorType.SCHEMA_TRANSFORM] = ProcessorType.SCHEMA_TRANSFORM
+
+    @field_validator("template")
+    def validate_template(cls, v: dict[str, Any]) -> dict[str, Any]:
+        try:
+            json.dumps(v)
+        except TypeError as e:
+            if "not JSON serializable" in str(e):
+                raise InvalidConfigError("Template must be JSON serializable")
+        return v

src/data_designer/config/utils/validation.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import json
 from enum import Enum
 from string import Formatter
 from typing import Optional
@@ -37,7 +36,6 @@ class ViolationType(str, Enum):
     INVALID_COLUMN = "invalid_column"
     INVALID_MODEL_CONFIG = "invalid_model_config"
     INVALID_REFERENCE = "invalid_reference"
-    INVALID_TEMPLATE = "invalid_template"
     PROMPT_WITHOUT_REFERENCES = "prompt_without_references"
 
 
@@ -303,20 +301,6 @@ def validate_schema_transform_processor(
     all_column_names = {c.name for c in columns}
     for processor_config in processor_configs:
         if processor_config.processor_type == ProcessorType.SCHEMA_TRANSFORM:
-            try:
-                json.dumps(processor_config.template)
-            except TypeError as e:
-                if "not JSON serializable" in str(e):
-                    violations.append(
-                        Violation(
-                            column=None,
-                            type=ViolationType.INVALID_TEMPLATE,
-                            message=f"Ancillary dataset processor {processor_config.name} template is not a valid JSON object.",
-                            level=ViolationLevel.ERROR,
-                        )
-                    )
-                    continue
-
             for col, template in processor_config.template.items():
                 template_keywords = get_prompt_template_keywords(template)
                 invalid_keywords = set(template_keywords) - all_column_names

src/data_designer/engine/dataset_builders/artifact_storage.py

@@ -178,8 +178,6 @@ def write_batch_to_parquet_file(
         batch_stage: BatchStage,
         subfolder: str | None = None,
     ) -> Path:
-        if subfolder is None:
-            subfolder = ""
         file_path = self.create_batch_file_path(batch_number, batch_stage=batch_stage)
         self.write_parquet_file(file_path.name, dataframe, batch_stage, subfolder=subfolder)
         return file_path
@@ -191,8 +189,7 @@ def write_parquet_file(
         batch_stage: BatchStage,
         subfolder: str | None = None,
     ) -> Path:
-        if subfolder is None:
-            subfolder = ""
+        subfolder = subfolder or ""
         self.mkdir_if_needed(self._get_stage_path(batch_stage) / subfolder)
         file_path = self._get_stage_path(batch_stage) / subfolder / parquet_file_name
         dataframe.to_parquet(file_path, index=False)

tests/config/test_processors.py

@@ -5,6 +5,7 @@
 from pydantic import ValidationError
 
 from data_designer.config.dataset_builders import BuildStage
+from data_designer.config.errors import InvalidConfigError
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
     ProcessorConfig,
@@ -79,10 +80,16 @@ def test_schema_transform_processor_config_validation():
     with pytest.raises(ValidationError, match="Field required"):
         SchemaTransformProcessorConfig(name="schema_transform_processor", build_stage=BuildStage.POST_BATCH)
 
+    # Test invalid template raises error
+    with pytest.raises(InvalidConfigError, match="Template must be JSON serializable"):
+        SchemaTransformProcessorConfig(
+            name="schema_transform_processor", build_stage=BuildStage.POST_BATCH, template={"text": {1, 2, 3}}
+        )
 
-def test_output_format_processor_config_serialization():
+
+def test_schema_transform_processor_config_serialization():
     config = SchemaTransformProcessorConfig(
-        name="output_format_processor",
+        name="schema_transform_processor",
         build_stage=BuildStage.POST_BATCH,
         template={"text": "{{ col1 }}"},
     )

tests/config/utils/test_validation.py

@@ -111,11 +111,6 @@
         template={"text": "{{ invalid_reference }}"},
         build_stage=BuildStage.POST_BATCH,
     ),
-    SchemaTransformProcessorConfig(
-        name="schema_transform_processor_invalid_template",
-        template={"text": {1, 2, 3}},
-        build_stage=BuildStage.POST_BATCH,
-    ),
 ]
 ALLOWED_REFERENCE = [c.name for c in COLUMNS]
 
@@ -180,17 +175,11 @@ def test_validate_data_designer_config(
             type=ViolationType.INVALID_REFERENCE,
             message="Ancillary dataset processor attempts to reference columns 'invalid_reference' in the template for 'text', but the columns are not defined in the dataset.",
             level=ViolationLevel.ERROR,
-        ),
-        Violation(
-            column="text",
-            type=ViolationType.INVALID_TEMPLATE,
-            message="Ancillary dataset processor template is not a valid JSON object.",
-            level=ViolationLevel.ERROR,
-        ),
+        )
     ]
 
     violations = validate_data_designer_config(COLUMNS, PROCESSOR_CONFIGS, ALLOWED_REFERENCE)
-    assert len(violations) == 7
+    assert len(violations) == 6
     mock_validate_columns_not_all_dropped.assert_called_once()
     mock_validate_expression_references.assert_called_once()
     mock_validate_code_validation.assert_called_once()
@@ -283,21 +272,14 @@ def test_validate_expression_references():
 
 def test_validate_schema_transform_processor():
     violations = validate_schema_transform_processor(COLUMNS, PROCESSOR_CONFIGS)
-    assert len(violations) == 2
+    assert len(violations) == 1
     assert violations[0].type == ViolationType.INVALID_REFERENCE
     assert violations[0].column is None
     assert (
         violations[0].message
         == "Ancillary dataset processor attempts to reference columns 'invalid_reference' in the template for 'text', but the columns are not defined in the dataset."
     )
     assert violations[0].level == ViolationLevel.ERROR
-    assert violations[1].type == ViolationType.INVALID_TEMPLATE
-    assert violations[1].column is None
-    assert (
-        violations[1].message
-        == "Ancillary dataset processor schema_transform_processor_invalid_template template is not a valid JSON object."
-    )
-    assert violations[1].level == ViolationLevel.ERROR
 
 
 @patch("data_designer.config.utils.validation.Console.print")