from pydantic import Field, model_validator, field_validator
from typing import List, Optional, Dict, Any
from ...core.base.hyperparameters_base import ModelHyperparameters
from ...core.base.config_base import BasePipelineConfig
[docs]
class PyTorchTrainingConfig(BasePipelineConfig):
"""
Configuration specific to the SageMaker PyTorch Training Step.
This version is streamlined to work with specification-driven architecture.
Input/output paths are now provided via step specifications and dependencies.
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
training_entry_point: str = Field(
description="Entry point script for Pytorch training."
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
# Instance configuration
training_instance_type: str = Field(
default="ml.g5.12xlarge", description="Instance type for training job."
)
training_instance_count: int = Field(
default=1, ge=1, description="Number of instances for training job."
)
training_volume_size: int = Field(
default=30, ge=1, description="Volume size (GB) for training instances."
)
# Framework versions for SageMaker PyTorch container
framework_version: str = Field(
default="2.1.2", description="SageMaker PyTorch framework version."
)
py_version: str = Field(
default="py310",
description="Python version for the SageMaker PyTorch container.",
)
ca_repository_arn: str = Field(
default="arn:aws:codeartifact:us-west-2:149122183214:repository/amazon/secure-pypi",
description="CodeArtifact repository ARN for secure PyPI access. Only used when use_secure_pypi=True.",
)
# Hyperparameters handling configuration
skip_hyperparameters_s3_uri: bool = Field(
default=True,
description="Whether to skip hyperparameters_s3_uri channel during _get_inputs. "
"If True (default), hyperparameters are loaded from script folder. "
"If False, hyperparameters_s3_uri channel is created as TrainingInput.",
)
# Hyperparameters object (optional for backward compatibility)
hyperparameters: Optional[ModelHyperparameters] = Field(
None,
description="Model hyperparameters (optional when using external JSON files)",
)
# Pre-computed artifact flags
use_precomputed_imputation: bool = Field(
default=False,
description="Controls whether to use pre-computed imputation artifacts. "
"If True, expects input data to be already imputed and loads impute_dict.pkl from model_artifacts_input, skipping inline computation. "
"If False (default), computes imputation inline and transforms data.",
)
use_precomputed_risk_tables: bool = Field(
default=False,
description="Controls whether to use pre-computed risk table artifacts. "
"If True, expects input data to be already risk-mapped and loads risk_table_map.pkl from model_artifacts_input, skipping inline computation. "
"If False (default), computes risk tables inline and transforms data.",
)
use_precomputed_features: bool = Field(
default=False,
description="Controls whether to use pre-computed feature selection. "
"If True, expects input data to be already feature-selected and loads selected_features.json from model_artifacts_input, skipping inline computation. "
"If False (default), uses all features without selection.",
)
enable_true_streaming: bool = Field(
default=False,
description="Controls whether to enable streaming mode with PipelineIterableDataset for memory-efficient data loading. "
"If True, uses PipelineIterableDataset which loads data incrementally from sharded files (part-*.parquet), enabling constant memory usage. "
"If False (default), uses PipelineDataset which loads entire dataset into memory. "
"Requires preprocessing to output sharded data (CONSOLIDATE_SHARDS=false). "
"Automatically falls back to batch mode if no shards detected.",
)
# DataLoader worker configuration (only used when enable_true_streaming=True)
# Defaults match batch mode (enable_true_streaming=False default)
num_workers_per_rank: int = Field(
default=0,
ge=0,
le=16,
description=(
"Number of DataLoader workers per GPU rank for parallel data loading. "
"Only used when enable_true_streaming=True. "
"Default: 0 (matches batch mode default). "
"Recommended for streaming mode: 2-8 depending on CPU cores. "
"Set to 0 to disable workers even in streaming mode."
),
)
prefetch_factor: Optional[int] = Field(
default=None,
ge=1,
le=10,
description=(
"Number of batches to prefetch per DataLoader worker. "
"Only used when enable_true_streaming=True and num_workers_per_rank > 0. "
"Default: None (matches batch mode default). "
"Recommended for streaming mode: 2. "
"Higher values use more memory but reduce waiting time."
),
)
use_persistent_workers: bool = Field(
default=False,
description=(
"Whether to keep DataLoader workers alive between epochs. "
"Only used when enable_true_streaming=True and num_workers_per_rank > 0. "
"Default: False (matches batch mode default). "
"Recommended for streaming mode: True (faster epoch transitions). "
"False: Workers restart each epoch (slower, less memory)."
),
)
# Semi-supervised learning support
job_type: Optional[str] = Field(
default=None,
description=(
"Training job type for semi-supervised learning workflows:\n"
"• None (default): Standard supervised learning - no step name suffix\n"
"• 'pretrain': SSL pretraining phase - adds '-Pretrain' suffix\n"
"• 'finetune': SSL fine-tuning phase - adds '-Finetune' suffix"
),
)
model_config = BasePipelineConfig.model_config
[docs]
@field_validator("job_type")
@classmethod
def validate_job_type(cls, v: Optional[str]) -> Optional[str]:
"""Validate job_type is open (any lowercase alphanumeric with underscores; None = standard)."""
if v is None:
return None # Standard supervised learning
if not v.replace("_", "").isalnum() or v != v.lower():
raise ValueError(
f"job_type must be lowercase alphanumeric (with underscores), got '{v}'"
)
return v
[docs]
@model_validator(mode="after")
def validate_dataloader_config(self) -> "PyTorchTrainingConfig":
"""
Validate DataLoader worker configuration (warnings only, no mutation).
Conditional logic is enforced in training script, not here.
This validator only checks for potentially problematic values.
"""
# Warn about very high worker counts
if self.num_workers_per_rank > 8:
print(
f"⚠️ WARNING: num_workers_per_rank={self.num_workers_per_rank} is high. "
f"Recommended: 2-8 for optimal performance."
)
# Warn if streaming enabled but no workers
if self.enable_true_streaming and self.num_workers_per_rank == 0:
print(
"⚠️ WARNING: enable_true_streaming=True but num_workers_per_rank=0. "
"Consider increasing for better parallel I/O performance."
)
# Warn about high prefetch factor (only check if not None)
if self.prefetch_factor is not None and self.prefetch_factor > 4:
print(
f"⚠️ WARNING: prefetch_factor={self.prefetch_factor} is high. "
f"May use excessive memory."
)
return self
@field_validator("training_instance_type")
@classmethod
def _validate_sagemaker_training_instance_type(cls, v: str) -> str:
valid_instances = [
"ml.m5.4xlarge",
"ml.m5.8xlarge",
"ml.m5.12xlarge",
"ml.m5.24xlarge",
"ml.g4dn.16xlarge",
"ml.g5.12xlarge",
"ml.g5.16xlarge",
"ml.g5.24xlarge",
"ml.g5.48xlarge",
"ml.p3.8xlarge",
"ml.p3.16xlarge",
"ml.p4d.24xlarge",
"ml.p4de.24xlarge",
]
if v not in valid_instances:
raise ValueError(
f"Invalid training instance type: {v}. "
f"Must be one of: {', '.join(valid_instances)}"
)
return v
[docs]
def get_environment_variables(self) -> Dict[str, str]:
"""
Get environment variables for the PyTorch training 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 PyTorch training specific environment variables
env_vars.update(
{
"REGION": self.region,
"USE_SECURE_PYPI": str(self.use_secure_pypi).lower(),
"USE_PRECOMPUTED_IMPUTATION": str(
self.use_precomputed_imputation
).lower(),
"USE_PRECOMPUTED_RISK_TABLES": str(
self.use_precomputed_risk_tables
).lower(),
"USE_PRECOMPUTED_FEATURES": str(self.use_precomputed_features).lower(),
"ENABLE_TRUE_STREAMING": str(self.enable_true_streaming).lower(),
# DataLoader worker configuration
"NUM_WORKERS_PER_RANK": str(self.num_workers_per_rank),
"PREFETCH_FACTOR": str(self.prefetch_factor),
"USE_PERSISTENT_WORKERS": str(self.use_persistent_workers).lower(),
}
)
return env_vars
[docs]
def get_public_init_fields(self) -> Dict[str, Any]:
"""
Override get_public_init_fields to include PyTorch training-specific fields.
Gets a dictionary of public fields suitable for initializing a child config.
Includes both base fields (from parent) and PyTorch training-specific fields.
Returns:
Dict[str, Any]: Dictionary of field names to values for child initialization
"""
# Get fields from parent class (BasePipelineConfig)
base_fields = super().get_public_init_fields()
# Add PyTorch training-specific fields (Tier 1 and Tier 2)
training_fields = {
"training_entry_point": self.training_entry_point,
"training_instance_type": self.training_instance_type,
"training_instance_count": self.training_instance_count,
"training_volume_size": self.training_volume_size,
"framework_version": self.framework_version,
"py_version": self.py_version,
"ca_repository_arn": self.ca_repository_arn,
"skip_hyperparameters_s3_uri": self.skip_hyperparameters_s3_uri,
"use_precomputed_imputation": self.use_precomputed_imputation,
"use_precomputed_risk_tables": self.use_precomputed_risk_tables,
"use_precomputed_features": self.use_precomputed_features,
"enable_true_streaming": self.enable_true_streaming,
"num_workers_per_rank": self.num_workers_per_rank,
"prefetch_factor": self.prefetch_factor,
"use_persistent_workers": self.use_persistent_workers,
"job_type": self.job_type,
}
# Add hyperparameters if present (use model_dump for Pydantic models)
if self.hyperparameters is not None:
training_fields["hyperparameters"] = self.hyperparameters.model_dump()
# Combine base fields and training fields (training fields take precedence if overlap)
init_fields = {**base_fields, **training_fields}
return init_fields
[docs]
def get_job_arguments(self) -> Optional[List[str]]:
"""CLI args — config is the single source (FZ 31e1d3h)."""
return self._job_type_arg()