Advanced Training with ModelTrainer

Introduction & Lesson Overview

Welcome back! In the previous lessons, you learned how to train and evaluate machine learning models in Amazon SageMaker using the classic Estimator pattern. You now know how to upload data to S3, launch a training job, retrieve the trained model, and evaluate its performance. These are essential skills for any machine learning workflow in the cloud. As you continue your journey, it is important to know that SageMaker is always evolving. For more advanced and production-ready workflows, SageMaker now offers a new, modular approach to training called ModelTrainer . This lesson will introduce you to ModelTrainer and show you how it builds on what you have already learned. By the end of this lesson, you will be able to set up and launch a training job using ModelTrainer, retrieve information about your training jobs, and understand the key differences between this modern approach and the classic Estimator pattern. This will prepare you for more sophisticated model development and help you take full advantage of SageMaker's advanced features.

Understanding ModelTrainer vs Estimators

Before diving into the implementation, it's important to understand what ModelTrainer is and how it differs from the Estimator pattern you've been using. Estimators remain a valuable and practical choice for many machine learning projects. They provide a simple, all-in-one approach where you pass all configuration parameters directly to the constructor - framework version, instance type, entry point, output path. The Estimator automatically retrieves the appropriate Docker container image based on your framework specifications, making it perfect for straightforward training workflows. ModelTrainer is designed for when you need more control and flexibility. It uses a modular approach with separate configuration objects: SourceCode for your training script Compute for compute resources OutputDataConfig for model output location InputData for training data A key difference is container handling. With Estimators, you specify framework_version='1.2-1' and it automatically finds the right container. With ModelTrainer, you explicitly retrieve the container image URI first using sagemaker.image_uris.retrieve(), giving you precise control over which container is used. Both approaches have their place. Estimators are excellent for getting started and for straightforward training jobs. ModelTrainer shines when you need better organization for complex projects, want to integrate with SageMaker's advanced features, or require fine-grained control over your training environment.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Python

import sagemaker

# Create a SageMaker session
sagemaker_session = sagemaker.Session()
# Get the default S3 bucket
default_bucket = sagemaker_session.default_bucket()
# Retrieve the AWS account ID
account_id = sagemaker_session.boto_session.client('sts').get_caller_identity()['Account']
# Get the default AWS region where resources will be created
region = sagemaker_session.boto_region_name

import sagemaker

# Create a SageMaker session
sagemaker_session = sagemaker.Session()
# Get the default S3 bucket
default_bucket = sagemaker_session.default_bucket()
# Retrieve the AWS account ID
account_id = sagemaker_session.boto_session.client('sts').get_caller_identity()['Account']
# Get the default AWS region where resources will be created
region = sagemaker_session.boto_region_name

Python

# S3 URI for the training dataset
S3_TRAIN_DATA_URI = f"s3://{default_bucket}/datasets/california_housing_train.csv"
# IAM role ARN for SageMaker to access AWS resources
SAGEMAKER_ROLE = f"arn:aws:iam::{account_id}:role/SageMakerDefaultExecution"
# S3 path to store the trained model artifacts
MODEL_OUTPUT_PATH = f"s3://{default_bucket}/models/california-housing/"
# Type of EC2 instance to use for training
INSTANCE_TYPE = "ml.m5.large"
# Number of instances to use for the training job
INSTANCE_COUNT = 1
# Size (in GB) of the EBS volume attached to the training instance(s)
VOLUME_SIZE_GB = 30

# S3 URI for the training dataset
S3_TRAIN_DATA_URI = f"s3://{default_bucket}/datasets/california_housing_train.csv"
# IAM role ARN for SageMaker to access AWS resources
SAGEMAKER_ROLE = f"arn:aws:iam::{account_id}:role/SageMakerDefaultExecution"
# S3 path to store the trained model artifacts
MODEL_OUTPUT_PATH = f"s3://{default_bucket}/models/california-housing/"
# Type of EC2 instance to use for training
INSTANCE_TYPE = "ml.m5.large"
# Number of instances to use for the training job
INSTANCE_COUNT = 1
# Size (in GB) of the EBS volume attached to the training instance(s)
VOLUME_SIZE_GB = 30

