from pydantic import (
Field,
model_validator,
field_validator,
ConfigDict,
field_serializer,
)
from typing import Optional, Dict, List, Union, TYPE_CHECKING, ClassVar
from pathlib import Path
import logging
logger = logging.getLogger(__name__)
from .config_processing_step_base import ProcessingStepConfigBase
# Import for type hints only
if TYPE_CHECKING:
pass
[docs]
class PayloadConfig(ProcessingStepConfigBase):
"""
Configuration for payload generation and testing.
This configuration follows the three-tier field categorization:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that users can override
3. Tier 3: Derived Fields - fields calculated from other fields, stored in private attributes
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
# NOTE: Variable lists removed - script gets all variable information from hyperparameters.json
# The script extracts field information from hyperparameters using:
# - full_field_list, tab_field_list, cat_field_list from hyperparameters
# - Creates variable list dynamically using create_model_variable_list()
# Config-based variable lists were redundant and unused.
# Performance metrics
expected_tps: int = Field(ge=1, description="Expected transactions per second")
max_latency_in_millisecond: int = Field(
ge=100, le=10000, description="Maximum acceptable latency in milliseconds"
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
# Entry point script
processing_entry_point: str = Field(
default="payload.py", description="Entry point script for payload generation"
)
# Content and response types
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']",
)
# Default values for payload generation
default_numeric_value: float = Field(
default=0.0, description="Default value for numeric fields"
)
default_text_value: str = Field(
default="DEFAULT_TEXT", description="Default value for text fields"
)
# NEW: Unified field defaults for multi-modal support
field_defaults: Optional[Dict[str, str]] = Field(
default=None,
description="""
Optional dictionary mapping field names to sample values for payload generation.
Works for all field types: text fields, numeric fields, categorical fields, etc.
Supports template expansion (e.g., {timestamp} → actual timestamp).
Examples:
- Text fields: {"chat": "Hello, I need help", "shiptrack": "Shipped|In Transit"}
- ID fields: {"order_id": "ORDER_{timestamp}"}
- Numeric fields: {"price": "99.99", "quantity": "5"}
- Any field: Maps directly to payload value
""",
)
# NEW: Custom payload path (S3 or local)
custom_payload_path: Optional[str] = Field(
default=None,
description="""
Optional path to user-provided custom payload sample file (JSON/CSV) or directory.
Supports both S3 paths and local file paths.
When provided, the script will use this instead of auto-generating payloads.
Examples:
- S3: "s3://my-bucket/custom_payload_samples/sample.json"
- Local: "/opt/ml/input/data/custom_payload/sample.json"
- Local: "file:///path/to/payload.json"
""",
)
# Performance thresholds
max_acceptable_error_rate: float = Field(
default=0.2, ge=0.0, le=1.0, description="Maximum acceptable error rate (0-1)"
)
# Load testing configuration for registered model endpoint
load_test_instance_type_list: List[str] = Field(
default=["ml.m5.4xlarge"], # Matches processing_instance_type_large default
description="""
List of instance types for load testing the registered model endpoint.
These instance types will be used for endpoint deployment during load testing.
Default: ["ml.m5.4xlarge"] (matches processing_instance_type_large)
Examples:
- Single instance: ["ml.m5.2xlarge"]
- Multiple instances: ["ml.m5.2xlarge", "ml.m5.4xlarge", "ml.m5.12xlarge"]
Common CPU instance types:
- ml.m5.2xlarge, ml.m5.4xlarge, ml.m5.12xlarge
Common GPU instance types (for deep learning models):
- ml.g4dn.xlarge, ml.g5.xlarge, ml.g5.2xlarge, ml.g5.12xlarge
""",
)
# ===== Derived Fields (Tier 3) =====
# These are fields calculated from other fields, stored in private attributes
# Valid types for validation
_VALID_TYPES: ClassVar[List[str]] = ["NUMERIC", "TEXT"]
# Update to Pydantic V2 style model_config
model_config = ConfigDict(
arbitrary_types_allowed=True,
validate_assignment=False, # Changed from True to False to prevent recursion
extra="allow", # Changed from 'forbid' to 'allow' to accept metadata fields during deserialization
)
# Custom serializer for Path fields (Pydantic V2 approach)
[docs]
@field_serializer("processing_source_dir", "source_dir", when_used="json")
def serialize_path_fields(self, value: Optional[Union[str, Path]]) -> Optional[str]:
"""Serialize Path objects to strings"""
if value is None:
return None
return str(value)
# Removed sample_payload_s3_key property - S3 path construction should happen in builders/scripts
# Removed validators for variable lists since those fields were removed
# Script gets all variable information from hyperparameters.json
# Field validators
[docs]
@field_validator("load_test_instance_type_list")
@classmethod
def validate_load_test_instance_types(cls, v: List[str]) -> List[str]:
"""Validate load test instance type format"""
if not v: # Empty list
raise ValueError("load_test_instance_type_list cannot be empty")
for instance_type in v:
if not instance_type.startswith("ml."):
raise ValueError(
f"Instance type '{instance_type}' must start with 'ml.'"
)
return v
# Model validators
[docs]
@model_validator(mode="after")
def initialize_derived_fields(self) -> "PayloadConfig":
"""Initialize all derived fields once after validation."""
# Call parent validator first
super().initialize_derived_fields()
# No additional derived fields to initialize for PayloadConfig
return self
def _env_overrides(self) -> Dict[str, Optional[str]]:
"""Computed env values for the payload script (FZ 31e1d3g) — these declared env vars are
not plain field passthroughs: CONTENT_TYPES is comma-joined, FIELD_DEFAULTS is JSON. The
base resolver consumes this map for the names the interface declares; ``None`` omits a key
(matches the builder's original conditional-add for an empty field_defaults)."""
import json
overrides: Dict[str, Optional[str]] = {
"CONTENT_TYPES": ",".join(self.source_model_inference_content_types),
"DEFAULT_NUMERIC_VALUE": str(self.default_numeric_value),
"DEFAULT_TEXT_VALUE": str(self.default_text_value),
"FIELD_DEFAULTS": json.dumps(self.field_defaults)
if self.field_defaults
else None,
}
return overrides
[docs]
def get_job_arguments(self) -> Optional[List[str]]:
"""CLI args — config is the single source (FZ 31e1d3h). Custom passthrough or None."""
if getattr(self, "processing_script_arguments", None):
return list(self.processing_script_arguments)
return None
# Removed validate_special_fields validator since it referenced the removed variable list fields
# Special field validation is no longer needed since the script gets variables from hyperparameters
# and special_field_values is optional with proper defaults
# Methods for payload generation and paths
# Removed ensure_payload_path() and get_full_payload_path() methods
# These are redundant and not portable - S3 path construction should happen in builders/scripts
# Removed get_field_default_value() method - this is processing logic that belongs in the script
# Config should only provide the default values, not compute them
# Removed payload generation and processing methods - these belong in the script, not config
# Config should only handle user input and configuration, not actual processing logic
# Removed get_script_path() method - using inherited implementation from base config
# The base config's implementation handles script path resolution properly
# Removed redundant input/output variable helper methods:
# - get_normalized_input_variables() - duplicates script logic
# - get_input_variables_as_dict() - duplicates script logic
# These methods duplicate processing logic that belongs in the script, not config.
# Config should only provide the raw data, not process it.
# Removed model_dump() method - it was redundant, just calling super().model_dump()
# Using inherited implementation from base config