Add Pydantic to SDG for stronger validation, error handling, and extensibility (#931)

Signed-off-by: Ian Hoang <[email protected]>

Author: IanHoang

.pylintrc

@@ -166,14 +166,21 @@ disable=print-statement,
         too-many-instance-attributes,
         too-many-statements,
         inconsistent-return-statements,
-        C0302,
+        too-many-lines,
         C4001,
-        R0916,
-        W0201,
-        W0613,
-        W0621,
+        too-many-boolean-expressions,
+        attribute-defined-outside-init,
+        unused-argument,
+        redefined-outer-name,
         invalid-docstring-quote,
-        raise-missing-from
+        raise-missing-from,
+        consider-using-with,
+        duplicate-code,
+        consider-using-from-import,
+        bad-option-value,
+        consider-using-dict-items,
+        unused-private-member,
+        use-a-generator
 
 
 # Enable the message, report, category or checker with the given id(s). You can

osbenchmark/synthetic_data_generator/helpers.py

@@ -17,15 +17,15 @@
 import yaml
 
 from osbenchmark.utils import console
-from osbenchmark.exceptions import SystemSetupError, ExecutorError
+from osbenchmark import exceptions
 from osbenchmark.synthetic_data_generator.strategies.strategy import DataGenerationStrategy
-from osbenchmark.synthetic_data_generator.types import DEFAULT_GENERATION_SETTINGS, SyntheticDataGeneratorMetadata, GB_TO_BYTES
+from osbenchmark.synthetic_data_generator.models import SyntheticDataGeneratorMetadata, SDGConfig, GB_TO_BYTES
 
 def load_user_module(file_path):
     allowed_extensions = ['.py']
     extension = os.path.splitext(file_path)[1]
     if extension not in allowed_extensions:
-        raise SystemSetupError(f"User provided module with file extension [{extension}]. Python modules must have {allowed_extensions} extension.")
+        raise exceptions.SystemSetupError(f"User provided module with file extension [{extension}]. Python modules must have {allowed_extensions} extension.")
 
     spec = importlib.util.spec_from_file_location("user_module", file_path)
     user_module = importlib.util.module_from_spec(spec)
@@ -82,7 +82,7 @@ def remove_existing_files(existing_files_found: list):
         for file in existing_files_found:
             os.remove(file)
     except OSError as e:
-        raise ExecutorError("OSB could not remove existing files for SDG: ", e)
+        raise exceptions.ExecutorError("OSB could not remove existing files for SDG: ", e)
 
 def host_has_available_disk_storage(sdg_metadata: SyntheticDataGeneratorMetadata) -> bool:
     logger = logging.getLogger(__name__)
@@ -98,18 +98,23 @@ def host_has_available_disk_storage(sdg_metadata: SyntheticDataGeneratorMetadata
         logger.error("Error checking disk space.")
         return False
 
-def load_config(config_path):
+def load_config(config_path: str) -> SDGConfig:
     try:
         allowed_extensions = ['.yml', '.yaml']
 
         extension = os.path.splitext(config_path)[1]
         if extension not in allowed_extensions:
-            raise SystemSetupError(f"User provided config with extension [{extension}]. Config must have a {allowed_extensions} extension.")
+            raise exceptions.ConfigError(f"User provided config with extension [{extension}]. Config must have a {allowed_extensions} extension.")
         else:
             with open(config_path, 'r') as file:
-                return yaml.safe_load(file)
+                config_details = yaml.safe_load(file)
+
+        return SDGConfig(**config_details) if config_details else SDGConfig()
+
+    except yaml.YAMLError as e:
+        raise exceptions.ConfigError(f"Error when loading config due to YAML error: {e}")
     except TypeError:
-        raise SystemSetupError("Error when loading config. Please ensure that the proper config was provided")
+        raise exceptions.SystemSetupError("Error when loading config. Please ensure that the proper config was provided")
 
 def write_chunk(data, file_path):
     written_bytes = 0
@@ -130,29 +135,6 @@ def calculate_avg_doc_size(strategy: DataGenerationStrategy):
 
     return size
 
-def get_generation_settings(input_config: dict) -> dict:
-    '''
-    Grabs the user's config's generation settings and compares it with the default generation settings.
-    If there are missing fields in the user's config, it populates it with the default values
-    '''
-    generation_settings = DEFAULT_GENERATION_SETTINGS
-    if input_config is None: # if user did not provide a custom config
-        return generation_settings
-
-    user_generation_settings = input_config.get('settings', {})
-
-    if not user_generation_settings: # If user provided custom config but did not include settings
-        return generation_settings
-    else:
-        # Traverse and update valid settings that user specified.
-        for k in generation_settings:
-            if k in user_generation_settings and user_generation_settings[k] is not None:
-                generation_settings[k] = user_generation_settings[k]
-            else:
-                continue
-
-        return generation_settings
-
 def format_size(bytes):
     for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
         if bytes < 1024:

osbenchmark/synthetic_data_generator/input_processor.py

@@ -8,7 +8,7 @@
 
 import logging
 
-from osbenchmark.synthetic_data_generator.types import SyntheticDataGeneratorMetadata
+from osbenchmark.synthetic_data_generator.models import SyntheticDataGeneratorMetadata
 from osbenchmark.exceptions import ConfigError
 
 logger = logging.getLogger(__name__)

osbenchmark/synthetic_data_generator/models.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+# Modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+
+import os
+from typing import Optional, Dict, List, Any, Union
+import re
+
+from pydantic import BaseModel, Field, field_validator
+
+GB_TO_BYTES = 1024 ** 3
+
+class SettingsConfig(BaseModel):
+    workers: Optional[int] = Field(default_factory=os.cpu_count) # Number of workers recommended to not exceed CPU count
+    max_file_size_gb: Optional[int] = 40 # Default because some CloudProviders limit the size of files stored
+    docs_per_chunk: Optional[int] = 10000 # Default based on testing
+
+    # pylint: disable = no-self-argument
+    @field_validator('workers', 'max_file_size_gb', 'docs_per_chunk')
+    def validate_values_are_positive_integers(cls, v):
+        if v is not None and v <= 0:
+            raise ValueError(f"Value '{v}' in Settings portion must be a positive integer.")
+
+        return v
+
+class CustomGenerationValuesConfig(BaseModel):
+    custom_lists: Optional[Dict[str, List[Any]]] = None
+    custom_providers: Optional[List[Any]] = None
+
+    # pylint: disable = no-self-argument
+    @field_validator('custom_lists')
+    def validate_custom_lists(cls, v):
+        if v is not None:
+            for key, value in v.items():
+                if not isinstance(key, str):
+                    raise ValueError(f"All keys within custom_lists of CustomGenerationValues section must be strings. '{key}' is not a string")
+                if not isinstance(value, list):
+                    raise ValueError(f"Value for key '{key}' must be a list.")
+        return v
+
+class GeneratorParams(BaseModel):
+    # Integer / Long Params
+    min: Optional[Union[int, float]] = None
+    max: Optional[Union[int, float]] = None
+
+    # Date Params
+    start_date: Optional[str] = None
+    end_date: Optional[str] = None
+    format: Optional[str] = None
+
+    # Text / Keywords Params
+    must_include: Optional[List[str]] = None
+    choices: Optional[List[str]] = None
+
+    class Config:
+        extra = 'forbid'
+
+class FieldOverride(BaseModel):
+    generator: str
+    params: GeneratorParams
+
+    # pylint: disable = no-self-argument
+    @field_validator('generator')
+    def validate_generator_name(cls, v):
+        valid_generators = [
+            'generate_text',
+            'generate_keyword',
+            'generate_integer',
+            'generate_long',
+            'generate_short',
+            'generate_byte',
+            'generate_float',
+            'generate_double',
+            'generate_boolean',
+            'generate_date',
+            'generate_ip',
+            'generate_geopoint',
+            'generate_object',
+            'generate_nested'
+        ]
+
+        if v not in valid_generators:
+            raise ValueError(f"Generator '{v}' mentioned in FieldOverrides not among valid generators: {valid_generators}")
+        return v
+
+class MappingGenerationValuesConfig(BaseModel):
+    generator_overrides: Optional[Dict[str, GeneratorParams]] = None
+    field_overrides: Optional[Dict[str, FieldOverride]] = None
+
+    # pylint: disable = no-self-argument
+    @field_validator('generator_overrides')
+    def validate_generator_types(cls, v):
+        if v is not None:
+            valid_generator_types = ['integer', 'long', 'float', 'double', 'date', 'text', 'keyword', 'short', 'byte', 'ip', 'geopoint', 'nested', 'boolean']
+
+            for generator_type in v.keys():
+                if generator_type not in valid_generator_types:
+                    raise ValueError(f"Invalid Generator Type '{generator_type}. Must be one of: {valid_generator_types}'")
+
+        return v
+
+    # pylint: disable = no-self-argument
+    @field_validator('field_overrides')
+    def validate_field_names(cls, v):
+        if v is not None:
+            for field_name in v.keys():
+                if not re.match(r'^[a-zA-Z][a-zA-Z0-9_.]*$', field_name):
+                    raise ValueError(f"Invalid Field Name '{field_name}' in FieldOverrides. Only alphanumeric characters, underscores and periods are allowed.")
+
+        return v
+
+class SyntheticDataGeneratorMetadata(BaseModel):
+    index_name: Optional[str] = None
+    index_mappings_path: Optional[str] = None
+    custom_module_path: Optional[str] = None
+    custom_config_path: Optional[str] = None
+    output_path: Optional[str] = None
+    total_size_gb: Optional[int] = None
+
+    class Config:
+        extra = 'forbid'
+
+class SDGConfig(BaseModel):
+    # If user does not provide YAML fil or provides YAML without all settings fields, it will use default generation settings.
+    settings: Optional[SettingsConfig] = Field(default_factory=SettingsConfig)
+    CustomGenerationValues: Optional[CustomGenerationValuesConfig] = None
+    MappingGenerationValues: Optional[MappingGenerationValuesConfig] = None
+
+    class Config:
+        extra = 'forbid'

osbenchmark/synthetic_data_generator/strategies/custom_module_strategy.py

@@ -16,29 +16,38 @@
 from mimesis.random import Random
 from mimesis.providers.base import BaseProvider
 
-from osbenchmark.exceptions import ConfigError
+from osbenchmark import exceptions
 from osbenchmark.synthetic_data_generator.strategies import DataGenerationStrategy
-from osbenchmark.synthetic_data_generator.types import SyntheticDataGeneratorMetadata
+from osbenchmark.synthetic_data_generator.models import SyntheticDataGeneratorMetadata, SDGConfig
 
 class CustomModuleStrategy(DataGenerationStrategy):
-    def __init__(self, sdg_metadata: SyntheticDataGeneratorMetadata,  sdg_config: dict, custom_module: ModuleType) -> None:
+    def __init__(self, sdg_metadata: SyntheticDataGeneratorMetadata,  sdg_config: SDGConfig, custom_module: ModuleType) -> None:
         self.sdg_metadata = sdg_metadata
         self.sdg_config = sdg_config
         self.custom_module = custom_module
         self.logger = logging.getLogger(__name__)
 
         if not hasattr(self.custom_module, 'generate_synthetic_document'):
             msg = f"Custom module at [{self.sdg_metadata.custom_module_path}] does not define a function called generate_synthetic_document(). Ensure that this method is defined."
-            raise ConfigError(msg)
+            raise exceptions.ConfigError(msg)
 
         # Fetch settings and custom module components from sdg-config.yml
-        custom_module_values = self.sdg_config.get('CustomGenerationValues', {})
-        try:
-            self.custom_lists = custom_module_values.get('custom_lists', {})
-            self.custom_providers = {name: getattr(self.custom_module, name) for name in custom_module_values.get('custom_providers', [])}
-        except TypeError:
-            msg = "Synthetic Data Generator Config has custom_lists and custom_providers pointing to null values. Either populate or remove."
-            raise ConfigError(msg)
+        if self.sdg_config.CustomGenerationValues is None:
+            self.custom_lists = {}
+            self.custom_providers = {}
+        else:
+            try:
+                self.custom_lists = self.sdg_config.CustomGenerationValues.custom_lists or {}
+                provider_names = self.sdg_config.CustomGenerationValues.custom_providers or []
+                self.custom_providers = {
+                    name: getattr(self.custom_module, name) for name in provider_names
+                }
+            except AttributeError as e:
+                msg = f"Error when setting up custom lists and custom providers: {e}"
+                raise exceptions.ConfigError(msg)
+            except TypeError:
+                msg = "Synthetic Data Generator Config has custom_lists and custom_providers pointing to null values. Either populate or remove."
+
 
     # pylint: disable=arguments-differ
     def generate_data_chunks_across_workers(self, dask_client: Client, docs_per_chunk: int, seeds: list ) -> list:
@@ -51,7 +60,7 @@ def generate_data_chunks_across_workers(self, dask_client: Client, docs_per_chun
             self.generate_data_chunk_from_worker, self.custom_module.generate_synthetic_document,
             docs_per_chunk, seed) for seed in seeds]
 
-
+    # pylint: disable=arguments-renamed
     def generate_data_chunk_from_worker(self, generate_synthetic_document: Callable, docs_per_chunk: int, seed: Optional[int]) -> list:
         """
         This method is submitted to Dask worker and can be thought of as the worker performing a job, which is calling the
@@ -83,7 +92,7 @@ def generate_test_document(self):
             msg = "Encountered AttributeError when setting up custom_providers and custom_lists. " + \
                     "It seems that your module might be using custom_lists and custom_providers." + \
                     f"Please ensure you have provided a custom config with custom_providers and custom_lists: {e}"
-            raise ConfigError(msg)
+            raise exceptions.ConfigError(msg)
         return document
 
     def _instantiate_all_providers(self, custom_providers):

osbenchmark/synthetic_data_generator/strategies/mapping_strategy.py

@@ -19,14 +19,14 @@
 
 from osbenchmark.exceptions import ConfigError, MappingsError
 from osbenchmark.synthetic_data_generator.strategies import DataGenerationStrategy
-from osbenchmark.synthetic_data_generator.types import SyntheticDataGeneratorMetadata
+from osbenchmark.synthetic_data_generator.models import SyntheticDataGeneratorMetadata, SDGConfig, MappingGenerationValuesConfig
 
 class MappingStrategy(DataGenerationStrategy):
-    def __init__(self, sdg_metadata: SyntheticDataGeneratorMetadata,  sdg_config: dict, index_mapping: dict) -> None:
+    def __init__(self, sdg_metadata: SyntheticDataGeneratorMetadata,  sdg_config: SDGConfig, index_mapping: dict) -> None:
         self.sdg_metadata = sdg_metadata
         self.sdg_config = sdg_config # Optional YAML-based config for value constraints
         self.index_mapping = index_mapping # OpenSearch Mapping
-        self.mapping_generation_values = self.sdg_config.get("MappingGenerationValues", {}) if self.sdg_config else {}
+        self.mapping_generation_values =  (self.sdg_config.MappingGenerationValues or {}) if self.sdg_config else {}
 
         self.logger = logging.getLogger(__name__)
 
@@ -224,8 +224,12 @@ def transform_mapping_to_generators(self, mapping_dict: Dict[str, Any], field_pa
         transformed_mapping = {}
 
         # Extract configuration settings (both default generators and field overrides) from config user provided
-        # TODO: Set self.mapping_config to automatically point to MappingConverter
-        config = self.mapping_config
+        # Convert the sdg_config's MappingGenerationValues section into a dictionary and access the overrides
+        if isinstance(self.mapping_config, MappingGenerationValuesConfig):
+            config = self.mapping_config.model_dump()
+        else:
+            config = self.mapping_config
+
         generator_overrides = config.get("generator_overrides", {})
         field_overrides = config.get("field_overrides", {})