Training Models with SageMaker Estimators

Introduction & Overview

Welcome back! In the previous lessons, you've built a solid foundation for working with SageMaker and successfully uploaded your training data to S3 . Now you're ready to take the next exciting step: actually training your first machine learning model in the cloud. A training job is simply the process of running your machine learning code on cloud infrastructure to create a trained model. Instead of running your training script on your local computer, SageMaker executes it on powerful AWS servers, automatically managing all the infrastructure details like spinning up compute resources, downloading your data, and saving your results back to S3. In this lesson, you'll learn how to use SageMaker estimators to launch training jobs. Since we've been working with scikit-learn for our machine learning code, we'll focus on the SKLearn estimator . Of course, if you're more comfortable with TensorFlow, PyTorch, or other frameworks, SageMaker has dedicated estimators for those too, following the same patterns you'll master here. By the end of this lesson, you'll have launched your first SageMaker training job, configured compute resources, and monitored the training process from start to finish. This represents a major milestone in your machine learning journey, as you'll be running real training jobs in AWS's cloud infrastructure using the same scikit-learn skills you've already developed.

Understanding SageMaker Estimators

An estimator in SageMaker is a high-level interface that handles the complexity of launching and managing training jobs in the cloud. Think of an estimator as our control center for training — it knows how to package our code, spin up the right computing resources, run our training script, and save the results back to S3. SageMaker provides several types of estimators to match different machine learning frameworks and use cases. It provides framework-specific estimators like: SKLearn (for scikit-learn) TensorFlow PyTorch XGBoost And many others Additionally, SageMaker offers generic estimators for custom Docker containers. Each estimator is optimized for its respective framework, providing the right environment and dependencies out of the box.

Working with the SKLearn Estimator

For this lesson, we'll focus on the SKLearn estimator since our machine learning code was developed using the popular scikit-learn library. However, if you prefer to work with other frameworks like TensorFlow or PyTorch, SageMaker will be there for you with dedicated estimators that follow the same patterns we'll learn here. The SKLearn estimator is specifically designed for training machine learning models using scikit-learn. What makes it particularly powerful is that it allows us to bring our own custom training code while SageMaker handles all the infrastructure management. Instead of being limited to pre-built algorithms, we can write our training logic exactly how we want it and let SageMaker execute it at scale. Behind the scenes, when we use the SKLearn estimator, SageMaker essentially creates a containerized environment with scikit-learn and Python pre-installed, similar to how we might use a Docker image with our dependencies already configured. However, we don't need to worry about container management, image building, or orchestration — SageMaker handles all of that complexity for us. It automatically sets up the training environment, runs our custom Python script inside that environment, and manages all the cloud resources needed for training. This gives us the flexibility of custom code with the power and convenience of cloud infrastructure. We get the best of both worlds: complete control over our machine learning logic and the ability to scale our training jobs without managing servers or containers ourselves.

Join the 1M+ learners on CodeSignal

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

Creating Your Custom Training Script

Before configuring SageMaker to run our training job, we need to create the training script that contains our machine learning logic. This script needs to be adapted to work within SageMaker's environment while maintaining the familiar machine learning patterns we already know.

In this case, we'll name our training script train.py. This file will contain our custom machine learning logic, and while the core concepts remain familiar, the script needs to be specifically adapted to work within SageMaker's environment:

Python

import os
import argparse
import joblib
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

def model_fn(model_dir):
    """Load model for inference - SageMaker will handle the rest"""
    # Load the saved model from the model directory using joblib
    # SageMaker expects the model to be saved as 'model.joblib'
    model = joblib.load(os.path.join(model_dir, 'model.joblib'))
    # Return the loaded model so SageMaker can use it for predictions
    return model

import os
import argparse
import joblib
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

def model_fn(model_dir):
    """Load model for inference - SageMaker will handle the rest"""
    # Load the saved model from the model directory using joblib
    # SageMaker expects the model to be saved as 'model.joblib'
    model = joblib.load(os.path.join(model_dir, 'model.joblib'))
    # Return the loaded model so SageMaker can use it for predictions
    return model

The model_fn function is a special function that SageMaker looks for when we later want to deploy our model for making predictions. This function tells SageMaker how to load our saved model from the artifacts directory. Even though we're not deploying the model in this lesson, including this function makes our training script deployment-ready for future use. SageMaker will call this function automatically when setting up inference endpoints, so it needs to know exactly how to reconstruct our trained model from the saved files.

Handling SageMaker Arguments and Environment

The script uses a specific pattern to receive information from SageMaker about where to find data and where to save results:

Python

if __name__ == '__main__':
    # Create argument parser to handle SageMaker's environment variables
    parser = argparse.ArgumentParser()
    # Add argument for model directory - where SageMaker expects us to save the trained model
    parser.add_argument('--model-dir', type=str, default=os.environ.get('SM_MODEL_DIR'))
    # Add argument for training data directory - where SageMaker downloads our training data
    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAIN'))
    # Parse all the arguments that SageMaker provides
    args = parser.parse_args()

