from pydantic import (
Field,
model_validator,
field_validator,
PrivateAttr,
ConfigDict,
field_serializer,
)
from typing import Union, Optional, Dict, List, Any
from enum import Enum
import logging
from ...core.base.config_base import BasePipelineConfig
logger = logging.getLogger(__name__)
[docs]
class VariableType(str, Enum):
NUMERIC = "NUMERIC"
TEXT = "TEXT"
@classmethod
def _missing_(cls, value: str) -> Optional["VariableType"]:
"""Handle string values"""
try:
return cls(value.upper())
except ValueError:
return None
def __str__(self) -> str:
"""String representation"""
return self.value
[docs]
def create_inference_variable_list(
numeric_fields: List[str] = None,
text_fields: List[str] = None,
output_format: str = "dict",
) -> Union[Dict[str, Union[VariableType, str]], List[List[str]]]:
"""
Create an inference variable list for model input variables using separate lists for numeric and text fields.
This is a helper function that can be used standalone or within RegistrationConfig.
Args:
numeric_fields: List of field names that should be treated as NUMERIC
text_fields: List of field names that should be treated as TEXT
output_format: Format for storing variable list - either 'dict' or 'list'
Returns:
A dictionary mapping variable names to their types, or a list of [name, type] pairs,
depending on the output_format parameter
"""
# Initialize with empty lists if not provided
numeric_fields = numeric_fields or []
text_fields = text_fields or []
# Validate inputs are lists of strings
for field_name in numeric_fields:
if not isinstance(field_name, str):
raise ValueError(
f"Field name must be string, got {type(field_name)} for: {field_name}"
)
for field_name in text_fields:
if not isinstance(field_name, str):
raise ValueError(
f"Field name must be string, got {type(field_name)} for: {field_name}"
)
# Check for duplicates between numeric and text fields
common_fields = set(numeric_fields) & set(text_fields)
if common_fields:
raise ValueError(f"Fields cannot be both numeric and text: {common_fields}")
# Validate output format
if output_format not in ["dict", "list"]:
raise ValueError(
f"Output format must be 'dict' or 'list', got: {output_format}"
)
# Create the variable list in the requested format
if output_format == "dict":
# Dictionary format - map field names to their types
result = {}
# Add numeric fields
for field_name in numeric_fields:
result[field_name] = VariableType.NUMERIC
# Add text fields
for field_name in text_fields:
result[field_name] = VariableType.TEXT
else: # 'list' format
# List format - create list of [name, type] pairs
result = []
# Add numeric fields
for field_name in numeric_fields:
result.append([field_name, VariableType.NUMERIC.value])
# Add text fields
for field_name in text_fields:
result.append([field_name, VariableType.TEXT.value])
return result
[docs]
class RegistrationConfig(BasePipelineConfig):
"""
Configuration for model registration step, following the three-tier categorization:
Tier 1: Essential User Inputs - fields that users must explicitly provide
Tier 2: System Inputs - fields with reasonable defaults that users can override
Tier 3: Derived Fields - private fields with read-only property access
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
model_owner: str = Field(description="Team ID of model owner")
model_domain: str = Field(description="Domain for model registration")
model_objective: str = Field(description="Objective of model registration")
framework: str = Field(description="ML framework used for the model")
inference_entry_point: str = Field(description="Entry point script for inference")
source_model_inference_input_variable_list: Union[
Dict[str, Union[VariableType, str]], List[List[str]]
] = Field(
default_factory=dict,
description="Input variables and their types. Can be either:\n"
"1. Dictionary: {'var1': 'NUMERIC', 'var2': 'TEXT'}\n"
"2. List of pairs: [['var1', 'NUMERIC'], ['var2', 'TEXT']]",
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
inference_instance_type: str = Field(
default="ml.m5.large",
description="Instance type for inference endpoint/transform job",
)
source_model_inference_content_types: List[str] = Field(
default=["text/csv"],
description="Content type for model inference input. Must be exactly ['text/csv'] or ['application/json']",
)
source_model_inference_response_types: List[str] = Field(
default=["application/json"],
description="Response type for model inference output. Must be exactly ['text/csv'] or ['application/json']",
)
source_model_inference_output_variable_list: Dict[str, VariableType] = Field(
default={"legacy-score": VariableType.NUMERIC},
description="Dictionary of output variables and their types (NUMERIC or TEXT)",
)
# ===== Derived Fields (Tier 3) =====
# These are fields calculated from other fields, stored in private attributes
# with public read-only properties for access
# Removed _source_model_inference_input_variable_list as it's now a field
_variable_schema: Optional[Dict[str, Dict[str, List[Dict[str, str]]]]] = (
PrivateAttr(default=None)
)
# Update to Pydantic V2 style model_config
model_config = ConfigDict(
arbitrary_types_allowed=True,
validate_assignment=True,
extra="allow", # Accept metadata fields during deserialization
)
# Custom serializer for VariableType fields (Pydantic V2 approach)
[docs]
@field_serializer("source_model_inference_output_variable_list")
def serialize_output_variable_list(
self, value: Dict[str, VariableType]
) -> Dict[str, str]:
"""Serialize VariableType enum values to strings"""
return {
k: v.value if isinstance(v, VariableType) else v for k, v in value.items()
}
# ===== Property Accessors for Derived Fields =====
# (No property accessor needed for source_model_inference_input_variable_list since it's now a field)
# ===== Validators =====
[docs]
@field_validator("inference_instance_type")
@classmethod
def validate_inference_instance_type(cls, v: str) -> str:
"""Validate the inference instance type"""
if not v.startswith("ml."):
raise ValueError(
f"Invalid inference instance type: {v}. Must start with 'ml.'"
)
return v
[docs]
@field_validator("framework")
@classmethod
def validate_framework(cls, v: str) -> str:
"""Validate the ML framework"""
valid_frameworks = ["xgboost", "sklearn", "pytorch", "tensorflow"]
if v.lower() not in valid_frameworks:
raise ValueError(f"Framework must be one of {valid_frameworks}")
return v.lower()
# ===== Model Validation =====
[docs]
@model_validator(mode="after")
def initialize_derived_fields(self) -> "RegistrationConfig":
"""Initialize all derived fields once after validation."""
# Call parent validator first
super().initialize_derived_fields()
return self
[docs]
@model_validator(mode="after")
def validate_registration_configs(self) -> "RegistrationConfig":
"""Validate registration-specific configurations (without file existence checks)"""
# Removed file existence validation to improve configuration portability
# File validation should happen at execution time in builders, not at config creation time
# Only validate that source_dir is provided if inference_entry_point is a relative path
if self.inference_entry_point and not self.inference_entry_point.startswith(
"s3://"
):
if not self.source_dir:
raise ValueError(
"source_dir must be provided when inference_entry_point is a relative path"
)
return self
[docs]
@field_validator(
"source_model_inference_content_types", "source_model_inference_response_types"
)
@classmethod
def validate_content_types(cls, v: List[str]) -> List[str]:
"""Validate content and response types"""
valid_types = [["text/csv"], ["application/json"]]
if v not in valid_types:
raise ValueError(f"Content/Response types must be one of {valid_types}")
return v
[docs]
@field_validator("source_model_inference_output_variable_list")
@classmethod
def validate_output_variable_list(
cls, v: Dict[str, Union[VariableType, str]]
) -> Dict[str, str]:
"""Validate variable lists and convert to string values"""
if not v: # If empty dictionary
return v
result = {}
for key, value in v.items():
# Validate key is a string
if not isinstance(key, str):
raise ValueError(f"Key must be string, got {type(key)} for key: {key}")
# Convert VariableType to string or validate string value
if isinstance(value, VariableType):
result[key] = value.value
elif isinstance(value, str) and value in [vt.value for vt in VariableType]:
result[key] = value
else:
raise ValueError(
f"Value must be either 'NUMERIC' or 'TEXT', got: {value}"
)
return result
# ===== Property Accessors for Derived Method Results =====
@property
def variable_schema(self) -> Dict[str, Dict[str, List[Dict[str, str]]]]:
"""Generate variable schema for model registration"""
if self._variable_schema is None:
schema = {"input": {"variables": []}, "output": {"variables": []}}
# Handle input variables in either format
input_vars = self.source_model_inference_input_variable_list
if isinstance(input_vars, dict):
# Dictionary format
for var_name, var_type in input_vars.items():
schema["input"]["variables"].append(
{
"name": var_name,
"type": (
var_type
if isinstance(var_type, str)
else var_type.value
),
}
)
elif isinstance(input_vars, list):
# List format
for var_name, var_type in input_vars:
schema["input"]["variables"].append(
{"name": var_name, "type": var_type}
)
# Add output variables
for (
name,
var_type,
) in self.source_model_inference_output_variable_list.items():
schema["output"]["variables"].append(
{
"name": name,
"type": (
var_type if isinstance(var_type, str) else var_type.value
),
}
)
self._variable_schema = schema
return self._variable_schema
# ===== Legacy Methods =====
[docs]
def get_variable_schema(self) -> Dict[str, Dict[str, List[Dict[str, str]]]]:
"""Legacy method that forwards to the property"""
return self.variable_schema
# Removed get_script_path override - now inherits modernized version from BasePipelineConfig
# which includes hybrid resolution and comprehensive fallbacks
# Note: This config uses inference_entry_point instead of processing_entry_point,
# but the modernized base method can handle this through the comprehensive fallback strategy
# ===== Serialization =====
[docs]
def model_dump(self, **kwargs) -> Dict[str, Any]:
"""Override model_dump to include derived properties."""
data = super().model_dump(**kwargs)
# Add derived properties to output
data["variable_schema"] = self.variable_schema
# Process variable lists for proper serialization
if "source_model_inference_output_variable_list" in data:
data["source_model_inference_output_variable_list"] = {
k: v.value if isinstance(v, VariableType) else v
for k, v in data["source_model_inference_output_variable_list"].items()
}
if "source_model_inference_input_variable_list" in data:
input_vars = data["source_model_inference_input_variable_list"]
if isinstance(input_vars, dict):
data["source_model_inference_input_variable_list"] = {
k: v.value if isinstance(v, VariableType) else v
for k, v in input_vars.items()
}
return data
# ===== Methods for working with input variable lists =====