Python

# Define the scikit-learn training image URI
sklearn_image = sagemaker.image_uris.retrieve(
    framework="sklearn",         # Specify the ML framework (scikit-learn)
    region=region,               # AWS region for the image
    version="1.2-1",             # scikit-learn version to use
    py_version="py3",            # Python version for the container
    instance_type=INSTANCE_TYPE  # Instance type for compatibility (required by image_uris.retrieve)
)

# Define the scikit-learn training image URI
sklearn_image = sagemaker.image_uris.retrieve(
    framework="sklearn",         # Specify the ML framework (scikit-learn)
    region=region,               # AWS region for the image
    version="1.2-1",             # scikit-learn version to use
    py_version="py3",            # Python version for the container
    instance_type=INSTANCE_TYPE  # Instance type for compatibility (required by image_uris.retrieve)
)

Python

from sagemaker.modules.configs import SourceCode

# Specify the source code configuration
source_code = SourceCode(
    source_dir=".",           # Directory containing the training script
    entry_script="train.py",  # The main script SageMaker will run for training
)

from sagemaker.modules.configs import SourceCode

# Specify the source code configuration
source_code = SourceCode(
    source_dir=".",           # Directory containing the training script
    entry_script="train.py",  # The main script SageMaker will run for training
)

Python

from sagemaker.modules.configs import Compute

# Define the compute configuration for ModelTrainer
compute_config = Compute(
    instance_type=INSTANCE_TYPE,      # Type of EC2 instance to use for training
    instance_count=INSTANCE_COUNT,    # Number of instances to launch for the training job
    volume_size_in_gb=VOLUME_SIZE_GB  # Size (in GB) of the EBS volume attached to the instance
)

from sagemaker.modules.configs import Compute

# Define the compute configuration for ModelTrainer
compute_config = Compute(
    instance_type=INSTANCE_TYPE,      # Type of EC2 instance to use for training
    instance_count=INSTANCE_COUNT,    # Number of instances to launch for the training job
    volume_size_in_gb=VOLUME_SIZE_GB  # Size (in GB) of the EBS volume attached to the instance
)

Python

from sagemaker.modules.configs import OutputDataConfig

# Define the output data configuration for ModelTrainer
output_config = OutputDataConfig(
    s3_output_path=MODEL_OUTPUT_PATH
)

from sagemaker.modules.configs import OutputDataConfig

# Define the output data configuration for ModelTrainer
output_config = OutputDataConfig(
    s3_output_path=MODEL_OUTPUT_PATH
)

Python

from sagemaker.modules.train import ModelTrainer

# Initialize the ModelTrainer with correct parameters
model_trainer = ModelTrainer(
    training_image=sklearn_image,         # Docker image URI for the scikit-learn training container
    source_code=source_code,              # Source code configuration (directory and entry script)
    base_job_name="sklearn-modeltrainer", # Base name for the SageMaker training job
    role=SAGEMAKER_ROLE,                  # IAM role for SageMaker to access AWS resources
    compute=compute_config,               # Compute configuration (instance type, count, volume size)
    output_data_config=output_config      # Output configuration (S3 path for model artifacts)
)

from sagemaker.modules.train import ModelTrainer

# Initialize the ModelTrainer with correct parameters
model_trainer = ModelTrainer(
    training_image=sklearn_image,         # Docker image URI for the scikit-learn training container
    source_code=source_code,              # Source code configuration (directory and entry script)
    base_job_name="sklearn-modeltrainer", # Base name for the SageMaker training job
    role=SAGEMAKER_ROLE,                  # IAM role for SageMaker to access AWS resources
    compute=compute_config,               # Compute configuration (instance type, count, volume size)
    output_data_config=output_config      # Output configuration (S3 path for model artifacts)
)

text

