"""
Bedrock Processing Step Configuration
This module implements the configuration class for the Bedrock Processing step
using the three-tier design pattern for optimal user experience and maintainability.
"""
from pydantic import Field, PrivateAttr, model_validator, field_validator
from typing import Dict, Any, Optional, List
import json
import logging
from .config_processing_step_base import ProcessingStepConfigBase
logger = logging.getLogger(__name__)
[docs]
class BedrockProcessingConfig(ProcessingStepConfigBase):
"""
Configuration for Bedrock Processing step using three-tier design.
This step processes input data through AWS Bedrock models using generated prompt
templates and validation schemas from the Bedrock Prompt Template Generation step.
Supports template-driven response processing with dynamic Pydantic model creation
and both sequential and concurrent processing modes.
Tier 1: Essential user inputs (required)
Tier 2: System inputs with defaults (optional)
Tier 3: Derived fields (private with property access)
"""
# ===== Tier 1: Essential User Inputs (Required) =====
# These fields must be provided by users with no defaults
bedrock_inference_profile_arn: str = Field(
description="Inference profile ARN for capacity management (e.g., 'arn:aws:bedrock:us-east-1:123456789012:inference-profile/abc123')"
)
# ===== Tier 2: System Inputs with Defaults (Optional) =====
# These fields have sensible defaults but can be overridden
job_type: str = Field(
default="training",
description="One of ['training','validation','testing','calibration'] - determines processing behavior and output naming",
)
# Model configuration
bedrock_primary_model_id: str = Field(
default="anthropic.claude-sonnet-4-5-20250929-v1:0",
description="Primary Bedrock model ID for processing (Claude Sonnet 4.5 default, latest stable model)",
)
bedrock_fallback_model_id: Optional[str] = Field(
default="anthropic.claude-sonnet-4-20250514-v1:0",
description="Fallback model ID for inference profile failures (Claude Sonnet 4.0 for production reliability)",
)
bedrock_inference_profile_required_models: List[str] = Field(
default_factory=lambda: [
"anthropic.claude-sonnet-4-5-20250929-v1:0", # Claude Sonnet 4.5
"anthropic.claude-sonnet-4-20250514-v1:0", # Claude Sonnet 4.0
"anthropic.claude-opus-4-1-20250805-v1:0", # Claude Opus 4.1
],
description="List of models requiring inference profiles (Claude 4.5, 4.0, and Opus 4.1 included by default)",
)
# API parameters (optimized for Claude 4.0)
bedrock_max_tokens: int = Field(
default=32000,
ge=1,
le=64000,
description="Maximum tokens for Bedrock responses (32K optimal default, 64K max for Claude 4.0)",
)
bedrock_temperature: float = Field(
default=1.0,
ge=0.0,
le=2.0,
description="Temperature for response generation (1.0 optimized for Claude 4.0)",
)
bedrock_top_p: float = Field(
default=0.999,
ge=0.0,
le=1.0,
description="Top-p sampling parameter (0.999 optimized for Claude 4.0)",
)
# Processing configuration
bedrock_batch_size: int = Field(
default=10, ge=1, le=100, description="Number of records per processing batch"
)
bedrock_max_retries: int = Field(
default=3,
ge=0,
le=10,
description="Maximum retries for failed Bedrock requests",
)
bedrock_output_column_prefix: str = Field(
default="llm_", description="Prefix for output columns in processed data"
)
bedrock_skip_error_records: bool = Field(
default=False,
description="Whether to exclude error records from output files (statistics still track all records)",
)
# Concurrency configuration
bedrock_concurrency_mode: str = Field(
default="sequential",
description="Processing mode: 'sequential' (safer, easier debugging) or 'concurrent' (faster, 3-10x speedup)",
)
bedrock_max_concurrent_workers: int = Field(
default=5,
ge=1,
le=100,
description="Number of concurrent threads for concurrent processing (recommended: 3-50)",
)
bedrock_rate_limit_per_second: int = Field(
default=10,
ge=1,
le=100,
description="API requests per second limit for concurrent processing",
)
# Input truncation configuration
bedrock_max_input_field_length: int = Field(
default=400000,
ge=100,
le=1000000,
description="Maximum length in characters for input fields before truncation (default: 400,000 chars ≈ 100,000 tokens)",
)
bedrock_truncation_enabled: bool = Field(
default=True,
description="Enable automatic truncation of oversized input fields to prevent error 413 'Input is too long'",
)
bedrock_log_truncations: bool = Field(
default=True,
description="Log detailed information about truncated fields for debugging and monitoring",
)
# Config-embedded template support (self-contained mode)
bedrock_user_prompt_template: Optional[str] = Field(
default=None,
description="User prompt template with {placeholder} syntax. If provided, BedrockPromptTemplateGeneration step is not needed.",
)
bedrock_system_prompt: Optional[str] = Field(
default=None,
description="System prompt for Bedrock API. Used only in config-embedded mode.",
)
bedrock_input_placeholders: Optional[List[str]] = Field(
default=None,
description="List of input placeholder names mapping to DataFrame columns.",
)
bedrock_validation_schema: Optional[Dict[str, Any]] = Field(
default=None,
description="JSON validation schema for Pydantic response model creation.",
)
# Structured output mode
bedrock_use_structured_output: bool = Field(
default=False,
description="Use tool_use for guaranteed schema compliance (0% parse failures).",
)
# Converse API mode
bedrock_use_converse_api: bool = Field(
default=False,
description="Use Converse API (model-agnostic) instead of invoke_model. Enables Nova/Llama/Mistral.",
)
# Adaptive rate limiting
bedrock_adaptive_rate_limiting: bool = Field(
default=False,
description="Auto-tune rate limit based on observed throttle rate.",
)
# Processing step overrides
processing_entry_point: str = Field(
default="bedrock_processing.py",
description="Entry point script for Bedrock processing",
)
# PyTorch framework configuration
framework_version: str = Field(
default="2.1.2",
description="PyTorch framework version for processing container",
)
py_version: str = Field(
default="py310",
description="Python version for PyTorch container (e.g., 'py310', 'py39')",
)
# ===== Tier 3: Derived Fields (Private with Property Access) =====
# These fields are calculated from other fields
_effective_inference_profile_required_models: Optional[List[str]] = PrivateAttr(
default=None
)
_bedrock_environment_variables: Optional[Dict[str, str]] = PrivateAttr(default=None)
_processing_metadata: Optional[Dict[str, Any]] = PrivateAttr(default=None)
_concurrency_configuration: Optional[Dict[str, Any]] = PrivateAttr(default=None)
# Public properties for derived fields
@property
def effective_inference_profile_required_models(self) -> List[str]:
"""Get effective list of models requiring inference profiles with auto-detection."""
if self._effective_inference_profile_required_models is None:
# Start with user-provided list
models = list(self.bedrock_inference_profile_required_models)
# Auto-detect known models that require inference profiles
known_profile_models = [
"anthropic.claude-sonnet-4-5-20250929-v1:0", # Claude Sonnet 4.5
"anthropic.claude-haiku-4-5-20251001-v1:0", # Claude Haiku 4.5
"anthropic.claude-sonnet-4-20250514-v1:0", # Claude Sonnet 4.0
"anthropic.claude-opus-4-1-20250805-v1:0", # Claude Opus 4.1
# Add other known models that require inference profiles
]
# Add primary model if it's known to require profiles and not already in list
if (
self.bedrock_primary_model_id in known_profile_models
and self.bedrock_primary_model_id not in models
):
models.append(self.bedrock_primary_model_id)
self._effective_inference_profile_required_models = models
return self._effective_inference_profile_required_models
@property
def bedrock_environment_variables(self) -> Dict[str, str]:
"""Get environment variables for the Bedrock processing step."""
if self._bedrock_environment_variables is None:
self._bedrock_environment_variables = {
# Model configuration (required)
"BEDROCK_PRIMARY_MODEL_ID": self.bedrock_primary_model_id,
# Model configuration (optional)
"BEDROCK_FALLBACK_MODEL_ID": self.bedrock_fallback_model_id or "",
"BEDROCK_INFERENCE_PROFILE_ARN": self.bedrock_inference_profile_arn,
"BEDROCK_INFERENCE_PROFILE_REQUIRED_MODELS": json.dumps(
self.effective_inference_profile_required_models
),
# AWS configuration
"AWS_DEFAULT_REGION": self.aws_region,
# API parameters
"BEDROCK_MAX_TOKENS": str(self.bedrock_max_tokens),
"BEDROCK_TEMPERATURE": str(self.bedrock_temperature),
"BEDROCK_TOP_P": str(self.bedrock_top_p),
# Processing configuration
"BEDROCK_BATCH_SIZE": str(self.bedrock_batch_size),
"BEDROCK_MAX_RETRIES": str(self.bedrock_max_retries),
"BEDROCK_OUTPUT_COLUMN_PREFIX": self.bedrock_output_column_prefix,
"BEDROCK_SKIP_ERROR_RECORDS": str(
self.bedrock_skip_error_records
).lower(),
# Concurrency configuration
"BEDROCK_CONCURRENCY_MODE": self.bedrock_concurrency_mode,
"BEDROCK_MAX_CONCURRENT_WORKERS": str(
self.bedrock_max_concurrent_workers
),
"BEDROCK_RATE_LIMIT_PER_SECOND": str(
self.bedrock_rate_limit_per_second
),
# Input truncation configuration
"BEDROCK_MAX_INPUT_FIELD_LENGTH": str(
self.bedrock_max_input_field_length
),
"BEDROCK_TRUNCATION_ENABLED": str(
self.bedrock_truncation_enabled
).lower(),
"BEDROCK_LOG_TRUNCATIONS": str(self.bedrock_log_truncations).lower(),
"USE_SECURE_PYPI": str(self.use_secure_pypi).lower(),
# Config-embedded template support (self-contained mode)
"BEDROCK_USER_PROMPT_TEMPLATE": self.bedrock_user_prompt_template or "",
"BEDROCK_SYSTEM_PROMPT": self.bedrock_system_prompt or "",
"BEDROCK_INPUT_PLACEHOLDERS": json.dumps(
self.bedrock_input_placeholders
)
if self.bedrock_input_placeholders
else "[]",
"BEDROCK_VALIDATION_SCHEMA": json.dumps(self.bedrock_validation_schema)
if self.bedrock_validation_schema
else "{}",
# Structured output mode
"BEDROCK_USE_STRUCTURED_OUTPUT": str(
self.bedrock_use_structured_output
).lower(),
# Converse API mode
"BEDROCK_USE_CONVERSE_API": str(self.bedrock_use_converse_api).lower(),
# Adaptive rate limiting
"BEDROCK_ADAPTIVE_RATE_LIMITING": str(
self.bedrock_adaptive_rate_limiting
).lower(),
}
return self._bedrock_environment_variables
@property
def processing_metadata(self) -> Dict[str, Any]:
"""Get processing step metadata."""
if self._processing_metadata is None:
self._processing_metadata = {
"step_type": "bedrock_processing",
"primary_model": self.bedrock_primary_model_id,
"fallback_model": self.bedrock_fallback_model_id,
"concurrency_mode": self.bedrock_concurrency_mode,
"batch_size": self.bedrock_batch_size,
"max_tokens": self.bedrock_max_tokens,
"temperature": self.bedrock_temperature,
"top_p": self.bedrock_top_p,
"uses_inference_profile": bool(self.bedrock_inference_profile_arn),
"inference_profile_required_models": self.effective_inference_profile_required_models,
"output_column_prefix": self.bedrock_output_column_prefix,
}
return self._processing_metadata
@property
def concurrency_configuration(self) -> Dict[str, Any]:
"""Get concurrency configuration details."""
if self._concurrency_configuration is None:
self._concurrency_configuration = {
"mode": self.bedrock_concurrency_mode,
"max_workers": self.bedrock_max_concurrent_workers,
"rate_limit_per_second": self.bedrock_rate_limit_per_second,
"is_concurrent": self.bedrock_concurrency_mode == "concurrent",
"expected_speedup": f"{self.bedrock_max_concurrent_workers}x"
if self.bedrock_concurrency_mode == "concurrent"
else "1x",
"recommended_for_production": self.bedrock_concurrency_mode
== "concurrent"
and self.bedrock_fallback_model_id is not None,
}
return self._concurrency_configuration
# Validators
[docs]
@field_validator("job_type")
@classmethod
def validate_job_type(cls, v: str) -> str:
"""Validate job_type is one of the allowed values."""
if not v.replace("_", "").isalnum() or v != v.lower():
raise ValueError(
"job_type must be lowercase alphanumeric (with underscores), got '{v}'"
)
return v
[docs]
@field_validator("bedrock_primary_model_id")
@classmethod
def validate_primary_model_id(cls, v: str) -> str:
"""Validate primary model ID format."""
if not v or not v.strip():
raise ValueError("bedrock_primary_model_id cannot be empty")
# Basic format validation for common Bedrock model patterns
valid_prefixes = [
"anthropic.",
"amazon.",
"ai21.",
"cohere.",
"meta.",
"mistral.",
"stability.",
"global.", # For inference profile IDs
]
if not any(v.startswith(prefix) for prefix in valid_prefixes):
logger.warning(
f"Model ID '{v}' doesn't match common Bedrock patterns. Ensure it's a valid Bedrock model ID."
)
return v.strip()
[docs]
@field_validator("bedrock_concurrency_mode")
@classmethod
def validate_concurrency_mode(cls, v: str) -> str:
"""Validate concurrency mode."""
valid_modes = ["sequential", "concurrent"]
if v not in valid_modes:
raise ValueError(
f"bedrock_concurrency_mode must be one of {valid_modes}, got: {v}"
)
return v
[docs]
@field_validator("bedrock_inference_profile_required_models")
@classmethod
def validate_inference_profile_models(cls, v: List[str]) -> List[str]:
"""Validate inference profile required models list."""
if v is None:
return []
# Remove empty strings and duplicates
cleaned = list(set(model.strip() for model in v if model and model.strip()))
return cleaned
# Custom model_dump method to include derived properties
[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["effective_inference_profile_required_models"] = (
self.effective_inference_profile_required_models
)
data["bedrock_environment_variables"] = self.bedrock_environment_variables
data["processing_metadata"] = self.processing_metadata
data["concurrency_configuration"] = self.concurrency_configuration
return data
# Initialize derived fields at creation time
[docs]
@model_validator(mode="after")
def initialize_derived_fields(self) -> "BedrockProcessingConfig":
"""Initialize all derived fields once after validation."""
# Call parent validator first
super().initialize_derived_fields()
# Initialize Bedrock-specific derived fields
_ = self.effective_inference_profile_required_models
_ = self.bedrock_environment_variables
_ = self.processing_metadata
_ = self.concurrency_configuration
return self
[docs]
@model_validator(mode="after")
def validate_production_readiness(self) -> "BedrockProcessingConfig":
"""Validate configuration for production readiness."""
# Warn if using concurrent mode without fallback model
if (
self.bedrock_concurrency_mode == "concurrent"
and not self.bedrock_fallback_model_id
):
logger.warning(
"Using concurrent processing without fallback model. "
"Consider setting bedrock_fallback_model_id for production reliability."
)
# Warn if using inference profile without fallback
if self.bedrock_inference_profile_arn and not self.bedrock_fallback_model_id:
logger.warning(
"Using inference profile without fallback model. "
"Consider setting bedrock_fallback_model_id for production reliability."
)
# Validate concurrent processing parameters
if self.bedrock_concurrency_mode == "concurrent":
if self.bedrock_max_concurrent_workers > 10:
logger.warning(
f"High concurrent worker count ({self.bedrock_max_concurrent_workers}). "
"Consider reducing to 3-10 workers to avoid rate limiting."
)
if self.bedrock_rate_limit_per_second > 50:
logger.warning(
f"High rate limit ({self.bedrock_rate_limit_per_second} req/sec). "
"Ensure this doesn't exceed your Bedrock API limits."
)
return self
[docs]
def get_script_path(self, default_path: Optional[str] = None) -> Optional[str]:
"""
Get script path for the Bedrock processing step.
Args:
default_path: Default script path to use if not found via other methods
Returns:
Script path resolved from processing_entry_point and source directories
"""
# Use the parent class implementation which handles hybrid resolution
return super().get_script_path(default_path)
[docs]
def get_public_init_fields(self) -> Dict[str, Any]:
"""
Override get_public_init_fields to include Bedrock-specific fields.
Gets a dictionary of public fields suitable for initializing a child config.
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 Bedrock-specific fields (Tier 1 + Tier 2)
bedrock_fields = {
# Tier 1: Essential fields
"job_type": self.job_type,
"bedrock_inference_profile_arn": self.bedrock_inference_profile_arn,
# Tier 2: System fields with defaults
"bedrock_primary_model_id": self.bedrock_primary_model_id,
"bedrock_max_tokens": self.bedrock_max_tokens,
"bedrock_temperature": self.bedrock_temperature,
"bedrock_top_p": self.bedrock_top_p,
"bedrock_batch_size": self.bedrock_batch_size,
"bedrock_max_retries": self.bedrock_max_retries,
"bedrock_output_column_prefix": self.bedrock_output_column_prefix,
"bedrock_concurrency_mode": self.bedrock_concurrency_mode,
"bedrock_max_concurrent_workers": self.bedrock_max_concurrent_workers,
"bedrock_rate_limit_per_second": self.bedrock_rate_limit_per_second,
"bedrock_inference_profile_required_models": self.bedrock_inference_profile_required_models,
}
# Only include optional fields if they're set
if self.bedrock_fallback_model_id is not None:
bedrock_fields["bedrock_fallback_model_id"] = self.bedrock_fallback_model_id
if self.bedrock_user_prompt_template is not None:
bedrock_fields["bedrock_user_prompt_template"] = (
self.bedrock_user_prompt_template
)
if self.bedrock_system_prompt is not None:
bedrock_fields["bedrock_system_prompt"] = self.bedrock_system_prompt
if self.bedrock_input_placeholders is not None:
bedrock_fields["bedrock_input_placeholders"] = (
self.bedrock_input_placeholders
)
if self.bedrock_validation_schema is not None:
bedrock_fields["bedrock_validation_schema"] = self.bedrock_validation_schema
bedrock_fields["bedrock_use_structured_output"] = (
self.bedrock_use_structured_output
)
bedrock_fields["bedrock_use_converse_api"] = self.bedrock_use_converse_api
bedrock_fields["bedrock_adaptive_rate_limiting"] = (
self.bedrock_adaptive_rate_limiting
)
# Combine base fields and Bedrock fields (Bedrock fields take precedence if overlap)
init_fields = {**base_fields, **bedrock_fields}
return init_fields
[docs]
def get_environment_variables(self) -> Dict[str, str]:
"""
Get all environment variables for the step builder.
Returns:
Dict[str, str]: Complete environment variables dictionary
"""
return self.bedrock_environment_variables
[docs]
def is_production_ready(self) -> bool:
"""
Check if configuration is production-ready.
Returns:
bool: True if configuration has production-ready settings
"""
return (
# Has fallback model for reliability
self.bedrock_fallback_model_id is not None
and
# Uses reasonable concurrency settings
(
self.bedrock_concurrency_mode == "sequential"
or (
self.bedrock_max_concurrent_workers <= 10
and self.bedrock_rate_limit_per_second <= 50
)
)
)
[docs]
def get_job_arguments(self) -> Optional[List[str]]:
"""CLI args — config is the single source (FZ 31e1d3h). job_type + batch/retry overrides."""
args = ["--job_type", self.job_type]
args.extend(["--batch-size", str(self.bedrock_batch_size)])
args.extend(["--max-retries", str(self.bedrock_max_retries)])
return args