if __name__ == '__main__':
    # Create argument parser to handle SageMaker's environment variables
    parser = argparse.ArgumentParser()
    # Add argument for model directory - where SageMaker expects us to save the trained model
    parser.add_argument('--model-dir', type=str, default=os.environ.get('SM_MODEL_DIR'))
    # Add argument for training data directory - where SageMaker downloads our training data
    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAIN'))
    # Parse all the arguments that SageMaker provides
    args = parser.parse_args()

The if __name__ == '__main__': block is essential because SageMaker will import our script as a module when it needs to use the model_fn function for inference. Without this guard, all our training code would execute every time SageMaker imports the script, which isn't what we want. The guard ensures that the training logic only runs when the script is executed directly during the training job.

The key difference from standalone machine learning scripts is how it receives information about data locations and where to save results. Instead of hardcoding file paths, the script uses argparse to read arguments that SageMaker automatically provides through environment variables:

SM_MODEL_DIR — SageMaker sets this environment variable to point to a directory (like /opt/ml/model/) on the training instance where our script should save the trained model. After training completes, SageMaker automatically uploads everything in this directory to S3 as model artifacts.
SM_CHANNEL_TRAIN — SageMaker sets this to point to a directory (like /opt/ml/input/data/train/) on the training instance where it has already downloaded our training data from S3. The "TRAIN" part must match the channel name we'll specify later when we launch the training job.

Think of these environment variables as SageMaker's way of leaving "notes" for our script about the local file system layout. When our training job runs in the cloud, SageMaker handles downloading data from S3 to local storage and sets up these paths so our script knows exactly where everything is located.

The Training Logic Implementation

The actual training logic within the if __name__ == '__main__': block follows familiar machine learning patterns, but adapted to work with SageMaker's file system:

Python

if __name__ == '__main__':
    # Parse SageMaker arguments
    # ...
    
    # Load the training data from S3
    train_data_path = os.path.join(args.train, 'california_housing_train.csv')
    df = pd.read_csv(train_data_path)
    
    # Separate features and target variable
    X_train = df.drop("MedHouseVal", axis=1)  # Features
    y_train = df["MedHouseVal"]               # Target: median house value
    
    print(f"Starting training with {len(X_train)} examples...")
    print(f"Features: {list(X_train.columns)}")
    
    # Create and train the model directly
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    # Evaluate model performance on training data
    print(f"Starting evaluation...")
    y_train_pred = model.predict(X_train)
    train_r2 = r2_score(y_train, y_train_pred)
    train_rmse = np.sqrt(mean_squared_error(y_train, y_train_pred))
    
    # Print training performance metrics
    print(f"R² Score: {train_r2:.4f}")
    print(f"RMSE: {train_rmse:.4f}")
    
    # Save the model for SageMaker deployment
    joblib.dump(model, os.path.join(args.model_dir, 'model.joblib'))
    print("Training completed!")

if __name__ == '__main__':
    # Parse SageMaker arguments
    # ...
    
    # Load the training data from S3
    train_data_path = os.path.join(args.train, 'california_housing_train.csv')
    df = pd.read_csv(train_data_path)
    
    # Separate features and target variable
    X_train = df.drop("MedHouseVal", axis=1)  # Features
    y_train = df["MedHouseVal"]               # Target: median house value
    
    print(f"Starting training with {len(X_train)} examples...")
    print(f"Features: {list(X_train.columns)}")
    
    # Create and train the model directly
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    # Evaluate model performance on training data
    print(f"Starting evaluation...")
    y_train_pred = model.predict(X_train)
    train_r2 = r2_score(y_train, y_train_pred)
    train_rmse = np.sqrt(mean_squared_error(y_train, y_train_pred))
    
    # Print training performance metrics
    print(f"R² Score: {train_r2:.4f}")
    print(f"RMSE: {train_rmse:.4f}")
    
    # Save the model for SageMaker deployment
    joblib.dump(model, os.path.join(args.model_dir, 'model.joblib'))
    print("Training completed!")

This section contains the familiar machine learning workflow we've seen before, but with key adaptations for the SageMaker environment:

Loading data — Uses the parsed args.train parameter to find where SageMaker downloaded our training data from S3
Training the model — Follows the standard scikit-learn pattern of creating a model instance and calling fit()
Evaluating performance — Calculates metrics to understand how well our model learned from the training data
Saving results — Uses args.model_dir to save the trained model where SageMaker expects to find it

The main difference from standalone machine learning scripts is using the parsed arguments to determine file paths rather than hardcoded locations. SageMaker automatically downloads our training data from S3 to the local file system and provides the path through the args.train parameter. Similarly, args.model_dir points to where SageMaker expects us to save our trained model artifacts, which it will then automatically upload back to S3 when training completes.

Creating the SKLearn Estimator

Now we're ready to create the SKLearn estimator that will bring together all the configuration parameters we've just defined. The estimator acts as the bridge between our setup (the session, data locations, compute resources, and permissions) and our actual training execution:

Python

from sagemaker.sklearn.estimator import SKLearn

# Create SKLearn estimator
sklearn_estimator = SKLearn(
    entry_point='train.py',              # Python script containing training code
    role=SAGEMAKER_ROLE,                 # IAM role for SageMaker to access AWS resources
    instance_type=INSTANCE_TYPE,         # Type of EC2 instance for training
    instance_count=INSTANCE_COUNT,       # Number of instances to use
    framework_version='1.2-1',           # Version of scikit-learn to use
    py_version='py3',                    # Python version for the training environment
    script_mode=True,                    # Use script mode for training (vs legacy mode)
    sagemaker_session=sagemaker_session, # Session for interacting with SageMaker
    output_path=MODEL_OUTPUT_PATH        # S3 location to save model artifacts
)

from sagemaker.sklearn.estimator import SKLearn

# Create SKLearn estimator
sklearn_estimator = SKLearn(
    entry_point='train.py',              # Python script containing training code
    role=SAGEMAKER_ROLE,                 # IAM role for SageMaker to access AWS resources
    instance_type=INSTANCE_TYPE,         # Type of EC2 instance for training
    instance_count=INSTANCE_COUNT,       # Number of instances to use
    framework_version='1.2-1',           # Version of scikit-learn to use
    py_version='py3',                    # Python version for the training environment
    script_mode=True,                    # Use script mode for training (vs legacy mode)
    sagemaker_session=sagemaker_session, # Session for interacting with SageMaker
    output_path=MODEL_OUTPUT_PATH        # S3 location to save model artifacts
)

Let's break down all the key parameters that configure our SKLearn estimator:

entry_point — Specifies the Python file containing our training logic
role — The IAM role that grants SageMaker permission to access our AWS resources like S3 buckets
instance_type — Defines the type of EC2 instance (virtual machine) that will run our training job
instance_count — Sets the number of instances to use for training (we'll start with 1 for learning)
framework_version — Ensures we're using a specific version of scikit-learn, helping maintain consistency across different training runs. (At the time this course was developed, '1.2-1' is the latest version available in SageMaker, even though newer versions may exist on PyPI or elsewhere.)
py_version — Specifies the Python version for the training environment
script_mode=True — Tells SageMaker to use the modern script mode approach, giving us more flexibility to write training scripts that feel like regular Python programs rather than following rigid templates
sagemaker_session — The session object that handles communication with the SageMaker service
output_path — The S3 location where SageMaker will save our trained model artifacts after training completes

Together, these parameters give SageMaker everything it needs to create the right environment, run our code, and manage the entire training process from start to finish.

Getting the Training Job Name

Once the training job starts, you'll want to know its unique name so you can track it later. Every SageMaker training job gets a unique identifier that includes a timestamp, making it easy to distinguish between different training runs.

To retrieve and display your training job's name, you can access it through the estimator's latest training job object:

Python

# Extract and display the training job name for tracking and later status checks
training_job_name = sklearn_estimator.latest_training_job.name
print(f"Training job started: {training_job_name}")

# Extract and display the training job name for tracking and later status checks
training_job_name = sklearn_estimator.latest_training_job.name
print(f"Training job started: {training_job_name}")

This will output something like:

text

Training job started: sagemaker-scikit-learn-2025-07-22-08-42-55-894

Training job started: sagemaker-scikit-learn-2025-07-22-08-42-55-894

Notice how SageMaker automatically generated this name following a specific template for estimator-based training jobs:

text

sagemaker-<framework>-<YYYY>-<MM>-<DD>-<HH>-<MM>-<SS>-<milliseconds>

sagemaker-<framework>-<YYYY>-<MM>-<DD>-<HH>-<MM>-<SS>-<milliseconds>

Looking at our example sagemaker-scikit-learn-2025-07-22-08-42-55-894:

sagemaker — A fixed prefix for all estimator-based training jobs
scikit-learn — The framework specified by your estimator (would be tensorflow, pytorch, xgboost, etc. for other framework estimators)
2025-07-22 — The date when the job was created (year-month-day)
08-42-55 — The time when the job was created (hour-minute-second in UTC)
894 — Additional milliseconds to ensure uniqueness

This naming convention makes it easy to identify when a job was created and what framework it used, while the timestamp ensures each job has a completely unique identifier. You can use this name to reference this specific training job later, even if you close your notebook and come back to it. Think of it as a receipt number that lets you check on your training job's progress at any time.

Introduction & Overview

Understanding SageMaker Estimators

Working with the SKLearn Estimator

Join the 1M+ learners on CodeSignal

Creating Your Custom Training Script

Handling SageMaker Arguments and Environment

The Training Logic Implementation

Setting Up Your SageMaker Session and Resources

Defining Data Locations and Paths

Configuring Compute Resources

Setting Up AWS Account Information and Permissions

Creating the SKLearn Estimator

Launching the Training Job

Getting the Training Job Name

Checking the Job Status

Understanding Model Artifacts

Summary & Next Steps