[07/23/25 10:02:49] WARNING  SageMaker session not          model_trainer.py:461
                             provided. Using default                            
                             Session.                                           
                    WARNING  StoppingCondition not          model_trainer.py:502
                             provided. Using default:                           
                             max_runtime_in_seconds=3600                        
                             max_wait_time_in_seconds=None                      
                             max_pending_time_in_seconds=No                     
                             ne                                                 
                    WARNING  OutputDataConfig compression   model_trainer.py:560
                             type not provided. Using                           
                             default:                                           
                             GZIP                                               
                    INFO     Training image URI:            model_trainer.py:566
                             257758044811.dkr.ecr.us-east-1                     
                             .amazonaws.com/sagemaker-sciki                     
                             t-learn:1.2-1-cpu-py3

[07/23/25 10:02:49] WARNING  SageMaker session not          model_trainer.py:461
                             provided. Using default                            
                             Session.                                           
                    WARNING  StoppingCondition not          model_trainer.py:502
                             provided. Using default:                           
                             max_runtime_in_seconds=3600                        
                             max_wait_time_in_seconds=None                      
                             max_pending_time_in_seconds=No                     
                             ne                                                 
                    WARNING  OutputDataConfig compression   model_trainer.py:560
                             type not provided. Using                           
                             default:                                           
                             GZIP                                               
                    INFO     Training image URI:            model_trainer.py:566
                             257758044811.dkr.ecr.us-east-1                     
                             .amazonaws.com/sagemaker-sciki                     
                             t-learn:1.2-1-cpu-py3

Python

from sagemaker.modules.configs import InputData

# Define the input data configuration
input_data = [
    InputData(
        channel_name="train",           # Name of the input channel for training data
        data_source=S3_TRAIN_DATA_URI,  # S3 URI where the training data is stored
    )
]

from sagemaker.modules.configs import InputData

# Define the input data configuration
input_data = [
    InputData(
        channel_name="train",           # Name of the input channel for training data
        data_source=S3_TRAIN_DATA_URI,  # S3 URI where the training data is stored
    )
]

Python

# Call train() to start the training job asynchronously (does not block)
model_trainer.train(input_data_config=input_data, wait=False)

# Call train() to start the training job asynchronously (does not block)
model_trainer.train(input_data_config=input_data, wait=False)

text

                    WARNING  key_prefix is only applicable  model_trainer.py:804
                             when data_source is a local                        
                             file path.                                         
[07/23/25 10:03:57] INFO     Creating training_job resource.  resources.py:28474
[07/23/25 10:03:57] WARNING  No region provided. Using default      utils.py:343
                             region.                                            
                    WARNING  No config provided. Using default      utils.py:347
                             config.                                            
[07/23/25 10:03:58] WARNING  Not displaing the training     model_trainer.py:744
                             container logs as 'wait' is                        
                             set to False.

                    WARNING  key_prefix is only applicable  model_trainer.py:804
                             when data_source is a local                        
                             file path.                                         
[07/23/25 10:03:57] INFO     Creating training_job resource.  resources.py:28474
[07/23/25 10:03:57] WARNING  No region provided. Using default      utils.py:343
                             region.                                            
                    WARNING  No config provided. Using default      utils.py:347
                             config.                                            
[07/23/25 10:03:58] WARNING  Not displaing the training     model_trainer.py:744
                             container logs as 'wait' is                        
                             set to False.

Python

# Access training job information
latest_job = model_trainer._latest_training_job

# Print the training job name and status
print(f"Training job name: {latest_job.training_job_name}")
print(f"Training job status: {latest_job.training_job_status}")

# Access training job information
latest_job = model_trainer._latest_training_job

# Print the training job name and status
print(f"Training job name: {latest_job.training_job_name}")
print(f"Training job status: {latest_job.training_job_status}")

text

Training job name: sklearn-modeltrainer-20250723100358
Training job status: InProgress

Training job name: sklearn-modeltrainer-20250723100358
Training job status: InProgress

Python

