"""
TSA Model Evaluation Step Configuration with Self-Contained Derivation Logic
This module implements the configuration class for the TSA model evaluation step
using a self-contained design where derived fields are private with read-only properties.
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private with properties)
Aligned with PyTorch model eval config structure.
"""
from pydantic import Field, model_validator, field_validator, PrivateAttr
from typing import Optional, Dict, Any, TYPE_CHECKING
from pathlib import Path
import logging
from .config_processing_step_base import ProcessingStepConfigBase
# Import for type hints only
if TYPE_CHECKING:
from ...core.base.contract_base import ScriptContract
logger = logging.getLogger(__name__)
[docs]
class TSAModelEvalConfig(ProcessingStepConfigBase):
"""
Configuration for TSA model evaluation step with self-contained derivation logic.
This class defines the configuration parameters for the TSA (Temporal Self-Attention)
model evaluation step, which calculates evaluation metrics for trained PyTorch models
with dual-task learning support. Computes comprehensive metrics including AUC-ROC,
precision-recall, dollar-weighted metrics, and generates visualizations.
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
id_name: str = Field(
default="objectId",
description="Name of the ID field in the dataset (required for evaluation).",
)
label_name: str = Field(
default="is_abusive_mdr",
description="Name of the Task 1 label field, used in riskband/percentile metrics (e.g., 'is_abusive_mdr').",
)
task2_label_name: str = Field(
default="is_abusive_flr",
description="Display name for Task 2 label in riskband/percentile metrics (e.g., 'is_abusive_flr').",
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
processing_entry_point: str = Field(
default="tsa_model_eval.py",
description="Entry point script for TSA model evaluation.",
)
job_type: str = Field(
default="evaluation",
description="Type of evaluation job to perform (e.g., 'evaluation', 'validation', 'testing').",
)
data_version: str = Field(
default="v0",
description="Version suffix for numpy array files (e.g., X_num_v0.npy). Used to locate input data files.",
)
# PyTorch specific fields
framework_version: str = Field(
default="2.1.2", description="PyTorch framework version for processing"
)
py_version: str = Field(
default="py310",
description="Python version for the SageMaker PyTorch container.",
)
# Note: use_secure_pypi is inherited from BasePipelineConfig
# Instance Configuration - Supports both GPU and CPU
# GPU recommended for optimal performance (2-10x faster), CPU supported for flexibility
processing_instance_type_large: str = Field(
default="ml.p3.16xlarge",
description="Large instance type for TSA model evaluation. "
"GPU (recommended): ml.p3.16xlarge (8x V100, 128GB total), ml.p3.2xlarge (1x V100, 16GB). "
"CPU (alternative): ml.m5.4xlarge (16 vCPU, cost-effective). "
"GPU provides 2-10x speedup with AMP and optimized PyTorch operations.",
)
processing_instance_type_small: str = Field(
default="ml.g5.16xlarge",
description="Small instance type for TSA model evaluation. "
"GPU (recommended): ml.g5.16xlarge (1x A10G, 24GB), ml.g4dn.xlarge (1x T4, 16GB). "
"CPU (alternative): ml.m5.2xlarge (8 vCPU), ml.m5.xlarge (4 vCPU). "
"GPU provides significant performance advantage for PyTorch inference.",
)
use_large_processing_instance: bool = Field(
default=True,
description="Whether to use large GPU instance type (ml.p3.2xlarge) for processing. "
"TSA evaluation benefits from V100 GPU for 4-6x faster evaluation with AMP.",
)
# Performance optimizations
enable_eval_streaming: bool = Field(
default=False,
description="Enable two-pass streaming evaluation mode for better performance. "
"When True: 30-40% faster evaluation with 50% lower memory usage. "
"Maintains 100% metric accuracy. Recommended for large datasets (>1M samples).",
)
enable_amp: bool = Field(
default=True,
description="Enable mixed precision (AMP) for faster GPU inference. "
"When True: 2-3x faster evaluation on GPU with no accuracy loss. "
"Automatically enabled on CUDA devices. Set to False to disable.",
)
num_workers: int = Field(
default=4,
ge=0,
le=32,
description="Number of parallel data loading workers (0-32). "
"4 workers recommended for production (30-50% faster I/O). "
"Set to 0 for single-threaded loading (debugging).",
)
enable_cpu_optimization: bool = Field(
default=True,
description="Enable CPU-specific optimizations (threading, JIT compilation, batch size tuning). "
"When True: 2-5x faster evaluation on CPU with optimized threading (Intel MKL), "
"TorchScript JIT compilation, and cache-friendly batch sizes. "
"Automatically detects CPU and applies optimizations. "
"Set to False to disable for baseline comparison or troubleshooting.",
)
eval_percentile: float = Field(
default=0.99,
gt=0.0,
le=1.0,
description="Quantile threshold for computing recall and precision metrics during evaluation. "
"Used by tsa_model_eval.py to compute recall/precision at this percentile of predicted scores.",
)
eval_batch_size: Optional[int] = Field(
default=None,
ge=1,
le=2048,
description="Batch size for model evaluation. "
"If None, auto-detects optimal size based on instance type, or falls back to "
"batch_size from hyperparameters.json (not recommended - designed for training). "
"Recommended values by instance type: "
"- ml.p3.16xlarge (8x V100): 512 per GPU, effective 4096 with DDP "
"- ml.p3.8xlarge (4x V100): 512 per GPU, effective 2048 with DDP "
"- ml.p3.2xlarge (1x V100): 256-512 "
"- ml.g5.16xlarge (1x A10G): 256-512 "
"- ml.g4dn.xlarge (1x T4): 128-256 "
"- CPU instances (ml.m5.*): 64-128 "
"Larger batches dramatically reduce overhead: batch_size=512 vs 2 = 256x fewer iterations, "
"resulting in 18-36x faster evaluation for large datasets (1M+ samples).",
)
model_config = ProcessingStepConfigBase.model_config
# ===== Derived Fields (Tier 3) =====
# These are fields calculated from other fields, stored in private attributes
# with public read-only properties for access
# Currently no derived fields specific to TSA model evaluation
# beyond what's inherited from the ProcessingStepConfigBase class
# Combine all model validators into one to guarantee execution order
[docs]
@model_validator(mode="after")
def validate_and_initialize(self) -> "TSAModelEvalConfig":
"""
Single unified validator that handles initialization, defaults, and validation.
Combining all validators ensures proper execution order:
1. Initialize derived fields (parent)
2. Set eval_batch_size defaults
3. Validate TSA-specific requirements
"""
# Step 1: Call parent initialization
super().initialize_derived_fields()
# Step 2: Set eval_batch_size defaults only if not explicitly provided by user
# Check model_fields_set to distinguish between user-provided values and default None
if (
"eval_batch_size" not in self.model_fields_set
and self.eval_batch_size is None
):
# User didn't provide eval_batch_size, so auto-set based on instance type
instance = (
self.processing_instance_type_large
if self.use_large_processing_instance
else self.processing_instance_type_small
)
# Smart defaults based on instance GPU capacity and memory
if "p3.16xlarge" in instance or "p4d.24xlarge" in instance:
# 8 GPUs with large memory (V100 16GB or A100 40GB each)
self.eval_batch_size = 512
logger.info(
f"Auto-set eval_batch_size=512 for {instance} "
f"(8 GPUs, effective batch 4096 with DDP)"
)
elif "p3.8xlarge" in instance:
# 4 GPUs with V100 16GB each
self.eval_batch_size = 512
logger.info(
f"Auto-set eval_batch_size=512 for {instance} "
f"(4 GPUs, effective batch 2048 with DDP)"
)
elif (
"p3.2xlarge" in instance
or "g5.16xlarge" in instance
or "g5.12xlarge" in instance
):
# 1 GPU with 16-24GB memory (V100 or A10G)
self.eval_batch_size = 256
logger.info(
f"Auto-set eval_batch_size=256 for {instance} (1 large GPU)"
)
elif "g5." in instance or "g4dn." in instance:
# Smaller GPUs (A10G, T4)
self.eval_batch_size = 128
logger.info(
f"Auto-set eval_batch_size=128 for {instance} (1 medium GPU)"
)
else:
# CPU instances or unknown - use conservative batch size
self.eval_batch_size = 64
logger.info(
f"Auto-set eval_batch_size=64 for {instance} "
f"(CPU or unknown instance - cache-friendly size)"
)
elif "eval_batch_size" in self.model_fields_set:
# User explicitly provided eval_batch_size - respect their value
logger.info(
f"Using user-provided eval_batch_size={self.eval_batch_size} "
f"(overriding instance-based defaults)"
)
# Step 3: Validate TSA-specific requirements
if not self.processing_entry_point:
raise ValueError("TSA evaluation step requires a processing_entry_point")
if not self.id_name:
raise ValueError(
"id_name must be provided (required by TSA model evaluation contract)"
)
if not self.label_name:
raise ValueError(
"label_name must be provided (required by TSA model evaluation contract)"
)
if not self.data_version or not self.data_version.strip():
raise ValueError("data_version must be a non-empty string")
logger.debug(
f"TSA evaluation config validated: "
f"job_type='{self.job_type}', data_version='{self.data_version}', "
f"id_name='{self.id_name}', label_name='{self.label_name}', "
f"task2_label='{self.task2_label_name}', eval_batch_size={self.eval_batch_size}"
)
return self
[docs]
@field_validator("processing_instance_type_large", "processing_instance_type_small")
@classmethod
def validate_instance_type_flexible(cls, v: str) -> str:
"""Validate instance type supports PyTorch (GPU or CPU)."""
# GPU instances for optimal performance (recommended)
valid_gpu_instances = [
"ml.p3.2xlarge",
"ml.p3.8xlarge",
"ml.p3.16xlarge",
"ml.p3dn.24xlarge",
"ml.g4dn.xlarge",
"ml.g4dn.2xlarge",
"ml.g4dn.4xlarge",
"ml.g4dn.8xlarge",
"ml.g4dn.12xlarge",
"ml.g4dn.16xlarge",
"ml.g5.xlarge",
"ml.g5.2xlarge",
"ml.g5.4xlarge",
"ml.g5.8xlarge",
"ml.g5.12xlarge",
"ml.g5.16xlarge",
"ml.g5.24xlarge",
"ml.g5.48xlarge",
"ml.p4d.24xlarge",
"ml.p5.48xlarge",
]
# CPU instances for flexibility/cost savings (alternative)
valid_cpu_instances = [
# M5 family - General purpose compute
"ml.m5.large",
"ml.m5.xlarge",
"ml.m5.2xlarge",
"ml.m5.4xlarge",
"ml.m5.8xlarge",
"ml.m5.12xlarge",
"ml.m5.16xlarge",
"ml.m5.24xlarge",
# C5 family - Compute optimized
"ml.c5.large",
"ml.c5.xlarge",
"ml.c5.2xlarge",
"ml.c5.4xlarge",
"ml.c5.9xlarge",
"ml.c5.18xlarge",
# R5 family - Memory optimized
"ml.r5.large",
"ml.r5.xlarge",
"ml.r5.2xlarge",
"ml.r5.4xlarge",
"ml.r5.8xlarge",
"ml.r5.12xlarge",
"ml.r5.16xlarge",
"ml.r5.24xlarge",
]
if v in valid_gpu_instances:
return v
elif v in valid_cpu_instances:
logger.warning(
f"Using CPU instance '{v}' for TSA evaluation. "
f"GPU instances are recommended for 2-10x faster evaluation. "
f"Consider ml.p3.2xlarge (production) or ml.g4dn.xlarge (dev/test) for better performance."
)
return v
else:
raise ValueError(
f"Invalid instance type for TSA evaluation: {v}. "
f"Must be a valid GPU or CPU instance. "
f"GPU (recommended): ml.p3.*, ml.g4dn.*, ml.g5.* "
f"CPU (alternative): ml.m5.*, ml.c5.*, ml.r5.*"
)
[docs]
@field_validator("job_type")
@classmethod
def validate_job_type(cls, v: str) -> str:
"""Validate job_type is a valid value."""
valid_job_types = {
"evaluation",
"training",
"calibration",
"validation",
"testing",
}
if v not in valid_job_types:
logger.warning(
f"job_type '{v}' not in standard set {valid_job_types}. "
f"Proceeding anyway to allow custom job types."
)
return v
[docs]
def get_environment_variables(self) -> Dict[str, str]:
"""
Get environment variables for the TSA model evaluation script.
Returns:
Dict[str, str]: Dictionary mapping environment variable names to values
"""
# Get base environment variables from parent class if available
env_vars = (
super().get_environment_variables()
if hasattr(super(), "get_environment_variables")
else {}
)
# Add TSA evaluation specific environment variables
# All align with the optional env_vars in the TSAModelEval .step.yaml contract
env_vars.update(
{
"DATA_VERSION": self.data_version,
"ID_FIELD": self.id_name, # Map id_name to ID_FIELD
"TASK1_LABEL_NAME": self.label_name, # Label name for Task 1 riskband metrics
"TASK2_LABEL_NAME": self.task2_label_name, # Label name for Task 2 riskband metrics
"USE_SECURE_PYPI": str(self.use_secure_pypi).lower(),
"LOCAL_RANK": "-1", # Default to single GPU/CPU mode, SageMaker overrides for distributed
"ENABLE_EVAL_STREAMING": str(self.enable_eval_streaming).lower(),
"ENABLE_AMP": str(self.enable_amp).lower(),
"NUM_WORKERS": str(self.num_workers),
"ENABLE_CPU_OPTIMIZATION": str(self.enable_cpu_optimization).lower(),
"EVAL_PERCENTILE": str(self.eval_percentile),
"EVAL_BATCH_SIZE": str(self.eval_batch_size)
if self.eval_batch_size
else "",
}
)
logger.debug(f"TSA evaluation environment variables: {env_vars}")
return env_vars
# get_script_contract() / get_script_path() are inherited from BasePipelineConfig, which loads
# the contract from the unified .step.yaml interface via the step catalog (Design B standard —
# matches ModelCalibration / XGBoostModelEval / TabularPreprocessing, which define no override).
[docs]
def get_public_init_fields(self) -> Dict[str, Any]:
"""
Override get_public_init_fields to include TSA evaluation-specific fields.
Gets a dictionary of public fields suitable for initializing a child config.
Includes both base fields (from parent) and evaluation-specific fields.
Returns:
Dict[str, Any]: Dictionary of field names to values for child initialization
"""
# Get fields from parent class (ProcessingStepConfigBase)
base_fields = super().get_public_init_fields()
# Add TSA model evaluation specific fields
eval_fields = {
# Tier 1 - Essential User Inputs
"id_name": self.id_name,
"label_name": self.label_name,
"task2_label_name": self.task2_label_name,
# Tier 2 - System Inputs with Defaults
"processing_entry_point": self.processing_entry_point,
"job_type": self.job_type,
"data_version": self.data_version,
"framework_version": self.framework_version,
"py_version": self.py_version,
"use_large_processing_instance": self.use_large_processing_instance,
# Performance optimizations
"enable_eval_streaming": self.enable_eval_streaming,
"enable_amp": self.enable_amp,
"num_workers": self.num_workers,
"enable_cpu_optimization": self.enable_cpu_optimization,
"eval_percentile": self.eval_percentile,
# Note: use_secure_pypi is inherited from base_fields, no need to add here
}
# Combine base fields and evaluation fields (evaluation fields take precedence if overlap)
init_fields = {**base_fields, **eval_fields}
return init_fields