# List most recent completed training job
training_jobs = sagemaker_session.sagemaker_client.list_training_jobs(
    SortBy='CreationTime',
    SortOrder='Descending',
    StatusEquals='Completed',
    NameContains='sklearn-modeltrainer'  # Filter for ModelTrainer jobs
)

# Extract the name of the latest training job of the list
TRAINING_JOB_NAME = training_jobs['TrainingJobSummaries'][0]['TrainingJobName']

# Path to test data
TEST_DATA_FILE = "data/california_housing_test.csv"

# Since ModelTrainer doesn't have attach() method,
# we use SageMaker Session to get training job details
training_job_details = sagemaker_session.describe_training_job(TRAINING_JOB_NAME)
    
# Get model S3 location
model_s3_uri = training_job_details['ModelArtifacts']['S3ModelArtifacts']

# List most recent completed training job
training_jobs = sagemaker_session.sagemaker_client.list_training_jobs(
    SortBy='CreationTime',
    SortOrder='Descending',
    StatusEquals='Completed',
    NameContains='sklearn-modeltrainer'  # Filter for ModelTrainer jobs
)

# Extract the name of the latest training job of the list
TRAINING_JOB_NAME = training_jobs['TrainingJobSummaries'][0]['TrainingJobName']

# Path to test data
TEST_DATA_FILE = "data/california_housing_test.csv"

# Since ModelTrainer doesn't have attach() method,
# we use SageMaker Session to get training job details
training_job_details = sagemaker_session.describe_training_job(TRAINING_JOB_NAME)
    
# Get model S3 location
model_s3_uri = training_job_details['ModelArtifacts']['S3ModelArtifacts']

Python

import sagemaker

# Create a SageMaker session
sagemaker_session = sagemaker.Session()
# Get the default S3 bucket
default_bucket = sagemaker_session.default_bucket()
# Retrieve the AWS account ID
account_id = sagemaker_session.boto_session.client('sts').get_caller_identity()['Account']
# Get the default AWS region where resources will be created
region = sagemaker_session.boto_region_name

Python

# S3 URI for the training dataset
S3_TRAIN_DATA_URI = f"s3://{default_bucket}/datasets/california_housing_train.csv"
# IAM role ARN for SageMaker to access AWS resources
SAGEMAKER_ROLE = f"arn:aws:iam::{account_id}:role/SageMakerDefaultExecution"
# S3 path to store the trained model artifacts
MODEL_OUTPUT_PATH = f"s3://{default_bucket}/models/california-housing/"
# Type of EC2 instance to use for training
INSTANCE_TYPE = "ml.m5.large"
# Number of instances to use for the training job
INSTANCE_COUNT = 1
# Size (in GB) of the EBS volume attached to the training instance(s)
VOLUME_SIZE_GB = 30

Python

# Define the scikit-learn training image URI
sklearn_image = sagemaker.image_uris.retrieve(
    framework="sklearn",         # Specify the ML framework (scikit-learn)
    region=region,               # AWS region for the image
    version="1.2-1",             # scikit-learn version to use
    py_version="py3",            # Python version for the container
    instance_type=INSTANCE_TYPE  # Instance type for compatibility (required by image_uris.retrieve)
)

Python

from sagemaker.modules.configs import SourceCode

# Specify the source code configuration
source_code = SourceCode(
    source_dir=".",           # Directory containing the training script
    entry_script="train.py",  # The main script SageMaker will run for training
)

Python

from sagemaker.modules.configs import Compute

# Define the compute configuration for ModelTrainer
compute_config = Compute(
    instance_type=INSTANCE_TYPE,      # Type of EC2 instance to use for training
    instance_count=INSTANCE_COUNT,    # Number of instances to launch for the training job
    volume_size_in_gb=VOLUME_SIZE_GB  # Size (in GB) of the EBS volume attached to the instance
)

Python

from sagemaker.modules.configs import OutputDataConfig

# Define the output data configuration for ModelTrainer
output_config = OutputDataConfig(
    s3_output_path=MODEL_OUTPUT_PATH
)

Python

from sagemaker.modules.train import ModelTrainer

# Initialize the ModelTrainer with correct parameters
model_trainer = ModelTrainer(
    training_image=sklearn_image,         # Docker image URI for the scikit-learn training container
    source_code=source_code,              # Source code configuration (directory and entry script)
    base_job_name="sklearn-modeltrainer", # Base name for the SageMaker training job
    role=SAGEMAKER_ROLE,                  # IAM role for SageMaker to access AWS resources
    compute=compute_config,               # Compute configuration (instance type, count, volume size)
    output_data_config=output_config      # Output configuration (S3 path for model artifacts)
)

text

[07/23/25 10:02:49] WARNING  SageMaker session not          model_trainer.py:461
                             provided. Using default                            
                             Session.                                           
                    WARNING  StoppingCondition not          model_trainer.py:502
                             provided. Using default:                           
                             max_runtime_in_seconds=3600                        
                             max_wait_time_in_seconds=None                      
                             max_pending_time_in_seconds=No                     
                             ne                                                 
                    WARNING  OutputDataConfig compression   model_trainer.py:560
                             type not provided. Using                           
                             default:                                           
                             GZIP                                               
                    INFO     Training image URI:            model_trainer.py:566
                             257758044811.dkr.ecr.us-east-1                     
                             .amazonaws.com/sagemaker-sciki                     
                             t-learn:1.2-1-cpu-py3

Python

from sagemaker.modules.configs import InputData

# Define the input data configuration
input_data = [
    InputData(
        channel_name="train",           # Name of the input channel for training data
        data_source=S3_TRAIN_DATA_URI,  # S3 URI where the training data is stored
    )
]

Python

# Call train() to start the training job asynchronously (does not block)
model_trainer.train(input_data_config=input_data, wait=False)

text

                    WARNING  key_prefix is only applicable  model_trainer.py:804
                             when data_source is a local                        
                             file path.                                         
[07/23/25 10:03:57] INFO     Creating training_job resource.  resources.py:28474
[07/23/25 10:03:57] WARNING  No region provided. Using default      utils.py:343
                             region.                                            
                    WARNING  No config provided. Using default      utils.py:347
                             config.                                            
[07/23/25 10:03:58] WARNING  Not displaing the training     model_trainer.py:744
                             container logs as 'wait' is                        
                             set to False.

Python

# Access training job information
latest_job = model_trainer._latest_training_job

# Print the training job name and status
print(f"Training job name: {latest_job.training_job_name}")
print(f"Training job status: {latest_job.training_job_status}")

text

Training job name: sklearn-modeltrainer-20250723100358
Training job status: InProgress

Python

# List most recent completed training job
training_jobs = sagemaker_session.sagemaker_client.list_training_jobs(
    SortBy='CreationTime',
    SortOrder='Descending',
    StatusEquals='Completed',
    NameContains='sklearn-modeltrainer'  # Filter for ModelTrainer jobs
)

# Extract the name of the latest training job of the list
TRAINING_JOB_NAME = training_jobs['TrainingJobSummaries'][0]['TrainingJobName']

# Path to test data
TEST_DATA_FILE = "data/california_housing_test.csv"

# Since ModelTrainer doesn't have attach() method,
# we use SageMaker Session to get training job details
training_job_details = sagemaker_session.describe_training_job(TRAINING_JOB_NAME)
    
# Get model S3 location
model_s3_uri = training_job_details['ModelArtifacts']['S3ModelArtifacts']

Introduction & Lesson Overview

Understanding ModelTrainer vs Estimators

Join the 1M+ learners on CodeSignal

Environment Setup and Configuration

Defining Configuration Constants

Defining the Training Container Image

Configuring Your Training Script

Setting Up Compute Resources

Configuring Model Output Location

Creating the ModelTrainer Instance

Defining Input Data Configuration

Launching the Asynchronous Training Job

Retrieving Training Job Information

Retrieving Model Artifacts from Completed Jobs

Summary & Next Steps